From b0feef8bb9226b2d47340b5ceaf2ee61b426f842 Mon Sep 17 00:00:00 2001 From: nihalpasham Date: Fri, 28 Aug 2026 22:18:56 +0530 Subject: [PATCH 01/18] refactor(dialect-mir): split tests/ops_test.rs into ops_test/ modules --- crates/dialect-mir/tests/ops_test.rs | 4251 ----------------- .../dialect-mir/tests/ops_test/aggregates.rs | 687 +++ .../tests/ops_test/cast_authority.rs | 1060 ++++ .../tests/ops_test/control_flow.rs | 175 + crates/dialect-mir/tests/ops_test/enums.rs | 406 ++ crates/dialect-mir/tests/ops_test/main.rs | 16 + crates/dialect-mir/tests/ops_test/memory.rs | 340 ++ .../tests/ops_test/pointer_kinds.rs | 1256 +++++ .../tests/ops_test/scalars_and_funcs.rs | 447 ++ 9 files changed, 4387 insertions(+), 4251 deletions(-) delete mode 100644 crates/dialect-mir/tests/ops_test.rs create mode 100644 crates/dialect-mir/tests/ops_test/aggregates.rs create mode 100644 crates/dialect-mir/tests/ops_test/cast_authority.rs create mode 100644 crates/dialect-mir/tests/ops_test/control_flow.rs create mode 100644 crates/dialect-mir/tests/ops_test/enums.rs create mode 100644 crates/dialect-mir/tests/ops_test/main.rs create mode 100644 crates/dialect-mir/tests/ops_test/memory.rs create mode 100644 crates/dialect-mir/tests/ops_test/pointer_kinds.rs create mode 100644 crates/dialect-mir/tests/ops_test/scalars_and_funcs.rs diff --git a/crates/dialect-mir/tests/ops_test.rs b/crates/dialect-mir/tests/ops_test.rs deleted file mode 100644 index b235dd9c09..0000000000 --- a/crates/dialect-mir/tests/ops_test.rs +++ /dev/null @@ -1,4251 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -// Tests build kinded fixture types directly; production minting lives in -// mir-importer's facts.rs (see the workspace clippy.toml disallowed-methods). -#![allow(clippy::disallowed_methods)] - -use dialect_mir::{ - attributes::{FieldIndexAttr, MirCastKindAttr, MirPointerKindAuthorityAttr, VariantIndexAttr}, - ops::{ - MirAddOp, MirAllocaOp, MirArrayElementAddrOp, MirAssertOp, MirAssignOp, MirCallOp, - MirCastOp, MirCheckedAddOp, MirCmpOp, MirCondBranchOp, MirConstantOp, - MirConstructDisjointSliceOp, MirConstructEnumOp, MirConstructSliceOp, MirDivOp, - MirEnumPayloadOp, MirEqOp, MirExtractFieldOp, MirFieldAddrOp, MirFuncOp, MirGeOp, - MirGetDiscriminantOp, MirGlobalAllocOp, MirGotoOp, MirGtOp, MirLeOp, MirLoadOp, MirLtOp, - MirMulOp, MirNeOp, MirNegOp, MirNotOp, MirPtrOffsetOp, MirRefOp, MirRemOp, MirReturnOp, - MirSetDiscriminantOp, MirSharedAllocOp, MirStoreOp, MirSubOp, - }, - types::{ - EnumVariant, MirArrayType, MirDisjointSliceType, MirEnumType, MirPointerKind, MirPtrType, - MirSliceType, MirStructType, MirTupleType, MirUnionType, - }, -}; -use pliron::{ - basic_block::BasicBlock, - builtin::{ - attributes::{IntegerAttr, StringAttr, TypeAttr}, - op_interfaces::OperandSegmentInterface, - types::{FP32Type, FunctionType, IntegerType, Signedness}, - }, - common_traits::Verify, - context::Context, - op::Op, - operation::Operation, - opts::mem2reg::{AllocInfo, PromotableOpInterface, PromotableOpKind}, - r#type::TypeHandle, - utils::apint::APInt, -}; -use std::num::NonZeroUsize; - -#[test] -fn test_mir_control_flow_verify() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); - - // 1. MirGotoOp - let target_block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); - let src_block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); - let arg_val = src_block.deref(&ctx).get_argument(0); - - let op = Operation::new( - &mut ctx, - MirGotoOp::get_concrete_op_info(), - vec![], - vec![arg_val], - vec![target_block], - 0, - ); - let goto_op = MirGotoOp::new(op); - assert!(goto_op.verify(&ctx).is_ok(), "Valid Goto"); - - let op_bad = Operation::new( - &mut ctx, - MirGotoOp::get_concrete_op_info(), - vec![], - vec![], - vec![target_block], - 0, - ); - assert!( - MirGotoOp::new(op_bad).verify(&ctx).is_err(), - "Goto missing operand" - ); - - // 2. MirCondBranchOp - let true_block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); - let false_block = BasicBlock::new(&mut ctx, None, vec![]); - let cond_block = BasicBlock::new(&mut ctx, None, vec![i1_ty.into(), i32_ty.into()]); - let cond_val = cond_block.deref(&ctx).get_argument(0); - let val = cond_block.deref(&ctx).get_argument(1); - - let (cond_flat, cond_sizes) = - MirCondBranchOp::compute_segment_sizes(vec![vec![cond_val], vec![val], vec![]]); - let op_cond = Operation::new( - &mut ctx, - MirCondBranchOp::get_concrete_op_info(), - vec![], - cond_flat, - vec![true_block, false_block], - 0, - ); - MirCondBranchOp::new(op_cond).set_operand_segment_sizes(&ctx, cond_sizes); - let cond_br = MirCondBranchOp::new(op_cond); - assert!(cond_br.verify(&ctx).is_ok(), "Valid CondBranch"); - - let (cond_bad_flat, cond_bad_sizes) = - MirCondBranchOp::compute_segment_sizes(vec![vec![cond_val], vec![], vec![]]); - let op_cond_bad = Operation::new( - &mut ctx, - MirCondBranchOp::get_concrete_op_info(), - vec![], - cond_bad_flat, - vec![true_block, false_block], - 0, - ); - MirCondBranchOp::new(op_cond_bad).set_operand_segment_sizes(&ctx, cond_bad_sizes); - assert!( - MirCondBranchOp::new(op_cond_bad).verify(&ctx).is_err(), - "CondBranch missing operand" - ); - - // 3. MirReturnOp - let func_ty = FunctionType::get(&ctx, vec![], vec![i32_ty.into()]); - let func_ty_attr = TypeAttr::new(func_ty.into()); - - let func_op_ptr = Operation::new( - &mut ctx, - MirFuncOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 1, - ); - let mir_func = MirFuncOp::new(&mut ctx, func_op_ptr, func_ty_attr); - let entry_block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); - let ret_val = entry_block.deref(&ctx).get_argument(0); - - let region = mir_func.get_operation().deref(&ctx).get_region(0); - entry_block.insert_at_front(region, &ctx); - - let ret_op = Operation::new( - &mut ctx, - MirReturnOp::get_concrete_op_info(), - vec![], - vec![ret_val], - vec![], - 0, - ); - ret_op.insert_at_back(entry_block, &ctx); - - let mir_ret = MirReturnOp::new(ret_op); - assert!(mir_ret.verify(&ctx).is_ok(), "Valid Return"); - - let f32_ty = FP32Type::get(&ctx); - let f32_block = BasicBlock::new(&mut ctx, None, vec![f32_ty.into()]); - let f32_val = f32_block.deref(&ctx).get_argument(0); - - let ret_op_bad = Operation::new( - &mut ctx, - MirReturnOp::get_concrete_op_info(), - vec![], - vec![f32_val], - vec![], - 0, - ); - ret_op_bad.insert_at_back(entry_block, &ctx); - - let mir_ret_bad = MirReturnOp::new(ret_op_bad); - assert!(mir_ret_bad.verify(&ctx).is_err(), "Return type mismatch"); - - // 4. MirAssertOp - let assert_succ = BasicBlock::new(&mut ctx, None, vec![]); - - let (assert_flat, assert_sizes) = - MirAssertOp::compute_segment_sizes(vec![vec![cond_val], vec![]]); - let op_assert = Operation::new( - &mut ctx, - MirAssertOp::get_concrete_op_info(), - vec![], - assert_flat, - vec![assert_succ], - 0, - ); - MirAssertOp::new(op_assert).set_operand_segment_sizes(&ctx, assert_sizes); - let assert_op = MirAssertOp::new(op_assert); - assert!(assert_op.verify(&ctx).is_ok(), "Valid Assert"); - - let (assert_bad_flat, assert_bad_sizes) = - MirAssertOp::compute_segment_sizes(vec![vec![val], vec![]]); - let op_assert_bad = Operation::new( - &mut ctx, - MirAssertOp::get_concrete_op_info(), - vec![], - assert_bad_flat, - vec![assert_succ], - 0, - ); - MirAssertOp::new(op_assert_bad).set_operand_segment_sizes(&ctx, assert_bad_sizes); - assert!( - MirAssertOp::new(op_assert_bad).verify(&ctx).is_err(), - "Assert cond type mismatch" - ); -} - -#[test] -fn test_mir_pointer_kind_distinguishes_references_from_raw_pointers() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let pointee: pliron::r#type::TypeHandle = i32_ty.into(); - - let shared_ref = - MirPtrType::get_generic_with_kind(&mut ctx, pointee, false, MirPointerKind::SharedRef); - let raw_const = - MirPtrType::get_generic_with_kind(&mut ctx, pointee, false, MirPointerKind::RawConst); - let unique_ref = - MirPtrType::get_generic_with_kind(&mut ctx, pointee, true, MirPointerKind::UniqueRef); - let raw_mut = - MirPtrType::get_generic_with_kind(&mut ctx, pointee, true, MirPointerKind::RawMut); - - assert_ne!( - shared_ref, raw_const, - "&T must remain distinct from *const T" - ); - assert_ne!( - unique_ref, raw_mut, - "&mut T must remain distinct from *mut T" - ); - assert_ne!( - shared_ref, unique_ref, - "&T must remain distinct from &mut T" - ); - assert_ne!( - raw_const, raw_mut, - "*const T must remain distinct from *mut T" - ); - - assert_eq!( - shared_ref.deref(&ctx).pointer_kind(), - MirPointerKind::SharedRef - ); - assert_eq!( - unique_ref.deref(&ctx).pointer_kind(), - MirPointerKind::UniqueRef - ); - assert_eq!( - raw_const.deref(&ctx).pointer_kind(), - MirPointerKind::RawConst - ); - assert_eq!(raw_mut.deref(&ctx).pointer_kind(), MirPointerKind::RawMut); -} - -#[test] -fn test_mir_slice_preserves_pointer_kind() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let u8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); - let element: pliron::r#type::TypeHandle = u8_ty.into(); - - let shared = MirSliceType::get_with_kind(&mut ctx, element, MirPointerKind::SharedRef); - let raw = MirSliceType::get_with_kind(&mut ctx, element, MirPointerKind::RawConst); - - assert_ne!(shared, raw, "&[T] must remain distinct from *const [T]"); - assert_eq!(shared.deref(&ctx).pointer_kind(), MirPointerKind::SharedRef); - assert_eq!(raw.deref(&ctx).pointer_kind(), MirPointerKind::RawConst); -} - -#[test] -fn test_mir_pointer_kind_mutability_consistency_verify() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let invalid_shared = MirPtrType { - pointee: i32_ty.into(), - is_mutable: true, - address_space: 0, - kind: MirPointerKind::SharedRef, - }; - let invalid_raw_mut = MirPtrType { - pointee: i32_ty.into(), - is_mutable: false, - address_space: 0, - kind: MirPointerKind::RawMut, - }; - - assert!(invalid_shared.verify(&ctx).is_err()); - assert!(invalid_raw_mut.verify(&ctx).is_err()); - - let invalid_shared_slice = MirSliceType { - element_ty: i32_ty.into(), - is_mutable: true, - kind: MirPointerKind::SharedRef, - }; - let valid_erased_mut_slice = MirSliceType { - element_ty: i32_ty.into(), - is_mutable: true, - kind: MirPointerKind::Erased, - }; - assert!(invalid_shared_slice.verify(&ctx).is_err()); - assert!(valid_erased_mut_slice.verify(&ctx).is_ok()); -} - -#[test] -fn test_alloca_cannot_claim_a_rust_pointer_kind() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let erased = MirPtrType::get_generic(&mut ctx, i32_ty.into(), true); - let erased_alloca = Operation::new( - &mut ctx, - MirAllocaOp::get_concrete_op_info(), - vec![erased.into()], - vec![], - vec![], - 0, - ); - assert!(MirAllocaOp::new(erased_alloca).verify(&ctx).is_ok()); - - let immutable_erased = MirPtrType::get_generic(&mut ctx, i32_ty.into(), false); - let immutable_alloca = Operation::new( - &mut ctx, - MirAllocaOp::get_concrete_op_info(), - vec![immutable_erased.into()], - vec![], - vec![], - 0, - ); - assert!( - MirAllocaOp::new(immutable_alloca).verify(&ctx).is_err(), - "an alloca cannot masquerade as the immutable canonical function-pointer carrier" - ); - - let shared_erased = MirPtrType::get_shared(&mut ctx, i32_ty.into(), true); - let shared_alloca = Operation::new( - &mut ctx, - MirAllocaOp::get_concrete_op_info(), - vec![shared_erased.into()], - vec![], - vec![], - 0, - ); - assert!( - MirAllocaOp::new(shared_alloca).verify(&ctx).is_err(), - "a stack allocation must remain in generic address space" - ); - - let unique = - MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::UniqueRef); - let unique_alloca = Operation::new( - &mut ctx, - MirAllocaOp::get_concrete_op_info(), - vec![unique.into()], - vec![], - vec![], - 0, - ); - assert!( - MirAllocaOp::new(unique_alloca).verify(&ctx).is_err(), - "compiler storage must not manufacture UniqueRef" - ); -} - -#[test] -fn test_shared_alloc_result_pointee_must_match_element_type() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let i8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Signless).into(); - let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signless).into(); - let usize_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let mismatched_result = MirPtrType::get_shared(&mut ctx, i64_ty, true); - let op = Operation::new( - &mut ctx, - MirSharedAllocOp::get_concrete_op_info(), - vec![mismatched_result.into()], - vec![], - vec![], - 0, - ); - let alloc = MirSharedAllocOp::new(op); - alloc.set_attr_elem_type(&ctx, TypeAttr::new(i8_ty)); - alloc.set_attr_size( - &ctx, - IntegerAttr::new(usize_ty, APInt::from_u64(1, NonZeroUsize::new(64).unwrap())), - ); - alloc.set_attr_alloc_key(&ctx, StringAttr::new("mismatched-shared".to_string())); - - let error = alloc - .verify(&ctx) - .expect_err("shared storage cannot claim an unrelated result pointee type"); - assert!( - error - .to_string() - .contains("pointee type must match elem_type"), - "{error}" - ); -} - -#[test] -fn test_mir_load_verify() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), false); - - let block = BasicBlock::new(&mut ctx, None, vec![ptr_ty.into()]); - let ptr_val = block.deref(&ctx).get_argument(0); - - let op = Operation::new( - &mut ctx, - MirLoadOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![ptr_val], - vec![], - 0, - ); - let mir_load = MirLoadOp::new(op); - assert!(mir_load.verify(&ctx).is_ok(), "Valid MirLoadOp"); - - let block_i32 = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); - let i32_val = block_i32.deref(&ctx).get_argument(0); - - let op_fail_operand = Operation::new( - &mut ctx, - MirLoadOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![i32_val], - vec![], - 0, - ); - let mir_load_fail_operand = MirLoadOp::new(op_fail_operand); - assert!( - mir_load_fail_operand.verify(&ctx).is_err(), - "MirLoadOp non-ptr operand" - ); - - let f32_ty = FP32Type::get(&ctx); - let op_fail_res = Operation::new( - &mut ctx, - MirLoadOp::get_concrete_op_info(), - vec![f32_ty.into()], - vec![ptr_val], - vec![], - 0, - ); - let mir_load_fail_res = MirLoadOp::new(op_fail_res); - assert!( - mir_load_fail_res.verify(&ctx).is_err(), - "MirLoadOp result mismatch" - ); -} - -#[test] -fn test_mir_load_volatile_is_not_promotable() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), false); - let block = BasicBlock::new(&mut ctx, None, vec![ptr_ty.into()]); - let ptr_val = block.deref(&ctx).get_argument(0); - - let op = Operation::new( - &mut ctx, - MirLoadOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![ptr_val], - vec![], - 0, - ); - let mir_load = MirLoadOp::new(op); - let alloc_info = AllocInfo { - ptr: ptr_val, - ty: i32_ty.into(), - }; - - assert!(!mir_load.is_volatile(&ctx)); - assert!(matches!( - mir_load.promotion_kind(&ctx, &alloc_info), - PromotableOpKind::Load - )); - - mir_load.set_volatile(&mut ctx, true); - - assert!(mir_load.is_volatile(&ctx)); - assert!(matches!( - mir_load.promotion_kind(&ctx, &alloc_info), - PromotableOpKind::NonPromotableUse - )); -} - -#[test] -fn test_mir_ptr_offset_verify() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), false); - let usize_ty = IntegerType::get(&ctx, 64, Signedness::Signless); - - let block = BasicBlock::new(&mut ctx, None, vec![ptr_ty.into(), usize_ty.into()]); - let ptr_val = block.deref(&ctx).get_argument(0); - let idx_val = block.deref(&ctx).get_argument(1); - - let op = Operation::new( - &mut ctx, - MirPtrOffsetOp::get_concrete_op_info(), - vec![ptr_ty.into()], - vec![ptr_val, idx_val], - vec![], - 0, - ); - let offset_op = MirPtrOffsetOp::new(op); - assert!(offset_op.verify(&ctx).is_ok(), "Valid MirPtrOffsetOp"); - assert!( - offset_op.is_inbounds(&ctx), - "ordinary pointer offsets default to inbounds" - ); - offset_op.set_inbounds(&mut ctx, false); - assert!( - !offset_op.is_inbounds(&ctx), - "wrapping pointer offsets retain their explicit semantics" - ); - - let block2 = BasicBlock::new(&mut ctx, None, vec![i32_ty.into(), usize_ty.into()]); - let i32_val = block2.deref(&ctx).get_argument(0); - let idx_val2 = block2.deref(&ctx).get_argument(1); - - let op_bad_base = Operation::new( - &mut ctx, - MirPtrOffsetOp::get_concrete_op_info(), - vec![ptr_ty.into()], - vec![i32_val, idx_val2], - vec![], - 0, - ); - assert!(MirPtrOffsetOp::new(op_bad_base).verify(&ctx).is_err()); - - let f32_ty = FP32Type::get(&ctx); - let ptr_f32_ty = MirPtrType::get_generic(&mut ctx, f32_ty.into(), false); - let op_bad_res = Operation::new( - &mut ctx, - MirPtrOffsetOp::get_concrete_op_info(), - vec![ptr_f32_ty.into()], - vec![ptr_val, idx_val], - vec![], - 0, - ); - assert!(MirPtrOffsetOp::new(op_bad_res).verify(&ctx).is_err()); -} - -#[test] -fn test_pointer_kind_laundering_requires_explicit_authority() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let usize_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let raw_mut_ty = - MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::RawMut); - let erased_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), true); - let erased_read_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), false); - let unique_ty = - MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::UniqueRef); - let shared_ty = MirPtrType::get_generic_with_kind( - &mut ctx, - i32_ty.into(), - false, - MirPointerKind::SharedRef, - ); - let raw_const_ty = - MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), false, MirPointerKind::RawConst); - let block = BasicBlock::new( - &mut ctx, - None, - vec![ - raw_mut_ty.into(), - erased_ty.into(), - usize_ty.into(), - shared_ty.into(), - raw_const_ty.into(), - erased_read_ty.into(), - ], - ); - let raw_mut = block.deref(&ctx).get_argument(0); - let erased = block.deref(&ctx).get_argument(1); - let offset = block.deref(&ctx).get_argument(2); - let shared = block.deref(&ctx).get_argument(3); - let raw_const = block.deref(&ctx).get_argument(4); - let erased_read = block.deref(&ctx).get_argument(5); - - // The concrete laundering example: pointer arithmetic is not a Rust - // reborrow and cannot manufacture `&mut T` from `*mut T`. - let raw_offset_to_unique = Operation::new( - &mut ctx, - MirPtrOffsetOp::get_concrete_op_info(), - vec![unique_ty.into()], - vec![raw_mut, offset], - vec![], - 0, - ); - assert!( - MirPtrOffsetOp::new(raw_offset_to_unique) - .verify(&ctx) - .is_err(), - "ptr_offset must not invent UniqueRef from RawMut" - ); - - let erased_offset_to_unique = Operation::new( - &mut ctx, - MirPtrOffsetOp::get_concrete_op_info(), - vec![unique_ty.into()], - vec![erased, offset], - vec![], - 0, - ); - assert!( - MirPtrOffsetOp::new(erased_offset_to_unique) - .verify(&ctx) - .is_err(), - "ptr_offset must not recover UniqueRef from Erased" - ); - - let preserving_offset = Operation::new( - &mut ctx, - MirPtrOffsetOp::get_concrete_op_info(), - vec![raw_mut_ty.into()], - vec![raw_mut, offset], - vec![], - 0, - ); - assert!(MirPtrOffsetOp::new(preserving_offset).verify(&ctx).is_ok()); - - let erasing_offset = Operation::new( - &mut ctx, - MirPtrOffsetOp::get_concrete_op_info(), - vec![erased_ty.into()], - vec![raw_mut, offset], - vec![], - 0, - ); - assert!(MirPtrOffsetOp::new(erasing_offset).verify(&ctx).is_ok()); - - let mutability_launder = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![erased_ty.into()], - vec![erased_read], - vec![], - 0, - ); - let mutability_launder_cast = MirCastOp::new(mutability_launder); - mutability_launder_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - assert!( - mutability_launder_cast.verify(&ctx).is_err(), - "an unmarked cast cannot manufacture writable Erased evidence" - ); - - let reborrow = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![unique_ty.into()], - vec![raw_mut], - vec![], - 0, - ); - let reborrow_cast = MirCastOp::new(reborrow); - reborrow_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - assert!( - reborrow_cast.verify(&ctx).is_err(), - "an unmarked cast must not manufacture UniqueRef" - ); - reborrow_cast.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); - assert!( - reborrow_cast.verify(&ctx).is_ok(), - "a rustc-declared reborrow is the explicit authority" - ); - - for (target, authority) in [ - (unique_ty.into(), MirPointerKindAuthorityAttr::Reborrow), - (raw_mut_ty.into(), MirPointerKindAuthorityAttr::RawAddress), - ] { - let mutable_storage_boundary = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![target], - vec![erased], - vec![], - 0, - ); - let mutable_storage_boundary_cast = MirCastOp::new(mutable_storage_boundary); - mutable_storage_boundary_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - mutable_storage_boundary_cast.set_pointer_kind_authority(&mut ctx, authority); - assert!( - mutable_storage_boundary_cast.verify(&ctx).is_ok(), - "mutable compiler storage is a valid source for a mutable Rust boundary" - ); - } - - let raw_to_shared_static = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![shared_ty.into()], - vec![raw_const], - vec![], - 0, - ); - let raw_to_shared_static_cast = MirCastOp::new(raw_to_shared_static); - raw_to_shared_static_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - raw_to_shared_static_cast - .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::StaticAddress); - assert!( - raw_to_shared_static_cast.verify(&ctx).is_err(), - "StaticAddress may establish a typed static value only from Erased storage, not relabel an arbitrary raw pointer as SharedRef" - ); - - let raw_to_erased = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![erased_read_ty.into()], - vec![raw_const], - vec![], - 0, - ); - let raw_to_erased_cast = MirCastOp::new(raw_to_erased); - raw_to_erased_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - assert!(raw_to_erased_cast.verify(&ctx).is_ok()); - let erased_from_raw = raw_to_erased.deref(&ctx).get_result(0); - let laundered_static = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![shared_ty.into()], - vec![erased_from_raw], - vec![], - 0, - ); - let laundered_static_cast = MirCastOp::new(laundered_static); - laundered_static_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - laundered_static_cast - .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::StaticAddress); - assert!( - laundered_static_cast.verify(&ctx).is_err(), - "erasing RawConst must not make it valid StaticAddress storage" - ); - - let raw_mut_to_erased = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![erased_ty.into()], - vec![raw_mut], - vec![], - 0, - ); - let raw_mut_to_erased_cast = MirCastOp::new(raw_mut_to_erased); - raw_mut_to_erased_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - assert!(raw_mut_to_erased_cast.verify(&ctx).is_ok()); - let erased_from_raw_mut = raw_mut_to_erased.deref(&ctx).get_result(0); - let laundered_abi = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![unique_ty.into()], - vec![erased_from_raw_mut], - vec![], - 0, - ); - let laundered_abi_cast = MirCastOp::new(laundered_abi); - laundered_abi_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - laundered_abi_cast - .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::AbiBoundary); - assert!( - laundered_abi_cast.verify(&ctx).is_err(), - "erasing RawMut must not let AbiBoundary manufacture UniqueRef" - ); - - let global_erased_ty = MirPtrType::get_with_kind( - &mut ctx, - i32_ty.into(), - false, - dialect_mir::types::address_space::GLOBAL, - MirPointerKind::Erased, - ); - let global_op = Operation::new( - &mut ctx, - MirGlobalAllocOp::get_concrete_op_info(), - vec![global_erased_ty.into()], - vec![], - vec![], - 0, - ); - let global = MirGlobalAllocOp::new(global_op); - global.set_attr_global_type(&ctx, TypeAttr::new(i32_ty.into())); - global.set_attr_global_key(&ctx, StringAttr::new("lineage-global".to_string())); - assert!(global.verify(&ctx).is_ok()); - let global_storage = global_op.deref(&ctx).get_result(0); - let typed_global = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![shared_ty.into()], - vec![global_storage], - vec![], - 0, - ); - let typed_global_cast = MirCastOp::new(typed_global); - typed_global_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - typed_global_cast - .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::StaticAddress); - assert!( - typed_global_cast.verify(&ctx).is_ok(), - "StaticAddress accepts a verified global-allocation root" - ); - - let shared_erased_ty = MirPtrType::get_shared(&mut ctx, i32_ty.into(), true); - let shared_op = Operation::new( - &mut ctx, - MirSharedAllocOp::get_concrete_op_info(), - vec![shared_erased_ty.into()], - vec![], - vec![], - 0, - ); - let shared_alloc = MirSharedAllocOp::new(shared_op); - shared_alloc.set_attr_elem_type(&ctx, TypeAttr::new(i32_ty.into())); - shared_alloc.set_attr_size( - &ctx, - IntegerAttr::new(usize_ty, APInt::from_u64(1, NonZeroUsize::new(64).unwrap())), - ); - shared_alloc.set_attr_alloc_key(&ctx, StringAttr::new("lineage-shared".to_string())); - assert!(shared_alloc.verify(&ctx).is_ok()); - let shared_storage = shared_op.deref(&ctx).get_result(0); - let typed_shared = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![raw_mut_ty.into()], - vec![shared_storage], - vec![], - 0, - ); - let typed_shared_cast = MirCastOp::new(typed_shared); - typed_shared_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - typed_shared_cast - .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::StaticAddress); - assert!( - typed_shared_cast.verify(&ctx).is_ok(), - "StaticAddress accepts a verified shared-allocation root" - ); - - let alloca_op = Operation::new( - &mut ctx, - MirAllocaOp::get_concrete_op_info(), - vec![erased_ty.into()], - vec![], - vec![], - 0, - ); - let alloca = MirAllocaOp::new(alloca_op); - assert!(alloca.verify(&ctx).is_ok()); - let compiler_storage = alloca_op.deref(&ctx).get_result(0); - let typed_abi = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![raw_mut_ty.into()], - vec![compiler_storage], - vec![], - 0, - ); - let typed_abi_cast = MirCastOp::new(typed_abi); - typed_abi_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - typed_abi_cast.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::AbiBoundary); - assert!( - typed_abi_cast.verify(&ctx).is_ok(), - "AbiBoundary accepts verified compiler-owned alloca storage" - ); - - for (target, authority) in [ - (unique_ty.into(), MirPointerKindAuthorityAttr::Reborrow), - (raw_mut_ty.into(), MirPointerKindAuthorityAttr::RawAddress), - ( - raw_mut_ty.into(), - MirPointerKindAuthorityAttr::StaticAddress, - ), - (raw_mut_ty.into(), MirPointerKindAuthorityAttr::AbiBoundary), - ] { - let immutable_storage_boundary = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![target], - vec![erased_read], - vec![], - 0, - ); - let immutable_storage_boundary_cast = MirCastOp::new(immutable_storage_boundary); - immutable_storage_boundary_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - immutable_storage_boundary_cast.set_pointer_kind_authority(&mut ctx, authority); - assert!( - immutable_storage_boundary_cast.verify(&ctx).is_err(), - "an immutable Erased thin pointer cannot establish a mutable Rust pointer kind" - ); - } - - let wrong_authority = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![unique_ty.into()], - vec![erased], - vec![], - 0, - ); - let wrong_authority_cast = MirCastOp::new(wrong_authority); - wrong_authority_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - wrong_authority_cast - .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::RawAddress); - assert!( - wrong_authority_cast.verify(&ctx).is_err(), - "RawAddress authority cannot establish a reference kind" - ); - - let inline_asm_authority = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![unique_ty.into()], - vec![erased], - vec![], - 0, - ); - let inline_asm_authority_cast = MirCastOp::new(inline_asm_authority); - inline_asm_authority_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - inline_asm_authority_cast - .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::InlineAsm); - assert!( - inline_asm_authority_cast.verify(&ctx).is_err(), - "InlineAsm is a producer authority and must never authorize a cast" - ); - - let integer_reborrow = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![unique_ty.into()], - vec![offset], - vec![], - 0, - ); - let integer_reborrow_cast = MirCastOp::new(integer_reborrow); - integer_reborrow_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - integer_reborrow_cast - .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); - assert!( - integer_reborrow_cast.verify(&ctx).is_err(), - "Reborrow authority cannot reinterpret an integer as a Rust reference" - ); - integer_reborrow_cast - .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::RustCast); - assert!( - integer_reborrow_cast.verify(&ctx).is_err(), - "RustCast authority cannot make an integer a PtrToPtr operand" - ); - integer_reborrow_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::FnPtrToPtr); - assert!( - integer_reborrow_cast.verify(&ctx).is_err(), - "FnPtrToPtr also requires a real pointer carrier" - ); - - let f32_ty = FP32Type::get(&ctx); - let wrong_pointee_ty = - MirPtrType::get_generic_with_kind(&mut ctx, f32_ty.into(), true, MirPointerKind::RawMut); - let wrong_pointee_block = BasicBlock::new(&mut ctx, None, vec![wrong_pointee_ty.into()]); - let wrong_pointee = wrong_pointee_block.deref(&ctx).get_argument(0); - let wrong_pointee_reborrow = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![unique_ty.into()], - vec![wrong_pointee], - vec![], - 0, - ); - let wrong_pointee_reborrow_cast = MirCastOp::new(wrong_pointee_reborrow); - wrong_pointee_reborrow_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - wrong_pointee_reborrow_cast - .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); - assert!( - wrong_pointee_reborrow_cast.verify(&ctx).is_err(), - "Reborrow authority must retain the pointee type" - ); - - for immutable_source in [shared, raw_const] { - let invalid_unique = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![unique_ty.into()], - vec![immutable_source], - vec![], - 0, - ); - let invalid_unique_cast = MirCastOp::new(invalid_unique); - invalid_unique_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - invalid_unique_cast - .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); - assert!( - invalid_unique_cast.verify(&ctx).is_err(), - "an immutable source cannot be relabelled as UniqueRef by Reborrow authority" - ); - - let invalid_raw_mut = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![raw_mut_ty.into()], - vec![immutable_source], - vec![], - 0, - ); - let invalid_raw_mut_cast = MirCastOp::new(invalid_raw_mut); - invalid_raw_mut_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - invalid_raw_mut_cast - .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::RawAddress); - assert!( - invalid_raw_mut_cast.verify(&ctx).is_err(), - "an immutable source cannot be relabelled as RawMut by RawAddress authority" - ); - } -} - -#[test] -fn test_promoted_empty_mutable_reference_is_a_narrow_static_exception() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let empty_array_ty = MirArrayType::get(&mut ctx, i32_ty.into(), 0); - let nonempty_array_ty = MirArrayType::get(&mut ctx, i32_ty.into(), 1); - - let empty_storage_ty = MirPtrType::get_with_kind( - &mut ctx, - empty_array_ty.into(), - false, - dialect_mir::types::address_space::GLOBAL, - MirPointerKind::Erased, - ); - let empty_unique_ty = MirPtrType::get_generic_with_kind( - &mut ctx, - empty_array_ty.into(), - true, - MirPointerKind::UniqueRef, - ); - let empty_global_op = Operation::new( - &mut ctx, - MirGlobalAllocOp::get_concrete_op_info(), - vec![empty_storage_ty.into()], - vec![], - vec![], - 0, - ); - let empty_global = MirGlobalAllocOp::new(empty_global_op); - empty_global.set_attr_global_type(&ctx, TypeAttr::new(empty_array_ty.into())); - empty_global.set_attr_global_key( - &ctx, - StringAttr::new("promoted-empty-mutable-reference".to_string()), - ); - empty_global.set_alignment_value(&mut ctx, 4); - empty_global_op.deref_mut(&ctx).attributes.set( - "global_initializer_hex".try_into().unwrap(), - StringAttr::new(String::new()), - ); - empty_global.mark_immutable(&mut ctx); - assert!(empty_global.verify(&ctx).is_ok()); - - let empty_global_storage = empty_global_op.deref(&ctx).get_result(0); - let empty_static_borrow_op = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![empty_unique_ty.into()], - vec![empty_global_storage], - vec![], - 0, - ); - let empty_static_borrow = MirCastOp::new(empty_static_borrow_op); - empty_static_borrow.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - empty_static_borrow - .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::StaticAddress); - - empty_global.set_alignment_value(&mut ctx, 1); - assert!( - empty_static_borrow.verify(&ctx).is_err(), - "even [i32; 0] must retain i32's natural four-byte alignment" - ); - empty_global.set_alignment_value(&mut ctx, 4); - assert!( - empty_static_borrow.verify(&ctx).is_ok(), - "an immutable promoted [T; 0] global may back rustc's vacuous &mut []" - ); - - empty_global_op.deref_mut(&ctx).attributes.set( - "global_initializer_hex".try_into().unwrap(), - StringAttr::new("00".to_string()), - ); - assert!( - empty_static_borrow.verify(&ctx).is_err(), - "the exception requires an actually empty initializer, not merely the attribute" - ); - empty_global_op.deref_mut(&ctx).attributes.set( - "global_initializer_hex".try_into().unwrap(), - StringAttr::new(String::new()), - ); - empty_global_op.deref_mut(&ctx).attributes.set( - "global_initializer_relocations".try_into().unwrap(), - StringAttr::new("unexpected-relocation".to_string()), - ); - assert!( - empty_static_borrow.verify(&ctx).is_err(), - "zero-byte promoted storage cannot carry a relocation" - ); - - let mutable_empty_global_op = Operation::new( - &mut ctx, - MirGlobalAllocOp::get_concrete_op_info(), - vec![empty_storage_ty.into()], - vec![], - vec![], - 0, - ); - let mutable_empty_global = MirGlobalAllocOp::new(mutable_empty_global_op); - mutable_empty_global.set_attr_global_type(&ctx, TypeAttr::new(empty_array_ty.into())); - mutable_empty_global.set_attr_global_key( - &ctx, - StringAttr::new("non-promoted-empty-global".to_string()), - ); - assert!(mutable_empty_global.verify(&ctx).is_ok()); - let mutable_empty_storage = mutable_empty_global_op.deref(&ctx).get_result(0); - let mutable_empty_borrow_op = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![empty_unique_ty.into()], - vec![mutable_empty_storage], - vec![], - 0, - ); - let mutable_empty_borrow = MirCastOp::new(mutable_empty_borrow_op); - mutable_empty_borrow.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - mutable_empty_borrow - .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::StaticAddress); - assert!( - mutable_empty_borrow.verify(&ctx).is_err(), - "the [T; 0] exception requires compiler-promoted immutable storage" - ); - - let aligned_element_ty: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "Align16".to_string(), - vec!["value".to_string()], - vec![i32_ty.into()], - vec![], - vec![0], - 16, - 16, - ) - .into(); - let aligned_empty_array_ty = MirArrayType::get(&mut ctx, aligned_element_ty, 0); - let underaligned_storage_ty = MirPtrType::get_with_kind( - &mut ctx, - aligned_empty_array_ty.into(), - false, - dialect_mir::types::address_space::GLOBAL, - MirPointerKind::Erased, - ); - let aligned_empty_unique_ty = MirPtrType::get_generic_with_kind( - &mut ctx, - aligned_empty_array_ty.into(), - true, - MirPointerKind::UniqueRef, - ); - let underaligned_global_op = Operation::new( - &mut ctx, - MirGlobalAllocOp::get_concrete_op_info(), - vec![underaligned_storage_ty.into()], - vec![], - vec![], - 0, - ); - let underaligned_global = MirGlobalAllocOp::new(underaligned_global_op); - underaligned_global.set_attr_global_type(&ctx, TypeAttr::new(aligned_empty_array_ty.into())); - underaligned_global.set_attr_global_key( - &ctx, - StringAttr::new("underaligned-empty-reference".to_string()), - ); - underaligned_global.set_alignment_value(&mut ctx, 1); - underaligned_global_op.deref_mut(&ctx).attributes.set( - "global_initializer_hex".try_into().unwrap(), - StringAttr::new(String::new()), - ); - underaligned_global.mark_immutable(&mut ctx); - assert!(underaligned_global.verify(&ctx).is_ok()); - let underaligned_storage = underaligned_global_op.deref(&ctx).get_result(0); - let underaligned_borrow_op = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![aligned_empty_unique_ty.into()], - vec![underaligned_storage], - vec![], - 0, - ); - let underaligned_borrow = MirCastOp::new(underaligned_borrow_op); - underaligned_borrow.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - underaligned_borrow - .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::StaticAddress); - assert!( - underaligned_borrow.verify(&ctx).is_err(), - "promoted &mut [Align16; 0] requires at least 16-byte global alignment" - ); - - let nonempty_storage_ty = MirPtrType::get_with_kind( - &mut ctx, - nonempty_array_ty.into(), - false, - dialect_mir::types::address_space::GLOBAL, - MirPointerKind::Erased, - ); - let nonempty_unique_ty = MirPtrType::get_generic_with_kind( - &mut ctx, - nonempty_array_ty.into(), - true, - MirPointerKind::UniqueRef, - ); - let nonempty_global_op = Operation::new( - &mut ctx, - MirGlobalAllocOp::get_concrete_op_info(), - vec![nonempty_storage_ty.into()], - vec![], - vec![], - 0, - ); - let nonempty_global = MirGlobalAllocOp::new(nonempty_global_op); - nonempty_global.set_attr_global_type(&ctx, TypeAttr::new(nonempty_array_ty.into())); - nonempty_global.set_attr_global_key( - &ctx, - StringAttr::new("promoted-nonempty-mutable-reference".to_string()), - ); - nonempty_global.mark_immutable(&mut ctx); - assert!(nonempty_global.verify(&ctx).is_ok()); - let nonempty_global_storage = nonempty_global_op.deref(&ctx).get_result(0); - let nonempty_static_borrow_op = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![nonempty_unique_ty.into()], - vec![nonempty_global_storage], - vec![], - 0, - ); - let nonempty_static_borrow = MirCastOp::new(nonempty_static_borrow_op); - nonempty_static_borrow.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - nonempty_static_borrow - .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::StaticAddress); - assert!( - nonempty_static_borrow.verify(&ctx).is_err(), - "StaticAddress must never manufacture UniqueRef for non-empty promoted storage" - ); - - let block = BasicBlock::new( - &mut ctx, - None, - vec![empty_storage_ty.into(), empty_unique_ty.into()], - ); - let erased_block_argument = block.deref(&ctx).get_argument(0); - let block_argument_borrow_op = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![empty_unique_ty.into()], - vec![erased_block_argument], - vec![], - 0, - ); - let block_argument_borrow = MirCastOp::new(block_argument_borrow_op); - block_argument_borrow.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - block_argument_borrow - .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::StaticAddress); - assert!( - block_argument_borrow.verify(&ctx).is_err(), - "an Erased [T; 0] block argument has no proven promoted-global lineage" - ); - - let raw_empty_ty = MirPtrType::get_generic_with_kind( - &mut ctx, - empty_array_ty.into(), - false, - MirPointerKind::RawConst, - ); - let raw_block = BasicBlock::new(&mut ctx, None, vec![raw_empty_ty.into()]); - let raw_empty = raw_block.deref(&ctx).get_argument(0); - let erase_raw_op = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![empty_storage_ty.into()], - vec![raw_empty], - vec![], - 0, - ); - let erase_raw = MirCastOp::new(erase_raw_op); - erase_raw.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - assert!(erase_raw.verify(&ctx).is_ok()); - let erased_from_raw = erase_raw_op.deref(&ctx).get_result(0); - let laundered_borrow_op = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![empty_unique_ty.into()], - vec![erased_from_raw], - vec![], - 0, - ); - let laundered_borrow = MirCastOp::new(laundered_borrow_op); - laundered_borrow.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - laundered_borrow - .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::StaticAddress); - assert!( - laundered_borrow.verify(&ctx).is_err(), - "RawConst -> Erased must not launder a zero-length pointer into UniqueRef" - ); - - let byte_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); - let byte_storage_ty = MirPtrType::get_with_kind( - &mut ctx, - byte_ty.into(), - false, - dialect_mir::types::address_space::GLOBAL, - MirPointerKind::Erased, - ); - let byte_global_op = Operation::new( - &mut ctx, - MirGlobalAllocOp::get_concrete_op_info(), - vec![byte_storage_ty.into()], - vec![], - vec![], - 0, - ); - let byte_global = MirGlobalAllocOp::new(byte_global_op); - byte_global.set_attr_global_type(&ctx, TypeAttr::new(byte_ty.into())); - byte_global.set_attr_global_key( - &ctx, - StringAttr::new("misaligned-empty-reference-root".to_string()), - ); - byte_global.mark_immutable(&mut ctx); - assert!(byte_global.verify(&ctx).is_ok()); - let byte_storage = byte_global_op.deref(&ctx).get_result(0); - let retype_root_op = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![empty_storage_ty.into()], - vec![byte_storage], - vec![], - 0, - ); - let retype_root = MirCastOp::new(retype_root_op); - retype_root.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - assert!(retype_root.verify(&ctx).is_ok()); - let retyped_storage = retype_root_op.deref(&ctx).get_result(0); - let misaligned_borrow_op = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![empty_unique_ty.into()], - vec![retyped_storage], - vec![], - 0, - ); - let misaligned_borrow = MirCastOp::new(misaligned_borrow_op); - misaligned_borrow.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - misaligned_borrow - .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::StaticAddress); - assert!( - misaligned_borrow.verify(&ctx).is_err(), - "an immutable byte global cannot be retyped into an aligned &mut [i32; 0] capability" - ); -} - -fn promoted_empty_unique_ref_verifies( - ctx: &mut Context, - element_ty: TypeHandle, - alignment: u64, -) -> bool { - let empty_array_ty: TypeHandle = MirArrayType::get(ctx, element_ty, 0).into(); - let storage_ty = MirPtrType::get_with_kind( - ctx, - empty_array_ty, - false, - dialect_mir::types::address_space::GLOBAL, - MirPointerKind::Erased, - ); - let unique_ty = - MirPtrType::get_generic_with_kind(ctx, empty_array_ty, true, MirPointerKind::UniqueRef); - let global_op = Operation::new( - ctx, - MirGlobalAllocOp::get_concrete_op_info(), - vec![storage_ty.into()], - vec![], - vec![], - 0, - ); - let global = MirGlobalAllocOp::new(global_op); - global.set_attr_global_type(ctx, TypeAttr::new(empty_array_ty)); - global.set_attr_global_key(ctx, StringAttr::new("promoted-empty-shape".to_string())); - global.set_alignment_value(ctx, alignment); - global_op.deref_mut(ctx).attributes.set( - "global_initializer_hex".try_into().unwrap(), - StringAttr::new(String::new()), - ); - global.mark_immutable(ctx); - - let storage = global_op.deref(ctx).get_result(0); - let cast_op = Operation::new( - ctx, - MirCastOp::get_concrete_op_info(), - vec![unique_ty.into()], - vec![storage], - vec![], - 0, - ); - let cast = MirCastOp::new(cast_op); - cast.set_attr_cast_kind(ctx, MirCastKindAttr::PtrToPtr); - cast.set_pointer_kind_authority(ctx, MirPointerKindAuthorityAttr::StaticAddress); - cast.verify(ctx).is_ok() -} - -#[test] -fn test_promoted_empty_alignment_covers_supported_fat_and_unit_elements() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let byte: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); - let word: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); - let unit: TypeHandle = MirTupleType::get(&mut ctx, vec![]).into(); - let slice: TypeHandle = - MirSliceType::get_with_kind(&mut ctx, byte, MirPointerKind::SharedRef).into(); - let disjoint: TypeHandle = MirDisjointSliceType::get(&mut ctx, byte).into(); - - assert!(promoted_empty_unique_ref_verifies(&mut ctx, unit, 1)); - assert!(promoted_empty_unique_ref_verifies(&mut ctx, slice, 8)); - assert!(promoted_empty_unique_ref_verifies(&mut ctx, disjoint, 8)); - assert!(promoted_empty_unique_ref_verifies(&mut ctx, word, 8)); - assert!( - !promoted_empty_unique_ref_verifies(&mut ctx, word, 12), - "a non-power-of-two numeric value is not a valid alignment guarantee" - ); - assert!( - !promoted_empty_unique_ref_verifies(&mut ctx, slice, 4), - "a stored slice value still requires pointer-word alignment" - ); -} - -#[test] -fn test_mir_ref_requires_exact_pointer_kind_authority() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let shared_ty = MirPtrType::get_generic_with_kind( - &mut ctx, - i32_ty.into(), - false, - MirPointerKind::SharedRef, - ); - let unique_ty = - MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::UniqueRef); - let raw_const_ty = - MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), false, MirPointerKind::RawConst); - let raw_mut_ty = - MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::RawMut); - let erased_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), false); - let global_shared_ty = MirPtrType::get_with_kind( - &mut ctx, - i32_ty.into(), - false, - dialect_mir::types::address_space::GLOBAL, - MirPointerKind::SharedRef, - ); - let block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); - let value = block.deref(&ctx).get_argument(0); - - let build = - |ctx: &mut Context, result_ty, mutable, authority: Option| { - let op = Operation::new( - ctx, - MirRefOp::get_concrete_op_info(), - vec![result_ty], - vec![value], - vec![], - 0, - ); - let reference = MirRefOp::new(op); - reference.set_mutable(ctx, mutable); - if let Some(authority) = authority { - reference.set_pointer_kind_authority(ctx, authority); - } - op - }; - - for (result_ty, mutable, authority) in [ - ( - shared_ty.into(), - false, - MirPointerKindAuthorityAttr::Reborrow, - ), - ( - unique_ty.into(), - true, - MirPointerKindAuthorityAttr::Reborrow, - ), - ( - raw_const_ty.into(), - false, - MirPointerKindAuthorityAttr::RawAddress, - ), - ( - raw_mut_ty.into(), - true, - MirPointerKindAuthorityAttr::RawAddress, - ), - ( - shared_ty.into(), - false, - MirPointerKindAuthorityAttr::StaticAddress, - ), - ] { - assert!( - MirRefOp::new(build(&mut ctx, result_ty, mutable, Some(authority))) - .verify(&ctx) - .is_ok() - ); - } - - assert!( - MirRefOp::new(build(&mut ctx, shared_ty.into(), false, None)) - .verify(&ctx) - .is_err(), - "mir.ref must visibly identify its Rust semantic origin" - ); - assert!( - MirRefOp::new(build(&mut ctx, erased_ty.into(), false, None)) - .verify(&ctx) - .is_err(), - "mir.ref is an explicit pointer-creation boundary, not generic Erased storage" - ); - assert!( - MirRefOp::new(build( - &mut ctx, - shared_ty.into(), - false, - Some(MirPointerKindAuthorityAttr::RawAddress), - )) - .verify(&ctx) - .is_err(), - "RawAddress cannot manufacture SharedRef" - ); - assert!( - MirRefOp::new(build( - &mut ctx, - unique_ty.into(), - true, - Some(MirPointerKindAuthorityAttr::StaticAddress), - )) - .verify(&ctx) - .is_err(), - "constant/static materialization cannot manufacture uniqueness" - ); - assert!( - MirRefOp::new(build( - &mut ctx, - unique_ty.into(), - false, - Some(MirPointerKindAuthorityAttr::Reborrow), - )) - .verify(&ctx) - .is_err(), - "reference kind must agree with the operation's mutability" - ); - assert!( - MirRefOp::new(build( - &mut ctx, - global_shared_ty.into(), - false, - Some(MirPointerKindAuthorityAttr::Reborrow), - )) - .verify(&ctx) - .is_err(), - "mir.ref materializes generic stack storage" - ); -} - -#[test] -fn test_aggregate_cast_cannot_hide_pointer_kind_laundering() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let raw_mut_ty = - MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::RawMut); - let unique_ty = - MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::UniqueRef); - let raw_tuple_ty = MirTupleType::get(&mut ctx, vec![raw_mut_ty.into()]); - let unique_tuple_ty = MirTupleType::get(&mut ctx, vec![unique_ty.into()]); - let block = BasicBlock::new(&mut ctx, None, vec![raw_tuple_ty.into()]); - let raw_tuple = block.deref(&ctx).get_argument(0); - - let build = |ctx: &mut Context, authority: Option| { - let op = Operation::new( - ctx, - MirCastOp::get_concrete_op_info(), - vec![unique_tuple_ty.into()], - vec![raw_tuple], - vec![], - 0, - ); - let cast = MirCastOp::new(op); - cast.set_attr_cast_kind(ctx, MirCastKindAttr::Transmute); - if let Some(authority) = authority { - cast.set_pointer_kind_authority(ctx, authority); - } - op - }; - - assert!( - MirCastOp::new(build(&mut ctx, None)).verify(&ctx).is_err(), - "nested RawMut -> UniqueRef laundering must be rejected" - ); - assert!( - MirCastOp::new(build(&mut ctx, Some(MirPointerKindAuthorityAttr::Reborrow),)) - .verify(&ctx) - .is_err(), - "Rvalue::Ref cannot authorize a nested aggregate transition" - ); - assert!( - MirCastOp::new(build(&mut ctx, Some(MirPointerKindAuthorityAttr::RustCast),)) - .verify(&ctx) - .is_ok(), - "an explicit rustc aggregate transmute is visible and authorized" - ); -} - -#[test] -fn test_aggregate_reinterpretation_requires_rust_cast_authority() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let i64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let shared_ty = MirPtrType::get_generic_with_kind( - &mut ctx, - i64_ty.into(), - false, - MirPointerKind::SharedRef, - ); - // Both tuple declarations list the same pointer at field 0, but its byte - // offset changes from 0 to 8. Pairing by declaration index alone would - // therefore mistake an integer's old bytes for a preserved reference. - let source_ty = MirTupleType::get_with_layout( - &mut ctx, - vec![shared_ty.into(), i64_ty.into()], - vec![0, 1], - vec![0, 8], - 16, - 8, - ); - let target_ty = MirTupleType::get_with_layout( - &mut ctx, - vec![shared_ty.into(), i64_ty.into()], - vec![1, 0], - vec![8, 0], - 16, - 8, - ); - let block = BasicBlock::new(&mut ctx, None, vec![source_ty.into()]); - let source = block.deref(&ctx).get_argument(0); - - let build = |ctx: &mut Context, authority: Option| { - let op = Operation::new( - ctx, - MirCastOp::get_concrete_op_info(), - vec![target_ty.into()], - vec![source], - vec![], - 0, - ); - let cast = MirCastOp::new(op); - cast.set_attr_cast_kind(ctx, MirCastKindAttr::Transmute); - if let Some(authority) = authority { - cast.set_pointer_kind_authority(ctx, authority); - } - op - }; - - assert!( - MirCastOp::new(build(&mut ctx, None)).verify(&ctx).is_err(), - "layout-changing aggregate casts must not infer pointer preservation by field index" - ); - assert!( - MirCastOp::new(build(&mut ctx, Some(MirPointerKindAuthorityAttr::RustCast),)) - .verify(&ctx) - .is_ok(), - "an explicit rustc transmute makes the representation reinterpretation auditable" - ); -} - -#[test] -fn test_generic_aggregate_cast_cannot_invent_erased_pointer_carriers() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let word_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let erased_mut = MirPtrType::get_generic(&mut ctx, word_ty.into(), true); - let source_ty = MirTupleType::get_with_layout( - &mut ctx, - vec![erased_mut.into(), word_ty.into()], - vec![0, 1], - vec![0, 8], - 16, - 8, - ); - let extra_pointer_ty = MirTupleType::get_with_layout( - &mut ctx, - vec![erased_mut.into(), erased_mut.into()], - vec![0, 1], - vec![0, 8], - 16, - 8, - ); - let moved_pointer_ty = MirTupleType::get_with_layout( - &mut ctx, - vec![erased_mut.into(), word_ty.into()], - vec![1, 0], - vec![8, 0], - 16, - 8, - ); - let one_pointer = MirArrayType::get(&mut ctx, erased_mut.into(), 1); - let two_pointers = MirArrayType::get(&mut ctx, erased_mut.into(), 2); - - let source_block = BasicBlock::new(&mut ctx, None, vec![source_ty.into()]); - let source = source_block.deref(&ctx).get_argument(0); - let array_block = BasicBlock::new(&mut ctx, None, vec![one_pointer.into()]); - let array_source = array_block.deref(&ctx).get_argument(0); - let build = |ctx: &mut Context, source, target| { - let op = Operation::new( - ctx, - MirCastOp::get_concrete_op_info(), - vec![target], - vec![source], - vec![], - 0, - ); - MirCastOp::new(op).set_attr_cast_kind(ctx, MirCastKindAttr::PtrToPtr); - op - }; - - assert!( - MirCastOp::new(build(&mut ctx, source, extra_pointer_ty.into())) - .verify(&ctx) - .is_err(), - "PtrToPtr cannot reinterpret an integer field as writable Erased pointer evidence" - ); - assert!( - MirCastOp::new(build(&mut ctx, source, moved_pointer_ty.into())) - .verify(&ctx) - .is_err(), - "field-index equality cannot hide an Erased pointer moving onto integer bytes" - ); - assert!( - MirCastOp::new(build(&mut ctx, array_source, two_pointers.into())) - .verify(&ctx) - .is_err(), - "homogeneous array traversal must retain pointer-carrier cardinality" - ); -} - -#[test] -fn test_rust_cast_authority_obeys_cast_kind_semantics() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let u8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); - let usize_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let array_ty = MirArrayType::get(&mut ctx, u8_ty.into(), 4); - let raw_mut = - MirPtrType::get_generic_with_kind(&mut ctx, u8_ty.into(), true, MirPointerKind::RawMut); - let raw_const = - MirPtrType::get_generic_with_kind(&mut ctx, u8_ty.into(), false, MirPointerKind::RawConst); - let unique = - MirPtrType::get_generic_with_kind(&mut ctx, u8_ty.into(), true, MirPointerKind::UniqueRef); - let fn_target = MirStructType::get_with_full_layout( - &mut ctx, - "FnPtrTarget".into(), - vec![], - vec![], - vec![], - vec![], - 0, - 0, - ); - let fn_target_ty: pliron::r#type::TypeHandle = fn_target.into(); - let fn_carrier = MirPtrType::get_generic(&mut ctx, fn_target_ty, false); - let raw_mut_array = - MirPtrType::get_generic_with_kind(&mut ctx, array_ty.into(), true, MirPointerKind::RawMut); - let raw_const_array = MirPtrType::get_generic_with_kind( - &mut ctx, - array_ty.into(), - false, - MirPointerKind::RawConst, - ); - let erased_array = MirPtrType::get_generic(&mut ctx, array_ty.into(), false); - let erased_slice = MirSliceType::get(&mut ctx, u8_ty.into()); - let erased_mut_slice = MirSliceType::get_with_mutability(&mut ctx, u8_ty.into(), true); - let raw_const_slice = - MirSliceType::get_with_kind(&mut ctx, u8_ty.into(), MirPointerKind::RawConst); - let block = BasicBlock::new( - &mut ctx, - None, - vec![ - raw_mut.into(), - raw_const.into(), - fn_carrier.into(), - raw_mut_array.into(), - raw_const_array.into(), - erased_array.into(), - usize_ty.into(), - ], - ); - let raw_mut_value = block.deref(&ctx).get_argument(0); - let raw_const_value = block.deref(&ctx).get_argument(1); - let fn_value = block.deref(&ctx).get_argument(2); - let raw_mut_array_value = block.deref(&ctx).get_argument(3); - let raw_const_array_value = block.deref(&ctx).get_argument(4); - let erased_array_value = block.deref(&ctx).get_argument(5); - let integer_value = block.deref(&ctx).get_argument(6); - - let build = |ctx: &mut Context, - source, - target, - kind, - authority: Option| { - let op = Operation::new( - ctx, - MirCastOp::get_concrete_op_info(), - vec![target], - vec![source], - vec![], - 0, - ); - let cast = MirCastOp::new(op); - cast.set_attr_cast_kind(ctx, kind); - if let Some(authority) = authority { - cast.set_pointer_kind_authority(ctx, authority); - } - op - }; - let rust_cast = Some(MirPointerKindAuthorityAttr::RustCast); - - assert!( - MirCastOp::new(build( - &mut ctx, - raw_mut_value, - unique.into(), - MirCastKindAttr::PtrToPtr, - rust_cast.clone(), - )) - .verify(&ctx) - .is_err(), - "PtrToPtr cannot use RustCast to invent a reference" - ); - assert!( - MirCastOp::new(build( - &mut ctx, - raw_mut_value, - raw_const.into(), - MirCastKindAttr::PtrToPtr, - rust_cast.clone(), - )) - .verify(&ctx) - .is_ok(), - "PtrToPtr may perform an explicit raw-to-raw cast" - ); - assert!( - MirCastOp::new(build( - &mut ctx, - fn_value, - raw_const.into(), - MirCastKindAttr::FnPtrToPtr, - rust_cast.clone(), - )) - .verify(&ctx) - .is_ok(), - "FnPtrToPtr may expose an opaque function pointer as a raw pointer" - ); - assert!( - MirCastOp::new(build( - &mut ctx, - erased_array_value, - raw_const.into(), - MirCastKindAttr::FnPtrToPtr, - rust_cast.clone(), - )) - .verify(&ctx) - .is_err(), - "FnPtrToPtr cannot relabel arbitrary Erased storage as a function pointer" - ); - assert!( - MirCastOp::new(build( - &mut ctx, - fn_value, - unique.into(), - MirCastKindAttr::FnPtrToPtr, - rust_cast.clone(), - )) - .verify(&ctx) - .is_err() - ); - assert!( - MirCastOp::new(build( - &mut ctx, - raw_mut_value, - raw_const.into(), - MirCastKindAttr::PointerCoercionMutToConst, - rust_cast.clone(), - )) - .verify(&ctx) - .is_ok() - ); - assert!( - MirCastOp::new(build( - &mut ctx, - raw_const_value, - raw_mut.into(), - MirCastKindAttr::PointerCoercionMutToConst, - rust_cast.clone(), - )) - .verify(&ctx) - .is_err() - ); - assert!( - MirCastOp::new(build( - &mut ctx, - raw_mut_value, - raw_mut.into(), - MirCastKindAttr::PointerCoercionMutToConst, - None, - )) - .verify(&ctx) - .is_err(), - "MutToConst cannot masquerade as a same-category pointer cast" - ); - assert!( - MirCastOp::new(build( - &mut ctx, - raw_const_array_value, - raw_const_slice.into(), - MirCastKindAttr::PointerCoercionUnsize, - rust_cast.clone(), - )) - .verify(&ctx) - .is_ok(), - "Unsize may change thin/fat shape while preserving its carrier" - ); - assert!( - MirCastOp::new(build( - &mut ctx, - erased_array_value, - erased_slice.into(), - MirCastKindAttr::PointerCoercionUnsize, - None, - )) - .verify(&ctx) - .is_ok(), - "an all-Erased unsize still preserves read-only carrier state" - ); - assert!( - MirCastOp::new(build( - &mut ctx, - erased_array_value, - erased_mut_slice.into(), - MirCastKindAttr::PointerCoercionUnsize, - None, - )) - .verify(&ctx) - .is_err(), - "Unsize cannot turn an immutable Erased thin pointer into writable fat evidence" - ); - - let prefix_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let sized_tail = MirStructType::get_with_full_layout( - &mut ctx, - "TailCarrier".into(), - vec!["prefix".into(), "tail".into()], - vec![prefix_ty.into(), array_ty.into()], - vec![0, 1], - vec![0, 8], - 16, - 8, - ); - let shifted_unsized_tail = MirStructType::get_with_full_layout( - &mut ctx, - "TailCarrier".into(), - vec!["prefix".into(), "tail".into()], - vec![prefix_ty.into(), u8_ty.into()], - vec![0, 1], - vec![0, 16], - 24, - 8, - ); - let sized_tail_ptr = MirPtrType::get_generic_with_kind( - &mut ctx, - sized_tail.into(), - false, - MirPointerKind::RawConst, - ); - let shifted_tail_slice = MirSliceType::get_with_kind( - &mut ctx, - shifted_unsized_tail.into(), - MirPointerKind::RawConst, - ); - let tail_block = BasicBlock::new(&mut ctx, None, vec![sized_tail_ptr.into()]); - let tail_value = tail_block.deref(&ctx).get_argument(0); - assert!( - MirCastOp::new(build( - &mut ctx, - tail_value, - shifted_tail_slice.into(), - MirCastKindAttr::PointerCoercionUnsize, - rust_cast.clone(), - )) - .verify(&ctx) - .is_err(), - "Unsize cannot move the trailing data behind a changed struct-field offset" - ); - - assert!( - MirCastOp::new(build( - &mut ctx, - raw_const_array_value, - raw_mut.into(), - MirCastKindAttr::PointerCoercionArrayToPointer, - rust_cast.clone(), - )) - .verify(&ctx) - .is_err(), - "ArrayToPointer cannot strengthen const raw storage to mutable" - ); - assert!( - MirCastOp::new(build( - &mut ctx, - raw_mut_value, - raw_const.into(), - MirCastKindAttr::PointerCoercionArrayToPointer, - rust_cast.clone(), - )) - .verify(&ctx) - .is_err(), - "ArrayToPointer requires a raw pointer to an array, not an arbitrary pointer" - ); - assert!( - MirCastOp::new(build( - &mut ctx, - raw_mut_value, - raw_mut.into(), - MirCastKindAttr::PointerCoercionArrayToPointer, - None, - )) - .verify(&ctx) - .is_err(), - "an unmarked same-kind ArrayToPointer still requires an array source" - ); - assert!( - MirCastOp::new(build( - &mut ctx, - raw_mut_array_value, - raw_const.into(), - MirCastKindAttr::PointerCoercionArrayToPointer, - rust_cast.clone(), - )) - .verify(&ctx) - .is_ok() - ); - assert!( - MirCastOp::new(build( - &mut ctx, - integer_value, - unique.into(), - MirCastKindAttr::Transmute, - rust_cast, - )) - .verify(&ctx) - .is_ok(), - "only an explicit Transmute may establish an arbitrary pointer category" - ); - - let word_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let source_wrapper = MirStructType::get_with_full_layout( - &mut ctx, - "CarrierWrapper".into(), - vec!["pointer".into(), "word".into()], - vec![unique.into(), word_ty.into()], - vec![0, 1], - vec![0, 8], - 16, - 8, - ); - let moved_wrapper = MirStructType::get_with_full_layout( - &mut ctx, - "CarrierWrapper".into(), - vec!["word".into(), "pointer".into()], - vec![word_ty.into(), unique.into()], - vec![0, 1], - vec![0, 8], - 16, - 8, - ); - let wrapper_block = BasicBlock::new(&mut ctx, None, vec![source_wrapper.into()]); - let wrapper_value = wrapper_block.deref(&ctx).get_argument(0); - for kind in [ - MirCastKindAttr::PointerCoercionUnsize, - MirCastKindAttr::Subtype, - ] { - assert!( - MirCastOp::new(build( - &mut ctx, - wrapper_value, - moved_wrapper.into(), - kind, - Some(MirPointerKindAuthorityAttr::RustCast), - )) - .verify(&ctx) - .is_err(), - "Unsize/Subtype cannot bless a pointer carrier moved onto integer bytes" - ); - } -} - -#[test] -fn test_pointer_projections_preserve_or_erase_kind_only() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let usize_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let tuple_ty = MirTupleType::get(&mut ctx, vec![i32_ty.into()]); - let array_ty = MirArrayType::get(&mut ctx, i32_ty.into(), 4); - let raw_tuple_ty = - MirPtrType::get_generic_with_kind(&mut ctx, tuple_ty.into(), true, MirPointerKind::RawMut); - let raw_array_ty = - MirPtrType::get_generic_with_kind(&mut ctx, array_ty.into(), true, MirPointerKind::RawMut); - let raw_field_ty = - MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::RawMut); - let erased_field_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), true); - let erased_read_field_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), false); - let unique_field_ty = - MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::UniqueRef); - let block = BasicBlock::new( - &mut ctx, - None, - vec![raw_tuple_ty.into(), raw_array_ty.into(), usize_ty.into()], - ); - let tuple = block.deref(&ctx).get_argument(0); - let array = block.deref(&ctx).get_argument(1); - let index = block.deref(&ctx).get_argument(2); - - let field = |ctx: &mut Context, result_ty| { - let op = Operation::new( - ctx, - MirFieldAddrOp::get_concrete_op_info(), - vec![result_ty], - vec![tuple], - vec![], - 0, - ); - MirFieldAddrOp::new(op).set_attr_field_index(ctx, FieldIndexAttr(0)); - MirFieldAddrOp::new(op).set_attr_aggregate_ty(ctx, TypeAttr::new(tuple_ty.into())); - op - }; - assert!( - MirFieldAddrOp::new(field(&mut ctx, raw_field_ty.into())) - .verify(&ctx) - .is_ok() - ); - assert!( - MirFieldAddrOp::new(field(&mut ctx, erased_field_ty.into())) - .verify(&ctx) - .is_ok() - ); - assert!( - MirFieldAddrOp::new(field(&mut ctx, unique_field_ty.into())) - .verify(&ctx) - .is_err() - ); - assert!( - MirFieldAddrOp::new(field(&mut ctx, erased_read_field_ty.into())) - .verify(&ctx) - .is_err(), - "field projection cannot flip an Erased address from writable to read-only" - ); - - let array_launder = Operation::new( - &mut ctx, - MirArrayElementAddrOp::get_concrete_op_info(), - vec![unique_field_ty.into()], - vec![array, index], - vec![], - 0, - ); - assert!( - MirArrayElementAddrOp::new(array_launder) - .verify(&ctx) - .is_err() - ); - let array_erase = Operation::new( - &mut ctx, - MirArrayElementAddrOp::get_concrete_op_info(), - vec![erased_field_ty.into()], - vec![array, index], - vec![], - 0, - ); - assert!( - MirArrayElementAddrOp::new(array_erase).verify(&ctx).is_ok(), - "projection may erase kind while preserving machine mutability" - ); - let array_mutability_flip = Operation::new( - &mut ctx, - MirArrayElementAddrOp::get_concrete_op_info(), - vec![erased_read_field_ty.into()], - vec![array, index], - vec![], - 0, - ); - assert!( - MirArrayElementAddrOp::new(array_mutability_flip) - .verify(&ctx) - .is_err(), - "array projection cannot change machine mutability" - ); -} - -#[test] -fn test_pointer_with_exposed_provenance_only_creates_raw_or_erased_pointer() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let usize_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let fn_target = MirStructType::get_with_full_layout( - &mut ctx, - "FnPtrTarget".into(), - vec![], - vec![], - vec![], - vec![], - 0, - 0, - ); - let fn_target_ty: pliron::r#type::TypeHandle = fn_target.into(); - let raw_ty = - MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::RawMut); - let opaque_fn_ty = MirPtrType::get_generic(&mut ctx, fn_target_ty, false); - let arbitrary_erased_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), false); - let writable_erased_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), true); - let unique_ty = - MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::UniqueRef); - let block = BasicBlock::new(&mut ctx, None, vec![usize_ty.into()]); - let address = block.deref(&ctx).get_argument(0); - - let build = |ctx: &mut Context, result_ty| { - let op = Operation::new( - ctx, - MirCastOp::get_concrete_op_info(), - vec![result_ty], - vec![address], - vec![], - 0, - ); - let cast = MirCastOp::new(op); - cast.set_attr_cast_kind(ctx, MirCastKindAttr::PointerWithExposedProvenance); - cast.set_pointer_kind_authority(ctx, MirPointerKindAuthorityAttr::RustCast); - op - }; - - assert!( - MirCastOp::new(build(&mut ctx, raw_ty.into())) - .verify(&ctx) - .is_ok() - ); - assert!( - MirCastOp::new(build(&mut ctx, unique_ty.into())) - .verify(&ctx) - .is_err(), - "integer provenance cannot directly materialize a Rust reference" - ); - - let build_unmarked = |ctx: &mut Context, result_ty| { - let op = Operation::new( - ctx, - MirCastOp::get_concrete_op_info(), - vec![result_ty], - vec![address], - vec![], - 0, - ); - MirCastOp::new(op).set_attr_cast_kind(ctx, MirCastKindAttr::PointerWithExposedProvenance); - op - }; - assert!( - MirCastOp::new(build_unmarked(&mut ctx, opaque_fn_ty.into())) - .verify(&ctx) - .is_ok(), - "function-pointer tokens may materialize only the canonical immutable Erased carrier" - ); - assert!( - MirCastOp::new(build_unmarked(&mut ctx, arbitrary_erased_ty.into())) - .verify(&ctx) - .is_err(), - "an integer cannot manufacture arbitrary Erased pointer evidence" - ); - assert!( - MirCastOp::new(build_unmarked(&mut ctx, writable_erased_ty.into())) - .verify(&ctx) - .is_err(), - "an integer cannot manufacture writable Erased evidence and then reborrow it as UniqueRef" - ); - - let fn_block = BasicBlock::new(&mut ctx, None, vec![opaque_fn_ty.into()]); - let opaque_fn = fn_block.deref(&ctx).get_argument(0); - let disguised_data_pointer = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![arbitrary_erased_ty.into()], - vec![opaque_fn], - vec![], - 0, - ); - MirCastOp::new(disguised_data_pointer) - .set_attr_cast_kind(&ctx, MirCastKindAttr::PointerCoercionArrayToPointer); - assert!( - MirCastOp::new(disguised_data_pointer).verify(&ctx).is_err(), - "an unrelated coercion cannot turn the opaque function token into Erased data storage" - ); - - let shared_fn_ty = - MirPtrType::get_generic_with_kind(&mut ctx, fn_target_ty, false, MirPointerKind::SharedRef); - let fake_reborrow = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![shared_fn_ty.into()], - vec![opaque_fn], - vec![], - 0, - ); - let fake_reborrow = MirCastOp::new(fake_reborrow); - fake_reborrow.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - fake_reborrow.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); - assert!( - fake_reborrow.verify(&ctx).is_err(), - "an opaque function-pointer value is not compiler storage that may be reborrowed" - ); - - // Keeping only kind+mutability while recursively pairing aggregate fields - // is insufficient: the canonical function token is also identified by - // its pointee. Otherwise a tuple cast can disguise data storage as a - // function pointer (or vice versa), after which individually legal casts - // can manufacture a reference. - let data_tuple_ty = MirTupleType::get(&mut ctx, vec![arbitrary_erased_ty.into()]); - let fn_tuple_ty = MirTupleType::get(&mut ctx, vec![opaque_fn_ty.into()]); - let nested_block = BasicBlock::new( - &mut ctx, - None, - vec![data_tuple_ty.into(), fn_tuple_ty.into()], - ); - let data_tuple = nested_block.deref(&ctx).get_argument(0); - let fn_tuple = nested_block.deref(&ctx).get_argument(1); - let nested_cast = |ctx: &mut Context, source, target| { - let op = Operation::new( - ctx, - MirCastOp::get_concrete_op_info(), - vec![target], - vec![source], - vec![], - 0, - ); - MirCastOp::new(op).set_attr_cast_kind(ctx, MirCastKindAttr::PtrToPtr); - op - }; - let data_to_fn = nested_cast(&mut ctx, data_tuple, fn_tuple_ty.into()); - assert!( - MirCastOp::new(data_to_fn).verify(&ctx).is_err(), - "an aggregate cast cannot manufacture a nested canonical function token" - ); - assert!( - MirCastOp::new(nested_cast(&mut ctx, fn_tuple, data_tuple_ty.into(),)) - .verify(&ctx) - .is_err(), - "an aggregate cast cannot disguise a nested function token as Erased data" - ); - - let forged_fn_tuple = data_to_fn.deref(&ctx).get_result(0); - let extract_forged_fn = Operation::new( - &mut ctx, - MirExtractFieldOp::get_concrete_op_info(), - vec![opaque_fn_ty.into()], - vec![forged_fn_tuple], - vec![], - 0, - ); - let extract_forged_fn = MirExtractFieldOp::new(extract_forged_fn); - extract_forged_fn.set_attr_index(&ctx, FieldIndexAttr(0)); - assert!(extract_forged_fn.verify(&ctx).is_ok()); - - let forged_fn = extract_forged_fn.get_operation().deref(&ctx).get_result(0); - let expose_forged_fn = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![raw_ty.into()], - vec![forged_fn], - vec![], - 0, - ); - let expose_forged_fn = MirCastOp::new(expose_forged_fn); - expose_forged_fn.set_attr_cast_kind(&ctx, MirCastKindAttr::FnPtrToPtr); - expose_forged_fn.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::RustCast); - assert!(expose_forged_fn.verify(&ctx).is_ok()); - - let forged_raw = expose_forged_fn.get_operation().deref(&ctx).get_result(0); - let reborrow_forged_fn = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![unique_ty.into()], - vec![forged_raw], - vec![], - 0, - ); - let reborrow_forged_fn = MirCastOp::new(reborrow_forged_fn); - reborrow_forged_fn.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - reborrow_forged_fn.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); - assert!( - reborrow_forged_fn.verify(&ctx).is_ok(), - "the laundering chain must be rejected at its aggregate reinterpretation" - ); - - let shared_marker_ty = - MirPtrType::get_generic_with_kind(&mut ctx, fn_target_ty, false, MirPointerKind::SharedRef); - let marker_tuple_ty = MirTupleType::get(&mut ctx, vec![fn_target_ty]); - let shared_marker_tuple_ty = MirPtrType::get_generic_with_kind( - &mut ctx, - marker_tuple_ty.into(), - false, - MirPointerKind::SharedRef, - ); - let marker_array_ty = MirArrayType::get(&mut ctx, fn_target_ty, 1); - let shared_marker_array_ty = MirPtrType::get_generic_with_kind( - &mut ctx, - marker_array_ty.into(), - false, - MirPointerKind::SharedRef, - ); - let shared_marker_slice_ty = - MirSliceType::get_with_kind(&mut ctx, fn_target_ty, MirPointerKind::SharedRef); - let projection_block = BasicBlock::new( - &mut ctx, - None, - vec![ - shared_marker_ty.into(), - shared_marker_tuple_ty.into(), - shared_marker_array_ty.into(), - shared_marker_slice_ty.into(), - usize_ty.into(), - ], - ); - let marker_address = projection_block.deref(&ctx).get_argument(0); - let marker_tuple_address = projection_block.deref(&ctx).get_argument(1); - let marker_array_address = projection_block.deref(&ctx).get_argument(2); - let marker_slice = projection_block.deref(&ctx).get_argument(3); - let zero = projection_block.deref(&ctx).get_argument(4); - - let offset_to_token = Operation::new( - &mut ctx, - MirPtrOffsetOp::get_concrete_op_info(), - vec![opaque_fn_ty.into()], - vec![marker_address, zero], - vec![], - 0, - ); - assert!( - MirPtrOffsetOp::new(offset_to_token).verify(&ctx).is_err(), - "pointer arithmetic produces a data address, never a function token" - ); - - let field_address_to_token = Operation::new( - &mut ctx, - MirFieldAddrOp::get_concrete_op_info(), - vec![opaque_fn_ty.into()], - vec![marker_tuple_address], - vec![], - 0, - ); - let field_address_to_token = MirFieldAddrOp::new(field_address_to_token); - field_address_to_token.set_attr_field_index(&ctx, FieldIndexAttr(0)); - assert!( - field_address_to_token.verify(&ctx).is_err(), - "a field address cannot masquerade as a function token" - ); - - let array_address_to_token = Operation::new( - &mut ctx, - MirArrayElementAddrOp::get_concrete_op_info(), - vec![opaque_fn_ty.into()], - vec![marker_array_address, zero], - vec![], - 0, - ); - assert!( - MirArrayElementAddrOp::new(array_address_to_token) - .verify(&ctx) - .is_err(), - "an array element address cannot masquerade as a function token" - ); - - let slice_data_to_token = Operation::new( - &mut ctx, - MirExtractFieldOp::get_concrete_op_info(), - vec![opaque_fn_ty.into()], - vec![marker_slice], - vec![], - 0, - ); - let slice_data_to_token = MirExtractFieldOp::new(slice_data_to_token); - slice_data_to_token.set_attr_index(&ctx, FieldIndexAttr(0)); - assert!( - slice_data_to_token.verify(&ctx).is_err(), - "a slice data address cannot masquerade as a function token" - ); - - // A ClosureFnPointer cast must not extract captured reference bits as a - // function pointer, expose them as RawMut, and then reborrow them as - // UniqueRef. Only a genuinely non-capturing, zero-sized closure may enter - // the opaque function-pointer path. - let captured_ref = MirPtrType::get_generic_with_kind( - &mut ctx, - i32_ty.into(), - false, - MirPointerKind::SharedRef, - ); - let captured_closure = MirStructType::get_with_full_layout( - &mut ctx, - "CapturedClosure".into(), - vec!["capture_0".into()], - vec![captured_ref.into()], - vec![0], - vec![0], - 8, - 8, - ); - let empty_closure = MirStructType::get_with_full_layout( - &mut ctx, - "NonCapturingClosure".into(), - vec![], - vec![], - vec![], - vec![], - 0, - 1, - ); - let closure_block = BasicBlock::new( - &mut ctx, - None, - vec![captured_closure.into(), empty_closure.into()], - ); - let captured_value = closure_block.deref(&ctx).get_argument(0); - let empty_value = closure_block.deref(&ctx).get_argument(1); - let closure_cast = |ctx: &mut Context, source| { - let op = Operation::new( - ctx, - MirCastOp::get_concrete_op_info(), - vec![opaque_fn_ty.into()], - vec![source], - vec![], - 0, - ); - MirCastOp::new(op) - .set_attr_cast_kind(ctx, MirCastKindAttr::PointerCoercionClosureFnPointer); - op - }; - let captured_to_fn = closure_cast(&mut ctx, captured_value); - assert!( - MirCastOp::new(captured_to_fn).verify(&ctx).is_err(), - "a closure carrying SharedRef bytes cannot become an opaque function pointer" - ); - assert!( - MirCastOp::new(closure_cast(&mut ctx, empty_value)) - .verify(&ctx) - .is_err(), - "the importer materializes a closure function token directly; the legacy cast is not lowerable" - ); - - let captured_fn_value = captured_to_fn.deref(&ctx).get_result(0); - let exposed_capture = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![raw_ty.into()], - vec![captured_fn_value], - vec![], - 0, - ); - let exposed_capture = MirCastOp::new(exposed_capture); - exposed_capture.set_attr_cast_kind(&ctx, MirCastKindAttr::FnPtrToPtr); - exposed_capture.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::RustCast); - assert!(exposed_capture.verify(&ctx).is_ok()); - - let exposed_capture_value = exposed_capture.get_operation().deref(&ctx).get_result(0); - let unique_capture = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![unique_ty.into()], - vec![exposed_capture_value], - vec![], - 0, - ); - let unique_capture = MirCastOp::new(unique_capture); - unique_capture.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - unique_capture.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); - assert!( - unique_capture.verify(&ctx).is_ok(), - "the chain must be stopped at ClosureFnPointer before later individually legal boundaries" - ); -} - -#[test] -fn test_mir_extract_field_verify() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let tuple_ty = MirTupleType::get(&mut ctx, vec![i32_ty.into(), i32_ty.into()]); - - let block = BasicBlock::new(&mut ctx, None, vec![tuple_ty.into()]); - let tuple_val = block.deref(&ctx).get_argument(0); - - let op = Operation::new( - &mut ctx, - MirExtractFieldOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![tuple_val], - vec![], - 0, - ); - let extract_op = MirExtractFieldOp::new(op); - extract_op.set_attr_index(&ctx, dialect_mir::attributes::FieldIndexAttr(0)); - assert!(extract_op.verify(&ctx).is_ok(), "Valid Tuple Extract"); - - let op_oob = Operation::new( - &mut ctx, - MirExtractFieldOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![tuple_val], - vec![], - 0, - ); - let extract_op_oob = MirExtractFieldOp::new(op_oob); - extract_op_oob.set_attr_index(&ctx, dialect_mir::attributes::FieldIndexAttr(2)); - assert!(extract_op_oob.verify(&ctx).is_err(), "OOB Index"); - - let union_ty = MirUnionType::get( - &mut ctx, - "Bits".into(), - vec!["word".into(), "alias".into()], - vec![i32_ty.into(), i32_ty.into()], - 4, - 4, - ); - let union_block = BasicBlock::new(&mut ctx, None, vec![union_ty.into()]); - let union_val = union_block.deref(&ctx).get_argument(0); - let union_extract = Operation::new( - &mut ctx, - MirExtractFieldOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![union_val], - vec![], - 0, - ); - let union_extract = MirExtractFieldOp::new(union_extract); - union_extract.set_attr_index(&ctx, dialect_mir::attributes::FieldIndexAttr(1)); - assert!(union_extract.verify(&ctx).is_ok(), "Valid union extract"); -} - -#[test] -fn test_mir_construct_disjoint_slice_verify() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let f32_ty = FP32Type::get(&ctx); - let usize_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let width_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let f32_ptr_ty = - MirPtrType::get_generic_with_kind(&mut ctx, f32_ty.into(), true, MirPointerKind::RawMut); - let plain_ty = MirDisjointSliceType::get(&mut ctx, f32_ty.into()); - let width_ty_handle: pliron::r#type::TypeHandle = width_ty.into(); - let row_width_ty = - MirDisjointSliceType::get_with_space(&mut ctx, f32_ty.into(), vec![width_ty_handle]); - - let block = BasicBlock::new( - &mut ctx, - None, - vec![f32_ptr_ty.into(), usize_ty.into(), width_ty.into()], - ); - let ptr_val = block.deref(&ctx).get_argument(0); - let len_val = block.deref(&ctx).get_argument(1); - let width_val = block.deref(&ctx).get_argument(2); - - // Valid: an index space with no runtime layout takes two operands. - let op = Operation::new( - &mut ctx, - MirConstructDisjointSliceOp::get_concrete_op_info(), - vec![plain_ty.into()], - vec![ptr_val, len_val], - vec![], - 0, - ); - assert!( - MirConstructDisjointSliceOp::new(op).verify(&ctx).is_ok(), - "Valid space-free disjoint slice construction" - ); - - // Valid: a runtime row width takes a third operand. - let op_width = Operation::new( - &mut ctx, - MirConstructDisjointSliceOp::get_concrete_op_info(), - vec![row_width_ty.into()], - vec![ptr_val, len_val, width_val], - vec![], - 0, - ); - assert!( - MirConstructDisjointSliceOp::new(op_width) - .verify(&ctx) - .is_ok(), - "Valid row-width disjoint slice construction" - ); - - let erased_ptr_ty = MirPtrType::get_generic(&mut ctx, f32_ty.into(), true); - let erased_block = BasicBlock::new(&mut ctx, None, vec![erased_ptr_ty.into()]); - let erased_ptr = erased_block.deref(&ctx).get_argument(0); - let op_erased_data = Operation::new( - &mut ctx, - MirConstructDisjointSliceOp::get_concrete_op_info(), - vec![plain_ty.into()], - vec![erased_ptr, len_val], - vec![], - 0, - ); - assert!( - MirConstructDisjointSliceOp::new(op_erased_data) - .verify(&ctx) - .is_err(), - "DisjointSlice's fixed RawMut field cannot be reconstructed from Erased" - ); - - // Invalid: the row width is missing, so the slice would carry whatever - // slot 2 held. - let op_missing_width = Operation::new( - &mut ctx, - MirConstructDisjointSliceOp::get_concrete_op_info(), - vec![row_width_ty.into()], - vec![ptr_val, len_val], - vec![], - 0, - ); - assert!( - MirConstructDisjointSliceOp::new(op_missing_width) - .verify(&ctx) - .is_err(), - "Missing index-space operand" - ); - - // Invalid: a space-free slice given a third operand. - let op_extra = Operation::new( - &mut ctx, - MirConstructDisjointSliceOp::get_concrete_op_info(), - vec![plain_ty.into()], - vec![ptr_val, len_val, width_val], - vec![], - 0, - ); - assert!( - MirConstructDisjointSliceOp::new(op_extra) - .verify(&ctx) - .is_err(), - "Index-space operand for a space-free slice" - ); - - // Invalid: the row width operand has the wrong width, which would write a - // 64-bit value into the 32-bit row width slot. - let op_wrong_width_ty = Operation::new( - &mut ctx, - MirConstructDisjointSliceOp::get_concrete_op_info(), - vec![row_width_ty.into()], - vec![ptr_val, len_val, len_val], - vec![], - 0, - ); - assert!( - MirConstructDisjointSliceOp::new(op_wrong_width_ty) - .verify(&ctx) - .is_err(), - "Index-space operand type mismatch" - ); - - // Invalid: result is a plain slice, which `mir.construct_slice` owns. - let plain_slice_ty = MirSliceType::get(&mut ctx, f32_ty.into()); - let op_bad_res = Operation::new( - &mut ctx, - MirConstructDisjointSliceOp::get_concrete_op_info(), - vec![plain_slice_ty.into()], - vec![ptr_val, len_val], - vec![], - 0, - ); - assert!( - MirConstructDisjointSliceOp::new(op_bad_res) - .verify(&ctx) - .is_err(), - "Result must be a disjoint slice type" - ); -} - -#[test] -fn test_mir_construct_slice_verify() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let u8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let usize_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let u8_ptr_ty = MirPtrType::get_generic(&mut ctx, u8_ty.into(), false); - let u8_slice_ty = MirSliceType::get(&mut ctx, u8_ty.into()); - let i32_slice_ty = MirSliceType::get(&mut ctx, i32_ty.into()); - - let block = BasicBlock::new(&mut ctx, None, vec![u8_ptr_ty.into(), usize_ty.into()]); - let ptr_val = block.deref(&ctx).get_argument(0); - let len_val = block.deref(&ctx).get_argument(1); - - // Valid: (ptr to u8, usize len) -> slice of u8 - let op = Operation::new( - &mut ctx, - MirConstructSliceOp::get_concrete_op_info(), - vec![u8_slice_ty.into()], - vec![ptr_val, len_val], - vec![], - 0, - ); - assert!( - MirConstructSliceOp::new(op).verify(&ctx).is_ok(), - "Valid slice construction" - ); - - // Invalid: data pointer pointee does not match slice element type - let op_bad_elem = Operation::new( - &mut ctx, - MirConstructSliceOp::get_concrete_op_info(), - vec![i32_slice_ty.into()], - vec![ptr_val, len_val], - vec![], - 0, - ); - assert!( - MirConstructSliceOp::new(op_bad_elem).verify(&ctx).is_err(), - "Pointee/element mismatch" - ); - - // Invalid: operands swapped (length where the pointer should be) - let op_swapped = Operation::new( - &mut ctx, - MirConstructSliceOp::get_concrete_op_info(), - vec![u8_slice_ty.into()], - vec![len_val, ptr_val], - vec![], - 0, - ); - assert!( - MirConstructSliceOp::new(op_swapped).verify(&ctx).is_err(), - "Swapped operands" - ); - - // Invalid: result is not a slice type - let op_bad_res = Operation::new( - &mut ctx, - MirConstructSliceOp::get_concrete_op_info(), - vec![u8_ptr_ty.into()], - vec![ptr_val, len_val], - vec![], - 0, - ); - assert!( - MirConstructSliceOp::new(op_bad_res).verify(&ctx).is_err(), - "Non-slice result type" - ); -} - -#[test] -fn test_slice_carrier_cannot_launder_pointer_kind() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let u8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); - let usize_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let raw_mut_ptr = - MirPtrType::get_generic_with_kind(&mut ctx, u8_ty.into(), true, MirPointerKind::RawMut); - let erased_ptr = MirPtrType::get_generic(&mut ctx, u8_ty.into(), true); - let erased_const_ptr = MirPtrType::get_generic(&mut ctx, u8_ty.into(), false); - let global_raw_mut_ptr = MirPtrType::get_with_kind( - &mut ctx, - u8_ty.into(), - true, - dialect_mir::types::address_space::GLOBAL, - MirPointerKind::RawMut, - ); - let unique_ptr = - MirPtrType::get_generic_with_kind(&mut ctx, u8_ty.into(), true, MirPointerKind::UniqueRef); - let raw_mut_slice = MirSliceType::get_with_kind(&mut ctx, u8_ty.into(), MirPointerKind::RawMut); - let unique_slice = - MirSliceType::get_with_kind(&mut ctx, u8_ty.into(), MirPointerKind::UniqueRef); - let erased_slice = MirSliceType::get(&mut ctx, u8_ty.into()); - let erased_mut_slice = MirSliceType::get_with_mutability(&mut ctx, u8_ty.into(), true); - let block = BasicBlock::new( - &mut ctx, - None, - vec![ - raw_mut_ptr.into(), - erased_ptr.into(), - erased_const_ptr.into(), - global_raw_mut_ptr.into(), - usize_ty.into(), - raw_mut_slice.into(), - erased_slice.into(), - erased_mut_slice.into(), - ], - ); - let raw_mut = block.deref(&ctx).get_argument(0); - let erased = block.deref(&ctx).get_argument(1); - let erased_const = block.deref(&ctx).get_argument(2); - let global_raw_mut = block.deref(&ctx).get_argument(3); - let len = block.deref(&ctx).get_argument(4); - let raw_slice = block.deref(&ctx).get_argument(5); - let erased_slice_value = block.deref(&ctx).get_argument(6); - let erased_mut_slice_value = block.deref(&ctx).get_argument(7); - - let construct = |ctx: &mut Context, data, result_ty| { - Operation::new( - ctx, - MirConstructSliceOp::get_concrete_op_info(), - vec![result_ty], - vec![data, len], - vec![], - 0, - ) - }; - assert!( - MirConstructSliceOp::new(construct(&mut ctx, raw_mut, raw_mut_slice.into())) - .verify(&ctx) - .is_ok() - ); - assert!( - MirConstructSliceOp::new(construct(&mut ctx, raw_mut, unique_slice.into())) - .verify(&ctx) - .is_err(), - "construct_slice must not turn RawMut into UniqueRef" - ); - assert!( - MirConstructSliceOp::new(construct(&mut ctx, erased, raw_mut_slice.into())) - .verify(&ctx) - .is_err(), - "construct_slice must not recover RawMut from Erased" - ); - assert!( - MirConstructSliceOp::new(construct(&mut ctx, global_raw_mut, raw_mut_slice.into())) - .verify(&ctx) - .is_err(), - "ordinary slice carriers always use generic address space" - ); - assert!( - MirConstructSliceOp::new(construct(&mut ctx, erased_const, erased_slice.into())) - .verify(&ctx) - .is_ok(), - "an immutable Erased data pointer constructs an immutable Erased slice" - ); - - let extract = |ctx: &mut Context, result_ty| { - let op = Operation::new( - ctx, - MirExtractFieldOp::get_concrete_op_info(), - vec![result_ty], - vec![raw_slice], - vec![], - 0, - ); - MirExtractFieldOp::new(op).set_attr_index(ctx, FieldIndexAttr(0)); - op - }; - assert!( - MirExtractFieldOp::new(extract(&mut ctx, raw_mut_ptr.into())) - .verify(&ctx) - .is_ok() - ); - assert!( - MirExtractFieldOp::new(extract(&mut ctx, erased_ptr.into())) - .verify(&ctx) - .is_ok(), - "slice extraction may deliberately erase a concrete kind" - ); - assert!( - MirExtractFieldOp::new(extract(&mut ctx, unique_ptr.into())) - .verify(&ctx) - .is_err(), - "slice extraction must not change RawMut into UniqueRef" - ); - assert!( - MirExtractFieldOp::new(extract(&mut ctx, global_raw_mut_ptr.into())) - .verify(&ctx) - .is_err(), - "ordinary slice extraction cannot invent a non-generic address space" - ); - assert!( - MirExtractFieldOp::new(extract(&mut ctx, erased_const_ptr.into())) - .verify(&ctx) - .is_err(), - "erasing a concrete slice kind must preserve machine mutability" - ); - - let extract_erased = |ctx: &mut Context, result_ty| { - let op = Operation::new( - ctx, - MirExtractFieldOp::get_concrete_op_info(), - vec![result_ty], - vec![erased_slice_value], - vec![], - 0, - ); - MirExtractFieldOp::new(op).set_attr_index(ctx, FieldIndexAttr(0)); - op - }; - assert!( - MirExtractFieldOp::new(extract_erased(&mut ctx, erased_const_ptr.into())) - .verify(&ctx) - .is_ok(), - "an immutable Erased slice extracts an immutable Erased data pointer" - ); - assert!( - MirExtractFieldOp::new(extract_erased(&mut ctx, erased_ptr.into())) - .verify(&ctx) - .is_err(), - "an immutable Erased slice cannot manufacture writable data-pointer evidence" - ); - - let extract_erased_mut = |ctx: &mut Context, result_ty| { - let op = Operation::new( - ctx, - MirExtractFieldOp::get_concrete_op_info(), - vec![result_ty], - vec![erased_mut_slice_value], - vec![], - 0, - ); - MirExtractFieldOp::new(op).set_attr_index(ctx, FieldIndexAttr(0)); - op - }; - assert!( - MirExtractFieldOp::new(extract_erased_mut(&mut ctx, erased_ptr.into())) - .verify(&ctx) - .is_ok(), - "a mutable Erased slice retains writable data-pointer evidence" - ); - - let immutable_slice_reborrow = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![unique_slice.into()], - vec![erased_slice_value], - vec![], - 0, - ); - let immutable_slice_reborrow = MirCastOp::new(immutable_slice_reborrow); - immutable_slice_reborrow.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - immutable_slice_reborrow - .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); - assert!( - immutable_slice_reborrow.verify(&ctx).is_err(), - "an immutable Erased slice cannot be reborrowed as UniqueRef" - ); - - let mutable_slice_reborrow = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![unique_slice.into()], - vec![erased_mut_slice_value], - vec![], - 0, - ); - let mutable_slice_reborrow = MirCastOp::new(mutable_slice_reborrow); - mutable_slice_reborrow.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - mutable_slice_reborrow - .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); - assert!( - mutable_slice_reborrow.verify(&ctx).is_ok(), - "a mutable Erased slice may establish UniqueRef at a real reborrow boundary" - ); -} - -#[test] -fn test_mir_arithmetic_verify() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into(), i32_ty.into()]); - let lhs = block.deref(&ctx).get_argument(0); - - let check_bin_op = |opid: ( - fn(pliron::context::Ptr) -> pliron::op::OpObj, - std::any::TypeId, - ), - name: &str| { - let mut context = Context::new(); - dialect_mir::register(&mut context); - let ty = IntegerType::get(&context, 32, Signedness::Signed); - let blk = BasicBlock::new(&mut context, None, vec![ty.into(), ty.into()]); - let l = blk.deref(&context).get_argument(0); - let r = blk.deref(&context).get_argument(1); - - let op = Operation::new(&mut context, opid, vec![ty.into()], vec![l, r], vec![], 0); - assert!(op.verify(&context).is_ok(), "Valid {}", name); - - let f32_t = FP32Type::get(&context); - let blk2 = BasicBlock::new(&mut context, None, vec![f32_t.into()]); - let f32_val = blk2.deref(&context).get_argument(0); - - let op_bad = Operation::new( - &mut context, - opid, - vec![ty.into()], - vec![l, f32_val], - vec![], - 0, - ); - assert!(op_bad.verify(&context).is_err(), "Type mismatch {}", name); - }; - - check_bin_op(MirAddOp::get_concrete_op_info(), "Add"); - check_bin_op(MirSubOp::get_concrete_op_info(), "Sub"); - check_bin_op(MirMulOp::get_concrete_op_info(), "Mul"); - check_bin_op(MirDivOp::get_concrete_op_info(), "Div"); - check_bin_op(MirRemOp::get_concrete_op_info(), "Rem"); - - let op_neg = Operation::new( - &mut ctx, - MirNegOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![lhs], - vec![], - 0, - ); - assert!(op_neg.verify(&ctx).is_ok(), "Valid Neg"); - - let f32_ty = FP32Type::get(&ctx); - let op_neg_bad = Operation::new( - &mut ctx, - MirNegOp::get_concrete_op_info(), - vec![f32_ty.into()], - vec![lhs], - vec![], - 0, - ); - assert!(op_neg_bad.verify(&ctx).is_err(), "Neg type mismatch"); - - let op_not = Operation::new( - &mut ctx, - MirNotOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![lhs], - vec![], - 0, - ); - assert!(op_not.verify(&ctx).is_ok(), "Valid Not"); -} - -#[test] -fn test_mir_misc_verify() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signed); - let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); - - // 1. MirConstantOp - let i32_signless = IntegerType::get(&ctx, 32, Signedness::Signless); - let width = NonZeroUsize::new(32).unwrap(); - let apint = APInt::from_u32(42, width); - let int_attr = IntegerAttr::new(i32_signless, apint); - - let const_op_ptr = Operation::new( - &mut ctx, - MirConstantOp::get_concrete_op_info(), - vec![i32_signless.into()], - vec![], - vec![], - 0, - ); - let const_op = MirConstantOp::new(const_op_ptr); - const_op.set_attr_value(&ctx, int_attr); - assert!(const_op.verify(&ctx).is_ok(), "Valid Constant"); - - // Mismatch type - let i64_signless = IntegerType::get(&ctx, 64, Signedness::Signless); - let i64_width = NonZeroUsize::new(64).unwrap(); - let i64_attr = IntegerAttr::new(i64_signless, APInt::from_u64(42, i64_width)); - const_op.set_attr_value(&ctx, i64_attr); - assert!(const_op.verify(&ctx).is_err(), "Constant type mismatch"); - - // 2. MirCastOp - let block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); - let arg = block.deref(&ctx).get_argument(0); - - let cast_op = Operation::new( - &mut ctx, - MirCastOp::get_concrete_op_info(), - vec![i64_ty.into()], - vec![arg], - vec![], - 0, - ); - MirCastOp::new(cast_op).set_attr_cast_kind(&ctx, MirCastKindAttr::IntToInt); - assert!(MirCastOp::new(cast_op).verify(&ctx).is_ok(), "Valid Cast"); - - // 3. MirCheckedAddOp - let tuple_ty = MirTupleType::get(&mut ctx, vec![i32_ty.into(), i1_ty.into()]); - let block2 = BasicBlock::new(&mut ctx, None, vec![i32_ty.into(), i32_ty.into()]); - let lhs = block2.deref(&ctx).get_argument(0); - let rhs = block2.deref(&ctx).get_argument(1); - - let checked_add = Operation::new( - &mut ctx, - MirCheckedAddOp::get_concrete_op_info(), - vec![tuple_ty.into()], - vec![lhs, rhs], - vec![], - 0, - ); - assert!( - MirCheckedAddOp::new(checked_add).verify(&ctx).is_ok(), - "Valid CheckedAdd" - ); - - // Invalid result type (not tuple) - let checked_add_bad = Operation::new( - &mut ctx, - MirCheckedAddOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![lhs, rhs], - vec![], - 0, - ); - assert!( - MirCheckedAddOp::new(checked_add_bad).verify(&ctx).is_err(), - "CheckedAdd bad result" - ); -} - -#[test] -fn test_mir_comparison_verify() { - let check_cmp = |opid: ( - fn(pliron::context::Ptr) -> pliron::op::OpObj, - std::any::TypeId, - ), - name: &str| { - let mut context = Context::new(); - dialect_mir::register(&mut context); - let ty = IntegerType::get(&context, 32, Signedness::Signed); - let res_ty = IntegerType::get(&context, 1, Signedness::Signless); - let blk = BasicBlock::new(&mut context, None, vec![ty.into(), ty.into()]); - let l = blk.deref(&context).get_argument(0); - let r = blk.deref(&context).get_argument(1); - - let op = Operation::new( - &mut context, - opid, - vec![res_ty.into()], - vec![l, r], - vec![], - 0, - ); - assert!(op.verify(&context).is_ok(), "Valid {}", name); - - // Invalid operand types - let f32_ty = FP32Type::get(&context); - let blk2 = BasicBlock::new(&mut context, None, vec![f32_ty.into()]); - let f32_val = blk2.deref(&context).get_argument(0); - let op_bad = Operation::new( - &mut context, - opid, - vec![res_ty.into()], - vec![l, f32_val], - vec![], - 0, - ); - assert!(op_bad.verify(&context).is_err(), "Type mismatch {}", name); - - // Invalid result type - let op_bad_res = Operation::new( - &mut context, - opid, - vec![ty.into()], // i32 result instead of i1 - vec![l, r], - vec![], - 0, - ); - assert!( - op_bad_res.verify(&context).is_err(), - "Result type mismatch {}", - name - ); - }; - - check_cmp(MirEqOp::get_concrete_op_info(), "Eq"); - check_cmp(MirNeOp::get_concrete_op_info(), "Ne"); - check_cmp(MirLtOp::get_concrete_op_info(), "Lt"); - check_cmp(MirLeOp::get_concrete_op_info(), "Le"); - check_cmp(MirGtOp::get_concrete_op_info(), "Gt"); - check_cmp(MirGeOp::get_concrete_op_info(), "Ge"); - - let mut context = Context::new(); - dialect_mir::register(&mut context); - let i8_ty = IntegerType::get(&context, 8, Signedness::Signed); - let i32_ty = IntegerType::get(&context, 32, Signedness::Signed); - let unit = |name: &str| EnumVariant::unit(name.to_string()); - let ordering_ty = MirEnumType::get( - &mut context, - "Ordering".to_string(), - i8_ty.into(), - vec![255, 0, 1], - vec![unit("Less"), unit("Equal"), unit("Greater")], - ); - let blk = BasicBlock::new(&mut context, None, vec![i32_ty.into(), i32_ty.into()]); - let lhs = blk.deref(&context).get_argument(0); - let rhs = blk.deref(&context).get_argument(1); - let two_variant_ty = MirEnumType::get( - &mut context, - "Two".to_string(), - i8_ty.into(), - vec![0, 1], - vec![unit("A"), unit("B")], - ); - // Payload variants disqualify the Ordering shape. - let payload_ty = MirEnumType::get( - &mut context, - "ThreeWithPayload".to_string(), - i8_ty.into(), - vec![0, 1, 2], - vec![ - unit("A"), - EnumVariant::new("B".to_string(), vec![i32_ty.into()]), - unit("C"), - ], - ); - let mut check_cmp_result = |result_ty, valid| { - let op = Operation::new( - &mut context, - MirCmpOp::get_concrete_op_info(), - vec![result_ty], - vec![lhs, rhs], - vec![], - 0, - ); - assert_eq!(op.verify(&context).is_ok(), valid); - }; - check_cmp_result(ordering_ty.into(), true); - check_cmp_result(i32_ty.into(), false); - check_cmp_result(two_variant_ty.into(), false); - check_cmp_result(payload_ty.into(), false); - - // Float operands are rejected: rustc never emits BinOp::Cmp on floats. - let f32_ty = FP32Type::get(&context); - let fblk = BasicBlock::new(&mut context, None, vec![f32_ty.into(), f32_ty.into()]); - let flhs = fblk.deref(&context).get_argument(0); - let frhs = fblk.deref(&context).get_argument(1); - let float_cmp = Operation::new( - &mut context, - MirCmpOp::get_concrete_op_info(), - vec![ordering_ty.into()], - vec![flhs, frhs], - vec![], - 0, - ); - assert!( - float_cmp.verify(&context).is_err(), - "float mir.cmp must be rejected" - ); -} - -#[test] -fn test_mir_func_verify() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let func_ty = FunctionType::get(&ctx, vec![i32_ty.into()], vec![]); - let func_ty_attr = TypeAttr::new(func_ty.into()); - - // Valid Function - let op_ptr = Operation::new( - &mut ctx, - MirFuncOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 1, - ); - let mir_func = MirFuncOp::new(&mut ctx, op_ptr, func_ty_attr.clone()); - - // Add entry block with correct argument - let region = mir_func.get_operation().deref(&ctx).get_region(0); - let entry_block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); - entry_block.insert_at_front(region, &ctx); - - assert!(mir_func.verify(&ctx).is_ok(), "Valid MirFuncOp"); - - // Invalid: Argument count mismatch - let op_ptr2 = Operation::new( - &mut ctx, - MirFuncOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 1, - ); - let mir_func2 = MirFuncOp::new(&mut ctx, op_ptr2, func_ty_attr.clone()); - let region2 = mir_func2.get_operation().deref(&ctx).get_region(0); - // Block with 0 args - let entry_block2 = BasicBlock::new(&mut ctx, None, vec![]); - entry_block2.insert_at_front(region2, &ctx); - - assert!( - mir_func2.verify(&ctx).is_err(), - "MirFuncOp argument count mismatch" - ); - - // Invalid: Argument type mismatch - let op_ptr3 = Operation::new( - &mut ctx, - MirFuncOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 1, - ); - let mir_func3 = MirFuncOp::new(&mut ctx, op_ptr3, func_ty_attr); - let region3 = mir_func3.get_operation().deref(&ctx).get_region(0); - let f32_ty = FP32Type::get(&ctx); - let entry_block3 = BasicBlock::new(&mut ctx, None, vec![f32_ty.into()]); - entry_block3.insert_at_front(region3, &ctx); - - assert!( - mir_func3.verify(&ctx).is_err(), - "MirFuncOp argument type mismatch" - ); -} - -#[test] -fn test_mir_assign_verify() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); - let val = block.deref(&ctx).get_argument(0); - - let op = Operation::new( - &mut ctx, - MirAssignOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![val], - vec![], - 0, - ); - assert!( - MirAssignOp::new(op).verify(&ctx).is_ok(), - "Valid MirAssignOp" - ); - - let f32_ty = FP32Type::get(&ctx); - let op_bad = Operation::new( - &mut ctx, - MirAssignOp::get_concrete_op_info(), - vec![f32_ty.into()], - vec![val], - vec![], - 0, - ); - assert!( - MirAssignOp::new(op_bad).verify(&ctx).is_err(), - "MirAssignOp type mismatch" - ); -} - -#[test] -fn test_mir_call_verify() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let op = Operation::new( - &mut ctx, - MirCallOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 0, - ); - let call_op = MirCallOp::new(op); - - // Missing attribute - assert!(call_op.verify(&ctx).is_err(), "MirCallOp missing attribute"); - - // With attribute - let name = StringAttr::new("my_func".to_string()); - call_op.set_attr_callee(&ctx, name); - assert!(call_op.verify(&ctx).is_ok(), "Valid MirCallOp"); -} - -#[test] -fn test_mir_store_verify() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), false); - let block = BasicBlock::new(&mut ctx, None, vec![ptr_ty.into(), i32_ty.into()]); - let ptr_val = block.deref(&ctx).get_argument(0); - let val = block.deref(&ctx).get_argument(1); - - let op = Operation::new( - &mut ctx, - MirStoreOp::get_concrete_op_info(), - vec![], - vec![ptr_val, val], - vec![], - 0, - ); - assert!(MirStoreOp::new(op).verify(&ctx).is_ok(), "Valid MirStoreOp"); - - // Invalid: store to non-ptr - let op_bad_ptr = Operation::new( - &mut ctx, - MirStoreOp::get_concrete_op_info(), - vec![], - vec![val, val], - vec![], - 0, - ); - assert!( - MirStoreOp::new(op_bad_ptr).verify(&ctx).is_err(), - "MirStoreOp non-ptr dest" - ); - - // Invalid: type mismatch - let f32_ty = FP32Type::get(&ctx); - let block2 = BasicBlock::new(&mut ctx, None, vec![f32_ty.into()]); - let f32_val = block2.deref(&ctx).get_argument(0); - let op_bad_type = Operation::new( - &mut ctx, - MirStoreOp::get_concrete_op_info(), - vec![], - vec![ptr_val, f32_val], - vec![], - 0, - ); - assert!( - MirStoreOp::new(op_bad_type).verify(&ctx).is_err(), - "MirStoreOp type mismatch" - ); -} - -#[test] -fn test_mir_store_volatile_is_not_promotable() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), false); - let block = BasicBlock::new(&mut ctx, None, vec![ptr_ty.into(), i32_ty.into()]); - let ptr_val = block.deref(&ctx).get_argument(0); - let val = block.deref(&ctx).get_argument(1); - - let op = Operation::new( - &mut ctx, - MirStoreOp::get_concrete_op_info(), - vec![], - vec![ptr_val, val], - vec![], - 0, - ); - let mir_store = MirStoreOp::new(op); - let alloc_info = AllocInfo { - ptr: ptr_val, - ty: i32_ty.into(), - }; - - assert!(!mir_store.is_volatile(&ctx)); - match mir_store.promotion_kind(&ctx, &alloc_info) { - PromotableOpKind::Store(stored) => assert!(stored == val), - _ => panic!("non-volatile store should be promotable"), - } - - mir_store.set_volatile(&mut ctx, true); - - assert!(mir_store.is_volatile(&ctx)); - assert!(matches!( - mir_store.promotion_kind(&ctx, &alloc_info), - PromotableOpKind::NonPromotableUse - )); -} - -#[test] -fn test_mir_global_alloc_verify() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let f32_ty = FP32Type::get(&ctx); - - // Helper: build a MirGlobalAllocOp whose result pointer is in `ptr_ty` - // address space, with valid attributes. - let build = |ctx: &mut Context, ptr_ty: pliron::r#type::TypedHandle| { - let op = Operation::new( - ctx, - MirGlobalAllocOp::get_concrete_op_info(), - vec![ptr_ty.into()], - vec![], - vec![], - 0, - ); - let alloc = MirGlobalAllocOp::new(op); - alloc.set_attr_global_type(ctx, TypeAttr::new(f32_ty.into())); - alloc.set_attr_global_key(ctx, StringAttr::new("k".to_string())); - alloc - }; - - // Global memory (addrspace 1) — the original allowed space. - let ptr_global = MirPtrType::get_global(&mut ctx, f32_ty.into(), true); - assert!( - build(&mut ctx, ptr_global).verify(&ctx).is_ok(), - "global addrspace accepted" - ); - - // Constant memory (addrspace 4) — added for `#[constant]` support. - let ptr_const = MirPtrType::get_constant(&mut ctx, f32_ty.into(), true); - assert!( - build(&mut ctx, ptr_const).verify(&ctx).is_ok(), - "constant addrspace accepted" - ); - - // Shared memory (addrspace 3) — must be rejected. - let ptr_shared = MirPtrType::get_shared(&mut ctx, f32_ty.into(), true); - assert!( - build(&mut ctx, ptr_shared).verify(&ctx).is_err(), - "shared addrspace rejected" - ); - - let ptr_unique_global = MirPtrType::get_with_kind( - &mut ctx, - f32_ty.into(), - true, - dialect_mir::types::address_space::GLOBAL, - MirPointerKind::UniqueRef, - ); - assert!( - build(&mut ctx, ptr_unique_global).verify(&ctx).is_err(), - "a storage allocation cannot directly claim UniqueRef" - ); - - // Missing required attributes. - let no_attrs = Operation::new( - &mut ctx, - MirGlobalAllocOp::get_concrete_op_info(), - vec![ptr_global.into()], - vec![], - vec![], - 0, - ); - assert!( - MirGlobalAllocOp::new(no_attrs).verify(&ctx).is_err(), - "missing attributes rejected" - ); -} - -#[test] -fn test_mir_set_discriminant_verify() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signed); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let unit = |name: &str| EnumVariant::unit(name.to_string()); - - let enum_ty = MirEnumType::get( - &mut ctx, - "DeviceState".to_string(), - i8_ty.into(), - vec![0, 1], - vec![ - unit("Empty"), - EnumVariant::new("Full".to_string(), vec![i32_ty.into()]), - ], - ); - - let enum_ptr_ty = MirPtrType::get_generic(&mut ctx, enum_ty.into(), true); - let blk = BasicBlock::new(&mut ctx, None, vec![enum_ptr_ty.into(), i8_ty.into()]); - let enum_ptr = blk.deref(&ctx).get_argument(0); - let discr_val = blk.deref(&ctx).get_argument(1); - - // Malformed IR must be diagnosed without panicking, regardless of whether - // the generated operand-count interface happens to run first. - let op_no_operands = Operation::new( - &mut ctx, - MirSetDiscriminantOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 0, - ); - let no_operands = MirSetDiscriminantOp::new(op_no_operands); - no_operands.set_attr_set_discriminant_variant_index(&ctx, VariantIndexAttr(0)); - assert!(no_operands.verify(&ctx).is_err(), "zero operands rejected"); - - let op_extra_operand = Operation::new( - &mut ctx, - MirSetDiscriminantOp::get_concrete_op_info(), - vec![], - vec![enum_ptr, discr_val], - vec![], - 0, - ); - let extra_operand = MirSetDiscriminantOp::new(op_extra_operand); - extra_operand.set_attr_set_discriminant_variant_index(&ctx, VariantIndexAttr(0)); - assert!( - extra_operand.verify(&ctx).is_err(), - "extra operand rejected" - ); - - // Valid: pointer to enum plus the semantic target variant attribute. - let op_valid = Operation::new( - &mut ctx, - MirSetDiscriminantOp::get_concrete_op_info(), - vec![], - vec![enum_ptr], - vec![], - 0, - ); - let valid = MirSetDiscriminantOp::new(op_valid); - valid.set_attr_set_discriminant_variant_index(&ctx, VariantIndexAttr(1)); - valid.set_attr_set_discriminant_enum_ty(&ctx, TypeAttr::new(enum_ty.into())); - assert!(valid.verify(&ctx).is_ok(), "Valid set_discriminant"); - - // Invalid: first operand is not a pointer. - let op_bad_ptr = Operation::new( - &mut ctx, - MirSetDiscriminantOp::get_concrete_op_info(), - vec![], - vec![discr_val], - vec![], - 0, - ); - assert!( - MirSetDiscriminantOp::new(op_bad_ptr).verify(&ctx).is_err(), - "Non-pointer enum operand rejected" - ); - - // Invalid: SetDiscriminant writes memory and therefore cannot accept an - // immutable pointer even when the pointee is the right enum. - let immutable_ptr_ty = MirPtrType::get_generic(&mut ctx, enum_ty.into(), false); - let immutable_block = BasicBlock::new(&mut ctx, None, vec![immutable_ptr_ty.into()]); - let immutable_ptr = immutable_block.deref(&ctx).get_argument(0); - let op_immutable = Operation::new( - &mut ctx, - MirSetDiscriminantOp::get_concrete_op_info(), - vec![], - vec![immutable_ptr], - vec![], - 0, - ); - let immutable = MirSetDiscriminantOp::new(op_immutable); - immutable.set_attr_set_discriminant_variant_index(&ctx, VariantIndexAttr(0)); - assert!( - immutable.verify(&ctx).is_err(), - "immutable pointer rejected" - ); - - // Invalid: pointer does not point to an enum. - let i32_ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), true); - let blk_i32 = BasicBlock::new(&mut ctx, None, vec![i32_ptr_ty.into(), i8_ty.into()]); - let i32_ptr = blk_i32.deref(&ctx).get_argument(0); - let op_bad_pointee = Operation::new( - &mut ctx, - MirSetDiscriminantOp::get_concrete_op_info(), - vec![], - vec![i32_ptr], - vec![], - 0, - ); - let bad_pointee = MirSetDiscriminantOp::new(op_bad_pointee); - bad_pointee.set_attr_set_discriminant_variant_index(&ctx, VariantIndexAttr(0)); - assert!( - bad_pointee.verify(&ctx).is_err(), - "Non-enum pointee rejected" - ); - - // Invalid: target attribute is required. - let op_missing_target = Operation::new( - &mut ctx, - MirSetDiscriminantOp::get_concrete_op_info(), - vec![], - vec![enum_ptr], - vec![], - 0, - ); - assert!( - MirSetDiscriminantOp::new(op_missing_target) - .verify(&ctx) - .is_err(), - "Missing target rejected" - ); - - // Invalid: target index is out of bounds. - let op_oob = Operation::new( - &mut ctx, - MirSetDiscriminantOp::get_concrete_op_info(), - vec![], - vec![enum_ptr], - vec![], - 0, - ); - let oob = MirSetDiscriminantOp::new(op_oob); - oob.set_attr_set_discriminant_variant_index(&ctx, VariantIndexAttr(2)); - assert!(oob.verify(&ctx).is_err(), "Out-of-bounds target rejected"); -} - -#[test] -fn test_mir_enum_ops_malformed_arity_is_diagnostic() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - let i8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let unit_enum = MirEnumType::get( - &mut ctx, - "Unit".into(), - i8_ty.into(), - vec![0], - vec![EnumVariant::unit("Only".into())], - ); - let payload_enum = MirEnumType::get( - &mut ctx, - "Payload".into(), - i8_ty.into(), - vec![0], - vec![EnumVariant::new("Only".into(), vec![i32_ty.into()])], - ); - let block = BasicBlock::new(&mut ctx, None, vec![unit_enum.into(), payload_enum.into()]); - let unit_value = block.deref(&ctx).get_argument(0); - let payload_value = block.deref(&ctx).get_argument(1); - - let construct_no_result = Operation::new( - &mut ctx, - MirConstructEnumOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 0, - ); - let construct_no_result = MirConstructEnumOp::new(construct_no_result); - construct_no_result.set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(0)); - assert!(construct_no_result.verify(&ctx).is_err()); - - let construct_extra_result = Operation::new( - &mut ctx, - MirConstructEnumOp::get_concrete_op_info(), - vec![unit_enum.into(), unit_enum.into()], - vec![], - vec![], - 0, - ); - let construct_extra_result = MirConstructEnumOp::new(construct_extra_result); - construct_extra_result.set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(0)); - assert!(construct_extra_result.verify(&ctx).is_err()); - - let get_empty = Operation::new( - &mut ctx, - MirGetDiscriminantOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 0, - ); - assert!(MirGetDiscriminantOp::new(get_empty).verify(&ctx).is_err()); - let get_extra = Operation::new( - &mut ctx, - MirGetDiscriminantOp::get_concrete_op_info(), - vec![i8_ty.into(), i8_ty.into()], - vec![unit_value, unit_value], - vec![], - 0, - ); - assert!(MirGetDiscriminantOp::new(get_extra).verify(&ctx).is_err()); - - let payload_empty = Operation::new( - &mut ctx, - MirEnumPayloadOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 0, - ); - let payload_empty = MirEnumPayloadOp::new(payload_empty); - payload_empty.set_attr_payload_variant_index(&ctx, VariantIndexAttr(0)); - payload_empty.set_attr_payload_field_index(&ctx, FieldIndexAttr(0)); - assert!(payload_empty.verify(&ctx).is_err()); - - let payload_extra = Operation::new( - &mut ctx, - MirEnumPayloadOp::get_concrete_op_info(), - vec![i32_ty.into(), i32_ty.into()], - vec![payload_value, payload_value], - vec![], - 0, - ); - let payload_extra = MirEnumPayloadOp::new(payload_extra); - payload_extra.set_attr_payload_variant_index(&ctx, VariantIndexAttr(0)); - payload_extra.set_attr_payload_field_index(&ctx, FieldIndexAttr(0)); - assert!(payload_extra.verify(&ctx).is_err()); -} - -#[test] -fn test_mir_enum_payload_rejects_uninhabited_variant() { - use dialect_mir::types::{EnumCarrierKind, EnumEncoding, EnumLayoutKind}; - - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - let i8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let enum_ty = MirEnumType::get_with_encoding( - &mut ctx, - "HasImpossibleField".into(), - i8_ty.into(), - vec![0, 1], - vec![ - EnumVariant::unit("Live".into()), - // An uninhabited variant's unused physical offsets need not fit - // the object. The verifier must stop them reaching lowering. - EnumVariant::new_with_layout( - "Impossible".into(), - vec![i32_ty.into()], - vec![64], - vec![4], - ), - ], - EnumEncoding { - tag_offset: 0, - total_size: 1, - abi_align: 1, - layout_kind: EnumLayoutKind::Direct, - carrier_kind: EnumCarrierKind::Integer, - carrier_width: 8, - variant_inhabited: vec![1, 0], - ..EnumEncoding::default() - }, - ); - assert!(enum_ty.verify(&ctx).is_ok()); - - let block = BasicBlock::new(&mut ctx, None, vec![enum_ty.into()]); - let value = block.deref(&ctx).get_argument(0); - let op = Operation::new( - &mut ctx, - MirEnumPayloadOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![value], - vec![], - 0, - ); - let payload = MirEnumPayloadOp::new(op); - payload.set_attr_payload_variant_index(&ctx, VariantIndexAttr(1)); - payload.set_attr_payload_field_index(&ctx, FieldIndexAttr(0)); - assert!(payload.verify(&ctx).is_err()); -} - -/// Reachable-but-dead MIR (e.g. the residual arms of `array::try_from_fn`) -/// can name uninhabited variants. The importer must lower such reads and -/// constructions to typed undefs; if it ever emits the real ops again, these -/// verifiers are the loud stop. -#[test] -fn test_uninhabited_enum_construct_and_discriminant_fail_verification() { - use dialect_mir::types::{EnumCarrierKind, EnumEncoding, EnumLayoutKind}; - - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - let i8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - - // Constructing an uninhabited variant must fail verification. - let partial = MirEnumType::get_with_encoding( - &mut ctx, - "HasImpossibleVariant".into(), - i8_ty.into(), - vec![0, 1], - vec![ - EnumVariant::unit("Live".into()), - EnumVariant::new_with_layout( - "Impossible".into(), - vec![i32_ty.into()], - vec![0], - vec![4], - ), - ], - EnumEncoding { - tag_offset: 0, - total_size: 8, - abi_align: 4, - layout_kind: EnumLayoutKind::Direct, - carrier_kind: EnumCarrierKind::Integer, - carrier_width: 8, - variant_inhabited: vec![1, 0], - ..EnumEncoding::default() - }, - ); - let block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); - let field = block.deref(&ctx).get_argument(0); - let construct = Operation::new( - &mut ctx, - MirConstructEnumOp::get_concrete_op_info(), - vec![partial.into()], - vec![field], - vec![], - 0, - ); - let construct = MirConstructEnumOp::new(construct); - construct.set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(1)); - assert!(construct.verify(&ctx).is_err()); - - // Reading the discriminant of a fully uninhabited enum must fail - // verification. - let never = MirEnumType::get_with_encoding( - &mut ctx, - "Never".into(), - i8_ty.into(), - vec![], - vec![], - EnumEncoding { - tag_offset: 0, - total_size: 0, - abi_align: 1, - layout_kind: EnumLayoutKind::Empty, - carrier_kind: EnumCarrierKind::None, - carrier_width: 0, - variant_inhabited: vec![], - ..EnumEncoding::default() - }, - ); - assert!(never.verify(&ctx).is_ok()); - let never_block = BasicBlock::new(&mut ctx, None, vec![never.into()]); - let never_value = never_block.deref(&ctx).get_argument(0); - let get_discriminant = Operation::new( - &mut ctx, - MirGetDiscriminantOp::get_concrete_op_info(), - vec![i8_ty.into()], - vec![never_value], - vec![], - 0, - ); - let get_discriminant = MirGetDiscriminantOp::new(get_discriminant); - assert!(get_discriminant.verify(&ctx).is_err()); -} - -#[test] -fn test_mir_field_addr_tuple_pointee_verify() { - // `(u8, u32)` laid out the way rustc actually places it: the u32 field - // first in memory for alignment, so declaration index 0 (`u8`) lands at - // byte offset 4 and declaration index 1 (`u32`) lands at byte offset 0. - // `field_addr`'s `field_index` attribute is a DECLARATION index (it names - // `.0`/`.1` as written), so this test only passes if the op resolves the - // field's type through `MirTupleType::get_types()` (declaration order) - // rather than assuming identity with memory order. - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let u8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - - let tuple_ty = MirTupleType::get_with_layout( - &mut ctx, - vec![u8_ty.into(), u32_ty.into()], - vec![1, 0], - vec![4, 0], - 8, - 4, - ); - - let tuple_ptr_ty = MirPtrType::get_generic(&mut ctx, tuple_ty.into(), false); - let blk = BasicBlock::new(&mut ctx, None, vec![tuple_ptr_ty.into()]); - let tuple_ptr = blk.deref(&ctx).get_argument(0); - - let u8_ptr_ty = MirPtrType::get_generic(&mut ctx, u8_ty.into(), false); - let op_field0 = Operation::new( - &mut ctx, - MirFieldAddrOp::get_concrete_op_info(), - vec![u8_ptr_ty.into()], - vec![tuple_ptr], - vec![], - 0, - ); - let field0 = MirFieldAddrOp::new(op_field0); - field0.set_attr_field_index(&ctx, FieldIndexAttr(0)); - field0.set_attr_aggregate_ty(&ctx, TypeAttr::new(tuple_ty.into())); - assert!( - field0.verify(&ctx).is_ok(), - "tuple field 0 (u8) address accepted" - ); - - let u32_ptr_ty = MirPtrType::get_generic(&mut ctx, u32_ty.into(), false); - let op_field1 = Operation::new( - &mut ctx, - MirFieldAddrOp::get_concrete_op_info(), - vec![u32_ptr_ty.into()], - vec![tuple_ptr], - vec![], - 0, - ); - let field1 = MirFieldAddrOp::new(op_field1); - field1.set_attr_field_index(&ctx, FieldIndexAttr(1)); - field1.set_attr_aggregate_ty(&ctx, TypeAttr::new(tuple_ty.into())); - assert!( - field1.verify(&ctx).is_ok(), - "tuple field 1 (u32) address accepted" - ); - - // Result pointee type must match the DECLARED field type, not whatever - // sits at that byte offset: pointing field 0's result at u32 (field 1's - // type) must be rejected even though both are in-bounds indices. - let op_wrong_result_ty = Operation::new( - &mut ctx, - MirFieldAddrOp::get_concrete_op_info(), - vec![u32_ptr_ty.into()], - vec![tuple_ptr], - vec![], - 0, - ); - let wrong_result_ty = MirFieldAddrOp::new(op_wrong_result_ty); - wrong_result_ty.set_attr_field_index(&ctx, FieldIndexAttr(0)); - wrong_result_ty.set_attr_aggregate_ty(&ctx, TypeAttr::new(tuple_ty.into())); - assert!( - wrong_result_ty.verify(&ctx).is_err(), - "result pointee type mismatch rejected" - ); - - let op_out_of_bounds = Operation::new( - &mut ctx, - MirFieldAddrOp::get_concrete_op_info(), - vec![u8_ptr_ty.into()], - vec![tuple_ptr], - vec![], - 0, - ); - let out_of_bounds = MirFieldAddrOp::new(op_out_of_bounds); - out_of_bounds.set_attr_field_index(&ctx, FieldIndexAttr(2)); - out_of_bounds.set_attr_aggregate_ty(&ctx, TypeAttr::new(tuple_ty.into())); - assert!( - out_of_bounds.verify(&ctx).is_err(), - "out-of-bounds tuple field index rejected" - ); -} - -#[test] -fn test_mir_field_addr_tuple_pointee_store_verify() { - // The WRITE side of the tuple-pointee unlock: `t.1 = x` / `arr[i].1 = x` - // lower to `mir.field_addr` + `mir.store` through the field's address, so - // a tuple-pointee field address used as a store destination must pass - // verification too. Same reordered `(u8, u32)` layout as above (the u32 - // field first in memory), so the store type-checks against the DECLARED - // field type, not whatever occupies that memory slot. - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - - let u8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - - let tuple_ty = MirTupleType::get_with_layout( - &mut ctx, - vec![u8_ty.into(), u32_ty.into()], - vec![1, 0], - vec![4, 0], - 8, - 4, - ); - - let tuple_ptr_ty = MirPtrType::get_generic(&mut ctx, tuple_ty.into(), false); - let blk = BasicBlock::new( - &mut ctx, - None, - vec![tuple_ptr_ty.into(), u32_ty.into(), u8_ty.into()], - ); - let tuple_ptr = blk.deref(&ctx).get_argument(0); - let u32_val = blk.deref(&ctx).get_argument(1); - let u8_val = blk.deref(&ctx).get_argument(2); - - // `.1 = x`: address declaration field 1 (u32, memory slot 0) and store a - // u32 through it. - let u32_ptr_ty = MirPtrType::get_generic(&mut ctx, u32_ty.into(), false); - let op_field1 = Operation::new( - &mut ctx, - MirFieldAddrOp::get_concrete_op_info(), - vec![u32_ptr_ty.into()], - vec![tuple_ptr], - vec![], - 0, - ); - let field1 = MirFieldAddrOp::new(op_field1); - field1.set_attr_field_index(&ctx, FieldIndexAttr(1)); - field1.set_attr_aggregate_ty(&ctx, TypeAttr::new(tuple_ty.into())); - assert!( - field1.verify(&ctx).is_ok(), - "tuple field 1 (u32) address accepted as a store destination" - ); - let field1_ptr = op_field1.deref(&ctx).get_result(0); - - let op_store = Operation::new( - &mut ctx, - MirStoreOp::get_concrete_op_info(), - vec![], - vec![field1_ptr, u32_val], - vec![], - 0, - ); - assert!( - MirStoreOp::new(op_store).verify(&ctx).is_ok(), - "store through a tuple field address verifies" - ); - - // The stored value must match the DECLARED field type (`u32` for `.1`), - // not the type of the field sharing the tuple: a u8 store through the - // `.1` pointer is a type mismatch. - let op_store_wrong_ty = Operation::new( - &mut ctx, - MirStoreOp::get_concrete_op_info(), - vec![], - vec![field1_ptr, u8_val], - vec![], - 0, - ); - assert!( - MirStoreOp::new(op_store_wrong_ty).verify(&ctx).is_err(), - "store of a mismatched value type through a tuple field address rejected" - ); -} diff --git a/crates/dialect-mir/tests/ops_test/aggregates.rs b/crates/dialect-mir/tests/ops_test/aggregates.rs new file mode 100644 index 0000000000..545f81addc --- /dev/null +++ b/crates/dialect-mir/tests/ops_test/aggregates.rs @@ -0,0 +1,687 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use dialect_mir::{ + attributes::{FieldIndexAttr, MirCastKindAttr, MirPointerKindAuthorityAttr}, + ops::{ + MirCastOp, MirConstructDisjointSliceOp, MirConstructSliceOp, MirExtractFieldOp, + MirFieldAddrOp, MirStoreOp, + }, + types::{ + MirDisjointSliceType, MirPointerKind, MirPtrType, MirSliceType, MirTupleType, MirUnionType, + }, +}; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::TypeAttr, + types::{FP32Type, IntegerType, Signedness}, + }, + common_traits::Verify, + context::Context, + op::Op, + operation::Operation, +}; + +#[test] +fn test_mir_extract_field_verify() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let tuple_ty = MirTupleType::get(&mut ctx, vec![i32_ty.into(), i32_ty.into()]); + + let block = BasicBlock::new(&mut ctx, None, vec![tuple_ty.into()]); + let tuple_val = block.deref(&ctx).get_argument(0); + + let op = Operation::new( + &mut ctx, + MirExtractFieldOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![tuple_val], + vec![], + 0, + ); + let extract_op = MirExtractFieldOp::new(op); + extract_op.set_attr_index(&ctx, dialect_mir::attributes::FieldIndexAttr(0)); + assert!(extract_op.verify(&ctx).is_ok(), "Valid Tuple Extract"); + + let op_oob = Operation::new( + &mut ctx, + MirExtractFieldOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![tuple_val], + vec![], + 0, + ); + let extract_op_oob = MirExtractFieldOp::new(op_oob); + extract_op_oob.set_attr_index(&ctx, dialect_mir::attributes::FieldIndexAttr(2)); + assert!(extract_op_oob.verify(&ctx).is_err(), "OOB Index"); + + let union_ty = MirUnionType::get( + &mut ctx, + "Bits".into(), + vec!["word".into(), "alias".into()], + vec![i32_ty.into(), i32_ty.into()], + 4, + 4, + ); + let union_block = BasicBlock::new(&mut ctx, None, vec![union_ty.into()]); + let union_val = union_block.deref(&ctx).get_argument(0); + let union_extract = Operation::new( + &mut ctx, + MirExtractFieldOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![union_val], + vec![], + 0, + ); + let union_extract = MirExtractFieldOp::new(union_extract); + union_extract.set_attr_index(&ctx, dialect_mir::attributes::FieldIndexAttr(1)); + assert!(union_extract.verify(&ctx).is_ok(), "Valid union extract"); +} + +#[test] +fn test_mir_construct_disjoint_slice_verify() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let f32_ty = FP32Type::get(&ctx); + let usize_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let width_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let f32_ptr_ty = + MirPtrType::get_generic_with_kind(&mut ctx, f32_ty.into(), true, MirPointerKind::RawMut); + let plain_ty = MirDisjointSliceType::get(&mut ctx, f32_ty.into()); + let width_ty_handle: pliron::r#type::TypeHandle = width_ty.into(); + let row_width_ty = + MirDisjointSliceType::get_with_space(&mut ctx, f32_ty.into(), vec![width_ty_handle]); + + let block = BasicBlock::new( + &mut ctx, + None, + vec![f32_ptr_ty.into(), usize_ty.into(), width_ty.into()], + ); + let ptr_val = block.deref(&ctx).get_argument(0); + let len_val = block.deref(&ctx).get_argument(1); + let width_val = block.deref(&ctx).get_argument(2); + + // Valid: an index space with no runtime layout takes two operands. + let op = Operation::new( + &mut ctx, + MirConstructDisjointSliceOp::get_concrete_op_info(), + vec![plain_ty.into()], + vec![ptr_val, len_val], + vec![], + 0, + ); + assert!( + MirConstructDisjointSliceOp::new(op).verify(&ctx).is_ok(), + "Valid space-free disjoint slice construction" + ); + + // Valid: a runtime row width takes a third operand. + let op_width = Operation::new( + &mut ctx, + MirConstructDisjointSliceOp::get_concrete_op_info(), + vec![row_width_ty.into()], + vec![ptr_val, len_val, width_val], + vec![], + 0, + ); + assert!( + MirConstructDisjointSliceOp::new(op_width) + .verify(&ctx) + .is_ok(), + "Valid row-width disjoint slice construction" + ); + + let erased_ptr_ty = MirPtrType::get_generic(&mut ctx, f32_ty.into(), true); + let erased_block = BasicBlock::new(&mut ctx, None, vec![erased_ptr_ty.into()]); + let erased_ptr = erased_block.deref(&ctx).get_argument(0); + let op_erased_data = Operation::new( + &mut ctx, + MirConstructDisjointSliceOp::get_concrete_op_info(), + vec![plain_ty.into()], + vec![erased_ptr, len_val], + vec![], + 0, + ); + assert!( + MirConstructDisjointSliceOp::new(op_erased_data) + .verify(&ctx) + .is_err(), + "DisjointSlice's fixed RawMut field cannot be reconstructed from Erased" + ); + + // Invalid: the row width is missing, so the slice would carry whatever + // slot 2 held. + let op_missing_width = Operation::new( + &mut ctx, + MirConstructDisjointSliceOp::get_concrete_op_info(), + vec![row_width_ty.into()], + vec![ptr_val, len_val], + vec![], + 0, + ); + assert!( + MirConstructDisjointSliceOp::new(op_missing_width) + .verify(&ctx) + .is_err(), + "Missing index-space operand" + ); + + // Invalid: a space-free slice given a third operand. + let op_extra = Operation::new( + &mut ctx, + MirConstructDisjointSliceOp::get_concrete_op_info(), + vec![plain_ty.into()], + vec![ptr_val, len_val, width_val], + vec![], + 0, + ); + assert!( + MirConstructDisjointSliceOp::new(op_extra) + .verify(&ctx) + .is_err(), + "Index-space operand for a space-free slice" + ); + + // Invalid: the row width operand has the wrong width, which would write a + // 64-bit value into the 32-bit row width slot. + let op_wrong_width_ty = Operation::new( + &mut ctx, + MirConstructDisjointSliceOp::get_concrete_op_info(), + vec![row_width_ty.into()], + vec![ptr_val, len_val, len_val], + vec![], + 0, + ); + assert!( + MirConstructDisjointSliceOp::new(op_wrong_width_ty) + .verify(&ctx) + .is_err(), + "Index-space operand type mismatch" + ); + + // Invalid: result is a plain slice, which `mir.construct_slice` owns. + let plain_slice_ty = MirSliceType::get(&mut ctx, f32_ty.into()); + let op_bad_res = Operation::new( + &mut ctx, + MirConstructDisjointSliceOp::get_concrete_op_info(), + vec![plain_slice_ty.into()], + vec![ptr_val, len_val], + vec![], + 0, + ); + assert!( + MirConstructDisjointSliceOp::new(op_bad_res) + .verify(&ctx) + .is_err(), + "Result must be a disjoint slice type" + ); +} + +#[test] +fn test_mir_construct_slice_verify() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let u8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let usize_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let u8_ptr_ty = MirPtrType::get_generic(&mut ctx, u8_ty.into(), false); + let u8_slice_ty = MirSliceType::get(&mut ctx, u8_ty.into()); + let i32_slice_ty = MirSliceType::get(&mut ctx, i32_ty.into()); + + let block = BasicBlock::new(&mut ctx, None, vec![u8_ptr_ty.into(), usize_ty.into()]); + let ptr_val = block.deref(&ctx).get_argument(0); + let len_val = block.deref(&ctx).get_argument(1); + + // Valid: (ptr to u8, usize len) -> slice of u8 + let op = Operation::new( + &mut ctx, + MirConstructSliceOp::get_concrete_op_info(), + vec![u8_slice_ty.into()], + vec![ptr_val, len_val], + vec![], + 0, + ); + assert!( + MirConstructSliceOp::new(op).verify(&ctx).is_ok(), + "Valid slice construction" + ); + + // Invalid: data pointer pointee does not match slice element type + let op_bad_elem = Operation::new( + &mut ctx, + MirConstructSliceOp::get_concrete_op_info(), + vec![i32_slice_ty.into()], + vec![ptr_val, len_val], + vec![], + 0, + ); + assert!( + MirConstructSliceOp::new(op_bad_elem).verify(&ctx).is_err(), + "Pointee/element mismatch" + ); + + // Invalid: operands swapped (length where the pointer should be) + let op_swapped = Operation::new( + &mut ctx, + MirConstructSliceOp::get_concrete_op_info(), + vec![u8_slice_ty.into()], + vec![len_val, ptr_val], + vec![], + 0, + ); + assert!( + MirConstructSliceOp::new(op_swapped).verify(&ctx).is_err(), + "Swapped operands" + ); + + // Invalid: result is not a slice type + let op_bad_res = Operation::new( + &mut ctx, + MirConstructSliceOp::get_concrete_op_info(), + vec![u8_ptr_ty.into()], + vec![ptr_val, len_val], + vec![], + 0, + ); + assert!( + MirConstructSliceOp::new(op_bad_res).verify(&ctx).is_err(), + "Non-slice result type" + ); +} + +#[test] +fn test_slice_carrier_cannot_launder_pointer_kind() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let u8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); + let usize_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let raw_mut_ptr = + MirPtrType::get_generic_with_kind(&mut ctx, u8_ty.into(), true, MirPointerKind::RawMut); + let erased_ptr = MirPtrType::get_generic(&mut ctx, u8_ty.into(), true); + let erased_const_ptr = MirPtrType::get_generic(&mut ctx, u8_ty.into(), false); + let global_raw_mut_ptr = MirPtrType::get_with_kind( + &mut ctx, + u8_ty.into(), + true, + dialect_mir::types::address_space::GLOBAL, + MirPointerKind::RawMut, + ); + let unique_ptr = + MirPtrType::get_generic_with_kind(&mut ctx, u8_ty.into(), true, MirPointerKind::UniqueRef); + let raw_mut_slice = MirSliceType::get_with_kind(&mut ctx, u8_ty.into(), MirPointerKind::RawMut); + let unique_slice = + MirSliceType::get_with_kind(&mut ctx, u8_ty.into(), MirPointerKind::UniqueRef); + let erased_slice = MirSliceType::get(&mut ctx, u8_ty.into()); + let erased_mut_slice = MirSliceType::get_with_mutability(&mut ctx, u8_ty.into(), true); + let block = BasicBlock::new( + &mut ctx, + None, + vec![ + raw_mut_ptr.into(), + erased_ptr.into(), + erased_const_ptr.into(), + global_raw_mut_ptr.into(), + usize_ty.into(), + raw_mut_slice.into(), + erased_slice.into(), + erased_mut_slice.into(), + ], + ); + let raw_mut = block.deref(&ctx).get_argument(0); + let erased = block.deref(&ctx).get_argument(1); + let erased_const = block.deref(&ctx).get_argument(2); + let global_raw_mut = block.deref(&ctx).get_argument(3); + let len = block.deref(&ctx).get_argument(4); + let raw_slice = block.deref(&ctx).get_argument(5); + let erased_slice_value = block.deref(&ctx).get_argument(6); + let erased_mut_slice_value = block.deref(&ctx).get_argument(7); + + let construct = |ctx: &mut Context, data, result_ty| { + Operation::new( + ctx, + MirConstructSliceOp::get_concrete_op_info(), + vec![result_ty], + vec![data, len], + vec![], + 0, + ) + }; + assert!( + MirConstructSliceOp::new(construct(&mut ctx, raw_mut, raw_mut_slice.into())) + .verify(&ctx) + .is_ok() + ); + assert!( + MirConstructSliceOp::new(construct(&mut ctx, raw_mut, unique_slice.into())) + .verify(&ctx) + .is_err(), + "construct_slice must not turn RawMut into UniqueRef" + ); + assert!( + MirConstructSliceOp::new(construct(&mut ctx, erased, raw_mut_slice.into())) + .verify(&ctx) + .is_err(), + "construct_slice must not recover RawMut from Erased" + ); + assert!( + MirConstructSliceOp::new(construct(&mut ctx, global_raw_mut, raw_mut_slice.into())) + .verify(&ctx) + .is_err(), + "ordinary slice carriers always use generic address space" + ); + assert!( + MirConstructSliceOp::new(construct(&mut ctx, erased_const, erased_slice.into())) + .verify(&ctx) + .is_ok(), + "an immutable Erased data pointer constructs an immutable Erased slice" + ); + + let extract = |ctx: &mut Context, result_ty| { + let op = Operation::new( + ctx, + MirExtractFieldOp::get_concrete_op_info(), + vec![result_ty], + vec![raw_slice], + vec![], + 0, + ); + MirExtractFieldOp::new(op).set_attr_index(ctx, FieldIndexAttr(0)); + op + }; + assert!( + MirExtractFieldOp::new(extract(&mut ctx, raw_mut_ptr.into())) + .verify(&ctx) + .is_ok() + ); + assert!( + MirExtractFieldOp::new(extract(&mut ctx, erased_ptr.into())) + .verify(&ctx) + .is_ok(), + "slice extraction may deliberately erase a concrete kind" + ); + assert!( + MirExtractFieldOp::new(extract(&mut ctx, unique_ptr.into())) + .verify(&ctx) + .is_err(), + "slice extraction must not change RawMut into UniqueRef" + ); + assert!( + MirExtractFieldOp::new(extract(&mut ctx, global_raw_mut_ptr.into())) + .verify(&ctx) + .is_err(), + "ordinary slice extraction cannot invent a non-generic address space" + ); + assert!( + MirExtractFieldOp::new(extract(&mut ctx, erased_const_ptr.into())) + .verify(&ctx) + .is_err(), + "erasing a concrete slice kind must preserve machine mutability" + ); + + let extract_erased = |ctx: &mut Context, result_ty| { + let op = Operation::new( + ctx, + MirExtractFieldOp::get_concrete_op_info(), + vec![result_ty], + vec![erased_slice_value], + vec![], + 0, + ); + MirExtractFieldOp::new(op).set_attr_index(ctx, FieldIndexAttr(0)); + op + }; + assert!( + MirExtractFieldOp::new(extract_erased(&mut ctx, erased_const_ptr.into())) + .verify(&ctx) + .is_ok(), + "an immutable Erased slice extracts an immutable Erased data pointer" + ); + assert!( + MirExtractFieldOp::new(extract_erased(&mut ctx, erased_ptr.into())) + .verify(&ctx) + .is_err(), + "an immutable Erased slice cannot manufacture writable data-pointer evidence" + ); + + let extract_erased_mut = |ctx: &mut Context, result_ty| { + let op = Operation::new( + ctx, + MirExtractFieldOp::get_concrete_op_info(), + vec![result_ty], + vec![erased_mut_slice_value], + vec![], + 0, + ); + MirExtractFieldOp::new(op).set_attr_index(ctx, FieldIndexAttr(0)); + op + }; + assert!( + MirExtractFieldOp::new(extract_erased_mut(&mut ctx, erased_ptr.into())) + .verify(&ctx) + .is_ok(), + "a mutable Erased slice retains writable data-pointer evidence" + ); + + let immutable_slice_reborrow = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![unique_slice.into()], + vec![erased_slice_value], + vec![], + 0, + ); + let immutable_slice_reborrow = MirCastOp::new(immutable_slice_reborrow); + immutable_slice_reborrow.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + immutable_slice_reborrow + .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); + assert!( + immutable_slice_reborrow.verify(&ctx).is_err(), + "an immutable Erased slice cannot be reborrowed as UniqueRef" + ); + + let mutable_slice_reborrow = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![unique_slice.into()], + vec![erased_mut_slice_value], + vec![], + 0, + ); + let mutable_slice_reborrow = MirCastOp::new(mutable_slice_reborrow); + mutable_slice_reborrow.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + mutable_slice_reborrow + .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); + assert!( + mutable_slice_reborrow.verify(&ctx).is_ok(), + "a mutable Erased slice may establish UniqueRef at a real reborrow boundary" + ); +} + +#[test] +fn test_mir_field_addr_tuple_pointee_verify() { + // `(u8, u32)` laid out the way rustc actually places it: the u32 field + // first in memory for alignment, so declaration index 0 (`u8`) lands at + // byte offset 4 and declaration index 1 (`u32`) lands at byte offset 0. + // `field_addr`'s `field_index` attribute is a DECLARATION index (it names + // `.0`/`.1` as written), so this test only passes if the op resolves the + // field's type through `MirTupleType::get_types()` (declaration order) + // rather than assuming identity with memory order. + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let u8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + + let tuple_ty = MirTupleType::get_with_layout( + &mut ctx, + vec![u8_ty.into(), u32_ty.into()], + vec![1, 0], + vec![4, 0], + 8, + 4, + ); + + let tuple_ptr_ty = MirPtrType::get_generic(&mut ctx, tuple_ty.into(), false); + let blk = BasicBlock::new(&mut ctx, None, vec![tuple_ptr_ty.into()]); + let tuple_ptr = blk.deref(&ctx).get_argument(0); + + let u8_ptr_ty = MirPtrType::get_generic(&mut ctx, u8_ty.into(), false); + let op_field0 = Operation::new( + &mut ctx, + MirFieldAddrOp::get_concrete_op_info(), + vec![u8_ptr_ty.into()], + vec![tuple_ptr], + vec![], + 0, + ); + let field0 = MirFieldAddrOp::new(op_field0); + field0.set_attr_field_index(&ctx, FieldIndexAttr(0)); + field0.set_attr_aggregate_ty(&ctx, TypeAttr::new(tuple_ty.into())); + assert!( + field0.verify(&ctx).is_ok(), + "tuple field 0 (u8) address accepted" + ); + + let u32_ptr_ty = MirPtrType::get_generic(&mut ctx, u32_ty.into(), false); + let op_field1 = Operation::new( + &mut ctx, + MirFieldAddrOp::get_concrete_op_info(), + vec![u32_ptr_ty.into()], + vec![tuple_ptr], + vec![], + 0, + ); + let field1 = MirFieldAddrOp::new(op_field1); + field1.set_attr_field_index(&ctx, FieldIndexAttr(1)); + field1.set_attr_aggregate_ty(&ctx, TypeAttr::new(tuple_ty.into())); + assert!( + field1.verify(&ctx).is_ok(), + "tuple field 1 (u32) address accepted" + ); + + // Result pointee type must match the DECLARED field type, not whatever + // sits at that byte offset: pointing field 0's result at u32 (field 1's + // type) must be rejected even though both are in-bounds indices. + let op_wrong_result_ty = Operation::new( + &mut ctx, + MirFieldAddrOp::get_concrete_op_info(), + vec![u32_ptr_ty.into()], + vec![tuple_ptr], + vec![], + 0, + ); + let wrong_result_ty = MirFieldAddrOp::new(op_wrong_result_ty); + wrong_result_ty.set_attr_field_index(&ctx, FieldIndexAttr(0)); + wrong_result_ty.set_attr_aggregate_ty(&ctx, TypeAttr::new(tuple_ty.into())); + assert!( + wrong_result_ty.verify(&ctx).is_err(), + "result pointee type mismatch rejected" + ); + + let op_out_of_bounds = Operation::new( + &mut ctx, + MirFieldAddrOp::get_concrete_op_info(), + vec![u8_ptr_ty.into()], + vec![tuple_ptr], + vec![], + 0, + ); + let out_of_bounds = MirFieldAddrOp::new(op_out_of_bounds); + out_of_bounds.set_attr_field_index(&ctx, FieldIndexAttr(2)); + out_of_bounds.set_attr_aggregate_ty(&ctx, TypeAttr::new(tuple_ty.into())); + assert!( + out_of_bounds.verify(&ctx).is_err(), + "out-of-bounds tuple field index rejected" + ); +} + +#[test] +fn test_mir_field_addr_tuple_pointee_store_verify() { + // The WRITE side of the tuple-pointee unlock: `t.1 = x` / `arr[i].1 = x` + // lower to `mir.field_addr` + `mir.store` through the field's address, so + // a tuple-pointee field address used as a store destination must pass + // verification too. Same reordered `(u8, u32)` layout as above (the u32 + // field first in memory), so the store type-checks against the DECLARED + // field type, not whatever occupies that memory slot. + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let u8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + + let tuple_ty = MirTupleType::get_with_layout( + &mut ctx, + vec![u8_ty.into(), u32_ty.into()], + vec![1, 0], + vec![4, 0], + 8, + 4, + ); + + let tuple_ptr_ty = MirPtrType::get_generic(&mut ctx, tuple_ty.into(), false); + let blk = BasicBlock::new( + &mut ctx, + None, + vec![tuple_ptr_ty.into(), u32_ty.into(), u8_ty.into()], + ); + let tuple_ptr = blk.deref(&ctx).get_argument(0); + let u32_val = blk.deref(&ctx).get_argument(1); + let u8_val = blk.deref(&ctx).get_argument(2); + + // `.1 = x`: address declaration field 1 (u32, memory slot 0) and store a + // u32 through it. + let u32_ptr_ty = MirPtrType::get_generic(&mut ctx, u32_ty.into(), false); + let op_field1 = Operation::new( + &mut ctx, + MirFieldAddrOp::get_concrete_op_info(), + vec![u32_ptr_ty.into()], + vec![tuple_ptr], + vec![], + 0, + ); + let field1 = MirFieldAddrOp::new(op_field1); + field1.set_attr_field_index(&ctx, FieldIndexAttr(1)); + field1.set_attr_aggregate_ty(&ctx, TypeAttr::new(tuple_ty.into())); + assert!( + field1.verify(&ctx).is_ok(), + "tuple field 1 (u32) address accepted as a store destination" + ); + let field1_ptr = op_field1.deref(&ctx).get_result(0); + + let op_store = Operation::new( + &mut ctx, + MirStoreOp::get_concrete_op_info(), + vec![], + vec![field1_ptr, u32_val], + vec![], + 0, + ); + assert!( + MirStoreOp::new(op_store).verify(&ctx).is_ok(), + "store through a tuple field address verifies" + ); + + // The stored value must match the DECLARED field type (`u32` for `.1`), + // not the type of the field sharing the tuple: a u8 store through the + // `.1` pointer is a type mismatch. + let op_store_wrong_ty = Operation::new( + &mut ctx, + MirStoreOp::get_concrete_op_info(), + vec![], + vec![field1_ptr, u8_val], + vec![], + 0, + ); + assert!( + MirStoreOp::new(op_store_wrong_ty).verify(&ctx).is_err(), + "store of a mismatched value type through a tuple field address rejected" + ); +} diff --git a/crates/dialect-mir/tests/ops_test/cast_authority.rs b/crates/dialect-mir/tests/ops_test/cast_authority.rs new file mode 100644 index 0000000000..b3dc7a8b47 --- /dev/null +++ b/crates/dialect-mir/tests/ops_test/cast_authority.rs @@ -0,0 +1,1060 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use dialect_mir::{ + attributes::{FieldIndexAttr, MirCastKindAttr, MirPointerKindAuthorityAttr}, + ops::{MirArrayElementAddrOp, MirCastOp, MirExtractFieldOp, MirFieldAddrOp, MirPtrOffsetOp}, + types::{MirArrayType, MirPointerKind, MirPtrType, MirSliceType, MirStructType, MirTupleType}, +}; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::TypeAttr, + types::{IntegerType, Signedness}, + }, + common_traits::Verify, + context::Context, + op::Op, + operation::Operation, +}; + +#[test] +fn test_aggregate_cast_cannot_hide_pointer_kind_laundering() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let raw_mut_ty = + MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::RawMut); + let unique_ty = + MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::UniqueRef); + let raw_tuple_ty = MirTupleType::get(&mut ctx, vec![raw_mut_ty.into()]); + let unique_tuple_ty = MirTupleType::get(&mut ctx, vec![unique_ty.into()]); + let block = BasicBlock::new(&mut ctx, None, vec![raw_tuple_ty.into()]); + let raw_tuple = block.deref(&ctx).get_argument(0); + + let build = |ctx: &mut Context, authority: Option| { + let op = Operation::new( + ctx, + MirCastOp::get_concrete_op_info(), + vec![unique_tuple_ty.into()], + vec![raw_tuple], + vec![], + 0, + ); + let cast = MirCastOp::new(op); + cast.set_attr_cast_kind(ctx, MirCastKindAttr::Transmute); + if let Some(authority) = authority { + cast.set_pointer_kind_authority(ctx, authority); + } + op + }; + + assert!( + MirCastOp::new(build(&mut ctx, None)).verify(&ctx).is_err(), + "nested RawMut -> UniqueRef laundering must be rejected" + ); + assert!( + MirCastOp::new(build(&mut ctx, Some(MirPointerKindAuthorityAttr::Reborrow),)) + .verify(&ctx) + .is_err(), + "Rvalue::Ref cannot authorize a nested aggregate transition" + ); + assert!( + MirCastOp::new(build(&mut ctx, Some(MirPointerKindAuthorityAttr::RustCast),)) + .verify(&ctx) + .is_ok(), + "an explicit rustc aggregate transmute is visible and authorized" + ); +} + +#[test] +fn test_aggregate_reinterpretation_requires_rust_cast_authority() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let shared_ty = MirPtrType::get_generic_with_kind( + &mut ctx, + i64_ty.into(), + false, + MirPointerKind::SharedRef, + ); + // Both tuple declarations list the same pointer at field 0, but its byte + // offset changes from 0 to 8. Pairing by declaration index alone would + // therefore mistake an integer's old bytes for a preserved reference. + let source_ty = MirTupleType::get_with_layout( + &mut ctx, + vec![shared_ty.into(), i64_ty.into()], + vec![0, 1], + vec![0, 8], + 16, + 8, + ); + let target_ty = MirTupleType::get_with_layout( + &mut ctx, + vec![shared_ty.into(), i64_ty.into()], + vec![1, 0], + vec![8, 0], + 16, + 8, + ); + let block = BasicBlock::new(&mut ctx, None, vec![source_ty.into()]); + let source = block.deref(&ctx).get_argument(0); + + let build = |ctx: &mut Context, authority: Option| { + let op = Operation::new( + ctx, + MirCastOp::get_concrete_op_info(), + vec![target_ty.into()], + vec![source], + vec![], + 0, + ); + let cast = MirCastOp::new(op); + cast.set_attr_cast_kind(ctx, MirCastKindAttr::Transmute); + if let Some(authority) = authority { + cast.set_pointer_kind_authority(ctx, authority); + } + op + }; + + assert!( + MirCastOp::new(build(&mut ctx, None)).verify(&ctx).is_err(), + "layout-changing aggregate casts must not infer pointer preservation by field index" + ); + assert!( + MirCastOp::new(build(&mut ctx, Some(MirPointerKindAuthorityAttr::RustCast),)) + .verify(&ctx) + .is_ok(), + "an explicit rustc transmute makes the representation reinterpretation auditable" + ); +} + +#[test] +fn test_generic_aggregate_cast_cannot_invent_erased_pointer_carriers() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let word_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let erased_mut = MirPtrType::get_generic(&mut ctx, word_ty.into(), true); + let source_ty = MirTupleType::get_with_layout( + &mut ctx, + vec![erased_mut.into(), word_ty.into()], + vec![0, 1], + vec![0, 8], + 16, + 8, + ); + let extra_pointer_ty = MirTupleType::get_with_layout( + &mut ctx, + vec![erased_mut.into(), erased_mut.into()], + vec![0, 1], + vec![0, 8], + 16, + 8, + ); + let moved_pointer_ty = MirTupleType::get_with_layout( + &mut ctx, + vec![erased_mut.into(), word_ty.into()], + vec![1, 0], + vec![8, 0], + 16, + 8, + ); + let one_pointer = MirArrayType::get(&mut ctx, erased_mut.into(), 1); + let two_pointers = MirArrayType::get(&mut ctx, erased_mut.into(), 2); + + let source_block = BasicBlock::new(&mut ctx, None, vec![source_ty.into()]); + let source = source_block.deref(&ctx).get_argument(0); + let array_block = BasicBlock::new(&mut ctx, None, vec![one_pointer.into()]); + let array_source = array_block.deref(&ctx).get_argument(0); + let build = |ctx: &mut Context, source, target| { + let op = Operation::new( + ctx, + MirCastOp::get_concrete_op_info(), + vec![target], + vec![source], + vec![], + 0, + ); + MirCastOp::new(op).set_attr_cast_kind(ctx, MirCastKindAttr::PtrToPtr); + op + }; + + assert!( + MirCastOp::new(build(&mut ctx, source, extra_pointer_ty.into())) + .verify(&ctx) + .is_err(), + "PtrToPtr cannot reinterpret an integer field as writable Erased pointer evidence" + ); + assert!( + MirCastOp::new(build(&mut ctx, source, moved_pointer_ty.into())) + .verify(&ctx) + .is_err(), + "field-index equality cannot hide an Erased pointer moving onto integer bytes" + ); + assert!( + MirCastOp::new(build(&mut ctx, array_source, two_pointers.into())) + .verify(&ctx) + .is_err(), + "homogeneous array traversal must retain pointer-carrier cardinality" + ); +} + +#[test] +fn test_rust_cast_authority_obeys_cast_kind_semantics() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let u8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); + let usize_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let array_ty = MirArrayType::get(&mut ctx, u8_ty.into(), 4); + let raw_mut = + MirPtrType::get_generic_with_kind(&mut ctx, u8_ty.into(), true, MirPointerKind::RawMut); + let raw_const = + MirPtrType::get_generic_with_kind(&mut ctx, u8_ty.into(), false, MirPointerKind::RawConst); + let unique = + MirPtrType::get_generic_with_kind(&mut ctx, u8_ty.into(), true, MirPointerKind::UniqueRef); + let fn_target = MirStructType::get_with_full_layout( + &mut ctx, + "FnPtrTarget".into(), + vec![], + vec![], + vec![], + vec![], + 0, + 0, + ); + let fn_target_ty: pliron::r#type::TypeHandle = fn_target.into(); + let fn_carrier = MirPtrType::get_generic(&mut ctx, fn_target_ty, false); + let raw_mut_array = + MirPtrType::get_generic_with_kind(&mut ctx, array_ty.into(), true, MirPointerKind::RawMut); + let raw_const_array = MirPtrType::get_generic_with_kind( + &mut ctx, + array_ty.into(), + false, + MirPointerKind::RawConst, + ); + let erased_array = MirPtrType::get_generic(&mut ctx, array_ty.into(), false); + let erased_slice = MirSliceType::get(&mut ctx, u8_ty.into()); + let erased_mut_slice = MirSliceType::get_with_mutability(&mut ctx, u8_ty.into(), true); + let raw_const_slice = + MirSliceType::get_with_kind(&mut ctx, u8_ty.into(), MirPointerKind::RawConst); + let block = BasicBlock::new( + &mut ctx, + None, + vec![ + raw_mut.into(), + raw_const.into(), + fn_carrier.into(), + raw_mut_array.into(), + raw_const_array.into(), + erased_array.into(), + usize_ty.into(), + ], + ); + let raw_mut_value = block.deref(&ctx).get_argument(0); + let raw_const_value = block.deref(&ctx).get_argument(1); + let fn_value = block.deref(&ctx).get_argument(2); + let raw_mut_array_value = block.deref(&ctx).get_argument(3); + let raw_const_array_value = block.deref(&ctx).get_argument(4); + let erased_array_value = block.deref(&ctx).get_argument(5); + let integer_value = block.deref(&ctx).get_argument(6); + + let build = |ctx: &mut Context, + source, + target, + kind, + authority: Option| { + let op = Operation::new( + ctx, + MirCastOp::get_concrete_op_info(), + vec![target], + vec![source], + vec![], + 0, + ); + let cast = MirCastOp::new(op); + cast.set_attr_cast_kind(ctx, kind); + if let Some(authority) = authority { + cast.set_pointer_kind_authority(ctx, authority); + } + op + }; + let rust_cast = Some(MirPointerKindAuthorityAttr::RustCast); + + assert!( + MirCastOp::new(build( + &mut ctx, + raw_mut_value, + unique.into(), + MirCastKindAttr::PtrToPtr, + rust_cast.clone(), + )) + .verify(&ctx) + .is_err(), + "PtrToPtr cannot use RustCast to invent a reference" + ); + assert!( + MirCastOp::new(build( + &mut ctx, + raw_mut_value, + raw_const.into(), + MirCastKindAttr::PtrToPtr, + rust_cast.clone(), + )) + .verify(&ctx) + .is_ok(), + "PtrToPtr may perform an explicit raw-to-raw cast" + ); + assert!( + MirCastOp::new(build( + &mut ctx, + fn_value, + raw_const.into(), + MirCastKindAttr::FnPtrToPtr, + rust_cast.clone(), + )) + .verify(&ctx) + .is_ok(), + "FnPtrToPtr may expose an opaque function pointer as a raw pointer" + ); + assert!( + MirCastOp::new(build( + &mut ctx, + erased_array_value, + raw_const.into(), + MirCastKindAttr::FnPtrToPtr, + rust_cast.clone(), + )) + .verify(&ctx) + .is_err(), + "FnPtrToPtr cannot relabel arbitrary Erased storage as a function pointer" + ); + assert!( + MirCastOp::new(build( + &mut ctx, + fn_value, + unique.into(), + MirCastKindAttr::FnPtrToPtr, + rust_cast.clone(), + )) + .verify(&ctx) + .is_err() + ); + assert!( + MirCastOp::new(build( + &mut ctx, + raw_mut_value, + raw_const.into(), + MirCastKindAttr::PointerCoercionMutToConst, + rust_cast.clone(), + )) + .verify(&ctx) + .is_ok() + ); + assert!( + MirCastOp::new(build( + &mut ctx, + raw_const_value, + raw_mut.into(), + MirCastKindAttr::PointerCoercionMutToConst, + rust_cast.clone(), + )) + .verify(&ctx) + .is_err() + ); + assert!( + MirCastOp::new(build( + &mut ctx, + raw_mut_value, + raw_mut.into(), + MirCastKindAttr::PointerCoercionMutToConst, + None, + )) + .verify(&ctx) + .is_err(), + "MutToConst cannot masquerade as a same-category pointer cast" + ); + assert!( + MirCastOp::new(build( + &mut ctx, + raw_const_array_value, + raw_const_slice.into(), + MirCastKindAttr::PointerCoercionUnsize, + rust_cast.clone(), + )) + .verify(&ctx) + .is_ok(), + "Unsize may change thin/fat shape while preserving its carrier" + ); + assert!( + MirCastOp::new(build( + &mut ctx, + erased_array_value, + erased_slice.into(), + MirCastKindAttr::PointerCoercionUnsize, + None, + )) + .verify(&ctx) + .is_ok(), + "an all-Erased unsize still preserves read-only carrier state" + ); + assert!( + MirCastOp::new(build( + &mut ctx, + erased_array_value, + erased_mut_slice.into(), + MirCastKindAttr::PointerCoercionUnsize, + None, + )) + .verify(&ctx) + .is_err(), + "Unsize cannot turn an immutable Erased thin pointer into writable fat evidence" + ); + + let prefix_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let sized_tail = MirStructType::get_with_full_layout( + &mut ctx, + "TailCarrier".into(), + vec!["prefix".into(), "tail".into()], + vec![prefix_ty.into(), array_ty.into()], + vec![0, 1], + vec![0, 8], + 16, + 8, + ); + let shifted_unsized_tail = MirStructType::get_with_full_layout( + &mut ctx, + "TailCarrier".into(), + vec!["prefix".into(), "tail".into()], + vec![prefix_ty.into(), u8_ty.into()], + vec![0, 1], + vec![0, 16], + 24, + 8, + ); + let sized_tail_ptr = MirPtrType::get_generic_with_kind( + &mut ctx, + sized_tail.into(), + false, + MirPointerKind::RawConst, + ); + let shifted_tail_slice = MirSliceType::get_with_kind( + &mut ctx, + shifted_unsized_tail.into(), + MirPointerKind::RawConst, + ); + let tail_block = BasicBlock::new(&mut ctx, None, vec![sized_tail_ptr.into()]); + let tail_value = tail_block.deref(&ctx).get_argument(0); + assert!( + MirCastOp::new(build( + &mut ctx, + tail_value, + shifted_tail_slice.into(), + MirCastKindAttr::PointerCoercionUnsize, + rust_cast.clone(), + )) + .verify(&ctx) + .is_err(), + "Unsize cannot move the trailing data behind a changed struct-field offset" + ); + + assert!( + MirCastOp::new(build( + &mut ctx, + raw_const_array_value, + raw_mut.into(), + MirCastKindAttr::PointerCoercionArrayToPointer, + rust_cast.clone(), + )) + .verify(&ctx) + .is_err(), + "ArrayToPointer cannot strengthen const raw storage to mutable" + ); + assert!( + MirCastOp::new(build( + &mut ctx, + raw_mut_value, + raw_const.into(), + MirCastKindAttr::PointerCoercionArrayToPointer, + rust_cast.clone(), + )) + .verify(&ctx) + .is_err(), + "ArrayToPointer requires a raw pointer to an array, not an arbitrary pointer" + ); + assert!( + MirCastOp::new(build( + &mut ctx, + raw_mut_value, + raw_mut.into(), + MirCastKindAttr::PointerCoercionArrayToPointer, + None, + )) + .verify(&ctx) + .is_err(), + "an unmarked same-kind ArrayToPointer still requires an array source" + ); + assert!( + MirCastOp::new(build( + &mut ctx, + raw_mut_array_value, + raw_const.into(), + MirCastKindAttr::PointerCoercionArrayToPointer, + rust_cast.clone(), + )) + .verify(&ctx) + .is_ok() + ); + assert!( + MirCastOp::new(build( + &mut ctx, + integer_value, + unique.into(), + MirCastKindAttr::Transmute, + rust_cast, + )) + .verify(&ctx) + .is_ok(), + "only an explicit Transmute may establish an arbitrary pointer category" + ); + + let word_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let source_wrapper = MirStructType::get_with_full_layout( + &mut ctx, + "CarrierWrapper".into(), + vec!["pointer".into(), "word".into()], + vec![unique.into(), word_ty.into()], + vec![0, 1], + vec![0, 8], + 16, + 8, + ); + let moved_wrapper = MirStructType::get_with_full_layout( + &mut ctx, + "CarrierWrapper".into(), + vec!["word".into(), "pointer".into()], + vec![word_ty.into(), unique.into()], + vec![0, 1], + vec![0, 8], + 16, + 8, + ); + let wrapper_block = BasicBlock::new(&mut ctx, None, vec![source_wrapper.into()]); + let wrapper_value = wrapper_block.deref(&ctx).get_argument(0); + for kind in [ + MirCastKindAttr::PointerCoercionUnsize, + MirCastKindAttr::Subtype, + ] { + assert!( + MirCastOp::new(build( + &mut ctx, + wrapper_value, + moved_wrapper.into(), + kind, + Some(MirPointerKindAuthorityAttr::RustCast), + )) + .verify(&ctx) + .is_err(), + "Unsize/Subtype cannot bless a pointer carrier moved onto integer bytes" + ); + } +} + +#[test] +fn test_pointer_projections_preserve_or_erase_kind_only() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let usize_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let tuple_ty = MirTupleType::get(&mut ctx, vec![i32_ty.into()]); + let array_ty = MirArrayType::get(&mut ctx, i32_ty.into(), 4); + let raw_tuple_ty = + MirPtrType::get_generic_with_kind(&mut ctx, tuple_ty.into(), true, MirPointerKind::RawMut); + let raw_array_ty = + MirPtrType::get_generic_with_kind(&mut ctx, array_ty.into(), true, MirPointerKind::RawMut); + let raw_field_ty = + MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::RawMut); + let erased_field_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), true); + let erased_read_field_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), false); + let unique_field_ty = + MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::UniqueRef); + let block = BasicBlock::new( + &mut ctx, + None, + vec![raw_tuple_ty.into(), raw_array_ty.into(), usize_ty.into()], + ); + let tuple = block.deref(&ctx).get_argument(0); + let array = block.deref(&ctx).get_argument(1); + let index = block.deref(&ctx).get_argument(2); + + let field = |ctx: &mut Context, result_ty| { + let op = Operation::new( + ctx, + MirFieldAddrOp::get_concrete_op_info(), + vec![result_ty], + vec![tuple], + vec![], + 0, + ); + MirFieldAddrOp::new(op).set_attr_field_index(ctx, FieldIndexAttr(0)); + MirFieldAddrOp::new(op).set_attr_aggregate_ty(ctx, TypeAttr::new(tuple_ty.into())); + op + }; + assert!( + MirFieldAddrOp::new(field(&mut ctx, raw_field_ty.into())) + .verify(&ctx) + .is_ok() + ); + assert!( + MirFieldAddrOp::new(field(&mut ctx, erased_field_ty.into())) + .verify(&ctx) + .is_ok() + ); + assert!( + MirFieldAddrOp::new(field(&mut ctx, unique_field_ty.into())) + .verify(&ctx) + .is_err() + ); + assert!( + MirFieldAddrOp::new(field(&mut ctx, erased_read_field_ty.into())) + .verify(&ctx) + .is_err(), + "field projection cannot flip an Erased address from writable to read-only" + ); + + let array_launder = Operation::new( + &mut ctx, + MirArrayElementAddrOp::get_concrete_op_info(), + vec![unique_field_ty.into()], + vec![array, index], + vec![], + 0, + ); + assert!( + MirArrayElementAddrOp::new(array_launder) + .verify(&ctx) + .is_err() + ); + let array_erase = Operation::new( + &mut ctx, + MirArrayElementAddrOp::get_concrete_op_info(), + vec![erased_field_ty.into()], + vec![array, index], + vec![], + 0, + ); + assert!( + MirArrayElementAddrOp::new(array_erase).verify(&ctx).is_ok(), + "projection may erase kind while preserving machine mutability" + ); + let array_mutability_flip = Operation::new( + &mut ctx, + MirArrayElementAddrOp::get_concrete_op_info(), + vec![erased_read_field_ty.into()], + vec![array, index], + vec![], + 0, + ); + assert!( + MirArrayElementAddrOp::new(array_mutability_flip) + .verify(&ctx) + .is_err(), + "array projection cannot change machine mutability" + ); +} + +#[test] +fn test_pointer_with_exposed_provenance_only_creates_raw_or_erased_pointer() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let usize_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let fn_target = MirStructType::get_with_full_layout( + &mut ctx, + "FnPtrTarget".into(), + vec![], + vec![], + vec![], + vec![], + 0, + 0, + ); + let fn_target_ty: pliron::r#type::TypeHandle = fn_target.into(); + let raw_ty = + MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::RawMut); + let opaque_fn_ty = MirPtrType::get_generic(&mut ctx, fn_target_ty, false); + let arbitrary_erased_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), false); + let writable_erased_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), true); + let unique_ty = + MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::UniqueRef); + let block = BasicBlock::new(&mut ctx, None, vec![usize_ty.into()]); + let address = block.deref(&ctx).get_argument(0); + + let build = |ctx: &mut Context, result_ty| { + let op = Operation::new( + ctx, + MirCastOp::get_concrete_op_info(), + vec![result_ty], + vec![address], + vec![], + 0, + ); + let cast = MirCastOp::new(op); + cast.set_attr_cast_kind(ctx, MirCastKindAttr::PointerWithExposedProvenance); + cast.set_pointer_kind_authority(ctx, MirPointerKindAuthorityAttr::RustCast); + op + }; + + assert!( + MirCastOp::new(build(&mut ctx, raw_ty.into())) + .verify(&ctx) + .is_ok() + ); + assert!( + MirCastOp::new(build(&mut ctx, unique_ty.into())) + .verify(&ctx) + .is_err(), + "integer provenance cannot directly materialize a Rust reference" + ); + + let build_unmarked = |ctx: &mut Context, result_ty| { + let op = Operation::new( + ctx, + MirCastOp::get_concrete_op_info(), + vec![result_ty], + vec![address], + vec![], + 0, + ); + MirCastOp::new(op).set_attr_cast_kind(ctx, MirCastKindAttr::PointerWithExposedProvenance); + op + }; + assert!( + MirCastOp::new(build_unmarked(&mut ctx, opaque_fn_ty.into())) + .verify(&ctx) + .is_ok(), + "function-pointer tokens may materialize only the canonical immutable Erased carrier" + ); + assert!( + MirCastOp::new(build_unmarked(&mut ctx, arbitrary_erased_ty.into())) + .verify(&ctx) + .is_err(), + "an integer cannot manufacture arbitrary Erased pointer evidence" + ); + assert!( + MirCastOp::new(build_unmarked(&mut ctx, writable_erased_ty.into())) + .verify(&ctx) + .is_err(), + "an integer cannot manufacture writable Erased evidence and then reborrow it as UniqueRef" + ); + + let fn_block = BasicBlock::new(&mut ctx, None, vec![opaque_fn_ty.into()]); + let opaque_fn = fn_block.deref(&ctx).get_argument(0); + let disguised_data_pointer = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![arbitrary_erased_ty.into()], + vec![opaque_fn], + vec![], + 0, + ); + MirCastOp::new(disguised_data_pointer) + .set_attr_cast_kind(&ctx, MirCastKindAttr::PointerCoercionArrayToPointer); + assert!( + MirCastOp::new(disguised_data_pointer).verify(&ctx).is_err(), + "an unrelated coercion cannot turn the opaque function token into Erased data storage" + ); + + let shared_fn_ty = + MirPtrType::get_generic_with_kind(&mut ctx, fn_target_ty, false, MirPointerKind::SharedRef); + let fake_reborrow = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![shared_fn_ty.into()], + vec![opaque_fn], + vec![], + 0, + ); + let fake_reborrow = MirCastOp::new(fake_reborrow); + fake_reborrow.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + fake_reborrow.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); + assert!( + fake_reborrow.verify(&ctx).is_err(), + "an opaque function-pointer value is not compiler storage that may be reborrowed" + ); + + // Keeping only kind+mutability while recursively pairing aggregate fields + // is insufficient: the canonical function token is also identified by + // its pointee. Otherwise a tuple cast can disguise data storage as a + // function pointer (or vice versa), after which individually legal casts + // can manufacture a reference. + let data_tuple_ty = MirTupleType::get(&mut ctx, vec![arbitrary_erased_ty.into()]); + let fn_tuple_ty = MirTupleType::get(&mut ctx, vec![opaque_fn_ty.into()]); + let nested_block = BasicBlock::new( + &mut ctx, + None, + vec![data_tuple_ty.into(), fn_tuple_ty.into()], + ); + let data_tuple = nested_block.deref(&ctx).get_argument(0); + let fn_tuple = nested_block.deref(&ctx).get_argument(1); + let nested_cast = |ctx: &mut Context, source, target| { + let op = Operation::new( + ctx, + MirCastOp::get_concrete_op_info(), + vec![target], + vec![source], + vec![], + 0, + ); + MirCastOp::new(op).set_attr_cast_kind(ctx, MirCastKindAttr::PtrToPtr); + op + }; + let data_to_fn = nested_cast(&mut ctx, data_tuple, fn_tuple_ty.into()); + assert!( + MirCastOp::new(data_to_fn).verify(&ctx).is_err(), + "an aggregate cast cannot manufacture a nested canonical function token" + ); + assert!( + MirCastOp::new(nested_cast(&mut ctx, fn_tuple, data_tuple_ty.into(),)) + .verify(&ctx) + .is_err(), + "an aggregate cast cannot disguise a nested function token as Erased data" + ); + + let forged_fn_tuple = data_to_fn.deref(&ctx).get_result(0); + let extract_forged_fn = Operation::new( + &mut ctx, + MirExtractFieldOp::get_concrete_op_info(), + vec![opaque_fn_ty.into()], + vec![forged_fn_tuple], + vec![], + 0, + ); + let extract_forged_fn = MirExtractFieldOp::new(extract_forged_fn); + extract_forged_fn.set_attr_index(&ctx, FieldIndexAttr(0)); + assert!(extract_forged_fn.verify(&ctx).is_ok()); + + let forged_fn = extract_forged_fn.get_operation().deref(&ctx).get_result(0); + let expose_forged_fn = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![raw_ty.into()], + vec![forged_fn], + vec![], + 0, + ); + let expose_forged_fn = MirCastOp::new(expose_forged_fn); + expose_forged_fn.set_attr_cast_kind(&ctx, MirCastKindAttr::FnPtrToPtr); + expose_forged_fn.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::RustCast); + assert!(expose_forged_fn.verify(&ctx).is_ok()); + + let forged_raw = expose_forged_fn.get_operation().deref(&ctx).get_result(0); + let reborrow_forged_fn = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![unique_ty.into()], + vec![forged_raw], + vec![], + 0, + ); + let reborrow_forged_fn = MirCastOp::new(reborrow_forged_fn); + reborrow_forged_fn.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + reborrow_forged_fn.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); + assert!( + reborrow_forged_fn.verify(&ctx).is_ok(), + "the laundering chain must be rejected at its aggregate reinterpretation" + ); + + let shared_marker_ty = + MirPtrType::get_generic_with_kind(&mut ctx, fn_target_ty, false, MirPointerKind::SharedRef); + let marker_tuple_ty = MirTupleType::get(&mut ctx, vec![fn_target_ty]); + let shared_marker_tuple_ty = MirPtrType::get_generic_with_kind( + &mut ctx, + marker_tuple_ty.into(), + false, + MirPointerKind::SharedRef, + ); + let marker_array_ty = MirArrayType::get(&mut ctx, fn_target_ty, 1); + let shared_marker_array_ty = MirPtrType::get_generic_with_kind( + &mut ctx, + marker_array_ty.into(), + false, + MirPointerKind::SharedRef, + ); + let shared_marker_slice_ty = + MirSliceType::get_with_kind(&mut ctx, fn_target_ty, MirPointerKind::SharedRef); + let projection_block = BasicBlock::new( + &mut ctx, + None, + vec![ + shared_marker_ty.into(), + shared_marker_tuple_ty.into(), + shared_marker_array_ty.into(), + shared_marker_slice_ty.into(), + usize_ty.into(), + ], + ); + let marker_address = projection_block.deref(&ctx).get_argument(0); + let marker_tuple_address = projection_block.deref(&ctx).get_argument(1); + let marker_array_address = projection_block.deref(&ctx).get_argument(2); + let marker_slice = projection_block.deref(&ctx).get_argument(3); + let zero = projection_block.deref(&ctx).get_argument(4); + + let offset_to_token = Operation::new( + &mut ctx, + MirPtrOffsetOp::get_concrete_op_info(), + vec![opaque_fn_ty.into()], + vec![marker_address, zero], + vec![], + 0, + ); + assert!( + MirPtrOffsetOp::new(offset_to_token).verify(&ctx).is_err(), + "pointer arithmetic produces a data address, never a function token" + ); + + let field_address_to_token = Operation::new( + &mut ctx, + MirFieldAddrOp::get_concrete_op_info(), + vec![opaque_fn_ty.into()], + vec![marker_tuple_address], + vec![], + 0, + ); + let field_address_to_token = MirFieldAddrOp::new(field_address_to_token); + field_address_to_token.set_attr_field_index(&ctx, FieldIndexAttr(0)); + assert!( + field_address_to_token.verify(&ctx).is_err(), + "a field address cannot masquerade as a function token" + ); + + let array_address_to_token = Operation::new( + &mut ctx, + MirArrayElementAddrOp::get_concrete_op_info(), + vec![opaque_fn_ty.into()], + vec![marker_array_address, zero], + vec![], + 0, + ); + assert!( + MirArrayElementAddrOp::new(array_address_to_token) + .verify(&ctx) + .is_err(), + "an array element address cannot masquerade as a function token" + ); + + let slice_data_to_token = Operation::new( + &mut ctx, + MirExtractFieldOp::get_concrete_op_info(), + vec![opaque_fn_ty.into()], + vec![marker_slice], + vec![], + 0, + ); + let slice_data_to_token = MirExtractFieldOp::new(slice_data_to_token); + slice_data_to_token.set_attr_index(&ctx, FieldIndexAttr(0)); + assert!( + slice_data_to_token.verify(&ctx).is_err(), + "a slice data address cannot masquerade as a function token" + ); + + // A ClosureFnPointer cast must not extract captured reference bits as a + // function pointer, expose them as RawMut, and then reborrow them as + // UniqueRef. Only a genuinely non-capturing, zero-sized closure may enter + // the opaque function-pointer path. + let captured_ref = MirPtrType::get_generic_with_kind( + &mut ctx, + i32_ty.into(), + false, + MirPointerKind::SharedRef, + ); + let captured_closure = MirStructType::get_with_full_layout( + &mut ctx, + "CapturedClosure".into(), + vec!["capture_0".into()], + vec![captured_ref.into()], + vec![0], + vec![0], + 8, + 8, + ); + let empty_closure = MirStructType::get_with_full_layout( + &mut ctx, + "NonCapturingClosure".into(), + vec![], + vec![], + vec![], + vec![], + 0, + 1, + ); + let closure_block = BasicBlock::new( + &mut ctx, + None, + vec![captured_closure.into(), empty_closure.into()], + ); + let captured_value = closure_block.deref(&ctx).get_argument(0); + let empty_value = closure_block.deref(&ctx).get_argument(1); + let closure_cast = |ctx: &mut Context, source| { + let op = Operation::new( + ctx, + MirCastOp::get_concrete_op_info(), + vec![opaque_fn_ty.into()], + vec![source], + vec![], + 0, + ); + MirCastOp::new(op) + .set_attr_cast_kind(ctx, MirCastKindAttr::PointerCoercionClosureFnPointer); + op + }; + let captured_to_fn = closure_cast(&mut ctx, captured_value); + assert!( + MirCastOp::new(captured_to_fn).verify(&ctx).is_err(), + "a closure carrying SharedRef bytes cannot become an opaque function pointer" + ); + assert!( + MirCastOp::new(closure_cast(&mut ctx, empty_value)) + .verify(&ctx) + .is_err(), + "the importer materializes a closure function token directly; the legacy cast is not lowerable" + ); + + let captured_fn_value = captured_to_fn.deref(&ctx).get_result(0); + let exposed_capture = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![raw_ty.into()], + vec![captured_fn_value], + vec![], + 0, + ); + let exposed_capture = MirCastOp::new(exposed_capture); + exposed_capture.set_attr_cast_kind(&ctx, MirCastKindAttr::FnPtrToPtr); + exposed_capture.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::RustCast); + assert!(exposed_capture.verify(&ctx).is_ok()); + + let exposed_capture_value = exposed_capture.get_operation().deref(&ctx).get_result(0); + let unique_capture = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![unique_ty.into()], + vec![exposed_capture_value], + vec![], + 0, + ); + let unique_capture = MirCastOp::new(unique_capture); + unique_capture.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + unique_capture.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); + assert!( + unique_capture.verify(&ctx).is_ok(), + "the chain must be stopped at ClosureFnPointer before later individually legal boundaries" + ); +} diff --git a/crates/dialect-mir/tests/ops_test/control_flow.rs b/crates/dialect-mir/tests/ops_test/control_flow.rs new file mode 100644 index 0000000000..5f2664f0b8 --- /dev/null +++ b/crates/dialect-mir/tests/ops_test/control_flow.rs @@ -0,0 +1,175 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use dialect_mir::ops::{MirAssertOp, MirCondBranchOp, MirFuncOp, MirGotoOp, MirReturnOp}; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::TypeAttr, + op_interfaces::OperandSegmentInterface, + types::{FP32Type, FunctionType, IntegerType, Signedness}, + }, + common_traits::Verify, + context::Context, + op::Op, + operation::Operation, +}; + +#[test] +fn test_mir_control_flow_verify() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); + + // 1. MirGotoOp + let target_block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); + let src_block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); + let arg_val = src_block.deref(&ctx).get_argument(0); + + let op = Operation::new( + &mut ctx, + MirGotoOp::get_concrete_op_info(), + vec![], + vec![arg_val], + vec![target_block], + 0, + ); + let goto_op = MirGotoOp::new(op); + assert!(goto_op.verify(&ctx).is_ok(), "Valid Goto"); + + let op_bad = Operation::new( + &mut ctx, + MirGotoOp::get_concrete_op_info(), + vec![], + vec![], + vec![target_block], + 0, + ); + assert!( + MirGotoOp::new(op_bad).verify(&ctx).is_err(), + "Goto missing operand" + ); + + // 2. MirCondBranchOp + let true_block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); + let false_block = BasicBlock::new(&mut ctx, None, vec![]); + let cond_block = BasicBlock::new(&mut ctx, None, vec![i1_ty.into(), i32_ty.into()]); + let cond_val = cond_block.deref(&ctx).get_argument(0); + let val = cond_block.deref(&ctx).get_argument(1); + + let (cond_flat, cond_sizes) = + MirCondBranchOp::compute_segment_sizes(vec![vec![cond_val], vec![val], vec![]]); + let op_cond = Operation::new( + &mut ctx, + MirCondBranchOp::get_concrete_op_info(), + vec![], + cond_flat, + vec![true_block, false_block], + 0, + ); + MirCondBranchOp::new(op_cond).set_operand_segment_sizes(&ctx, cond_sizes); + let cond_br = MirCondBranchOp::new(op_cond); + assert!(cond_br.verify(&ctx).is_ok(), "Valid CondBranch"); + + let (cond_bad_flat, cond_bad_sizes) = + MirCondBranchOp::compute_segment_sizes(vec![vec![cond_val], vec![], vec![]]); + let op_cond_bad = Operation::new( + &mut ctx, + MirCondBranchOp::get_concrete_op_info(), + vec![], + cond_bad_flat, + vec![true_block, false_block], + 0, + ); + MirCondBranchOp::new(op_cond_bad).set_operand_segment_sizes(&ctx, cond_bad_sizes); + assert!( + MirCondBranchOp::new(op_cond_bad).verify(&ctx).is_err(), + "CondBranch missing operand" + ); + + // 3. MirReturnOp + let func_ty = FunctionType::get(&ctx, vec![], vec![i32_ty.into()]); + let func_ty_attr = TypeAttr::new(func_ty.into()); + + let func_op_ptr = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let mir_func = MirFuncOp::new(&mut ctx, func_op_ptr, func_ty_attr); + let entry_block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); + let ret_val = entry_block.deref(&ctx).get_argument(0); + + let region = mir_func.get_operation().deref(&ctx).get_region(0); + entry_block.insert_at_front(region, &ctx); + + let ret_op = Operation::new( + &mut ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + vec![ret_val], + vec![], + 0, + ); + ret_op.insert_at_back(entry_block, &ctx); + + let mir_ret = MirReturnOp::new(ret_op); + assert!(mir_ret.verify(&ctx).is_ok(), "Valid Return"); + + let f32_ty = FP32Type::get(&ctx); + let f32_block = BasicBlock::new(&mut ctx, None, vec![f32_ty.into()]); + let f32_val = f32_block.deref(&ctx).get_argument(0); + + let ret_op_bad = Operation::new( + &mut ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + vec![f32_val], + vec![], + 0, + ); + ret_op_bad.insert_at_back(entry_block, &ctx); + + let mir_ret_bad = MirReturnOp::new(ret_op_bad); + assert!(mir_ret_bad.verify(&ctx).is_err(), "Return type mismatch"); + + // 4. MirAssertOp + let assert_succ = BasicBlock::new(&mut ctx, None, vec![]); + + let (assert_flat, assert_sizes) = + MirAssertOp::compute_segment_sizes(vec![vec![cond_val], vec![]]); + let op_assert = Operation::new( + &mut ctx, + MirAssertOp::get_concrete_op_info(), + vec![], + assert_flat, + vec![assert_succ], + 0, + ); + MirAssertOp::new(op_assert).set_operand_segment_sizes(&ctx, assert_sizes); + let assert_op = MirAssertOp::new(op_assert); + assert!(assert_op.verify(&ctx).is_ok(), "Valid Assert"); + + let (assert_bad_flat, assert_bad_sizes) = + MirAssertOp::compute_segment_sizes(vec![vec![val], vec![]]); + let op_assert_bad = Operation::new( + &mut ctx, + MirAssertOp::get_concrete_op_info(), + vec![], + assert_bad_flat, + vec![assert_succ], + 0, + ); + MirAssertOp::new(op_assert_bad).set_operand_segment_sizes(&ctx, assert_bad_sizes); + assert!( + MirAssertOp::new(op_assert_bad).verify(&ctx).is_err(), + "Assert cond type mismatch" + ); +} diff --git a/crates/dialect-mir/tests/ops_test/enums.rs b/crates/dialect-mir/tests/ops_test/enums.rs new file mode 100644 index 0000000000..09213860fd --- /dev/null +++ b/crates/dialect-mir/tests/ops_test/enums.rs @@ -0,0 +1,406 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use dialect_mir::{ + attributes::{FieldIndexAttr, VariantIndexAttr}, + ops::{MirConstructEnumOp, MirEnumPayloadOp, MirGetDiscriminantOp, MirSetDiscriminantOp}, + types::{EnumVariant, MirEnumType, MirPtrType}, +}; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::TypeAttr, + types::{IntegerType, Signedness}, + }, + common_traits::Verify, + context::Context, + op::Op, + operation::Operation, +}; + +#[test] +fn test_mir_set_discriminant_verify() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signed); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let unit = |name: &str| EnumVariant::unit(name.to_string()); + + let enum_ty = MirEnumType::get( + &mut ctx, + "DeviceState".to_string(), + i8_ty.into(), + vec![0, 1], + vec![ + unit("Empty"), + EnumVariant::new("Full".to_string(), vec![i32_ty.into()]), + ], + ); + + let enum_ptr_ty = MirPtrType::get_generic(&mut ctx, enum_ty.into(), true); + let blk = BasicBlock::new(&mut ctx, None, vec![enum_ptr_ty.into(), i8_ty.into()]); + let enum_ptr = blk.deref(&ctx).get_argument(0); + let discr_val = blk.deref(&ctx).get_argument(1); + + // Malformed IR must be diagnosed without panicking, regardless of whether + // the generated operand-count interface happens to run first. + let op_no_operands = Operation::new( + &mut ctx, + MirSetDiscriminantOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + let no_operands = MirSetDiscriminantOp::new(op_no_operands); + no_operands.set_attr_set_discriminant_variant_index(&ctx, VariantIndexAttr(0)); + assert!(no_operands.verify(&ctx).is_err(), "zero operands rejected"); + + let op_extra_operand = Operation::new( + &mut ctx, + MirSetDiscriminantOp::get_concrete_op_info(), + vec![], + vec![enum_ptr, discr_val], + vec![], + 0, + ); + let extra_operand = MirSetDiscriminantOp::new(op_extra_operand); + extra_operand.set_attr_set_discriminant_variant_index(&ctx, VariantIndexAttr(0)); + assert!( + extra_operand.verify(&ctx).is_err(), + "extra operand rejected" + ); + + // Valid: pointer to enum plus the semantic target variant attribute. + let op_valid = Operation::new( + &mut ctx, + MirSetDiscriminantOp::get_concrete_op_info(), + vec![], + vec![enum_ptr], + vec![], + 0, + ); + let valid = MirSetDiscriminantOp::new(op_valid); + valid.set_attr_set_discriminant_variant_index(&ctx, VariantIndexAttr(1)); + valid.set_attr_set_discriminant_enum_ty(&ctx, TypeAttr::new(enum_ty.into())); + assert!(valid.verify(&ctx).is_ok(), "Valid set_discriminant"); + + // Invalid: first operand is not a pointer. + let op_bad_ptr = Operation::new( + &mut ctx, + MirSetDiscriminantOp::get_concrete_op_info(), + vec![], + vec![discr_val], + vec![], + 0, + ); + assert!( + MirSetDiscriminantOp::new(op_bad_ptr).verify(&ctx).is_err(), + "Non-pointer enum operand rejected" + ); + + // Invalid: SetDiscriminant writes memory and therefore cannot accept an + // immutable pointer even when the pointee is the right enum. + let immutable_ptr_ty = MirPtrType::get_generic(&mut ctx, enum_ty.into(), false); + let immutable_block = BasicBlock::new(&mut ctx, None, vec![immutable_ptr_ty.into()]); + let immutable_ptr = immutable_block.deref(&ctx).get_argument(0); + let op_immutable = Operation::new( + &mut ctx, + MirSetDiscriminantOp::get_concrete_op_info(), + vec![], + vec![immutable_ptr], + vec![], + 0, + ); + let immutable = MirSetDiscriminantOp::new(op_immutable); + immutable.set_attr_set_discriminant_variant_index(&ctx, VariantIndexAttr(0)); + assert!( + immutable.verify(&ctx).is_err(), + "immutable pointer rejected" + ); + + // Invalid: pointer does not point to an enum. + let i32_ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), true); + let blk_i32 = BasicBlock::new(&mut ctx, None, vec![i32_ptr_ty.into(), i8_ty.into()]); + let i32_ptr = blk_i32.deref(&ctx).get_argument(0); + let op_bad_pointee = Operation::new( + &mut ctx, + MirSetDiscriminantOp::get_concrete_op_info(), + vec![], + vec![i32_ptr], + vec![], + 0, + ); + let bad_pointee = MirSetDiscriminantOp::new(op_bad_pointee); + bad_pointee.set_attr_set_discriminant_variant_index(&ctx, VariantIndexAttr(0)); + assert!( + bad_pointee.verify(&ctx).is_err(), + "Non-enum pointee rejected" + ); + + // Invalid: target attribute is required. + let op_missing_target = Operation::new( + &mut ctx, + MirSetDiscriminantOp::get_concrete_op_info(), + vec![], + vec![enum_ptr], + vec![], + 0, + ); + assert!( + MirSetDiscriminantOp::new(op_missing_target) + .verify(&ctx) + .is_err(), + "Missing target rejected" + ); + + // Invalid: target index is out of bounds. + let op_oob = Operation::new( + &mut ctx, + MirSetDiscriminantOp::get_concrete_op_info(), + vec![], + vec![enum_ptr], + vec![], + 0, + ); + let oob = MirSetDiscriminantOp::new(op_oob); + oob.set_attr_set_discriminant_variant_index(&ctx, VariantIndexAttr(2)); + assert!(oob.verify(&ctx).is_err(), "Out-of-bounds target rejected"); +} + +#[test] +fn test_mir_enum_ops_malformed_arity_is_diagnostic() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let i8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let unit_enum = MirEnumType::get( + &mut ctx, + "Unit".into(), + i8_ty.into(), + vec![0], + vec![EnumVariant::unit("Only".into())], + ); + let payload_enum = MirEnumType::get( + &mut ctx, + "Payload".into(), + i8_ty.into(), + vec![0], + vec![EnumVariant::new("Only".into(), vec![i32_ty.into()])], + ); + let block = BasicBlock::new(&mut ctx, None, vec![unit_enum.into(), payload_enum.into()]); + let unit_value = block.deref(&ctx).get_argument(0); + let payload_value = block.deref(&ctx).get_argument(1); + + let construct_no_result = Operation::new( + &mut ctx, + MirConstructEnumOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + let construct_no_result = MirConstructEnumOp::new(construct_no_result); + construct_no_result.set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(0)); + assert!(construct_no_result.verify(&ctx).is_err()); + + let construct_extra_result = Operation::new( + &mut ctx, + MirConstructEnumOp::get_concrete_op_info(), + vec![unit_enum.into(), unit_enum.into()], + vec![], + vec![], + 0, + ); + let construct_extra_result = MirConstructEnumOp::new(construct_extra_result); + construct_extra_result.set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(0)); + assert!(construct_extra_result.verify(&ctx).is_err()); + + let get_empty = Operation::new( + &mut ctx, + MirGetDiscriminantOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + assert!(MirGetDiscriminantOp::new(get_empty).verify(&ctx).is_err()); + let get_extra = Operation::new( + &mut ctx, + MirGetDiscriminantOp::get_concrete_op_info(), + vec![i8_ty.into(), i8_ty.into()], + vec![unit_value, unit_value], + vec![], + 0, + ); + assert!(MirGetDiscriminantOp::new(get_extra).verify(&ctx).is_err()); + + let payload_empty = Operation::new( + &mut ctx, + MirEnumPayloadOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + let payload_empty = MirEnumPayloadOp::new(payload_empty); + payload_empty.set_attr_payload_variant_index(&ctx, VariantIndexAttr(0)); + payload_empty.set_attr_payload_field_index(&ctx, FieldIndexAttr(0)); + assert!(payload_empty.verify(&ctx).is_err()); + + let payload_extra = Operation::new( + &mut ctx, + MirEnumPayloadOp::get_concrete_op_info(), + vec![i32_ty.into(), i32_ty.into()], + vec![payload_value, payload_value], + vec![], + 0, + ); + let payload_extra = MirEnumPayloadOp::new(payload_extra); + payload_extra.set_attr_payload_variant_index(&ctx, VariantIndexAttr(0)); + payload_extra.set_attr_payload_field_index(&ctx, FieldIndexAttr(0)); + assert!(payload_extra.verify(&ctx).is_err()); +} + +#[test] +fn test_mir_enum_payload_rejects_uninhabited_variant() { + use dialect_mir::types::{EnumCarrierKind, EnumEncoding, EnumLayoutKind}; + + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let i8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let enum_ty = MirEnumType::get_with_encoding( + &mut ctx, + "HasImpossibleField".into(), + i8_ty.into(), + vec![0, 1], + vec![ + EnumVariant::unit("Live".into()), + // An uninhabited variant's unused physical offsets need not fit + // the object. The verifier must stop them reaching lowering. + EnumVariant::new_with_layout( + "Impossible".into(), + vec![i32_ty.into()], + vec![64], + vec![4], + ), + ], + EnumEncoding { + tag_offset: 0, + total_size: 1, + abi_align: 1, + layout_kind: EnumLayoutKind::Direct, + carrier_kind: EnumCarrierKind::Integer, + carrier_width: 8, + variant_inhabited: vec![1, 0], + ..EnumEncoding::default() + }, + ); + assert!(enum_ty.verify(&ctx).is_ok()); + + let block = BasicBlock::new(&mut ctx, None, vec![enum_ty.into()]); + let value = block.deref(&ctx).get_argument(0); + let op = Operation::new( + &mut ctx, + MirEnumPayloadOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![value], + vec![], + 0, + ); + let payload = MirEnumPayloadOp::new(op); + payload.set_attr_payload_variant_index(&ctx, VariantIndexAttr(1)); + payload.set_attr_payload_field_index(&ctx, FieldIndexAttr(0)); + assert!(payload.verify(&ctx).is_err()); +} + +/// Reachable-but-dead MIR (e.g. the residual arms of `array::try_from_fn`) +/// can name uninhabited variants. The importer must lower such reads and +/// constructions to typed undefs; if it ever emits the real ops again, these +/// verifiers are the loud stop. +#[test] +fn test_uninhabited_enum_construct_and_discriminant_fail_verification() { + use dialect_mir::types::{EnumCarrierKind, EnumEncoding, EnumLayoutKind}; + + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let i8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + + // Constructing an uninhabited variant must fail verification. + let partial = MirEnumType::get_with_encoding( + &mut ctx, + "HasImpossibleVariant".into(), + i8_ty.into(), + vec![0, 1], + vec![ + EnumVariant::unit("Live".into()), + EnumVariant::new_with_layout( + "Impossible".into(), + vec![i32_ty.into()], + vec![0], + vec![4], + ), + ], + EnumEncoding { + tag_offset: 0, + total_size: 8, + abi_align: 4, + layout_kind: EnumLayoutKind::Direct, + carrier_kind: EnumCarrierKind::Integer, + carrier_width: 8, + variant_inhabited: vec![1, 0], + ..EnumEncoding::default() + }, + ); + let block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); + let field = block.deref(&ctx).get_argument(0); + let construct = Operation::new( + &mut ctx, + MirConstructEnumOp::get_concrete_op_info(), + vec![partial.into()], + vec![field], + vec![], + 0, + ); + let construct = MirConstructEnumOp::new(construct); + construct.set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(1)); + assert!(construct.verify(&ctx).is_err()); + + // Reading the discriminant of a fully uninhabited enum must fail + // verification. + let never = MirEnumType::get_with_encoding( + &mut ctx, + "Never".into(), + i8_ty.into(), + vec![], + vec![], + EnumEncoding { + tag_offset: 0, + total_size: 0, + abi_align: 1, + layout_kind: EnumLayoutKind::Empty, + carrier_kind: EnumCarrierKind::None, + carrier_width: 0, + variant_inhabited: vec![], + ..EnumEncoding::default() + }, + ); + assert!(never.verify(&ctx).is_ok()); + let never_block = BasicBlock::new(&mut ctx, None, vec![never.into()]); + let never_value = never_block.deref(&ctx).get_argument(0); + let get_discriminant = Operation::new( + &mut ctx, + MirGetDiscriminantOp::get_concrete_op_info(), + vec![i8_ty.into()], + vec![never_value], + vec![], + 0, + ); + let get_discriminant = MirGetDiscriminantOp::new(get_discriminant); + assert!(get_discriminant.verify(&ctx).is_err()); +} diff --git a/crates/dialect-mir/tests/ops_test/main.rs b/crates/dialect-mir/tests/ops_test/main.rs new file mode 100644 index 0000000000..a7d126d046 --- /dev/null +++ b/crates/dialect-mir/tests/ops_test/main.rs @@ -0,0 +1,16 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +// Tests build kinded fixture types directly; production minting lives in +// mir-importer's facts.rs (see the workspace clippy.toml disallowed-methods). +#![allow(clippy::disallowed_methods)] + +mod aggregates; +mod cast_authority; +mod control_flow; +mod enums; +mod memory; +mod pointer_kinds; +mod scalars_and_funcs; diff --git a/crates/dialect-mir/tests/ops_test/memory.rs b/crates/dialect-mir/tests/ops_test/memory.rs new file mode 100644 index 0000000000..ac740b48e2 --- /dev/null +++ b/crates/dialect-mir/tests/ops_test/memory.rs @@ -0,0 +1,340 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use dialect_mir::{ + ops::{MirGlobalAllocOp, MirLoadOp, MirPtrOffsetOp, MirStoreOp}, + types::{MirPointerKind, MirPtrType}, +}; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::{StringAttr, TypeAttr}, + types::{FP32Type, IntegerType, Signedness}, + }, + common_traits::Verify, + context::Context, + op::Op, + operation::Operation, + opts::mem2reg::{AllocInfo, PromotableOpInterface, PromotableOpKind}, +}; + +#[test] +fn test_mir_load_verify() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), false); + + let block = BasicBlock::new(&mut ctx, None, vec![ptr_ty.into()]); + let ptr_val = block.deref(&ctx).get_argument(0); + + let op = Operation::new( + &mut ctx, + MirLoadOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![ptr_val], + vec![], + 0, + ); + let mir_load = MirLoadOp::new(op); + assert!(mir_load.verify(&ctx).is_ok(), "Valid MirLoadOp"); + + let block_i32 = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); + let i32_val = block_i32.deref(&ctx).get_argument(0); + + let op_fail_operand = Operation::new( + &mut ctx, + MirLoadOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![i32_val], + vec![], + 0, + ); + let mir_load_fail_operand = MirLoadOp::new(op_fail_operand); + assert!( + mir_load_fail_operand.verify(&ctx).is_err(), + "MirLoadOp non-ptr operand" + ); + + let f32_ty = FP32Type::get(&ctx); + let op_fail_res = Operation::new( + &mut ctx, + MirLoadOp::get_concrete_op_info(), + vec![f32_ty.into()], + vec![ptr_val], + vec![], + 0, + ); + let mir_load_fail_res = MirLoadOp::new(op_fail_res); + assert!( + mir_load_fail_res.verify(&ctx).is_err(), + "MirLoadOp result mismatch" + ); +} + +#[test] +fn test_mir_load_volatile_is_not_promotable() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), false); + let block = BasicBlock::new(&mut ctx, None, vec![ptr_ty.into()]); + let ptr_val = block.deref(&ctx).get_argument(0); + + let op = Operation::new( + &mut ctx, + MirLoadOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![ptr_val], + vec![], + 0, + ); + let mir_load = MirLoadOp::new(op); + let alloc_info = AllocInfo { + ptr: ptr_val, + ty: i32_ty.into(), + }; + + assert!(!mir_load.is_volatile(&ctx)); + assert!(matches!( + mir_load.promotion_kind(&ctx, &alloc_info), + PromotableOpKind::Load + )); + + mir_load.set_volatile(&mut ctx, true); + + assert!(mir_load.is_volatile(&ctx)); + assert!(matches!( + mir_load.promotion_kind(&ctx, &alloc_info), + PromotableOpKind::NonPromotableUse + )); +} + +#[test] +fn test_mir_ptr_offset_verify() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), false); + let usize_ty = IntegerType::get(&ctx, 64, Signedness::Signless); + + let block = BasicBlock::new(&mut ctx, None, vec![ptr_ty.into(), usize_ty.into()]); + let ptr_val = block.deref(&ctx).get_argument(0); + let idx_val = block.deref(&ctx).get_argument(1); + + let op = Operation::new( + &mut ctx, + MirPtrOffsetOp::get_concrete_op_info(), + vec![ptr_ty.into()], + vec![ptr_val, idx_val], + vec![], + 0, + ); + let offset_op = MirPtrOffsetOp::new(op); + assert!(offset_op.verify(&ctx).is_ok(), "Valid MirPtrOffsetOp"); + assert!( + offset_op.is_inbounds(&ctx), + "ordinary pointer offsets default to inbounds" + ); + offset_op.set_inbounds(&mut ctx, false); + assert!( + !offset_op.is_inbounds(&ctx), + "wrapping pointer offsets retain their explicit semantics" + ); + + let block2 = BasicBlock::new(&mut ctx, None, vec![i32_ty.into(), usize_ty.into()]); + let i32_val = block2.deref(&ctx).get_argument(0); + let idx_val2 = block2.deref(&ctx).get_argument(1); + + let op_bad_base = Operation::new( + &mut ctx, + MirPtrOffsetOp::get_concrete_op_info(), + vec![ptr_ty.into()], + vec![i32_val, idx_val2], + vec![], + 0, + ); + assert!(MirPtrOffsetOp::new(op_bad_base).verify(&ctx).is_err()); + + let f32_ty = FP32Type::get(&ctx); + let ptr_f32_ty = MirPtrType::get_generic(&mut ctx, f32_ty.into(), false); + let op_bad_res = Operation::new( + &mut ctx, + MirPtrOffsetOp::get_concrete_op_info(), + vec![ptr_f32_ty.into()], + vec![ptr_val, idx_val], + vec![], + 0, + ); + assert!(MirPtrOffsetOp::new(op_bad_res).verify(&ctx).is_err()); +} + +#[test] +fn test_mir_store_verify() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), false); + let block = BasicBlock::new(&mut ctx, None, vec![ptr_ty.into(), i32_ty.into()]); + let ptr_val = block.deref(&ctx).get_argument(0); + let val = block.deref(&ctx).get_argument(1); + + let op = Operation::new( + &mut ctx, + MirStoreOp::get_concrete_op_info(), + vec![], + vec![ptr_val, val], + vec![], + 0, + ); + assert!(MirStoreOp::new(op).verify(&ctx).is_ok(), "Valid MirStoreOp"); + + // Invalid: store to non-ptr + let op_bad_ptr = Operation::new( + &mut ctx, + MirStoreOp::get_concrete_op_info(), + vec![], + vec![val, val], + vec![], + 0, + ); + assert!( + MirStoreOp::new(op_bad_ptr).verify(&ctx).is_err(), + "MirStoreOp non-ptr dest" + ); + + // Invalid: type mismatch + let f32_ty = FP32Type::get(&ctx); + let block2 = BasicBlock::new(&mut ctx, None, vec![f32_ty.into()]); + let f32_val = block2.deref(&ctx).get_argument(0); + let op_bad_type = Operation::new( + &mut ctx, + MirStoreOp::get_concrete_op_info(), + vec![], + vec![ptr_val, f32_val], + vec![], + 0, + ); + assert!( + MirStoreOp::new(op_bad_type).verify(&ctx).is_err(), + "MirStoreOp type mismatch" + ); +} + +#[test] +fn test_mir_store_volatile_is_not_promotable() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), false); + let block = BasicBlock::new(&mut ctx, None, vec![ptr_ty.into(), i32_ty.into()]); + let ptr_val = block.deref(&ctx).get_argument(0); + let val = block.deref(&ctx).get_argument(1); + + let op = Operation::new( + &mut ctx, + MirStoreOp::get_concrete_op_info(), + vec![], + vec![ptr_val, val], + vec![], + 0, + ); + let mir_store = MirStoreOp::new(op); + let alloc_info = AllocInfo { + ptr: ptr_val, + ty: i32_ty.into(), + }; + + assert!(!mir_store.is_volatile(&ctx)); + match mir_store.promotion_kind(&ctx, &alloc_info) { + PromotableOpKind::Store(stored) => assert!(stored == val), + _ => panic!("non-volatile store should be promotable"), + } + + mir_store.set_volatile(&mut ctx, true); + + assert!(mir_store.is_volatile(&ctx)); + assert!(matches!( + mir_store.promotion_kind(&ctx, &alloc_info), + PromotableOpKind::NonPromotableUse + )); +} + +#[test] +fn test_mir_global_alloc_verify() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let f32_ty = FP32Type::get(&ctx); + + // Helper: build a MirGlobalAllocOp whose result pointer is in `ptr_ty` + // address space, with valid attributes. + let build = |ctx: &mut Context, ptr_ty: pliron::r#type::TypedHandle| { + let op = Operation::new( + ctx, + MirGlobalAllocOp::get_concrete_op_info(), + vec![ptr_ty.into()], + vec![], + vec![], + 0, + ); + let alloc = MirGlobalAllocOp::new(op); + alloc.set_attr_global_type(ctx, TypeAttr::new(f32_ty.into())); + alloc.set_attr_global_key(ctx, StringAttr::new("k".to_string())); + alloc + }; + + // Global memory (addrspace 1) — the original allowed space. + let ptr_global = MirPtrType::get_global(&mut ctx, f32_ty.into(), true); + assert!( + build(&mut ctx, ptr_global).verify(&ctx).is_ok(), + "global addrspace accepted" + ); + + // Constant memory (addrspace 4) — added for `#[constant]` support. + let ptr_const = MirPtrType::get_constant(&mut ctx, f32_ty.into(), true); + assert!( + build(&mut ctx, ptr_const).verify(&ctx).is_ok(), + "constant addrspace accepted" + ); + + // Shared memory (addrspace 3) — must be rejected. + let ptr_shared = MirPtrType::get_shared(&mut ctx, f32_ty.into(), true); + assert!( + build(&mut ctx, ptr_shared).verify(&ctx).is_err(), + "shared addrspace rejected" + ); + + let ptr_unique_global = MirPtrType::get_with_kind( + &mut ctx, + f32_ty.into(), + true, + dialect_mir::types::address_space::GLOBAL, + MirPointerKind::UniqueRef, + ); + assert!( + build(&mut ctx, ptr_unique_global).verify(&ctx).is_err(), + "a storage allocation cannot directly claim UniqueRef" + ); + + // Missing required attributes. + let no_attrs = Operation::new( + &mut ctx, + MirGlobalAllocOp::get_concrete_op_info(), + vec![ptr_global.into()], + vec![], + vec![], + 0, + ); + assert!( + MirGlobalAllocOp::new(no_attrs).verify(&ctx).is_err(), + "missing attributes rejected" + ); +} diff --git a/crates/dialect-mir/tests/ops_test/pointer_kinds.rs b/crates/dialect-mir/tests/ops_test/pointer_kinds.rs new file mode 100644 index 0000000000..608b51ba56 --- /dev/null +++ b/crates/dialect-mir/tests/ops_test/pointer_kinds.rs @@ -0,0 +1,1256 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use dialect_mir::{ + attributes::{MirCastKindAttr, MirPointerKindAuthorityAttr}, + ops::{MirAllocaOp, MirCastOp, MirGlobalAllocOp, MirPtrOffsetOp, MirRefOp, MirSharedAllocOp}, + types::{ + MirArrayType, MirDisjointSliceType, MirPointerKind, MirPtrType, MirSliceType, + MirStructType, MirTupleType, + }, +}; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::{IntegerAttr, StringAttr, TypeAttr}, + types::{FP32Type, IntegerType, Signedness}, + }, + common_traits::Verify, + context::Context, + op::Op, + operation::Operation, + r#type::TypeHandle, + utils::apint::APInt, +}; +use std::num::NonZeroUsize; + +#[test] +fn test_mir_pointer_kind_distinguishes_references_from_raw_pointers() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let pointee: pliron::r#type::TypeHandle = i32_ty.into(); + + let shared_ref = + MirPtrType::get_generic_with_kind(&mut ctx, pointee, false, MirPointerKind::SharedRef); + let raw_const = + MirPtrType::get_generic_with_kind(&mut ctx, pointee, false, MirPointerKind::RawConst); + let unique_ref = + MirPtrType::get_generic_with_kind(&mut ctx, pointee, true, MirPointerKind::UniqueRef); + let raw_mut = + MirPtrType::get_generic_with_kind(&mut ctx, pointee, true, MirPointerKind::RawMut); + + assert_ne!( + shared_ref, raw_const, + "&T must remain distinct from *const T" + ); + assert_ne!( + unique_ref, raw_mut, + "&mut T must remain distinct from *mut T" + ); + assert_ne!( + shared_ref, unique_ref, + "&T must remain distinct from &mut T" + ); + assert_ne!( + raw_const, raw_mut, + "*const T must remain distinct from *mut T" + ); + + assert_eq!( + shared_ref.deref(&ctx).pointer_kind(), + MirPointerKind::SharedRef + ); + assert_eq!( + unique_ref.deref(&ctx).pointer_kind(), + MirPointerKind::UniqueRef + ); + assert_eq!( + raw_const.deref(&ctx).pointer_kind(), + MirPointerKind::RawConst + ); + assert_eq!(raw_mut.deref(&ctx).pointer_kind(), MirPointerKind::RawMut); +} + +#[test] +fn test_mir_slice_preserves_pointer_kind() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let u8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); + let element: pliron::r#type::TypeHandle = u8_ty.into(); + + let shared = MirSliceType::get_with_kind(&mut ctx, element, MirPointerKind::SharedRef); + let raw = MirSliceType::get_with_kind(&mut ctx, element, MirPointerKind::RawConst); + + assert_ne!(shared, raw, "&[T] must remain distinct from *const [T]"); + assert_eq!(shared.deref(&ctx).pointer_kind(), MirPointerKind::SharedRef); + assert_eq!(raw.deref(&ctx).pointer_kind(), MirPointerKind::RawConst); +} + +#[test] +fn test_mir_pointer_kind_mutability_consistency_verify() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let invalid_shared = MirPtrType { + pointee: i32_ty.into(), + is_mutable: true, + address_space: 0, + kind: MirPointerKind::SharedRef, + }; + let invalid_raw_mut = MirPtrType { + pointee: i32_ty.into(), + is_mutable: false, + address_space: 0, + kind: MirPointerKind::RawMut, + }; + + assert!(invalid_shared.verify(&ctx).is_err()); + assert!(invalid_raw_mut.verify(&ctx).is_err()); + + let invalid_shared_slice = MirSliceType { + element_ty: i32_ty.into(), + is_mutable: true, + kind: MirPointerKind::SharedRef, + }; + let valid_erased_mut_slice = MirSliceType { + element_ty: i32_ty.into(), + is_mutable: true, + kind: MirPointerKind::Erased, + }; + assert!(invalid_shared_slice.verify(&ctx).is_err()); + assert!(valid_erased_mut_slice.verify(&ctx).is_ok()); +} + +#[test] +fn test_alloca_cannot_claim_a_rust_pointer_kind() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let erased = MirPtrType::get_generic(&mut ctx, i32_ty.into(), true); + let erased_alloca = Operation::new( + &mut ctx, + MirAllocaOp::get_concrete_op_info(), + vec![erased.into()], + vec![], + vec![], + 0, + ); + assert!(MirAllocaOp::new(erased_alloca).verify(&ctx).is_ok()); + + let immutable_erased = MirPtrType::get_generic(&mut ctx, i32_ty.into(), false); + let immutable_alloca = Operation::new( + &mut ctx, + MirAllocaOp::get_concrete_op_info(), + vec![immutable_erased.into()], + vec![], + vec![], + 0, + ); + assert!( + MirAllocaOp::new(immutable_alloca).verify(&ctx).is_err(), + "an alloca cannot masquerade as the immutable canonical function-pointer carrier" + ); + + let shared_erased = MirPtrType::get_shared(&mut ctx, i32_ty.into(), true); + let shared_alloca = Operation::new( + &mut ctx, + MirAllocaOp::get_concrete_op_info(), + vec![shared_erased.into()], + vec![], + vec![], + 0, + ); + assert!( + MirAllocaOp::new(shared_alloca).verify(&ctx).is_err(), + "a stack allocation must remain in generic address space" + ); + + let unique = + MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::UniqueRef); + let unique_alloca = Operation::new( + &mut ctx, + MirAllocaOp::get_concrete_op_info(), + vec![unique.into()], + vec![], + vec![], + 0, + ); + assert!( + MirAllocaOp::new(unique_alloca).verify(&ctx).is_err(), + "compiler storage must not manufacture UniqueRef" + ); +} + +#[test] +fn test_shared_alloc_result_pointee_must_match_element_type() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let i8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Signless).into(); + let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signless).into(); + let usize_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let mismatched_result = MirPtrType::get_shared(&mut ctx, i64_ty, true); + let op = Operation::new( + &mut ctx, + MirSharedAllocOp::get_concrete_op_info(), + vec![mismatched_result.into()], + vec![], + vec![], + 0, + ); + let alloc = MirSharedAllocOp::new(op); + alloc.set_attr_elem_type(&ctx, TypeAttr::new(i8_ty)); + alloc.set_attr_size( + &ctx, + IntegerAttr::new(usize_ty, APInt::from_u64(1, NonZeroUsize::new(64).unwrap())), + ); + alloc.set_attr_alloc_key(&ctx, StringAttr::new("mismatched-shared".to_string())); + + let error = alloc + .verify(&ctx) + .expect_err("shared storage cannot claim an unrelated result pointee type"); + assert!( + error + .to_string() + .contains("pointee type must match elem_type"), + "{error}" + ); +} + +#[test] +fn test_pointer_kind_laundering_requires_explicit_authority() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let usize_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let raw_mut_ty = + MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::RawMut); + let erased_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), true); + let erased_read_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), false); + let unique_ty = + MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::UniqueRef); + let shared_ty = MirPtrType::get_generic_with_kind( + &mut ctx, + i32_ty.into(), + false, + MirPointerKind::SharedRef, + ); + let raw_const_ty = + MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), false, MirPointerKind::RawConst); + let block = BasicBlock::new( + &mut ctx, + None, + vec![ + raw_mut_ty.into(), + erased_ty.into(), + usize_ty.into(), + shared_ty.into(), + raw_const_ty.into(), + erased_read_ty.into(), + ], + ); + let raw_mut = block.deref(&ctx).get_argument(0); + let erased = block.deref(&ctx).get_argument(1); + let offset = block.deref(&ctx).get_argument(2); + let shared = block.deref(&ctx).get_argument(3); + let raw_const = block.deref(&ctx).get_argument(4); + let erased_read = block.deref(&ctx).get_argument(5); + + // The concrete laundering example: pointer arithmetic is not a Rust + // reborrow and cannot manufacture `&mut T` from `*mut T`. + let raw_offset_to_unique = Operation::new( + &mut ctx, + MirPtrOffsetOp::get_concrete_op_info(), + vec![unique_ty.into()], + vec![raw_mut, offset], + vec![], + 0, + ); + assert!( + MirPtrOffsetOp::new(raw_offset_to_unique) + .verify(&ctx) + .is_err(), + "ptr_offset must not invent UniqueRef from RawMut" + ); + + let erased_offset_to_unique = Operation::new( + &mut ctx, + MirPtrOffsetOp::get_concrete_op_info(), + vec![unique_ty.into()], + vec![erased, offset], + vec![], + 0, + ); + assert!( + MirPtrOffsetOp::new(erased_offset_to_unique) + .verify(&ctx) + .is_err(), + "ptr_offset must not recover UniqueRef from Erased" + ); + + let preserving_offset = Operation::new( + &mut ctx, + MirPtrOffsetOp::get_concrete_op_info(), + vec![raw_mut_ty.into()], + vec![raw_mut, offset], + vec![], + 0, + ); + assert!(MirPtrOffsetOp::new(preserving_offset).verify(&ctx).is_ok()); + + let erasing_offset = Operation::new( + &mut ctx, + MirPtrOffsetOp::get_concrete_op_info(), + vec![erased_ty.into()], + vec![raw_mut, offset], + vec![], + 0, + ); + assert!(MirPtrOffsetOp::new(erasing_offset).verify(&ctx).is_ok()); + + let mutability_launder = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![erased_ty.into()], + vec![erased_read], + vec![], + 0, + ); + let mutability_launder_cast = MirCastOp::new(mutability_launder); + mutability_launder_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + assert!( + mutability_launder_cast.verify(&ctx).is_err(), + "an unmarked cast cannot manufacture writable Erased evidence" + ); + + let reborrow = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![unique_ty.into()], + vec![raw_mut], + vec![], + 0, + ); + let reborrow_cast = MirCastOp::new(reborrow); + reborrow_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + assert!( + reborrow_cast.verify(&ctx).is_err(), + "an unmarked cast must not manufacture UniqueRef" + ); + reborrow_cast.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); + assert!( + reborrow_cast.verify(&ctx).is_ok(), + "a rustc-declared reborrow is the explicit authority" + ); + + for (target, authority) in [ + (unique_ty.into(), MirPointerKindAuthorityAttr::Reborrow), + (raw_mut_ty.into(), MirPointerKindAuthorityAttr::RawAddress), + ] { + let mutable_storage_boundary = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![target], + vec![erased], + vec![], + 0, + ); + let mutable_storage_boundary_cast = MirCastOp::new(mutable_storage_boundary); + mutable_storage_boundary_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + mutable_storage_boundary_cast.set_pointer_kind_authority(&mut ctx, authority); + assert!( + mutable_storage_boundary_cast.verify(&ctx).is_ok(), + "mutable compiler storage is a valid source for a mutable Rust boundary" + ); + } + + let raw_to_shared_static = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![shared_ty.into()], + vec![raw_const], + vec![], + 0, + ); + let raw_to_shared_static_cast = MirCastOp::new(raw_to_shared_static); + raw_to_shared_static_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + raw_to_shared_static_cast + .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::StaticAddress); + assert!( + raw_to_shared_static_cast.verify(&ctx).is_err(), + "StaticAddress may establish a typed static value only from Erased storage, not relabel an arbitrary raw pointer as SharedRef" + ); + + let raw_to_erased = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![erased_read_ty.into()], + vec![raw_const], + vec![], + 0, + ); + let raw_to_erased_cast = MirCastOp::new(raw_to_erased); + raw_to_erased_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + assert!(raw_to_erased_cast.verify(&ctx).is_ok()); + let erased_from_raw = raw_to_erased.deref(&ctx).get_result(0); + let laundered_static = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![shared_ty.into()], + vec![erased_from_raw], + vec![], + 0, + ); + let laundered_static_cast = MirCastOp::new(laundered_static); + laundered_static_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + laundered_static_cast + .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::StaticAddress); + assert!( + laundered_static_cast.verify(&ctx).is_err(), + "erasing RawConst must not make it valid StaticAddress storage" + ); + + let raw_mut_to_erased = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![erased_ty.into()], + vec![raw_mut], + vec![], + 0, + ); + let raw_mut_to_erased_cast = MirCastOp::new(raw_mut_to_erased); + raw_mut_to_erased_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + assert!(raw_mut_to_erased_cast.verify(&ctx).is_ok()); + let erased_from_raw_mut = raw_mut_to_erased.deref(&ctx).get_result(0); + let laundered_abi = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![unique_ty.into()], + vec![erased_from_raw_mut], + vec![], + 0, + ); + let laundered_abi_cast = MirCastOp::new(laundered_abi); + laundered_abi_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + laundered_abi_cast + .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::AbiBoundary); + assert!( + laundered_abi_cast.verify(&ctx).is_err(), + "erasing RawMut must not let AbiBoundary manufacture UniqueRef" + ); + + let global_erased_ty = MirPtrType::get_with_kind( + &mut ctx, + i32_ty.into(), + false, + dialect_mir::types::address_space::GLOBAL, + MirPointerKind::Erased, + ); + let global_op = Operation::new( + &mut ctx, + MirGlobalAllocOp::get_concrete_op_info(), + vec![global_erased_ty.into()], + vec![], + vec![], + 0, + ); + let global = MirGlobalAllocOp::new(global_op); + global.set_attr_global_type(&ctx, TypeAttr::new(i32_ty.into())); + global.set_attr_global_key(&ctx, StringAttr::new("lineage-global".to_string())); + assert!(global.verify(&ctx).is_ok()); + let global_storage = global_op.deref(&ctx).get_result(0); + let typed_global = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![shared_ty.into()], + vec![global_storage], + vec![], + 0, + ); + let typed_global_cast = MirCastOp::new(typed_global); + typed_global_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + typed_global_cast + .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::StaticAddress); + assert!( + typed_global_cast.verify(&ctx).is_ok(), + "StaticAddress accepts a verified global-allocation root" + ); + + let shared_erased_ty = MirPtrType::get_shared(&mut ctx, i32_ty.into(), true); + let shared_op = Operation::new( + &mut ctx, + MirSharedAllocOp::get_concrete_op_info(), + vec![shared_erased_ty.into()], + vec![], + vec![], + 0, + ); + let shared_alloc = MirSharedAllocOp::new(shared_op); + shared_alloc.set_attr_elem_type(&ctx, TypeAttr::new(i32_ty.into())); + shared_alloc.set_attr_size( + &ctx, + IntegerAttr::new(usize_ty, APInt::from_u64(1, NonZeroUsize::new(64).unwrap())), + ); + shared_alloc.set_attr_alloc_key(&ctx, StringAttr::new("lineage-shared".to_string())); + assert!(shared_alloc.verify(&ctx).is_ok()); + let shared_storage = shared_op.deref(&ctx).get_result(0); + let typed_shared = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![raw_mut_ty.into()], + vec![shared_storage], + vec![], + 0, + ); + let typed_shared_cast = MirCastOp::new(typed_shared); + typed_shared_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + typed_shared_cast + .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::StaticAddress); + assert!( + typed_shared_cast.verify(&ctx).is_ok(), + "StaticAddress accepts a verified shared-allocation root" + ); + + let alloca_op = Operation::new( + &mut ctx, + MirAllocaOp::get_concrete_op_info(), + vec![erased_ty.into()], + vec![], + vec![], + 0, + ); + let alloca = MirAllocaOp::new(alloca_op); + assert!(alloca.verify(&ctx).is_ok()); + let compiler_storage = alloca_op.deref(&ctx).get_result(0); + let typed_abi = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![raw_mut_ty.into()], + vec![compiler_storage], + vec![], + 0, + ); + let typed_abi_cast = MirCastOp::new(typed_abi); + typed_abi_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + typed_abi_cast.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::AbiBoundary); + assert!( + typed_abi_cast.verify(&ctx).is_ok(), + "AbiBoundary accepts verified compiler-owned alloca storage" + ); + + for (target, authority) in [ + (unique_ty.into(), MirPointerKindAuthorityAttr::Reborrow), + (raw_mut_ty.into(), MirPointerKindAuthorityAttr::RawAddress), + ( + raw_mut_ty.into(), + MirPointerKindAuthorityAttr::StaticAddress, + ), + (raw_mut_ty.into(), MirPointerKindAuthorityAttr::AbiBoundary), + ] { + let immutable_storage_boundary = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![target], + vec![erased_read], + vec![], + 0, + ); + let immutable_storage_boundary_cast = MirCastOp::new(immutable_storage_boundary); + immutable_storage_boundary_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + immutable_storage_boundary_cast.set_pointer_kind_authority(&mut ctx, authority); + assert!( + immutable_storage_boundary_cast.verify(&ctx).is_err(), + "an immutable Erased thin pointer cannot establish a mutable Rust pointer kind" + ); + } + + let wrong_authority = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![unique_ty.into()], + vec![erased], + vec![], + 0, + ); + let wrong_authority_cast = MirCastOp::new(wrong_authority); + wrong_authority_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + wrong_authority_cast + .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::RawAddress); + assert!( + wrong_authority_cast.verify(&ctx).is_err(), + "RawAddress authority cannot establish a reference kind" + ); + + let inline_asm_authority = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![unique_ty.into()], + vec![erased], + vec![], + 0, + ); + let inline_asm_authority_cast = MirCastOp::new(inline_asm_authority); + inline_asm_authority_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + inline_asm_authority_cast + .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::InlineAsm); + assert!( + inline_asm_authority_cast.verify(&ctx).is_err(), + "InlineAsm is a producer authority and must never authorize a cast" + ); + + let integer_reborrow = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![unique_ty.into()], + vec![offset], + vec![], + 0, + ); + let integer_reborrow_cast = MirCastOp::new(integer_reborrow); + integer_reborrow_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + integer_reborrow_cast + .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); + assert!( + integer_reborrow_cast.verify(&ctx).is_err(), + "Reborrow authority cannot reinterpret an integer as a Rust reference" + ); + integer_reborrow_cast + .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::RustCast); + assert!( + integer_reborrow_cast.verify(&ctx).is_err(), + "RustCast authority cannot make an integer a PtrToPtr operand" + ); + integer_reborrow_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::FnPtrToPtr); + assert!( + integer_reborrow_cast.verify(&ctx).is_err(), + "FnPtrToPtr also requires a real pointer carrier" + ); + + let f32_ty = FP32Type::get(&ctx); + let wrong_pointee_ty = + MirPtrType::get_generic_with_kind(&mut ctx, f32_ty.into(), true, MirPointerKind::RawMut); + let wrong_pointee_block = BasicBlock::new(&mut ctx, None, vec![wrong_pointee_ty.into()]); + let wrong_pointee = wrong_pointee_block.deref(&ctx).get_argument(0); + let wrong_pointee_reborrow = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![unique_ty.into()], + vec![wrong_pointee], + vec![], + 0, + ); + let wrong_pointee_reborrow_cast = MirCastOp::new(wrong_pointee_reborrow); + wrong_pointee_reborrow_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + wrong_pointee_reborrow_cast + .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); + assert!( + wrong_pointee_reborrow_cast.verify(&ctx).is_err(), + "Reborrow authority must retain the pointee type" + ); + + for immutable_source in [shared, raw_const] { + let invalid_unique = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![unique_ty.into()], + vec![immutable_source], + vec![], + 0, + ); + let invalid_unique_cast = MirCastOp::new(invalid_unique); + invalid_unique_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + invalid_unique_cast + .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); + assert!( + invalid_unique_cast.verify(&ctx).is_err(), + "an immutable source cannot be relabelled as UniqueRef by Reborrow authority" + ); + + let invalid_raw_mut = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![raw_mut_ty.into()], + vec![immutable_source], + vec![], + 0, + ); + let invalid_raw_mut_cast = MirCastOp::new(invalid_raw_mut); + invalid_raw_mut_cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + invalid_raw_mut_cast + .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::RawAddress); + assert!( + invalid_raw_mut_cast.verify(&ctx).is_err(), + "an immutable source cannot be relabelled as RawMut by RawAddress authority" + ); + } +} + +#[test] +fn test_promoted_empty_mutable_reference_is_a_narrow_static_exception() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let empty_array_ty = MirArrayType::get(&mut ctx, i32_ty.into(), 0); + let nonempty_array_ty = MirArrayType::get(&mut ctx, i32_ty.into(), 1); + + let empty_storage_ty = MirPtrType::get_with_kind( + &mut ctx, + empty_array_ty.into(), + false, + dialect_mir::types::address_space::GLOBAL, + MirPointerKind::Erased, + ); + let empty_unique_ty = MirPtrType::get_generic_with_kind( + &mut ctx, + empty_array_ty.into(), + true, + MirPointerKind::UniqueRef, + ); + let empty_global_op = Operation::new( + &mut ctx, + MirGlobalAllocOp::get_concrete_op_info(), + vec![empty_storage_ty.into()], + vec![], + vec![], + 0, + ); + let empty_global = MirGlobalAllocOp::new(empty_global_op); + empty_global.set_attr_global_type(&ctx, TypeAttr::new(empty_array_ty.into())); + empty_global.set_attr_global_key( + &ctx, + StringAttr::new("promoted-empty-mutable-reference".to_string()), + ); + empty_global.set_alignment_value(&mut ctx, 4); + empty_global_op.deref_mut(&ctx).attributes.set( + "global_initializer_hex".try_into().unwrap(), + StringAttr::new(String::new()), + ); + empty_global.mark_immutable(&mut ctx); + assert!(empty_global.verify(&ctx).is_ok()); + + let empty_global_storage = empty_global_op.deref(&ctx).get_result(0); + let empty_static_borrow_op = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![empty_unique_ty.into()], + vec![empty_global_storage], + vec![], + 0, + ); + let empty_static_borrow = MirCastOp::new(empty_static_borrow_op); + empty_static_borrow.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + empty_static_borrow + .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::StaticAddress); + + empty_global.set_alignment_value(&mut ctx, 1); + assert!( + empty_static_borrow.verify(&ctx).is_err(), + "even [i32; 0] must retain i32's natural four-byte alignment" + ); + empty_global.set_alignment_value(&mut ctx, 4); + assert!( + empty_static_borrow.verify(&ctx).is_ok(), + "an immutable promoted [T; 0] global may back rustc's vacuous &mut []" + ); + + empty_global_op.deref_mut(&ctx).attributes.set( + "global_initializer_hex".try_into().unwrap(), + StringAttr::new("00".to_string()), + ); + assert!( + empty_static_borrow.verify(&ctx).is_err(), + "the exception requires an actually empty initializer, not merely the attribute" + ); + empty_global_op.deref_mut(&ctx).attributes.set( + "global_initializer_hex".try_into().unwrap(), + StringAttr::new(String::new()), + ); + empty_global_op.deref_mut(&ctx).attributes.set( + "global_initializer_relocations".try_into().unwrap(), + StringAttr::new("unexpected-relocation".to_string()), + ); + assert!( + empty_static_borrow.verify(&ctx).is_err(), + "zero-byte promoted storage cannot carry a relocation" + ); + + let mutable_empty_global_op = Operation::new( + &mut ctx, + MirGlobalAllocOp::get_concrete_op_info(), + vec![empty_storage_ty.into()], + vec![], + vec![], + 0, + ); + let mutable_empty_global = MirGlobalAllocOp::new(mutable_empty_global_op); + mutable_empty_global.set_attr_global_type(&ctx, TypeAttr::new(empty_array_ty.into())); + mutable_empty_global.set_attr_global_key( + &ctx, + StringAttr::new("non-promoted-empty-global".to_string()), + ); + assert!(mutable_empty_global.verify(&ctx).is_ok()); + let mutable_empty_storage = mutable_empty_global_op.deref(&ctx).get_result(0); + let mutable_empty_borrow_op = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![empty_unique_ty.into()], + vec![mutable_empty_storage], + vec![], + 0, + ); + let mutable_empty_borrow = MirCastOp::new(mutable_empty_borrow_op); + mutable_empty_borrow.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + mutable_empty_borrow + .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::StaticAddress); + assert!( + mutable_empty_borrow.verify(&ctx).is_err(), + "the [T; 0] exception requires compiler-promoted immutable storage" + ); + + let aligned_element_ty: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "Align16".to_string(), + vec!["value".to_string()], + vec![i32_ty.into()], + vec![], + vec![0], + 16, + 16, + ) + .into(); + let aligned_empty_array_ty = MirArrayType::get(&mut ctx, aligned_element_ty, 0); + let underaligned_storage_ty = MirPtrType::get_with_kind( + &mut ctx, + aligned_empty_array_ty.into(), + false, + dialect_mir::types::address_space::GLOBAL, + MirPointerKind::Erased, + ); + let aligned_empty_unique_ty = MirPtrType::get_generic_with_kind( + &mut ctx, + aligned_empty_array_ty.into(), + true, + MirPointerKind::UniqueRef, + ); + let underaligned_global_op = Operation::new( + &mut ctx, + MirGlobalAllocOp::get_concrete_op_info(), + vec![underaligned_storage_ty.into()], + vec![], + vec![], + 0, + ); + let underaligned_global = MirGlobalAllocOp::new(underaligned_global_op); + underaligned_global.set_attr_global_type(&ctx, TypeAttr::new(aligned_empty_array_ty.into())); + underaligned_global.set_attr_global_key( + &ctx, + StringAttr::new("underaligned-empty-reference".to_string()), + ); + underaligned_global.set_alignment_value(&mut ctx, 1); + underaligned_global_op.deref_mut(&ctx).attributes.set( + "global_initializer_hex".try_into().unwrap(), + StringAttr::new(String::new()), + ); + underaligned_global.mark_immutable(&mut ctx); + assert!(underaligned_global.verify(&ctx).is_ok()); + let underaligned_storage = underaligned_global_op.deref(&ctx).get_result(0); + let underaligned_borrow_op = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![aligned_empty_unique_ty.into()], + vec![underaligned_storage], + vec![], + 0, + ); + let underaligned_borrow = MirCastOp::new(underaligned_borrow_op); + underaligned_borrow.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + underaligned_borrow + .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::StaticAddress); + assert!( + underaligned_borrow.verify(&ctx).is_err(), + "promoted &mut [Align16; 0] requires at least 16-byte global alignment" + ); + + let nonempty_storage_ty = MirPtrType::get_with_kind( + &mut ctx, + nonempty_array_ty.into(), + false, + dialect_mir::types::address_space::GLOBAL, + MirPointerKind::Erased, + ); + let nonempty_unique_ty = MirPtrType::get_generic_with_kind( + &mut ctx, + nonempty_array_ty.into(), + true, + MirPointerKind::UniqueRef, + ); + let nonempty_global_op = Operation::new( + &mut ctx, + MirGlobalAllocOp::get_concrete_op_info(), + vec![nonempty_storage_ty.into()], + vec![], + vec![], + 0, + ); + let nonempty_global = MirGlobalAllocOp::new(nonempty_global_op); + nonempty_global.set_attr_global_type(&ctx, TypeAttr::new(nonempty_array_ty.into())); + nonempty_global.set_attr_global_key( + &ctx, + StringAttr::new("promoted-nonempty-mutable-reference".to_string()), + ); + nonempty_global.mark_immutable(&mut ctx); + assert!(nonempty_global.verify(&ctx).is_ok()); + let nonempty_global_storage = nonempty_global_op.deref(&ctx).get_result(0); + let nonempty_static_borrow_op = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![nonempty_unique_ty.into()], + vec![nonempty_global_storage], + vec![], + 0, + ); + let nonempty_static_borrow = MirCastOp::new(nonempty_static_borrow_op); + nonempty_static_borrow.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + nonempty_static_borrow + .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::StaticAddress); + assert!( + nonempty_static_borrow.verify(&ctx).is_err(), + "StaticAddress must never manufacture UniqueRef for non-empty promoted storage" + ); + + let block = BasicBlock::new( + &mut ctx, + None, + vec![empty_storage_ty.into(), empty_unique_ty.into()], + ); + let erased_block_argument = block.deref(&ctx).get_argument(0); + let block_argument_borrow_op = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![empty_unique_ty.into()], + vec![erased_block_argument], + vec![], + 0, + ); + let block_argument_borrow = MirCastOp::new(block_argument_borrow_op); + block_argument_borrow.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + block_argument_borrow + .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::StaticAddress); + assert!( + block_argument_borrow.verify(&ctx).is_err(), + "an Erased [T; 0] block argument has no proven promoted-global lineage" + ); + + let raw_empty_ty = MirPtrType::get_generic_with_kind( + &mut ctx, + empty_array_ty.into(), + false, + MirPointerKind::RawConst, + ); + let raw_block = BasicBlock::new(&mut ctx, None, vec![raw_empty_ty.into()]); + let raw_empty = raw_block.deref(&ctx).get_argument(0); + let erase_raw_op = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![empty_storage_ty.into()], + vec![raw_empty], + vec![], + 0, + ); + let erase_raw = MirCastOp::new(erase_raw_op); + erase_raw.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + assert!(erase_raw.verify(&ctx).is_ok()); + let erased_from_raw = erase_raw_op.deref(&ctx).get_result(0); + let laundered_borrow_op = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![empty_unique_ty.into()], + vec![erased_from_raw], + vec![], + 0, + ); + let laundered_borrow = MirCastOp::new(laundered_borrow_op); + laundered_borrow.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + laundered_borrow + .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::StaticAddress); + assert!( + laundered_borrow.verify(&ctx).is_err(), + "RawConst -> Erased must not launder a zero-length pointer into UniqueRef" + ); + + let byte_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); + let byte_storage_ty = MirPtrType::get_with_kind( + &mut ctx, + byte_ty.into(), + false, + dialect_mir::types::address_space::GLOBAL, + MirPointerKind::Erased, + ); + let byte_global_op = Operation::new( + &mut ctx, + MirGlobalAllocOp::get_concrete_op_info(), + vec![byte_storage_ty.into()], + vec![], + vec![], + 0, + ); + let byte_global = MirGlobalAllocOp::new(byte_global_op); + byte_global.set_attr_global_type(&ctx, TypeAttr::new(byte_ty.into())); + byte_global.set_attr_global_key( + &ctx, + StringAttr::new("misaligned-empty-reference-root".to_string()), + ); + byte_global.mark_immutable(&mut ctx); + assert!(byte_global.verify(&ctx).is_ok()); + let byte_storage = byte_global_op.deref(&ctx).get_result(0); + let retype_root_op = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![empty_storage_ty.into()], + vec![byte_storage], + vec![], + 0, + ); + let retype_root = MirCastOp::new(retype_root_op); + retype_root.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + assert!(retype_root.verify(&ctx).is_ok()); + let retyped_storage = retype_root_op.deref(&ctx).get_result(0); + let misaligned_borrow_op = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![empty_unique_ty.into()], + vec![retyped_storage], + vec![], + 0, + ); + let misaligned_borrow = MirCastOp::new(misaligned_borrow_op); + misaligned_borrow.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + misaligned_borrow + .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::StaticAddress); + assert!( + misaligned_borrow.verify(&ctx).is_err(), + "an immutable byte global cannot be retyped into an aligned &mut [i32; 0] capability" + ); +} + +fn promoted_empty_unique_ref_verifies( + ctx: &mut Context, + element_ty: TypeHandle, + alignment: u64, +) -> bool { + let empty_array_ty: TypeHandle = MirArrayType::get(ctx, element_ty, 0).into(); + let storage_ty = MirPtrType::get_with_kind( + ctx, + empty_array_ty, + false, + dialect_mir::types::address_space::GLOBAL, + MirPointerKind::Erased, + ); + let unique_ty = + MirPtrType::get_generic_with_kind(ctx, empty_array_ty, true, MirPointerKind::UniqueRef); + let global_op = Operation::new( + ctx, + MirGlobalAllocOp::get_concrete_op_info(), + vec![storage_ty.into()], + vec![], + vec![], + 0, + ); + let global = MirGlobalAllocOp::new(global_op); + global.set_attr_global_type(ctx, TypeAttr::new(empty_array_ty)); + global.set_attr_global_key(ctx, StringAttr::new("promoted-empty-shape".to_string())); + global.set_alignment_value(ctx, alignment); + global_op.deref_mut(ctx).attributes.set( + "global_initializer_hex".try_into().unwrap(), + StringAttr::new(String::new()), + ); + global.mark_immutable(ctx); + + let storage = global_op.deref(ctx).get_result(0); + let cast_op = Operation::new( + ctx, + MirCastOp::get_concrete_op_info(), + vec![unique_ty.into()], + vec![storage], + vec![], + 0, + ); + let cast = MirCastOp::new(cast_op); + cast.set_attr_cast_kind(ctx, MirCastKindAttr::PtrToPtr); + cast.set_pointer_kind_authority(ctx, MirPointerKindAuthorityAttr::StaticAddress); + cast.verify(ctx).is_ok() +} + +#[test] +fn test_promoted_empty_alignment_covers_supported_fat_and_unit_elements() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let byte: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let word: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let unit: TypeHandle = MirTupleType::get(&mut ctx, vec![]).into(); + let slice: TypeHandle = + MirSliceType::get_with_kind(&mut ctx, byte, MirPointerKind::SharedRef).into(); + let disjoint: TypeHandle = MirDisjointSliceType::get(&mut ctx, byte).into(); + + assert!(promoted_empty_unique_ref_verifies(&mut ctx, unit, 1)); + assert!(promoted_empty_unique_ref_verifies(&mut ctx, slice, 8)); + assert!(promoted_empty_unique_ref_verifies(&mut ctx, disjoint, 8)); + assert!(promoted_empty_unique_ref_verifies(&mut ctx, word, 8)); + assert!( + !promoted_empty_unique_ref_verifies(&mut ctx, word, 12), + "a non-power-of-two numeric value is not a valid alignment guarantee" + ); + assert!( + !promoted_empty_unique_ref_verifies(&mut ctx, slice, 4), + "a stored slice value still requires pointer-word alignment" + ); +} + +#[test] +fn test_mir_ref_requires_exact_pointer_kind_authority() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let shared_ty = MirPtrType::get_generic_with_kind( + &mut ctx, + i32_ty.into(), + false, + MirPointerKind::SharedRef, + ); + let unique_ty = + MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::UniqueRef); + let raw_const_ty = + MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), false, MirPointerKind::RawConst); + let raw_mut_ty = + MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::RawMut); + let erased_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), false); + let global_shared_ty = MirPtrType::get_with_kind( + &mut ctx, + i32_ty.into(), + false, + dialect_mir::types::address_space::GLOBAL, + MirPointerKind::SharedRef, + ); + let block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); + let value = block.deref(&ctx).get_argument(0); + + let build = + |ctx: &mut Context, result_ty, mutable, authority: Option| { + let op = Operation::new( + ctx, + MirRefOp::get_concrete_op_info(), + vec![result_ty], + vec![value], + vec![], + 0, + ); + let reference = MirRefOp::new(op); + reference.set_mutable(ctx, mutable); + if let Some(authority) = authority { + reference.set_pointer_kind_authority(ctx, authority); + } + op + }; + + for (result_ty, mutable, authority) in [ + ( + shared_ty.into(), + false, + MirPointerKindAuthorityAttr::Reborrow, + ), + ( + unique_ty.into(), + true, + MirPointerKindAuthorityAttr::Reborrow, + ), + ( + raw_const_ty.into(), + false, + MirPointerKindAuthorityAttr::RawAddress, + ), + ( + raw_mut_ty.into(), + true, + MirPointerKindAuthorityAttr::RawAddress, + ), + ( + shared_ty.into(), + false, + MirPointerKindAuthorityAttr::StaticAddress, + ), + ] { + assert!( + MirRefOp::new(build(&mut ctx, result_ty, mutable, Some(authority))) + .verify(&ctx) + .is_ok() + ); + } + + assert!( + MirRefOp::new(build(&mut ctx, shared_ty.into(), false, None)) + .verify(&ctx) + .is_err(), + "mir.ref must visibly identify its Rust semantic origin" + ); + assert!( + MirRefOp::new(build(&mut ctx, erased_ty.into(), false, None)) + .verify(&ctx) + .is_err(), + "mir.ref is an explicit pointer-creation boundary, not generic Erased storage" + ); + assert!( + MirRefOp::new(build( + &mut ctx, + shared_ty.into(), + false, + Some(MirPointerKindAuthorityAttr::RawAddress), + )) + .verify(&ctx) + .is_err(), + "RawAddress cannot manufacture SharedRef" + ); + assert!( + MirRefOp::new(build( + &mut ctx, + unique_ty.into(), + true, + Some(MirPointerKindAuthorityAttr::StaticAddress), + )) + .verify(&ctx) + .is_err(), + "constant/static materialization cannot manufacture uniqueness" + ); + assert!( + MirRefOp::new(build( + &mut ctx, + unique_ty.into(), + false, + Some(MirPointerKindAuthorityAttr::Reborrow), + )) + .verify(&ctx) + .is_err(), + "reference kind must agree with the operation's mutability" + ); + assert!( + MirRefOp::new(build( + &mut ctx, + global_shared_ty.into(), + false, + Some(MirPointerKindAuthorityAttr::Reborrow), + )) + .verify(&ctx) + .is_err(), + "mir.ref materializes generic stack storage" + ); +} diff --git a/crates/dialect-mir/tests/ops_test/scalars_and_funcs.rs b/crates/dialect-mir/tests/ops_test/scalars_and_funcs.rs new file mode 100644 index 0000000000..15be9832ef --- /dev/null +++ b/crates/dialect-mir/tests/ops_test/scalars_and_funcs.rs @@ -0,0 +1,447 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use dialect_mir::{ + attributes::MirCastKindAttr, + ops::{ + MirAddOp, MirAssignOp, MirCallOp, MirCastOp, MirCheckedAddOp, MirCmpOp, MirConstantOp, + MirDivOp, MirEqOp, MirFuncOp, MirGeOp, MirGtOp, MirLeOp, MirLtOp, MirMulOp, MirNeOp, + MirNegOp, MirNotOp, MirRemOp, MirSubOp, + }, + types::{EnumVariant, MirEnumType, MirTupleType}, +}; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::{IntegerAttr, StringAttr, TypeAttr}, + types::{FP32Type, FunctionType, IntegerType, Signedness}, + }, + common_traits::Verify, + context::Context, + op::Op, + operation::Operation, + utils::apint::APInt, +}; +use std::num::NonZeroUsize; + +#[test] +fn test_mir_arithmetic_verify() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into(), i32_ty.into()]); + let lhs = block.deref(&ctx).get_argument(0); + + let check_bin_op = |opid: ( + fn(pliron::context::Ptr) -> pliron::op::OpObj, + std::any::TypeId, + ), + name: &str| { + let mut context = Context::new(); + dialect_mir::register(&mut context); + let ty = IntegerType::get(&context, 32, Signedness::Signed); + let blk = BasicBlock::new(&mut context, None, vec![ty.into(), ty.into()]); + let l = blk.deref(&context).get_argument(0); + let r = blk.deref(&context).get_argument(1); + + let op = Operation::new(&mut context, opid, vec![ty.into()], vec![l, r], vec![], 0); + assert!(op.verify(&context).is_ok(), "Valid {}", name); + + let f32_t = FP32Type::get(&context); + let blk2 = BasicBlock::new(&mut context, None, vec![f32_t.into()]); + let f32_val = blk2.deref(&context).get_argument(0); + + let op_bad = Operation::new( + &mut context, + opid, + vec![ty.into()], + vec![l, f32_val], + vec![], + 0, + ); + assert!(op_bad.verify(&context).is_err(), "Type mismatch {}", name); + }; + + check_bin_op(MirAddOp::get_concrete_op_info(), "Add"); + check_bin_op(MirSubOp::get_concrete_op_info(), "Sub"); + check_bin_op(MirMulOp::get_concrete_op_info(), "Mul"); + check_bin_op(MirDivOp::get_concrete_op_info(), "Div"); + check_bin_op(MirRemOp::get_concrete_op_info(), "Rem"); + + let op_neg = Operation::new( + &mut ctx, + MirNegOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![lhs], + vec![], + 0, + ); + assert!(op_neg.verify(&ctx).is_ok(), "Valid Neg"); + + let f32_ty = FP32Type::get(&ctx); + let op_neg_bad = Operation::new( + &mut ctx, + MirNegOp::get_concrete_op_info(), + vec![f32_ty.into()], + vec![lhs], + vec![], + 0, + ); + assert!(op_neg_bad.verify(&ctx).is_err(), "Neg type mismatch"); + + let op_not = Operation::new( + &mut ctx, + MirNotOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![lhs], + vec![], + 0, + ); + assert!(op_not.verify(&ctx).is_ok(), "Valid Not"); +} + +#[test] +fn test_mir_misc_verify() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signed); + let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); + + // 1. MirConstantOp + let i32_signless = IntegerType::get(&ctx, 32, Signedness::Signless); + let width = NonZeroUsize::new(32).unwrap(); + let apint = APInt::from_u32(42, width); + let int_attr = IntegerAttr::new(i32_signless, apint); + + let const_op_ptr = Operation::new( + &mut ctx, + MirConstantOp::get_concrete_op_info(), + vec![i32_signless.into()], + vec![], + vec![], + 0, + ); + let const_op = MirConstantOp::new(const_op_ptr); + const_op.set_attr_value(&ctx, int_attr); + assert!(const_op.verify(&ctx).is_ok(), "Valid Constant"); + + // Mismatch type + let i64_signless = IntegerType::get(&ctx, 64, Signedness::Signless); + let i64_width = NonZeroUsize::new(64).unwrap(); + let i64_attr = IntegerAttr::new(i64_signless, APInt::from_u64(42, i64_width)); + const_op.set_attr_value(&ctx, i64_attr); + assert!(const_op.verify(&ctx).is_err(), "Constant type mismatch"); + + // 2. MirCastOp + let block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); + let arg = block.deref(&ctx).get_argument(0); + + let cast_op = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![i64_ty.into()], + vec![arg], + vec![], + 0, + ); + MirCastOp::new(cast_op).set_attr_cast_kind(&ctx, MirCastKindAttr::IntToInt); + assert!(MirCastOp::new(cast_op).verify(&ctx).is_ok(), "Valid Cast"); + + // 3. MirCheckedAddOp + let tuple_ty = MirTupleType::get(&mut ctx, vec![i32_ty.into(), i1_ty.into()]); + let block2 = BasicBlock::new(&mut ctx, None, vec![i32_ty.into(), i32_ty.into()]); + let lhs = block2.deref(&ctx).get_argument(0); + let rhs = block2.deref(&ctx).get_argument(1); + + let checked_add = Operation::new( + &mut ctx, + MirCheckedAddOp::get_concrete_op_info(), + vec![tuple_ty.into()], + vec![lhs, rhs], + vec![], + 0, + ); + assert!( + MirCheckedAddOp::new(checked_add).verify(&ctx).is_ok(), + "Valid CheckedAdd" + ); + + // Invalid result type (not tuple) + let checked_add_bad = Operation::new( + &mut ctx, + MirCheckedAddOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![lhs, rhs], + vec![], + 0, + ); + assert!( + MirCheckedAddOp::new(checked_add_bad).verify(&ctx).is_err(), + "CheckedAdd bad result" + ); +} + +#[test] +fn test_mir_comparison_verify() { + let check_cmp = |opid: ( + fn(pliron::context::Ptr) -> pliron::op::OpObj, + std::any::TypeId, + ), + name: &str| { + let mut context = Context::new(); + dialect_mir::register(&mut context); + let ty = IntegerType::get(&context, 32, Signedness::Signed); + let res_ty = IntegerType::get(&context, 1, Signedness::Signless); + let blk = BasicBlock::new(&mut context, None, vec![ty.into(), ty.into()]); + let l = blk.deref(&context).get_argument(0); + let r = blk.deref(&context).get_argument(1); + + let op = Operation::new( + &mut context, + opid, + vec![res_ty.into()], + vec![l, r], + vec![], + 0, + ); + assert!(op.verify(&context).is_ok(), "Valid {}", name); + + // Invalid operand types + let f32_ty = FP32Type::get(&context); + let blk2 = BasicBlock::new(&mut context, None, vec![f32_ty.into()]); + let f32_val = blk2.deref(&context).get_argument(0); + let op_bad = Operation::new( + &mut context, + opid, + vec![res_ty.into()], + vec![l, f32_val], + vec![], + 0, + ); + assert!(op_bad.verify(&context).is_err(), "Type mismatch {}", name); + + // Invalid result type + let op_bad_res = Operation::new( + &mut context, + opid, + vec![ty.into()], // i32 result instead of i1 + vec![l, r], + vec![], + 0, + ); + assert!( + op_bad_res.verify(&context).is_err(), + "Result type mismatch {}", + name + ); + }; + + check_cmp(MirEqOp::get_concrete_op_info(), "Eq"); + check_cmp(MirNeOp::get_concrete_op_info(), "Ne"); + check_cmp(MirLtOp::get_concrete_op_info(), "Lt"); + check_cmp(MirLeOp::get_concrete_op_info(), "Le"); + check_cmp(MirGtOp::get_concrete_op_info(), "Gt"); + check_cmp(MirGeOp::get_concrete_op_info(), "Ge"); + + let mut context = Context::new(); + dialect_mir::register(&mut context); + let i8_ty = IntegerType::get(&context, 8, Signedness::Signed); + let i32_ty = IntegerType::get(&context, 32, Signedness::Signed); + let unit = |name: &str| EnumVariant::unit(name.to_string()); + let ordering_ty = MirEnumType::get( + &mut context, + "Ordering".to_string(), + i8_ty.into(), + vec![255, 0, 1], + vec![unit("Less"), unit("Equal"), unit("Greater")], + ); + let blk = BasicBlock::new(&mut context, None, vec![i32_ty.into(), i32_ty.into()]); + let lhs = blk.deref(&context).get_argument(0); + let rhs = blk.deref(&context).get_argument(1); + let two_variant_ty = MirEnumType::get( + &mut context, + "Two".to_string(), + i8_ty.into(), + vec![0, 1], + vec![unit("A"), unit("B")], + ); + // Payload variants disqualify the Ordering shape. + let payload_ty = MirEnumType::get( + &mut context, + "ThreeWithPayload".to_string(), + i8_ty.into(), + vec![0, 1, 2], + vec![ + unit("A"), + EnumVariant::new("B".to_string(), vec![i32_ty.into()]), + unit("C"), + ], + ); + let mut check_cmp_result = |result_ty, valid| { + let op = Operation::new( + &mut context, + MirCmpOp::get_concrete_op_info(), + vec![result_ty], + vec![lhs, rhs], + vec![], + 0, + ); + assert_eq!(op.verify(&context).is_ok(), valid); + }; + check_cmp_result(ordering_ty.into(), true); + check_cmp_result(i32_ty.into(), false); + check_cmp_result(two_variant_ty.into(), false); + check_cmp_result(payload_ty.into(), false); + + // Float operands are rejected: rustc never emits BinOp::Cmp on floats. + let f32_ty = FP32Type::get(&context); + let fblk = BasicBlock::new(&mut context, None, vec![f32_ty.into(), f32_ty.into()]); + let flhs = fblk.deref(&context).get_argument(0); + let frhs = fblk.deref(&context).get_argument(1); + let float_cmp = Operation::new( + &mut context, + MirCmpOp::get_concrete_op_info(), + vec![ordering_ty.into()], + vec![flhs, frhs], + vec![], + 0, + ); + assert!( + float_cmp.verify(&context).is_err(), + "float mir.cmp must be rejected" + ); +} + +#[test] +fn test_mir_func_verify() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let func_ty = FunctionType::get(&ctx, vec![i32_ty.into()], vec![]); + let func_ty_attr = TypeAttr::new(func_ty.into()); + + // Valid Function + let op_ptr = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let mir_func = MirFuncOp::new(&mut ctx, op_ptr, func_ty_attr.clone()); + + // Add entry block with correct argument + let region = mir_func.get_operation().deref(&ctx).get_region(0); + let entry_block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); + entry_block.insert_at_front(region, &ctx); + + assert!(mir_func.verify(&ctx).is_ok(), "Valid MirFuncOp"); + + // Invalid: Argument count mismatch + let op_ptr2 = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let mir_func2 = MirFuncOp::new(&mut ctx, op_ptr2, func_ty_attr.clone()); + let region2 = mir_func2.get_operation().deref(&ctx).get_region(0); + // Block with 0 args + let entry_block2 = BasicBlock::new(&mut ctx, None, vec![]); + entry_block2.insert_at_front(region2, &ctx); + + assert!( + mir_func2.verify(&ctx).is_err(), + "MirFuncOp argument count mismatch" + ); + + // Invalid: Argument type mismatch + let op_ptr3 = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let mir_func3 = MirFuncOp::new(&mut ctx, op_ptr3, func_ty_attr); + let region3 = mir_func3.get_operation().deref(&ctx).get_region(0); + let f32_ty = FP32Type::get(&ctx); + let entry_block3 = BasicBlock::new(&mut ctx, None, vec![f32_ty.into()]); + entry_block3.insert_at_front(region3, &ctx); + + assert!( + mir_func3.verify(&ctx).is_err(), + "MirFuncOp argument type mismatch" + ); +} + +#[test] +fn test_mir_assign_verify() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); + let val = block.deref(&ctx).get_argument(0); + + let op = Operation::new( + &mut ctx, + MirAssignOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![val], + vec![], + 0, + ); + assert!( + MirAssignOp::new(op).verify(&ctx).is_ok(), + "Valid MirAssignOp" + ); + + let f32_ty = FP32Type::get(&ctx); + let op_bad = Operation::new( + &mut ctx, + MirAssignOp::get_concrete_op_info(), + vec![f32_ty.into()], + vec![val], + vec![], + 0, + ); + assert!( + MirAssignOp::new(op_bad).verify(&ctx).is_err(), + "MirAssignOp type mismatch" + ); +} + +#[test] +fn test_mir_call_verify() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let op = Operation::new( + &mut ctx, + MirCallOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + let call_op = MirCallOp::new(op); + + // Missing attribute + assert!(call_op.verify(&ctx).is_err(), "MirCallOp missing attribute"); + + // With attribute + let name = StringAttr::new("my_func".to_string()); + call_op.set_attr_callee(&ctx, name); + assert!(call_op.verify(&ctx).is_ok(), "Valid MirCallOp"); +} From fb66673356cd3f26ecec89cf66e201225f70ebb5 Mon Sep 17 00:00:00 2001 From: nihalpasham Date: Fri, 28 Aug 2026 22:06:47 +0530 Subject: [PATCH 02/18] refactor(dialect-nvvm): split tests/ops_test.rs into ops_test/ modules --- crates/dialect-nvvm/tests/ops_test.rs | 5575 ----------------- .../dialect-nvvm/tests/ops_test/allowlist.rs | 74 + crates/dialect-nvvm/tests/ops_test/atomics.rs | 516 ++ .../dialect-nvvm/tests/ops_test/carriers.rs | 923 +++ .../tests/ops_test/cluster_and_memory.rs | 761 +++ .../dialect-nvvm/tests/ops_test/inline_ptx.rs | 155 + crates/dialect-nvvm/tests/ops_test/main.rs | 16 + crates/dialect-nvvm/tests/ops_test/matrix.rs | 2278 +++++++ .../tests/ops_test/thread_and_warp.rs | 953 +++ 9 files changed, 5676 insertions(+), 5575 deletions(-) delete mode 100644 crates/dialect-nvvm/tests/ops_test.rs create mode 100644 crates/dialect-nvvm/tests/ops_test/allowlist.rs create mode 100644 crates/dialect-nvvm/tests/ops_test/atomics.rs create mode 100644 crates/dialect-nvvm/tests/ops_test/carriers.rs create mode 100644 crates/dialect-nvvm/tests/ops_test/cluster_and_memory.rs create mode 100644 crates/dialect-nvvm/tests/ops_test/inline_ptx.rs create mode 100644 crates/dialect-nvvm/tests/ops_test/main.rs create mode 100644 crates/dialect-nvvm/tests/ops_test/matrix.rs create mode 100644 crates/dialect-nvvm/tests/ops_test/thread_and_warp.rs diff --git a/crates/dialect-nvvm/tests/ops_test.rs b/crates/dialect-nvvm/tests/ops_test.rs deleted file mode 100644 index d8a7099fcf..0000000000 --- a/crates/dialect-nvvm/tests/ops_test.rs +++ /dev/null @@ -1,5575 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -// Tests build kinded fixture types directly; production minting lives in -// mir-importer's facts.rs (see the workspace clippy.toml disallowed-methods). -#![allow(clippy::disallowed_methods)] - -use dialect_mir::{ - attributes::MirPointerKindAuthorityAttr, - types::{MirPointerKind, MirPtrType, MirTupleType, address_space}, -}; -use dialect_nvvm::ops::{ - ActiveMaskOp, AssertFailOp, AtomicOrdering, AtomicRmwKind, AtomicScope, BarWarpSyncOp, - Barrier0Op, ClcQueryIsCanceledOp, ClusterBarrierModeAttr, ClusterBarrierOp, CpAsyncCa4Op, - CpAsyncCaZfill4Op, CpAsyncMbarrierArriveNoIncOp, CpAsyncMbarrierArriveNoIncSharedOp, - CpAsyncMbarrierArriveOp, CpAsyncMbarrierArriveSharedOp, CpAsyncWaitGroupOp, - CvtaGenericToSharedOffsetOp, Dp2aS32Op, Dp2aU32Op, Dp4aS32Op, Dp4aU32Op, ElectSyncOp, - FmaBf16x2Op, InlinePtxOp, LdmatrixElementAttr, LdmatrixLayoutAttr, LdmatrixMultiplicityAttr, - LdmatrixOp, LdmatrixShapeAttr, LdmatrixStateSpaceAttr, LdmatrixX1Op, LdmatrixX1TransOp, - LdmatrixX2Op, LdmatrixX2TransOp, LdmatrixX4Op, LdmatrixX4TransOp, MatchAllSyncI32Op, - MatchAllSyncI64Op, MatchAnySyncI32Op, MatchAnySyncI64Op, MbarrierArriveSharedOp, - MbarrierInitSharedOp, MbarrierInvalSharedOp, MbarrierTestWaitSharedOp, MmaM8N8K4F64Op, - MmaM16N8K8F32Tf32Op, MmaM16N8K16F32Bf16Op, MmaM16N8K16F32F16Op, MmaM16N8K32S32S8Op, - MovmatrixTransB16Op, NvvmAtomAddBf16x2Op, NvvmAtomAddF16x2Op, NvvmAtomicCmpxchgOp, - NvvmAtomicLoadOp, NvvmAtomicRmwOp, NvvmAtomicStoreOp, PackedAtomicAddOp, - PackedAtomicAtomicityAttr, PackedAtomicFormatAttr, PackedAtomicOrderingAttr, - PackedAtomicRoundingAttr, PackedAtomicScopeAttr, PackedAtomicStateSpaceAttr, - PackedAtomicSubnormalAttr, ReadPtxSregClusterIdxOp, ReadPtxSregDynamicSmemSizeOp, - ReadPtxSregGridIdOp, ReadPtxSregLaneIdOp, ReadPtxSregLanemaskEqOp, ReadPtxSregLanemaskGeOp, - ReadPtxSregLanemaskGtOp, ReadPtxSregLanemaskLeOp, ReadPtxSregLanemaskLtOp, - ReadPtxSregNclusterIdOp, ReadPtxSregNsmIdOp, ReadPtxSregNwarpIdOp, ReadPtxSregSmIdOp, - ReadPtxSregTidXOp, ReadPtxSregTotalSmemSizeOp, ReadPtxSregWarpIdOp, ReduxSyncAddOp, - ReduxSyncAndOp, ReduxSyncFmaxAbsNanOp, ReduxSyncFmaxAbsOp, ReduxSyncFmaxNanOp, ReduxSyncFmaxOp, - ReduxSyncFminAbsNanOp, ReduxSyncFminAbsOp, ReduxSyncFminNanOp, ReduxSyncFminOp, ReduxSyncMaxOp, - ReduxSyncMinOp, ReduxSyncOrOp, ReduxSyncUmaxOp, ReduxSyncUminOp, ReduxSyncXorOp, - RegisterMmaAccumulatorAttr, RegisterMmaElementAttr, RegisterMmaLayoutAttr, RegisterMmaOp, - RegisterMmaOperationAttr, RegisterMmaOverflowAttr, RegisterMmaShapeAttr, - ScalarArithmeticFormatAttr, ScalarArithmeticOp, ScalarArithmeticOperationAttr, - ScalarArithmeticRoundingAttr, ScalarArithmeticSaturationAttr, ScalarArithmeticSubnormalAttr, - ScalarConversionOp, ScalarConversionRoundingAttr, ScalarConversionSaturationAttr, - ShflSyncBflyI64Op, ShflSyncDownI64Op, ShflSyncIdxI64Op, ShflSyncUpI64Op, - SparseMmaAccumulatorAttr, SparseMmaElementAttr, SparseMmaLayoutAttr, SparseMmaMetadataAttr, - SparseMmaOp, SparseMmaOverflowAttr, SparseMmaSelectorAttr, SparseMmaShapeAttr, - StmatrixM8n8X4Op, Tcgen05AllocOp, Tcgen05CommitMulticastCg2Op, Tcgen05Ld16x32bx2X1RawOp, - Tcgen05Ld16x256bPureOp, Tcgen05MmaF16Op, ThreadfenceBlockOp, ThreadfenceOp, - ThreadfenceSystemOp, VoteSyncAllOp, VoteSyncAnyOp, VoteSyncBallotOp, VoteSyncUniOp, VprintfOp, - WgmmaMakeSmemDescOp, WgmmaMaxPendingAttr, WgmmaMmaGroupM64N64K16F32Bf16Op, - WgmmaMmaGroupValuesM64N64K8F32Tf32Op, WgmmaMmaGroupValuesM64N64K16F32Bf16Op, - WgmmaMmaGroupValuesM64N64K16F32F16Op, WgmmaMmaGroupValuesM64N128K16F32Bf16Op, - WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op, WgmmaMmaLoopValuesM64N64K16F32Bf16Op, - WgmmaMmaLoopValuesM64N64K16F32F16Op, WgmmaMmaM64N64K8F32Tf32Op, WgmmaMmaM64N64K16F32Bf16Op, - WgmmaMmaM64N64K16F32F16Op, WgmmaMmaM64N128K16F32Bf16Op, - WgmmaMmaPipelineValuesM64N64K16F32Bf16Op, -}; - -#[test] -fn handwritten_ops_match_reviewed_allowlist() { - use std::{fs, path::Path}; - - let ops_dir = Path::new(env!("CARGO_MANIFEST_DIR")).join("src/ops"); - let mut found = Vec::new(); - - for entry in fs::read_dir(&ops_dir).expect("read top-level NVVM ops directory") { - let path = entry.expect("read NVVM ops entry").path(); - if path.extension().and_then(|extension| extension.to_str()) != Some("rs") { - continue; - } - - let source = fs::read_to_string(&path).expect("read handwritten NVVM ops source"); - let file = path - .file_name() - .and_then(|name| name.to_str()) - .expect("NVVM ops file name"); - - for (marker, _) in source.match_indices("#[pliron_op(") { - let declaration = source[marker..] - .lines() - .find_map(|line| line.trim().strip_prefix("pub struct ")) - .expect("pliron_op must be followed by a public struct"); - let name = declaration - .split(|character: char| !character.is_ascii_alphanumeric() && character != '_') - .next() - .expect("NVVM op struct name"); - found.push((file.to_owned(), name.to_owned())); - } - } - - let mut expected = [ - ("asm.rs", "InlinePtxOp"), - ("atomic.rs", "NvvmAtomicLoadOp"), - ("atomic.rs", "NvvmAtomicStoreOp"), - ("atomic.rs", "NvvmAtomicFenceOp"), - ("atomic.rs", "NvvmAtomicRmwOp"), - ("atomic.rs", "NvvmAtomicCmpxchgOp"), - ("cluster.rs", "ReadPtxSregClusterIdxOp"), - ("cluster.rs", "ReadPtxSregNclusterIdOp"), - ("debug.rs", "AssertFailOp"), - ("debug.rs", "VprintfOp"), - ("grid.rs", "GridSyncOp"), - ("memory.rs", "CvtaGenericToSharedOffsetOp"), - ("wgmma.rs", "WgmmaMakeSmemDescOp"), - ("wgmma.rs", "WgmmaMmaM64N64K16F32Bf16Op"), - ("wgmma.rs", "WgmmaMmaM64N128K16F32Bf16Op"), - ("wgmma.rs", "WgmmaMmaM64N64K16F32F16Op"), - ("wgmma.rs", "WgmmaMmaM64N64K8F32Tf32Op"), - ("wgmma.rs", "WgmmaMmaGroupM64N64K16F32Bf16Op"), - ("wgmma.rs", "WgmmaMmaGroupValuesM64N64K16F32Bf16Op"), - ("wgmma.rs", "WgmmaMmaGroupValuesM64N128K16F32Bf16Op"), - ("wgmma.rs", "WgmmaMmaGroupValuesM64N64K16F32F16Op"), - ("wgmma.rs", "WgmmaMmaGroupValuesM64N64K8F32Tf32Op"), - ("wgmma.rs", "WgmmaMmaLoopValuesM64N64K16F32Bf16Op"), - ("wgmma.rs", "WgmmaMmaLoopValuesM64N64K16F32F16Op"), - ("wgmma.rs", "WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op"), - ("wgmma.rs", "WgmmaMmaPipelineValuesM64N64K16F32Bf16Op"), - ]; - expected.sort_unstable(); - found.sort_unstable(); - - assert_eq!( - found, - expected.map(|(file, op)| (file.to_owned(), op.to_owned())), - "top-level handwritten NVVM ops changed; generate leaf ops or review this allowlist" - ); -} - -#[test] -fn cluster_grid_compatibility_ops_keep_names_and_i32_shape() { - use pliron::builtin::types::{IntegerType, Signedness}; - use pliron::common_traits::Verify; - use pliron::context::Context; - use pliron::op::Op; - use pliron::operation::Operation; - - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - assert_eq!( - ReadPtxSregClusterIdxOp::get_opid_static().to_string(), - "nvvm.read_ptx_sreg_cluster_idx" - ); - assert_eq!( - ReadPtxSregNclusterIdOp::get_opid_static().to_string(), - "nvvm.read_ptx_sreg_nclusterid" - ); - - let i32_type = IntegerType::get(&ctx, 32, Signedness::Signless); - for op_info in [ - ReadPtxSregClusterIdxOp::get_concrete_op_info(), - ReadPtxSregNclusterIdOp::get_concrete_op_info(), - ] { - let op = Operation::new(&mut ctx, op_info, vec![i32_type.into()], vec![], vec![], 0); - assert!(op.deref(&ctx).verify(&ctx).is_ok()); - } -} - -#[test] -fn mapa_shared_cluster_requires_and_preserves_raw_pointer_kind() { - use dialect_mir::types::{MirPointerKind, MirPtrType, address_space}; - use dialect_nvvm::ops::MapaSharedClusterOp; - use pliron::{ - basic_block::BasicBlock, - builtin::types::{IntegerType, Signedness}, - common_traits::Verify, - context::Context, - op::Op, - operation::Operation, - }; - - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let raw_const_ty = - MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), false, MirPointerKind::RawConst); - let raw_mut_ty = - MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::RawMut); - let erased_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), false); - let block = BasicBlock::new( - &mut ctx, - None, - vec![ - raw_const_ty.into(), - raw_mut_ty.into(), - erased_ty.into(), - i32_ty.into(), - ], - ); - let raw_const = block.deref(&ctx).get_argument(0); - let raw_mut = block.deref(&ctx).get_argument(1); - let erased = block.deref(&ctx).get_argument(2); - let rank = block.deref(&ctx).get_argument(3); - - for (source, expected_kind) in [ - (raw_const, MirPointerKind::RawConst), - (raw_mut, MirPointerKind::RawMut), - ] { - let operation = MapaSharedClusterOp::build(&mut ctx, source, rank); - let result_ty = operation.deref(&ctx).get_result(0).get_type(&ctx); - let result_ty = result_ty.deref(&ctx); - let result_ptr = result_ty - .downcast_ref::() - .expect("mapa result must be a MIR pointer"); - assert_eq!(result_ptr.pointer_kind(), expected_kind); - assert_eq!(result_ptr.address_space(), address_space::CLUSTER_SHARED); - assert!(MapaSharedClusterOp::new(operation).verify(&ctx).is_ok()); - } - - let erased_result = MirPtrType::get_with_kind( - &mut ctx, - i32_ty.into(), - true, - address_space::CLUSTER_SHARED, - MirPointerKind::Erased, - ); - let erasing = Operation::new( - &mut ctx, - MapaSharedClusterOp::get_concrete_op_info(), - vec![erased_result.into()], - vec![raw_mut, rank], - vec![], - 0, - ); - assert!( - MapaSharedClusterOp::new(erasing).verify(&ctx).is_err(), - "cluster address mapping must preserve the source raw-pointer kind" - ); - - let erased_source = MapaSharedClusterOp::build(&mut ctx, erased, rank); - assert!( - MapaSharedClusterOp::new(erased_source) - .verify(&ctx) - .is_err(), - "an Erased data/function carrier cannot enter the raw-pointer intrinsic ABI" - ); - - let invented_unique = MirPtrType::get_with_kind( - &mut ctx, - i32_ty.into(), - true, - address_space::CLUSTER_SHARED, - MirPointerKind::UniqueRef, - ); - let laundering = Operation::new( - &mut ctx, - MapaSharedClusterOp::get_concrete_op_info(), - vec![invented_unique.into()], - vec![raw_mut, rank], - vec![], - 0, - ); - assert!( - MapaSharedClusterOp::new(laundering).verify(&ctx).is_err(), - "cluster address mapping must not turn RawMut into UniqueRef" - ); - - let invented_shared = MirPtrType::get_with_kind( - &mut ctx, - i32_ty.into(), - false, - address_space::CLUSTER_SHARED, - MirPointerKind::SharedRef, - ); - let laundering = Operation::new( - &mut ctx, - MapaSharedClusterOp::get_concrete_op_info(), - vec![invented_shared.into()], - vec![raw_const, rank], - vec![], - 0, - ); - assert!( - MapaSharedClusterOp::new(laundering).verify(&ctx).is_err(), - "cluster address mapping must not turn RawConst into SharedRef" - ); -} - -#[test] -fn generated_cluster_barrier_requires_one_closed_mode_and_no_values() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - for mode in [ - ClusterBarrierModeAttr::Arrive, - ClusterBarrierModeAttr::ArriveAligned, - ClusterBarrierModeAttr::ArriveRelaxed, - ClusterBarrierModeAttr::ArriveRelaxedAligned, - ClusterBarrierModeAttr::Wait, - ClusterBarrierModeAttr::WaitAligned, - ] { - let op = ClusterBarrierOp::build(&mut ctx, mode); - assert!(verify_op(&ClusterBarrierOp::new(op), &ctx).is_ok()); - } - - let missing_mode = Operation::new( - &mut ctx, - ClusterBarrierOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 0, - ); - assert!(verify_op(&ClusterBarrierOp::new(missing_mode), &ctx).is_err()); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let wrong_shape = Operation::new( - &mut ctx, - ClusterBarrierOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![], - vec![], - 0, - ); - ClusterBarrierOp::new(wrong_shape) - .set_attr_nvvm_cluster_barrier_mode(&ctx, ClusterBarrierModeAttr::Wait); - assert!(verify_op(&ClusterBarrierOp::new(wrong_shape), &ctx).is_err()); -} - -#[test] -fn generated_scalar_conversion_accepts_only_reviewed_f32_to_i32_variants() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - let f32_ty = FP32Type::get(&ctx); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); - let block = BasicBlock::new(&mut ctx, None, vec![f32_ty.into(), i32_ty.into()]); - let f32_value = block.deref(&ctx).get_argument(0); - let i32_value = block.deref(&ctx).get_argument(1); - - for (rounding, saturation) in [ - ( - ScalarConversionRoundingAttr::NearestAway, - ScalarConversionSaturationAttr::None, - ), - ( - ScalarConversionRoundingAttr::NearestAway, - ScalarConversionSaturationAttr::Satfinite, - ), - ( - ScalarConversionRoundingAttr::NearestEven, - ScalarConversionSaturationAttr::None, - ), - ( - ScalarConversionRoundingAttr::NearestEven, - ScalarConversionSaturationAttr::Relu, - ), - ( - ScalarConversionRoundingAttr::NearestEven, - ScalarConversionSaturationAttr::Satfinite, - ), - ( - ScalarConversionRoundingAttr::NearestEven, - ScalarConversionSaturationAttr::ReluSatfinite, - ), - ( - ScalarConversionRoundingAttr::TowardZero, - ScalarConversionSaturationAttr::None, - ), - ( - ScalarConversionRoundingAttr::TowardZero, - ScalarConversionSaturationAttr::Relu, - ), - ( - ScalarConversionRoundingAttr::TowardZero, - ScalarConversionSaturationAttr::Satfinite, - ), - ( - ScalarConversionRoundingAttr::TowardZero, - ScalarConversionSaturationAttr::ReluSatfinite, - ), - ] { - let op = ScalarConversionOp::build(&mut ctx, f32_value, rounding, saturation); - assert!(verify_op(&ScalarConversionOp::new(op), &ctx).is_ok()); - } - - let invalid_variant = ScalarConversionOp::build( - &mut ctx, - f32_value, - ScalarConversionRoundingAttr::NearestAway, - ScalarConversionSaturationAttr::Relu, - ); - assert!(verify_op(&ScalarConversionOp::new(invalid_variant), &ctx).is_err()); - - let wrong_operand = Operation::new( - &mut ctx, - ScalarConversionOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![i32_value], - vec![], - 0, - ); - let wrong_operand = ScalarConversionOp::new(wrong_operand); - wrong_operand - .set_attr_nvvm_scalar_conversion_rounding(&ctx, ScalarConversionRoundingAttr::NearestEven); - wrong_operand - .set_attr_nvvm_scalar_conversion_saturation(&ctx, ScalarConversionSaturationAttr::None); - assert!(verify_op(&wrong_operand, &ctx).is_err()); - - let wrong_result = Operation::new( - &mut ctx, - ScalarConversionOp::get_concrete_op_info(), - vec![i64_ty.into()], - vec![f32_value], - vec![], - 0, - ); - let wrong_result = ScalarConversionOp::new(wrong_result); - wrong_result - .set_attr_nvvm_scalar_conversion_rounding(&ctx, ScalarConversionRoundingAttr::TowardZero); - wrong_result - .set_attr_nvvm_scalar_conversion_saturation(&ctx, ScalarConversionSaturationAttr::None); - assert!(verify_op(&wrong_result, &ctx).is_err()); -} - -#[test] -fn generated_scalar_arithmetic_accepts_only_admitted_shapes_and_types() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - let f32_ty = FP32Type::get(&ctx); - let f64_ty = FP64Type::get(&ctx); - let block = BasicBlock::new(&mut ctx, None, vec![f32_ty.into(), f64_ty.into()]); - let f32_value = block.deref(&ctx).get_argument(0); - let f64_value = block.deref(&ctx).get_argument(1); - - let valid_f32 = ScalarArithmeticOp::build( - &mut ctx, - vec![f32_value, f32_value], - ScalarArithmeticFormatAttr::F32, - ScalarArithmeticOperationAttr::Mul, - ScalarArithmeticRoundingAttr::Rn, - ScalarArithmeticSubnormalAttr::Preserve, - ScalarArithmeticSaturationAttr::None, - ); - assert!(verify_op(&ScalarArithmeticOp::new(valid_f32), &ctx).is_ok()); - - let valid_f64 = ScalarArithmeticOp::build( - &mut ctx, - vec![f64_value, f64_value, f64_value], - ScalarArithmeticFormatAttr::F64, - ScalarArithmeticOperationAttr::Fma, - ScalarArithmeticRoundingAttr::Rz, - ScalarArithmeticSubnormalAttr::Preserve, - ScalarArithmeticSaturationAttr::None, - ); - assert!(verify_op(&ScalarArithmeticOp::new(valid_f64), &ctx).is_ok()); - - let valid_add = ScalarArithmeticOp::build( - &mut ctx, - vec![f32_value, f32_value], - ScalarArithmeticFormatAttr::F32, - ScalarArithmeticOperationAttr::Add, - ScalarArithmeticRoundingAttr::Rp, - ScalarArithmeticSubnormalAttr::Ftz, - ScalarArithmeticSaturationAttr::Sat, - ); - assert!(verify_op(&ScalarArithmeticOp::new(valid_add), &ctx).is_ok()); - - let f64_ftz = ScalarArithmeticOp::build( - &mut ctx, - vec![f64_value, f64_value], - ScalarArithmeticFormatAttr::F64, - ScalarArithmeticOperationAttr::Mul, - ScalarArithmeticRoundingAttr::Rn, - ScalarArithmeticSubnormalAttr::Ftz, - ScalarArithmeticSaturationAttr::None, - ); - assert!(verify_op(&ScalarArithmeticOp::new(f64_ftz), &ctx).is_err()); - - let wrong_arity = ScalarArithmeticOp::build( - &mut ctx, - vec![f32_value, f32_value], - ScalarArithmeticFormatAttr::F32, - ScalarArithmeticOperationAttr::Fma, - ScalarArithmeticRoundingAttr::Rn, - ScalarArithmeticSubnormalAttr::Preserve, - ScalarArithmeticSaturationAttr::None, - ); - assert!(verify_op(&ScalarArithmeticOp::new(wrong_arity), &ctx).is_err()); - - let wrong_type = ScalarArithmeticOp::build( - &mut ctx, - vec![f32_value, f64_value], - ScalarArithmeticFormatAttr::F32, - ScalarArithmeticOperationAttr::Mul, - ScalarArithmeticRoundingAttr::Rn, - ScalarArithmeticSubnormalAttr::Preserve, - ScalarArithmeticSaturationAttr::None, - ); - assert!(verify_op(&ScalarArithmeticOp::new(wrong_type), &ctx).is_err()); - - let missing_attrs = Operation::new( - &mut ctx, - ScalarArithmeticOp::get_concrete_op_info(), - vec![f32_ty.into()], - vec![f32_value, f32_value], - vec![], - 0, - ); - assert!(verify_op(&ScalarArithmeticOp::new(missing_attrs), &ctx).is_err()); -} - -#[test] -fn test_generated_cp_async_accepts_pointer_shapes_and_both_constant_kinds() { - use dialect_mir::ops::MirConstantOp; - use pliron::builtin::{attributes::IntegerAttr, ops::ConstantOp}; - use pliron::utils::apint::APInt; - use std::num::NonZeroUsize; - - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - - let u8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let dst_ty = MirPtrType::get_generic(&mut ctx, u8_ty.into(), false); - let src_ty = MirPtrType::get(&mut ctx, u8_ty.into(), true, address_space::GLOBAL); - let wrong_dst_ty = MirPtrType::get(&mut ctx, u8_ty.into(), true, address_space::GLOBAL); - let block = BasicBlock::new( - &mut ctx, - None, - vec![ - dst_ty.into(), - src_ty.into(), - wrong_dst_ty.into(), - u32_ty.into(), - ], - ); - let dst = block.deref(&ctx).get_argument(0); - let src = block.deref(&ctx).get_argument(1); - let wrong_dst = block.deref(&ctx).get_argument(2); - let dynamic = block.deref(&ctx).get_argument(3); - - let copy = CpAsyncCa4Op::build(&mut ctx, dst, src); - assert!(verify_op(&CpAsyncCa4Op::new(copy), &ctx).is_ok()); - let zfill = CpAsyncCaZfill4Op::build(&mut ctx, dst, src, dynamic); - assert!(verify_op(&CpAsyncCaZfill4Op::new(zfill), &ctx).is_ok()); - let wrong_space = CpAsyncCa4Op::build(&mut ctx, wrong_dst, src); - assert!(verify_op(&CpAsyncCa4Op::new(wrong_space), &ctx).is_err()); - - let value = IntegerAttr::new(u32_ty, APInt::from_u32(0, NonZeroUsize::new(32).unwrap())); - let builtin = ConstantOp::new(&mut ctx, value.clone().into()); - let builtin_value = builtin.get_operation().deref(&ctx).get_result(0); - let builtin_wait = CpAsyncWaitGroupOp::build(&mut ctx, builtin_value); - assert!(verify_op(&CpAsyncWaitGroupOp::new(builtin_wait), &ctx).is_ok()); - - let mir_constant = Operation::new( - &mut ctx, - MirConstantOp::get_concrete_op_info(), - vec![u32_ty.into()], - vec![], - vec![], - 0, - ); - MirConstantOp::new(mir_constant).set_attr_value(&ctx, value); - let mir_value = mir_constant.deref(&ctx).get_result(0); - let mir_wait = CpAsyncWaitGroupOp::build(&mut ctx, mir_value); - assert!(verify_op(&CpAsyncWaitGroupOp::new(mir_wait), &ctx).is_ok()); - - let dynamic_wait = CpAsyncWaitGroupOp::build(&mut ctx, dynamic); - assert!(verify_op(&CpAsyncWaitGroupOp::new(dynamic_wait), &ctx).is_err()); -} - -#[test] -fn generated_tcgen05_verifies_carriers_and_half_split_constants() { - use dialect_mir::ops::MirConstantOp; - use pliron::builtin::{attributes::IntegerAttr, ops::ConstantOp}; - use pliron::utils::apint::APInt; - use std::num::NonZeroUsize; - - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - - let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); - 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 f32_ty = FP32Type::get(&ctx); - let pointer_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), true); - let block = BasicBlock::new( - &mut ctx, - None, - vec![ - pointer_ty.into(), - i1_ty.into(), - i16_ty.into(), - i32_ty.into(), - i64_ty.into(), - f32_ty.into(), - ], - ); - let pointer = block.deref(&ctx).get_argument(0); - let predicate = block.deref(&ctx).get_argument(1); - let mask = block.deref(&ctx).get_argument(2); - let address = block.deref(&ctx).get_argument(3); - let dynamic_offset = block.deref(&ctx).get_argument(4); - - let offset_attr = IntegerAttr::new(i64_ty, APInt::from_i64(16, NonZeroUsize::new(64).unwrap())); - let builtin_offset = ConstantOp::new(&mut ctx, offset_attr.clone().into()) - .get_operation() - .deref(&ctx) - .get_result(0); - let mir_constant = Operation::new( - &mut ctx, - MirConstantOp::get_concrete_op_info(), - vec![i64_ty.into()], - vec![], - vec![], - 0, - ); - MirConstantOp::new(mir_constant).set_attr_value(&ctx, offset_attr); - let mir_offset = mir_constant.deref(&ctx).get_result(0); - - for offset in [builtin_offset, mir_offset] { - let load = Operation::new( - &mut ctx, - Tcgen05Ld16x32bx2X1RawOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![address, offset], - vec![], - 0, - ); - assert!(verify_op(&Tcgen05Ld16x32bx2X1RawOp::new(load), &ctx).is_ok()); - } - - let dynamic = Operation::new( - &mut ctx, - Tcgen05Ld16x32bx2X1RawOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![address, dynamic_offset], - vec![], - 0, - ); - assert!(verify_op(&Tcgen05Ld16x32bx2X1RawOp::new(dynamic), &ctx).is_err()); - - let wrong_offset_type = Operation::new( - &mut ctx, - Tcgen05Ld16x32bx2X1RawOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![address, address], - vec![], - 0, - ); - assert!(verify_op(&Tcgen05Ld16x32bx2X1RawOp::new(wrong_offset_type), &ctx).is_err()); - - let wrong_result_type = Operation::new( - &mut ctx, - Tcgen05Ld16x32bx2X1RawOp::get_concrete_op_info(), - vec![i64_ty.into()], - vec![address, builtin_offset], - vec![], - 0, - ); - assert!(verify_op(&Tcgen05Ld16x32bx2X1RawOp::new(wrong_result_type), &ctx).is_err()); - - let alloc = Operation::new( - &mut ctx, - Tcgen05AllocOp::get_concrete_op_info(), - vec![], - vec![pointer, address], - vec![], - 0, - ); - assert!(verify_op(&Tcgen05AllocOp::new(alloc), &ctx).is_ok()); - - let multicast = Operation::new( - &mut ctx, - Tcgen05CommitMulticastCg2Op::get_concrete_op_info(), - vec![], - vec![pointer, mask], - vec![], - 0, - ); - assert!(verify_op(&Tcgen05CommitMulticastCg2Op::new(multicast), &ctx).is_ok()); - - let mma = Operation::new( - &mut ctx, - Tcgen05MmaF16Op::get_concrete_op_info(), - vec![], - vec![address, dynamic_offset, dynamic_offset, address, predicate], - vec![], - 0, - ); - assert!(verify_op(&Tcgen05MmaF16Op::new(mma), &ctx).is_ok()); - - let pure_load = Operation::new( - &mut ctx, - Tcgen05Ld16x256bPureOp::get_concrete_op_info(), - vec![f32_ty.into(); 4], - vec![address], - vec![], - 0, - ); - assert!(verify_op(&Tcgen05Ld16x256bPureOp::new(pure_load), &ctx).is_ok()); -} - -#[test] -fn generated_cp_async_mbarrier_requires_mutable_generic_or_shared_u64() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let generic_u64 = MirPtrType::get_generic(&mut ctx, u64_ty.into(), true); - let shared_u64 = MirPtrType::get_shared(&mut ctx, u64_ty.into(), true); - let global_u64 = MirPtrType::get_global(&mut ctx, u64_ty.into(), true); - let immutable_u64 = MirPtrType::get_generic(&mut ctx, u64_ty.into(), false); - let generic_u32 = MirPtrType::get_generic(&mut ctx, u32_ty.into(), true); - let block = BasicBlock::new( - &mut ctx, - None, - vec![ - generic_u64.into(), - shared_u64.into(), - global_u64.into(), - immutable_u64.into(), - generic_u32.into(), - u64_ty.into(), - ], - ); - let generic = block.deref(&ctx).get_argument(0); - let shared = block.deref(&ctx).get_argument(1); - let global = block.deref(&ctx).get_argument(2); - let immutable = block.deref(&ctx).get_argument(3); - let wrong_pointee = block.deref(&ctx).get_argument(4); - let scalar = block.deref(&ctx).get_argument(5); - - macro_rules! check_bridge { - ($op:ty) => {{ - for barrier in [generic, shared] { - let valid = <$op>::build(&mut ctx, barrier); - assert!(verify_op(&<$op>::new(valid), &ctx).is_ok()); - } - for barrier in [global, immutable, wrong_pointee, scalar] { - let invalid = <$op>::build(&mut ctx, barrier); - assert!(verify_op(&<$op>::new(invalid), &ctx).is_err()); - } - let wrong_shape = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - vec![u64_ty.into()], - vec![generic], - vec![], - 0, - ); - assert!(verify_op(&<$op>::new(wrong_shape), &ctx).is_err()); - }}; - } - - check_bridge!(CpAsyncMbarrierArriveOp); - check_bridge!(CpAsyncMbarrierArriveSharedOp); - check_bridge!(CpAsyncMbarrierArriveNoIncOp); - check_bridge!(CpAsyncMbarrierArriveNoIncSharedOp); -} - -#[test] -fn generated_mbarrier_builders_and_verifiers_are_closed_over_their_shapes() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - - let u1_ty = IntegerType::get(&ctx, 1, Signedness::Unsigned); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let generic_ptr = MirPtrType::get_generic(&mut ctx, u64_ty.into(), false); - let shared_ptr = MirPtrType::get_shared(&mut ctx, u64_ty.into(), false); - let global_ptr = MirPtrType::get_global(&mut ctx, u64_ty.into(), false); - let block = BasicBlock::new( - &mut ctx, - None, - vec![ - generic_ptr.into(), - shared_ptr.into(), - global_ptr.into(), - i32_ty.into(), - u32_ty.into(), - u64_ty.into(), - ], - ); - let generic = block.deref(&ctx).get_argument(0); - let shared = block.deref(&ctx).get_argument(1); - let global = block.deref(&ctx).get_argument(2); - let signless_i32 = block.deref(&ctx).get_argument(3); - let count = block.deref(&ctx).get_argument(4); - let token = block.deref(&ctx).get_argument(5); - - for barrier in [generic, shared] { - let init = MbarrierInitSharedOp::build(&mut ctx, barrier, count); - assert!(MbarrierInitSharedOp::new(init).verify(&ctx).is_ok()); - let arrive = MbarrierArriveSharedOp::build(&mut ctx, barrier); - assert!(MbarrierArriveSharedOp::new(arrive).verify(&ctx).is_ok()); - let test_wait = MbarrierTestWaitSharedOp::build(&mut ctx, barrier, token); - assert!( - MbarrierTestWaitSharedOp::new(test_wait) - .verify(&ctx) - .is_ok() - ); - let inval = MbarrierInvalSharedOp::build(&mut ctx, barrier); - assert!(MbarrierInvalSharedOp::new(inval).verify(&ctx).is_ok()); - } - - for barrier in [token, global] { - let inval = MbarrierInvalSharedOp::build(&mut ctx, barrier); - assert!(MbarrierInvalSharedOp::new(inval).verify(&ctx).is_err()); - } - - let bad_count = MbarrierInitSharedOp::build(&mut ctx, shared, signless_i32); - assert!(MbarrierInitSharedOp::new(bad_count).verify(&ctx).is_err()); - let bad_arrive_result = Operation::new( - &mut ctx, - MbarrierArriveSharedOp::get_concrete_op_info(), - vec![u32_ty.into()], - vec![shared], - vec![], - 0, - ); - assert!( - MbarrierArriveSharedOp::new(bad_arrive_result) - .verify(&ctx) - .is_err() - ); - let bad_token = MbarrierTestWaitSharedOp::build(&mut ctx, shared, count); - assert!( - MbarrierTestWaitSharedOp::new(bad_token) - .verify(&ctx) - .is_err() - ); - let bad_predicate = Operation::new( - &mut ctx, - MbarrierTestWaitSharedOp::get_concrete_op_info(), - vec![u1_ty.into()], - vec![shared, token], - vec![], - 0, - ); - assert!( - MbarrierTestWaitSharedOp::new(bad_predicate) - .verify(&ctx) - .is_err() - ); - - let missing_operands = Operation::new( - &mut ctx, - MbarrierInitSharedOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 0, - ); - assert!( - MbarrierInitSharedOp::new(missing_operands) - .verify(&ctx) - .is_err() - ); -} -use pliron::{ - basic_block::BasicBlock, - builtin::types::{FP32Type, FP64Type, IntegerType, Signedness}, - common_traits::Verify, - context::Context, - op::{Op, verify_op}, - operation::Operation, - r#type::Typed, -}; - -#[test] -fn test_mma_m8n8k4_f64_requires_four_f64_operands_and_two_f64_results() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - let f64_ty = FP64Type::get(&ctx); - let f32_ty = FP32Type::get(&ctx); - let block = BasicBlock::new( - &mut ctx, - None, - vec![ - f64_ty.into(), - f64_ty.into(), - f64_ty.into(), - f64_ty.into(), - f32_ty.into(), - ], - ); - let f64_operands = (0..4) - .map(|index| block.deref(&ctx).get_argument(index)) - .collect::>(); - let f32_value = block.deref(&ctx).get_argument(4); - - let valid = Operation::new( - &mut ctx, - MmaM8N8K4F64Op::get_concrete_op_info(), - vec![f64_ty.into(), f64_ty.into()], - f64_operands.clone(), - vec![], - 0, - ); - assert!(verify_op(&MmaM8N8K4F64Op::new(valid), &ctx).is_ok()); - - let mut bad_operands = f64_operands.clone(); - bad_operands[2] = f32_value; - let invalid_operand = Operation::new( - &mut ctx, - MmaM8N8K4F64Op::get_concrete_op_info(), - vec![f64_ty.into(), f64_ty.into()], - bad_operands, - vec![], - 0, - ); - assert!(verify_op(&MmaM8N8K4F64Op::new(invalid_operand), &ctx).is_err()); - - let invalid_result = Operation::new( - &mut ctx, - MmaM8N8K4F64Op::get_concrete_op_info(), - vec![f64_ty.into(), f32_ty.into()], - f64_operands, - vec![], - 0, - ); - assert!(verify_op(&MmaM8N8K4F64Op::new(invalid_result), &ctx).is_err()); -} - -#[test] -fn test_movmatrix_requires_one_i32_operand_and_result() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); - let f32_ty = FP32Type::get(&ctx); - let block = BasicBlock::new( - &mut ctx, - None, - vec![i32_ty.into(), i64_ty.into(), f32_ty.into()], - ); - let i32_value = block.deref(&ctx).get_argument(0); - let i64_value = block.deref(&ctx).get_argument(1); - let f32_value = block.deref(&ctx).get_argument(2); - - let valid = Operation::new( - &mut ctx, - MovmatrixTransB16Op::get_concrete_op_info(), - vec![i32_ty.into()], - vec![i32_value], - vec![], - 0, - ); - assert!(verify_op(&MovmatrixTransB16Op::new(valid), &ctx).is_ok()); - - for (operand, result_type) in [ - (i64_value, i32_ty.into()), - (f32_value, i32_ty.into()), - (i32_value, i64_ty.into()), - (i32_value, f32_ty.into()), - ] { - let invalid = Operation::new( - &mut ctx, - MovmatrixTransB16Op::get_concrete_op_info(), - vec![result_type], - vec![operand], - vec![], - 0, - ); - assert!( - verify_op(&MovmatrixTransB16Op::new(invalid), &ctx).is_err(), - "movmatrix must reject non-i32 carriers" - ); - } -} - -fn make_ldmatrix_x2( - ctx: &mut Context, - pointer: pliron::value::Value, - result_types: Vec, -) -> LdmatrixOp { - let operation = Operation::new( - ctx, - LdmatrixOp::get_concrete_op_info(), - result_types, - vec![pointer], - vec![], - 0, - ); - let ldmatrix = LdmatrixOp::new(operation); - ldmatrix.set_attr_nvvm_ldmatrix_shape(ctx, LdmatrixShapeAttr::M8n8); - ldmatrix.set_attr_nvvm_ldmatrix_multiplicity(ctx, LdmatrixMultiplicityAttr::X2); - ldmatrix.set_attr_nvvm_ldmatrix_layout(ctx, LdmatrixLayoutAttr::Normal); - ldmatrix.set_attr_nvvm_ldmatrix_element(ctx, LdmatrixElementAttr::B16); - ldmatrix.set_attr_nvvm_ldmatrix_state_space(ctx, LdmatrixStateSpaceAttr::Shared); - ldmatrix -} - -#[test] -fn test_ldmatrix_accepts_only_generic_or_shared_u32_pointers() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - let u16_ty = IntegerType::get(&ctx, 16, Signedness::Unsigned); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let pointer_types = [ - MirPtrType::get_generic(&mut ctx, u32_ty.into(), false), - MirPtrType::get_shared(&mut ctx, u32_ty.into(), false), - MirPtrType::get_global(&mut ctx, u32_ty.into(), false), - MirPtrType::get_constant(&mut ctx, u32_ty.into(), false), - MirPtrType::get(&mut ctx, u32_ty.into(), false, address_space::LOCAL), - MirPtrType::get_generic(&mut ctx, u16_ty.into(), false), - ]; - let block = BasicBlock::new( - &mut ctx, - None, - pointer_types - .iter() - .map(|pointer| (*pointer).into()) - .collect(), - ); - - for index in 0..pointer_types.len() { - let pointer = block.deref(&ctx).get_argument(index); - let operation = make_ldmatrix_x2(&mut ctx, pointer, vec![u32_ty.into(), u32_ty.into()]); - let verified = verify_op(&operation, &ctx); - if index < 2 { - assert!(verified.is_ok(), "pointer case {index} should be accepted"); - } else { - assert!(verified.is_err(), "pointer case {index} should be rejected"); - } - } -} - -#[test] -fn generated_ldmatrix_verifier_rejects_zero_or_two_operands_without_panicking() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let pointer_ty = MirPtrType::get_shared(&mut ctx, u32_ty.into(), false); - let block = BasicBlock::new(&mut ctx, None, vec![pointer_ty.into(), pointer_ty.into()]); - let pointer0 = block.deref(&ctx).get_argument(0); - let pointer1 = block.deref(&ctx).get_argument(1); - - let zero = Operation::new( - &mut ctx, - LdmatrixOp::get_concrete_op_info(), - vec![u32_ty.into(); 4], - vec![], - vec![], - 0, - ); - assert!(LdmatrixOp::new(zero).verify(&ctx).is_err()); - - let two = Operation::new( - &mut ctx, - LdmatrixOp::get_concrete_op_info(), - vec![u32_ty.into(); 4], - vec![pointer0, pointer1], - vec![], - 0, - ); - assert!(LdmatrixOp::new(two).verify(&ctx).is_err()); - - let valid = LdmatrixOp::build( - &mut ctx, - pointer0, - LdmatrixShapeAttr::M8n8, - LdmatrixMultiplicityAttr::X4, - LdmatrixLayoutAttr::Normal, - LdmatrixElementAttr::B16, - LdmatrixStateSpaceAttr::Shared, - ); - assert!(LdmatrixOp::new(valid).verify(&ctx).is_ok()); -} - -#[test] -fn blackwell_ldmatrix_verifier_accepts_only_reviewed_shapes() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - let u8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let shared_u8 = MirPtrType::get_shared(&mut ctx, u8_ty.into(), false); - let generic_u8 = MirPtrType::get_generic(&mut ctx, u8_ty.into(), false); - let global_u8 = MirPtrType::get_global(&mut ctx, u8_ty.into(), false); - let shared_u32 = MirPtrType::get_shared(&mut ctx, u32_ty.into(), false); - let block = BasicBlock::new( - &mut ctx, - None, - vec![ - shared_u8.into(), - generic_u8.into(), - global_u8.into(), - shared_u32.into(), - ], - ); - let shared_pointer = block.deref(&ctx).get_argument(0); - - for pointer_index in [0, 1] { - let pointer = block.deref(&ctx).get_argument(pointer_index); - for (shape, multiplicity, layout, element, result_count) in [ - ( - LdmatrixShapeAttr::M16n16, - LdmatrixMultiplicityAttr::X1, - LdmatrixLayoutAttr::Transposed, - LdmatrixElementAttr::B8, - 2, - ), - ( - LdmatrixShapeAttr::M16n16, - LdmatrixMultiplicityAttr::X1, - LdmatrixLayoutAttr::Transposed, - LdmatrixElementAttr::B8x16B4x16P64, - 2, - ), - ( - LdmatrixShapeAttr::M16n16, - LdmatrixMultiplicityAttr::X1, - LdmatrixLayoutAttr::Transposed, - LdmatrixElementAttr::B8x16B6x16P32, - 2, - ), - ( - LdmatrixShapeAttr::M16n16, - LdmatrixMultiplicityAttr::X2, - LdmatrixLayoutAttr::Transposed, - LdmatrixElementAttr::B8, - 4, - ), - ( - LdmatrixShapeAttr::M16n16, - LdmatrixMultiplicityAttr::X2, - LdmatrixLayoutAttr::Transposed, - LdmatrixElementAttr::B8x16B4x16P64, - 4, - ), - ( - LdmatrixShapeAttr::M16n16, - LdmatrixMultiplicityAttr::X2, - LdmatrixLayoutAttr::Transposed, - LdmatrixElementAttr::B8x16B6x16P32, - 4, - ), - ( - LdmatrixShapeAttr::M8n16, - LdmatrixMultiplicityAttr::X1, - LdmatrixLayoutAttr::Normal, - LdmatrixElementAttr::B8x16B4x16P64, - 1, - ), - ( - LdmatrixShapeAttr::M8n16, - LdmatrixMultiplicityAttr::X1, - LdmatrixLayoutAttr::Normal, - LdmatrixElementAttr::B8x16B6x16P32, - 1, - ), - ( - LdmatrixShapeAttr::M8n16, - LdmatrixMultiplicityAttr::X2, - LdmatrixLayoutAttr::Normal, - LdmatrixElementAttr::B8x16B4x16P64, - 2, - ), - ( - LdmatrixShapeAttr::M8n16, - LdmatrixMultiplicityAttr::X2, - LdmatrixLayoutAttr::Normal, - LdmatrixElementAttr::B8x16B6x16P32, - 2, - ), - ( - LdmatrixShapeAttr::M8n16, - LdmatrixMultiplicityAttr::X4, - LdmatrixLayoutAttr::Normal, - LdmatrixElementAttr::B8x16B4x16P64, - 4, - ), - ( - LdmatrixShapeAttr::M8n16, - LdmatrixMultiplicityAttr::X4, - LdmatrixLayoutAttr::Normal, - LdmatrixElementAttr::B8x16B6x16P32, - 4, - ), - ] { - let op = LdmatrixOp::build( - &mut ctx, - pointer, - shape, - multiplicity, - layout, - element, - LdmatrixStateSpaceAttr::Shared, - ); - assert_eq!(op.deref(&ctx).get_num_results(), result_count); - assert!(LdmatrixOp::new(op).verify(&ctx).is_ok()); - } - } - - for (shape, multiplicity, layout, element) in [ - ( - LdmatrixShapeAttr::M16n16, - LdmatrixMultiplicityAttr::X4, - LdmatrixLayoutAttr::Transposed, - LdmatrixElementAttr::B8, - ), - ( - LdmatrixShapeAttr::M16n16, - LdmatrixMultiplicityAttr::X1, - LdmatrixLayoutAttr::Normal, - LdmatrixElementAttr::B8, - ), - ( - LdmatrixShapeAttr::M8n16, - LdmatrixMultiplicityAttr::X1, - LdmatrixLayoutAttr::Transposed, - LdmatrixElementAttr::B8x16B4x16P64, - ), - ( - LdmatrixShapeAttr::M8n16, - LdmatrixMultiplicityAttr::X1, - LdmatrixLayoutAttr::Normal, - LdmatrixElementAttr::B8, - ), - ] { - let op = LdmatrixOp::build( - &mut ctx, - shared_pointer, - shape, - multiplicity, - layout, - element, - LdmatrixStateSpaceAttr::Shared, - ); - assert!(LdmatrixOp::new(op).verify(&ctx).is_err()); - } - - for pointer_index in [2, 3] { - let pointer = block.deref(&ctx).get_argument(pointer_index); - let op = LdmatrixOp::build( - &mut ctx, - pointer, - LdmatrixShapeAttr::M16n16, - LdmatrixMultiplicityAttr::X1, - LdmatrixLayoutAttr::Transposed, - LdmatrixElementAttr::B8, - LdmatrixStateSpaceAttr::Shared, - ); - assert!(LdmatrixOp::new(op).verify(&ctx).is_err()); - } -} - -#[test] -fn classic_ldmatrix_compatibility_ops_keep_names_and_register_shapes() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let pointer_ty = MirPtrType::get_shared(&mut ctx, u32_ty.into(), false); - let block = BasicBlock::new(&mut ctx, None, vec![pointer_ty.into()]); - let pointer = block.deref(&ctx).get_argument(0); - - macro_rules! check_compat { - ($op:ty, $name:literal, $results:literal) => {{ - assert_eq!(<$op>::get_opid_static().to_string(), $name); - let valid = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - vec![u32_ty.into(); $results], - vec![pointer], - vec![], - 0, - ); - assert!(verify_op(&<$op>::new(valid), &ctx).is_ok()); - - let wrong_shape = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - vec![u32_ty.into(); $results + 1], - vec![pointer], - vec![], - 0, - ); - assert!(verify_op(&<$op>::new(wrong_shape), &ctx).is_err()); - }}; - } - - check_compat!(LdmatrixX1Op, "nvvm.ldmatrix_x1", 1); - check_compat!(LdmatrixX1TransOp, "nvvm.ldmatrix_x1_trans", 1); - check_compat!(LdmatrixX2Op, "nvvm.ldmatrix_x2", 2); - check_compat!(LdmatrixX2TransOp, "nvvm.ldmatrix_x2_trans", 2); - check_compat!(LdmatrixX4Op, "nvvm.ldmatrix_x4", 4); - check_compat!(LdmatrixX4TransOp, "nvvm.ldmatrix_x4_trans", 4); -} - -#[test] -fn test_mma_m16n8k16_bf16_verifies_exact_register_signature() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - let f32_ty = FP32Type::get(&ctx); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); - let block = BasicBlock::new( - &mut ctx, - None, - vec![f32_ty.into(), i32_ty.into(), i64_ty.into()], - ); - let f32_value = block.deref(&ctx).get_argument(0); - let i32_value = block.deref(&ctx).get_argument(1); - let i64_value = block.deref(&ctx).get_argument(2); - - let valid_operands = (0..4) - .map(|_| f32_value) - .chain((0..6).map(|_| i32_value)) - .collect(); - let valid = Operation::new( - &mut ctx, - MmaM16N8K16F32Bf16Op::get_concrete_op_info(), - vec![f32_ty.into(); 4], - valid_operands, - vec![], - 0, - ); - assert!(verify_op(&MmaM16N8K16F32Bf16Op::new(valid), &ctx).is_ok()); - - let bad_c_operands = (0..4) - .map(|index| if index == 0 { i32_value } else { f32_value }) - .chain((0..6).map(|_| i32_value)) - .collect(); - let bad_c = Operation::new( - &mut ctx, - MmaM16N8K16F32Bf16Op::get_concrete_op_info(), - vec![f32_ty.into(); 4], - bad_c_operands, - vec![], - 0, - ); - assert!(verify_op(&MmaM16N8K16F32Bf16Op::new(bad_c), &ctx).is_err()); - - let bad_a_operands = (0..4) - .map(|_| f32_value) - .chain((0..6).map(|index| if index == 0 { i64_value } else { i32_value })) - .collect(); - let bad_a = Operation::new( - &mut ctx, - MmaM16N8K16F32Bf16Op::get_concrete_op_info(), - vec![f32_ty.into(); 4], - bad_a_operands, - vec![], - 0, - ); - assert!(verify_op(&MmaM16N8K16F32Bf16Op::new(bad_a), &ctx).is_err()); - - let bad_result = Operation::new( - &mut ctx, - MmaM16N8K16F32Bf16Op::get_concrete_op_info(), - vec![f32_ty.into(), f32_ty.into(), f32_ty.into(), i32_ty.into()], - (0..4) - .map(|_| f32_value) - .chain((0..6).map(|_| i32_value)) - .collect(), - vec![], - 0, - ); - assert!(verify_op(&MmaM16N8K16F32Bf16Op::new(bad_result), &ctx).is_err()); - - let bad_arity = Operation::new( - &mut ctx, - MmaM16N8K16F32Bf16Op::get_concrete_op_info(), - vec![f32_ty.into(); 4], - (0..4) - .map(|_| f32_value) - .chain((0..5).map(|_| i32_value)) - .collect(), - vec![], - 0, - ); - assert!(verify_op(&MmaM16N8K16F32Bf16Op::new(bad_arity), &ctx).is_err()); -} - -#[test] -fn test_mma_m16n8k16_f16_verifies_exact_register_signature() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - let f32_ty = FP32Type::get(&ctx); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); - let block = BasicBlock::new( - &mut ctx, - None, - vec![f32_ty.into(), i32_ty.into(), i64_ty.into()], - ); - let f32_value = block.deref(&ctx).get_argument(0); - let i32_value = block.deref(&ctx).get_argument(1); - let i64_value = block.deref(&ctx).get_argument(2); - - let operands = || { - (0..4) - .map(|_| f32_value) - .chain((0..6).map(|_| i32_value)) - .collect() - }; - let valid = Operation::new( - &mut ctx, - MmaM16N8K16F32F16Op::get_concrete_op_info(), - vec![f32_ty.into(); 4], - operands(), - vec![], - 0, - ); - assert!(verify_op(&MmaM16N8K16F32F16Op::new(valid), &ctx).is_ok()); - - let bad_c_operands = (0..4) - .map(|index| if index == 0 { i32_value } else { f32_value }) - .chain((0..6).map(|_| i32_value)) - .collect(); - let bad_c = Operation::new( - &mut ctx, - MmaM16N8K16F32F16Op::get_concrete_op_info(), - vec![f32_ty.into(); 4], - bad_c_operands, - vec![], - 0, - ); - assert!(verify_op(&MmaM16N8K16F32F16Op::new(bad_c), &ctx).is_err()); - - let bad_packed_operands = (0..4) - .map(|_| f32_value) - .chain((0..6).map(|index| if index == 0 { i64_value } else { i32_value })) - .collect(); - let bad_packed = Operation::new( - &mut ctx, - MmaM16N8K16F32F16Op::get_concrete_op_info(), - vec![f32_ty.into(); 4], - bad_packed_operands, - vec![], - 0, - ); - assert!(verify_op(&MmaM16N8K16F32F16Op::new(bad_packed), &ctx).is_err()); - - let bad_result_type = Operation::new( - &mut ctx, - MmaM16N8K16F32F16Op::get_concrete_op_info(), - vec![f32_ty.into(), f32_ty.into(), f32_ty.into(), i32_ty.into()], - operands(), - vec![], - 0, - ); - assert!(verify_op(&MmaM16N8K16F32F16Op::new(bad_result_type), &ctx).is_err()); - - let bad_operand_arity = Operation::new( - &mut ctx, - MmaM16N8K16F32F16Op::get_concrete_op_info(), - vec![f32_ty.into(); 4], - (0..4) - .map(|_| f32_value) - .chain((0..5).map(|_| i32_value)) - .collect(), - vec![], - 0, - ); - assert!(verify_op(&MmaM16N8K16F32F16Op::new(bad_operand_arity), &ctx).is_err()); - - let bad_result_arity = Operation::new( - &mut ctx, - MmaM16N8K16F32F16Op::get_concrete_op_info(), - vec![f32_ty.into(); 3], - operands(), - vec![], - 0, - ); - assert!(verify_op(&MmaM16N8K16F32F16Op::new(bad_result_arity), &ctx).is_err()); -} - -#[test] -fn test_mma_m16n8k8_tf32_verifies_exact_register_signature() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - let f32_ty = FP32Type::get(&ctx); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); - let block = BasicBlock::new( - &mut ctx, - None, - vec![f32_ty.into(), i32_ty.into(), i64_ty.into()], - ); - let f32_value = block.deref(&ctx).get_argument(0); - let i32_value = block.deref(&ctx).get_argument(1); - let i64_value = block.deref(&ctx).get_argument(2); - - let operands = || { - (0..4) - .map(|_| f32_value) - .chain((0..6).map(|_| i32_value)) - .collect() - }; - let valid = Operation::new( - &mut ctx, - MmaM16N8K8F32Tf32Op::get_concrete_op_info(), - vec![f32_ty.into(); 4], - operands(), - vec![], - 0, - ); - assert!(verify_op(&MmaM16N8K8F32Tf32Op::new(valid), &ctx).is_ok()); - - let bad_c_operands = (0..4) - .map(|index| if index == 0 { i32_value } else { f32_value }) - .chain((0..6).map(|_| i32_value)) - .collect(); - let bad_c = Operation::new( - &mut ctx, - MmaM16N8K8F32Tf32Op::get_concrete_op_info(), - vec![f32_ty.into(); 4], - bad_c_operands, - vec![], - 0, - ); - assert!(verify_op(&MmaM16N8K8F32Tf32Op::new(bad_c), &ctx).is_err()); - - let bad_packed_operands = (0..4) - .map(|_| f32_value) - .chain((0..6).map(|index| if index == 0 { i64_value } else { i32_value })) - .collect(); - let bad_packed = Operation::new( - &mut ctx, - MmaM16N8K8F32Tf32Op::get_concrete_op_info(), - vec![f32_ty.into(); 4], - bad_packed_operands, - vec![], - 0, - ); - assert!(verify_op(&MmaM16N8K8F32Tf32Op::new(bad_packed), &ctx).is_err()); - - let bad_result_type = Operation::new( - &mut ctx, - MmaM16N8K8F32Tf32Op::get_concrete_op_info(), - vec![f32_ty.into(), f32_ty.into(), f32_ty.into(), i32_ty.into()], - operands(), - vec![], - 0, - ); - assert!(verify_op(&MmaM16N8K8F32Tf32Op::new(bad_result_type), &ctx).is_err()); - - let bad_operand_arity = Operation::new( - &mut ctx, - MmaM16N8K8F32Tf32Op::get_concrete_op_info(), - vec![f32_ty.into(); 4], - (0..4) - .map(|_| f32_value) - .chain((0..5).map(|_| i32_value)) - .collect(), - vec![], - 0, - ); - assert!(verify_op(&MmaM16N8K8F32Tf32Op::new(bad_operand_arity), &ctx).is_err()); - - let bad_result_arity = Operation::new( - &mut ctx, - MmaM16N8K8F32Tf32Op::get_concrete_op_info(), - vec![f32_ty.into(); 3], - operands(), - vec![], - 0, - ); - assert!(verify_op(&MmaM16N8K8F32Tf32Op::new(bad_result_arity), &ctx).is_err()); -} - -#[test] -fn test_mma_m16n8k32_s8_verifies_exact_register_signature() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); - let f32_ty = FP32Type::get(&ctx); - let block = BasicBlock::new( - &mut ctx, - None, - vec![i32_ty.into(), i64_ty.into(), f32_ty.into()], - ); - let i32_value = block.deref(&ctx).get_argument(0); - let i64_value = block.deref(&ctx).get_argument(1); - let f32_value = block.deref(&ctx).get_argument(2); - let valid_operands = vec![i32_value; 10]; - - let valid = Operation::new( - &mut ctx, - MmaM16N8K32S32S8Op::get_concrete_op_info(), - vec![i32_ty.into(); 4], - valid_operands.clone(), - vec![], - 0, - ); - assert!(verify_op(&MmaM16N8K32S32S8Op::new(valid), &ctx).is_ok()); - - for bad_value in [i64_value, f32_value] { - let mut bad_operands = valid_operands.clone(); - bad_operands[4] = bad_value; - let invalid = Operation::new( - &mut ctx, - MmaM16N8K32S32S8Op::get_concrete_op_info(), - vec![i32_ty.into(); 4], - bad_operands, - vec![], - 0, - ); - assert!( - verify_op(&MmaM16N8K32S32S8Op::new(invalid), &ctx).is_err(), - "MMA must reject non-i32 register operands" - ); - } - - for bad_results in [ - vec![i32_ty.into(), i32_ty.into(), i32_ty.into(), i64_ty.into()], - vec![i32_ty.into(), i32_ty.into(), i32_ty.into(), f32_ty.into()], - vec![i32_ty.into(); 3], - ] { - let invalid = Operation::new( - &mut ctx, - MmaM16N8K32S32S8Op::get_concrete_op_info(), - bad_results, - valid_operands.clone(), - vec![], - 0, - ); - assert!( - verify_op(&MmaM16N8K32S32S8Op::new(invalid), &ctx).is_err(), - "MMA must reject the wrong result register signature" - ); - } - - let invalid_arity = Operation::new( - &mut ctx, - MmaM16N8K32S32S8Op::get_concrete_op_info(), - vec![i32_ty.into(); 4], - vec![i32_value; 9], - vec![], - 0, - ); - assert!(verify_op(&MmaM16N8K32S32S8Op::new(invalid_arity), &ctx).is_err()); -} - -#[test] -fn generated_register_mma_verifier_rejects_crossed_variants_and_carriers() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - macro_rules! set_variant { - ($op:expr, $shape:expr, $acc:expr, $a:expr, $b:expr, $overflow:expr) => {{ - $op.set_attr_nvvm_register_mma_shape(&ctx, $shape); - $op.set_attr_nvvm_register_mma_accumulator(&ctx, $acc); - $op.set_attr_nvvm_register_mma_a_element(&ctx, $a); - $op.set_attr_nvvm_register_mma_b_element(&ctx, $b); - $op.set_attr_nvvm_register_mma_a_layout(&ctx, RegisterMmaLayoutAttr::Row); - $op.set_attr_nvvm_register_mma_b_layout(&ctx, RegisterMmaLayoutAttr::Col); - $op.set_attr_nvvm_register_mma_overflow(&ctx, $overflow); - }}; - } - - let f32_ty = FP32Type::get(&ctx); - let f64_ty = FP64Type::get(&ctx); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let signless_i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let block = BasicBlock::new( - &mut ctx, - None, - vec![f32_ty.into(), f64_ty.into(), i32_ty.into(), u32_ty.into()], - ); - let f32_value = block.deref(&ctx).get_argument(0); - let f64_value = block.deref(&ctx).get_argument(1); - let i32_value = block.deref(&ctx).get_argument(2); - let u32_value = block.deref(&ctx).get_argument(3); - - let bf16_operation = Operation::new( - &mut ctx, - RegisterMmaOp::get_concrete_op_info(), - vec![f32_ty.into(); 4], - [vec![f32_value; 4], vec![u32_value; 6]].concat(), - vec![], - 0, - ); - let bf16 = RegisterMmaOp::new(bf16_operation); - set_variant!( - bf16, - RegisterMmaShapeAttr::M16n8k16, - RegisterMmaAccumulatorAttr::F32, - RegisterMmaElementAttr::Bf16, - RegisterMmaElementAttr::Bf16, - RegisterMmaOverflowAttr::NotApplicable - ); - assert!(bf16.get_attr_nvvm_register_mma_operation(&ctx).is_none()); - assert!(verify_op(&bf16, &ctx).is_ok()); - bf16.set_attr_nvvm_register_mma_b_element(&ctx, RegisterMmaElementAttr::F16); - assert!(verify_op(&bf16, &ctx).is_err()); - - let f64_operation = Operation::new( - &mut ctx, - RegisterMmaOp::get_concrete_op_info(), - vec![f64_ty.into(); 2], - vec![f64_value; 4], - vec![], - 0, - ); - let f64_mma = RegisterMmaOp::new(f64_operation); - set_variant!( - f64_mma, - RegisterMmaShapeAttr::M8n8k4, - RegisterMmaAccumulatorAttr::F64, - RegisterMmaElementAttr::F64, - RegisterMmaElementAttr::F64, - RegisterMmaOverflowAttr::NotApplicable - ); - assert!(verify_op(&f64_mma, &ctx).is_ok()); - - let int_operands = [vec![i32_value; 4], vec![u32_value; 6]].concat(); - let int_operation = Operation::new( - &mut ctx, - RegisterMmaOp::get_concrete_op_info(), - vec![i32_ty.into(); 4], - int_operands.clone(), - vec![], - 0, - ); - let int_mma = RegisterMmaOp::new(int_operation); - set_variant!( - int_mma, - RegisterMmaShapeAttr::M16n8k32, - RegisterMmaAccumulatorAttr::S32, - RegisterMmaElementAttr::S8, - RegisterMmaElementAttr::S8, - RegisterMmaOverflowAttr::Wrapping - ); - assert!(verify_op(&int_mma, &ctx).is_ok()); - - let wrong_signedness = Operation::new( - &mut ctx, - RegisterMmaOp::get_concrete_op_info(), - vec![signless_i32_ty.into(); 4], - int_operands, - vec![], - 0, - ); - let wrong_signedness = RegisterMmaOp::new(wrong_signedness); - set_variant!( - wrong_signedness, - RegisterMmaShapeAttr::M16n8k32, - RegisterMmaAccumulatorAttr::S32, - RegisterMmaElementAttr::S8, - RegisterMmaElementAttr::S8, - RegisterMmaOverflowAttr::Wrapping - ); - assert!(verify_op(&wrong_signedness, &ctx).is_err()); - - let missing_attributes = Operation::new( - &mut ctx, - RegisterMmaOp::get_concrete_op_info(), - vec![f64_ty.into(); 2], - vec![f64_value; 4], - vec![], - 0, - ); - assert!(verify_op(&RegisterMmaOp::new(missing_attributes), &ctx).is_err()); -} - -#[test] -fn generated_register_mma_verifies_dense_integer_families() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let signless_i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let block = BasicBlock::new( - &mut ctx, - None, - vec![i32_ty.into(), u32_ty.into(), signless_i32_ty.into()], - ); - let i32_value = block.deref(&ctx).get_argument(0); - let u32_value = block.deref(&ctx).get_argument(1); - - macro_rules! int_mma { - ($shape:expr, $a:expr, $b:expr, $overflow:expr, $operands:expr, $results:expr) => {{ - let operation = Operation::new( - &mut ctx, - RegisterMmaOp::get_concrete_op_info(), - $results, - $operands, - vec![], - 0, - ); - let mma = RegisterMmaOp::new(operation); - mma.set_attr_nvvm_register_mma_shape(&ctx, $shape); - mma.set_attr_nvvm_register_mma_operation(&ctx, RegisterMmaOperationAttr::Multiply); - mma.set_attr_nvvm_register_mma_accumulator(&ctx, RegisterMmaAccumulatorAttr::S32); - mma.set_attr_nvvm_register_mma_a_element(&ctx, $a); - mma.set_attr_nvvm_register_mma_b_element(&ctx, $b); - mma.set_attr_nvvm_register_mma_a_layout(&ctx, RegisterMmaLayoutAttr::Row); - mma.set_attr_nvvm_register_mma_b_layout(&ctx, RegisterMmaLayoutAttr::Col); - mma.set_attr_nvvm_register_mma_overflow(&ctx, $overflow); - mma - }}; - } - - let mut accepted = 0; - for (shape, accumulator_count, operand_count, result_count) in [ - (RegisterMmaShapeAttr::M8n8k16, 2, 4, 2), - (RegisterMmaShapeAttr::M16n8k16, 4, 7, 4), - (RegisterMmaShapeAttr::M16n8k32, 4, 10, 4), - ] { - for (a_element, b_element) in [ - (RegisterMmaElementAttr::S8, RegisterMmaElementAttr::S8), - (RegisterMmaElementAttr::S8, RegisterMmaElementAttr::U8), - (RegisterMmaElementAttr::U8, RegisterMmaElementAttr::S8), - (RegisterMmaElementAttr::U8, RegisterMmaElementAttr::U8), - ] { - for overflow in [ - RegisterMmaOverflowAttr::Wrapping, - RegisterMmaOverflowAttr::Satfinite, - ] { - let operands = [ - vec![i32_value; accumulator_count], - vec![u32_value; operand_count - accumulator_count], - ] - .concat(); - let mma = int_mma!( - shape.clone(), - a_element.clone(), - b_element.clone(), - overflow.clone(), - operands, - vec![i32_ty.into(); result_count] - ); - assert_eq!( - mma.get_operation().deref(&ctx).get_num_operands(), - operand_count - ); - assert_eq!( - mma.get_operation().deref(&ctx).get_num_results(), - result_count - ); - assert!( - verify_op(&mma, &ctx).is_ok(), - "rejected {shape:?} {a_element:?}x{b_element:?} {overflow:?}" - ); - accepted += 1; - } - } - } - assert_eq!(accepted, 24); - - let mut int4_accepted = 0; - for (shape, accumulator_count, a_count, b_count, result_count) in [ - (RegisterMmaShapeAttr::M8n8k32, 2, 1, 1, 2), - (RegisterMmaShapeAttr::M16n8k32, 4, 2, 1, 4), - (RegisterMmaShapeAttr::M16n8k64, 4, 4, 2, 4), - ] { - for (a_element, b_element) in [ - (RegisterMmaElementAttr::S4, RegisterMmaElementAttr::S4), - (RegisterMmaElementAttr::S4, RegisterMmaElementAttr::U4), - (RegisterMmaElementAttr::U4, RegisterMmaElementAttr::S4), - (RegisterMmaElementAttr::U4, RegisterMmaElementAttr::U4), - ] { - for overflow in [ - RegisterMmaOverflowAttr::Wrapping, - RegisterMmaOverflowAttr::Satfinite, - ] { - let operands = [ - vec![i32_value; accumulator_count], - vec![u32_value; a_count], - vec![u32_value; b_count], - ] - .concat(); - let expected_operand_types = [ - vec![i32_ty.into(); accumulator_count], - vec![u32_ty.into(); a_count], - vec![u32_ty.into(); b_count], - ] - .concat(); - let mma = int_mma!( - shape.clone(), - a_element.clone(), - b_element.clone(), - overflow.clone(), - operands, - vec![i32_ty.into(); result_count] - ); - let operation = mma.get_operation().deref(&ctx); - assert_eq!( - operation.get_num_operands(), - accumulator_count + a_count + b_count - ); - assert_eq!(operation.get_num_results(), result_count); - assert_eq!( - operation - .operands() - .map(|operand| operand.get_type(&ctx)) - .collect::>(), - expected_operand_types - ); - assert_eq!( - (0..operation.get_num_results()) - .map(|index| operation.get_result(index).get_type(&ctx)) - .collect::>(), - vec![i32_ty.into(); result_count] - ); - assert_eq!( - mma.get_attr_nvvm_register_mma_shape(&ctx).as_deref(), - Some(&shape) - ); - assert_eq!( - mma.get_attr_nvvm_register_mma_accumulator(&ctx).as_deref(), - Some(&RegisterMmaAccumulatorAttr::S32) - ); - assert_eq!( - mma.get_attr_nvvm_register_mma_a_element(&ctx).as_deref(), - Some(&a_element) - ); - assert_eq!( - mma.get_attr_nvvm_register_mma_b_element(&ctx).as_deref(), - Some(&b_element) - ); - assert_eq!( - mma.get_attr_nvvm_register_mma_a_layout(&ctx).as_deref(), - Some(&RegisterMmaLayoutAttr::Row) - ); - assert_eq!( - mma.get_attr_nvvm_register_mma_b_layout(&ctx).as_deref(), - Some(&RegisterMmaLayoutAttr::Col) - ); - assert_eq!( - mma.get_attr_nvvm_register_mma_overflow(&ctx).as_deref(), - Some(&overflow) - ); - assert!( - verify_op(&mma, &ctx).is_ok(), - "rejected {shape:?} {a_element:?}x{b_element:?} {overflow:?}" - ); - int4_accepted += 1; - } - } - } - assert_eq!(int4_accepted, 24); - - for (shape, accumulator_count, operand_count, result_count) in [ - (RegisterMmaShapeAttr::M8n8k32, 2, 4, 2), - (RegisterMmaShapeAttr::M16n8k32, 4, 7, 4), - (RegisterMmaShapeAttr::M16n8k64, 4, 10, 4), - ] { - for wrong_operand_count in [operand_count - 1, operand_count + 1] { - let mma = int_mma!( - shape.clone(), - RegisterMmaElementAttr::S4, - RegisterMmaElementAttr::U4, - RegisterMmaOverflowAttr::Wrapping, - [ - vec![i32_value; accumulator_count], - vec![u32_value; wrong_operand_count - accumulator_count], - ] - .concat(), - vec![i32_ty.into(); result_count] - ); - assert!(verify_op(&mma, &ctx).is_err()); - } - - for wrong_result_count in [result_count - 1, result_count + 1] { - let mma = int_mma!( - shape.clone(), - RegisterMmaElementAttr::U4, - RegisterMmaElementAttr::S4, - RegisterMmaOverflowAttr::Satfinite, - [ - vec![i32_value; accumulator_count], - vec![u32_value; operand_count - accumulator_count], - ] - .concat(), - vec![i32_ty.into(); wrong_result_count] - ); - assert!(verify_op(&mma, &ctx).is_err()); - } - } - - let int4_on_int8_shape = int_mma!( - RegisterMmaShapeAttr::M8n8k16, - RegisterMmaElementAttr::S4, - RegisterMmaElementAttr::U4, - RegisterMmaOverflowAttr::Wrapping, - [vec![i32_value; 2], vec![u32_value; 2]].concat(), - vec![i32_ty.into(); 2] - ); - assert!(verify_op(&int4_on_int8_shape, &ctx).is_err()); - - let int8_on_int4_shape = int_mma!( - RegisterMmaShapeAttr::M8n8k32, - RegisterMmaElementAttr::S8, - RegisterMmaElementAttr::U8, - RegisterMmaOverflowAttr::Wrapping, - [vec![i32_value; 2], vec![u32_value; 2]].concat(), - vec![i32_ty.into(); 2] - ); - assert!(verify_op(&int8_on_int4_shape, &ctx).is_err()); - - let crossed_integer_width = int_mma!( - RegisterMmaShapeAttr::M16n8k32, - RegisterMmaElementAttr::S4, - RegisterMmaElementAttr::U8, - RegisterMmaOverflowAttr::Satfinite, - [vec![i32_value; 4], vec![u32_value; 3]].concat(), - vec![i32_ty.into(); 4] - ); - assert!(verify_op(&crossed_integer_width, &ctx).is_err()); - - let m16k32_int4_with_int8_carriers = int_mma!( - RegisterMmaShapeAttr::M16n8k32, - RegisterMmaElementAttr::U4, - RegisterMmaElementAttr::S4, - RegisterMmaOverflowAttr::Wrapping, - [vec![i32_value; 4], vec![u32_value; 6]].concat(), - vec![i32_ty.into(); 4] - ); - assert!(verify_op(&m16k32_int4_with_int8_carriers, &ctx).is_err()); - - let m16k32_int8_with_int4_carriers = int_mma!( - RegisterMmaShapeAttr::M16n8k32, - RegisterMmaElementAttr::S8, - RegisterMmaElementAttr::U8, - RegisterMmaOverflowAttr::Satfinite, - [vec![i32_value; 4], vec![u32_value; 3]].concat(), - vec![i32_ty.into(); 4] - ); - assert!(verify_op(&m16k32_int8_with_int4_carriers, &ctx).is_err()); - - let m16k64_int4_with_k32_carriers = int_mma!( - RegisterMmaShapeAttr::M16n8k64, - RegisterMmaElementAttr::S4, - RegisterMmaElementAttr::U4, - RegisterMmaOverflowAttr::Wrapping, - [vec![i32_value; 4], vec![u32_value; 3]].concat(), - vec![i32_ty.into(); 4] - ); - assert!(verify_op(&m16k64_int4_with_k32_carriers, &ctx).is_err()); - - let m16k32_int4_with_k64_carriers = int_mma!( - RegisterMmaShapeAttr::M16n8k32, - RegisterMmaElementAttr::U4, - RegisterMmaElementAttr::S4, - RegisterMmaOverflowAttr::Satfinite, - [vec![i32_value; 4], vec![u32_value; 6]].concat(), - vec![i32_ty.into(); 4] - ); - assert!(verify_op(&m16k32_int4_with_k64_carriers, &ctx).is_err()); - - let int8_on_m16k64_shape = int_mma!( - RegisterMmaShapeAttr::M16n8k64, - RegisterMmaElementAttr::S8, - RegisterMmaElementAttr::U8, - RegisterMmaOverflowAttr::Wrapping, - [vec![i32_value; 4], vec![u32_value; 6]].concat(), - vec![i32_ty.into(); 4] - ); - assert!(verify_op(&int8_on_m16k64_shape, &ctx).is_err()); - - for (shape, accumulator_count, operand_count, result_count) in [ - (RegisterMmaShapeAttr::M8n8k16, 2, 4, 2), - (RegisterMmaShapeAttr::M16n8k16, 4, 7, 4), - (RegisterMmaShapeAttr::M16n8k32, 4, 10, 4), - ] { - for wrong_count in [operand_count - 1, operand_count + 1] { - let operands = [ - vec![i32_value; accumulator_count], - vec![u32_value; wrong_count - accumulator_count], - ] - .concat(); - let mma = int_mma!( - shape.clone(), - RegisterMmaElementAttr::S8, - RegisterMmaElementAttr::U8, - RegisterMmaOverflowAttr::Wrapping, - operands, - vec![i32_ty.into(); result_count] - ); - assert!(verify_op(&mma, &ctx).is_err()); - } - } - - let m8_wrong_result_signedness = int_mma!( - RegisterMmaShapeAttr::M8n8k16, - RegisterMmaElementAttr::S8, - RegisterMmaElementAttr::U8, - RegisterMmaOverflowAttr::Wrapping, - [vec![i32_value; 2], vec![u32_value; 2]].concat(), - vec![signless_i32_ty.into(); 2] - ); - assert!(verify_op(&m8_wrong_result_signedness, &ctx).is_err()); - - let m8_wrong_accumulator_signedness = int_mma!( - RegisterMmaShapeAttr::M8n8k16, - RegisterMmaElementAttr::S8, - RegisterMmaElementAttr::U8, - RegisterMmaOverflowAttr::Wrapping, - vec![u32_value; 4], - vec![i32_ty.into(); 2] - ); - assert!(verify_op(&m8_wrong_accumulator_signedness, &ctx).is_err()); - - let m8_wrong_fragment_signedness = int_mma!( - RegisterMmaShapeAttr::M8n8k16, - RegisterMmaElementAttr::U8, - RegisterMmaElementAttr::S8, - RegisterMmaOverflowAttr::Satfinite, - vec![i32_value; 4], - vec![i32_ty.into(); 2] - ); - assert!(verify_op(&m8_wrong_fragment_signedness, &ctx).is_err()); - - let m8_crossed_element = int_mma!( - RegisterMmaShapeAttr::M8n8k16, - RegisterMmaElementAttr::S8, - RegisterMmaElementAttr::F16, - RegisterMmaOverflowAttr::Wrapping, - [vec![i32_value; 2], vec![u32_value; 2]].concat(), - vec![i32_ty.into(); 2] - ); - assert!(verify_op(&m8_crossed_element, &ctx).is_err()); - - let m8_crossed_overflow = int_mma!( - RegisterMmaShapeAttr::M8n8k16, - RegisterMmaElementAttr::U8, - RegisterMmaElementAttr::U8, - RegisterMmaOverflowAttr::NotApplicable, - [vec![i32_value; 2], vec![u32_value; 2]].concat(), - vec![i32_ty.into(); 2] - ); - assert!(verify_op(&m8_crossed_overflow, &ctx).is_err()); - - let m8_crossed_carrier_shape = int_mma!( - RegisterMmaShapeAttr::M8n8k16, - RegisterMmaElementAttr::S8, - RegisterMmaElementAttr::S8, - RegisterMmaOverflowAttr::Wrapping, - [vec![i32_value; 4], vec![u32_value; 3]].concat(), - vec![i32_ty.into(); 4] - ); - assert!(verify_op(&m8_crossed_carrier_shape, &ctx).is_err()); - - let m8_crossed_shape = int_mma!( - RegisterMmaShapeAttr::M8n8k4, - RegisterMmaElementAttr::S8, - RegisterMmaElementAttr::S8, - RegisterMmaOverflowAttr::Wrapping, - [vec![i32_value; 2], vec![u32_value; 2]].concat(), - vec![i32_ty.into(); 2] - ); - assert!(verify_op(&m8_crossed_shape, &ctx).is_err()); - - let wrong_result_signedness = int_mma!( - RegisterMmaShapeAttr::M16n8k16, - RegisterMmaElementAttr::S8, - RegisterMmaElementAttr::U8, - RegisterMmaOverflowAttr::Wrapping, - [vec![i32_value; 4], vec![u32_value; 3]].concat(), - vec![signless_i32_ty.into(); 4] - ); - assert!(verify_op(&wrong_result_signedness, &ctx).is_err()); - - let wrong_accumulator_signedness = int_mma!( - RegisterMmaShapeAttr::M16n8k16, - RegisterMmaElementAttr::S8, - RegisterMmaElementAttr::U8, - RegisterMmaOverflowAttr::Wrapping, - vec![u32_value; 7], - vec![i32_ty.into(); 4] - ); - assert!(verify_op(&wrong_accumulator_signedness, &ctx).is_err()); - - let wrong_fragment_signedness = int_mma!( - RegisterMmaShapeAttr::M16n8k32, - RegisterMmaElementAttr::U8, - RegisterMmaElementAttr::S8, - RegisterMmaOverflowAttr::Satfinite, - vec![i32_value; 10], - vec![i32_ty.into(); 4] - ); - assert!(verify_op(&wrong_fragment_signedness, &ctx).is_err()); - - let crossed_element = int_mma!( - RegisterMmaShapeAttr::M16n8k16, - RegisterMmaElementAttr::S8, - RegisterMmaElementAttr::F16, - RegisterMmaOverflowAttr::Wrapping, - [vec![i32_value; 4], vec![u32_value; 3]].concat(), - vec![i32_ty.into(); 4] - ); - assert!(verify_op(&crossed_element, &ctx).is_err()); - - let crossed_overflow = int_mma!( - RegisterMmaShapeAttr::M16n8k32, - RegisterMmaElementAttr::U8, - RegisterMmaElementAttr::U8, - RegisterMmaOverflowAttr::NotApplicable, - [vec![i32_value; 4], vec![u32_value; 6]].concat(), - vec![i32_ty.into(); 4] - ); - assert!(verify_op(&crossed_overflow, &ctx).is_err()); - - let crossed_shape = int_mma!( - RegisterMmaShapeAttr::M16n8k8, - RegisterMmaElementAttr::S8, - RegisterMmaElementAttr::S8, - RegisterMmaOverflowAttr::Satfinite, - [vec![i32_value; 4], vec![u32_value; 3]].concat(), - vec![i32_ty.into(); 4] - ); - assert!(verify_op(&crossed_shape, &ctx).is_err()); -} - -#[test] -fn generated_register_mma_verifies_dense_b1_families() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - fn b1_mma( - ctx: &mut Context, - shape: RegisterMmaShapeAttr, - operation: Option, - ) -> RegisterMmaOp { - let (accumulator_count, a_count, b_count, result_count) = match shape { - RegisterMmaShapeAttr::M8n8k128 => (2, 1, 1, 2), - RegisterMmaShapeAttr::M16n8k128 => (4, 2, 1, 4), - RegisterMmaShapeAttr::M16n8k256 => (4, 4, 2, 4), - _ => panic!("unsupported B1 MMA shape"), - }; - let i32_ty = IntegerType::get(ctx, 32, Signedness::Signed); - let u32_ty = IntegerType::get(ctx, 32, Signedness::Unsigned); - let argument_types = (0..accumulator_count) - .map(|_| i32_ty.into()) - .chain((0..a_count + b_count).map(|_| u32_ty.into())) - .collect(); - let block = BasicBlock::new(ctx, None, argument_types); - let operands = (0..accumulator_count + a_count + b_count) - .map(|index| block.deref(ctx).get_argument(index)) - .collect(); - let op = Operation::new( - ctx, - RegisterMmaOp::get_concrete_op_info(), - vec![i32_ty.into(); result_count], - operands, - vec![], - 0, - ); - let mma = RegisterMmaOp::new(op); - mma.set_attr_nvvm_register_mma_shape(ctx, shape); - if let Some(operation) = operation { - mma.set_attr_nvvm_register_mma_operation(ctx, operation); - } - mma.set_attr_nvvm_register_mma_accumulator(ctx, RegisterMmaAccumulatorAttr::S32); - mma.set_attr_nvvm_register_mma_a_element(ctx, RegisterMmaElementAttr::B1); - mma.set_attr_nvvm_register_mma_b_element(ctx, RegisterMmaElementAttr::B1); - mma.set_attr_nvvm_register_mma_a_layout(ctx, RegisterMmaLayoutAttr::Row); - mma.set_attr_nvvm_register_mma_b_layout(ctx, RegisterMmaLayoutAttr::Col); - mma.set_attr_nvvm_register_mma_overflow(ctx, RegisterMmaOverflowAttr::Wrapping); - mma - } - - let mut accepted = 0; - for shape in [ - RegisterMmaShapeAttr::M8n8k128, - RegisterMmaShapeAttr::M16n8k128, - RegisterMmaShapeAttr::M16n8k256, - ] { - for operation in [ - RegisterMmaOperationAttr::XorPopc, - RegisterMmaOperationAttr::AndPopc, - ] { - let mma = b1_mma(&mut ctx, shape.clone(), Some(operation)); - assert!(verify_op(&mma, &ctx).is_ok(), "rejected {shape:?}"); - accepted += 1; - } - } - assert_eq!(accepted, 6); - - let multiply = b1_mma( - &mut ctx, - RegisterMmaShapeAttr::M8n8k128, - Some(RegisterMmaOperationAttr::Multiply), - ); - assert!(verify_op(&multiply, &ctx).is_err()); - - let wrong_shape = b1_mma( - &mut ctx, - RegisterMmaShapeAttr::M16n8k128, - Some(RegisterMmaOperationAttr::XorPopc), - ); - wrong_shape.set_attr_nvvm_register_mma_shape(&ctx, RegisterMmaShapeAttr::M16n8k64); - assert!(verify_op(&wrong_shape, &ctx).is_err()); - - let missing_operation = b1_mma(&mut ctx, RegisterMmaShapeAttr::M8n8k128, None); - assert!(verify_op(&missing_operation, &ctx).is_err()); -} - -#[test] -fn generated_sparse_mma_verifies_all_int8_variants_and_metadata_modes() { - use dialect_mir::ops::MirConstantOp; - use pliron::builtin::{attributes::IntegerAttr, ops::ConstantOp}; - use pliron::utils::apint::APInt; - use std::num::NonZeroUsize; - - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let signless_i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let block = BasicBlock::new( - &mut ctx, - None, - vec![i32_ty.into(), u32_ty.into(), signless_i32_ty.into()], - ); - let i32_value = block.deref(&ctx).get_argument(0); - let u32_value = block.deref(&ctx).get_argument(1); - let signless_i32_value = block.deref(&ctx).get_argument(2); - - let integer = |value| { - IntegerAttr::new( - u32_ty, - APInt::from_u32(value, NonZeroUsize::new(32).unwrap()), - ) - }; - let builtin_zero = ConstantOp::new(&mut ctx, integer(0).into()); - let builtin_zero = builtin_zero.get_operation().deref(&ctx).get_result(0); - let builtin_two = ConstantOp::new(&mut ctx, integer(2).into()); - let builtin_two = builtin_two.get_operation().deref(&ctx).get_result(0); - let mir_one = Operation::new( - &mut ctx, - MirConstantOp::get_concrete_op_info(), - vec![u32_ty.into()], - vec![], - vec![], - 0, - ); - MirConstantOp::new(mir_one).set_attr_value(&ctx, integer(1)); - let mir_one = mir_one.deref(&ctx).get_result(0); - - macro_rules! sparse_mma { - ($operands:expr, $results:expr, $a:expr, $b:expr, $overflow:expr, $metadata:expr) => {{ - let operation = Operation::new( - &mut ctx, - SparseMmaOp::get_concrete_op_info(), - $results, - $operands, - vec![], - 0, - ); - let mma = SparseMmaOp::new(operation); - mma.set_attr_nvvm_sparse_mma_shape(&ctx, SparseMmaShapeAttr::M16n8k32); - mma.set_attr_nvvm_sparse_mma_accumulator(&ctx, SparseMmaAccumulatorAttr::S32); - mma.set_attr_nvvm_sparse_mma_a_element(&ctx, $a); - mma.set_attr_nvvm_sparse_mma_b_element(&ctx, $b); - mma.set_attr_nvvm_sparse_mma_a_layout(&ctx, SparseMmaLayoutAttr::Row); - mma.set_attr_nvvm_sparse_mma_b_layout(&ctx, SparseMmaLayoutAttr::Col); - mma.set_attr_nvvm_sparse_mma_overflow(&ctx, $overflow); - mma.set_attr_nvvm_sparse_mma_metadata(&ctx, $metadata); - mma.set_attr_nvvm_sparse_mma_selector(&ctx, SparseMmaSelectorAttr::ImmediateZeroOrOne); - mma - }}; - } - - let variants = [ - ( - SparseMmaElementAttr::S8, - SparseMmaElementAttr::S8, - SparseMmaOverflowAttr::Wrapping, - ), - ( - SparseMmaElementAttr::S8, - SparseMmaElementAttr::U8, - SparseMmaOverflowAttr::Wrapping, - ), - ( - SparseMmaElementAttr::U8, - SparseMmaElementAttr::U8, - SparseMmaOverflowAttr::Wrapping, - ), - ( - SparseMmaElementAttr::U8, - SparseMmaElementAttr::S8, - SparseMmaOverflowAttr::Wrapping, - ), - ( - SparseMmaElementAttr::S8, - SparseMmaElementAttr::S8, - SparseMmaOverflowAttr::Satfinite, - ), - ( - SparseMmaElementAttr::S8, - SparseMmaElementAttr::U8, - SparseMmaOverflowAttr::Satfinite, - ), - ( - SparseMmaElementAttr::U8, - SparseMmaElementAttr::U8, - SparseMmaOverflowAttr::Satfinite, - ), - ( - SparseMmaElementAttr::U8, - SparseMmaElementAttr::S8, - SparseMmaOverflowAttr::Satfinite, - ), - ]; - for metadata in [ - SparseMmaMetadataAttr::Standard, - SparseMmaMetadataAttr::Ordered, - ] { - for (index, (a_element, b_element, overflow)) in variants.iter().enumerate() { - let selector = if index % 2 == 0 { - builtin_zero - } else { - mir_one - }; - let operands = [vec![i32_value; 4], vec![u32_value; 5], vec![selector]].concat(); - let mma = sparse_mma!( - operands, - vec![i32_ty.into(); 4], - a_element.clone(), - b_element.clone(), - overflow.clone(), - metadata.clone() - ); - assert_eq!( - mma.get_attr_nvvm_sparse_mma_a_element(&ctx).as_deref(), - Some(a_element) - ); - assert_eq!( - mma.get_attr_nvvm_sparse_mma_b_element(&ctx).as_deref(), - Some(b_element) - ); - assert_eq!( - mma.get_attr_nvvm_sparse_mma_overflow(&ctx).as_deref(), - Some(overflow) - ); - assert_eq!( - mma.get_attr_nvvm_sparse_mma_metadata(&ctx).as_deref(), - Some(&metadata) - ); - assert!( - verify_op(&mma, &ctx).is_ok(), - "rejected sparse {metadata:?} {a_element:?}x{b_element:?} {overflow:?}" - ); - } - } - - for metadata in [ - SparseMmaMetadataAttr::Standard, - SparseMmaMetadataAttr::Ordered, - ] { - for selector in [u32_value, builtin_two] { - let invalid = sparse_mma!( - [vec![i32_value; 4], vec![u32_value; 5], vec![selector],].concat(), - vec![i32_ty.into(); 4], - SparseMmaElementAttr::S8, - SparseMmaElementAttr::U8, - SparseMmaOverflowAttr::Wrapping, - metadata.clone() - ); - assert!(verify_op(&invalid, &ctx).is_err()); - } - } - - let wrong_accumulator_type = sparse_mma!( - [ - vec![signless_i32_value; 4], - vec![u32_value; 5], - vec![builtin_zero], - ] - .concat(), - vec![i32_ty.into(); 4], - SparseMmaElementAttr::U8, - SparseMmaElementAttr::S8, - SparseMmaOverflowAttr::Satfinite, - SparseMmaMetadataAttr::Standard - ); - assert!(verify_op(&wrong_accumulator_type, &ctx).is_err()); - - let wrong_count = sparse_mma!( - [vec![i32_value; 4], vec![u32_value; 4], vec![builtin_zero],].concat(), - vec![i32_ty.into(); 4], - SparseMmaElementAttr::S8, - SparseMmaElementAttr::S8, - SparseMmaOverflowAttr::Wrapping, - SparseMmaMetadataAttr::Standard - ); - assert!(verify_op(&wrong_count, &ctx).is_err()); - - let wrong_results = sparse_mma!( - [vec![i32_value; 4], vec![u32_value; 5], vec![mir_one],].concat(), - vec![signless_i32_ty.into(); 4], - SparseMmaElementAttr::U8, - SparseMmaElementAttr::U8, - SparseMmaOverflowAttr::Wrapping, - SparseMmaMetadataAttr::Standard - ); - assert!(verify_op(&wrong_results, &ctx).is_err()); - - let wrong_layout = sparse_mma!( - [vec![i32_value; 4], vec![u32_value; 5], vec![builtin_zero],].concat(), - vec![i32_ty.into(); 4], - SparseMmaElementAttr::S8, - SparseMmaElementAttr::U8, - SparseMmaOverflowAttr::Satfinite, - SparseMmaMetadataAttr::Standard - ); - wrong_layout.set_attr_nvvm_sparse_mma_b_layout(&ctx, SparseMmaLayoutAttr::Row); - assert!(verify_op(&wrong_layout, &ctx).is_err()); -} - -#[test] -fn generated_sparse_mma_m16n8k64_verifies_selector_and_carriers() { - use pliron::builtin::{attributes::IntegerAttr, ops::ConstantOp}; - use pliron::utils::apint::APInt; - use std::num::NonZeroUsize; - - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let signless_i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let block = BasicBlock::new( - &mut ctx, - None, - vec![i32_ty.into(), u32_ty.into(), signless_i32_ty.into()], - ); - let i32_value = block.deref(&ctx).get_argument(0); - let u32_value = block.deref(&ctx).get_argument(1); - let signless_i32_value = block.deref(&ctx).get_argument(2); - let integer = |value| { - IntegerAttr::new( - u32_ty, - APInt::from_u32(value, NonZeroUsize::new(32).unwrap()), - ) - }; - let zero = ConstantOp::new(&mut ctx, integer(0).into()) - .get_operation() - .deref(&ctx) - .get_result(0); - let one = ConstantOp::new(&mut ctx, integer(1).into()) - .get_operation() - .deref(&ctx) - .get_result(0); - let two = ConstantOp::new(&mut ctx, integer(2).into()) - .get_operation() - .deref(&ctx) - .get_result(0); - - macro_rules! k64_mma { - ($operands:expr, $metadata:expr) => {{ - let operation = Operation::new( - &mut ctx, - SparseMmaOp::get_concrete_op_info(), - vec![i32_ty.into(); 4], - $operands, - vec![], - 0, - ); - let mma = SparseMmaOp::new(operation); - mma.set_attr_nvvm_sparse_mma_shape(&ctx, SparseMmaShapeAttr::M16n8k64); - mma.set_attr_nvvm_sparse_mma_accumulator(&ctx, SparseMmaAccumulatorAttr::S32); - mma.set_attr_nvvm_sparse_mma_a_element(&ctx, SparseMmaElementAttr::S8); - mma.set_attr_nvvm_sparse_mma_b_element(&ctx, SparseMmaElementAttr::U8); - mma.set_attr_nvvm_sparse_mma_a_layout(&ctx, SparseMmaLayoutAttr::Row); - mma.set_attr_nvvm_sparse_mma_b_layout(&ctx, SparseMmaLayoutAttr::Col); - mma.set_attr_nvvm_sparse_mma_overflow(&ctx, SparseMmaOverflowAttr::Wrapping); - mma.set_attr_nvvm_sparse_mma_metadata(&ctx, $metadata); - mma.set_attr_nvvm_sparse_mma_selector(&ctx, SparseMmaSelectorAttr::ImmediateZero); - mma - }}; - } - - let operands = |selector| [vec![i32_value; 4], vec![u32_value; 9], vec![selector]].concat(); - assert!( - verify_op( - &k64_mma!(operands(zero), SparseMmaMetadataAttr::Standard), - &ctx - ) - .is_ok() - ); - assert!( - verify_op( - &k64_mma!(operands(zero), SparseMmaMetadataAttr::Ordered), - &ctx - ) - .is_ok() - ); - assert!( - verify_op( - &k64_mma!(operands(one), SparseMmaMetadataAttr::Standard), - &ctx - ) - .is_err() - ); - assert!( - verify_op( - &k64_mma!(operands(one), SparseMmaMetadataAttr::Ordered), - &ctx - ) - .is_err() - ); - assert!( - verify_op( - &k64_mma!(operands(u32_value), SparseMmaMetadataAttr::Standard), - &ctx - ) - .is_err() - ); - assert!( - verify_op( - &k64_mma!(operands(u32_value), SparseMmaMetadataAttr::Ordered), - &ctx - ) - .is_err() - ); - - let wrong_count = [vec![i32_value; 4], vec![u32_value; 8], vec![zero]].concat(); - assert!( - verify_op( - &k64_mma!(wrong_count, SparseMmaMetadataAttr::Standard), - &ctx - ) - .is_err() - ); - - let mut wrong_type = operands(zero); - wrong_type[4] = signless_i32_value; - assert!(verify_op(&k64_mma!(wrong_type, SparseMmaMetadataAttr::Standard), &ctx).is_err()); - - macro_rules! int4_mma { - ($operands:expr, $a:expr, $b:expr, $overflow:expr, $metadata:expr) => {{ - let operation = Operation::new( - &mut ctx, - SparseMmaOp::get_concrete_op_info(), - vec![i32_ty.into(); 4], - $operands, - vec![], - 0, - ); - let mma = SparseMmaOp::new(operation); - mma.set_attr_nvvm_sparse_mma_shape(&ctx, SparseMmaShapeAttr::M16n8k64); - mma.set_attr_nvvm_sparse_mma_accumulator(&ctx, SparseMmaAccumulatorAttr::S32); - mma.set_attr_nvvm_sparse_mma_a_element(&ctx, $a); - mma.set_attr_nvvm_sparse_mma_b_element(&ctx, $b); - mma.set_attr_nvvm_sparse_mma_a_layout(&ctx, SparseMmaLayoutAttr::Row); - mma.set_attr_nvvm_sparse_mma_b_layout(&ctx, SparseMmaLayoutAttr::Col); - mma.set_attr_nvvm_sparse_mma_overflow(&ctx, $overflow); - mma.set_attr_nvvm_sparse_mma_metadata(&ctx, $metadata); - mma.set_attr_nvvm_sparse_mma_selector(&ctx, SparseMmaSelectorAttr::ImmediateZeroOrOne); - mma - }}; - } - - let int4_variants = [ - ( - SparseMmaElementAttr::S4, - SparseMmaElementAttr::S4, - SparseMmaOverflowAttr::Wrapping, - ), - ( - SparseMmaElementAttr::S4, - SparseMmaElementAttr::U4, - SparseMmaOverflowAttr::Wrapping, - ), - ( - SparseMmaElementAttr::U4, - SparseMmaElementAttr::U4, - SparseMmaOverflowAttr::Wrapping, - ), - ( - SparseMmaElementAttr::U4, - SparseMmaElementAttr::S4, - SparseMmaOverflowAttr::Wrapping, - ), - ( - SparseMmaElementAttr::S4, - SparseMmaElementAttr::S4, - SparseMmaOverflowAttr::Satfinite, - ), - ( - SparseMmaElementAttr::S4, - SparseMmaElementAttr::U4, - SparseMmaOverflowAttr::Satfinite, - ), - ( - SparseMmaElementAttr::U4, - SparseMmaElementAttr::U4, - SparseMmaOverflowAttr::Satfinite, - ), - ( - SparseMmaElementAttr::U4, - SparseMmaElementAttr::S4, - SparseMmaOverflowAttr::Satfinite, - ), - ]; - let int4_operands = - |selector| [vec![i32_value; 4], vec![u32_value; 5], vec![selector]].concat(); - for metadata in [ - SparseMmaMetadataAttr::Standard, - SparseMmaMetadataAttr::Ordered, - ] { - for (index, (a, b, overflow)) in int4_variants.iter().enumerate() { - let selector = if index % 2 == 0 { zero } else { one }; - assert!( - verify_op( - &int4_mma!( - int4_operands(selector), - a.clone(), - b.clone(), - overflow.clone(), - metadata.clone() - ), - &ctx, - ) - .is_ok() - ); - } - } - - assert!( - verify_op( - &int4_mma!( - int4_operands(zero), - SparseMmaElementAttr::S4, - SparseMmaElementAttr::U8, - SparseMmaOverflowAttr::Wrapping, - SparseMmaMetadataAttr::Standard - ), - &ctx, - ) - .is_err() - ); - for metadata in [ - SparseMmaMetadataAttr::Standard, - SparseMmaMetadataAttr::Ordered, - ] { - for selector in [two, u32_value] { - assert!( - verify_op( - &int4_mma!( - int4_operands(selector), - SparseMmaElementAttr::S4, - SparseMmaElementAttr::U4, - SparseMmaOverflowAttr::Wrapping, - metadata.clone() - ), - &ctx, - ) - .is_err() - ); - } - } -} - -#[test] -fn generated_sparse_mma_m16n8k128_int4_verifies_metadata_selector_and_widths() { - use pliron::builtin::{attributes::IntegerAttr, ops::ConstantOp}; - use pliron::utils::apint::APInt; - use std::num::NonZeroUsize; - - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into(), u32_ty.into()]); - let i32_value = block.deref(&ctx).get_argument(0); - let u32_value = block.deref(&ctx).get_argument(1); - let integer = |value| { - IntegerAttr::new( - u32_ty, - APInt::from_u32(value, NonZeroUsize::new(32).unwrap()), - ) - }; - let zero = ConstantOp::new(&mut ctx, integer(0).into()) - .get_operation() - .deref(&ctx) - .get_result(0); - let one = ConstantOp::new(&mut ctx, integer(1).into()) - .get_operation() - .deref(&ctx) - .get_result(0); - - macro_rules! k128_mma { - ($operands:expr, $a:expr, $b:expr, $overflow:expr, $metadata:expr) => {{ - let operation = Operation::new( - &mut ctx, - SparseMmaOp::get_concrete_op_info(), - vec![i32_ty.into(); 4], - $operands, - vec![], - 0, - ); - let mma = SparseMmaOp::new(operation); - mma.set_attr_nvvm_sparse_mma_shape(&ctx, SparseMmaShapeAttr::M16n8k128); - mma.set_attr_nvvm_sparse_mma_accumulator(&ctx, SparseMmaAccumulatorAttr::S32); - mma.set_attr_nvvm_sparse_mma_a_element(&ctx, $a); - mma.set_attr_nvvm_sparse_mma_b_element(&ctx, $b); - mma.set_attr_nvvm_sparse_mma_a_layout(&ctx, SparseMmaLayoutAttr::Row); - mma.set_attr_nvvm_sparse_mma_b_layout(&ctx, SparseMmaLayoutAttr::Col); - mma.set_attr_nvvm_sparse_mma_overflow(&ctx, $overflow); - mma.set_attr_nvvm_sparse_mma_metadata(&ctx, $metadata); - mma.set_attr_nvvm_sparse_mma_selector(&ctx, SparseMmaSelectorAttr::ImmediateZero); - mma - }}; - } - - let variants = [ - ( - SparseMmaElementAttr::S4, - SparseMmaElementAttr::S4, - SparseMmaOverflowAttr::Wrapping, - ), - ( - SparseMmaElementAttr::S4, - SparseMmaElementAttr::U4, - SparseMmaOverflowAttr::Wrapping, - ), - ( - SparseMmaElementAttr::U4, - SparseMmaElementAttr::U4, - SparseMmaOverflowAttr::Wrapping, - ), - ( - SparseMmaElementAttr::U4, - SparseMmaElementAttr::S4, - SparseMmaOverflowAttr::Wrapping, - ), - ( - SparseMmaElementAttr::S4, - SparseMmaElementAttr::S4, - SparseMmaOverflowAttr::Satfinite, - ), - ( - SparseMmaElementAttr::S4, - SparseMmaElementAttr::U4, - SparseMmaOverflowAttr::Satfinite, - ), - ( - SparseMmaElementAttr::U4, - SparseMmaElementAttr::U4, - SparseMmaOverflowAttr::Satfinite, - ), - ( - SparseMmaElementAttr::U4, - SparseMmaElementAttr::S4, - SparseMmaOverflowAttr::Satfinite, - ), - ]; - let operands = |selector| [vec![i32_value; 4], vec![u32_value; 9], vec![selector]].concat(); - for metadata in [ - SparseMmaMetadataAttr::Standard, - SparseMmaMetadataAttr::Ordered, - ] { - for (a, b, overflow) in &variants { - assert!( - verify_op( - &k128_mma!( - operands(zero), - a.clone(), - b.clone(), - overflow.clone(), - metadata.clone() - ), - &ctx, - ) - .is_ok() - ); - } - } - - for metadata in [ - SparseMmaMetadataAttr::Standard, - SparseMmaMetadataAttr::Ordered, - ] { - for selector in [one, u32_value] { - assert!( - verify_op( - &k128_mma!( - operands(selector), - SparseMmaElementAttr::S4, - SparseMmaElementAttr::U4, - SparseMmaOverflowAttr::Wrapping, - metadata.clone() - ), - &ctx, - ) - .is_err() - ); - } - } - assert!( - verify_op( - &k128_mma!( - operands(zero), - SparseMmaElementAttr::S4, - SparseMmaElementAttr::U8, - SparseMmaOverflowAttr::Wrapping, - SparseMmaMetadataAttr::Standard - ), - &ctx, - ) - .is_err() - ); - assert!( - verify_op( - &k128_mma!( - [vec![i32_value; 4], vec![u32_value; 8], vec![zero]].concat(), - SparseMmaElementAttr::S4, - SparseMmaElementAttr::S4, - SparseMmaOverflowAttr::Wrapping, - SparseMmaMetadataAttr::Standard - ), - &ctx, - ) - .is_err() - ); -} - -#[test] -fn generated_sparse_mma_f8f6f4_verifies_all_formats_and_closed_shape() { - use pliron::builtin::{attributes::IntegerAttr, ops::ConstantOp}; - use pliron::utils::apint::APInt; - use std::num::NonZeroUsize; - - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - - let f32_ty = FP32Type::get(&ctx); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let block = BasicBlock::new(&mut ctx, None, vec![f32_ty.into(), u32_ty.into()]); - let f32_value = block.deref(&ctx).get_argument(0); - let u32_value = block.deref(&ctx).get_argument(1); - let integer = |value| { - IntegerAttr::new( - u32_ty, - APInt::from_u32(value, NonZeroUsize::new(32).unwrap()), - ) - }; - let zero = ConstantOp::new(&mut ctx, integer(0).into()) - .get_operation() - .deref(&ctx) - .get_result(0); - let one = ConstantOp::new(&mut ctx, integer(1).into()) - .get_operation() - .deref(&ctx) - .get_result(0); - - macro_rules! f8f6f4_mma { - ($operands:expr, $results:expr, $a:expr, $b:expr, $overflow:expr, $metadata:expr) => {{ - let operation = Operation::new( - &mut ctx, - SparseMmaOp::get_concrete_op_info(), - $results, - $operands, - vec![], - 0, - ); - let mma = SparseMmaOp::new(operation); - mma.set_attr_nvvm_sparse_mma_shape(&ctx, SparseMmaShapeAttr::M16n8k64); - mma.set_attr_nvvm_sparse_mma_accumulator(&ctx, SparseMmaAccumulatorAttr::F32); - mma.set_attr_nvvm_sparse_mma_a_element(&ctx, $a); - mma.set_attr_nvvm_sparse_mma_b_element(&ctx, $b); - mma.set_attr_nvvm_sparse_mma_a_layout(&ctx, SparseMmaLayoutAttr::Row); - mma.set_attr_nvvm_sparse_mma_b_layout(&ctx, SparseMmaLayoutAttr::Col); - mma.set_attr_nvvm_sparse_mma_overflow(&ctx, $overflow); - mma.set_attr_nvvm_sparse_mma_metadata(&ctx, $metadata); - mma.set_attr_nvvm_sparse_mma_selector(&ctx, SparseMmaSelectorAttr::ImmediateZero); - mma - }}; - } - - let elements = [ - SparseMmaElementAttr::E2m1, - SparseMmaElementAttr::E2m3, - SparseMmaElementAttr::E3m2, - SparseMmaElementAttr::E4m3, - SparseMmaElementAttr::E5m2, - ]; - let operands = |selector| [vec![f32_value; 4], vec![u32_value; 9], vec![selector]].concat(); - for a in &elements { - for b in &elements { - let mma = f8f6f4_mma!( - operands(zero), - vec![f32_ty.into(); 4], - a.clone(), - b.clone(), - SparseMmaOverflowAttr::NotApplicable, - SparseMmaMetadataAttr::Ordered - ); - assert!(verify_op(&mma, &ctx).is_ok(), "rejected {a:?}x{b:?}"); - } - } - - for invalid in [ - f8f6f4_mma!( - operands(one), - vec![f32_ty.into(); 4], - SparseMmaElementAttr::E2m1, - SparseMmaElementAttr::E2m1, - SparseMmaOverflowAttr::NotApplicable, - SparseMmaMetadataAttr::Ordered - ), - f8f6f4_mma!( - operands(zero), - vec![f32_ty.into(); 4], - SparseMmaElementAttr::E2m1, - SparseMmaElementAttr::E2m1, - SparseMmaOverflowAttr::Wrapping, - SparseMmaMetadataAttr::Ordered - ), - f8f6f4_mma!( - operands(zero), - vec![f32_ty.into(); 4], - SparseMmaElementAttr::E2m1, - SparseMmaElementAttr::E2m1, - SparseMmaOverflowAttr::NotApplicable, - SparseMmaMetadataAttr::Standard - ), - f8f6f4_mma!( - operands(zero), - vec![u32_ty.into(); 4], - SparseMmaElementAttr::E2m1, - SparseMmaElementAttr::E2m1, - SparseMmaOverflowAttr::NotApplicable, - SparseMmaMetadataAttr::Ordered - ), - ] { - assert!(verify_op(&invalid, &ctx).is_err()); - } -} - -#[test] -fn test_matrix_memory_ops_verify_pointer_and_packed_register_types() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); - let f32_ty = FP32Type::get(&ctx); - let ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), true); - let load_ptr_ty = MirPtrType::get_generic(&mut ctx, u32_ty.into(), true); - - let load_block = BasicBlock::new(&mut ctx, None, vec![load_ptr_ty.into()]); - let load_pointer = load_block.deref(&ctx).get_argument(0); - let load = make_ldmatrix_x2(&mut ctx, load_pointer, vec![u32_ty.into(), u32_ty.into()]); - assert!(load.verify(&ctx).is_ok()); - - let bad_load_pointer_block = BasicBlock::new(&mut ctx, None, vec![i64_ty.into()]); - let bad_pointer = bad_load_pointer_block.deref(&ctx).get_argument(0); - let bad_load_pointer = - make_ldmatrix_x2(&mut ctx, bad_pointer, vec![u32_ty.into(), u32_ty.into()]); - assert!(bad_load_pointer.verify(&ctx).is_err()); - - let bad_load_result = - make_ldmatrix_x2(&mut ctx, load_pointer, vec![u32_ty.into(), f32_ty.into()]); - assert!(bad_load_result.verify(&ctx).is_err()); - - let store_block = BasicBlock::new( - &mut ctx, - None, - vec![ - ptr_ty.into(), - i32_ty.into(), - i32_ty.into(), - i32_ty.into(), - i32_ty.into(), - ], - ); - let store_operands = (0..5) - .map(|index| store_block.deref(&ctx).get_argument(index)) - .collect(); - let store = Operation::new( - &mut ctx, - StmatrixM8n8X4Op::get_concrete_op_info(), - vec![], - store_operands, - vec![], - 0, - ); - assert!(StmatrixM8n8X4Op::new(store).verify(&ctx).is_ok()); - - let bad_store_block = BasicBlock::new( - &mut ctx, - None, - vec![ - ptr_ty.into(), - f32_ty.into(), - i32_ty.into(), - i32_ty.into(), - i32_ty.into(), - ], - ); - let bad_store_operands = (0..5) - .map(|index| bad_store_block.deref(&ctx).get_argument(index)) - .collect(); - let bad_store = Operation::new( - &mut ctx, - StmatrixM8n8X4Op::get_concrete_op_info(), - vec![], - bad_store_operands, - vec![], - 0, - ); - assert!(StmatrixM8n8X4Op::new(bad_store).verify(&ctx).is_err()); -} - -#[test] -fn test_packed_atomic_add_verifies_exact_raw_u32_shape() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let signless_i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let generic_u32_ptr = MirPtrType::get_generic(&mut ctx, u32_ty.into(), true); - let global_u32_ptr = MirPtrType::get_global(&mut ctx, u32_ty.into(), true); - let immutable_u32_ptr = MirPtrType::get_generic(&mut ctx, u32_ty.into(), false); - let shared_u32_ptr = MirPtrType::get_shared(&mut ctx, u32_ty.into(), true); - let generic_u64_ptr = MirPtrType::get_generic(&mut ctx, u64_ty.into(), true); - - let block = BasicBlock::new( - &mut ctx, - None, - vec![ - generic_u32_ptr.into(), - global_u32_ptr.into(), - immutable_u32_ptr.into(), - shared_u32_ptr.into(), - generic_u64_ptr.into(), - u32_ty.into(), - signless_i32_ty.into(), - ], - ); - let generic_ptr = block.deref(&ctx).get_argument(0); - let global_ptr = block.deref(&ctx).get_argument(1); - let immutable_ptr = block.deref(&ctx).get_argument(2); - let shared_ptr = block.deref(&ctx).get_argument(3); - let wrong_pointee_ptr = block.deref(&ctx).get_argument(4); - let addend = block.deref(&ctx).get_argument(5); - let signless_addend = block.deref(&ctx).get_argument(6); - - macro_rules! check_variant { - ($op:ty) => {{ - for pointer in [generic_ptr, global_ptr] { - let valid = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - vec![u32_ty.into()], - vec![pointer, addend], - vec![], - 0, - ); - assert!(verify_op(&<$op>::new(valid), &ctx).is_ok()); - } - - for pointer in [immutable_ptr, shared_ptr, wrong_pointee_ptr] { - let invalid = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - vec![u32_ty.into()], - vec![pointer, addend], - vec![], - 0, - ); - assert!(verify_op(&<$op>::new(invalid), &ctx).is_err()); - } - - let bad_address = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - vec![u32_ty.into()], - vec![addend, addend], - vec![], - 0, - ); - assert!(verify_op(&<$op>::new(bad_address), &ctx).is_err()); - - let bad_addend = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - vec![u32_ty.into()], - vec![generic_ptr, signless_addend], - vec![], - 0, - ); - assert!(verify_op(&<$op>::new(bad_addend), &ctx).is_err()); - - let bad_result = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - vec![signless_i32_ty.into()], - vec![generic_ptr, addend], - vec![], - 0, - ); - assert!(verify_op(&<$op>::new(bad_result), &ctx).is_err()); - - let bad_counts = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - vec![], - vec![generic_ptr], - vec![], - 0, - ); - assert!(verify_op(&<$op>::new(bad_counts), &ctx).is_err()); - }}; - } - - check_variant!(NvvmAtomAddF16x2Op); - check_variant!(NvvmAtomAddBf16x2Op); -} - -#[test] -fn test_generated_packed_atomic_add_requires_closed_attributes_and_raw_u32_shape() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let signless_i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let signed_i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let global_u32_ptr = MirPtrType::get_global(&mut ctx, u32_ty.into(), true); - let immutable_u32_ptr = MirPtrType::get_global(&mut ctx, u32_ty.into(), false); - let shared_u32_ptr = MirPtrType::get_shared(&mut ctx, u32_ty.into(), true); - let local_u32_ptr = MirPtrType::get(&mut ctx, u32_ty.into(), true, address_space::LOCAL); - let constant_u32_ptr = MirPtrType::get_constant(&mut ctx, u32_ty.into(), true); - let global_u64_ptr = MirPtrType::get_global(&mut ctx, u64_ty.into(), true); - let block = BasicBlock::new( - &mut ctx, - None, - vec![ - global_u32_ptr.into(), - immutable_u32_ptr.into(), - shared_u32_ptr.into(), - local_u32_ptr.into(), - constant_u32_ptr.into(), - global_u64_ptr.into(), - u32_ty.into(), - signless_i32_ty.into(), - ], - ); - let global = block.deref(&ctx).get_argument(0); - let immutable = block.deref(&ctx).get_argument(1); - let shared = block.deref(&ctx).get_argument(2); - let local = block.deref(&ctx).get_argument(3); - let constant = block.deref(&ctx).get_argument(4); - let wrong_pointee = block.deref(&ctx).get_argument(5); - let addend = block.deref(&ctx).get_argument(6); - let signless = block.deref(&ctx).get_argument(7); - - for format in [ - PackedAtomicFormatAttr::F16x2, - PackedAtomicFormatAttr::Bf16x2, - ] { - let valid = PackedAtomicAddOp::build(&mut ctx, global, addend, format); - assert!(verify_op(&PackedAtomicAddOp::new(valid), &ctx).is_ok()); - } - - for pointer in [immutable, shared, local, constant, wrong_pointee, addend] { - let invalid = - PackedAtomicAddOp::build(&mut ctx, pointer, addend, PackedAtomicFormatAttr::F16x2); - assert!(verify_op(&PackedAtomicAddOp::new(invalid), &ctx).is_err()); - } - let bad_addend = Operation::new( - &mut ctx, - PackedAtomicAddOp::get_concrete_op_info(), - vec![u32_ty.into()], - vec![global, signless], - vec![], - 0, - ); - assert!(verify_op(&PackedAtomicAddOp::new(bad_addend), &ctx).is_err()); - - fn set_closed_attributes(ctx: &Context, op: pliron::context::Ptr) { - let packed = PackedAtomicAddOp::new(op); - packed.set_attr_nvvm_packed_atomic_format(ctx, PackedAtomicFormatAttr::F16x2); - packed.set_attr_nvvm_packed_atomic_state_space(ctx, PackedAtomicStateSpaceAttr::Global); - packed.set_attr_nvvm_packed_atomic_ordering(ctx, PackedAtomicOrderingAttr::Relaxed); - packed.set_attr_nvvm_packed_atomic_scope(ctx, PackedAtomicScopeAttr::Gpu); - packed.set_attr_nvvm_packed_atomic_rounding(ctx, PackedAtomicRoundingAttr::Rn); - packed.set_attr_nvvm_packed_atomic_subnormal(ctx, PackedAtomicSubnormalAttr::NoFtz); - packed.set_attr_nvvm_packed_atomic_atomicity(ctx, PackedAtomicAtomicityAttr::PerElement); - } - - for result_ty in [signless_i32_ty.into(), signed_i32_ty.into(), u64_ty.into()] { - let bad_result = Operation::new( - &mut ctx, - PackedAtomicAddOp::get_concrete_op_info(), - vec![result_ty], - vec![global, addend], - vec![], - 0, - ); - set_closed_attributes(&ctx, bad_result); - assert!(verify_op(&PackedAtomicAddOp::new(bad_result), &ctx).is_err()); - } - - for (results, operands) in [ - (vec![], vec![global, addend]), - (vec![u32_ty.into(), u32_ty.into()], vec![global, addend]), - (vec![u32_ty.into()], vec![global]), - (vec![u32_ty.into()], vec![global, addend, addend]), - ] { - let bad_counts = Operation::new( - &mut ctx, - PackedAtomicAddOp::get_concrete_op_info(), - results, - operands, - vec![], - 0, - ); - set_closed_attributes(&ctx, bad_counts); - assert!(verify_op(&PackedAtomicAddOp::new(bad_counts), &ctx).is_err()); - } - - // A structurally correct operation without the closed semantic attributes - // must fail instead of inheriting implicit defaults. - let missing_attributes = Operation::new( - &mut ctx, - PackedAtomicAddOp::get_concrete_op_info(), - vec![u32_ty.into()], - vec![global, addend], - vec![], - 0, - ); - assert!(verify_op(&PackedAtomicAddOp::new(missing_attributes), &ctx).is_err()); -} - -#[test] -fn test_thread_register_ops_verify_i32_results() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - - let tid_x = Operation::new( - &mut ctx, - ReadPtxSregTidXOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![], - vec![], - 0, - ); - assert!(ReadPtxSregTidXOp::new(tid_x).verify(&ctx).is_ok()); - - let lane_id = Operation::new( - &mut ctx, - ReadPtxSregLaneIdOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![], - vec![], - 0, - ); - assert!(ReadPtxSregLaneIdOp::new(lane_id).verify(&ctx).is_ok()); -} - -#[test] -fn test_thread_register_ops_reject_non_i32_results() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); - let op = Operation::new( - &mut ctx, - ReadPtxSregTidXOp::get_concrete_op_info(), - vec![i64_ty.into()], - vec![], - vec![], - 0, - ); - - assert!(ReadPtxSregTidXOp::new(op).verify(&ctx).is_err()); -} - -#[test] -fn test_lanemask_ops_verify_i32_results() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - - // Each lane-position mask is a zero-operand, single-i32-result sreg read. - let lt = Operation::new( - &mut ctx, - ReadPtxSregLanemaskLtOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![], - vec![], - 0, - ); - assert!(ReadPtxSregLanemaskLtOp::new(lt).verify(&ctx).is_ok()); - - let le = Operation::new( - &mut ctx, - ReadPtxSregLanemaskLeOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![], - vec![], - 0, - ); - assert!(ReadPtxSregLanemaskLeOp::new(le).verify(&ctx).is_ok()); - - let eq = Operation::new( - &mut ctx, - ReadPtxSregLanemaskEqOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![], - vec![], - 0, - ); - assert!(ReadPtxSregLanemaskEqOp::new(eq).verify(&ctx).is_ok()); - - let ge = Operation::new( - &mut ctx, - ReadPtxSregLanemaskGeOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![], - vec![], - 0, - ); - assert!(ReadPtxSregLanemaskGeOp::new(ge).verify(&ctx).is_ok()); - - let gt = Operation::new( - &mut ctx, - ReadPtxSregLanemaskGtOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![], - vec![], - 0, - ); - assert!(ReadPtxSregLanemaskGtOp::new(gt).verify(&ctx).is_ok()); -} - -#[test] -fn test_lanemask_op_rejects_non_i32_result() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - // A 64-bit result must fail the shared lane-position mask verifier. - let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); - let op = Operation::new( - &mut ctx, - ReadPtxSregLanemaskLtOp::get_concrete_op_info(), - vec![i64_ty.into()], - vec![], - vec![], - 0, - ); - assert!(ReadPtxSregLanemaskLtOp::new(op).verify(&ctx).is_err()); -} - -#[test] -fn test_generated_vote_sync_family_requires_exact_mask_predicate_and_result_types() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); - let block = BasicBlock::new( - &mut ctx, - None, - vec![i32_ty.into(), i1_ty.into(), i64_ty.into()], - ); - let mask = block.deref(&ctx).get_argument(0); - let predicate = block.deref(&ctx).get_argument(1); - let wide_mask = block.deref(&ctx).get_argument(2); - - macro_rules! check_vote { - ($op:ty, $result_ty:expr, $wrong_result_ty:expr) => {{ - let valid = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - vec![$result_ty.into()], - vec![mask, predicate], - vec![], - 0, - ); - assert!(verify_op(&<$op>::new(valid), &ctx).is_ok()); - - for operands in [vec![], vec![mask], vec![mask, predicate, predicate]] { - let wrong_arity = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - vec![$result_ty.into()], - operands, - vec![], - 0, - ); - assert!(verify_op(&<$op>::new(wrong_arity), &ctx).is_err()); - } - - for results in [vec![], vec![$result_ty.into(), $result_ty.into()]] { - let wrong_arity = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - results, - vec![mask, predicate], - vec![], - 0, - ); - assert!(verify_op(&<$op>::new(wrong_arity), &ctx).is_err()); - } - - let wrong_mask = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - vec![$result_ty.into()], - vec![wide_mask, predicate], - vec![], - 0, - ); - assert!(verify_op(&<$op>::new(wrong_mask), &ctx).is_err()); - - let wrong_predicate = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - vec![$result_ty.into()], - vec![mask, mask], - vec![], - 0, - ); - assert!(verify_op(&<$op>::new(wrong_predicate), &ctx).is_err()); - - let wrong_result = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - vec![$wrong_result_ty.into()], - vec![mask, predicate], - vec![], - 0, - ); - assert!(verify_op(&<$op>::new(wrong_result), &ctx).is_err()); - }}; - } - - check_vote!(VoteSyncAllOp, i1_ty, u32_ty); - check_vote!(VoteSyncAnyOp, i1_ty, u32_ty); - check_vote!(VoteSyncBallotOp, u32_ty, i1_ty); - check_vote!(VoteSyncUniOp, i1_ty, u32_ty); -} - -#[test] -fn test_generated_active_mask_requires_no_operands_and_i32_result() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); - let block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); - let unexpected_operand = block.deref(&ctx).get_argument(0); - - let valid = Operation::new( - &mut ctx, - ActiveMaskOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![], - vec![], - 0, - ); - assert!(verify_op(&ActiveMaskOp::new(valid), &ctx).is_ok()); - - let wrong_operand_count = Operation::new( - &mut ctx, - ActiveMaskOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![unexpected_operand], - vec![], - 0, - ); - assert!(verify_op(&ActiveMaskOp::new(wrong_operand_count), &ctx).is_err()); - - for results in [vec![], vec![i32_ty.into(), i32_ty.into()]] { - let wrong_result_count = Operation::new( - &mut ctx, - ActiveMaskOp::get_concrete_op_info(), - results, - vec![], - vec![], - 0, - ); - assert!(verify_op(&ActiveMaskOp::new(wrong_result_count), &ctx).is_err()); - } - - let wrong_result_width = Operation::new( - &mut ctx, - ActiveMaskOp::get_concrete_op_info(), - vec![i64_ty.into()], - vec![], - vec![], - 0, - ); - assert!(verify_op(&ActiveMaskOp::new(wrong_result_width), &ctx).is_err()); -} - -#[test] -fn test_generated_match_family_requires_exact_mask_value_and_result_widths() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); - let block = BasicBlock::new( - &mut ctx, - None, - vec![i32_ty.into(), i32_ty.into(), i64_ty.into()], - ); - let mask = block.deref(&ctx).get_argument(0); - let value32 = block.deref(&ctx).get_argument(1); - let value64 = block.deref(&ctx).get_argument(2); - - macro_rules! check_match { - ($op:ty, $value:expr, $wrong_value:expr) => {{ - let valid = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - vec![i32_ty.into()], - vec![mask, $value], - vec![], - 0, - ); - assert!(verify_op(&<$op>::new(valid), &ctx).is_ok()); - - for operands in [vec![], vec![mask], vec![mask, $value, $value]] { - let wrong_operand_count = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - vec![i32_ty.into()], - operands, - vec![], - 0, - ); - assert!(verify_op(&<$op>::new(wrong_operand_count), &ctx).is_err()); - } - - for results in [vec![], vec![i32_ty.into(), i32_ty.into()]] { - let wrong_result_count = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - results, - vec![mask, $value], - vec![], - 0, - ); - assert!(verify_op(&<$op>::new(wrong_result_count), &ctx).is_err()); - } - - let wrong_mask_width = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - vec![i32_ty.into()], - vec![value64, $value], - vec![], - 0, - ); - assert!(verify_op(&<$op>::new(wrong_mask_width), &ctx).is_err()); - - let wrong_value_width = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - vec![i32_ty.into()], - vec![mask, $wrong_value], - vec![], - 0, - ); - assert!(verify_op(&<$op>::new(wrong_value_width), &ctx).is_err()); - - let wrong_result_width = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - vec![i64_ty.into()], - vec![mask, $value], - vec![], - 0, - ); - assert!(verify_op(&<$op>::new(wrong_result_width), &ctx).is_err()); - }}; - } - - check_match!(MatchAnySyncI32Op, value32, value64); - check_match!(MatchAllSyncI32Op, value32, value64); - check_match!(MatchAnySyncI64Op, value64, value32); - check_match!(MatchAllSyncI64Op, value64, value32); -} - -#[test] -fn test_special_register_ops_verify_authoritative_widths() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); - - macro_rules! check_width { - ($op:ty, $good:expr, $bad:expr) => {{ - let good = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - vec![$good.into()], - vec![], - vec![], - 0, - ); - assert!( - verify_op(&<$op>::new(good), &ctx).is_ok(), - "{} must accept its PTX register width", - stringify!($op) - ); - - let bad = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - vec![$bad.into()], - vec![], - vec![], - 0, - ); - assert!( - verify_op(&<$op>::new(bad), &ctx).is_err(), - "{} must reject the other integer width", - stringify!($op) - ); - }}; - } - - check_width!(ReadPtxSregWarpIdOp, i32_ty, i64_ty); - check_width!(ReadPtxSregNwarpIdOp, i32_ty, i64_ty); - check_width!(ReadPtxSregSmIdOp, i32_ty, i64_ty); - check_width!(ReadPtxSregNsmIdOp, i32_ty, i64_ty); - check_width!(ReadPtxSregDynamicSmemSizeOp, i32_ty, i64_ty); - check_width!(ReadPtxSregTotalSmemSizeOp, i32_ty, i64_ty); - check_width!(ReadPtxSregGridIdOp, i64_ty, i32_ty); -} - -#[test] -fn test_sync_ops_construct_and_verify() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - let barrier = Operation::new( - &mut ctx, - Barrier0Op::get_concrete_op_info(), - vec![], - vec![], - vec![], - 0, - ); - assert!(Barrier0Op::new(barrier).verify(&ctx).is_ok()); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); - let unexpected_operand = block.deref(&ctx).get_argument(0); - let bad_operand = Operation::new( - &mut ctx, - Barrier0Op::get_concrete_op_info(), - vec![], - vec![unexpected_operand], - vec![], - 0, - ); - assert!(verify_op(&Barrier0Op::new(bad_operand), &ctx).is_err()); - - let bad_result = Operation::new( - &mut ctx, - Barrier0Op::get_concrete_op_info(), - vec![i32_ty.into()], - vec![], - vec![], - 0, - ); - assert!(verify_op(&Barrier0Op::new(bad_result), &ctx).is_err()); - - let block_fence = Operation::new( - &mut ctx, - ThreadfenceBlockOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 0, - ); - assert!(ThreadfenceBlockOp::new(block_fence).verify(&ctx).is_ok()); - - let device_fence = Operation::new( - &mut ctx, - ThreadfenceOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 0, - ); - assert!(ThreadfenceOp::new(device_fence).verify(&ctx).is_ok()); - - let system_fence = Operation::new( - &mut ctx, - ThreadfenceSystemOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 0, - ); - assert!(ThreadfenceSystemOp::new(system_fence).verify(&ctx).is_ok()); -} - -#[test] -fn test_bf16x2_fma_constructs_and_verifies_three_operands() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - - let a = Operation::new( - &mut ctx, - ReadPtxSregTidXOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![], - vec![], - 0, - ); - let b = Operation::new( - &mut ctx, - ReadPtxSregTidXOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![], - vec![], - 0, - ); - let c = Operation::new( - &mut ctx, - ReadPtxSregTidXOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![], - vec![], - 0, - ); - - let operands = vec![ - a.deref(&ctx).get_result(0), - b.deref(&ctx).get_result(0), - c.deref(&ctx).get_result(0), - ]; - - let fma = Operation::new( - &mut ctx, - FmaBf16x2Op::get_concrete_op_info(), - vec![u32_ty.into()], - operands, - vec![], - 0, - ); - - assert!(FmaBf16x2Op::new(fma).verify(&ctx).is_ok()); -} - -#[test] -fn test_generated_dot_products_require_three_i32_operands_and_one_i32_result() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); - let block = BasicBlock::new( - &mut ctx, - None, - vec![i32_ty.into(), i32_ty.into(), i32_ty.into(), i64_ty.into()], - ); - let a = block.deref(&ctx).get_argument(0); - let b = block.deref(&ctx).get_argument(1); - let c = block.deref(&ctx).get_argument(2); - let wide = block.deref(&ctx).get_argument(3); - - macro_rules! check_variant { - ($op:ty) => {{ - let valid = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - vec![i32_ty.into()], - vec![a, b, c], - vec![], - 0, - ); - assert!(verify_op(&<$op>::new(valid), &ctx).is_ok()); - - let wrong_width = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - vec![i32_ty.into()], - vec![a, b, wide], - vec![], - 0, - ); - assert!(verify_op(&<$op>::new(wrong_width), &ctx).is_err()); - - for (results, operands) in [ - (vec![], vec![a, b, c]), - (vec![i32_ty.into(), i32_ty.into()], vec![a, b, c]), - (vec![i32_ty.into()], vec![a, b]), - (vec![i32_ty.into()], vec![a, b, c, c]), - ] { - let wrong_count = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - results, - operands, - vec![], - 0, - ); - assert!(verify_op(&<$op>::new(wrong_count), &ctx).is_err()); - } - }}; - } - - check_variant!(Dp4aS32Op); - check_variant!(Dp4aU32Op); - check_variant!(Dp2aS32Op); - check_variant!(Dp2aU32Op); -} - -#[test] -fn test_redux_sync_add_construct_and_verify() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - - // A block supplies the two operands [mask, value]. - let block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into(), i32_ty.into()]); - let mask = block.deref(&ctx).get_argument(0); - let value = block.deref(&ctx).get_argument(1); - - // Valid: 2 operands, 1 result (matches NOpdsInterface<2>/NResultsInterface<1>). - let op = Operation::new( - &mut ctx, - ReduxSyncAddOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![mask, value], - vec![], - 0, - ); - assert!(verify_op(&ReduxSyncAddOp::new(op), &ctx).is_ok()); - - // Invalid: wrong operand count (1 instead of 2) must fail verification. - let bad_opnds = Operation::new( - &mut ctx, - ReduxSyncAddOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![mask], - vec![], - 0, - ); - assert!(verify_op(&ReduxSyncAddOp::new(bad_opnds), &ctx).is_err()); - - // Invalid: wrong result count (0 instead of 1) must fail verification. - let bad_results = Operation::new( - &mut ctx, - ReduxSyncAddOp::get_concrete_op_info(), - vec![], - vec![mask, value], - vec![], - 0, - ); - assert!(verify_op(&ReduxSyncAddOp::new(bad_results), &ctx).is_err()); -} - -#[test] -fn test_redux_sync_integer_family_construct_and_verify() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into(), i32_ty.into()]); - let mask = block.deref(&ctx).get_argument(0); - let value = block.deref(&ctx).get_argument(1); - - // Every integer-family variant has the same 2-operand/1-result shape. A - // valid build of each must verify; a wrong operand count must not. The - // `new` wrapper is invoked so each concrete op type is exercised. - macro_rules! check_variant { - ($op:ty) => {{ - let good = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - vec![i32_ty.into()], - vec![mask, value], - vec![], - 0, - ); - assert!( - verify_op(&<$op>::new(good), &ctx).is_ok(), - "{} should verify with [mask, value] -> i32", - stringify!($op) - ); - - let bad = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - vec![i32_ty.into()], - vec![mask], - vec![], - 0, - ); - assert!( - verify_op(&<$op>::new(bad), &ctx).is_err(), - "{} must reject a single operand", - stringify!($op) - ); - }}; - } - - check_variant!(ReduxSyncUminOp); - check_variant!(ReduxSyncMinOp); - check_variant!(ReduxSyncUmaxOp); - check_variant!(ReduxSyncMaxOp); - check_variant!(ReduxSyncAndOp); - check_variant!(ReduxSyncOrOp); - check_variant!(ReduxSyncXorOp); -} - -#[test] -fn test_redux_sync_f32_family_construct_and_verify() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let f32_ty = FP32Type::get(&ctx); - let block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into(), f32_ty.into()]); - let mask = block.deref(&ctx).get_argument(0); - let value = block.deref(&ctx).get_argument(1); - - // The good case goes through the generated `build()` so its f32 result - // construction is exercised; the bad case's i32 value and result must - // fail the `is_f32` verifier path. - macro_rules! check_variant { - ($op:ty) => {{ - let good = <$op>::build(&mut ctx, mask, value); - assert!( - verify_op(&<$op>::new(good), &ctx).is_ok(), - "{} should verify with [mask, value] -> f32", - stringify!($op) - ); - - let bad = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - vec![i32_ty.into()], - vec![mask, mask], - vec![], - 0, - ); - assert!( - verify_op(&<$op>::new(bad), &ctx).is_err(), - "{} must reject an i32 value and result", - stringify!($op) - ); - }}; - } - - check_variant!(ReduxSyncFminOp); - check_variant!(ReduxSyncFminNanOp); - check_variant!(ReduxSyncFminAbsOp); - check_variant!(ReduxSyncFminAbsNanOp); - check_variant!(ReduxSyncFmaxOp); - check_variant!(ReduxSyncFmaxNanOp); - check_variant!(ReduxSyncFmaxAbsOp); - check_variant!(ReduxSyncFmaxAbsNanOp); -} - -#[test] -fn test_bar_warp_sync_construct_and_verify() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); - let block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into(), i64_ty.into()]); - let mask = block.deref(&ctx).get_argument(0); - let wrong_mask = block.deref(&ctx).get_argument(1); - - let valid = BarWarpSyncOp::build(&mut ctx, mask); - assert!(verify_op(&BarWarpSyncOp::new(valid), &ctx).is_ok()); - - let wrong_type = BarWarpSyncOp::build(&mut ctx, wrong_mask); - assert!(verify_op(&BarWarpSyncOp::new(wrong_type), &ctx).is_err()); - - let wrong_arity = Operation::new( - &mut ctx, - BarWarpSyncOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 0, - ); - assert!(verify_op(&BarWarpSyncOp::new(wrong_arity), &ctx).is_err()); -} - -#[test] -fn test_elect_sync_construct_and_verify() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); - - // A block supplies the single `mask` operand. - let block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); - let mask = block.deref(&ctx).get_argument(0); - - // Valid: 1 operand [mask], 2 results [leader (i32), is_elected (i1)] - // (matches NOpdsInterface<1>/NResultsInterface<2>). - let op = Operation::new( - &mut ctx, - ElectSyncOp::get_concrete_op_info(), - vec![i32_ty.into(), i1_ty.into()], - vec![mask], - vec![], - 0, - ); - assert!(verify_op(&ElectSyncOp::new(op), &ctx).is_ok()); - - // Invalid: wrong operand count (0 instead of 1) must fail verification. - let bad_opnds = Operation::new( - &mut ctx, - ElectSyncOp::get_concrete_op_info(), - vec![i32_ty.into(), i1_ty.into()], - vec![], - vec![], - 0, - ); - assert!(verify_op(&ElectSyncOp::new(bad_opnds), &ctx).is_err()); - - // Invalid: wrong result count (1 instead of 2) must fail verification. - let bad_results = Operation::new( - &mut ctx, - ElectSyncOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![mask], - vec![], - 0, - ); - assert!(verify_op(&ElectSyncOp::new(bad_results), &ctx).is_err()); - - // Invalid: a correctly-sized result list still cannot use a pointer as - // either result and thereby manufacture a MIR pointer kind. - let pointer_ty = - MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::UniqueRef); - let bad_result_type = Operation::new( - &mut ctx, - ElectSyncOp::get_concrete_op_info(), - vec![i32_ty.into(), pointer_ty.into()], - vec![mask], - vec![], - 0, - ); - assert!(verify_op(&ElectSyncOp::new(bad_result_type), &ctx).is_err()); -} - -#[test] -fn test_clc_query_rejects_pointer_results() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - - let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let block = BasicBlock::new(&mut ctx, None, vec![u64_ty.into(), u64_ty.into()]); - let operands = vec![ - block.deref(&ctx).get_argument(0), - block.deref(&ctx).get_argument(1), - ]; - - let valid = Operation::new( - &mut ctx, - ClcQueryIsCanceledOp::get_concrete_op_info(), - vec![u32_ty.into()], - operands.clone(), - vec![], - 0, - ); - assert!(verify_op(&ClcQueryIsCanceledOp::new(valid), &ctx).is_ok()); - - let pointer_ty = - MirPtrType::get_generic_with_kind(&mut ctx, u32_ty.into(), true, MirPointerKind::UniqueRef); - let bad = Operation::new( - &mut ctx, - ClcQueryIsCanceledOp::get_concrete_op_info(), - vec![pointer_ty.into()], - operands, - vec![], - 0, - ); - assert!(verify_op(&ClcQueryIsCanceledOp::new(bad), &ctx).is_err()); -} - -#[test] -fn test_shfl_sync_i64_construct_and_verify() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); - - let block = BasicBlock::new( - &mut ctx, - None, - vec![i32_ty.into(), i64_ty.into(), i32_ty.into()], - ); - let mask = block.deref(&ctx).get_argument(0); - let value = block.deref(&ctx).get_argument(1); - let lane = block.deref(&ctx).get_argument(2); - - macro_rules! check_mode { - ($op:ty) => {{ - let valid = <$op>::build(&mut ctx, mask, value, lane); - assert!(verify_op(&<$op>::new(valid), &ctx).is_ok()); - - for (operands, result_ty) in [ - (vec![mask, value], i64_ty.into()), - (vec![value, value, lane], i64_ty.into()), - (vec![mask, mask, lane], i64_ty.into()), - (vec![mask, value, value], i64_ty.into()), - (vec![mask, value, lane], i32_ty.into()), - ] { - let invalid = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - vec![result_ty], - operands, - vec![], - 0, - ); - assert!(verify_op(&<$op>::new(invalid), &ctx).is_err()); - } - - let no_result = Operation::new( - &mut ctx, - <$op>::get_concrete_op_info(), - vec![], - vec![mask, value, lane], - vec![], - 0, - ); - assert!(verify_op(&<$op>::new(no_result), &ctx).is_err()); - }}; - } - - check_mode!(ShflSyncIdxI64Op); - check_mode!(ShflSyncBflyI64Op); - check_mode!(ShflSyncDownI64Op); - check_mode!(ShflSyncUpI64Op); -} - -#[test] -fn handwritten_atomic_carriers_reject_malformed_ir() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let f32_ty = FP32Type::get(&ctx); - let generic_ptr = MirPtrType::get_generic(&mut ctx, u32_ty.into(), false); - let block = BasicBlock::new( - &mut ctx, - None, - vec![ - generic_ptr.into(), - u32_ty.into(), - i32_ty.into(), - f32_ty.into(), - ], - ); - let pointer = block.deref(&ctx).get_argument(0); - let u32_value = block.deref(&ctx).get_argument(1); - let i32_value = block.deref(&ctx).get_argument(2); - let f32_value = block.deref(&ctx).get_argument(3); - - assert!( - NvvmAtomicLoadOp::build( - &mut ctx, - pointer, - u32_ty.into(), - AtomicOrdering::Acquire, - AtomicScope::Device, - ) - .verify(&ctx) - .is_ok() - ); - assert!( - NvvmAtomicStoreOp::build( - &mut ctx, - u32_value, - pointer, - AtomicOrdering::Release, - AtomicScope::Device, - ) - .verify(&ctx) - .is_ok() - ); - assert!( - NvvmAtomicRmwOp::build( - &mut ctx, - pointer, - u32_value, - u32_ty.into(), - AtomicRmwKind::Add, - AtomicOrdering::AcqRel, - AtomicScope::Device, - ) - .verify(&ctx) - .is_ok() - ); - assert!( - NvvmAtomicRmwOp::build( - &mut ctx, - pointer, - f32_value, - f32_ty.into(), - AtomicRmwKind::FAdd, - AtomicOrdering::Relaxed, - AtomicScope::Device, - ) - .verify(&ctx) - .is_ok() - ); - assert!( - NvvmAtomicCmpxchgOp::build( - &mut ctx, - pointer, - u32_value, - u32_value, - u32_ty.into(), - AtomicOrdering::Relaxed, - AtomicOrdering::Acquire, - AtomicScope::System, - ) - .verify(&ctx) - .is_ok() - ); - - for invalid in [ - NvvmAtomicLoadOp::build( - &mut ctx, - pointer, - u32_ty.into(), - AtomicOrdering::Release, - AtomicScope::Device, - ), - NvvmAtomicLoadOp::build( - &mut ctx, - u32_value, - u32_ty.into(), - AtomicOrdering::Relaxed, - AtomicScope::Device, - ), - ] { - assert!(invalid.verify(&ctx).is_err()); - } - assert!( - NvvmAtomicStoreOp::build( - &mut ctx, - u32_value, - pointer, - AtomicOrdering::Acquire, - AtomicScope::Device, - ) - .verify(&ctx) - .is_err() - ); - assert!( - NvvmAtomicRmwOp::build( - &mut ctx, - pointer, - u32_value, - u64_ty.into(), - AtomicRmwKind::Add, - AtomicOrdering::Relaxed, - AtomicScope::Device, - ) - .verify(&ctx) - .is_err() - ); - assert!( - NvvmAtomicRmwOp::build( - &mut ctx, - pointer, - u32_value, - u32_ty.into(), - AtomicRmwKind::FAdd, - AtomicOrdering::Relaxed, - AtomicScope::Device, - ) - .verify(&ctx) - .is_err() - ); - assert!( - NvvmAtomicRmwOp::build( - &mut ctx, - pointer, - i32_value, - i32_ty.into(), - AtomicRmwKind::UMin, - AtomicOrdering::Relaxed, - AtomicScope::Device, - ) - .verify(&ctx) - .is_err() - ); - assert!( - NvvmAtomicCmpxchgOp::build( - &mut ctx, - pointer, - f32_value, - f32_value, - f32_ty.into(), - AtomicOrdering::SeqCst, - AtomicOrdering::Relaxed, - AtomicScope::Device, - ) - .verify(&ctx) - .is_err() - ); - assert!( - NvvmAtomicCmpxchgOp::build( - &mut ctx, - pointer, - u32_value, - i32_value, - u32_ty.into(), - AtomicOrdering::SeqCst, - AtomicOrdering::Relaxed, - AtomicScope::Device, - ) - .verify(&ctx) - .is_err() - ); - assert!( - NvvmAtomicCmpxchgOp::build( - &mut ctx, - pointer, - u32_value, - u32_value, - u32_ty.into(), - AtomicOrdering::SeqCst, - AtomicOrdering::Release, - AtomicScope::Device, - ) - .verify(&ctx) - .is_err() - ); - - let missing_attributes = Operation::new( - &mut ctx, - NvvmAtomicRmwOp::get_concrete_op_info(), - vec![u32_ty.into()], - vec![pointer, u32_value], - vec![], - 0, - ); - assert!( - NvvmAtomicRmwOp::new(missing_attributes) - .verify(&ctx) - .is_err() - ); - let bad_count = Operation::new( - &mut ctx, - NvvmAtomicCmpxchgOp::get_concrete_op_info(), - vec![], - vec![pointer, u32_value], - vec![], - 0, - ); - assert!(NvvmAtomicCmpxchgOp::new(bad_count).verify(&ctx).is_err()); -} - -#[test] -fn atomic_cmpxchg_accepts_exactly_llvm_ordering_pairs() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let pointer_ty = MirPtrType::get_generic(&mut ctx, u32_ty.into(), false); - let block = BasicBlock::new(&mut ctx, None, vec![pointer_ty.into(), u32_ty.into()]); - let pointer = block.deref(&ctx).get_argument(0); - let value = block.deref(&ctx).get_argument(1); - - let orderings = [ - AtomicOrdering::Relaxed, - AtomicOrdering::Acquire, - AtomicOrdering::Release, - AtomicOrdering::AcqRel, - AtomicOrdering::SeqCst, - ]; - for success in &orderings { - for failure in &orderings { - let expected = matches!( - failure, - AtomicOrdering::Relaxed | AtomicOrdering::Acquire | AtomicOrdering::SeqCst - ); - let actual = NvvmAtomicCmpxchgOp::build( - &mut ctx, - pointer, - value, - value, - u32_ty.into(), - success.clone(), - failure.clone(), - AtomicScope::Device, - ) - .verify(&ctx) - .is_ok(); - assert_eq!( - actual, expected, - "unexpected cmpxchg ordering result for success={success:?}, failure={failure:?}" - ); - } - } -} - -#[test] -fn handwritten_ffi_and_wgmma_carriers_verify_exact_shapes() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - - let u8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let f32_ty = FP32Type::get(&ctx); - let pointer_ty = MirPtrType::get_generic(&mut ctx, u8_ty.into(), false); - let accumulator_pointer_ty = MirPtrType::get_generic(&mut ctx, u8_ty.into(), true); - let global_pointer_ty = MirPtrType::get_global(&mut ctx, u8_ty.into(), false); - let mutable_global_pointer_ty = MirPtrType::get_global(&mut ctx, u8_ty.into(), true); - let block = BasicBlock::new( - &mut ctx, - None, - vec![ - pointer_ty.into(), - accumulator_pointer_ty.into(), - global_pointer_ty.into(), - mutable_global_pointer_ty.into(), - u32_ty.into(), - u64_ty.into(), - f32_ty.into(), - ], - ); - let pointer = block.deref(&ctx).get_argument(0); - let accumulator_pointer = block.deref(&ctx).get_argument(1); - let global_pointer = block.deref(&ctx).get_argument(2); - let mutable_global_pointer = block.deref(&ctx).get_argument(3); - let u32_value = block.deref(&ctx).get_argument(4); - let u64_value = block.deref(&ctx).get_argument(5); - let f32_value = block.deref(&ctx).get_argument(6); - - let vprintf = VprintfOp::build(&mut ctx, pointer, pointer); - assert!(VprintfOp::new(vprintf).verify(&ctx).is_ok()); - - let i64_signless_ty = IntegerType::get(&ctx, 64, Signedness::Signless); - let assertfail = AssertFailOp::build(&mut ctx, pointer, pointer, u32_value, pointer, u64_value); - assert!(AssertFailOp::new(assertfail).verify(&ctx).is_ok()); - for (operands, results) in [ - // message must be a MIR pointer - ( - vec![u32_value, pointer, u32_value, pointer, u64_value], - vec![], - ), - // line must be a 32-bit integer - ( - vec![pointer, pointer, u64_value, pointer, u64_value], - vec![], - ), - // char size must be a 64-bit integer - ( - vec![pointer, pointer, u32_value, pointer, u32_value], - vec![], - ), - // wrong operand count - (vec![pointer, pointer, u32_value], vec![]), - // must have no results - ( - vec![pointer, pointer, u32_value, pointer, u64_value], - vec![i64_signless_ty.into()], - ), - ] { - let invalid = Operation::new( - &mut ctx, - AssertFailOp::get_concrete_op_info(), - results, - operands, - vec![], - 0, - ); - assert!(AssertFailOp::new(invalid).verify(&ctx).is_err()); - } - let bad_vprintf = Operation::new( - &mut ctx, - VprintfOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![pointer, u32_value], - vec![], - 0, - ); - assert!(VprintfOp::new(bad_vprintf).verify(&ctx).is_err()); - - let descriptor = Operation::new( - &mut ctx, - WgmmaMakeSmemDescOp::get_concrete_op_info(), - vec![u64_ty.into()], - vec![pointer], - vec![], - 0, - ); - assert!(WgmmaMakeSmemDescOp::new(descriptor).verify(&ctx).is_ok()); - for (operands, results) in [ - (vec![u32_value], vec![u64_ty.into()]), - (vec![global_pointer], vec![u64_ty.into()]), - (vec![pointer], vec![u32_ty.into()]), - (vec![], vec![u64_ty.into()]), - ] { - let invalid = Operation::new( - &mut ctx, - WgmmaMakeSmemDescOp::get_concrete_op_info(), - results, - operands, - vec![], - 0, - ); - assert!(WgmmaMakeSmemDescOp::new(invalid).verify(&ctx).is_err()); - } - - let cvta = Operation::new( - &mut ctx, - CvtaGenericToSharedOffsetOp::get_concrete_op_info(), - vec![u64_ty.into()], - vec![pointer], - vec![], - 0, - ); - assert!(CvtaGenericToSharedOffsetOp::new(cvta).verify(&ctx).is_ok()); - for (operands, results) in [ - // Operand must be a MIR pointer. - (vec![u32_value], vec![u64_ty.into()]), - // Operand must point to generic or shared memory. - (vec![global_pointer], vec![u64_ty.into()]), - // Result must be u64. - (vec![pointer], vec![u32_ty.into()]), - // Exact arity is required. - (vec![], vec![u64_ty.into()]), - ] { - let invalid = Operation::new( - &mut ctx, - CvtaGenericToSharedOffsetOp::get_concrete_op_info(), - results, - operands, - vec![], - 0, - ); - assert!( - CvtaGenericToSharedOffsetOp::new(invalid) - .verify(&ctx) - .is_err() - ); - } - - let mma = Operation::new( - &mut ctx, - WgmmaMmaM64N64K16F32Bf16Op::get_concrete_op_info(), - vec![], - vec![accumulator_pointer, u64_value, u64_value], - vec![], - 0, - ); - assert!(WgmmaMmaM64N64K16F32Bf16Op::new(mma).verify(&ctx).is_ok()); - for (operands, results) in [ - // Accumulator must be mutable. - (vec![pointer, u64_value, u64_value], vec![]), - // Accumulator must use generic address space. - (vec![mutable_global_pointer, u64_value, u64_value], vec![]), - // Accumulator must be a MIR pointer. - (vec![u32_value, u64_value, u64_value], vec![]), - // Descriptors must be u64. - (vec![accumulator_pointer, u32_value, u64_value], vec![]), - // Exact arity is required. - (vec![accumulator_pointer, u64_value], vec![]), - // No results are permitted. - ( - vec![accumulator_pointer, u64_value, u64_value], - vec![u32_ty.into()], - ), - ] { - let invalid = Operation::new( - &mut ctx, - WgmmaMmaM64N64K16F32Bf16Op::get_concrete_op_info(), - results, - operands, - vec![], - 0, - ); - assert!( - WgmmaMmaM64N64K16F32Bf16Op::new(invalid) - .verify(&ctx) - .is_err() - ); - } - - let group = WgmmaMmaGroupM64N64K16F32Bf16Op::build( - &mut ctx, - accumulator_pointer, - vec![u64_value, u64_value, u64_value, u64_value], - ); - assert!( - WgmmaMmaGroupM64N64K16F32Bf16Op::new(group) - .verify(&ctx) - .is_ok() - ); - for (operands, results) in [ - // Missing descriptor. - (vec![accumulator_pointer, u64_value], vec![]), - // Incomplete descriptor pair. - ( - vec![accumulator_pointer, u64_value, u64_value, u64_value], - vec![], - ), - // Accumulator must be mutable. - (vec![pointer, u64_value, u64_value], vec![]), - // Accumulator must use generic address space. - (vec![mutable_global_pointer, u64_value, u64_value], vec![]), - // Accumulator must be a MIR pointer. - (vec![u32_value, u64_value, u64_value], vec![]), - // Descriptors must be u64. - (vec![accumulator_pointer, u64_value, u32_value], vec![]), - // No results are permitted. - ( - vec![accumulator_pointer, u64_value, u64_value], - vec![u32_ty.into()], - ), - ] { - let invalid = Operation::new( - &mut ctx, - WgmmaMmaGroupM64N64K16F32Bf16Op::get_concrete_op_info(), - results, - operands, - vec![], - 0, - ); - assert!( - WgmmaMmaGroupM64N64K16F32Bf16Op::new(invalid) - .verify(&ctx) - .is_err() - ); - } - - let value_group = WgmmaMmaGroupValuesM64N64K16F32Bf16Op::build( - &mut ctx, - vec![f32_value; 32], - vec![u64_value, u64_value, u64_value, u64_value], - ); - { - let value_group_ref = value_group.deref(&ctx); - assert_eq!(value_group_ref.get_num_operands(), 36); - assert_eq!(value_group_ref.get_num_results(), 32); - } - assert!( - WgmmaMmaGroupValuesM64N64K16F32Bf16Op::new(value_group) - .verify(&ctx) - .is_ok() - ); - - let too_few_accumulators = WgmmaMmaGroupValuesM64N64K16F32Bf16Op::build( - &mut ctx, - vec![f32_value; 31], - vec![u64_value, u64_value], - ); - assert!( - WgmmaMmaGroupValuesM64N64K16F32Bf16Op::new(too_few_accumulators) - .verify(&ctx) - .is_err() - ); - - let missing_descriptors = - WgmmaMmaGroupValuesM64N64K16F32Bf16Op::build(&mut ctx, vec![f32_value; 32], vec![]); - assert!( - WgmmaMmaGroupValuesM64N64K16F32Bf16Op::new(missing_descriptors) - .verify(&ctx) - .is_err() - ); - - let incomplete_descriptor_pair = WgmmaMmaGroupValuesM64N64K16F32Bf16Op::build( - &mut ctx, - vec![f32_value; 32], - vec![u64_value, u64_value, u64_value], - ); - assert!( - WgmmaMmaGroupValuesM64N64K16F32Bf16Op::new(incomplete_descriptor_pair) - .verify(&ctx) - .is_err() - ); - - let mut wrong_accumulator_operands = vec![f32_value; 32]; - wrong_accumulator_operands[0] = u32_value; - wrong_accumulator_operands.extend([u64_value, u64_value]); - let wrong_accumulator = Operation::new( - &mut ctx, - WgmmaMmaGroupValuesM64N64K16F32Bf16Op::get_concrete_op_info(), - vec![f32_ty.into(); 32], - wrong_accumulator_operands, - vec![], - 0, - ); - assert!( - WgmmaMmaGroupValuesM64N64K16F32Bf16Op::new(wrong_accumulator) - .verify(&ctx) - .is_err() - ); - - let mut wrong_descriptor_operands = vec![f32_value; 32]; - wrong_descriptor_operands.extend([u32_value, u64_value]); - let wrong_descriptor = Operation::new( - &mut ctx, - WgmmaMmaGroupValuesM64N64K16F32Bf16Op::get_concrete_op_info(), - vec![f32_ty.into(); 32], - wrong_descriptor_operands, - vec![], - 0, - ); - assert!( - WgmmaMmaGroupValuesM64N64K16F32Bf16Op::new(wrong_descriptor) - .verify(&ctx) - .is_err() - ); - - let mut valid_value_operands = vec![f32_value; 32]; - valid_value_operands.extend([u64_value, u64_value]); - - let wrong_result_count = Operation::new( - &mut ctx, - WgmmaMmaGroupValuesM64N64K16F32Bf16Op::get_concrete_op_info(), - vec![f32_ty.into(); 31], - valid_value_operands.clone(), - vec![], - 0, - ); - assert!( - WgmmaMmaGroupValuesM64N64K16F32Bf16Op::new(wrong_result_count) - .verify(&ctx) - .is_err() - ); - - let mut wrong_result_types = vec![f32_ty.into(); 32]; - wrong_result_types[0] = u32_ty.into(); - let wrong_result_type = Operation::new( - &mut ctx, - WgmmaMmaGroupValuesM64N64K16F32Bf16Op::get_concrete_op_info(), - wrong_result_types, - valid_value_operands, - vec![], - 0, - ); - assert!( - WgmmaMmaGroupValuesM64N64K16F32Bf16Op::new(wrong_result_type) - .verify(&ctx) - .is_err() - ); - - let wide_mma = Operation::new( - &mut ctx, - WgmmaMmaM64N128K16F32Bf16Op::get_concrete_op_info(), - vec![], - vec![accumulator_pointer, u64_value, u64_value], - vec![], - 0, - ); - assert!( - WgmmaMmaM64N128K16F32Bf16Op::new(wide_mma) - .verify(&ctx) - .is_ok() - ); - - let wide_value_group = WgmmaMmaGroupValuesM64N128K16F32Bf16Op::build( - &mut ctx, - vec![f32_value; 64], - vec![u64_value, u64_value], - ); - { - let group_ref = wide_value_group.deref(&ctx); - assert_eq!(group_ref.get_num_operands(), 66); - assert_eq!(group_ref.get_num_results(), 64); - } - assert!( - WgmmaMmaGroupValuesM64N128K16F32Bf16Op::new(wide_value_group) - .verify(&ctx) - .is_ok() - ); - - let wide_too_few = WgmmaMmaGroupValuesM64N128K16F32Bf16Op::build( - &mut ctx, - vec![f32_value; 63], - vec![u64_value, u64_value], - ); - assert!( - WgmmaMmaGroupValuesM64N128K16F32Bf16Op::new(wide_too_few) - .verify(&ctx) - .is_err() - ); - - let wide_incomplete_pair = WgmmaMmaGroupValuesM64N128K16F32Bf16Op::build( - &mut ctx, - vec![f32_value; 64], - vec![u64_value, u64_value, u64_value], - ); - assert!( - WgmmaMmaGroupValuesM64N128K16F32Bf16Op::new(wide_incomplete_pair) - .verify(&ctx) - .is_err() - ); - - let f16_mma = Operation::new( - &mut ctx, - WgmmaMmaM64N64K16F32F16Op::get_concrete_op_info(), - vec![], - vec![accumulator_pointer, u64_value, u64_value], - vec![], - 0, - ); - assert!(WgmmaMmaM64N64K16F32F16Op::new(f16_mma).verify(&ctx).is_ok()); - for operands in [ - vec![pointer, u64_value, u64_value], - vec![mutable_global_pointer, u64_value, u64_value], - vec![accumulator_pointer, u32_value, u64_value], - ] { - let invalid = Operation::new( - &mut ctx, - WgmmaMmaM64N64K16F32F16Op::get_concrete_op_info(), - vec![], - operands, - vec![], - 0, - ); - assert!( - WgmmaMmaM64N64K16F32F16Op::new(invalid) - .verify(&ctx) - .is_err() - ); - } - - let f16_value_group = WgmmaMmaGroupValuesM64N64K16F32F16Op::build( - &mut ctx, - vec![f32_value; 32], - vec![u64_value, u64_value], - ); - assert!( - WgmmaMmaGroupValuesM64N64K16F32F16Op::new(f16_value_group) - .verify(&ctx) - .is_ok() - ); - - let f16_too_few_accumulators = WgmmaMmaGroupValuesM64N64K16F32F16Op::build( - &mut ctx, - vec![f32_value; 31], - vec![u64_value, u64_value], - ); - assert!( - WgmmaMmaGroupValuesM64N64K16F32F16Op::new(f16_too_few_accumulators) - .verify(&ctx) - .is_err() - ); - - let f16_incomplete_descriptor_pair = WgmmaMmaGroupValuesM64N64K16F32F16Op::build( - &mut ctx, - vec![f32_value; 32], - vec![u64_value, u64_value, u64_value], - ); - assert!( - WgmmaMmaGroupValuesM64N64K16F32F16Op::new(f16_incomplete_descriptor_pair) - .verify(&ctx) - .is_err() - ); - - let mut f16_wrong_accumulator_operands = vec![f32_value; 32]; - f16_wrong_accumulator_operands[0] = u32_value; - f16_wrong_accumulator_operands.extend([u64_value, u64_value]); - let f16_wrong_accumulator = Operation::new( - &mut ctx, - WgmmaMmaGroupValuesM64N64K16F32F16Op::get_concrete_op_info(), - vec![f32_ty.into(); 32], - f16_wrong_accumulator_operands, - vec![], - 0, - ); - assert!( - WgmmaMmaGroupValuesM64N64K16F32F16Op::new(f16_wrong_accumulator) - .verify(&ctx) - .is_err() - ); - - let tf32_mma = Operation::new( - &mut ctx, - WgmmaMmaM64N64K8F32Tf32Op::get_concrete_op_info(), - vec![], - vec![accumulator_pointer, u64_value, u64_value], - vec![], - 0, - ); - assert!( - WgmmaMmaM64N64K8F32Tf32Op::new(tf32_mma) - .verify(&ctx) - .is_ok() - ); - for operands in [ - vec![pointer, u64_value, u64_value], - vec![mutable_global_pointer, u64_value, u64_value], - vec![accumulator_pointer, u32_value, u64_value], - ] { - let invalid = Operation::new( - &mut ctx, - WgmmaMmaM64N64K8F32Tf32Op::get_concrete_op_info(), - vec![], - operands, - vec![], - 0, - ); - assert!( - WgmmaMmaM64N64K8F32Tf32Op::new(invalid) - .verify(&ctx) - .is_err() - ); - } - - let tf32_value_group = WgmmaMmaGroupValuesM64N64K8F32Tf32Op::build( - &mut ctx, - vec![f32_value; 32], - vec![u64_value, u64_value], - ); - assert!( - WgmmaMmaGroupValuesM64N64K8F32Tf32Op::new(tf32_value_group) - .verify(&ctx) - .is_ok() - ); - - let tf32_too_few_accumulators = WgmmaMmaGroupValuesM64N64K8F32Tf32Op::build( - &mut ctx, - vec![f32_value; 31], - vec![u64_value, u64_value], - ); - assert!( - WgmmaMmaGroupValuesM64N64K8F32Tf32Op::new(tf32_too_few_accumulators) - .verify(&ctx) - .is_err() - ); - - let tf32_incomplete_descriptor_pair = WgmmaMmaGroupValuesM64N64K8F32Tf32Op::build( - &mut ctx, - vec![f32_value; 32], - vec![u64_value, u64_value, u64_value], - ); - assert!( - WgmmaMmaGroupValuesM64N64K8F32Tf32Op::new(tf32_incomplete_descriptor_pair) - .verify(&ctx) - .is_err() - ); - - let mut tf32_wrong_accumulator_operands = vec![f32_value; 32]; - tf32_wrong_accumulator_operands[0] = u32_value; - tf32_wrong_accumulator_operands.extend([u64_value, u64_value]); - let tf32_wrong_accumulator = Operation::new( - &mut ctx, - WgmmaMmaGroupValuesM64N64K8F32Tf32Op::get_concrete_op_info(), - vec![f32_ty.into(); 32], - tf32_wrong_accumulator_operands, - vec![], - 0, - ); - assert!( - WgmmaMmaGroupValuesM64N64K8F32Tf32Op::new(tf32_wrong_accumulator) - .verify(&ctx) - .is_err() - ); - - let loop_group = WgmmaMmaLoopValuesM64N64K16F32Bf16Op::build( - &mut ctx, - vec![f32_value; 32], - u64_value, - u64_value, - u64_value, - u64_value, - u64_value, - ); - { - let loop_group_ref = loop_group.deref(&ctx); - assert_eq!(loop_group_ref.get_num_operands(), 37); - assert_eq!(loop_group_ref.get_num_results(), 32); - } - assert!( - WgmmaMmaLoopValuesM64N64K16F32Bf16Op::new(loop_group) - .verify(&ctx) - .is_ok() - ); - - let f16_loop_group = WgmmaMmaLoopValuesM64N64K16F32F16Op::build( - &mut ctx, - vec![f32_value; 32], - u64_value, - u64_value, - u64_value, - u64_value, - u64_value, - ); - { - let loop_group_ref = f16_loop_group.deref(&ctx); - assert_eq!(loop_group_ref.get_num_operands(), 37); - assert_eq!(loop_group_ref.get_num_results(), 32); - } - assert!( - WgmmaMmaLoopValuesM64N64K16F32F16Op::new(f16_loop_group) - .verify(&ctx) - .is_ok() - ); - - let mut f16_wrong_loop_control_operands = vec![f32_value; 32]; - f16_wrong_loop_control_operands.extend([u64_value, u64_value, u32_value, u64_value, u64_value]); - let f16_wrong_loop_control = Operation::new( - &mut ctx, - WgmmaMmaLoopValuesM64N64K16F32F16Op::get_concrete_op_info(), - vec![f32_ty.into(); 32], - f16_wrong_loop_control_operands, - vec![], - 0, - ); - assert!( - WgmmaMmaLoopValuesM64N64K16F32F16Op::new(f16_wrong_loop_control) - .verify(&ctx) - .is_err() - ); - - let too_few_loop_accumulators = WgmmaMmaLoopValuesM64N64K16F32Bf16Op::build( - &mut ctx, - vec![f32_value; 31], - u64_value, - u64_value, - u64_value, - u64_value, - u64_value, - ); - assert!( - WgmmaMmaLoopValuesM64N64K16F32Bf16Op::new(too_few_loop_accumulators) - .verify(&ctx) - .is_err() - ); - - let mut wrong_loop_control_operands = vec![f32_value; 32]; - wrong_loop_control_operands.extend([u64_value, u64_value, u32_value, u64_value, u64_value]); - let wrong_loop_control = Operation::new( - &mut ctx, - WgmmaMmaLoopValuesM64N64K16F32Bf16Op::get_concrete_op_info(), - vec![f32_ty.into(); 32], - wrong_loop_control_operands, - vec![], - 0, - ); - assert!( - WgmmaMmaLoopValuesM64N64K16F32Bf16Op::new(wrong_loop_control) - .verify(&ctx) - .is_err() - ); - - let mut valid_loop_operands = vec![f32_value; 32]; - valid_loop_operands.extend([u64_value; 5]); - - let wrong_loop_result_count = Operation::new( - &mut ctx, - WgmmaMmaLoopValuesM64N64K16F32Bf16Op::get_concrete_op_info(), - vec![f32_ty.into(); 31], - valid_loop_operands.clone(), - vec![], - 0, - ); - assert!( - WgmmaMmaLoopValuesM64N64K16F32Bf16Op::new(wrong_loop_result_count) - .verify(&ctx) - .is_err() - ); - - let mut wrong_loop_result_types = vec![f32_ty.into(); 32]; - wrong_loop_result_types[0] = u32_ty.into(); - let wrong_loop_result_type = Operation::new( - &mut ctx, - WgmmaMmaLoopValuesM64N64K16F32Bf16Op::get_concrete_op_info(), - wrong_loop_result_types, - valid_loop_operands, - vec![], - 0, - ); - assert!( - WgmmaMmaLoopValuesM64N64K16F32Bf16Op::new(wrong_loop_result_type) - .verify(&ctx) - .is_err() - ); - - let counted_pipeline_group = WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::build( - &mut ctx, - vec![f32_value; 64], - vec![u64_value; 4], - vec![u64_value; 4], - u64_value, - 1, - ); - { - let counted_pipeline_ref = counted_pipeline_group.deref(&ctx); - assert_eq!(counted_pipeline_ref.get_num_operands(), 73); - assert_eq!(counted_pipeline_ref.get_num_results(), 64); - } - let counted_pipeline = - WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::new(counted_pipeline_group); - assert_eq!(counted_pipeline.max_pending_groups(&ctx), Some(1)); - assert!(counted_pipeline.verify(&ctx).is_ok()); - - let three_slot_counted_pipeline = WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::build( - &mut ctx, - vec![f32_value; 96], - vec![u64_value; 6], - vec![u64_value; 6], - u64_value, - 2, - ); - { - let counted_pipeline_ref = three_slot_counted_pipeline.deref(&ctx); - assert_eq!(counted_pipeline_ref.get_num_operands(), 109); - assert_eq!(counted_pipeline_ref.get_num_results(), 96); - } - let three_slot_counted_pipeline = - WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::new(three_slot_counted_pipeline); - assert_eq!( - three_slot_counted_pipeline.max_pending_groups(&ctx), - Some(2) - ); - assert!(three_slot_counted_pipeline.verify(&ctx).is_ok()); - - let too_many_accumulators_for_wait_one = WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::build( - &mut ctx, - vec![f32_value; 96], - vec![u64_value; 4], - vec![u64_value; 4], - u64_value, - 1, - ); - assert!( - WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::new(too_many_accumulators_for_wait_one) - .verify(&ctx) - .is_err() - ); - - let too_few_accumulators_for_wait_two = WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::build( - &mut ctx, - vec![f32_value; 64], - vec![u64_value; 6], - vec![u64_value; 6], - u64_value, - 2, - ); - assert!( - WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::new(too_few_accumulators_for_wait_two) - .verify(&ctx) - .is_err() - ); - - let unsupported_wait_three = WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::build( - &mut ctx, - vec![f32_value; 128], - vec![u64_value; 8], - vec![u64_value; 8], - u64_value, - 3, - ); - assert!( - WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::new(unsupported_wait_three) - .verify(&ctx) - .is_err() - ); - - let wrong_counted_pipeline_control = WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::build( - &mut ctx, - vec![f32_value; 64], - vec![u64_value; 4], - vec![u64_value, u64_value, u32_value, u64_value], - u64_value, - 1, - ); - assert!( - WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::new(wrong_counted_pipeline_control) - .verify(&ctx) - .is_err() - ); - - let mut valid_counted_pipeline_operands = vec![f32_value; 64]; - valid_counted_pipeline_operands.extend([u64_value; 9]); - let wrong_counted_pipeline_result_count = Operation::new( - &mut ctx, - WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::get_concrete_op_info(), - vec![f32_ty.into(); 63], - valid_counted_pipeline_operands.clone(), - vec![], - 0, - ); - let wrong_counted_pipeline_result_count = - WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::new(wrong_counted_pipeline_result_count); - wrong_counted_pipeline_result_count - .set_attr_counted_max_pending_groups(&ctx, WgmmaMaxPendingAttr(1)); - assert!(wrong_counted_pipeline_result_count.verify(&ctx).is_err()); - - let mut wrong_counted_pipeline_result_types = vec![f32_ty.into(); 64]; - wrong_counted_pipeline_result_types[0] = u32_ty.into(); - let wrong_counted_pipeline_result_type = Operation::new( - &mut ctx, - WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::get_concrete_op_info(), - wrong_counted_pipeline_result_types, - valid_counted_pipeline_operands, - vec![], - 0, - ); - let wrong_counted_pipeline_result_type = - WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::new(wrong_counted_pipeline_result_type); - wrong_counted_pipeline_result_type - .set_attr_counted_max_pending_groups(&ctx, WgmmaMaxPendingAttr(1)); - assert!(wrong_counted_pipeline_result_type.verify(&ctx).is_err()); - - let pipeline_group = WgmmaMmaPipelineValuesM64N64K16F32Bf16Op::build( - &mut ctx, - vec![f32_value; 64], - vec![u64_value; 8], - 1, - ); - { - let pipeline_ref = pipeline_group.deref(&ctx); - assert_eq!(pipeline_ref.get_num_operands(), 72); - assert_eq!(pipeline_ref.get_num_results(), 64); - } - let pipeline = WgmmaMmaPipelineValuesM64N64K16F32Bf16Op::new(pipeline_group); - assert_eq!(pipeline.max_pending_groups(&ctx), Some(1)); - assert!(pipeline.verify(&ctx).is_ok()); - - let too_few_pipeline_slots = WgmmaMmaPipelineValuesM64N64K16F32Bf16Op::build( - &mut ctx, - vec![f32_value; 32], - vec![u64_value; 4], - 1, - ); - assert!( - WgmmaMmaPipelineValuesM64N64K16F32Bf16Op::new(too_few_pipeline_slots) - .verify(&ctx) - .is_err() - ); - - assert!(WgmmaMaxPendingAttr(1).verify(&ctx).is_ok()); - assert!(WgmmaMaxPendingAttr(7).verify(&ctx).is_ok()); - assert!(WgmmaMaxPendingAttr(0).verify(&ctx).is_err()); - assert!(WgmmaMaxPendingAttr(8).verify(&ctx).is_err()); - - let zero_pending_pipeline = WgmmaMmaPipelineValuesM64N64K16F32Bf16Op::build( - &mut ctx, - vec![f32_value; 32], - vec![u64_value; 2], - 0, - ); - let zero_pending_error = WgmmaMmaPipelineValuesM64N64K16F32Bf16Op::new(zero_pending_pipeline) - .verify(&ctx) - .unwrap_err(); - assert!( - zero_pending_error - .err - .to_string() - .contains("nvvm.wgmma_max_pending_groups must be in 1..=7, got 0"), - "unexpected zero-pending error: {}", - zero_pending_error.err - ); - - let excess_pending_pipeline = WgmmaMmaPipelineValuesM64N64K16F32Bf16Op::build( - &mut ctx, - vec![f32_value; 32], - vec![u64_value; 2], - 8, - ); - let excess_pending_error = - WgmmaMmaPipelineValuesM64N64K16F32Bf16Op::new(excess_pending_pipeline) - .verify(&ctx) - .unwrap_err(); - assert!( - excess_pending_error - .err - .to_string() - .contains("nvvm.wgmma_max_pending_groups must be in 1..=7, got 8"), - "unexpected excess-pending error: {}", - excess_pending_error.err - ); - - let mut missing_attr_operands = vec![f32_value; 64]; - missing_attr_operands.extend([u64_value; 8]); - let missing_attr_pipeline = Operation::new( - &mut ctx, - WgmmaMmaPipelineValuesM64N64K16F32Bf16Op::get_concrete_op_info(), - vec![f32_ty.into(); 64], - missing_attr_operands, - vec![], - 0, - ); - let missing_attr_error = WgmmaMmaPipelineValuesM64N64K16F32Bf16Op::new(missing_attr_pipeline) - .verify(&ctx) - .unwrap_err(); - assert!( - missing_attr_error - .err - .to_string() - .contains("requires an nvvm.wgmma_max_pending_groups attribute"), - "unexpected missing-attribute error: {}", - missing_attr_error.err - ); -} - -#[test] -fn test_inline_ptx_results_must_match_output_constraints() { - let mut ctx = Context::new(); - dialect_nvvm::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); - let input = block.deref(&ctx).get_argument(0); - - let void = InlinePtxOp::build(&mut ctx, vec![], vec![], "membar.gl;", "", true, false); - assert!(verify_op(&InlinePtxOp::new(void), &ctx).is_ok()); - - let single = InlinePtxOp::build( - &mut ctx, - vec![i32_ty.into()], - vec![input], - "add.u32 $0, $1, $1;", - "=r,r", - false, - false, - ); - assert!(verify_op(&InlinePtxOp::new(single), &ctx).is_ok()); - - let multi = InlinePtxOp::build( - &mut ctx, - vec![i32_ty.into(), i32_ty.into()], - vec![input], - "add.u32 $0, $2, $2; mul.lo.u32 $1, $2, $2;", - "=r,=r,r", - false, - false, - ); - assert!(verify_op(&InlinePtxOp::new(multi), &ctx).is_ok()); - - let missing_result = InlinePtxOp::build( - &mut ctx, - vec![i32_ty.into()], - vec![input], - "add.u32 $0, $2, $2; mul.lo.u32 $1, $2, $2;", - "=r,=r,r", - false, - false, - ); - assert!(verify_op(&InlinePtxOp::new(missing_result), &ctx).is_err()); - - let extra_result = InlinePtxOp::build( - &mut ctx, - vec![i32_ty.into()], - vec![input], - "prefetch.global.L1 [$0];", - "r", - true, - false, - ); - assert!(verify_op(&InlinePtxOp::new(extra_result), &ctx).is_err()); -} - -#[test] -fn test_inline_ptx_pointer_results_require_exact_inline_asm_authority() { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let unique_ty = - MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::UniqueRef); - let erased_ty = - MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::Erased); - let nested_erased_ty = MirTupleType::get(&mut ctx, vec![erased_ty.into()]); - - let build = |ctx: &mut Context, result_ty| { - InlinePtxOp::build( - ctx, - vec![result_ty], - vec![], - "mov.u64 $0, 0;", - "=l", - false, - false, - ) - }; - - for (name, result_ty) in [ - ("UniqueRef", unique_ty.into()), - ("Erased", erased_ty.into()), - ("nested Erased", nested_erased_ty.into()), - ] { - let unmarked = InlinePtxOp::new(build(&mut ctx, result_ty)); - let error = verify_op(&unmarked, &ctx).unwrap_err(); - assert!( - error - .err - .to_string() - .contains("pointer-carrying results require InlineAsm pointer-kind authority"), - "unexpected unmarked {name} error: {}", - error.err - ); - - let marked = InlinePtxOp::new(build(&mut ctx, result_ty)); - marked.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::InlineAsm); - assert!( - verify_op(&marked, &ctx).is_ok(), - "marked {name} pointer result should verify" - ); - } - - let wrong_authority = InlinePtxOp::new(build(&mut ctx, unique_ty.into())); - wrong_authority.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::RustCast); - let wrong_error = verify_op(&wrong_authority, &ctx).unwrap_err(); - assert!( - wrong_error.err.to_string().contains("found RustCast"), - "unexpected wrong-authority error: {}", - wrong_error.err - ); - - let spurious = InlinePtxOp::new(build(&mut ctx, i32_ty.into())); - spurious.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::InlineAsm); - let spurious_error = verify_op(&spurious, &ctx).unwrap_err(); - assert!( - spurious_error - .err - .to_string() - .contains("spurious for pointer-free results"), - "unexpected spurious-authority error: {}", - spurious_error.err - ); -} - -#[test] -fn test_inline_ptx_count_output_constraints() { - assert_eq!(InlinePtxOp::count_output_constraints("=r,r,r"), 1); - assert_eq!(InlinePtxOp::count_output_constraints("=r,=r,=f,=d,r,l"), 4); - assert_eq!(InlinePtxOp::count_output_constraints("r,l,~{memory}"), 0); - assert_eq!(InlinePtxOp::count_output_constraints(""), 0); - // `=` only counts as an output marker at the start of a token. - assert_eq!(InlinePtxOp::count_output_constraints("r,r=f"), 0); -} diff --git a/crates/dialect-nvvm/tests/ops_test/allowlist.rs b/crates/dialect-nvvm/tests/ops_test/allowlist.rs new file mode 100644 index 0000000000..bc48d66501 --- /dev/null +++ b/crates/dialect-nvvm/tests/ops_test/allowlist.rs @@ -0,0 +1,74 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#[test] +fn handwritten_ops_match_reviewed_allowlist() { + use std::{fs, path::Path}; + + let ops_dir = Path::new(env!("CARGO_MANIFEST_DIR")).join("src/ops"); + let mut found = Vec::new(); + + for entry in fs::read_dir(&ops_dir).expect("read top-level NVVM ops directory") { + let path = entry.expect("read NVVM ops entry").path(); + if path.extension().and_then(|extension| extension.to_str()) != Some("rs") { + continue; + } + + let source = fs::read_to_string(&path).expect("read handwritten NVVM ops source"); + let file = path + .file_name() + .and_then(|name| name.to_str()) + .expect("NVVM ops file name"); + + for (marker, _) in source.match_indices("#[pliron_op(") { + let declaration = source[marker..] + .lines() + .find_map(|line| line.trim().strip_prefix("pub struct ")) + .expect("pliron_op must be followed by a public struct"); + let name = declaration + .split(|character: char| !character.is_ascii_alphanumeric() && character != '_') + .next() + .expect("NVVM op struct name"); + found.push((file.to_owned(), name.to_owned())); + } + } + + let mut expected = [ + ("asm.rs", "InlinePtxOp"), + ("atomic.rs", "NvvmAtomicLoadOp"), + ("atomic.rs", "NvvmAtomicStoreOp"), + ("atomic.rs", "NvvmAtomicFenceOp"), + ("atomic.rs", "NvvmAtomicRmwOp"), + ("atomic.rs", "NvvmAtomicCmpxchgOp"), + ("cluster.rs", "ReadPtxSregClusterIdxOp"), + ("cluster.rs", "ReadPtxSregNclusterIdOp"), + ("debug.rs", "AssertFailOp"), + ("debug.rs", "VprintfOp"), + ("grid.rs", "GridSyncOp"), + ("memory.rs", "CvtaGenericToSharedOffsetOp"), + ("wgmma.rs", "WgmmaMakeSmemDescOp"), + ("wgmma.rs", "WgmmaMmaM64N64K16F32Bf16Op"), + ("wgmma.rs", "WgmmaMmaM64N128K16F32Bf16Op"), + ("wgmma.rs", "WgmmaMmaM64N64K16F32F16Op"), + ("wgmma.rs", "WgmmaMmaM64N64K8F32Tf32Op"), + ("wgmma.rs", "WgmmaMmaGroupM64N64K16F32Bf16Op"), + ("wgmma.rs", "WgmmaMmaGroupValuesM64N64K16F32Bf16Op"), + ("wgmma.rs", "WgmmaMmaGroupValuesM64N128K16F32Bf16Op"), + ("wgmma.rs", "WgmmaMmaGroupValuesM64N64K16F32F16Op"), + ("wgmma.rs", "WgmmaMmaGroupValuesM64N64K8F32Tf32Op"), + ("wgmma.rs", "WgmmaMmaLoopValuesM64N64K16F32Bf16Op"), + ("wgmma.rs", "WgmmaMmaLoopValuesM64N64K16F32F16Op"), + ("wgmma.rs", "WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op"), + ("wgmma.rs", "WgmmaMmaPipelineValuesM64N64K16F32Bf16Op"), + ]; + expected.sort_unstable(); + found.sort_unstable(); + + assert_eq!( + found, + expected.map(|(file, op)| (file.to_owned(), op.to_owned())), + "top-level handwritten NVVM ops changed; generate leaf ops or review this allowlist" + ); +} diff --git a/crates/dialect-nvvm/tests/ops_test/atomics.rs b/crates/dialect-nvvm/tests/ops_test/atomics.rs new file mode 100644 index 0000000000..5680a27e9f --- /dev/null +++ b/crates/dialect-nvvm/tests/ops_test/atomics.rs @@ -0,0 +1,516 @@ +/* + * 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::{ + AtomicOrdering, AtomicRmwKind, AtomicScope, NvvmAtomAddBf16x2Op, NvvmAtomAddF16x2Op, + NvvmAtomicCmpxchgOp, NvvmAtomicLoadOp, NvvmAtomicRmwOp, NvvmAtomicStoreOp, PackedAtomicAddOp, + PackedAtomicAtomicityAttr, PackedAtomicFormatAttr, PackedAtomicOrderingAttr, + PackedAtomicRoundingAttr, PackedAtomicScopeAttr, PackedAtomicStateSpaceAttr, + PackedAtomicSubnormalAttr, +}; + +use pliron::{ + basic_block::BasicBlock, + builtin::types::{FP32Type, IntegerType, Signedness}, + common_traits::Verify, + context::Context, + op::{Op, verify_op}, + operation::Operation, +}; + +#[test] +fn test_packed_atomic_add_verifies_exact_raw_u32_shape() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let signless_i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let generic_u32_ptr = MirPtrType::get_generic(&mut ctx, u32_ty.into(), true); + let global_u32_ptr = MirPtrType::get_global(&mut ctx, u32_ty.into(), true); + let immutable_u32_ptr = MirPtrType::get_generic(&mut ctx, u32_ty.into(), false); + let shared_u32_ptr = MirPtrType::get_shared(&mut ctx, u32_ty.into(), true); + let generic_u64_ptr = MirPtrType::get_generic(&mut ctx, u64_ty.into(), true); + + let block = BasicBlock::new( + &mut ctx, + None, + vec![ + generic_u32_ptr.into(), + global_u32_ptr.into(), + immutable_u32_ptr.into(), + shared_u32_ptr.into(), + generic_u64_ptr.into(), + u32_ty.into(), + signless_i32_ty.into(), + ], + ); + let generic_ptr = block.deref(&ctx).get_argument(0); + let global_ptr = block.deref(&ctx).get_argument(1); + let immutable_ptr = block.deref(&ctx).get_argument(2); + let shared_ptr = block.deref(&ctx).get_argument(3); + let wrong_pointee_ptr = block.deref(&ctx).get_argument(4); + let addend = block.deref(&ctx).get_argument(5); + let signless_addend = block.deref(&ctx).get_argument(6); + + macro_rules! check_variant { + ($op:ty) => {{ + for pointer in [generic_ptr, global_ptr] { + let valid = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + vec![u32_ty.into()], + vec![pointer, addend], + vec![], + 0, + ); + assert!(verify_op(&<$op>::new(valid), &ctx).is_ok()); + } + + for pointer in [immutable_ptr, shared_ptr, wrong_pointee_ptr] { + let invalid = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + vec![u32_ty.into()], + vec![pointer, addend], + vec![], + 0, + ); + assert!(verify_op(&<$op>::new(invalid), &ctx).is_err()); + } + + let bad_address = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + vec![u32_ty.into()], + vec![addend, addend], + vec![], + 0, + ); + assert!(verify_op(&<$op>::new(bad_address), &ctx).is_err()); + + let bad_addend = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + vec![u32_ty.into()], + vec![generic_ptr, signless_addend], + vec![], + 0, + ); + assert!(verify_op(&<$op>::new(bad_addend), &ctx).is_err()); + + let bad_result = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + vec![signless_i32_ty.into()], + vec![generic_ptr, addend], + vec![], + 0, + ); + assert!(verify_op(&<$op>::new(bad_result), &ctx).is_err()); + + let bad_counts = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + vec![], + vec![generic_ptr], + vec![], + 0, + ); + assert!(verify_op(&<$op>::new(bad_counts), &ctx).is_err()); + }}; + } + + check_variant!(NvvmAtomAddF16x2Op); + check_variant!(NvvmAtomAddBf16x2Op); +} + +#[test] +fn test_generated_packed_atomic_add_requires_closed_attributes_and_raw_u32_shape() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let signless_i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let signed_i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let global_u32_ptr = MirPtrType::get_global(&mut ctx, u32_ty.into(), true); + let immutable_u32_ptr = MirPtrType::get_global(&mut ctx, u32_ty.into(), false); + let shared_u32_ptr = MirPtrType::get_shared(&mut ctx, u32_ty.into(), true); + let local_u32_ptr = MirPtrType::get(&mut ctx, u32_ty.into(), true, address_space::LOCAL); + let constant_u32_ptr = MirPtrType::get_constant(&mut ctx, u32_ty.into(), true); + let global_u64_ptr = MirPtrType::get_global(&mut ctx, u64_ty.into(), true); + let block = BasicBlock::new( + &mut ctx, + None, + vec![ + global_u32_ptr.into(), + immutable_u32_ptr.into(), + shared_u32_ptr.into(), + local_u32_ptr.into(), + constant_u32_ptr.into(), + global_u64_ptr.into(), + u32_ty.into(), + signless_i32_ty.into(), + ], + ); + let global = block.deref(&ctx).get_argument(0); + let immutable = block.deref(&ctx).get_argument(1); + let shared = block.deref(&ctx).get_argument(2); + let local = block.deref(&ctx).get_argument(3); + let constant = block.deref(&ctx).get_argument(4); + let wrong_pointee = block.deref(&ctx).get_argument(5); + let addend = block.deref(&ctx).get_argument(6); + let signless = block.deref(&ctx).get_argument(7); + + for format in [ + PackedAtomicFormatAttr::F16x2, + PackedAtomicFormatAttr::Bf16x2, + ] { + let valid = PackedAtomicAddOp::build(&mut ctx, global, addend, format); + assert!(verify_op(&PackedAtomicAddOp::new(valid), &ctx).is_ok()); + } + + for pointer in [immutable, shared, local, constant, wrong_pointee, addend] { + let invalid = + PackedAtomicAddOp::build(&mut ctx, pointer, addend, PackedAtomicFormatAttr::F16x2); + assert!(verify_op(&PackedAtomicAddOp::new(invalid), &ctx).is_err()); + } + let bad_addend = Operation::new( + &mut ctx, + PackedAtomicAddOp::get_concrete_op_info(), + vec![u32_ty.into()], + vec![global, signless], + vec![], + 0, + ); + assert!(verify_op(&PackedAtomicAddOp::new(bad_addend), &ctx).is_err()); + + fn set_closed_attributes(ctx: &Context, op: pliron::context::Ptr) { + let packed = PackedAtomicAddOp::new(op); + packed.set_attr_nvvm_packed_atomic_format(ctx, PackedAtomicFormatAttr::F16x2); + packed.set_attr_nvvm_packed_atomic_state_space(ctx, PackedAtomicStateSpaceAttr::Global); + packed.set_attr_nvvm_packed_atomic_ordering(ctx, PackedAtomicOrderingAttr::Relaxed); + packed.set_attr_nvvm_packed_atomic_scope(ctx, PackedAtomicScopeAttr::Gpu); + packed.set_attr_nvvm_packed_atomic_rounding(ctx, PackedAtomicRoundingAttr::Rn); + packed.set_attr_nvvm_packed_atomic_subnormal(ctx, PackedAtomicSubnormalAttr::NoFtz); + packed.set_attr_nvvm_packed_atomic_atomicity(ctx, PackedAtomicAtomicityAttr::PerElement); + } + + for result_ty in [signless_i32_ty.into(), signed_i32_ty.into(), u64_ty.into()] { + let bad_result = Operation::new( + &mut ctx, + PackedAtomicAddOp::get_concrete_op_info(), + vec![result_ty], + vec![global, addend], + vec![], + 0, + ); + set_closed_attributes(&ctx, bad_result); + assert!(verify_op(&PackedAtomicAddOp::new(bad_result), &ctx).is_err()); + } + + for (results, operands) in [ + (vec![], vec![global, addend]), + (vec![u32_ty.into(), u32_ty.into()], vec![global, addend]), + (vec![u32_ty.into()], vec![global]), + (vec![u32_ty.into()], vec![global, addend, addend]), + ] { + let bad_counts = Operation::new( + &mut ctx, + PackedAtomicAddOp::get_concrete_op_info(), + results, + operands, + vec![], + 0, + ); + set_closed_attributes(&ctx, bad_counts); + assert!(verify_op(&PackedAtomicAddOp::new(bad_counts), &ctx).is_err()); + } + + // A structurally correct operation without the closed semantic attributes + // must fail instead of inheriting implicit defaults. + let missing_attributes = Operation::new( + &mut ctx, + PackedAtomicAddOp::get_concrete_op_info(), + vec![u32_ty.into()], + vec![global, addend], + vec![], + 0, + ); + assert!(verify_op(&PackedAtomicAddOp::new(missing_attributes), &ctx).is_err()); +} + +#[test] +fn handwritten_atomic_carriers_reject_malformed_ir() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let f32_ty = FP32Type::get(&ctx); + let generic_ptr = MirPtrType::get_generic(&mut ctx, u32_ty.into(), false); + let block = BasicBlock::new( + &mut ctx, + None, + vec![ + generic_ptr.into(), + u32_ty.into(), + i32_ty.into(), + f32_ty.into(), + ], + ); + let pointer = block.deref(&ctx).get_argument(0); + let u32_value = block.deref(&ctx).get_argument(1); + let i32_value = block.deref(&ctx).get_argument(2); + let f32_value = block.deref(&ctx).get_argument(3); + + assert!( + NvvmAtomicLoadOp::build( + &mut ctx, + pointer, + u32_ty.into(), + AtomicOrdering::Acquire, + AtomicScope::Device, + ) + .verify(&ctx) + .is_ok() + ); + assert!( + NvvmAtomicStoreOp::build( + &mut ctx, + u32_value, + pointer, + AtomicOrdering::Release, + AtomicScope::Device, + ) + .verify(&ctx) + .is_ok() + ); + assert!( + NvvmAtomicRmwOp::build( + &mut ctx, + pointer, + u32_value, + u32_ty.into(), + AtomicRmwKind::Add, + AtomicOrdering::AcqRel, + AtomicScope::Device, + ) + .verify(&ctx) + .is_ok() + ); + assert!( + NvvmAtomicRmwOp::build( + &mut ctx, + pointer, + f32_value, + f32_ty.into(), + AtomicRmwKind::FAdd, + AtomicOrdering::Relaxed, + AtomicScope::Device, + ) + .verify(&ctx) + .is_ok() + ); + assert!( + NvvmAtomicCmpxchgOp::build( + &mut ctx, + pointer, + u32_value, + u32_value, + u32_ty.into(), + AtomicOrdering::Relaxed, + AtomicOrdering::Acquire, + AtomicScope::System, + ) + .verify(&ctx) + .is_ok() + ); + + for invalid in [ + NvvmAtomicLoadOp::build( + &mut ctx, + pointer, + u32_ty.into(), + AtomicOrdering::Release, + AtomicScope::Device, + ), + NvvmAtomicLoadOp::build( + &mut ctx, + u32_value, + u32_ty.into(), + AtomicOrdering::Relaxed, + AtomicScope::Device, + ), + ] { + assert!(invalid.verify(&ctx).is_err()); + } + assert!( + NvvmAtomicStoreOp::build( + &mut ctx, + u32_value, + pointer, + AtomicOrdering::Acquire, + AtomicScope::Device, + ) + .verify(&ctx) + .is_err() + ); + assert!( + NvvmAtomicRmwOp::build( + &mut ctx, + pointer, + u32_value, + u64_ty.into(), + AtomicRmwKind::Add, + AtomicOrdering::Relaxed, + AtomicScope::Device, + ) + .verify(&ctx) + .is_err() + ); + assert!( + NvvmAtomicRmwOp::build( + &mut ctx, + pointer, + u32_value, + u32_ty.into(), + AtomicRmwKind::FAdd, + AtomicOrdering::Relaxed, + AtomicScope::Device, + ) + .verify(&ctx) + .is_err() + ); + assert!( + NvvmAtomicRmwOp::build( + &mut ctx, + pointer, + i32_value, + i32_ty.into(), + AtomicRmwKind::UMin, + AtomicOrdering::Relaxed, + AtomicScope::Device, + ) + .verify(&ctx) + .is_err() + ); + assert!( + NvvmAtomicCmpxchgOp::build( + &mut ctx, + pointer, + f32_value, + f32_value, + f32_ty.into(), + AtomicOrdering::SeqCst, + AtomicOrdering::Relaxed, + AtomicScope::Device, + ) + .verify(&ctx) + .is_err() + ); + assert!( + NvvmAtomicCmpxchgOp::build( + &mut ctx, + pointer, + u32_value, + i32_value, + u32_ty.into(), + AtomicOrdering::SeqCst, + AtomicOrdering::Relaxed, + AtomicScope::Device, + ) + .verify(&ctx) + .is_err() + ); + assert!( + NvvmAtomicCmpxchgOp::build( + &mut ctx, + pointer, + u32_value, + u32_value, + u32_ty.into(), + AtomicOrdering::SeqCst, + AtomicOrdering::Release, + AtomicScope::Device, + ) + .verify(&ctx) + .is_err() + ); + + let missing_attributes = Operation::new( + &mut ctx, + NvvmAtomicRmwOp::get_concrete_op_info(), + vec![u32_ty.into()], + vec![pointer, u32_value], + vec![], + 0, + ); + assert!( + NvvmAtomicRmwOp::new(missing_attributes) + .verify(&ctx) + .is_err() + ); + let bad_count = Operation::new( + &mut ctx, + NvvmAtomicCmpxchgOp::get_concrete_op_info(), + vec![], + vec![pointer, u32_value], + vec![], + 0, + ); + assert!(NvvmAtomicCmpxchgOp::new(bad_count).verify(&ctx).is_err()); +} + +#[test] +fn atomic_cmpxchg_accepts_exactly_llvm_ordering_pairs() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let pointer_ty = MirPtrType::get_generic(&mut ctx, u32_ty.into(), false); + let block = BasicBlock::new(&mut ctx, None, vec![pointer_ty.into(), u32_ty.into()]); + let pointer = block.deref(&ctx).get_argument(0); + let value = block.deref(&ctx).get_argument(1); + + let orderings = [ + AtomicOrdering::Relaxed, + AtomicOrdering::Acquire, + AtomicOrdering::Release, + AtomicOrdering::AcqRel, + AtomicOrdering::SeqCst, + ]; + for success in &orderings { + for failure in &orderings { + let expected = matches!( + failure, + AtomicOrdering::Relaxed | AtomicOrdering::Acquire | AtomicOrdering::SeqCst + ); + let actual = NvvmAtomicCmpxchgOp::build( + &mut ctx, + pointer, + value, + value, + u32_ty.into(), + success.clone(), + failure.clone(), + AtomicScope::Device, + ) + .verify(&ctx) + .is_ok(); + assert_eq!( + actual, expected, + "unexpected cmpxchg ordering result for success={success:?}, failure={failure:?}" + ); + } + } +} diff --git a/crates/dialect-nvvm/tests/ops_test/carriers.rs b/crates/dialect-nvvm/tests/ops_test/carriers.rs new file mode 100644 index 0000000000..9555f85cab --- /dev/null +++ b/crates/dialect-nvvm/tests/ops_test/carriers.rs @@ -0,0 +1,923 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use dialect_mir::types::MirPtrType; +use dialect_nvvm::ops::{ + AssertFailOp, CvtaGenericToSharedOffsetOp, VprintfOp, WgmmaMakeSmemDescOp, WgmmaMaxPendingAttr, + WgmmaMmaGroupM64N64K16F32Bf16Op, WgmmaMmaGroupValuesM64N64K8F32Tf32Op, + WgmmaMmaGroupValuesM64N64K16F32Bf16Op, WgmmaMmaGroupValuesM64N64K16F32F16Op, + WgmmaMmaGroupValuesM64N128K16F32Bf16Op, WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op, + WgmmaMmaLoopValuesM64N64K16F32Bf16Op, WgmmaMmaLoopValuesM64N64K16F32F16Op, + WgmmaMmaM64N64K8F32Tf32Op, WgmmaMmaM64N64K16F32Bf16Op, WgmmaMmaM64N64K16F32F16Op, + WgmmaMmaM64N128K16F32Bf16Op, WgmmaMmaPipelineValuesM64N64K16F32Bf16Op, +}; + +use pliron::{ + basic_block::BasicBlock, + builtin::types::{FP32Type, IntegerType, Signedness}, + common_traits::Verify, + context::Context, + op::Op, + operation::Operation, +}; + +#[test] +fn handwritten_ffi_and_wgmma_carriers_verify_exact_shapes() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + + let u8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let f32_ty = FP32Type::get(&ctx); + let pointer_ty = MirPtrType::get_generic(&mut ctx, u8_ty.into(), false); + let accumulator_pointer_ty = MirPtrType::get_generic(&mut ctx, u8_ty.into(), true); + let global_pointer_ty = MirPtrType::get_global(&mut ctx, u8_ty.into(), false); + let mutable_global_pointer_ty = MirPtrType::get_global(&mut ctx, u8_ty.into(), true); + let block = BasicBlock::new( + &mut ctx, + None, + vec![ + pointer_ty.into(), + accumulator_pointer_ty.into(), + global_pointer_ty.into(), + mutable_global_pointer_ty.into(), + u32_ty.into(), + u64_ty.into(), + f32_ty.into(), + ], + ); + let pointer = block.deref(&ctx).get_argument(0); + let accumulator_pointer = block.deref(&ctx).get_argument(1); + let global_pointer = block.deref(&ctx).get_argument(2); + let mutable_global_pointer = block.deref(&ctx).get_argument(3); + let u32_value = block.deref(&ctx).get_argument(4); + let u64_value = block.deref(&ctx).get_argument(5); + let f32_value = block.deref(&ctx).get_argument(6); + + let vprintf = VprintfOp::build(&mut ctx, pointer, pointer); + assert!(VprintfOp::new(vprintf).verify(&ctx).is_ok()); + + let i64_signless_ty = IntegerType::get(&ctx, 64, Signedness::Signless); + let assertfail = AssertFailOp::build(&mut ctx, pointer, pointer, u32_value, pointer, u64_value); + assert!(AssertFailOp::new(assertfail).verify(&ctx).is_ok()); + for (operands, results) in [ + // message must be a MIR pointer + ( + vec![u32_value, pointer, u32_value, pointer, u64_value], + vec![], + ), + // line must be a 32-bit integer + ( + vec![pointer, pointer, u64_value, pointer, u64_value], + vec![], + ), + // char size must be a 64-bit integer + ( + vec![pointer, pointer, u32_value, pointer, u32_value], + vec![], + ), + // wrong operand count + (vec![pointer, pointer, u32_value], vec![]), + // must have no results + ( + vec![pointer, pointer, u32_value, pointer, u64_value], + vec![i64_signless_ty.into()], + ), + ] { + let invalid = Operation::new( + &mut ctx, + AssertFailOp::get_concrete_op_info(), + results, + operands, + vec![], + 0, + ); + assert!(AssertFailOp::new(invalid).verify(&ctx).is_err()); + } + let bad_vprintf = Operation::new( + &mut ctx, + VprintfOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![pointer, u32_value], + vec![], + 0, + ); + assert!(VprintfOp::new(bad_vprintf).verify(&ctx).is_err()); + + let descriptor = Operation::new( + &mut ctx, + WgmmaMakeSmemDescOp::get_concrete_op_info(), + vec![u64_ty.into()], + vec![pointer], + vec![], + 0, + ); + assert!(WgmmaMakeSmemDescOp::new(descriptor).verify(&ctx).is_ok()); + for (operands, results) in [ + (vec![u32_value], vec![u64_ty.into()]), + (vec![global_pointer], vec![u64_ty.into()]), + (vec![pointer], vec![u32_ty.into()]), + (vec![], vec![u64_ty.into()]), + ] { + let invalid = Operation::new( + &mut ctx, + WgmmaMakeSmemDescOp::get_concrete_op_info(), + results, + operands, + vec![], + 0, + ); + assert!(WgmmaMakeSmemDescOp::new(invalid).verify(&ctx).is_err()); + } + + let cvta = Operation::new( + &mut ctx, + CvtaGenericToSharedOffsetOp::get_concrete_op_info(), + vec![u64_ty.into()], + vec![pointer], + vec![], + 0, + ); + assert!(CvtaGenericToSharedOffsetOp::new(cvta).verify(&ctx).is_ok()); + for (operands, results) in [ + // Operand must be a MIR pointer. + (vec![u32_value], vec![u64_ty.into()]), + // Operand must point to generic or shared memory. + (vec![global_pointer], vec![u64_ty.into()]), + // Result must be u64. + (vec![pointer], vec![u32_ty.into()]), + // Exact arity is required. + (vec![], vec![u64_ty.into()]), + ] { + let invalid = Operation::new( + &mut ctx, + CvtaGenericToSharedOffsetOp::get_concrete_op_info(), + results, + operands, + vec![], + 0, + ); + assert!( + CvtaGenericToSharedOffsetOp::new(invalid) + .verify(&ctx) + .is_err() + ); + } + + let mma = Operation::new( + &mut ctx, + WgmmaMmaM64N64K16F32Bf16Op::get_concrete_op_info(), + vec![], + vec![accumulator_pointer, u64_value, u64_value], + vec![], + 0, + ); + assert!(WgmmaMmaM64N64K16F32Bf16Op::new(mma).verify(&ctx).is_ok()); + for (operands, results) in [ + // Accumulator must be mutable. + (vec![pointer, u64_value, u64_value], vec![]), + // Accumulator must use generic address space. + (vec![mutable_global_pointer, u64_value, u64_value], vec![]), + // Accumulator must be a MIR pointer. + (vec![u32_value, u64_value, u64_value], vec![]), + // Descriptors must be u64. + (vec![accumulator_pointer, u32_value, u64_value], vec![]), + // Exact arity is required. + (vec![accumulator_pointer, u64_value], vec![]), + // No results are permitted. + ( + vec![accumulator_pointer, u64_value, u64_value], + vec![u32_ty.into()], + ), + ] { + let invalid = Operation::new( + &mut ctx, + WgmmaMmaM64N64K16F32Bf16Op::get_concrete_op_info(), + results, + operands, + vec![], + 0, + ); + assert!( + WgmmaMmaM64N64K16F32Bf16Op::new(invalid) + .verify(&ctx) + .is_err() + ); + } + + let group = WgmmaMmaGroupM64N64K16F32Bf16Op::build( + &mut ctx, + accumulator_pointer, + vec![u64_value, u64_value, u64_value, u64_value], + ); + assert!( + WgmmaMmaGroupM64N64K16F32Bf16Op::new(group) + .verify(&ctx) + .is_ok() + ); + for (operands, results) in [ + // Missing descriptor. + (vec![accumulator_pointer, u64_value], vec![]), + // Incomplete descriptor pair. + ( + vec![accumulator_pointer, u64_value, u64_value, u64_value], + vec![], + ), + // Accumulator must be mutable. + (vec![pointer, u64_value, u64_value], vec![]), + // Accumulator must use generic address space. + (vec![mutable_global_pointer, u64_value, u64_value], vec![]), + // Accumulator must be a MIR pointer. + (vec![u32_value, u64_value, u64_value], vec![]), + // Descriptors must be u64. + (vec![accumulator_pointer, u64_value, u32_value], vec![]), + // No results are permitted. + ( + vec![accumulator_pointer, u64_value, u64_value], + vec![u32_ty.into()], + ), + ] { + let invalid = Operation::new( + &mut ctx, + WgmmaMmaGroupM64N64K16F32Bf16Op::get_concrete_op_info(), + results, + operands, + vec![], + 0, + ); + assert!( + WgmmaMmaGroupM64N64K16F32Bf16Op::new(invalid) + .verify(&ctx) + .is_err() + ); + } + + let value_group = WgmmaMmaGroupValuesM64N64K16F32Bf16Op::build( + &mut ctx, + vec![f32_value; 32], + vec![u64_value, u64_value, u64_value, u64_value], + ); + { + let value_group_ref = value_group.deref(&ctx); + assert_eq!(value_group_ref.get_num_operands(), 36); + assert_eq!(value_group_ref.get_num_results(), 32); + } + assert!( + WgmmaMmaGroupValuesM64N64K16F32Bf16Op::new(value_group) + .verify(&ctx) + .is_ok() + ); + + let too_few_accumulators = WgmmaMmaGroupValuesM64N64K16F32Bf16Op::build( + &mut ctx, + vec![f32_value; 31], + vec![u64_value, u64_value], + ); + assert!( + WgmmaMmaGroupValuesM64N64K16F32Bf16Op::new(too_few_accumulators) + .verify(&ctx) + .is_err() + ); + + let missing_descriptors = + WgmmaMmaGroupValuesM64N64K16F32Bf16Op::build(&mut ctx, vec![f32_value; 32], vec![]); + assert!( + WgmmaMmaGroupValuesM64N64K16F32Bf16Op::new(missing_descriptors) + .verify(&ctx) + .is_err() + ); + + let incomplete_descriptor_pair = WgmmaMmaGroupValuesM64N64K16F32Bf16Op::build( + &mut ctx, + vec![f32_value; 32], + vec![u64_value, u64_value, u64_value], + ); + assert!( + WgmmaMmaGroupValuesM64N64K16F32Bf16Op::new(incomplete_descriptor_pair) + .verify(&ctx) + .is_err() + ); + + let mut wrong_accumulator_operands = vec![f32_value; 32]; + wrong_accumulator_operands[0] = u32_value; + wrong_accumulator_operands.extend([u64_value, u64_value]); + let wrong_accumulator = Operation::new( + &mut ctx, + WgmmaMmaGroupValuesM64N64K16F32Bf16Op::get_concrete_op_info(), + vec![f32_ty.into(); 32], + wrong_accumulator_operands, + vec![], + 0, + ); + assert!( + WgmmaMmaGroupValuesM64N64K16F32Bf16Op::new(wrong_accumulator) + .verify(&ctx) + .is_err() + ); + + let mut wrong_descriptor_operands = vec![f32_value; 32]; + wrong_descriptor_operands.extend([u32_value, u64_value]); + let wrong_descriptor = Operation::new( + &mut ctx, + WgmmaMmaGroupValuesM64N64K16F32Bf16Op::get_concrete_op_info(), + vec![f32_ty.into(); 32], + wrong_descriptor_operands, + vec![], + 0, + ); + assert!( + WgmmaMmaGroupValuesM64N64K16F32Bf16Op::new(wrong_descriptor) + .verify(&ctx) + .is_err() + ); + + let mut valid_value_operands = vec![f32_value; 32]; + valid_value_operands.extend([u64_value, u64_value]); + + let wrong_result_count = Operation::new( + &mut ctx, + WgmmaMmaGroupValuesM64N64K16F32Bf16Op::get_concrete_op_info(), + vec![f32_ty.into(); 31], + valid_value_operands.clone(), + vec![], + 0, + ); + assert!( + WgmmaMmaGroupValuesM64N64K16F32Bf16Op::new(wrong_result_count) + .verify(&ctx) + .is_err() + ); + + let mut wrong_result_types = vec![f32_ty.into(); 32]; + wrong_result_types[0] = u32_ty.into(); + let wrong_result_type = Operation::new( + &mut ctx, + WgmmaMmaGroupValuesM64N64K16F32Bf16Op::get_concrete_op_info(), + wrong_result_types, + valid_value_operands, + vec![], + 0, + ); + assert!( + WgmmaMmaGroupValuesM64N64K16F32Bf16Op::new(wrong_result_type) + .verify(&ctx) + .is_err() + ); + + let wide_mma = Operation::new( + &mut ctx, + WgmmaMmaM64N128K16F32Bf16Op::get_concrete_op_info(), + vec![], + vec![accumulator_pointer, u64_value, u64_value], + vec![], + 0, + ); + assert!( + WgmmaMmaM64N128K16F32Bf16Op::new(wide_mma) + .verify(&ctx) + .is_ok() + ); + + let wide_value_group = WgmmaMmaGroupValuesM64N128K16F32Bf16Op::build( + &mut ctx, + vec![f32_value; 64], + vec![u64_value, u64_value], + ); + { + let group_ref = wide_value_group.deref(&ctx); + assert_eq!(group_ref.get_num_operands(), 66); + assert_eq!(group_ref.get_num_results(), 64); + } + assert!( + WgmmaMmaGroupValuesM64N128K16F32Bf16Op::new(wide_value_group) + .verify(&ctx) + .is_ok() + ); + + let wide_too_few = WgmmaMmaGroupValuesM64N128K16F32Bf16Op::build( + &mut ctx, + vec![f32_value; 63], + vec![u64_value, u64_value], + ); + assert!( + WgmmaMmaGroupValuesM64N128K16F32Bf16Op::new(wide_too_few) + .verify(&ctx) + .is_err() + ); + + let wide_incomplete_pair = WgmmaMmaGroupValuesM64N128K16F32Bf16Op::build( + &mut ctx, + vec![f32_value; 64], + vec![u64_value, u64_value, u64_value], + ); + assert!( + WgmmaMmaGroupValuesM64N128K16F32Bf16Op::new(wide_incomplete_pair) + .verify(&ctx) + .is_err() + ); + + let f16_mma = Operation::new( + &mut ctx, + WgmmaMmaM64N64K16F32F16Op::get_concrete_op_info(), + vec![], + vec![accumulator_pointer, u64_value, u64_value], + vec![], + 0, + ); + assert!(WgmmaMmaM64N64K16F32F16Op::new(f16_mma).verify(&ctx).is_ok()); + for operands in [ + vec![pointer, u64_value, u64_value], + vec![mutable_global_pointer, u64_value, u64_value], + vec![accumulator_pointer, u32_value, u64_value], + ] { + let invalid = Operation::new( + &mut ctx, + WgmmaMmaM64N64K16F32F16Op::get_concrete_op_info(), + vec![], + operands, + vec![], + 0, + ); + assert!( + WgmmaMmaM64N64K16F32F16Op::new(invalid) + .verify(&ctx) + .is_err() + ); + } + + let f16_value_group = WgmmaMmaGroupValuesM64N64K16F32F16Op::build( + &mut ctx, + vec![f32_value; 32], + vec![u64_value, u64_value], + ); + assert!( + WgmmaMmaGroupValuesM64N64K16F32F16Op::new(f16_value_group) + .verify(&ctx) + .is_ok() + ); + + let f16_too_few_accumulators = WgmmaMmaGroupValuesM64N64K16F32F16Op::build( + &mut ctx, + vec![f32_value; 31], + vec![u64_value, u64_value], + ); + assert!( + WgmmaMmaGroupValuesM64N64K16F32F16Op::new(f16_too_few_accumulators) + .verify(&ctx) + .is_err() + ); + + let f16_incomplete_descriptor_pair = WgmmaMmaGroupValuesM64N64K16F32F16Op::build( + &mut ctx, + vec![f32_value; 32], + vec![u64_value, u64_value, u64_value], + ); + assert!( + WgmmaMmaGroupValuesM64N64K16F32F16Op::new(f16_incomplete_descriptor_pair) + .verify(&ctx) + .is_err() + ); + + let mut f16_wrong_accumulator_operands = vec![f32_value; 32]; + f16_wrong_accumulator_operands[0] = u32_value; + f16_wrong_accumulator_operands.extend([u64_value, u64_value]); + let f16_wrong_accumulator = Operation::new( + &mut ctx, + WgmmaMmaGroupValuesM64N64K16F32F16Op::get_concrete_op_info(), + vec![f32_ty.into(); 32], + f16_wrong_accumulator_operands, + vec![], + 0, + ); + assert!( + WgmmaMmaGroupValuesM64N64K16F32F16Op::new(f16_wrong_accumulator) + .verify(&ctx) + .is_err() + ); + + let tf32_mma = Operation::new( + &mut ctx, + WgmmaMmaM64N64K8F32Tf32Op::get_concrete_op_info(), + vec![], + vec![accumulator_pointer, u64_value, u64_value], + vec![], + 0, + ); + assert!( + WgmmaMmaM64N64K8F32Tf32Op::new(tf32_mma) + .verify(&ctx) + .is_ok() + ); + for operands in [ + vec![pointer, u64_value, u64_value], + vec![mutable_global_pointer, u64_value, u64_value], + vec![accumulator_pointer, u32_value, u64_value], + ] { + let invalid = Operation::new( + &mut ctx, + WgmmaMmaM64N64K8F32Tf32Op::get_concrete_op_info(), + vec![], + operands, + vec![], + 0, + ); + assert!( + WgmmaMmaM64N64K8F32Tf32Op::new(invalid) + .verify(&ctx) + .is_err() + ); + } + + let tf32_value_group = WgmmaMmaGroupValuesM64N64K8F32Tf32Op::build( + &mut ctx, + vec![f32_value; 32], + vec![u64_value, u64_value], + ); + assert!( + WgmmaMmaGroupValuesM64N64K8F32Tf32Op::new(tf32_value_group) + .verify(&ctx) + .is_ok() + ); + + let tf32_too_few_accumulators = WgmmaMmaGroupValuesM64N64K8F32Tf32Op::build( + &mut ctx, + vec![f32_value; 31], + vec![u64_value, u64_value], + ); + assert!( + WgmmaMmaGroupValuesM64N64K8F32Tf32Op::new(tf32_too_few_accumulators) + .verify(&ctx) + .is_err() + ); + + let tf32_incomplete_descriptor_pair = WgmmaMmaGroupValuesM64N64K8F32Tf32Op::build( + &mut ctx, + vec![f32_value; 32], + vec![u64_value, u64_value, u64_value], + ); + assert!( + WgmmaMmaGroupValuesM64N64K8F32Tf32Op::new(tf32_incomplete_descriptor_pair) + .verify(&ctx) + .is_err() + ); + + let mut tf32_wrong_accumulator_operands = vec![f32_value; 32]; + tf32_wrong_accumulator_operands[0] = u32_value; + tf32_wrong_accumulator_operands.extend([u64_value, u64_value]); + let tf32_wrong_accumulator = Operation::new( + &mut ctx, + WgmmaMmaGroupValuesM64N64K8F32Tf32Op::get_concrete_op_info(), + vec![f32_ty.into(); 32], + tf32_wrong_accumulator_operands, + vec![], + 0, + ); + assert!( + WgmmaMmaGroupValuesM64N64K8F32Tf32Op::new(tf32_wrong_accumulator) + .verify(&ctx) + .is_err() + ); + + let loop_group = WgmmaMmaLoopValuesM64N64K16F32Bf16Op::build( + &mut ctx, + vec![f32_value; 32], + u64_value, + u64_value, + u64_value, + u64_value, + u64_value, + ); + { + let loop_group_ref = loop_group.deref(&ctx); + assert_eq!(loop_group_ref.get_num_operands(), 37); + assert_eq!(loop_group_ref.get_num_results(), 32); + } + assert!( + WgmmaMmaLoopValuesM64N64K16F32Bf16Op::new(loop_group) + .verify(&ctx) + .is_ok() + ); + + let f16_loop_group = WgmmaMmaLoopValuesM64N64K16F32F16Op::build( + &mut ctx, + vec![f32_value; 32], + u64_value, + u64_value, + u64_value, + u64_value, + u64_value, + ); + { + let loop_group_ref = f16_loop_group.deref(&ctx); + assert_eq!(loop_group_ref.get_num_operands(), 37); + assert_eq!(loop_group_ref.get_num_results(), 32); + } + assert!( + WgmmaMmaLoopValuesM64N64K16F32F16Op::new(f16_loop_group) + .verify(&ctx) + .is_ok() + ); + + let mut f16_wrong_loop_control_operands = vec![f32_value; 32]; + f16_wrong_loop_control_operands.extend([u64_value, u64_value, u32_value, u64_value, u64_value]); + let f16_wrong_loop_control = Operation::new( + &mut ctx, + WgmmaMmaLoopValuesM64N64K16F32F16Op::get_concrete_op_info(), + vec![f32_ty.into(); 32], + f16_wrong_loop_control_operands, + vec![], + 0, + ); + assert!( + WgmmaMmaLoopValuesM64N64K16F32F16Op::new(f16_wrong_loop_control) + .verify(&ctx) + .is_err() + ); + + let too_few_loop_accumulators = WgmmaMmaLoopValuesM64N64K16F32Bf16Op::build( + &mut ctx, + vec![f32_value; 31], + u64_value, + u64_value, + u64_value, + u64_value, + u64_value, + ); + assert!( + WgmmaMmaLoopValuesM64N64K16F32Bf16Op::new(too_few_loop_accumulators) + .verify(&ctx) + .is_err() + ); + + let mut wrong_loop_control_operands = vec![f32_value; 32]; + wrong_loop_control_operands.extend([u64_value, u64_value, u32_value, u64_value, u64_value]); + let wrong_loop_control = Operation::new( + &mut ctx, + WgmmaMmaLoopValuesM64N64K16F32Bf16Op::get_concrete_op_info(), + vec![f32_ty.into(); 32], + wrong_loop_control_operands, + vec![], + 0, + ); + assert!( + WgmmaMmaLoopValuesM64N64K16F32Bf16Op::new(wrong_loop_control) + .verify(&ctx) + .is_err() + ); + + let mut valid_loop_operands = vec![f32_value; 32]; + valid_loop_operands.extend([u64_value; 5]); + + let wrong_loop_result_count = Operation::new( + &mut ctx, + WgmmaMmaLoopValuesM64N64K16F32Bf16Op::get_concrete_op_info(), + vec![f32_ty.into(); 31], + valid_loop_operands.clone(), + vec![], + 0, + ); + assert!( + WgmmaMmaLoopValuesM64N64K16F32Bf16Op::new(wrong_loop_result_count) + .verify(&ctx) + .is_err() + ); + + let mut wrong_loop_result_types = vec![f32_ty.into(); 32]; + wrong_loop_result_types[0] = u32_ty.into(); + let wrong_loop_result_type = Operation::new( + &mut ctx, + WgmmaMmaLoopValuesM64N64K16F32Bf16Op::get_concrete_op_info(), + wrong_loop_result_types, + valid_loop_operands, + vec![], + 0, + ); + assert!( + WgmmaMmaLoopValuesM64N64K16F32Bf16Op::new(wrong_loop_result_type) + .verify(&ctx) + .is_err() + ); + + let counted_pipeline_group = WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::build( + &mut ctx, + vec![f32_value; 64], + vec![u64_value; 4], + vec![u64_value; 4], + u64_value, + 1, + ); + { + let counted_pipeline_ref = counted_pipeline_group.deref(&ctx); + assert_eq!(counted_pipeline_ref.get_num_operands(), 73); + assert_eq!(counted_pipeline_ref.get_num_results(), 64); + } + let counted_pipeline = + WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::new(counted_pipeline_group); + assert_eq!(counted_pipeline.max_pending_groups(&ctx), Some(1)); + assert!(counted_pipeline.verify(&ctx).is_ok()); + + let three_slot_counted_pipeline = WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::build( + &mut ctx, + vec![f32_value; 96], + vec![u64_value; 6], + vec![u64_value; 6], + u64_value, + 2, + ); + { + let counted_pipeline_ref = three_slot_counted_pipeline.deref(&ctx); + assert_eq!(counted_pipeline_ref.get_num_operands(), 109); + assert_eq!(counted_pipeline_ref.get_num_results(), 96); + } + let three_slot_counted_pipeline = + WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::new(three_slot_counted_pipeline); + assert_eq!( + three_slot_counted_pipeline.max_pending_groups(&ctx), + Some(2) + ); + assert!(three_slot_counted_pipeline.verify(&ctx).is_ok()); + + let too_many_accumulators_for_wait_one = WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::build( + &mut ctx, + vec![f32_value; 96], + vec![u64_value; 4], + vec![u64_value; 4], + u64_value, + 1, + ); + assert!( + WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::new(too_many_accumulators_for_wait_one) + .verify(&ctx) + .is_err() + ); + + let too_few_accumulators_for_wait_two = WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::build( + &mut ctx, + vec![f32_value; 64], + vec![u64_value; 6], + vec![u64_value; 6], + u64_value, + 2, + ); + assert!( + WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::new(too_few_accumulators_for_wait_two) + .verify(&ctx) + .is_err() + ); + + let unsupported_wait_three = WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::build( + &mut ctx, + vec![f32_value; 128], + vec![u64_value; 8], + vec![u64_value; 8], + u64_value, + 3, + ); + assert!( + WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::new(unsupported_wait_three) + .verify(&ctx) + .is_err() + ); + + let wrong_counted_pipeline_control = WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::build( + &mut ctx, + vec![f32_value; 64], + vec![u64_value; 4], + vec![u64_value, u64_value, u32_value, u64_value], + u64_value, + 1, + ); + assert!( + WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::new(wrong_counted_pipeline_control) + .verify(&ctx) + .is_err() + ); + + let mut valid_counted_pipeline_operands = vec![f32_value; 64]; + valid_counted_pipeline_operands.extend([u64_value; 9]); + let wrong_counted_pipeline_result_count = Operation::new( + &mut ctx, + WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::get_concrete_op_info(), + vec![f32_ty.into(); 63], + valid_counted_pipeline_operands.clone(), + vec![], + 0, + ); + let wrong_counted_pipeline_result_count = + WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::new(wrong_counted_pipeline_result_count); + wrong_counted_pipeline_result_count + .set_attr_counted_max_pending_groups(&ctx, WgmmaMaxPendingAttr(1)); + assert!(wrong_counted_pipeline_result_count.verify(&ctx).is_err()); + + let mut wrong_counted_pipeline_result_types = vec![f32_ty.into(); 64]; + wrong_counted_pipeline_result_types[0] = u32_ty.into(); + let wrong_counted_pipeline_result_type = Operation::new( + &mut ctx, + WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::get_concrete_op_info(), + wrong_counted_pipeline_result_types, + valid_counted_pipeline_operands, + vec![], + 0, + ); + let wrong_counted_pipeline_result_type = + WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::new(wrong_counted_pipeline_result_type); + wrong_counted_pipeline_result_type + .set_attr_counted_max_pending_groups(&ctx, WgmmaMaxPendingAttr(1)); + assert!(wrong_counted_pipeline_result_type.verify(&ctx).is_err()); + + let pipeline_group = WgmmaMmaPipelineValuesM64N64K16F32Bf16Op::build( + &mut ctx, + vec![f32_value; 64], + vec![u64_value; 8], + 1, + ); + { + let pipeline_ref = pipeline_group.deref(&ctx); + assert_eq!(pipeline_ref.get_num_operands(), 72); + assert_eq!(pipeline_ref.get_num_results(), 64); + } + let pipeline = WgmmaMmaPipelineValuesM64N64K16F32Bf16Op::new(pipeline_group); + assert_eq!(pipeline.max_pending_groups(&ctx), Some(1)); + assert!(pipeline.verify(&ctx).is_ok()); + + let too_few_pipeline_slots = WgmmaMmaPipelineValuesM64N64K16F32Bf16Op::build( + &mut ctx, + vec![f32_value; 32], + vec![u64_value; 4], + 1, + ); + assert!( + WgmmaMmaPipelineValuesM64N64K16F32Bf16Op::new(too_few_pipeline_slots) + .verify(&ctx) + .is_err() + ); + + assert!(WgmmaMaxPendingAttr(1).verify(&ctx).is_ok()); + assert!(WgmmaMaxPendingAttr(7).verify(&ctx).is_ok()); + assert!(WgmmaMaxPendingAttr(0).verify(&ctx).is_err()); + assert!(WgmmaMaxPendingAttr(8).verify(&ctx).is_err()); + + let zero_pending_pipeline = WgmmaMmaPipelineValuesM64N64K16F32Bf16Op::build( + &mut ctx, + vec![f32_value; 32], + vec![u64_value; 2], + 0, + ); + let zero_pending_error = WgmmaMmaPipelineValuesM64N64K16F32Bf16Op::new(zero_pending_pipeline) + .verify(&ctx) + .unwrap_err(); + assert!( + zero_pending_error + .err + .to_string() + .contains("nvvm.wgmma_max_pending_groups must be in 1..=7, got 0"), + "unexpected zero-pending error: {}", + zero_pending_error.err + ); + + let excess_pending_pipeline = WgmmaMmaPipelineValuesM64N64K16F32Bf16Op::build( + &mut ctx, + vec![f32_value; 32], + vec![u64_value; 2], + 8, + ); + let excess_pending_error = + WgmmaMmaPipelineValuesM64N64K16F32Bf16Op::new(excess_pending_pipeline) + .verify(&ctx) + .unwrap_err(); + assert!( + excess_pending_error + .err + .to_string() + .contains("nvvm.wgmma_max_pending_groups must be in 1..=7, got 8"), + "unexpected excess-pending error: {}", + excess_pending_error.err + ); + + let mut missing_attr_operands = vec![f32_value; 64]; + missing_attr_operands.extend([u64_value; 8]); + let missing_attr_pipeline = Operation::new( + &mut ctx, + WgmmaMmaPipelineValuesM64N64K16F32Bf16Op::get_concrete_op_info(), + vec![f32_ty.into(); 64], + missing_attr_operands, + vec![], + 0, + ); + let missing_attr_error = WgmmaMmaPipelineValuesM64N64K16F32Bf16Op::new(missing_attr_pipeline) + .verify(&ctx) + .unwrap_err(); + assert!( + missing_attr_error + .err + .to_string() + .contains("requires an nvvm.wgmma_max_pending_groups attribute"), + "unexpected missing-attribute error: {}", + missing_attr_error.err + ); +} diff --git a/crates/dialect-nvvm/tests/ops_test/cluster_and_memory.rs b/crates/dialect-nvvm/tests/ops_test/cluster_and_memory.rs new file mode 100644 index 0000000000..db490524aa --- /dev/null +++ b/crates/dialect-nvvm/tests/ops_test/cluster_and_memory.rs @@ -0,0 +1,761 @@ +/* + * 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::{ + ClusterBarrierModeAttr, ClusterBarrierOp, CpAsyncCa4Op, CpAsyncCaZfill4Op, + CpAsyncMbarrierArriveNoIncOp, CpAsyncMbarrierArriveNoIncSharedOp, CpAsyncMbarrierArriveOp, + CpAsyncMbarrierArriveSharedOp, CpAsyncWaitGroupOp, MbarrierArriveSharedOp, + MbarrierInitSharedOp, MbarrierInvalSharedOp, MbarrierTestWaitSharedOp, ReadPtxSregClusterIdxOp, + ReadPtxSregNclusterIdOp, ScalarArithmeticFormatAttr, ScalarArithmeticOp, + ScalarArithmeticOperationAttr, ScalarArithmeticRoundingAttr, ScalarArithmeticSaturationAttr, + ScalarArithmeticSubnormalAttr, ScalarConversionOp, ScalarConversionRoundingAttr, + ScalarConversionSaturationAttr, Tcgen05AllocOp, Tcgen05CommitMulticastCg2Op, + Tcgen05Ld16x32bx2X1RawOp, Tcgen05Ld16x256bPureOp, Tcgen05MmaF16Op, +}; + +use pliron::{ + basic_block::BasicBlock, + builtin::types::{FP32Type, FP64Type, IntegerType, Signedness}, + common_traits::Verify, + context::Context, + op::{Op, verify_op}, + operation::Operation, + r#type::Typed, +}; + +#[test] +fn cluster_grid_compatibility_ops_keep_names_and_i32_shape() { + use pliron::builtin::types::{IntegerType, Signedness}; + use pliron::common_traits::Verify; + use pliron::context::Context; + use pliron::op::Op; + use pliron::operation::Operation; + + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + assert_eq!( + ReadPtxSregClusterIdxOp::get_opid_static().to_string(), + "nvvm.read_ptx_sreg_cluster_idx" + ); + assert_eq!( + ReadPtxSregNclusterIdOp::get_opid_static().to_string(), + "nvvm.read_ptx_sreg_nclusterid" + ); + + let i32_type = IntegerType::get(&ctx, 32, Signedness::Signless); + for op_info in [ + ReadPtxSregClusterIdxOp::get_concrete_op_info(), + ReadPtxSregNclusterIdOp::get_concrete_op_info(), + ] { + let op = Operation::new(&mut ctx, op_info, vec![i32_type.into()], vec![], vec![], 0); + assert!(op.deref(&ctx).verify(&ctx).is_ok()); + } +} + +#[test] +fn mapa_shared_cluster_requires_and_preserves_raw_pointer_kind() { + use dialect_mir::types::{MirPointerKind, MirPtrType, address_space}; + use dialect_nvvm::ops::MapaSharedClusterOp; + use pliron::{ + basic_block::BasicBlock, + builtin::types::{IntegerType, Signedness}, + common_traits::Verify, + context::Context, + op::Op, + operation::Operation, + }; + + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let raw_const_ty = + MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), false, MirPointerKind::RawConst); + let raw_mut_ty = + MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::RawMut); + let erased_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), false); + let block = BasicBlock::new( + &mut ctx, + None, + vec![ + raw_const_ty.into(), + raw_mut_ty.into(), + erased_ty.into(), + i32_ty.into(), + ], + ); + let raw_const = block.deref(&ctx).get_argument(0); + let raw_mut = block.deref(&ctx).get_argument(1); + let erased = block.deref(&ctx).get_argument(2); + let rank = block.deref(&ctx).get_argument(3); + + for (source, expected_kind) in [ + (raw_const, MirPointerKind::RawConst), + (raw_mut, MirPointerKind::RawMut), + ] { + let operation = MapaSharedClusterOp::build(&mut ctx, source, rank); + let result_ty = operation.deref(&ctx).get_result(0).get_type(&ctx); + let result_ty = result_ty.deref(&ctx); + let result_ptr = result_ty + .downcast_ref::() + .expect("mapa result must be a MIR pointer"); + assert_eq!(result_ptr.pointer_kind(), expected_kind); + assert_eq!(result_ptr.address_space(), address_space::CLUSTER_SHARED); + assert!(MapaSharedClusterOp::new(operation).verify(&ctx).is_ok()); + } + + let erased_result = MirPtrType::get_with_kind( + &mut ctx, + i32_ty.into(), + true, + address_space::CLUSTER_SHARED, + MirPointerKind::Erased, + ); + let erasing = Operation::new( + &mut ctx, + MapaSharedClusterOp::get_concrete_op_info(), + vec![erased_result.into()], + vec![raw_mut, rank], + vec![], + 0, + ); + assert!( + MapaSharedClusterOp::new(erasing).verify(&ctx).is_err(), + "cluster address mapping must preserve the source raw-pointer kind" + ); + + let erased_source = MapaSharedClusterOp::build(&mut ctx, erased, rank); + assert!( + MapaSharedClusterOp::new(erased_source) + .verify(&ctx) + .is_err(), + "an Erased data/function carrier cannot enter the raw-pointer intrinsic ABI" + ); + + let invented_unique = MirPtrType::get_with_kind( + &mut ctx, + i32_ty.into(), + true, + address_space::CLUSTER_SHARED, + MirPointerKind::UniqueRef, + ); + let laundering = Operation::new( + &mut ctx, + MapaSharedClusterOp::get_concrete_op_info(), + vec![invented_unique.into()], + vec![raw_mut, rank], + vec![], + 0, + ); + assert!( + MapaSharedClusterOp::new(laundering).verify(&ctx).is_err(), + "cluster address mapping must not turn RawMut into UniqueRef" + ); + + let invented_shared = MirPtrType::get_with_kind( + &mut ctx, + i32_ty.into(), + false, + address_space::CLUSTER_SHARED, + MirPointerKind::SharedRef, + ); + let laundering = Operation::new( + &mut ctx, + MapaSharedClusterOp::get_concrete_op_info(), + vec![invented_shared.into()], + vec![raw_const, rank], + vec![], + 0, + ); + assert!( + MapaSharedClusterOp::new(laundering).verify(&ctx).is_err(), + "cluster address mapping must not turn RawConst into SharedRef" + ); +} + +#[test] +fn generated_cluster_barrier_requires_one_closed_mode_and_no_values() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + for mode in [ + ClusterBarrierModeAttr::Arrive, + ClusterBarrierModeAttr::ArriveAligned, + ClusterBarrierModeAttr::ArriveRelaxed, + ClusterBarrierModeAttr::ArriveRelaxedAligned, + ClusterBarrierModeAttr::Wait, + ClusterBarrierModeAttr::WaitAligned, + ] { + let op = ClusterBarrierOp::build(&mut ctx, mode); + assert!(verify_op(&ClusterBarrierOp::new(op), &ctx).is_ok()); + } + + let missing_mode = Operation::new( + &mut ctx, + ClusterBarrierOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + assert!(verify_op(&ClusterBarrierOp::new(missing_mode), &ctx).is_err()); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let wrong_shape = Operation::new( + &mut ctx, + ClusterBarrierOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![], + vec![], + 0, + ); + ClusterBarrierOp::new(wrong_shape) + .set_attr_nvvm_cluster_barrier_mode(&ctx, ClusterBarrierModeAttr::Wait); + assert!(verify_op(&ClusterBarrierOp::new(wrong_shape), &ctx).is_err()); +} + +#[test] +fn generated_scalar_conversion_accepts_only_reviewed_f32_to_i32_variants() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + let f32_ty = FP32Type::get(&ctx); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); + let block = BasicBlock::new(&mut ctx, None, vec![f32_ty.into(), i32_ty.into()]); + let f32_value = block.deref(&ctx).get_argument(0); + let i32_value = block.deref(&ctx).get_argument(1); + + for (rounding, saturation) in [ + ( + ScalarConversionRoundingAttr::NearestAway, + ScalarConversionSaturationAttr::None, + ), + ( + ScalarConversionRoundingAttr::NearestAway, + ScalarConversionSaturationAttr::Satfinite, + ), + ( + ScalarConversionRoundingAttr::NearestEven, + ScalarConversionSaturationAttr::None, + ), + ( + ScalarConversionRoundingAttr::NearestEven, + ScalarConversionSaturationAttr::Relu, + ), + ( + ScalarConversionRoundingAttr::NearestEven, + ScalarConversionSaturationAttr::Satfinite, + ), + ( + ScalarConversionRoundingAttr::NearestEven, + ScalarConversionSaturationAttr::ReluSatfinite, + ), + ( + ScalarConversionRoundingAttr::TowardZero, + ScalarConversionSaturationAttr::None, + ), + ( + ScalarConversionRoundingAttr::TowardZero, + ScalarConversionSaturationAttr::Relu, + ), + ( + ScalarConversionRoundingAttr::TowardZero, + ScalarConversionSaturationAttr::Satfinite, + ), + ( + ScalarConversionRoundingAttr::TowardZero, + ScalarConversionSaturationAttr::ReluSatfinite, + ), + ] { + let op = ScalarConversionOp::build(&mut ctx, f32_value, rounding, saturation); + assert!(verify_op(&ScalarConversionOp::new(op), &ctx).is_ok()); + } + + let invalid_variant = ScalarConversionOp::build( + &mut ctx, + f32_value, + ScalarConversionRoundingAttr::NearestAway, + ScalarConversionSaturationAttr::Relu, + ); + assert!(verify_op(&ScalarConversionOp::new(invalid_variant), &ctx).is_err()); + + let wrong_operand = Operation::new( + &mut ctx, + ScalarConversionOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![i32_value], + vec![], + 0, + ); + let wrong_operand = ScalarConversionOp::new(wrong_operand); + wrong_operand + .set_attr_nvvm_scalar_conversion_rounding(&ctx, ScalarConversionRoundingAttr::NearestEven); + wrong_operand + .set_attr_nvvm_scalar_conversion_saturation(&ctx, ScalarConversionSaturationAttr::None); + assert!(verify_op(&wrong_operand, &ctx).is_err()); + + let wrong_result = Operation::new( + &mut ctx, + ScalarConversionOp::get_concrete_op_info(), + vec![i64_ty.into()], + vec![f32_value], + vec![], + 0, + ); + let wrong_result = ScalarConversionOp::new(wrong_result); + wrong_result + .set_attr_nvvm_scalar_conversion_rounding(&ctx, ScalarConversionRoundingAttr::TowardZero); + wrong_result + .set_attr_nvvm_scalar_conversion_saturation(&ctx, ScalarConversionSaturationAttr::None); + assert!(verify_op(&wrong_result, &ctx).is_err()); +} + +#[test] +fn generated_scalar_arithmetic_accepts_only_admitted_shapes_and_types() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + let f32_ty = FP32Type::get(&ctx); + let f64_ty = FP64Type::get(&ctx); + let block = BasicBlock::new(&mut ctx, None, vec![f32_ty.into(), f64_ty.into()]); + let f32_value = block.deref(&ctx).get_argument(0); + let f64_value = block.deref(&ctx).get_argument(1); + + let valid_f32 = ScalarArithmeticOp::build( + &mut ctx, + vec![f32_value, f32_value], + ScalarArithmeticFormatAttr::F32, + ScalarArithmeticOperationAttr::Mul, + ScalarArithmeticRoundingAttr::Rn, + ScalarArithmeticSubnormalAttr::Preserve, + ScalarArithmeticSaturationAttr::None, + ); + assert!(verify_op(&ScalarArithmeticOp::new(valid_f32), &ctx).is_ok()); + + let valid_f64 = ScalarArithmeticOp::build( + &mut ctx, + vec![f64_value, f64_value, f64_value], + ScalarArithmeticFormatAttr::F64, + ScalarArithmeticOperationAttr::Fma, + ScalarArithmeticRoundingAttr::Rz, + ScalarArithmeticSubnormalAttr::Preserve, + ScalarArithmeticSaturationAttr::None, + ); + assert!(verify_op(&ScalarArithmeticOp::new(valid_f64), &ctx).is_ok()); + + let valid_add = ScalarArithmeticOp::build( + &mut ctx, + vec![f32_value, f32_value], + ScalarArithmeticFormatAttr::F32, + ScalarArithmeticOperationAttr::Add, + ScalarArithmeticRoundingAttr::Rp, + ScalarArithmeticSubnormalAttr::Ftz, + ScalarArithmeticSaturationAttr::Sat, + ); + assert!(verify_op(&ScalarArithmeticOp::new(valid_add), &ctx).is_ok()); + + let f64_ftz = ScalarArithmeticOp::build( + &mut ctx, + vec![f64_value, f64_value], + ScalarArithmeticFormatAttr::F64, + ScalarArithmeticOperationAttr::Mul, + ScalarArithmeticRoundingAttr::Rn, + ScalarArithmeticSubnormalAttr::Ftz, + ScalarArithmeticSaturationAttr::None, + ); + assert!(verify_op(&ScalarArithmeticOp::new(f64_ftz), &ctx).is_err()); + + let wrong_arity = ScalarArithmeticOp::build( + &mut ctx, + vec![f32_value, f32_value], + ScalarArithmeticFormatAttr::F32, + ScalarArithmeticOperationAttr::Fma, + ScalarArithmeticRoundingAttr::Rn, + ScalarArithmeticSubnormalAttr::Preserve, + ScalarArithmeticSaturationAttr::None, + ); + assert!(verify_op(&ScalarArithmeticOp::new(wrong_arity), &ctx).is_err()); + + let wrong_type = ScalarArithmeticOp::build( + &mut ctx, + vec![f32_value, f64_value], + ScalarArithmeticFormatAttr::F32, + ScalarArithmeticOperationAttr::Mul, + ScalarArithmeticRoundingAttr::Rn, + ScalarArithmeticSubnormalAttr::Preserve, + ScalarArithmeticSaturationAttr::None, + ); + assert!(verify_op(&ScalarArithmeticOp::new(wrong_type), &ctx).is_err()); + + let missing_attrs = Operation::new( + &mut ctx, + ScalarArithmeticOp::get_concrete_op_info(), + vec![f32_ty.into()], + vec![f32_value, f32_value], + vec![], + 0, + ); + assert!(verify_op(&ScalarArithmeticOp::new(missing_attrs), &ctx).is_err()); +} + +#[test] +fn test_generated_cp_async_accepts_pointer_shapes_and_both_constant_kinds() { + use dialect_mir::ops::MirConstantOp; + use pliron::builtin::{attributes::IntegerAttr, ops::ConstantOp}; + use pliron::utils::apint::APInt; + use std::num::NonZeroUsize; + + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + + let u8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let dst_ty = MirPtrType::get_generic(&mut ctx, u8_ty.into(), false); + let src_ty = MirPtrType::get(&mut ctx, u8_ty.into(), true, address_space::GLOBAL); + let wrong_dst_ty = MirPtrType::get(&mut ctx, u8_ty.into(), true, address_space::GLOBAL); + let block = BasicBlock::new( + &mut ctx, + None, + vec![ + dst_ty.into(), + src_ty.into(), + wrong_dst_ty.into(), + u32_ty.into(), + ], + ); + let dst = block.deref(&ctx).get_argument(0); + let src = block.deref(&ctx).get_argument(1); + let wrong_dst = block.deref(&ctx).get_argument(2); + let dynamic = block.deref(&ctx).get_argument(3); + + let copy = CpAsyncCa4Op::build(&mut ctx, dst, src); + assert!(verify_op(&CpAsyncCa4Op::new(copy), &ctx).is_ok()); + let zfill = CpAsyncCaZfill4Op::build(&mut ctx, dst, src, dynamic); + assert!(verify_op(&CpAsyncCaZfill4Op::new(zfill), &ctx).is_ok()); + let wrong_space = CpAsyncCa4Op::build(&mut ctx, wrong_dst, src); + assert!(verify_op(&CpAsyncCa4Op::new(wrong_space), &ctx).is_err()); + + let value = IntegerAttr::new(u32_ty, APInt::from_u32(0, NonZeroUsize::new(32).unwrap())); + let builtin = ConstantOp::new(&mut ctx, value.clone().into()); + let builtin_value = builtin.get_operation().deref(&ctx).get_result(0); + let builtin_wait = CpAsyncWaitGroupOp::build(&mut ctx, builtin_value); + assert!(verify_op(&CpAsyncWaitGroupOp::new(builtin_wait), &ctx).is_ok()); + + let mir_constant = Operation::new( + &mut ctx, + MirConstantOp::get_concrete_op_info(), + vec![u32_ty.into()], + vec![], + vec![], + 0, + ); + MirConstantOp::new(mir_constant).set_attr_value(&ctx, value); + let mir_value = mir_constant.deref(&ctx).get_result(0); + let mir_wait = CpAsyncWaitGroupOp::build(&mut ctx, mir_value); + assert!(verify_op(&CpAsyncWaitGroupOp::new(mir_wait), &ctx).is_ok()); + + let dynamic_wait = CpAsyncWaitGroupOp::build(&mut ctx, dynamic); + assert!(verify_op(&CpAsyncWaitGroupOp::new(dynamic_wait), &ctx).is_err()); +} + +#[test] +fn generated_tcgen05_verifies_carriers_and_half_split_constants() { + use dialect_mir::ops::MirConstantOp; + use pliron::builtin::{attributes::IntegerAttr, ops::ConstantOp}; + use pliron::utils::apint::APInt; + use std::num::NonZeroUsize; + + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + + let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); + 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 f32_ty = FP32Type::get(&ctx); + let pointer_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), true); + let block = BasicBlock::new( + &mut ctx, + None, + vec![ + pointer_ty.into(), + i1_ty.into(), + i16_ty.into(), + i32_ty.into(), + i64_ty.into(), + f32_ty.into(), + ], + ); + let pointer = block.deref(&ctx).get_argument(0); + let predicate = block.deref(&ctx).get_argument(1); + let mask = block.deref(&ctx).get_argument(2); + let address = block.deref(&ctx).get_argument(3); + let dynamic_offset = block.deref(&ctx).get_argument(4); + + let offset_attr = IntegerAttr::new(i64_ty, APInt::from_i64(16, NonZeroUsize::new(64).unwrap())); + let builtin_offset = ConstantOp::new(&mut ctx, offset_attr.clone().into()) + .get_operation() + .deref(&ctx) + .get_result(0); + let mir_constant = Operation::new( + &mut ctx, + MirConstantOp::get_concrete_op_info(), + vec![i64_ty.into()], + vec![], + vec![], + 0, + ); + MirConstantOp::new(mir_constant).set_attr_value(&ctx, offset_attr); + let mir_offset = mir_constant.deref(&ctx).get_result(0); + + for offset in [builtin_offset, mir_offset] { + let load = Operation::new( + &mut ctx, + Tcgen05Ld16x32bx2X1RawOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![address, offset], + vec![], + 0, + ); + assert!(verify_op(&Tcgen05Ld16x32bx2X1RawOp::new(load), &ctx).is_ok()); + } + + let dynamic = Operation::new( + &mut ctx, + Tcgen05Ld16x32bx2X1RawOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![address, dynamic_offset], + vec![], + 0, + ); + assert!(verify_op(&Tcgen05Ld16x32bx2X1RawOp::new(dynamic), &ctx).is_err()); + + let wrong_offset_type = Operation::new( + &mut ctx, + Tcgen05Ld16x32bx2X1RawOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![address, address], + vec![], + 0, + ); + assert!(verify_op(&Tcgen05Ld16x32bx2X1RawOp::new(wrong_offset_type), &ctx).is_err()); + + let wrong_result_type = Operation::new( + &mut ctx, + Tcgen05Ld16x32bx2X1RawOp::get_concrete_op_info(), + vec![i64_ty.into()], + vec![address, builtin_offset], + vec![], + 0, + ); + assert!(verify_op(&Tcgen05Ld16x32bx2X1RawOp::new(wrong_result_type), &ctx).is_err()); + + let alloc = Operation::new( + &mut ctx, + Tcgen05AllocOp::get_concrete_op_info(), + vec![], + vec![pointer, address], + vec![], + 0, + ); + assert!(verify_op(&Tcgen05AllocOp::new(alloc), &ctx).is_ok()); + + let multicast = Operation::new( + &mut ctx, + Tcgen05CommitMulticastCg2Op::get_concrete_op_info(), + vec![], + vec![pointer, mask], + vec![], + 0, + ); + assert!(verify_op(&Tcgen05CommitMulticastCg2Op::new(multicast), &ctx).is_ok()); + + let mma = Operation::new( + &mut ctx, + Tcgen05MmaF16Op::get_concrete_op_info(), + vec![], + vec![address, dynamic_offset, dynamic_offset, address, predicate], + vec![], + 0, + ); + assert!(verify_op(&Tcgen05MmaF16Op::new(mma), &ctx).is_ok()); + + let pure_load = Operation::new( + &mut ctx, + Tcgen05Ld16x256bPureOp::get_concrete_op_info(), + vec![f32_ty.into(); 4], + vec![address], + vec![], + 0, + ); + assert!(verify_op(&Tcgen05Ld16x256bPureOp::new(pure_load), &ctx).is_ok()); +} + +#[test] +fn generated_cp_async_mbarrier_requires_mutable_generic_or_shared_u64() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let generic_u64 = MirPtrType::get_generic(&mut ctx, u64_ty.into(), true); + let shared_u64 = MirPtrType::get_shared(&mut ctx, u64_ty.into(), true); + let global_u64 = MirPtrType::get_global(&mut ctx, u64_ty.into(), true); + let immutable_u64 = MirPtrType::get_generic(&mut ctx, u64_ty.into(), false); + let generic_u32 = MirPtrType::get_generic(&mut ctx, u32_ty.into(), true); + let block = BasicBlock::new( + &mut ctx, + None, + vec![ + generic_u64.into(), + shared_u64.into(), + global_u64.into(), + immutable_u64.into(), + generic_u32.into(), + u64_ty.into(), + ], + ); + let generic = block.deref(&ctx).get_argument(0); + let shared = block.deref(&ctx).get_argument(1); + let global = block.deref(&ctx).get_argument(2); + let immutable = block.deref(&ctx).get_argument(3); + let wrong_pointee = block.deref(&ctx).get_argument(4); + let scalar = block.deref(&ctx).get_argument(5); + + macro_rules! check_bridge { + ($op:ty) => {{ + for barrier in [generic, shared] { + let valid = <$op>::build(&mut ctx, barrier); + assert!(verify_op(&<$op>::new(valid), &ctx).is_ok()); + } + for barrier in [global, immutable, wrong_pointee, scalar] { + let invalid = <$op>::build(&mut ctx, barrier); + assert!(verify_op(&<$op>::new(invalid), &ctx).is_err()); + } + let wrong_shape = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + vec![u64_ty.into()], + vec![generic], + vec![], + 0, + ); + assert!(verify_op(&<$op>::new(wrong_shape), &ctx).is_err()); + }}; + } + + check_bridge!(CpAsyncMbarrierArriveOp); + check_bridge!(CpAsyncMbarrierArriveSharedOp); + check_bridge!(CpAsyncMbarrierArriveNoIncOp); + check_bridge!(CpAsyncMbarrierArriveNoIncSharedOp); +} + +#[test] +fn generated_mbarrier_builders_and_verifiers_are_closed_over_their_shapes() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + + let u1_ty = IntegerType::get(&ctx, 1, Signedness::Unsigned); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let generic_ptr = MirPtrType::get_generic(&mut ctx, u64_ty.into(), false); + let shared_ptr = MirPtrType::get_shared(&mut ctx, u64_ty.into(), false); + let global_ptr = MirPtrType::get_global(&mut ctx, u64_ty.into(), false); + let block = BasicBlock::new( + &mut ctx, + None, + vec![ + generic_ptr.into(), + shared_ptr.into(), + global_ptr.into(), + i32_ty.into(), + u32_ty.into(), + u64_ty.into(), + ], + ); + let generic = block.deref(&ctx).get_argument(0); + let shared = block.deref(&ctx).get_argument(1); + let global = block.deref(&ctx).get_argument(2); + let signless_i32 = block.deref(&ctx).get_argument(3); + let count = block.deref(&ctx).get_argument(4); + let token = block.deref(&ctx).get_argument(5); + + for barrier in [generic, shared] { + let init = MbarrierInitSharedOp::build(&mut ctx, barrier, count); + assert!(MbarrierInitSharedOp::new(init).verify(&ctx).is_ok()); + let arrive = MbarrierArriveSharedOp::build(&mut ctx, barrier); + assert!(MbarrierArriveSharedOp::new(arrive).verify(&ctx).is_ok()); + let test_wait = MbarrierTestWaitSharedOp::build(&mut ctx, barrier, token); + assert!( + MbarrierTestWaitSharedOp::new(test_wait) + .verify(&ctx) + .is_ok() + ); + let inval = MbarrierInvalSharedOp::build(&mut ctx, barrier); + assert!(MbarrierInvalSharedOp::new(inval).verify(&ctx).is_ok()); + } + + for barrier in [token, global] { + let inval = MbarrierInvalSharedOp::build(&mut ctx, barrier); + assert!(MbarrierInvalSharedOp::new(inval).verify(&ctx).is_err()); + } + + let bad_count = MbarrierInitSharedOp::build(&mut ctx, shared, signless_i32); + assert!(MbarrierInitSharedOp::new(bad_count).verify(&ctx).is_err()); + let bad_arrive_result = Operation::new( + &mut ctx, + MbarrierArriveSharedOp::get_concrete_op_info(), + vec![u32_ty.into()], + vec![shared], + vec![], + 0, + ); + assert!( + MbarrierArriveSharedOp::new(bad_arrive_result) + .verify(&ctx) + .is_err() + ); + let bad_token = MbarrierTestWaitSharedOp::build(&mut ctx, shared, count); + assert!( + MbarrierTestWaitSharedOp::new(bad_token) + .verify(&ctx) + .is_err() + ); + let bad_predicate = Operation::new( + &mut ctx, + MbarrierTestWaitSharedOp::get_concrete_op_info(), + vec![u1_ty.into()], + vec![shared, token], + vec![], + 0, + ); + assert!( + MbarrierTestWaitSharedOp::new(bad_predicate) + .verify(&ctx) + .is_err() + ); + + let missing_operands = Operation::new( + &mut ctx, + MbarrierInitSharedOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + assert!( + MbarrierInitSharedOp::new(missing_operands) + .verify(&ctx) + .is_err() + ); +} diff --git a/crates/dialect-nvvm/tests/ops_test/inline_ptx.rs b/crates/dialect-nvvm/tests/ops_test/inline_ptx.rs new file mode 100644 index 0000000000..b11bf9866b --- /dev/null +++ b/crates/dialect-nvvm/tests/ops_test/inline_ptx.rs @@ -0,0 +1,155 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use dialect_mir::{ + attributes::MirPointerKindAuthorityAttr, + types::{MirPointerKind, MirPtrType, MirTupleType}, +}; +use dialect_nvvm::ops::InlinePtxOp; + +use pliron::{ + basic_block::BasicBlock, + builtin::types::{IntegerType, Signedness}, + context::Context, + op::verify_op, +}; + +#[test] +fn test_inline_ptx_results_must_match_output_constraints() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); + let input = block.deref(&ctx).get_argument(0); + + let void = InlinePtxOp::build(&mut ctx, vec![], vec![], "membar.gl;", "", true, false); + assert!(verify_op(&InlinePtxOp::new(void), &ctx).is_ok()); + + let single = InlinePtxOp::build( + &mut ctx, + vec![i32_ty.into()], + vec![input], + "add.u32 $0, $1, $1;", + "=r,r", + false, + false, + ); + assert!(verify_op(&InlinePtxOp::new(single), &ctx).is_ok()); + + let multi = InlinePtxOp::build( + &mut ctx, + vec![i32_ty.into(), i32_ty.into()], + vec![input], + "add.u32 $0, $2, $2; mul.lo.u32 $1, $2, $2;", + "=r,=r,r", + false, + false, + ); + assert!(verify_op(&InlinePtxOp::new(multi), &ctx).is_ok()); + + let missing_result = InlinePtxOp::build( + &mut ctx, + vec![i32_ty.into()], + vec![input], + "add.u32 $0, $2, $2; mul.lo.u32 $1, $2, $2;", + "=r,=r,r", + false, + false, + ); + assert!(verify_op(&InlinePtxOp::new(missing_result), &ctx).is_err()); + + let extra_result = InlinePtxOp::build( + &mut ctx, + vec![i32_ty.into()], + vec![input], + "prefetch.global.L1 [$0];", + "r", + true, + false, + ); + assert!(verify_op(&InlinePtxOp::new(extra_result), &ctx).is_err()); +} + +#[test] +fn test_inline_ptx_pointer_results_require_exact_inline_asm_authority() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let unique_ty = + MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::UniqueRef); + let erased_ty = + MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::Erased); + let nested_erased_ty = MirTupleType::get(&mut ctx, vec![erased_ty.into()]); + + let build = |ctx: &mut Context, result_ty| { + InlinePtxOp::build( + ctx, + vec![result_ty], + vec![], + "mov.u64 $0, 0;", + "=l", + false, + false, + ) + }; + + for (name, result_ty) in [ + ("UniqueRef", unique_ty.into()), + ("Erased", erased_ty.into()), + ("nested Erased", nested_erased_ty.into()), + ] { + let unmarked = InlinePtxOp::new(build(&mut ctx, result_ty)); + let error = verify_op(&unmarked, &ctx).unwrap_err(); + assert!( + error + .err + .to_string() + .contains("pointer-carrying results require InlineAsm pointer-kind authority"), + "unexpected unmarked {name} error: {}", + error.err + ); + + let marked = InlinePtxOp::new(build(&mut ctx, result_ty)); + marked.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::InlineAsm); + assert!( + verify_op(&marked, &ctx).is_ok(), + "marked {name} pointer result should verify" + ); + } + + let wrong_authority = InlinePtxOp::new(build(&mut ctx, unique_ty.into())); + wrong_authority.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::RustCast); + let wrong_error = verify_op(&wrong_authority, &ctx).unwrap_err(); + assert!( + wrong_error.err.to_string().contains("found RustCast"), + "unexpected wrong-authority error: {}", + wrong_error.err + ); + + let spurious = InlinePtxOp::new(build(&mut ctx, i32_ty.into())); + spurious.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::InlineAsm); + let spurious_error = verify_op(&spurious, &ctx).unwrap_err(); + assert!( + spurious_error + .err + .to_string() + .contains("spurious for pointer-free results"), + "unexpected spurious-authority error: {}", + spurious_error.err + ); +} + +#[test] +fn test_inline_ptx_count_output_constraints() { + assert_eq!(InlinePtxOp::count_output_constraints("=r,r,r"), 1); + assert_eq!(InlinePtxOp::count_output_constraints("=r,=r,=f,=d,r,l"), 4); + assert_eq!(InlinePtxOp::count_output_constraints("r,l,~{memory}"), 0); + assert_eq!(InlinePtxOp::count_output_constraints(""), 0); + // `=` only counts as an output marker at the start of a token. + assert_eq!(InlinePtxOp::count_output_constraints("r,r=f"), 0); +} diff --git a/crates/dialect-nvvm/tests/ops_test/main.rs b/crates/dialect-nvvm/tests/ops_test/main.rs new file mode 100644 index 0000000000..2645507245 --- /dev/null +++ b/crates/dialect-nvvm/tests/ops_test/main.rs @@ -0,0 +1,16 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +// Tests build kinded fixture types directly; production minting lives in +// mir-importer's facts.rs (see the workspace clippy.toml disallowed-methods). +#![allow(clippy::disallowed_methods)] + +mod allowlist; +mod atomics; +mod carriers; +mod cluster_and_memory; +mod inline_ptx; +mod matrix; +mod thread_and_warp; diff --git a/crates/dialect-nvvm/tests/ops_test/matrix.rs b/crates/dialect-nvvm/tests/ops_test/matrix.rs new file mode 100644 index 0000000000..13aee19a73 --- /dev/null +++ b/crates/dialect-nvvm/tests/ops_test/matrix.rs @@ -0,0 +1,2278 @@ +/* + * 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::{ + LdmatrixElementAttr, LdmatrixLayoutAttr, LdmatrixMultiplicityAttr, LdmatrixOp, + LdmatrixShapeAttr, LdmatrixStateSpaceAttr, LdmatrixX1Op, LdmatrixX1TransOp, LdmatrixX2Op, + LdmatrixX2TransOp, LdmatrixX4Op, LdmatrixX4TransOp, MmaM8N8K4F64Op, MmaM16N8K8F32Tf32Op, + MmaM16N8K16F32Bf16Op, MmaM16N8K16F32F16Op, MmaM16N8K32S32S8Op, MovmatrixTransB16Op, + RegisterMmaAccumulatorAttr, RegisterMmaElementAttr, RegisterMmaLayoutAttr, RegisterMmaOp, + RegisterMmaOperationAttr, RegisterMmaOverflowAttr, RegisterMmaShapeAttr, + SparseMmaAccumulatorAttr, SparseMmaElementAttr, SparseMmaLayoutAttr, SparseMmaMetadataAttr, + SparseMmaOp, SparseMmaOverflowAttr, SparseMmaSelectorAttr, SparseMmaShapeAttr, + StmatrixM8n8X4Op, +}; + +use pliron::{ + basic_block::BasicBlock, + builtin::types::{FP32Type, FP64Type, IntegerType, Signedness}, + common_traits::Verify, + context::Context, + op::{Op, verify_op}, + operation::Operation, + r#type::Typed, +}; + +#[test] +fn test_mma_m8n8k4_f64_requires_four_f64_operands_and_two_f64_results() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + let f64_ty = FP64Type::get(&ctx); + let f32_ty = FP32Type::get(&ctx); + let block = BasicBlock::new( + &mut ctx, + None, + vec![ + f64_ty.into(), + f64_ty.into(), + f64_ty.into(), + f64_ty.into(), + f32_ty.into(), + ], + ); + let f64_operands = (0..4) + .map(|index| block.deref(&ctx).get_argument(index)) + .collect::>(); + let f32_value = block.deref(&ctx).get_argument(4); + + let valid = Operation::new( + &mut ctx, + MmaM8N8K4F64Op::get_concrete_op_info(), + vec![f64_ty.into(), f64_ty.into()], + f64_operands.clone(), + vec![], + 0, + ); + assert!(verify_op(&MmaM8N8K4F64Op::new(valid), &ctx).is_ok()); + + let mut bad_operands = f64_operands.clone(); + bad_operands[2] = f32_value; + let invalid_operand = Operation::new( + &mut ctx, + MmaM8N8K4F64Op::get_concrete_op_info(), + vec![f64_ty.into(), f64_ty.into()], + bad_operands, + vec![], + 0, + ); + assert!(verify_op(&MmaM8N8K4F64Op::new(invalid_operand), &ctx).is_err()); + + let invalid_result = Operation::new( + &mut ctx, + MmaM8N8K4F64Op::get_concrete_op_info(), + vec![f64_ty.into(), f32_ty.into()], + f64_operands, + vec![], + 0, + ); + assert!(verify_op(&MmaM8N8K4F64Op::new(invalid_result), &ctx).is_err()); +} + +#[test] +fn test_movmatrix_requires_one_i32_operand_and_result() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); + let f32_ty = FP32Type::get(&ctx); + let block = BasicBlock::new( + &mut ctx, + None, + vec![i32_ty.into(), i64_ty.into(), f32_ty.into()], + ); + let i32_value = block.deref(&ctx).get_argument(0); + let i64_value = block.deref(&ctx).get_argument(1); + let f32_value = block.deref(&ctx).get_argument(2); + + let valid = Operation::new( + &mut ctx, + MovmatrixTransB16Op::get_concrete_op_info(), + vec![i32_ty.into()], + vec![i32_value], + vec![], + 0, + ); + assert!(verify_op(&MovmatrixTransB16Op::new(valid), &ctx).is_ok()); + + for (operand, result_type) in [ + (i64_value, i32_ty.into()), + (f32_value, i32_ty.into()), + (i32_value, i64_ty.into()), + (i32_value, f32_ty.into()), + ] { + let invalid = Operation::new( + &mut ctx, + MovmatrixTransB16Op::get_concrete_op_info(), + vec![result_type], + vec![operand], + vec![], + 0, + ); + assert!( + verify_op(&MovmatrixTransB16Op::new(invalid), &ctx).is_err(), + "movmatrix must reject non-i32 carriers" + ); + } +} + +fn make_ldmatrix_x2( + ctx: &mut Context, + pointer: pliron::value::Value, + result_types: Vec, +) -> LdmatrixOp { + let operation = Operation::new( + ctx, + LdmatrixOp::get_concrete_op_info(), + result_types, + vec![pointer], + vec![], + 0, + ); + let ldmatrix = LdmatrixOp::new(operation); + ldmatrix.set_attr_nvvm_ldmatrix_shape(ctx, LdmatrixShapeAttr::M8n8); + ldmatrix.set_attr_nvvm_ldmatrix_multiplicity(ctx, LdmatrixMultiplicityAttr::X2); + ldmatrix.set_attr_nvvm_ldmatrix_layout(ctx, LdmatrixLayoutAttr::Normal); + ldmatrix.set_attr_nvvm_ldmatrix_element(ctx, LdmatrixElementAttr::B16); + ldmatrix.set_attr_nvvm_ldmatrix_state_space(ctx, LdmatrixStateSpaceAttr::Shared); + ldmatrix +} + +#[test] +fn test_ldmatrix_accepts_only_generic_or_shared_u32_pointers() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + let u16_ty = IntegerType::get(&ctx, 16, Signedness::Unsigned); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let pointer_types = [ + MirPtrType::get_generic(&mut ctx, u32_ty.into(), false), + MirPtrType::get_shared(&mut ctx, u32_ty.into(), false), + MirPtrType::get_global(&mut ctx, u32_ty.into(), false), + MirPtrType::get_constant(&mut ctx, u32_ty.into(), false), + MirPtrType::get(&mut ctx, u32_ty.into(), false, address_space::LOCAL), + MirPtrType::get_generic(&mut ctx, u16_ty.into(), false), + ]; + let block = BasicBlock::new( + &mut ctx, + None, + pointer_types + .iter() + .map(|pointer| (*pointer).into()) + .collect(), + ); + + for index in 0..pointer_types.len() { + let pointer = block.deref(&ctx).get_argument(index); + let operation = make_ldmatrix_x2(&mut ctx, pointer, vec![u32_ty.into(), u32_ty.into()]); + let verified = verify_op(&operation, &ctx); + if index < 2 { + assert!(verified.is_ok(), "pointer case {index} should be accepted"); + } else { + assert!(verified.is_err(), "pointer case {index} should be rejected"); + } + } +} + +#[test] +fn generated_ldmatrix_verifier_rejects_zero_or_two_operands_without_panicking() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let pointer_ty = MirPtrType::get_shared(&mut ctx, u32_ty.into(), false); + let block = BasicBlock::new(&mut ctx, None, vec![pointer_ty.into(), pointer_ty.into()]); + let pointer0 = block.deref(&ctx).get_argument(0); + let pointer1 = block.deref(&ctx).get_argument(1); + + let zero = Operation::new( + &mut ctx, + LdmatrixOp::get_concrete_op_info(), + vec![u32_ty.into(); 4], + vec![], + vec![], + 0, + ); + assert!(LdmatrixOp::new(zero).verify(&ctx).is_err()); + + let two = Operation::new( + &mut ctx, + LdmatrixOp::get_concrete_op_info(), + vec![u32_ty.into(); 4], + vec![pointer0, pointer1], + vec![], + 0, + ); + assert!(LdmatrixOp::new(two).verify(&ctx).is_err()); + + let valid = LdmatrixOp::build( + &mut ctx, + pointer0, + LdmatrixShapeAttr::M8n8, + LdmatrixMultiplicityAttr::X4, + LdmatrixLayoutAttr::Normal, + LdmatrixElementAttr::B16, + LdmatrixStateSpaceAttr::Shared, + ); + assert!(LdmatrixOp::new(valid).verify(&ctx).is_ok()); +} + +#[test] +fn blackwell_ldmatrix_verifier_accepts_only_reviewed_shapes() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + let u8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let shared_u8 = MirPtrType::get_shared(&mut ctx, u8_ty.into(), false); + let generic_u8 = MirPtrType::get_generic(&mut ctx, u8_ty.into(), false); + let global_u8 = MirPtrType::get_global(&mut ctx, u8_ty.into(), false); + let shared_u32 = MirPtrType::get_shared(&mut ctx, u32_ty.into(), false); + let block = BasicBlock::new( + &mut ctx, + None, + vec![ + shared_u8.into(), + generic_u8.into(), + global_u8.into(), + shared_u32.into(), + ], + ); + let shared_pointer = block.deref(&ctx).get_argument(0); + + for pointer_index in [0, 1] { + let pointer = block.deref(&ctx).get_argument(pointer_index); + for (shape, multiplicity, layout, element, result_count) in [ + ( + LdmatrixShapeAttr::M16n16, + LdmatrixMultiplicityAttr::X1, + LdmatrixLayoutAttr::Transposed, + LdmatrixElementAttr::B8, + 2, + ), + ( + LdmatrixShapeAttr::M16n16, + LdmatrixMultiplicityAttr::X1, + LdmatrixLayoutAttr::Transposed, + LdmatrixElementAttr::B8x16B4x16P64, + 2, + ), + ( + LdmatrixShapeAttr::M16n16, + LdmatrixMultiplicityAttr::X1, + LdmatrixLayoutAttr::Transposed, + LdmatrixElementAttr::B8x16B6x16P32, + 2, + ), + ( + LdmatrixShapeAttr::M16n16, + LdmatrixMultiplicityAttr::X2, + LdmatrixLayoutAttr::Transposed, + LdmatrixElementAttr::B8, + 4, + ), + ( + LdmatrixShapeAttr::M16n16, + LdmatrixMultiplicityAttr::X2, + LdmatrixLayoutAttr::Transposed, + LdmatrixElementAttr::B8x16B4x16P64, + 4, + ), + ( + LdmatrixShapeAttr::M16n16, + LdmatrixMultiplicityAttr::X2, + LdmatrixLayoutAttr::Transposed, + LdmatrixElementAttr::B8x16B6x16P32, + 4, + ), + ( + LdmatrixShapeAttr::M8n16, + LdmatrixMultiplicityAttr::X1, + LdmatrixLayoutAttr::Normal, + LdmatrixElementAttr::B8x16B4x16P64, + 1, + ), + ( + LdmatrixShapeAttr::M8n16, + LdmatrixMultiplicityAttr::X1, + LdmatrixLayoutAttr::Normal, + LdmatrixElementAttr::B8x16B6x16P32, + 1, + ), + ( + LdmatrixShapeAttr::M8n16, + LdmatrixMultiplicityAttr::X2, + LdmatrixLayoutAttr::Normal, + LdmatrixElementAttr::B8x16B4x16P64, + 2, + ), + ( + LdmatrixShapeAttr::M8n16, + LdmatrixMultiplicityAttr::X2, + LdmatrixLayoutAttr::Normal, + LdmatrixElementAttr::B8x16B6x16P32, + 2, + ), + ( + LdmatrixShapeAttr::M8n16, + LdmatrixMultiplicityAttr::X4, + LdmatrixLayoutAttr::Normal, + LdmatrixElementAttr::B8x16B4x16P64, + 4, + ), + ( + LdmatrixShapeAttr::M8n16, + LdmatrixMultiplicityAttr::X4, + LdmatrixLayoutAttr::Normal, + LdmatrixElementAttr::B8x16B6x16P32, + 4, + ), + ] { + let op = LdmatrixOp::build( + &mut ctx, + pointer, + shape, + multiplicity, + layout, + element, + LdmatrixStateSpaceAttr::Shared, + ); + assert_eq!(op.deref(&ctx).get_num_results(), result_count); + assert!(LdmatrixOp::new(op).verify(&ctx).is_ok()); + } + } + + for (shape, multiplicity, layout, element) in [ + ( + LdmatrixShapeAttr::M16n16, + LdmatrixMultiplicityAttr::X4, + LdmatrixLayoutAttr::Transposed, + LdmatrixElementAttr::B8, + ), + ( + LdmatrixShapeAttr::M16n16, + LdmatrixMultiplicityAttr::X1, + LdmatrixLayoutAttr::Normal, + LdmatrixElementAttr::B8, + ), + ( + LdmatrixShapeAttr::M8n16, + LdmatrixMultiplicityAttr::X1, + LdmatrixLayoutAttr::Transposed, + LdmatrixElementAttr::B8x16B4x16P64, + ), + ( + LdmatrixShapeAttr::M8n16, + LdmatrixMultiplicityAttr::X1, + LdmatrixLayoutAttr::Normal, + LdmatrixElementAttr::B8, + ), + ] { + let op = LdmatrixOp::build( + &mut ctx, + shared_pointer, + shape, + multiplicity, + layout, + element, + LdmatrixStateSpaceAttr::Shared, + ); + assert!(LdmatrixOp::new(op).verify(&ctx).is_err()); + } + + for pointer_index in [2, 3] { + let pointer = block.deref(&ctx).get_argument(pointer_index); + let op = LdmatrixOp::build( + &mut ctx, + pointer, + LdmatrixShapeAttr::M16n16, + LdmatrixMultiplicityAttr::X1, + LdmatrixLayoutAttr::Transposed, + LdmatrixElementAttr::B8, + LdmatrixStateSpaceAttr::Shared, + ); + assert!(LdmatrixOp::new(op).verify(&ctx).is_err()); + } +} + +#[test] +fn classic_ldmatrix_compatibility_ops_keep_names_and_register_shapes() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let pointer_ty = MirPtrType::get_shared(&mut ctx, u32_ty.into(), false); + let block = BasicBlock::new(&mut ctx, None, vec![pointer_ty.into()]); + let pointer = block.deref(&ctx).get_argument(0); + + macro_rules! check_compat { + ($op:ty, $name:literal, $results:literal) => {{ + assert_eq!(<$op>::get_opid_static().to_string(), $name); + let valid = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + vec![u32_ty.into(); $results], + vec![pointer], + vec![], + 0, + ); + assert!(verify_op(&<$op>::new(valid), &ctx).is_ok()); + + let wrong_shape = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + vec![u32_ty.into(); $results + 1], + vec![pointer], + vec![], + 0, + ); + assert!(verify_op(&<$op>::new(wrong_shape), &ctx).is_err()); + }}; + } + + check_compat!(LdmatrixX1Op, "nvvm.ldmatrix_x1", 1); + check_compat!(LdmatrixX1TransOp, "nvvm.ldmatrix_x1_trans", 1); + check_compat!(LdmatrixX2Op, "nvvm.ldmatrix_x2", 2); + check_compat!(LdmatrixX2TransOp, "nvvm.ldmatrix_x2_trans", 2); + check_compat!(LdmatrixX4Op, "nvvm.ldmatrix_x4", 4); + check_compat!(LdmatrixX4TransOp, "nvvm.ldmatrix_x4_trans", 4); +} + +#[test] +fn test_mma_m16n8k16_bf16_verifies_exact_register_signature() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + let f32_ty = FP32Type::get(&ctx); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); + let block = BasicBlock::new( + &mut ctx, + None, + vec![f32_ty.into(), i32_ty.into(), i64_ty.into()], + ); + let f32_value = block.deref(&ctx).get_argument(0); + let i32_value = block.deref(&ctx).get_argument(1); + let i64_value = block.deref(&ctx).get_argument(2); + + let valid_operands = (0..4) + .map(|_| f32_value) + .chain((0..6).map(|_| i32_value)) + .collect(); + let valid = Operation::new( + &mut ctx, + MmaM16N8K16F32Bf16Op::get_concrete_op_info(), + vec![f32_ty.into(); 4], + valid_operands, + vec![], + 0, + ); + assert!(verify_op(&MmaM16N8K16F32Bf16Op::new(valid), &ctx).is_ok()); + + let bad_c_operands = (0..4) + .map(|index| if index == 0 { i32_value } else { f32_value }) + .chain((0..6).map(|_| i32_value)) + .collect(); + let bad_c = Operation::new( + &mut ctx, + MmaM16N8K16F32Bf16Op::get_concrete_op_info(), + vec![f32_ty.into(); 4], + bad_c_operands, + vec![], + 0, + ); + assert!(verify_op(&MmaM16N8K16F32Bf16Op::new(bad_c), &ctx).is_err()); + + let bad_a_operands = (0..4) + .map(|_| f32_value) + .chain((0..6).map(|index| if index == 0 { i64_value } else { i32_value })) + .collect(); + let bad_a = Operation::new( + &mut ctx, + MmaM16N8K16F32Bf16Op::get_concrete_op_info(), + vec![f32_ty.into(); 4], + bad_a_operands, + vec![], + 0, + ); + assert!(verify_op(&MmaM16N8K16F32Bf16Op::new(bad_a), &ctx).is_err()); + + let bad_result = Operation::new( + &mut ctx, + MmaM16N8K16F32Bf16Op::get_concrete_op_info(), + vec![f32_ty.into(), f32_ty.into(), f32_ty.into(), i32_ty.into()], + (0..4) + .map(|_| f32_value) + .chain((0..6).map(|_| i32_value)) + .collect(), + vec![], + 0, + ); + assert!(verify_op(&MmaM16N8K16F32Bf16Op::new(bad_result), &ctx).is_err()); + + let bad_arity = Operation::new( + &mut ctx, + MmaM16N8K16F32Bf16Op::get_concrete_op_info(), + vec![f32_ty.into(); 4], + (0..4) + .map(|_| f32_value) + .chain((0..5).map(|_| i32_value)) + .collect(), + vec![], + 0, + ); + assert!(verify_op(&MmaM16N8K16F32Bf16Op::new(bad_arity), &ctx).is_err()); +} + +#[test] +fn test_mma_m16n8k16_f16_verifies_exact_register_signature() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + let f32_ty = FP32Type::get(&ctx); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); + let block = BasicBlock::new( + &mut ctx, + None, + vec![f32_ty.into(), i32_ty.into(), i64_ty.into()], + ); + let f32_value = block.deref(&ctx).get_argument(0); + let i32_value = block.deref(&ctx).get_argument(1); + let i64_value = block.deref(&ctx).get_argument(2); + + let operands = || { + (0..4) + .map(|_| f32_value) + .chain((0..6).map(|_| i32_value)) + .collect() + }; + let valid = Operation::new( + &mut ctx, + MmaM16N8K16F32F16Op::get_concrete_op_info(), + vec![f32_ty.into(); 4], + operands(), + vec![], + 0, + ); + assert!(verify_op(&MmaM16N8K16F32F16Op::new(valid), &ctx).is_ok()); + + let bad_c_operands = (0..4) + .map(|index| if index == 0 { i32_value } else { f32_value }) + .chain((0..6).map(|_| i32_value)) + .collect(); + let bad_c = Operation::new( + &mut ctx, + MmaM16N8K16F32F16Op::get_concrete_op_info(), + vec![f32_ty.into(); 4], + bad_c_operands, + vec![], + 0, + ); + assert!(verify_op(&MmaM16N8K16F32F16Op::new(bad_c), &ctx).is_err()); + + let bad_packed_operands = (0..4) + .map(|_| f32_value) + .chain((0..6).map(|index| if index == 0 { i64_value } else { i32_value })) + .collect(); + let bad_packed = Operation::new( + &mut ctx, + MmaM16N8K16F32F16Op::get_concrete_op_info(), + vec![f32_ty.into(); 4], + bad_packed_operands, + vec![], + 0, + ); + assert!(verify_op(&MmaM16N8K16F32F16Op::new(bad_packed), &ctx).is_err()); + + let bad_result_type = Operation::new( + &mut ctx, + MmaM16N8K16F32F16Op::get_concrete_op_info(), + vec![f32_ty.into(), f32_ty.into(), f32_ty.into(), i32_ty.into()], + operands(), + vec![], + 0, + ); + assert!(verify_op(&MmaM16N8K16F32F16Op::new(bad_result_type), &ctx).is_err()); + + let bad_operand_arity = Operation::new( + &mut ctx, + MmaM16N8K16F32F16Op::get_concrete_op_info(), + vec![f32_ty.into(); 4], + (0..4) + .map(|_| f32_value) + .chain((0..5).map(|_| i32_value)) + .collect(), + vec![], + 0, + ); + assert!(verify_op(&MmaM16N8K16F32F16Op::new(bad_operand_arity), &ctx).is_err()); + + let bad_result_arity = Operation::new( + &mut ctx, + MmaM16N8K16F32F16Op::get_concrete_op_info(), + vec![f32_ty.into(); 3], + operands(), + vec![], + 0, + ); + assert!(verify_op(&MmaM16N8K16F32F16Op::new(bad_result_arity), &ctx).is_err()); +} + +#[test] +fn test_mma_m16n8k8_tf32_verifies_exact_register_signature() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + let f32_ty = FP32Type::get(&ctx); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); + let block = BasicBlock::new( + &mut ctx, + None, + vec![f32_ty.into(), i32_ty.into(), i64_ty.into()], + ); + let f32_value = block.deref(&ctx).get_argument(0); + let i32_value = block.deref(&ctx).get_argument(1); + let i64_value = block.deref(&ctx).get_argument(2); + + let operands = || { + (0..4) + .map(|_| f32_value) + .chain((0..6).map(|_| i32_value)) + .collect() + }; + let valid = Operation::new( + &mut ctx, + MmaM16N8K8F32Tf32Op::get_concrete_op_info(), + vec![f32_ty.into(); 4], + operands(), + vec![], + 0, + ); + assert!(verify_op(&MmaM16N8K8F32Tf32Op::new(valid), &ctx).is_ok()); + + let bad_c_operands = (0..4) + .map(|index| if index == 0 { i32_value } else { f32_value }) + .chain((0..6).map(|_| i32_value)) + .collect(); + let bad_c = Operation::new( + &mut ctx, + MmaM16N8K8F32Tf32Op::get_concrete_op_info(), + vec![f32_ty.into(); 4], + bad_c_operands, + vec![], + 0, + ); + assert!(verify_op(&MmaM16N8K8F32Tf32Op::new(bad_c), &ctx).is_err()); + + let bad_packed_operands = (0..4) + .map(|_| f32_value) + .chain((0..6).map(|index| if index == 0 { i64_value } else { i32_value })) + .collect(); + let bad_packed = Operation::new( + &mut ctx, + MmaM16N8K8F32Tf32Op::get_concrete_op_info(), + vec![f32_ty.into(); 4], + bad_packed_operands, + vec![], + 0, + ); + assert!(verify_op(&MmaM16N8K8F32Tf32Op::new(bad_packed), &ctx).is_err()); + + let bad_result_type = Operation::new( + &mut ctx, + MmaM16N8K8F32Tf32Op::get_concrete_op_info(), + vec![f32_ty.into(), f32_ty.into(), f32_ty.into(), i32_ty.into()], + operands(), + vec![], + 0, + ); + assert!(verify_op(&MmaM16N8K8F32Tf32Op::new(bad_result_type), &ctx).is_err()); + + let bad_operand_arity = Operation::new( + &mut ctx, + MmaM16N8K8F32Tf32Op::get_concrete_op_info(), + vec![f32_ty.into(); 4], + (0..4) + .map(|_| f32_value) + .chain((0..5).map(|_| i32_value)) + .collect(), + vec![], + 0, + ); + assert!(verify_op(&MmaM16N8K8F32Tf32Op::new(bad_operand_arity), &ctx).is_err()); + + let bad_result_arity = Operation::new( + &mut ctx, + MmaM16N8K8F32Tf32Op::get_concrete_op_info(), + vec![f32_ty.into(); 3], + operands(), + vec![], + 0, + ); + assert!(verify_op(&MmaM16N8K8F32Tf32Op::new(bad_result_arity), &ctx).is_err()); +} + +#[test] +fn test_mma_m16n8k32_s8_verifies_exact_register_signature() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); + let f32_ty = FP32Type::get(&ctx); + let block = BasicBlock::new( + &mut ctx, + None, + vec![i32_ty.into(), i64_ty.into(), f32_ty.into()], + ); + let i32_value = block.deref(&ctx).get_argument(0); + let i64_value = block.deref(&ctx).get_argument(1); + let f32_value = block.deref(&ctx).get_argument(2); + let valid_operands = vec![i32_value; 10]; + + let valid = Operation::new( + &mut ctx, + MmaM16N8K32S32S8Op::get_concrete_op_info(), + vec![i32_ty.into(); 4], + valid_operands.clone(), + vec![], + 0, + ); + assert!(verify_op(&MmaM16N8K32S32S8Op::new(valid), &ctx).is_ok()); + + for bad_value in [i64_value, f32_value] { + let mut bad_operands = valid_operands.clone(); + bad_operands[4] = bad_value; + let invalid = Operation::new( + &mut ctx, + MmaM16N8K32S32S8Op::get_concrete_op_info(), + vec![i32_ty.into(); 4], + bad_operands, + vec![], + 0, + ); + assert!( + verify_op(&MmaM16N8K32S32S8Op::new(invalid), &ctx).is_err(), + "MMA must reject non-i32 register operands" + ); + } + + for bad_results in [ + vec![i32_ty.into(), i32_ty.into(), i32_ty.into(), i64_ty.into()], + vec![i32_ty.into(), i32_ty.into(), i32_ty.into(), f32_ty.into()], + vec![i32_ty.into(); 3], + ] { + let invalid = Operation::new( + &mut ctx, + MmaM16N8K32S32S8Op::get_concrete_op_info(), + bad_results, + valid_operands.clone(), + vec![], + 0, + ); + assert!( + verify_op(&MmaM16N8K32S32S8Op::new(invalid), &ctx).is_err(), + "MMA must reject the wrong result register signature" + ); + } + + let invalid_arity = Operation::new( + &mut ctx, + MmaM16N8K32S32S8Op::get_concrete_op_info(), + vec![i32_ty.into(); 4], + vec![i32_value; 9], + vec![], + 0, + ); + assert!(verify_op(&MmaM16N8K32S32S8Op::new(invalid_arity), &ctx).is_err()); +} + +#[test] +fn generated_register_mma_verifier_rejects_crossed_variants_and_carriers() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + macro_rules! set_variant { + ($op:expr, $shape:expr, $acc:expr, $a:expr, $b:expr, $overflow:expr) => {{ + $op.set_attr_nvvm_register_mma_shape(&ctx, $shape); + $op.set_attr_nvvm_register_mma_accumulator(&ctx, $acc); + $op.set_attr_nvvm_register_mma_a_element(&ctx, $a); + $op.set_attr_nvvm_register_mma_b_element(&ctx, $b); + $op.set_attr_nvvm_register_mma_a_layout(&ctx, RegisterMmaLayoutAttr::Row); + $op.set_attr_nvvm_register_mma_b_layout(&ctx, RegisterMmaLayoutAttr::Col); + $op.set_attr_nvvm_register_mma_overflow(&ctx, $overflow); + }}; + } + + let f32_ty = FP32Type::get(&ctx); + let f64_ty = FP64Type::get(&ctx); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let signless_i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let block = BasicBlock::new( + &mut ctx, + None, + vec![f32_ty.into(), f64_ty.into(), i32_ty.into(), u32_ty.into()], + ); + let f32_value = block.deref(&ctx).get_argument(0); + let f64_value = block.deref(&ctx).get_argument(1); + let i32_value = block.deref(&ctx).get_argument(2); + let u32_value = block.deref(&ctx).get_argument(3); + + let bf16_operation = Operation::new( + &mut ctx, + RegisterMmaOp::get_concrete_op_info(), + vec![f32_ty.into(); 4], + [vec![f32_value; 4], vec![u32_value; 6]].concat(), + vec![], + 0, + ); + let bf16 = RegisterMmaOp::new(bf16_operation); + set_variant!( + bf16, + RegisterMmaShapeAttr::M16n8k16, + RegisterMmaAccumulatorAttr::F32, + RegisterMmaElementAttr::Bf16, + RegisterMmaElementAttr::Bf16, + RegisterMmaOverflowAttr::NotApplicable + ); + assert!(bf16.get_attr_nvvm_register_mma_operation(&ctx).is_none()); + assert!(verify_op(&bf16, &ctx).is_ok()); + bf16.set_attr_nvvm_register_mma_b_element(&ctx, RegisterMmaElementAttr::F16); + assert!(verify_op(&bf16, &ctx).is_err()); + + let f64_operation = Operation::new( + &mut ctx, + RegisterMmaOp::get_concrete_op_info(), + vec![f64_ty.into(); 2], + vec![f64_value; 4], + vec![], + 0, + ); + let f64_mma = RegisterMmaOp::new(f64_operation); + set_variant!( + f64_mma, + RegisterMmaShapeAttr::M8n8k4, + RegisterMmaAccumulatorAttr::F64, + RegisterMmaElementAttr::F64, + RegisterMmaElementAttr::F64, + RegisterMmaOverflowAttr::NotApplicable + ); + assert!(verify_op(&f64_mma, &ctx).is_ok()); + + let int_operands = [vec![i32_value; 4], vec![u32_value; 6]].concat(); + let int_operation = Operation::new( + &mut ctx, + RegisterMmaOp::get_concrete_op_info(), + vec![i32_ty.into(); 4], + int_operands.clone(), + vec![], + 0, + ); + let int_mma = RegisterMmaOp::new(int_operation); + set_variant!( + int_mma, + RegisterMmaShapeAttr::M16n8k32, + RegisterMmaAccumulatorAttr::S32, + RegisterMmaElementAttr::S8, + RegisterMmaElementAttr::S8, + RegisterMmaOverflowAttr::Wrapping + ); + assert!(verify_op(&int_mma, &ctx).is_ok()); + + let wrong_signedness = Operation::new( + &mut ctx, + RegisterMmaOp::get_concrete_op_info(), + vec![signless_i32_ty.into(); 4], + int_operands, + vec![], + 0, + ); + let wrong_signedness = RegisterMmaOp::new(wrong_signedness); + set_variant!( + wrong_signedness, + RegisterMmaShapeAttr::M16n8k32, + RegisterMmaAccumulatorAttr::S32, + RegisterMmaElementAttr::S8, + RegisterMmaElementAttr::S8, + RegisterMmaOverflowAttr::Wrapping + ); + assert!(verify_op(&wrong_signedness, &ctx).is_err()); + + let missing_attributes = Operation::new( + &mut ctx, + RegisterMmaOp::get_concrete_op_info(), + vec![f64_ty.into(); 2], + vec![f64_value; 4], + vec![], + 0, + ); + assert!(verify_op(&RegisterMmaOp::new(missing_attributes), &ctx).is_err()); +} + +#[test] +fn generated_register_mma_verifies_dense_integer_families() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let signless_i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let block = BasicBlock::new( + &mut ctx, + None, + vec![i32_ty.into(), u32_ty.into(), signless_i32_ty.into()], + ); + let i32_value = block.deref(&ctx).get_argument(0); + let u32_value = block.deref(&ctx).get_argument(1); + + macro_rules! int_mma { + ($shape:expr, $a:expr, $b:expr, $overflow:expr, $operands:expr, $results:expr) => {{ + let operation = Operation::new( + &mut ctx, + RegisterMmaOp::get_concrete_op_info(), + $results, + $operands, + vec![], + 0, + ); + let mma = RegisterMmaOp::new(operation); + mma.set_attr_nvvm_register_mma_shape(&ctx, $shape); + mma.set_attr_nvvm_register_mma_operation(&ctx, RegisterMmaOperationAttr::Multiply); + mma.set_attr_nvvm_register_mma_accumulator(&ctx, RegisterMmaAccumulatorAttr::S32); + mma.set_attr_nvvm_register_mma_a_element(&ctx, $a); + mma.set_attr_nvvm_register_mma_b_element(&ctx, $b); + mma.set_attr_nvvm_register_mma_a_layout(&ctx, RegisterMmaLayoutAttr::Row); + mma.set_attr_nvvm_register_mma_b_layout(&ctx, RegisterMmaLayoutAttr::Col); + mma.set_attr_nvvm_register_mma_overflow(&ctx, $overflow); + mma + }}; + } + + let mut accepted = 0; + for (shape, accumulator_count, operand_count, result_count) in [ + (RegisterMmaShapeAttr::M8n8k16, 2, 4, 2), + (RegisterMmaShapeAttr::M16n8k16, 4, 7, 4), + (RegisterMmaShapeAttr::M16n8k32, 4, 10, 4), + ] { + for (a_element, b_element) in [ + (RegisterMmaElementAttr::S8, RegisterMmaElementAttr::S8), + (RegisterMmaElementAttr::S8, RegisterMmaElementAttr::U8), + (RegisterMmaElementAttr::U8, RegisterMmaElementAttr::S8), + (RegisterMmaElementAttr::U8, RegisterMmaElementAttr::U8), + ] { + for overflow in [ + RegisterMmaOverflowAttr::Wrapping, + RegisterMmaOverflowAttr::Satfinite, + ] { + let operands = [ + vec![i32_value; accumulator_count], + vec![u32_value; operand_count - accumulator_count], + ] + .concat(); + let mma = int_mma!( + shape.clone(), + a_element.clone(), + b_element.clone(), + overflow.clone(), + operands, + vec![i32_ty.into(); result_count] + ); + assert_eq!( + mma.get_operation().deref(&ctx).get_num_operands(), + operand_count + ); + assert_eq!( + mma.get_operation().deref(&ctx).get_num_results(), + result_count + ); + assert!( + verify_op(&mma, &ctx).is_ok(), + "rejected {shape:?} {a_element:?}x{b_element:?} {overflow:?}" + ); + accepted += 1; + } + } + } + assert_eq!(accepted, 24); + + let mut int4_accepted = 0; + for (shape, accumulator_count, a_count, b_count, result_count) in [ + (RegisterMmaShapeAttr::M8n8k32, 2, 1, 1, 2), + (RegisterMmaShapeAttr::M16n8k32, 4, 2, 1, 4), + (RegisterMmaShapeAttr::M16n8k64, 4, 4, 2, 4), + ] { + for (a_element, b_element) in [ + (RegisterMmaElementAttr::S4, RegisterMmaElementAttr::S4), + (RegisterMmaElementAttr::S4, RegisterMmaElementAttr::U4), + (RegisterMmaElementAttr::U4, RegisterMmaElementAttr::S4), + (RegisterMmaElementAttr::U4, RegisterMmaElementAttr::U4), + ] { + for overflow in [ + RegisterMmaOverflowAttr::Wrapping, + RegisterMmaOverflowAttr::Satfinite, + ] { + let operands = [ + vec![i32_value; accumulator_count], + vec![u32_value; a_count], + vec![u32_value; b_count], + ] + .concat(); + let expected_operand_types = [ + vec![i32_ty.into(); accumulator_count], + vec![u32_ty.into(); a_count], + vec![u32_ty.into(); b_count], + ] + .concat(); + let mma = int_mma!( + shape.clone(), + a_element.clone(), + b_element.clone(), + overflow.clone(), + operands, + vec![i32_ty.into(); result_count] + ); + let operation = mma.get_operation().deref(&ctx); + assert_eq!( + operation.get_num_operands(), + accumulator_count + a_count + b_count + ); + assert_eq!(operation.get_num_results(), result_count); + assert_eq!( + operation + .operands() + .map(|operand| operand.get_type(&ctx)) + .collect::>(), + expected_operand_types + ); + assert_eq!( + (0..operation.get_num_results()) + .map(|index| operation.get_result(index).get_type(&ctx)) + .collect::>(), + vec![i32_ty.into(); result_count] + ); + assert_eq!( + mma.get_attr_nvvm_register_mma_shape(&ctx).as_deref(), + Some(&shape) + ); + assert_eq!( + mma.get_attr_nvvm_register_mma_accumulator(&ctx).as_deref(), + Some(&RegisterMmaAccumulatorAttr::S32) + ); + assert_eq!( + mma.get_attr_nvvm_register_mma_a_element(&ctx).as_deref(), + Some(&a_element) + ); + assert_eq!( + mma.get_attr_nvvm_register_mma_b_element(&ctx).as_deref(), + Some(&b_element) + ); + assert_eq!( + mma.get_attr_nvvm_register_mma_a_layout(&ctx).as_deref(), + Some(&RegisterMmaLayoutAttr::Row) + ); + assert_eq!( + mma.get_attr_nvvm_register_mma_b_layout(&ctx).as_deref(), + Some(&RegisterMmaLayoutAttr::Col) + ); + assert_eq!( + mma.get_attr_nvvm_register_mma_overflow(&ctx).as_deref(), + Some(&overflow) + ); + assert!( + verify_op(&mma, &ctx).is_ok(), + "rejected {shape:?} {a_element:?}x{b_element:?} {overflow:?}" + ); + int4_accepted += 1; + } + } + } + assert_eq!(int4_accepted, 24); + + for (shape, accumulator_count, operand_count, result_count) in [ + (RegisterMmaShapeAttr::M8n8k32, 2, 4, 2), + (RegisterMmaShapeAttr::M16n8k32, 4, 7, 4), + (RegisterMmaShapeAttr::M16n8k64, 4, 10, 4), + ] { + for wrong_operand_count in [operand_count - 1, operand_count + 1] { + let mma = int_mma!( + shape.clone(), + RegisterMmaElementAttr::S4, + RegisterMmaElementAttr::U4, + RegisterMmaOverflowAttr::Wrapping, + [ + vec![i32_value; accumulator_count], + vec![u32_value; wrong_operand_count - accumulator_count], + ] + .concat(), + vec![i32_ty.into(); result_count] + ); + assert!(verify_op(&mma, &ctx).is_err()); + } + + for wrong_result_count in [result_count - 1, result_count + 1] { + let mma = int_mma!( + shape.clone(), + RegisterMmaElementAttr::U4, + RegisterMmaElementAttr::S4, + RegisterMmaOverflowAttr::Satfinite, + [ + vec![i32_value; accumulator_count], + vec![u32_value; operand_count - accumulator_count], + ] + .concat(), + vec![i32_ty.into(); wrong_result_count] + ); + assert!(verify_op(&mma, &ctx).is_err()); + } + } + + let int4_on_int8_shape = int_mma!( + RegisterMmaShapeAttr::M8n8k16, + RegisterMmaElementAttr::S4, + RegisterMmaElementAttr::U4, + RegisterMmaOverflowAttr::Wrapping, + [vec![i32_value; 2], vec![u32_value; 2]].concat(), + vec![i32_ty.into(); 2] + ); + assert!(verify_op(&int4_on_int8_shape, &ctx).is_err()); + + let int8_on_int4_shape = int_mma!( + RegisterMmaShapeAttr::M8n8k32, + RegisterMmaElementAttr::S8, + RegisterMmaElementAttr::U8, + RegisterMmaOverflowAttr::Wrapping, + [vec![i32_value; 2], vec![u32_value; 2]].concat(), + vec![i32_ty.into(); 2] + ); + assert!(verify_op(&int8_on_int4_shape, &ctx).is_err()); + + let crossed_integer_width = int_mma!( + RegisterMmaShapeAttr::M16n8k32, + RegisterMmaElementAttr::S4, + RegisterMmaElementAttr::U8, + RegisterMmaOverflowAttr::Satfinite, + [vec![i32_value; 4], vec![u32_value; 3]].concat(), + vec![i32_ty.into(); 4] + ); + assert!(verify_op(&crossed_integer_width, &ctx).is_err()); + + let m16k32_int4_with_int8_carriers = int_mma!( + RegisterMmaShapeAttr::M16n8k32, + RegisterMmaElementAttr::U4, + RegisterMmaElementAttr::S4, + RegisterMmaOverflowAttr::Wrapping, + [vec![i32_value; 4], vec![u32_value; 6]].concat(), + vec![i32_ty.into(); 4] + ); + assert!(verify_op(&m16k32_int4_with_int8_carriers, &ctx).is_err()); + + let m16k32_int8_with_int4_carriers = int_mma!( + RegisterMmaShapeAttr::M16n8k32, + RegisterMmaElementAttr::S8, + RegisterMmaElementAttr::U8, + RegisterMmaOverflowAttr::Satfinite, + [vec![i32_value; 4], vec![u32_value; 3]].concat(), + vec![i32_ty.into(); 4] + ); + assert!(verify_op(&m16k32_int8_with_int4_carriers, &ctx).is_err()); + + let m16k64_int4_with_k32_carriers = int_mma!( + RegisterMmaShapeAttr::M16n8k64, + RegisterMmaElementAttr::S4, + RegisterMmaElementAttr::U4, + RegisterMmaOverflowAttr::Wrapping, + [vec![i32_value; 4], vec![u32_value; 3]].concat(), + vec![i32_ty.into(); 4] + ); + assert!(verify_op(&m16k64_int4_with_k32_carriers, &ctx).is_err()); + + let m16k32_int4_with_k64_carriers = int_mma!( + RegisterMmaShapeAttr::M16n8k32, + RegisterMmaElementAttr::U4, + RegisterMmaElementAttr::S4, + RegisterMmaOverflowAttr::Satfinite, + [vec![i32_value; 4], vec![u32_value; 6]].concat(), + vec![i32_ty.into(); 4] + ); + assert!(verify_op(&m16k32_int4_with_k64_carriers, &ctx).is_err()); + + let int8_on_m16k64_shape = int_mma!( + RegisterMmaShapeAttr::M16n8k64, + RegisterMmaElementAttr::S8, + RegisterMmaElementAttr::U8, + RegisterMmaOverflowAttr::Wrapping, + [vec![i32_value; 4], vec![u32_value; 6]].concat(), + vec![i32_ty.into(); 4] + ); + assert!(verify_op(&int8_on_m16k64_shape, &ctx).is_err()); + + for (shape, accumulator_count, operand_count, result_count) in [ + (RegisterMmaShapeAttr::M8n8k16, 2, 4, 2), + (RegisterMmaShapeAttr::M16n8k16, 4, 7, 4), + (RegisterMmaShapeAttr::M16n8k32, 4, 10, 4), + ] { + for wrong_count in [operand_count - 1, operand_count + 1] { + let operands = [ + vec![i32_value; accumulator_count], + vec![u32_value; wrong_count - accumulator_count], + ] + .concat(); + let mma = int_mma!( + shape.clone(), + RegisterMmaElementAttr::S8, + RegisterMmaElementAttr::U8, + RegisterMmaOverflowAttr::Wrapping, + operands, + vec![i32_ty.into(); result_count] + ); + assert!(verify_op(&mma, &ctx).is_err()); + } + } + + let m8_wrong_result_signedness = int_mma!( + RegisterMmaShapeAttr::M8n8k16, + RegisterMmaElementAttr::S8, + RegisterMmaElementAttr::U8, + RegisterMmaOverflowAttr::Wrapping, + [vec![i32_value; 2], vec![u32_value; 2]].concat(), + vec![signless_i32_ty.into(); 2] + ); + assert!(verify_op(&m8_wrong_result_signedness, &ctx).is_err()); + + let m8_wrong_accumulator_signedness = int_mma!( + RegisterMmaShapeAttr::M8n8k16, + RegisterMmaElementAttr::S8, + RegisterMmaElementAttr::U8, + RegisterMmaOverflowAttr::Wrapping, + vec![u32_value; 4], + vec![i32_ty.into(); 2] + ); + assert!(verify_op(&m8_wrong_accumulator_signedness, &ctx).is_err()); + + let m8_wrong_fragment_signedness = int_mma!( + RegisterMmaShapeAttr::M8n8k16, + RegisterMmaElementAttr::U8, + RegisterMmaElementAttr::S8, + RegisterMmaOverflowAttr::Satfinite, + vec![i32_value; 4], + vec![i32_ty.into(); 2] + ); + assert!(verify_op(&m8_wrong_fragment_signedness, &ctx).is_err()); + + let m8_crossed_element = int_mma!( + RegisterMmaShapeAttr::M8n8k16, + RegisterMmaElementAttr::S8, + RegisterMmaElementAttr::F16, + RegisterMmaOverflowAttr::Wrapping, + [vec![i32_value; 2], vec![u32_value; 2]].concat(), + vec![i32_ty.into(); 2] + ); + assert!(verify_op(&m8_crossed_element, &ctx).is_err()); + + let m8_crossed_overflow = int_mma!( + RegisterMmaShapeAttr::M8n8k16, + RegisterMmaElementAttr::U8, + RegisterMmaElementAttr::U8, + RegisterMmaOverflowAttr::NotApplicable, + [vec![i32_value; 2], vec![u32_value; 2]].concat(), + vec![i32_ty.into(); 2] + ); + assert!(verify_op(&m8_crossed_overflow, &ctx).is_err()); + + let m8_crossed_carrier_shape = int_mma!( + RegisterMmaShapeAttr::M8n8k16, + RegisterMmaElementAttr::S8, + RegisterMmaElementAttr::S8, + RegisterMmaOverflowAttr::Wrapping, + [vec![i32_value; 4], vec![u32_value; 3]].concat(), + vec![i32_ty.into(); 4] + ); + assert!(verify_op(&m8_crossed_carrier_shape, &ctx).is_err()); + + let m8_crossed_shape = int_mma!( + RegisterMmaShapeAttr::M8n8k4, + RegisterMmaElementAttr::S8, + RegisterMmaElementAttr::S8, + RegisterMmaOverflowAttr::Wrapping, + [vec![i32_value; 2], vec![u32_value; 2]].concat(), + vec![i32_ty.into(); 2] + ); + assert!(verify_op(&m8_crossed_shape, &ctx).is_err()); + + let wrong_result_signedness = int_mma!( + RegisterMmaShapeAttr::M16n8k16, + RegisterMmaElementAttr::S8, + RegisterMmaElementAttr::U8, + RegisterMmaOverflowAttr::Wrapping, + [vec![i32_value; 4], vec![u32_value; 3]].concat(), + vec![signless_i32_ty.into(); 4] + ); + assert!(verify_op(&wrong_result_signedness, &ctx).is_err()); + + let wrong_accumulator_signedness = int_mma!( + RegisterMmaShapeAttr::M16n8k16, + RegisterMmaElementAttr::S8, + RegisterMmaElementAttr::U8, + RegisterMmaOverflowAttr::Wrapping, + vec![u32_value; 7], + vec![i32_ty.into(); 4] + ); + assert!(verify_op(&wrong_accumulator_signedness, &ctx).is_err()); + + let wrong_fragment_signedness = int_mma!( + RegisterMmaShapeAttr::M16n8k32, + RegisterMmaElementAttr::U8, + RegisterMmaElementAttr::S8, + RegisterMmaOverflowAttr::Satfinite, + vec![i32_value; 10], + vec![i32_ty.into(); 4] + ); + assert!(verify_op(&wrong_fragment_signedness, &ctx).is_err()); + + let crossed_element = int_mma!( + RegisterMmaShapeAttr::M16n8k16, + RegisterMmaElementAttr::S8, + RegisterMmaElementAttr::F16, + RegisterMmaOverflowAttr::Wrapping, + [vec![i32_value; 4], vec![u32_value; 3]].concat(), + vec![i32_ty.into(); 4] + ); + assert!(verify_op(&crossed_element, &ctx).is_err()); + + let crossed_overflow = int_mma!( + RegisterMmaShapeAttr::M16n8k32, + RegisterMmaElementAttr::U8, + RegisterMmaElementAttr::U8, + RegisterMmaOverflowAttr::NotApplicable, + [vec![i32_value; 4], vec![u32_value; 6]].concat(), + vec![i32_ty.into(); 4] + ); + assert!(verify_op(&crossed_overflow, &ctx).is_err()); + + let crossed_shape = int_mma!( + RegisterMmaShapeAttr::M16n8k8, + RegisterMmaElementAttr::S8, + RegisterMmaElementAttr::S8, + RegisterMmaOverflowAttr::Satfinite, + [vec![i32_value; 4], vec![u32_value; 3]].concat(), + vec![i32_ty.into(); 4] + ); + assert!(verify_op(&crossed_shape, &ctx).is_err()); +} + +#[test] +fn generated_register_mma_verifies_dense_b1_families() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + fn b1_mma( + ctx: &mut Context, + shape: RegisterMmaShapeAttr, + operation: Option, + ) -> RegisterMmaOp { + let (accumulator_count, a_count, b_count, result_count) = match shape { + RegisterMmaShapeAttr::M8n8k128 => (2, 1, 1, 2), + RegisterMmaShapeAttr::M16n8k128 => (4, 2, 1, 4), + RegisterMmaShapeAttr::M16n8k256 => (4, 4, 2, 4), + _ => panic!("unsupported B1 MMA shape"), + }; + let i32_ty = IntegerType::get(ctx, 32, Signedness::Signed); + let u32_ty = IntegerType::get(ctx, 32, Signedness::Unsigned); + let argument_types = (0..accumulator_count) + .map(|_| i32_ty.into()) + .chain((0..a_count + b_count).map(|_| u32_ty.into())) + .collect(); + let block = BasicBlock::new(ctx, None, argument_types); + let operands = (0..accumulator_count + a_count + b_count) + .map(|index| block.deref(ctx).get_argument(index)) + .collect(); + let op = Operation::new( + ctx, + RegisterMmaOp::get_concrete_op_info(), + vec![i32_ty.into(); result_count], + operands, + vec![], + 0, + ); + let mma = RegisterMmaOp::new(op); + mma.set_attr_nvvm_register_mma_shape(ctx, shape); + if let Some(operation) = operation { + mma.set_attr_nvvm_register_mma_operation(ctx, operation); + } + mma.set_attr_nvvm_register_mma_accumulator(ctx, RegisterMmaAccumulatorAttr::S32); + mma.set_attr_nvvm_register_mma_a_element(ctx, RegisterMmaElementAttr::B1); + mma.set_attr_nvvm_register_mma_b_element(ctx, RegisterMmaElementAttr::B1); + mma.set_attr_nvvm_register_mma_a_layout(ctx, RegisterMmaLayoutAttr::Row); + mma.set_attr_nvvm_register_mma_b_layout(ctx, RegisterMmaLayoutAttr::Col); + mma.set_attr_nvvm_register_mma_overflow(ctx, RegisterMmaOverflowAttr::Wrapping); + mma + } + + let mut accepted = 0; + for shape in [ + RegisterMmaShapeAttr::M8n8k128, + RegisterMmaShapeAttr::M16n8k128, + RegisterMmaShapeAttr::M16n8k256, + ] { + for operation in [ + RegisterMmaOperationAttr::XorPopc, + RegisterMmaOperationAttr::AndPopc, + ] { + let mma = b1_mma(&mut ctx, shape.clone(), Some(operation)); + assert!(verify_op(&mma, &ctx).is_ok(), "rejected {shape:?}"); + accepted += 1; + } + } + assert_eq!(accepted, 6); + + let multiply = b1_mma( + &mut ctx, + RegisterMmaShapeAttr::M8n8k128, + Some(RegisterMmaOperationAttr::Multiply), + ); + assert!(verify_op(&multiply, &ctx).is_err()); + + let wrong_shape = b1_mma( + &mut ctx, + RegisterMmaShapeAttr::M16n8k128, + Some(RegisterMmaOperationAttr::XorPopc), + ); + wrong_shape.set_attr_nvvm_register_mma_shape(&ctx, RegisterMmaShapeAttr::M16n8k64); + assert!(verify_op(&wrong_shape, &ctx).is_err()); + + let missing_operation = b1_mma(&mut ctx, RegisterMmaShapeAttr::M8n8k128, None); + assert!(verify_op(&missing_operation, &ctx).is_err()); +} + +#[test] +fn generated_sparse_mma_verifies_all_int8_variants_and_metadata_modes() { + use dialect_mir::ops::MirConstantOp; + use pliron::builtin::{attributes::IntegerAttr, ops::ConstantOp}; + use pliron::utils::apint::APInt; + use std::num::NonZeroUsize; + + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let signless_i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let block = BasicBlock::new( + &mut ctx, + None, + vec![i32_ty.into(), u32_ty.into(), signless_i32_ty.into()], + ); + let i32_value = block.deref(&ctx).get_argument(0); + let u32_value = block.deref(&ctx).get_argument(1); + let signless_i32_value = block.deref(&ctx).get_argument(2); + + let integer = |value| { + IntegerAttr::new( + u32_ty, + APInt::from_u32(value, NonZeroUsize::new(32).unwrap()), + ) + }; + let builtin_zero = ConstantOp::new(&mut ctx, integer(0).into()); + let builtin_zero = builtin_zero.get_operation().deref(&ctx).get_result(0); + let builtin_two = ConstantOp::new(&mut ctx, integer(2).into()); + let builtin_two = builtin_two.get_operation().deref(&ctx).get_result(0); + let mir_one = Operation::new( + &mut ctx, + MirConstantOp::get_concrete_op_info(), + vec![u32_ty.into()], + vec![], + vec![], + 0, + ); + MirConstantOp::new(mir_one).set_attr_value(&ctx, integer(1)); + let mir_one = mir_one.deref(&ctx).get_result(0); + + macro_rules! sparse_mma { + ($operands:expr, $results:expr, $a:expr, $b:expr, $overflow:expr, $metadata:expr) => {{ + let operation = Operation::new( + &mut ctx, + SparseMmaOp::get_concrete_op_info(), + $results, + $operands, + vec![], + 0, + ); + let mma = SparseMmaOp::new(operation); + mma.set_attr_nvvm_sparse_mma_shape(&ctx, SparseMmaShapeAttr::M16n8k32); + mma.set_attr_nvvm_sparse_mma_accumulator(&ctx, SparseMmaAccumulatorAttr::S32); + mma.set_attr_nvvm_sparse_mma_a_element(&ctx, $a); + mma.set_attr_nvvm_sparse_mma_b_element(&ctx, $b); + mma.set_attr_nvvm_sparse_mma_a_layout(&ctx, SparseMmaLayoutAttr::Row); + mma.set_attr_nvvm_sparse_mma_b_layout(&ctx, SparseMmaLayoutAttr::Col); + mma.set_attr_nvvm_sparse_mma_overflow(&ctx, $overflow); + mma.set_attr_nvvm_sparse_mma_metadata(&ctx, $metadata); + mma.set_attr_nvvm_sparse_mma_selector(&ctx, SparseMmaSelectorAttr::ImmediateZeroOrOne); + mma + }}; + } + + let variants = [ + ( + SparseMmaElementAttr::S8, + SparseMmaElementAttr::S8, + SparseMmaOverflowAttr::Wrapping, + ), + ( + SparseMmaElementAttr::S8, + SparseMmaElementAttr::U8, + SparseMmaOverflowAttr::Wrapping, + ), + ( + SparseMmaElementAttr::U8, + SparseMmaElementAttr::U8, + SparseMmaOverflowAttr::Wrapping, + ), + ( + SparseMmaElementAttr::U8, + SparseMmaElementAttr::S8, + SparseMmaOverflowAttr::Wrapping, + ), + ( + SparseMmaElementAttr::S8, + SparseMmaElementAttr::S8, + SparseMmaOverflowAttr::Satfinite, + ), + ( + SparseMmaElementAttr::S8, + SparseMmaElementAttr::U8, + SparseMmaOverflowAttr::Satfinite, + ), + ( + SparseMmaElementAttr::U8, + SparseMmaElementAttr::U8, + SparseMmaOverflowAttr::Satfinite, + ), + ( + SparseMmaElementAttr::U8, + SparseMmaElementAttr::S8, + SparseMmaOverflowAttr::Satfinite, + ), + ]; + for metadata in [ + SparseMmaMetadataAttr::Standard, + SparseMmaMetadataAttr::Ordered, + ] { + for (index, (a_element, b_element, overflow)) in variants.iter().enumerate() { + let selector = if index % 2 == 0 { + builtin_zero + } else { + mir_one + }; + let operands = [vec![i32_value; 4], vec![u32_value; 5], vec![selector]].concat(); + let mma = sparse_mma!( + operands, + vec![i32_ty.into(); 4], + a_element.clone(), + b_element.clone(), + overflow.clone(), + metadata.clone() + ); + assert_eq!( + mma.get_attr_nvvm_sparse_mma_a_element(&ctx).as_deref(), + Some(a_element) + ); + assert_eq!( + mma.get_attr_nvvm_sparse_mma_b_element(&ctx).as_deref(), + Some(b_element) + ); + assert_eq!( + mma.get_attr_nvvm_sparse_mma_overflow(&ctx).as_deref(), + Some(overflow) + ); + assert_eq!( + mma.get_attr_nvvm_sparse_mma_metadata(&ctx).as_deref(), + Some(&metadata) + ); + assert!( + verify_op(&mma, &ctx).is_ok(), + "rejected sparse {metadata:?} {a_element:?}x{b_element:?} {overflow:?}" + ); + } + } + + for metadata in [ + SparseMmaMetadataAttr::Standard, + SparseMmaMetadataAttr::Ordered, + ] { + for selector in [u32_value, builtin_two] { + let invalid = sparse_mma!( + [vec![i32_value; 4], vec![u32_value; 5], vec![selector],].concat(), + vec![i32_ty.into(); 4], + SparseMmaElementAttr::S8, + SparseMmaElementAttr::U8, + SparseMmaOverflowAttr::Wrapping, + metadata.clone() + ); + assert!(verify_op(&invalid, &ctx).is_err()); + } + } + + let wrong_accumulator_type = sparse_mma!( + [ + vec![signless_i32_value; 4], + vec![u32_value; 5], + vec![builtin_zero], + ] + .concat(), + vec![i32_ty.into(); 4], + SparseMmaElementAttr::U8, + SparseMmaElementAttr::S8, + SparseMmaOverflowAttr::Satfinite, + SparseMmaMetadataAttr::Standard + ); + assert!(verify_op(&wrong_accumulator_type, &ctx).is_err()); + + let wrong_count = sparse_mma!( + [vec![i32_value; 4], vec![u32_value; 4], vec![builtin_zero],].concat(), + vec![i32_ty.into(); 4], + SparseMmaElementAttr::S8, + SparseMmaElementAttr::S8, + SparseMmaOverflowAttr::Wrapping, + SparseMmaMetadataAttr::Standard + ); + assert!(verify_op(&wrong_count, &ctx).is_err()); + + let wrong_results = sparse_mma!( + [vec![i32_value; 4], vec![u32_value; 5], vec![mir_one],].concat(), + vec![signless_i32_ty.into(); 4], + SparseMmaElementAttr::U8, + SparseMmaElementAttr::U8, + SparseMmaOverflowAttr::Wrapping, + SparseMmaMetadataAttr::Standard + ); + assert!(verify_op(&wrong_results, &ctx).is_err()); + + let wrong_layout = sparse_mma!( + [vec![i32_value; 4], vec![u32_value; 5], vec![builtin_zero],].concat(), + vec![i32_ty.into(); 4], + SparseMmaElementAttr::S8, + SparseMmaElementAttr::U8, + SparseMmaOverflowAttr::Satfinite, + SparseMmaMetadataAttr::Standard + ); + wrong_layout.set_attr_nvvm_sparse_mma_b_layout(&ctx, SparseMmaLayoutAttr::Row); + assert!(verify_op(&wrong_layout, &ctx).is_err()); +} + +#[test] +fn generated_sparse_mma_m16n8k64_verifies_selector_and_carriers() { + use pliron::builtin::{attributes::IntegerAttr, ops::ConstantOp}; + use pliron::utils::apint::APInt; + use std::num::NonZeroUsize; + + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let signless_i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let block = BasicBlock::new( + &mut ctx, + None, + vec![i32_ty.into(), u32_ty.into(), signless_i32_ty.into()], + ); + let i32_value = block.deref(&ctx).get_argument(0); + let u32_value = block.deref(&ctx).get_argument(1); + let signless_i32_value = block.deref(&ctx).get_argument(2); + let integer = |value| { + IntegerAttr::new( + u32_ty, + APInt::from_u32(value, NonZeroUsize::new(32).unwrap()), + ) + }; + let zero = ConstantOp::new(&mut ctx, integer(0).into()) + .get_operation() + .deref(&ctx) + .get_result(0); + let one = ConstantOp::new(&mut ctx, integer(1).into()) + .get_operation() + .deref(&ctx) + .get_result(0); + let two = ConstantOp::new(&mut ctx, integer(2).into()) + .get_operation() + .deref(&ctx) + .get_result(0); + + macro_rules! k64_mma { + ($operands:expr, $metadata:expr) => {{ + let operation = Operation::new( + &mut ctx, + SparseMmaOp::get_concrete_op_info(), + vec![i32_ty.into(); 4], + $operands, + vec![], + 0, + ); + let mma = SparseMmaOp::new(operation); + mma.set_attr_nvvm_sparse_mma_shape(&ctx, SparseMmaShapeAttr::M16n8k64); + mma.set_attr_nvvm_sparse_mma_accumulator(&ctx, SparseMmaAccumulatorAttr::S32); + mma.set_attr_nvvm_sparse_mma_a_element(&ctx, SparseMmaElementAttr::S8); + mma.set_attr_nvvm_sparse_mma_b_element(&ctx, SparseMmaElementAttr::U8); + mma.set_attr_nvvm_sparse_mma_a_layout(&ctx, SparseMmaLayoutAttr::Row); + mma.set_attr_nvvm_sparse_mma_b_layout(&ctx, SparseMmaLayoutAttr::Col); + mma.set_attr_nvvm_sparse_mma_overflow(&ctx, SparseMmaOverflowAttr::Wrapping); + mma.set_attr_nvvm_sparse_mma_metadata(&ctx, $metadata); + mma.set_attr_nvvm_sparse_mma_selector(&ctx, SparseMmaSelectorAttr::ImmediateZero); + mma + }}; + } + + let operands = |selector| [vec![i32_value; 4], vec![u32_value; 9], vec![selector]].concat(); + assert!( + verify_op( + &k64_mma!(operands(zero), SparseMmaMetadataAttr::Standard), + &ctx + ) + .is_ok() + ); + assert!( + verify_op( + &k64_mma!(operands(zero), SparseMmaMetadataAttr::Ordered), + &ctx + ) + .is_ok() + ); + assert!( + verify_op( + &k64_mma!(operands(one), SparseMmaMetadataAttr::Standard), + &ctx + ) + .is_err() + ); + assert!( + verify_op( + &k64_mma!(operands(one), SparseMmaMetadataAttr::Ordered), + &ctx + ) + .is_err() + ); + assert!( + verify_op( + &k64_mma!(operands(u32_value), SparseMmaMetadataAttr::Standard), + &ctx + ) + .is_err() + ); + assert!( + verify_op( + &k64_mma!(operands(u32_value), SparseMmaMetadataAttr::Ordered), + &ctx + ) + .is_err() + ); + + let wrong_count = [vec![i32_value; 4], vec![u32_value; 8], vec![zero]].concat(); + assert!( + verify_op( + &k64_mma!(wrong_count, SparseMmaMetadataAttr::Standard), + &ctx + ) + .is_err() + ); + + let mut wrong_type = operands(zero); + wrong_type[4] = signless_i32_value; + assert!(verify_op(&k64_mma!(wrong_type, SparseMmaMetadataAttr::Standard), &ctx).is_err()); + + macro_rules! int4_mma { + ($operands:expr, $a:expr, $b:expr, $overflow:expr, $metadata:expr) => {{ + let operation = Operation::new( + &mut ctx, + SparseMmaOp::get_concrete_op_info(), + vec![i32_ty.into(); 4], + $operands, + vec![], + 0, + ); + let mma = SparseMmaOp::new(operation); + mma.set_attr_nvvm_sparse_mma_shape(&ctx, SparseMmaShapeAttr::M16n8k64); + mma.set_attr_nvvm_sparse_mma_accumulator(&ctx, SparseMmaAccumulatorAttr::S32); + mma.set_attr_nvvm_sparse_mma_a_element(&ctx, $a); + mma.set_attr_nvvm_sparse_mma_b_element(&ctx, $b); + mma.set_attr_nvvm_sparse_mma_a_layout(&ctx, SparseMmaLayoutAttr::Row); + mma.set_attr_nvvm_sparse_mma_b_layout(&ctx, SparseMmaLayoutAttr::Col); + mma.set_attr_nvvm_sparse_mma_overflow(&ctx, $overflow); + mma.set_attr_nvvm_sparse_mma_metadata(&ctx, $metadata); + mma.set_attr_nvvm_sparse_mma_selector(&ctx, SparseMmaSelectorAttr::ImmediateZeroOrOne); + mma + }}; + } + + let int4_variants = [ + ( + SparseMmaElementAttr::S4, + SparseMmaElementAttr::S4, + SparseMmaOverflowAttr::Wrapping, + ), + ( + SparseMmaElementAttr::S4, + SparseMmaElementAttr::U4, + SparseMmaOverflowAttr::Wrapping, + ), + ( + SparseMmaElementAttr::U4, + SparseMmaElementAttr::U4, + SparseMmaOverflowAttr::Wrapping, + ), + ( + SparseMmaElementAttr::U4, + SparseMmaElementAttr::S4, + SparseMmaOverflowAttr::Wrapping, + ), + ( + SparseMmaElementAttr::S4, + SparseMmaElementAttr::S4, + SparseMmaOverflowAttr::Satfinite, + ), + ( + SparseMmaElementAttr::S4, + SparseMmaElementAttr::U4, + SparseMmaOverflowAttr::Satfinite, + ), + ( + SparseMmaElementAttr::U4, + SparseMmaElementAttr::U4, + SparseMmaOverflowAttr::Satfinite, + ), + ( + SparseMmaElementAttr::U4, + SparseMmaElementAttr::S4, + SparseMmaOverflowAttr::Satfinite, + ), + ]; + let int4_operands = + |selector| [vec![i32_value; 4], vec![u32_value; 5], vec![selector]].concat(); + for metadata in [ + SparseMmaMetadataAttr::Standard, + SparseMmaMetadataAttr::Ordered, + ] { + for (index, (a, b, overflow)) in int4_variants.iter().enumerate() { + let selector = if index % 2 == 0 { zero } else { one }; + assert!( + verify_op( + &int4_mma!( + int4_operands(selector), + a.clone(), + b.clone(), + overflow.clone(), + metadata.clone() + ), + &ctx, + ) + .is_ok() + ); + } + } + + assert!( + verify_op( + &int4_mma!( + int4_operands(zero), + SparseMmaElementAttr::S4, + SparseMmaElementAttr::U8, + SparseMmaOverflowAttr::Wrapping, + SparseMmaMetadataAttr::Standard + ), + &ctx, + ) + .is_err() + ); + for metadata in [ + SparseMmaMetadataAttr::Standard, + SparseMmaMetadataAttr::Ordered, + ] { + for selector in [two, u32_value] { + assert!( + verify_op( + &int4_mma!( + int4_operands(selector), + SparseMmaElementAttr::S4, + SparseMmaElementAttr::U4, + SparseMmaOverflowAttr::Wrapping, + metadata.clone() + ), + &ctx, + ) + .is_err() + ); + } + } +} + +#[test] +fn generated_sparse_mma_m16n8k128_int4_verifies_metadata_selector_and_widths() { + use pliron::builtin::{attributes::IntegerAttr, ops::ConstantOp}; + use pliron::utils::apint::APInt; + use std::num::NonZeroUsize; + + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into(), u32_ty.into()]); + let i32_value = block.deref(&ctx).get_argument(0); + let u32_value = block.deref(&ctx).get_argument(1); + let integer = |value| { + IntegerAttr::new( + u32_ty, + APInt::from_u32(value, NonZeroUsize::new(32).unwrap()), + ) + }; + let zero = ConstantOp::new(&mut ctx, integer(0).into()) + .get_operation() + .deref(&ctx) + .get_result(0); + let one = ConstantOp::new(&mut ctx, integer(1).into()) + .get_operation() + .deref(&ctx) + .get_result(0); + + macro_rules! k128_mma { + ($operands:expr, $a:expr, $b:expr, $overflow:expr, $metadata:expr) => {{ + let operation = Operation::new( + &mut ctx, + SparseMmaOp::get_concrete_op_info(), + vec![i32_ty.into(); 4], + $operands, + vec![], + 0, + ); + let mma = SparseMmaOp::new(operation); + mma.set_attr_nvvm_sparse_mma_shape(&ctx, SparseMmaShapeAttr::M16n8k128); + mma.set_attr_nvvm_sparse_mma_accumulator(&ctx, SparseMmaAccumulatorAttr::S32); + mma.set_attr_nvvm_sparse_mma_a_element(&ctx, $a); + mma.set_attr_nvvm_sparse_mma_b_element(&ctx, $b); + mma.set_attr_nvvm_sparse_mma_a_layout(&ctx, SparseMmaLayoutAttr::Row); + mma.set_attr_nvvm_sparse_mma_b_layout(&ctx, SparseMmaLayoutAttr::Col); + mma.set_attr_nvvm_sparse_mma_overflow(&ctx, $overflow); + mma.set_attr_nvvm_sparse_mma_metadata(&ctx, $metadata); + mma.set_attr_nvvm_sparse_mma_selector(&ctx, SparseMmaSelectorAttr::ImmediateZero); + mma + }}; + } + + let variants = [ + ( + SparseMmaElementAttr::S4, + SparseMmaElementAttr::S4, + SparseMmaOverflowAttr::Wrapping, + ), + ( + SparseMmaElementAttr::S4, + SparseMmaElementAttr::U4, + SparseMmaOverflowAttr::Wrapping, + ), + ( + SparseMmaElementAttr::U4, + SparseMmaElementAttr::U4, + SparseMmaOverflowAttr::Wrapping, + ), + ( + SparseMmaElementAttr::U4, + SparseMmaElementAttr::S4, + SparseMmaOverflowAttr::Wrapping, + ), + ( + SparseMmaElementAttr::S4, + SparseMmaElementAttr::S4, + SparseMmaOverflowAttr::Satfinite, + ), + ( + SparseMmaElementAttr::S4, + SparseMmaElementAttr::U4, + SparseMmaOverflowAttr::Satfinite, + ), + ( + SparseMmaElementAttr::U4, + SparseMmaElementAttr::U4, + SparseMmaOverflowAttr::Satfinite, + ), + ( + SparseMmaElementAttr::U4, + SparseMmaElementAttr::S4, + SparseMmaOverflowAttr::Satfinite, + ), + ]; + let operands = |selector| [vec![i32_value; 4], vec![u32_value; 9], vec![selector]].concat(); + for metadata in [ + SparseMmaMetadataAttr::Standard, + SparseMmaMetadataAttr::Ordered, + ] { + for (a, b, overflow) in &variants { + assert!( + verify_op( + &k128_mma!( + operands(zero), + a.clone(), + b.clone(), + overflow.clone(), + metadata.clone() + ), + &ctx, + ) + .is_ok() + ); + } + } + + for metadata in [ + SparseMmaMetadataAttr::Standard, + SparseMmaMetadataAttr::Ordered, + ] { + for selector in [one, u32_value] { + assert!( + verify_op( + &k128_mma!( + operands(selector), + SparseMmaElementAttr::S4, + SparseMmaElementAttr::U4, + SparseMmaOverflowAttr::Wrapping, + metadata.clone() + ), + &ctx, + ) + .is_err() + ); + } + } + assert!( + verify_op( + &k128_mma!( + operands(zero), + SparseMmaElementAttr::S4, + SparseMmaElementAttr::U8, + SparseMmaOverflowAttr::Wrapping, + SparseMmaMetadataAttr::Standard + ), + &ctx, + ) + .is_err() + ); + assert!( + verify_op( + &k128_mma!( + [vec![i32_value; 4], vec![u32_value; 8], vec![zero]].concat(), + SparseMmaElementAttr::S4, + SparseMmaElementAttr::S4, + SparseMmaOverflowAttr::Wrapping, + SparseMmaMetadataAttr::Standard + ), + &ctx, + ) + .is_err() + ); +} + +#[test] +fn generated_sparse_mma_f8f6f4_verifies_all_formats_and_closed_shape() { + use pliron::builtin::{attributes::IntegerAttr, ops::ConstantOp}; + use pliron::utils::apint::APInt; + use std::num::NonZeroUsize; + + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + + let f32_ty = FP32Type::get(&ctx); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let block = BasicBlock::new(&mut ctx, None, vec![f32_ty.into(), u32_ty.into()]); + let f32_value = block.deref(&ctx).get_argument(0); + let u32_value = block.deref(&ctx).get_argument(1); + let integer = |value| { + IntegerAttr::new( + u32_ty, + APInt::from_u32(value, NonZeroUsize::new(32).unwrap()), + ) + }; + let zero = ConstantOp::new(&mut ctx, integer(0).into()) + .get_operation() + .deref(&ctx) + .get_result(0); + let one = ConstantOp::new(&mut ctx, integer(1).into()) + .get_operation() + .deref(&ctx) + .get_result(0); + + macro_rules! f8f6f4_mma { + ($operands:expr, $results:expr, $a:expr, $b:expr, $overflow:expr, $metadata:expr) => {{ + let operation = Operation::new( + &mut ctx, + SparseMmaOp::get_concrete_op_info(), + $results, + $operands, + vec![], + 0, + ); + let mma = SparseMmaOp::new(operation); + mma.set_attr_nvvm_sparse_mma_shape(&ctx, SparseMmaShapeAttr::M16n8k64); + mma.set_attr_nvvm_sparse_mma_accumulator(&ctx, SparseMmaAccumulatorAttr::F32); + mma.set_attr_nvvm_sparse_mma_a_element(&ctx, $a); + mma.set_attr_nvvm_sparse_mma_b_element(&ctx, $b); + mma.set_attr_nvvm_sparse_mma_a_layout(&ctx, SparseMmaLayoutAttr::Row); + mma.set_attr_nvvm_sparse_mma_b_layout(&ctx, SparseMmaLayoutAttr::Col); + mma.set_attr_nvvm_sparse_mma_overflow(&ctx, $overflow); + mma.set_attr_nvvm_sparse_mma_metadata(&ctx, $metadata); + mma.set_attr_nvvm_sparse_mma_selector(&ctx, SparseMmaSelectorAttr::ImmediateZero); + mma + }}; + } + + let elements = [ + SparseMmaElementAttr::E2m1, + SparseMmaElementAttr::E2m3, + SparseMmaElementAttr::E3m2, + SparseMmaElementAttr::E4m3, + SparseMmaElementAttr::E5m2, + ]; + let operands = |selector| [vec![f32_value; 4], vec![u32_value; 9], vec![selector]].concat(); + for a in &elements { + for b in &elements { + let mma = f8f6f4_mma!( + operands(zero), + vec![f32_ty.into(); 4], + a.clone(), + b.clone(), + SparseMmaOverflowAttr::NotApplicable, + SparseMmaMetadataAttr::Ordered + ); + assert!(verify_op(&mma, &ctx).is_ok(), "rejected {a:?}x{b:?}"); + } + } + + for invalid in [ + f8f6f4_mma!( + operands(one), + vec![f32_ty.into(); 4], + SparseMmaElementAttr::E2m1, + SparseMmaElementAttr::E2m1, + SparseMmaOverflowAttr::NotApplicable, + SparseMmaMetadataAttr::Ordered + ), + f8f6f4_mma!( + operands(zero), + vec![f32_ty.into(); 4], + SparseMmaElementAttr::E2m1, + SparseMmaElementAttr::E2m1, + SparseMmaOverflowAttr::Wrapping, + SparseMmaMetadataAttr::Ordered + ), + f8f6f4_mma!( + operands(zero), + vec![f32_ty.into(); 4], + SparseMmaElementAttr::E2m1, + SparseMmaElementAttr::E2m1, + SparseMmaOverflowAttr::NotApplicable, + SparseMmaMetadataAttr::Standard + ), + f8f6f4_mma!( + operands(zero), + vec![u32_ty.into(); 4], + SparseMmaElementAttr::E2m1, + SparseMmaElementAttr::E2m1, + SparseMmaOverflowAttr::NotApplicable, + SparseMmaMetadataAttr::Ordered + ), + ] { + assert!(verify_op(&invalid, &ctx).is_err()); + } +} + +#[test] +fn test_matrix_memory_ops_verify_pointer_and_packed_register_types() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); + let f32_ty = FP32Type::get(&ctx); + let ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), true); + let load_ptr_ty = MirPtrType::get_generic(&mut ctx, u32_ty.into(), true); + + let load_block = BasicBlock::new(&mut ctx, None, vec![load_ptr_ty.into()]); + let load_pointer = load_block.deref(&ctx).get_argument(0); + let load = make_ldmatrix_x2(&mut ctx, load_pointer, vec![u32_ty.into(), u32_ty.into()]); + assert!(load.verify(&ctx).is_ok()); + + let bad_load_pointer_block = BasicBlock::new(&mut ctx, None, vec![i64_ty.into()]); + let bad_pointer = bad_load_pointer_block.deref(&ctx).get_argument(0); + let bad_load_pointer = + make_ldmatrix_x2(&mut ctx, bad_pointer, vec![u32_ty.into(), u32_ty.into()]); + assert!(bad_load_pointer.verify(&ctx).is_err()); + + let bad_load_result = + make_ldmatrix_x2(&mut ctx, load_pointer, vec![u32_ty.into(), f32_ty.into()]); + assert!(bad_load_result.verify(&ctx).is_err()); + + let store_block = BasicBlock::new( + &mut ctx, + None, + vec![ + ptr_ty.into(), + i32_ty.into(), + i32_ty.into(), + i32_ty.into(), + i32_ty.into(), + ], + ); + let store_operands = (0..5) + .map(|index| store_block.deref(&ctx).get_argument(index)) + .collect(); + let store = Operation::new( + &mut ctx, + StmatrixM8n8X4Op::get_concrete_op_info(), + vec![], + store_operands, + vec![], + 0, + ); + assert!(StmatrixM8n8X4Op::new(store).verify(&ctx).is_ok()); + + let bad_store_block = BasicBlock::new( + &mut ctx, + None, + vec![ + ptr_ty.into(), + f32_ty.into(), + i32_ty.into(), + i32_ty.into(), + i32_ty.into(), + ], + ); + let bad_store_operands = (0..5) + .map(|index| bad_store_block.deref(&ctx).get_argument(index)) + .collect(); + let bad_store = Operation::new( + &mut ctx, + StmatrixM8n8X4Op::get_concrete_op_info(), + vec![], + bad_store_operands, + vec![], + 0, + ); + assert!(StmatrixM8n8X4Op::new(bad_store).verify(&ctx).is_err()); +} diff --git a/crates/dialect-nvvm/tests/ops_test/thread_and_warp.rs b/crates/dialect-nvvm/tests/ops_test/thread_and_warp.rs new file mode 100644 index 0000000000..4b37d837ff --- /dev/null +++ b/crates/dialect-nvvm/tests/ops_test/thread_and_warp.rs @@ -0,0 +1,953 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use dialect_mir::types::{MirPointerKind, MirPtrType}; +use dialect_nvvm::ops::{ + ActiveMaskOp, BarWarpSyncOp, Barrier0Op, ClcQueryIsCanceledOp, Dp2aS32Op, Dp2aU32Op, Dp4aS32Op, + Dp4aU32Op, ElectSyncOp, FmaBf16x2Op, MatchAllSyncI32Op, MatchAllSyncI64Op, MatchAnySyncI32Op, + MatchAnySyncI64Op, ReadPtxSregDynamicSmemSizeOp, ReadPtxSregGridIdOp, ReadPtxSregLaneIdOp, + ReadPtxSregLanemaskEqOp, ReadPtxSregLanemaskGeOp, ReadPtxSregLanemaskGtOp, + ReadPtxSregLanemaskLeOp, ReadPtxSregLanemaskLtOp, ReadPtxSregNsmIdOp, ReadPtxSregNwarpIdOp, + ReadPtxSregSmIdOp, ReadPtxSregTidXOp, ReadPtxSregTotalSmemSizeOp, ReadPtxSregWarpIdOp, + ReduxSyncAddOp, ReduxSyncAndOp, ReduxSyncFmaxAbsNanOp, ReduxSyncFmaxAbsOp, ReduxSyncFmaxNanOp, + ReduxSyncFmaxOp, ReduxSyncFminAbsNanOp, ReduxSyncFminAbsOp, ReduxSyncFminNanOp, + ReduxSyncFminOp, ReduxSyncMaxOp, ReduxSyncMinOp, ReduxSyncOrOp, ReduxSyncUmaxOp, + ReduxSyncUminOp, ReduxSyncXorOp, ShflSyncBflyI64Op, ShflSyncDownI64Op, ShflSyncIdxI64Op, + ShflSyncUpI64Op, ThreadfenceBlockOp, ThreadfenceOp, ThreadfenceSystemOp, VoteSyncAllOp, + VoteSyncAnyOp, VoteSyncBallotOp, VoteSyncUniOp, +}; + +use pliron::{ + basic_block::BasicBlock, + builtin::types::{FP32Type, IntegerType, Signedness}, + common_traits::Verify, + context::Context, + op::{Op, verify_op}, + operation::Operation, +}; + +#[test] +fn test_thread_register_ops_verify_i32_results() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + + let tid_x = Operation::new( + &mut ctx, + ReadPtxSregTidXOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![], + vec![], + 0, + ); + assert!(ReadPtxSregTidXOp::new(tid_x).verify(&ctx).is_ok()); + + let lane_id = Operation::new( + &mut ctx, + ReadPtxSregLaneIdOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![], + vec![], + 0, + ); + assert!(ReadPtxSregLaneIdOp::new(lane_id).verify(&ctx).is_ok()); +} + +#[test] +fn test_thread_register_ops_reject_non_i32_results() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); + let op = Operation::new( + &mut ctx, + ReadPtxSregTidXOp::get_concrete_op_info(), + vec![i64_ty.into()], + vec![], + vec![], + 0, + ); + + assert!(ReadPtxSregTidXOp::new(op).verify(&ctx).is_err()); +} + +#[test] +fn test_lanemask_ops_verify_i32_results() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + + // Each lane-position mask is a zero-operand, single-i32-result sreg read. + let lt = Operation::new( + &mut ctx, + ReadPtxSregLanemaskLtOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![], + vec![], + 0, + ); + assert!(ReadPtxSregLanemaskLtOp::new(lt).verify(&ctx).is_ok()); + + let le = Operation::new( + &mut ctx, + ReadPtxSregLanemaskLeOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![], + vec![], + 0, + ); + assert!(ReadPtxSregLanemaskLeOp::new(le).verify(&ctx).is_ok()); + + let eq = Operation::new( + &mut ctx, + ReadPtxSregLanemaskEqOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![], + vec![], + 0, + ); + assert!(ReadPtxSregLanemaskEqOp::new(eq).verify(&ctx).is_ok()); + + let ge = Operation::new( + &mut ctx, + ReadPtxSregLanemaskGeOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![], + vec![], + 0, + ); + assert!(ReadPtxSregLanemaskGeOp::new(ge).verify(&ctx).is_ok()); + + let gt = Operation::new( + &mut ctx, + ReadPtxSregLanemaskGtOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![], + vec![], + 0, + ); + assert!(ReadPtxSregLanemaskGtOp::new(gt).verify(&ctx).is_ok()); +} + +#[test] +fn test_lanemask_op_rejects_non_i32_result() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + // A 64-bit result must fail the shared lane-position mask verifier. + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); + let op = Operation::new( + &mut ctx, + ReadPtxSregLanemaskLtOp::get_concrete_op_info(), + vec![i64_ty.into()], + vec![], + vec![], + 0, + ); + assert!(ReadPtxSregLanemaskLtOp::new(op).verify(&ctx).is_err()); +} + +#[test] +fn test_generated_vote_sync_family_requires_exact_mask_predicate_and_result_types() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); + let block = BasicBlock::new( + &mut ctx, + None, + vec![i32_ty.into(), i1_ty.into(), i64_ty.into()], + ); + let mask = block.deref(&ctx).get_argument(0); + let predicate = block.deref(&ctx).get_argument(1); + let wide_mask = block.deref(&ctx).get_argument(2); + + macro_rules! check_vote { + ($op:ty, $result_ty:expr, $wrong_result_ty:expr) => {{ + let valid = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + vec![$result_ty.into()], + vec![mask, predicate], + vec![], + 0, + ); + assert!(verify_op(&<$op>::new(valid), &ctx).is_ok()); + + for operands in [vec![], vec![mask], vec![mask, predicate, predicate]] { + let wrong_arity = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + vec![$result_ty.into()], + operands, + vec![], + 0, + ); + assert!(verify_op(&<$op>::new(wrong_arity), &ctx).is_err()); + } + + for results in [vec![], vec![$result_ty.into(), $result_ty.into()]] { + let wrong_arity = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + results, + vec![mask, predicate], + vec![], + 0, + ); + assert!(verify_op(&<$op>::new(wrong_arity), &ctx).is_err()); + } + + let wrong_mask = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + vec![$result_ty.into()], + vec![wide_mask, predicate], + vec![], + 0, + ); + assert!(verify_op(&<$op>::new(wrong_mask), &ctx).is_err()); + + let wrong_predicate = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + vec![$result_ty.into()], + vec![mask, mask], + vec![], + 0, + ); + assert!(verify_op(&<$op>::new(wrong_predicate), &ctx).is_err()); + + let wrong_result = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + vec![$wrong_result_ty.into()], + vec![mask, predicate], + vec![], + 0, + ); + assert!(verify_op(&<$op>::new(wrong_result), &ctx).is_err()); + }}; + } + + check_vote!(VoteSyncAllOp, i1_ty, u32_ty); + check_vote!(VoteSyncAnyOp, i1_ty, u32_ty); + check_vote!(VoteSyncBallotOp, u32_ty, i1_ty); + check_vote!(VoteSyncUniOp, i1_ty, u32_ty); +} + +#[test] +fn test_generated_active_mask_requires_no_operands_and_i32_result() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); + let block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); + let unexpected_operand = block.deref(&ctx).get_argument(0); + + let valid = Operation::new( + &mut ctx, + ActiveMaskOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![], + vec![], + 0, + ); + assert!(verify_op(&ActiveMaskOp::new(valid), &ctx).is_ok()); + + let wrong_operand_count = Operation::new( + &mut ctx, + ActiveMaskOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![unexpected_operand], + vec![], + 0, + ); + assert!(verify_op(&ActiveMaskOp::new(wrong_operand_count), &ctx).is_err()); + + for results in [vec![], vec![i32_ty.into(), i32_ty.into()]] { + let wrong_result_count = Operation::new( + &mut ctx, + ActiveMaskOp::get_concrete_op_info(), + results, + vec![], + vec![], + 0, + ); + assert!(verify_op(&ActiveMaskOp::new(wrong_result_count), &ctx).is_err()); + } + + let wrong_result_width = Operation::new( + &mut ctx, + ActiveMaskOp::get_concrete_op_info(), + vec![i64_ty.into()], + vec![], + vec![], + 0, + ); + assert!(verify_op(&ActiveMaskOp::new(wrong_result_width), &ctx).is_err()); +} + +#[test] +fn test_generated_match_family_requires_exact_mask_value_and_result_widths() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); + let block = BasicBlock::new( + &mut ctx, + None, + vec![i32_ty.into(), i32_ty.into(), i64_ty.into()], + ); + let mask = block.deref(&ctx).get_argument(0); + let value32 = block.deref(&ctx).get_argument(1); + let value64 = block.deref(&ctx).get_argument(2); + + macro_rules! check_match { + ($op:ty, $value:expr, $wrong_value:expr) => {{ + let valid = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + vec![i32_ty.into()], + vec![mask, $value], + vec![], + 0, + ); + assert!(verify_op(&<$op>::new(valid), &ctx).is_ok()); + + for operands in [vec![], vec![mask], vec![mask, $value, $value]] { + let wrong_operand_count = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + vec![i32_ty.into()], + operands, + vec![], + 0, + ); + assert!(verify_op(&<$op>::new(wrong_operand_count), &ctx).is_err()); + } + + for results in [vec![], vec![i32_ty.into(), i32_ty.into()]] { + let wrong_result_count = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + results, + vec![mask, $value], + vec![], + 0, + ); + assert!(verify_op(&<$op>::new(wrong_result_count), &ctx).is_err()); + } + + let wrong_mask_width = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + vec![i32_ty.into()], + vec![value64, $value], + vec![], + 0, + ); + assert!(verify_op(&<$op>::new(wrong_mask_width), &ctx).is_err()); + + let wrong_value_width = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + vec![i32_ty.into()], + vec![mask, $wrong_value], + vec![], + 0, + ); + assert!(verify_op(&<$op>::new(wrong_value_width), &ctx).is_err()); + + let wrong_result_width = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + vec![i64_ty.into()], + vec![mask, $value], + vec![], + 0, + ); + assert!(verify_op(&<$op>::new(wrong_result_width), &ctx).is_err()); + }}; + } + + check_match!(MatchAnySyncI32Op, value32, value64); + check_match!(MatchAllSyncI32Op, value32, value64); + check_match!(MatchAnySyncI64Op, value64, value32); + check_match!(MatchAllSyncI64Op, value64, value32); +} + +#[test] +fn test_special_register_ops_verify_authoritative_widths() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); + + macro_rules! check_width { + ($op:ty, $good:expr, $bad:expr) => {{ + let good = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + vec![$good.into()], + vec![], + vec![], + 0, + ); + assert!( + verify_op(&<$op>::new(good), &ctx).is_ok(), + "{} must accept its PTX register width", + stringify!($op) + ); + + let bad = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + vec![$bad.into()], + vec![], + vec![], + 0, + ); + assert!( + verify_op(&<$op>::new(bad), &ctx).is_err(), + "{} must reject the other integer width", + stringify!($op) + ); + }}; + } + + check_width!(ReadPtxSregWarpIdOp, i32_ty, i64_ty); + check_width!(ReadPtxSregNwarpIdOp, i32_ty, i64_ty); + check_width!(ReadPtxSregSmIdOp, i32_ty, i64_ty); + check_width!(ReadPtxSregNsmIdOp, i32_ty, i64_ty); + check_width!(ReadPtxSregDynamicSmemSizeOp, i32_ty, i64_ty); + check_width!(ReadPtxSregTotalSmemSizeOp, i32_ty, i64_ty); + check_width!(ReadPtxSregGridIdOp, i64_ty, i32_ty); +} + +#[test] +fn test_sync_ops_construct_and_verify() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + let barrier = Operation::new( + &mut ctx, + Barrier0Op::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + assert!(Barrier0Op::new(barrier).verify(&ctx).is_ok()); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); + let unexpected_operand = block.deref(&ctx).get_argument(0); + let bad_operand = Operation::new( + &mut ctx, + Barrier0Op::get_concrete_op_info(), + vec![], + vec![unexpected_operand], + vec![], + 0, + ); + assert!(verify_op(&Barrier0Op::new(bad_operand), &ctx).is_err()); + + let bad_result = Operation::new( + &mut ctx, + Barrier0Op::get_concrete_op_info(), + vec![i32_ty.into()], + vec![], + vec![], + 0, + ); + assert!(verify_op(&Barrier0Op::new(bad_result), &ctx).is_err()); + + let block_fence = Operation::new( + &mut ctx, + ThreadfenceBlockOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + assert!(ThreadfenceBlockOp::new(block_fence).verify(&ctx).is_ok()); + + let device_fence = Operation::new( + &mut ctx, + ThreadfenceOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + assert!(ThreadfenceOp::new(device_fence).verify(&ctx).is_ok()); + + let system_fence = Operation::new( + &mut ctx, + ThreadfenceSystemOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + assert!(ThreadfenceSystemOp::new(system_fence).verify(&ctx).is_ok()); +} + +#[test] +fn test_bf16x2_fma_constructs_and_verifies_three_operands() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + + let a = Operation::new( + &mut ctx, + ReadPtxSregTidXOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![], + vec![], + 0, + ); + let b = Operation::new( + &mut ctx, + ReadPtxSregTidXOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![], + vec![], + 0, + ); + let c = Operation::new( + &mut ctx, + ReadPtxSregTidXOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![], + vec![], + 0, + ); + + let operands = vec![ + a.deref(&ctx).get_result(0), + b.deref(&ctx).get_result(0), + c.deref(&ctx).get_result(0), + ]; + + let fma = Operation::new( + &mut ctx, + FmaBf16x2Op::get_concrete_op_info(), + vec![u32_ty.into()], + operands, + vec![], + 0, + ); + + assert!(FmaBf16x2Op::new(fma).verify(&ctx).is_ok()); +} + +#[test] +fn test_generated_dot_products_require_three_i32_operands_and_one_i32_result() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); + let block = BasicBlock::new( + &mut ctx, + None, + vec![i32_ty.into(), i32_ty.into(), i32_ty.into(), i64_ty.into()], + ); + let a = block.deref(&ctx).get_argument(0); + let b = block.deref(&ctx).get_argument(1); + let c = block.deref(&ctx).get_argument(2); + let wide = block.deref(&ctx).get_argument(3); + + macro_rules! check_variant { + ($op:ty) => {{ + let valid = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + vec![i32_ty.into()], + vec![a, b, c], + vec![], + 0, + ); + assert!(verify_op(&<$op>::new(valid), &ctx).is_ok()); + + let wrong_width = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + vec![i32_ty.into()], + vec![a, b, wide], + vec![], + 0, + ); + assert!(verify_op(&<$op>::new(wrong_width), &ctx).is_err()); + + for (results, operands) in [ + (vec![], vec![a, b, c]), + (vec![i32_ty.into(), i32_ty.into()], vec![a, b, c]), + (vec![i32_ty.into()], vec![a, b]), + (vec![i32_ty.into()], vec![a, b, c, c]), + ] { + let wrong_count = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + results, + operands, + vec![], + 0, + ); + assert!(verify_op(&<$op>::new(wrong_count), &ctx).is_err()); + } + }}; + } + + check_variant!(Dp4aS32Op); + check_variant!(Dp4aU32Op); + check_variant!(Dp2aS32Op); + check_variant!(Dp2aU32Op); +} + +#[test] +fn test_redux_sync_add_construct_and_verify() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + + // A block supplies the two operands [mask, value]. + let block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into(), i32_ty.into()]); + let mask = block.deref(&ctx).get_argument(0); + let value = block.deref(&ctx).get_argument(1); + + // Valid: 2 operands, 1 result (matches NOpdsInterface<2>/NResultsInterface<1>). + let op = Operation::new( + &mut ctx, + ReduxSyncAddOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![mask, value], + vec![], + 0, + ); + assert!(verify_op(&ReduxSyncAddOp::new(op), &ctx).is_ok()); + + // Invalid: wrong operand count (1 instead of 2) must fail verification. + let bad_opnds = Operation::new( + &mut ctx, + ReduxSyncAddOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![mask], + vec![], + 0, + ); + assert!(verify_op(&ReduxSyncAddOp::new(bad_opnds), &ctx).is_err()); + + // Invalid: wrong result count (0 instead of 1) must fail verification. + let bad_results = Operation::new( + &mut ctx, + ReduxSyncAddOp::get_concrete_op_info(), + vec![], + vec![mask, value], + vec![], + 0, + ); + assert!(verify_op(&ReduxSyncAddOp::new(bad_results), &ctx).is_err()); +} + +#[test] +fn test_redux_sync_integer_family_construct_and_verify() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into(), i32_ty.into()]); + let mask = block.deref(&ctx).get_argument(0); + let value = block.deref(&ctx).get_argument(1); + + // Every integer-family variant has the same 2-operand/1-result shape. A + // valid build of each must verify; a wrong operand count must not. The + // `new` wrapper is invoked so each concrete op type is exercised. + macro_rules! check_variant { + ($op:ty) => {{ + let good = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + vec![i32_ty.into()], + vec![mask, value], + vec![], + 0, + ); + assert!( + verify_op(&<$op>::new(good), &ctx).is_ok(), + "{} should verify with [mask, value] -> i32", + stringify!($op) + ); + + let bad = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + vec![i32_ty.into()], + vec![mask], + vec![], + 0, + ); + assert!( + verify_op(&<$op>::new(bad), &ctx).is_err(), + "{} must reject a single operand", + stringify!($op) + ); + }}; + } + + check_variant!(ReduxSyncUminOp); + check_variant!(ReduxSyncMinOp); + check_variant!(ReduxSyncUmaxOp); + check_variant!(ReduxSyncMaxOp); + check_variant!(ReduxSyncAndOp); + check_variant!(ReduxSyncOrOp); + check_variant!(ReduxSyncXorOp); +} + +#[test] +fn test_redux_sync_f32_family_construct_and_verify() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let f32_ty = FP32Type::get(&ctx); + let block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into(), f32_ty.into()]); + let mask = block.deref(&ctx).get_argument(0); + let value = block.deref(&ctx).get_argument(1); + + // The good case goes through the generated `build()` so its f32 result + // construction is exercised; the bad case's i32 value and result must + // fail the `is_f32` verifier path. + macro_rules! check_variant { + ($op:ty) => {{ + let good = <$op>::build(&mut ctx, mask, value); + assert!( + verify_op(&<$op>::new(good), &ctx).is_ok(), + "{} should verify with [mask, value] -> f32", + stringify!($op) + ); + + let bad = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + vec![i32_ty.into()], + vec![mask, mask], + vec![], + 0, + ); + assert!( + verify_op(&<$op>::new(bad), &ctx).is_err(), + "{} must reject an i32 value and result", + stringify!($op) + ); + }}; + } + + check_variant!(ReduxSyncFminOp); + check_variant!(ReduxSyncFminNanOp); + check_variant!(ReduxSyncFminAbsOp); + check_variant!(ReduxSyncFminAbsNanOp); + check_variant!(ReduxSyncFmaxOp); + check_variant!(ReduxSyncFmaxNanOp); + check_variant!(ReduxSyncFmaxAbsOp); + check_variant!(ReduxSyncFmaxAbsNanOp); +} + +#[test] +fn test_bar_warp_sync_construct_and_verify() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); + let block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into(), i64_ty.into()]); + let mask = block.deref(&ctx).get_argument(0); + let wrong_mask = block.deref(&ctx).get_argument(1); + + let valid = BarWarpSyncOp::build(&mut ctx, mask); + assert!(verify_op(&BarWarpSyncOp::new(valid), &ctx).is_ok()); + + let wrong_type = BarWarpSyncOp::build(&mut ctx, wrong_mask); + assert!(verify_op(&BarWarpSyncOp::new(wrong_type), &ctx).is_err()); + + let wrong_arity = Operation::new( + &mut ctx, + BarWarpSyncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + assert!(verify_op(&BarWarpSyncOp::new(wrong_arity), &ctx).is_err()); +} + +#[test] +fn test_elect_sync_construct_and_verify() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); + + // A block supplies the single `mask` operand. + let block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); + let mask = block.deref(&ctx).get_argument(0); + + // Valid: 1 operand [mask], 2 results [leader (i32), is_elected (i1)] + // (matches NOpdsInterface<1>/NResultsInterface<2>). + let op = Operation::new( + &mut ctx, + ElectSyncOp::get_concrete_op_info(), + vec![i32_ty.into(), i1_ty.into()], + vec![mask], + vec![], + 0, + ); + assert!(verify_op(&ElectSyncOp::new(op), &ctx).is_ok()); + + // Invalid: wrong operand count (0 instead of 1) must fail verification. + let bad_opnds = Operation::new( + &mut ctx, + ElectSyncOp::get_concrete_op_info(), + vec![i32_ty.into(), i1_ty.into()], + vec![], + vec![], + 0, + ); + assert!(verify_op(&ElectSyncOp::new(bad_opnds), &ctx).is_err()); + + // Invalid: wrong result count (1 instead of 2) must fail verification. + let bad_results = Operation::new( + &mut ctx, + ElectSyncOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![mask], + vec![], + 0, + ); + assert!(verify_op(&ElectSyncOp::new(bad_results), &ctx).is_err()); + + // Invalid: a correctly-sized result list still cannot use a pointer as + // either result and thereby manufacture a MIR pointer kind. + let pointer_ty = + MirPtrType::get_generic_with_kind(&mut ctx, i32_ty.into(), true, MirPointerKind::UniqueRef); + let bad_result_type = Operation::new( + &mut ctx, + ElectSyncOp::get_concrete_op_info(), + vec![i32_ty.into(), pointer_ty.into()], + vec![mask], + vec![], + 0, + ); + assert!(verify_op(&ElectSyncOp::new(bad_result_type), &ctx).is_err()); +} + +#[test] +fn test_clc_query_rejects_pointer_results() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + + let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let block = BasicBlock::new(&mut ctx, None, vec![u64_ty.into(), u64_ty.into()]); + let operands = vec![ + block.deref(&ctx).get_argument(0), + block.deref(&ctx).get_argument(1), + ]; + + let valid = Operation::new( + &mut ctx, + ClcQueryIsCanceledOp::get_concrete_op_info(), + vec![u32_ty.into()], + operands.clone(), + vec![], + 0, + ); + assert!(verify_op(&ClcQueryIsCanceledOp::new(valid), &ctx).is_ok()); + + let pointer_ty = + MirPtrType::get_generic_with_kind(&mut ctx, u32_ty.into(), true, MirPointerKind::UniqueRef); + let bad = Operation::new( + &mut ctx, + ClcQueryIsCanceledOp::get_concrete_op_info(), + vec![pointer_ty.into()], + operands, + vec![], + 0, + ); + assert!(verify_op(&ClcQueryIsCanceledOp::new(bad), &ctx).is_err()); +} + +#[test] +fn test_shfl_sync_i64_construct_and_verify() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); + + let block = BasicBlock::new( + &mut ctx, + None, + vec![i32_ty.into(), i64_ty.into(), i32_ty.into()], + ); + let mask = block.deref(&ctx).get_argument(0); + let value = block.deref(&ctx).get_argument(1); + let lane = block.deref(&ctx).get_argument(2); + + macro_rules! check_mode { + ($op:ty) => {{ + let valid = <$op>::build(&mut ctx, mask, value, lane); + assert!(verify_op(&<$op>::new(valid), &ctx).is_ok()); + + for (operands, result_ty) in [ + (vec![mask, value], i64_ty.into()), + (vec![value, value, lane], i64_ty.into()), + (vec![mask, mask, lane], i64_ty.into()), + (vec![mask, value, value], i64_ty.into()), + (vec![mask, value, lane], i32_ty.into()), + ] { + let invalid = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + vec![result_ty], + operands, + vec![], + 0, + ); + assert!(verify_op(&<$op>::new(invalid), &ctx).is_err()); + } + + let no_result = Operation::new( + &mut ctx, + <$op>::get_concrete_op_info(), + vec![], + vec![mask, value, lane], + vec![], + 0, + ); + assert!(verify_op(&<$op>::new(no_result), &ctx).is_err()); + }}; + } + + check_mode!(ShflSyncIdxI64Op); + check_mode!(ShflSyncBflyI64Op); + check_mode!(ShflSyncDownI64Op); + check_mode!(ShflSyncUpI64Op); +} From 3e2b8dbc4f4ee3043ea92d9437b3cb41158f16eb Mon Sep 17 00:00:00 2001 From: nihalpasham Date: Fri, 28 Aug 2026 22:08:19 +0530 Subject: [PATCH 03/18] refactor(llvm-export): split tests/export_test.rs into export_test/ modules --- crates/llvm-export/tests/export_test.rs | 5373 ----------------- .../tests/export_test/calls_and_addresses.rs | 377 ++ .../llvm-export/tests/export_test/common.rs | 132 + .../tests/export_test/control_flow.rs | 234 + .../tests/export_test/debug_info.rs | 2458 ++++++++ .../tests/export_test/device_externs.rs | 800 +++ .../llvm-export/tests/export_test/globals.rs | 478 ++ .../tests/export_test/instructions.rs | 578 ++ .../tests/export_test/kernels_and_roots.rs | 283 + crates/llvm-export/tests/export_test/main.rs | 14 + .../llvm-export/tests/export_test/metadata.rs | 231 + 11 files changed, 5585 insertions(+), 5373 deletions(-) delete mode 100644 crates/llvm-export/tests/export_test.rs create mode 100644 crates/llvm-export/tests/export_test/calls_and_addresses.rs create mode 100644 crates/llvm-export/tests/export_test/common.rs create mode 100644 crates/llvm-export/tests/export_test/control_flow.rs create mode 100644 crates/llvm-export/tests/export_test/debug_info.rs create mode 100644 crates/llvm-export/tests/export_test/device_externs.rs create mode 100644 crates/llvm-export/tests/export_test/globals.rs create mode 100644 crates/llvm-export/tests/export_test/instructions.rs create mode 100644 crates/llvm-export/tests/export_test/kernels_and_roots.rs create mode 100644 crates/llvm-export/tests/export_test/main.rs create mode 100644 crates/llvm-export/tests/export_test/metadata.rs diff --git a/crates/llvm-export/tests/export_test.rs b/crates/llvm-export/tests/export_test.rs deleted file mode 100644 index 2ae0d00273..0000000000 --- a/crates/llvm-export/tests/export_test.rs +++ /dev/null @@ -1,5373 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -use combine::stream::position::SourcePosition; -use llvm_export::{ - export::{ - DebugKind, DeviceExternAttrs, DeviceExternDecl, DeviceExternType, ExportBackendConfig, - NvvmExportConfig, NvvmIrDialect, PtxExportConfig, export_module_to_string, - export_module_to_string_with_config, export_module_with_externs, - export_module_with_externs_and_roots, - }, - op_interfaces::CastOpInterface, - ops::{ - AddrSpaceCastOp, AddressOfOp, AllocaOp, BitcastOp, BrOp, CallOp, CondBrOp, ConstantOp, - DebugEnumDiscriminant, DebugEnumVariant, DebugFragment, DebugFragmentVariableInfo, - DebugGlobalVariableInfo, DebugLocalTypeKind, DebugLocalVariableInfo, - DebugProjectedVariableInfo, DebugSourcePosition, DebugSourceScope, - DebugSourceScopeLocation, DebugSourceScopeMap, DebugValueExpression, - DebugValueExpressionOp, DebugValueListOp, DebugValueOp, FuncOp, GepIndex, GetElementPtrOp, - GlobalInitializerRelocation, GlobalOp, GlobalOpExt, InlineAsmOp, LoadOp, ReturnOp, - SelectOp, StoreOp, UndefOp, encode_global_initializer_relocations, - }, - types::{ArrayType, FuncType, HalfType, PointerType, StructLayout, StructType, VoidType}, -}; -use pliron::{ - basic_block::BasicBlock, - builtin::{ - attributes::{IntegerAttr, StringAttr}, - op_interfaces::CallOpCallable, - ops::ModuleOp, - types::{FP32Type, IntegerType, Signedness}, - }, - common_traits::Verify, - context::{Context, Ptr}, - identifier::Identifier, - linked_list::ContainsLinkedList, - location::{Located, Location, Source}, - op::Op, - utils::apint::APInt, -}; -use std::{num::NonZero, path::PathBuf}; - -struct DebugConfig { - inner: C, - debug_kind: DebugKind, -} - -impl ExportBackendConfig for DebugConfig { - fn datalayout(&self) -> &str { - self.inner.datalayout() - } - - fn emit_llvm_used(&self) -> bool { - self.inner.emit_llvm_used() - } - - fn emit_nvvmir_version(&self) -> bool { - self.inner.emit_nvvmir_version() - } - - fn nvvmir_version(&self) -> [i32; 4] { - self.inner.nvvmir_version() - } - - fn emit_all_kernel_annotations(&self) -> bool { - self.inner.emit_all_kernel_annotations() - } - - fn emit_ptx_kernel_keyword(&self) -> bool { - self.inner.emit_ptx_kernel_keyword() - } - - fn nvvm_ir_dialect(&self) -> Option { - self.inner.nvvm_ir_dialect() - } - - fn debug_kind(&self) -> DebugKind { - self.debug_kind - } -} - -fn src_location(ctx: &mut Context, file: &str, line: i32, column: i32) -> Location { - Location::SrcPos { - src: Source::new_from_file(ctx, PathBuf::from(file)), - pos: SourcePosition { line, column }, - } -} - -fn module_top_block(ctx: &mut Context, module: &ModuleOp) -> Ptr { - let module_region = module.get_operation().deref(ctx).get_region(0); - { - let region = module_region.deref(ctx); - if let Some(block) = region.iter(ctx).next() { - return block; - } - } - - let block = BasicBlock::new(ctx, None, vec![]); - block.insert_at_back(module_region, ctx); - block -} - -fn metadata_id<'a>(ir: &'a str, needle: &str) -> &'a str { - ir.lines() - .find(|line| line.contains(needle)) - .and_then(|line| line.split_once(" = ").map(|(id, _)| id)) - .unwrap_or_else(|| panic!("missing metadata node containing {needle:?}:\n{ir}")) -} - -#[test] -fn export_volatile_load_prints_keyword() { - let mut ctx = Context::new(); - - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let ptr_ty = PointerType::get(&ctx, 0); - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![ptr_ty.to_handle()], false); - let func = FuncOp::new(&mut ctx, "volatile_load_test".try_into().unwrap(), func_ty); - let entry = func.get_or_create_entry_block(&mut ctx); - let ptr = entry.deref(&ctx).get_argument(0); - - let load = LoadOp::new(&mut ctx, ptr, i32_ty.to_handle()); - llvm_export::ops::set_op_volatile(&mut ctx, load.get_operation(), true); - load.get_operation().insert_at_back(entry, &ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - let ir = export_module_to_string(&ctx, &module).expect("export succeeds"); - let line = ir - .lines() - .find(|line| line.contains("load volatile")) - .expect("volatile load line"); - - assert!( - line.trim_start().contains(" = load volatile i32, ptr "), - "volatile load keyword must appear immediately after load:\n{ir}" - ); -} - -#[test] -fn export_volatile_store_prints_keyword() { - let mut ctx = Context::new(); - - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let ptr_ty = PointerType::get(&ctx, 0); - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get( - &ctx, - void_ty.to_handle(), - vec![ptr_ty.to_handle(), i32_ty.to_handle()], - false, - ); - let func = FuncOp::new(&mut ctx, "volatile_store_test".try_into().unwrap(), func_ty); - let entry = func.get_or_create_entry_block(&mut ctx); - let ptr = entry.deref(&ctx).get_argument(0); - let val = entry.deref(&ctx).get_argument(1); - - let store = StoreOp::new(&mut ctx, val, ptr); - llvm_export::ops::set_op_volatile(&mut ctx, store.get_operation(), true); - store.get_operation().insert_at_back(entry, &ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - let ir = export_module_to_string(&ctx, &module).expect("export succeeds"); - let line = ir - .lines() - .find(|line| line.contains("store volatile")) - .expect("volatile store line"); - - assert!( - line.trim_start().starts_with("store volatile i32 "), - "volatile store keyword must appear immediately after store:\n{ir}" - ); -} - -#[test] -fn legacy_export_uses_one_canonical_pointer_with_multiple_typed_views() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "legacy_views".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - - let ptr_ty = PointerType::get(&ctx, 0); - let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signless); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let f32_ty = FP32Type::get(&ctx); - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.into(), vec![ptr_ty.into()], false); - let func = FuncOp::new(&mut ctx, "multiple_views".try_into().unwrap(), func_ty); - let entry = func.get_or_create_entry_block(&mut ctx); - let pointer = entry.deref(&ctx).get_argument(0); - - LoadOp::new(&mut ctx, pointer, i32_ty.into()) - .get_operation() - .insert_at_back(entry, &ctx); - LoadOp::new(&mut ctx, pointer, f32_ty.into()) - .get_operation() - .insert_at_back(entry, &ctx); - LoadOp::new(&mut ctx, pointer, i8_ty.into()) - .get_operation() - .insert_at_back(entry, &ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - let config = NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7); - let ir = export_module_to_string_with_config(&ctx, &module, &config) - .expect("legacy export succeeds"); - - assert!(ir.contains("define void @multiple_views(i8* %v0)"), "{ir}"); - assert!(ir.contains("bitcast i8* %v0 to i32*"), "{ir}"); - assert!(ir.contains("load i32, i32*"), "{ir}"); - assert!(ir.contains("bitcast i8* %v0 to float*"), "{ir}"); - assert!(ir.contains("load float, float*"), "{ir}"); - assert!(ir.contains("load i8, i8* %v0"), "{ir}"); - assert!(!ir.contains("bitcast i8* %v0 to i8*"), "{ir}"); - assert!( - !ir.split(|c: char| !c.is_ascii_alphanumeric()) - .any(|t| t == "ptr") - ); -} - -#[test] -fn legacy_alloca_rejects_a_non_default_result_address_space() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "invalid_alloca_as".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.into(), vec![], false); - let func = FuncOp::new(&mut ctx, "invalid_alloca".try_into().unwrap(), func_ty); - let entry = func.get_or_create_entry_block(&mut ctx); - - let one_attr = IntegerAttr::new(i32_ty, APInt::from_u32(1, NonZero::new(32).unwrap())); - let one = ConstantOp::new(&mut ctx, one_attr.into()); - let one_value = one.get_operation().deref(&ctx).get_result(0); - one.get_operation().insert_at_back(entry, &ctx); - let alloca = AllocaOp::new(&mut ctx, i32_ty.into(), one_value); - let alloca_result = alloca.get_operation().deref(&ctx).get_result(0); - let shared_pointer = PointerType::get(&ctx, 3); - alloca_result.set_type(&ctx, shared_pointer.into()); - alloca.get_operation().insert_at_back(entry, &ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - module - .get_operation() - .deref(&ctx) - .verify(&ctx) - .expect("upstream verification currently does not enforce alloca result AS0"); - let error = export_module_to_string_with_config( - &ctx, - &module, - &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), - ) - .expect_err("legacy export must reject an alloca address-space mismatch"); - assert!( - error.contains("alloca result uses address space 3"), - "{error}" - ); -} - -#[test] -fn legacy_gep_rejects_a_result_address_space_different_from_its_base() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "invalid_gep_as".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - - let global_pointer = PointerType::get(&ctx, 1); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.into(), vec![global_pointer.into()], false); - let func = FuncOp::new(&mut ctx, "invalid_gep".try_into().unwrap(), func_ty); - let entry = func.get_or_create_entry_block(&mut ctx); - let base = entry.deref(&ctx).get_argument(0); - let gep = GetElementPtrOp::new(&mut ctx, base, vec![GepIndex::Constant(0)], i32_ty.into()); - let gep_result = gep.get_operation().deref(&ctx).get_result(0); - let shared_pointer = PointerType::get(&ctx, 3); - gep_result.set_type(&ctx, shared_pointer.into()); - gep.get_operation().insert_at_back(entry, &ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - module - .get_operation() - .deref(&ctx) - .verify(&ctx) - .expect("upstream verification currently does not enforce GEP result/base AS equality"); - let error = export_module_to_string_with_config( - &ctx, - &module, - &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), - ) - .expect_err("legacy export must reject a GEP address-space mismatch"); - assert!( - error.contains("GEP result address-space mismatch: base is 1, result is 3"), - "{error}" - ); -} - -#[test] -fn gep_inbounds_marker_controls_exported_pointer_semantics() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "gep_semantics".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - - let pointer = PointerType::get(&ctx, 0); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.into(), vec![pointer.into()], false); - let func = FuncOp::new(&mut ctx, "offsets".try_into().unwrap(), func_ty); - let entry = func.get_or_create_entry_block(&mut ctx); - let base = entry.deref(&ctx).get_argument(0); - - let ordinary = GetElementPtrOp::new(&mut ctx, base, vec![GepIndex::Constant(1)], i32_ty.into()); - ordinary.get_operation().insert_at_back(entry, &ctx); - - let wrapping = GetElementPtrOp::new(&mut ctx, base, vec![GepIndex::Constant(2)], i32_ty.into()); - llvm_export::ops::set_gep_inbounds(&mut ctx, wrapping.get_operation(), false); - wrapping.get_operation().insert_at_back(entry, &ctx); - - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - let ir = export_module_to_string(&ctx, &module).expect("GEP export succeeds"); - let gep_lines: Vec<_> = ir - .lines() - .filter(|line| line.contains("getelementptr")) - .collect(); - assert_eq!(gep_lines.len(), 2, "{ir}"); - assert!(gep_lines[0].contains("getelementptr inbounds"), "{ir}"); - assert!( - gep_lines[1].contains("getelementptr i32") - && !gep_lines[1].contains("getelementptr inbounds"), - "{ir}" - ); -} - -#[test] -fn legacy_pointer_select_keeps_one_canonical_type() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "legacy_pointer_select".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - - let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); - let ptr_ty = PointerType::get(&ctx, 0); - let func_ty = FuncType::get( - &ctx, - ptr_ty.into(), - vec![i1_ty.into(), ptr_ty.into(), ptr_ty.into()], - false, - ); - let func = FuncOp::new(&mut ctx, "choose_pointer".try_into().unwrap(), func_ty); - let entry = func.get_or_create_entry_block(&mut ctx); - let condition = entry.deref(&ctx).get_argument(0); - let if_true = entry.deref(&ctx).get_argument(1); - let if_false = entry.deref(&ctx).get_argument(2); - let select = SelectOp::new(&mut ctx, condition, if_true, if_false); - let selected = select.get_operation().deref(&ctx).get_result(0); - select.get_operation().insert_at_back(entry, &ctx); - ReturnOp::new(&mut ctx, Some(selected)) - .get_operation() - .insert_at_back(entry, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - let ir = export_module_to_string_with_config( - &ctx, - &module, - &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), - ) - .expect("legacy pointer select export succeeds"); - assert!( - ir.contains("select i1 %v0, i8* %v1, i8* %v2"), - "pointer select must use the canonical byte-pointer type:\n{ir}" - ); - assert!(ir.contains("ret i8*"), "{ir}"); -} - -#[test] -fn exporter_rejects_extra_predecessor_values_before_emitting_phis() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "invalid_branch_arity".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.into(), vec![], false); - let func = FuncOp::new(&mut ctx, "invalid_branch".try_into().unwrap(), func_ty); - let entry = func.get_or_create_entry_block(&mut ctx); - let region = func.get_operation().deref(&ctx).get_region(0); - let destination = BasicBlock::new(&mut ctx, None, vec![]); - destination.insert_at_back(region, &ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let one_attr = IntegerAttr::new(i32_ty, APInt::from_u32(1, NonZero::new(32).unwrap())); - let one = ConstantOp::new(&mut ctx, one_attr.into()); - let one_value = one.get_operation().deref(&ctx).get_result(0); - one.get_operation().insert_at_back(entry, &ctx); - BrOp::new(&mut ctx, destination, vec![one_value]) - .get_operation() - .insert_at_back(entry, &ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(destination, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - let error = export_module_to_string_with_config( - &ctx, - &module, - &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), - ) - .expect_err("extra predecessor values must be rejected"); - assert!( - error.contains("supplies 1 values") && error.contains("expects 0 block arguments"), - "{error}" - ); -} - -#[test] -fn exporter_rejects_distinct_values_on_duplicate_conditional_edges() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "duplicate_conditional_edge".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let func_ty = FuncType::get( - &ctx, - i32_ty.into(), - vec![i1_ty.into(), i32_ty.into(), i32_ty.into()], - false, - ); - let func = FuncOp::new(&mut ctx, "duplicate_edge".try_into().unwrap(), func_ty); - let entry = func.get_or_create_entry_block(&mut ctx); - let condition = entry.deref(&ctx).get_argument(0); - let if_true = entry.deref(&ctx).get_argument(1); - let if_false = entry.deref(&ctx).get_argument(2); - let region = func.get_operation().deref(&ctx).get_region(0); - let destination = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); - destination.insert_at_back(region, &ctx); - CondBrOp::new( - &mut ctx, - condition, - destination, - vec![if_true], - destination, - vec![if_false], - ) - .get_operation() - .insert_at_back(entry, &ctx); - let result = destination.deref(&ctx).get_argument(0); - ReturnOp::new(&mut ctx, Some(result)) - .get_operation() - .insert_at_back(destination, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - module - .get_operation() - .deref(&ctx) - .verify(&ctx) - .expect("pliron permits same-destination conditional edges with distinct values"); - let error = export_module_to_string_with_config( - &ctx, - &module, - &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), - ) - .expect_err("LLVM PHIs cannot distinguish duplicate predecessor edges"); - assert!( - error.contains("both edges with different forwarded values"), - "{error}" - ); -} - -#[test] -fn exporter_deduplicates_identical_values_on_duplicate_conditional_edges() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "identical_conditional_edge".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let func_ty = FuncType::get( - &ctx, - i32_ty.into(), - vec![i1_ty.into(), i32_ty.into()], - false, - ); - let func = FuncOp::new(&mut ctx, "identical_edge".try_into().unwrap(), func_ty); - let entry = func.get_or_create_entry_block(&mut ctx); - let condition = entry.deref(&ctx).get_argument(0); - let value = entry.deref(&ctx).get_argument(1); - let region = func.get_operation().deref(&ctx).get_region(0); - let destination = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); - destination.insert_at_back(region, &ctx); - CondBrOp::new( - &mut ctx, - condition, - destination, - vec![value], - destination, - vec![value], - ) - .get_operation() - .insert_at_back(entry, &ctx); - let result = destination.deref(&ctx).get_argument(0); - ReturnOp::new(&mut ctx, Some(result)) - .get_operation() - .insert_at_back(destination, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - let ir = export_module_to_string_with_config( - &ctx, - &module, - &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), - ) - .expect("identical duplicate-edge values can use one PHI predecessor"); - let phi = ir - .lines() - .find(|line| line.contains(" = phi i32 ")) - .expect("destination block must contain a PHI"); - assert_eq!(phi.matches("%entry").count(), 1, "{phi}"); -} - -#[test] -fn phi_can_reference_undef_from_a_later_block() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "later_undef_phi".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let func_ty = FuncType::get( - &ctx, - i32_ty.into(), - vec![i1_ty.into(), i32_ty.into()], - false, - ); - let func = FuncOp::new(&mut ctx, "choose_undef".try_into().unwrap(), func_ty); - let entry = func.get_or_create_entry_block(&mut ctx); - let condition = entry.deref(&ctx).get_argument(0); - let fallback = entry.deref(&ctx).get_argument(1); - let region = func.get_operation().deref(&ctx).get_region(0); - - // The join precedes both predecessors in print order, so its PHI depends - // on the exporter's whole-function value-name pre-pass. - let join = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); - join.insert_at_back(region, &ctx); - let undef_block = BasicBlock::new(&mut ctx, None, vec![]); - undef_block.insert_at_back(region, &ctx); - let value_block = BasicBlock::new(&mut ctx, None, vec![]); - value_block.insert_at_back(region, &ctx); - - CondBrOp::new( - &mut ctx, - condition, - undef_block, - vec![], - value_block, - vec![], - ) - .get_operation() - .insert_at_back(entry, &ctx); - - let undef = UndefOp::new(&mut ctx, i32_ty.into()); - let undef_value = undef.get_operation().deref(&ctx).get_result(0); - undef.get_operation().insert_at_back(undef_block, &ctx); - BrOp::new(&mut ctx, join, vec![undef_value]) - .get_operation() - .insert_at_back(undef_block, &ctx); - BrOp::new(&mut ctx, join, vec![fallback]) - .get_operation() - .insert_at_back(value_block, &ctx); - - let result = join.deref(&ctx).get_argument(0); - ReturnOp::new(&mut ctx, Some(result)) - .get_operation() - .insert_at_back(join, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - let ir = export_module_to_string_with_config( - &ctx, - &module, - &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), - ) - .expect("later-block undef must be available while exporting an earlier PHI"); - assert!( - ir.lines() - .any(|line| line.contains(" = phi i32 ") && line.contains("[ undef,")), - "{ir}" - ); -} - -#[test] -fn indirect_call_rejects_non_program_address_space() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "invalid_indirect_callee".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - - let shared_ptr_ty = PointerType::get(&ctx, 3); - let void_ty = VoidType::get(&ctx); - let callee_ty = FuncType::get(&ctx, void_ty.into(), vec![], false); - let caller_ty = FuncType::get(&ctx, void_ty.into(), vec![shared_ptr_ty.into()], false); - let caller = FuncOp::new(&mut ctx, "caller".try_into().unwrap(), caller_ty); - let entry = caller.get_or_create_entry_block(&mut ctx); - let callee = entry.deref(&ctx).get_argument(0); - CallOp::new( - &mut ctx, - CallOpCallable::Indirect(callee), - callee_ty, - vec![], - ) - .get_operation() - .insert_at_back(entry, &ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - caller.get_operation().insert_at_back(module_block, &ctx); - - for dialect in [NvvmIrDialect::LegacyLlvm7, NvvmIrDialect::Modern] { - let error = - export_module_to_string_with_config(&ctx, &module, &NvvmExportConfig::new(dialect)) - .expect_err("NVPTX must reject a shared-memory function pointer"); - assert!(error.contains("address space 3"), "{error}"); - assert!(error.contains("function pointers"), "{error}"); - } -} - -#[test] -fn intrinsic_export_preserves_legacy_dots_and_literal_underscores() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "intrinsic_names".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let void_ty = VoidType::get(&ctx); - let function_ty = FuncType::get(&ctx, void_ty.into(), vec![], false); - let legacy = "llvm_nvvm_wgmma_fence_sync_aligned"; - let escaped = "llvm__nvvm_dwgmma_dcommit_ugroup_dsync_daligned"; - - for name in [legacy, escaped] { - FuncOp::new(&mut ctx, name.try_into().unwrap(), function_ty) - .get_operation() - .insert_at_back(module_block, &ctx); - } - - let caller = FuncOp::new(&mut ctx, "caller".try_into().unwrap(), function_ty); - let entry = caller.get_or_create_entry_block(&mut ctx); - for name in [legacy, escaped] { - CallOp::new( - &mut ctx, - CallOpCallable::Direct(name.try_into().unwrap()), - function_ty, - vec![], - ) - .get_operation() - .insert_at_back(entry, &ctx); - } - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - caller.get_operation().insert_at_back(module_block, &ctx); - - let ir = export_module_to_string(&ctx, &module).expect("intrinsic export succeeds"); - assert!(ir.contains("@llvm.nvvm.wgmma.fence.sync.aligned"), "{ir}"); - assert!( - ir.contains("@llvm.nvvm.wgmma.commit_group.sync.aligned"), - "{ir}" - ); - assert!(!ir.contains("@llvm.nvvm.wgmma.commit.group"), "{ir}"); -} - -#[test] -fn pointer_bitcast_cannot_cross_address_spaces_in_either_nvvm_dialect() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "invalid_pointer_bitcast".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let global_pointer = PointerType::get(&ctx, 1); - let shared_pointer = PointerType::get(&ctx, 3); - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.into(), vec![global_pointer.into()], false); - let func = FuncOp::new(&mut ctx, "invalid_cast".try_into().unwrap(), func_ty); - let entry = func.get_or_create_entry_block(&mut ctx); - let input = entry.deref(&ctx).get_argument(0); - BitcastOp::new(&mut ctx, input, shared_pointer.into()) - .get_operation() - .insert_at_back(entry, &ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - module - .get_operation() - .deref(&ctx) - .verify(&ctx) - .expect("upstream bitcast verification currently does not enforce pointer AS equality"); - for dialect in [NvvmIrDialect::LegacyLlvm7, NvvmIrDialect::Modern] { - let error = - export_module_to_string_with_config(&ctx, &module, &NvvmExportConfig::new(dialect)) - .expect_err("a cross-address-space pointer bitcast must be rejected"); - assert!( - error.contains("pointer bitcast cannot cross address spaces 1 -> 3"), - "{dialect:?}: {error}" - ); - } -} - -#[test] -fn addrspacecast_must_change_address_spaces_in_either_nvvm_dialect() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "invalid_addrspacecast".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let shared_pointer = PointerType::get(&ctx, 3); - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.into(), vec![shared_pointer.into()], false); - let func = FuncOp::new(&mut ctx, "invalid_cast".try_into().unwrap(), func_ty); - let entry = func.get_or_create_entry_block(&mut ctx); - let input = entry.deref(&ctx).get_argument(0); - AddrSpaceCastOp::new(&mut ctx, input, shared_pointer.into()) - .get_operation() - .insert_at_back(entry, &ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - module - .get_operation() - .deref(&ctx) - .verify(&ctx) - .expect("upstream addrspacecast verification currently permits equal address spaces"); - for dialect in [NvvmIrDialect::LegacyLlvm7, NvvmIrDialect::Modern] { - let error = - export_module_to_string_with_config(&ctx, &module, &NvvmExportConfig::new(dialect)) - .expect_err("addrspacecast must not encode a no-op address-space conversion"); - assert!( - error.contains( - "addrspacecast must change address spaces; source and result are both address space 3" - ), - "{dialect:?}: {error}" - ); - } -} - -#[test] -fn legacy_function_address_defined_later_round_trips_through_indirect_call() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "function_address".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let void_ty = VoidType::get(&ctx); - let callee_ty = FuncType::get(&ctx, void_ty.into(), vec![], false); - - // Print the caller first to prove symbol typing is a module pre-pass, not - // an accidental dependency on textual definition order. - let caller = FuncOp::new( - &mut ctx, - "call_function_pointer".try_into().unwrap(), - callee_ty, - ); - let caller_entry = caller.get_or_create_entry_block(&mut ctx); - let address = AddressOfOp::new(&mut ctx, "target".try_into().unwrap(), 0); - let address_value = address.get_operation().deref(&ctx).get_result(0); - address.get_operation().insert_at_back(caller_entry, &ctx); - CallOp::new( - &mut ctx, - CallOpCallable::Indirect(address_value), - callee_ty, - vec![], - ) - .get_operation() - .insert_at_back(caller_entry, &ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(caller_entry, &ctx); - caller.get_operation().insert_at_back(module_block, &ctx); - - let target = FuncOp::new(&mut ctx, "target".try_into().unwrap(), callee_ty); - let target_entry = target.get_or_create_entry_block(&mut ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(target_entry, &ctx); - target.get_operation().insert_at_back(module_block, &ctx); - - let ir = export_module_to_string_with_config( - &ctx, - &module, - &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), - ) - .expect("legacy function-address export succeeds"); - assert!( - ir.contains("bitcast void ()* @target to i8*"), - "function address must normalize to the canonical byte pointer:\n{ir}" - ); - assert!( - ir.contains("bitcast i8*") && ir.contains("to void ()*"), - "indirect call must restore the exact function pointer type:\n{ir}" - ); - assert!(ir.contains("call void %"), "{ir}"); -} - -#[test] -fn modern_function_address_uses_the_normalized_definition_name() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "modern_function_address".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let void_ty = VoidType::get(&ctx); - let callee_ty = FuncType::get(&ctx, void_ty.into(), vec![], false); - let prefixed_name = reserved_oxide_symbols::device_symbol("target"); - - let caller = FuncOp::new(&mut ctx, "caller".try_into().unwrap(), callee_ty); - let caller_entry = caller.get_or_create_entry_block(&mut ctx); - let address = AddressOfOp::new(&mut ctx, prefixed_name.as_str().try_into().unwrap(), 0); - let address_value = address.get_operation().deref(&ctx).get_result(0); - address.get_operation().insert_at_back(caller_entry, &ctx); - CallOp::new( - &mut ctx, - CallOpCallable::Indirect(address_value), - callee_ty, - vec![], - ) - .get_operation() - .insert_at_back(caller_entry, &ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(caller_entry, &ctx); - caller.get_operation().insert_at_back(module_block, &ctx); - - let target = FuncOp::new( - &mut ctx, - prefixed_name.as_str().try_into().unwrap(), - callee_ty, - ); - let target_entry = target.get_or_create_entry_block(&mut ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(target_entry, &ctx); - target.get_operation().insert_at_back(module_block, &ctx); - - let ir = export_module_to_string_with_config( - &ctx, - &module, - &NvvmExportConfig::new(NvvmIrDialect::Modern), - ) - .expect("modern function-address export succeeds"); - assert!(ir.contains("define void @target()"), "{ir}"); - assert!(ir.contains("call void @target()"), "{ir}"); - assert!(!ir.contains(&prefixed_name), "{ir}"); -} - -#[test] -fn modern_addressof_rejects_global_and_function_address_space_mismatches() { - let mut ctx = Context::new(); - let void_ty = VoidType::get(&ctx); - let no_args = FuncType::get(&ctx, void_ty.into(), vec![], false); - - let global_module = ModuleOp::new(&mut ctx, "bad_global_address".try_into().unwrap()); - let global_module_block = module_top_block(&mut ctx, &global_module); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let global = GlobalOp::new(&mut ctx, "shared_value".try_into().unwrap(), i32_ty.into()); - global.set_address_space(&mut ctx, 3); - global - .get_operation() - .insert_at_back(global_module_block, &ctx); - let global_user = FuncOp::new(&mut ctx, "global_user".try_into().unwrap(), no_args); - let global_entry = global_user.get_or_create_entry_block(&mut ctx); - AddressOfOp::new(&mut ctx, "shared_value".try_into().unwrap(), 0) - .get_operation() - .insert_at_back(global_entry, &ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(global_entry, &ctx); - global_user - .get_operation() - .insert_at_back(global_module_block, &ctx); - let error = export_module_to_string_with_config( - &ctx, - &global_module, - &NvvmExportConfig::new(NvvmIrDialect::Modern), - ) - .expect_err("modern global addressof must preserve address spaces"); - assert!( - error.contains("result is 0, global is 3"), - "unexpected global error: {error}" - ); - - let function_module = ModuleOp::new(&mut ctx, "bad_function_address".try_into().unwrap()); - let function_module_block = module_top_block(&mut ctx, &function_module); - FuncOp::new(&mut ctx, "target".try_into().unwrap(), no_args) - .get_operation() - .insert_at_back(function_module_block, &ctx); - let function_user = FuncOp::new(&mut ctx, "function_user".try_into().unwrap(), no_args); - let function_entry = function_user.get_or_create_entry_block(&mut ctx); - AddressOfOp::new(&mut ctx, "target".try_into().unwrap(), 3) - .get_operation() - .insert_at_back(function_entry, &ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(function_entry, &ctx); - function_user - .get_operation() - .insert_at_back(function_module_block, &ctx); - let error = export_module_to_string_with_config( - &ctx, - &function_module, - &NvvmExportConfig::new(NvvmIrDialect::Modern), - ) - .expect_err("modern function addressof must use program address space"); - assert!(error.contains("program-address-space (0)"), "{error}"); -} - -#[test] -fn legacy_device_extern_adapts_exact_pointer_arguments_and_results() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "legacy_extern".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - - let ptr_ty = PointerType::get(&ctx, 0); - let external_ty = FuncType::get(&ctx, ptr_ty.into(), vec![ptr_ty.into()], false); - FuncOp::new(&mut ctx, "float_roundtrip".try_into().unwrap(), external_ty) - .get_operation() - .insert_at_back(module_block, &ctx); - - let void_ty = VoidType::get(&ctx); - let caller_ty = FuncType::get(&ctx, void_ty.into(), vec![ptr_ty.into()], false); - let caller = FuncOp::new(&mut ctx, "caller".try_into().unwrap(), caller_ty); - let entry = caller.get_or_create_entry_block(&mut ctx); - let pointer = entry.deref(&ctx).get_argument(0); - CallOp::new( - &mut ctx, - CallOpCallable::Direct("float_roundtrip".try_into().unwrap()), - external_ty, - vec![pointer], - ) - .get_operation() - .insert_at_back(entry, &ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - caller.get_operation().insert_at_back(module_block, &ctx); - - let externs = [DeviceExternDecl { - export_name: "float_roundtrip".to_string(), - param_types: vec![DeviceExternType::pointer_to(DeviceExternType::Float32, 0)], - return_type: DeviceExternType::pointer_to(DeviceExternType::Float32, 0), - attrs: DeviceExternAttrs::default(), - }]; - let ir = export_module_with_externs( - &ctx, - &module, - &externs, - &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), - ) - .expect("legacy extern export succeeds"); - - assert!( - ir.contains("declare float* @float_roundtrip(float*)"), - "{ir}" - ); - assert!(ir.contains("bitcast i8* %v0 to float*"), "{ir}"); - assert!(ir.contains("call float* @float_roundtrip(float*"), "{ir}"); - assert!( - ir.contains(" = bitcast float* ") && ir.contains(" to i8*"), - "{ir}" - ); - assert!( - !ir.split(|c: char| !c.is_ascii_alphanumeric()) - .any(|token| token == "ptr"), - "{ir}" - ); -} - -#[test] -fn legacy_device_extern_preserves_pointer_address_spaces() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "legacy_extern_as".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let ptr_ty = PointerType::get(&ctx, 3); - let void_ty = VoidType::get(&ctx); - let external_ty = FuncType::get(&ctx, void_ty.into(), vec![ptr_ty.into()], false); - FuncOp::new(&mut ctx, "shared_float".try_into().unwrap(), external_ty) - .get_operation() - .insert_at_back(module_block, &ctx); - let caller = FuncOp::new(&mut ctx, "caller".try_into().unwrap(), external_ty); - let entry = caller.get_or_create_entry_block(&mut ctx); - let pointer = entry.deref(&ctx).get_argument(0); - CallOp::new( - &mut ctx, - CallOpCallable::Direct("shared_float".try_into().unwrap()), - external_ty, - vec![pointer], - ) - .get_operation() - .insert_at_back(entry, &ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - caller.get_operation().insert_at_back(module_block, &ctx); - - let externs = [DeviceExternDecl { - export_name: "shared_float".to_string(), - param_types: vec![DeviceExternType::pointer_to(DeviceExternType::Float32, 3)], - return_type: DeviceExternType::Void, - attrs: DeviceExternAttrs::default(), - }]; - let ir = export_module_with_externs( - &ctx, - &module, - &externs, - &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), - ) - .expect("legacy extern export succeeds"); - assert!( - ir.contains("declare void @shared_float(float addrspace(3)*)"), - "{ir}" - ); - assert!( - ir.contains("bitcast i8 addrspace(3)* %v0 to float addrspace(3)*"), - "{ir}" - ); -} - -#[test] -fn modern_device_extern_erases_pointee_without_boundary_casts() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "modern_extern".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let ptr_ty = PointerType::get(&ctx, 0); - let void_ty = VoidType::get(&ctx); - let external_ty = FuncType::get(&ctx, void_ty.into(), vec![ptr_ty.into()], false); - FuncOp::new(&mut ctx, "takes_float".try_into().unwrap(), external_ty) - .get_operation() - .insert_at_back(module_block, &ctx); - let caller = FuncOp::new(&mut ctx, "caller".try_into().unwrap(), external_ty); - let entry = caller.get_or_create_entry_block(&mut ctx); - let pointer = entry.deref(&ctx).get_argument(0); - CallOp::new( - &mut ctx, - CallOpCallable::Direct("takes_float".try_into().unwrap()), - external_ty, - vec![pointer], - ) - .get_operation() - .insert_at_back(entry, &ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - caller.get_operation().insert_at_back(module_block, &ctx); - - let externs = [DeviceExternDecl { - export_name: "takes_float".to_string(), - param_types: vec![DeviceExternType::pointer_to(DeviceExternType::Float32, 0)], - return_type: DeviceExternType::Void, - attrs: DeviceExternAttrs::default(), - }]; - let ir = export_module_with_externs( - &ctx, - &module, - &externs, - &NvvmExportConfig::new(NvvmIrDialect::Modern), - ) - .expect("modern extern export succeeds"); - assert!(ir.contains("declare void @takes_float(ptr)"), "{ir}"); - assert!(ir.contains("call void @takes_float(ptr %v0)"), "{ir}"); - assert!(!ir.contains("bitcast"), "{ir}"); -} - -#[test] -fn device_extern_rejects_invalid_symbol_and_address_space_mismatch() { - let mut ctx = Context::new(); - let empty = ModuleOp::new(&mut ctx, "empty".try_into().unwrap()); - let invalid = [DeviceExternDecl { - export_name: "bad.name".to_string(), - param_types: vec![], - return_type: DeviceExternType::Void, - attrs: DeviceExternAttrs::default(), - }]; - let err = export_module_with_externs( - &ctx, - &empty, - &invalid, - &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), - ) - .expect_err("invalid NVVM symbol must fail"); - assert!(err.contains("global-identifier subset"), "{err}"); - - let reserved_intrinsic_prefix = [DeviceExternDecl { - export_name: "llvm_external".to_string(), - param_types: vec![], - return_type: DeviceExternType::Void, - attrs: DeviceExternAttrs::default(), - }]; - let err = export_module_with_externs( - &ctx, - &empty, - &reserved_intrinsic_prefix, - &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), - ) - .expect_err("the reserved intrinsic namespace must not be ambiguous"); - assert!(err.contains("reserves for LLVM intrinsics"), "{err}"); - - let by_value_array = [DeviceExternDecl { - export_name: "array_by_value".to_string(), - param_types: vec![DeviceExternType::Array { - element: Box::new(DeviceExternType::Float32), - len: 4, - }], - return_type: DeviceExternType::Void, - attrs: DeviceExternAttrs::default(), - }]; - let err = export_module_with_externs( - &ctx, - &empty, - &by_value_array, - &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), - ) - .expect_err("by-value aggregate externs must be rejected"); - assert!(err.contains("passes an array by value"), "{err}"); - - let nested_half = [DeviceExternDecl { - export_name: "half_buffer".to_string(), - param_types: vec![DeviceExternType::pointer_to( - DeviceExternType::Array { - element: Box::new(DeviceExternType::Float16), - len: 4, - }, - 0, - )], - return_type: DeviceExternType::Void, - attrs: DeviceExternAttrs::default(), - }]; - let err = export_module_with_externs( - &ctx, - &empty, - &nested_half, - &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), - ) - .expect_err("legacy half nested in a pointer must fail"); - assert!( - err.contains("CUDA 12 legacy") && err.contains("half"), - "{err}" - ); - let modern = export_module_with_externs( - &ctx, - &empty, - &nested_half, - &NvvmExportConfig::new(NvvmIrDialect::Modern), - ) - .expect("modern opaque-pointer extern may use half pointees"); - assert!( - modern.contains("declare void @half_buffer(ptr)"), - "{modern}" - ); - - let module = ModuleOp::new(&mut ctx, "mismatch".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let ptr0 = PointerType::get(&ctx, 0); - let void_ty = VoidType::get(&ctx); - let external_ty = FuncType::get(&ctx, void_ty.into(), vec![ptr0.into()], false); - FuncOp::new(&mut ctx, "shared_only".try_into().unwrap(), external_ty) - .get_operation() - .insert_at_back(module_block, &ctx); - let caller = FuncOp::new(&mut ctx, "caller".try_into().unwrap(), external_ty); - let entry = caller.get_or_create_entry_block(&mut ctx); - let pointer = entry.deref(&ctx).get_argument(0); - CallOp::new( - &mut ctx, - CallOpCallable::Direct("shared_only".try_into().unwrap()), - external_ty, - vec![pointer], - ) - .get_operation() - .insert_at_back(entry, &ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - caller.get_operation().insert_at_back(module_block, &ctx); - let mismatch = [DeviceExternDecl { - export_name: "shared_only".to_string(), - param_types: vec![DeviceExternType::pointer_to(DeviceExternType::Float32, 3)], - return_type: DeviceExternType::Void, - attrs: DeviceExternAttrs::default(), - }]; - let err = export_module_with_externs( - &ctx, - &module, - &mismatch, - &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), - ) - .expect_err("address-space mismatch must fail"); - assert!( - err.contains("parameter, result, or pointer address-space types"), - "{err}" - ); -} - -#[test] -fn device_extern_rejects_same_name_declaration_shape_without_a_call() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "extern_decl_conflict".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let void_ty = VoidType::get(&ctx); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let lowered_type = FuncType::get(&ctx, void_ty.into(), vec![i32_ty.into()], false); - FuncOp::new( - &mut ctx, - "conflicting_decl".try_into().unwrap(), - lowered_type, - ) - .get_operation() - .insert_at_back(module_block, &ctx); - - let externs = [DeviceExternDecl { - export_name: "conflicting_decl".to_string(), - param_types: vec![DeviceExternType::Float32], - return_type: DeviceExternType::Void, - attrs: DeviceExternAttrs::default(), - }]; - let error = export_module_with_externs( - &ctx, - &module, - &externs, - &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), - ) - .expect_err("the exporter must independently reject a conflicting declaration"); - assert!( - error.contains("parameter, result, or pointer address-space types"), - "{error}" - ); -} - -#[test] -fn device_extern_rejects_definition_and_address_taken_shape_conflicts() { - let mut ctx = Context::new(); - let void_ty = VoidType::get(&ctx); - - let definition_module = ModuleOp::new(&mut ctx, "extern_definition".try_into().unwrap()); - let definition_block = module_top_block(&mut ctx, &definition_module); - let no_args = FuncType::get(&ctx, void_ty.into(), vec![], false); - let definition = FuncOp::new(&mut ctx, "defined_external".try_into().unwrap(), no_args); - let definition_entry = definition.get_or_create_entry_block(&mut ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(definition_entry, &ctx); - definition - .get_operation() - .insert_at_back(definition_block, &ctx); - let definition_extern = [DeviceExternDecl { - export_name: "defined_external".to_string(), - param_types: vec![], - return_type: DeviceExternType::Void, - attrs: DeviceExternAttrs::default(), - }]; - let error = export_module_with_externs( - &ctx, - &definition_module, - &definition_extern, - &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), - ) - .expect_err("a side-table extern must not collide with a definition"); - assert!(error.contains("function definition"), "{error}"); - - let address_module = ModuleOp::new(&mut ctx, "extern_address".try_into().unwrap()); - let address_block = module_top_block(&mut ctx, &address_module); - let generic_pointer = PointerType::get(&ctx, 0); - let lowered_type = FuncType::get(&ctx, void_ty.into(), vec![generic_pointer.into()], false); - FuncOp::new( - &mut ctx, - "addressed_external".try_into().unwrap(), - lowered_type, - ) - .get_operation() - .insert_at_back(address_block, &ctx); - let caller = FuncOp::new(&mut ctx, "caller".try_into().unwrap(), lowered_type); - let caller_entry = caller.get_or_create_entry_block(&mut ctx); - let argument = caller_entry.deref(&ctx).get_argument(0); - let address = AddressOfOp::new(&mut ctx, "addressed_external".try_into().unwrap(), 0); - let address_value = address.get_operation().deref(&ctx).get_result(0); - address.get_operation().insert_at_back(caller_entry, &ctx); - CallOp::new( - &mut ctx, - CallOpCallable::Indirect(address_value), - lowered_type, - vec![argument], - ) - .get_operation() - .insert_at_back(caller_entry, &ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(caller_entry, &ctx); - caller.get_operation().insert_at_back(address_block, &ctx); - let address_extern = [DeviceExternDecl { - export_name: "addressed_external".to_string(), - param_types: vec![DeviceExternType::pointer_to(DeviceExternType::Float32, 3)], - return_type: DeviceExternType::Void, - attrs: DeviceExternAttrs::default(), - }]; - let error = export_module_with_externs( - &ctx, - &address_module, - &address_extern, - &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), - ) - .expect_err("address-taking must not bypass exact extern shape validation"); - assert!( - error.contains("parameter, result, or pointer address-space types"), - "{error}" - ); -} - -#[test] -fn legacy_kernel_metadata_uses_typed_function_references() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "legacy_metadata".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let ptr_ty = PointerType::get(&ctx, 0); - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.into(), vec![ptr_ty.into()], false); - let func = FuncOp::new(&mut ctx, "metadata_kernel".try_into().unwrap(), func_ty); - func.get_operation().deref_mut(&ctx).attributes.set( - "gpu_kernel".try_into().unwrap(), - StringAttr::new("true".into()), - ); - let entry = func.get_or_create_entry_block(&mut ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - let config = NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7); - let ir = export_module_to_string_with_config(&ctx, &module, &config) - .expect("legacy metadata export succeeds"); - assert!( - ir.contains( - "@llvm.used = appending global [1 x i8*] [i8* bitcast (void (i8*)* @metadata_kernel to i8*)]" - ), - "{ir}" - ); - assert!( - ir.contains("!{void (i8*)* @metadata_kernel, !\"kernel\", i32 1}"), - "{ir}" - ); -} - -#[test] -fn llvm_used_roots_only_explicitly_retained_globals() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "retained_globals".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - - let retained = GlobalOp::new(&mut ctx, "IKET_META".try_into().unwrap(), i32_ty.into()); - retained.set_address_space(&mut ctx, 1); - retained.mark_retained(&mut ctx); - retained.get_operation().insert_at_back(module_block, &ctx); - - let ordinary = GlobalOp::new(&mut ctx, "ORDINARY".try_into().unwrap(), i32_ty.into()); - ordinary.set_address_space(&mut ctx, 1); - ordinary.get_operation().insert_at_back(module_block, &ctx); - - let ir = export_module_to_string_with_config(&ctx, &module, &PtxExportConfig) - .expect("retained global export succeeds"); - assert!( - ir.contains( - "@llvm.used = appending global [1 x ptr] [ptr addrspacecast (ptr addrspace(1) @IKET_META to ptr)], section \"llvm.metadata\"" - ), - "{ir}" - ); - assert!(!ir.contains("@ORDINARY to ptr"), "{ir}"); -} - -#[test] -fn legacy_llvm_used_roots_retained_address_space_global() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "legacy_retained_global".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let retained = GlobalOp::new(&mut ctx, "IKET_META".try_into().unwrap(), i32_ty.into()); - retained.set_address_space(&mut ctx, 1); - retained.mark_retained(&mut ctx); - retained.get_operation().insert_at_back(module_block, &ctx); - - let config = NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7); - let ir = export_module_to_string_with_config(&ctx, &module, &config) - .expect("legacy retained global export succeeds"); - assert!( - ir.contains( - "@llvm.used = appending global [1 x i8*] [i8* addrspacecast (i32 addrspace(1)* @IKET_META to i8*)], section \"llvm.metadata\"" - ), - "{ir}" - ); -} - -#[test] -fn ptx_export_records_kernel_roots_for_internalization() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "ptx_roots".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let global = GlobalOp::new(&mut ctx, "COEFFS".try_into().unwrap(), i32_ty.into()); - global.set_address_space(&mut ctx, 4); - global.get_operation().insert_at_back(module_block, &ctx); - let func_ty = FuncType::get(&ctx, VoidType::get(&ctx).into(), vec![], false); - let func = FuncOp::new(&mut ctx, "entry_kernel".try_into().unwrap(), func_ty); - func.get_operation().deref_mut(&ctx).attributes.set( - "gpu_kernel".try_into().unwrap(), - StringAttr::new("true".into()), - ); - let entry = func.get_or_create_entry_block(&mut ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - let exported = export_module_with_externs_and_roots::( - &ctx, - &module, - &[], - &PtxExportConfig, - ) - .expect("PTX export succeeds"); - let ir = exported.llvm_ir; - assert!( - ir.contains( - "@llvm.used = appending global [1 x ptr] [ptr @entry_kernel], section \"llvm.metadata\"" - ), - "{ir}" - ); - assert_eq!(exported.public_symbols, ["COEFFS", "entry_kernel"]); -} - -#[test] -fn ptx_export_records_standalone_device_function_roots_for_internalization() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "ptx_device_root".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let func_ty = FuncType::get(&ctx, VoidType::get(&ctx).into(), vec![], false); - let prefixed_name = format!( - "{}standalone_export", - reserved_oxide_symbols::LEGACY_DEVICE_PREFIX - ); - let func = FuncOp::new( - &mut ctx, - prefixed_name.as_str().try_into().unwrap(), - func_ty, - ); - let entry = func.get_or_create_entry_block(&mut ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - let exported = export_module_with_externs_and_roots::( - &ctx, - &module, - &[], - &PtxExportConfig, - ) - .expect("standalone PTX export succeeds"); - assert_eq!(exported.public_symbols, ["standalone_export"]); - assert!( - exported.llvm_ir.contains( - "@llvm.used = appending global [1 x ptr] [ptr @standalone_export], section \"llvm.metadata\"" - ), - "{}", - exported.llvm_ir - ); -} - -#[test] -fn legacy_export_rejects_debug_metadata() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "legacy_debug".try_into().unwrap()); - let config = DebugConfig { - inner: NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), - debug_kind: DebugKind::LineTables, - }; - - let error = export_module_to_string_with_config(&ctx, &module, &config) - .expect_err("legacy debug output must be rejected"); - assert!(error.contains("legacy LLVM 7"), "{error}"); - assert!(error.contains("debug"), "{error}"); -} - -#[test] -fn legacy_pointer_slot_is_recursively_canonical() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "legacy_pointer_slot".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let ptr_ty = PointerType::get(&ctx, 0); - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.into(), vec![ptr_ty.into()], false); - let func = FuncOp::new(&mut ctx, "pointer_slot".try_into().unwrap(), func_ty); - let entry = func.get_or_create_entry_block(&mut ctx); - let incoming = entry.deref(&ctx).get_argument(0); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let one_attr = IntegerAttr::new(i32_ty, APInt::from_u32(1, NonZero::new(32).unwrap())); - let one = ConstantOp::new(&mut ctx, one_attr.into()); - let one_value = one.get_operation().deref(&ctx).get_result(0); - one.get_operation().insert_at_back(entry, &ctx); - - let slot = AllocaOp::new(&mut ctx, ptr_ty.into(), one_value); - let slot_value = slot.get_operation().deref(&ctx).get_result(0); - slot.get_operation().insert_at_back(entry, &ctx); - StoreOp::new(&mut ctx, incoming, slot_value) - .get_operation() - .insert_at_back(entry, &ctx); - LoadOp::new(&mut ctx, slot_value, ptr_ty.into()) - .get_operation() - .insert_at_back(entry, &ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - let ir = export_module_to_string_with_config( - &ctx, - &module, - &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), - ) - .expect("legacy pointer-slot export succeeds"); - assert!(ir.contains("alloca i8*"), "{ir}"); - assert!(ir.contains("bitcast i8**"), "{ir}"); - assert!(ir.matches("to i8**").count() >= 2, "{ir}"); - assert!(ir.contains("store i8* %v0, i8**"), "{ir}"); - assert!(ir.contains("load i8*, i8**"), "{ir}"); -} - -#[test] -fn export_addressof_uses_symbol_when_definition_block_prints_later() { - let mut ctx = Context::new(); - - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_region = module.get_operation().deref(&ctx).get_region(0); - let module_block = { - let existing = { - let region = module_region.deref(&ctx); - region.iter(&ctx).next() - }; - if let Some(block) = existing { - block - } else { - let block = BasicBlock::new(&mut ctx, None, vec![]); - block.insert_at_back(module_region, &ctx); - block - } - }; - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let global = GlobalOp::new( - &mut ctx, - "__shared_mem_20".try_into().unwrap(), - i32_ty.to_handle(), - ); - global.set_address_space(&mut ctx, 3); - global.get_operation().insert_at_back(module_block, &ctx); - - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); - let func = FuncOp::new(&mut ctx, "uses_late_addressof".try_into().unwrap(), func_ty); - let entry = func.get_or_create_entry_block(&mut ctx); - let func_region = func.get_operation().deref(&ctx).get_region(0); - let use_block = BasicBlock::new(&mut ctx, None, vec![]); - use_block.insert_at_back(func_region, &ctx); - let address_block = BasicBlock::new(&mut ctx, None, vec![]); - address_block.insert_at_back(func_region, &ctx); - - BrOp::new(&mut ctx, address_block, vec![]) - .get_operation() - .insert_at_back(entry, &ctx); - - let address = AddressOfOp::new(&mut ctx, "__shared_mem_20".try_into().unwrap(), 3); - let address_value = address.get_operation().deref(&ctx).get_result(0); - address.get_operation().insert_at_back(address_block, &ctx); - BrOp::new(&mut ctx, use_block, vec![]) - .get_operation() - .insert_at_back(address_block, &ctx); - - let gep = GetElementPtrOp::new( - &mut ctx, - address_value, - vec![GepIndex::Constant(0)], - i32_ty.to_handle(), - ); - gep.get_operation().insert_at_back(use_block, &ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(use_block, &ctx); - - func.get_operation().insert_at_back(module_block, &ctx); - - let ir = export_module_to_string(&ctx, &module).expect("export succeeds"); - - // The shared global must be declared at module scope. - assert!( - ir.contains("@__shared_mem_20 = addrspace(3) global"), - "module must declare the shared global:\n{ir}" - ); - - // The GEP base operand must be the global symbol, not a stale `%vN`. - let gep_line = ir - .lines() - .find(|line| line.contains("getelementptr inbounds")) - .expect("exported GEP line"); - assert!( - gep_line.contains("@__shared_mem_20"), - "GEP must use the global symbol, not a stale temporary:\n{ir}" - ); - - // Bug class from issue #54: every `%vN` reference in the IR must have a - // matching `%vN = ...` definition. With the bug present the addressof - // result was named `%v1` but never defined; this catches that and any - // future regression that re-introduces a dangling SSA reference. - assert_no_undefined_temporaries(&ir); - - let legacy = export_module_to_string_with_config( - &ctx, - &module, - &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), - ) - .expect("legacy addressof export succeeds"); - assert!( - legacy.contains("@__shared_mem_20 = addrspace(3) global i32 undef"), - "NVVM shared globals must be uninitialized:\n{legacy}" - ); - assert!( - legacy.contains("bitcast i32 addrspace(3)* @__shared_mem_20 to i8 addrspace(3)*"), - "legacy addressof must normalize the global pointer:\n{legacy}" - ); - assert!( - legacy.contains("bitcast i8 addrspace(3)*") && legacy.contains("to i32 addrspace(3)*"), - "legacy GEP must repair its canonical base pointer:\n{legacy}" - ); - assert_no_undefined_temporaries(&legacy); -} - -/// Export a module holding one shared global, optionally labelled with the -/// Rust path of the `static` it came from. -fn export_shared_global_with_source_name(source_name: Option<&str>) -> String { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let array_ty = ArrayType::get(&ctx, i32_ty.to_handle(), 64); - let global = GlobalOp::new( - &mut ctx, - "__shared_mem_7".try_into().unwrap(), - array_ty.to_handle(), - ); - global.set_address_space(&mut ctx, 3); - if let Some(source_name) = source_name { - global.set_shared_source_name(&mut ctx, source_name); - } - global.get_operation().insert_at_back(module_block, &ctx); - - export_module_to_string(&ctx, &module).expect("export succeeds") -} - -#[test] -fn shared_global_source_name_is_exported_as_a_comment_above_the_definition() { - let ir = export_shared_global_with_source_name(Some("my_kernel::TILE")); - - let definition_index = ir - .find("@__shared_mem_7 = addrspace(3) global") - .expect("module must declare the shared global"); - let comment_index = ir - .find("; shared source: my_kernel::TILE") - .unwrap_or_else(|| panic!("shared global must name its Rust source:\n{ir}")); - assert!( - comment_index < definition_index, - "the source comment must precede the definition it describes:\n{ir}" - ); -} - -#[test] -fn shared_global_without_a_source_name_exports_no_comment() { - let ir = export_shared_global_with_source_name(None); - - assert!( - ir.contains("@__shared_mem_7 = addrspace(3) global"), - "module must declare the shared global:\n{ir}" - ); - assert!( - !ir.contains("; shared source:"), - "an unlabelled global must not gain a comment:\n{ir}" - ); -} - -#[test] -fn shared_global_source_name_cannot_escape_its_comment_line() { - // A newline in the label would end the comment and leave the remainder to - // be parsed as IR. Nothing in the current pipeline produces such a name, - // so this pins the exporter's own guarantee rather than a live bug. - let ir = export_shared_global_with_source_name(Some("EVIL\n@injected = addrspace(3) global")); - - assert!( - ir.lines().all(|line| !line.starts_with("@injected")), - "a control character in the label must not open a new IR line:\n{ir}" - ); - assert!( - ir.contains("; shared source: EVIL @injected = addrspace(3) global"), - "the label must survive on one line with controls flattened:\n{ir}" - ); -} - -#[test] -fn nvvm_export_rejects_invalid_global_address_spaces() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let global = GlobalOp::new( - &mut ctx, - "thread_local_global".try_into().unwrap(), - i32_ty.to_handle(), - ); - global.set_address_space(&mut ctx, 5); - global.get_operation().insert_at_back(module_block, &ctx); - - let error = export_module_to_string_with_config( - &ctx, - &module, - &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), - ) - .expect_err("NVVM module-scope local-memory global must be rejected"); - assert!(error.contains("unsupported address space 5"), "{error}"); - - // The ordinary LLVM/PTX exporter retains its prior behavior; this - // restriction is specifically part of the NVVM IR contract. - assert!(export_module_to_string(&ctx, &module).is_ok()); -} - -#[test] -fn initialized_globals_export_exact_bytes() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "exact_global_bytes".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signless); - - // 0x7fc01234 is a quiet f32 NaN with a non-canonical payload. Treating - // these bytes as an f32 before printing would collapse it to 0x7fc00000. - let nan_ty = ArrayType::get(&ctx, i8_ty.into(), 4); - let nan = GlobalOp::new_with_alignment( - &mut ctx, - "nan_payload".try_into().unwrap(), - nan_ty.into(), - 4, - ); - nan.set_address_space(&mut ctx, 1); - nan.set_initializer_hex(&mut ctx, "3412c07f"); - nan.get_operation().insert_at_back(module_block, &ctx); - - // Byte 0 is a u8, bytes 1..4 are zeroed repr(C) padding, and bytes 4..8 - // are a little-endian u32. The exporter must not recompute those offsets. - let padded_ty = ArrayType::get(&ctx, i8_ty.into(), 8); - let padded = GlobalOp::new_with_alignment( - &mut ctx, - "padded_struct".try_into().unwrap(), - padded_ty.into(), - 4, - ); - padded.set_address_space(&mut ctx, 1); - padded.set_initializer_hex(&mut ctx, "ab00000078563412"); - padded.get_operation().insert_at_back(module_block, &ctx); - - for config in [ - NvvmExportConfig::new(NvvmIrDialect::Modern), - NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), - ] { - let ir = export_module_to_string_with_config(&ctx, &module, &config) - .expect("byte-exact global export succeeds"); - assert!( - ir.contains(r#"@nan_payload = addrspace(1) global [4 x i8] c"\34\12\C0\7F", align 4"#), - "NaN payload bytes changed:\n{ir}" - ); - assert!( - ir.contains( - r#"@padded_struct = addrspace(1) global [8 x i8] c"\AB\00\00\00\78\56\34\12", align 4"# - ), - "repr(C) layout bytes changed:\n{ir}" - ); - } -} - -#[test] -fn immutable_globals_export_the_constant_keyword() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "constant_keyword".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signless); - let table_ty = ArrayType::get(&ctx, i8_ty.into(), 4); - - // The compiler's own promoted table: marked never-written, so it must - // export as `constant`. That keyword is the whole point of the marker: - // it is what lets `opt` treat reads as invariant (deleting a copy into a - // stack slot) and what makes `llc` select `ld.global.nc`. - let promoted = GlobalOp::new_with_alignment( - &mut ctx, - "promoted_table".try_into().unwrap(), - table_ty.into(), - 4, - ); - promoted.set_address_space(&mut ctx, 1); - promoted.set_initializer_hex(&mut ctx, "01020304"); - promoted.mark_immutable(&mut ctx); - promoted.get_operation().insert_at_back(module_block, &ctx); - - // An identically shaped global without the marker: the host may still - // write such storage by symbol, so it must keep `global`. Immutability is - // opt-in per global, never inferred from the shape of the initializer. - let plain = GlobalOp::new_with_alignment( - &mut ctx, - "plain_static".try_into().unwrap(), - table_ty.into(), - 4, - ); - plain.set_address_space(&mut ctx, 1); - plain.set_initializer_hex(&mut ctx, "01020304"); - plain.get_operation().insert_at_back(module_block, &ctx); - - for config in [ - NvvmExportConfig::new(NvvmIrDialect::Modern), - NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), - ] { - let ir = export_module_to_string_with_config(&ctx, &module, &config) - .expect("immutable global export succeeds"); - assert!( - ir.contains( - r#"@promoted_table = addrspace(1) constant [4 x i8] c"\01\02\03\04", align 4"# - ), - "promoted global lost the constant keyword:\n{ir}" - ); - assert!( - ir.contains(r#"@plain_static = addrspace(1) global [4 x i8] c"\01\02\03\04", align 4"#), - "unmarked global must not become constant:\n{ir}" - ); - } -} - -#[test] -fn initialized_global_exports_static_pointer_relocation() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "static_relocation".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signless); - let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); - - // Insert the reference first. Module symbol indexing must make relocation - // resolution independent of textual global order. - let reference_ty = StructType::get_unnamed(&ctx, (vec![i64_ty.into()], StructLayout::Unpacked)); - let reference = GlobalOp::new_with_alignment( - &mut ctx, - "reference".try_into().unwrap(), - reference_ty.into(), - 8, - ); - reference.set_address_space(&mut ctx, 1); - reference.set_source_global_key(&mut ctx, "REFERENCE"); - reference.set_initializer_hex(&mut ctx, "0000000000000000"); - let encoded = encode_global_initializer_relocations(&[GlobalInitializerRelocation { - source_offset: 0, - width_bytes: 8, - target_address_space: 1, - target_addend: 0, - target_key: "TARGET".to_string(), - }]); - reference.set_initializer_relocations(&mut ctx, &encoded); - reference.get_operation().insert_at_back(module_block, &ctx); - - let target_ty = ArrayType::get(&ctx, i8_ty.into(), 4); - let target = - GlobalOp::new_with_alignment(&mut ctx, "target".try_into().unwrap(), target_ty.into(), 4); - target.set_address_space(&mut ctx, 1); - target.set_source_global_key(&mut ctx, "TARGET"); - target.set_initializer_hex(&mut ctx, "78563412"); - target.get_operation().insert_at_back(module_block, &ctx); - - let modern = export_module_to_string_with_config( - &ctx, - &module, - &NvvmExportConfig::new(NvvmIrDialect::Modern), - ) - .expect("modern relocated initializer export succeeds"); - assert!( - modern.contains( - "@reference = addrspace(1) global { i64 } { i64 ptrtoint (ptr addrspacecast (ptr addrspace(1) @target to ptr) to i64) }, align 8" - ), - "{modern}" - ); - - let legacy = export_module_to_string_with_config( - &ctx, - &module, - &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), - ) - .expect("legacy relocated initializer export succeeds"); - assert!( - legacy.contains( - "@reference = addrspace(1) global { i64 } { i64 ptrtoint (i8* addrspacecast (i8 addrspace(1)* bitcast ([4 x i8] addrspace(1)* @target to i8 addrspace(1)*) to i8*) to i64) }, align 8" - ), - "{legacy}" - ); -} - -#[test] -fn initialized_global_exports_unaligned_pointer_relocation_as_packed_storage() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "packed_static_relocation".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signless); - let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); - - let literal_ty = ArrayType::get(&ctx, i8_ty.into(), 1); - let reference_ty = StructType::get_unnamed( - &ctx, - (vec![literal_ty.into(), i64_ty.into()], StructLayout::Packed), - ); - let reference = GlobalOp::new_with_alignment( - &mut ctx, - "packed_reference".try_into().unwrap(), - reference_ty.into(), - 1, - ); - reference.set_address_space(&mut ctx, 1); - reference.set_source_global_key(&mut ctx, "PACKED_REFERENCE"); - reference.set_initializer_hex(&mut ctx, "7b0000000000000000"); - let encoded = encode_global_initializer_relocations(&[GlobalInitializerRelocation { - source_offset: 1, - width_bytes: 8, - target_address_space: 1, - target_addend: 0, - target_key: "TARGET".to_string(), - }]); - reference.set_initializer_relocations(&mut ctx, &encoded); - reference.get_operation().insert_at_back(module_block, &ctx); - - let target_ty = ArrayType::get(&ctx, i8_ty.into(), 4); - let target = - GlobalOp::new_with_alignment(&mut ctx, "target".try_into().unwrap(), target_ty.into(), 4); - target.set_address_space(&mut ctx, 1); - target.set_source_global_key(&mut ctx, "TARGET"); - target.set_initializer_hex(&mut ctx, "78563412"); - target.get_operation().insert_at_back(module_block, &ctx); - - let modern = export_module_to_string_with_config( - &ctx, - &module, - &NvvmExportConfig::new(NvvmIrDialect::Modern), - ) - .expect("modern packed relocated initializer export succeeds"); - assert!( - modern.contains( - r#"@packed_reference = addrspace(1) global <{ [1 x i8], i64 }> <{ [1 x i8] c"\7B", i64 ptrtoint (ptr addrspacecast (ptr addrspace(1) @target to ptr) to i64) }>, align 1"# - ), - "{modern}" - ); - - let legacy = export_module_to_string_with_config( - &ctx, - &module, - &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), - ) - .expect("legacy packed relocated initializer export succeeds"); - assert!( - legacy.contains( - r#"@packed_reference = addrspace(1) global <{ [1 x i8], i64 }> <{ [1 x i8] c"\7B", i64 ptrtoint (i8* addrspacecast (i8 addrspace(1)* bitcast ([4 x i8] addrspace(1)* @target to i8 addrspace(1)*) to i8*) to i64) }>, align 1"# - ), - "{legacy}" - ); -} - -#[test] -fn initialized_global_exports_multiple_relocations_and_addends() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "multiple_static_relocations".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signless); - let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); - - let target_a_ty = ArrayType::get(&ctx, i8_ty.into(), 16); - let target_a = GlobalOp::new_with_alignment( - &mut ctx, - "target_a".try_into().unwrap(), - target_a_ty.into(), - 8, - ); - target_a.set_address_space(&mut ctx, 1); - target_a.set_source_global_key(&mut ctx, "TARGET_A"); - target_a.set_initializer_hex(&mut ctx, "000102030405060708090a0b0c0d0e0f"); - target_a.get_operation().insert_at_back(module_block, &ctx); - - let target_b_ty = ArrayType::get(&ctx, i8_ty.into(), 8); - let target_b = GlobalOp::new_with_alignment( - &mut ctx, - "target_b".try_into().unwrap(), - target_b_ty.into(), - 8, - ); - target_b.set_address_space(&mut ctx, 4); - target_b.set_source_global_key(&mut ctx, "TARGET_B"); - target_b.set_initializer_hex(&mut ctx, "1011121314151617"); - target_b.get_operation().insert_at_back(module_block, &ctx); - - let table_ty = StructType::get_unnamed( - &ctx, - (vec![i64_ty.into(), i64_ty.into()], StructLayout::Unpacked), - ); - let table = GlobalOp::new_with_alignment( - &mut ctx, - "reference_table".try_into().unwrap(), - table_ty.into(), - 8, - ); - table.set_address_space(&mut ctx, 1); - table.set_source_global_key(&mut ctx, "REFERENCE_TABLE"); - table.set_initializer_hex(&mut ctx, "00000000000000000000000000000000"); - let encoded = encode_global_initializer_relocations(&[ - GlobalInitializerRelocation { - source_offset: 0, - width_bytes: 8, - target_address_space: 1, - target_addend: 4, - target_key: "TARGET_A".to_string(), - }, - GlobalInitializerRelocation { - source_offset: 8, - width_bytes: 8, - target_address_space: 4, - target_addend: 0, - target_key: "TARGET_B".to_string(), - }, - ]); - table.set_initializer_relocations(&mut ctx, &encoded); - table.get_operation().insert_at_back(module_block, &ctx); - - for dialect in [NvvmIrDialect::Modern, NvvmIrDialect::LegacyLlvm7] { - let ir = - export_module_to_string_with_config(&ctx, &module, &NvvmExportConfig::new(dialect)) - .expect("relocated initializer export succeeds"); - assert!( - ir.contains("@reference_table = addrspace(1) global { i64, i64 }"), - "{ir}" - ); - assert!(ir.contains("getelementptr (i8"), "{ir}"); - assert!(ir.contains("@target_a"), "{ir}"); - assert!(ir.contains("@target_b"), "{ir}"); - assert!(!ir.contains("inttoptr"), "{ir}"); - } -} - -#[test] -fn initialized_global_relocation_rejects_unknown_target_key() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "unknown_relocation_target".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); - let reference_ty = StructType::get_unnamed(&ctx, (vec![i64_ty.into()], StructLayout::Unpacked)); - let reference = GlobalOp::new_with_alignment( - &mut ctx, - "reference".try_into().unwrap(), - reference_ty.into(), - 8, - ); - reference.set_address_space(&mut ctx, 1); - reference.set_source_global_key(&mut ctx, "REFERENCE"); - reference.set_initializer_hex(&mut ctx, "0000000000000000"); - let encoded = encode_global_initializer_relocations(&[GlobalInitializerRelocation { - source_offset: 0, - width_bytes: 8, - target_address_space: 1, - target_addend: 0, - target_key: "MISSING".to_string(), - }]); - reference.set_initializer_relocations(&mut ctx, &encoded); - reference.get_operation().insert_at_back(module_block, &ctx); - - let error = export_module_to_string_with_config( - &ctx, - &module, - &NvvmExportConfig::new(NvvmIrDialect::Modern), - ) - .expect_err("unknown relocation target must fail"); - assert!( - error.contains("unknown rustc global key `MISSING`"), - "{error}" - ); -} - -#[test] -fn export_inline_asm_respects_sideeffect_marker() { - let mut ctx = Context::new(); - - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_region = module.get_operation().deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); - let func = FuncOp::new(&mut ctx, "has_inline_asm".try_into().unwrap(), func_ty); - let entry = func.get_or_create_entry_block(&mut ctx); - - let default_asm = InlineAsmOp::new(&mut ctx, void_ty.into(), vec![], "bar.sync 0;", "", false); - default_asm.get_operation().insert_at_back(entry, &ctx); - - let register_only_asm = InlineAsmOp::new(&mut ctx, void_ty.into(), vec![], "nop;", "", true); - llvm_export::ops::set_inline_asm_sideeffect(&mut ctx, register_only_asm.get_operation(), false); - register_only_asm - .get_operation() - .insert_at_back(entry, &ctx); - - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - let ir = export_module_to_string(&ctx, &module).expect("export succeeds"); - - assert!( - ir.contains("call void asm sideeffect \"bar.sync 0;\", \"\"()"), - "inline asm without an explicit marker should remain conservative:\n{ir}" - ); - assert!( - ir.contains("call void asm \"nop;\", \"\"() #0"), - "inline asm marked sideeffect=false should omit the keyword while preserving convergent:\n{ir}" - ); - assert!( - ir.contains("attributes #0 = { convergent }"), - "convergent inline asm must emit the convergent attr group:\n{ir}" - ); -} - -#[test] -fn export_inline_asm_escapes_llvm_string_literals() { - let mut ctx = Context::new(); - - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_region = module.get_operation().deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); - let func = FuncOp::new( - &mut ctx, - "has_escaped_inline_asm".try_into().unwrap(), - func_ty, - ); - let entry = func.get_or_create_entry_block(&mut ctx); - - let asm = InlineAsmOp::new( - &mut ctx, - void_ty.into(), - vec![], - "mov.u32 $0, %laneid;\n// \"quoted\" \\22", - "~{memory}\\raw", - false, - ); - asm.get_operation().insert_at_back(entry, &ctx); - - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - let ir = export_module_to_string(&ctx, &module).expect("export succeeds"); - - assert!( - ir.contains( - "call void asm sideeffect \"mov.u32 $0, %laneid;\\0A// \\22quoted\\22 \\5C22\", \"~{memory}\\5Craw\"()" - ), - "inline asm template and constraints must be escaped as LLVM string literals:\n{ir}" - ); -} - -#[test] -fn nvvm_metadata_version_uses_next_allocated_metadata_id() { - let mut ctx = Context::new(); - - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_region = module.get_operation().deref(&ctx).get_region(0); - let module_block = { - let region = module_region.deref(&ctx); - region.iter(&ctx).next().unwrap() - }; - - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); - let func = FuncOp::new(&mut ctx, "bounded_kernel".try_into().unwrap(), func_ty); - let entry = func.get_or_create_entry_block(&mut ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let width = NonZero::new(32).unwrap(); - let max_threads = IntegerAttr::new(u32_ty, APInt::from_u32(256, width)); - let min_blocks = IntegerAttr::new(u32_ty, APInt::from_u32(2, width)); - - { - let attrs = &mut func.get_operation().deref_mut(&ctx).attributes; - attrs.set( - Identifier::try_from("gpu_kernel").unwrap(), - StringAttr::new("true".into()), - ); - attrs.set(Identifier::try_from("maxntid").unwrap(), max_threads); - attrs.set(Identifier::try_from("minctasm").unwrap(), min_blocks); - } - - func.get_operation().insert_at_back(module_block, &ctx); - - let ir = export_module_to_string_with_config(&ctx, &module, &NvvmExportConfig::default()) - .expect("NVVM export succeeds"); - - assert!( - ir.contains("!0 = !{ptr @bounded_kernel, !\"kernel\", i32 1}"), - "a launch-bounded kernel still needs its kernel annotation:\n{ir}" - ); - assert!( - ir.contains("!nvvm.annotations = !{!0, !1, !2, !3, !4}"), - "kernel identity plus launch-bounds annotations should occupy !0..!4:\n{ir}" - ); - assert!( - ir.contains("!nvvmir.version = !{!5}\n!5 = !{i32 2, i32 0, i32 3, i32 2}"), - "version metadata should use the next allocated ID:\n{ir}" - ); -} - -#[test] -fn line_table_debug_metadata_emits_function_scope_and_instruction_locations() { - let mut ctx = Context::new(); - - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_region = module.get_operation().deref(&ctx).get_region(0); - let module_block = { - let region = module_region.deref(&ctx); - region.iter(&ctx).next().unwrap() - }; - - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); - let func = FuncOp::new(&mut ctx, "debug_kernel".try_into().unwrap(), func_ty); - llvm_export::ops::set_debug_function_name( - &mut ctx, - func.get_operation(), - "source_crate::debug_kernel", - ); - let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 7, 1); - func.get_operation().deref_mut(&ctx).set_loc(func_loc); - - let entry = func.get_or_create_entry_block(&mut ctx); - let ret = ReturnOp::new(&mut ctx, None); - let ret_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 8, 5); - ret.get_operation().deref_mut(&ctx).set_loc(ret_loc); - ret.get_operation().insert_at_back(entry, &ctx); - - func.get_operation().insert_at_back(module_block, &ctx); - - let config = DebugConfig { - inner: PtxExportConfig, - debug_kind: DebugKind::LineTables, - }; - let ir = - export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); - - let define_line = ir - .lines() - .find(|line| line.starts_with("define ")) - .expect("function definition"); - assert!( - define_line.contains("!dbg !"), - "function definition should reference its DISubprogram:\n{ir}" - ); - - let ret_line = ir - .lines() - .find(|line| line.trim_start().starts_with("ret void")) - .expect("return instruction"); - assert!( - ret_line.contains(", !dbg !"), - "real instructions should carry DILocation attachments:\n{ir}" - ); - - assert!( - ir.contains("!llvm.dbg.cu = !{!"), - "module should reference a compile unit:\n{ir}" - ); - assert!( - ir.contains("!llvm.module.flags = !{!"), - "module should declare debug-info flags:\n{ir}" - ); - assert!( - ir.contains("!DIFile(filename: \"kernel.rs\", directory: \"/tmp/cuda-oxide/tests\")"), - "source path should be split into DIFile filename and directory:\n{ir}" - ); - assert!( - ir.contains("distinct !DICompileUnit(language: DW_LANG_Rust"), - "debug export should describe the Rust compile unit:\n{ir}" - ); - assert!( - ir.contains( - "distinct !DISubprogram(name: \"source_crate::debug_kernel\", linkageName: \"debug_kernel\"" - ), - "function definition should separate its source name from its physical linkage name:\n{ir}" - ); - assert!( - ir.contains("!DILocation(line: 8, column: 5, scope: !"), - "instruction location should preserve the source line and column:\n{ir}" - ); -} - -#[test] -fn full_debug_metadata_keeps_generic_source_name_separate_from_mangled_symbol() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); - let linkage_name = "_RNvMCexampleINtCsource7WrapperjE7get_mut"; - let func = FuncOp::new(&mut ctx, linkage_name.try_into().unwrap(), func_ty); - llvm_export::ops::set_debug_function_name( - &mut ctx, - func.get_operation(), - "source_crate::Wrapper::::get_mut", - ); - let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/generic.rs", 19, 1); - func.get_operation().deref_mut(&ctx).set_loc(func_loc); - - let entry = func.get_or_create_entry_block(&mut ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - let config = DebugConfig { - inner: PtxExportConfig, - debug_kind: DebugKind::Full, - }; - let ir = - export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); - - assert!( - ir.contains(&format!("define void @{linkage_name}()")), - "debug naming must not change the physical function symbol:\n{ir}" - ); - assert!( - ir.contains(&format!( - "distinct !DISubprogram(name: \"source_crate::Wrapper::::get_mut\", linkageName: \"{linkage_name}\"" - )), - "generic debug metadata must carry both source and linkage spellings:\n{ir}" - ); -} - -#[test] -fn export_alwaysinline_function_attribute_uses_llvm_define_syntax() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); - let func = FuncOp::new(&mut ctx, "inline_helper".try_into().unwrap(), func_ty); - let entry = func.get_or_create_entry_block(&mut ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - - let key: pliron::identifier::Identifier = "alwaysinline".try_into().unwrap(); - func.get_operation() - .deref_mut(&ctx) - .attributes - .set(key, StringAttr::new("true".to_string())); - func.get_operation().insert_at_back(module_block, &ctx); - - let ir = export_module_to_string(&ctx, &module).expect("export succeeds"); - let define_line = ir - .lines() - .find(|line| line.starts_with("define void @inline_helper(")) - .expect("inline helper definition"); - assert_eq!( - define_line, "define void @inline_helper() alwaysinline #0 {", - "`alwaysinline` must be emitted after the parameter list, before attr group #0:\n{ir}" - ); - assert!( - ir.contains("attributes #0 = { convergent }"), - "convergent attribute group must still be emitted:\n{ir}" - ); -} - -#[test] -fn export_alwaysinline_coexists_with_debug_scope() { - // alwaysinline and the !dbg scope are emitted on the same define line and - // must not crowd each other out. This guards the 4-way emission: a future - // change that drops either one when both are present fails here. - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); - let func = FuncOp::new(&mut ctx, "inline_helper".try_into().unwrap(), func_ty); - let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 7, 1); - func.get_operation().deref_mut(&ctx).set_loc(func_loc); - let entry = func.get_or_create_entry_block(&mut ctx); - let ret = ReturnOp::new(&mut ctx, None); - let ret_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 8, 5); - ret.get_operation().deref_mut(&ctx).set_loc(ret_loc); - ret.get_operation().insert_at_back(entry, &ctx); - - let key: pliron::identifier::Identifier = "alwaysinline".try_into().unwrap(); - func.get_operation() - .deref_mut(&ctx) - .attributes - .set(key, StringAttr::new("true".to_string())); - func.get_operation().insert_at_back(module_block, &ctx); - - let config = DebugConfig { - inner: PtxExportConfig, - debug_kind: DebugKind::LineTables, - }; - let ir = - export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); - let define_line = ir - .lines() - .find(|line| line.starts_with("define void @inline_helper(")) - .expect("inline helper definition"); - assert!( - define_line.contains("alwaysinline"), - "alwaysinline must survive when debug info is on:\n{ir}" - ); - assert!( - define_line.contains("!dbg !"), - "!dbg scope must survive when alwaysinline is present:\n{ir}" - ); -} - -#[test] -fn line_table_debug_metadata_uses_file_scope_for_cross_file_locations() { - let mut ctx = Context::new(); - - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_region = module.get_operation().deref(&ctx).get_region(0); - let module_block = { - let region = module_region.deref(&ctx); - region.iter(&ctx).next().unwrap() - }; - - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); - let func = FuncOp::new(&mut ctx, "debug_kernel".try_into().unwrap(), func_ty); - let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 38, 1); - func.get_operation().deref_mut(&ctx).set_loc(func_loc); - - let entry = func.get_or_create_entry_block(&mut ctx); - let ret = ReturnOp::new(&mut ctx, None); - let ret_loc = src_location( - &mut ctx, - "/tmp/cuda-oxide/crates/cuda-device/src/thread.rs", - 292, - 19, - ); - ret.get_operation().deref_mut(&ctx).set_loc(ret_loc); - ret.get_operation().insert_at_back(entry, &ctx); - - func.get_operation().insert_at_back(module_block, &ctx); - - let config = DebugConfig { - inner: PtxExportConfig, - debug_kind: DebugKind::LineTables, - }; - let ir = - export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); - - let ret_line = ir - .lines() - .find(|line| line.trim_start().starts_with("ret void")) - .expect("return instruction"); - assert!( - ret_line.contains(", !dbg !"), - "cross-file instructions should keep source locations:\n{ir}" - ); - assert!( - ir.contains( - "!DIFile(filename: \"thread.rs\", directory: \"/tmp/cuda-oxide/crates/cuda-device/src\")" - ), - "cross-file locations should get their own DIFile:\n{ir}" - ); - assert!( - ir.contains("!DILexicalBlockFile(scope: !"), - "cross-file locations should use a file-specific debug scope:\n{ir}" - ); - assert!( - ir.contains("!DILocation(line: 292, column: 19, scope: !"), - "cross-file locations should preserve their real source line:\n{ir}" - ); -} - -#[test] -fn line_table_debug_metadata_emits_inlined_at_for_callsite_locations() { - let mut ctx = Context::new(); - - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_region = module.get_operation().deref(&ctx).get_region(0); - let module_block = { - let region = module_region.deref(&ctx); - region.iter(&ctx).next().unwrap() - }; - - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); - let func = FuncOp::new(&mut ctx, "debug_kernel".try_into().unwrap(), func_ty); - let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 38, 1); - func.get_operation().deref_mut(&ctx).set_loc(func_loc); - - let entry = func.get_or_create_entry_block(&mut ctx); - let ret = ReturnOp::new(&mut ctx, None); - let callee = src_location( - &mut ctx, - "/tmp/cuda-oxide/crates/cuda-device/src/thread.rs", - 292, - 19, - ); - let caller = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 39, 13); - ret.get_operation() - .deref_mut(&ctx) - .set_loc(Location::CallSite { - callee: Box::new(callee), - caller: Box::new(caller), - }); - ret.get_operation().insert_at_back(entry, &ctx); - - func.get_operation().insert_at_back(module_block, &ctx); - - let config = DebugConfig { - inner: PtxExportConfig, - debug_kind: DebugKind::LineTables, - }; - let ir = - export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); - - assert!( - ir.contains("!DILocation(line: 39, column: 13, scope: !"), - "callsite metadata should preserve the caller location:\n{ir}" - ); - assert!( - ir.contains("!DILocation(line: 292, column: 19, scope: !") && ir.contains(", inlinedAt: !"), - "callsite metadata should describe the callee location as inlined at the caller:\n{ir}" - ); -} - -#[test] -fn debug_metadata_shares_allocator_with_nvvm_metadata() { - let mut ctx = Context::new(); - - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_region = module.get_operation().deref(&ctx).get_region(0); - let module_block = { - let region = module_region.deref(&ctx); - region.iter(&ctx).next().unwrap() - }; - - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); - let func = FuncOp::new(&mut ctx, "debug_kernel".try_into().unwrap(), func_ty); - let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 10, 1); - func.get_operation().deref_mut(&ctx).set_loc(func_loc); - - { - let attrs = &mut func.get_operation().deref_mut(&ctx).attributes; - attrs.set( - Identifier::try_from("gpu_kernel").unwrap(), - StringAttr::new("true".into()), - ); - } - - let entry = func.get_or_create_entry_block(&mut ctx); - let ret = ReturnOp::new(&mut ctx, None); - let ret_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 11, 5); - ret.get_operation().deref_mut(&ctx).set_loc(ret_loc); - ret.get_operation().insert_at_back(entry, &ctx); - - func.get_operation().insert_at_back(module_block, &ctx); - - let config = DebugConfig { - inner: NvvmExportConfig::default(), - debug_kind: DebugKind::LineTables, - }; - let ir = export_module_to_string_with_config(&ctx, &module, &config) - .expect("debug NVVM export succeeds"); - - assert!( - ir.contains("!0 = !DIFile(filename: \"kernel.rs\", directory: \"/tmp/cuda-oxide/tests\")"), - "debug file node should take the first metadata ID:\n{ir}" - ); - assert!( - ir.contains("!4 = !DILocation(line: 11, column: 5, scope: !3)"), - "instruction location should be allocated before NVVM metadata:\n{ir}" - ); - assert!( - ir.contains("!5 = !{ptr @debug_kernel, !\"kernel\", i32 1}"), - "NVVM annotations should continue after debug metadata:\n{ir}" - ); - assert!( - ir.contains("!nvvm.annotations = !{!5}"), - "named NVVM metadata should reference its allocated node:\n{ir}" - ); - assert!( - ir.contains("!nvvmir.version = !{!6}\n!6 = !{i32 2, i32 0, i32 3, i32 2}"), - "NVVM version should use the next free metadata ID:\n{ir}" - ); - assert!( - ir.contains("!llvm.module.flags = !{!7, !8}"), - "debug module flags should also use the shared allocator:\n{ir}" - ); -} - -#[test] -fn debug_locations_use_rustc_source_scope_positions() { - let mut ctx = Context::new(); - - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_region = module.get_operation().deref(&ctx).get_region(0); - let module_block = { - let region = module_region.deref(&ctx); - region.iter(&ctx).next().unwrap() - }; - - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); - let func = FuncOp::new(&mut ctx, "debug_kernel".try_into().unwrap(), func_ty); - let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 10, 1); - func.get_operation().deref_mut(&ctx).set_loc(func_loc); - llvm_export::ops::set_debug_source_scope_map( - &mut ctx, - func.get_operation(), - &DebugSourceScopeMap { - scopes: vec![ - DebugSourceScope { - id: 0, - parent: None, - span: Some(DebugSourcePosition { - file: PathBuf::from("/tmp/cuda-oxide/tests/kernel.rs"), - line: 10, - column: 1, - }), - inlined: None, - }, - DebugSourceScope { - id: 1, - parent: Some(0), - span: Some(DebugSourcePosition { - file: PathBuf::from("/tmp/cuda-oxide/tests/kernel.rs"), - line: 12, - column: 9, - }), - inlined: None, - }, - ], - locations: vec![DebugSourceScopeLocation { - pos: DebugSourcePosition { - file: PathBuf::from("/tmp/cuda-oxide/tests/kernel.rs"), - line: 12, - column: 9, - }, - scope: 1, - }], - }, - ); - - let entry = func.get_or_create_entry_block(&mut ctx); - let ret = ReturnOp::new(&mut ctx, None); - let ret_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 12, 9); - ret.get_operation().deref_mut(&ctx).set_loc(ret_loc); - ret.get_operation().insert_at_back(entry, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - let config = DebugConfig { - inner: PtxExportConfig, - debug_kind: DebugKind::LineTables, - }; - let ir = - export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); - - let block_id = ir - .lines() - .find_map(|line| { - if line.contains("!DILexicalBlock(scope: !") && line.contains("line: 12, column: 9") { - line.split_once(" = ") - .map(|(id, _)| id.trim_start_matches('!').to_string()) - } else { - None - } - }) - .expect("nested lexical block should be emitted"); - - assert!( - ir.contains(&format!( - "!DILocation(line: 12, column: 9, scope: !{block_id})" - )), - "instruction location should use the exact rustc source scope, not the function scope:\n{ir}" - ); -} - -#[test] -fn full_debug_metadata_emits_dbg_declare_for_tagged_allocas() { - let mut ctx = Context::new(); - - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_region = module.get_operation().deref(&ctx).get_region(0); - let module_block = { - let region = module_region.deref(&ctx); - region.iter(&ctx).next().unwrap() - }; - - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); - let func = FuncOp::new(&mut ctx, "debug_kernel".try_into().unwrap(), func_ty); - let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 30, 1); - func.get_operation().deref_mut(&ctx).set_loc(func_loc); - - let entry = func.get_or_create_entry_block(&mut ctx); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let one_attr = IntegerAttr::new(i32_ty, APInt::from_u32(1, NonZero::new(32).unwrap())); - let one = ConstantOp::new(&mut ctx, one_attr.into()); - one.get_operation().insert_at_back(entry, &ctx); - let one_val = one.get_operation().deref(&ctx).get_result(0); - - let tid = AllocaOp::new(&mut ctx, i32_ty.into(), one_val); - let tid_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 31, 9); - tid.get_operation().deref_mut(&ctx).set_loc(tid_loc); - llvm_export::ops::set_debug_local_variable( - &mut ctx, - tid.get_operation(), - DebugLocalVariableInfo { - name: "tid".to_string(), - argument_index: Some(1), - ty: DebugLocalTypeKind::Basic { - name: "u32".to_string(), - size_bits: 32, - encoding: "DW_ATE_unsigned", - }, - }, - ); - tid.get_operation().insert_at_back(entry, &ctx); - - let ptr_ty = PointerType::get(&ctx, 0); - let ptr = AllocaOp::new(&mut ctx, ptr_ty.into(), one_val); - let ptr_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 32, 9); - ptr.get_operation().deref_mut(&ctx).set_loc(ptr_loc); - llvm_export::ops::set_debug_local_variable( - &mut ctx, - ptr.get_operation(), - DebugLocalVariableInfo { - name: "ptr".to_string(), - argument_index: None, - ty: DebugLocalTypeKind::TypedPointer { - name: "*mut f32".to_string(), - size_bits: 64, - pointee: Box::new(DebugLocalTypeKind::Basic { - name: "f32".to_string(), - size_bits: 32, - encoding: "DW_ATE_float", - }), - }, - }, - ); - ptr.get_operation().insert_at_back(entry, &ctx); - - let nested_ptr = AllocaOp::new(&mut ctx, ptr_ty.into(), one_val); - let nested_ptr_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 33, 9); - nested_ptr - .get_operation() - .deref_mut(&ctx) - .set_loc(nested_ptr_loc); - llvm_export::ops::set_debug_local_variable( - &mut ctx, - nested_ptr.get_operation(), - DebugLocalVariableInfo { - name: "nested_ptr".to_string(), - argument_index: None, - ty: DebugLocalTypeKind::TypedPointer { - name: "*const *mut i32".to_string(), - size_bits: 64, - pointee: Box::new(DebugLocalTypeKind::TypedPointer { - name: "*mut i32".to_string(), - size_bits: 64, - pointee: Box::new(DebugLocalTypeKind::Basic { - name: "i32".to_string(), - size_bits: 32, - encoding: "DW_ATE_signed", - }), - }), - }, - }, - ); - nested_ptr.get_operation().insert_at_back(entry, &ctx); - - let array_ptr = AllocaOp::new(&mut ctx, ptr_ty.into(), one_val); - let array_ptr_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 34, 9); - array_ptr - .get_operation() - .deref_mut(&ctx) - .set_loc(array_ptr_loc); - llvm_export::ops::set_debug_local_variable( - &mut ctx, - array_ptr.get_operation(), - DebugLocalVariableInfo { - name: "array_ptr".to_string(), - argument_index: None, - ty: DebugLocalTypeKind::TypedPointer { - name: "*const [u16; 4]".to_string(), - size_bits: 64, - pointee: Box::new(DebugLocalTypeKind::Array { - name: "[u16; 4]".to_string(), - size_bits: 64, - element: Box::new(DebugLocalTypeKind::Basic { - name: "u16".to_string(), - size_bits: 16, - encoding: "DW_ATE_unsigned", - }), - count: 4, - }), - }, - }, - ); - array_ptr.get_operation().insert_at_back(entry, &ctx); - - let opaque_ptr = AllocaOp::new(&mut ctx, ptr_ty.into(), one_val); - let opaque_ptr_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 35, 9); - opaque_ptr - .get_operation() - .deref_mut(&ctx) - .set_loc(opaque_ptr_loc); - llvm_export::ops::set_debug_local_variable( - &mut ctx, - opaque_ptr.get_operation(), - DebugLocalVariableInfo { - name: "opaque_ptr".to_string(), - argument_index: None, - ty: DebugLocalTypeKind::Pointer { - name: "*const _".to_string(), - size_bits: 64, - }, - }, - ); - opaque_ptr.get_operation().insert_at_back(entry, &ctx); - - let ret = ReturnOp::new(&mut ctx, None); - let ret_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 36, 1); - ret.get_operation().deref_mut(&ctx).set_loc(ret_loc); - ret.get_operation().insert_at_back(entry, &ctx); - - func.get_operation().insert_at_back(module_block, &ctx); - - let config = DebugConfig { - inner: PtxExportConfig, - debug_kind: DebugKind::Full, - }; - let ir = - export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); - - assert!( - ir.contains("emissionKind: FullDebug"), - "full debug should request full DWARF metadata:\n{ir}" - ); - assert!( - ir.contains("isOptimized: false"), - "full debug export should describe the unoptimized debug path:\n{ir}" - ); - assert!( - ir.contains("declare void @llvm.dbg.declare(metadata, metadata, metadata)"), - "full debug should declare the debug intrinsic it calls:\n{ir}" - ); - assert!( - ir.contains("call void @llvm.dbg.declare(metadata ptr %"), - "tagged allocas should be bound to variables with dbg.declare:\n{ir}" - ); - assert!( - ir.contains("!DILocalVariable(name: \"tid\", arg: 1, scope: !"), - "argument debug metadata should preserve the argument number:\n{ir}" - ); - assert!( - ir.contains("!DILocalVariable(name: \"ptr\", scope: !"), - "local debug metadata should omit the arg field:\n{ir}" - ); - assert!( - ir.contains("!DIBasicType(name: \"u32\", size: 32, encoding: DW_ATE_unsigned)"), - "basic integer variables should get DIBasicType metadata:\n{ir}" - ); - let f32_id = metadata_id( - &ir, - "!DIBasicType(name: \"f32\", size: 32, encoding: DW_ATE_float)", - ); - assert!( - ir.contains(&format!( - "!DIDerivedType(tag: DW_TAG_pointer_type, name: \"*mut f32\", baseType: {f32_id}, size: 64)" - )), - "pointer variables should reference their pointee DIType:\n{ir}" - ); - let i32_id = metadata_id( - &ir, - "!DIBasicType(name: \"i32\", size: 32, encoding: DW_ATE_signed)", - ); - let inner_pointer = format!( - "!DIDerivedType(tag: DW_TAG_pointer_type, name: \"*mut i32\", baseType: {i32_id}, size: 64)" - ); - let inner_pointer_id = metadata_id(&ir, &inner_pointer); - assert!( - ir.contains(&format!( - "!DIDerivedType(tag: DW_TAG_pointer_type, name: \"*const *mut i32\", baseType: {inner_pointer_id}, size: 64)" - )), - "nested pointer metadata must preserve the complete base-type chain:\n{ir}" - ); - - let u16_id = metadata_id( - &ir, - "!DIBasicType(name: \"u16\", size: 16, encoding: DW_ATE_unsigned)", - ); - let subrange_id = metadata_id(&ir, "!DISubrange(count: 4)"); - let subrange_tuple_id = metadata_id(&ir, &format!("!{{{subrange_id}}}")); - let array_type = format!( - "!DICompositeType(tag: DW_TAG_array_type, baseType: {u16_id}, size: 64, elements: {subrange_tuple_id})" - ); - let array_type_id = metadata_id(&ir, &array_type); - assert!( - ir.contains(&format!( - "!DIDerivedType(tag: DW_TAG_pointer_type, name: \"*const [u16; 4]\", baseType: {array_type_id}, size: 64)" - )), - "array pointer metadata must reference the exact element/subrange graph:\n{ir}" - ); - - assert!( - ir.contains("!DILocalVariable(name: \"opaque_ptr\", scope: !"), - "unsupported pointers must remain visible as source variables:\n{ir}" - ); - assert!( - ir.contains( - "!DIDerivedType(tag: DW_TAG_pointer_type, name: \"*const _\", baseType: null, size: 64)" - ), - "the compatibility pointer must retain its legacy null-base metadata:\n{ir}" - ); - let null_base_pointer_count = ir - .lines() - .filter(|line| line.contains("DW_TAG_pointer_type") && line.contains("baseType: null")) - .count(); - assert!( - null_base_pointer_count == 1, - "only the explicit opaque compatibility pointer may have a null baseType:\n{ir}" - ); - - let Some(llvm_as) = ["llvm-as-22", "llvm-as-21", "llvm-as"] - .into_iter() - .find(|tool| { - std::process::Command::new(tool) - .arg("--version") - .output() - .is_ok_and(|out| out.status.success()) - }) - else { - eprintln!("skipping pointer-debug llvm-as parse gate: no supported llvm-as on PATH"); - return; - }; - let ll_path = std::env::temp_dir().join(format!( - "cuda_oxide_pointer_debug_parse_gate_{}.ll", - std::process::id() - )); - std::fs::write(&ll_path, &ir).expect("write pointer-debug temp .ll"); - let output = std::process::Command::new(llvm_as) - .arg("-o") - .arg("/dev/null") - .arg(&ll_path) - .output() - .expect("run llvm-as for pointer debug metadata"); - let stderr = String::from_utf8_lossy(&output.stderr).into_owned(); - let _ = std::fs::remove_file(&ll_path); - assert!( - output.status.success(), - "{llvm_as} rejected pointer debug metadata:\n{stderr}\n--- module ---\n{ir}" - ); -} - -#[test] -fn full_debug_metadata_describes_as1_global_with_semantic_rust_type() { - let mut ctx = Context::new(); - - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signless); - let storage_ty = ArrayType::get(&ctx, i8_ty.into(), 8); - let global = GlobalOp::new_with_alignment( - &mut ctx, - "__device_global_7".try_into().unwrap(), - storage_ty.into(), - 8, - ); - global.set_address_space(&mut ctx, llvm_export::types::address_space::GLOBAL); - global.set_initializer_hex(&mut ctx, "0000000000000000"); - llvm_export::ops::set_debug_global_variable( - &mut ctx, - global.get_operation(), - &DebugGlobalVariableInfo { - name: "GLOBAL_COUNTER".to_string(), - namespace: vec!["debuginfo".to_string(), "state".to_string()], - ty: DebugLocalTypeKind::Basic { - name: "u64".to_string(), - size_bits: 64, - encoding: "DW_ATE_unsigned", - }, - declaration: DebugSourcePosition { - file: PathBuf::from("/tmp/cuda-oxide/tests/kernel.rs"), - line: 90, - column: 1, - }, - is_local_to_unit: true, - is_function_local: false, - }, - ); - global.get_operation().insert_at_back(module_block, &ctx); - - let full = export_module_to_string_with_config( - &ctx, - &module, - &DebugConfig { - inner: PtxExportConfig, - debug_kind: DebugKind::Full, - }, - ) - .expect("full debug export succeeds"); - - let definition = full - .lines() - .find(|line| line.starts_with("@__device_global_7 = ")) - .expect("device-global definition"); - assert!( - definition.contains("addrspace(1) global [8 x i8] c\"\\00\\00\\00\\00\\00\\00\\00\\00\"") - && definition.contains(", align 8, !dbg !"), - "the physical byte storage should retain a global debug attachment:\n{full}" - ); - assert!( - full.contains( - "distinct !DIGlobalVariable(name: \"GLOBAL_COUNTER\", linkageName: \"__device_global_7\"" - ), - "DWARF should separate the source name from the generated linkage name:\n{full}" - ); - assert!( - full.contains("!DINamespace(name: \"debuginfo\", scope: null)") - && full.contains("!DINamespace(name: \"state\", scope: !") - && full.contains("scope: !"), - "the source crate/module hierarchy should scope the leaf name:\n{full}" - ); - assert!( - full.contains("isLocal: true, isDefinition: true"), - "private Rust statics should remain local to the compile unit:\n{full}" - ); - assert!( - full.contains("file: !") - && full.contains("line: 90, type: !") - && full.contains("align: 64)"), - "the declaration location and source alignment should be retained:\n{full}" - ); - assert!( - full.contains("!DIBasicType(name: \"u64\", size: 64, encoding: DW_ATE_unsigned)"), - "the debug type must be semantic u64, not the physical [8 x i8] storage:\n{full}" - ); - assert!( - full.contains("!DIGlobalVariableExpression(var: !") - && full.contains("expr: !DIExpression())") - && full.contains("emissionKind: FullDebug, globals: !"), - "the global expression should be retained by the compile unit:\n{full}" - ); - - let off = export_module_to_string_with_config( - &ctx, - &module, - &DebugConfig { - inner: PtxExportConfig, - debug_kind: DebugKind::Off, - }, - ) - .expect("non-debug export succeeds"); - let line_tables = export_module_to_string_with_config( - &ctx, - &module, - &DebugConfig { - inner: PtxExportConfig, - debug_kind: DebugKind::LineTables, - }, - ) - .expect("line-table export succeeds"); - assert_eq!( - line_tables, off, - "a global-only module must not gain variable metadata outside full debug" - ); -} - -#[test] -fn full_debug_metadata_describes_function_local_as3_array_with_shared_address_class() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "shared_debug".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let raw_owner = reserved_oxide_symbols::device_symbol("shared_kernel"); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let storage_ty = ArrayType::get(&ctx, i32_ty.into(), 32); - let global = GlobalOp::new_with_alignment( - &mut ctx, - "__shared_mem_0".try_into().unwrap(), - storage_ty.into(), - 4, - ); - global.set_address_space(&mut ctx, llvm_export::types::address_space::SHARED); - let info = DebugGlobalVariableInfo { - name: "TILE".to_string(), - namespace: vec![ - "debuginfo".to_string(), - "kernels".to_string(), - "shared_kernel".to_string(), - ], - ty: DebugLocalTypeKind::Array { - name: "[i32; 32]".to_string(), - size_bits: 1024, - element: Box::new(DebugLocalTypeKind::Basic { - name: "i32".to_string(), - size_bits: 32, - encoding: "DW_ATE_signed", - }), - count: 32, - }, - declaration: DebugSourcePosition { - file: PathBuf::from("/tmp/cuda-oxide/tests/shared.rs"), - line: 40, - column: 9, - }, - is_local_to_unit: true, - is_function_local: true, - }; - llvm_export::ops::set_debug_global_variable(&mut ctx, global.get_operation(), &info); - llvm_export::ops::set_debug_global_owner_function(&mut ctx, global.get_operation(), &raw_owner); - global.get_operation().insert_at_back(module_block, &ctx); - - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); - let function = FuncOp::new(&mut ctx, raw_owner.as_str().try_into().unwrap(), func_ty); - let function_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/shared.rs", 35, 1); - function - .get_operation() - .deref_mut(&ctx) - .set_loc(function_loc); - let entry = function.get_or_create_entry_block(&mut ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - function.get_operation().insert_at_back(module_block, &ctx); - - let full = export_module_to_string_with_config( - &ctx, - &module, - &DebugConfig { - inner: PtxExportConfig, - debug_kind: DebugKind::Full, - }, - ) - .expect("full shared debug export succeeds"); - - let definition = full - .lines() - .find(|line| line.starts_with("@__shared_mem_0 = ")) - .expect("shared definition"); - assert!( - definition.contains("addrspace(3) global [32 x i32]") - && definition.contains(", align 4, !dbg !"), - "physical AS3 array must retain the debug attachment:\n{full}" - ); - assert!( - full.contains("!DIExpression(DW_OP_constu, 8, DW_OP_swap, DW_OP_xderef)"), - "AS3 must carry CUDA DWARF shared address class 8:\n{full}" - ); - assert!( - full.contains("!DINamespace(name: \"debuginfo\", scope: null)") - && full.contains("!DINamespace(name: \"kernels\", scope: !") - && full.contains( - "distinct !DISubprogram(name: \"shared_kernel\", \ - linkageName: \"shared_kernel\", scope: !" - ), - "the owning subprogram must be nested under the structured namespace:\n{full}" - ); - assert_eq!( - full.matches("distinct !DISubprogram(name: \"shared_kernel\"") - .count(), - 1, - "pre-reservation and function export must reuse one DISubprogram:\n{full}" - ); - assert!( - full.contains("distinct !DIGlobalVariable(name: \"TILE\", linkageName: \"__shared_mem_0\"") - && full.contains("!DISubrange(count: 32)") - && full.contains("!DICompositeType(tag: DW_TAG_array_type, baseType: !") - && full.contains("size: 1024, elements: !"), - "the DIE must keep leaf/linkage and the logical array type:\n{full}" - ); - assert!( - full.contains("emissionKind: FullDebug, globals: !"), - "the CU must retain the one shared expression:\n{full}" - ); -} - -fn export_mixed_shared_owner_order(malformed_first: bool) -> String { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "shared_owner_validation".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let raw_owner = reserved_oxide_symbols::device_symbol("shared_kernel"); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let storage_ty = ArrayType::get(&ctx, i32_ty.into(), 4); - - let valid = GlobalOp::new_with_alignment( - &mut ctx, - "__shared_mem_valid".try_into().unwrap(), - storage_ty.into(), - 4, - ); - valid.set_address_space(&mut ctx, llvm_export::types::address_space::SHARED); - llvm_export::ops::set_debug_global_variable( - &mut ctx, - valid.get_operation(), - &DebugGlobalVariableInfo { - name: "VALID".to_string(), - namespace: vec!["owner_fixture".to_string(), "shared_kernel".to_string()], - ty: DebugLocalTypeKind::Array { - name: "[i32; 4]".to_string(), - size_bits: 128, - element: Box::new(DebugLocalTypeKind::Basic { - name: "i32".to_string(), - size_bits: 32, - encoding: "DW_ATE_signed", - }), - count: 4, - }, - declaration: DebugSourcePosition { - file: PathBuf::from("/tmp/cuda-oxide/tests/shared-owner.rs"), - line: 20, - column: 5, - }, - is_local_to_unit: true, - is_function_local: true, - }, - ); - llvm_export::ops::set_debug_global_owner_function(&mut ctx, valid.get_operation(), &raw_owner); - - let malformed = GlobalOp::new_with_alignment( - &mut ctx, - "__shared_mem_malformed".try_into().unwrap(), - storage_ty.into(), - 4, - ); - malformed.set_address_space(&mut ctx, llvm_export::types::address_space::SHARED); - llvm_export::ops::set_debug_global_variable( - &mut ctx, - malformed.get_operation(), - &DebugGlobalVariableInfo { - name: "MALFORMED".to_string(), - namespace: vec!["owner_fixture".to_string(), "shared_kernel".to_string()], - ty: DebugLocalTypeKind::Array { - name: "[i32; 4]".to_string(), - size_bits: 128, - element: Box::new(DebugLocalTypeKind::Basic { - name: "i32".to_string(), - size_bits: 32, - encoding: "DW_ATE_signed", - }), - count: 4, - }, - declaration: DebugSourcePosition { - file: PathBuf::from("/tmp/cuda-oxide/tests/shared-owner.rs"), - line: 24, - column: 5, - }, - is_local_to_unit: true, - is_function_local: true, - }, - ); - // This raw spelling normalizes to `shared_kernel`, but it is not the raw - // symbol indexed for that function definition. It must not borrow the - // valid sibling's pre-reserved DISubprogram. - llvm_export::ops::set_debug_global_owner_function( - &mut ctx, - malformed.get_operation(), - "shared_kernel", - ); - - if malformed_first { - malformed.get_operation().insert_at_back(module_block, &ctx); - valid.get_operation().insert_at_back(module_block, &ctx); - } else { - valid.get_operation().insert_at_back(module_block, &ctx); - malformed.get_operation().insert_at_back(module_block, &ctx); - } - - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); - let function = FuncOp::new(&mut ctx, raw_owner.as_str().try_into().unwrap(), func_ty); - let function_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/shared-owner.rs", 16, 1); - function - .get_operation() - .deref_mut(&ctx) - .set_loc(function_loc); - let entry = function.get_or_create_entry_block(&mut ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - function.get_operation().insert_at_back(module_block, &ctx); - - export_module_to_string_with_config( - &ctx, - &module, - &DebugConfig { - inner: PtxExportConfig, - debug_kind: DebugKind::Full, - }, - ) - .expect("a malformed shared owner must not break its valid sibling") -} - -fn assert_malformed_shared_owner_fails_closed(full: &str) { - let valid_definition = full - .lines() - .find(|line| line.starts_with("@__shared_mem_valid = ")) - .expect("valid shared definition"); - let malformed_definition = full - .lines() - .find(|line| line.starts_with("@__shared_mem_malformed = ")) - .expect("malformed-owner shared definition"); - assert!( - valid_definition.contains(", !dbg !"), - "the valid sibling must retain its debug attachment:\n{full}" - ); - assert!( - !malformed_definition.contains(", !dbg !"), - "the malformed owner must not borrow the valid sibling's scope:\n{full}" - ); - assert!( - full.contains("distinct !DIGlobalVariable(name: \"VALID\"") - && !full.contains("distinct !DIGlobalVariable(name: \"MALFORMED\""), - "only the valid shared source identity may become a DIE:\n{full}" - ); - assert_eq!( - full.matches("!DIGlobalVariableExpression(var:").count(), - 1, - "only the valid sibling may be retained by the compile unit:\n{full}" - ); - assert_eq!( - full.matches("distinct !DISubprogram(name: \"shared_kernel\"") - .count(), - 1, - "the valid owner must still reuse exactly one DISubprogram:\n{full}" - ); -} - -#[test] -fn malformed_shared_owner_before_valid_sibling_cannot_borrow_its_scope() { - assert_malformed_shared_owner_fails_closed(&export_mixed_shared_owner_order(true)); -} - -#[test] -fn malformed_shared_owner_after_valid_sibling_cannot_borrow_its_scope() { - assert_malformed_shared_owner_fails_closed(&export_mixed_shared_owner_order(false)); -} - -#[test] -fn full_debug_globals_preserve_qualified_identity_visibility_and_relocations() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "global_debug_adversarial".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signless); - let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); - - let left_ty = ArrayType::get(&ctx, i8_ty.into(), 4); - let left = GlobalOp::new_with_alignment( - &mut ctx, - "__device_global_20".try_into().unwrap(), - left_ty.into(), - 4, - ); - left.set_address_space(&mut ctx, llvm_export::types::address_space::GLOBAL); - left.set_initializer_hex(&mut ctx, "01000000"); - llvm_export::ops::set_debug_global_variable( - &mut ctx, - left.get_operation(), - &DebugGlobalVariableInfo { - name: "SAME_LEAF".to_string(), - namespace: vec!["qualified_fixture".to_string(), "left".to_string()], - ty: DebugLocalTypeKind::Basic { - name: "u32".to_string(), - size_bits: 32, - encoding: "DW_ATE_unsigned", - }, - declaration: DebugSourcePosition { - file: PathBuf::from("/tmp/cuda-oxide/tests/qualified.rs"), - line: 20, - column: 5, - }, - is_local_to_unit: true, - is_function_local: false, - }, - ); - left.get_operation().insert_at_back(module_block, &ctx); - - let right_ty = ArrayType::get(&ctx, i8_ty.into(), 8); - let right = GlobalOp::new_with_alignment( - &mut ctx, - "__device_global_21".try_into().unwrap(), - right_ty.into(), - 8, - ); - right.set_address_space(&mut ctx, llvm_export::types::address_space::GLOBAL); - right.set_initializer_hex(&mut ctx, "0200000000000000"); - llvm_export::ops::set_debug_global_variable( - &mut ctx, - right.get_operation(), - &DebugGlobalVariableInfo { - name: "SAME_LEAF".to_string(), - namespace: vec!["qualified_fixture".to_string(), "right".to_string()], - ty: DebugLocalTypeKind::Basic { - name: "u64".to_string(), - size_bits: 64, - encoding: "DW_ATE_unsigned", - }, - declaration: DebugSourcePosition { - file: PathBuf::from("/tmp/cuda-oxide/tests/qualified.rs"), - line: 24, - column: 5, - }, - is_local_to_unit: false, - is_function_local: false, - }, - ); - right.get_operation().insert_at_back(module_block, &ctx); - - let target_ty = ArrayType::get(&ctx, i8_ty.into(), 4); - let target = GlobalOp::new_with_alignment( - &mut ctx, - "__device_global_22".try_into().unwrap(), - target_ty.into(), - 4, - ); - target.set_address_space(&mut ctx, llvm_export::types::address_space::GLOBAL); - target.set_source_global_key(&mut ctx, "qualified_fixture::RELOCATION_TARGET"); - target.set_initializer_hex(&mut ctx, "78563412"); - llvm_export::ops::set_debug_global_variable( - &mut ctx, - target.get_operation(), - &DebugGlobalVariableInfo { - name: "RELOCATION_TARGET".to_string(), - namespace: vec!["qualified_fixture".to_string()], - ty: DebugLocalTypeKind::Basic { - name: "u32".to_string(), - size_bits: 32, - encoding: "DW_ATE_unsigned", - }, - declaration: DebugSourcePosition { - file: PathBuf::from("/tmp/cuda-oxide/tests/qualified.rs"), - line: 30, - column: 1, - }, - is_local_to_unit: true, - is_function_local: false, - }, - ); - target.get_operation().insert_at_back(module_block, &ctx); - - let references_ty = StructType::get_unnamed( - &ctx, - (vec![i64_ty.into(), i64_ty.into()], StructLayout::Unpacked), - ); - let references = GlobalOp::new_with_alignment( - &mut ctx, - "__device_global_23".try_into().unwrap(), - references_ty.into(), - 8, - ); - references.set_address_space(&mut ctx, llvm_export::types::address_space::GLOBAL); - references.set_source_global_key(&mut ctx, "qualified_fixture::REFERENCES"); - references.set_initializer_hex(&mut ctx, "00000000000000000000000000000000"); - references.set_initializer_relocations( - &mut ctx, - &encode_global_initializer_relocations(&[ - GlobalInitializerRelocation { - source_offset: 0, - width_bytes: 8, - target_address_space: llvm_export::types::address_space::GLOBAL, - target_addend: 0, - target_key: "qualified_fixture::RELOCATION_TARGET".to_string(), - }, - GlobalInitializerRelocation { - source_offset: 8, - width_bytes: 8, - target_address_space: llvm_export::types::address_space::GLOBAL, - target_addend: 0, - target_key: "qualified_fixture::RELOCATION_TARGET".to_string(), - }, - ]), - ); - llvm_export::ops::set_debug_global_variable( - &mut ctx, - references.get_operation(), - &DebugGlobalVariableInfo { - name: "REFERENCES".to_string(), - namespace: vec!["qualified_fixture".to_string()], - ty: DebugLocalTypeKind::Array { - name: "[&u32; 2]".to_string(), - size_bits: 128, - element: Box::new(DebugLocalTypeKind::Pointer { - name: "&u32".to_string(), - size_bits: 64, - }), - count: 2, - }, - declaration: DebugSourcePosition { - file: PathBuf::from("/tmp/cuda-oxide/tests/qualified.rs"), - line: 31, - column: 1, - }, - is_local_to_unit: true, - is_function_local: false, - }, - ); - references - .get_operation() - .insert_at_back(module_block, &ctx); - - let export = |debug_kind| { - export_module_to_string_with_config( - &ctx, - &module, - &DebugConfig { - inner: PtxExportConfig, - debug_kind, - }, - ) - .expect("global debug export succeeds") - }; - let full = export(DebugKind::Full); - - let metadata_id = |needle: &str| { - full.lines() - .find(|line| line.contains(needle)) - .unwrap_or_else(|| panic!("missing `{needle}` in:\n{full}")) - .split_once(" = ") - .expect("metadata definition") - .0 - .trim_start_matches('!') - .to_string() - }; - let left_scope = metadata_id("!DINamespace(name: \"left\""); - let right_scope = metadata_id("!DINamespace(name: \"right\""); - let left_variable = full - .lines() - .find(|line| { - line.contains("!DIGlobalVariable(name: \"SAME_LEAF\"") - && line.contains("linkageName: \"__device_global_20\"") - }) - .expect("left source global"); - let right_variable = full - .lines() - .find(|line| { - line.contains("!DIGlobalVariable(name: \"SAME_LEAF\"") - && line.contains("linkageName: \"__device_global_21\"") - }) - .expect("right source global"); - assert!(left_variable.contains(&format!("scope: !{left_scope}"))); - assert!(left_variable.contains("line: 20") && left_variable.contains("isLocal: true")); - assert!(right_variable.contains(&format!("scope: !{right_scope}"))); - assert!(right_variable.contains("line: 24") && right_variable.contains("isLocal: false")); - - let target_definition = full - .lines() - .find(|line| line.starts_with("@__device_global_22 = ")) - .expect("relocation-only target definition"); - assert!(target_definition.contains("[4 x i8]") && target_definition.contains("!dbg !")); - let target_variable = full - .lines() - .find(|line| line.contains("!DIGlobalVariable(name: \"RELOCATION_TARGET\"")) - .expect("relocation-only target source identity"); - assert!( - target_variable.contains("linkageName: \"__device_global_22\"") - && target_variable.contains("line: 30") - ); - - let reference_definition = full - .lines() - .find(|line| line.starts_with("@__device_global_23 = ")) - .expect("relocation-backed definition"); - assert_eq!( - reference_definition.matches("@__device_global_22").count(), - 2, - "both initializer references must target the same one physical global" - ); - assert!( - full.contains("!DICompositeType(tag: DW_TAG_array_type, baseType: !") - && full.contains("size: 128, elements: !") - && full.contains("!DIDerivedType(tag: DW_TAG_pointer_type, name: \"&u32\"") - ); - - assert_eq!( - full.matches("!DIGlobalVariableExpression(var:").count(), - 4, - "one expression per physical source global, despite repeated relocations" - ); - let compile_unit = full - .lines() - .find(|line| line.contains("!DICompileUnit(")) - .expect("debug compile unit"); - let globals_id = compile_unit - .split("globals: !") - .nth(1) - .and_then(|tail| tail.strip_suffix(')')) - .expect("compile-unit globals tuple"); - let globals_tuple = full - .lines() - .find(|line| line.starts_with(&format!("!{globals_id} = !{{"))) - .expect("compile-unit globals tuple definition"); - assert_eq!(globals_tuple.matches('!').count() - 2, 4, "{globals_tuple}"); - - assert_eq!( - export(DebugKind::LineTables), - export(DebugKind::Off), - "initialized and relocation-backed globals must not gain metadata outside Full debug" - ); - - let Some(llvm_as) = ["llvm-as-22", "llvm-as-21", "llvm-as"] - .into_iter() - .find(|tool| { - std::process::Command::new(tool) - .arg("--version") - .output() - .is_ok_and(|out| out.status.success()) - }) - else { - eprintln!("skipping global-debug parse gate: no llvm-as on PATH"); - return; - }; - let ll_path = std::env::temp_dir().join(format!( - "cuda_oxide_global_debug_parse_gate_{}.ll", - std::process::id() - )); - std::fs::write(&ll_path, &full).expect("write temp .ll"); - let output = std::process::Command::new(llvm_as) - .arg("-o") - .arg("/dev/null") - .arg(&ll_path) - .output() - .expect("run llvm-as"); - let stderr = String::from_utf8_lossy(&output.stderr).into_owned(); - let _ = std::fs::remove_file(&ll_path); - assert!( - output.status.success(), - "{llvm_as} rejected global debug metadata:\n{stderr}\n--- module ---\n{full}" - ); -} - -#[test] -fn full_debug_metadata_emits_rust_enum_variant_parts() { - let mut ctx = Context::new(); - - let module = ModuleOp::new(&mut ctx, "enum_debug".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); - let func = FuncOp::new(&mut ctx, "enum_debug_kernel".try_into().unwrap(), func_ty); - let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/enum.rs", 10, 1); - func.get_operation().deref_mut(&ctx).set_loc(func_loc); - let entry = func.get_or_create_entry_block(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); - let one_attr = IntegerAttr::new(i32_ty, APInt::from_u32(1, NonZero::new(32).unwrap())); - let one = ConstantOp::new(&mut ctx, one_attr.into()); - one.get_operation().insert_at_back(entry, &ctx); - let one_val = one.get_operation().deref(&ctx).get_result(0); - - let direct = AllocaOp::new(&mut ctx, i64_ty.into(), one_val); - let direct_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/enum.rs", 11, 9); - direct.get_operation().deref_mut(&ctx).set_loc(direct_loc); - llvm_export::ops::set_debug_local_variable( - &mut ctx, - direct.get_operation(), - DebugLocalVariableInfo { - name: "direct".to_string(), - argument_index: None, - ty: DebugLocalTypeKind::Enum { - name: "Direct".to_string(), - size_bits: 64, - discriminant: Some(DebugEnumDiscriminant { - offset_bits: 0, - ty: Box::new(DebugLocalTypeKind::Basic { - name: "u8".to_string(), - size_bits: 8, - encoding: "DW_ATE_unsigned", - }), - }), - variants: vec![ - DebugEnumVariant { - name: "Small".to_string(), - discriminant: Some(3), - members: vec![llvm_export::ops::DebugTypeMember { - name: "0".to_string(), - offset_bits: 32, - ty: DebugLocalTypeKind::Basic { - name: "u32".to_string(), - size_bits: 32, - encoding: "DW_ATE_unsigned", - }, - }], - }, - DebugEnumVariant { - name: "Empty".to_string(), - discriminant: Some(9), - members: vec![], - }, - ], - }, - }, - ); - direct.get_operation().insert_at_back(entry, &ctx); - - let niche = AllocaOp::new(&mut ctx, i64_ty.into(), one_val); - let niche_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/enum.rs", 12, 9); - niche.get_operation().deref_mut(&ctx).set_loc(niche_loc); - llvm_export::ops::set_debug_local_variable( - &mut ctx, - niche.get_operation(), - DebugLocalVariableInfo { - name: "niche".to_string(), - argument_index: None, - ty: DebugLocalTypeKind::Enum { - name: "OptionRef".to_string(), - size_bits: 64, - discriminant: Some(DebugEnumDiscriminant { - offset_bits: 0, - ty: Box::new(DebugLocalTypeKind::Basic { - name: "usize".to_string(), - size_bits: 64, - encoding: "DW_ATE_unsigned", - }), - }), - variants: vec![ - DebugEnumVariant { - name: "None".to_string(), - discriminant: Some(0), - members: vec![], - }, - DebugEnumVariant { - name: "Some".to_string(), - discriminant: None, - members: vec![llvm_export::ops::DebugTypeMember { - name: "0".to_string(), - offset_bits: 0, - ty: DebugLocalTypeKind::TypedPointer { - name: "&u32".to_string(), - size_bits: 64, - pointee: Box::new(DebugLocalTypeKind::Basic { - name: "u32".to_string(), - size_bits: 32, - encoding: "DW_ATE_unsigned", - }), - }, - }], - }, - ], - }, - }, - ); - niche.get_operation().insert_at_back(entry, &ctx); - - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - let config = DebugConfig { - inner: PtxExportConfig, - debug_kind: DebugKind::Full, - }; - let ir = - export_module_to_string_with_config(&ctx, &module, &config).expect("enum debug export"); - - assert!( - ir.contains("!DICompositeType(tag: DW_TAG_variant_part"), - "Rust enums must contain a DW_TAG_variant_part:\n{ir}" - ); - assert!( - ir.contains("flags: DIFlagArtificial"), - "the physical enum carrier must be marked artificial:\n{ir}" - ); - let discriminator_member = ir - .lines() - .find(|line| { - line.contains("!DIDerivedType(tag: DW_TAG_member") - && line.contains("flags: DIFlagArtificial") - }) - .expect("enum discriminator member"); - assert!( - discriminator_member.contains("scope: !"), - "the physical enum carrier must be scoped to its enum object:\n{discriminator_member}\n{ir}" - ); - let variant_part = ir - .lines() - .find(|line| line.contains("!DICompositeType(tag: DW_TAG_variant_part")) - .expect("enum variant part"); - assert!( - variant_part.contains("scope: !"), - "the variant part must be scoped to its enum object:\n{variant_part}\n{ir}" - ); - let direct_variant_member = ir - .lines() - .find(|line| { - line.contains("!DIDerivedType(tag: DW_TAG_member, name: \"Small\"") - && line.contains("extraData: i8 3") - }) - .expect("direct enum variant member"); - assert!( - direct_variant_member.contains("scope: !"), - "variant members must be scoped to their DW_TAG_variant_part:\n{direct_variant_member}\n{ir}" - ); - assert!( - ir.contains("extraData: i8 3") && ir.contains("extraData: i8 9"), - "direct-tag variants must carry their physical discriminant values:\n{ir}" - ); - assert!( - ir.contains("extraData: i64 0"), - "the tagged niche variant must carry the niche value:\n{ir}" - ); - - let some_variant_member = ir - .lines() - .find(|line| { - line.contains("!DIDerivedType(tag: DW_TAG_member, name: \"Some\"") - && line.contains("baseType:") - }) - .expect("Some variant member"); - assert!( - !some_variant_member.contains("extraData:"), - "the untagged niche variant must be the default branch:\n{some_variant_member}\n{ir}" - ); -} - -#[test] -fn full_debug_metadata_uses_file_scope_for_cross_file_local_variables() { - let mut ctx = Context::new(); - - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_region = module.get_operation().deref(&ctx).get_region(0); - let module_block = { - let region = module_region.deref(&ctx); - region.iter(&ctx).next().unwrap() - }; - - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); - let func = FuncOp::new(&mut ctx, "debug_kernel".try_into().unwrap(), func_ty); - let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 30, 1); - func.get_operation().deref_mut(&ctx).set_loc(func_loc); - - let entry = func.get_or_create_entry_block(&mut ctx); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let one_attr = IntegerAttr::new(i32_ty, APInt::from_u32(1, NonZero::new(32).unwrap())); - let one = ConstantOp::new(&mut ctx, one_attr.into()); - one.get_operation().insert_at_back(entry, &ctx); - let one_val = one.get_operation().deref(&ctx).get_result(0); - - let tid = AllocaOp::new(&mut ctx, i32_ty.into(), one_val); - let tid_loc = src_location( - &mut ctx, - "/tmp/cuda-oxide/crates/cuda-device/src/thread.rs", - 292, - 19, - ); - tid.get_operation().deref_mut(&ctx).set_loc(tid_loc); - llvm_export::ops::set_debug_local_variable( - &mut ctx, - tid.get_operation(), - DebugLocalVariableInfo { - name: "tid".to_string(), - argument_index: None, - ty: DebugLocalTypeKind::Basic { - name: "u32".to_string(), - size_bits: 32, - encoding: "DW_ATE_unsigned", - }, - }, - ); - tid.get_operation().insert_at_back(entry, &ctx); - - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - let config = DebugConfig { - inner: PtxExportConfig, - debug_kind: DebugKind::Full, - }; - let ir = - export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); - - assert!( - ir.contains("!DILexicalBlockFile(scope: !"), - "cross-file local variables should get a file-specific debug scope:\n{ir}" - ); - assert!( - ir.contains( - "!DIFile(filename: \"thread.rs\", directory: \"/tmp/cuda-oxide/crates/cuda-device/src\")" - ), - "cross-file local variables should reference their source file:\n{ir}" - ); - assert!( - ir.contains("!DILocalVariable(name: \"tid\", scope: !") && ir.contains("line: 292"), - "cross-file local variables should preserve the variable file scope and line:\n{ir}" - ); - assert!( - ir.contains("call void @llvm.dbg.declare"), - "cross-file local variables should still get dbg.declare bindings:\n{ir}" - ); -} - -#[test] -fn full_debug_metadata_emits_dbg_value_for_promoted_locals() { - let mut ctx = Context::new(); - - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_region = module.get_operation().deref(&ctx).get_region(0); - let module_block = { - let region = module_region.deref(&ctx); - region.iter(&ctx).next().unwrap() - }; - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![i32_ty.into()], false); - let func = FuncOp::new(&mut ctx, "debug_kernel".try_into().unwrap(), func_ty); - let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 30, 1); - func.get_operation().deref_mut(&ctx).set_loc(func_loc); - - let entry = func.get_or_create_entry_block(&mut ctx); - let arg = entry.deref(&ctx).get_argument(0); - let dbg_value = DebugValueOp::new(&mut ctx, arg); - let dbg_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 31, 13); - dbg_value.get_operation().deref_mut(&ctx).set_loc(dbg_loc); - llvm_export::ops::set_debug_local_variable( - &mut ctx, - dbg_value.get_operation(), - DebugLocalVariableInfo { - name: "x".to_string(), - argument_index: Some(1), - ty: DebugLocalTypeKind::Basic { - name: "i32".to_string(), - size_bits: 32, - encoding: "DW_ATE_signed", - }, - }, - ); - llvm_export::ops::set_debug_local_declaration_location( - &mut ctx, - dbg_value.get_operation(), - PathBuf::from("/tmp/cuda-oxide/tests/declarations.rs"), - 12, - 5, - ); - llvm_export::ops::set_debug_fragment_variables( - &mut ctx, - dbg_value.get_operation(), - &[DebugFragmentVariableInfo { - variable: DebugLocalVariableInfo { - name: "pair".to_string(), - argument_index: None, - ty: DebugLocalTypeKind::Array { - name: "[u32; 2]".to_string(), - size_bits: 64, - element: Box::new(DebugLocalTypeKind::Basic { - name: "u32".to_string(), - size_bits: 32, - encoding: "DW_ATE_unsigned", - }), - count: 2, - }, - }, - fragment: DebugFragment { - offset_bits: 32, - size_bits: 32, - }, - source_scope: None, - declaration: Some(DebugSourcePosition { - file: PathBuf::from("/tmp/cuda-oxide/tests/declarations.rs"), - line: 13, - column: 5, - }), - }], - ); - dbg_value.get_operation().insert_at_back(entry, &ctx); - - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - let config = DebugConfig { - inner: PtxExportConfig, - debug_kind: DebugKind::Full, - }; - let ir = - export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); - - assert!( - ir.contains("declare void @llvm.dbg.value(metadata, metadata, metadata)"), - "full debug should declare dbg.value when it emits one:\n{ir}" - ); - assert!( - ir.contains("call void @llvm.dbg.value(metadata i32 %v0, metadata !"), - "dbg.value should describe the local as the current SSA value:\n{ir}" - ); - assert!( - ir.contains("!DILocalVariable(name: \"x\", arg: 1, scope: !"), - "dbg.value should preserve formal-argument metadata when the source local is an argument:\n{ir}" - ); - assert!( - ir.contains("!DILocalVariable(name: \"x\", arg: 1, scope: !") - && ir.contains("file: !") - && ir.contains("line: 12"), - "DILocalVariable should use the source declaration line, not the dbg.value line:\n{ir}" - ); - assert!( - ir.contains("!DILocation(line: 31, column: 13, scope: !"), - "dbg.value should still be located at the value's current source point:\n{ir}" - ); - assert!( - ir.contains("DW_OP_LLVM_fragment, 32, 32"), - "dbg.value should preserve scalarized source-variable fragments:\n{ir}" - ); - assert!( - ir.contains("!DILocalVariable(name: \"pair\", scope: !") && ir.contains("line: 13"), - "fragment dbg.value should describe the complete source variable:\n{ir}" - ); - assert!( - !ir.contains("llvm.dbg.declare"), - "a value-only debug record should not force dbg.declare:\n{ir}" - ); -} - -#[test] -fn full_debug_metadata_emits_diarglist_for_multi_value_locations() { - let mut ctx = Context::new(); - - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_region = module.get_operation().deref(&ctx).get_region(0); - let module_block = { - let region = module_region.deref(&ctx); - region.iter(&ctx).next().unwrap() - }; - - let ptr_ty = PointerType::get(&ctx, 0); - let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get( - &ctx, - void_ty.to_handle(), - vec![ptr_ty.into(), i64_ty.into()], - false, - ); - let func = FuncOp::new(&mut ctx, "debug_kernel".try_into().unwrap(), func_ty); - let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 40, 1); - func.get_operation().deref_mut(&ctx).set_loc(func_loc); - - let entry = func.get_or_create_entry_block(&mut ctx); - let base = entry.deref(&ctx).get_argument(0); - let index = entry.deref(&ctx).get_argument(1); - let dbg_value = DebugValueListOp::new(&mut ctx, vec![base, index]); - let dbg_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 41, 17); - dbg_value.get_operation().deref_mut(&ctx).set_loc(dbg_loc); - llvm_export::ops::set_debug_local_variable( - &mut ctx, - dbg_value.get_operation(), - DebugLocalVariableInfo { - name: "item".to_string(), - argument_index: None, - ty: DebugLocalTypeKind::Basic { - name: "u32".to_string(), - size_bits: 32, - encoding: "DW_ATE_unsigned", - }, - }, - ); - llvm_export::ops::set_debug_value_expression( - &mut ctx, - dbg_value.get_operation(), - &DebugValueExpression::new(vec![ - DebugValueExpressionOp::Arg(0), - DebugValueExpressionOp::Arg(1), - DebugValueExpressionOp::ConstU(4), - DebugValueExpressionOp::Mul, - DebugValueExpressionOp::Plus, - DebugValueExpressionOp::Deref, - ]), - ); - dbg_value.get_operation().insert_at_back(entry, &ctx); - - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - let config = DebugConfig { - inner: PtxExportConfig, - debug_kind: DebugKind::Full, - }; - let ir = - export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); - - assert!( - ir.contains("call void @llvm.dbg.value(metadata !DIArgList(ptr %v0, i64 %v1), metadata !"), - "multi-value debug records should emit an inline DIArgList:\n{ir}" - ); - assert!( - ir.contains( - "!DIExpression(DW_OP_LLVM_arg, 0, DW_OP_LLVM_arg, 1, DW_OP_constu, 4, DW_OP_mul, DW_OP_plus, DW_OP_deref)" - ), - "multi-value debug records should emit the typed location recipe:\n{ir}" - ); - assert!( - ir.contains("declare void @llvm.dbg.value(metadata, metadata, metadata)"), - "multi-value records should reuse the ordinary dbg.value intrinsic declaration:\n{ir}" - ); - - let Some(llvm_as) = ["llvm-as-22", "llvm-as-21", "llvm-as"] - .into_iter() - .find(|tool| { - std::process::Command::new(tool) - .arg("--version") - .output() - .is_ok_and(|out| out.status.success()) - }) - else { - eprintln!("skipping llvm-as parse gate: no llvm-as-22/llvm-as-21/llvm-as on PATH"); - return; - }; - - let ll_path = std::env::temp_dir().join(format!( - "cuda_oxide_diarglist_parse_gate_{}.ll", - std::process::id() - )); - std::fs::write(&ll_path, &ir).expect("write temp .ll"); - let output = std::process::Command::new(llvm_as) - .arg("-o") - .arg("/dev/null") - .arg(&ll_path) - .output() - .expect("run llvm-as"); - let stderr = String::from_utf8_lossy(&output.stderr).into_owned(); - let _ = std::fs::remove_file(&ll_path); - assert!( - output.status.success(), - "{llvm_as} rejected the emitted multi-value debug module:\n{stderr}\n--- module ---\n{ir}" - ); -} - -#[test] -fn full_debug_metadata_uses_inlined_callee_scope_for_inlined_arguments() { - let mut ctx = Context::new(); - - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_region = module.get_operation().deref(&ctx).get_region(0); - let module_block = { - let region = module_region.deref(&ctx); - region.iter(&ctx).next().unwrap() - }; - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![i32_ty.into()], false); - let func = FuncOp::new(&mut ctx, "caller_kernel".try_into().unwrap(), func_ty); - let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 30, 1); - func.get_operation().deref_mut(&ctx).set_loc(func_loc); - llvm_export::ops::set_debug_source_scope_map( - &mut ctx, - func.get_operation(), - &DebugSourceScopeMap { - scopes: vec![ - DebugSourceScope { - id: 0, - parent: None, - span: Some(DebugSourcePosition { - file: PathBuf::from("/tmp/cuda-oxide/tests/kernel.rs"), - line: 30, - column: 1, - }), - inlined: None, - }, - DebugSourceScope { - id: 1, - parent: Some(0), - span: Some(DebugSourcePosition { - file: PathBuf::from("/tmp/cuda-oxide/tests/helper.rs"), - line: 7, - column: 1, - }), - inlined: Some(llvm_export::ops::DebugInlinedScope { - callee_name: "helper::next".to_string(), - callsite: Some(DebugSourcePosition { - file: PathBuf::from("/tmp/cuda-oxide/tests/kernel.rs"), - line: 41, - column: 13, - }), - }), - }, - ], - locations: vec![], - }, - ); - - let entry = func.get_or_create_entry_block(&mut ctx); - let arg = entry.deref(&ctx).get_argument(0); - - let caller_value = DebugValueOp::new(&mut ctx, arg); - let caller_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 31, 9); - caller_value - .get_operation() - .deref_mut(&ctx) - .set_loc(caller_loc); - llvm_export::ops::set_debug_local_variable( - &mut ctx, - caller_value.get_operation(), - DebugLocalVariableInfo { - name: "data".to_string(), - argument_index: Some(1), - ty: DebugLocalTypeKind::Basic { - name: "i32".to_string(), - size_bits: 32, - encoding: "DW_ATE_signed", - }, - }, - ); - llvm_export::ops::set_debug_local_source_scope(&mut ctx, caller_value.get_operation(), 0); - caller_value.get_operation().insert_at_back(entry, &ctx); - - let inlined_value = DebugValueOp::new(&mut ctx, arg); - let inlined_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/helper.rs", 8, 17); - inlined_value - .get_operation() - .deref_mut(&ctx) - .set_loc(inlined_loc); - llvm_export::ops::set_debug_local_variable( - &mut ctx, - inlined_value.get_operation(), - DebugLocalVariableInfo { - name: "self".to_string(), - argument_index: Some(1), - ty: DebugLocalTypeKind::Basic { - name: "i32".to_string(), - size_bits: 32, - encoding: "DW_ATE_signed", - }, - }, - ); - llvm_export::ops::set_debug_local_source_scope(&mut ctx, inlined_value.get_operation(), 1); - inlined_value.get_operation().insert_at_back(entry, &ctx); - - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - let config = DebugConfig { - inner: PtxExportConfig, - debug_kind: DebugKind::Full, - }; - let ir = - export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); - - assert!( - ir.contains("distinct !DISubprogram(name: \"caller_kernel\""), - "caller should keep its own function debug scope:\n{ir}" - ); - assert!( - ir.contains("distinct !DISubprogram(name: \"helper::next\""), - "inlined callee should get its own DISubprogram scope:\n{ir}" - ); - assert!( - ir.contains("!DILocalVariable(name: \"data\", arg: 1, scope: !"), - "caller argument should remain arg #1 in the caller scope:\n{ir}" - ); - assert!( - ir.contains("!DILocalVariable(name: \"self\", arg: 1, scope: !"), - "inlined callee argument should remain arg #1 in the callee scope:\n{ir}" - ); - assert!( - ir.contains("!DILocation(line: 8, column: 17, scope: !") && ir.contains("inlinedAt: !"), - "inlined dbg.value location should point at the callee line and caller callsite:\n{ir}" - ); -} - -#[test] -fn line_table_debug_metadata_ignores_tagged_alloca_variables() { - let mut ctx = Context::new(); - - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_region = module.get_operation().deref(&ctx).get_region(0); - let module_block = { - let region = module_region.deref(&ctx); - region.iter(&ctx).next().unwrap() - }; - - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); - let func = FuncOp::new(&mut ctx, "debug_kernel".try_into().unwrap(), func_ty); - let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 40, 1); - func.get_operation().deref_mut(&ctx).set_loc(func_loc); - - let entry = func.get_or_create_entry_block(&mut ctx); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let one_attr = IntegerAttr::new(i32_ty, APInt::from_u32(1, NonZero::new(32).unwrap())); - let one = ConstantOp::new(&mut ctx, one_attr.into()); - one.get_operation().insert_at_back(entry, &ctx); - let one_val = one.get_operation().deref(&ctx).get_result(0); - - let local = AllocaOp::new(&mut ctx, i32_ty.into(), one_val); - let local_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 41, 9); - local.get_operation().deref_mut(&ctx).set_loc(local_loc); - llvm_export::ops::set_debug_local_variable( - &mut ctx, - local.get_operation(), - DebugLocalVariableInfo { - name: "x".to_string(), - argument_index: None, - ty: DebugLocalTypeKind::Basic { - name: "i32".to_string(), - size_bits: 32, - encoding: "DW_ATE_signed", - }, - }, - ); - local.get_operation().insert_at_back(entry, &ctx); - - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - let config = DebugConfig { - inner: PtxExportConfig, - debug_kind: DebugKind::LineTables, - }; - let ir = - export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); - - assert!( - ir.contains("emissionKind: LineTablesOnly"), - "line-table mode should stay line-table-only:\n{ir}" - ); - assert!( - !ir.contains("llvm.dbg.declare"), - "line-table mode should not emit variable bindings:\n{ir}" - ); - assert!( - !ir.contains("DILocalVariable"), - "line-table mode should not emit local-variable metadata:\n{ir}" - ); -} - -#[test] -fn line_table_debug_metadata_adds_fallback_locations_to_calls() { - let mut ctx = Context::new(); - - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_region = module.get_operation().deref(&ctx).get_region(0); - let module_block = { - let region = module_region.deref(&ctx); - region.iter(&ctx).next().unwrap() - }; - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let void_ty = VoidType::get(&ctx); - let helper_ty = FuncType::get(&ctx, i32_ty.to_handle(), vec![], false); - let helper = FuncOp::new(&mut ctx, "helper".try_into().unwrap(), helper_ty); - helper.get_operation().insert_at_back(module_block, &ctx); - - let caller_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); - let caller = FuncOp::new(&mut ctx, "debug_kernel".try_into().unwrap(), caller_ty); - let caller_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 20, 3); - caller.get_operation().deref_mut(&ctx).set_loc(caller_loc); - - let entry = caller.get_or_create_entry_block(&mut ctx); - let call = CallOp::new( - &mut ctx, - CallOpCallable::Direct("helper".try_into().unwrap()), - helper_ty, - vec![], - ); - call.get_operation().insert_at_back(entry, &ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - - caller.get_operation().insert_at_back(module_block, &ctx); - - let config = DebugConfig { - inner: PtxExportConfig, - debug_kind: DebugKind::LineTables, - }; - let ir = - export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); - - let call_line = ir - .lines() - .find(|line| line.contains("call i32 @helper()")) - .expect("call instruction"); - assert!( - call_line.contains(", !dbg !"), - "calls without their own source span should use the function fallback location:\n{ir}" - ); - assert!( - ir.contains("!DILocation(line: 20, column: 3, scope: !"), - "fallback call location should point at the caller's function line:\n{ir}" - ); -} - -#[test] -fn line_table_debug_metadata_emits_explicit_artificial_calls_at_line_zero() { - let mut ctx = Context::new(); - - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_region = module.get_operation().deref(&ctx).get_region(0); - let module_block = { - let region = module_region.deref(&ctx); - region.iter(&ctx).next().unwrap() - }; - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let void_ty = VoidType::get(&ctx); - let helper_ty = FuncType::get(&ctx, i32_ty.to_handle(), vec![], false); - let helper = FuncOp::new(&mut ctx, "helper".try_into().unwrap(), helper_ty); - helper.get_operation().insert_at_back(module_block, &ctx); - - let caller_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); - let caller = FuncOp::new(&mut ctx, "debug_kernel".try_into().unwrap(), caller_ty); - let caller_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 20, 3); - caller.get_operation().deref_mut(&ctx).set_loc(caller_loc); - - let entry = caller.get_or_create_entry_block(&mut ctx); - let call = CallOp::new( - &mut ctx, - CallOpCallable::Direct("helper".try_into().unwrap()), - helper_ty, - vec![], - ); - call.get_operation() - .deref_mut(&ctx) - .set_loc(llvm_export::artificial_debug_location()); - call.get_operation().insert_at_back(entry, &ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - - caller.get_operation().insert_at_back(module_block, &ctx); - - let config = DebugConfig { - inner: PtxExportConfig, - debug_kind: DebugKind::LineTables, - }; - let ir = - export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); - - let call_line = ir - .lines() - .find(|line| line.contains("call i32 @helper()")) - .expect("call instruction"); - assert!( - call_line.contains(", !dbg !"), - "LLVM still requires an artificial call location:\n{ir}" - ); - assert!( - ir.contains("!DILocation(line: 0, column: 0, scope: !"), - "artificial setup must not reuse the caller's user line:\n{ir}" - ); -} - -/// Scans the textual LLVM IR and asserts that every `%vN` token appearing in -/// an operand position has a corresponding `%vN = ...` definition somewhere -/// in the module. Operates on `%v` temporaries only because that's the -/// exporter's naming scheme; named values like `%entry` (block labels) are -/// ignored by construction. -fn assert_no_undefined_temporaries(ir: &str) { - use std::collections::HashSet; - - let mut defined: HashSet = HashSet::new(); - for line in ir.lines() { - let trimmed = line.trim_start(); - if !trimmed.starts_with("%v") { - continue; - } - let Some((lhs, _)) = trimmed.split_once('=') else { - continue; - }; - defined.insert(lhs.trim().to_string()); - } - - let mut referenced: HashSet = HashSet::new(); - for line in ir.lines() { - let trimmed = line.trim_start(); - // Skip the lhs of a definition; only operand positions can be stale. - let body = if trimmed.starts_with("%v") - && let Some(eq) = trimmed.find('=') - { - &trimmed[eq + 1..] - } else { - trimmed - }; - for tok in body.split(|c: char| !c.is_alphanumeric() && c != '%' && c != '_') { - if let Some(num) = tok.strip_prefix("%v") - && !num.is_empty() - && num.chars().all(|c| c.is_ascii_digit()) - { - referenced.insert(format!("%v{num}")); - } - } - } - - let mut undefined: Vec<&String> = referenced.difference(&defined).collect(); - undefined.sort(); - assert!( - undefined.is_empty(), - "IR references undefined SSA temporaries: {undefined:?}\nIR:\n{ir}" - ); -} - -/// A float binop carrying `FastmathFlags` must export with the matching LLVM -/// fast-math keyword (`fast` for the all-bits set), while a float binop with no -/// flags must export with none. Regression guard: the textual exporter -/// previously dropped fast-math flags entirely, making the `f*_fast` intrinsic -/// lowering inert end to end. -#[test] -fn export_emits_fast_math_flags_only_on_flagged_float_ops() { - use llvm_export::attributes::FastmathFlags; - use llvm_export::op_interfaces::{BinArithOp, FloatBinArithOpWithFastMathFlags}; - use llvm_export::ops::{FAddOp, FMulOp}; - use pliron::builtin::types::FP32Type; - - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_region = module.get_operation().deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - - let f32_ty = FP32Type::get(&ctx); - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get( - &ctx, - void_ty.to_handle(), - vec![f32_ty.into(), f32_ty.into()], - false, - ); - let func = FuncOp::new(&mut ctx, "fast_math".try_into().unwrap(), func_ty); - let entry = func.get_or_create_entry_block(&mut ctx); - let a = entry.deref(&ctx).get_argument(0); - let b = entry.deref(&ctx).get_argument(1); - - // fadd with the full fast-math set. - let fadd = FAddOp::new_with_fast_math_flags(&mut ctx, a, b, FastmathFlags::FAST.into()); - fadd.get_operation().insert_at_back(entry, &ctx); - // fmul with no flags: must stay flag-free. - let fmul = FMulOp::new(&mut ctx, a, b); - fmul.get_operation().insert_at_back(entry, &ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - let ir = export_module_to_string(&ctx, &module).expect("export succeeds"); - - assert!( - ir.contains("fadd fast float"), - "fast-math fadd must export the `fast` keyword:\n{ir}" - ); - let fmul_line = ir - .lines() - .find(|line| line.contains("fmul")) - .expect("exported fmul line"); - assert!( - !fmul_line.contains("fast"), - "a float binop with no fast-math flags must not gain them:\n{ir}" - ); -} - -#[test] -fn modern_device_extern_emits_signext_zeroext_for_small_integer_params() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "small_types_extern".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - - // The extern takes (i8 signext, i16 zeroext, i1 zeroext, half) and - // returns void. The declared IR types stay NARROW (matching cuda-oxide's - // own i8/i16/i1 SSA values and the clang-compiled LTOIR definition); the - // ABI extension attributes live only in the DeviceExternType metadata. - let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signless); - let i16_ty = IntegerType::get(&ctx, 16, Signedness::Signless); - let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); - let half_ty = HalfType::get(&ctx); - let void_ty = VoidType::get(&ctx); - - let external_ty = FuncType::get( - &ctx, - void_ty.into(), - vec![i8_ty.into(), i16_ty.into(), i1_ty.into(), half_ty.into()], - false, - ); - FuncOp::new(&mut ctx, "small_types_fn".try_into().unwrap(), external_ty) - .get_operation() - .insert_at_back(module_block, &ctx); - - // Caller that forwards its own narrow parameters to the extern with no - // value conversions. - let caller_ty = FuncType::get( - &ctx, - void_ty.into(), - vec![i8_ty.into(), i16_ty.into(), i1_ty.into(), half_ty.into()], - false, - ); - let caller = FuncOp::new(&mut ctx, "caller".try_into().unwrap(), caller_ty); - let entry = caller.get_or_create_entry_block(&mut ctx); - let arg0 = entry.deref(&ctx).get_argument(0); - let arg1 = entry.deref(&ctx).get_argument(1); - let arg2 = entry.deref(&ctx).get_argument(2); - let arg3 = entry.deref(&ctx).get_argument(3); - CallOp::new( - &mut ctx, - CallOpCallable::Direct("small_types_fn".try_into().unwrap()), - external_ty, - vec![arg0, arg1, arg2, arg3], - ) - .get_operation() - .insert_at_back(entry, &ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - caller.get_operation().insert_at_back(module_block, &ctx); - - let externs = [DeviceExternDecl { - export_name: "small_types_fn".to_string(), - param_types: vec![ - DeviceExternType::SignExtInteger(8), // i8, sign-extended by NVPTX - DeviceExternType::ZeroExtInteger(16), // u16, zero-extended by NVPTX - DeviceExternType::ZeroExtInteger(1), // bool - DeviceExternType::Float16, // f16 as native half - ], - return_type: DeviceExternType::Void, - attrs: DeviceExternAttrs::default(), - }]; - - let ir = export_module_with_externs( - &ctx, - &module, - &externs, - &NvvmExportConfig::new(NvvmIrDialect::Modern), - ) - .expect("modern extern with small types export succeeds"); - - // Declaration: narrow types with parameter-position attributes. - assert!( - ir.contains("declare void @small_types_fn(i8 signext, i16 zeroext, i1 zeroext, half)"), - "declaration should keep narrow types with signext/zeroext attributes:\n{ir}" - ); - - // Call site keeps the narrow types and attributes too. - assert!( - ir.contains("i8 signext %v0"), - "call should use signext on first arg:\n{ir}" - ); - assert!( - ir.contains("i16 zeroext %v1"), - "call should use zeroext on second arg:\n{ir}" - ); - assert!( - ir.contains("i1 zeroext %v2"), - "call should use zeroext on the bool arg:\n{ir}" - ); - assert!( - ir.contains("half %v3"), - "call should use half for the f16 arg:\n{ir}" - ); -} - -#[test] -fn modern_device_extern_signext_return_type() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "signext_return".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - - let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signless); - let void_ty = VoidType::get(&ctx); - - // Extern returns i8 with signext. The declared type stays i8; only the - // attribute marks the NVPTX widening, and in return position LLVM's - // grammar requires the attribute BEFORE the type. - let external_ty = FuncType::get(&ctx, i8_ty.into(), vec![], false); - FuncOp::new(&mut ctx, "get_small_val".try_into().unwrap(), external_ty) - .get_operation() - .insert_at_back(module_block, &ctx); - - // Caller that calls the extern and discards the result. - let caller_ty = FuncType::get(&ctx, void_ty.into(), vec![], false); - let caller = FuncOp::new(&mut ctx, "caller".try_into().unwrap(), caller_ty); - let entry = caller.get_or_create_entry_block(&mut ctx); - CallOp::new( - &mut ctx, - CallOpCallable::Direct("get_small_val".try_into().unwrap()), - external_ty, - vec![], - ) - .get_operation() - .insert_at_back(entry, &ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - caller.get_operation().insert_at_back(module_block, &ctx); - - let externs = [DeviceExternDecl { - export_name: "get_small_val".to_string(), - param_types: vec![], - return_type: DeviceExternType::SignExtInteger(8), - attrs: DeviceExternAttrs::default(), - }]; - - let ir = export_module_with_externs( - &ctx, - &module, - &externs, - &NvvmExportConfig::new(NvvmIrDialect::Modern), - ) - .expect("modern extern with signext return export succeeds"); - - // Declaration: attribute precedes the narrow return type. - assert!( - ir.contains("declare signext i8 @get_small_val()"), - "declaration should have signext BEFORE the return type:\n{ir}" - ); - - // Call site uses the same return-position placement. - assert!( - ir.contains("call signext i8 @get_small_val()"), - "call should have signext BEFORE the return type:\n{ir}" - ); -} - -#[test] -fn plain_integer_device_extern_has_no_extension_attributes() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "plain_int".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let void_ty = VoidType::get(&ctx); - - let external_ty = FuncType::get(&ctx, void_ty.into(), vec![i32_ty.into()], false); - FuncOp::new(&mut ctx, "plain_int_fn".try_into().unwrap(), external_ty) - .get_operation() - .insert_at_back(module_block, &ctx); - - let caller = FuncOp::new(&mut ctx, "caller".try_into().unwrap(), external_ty); - let entry = caller.get_or_create_entry_block(&mut ctx); - let arg0 = entry.deref(&ctx).get_argument(0); - CallOp::new( - &mut ctx, - CallOpCallable::Direct("plain_int_fn".try_into().unwrap()), - external_ty, - vec![arg0], - ) - .get_operation() - .insert_at_back(entry, &ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - caller.get_operation().insert_at_back(module_block, &ctx); - - let externs = [DeviceExternDecl { - export_name: "plain_int_fn".to_string(), - param_types: vec![DeviceExternType::Integer(32)], - return_type: DeviceExternType::Void, - attrs: DeviceExternAttrs::default(), - }]; - - let ir = export_module_with_externs( - &ctx, - &module, - &externs, - &NvvmExportConfig::new(NvvmIrDialect::Modern), - ) - .expect("plain int extern export succeeds"); - - // Plain i32 should have no signext/zeroext. - assert!( - ir.contains("declare void @plain_int_fn(i32)"), - "declaration should have plain i32 without attributes:\n{ir}" - ); - assert!( - !ir.contains("signext") && !ir.contains("zeroext"), - "plain i32 should have no extension attributes:\n{ir}" - ); -} - -#[test] -fn legacy_device_extern_rejects_small_integers_by_value() { - let mut ctx = Context::new(); - let empty = ModuleOp::new(&mut ctx, "empty".try_into().unwrap()); - let externs = [DeviceExternDecl { - export_name: "takes_small".to_string(), - param_types: vec![DeviceExternType::SignExtInteger(8)], - return_type: DeviceExternType::ZeroExtInteger(16), - attrs: DeviceExternAttrs::default(), - }]; - let err = export_module_with_externs( - &ctx, - &empty, - &externs, - &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), - ) - .expect_err("legacy sub-32-bit by-value externs must fail cleanly"); - assert!( - err.contains("CUDA 12 legacy") && err.contains("sub-32-bit"), - "{err}" - ); -} - -/// Build a module whose externs carry small params AND a small return, in -/// both declare and call positions, so the emitted attribute placement can -/// be checked against a real LLVM parser. -fn small_type_extern_module(ctx: &mut Context) -> (ModuleOp, Vec) { - let module = ModuleOp::new(ctx, "small_types_parse_gate".try_into().unwrap()); - let module_block = module_top_block(ctx, &module); - - let i8_ty = IntegerType::get(ctx, 8, Signedness::Signless); - let i16_ty = IntegerType::get(ctx, 16, Signedness::Signless); - let i1_ty = IntegerType::get(ctx, 1, Signedness::Signless); - let half_ty = HalfType::get(ctx); - let void_ty = VoidType::get(ctx); - - let take_small_ty = FuncType::get( - ctx, - void_ty.into(), - vec![i8_ty.into(), i16_ty.into(), i1_ty.into(), half_ty.into()], - false, - ); - FuncOp::new(ctx, "take_small".try_into().unwrap(), take_small_ty) - .get_operation() - .insert_at_back(module_block, &*ctx); - - let give_small_ty = FuncType::get(ctx, i8_ty.into(), vec![], false); - FuncOp::new(ctx, "give_small".try_into().unwrap(), give_small_ty) - .get_operation() - .insert_at_back(module_block, &*ctx); - - let caller_ty = FuncType::get( - ctx, - void_ty.into(), - vec![i8_ty.into(), i16_ty.into(), i1_ty.into(), half_ty.into()], - false, - ); - let caller = FuncOp::new(ctx, "caller".try_into().unwrap(), caller_ty); - let entry = caller.get_or_create_entry_block(ctx); - let arg0 = entry.deref(ctx).get_argument(0); - let arg1 = entry.deref(ctx).get_argument(1); - let arg2 = entry.deref(ctx).get_argument(2); - let arg3 = entry.deref(ctx).get_argument(3); - CallOp::new( - ctx, - CallOpCallable::Direct("take_small".try_into().unwrap()), - take_small_ty, - vec![arg0, arg1, arg2, arg3], - ) - .get_operation() - .insert_at_back(entry, &*ctx); - CallOp::new( - ctx, - CallOpCallable::Direct("give_small".try_into().unwrap()), - give_small_ty, - vec![], - ) - .get_operation() - .insert_at_back(entry, &*ctx); - ReturnOp::new(ctx, None) - .get_operation() - .insert_at_back(entry, &*ctx); - caller.get_operation().insert_at_back(module_block, &*ctx); - - let externs = vec![ - DeviceExternDecl { - export_name: "take_small".to_string(), - param_types: vec![ - DeviceExternType::SignExtInteger(8), - DeviceExternType::ZeroExtInteger(16), - DeviceExternType::ZeroExtInteger(1), - DeviceExternType::Float16, - ], - return_type: DeviceExternType::Void, - attrs: DeviceExternAttrs::default(), - }, - DeviceExternDecl { - export_name: "give_small".to_string(), - param_types: vec![], - return_type: DeviceExternType::SignExtInteger(8), - attrs: DeviceExternAttrs::default(), - }, - ]; - (module, externs) -} - -/// Parse gate: the emitted attribute placement (`i8 signext` in parameter -/// position, `signext i8` in return position, for both `declare` and `call`) -/// must be accepted by a real LLVM parser, not just string-matched. -#[test] -fn small_type_extern_module_parses_with_llvm_as() { - let mut ctx = Context::new(); - let (module, externs) = small_type_extern_module(&mut ctx); - let ir = export_module_with_externs( - &ctx, - &module, - &externs, - &NvvmExportConfig::new(NvvmIrDialect::Modern), - ) - .expect("modern small-type extern export succeeds"); - - // Belt-and-braces string checks before the external parse. - assert!( - ir.contains("declare void @take_small(i8 signext, i16 zeroext, i1 zeroext, half)"), - "{ir}" - ); - assert!(ir.contains("declare signext i8 @give_small()"), "{ir}"); - assert!(ir.contains("call signext i8 @give_small()"), "{ir}"); - - let Some(llvm_as) = ["llvm-as-22", "llvm-as-21", "llvm-as"] - .into_iter() - .find(|tool| { - std::process::Command::new(tool) - .arg("--version") - .output() - .is_ok_and(|out| out.status.success()) - }) - else { - eprintln!("skipping llvm-as parse gate: no llvm-as-22/llvm-as-21/llvm-as on PATH"); - return; - }; - - let ll_path = std::env::temp_dir().join(format!( - "cuda_oxide_small_type_parse_gate_{}.ll", - std::process::id() - )); - std::fs::write(&ll_path, &ir).expect("write temp .ll"); - let output = std::process::Command::new(llvm_as) - .arg("-o") - .arg("/dev/null") - .arg(&ll_path) - .output() - .expect("run llvm-as"); - let stderr = String::from_utf8_lossy(&output.stderr).into_owned(); - let _ = std::fs::remove_file(&ll_path); - assert!( - output.status.success(), - "{llvm_as} rejected the emitted module:\n{stderr}\n--- module ---\n{ir}" - ); -} - -/// Builds a `void` function taking pointer parameters in the given address -/// spaces, with an empty body, optionally carrying the `gpu_kernel` marker -/// the exporter uses to distinguish an entry from a device function. -fn pointer_param_function( - ctx: &mut Context, - module_block: Ptr, - name: &str, - address_spaces: &[u32], - is_kernel: bool, -) { - let params: Vec<_> = address_spaces - .iter() - .map(|space| PointerType::get(ctx, *space).into()) - .collect(); - let void_ty = VoidType::get(ctx); - let func_ty = FuncType::get(ctx, void_ty.into(), params, false); - let func = FuncOp::new(ctx, name.try_into().unwrap(), func_ty); - if is_kernel { - let kernel_key: Identifier = "gpu_kernel".try_into().unwrap(); - func.get_operation() - .deref_mut(ctx) - .attributes - .set(kernel_key, StringAttr::new("true".to_string())); - } - let entry = func.get_or_create_entry_block(ctx); - ReturnOp::new(ctx, None) - .get_operation() - .insert_at_back(entry, ctx); - func.get_operation().insert_at_back(module_block, ctx); -} - -/// A kernel receives its parameters in `.param` space from the host, which -/// holds no shared-memory address to pass. The exporter must refuse the -/// signature instead of emitting `.ptr .shared`, which ptxas accepts but the -/// driver rejects at module load, taking every kernel in the module down. -#[test] -fn kernel_rejects_a_shared_memory_pointer_parameter() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "kernel_shared_param".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - pointer_param_function(&mut ctx, module_block, "shared_param", &[0, 3], true); - - let error = export_module_to_string(&ctx, &module) - .expect_err("export must refuse a kernel parameter in shared memory"); - assert!( - error.contains( - "kernel `@shared_param` parameter 1 is a pointer into shared memory (address space 3)" - ), - "{error}" - ); -} - -/// Local memory is per-thread, so the host has no address in it to pass -/// either; the exporter refuses it on the same grounds as shared. -#[test] -fn kernel_rejects_a_local_memory_pointer_parameter() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "kernel_local_param".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - pointer_param_function(&mut ctx, module_block, "local_param", &[5], true); - - let error = export_module_to_string(&ctx, &module) - .expect_err("export must refuse a kernel parameter in local memory"); - assert!( - error.contains( - "kernel `@local_param` parameter 0 is a pointer into local memory (address space 5)" - ), - "{error}" - ); -} - -/// The refusal is scoped to what the host cannot supply. Generic, global, -/// and constant pointer parameters are host-addressable and stay accepted on -/// kernels, and a device (non-kernel) function may carry any state space on -/// its parameters, shared included. -#[test] -fn kernel_keeps_host_addressable_pointer_parameters_and_device_functions_keep_shared() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "kernel_allowed_params".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - pointer_param_function(&mut ctx, module_block, "allowed_param", &[0, 1, 4], true); - pointer_param_function(&mut ctx, module_block, "device_shared", &[3], false); - - let ir = export_module_to_string(&ctx, &module) - .expect("generic/global/constant kernel parameters and device shared must export"); - assert!( - ir.contains( - "define ptx_kernel void @allowed_param(ptr %v0, ptr addrspace(1) %v1, ptr addrspace(4) %v2)" - ), - "{ir}" - ); - assert!( - ir.contains("define void @device_shared(ptr addrspace(3) %v0)"), - "{ir}" - ); -} - -#[test] -fn full_debug_metadata_emits_scalarized_fragment_dbg_declares() { - let mut ctx = Context::new(); - - let module = ModuleOp::new(&mut ctx, "fragment_debug".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.into(), vec![], false); - let func = FuncOp::new( - &mut ctx, - "fragment_debug_kernel".try_into().unwrap(), - func_ty, - ); - let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/fragments.rs", 10, 1); - func.get_operation().deref_mut(&ctx).set_loc(func_loc); - let entry = func.get_or_create_entry_block(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let one_attr = IntegerAttr::new(i32_ty, APInt::from_u32(1, NonZero::new(32).unwrap())); - let one = ConstantOp::new(&mut ctx, one_attr.into()); - let one_value = one.get_operation().deref(&ctx).get_result(0); - one.get_operation().insert_at_back(entry, &ctx); - - let whole_ty = DebugLocalTypeKind::Array { - name: "[u32; 2]".to_string(), - size_bits: 64, - element: Box::new(DebugLocalTypeKind::Basic { - name: "u32".to_string(), - size_bits: 32, - encoding: "DW_ATE_unsigned", - }), - count: 2, - }; - for (index, offset_bits) in [0u64, 32].into_iter().enumerate() { - let alloca = AllocaOp::new(&mut ctx, i32_ty.into(), one_value); - llvm_export::ops::set_debug_fragment_variables( - &mut ctx, - alloca.get_operation(), - &[DebugFragmentVariableInfo { - variable: DebugLocalVariableInfo { - name: "pair".to_string(), - argument_index: None, - ty: whole_ty.clone(), - }, - fragment: DebugFragment { - offset_bits, - size_bits: 32, - }, - source_scope: None, - declaration: Some(DebugSourcePosition { - file: PathBuf::from("/tmp/cuda-oxide/tests/fragments.rs"), - line: 11, - column: 9, - }), - }], - ); - let loc = src_location( - &mut ctx, - "/tmp/cuda-oxide/tests/fragments.rs", - 12 + index as i32, - 9, - ); - alloca.get_operation().deref_mut(&ctx).set_loc(loc); - alloca.get_operation().insert_at_back(entry, &ctx); - } - - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - let config = DebugConfig { - inner: PtxExportConfig, - debug_kind: DebugKind::Full, - }; - let ir = - export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); - - assert!( - ir.contains("DW_OP_LLVM_fragment, 0, 32"), - "first scalarized piece should describe the low fragment:\n{ir}" - ); - assert!( - ir.contains("DW_OP_LLVM_fragment, 32, 32"), - "second scalarized piece should describe the high fragment:\n{ir}" - ); - assert_eq!( - ir.matches("!DILocalVariable(name: \"pair\"").count(), - 1, - "all pieces must share one DILocalVariable identity:\n{ir}" - ); - assert_eq!( - ir.matches("call void @llvm.dbg.declare").count(), - 2, - "each scalarized storage piece should emit one dbg.declare:\n{ir}" - ); - - let Some(llvm_as) = ["llvm-as-22", "llvm-as-21", "llvm-as"] - .into_iter() - .find(|tool| { - std::process::Command::new(tool) - .arg("--version") - .output() - .is_ok_and(|out| out.status.success()) - }) - else { - eprintln!("skipping fragment llvm-as parse gate: no llvm-as on PATH"); - return; - }; - let ll_path = std::env::temp_dir().join(format!( - "cuda_oxide_fragment_parse_gate_{}.ll", - std::process::id() - )); - std::fs::write(&ll_path, &ir).expect("write fragment temp .ll"); - let output = std::process::Command::new(llvm_as) - .arg("-o") - .arg("/dev/null") - .arg(&ll_path) - .output() - .expect("run llvm-as for fragments"); - let stderr = String::from_utf8_lossy(&output.stderr).into_owned(); - let _ = std::fs::remove_file(&ll_path); - assert!( - output.status.success(), - "{llvm_as} rejected scalarized fragment debug metadata:\n{stderr}\n--- module ---\n{ir}" - ); -} - -#[test] -fn full_debug_metadata_emits_projected_dbg_declares() { - let mut ctx = Context::new(); - - let module = ModuleOp::new(&mut ctx, "projected_debug".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.into(), vec![], false); - let func = FuncOp::new( - &mut ctx, - "projected_debug_kernel".try_into().unwrap(), - func_ty, - ); - let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/projected.rs", 10, 1); - func.get_operation().deref_mut(&ctx).set_loc(func_loc); - let entry = func.get_or_create_entry_block(&mut ctx); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let one_attr = IntegerAttr::new(i32_ty, APInt::from_u32(1, NonZero::new(32).unwrap())); - let one = ConstantOp::new(&mut ctx, one_attr.into()); - let one_value = one.get_operation().deref(&ctx).get_result(0); - one.get_operation().insert_at_back(entry, &ctx); - - let storage_ty = ArrayType::get(&ctx, i32_ty.into(), 8); - let alloca = AllocaOp::new(&mut ctx, storage_ty.into(), one_value); - llvm_export::ops::set_debug_projected_variables( - &mut ctx, - alloca.get_operation(), - &[ - DebugProjectedVariableInfo { - variable: DebugLocalVariableInfo { - name: "field_value".to_string(), - argument_index: None, - ty: DebugLocalTypeKind::Basic { - name: "u32".to_string(), - size_bits: 32, - encoding: "DW_ATE_unsigned", - }, - }, - dereference_base: false, - offset_bytes: 8, - source_scope: None, - declaration: Some(DebugSourcePosition { - file: PathBuf::from("/tmp/cuda-oxide/tests/projected.rs"), - line: 11, - column: 9, - }), - }, - DebugProjectedVariableInfo { - variable: DebugLocalVariableInfo { - name: "tuple_value".to_string(), - argument_index: None, - ty: DebugLocalTypeKind::Basic { - name: "u64".to_string(), - size_bits: 64, - encoding: "DW_ATE_unsigned", - }, - }, - dereference_base: false, - offset_bytes: 16, - source_scope: None, - declaration: Some(DebugSourcePosition { - file: PathBuf::from("/tmp/cuda-oxide/tests/projected.rs"), - line: 12, - column: 9, - }), - }, - DebugProjectedVariableInfo { - variable: DebugLocalVariableInfo { - name: "array_value".to_string(), - argument_index: None, - ty: DebugLocalTypeKind::Basic { - name: "u32".to_string(), - size_bits: 32, - encoding: "DW_ATE_unsigned", - }, - }, - dereference_base: false, - offset_bytes: 24, - source_scope: None, - declaration: Some(DebugSourcePosition { - file: PathBuf::from("/tmp/cuda-oxide/tests/projected.rs"), - line: 13, - column: 9, - }), - }, - DebugProjectedVariableInfo { - variable: DebugLocalVariableInfo { - name: "deref_value".to_string(), - argument_index: None, - ty: DebugLocalTypeKind::Basic { - name: "u32".to_string(), - size_bits: 32, - encoding: "DW_ATE_unsigned", - }, - }, - dereference_base: true, - offset_bytes: 0, - source_scope: None, - declaration: Some(DebugSourcePosition { - file: PathBuf::from("/tmp/cuda-oxide/tests/projected.rs"), - line: 14, - column: 9, - }), - }, - DebugProjectedVariableInfo { - variable: DebugLocalVariableInfo { - name: "deref_field_value".to_string(), - argument_index: None, - ty: DebugLocalTypeKind::Basic { - name: "u64".to_string(), - size_bits: 64, - encoding: "DW_ATE_unsigned", - }, - }, - dereference_base: true, - offset_bytes: 32, - source_scope: None, - declaration: Some(DebugSourcePosition { - file: PathBuf::from("/tmp/cuda-oxide/tests/projected.rs"), - line: 15, - column: 9, - }), - }, - ], - ); - alloca.get_operation().insert_at_back(entry, &ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); - func.get_operation().insert_at_back(module_block, &ctx); - - let config = DebugConfig { - inner: PtxExportConfig, - debug_kind: DebugKind::Full, - }; - let ir = export_module_to_string_with_config(&ctx, &module, &config) - .expect("projected debug export succeeds"); - - assert_eq!( - ir.matches("call void @llvm.dbg.declare").count(), - 5, - "each projected source variable should get its own dbg.declare:\n{ir}" - ); - for name in [ - "field_value", - "tuple_value", - "array_value", - "deref_value", - "deref_field_value", - ] { - assert!( - ir.contains(&format!("!DILocalVariable(name: \"{name}\"")), - "missing projected variable {name}:\n{ir}" - ); - } - for offset in [8u64, 16, 24] { - assert!( - ir.contains(&format!("!DIExpression(DW_OP_plus_uconst, {offset})")), - "missing static projection offset {offset}:\n{ir}" - ); - } - assert!( - ir.contains("!DIExpression(DW_OP_deref)"), - "missing dereference-only debug expression:\n{ir}" - ); - assert!( - ir.contains("!DIExpression(DW_OP_deref, DW_OP_plus_uconst, 32)"), - "missing dereference-plus-field debug expression:\n{ir}" - ); -} diff --git a/crates/llvm-export/tests/export_test/calls_and_addresses.rs b/crates/llvm-export/tests/export_test/calls_and_addresses.rs new file mode 100644 index 0000000000..7535b36f12 --- /dev/null +++ b/crates/llvm-export/tests/export_test/calls_and_addresses.rs @@ -0,0 +1,377 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use llvm_export::{ + export::{ + NvvmExportConfig, NvvmIrDialect, export_module_to_string, + export_module_to_string_with_config, + }, + ops::{AddressOfOp, BrOp, CallOp, FuncOp, GepIndex, GetElementPtrOp, GlobalOp, ReturnOp}, + types::{FuncType, PointerType, VoidType}, +}; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + op_interfaces::CallOpCallable, + ops::ModuleOp, + types::{IntegerType, Signedness}, + }, + context::Context, + linked_list::ContainsLinkedList, + op::Op, +}; + +use crate::common::{assert_no_undefined_temporaries, module_top_block}; + +#[test] +fn indirect_call_rejects_non_program_address_space() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "invalid_indirect_callee".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + + let shared_ptr_ty = PointerType::get(&ctx, 3); + let void_ty = VoidType::get(&ctx); + let callee_ty = FuncType::get(&ctx, void_ty.into(), vec![], false); + let caller_ty = FuncType::get(&ctx, void_ty.into(), vec![shared_ptr_ty.into()], false); + let caller = FuncOp::new(&mut ctx, "caller".try_into().unwrap(), caller_ty); + let entry = caller.get_or_create_entry_block(&mut ctx); + let callee = entry.deref(&ctx).get_argument(0); + CallOp::new( + &mut ctx, + CallOpCallable::Indirect(callee), + callee_ty, + vec![], + ) + .get_operation() + .insert_at_back(entry, &ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + caller.get_operation().insert_at_back(module_block, &ctx); + + for dialect in [NvvmIrDialect::LegacyLlvm7, NvvmIrDialect::Modern] { + let error = + export_module_to_string_with_config(&ctx, &module, &NvvmExportConfig::new(dialect)) + .expect_err("NVPTX must reject a shared-memory function pointer"); + assert!(error.contains("address space 3"), "{error}"); + assert!(error.contains("function pointers"), "{error}"); + } +} + +#[test] +fn intrinsic_export_preserves_legacy_dots_and_literal_underscores() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "intrinsic_names".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let void_ty = VoidType::get(&ctx); + let function_ty = FuncType::get(&ctx, void_ty.into(), vec![], false); + let legacy = "llvm_nvvm_wgmma_fence_sync_aligned"; + let escaped = "llvm__nvvm_dwgmma_dcommit_ugroup_dsync_daligned"; + + for name in [legacy, escaped] { + FuncOp::new(&mut ctx, name.try_into().unwrap(), function_ty) + .get_operation() + .insert_at_back(module_block, &ctx); + } + + let caller = FuncOp::new(&mut ctx, "caller".try_into().unwrap(), function_ty); + let entry = caller.get_or_create_entry_block(&mut ctx); + for name in [legacy, escaped] { + CallOp::new( + &mut ctx, + CallOpCallable::Direct(name.try_into().unwrap()), + function_ty, + vec![], + ) + .get_operation() + .insert_at_back(entry, &ctx); + } + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + caller.get_operation().insert_at_back(module_block, &ctx); + + let ir = export_module_to_string(&ctx, &module).expect("intrinsic export succeeds"); + assert!(ir.contains("@llvm.nvvm.wgmma.fence.sync.aligned"), "{ir}"); + assert!( + ir.contains("@llvm.nvvm.wgmma.commit_group.sync.aligned"), + "{ir}" + ); + assert!(!ir.contains("@llvm.nvvm.wgmma.commit.group"), "{ir}"); +} + +#[test] +fn legacy_function_address_defined_later_round_trips_through_indirect_call() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "function_address".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let void_ty = VoidType::get(&ctx); + let callee_ty = FuncType::get(&ctx, void_ty.into(), vec![], false); + + // Print the caller first to prove symbol typing is a module pre-pass, not + // an accidental dependency on textual definition order. + let caller = FuncOp::new( + &mut ctx, + "call_function_pointer".try_into().unwrap(), + callee_ty, + ); + let caller_entry = caller.get_or_create_entry_block(&mut ctx); + let address = AddressOfOp::new(&mut ctx, "target".try_into().unwrap(), 0); + let address_value = address.get_operation().deref(&ctx).get_result(0); + address.get_operation().insert_at_back(caller_entry, &ctx); + CallOp::new( + &mut ctx, + CallOpCallable::Indirect(address_value), + callee_ty, + vec![], + ) + .get_operation() + .insert_at_back(caller_entry, &ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(caller_entry, &ctx); + caller.get_operation().insert_at_back(module_block, &ctx); + + let target = FuncOp::new(&mut ctx, "target".try_into().unwrap(), callee_ty); + let target_entry = target.get_or_create_entry_block(&mut ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(target_entry, &ctx); + target.get_operation().insert_at_back(module_block, &ctx); + + let ir = export_module_to_string_with_config( + &ctx, + &module, + &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + ) + .expect("legacy function-address export succeeds"); + assert!( + ir.contains("bitcast void ()* @target to i8*"), + "function address must normalize to the canonical byte pointer:\n{ir}" + ); + assert!( + ir.contains("bitcast i8*") && ir.contains("to void ()*"), + "indirect call must restore the exact function pointer type:\n{ir}" + ); + assert!(ir.contains("call void %"), "{ir}"); +} + +#[test] +fn modern_function_address_uses_the_normalized_definition_name() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "modern_function_address".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let void_ty = VoidType::get(&ctx); + let callee_ty = FuncType::get(&ctx, void_ty.into(), vec![], false); + let prefixed_name = reserved_oxide_symbols::device_symbol("target"); + + let caller = FuncOp::new(&mut ctx, "caller".try_into().unwrap(), callee_ty); + let caller_entry = caller.get_or_create_entry_block(&mut ctx); + let address = AddressOfOp::new(&mut ctx, prefixed_name.as_str().try_into().unwrap(), 0); + let address_value = address.get_operation().deref(&ctx).get_result(0); + address.get_operation().insert_at_back(caller_entry, &ctx); + CallOp::new( + &mut ctx, + CallOpCallable::Indirect(address_value), + callee_ty, + vec![], + ) + .get_operation() + .insert_at_back(caller_entry, &ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(caller_entry, &ctx); + caller.get_operation().insert_at_back(module_block, &ctx); + + let target = FuncOp::new( + &mut ctx, + prefixed_name.as_str().try_into().unwrap(), + callee_ty, + ); + let target_entry = target.get_or_create_entry_block(&mut ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(target_entry, &ctx); + target.get_operation().insert_at_back(module_block, &ctx); + + let ir = export_module_to_string_with_config( + &ctx, + &module, + &NvvmExportConfig::new(NvvmIrDialect::Modern), + ) + .expect("modern function-address export succeeds"); + assert!(ir.contains("define void @target()"), "{ir}"); + assert!(ir.contains("call void @target()"), "{ir}"); + assert!(!ir.contains(&prefixed_name), "{ir}"); +} + +#[test] +fn modern_addressof_rejects_global_and_function_address_space_mismatches() { + let mut ctx = Context::new(); + let void_ty = VoidType::get(&ctx); + let no_args = FuncType::get(&ctx, void_ty.into(), vec![], false); + + let global_module = ModuleOp::new(&mut ctx, "bad_global_address".try_into().unwrap()); + let global_module_block = module_top_block(&mut ctx, &global_module); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let global = GlobalOp::new(&mut ctx, "shared_value".try_into().unwrap(), i32_ty.into()); + global.set_address_space(&mut ctx, 3); + global + .get_operation() + .insert_at_back(global_module_block, &ctx); + let global_user = FuncOp::new(&mut ctx, "global_user".try_into().unwrap(), no_args); + let global_entry = global_user.get_or_create_entry_block(&mut ctx); + AddressOfOp::new(&mut ctx, "shared_value".try_into().unwrap(), 0) + .get_operation() + .insert_at_back(global_entry, &ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(global_entry, &ctx); + global_user + .get_operation() + .insert_at_back(global_module_block, &ctx); + let error = export_module_to_string_with_config( + &ctx, + &global_module, + &NvvmExportConfig::new(NvvmIrDialect::Modern), + ) + .expect_err("modern global addressof must preserve address spaces"); + assert!( + error.contains("result is 0, global is 3"), + "unexpected global error: {error}" + ); + + let function_module = ModuleOp::new(&mut ctx, "bad_function_address".try_into().unwrap()); + let function_module_block = module_top_block(&mut ctx, &function_module); + FuncOp::new(&mut ctx, "target".try_into().unwrap(), no_args) + .get_operation() + .insert_at_back(function_module_block, &ctx); + let function_user = FuncOp::new(&mut ctx, "function_user".try_into().unwrap(), no_args); + let function_entry = function_user.get_or_create_entry_block(&mut ctx); + AddressOfOp::new(&mut ctx, "target".try_into().unwrap(), 3) + .get_operation() + .insert_at_back(function_entry, &ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(function_entry, &ctx); + function_user + .get_operation() + .insert_at_back(function_module_block, &ctx); + let error = export_module_to_string_with_config( + &ctx, + &function_module, + &NvvmExportConfig::new(NvvmIrDialect::Modern), + ) + .expect_err("modern function addressof must use program address space"); + assert!(error.contains("program-address-space (0)"), "{error}"); +} + +#[test] +fn export_addressof_uses_symbol_when_definition_block_prints_later() { + let mut ctx = Context::new(); + + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_region = module.get_operation().deref(&ctx).get_region(0); + let module_block = { + let existing = { + let region = module_region.deref(&ctx); + region.iter(&ctx).next() + }; + if let Some(block) = existing { + block + } else { + let block = BasicBlock::new(&mut ctx, None, vec![]); + block.insert_at_back(module_region, &ctx); + block + } + }; + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let global = GlobalOp::new( + &mut ctx, + "__shared_mem_20".try_into().unwrap(), + i32_ty.to_handle(), + ); + global.set_address_space(&mut ctx, 3); + global.get_operation().insert_at_back(module_block, &ctx); + + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); + let func = FuncOp::new(&mut ctx, "uses_late_addressof".try_into().unwrap(), func_ty); + let entry = func.get_or_create_entry_block(&mut ctx); + let func_region = func.get_operation().deref(&ctx).get_region(0); + let use_block = BasicBlock::new(&mut ctx, None, vec![]); + use_block.insert_at_back(func_region, &ctx); + let address_block = BasicBlock::new(&mut ctx, None, vec![]); + address_block.insert_at_back(func_region, &ctx); + + BrOp::new(&mut ctx, address_block, vec![]) + .get_operation() + .insert_at_back(entry, &ctx); + + let address = AddressOfOp::new(&mut ctx, "__shared_mem_20".try_into().unwrap(), 3); + let address_value = address.get_operation().deref(&ctx).get_result(0); + address.get_operation().insert_at_back(address_block, &ctx); + BrOp::new(&mut ctx, use_block, vec![]) + .get_operation() + .insert_at_back(address_block, &ctx); + + let gep = GetElementPtrOp::new( + &mut ctx, + address_value, + vec![GepIndex::Constant(0)], + i32_ty.to_handle(), + ); + gep.get_operation().insert_at_back(use_block, &ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(use_block, &ctx); + + func.get_operation().insert_at_back(module_block, &ctx); + + let ir = export_module_to_string(&ctx, &module).expect("export succeeds"); + + // The shared global must be declared at module scope. + assert!( + ir.contains("@__shared_mem_20 = addrspace(3) global"), + "module must declare the shared global:\n{ir}" + ); + + // The GEP base operand must be the global symbol, not a stale `%vN`. + let gep_line = ir + .lines() + .find(|line| line.contains("getelementptr inbounds")) + .expect("exported GEP line"); + assert!( + gep_line.contains("@__shared_mem_20"), + "GEP must use the global symbol, not a stale temporary:\n{ir}" + ); + + // Bug class from issue #54: every `%vN` reference in the IR must have a + // matching `%vN = ...` definition. With the bug present the addressof + // result was named `%v1` but never defined; this catches that and any + // future regression that re-introduces a dangling SSA reference. + assert_no_undefined_temporaries(&ir); + + let legacy = export_module_to_string_with_config( + &ctx, + &module, + &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + ) + .expect("legacy addressof export succeeds"); + assert!( + legacy.contains("@__shared_mem_20 = addrspace(3) global i32 undef"), + "NVVM shared globals must be uninitialized:\n{legacy}" + ); + assert!( + legacy.contains("bitcast i32 addrspace(3)* @__shared_mem_20 to i8 addrspace(3)*"), + "legacy addressof must normalize the global pointer:\n{legacy}" + ); + assert!( + legacy.contains("bitcast i8 addrspace(3)*") && legacy.contains("to i32 addrspace(3)*"), + "legacy GEP must repair its canonical base pointer:\n{legacy}" + ); + assert_no_undefined_temporaries(&legacy); +} diff --git a/crates/llvm-export/tests/export_test/common.rs b/crates/llvm-export/tests/export_test/common.rs new file mode 100644 index 0000000000..40583f911b --- /dev/null +++ b/crates/llvm-export/tests/export_test/common.rs @@ -0,0 +1,132 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use combine::stream::position::SourcePosition; +use llvm_export::export::{DebugKind, ExportBackendConfig}; +use pliron::{ + basic_block::BasicBlock, + builtin::ops::ModuleOp, + context::{Context, Ptr}, + linked_list::ContainsLinkedList, + location::{Location, Source}, + op::Op, +}; +use std::path::PathBuf; + +pub(super) struct DebugConfig { + pub(super) inner: C, + pub(super) debug_kind: DebugKind, +} + +impl ExportBackendConfig for DebugConfig { + fn datalayout(&self) -> &str { + self.inner.datalayout() + } + + fn emit_llvm_used(&self) -> bool { + self.inner.emit_llvm_used() + } + + fn emit_nvvmir_version(&self) -> bool { + self.inner.emit_nvvmir_version() + } + + fn nvvmir_version(&self) -> [i32; 4] { + self.inner.nvvmir_version() + } + + fn emit_all_kernel_annotations(&self) -> bool { + self.inner.emit_all_kernel_annotations() + } + + fn emit_ptx_kernel_keyword(&self) -> bool { + self.inner.emit_ptx_kernel_keyword() + } + + fn nvvm_ir_dialect(&self) -> Option { + self.inner.nvvm_ir_dialect() + } + + fn debug_kind(&self) -> DebugKind { + self.debug_kind + } +} + +pub(super) fn src_location(ctx: &mut Context, file: &str, line: i32, column: i32) -> Location { + Location::SrcPos { + src: Source::new_from_file(ctx, PathBuf::from(file)), + pos: SourcePosition { line, column }, + } +} + +pub(super) fn module_top_block(ctx: &mut Context, module: &ModuleOp) -> Ptr { + let module_region = module.get_operation().deref(ctx).get_region(0); + { + let region = module_region.deref(ctx); + if let Some(block) = region.iter(ctx).next() { + return block; + } + } + + let block = BasicBlock::new(ctx, None, vec![]); + block.insert_at_back(module_region, ctx); + block +} + +pub(super) fn metadata_id<'a>(ir: &'a str, needle: &str) -> &'a str { + ir.lines() + .find(|line| line.contains(needle)) + .and_then(|line| line.split_once(" = ").map(|(id, _)| id)) + .unwrap_or_else(|| panic!("missing metadata node containing {needle:?}:\n{ir}")) +} + +/// Scans the textual LLVM IR and asserts that every `%vN` token appearing in +/// an operand position has a corresponding `%vN = ...` definition somewhere +/// in the module. Operates on `%v` temporaries only because that's the +/// exporter's naming scheme; named values like `%entry` (block labels) are +/// ignored by construction. +pub(super) fn assert_no_undefined_temporaries(ir: &str) { + use std::collections::HashSet; + + let mut defined: HashSet = HashSet::new(); + for line in ir.lines() { + let trimmed = line.trim_start(); + if !trimmed.starts_with("%v") { + continue; + } + let Some((lhs, _)) = trimmed.split_once('=') else { + continue; + }; + defined.insert(lhs.trim().to_string()); + } + + let mut referenced: HashSet = HashSet::new(); + for line in ir.lines() { + let trimmed = line.trim_start(); + // Skip the lhs of a definition; only operand positions can be stale. + let body = if trimmed.starts_with("%v") + && let Some(eq) = trimmed.find('=') + { + &trimmed[eq + 1..] + } else { + trimmed + }; + for tok in body.split(|c: char| !c.is_alphanumeric() && c != '%' && c != '_') { + if let Some(num) = tok.strip_prefix("%v") + && !num.is_empty() + && num.chars().all(|c| c.is_ascii_digit()) + { + referenced.insert(format!("%v{num}")); + } + } + } + + let mut undefined: Vec<&String> = referenced.difference(&defined).collect(); + undefined.sort(); + assert!( + undefined.is_empty(), + "IR references undefined SSA temporaries: {undefined:?}\nIR:\n{ir}" + ); +} diff --git a/crates/llvm-export/tests/export_test/control_flow.rs b/crates/llvm-export/tests/export_test/control_flow.rs new file mode 100644 index 0000000000..ceb24956ce --- /dev/null +++ b/crates/llvm-export/tests/export_test/control_flow.rs @@ -0,0 +1,234 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use llvm_export::{ + export::{NvvmExportConfig, NvvmIrDialect, export_module_to_string_with_config}, + ops::{BrOp, CondBrOp, ConstantOp, FuncOp, ReturnOp, UndefOp}, + types::{FuncType, VoidType}, +}; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::IntegerAttr, + ops::ModuleOp, + types::{IntegerType, Signedness}, + }, + common_traits::Verify, + context::Context, + op::Op, + utils::apint::APInt, +}; +use std::num::NonZero; + +use crate::common::module_top_block; + +#[test] +fn exporter_rejects_extra_predecessor_values_before_emitting_phis() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "invalid_branch_arity".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.into(), vec![], false); + let func = FuncOp::new(&mut ctx, "invalid_branch".try_into().unwrap(), func_ty); + let entry = func.get_or_create_entry_block(&mut ctx); + let region = func.get_operation().deref(&ctx).get_region(0); + let destination = BasicBlock::new(&mut ctx, None, vec![]); + destination.insert_at_back(region, &ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let one_attr = IntegerAttr::new(i32_ty, APInt::from_u32(1, NonZero::new(32).unwrap())); + let one = ConstantOp::new(&mut ctx, one_attr.into()); + let one_value = one.get_operation().deref(&ctx).get_result(0); + one.get_operation().insert_at_back(entry, &ctx); + BrOp::new(&mut ctx, destination, vec![one_value]) + .get_operation() + .insert_at_back(entry, &ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(destination, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + let error = export_module_to_string_with_config( + &ctx, + &module, + &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + ) + .expect_err("extra predecessor values must be rejected"); + assert!( + error.contains("supplies 1 values") && error.contains("expects 0 block arguments"), + "{error}" + ); +} + +#[test] +fn exporter_rejects_distinct_values_on_duplicate_conditional_edges() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "duplicate_conditional_edge".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let func_ty = FuncType::get( + &ctx, + i32_ty.into(), + vec![i1_ty.into(), i32_ty.into(), i32_ty.into()], + false, + ); + let func = FuncOp::new(&mut ctx, "duplicate_edge".try_into().unwrap(), func_ty); + let entry = func.get_or_create_entry_block(&mut ctx); + let condition = entry.deref(&ctx).get_argument(0); + let if_true = entry.deref(&ctx).get_argument(1); + let if_false = entry.deref(&ctx).get_argument(2); + let region = func.get_operation().deref(&ctx).get_region(0); + let destination = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); + destination.insert_at_back(region, &ctx); + CondBrOp::new( + &mut ctx, + condition, + destination, + vec![if_true], + destination, + vec![if_false], + ) + .get_operation() + .insert_at_back(entry, &ctx); + let result = destination.deref(&ctx).get_argument(0); + ReturnOp::new(&mut ctx, Some(result)) + .get_operation() + .insert_at_back(destination, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + module + .get_operation() + .deref(&ctx) + .verify(&ctx) + .expect("pliron permits same-destination conditional edges with distinct values"); + let error = export_module_to_string_with_config( + &ctx, + &module, + &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + ) + .expect_err("LLVM PHIs cannot distinguish duplicate predecessor edges"); + assert!( + error.contains("both edges with different forwarded values"), + "{error}" + ); +} + +#[test] +fn exporter_deduplicates_identical_values_on_duplicate_conditional_edges() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "identical_conditional_edge".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let func_ty = FuncType::get( + &ctx, + i32_ty.into(), + vec![i1_ty.into(), i32_ty.into()], + false, + ); + let func = FuncOp::new(&mut ctx, "identical_edge".try_into().unwrap(), func_ty); + let entry = func.get_or_create_entry_block(&mut ctx); + let condition = entry.deref(&ctx).get_argument(0); + let value = entry.deref(&ctx).get_argument(1); + let region = func.get_operation().deref(&ctx).get_region(0); + let destination = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); + destination.insert_at_back(region, &ctx); + CondBrOp::new( + &mut ctx, + condition, + destination, + vec![value], + destination, + vec![value], + ) + .get_operation() + .insert_at_back(entry, &ctx); + let result = destination.deref(&ctx).get_argument(0); + ReturnOp::new(&mut ctx, Some(result)) + .get_operation() + .insert_at_back(destination, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + let ir = export_module_to_string_with_config( + &ctx, + &module, + &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + ) + .expect("identical duplicate-edge values can use one PHI predecessor"); + let phi = ir + .lines() + .find(|line| line.contains(" = phi i32 ")) + .expect("destination block must contain a PHI"); + assert_eq!(phi.matches("%entry").count(), 1, "{phi}"); +} + +#[test] +fn phi_can_reference_undef_from_a_later_block() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "later_undef_phi".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let func_ty = FuncType::get( + &ctx, + i32_ty.into(), + vec![i1_ty.into(), i32_ty.into()], + false, + ); + let func = FuncOp::new(&mut ctx, "choose_undef".try_into().unwrap(), func_ty); + let entry = func.get_or_create_entry_block(&mut ctx); + let condition = entry.deref(&ctx).get_argument(0); + let fallback = entry.deref(&ctx).get_argument(1); + let region = func.get_operation().deref(&ctx).get_region(0); + + // The join precedes both predecessors in print order, so its PHI depends + // on the exporter's whole-function value-name pre-pass. + let join = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); + join.insert_at_back(region, &ctx); + let undef_block = BasicBlock::new(&mut ctx, None, vec![]); + undef_block.insert_at_back(region, &ctx); + let value_block = BasicBlock::new(&mut ctx, None, vec![]); + value_block.insert_at_back(region, &ctx); + + CondBrOp::new( + &mut ctx, + condition, + undef_block, + vec![], + value_block, + vec![], + ) + .get_operation() + .insert_at_back(entry, &ctx); + + let undef = UndefOp::new(&mut ctx, i32_ty.into()); + let undef_value = undef.get_operation().deref(&ctx).get_result(0); + undef.get_operation().insert_at_back(undef_block, &ctx); + BrOp::new(&mut ctx, join, vec![undef_value]) + .get_operation() + .insert_at_back(undef_block, &ctx); + BrOp::new(&mut ctx, join, vec![fallback]) + .get_operation() + .insert_at_back(value_block, &ctx); + + let result = join.deref(&ctx).get_argument(0); + ReturnOp::new(&mut ctx, Some(result)) + .get_operation() + .insert_at_back(join, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + let ir = export_module_to_string_with_config( + &ctx, + &module, + &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + ) + .expect("later-block undef must be available while exporting an earlier PHI"); + assert!( + ir.lines() + .any(|line| line.contains(" = phi i32 ") && line.contains("[ undef,")), + "{ir}" + ); +} diff --git a/crates/llvm-export/tests/export_test/debug_info.rs b/crates/llvm-export/tests/export_test/debug_info.rs new file mode 100644 index 0000000000..6b941ff8c7 --- /dev/null +++ b/crates/llvm-export/tests/export_test/debug_info.rs @@ -0,0 +1,2458 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use llvm_export::{ + export::{ + DebugKind, NvvmExportConfig, NvvmIrDialect, PtxExportConfig, + export_module_to_string_with_config, + }, + ops::{ + AllocaOp, CallOp, ConstantOp, DebugEnumDiscriminant, DebugEnumVariant, DebugFragment, + DebugFragmentVariableInfo, DebugGlobalVariableInfo, DebugLocalTypeKind, + DebugLocalVariableInfo, DebugProjectedVariableInfo, DebugSourcePosition, DebugSourceScope, + DebugSourceScopeLocation, DebugSourceScopeMap, DebugValueExpression, + DebugValueExpressionOp, DebugValueListOp, DebugValueOp, FuncOp, + GlobalInitializerRelocation, GlobalOp, GlobalOpExt, ReturnOp, + encode_global_initializer_relocations, + }, + types::{ArrayType, FuncType, PointerType, StructLayout, StructType, VoidType}, +}; +use pliron::{ + builtin::{ + attributes::IntegerAttr, + op_interfaces::CallOpCallable, + ops::ModuleOp, + types::{IntegerType, Signedness}, + }, + context::Context, + linked_list::ContainsLinkedList, + location::{Located, Location}, + op::Op, + utils::apint::APInt, +}; +use std::{num::NonZero, path::PathBuf}; + +use crate::common::{DebugConfig, metadata_id, module_top_block, src_location}; + +#[test] +fn legacy_export_rejects_debug_metadata() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "legacy_debug".try_into().unwrap()); + let config = DebugConfig { + inner: NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + debug_kind: DebugKind::LineTables, + }; + + let error = export_module_to_string_with_config(&ctx, &module, &config) + .expect_err("legacy debug output must be rejected"); + assert!(error.contains("legacy LLVM 7"), "{error}"); + assert!(error.contains("debug"), "{error}"); +} + +#[test] +fn line_table_debug_metadata_emits_function_scope_and_instruction_locations() { + let mut ctx = Context::new(); + + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_region = module.get_operation().deref(&ctx).get_region(0); + let module_block = { + let region = module_region.deref(&ctx); + region.iter(&ctx).next().unwrap() + }; + + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); + let func = FuncOp::new(&mut ctx, "debug_kernel".try_into().unwrap(), func_ty); + llvm_export::ops::set_debug_function_name( + &mut ctx, + func.get_operation(), + "source_crate::debug_kernel", + ); + let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 7, 1); + func.get_operation().deref_mut(&ctx).set_loc(func_loc); + + let entry = func.get_or_create_entry_block(&mut ctx); + let ret = ReturnOp::new(&mut ctx, None); + let ret_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 8, 5); + ret.get_operation().deref_mut(&ctx).set_loc(ret_loc); + ret.get_operation().insert_at_back(entry, &ctx); + + func.get_operation().insert_at_back(module_block, &ctx); + + let config = DebugConfig { + inner: PtxExportConfig, + debug_kind: DebugKind::LineTables, + }; + let ir = + export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); + + let define_line = ir + .lines() + .find(|line| line.starts_with("define ")) + .expect("function definition"); + assert!( + define_line.contains("!dbg !"), + "function definition should reference its DISubprogram:\n{ir}" + ); + + let ret_line = ir + .lines() + .find(|line| line.trim_start().starts_with("ret void")) + .expect("return instruction"); + assert!( + ret_line.contains(", !dbg !"), + "real instructions should carry DILocation attachments:\n{ir}" + ); + + assert!( + ir.contains("!llvm.dbg.cu = !{!"), + "module should reference a compile unit:\n{ir}" + ); + assert!( + ir.contains("!llvm.module.flags = !{!"), + "module should declare debug-info flags:\n{ir}" + ); + assert!( + ir.contains("!DIFile(filename: \"kernel.rs\", directory: \"/tmp/cuda-oxide/tests\")"), + "source path should be split into DIFile filename and directory:\n{ir}" + ); + assert!( + ir.contains("distinct !DICompileUnit(language: DW_LANG_Rust"), + "debug export should describe the Rust compile unit:\n{ir}" + ); + assert!( + ir.contains( + "distinct !DISubprogram(name: \"source_crate::debug_kernel\", linkageName: \"debug_kernel\"" + ), + "function definition should separate its source name from its physical linkage name:\n{ir}" + ); + assert!( + ir.contains("!DILocation(line: 8, column: 5, scope: !"), + "instruction location should preserve the source line and column:\n{ir}" + ); +} + +#[test] +fn full_debug_metadata_keeps_generic_source_name_separate_from_mangled_symbol() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); + let linkage_name = "_RNvMCexampleINtCsource7WrapperjE7get_mut"; + let func = FuncOp::new(&mut ctx, linkage_name.try_into().unwrap(), func_ty); + llvm_export::ops::set_debug_function_name( + &mut ctx, + func.get_operation(), + "source_crate::Wrapper::::get_mut", + ); + let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/generic.rs", 19, 1); + func.get_operation().deref_mut(&ctx).set_loc(func_loc); + + let entry = func.get_or_create_entry_block(&mut ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + let config = DebugConfig { + inner: PtxExportConfig, + debug_kind: DebugKind::Full, + }; + let ir = + export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); + + assert!( + ir.contains(&format!("define void @{linkage_name}()")), + "debug naming must not change the physical function symbol:\n{ir}" + ); + assert!( + ir.contains(&format!( + "distinct !DISubprogram(name: \"source_crate::Wrapper::::get_mut\", linkageName: \"{linkage_name}\"" + )), + "generic debug metadata must carry both source and linkage spellings:\n{ir}" + ); +} + +#[test] +fn line_table_debug_metadata_uses_file_scope_for_cross_file_locations() { + let mut ctx = Context::new(); + + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_region = module.get_operation().deref(&ctx).get_region(0); + let module_block = { + let region = module_region.deref(&ctx); + region.iter(&ctx).next().unwrap() + }; + + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); + let func = FuncOp::new(&mut ctx, "debug_kernel".try_into().unwrap(), func_ty); + let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 38, 1); + func.get_operation().deref_mut(&ctx).set_loc(func_loc); + + let entry = func.get_or_create_entry_block(&mut ctx); + let ret = ReturnOp::new(&mut ctx, None); + let ret_loc = src_location( + &mut ctx, + "/tmp/cuda-oxide/crates/cuda-device/src/thread.rs", + 292, + 19, + ); + ret.get_operation().deref_mut(&ctx).set_loc(ret_loc); + ret.get_operation().insert_at_back(entry, &ctx); + + func.get_operation().insert_at_back(module_block, &ctx); + + let config = DebugConfig { + inner: PtxExportConfig, + debug_kind: DebugKind::LineTables, + }; + let ir = + export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); + + let ret_line = ir + .lines() + .find(|line| line.trim_start().starts_with("ret void")) + .expect("return instruction"); + assert!( + ret_line.contains(", !dbg !"), + "cross-file instructions should keep source locations:\n{ir}" + ); + assert!( + ir.contains( + "!DIFile(filename: \"thread.rs\", directory: \"/tmp/cuda-oxide/crates/cuda-device/src\")" + ), + "cross-file locations should get their own DIFile:\n{ir}" + ); + assert!( + ir.contains("!DILexicalBlockFile(scope: !"), + "cross-file locations should use a file-specific debug scope:\n{ir}" + ); + assert!( + ir.contains("!DILocation(line: 292, column: 19, scope: !"), + "cross-file locations should preserve their real source line:\n{ir}" + ); +} + +#[test] +fn line_table_debug_metadata_emits_inlined_at_for_callsite_locations() { + let mut ctx = Context::new(); + + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_region = module.get_operation().deref(&ctx).get_region(0); + let module_block = { + let region = module_region.deref(&ctx); + region.iter(&ctx).next().unwrap() + }; + + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); + let func = FuncOp::new(&mut ctx, "debug_kernel".try_into().unwrap(), func_ty); + let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 38, 1); + func.get_operation().deref_mut(&ctx).set_loc(func_loc); + + let entry = func.get_or_create_entry_block(&mut ctx); + let ret = ReturnOp::new(&mut ctx, None); + let callee = src_location( + &mut ctx, + "/tmp/cuda-oxide/crates/cuda-device/src/thread.rs", + 292, + 19, + ); + let caller = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 39, 13); + ret.get_operation() + .deref_mut(&ctx) + .set_loc(Location::CallSite { + callee: Box::new(callee), + caller: Box::new(caller), + }); + ret.get_operation().insert_at_back(entry, &ctx); + + func.get_operation().insert_at_back(module_block, &ctx); + + let config = DebugConfig { + inner: PtxExportConfig, + debug_kind: DebugKind::LineTables, + }; + let ir = + export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); + + assert!( + ir.contains("!DILocation(line: 39, column: 13, scope: !"), + "callsite metadata should preserve the caller location:\n{ir}" + ); + assert!( + ir.contains("!DILocation(line: 292, column: 19, scope: !") && ir.contains(", inlinedAt: !"), + "callsite metadata should describe the callee location as inlined at the caller:\n{ir}" + ); +} + +#[test] +fn debug_locations_use_rustc_source_scope_positions() { + let mut ctx = Context::new(); + + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_region = module.get_operation().deref(&ctx).get_region(0); + let module_block = { + let region = module_region.deref(&ctx); + region.iter(&ctx).next().unwrap() + }; + + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); + let func = FuncOp::new(&mut ctx, "debug_kernel".try_into().unwrap(), func_ty); + let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 10, 1); + func.get_operation().deref_mut(&ctx).set_loc(func_loc); + llvm_export::ops::set_debug_source_scope_map( + &mut ctx, + func.get_operation(), + &DebugSourceScopeMap { + scopes: vec![ + DebugSourceScope { + id: 0, + parent: None, + span: Some(DebugSourcePosition { + file: PathBuf::from("/tmp/cuda-oxide/tests/kernel.rs"), + line: 10, + column: 1, + }), + inlined: None, + }, + DebugSourceScope { + id: 1, + parent: Some(0), + span: Some(DebugSourcePosition { + file: PathBuf::from("/tmp/cuda-oxide/tests/kernel.rs"), + line: 12, + column: 9, + }), + inlined: None, + }, + ], + locations: vec![DebugSourceScopeLocation { + pos: DebugSourcePosition { + file: PathBuf::from("/tmp/cuda-oxide/tests/kernel.rs"), + line: 12, + column: 9, + }, + scope: 1, + }], + }, + ); + + let entry = func.get_or_create_entry_block(&mut ctx); + let ret = ReturnOp::new(&mut ctx, None); + let ret_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 12, 9); + ret.get_operation().deref_mut(&ctx).set_loc(ret_loc); + ret.get_operation().insert_at_back(entry, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + let config = DebugConfig { + inner: PtxExportConfig, + debug_kind: DebugKind::LineTables, + }; + let ir = + export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); + + let block_id = ir + .lines() + .find_map(|line| { + if line.contains("!DILexicalBlock(scope: !") && line.contains("line: 12, column: 9") { + line.split_once(" = ") + .map(|(id, _)| id.trim_start_matches('!').to_string()) + } else { + None + } + }) + .expect("nested lexical block should be emitted"); + + assert!( + ir.contains(&format!( + "!DILocation(line: 12, column: 9, scope: !{block_id})" + )), + "instruction location should use the exact rustc source scope, not the function scope:\n{ir}" + ); +} + +#[test] +fn full_debug_metadata_emits_dbg_declare_for_tagged_allocas() { + let mut ctx = Context::new(); + + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_region = module.get_operation().deref(&ctx).get_region(0); + let module_block = { + let region = module_region.deref(&ctx); + region.iter(&ctx).next().unwrap() + }; + + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); + let func = FuncOp::new(&mut ctx, "debug_kernel".try_into().unwrap(), func_ty); + let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 30, 1); + func.get_operation().deref_mut(&ctx).set_loc(func_loc); + + let entry = func.get_or_create_entry_block(&mut ctx); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let one_attr = IntegerAttr::new(i32_ty, APInt::from_u32(1, NonZero::new(32).unwrap())); + let one = ConstantOp::new(&mut ctx, one_attr.into()); + one.get_operation().insert_at_back(entry, &ctx); + let one_val = one.get_operation().deref(&ctx).get_result(0); + + let tid = AllocaOp::new(&mut ctx, i32_ty.into(), one_val); + let tid_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 31, 9); + tid.get_operation().deref_mut(&ctx).set_loc(tid_loc); + llvm_export::ops::set_debug_local_variable( + &mut ctx, + tid.get_operation(), + DebugLocalVariableInfo { + name: "tid".to_string(), + argument_index: Some(1), + ty: DebugLocalTypeKind::Basic { + name: "u32".to_string(), + size_bits: 32, + encoding: "DW_ATE_unsigned", + }, + }, + ); + tid.get_operation().insert_at_back(entry, &ctx); + + let ptr_ty = PointerType::get(&ctx, 0); + let ptr = AllocaOp::new(&mut ctx, ptr_ty.into(), one_val); + let ptr_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 32, 9); + ptr.get_operation().deref_mut(&ctx).set_loc(ptr_loc); + llvm_export::ops::set_debug_local_variable( + &mut ctx, + ptr.get_operation(), + DebugLocalVariableInfo { + name: "ptr".to_string(), + argument_index: None, + ty: DebugLocalTypeKind::TypedPointer { + name: "*mut f32".to_string(), + size_bits: 64, + pointee: Box::new(DebugLocalTypeKind::Basic { + name: "f32".to_string(), + size_bits: 32, + encoding: "DW_ATE_float", + }), + }, + }, + ); + ptr.get_operation().insert_at_back(entry, &ctx); + + let nested_ptr = AllocaOp::new(&mut ctx, ptr_ty.into(), one_val); + let nested_ptr_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 33, 9); + nested_ptr + .get_operation() + .deref_mut(&ctx) + .set_loc(nested_ptr_loc); + llvm_export::ops::set_debug_local_variable( + &mut ctx, + nested_ptr.get_operation(), + DebugLocalVariableInfo { + name: "nested_ptr".to_string(), + argument_index: None, + ty: DebugLocalTypeKind::TypedPointer { + name: "*const *mut i32".to_string(), + size_bits: 64, + pointee: Box::new(DebugLocalTypeKind::TypedPointer { + name: "*mut i32".to_string(), + size_bits: 64, + pointee: Box::new(DebugLocalTypeKind::Basic { + name: "i32".to_string(), + size_bits: 32, + encoding: "DW_ATE_signed", + }), + }), + }, + }, + ); + nested_ptr.get_operation().insert_at_back(entry, &ctx); + + let array_ptr = AllocaOp::new(&mut ctx, ptr_ty.into(), one_val); + let array_ptr_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 34, 9); + array_ptr + .get_operation() + .deref_mut(&ctx) + .set_loc(array_ptr_loc); + llvm_export::ops::set_debug_local_variable( + &mut ctx, + array_ptr.get_operation(), + DebugLocalVariableInfo { + name: "array_ptr".to_string(), + argument_index: None, + ty: DebugLocalTypeKind::TypedPointer { + name: "*const [u16; 4]".to_string(), + size_bits: 64, + pointee: Box::new(DebugLocalTypeKind::Array { + name: "[u16; 4]".to_string(), + size_bits: 64, + element: Box::new(DebugLocalTypeKind::Basic { + name: "u16".to_string(), + size_bits: 16, + encoding: "DW_ATE_unsigned", + }), + count: 4, + }), + }, + }, + ); + array_ptr.get_operation().insert_at_back(entry, &ctx); + + let opaque_ptr = AllocaOp::new(&mut ctx, ptr_ty.into(), one_val); + let opaque_ptr_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 35, 9); + opaque_ptr + .get_operation() + .deref_mut(&ctx) + .set_loc(opaque_ptr_loc); + llvm_export::ops::set_debug_local_variable( + &mut ctx, + opaque_ptr.get_operation(), + DebugLocalVariableInfo { + name: "opaque_ptr".to_string(), + argument_index: None, + ty: DebugLocalTypeKind::Pointer { + name: "*const _".to_string(), + size_bits: 64, + }, + }, + ); + opaque_ptr.get_operation().insert_at_back(entry, &ctx); + + let ret = ReturnOp::new(&mut ctx, None); + let ret_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 36, 1); + ret.get_operation().deref_mut(&ctx).set_loc(ret_loc); + ret.get_operation().insert_at_back(entry, &ctx); + + func.get_operation().insert_at_back(module_block, &ctx); + + let config = DebugConfig { + inner: PtxExportConfig, + debug_kind: DebugKind::Full, + }; + let ir = + export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); + + assert!( + ir.contains("emissionKind: FullDebug"), + "full debug should request full DWARF metadata:\n{ir}" + ); + assert!( + ir.contains("isOptimized: false"), + "full debug export should describe the unoptimized debug path:\n{ir}" + ); + assert!( + ir.contains("declare void @llvm.dbg.declare(metadata, metadata, metadata)"), + "full debug should declare the debug intrinsic it calls:\n{ir}" + ); + assert!( + ir.contains("call void @llvm.dbg.declare(metadata ptr %"), + "tagged allocas should be bound to variables with dbg.declare:\n{ir}" + ); + assert!( + ir.contains("!DILocalVariable(name: \"tid\", arg: 1, scope: !"), + "argument debug metadata should preserve the argument number:\n{ir}" + ); + assert!( + ir.contains("!DILocalVariable(name: \"ptr\", scope: !"), + "local debug metadata should omit the arg field:\n{ir}" + ); + assert!( + ir.contains("!DIBasicType(name: \"u32\", size: 32, encoding: DW_ATE_unsigned)"), + "basic integer variables should get DIBasicType metadata:\n{ir}" + ); + let f32_id = metadata_id( + &ir, + "!DIBasicType(name: \"f32\", size: 32, encoding: DW_ATE_float)", + ); + assert!( + ir.contains(&format!( + "!DIDerivedType(tag: DW_TAG_pointer_type, name: \"*mut f32\", baseType: {f32_id}, size: 64)" + )), + "pointer variables should reference their pointee DIType:\n{ir}" + ); + let i32_id = metadata_id( + &ir, + "!DIBasicType(name: \"i32\", size: 32, encoding: DW_ATE_signed)", + ); + let inner_pointer = format!( + "!DIDerivedType(tag: DW_TAG_pointer_type, name: \"*mut i32\", baseType: {i32_id}, size: 64)" + ); + let inner_pointer_id = metadata_id(&ir, &inner_pointer); + assert!( + ir.contains(&format!( + "!DIDerivedType(tag: DW_TAG_pointer_type, name: \"*const *mut i32\", baseType: {inner_pointer_id}, size: 64)" + )), + "nested pointer metadata must preserve the complete base-type chain:\n{ir}" + ); + + let u16_id = metadata_id( + &ir, + "!DIBasicType(name: \"u16\", size: 16, encoding: DW_ATE_unsigned)", + ); + let subrange_id = metadata_id(&ir, "!DISubrange(count: 4)"); + let subrange_tuple_id = metadata_id(&ir, &format!("!{{{subrange_id}}}")); + let array_type = format!( + "!DICompositeType(tag: DW_TAG_array_type, baseType: {u16_id}, size: 64, elements: {subrange_tuple_id})" + ); + let array_type_id = metadata_id(&ir, &array_type); + assert!( + ir.contains(&format!( + "!DIDerivedType(tag: DW_TAG_pointer_type, name: \"*const [u16; 4]\", baseType: {array_type_id}, size: 64)" + )), + "array pointer metadata must reference the exact element/subrange graph:\n{ir}" + ); + + assert!( + ir.contains("!DILocalVariable(name: \"opaque_ptr\", scope: !"), + "unsupported pointers must remain visible as source variables:\n{ir}" + ); + assert!( + ir.contains( + "!DIDerivedType(tag: DW_TAG_pointer_type, name: \"*const _\", baseType: null, size: 64)" + ), + "the compatibility pointer must retain its legacy null-base metadata:\n{ir}" + ); + let null_base_pointer_count = ir + .lines() + .filter(|line| line.contains("DW_TAG_pointer_type") && line.contains("baseType: null")) + .count(); + assert!( + null_base_pointer_count == 1, + "only the explicit opaque compatibility pointer may have a null baseType:\n{ir}" + ); + + let Some(llvm_as) = ["llvm-as-22", "llvm-as-21", "llvm-as"] + .into_iter() + .find(|tool| { + std::process::Command::new(tool) + .arg("--version") + .output() + .is_ok_and(|out| out.status.success()) + }) + else { + eprintln!("skipping pointer-debug llvm-as parse gate: no supported llvm-as on PATH"); + return; + }; + let ll_path = std::env::temp_dir().join(format!( + "cuda_oxide_pointer_debug_parse_gate_{}.ll", + std::process::id() + )); + std::fs::write(&ll_path, &ir).expect("write pointer-debug temp .ll"); + let output = std::process::Command::new(llvm_as) + .arg("-o") + .arg("/dev/null") + .arg(&ll_path) + .output() + .expect("run llvm-as for pointer debug metadata"); + let stderr = String::from_utf8_lossy(&output.stderr).into_owned(); + let _ = std::fs::remove_file(&ll_path); + assert!( + output.status.success(), + "{llvm_as} rejected pointer debug metadata:\n{stderr}\n--- module ---\n{ir}" + ); +} + +#[test] +fn full_debug_metadata_describes_as1_global_with_semantic_rust_type() { + let mut ctx = Context::new(); + + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signless); + let storage_ty = ArrayType::get(&ctx, i8_ty.into(), 8); + let global = GlobalOp::new_with_alignment( + &mut ctx, + "__device_global_7".try_into().unwrap(), + storage_ty.into(), + 8, + ); + global.set_address_space(&mut ctx, llvm_export::types::address_space::GLOBAL); + global.set_initializer_hex(&mut ctx, "0000000000000000"); + llvm_export::ops::set_debug_global_variable( + &mut ctx, + global.get_operation(), + &DebugGlobalVariableInfo { + name: "GLOBAL_COUNTER".to_string(), + namespace: vec!["debuginfo".to_string(), "state".to_string()], + ty: DebugLocalTypeKind::Basic { + name: "u64".to_string(), + size_bits: 64, + encoding: "DW_ATE_unsigned", + }, + declaration: DebugSourcePosition { + file: PathBuf::from("/tmp/cuda-oxide/tests/kernel.rs"), + line: 90, + column: 1, + }, + is_local_to_unit: true, + is_function_local: false, + }, + ); + global.get_operation().insert_at_back(module_block, &ctx); + + let full = export_module_to_string_with_config( + &ctx, + &module, + &DebugConfig { + inner: PtxExportConfig, + debug_kind: DebugKind::Full, + }, + ) + .expect("full debug export succeeds"); + + let definition = full + .lines() + .find(|line| line.starts_with("@__device_global_7 = ")) + .expect("device-global definition"); + assert!( + definition.contains("addrspace(1) global [8 x i8] c\"\\00\\00\\00\\00\\00\\00\\00\\00\"") + && definition.contains(", align 8, !dbg !"), + "the physical byte storage should retain a global debug attachment:\n{full}" + ); + assert!( + full.contains( + "distinct !DIGlobalVariable(name: \"GLOBAL_COUNTER\", linkageName: \"__device_global_7\"" + ), + "DWARF should separate the source name from the generated linkage name:\n{full}" + ); + assert!( + full.contains("!DINamespace(name: \"debuginfo\", scope: null)") + && full.contains("!DINamespace(name: \"state\", scope: !") + && full.contains("scope: !"), + "the source crate/module hierarchy should scope the leaf name:\n{full}" + ); + assert!( + full.contains("isLocal: true, isDefinition: true"), + "private Rust statics should remain local to the compile unit:\n{full}" + ); + assert!( + full.contains("file: !") + && full.contains("line: 90, type: !") + && full.contains("align: 64)"), + "the declaration location and source alignment should be retained:\n{full}" + ); + assert!( + full.contains("!DIBasicType(name: \"u64\", size: 64, encoding: DW_ATE_unsigned)"), + "the debug type must be semantic u64, not the physical [8 x i8] storage:\n{full}" + ); + assert!( + full.contains("!DIGlobalVariableExpression(var: !") + && full.contains("expr: !DIExpression())") + && full.contains("emissionKind: FullDebug, globals: !"), + "the global expression should be retained by the compile unit:\n{full}" + ); + + let off = export_module_to_string_with_config( + &ctx, + &module, + &DebugConfig { + inner: PtxExportConfig, + debug_kind: DebugKind::Off, + }, + ) + .expect("non-debug export succeeds"); + let line_tables = export_module_to_string_with_config( + &ctx, + &module, + &DebugConfig { + inner: PtxExportConfig, + debug_kind: DebugKind::LineTables, + }, + ) + .expect("line-table export succeeds"); + assert_eq!( + line_tables, off, + "a global-only module must not gain variable metadata outside full debug" + ); +} + +#[test] +fn full_debug_metadata_describes_function_local_as3_array_with_shared_address_class() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "shared_debug".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let raw_owner = reserved_oxide_symbols::device_symbol("shared_kernel"); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let storage_ty = ArrayType::get(&ctx, i32_ty.into(), 32); + let global = GlobalOp::new_with_alignment( + &mut ctx, + "__shared_mem_0".try_into().unwrap(), + storage_ty.into(), + 4, + ); + global.set_address_space(&mut ctx, llvm_export::types::address_space::SHARED); + let info = DebugGlobalVariableInfo { + name: "TILE".to_string(), + namespace: vec![ + "debuginfo".to_string(), + "kernels".to_string(), + "shared_kernel".to_string(), + ], + ty: DebugLocalTypeKind::Array { + name: "[i32; 32]".to_string(), + size_bits: 1024, + element: Box::new(DebugLocalTypeKind::Basic { + name: "i32".to_string(), + size_bits: 32, + encoding: "DW_ATE_signed", + }), + count: 32, + }, + declaration: DebugSourcePosition { + file: PathBuf::from("/tmp/cuda-oxide/tests/shared.rs"), + line: 40, + column: 9, + }, + is_local_to_unit: true, + is_function_local: true, + }; + llvm_export::ops::set_debug_global_variable(&mut ctx, global.get_operation(), &info); + llvm_export::ops::set_debug_global_owner_function(&mut ctx, global.get_operation(), &raw_owner); + global.get_operation().insert_at_back(module_block, &ctx); + + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); + let function = FuncOp::new(&mut ctx, raw_owner.as_str().try_into().unwrap(), func_ty); + let function_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/shared.rs", 35, 1); + function + .get_operation() + .deref_mut(&ctx) + .set_loc(function_loc); + let entry = function.get_or_create_entry_block(&mut ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + function.get_operation().insert_at_back(module_block, &ctx); + + let full = export_module_to_string_with_config( + &ctx, + &module, + &DebugConfig { + inner: PtxExportConfig, + debug_kind: DebugKind::Full, + }, + ) + .expect("full shared debug export succeeds"); + + let definition = full + .lines() + .find(|line| line.starts_with("@__shared_mem_0 = ")) + .expect("shared definition"); + assert!( + definition.contains("addrspace(3) global [32 x i32]") + && definition.contains(", align 4, !dbg !"), + "physical AS3 array must retain the debug attachment:\n{full}" + ); + assert!( + full.contains("!DIExpression(DW_OP_constu, 8, DW_OP_swap, DW_OP_xderef)"), + "AS3 must carry CUDA DWARF shared address class 8:\n{full}" + ); + assert!( + full.contains("!DINamespace(name: \"debuginfo\", scope: null)") + && full.contains("!DINamespace(name: \"kernels\", scope: !") + && full.contains( + "distinct !DISubprogram(name: \"shared_kernel\", \ + linkageName: \"shared_kernel\", scope: !" + ), + "the owning subprogram must be nested under the structured namespace:\n{full}" + ); + assert_eq!( + full.matches("distinct !DISubprogram(name: \"shared_kernel\"") + .count(), + 1, + "pre-reservation and function export must reuse one DISubprogram:\n{full}" + ); + assert!( + full.contains("distinct !DIGlobalVariable(name: \"TILE\", linkageName: \"__shared_mem_0\"") + && full.contains("!DISubrange(count: 32)") + && full.contains("!DICompositeType(tag: DW_TAG_array_type, baseType: !") + && full.contains("size: 1024, elements: !"), + "the DIE must keep leaf/linkage and the logical array type:\n{full}" + ); + assert!( + full.contains("emissionKind: FullDebug, globals: !"), + "the CU must retain the one shared expression:\n{full}" + ); +} + +fn export_mixed_shared_owner_order(malformed_first: bool) -> String { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "shared_owner_validation".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let raw_owner = reserved_oxide_symbols::device_symbol("shared_kernel"); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let storage_ty = ArrayType::get(&ctx, i32_ty.into(), 4); + + let valid = GlobalOp::new_with_alignment( + &mut ctx, + "__shared_mem_valid".try_into().unwrap(), + storage_ty.into(), + 4, + ); + valid.set_address_space(&mut ctx, llvm_export::types::address_space::SHARED); + llvm_export::ops::set_debug_global_variable( + &mut ctx, + valid.get_operation(), + &DebugGlobalVariableInfo { + name: "VALID".to_string(), + namespace: vec!["owner_fixture".to_string(), "shared_kernel".to_string()], + ty: DebugLocalTypeKind::Array { + name: "[i32; 4]".to_string(), + size_bits: 128, + element: Box::new(DebugLocalTypeKind::Basic { + name: "i32".to_string(), + size_bits: 32, + encoding: "DW_ATE_signed", + }), + count: 4, + }, + declaration: DebugSourcePosition { + file: PathBuf::from("/tmp/cuda-oxide/tests/shared-owner.rs"), + line: 20, + column: 5, + }, + is_local_to_unit: true, + is_function_local: true, + }, + ); + llvm_export::ops::set_debug_global_owner_function(&mut ctx, valid.get_operation(), &raw_owner); + + let malformed = GlobalOp::new_with_alignment( + &mut ctx, + "__shared_mem_malformed".try_into().unwrap(), + storage_ty.into(), + 4, + ); + malformed.set_address_space(&mut ctx, llvm_export::types::address_space::SHARED); + llvm_export::ops::set_debug_global_variable( + &mut ctx, + malformed.get_operation(), + &DebugGlobalVariableInfo { + name: "MALFORMED".to_string(), + namespace: vec!["owner_fixture".to_string(), "shared_kernel".to_string()], + ty: DebugLocalTypeKind::Array { + name: "[i32; 4]".to_string(), + size_bits: 128, + element: Box::new(DebugLocalTypeKind::Basic { + name: "i32".to_string(), + size_bits: 32, + encoding: "DW_ATE_signed", + }), + count: 4, + }, + declaration: DebugSourcePosition { + file: PathBuf::from("/tmp/cuda-oxide/tests/shared-owner.rs"), + line: 24, + column: 5, + }, + is_local_to_unit: true, + is_function_local: true, + }, + ); + // This raw spelling normalizes to `shared_kernel`, but it is not the raw + // symbol indexed for that function definition. It must not borrow the + // valid sibling's pre-reserved DISubprogram. + llvm_export::ops::set_debug_global_owner_function( + &mut ctx, + malformed.get_operation(), + "shared_kernel", + ); + + if malformed_first { + malformed.get_operation().insert_at_back(module_block, &ctx); + valid.get_operation().insert_at_back(module_block, &ctx); + } else { + valid.get_operation().insert_at_back(module_block, &ctx); + malformed.get_operation().insert_at_back(module_block, &ctx); + } + + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); + let function = FuncOp::new(&mut ctx, raw_owner.as_str().try_into().unwrap(), func_ty); + let function_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/shared-owner.rs", 16, 1); + function + .get_operation() + .deref_mut(&ctx) + .set_loc(function_loc); + let entry = function.get_or_create_entry_block(&mut ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + function.get_operation().insert_at_back(module_block, &ctx); + + export_module_to_string_with_config( + &ctx, + &module, + &DebugConfig { + inner: PtxExportConfig, + debug_kind: DebugKind::Full, + }, + ) + .expect("a malformed shared owner must not break its valid sibling") +} + +fn assert_malformed_shared_owner_fails_closed(full: &str) { + let valid_definition = full + .lines() + .find(|line| line.starts_with("@__shared_mem_valid = ")) + .expect("valid shared definition"); + let malformed_definition = full + .lines() + .find(|line| line.starts_with("@__shared_mem_malformed = ")) + .expect("malformed-owner shared definition"); + assert!( + valid_definition.contains(", !dbg !"), + "the valid sibling must retain its debug attachment:\n{full}" + ); + assert!( + !malformed_definition.contains(", !dbg !"), + "the malformed owner must not borrow the valid sibling's scope:\n{full}" + ); + assert!( + full.contains("distinct !DIGlobalVariable(name: \"VALID\"") + && !full.contains("distinct !DIGlobalVariable(name: \"MALFORMED\""), + "only the valid shared source identity may become a DIE:\n{full}" + ); + assert_eq!( + full.matches("!DIGlobalVariableExpression(var:").count(), + 1, + "only the valid sibling may be retained by the compile unit:\n{full}" + ); + assert_eq!( + full.matches("distinct !DISubprogram(name: \"shared_kernel\"") + .count(), + 1, + "the valid owner must still reuse exactly one DISubprogram:\n{full}" + ); +} + +#[test] +fn malformed_shared_owner_before_valid_sibling_cannot_borrow_its_scope() { + assert_malformed_shared_owner_fails_closed(&export_mixed_shared_owner_order(true)); +} + +#[test] +fn malformed_shared_owner_after_valid_sibling_cannot_borrow_its_scope() { + assert_malformed_shared_owner_fails_closed(&export_mixed_shared_owner_order(false)); +} + +#[test] +fn full_debug_globals_preserve_qualified_identity_visibility_and_relocations() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "global_debug_adversarial".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signless); + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); + + let left_ty = ArrayType::get(&ctx, i8_ty.into(), 4); + let left = GlobalOp::new_with_alignment( + &mut ctx, + "__device_global_20".try_into().unwrap(), + left_ty.into(), + 4, + ); + left.set_address_space(&mut ctx, llvm_export::types::address_space::GLOBAL); + left.set_initializer_hex(&mut ctx, "01000000"); + llvm_export::ops::set_debug_global_variable( + &mut ctx, + left.get_operation(), + &DebugGlobalVariableInfo { + name: "SAME_LEAF".to_string(), + namespace: vec!["qualified_fixture".to_string(), "left".to_string()], + ty: DebugLocalTypeKind::Basic { + name: "u32".to_string(), + size_bits: 32, + encoding: "DW_ATE_unsigned", + }, + declaration: DebugSourcePosition { + file: PathBuf::from("/tmp/cuda-oxide/tests/qualified.rs"), + line: 20, + column: 5, + }, + is_local_to_unit: true, + is_function_local: false, + }, + ); + left.get_operation().insert_at_back(module_block, &ctx); + + let right_ty = ArrayType::get(&ctx, i8_ty.into(), 8); + let right = GlobalOp::new_with_alignment( + &mut ctx, + "__device_global_21".try_into().unwrap(), + right_ty.into(), + 8, + ); + right.set_address_space(&mut ctx, llvm_export::types::address_space::GLOBAL); + right.set_initializer_hex(&mut ctx, "0200000000000000"); + llvm_export::ops::set_debug_global_variable( + &mut ctx, + right.get_operation(), + &DebugGlobalVariableInfo { + name: "SAME_LEAF".to_string(), + namespace: vec!["qualified_fixture".to_string(), "right".to_string()], + ty: DebugLocalTypeKind::Basic { + name: "u64".to_string(), + size_bits: 64, + encoding: "DW_ATE_unsigned", + }, + declaration: DebugSourcePosition { + file: PathBuf::from("/tmp/cuda-oxide/tests/qualified.rs"), + line: 24, + column: 5, + }, + is_local_to_unit: false, + is_function_local: false, + }, + ); + right.get_operation().insert_at_back(module_block, &ctx); + + let target_ty = ArrayType::get(&ctx, i8_ty.into(), 4); + let target = GlobalOp::new_with_alignment( + &mut ctx, + "__device_global_22".try_into().unwrap(), + target_ty.into(), + 4, + ); + target.set_address_space(&mut ctx, llvm_export::types::address_space::GLOBAL); + target.set_source_global_key(&mut ctx, "qualified_fixture::RELOCATION_TARGET"); + target.set_initializer_hex(&mut ctx, "78563412"); + llvm_export::ops::set_debug_global_variable( + &mut ctx, + target.get_operation(), + &DebugGlobalVariableInfo { + name: "RELOCATION_TARGET".to_string(), + namespace: vec!["qualified_fixture".to_string()], + ty: DebugLocalTypeKind::Basic { + name: "u32".to_string(), + size_bits: 32, + encoding: "DW_ATE_unsigned", + }, + declaration: DebugSourcePosition { + file: PathBuf::from("/tmp/cuda-oxide/tests/qualified.rs"), + line: 30, + column: 1, + }, + is_local_to_unit: true, + is_function_local: false, + }, + ); + target.get_operation().insert_at_back(module_block, &ctx); + + let references_ty = StructType::get_unnamed( + &ctx, + (vec![i64_ty.into(), i64_ty.into()], StructLayout::Unpacked), + ); + let references = GlobalOp::new_with_alignment( + &mut ctx, + "__device_global_23".try_into().unwrap(), + references_ty.into(), + 8, + ); + references.set_address_space(&mut ctx, llvm_export::types::address_space::GLOBAL); + references.set_source_global_key(&mut ctx, "qualified_fixture::REFERENCES"); + references.set_initializer_hex(&mut ctx, "00000000000000000000000000000000"); + references.set_initializer_relocations( + &mut ctx, + &encode_global_initializer_relocations(&[ + GlobalInitializerRelocation { + source_offset: 0, + width_bytes: 8, + target_address_space: llvm_export::types::address_space::GLOBAL, + target_addend: 0, + target_key: "qualified_fixture::RELOCATION_TARGET".to_string(), + }, + GlobalInitializerRelocation { + source_offset: 8, + width_bytes: 8, + target_address_space: llvm_export::types::address_space::GLOBAL, + target_addend: 0, + target_key: "qualified_fixture::RELOCATION_TARGET".to_string(), + }, + ]), + ); + llvm_export::ops::set_debug_global_variable( + &mut ctx, + references.get_operation(), + &DebugGlobalVariableInfo { + name: "REFERENCES".to_string(), + namespace: vec!["qualified_fixture".to_string()], + ty: DebugLocalTypeKind::Array { + name: "[&u32; 2]".to_string(), + size_bits: 128, + element: Box::new(DebugLocalTypeKind::Pointer { + name: "&u32".to_string(), + size_bits: 64, + }), + count: 2, + }, + declaration: DebugSourcePosition { + file: PathBuf::from("/tmp/cuda-oxide/tests/qualified.rs"), + line: 31, + column: 1, + }, + is_local_to_unit: true, + is_function_local: false, + }, + ); + references + .get_operation() + .insert_at_back(module_block, &ctx); + + let export = |debug_kind| { + export_module_to_string_with_config( + &ctx, + &module, + &DebugConfig { + inner: PtxExportConfig, + debug_kind, + }, + ) + .expect("global debug export succeeds") + }; + let full = export(DebugKind::Full); + + let metadata_id = |needle: &str| { + full.lines() + .find(|line| line.contains(needle)) + .unwrap_or_else(|| panic!("missing `{needle}` in:\n{full}")) + .split_once(" = ") + .expect("metadata definition") + .0 + .trim_start_matches('!') + .to_string() + }; + let left_scope = metadata_id("!DINamespace(name: \"left\""); + let right_scope = metadata_id("!DINamespace(name: \"right\""); + let left_variable = full + .lines() + .find(|line| { + line.contains("!DIGlobalVariable(name: \"SAME_LEAF\"") + && line.contains("linkageName: \"__device_global_20\"") + }) + .expect("left source global"); + let right_variable = full + .lines() + .find(|line| { + line.contains("!DIGlobalVariable(name: \"SAME_LEAF\"") + && line.contains("linkageName: \"__device_global_21\"") + }) + .expect("right source global"); + assert!(left_variable.contains(&format!("scope: !{left_scope}"))); + assert!(left_variable.contains("line: 20") && left_variable.contains("isLocal: true")); + assert!(right_variable.contains(&format!("scope: !{right_scope}"))); + assert!(right_variable.contains("line: 24") && right_variable.contains("isLocal: false")); + + let target_definition = full + .lines() + .find(|line| line.starts_with("@__device_global_22 = ")) + .expect("relocation-only target definition"); + assert!(target_definition.contains("[4 x i8]") && target_definition.contains("!dbg !")); + let target_variable = full + .lines() + .find(|line| line.contains("!DIGlobalVariable(name: \"RELOCATION_TARGET\"")) + .expect("relocation-only target source identity"); + assert!( + target_variable.contains("linkageName: \"__device_global_22\"") + && target_variable.contains("line: 30") + ); + + let reference_definition = full + .lines() + .find(|line| line.starts_with("@__device_global_23 = ")) + .expect("relocation-backed definition"); + assert_eq!( + reference_definition.matches("@__device_global_22").count(), + 2, + "both initializer references must target the same one physical global" + ); + assert!( + full.contains("!DICompositeType(tag: DW_TAG_array_type, baseType: !") + && full.contains("size: 128, elements: !") + && full.contains("!DIDerivedType(tag: DW_TAG_pointer_type, name: \"&u32\"") + ); + + assert_eq!( + full.matches("!DIGlobalVariableExpression(var:").count(), + 4, + "one expression per physical source global, despite repeated relocations" + ); + let compile_unit = full + .lines() + .find(|line| line.contains("!DICompileUnit(")) + .expect("debug compile unit"); + let globals_id = compile_unit + .split("globals: !") + .nth(1) + .and_then(|tail| tail.strip_suffix(')')) + .expect("compile-unit globals tuple"); + let globals_tuple = full + .lines() + .find(|line| line.starts_with(&format!("!{globals_id} = !{{"))) + .expect("compile-unit globals tuple definition"); + assert_eq!(globals_tuple.matches('!').count() - 2, 4, "{globals_tuple}"); + + assert_eq!( + export(DebugKind::LineTables), + export(DebugKind::Off), + "initialized and relocation-backed globals must not gain metadata outside Full debug" + ); + + let Some(llvm_as) = ["llvm-as-22", "llvm-as-21", "llvm-as"] + .into_iter() + .find(|tool| { + std::process::Command::new(tool) + .arg("--version") + .output() + .is_ok_and(|out| out.status.success()) + }) + else { + eprintln!("skipping global-debug parse gate: no llvm-as on PATH"); + return; + }; + let ll_path = std::env::temp_dir().join(format!( + "cuda_oxide_global_debug_parse_gate_{}.ll", + std::process::id() + )); + std::fs::write(&ll_path, &full).expect("write temp .ll"); + let output = std::process::Command::new(llvm_as) + .arg("-o") + .arg("/dev/null") + .arg(&ll_path) + .output() + .expect("run llvm-as"); + let stderr = String::from_utf8_lossy(&output.stderr).into_owned(); + let _ = std::fs::remove_file(&ll_path); + assert!( + output.status.success(), + "{llvm_as} rejected global debug metadata:\n{stderr}\n--- module ---\n{full}" + ); +} + +#[test] +fn full_debug_metadata_emits_rust_enum_variant_parts() { + let mut ctx = Context::new(); + + let module = ModuleOp::new(&mut ctx, "enum_debug".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); + let func = FuncOp::new(&mut ctx, "enum_debug_kernel".try_into().unwrap(), func_ty); + let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/enum.rs", 10, 1); + func.get_operation().deref_mut(&ctx).set_loc(func_loc); + let entry = func.get_or_create_entry_block(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); + let one_attr = IntegerAttr::new(i32_ty, APInt::from_u32(1, NonZero::new(32).unwrap())); + let one = ConstantOp::new(&mut ctx, one_attr.into()); + one.get_operation().insert_at_back(entry, &ctx); + let one_val = one.get_operation().deref(&ctx).get_result(0); + + let direct = AllocaOp::new(&mut ctx, i64_ty.into(), one_val); + let direct_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/enum.rs", 11, 9); + direct.get_operation().deref_mut(&ctx).set_loc(direct_loc); + llvm_export::ops::set_debug_local_variable( + &mut ctx, + direct.get_operation(), + DebugLocalVariableInfo { + name: "direct".to_string(), + argument_index: None, + ty: DebugLocalTypeKind::Enum { + name: "Direct".to_string(), + size_bits: 64, + discriminant: Some(DebugEnumDiscriminant { + offset_bits: 0, + ty: Box::new(DebugLocalTypeKind::Basic { + name: "u8".to_string(), + size_bits: 8, + encoding: "DW_ATE_unsigned", + }), + }), + variants: vec![ + DebugEnumVariant { + name: "Small".to_string(), + discriminant: Some(3), + members: vec![llvm_export::ops::DebugTypeMember { + name: "0".to_string(), + offset_bits: 32, + ty: DebugLocalTypeKind::Basic { + name: "u32".to_string(), + size_bits: 32, + encoding: "DW_ATE_unsigned", + }, + }], + }, + DebugEnumVariant { + name: "Empty".to_string(), + discriminant: Some(9), + members: vec![], + }, + ], + }, + }, + ); + direct.get_operation().insert_at_back(entry, &ctx); + + let niche = AllocaOp::new(&mut ctx, i64_ty.into(), one_val); + let niche_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/enum.rs", 12, 9); + niche.get_operation().deref_mut(&ctx).set_loc(niche_loc); + llvm_export::ops::set_debug_local_variable( + &mut ctx, + niche.get_operation(), + DebugLocalVariableInfo { + name: "niche".to_string(), + argument_index: None, + ty: DebugLocalTypeKind::Enum { + name: "OptionRef".to_string(), + size_bits: 64, + discriminant: Some(DebugEnumDiscriminant { + offset_bits: 0, + ty: Box::new(DebugLocalTypeKind::Basic { + name: "usize".to_string(), + size_bits: 64, + encoding: "DW_ATE_unsigned", + }), + }), + variants: vec![ + DebugEnumVariant { + name: "None".to_string(), + discriminant: Some(0), + members: vec![], + }, + DebugEnumVariant { + name: "Some".to_string(), + discriminant: None, + members: vec![llvm_export::ops::DebugTypeMember { + name: "0".to_string(), + offset_bits: 0, + ty: DebugLocalTypeKind::TypedPointer { + name: "&u32".to_string(), + size_bits: 64, + pointee: Box::new(DebugLocalTypeKind::Basic { + name: "u32".to_string(), + size_bits: 32, + encoding: "DW_ATE_unsigned", + }), + }, + }], + }, + ], + }, + }, + ); + niche.get_operation().insert_at_back(entry, &ctx); + + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + let config = DebugConfig { + inner: PtxExportConfig, + debug_kind: DebugKind::Full, + }; + let ir = + export_module_to_string_with_config(&ctx, &module, &config).expect("enum debug export"); + + assert!( + ir.contains("!DICompositeType(tag: DW_TAG_variant_part"), + "Rust enums must contain a DW_TAG_variant_part:\n{ir}" + ); + assert!( + ir.contains("flags: DIFlagArtificial"), + "the physical enum carrier must be marked artificial:\n{ir}" + ); + let discriminator_member = ir + .lines() + .find(|line| { + line.contains("!DIDerivedType(tag: DW_TAG_member") + && line.contains("flags: DIFlagArtificial") + }) + .expect("enum discriminator member"); + assert!( + discriminator_member.contains("scope: !"), + "the physical enum carrier must be scoped to its enum object:\n{discriminator_member}\n{ir}" + ); + let variant_part = ir + .lines() + .find(|line| line.contains("!DICompositeType(tag: DW_TAG_variant_part")) + .expect("enum variant part"); + assert!( + variant_part.contains("scope: !"), + "the variant part must be scoped to its enum object:\n{variant_part}\n{ir}" + ); + let direct_variant_member = ir + .lines() + .find(|line| { + line.contains("!DIDerivedType(tag: DW_TAG_member, name: \"Small\"") + && line.contains("extraData: i8 3") + }) + .expect("direct enum variant member"); + assert!( + direct_variant_member.contains("scope: !"), + "variant members must be scoped to their DW_TAG_variant_part:\n{direct_variant_member}\n{ir}" + ); + assert!( + ir.contains("extraData: i8 3") && ir.contains("extraData: i8 9"), + "direct-tag variants must carry their physical discriminant values:\n{ir}" + ); + assert!( + ir.contains("extraData: i64 0"), + "the tagged niche variant must carry the niche value:\n{ir}" + ); + + let some_variant_member = ir + .lines() + .find(|line| { + line.contains("!DIDerivedType(tag: DW_TAG_member, name: \"Some\"") + && line.contains("baseType:") + }) + .expect("Some variant member"); + assert!( + !some_variant_member.contains("extraData:"), + "the untagged niche variant must be the default branch:\n{some_variant_member}\n{ir}" + ); +} + +#[test] +fn full_debug_metadata_uses_file_scope_for_cross_file_local_variables() { + let mut ctx = Context::new(); + + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_region = module.get_operation().deref(&ctx).get_region(0); + let module_block = { + let region = module_region.deref(&ctx); + region.iter(&ctx).next().unwrap() + }; + + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); + let func = FuncOp::new(&mut ctx, "debug_kernel".try_into().unwrap(), func_ty); + let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 30, 1); + func.get_operation().deref_mut(&ctx).set_loc(func_loc); + + let entry = func.get_or_create_entry_block(&mut ctx); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let one_attr = IntegerAttr::new(i32_ty, APInt::from_u32(1, NonZero::new(32).unwrap())); + let one = ConstantOp::new(&mut ctx, one_attr.into()); + one.get_operation().insert_at_back(entry, &ctx); + let one_val = one.get_operation().deref(&ctx).get_result(0); + + let tid = AllocaOp::new(&mut ctx, i32_ty.into(), one_val); + let tid_loc = src_location( + &mut ctx, + "/tmp/cuda-oxide/crates/cuda-device/src/thread.rs", + 292, + 19, + ); + tid.get_operation().deref_mut(&ctx).set_loc(tid_loc); + llvm_export::ops::set_debug_local_variable( + &mut ctx, + tid.get_operation(), + DebugLocalVariableInfo { + name: "tid".to_string(), + argument_index: None, + ty: DebugLocalTypeKind::Basic { + name: "u32".to_string(), + size_bits: 32, + encoding: "DW_ATE_unsigned", + }, + }, + ); + tid.get_operation().insert_at_back(entry, &ctx); + + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + let config = DebugConfig { + inner: PtxExportConfig, + debug_kind: DebugKind::Full, + }; + let ir = + export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); + + assert!( + ir.contains("!DILexicalBlockFile(scope: !"), + "cross-file local variables should get a file-specific debug scope:\n{ir}" + ); + assert!( + ir.contains( + "!DIFile(filename: \"thread.rs\", directory: \"/tmp/cuda-oxide/crates/cuda-device/src\")" + ), + "cross-file local variables should reference their source file:\n{ir}" + ); + assert!( + ir.contains("!DILocalVariable(name: \"tid\", scope: !") && ir.contains("line: 292"), + "cross-file local variables should preserve the variable file scope and line:\n{ir}" + ); + assert!( + ir.contains("call void @llvm.dbg.declare"), + "cross-file local variables should still get dbg.declare bindings:\n{ir}" + ); +} + +#[test] +fn full_debug_metadata_emits_dbg_value_for_promoted_locals() { + let mut ctx = Context::new(); + + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_region = module.get_operation().deref(&ctx).get_region(0); + let module_block = { + let region = module_region.deref(&ctx); + region.iter(&ctx).next().unwrap() + }; + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![i32_ty.into()], false); + let func = FuncOp::new(&mut ctx, "debug_kernel".try_into().unwrap(), func_ty); + let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 30, 1); + func.get_operation().deref_mut(&ctx).set_loc(func_loc); + + let entry = func.get_or_create_entry_block(&mut ctx); + let arg = entry.deref(&ctx).get_argument(0); + let dbg_value = DebugValueOp::new(&mut ctx, arg); + let dbg_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 31, 13); + dbg_value.get_operation().deref_mut(&ctx).set_loc(dbg_loc); + llvm_export::ops::set_debug_local_variable( + &mut ctx, + dbg_value.get_operation(), + DebugLocalVariableInfo { + name: "x".to_string(), + argument_index: Some(1), + ty: DebugLocalTypeKind::Basic { + name: "i32".to_string(), + size_bits: 32, + encoding: "DW_ATE_signed", + }, + }, + ); + llvm_export::ops::set_debug_local_declaration_location( + &mut ctx, + dbg_value.get_operation(), + PathBuf::from("/tmp/cuda-oxide/tests/declarations.rs"), + 12, + 5, + ); + llvm_export::ops::set_debug_fragment_variables( + &mut ctx, + dbg_value.get_operation(), + &[DebugFragmentVariableInfo { + variable: DebugLocalVariableInfo { + name: "pair".to_string(), + argument_index: None, + ty: DebugLocalTypeKind::Array { + name: "[u32; 2]".to_string(), + size_bits: 64, + element: Box::new(DebugLocalTypeKind::Basic { + name: "u32".to_string(), + size_bits: 32, + encoding: "DW_ATE_unsigned", + }), + count: 2, + }, + }, + fragment: DebugFragment { + offset_bits: 32, + size_bits: 32, + }, + source_scope: None, + declaration: Some(DebugSourcePosition { + file: PathBuf::from("/tmp/cuda-oxide/tests/declarations.rs"), + line: 13, + column: 5, + }), + }], + ); + dbg_value.get_operation().insert_at_back(entry, &ctx); + + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + let config = DebugConfig { + inner: PtxExportConfig, + debug_kind: DebugKind::Full, + }; + let ir = + export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); + + assert!( + ir.contains("declare void @llvm.dbg.value(metadata, metadata, metadata)"), + "full debug should declare dbg.value when it emits one:\n{ir}" + ); + assert!( + ir.contains("call void @llvm.dbg.value(metadata i32 %v0, metadata !"), + "dbg.value should describe the local as the current SSA value:\n{ir}" + ); + assert!( + ir.contains("!DILocalVariable(name: \"x\", arg: 1, scope: !"), + "dbg.value should preserve formal-argument metadata when the source local is an argument:\n{ir}" + ); + assert!( + ir.contains("!DILocalVariable(name: \"x\", arg: 1, scope: !") + && ir.contains("file: !") + && ir.contains("line: 12"), + "DILocalVariable should use the source declaration line, not the dbg.value line:\n{ir}" + ); + assert!( + ir.contains("!DILocation(line: 31, column: 13, scope: !"), + "dbg.value should still be located at the value's current source point:\n{ir}" + ); + assert!( + ir.contains("DW_OP_LLVM_fragment, 32, 32"), + "dbg.value should preserve scalarized source-variable fragments:\n{ir}" + ); + assert!( + ir.contains("!DILocalVariable(name: \"pair\", scope: !") && ir.contains("line: 13"), + "fragment dbg.value should describe the complete source variable:\n{ir}" + ); + assert!( + !ir.contains("llvm.dbg.declare"), + "a value-only debug record should not force dbg.declare:\n{ir}" + ); +} + +#[test] +fn full_debug_metadata_emits_diarglist_for_multi_value_locations() { + let mut ctx = Context::new(); + + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_region = module.get_operation().deref(&ctx).get_region(0); + let module_block = { + let region = module_region.deref(&ctx); + region.iter(&ctx).next().unwrap() + }; + + let ptr_ty = PointerType::get(&ctx, 0); + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get( + &ctx, + void_ty.to_handle(), + vec![ptr_ty.into(), i64_ty.into()], + false, + ); + let func = FuncOp::new(&mut ctx, "debug_kernel".try_into().unwrap(), func_ty); + let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 40, 1); + func.get_operation().deref_mut(&ctx).set_loc(func_loc); + + let entry = func.get_or_create_entry_block(&mut ctx); + let base = entry.deref(&ctx).get_argument(0); + let index = entry.deref(&ctx).get_argument(1); + let dbg_value = DebugValueListOp::new(&mut ctx, vec![base, index]); + let dbg_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 41, 17); + dbg_value.get_operation().deref_mut(&ctx).set_loc(dbg_loc); + llvm_export::ops::set_debug_local_variable( + &mut ctx, + dbg_value.get_operation(), + DebugLocalVariableInfo { + name: "item".to_string(), + argument_index: None, + ty: DebugLocalTypeKind::Basic { + name: "u32".to_string(), + size_bits: 32, + encoding: "DW_ATE_unsigned", + }, + }, + ); + llvm_export::ops::set_debug_value_expression( + &mut ctx, + dbg_value.get_operation(), + &DebugValueExpression::new(vec![ + DebugValueExpressionOp::Arg(0), + DebugValueExpressionOp::Arg(1), + DebugValueExpressionOp::ConstU(4), + DebugValueExpressionOp::Mul, + DebugValueExpressionOp::Plus, + DebugValueExpressionOp::Deref, + ]), + ); + dbg_value.get_operation().insert_at_back(entry, &ctx); + + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + let config = DebugConfig { + inner: PtxExportConfig, + debug_kind: DebugKind::Full, + }; + let ir = + export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); + + assert!( + ir.contains("call void @llvm.dbg.value(metadata !DIArgList(ptr %v0, i64 %v1), metadata !"), + "multi-value debug records should emit an inline DIArgList:\n{ir}" + ); + assert!( + ir.contains( + "!DIExpression(DW_OP_LLVM_arg, 0, DW_OP_LLVM_arg, 1, DW_OP_constu, 4, DW_OP_mul, DW_OP_plus, DW_OP_deref)" + ), + "multi-value debug records should emit the typed location recipe:\n{ir}" + ); + assert!( + ir.contains("declare void @llvm.dbg.value(metadata, metadata, metadata)"), + "multi-value records should reuse the ordinary dbg.value intrinsic declaration:\n{ir}" + ); + + let Some(llvm_as) = ["llvm-as-22", "llvm-as-21", "llvm-as"] + .into_iter() + .find(|tool| { + std::process::Command::new(tool) + .arg("--version") + .output() + .is_ok_and(|out| out.status.success()) + }) + else { + eprintln!("skipping llvm-as parse gate: no llvm-as-22/llvm-as-21/llvm-as on PATH"); + return; + }; + + let ll_path = std::env::temp_dir().join(format!( + "cuda_oxide_diarglist_parse_gate_{}.ll", + std::process::id() + )); + std::fs::write(&ll_path, &ir).expect("write temp .ll"); + let output = std::process::Command::new(llvm_as) + .arg("-o") + .arg("/dev/null") + .arg(&ll_path) + .output() + .expect("run llvm-as"); + let stderr = String::from_utf8_lossy(&output.stderr).into_owned(); + let _ = std::fs::remove_file(&ll_path); + assert!( + output.status.success(), + "{llvm_as} rejected the emitted multi-value debug module:\n{stderr}\n--- module ---\n{ir}" + ); +} + +#[test] +fn full_debug_metadata_uses_inlined_callee_scope_for_inlined_arguments() { + let mut ctx = Context::new(); + + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_region = module.get_operation().deref(&ctx).get_region(0); + let module_block = { + let region = module_region.deref(&ctx); + region.iter(&ctx).next().unwrap() + }; + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![i32_ty.into()], false); + let func = FuncOp::new(&mut ctx, "caller_kernel".try_into().unwrap(), func_ty); + let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 30, 1); + func.get_operation().deref_mut(&ctx).set_loc(func_loc); + llvm_export::ops::set_debug_source_scope_map( + &mut ctx, + func.get_operation(), + &DebugSourceScopeMap { + scopes: vec![ + DebugSourceScope { + id: 0, + parent: None, + span: Some(DebugSourcePosition { + file: PathBuf::from("/tmp/cuda-oxide/tests/kernel.rs"), + line: 30, + column: 1, + }), + inlined: None, + }, + DebugSourceScope { + id: 1, + parent: Some(0), + span: Some(DebugSourcePosition { + file: PathBuf::from("/tmp/cuda-oxide/tests/helper.rs"), + line: 7, + column: 1, + }), + inlined: Some(llvm_export::ops::DebugInlinedScope { + callee_name: "helper::next".to_string(), + callsite: Some(DebugSourcePosition { + file: PathBuf::from("/tmp/cuda-oxide/tests/kernel.rs"), + line: 41, + column: 13, + }), + }), + }, + ], + locations: vec![], + }, + ); + + let entry = func.get_or_create_entry_block(&mut ctx); + let arg = entry.deref(&ctx).get_argument(0); + + let caller_value = DebugValueOp::new(&mut ctx, arg); + let caller_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 31, 9); + caller_value + .get_operation() + .deref_mut(&ctx) + .set_loc(caller_loc); + llvm_export::ops::set_debug_local_variable( + &mut ctx, + caller_value.get_operation(), + DebugLocalVariableInfo { + name: "data".to_string(), + argument_index: Some(1), + ty: DebugLocalTypeKind::Basic { + name: "i32".to_string(), + size_bits: 32, + encoding: "DW_ATE_signed", + }, + }, + ); + llvm_export::ops::set_debug_local_source_scope(&mut ctx, caller_value.get_operation(), 0); + caller_value.get_operation().insert_at_back(entry, &ctx); + + let inlined_value = DebugValueOp::new(&mut ctx, arg); + let inlined_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/helper.rs", 8, 17); + inlined_value + .get_operation() + .deref_mut(&ctx) + .set_loc(inlined_loc); + llvm_export::ops::set_debug_local_variable( + &mut ctx, + inlined_value.get_operation(), + DebugLocalVariableInfo { + name: "self".to_string(), + argument_index: Some(1), + ty: DebugLocalTypeKind::Basic { + name: "i32".to_string(), + size_bits: 32, + encoding: "DW_ATE_signed", + }, + }, + ); + llvm_export::ops::set_debug_local_source_scope(&mut ctx, inlined_value.get_operation(), 1); + inlined_value.get_operation().insert_at_back(entry, &ctx); + + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + let config = DebugConfig { + inner: PtxExportConfig, + debug_kind: DebugKind::Full, + }; + let ir = + export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); + + assert!( + ir.contains("distinct !DISubprogram(name: \"caller_kernel\""), + "caller should keep its own function debug scope:\n{ir}" + ); + assert!( + ir.contains("distinct !DISubprogram(name: \"helper::next\""), + "inlined callee should get its own DISubprogram scope:\n{ir}" + ); + assert!( + ir.contains("!DILocalVariable(name: \"data\", arg: 1, scope: !"), + "caller argument should remain arg #1 in the caller scope:\n{ir}" + ); + assert!( + ir.contains("!DILocalVariable(name: \"self\", arg: 1, scope: !"), + "inlined callee argument should remain arg #1 in the callee scope:\n{ir}" + ); + assert!( + ir.contains("!DILocation(line: 8, column: 17, scope: !") && ir.contains("inlinedAt: !"), + "inlined dbg.value location should point at the callee line and caller callsite:\n{ir}" + ); +} + +#[test] +fn line_table_debug_metadata_ignores_tagged_alloca_variables() { + let mut ctx = Context::new(); + + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_region = module.get_operation().deref(&ctx).get_region(0); + let module_block = { + let region = module_region.deref(&ctx); + region.iter(&ctx).next().unwrap() + }; + + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); + let func = FuncOp::new(&mut ctx, "debug_kernel".try_into().unwrap(), func_ty); + let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 40, 1); + func.get_operation().deref_mut(&ctx).set_loc(func_loc); + + let entry = func.get_or_create_entry_block(&mut ctx); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let one_attr = IntegerAttr::new(i32_ty, APInt::from_u32(1, NonZero::new(32).unwrap())); + let one = ConstantOp::new(&mut ctx, one_attr.into()); + one.get_operation().insert_at_back(entry, &ctx); + let one_val = one.get_operation().deref(&ctx).get_result(0); + + let local = AllocaOp::new(&mut ctx, i32_ty.into(), one_val); + let local_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 41, 9); + local.get_operation().deref_mut(&ctx).set_loc(local_loc); + llvm_export::ops::set_debug_local_variable( + &mut ctx, + local.get_operation(), + DebugLocalVariableInfo { + name: "x".to_string(), + argument_index: None, + ty: DebugLocalTypeKind::Basic { + name: "i32".to_string(), + size_bits: 32, + encoding: "DW_ATE_signed", + }, + }, + ); + local.get_operation().insert_at_back(entry, &ctx); + + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + let config = DebugConfig { + inner: PtxExportConfig, + debug_kind: DebugKind::LineTables, + }; + let ir = + export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); + + assert!( + ir.contains("emissionKind: LineTablesOnly"), + "line-table mode should stay line-table-only:\n{ir}" + ); + assert!( + !ir.contains("llvm.dbg.declare"), + "line-table mode should not emit variable bindings:\n{ir}" + ); + assert!( + !ir.contains("DILocalVariable"), + "line-table mode should not emit local-variable metadata:\n{ir}" + ); +} + +#[test] +fn line_table_debug_metadata_adds_fallback_locations_to_calls() { + let mut ctx = Context::new(); + + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_region = module.get_operation().deref(&ctx).get_region(0); + let module_block = { + let region = module_region.deref(&ctx); + region.iter(&ctx).next().unwrap() + }; + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let void_ty = VoidType::get(&ctx); + let helper_ty = FuncType::get(&ctx, i32_ty.to_handle(), vec![], false); + let helper = FuncOp::new(&mut ctx, "helper".try_into().unwrap(), helper_ty); + helper.get_operation().insert_at_back(module_block, &ctx); + + let caller_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); + let caller = FuncOp::new(&mut ctx, "debug_kernel".try_into().unwrap(), caller_ty); + let caller_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 20, 3); + caller.get_operation().deref_mut(&ctx).set_loc(caller_loc); + + let entry = caller.get_or_create_entry_block(&mut ctx); + let call = CallOp::new( + &mut ctx, + CallOpCallable::Direct("helper".try_into().unwrap()), + helper_ty, + vec![], + ); + call.get_operation().insert_at_back(entry, &ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + + caller.get_operation().insert_at_back(module_block, &ctx); + + let config = DebugConfig { + inner: PtxExportConfig, + debug_kind: DebugKind::LineTables, + }; + let ir = + export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); + + let call_line = ir + .lines() + .find(|line| line.contains("call i32 @helper()")) + .expect("call instruction"); + assert!( + call_line.contains(", !dbg !"), + "calls without their own source span should use the function fallback location:\n{ir}" + ); + assert!( + ir.contains("!DILocation(line: 20, column: 3, scope: !"), + "fallback call location should point at the caller's function line:\n{ir}" + ); +} + +#[test] +fn line_table_debug_metadata_emits_explicit_artificial_calls_at_line_zero() { + let mut ctx = Context::new(); + + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_region = module.get_operation().deref(&ctx).get_region(0); + let module_block = { + let region = module_region.deref(&ctx); + region.iter(&ctx).next().unwrap() + }; + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let void_ty = VoidType::get(&ctx); + let helper_ty = FuncType::get(&ctx, i32_ty.to_handle(), vec![], false); + let helper = FuncOp::new(&mut ctx, "helper".try_into().unwrap(), helper_ty); + helper.get_operation().insert_at_back(module_block, &ctx); + + let caller_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); + let caller = FuncOp::new(&mut ctx, "debug_kernel".try_into().unwrap(), caller_ty); + let caller_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 20, 3); + caller.get_operation().deref_mut(&ctx).set_loc(caller_loc); + + let entry = caller.get_or_create_entry_block(&mut ctx); + let call = CallOp::new( + &mut ctx, + CallOpCallable::Direct("helper".try_into().unwrap()), + helper_ty, + vec![], + ); + call.get_operation() + .deref_mut(&ctx) + .set_loc(llvm_export::artificial_debug_location()); + call.get_operation().insert_at_back(entry, &ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + + caller.get_operation().insert_at_back(module_block, &ctx); + + let config = DebugConfig { + inner: PtxExportConfig, + debug_kind: DebugKind::LineTables, + }; + let ir = + export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); + + let call_line = ir + .lines() + .find(|line| line.contains("call i32 @helper()")) + .expect("call instruction"); + assert!( + call_line.contains(", !dbg !"), + "LLVM still requires an artificial call location:\n{ir}" + ); + assert!( + ir.contains("!DILocation(line: 0, column: 0, scope: !"), + "artificial setup must not reuse the caller's user line:\n{ir}" + ); +} + +#[test] +fn full_debug_metadata_emits_scalarized_fragment_dbg_declares() { + let mut ctx = Context::new(); + + let module = ModuleOp::new(&mut ctx, "fragment_debug".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.into(), vec![], false); + let func = FuncOp::new( + &mut ctx, + "fragment_debug_kernel".try_into().unwrap(), + func_ty, + ); + let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/fragments.rs", 10, 1); + func.get_operation().deref_mut(&ctx).set_loc(func_loc); + let entry = func.get_or_create_entry_block(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let one_attr = IntegerAttr::new(i32_ty, APInt::from_u32(1, NonZero::new(32).unwrap())); + let one = ConstantOp::new(&mut ctx, one_attr.into()); + let one_value = one.get_operation().deref(&ctx).get_result(0); + one.get_operation().insert_at_back(entry, &ctx); + + let whole_ty = DebugLocalTypeKind::Array { + name: "[u32; 2]".to_string(), + size_bits: 64, + element: Box::new(DebugLocalTypeKind::Basic { + name: "u32".to_string(), + size_bits: 32, + encoding: "DW_ATE_unsigned", + }), + count: 2, + }; + for (index, offset_bits) in [0u64, 32].into_iter().enumerate() { + let alloca = AllocaOp::new(&mut ctx, i32_ty.into(), one_value); + llvm_export::ops::set_debug_fragment_variables( + &mut ctx, + alloca.get_operation(), + &[DebugFragmentVariableInfo { + variable: DebugLocalVariableInfo { + name: "pair".to_string(), + argument_index: None, + ty: whole_ty.clone(), + }, + fragment: DebugFragment { + offset_bits, + size_bits: 32, + }, + source_scope: None, + declaration: Some(DebugSourcePosition { + file: PathBuf::from("/tmp/cuda-oxide/tests/fragments.rs"), + line: 11, + column: 9, + }), + }], + ); + let loc = src_location( + &mut ctx, + "/tmp/cuda-oxide/tests/fragments.rs", + 12 + index as i32, + 9, + ); + alloca.get_operation().deref_mut(&ctx).set_loc(loc); + alloca.get_operation().insert_at_back(entry, &ctx); + } + + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + let config = DebugConfig { + inner: PtxExportConfig, + debug_kind: DebugKind::Full, + }; + let ir = + export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); + + assert!( + ir.contains("DW_OP_LLVM_fragment, 0, 32"), + "first scalarized piece should describe the low fragment:\n{ir}" + ); + assert!( + ir.contains("DW_OP_LLVM_fragment, 32, 32"), + "second scalarized piece should describe the high fragment:\n{ir}" + ); + assert_eq!( + ir.matches("!DILocalVariable(name: \"pair\"").count(), + 1, + "all pieces must share one DILocalVariable identity:\n{ir}" + ); + assert_eq!( + ir.matches("call void @llvm.dbg.declare").count(), + 2, + "each scalarized storage piece should emit one dbg.declare:\n{ir}" + ); + + let Some(llvm_as) = ["llvm-as-22", "llvm-as-21", "llvm-as"] + .into_iter() + .find(|tool| { + std::process::Command::new(tool) + .arg("--version") + .output() + .is_ok_and(|out| out.status.success()) + }) + else { + eprintln!("skipping fragment llvm-as parse gate: no llvm-as on PATH"); + return; + }; + let ll_path = std::env::temp_dir().join(format!( + "cuda_oxide_fragment_parse_gate_{}.ll", + std::process::id() + )); + std::fs::write(&ll_path, &ir).expect("write fragment temp .ll"); + let output = std::process::Command::new(llvm_as) + .arg("-o") + .arg("/dev/null") + .arg(&ll_path) + .output() + .expect("run llvm-as for fragments"); + let stderr = String::from_utf8_lossy(&output.stderr).into_owned(); + let _ = std::fs::remove_file(&ll_path); + assert!( + output.status.success(), + "{llvm_as} rejected scalarized fragment debug metadata:\n{stderr}\n--- module ---\n{ir}" + ); +} + +#[test] +fn full_debug_metadata_emits_projected_dbg_declares() { + let mut ctx = Context::new(); + + let module = ModuleOp::new(&mut ctx, "projected_debug".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.into(), vec![], false); + let func = FuncOp::new( + &mut ctx, + "projected_debug_kernel".try_into().unwrap(), + func_ty, + ); + let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/projected.rs", 10, 1); + func.get_operation().deref_mut(&ctx).set_loc(func_loc); + let entry = func.get_or_create_entry_block(&mut ctx); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let one_attr = IntegerAttr::new(i32_ty, APInt::from_u32(1, NonZero::new(32).unwrap())); + let one = ConstantOp::new(&mut ctx, one_attr.into()); + let one_value = one.get_operation().deref(&ctx).get_result(0); + one.get_operation().insert_at_back(entry, &ctx); + + let storage_ty = ArrayType::get(&ctx, i32_ty.into(), 8); + let alloca = AllocaOp::new(&mut ctx, storage_ty.into(), one_value); + llvm_export::ops::set_debug_projected_variables( + &mut ctx, + alloca.get_operation(), + &[ + DebugProjectedVariableInfo { + variable: DebugLocalVariableInfo { + name: "field_value".to_string(), + argument_index: None, + ty: DebugLocalTypeKind::Basic { + name: "u32".to_string(), + size_bits: 32, + encoding: "DW_ATE_unsigned", + }, + }, + dereference_base: false, + offset_bytes: 8, + source_scope: None, + declaration: Some(DebugSourcePosition { + file: PathBuf::from("/tmp/cuda-oxide/tests/projected.rs"), + line: 11, + column: 9, + }), + }, + DebugProjectedVariableInfo { + variable: DebugLocalVariableInfo { + name: "tuple_value".to_string(), + argument_index: None, + ty: DebugLocalTypeKind::Basic { + name: "u64".to_string(), + size_bits: 64, + encoding: "DW_ATE_unsigned", + }, + }, + dereference_base: false, + offset_bytes: 16, + source_scope: None, + declaration: Some(DebugSourcePosition { + file: PathBuf::from("/tmp/cuda-oxide/tests/projected.rs"), + line: 12, + column: 9, + }), + }, + DebugProjectedVariableInfo { + variable: DebugLocalVariableInfo { + name: "array_value".to_string(), + argument_index: None, + ty: DebugLocalTypeKind::Basic { + name: "u32".to_string(), + size_bits: 32, + encoding: "DW_ATE_unsigned", + }, + }, + dereference_base: false, + offset_bytes: 24, + source_scope: None, + declaration: Some(DebugSourcePosition { + file: PathBuf::from("/tmp/cuda-oxide/tests/projected.rs"), + line: 13, + column: 9, + }), + }, + DebugProjectedVariableInfo { + variable: DebugLocalVariableInfo { + name: "deref_value".to_string(), + argument_index: None, + ty: DebugLocalTypeKind::Basic { + name: "u32".to_string(), + size_bits: 32, + encoding: "DW_ATE_unsigned", + }, + }, + dereference_base: true, + offset_bytes: 0, + source_scope: None, + declaration: Some(DebugSourcePosition { + file: PathBuf::from("/tmp/cuda-oxide/tests/projected.rs"), + line: 14, + column: 9, + }), + }, + DebugProjectedVariableInfo { + variable: DebugLocalVariableInfo { + name: "deref_field_value".to_string(), + argument_index: None, + ty: DebugLocalTypeKind::Basic { + name: "u64".to_string(), + size_bits: 64, + encoding: "DW_ATE_unsigned", + }, + }, + dereference_base: true, + offset_bytes: 32, + source_scope: None, + declaration: Some(DebugSourcePosition { + file: PathBuf::from("/tmp/cuda-oxide/tests/projected.rs"), + line: 15, + column: 9, + }), + }, + ], + ); + alloca.get_operation().insert_at_back(entry, &ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + let config = DebugConfig { + inner: PtxExportConfig, + debug_kind: DebugKind::Full, + }; + let ir = export_module_to_string_with_config(&ctx, &module, &config) + .expect("projected debug export succeeds"); + + assert_eq!( + ir.matches("call void @llvm.dbg.declare").count(), + 5, + "each projected source variable should get its own dbg.declare:\n{ir}" + ); + for name in [ + "field_value", + "tuple_value", + "array_value", + "deref_value", + "deref_field_value", + ] { + assert!( + ir.contains(&format!("!DILocalVariable(name: \"{name}\"")), + "missing projected variable {name}:\n{ir}" + ); + } + for offset in [8u64, 16, 24] { + assert!( + ir.contains(&format!("!DIExpression(DW_OP_plus_uconst, {offset})")), + "missing static projection offset {offset}:\n{ir}" + ); + } + assert!( + ir.contains("!DIExpression(DW_OP_deref)"), + "missing dereference-only debug expression:\n{ir}" + ); + assert!( + ir.contains("!DIExpression(DW_OP_deref, DW_OP_plus_uconst, 32)"), + "missing dereference-plus-field debug expression:\n{ir}" + ); +} diff --git a/crates/llvm-export/tests/export_test/device_externs.rs b/crates/llvm-export/tests/export_test/device_externs.rs new file mode 100644 index 0000000000..ee928c3819 --- /dev/null +++ b/crates/llvm-export/tests/export_test/device_externs.rs @@ -0,0 +1,800 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use llvm_export::{ + export::{ + DeviceExternAttrs, DeviceExternDecl, DeviceExternType, NvvmExportConfig, NvvmIrDialect, + export_module_with_externs, + }, + ops::{AddressOfOp, CallOp, FuncOp, ReturnOp}, + types::{FuncType, HalfType, PointerType, VoidType}, +}; +use pliron::{ + builtin::{ + op_interfaces::CallOpCallable, + ops::ModuleOp, + types::{IntegerType, Signedness}, + }, + context::Context, + op::Op, +}; + +use crate::common::module_top_block; + +#[test] +fn legacy_device_extern_adapts_exact_pointer_arguments_and_results() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "legacy_extern".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + + let ptr_ty = PointerType::get(&ctx, 0); + let external_ty = FuncType::get(&ctx, ptr_ty.into(), vec![ptr_ty.into()], false); + FuncOp::new(&mut ctx, "float_roundtrip".try_into().unwrap(), external_ty) + .get_operation() + .insert_at_back(module_block, &ctx); + + let void_ty = VoidType::get(&ctx); + let caller_ty = FuncType::get(&ctx, void_ty.into(), vec![ptr_ty.into()], false); + let caller = FuncOp::new(&mut ctx, "caller".try_into().unwrap(), caller_ty); + let entry = caller.get_or_create_entry_block(&mut ctx); + let pointer = entry.deref(&ctx).get_argument(0); + CallOp::new( + &mut ctx, + CallOpCallable::Direct("float_roundtrip".try_into().unwrap()), + external_ty, + vec![pointer], + ) + .get_operation() + .insert_at_back(entry, &ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + caller.get_operation().insert_at_back(module_block, &ctx); + + let externs = [DeviceExternDecl { + export_name: "float_roundtrip".to_string(), + param_types: vec![DeviceExternType::pointer_to(DeviceExternType::Float32, 0)], + return_type: DeviceExternType::pointer_to(DeviceExternType::Float32, 0), + attrs: DeviceExternAttrs::default(), + }]; + let ir = export_module_with_externs( + &ctx, + &module, + &externs, + &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + ) + .expect("legacy extern export succeeds"); + + assert!( + ir.contains("declare float* @float_roundtrip(float*)"), + "{ir}" + ); + assert!(ir.contains("bitcast i8* %v0 to float*"), "{ir}"); + assert!(ir.contains("call float* @float_roundtrip(float*"), "{ir}"); + assert!( + ir.contains(" = bitcast float* ") && ir.contains(" to i8*"), + "{ir}" + ); + assert!( + !ir.split(|c: char| !c.is_ascii_alphanumeric()) + .any(|token| token == "ptr"), + "{ir}" + ); +} + +#[test] +fn legacy_device_extern_preserves_pointer_address_spaces() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "legacy_extern_as".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let ptr_ty = PointerType::get(&ctx, 3); + let void_ty = VoidType::get(&ctx); + let external_ty = FuncType::get(&ctx, void_ty.into(), vec![ptr_ty.into()], false); + FuncOp::new(&mut ctx, "shared_float".try_into().unwrap(), external_ty) + .get_operation() + .insert_at_back(module_block, &ctx); + let caller = FuncOp::new(&mut ctx, "caller".try_into().unwrap(), external_ty); + let entry = caller.get_or_create_entry_block(&mut ctx); + let pointer = entry.deref(&ctx).get_argument(0); + CallOp::new( + &mut ctx, + CallOpCallable::Direct("shared_float".try_into().unwrap()), + external_ty, + vec![pointer], + ) + .get_operation() + .insert_at_back(entry, &ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + caller.get_operation().insert_at_back(module_block, &ctx); + + let externs = [DeviceExternDecl { + export_name: "shared_float".to_string(), + param_types: vec![DeviceExternType::pointer_to(DeviceExternType::Float32, 3)], + return_type: DeviceExternType::Void, + attrs: DeviceExternAttrs::default(), + }]; + let ir = export_module_with_externs( + &ctx, + &module, + &externs, + &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + ) + .expect("legacy extern export succeeds"); + assert!( + ir.contains("declare void @shared_float(float addrspace(3)*)"), + "{ir}" + ); + assert!( + ir.contains("bitcast i8 addrspace(3)* %v0 to float addrspace(3)*"), + "{ir}" + ); +} + +#[test] +fn modern_device_extern_erases_pointee_without_boundary_casts() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "modern_extern".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let ptr_ty = PointerType::get(&ctx, 0); + let void_ty = VoidType::get(&ctx); + let external_ty = FuncType::get(&ctx, void_ty.into(), vec![ptr_ty.into()], false); + FuncOp::new(&mut ctx, "takes_float".try_into().unwrap(), external_ty) + .get_operation() + .insert_at_back(module_block, &ctx); + let caller = FuncOp::new(&mut ctx, "caller".try_into().unwrap(), external_ty); + let entry = caller.get_or_create_entry_block(&mut ctx); + let pointer = entry.deref(&ctx).get_argument(0); + CallOp::new( + &mut ctx, + CallOpCallable::Direct("takes_float".try_into().unwrap()), + external_ty, + vec![pointer], + ) + .get_operation() + .insert_at_back(entry, &ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + caller.get_operation().insert_at_back(module_block, &ctx); + + let externs = [DeviceExternDecl { + export_name: "takes_float".to_string(), + param_types: vec![DeviceExternType::pointer_to(DeviceExternType::Float32, 0)], + return_type: DeviceExternType::Void, + attrs: DeviceExternAttrs::default(), + }]; + let ir = export_module_with_externs( + &ctx, + &module, + &externs, + &NvvmExportConfig::new(NvvmIrDialect::Modern), + ) + .expect("modern extern export succeeds"); + assert!(ir.contains("declare void @takes_float(ptr)"), "{ir}"); + assert!(ir.contains("call void @takes_float(ptr %v0)"), "{ir}"); + assert!(!ir.contains("bitcast"), "{ir}"); +} + +#[test] +fn device_extern_rejects_invalid_symbol_and_address_space_mismatch() { + let mut ctx = Context::new(); + let empty = ModuleOp::new(&mut ctx, "empty".try_into().unwrap()); + let invalid = [DeviceExternDecl { + export_name: "bad.name".to_string(), + param_types: vec![], + return_type: DeviceExternType::Void, + attrs: DeviceExternAttrs::default(), + }]; + let err = export_module_with_externs( + &ctx, + &empty, + &invalid, + &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + ) + .expect_err("invalid NVVM symbol must fail"); + assert!(err.contains("global-identifier subset"), "{err}"); + + let reserved_intrinsic_prefix = [DeviceExternDecl { + export_name: "llvm_external".to_string(), + param_types: vec![], + return_type: DeviceExternType::Void, + attrs: DeviceExternAttrs::default(), + }]; + let err = export_module_with_externs( + &ctx, + &empty, + &reserved_intrinsic_prefix, + &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + ) + .expect_err("the reserved intrinsic namespace must not be ambiguous"); + assert!(err.contains("reserves for LLVM intrinsics"), "{err}"); + + let by_value_array = [DeviceExternDecl { + export_name: "array_by_value".to_string(), + param_types: vec![DeviceExternType::Array { + element: Box::new(DeviceExternType::Float32), + len: 4, + }], + return_type: DeviceExternType::Void, + attrs: DeviceExternAttrs::default(), + }]; + let err = export_module_with_externs( + &ctx, + &empty, + &by_value_array, + &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + ) + .expect_err("by-value aggregate externs must be rejected"); + assert!(err.contains("passes an array by value"), "{err}"); + + let nested_half = [DeviceExternDecl { + export_name: "half_buffer".to_string(), + param_types: vec![DeviceExternType::pointer_to( + DeviceExternType::Array { + element: Box::new(DeviceExternType::Float16), + len: 4, + }, + 0, + )], + return_type: DeviceExternType::Void, + attrs: DeviceExternAttrs::default(), + }]; + let err = export_module_with_externs( + &ctx, + &empty, + &nested_half, + &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + ) + .expect_err("legacy half nested in a pointer must fail"); + assert!( + err.contains("CUDA 12 legacy") && err.contains("half"), + "{err}" + ); + let modern = export_module_with_externs( + &ctx, + &empty, + &nested_half, + &NvvmExportConfig::new(NvvmIrDialect::Modern), + ) + .expect("modern opaque-pointer extern may use half pointees"); + assert!( + modern.contains("declare void @half_buffer(ptr)"), + "{modern}" + ); + + let module = ModuleOp::new(&mut ctx, "mismatch".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let ptr0 = PointerType::get(&ctx, 0); + let void_ty = VoidType::get(&ctx); + let external_ty = FuncType::get(&ctx, void_ty.into(), vec![ptr0.into()], false); + FuncOp::new(&mut ctx, "shared_only".try_into().unwrap(), external_ty) + .get_operation() + .insert_at_back(module_block, &ctx); + let caller = FuncOp::new(&mut ctx, "caller".try_into().unwrap(), external_ty); + let entry = caller.get_or_create_entry_block(&mut ctx); + let pointer = entry.deref(&ctx).get_argument(0); + CallOp::new( + &mut ctx, + CallOpCallable::Direct("shared_only".try_into().unwrap()), + external_ty, + vec![pointer], + ) + .get_operation() + .insert_at_back(entry, &ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + caller.get_operation().insert_at_back(module_block, &ctx); + let mismatch = [DeviceExternDecl { + export_name: "shared_only".to_string(), + param_types: vec![DeviceExternType::pointer_to(DeviceExternType::Float32, 3)], + return_type: DeviceExternType::Void, + attrs: DeviceExternAttrs::default(), + }]; + let err = export_module_with_externs( + &ctx, + &module, + &mismatch, + &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + ) + .expect_err("address-space mismatch must fail"); + assert!( + err.contains("parameter, result, or pointer address-space types"), + "{err}" + ); +} + +#[test] +fn device_extern_rejects_same_name_declaration_shape_without_a_call() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "extern_decl_conflict".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let void_ty = VoidType::get(&ctx); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let lowered_type = FuncType::get(&ctx, void_ty.into(), vec![i32_ty.into()], false); + FuncOp::new( + &mut ctx, + "conflicting_decl".try_into().unwrap(), + lowered_type, + ) + .get_operation() + .insert_at_back(module_block, &ctx); + + let externs = [DeviceExternDecl { + export_name: "conflicting_decl".to_string(), + param_types: vec![DeviceExternType::Float32], + return_type: DeviceExternType::Void, + attrs: DeviceExternAttrs::default(), + }]; + let error = export_module_with_externs( + &ctx, + &module, + &externs, + &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + ) + .expect_err("the exporter must independently reject a conflicting declaration"); + assert!( + error.contains("parameter, result, or pointer address-space types"), + "{error}" + ); +} + +#[test] +fn device_extern_rejects_definition_and_address_taken_shape_conflicts() { + let mut ctx = Context::new(); + let void_ty = VoidType::get(&ctx); + + let definition_module = ModuleOp::new(&mut ctx, "extern_definition".try_into().unwrap()); + let definition_block = module_top_block(&mut ctx, &definition_module); + let no_args = FuncType::get(&ctx, void_ty.into(), vec![], false); + let definition = FuncOp::new(&mut ctx, "defined_external".try_into().unwrap(), no_args); + let definition_entry = definition.get_or_create_entry_block(&mut ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(definition_entry, &ctx); + definition + .get_operation() + .insert_at_back(definition_block, &ctx); + let definition_extern = [DeviceExternDecl { + export_name: "defined_external".to_string(), + param_types: vec![], + return_type: DeviceExternType::Void, + attrs: DeviceExternAttrs::default(), + }]; + let error = export_module_with_externs( + &ctx, + &definition_module, + &definition_extern, + &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + ) + .expect_err("a side-table extern must not collide with a definition"); + assert!(error.contains("function definition"), "{error}"); + + let address_module = ModuleOp::new(&mut ctx, "extern_address".try_into().unwrap()); + let address_block = module_top_block(&mut ctx, &address_module); + let generic_pointer = PointerType::get(&ctx, 0); + let lowered_type = FuncType::get(&ctx, void_ty.into(), vec![generic_pointer.into()], false); + FuncOp::new( + &mut ctx, + "addressed_external".try_into().unwrap(), + lowered_type, + ) + .get_operation() + .insert_at_back(address_block, &ctx); + let caller = FuncOp::new(&mut ctx, "caller".try_into().unwrap(), lowered_type); + let caller_entry = caller.get_or_create_entry_block(&mut ctx); + let argument = caller_entry.deref(&ctx).get_argument(0); + let address = AddressOfOp::new(&mut ctx, "addressed_external".try_into().unwrap(), 0); + let address_value = address.get_operation().deref(&ctx).get_result(0); + address.get_operation().insert_at_back(caller_entry, &ctx); + CallOp::new( + &mut ctx, + CallOpCallable::Indirect(address_value), + lowered_type, + vec![argument], + ) + .get_operation() + .insert_at_back(caller_entry, &ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(caller_entry, &ctx); + caller.get_operation().insert_at_back(address_block, &ctx); + let address_extern = [DeviceExternDecl { + export_name: "addressed_external".to_string(), + param_types: vec![DeviceExternType::pointer_to(DeviceExternType::Float32, 3)], + return_type: DeviceExternType::Void, + attrs: DeviceExternAttrs::default(), + }]; + let error = export_module_with_externs( + &ctx, + &address_module, + &address_extern, + &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + ) + .expect_err("address-taking must not bypass exact extern shape validation"); + assert!( + error.contains("parameter, result, or pointer address-space types"), + "{error}" + ); +} + +#[test] +fn modern_device_extern_emits_signext_zeroext_for_small_integer_params() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "small_types_extern".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + + // The extern takes (i8 signext, i16 zeroext, i1 zeroext, half) and + // returns void. The declared IR types stay NARROW (matching cuda-oxide's + // own i8/i16/i1 SSA values and the clang-compiled LTOIR definition); the + // ABI extension attributes live only in the DeviceExternType metadata. + let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signless); + let i16_ty = IntegerType::get(&ctx, 16, Signedness::Signless); + let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); + let half_ty = HalfType::get(&ctx); + let void_ty = VoidType::get(&ctx); + + let external_ty = FuncType::get( + &ctx, + void_ty.into(), + vec![i8_ty.into(), i16_ty.into(), i1_ty.into(), half_ty.into()], + false, + ); + FuncOp::new(&mut ctx, "small_types_fn".try_into().unwrap(), external_ty) + .get_operation() + .insert_at_back(module_block, &ctx); + + // Caller that forwards its own narrow parameters to the extern with no + // value conversions. + let caller_ty = FuncType::get( + &ctx, + void_ty.into(), + vec![i8_ty.into(), i16_ty.into(), i1_ty.into(), half_ty.into()], + false, + ); + let caller = FuncOp::new(&mut ctx, "caller".try_into().unwrap(), caller_ty); + let entry = caller.get_or_create_entry_block(&mut ctx); + let arg0 = entry.deref(&ctx).get_argument(0); + let arg1 = entry.deref(&ctx).get_argument(1); + let arg2 = entry.deref(&ctx).get_argument(2); + let arg3 = entry.deref(&ctx).get_argument(3); + CallOp::new( + &mut ctx, + CallOpCallable::Direct("small_types_fn".try_into().unwrap()), + external_ty, + vec![arg0, arg1, arg2, arg3], + ) + .get_operation() + .insert_at_back(entry, &ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + caller.get_operation().insert_at_back(module_block, &ctx); + + let externs = [DeviceExternDecl { + export_name: "small_types_fn".to_string(), + param_types: vec![ + DeviceExternType::SignExtInteger(8), // i8, sign-extended by NVPTX + DeviceExternType::ZeroExtInteger(16), // u16, zero-extended by NVPTX + DeviceExternType::ZeroExtInteger(1), // bool + DeviceExternType::Float16, // f16 as native half + ], + return_type: DeviceExternType::Void, + attrs: DeviceExternAttrs::default(), + }]; + + let ir = export_module_with_externs( + &ctx, + &module, + &externs, + &NvvmExportConfig::new(NvvmIrDialect::Modern), + ) + .expect("modern extern with small types export succeeds"); + + // Declaration: narrow types with parameter-position attributes. + assert!( + ir.contains("declare void @small_types_fn(i8 signext, i16 zeroext, i1 zeroext, half)"), + "declaration should keep narrow types with signext/zeroext attributes:\n{ir}" + ); + + // Call site keeps the narrow types and attributes too. + assert!( + ir.contains("i8 signext %v0"), + "call should use signext on first arg:\n{ir}" + ); + assert!( + ir.contains("i16 zeroext %v1"), + "call should use zeroext on second arg:\n{ir}" + ); + assert!( + ir.contains("i1 zeroext %v2"), + "call should use zeroext on the bool arg:\n{ir}" + ); + assert!( + ir.contains("half %v3"), + "call should use half for the f16 arg:\n{ir}" + ); +} + +#[test] +fn modern_device_extern_signext_return_type() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "signext_return".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + + let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signless); + let void_ty = VoidType::get(&ctx); + + // Extern returns i8 with signext. The declared type stays i8; only the + // attribute marks the NVPTX widening, and in return position LLVM's + // grammar requires the attribute BEFORE the type. + let external_ty = FuncType::get(&ctx, i8_ty.into(), vec![], false); + FuncOp::new(&mut ctx, "get_small_val".try_into().unwrap(), external_ty) + .get_operation() + .insert_at_back(module_block, &ctx); + + // Caller that calls the extern and discards the result. + let caller_ty = FuncType::get(&ctx, void_ty.into(), vec![], false); + let caller = FuncOp::new(&mut ctx, "caller".try_into().unwrap(), caller_ty); + let entry = caller.get_or_create_entry_block(&mut ctx); + CallOp::new( + &mut ctx, + CallOpCallable::Direct("get_small_val".try_into().unwrap()), + external_ty, + vec![], + ) + .get_operation() + .insert_at_back(entry, &ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + caller.get_operation().insert_at_back(module_block, &ctx); + + let externs = [DeviceExternDecl { + export_name: "get_small_val".to_string(), + param_types: vec![], + return_type: DeviceExternType::SignExtInteger(8), + attrs: DeviceExternAttrs::default(), + }]; + + let ir = export_module_with_externs( + &ctx, + &module, + &externs, + &NvvmExportConfig::new(NvvmIrDialect::Modern), + ) + .expect("modern extern with signext return export succeeds"); + + // Declaration: attribute precedes the narrow return type. + assert!( + ir.contains("declare signext i8 @get_small_val()"), + "declaration should have signext BEFORE the return type:\n{ir}" + ); + + // Call site uses the same return-position placement. + assert!( + ir.contains("call signext i8 @get_small_val()"), + "call should have signext BEFORE the return type:\n{ir}" + ); +} + +#[test] +fn plain_integer_device_extern_has_no_extension_attributes() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "plain_int".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let void_ty = VoidType::get(&ctx); + + let external_ty = FuncType::get(&ctx, void_ty.into(), vec![i32_ty.into()], false); + FuncOp::new(&mut ctx, "plain_int_fn".try_into().unwrap(), external_ty) + .get_operation() + .insert_at_back(module_block, &ctx); + + let caller = FuncOp::new(&mut ctx, "caller".try_into().unwrap(), external_ty); + let entry = caller.get_or_create_entry_block(&mut ctx); + let arg0 = entry.deref(&ctx).get_argument(0); + CallOp::new( + &mut ctx, + CallOpCallable::Direct("plain_int_fn".try_into().unwrap()), + external_ty, + vec![arg0], + ) + .get_operation() + .insert_at_back(entry, &ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + caller.get_operation().insert_at_back(module_block, &ctx); + + let externs = [DeviceExternDecl { + export_name: "plain_int_fn".to_string(), + param_types: vec![DeviceExternType::Integer(32)], + return_type: DeviceExternType::Void, + attrs: DeviceExternAttrs::default(), + }]; + + let ir = export_module_with_externs( + &ctx, + &module, + &externs, + &NvvmExportConfig::new(NvvmIrDialect::Modern), + ) + .expect("plain int extern export succeeds"); + + // Plain i32 should have no signext/zeroext. + assert!( + ir.contains("declare void @plain_int_fn(i32)"), + "declaration should have plain i32 without attributes:\n{ir}" + ); + assert!( + !ir.contains("signext") && !ir.contains("zeroext"), + "plain i32 should have no extension attributes:\n{ir}" + ); +} + +#[test] +fn legacy_device_extern_rejects_small_integers_by_value() { + let mut ctx = Context::new(); + let empty = ModuleOp::new(&mut ctx, "empty".try_into().unwrap()); + let externs = [DeviceExternDecl { + export_name: "takes_small".to_string(), + param_types: vec![DeviceExternType::SignExtInteger(8)], + return_type: DeviceExternType::ZeroExtInteger(16), + attrs: DeviceExternAttrs::default(), + }]; + let err = export_module_with_externs( + &ctx, + &empty, + &externs, + &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + ) + .expect_err("legacy sub-32-bit by-value externs must fail cleanly"); + assert!( + err.contains("CUDA 12 legacy") && err.contains("sub-32-bit"), + "{err}" + ); +} + +/// Build a module whose externs carry small params AND a small return, in +/// both declare and call positions, so the emitted attribute placement can +/// be checked against a real LLVM parser. +fn small_type_extern_module(ctx: &mut Context) -> (ModuleOp, Vec) { + let module = ModuleOp::new(ctx, "small_types_parse_gate".try_into().unwrap()); + let module_block = module_top_block(ctx, &module); + + let i8_ty = IntegerType::get(ctx, 8, Signedness::Signless); + let i16_ty = IntegerType::get(ctx, 16, Signedness::Signless); + let i1_ty = IntegerType::get(ctx, 1, Signedness::Signless); + let half_ty = HalfType::get(ctx); + let void_ty = VoidType::get(ctx); + + let take_small_ty = FuncType::get( + ctx, + void_ty.into(), + vec![i8_ty.into(), i16_ty.into(), i1_ty.into(), half_ty.into()], + false, + ); + FuncOp::new(ctx, "take_small".try_into().unwrap(), take_small_ty) + .get_operation() + .insert_at_back(module_block, &*ctx); + + let give_small_ty = FuncType::get(ctx, i8_ty.into(), vec![], false); + FuncOp::new(ctx, "give_small".try_into().unwrap(), give_small_ty) + .get_operation() + .insert_at_back(module_block, &*ctx); + + let caller_ty = FuncType::get( + ctx, + void_ty.into(), + vec![i8_ty.into(), i16_ty.into(), i1_ty.into(), half_ty.into()], + false, + ); + let caller = FuncOp::new(ctx, "caller".try_into().unwrap(), caller_ty); + let entry = caller.get_or_create_entry_block(ctx); + let arg0 = entry.deref(ctx).get_argument(0); + let arg1 = entry.deref(ctx).get_argument(1); + let arg2 = entry.deref(ctx).get_argument(2); + let arg3 = entry.deref(ctx).get_argument(3); + CallOp::new( + ctx, + CallOpCallable::Direct("take_small".try_into().unwrap()), + take_small_ty, + vec![arg0, arg1, arg2, arg3], + ) + .get_operation() + .insert_at_back(entry, &*ctx); + CallOp::new( + ctx, + CallOpCallable::Direct("give_small".try_into().unwrap()), + give_small_ty, + vec![], + ) + .get_operation() + .insert_at_back(entry, &*ctx); + ReturnOp::new(ctx, None) + .get_operation() + .insert_at_back(entry, &*ctx); + caller.get_operation().insert_at_back(module_block, &*ctx); + + let externs = vec![ + DeviceExternDecl { + export_name: "take_small".to_string(), + param_types: vec![ + DeviceExternType::SignExtInteger(8), + DeviceExternType::ZeroExtInteger(16), + DeviceExternType::ZeroExtInteger(1), + DeviceExternType::Float16, + ], + return_type: DeviceExternType::Void, + attrs: DeviceExternAttrs::default(), + }, + DeviceExternDecl { + export_name: "give_small".to_string(), + param_types: vec![], + return_type: DeviceExternType::SignExtInteger(8), + attrs: DeviceExternAttrs::default(), + }, + ]; + (module, externs) +} + +/// Parse gate: the emitted attribute placement (`i8 signext` in parameter +/// position, `signext i8` in return position, for both `declare` and `call`) +/// must be accepted by a real LLVM parser, not just string-matched. +#[test] +fn small_type_extern_module_parses_with_llvm_as() { + let mut ctx = Context::new(); + let (module, externs) = small_type_extern_module(&mut ctx); + let ir = export_module_with_externs( + &ctx, + &module, + &externs, + &NvvmExportConfig::new(NvvmIrDialect::Modern), + ) + .expect("modern small-type extern export succeeds"); + + // Belt-and-braces string checks before the external parse. + assert!( + ir.contains("declare void @take_small(i8 signext, i16 zeroext, i1 zeroext, half)"), + "{ir}" + ); + assert!(ir.contains("declare signext i8 @give_small()"), "{ir}"); + assert!(ir.contains("call signext i8 @give_small()"), "{ir}"); + + let Some(llvm_as) = ["llvm-as-22", "llvm-as-21", "llvm-as"] + .into_iter() + .find(|tool| { + std::process::Command::new(tool) + .arg("--version") + .output() + .is_ok_and(|out| out.status.success()) + }) + else { + eprintln!("skipping llvm-as parse gate: no llvm-as-22/llvm-as-21/llvm-as on PATH"); + return; + }; + + let ll_path = std::env::temp_dir().join(format!( + "cuda_oxide_small_type_parse_gate_{}.ll", + std::process::id() + )); + std::fs::write(&ll_path, &ir).expect("write temp .ll"); + let output = std::process::Command::new(llvm_as) + .arg("-o") + .arg("/dev/null") + .arg(&ll_path) + .output() + .expect("run llvm-as"); + let stderr = String::from_utf8_lossy(&output.stderr).into_owned(); + let _ = std::fs::remove_file(&ll_path); + assert!( + output.status.success(), + "{llvm_as} rejected the emitted module:\n{stderr}\n--- module ---\n{ir}" + ); +} diff --git a/crates/llvm-export/tests/export_test/globals.rs b/crates/llvm-export/tests/export_test/globals.rs new file mode 100644 index 0000000000..f79923cbcf --- /dev/null +++ b/crates/llvm-export/tests/export_test/globals.rs @@ -0,0 +1,478 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use llvm_export::{ + export::{ + NvvmExportConfig, NvvmIrDialect, export_module_to_string, + export_module_to_string_with_config, + }, + ops::{ + GlobalInitializerRelocation, GlobalOp, GlobalOpExt, encode_global_initializer_relocations, + }, + types::{ArrayType, StructLayout, StructType}, +}; +use pliron::{ + builtin::{ + ops::ModuleOp, + types::{IntegerType, Signedness}, + }, + context::Context, + op::Op, +}; + +use crate::common::module_top_block; + +/// Export a module holding one shared global, optionally labelled with the +/// Rust path of the `static` it came from. +fn export_shared_global_with_source_name(source_name: Option<&str>) -> String { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let array_ty = ArrayType::get(&ctx, i32_ty.to_handle(), 64); + let global = GlobalOp::new( + &mut ctx, + "__shared_mem_7".try_into().unwrap(), + array_ty.to_handle(), + ); + global.set_address_space(&mut ctx, 3); + if let Some(source_name) = source_name { + global.set_shared_source_name(&mut ctx, source_name); + } + global.get_operation().insert_at_back(module_block, &ctx); + + export_module_to_string(&ctx, &module).expect("export succeeds") +} + +#[test] +fn shared_global_source_name_is_exported_as_a_comment_above_the_definition() { + let ir = export_shared_global_with_source_name(Some("my_kernel::TILE")); + + let definition_index = ir + .find("@__shared_mem_7 = addrspace(3) global") + .expect("module must declare the shared global"); + let comment_index = ir + .find("; shared source: my_kernel::TILE") + .unwrap_or_else(|| panic!("shared global must name its Rust source:\n{ir}")); + assert!( + comment_index < definition_index, + "the source comment must precede the definition it describes:\n{ir}" + ); +} + +#[test] +fn shared_global_without_a_source_name_exports_no_comment() { + let ir = export_shared_global_with_source_name(None); + + assert!( + ir.contains("@__shared_mem_7 = addrspace(3) global"), + "module must declare the shared global:\n{ir}" + ); + assert!( + !ir.contains("; shared source:"), + "an unlabelled global must not gain a comment:\n{ir}" + ); +} + +#[test] +fn shared_global_source_name_cannot_escape_its_comment_line() { + // A newline in the label would end the comment and leave the remainder to + // be parsed as IR. Nothing in the current pipeline produces such a name, + // so this pins the exporter's own guarantee rather than a live bug. + let ir = export_shared_global_with_source_name(Some("EVIL\n@injected = addrspace(3) global")); + + assert!( + ir.lines().all(|line| !line.starts_with("@injected")), + "a control character in the label must not open a new IR line:\n{ir}" + ); + assert!( + ir.contains("; shared source: EVIL @injected = addrspace(3) global"), + "the label must survive on one line with controls flattened:\n{ir}" + ); +} + +#[test] +fn nvvm_export_rejects_invalid_global_address_spaces() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let global = GlobalOp::new( + &mut ctx, + "thread_local_global".try_into().unwrap(), + i32_ty.to_handle(), + ); + global.set_address_space(&mut ctx, 5); + global.get_operation().insert_at_back(module_block, &ctx); + + let error = export_module_to_string_with_config( + &ctx, + &module, + &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + ) + .expect_err("NVVM module-scope local-memory global must be rejected"); + assert!(error.contains("unsupported address space 5"), "{error}"); + + // The ordinary LLVM/PTX exporter retains its prior behavior; this + // restriction is specifically part of the NVVM IR contract. + assert!(export_module_to_string(&ctx, &module).is_ok()); +} + +#[test] +fn initialized_globals_export_exact_bytes() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "exact_global_bytes".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signless); + + // 0x7fc01234 is a quiet f32 NaN with a non-canonical payload. Treating + // these bytes as an f32 before printing would collapse it to 0x7fc00000. + let nan_ty = ArrayType::get(&ctx, i8_ty.into(), 4); + let nan = GlobalOp::new_with_alignment( + &mut ctx, + "nan_payload".try_into().unwrap(), + nan_ty.into(), + 4, + ); + nan.set_address_space(&mut ctx, 1); + nan.set_initializer_hex(&mut ctx, "3412c07f"); + nan.get_operation().insert_at_back(module_block, &ctx); + + // Byte 0 is a u8, bytes 1..4 are zeroed repr(C) padding, and bytes 4..8 + // are a little-endian u32. The exporter must not recompute those offsets. + let padded_ty = ArrayType::get(&ctx, i8_ty.into(), 8); + let padded = GlobalOp::new_with_alignment( + &mut ctx, + "padded_struct".try_into().unwrap(), + padded_ty.into(), + 4, + ); + padded.set_address_space(&mut ctx, 1); + padded.set_initializer_hex(&mut ctx, "ab00000078563412"); + padded.get_operation().insert_at_back(module_block, &ctx); + + for config in [ + NvvmExportConfig::new(NvvmIrDialect::Modern), + NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + ] { + let ir = export_module_to_string_with_config(&ctx, &module, &config) + .expect("byte-exact global export succeeds"); + assert!( + ir.contains(r#"@nan_payload = addrspace(1) global [4 x i8] c"\34\12\C0\7F", align 4"#), + "NaN payload bytes changed:\n{ir}" + ); + assert!( + ir.contains( + r#"@padded_struct = addrspace(1) global [8 x i8] c"\AB\00\00\00\78\56\34\12", align 4"# + ), + "repr(C) layout bytes changed:\n{ir}" + ); + } +} + +#[test] +fn immutable_globals_export_the_constant_keyword() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "constant_keyword".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signless); + let table_ty = ArrayType::get(&ctx, i8_ty.into(), 4); + + // The compiler's own promoted table: marked never-written, so it must + // export as `constant`. That keyword is the whole point of the marker: + // it is what lets `opt` treat reads as invariant (deleting a copy into a + // stack slot) and what makes `llc` select `ld.global.nc`. + let promoted = GlobalOp::new_with_alignment( + &mut ctx, + "promoted_table".try_into().unwrap(), + table_ty.into(), + 4, + ); + promoted.set_address_space(&mut ctx, 1); + promoted.set_initializer_hex(&mut ctx, "01020304"); + promoted.mark_immutable(&mut ctx); + promoted.get_operation().insert_at_back(module_block, &ctx); + + // An identically shaped global without the marker: the host may still + // write such storage by symbol, so it must keep `global`. Immutability is + // opt-in per global, never inferred from the shape of the initializer. + let plain = GlobalOp::new_with_alignment( + &mut ctx, + "plain_static".try_into().unwrap(), + table_ty.into(), + 4, + ); + plain.set_address_space(&mut ctx, 1); + plain.set_initializer_hex(&mut ctx, "01020304"); + plain.get_operation().insert_at_back(module_block, &ctx); + + for config in [ + NvvmExportConfig::new(NvvmIrDialect::Modern), + NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + ] { + let ir = export_module_to_string_with_config(&ctx, &module, &config) + .expect("immutable global export succeeds"); + assert!( + ir.contains( + r#"@promoted_table = addrspace(1) constant [4 x i8] c"\01\02\03\04", align 4"# + ), + "promoted global lost the constant keyword:\n{ir}" + ); + assert!( + ir.contains(r#"@plain_static = addrspace(1) global [4 x i8] c"\01\02\03\04", align 4"#), + "unmarked global must not become constant:\n{ir}" + ); + } +} + +#[test] +fn initialized_global_exports_static_pointer_relocation() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "static_relocation".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signless); + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); + + // Insert the reference first. Module symbol indexing must make relocation + // resolution independent of textual global order. + let reference_ty = StructType::get_unnamed(&ctx, (vec![i64_ty.into()], StructLayout::Unpacked)); + let reference = GlobalOp::new_with_alignment( + &mut ctx, + "reference".try_into().unwrap(), + reference_ty.into(), + 8, + ); + reference.set_address_space(&mut ctx, 1); + reference.set_source_global_key(&mut ctx, "REFERENCE"); + reference.set_initializer_hex(&mut ctx, "0000000000000000"); + let encoded = encode_global_initializer_relocations(&[GlobalInitializerRelocation { + source_offset: 0, + width_bytes: 8, + target_address_space: 1, + target_addend: 0, + target_key: "TARGET".to_string(), + }]); + reference.set_initializer_relocations(&mut ctx, &encoded); + reference.get_operation().insert_at_back(module_block, &ctx); + + let target_ty = ArrayType::get(&ctx, i8_ty.into(), 4); + let target = + GlobalOp::new_with_alignment(&mut ctx, "target".try_into().unwrap(), target_ty.into(), 4); + target.set_address_space(&mut ctx, 1); + target.set_source_global_key(&mut ctx, "TARGET"); + target.set_initializer_hex(&mut ctx, "78563412"); + target.get_operation().insert_at_back(module_block, &ctx); + + let modern = export_module_to_string_with_config( + &ctx, + &module, + &NvvmExportConfig::new(NvvmIrDialect::Modern), + ) + .expect("modern relocated initializer export succeeds"); + assert!( + modern.contains( + "@reference = addrspace(1) global { i64 } { i64 ptrtoint (ptr addrspacecast (ptr addrspace(1) @target to ptr) to i64) }, align 8" + ), + "{modern}" + ); + + let legacy = export_module_to_string_with_config( + &ctx, + &module, + &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + ) + .expect("legacy relocated initializer export succeeds"); + assert!( + legacy.contains( + "@reference = addrspace(1) global { i64 } { i64 ptrtoint (i8* addrspacecast (i8 addrspace(1)* bitcast ([4 x i8] addrspace(1)* @target to i8 addrspace(1)*) to i8*) to i64) }, align 8" + ), + "{legacy}" + ); +} + +#[test] +fn initialized_global_exports_unaligned_pointer_relocation_as_packed_storage() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "packed_static_relocation".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signless); + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); + + let literal_ty = ArrayType::get(&ctx, i8_ty.into(), 1); + let reference_ty = StructType::get_unnamed( + &ctx, + (vec![literal_ty.into(), i64_ty.into()], StructLayout::Packed), + ); + let reference = GlobalOp::new_with_alignment( + &mut ctx, + "packed_reference".try_into().unwrap(), + reference_ty.into(), + 1, + ); + reference.set_address_space(&mut ctx, 1); + reference.set_source_global_key(&mut ctx, "PACKED_REFERENCE"); + reference.set_initializer_hex(&mut ctx, "7b0000000000000000"); + let encoded = encode_global_initializer_relocations(&[GlobalInitializerRelocation { + source_offset: 1, + width_bytes: 8, + target_address_space: 1, + target_addend: 0, + target_key: "TARGET".to_string(), + }]); + reference.set_initializer_relocations(&mut ctx, &encoded); + reference.get_operation().insert_at_back(module_block, &ctx); + + let target_ty = ArrayType::get(&ctx, i8_ty.into(), 4); + let target = + GlobalOp::new_with_alignment(&mut ctx, "target".try_into().unwrap(), target_ty.into(), 4); + target.set_address_space(&mut ctx, 1); + target.set_source_global_key(&mut ctx, "TARGET"); + target.set_initializer_hex(&mut ctx, "78563412"); + target.get_operation().insert_at_back(module_block, &ctx); + + let modern = export_module_to_string_with_config( + &ctx, + &module, + &NvvmExportConfig::new(NvvmIrDialect::Modern), + ) + .expect("modern packed relocated initializer export succeeds"); + assert!( + modern.contains( + r#"@packed_reference = addrspace(1) global <{ [1 x i8], i64 }> <{ [1 x i8] c"\7B", i64 ptrtoint (ptr addrspacecast (ptr addrspace(1) @target to ptr) to i64) }>, align 1"# + ), + "{modern}" + ); + + let legacy = export_module_to_string_with_config( + &ctx, + &module, + &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + ) + .expect("legacy packed relocated initializer export succeeds"); + assert!( + legacy.contains( + r#"@packed_reference = addrspace(1) global <{ [1 x i8], i64 }> <{ [1 x i8] c"\7B", i64 ptrtoint (i8* addrspacecast (i8 addrspace(1)* bitcast ([4 x i8] addrspace(1)* @target to i8 addrspace(1)*) to i8*) to i64) }>, align 1"# + ), + "{legacy}" + ); +} + +#[test] +fn initialized_global_exports_multiple_relocations_and_addends() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "multiple_static_relocations".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signless); + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); + + let target_a_ty = ArrayType::get(&ctx, i8_ty.into(), 16); + let target_a = GlobalOp::new_with_alignment( + &mut ctx, + "target_a".try_into().unwrap(), + target_a_ty.into(), + 8, + ); + target_a.set_address_space(&mut ctx, 1); + target_a.set_source_global_key(&mut ctx, "TARGET_A"); + target_a.set_initializer_hex(&mut ctx, "000102030405060708090a0b0c0d0e0f"); + target_a.get_operation().insert_at_back(module_block, &ctx); + + let target_b_ty = ArrayType::get(&ctx, i8_ty.into(), 8); + let target_b = GlobalOp::new_with_alignment( + &mut ctx, + "target_b".try_into().unwrap(), + target_b_ty.into(), + 8, + ); + target_b.set_address_space(&mut ctx, 4); + target_b.set_source_global_key(&mut ctx, "TARGET_B"); + target_b.set_initializer_hex(&mut ctx, "1011121314151617"); + target_b.get_operation().insert_at_back(module_block, &ctx); + + let table_ty = StructType::get_unnamed( + &ctx, + (vec![i64_ty.into(), i64_ty.into()], StructLayout::Unpacked), + ); + let table = GlobalOp::new_with_alignment( + &mut ctx, + "reference_table".try_into().unwrap(), + table_ty.into(), + 8, + ); + table.set_address_space(&mut ctx, 1); + table.set_source_global_key(&mut ctx, "REFERENCE_TABLE"); + table.set_initializer_hex(&mut ctx, "00000000000000000000000000000000"); + let encoded = encode_global_initializer_relocations(&[ + GlobalInitializerRelocation { + source_offset: 0, + width_bytes: 8, + target_address_space: 1, + target_addend: 4, + target_key: "TARGET_A".to_string(), + }, + GlobalInitializerRelocation { + source_offset: 8, + width_bytes: 8, + target_address_space: 4, + target_addend: 0, + target_key: "TARGET_B".to_string(), + }, + ]); + table.set_initializer_relocations(&mut ctx, &encoded); + table.get_operation().insert_at_back(module_block, &ctx); + + for dialect in [NvvmIrDialect::Modern, NvvmIrDialect::LegacyLlvm7] { + let ir = + export_module_to_string_with_config(&ctx, &module, &NvvmExportConfig::new(dialect)) + .expect("relocated initializer export succeeds"); + assert!( + ir.contains("@reference_table = addrspace(1) global { i64, i64 }"), + "{ir}" + ); + assert!(ir.contains("getelementptr (i8"), "{ir}"); + assert!(ir.contains("@target_a"), "{ir}"); + assert!(ir.contains("@target_b"), "{ir}"); + assert!(!ir.contains("inttoptr"), "{ir}"); + } +} + +#[test] +fn initialized_global_relocation_rejects_unknown_target_key() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "unknown_relocation_target".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); + let reference_ty = StructType::get_unnamed(&ctx, (vec![i64_ty.into()], StructLayout::Unpacked)); + let reference = GlobalOp::new_with_alignment( + &mut ctx, + "reference".try_into().unwrap(), + reference_ty.into(), + 8, + ); + reference.set_address_space(&mut ctx, 1); + reference.set_source_global_key(&mut ctx, "REFERENCE"); + reference.set_initializer_hex(&mut ctx, "0000000000000000"); + let encoded = encode_global_initializer_relocations(&[GlobalInitializerRelocation { + source_offset: 0, + width_bytes: 8, + target_address_space: 1, + target_addend: 0, + target_key: "MISSING".to_string(), + }]); + reference.set_initializer_relocations(&mut ctx, &encoded); + reference.get_operation().insert_at_back(module_block, &ctx); + + let error = export_module_to_string_with_config( + &ctx, + &module, + &NvvmExportConfig::new(NvvmIrDialect::Modern), + ) + .expect_err("unknown relocation target must fail"); + assert!( + error.contains("unknown rustc global key `MISSING`"), + "{error}" + ); +} diff --git a/crates/llvm-export/tests/export_test/instructions.rs b/crates/llvm-export/tests/export_test/instructions.rs new file mode 100644 index 0000000000..021ddd5364 --- /dev/null +++ b/crates/llvm-export/tests/export_test/instructions.rs @@ -0,0 +1,578 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use llvm_export::{ + export::{ + NvvmExportConfig, NvvmIrDialect, export_module_to_string, + export_module_to_string_with_config, + }, + op_interfaces::CastOpInterface, + ops::{ + AddrSpaceCastOp, AllocaOp, BitcastOp, ConstantOp, FuncOp, GepIndex, GetElementPtrOp, + InlineAsmOp, LoadOp, ReturnOp, SelectOp, StoreOp, + }, + types::{FuncType, PointerType, VoidType}, +}; +use pliron::{ + builtin::{ + attributes::IntegerAttr, + ops::ModuleOp, + types::{FP32Type, IntegerType, Signedness}, + }, + common_traits::Verify, + context::Context, + linked_list::ContainsLinkedList, + op::Op, + utils::apint::APInt, +}; +use std::num::NonZero; + +use crate::common::module_top_block; + +#[test] +fn export_volatile_load_prints_keyword() { + let mut ctx = Context::new(); + + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let ptr_ty = PointerType::get(&ctx, 0); + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![ptr_ty.to_handle()], false); + let func = FuncOp::new(&mut ctx, "volatile_load_test".try_into().unwrap(), func_ty); + let entry = func.get_or_create_entry_block(&mut ctx); + let ptr = entry.deref(&ctx).get_argument(0); + + let load = LoadOp::new(&mut ctx, ptr, i32_ty.to_handle()); + llvm_export::ops::set_op_volatile(&mut ctx, load.get_operation(), true); + load.get_operation().insert_at_back(entry, &ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + let ir = export_module_to_string(&ctx, &module).expect("export succeeds"); + let line = ir + .lines() + .find(|line| line.contains("load volatile")) + .expect("volatile load line"); + + assert!( + line.trim_start().contains(" = load volatile i32, ptr "), + "volatile load keyword must appear immediately after load:\n{ir}" + ); +} + +#[test] +fn export_volatile_store_prints_keyword() { + let mut ctx = Context::new(); + + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let ptr_ty = PointerType::get(&ctx, 0); + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get( + &ctx, + void_ty.to_handle(), + vec![ptr_ty.to_handle(), i32_ty.to_handle()], + false, + ); + let func = FuncOp::new(&mut ctx, "volatile_store_test".try_into().unwrap(), func_ty); + let entry = func.get_or_create_entry_block(&mut ctx); + let ptr = entry.deref(&ctx).get_argument(0); + let val = entry.deref(&ctx).get_argument(1); + + let store = StoreOp::new(&mut ctx, val, ptr); + llvm_export::ops::set_op_volatile(&mut ctx, store.get_operation(), true); + store.get_operation().insert_at_back(entry, &ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + let ir = export_module_to_string(&ctx, &module).expect("export succeeds"); + let line = ir + .lines() + .find(|line| line.contains("store volatile")) + .expect("volatile store line"); + + assert!( + line.trim_start().starts_with("store volatile i32 "), + "volatile store keyword must appear immediately after store:\n{ir}" + ); +} + +#[test] +fn legacy_export_uses_one_canonical_pointer_with_multiple_typed_views() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "legacy_views".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + + let ptr_ty = PointerType::get(&ctx, 0); + let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signless); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let f32_ty = FP32Type::get(&ctx); + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.into(), vec![ptr_ty.into()], false); + let func = FuncOp::new(&mut ctx, "multiple_views".try_into().unwrap(), func_ty); + let entry = func.get_or_create_entry_block(&mut ctx); + let pointer = entry.deref(&ctx).get_argument(0); + + LoadOp::new(&mut ctx, pointer, i32_ty.into()) + .get_operation() + .insert_at_back(entry, &ctx); + LoadOp::new(&mut ctx, pointer, f32_ty.into()) + .get_operation() + .insert_at_back(entry, &ctx); + LoadOp::new(&mut ctx, pointer, i8_ty.into()) + .get_operation() + .insert_at_back(entry, &ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + let config = NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7); + let ir = export_module_to_string_with_config(&ctx, &module, &config) + .expect("legacy export succeeds"); + + assert!(ir.contains("define void @multiple_views(i8* %v0)"), "{ir}"); + assert!(ir.contains("bitcast i8* %v0 to i32*"), "{ir}"); + assert!(ir.contains("load i32, i32*"), "{ir}"); + assert!(ir.contains("bitcast i8* %v0 to float*"), "{ir}"); + assert!(ir.contains("load float, float*"), "{ir}"); + assert!(ir.contains("load i8, i8* %v0"), "{ir}"); + assert!(!ir.contains("bitcast i8* %v0 to i8*"), "{ir}"); + assert!( + !ir.split(|c: char| !c.is_ascii_alphanumeric()) + .any(|t| t == "ptr") + ); +} + +#[test] +fn legacy_alloca_rejects_a_non_default_result_address_space() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "invalid_alloca_as".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.into(), vec![], false); + let func = FuncOp::new(&mut ctx, "invalid_alloca".try_into().unwrap(), func_ty); + let entry = func.get_or_create_entry_block(&mut ctx); + + let one_attr = IntegerAttr::new(i32_ty, APInt::from_u32(1, NonZero::new(32).unwrap())); + let one = ConstantOp::new(&mut ctx, one_attr.into()); + let one_value = one.get_operation().deref(&ctx).get_result(0); + one.get_operation().insert_at_back(entry, &ctx); + let alloca = AllocaOp::new(&mut ctx, i32_ty.into(), one_value); + let alloca_result = alloca.get_operation().deref(&ctx).get_result(0); + let shared_pointer = PointerType::get(&ctx, 3); + alloca_result.set_type(&ctx, shared_pointer.into()); + alloca.get_operation().insert_at_back(entry, &ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + module + .get_operation() + .deref(&ctx) + .verify(&ctx) + .expect("upstream verification currently does not enforce alloca result AS0"); + let error = export_module_to_string_with_config( + &ctx, + &module, + &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + ) + .expect_err("legacy export must reject an alloca address-space mismatch"); + assert!( + error.contains("alloca result uses address space 3"), + "{error}" + ); +} + +#[test] +fn legacy_gep_rejects_a_result_address_space_different_from_its_base() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "invalid_gep_as".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + + let global_pointer = PointerType::get(&ctx, 1); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.into(), vec![global_pointer.into()], false); + let func = FuncOp::new(&mut ctx, "invalid_gep".try_into().unwrap(), func_ty); + let entry = func.get_or_create_entry_block(&mut ctx); + let base = entry.deref(&ctx).get_argument(0); + let gep = GetElementPtrOp::new(&mut ctx, base, vec![GepIndex::Constant(0)], i32_ty.into()); + let gep_result = gep.get_operation().deref(&ctx).get_result(0); + let shared_pointer = PointerType::get(&ctx, 3); + gep_result.set_type(&ctx, shared_pointer.into()); + gep.get_operation().insert_at_back(entry, &ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + module + .get_operation() + .deref(&ctx) + .verify(&ctx) + .expect("upstream verification currently does not enforce GEP result/base AS equality"); + let error = export_module_to_string_with_config( + &ctx, + &module, + &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + ) + .expect_err("legacy export must reject a GEP address-space mismatch"); + assert!( + error.contains("GEP result address-space mismatch: base is 1, result is 3"), + "{error}" + ); +} + +#[test] +fn gep_inbounds_marker_controls_exported_pointer_semantics() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "gep_semantics".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + + let pointer = PointerType::get(&ctx, 0); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.into(), vec![pointer.into()], false); + let func = FuncOp::new(&mut ctx, "offsets".try_into().unwrap(), func_ty); + let entry = func.get_or_create_entry_block(&mut ctx); + let base = entry.deref(&ctx).get_argument(0); + + let ordinary = GetElementPtrOp::new(&mut ctx, base, vec![GepIndex::Constant(1)], i32_ty.into()); + ordinary.get_operation().insert_at_back(entry, &ctx); + + let wrapping = GetElementPtrOp::new(&mut ctx, base, vec![GepIndex::Constant(2)], i32_ty.into()); + llvm_export::ops::set_gep_inbounds(&mut ctx, wrapping.get_operation(), false); + wrapping.get_operation().insert_at_back(entry, &ctx); + + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + let ir = export_module_to_string(&ctx, &module).expect("GEP export succeeds"); + let gep_lines: Vec<_> = ir + .lines() + .filter(|line| line.contains("getelementptr")) + .collect(); + assert_eq!(gep_lines.len(), 2, "{ir}"); + assert!(gep_lines[0].contains("getelementptr inbounds"), "{ir}"); + assert!( + gep_lines[1].contains("getelementptr i32") + && !gep_lines[1].contains("getelementptr inbounds"), + "{ir}" + ); +} + +#[test] +fn legacy_pointer_select_keeps_one_canonical_type() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "legacy_pointer_select".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + + let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); + let ptr_ty = PointerType::get(&ctx, 0); + let func_ty = FuncType::get( + &ctx, + ptr_ty.into(), + vec![i1_ty.into(), ptr_ty.into(), ptr_ty.into()], + false, + ); + let func = FuncOp::new(&mut ctx, "choose_pointer".try_into().unwrap(), func_ty); + let entry = func.get_or_create_entry_block(&mut ctx); + let condition = entry.deref(&ctx).get_argument(0); + let if_true = entry.deref(&ctx).get_argument(1); + let if_false = entry.deref(&ctx).get_argument(2); + let select = SelectOp::new(&mut ctx, condition, if_true, if_false); + let selected = select.get_operation().deref(&ctx).get_result(0); + select.get_operation().insert_at_back(entry, &ctx); + ReturnOp::new(&mut ctx, Some(selected)) + .get_operation() + .insert_at_back(entry, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + let ir = export_module_to_string_with_config( + &ctx, + &module, + &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + ) + .expect("legacy pointer select export succeeds"); + assert!( + ir.contains("select i1 %v0, i8* %v1, i8* %v2"), + "pointer select must use the canonical byte-pointer type:\n{ir}" + ); + assert!(ir.contains("ret i8*"), "{ir}"); +} + +#[test] +fn pointer_bitcast_cannot_cross_address_spaces_in_either_nvvm_dialect() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "invalid_pointer_bitcast".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let global_pointer = PointerType::get(&ctx, 1); + let shared_pointer = PointerType::get(&ctx, 3); + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.into(), vec![global_pointer.into()], false); + let func = FuncOp::new(&mut ctx, "invalid_cast".try_into().unwrap(), func_ty); + let entry = func.get_or_create_entry_block(&mut ctx); + let input = entry.deref(&ctx).get_argument(0); + BitcastOp::new(&mut ctx, input, shared_pointer.into()) + .get_operation() + .insert_at_back(entry, &ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + module + .get_operation() + .deref(&ctx) + .verify(&ctx) + .expect("upstream bitcast verification currently does not enforce pointer AS equality"); + for dialect in [NvvmIrDialect::LegacyLlvm7, NvvmIrDialect::Modern] { + let error = + export_module_to_string_with_config(&ctx, &module, &NvvmExportConfig::new(dialect)) + .expect_err("a cross-address-space pointer bitcast must be rejected"); + assert!( + error.contains("pointer bitcast cannot cross address spaces 1 -> 3"), + "{dialect:?}: {error}" + ); + } +} + +#[test] +fn addrspacecast_must_change_address_spaces_in_either_nvvm_dialect() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "invalid_addrspacecast".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let shared_pointer = PointerType::get(&ctx, 3); + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.into(), vec![shared_pointer.into()], false); + let func = FuncOp::new(&mut ctx, "invalid_cast".try_into().unwrap(), func_ty); + let entry = func.get_or_create_entry_block(&mut ctx); + let input = entry.deref(&ctx).get_argument(0); + AddrSpaceCastOp::new(&mut ctx, input, shared_pointer.into()) + .get_operation() + .insert_at_back(entry, &ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + module + .get_operation() + .deref(&ctx) + .verify(&ctx) + .expect("upstream addrspacecast verification currently permits equal address spaces"); + for dialect in [NvvmIrDialect::LegacyLlvm7, NvvmIrDialect::Modern] { + let error = + export_module_to_string_with_config(&ctx, &module, &NvvmExportConfig::new(dialect)) + .expect_err("addrspacecast must not encode a no-op address-space conversion"); + assert!( + error.contains( + "addrspacecast must change address spaces; source and result are both address space 3" + ), + "{dialect:?}: {error}" + ); + } +} + +#[test] +fn legacy_pointer_slot_is_recursively_canonical() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "legacy_pointer_slot".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let ptr_ty = PointerType::get(&ctx, 0); + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.into(), vec![ptr_ty.into()], false); + let func = FuncOp::new(&mut ctx, "pointer_slot".try_into().unwrap(), func_ty); + let entry = func.get_or_create_entry_block(&mut ctx); + let incoming = entry.deref(&ctx).get_argument(0); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let one_attr = IntegerAttr::new(i32_ty, APInt::from_u32(1, NonZero::new(32).unwrap())); + let one = ConstantOp::new(&mut ctx, one_attr.into()); + let one_value = one.get_operation().deref(&ctx).get_result(0); + one.get_operation().insert_at_back(entry, &ctx); + + let slot = AllocaOp::new(&mut ctx, ptr_ty.into(), one_value); + let slot_value = slot.get_operation().deref(&ctx).get_result(0); + slot.get_operation().insert_at_back(entry, &ctx); + StoreOp::new(&mut ctx, incoming, slot_value) + .get_operation() + .insert_at_back(entry, &ctx); + LoadOp::new(&mut ctx, slot_value, ptr_ty.into()) + .get_operation() + .insert_at_back(entry, &ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + let ir = export_module_to_string_with_config( + &ctx, + &module, + &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + ) + .expect("legacy pointer-slot export succeeds"); + assert!(ir.contains("alloca i8*"), "{ir}"); + assert!(ir.contains("bitcast i8**"), "{ir}"); + assert!(ir.matches("to i8**").count() >= 2, "{ir}"); + assert!(ir.contains("store i8* %v0, i8**"), "{ir}"); + assert!(ir.contains("load i8*, i8**"), "{ir}"); +} + +#[test] +fn export_inline_asm_respects_sideeffect_marker() { + let mut ctx = Context::new(); + + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_region = module.get_operation().deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); + let func = FuncOp::new(&mut ctx, "has_inline_asm".try_into().unwrap(), func_ty); + let entry = func.get_or_create_entry_block(&mut ctx); + + let default_asm = InlineAsmOp::new(&mut ctx, void_ty.into(), vec![], "bar.sync 0;", "", false); + default_asm.get_operation().insert_at_back(entry, &ctx); + + let register_only_asm = InlineAsmOp::new(&mut ctx, void_ty.into(), vec![], "nop;", "", true); + llvm_export::ops::set_inline_asm_sideeffect(&mut ctx, register_only_asm.get_operation(), false); + register_only_asm + .get_operation() + .insert_at_back(entry, &ctx); + + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + let ir = export_module_to_string(&ctx, &module).expect("export succeeds"); + + assert!( + ir.contains("call void asm sideeffect \"bar.sync 0;\", \"\"()"), + "inline asm without an explicit marker should remain conservative:\n{ir}" + ); + assert!( + ir.contains("call void asm \"nop;\", \"\"() #0"), + "inline asm marked sideeffect=false should omit the keyword while preserving convergent:\n{ir}" + ); + assert!( + ir.contains("attributes #0 = { convergent }"), + "convergent inline asm must emit the convergent attr group:\n{ir}" + ); +} + +#[test] +fn export_inline_asm_escapes_llvm_string_literals() { + let mut ctx = Context::new(); + + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_region = module.get_operation().deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); + let func = FuncOp::new( + &mut ctx, + "has_escaped_inline_asm".try_into().unwrap(), + func_ty, + ); + let entry = func.get_or_create_entry_block(&mut ctx); + + let asm = InlineAsmOp::new( + &mut ctx, + void_ty.into(), + vec![], + "mov.u32 $0, %laneid;\n// \"quoted\" \\22", + "~{memory}\\raw", + false, + ); + asm.get_operation().insert_at_back(entry, &ctx); + + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + let ir = export_module_to_string(&ctx, &module).expect("export succeeds"); + + assert!( + ir.contains( + "call void asm sideeffect \"mov.u32 $0, %laneid;\\0A// \\22quoted\\22 \\5C22\", \"~{memory}\\5Craw\"()" + ), + "inline asm template and constraints must be escaped as LLVM string literals:\n{ir}" + ); +} + +/// A float binop carrying `FastmathFlags` must export with the matching LLVM +/// fast-math keyword (`fast` for the all-bits set), while a float binop with no +/// flags must export with none. Regression guard: the textual exporter +/// previously dropped fast-math flags entirely, making the `f*_fast` intrinsic +/// lowering inert end to end. +#[test] +fn export_emits_fast_math_flags_only_on_flagged_float_ops() { + use llvm_export::attributes::FastmathFlags; + use llvm_export::op_interfaces::{BinArithOp, FloatBinArithOpWithFastMathFlags}; + use llvm_export::ops::{FAddOp, FMulOp}; + use pliron::builtin::types::FP32Type; + + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_region = module.get_operation().deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + + let f32_ty = FP32Type::get(&ctx); + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get( + &ctx, + void_ty.to_handle(), + vec![f32_ty.into(), f32_ty.into()], + false, + ); + let func = FuncOp::new(&mut ctx, "fast_math".try_into().unwrap(), func_ty); + let entry = func.get_or_create_entry_block(&mut ctx); + let a = entry.deref(&ctx).get_argument(0); + let b = entry.deref(&ctx).get_argument(1); + + // fadd with the full fast-math set. + let fadd = FAddOp::new_with_fast_math_flags(&mut ctx, a, b, FastmathFlags::FAST.into()); + fadd.get_operation().insert_at_back(entry, &ctx); + // fmul with no flags: must stay flag-free. + let fmul = FMulOp::new(&mut ctx, a, b); + fmul.get_operation().insert_at_back(entry, &ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + let ir = export_module_to_string(&ctx, &module).expect("export succeeds"); + + assert!( + ir.contains("fadd fast float"), + "fast-math fadd must export the `fast` keyword:\n{ir}" + ); + let fmul_line = ir + .lines() + .find(|line| line.contains("fmul")) + .expect("exported fmul line"); + assert!( + !fmul_line.contains("fast"), + "a float binop with no fast-math flags must not gain them:\n{ir}" + ); +} diff --git a/crates/llvm-export/tests/export_test/kernels_and_roots.rs b/crates/llvm-export/tests/export_test/kernels_and_roots.rs new file mode 100644 index 0000000000..e6490122cf --- /dev/null +++ b/crates/llvm-export/tests/export_test/kernels_and_roots.rs @@ -0,0 +1,283 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use llvm_export::{ + export::{ + DeviceExternDecl, NvvmExportConfig, NvvmIrDialect, PtxExportConfig, + export_module_to_string, export_module_to_string_with_config, + export_module_with_externs_and_roots, + }, + ops::{FuncOp, GlobalOp, GlobalOpExt, ReturnOp}, + types::{FuncType, PointerType, VoidType}, +}; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::StringAttr, + ops::ModuleOp, + types::{IntegerType, Signedness}, + }, + context::{Context, Ptr}, + identifier::Identifier, + op::Op, +}; + +use crate::common::module_top_block; + +#[test] +fn legacy_kernel_metadata_uses_typed_function_references() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "legacy_metadata".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let ptr_ty = PointerType::get(&ctx, 0); + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.into(), vec![ptr_ty.into()], false); + let func = FuncOp::new(&mut ctx, "metadata_kernel".try_into().unwrap(), func_ty); + func.get_operation().deref_mut(&ctx).attributes.set( + "gpu_kernel".try_into().unwrap(), + StringAttr::new("true".into()), + ); + let entry = func.get_or_create_entry_block(&mut ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + let config = NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7); + let ir = export_module_to_string_with_config(&ctx, &module, &config) + .expect("legacy metadata export succeeds"); + assert!( + ir.contains( + "@llvm.used = appending global [1 x i8*] [i8* bitcast (void (i8*)* @metadata_kernel to i8*)]" + ), + "{ir}" + ); + assert!( + ir.contains("!{void (i8*)* @metadata_kernel, !\"kernel\", i32 1}"), + "{ir}" + ); +} + +#[test] +fn llvm_used_roots_only_explicitly_retained_globals() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "retained_globals".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + + let retained = GlobalOp::new(&mut ctx, "IKET_META".try_into().unwrap(), i32_ty.into()); + retained.set_address_space(&mut ctx, 1); + retained.mark_retained(&mut ctx); + retained.get_operation().insert_at_back(module_block, &ctx); + + let ordinary = GlobalOp::new(&mut ctx, "ORDINARY".try_into().unwrap(), i32_ty.into()); + ordinary.set_address_space(&mut ctx, 1); + ordinary.get_operation().insert_at_back(module_block, &ctx); + + let ir = export_module_to_string_with_config(&ctx, &module, &PtxExportConfig) + .expect("retained global export succeeds"); + assert!( + ir.contains( + "@llvm.used = appending global [1 x ptr] [ptr addrspacecast (ptr addrspace(1) @IKET_META to ptr)], section \"llvm.metadata\"" + ), + "{ir}" + ); + assert!(!ir.contains("@ORDINARY to ptr"), "{ir}"); +} + +#[test] +fn legacy_llvm_used_roots_retained_address_space_global() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "legacy_retained_global".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let retained = GlobalOp::new(&mut ctx, "IKET_META".try_into().unwrap(), i32_ty.into()); + retained.set_address_space(&mut ctx, 1); + retained.mark_retained(&mut ctx); + retained.get_operation().insert_at_back(module_block, &ctx); + + let config = NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7); + let ir = export_module_to_string_with_config(&ctx, &module, &config) + .expect("legacy retained global export succeeds"); + assert!( + ir.contains( + "@llvm.used = appending global [1 x i8*] [i8* addrspacecast (i32 addrspace(1)* @IKET_META to i8*)], section \"llvm.metadata\"" + ), + "{ir}" + ); +} + +#[test] +fn ptx_export_records_kernel_roots_for_internalization() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "ptx_roots".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let global = GlobalOp::new(&mut ctx, "COEFFS".try_into().unwrap(), i32_ty.into()); + global.set_address_space(&mut ctx, 4); + global.get_operation().insert_at_back(module_block, &ctx); + let func_ty = FuncType::get(&ctx, VoidType::get(&ctx).into(), vec![], false); + let func = FuncOp::new(&mut ctx, "entry_kernel".try_into().unwrap(), func_ty); + func.get_operation().deref_mut(&ctx).attributes.set( + "gpu_kernel".try_into().unwrap(), + StringAttr::new("true".into()), + ); + let entry = func.get_or_create_entry_block(&mut ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + let exported = export_module_with_externs_and_roots::( + &ctx, + &module, + &[], + &PtxExportConfig, + ) + .expect("PTX export succeeds"); + let ir = exported.llvm_ir; + assert!( + ir.contains( + "@llvm.used = appending global [1 x ptr] [ptr @entry_kernel], section \"llvm.metadata\"" + ), + "{ir}" + ); + assert_eq!(exported.public_symbols, ["COEFFS", "entry_kernel"]); +} + +#[test] +fn ptx_export_records_standalone_device_function_roots_for_internalization() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "ptx_device_root".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let func_ty = FuncType::get(&ctx, VoidType::get(&ctx).into(), vec![], false); + let prefixed_name = format!( + "{}standalone_export", + reserved_oxide_symbols::LEGACY_DEVICE_PREFIX + ); + let func = FuncOp::new( + &mut ctx, + prefixed_name.as_str().try_into().unwrap(), + func_ty, + ); + let entry = func.get_or_create_entry_block(&mut ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + func.get_operation().insert_at_back(module_block, &ctx); + + let exported = export_module_with_externs_and_roots::( + &ctx, + &module, + &[], + &PtxExportConfig, + ) + .expect("standalone PTX export succeeds"); + assert_eq!(exported.public_symbols, ["standalone_export"]); + assert!( + exported.llvm_ir.contains( + "@llvm.used = appending global [1 x ptr] [ptr @standalone_export], section \"llvm.metadata\"" + ), + "{}", + exported.llvm_ir + ); +} + +/// Builds a `void` function taking pointer parameters in the given address +/// spaces, with an empty body, optionally carrying the `gpu_kernel` marker +/// the exporter uses to distinguish an entry from a device function. +fn pointer_param_function( + ctx: &mut Context, + module_block: Ptr, + name: &str, + address_spaces: &[u32], + is_kernel: bool, +) { + let params: Vec<_> = address_spaces + .iter() + .map(|space| PointerType::get(ctx, *space).into()) + .collect(); + let void_ty = VoidType::get(ctx); + let func_ty = FuncType::get(ctx, void_ty.into(), params, false); + let func = FuncOp::new(ctx, name.try_into().unwrap(), func_ty); + if is_kernel { + let kernel_key: Identifier = "gpu_kernel".try_into().unwrap(); + func.get_operation() + .deref_mut(ctx) + .attributes + .set(kernel_key, StringAttr::new("true".to_string())); + } + let entry = func.get_or_create_entry_block(ctx); + ReturnOp::new(ctx, None) + .get_operation() + .insert_at_back(entry, ctx); + func.get_operation().insert_at_back(module_block, ctx); +} + +/// A kernel receives its parameters in `.param` space from the host, which +/// holds no shared-memory address to pass. The exporter must refuse the +/// signature instead of emitting `.ptr .shared`, which ptxas accepts but the +/// driver rejects at module load, taking every kernel in the module down. +#[test] +fn kernel_rejects_a_shared_memory_pointer_parameter() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "kernel_shared_param".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + pointer_param_function(&mut ctx, module_block, "shared_param", &[0, 3], true); + + let error = export_module_to_string(&ctx, &module) + .expect_err("export must refuse a kernel parameter in shared memory"); + assert!( + error.contains( + "kernel `@shared_param` parameter 1 is a pointer into shared memory (address space 3)" + ), + "{error}" + ); +} + +/// Local memory is per-thread, so the host has no address in it to pass +/// either; the exporter refuses it on the same grounds as shared. +#[test] +fn kernel_rejects_a_local_memory_pointer_parameter() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "kernel_local_param".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + pointer_param_function(&mut ctx, module_block, "local_param", &[5], true); + + let error = export_module_to_string(&ctx, &module) + .expect_err("export must refuse a kernel parameter in local memory"); + assert!( + error.contains( + "kernel `@local_param` parameter 0 is a pointer into local memory (address space 5)" + ), + "{error}" + ); +} + +/// The refusal is scoped to what the host cannot supply. Generic, global, +/// and constant pointer parameters are host-addressable and stay accepted on +/// kernels, and a device (non-kernel) function may carry any state space on +/// its parameters, shared included. +#[test] +fn kernel_keeps_host_addressable_pointer_parameters_and_device_functions_keep_shared() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "kernel_allowed_params".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + pointer_param_function(&mut ctx, module_block, "allowed_param", &[0, 1, 4], true); + pointer_param_function(&mut ctx, module_block, "device_shared", &[3], false); + + let ir = export_module_to_string(&ctx, &module) + .expect("generic/global/constant kernel parameters and device shared must export"); + assert!( + ir.contains( + "define ptx_kernel void @allowed_param(ptr %v0, ptr addrspace(1) %v1, ptr addrspace(4) %v2)" + ), + "{ir}" + ); + assert!( + ir.contains("define void @device_shared(ptr addrspace(3) %v0)"), + "{ir}" + ); +} diff --git a/crates/llvm-export/tests/export_test/main.rs b/crates/llvm-export/tests/export_test/main.rs new file mode 100644 index 0000000000..9b77233528 --- /dev/null +++ b/crates/llvm-export/tests/export_test/main.rs @@ -0,0 +1,14 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +mod calls_and_addresses; +mod common; +mod control_flow; +mod debug_info; +mod device_externs; +mod globals; +mod instructions; +mod kernels_and_roots; +mod metadata; diff --git a/crates/llvm-export/tests/export_test/metadata.rs b/crates/llvm-export/tests/export_test/metadata.rs new file mode 100644 index 0000000000..0330a9f420 --- /dev/null +++ b/crates/llvm-export/tests/export_test/metadata.rs @@ -0,0 +1,231 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use llvm_export::{ + export::{ + DebugKind, NvvmExportConfig, PtxExportConfig, export_module_to_string, + export_module_to_string_with_config, + }, + ops::{FuncOp, ReturnOp}, + types::{FuncType, VoidType}, +}; +use pliron::{ + builtin::{ + attributes::{IntegerAttr, StringAttr}, + ops::ModuleOp, + types::{IntegerType, Signedness}, + }, + context::Context, + identifier::Identifier, + linked_list::ContainsLinkedList, + location::Located, + op::Op, + utils::apint::APInt, +}; +use std::num::NonZero; + +use crate::common::{DebugConfig, module_top_block, src_location}; + +#[test] +fn nvvm_metadata_version_uses_next_allocated_metadata_id() { + let mut ctx = Context::new(); + + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_region = module.get_operation().deref(&ctx).get_region(0); + let module_block = { + let region = module_region.deref(&ctx); + region.iter(&ctx).next().unwrap() + }; + + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); + let func = FuncOp::new(&mut ctx, "bounded_kernel".try_into().unwrap(), func_ty); + let entry = func.get_or_create_entry_block(&mut ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let width = NonZero::new(32).unwrap(); + let max_threads = IntegerAttr::new(u32_ty, APInt::from_u32(256, width)); + let min_blocks = IntegerAttr::new(u32_ty, APInt::from_u32(2, width)); + + { + let attrs = &mut func.get_operation().deref_mut(&ctx).attributes; + attrs.set( + Identifier::try_from("gpu_kernel").unwrap(), + StringAttr::new("true".into()), + ); + attrs.set(Identifier::try_from("maxntid").unwrap(), max_threads); + attrs.set(Identifier::try_from("minctasm").unwrap(), min_blocks); + } + + func.get_operation().insert_at_back(module_block, &ctx); + + let ir = export_module_to_string_with_config(&ctx, &module, &NvvmExportConfig::default()) + .expect("NVVM export succeeds"); + + assert!( + ir.contains("!0 = !{ptr @bounded_kernel, !\"kernel\", i32 1}"), + "a launch-bounded kernel still needs its kernel annotation:\n{ir}" + ); + assert!( + ir.contains("!nvvm.annotations = !{!0, !1, !2, !3, !4}"), + "kernel identity plus launch-bounds annotations should occupy !0..!4:\n{ir}" + ); + assert!( + ir.contains("!nvvmir.version = !{!5}\n!5 = !{i32 2, i32 0, i32 3, i32 2}"), + "version metadata should use the next allocated ID:\n{ir}" + ); +} + +#[test] +fn export_alwaysinline_function_attribute_uses_llvm_define_syntax() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); + let func = FuncOp::new(&mut ctx, "inline_helper".try_into().unwrap(), func_ty); + let entry = func.get_or_create_entry_block(&mut ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + + let key: pliron::identifier::Identifier = "alwaysinline".try_into().unwrap(); + func.get_operation() + .deref_mut(&ctx) + .attributes + .set(key, StringAttr::new("true".to_string())); + func.get_operation().insert_at_back(module_block, &ctx); + + let ir = export_module_to_string(&ctx, &module).expect("export succeeds"); + let define_line = ir + .lines() + .find(|line| line.starts_with("define void @inline_helper(")) + .expect("inline helper definition"); + assert_eq!( + define_line, "define void @inline_helper() alwaysinline #0 {", + "`alwaysinline` must be emitted after the parameter list, before attr group #0:\n{ir}" + ); + assert!( + ir.contains("attributes #0 = { convergent }"), + "convergent attribute group must still be emitted:\n{ir}" + ); +} + +#[test] +fn export_alwaysinline_coexists_with_debug_scope() { + // alwaysinline and the !dbg scope are emitted on the same define line and + // must not crowd each other out. This guards the 4-way emission: a future + // change that drops either one when both are present fails here. + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); + let func = FuncOp::new(&mut ctx, "inline_helper".try_into().unwrap(), func_ty); + let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 7, 1); + func.get_operation().deref_mut(&ctx).set_loc(func_loc); + let entry = func.get_or_create_entry_block(&mut ctx); + let ret = ReturnOp::new(&mut ctx, None); + let ret_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 8, 5); + ret.get_operation().deref_mut(&ctx).set_loc(ret_loc); + ret.get_operation().insert_at_back(entry, &ctx); + + let key: pliron::identifier::Identifier = "alwaysinline".try_into().unwrap(); + func.get_operation() + .deref_mut(&ctx) + .attributes + .set(key, StringAttr::new("true".to_string())); + func.get_operation().insert_at_back(module_block, &ctx); + + let config = DebugConfig { + inner: PtxExportConfig, + debug_kind: DebugKind::LineTables, + }; + let ir = + export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); + let define_line = ir + .lines() + .find(|line| line.starts_with("define void @inline_helper(")) + .expect("inline helper definition"); + assert!( + define_line.contains("alwaysinline"), + "alwaysinline must survive when debug info is on:\n{ir}" + ); + assert!( + define_line.contains("!dbg !"), + "!dbg scope must survive when alwaysinline is present:\n{ir}" + ); +} + +#[test] +fn debug_metadata_shares_allocator_with_nvvm_metadata() { + let mut ctx = Context::new(); + + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_region = module.get_operation().deref(&ctx).get_region(0); + let module_block = { + let region = module_region.deref(&ctx); + region.iter(&ctx).next().unwrap() + }; + + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); + let func = FuncOp::new(&mut ctx, "debug_kernel".try_into().unwrap(), func_ty); + let func_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 10, 1); + func.get_operation().deref_mut(&ctx).set_loc(func_loc); + + { + let attrs = &mut func.get_operation().deref_mut(&ctx).attributes; + attrs.set( + Identifier::try_from("gpu_kernel").unwrap(), + StringAttr::new("true".into()), + ); + } + + let entry = func.get_or_create_entry_block(&mut ctx); + let ret = ReturnOp::new(&mut ctx, None); + let ret_loc = src_location(&mut ctx, "/tmp/cuda-oxide/tests/kernel.rs", 11, 5); + ret.get_operation().deref_mut(&ctx).set_loc(ret_loc); + ret.get_operation().insert_at_back(entry, &ctx); + + func.get_operation().insert_at_back(module_block, &ctx); + + let config = DebugConfig { + inner: NvvmExportConfig::default(), + debug_kind: DebugKind::LineTables, + }; + let ir = export_module_to_string_with_config(&ctx, &module, &config) + .expect("debug NVVM export succeeds"); + + assert!( + ir.contains("!0 = !DIFile(filename: \"kernel.rs\", directory: \"/tmp/cuda-oxide/tests\")"), + "debug file node should take the first metadata ID:\n{ir}" + ); + assert!( + ir.contains("!4 = !DILocation(line: 11, column: 5, scope: !3)"), + "instruction location should be allocated before NVVM metadata:\n{ir}" + ); + assert!( + ir.contains("!5 = !{ptr @debug_kernel, !\"kernel\", i32 1}"), + "NVVM annotations should continue after debug metadata:\n{ir}" + ); + assert!( + ir.contains("!nvvm.annotations = !{!5}"), + "named NVVM metadata should reference its allocated node:\n{ir}" + ); + assert!( + ir.contains("!nvvmir.version = !{!6}\n!6 = !{i32 2, i32 0, i32 3, i32 2}"), + "NVVM version should use the next free metadata ID:\n{ir}" + ); + assert!( + ir.contains("!llvm.module.flags = !{!7, !8}"), + "debug module flags should also use the shared allocator:\n{ir}" + ); +} From f69ba40ae5194b2cddf823a6784a70825c31878e Mon Sep 17 00:00:00 2001 From: nihalpasham Date: Fri, 28 Aug 2026 22:19:56 +0530 Subject: [PATCH 04/18] refactor(mir-lower): split tests/lowering_test.rs into lowering_test/ modules --- crates/mir-lower/tests/lowering_test.rs | 13177 ---------------- .../tests/lowering_test/atomics_fences.rs | 799 + .../tests/lowering_test/barriers_sync.rs | 743 + .../tests/lowering_test/calls_and_values.rs | 897 ++ .../mir-lower/tests/lowering_test/common.rs | 101 + .../mir-lower/tests/lowering_test/cp_async.rs | 727 + .../tests/lowering_test/inline_ptx.rs | 374 + crates/mir-lower/tests/lowering_test/main.rs | 21 + .../tests/lowering_test/math_conversions.rs | 1712 ++ .../tests/lowering_test/matrix_memory.rs | 1214 ++ crates/mir-lower/tests/lowering_test/mma.rs | 1754 ++ .../tests/lowering_test/sregs_and_warp.rs | 1618 ++ .../tests/lowering_test/wgmma_lowering.rs | 2658 ++++ .../tests/lowering_test/wgmma_rejections.rs | 747 + 14 files changed, 13365 insertions(+), 13177 deletions(-) delete mode 100644 crates/mir-lower/tests/lowering_test.rs create mode 100644 crates/mir-lower/tests/lowering_test/atomics_fences.rs create mode 100644 crates/mir-lower/tests/lowering_test/barriers_sync.rs create mode 100644 crates/mir-lower/tests/lowering_test/calls_and_values.rs create mode 100644 crates/mir-lower/tests/lowering_test/common.rs create mode 100644 crates/mir-lower/tests/lowering_test/cp_async.rs create mode 100644 crates/mir-lower/tests/lowering_test/inline_ptx.rs create mode 100644 crates/mir-lower/tests/lowering_test/main.rs create mode 100644 crates/mir-lower/tests/lowering_test/math_conversions.rs create mode 100644 crates/mir-lower/tests/lowering_test/matrix_memory.rs create mode 100644 crates/mir-lower/tests/lowering_test/mma.rs create mode 100644 crates/mir-lower/tests/lowering_test/sregs_and_warp.rs create mode 100644 crates/mir-lower/tests/lowering_test/wgmma_lowering.rs create mode 100644 crates/mir-lower/tests/lowering_test/wgmma_rejections.rs diff --git a/crates/mir-lower/tests/lowering_test.rs b/crates/mir-lower/tests/lowering_test.rs deleted file mode 100644 index e393cee672..0000000000 --- a/crates/mir-lower/tests/lowering_test.rs +++ /dev/null @@ -1,13177 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -// Tests build kinded fixture types directly; production minting lives in -// mir-importer's facts.rs (see the workspace clippy.toml disallowed-methods). -#![allow(clippy::disallowed_methods)] - -use dialect_mir::ops as mir; -use dialect_nvvm::ops as nvvm; -use llvm_export::ops as llvm; -use pliron::builtin::op_interfaces::{ - CallOpCallable, CallOpInterface, OneRegionInterface, SymbolOpInterface, -}; -use pliron::builtin::ops::ModuleOp; -use pliron::context::Context; -use pliron::linked_list::ContainsLinkedList; -use pliron::location::{Located, Location}; -use pliron::op::Op; -use pliron::operation::Operation; -use pliron::r#type::Typed; - -#[test] -fn test_standalone_lowering_rejects_builtin_pointer_constant() { - use dialect_mir::types::{MirPointerKind, MirPtrType}; - use pliron::builtin::{ - attributes::IntegerAttr, - ops::ConstantOp, - types::{IntegerType, Signedness}, - }; - use pliron::utils::apint::APInt; - use std::num::NonZeroUsize; - - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - mir_lower::register(&mut ctx); - - let module = ModuleOp::new(&mut ctx, "pointer_constant".try_into().unwrap()); - let block = module - .get_region(&ctx) - .deref(&ctx) - .iter(&ctx) - .next() - .unwrap(); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let pointer_ty = - MirPtrType::get_generic_with_kind(&mut ctx, u32_ty.into(), true, MirPointerKind::UniqueRef); - let value = APInt::from_u64(0, NonZeroUsize::new(32).unwrap()); - let constant = ConstantOp::new(&mut ctx, IntegerAttr::new(u32_ty, value).into()); - let result = constant.get_operation().deref(&ctx).get_result(0); - result.set_type(&ctx, pointer_ty.into()); - constant.get_operation().insert_at_back(block, &ctx); - - let error = mir_lower::lower_mir_to_llvm(&mut ctx, module.get_operation()).unwrap_err(); - assert!( - error - .to_string() - .contains("builtin.constant cannot produce a MIR pointer carrier") - ); -} - -#[test] -fn test_intrinsic_insertion() -> Result<(), anyhow::Error> { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - mir_lower::register(&mut ctx); - - // Create Module - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_ptr = module.get_operation(); - - // Create MirFunc - let func_name = "kernel_func"; - let func_ty = pliron::builtin::types::FunctionType::get(&ctx, vec![], vec![]); - - // Manual construction of MirFuncOp - let func_op_ptr = Operation::new( - &mut ctx, - mir::MirFuncOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 1, // 1 region - ); - let func_ty_attr = pliron::builtin::attributes::TypeAttr::new(func_ty.into()); - let func = mir::MirFuncOp::new(&mut ctx, func_op_ptr, func_ty_attr); - func.set_symbol_name(&mut ctx, func_name.try_into().unwrap()); - - // Add body - MirFuncOp has 1 region - let region = func.get_operation().deref(&ctx).get_region(0); - - // Create block if empty (it is empty by default from Operation::new) - let block = { - let b = pliron::basic_block::BasicBlock::new(&mut ctx, None, vec![]); - b.insert_at_back(region, &ctx); - b - }; - - // Add ReadPtxSregTidXOp - let int32_ty = pliron::builtin::types::IntegerType::get( - &ctx, - 32, - pliron::builtin::types::Signedness::Signless, - ); - - let tid_op_ptr = Operation::new( - &mut ctx, - nvvm::ReadPtxSregTidXOp::get_concrete_op_info(), - vec![int32_ty.into()], - vec![], - vec![], - 0, - ); - let tid_op = nvvm::ReadPtxSregTidXOp::new(tid_op_ptr); - let expected_location = Location::Named { - name: "source-tid-x".to_string(), - child_loc: Box::new(Location::Unknown), - }; - tid_op - .get_operation() - .deref_mut(&ctx) - .set_loc(expected_location.clone()); - tid_op.get_operation().insert_at_back(block, &ctx); - - // Add Return - let ret_op_ptr = Operation::new( - &mut ctx, - mir::MirReturnOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 0, - ); - let ret_op = mir::MirReturnOp::new(ret_op_ptr); - ret_op.get_operation().insert_at_back(block, &ctx); - - // Add Func to Module - let module_region = module.get_operation().deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - func.get_operation().insert_at_back(module_block, &ctx); - - // Run DialectConversion-based lowering - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{}", e))?; - - // Verify result - let mut found_intrinsic = false; - let mut found_intrinsic_call = false; - let mut found_kernel = false; - - let module_op = module_ptr.deref(&ctx); - let region = module_op.get_region(0); - let block = region.deref(&ctx).iter(&ctx).next().unwrap(); - - for op in block.deref(&ctx).iter(&ctx) { - if let Some(func_op) = Operation::get_op::(op, &ctx) { - let name = func_op.get_symbol_name(&ctx).to_string(); - if name == "llvm_nvvm_read_ptx_sreg_tid_x" { - found_intrinsic = true; - // Intrinsic (declaration) should have 0 regions or empty region - let num_regions = func_op.get_operation().deref(&ctx).regions().count(); - if num_regions > 0 { - assert!( - func_op - .get_operation() - .deref(&ctx) - .get_region(0) - .deref(&ctx) - .iter(&ctx) - .next() - .is_none() - ); - } - } else if name == "kernel_func" { - found_kernel = true; - // Kernel should have body (1 region, not empty) - assert!(func_op.get_operation().deref(&ctx).regions().count() > 0); - assert!( - func_op - .get_operation() - .deref(&ctx) - .get_region(0) - .deref(&ctx) - .iter(&ctx) - .next() - .is_some() - ); - let kernel_region = func_op.get_operation().deref(&ctx).get_region(0); - for kernel_block in kernel_region.deref(&ctx).iter(&ctx) { - for body_op in kernel_block.deref(&ctx).iter(&ctx) { - let Some(call) = Operation::get_op::(body_op, &ctx) else { - continue; - }; - if matches!( - call.callee(&ctx), - CallOpCallable::Direct(symbol) - if symbol.to_string() == "llvm_nvvm_read_ptx_sreg_tid_x" - ) { - found_intrinsic_call = true; - assert_eq!(call.get_operation().deref(&ctx).loc(), expected_location); - } - } - } - } - } - } - - assert!(found_intrinsic, "Intrinsic function declaration not found"); - assert!(found_intrinsic_call, "Intrinsic call not found in kernel"); - assert!(found_kernel, "Kernel function not found"); - - Ok(()) -} - -#[test] -fn test_assertfail_orphaned_successor_block_is_removed() -> Result<(), anyhow::Error> { - use dialect_mir::types::MirPtrType; - - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - mir_lower::register(&mut ctx); - - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_ptr = module.get_operation(); - - let func_name = "kernel_func"; - let u8_ty = pliron::builtin::types::IntegerType::get( - &ctx, - 8, - pliron::builtin::types::Signedness::Unsigned, - ); - let ptr_ty = MirPtrType::get_generic(&mut ctx, u8_ty.into(), false); - let u32_ty = pliron::builtin::types::IntegerType::get( - &ctx, - 32, - pliron::builtin::types::Signedness::Unsigned, - ); - let u64_ty = pliron::builtin::types::IntegerType::get( - &ctx, - 64, - pliron::builtin::types::Signedness::Unsigned, - ); - let arg_tys: Vec = vec![ - ptr_ty.into(), - ptr_ty.into(), - u32_ty.into(), - ptr_ty.into(), - u64_ty.into(), - ]; - let func_ty = pliron::builtin::types::FunctionType::get(&ctx, arg_tys.clone(), vec![]); - - let func_op_ptr = Operation::new( - &mut ctx, - mir::MirFuncOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 1, - ); - let func_ty_attr = pliron::builtin::attributes::TypeAttr::new(func_ty.into()); - let func = mir::MirFuncOp::new(&mut ctx, func_op_ptr, func_ty_attr); - func.set_symbol_name(&mut ctx, func_name.try_into().unwrap()); - - let region = func.get_operation().deref(&ctx).get_region(0); - let entry = { - let b = pliron::basic_block::BasicBlock::new(&mut ctx, None, arg_tys); - b.insert_at_back(region, &ctx); - b - }; - // A block reachable only through the assert's success edge, carrying a - // block argument: the CFG a constant-false `gpu_assert!` leaves behind - // after MIR optimization folds the direct edge away. - let tail = { - let b = pliron::basic_block::BasicBlock::new(&mut ctx, None, vec![u32_ty.into()]); - b.insert_at_back(region, &ctx); - b - }; - - let message = entry.deref(&ctx).get_argument(0); - let file = entry.deref(&ctx).get_argument(1); - let line = entry.deref(&ctx).get_argument(2); - let function = entry.deref(&ctx).get_argument(3); - let char_size = entry.deref(&ctx).get_argument(4); - - let assertfail_op = - nvvm::AssertFailOp::build(&mut ctx, message, file, line, function, char_size); - assertfail_op.insert_at_back(entry, &ctx); - - // The success edge forwards `line` into the tail block's argument. - // Lowering erases this branch (the call never returns), orphaning the - // tail block, which then cannot be exported: a block argument needs one - // PHI incoming value per predecessor and there are none left. - let goto_op = Operation::new( - &mut ctx, - mir::MirGotoOp::get_concrete_op_info(), - vec![], - vec![line], - vec![tail], - 0, - ); - goto_op.insert_at_back(entry, &ctx); - - let ret_op = Operation::new( - &mut ctx, - mir::MirReturnOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 0, - ); - ret_op.insert_at_back(tail, &ctx); - - let module_region = module.get_operation().deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - func.get_operation().insert_at_back(module_block, &ctx); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let mut found_kernel = false; - let module_op = module_ptr.deref(&ctx); - let region = module_op.get_region(0); - let block = region.deref(&ctx).iter(&ctx).next().unwrap(); - for op in block.deref(&ctx).iter(&ctx) { - let Some(func_op) = Operation::get_op::(op, &ctx) else { - continue; - }; - if func_op.get_symbol_name(&ctx).to_string() != func_name { - continue; - } - found_kernel = true; - let func_region = func_op.get_operation().deref(&ctx).get_region(0); - let blocks: Vec<_> = func_region.deref(&ctx).iter(&ctx).collect(); - for block in blocks.iter().skip(1) { - assert!( - !block.preds(&ctx).is_empty(), - "an unreachable block survived lowering; with block arguments \ - it cannot be exported (a PHI needs an incoming value per \ - predecessor and there are none)" - ); - } - } - assert!(found_kernel, "Kernel function not found"); - - let module = Operation::get_op::(module_ptr, &ctx).unwrap(); - let ir = - llvm_export::export::export_module_to_string(&ctx, &module).map_err(anyhow::Error::msg)?; - let lines: Vec<_> = ir.lines().collect(); - let call_line = lines - .iter() - .position(|line| line.contains("call void @__assertfail(")) - .expect("expected exported __assertfail call"); - assert_eq!(lines[call_line + 1].trim(), "unreachable", "{ir}"); - - Ok(()) -} - -#[test] -fn test_globaltimer_lowers_to_intrinsic_call() -> Result<(), anyhow::Error> { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - mir_lower::register(&mut ctx); - - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_ptr = module.get_operation(); - - let func_name = "kernel_func"; - let func_ty = pliron::builtin::types::FunctionType::get(&ctx, vec![], vec![]); - - let func_op_ptr = Operation::new( - &mut ctx, - mir::MirFuncOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 1, - ); - let func_ty_attr = pliron::builtin::attributes::TypeAttr::new(func_ty.into()); - let func = mir::MirFuncOp::new(&mut ctx, func_op_ptr, func_ty_attr); - func.set_symbol_name(&mut ctx, func_name.try_into().unwrap()); - - let region = func.get_operation().deref(&ctx).get_region(0); - let block = { - let b = pliron::basic_block::BasicBlock::new(&mut ctx, None, vec![]); - b.insert_at_back(region, &ctx); - b - }; - - let i64_ty = pliron::builtin::types::IntegerType::get( - &ctx, - 64, - pliron::builtin::types::Signedness::Signless, - ); - let timer_op = Operation::new( - &mut ctx, - nvvm::ReadPtxSregGlobaltimerOp::get_concrete_op_info(), - vec![i64_ty.into()], - vec![], - vec![], - 0, - ); - timer_op.insert_at_back(block, &ctx); - - let ret_op_ptr = Operation::new( - &mut ctx, - mir::MirReturnOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 0, - ); - let ret_op = mir::MirReturnOp::new(ret_op_ptr); - ret_op.get_operation().insert_at_back(block, &ctx); - - let module_region = module.get_operation().deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - func.get_operation().insert_at_back(module_block, &ctx); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{}", e))?; - - const INTRINSIC: &str = "llvm_nvvm_read_ptx_sreg_globaltimer"; - - let mut found_decl = false; - let mut found_call = false; - let module_op = module_ptr.deref(&ctx); - let region = module_op.get_region(0); - let block = region.deref(&ctx).iter(&ctx).next().unwrap(); - - for op in block.deref(&ctx).iter(&ctx) { - let Some(func_op) = Operation::get_op::(op, &ctx) else { - continue; - }; - let name = func_op.get_symbol_name(&ctx).to_string(); - - if name == INTRINSIC { - // Intrinsic declaration: present with empty body. - found_decl = true; - let num_regions = func_op.get_operation().deref(&ctx).regions().count(); - if num_regions > 0 { - assert!( - func_op - .get_operation() - .deref(&ctx) - .get_region(0) - .deref(&ctx) - .iter(&ctx) - .next() - .is_none(), - "intrinsic declaration must have empty body" - ); - } - } else if name == func_name { - let func_region = func_op.get_operation().deref(&ctx).get_region(0); - for func_block in func_region.deref(&ctx).iter(&ctx) { - for body_op in func_block.deref(&ctx).iter(&ctx) { - if let Some(call) = Operation::get_op::(body_op, &ctx) - && let CallOpCallable::Direct(sym) = call.callee(&ctx) - && sym.to_string() == INTRINSIC - { - found_call = true; - } - assert!( - Operation::get_op::(body_op, &ctx).is_none(), - "globaltimer must not lower to inline asm" - ); - } - } - } - } - - assert!( - found_decl, - "Expected `{INTRINSIC}` declaration in lowered module" - ); - assert!( - found_call, - "Expected call to `{INTRINSIC}` in lowered kernel body" - ); - Ok(()) -} - -#[test] -fn test_assertfail_lowers_to_noreturn_call_and_unreachable() -> Result<(), anyhow::Error> { - use dialect_mir::types::MirPtrType; - - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - mir_lower::register(&mut ctx); - - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_ptr = module.get_operation(); - - let func_name = "kernel_func"; - let u8_ty = pliron::builtin::types::IntegerType::get( - &ctx, - 8, - pliron::builtin::types::Signedness::Unsigned, - ); - let ptr_ty = MirPtrType::get_generic(&mut ctx, u8_ty.into(), false); - let u32_ty = pliron::builtin::types::IntegerType::get( - &ctx, - 32, - pliron::builtin::types::Signedness::Unsigned, - ); - let u64_ty = pliron::builtin::types::IntegerType::get( - &ctx, - 64, - pliron::builtin::types::Signedness::Unsigned, - ); - let arg_tys: Vec = vec![ - ptr_ty.into(), - ptr_ty.into(), - u32_ty.into(), - ptr_ty.into(), - u64_ty.into(), - ]; - let func_ty = pliron::builtin::types::FunctionType::get(&ctx, arg_tys.clone(), vec![]); - - let func_op_ptr = Operation::new( - &mut ctx, - mir::MirFuncOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 1, - ); - let func_ty_attr = pliron::builtin::attributes::TypeAttr::new(func_ty.into()); - let func = mir::MirFuncOp::new(&mut ctx, func_op_ptr, func_ty_attr); - func.set_symbol_name(&mut ctx, func_name.try_into().unwrap()); - - let region = func.get_operation().deref(&ctx).get_region(0); - let block = { - let b = pliron::basic_block::BasicBlock::new(&mut ctx, None, arg_tys); - b.insert_at_back(region, &ctx); - b - }; - - let message = block.deref(&ctx).get_argument(0); - let file = block.deref(&ctx).get_argument(1); - let line = block.deref(&ctx).get_argument(2); - let function = block.deref(&ctx).get_argument(3); - let char_size = block.deref(&ctx).get_argument(4); - - let assertfail_op = - nvvm::AssertFailOp::build(&mut ctx, message, file, line, function, char_size); - assertfail_op.insert_at_back(block, &ctx); - - // This return must be removed by lowering because __assertfail never returns. - let ret_op_ptr = Operation::new( - &mut ctx, - mir::MirReturnOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 0, - ); - ret_op_ptr.insert_at_back(block, &ctx); - - let module_region = module.get_operation().deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - func.get_operation().insert_at_back(module_block, &ctx); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - const EXTERN: &str = "__assertfail"; - - let mut found_decl = false; - let mut found_call = false; - let mut found_unreachable = false; - - let module_op = module_ptr.deref(&ctx); - let region = module_op.get_region(0); - let block = region.deref(&ctx).iter(&ctx).next().unwrap(); - - for op in block.deref(&ctx).iter(&ctx) { - let Some(func_op) = Operation::get_op::(op, &ctx) else { - continue; - }; - let name = func_op.get_symbol_name(&ctx).to_string(); - - if name == EXTERN { - found_decl = true; - assert!( - llvm::op_noreturn(&ctx, func_op.get_operation()), - "__assertfail declaration must be marked noreturn" - ); - continue; - } - - if name != func_name { - continue; - } - - let func_region = func_op.get_operation().deref(&ctx).get_region(0); - for func_block in func_region.deref(&ctx).iter(&ctx) { - let body_ops: Vec<_> = func_block.deref(&ctx).iter(&ctx).collect(); - - for (index, body_op) in body_ops.iter().copied().enumerate() { - if let Some(call) = Operation::get_op::(body_op, &ctx) - && let CallOpCallable::Direct(sym) = call.callee(&ctx) - && sym.to_string() == EXTERN - { - found_call = true; - - assert!( - llvm::op_noreturn(&ctx, body_op), - "__assertfail call must be marked noreturn" - ); - assert_eq!( - body_op.deref(&ctx).get_num_operands(), - 5, - "__assertfail call must forward all five operands" - ); - assert!( - body_ops.get(index + 1).is_some_and(|next| { - Operation::get_op::(*next, &ctx).is_some() - }), - "llvm.unreachable must immediately follow __assertfail" - ); - assert_eq!( - index + 2, - body_ops.len(), - "no operations may remain after llvm.unreachable" - ); - } - - if Operation::get_op::(body_op, &ctx).is_some() { - found_unreachable = true; - } - - assert!( - Operation::get_op::(body_op, &ctx).is_none(), - "nvvm.assertfail must be fully consumed by the lowering" - ); - } - } - } - - assert!( - found_decl, - "Expected `{EXTERN}` declaration in lowered module" - ); - assert!( - found_call, - "Expected call to `{EXTERN}` in lowered kernel body" - ); - assert!( - found_unreachable, - "__assertfail must terminate the block with llvm.unreachable" - ); - - let module = Operation::get_op::(module_ptr, &ctx).unwrap(); - let ir = - llvm_export::export::export_module_to_string(&ctx, &module).map_err(anyhow::Error::msg)?; - - let declaration = ir - .lines() - .find(|line| line.contains("declare void @__assertfail(")) - .expect("expected exported __assertfail declaration"); - assert!(declaration.contains("noreturn"), "{ir}"); - - let lines: Vec<_> = ir.lines().collect(); - let call_line = lines - .iter() - .position(|line| line.contains("call void @__assertfail(")) - .expect("expected exported __assertfail call"); - assert!(lines[call_line].contains("noreturn"), "{ir}"); - assert_eq!(lines[call_line + 1].trim(), "unreachable", "{ir}"); - - Ok(()) -} - -/// Lower a single zero-operand, i32-result special-register op and assert it -/// emits a declaration of and direct call to `intrinsic` (and no inline asm). -fn assert_sreg_i32_lowers_to_intrinsic( - op_info: ( - fn(pliron::context::Ptr) -> pliron::op::OpObj, - std::any::TypeId, - ), - intrinsic: &str, -) -> Result<(), anyhow::Error> { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - mir_lower::register(&mut ctx); - - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_ptr = module.get_operation(); - - let func_name = "kernel_func"; - let func_ty = pliron::builtin::types::FunctionType::get(&ctx, vec![], vec![]); - - let func_op_ptr = Operation::new( - &mut ctx, - mir::MirFuncOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 1, - ); - let func_ty_attr = pliron::builtin::attributes::TypeAttr::new(func_ty.into()); - let func = mir::MirFuncOp::new(&mut ctx, func_op_ptr, func_ty_attr); - func.set_symbol_name(&mut ctx, func_name.try_into().unwrap()); - - let region = func.get_operation().deref(&ctx).get_region(0); - let block = { - let b = pliron::basic_block::BasicBlock::new(&mut ctx, None, vec![]); - b.insert_at_back(region, &ctx); - b - }; - - let i32_ty = pliron::builtin::types::IntegerType::get( - &ctx, - 32, - pliron::builtin::types::Signedness::Signless, - ); - let sreg_op = Operation::new(&mut ctx, op_info, vec![i32_ty.into()], vec![], vec![], 0); - sreg_op.insert_at_back(block, &ctx); - - let ret_op_ptr = Operation::new( - &mut ctx, - mir::MirReturnOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 0, - ); - let ret_op = mir::MirReturnOp::new(ret_op_ptr); - ret_op.get_operation().insert_at_back(block, &ctx); - - let module_region = module.get_operation().deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - func.get_operation().insert_at_back(module_block, &ctx); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{}", e))?; - - let mut found_decl = false; - let mut found_call = false; - let module_op = module_ptr.deref(&ctx); - let region = module_op.get_region(0); - let block = region.deref(&ctx).iter(&ctx).next().unwrap(); - - for op in block.deref(&ctx).iter(&ctx) { - let Some(func_op) = Operation::get_op::(op, &ctx) else { - continue; - }; - let name = func_op.get_symbol_name(&ctx).to_string(); - - if name == intrinsic { - found_decl = true; - } else if name == func_name { - let func_region = func_op.get_operation().deref(&ctx).get_region(0); - for func_block in func_region.deref(&ctx).iter(&ctx) { - for body_op in func_block.deref(&ctx).iter(&ctx) { - if let Some(call) = Operation::get_op::(body_op, &ctx) - && let CallOpCallable::Direct(sym) = call.callee(&ctx) - && sym.to_string() == intrinsic - { - found_call = true; - } - assert!( - Operation::get_op::(body_op, &ctx).is_none(), - "{intrinsic} must not lower to inline asm" - ); - } - } - } - } - - assert!( - found_decl, - "Expected `{intrinsic}` declaration in lowered module" - ); - assert!( - found_call, - "Expected call to `{intrinsic}` in lowered kernel body" - ); - Ok(()) -} - -fn assert_sreg_lowers_to_inline_asm( - op_info: ( - fn(pliron::context::Ptr) -> pliron::op::OpObj, - std::any::TypeId, - ), - result_width: u32, - expected_template: &str, - expected_constraints: &str, - expected_kind: llvm::AsmKind, -) -> Result<(), anyhow::Error> { - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let result_ty = IntegerType::get(&ctx, result_width, Signedness::Signless); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![]); - let sreg_op = Operation::new(&mut ctx, op_info, vec![result_ty.into()], vec![], vec![], 0); - sreg_op.insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let mut matches = 0usize; - let module_region = module_ptr.deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - for op in module_block.deref(&ctx).iter(&ctx) { - let Some(func_op) = Operation::get_op::(op, &ctx) else { - continue; - }; - if func_op.get_symbol_name(&ctx).to_string() != "kernel_func" { - continue; - } - - let func_region = func_op.get_operation().deref(&ctx).get_region(0); - for func_block in func_region.deref(&ctx).iter(&ctx) { - for body_op in func_block.deref(&ctx).iter(&ctx) { - let Some(inline_asm) = Operation::get_op::(body_op, &ctx) else { - continue; - }; - let template = inline_asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())); - if template.as_deref() != Some(expected_template) { - continue; - } - - matches += 1; - assert_eq!( - inline_asm - .get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some(expected_constraints) - ); - assert_eq!(llvm::asm_kind(&ctx, &inline_asm), expected_kind); - } - } - } - - assert_eq!(matches, 1, "expected one exact `{expected_template}` read"); - Ok(()) -} - -#[test] -fn test_lanemask_ops_lower_to_sreg_intrinsic_calls() -> Result<(), anyhow::Error> { - // Each lane-position mask op lowers to its matching read-only sreg intrinsic - // (underscores become dots on export: `..._lanemask_lt` -> `...lanemask.lt`). - assert_sreg_i32_lowers_to_intrinsic( - nvvm::ReadPtxSregLanemaskLtOp::get_concrete_op_info(), - "llvm_nvvm_read_ptx_sreg_lanemask_lt", - )?; - assert_sreg_i32_lowers_to_intrinsic( - nvvm::ReadPtxSregLanemaskLeOp::get_concrete_op_info(), - "llvm_nvvm_read_ptx_sreg_lanemask_le", - )?; - assert_sreg_i32_lowers_to_intrinsic( - nvvm::ReadPtxSregLanemaskEqOp::get_concrete_op_info(), - "llvm_nvvm_read_ptx_sreg_lanemask_eq", - )?; - assert_sreg_i32_lowers_to_intrinsic( - nvvm::ReadPtxSregLanemaskGeOp::get_concrete_op_info(), - "llvm_nvvm_read_ptx_sreg_lanemask_ge", - )?; - assert_sreg_i32_lowers_to_intrinsic( - nvvm::ReadPtxSregLanemaskGtOp::get_concrete_op_info(), - "llvm_nvvm_read_ptx_sreg_lanemask_gt", - )?; - Ok(()) -} - -#[test] -fn test_generated_vote_sync_family_lowers_to_exact_typed_intrinsics() -> Result<(), anyhow::Error> { - use pliron::builtin::types::{IntegerType, Signedness}; - use pliron::r#type::Typed; - - let mut ctx = make_test_ctx(); - let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![i32_ty.into(), i1_ty.into()]); - let mask = entry.deref(&ctx).get_argument(0); - let predicate = entry.deref(&ctx).get_argument(1); - - for vote in [ - nvvm::VoteSyncAllOp::build(&mut ctx, mask, predicate), - nvvm::VoteSyncAnyOp::build(&mut ctx, mask, predicate), - nvvm::VoteSyncBallotOp::build(&mut ctx, mask, predicate), - nvvm::VoteSyncUniOp::build(&mut ctx, mask, predicate), - ] { - vote.insert_at_back(entry, &ctx); - } - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let expected = [ - ("llvm_nvvm_vote_all_sync", 1), - ("llvm_nvvm_vote_any_sync", 1), - ("llvm_nvvm_vote_ballot_sync", 32), - ("llvm_nvvm_vote_uni_sync", 1), - ]; - let mut found = Vec::new(); - for op in lowered_kernel_body(&ctx, module_ptr) { - assert!( - Operation::get_op::(op, &ctx).is_none(), - "generated vote.sync operations must use typed LLVM intrinsics" - ); - let Some(call) = Operation::get_op::(op, &ctx) else { - continue; - }; - let CallOpCallable::Direct(callee) = call.callee(&ctx) else { - continue; - }; - let callee = callee.to_string(); - let Some((_, result_width)) = expected.iter().find(|(name, _)| *name == callee) else { - continue; - }; - - let call = call.get_operation().deref(&ctx); - assert_eq!(call.get_num_operands(), 2); - assert_eq!(call.get_num_results(), 1); - - let integer_shape = |value: pliron::value::Value| { - let ty = value.get_type(&ctx); - let ty = ty.deref(&ctx); - let integer = ty - .downcast_ref::() - .expect("vote.sync operands and results are integers"); - (integer.width(), integer.signedness()) - }; - assert_eq!( - [ - integer_shape(call.get_operand(0)), - integer_shape(call.get_operand(1)), - ], - [(32, Signedness::Signless), (1, Signedness::Signless),], - "{callee} must preserve [mask, predicate] operand order" - ); - assert_eq!( - integer_shape(call.get_result(0)), - (*result_width, Signedness::Signless), - "{callee} returned the wrong LLVM integer type" - ); - found.push((callee, *result_width)); - } - - found.sort(); - let mut expected = expected - .map(|(name, width)| (name.to_owned(), width)) - .to_vec(); - expected.sort(); - assert_eq!(found, expected); - Ok(()) -} - -fn lower_generated_active_mask( - backend: mir_lower::IntrinsicBackend, -) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { - let mut ctx = make_test_ctx(); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![]); - nvvm::ActiveMaskOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm_with_options( - &mut ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: backend, - ..Default::default() - }, - ) - .map_err(|error| anyhow::anyhow!("{error}"))?; - Ok((ctx, module_ptr)) -} - -#[test] -fn test_generated_active_mask_llvm_nvptx_uses_typed_intrinsic() -> Result<(), anyhow::Error> { - use pliron::builtin::types::IntegerType; - use pliron::r#type::Typed; - - let (ctx, module_ptr) = lower_generated_active_mask(mir_lower::IntrinsicBackend::LlvmNvptx)?; - let mut found = 0; - for op in lowered_kernel_body(&ctx, module_ptr) { - assert!( - Operation::get_op::(op, &ctx).is_none(), - "LLVM-NVPTX active_mask must use the typed intrinsic" - ); - let Some(call) = Operation::get_op::(op, &ctx) else { - continue; - }; - let CallOpCallable::Direct(callee) = call.callee(&ctx) else { - continue; - }; - if callee.to_string() != "llvm_nvvm_activemask" { - continue; - } - - found += 1; - let call = op.deref(&ctx); - assert_eq!(call.get_num_operands(), 0); - assert_eq!(call.get_num_results(), 1); - let result_ty = call.get_result(0).get_type(&ctx); - let result_ty = result_ty.deref(&ctx); - let result_ty = result_ty - .downcast_ref::() - .expect("active_mask returns an integer"); - assert_eq!(result_ty.width(), 32); - } - assert_eq!(found, 1, "expected one typed active_mask call"); - - let module = Operation::get_op::(module_ptr, &ctx).unwrap(); - let ir = llvm_export::export::export_module_to_string(&ctx, &module) - .map_err(|error| anyhow::anyhow!(error))?; - assert!(ir.contains("call i32 @llvm.nvvm.activemask()"), "{ir}"); - Ok(()) -} - -#[test] -fn test_generated_active_mask_libnvvm_uses_convergent_sideeffect_asm() -> Result<(), anyhow::Error> -{ - use pliron::builtin::types::IntegerType; - use pliron::r#type::Typed; - - let (ctx, module_ptr) = lower_generated_active_mask(mir_lower::IntrinsicBackend::LibNvvm)?; - let mut found = 0; - for op in lowered_kernel_body(&ctx, module_ptr) { - if let Some(call) = Operation::get_op::(op, &ctx) - && let CallOpCallable::Direct(callee) = call.callee(&ctx) - { - assert_ne!(callee.to_string(), "llvm_nvvm_activemask"); - } - let Some(asm) = Operation::get_op::(op, &ctx) else { - continue; - }; - - found += 1; - assert_eq!( - asm.get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some("activemask.b32 $0;") - ); - assert_eq!( - asm.get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some("=r,~{memory}") - ); - assert_eq!(llvm::asm_kind(&ctx, &asm), llvm::AsmKind::Convergent); - assert!( - asm.get_attr_inline_asm_convergent(&ctx) - .is_some_and(|value| bool::from((*value).clone())) - ); - let asm = op.deref(&ctx); - assert_eq!(asm.get_num_operands(), 0); - assert_eq!(asm.get_num_results(), 1); - let result_ty = asm.get_result(0).get_type(&ctx); - let result_ty = result_ty.deref(&ctx); - let result_ty = result_ty - .downcast_ref::() - .expect("active_mask inline asm returns an integer"); - assert_eq!(result_ty.width(), 32); - } - assert_eq!(found, 1, "expected one exact active_mask asm block"); - - let module = Operation::get_op::(module_ptr, &ctx).unwrap(); - let ir = llvm_export::export::export_module_to_string(&ctx, &module) - .map_err(|error| anyhow::anyhow!(error))?; - assert!( - ir.contains("call i32 asm sideeffect \"activemask.b32 $0;\", \"=r,~{memory}\"()"), - "{ir}" - ); - assert!(ir.contains("attributes #0 = { convergent }"), "{ir}"); - Ok(()) -} - -#[test] -fn test_generated_warp_match_family_uses_exact_typed_calls_and_mask_projection() --> Result<(), anyhow::Error> { - use llvm_export::types::StructType; - use pliron::builtin::types::{IntegerType, Signedness}; - use pliron::r#type::Typed; - - let mut ctx = make_test_ctx(); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); - let (module_ptr, entry) = - build_test_kernel(&mut ctx, vec![i32_ty.into(), i32_ty.into(), i64_ty.into()]); - let mask = entry.deref(&ctx).get_argument(0); - let value32 = entry.deref(&ctx).get_argument(1); - let value64 = entry.deref(&ctx).get_argument(2); - for warp_match in [ - nvvm::MatchAnySyncI32Op::build(&mut ctx, mask, value32), - nvvm::MatchAnySyncI64Op::build(&mut ctx, mask, value64), - nvvm::MatchAllSyncI32Op::build(&mut ctx, mask, value32), - nvvm::MatchAllSyncI64Op::build(&mut ctx, mask, value64), - ] { - warp_match.insert_at_back(entry, &ctx); - } - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm_with_options( - &mut ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: mir_lower::IntrinsicBackend::LlvmNvptx, - ..Default::default() - }, - ) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let expected = [ - ("llvm_nvvm_match_any_sync_i32", 32, false), - ("llvm_nvvm_match_any_sync_i64", 64, false), - ("llvm_nvvm_match_all_sync_i32p", 32, true), - ("llvm_nvvm_match_all_sync_i64p", 64, true), - ]; - let body = lowered_kernel_body(&ctx, module_ptr); - let integer_width = |value: pliron::value::Value| { - let ty = value.get_type(&ctx); - let ty = ty.deref(&ctx); - ty.downcast_ref::() - .expect("warp-match value is an integer") - .width() - }; - let mut found = Vec::new(); - let mut aggregate_results = Vec::new(); - for &op in &body { - assert!( - Operation::get_op::(op, &ctx).is_none(), - "warp match must use typed LLVM intrinsics" - ); - let Some(call) = Operation::get_op::(op, &ctx) else { - continue; - }; - let CallOpCallable::Direct(callee) = call.callee(&ctx) else { - continue; - }; - let callee = callee.to_string(); - let Some((_, value_width, aggregate)) = - expected.iter().find(|(name, _, _)| *name == callee) - else { - continue; - }; - - let call = op.deref(&ctx); - assert_eq!(call.get_num_operands(), 2); - assert_eq!( - [ - integer_width(call.get_operand(0)), - integer_width(call.get_operand(1)), - ], - [32, *value_width], - "{callee} has the wrong typed signature" - ); - assert_eq!(call.get_num_results(), 1); - let result = call.get_result(0); - if *aggregate { - let result_ty = result.get_type(&ctx); - let result_ty = result_ty.deref(&ctx); - let result_ty = result_ty - .downcast_ref::() - .expect("match.all returns an LLVM aggregate"); - assert_eq!(result_ty.num_fields(), 2); - let field_widths = (0..2) - .map(|index| { - let field = result_ty.field_type(index); - let field = field.deref(&ctx); - field - .downcast_ref::() - .expect("match.all aggregate fields are integers") - .width() - }) - .collect::>(); - assert_eq!(field_widths, [32, 1]); - aggregate_results.push((callee.clone(), result)); - } else { - assert_eq!(integer_width(result), 32); - } - found.push(callee); - } - - found.sort(); - let mut expected_calls = expected.map(|(name, _, _)| name.to_owned()); - expected_calls.sort(); - assert_eq!(found, expected_calls); - - let mut projected = Vec::new(); - for &op in &body { - let Some(extract) = Operation::get_op::(op, &ctx) else { - continue; - }; - assert_eq!(extract.indices(&ctx), vec![0]); - let extract = op.deref(&ctx); - assert_eq!(extract.get_num_operands(), 1); - assert_eq!(extract.get_num_results(), 1); - assert_eq!(integer_width(extract.get_result(0)), 32); - let aggregate = extract.get_operand(0); - let callee = aggregate_results - .iter() - .find_map(|(callee, result)| (*result == aggregate).then(|| callee.clone())) - .expect("match.all must extract from its aggregate call result"); - projected.push(callee); - } - projected.sort(); - assert_eq!( - projected, - [ - "llvm_nvvm_match_all_sync_i32p".to_owned(), - "llvm_nvvm_match_all_sync_i64p".to_owned(), - ] - ); - Ok(()) -} - -#[test] -fn test_warpid_ops_preserve_snapshot_semantics() -> Result<(), anyhow::Error> { - assert_sreg_lowers_to_inline_asm( - nvvm::ReadPtxSregWarpIdOp::get_concrete_op_info(), - 32, - "mov.u32 $0, %warpid;", - "=r", - llvm::AsmKind::SideEffect, - )?; - assert_sreg_i32_lowers_to_intrinsic( - nvvm::ReadPtxSregNwarpIdOp::get_concrete_op_info(), - "llvm_nvvm_read_ptx_sreg_nwarpid", - )?; - Ok(()) -} - -#[test] -fn test_smid_ops_preserve_snapshot_semantics() -> Result<(), anyhow::Error> { - assert_sreg_lowers_to_inline_asm( - nvvm::ReadPtxSregSmIdOp::get_concrete_op_info(), - 32, - "mov.u32 $0, %smid;", - "=r", - llvm::AsmKind::SideEffect, - )?; - assert_sreg_i32_lowers_to_intrinsic( - nvvm::ReadPtxSregNsmIdOp::get_concrete_op_info(), - "llvm_nvvm_read_ptx_sreg_nsmid", - )?; - Ok(()) -} - -#[test] -fn test_repeated_location_samples_remain_side_effecting_reads() -> Result<(), anyhow::Error> { - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![]); - - for op_info in [ - nvvm::ReadPtxSregWarpIdOp::get_concrete_op_info(), - nvvm::ReadPtxSregSmIdOp::get_concrete_op_info(), - nvvm::ReadPtxSregWarpIdOp::get_concrete_op_info(), - nvvm::ReadPtxSregSmIdOp::get_concrete_op_info(), - ] { - let op = Operation::new(&mut ctx, op_info, vec![i32_ty.into()], vec![], vec![], 0); - op.insert_at_back(entry, &ctx); - } - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let mut warpid_reads = 0usize; - let mut smid_reads = 0usize; - let module_region = module_ptr.deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - for op in module_block.deref(&ctx).iter(&ctx) { - let Some(func_op) = Operation::get_op::(op, &ctx) else { - continue; - }; - if func_op.get_symbol_name(&ctx).to_string() != "kernel_func" { - continue; - } - let func_region = func_op.get_operation().deref(&ctx).get_region(0); - for func_block in func_region.deref(&ctx).iter(&ctx) { - for body_op in func_block.deref(&ctx).iter(&ctx) { - let Some(inline_asm) = Operation::get_op::(body_op, &ctx) else { - continue; - }; - let template = inline_asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())); - match template.as_deref() { - Some("mov.u32 $0, %warpid;") => warpid_reads += 1, - Some("mov.u32 $0, %smid;") => smid_reads += 1, - _ => continue, - } - assert_eq!( - llvm::asm_kind(&ctx, &inline_asm), - llvm::AsmKind::SideEffect, - "location snapshots must survive LLVM CSE" - ); - } - } - } - - assert_eq!(warpid_reads, 2); - assert_eq!(smid_reads, 2); - Ok(()) -} - -#[test] -fn test_gridid_op_lowers_to_full_width_inline_ptx() -> Result<(), anyhow::Error> { - assert_sreg_lowers_to_inline_asm( - nvvm::ReadPtxSregGridIdOp::get_concrete_op_info(), - 64, - "mov.u64 $0, %gridid;", - "=l", - llvm::AsmKind::Pure, - ) -} - -#[test] -fn test_smem_size_ops_lower_to_portable_inline_ptx() -> Result<(), anyhow::Error> { - assert_sreg_lowers_to_inline_asm( - nvvm::ReadPtxSregDynamicSmemSizeOp::get_concrete_op_info(), - 32, - "mov.u32 $0, %dynamic_smem_size;", - "=r", - llvm::AsmKind::Pure, - )?; - assert_sreg_lowers_to_inline_asm( - nvvm::ReadPtxSregTotalSmemSizeOp::get_concrete_op_info(), - 32, - "mov.u32 $0, %total_smem_size;", - "=r", - llvm::AsmKind::Pure, - )?; - Ok(()) -} - -#[test] -fn generated_threadfences_use_typed_intrinsics_on_both_backends() -> Result<(), anyhow::Error> { - const EXPECTED: [&str; 3] = [ - "llvm_nvvm_membar_cta", - "llvm_nvvm_membar_gl", - "llvm_nvvm_membar_sys", - ]; - - for backend in [ - mir_lower::IntrinsicBackend::LlvmNvptx, - mir_lower::IntrinsicBackend::LibNvvm, - ] { - let mut ctx = make_test_ctx(); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![]); - for op_info in [ - nvvm::ThreadfenceBlockOp::get_concrete_op_info(), - nvvm::ThreadfenceOp::get_concrete_op_info(), - nvvm::ThreadfenceSystemOp::get_concrete_op_info(), - ] { - Operation::new(&mut ctx, op_info, vec![], vec![], vec![], 0) - .insert_at_back(entry, &ctx); - } - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm_with_options( - &mut ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: backend, - ..Default::default() - }, - ) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let mut found = Vec::new(); - for op in lowered_kernel_body(&ctx, module_ptr) { - assert!( - Operation::get_op::(op, &ctx).is_none(), - "threadfences must use their reviewed typed route" - ); - let Some(call) = Operation::get_op::(op, &ctx) else { - continue; - }; - let CallOpCallable::Direct(callee) = call.callee(&ctx) else { - continue; - }; - let callee = callee.to_string(); - if EXPECTED.contains(&callee.as_str()) { - assert_eq!(op.deref(&ctx).get_num_operands(), 0); - found.push(callee); - } - } - found.sort(); - assert_eq!(found, EXPECTED); - - let module = Operation::get_op::(module_ptr, &ctx).unwrap(); - let ir = llvm_export::export::export_module_to_string(&ctx, &module) - .map_err(|error| anyhow::anyhow!(error))?; - for symbol in [ - "llvm.nvvm.membar.cta", - "llvm.nvvm.membar.gl", - "llvm.nvvm.membar.sys", - ] { - assert!(ir.contains(&format!("call void @{symbol}()")), "{ir}"); - } - } - Ok(()) -} - -/// LLVM uses its typed intrinsic. libNVVM uses the reviewed inline-PTX fallback. -#[test] -fn generated_elect_sync_uses_the_selected_backend_route() -> Result<(), anyhow::Error> { - use pliron::builtin::types::{IntegerType, Signedness}; - - for backend in [ - mir_lower::IntrinsicBackend::LlvmNvptx, - mir_lower::IntrinsicBackend::LibNvvm, - ] { - let mut ctx = make_test_ctx(); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![i32_ty.into()]); - let mask = entry.deref(&ctx).get_argument(0); - Operation::new( - &mut ctx, - nvvm::ElectSyncOp::get_concrete_op_info(), - vec![i32_ty.into(), i1_ty.into()], - vec![mask], - vec![], - 0, - ) - .insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm_with_options( - &mut ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: backend, - ..Default::default() - }, - ) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let mut inline_asm_count = 0; - let mut typed_call_count = 0; - let mut extract_count = 0; - let mut trunc_count = 0; - for body_op in lowered_kernel_body(&ctx, module_ptr) { - if let Some(inline_asm) = Operation::get_op::(body_op, &ctx) { - assert_eq!( - inline_asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some("{ .reg .pred p; elect.sync $0|p, $2; selp.b32 $1, 1, 0, p; }") - ); - assert_eq!( - inline_asm - .get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some("=r,=r,r") - ); - assert!( - inline_asm - .get_attr_inline_asm_convergent(&ctx) - .is_some_and(|value| bool::from((*value).clone())) - ); - inline_asm_count += 1; - } - if let Some(call) = Operation::get_op::(body_op, &ctx) - && let CallOpCallable::Direct(callee) = call.callee(&ctx) - && callee.to_string() == "llvm_nvvm_elect_sync" - { - typed_call_count += 1; - } - if Operation::get_op::(body_op, &ctx).is_some() { - extract_count += 1; - } - if Operation::get_op::(body_op, &ctx).is_some() { - trunc_count += 1; - } - } - - assert_eq!(extract_count, 2); - match backend { - mir_lower::IntrinsicBackend::LlvmNvptx => { - assert_eq!(typed_call_count, 1); - assert_eq!(inline_asm_count, 0); - assert_eq!(trunc_count, 0); - } - mir_lower::IntrinsicBackend::LibNvvm => { - assert_eq!(typed_call_count, 0); - assert_eq!(inline_asm_count, 1); - assert_eq!(trunc_count, 1); - } - } - } - Ok(()) -} - -/// The exact inline-PTX template `convert_shuffle_i64` must emit for `mode`/`clamp`. -/// Mirrors the production `format!` so a drift in either side fails the test. -fn expected_shfl_i64_template(mode: &str, clamp: i32) -> String { - format!( - "{{ .reg .b32 lo; .reg .b32 hi; mov.b64 {{lo, hi}}, $1; \ - shfl.sync.{mode}.b32 lo, lo, $2, {clamp}, $3; \ - shfl.sync.{mode}.b32 hi, hi, $2, {clamp}, $3; \ - mov.b64 $0, {{lo, hi}}; }}" - ) -} - -/// 64-bit warp shuffle has no LLVM intrinsic (`shfl.sync` is 32-bit only), so it -/// lowers to convergent inline PTX that splits the value into two halves and runs -/// two `shfl.sync.*.b32`. Inline asm is opaque to LLVM, so a wrong mnemonic, -/// swapped operand order, wrong clamp, or missing `convergent` would only surface -/// as bad PTX downstream. This pins, for every mode, the exact template (incl. the -/// per-mode clamp: 31 for idx/bfly/down, 0 for up), the `=l,l,r,r` constraints, -/// and the convergent flag. -#[test] -fn test_shuffle_i64_lowers_to_inline_asm() -> Result<(), anyhow::Error> { - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); - // Kernel args: [mask (i32), value (i64), lane/delta (i32)]. - let (module_ptr, entry) = - build_test_kernel(&mut ctx, vec![i32_ty.into(), i64_ty.into(), i32_ty.into()]); - let mask = entry.deref(&ctx).get_argument(0); - let value = entry.deref(&ctx).get_argument(1); - let lane = entry.deref(&ctx).get_argument(2); - - // One op per mode, all sharing the same [mask, value, lane] operands. - type OpInfo = ( - fn(pliron::context::Ptr) -> pliron::op::OpObj, - std::any::TypeId, - ); - let modes: [(OpInfo, &str, i32); 4] = [ - (nvvm::ShflSyncIdxI64Op::get_concrete_op_info(), "idx", 31), - (nvvm::ShflSyncBflyI64Op::get_concrete_op_info(), "bfly", 31), - (nvvm::ShflSyncDownI64Op::get_concrete_op_info(), "down", 31), - (nvvm::ShflSyncUpI64Op::get_concrete_op_info(), "up", 0), - ]; - for (opid, _, _) in modes { - let op = Operation::new( - &mut ctx, - opid, - vec![i64_ty.into()], - vec![mask, value, lane], - vec![], - 0, - ); - op.insert_at_back(entry, &ctx); - } - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{}", e))?; - - // Collect every inline-asm template emitted into the kernel body. - let mut templates: Vec = Vec::new(); - let module_op = module_ptr.deref(&ctx); - let region = module_op.get_region(0); - let block = region.deref(&ctx).iter(&ctx).next().unwrap(); - for op in block.deref(&ctx).iter(&ctx) { - let Some(func_op) = Operation::get_op::(op, &ctx) else { - continue; - }; - if func_op.get_symbol_name(&ctx).to_string() != "kernel_func" { - continue; - } - let func_region = func_op.get_operation().deref(&ctx).get_region(0); - for func_block in func_region.deref(&ctx).iter(&ctx) { - for body_op in func_block.deref(&ctx).iter(&ctx) { - let Some(inline_asm) = Operation::get_op::(body_op, &ctx) else { - continue; - }; - assert_eq!( - inline_asm - .get_attr_inline_asm_constraints(&ctx) - .map(|s| String::from((*s).clone())) - .as_deref(), - Some("=l,l,r,r"), - "shfl.b64 constraints must be [out i64, value i64, lane i32, mask i32]" - ); - assert!( - inline_asm - .get_attr_inline_asm_convergent(&ctx) - .is_some_and(|b| bool::from((*b).clone())), - "shfl.b64 inline asm must be convergent" - ); - templates.push( - inline_asm - .get_attr_inline_asm_template(&ctx) - .map(|s| String::from((*s).clone())) - .unwrap_or_default(), - ); - } - } - } - - assert_eq!( - templates.len(), - 4, - "each of the 4 shfl.b64 modes must lower to one inline-asm op" - ); - for (_, mode, clamp) in modes { - let want = expected_shfl_i64_template(mode, clamp); - assert!( - templates.contains(&want), - "missing inline PTX for shfl.sync.{mode}.b32 (clamp {clamp}); got {templates:?}" - ); - } - - Ok(()) -} - -/// Regression cover for explicit MIR address-space coercion before a call. -/// -/// When a caller passes a pointer in one address space to a callee whose -/// declared parameter lives in a different address space (the -/// `*mut SharedArray` / `addrspace(3)` case that surfaces from -/// `block_reduce` and friends), MIR must record the representational cast -/// before the exact call-signature verifier sees the argument. Lowering then -/// emits the corresponding `llvm.addrspacecast`. -/// -/// This test builds two MIR functions in one module: -/// - `callee(p: *mut i32 in addrspace(3))` -/// - `caller(p: *mut i32 in addrspace(0)) { callee(p) }` -/// -/// and asserts the lowered `caller` body contains an `AddrSpaceCastOp`. -#[test] -fn addrspace_coercion_inserts_addrspacecast_at_call_site() -> Result<(), anyhow::Error> { - use dialect_mir::types::MirPtrType; - use llvm_export::ops::AddrSpaceCastOp; - use pliron::basic_block::BasicBlock; - use pliron::builtin::attributes::{StringAttr, TypeAttr}; - use pliron::builtin::types::{FunctionType, IntegerType, Signedness}; - - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - mir_lower::register(&mut ctx); - - let module = ModuleOp::new(&mut ctx, "test_addrspace_coercion".try_into().unwrap()); - let module_ptr = module.get_operation(); - let module_region = module_ptr.deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let shared_ptr_ty = MirPtrType::get_shared(&mut ctx, i32_ty.into(), true); - let generic_ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), true); - let expected_call_location = Location::Named { - name: "source-device-call".to_string(), - child_loc: Box::new(Location::Unknown), - }; - - // Callee: takes a *mut i32 in addrspace(3), returns (). - let callee_func_ty = FunctionType::get(&ctx, vec![shared_ptr_ty.into()], vec![]); - let callee_func_op = Operation::new( - &mut ctx, - mir::MirFuncOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 1, - ); - let callee_func = mir::MirFuncOp::new( - &mut ctx, - callee_func_op, - TypeAttr::new(callee_func_ty.into()), - ); - callee_func.set_symbol_name(&mut ctx, "callee".try_into().unwrap()); - { - let region = callee_func.get_operation().deref(&ctx).get_region(0); - let block = BasicBlock::new(&mut ctx, None, vec![shared_ptr_ty.into()]); - block.insert_at_back(region, &ctx); - - let ret_op = Operation::new( - &mut ctx, - mir::MirReturnOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 0, - ); - ret_op.insert_at_back(block, &ctx); - } - callee_func - .get_operation() - .insert_at_back(module_block, &ctx); - - // Caller: takes a *mut i32 in addrspace(0), calls `callee` with that - // pointer. The lowerer is responsible for inserting an addrspacecast - // since the callee's declared addrspace differs. - let caller_func_ty = FunctionType::get(&ctx, vec![generic_ptr_ty.into()], vec![]); - let caller_func_op = Operation::new( - &mut ctx, - mir::MirFuncOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 1, - ); - let caller_func = mir::MirFuncOp::new( - &mut ctx, - caller_func_op, - TypeAttr::new(caller_func_ty.into()), - ); - caller_func.set_symbol_name(&mut ctx, "caller".try_into().unwrap()); - { - let region = caller_func.get_operation().deref(&ctx).get_region(0); - let block = BasicBlock::new(&mut ctx, None, vec![generic_ptr_ty.into()]); - block.insert_at_back(region, &ctx); - let arg = block.deref(&ctx).get_argument(0); - let cast_op_ptr = Operation::new( - &mut ctx, - mir::MirCastOp::get_concrete_op_info(), - vec![shared_ptr_ty.into()], - vec![arg], - vec![], - 0, - ); - mir::MirCastOp::new(cast_op_ptr) - .set_attr_cast_kind(&ctx, dialect_mir::attributes::MirCastKindAttr::PtrToPtr); - cast_op_ptr.insert_at_back(block, &ctx); - let coerced_arg = cast_op_ptr.deref(&ctx).get_result(0); - - let call_op_ptr = Operation::new( - &mut ctx, - mir::MirCallOp::get_concrete_op_info(), - vec![], - vec![coerced_arg], - vec![], - 0, - ); - let call_op = mir::MirCallOp::new(call_op_ptr); - call_op.set_attr_callee(&ctx, StringAttr::new("callee".to_string())); - call_op_ptr - .deref_mut(&ctx) - .set_loc(expected_call_location.clone()); - call_op_ptr.insert_at_back(block, &ctx); - - let ret_op = Operation::new( - &mut ctx, - mir::MirReturnOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 0, - ); - ret_op.insert_at_back(block, &ctx); - } - caller_func - .get_operation() - .insert_at_back(module_block, &ctx); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{}", e))?; - - let mut found_addrspace_cast = false; - let mut found_call = false; - let module_op = module_ptr.deref(&ctx); - let region = module_op.get_region(0); - let block = region.deref(&ctx).iter(&ctx).next().unwrap(); - for op in block.deref(&ctx).iter(&ctx) { - let Some(func_op) = Operation::get_op::(op, &ctx) else { - continue; - }; - if func_op.get_symbol_name(&ctx).to_string() != "caller" { - continue; - } - let func_region = func_op.get_operation().deref(&ctx).get_region(0); - for func_block in func_region.deref(&ctx).iter(&ctx) { - for body_op in func_block.deref(&ctx).iter(&ctx) { - if Operation::get_op::(body_op, &ctx).is_some() { - found_addrspace_cast = true; - } - if let Some(call) = Operation::get_op::(body_op, &ctx) - && matches!( - call.callee(&ctx), - CallOpCallable::Direct(symbol) if symbol.to_string() == "callee" - ) - { - found_call = true; - assert_eq!( - call.get_operation().deref(&ctx).loc(), - expected_call_location - ); - } - } - } - } - - assert!( - found_addrspace_cast, - "caller body must contain llvm.addrspacecast for the addrspace(0) -> (3) coercion at the call site", - ); - assert!(found_call, "caller body must contain the lowered call"); - Ok(()) -} - -/// A zero-sized MIR result is erased from the NVPTX function ABI, but its -/// typed value can remain live inside MIR (for example, when one ZST-returning -/// function returns the result of another). Lowering must keep the void call -/// for side effects and replace only its value result with a typed LLVM undef. -#[test] -fn zst_union_call_result_keeps_void_call_and_replaces_live_uses() -> Result<(), anyhow::Error> { - use dialect_mir::types::{MirTupleType, MirUnionType}; - use pliron::basic_block::BasicBlock; - use pliron::builtin::attributes::{StringAttr, TypeAttr}; - use pliron::builtin::types::FunctionType; - use pliron::r#type::{TypeHandle, Typed}; - - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - mir_lower::register(&mut ctx); - - let unit_ty: TypeHandle = MirTupleType::get(&mut ctx, vec![]).into(); - let union_ty: TypeHandle = MirUnionType::get( - &mut ctx, - "AlignedZeroUnion".into(), - vec!["unit".into()], - vec![unit_ty], - 0, - 16, - ) - .into(); - - let module = ModuleOp::new(&mut ctx, "test_zst_union_call".try_into().unwrap()); - let module_ptr = module.get_operation(); - let module_region = module_ptr.deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - - let callee_ty = FunctionType::get(&ctx, vec![], vec![union_ty]); - let callee_ptr = Operation::new( - &mut ctx, - mir::MirFuncOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 1, - ); - let callee = mir::MirFuncOp::new(&mut ctx, callee_ptr, TypeAttr::new(callee_ty.into())); - callee.set_symbol_name(&mut ctx, "make_zero".try_into().unwrap()); - { - let region = callee.get_operation().deref(&ctx).get_region(0); - let block = BasicBlock::new(&mut ctx, None, vec![]); - block.insert_at_back(region, &ctx); - - let undef = mir::MirUndefOp::new(&mut ctx, union_ty); - undef.get_operation().insert_at_back(block, &ctx); - let value = undef.get_operation().deref(&ctx).get_result(0); - - let ret = Operation::new( - &mut ctx, - mir::MirReturnOp::get_concrete_op_info(), - vec![], - vec![value], - vec![], - 0, - ); - ret.insert_at_back(block, &ctx); - } - callee.get_operation().insert_at_back(module_block, &ctx); - - let caller_ty = FunctionType::get(&ctx, vec![], vec![union_ty]); - let caller_ptr = Operation::new( - &mut ctx, - mir::MirFuncOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 1, - ); - let caller = mir::MirFuncOp::new(&mut ctx, caller_ptr, TypeAttr::new(caller_ty.into())); - caller.set_symbol_name(&mut ctx, "return_called_zero".try_into().unwrap()); - { - let region = caller.get_operation().deref(&ctx).get_region(0); - let block = BasicBlock::new(&mut ctx, None, vec![]); - block.insert_at_back(region, &ctx); - - let call_ptr = Operation::new( - &mut ctx, - mir::MirCallOp::get_concrete_op_info(), - vec![union_ty], - vec![], - vec![], - 0, - ); - let call = mir::MirCallOp::new(call_ptr); - call.set_attr_callee(&ctx, StringAttr::new("make_zero".to_string())); - call_ptr.insert_at_back(block, &ctx); - let value = call_ptr.deref(&ctx).get_result(0); - - let ret = Operation::new( - &mut ctx, - mir::MirReturnOp::get_concrete_op_info(), - vec![], - vec![value], - vec![], - 0, - ); - ret.insert_at_back(block, &ctx); - } - caller.get_operation().insert_at_back(module_block, &ctx); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{e}"))?; - - let mut found_void_call = false; - let mut caller_undefs = 0; - let mut found_void_return = false; - for op in module_block.deref(&ctx).iter(&ctx) { - let Some(func) = Operation::get_op::(op, &ctx) else { - continue; - }; - if func.get_symbol_name(&ctx).to_string() != "return_called_zero" { - continue; - } - let region = func.get_operation().deref(&ctx).get_region(0); - for block in region.deref(&ctx).iter(&ctx) { - for body_op in block.deref(&ctx).iter(&ctx) { - if let Some(call) = Operation::get_op::(body_op, &ctx) - && let CallOpCallable::Direct(callee) = call.callee(&ctx) - && callee.to_string() == "make_zero" - { - let call_op = call.get_operation().deref(&ctx); - assert_eq!(call_op.get_num_results(), 1); - assert!( - call_op - .get_result(0) - .get_type(&ctx) - .deref(&ctx) - .is::(), - "the ZST-returning call must use the void ABI" - ); - found_void_call = true; - } - if Operation::get_op::(body_op, &ctx).is_some() { - caller_undefs += 1; - } - if let Some(ret) = Operation::get_op::(body_op, &ctx) { - assert_eq!( - ret.get_operation().deref(&ctx).get_num_operands(), - 0, - "the caller's ZST return must also use the void ABI" - ); - found_void_return = true; - } - } - } - } - - assert!( - found_void_call, - "the LLVM call must be retained because the callee may have side effects" - ); - assert_eq!( - caller_undefs, 1, - "the live MIR result must be replaced by one typed LLVM undef" - ); - assert!( - found_void_return, - "the caller must retain its return terminator" - ); - Ok(()) -} - -/// Lock the comparison-predicate lowering table to the rustc_codegen_ssa -/// reference (`bin_op_to_fcmp_predicate` / `bin_op_to_icmp_predicate`): -/// -/// | MIR op | float `fcmp` | signed `icmp` | unsigned `icmp` | -/// |----------|-------------------|---------------|-----------------| -/// | `mir.eq` | `oeq` (ordered) | `eq` | `eq` | -/// | `mir.ne` | `une` (UNordered) | `ne` | `ne` | -/// | `mir.lt` | `olt` | `slt` | `ult` | -/// | `mir.le` | `ole` | `sle` | `ule` | -/// | `mir.gt` | `ogt` | `sgt` | `ugt` | -/// | `mir.ge` | `oge` | `sge` | `uge` | -/// -/// `ne` is the one float predicate that must be UNordered: Rust requires -/// `a != b == !(a == b)`, so `x != x` must be true for NaN (issue #123; -/// the ordered `one` folds the canonical NaN check to `false`). -/// -/// The test also locks fastmath flags to *empty* on every lowered `fcmp`: -/// a future `nnan` default would make `fcmp nnan une x, x` poison for NaN -/// and silently re-break NaN detection while the predicate assertion above -/// stays green. -#[test] -fn test_cmp_predicate_lowering() -> Result<(), anyhow::Error> { - use llvm_export::attributes::{FCmpPredicateAttr, FastmathFlagsAttr, ICmpPredicateAttr}; - use llvm_export::op_interfaces::FastMathFlags; - - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - mir_lower::register(&mut ctx); - - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_ptr = module.get_operation(); - - let f32_ty = pliron::builtin::types::FP32Type::get(&ctx); - let i32_signed = pliron::builtin::types::IntegerType::get( - &ctx, - 32, - pliron::builtin::types::Signedness::Signed, - ); - let u32_unsigned = pliron::builtin::types::IntegerType::get( - &ctx, - 32, - pliron::builtin::types::Signedness::Unsigned, - ); - let bool_ty = pliron::builtin::types::IntegerType::get( - &ctx, - 1, - pliron::builtin::types::Signedness::Signless, - ); - - // Args: (f32, f32, i32, u32). The integer args carry pre-conversion - // signedness, which is what selects signed vs unsigned icmp predicates. - let arg_tys: Vec = vec![ - f32_ty.into(), - f32_ty.into(), - i32_signed.into(), - u32_unsigned.into(), - ]; - let func_name = "cmp_func"; - let func_ty = pliron::builtin::types::FunctionType::get(&ctx, arg_tys.clone(), vec![]); - - let func_op_ptr = Operation::new( - &mut ctx, - mir::MirFuncOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 1, - ); - let func_ty_attr = pliron::builtin::attributes::TypeAttr::new(func_ty.into()); - let func = mir::MirFuncOp::new(&mut ctx, func_op_ptr, func_ty_attr); - func.set_symbol_name(&mut ctx, func_name.try_into().unwrap()); - - let region = func.get_operation().deref(&ctx).get_region(0); - let block = { - let b = pliron::basic_block::BasicBlock::new(&mut ctx, None, arg_tys); - b.insert_at_back(region, &ctx); - b - }; - let fa = block.deref(&ctx).get_argument(0); - let fb = block.deref(&ctx).get_argument(1); - let si = block.deref(&ctx).get_argument(2); - let ui = block.deref(&ctx).get_argument(3); - - // One comparison op per table row, in a fixed program order. The raw - // `Operation::new` construction mirrors how the importer builds these - // ops (mir-importer translator/rvalue.rs BinaryOp arm). - let cmp_infos = [ - // Floats: all six predicates. - (mir::MirEqOp::get_concrete_op_info(), fa, fb), - (mir::MirNeOp::get_concrete_op_info(), fa, fb), - (mir::MirLtOp::get_concrete_op_info(), fa, fb), - (mir::MirLeOp::get_concrete_op_info(), fa, fb), - (mir::MirGtOp::get_concrete_op_info(), fa, fb), - (mir::MirGeOp::get_concrete_op_info(), fa, fb), - // Signed integers: eq/ne are sign-agnostic, the rest must be s*. - (mir::MirEqOp::get_concrete_op_info(), si, si), - (mir::MirNeOp::get_concrete_op_info(), si, si), - (mir::MirLtOp::get_concrete_op_info(), si, si), - (mir::MirLeOp::get_concrete_op_info(), si, si), - (mir::MirGtOp::get_concrete_op_info(), si, si), - (mir::MirGeOp::get_concrete_op_info(), si, si), - // Unsigned integers: the relational predicates must be u*. - (mir::MirLtOp::get_concrete_op_info(), ui, ui), - (mir::MirLeOp::get_concrete_op_info(), ui, ui), - (mir::MirGtOp::get_concrete_op_info(), ui, ui), - (mir::MirGeOp::get_concrete_op_info(), ui, ui), - ]; - for (info, lhs, rhs) in cmp_infos { - let op = Operation::new( - &mut ctx, - info, - vec![bool_ty.into()], - vec![lhs, rhs], - vec![], - 0, - ); - op.insert_at_back(block, &ctx); - } - - let ret_op_ptr = Operation::new( - &mut ctx, - mir::MirReturnOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 0, - ); - ret_op_ptr.insert_at_back(block, &ctx); - - let module_region = module.get_operation().deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - func.get_operation().insert_at_back(module_block, &ctx); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{}", e))?; - - // Collect lowered predicates in program order. - let mut fcmp_preds = Vec::new(); - let mut icmp_preds = Vec::new(); - let module_op = module_ptr.deref(&ctx); - let region = module_op.get_region(0); - let block = region.deref(&ctx).iter(&ctx).next().unwrap(); - for op in block.deref(&ctx).iter(&ctx) { - let Some(func_op) = Operation::get_op::(op, &ctx) else { - continue; - }; - if func_op.get_symbol_name(&ctx).to_string() != func_name { - continue; - } - let func_region = func_op.get_operation().deref(&ctx).get_region(0); - for func_block in func_region.deref(&ctx).iter(&ctx) { - for body_op in func_block.deref(&ctx).iter(&ctx) { - if let Some(fcmp) = Operation::get_op::(body_op, &ctx) { - fcmp_preds.push(fcmp.predicate(&ctx)); - // fcmp carries `contract` (set by add_fastmath_flags) which is a - // no-op for comparisons at the LLVM / PTX level. Critically, nnan - // is NOT set, so NaN checks like `x != x` still evaluate correctly. - let expected: FastmathFlagsAttr = - llvm_export::attributes::FastmathFlags::CONTRACT.into(); - assert_eq!( - fcmp.fast_math_flags(&ctx), - expected, - "fcmp must carry only the contract flag (nnan would poison NaN checks)" - ); - } - if let Some(icmp) = Operation::get_op::(body_op, &ctx) { - icmp_preds.push(icmp.predicate(&ctx)); - } - } - } - } - - assert_eq!( - fcmp_preds, - vec![ - FCmpPredicateAttr::OEQ, - FCmpPredicateAttr::UNE, - FCmpPredicateAttr::OLT, - FCmpPredicateAttr::OLE, - FCmpPredicateAttr::OGT, - FCmpPredicateAttr::OGE, - ], - "float comparison predicates must mirror rustc: ordered except Ne (une)" - ); - assert_eq!( - icmp_preds, - vec![ - ICmpPredicateAttr::EQ, - ICmpPredicateAttr::NE, - ICmpPredicateAttr::SLT, - ICmpPredicateAttr::SLE, - ICmpPredicateAttr::SGT, - ICmpPredicateAttr::SGE, - ICmpPredicateAttr::ULT, - ICmpPredicateAttr::ULE, - ICmpPredicateAttr::UGT, - ICmpPredicateAttr::UGE, - ], - "integer comparison predicates must respect pre-conversion signedness" - ); - Ok(()) -} - -/// Helper: fresh context with all dialects registered. -fn make_test_ctx() -> Context { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - mir_lower::register(&mut ctx); - ctx -} - -/// Helper: build a module + MirFuncOp("kernel_func") with given arg types, -/// returning the module ptr and entry block. -fn build_test_kernel( - ctx: &mut Context, - arg_tys: Vec, -) -> ( - pliron::context::Ptr, - pliron::context::Ptr, -) { - use pliron::basic_block::BasicBlock; - use pliron::builtin::attributes::TypeAttr; - use pliron::builtin::types::FunctionType; - - let module = ModuleOp::new(ctx, "test_module".try_into().unwrap()); - let module_ptr = module.get_operation(); - - let func_ty = FunctionType::get(ctx, arg_tys.clone(), vec![]); - let func_op_ptr = Operation::new( - ctx, - mir::MirFuncOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 1, - ); - let func = mir::MirFuncOp::new(ctx, func_op_ptr, TypeAttr::new(func_ty.into())); - func.set_symbol_name(ctx, "kernel_func".try_into().unwrap()); - - let region = func.get_operation().deref(ctx).get_region(0); - let entry = BasicBlock::new(ctx, None, arg_tys); - entry.insert_at_back(region, ctx); - - let module_region = module_ptr.deref(ctx).get_region(0); - let module_block = module_region.deref(ctx).iter(ctx).next().unwrap(); - func.get_operation().insert_at_back(module_block, ctx); - - (module_ptr, entry) -} - -/// Helper: append a mir.return (void) to a block. -fn append_return(ctx: &mut Context, block: pliron::context::Ptr) { - let ret = Operation::new( - ctx, - mir::MirReturnOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 0, - ); - ret.insert_at_back(block, ctx); -} - -fn lower_basic_mbarrier( - backend: mir_lower::IntrinsicBackend, -) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { - use dialect_mir::types::MirPtrType; - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let bar_ptr_ty = MirPtrType::get_shared(&mut ctx, u64_ty.into(), false); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![bar_ptr_ty.into(), u32_ty.into()]); - let barrier = entry.deref(&ctx).get_argument(0); - let count = entry.deref(&ctx).get_argument(1); - - nvvm::MbarrierInitSharedOp::build(&mut ctx, barrier, count).insert_at_back(entry, &ctx); - let arrive = nvvm::MbarrierArriveSharedOp::build(&mut ctx, barrier); - let token = arrive.deref(&ctx).get_result(0); - arrive.insert_at_back(entry, &ctx); - nvvm::MbarrierTestWaitSharedOp::build(&mut ctx, barrier, token).insert_at_back(entry, &ctx); - nvvm::MbarrierInvalSharedOp::build(&mut ctx, barrier).insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm_with_options( - &mut ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: backend, - ..Default::default() - }, - ) - .map_err(|error| anyhow::anyhow!("{error}"))?; - Ok((ctx, module_ptr)) -} - -#[test] -fn test_generated_basic_mbarrier_uses_shared_lowering_on_both_backends() -> Result<(), anyhow::Error> -{ - let expected_calls = [ - ("llvm_nvvm_mbarrier_init_shared", 2), - ("llvm_nvvm_mbarrier_arrive_shared", 1), - ("llvm_nvvm_mbarrier_inval_shared", 1), - ]; - let init_template = "mbarrier.init.shared.b64 [$0], $1;"; - let arrive_template = "mbarrier.arrive.shared.b64 $0, [$1];"; - let test_wait_template = - "{ .reg .pred %p0; mbarrier.test_wait.shared.b64 %p0, [$1], $2; selp.b32 $0, 1, 0, %p0; }"; - let inval_template = "mbarrier.inval.shared.b64 [$0];"; - - for backend in [ - mir_lower::IntrinsicBackend::LlvmNvptx, - mir_lower::IntrinsicBackend::LibNvvm, - ] { - let (ctx, module_ptr) = lower_basic_mbarrier(backend)?; - let mut call_counts = [0usize; 3]; - let expected_asm = match backend { - mir_lower::IntrinsicBackend::LlvmNvptx => { - vec![(test_wait_template, "=r,l,l,~{memory}", 2)] - } - mir_lower::IntrinsicBackend::LibNvvm => vec![ - (init_template, "l,r,~{memory}", 2), - (arrive_template, "=l,l,~{memory}", 1), - (test_wait_template, "=r,l,l,~{memory}", 2), - (inval_template, "l,~{memory}", 1), - ], - }; - let mut asm_counts = vec![0usize; expected_asm.len()]; - let mut trunc_count = 0; - - for op in lowered_kernel_body(&ctx, module_ptr) { - if let Some(call) = Operation::get_op::(op, &ctx) { - let CallOpCallable::Direct(callee) = call.callee(&ctx) else { - continue; - }; - let callee = callee.to_string(); - let Some(index) = expected_calls - .iter() - .position(|(expected, _)| callee == *expected) - else { - assert_ne!(callee, "llvm_nvvm_mbarrier_test_wait_shared"); - continue; - }; - call_counts[index] += 1; - assert_eq!(op.deref(&ctx).get_num_operands(), expected_calls[index].1); - } - - if let Some(inline_asm) = Operation::get_op::(op, &ctx) { - let template = inline_asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())); - let index = expected_asm - .iter() - .position(|(expected, _, _)| template.as_deref() == Some(*expected)) - .unwrap_or_else(|| panic!("unexpected {backend:?} inline PTX: {template:?}")); - let (_, constraints, operand_count) = expected_asm[index]; - asm_counts[index] += 1; - assert_eq!( - inline_asm - .get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some(constraints) - ); - assert_eq!(llvm::asm_kind(&ctx, &inline_asm), llvm::AsmKind::Convergent); - let asm = op.deref(&ctx); - assert_eq!(asm.get_num_operands(), operand_count); - assert_eq!(asm.get_num_results(), 1); - } - - if Operation::get_op::(op, &ctx).is_some() { - trunc_count += 1; - } - } - - match backend { - mir_lower::IntrinsicBackend::LlvmNvptx => assert_eq!(call_counts, [1; 3]), - mir_lower::IntrinsicBackend::LibNvvm => assert_eq!(call_counts, [0; 3]), - } - assert_eq!(asm_counts, vec![1; expected_asm.len()]); - assert_eq!( - trunc_count, 1, - "test-wait must adapt its i32 predicate to i1" - ); - - let module = Operation::get_op::(module_ptr, &ctx).unwrap(); - let ir = llvm_export::export::export_module_to_string(&ctx, &module) - .map_err(|error| anyhow::anyhow!(error))?; - match backend { - mir_lower::IntrinsicBackend::LlvmNvptx => { - assert!( - ir.contains("call void @llvm.nvvm.mbarrier.init.shared(ptr addrspace(3)"), - "{ir}" - ); - assert!( - ir.contains("call i64 @llvm.nvvm.mbarrier.arrive.shared(ptr addrspace(3)"), - "{ir}" - ); - assert!( - ir.contains("call void @llvm.nvvm.mbarrier.inval.shared(ptr addrspace(3)"), - "{ir}" - ); - assert!(!ir.contains(init_template), "{ir}"); - assert!(!ir.contains(arrive_template), "{ir}"); - assert!(!ir.contains(inval_template), "{ir}"); - } - mir_lower::IntrinsicBackend::LibNvvm => { - for symbol in [ - "@llvm.nvvm.mbarrier.init.shared", - "@llvm.nvvm.mbarrier.arrive.shared", - "@llvm.nvvm.mbarrier.inval.shared", - ] { - assert!( - !ir.contains(symbol), - "libNVVM route retained {symbol}:\n{ir}" - ); - } - for template in [init_template, arrive_template, inval_template] { - assert!(ir.contains(template), "{ir}"); - } - } - } - assert!(ir.contains(test_wait_template), "{ir}"); - assert!(ir.contains("asm sideeffect"), "{ir}"); - assert!(ir.contains("trunc i32") && ir.contains("to i1"), "{ir}"); - assert!(ir.contains("attributes #0 = { convergent }"), "{ir}"); - } - Ok(()) -} - -fn lower_cluster_barriers( - backend: mir_lower::IntrinsicBackend, -) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { - let mut ctx = make_test_ctx(); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![]); - for mode in [ - nvvm::ClusterBarrierModeAttr::Arrive, - nvvm::ClusterBarrierModeAttr::ArriveAligned, - nvvm::ClusterBarrierModeAttr::ArriveRelaxed, - nvvm::ClusterBarrierModeAttr::ArriveRelaxedAligned, - nvvm::ClusterBarrierModeAttr::Wait, - nvvm::ClusterBarrierModeAttr::WaitAligned, - ] { - nvvm::ClusterBarrierOp::build(&mut ctx, mode).insert_at_back(entry, &ctx); - } - append_return(&mut ctx, entry); - mir_lower::lower_mir_to_llvm_with_options( - &mut ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: backend, - ..Default::default() - }, - ) - .map_err(|error| anyhow::anyhow!("{error}"))?; - Ok((ctx, module_ptr)) -} - -#[test] -fn generated_cluster_barriers_lower_exactly_on_both_backends() -> Result<(), anyhow::Error> { - let recipes = [ - ( - "llvm_nvvm_barrier_cluster_arrive", - "barrier.cluster.arrive;", - ), - ( - "llvm_nvvm_barrier_cluster_arrive_aligned", - "barrier.cluster.arrive.aligned;", - ), - ( - "llvm_nvvm_barrier_cluster_arrive_relaxed", - "barrier.cluster.arrive.relaxed;", - ), - ( - "llvm_nvvm_barrier_cluster_arrive_relaxed_aligned", - "barrier.cluster.arrive.relaxed.aligned;", - ), - ("llvm_nvvm_barrier_cluster_wait", "barrier.cluster.wait;"), - ( - "llvm_nvvm_barrier_cluster_wait_aligned", - "barrier.cluster.wait.aligned;", - ), - ]; - - for backend in [ - mir_lower::IntrinsicBackend::LlvmNvptx, - mir_lower::IntrinsicBackend::LibNvvm, - ] { - let (ctx, module_ptr) = lower_cluster_barriers(backend)?; - let mut calls = [0usize; 6]; - let mut asm = [0usize; 6]; - for op in lowered_kernel_body(&ctx, module_ptr) { - if let Some(call) = Operation::get_op::(op, &ctx) { - let CallOpCallable::Direct(callee) = call.callee(&ctx) else { - continue; - }; - if let Some(index) = recipes - .iter() - .position(|(symbol, _)| callee.to_string() == *symbol) - { - calls[index] += 1; - } - } - if let Some(inline_asm) = Operation::get_op::(op, &ctx) { - let template = inline_asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())); - if let Some(index) = recipes - .iter() - .position(|(_, expected)| template.as_deref() == Some(*expected)) - { - asm[index] += 1; - assert_eq!( - inline_asm - .get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some("~{memory}") - ); - assert_eq!(llvm::asm_kind(&ctx, &inline_asm), llvm::AsmKind::Convergent); - } - } - } - match backend { - mir_lower::IntrinsicBackend::LlvmNvptx => { - assert_eq!(calls, [1; 6]); - assert_eq!(asm, [0; 6]); - } - mir_lower::IntrinsicBackend::LibNvvm => { - assert_eq!(calls, [0; 6]); - assert_eq!(asm, [1; 6]); - } - } - } - Ok(()) -} - -#[test] -fn test_cluster_mbarrier_and_fences_lower_to_exact_inline_ptx() -> Result<(), anyhow::Error> { - use dialect_mir::types::MirPtrType; - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let i64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let bar_ptr_ty = MirPtrType::get_shared(&mut ctx, i64_ty.into(), false); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![bar_ptr_ty.into(), i32_ty.into()]); - - let bar_ptr = entry.deref(&ctx).get_argument(0); - let bytes_or_parity = entry.deref(&ctx).get_argument(1); - - let arrive = Operation::new( - &mut ctx, - nvvm::MbarrierArriveExpectTxClusterOp::get_concrete_op_info(), - vec![i64_ty.into()], - vec![bar_ptr, bytes_or_parity], - vec![], - 0, - ); - arrive.insert_at_back(entry, &ctx); - - let try_wait = Operation::new( - &mut ctx, - nvvm::MbarrierTryWaitParityClusterOp::get_concrete_op_info(), - vec![i1_ty.into()], - vec![bar_ptr, bytes_or_parity], - vec![], - 0, - ); - try_wait.insert_at_back(entry, &ctx); - - let mbarrier_fence = Operation::new( - &mut ctx, - nvvm::FenceMbarrierInitReleaseClusterOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 0, - ); - mbarrier_fence.insert_at_back(entry, &ctx); - - let proxy_release_fence = Operation::new( - &mut ctx, - nvvm::FenceProxyAsyncGenericReleaseSharedCtaClusterOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 0, - ); - proxy_release_fence.insert_at_back(entry, &ctx); - - let proxy_acquire_fence = Operation::new( - &mut ctx, - nvvm::FenceProxyAsyncGenericAcquireSharedClusterClusterOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 0, - ); - proxy_acquire_fence.insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let expected = [ - ( - "mbarrier.arrive.expect_tx.relaxed.cluster.shared::cta.b64 $0, [$1], $2;", - "=l,l,r,~{memory}", - ), - ( - "{ .reg .pred %p0; mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 %p0, [$1], $2; selp.b32 $0, 1, 0, %p0; }", - "=r,l,r,~{memory}", - ), - ("fence.mbarrier_init.release.cluster;", "~{memory}"), - ( - "fence.proxy.async::generic.release.sync_restrict::shared::cta.cluster;", - "~{memory}", - ), - ( - "fence.proxy.async::generic.acquire.sync_restrict::shared::cluster.cluster;", - "~{memory}", - ), - ]; - let mut matches = [0usize; 5]; - - let module_region = module_ptr.deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - for op in module_block.deref(&ctx).iter(&ctx) { - let Some(func_op) = Operation::get_op::(op, &ctx) else { - continue; - }; - if func_op.get_symbol_name(&ctx).to_string() != "kernel_func" { - continue; - } - - let func_region = func_op.get_operation().deref(&ctx).get_region(0); - for func_block in func_region.deref(&ctx).iter(&ctx) { - for body_op in func_block.deref(&ctx).iter(&ctx) { - let Some(inline_asm) = Operation::get_op::(body_op, &ctx) else { - continue; - }; - let template = inline_asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())); - let Some(index) = expected.iter().position(|(expected_template, _)| { - template.as_deref() == Some(*expected_template) - }) else { - continue; - }; - - matches[index] += 1; - assert_eq!( - inline_asm - .get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some(expected[index].1) - ); - assert_eq!( - llvm::asm_kind_opt(&ctx, &inline_asm), - Some(llvm::AsmKind::Convergent) - ); - } - } - } - - assert_eq!( - matches, [1; 5], - "each cluster barrier/fence must lower to its exact PTX template once" - ); - Ok(()) -} - -#[test] -fn test_fast_float_intrinsics_lower_to_explicit_fast_binops() -> Result<(), anyhow::Error> { - use dialect_mir::rust_intrinsics; - use llvm_export::attributes::{FastmathFlags, FastmathFlagsAttr}; - use llvm_export::op_interfaces::FastMathFlags; - use pliron::builtin::attributes::StringAttr; - use pliron::builtin::op_interfaces::CallOpInterface; - use pliron::builtin::types::{FP32Type, FP64Type}; - use pliron::r#type::{TypeHandle, Typed}; - - let mut ctx = make_test_ctx(); - let f32_ty = FP32Type::get(&ctx); - let f64_ty = FP64Type::get(&ctx); - let f32_ty_obj: TypeHandle = f32_ty.into(); - let f64_ty_obj: TypeHandle = f64_ty.into(); - let (module_ptr, entry) = build_test_kernel( - &mut ctx, - vec![f32_ty_obj, f32_ty_obj, f64_ty_obj, f64_ty_obj], - ); - let f32_lhs = entry.deref(&ctx).get_argument(0); - let f32_rhs = entry.deref(&ctx).get_argument(1); - let f64_lhs = entry.deref(&ctx).get_argument(2); - let f64_rhs = entry.deref(&ctx).get_argument(3); - - for (callee, lhs, rhs, result_ty) in [ - ( - rust_intrinsics::CALLEE_FADD_FAST, - f32_lhs, - f32_rhs, - f32_ty_obj, - ), - ( - rust_intrinsics::CALLEE_FSUB_FAST, - f32_lhs, - f32_rhs, - f32_ty_obj, - ), - ( - rust_intrinsics::CALLEE_FMUL_FAST, - f32_lhs, - f32_rhs, - f32_ty_obj, - ), - ( - rust_intrinsics::CALLEE_FDIV_FAST, - f32_lhs, - f32_rhs, - f32_ty_obj, - ), - ( - rust_intrinsics::CALLEE_FREM_FAST, - f32_lhs, - f32_rhs, - f32_ty_obj, - ), - ( - rust_intrinsics::CALLEE_FADD_FAST, - f64_lhs, - f64_rhs, - f64_ty_obj, - ), - ( - rust_intrinsics::CALLEE_FSUB_FAST, - f64_lhs, - f64_rhs, - f64_ty_obj, - ), - ( - rust_intrinsics::CALLEE_FMUL_FAST, - f64_lhs, - f64_rhs, - f64_ty_obj, - ), - ( - rust_intrinsics::CALLEE_FDIV_FAST, - f64_lhs, - f64_rhs, - f64_ty_obj, - ), - ( - rust_intrinsics::CALLEE_FREM_FAST, - f64_lhs, - f64_rhs, - f64_ty_obj, - ), - ] { - let call_ptr = Operation::new( - &mut ctx, - mir::MirCallOp::get_concrete_op_info(), - vec![result_ty], - vec![lhs, rhs], - vec![], - 0, - ); - let call = mir::MirCallOp::new(call_ptr); - call.set_attr_callee(&ctx, StringAttr::new(callee.to_string())); - call_ptr.insert_at_back(entry, &ctx); - } - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{}", e))?; - - let explicit_fast_flags: FastmathFlagsAttr = FastmathFlags::FAST.into(); - assert_ne!( - explicit_fast_flags, - FastmathFlagsAttr::default(), - "FastmathFlagsAttr::default() is empty; f*_fast must use explicit fast flags" - ); - - let mut fadd_counts = [0usize; 2]; - let mut fsub_counts = [0usize; 2]; - let mut fmul_counts = [0usize; 2]; - let mut fdiv_counts = [0usize; 2]; - let mut frem_counts = [0usize; 2]; - - macro_rules! count_fast_binop { - ($body_op:expr, $op_ty:ty, $counts:ident, $name:literal) => { - if let Some(op) = Operation::get_op::<$op_ty>($body_op, &ctx) { - assert_eq!( - op.fast_math_flags(&ctx), - explicit_fast_flags, - concat!($name, " must carry explicit LLVM fast-math flags") - ); - let result_ty = op.get_operation().deref(&ctx).get_result(0).get_type(&ctx); - if result_ty == f32_ty_obj { - $counts[0] += 1; - } else if result_ty == f64_ty_obj { - $counts[1] += 1; - } else { - panic!(concat!($name, " lowered to an unexpected result type")); - } - } - }; - } - - let module_op = module_ptr.deref(&ctx); - let region = module_op.get_region(0); - let block = region.deref(&ctx).iter(&ctx).next().unwrap(); - for op in block.deref(&ctx).iter(&ctx) { - let Some(func_op) = Operation::get_op::(op, &ctx) else { - continue; - }; - if func_op.get_symbol_name(&ctx).to_string() != "kernel_func" { - continue; - } - let func_region = func_op.get_operation().deref(&ctx).get_region(0); - for func_block in func_region.deref(&ctx).iter(&ctx) { - for body_op in func_block.deref(&ctx).iter(&ctx) { - assert!( - Operation::get_op::(body_op, &ctx).is_none(), - "f*_fast placeholder mir.call must not survive MIR lowering" - ); - if let Some(call) = Operation::get_op::(body_op, &ctx) - && let CallOpCallable::Direct(sym) = call.callee(&ctx) - { - let callee = sym.to_string(); - assert!( - !callee.starts_with(rust_intrinsics::PLACEHOLDER_PREFIX), - "lowered LLVM must not call unresolved Rust intrinsic placeholder `{callee}`" - ); - } - count_fast_binop!(body_op, llvm::FAddOp, fadd_counts, "fadd_fast"); - count_fast_binop!(body_op, llvm::FSubOp, fsub_counts, "fsub_fast"); - count_fast_binop!(body_op, llvm::FMulOp, fmul_counts, "fmul_fast"); - count_fast_binop!(body_op, llvm::FDivOp, fdiv_counts, "fdiv_fast"); - count_fast_binop!(body_op, llvm::FRemOp, frem_counts, "frem_fast"); - } - } - } - - assert_eq!(fadd_counts, [1, 1], "fadd_fast must lower for f32 and f64"); - assert_eq!(fsub_counts, [1, 1], "fsub_fast must lower for f32 and f64"); - assert_eq!(fmul_counts, [1, 1], "fmul_fast must lower for f32 and f64"); - assert_eq!(fdiv_counts, [1, 1], "fdiv_fast must lower for f32 and f64"); - assert_eq!(frem_counts, [1, 1], "frem_fast must lower for f32 and f64"); - - Ok(()) -} - -#[test] -fn test_inline_ptx_op_lowers_to_inline_asm_attrs() -> Result<(), anyhow::Error> { - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![i32_ty.into()]); - let input = entry.deref(&ctx).get_argument(0); - - let inline_ptx = nvvm::InlinePtxOp::build( - &mut ctx, - vec![i32_ty.into()], - vec![input], - "add.u32 $0, $1, $1;", - "=r,r", - true, - true, - ); - inline_ptx.insert_at_back(entry, &ctx); - let register_only_ptx = nvvm::InlinePtxOp::build( - &mut ctx, - vec![i32_ty.into()], - vec![input], - "mul.lo.u32 $0, $1, $1;", - "=r,r", - false, - true, - ); - register_only_ptx.insert_at_back(entry, &ctx); - let may_diverge_ptx = nvvm::InlinePtxOp::build( - &mut ctx, - vec![i32_ty.into()], - vec![input], - "cvt.u32.u32 $0, $1;", - "=r,r", - false, - false, - ); - may_diverge_ptx.insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{}", e))?; - - let mut found_conservative = false; - let mut found_register_only = false; - let mut found_may_diverge = false; - let module_op = module_ptr.deref(&ctx); - let region = module_op.get_region(0); - let block = region.deref(&ctx).iter(&ctx).next().unwrap(); - - for op in block.deref(&ctx).iter(&ctx) { - let Some(func_op) = Operation::get_op::(op, &ctx) else { - continue; - }; - if func_op.get_symbol_name(&ctx).to_string() != "kernel_func" { - continue; - } - let func_region = func_op.get_operation().deref(&ctx).get_region(0); - for func_block in func_region.deref(&ctx).iter(&ctx) { - for body_op in func_block.deref(&ctx).iter(&ctx) { - let Some(inline_asm) = Operation::get_op::(body_op, &ctx) else { - continue; - }; - let template = inline_asm - .get_attr_inline_asm_template(&ctx) - .map(|s| String::from((*s).clone())); - match template.as_deref() { - Some("add.u32 $0, $1, $1;") => { - found_conservative = true; - assert_eq!( - inline_asm - .get_attr_inline_asm_constraints(&ctx) - .map(|s| String::from((*s).clone())) - .as_deref(), - Some("=r,r") - ); - assert!( - inline_asm - .get_attr_inline_asm_convergent(&ctx) - .is_some_and(|b| bool::from((*b).clone())) - ); - assert!(llvm::inline_asm_sideeffect( - &ctx, - inline_asm.get_operation() - )); - } - Some("mul.lo.u32 $0, $1, $1;") => { - found_register_only = true; - assert_eq!( - inline_asm - .get_attr_inline_asm_constraints(&ctx) - .map(|s| String::from((*s).clone())) - .as_deref(), - Some("=r,r") - ); - assert!( - inline_asm - .get_attr_inline_asm_convergent(&ctx) - .is_some_and(|b| bool::from((*b).clone())) - ); - assert!(!llvm::inline_asm_sideeffect( - &ctx, - inline_asm.get_operation() - )); - } - Some("cvt.u32.u32 $0, $1;") => { - found_may_diverge = true; - assert_eq!( - inline_asm - .get_attr_inline_asm_constraints(&ctx) - .map(|s| String::from((*s).clone())) - .as_deref(), - Some("=r,r") - ); - assert!( - inline_asm - .get_attr_inline_asm_convergent(&ctx) - .is_some_and(|b| !bool::from((*b).clone())) - ); - assert!(!llvm::inline_asm_sideeffect( - &ctx, - inline_asm.get_operation() - )); - } - _ => continue, - } - } - } - } - - assert!( - found_conservative, - "Expected conservative inline PTX asm op" - ); - assert!( - found_register_only, - "Expected register-only inline PTX asm op" - ); - assert!(found_may_diverge, "Expected may-diverge inline PTX asm op"); - Ok(()) -} - -#[test] -fn test_multi_result_inline_ptx_lowers_to_struct_asm_and_extractvalues() -> Result<(), anyhow::Error> -{ - use llvm_export::types as llvm_types; - use pliron::builtin::types::{IntegerType, Signedness}; - use pliron::r#type::Typed; - - let mut ctx = make_test_ctx(); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![i32_ty.into()]); - let input = entry.deref(&ctx).get_argument(0); - - let inline_ptx = nvvm::InlinePtxOp::build( - &mut ctx, - vec![i32_ty.into(), i32_ty.into()], - vec![input], - "add.u32 $0, $2, $2; mul.lo.u32 $1, $2, $2;", - "=r,=r,r", - true, - false, - ); - let expected_location = Location::Named { - name: "source-inline-ptx".to_string(), - child_loc: Box::new(Location::Unknown), - }; - inline_ptx - .deref_mut(&ctx) - .set_loc(expected_location.clone()); - inline_ptx.insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{}", e))?; - - let mut asm_result = None; - let mut extract_indices = Vec::new(); - for op in lowered_kernel_body(&ctx, module_ptr) { - if let Some(inline_asm) = Operation::get_op::(op, &ctx) { - assert_eq!( - inline_asm.get_operation().deref(&ctx).loc(), - expected_location - ); - assert_eq!( - inline_asm - .get_attr_inline_asm_template(&ctx) - .map(|s| String::from((*s).clone())) - .as_deref(), - Some("add.u32 $0, $2, $2; mul.lo.u32 $1, $2, $2;") - ); - assert_eq!( - inline_asm - .get_attr_inline_asm_constraints(&ctx) - .map(|s| String::from((*s).clone())) - .as_deref(), - Some("=r,=r,r") - ); - let result = inline_asm.get_operation().deref(&ctx).get_result(0); - let result_ty = result.get_type(&ctx); - let result_ty = result_ty.deref(&ctx); - let struct_ty = result_ty - .downcast_ref::() - .expect("multi-output inline PTX must return an LLVM struct"); - assert_eq!(struct_ty.num_fields(), 2); - for index in 0..2 { - assert_eq!( - struct_ty - .field_type(index) - .deref(&ctx) - .downcast_ref::() - .expect("multi-output inline PTX struct field must stay i32") - .width(), - 32 - ); - } - asm_result = Some(result); - } else if let Some(extract) = Operation::get_op::(op, &ctx) { - assert_eq!(extract.get_operation().deref(&ctx).loc(), expected_location); - let aggregate = extract.get_operation().deref(&ctx).get_operand(0); - assert_eq!( - Some(aggregate), - asm_result, - "extractvalue must consume the struct-returning asm result" - ); - extract_indices.push(extract.indices(&ctx)); - } - } - - assert!( - asm_result.is_some(), - "Expected struct-returning inline PTX asm op" - ); - assert_eq!( - extract_indices, - vec![vec![0], vec![1]], - "each output must be extracted once, in constraint order" - ); - Ok(()) -} - -#[test] -fn test_inline_ptx_supports_thirty_two_tied_f32_results() -> Result<(), anyhow::Error> { - use llvm_export::types as llvm_types; - use pliron::builtin::types::FP32Type; - use pliron::r#type::Typed; - - const ACCUMULATOR_LEN: usize = 32; - - let mut ctx = make_test_ctx(); - let f32_ty = FP32Type::get(&ctx); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![f32_ty.into(); ACCUMULATOR_LEN]); - - let inputs = (0..ACCUMULATOR_LEN) - .map(|index| entry.deref(&ctx).get_argument(index)) - .collect::>(); - - let template = (0..ACCUMULATOR_LEN) - .map(|index| format!("mov.f32 ${index}, ${index};")) - .collect::>() - .join(" "); - - let mut constraints = vec!["=f".to_string(); ACCUMULATOR_LEN]; - constraints.extend((0..ACCUMULATOR_LEN).map(|index| index.to_string())); - let constraints = constraints.join(","); - - let inline_ptx = nvvm::InlinePtxOp::build( - &mut ctx, - vec![f32_ty.into(); ACCUMULATOR_LEN], - inputs, - &template, - &constraints, - true, - true, - ); - inline_ptx.insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let mut asm_result = None; - let mut extract_indices = Vec::new(); - - for op in lowered_kernel_body(&ctx, module_ptr) { - if let Some(inline_asm) = Operation::get_op::(op, &ctx) { - assert_eq!( - inline_asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some(template.as_str()), - ); - - assert_eq!( - inline_asm - .get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some(constraints.as_str()), - ); - - let result = inline_asm.get_operation().deref(&ctx).get_result(0); - let result_ty = result.get_type(&ctx); - let result_ty = result_ty.deref(&ctx); - - let struct_ty = result_ty - .downcast_ref::() - .expect("32-output inline PTX must return an LLVM struct"); - - assert_eq!( - struct_ty.num_fields(), - ACCUMULATOR_LEN, - "inline PTX must return exactly 32 accumulator values", - ); - - for index in 0..ACCUMULATOR_LEN { - assert!( - struct_ty - .field_type(index) - .deref(&ctx) - .downcast_ref::() - .is_some(), - "inline PTX result field {index} must remain f32", - ); - } - - assert!( - asm_result.replace(result).is_none(), - "expected exactly one inline PTX operation", - ); - } else if let Some(extract) = Operation::get_op::(op, &ctx) { - let aggregate = extract.get_operation().deref(&ctx).get_operand(0); - - assert_eq!( - Some(aggregate), - asm_result, - "extractvalue must consume the struct-returning asm result", - ); - - extract_indices.push(extract.indices(&ctx)); - } - } - - assert!( - asm_result.is_some(), - "expected a struct-returning inline PTX asm operation", - ); - - let expected_indices = (0..ACCUMULATOR_LEN) - .map(|index| vec![index as u32]) - .collect::>(); - - assert_eq!( - extract_indices, expected_indices, - "each accumulator output must be extracted exactly once and in constraint order", - ); - - Ok(()) -} - -/// The mir-importer encodes `core::sync::atomic::compiler_fence` (issue #781) -/// as an empty, volatile, non-convergent inline-PTX block whose only content -/// is a `~{memory}` clobber. It must lower to a void side-effecting inline asm -/// call with the clobber intact and no instruction text, so no hardware -/// `fence` or `membar` can reach the emitted PTX. -#[test] -fn test_compiler_fence_encoding_lowers_to_empty_sideeffect_asm() -> Result<(), anyhow::Error> { - use llvm_export::types as llvm_types; - use pliron::r#type::Typed; - - let mut ctx = make_test_ctx(); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![]); - - let barrier = nvvm::InlinePtxOp::build(&mut ctx, vec![], vec![], "", "~{memory}", true, false); - barrier.insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{}", e))?; - - let mut found_barrier = false; - for op in lowered_kernel_body(&ctx, module_ptr) { - let Some(inline_asm) = Operation::get_op::(op, &ctx) else { - continue; - }; - assert!( - !found_barrier, - "expected exactly one inline asm op in the lowered kernel" - ); - found_barrier = true; - - let template = inline_asm - .get_attr_inline_asm_template(&ctx) - .map(|s| String::from((*s).clone())); - assert_eq!( - template.as_deref(), - Some(""), - "compiler fence must emit no PTX instruction" - ); - assert_eq!( - inline_asm - .get_attr_inline_asm_constraints(&ctx) - .map(|s| String::from((*s).clone())) - .as_deref(), - Some("~{memory}"), - "compiler fence must keep its memory clobber" - ); - assert!( - llvm::inline_asm_sideeffect(&ctx, inline_asm.get_operation()), - "compiler fence must stay side-effecting so the optimizer cannot drop it" - ); - assert!( - inline_asm - .get_attr_inline_asm_convergent(&ctx) - .is_some_and(|b| !bool::from((*b).clone())), - "compiler fence must not be convergent" - ); - // The zero-result lowering path models the void call as a single - // result of `llvm.void` type, so assert on the type rather than on - // the result count. - let result_ty = inline_asm - .get_operation() - .deref(&ctx) - .get_result(0) - .get_type(&ctx); - assert!( - result_ty - .deref(&ctx) - .downcast_ref::() - .is_some(), - "compiler fence lowers to a void inline asm call" - ); - } - - assert!( - found_barrier, - "expected the compiler-fence inline asm op in the lowered kernel" - ); - Ok(()) -} - -#[test] -fn test_cluster_grid_compatibility_ops_keep_original_lowering() -> Result<(), anyhow::Error> { - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let i32_type = IntegerType::get(&ctx, 32, Signedness::Signless); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![]); - for op_info in [ - nvvm::ReadPtxSregClusterIdxOp::get_concrete_op_info(), - nvvm::ReadPtxSregNclusterIdOp::get_concrete_op_info(), - ] { - Operation::new(&mut ctx, op_info, vec![i32_type.into()], vec![], vec![], 0) - .insert_at_back(entry, &ctx); - } - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let lowered = lowered_kernel_body(&ctx, module_ptr) - .into_iter() - .filter_map(|op| Operation::get_op::(op, &ctx)) - .filter_map(|asm| { - let template = asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone()))?; - (template.contains("%clusterid") || template.contains("%nclusterid")) - .then_some((template, asm)) - }) - .collect::>(); - assert_eq!(lowered.len(), 2); - assert!(lowered.iter().any(|(template, _)| { - template - == "{ .reg .u32 %cx, %cy, %cz, %nx, %ny, %nxy, %xy; mov.u32 %cx, %clusterid.x; mov.u32 %cy, %clusterid.y; mov.u32 %cz, %clusterid.z; mov.u32 %nx, %nclusterid.x; mov.u32 %ny, %nclusterid.y; mul.lo.u32 %nxy, %nx, %ny; mad.lo.u32 %xy, %cy, %nx, %cx; mad.lo.u32 $0, %cz, %nxy, %xy; }" - })); - assert!(lowered.iter().any(|(template, _)| { - template - == "{ .reg .u32 %nx, %ny, %nz, %nxy; mov.u32 %nx, %nclusterid.x; mov.u32 %ny, %nclusterid.y; mov.u32 %nz, %nclusterid.z; mul.lo.u32 %nxy, %nx, %ny; mul.lo.u32 $0, %nxy, %nz; }" - })); - for (_, asm) in lowered { - assert_eq!( - asm.get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some("=r") - ); - assert_eq!(llvm::asm_kind(&ctx, &asm), llvm::AsmKind::Convergent); - } - Ok(()) -} - -/// Regression cover for PR #141: comparisons whose operand is a bool phi. -/// -/// Bools are signless i1, which `can_convert_type` rejects (signless is -/// already the LLVM form), so DialectConversion records no type history for -/// a bool block argument. `is_signed_int_op` used to error out for such -/// operands ("expected IntegerType or MirPtrType operand in arithmetic op"); -/// it must instead fall back to the live operand type and lower the -/// comparison as unsigned. -/// -/// The function mirrors the MIR of a short-circuit kernel: -/// -/// ```text -/// let p = a || b; // bool phi: merge block argument -/// out = (p == q, p < q); // icmp eq i1 / icmp ult i1 -/// ``` -/// -/// ```text -/// bb0(a: i1, b: i1, q: i1): mir.cond_br a, bb2(a), bb1() -/// bb1(): mir.goto bb2(b) -/// bb2(p: i1): mir.eq p, q ; mir.lt p, q ; mir.return -/// ``` -#[test] -fn test_bool_phi_cmp_lowers_to_unsigned_i1_icmp() -> Result<(), anyhow::Error> { - use llvm_export::attributes::ICmpPredicateAttr; - use pliron::basic_block::BasicBlock; - use pliron::builtin::op_interfaces::OperandSegmentInterface; - use pliron::builtin::types::{FunctionType, IntegerType, Signedness}; - use pliron::r#type::Typed; - - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - mir_lower::register(&mut ctx); - - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_ptr = module.get_operation(); - - let bool_ty = IntegerType::get(&ctx, 1, Signedness::Signless); - let arg_tys: Vec = - vec![bool_ty.into(), bool_ty.into(), bool_ty.into()]; - let func_name = "bool_phi_cmp"; - let func_ty = FunctionType::get(&ctx, arg_tys.clone(), vec![]); - - let func_op_ptr = Operation::new( - &mut ctx, - mir::MirFuncOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 1, - ); - let func_ty_attr = pliron::builtin::attributes::TypeAttr::new(func_ty.into()); - let func = mir::MirFuncOp::new(&mut ctx, func_op_ptr, func_ty_attr); - func.set_symbol_name(&mut ctx, func_name.try_into().unwrap()); - - let region = func.get_operation().deref(&ctx).get_region(0); - - // bb0(a, b, q): the function entry. - let bb0 = BasicBlock::new(&mut ctx, None, arg_tys); - bb0.insert_at_back(region, &ctx); - let a = bb0.deref(&ctx).get_argument(0); - let b = bb0.deref(&ctx).get_argument(1); - let q = bb0.deref(&ctx).get_argument(2); - - // bb1(): the short-circuit "evaluate b" block. - let bb1 = BasicBlock::new(&mut ctx, None, vec![]); - bb1.insert_at_back(region, &ctx); - - // bb2(p): the merge block; `p` is the bool phi. - let bb2 = BasicBlock::new(&mut ctx, None, vec![bool_ty.into()]); - bb2.insert_at_back(region, &ctx); - let p = bb2.deref(&ctx).get_argument(0); - - // bb0: cond_br a, bb2(a), bb1(). On the true edge `a` is true, so - // passing `a` itself is `a || b` without needing a constant. - let (flat_operands, segment_sizes) = - mir::MirCondBranchOp::compute_segment_sizes(vec![vec![a], vec![a], vec![]]); - let cond_br = Operation::new( - &mut ctx, - mir::MirCondBranchOp::get_concrete_op_info(), - vec![], - flat_operands, - vec![bb2, bb1], - 0, - ); - Operation::get_op::(cond_br, &ctx) - .expect("MirCondBranchOp") - .set_operand_segment_sizes(&ctx, segment_sizes); - cond_br.insert_at_back(bb0, &ctx); - - // bb1: goto bb2(b). - let goto = Operation::new( - &mut ctx, - mir::MirGotoOp::get_concrete_op_info(), - vec![], - vec![b], - vec![bb2], - 0, - ); - goto.insert_at_back(bb1, &ctx); - - // bb2: p == q, then p < q. - for info in [ - mir::MirEqOp::get_concrete_op_info(), - mir::MirLtOp::get_concrete_op_info(), - ] { - let cmp = Operation::new(&mut ctx, info, vec![bool_ty.into()], vec![p, q], vec![], 0); - cmp.insert_at_back(bb2, &ctx); - } - let ret_op = Operation::new( - &mut ctx, - mir::MirReturnOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 0, - ); - ret_op.insert_at_back(bb2, &ctx); - - let module_region = module.get_operation().deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - func.get_operation().insert_at_back(module_block, &ctx); - - // Before the fallback, this failed with "expected IntegerType or - // MirPtrType operand in arithmetic op". - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{}", e))?; - - let mut icmps = Vec::new(); - let module_op = module_ptr.deref(&ctx); - let region = module_op.get_region(0); - let block = region.deref(&ctx).iter(&ctx).next().unwrap(); - for op in block.deref(&ctx).iter(&ctx) { - let Some(func_op) = Operation::get_op::(op, &ctx) else { - continue; - }; - if func_op.get_symbol_name(&ctx).to_string() != func_name { - continue; - } - let func_region = func_op.get_operation().deref(&ctx).get_region(0); - for func_block in func_region.deref(&ctx).iter(&ctx) { - for body_op in func_block.deref(&ctx).iter(&ctx) { - if let Some(icmp) = Operation::get_op::(body_op, &ctx) { - let lhs_ty = body_op.deref(&ctx).get_operand(0).get_type(&ctx); - icmps.push((icmp.predicate(&ctx), lhs_ty)); - } - } - } - } - - let i1: pliron::r#type::TypeHandle = bool_ty.into(); - assert_eq!( - icmps, - vec![(ICmpPredicateAttr::EQ, i1), (ICmpPredicateAttr::ULT, i1),], - "bool-phi comparisons must lower to `icmp eq i1` and `icmp ult i1`" - ); - Ok(()) -} - -// ============================================================================= -// Scalar math lowering tests -// ============================================================================= - -#[test] -fn test_ex2_approx_f16_uses_exact_pure_i16_inline_ptx_on_both_backends() -> Result<(), anyhow::Error> -{ - use pliron::builtin::types::{IntegerType, Signedness}; - use pliron::r#type::Typed; - - for backend in [ - mir_lower::IntrinsicBackend::LlvmNvptx, - mir_lower::IntrinsicBackend::LibNvvm, - ] { - let mut ctx = make_test_ctx(); - let i16_ty = IntegerType::get(&ctx, 16, Signedness::Unsigned); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![i16_ty.into()]); - let operand = entry.deref(&ctx).get_argument(0); - nvvm::ScalarMathOp::build( - &mut ctx, - operand, - nvvm::ScalarMathFormatAttr::F16, - nvvm::ScalarMathOperationAttr::Ex2, - nvvm::ScalarMathPrecisionAttr::Approx, - nvvm::ScalarMathSubnormalAttr::Preserve, - ) - .insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm_with_options( - &mut ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: backend, - ..Default::default() - }, - ) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let inline_asm = lowered_kernel_body(&ctx, module_ptr) - .into_iter() - .filter_map(|op| Operation::get_op::(op, &ctx)) - .collect::>(); - assert_eq!(inline_asm.len(), 1); - let inline_asm = &inline_asm[0]; - assert_eq!( - inline_asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some("ex2.approx.f16 $0, $1;") - ); - assert_eq!( - inline_asm - .get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some("=h,h") - ); - assert_eq!(llvm::asm_kind(&ctx, inline_asm), llvm::AsmKind::Pure); - assert_eq!( - inline_asm - .get_attr_inline_asm_convergent(&ctx) - .map(|value| bool::from((*value).clone())), - Some(false) - ); - - let op = inline_asm.get_operation().deref(&ctx); - assert_eq!(op.get_num_operands(), 1); - assert_eq!(op.get_num_results(), 1); - for ty in [ - op.get_operand(0).get_type(&ctx), - op.get_result(0).get_type(&ctx), - ] { - let ty = ty.deref(&ctx); - let integer = ty - .downcast_ref::() - .expect("f16 values use an i16 transport type"); - assert_eq!(integer.width(), 16); - } - - let module = Operation::get_op::(module_ptr, &ctx).unwrap(); - let ir = llvm_export::export::export_module_to_string(&ctx, &module) - .expect("ex2.approx.f16 module exports to LLVM IR"); - assert!( - ir.contains("call i16 asm \"ex2.approx.f16 $0, $1;\", \"=h,h\"(i16"), - "{ir}" - ); - assert!(!ir.contains("asm sideeffect"), "{ir}"); - assert!(!ir.contains("~{memory}"), "{ir}"); - } - Ok(()) -} - -#[test] -fn test_ex2_approx_f16_lowering_rejects_wrong_type_and_format() { - use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; - - for backend in [ - mir_lower::IntrinsicBackend::LlvmNvptx, - mir_lower::IntrinsicBackend::LibNvvm, - ] { - let mut wrong_type_ctx = make_test_ctx(); - let f32_ty = FP32Type::get(&wrong_type_ctx); - let (module_ptr, entry) = build_test_kernel(&mut wrong_type_ctx, vec![f32_ty.into()]); - let operand = entry.deref(&wrong_type_ctx).get_argument(0); - nvvm::ScalarMathOp::build( - &mut wrong_type_ctx, - operand, - nvvm::ScalarMathFormatAttr::F16, - nvvm::ScalarMathOperationAttr::Ex2, - nvvm::ScalarMathPrecisionAttr::Approx, - nvvm::ScalarMathSubnormalAttr::Preserve, - ) - .insert_at_back(entry, &wrong_type_ctx); - append_return(&mut wrong_type_ctx, entry); - let error = mir_lower::lower_mir_to_llvm_with_options( - &mut wrong_type_ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: backend, - ..Default::default() - }, - ) - .expect_err("an f16 operation with an f32 operand must fail closed") - .to_string(); - assert!( - error.contains("nvvm.scalar_math types do not match its format"), - "{error}" - ); - - let mut wrong_format_ctx = make_test_ctx(); - let i16_ty = IntegerType::get(&wrong_format_ctx, 16, Signedness::Unsigned); - let (module_ptr, entry) = build_test_kernel(&mut wrong_format_ctx, vec![i16_ty.into()]); - let operand = entry.deref(&wrong_format_ctx).get_argument(0); - nvvm::ScalarMathOp::build( - &mut wrong_format_ctx, - operand, - nvvm::ScalarMathFormatAttr::F16, - nvvm::ScalarMathOperationAttr::Ex2, - nvvm::ScalarMathPrecisionAttr::Approx, - nvvm::ScalarMathSubnormalAttr::Ftz, - ) - .insert_at_back(entry, &wrong_format_ctx); - append_return(&mut wrong_format_ctx, entry); - let error = mir_lower::lower_mir_to_llvm_with_options( - &mut wrong_format_ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: backend, - ..Default::default() - }, - ) - .expect_err("unadmitted ex2.approx.ftz.f16 must fail closed") - .to_string(); - assert!( - error.contains("nvvm.scalar_math variant is not admitted"), - "{error}" - ); - } -} - -// ============================================================================= -// Integer dot product (dp4a / dp2a) lowering tests -// ============================================================================= - -const DOT_PRODUCT_TYPED_INTRINSICS: [&str; 4] = [ - "llvm_nvvm_idp4a_s_s", - "llvm_nvvm_idp4a_u_u", - "llvm_nvvm_idp2a_s_s", - "llvm_nvvm_idp2a_u_u", -]; - -const DOT_PRODUCT_PTX: [&str; 4] = [ - "dp4a.s32.s32 $0, $1, $2, $3;", - "dp4a.u32.u32 $0, $1, $2, $3;", - "dp2a.lo.s32.s32 $0, $1, $2, $3;", - "dp2a.lo.u32.u32 $0, $1, $2, $3;", -]; - -fn lower_all_dot_product_forms( - backend: mir_lower::IntrinsicBackend, -) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let (module_ptr, entry) = - build_test_kernel(&mut ctx, vec![i32_ty.into(), i32_ty.into(), i32_ty.into()]); - let operands = (0..3) - .map(|index| entry.deref(&ctx).get_argument(index)) - .collect::>(); - for op_info in [ - nvvm::Dp4aS32Op::get_concrete_op_info(), - nvvm::Dp4aU32Op::get_concrete_op_info(), - nvvm::Dp2aS32Op::get_concrete_op_info(), - nvvm::Dp2aU32Op::get_concrete_op_info(), - ] { - Operation::new( - &mut ctx, - op_info, - vec![i32_ty.into()], - operands.clone(), - vec![], - 0, - ) - .insert_at_back(entry, &ctx); - } - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm_with_options( - &mut ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: backend, - ..Default::default() - }, - ) - .map_err(|error| anyhow::anyhow!("{error}"))?; - Ok((ctx, module_ptr)) -} - -#[test] -fn test_dot_product_llvm_nvptx_uses_typed_intrinsics_and_low_selector() -> Result<(), anyhow::Error> -{ - use pliron::builtin::attributes::IntegerAttr; - - let (ctx, module_ptr) = lower_all_dot_product_forms(mir_lower::IntrinsicBackend::LlvmNvptx)?; - let body = lowered_kernel_body(&ctx, module_ptr); - let mut calls = Vec::new(); - for op in body { - assert!( - Operation::get_op::(op, &ctx).is_none(), - "LLVM-NVPTX dot products must use typed intrinsics" - ); - let Some(call) = Operation::get_op::(op, &ctx) else { - continue; - }; - let CallOpCallable::Direct(callee) = call.callee(&ctx) else { - continue; - }; - let callee = callee.to_string(); - if !callee.starts_with("llvm_nvvm_idp") { - continue; - } - let expected_arity = if callee.contains("idp2a") { 4 } else { 3 }; - assert_eq!(op.deref(&ctx).get_num_operands(), expected_arity); - if expected_arity == 4 { - let selector = op.deref(&ctx).get_operand(2); - let defining_op = selector.defining_op().expect("selector is a constant"); - let constant = Operation::get_op::(defining_op, &ctx) - .expect("selector is an LLVM constant"); - let attribute = constant.get_value(&ctx); - let integer = attribute - .downcast_ref::() - .expect("selector constant is an integer"); - assert_eq!(integer.value().bw(), 1); - assert_eq!(integer.value().to_u64(), 0, "dp2a must select `.lo`"); - } - calls.push(callee); - } - calls.sort(); - let mut expected = DOT_PRODUCT_TYPED_INTRINSICS.map(str::to_owned); - expected.sort(); - assert_eq!(calls, expected); - Ok(()) -} - -#[test] -fn test_dot_product_libnvvm_uses_exact_pure_inline_ptx() -> Result<(), anyhow::Error> { - let (ctx, module_ptr) = lower_all_dot_product_forms(mir_lower::IntrinsicBackend::LibNvvm)?; - let body = lowered_kernel_body(&ctx, module_ptr); - let mut inline_ptx = Vec::new(); - for op in body { - assert!( - Operation::get_op::(op, &ctx).is_none(), - "libNVVM dot products must not use typed intrinsic calls" - ); - let Some(inline_asm) = Operation::get_op::(op, &ctx) else { - continue; - }; - inline_ptx.push( - inline_asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .unwrap_or_default(), - ); - assert_eq!( - inline_asm - .get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some("=r,r,r,r") - ); - assert_eq!(llvm::asm_kind(&ctx, &inline_asm), llvm::AsmKind::Pure); - assert_eq!(op.deref(&ctx).get_num_operands(), 3); - assert_eq!(op.deref(&ctx).get_num_results(), 1); - } - inline_ptx.sort(); - let mut expected = DOT_PRODUCT_PTX.map(str::to_owned); - expected.sort(); - assert_eq!(inline_ptx, expected); - Ok(()) -} - -// ============================================================================= -// Byte permutation lowering tests -// ============================================================================= - -const PRMT_TYPED_INTRINSICS: [(&str, usize); 7] = [ - ("llvm_nvvm_prmt", 3), - ("llvm_nvvm_prmt_f4e", 3), - ("llvm_nvvm_prmt_b4e", 3), - ("llvm_nvvm_prmt_rc8", 2), - ("llvm_nvvm_prmt_ecl", 2), - ("llvm_nvvm_prmt_ecr", 2), - ("llvm_nvvm_prmt_rc16", 2), -]; - -const PRMT_INLINE_PTX: [(&str, &str, usize); 7] = [ - ("prmt.b32 $0, $1, $2, $3;", "=r,r,r,r", 3), - ("prmt.b32.f4e $0, $1, $2, $3;", "=r,r,r,r", 3), - ("prmt.b32.b4e $0, $1, $2, $3;", "=r,r,r,r", 3), - ("prmt.b32.rc8 $0, $1, 0, $2;", "=r,r,r", 2), - ("prmt.b32.ecl $0, $1, 0, $2;", "=r,r,r", 2), - ("prmt.b32.ecr $0, $1, 0, $2;", "=r,r,r", 2), - ("prmt.b32.rc16 $0, $1, 0, $2;", "=r,r,r", 2), -]; - -fn lower_all_prmt_modes( - backend: mir_lower::IntrinsicBackend, -) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let (module_ptr, entry) = - build_test_kernel(&mut ctx, vec![i32_ty.into(), i32_ty.into(), i32_ty.into()]); - let operands = (0..3) - .map(|index| entry.deref(&ctx).get_argument(index)) - .collect::>(); - - for mode in [ - nvvm::PrmtModeAttr::Generic, - nvvm::PrmtModeAttr::F4e, - nvvm::PrmtModeAttr::B4e, - ] { - nvvm::PrmtOp::build(&mut ctx, operands.clone(), mode).insert_at_back(entry, &ctx); - } - for mode in [ - nvvm::PrmtModeAttr::Rc8, - nvvm::PrmtModeAttr::Ecl, - nvvm::PrmtModeAttr::Ecr, - nvvm::PrmtModeAttr::Rc16, - ] { - nvvm::PrmtOp::build(&mut ctx, vec![operands[0], operands[2]], mode) - .insert_at_back(entry, &ctx); - } - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm_with_options( - &mut ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: backend, - ..Default::default() - }, - ) - .map_err(|error| anyhow::anyhow!("{error}"))?; - Ok((ctx, module_ptr)) -} - -#[test] -fn test_prmt_llvm_nvptx_uses_exact_typed_intrinsics() -> Result<(), anyhow::Error> { - let (ctx, module_ptr) = lower_all_prmt_modes(mir_lower::IntrinsicBackend::LlvmNvptx)?; - let body = lowered_kernel_body(&ctx, module_ptr); - let mut calls = Vec::new(); - for op in body { - assert!( - Operation::get_op::(op, &ctx).is_none(), - "LLVM-NVPTX byte permutations must use typed intrinsics" - ); - let Some(call) = Operation::get_op::(op, &ctx) else { - continue; - }; - let CallOpCallable::Direct(callee) = call.callee(&ctx) else { - continue; - }; - let callee = callee.to_string(); - let Some((_, expected_arity)) = PRMT_TYPED_INTRINSICS - .iter() - .find(|(expected, _)| *expected == callee) - else { - continue; - }; - assert_eq!(op.deref(&ctx).get_num_operands(), *expected_arity); - calls.push((callee, *expected_arity)); - } - calls.sort(); - let mut expected = PRMT_TYPED_INTRINSICS.map(|(name, arity)| (name.to_owned(), arity)); - expected.sort(); - assert_eq!(calls, expected); - Ok(()) -} - -#[test] -fn test_prmt_libnvvm_uses_exact_pure_inline_ptx() -> Result<(), anyhow::Error> { - let (ctx, module_ptr) = lower_all_prmt_modes(mir_lower::IntrinsicBackend::LibNvvm)?; - let body = lowered_kernel_body(&ctx, module_ptr); - let mut inline_ptx = Vec::new(); - for op in body { - if let Some(call) = Operation::get_op::(op, &ctx) - && let CallOpCallable::Direct(callee) = call.callee(&ctx) - { - assert!( - !callee.to_string().starts_with("llvm_nvvm_prmt"), - "libNVVM byte permutations must not use typed intrinsic calls" - ); - } - let Some(inline_asm) = Operation::get_op::(op, &ctx) else { - continue; - }; - let template = inline_asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .unwrap_or_default(); - let Some((_, expected_constraints, expected_arity)) = PRMT_INLINE_PTX - .iter() - .find(|(expected, _, _)| *expected == template) - else { - continue; - }; - assert_eq!( - inline_asm - .get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some(*expected_constraints) - ); - assert_eq!(llvm::asm_kind(&ctx, &inline_asm), llvm::AsmKind::Pure); - assert_eq!(op.deref(&ctx).get_num_operands(), *expected_arity); - assert_eq!(op.deref(&ctx).get_num_results(), 1); - inline_ptx.push((template, *expected_constraints, *expected_arity)); - } - inline_ptx.sort(); - let mut expected = PRMT_INLINE_PTX - .map(|(template, constraints, arity)| (template.to_owned(), constraints, arity)); - expected.sort(); - assert_eq!(inline_ptx, expected); - Ok(()) -} - -fn lower_sync_threads( - backend: mir_lower::IntrinsicBackend, -) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { - let mut ctx = make_test_ctx(); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![]); - Operation::new( - &mut ctx, - nvvm::Barrier0Op::get_concrete_op_info(), - vec![], - vec![], - vec![], - 0, - ) - .insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm_with_options( - &mut ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: backend, - ..Default::default() - }, - ) - .map_err(|error| anyhow::anyhow!("{error}"))?; - Ok((ctx, module_ptr)) -} - -#[test] -fn test_sync_threads_llvm_nvptx_uses_typed_intrinsic_with_fixed_zero() -> Result<(), anyhow::Error> -{ - use pliron::builtin::attributes::IntegerAttr; - - let (ctx, module_ptr) = lower_sync_threads(mir_lower::IntrinsicBackend::LlvmNvptx)?; - let body = lowered_kernel_body(&ctx, module_ptr); - let mut found = false; - for op in body { - assert!( - Operation::get_op::(op, &ctx).is_none(), - "LLVM-NVPTX sync_threads must use the typed intrinsic" - ); - let Some(call) = Operation::get_op::(op, &ctx) else { - continue; - }; - let CallOpCallable::Direct(callee) = call.callee(&ctx) else { - continue; - }; - if callee.to_string() != "llvm_nvvm_barrier_cta_sync_aligned_all" { - continue; - } - let call = op.deref(&ctx); - assert_eq!(call.get_num_operands(), 1); - let barrier_id = call.get_operand(0); - let defining_op = barrier_id.defining_op().expect("barrier ID is constant"); - let constant = Operation::get_op::(defining_op, &ctx) - .expect("barrier ID is an LLVM constant"); - let value = constant.get_value(&ctx); - let integer = value - .downcast_ref::() - .expect("barrier ID is an integer"); - assert_eq!(integer.value().bw(), 32); - assert_eq!(integer.value().to_u64(), 0); - found = true; - } - assert!(found, "modern typed CTA barrier call was not emitted"); - - let module = Operation::get_op::(module_ptr, &ctx).unwrap(); - let ir = llvm_export::export::export_module_to_string(&ctx, &module) - .map_err(|error| anyhow::anyhow!(error))?; - assert!(ir.contains("@llvm.nvvm.barrier.cta.sync.aligned.all(i32 0)")); - assert!(!ir.contains("@llvm.nvvm.barrier0")); - Ok(()) -} - -#[test] -fn test_sync_threads_libnvvm_uses_exact_convergent_inline_ptx() -> Result<(), anyhow::Error> { - let (ctx, module_ptr) = lower_sync_threads(mir_lower::IntrinsicBackend::LibNvvm)?; - let body = lowered_kernel_body(&ctx, module_ptr); - let mut found = false; - for op in body { - let Some(inline_asm) = Operation::get_op::(op, &ctx) else { - if let Some(call) = Operation::get_op::(op, &ctx) - && let CallOpCallable::Direct(callee) = call.callee(&ctx) - { - assert_ne!(callee.to_string(), "llvm_nvvm_barrier_cta_sync_aligned_all"); - } - continue; - }; - assert_eq!( - inline_asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some("bar.sync 0;") - ); - assert_eq!( - inline_asm - .get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some("~{memory}") - ); - assert_eq!(llvm::asm_kind(&ctx, &inline_asm), llvm::AsmKind::Convergent); - assert_eq!(op.deref(&ctx).get_num_operands(), 0); - found = true; - } - assert!(found, "exact libNVVM barrier inline PTX was not emitted"); - - let module = Operation::get_op::(module_ptr, &ctx).unwrap(); - let ir = llvm_export::export::export_module_to_string(&ctx, &module) - .map_err(|error| anyhow::anyhow!(error))?; - assert!( - ir.contains("call void asm sideeffect \"bar.sync 0;\", \"~{memory}\"() #0"), - "{ir}" - ); - assert!(ir.contains("attributes #0 = { convergent }"), "{ir}"); - assert!(!ir.contains("@llvm.nvvm.barrier.cta.sync.aligned.all")); - Ok(()) -} - -fn lower_warp_barrier( - backend: mir_lower::IntrinsicBackend, -) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![i32_ty.into()]); - let member_mask = entry.deref(&ctx).get_argument(0); - nvvm::BarWarpSyncOp::build(&mut ctx, member_mask).insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm_with_options( - &mut ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: backend, - ..Default::default() - }, - ) - .map_err(|error| anyhow::anyhow!("{error}"))?; - Ok((ctx, module_ptr)) -} - -#[test] -fn test_warp_barrier_uses_typed_intrinsic_on_both_backends() -> Result<(), anyhow::Error> { - for backend in [ - mir_lower::IntrinsicBackend::LlvmNvptx, - mir_lower::IntrinsicBackend::LibNvvm, - ] { - let (ctx, module_ptr) = lower_warp_barrier(backend)?; - let body = lowered_kernel_body(&ctx, module_ptr); - let mut calls = 0; - for op in body { - assert!( - Operation::get_op::(op, &ctx).is_none(), - "warp barrier must use its typed intrinsic" - ); - let Some(call) = Operation::get_op::(op, &ctx) else { - continue; - }; - let CallOpCallable::Direct(callee) = call.callee(&ctx) else { - continue; - }; - if callee.to_string() == "llvm_nvvm_bar_warp_sync" { - assert_eq!(op.deref(&ctx).get_num_operands(), 1); - assert_eq!(op.deref(&ctx).get_num_results(), 1); - calls += 1; - } - } - assert_eq!(calls, 1, "expected one typed warp-barrier call"); - - let module = Operation::get_op::(module_ptr, &ctx).unwrap(); - let ir = llvm_export::export::export_module_to_string(&ctx, &module) - .map_err(|error| anyhow::anyhow!(error))?; - assert!(ir.contains("@llvm.nvvm.bar.warp.sync(i32"), "{ir}"); - } - Ok(()) -} - -// ============================================================================= -// cp.async lowering tests -// ============================================================================= - -fn lower_all_classic_cp_async( - backend: mir_lower::IntrinsicBackend, -) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { - use dialect_mir::types::MirPtrType; - use pliron::builtin::attributes::IntegerAttr; - use pliron::builtin::types::{IntegerType, Signedness}; - use pliron::utils::apint::APInt; - use std::num::NonZeroUsize; - - let mut ctx = make_test_ctx(); - let u8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let dst_ty = MirPtrType::get_generic(&mut ctx, u32_ty.into(), true); - let src32_ty = MirPtrType::get_generic(&mut ctx, u32_ty.into(), false); - let src8_ty = MirPtrType::get_generic(&mut ctx, u8_ty.into(), false); - let (module_ptr, entry) = build_test_kernel( - &mut ctx, - vec![ - dst_ty.into(), - src32_ty.into(), - src8_ty.into(), - u32_ty.into(), - ], - ); - let dst = entry.deref(&ctx).get_argument(0); - let src32 = entry.deref(&ctx).get_argument(1); - let src8 = entry.deref(&ctx).get_argument(2); - let source_size = entry.deref(&ctx).get_argument(3); - - let zero_op = Operation::new( - &mut ctx, - mir::MirConstantOp::get_concrete_op_info(), - vec![u32_ty.into()], - vec![], - vec![], - 0, - ); - mir::MirConstantOp::new(zero_op).set_attr_value( - &ctx, - IntegerAttr::new(u32_ty, APInt::from_u32(0, NonZeroUsize::new(32).unwrap())), - ); - zero_op.insert_at_back(entry, &ctx); - let zero = zero_op.deref(&ctx).get_result(0); - - for copy in [ - nvvm::CpAsyncCa4Op::build(&mut ctx, dst, src32), - nvvm::CpAsyncCa8Op::build(&mut ctx, dst, src32), - nvvm::CpAsyncCa16Op::build(&mut ctx, dst, src32), - nvvm::CpAsyncCaZfill4Op::build(&mut ctx, dst, src8, source_size), - nvvm::CpAsyncCaZfill8Op::build(&mut ctx, dst, src8, source_size), - nvvm::CpAsyncCaZfill16Op::build(&mut ctx, dst, src8, source_size), - nvvm::CpAsyncCg16Op::build(&mut ctx, dst, src32), - nvvm::CpAsyncCgZfill16Op::build(&mut ctx, dst, src8, source_size), - ] { - copy.insert_at_back(entry, &ctx); - } - for control in [ - nvvm::CpAsyncCommitGroupOp::build(&mut ctx), - nvvm::CpAsyncWaitGroupOp::build(&mut ctx, zero), - nvvm::CpAsyncWaitAllOp::build(&mut ctx), - ] { - control.insert_at_back(entry, &ctx); - } - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm_with_options( - &mut ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: backend, - ..Default::default() - }, - ) - .map_err(|error| anyhow::anyhow!("{error}"))?; - Ok((ctx, module_ptr)) -} - -fn lower_all_cp_async_mbarrier( - backend: mir_lower::IntrinsicBackend, -) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { - use dialect_mir::types::MirPtrType; - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let generic_ty = MirPtrType::get_generic(&mut ctx, u64_ty.into(), true); - let shared_ty = MirPtrType::get_shared(&mut ctx, u64_ty.into(), true); - let (module_ptr, entry) = - build_test_kernel(&mut ctx, vec![generic_ty.into(), shared_ty.into()]); - let generic = entry.deref(&ctx).get_argument(0); - let shared = entry.deref(&ctx).get_argument(1); - - for bridge in [ - nvvm::CpAsyncMbarrierArriveOp::build(&mut ctx, shared), - nvvm::CpAsyncMbarrierArriveSharedOp::build(&mut ctx, generic), - nvvm::CpAsyncMbarrierArriveNoIncOp::build(&mut ctx, shared), - nvvm::CpAsyncMbarrierArriveNoIncSharedOp::build(&mut ctx, generic), - ] { - bridge.insert_at_back(entry, &ctx); - } - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm_with_options( - &mut ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: backend, - ..Default::default() - }, - ) - .map_err(|error| anyhow::anyhow!("{error}"))?; - Ok((ctx, module_ptr)) -} - -#[test] -fn test_generated_cp_async_mbarrier_preserves_backend_and_address_routes() --> Result<(), anyhow::Error> { - use llvm_export::types::PointerType; - use pliron::r#type::Typed; - - let typed = [ - ("llvm_nvvm_cp_async_mbarrier_arrive", 0), - ("llvm_nvvm_cp_async_mbarrier_arrive_shared", 3), - ("llvm_nvvm_cp_async_mbarrier_arrive_noinc", 0), - ("llvm_nvvm_cp_async_mbarrier_arrive_noinc_shared", 3), - ]; - let templates = [ - ("cp.async.mbarrier.arrive.b64 [$0];", 0), - ("cp.async.mbarrier.arrive.shared.b64 [$0];", 3), - ("cp.async.mbarrier.arrive.noinc.b64 [$0];", 0), - ("cp.async.mbarrier.arrive.noinc.shared.b64 [$0];", 3), - ]; - - for backend in [ - mir_lower::IntrinsicBackend::LlvmNvptx, - mir_lower::IntrinsicBackend::LibNvvm, - ] { - let (ctx, module_ptr) = lower_all_cp_async_mbarrier(backend)?; - let mut call_counts = [0usize; 4]; - let mut asm_counts = [0usize; 4]; - - for op in lowered_kernel_body(&ctx, module_ptr) { - if let Some(call) = Operation::get_op::(op, &ctx) { - let CallOpCallable::Direct(callee) = call.callee(&ctx) else { - continue; - }; - let callee = callee.to_string(); - let Some(index) = typed.iter().position(|(name, _)| callee == *name) else { - continue; - }; - call_counts[index] += 1; - let pointer = op.deref(&ctx).get_operand(0).get_type(&ctx); - assert_eq!( - pointer - .deref(&ctx) - .downcast_ref::() - .unwrap() - .address_space(), - typed[index].1 - ); - } - if let Some(inline_asm) = Operation::get_op::(op, &ctx) { - let template = inline_asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .unwrap(); - let Some(index) = templates - .iter() - .position(|(expected, _)| template == *expected) - else { - continue; - }; - asm_counts[index] += 1; - assert_eq!(llvm::asm_kind(&ctx, &inline_asm), llvm::AsmKind::Convergent); - assert!( - inline_asm - .get_attr_inline_asm_convergent(&ctx) - .is_some_and(|value| bool::from((*value).clone())) - ); - assert_eq!( - inline_asm - .get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some("l,~{memory}") - ); - let pointer = op.deref(&ctx).get_operand(0).get_type(&ctx); - assert_eq!( - pointer - .deref(&ctx) - .downcast_ref::() - .unwrap() - .address_space(), - templates[index].1 - ); - } - } - - let module = Operation::get_op::(module_ptr, &ctx).unwrap(); - let ir = llvm_export::export::export_module_to_string(&ctx, &module) - .map_err(|error| anyhow::anyhow!(error))?; - match backend { - mir_lower::IntrinsicBackend::LlvmNvptx => { - assert_eq!(call_counts, [1; 4]); - assert_eq!(asm_counts, [0; 4]); - assert!( - ir.contains("@llvm.nvvm.cp.async.mbarrier.arrive(ptr"), - "{ir}" - ); - assert!( - ir.contains("@llvm.nvvm.cp.async.mbarrier.arrive.shared(ptr addrspace(3)"), - "{ir}" - ); - assert!( - ir.contains("@llvm.nvvm.cp.async.mbarrier.arrive.noinc(ptr"), - "{ir}" - ); - assert!( - ir.contains( - "@llvm.nvvm.cp.async.mbarrier.arrive.noinc.shared(ptr addrspace(3)" - ), - "{ir}" - ); - assert!(!ir.contains("cp.async.mbarrier.arrive.b64 [$0]"), "{ir}"); - } - mir_lower::IntrinsicBackend::LibNvvm => { - assert_eq!(call_counts, [0; 4]); - assert_eq!(asm_counts, [1; 4]); - assert!(!ir.contains("@llvm.nvvm.cp.async.mbarrier"), "{ir}"); - for (template, _) in templates { - assert!(ir.contains(template), "{ir}"); - } - assert!(ir.contains("asm sideeffect"), "{ir}"); - assert!(ir.contains("convergent"), "{ir}"); - } - } - } - Ok(()) -} - -#[test] -fn test_generated_cp_async_llvm_nvptx_uses_all_typed_intrinsics() -> Result<(), anyhow::Error> { - use llvm_export::types::PointerType; - use pliron::r#type::Typed; - - let (ctx, module_ptr) = lower_all_classic_cp_async(mir_lower::IntrinsicBackend::LlvmNvptx)?; - let mut found = Vec::new(); - for op in lowered_kernel_body(&ctx, module_ptr) { - assert!( - Operation::get_op::(op, &ctx).is_none(), - "LLVM-NVPTX cp.async must use typed intrinsics" - ); - let Some(call) = Operation::get_op::(op, &ctx) else { - continue; - }; - let CallOpCallable::Direct(callee) = call.callee(&ctx) else { - continue; - }; - let callee = callee.to_string(); - if !callee.starts_with("llvm_nvvm_cp_async") { - continue; - } - if callee.contains("shared_global") { - let call = op.deref(&ctx); - let destination_ty = call.get_operand(0).get_type(&ctx); - let source_ty = call.get_operand(1).get_type(&ctx); - assert_eq!( - destination_ty - .deref(&ctx) - .downcast_ref::() - .unwrap() - .address_space(), - 3 - ); - assert_eq!( - source_ty - .deref(&ctx) - .downcast_ref::() - .unwrap() - .address_space(), - 1 - ); - } - found.push(callee); - } - found.sort(); - let mut expected = [ - "llvm_nvvm_cp_async_ca_shared_global_4", - "llvm_nvvm_cp_async_ca_shared_global_4_s", - "llvm_nvvm_cp_async_ca_shared_global_8", - "llvm_nvvm_cp_async_ca_shared_global_8_s", - "llvm_nvvm_cp_async_ca_shared_global_16", - "llvm_nvvm_cp_async_ca_shared_global_16_s", - "llvm_nvvm_cp_async_cg_shared_global_16", - "llvm_nvvm_cp_async_cg_shared_global_16_s", - "llvm_nvvm_cp_async_commit_group", - "llvm_nvvm_cp_async_wait_all", - "llvm_nvvm_cp_async_wait_group", - ] - .map(str::to_owned); - expected.sort(); - assert_eq!(found, expected); - Ok(()) -} - -#[test] -fn test_generated_cp_async_libnvvm_uses_all_exact_inline_ptx() -> Result<(), anyhow::Error> { - let (ctx, module_ptr) = lower_all_classic_cp_async(mir_lower::IntrinsicBackend::LibNvvm)?; - let mut lowered = Vec::new(); - for op in lowered_kernel_body(&ctx, module_ptr) { - let Some(asm) = Operation::get_op::(op, &ctx) else { - continue; - }; - assert_eq!(llvm::asm_kind(&ctx, &asm), llvm::AsmKind::SideEffect); - assert!( - asm.get_attr_inline_asm_convergent(&ctx) - .is_some_and(|value| !bool::from((*value).clone())) - ); - lowered.push(( - asm.get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .unwrap(), - asm.get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .unwrap(), - )); - } - - let expected = [ - ( - "cp.async.ca.shared.global [%smem32], [%gmem64], 4;", - "l,l,~{memory}", - ), - ( - "cp.async.ca.shared.global [%smem32], [%gmem64], 8;", - "l,l,~{memory}", - ), - ( - "cp.async.ca.shared.global [%smem32], [%gmem64], 16;", - "l,l,~{memory}", - ), - ( - "cp.async.ca.shared.global [%smem32], [%gmem64], 4, $2;", - "l,l,r,~{memory}", - ), - ( - "cp.async.ca.shared.global [%smem32], [%gmem64], 8, $2;", - "l,l,r,~{memory}", - ), - ( - "cp.async.ca.shared.global [%smem32], [%gmem64], 16, $2;", - "l,l,r,~{memory}", - ), - ( - "cp.async.cg.shared.global [%smem32], [%gmem64], 16;", - "l,l,~{memory}", - ), - ( - "cp.async.cg.shared.global [%smem32], [%gmem64], 16, $2;", - "l,l,r,~{memory}", - ), - ("cp.async.commit_group;", "~{memory}"), - ("cp.async.wait_group $0;", "n,~{memory}"), - ("cp.async.wait_all;", "~{memory}"), - ]; - for (instruction, constraints) in expected { - assert_eq!( - lowered - .iter() - .filter(|(template, actual_constraints)| { - template.contains(instruction) && actual_constraints == constraints - }) - .count(), - 1, - "missing exact `{instruction}`" - ); - } - assert_eq!(lowered.len(), expected.len()); - Ok(()) -} - -#[test] -fn test_cp_async_ca_4_lowers_to_inline_asm() -> Result<(), anyhow::Error> { - use dialect_mir::types::MirPtrType; - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let dst_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), true); - let src_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), false); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![dst_ty.into(), src_ty.into()]); - - let dst = entry.deref(&ctx).get_argument(0); - let src = entry.deref(&ctx).get_argument(1); - - let op = Operation::new( - &mut ctx, - nvvm::CpAsyncCa4Op::get_concrete_op_info(), - vec![], - vec![dst, src], - vec![], - 0, - ); - op.insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - assert_cp_async_inline_asm_lowering(&mut ctx, module_ptr, 4) -} - -#[test] -fn test_cp_async_ca_8_lowers_to_inline_asm() -> Result<(), anyhow::Error> { - use dialect_mir::types::MirPtrType; - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let dst_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), true); - let src_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), false); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![dst_ty.into(), src_ty.into()]); - - let dst = entry.deref(&ctx).get_argument(0); - let src = entry.deref(&ctx).get_argument(1); - - let op = Operation::new( - &mut ctx, - nvvm::CpAsyncCa8Op::get_concrete_op_info(), - vec![], - vec![dst, src], - vec![], - 0, - ); - op.insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - assert_cp_async_inline_asm_lowering(&mut ctx, module_ptr, 8) -} - -fn assert_cp_async_inline_asm_lowering( - ctx: &mut Context, - module_ptr: pliron::context::Ptr, - copy_size: u32, -) -> Result<(), anyhow::Error> { - use pliron::r#type::Typed; - - mir_lower::lower_mir_to_llvm_with_options( - ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: mir_lower::IntrinsicBackend::LibNvvm, - ..Default::default() - }, - ) - .map_err(|e| anyhow::anyhow!("{e}"))?; - - let expected_template = format!( - "{{ .reg .u64 %smem64; .reg .u32 %smem32; .reg .u64 %gmem64; \ - cvta.to.shared.u64 %smem64, $0; cvt.u32.u64 %smem32, %smem64; \ - cvta.to.global.u64 %gmem64, $1; \ - cp.async.ca.shared.global [%smem32], [%gmem64], {copy_size}; }}" - ); - let mut matches = 0; - let module_region = module_ptr.deref(ctx).get_region(0); - let module_block = module_region.deref(ctx).iter(ctx).next().unwrap(); - - for op in module_block.deref(ctx).iter(ctx) { - let Some(func_op) = Operation::get_op::(op, ctx) else { - continue; - }; - if func_op.get_symbol_name(ctx).to_string() != "kernel_func" { - continue; - } - - let func_region = func_op.get_operation().deref(ctx).get_region(0); - for func_block in func_region.deref(ctx).iter(ctx) { - for body_op in func_block.deref(ctx).iter(ctx) { - let Some(inline_asm) = Operation::get_op::(body_op, ctx) else { - continue; - }; - let template = inline_asm - .get_attr_inline_asm_template(ctx) - .map(|s| String::from((*s).clone())); - if template.as_deref() != Some(expected_template.as_str()) { - continue; - } - - matches += 1; - assert_eq!( - inline_asm - .get_attr_inline_asm_constraints(ctx) - .map(|s| String::from((*s).clone())) - .as_deref(), - Some("l,l,~{memory}") - ); - assert_eq!(llvm::asm_kind(ctx, &inline_asm), llvm::AsmKind::SideEffect); - assert!( - inline_asm - .get_attr_inline_asm_convergent(ctx) - .is_some_and(|value| !bool::from((*value).clone())) - ); - - let operands: Vec<_> = inline_asm.get_operation().deref(ctx).operands().collect(); - assert_eq!(operands.len(), 2); - for operand in operands { - let ty = operand.get_type(ctx); - let ty = ty.deref(ctx); - let ptr_ty = ty - .downcast_ref::() - .expect("cp.async operands must lower to LLVM pointers"); - assert_eq!(ptr_ty.address_space(), 0); - } - } - } - } - - assert_eq!(matches, 1, "missing exact {copy_size}-byte cp.async asm"); - Ok(()) -} - -// ============================================================================= -// cp.async zero-fill lowering tests -// ============================================================================= - -#[test] -fn test_cp_async_ca_zfill_4_lowers_to_inline_asm() -> Result<(), anyhow::Error> { - use dialect_mir::types::MirPtrType; - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signless); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let dst_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), true); - let src_ty = MirPtrType::get_generic(&mut ctx, i8_ty.into(), false); - let (module_ptr, entry) = - build_test_kernel(&mut ctx, vec![dst_ty.into(), src_ty.into(), i32_ty.into()]); - - let dst = entry.deref(&ctx).get_argument(0); - let src = entry.deref(&ctx).get_argument(1); - let src_size = entry.deref(&ctx).get_argument(2); - - let op = Operation::new( - &mut ctx, - nvvm::CpAsyncCaZfill4Op::get_concrete_op_info(), - vec![], - vec![dst, src, src_size], - vec![], - 0, - ); - op.insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - assert_cp_async_zfill_inline_asm_lowering(&mut ctx, module_ptr, 4) -} - -#[test] -fn test_cp_async_ca_zfill_8_lowers_to_inline_asm() -> Result<(), anyhow::Error> { - use dialect_mir::types::MirPtrType; - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signless); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let dst_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), true); - let src_ty = MirPtrType::get_generic(&mut ctx, i8_ty.into(), false); - let (module_ptr, entry) = - build_test_kernel(&mut ctx, vec![dst_ty.into(), src_ty.into(), i32_ty.into()]); - - let dst = entry.deref(&ctx).get_argument(0); - let src = entry.deref(&ctx).get_argument(1); - let src_size = entry.deref(&ctx).get_argument(2); - - let op = Operation::new( - &mut ctx, - nvvm::CpAsyncCaZfill8Op::get_concrete_op_info(), - vec![], - vec![dst, src, src_size], - vec![], - 0, - ); - op.insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - assert_cp_async_zfill_inline_asm_lowering(&mut ctx, module_ptr, 8) -} - -#[test] -fn test_cp_async_ca_zfill_16_lowers_to_inline_asm() -> Result<(), anyhow::Error> { - use dialect_mir::types::MirPtrType; - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signless); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let dst_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), true); - let src_ty = MirPtrType::get_generic(&mut ctx, i8_ty.into(), false); - let (module_ptr, entry) = - build_test_kernel(&mut ctx, vec![dst_ty.into(), src_ty.into(), i32_ty.into()]); - - let dst = entry.deref(&ctx).get_argument(0); - let src = entry.deref(&ctx).get_argument(1); - let src_size = entry.deref(&ctx).get_argument(2); - - let op = Operation::new( - &mut ctx, - nvvm::CpAsyncCaZfill16Op::get_concrete_op_info(), - vec![], - vec![dst, src, src_size], - vec![], - 0, - ); - op.insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - assert_cp_async_zfill_inline_asm_lowering(&mut ctx, module_ptr, 16) -} - -fn assert_cp_async_zfill_inline_asm_lowering( - ctx: &mut Context, - module_ptr: pliron::context::Ptr, - copy_size: u32, -) -> Result<(), anyhow::Error> { - use pliron::builtin::types::IntegerType; - use pliron::r#type::Typed; - - mir_lower::lower_mir_to_llvm_with_options( - ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: mir_lower::IntrinsicBackend::LibNvvm, - ..Default::default() - }, - ) - .map_err(|e| anyhow::anyhow!("{e}"))?; - - let expected_template = format!( - "{{ .reg .u64 %smem64; .reg .u32 %smem32; .reg .u64 %gmem64; \ - cvta.to.shared.u64 %smem64, $0; cvt.u32.u64 %smem32, %smem64; \ - cvta.to.global.u64 %gmem64, $1; \ - cp.async.ca.shared.global [%smem32], [%gmem64], {copy_size}, $2; }}" - ); - let mut matches = 0; - let module_region = module_ptr.deref(ctx).get_region(0); - let module_block = module_region.deref(ctx).iter(ctx).next().unwrap(); - - for op in module_block.deref(ctx).iter(ctx) { - let Some(func_op) = Operation::get_op::(op, ctx) else { - continue; - }; - if func_op.get_symbol_name(ctx).to_string() != "kernel_func" { - continue; - } - - let func_region = func_op.get_operation().deref(ctx).get_region(0); - for func_block in func_region.deref(ctx).iter(ctx) { - for body_op in func_block.deref(ctx).iter(ctx) { - let Some(inline_asm) = Operation::get_op::(body_op, ctx) else { - continue; - }; - let template = inline_asm - .get_attr_inline_asm_template(ctx) - .map(|s| String::from((*s).clone())); - if template.as_deref() != Some(expected_template.as_str()) { - continue; - } - - matches += 1; - assert_eq!( - inline_asm - .get_attr_inline_asm_constraints(ctx) - .map(|s| String::from((*s).clone())) - .as_deref(), - Some("l,l,r,~{memory}") - ); - assert_eq!(llvm::asm_kind(ctx, &inline_asm), llvm::AsmKind::SideEffect); - assert!( - inline_asm - .get_attr_inline_asm_convergent(ctx) - .is_some_and(|value| !bool::from((*value).clone())) - ); - - let operands: Vec<_> = inline_asm.get_operation().deref(ctx).operands().collect(); - assert_eq!(operands.len(), 3); - for operand in &operands[..2] { - let ty = operand.get_type(ctx); - let ty = ty.deref(ctx); - let ptr_ty = ty - .downcast_ref::() - .expect("cp.async pointer operands must lower to LLVM pointers"); - assert_eq!(ptr_ty.address_space(), 0); - } - - let src_size_ty = operands[2].get_type(ctx); - let src_size_ty = src_size_ty.deref(ctx); - let src_size_ty = src_size_ty - .downcast_ref::() - .expect("cp.async src_size must lower to an integer"); - assert_eq!(src_size_ty.width(), 32); - } - } - } - - assert_eq!( - matches, 1, - "missing exact {copy_size}-byte zero-fill cp.async asm" - ); - Ok(()) -} - -// ============================================================================= -// Generated packed arithmetic lowering tests -// ============================================================================= - -#[test] -fn test_generated_packed_arithmetic_lowers_to_exact_pure_inline_asm() -> Result<(), anyhow::Error> { - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let (module_ptr, entry) = - build_test_kernel(&mut ctx, vec![i32_ty.into(), i32_ty.into(), i32_ty.into()]); - - type OpInfo = ( - fn(pliron::context::Ptr) -> pliron::op::OpObj, - std::any::TypeId, - ); - let cases: [(OpInfo, usize, &str, &str); 18] = [ - ( - nvvm::FmaBf16x2Op::get_concrete_op_info(), - 3, - "fma.rn.bf16x2 $0, $1, $2, $3;", - "=r,r,r,r", - ), - ( - nvvm::FmaReluBf16x2Op::get_concrete_op_info(), - 3, - "fma.rn.relu.bf16x2 $0, $1, $2, $3;", - "=r,r,r,r", - ), - ( - nvvm::AddBf16x2Op::get_concrete_op_info(), - 2, - "add.rn.bf16x2 $0, $1, $2;", - "=r,r,r", - ), - ( - nvvm::SubBf16x2Op::get_concrete_op_info(), - 2, - "sub.rn.bf16x2 $0, $1, $2;", - "=r,r,r", - ), - ( - nvvm::MulBf16x2Op::get_concrete_op_info(), - 2, - "mul.rn.bf16x2 $0, $1, $2;", - "=r,r,r", - ), - ( - nvvm::MinBf16x2Op::get_concrete_op_info(), - 2, - "min.bf16x2 $0, $1, $2;", - "=r,r,r", - ), - ( - nvvm::MaxBf16x2Op::get_concrete_op_info(), - 2, - "max.bf16x2 $0, $1, $2;", - "=r,r,r", - ), - ( - nvvm::NegBf16x2Op::get_concrete_op_info(), - 1, - "neg.bf16x2 $0, $1;", - "=r,r", - ), - ( - nvvm::AbsBf16x2Op::get_concrete_op_info(), - 1, - "abs.bf16x2 $0, $1;", - "=r,r", - ), - ( - nvvm::FmaF16x2Op::get_concrete_op_info(), - 3, - "fma.rn.f16x2 $0, $1, $2, $3;", - "=r,r,r,r", - ), - ( - nvvm::FmaReluF16x2Op::get_concrete_op_info(), - 3, - "fma.rn.relu.f16x2 $0, $1, $2, $3;", - "=r,r,r,r", - ), - ( - nvvm::AddF16x2Op::get_concrete_op_info(), - 2, - "add.rn.f16x2 $0, $1, $2;", - "=r,r,r", - ), - ( - nvvm::SubF16x2Op::get_concrete_op_info(), - 2, - "sub.rn.f16x2 $0, $1, $2;", - "=r,r,r", - ), - ( - nvvm::MulF16x2Op::get_concrete_op_info(), - 2, - "mul.rn.f16x2 $0, $1, $2;", - "=r,r,r", - ), - ( - nvvm::MinF16x2Op::get_concrete_op_info(), - 2, - "min.f16x2 $0, $1, $2;", - "=r,r,r", - ), - ( - nvvm::MaxF16x2Op::get_concrete_op_info(), - 2, - "max.f16x2 $0, $1, $2;", - "=r,r,r", - ), - ( - nvvm::NegF16x2Op::get_concrete_op_info(), - 1, - "neg.f16x2 $0, $1;", - "=r,r", - ), - ( - nvvm::AbsF16x2Op::get_concrete_op_info(), - 1, - "abs.f16x2 $0, $1;", - "=r,r", - ), - ]; - - let operands = [ - entry.deref(&ctx).get_argument(0), - entry.deref(&ctx).get_argument(1), - entry.deref(&ctx).get_argument(2), - ]; - for &(op_info, operand_count, _, _) in &cases { - let op = Operation::new( - &mut ctx, - op_info, - vec![i32_ty.into()], - operands[..operand_count].to_vec(), - vec![], - 0, - ); - op.insert_at_back(entry, &ctx); - } - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{}", e))?; - - let mut lowered = Vec::new(); - let module_op = module_ptr.deref(&ctx); - let region = module_op.get_region(0); - let block = region.deref(&ctx).iter(&ctx).next().unwrap(); - for op in block.deref(&ctx).iter(&ctx) { - let Some(func_op) = Operation::get_op::(op, &ctx) else { - continue; - }; - if func_op.get_symbol_name(&ctx).to_string() != "kernel_func" { - continue; - } - let func_region = func_op.get_operation().deref(&ctx).get_region(0); - for func_block in func_region.deref(&ctx).iter(&ctx) { - for body_op in func_block.deref(&ctx).iter(&ctx) { - let Some(inline_asm) = Operation::get_op::(body_op, &ctx) else { - continue; - }; - let inline_asm_op = inline_asm.get_operation(); - let operand_count = inline_asm_op.deref(&ctx).operands().count(); - let result_count = inline_asm_op.deref(&ctx).get_num_results(); - lowered.push(( - inline_asm - .get_attr_inline_asm_template(&ctx) - .map(|s| String::from((*s).clone())) - .expect("packed inline asm must have a template"), - inline_asm - .get_attr_inline_asm_constraints(&ctx) - .map(|s| String::from((*s).clone())) - .expect("packed inline asm must have constraints"), - llvm::asm_kind_opt(&ctx, &inline_asm), - inline_asm - .get_attr_inline_asm_convergent(&ctx) - .map(|b| bool::from((*b).clone())), - operand_count, - result_count, - )); - } - } - } - - assert_eq!( - lowered.len(), - cases.len(), - "each packed operation must lower to exactly one inline-asm op" - ); - for &(_, expected_operand_count, expected_template, expected_constraints) in &cases { - let matches: Vec<_> = lowered - .iter() - .filter(|(template, _, _, _, _, _)| template == expected_template) - .collect(); - assert_eq!( - matches.len(), - 1, - "expected one exact `{expected_template}` lowering" - ); - let (_, constraints, kind, convergent, operand_count, result_count) = matches[0]; - assert_eq!(constraints, expected_constraints, "{expected_template}"); - assert_eq!(*kind, Some(llvm::AsmKind::Pure), "{expected_template}"); - assert_eq!(*convergent, Some(false), "{expected_template}"); - assert_eq!( - *operand_count, expected_operand_count, - "{expected_template} input arity" - ); - assert_eq!(*result_count, 1, "{expected_template} result arity"); - } - - Ok(()) -} - -#[test] -fn test_generated_integer_minmax_lowers_to_exact_pure_inline_asm() -> Result<(), anyhow::Error> { - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![i32_ty.into(), i32_ty.into()]); - - type OpInfo = ( - fn(pliron::context::Ptr) -> pliron::op::OpObj, - std::any::TypeId, - ); - let cases: [(OpInfo, &str); 8] = [ - ( - nvvm::MinReluS32Op::get_concrete_op_info(), - "min.relu.s32 $0, $1, $2;", - ), - ( - nvvm::MaxReluS32Op::get_concrete_op_info(), - "max.relu.s32 $0, $1, $2;", - ), - ( - nvvm::MinS16x2Op::get_concrete_op_info(), - "min.s16x2 $0, $1, $2;", - ), - ( - nvvm::MaxS16x2Op::get_concrete_op_info(), - "max.s16x2 $0, $1, $2;", - ), - ( - nvvm::MinU16x2Op::get_concrete_op_info(), - "min.u16x2 $0, $1, $2;", - ), - ( - nvvm::MaxU16x2Op::get_concrete_op_info(), - "max.u16x2 $0, $1, $2;", - ), - ( - nvvm::MinReluS16x2Op::get_concrete_op_info(), - "min.relu.s16x2 $0, $1, $2;", - ), - ( - nvvm::MaxReluS16x2Op::get_concrete_op_info(), - "max.relu.s16x2 $0, $1, $2;", - ), - ]; - - let operands = [ - entry.deref(&ctx).get_argument(0), - entry.deref(&ctx).get_argument(1), - ]; - for &(op_info, _) in &cases { - let op = Operation::new( - &mut ctx, - op_info, - vec![i32_ty.into()], - operands.to_vec(), - vec![], - 0, - ); - op.insert_at_back(entry, &ctx); - } - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{}", e))?; - - let mut lowered = Vec::new(); - let module_op = module_ptr.deref(&ctx); - let region = module_op.get_region(0); - let block = region.deref(&ctx).iter(&ctx).next().unwrap(); - for op in block.deref(&ctx).iter(&ctx) { - let Some(func_op) = Operation::get_op::(op, &ctx) else { - continue; - }; - if func_op.get_symbol_name(&ctx).to_string() != "kernel_func" { - continue; - } - let func_region = func_op.get_operation().deref(&ctx).get_region(0); - for func_block in func_region.deref(&ctx).iter(&ctx) { - for body_op in func_block.deref(&ctx).iter(&ctx) { - let Some(inline_asm) = Operation::get_op::(body_op, &ctx) else { - continue; - }; - let inline_asm_op = inline_asm.get_operation(); - let operand_count = inline_asm_op.deref(&ctx).operands().count(); - let result_count = inline_asm_op.deref(&ctx).get_num_results(); - lowered.push(( - inline_asm - .get_attr_inline_asm_template(&ctx) - .map(|s| String::from((*s).clone())) - .expect("integer min/max inline asm must have a template"), - inline_asm - .get_attr_inline_asm_constraints(&ctx) - .map(|s| String::from((*s).clone())) - .expect("integer min/max inline asm must have constraints"), - llvm::asm_kind_opt(&ctx, &inline_asm), - inline_asm - .get_attr_inline_asm_convergent(&ctx) - .map(|b| bool::from((*b).clone())), - operand_count, - result_count, - )); - } - } - } - - assert_eq!( - lowered.len(), - cases.len(), - "each integer min/max operation must lower to exactly one inline-asm op" - ); - for &(_, expected_template) in &cases { - let matches: Vec<_> = lowered - .iter() - .filter(|(template, _, _, _, _, _)| template == expected_template) - .collect(); - assert_eq!( - matches.len(), - 1, - "expected one exact `{expected_template}` lowering" - ); - let (_, constraints, kind, convergent, operand_count, result_count) = matches[0]; - assert_eq!(constraints, "=r,r,r", "{expected_template}"); - assert_eq!(*kind, Some(llvm::AsmKind::Pure), "{expected_template}"); - assert_eq!(*convergent, Some(false), "{expected_template}"); - assert_eq!(*operand_count, 2, "{expected_template} input arity"); - assert_eq!(*result_count, 1, "{expected_template} result arity"); - } - - Ok(()) -} - -#[test] -fn test_generated_packed_conversions_lower_to_exact_pure_inline_asm() -> Result<(), anyhow::Error> { - use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let f32_ty = FP32Type::get(&ctx); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![f32_ty.into(), f32_ty.into()]); - let low = entry.deref(&ctx).get_argument(0); - let high = entry.deref(&ctx).get_argument(1); - - type OpInfo = ( - fn(pliron::context::Ptr) -> pliron::op::OpObj, - std::any::TypeId, - ); - let cases: [(OpInfo, &str); 6] = [ - ( - nvvm::CvtF32x2Bf16x2Op::get_concrete_op_info(), - "cvt.rn.bf16x2.f32 $0, $2, $1;", - ), - ( - nvvm::CvtF16x2F32Op::get_concrete_op_info(), - "cvt.rn.f16x2.f32 $0, $2, $1;", - ), - ( - nvvm::CvtRzF16x2F32Op::get_concrete_op_info(), - "cvt.rz.f16x2.f32 $0, $2, $1;", - ), - ( - nvvm::CvtRnReluF16x2F32Op::get_concrete_op_info(), - "cvt.rn.relu.f16x2.f32 $0, $2, $1;", - ), - ( - nvvm::CvtRnReluBf16x2F32Op::get_concrete_op_info(), - "cvt.rn.relu.bf16x2.f32 $0, $2, $1;", - ), - ( - nvvm::CvtRzBf16x2F32Op::get_concrete_op_info(), - "cvt.rz.bf16x2.f32 $0, $2, $1;", - ), - ]; - for &(op_info, _) in &cases { - let op = Operation::new( - &mut ctx, - op_info, - vec![i32_ty.into()], - vec![low, high], - vec![], - 0, - ); - op.insert_at_back(entry, &ctx); - } - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{}", e))?; - - let mut lowered = Vec::new(); - let module_region = module_ptr.deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - for op in module_block.deref(&ctx).iter(&ctx) { - let Some(func) = Operation::get_op::(op, &ctx) else { - continue; - }; - if func.get_symbol_name(&ctx).to_string() != "kernel_func" { - continue; - } - let region = func.get_operation().deref(&ctx).get_region(0); - for block in region.deref(&ctx).iter(&ctx) { - for body_op in block.deref(&ctx).iter(&ctx) { - let Some(asm) = Operation::get_op::(body_op, &ctx) else { - continue; - }; - let template = asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .expect("packed conversion must have an asm template"); - if !template.starts_with("cvt.") { - continue; - } - lowered.push(( - template, - asm.get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())), - asm.get_attr_inline_asm_convergent(&ctx) - .map(|value| bool::from((*value).clone())), - llvm::asm_kind_opt(&ctx, &asm), - asm.get_operation().deref(&ctx).operands().count(), - asm.get_operation().deref(&ctx).get_num_results(), - )); - } - } - } - - assert_eq!( - lowered.len(), - cases.len(), - "each packed conversion must lower to one inline-asm op" - ); - for &(_, expected_template) in &cases { - let matches: Vec<_> = lowered - .iter() - .filter(|(template, _, _, _, _, _)| template == expected_template) - .collect(); - assert_eq!( - matches.len(), - 1, - "expected one exact `{expected_template}` lowering" - ); - let (_, constraints, convergent, kind, operands, results) = matches[0]; - assert_eq!( - constraints.as_deref(), - Some("=r,f,f"), - "{expected_template}" - ); - assert_eq!(*convergent, Some(false), "{expected_template}"); - assert_eq!(*kind, Some(llvm::AsmKind::Pure), "{expected_template}"); - assert_eq!(*operands, 2, "{expected_template} input arity"); - assert_eq!(*results, 1, "{expected_template} result arity"); - } - - Ok(()) -} - -const FP8_CONVERSION_INTRINSICS: [&str; 4] = [ - "llvm_nvvm_ff_to_e4m3x2_rn", - "llvm_nvvm_ff_to_e4m3x2_rn_relu", - "llvm_nvvm_ff_to_e5m2x2_rn", - "llvm_nvvm_ff_to_e5m2x2_rn_relu", -]; - -const FP8_CONVERSION_PTX: [&str; 4] = [ - "cvt.rn.satfinite.e4m3x2.f32 $0, $2, $1;", - "cvt.rn.satfinite.relu.e4m3x2.f32 $0, $2, $1;", - "cvt.rn.satfinite.e5m2x2.f32 $0, $2, $1;", - "cvt.rn.satfinite.relu.e5m2x2.f32 $0, $2, $1;", -]; - -fn lower_all_fp8_conversions( - backend: mir_lower::IntrinsicBackend, -) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { - use pliron::builtin::types::FP32Type; - - let mut ctx = make_test_ctx(); - let f32_ty = FP32Type::get(&ctx); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![f32_ty.into(), f32_ty.into()]); - let low = entry.deref(&ctx).get_argument(0); - let high = entry.deref(&ctx).get_argument(1); - - nvvm::CvtRnSatfiniteE4m3x2F32Op::build(&mut ctx, low, high).insert_at_back(entry, &ctx); - nvvm::CvtRnSatfiniteReluE4m3x2F32Op::build(&mut ctx, low, high).insert_at_back(entry, &ctx); - nvvm::CvtRnSatfiniteE5m2x2F32Op::build(&mut ctx, low, high).insert_at_back(entry, &ctx); - nvvm::CvtRnSatfiniteReluE5m2x2F32Op::build(&mut ctx, low, high).insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm_with_options( - &mut ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: backend, - ..Default::default() - }, - ) - .map_err(|error| anyhow::anyhow!("{error}"))?; - Ok((ctx, module_ptr)) -} - -#[test] -fn test_fp8_conversions_llvm_nvptx_use_exact_typed_calls() -> Result<(), anyhow::Error> { - use pliron::builtin::types::IntegerType; - use pliron::r#type::Typed; - - let (ctx, module_ptr) = lower_all_fp8_conversions(mir_lower::IntrinsicBackend::LlvmNvptx)?; - let mut calls = Vec::new(); - for op in lowered_kernel_body(&ctx, module_ptr) { - assert!( - Operation::get_op::(op, &ctx).is_none(), - "LLVM-NVPTX FP8 conversions must use typed intrinsics" - ); - let Some(call) = Operation::get_op::(op, &ctx) else { - continue; - }; - let CallOpCallable::Direct(callee) = call.callee(&ctx) else { - continue; - }; - let callee = callee.to_string(); - if !FP8_CONVERSION_INTRINSICS.contains(&callee.as_str()) { - continue; - } - let call_op = call.get_operation(); - assert_eq!(call_op.deref(&ctx).get_num_operands(), 2); - assert_eq!(call_op.deref(&ctx).get_num_results(), 1); - let block = call_op.deref(&ctx).get_parent_block().unwrap(); - assert_eq!( - call_op.deref(&ctx).get_operand(0), - block.deref(&ctx).get_argument(1) - ); - assert_eq!( - call_op.deref(&ctx).get_operand(1), - block.deref(&ctx).get_argument(0) - ); - let result_ty = call_op.deref(&ctx).get_result(0).get_type(&ctx); - assert_eq!( - result_ty - .deref(&ctx) - .downcast_ref::() - .expect("FP8 conversion result is an integer") - .width(), - 16 - ); - calls.push(callee); - } - calls.sort(); - let mut expected = FP8_CONVERSION_INTRINSICS.map(str::to_owned); - expected.sort(); - assert_eq!(calls, expected); - Ok(()) -} - -#[test] -fn test_fp8_conversions_libnvvm_use_exact_pure_inline_ptx() -> Result<(), anyhow::Error> { - use pliron::builtin::types::IntegerType; - use pliron::r#type::Typed; - - let (ctx, module_ptr) = lower_all_fp8_conversions(mir_lower::IntrinsicBackend::LibNvvm)?; - let mut templates = Vec::new(); - for op in lowered_kernel_body(&ctx, module_ptr) { - if let Some(call) = Operation::get_op::(op, &ctx) - && let CallOpCallable::Direct(callee) = call.callee(&ctx) - { - assert!( - !FP8_CONVERSION_INTRINSICS.contains(&callee.as_ref()), - "libNVVM FP8 conversions must not use typed intrinsics" - ); - } - let Some(inline_asm) = Operation::get_op::(op, &ctx) else { - continue; - }; - let template = inline_asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .unwrap_or_default(); - if !FP8_CONVERSION_PTX.contains(&template.as_str()) { - continue; - } - assert_eq!( - inline_asm - .get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some("=h,f,f") - ); - assert_eq!(llvm::asm_kind(&ctx, &inline_asm), llvm::AsmKind::Pure); - assert_eq!( - inline_asm - .get_attr_inline_asm_convergent(&ctx) - .map(|value| bool::from((*value).clone())), - Some(false) - ); - let asm_op = inline_asm.get_operation(); - assert_eq!(asm_op.deref(&ctx).get_num_operands(), 2); - assert_eq!(asm_op.deref(&ctx).get_num_results(), 1); - let result_ty = asm_op.deref(&ctx).get_result(0).get_type(&ctx); - assert_eq!( - result_ty - .deref(&ctx) - .downcast_ref::() - .expect("FP8 conversion result is an integer") - .width(), - 16 - ); - templates.push(template); - } - templates.sort(); - let mut expected = FP8_CONVERSION_PTX.map(str::to_owned); - expected.sort(); - assert_eq!(templates, expected); - Ok(()) -} - -const SCALAR_CONVERSION_INTRINSICS: [&str; 10] = [ - "llvm_nvvm_f2tf32_rna", - "llvm_nvvm_f2tf32_rna_satfinite", - "llvm_nvvm_f2tf32_rn", - "llvm_nvvm_f2tf32_rn_relu", - "llvm_nvvm_f2tf32_rn_satfinite", - "llvm_nvvm_f2tf32_rn_relu_satfinite", - "llvm_nvvm_f2tf32_rz", - "llvm_nvvm_f2tf32_rz_relu", - "llvm_nvvm_f2tf32_rz_satfinite", - "llvm_nvvm_f2tf32_rz_relu_satfinite", -]; - -const SCALAR_CONVERSION_PTX: [&str; 10] = [ - "cvt.rna.tf32.f32 $0, $1;", - "cvt.rna.satfinite.tf32.f32 $0, $1;", - "cvt.rn.tf32.f32 $0, $1;", - "cvt.rn.relu.tf32.f32 $0, $1;", - "cvt.rn.satfinite.tf32.f32 $0, $1;", - "cvt.rn.relu.satfinite.tf32.f32 $0, $1;", - "cvt.rz.tf32.f32 $0, $1;", - "cvt.rz.relu.tf32.f32 $0, $1;", - "cvt.rz.satfinite.tf32.f32 $0, $1;", - "cvt.rz.relu.satfinite.tf32.f32 $0, $1;", -]; - -fn lower_all_scalar_conversions( - backend: mir_lower::IntrinsicBackend, -) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { - use pliron::builtin::types::FP32Type; - - let mut ctx = make_test_ctx(); - let f32_ty = FP32Type::get(&ctx); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![f32_ty.into()]); - let value = entry.deref(&ctx).get_argument(0); - - let variants = [ - ( - nvvm::ScalarConversionRoundingAttr::NearestAway, - nvvm::ScalarConversionSaturationAttr::None, - ), - ( - nvvm::ScalarConversionRoundingAttr::NearestAway, - nvvm::ScalarConversionSaturationAttr::Satfinite, - ), - ( - nvvm::ScalarConversionRoundingAttr::NearestEven, - nvvm::ScalarConversionSaturationAttr::None, - ), - ( - nvvm::ScalarConversionRoundingAttr::NearestEven, - nvvm::ScalarConversionSaturationAttr::Relu, - ), - ( - nvvm::ScalarConversionRoundingAttr::NearestEven, - nvvm::ScalarConversionSaturationAttr::Satfinite, - ), - ( - nvvm::ScalarConversionRoundingAttr::NearestEven, - nvvm::ScalarConversionSaturationAttr::ReluSatfinite, - ), - ( - nvvm::ScalarConversionRoundingAttr::TowardZero, - nvvm::ScalarConversionSaturationAttr::None, - ), - ( - nvvm::ScalarConversionRoundingAttr::TowardZero, - nvvm::ScalarConversionSaturationAttr::Relu, - ), - ( - nvvm::ScalarConversionRoundingAttr::TowardZero, - nvvm::ScalarConversionSaturationAttr::Satfinite, - ), - ( - nvvm::ScalarConversionRoundingAttr::TowardZero, - nvvm::ScalarConversionSaturationAttr::ReluSatfinite, - ), - ]; - for (rounding, saturation) in variants { - nvvm::ScalarConversionOp::build(&mut ctx, value, rounding, saturation) - .insert_at_back(entry, &ctx); - } - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm_with_options( - &mut ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: backend, - ..Default::default() - }, - ) - .map_err(|error| anyhow::anyhow!("{error}"))?; - Ok((ctx, module_ptr)) -} - -#[test] -fn test_scalar_conversions_llvm_nvptx_use_exact_typed_calls() -> Result<(), anyhow::Error> { - use llvm_export::types as llvm_types; - use pliron::builtin::type_interfaces::FunctionTypeInterface; - use pliron::builtin::types::{FP32Type, IntegerType}; - - let (ctx, module_ptr) = lower_all_scalar_conversions(mir_lower::IntrinsicBackend::LlvmNvptx)?; - let mut calls = Vec::new(); - for op in lowered_kernel_body(&ctx, module_ptr) { - assert!( - Operation::get_op::(op, &ctx).is_none(), - "LLVM-NVPTX scalar conversions must use typed intrinsics" - ); - let Some(call) = Operation::get_op::(op, &ctx) else { - continue; - }; - let CallOpCallable::Direct(callee) = call.callee(&ctx) else { - continue; - }; - let callee = callee.to_string(); - if !SCALAR_CONVERSION_INTRINSICS.contains(&callee.as_str()) { - continue; - } - - let call_op = call.get_operation(); - assert_eq!(call_op.deref(&ctx).get_num_operands(), 1, "{callee}"); - assert_eq!(call_op.deref(&ctx).get_num_results(), 1, "{callee}"); - let block = call_op.deref(&ctx).get_parent_block().unwrap(); - assert_eq!( - call_op.deref(&ctx).get_operand(0), - block.deref(&ctx).get_argument(0), - "{callee} source operand" - ); - - let function_ty = call.callee_type(&ctx); - let function_ty = function_ty.deref(&ctx); - let function_ty = function_ty - .downcast_ref::() - .expect("scalar conversion callee has an LLVM function type"); - assert_eq!(function_ty.arg_types().len(), 1, "{callee}"); - assert!( - function_ty.arg_types()[0] - .deref(&ctx) - .downcast_ref::() - .is_some(), - "{callee} source type" - ); - assert_eq!( - function_ty - .result_type() - .deref(&ctx) - .downcast_ref::() - .expect("scalar conversion result is an integer") - .width(), - 32, - "{callee} result type" - ); - calls.push(callee); - } - calls.sort(); - let mut expected = SCALAR_CONVERSION_INTRINSICS.map(str::to_owned); - expected.sort(); - assert_eq!(calls, expected); - - let module_region = module_ptr.deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - let mut declarations: Vec<_> = module_block - .deref(&ctx) - .iter(&ctx) - .filter_map(|op| Operation::get_op::(op, &ctx)) - .map(|func| func.get_symbol_name(&ctx).to_string()) - .filter(|name| SCALAR_CONVERSION_INTRINSICS.contains(&name.as_str())) - .collect(); - declarations.sort(); - assert_eq!(declarations, expected); - Ok(()) -} - -#[test] -fn test_scalar_conversions_libnvvm_use_exact_pure_inline_ptx() -> Result<(), anyhow::Error> { - use pliron::builtin::types::IntegerType; - use pliron::r#type::Typed; - - let (ctx, module_ptr) = lower_all_scalar_conversions(mir_lower::IntrinsicBackend::LibNvvm)?; - let mut templates = Vec::new(); - for op in lowered_kernel_body(&ctx, module_ptr) { - if let Some(call) = Operation::get_op::(op, &ctx) - && let CallOpCallable::Direct(callee) = call.callee(&ctx) - { - assert!( - !SCALAR_CONVERSION_INTRINSICS.contains(&callee.as_ref()), - "libNVVM scalar conversions must not use typed intrinsics" - ); - } - let Some(inline_asm) = Operation::get_op::(op, &ctx) else { - continue; - }; - let template = inline_asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .unwrap_or_default(); - assert!( - SCALAR_CONVERSION_PTX.contains(&template.as_str()), - "unexpected scalar conversion template `{template}`" - ); - assert_eq!( - inline_asm - .get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some("=r,f"), - "{template}" - ); - assert_eq!(llvm::asm_kind(&ctx, &inline_asm), llvm::AsmKind::Pure); - assert_eq!( - inline_asm - .get_attr_inline_asm_convergent(&ctx) - .map(|value| bool::from((*value).clone())), - Some(false), - "{template}" - ); - - let asm_op = inline_asm.get_operation(); - assert_eq!(asm_op.deref(&ctx).get_num_operands(), 1, "{template}"); - assert_eq!(asm_op.deref(&ctx).get_num_results(), 1, "{template}"); - let block = asm_op.deref(&ctx).get_parent_block().unwrap(); - assert_eq!( - asm_op.deref(&ctx).get_operand(0), - block.deref(&ctx).get_argument(0), - "{template} source operand" - ); - assert_eq!( - asm_op - .deref(&ctx) - .get_result(0) - .get_type(&ctx) - .deref(&ctx) - .downcast_ref::() - .expect("scalar conversion result is an integer") - .width(), - 32, - "{template} result type" - ); - templates.push(template); - } - templates.sort(); - let mut expected = SCALAR_CONVERSION_PTX.map(str::to_owned); - expected.sort(); - assert_eq!(templates, expected); - Ok(()) -} - -#[test] -fn test_scalar_conversion_invalid_variant_fails_closed() { - use pliron::builtin::types::FP32Type; - - let mut ctx = make_test_ctx(); - let f32_ty = FP32Type::get(&ctx); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![f32_ty.into()]); - let value = entry.deref(&ctx).get_argument(0); - nvvm::ScalarConversionOp::build( - &mut ctx, - value, - nvvm::ScalarConversionRoundingAttr::NearestAway, - nvvm::ScalarConversionSaturationAttr::Relu, - ) - .insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - let result = mir_lower::lower_mir_to_llvm_with_options( - &mut ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: mir_lower::IntrinsicBackend::LlvmNvptx, - ..Default::default() - }, - ); - let error = result.expect_err("unadmitted scalar conversion must not lower"); - assert!( - error.to_string().contains("scalar_conversion"), - "unexpected error: {error}" - ); -} - -fn lower_representative_scalar_arithmetic( - backend: mir_lower::IntrinsicBackend, -) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { - use pliron::builtin::types::FP32Type; - - let mut ctx = make_test_ctx(); - let f32_ty = FP32Type::get(&ctx); - let (module_ptr, entry) = - build_test_kernel(&mut ctx, vec![f32_ty.into(), f32_ty.into(), f32_ty.into()]); - let args: Vec<_> = (0..3) - .map(|index| entry.deref(&ctx).get_argument(index)) - .collect(); - - for (operation, saturation, operands) in [ - ( - nvvm::ScalarArithmeticOperationAttr::Add, - nvvm::ScalarArithmeticSaturationAttr::None, - vec![args[0], args[1]], - ), - ( - nvvm::ScalarArithmeticOperationAttr::Add, - nvvm::ScalarArithmeticSaturationAttr::Sat, - vec![args[0], args[1]], - ), - ( - nvvm::ScalarArithmeticOperationAttr::Fma, - nvvm::ScalarArithmeticSaturationAttr::None, - args.clone(), - ), - ( - nvvm::ScalarArithmeticOperationAttr::Fma, - nvvm::ScalarArithmeticSaturationAttr::Sat, - args, - ), - ] { - nvvm::ScalarArithmeticOp::build( - &mut ctx, - operands, - nvvm::ScalarArithmeticFormatAttr::F32, - operation, - nvvm::ScalarArithmeticRoundingAttr::Rn, - nvvm::ScalarArithmeticSubnormalAttr::Preserve, - saturation, - ) - .insert_at_back(entry, &ctx); - } - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm_with_options( - &mut ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: backend, - ..Default::default() - }, - ) - .map_err(|error| anyhow::anyhow!("{error}"))?; - Ok((ctx, module_ptr)) -} - -fn is_representative_scalar_arithmetic_intrinsic(name: &str) -> bool { - name.starts_with("llvm_nvvm_add_rn") || name.starts_with("llvm_nvvm_fma_rn") -} - -#[test] -fn test_scalar_arithmetic_llvm_uses_inline_ptx_only_for_saturation() -> Result<(), anyhow::Error> { - let (ctx, module_ptr) = - lower_representative_scalar_arithmetic(mir_lower::IntrinsicBackend::LlvmNvptx)?; - let mut calls = Vec::new(); - let mut inline_ptx = Vec::new(); - for op in lowered_kernel_body(&ctx, module_ptr) { - if let Some(call) = Operation::get_op::(op, &ctx) - && let CallOpCallable::Direct(callee) = call.callee(&ctx) - && is_representative_scalar_arithmetic_intrinsic(callee.as_ref()) - { - calls.push(callee.to_string()); - } - let Some(inline_asm) = Operation::get_op::(op, &ctx) else { - continue; - }; - let template = inline_asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .unwrap_or_default(); - let constraints = inline_asm - .get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .unwrap_or_default(); - assert_eq!(llvm::asm_kind(&ctx, &inline_asm), llvm::AsmKind::Pure); - assert_eq!( - inline_asm - .get_attr_inline_asm_convergent(&ctx) - .map(|value| bool::from((*value).clone())), - Some(false) - ); - inline_ptx.push((template, constraints)); - } - - calls.sort(); - assert_eq!(calls, ["llvm_nvvm_add_rn_f", "llvm_nvvm_fma_rn_f"]); - inline_ptx.sort(); - assert_eq!( - inline_ptx, - [ - ("add.rn.sat.f32 $0, $1, $2;".into(), "=f,f,f".into()), - ("fma.rn.sat.f32 $0, $1, $2, $3;".into(), "=f,f,f,f".into(),), - ] - ); - - let module_region = module_ptr.deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - let mut declarations: Vec<_> = module_block - .deref(&ctx) - .iter(&ctx) - .filter_map(|op| Operation::get_op::(op, &ctx)) - .map(|func| func.get_symbol_name(&ctx).to_string()) - .filter(|name| is_representative_scalar_arithmetic_intrinsic(name)) - .collect(); - declarations.sort(); - assert_eq!(declarations, calls); - Ok(()) -} - -#[test] -fn test_scalar_arithmetic_libnvvm_uses_exact_inline_ptx() -> Result<(), anyhow::Error> { - let (ctx, module_ptr) = - lower_representative_scalar_arithmetic(mir_lower::IntrinsicBackend::LibNvvm)?; - let mut calls = Vec::new(); - let mut inline_ptx = Vec::new(); - for op in lowered_kernel_body(&ctx, module_ptr) { - if let Some(call) = Operation::get_op::(op, &ctx) - && let CallOpCallable::Direct(callee) = call.callee(&ctx) - && is_representative_scalar_arithmetic_intrinsic(callee.as_ref()) - { - calls.push(callee.to_string()); - } - let Some(inline_asm) = Operation::get_op::(op, &ctx) else { - continue; - }; - let template = inline_asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .unwrap_or_default(); - let constraints = inline_asm - .get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .unwrap_or_default(); - assert_eq!(llvm::asm_kind(&ctx, &inline_asm), llvm::AsmKind::Pure); - inline_ptx.push((template, constraints)); - } - - assert!(calls.is_empty()); - inline_ptx.sort(); - assert_eq!( - inline_ptx, - [ - ("add.rn.f32 $0, $1, $2;".into(), "=f,f,f".into()), - ("add.rn.sat.f32 $0, $1, $2;".into(), "=f,f,f".into()), - ("fma.rn.f32 $0, $1, $2, $3;".into(), "=f,f,f,f".into(),), - ("fma.rn.sat.f32 $0, $1, $2, $3;".into(), "=f,f,f,f".into(),), - ] - ); - Ok(()) -} - -const STMATRIX_TYPED_INTRINSICS: [&str; 4] = [ - "llvm_nvvm_stmatrix_sync_aligned_m8n8_x2_b16_p3", - "llvm_nvvm_stmatrix_sync_aligned_m8n8_x2_trans_b16_p3", - "llvm_nvvm_stmatrix_sync_aligned_m8n8_x4_b16_p3", - "llvm_nvvm_stmatrix_sync_aligned_m8n8_x4_trans_b16_p3", -]; - -const STMATRIX_PTX: [(&str, &str); 4] = [ - ( - "{ .reg .u64 %ptr64; .reg .u32 %ptr32; cvta.to.shared.u64 %ptr64, $0; cvt.u32.u64 %ptr32, %ptr64; stmatrix.sync.aligned.m8n8.x2.shared.b16 [%ptr32], {$1, $2}; }", - "l,r,r,~{memory}", - ), - ( - "{ .reg .u64 %ptr64; .reg .u32 %ptr32; cvta.to.shared.u64 %ptr64, $0; cvt.u32.u64 %ptr32, %ptr64; stmatrix.sync.aligned.m8n8.x2.trans.shared.b16 [%ptr32], {$1, $2}; }", - "l,r,r,~{memory}", - ), - ( - "{ .reg .u64 %ptr64; .reg .u32 %ptr32; cvta.to.shared.u64 %ptr64, $0; cvt.u32.u64 %ptr32, %ptr64; stmatrix.sync.aligned.m8n8.x4.shared.b16 [%ptr32], {$1, $2, $3, $4}; }", - "l,r,r,r,r,~{memory}", - ), - ( - "{ .reg .u64 %ptr64; .reg .u32 %ptr32; cvta.to.shared.u64 %ptr64, $0; cvt.u32.u64 %ptr32, %ptr64; stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 [%ptr32], {$1, $2, $3, $4}; }", - "l,r,r,r,r,~{memory}", - ), -]; - -fn lower_all_stmatrix_forms( - backend: mir_lower::IntrinsicBackend, -) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { - use dialect_mir::types::MirPtrType; - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signless); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let ptr_ty = MirPtrType::get_generic(&mut ctx, i8_ty.into(), true); - let (module_ptr, entry) = build_test_kernel( - &mut ctx, - vec![ - ptr_ty.into(), - i32_ty.into(), - i32_ty.into(), - i32_ty.into(), - i32_ty.into(), - ], - ); - let args: Vec<_> = (0..5) - .map(|index| entry.deref(&ctx).get_argument(index)) - .collect(); - - for (op_info, operands) in [ - ( - nvvm::StmatrixM8n8X2Op::get_concrete_op_info(), - args[..3].to_vec(), - ), - ( - nvvm::StmatrixM8n8X2TransOp::get_concrete_op_info(), - args[..3].to_vec(), - ), - (nvvm::StmatrixM8n8X4Op::get_concrete_op_info(), args.clone()), - (nvvm::StmatrixM8n8X4TransOp::get_concrete_op_info(), args), - ] { - Operation::new(&mut ctx, op_info, vec![], operands, vec![], 0).insert_at_back(entry, &ctx); - } - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm_with_options( - &mut ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: backend, - ..Default::default() - }, - ) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - Ok((ctx, module_ptr)) -} - -#[test] -fn test_stmatrix_llvm_nvptx_uses_exact_typed_p3_intrinsics() -> Result<(), anyhow::Error> { - use llvm_export::types as llvm_types; - use pliron::r#type::Typed; - - let (ctx, module_ptr) = lower_all_stmatrix_forms(mir_lower::IntrinsicBackend::LlvmNvptx)?; - let mut callees = Vec::new(); - - for op in lowered_kernel_body(&ctx, module_ptr) { - assert!( - Operation::get_op::(op, &ctx).is_none(), - "LLVM-NVPTX stmatrix lowering must not emit inline PTX" - ); - let Some(call) = Operation::get_op::(op, &ctx) else { - continue; - }; - let CallOpCallable::Direct(callee) = call.callee(&ctx) else { - continue; - }; - let callee = callee.to_string(); - if !STMATRIX_TYPED_INTRINSICS.contains(&callee.as_str()) { - continue; - } - - let call_op = call.get_operation().deref(&ctx); - assert!(matches!(call_op.get_num_operands(), 3 | 5)); - assert_eq!(call_op.get_num_results(), 1); - let pointer_ty = call_op.get_operand(0).get_type(&ctx); - let pointer_ty = pointer_ty.deref(&ctx); - let pointer_ty = pointer_ty - .downcast_ref::() - .expect("stmatrix first argument is a pointer"); - assert_eq!(pointer_ty.address_space(), 3); - callees.push(callee); - } - - callees.sort(); - let mut expected = STMATRIX_TYPED_INTRINSICS.map(str::to_owned); - expected.sort(); - assert_eq!(callees, expected); - - let module = Operation::get_op::(module_ptr, &ctx).unwrap(); - let ir = llvm_export::export::export_module_to_string(&ctx, &module) - .expect("typed stmatrix module exports to LLVM IR"); - for intrinsic in STMATRIX_TYPED_INTRINSICS { - let dotted = intrinsic.replace('_', "."); - assert!( - ir.contains(&format!("@{dotted}(ptr addrspace(3)")), - "missing exact typed stmatrix declaration:\n{ir}" - ); - } - assert!(!ir.contains("asm sideeffect"), "{ir}"); - Ok(()) -} - -#[test] -fn test_stmatrix_libnvvm_uses_exact_convergent_memory_asm() -> Result<(), anyhow::Error> { - let (ctx, module_ptr) = lower_all_stmatrix_forms(mir_lower::IntrinsicBackend::LibNvvm)?; - let mut lowered = Vec::new(); - - for op in lowered_kernel_body(&ctx, module_ptr) { - if let Some(call) = Operation::get_op::(op, &ctx) - && let CallOpCallable::Direct(callee) = call.callee(&ctx) - { - assert!( - !STMATRIX_TYPED_INTRINSICS.contains(&callee.as_ref()), - "libNVVM stmatrix lowering must not emit typed intrinsic calls" - ); - } - let Some(inline_asm) = Operation::get_op::(op, &ctx) else { - continue; - }; - lowered.push(( - inline_asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .unwrap_or_default(), - inline_asm - .get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .unwrap_or_default(), - llvm::asm_kind(&ctx, &inline_asm), - op.deref(&ctx).get_num_operands(), - op.deref(&ctx).get_num_results(), - )); - } - - assert_eq!(lowered.len(), STMATRIX_PTX.len()); - for (template, constraints) in STMATRIX_PTX { - let matches: Vec<_> = lowered - .iter() - .filter(|(actual, _, _, _, _)| actual == template) - .collect(); - assert_eq!(matches.len(), 1, "missing exact PTX {template}"); - let (_, actual_constraints, kind, operands, results) = matches[0]; - assert_eq!(actual_constraints, constraints); - assert_eq!(*kind, llvm::AsmKind::Convergent); - assert!(matches!(*operands, 3 | 5)); - assert_eq!(*results, 1); - } - - let module = Operation::get_op::(module_ptr, &ctx).unwrap(); - let ir = llvm_export::export::export_module_to_string(&ctx, &module) - .expect("inline stmatrix module exports to LLVM IR"); - assert_eq!(ir.matches("asm sideeffect").count(), 4, "{ir}"); - assert_eq!(ir.matches("~{memory}").count(), 4, "{ir}"); - assert!(ir.contains("attributes #0 = { convergent }"), "{ir}"); - assert!(!ir.contains("@llvm.nvvm.stmatrix"), "{ir}"); - Ok(()) -} - -// ============================================================================= -// Warp-level matrix (`movmatrix`) lowering test -// ============================================================================= - -#[test] -fn test_movmatrix_trans_b16_lowers_to_inline_asm() -> Result<(), anyhow::Error> { - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![i32_ty.into()]); - - let a_val = entry.deref(&ctx).get_argument(0); - - let op = Operation::new( - &mut ctx, - nvvm::MovmatrixTransB16Op::get_concrete_op_info(), - vec![i32_ty.into()], - vec![a_val], - vec![], - 0, - ); - op.insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let mut found = 0; - let module_region = module_ptr.deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - for module_op in module_block.deref(&ctx).iter(&ctx) { - let Some(function) = Operation::get_op::(module_op, &ctx) else { - continue; - }; - if function.get_symbol_name(&ctx).to_string() != "kernel_func" { - continue; - } - let body = function.get_operation().deref(&ctx).get_region(0); - for block in body.deref(&ctx).iter(&ctx) { - for body_op in block.deref(&ctx).iter(&ctx) { - let Some(asm) = Operation::get_op::(body_op, &ctx) else { - continue; - }; - found += 1; - let template = asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())); - let constraints = asm - .get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())); - assert_eq!( - template.as_deref(), - Some("movmatrix.sync.aligned.m8n8.trans.b16 $0, $1;") - ); - assert_eq!(constraints.as_deref(), Some("=r,r")); - assert_eq!( - llvm::asm_kind_opt(&ctx, &asm), - Some(llvm::AsmKind::Convergent) - ); - assert!( - asm.get_attr_inline_asm_convergent(&ctx) - .is_some_and(|value| bool::from((*value).clone())) - ); - assert!( - !constraints.as_deref().unwrap().contains("memory"), - "register-only movmatrix must not claim a memory clobber" - ); - } - } - } - - assert_eq!(found, 1, "expected one movmatrix inline-asm operation"); - Ok(()) -} - -// ============================================================================= -// ldmatrix lowering tests -// ============================================================================= - -const LDMATRIX_TYPED_INTRINSICS: [&str; 6] = [ - "llvm_nvvm_ldmatrix_sync_aligned_m8n8_x1_b16_p3", - "llvm_nvvm_ldmatrix_sync_aligned_m8n8_x1_trans_b16_p3", - "llvm_nvvm_ldmatrix_sync_aligned_m8n8_x2_b16_p3", - "llvm_nvvm_ldmatrix_sync_aligned_m8n8_x2_trans_b16_p3", - "llvm_nvvm_ldmatrix_sync_aligned_m8n8_x4_b16_p3", - "llvm_nvvm_ldmatrix_sync_aligned_m8n8_x4_trans_b16_p3", -]; - -const LDMATRIX_PTX_TEMPLATES: [&str; 6] = [ - "ldmatrix.sync.aligned.m8n8.x1.shared.b16 {$0}, [$1];", - "ldmatrix.sync.aligned.m8n8.x1.trans.shared.b16 {$0}, [$1];", - "ldmatrix.sync.aligned.m8n8.x2.shared.b16 {$0, $1}, [$2];", - "ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 {$0, $1}, [$2];", - "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {$0, $1, $2, $3}, [$4];", - "ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {$0, $1, $2, $3}, [$4];", -]; - -const LDMATRIX_PTX_CONSTRAINTS: [&str; 6] = [ - "=r,r,~{memory}", - "=r,r,~{memory}", - "=r,=r,r,~{memory}", - "=r,=r,r,~{memory}", - "=r,=r,=r,=r,r,~{memory}", - "=r,=r,=r,=r,r,~{memory}", -]; - -const BLACKWELL_LDMATRIX_CASES: [(&str, &str, &str, usize); 12] = [ - ( - "llvm_nvvm_ldmatrix_sync_aligned_m16n16_x1_trans_b8_p3", - "llvm.nvvm.ldmatrix.sync.aligned.m16n16.x1.trans.b8.p3", - "ldmatrix.sync.aligned.m16n16.x1.trans.shared.b8 {$0, $1}, [$2];", - 2, - ), - ( - "llvm__nvvm_dldmatrix_dsync_daligned_dm16n16_dx1_dtrans_db8x16_db4x16_up64_dp3", - "llvm.nvvm.ldmatrix.sync.aligned.m16n16.x1.trans.b8x16.b4x16_p64.p3", - "ldmatrix.sync.aligned.m16n16.x1.trans.shared.b8x16.b4x16_p64 {$0, $1}, [$2];", - 2, - ), - ( - "llvm__nvvm_dldmatrix_dsync_daligned_dm16n16_dx1_dtrans_db8x16_db6x16_up32_dp3", - "llvm.nvvm.ldmatrix.sync.aligned.m16n16.x1.trans.b8x16.b6x16_p32.p3", - "ldmatrix.sync.aligned.m16n16.x1.trans.shared.b8x16.b6x16_p32 {$0, $1}, [$2];", - 2, - ), - ( - "llvm_nvvm_ldmatrix_sync_aligned_m16n16_x2_trans_b8_p3", - "llvm.nvvm.ldmatrix.sync.aligned.m16n16.x2.trans.b8.p3", - "ldmatrix.sync.aligned.m16n16.x2.trans.shared.b8 {$0, $1, $2, $3}, [$4];", - 4, - ), - ( - "llvm__nvvm_dldmatrix_dsync_daligned_dm16n16_dx2_dtrans_db8x16_db4x16_up64_dp3", - "llvm.nvvm.ldmatrix.sync.aligned.m16n16.x2.trans.b8x16.b4x16_p64.p3", - "ldmatrix.sync.aligned.m16n16.x2.trans.shared.b8x16.b4x16_p64 {$0, $1, $2, $3}, [$4];", - 4, - ), - ( - "llvm__nvvm_dldmatrix_dsync_daligned_dm16n16_dx2_dtrans_db8x16_db6x16_up32_dp3", - "llvm.nvvm.ldmatrix.sync.aligned.m16n16.x2.trans.b8x16.b6x16_p32.p3", - "ldmatrix.sync.aligned.m16n16.x2.trans.shared.b8x16.b6x16_p32 {$0, $1, $2, $3}, [$4];", - 4, - ), - ( - "llvm__nvvm_dldmatrix_dsync_daligned_dm8n16_dx1_db8x16_db4x16_up64_dp3", - "llvm.nvvm.ldmatrix.sync.aligned.m8n16.x1.b8x16.b4x16_p64.p3", - "ldmatrix.sync.aligned.m8n16.x1.shared.b8x16.b4x16_p64 {$0}, [$1];", - 1, - ), - ( - "llvm__nvvm_dldmatrix_dsync_daligned_dm8n16_dx1_db8x16_db6x16_up32_dp3", - "llvm.nvvm.ldmatrix.sync.aligned.m8n16.x1.b8x16.b6x16_p32.p3", - "ldmatrix.sync.aligned.m8n16.x1.shared.b8x16.b6x16_p32 {$0}, [$1];", - 1, - ), - ( - "llvm__nvvm_dldmatrix_dsync_daligned_dm8n16_dx2_db8x16_db4x16_up64_dp3", - "llvm.nvvm.ldmatrix.sync.aligned.m8n16.x2.b8x16.b4x16_p64.p3", - "ldmatrix.sync.aligned.m8n16.x2.shared.b8x16.b4x16_p64 {$0, $1}, [$2];", - 2, - ), - ( - "llvm__nvvm_dldmatrix_dsync_daligned_dm8n16_dx2_db8x16_db6x16_up32_dp3", - "llvm.nvvm.ldmatrix.sync.aligned.m8n16.x2.b8x16.b6x16_p32.p3", - "ldmatrix.sync.aligned.m8n16.x2.shared.b8x16.b6x16_p32 {$0, $1}, [$2];", - 2, - ), - ( - "llvm__nvvm_dldmatrix_dsync_daligned_dm8n16_dx4_db8x16_db4x16_up64_dp3", - "llvm.nvvm.ldmatrix.sync.aligned.m8n16.x4.b8x16.b4x16_p64.p3", - "ldmatrix.sync.aligned.m8n16.x4.shared.b8x16.b4x16_p64 {$0, $1, $2, $3}, [$4];", - 4, - ), - ( - "llvm__nvvm_dldmatrix_dsync_daligned_dm8n16_dx4_db8x16_db6x16_up32_dp3", - "llvm.nvvm.ldmatrix.sync.aligned.m8n16.x4.b8x16.b6x16_p32.p3", - "ldmatrix.sync.aligned.m8n16.x4.shared.b8x16.b6x16_p32 {$0, $1, $2, $3}, [$4];", - 4, - ), -]; - -fn ldmatrix_constraints(register_count: usize) -> &'static str { - match register_count { - 1 => "=r,r,~{memory}", - 2 => "=r,=r,r,~{memory}", - 4 => "=r,=r,=r,=r,r,~{memory}", - _ => unreachable!("closed ldmatrix register count"), - } -} - -fn lower_all_ldmatrix_forms( - address_space: u32, - backend: mir_lower::IntrinsicBackend, - compatibility: bool, -) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { - use dialect_mir::types::MirPtrType; - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let ptr_ty = MirPtrType::get(&mut ctx, u32_ty.into(), true, address_space); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![ptr_ty.into()]); - let pointer = entry.deref(&ctx).get_argument(0); - if compatibility { - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - for (op_info, result_count) in [ - (nvvm::LdmatrixX1Op::get_concrete_op_info(), 1), - (nvvm::LdmatrixX1TransOp::get_concrete_op_info(), 1), - (nvvm::LdmatrixX2Op::get_concrete_op_info(), 2), - (nvvm::LdmatrixX2TransOp::get_concrete_op_info(), 2), - (nvvm::LdmatrixX4Op::get_concrete_op_info(), 4), - (nvvm::LdmatrixX4TransOp::get_concrete_op_info(), 4), - ] { - Operation::new( - &mut ctx, - op_info, - vec![u32_ty.into(); result_count], - vec![pointer], - vec![], - 0, - ) - .insert_at_back(entry, &ctx); - } - } else { - for (multiplicity, layout) in [ - ( - nvvm::LdmatrixMultiplicityAttr::X1, - nvvm::LdmatrixLayoutAttr::Normal, - ), - ( - nvvm::LdmatrixMultiplicityAttr::X1, - nvvm::LdmatrixLayoutAttr::Transposed, - ), - ( - nvvm::LdmatrixMultiplicityAttr::X2, - nvvm::LdmatrixLayoutAttr::Normal, - ), - ( - nvvm::LdmatrixMultiplicityAttr::X2, - nvvm::LdmatrixLayoutAttr::Transposed, - ), - ( - nvvm::LdmatrixMultiplicityAttr::X4, - nvvm::LdmatrixLayoutAttr::Normal, - ), - ( - nvvm::LdmatrixMultiplicityAttr::X4, - nvvm::LdmatrixLayoutAttr::Transposed, - ), - ] { - nvvm::LdmatrixOp::build( - &mut ctx, - pointer, - nvvm::LdmatrixShapeAttr::M8n8, - multiplicity, - layout, - nvvm::LdmatrixElementAttr::B16, - nvvm::LdmatrixStateSpaceAttr::Shared, - ) - .insert_at_back(entry, &ctx); - } - } - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm_with_options( - &mut ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: backend, - ..Default::default() - }, - ) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - Ok((ctx, module_ptr)) -} - -fn lower_all_blackwell_ldmatrix_forms( - address_space: u32, - backend: mir_lower::IntrinsicBackend, -) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { - use dialect_mir::types::MirPtrType; - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let u8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); - let ptr_ty = MirPtrType::get(&mut ctx, u8_ty.into(), true, address_space); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![ptr_ty.into()]); - let pointer = entry.deref(&ctx).get_argument(0); - - for (shape, multiplicity, layout, element) in [ - ( - nvvm::LdmatrixShapeAttr::M16n16, - nvvm::LdmatrixMultiplicityAttr::X1, - nvvm::LdmatrixLayoutAttr::Transposed, - nvvm::LdmatrixElementAttr::B8, - ), - ( - nvvm::LdmatrixShapeAttr::M16n16, - nvvm::LdmatrixMultiplicityAttr::X1, - nvvm::LdmatrixLayoutAttr::Transposed, - nvvm::LdmatrixElementAttr::B8x16B4x16P64, - ), - ( - nvvm::LdmatrixShapeAttr::M16n16, - nvvm::LdmatrixMultiplicityAttr::X1, - nvvm::LdmatrixLayoutAttr::Transposed, - nvvm::LdmatrixElementAttr::B8x16B6x16P32, - ), - ( - nvvm::LdmatrixShapeAttr::M16n16, - nvvm::LdmatrixMultiplicityAttr::X2, - nvvm::LdmatrixLayoutAttr::Transposed, - nvvm::LdmatrixElementAttr::B8, - ), - ( - nvvm::LdmatrixShapeAttr::M16n16, - nvvm::LdmatrixMultiplicityAttr::X2, - nvvm::LdmatrixLayoutAttr::Transposed, - nvvm::LdmatrixElementAttr::B8x16B4x16P64, - ), - ( - nvvm::LdmatrixShapeAttr::M16n16, - nvvm::LdmatrixMultiplicityAttr::X2, - nvvm::LdmatrixLayoutAttr::Transposed, - nvvm::LdmatrixElementAttr::B8x16B6x16P32, - ), - ( - nvvm::LdmatrixShapeAttr::M8n16, - nvvm::LdmatrixMultiplicityAttr::X1, - nvvm::LdmatrixLayoutAttr::Normal, - nvvm::LdmatrixElementAttr::B8x16B4x16P64, - ), - ( - nvvm::LdmatrixShapeAttr::M8n16, - nvvm::LdmatrixMultiplicityAttr::X1, - nvvm::LdmatrixLayoutAttr::Normal, - nvvm::LdmatrixElementAttr::B8x16B6x16P32, - ), - ( - nvvm::LdmatrixShapeAttr::M8n16, - nvvm::LdmatrixMultiplicityAttr::X2, - nvvm::LdmatrixLayoutAttr::Normal, - nvvm::LdmatrixElementAttr::B8x16B4x16P64, - ), - ( - nvvm::LdmatrixShapeAttr::M8n16, - nvvm::LdmatrixMultiplicityAttr::X2, - nvvm::LdmatrixLayoutAttr::Normal, - nvvm::LdmatrixElementAttr::B8x16B6x16P32, - ), - ( - nvvm::LdmatrixShapeAttr::M8n16, - nvvm::LdmatrixMultiplicityAttr::X4, - nvvm::LdmatrixLayoutAttr::Normal, - nvvm::LdmatrixElementAttr::B8x16B4x16P64, - ), - ( - nvvm::LdmatrixShapeAttr::M8n16, - nvvm::LdmatrixMultiplicityAttr::X4, - nvvm::LdmatrixLayoutAttr::Normal, - nvvm::LdmatrixElementAttr::B8x16B6x16P32, - ), - ] { - nvvm::LdmatrixOp::build( - &mut ctx, - pointer, - shape, - multiplicity, - layout, - element, - nvvm::LdmatrixStateSpaceAttr::Shared, - ) - .insert_at_back(entry, &ctx); - } - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm_with_options( - &mut ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: backend, - ..Default::default() - }, - ) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - Ok((ctx, module_ptr)) -} - -fn lowered_kernel_body( - ctx: &Context, - module_ptr: pliron::context::Ptr, -) -> Vec> { - let module_region = module_ptr.deref(ctx).get_region(0); - let module_block = module_region.deref(ctx).iter(ctx).next().unwrap(); - for op in module_block.deref(ctx).iter(ctx) { - let Some(func_op) = Operation::get_op::(op, ctx) else { - continue; - }; - if func_op.get_symbol_name(ctx).to_string() != "kernel_func" { - continue; - } - - let func_region = func_op.get_operation().deref(ctx).get_region(0); - return func_region - .deref(ctx) - .iter(ctx) - .flat_map(|block| block.deref(ctx).iter(ctx)) - .collect(); - } - panic!("lowered kernel function not found") -} - -fn assert_ldmatrix_producer_result_shape( - ctx: &Context, - op: pliron::context::Ptr, - register_count: usize, -) { - use llvm_export::types as llvm_types; - use pliron::builtin::types::IntegerType; - use pliron::r#type::Typed; - - let ty = op.deref(ctx).get_result(0).get_type(ctx); - let ty = ty.deref(ctx); - if register_count == 1 { - assert_eq!( - ty.downcast_ref::() - .expect("single-register ldmatrix returns i32") - .width(), - 32 - ); - } else { - let ty = ty - .downcast_ref::() - .expect("multi-register ldmatrix returns an LLVM struct"); - assert_eq!(ty.num_fields(), register_count); - for index in 0..register_count { - assert_eq!( - ty.field_type(index) - .deref(ctx) - .downcast_ref::() - .expect("ldmatrix fragment field is i32") - .width(), - 32 - ); - } - } -} - -#[test] -fn test_ldmatrix_llvm_nvptx_uses_exact_typed_p3_intrinsics() -> Result<(), anyhow::Error> { - use llvm_export::types as llvm_types; - use pliron::builtin::type_interfaces::FunctionTypeInterface; - use pliron::builtin::types::IntegerType; - use pliron::r#type::Typed; - - for address_space in [0, 3] { - let (ctx, module_ptr) = - lower_all_ldmatrix_forms(address_space, mir_lower::IntrinsicBackend::LlvmNvptx, false)?; - let body = lowered_kernel_body(&ctx, module_ptr); - let mut callees = Vec::new(); - let mut cast_count = 0; - let mut extract_count = 0; - - for op in body { - if let Some(call) = Operation::get_op::(op, &ctx) - && let CallOpCallable::Direct(callee) = call.callee(&ctx) - { - let callee = callee.to_string(); - let register_count = if callee.contains("_x1_") { - 1 - } else if callee.contains("_x2_") { - 2 - } else { - 4 - }; - let function_ty = call.callee_type(&ctx); - let function_ty = function_ty.deref(&ctx); - let function_ty = function_ty - .downcast_ref::() - .expect("ldmatrix callee has an LLVM function type"); - assert_eq!(function_ty.arg_types().len(), 1); - let argument_ty = function_ty.arg_types()[0].deref(&ctx); - let argument_ty = argument_ty - .downcast_ref::() - .expect("ldmatrix argument is a pointer"); - assert_eq!(argument_ty.address_space(), 3); - - let result_ty = function_ty.result_type(); - let result_ty = result_ty.deref(&ctx); - if register_count == 1 { - let result_ty = result_ty - .downcast_ref::() - .expect("x1 returns i32"); - assert_eq!(result_ty.width(), 32); - } else { - let result_ty = result_ty - .downcast_ref::() - .expect("x2/x4 return an LLVM struct"); - assert_eq!(result_ty.num_fields(), register_count); - for index in 0..result_ty.num_fields() { - let field = result_ty.field_type(index); - let field = field.deref(&ctx); - let field = field - .downcast_ref::() - .expect("fragment register is i32"); - assert_eq!(field.width(), 32); - } - } - - callees.push(callee); - assert_eq!(op.deref(&ctx).get_num_operands(), 1); - assert_eq!(op.deref(&ctx).get_num_results(), 1); - } - if Operation::get_op::(op, &ctx).is_some() { - cast_count += 1; - let cast = op.deref(&ctx); - let source_ty = cast.get_operand(0).get_type(&ctx); - let source_ty = source_ty.deref(&ctx); - let source_ty = source_ty - .downcast_ref::() - .expect("addrspacecast source is a pointer"); - let result_ty = cast.get_result(0).get_type(&ctx); - let result_ty = result_ty.deref(&ctx); - let result_ty = result_ty - .downcast_ref::() - .expect("addrspacecast result is a pointer"); - assert_eq!( - (source_ty.address_space(), result_ty.address_space()), - (0, 3) - ); - } - extract_count += - usize::from(Operation::get_op::(op, &ctx).is_some()); - assert!( - Operation::get_op::(op, &ctx).is_none(), - "LLVM-NVPTX ldmatrix lowering must not emit inline PTX" - ); - assert!( - Operation::get_op::(op, &ctx).is_none(), - "the typed intrinsic consumes the shared pointer directly" - ); - } - - callees.sort(); - let mut expected = LDMATRIX_TYPED_INTRINSICS.map(str::to_owned); - expected.sort(); - assert_eq!(callees, expected); - assert_eq!(cast_count, if address_space == 0 { 6 } else { 0 }); - assert_eq!( - extract_count, 12, - "x2/x4 structs must preserve result order" - ); - - let module = Operation::get_op::(module_ptr, &ctx).unwrap(); - let ir = llvm_export::export::export_module_to_string(&ctx, &module) - .expect("typed ldmatrix module exports to LLVM IR"); - for intrinsic in LDMATRIX_TYPED_INTRINSICS { - let dotted = intrinsic.replace('_', "."); - assert!( - ir.contains(&format!("@{dotted}(ptr addrspace(3)")), - "underscore symbol must export as exact dotted p3 intrinsic:\n{ir}" - ); - } - assert!(!ir.contains("@llvm_nvvm_ldmatrix")); - } - Ok(()) -} - -#[test] -fn test_ldmatrix_libnvvm_uses_exact_convergent_shared_ptx() -> Result<(), anyhow::Error> { - use llvm_export::types as llvm_types; - use pliron::r#type::Typed; - - for address_space in [0, 3] { - let (ctx, module_ptr) = - lower_all_ldmatrix_forms(address_space, mir_lower::IntrinsicBackend::LibNvvm, false)?; - let body = lowered_kernel_body(&ctx, module_ptr); - let mut lowered = Vec::new(); - let mut cast_count = 0; - let mut ptrtoint_count = 0; - - for op in body { - assert!( - Operation::get_op::(op, &ctx).is_none(), - "libNVVM ldmatrix lowering must not emit typed intrinsic calls" - ); - cast_count += - usize::from(Operation::get_op::(op, &ctx).is_some()); - if Operation::get_op::(op, &ctx).is_some() { - ptrtoint_count += 1; - let cast = op.deref(&ctx); - let source_ty = cast.get_operand(0).get_type(&ctx); - let source_ty = source_ty.deref(&ctx); - let source_ty = source_ty - .downcast_ref::() - .expect("ptrtoint source is a pointer"); - assert_eq!(source_ty.address_space(), 3); - } - - let Some(inline_asm) = Operation::get_op::(op, &ctx) else { - continue; - }; - lowered.push(( - inline_asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .unwrap_or_default(), - inline_asm - .get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .unwrap_or_default(), - llvm::asm_kind(&ctx, &inline_asm), - op.deref(&ctx).get_num_operands(), - op.deref(&ctx).get_num_results(), - )); - } - - assert_eq!(lowered.len(), 6); - for (index, expected_template) in LDMATRIX_PTX_TEMPLATES.iter().enumerate() { - let matching: Vec<_> = lowered - .iter() - .filter(|(template, _, _, _, _)| template == expected_template) - .collect(); - assert_eq!(matching.len(), 1, "missing exact PTX `{expected_template}`"); - let (_, constraints, kind, operands, results) = matching[0]; - assert_eq!(constraints, LDMATRIX_PTX_CONSTRAINTS[index]); - assert_eq!(*kind, llvm::AsmKind::Convergent); - assert_eq!(*operands, 1, "inline PTX consumes one i32 shared address"); - assert_eq!(*results, 1, "inline PTX returns one scalar or struct"); - assert!(!expected_template.contains("cvta.to.shared")); - } - assert_eq!(cast_count, if address_space == 0 { 6 } else { 0 }); - assert_eq!(ptrtoint_count, 6); - - let module = Operation::get_op::(module_ptr, &ctx).unwrap(); - let ir = llvm_export::export::export_module_to_string(&ctx, &module) - .expect("inline ldmatrix module exports to LLVM IR"); - assert_eq!(ir.matches("asm sideeffect").count(), 6, "{ir}"); - assert_eq!(ir.matches("~{memory}").count(), 6, "{ir}"); - assert!(ir.contains("attributes #0 = { convergent }"), "{ir}"); - assert!(!ir.contains("@llvm.nvvm.ldmatrix"), "{ir}"); - } - Ok(()) -} - -#[test] -fn test_blackwell_ldmatrix_llvm_uses_all_exact_lossless_p3_intrinsics() -> Result<(), anyhow::Error> -{ - use llvm_export::types as llvm_types; - use pliron::builtin::type_interfaces::FunctionTypeInterface; - - for address_space in [0, 3] { - let (ctx, module_ptr) = lower_all_blackwell_ldmatrix_forms( - address_space, - mir_lower::IntrinsicBackend::LlvmNvptx, - )?; - let mut seen = [0; 12]; - let mut cast_count = 0; - let mut extract_count = 0; - - for op in lowered_kernel_body(&ctx, module_ptr) { - assert!(Operation::get_op::(op, &ctx).is_none()); - assert!(Operation::get_op::(op, &ctx).is_none()); - cast_count += - usize::from(Operation::get_op::(op, &ctx).is_some()); - extract_count += - usize::from(Operation::get_op::(op, &ctx).is_some()); - - let Some(call) = Operation::get_op::(op, &ctx) else { - continue; - }; - let CallOpCallable::Direct(callee) = call.callee(&ctx) else { - panic!("Blackwell ldmatrix intrinsic call must be direct"); - }; - let callee = callee.to_string(); - let index = BLACKWELL_LDMATRIX_CASES - .iter() - .position(|(identifier, _, _, _)| *identifier == callee) - .expect("exact lossless Blackwell ldmatrix identifier"); - seen[index] += 1; - let register_count = BLACKWELL_LDMATRIX_CASES[index].3; - assert_ldmatrix_producer_result_shape(&ctx, op, register_count); - - let function_ty = call.callee_type(&ctx); - let function_ty = function_ty.deref(&ctx); - let function_ty = function_ty - .downcast_ref::() - .expect("Blackwell ldmatrix has an LLVM function type"); - assert_eq!(function_ty.arg_types().len(), 1); - assert_eq!( - function_ty.arg_types()[0] - .deref(&ctx) - .downcast_ref::() - .expect("Blackwell ldmatrix argument is a pointer") - .address_space(), - 3 - ); - } - - assert_eq!(seen, [1; 12]); - assert_eq!(cast_count, if address_space == 0 { 12 } else { 0 }); - assert_eq!(extract_count, 30, "all multi-register results are unpacked"); - - let module = Operation::get_op::(module_ptr, &ctx).unwrap(); - let ir = llvm_export::export::export_module_to_string(&ctx, &module) - .expect("typed Blackwell ldmatrix module exports to LLVM IR"); - for (identifier, symbol, _, _) in BLACKWELL_LDMATRIX_CASES { - assert!( - ir.contains(&format!("@{symbol}(ptr addrspace(3)")), - "missing exact intrinsic symbol {symbol}:\n{ir}" - ); - assert!( - !ir.contains(&format!("@{identifier}(")), - "encoded Rust identifier leaked into LLVM IR: {identifier}" - ); - } - assert!(ir.contains("b4x16_p64.p3"), "literal _p64 was lost: {ir}"); - assert!(ir.contains("b6x16_p32.p3"), "literal _p32 was lost: {ir}"); - } - Ok(()) -} - -#[test] -fn test_blackwell_ldmatrix_libnvvm_uses_all_exact_convergent_templates_without_externs() --> Result<(), anyhow::Error> { - for address_space in [0, 3] { - let (ctx, module_ptr) = lower_all_blackwell_ldmatrix_forms( - address_space, - mir_lower::IntrinsicBackend::LibNvvm, - )?; - let mut seen = [0; 12]; - let mut cast_count = 0; - let mut ptrtoint_count = 0; - let mut extract_count = 0; - - for op in lowered_kernel_body(&ctx, module_ptr) { - assert!( - Operation::get_op::(op, &ctx).is_none(), - "libNVVM Blackwell ldmatrix must not emit an extern call" - ); - cast_count += - usize::from(Operation::get_op::(op, &ctx).is_some()); - ptrtoint_count += - usize::from(Operation::get_op::(op, &ctx).is_some()); - extract_count += - usize::from(Operation::get_op::(op, &ctx).is_some()); - - let Some(asm) = Operation::get_op::(op, &ctx) else { - continue; - }; - let template = asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .unwrap_or_default(); - let index = BLACKWELL_LDMATRIX_CASES - .iter() - .position(|(_, _, expected, _)| *expected == template) - .expect("exact Blackwell ldmatrix inline-PTX template"); - seen[index] += 1; - let register_count = BLACKWELL_LDMATRIX_CASES[index].3; - assert_eq!( - asm.get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some(ldmatrix_constraints(register_count)) - ); - assert_eq!(llvm::asm_kind(&ctx, &asm), llvm::AsmKind::Convergent); - assert_ldmatrix_producer_result_shape(&ctx, op, register_count); - } - - assert_eq!(seen, [1; 12]); - assert_eq!(cast_count, if address_space == 0 { 12 } else { 0 }); - assert_eq!(ptrtoint_count, 12); - assert_eq!(extract_count, 30, "all multi-register results are unpacked"); - - let module = Operation::get_op::(module_ptr, &ctx).unwrap(); - let ir = llvm_export::export::export_module_to_string(&ctx, &module) - .expect("inline Blackwell ldmatrix module exports to LLVM IR"); - assert_eq!(ir.matches("asm sideeffect").count(), 12, "{ir}"); - assert_eq!(ir.matches("~{memory}").count(), 12, "{ir}"); - assert!(!ir.contains("@llvm.nvvm.ldmatrix"), "{ir}"); - assert!(!ir.contains("llvm__nvvm_dldmatrix"), "{ir}"); - } - Ok(()) -} - -#[test] -fn test_blackwell_ldmatrix_rejects_unadmitted_m16n16_x4() { - use dialect_mir::types::MirPtrType; - use pliron::builtin::types::{IntegerType, Signedness}; - - for backend in [ - mir_lower::IntrinsicBackend::LlvmNvptx, - mir_lower::IntrinsicBackend::LibNvvm, - ] { - let mut ctx = make_test_ctx(); - let u8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); - let ptr_ty = MirPtrType::get_shared(&mut ctx, u8_ty.into(), false); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![ptr_ty.into()]); - let pointer = entry.deref(&ctx).get_argument(0); - nvvm::LdmatrixOp::build( - &mut ctx, - pointer, - nvvm::LdmatrixShapeAttr::M16n16, - nvvm::LdmatrixMultiplicityAttr::X4, - nvvm::LdmatrixLayoutAttr::Transposed, - nvvm::LdmatrixElementAttr::B8, - nvvm::LdmatrixStateSpaceAttr::Shared, - ) - .insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - let error = mir_lower::lower_mir_to_llvm_with_options( - &mut ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: backend, - ..Default::default() - }, - ) - .expect_err("m16n16.x4 must fail closed") - .to_string(); - assert!(error.contains("missing or unsupported variant"), "{error}"); - } -} - -#[test] -fn test_classic_ldmatrix_compatibility_ops_keep_exact_lowering() -> Result<(), anyhow::Error> { - use llvm_export::types as llvm_types; - use pliron::builtin::types::IntegerType; - use pliron::r#type::Typed; - - fn assert_result_shape( - ctx: &Context, - op: pliron::context::Ptr, - register_count: usize, - ) { - let operation = op.deref(ctx); - assert_eq!(operation.get_num_operands(), 1); - assert_eq!(operation.get_num_results(), 1); - let result_ty = operation.get_result(0).get_type(ctx); - let result_ty = result_ty.deref(ctx); - if register_count == 1 { - let result_ty = result_ty - .downcast_ref::() - .expect("x1 returns i32"); - assert_eq!(result_ty.width(), 32); - } else { - let result_ty = result_ty - .downcast_ref::() - .expect("x2/x4 return an LLVM struct"); - assert_eq!(result_ty.num_fields(), register_count); - for index in 0..result_ty.num_fields() { - let field = result_ty.field_type(index); - let field = field.deref(ctx); - let field = field - .downcast_ref::() - .expect("fragment register is i32"); - assert_eq!(field.width(), 32); - } - } - } - - let register_counts = [1, 1, 2, 2, 4, 4]; - for backend in [ - mir_lower::IntrinsicBackend::LlvmNvptx, - mir_lower::IntrinsicBackend::LibNvvm, - ] { - let (ctx, module_ptr) = lower_all_ldmatrix_forms(3, backend, true)?; - let mut seen = [0; 6]; - let mut extract_count = 0; - - for op in lowered_kernel_body(&ctx, module_ptr) { - extract_count += - usize::from(Operation::get_op::(op, &ctx).is_some()); - match backend { - mir_lower::IntrinsicBackend::LlvmNvptx => { - assert!(Operation::get_op::(op, &ctx).is_none()); - let Some(call) = Operation::get_op::(op, &ctx) else { - continue; - }; - let CallOpCallable::Direct(callee) = call.callee(&ctx) else { - panic!("ldmatrix intrinsic call must be direct"); - }; - let callee = callee.to_string(); - let index = LDMATRIX_TYPED_INTRINSICS - .iter() - .position(|expected| *expected == callee) - .expect("exact typed ldmatrix intrinsic"); - seen[index] += 1; - assert_result_shape(&ctx, op, register_counts[index]); - } - mir_lower::IntrinsicBackend::LibNvvm => { - assert!(Operation::get_op::(op, &ctx).is_none()); - let Some(inline_asm) = Operation::get_op::(op, &ctx) else { - continue; - }; - let template = inline_asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .unwrap_or_default(); - let index = LDMATRIX_PTX_TEMPLATES - .iter() - .position(|expected| *expected == template) - .expect("exact ldmatrix PTX template"); - seen[index] += 1; - assert_eq!( - inline_asm - .get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some(LDMATRIX_PTX_CONSTRAINTS[index]) - ); - assert_eq!(llvm::asm_kind(&ctx, &inline_asm), llvm::AsmKind::Convergent); - assert_result_shape(&ctx, op, register_counts[index]); - } - } - } - - assert_eq!(seen, [1; 6]); - assert_eq!(extract_count, 12, "x2/x4 results keep their order"); - } - Ok(()) -} - -#[test] -fn test_ldmatrix_rejects_non_shared_pointer_spaces() { - use dialect_mir::types::MirPtrType; - use pliron::builtin::types::{IntegerType, Signedness}; - - for backend in [ - mir_lower::IntrinsicBackend::LlvmNvptx, - mir_lower::IntrinsicBackend::LibNvvm, - ] { - for address_space in [1, 4, 5] { - let mut ctx = make_test_ctx(); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let ptr_ty = MirPtrType::get(&mut ctx, u32_ty.into(), false, address_space); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![ptr_ty.into()]); - let pointer = entry.deref(&ctx).get_argument(0); - nvvm::LdmatrixOp::build( - &mut ctx, - pointer, - nvvm::LdmatrixShapeAttr::M8n8, - nvvm::LdmatrixMultiplicityAttr::X1, - nvvm::LdmatrixLayoutAttr::Normal, - nvvm::LdmatrixElementAttr::B16, - nvvm::LdmatrixStateSpaceAttr::Shared, - ) - .insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - let error = mir_lower::lower_mir_to_llvm_with_options( - &mut ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: backend, - ..Default::default() - }, - ) - .expect_err("global/constant pointers must fail closed") - .to_string(); - assert!( - error.contains(&format!("not address space {address_space}")), - "{error}" - ); - } - } -} - -#[test] -fn test_ldmatrix_rejects_non_pointer_operand() { - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![i32_ty.into()]); - let not_a_pointer = entry.deref(&ctx).get_argument(0); - nvvm::LdmatrixOp::build( - &mut ctx, - not_a_pointer, - nvvm::LdmatrixShapeAttr::M8n8, - nvvm::LdmatrixMultiplicityAttr::X1, - nvvm::LdmatrixLayoutAttr::Normal, - nvvm::LdmatrixElementAttr::B16, - nvvm::LdmatrixStateSpaceAttr::Shared, - ) - .insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - let error = mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .expect_err("non-pointer ldmatrix input must fail closed") - .to_string(); - assert!(error.contains("operand must be a MIR pointer"), "{error}"); -} - -#[test] -fn test_ldmatrix_rejects_wrong_result_arity() { - use dialect_mir::types::MirPtrType; - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - let ptr_ty = MirPtrType::get_shared(&mut ctx, u32_ty.into(), false); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![ptr_ty.into()]); - let pointer = entry.deref(&ctx).get_argument(0); - let op = Operation::new( - &mut ctx, - nvvm::LdmatrixOp::get_concrete_op_info(), - vec![i32_ty.into(), i32_ty.into()], - vec![pointer], - vec![], - 0, - ); - let ldmatrix = nvvm::LdmatrixOp::new(op); - ldmatrix.set_attr_nvvm_ldmatrix_shape(&ctx, nvvm::LdmatrixShapeAttr::M8n8); - ldmatrix.set_attr_nvvm_ldmatrix_multiplicity(&ctx, nvvm::LdmatrixMultiplicityAttr::X1); - ldmatrix.set_attr_nvvm_ldmatrix_layout(&ctx, nvvm::LdmatrixLayoutAttr::Normal); - ldmatrix.set_attr_nvvm_ldmatrix_element(&ctx, nvvm::LdmatrixElementAttr::B16); - ldmatrix.set_attr_nvvm_ldmatrix_state_space(&ctx, nvvm::LdmatrixStateSpaceAttr::Shared); - op.insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - let error = mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .expect_err("x1 must return exactly one register") - .to_string(); - assert!(error.contains("requires 1 u32 results"), "{error}"); -} - -#[test] -fn test_generated_register_mma_variants_lower_to_exact_convergent_inline_ptx() --> Result<(), anyhow::Error> { - use pliron::builtin::types::{FP32Type, FP64Type, IntegerType, Signedness}; - use pliron::r#type::TypeHandle; - - #[derive(Clone, Copy)] - enum Carrier { - F32, - F64, - I32, - U32, - } - - struct Case { - shape: nvvm::RegisterMmaShapeAttr, - operation: nvvm::RegisterMmaOperationAttr, - accumulator: nvvm::RegisterMmaAccumulatorAttr, - a_element: nvvm::RegisterMmaElementAttr, - b_element: nvvm::RegisterMmaElementAttr, - overflow: nvvm::RegisterMmaOverflowAttr, - operands: &'static [Carrier], - results: &'static [Carrier], - template: String, - constraints: &'static str, - } - - let mut cases = vec![ - Case { - shape: nvvm::RegisterMmaShapeAttr::M16n8k16, - operation: nvvm::RegisterMmaOperationAttr::Multiply, - accumulator: nvvm::RegisterMmaAccumulatorAttr::F32, - a_element: nvvm::RegisterMmaElementAttr::Bf16, - b_element: nvvm::RegisterMmaElementAttr::Bf16, - overflow: nvvm::RegisterMmaOverflowAttr::NotApplicable, - operands: &[ - Carrier::F32, - Carrier::F32, - Carrier::F32, - Carrier::F32, - Carrier::U32, - Carrier::U32, - Carrier::U32, - Carrier::U32, - Carrier::U32, - Carrier::U32, - ], - results: &[Carrier::F32; 4], - template: concat!( - "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 ", - "{$0, $1, $2, $3}, {$8, $9, $10, $11}, ", - "{$12, $13}, {$4, $5, $6, $7};" - ) - .into(), - constraints: "=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r", - }, - Case { - shape: nvvm::RegisterMmaShapeAttr::M16n8k16, - operation: nvvm::RegisterMmaOperationAttr::Multiply, - accumulator: nvvm::RegisterMmaAccumulatorAttr::F32, - a_element: nvvm::RegisterMmaElementAttr::F16, - b_element: nvvm::RegisterMmaElementAttr::F16, - overflow: nvvm::RegisterMmaOverflowAttr::NotApplicable, - operands: &[ - Carrier::F32, - Carrier::F32, - Carrier::F32, - Carrier::F32, - Carrier::U32, - Carrier::U32, - Carrier::U32, - Carrier::U32, - Carrier::U32, - Carrier::U32, - ], - results: &[Carrier::F32; 4], - template: concat!( - "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 ", - "{$0, $1, $2, $3}, {$8, $9, $10, $11}, ", - "{$12, $13}, {$4, $5, $6, $7};" - ) - .into(), - constraints: "=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r", - }, - Case { - shape: nvvm::RegisterMmaShapeAttr::M16n8k8, - operation: nvvm::RegisterMmaOperationAttr::Multiply, - accumulator: nvvm::RegisterMmaAccumulatorAttr::F32, - a_element: nvvm::RegisterMmaElementAttr::Tf32, - b_element: nvvm::RegisterMmaElementAttr::Tf32, - overflow: nvvm::RegisterMmaOverflowAttr::NotApplicable, - operands: &[ - Carrier::F32, - Carrier::F32, - Carrier::F32, - Carrier::F32, - Carrier::U32, - Carrier::U32, - Carrier::U32, - Carrier::U32, - Carrier::U32, - Carrier::U32, - ], - results: &[Carrier::F32; 4], - template: concat!( - "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 ", - "{$0, $1, $2, $3}, {$8, $9, $10, $11}, ", - "{$12, $13}, {$4, $5, $6, $7};" - ) - .into(), - constraints: "=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r", - }, - Case { - shape: nvvm::RegisterMmaShapeAttr::M8n8k4, - operation: nvvm::RegisterMmaOperationAttr::Multiply, - accumulator: nvvm::RegisterMmaAccumulatorAttr::F64, - a_element: nvvm::RegisterMmaElementAttr::F64, - b_element: nvvm::RegisterMmaElementAttr::F64, - overflow: nvvm::RegisterMmaOverflowAttr::NotApplicable, - operands: &[Carrier::F64; 4], - results: &[Carrier::F64; 2], - template: concat!( - "mma.sync.aligned.m8n8k4.row.col.f64.f64.f64.f64 ", - "{$0, $1}, {$4}, {$5}, {$2, $3};" - ) - .into(), - constraints: "=d,=d,d,d,d,d", - }, - ]; - let c2_a1_b1: &'static [Carrier] = &[Carrier::I32, Carrier::I32, Carrier::U32, Carrier::U32]; - let c4_a2_b1: &'static [Carrier] = &[ - Carrier::I32, - Carrier::I32, - Carrier::I32, - Carrier::I32, - Carrier::U32, - Carrier::U32, - Carrier::U32, - ]; - let c4_a4_b2: &'static [Carrier] = &[ - Carrier::I32, - Carrier::I32, - Carrier::I32, - Carrier::I32, - Carrier::U32, - Carrier::U32, - Carrier::U32, - Carrier::U32, - Carrier::U32, - Carrier::U32, - ]; - let d2_i32: &'static [Carrier] = &[Carrier::I32; 2]; - let d4_i32: &'static [Carrier] = &[Carrier::I32; 4]; - let register_list = |first, count| { - format!( - "{{{}}}", - (first..first + count) - .map(|index| format!("${index}")) - .collect::>() - .join(", ") - ) - }; - - for (shape, shape_attr, operands, results, accumulator_count, a_count, b_count, constraints) in [ - ( - "m8n8k16", - nvvm::RegisterMmaShapeAttr::M8n8k16, - c2_a1_b1, - d2_i32, - 2, - 1, - 1, - "=r,=r,r,r,r,r", - ), - ( - "m16n8k16", - nvvm::RegisterMmaShapeAttr::M16n8k16, - c4_a2_b1, - d4_i32, - 4, - 2, - 1, - "=r,=r,=r,=r,r,r,r,r,r,r,r", - ), - ( - "m16n8k32", - nvvm::RegisterMmaShapeAttr::M16n8k32, - c4_a4_b2, - d4_i32, - 4, - 4, - 2, - "=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,r", - ), - ] { - for (a_name, a_element) in [ - ("s8", nvvm::RegisterMmaElementAttr::S8), - ("u8", nvvm::RegisterMmaElementAttr::U8), - ] { - for (b_name, b_element) in [ - ("s8", nvvm::RegisterMmaElementAttr::S8), - ("u8", nvvm::RegisterMmaElementAttr::U8), - ] { - for (overflow_name, overflow) in [ - ("", nvvm::RegisterMmaOverflowAttr::Wrapping), - (".satfinite", nvvm::RegisterMmaOverflowAttr::Satfinite), - ] { - let result_count = results.len(); - let d = register_list(0, result_count); - let c = register_list(result_count, accumulator_count); - let a = register_list(result_count + accumulator_count, a_count); - let b = register_list(result_count + accumulator_count + a_count, b_count); - cases.push(Case { - shape: shape_attr.clone(), - operation: nvvm::RegisterMmaOperationAttr::Multiply, - accumulator: nvvm::RegisterMmaAccumulatorAttr::S32, - a_element: a_element.clone(), - b_element: b_element.clone(), - overflow, - operands, - results, - template: format!( - "mma.sync.aligned.{shape}.row.col{overflow_name}.s32.{a_name}.{b_name}.s32 {d}, {a}, {b}, {c};" - ), - constraints, - }); - } - } - } - } - - for (shape, shape_attr, operands, results, accumulator_count, a_count, b_count, constraints) in [ - ( - "m8n8k32", - nvvm::RegisterMmaShapeAttr::M8n8k32, - c2_a1_b1, - d2_i32, - 2, - 1, - 1, - "=r,=r,r,r,r,r", - ), - ( - "m16n8k32", - nvvm::RegisterMmaShapeAttr::M16n8k32, - c4_a2_b1, - d4_i32, - 4, - 2, - 1, - "=r,=r,=r,=r,r,r,r,r,r,r,r", - ), - ( - "m16n8k64", - nvvm::RegisterMmaShapeAttr::M16n8k64, - c4_a4_b2, - d4_i32, - 4, - 4, - 2, - "=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,r", - ), - ] { - for (a_name, a_element) in [ - ("s4", nvvm::RegisterMmaElementAttr::S4), - ("u4", nvvm::RegisterMmaElementAttr::U4), - ] { - for (b_name, b_element) in [ - ("s4", nvvm::RegisterMmaElementAttr::S4), - ("u4", nvvm::RegisterMmaElementAttr::U4), - ] { - for (overflow_name, overflow) in [ - ("", nvvm::RegisterMmaOverflowAttr::Wrapping), - (".satfinite", nvvm::RegisterMmaOverflowAttr::Satfinite), - ] { - let result_count = results.len(); - let d = register_list(0, result_count); - let c = register_list(result_count, accumulator_count); - let a = register_list(result_count + accumulator_count, a_count); - let b = register_list(result_count + accumulator_count + a_count, b_count); - cases.push(Case { - shape: shape_attr.clone(), - operation: nvvm::RegisterMmaOperationAttr::Multiply, - accumulator: nvvm::RegisterMmaAccumulatorAttr::S32, - a_element: a_element.clone(), - b_element: b_element.clone(), - overflow, - operands, - results, - template: format!( - "mma.sync.aligned.{shape}.row.col{overflow_name}.s32.{a_name}.{b_name}.s32 {d}, {a}, {b}, {c};" - ), - constraints, - }); - } - } - } - } - - for (shape, shape_attr, operands, results, accumulator_count, a_count, b_count, constraints) in [ - ( - "m8n8k128", - nvvm::RegisterMmaShapeAttr::M8n8k128, - c2_a1_b1, - d2_i32, - 2, - 1, - 1, - "=r,=r,r,r,r,r", - ), - ( - "m16n8k128", - nvvm::RegisterMmaShapeAttr::M16n8k128, - c4_a2_b1, - d4_i32, - 4, - 2, - 1, - "=r,=r,=r,=r,r,r,r,r,r,r,r", - ), - ( - "m16n8k256", - nvvm::RegisterMmaShapeAttr::M16n8k256, - c4_a4_b2, - d4_i32, - 4, - 4, - 2, - "=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,r", - ), - ] { - let result_count = results.len(); - let d = register_list(0, result_count); - let c = register_list(result_count, accumulator_count); - let a = register_list(result_count + accumulator_count, a_count); - let b = register_list(result_count + accumulator_count + a_count, b_count); - for (operation_name, operation) in [ - ("xor", nvvm::RegisterMmaOperationAttr::XorPopc), - ("and", nvvm::RegisterMmaOperationAttr::AndPopc), - ] { - cases.push(Case { - shape: shape_attr.clone(), - operation, - accumulator: nvvm::RegisterMmaAccumulatorAttr::S32, - a_element: nvvm::RegisterMmaElementAttr::B1, - b_element: nvvm::RegisterMmaElementAttr::B1, - overflow: nvvm::RegisterMmaOverflowAttr::Wrapping, - operands, - results, - template: format!( - "mma.sync.aligned.{shape}.row.col.s32.b1.b1.s32.{operation_name}.popc {d}, {a}, {b}, {c};" - ), - constraints, - }); - } - } - assert_eq!(cases.len(), 58); - - let carrier_type = |ctx: &Context, carrier: Carrier| -> TypeHandle { - match carrier { - Carrier::F32 => FP32Type::get(ctx).into(), - Carrier::F64 => FP64Type::get(ctx).into(), - Carrier::I32 => IntegerType::get(ctx, 32, Signedness::Signed).into(), - Carrier::U32 => IntegerType::get(ctx, 32, Signedness::Unsigned).into(), - } - }; - - for backend in [ - mir_lower::IntrinsicBackend::LlvmNvptx, - mir_lower::IntrinsicBackend::LibNvvm, - ] { - for case in &cases { - let mut ctx = make_test_ctx(); - let argument_types = case - .operands - .iter() - .map(|carrier| carrier_type(&ctx, *carrier)) - .collect(); - let result_types = case - .results - .iter() - .map(|carrier| carrier_type(&ctx, *carrier)) - .collect(); - let (module_ptr, entry) = build_test_kernel(&mut ctx, argument_types); - let operands = (0..case.operands.len()) - .map(|index| entry.deref(&ctx).get_argument(index)) - .collect(); - let operation = Operation::new( - &mut ctx, - nvvm::RegisterMmaOp::get_concrete_op_info(), - result_types, - operands, - vec![], - 0, - ); - let mma = nvvm::RegisterMmaOp::new(operation); - mma.set_attr_nvvm_register_mma_shape(&ctx, case.shape.clone()); - let uses_legacy_default = matches!( - (&case.operation, &case.a_element), - ( - nvvm::RegisterMmaOperationAttr::Multiply, - nvvm::RegisterMmaElementAttr::Bf16 - ) - ); - if !uses_legacy_default { - mma.set_attr_nvvm_register_mma_operation(&ctx, case.operation.clone()); - } - mma.set_attr_nvvm_register_mma_accumulator(&ctx, case.accumulator.clone()); - mma.set_attr_nvvm_register_mma_a_element(&ctx, case.a_element.clone()); - mma.set_attr_nvvm_register_mma_b_element(&ctx, case.b_element.clone()); - mma.set_attr_nvvm_register_mma_a_layout(&ctx, nvvm::RegisterMmaLayoutAttr::Row); - mma.set_attr_nvvm_register_mma_b_layout(&ctx, nvvm::RegisterMmaLayoutAttr::Col); - mma.set_attr_nvvm_register_mma_overflow(&ctx, case.overflow.clone()); - operation.insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm_with_options( - &mut ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: backend, - ..Default::default() - }, - ) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let body = lowered_kernel_body(&ctx, module_ptr); - let lowered = body - .iter() - .filter_map(|op| Operation::get_op::(*op, &ctx)) - .collect::>(); - assert_eq!(lowered.len(), 1, "{:?}", backend); - let asm = &lowered[0]; - assert_eq!( - asm.get_attr_inline_asm_template(&ctx) - .as_deref() - .map(|value| String::from(value.clone())), - Some(case.template.clone()) - ); - assert_eq!( - asm.get_attr_inline_asm_constraints(&ctx) - .as_deref() - .map(|value| String::from(value.clone())), - Some(case.constraints.to_string()) - ); - assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); - assert_eq!( - asm.get_operation().deref(&ctx).get_num_operands(), - case.operands.len() - ); - assert_eq!(asm.get_operation().deref(&ctx).get_num_results(), 1); - assert_eq!( - body.iter() - .filter(|op| Operation::get_op::(**op, &ctx).is_some()) - .count(), - case.results.len() - ); - - let module = Operation::get_op::(module_ptr, &ctx).unwrap(); - let ir = llvm_export::export::export_module_to_string(&ctx, &module) - .expect("generated register MMA exports to LLVM IR"); - assert!(ir.contains("asm sideeffect"), "{ir}"); - assert!(ir.contains("{ convergent }"), "{ir}"); - } - } - - Ok(()) -} - -#[test] -fn test_generated_sparse_mma_variants_lower_to_exact_convergent_inline_ptx() --> Result<(), anyhow::Error> { - use pliron::builtin::attributes::IntegerAttr; - use pliron::builtin::types::{IntegerType, Signedness}; - use pliron::utils::apint::APInt; - use std::num::NonZeroUsize; - - let cases = [ - ( - "s8", - "s8", - "", - nvvm::SparseMmaElementAttr::S8, - nvvm::SparseMmaElementAttr::S8, - nvvm::SparseMmaOverflowAttr::Wrapping, - ), - ( - "s8", - "u8", - "", - nvvm::SparseMmaElementAttr::S8, - nvvm::SparseMmaElementAttr::U8, - nvvm::SparseMmaOverflowAttr::Wrapping, - ), - ( - "u8", - "u8", - "", - nvvm::SparseMmaElementAttr::U8, - nvvm::SparseMmaElementAttr::U8, - nvvm::SparseMmaOverflowAttr::Wrapping, - ), - ( - "u8", - "s8", - "", - nvvm::SparseMmaElementAttr::U8, - nvvm::SparseMmaElementAttr::S8, - nvvm::SparseMmaOverflowAttr::Wrapping, - ), - ( - "s8", - "s8", - ".satfinite", - nvvm::SparseMmaElementAttr::S8, - nvvm::SparseMmaElementAttr::S8, - nvvm::SparseMmaOverflowAttr::Satfinite, - ), - ( - "s8", - "u8", - ".satfinite", - nvvm::SparseMmaElementAttr::S8, - nvvm::SparseMmaElementAttr::U8, - nvvm::SparseMmaOverflowAttr::Satfinite, - ), - ( - "u8", - "u8", - ".satfinite", - nvvm::SparseMmaElementAttr::U8, - nvvm::SparseMmaElementAttr::U8, - nvvm::SparseMmaOverflowAttr::Satfinite, - ), - ( - "u8", - "s8", - ".satfinite", - nvvm::SparseMmaElementAttr::U8, - nvvm::SparseMmaElementAttr::S8, - nvvm::SparseMmaOverflowAttr::Satfinite, - ), - ]; - let metadata_modes = [ - ("sp", nvvm::SparseMmaMetadataAttr::Standard), - ("sp::ordered_metadata", nvvm::SparseMmaMetadataAttr::Ordered), - ]; - - for backend in [ - mir_lower::IntrinsicBackend::LlvmNvptx, - mir_lower::IntrinsicBackend::LibNvvm, - ] { - for (metadata_name, metadata) in &metadata_modes { - for (index, (a_name, b_name, overflow_name, a_element, b_element, overflow)) in - cases.iter().enumerate() - { - let mut ctx = make_test_ctx(); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let argument_types = (0..4) - .map(|_| i32_ty.into()) - .chain((0..5).map(|_| u32_ty.into())) - .collect(); - let (module_ptr, entry) = build_test_kernel(&mut ctx, argument_types); - - let selector_value = (index % 2) as u32; - let selector_op = Operation::new( - &mut ctx, - mir::MirConstantOp::get_concrete_op_info(), - vec![u32_ty.into()], - vec![], - vec![], - 0, - ); - mir::MirConstantOp::new(selector_op).set_attr_value( - &ctx, - IntegerAttr::new( - u32_ty, - APInt::from_u32(selector_value, NonZeroUsize::new(32).unwrap()), - ), - ); - selector_op.insert_at_back(entry, &ctx); - let selector = selector_op.deref(&ctx).get_result(0); - - let operands = (0..9) - .map(|operand| entry.deref(&ctx).get_argument(operand)) - .chain(std::iter::once(selector)) - .collect(); - let operation = Operation::new( - &mut ctx, - nvvm::SparseMmaOp::get_concrete_op_info(), - vec![i32_ty.into(); 4], - operands, - vec![], - 0, - ); - let mma = nvvm::SparseMmaOp::new(operation); - mma.set_attr_nvvm_sparse_mma_shape(&ctx, nvvm::SparseMmaShapeAttr::M16n8k32); - mma.set_attr_nvvm_sparse_mma_accumulator(&ctx, nvvm::SparseMmaAccumulatorAttr::S32); - mma.set_attr_nvvm_sparse_mma_a_element(&ctx, a_element.clone()); - mma.set_attr_nvvm_sparse_mma_b_element(&ctx, b_element.clone()); - mma.set_attr_nvvm_sparse_mma_a_layout(&ctx, nvvm::SparseMmaLayoutAttr::Row); - mma.set_attr_nvvm_sparse_mma_b_layout(&ctx, nvvm::SparseMmaLayoutAttr::Col); - mma.set_attr_nvvm_sparse_mma_overflow(&ctx, overflow.clone()); - mma.set_attr_nvvm_sparse_mma_metadata(&ctx, metadata.clone()); - mma.set_attr_nvvm_sparse_mma_selector( - &ctx, - nvvm::SparseMmaSelectorAttr::ImmediateZeroOrOne, - ); - operation.insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm_with_options( - &mut ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: backend, - ..Default::default() - }, - ) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let body = lowered_kernel_body(&ctx, module_ptr); - let lowered = body - .iter() - .filter_map(|op| Operation::get_op::(*op, &ctx)) - .collect::>(); - assert_eq!(lowered.len(), 1, "{backend:?}"); - let asm = &lowered[0]; - let expected_template = format!( - "mma.{metadata_name}.sync.aligned.m16n8k32.row.col{overflow_name}.s32.{a_name}.{b_name}.s32 {{$0, $1, $2, $3}}, {{$8, $9}}, {{$10, $11}}, {{$4, $5, $6, $7}}, $12, $13;" - ); - assert_eq!( - asm.get_attr_inline_asm_template(&ctx) - .as_deref() - .map(|value| String::from(value.clone())), - Some(expected_template) - ); - assert_eq!( - asm.get_attr_inline_asm_constraints(&ctx) - .as_deref() - .map(|value| String::from(value.clone())), - Some("=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,n".to_string()) - ); - assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); - let asm_operation = asm.get_operation().deref(&ctx); - assert_eq!(asm_operation.get_num_operands(), 10); - assert_eq!(asm_operation.get_num_results(), 1); - let lowered_selector = asm_operation.get_operand(9); - let defining_op = lowered_selector - .defining_op() - .expect("sparse MMA selector remains an LLVM constant"); - let constant = Operation::get_op::(defining_op, &ctx) - .expect("sparse MMA selector remains an LLVM integer constant"); - let attribute = constant.get_value(&ctx); - let integer = attribute - .downcast_ref::() - .expect("sparse MMA selector is an integer"); - assert_eq!(integer.value().bw(), 32); - assert_eq!(integer.value().to_u64(), selector_value as u64); - assert_eq!( - body.iter() - .filter(|op| { - Operation::get_op::(**op, &ctx).is_some() - }) - .count(), - 4 - ); - - let module = Operation::get_op::(module_ptr, &ctx).unwrap(); - let ir = llvm_export::export::export_module_to_string(&ctx, &module) - .expect("generated sparse MMA exports to LLVM IR"); - assert!(ir.contains("asm sideeffect"), "{ir}"); - assert!(ir.contains("{ convergent }"), "{ir}"); - } - } - } - - Ok(()) -} - -#[test] -fn test_generated_sparse_mma_m16n8k64_lowers_to_exact_convergent_inline_ptx() --> Result<(), anyhow::Error> { - use pliron::builtin::attributes::IntegerAttr; - use pliron::builtin::types::{IntegerType, Signedness}; - use pliron::utils::apint::APInt; - use std::num::NonZeroUsize; - - let cases = [ - ( - "s8", - "u8", - "", - nvvm::SparseMmaElementAttr::S8, - nvvm::SparseMmaElementAttr::U8, - nvvm::SparseMmaOverflowAttr::Wrapping, - ), - ( - "u8", - "s8", - ".satfinite", - nvvm::SparseMmaElementAttr::U8, - nvvm::SparseMmaElementAttr::S8, - nvvm::SparseMmaOverflowAttr::Satfinite, - ), - ]; - let metadata_cases = [ - ("sp", nvvm::SparseMmaMetadataAttr::Standard), - ("sp::ordered_metadata", nvvm::SparseMmaMetadataAttr::Ordered), - ]; - - for backend in [ - mir_lower::IntrinsicBackend::LlvmNvptx, - mir_lower::IntrinsicBackend::LibNvvm, - ] { - for (metadata_name, metadata) in &metadata_cases { - for (a_name, b_name, overflow_name, a_element, b_element, overflow) in &cases { - let mut ctx = make_test_ctx(); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let argument_types = (0..4) - .map(|_| i32_ty.into()) - .chain((0..9).map(|_| u32_ty.into())) - .collect(); - let (module_ptr, entry) = build_test_kernel(&mut ctx, argument_types); - - let selector_op = Operation::new( - &mut ctx, - mir::MirConstantOp::get_concrete_op_info(), - vec![u32_ty.into()], - vec![], - vec![], - 0, - ); - mir::MirConstantOp::new(selector_op).set_attr_value( - &ctx, - IntegerAttr::new(u32_ty, APInt::from_u32(0, NonZeroUsize::new(32).unwrap())), - ); - selector_op.insert_at_back(entry, &ctx); - let selector = selector_op.deref(&ctx).get_result(0); - - let operands = (0..13) - .map(|index| entry.deref(&ctx).get_argument(index)) - .chain(std::iter::once(selector)) - .collect(); - let operation = Operation::new( - &mut ctx, - nvvm::SparseMmaOp::get_concrete_op_info(), - vec![i32_ty.into(); 4], - operands, - vec![], - 0, - ); - let mma = nvvm::SparseMmaOp::new(operation); - mma.set_attr_nvvm_sparse_mma_shape(&ctx, nvvm::SparseMmaShapeAttr::M16n8k64); - mma.set_attr_nvvm_sparse_mma_accumulator(&ctx, nvvm::SparseMmaAccumulatorAttr::S32); - mma.set_attr_nvvm_sparse_mma_a_element(&ctx, a_element.clone()); - mma.set_attr_nvvm_sparse_mma_b_element(&ctx, b_element.clone()); - mma.set_attr_nvvm_sparse_mma_a_layout(&ctx, nvvm::SparseMmaLayoutAttr::Row); - mma.set_attr_nvvm_sparse_mma_b_layout(&ctx, nvvm::SparseMmaLayoutAttr::Col); - mma.set_attr_nvvm_sparse_mma_overflow(&ctx, overflow.clone()); - mma.set_attr_nvvm_sparse_mma_metadata(&ctx, metadata.clone()); - mma.set_attr_nvvm_sparse_mma_selector( - &ctx, - nvvm::SparseMmaSelectorAttr::ImmediateZero, - ); - operation.insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm_with_options( - &mut ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: backend, - ..Default::default() - }, - ) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let body = lowered_kernel_body(&ctx, module_ptr); - let lowered = body - .iter() - .filter_map(|op| Operation::get_op::(*op, &ctx)) - .collect::>(); - assert_eq!(lowered.len(), 1, "{backend:?}"); - let asm = &lowered[0]; - let expected_template = format!( - "mma.{metadata_name}.sync.aligned.m16n8k64.row.col{overflow_name}.s32.{a_name}.{b_name}.s32 {{$0, $1, $2, $3}}, {{$8, $9, $10, $11}}, {{$12, $13, $14, $15}}, {{$4, $5, $6, $7}}, $16, $17;" - ); - assert_eq!( - asm.get_attr_inline_asm_template(&ctx) - .as_deref() - .map(|value| String::from(value.clone())), - Some(expected_template) - ); - assert_eq!( - asm.get_attr_inline_asm_constraints(&ctx) - .as_deref() - .map(|value| String::from(value.clone())), - Some("=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,r,r,r,r,n".to_string()) - ); - assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); - let asm_operation = asm.get_operation().deref(&ctx); - assert_eq!(asm_operation.get_num_operands(), 14); - assert_eq!(asm_operation.get_num_results(), 1); - let lowered_selector = asm_operation.get_operand(13); - let defining_op = lowered_selector - .defining_op() - .expect("sparse MMA selector remains an LLVM constant"); - let constant = Operation::get_op::(defining_op, &ctx) - .expect("sparse MMA selector remains an LLVM integer constant"); - let attribute = constant.get_value(&ctx); - let integer = attribute - .downcast_ref::() - .expect("sparse MMA selector is an integer"); - assert_eq!(integer.value().bw(), 32); - assert_eq!(integer.value().to_u64(), 0); - assert_eq!( - body.iter() - .filter(|op| { - Operation::get_op::(**op, &ctx).is_some() - }) - .count(), - 4 - ); - - let module = Operation::get_op::(module_ptr, &ctx).unwrap(); - let ir = llvm_export::export::export_module_to_string(&ctx, &module) - .expect("generated sparse MMA exports to LLVM IR"); - assert!(ir.contains("asm sideeffect"), "{ir}"); - assert!(ir.contains("{ convergent }"), "{ir}"); - } - } - } - - Ok(()) -} - -#[test] -fn test_generated_sparse_mma_m16n8k64_int4_lowers_both_metadata_modes() -> Result<(), anyhow::Error> -{ - use pliron::builtin::attributes::IntegerAttr; - use pliron::builtin::types::{IntegerType, Signedness}; - use pliron::utils::apint::APInt; - use std::num::NonZeroUsize; - - let cases = [ - ( - "s4", - "s4", - "", - nvvm::SparseMmaElementAttr::S4, - nvvm::SparseMmaElementAttr::S4, - nvvm::SparseMmaOverflowAttr::Wrapping, - ), - ( - "s4", - "u4", - "", - nvvm::SparseMmaElementAttr::S4, - nvvm::SparseMmaElementAttr::U4, - nvvm::SparseMmaOverflowAttr::Wrapping, - ), - ( - "u4", - "u4", - "", - nvvm::SparseMmaElementAttr::U4, - nvvm::SparseMmaElementAttr::U4, - nvvm::SparseMmaOverflowAttr::Wrapping, - ), - ( - "u4", - "s4", - "", - nvvm::SparseMmaElementAttr::U4, - nvvm::SparseMmaElementAttr::S4, - nvvm::SparseMmaOverflowAttr::Wrapping, - ), - ( - "s4", - "s4", - ".satfinite", - nvvm::SparseMmaElementAttr::S4, - nvvm::SparseMmaElementAttr::S4, - nvvm::SparseMmaOverflowAttr::Satfinite, - ), - ( - "s4", - "u4", - ".satfinite", - nvvm::SparseMmaElementAttr::S4, - nvvm::SparseMmaElementAttr::U4, - nvvm::SparseMmaOverflowAttr::Satfinite, - ), - ( - "u4", - "u4", - ".satfinite", - nvvm::SparseMmaElementAttr::U4, - nvvm::SparseMmaElementAttr::U4, - nvvm::SparseMmaOverflowAttr::Satfinite, - ), - ( - "u4", - "s4", - ".satfinite", - nvvm::SparseMmaElementAttr::U4, - nvvm::SparseMmaElementAttr::S4, - nvvm::SparseMmaOverflowAttr::Satfinite, - ), - ]; - - for backend in [ - mir_lower::IntrinsicBackend::LlvmNvptx, - mir_lower::IntrinsicBackend::LibNvvm, - ] { - for (metadata_name, metadata) in [ - ("sp", nvvm::SparseMmaMetadataAttr::Standard), - ("sp::ordered_metadata", nvvm::SparseMmaMetadataAttr::Ordered), - ] { - for (index, (a_name, b_name, overflow_name, a_element, b_element, overflow)) in - cases.iter().enumerate() - { - let mut ctx = make_test_ctx(); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let argument_types = (0..4) - .map(|_| i32_ty.into()) - .chain((0..5).map(|_| u32_ty.into())) - .collect(); - let (module_ptr, entry) = build_test_kernel(&mut ctx, argument_types); - - let selector_value = (index % 2) as u32; - let selector_op = Operation::new( - &mut ctx, - mir::MirConstantOp::get_concrete_op_info(), - vec![u32_ty.into()], - vec![], - vec![], - 0, - ); - mir::MirConstantOp::new(selector_op).set_attr_value( - &ctx, - IntegerAttr::new( - u32_ty, - APInt::from_u32(selector_value, NonZeroUsize::new(32).unwrap()), - ), - ); - selector_op.insert_at_back(entry, &ctx); - let selector = selector_op.deref(&ctx).get_result(0); - - let operands = (0..9) - .map(|operand| entry.deref(&ctx).get_argument(operand)) - .chain(std::iter::once(selector)) - .collect(); - let operation = Operation::new( - &mut ctx, - nvvm::SparseMmaOp::get_concrete_op_info(), - vec![i32_ty.into(); 4], - operands, - vec![], - 0, - ); - let mma = nvvm::SparseMmaOp::new(operation); - mma.set_attr_nvvm_sparse_mma_shape(&ctx, nvvm::SparseMmaShapeAttr::M16n8k64); - mma.set_attr_nvvm_sparse_mma_accumulator(&ctx, nvvm::SparseMmaAccumulatorAttr::S32); - mma.set_attr_nvvm_sparse_mma_a_element(&ctx, a_element.clone()); - mma.set_attr_nvvm_sparse_mma_b_element(&ctx, b_element.clone()); - mma.set_attr_nvvm_sparse_mma_a_layout(&ctx, nvvm::SparseMmaLayoutAttr::Row); - mma.set_attr_nvvm_sparse_mma_b_layout(&ctx, nvvm::SparseMmaLayoutAttr::Col); - mma.set_attr_nvvm_sparse_mma_overflow(&ctx, overflow.clone()); - mma.set_attr_nvvm_sparse_mma_metadata(&ctx, metadata.clone()); - mma.set_attr_nvvm_sparse_mma_selector( - &ctx, - nvvm::SparseMmaSelectorAttr::ImmediateZeroOrOne, - ); - operation.insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm_with_options( - &mut ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: backend, - ..Default::default() - }, - ) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let body = lowered_kernel_body(&ctx, module_ptr); - let lowered = body - .iter() - .filter_map(|op| Operation::get_op::(*op, &ctx)) - .collect::>(); - assert_eq!(lowered.len(), 1, "{backend:?}"); - let asm = &lowered[0]; - let expected_template = format!( - "mma.{metadata_name}.sync.aligned.m16n8k64.row.col{overflow_name}.s32.{a_name}.{b_name}.s32 {{$0, $1, $2, $3}}, {{$8, $9}}, {{$10, $11}}, {{$4, $5, $6, $7}}, $12, $13;" - ); - assert_eq!( - asm.get_attr_inline_asm_template(&ctx) - .as_deref() - .map(|value| String::from(value.clone())), - Some(expected_template) - ); - assert_eq!( - asm.get_attr_inline_asm_constraints(&ctx) - .as_deref() - .map(|value| String::from(value.clone())), - Some("=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,n".to_string()) - ); - assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); - let asm_operation = asm.get_operation().deref(&ctx); - assert_eq!(asm_operation.get_num_operands(), 10); - assert_eq!(asm_operation.get_num_results(), 1); - let lowered_selector = asm_operation.get_operand(9); - let defining_op = lowered_selector - .defining_op() - .expect("sparse MMA selector remains an LLVM constant"); - let constant = Operation::get_op::(defining_op, &ctx) - .expect("sparse MMA selector remains an LLVM integer constant"); - let attribute = constant.get_value(&ctx); - let integer = attribute - .downcast_ref::() - .expect("sparse MMA selector is an integer"); - assert_eq!(integer.value().to_u64(), selector_value as u64); - - let module = Operation::get_op::(module_ptr, &ctx).unwrap(); - let ir = llvm_export::export::export_module_to_string(&ctx, &module) - .expect("generated sparse MMA exports to LLVM IR"); - assert!(ir.contains("asm sideeffect"), "{ir}"); - assert!(ir.contains("{ convergent }"), "{ir}"); - } - } - } - - Ok(()) -} - -#[test] -fn test_generated_sparse_mma_m16n8k128_int4_lowers_both_metadata_modes() -> Result<(), anyhow::Error> -{ - use pliron::builtin::attributes::IntegerAttr; - use pliron::builtin::types::{IntegerType, Signedness}; - use pliron::utils::apint::APInt; - use std::num::NonZeroUsize; - - let cases = [ - ( - "s4", - "s4", - "", - nvvm::SparseMmaElementAttr::S4, - nvvm::SparseMmaElementAttr::S4, - nvvm::SparseMmaOverflowAttr::Wrapping, - ), - ( - "s4", - "u4", - "", - nvvm::SparseMmaElementAttr::S4, - nvvm::SparseMmaElementAttr::U4, - nvvm::SparseMmaOverflowAttr::Wrapping, - ), - ( - "u4", - "u4", - "", - nvvm::SparseMmaElementAttr::U4, - nvvm::SparseMmaElementAttr::U4, - nvvm::SparseMmaOverflowAttr::Wrapping, - ), - ( - "u4", - "s4", - "", - nvvm::SparseMmaElementAttr::U4, - nvvm::SparseMmaElementAttr::S4, - nvvm::SparseMmaOverflowAttr::Wrapping, - ), - ( - "s4", - "s4", - ".satfinite", - nvvm::SparseMmaElementAttr::S4, - nvvm::SparseMmaElementAttr::S4, - nvvm::SparseMmaOverflowAttr::Satfinite, - ), - ( - "s4", - "u4", - ".satfinite", - nvvm::SparseMmaElementAttr::S4, - nvvm::SparseMmaElementAttr::U4, - nvvm::SparseMmaOverflowAttr::Satfinite, - ), - ( - "u4", - "u4", - ".satfinite", - nvvm::SparseMmaElementAttr::U4, - nvvm::SparseMmaElementAttr::U4, - nvvm::SparseMmaOverflowAttr::Satfinite, - ), - ( - "u4", - "s4", - ".satfinite", - nvvm::SparseMmaElementAttr::U4, - nvvm::SparseMmaElementAttr::S4, - nvvm::SparseMmaOverflowAttr::Satfinite, - ), - ]; - - for backend in [ - mir_lower::IntrinsicBackend::LlvmNvptx, - mir_lower::IntrinsicBackend::LibNvvm, - ] { - for (metadata_name, metadata) in [ - ("sp", nvvm::SparseMmaMetadataAttr::Standard), - ("sp::ordered_metadata", nvvm::SparseMmaMetadataAttr::Ordered), - ] { - for (a_name, b_name, overflow_name, a_element, b_element, overflow) in &cases { - let mut ctx = make_test_ctx(); - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let argument_types = (0..4) - .map(|_| i32_ty.into()) - .chain((0..9).map(|_| u32_ty.into())) - .collect(); - let (module_ptr, entry) = build_test_kernel(&mut ctx, argument_types); - - let selector_op = Operation::new( - &mut ctx, - mir::MirConstantOp::get_concrete_op_info(), - vec![u32_ty.into()], - vec![], - vec![], - 0, - ); - mir::MirConstantOp::new(selector_op).set_attr_value( - &ctx, - IntegerAttr::new(u32_ty, APInt::from_u32(0, NonZeroUsize::new(32).unwrap())), - ); - selector_op.insert_at_back(entry, &ctx); - let selector = selector_op.deref(&ctx).get_result(0); - - let operands = (0..13) - .map(|operand| entry.deref(&ctx).get_argument(operand)) - .chain(std::iter::once(selector)) - .collect(); - let operation = Operation::new( - &mut ctx, - nvvm::SparseMmaOp::get_concrete_op_info(), - vec![i32_ty.into(); 4], - operands, - vec![], - 0, - ); - let mma = nvvm::SparseMmaOp::new(operation); - mma.set_attr_nvvm_sparse_mma_shape(&ctx, nvvm::SparseMmaShapeAttr::M16n8k128); - mma.set_attr_nvvm_sparse_mma_accumulator(&ctx, nvvm::SparseMmaAccumulatorAttr::S32); - mma.set_attr_nvvm_sparse_mma_a_element(&ctx, a_element.clone()); - mma.set_attr_nvvm_sparse_mma_b_element(&ctx, b_element.clone()); - mma.set_attr_nvvm_sparse_mma_a_layout(&ctx, nvvm::SparseMmaLayoutAttr::Row); - mma.set_attr_nvvm_sparse_mma_b_layout(&ctx, nvvm::SparseMmaLayoutAttr::Col); - mma.set_attr_nvvm_sparse_mma_overflow(&ctx, overflow.clone()); - mma.set_attr_nvvm_sparse_mma_metadata(&ctx, metadata.clone()); - mma.set_attr_nvvm_sparse_mma_selector( - &ctx, - nvvm::SparseMmaSelectorAttr::ImmediateZero, - ); - operation.insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm_with_options( - &mut ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: backend, - ..Default::default() - }, - ) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let body = lowered_kernel_body(&ctx, module_ptr); - let lowered = body - .iter() - .filter_map(|op| Operation::get_op::(*op, &ctx)) - .collect::>(); - assert_eq!(lowered.len(), 1, "{backend:?}"); - let asm = &lowered[0]; - let expected_template = format!( - "mma.{metadata_name}.sync.aligned.m16n8k128.row.col{overflow_name}.s32.{a_name}.{b_name}.s32 {{$0, $1, $2, $3}}, {{$8, $9, $10, $11}}, {{$12, $13, $14, $15}}, {{$4, $5, $6, $7}}, $16, $17;" - ); - assert_eq!( - asm.get_attr_inline_asm_template(&ctx) - .as_deref() - .map(|value| String::from(value.clone())), - Some(expected_template) - ); - assert_eq!( - asm.get_attr_inline_asm_constraints(&ctx) - .as_deref() - .map(|value| String::from(value.clone())), - Some("=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,r,r,r,r,n".to_string()) - ); - assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); - let asm_operation = asm.get_operation().deref(&ctx); - assert_eq!(asm_operation.get_num_operands(), 14); - assert_eq!(asm_operation.get_num_results(), 1); - let lowered_selector = asm_operation.get_operand(13); - let defining_op = lowered_selector - .defining_op() - .expect("sparse MMA selector remains an LLVM constant"); - let constant = Operation::get_op::(defining_op, &ctx) - .expect("sparse MMA selector remains an LLVM integer constant"); - let attribute = constant.get_value(&ctx); - let integer = attribute - .downcast_ref::() - .expect("sparse MMA selector is an integer"); - assert_eq!(integer.value().to_u64(), 0); - - let module = Operation::get_op::(module_ptr, &ctx).unwrap(); - let ir = llvm_export::export::export_module_to_string(&ctx, &module) - .expect("generated sparse MMA exports to LLVM IR"); - assert!(ir.contains("asm sideeffect"), "{ir}"); - assert!(ir.contains("{ convergent }"), "{ir}"); - } - } - } - - Ok(()) -} - -// --------------------------------------------------------------------------- -// mma.sync m16n8k16 bf16 intrinsic lowering test -// --------------------------------------------------------------------------- - -#[test] -fn test_mma_m16n8k16_f32_bf16_lowers_to_inline_asm() -> Result<(), anyhow::Error> { - let mut ctx = make_test_ctx(); - let f32_ty = pliron::builtin::types::FP32Type::get(&ctx); - let i32_ty = pliron::builtin::types::IntegerType::get( - &ctx, - 32, - pliron::builtin::types::Signedness::Signless, - ); - let argument_types = (0..4) - .map(|_| f32_ty.into()) - .chain((0..6).map(|_| i32_ty.into())) - .collect(); - let (module_ptr, entry) = build_test_kernel(&mut ctx, argument_types); - let operands = (0..10) - .map(|index| entry.deref(&ctx).get_argument(index)) - .collect(); - - let op = Operation::new( - &mut ctx, - nvvm::MmaM16N8K16F32Bf16Op::get_concrete_op_info(), - vec![f32_ty.into(); 4], - operands, - vec![], - 0, - ); - op.insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let mut found = 0; - let module_region = module_ptr.deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - for module_op in module_block.deref(&ctx).iter(&ctx) { - let Some(function) = Operation::get_op::(module_op, &ctx) else { - continue; - }; - if function.get_symbol_name(&ctx).to_string() != "kernel_func" { - continue; - } - let body = function.get_operation().deref(&ctx).get_region(0); - for block in body.deref(&ctx).iter(&ctx) { - for body_op in block.deref(&ctx).iter(&ctx) { - let Some(asm) = Operation::get_op::(body_op, &ctx) else { - continue; - }; - let template = asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())); - let constraints = asm - .get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())); - if !template.as_deref().is_some_and(|t| { - t.contains("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32") - }) { - continue; - } - found += 1; - let template = template.expect("MMA inline asm must have a template"); - assert_eq!( - template, - concat!( - "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 ", - "{$0, $1, $2, $3}, ", - "{$8, $9, $10, $11}, ", - "{$12, $13}, ", - "{$4, $5, $6, $7};" - ) - ); - for forbidden in [".reg", "ld.", "st.", "["] { - assert!( - !template.contains(forbidden), - "register-only MMA must not contain {forbidden:?}: {template}" - ); - } - assert_eq!( - constraints.as_deref(), - Some("=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r") - ); - assert_eq!( - llvm::asm_kind_opt(&ctx, &asm), - Some(llvm::AsmKind::Convergent) - ); - assert_eq!( - body_op.deref(&ctx).get_num_operands(), - 10, - "expected C, A, and B scalar register operands" - ); - assert_eq!( - body_op.deref(&ctx).get_num_results(), - 1, - "LLVM inline asm returns the four D registers as one struct" - ); - } - } - } - - assert_eq!(found, 1, "expected one mma.sync inline-asm operation"); - Ok(()) -} - -fn build_wgmma_pointer_test_kernel( - ctx: &mut Context, - accumulator_count: usize, - trailing_arg_types: Vec, -) -> ( - pliron::context::Ptr, - pliron::context::Ptr, - Vec, - pliron::value::Value, - pliron::value::Value, - Vec, -) { - use dialect_mir::types::MirPtrType; - use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; - use pliron::r#type::TypeHandle; - - let f32_ty = FP32Type::get(ctx); - let accumulator_ptr_ty: TypeHandle = MirPtrType::get_generic(ctx, f32_ty.into(), true).into(); - let u64_ty: TypeHandle = IntegerType::get(ctx, 64, Signedness::Unsigned).into(); - - let trailing_count = trailing_arg_types.len(); - let mut argument_types = vec![accumulator_ptr_ty; accumulator_count]; - argument_types.push(u64_ty); - argument_types.push(u64_ty); - argument_types.extend(trailing_arg_types); - - let (module_ptr, entry) = build_test_kernel(ctx, argument_types); - - let accumulators = (0..accumulator_count) - .map(|index| entry.deref(ctx).get_argument(index)) - .collect::>(); - let desc_a = entry.deref(ctx).get_argument(accumulator_count); - let desc_b = entry.deref(ctx).get_argument(accumulator_count + 1); - let trailing_arguments = (0..trailing_count) - .map(|index| entry.deref(ctx).get_argument(accumulator_count + 2 + index)) - .collect::>(); - - ( - module_ptr, - entry, - accumulators, - desc_a, - desc_b, - trailing_arguments, - ) -} - -fn build_wgmma_canonical_pointer_test_kernel_with_rows( - ctx: &mut Context, - accumulator_rows: u64, - accumulator_count: usize, - descriptor_count: usize, -) -> ( - pliron::context::Ptr, - pliron::context::Ptr, - Vec, - Vec, -) { - use dialect_mir::types::{MirArrayType, MirPtrType}; - use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; - - let f32_ty = FP32Type::get(ctx); - let row_ty = MirArrayType::get(ctx, f32_ty.into(), 8); - let accumulator_ty = MirArrayType::get(ctx, row_ty.into(), accumulator_rows); - let accumulator_ptr_ty = MirPtrType::get_generic(ctx, accumulator_ty.into(), true); - let u64_ty = IntegerType::get(ctx, 64, Signedness::Unsigned); - - let accumulator_arg_ty: pliron::r#type::TypeHandle = accumulator_ptr_ty.into(); - let u64_arg_ty: pliron::r#type::TypeHandle = u64_ty.into(); - let mut argument_types = vec![accumulator_arg_ty; accumulator_count]; - argument_types.extend(vec![u64_arg_ty; descriptor_count]); - let (module_ptr, entry) = build_test_kernel(ctx, argument_types); - - let accumulators = (0..accumulator_count) - .map(|index| entry.deref(ctx).get_argument(index)) - .collect::>(); - let descriptors = (0..descriptor_count) - .map(|index| entry.deref(ctx).get_argument(accumulator_count + index)) - .collect::>(); - - (module_ptr, entry, accumulators, descriptors) -} - -fn build_wgmma_canonical_pointer_test_kernel( - ctx: &mut Context, - accumulator_count: usize, - descriptor_count: usize, -) -> ( - pliron::context::Ptr, - pliron::context::Ptr, - Vec, - Vec, -) { - build_wgmma_canonical_pointer_test_kernel_with_rows(ctx, 4, accumulator_count, descriptor_count) -} - -fn build_wgmma_m64n128_canonical_pointer_test_kernel( - ctx: &mut Context, - accumulator_count: usize, - descriptor_count: usize, -) -> ( - pliron::context::Ptr, - pliron::context::Ptr, - Vec, - Vec, -) { - build_wgmma_canonical_pointer_test_kernel_with_rows(ctx, 8, accumulator_count, descriptor_count) -} - -fn append_pointer_wgmma_mma( - ctx: &mut Context, - block: pliron::context::Ptr, - accumulator: pliron::value::Value, - desc_a: pliron::value::Value, - desc_b: pliron::value::Value, -) { - Operation::new( - ctx, - nvvm::WgmmaMmaM64N64K16F32Bf16Op::get_concrete_op_info(), - vec![], - vec![accumulator, desc_a, desc_b], - vec![], - 0, - ) - .insert_at_back(block, ctx); -} - -fn append_pointer_wgmma_mma_m64n128( - ctx: &mut Context, - block: pliron::context::Ptr, - accumulator: pliron::value::Value, - desc_a: pliron::value::Value, - desc_b: pliron::value::Value, -) { - Operation::new( - ctx, - nvvm::WgmmaMmaM64N128K16F32Bf16Op::get_concrete_op_info(), - vec![], - vec![accumulator, desc_a, desc_b], - vec![], - 0, - ) - .insert_at_back(block, ctx); -} - -fn append_pointer_wgmma_mma_f16( - ctx: &mut Context, - block: pliron::context::Ptr, - accumulator: pliron::value::Value, - desc_a: pliron::value::Value, - desc_b: pliron::value::Value, -) { - Operation::new( - ctx, - nvvm::WgmmaMmaM64N64K16F32F16Op::get_concrete_op_info(), - vec![], - vec![accumulator, desc_a, desc_b], - vec![], - 0, - ) - .insert_at_back(block, ctx); -} - -fn append_pointer_wgmma_mma_tf32( - ctx: &mut Context, - block: pliron::context::Ptr, - accumulator: pliron::value::Value, - desc_a: pliron::value::Value, - desc_b: pliron::value::Value, -) { - Operation::new( - ctx, - nvvm::WgmmaMmaM64N64K8F32Tf32Op::get_concrete_op_info(), - vec![], - vec![accumulator, desc_a, desc_b], - vec![], - 0, - ) - .insert_at_back(block, ctx); -} - -fn append_wgmma_wait_group_constant( - ctx: &mut Context, - block: pliron::context::Ptr, - pending_groups: i64, -) { - use pliron::builtin::attributes::IntegerAttr; - use pliron::builtin::types::{IntegerType, Signedness}; - use pliron::utils::apint::APInt; - use std::num::NonZeroUsize; - - let u64_ty = IntegerType::get(ctx, 64, Signedness::Unsigned); - let constant_attr = IntegerAttr::new( - u64_ty, - APInt::from_i64(pending_groups, NonZeroUsize::new(64).unwrap()), - ); - let constant = Operation::new( - ctx, - mir::MirConstantOp::get_concrete_op_info(), - vec![u64_ty.into()], - vec![], - vec![], - 0, - ); - mir::MirConstantOp::new(constant).set_attr_value(ctx, constant_attr); - constant.insert_at_back(block, ctx); - - let value = constant.deref(ctx).get_result(0); - nvvm::WgmmaWaitGroupSyncAlignedOp::build(ctx, value).insert_at_back(block, ctx); -} - -fn append_mir_unsigned_constant( - ctx: &mut Context, - block: pliron::context::Ptr, - ty: pliron::r#type::TypedHandle, - value: u64, -) -> pliron::value::Value { - use pliron::builtin::attributes::IntegerAttr; - use pliron::utils::apint::APInt; - use std::num::NonZeroUsize; - - let width = usize::try_from(ty.deref(ctx).width()).expect("integer width must fit usize"); - let constant = Operation::new( - ctx, - mir::MirConstantOp::get_concrete_op_info(), - vec![ty.into()], - vec![], - vec![], - 0, - ); - mir::MirConstantOp::new(constant).set_attr_value( - ctx, - IntegerAttr::new( - ty, - APInt::from_u64( - value, - NonZeroUsize::new(width).expect("nonzero integer width"), - ), - ), - ); - constant.insert_at_back(block, ctx); - constant.deref(ctx).get_result(0) -} - -fn assert_wgmma_lowering_rejected( - ctx: &mut Context, - module_ptr: pliron::context::Ptr, - expected_diagnostic: &str, -) { - let error = mir_lower::lower_mir_to_llvm(ctx, module_ptr) - .expect_err("invalid deferred WGMMA sequence must fail closed") - .to_string(); - - assert!( - error.contains(expected_diagnostic), - "expected diagnostic containing `{expected_diagnostic}`, got:\n{error}" - ); -} - -// --------------------------------------------------------------------------- -// mma.sync m16n8k16 f16 intrinsic lowering test -// --------------------------------------------------------------------------- - -#[test] -fn test_mma_m16n8k16_f32_f16_lowers_to_inline_asm() -> Result<(), anyhow::Error> { - let mut ctx = make_test_ctx(); - let f32_ty = pliron::builtin::types::FP32Type::get(&ctx); - let i32_ty = pliron::builtin::types::IntegerType::get( - &ctx, - 32, - pliron::builtin::types::Signedness::Signless, - ); - let argument_types = (0..4) - .map(|_| f32_ty.into()) - .chain((0..6).map(|_| i32_ty.into())) - .collect(); - let (module_ptr, entry) = build_test_kernel(&mut ctx, argument_types); - let operands = (0..10) - .map(|index| entry.deref(&ctx).get_argument(index)) - .collect(); - - let op = Operation::new( - &mut ctx, - nvvm::MmaM16N8K16F32F16Op::get_concrete_op_info(), - vec![f32_ty.into(); 4], - operands, - vec![], - 0, - ); - op.insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let mut found = 0; - let module_region = module_ptr.deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - for module_op in module_block.deref(&ctx).iter(&ctx) { - let Some(function) = Operation::get_op::(module_op, &ctx) else { - continue; - }; - if function.get_symbol_name(&ctx).to_string() != "kernel_func" { - continue; - } - let body = function.get_operation().deref(&ctx).get_region(0); - for block in body.deref(&ctx).iter(&ctx) { - for body_op in block.deref(&ctx).iter(&ctx) { - let Some(asm) = Operation::get_op::(body_op, &ctx) else { - continue; - }; - let template = asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())); - let constraints = asm - .get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())); - if !template.as_deref().is_some_and(|t| { - t.contains("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32") - }) { - continue; - } - found += 1; - let template = template.expect("MMA inline asm must have a template"); - assert_eq!( - template, - concat!( - "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 ", - "{$0, $1, $2, $3}, ", - "{$8, $9, $10, $11}, ", - "{$12, $13}, ", - "{$4, $5, $6, $7};" - ) - ); - for forbidden in [".reg", "ld.", "st.", "["] { - assert!( - !template.contains(forbidden), - "register-only MMA must not contain {forbidden:?}: {template}" - ); - } - assert_eq!( - constraints.as_deref(), - Some("=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r") - ); - assert_eq!( - llvm::asm_kind_opt(&ctx, &asm), - Some(llvm::AsmKind::Convergent) - ); - assert_eq!( - body_op.deref(&ctx).get_num_operands(), - 10, - "expected C, A, and B scalar register operands" - ); - assert_eq!( - body_op.deref(&ctx).get_num_results(), - 1, - "LLVM inline asm returns the four D registers as one struct" - ); - } - } - } - - assert_eq!(found, 1, "expected one mma.sync inline-asm operation"); - Ok(()) -} - -#[test] -fn test_mma_m16n8k8_f32_tf32_lowers_to_inline_asm() -> Result<(), anyhow::Error> { - let mut ctx = make_test_ctx(); - let f32_ty = pliron::builtin::types::FP32Type::get(&ctx); - let i32_ty = pliron::builtin::types::IntegerType::get( - &ctx, - 32, - pliron::builtin::types::Signedness::Signless, - ); - let argument_types = (0..4) - .map(|_| f32_ty.into()) - .chain((0..6).map(|_| i32_ty.into())) - .collect(); - let (module_ptr, entry) = build_test_kernel(&mut ctx, argument_types); - let operands = (0..10) - .map(|index| entry.deref(&ctx).get_argument(index)) - .collect(); - - let op = Operation::new( - &mut ctx, - nvvm::MmaM16N8K8F32Tf32Op::get_concrete_op_info(), - vec![f32_ty.into(); 4], - operands, - vec![], - 0, - ); - op.insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let mut found = 0; - let module_region = module_ptr.deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - for module_op in module_block.deref(&ctx).iter(&ctx) { - let Some(function) = Operation::get_op::(module_op, &ctx) else { - continue; - }; - if function.get_symbol_name(&ctx).to_string() != "kernel_func" { - continue; - } - let body = function.get_operation().deref(&ctx).get_region(0); - for block in body.deref(&ctx).iter(&ctx) { - for body_op in block.deref(&ctx).iter(&ctx) { - let Some(asm) = Operation::get_op::(body_op, &ctx) else { - continue; - }; - let template = asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())); - let constraints = asm - .get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())); - if !template.as_deref().is_some_and(|t| { - t.contains("mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32") - }) { - continue; - } - found += 1; - let template = template.expect("MMA inline asm must have a template"); - assert_eq!( - template, - concat!( - "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 ", - "{$0, $1, $2, $3}, ", - "{$8, $9, $10, $11}, ", - "{$12, $13}, ", - "{$4, $5, $6, $7};" - ) - ); - for forbidden in [".reg", "ld.", "st.", "["] { - assert!( - !template.contains(forbidden), - "register-only MMA must not contain {forbidden:?}: {template}" - ); - } - assert_eq!( - constraints.as_deref(), - Some("=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r") - ); - assert_eq!( - llvm::asm_kind_opt(&ctx, &asm), - Some(llvm::AsmKind::Convergent) - ); - assert_eq!( - body_op.deref(&ctx).get_num_operands(), - 10, - "expected C, A, and B scalar register operands" - ); - assert_eq!( - body_op.deref(&ctx).get_num_results(), - 1, - "LLVM inline asm returns the four D registers as one struct" - ); - } - } - } - - assert_eq!(found, 1, "expected one mma.sync inline-asm operation"); - Ok(()) -} - -#[test] -fn test_packed_atomic_add_lowers_to_exact_side_effecting_ptx() -> Result<(), anyhow::Error> { - use dialect_mir::types::MirPtrType; - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let ptr_ty = MirPtrType::get_generic(&mut ctx, u32_ty.into(), true); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![ptr_ty.into(), u32_ty.into()]); - let address = entry.deref(&ctx).get_argument(0); - let addend = entry.deref(&ctx).get_argument(1); - - for op_info in [ - nvvm::NvvmAtomAddF16x2Op::get_concrete_op_info(), - nvvm::NvvmAtomAddBf16x2Op::get_concrete_op_info(), - ] { - Operation::new( - &mut ctx, - op_info, - vec![u32_ty.into()], - vec![address, addend], - vec![], - 0, - ) - .insert_at_back(entry, &ctx); - } - for format in [ - nvvm::PackedAtomicFormatAttr::F16x2, - nvvm::PackedAtomicFormatAttr::Bf16x2, - ] { - nvvm::PackedAtomicAddOp::build(&mut ctx, address, addend, format) - .insert_at_back(entry, &ctx); - } - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let expected = [ - "atom.global.add.noftz.f16x2 $0, [$1], $2;", - "atom.global.add.noftz.bf16x2 $0, [$1], $2;", - ]; - let module_region = module_ptr.deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - let mut lowered = Vec::new(); - - for op in module_block.deref(&ctx).iter(&ctx) { - let Some(function) = Operation::get_op::(op, &ctx) else { - continue; - }; - if function.get_symbol_name(&ctx).to_string() != "kernel_func" { - continue; - } - let body = function.get_operation().deref(&ctx).get_region(0); - for block in body.deref(&ctx).iter(&ctx) { - for body_op in block.deref(&ctx).iter(&ctx) { - let Some(asm) = Operation::get_op::(body_op, &ctx) else { - continue; - }; - let template = asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .unwrap_or_default(); - assert!( - !template.contains("atom.cas"), - "packed atomic exact-native lowering must not use a CAS loop: {template}" - ); - if !template.starts_with("atom.global.add.noftz.") { - continue; - } - lowered.push(( - template, - asm.get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .unwrap_or_default(), - llvm::asm_kind(&ctx, &asm), - body_op.deref(&ctx).get_num_operands(), - body_op.deref(&ctx).get_num_results(), - )); - } - } - } - - assert_eq!(lowered.len(), expected.len() * 2); - for instruction in expected { - let matches: Vec<_> = lowered - .iter() - .filter(|(template, _, _, _, _)| template == instruction) - .collect(); - assert_eq!( - matches.len(), - 2, - "legacy and generated paths must have exact PTX parity for {instruction}" - ); - for (_, constraints, kind, operands, results) in matches { - assert_eq!(constraints, "=r,l,r,~{memory}"); - assert_eq!(*kind, llvm::AsmKind::SideEffect); - assert_eq!(*operands, 2); - assert_eq!(*results, 1); - } - } - - Ok(()) -} - -#[test] -fn test_generated_packed_atomic_add_libnvvm_route_is_exact() -> Result<(), anyhow::Error> { - use dialect_mir::types::MirPtrType; - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let ptr_ty = MirPtrType::get_generic(&mut ctx, u32_ty.into(), true); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![ptr_ty.into(), u32_ty.into()]); - let address = entry.deref(&ctx).get_argument(0); - let addend = entry.deref(&ctx).get_argument(1); - for format in [ - nvvm::PackedAtomicFormatAttr::F16x2, - nvvm::PackedAtomicFormatAttr::Bf16x2, - ] { - nvvm::PackedAtomicAddOp::build(&mut ctx, address, addend, format) - .insert_at_back(entry, &ctx); - } - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm_with_options( - &mut ctx, - module_ptr, - mir_lower::LoweringOptions { - intrinsic_backend: mir_lower::IntrinsicBackend::LibNvvm, - ..Default::default() - }, - ) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let lowered = lowered_kernel_body(&ctx, module_ptr) - .into_iter() - .filter_map(|op| { - let asm = Operation::get_op::(op, &ctx)?; - let template = asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone()))?; - template.starts_with("atom.global.add.noftz.").then(|| { - ( - template, - asm.get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())), - llvm::asm_kind(&ctx, &asm), - ) - }) - }) - .collect::>(); - assert_eq!(lowered.len(), 2); - assert!( - lowered - .iter() - .any(|(template, _, _)| { template == "atom.global.add.noftz.f16x2 $0, [$1], $2;" }) - ); - assert!( - lowered - .iter() - .any(|(template, _, _)| { template == "atom.global.add.noftz.bf16x2 $0, [$1], $2;" }) - ); - for (_, constraints, kind) in lowered { - assert_eq!(constraints.as_deref(), Some("=r,l,r,~{memory}")); - assert_eq!(kind, llvm::AsmKind::SideEffect); - } - Ok(()) -} - -#[test] -fn test_mma_m8n8k4_f64_lowers_to_inline_asm() -> Result<(), anyhow::Error> { - use pliron::builtin::types::FP64Type; - - let mut ctx = make_test_ctx(); - let f64_ty = FP64Type::get(&ctx); - let (module_ptr, entry) = build_test_kernel( - &mut ctx, - vec![f64_ty.into(), f64_ty.into(), f64_ty.into(), f64_ty.into()], - ); - let operands = (0..4) - .map(|index| entry.deref(&ctx).get_argument(index)) - .collect(); - - let op = Operation::new( - &mut ctx, - nvvm::MmaM8N8K4F64Op::get_concrete_op_info(), - vec![f64_ty.into(), f64_ty.into()], - operands, - vec![], - 0, - ); - op.insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let mut found = 0; - let module_region = module_ptr.deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - for module_op in module_block.deref(&ctx).iter(&ctx) { - let Some(function) = Operation::get_op::(module_op, &ctx) else { - continue; - }; - if function.get_symbol_name(&ctx).to_string() != "kernel_func" { - continue; - } - let body = function.get_operation().deref(&ctx).get_region(0); - for block in body.deref(&ctx).iter(&ctx) { - for body_op in block.deref(&ctx).iter(&ctx) { - let Some(asm) = Operation::get_op::(body_op, &ctx) else { - continue; - }; - let template = asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())); - let constraints = asm - .get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())); - if !template - .as_deref() - .is_some_and(|t| t.contains("mma.sync.aligned.m8n8k4.row.col.f64.f64.f64.f64")) - { - continue; - } - found += 1; - let template = template.unwrap(); - assert_eq!( - template, - "mma.sync.aligned.m8n8k4.row.col.f64.f64.f64.f64 {$0, $1}, {$4}, {$5}, {$2, $3};" - ); - assert!(!template.contains(".reg")); - assert!(!template.contains("ld.")); - assert!(!template.contains("st.")); - assert_eq!(constraints.as_deref(), Some("=d,=d,d,d,d,d")); - assert_eq!( - llvm::asm_kind_opt(&ctx, &asm), - Some(llvm::AsmKind::Convergent) - ); - assert_eq!( - body_op.deref(&ctx).get_num_operands(), - 4, - "expected four register inputs (c0, c1, a, b)" - ); - assert_eq!( - body_op.deref(&ctx).get_num_results(), - 1, - "LLVM inline asm returns one aggregate containing d0 and d1" - ); - } - } - } - - assert_eq!(found, 1, "expected one mma.sync inline-asm operation"); - Ok(()) -} - -#[test] -fn test_mma_m16n8k32_s32_s8_lowers_to_inline_asm() -> Result<(), anyhow::Error> { - let mut ctx = make_test_ctx(); - let i32_ty = pliron::builtin::types::IntegerType::get( - &ctx, - 32, - pliron::builtin::types::Signedness::Signless, - ); - let argument_types = (0..10).map(|_| i32_ty.into()).collect(); - let (module_ptr, entry) = build_test_kernel(&mut ctx, argument_types); - let operands = (0..10) - .map(|index| entry.deref(&ctx).get_argument(index)) - .collect(); - - let op = Operation::new( - &mut ctx, - nvvm::MmaM16N8K32S32S8Op::get_concrete_op_info(), - vec![i32_ty.into(); 4], - operands, - vec![], - 0, - ); - op.insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let mut found = 0; - let module_region = module_ptr.deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - for module_op in module_block.deref(&ctx).iter(&ctx) { - let Some(function) = Operation::get_op::(module_op, &ctx) else { - continue; - }; - if function.get_symbol_name(&ctx).to_string() != "kernel_func" { - continue; - } - let body = function.get_operation().deref(&ctx).get_region(0); - for block in body.deref(&ctx).iter(&ctx) { - for body_op in block.deref(&ctx).iter(&ctx) { - let Some(asm) = Operation::get_op::(body_op, &ctx) else { - continue; - }; - let template = asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())); - let constraints = asm - .get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())); - if !template - .as_deref() - .is_some_and(|t| t.contains("mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32")) - { - continue; - } - found += 1; - let template = template.expect("MMA inline asm must have a template"); - assert_eq!( - template, - concat!( - "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 ", - "{$0, $1, $2, $3}, ", - "{$8, $9, $10, $11}, ", - "{$12, $13}, ", - "{$4, $5, $6, $7};" - ) - ); - for forbidden in [".reg", "ld.", "st.", "["] { - assert!( - !template.contains(forbidden), - "register-only MMA must not contain {forbidden:?}: {template}" - ); - } - assert_eq!( - constraints.as_deref(), - Some("=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,r") - ); - assert_eq!( - llvm::asm_kind_opt(&ctx, &asm), - Some(llvm::AsmKind::Convergent) - ); - assert_eq!( - body_op.deref(&ctx).get_num_operands(), - 10, - "expected C, A, and B scalar register operands" - ); - assert_eq!( - body_op.deref(&ctx).get_num_results(), - 1, - "LLVM inline asm returns the four D registers as one struct" - ); - } - } - } - - assert_eq!(found, 1, "expected one mma.sync inline-asm operation"); - Ok(()) -} - -#[test] -fn test_deferred_wgmma_group_lowers_to_one_register_lifetime_scope() -> Result<(), anyhow::Error> { - use dialect_mir::types::MirPtrType; - use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let f32_ty = FP32Type::get(&ctx); - let accumulator_ptr_ty = MirPtrType::get_generic(&mut ctx, f32_ty.into(), true); - let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let (module_ptr, entry) = build_test_kernel( - &mut ctx, - vec![accumulator_ptr_ty.into(), u64_ty.into(), u64_ty.into()], - ); - let accumulator = entry.deref(&ctx).get_argument(0); - let desc_a = entry.deref(&ctx).get_argument(1); - let desc_b = entry.deref(&ctx).get_argument(2); - - nvvm::WgmmaMmaGroupM64N64K16F32Bf16Op::build(&mut ctx, accumulator, vec![desc_a, desc_b]) - .insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let matching = lowered_kernel_body(&ctx, module_ptr) - .into_iter() - .filter_map(|operation| Operation::get_op::(operation, &ctx)) - .filter(|asm| { - asm.get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .is_some_and(|template| { - template.contains("wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16") - }) - }) - .collect::>(); - assert_eq!(matching.len(), 1); - - let asm = &matching[0]; - let template = asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .expect("WGMMA template"); - assert_eq!(template.matches("ld.f32 %acc").count(), 32); - assert_eq!(template.matches("st.f32 [$0").count(), 32); - assert_eq!(template.matches("wgmma.mma_async").count(), 1); - let wait = template.find("wgmma.wait_group.sync.aligned 0").unwrap(); - let first_store = template.find("st.f32 [$0").unwrap(); - assert!( - wait < first_store, - "accumulator stores must follow wait_group<0>" - ); - assert_eq!( - asm.get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some("l,l,l,~{memory}") - ); - assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); - Ok(()) -} - -#[test] -fn test_value_form_wgmma_group_lowers_to_tied_register_inline_ptx() -> Result<(), anyhow::Error> { - use llvm_export::types as llvm_types; - use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; - use pliron::r#type::Typed; - - const ACCUMULATOR_LEN: usize = 32; - const DESCRIPTOR_COUNT: usize = 4; - - let mut ctx = make_test_ctx(); - let f32_ty = FP32Type::get(&ctx); - let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - - let argument_types = (0..ACCUMULATOR_LEN) - .map(|_| f32_ty.into()) - .chain((0..DESCRIPTOR_COUNT).map(|_| u64_ty.into())) - .collect::>(); - - let (module_ptr, entry) = build_test_kernel(&mut ctx, argument_types); - - let accumulators = (0..ACCUMULATOR_LEN) - .map(|index| entry.deref(&ctx).get_argument(index)) - .collect::>(); - let descriptors = (0..DESCRIPTOR_COUNT) - .map(|index| entry.deref(&ctx).get_argument(ACCUMULATOR_LEN + index)) - .collect::>(); - - nvvm::WgmmaMmaGroupValuesM64N64K16F32Bf16Op::build(&mut ctx, accumulators, descriptors) - .insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let matching = lowered_kernel_body(&ctx, module_ptr) - .into_iter() - .filter_map(|operation| Operation::get_op::(operation, &ctx)) - .filter(|asm| { - asm.get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .is_some_and(|template| { - template.contains("wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16") - }) - }) - .collect::>(); - assert_eq!(matching.len(), 1); - - let asm = &matching[0]; - let asm_op = asm.get_operation(); - let template = asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .expect("value-form WGMMA template"); - - assert_eq!(template.matches("wgmma.fence.sync.aligned").count(), 1); - assert_eq!(template.matches("wgmma.mma_async").count(), 2); - assert_eq!( - template.matches("wgmma.commit_group.sync.aligned").count(), - 1 - ); - assert_eq!( - template.matches("wgmma.wait_group.sync.aligned 0").count(), - 1 - ); - for forbidden in [".reg .f32", "ld.f32", "st.f32"] { - assert!( - !template.contains(forbidden), - "value-form WGMMA must not materialize accumulator memory via {forbidden:?}: {template}" - ); - } - - let accumulator_operands = (0..ACCUMULATOR_LEN) - .map(|index| format!("${index}")) - .collect::>() - .join(", "); - assert!( - template.contains(&format!("{{{accumulator_operands}}}")), - "WGMMA must consume the 32 output operands as one accumulator tuple: {template}" - ); - assert!(template.contains("$64, $65"), "{template}"); - assert!(template.contains("$66, $67"), "{template}"); - - let mut expected_constraints = vec!["=f".to_owned(); ACCUMULATOR_LEN]; - expected_constraints.extend((0..ACCUMULATOR_LEN).map(|index| index.to_string())); - expected_constraints.extend((0..DESCRIPTOR_COUNT).map(|_| "l".to_owned())); - expected_constraints.push("~{memory}".to_owned()); - let expected_constraints = expected_constraints.join(","); - - assert_eq!( - asm.get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some(expected_constraints.as_str()) - ); - assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); - - let asm_ref = asm_op.deref(&ctx); - assert_eq!( - asm_ref.get_num_operands(), - ACCUMULATOR_LEN + DESCRIPTOR_COUNT - ); - assert_eq!( - asm_ref.get_num_results(), - 1, - "LLVM inline asm must return the 32 WGMMA accumulator values as one struct" - ); - - let aggregate = asm_ref.get_result(0); - let aggregate_ty = aggregate.get_type(&ctx); - let aggregate_ty = aggregate_ty.deref(&ctx); - let struct_ty = aggregate_ty - .downcast_ref::() - .expect("value-form WGMMA inline asm must return an LLVM struct"); - assert_eq!(struct_ty.num_fields(), ACCUMULATOR_LEN); - for index in 0..ACCUMULATOR_LEN { - assert!( - struct_ty - .field_type(index) - .deref(&ctx) - .downcast_ref::() - .is_some(), - "WGMMA result field {index} must remain f32" - ); - } - - let extract_indices = lowered_kernel_body(&ctx, module_ptr) - .into_iter() - .filter_map(|operation| Operation::get_op::(operation, &ctx)) - .filter_map(|extract| { - let extract_op = extract.get_operation().deref(&ctx); - (extract_op.get_operand(0) == aggregate).then(|| extract.indices(&ctx)) - }) - .collect::>(); - - let expected_indices = (0..ACCUMULATOR_LEN) - .map(|index| vec![index as u32]) - .collect::>(); - assert_eq!( - extract_indices, expected_indices, - "all 32 value-form WGMMA results must be extracted in constraint order" - ); - - Ok(()) -} - -#[test] -fn test_value_form_m64n128_bf16_wgmma_group_lowers_to_sixty_four_tied_registers() --> Result<(), anyhow::Error> { - use llvm_export::types as llvm_types; - use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; - use pliron::r#type::Typed; - - const ACCUMULATOR_LEN: usize = 64; - const DESCRIPTOR_COUNT: usize = 4; - - let mut ctx = make_test_ctx(); - let f32_ty = FP32Type::get(&ctx); - let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let argument_types = (0..ACCUMULATOR_LEN) - .map(|_| f32_ty.into()) - .chain((0..DESCRIPTOR_COUNT).map(|_| u64_ty.into())) - .collect::>(); - let (module_ptr, entry) = build_test_kernel(&mut ctx, argument_types); - - let accumulators = (0..ACCUMULATOR_LEN) - .map(|index| entry.deref(&ctx).get_argument(index)) - .collect::>(); - let descriptors = (0..DESCRIPTOR_COUNT) - .map(|index| entry.deref(&ctx).get_argument(ACCUMULATOR_LEN + index)) - .collect::>(); - - nvvm::WgmmaMmaGroupValuesM64N128K16F32Bf16Op::build(&mut ctx, accumulators, descriptors) - .insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let matching = lowered_kernel_body(&ctx, module_ptr) - .into_iter() - .filter_map(|operation| Operation::get_op::(operation, &ctx)) - .filter(|asm| { - asm.get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .is_some_and(|template| template.contains("m64n128k16.f32.bf16.bf16")) - }) - .collect::>(); - assert_eq!(matching.len(), 1); - - let asm = &matching[0]; - let template = asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .expect("m64n128 value-form WGMMA template"); - assert_eq!(template.matches("wgmma.mma_async").count(), 2); - assert!(template.contains("m64n128k16.f32.bf16.bf16")); - assert!(template.contains("$128, $129")); - assert!(template.contains("$130, $131")); - assert!(!template.contains("ld.f32")); - assert!(!template.contains("st.f32")); - - let constraints = asm - .get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .expect("m64n128 value-form WGMMA constraints"); - assert_eq!( - constraints - .split(',') - .filter(|value| *value == "=f") - .count(), - 64 - ); - assert_eq!( - constraints.split(',').filter(|value| *value == "l").count(), - 4 - ); - assert!(constraints.ends_with("~{memory}")); - - let aggregate = asm.get_operation().deref(&ctx).get_result(0); - let aggregate_ty = aggregate.get_type(&ctx); - let aggregate_ty = aggregate_ty.deref(&ctx); - let struct_ty = aggregate_ty - .downcast_ref::() - .expect("m64n128 inline asm must return an LLVM struct"); - assert_eq!(struct_ty.num_fields(), ACCUMULATOR_LEN); - assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); - - Ok(()) -} - -#[test] -fn test_value_form_f16_wgmma_group_lowers_to_tied_register_inline_ptx() -> Result<(), anyhow::Error> -{ - use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; - - const ACCUMULATOR_LEN: usize = 32; - const DESCRIPTOR_COUNT: usize = 2; - - let mut ctx = make_test_ctx(); - let f32_ty = FP32Type::get(&ctx); - let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let argument_types = (0..ACCUMULATOR_LEN) - .map(|_| f32_ty.into()) - .chain((0..DESCRIPTOR_COUNT).map(|_| u64_ty.into())) - .collect::>(); - let (module_ptr, entry) = build_test_kernel(&mut ctx, argument_types); - - let accumulators = (0..ACCUMULATOR_LEN) - .map(|index| entry.deref(&ctx).get_argument(index)) - .collect::>(); - let descriptors = (0..DESCRIPTOR_COUNT) - .map(|index| entry.deref(&ctx).get_argument(ACCUMULATOR_LEN + index)) - .collect::>(); - - nvvm::WgmmaMmaGroupValuesM64N64K16F32F16Op::build(&mut ctx, accumulators, descriptors) - .insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let matching = lowered_kernel_body(&ctx, module_ptr) - .into_iter() - .filter_map(|operation| Operation::get_op::(operation, &ctx)) - .filter(|asm| { - asm.get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .is_some_and(|template| template.contains("m64n64k16.f32.f16.f16")) - }) - .collect::>(); - assert_eq!(matching.len(), 1); - - let asm = &matching[0]; - let template = asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .expect("F16 value-form WGMMA template"); - assert_eq!(template.matches("wgmma.mma_async").count(), 1); - assert!(template.contains("m64n64k16.f32.f16.f16")); - assert!(!template.contains(".bf16.bf16")); - assert!(template.contains("wgmma.fence.sync.aligned")); - assert!(template.contains("wgmma.commit_group.sync.aligned")); - assert!(template.contains("wgmma.wait_group.sync.aligned 0")); - assert!(!template.contains("ld.f32")); - assert!(!template.contains("st.f32")); - - let constraints = asm - .get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .expect("F16 value-form WGMMA constraints"); - assert_eq!( - constraints - .split(',') - .filter(|value| *value == "=f") - .count(), - 32 - ); - assert_eq!( - constraints.split(',').filter(|value| *value == "l").count(), - 2 - ); - assert!(constraints.ends_with("~{memory}")); - assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); - - Ok(()) -} - -#[test] -fn test_value_form_tf32_wgmma_group_lowers_to_tied_register_inline_ptx() -> Result<(), anyhow::Error> -{ - use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; - - const ACCUMULATOR_LEN: usize = 32; - const DESCRIPTOR_COUNT: usize = 2; - - let mut ctx = make_test_ctx(); - let f32_ty = FP32Type::get(&ctx); - let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let argument_types = (0..ACCUMULATOR_LEN) - .map(|_| f32_ty.into()) - .chain((0..DESCRIPTOR_COUNT).map(|_| u64_ty.into())) - .collect::>(); - let (module_ptr, entry) = build_test_kernel(&mut ctx, argument_types); - - let accumulators = (0..ACCUMULATOR_LEN) - .map(|index| entry.deref(&ctx).get_argument(index)) - .collect::>(); - let descriptors = (0..DESCRIPTOR_COUNT) - .map(|index| entry.deref(&ctx).get_argument(ACCUMULATOR_LEN + index)) - .collect::>(); - - nvvm::WgmmaMmaGroupValuesM64N64K8F32Tf32Op::build(&mut ctx, accumulators, descriptors) - .insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let matching = lowered_kernel_body(&ctx, module_ptr) - .into_iter() - .filter_map(|operation| Operation::get_op::(operation, &ctx)) - .filter(|asm| { - asm.get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .is_some_and(|template| template.contains("m64n64k8.f32.tf32.tf32")) - }) - .collect::>(); - assert_eq!(matching.len(), 1); - - let asm = &matching[0]; - let template = asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .expect("TF32 value-form WGMMA template"); - assert_eq!(template.matches("wgmma.mma_async").count(), 1); - assert!(template.contains("m64n64k8.f32.tf32.tf32")); - assert!(!template.contains(".bf16.bf16")); - assert!(!template.contains(".f16.f16")); - assert!(!template.contains("m64n64k16")); - assert!(template.contains("$64, $65, 1, 1, 1;")); - assert!(!template.contains("$64, $65, 1, 1, 1, 0, 0;")); - assert!(template.contains("wgmma.fence.sync.aligned")); - assert!(template.contains("wgmma.commit_group.sync.aligned")); - assert!(template.contains("wgmma.wait_group.sync.aligned 0")); - assert!(!template.contains("ld.f32")); - assert!(!template.contains("st.f32")); - - let constraints = asm - .get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .expect("TF32 value-form WGMMA constraints"); - assert_eq!( - constraints - .split(',') - .filter(|value| *value == "=f") - .count(), - 32 - ); - assert_eq!( - constraints.split(',').filter(|value| *value == "l").count(), - 2 - ); - assert!(constraints.ends_with("~{memory}")); - assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); - - Ok(()) -} - -#[test] -fn test_pointer_form_wgmma_sequence_preserves_deferred_fallback() -> Result<(), anyhow::Error> { - use dialect_mir::types::MirPtrType; - use pliron::builtin::attributes::IntegerAttr; - use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; - use pliron::utils::apint::APInt; - use std::num::NonZeroUsize; - - let mut ctx = make_test_ctx(); - let f32_ty = FP32Type::get(&ctx); - // A direct *mut f32 is intentionally not the public [[f32; 8]; 4] - // accumulator shape. It must therefore retain the deferred pointer path. - let accumulator_ptr_ty = MirPtrType::get_generic(&mut ctx, f32_ty.into(), true); - let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let (module_ptr, entry) = build_test_kernel( - &mut ctx, - vec![accumulator_ptr_ty.into(), u64_ty.into(), u64_ty.into()], - ); - let accumulator = entry.deref(&ctx).get_argument(0); - let desc_a = entry.deref(&ctx).get_argument(1); - let desc_b = entry.deref(&ctx).get_argument(2); - - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - Operation::new( - &mut ctx, - nvvm::WgmmaMmaM64N64K16F32Bf16Op::get_concrete_op_info(), - vec![], - vec![accumulator, desc_a, desc_b], - vec![], - 0, - ) - .insert_at_back(entry, &ctx); - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - - let zero_attr = IntegerAttr::new(u64_ty, APInt::from_i64(0, NonZeroUsize::new(64).unwrap())); - let zero = Operation::new( - &mut ctx, - mir::MirConstantOp::get_concrete_op_info(), - vec![u64_ty.into()], - vec![], - vec![], - 0, - ); - mir::MirConstantOp::new(zero).set_attr_value(&ctx, zero_attr); - zero.insert_at_back(entry, &ctx); - let zero_value = zero.deref(&ctx).get_result(0); - nvvm::WgmmaWaitGroupSyncAlignedOp::build(&mut ctx, zero_value).insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let matching = lowered_kernel_body(&ctx, module_ptr) - .into_iter() - .filter_map(|operation| Operation::get_op::(operation, &ctx)) - .filter(|asm| { - asm.get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .is_some_and(|template| template.contains("wgmma.mma_async")) - }) - .collect::>(); - assert_eq!(matching.len(), 1); - - let asm = &matching[0]; - let template = asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .expect("WGMMA template"); - assert_eq!(template.matches("ld.f32 %acc").count(), 32); - assert_eq!(template.matches("st.f32 [$0").count(), 32); - assert_eq!(template.matches("wgmma.mma_async").count(), 1); - assert!(template.contains("wgmma.fence.sync.aligned")); - assert!(template.contains("wgmma.commit_group.sync.aligned")); - assert!(template.contains("wgmma.wait_group.sync.aligned 0")); - assert_eq!( - asm.get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some("l,l,l,~{memory}") - ); - assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); - - Ok(()) -} - -#[test] -fn test_pointer_form_wgmma_region_canonicalizes_reborrow_identity() -> Result<(), anyhow::Error> { - use dialect_mir::attributes::{MirCastKindAttr, MirPointerKindAuthorityAttr}; - use dialect_mir::types::{MirArrayType, MirPointerKind, MirPtrType}; - use pliron::builtin::attributes::IntegerAttr; - use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; - use pliron::utils::apint::APInt; - use std::num::NonZeroUsize; - - let mut ctx = make_test_ctx(); - let f32_ty = FP32Type::get(&ctx); - let row_ty = MirArrayType::get(&mut ctx, f32_ty.into(), 8); - let accumulator_ty = MirArrayType::get(&mut ctx, row_ty.into(), 4); - let erased_ptr_ty = MirPtrType::get_generic(&mut ctx, accumulator_ty.into(), true); - let unique_ptr_ty: pliron::r#type::TypeHandle = MirPtrType::get_generic_with_kind( - &mut ctx, - accumulator_ty.into(), - true, - MirPointerKind::UniqueRef, - ) - .into(); - let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - - let (module_ptr, entry) = build_test_kernel( - &mut ctx, - vec![erased_ptr_ty.into(), u64_ty.into(), u64_ty.into()], - ); - let accumulator = entry.deref(&ctx).get_argument(0); - let desc_a = entry.deref(&ctx).get_argument(1); - let desc_b = entry.deref(&ctx).get_argument(2); - - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - - // The importer emits a fresh kind-only retype for each `Rvalue::Ref`. - // Both typed operands must canonicalize to the same storage identity, while - // the first typed operand remains available to the linear lowering plan. - let first_retype = Operation::new( - &mut ctx, - mir::MirCastOp::get_concrete_op_info(), - vec![unique_ptr_ty], - vec![accumulator], - vec![], - 0, - ); - mir::MirCastOp::new(first_retype).set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - mir::MirCastOp::new(first_retype) - .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); - first_retype.insert_at_back(entry, &ctx); - let first_reborrow = first_retype.deref(&ctx).get_result(0); - - Operation::new( - &mut ctx, - nvvm::WgmmaMmaM64N64K16F32Bf16Op::get_concrete_op_info(), - vec![], - vec![first_reborrow, desc_a, desc_b], - vec![], - 0, - ) - .insert_at_back(entry, &ctx); - - let second_retype = Operation::new( - &mut ctx, - mir::MirCastOp::get_concrete_op_info(), - vec![unique_ptr_ty], - vec![accumulator], - vec![], - 0, - ); - mir::MirCastOp::new(second_retype).set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - mir::MirCastOp::new(second_retype) - .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); - second_retype.insert_at_back(entry, &ctx); - let second_reborrow = second_retype.deref(&ctx).get_result(0); - assert_ne!( - first_reborrow, second_reborrow, - "the regression requires two distinct reborrow SSA values" - ); - let first_reborrow_ty = first_reborrow.get_type(&ctx); - assert_eq!( - first_reborrow_ty - .deref(&ctx) - .downcast_ref::() - .expect("first reborrow must remain a typed MIR pointer") - .pointer_kind(), - MirPointerKind::UniqueRef, - "the regression must provide a typed UniqueRef for the linear plan to retain" - ); - - Operation::new( - &mut ctx, - nvvm::WgmmaMmaM64N64K16F32Bf16Op::get_concrete_op_info(), - vec![], - vec![second_reborrow, desc_a, desc_b], - vec![], - 0, - ) - .insert_at_back(entry, &ctx); - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - - let zero_attr = IntegerAttr::new(u64_ty, APInt::from_i64(0, NonZeroUsize::new(64).unwrap())); - let zero = Operation::new( - &mut ctx, - mir::MirConstantOp::get_concrete_op_info(), - vec![u64_ty.into()], - vec![], - vec![], - 0, - ); - mir::MirConstantOp::new(zero).set_attr_value(&ctx, zero_attr); - zero.insert_at_back(entry, &ctx); - let zero_value = zero.deref(&ctx).get_result(0); - nvvm::WgmmaWaitGroupSyncAlignedOp::build(&mut ctx, zero_value).insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let matching = lowered_kernel_body(&ctx, module_ptr) - .into_iter() - .filter_map(|operation| Operation::get_op::(operation, &ctx)) - .filter(|asm| { - asm.get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .is_some_and(|template| { - template.contains("wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16") - }) - }) - .collect::>(); - assert_eq!( - matching.len(), - 1, - "distinct reborrow SSA values for one accumulator must not break deferred fusion" - ); - let template = matching[0] - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .expect("WGMMA template"); - assert_eq!( - template.matches("wgmma.mma_async").count(), - 2, - "both reborrows must remain in one fused accumulator region" - ); - Ok(()) -} - -#[test] -fn test_pointer_form_wgmma_sequence_uses_value_adapter_before_lowering() -> Result<(), anyhow::Error> -{ - use dialect_mir::types::{MirArrayType, MirPtrType}; - use pliron::builtin::attributes::IntegerAttr; - use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; - use pliron::utils::apint::APInt; - use std::num::NonZeroUsize; - - const ACCUMULATOR_LEN: usize = 32; - const DESCRIPTOR_COUNT: usize = 4; - - let mut ctx = make_test_ctx(); - let f32_ty = FP32Type::get(&ctx); - let row_ty = MirArrayType::get(&mut ctx, f32_ty.into(), 8); - let accumulator_ty = MirArrayType::get(&mut ctx, row_ty.into(), 4); - let accumulator_ptr_ty = MirPtrType::get_generic(&mut ctx, accumulator_ty.into(), true); - let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - - let (module_ptr, entry) = build_test_kernel( - &mut ctx, - vec![ - accumulator_ptr_ty.into(), - u64_ty.into(), - u64_ty.into(), - u64_ty.into(), - u64_ty.into(), - ], - ); - let accumulator = entry.deref(&ctx).get_argument(0); - let descriptors = (1..=DESCRIPTOR_COUNT) - .map(|index| entry.deref(&ctx).get_argument(index)) - .collect::>(); - - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - for pair in descriptors.chunks_exact(2) { - Operation::new( - &mut ctx, - nvvm::WgmmaMmaM64N64K16F32Bf16Op::get_concrete_op_info(), - vec![], - vec![accumulator, pair[0], pair[1]], - vec![], - 0, - ) - .insert_at_back(entry, &ctx); - } - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - - let zero_attr = IntegerAttr::new(u64_ty, APInt::from_i64(0, NonZeroUsize::new(64).unwrap())); - let zero = Operation::new( - &mut ctx, - mir::MirConstantOp::get_concrete_op_info(), - vec![u64_ty.into()], - vec![], - vec![], - 0, - ); - mir::MirConstantOp::new(zero).set_attr_value(&ctx, zero_attr); - zero.insert_at_back(entry, &ctx); - let zero_value = zero.deref(&ctx).get_result(0); - nvvm::WgmmaWaitGroupSyncAlignedOp::build(&mut ctx, zero_value).insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let body = lowered_kernel_body(&ctx, module_ptr); - let matching = body - .iter() - .copied() - .filter_map(|operation| { - Operation::get_op::(operation, &ctx) - .map(|inline_asm| (operation, inline_asm)) - }) - .filter(|(_, asm)| { - asm.get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .is_some_and(|template| template.contains("wgmma.mma_async")) - }) - .collect::>(); - assert_eq!(matching.len(), 1); - - let (asm_operation, asm) = &matching[0]; - let template = asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .expect("WGMMA template"); - - assert_eq!(template.matches("wgmma.fence.sync.aligned").count(), 1); - assert_eq!(template.matches("wgmma.mma_async").count(), 2); - assert_eq!( - template.matches("wgmma.commit_group.sync.aligned").count(), - 1 - ); - assert_eq!( - template.matches("wgmma.wait_group.sync.aligned 0").count(), - 1 - ); - assert!( - !template.contains(".reg .f32") - && !template.contains("ld.f32") - && !template.contains("st.f32"), - "value-form WGMMA must not materialize accumulator memory inside asm: {template}" - ); - - let mut expected_constraints = vec!["=f".to_owned(); ACCUMULATOR_LEN]; - expected_constraints.extend((0..ACCUMULATOR_LEN).map(|index| index.to_string())); - expected_constraints.extend((0..DESCRIPTOR_COUNT).map(|_| "l".to_owned())); - expected_constraints.push("~{memory}".to_owned()); - let expected_constraints = expected_constraints.join(","); - - assert_eq!( - asm.get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some(expected_constraints.as_str()) - ); - assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); - - let asm_ref = asm_operation.deref(&ctx); - assert_eq!( - asm_ref.get_num_operands(), - ACCUMULATOR_LEN + DESCRIPTOR_COUNT - ); - assert_eq!(asm_ref.get_num_results(), 1); - - let asm_position = body - .iter() - .position(|operation| operation == asm_operation) - .expect("WGMMA asm must be in the lowered kernel body"); - - let load_positions = body - .iter() - .enumerate() - .filter_map(|(index, operation)| { - Operation::get_op::(*operation, &ctx) - .is_some() - .then_some(index) - }) - .collect::>(); - assert_eq!( - load_positions.len(), - ACCUMULATOR_LEN, - "the linear adapter must load each accumulator value exactly once" - ); - assert!( - load_positions.iter().all(|index| *index < asm_position), - "all accumulator loads must happen before the WGMMA region" - ); - - let store_positions = body - .iter() - .enumerate() - .filter_map(|(index, operation)| { - Operation::get_op::(*operation, &ctx) - .is_some() - .then_some(index) - }) - .collect::>(); - assert_eq!( - store_positions.len(), - ACCUMULATOR_LEN, - "the linear adapter must store each accumulator value exactly once" - ); - assert!( - store_positions.iter().all(|index| *index > asm_position), - "all accumulator stores must happen after the final wait" - ); - - assert_eq!( - body.iter() - .filter( - |operation| Operation::get_op::(**operation, &ctx).is_some() - ) - .count(), - ACCUMULATOR_LEN, - "all WGMMA accumulator results must be recovered as scalar SSA values" - ); - - Ok(()) -} - -#[test] -fn test_pointer_form_m64n128_bf16_linear_full_drain_uses_sixty_four_value_adapter() --> Result<(), anyhow::Error> { - let mut ctx = make_test_ctx(); - let (module_ptr, entry, accumulators, descriptors) = - build_wgmma_m64n128_canonical_pointer_test_kernel(&mut ctx, 1, 4); - let accumulator = accumulators[0]; - - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_pointer_wgmma_mma_m64n128(&mut ctx, entry, accumulator, descriptors[0], descriptors[1]); - append_pointer_wgmma_mma_m64n128(&mut ctx, entry, accumulator, descriptors[2], descriptors[3]); - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_wgmma_wait_group_constant(&mut ctx, entry, 0); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let matching = lowered_kernel_body(&ctx, module_ptr) - .into_iter() - .filter_map(|operation| Operation::get_op::(operation, &ctx)) - .filter(|asm| { - asm.get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .is_some_and(|template| template.contains("m64n128k16.f32.bf16.bf16")) - }) - .collect::>(); - assert_eq!(matching.len(), 1); - - let asm = &matching[0]; - let template = asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .expect("m64n128 pointer-form WGMMA template"); - assert_eq!(template.matches("wgmma.mma_async").count(), 2); - assert!(template.contains("m64n128k16.f32.bf16.bf16")); - assert!(!template.contains("ld.f32")); - assert!(!template.contains("st.f32")); - - let constraints = asm - .get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .expect("m64n128 pointer-form WGMMA constraints"); - assert_eq!( - constraints - .split(',') - .filter(|value| *value == "=f") - .count(), - 64 - ); - assert_eq!( - constraints.split(',').filter(|value| *value == "l").count(), - 4 - ); - assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); - - Ok(()) -} - -#[test] -fn test_pointer_form_f16_wgmma_linear_full_drain_uses_value_adapter() -> Result<(), anyhow::Error> { - let mut ctx = make_test_ctx(); - let (module_ptr, entry, accumulators, descriptors) = - build_wgmma_canonical_pointer_test_kernel(&mut ctx, 1, 2); - let accumulator = accumulators[0]; - - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_pointer_wgmma_mma_f16(&mut ctx, entry, accumulator, descriptors[0], descriptors[1]); - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_wgmma_wait_group_constant(&mut ctx, entry, 0); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let matching = lowered_kernel_body(&ctx, module_ptr) - .into_iter() - .filter_map(|operation| Operation::get_op::(operation, &ctx)) - .filter(|asm| { - asm.get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .is_some_and(|template| template.contains("m64n64k16.f32.f16.f16")) - }) - .collect::>(); - assert_eq!(matching.len(), 1); - - let asm = &matching[0]; - let template = asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .expect("F16 pointer-form WGMMA template"); - assert_eq!(template.matches("wgmma.fence.sync.aligned").count(), 1); - assert_eq!(template.matches("wgmma.mma_async").count(), 1); - assert_eq!( - template.matches("wgmma.commit_group.sync.aligned").count(), - 1 - ); - assert_eq!( - template.matches("wgmma.wait_group.sync.aligned 0").count(), - 1 - ); - assert!(template.contains("m64n64k16.f32.f16.f16")); - assert!(!template.contains(".bf16.bf16")); - assert!(!template.contains("ld.f32")); - assert!(!template.contains("st.f32")); - assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); - - Ok(()) -} - -#[test] -fn test_pointer_form_tf32_wgmma_linear_full_drain_uses_value_adapter() -> Result<(), anyhow::Error> -{ - let mut ctx = make_test_ctx(); - let (module_ptr, entry, accumulators, descriptors) = - build_wgmma_canonical_pointer_test_kernel(&mut ctx, 1, 2); - let accumulator = accumulators[0]; - - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_pointer_wgmma_mma_tf32(&mut ctx, entry, accumulator, descriptors[0], descriptors[1]); - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_wgmma_wait_group_constant(&mut ctx, entry, 0); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let matching = lowered_kernel_body(&ctx, module_ptr) - .into_iter() - .filter_map(|operation| Operation::get_op::(operation, &ctx)) - .filter(|asm| { - asm.get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .is_some_and(|template| template.contains("m64n64k8.f32.tf32.tf32")) - }) - .collect::>(); - assert_eq!(matching.len(), 1); - - let asm = &matching[0]; - let template = asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .expect("TF32 pointer-form WGMMA template"); - assert_eq!(template.matches("wgmma.fence.sync.aligned").count(), 1); - assert_eq!(template.matches("wgmma.mma_async").count(), 1); - assert_eq!( - template.matches("wgmma.commit_group.sync.aligned").count(), - 1 - ); - assert_eq!( - template.matches("wgmma.wait_group.sync.aligned 0").count(), - 1 - ); - assert!(template.contains("m64n64k8.f32.tf32.tf32")); - assert!(!template.contains(".bf16.bf16")); - assert!(!template.contains(".f16.f16")); - assert!(!template.contains("m64n64k16")); - assert!(template.contains("$64, $65, 1, 1, 1;")); - assert!(!template.contains("$64, $65, 1, 1, 1, 0, 0;")); - assert!(!template.contains("ld.f32")); - assert!(!template.contains("st.f32")); - assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); - - Ok(()) -} - -#[test] -fn test_pointer_form_wgmma_partial_wait_pipeline_keeps_multiple_groups_in_flight() --> Result<(), anyhow::Error> { - const SLOT_COUNT: usize = 2; - const ACCUMULATOR_LEN: usize = 32; - const GROUP_COUNT: usize = 4; - const DESCRIPTOR_COUNT: usize = GROUP_COUNT * 2; - const MAX_PENDING_GROUPS: i64 = 1; - const RESULT_COUNT: usize = SLOT_COUNT * ACCUMULATOR_LEN; - - let mut ctx = make_test_ctx(); - let (module_ptr, entry, accumulators, descriptors) = - build_wgmma_canonical_pointer_test_kernel(&mut ctx, SLOT_COUNT, DESCRIPTOR_COUNT); - - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - for group in 0..GROUP_COUNT { - append_pointer_wgmma_mma( - &mut ctx, - entry, - accumulators[group % SLOT_COUNT], - descriptors[group * 2], - descriptors[group * 2 + 1], - ); - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - if group + 1 >= SLOT_COUNT { - append_wgmma_wait_group_constant(&mut ctx, entry, MAX_PENDING_GROUPS); - } - } - append_wgmma_wait_group_constant(&mut ctx, entry, 0); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let body = lowered_kernel_body(&ctx, module_ptr); - let matching = body - .iter() - .copied() - .filter_map(|operation| { - Operation::get_op::(operation, &ctx) - .map(|inline_asm| (operation, inline_asm)) - }) - .filter(|(_, asm)| { - asm.get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .is_some_and(|template| { - template.contains("wgmma.mma_async") - && template.contains("wgmma.wait_group.sync.aligned 1") - }) - }) - .collect::>(); - assert_eq!( - matching.len(), - 1, - "expected one fused partial-wait WGMMA pipeline" - ); - - let (asm_operation, asm) = &matching[0]; - let template = asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .expect("pipeline WGMMA template"); - - assert_eq!(template.matches("wgmma.fence.sync.aligned").count(), 1); - assert_eq!(template.matches("wgmma.mma_async").count(), GROUP_COUNT); - assert_eq!( - template.matches("wgmma.commit_group.sync.aligned").count(), - GROUP_COUNT - ); - assert_eq!( - template.matches("wgmma.wait_group.sync.aligned 1").count(), - GROUP_COUNT - SLOT_COUNT + 1 - ); - assert_eq!( - template.matches("wgmma.wait_group.sync.aligned 0").count(), - 1 - ); - assert!( - template.contains("{$0, $1"), - "slot 0 must use the first accumulator tuple" - ); - assert!( - template.contains("{$32, $33"), - "slot 1 must use the second accumulator tuple" - ); - assert!(template.contains("$128, $129")); - assert!(template.contains("$134, $135")); - assert!( - !template.contains(".reg .f32") - && !template.contains("ld.f32") - && !template.contains("st.f32"), - "pipeline WGMMA must not materialize accumulator memory inside asm: {template}" - ); - - let mut expected_constraints = vec!["=f".to_owned(); RESULT_COUNT]; - expected_constraints.extend((0..RESULT_COUNT).map(|index| index.to_string())); - expected_constraints.extend((0..DESCRIPTOR_COUNT).map(|_| "l".to_owned())); - expected_constraints.push("~{memory}".to_owned()); - let expected_constraints = expected_constraints.join(","); - assert_eq!( - asm.get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some(expected_constraints.as_str()) - ); - assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); - assert_eq!( - asm_operation.deref(&ctx).get_num_operands(), - RESULT_COUNT + DESCRIPTOR_COUNT - ); - assert_eq!(asm_operation.deref(&ctx).get_num_results(), 1); - - let asm_position = body - .iter() - .position(|operation| operation == asm_operation) - .expect("pipeline asm must be in the lowered kernel body"); - let load_positions = body - .iter() - .enumerate() - .filter_map(|(index, operation)| { - Operation::get_op::(*operation, &ctx) - .is_some() - .then_some(index) - }) - .collect::>(); - assert_eq!(load_positions.len(), RESULT_COUNT); - assert!(load_positions.iter().all(|index| *index < asm_position)); - - let store_positions = body - .iter() - .enumerate() - .filter_map(|(index, operation)| { - Operation::get_op::(*operation, &ctx) - .is_some() - .then_some(index) - }) - .collect::>(); - assert_eq!(store_positions.len(), RESULT_COUNT); - assert!(store_positions.iter().all(|index| *index > asm_position)); - assert_eq!( - body.iter() - .filter(|operation| { - Operation::get_op::(**operation, &ctx).is_some() - }) - .count(), - RESULT_COUNT - ); - - Ok(()) -} - -fn build_pointer_form_wgmma_counted_pipeline_case( - ctx: &mut Context, - slot_count: usize, - wait_depths: &[i64], - repeat_last_accumulator: bool, -) -> pliron::context::Ptr { - use dialect_mir::types::{MirArrayType, MirPtrType}; - use pliron::basic_block::BasicBlock; - use pliron::builtin::op_interfaces::OperandSegmentInterface; - use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; - - assert_eq!(wait_depths.len(), slot_count); - - let f32_ty = FP32Type::get(ctx); - let row_ty = MirArrayType::get(ctx, f32_ty.into(), 8); - let accumulator_ty = MirArrayType::get(ctx, row_ty.into(), 4); - let accumulator_ptr_ty = MirPtrType::get_generic(ctx, accumulator_ty.into(), true); - let u32_ty = IntegerType::get(ctx, 32, Signedness::Unsigned); - let u64_ty = IntegerType::get(ctx, 64, Signedness::Unsigned); - let i1_ty = IntegerType::get(ctx, 1, Signedness::Signless); - let u64_type: pliron::r#type::TypeHandle = u64_ty.into(); - - let mut argument_types: Vec = - vec![accumulator_ptr_ty.into(); slot_count]; - argument_types.extend(vec![u64_type; slot_count * 2]); - let (module_ptr, preheader) = build_test_kernel(ctx, argument_types); - - let accumulators = (0..slot_count) - .map(|slot| preheader.deref(ctx).get_argument(slot)) - .collect::>(); - let desc_bases = (0..slot_count * 2) - .map(|index| preheader.deref(ctx).get_argument(slot_count + index)) - .collect::>(); - - let module_region = module_ptr.deref(ctx).get_region(0); - let module_block = module_region.deref(ctx).iter(ctx).next().unwrap(); - let function = module_block.deref(ctx).iter(ctx).next().unwrap(); - let function_region = function.deref(ctx).get_region(0); - - let mut header_types: Vec = vec![u32_ty.into()]; - header_types.extend(vec![u64_type; slot_count * 2]); - let header = BasicBlock::new(ctx, None, header_types); - header.insert_at_back(function_region, ctx); - let latch = BasicBlock::new(ctx, None, vec![]); - latch.insert_at_back(function_region, ctx); - let exit = BasicBlock::new(ctx, None, vec![]); - exit.insert_at_back(function_region, ctx); - - nvvm::WgmmaFenceSyncAlignedOp::build(ctx).insert_at_back(preheader, ctx); - let i0 = append_mir_unsigned_constant(ctx, preheader, u32_ty, 0); - let mut initial_values = vec![i0]; - initial_values.extend(desc_bases.iter().copied()); - Operation::new( - ctx, - mir::MirGotoOp::get_concrete_op_info(), - vec![], - initial_values, - vec![header], - 0, - ) - .insert_at_back(preheader, ctx); - - let i = header.deref(ctx).get_argument(0); - let descriptors = (0..slot_count * 2) - .map(|index| header.deref(ctx).get_argument(1 + index)) - .collect::>(); - let bound = append_mir_unsigned_constant(ctx, header, u32_ty, 4); - let lt = Operation::new( - ctx, - mir::MirLtOp::get_concrete_op_info(), - vec![i1_ty.into()], - vec![i, bound], - vec![], - 0, - ); - lt.insert_at_back(header, ctx); - let lt_value = lt.deref(ctx).get_result(0); - let not_lt = Operation::new( - ctx, - mir::MirNotOp::get_concrete_op_info(), - vec![i1_ty.into()], - vec![lt_value], - vec![], - 0, - ); - not_lt.insert_at_back(header, ctx); - let not_lt_value = not_lt.deref(ctx).get_result(0); - let (branch_operands, segment_sizes) = - mir::MirCondBranchOp::compute_segment_sizes(vec![vec![not_lt_value], vec![], vec![]]); - let branch = Operation::new( - ctx, - mir::MirCondBranchOp::get_concrete_op_info(), - vec![], - branch_operands, - vec![exit, latch], - 0, - ); - Operation::get_op::(branch, ctx) - .expect("MirCondBranchOp") - .set_operand_segment_sizes(ctx, segment_sizes); - branch.insert_at_back(header, ctx); - - for slot in 0..slot_count { - let accumulator = if repeat_last_accumulator && slot + 1 == slot_count { - accumulators[0] - } else { - accumulators[slot] - }; - append_pointer_wgmma_mma( - ctx, - latch, - accumulator, - descriptors[slot * 2], - descriptors[slot * 2 + 1], - ); - nvvm::WgmmaCommitGroupSyncAlignedOp::build(ctx).insert_at_back(latch, ctx); - append_wgmma_wait_group_constant(ctx, latch, wait_depths[slot]); - } - - let one = append_mir_unsigned_constant(ctx, latch, u32_ty, 1); - let i_next = Operation::new( - ctx, - mir::MirAddOp::get_concrete_op_info(), - vec![u32_ty.into()], - vec![i, one], - vec![], - 0, - ); - i_next.insert_at_back(latch, ctx); - let i_next = i_next.deref(ctx).get_result(0); - - let mut next_values = vec![i_next]; - for (index, descriptor) in descriptors.iter().copied().enumerate() { - let step = append_mir_unsigned_constant(ctx, latch, u64_ty, 16 * (index as u64 + 1)); - let next = Operation::new( - ctx, - mir::MirAddOp::get_concrete_op_info(), - vec![u64_ty.into()], - vec![descriptor, step], - vec![], - 0, - ); - next.insert_at_back(latch, ctx); - next_values.push(next.deref(ctx).get_result(0)); - } - - Operation::new( - ctx, - mir::MirGotoOp::get_concrete_op_info(), - vec![], - next_values, - vec![header], - 0, - ) - .insert_at_back(latch, ctx); - - append_wgmma_wait_group_constant(ctx, exit, 0); - append_return(ctx, exit); - - module_ptr -} - -#[test] -fn test_pointer_form_wgmma_two_slot_counted_pipeline_stays_register_resident() --> Result<(), anyhow::Error> { - use dialect_mir::types::{MirArrayType, MirPtrType}; - use pliron::basic_block::BasicBlock; - use pliron::builtin::op_interfaces::OperandSegmentInterface; - use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; - - const SLOT_COUNT: usize = 2; - const ACCUMULATOR_LEN: usize = 32; - const RESULT_COUNT: usize = SLOT_COUNT * ACCUMULATOR_LEN; - const LOOP_CONTROL_COUNT: usize = 9; - const TRIP_COUNT: u64 = 4; - const DESC_A0_STEP: u64 = 16; - const DESC_B0_STEP: u64 = 32; - const DESC_A1_STEP: u64 = 48; - const DESC_B1_STEP: u64 = 64; - - let mut ctx = make_test_ctx(); - let f32_ty = FP32Type::get(&ctx); - let row_ty = MirArrayType::get(&mut ctx, f32_ty.into(), 8); - let accumulator_ty = MirArrayType::get(&mut ctx, row_ty.into(), 4); - let accumulator_ptr_ty = MirPtrType::get_generic(&mut ctx, accumulator_ty.into(), true); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); - - let (module_ptr, preheader) = build_test_kernel( - &mut ctx, - vec![ - accumulator_ptr_ty.into(), - accumulator_ptr_ty.into(), - u64_ty.into(), - u64_ty.into(), - u64_ty.into(), - u64_ty.into(), - ], - ); - let accumulator0 = preheader.deref(&ctx).get_argument(0); - let accumulator1 = preheader.deref(&ctx).get_argument(1); - let desc_a0_base = preheader.deref(&ctx).get_argument(2); - let desc_b0_base = preheader.deref(&ctx).get_argument(3); - let desc_a1_base = preheader.deref(&ctx).get_argument(4); - let desc_b1_base = preheader.deref(&ctx).get_argument(5); - - let module_region = module_ptr.deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - let function = module_block.deref(&ctx).iter(&ctx).next().unwrap(); - let function_region = function.deref(&ctx).get_region(0); - - let header = BasicBlock::new( - &mut ctx, - None, - vec![ - u32_ty.into(), - u64_ty.into(), - u64_ty.into(), - u64_ty.into(), - u64_ty.into(), - ], - ); - header.insert_at_back(function_region, &ctx); - let latch = BasicBlock::new(&mut ctx, None, vec![]); - latch.insert_at_back(function_region, &ctx); - let exit = BasicBlock::new(&mut ctx, None, vec![]); - exit.insert_at_back(function_region, &ctx); - - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(preheader, &ctx); - let i0 = append_mir_unsigned_constant(&mut ctx, preheader, u32_ty, 0); - Operation::new( - &mut ctx, - mir::MirGotoOp::get_concrete_op_info(), - vec![], - vec![i0, desc_a0_base, desc_b0_base, desc_a1_base, desc_b1_base], - vec![header], - 0, - ) - .insert_at_back(preheader, &ctx); - - let i = header.deref(&ctx).get_argument(0); - let desc_a0 = header.deref(&ctx).get_argument(1); - let desc_b0 = header.deref(&ctx).get_argument(2); - let desc_a1 = header.deref(&ctx).get_argument(3); - let desc_b1 = header.deref(&ctx).get_argument(4); - let bound = append_mir_unsigned_constant(&mut ctx, header, u32_ty, TRIP_COUNT); - let lt = Operation::new( - &mut ctx, - mir::MirLtOp::get_concrete_op_info(), - vec![i1_ty.into()], - vec![i, bound], - vec![], - 0, - ); - lt.insert_at_back(header, &ctx); - let lt_value = lt.deref(&ctx).get_result(0); - let not_lt = Operation::new( - &mut ctx, - mir::MirNotOp::get_concrete_op_info(), - vec![i1_ty.into()], - vec![lt_value], - vec![], - 0, - ); - not_lt.insert_at_back(header, &ctx); - let not_lt_value = not_lt.deref(&ctx).get_result(0); - let (branch_operands, segment_sizes) = - mir::MirCondBranchOp::compute_segment_sizes(vec![vec![not_lt_value], vec![], vec![]]); - let branch = Operation::new( - &mut ctx, - mir::MirCondBranchOp::get_concrete_op_info(), - vec![], - branch_operands, - vec![exit, latch], - 0, - ); - Operation::get_op::(branch, &ctx) - .expect("MirCondBranchOp") - .set_operand_segment_sizes(&ctx, segment_sizes); - branch.insert_at_back(header, &ctx); - - append_pointer_wgmma_mma(&mut ctx, latch, accumulator0, desc_a0, desc_b0); - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(latch, &ctx); - append_wgmma_wait_group_constant(&mut ctx, latch, 1); - append_pointer_wgmma_mma(&mut ctx, latch, accumulator1, desc_a1, desc_b1); - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(latch, &ctx); - append_wgmma_wait_group_constant(&mut ctx, latch, 1); - - let one = append_mir_unsigned_constant(&mut ctx, latch, u32_ty, 1); - let i_next = Operation::new( - &mut ctx, - mir::MirAddOp::get_concrete_op_info(), - vec![u32_ty.into()], - vec![i, one], - vec![], - 0, - ); - i_next.insert_at_back(latch, &ctx); - let i_next = i_next.deref(&ctx).get_result(0); - - let desc_a0_step = append_mir_unsigned_constant(&mut ctx, latch, u64_ty, DESC_A0_STEP); - let desc_a0_next = Operation::new( - &mut ctx, - mir::MirAddOp::get_concrete_op_info(), - vec![u64_ty.into()], - vec![desc_a0, desc_a0_step], - vec![], - 0, - ); - desc_a0_next.insert_at_back(latch, &ctx); - let desc_a0_next = desc_a0_next.deref(&ctx).get_result(0); - - let desc_b0_step = append_mir_unsigned_constant(&mut ctx, latch, u64_ty, DESC_B0_STEP); - let desc_b0_next = Operation::new( - &mut ctx, - mir::MirAddOp::get_concrete_op_info(), - vec![u64_ty.into()], - vec![desc_b0, desc_b0_step], - vec![], - 0, - ); - desc_b0_next.insert_at_back(latch, &ctx); - let desc_b0_next = desc_b0_next.deref(&ctx).get_result(0); - - let desc_a1_step = append_mir_unsigned_constant(&mut ctx, latch, u64_ty, DESC_A1_STEP); - let desc_a1_next = Operation::new( - &mut ctx, - mir::MirAddOp::get_concrete_op_info(), - vec![u64_ty.into()], - vec![desc_a1, desc_a1_step], - vec![], - 0, - ); - desc_a1_next.insert_at_back(latch, &ctx); - let desc_a1_next = desc_a1_next.deref(&ctx).get_result(0); - - let desc_b1_step = append_mir_unsigned_constant(&mut ctx, latch, u64_ty, DESC_B1_STEP); - let desc_b1_next = Operation::new( - &mut ctx, - mir::MirAddOp::get_concrete_op_info(), - vec![u64_ty.into()], - vec![desc_b1, desc_b1_step], - vec![], - 0, - ); - desc_b1_next.insert_at_back(latch, &ctx); - let desc_b1_next = desc_b1_next.deref(&ctx).get_result(0); - - Operation::new( - &mut ctx, - mir::MirGotoOp::get_concrete_op_info(), - vec![], - vec![ - i_next, - desc_a0_next, - desc_b0_next, - desc_a1_next, - desc_b1_next, - ], - vec![header], - 0, - ) - .insert_at_back(latch, &ctx); - - append_wgmma_wait_group_constant(&mut ctx, exit, 0); - append_return(&mut ctx, exit); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let body = lowered_kernel_body(&ctx, module_ptr); - let matching = body - .iter() - .copied() - .filter_map(|operation| { - Operation::get_op::(operation, &ctx) - .map(|inline_asm| (operation, inline_asm)) - }) - .filter(|(_, asm)| { - asm.get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .is_some_and(|template| template.contains("L__wgmma_pipeline_loop_${:uid}:")) - }) - .collect::>(); - assert_eq!( - matching.len(), - 1, - "expected one fused two-slot counted WGMMA pipeline" - ); - - let (asm_operation, asm) = &matching[0]; - let template = asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .expect("counted-pipeline WGMMA template"); - - assert_eq!(template.matches("wgmma.fence.sync.aligned").count(), 1); - assert_eq!(template.matches("wgmma.mma_async").count(), 2); - assert_eq!( - template.matches("wgmma.commit_group.sync.aligned").count(), - 2 - ); - assert_eq!( - template.matches("wgmma.wait_group.sync.aligned 1").count(), - 2 - ); - assert_eq!( - template.matches("wgmma.wait_group.sync.aligned 0").count(), - 1 - ); - assert!(template.contains("{$0, $1")); - assert!(template.contains("{$32, $33")); - assert!(template.contains("mov.u64 %desc_a0, $128;")); - assert!(template.contains("mov.u64 %desc_b0, $129;")); - assert!(template.contains("mov.u64 %desc_a1, $130;")); - assert!(template.contains("mov.u64 %desc_b1, $131;")); - assert!(template.contains("add.u64 %desc_a0, %desc_a0, $132;")); - assert!(template.contains("add.u64 %desc_b0, %desc_b0, $133;")); - assert!(template.contains("add.u64 %desc_a1, %desc_a1, $134;")); - assert!(template.contains("add.u64 %desc_b1, %desc_b1, $135;")); - assert!(template.contains("mov.u64 %remaining, $136;")); - assert!( - !template.contains(".reg .f32") - && !template.contains("ld.f32") - && !template.contains("st.f32"), - "counted pipeline must keep accumulator memory outside asm: {template}" - ); - - let mut expected_constraints = vec!["=f".to_owned(); RESULT_COUNT]; - expected_constraints.extend((0..RESULT_COUNT).map(|index| index.to_string())); - expected_constraints.extend((0..LOOP_CONTROL_COUNT).map(|_| "l".to_owned())); - expected_constraints.push("~{memory}".to_owned()); - let expected_constraints = expected_constraints.join(","); - assert_eq!( - asm.get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some(expected_constraints.as_str()) - ); - assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); - assert_eq!( - asm_operation.deref(&ctx).get_num_operands(), - RESULT_COUNT + LOOP_CONTROL_COUNT - ); - - let asm_position = body - .iter() - .position(|operation| operation == asm_operation) - .expect("counted-pipeline WGMMA asm must be in the lowered kernel body"); - let load_positions = body - .iter() - .enumerate() - .filter_map(|(index, operation)| { - Operation::get_op::(*operation, &ctx) - .is_some() - .then_some(index) - }) - .collect::>(); - assert_eq!(load_positions.len(), RESULT_COUNT); - assert!(load_positions.iter().all(|index| *index < asm_position)); - - let store_positions = body - .iter() - .enumerate() - .filter_map(|(index, operation)| { - Operation::get_op::(*operation, &ctx) - .is_some() - .then_some(index) - }) - .collect::>(); - assert_eq!(store_positions.len(), RESULT_COUNT); - assert!(store_positions.iter().all(|index| *index > asm_position)); - assert_eq!( - body.iter() - .filter(|operation| { - Operation::get_op::(**operation, &ctx).is_some() - }) - .count(), - RESULT_COUNT - ); - - Ok(()) -} - -#[test] -fn test_pointer_form_wgmma_three_slot_counted_pipeline_stays_register_resident() --> Result<(), anyhow::Error> { - const RESULT_COUNT: usize = 96; - const LOOP_CONTROL_COUNT: usize = 13; - - let mut ctx = make_test_ctx(); - let module_ptr = build_pointer_form_wgmma_counted_pipeline_case(&mut ctx, 3, &[2, 2, 2], false); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let body = lowered_kernel_body(&ctx, module_ptr); - let matching = body - .iter() - .copied() - .filter_map(|operation| { - Operation::get_op::(operation, &ctx) - .map(|inline_asm| (operation, inline_asm)) - }) - .filter(|(_, asm)| { - asm.get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .is_some_and(|template| template.contains("L__wgmma_pipeline_loop_${:uid}:")) - }) - .collect::>(); - assert_eq!( - matching.len(), - 1, - "expected one fused three-slot counted WGMMA pipeline" - ); - - let (asm_operation, asm) = &matching[0]; - let template = asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .expect("three-slot counted-pipeline WGMMA template"); - - assert_eq!(template.matches("wgmma.fence.sync.aligned").count(), 1); - assert_eq!(template.matches("wgmma.mma_async").count(), 3); - assert_eq!( - template.matches("wgmma.commit_group.sync.aligned").count(), - 3 - ); - assert_eq!( - template.matches("wgmma.wait_group.sync.aligned 2").count(), - 3 - ); - assert_eq!( - template.matches("wgmma.wait_group.sync.aligned 0").count(), - 1 - ); - assert!(template.contains("{$0, $1")); - assert!(template.contains("{$32, $33")); - assert!(template.contains("{$64, $65")); - assert!(template.contains("mov.u64 %desc_a0, $192;")); - assert!(template.contains("mov.u64 %desc_b0, $193;")); - assert!(template.contains("mov.u64 %desc_a1, $194;")); - assert!(template.contains("mov.u64 %desc_b1, $195;")); - assert!(template.contains("mov.u64 %desc_a2, $196;")); - assert!(template.contains("mov.u64 %desc_b2, $197;")); - assert!(template.contains("add.u64 %desc_a0, %desc_a0, $198;")); - assert!(template.contains("add.u64 %desc_b0, %desc_b0, $199;")); - assert!(template.contains("add.u64 %desc_a1, %desc_a1, $200;")); - assert!(template.contains("add.u64 %desc_b1, %desc_b1, $201;")); - assert!(template.contains("add.u64 %desc_a2, %desc_a2, $202;")); - assert!(template.contains("add.u64 %desc_b2, %desc_b2, $203;")); - assert!(template.contains("mov.u64 %remaining, $204;")); - assert!( - !template.contains(".reg .f32") - && !template.contains("ld.f32") - && !template.contains("st.f32"), - "three-slot counted pipeline must keep accumulator memory outside asm: {template}" - ); - - let mut expected_constraints = vec!["=f".to_owned(); RESULT_COUNT]; - expected_constraints.extend((0..RESULT_COUNT).map(|index| index.to_string())); - expected_constraints.extend((0..LOOP_CONTROL_COUNT).map(|_| "l".to_owned())); - expected_constraints.push("~{memory}".to_owned()); - let expected_constraints = expected_constraints.join(","); - assert_eq!( - asm.get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some(expected_constraints.as_str()) - ); - assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); - assert_eq!( - asm_operation.deref(&ctx).get_num_operands(), - RESULT_COUNT + LOOP_CONTROL_COUNT - ); - - let asm_position = body - .iter() - .position(|operation| operation == asm_operation) - .expect("three-slot counted-pipeline WGMMA asm must be in the lowered kernel body"); - let load_positions = body - .iter() - .enumerate() - .filter_map(|(index, operation)| { - Operation::get_op::(*operation, &ctx) - .is_some() - .then_some(index) - }) - .collect::>(); - assert_eq!(load_positions.len(), RESULT_COUNT); - assert!(load_positions.iter().all(|index| *index < asm_position)); - - let store_positions = body - .iter() - .enumerate() - .filter_map(|(index, operation)| { - Operation::get_op::(*operation, &ctx) - .is_some() - .then_some(index) - }) - .collect::>(); - assert_eq!(store_positions.len(), RESULT_COUNT); - assert!(store_positions.iter().all(|index| *index > asm_position)); - assert_eq!( - body.iter() - .filter(|operation| { - Operation::get_op::(**operation, &ctx).is_some() - }) - .count(), - RESULT_COUNT - ); - - Ok(()) -} - -#[test] -fn test_pointer_form_wgmma_counted_pipeline_rejects_wait_two_with_two_slots() { - let mut ctx = make_test_ctx(); - let module_ptr = build_pointer_form_wgmma_counted_pipeline_case(&mut ctx, 2, &[2, 2], false); - - assert!( - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).is_err(), - "wait_group<2> must require three counted-pipeline accumulator slots" - ); -} - -#[test] -fn test_pointer_form_wgmma_counted_pipeline_rejects_mixed_partial_waits() { - let mut ctx = make_test_ctx(); - let module_ptr = build_pointer_form_wgmma_counted_pipeline_case(&mut ctx, 3, &[2, 1, 2], false); - - assert!( - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).is_err(), - "all counted-pipeline stages must use the same partial-wait depth" - ); -} - -#[test] -fn test_pointer_form_wgmma_counted_pipeline_rejects_reused_accumulator_slot() { - let mut ctx = make_test_ctx(); - let module_ptr = build_pointer_form_wgmma_counted_pipeline_case(&mut ctx, 3, &[2, 2, 2], true); - - assert!( - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).is_err(), - "counted-pipeline accumulator slots must be pairwise distinct" - ); -} - -#[test] -fn test_pointer_form_wgmma_counted_k_loop_stays_register_resident() -> Result<(), anyhow::Error> { - use dialect_mir::attributes::{MirCastKindAttr, MirPointerKindAuthorityAttr}; - use dialect_mir::types::{MirArrayType, MirPointerKind, MirPtrType}; - use pliron::basic_block::BasicBlock; - use pliron::builtin::op_interfaces::OperandSegmentInterface; - use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; - - const ACCUMULATOR_LEN: usize = 32; - const LOOP_CONTROL_COUNT: usize = 5; - const TRIP_COUNT: u64 = 4; - const DESC_A_STEP: u64 = 16; - const DESC_B_STEP: u64 = 32; - - let mut ctx = make_test_ctx(); - let f32_ty = FP32Type::get(&ctx); - let row_ty = MirArrayType::get(&mut ctx, f32_ty.into(), 8); - let accumulator_ty = MirArrayType::get(&mut ctx, row_ty.into(), 4); - let accumulator_ptr_ty = MirPtrType::get_generic(&mut ctx, accumulator_ty.into(), true); - let unique_accumulator_ptr_ty: pliron::r#type::TypeHandle = MirPtrType::get_generic_with_kind( - &mut ctx, - accumulator_ty.into(), - true, - MirPointerKind::UniqueRef, - ) - .into(); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); - - let (module_ptr, preheader) = build_test_kernel( - &mut ctx, - vec![accumulator_ptr_ty.into(), u64_ty.into(), u64_ty.into()], - ); - let accumulator = preheader.deref(&ctx).get_argument(0); - let desc_a_base = preheader.deref(&ctx).get_argument(1); - let desc_b_base = preheader.deref(&ctx).get_argument(2); - - let module_region = module_ptr.deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - let function = module_block.deref(&ctx).iter(&ctx).next().unwrap(); - let function_region = function.deref(&ctx).get_region(0); - - let header = BasicBlock::new( - &mut ctx, - None, - vec![u32_ty.into(), u64_ty.into(), u64_ty.into()], - ); - header.insert_at_back(function_region, &ctx); - let body = BasicBlock::new(&mut ctx, None, vec![]); - body.insert_at_back(function_region, &ctx); - let latch = BasicBlock::new(&mut ctx, None, vec![]); - latch.insert_at_back(function_region, &ctx); - let exit = BasicBlock::new(&mut ctx, None, vec![]); - exit.insert_at_back(function_region, &ctx); - let wait_block = BasicBlock::new(&mut ctx, None, vec![]); - wait_block.insert_at_back(function_region, &ctx); - - // preheader: fence; i0 = 0; goto header(i0, desc_a_base, desc_b_base) - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(preheader, &ctx); - let i0 = append_mir_unsigned_constant(&mut ctx, preheader, u32_ty, 0); - Operation::new( - &mut ctx, - mir::MirGotoOp::get_concrete_op_info(), - vec![], - vec![i0, desc_a_base, desc_b_base], - vec![header], - 0, - ) - .insert_at_back(preheader, &ctx); - - // header(i, desc_a, desc_b): if !(i < 4) exit else body. - let i = header.deref(&ctx).get_argument(0); - let desc_a = header.deref(&ctx).get_argument(1); - let desc_b = header.deref(&ctx).get_argument(2); - let bound = append_mir_unsigned_constant(&mut ctx, header, u32_ty, TRIP_COUNT); - let lt = Operation::new( - &mut ctx, - mir::MirLtOp::get_concrete_op_info(), - vec![i1_ty.into()], - vec![i, bound], - vec![], - 0, - ); - lt.insert_at_back(header, &ctx); - let lt_value = lt.deref(&ctx).get_result(0); - let not_lt = Operation::new( - &mut ctx, - mir::MirNotOp::get_concrete_op_info(), - vec![i1_ty.into()], - vec![lt_value], - vec![], - 0, - ); - not_lt.insert_at_back(header, &ctx); - let not_lt_value = not_lt.deref(&ctx).get_result(0); - let (branch_operands, segment_sizes) = - mir::MirCondBranchOp::compute_segment_sizes(vec![vec![not_lt_value], vec![], vec![]]); - let branch = Operation::new( - &mut ctx, - mir::MirCondBranchOp::get_concrete_op_info(), - vec![], - branch_operands, - vec![exit, body], - 0, - ); - Operation::get_op::(branch, &ctx) - .expect("MirCondBranchOp") - .set_operand_segment_sizes(&ctx, segment_sizes); - branch.insert_at_back(header, &ctx); - - // body: a fresh Rust `&mut` reborrow and one WGMMA per K iteration. Rust - // calls end a MIR block, so the post-call arithmetic lives in the latch. - // The reborrow is loop-local, but its canonical storage identity is the - // preheader accumulator argument. - let retype = Operation::new( - &mut ctx, - mir::MirCastOp::get_concrete_op_info(), - vec![unique_accumulator_ptr_ty], - vec![accumulator], - vec![], - 0, - ); - mir::MirCastOp::new(retype).set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); - mir::MirCastOp::new(retype) - .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); - retype.insert_at_back(body, &ctx); - let reborrowed_accumulator = retype.deref(&ctx).get_result(0); - append_pointer_wgmma_mma(&mut ctx, body, reborrowed_accumulator, desc_a, desc_b); - Operation::new( - &mut ctx, - mir::MirGotoOp::get_concrete_op_info(), - vec![], - vec![], - vec![latch], - 0, - ) - .insert_at_back(body, &ctx); - - // latch: affine induction and descriptor recurrences, then the back edge. - let one = append_mir_unsigned_constant(&mut ctx, latch, u32_ty, 1); - let i_next = Operation::new( - &mut ctx, - mir::MirAddOp::get_concrete_op_info(), - vec![u32_ty.into()], - vec![i, one], - vec![], - 0, - ); - i_next.insert_at_back(latch, &ctx); - let i_next = i_next.deref(&ctx).get_result(0); - - let desc_a_step = append_mir_unsigned_constant(&mut ctx, latch, u64_ty, DESC_A_STEP); - let desc_a_next = Operation::new( - &mut ctx, - mir::MirAddOp::get_concrete_op_info(), - vec![u64_ty.into()], - vec![desc_a, desc_a_step], - vec![], - 0, - ); - desc_a_next.insert_at_back(latch, &ctx); - let desc_a_next = desc_a_next.deref(&ctx).get_result(0); - - let desc_b_step = append_mir_unsigned_constant(&mut ctx, latch, u64_ty, DESC_B_STEP); - let desc_b_next = Operation::new( - &mut ctx, - mir::MirAddOp::get_concrete_op_info(), - vec![u64_ty.into()], - vec![desc_b, desc_b_step], - vec![], - 0, - ); - desc_b_next.insert_at_back(latch, &ctx); - let desc_b_next = desc_b_next.deref(&ctx).get_result(0); - - Operation::new( - &mut ctx, - mir::MirGotoOp::get_concrete_op_info(), - vec![], - vec![i_next, desc_a_next, desc_b_next], - vec![header], - 0, - ) - .insert_at_back(latch, &ctx); - - // Rust calls split the exit sequence too: commit in the loop exit, then - // wait_group<0> in its unique linear successor. - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(exit, &ctx); - Operation::new( - &mut ctx, - mir::MirGotoOp::get_concrete_op_info(), - vec![], - vec![], - vec![wait_block], - 0, - ) - .insert_at_back(exit, &ctx); - append_wgmma_wait_group_constant(&mut ctx, wait_block, 0); - append_return(&mut ctx, wait_block); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let body = lowered_kernel_body(&ctx, module_ptr); - let matching = body - .iter() - .copied() - .filter_map(|operation| { - Operation::get_op::(operation, &ctx) - .map(|inline_asm| (operation, inline_asm)) - }) - .filter(|(_, asm)| { - asm.get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .is_some_and(|template| template.contains("L__wgmma_loop_${:uid}:")) - }) - .collect::>(); - assert_eq!( - matching.len(), - 1, - "expected one fused counted-loop WGMMA asm" - ); - - let (asm_operation, asm) = &matching[0]; - let template = asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .expect("counted-loop WGMMA template"); - - assert_eq!(template.matches("wgmma.fence.sync.aligned").count(), 1); - assert_eq!( - template.matches("wgmma.mma_async").count(), - 1, - "the PTX template contains one MMA instruction controlled by the internal K-loop" - ); - assert_eq!( - template.matches("wgmma.commit_group.sync.aligned").count(), - 1 - ); - assert_eq!( - template.matches("wgmma.wait_group.sync.aligned 0").count(), - 1 - ); - assert!(template.contains("mov.u64 %desc_a, $64;")); - assert!(template.contains("mov.u64 %desc_b, $65;")); - assert!(template.contains("add.u64 %desc_a, %desc_a, $66;")); - assert!(template.contains("add.u64 %desc_b, %desc_b, $67;")); - assert!(template.contains("mov.u64 %remaining, $68;")); - assert!(template.contains("@%loop_more bra.uni L__wgmma_done_${:uid};")); - assert!(template.contains("@%loop_more bra.uni L__wgmma_loop_${:uid};")); - assert!(template.contains("L__wgmma_done_${:uid}:")); - assert!( - !template.contains(".reg .f32") - && !template.contains("ld.f32") - && !template.contains("st.f32"), - "counted-loop WGMMA must keep accumulator memory outside asm: {template}" - ); - - let mut expected_constraints = vec!["=f".to_owned(); ACCUMULATOR_LEN]; - expected_constraints.extend((0..ACCUMULATOR_LEN).map(|index| index.to_string())); - expected_constraints.extend((0..LOOP_CONTROL_COUNT).map(|_| "l".to_owned())); - expected_constraints.push("~{memory}".to_owned()); - let expected_constraints = expected_constraints.join(","); - assert_eq!( - asm.get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some(expected_constraints.as_str()) - ); - assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); - assert_eq!( - asm_operation.deref(&ctx).get_num_operands(), - ACCUMULATOR_LEN + LOOP_CONTROL_COUNT - ); - - let asm_position = body - .iter() - .position(|operation| operation == asm_operation) - .expect("counted-loop WGMMA asm must be in the lowered kernel body"); - - let load_positions = body - .iter() - .enumerate() - .filter_map(|(index, operation)| { - Operation::get_op::(*operation, &ctx) - .is_some() - .then_some(index) - }) - .collect::>(); - assert_eq!( - load_positions.len(), - ACCUMULATOR_LEN, - "K-loop accumulator must be loaded exactly once, not once per iteration" - ); - assert!(load_positions.iter().all(|index| *index < asm_position)); - - let store_positions = body - .iter() - .enumerate() - .filter_map(|(index, operation)| { - Operation::get_op::(*operation, &ctx) - .is_some() - .then_some(index) - }) - .collect::>(); - assert_eq!( - store_positions.len(), - ACCUMULATOR_LEN, - "K-loop accumulator must be stored exactly once after the final wait" - ); - assert!(store_positions.iter().all(|index| *index > asm_position)); - - assert_eq!( - body.iter() - .filter( - |operation| Operation::get_op::(**operation, &ctx).is_some() - ) - .count(), - ACCUMULATOR_LEN - ); - - Ok(()) -} - -#[test] -fn test_f16_wgmma_counted_k_loop_stays_register_resident() -> Result<(), anyhow::Error> { - use dialect_mir::types::{MirArrayType, MirPtrType}; - use pliron::basic_block::BasicBlock; - use pliron::builtin::op_interfaces::OperandSegmentInterface; - use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; - - const ACCUMULATOR_LEN: usize = 32; - const LOOP_CONTROL_COUNT: usize = 5; - const TRIP_COUNT: u64 = 4; - const DESC_A_STEP: u64 = 16; - const DESC_B_STEP: u64 = 32; - - let mut ctx = make_test_ctx(); - let f32_ty = FP32Type::get(&ctx); - let row_ty = MirArrayType::get(&mut ctx, f32_ty.into(), 8); - let accumulator_ty = MirArrayType::get(&mut ctx, row_ty.into(), 4); - let accumulator_ptr_ty = MirPtrType::get_generic(&mut ctx, accumulator_ty.into(), true); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); - - let (module_ptr, preheader) = build_test_kernel( - &mut ctx, - vec![accumulator_ptr_ty.into(), u64_ty.into(), u64_ty.into()], - ); - let accumulator = preheader.deref(&ctx).get_argument(0); - let desc_a_base = preheader.deref(&ctx).get_argument(1); - let desc_b_base = preheader.deref(&ctx).get_argument(2); - - let module_region = module_ptr.deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - let function = module_block.deref(&ctx).iter(&ctx).next().unwrap(); - let function_region = function.deref(&ctx).get_region(0); - - let header = BasicBlock::new( - &mut ctx, - None, - vec![u32_ty.into(), u64_ty.into(), u64_ty.into()], - ); - header.insert_at_back(function_region, &ctx); - let latch = BasicBlock::new(&mut ctx, None, vec![]); - latch.insert_at_back(function_region, &ctx); - let exit = BasicBlock::new(&mut ctx, None, vec![]); - exit.insert_at_back(function_region, &ctx); - - // preheader: fence; i0 = 0; goto header(i0, desc_a_base, desc_b_base) - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(preheader, &ctx); - let i0 = append_mir_unsigned_constant(&mut ctx, preheader, u32_ty, 0); - Operation::new( - &mut ctx, - mir::MirGotoOp::get_concrete_op_info(), - vec![], - vec![i0, desc_a_base, desc_b_base], - vec![header], - 0, - ) - .insert_at_back(preheader, &ctx); - - // header(i, desc_a, desc_b): if !(i < 4) exit else latch. - let i = header.deref(&ctx).get_argument(0); - let desc_a = header.deref(&ctx).get_argument(1); - let desc_b = header.deref(&ctx).get_argument(2); - let bound = append_mir_unsigned_constant(&mut ctx, header, u32_ty, TRIP_COUNT); - let lt = Operation::new( - &mut ctx, - mir::MirLtOp::get_concrete_op_info(), - vec![i1_ty.into()], - vec![i, bound], - vec![], - 0, - ); - lt.insert_at_back(header, &ctx); - let lt_value = lt.deref(&ctx).get_result(0); - let not_lt = Operation::new( - &mut ctx, - mir::MirNotOp::get_concrete_op_info(), - vec![i1_ty.into()], - vec![lt_value], - vec![], - 0, - ); - not_lt.insert_at_back(header, &ctx); - let not_lt_value = not_lt.deref(&ctx).get_result(0); - let (branch_operands, segment_sizes) = - mir::MirCondBranchOp::compute_segment_sizes(vec![vec![not_lt_value], vec![], vec![]]); - let branch = Operation::new( - &mut ctx, - mir::MirCondBranchOp::get_concrete_op_info(), - vec![], - branch_operands, - vec![exit, latch], - 0, - ); - Operation::get_op::(branch, &ctx) - .expect("MirCondBranchOp") - .set_operand_segment_sizes(&ctx, segment_sizes); - branch.insert_at_back(header, &ctx); - - // latch: one WGMMA per K iteration and affine descriptor recurrences. - append_pointer_wgmma_mma_f16(&mut ctx, latch, accumulator, desc_a, desc_b); - - let one = append_mir_unsigned_constant(&mut ctx, latch, u32_ty, 1); - let i_next = Operation::new( - &mut ctx, - mir::MirAddOp::get_concrete_op_info(), - vec![u32_ty.into()], - vec![i, one], - vec![], - 0, - ); - i_next.insert_at_back(latch, &ctx); - let i_next = i_next.deref(&ctx).get_result(0); - - let desc_a_step = append_mir_unsigned_constant(&mut ctx, latch, u64_ty, DESC_A_STEP); - let desc_a_next = Operation::new( - &mut ctx, - mir::MirAddOp::get_concrete_op_info(), - vec![u64_ty.into()], - vec![desc_a, desc_a_step], - vec![], - 0, - ); - desc_a_next.insert_at_back(latch, &ctx); - let desc_a_next = desc_a_next.deref(&ctx).get_result(0); - - let desc_b_step = append_mir_unsigned_constant(&mut ctx, latch, u64_ty, DESC_B_STEP); - let desc_b_next = Operation::new( - &mut ctx, - mir::MirAddOp::get_concrete_op_info(), - vec![u64_ty.into()], - vec![desc_b, desc_b_step], - vec![], - 0, - ); - desc_b_next.insert_at_back(latch, &ctx); - let desc_b_next = desc_b_next.deref(&ctx).get_result(0); - - Operation::new( - &mut ctx, - mir::MirGotoOp::get_concrete_op_info(), - vec![], - vec![i_next, desc_a_next, desc_b_next], - vec![header], - 0, - ) - .insert_at_back(latch, &ctx); - - // exit: the only place where the asynchronous lifetime may become visible. - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(exit, &ctx); - append_wgmma_wait_group_constant(&mut ctx, exit, 0); - append_return(&mut ctx, exit); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let body = lowered_kernel_body(&ctx, module_ptr); - let matching = body - .iter() - .copied() - .filter_map(|operation| { - Operation::get_op::(operation, &ctx) - .map(|inline_asm| (operation, inline_asm)) - }) - .filter(|(_, asm)| { - asm.get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .is_some_and(|template| template.contains("L__wgmma_loop_${:uid}:")) - }) - .collect::>(); - assert_eq!( - matching.len(), - 1, - "expected one fused F16 counted-loop WGMMA asm" - ); - - let (asm_operation, asm) = &matching[0]; - let template = asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .expect("F16 counted-loop WGMMA template"); - - assert_eq!(template.matches("wgmma.fence.sync.aligned").count(), 1); - assert_eq!( - template - .matches("wgmma.mma_async.sync.aligned.m64n64k16.f32.f16.f16") - .count(), - 1, - "the F16 counted-loop template must contain exactly one F16 MMA" - ); - assert!(!template.contains(".bf16.bf16")); - assert_eq!( - template.matches("wgmma.commit_group.sync.aligned").count(), - 1 - ); - assert_eq!( - template.matches("wgmma.wait_group.sync.aligned 0").count(), - 1 - ); - assert!(template.contains("mov.u64 %desc_a, $64;")); - assert!(template.contains("mov.u64 %desc_b, $65;")); - assert!(template.contains("add.u64 %desc_a, %desc_a, $66;")); - assert!(template.contains("add.u64 %desc_b, %desc_b, $67;")); - assert!(template.contains("mov.u64 %remaining, $68;")); - assert!(template.contains("@%loop_more bra.uni L__wgmma_done_${:uid};")); - assert!(template.contains("@%loop_more bra.uni L__wgmma_loop_${:uid};")); - assert!(template.contains("L__wgmma_done_${:uid}:")); - assert!( - !template.contains(".reg .f32") - && !template.contains("ld.f32") - && !template.contains("st.f32"), - "F16 counted-loop WGMMA must keep accumulator memory outside asm: {template}" - ); - - let mut expected_constraints = vec!["=f".to_owned(); ACCUMULATOR_LEN]; - expected_constraints.extend((0..ACCUMULATOR_LEN).map(|index| index.to_string())); - expected_constraints.extend((0..LOOP_CONTROL_COUNT).map(|_| "l".to_owned())); - expected_constraints.push("~{memory}".to_owned()); - let expected_constraints = expected_constraints.join(","); - assert_eq!( - asm.get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some(expected_constraints.as_str()) - ); - assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); - assert_eq!( - asm_operation.deref(&ctx).get_num_operands(), - ACCUMULATOR_LEN + LOOP_CONTROL_COUNT - ); - - let asm_position = body - .iter() - .position(|operation| operation == asm_operation) - .expect("F16 counted-loop WGMMA asm must be in the lowered kernel body"); - - let load_positions = body - .iter() - .enumerate() - .filter_map(|(index, operation)| { - Operation::get_op::(*operation, &ctx) - .is_some() - .then_some(index) - }) - .collect::>(); - assert_eq!( - load_positions.len(), - ACCUMULATOR_LEN, - "F16 K-loop accumulator must be loaded exactly once" - ); - assert!(load_positions.iter().all(|index| *index < asm_position)); - - let store_positions = body - .iter() - .enumerate() - .filter_map(|(index, operation)| { - Operation::get_op::(*operation, &ctx) - .is_some() - .then_some(index) - }) - .collect::>(); - assert_eq!( - store_positions.len(), - ACCUMULATOR_LEN, - "F16 K-loop accumulator must be stored exactly once after the final wait" - ); - assert!(store_positions.iter().all(|index| *index > asm_position)); - - assert_eq!( - body.iter() - .filter( - |operation| Operation::get_op::(**operation, &ctx).is_some() - ) - .count(), - ACCUMULATOR_LEN - ); - - Ok(()) -} - -#[test] -fn test_tf32_wgmma_counted_k_loop_remains_unsupported() -> Result<(), anyhow::Error> { - use dialect_mir::types::{MirArrayType, MirPtrType}; - use pliron::basic_block::BasicBlock; - use pliron::builtin::op_interfaces::OperandSegmentInterface; - use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; - - const TRIP_COUNT: u64 = 4; - const DESC_A_STEP: u64 = 16; - const DESC_B_STEP: u64 = 32; - - let mut ctx = make_test_ctx(); - let f32_ty = FP32Type::get(&ctx); - let row_ty = MirArrayType::get(&mut ctx, f32_ty.into(), 8); - let accumulator_ty = MirArrayType::get(&mut ctx, row_ty.into(), 4); - let accumulator_ptr_ty = MirPtrType::get_generic(&mut ctx, accumulator_ty.into(), true); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); - - let (module_ptr, preheader) = build_test_kernel( - &mut ctx, - vec![accumulator_ptr_ty.into(), u64_ty.into(), u64_ty.into()], - ); - let accumulator = preheader.deref(&ctx).get_argument(0); - let desc_a_base = preheader.deref(&ctx).get_argument(1); - let desc_b_base = preheader.deref(&ctx).get_argument(2); - - let module_region = module_ptr.deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - let function = module_block.deref(&ctx).iter(&ctx).next().unwrap(); - let function_region = function.deref(&ctx).get_region(0); - - let header = BasicBlock::new( - &mut ctx, - None, - vec![u32_ty.into(), u64_ty.into(), u64_ty.into()], - ); - header.insert_at_back(function_region, &ctx); - let latch = BasicBlock::new(&mut ctx, None, vec![]); - latch.insert_at_back(function_region, &ctx); - let exit = BasicBlock::new(&mut ctx, None, vec![]); - exit.insert_at_back(function_region, &ctx); - - // preheader: fence; i0 = 0; goto header(i0, desc_a_base, desc_b_base) - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(preheader, &ctx); - let i0 = append_mir_unsigned_constant(&mut ctx, preheader, u32_ty, 0); - Operation::new( - &mut ctx, - mir::MirGotoOp::get_concrete_op_info(), - vec![], - vec![i0, desc_a_base, desc_b_base], - vec![header], - 0, - ) - .insert_at_back(preheader, &ctx); - - // header(i, desc_a, desc_b): if !(i < 4) exit else latch. - let i = header.deref(&ctx).get_argument(0); - let desc_a = header.deref(&ctx).get_argument(1); - let desc_b = header.deref(&ctx).get_argument(2); - let bound = append_mir_unsigned_constant(&mut ctx, header, u32_ty, TRIP_COUNT); - let lt = Operation::new( - &mut ctx, - mir::MirLtOp::get_concrete_op_info(), - vec![i1_ty.into()], - vec![i, bound], - vec![], - 0, - ); - lt.insert_at_back(header, &ctx); - let lt_value = lt.deref(&ctx).get_result(0); - let not_lt = Operation::new( - &mut ctx, - mir::MirNotOp::get_concrete_op_info(), - vec![i1_ty.into()], - vec![lt_value], - vec![], - 0, - ); - not_lt.insert_at_back(header, &ctx); - let not_lt_value = not_lt.deref(&ctx).get_result(0); - let (branch_operands, segment_sizes) = - mir::MirCondBranchOp::compute_segment_sizes(vec![vec![not_lt_value], vec![], vec![]]); - let branch = Operation::new( - &mut ctx, - mir::MirCondBranchOp::get_concrete_op_info(), - vec![], - branch_operands, - vec![exit, latch], - 0, - ); - Operation::get_op::(branch, &ctx) - .expect("MirCondBranchOp") - .set_operand_segment_sizes(&ctx, segment_sizes); - branch.insert_at_back(header, &ctx); - - // latch: one WGMMA per K iteration and affine descriptor recurrences. - append_pointer_wgmma_mma_tf32(&mut ctx, latch, accumulator, desc_a, desc_b); - - let one = append_mir_unsigned_constant(&mut ctx, latch, u32_ty, 1); - let i_next = Operation::new( - &mut ctx, - mir::MirAddOp::get_concrete_op_info(), - vec![u32_ty.into()], - vec![i, one], - vec![], - 0, - ); - i_next.insert_at_back(latch, &ctx); - let i_next = i_next.deref(&ctx).get_result(0); - - let desc_a_step = append_mir_unsigned_constant(&mut ctx, latch, u64_ty, DESC_A_STEP); - let desc_a_next = Operation::new( - &mut ctx, - mir::MirAddOp::get_concrete_op_info(), - vec![u64_ty.into()], - vec![desc_a, desc_a_step], - vec![], - 0, - ); - desc_a_next.insert_at_back(latch, &ctx); - let desc_a_next = desc_a_next.deref(&ctx).get_result(0); - - let desc_b_step = append_mir_unsigned_constant(&mut ctx, latch, u64_ty, DESC_B_STEP); - let desc_b_next = Operation::new( - &mut ctx, - mir::MirAddOp::get_concrete_op_info(), - vec![u64_ty.into()], - vec![desc_b, desc_b_step], - vec![], - 0, - ); - desc_b_next.insert_at_back(latch, &ctx); - let desc_b_next = desc_b_next.deref(&ctx).get_result(0); - - Operation::new( - &mut ctx, - mir::MirGotoOp::get_concrete_op_info(), - vec![], - vec![i_next, desc_a_next, desc_b_next], - vec![header], - 0, - ) - .insert_at_back(latch, &ctx); - - // exit: the only place where the asynchronous lifetime may become visible. - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(exit, &ctx); - append_wgmma_wait_group_constant(&mut ctx, exit, 0); - append_return(&mut ctx, exit); - - assert_wgmma_lowering_rejected( - &mut ctx, - module_ptr, - "WGMMA MMA reached lowering without deferred accumulator fusion", - ); - Ok(()) -} - -#[test] -fn test_m64n128_wgmma_noncanonical_accumulator_has_no_pointer_fallback() -> Result<(), anyhow::Error> -{ - let mut ctx = make_test_ctx(); - let (module_ptr, entry, accumulators, desc_a, desc_b, _) = - build_wgmma_pointer_test_kernel(&mut ctx, 1, vec![]); - - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_pointer_wgmma_mma_m64n128(&mut ctx, entry, accumulators[0], desc_a, desc_b); - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_wgmma_wait_group_constant(&mut ctx, entry, 0); - append_return(&mut ctx, entry); - - assert_wgmma_lowering_rejected( - &mut ctx, - module_ptr, - "WGMMA linear full-drain lowering for this variant requires a canonical [[f32; 8]; 8] accumulator", - ); - Ok(()) -} - -#[test] -fn test_f16_wgmma_noncanonical_accumulator_has_no_pointer_fallback() -> Result<(), anyhow::Error> { - let mut ctx = make_test_ctx(); - let (module_ptr, entry, accumulators, desc_a, desc_b, _) = - build_wgmma_pointer_test_kernel(&mut ctx, 1, vec![]); - - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_pointer_wgmma_mma_f16(&mut ctx, entry, accumulators[0], desc_a, desc_b); - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_wgmma_wait_group_constant(&mut ctx, entry, 0); - append_return(&mut ctx, entry); - - assert_wgmma_lowering_rejected( - &mut ctx, - module_ptr, - "WGMMA linear full-drain lowering for this variant requires a canonical [[f32; 8]; 4] accumulator", - ); - Ok(()) -} - -#[test] -fn test_tf32_wgmma_noncanonical_accumulator_has_no_pointer_fallback() -> Result<(), anyhow::Error> { - let mut ctx = make_test_ctx(); - let (module_ptr, entry, accumulators, desc_a, desc_b, _) = - build_wgmma_pointer_test_kernel(&mut ctx, 1, vec![]); - - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_pointer_wgmma_mma_tf32(&mut ctx, entry, accumulators[0], desc_a, desc_b); - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_wgmma_wait_group_constant(&mut ctx, entry, 0); - append_return(&mut ctx, entry); - - assert_wgmma_lowering_rejected( - &mut ctx, - module_ptr, - "WGMMA linear full-drain lowering for this variant requires a canonical [[f32; 8]; 4] accumulator", - ); - Ok(()) -} - -#[test] -fn test_linear_wgmma_full_drain_rejects_mixed_bf16_and_f16() -> Result<(), anyhow::Error> { - let mut ctx = make_test_ctx(); - let (module_ptr, entry, accumulators, descriptors) = - build_wgmma_canonical_pointer_test_kernel(&mut ctx, 1, 4); - let accumulator = accumulators[0]; - - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_pointer_wgmma_mma(&mut ctx, entry, accumulator, descriptors[0], descriptors[1]); - append_pointer_wgmma_mma_f16(&mut ctx, entry, accumulator, descriptors[2], descriptors[3]); - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_wgmma_wait_group_constant(&mut ctx, entry, 0); - append_return(&mut ctx, entry); - - assert_wgmma_lowering_rejected( - &mut ctx, - module_ptr, - "one linear WGMMA full-drain region cannot mix MMA variants or shapes", - ); - Ok(()) -} - -#[test] -fn test_linear_wgmma_full_drain_rejects_mixed_bf16_and_tf32() -> Result<(), anyhow::Error> { - let mut ctx = make_test_ctx(); - let (module_ptr, entry, accumulators, descriptors) = - build_wgmma_canonical_pointer_test_kernel(&mut ctx, 1, 4); - let accumulator = accumulators[0]; - - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_pointer_wgmma_mma(&mut ctx, entry, accumulator, descriptors[0], descriptors[1]); - append_pointer_wgmma_mma_tf32(&mut ctx, entry, accumulator, descriptors[2], descriptors[3]); - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_wgmma_wait_group_constant(&mut ctx, entry, 0); - append_return(&mut ctx, entry); - - assert_wgmma_lowering_rejected( - &mut ctx, - module_ptr, - "one linear WGMMA full-drain region cannot mix MMA variants or shapes", - ); - Ok(()) -} - -#[test] -fn test_linear_wgmma_full_drain_rejects_mixed_f16_and_tf32() -> Result<(), anyhow::Error> { - let mut ctx = make_test_ctx(); - let (module_ptr, entry, accumulators, descriptors) = - build_wgmma_canonical_pointer_test_kernel(&mut ctx, 1, 4); - let accumulator = accumulators[0]; - - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_pointer_wgmma_mma_f16(&mut ctx, entry, accumulator, descriptors[0], descriptors[1]); - append_pointer_wgmma_mma_tf32(&mut ctx, entry, accumulator, descriptors[2], descriptors[3]); - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_wgmma_wait_group_constant(&mut ctx, entry, 0); - append_return(&mut ctx, entry); - - assert_wgmma_lowering_rejected( - &mut ctx, - module_ptr, - "one linear WGMMA full-drain region cannot mix MMA variants or shapes", - ); - Ok(()) -} - -#[test] -fn test_f16_wgmma_partial_wait_remains_unsupported() -> Result<(), anyhow::Error> { - let mut ctx = make_test_ctx(); - let (module_ptr, entry, accumulators, descriptors) = - build_wgmma_canonical_pointer_test_kernel(&mut ctx, 1, 2); - - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_pointer_wgmma_mma_f16( - &mut ctx, - entry, - accumulators[0], - descriptors[0], - descriptors[1], - ); - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_wgmma_wait_group_constant(&mut ctx, entry, 1); - append_wgmma_wait_group_constant(&mut ctx, entry, 0); - append_return(&mut ctx, entry); - - assert_wgmma_lowering_rejected( - &mut ctx, - module_ptr, - "WGMMA deferred accumulator lowering requires wait_group<0>", - ); - Ok(()) -} - -#[test] -fn test_tf32_wgmma_partial_wait_remains_unsupported() -> Result<(), anyhow::Error> { - let mut ctx = make_test_ctx(); - let (module_ptr, entry, accumulators, descriptors) = - build_wgmma_canonical_pointer_test_kernel(&mut ctx, 1, 2); - - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_pointer_wgmma_mma_tf32( - &mut ctx, - entry, - accumulators[0], - descriptors[0], - descriptors[1], - ); - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_wgmma_wait_group_constant(&mut ctx, entry, 1); - append_wgmma_wait_group_constant(&mut ctx, entry, 0); - append_return(&mut ctx, entry); - - assert_wgmma_lowering_rejected( - &mut ctx, - module_ptr, - "WGMMA deferred accumulator lowering requires wait_group<0>", - ); - Ok(()) -} - -#[test] -fn test_pointer_form_wgmma_without_fence_is_rejected() -> Result<(), anyhow::Error> { - let mut ctx = make_test_ctx(); - let (module_ptr, entry, accumulators, desc_a, desc_b, _) = - build_wgmma_pointer_test_kernel(&mut ctx, 1, vec![]); - - append_pointer_wgmma_mma(&mut ctx, entry, accumulators[0], desc_a, desc_b); - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_wgmma_wait_group_constant(&mut ctx, entry, 0); - append_return(&mut ctx, entry); - - assert_wgmma_lowering_rejected( - &mut ctx, - module_ptr, - "WGMMA MMA reached lowering without deferred accumulator fusion", - ); - Ok(()) -} - -#[test] -fn test_deferred_wgmma_without_commit_is_rejected() -> Result<(), anyhow::Error> { - let mut ctx = make_test_ctx(); - let (module_ptr, entry, accumulators, desc_a, desc_b, _) = - build_wgmma_pointer_test_kernel(&mut ctx, 1, vec![]); - - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_pointer_wgmma_mma(&mut ctx, entry, accumulators[0], desc_a, desc_b); - append_wgmma_wait_group_constant(&mut ctx, entry, 0); - append_return(&mut ctx, entry); - - assert_wgmma_lowering_rejected( - &mut ctx, - module_ptr, - "WGMMA wait_group requires a preceding commit_group", - ); - Ok(()) -} - -#[test] -fn test_wgmma_partial_wait_without_final_wait_zero_is_rejected() -> Result<(), anyhow::Error> { - let mut ctx = make_test_ctx(); - let (module_ptr, entry, accumulators, desc_a, desc_b, _) = - build_wgmma_pointer_test_kernel(&mut ctx, 1, vec![]); - - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_pointer_wgmma_mma(&mut ctx, entry, accumulators[0], desc_a, desc_b); - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_wgmma_wait_group_constant(&mut ctx, entry, 1); - append_return(&mut ctx, entry); - - assert_wgmma_lowering_rejected( - &mut ctx, - module_ptr, - "WGMMA partial-wait pipeline requires a final wait_group<0>", - ); - Ok(()) -} - -#[test] -fn test_wgmma_partial_wait_pipeline_rejects_unsafe_accumulator_reuse() -> Result<(), anyhow::Error> -{ - let mut ctx = make_test_ctx(); - let (module_ptr, entry, accumulators, descriptors) = - build_wgmma_canonical_pointer_test_kernel(&mut ctx, 1, 4); - - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_pointer_wgmma_mma( - &mut ctx, - entry, - accumulators[0], - descriptors[0], - descriptors[1], - ); - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_pointer_wgmma_mma( - &mut ctx, - entry, - accumulators[0], - descriptors[2], - descriptors[3], - ); - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_wgmma_wait_group_constant(&mut ctx, entry, 1); - append_wgmma_wait_group_constant(&mut ctx, entry, 0); - append_return(&mut ctx, entry); - - assert_wgmma_lowering_rejected( - &mut ctx, - module_ptr, - "WGMMA partial-wait pipeline requires max_pending_groups + 1 distinct accumulator slots", - ); - Ok(()) -} - -#[test] -fn test_wgmma_wait_group_eight_is_rejected() -> Result<(), anyhow::Error> { - let mut ctx = make_test_ctx(); - let (module_ptr, entry, accumulators, desc_a, desc_b, _) = - build_wgmma_pointer_test_kernel(&mut ctx, 1, vec![]); - - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_pointer_wgmma_mma(&mut ctx, entry, accumulators[0], desc_a, desc_b); - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_wgmma_wait_group_constant(&mut ctx, entry, 8); - append_return(&mut ctx, entry); - - assert_wgmma_lowering_rejected( - &mut ctx, - module_ptr, - "WGMMA wait_group immediate must be in 0..=7", - ); - Ok(()) -} - -#[test] -fn test_deferred_wgmma_two_commits_are_rejected() -> Result<(), anyhow::Error> { - let mut ctx = make_test_ctx(); - let (module_ptr, entry, accumulators, desc_a, desc_b, _) = - build_wgmma_pointer_test_kernel(&mut ctx, 1, vec![]); - - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_pointer_wgmma_mma(&mut ctx, entry, accumulators[0], desc_a, desc_b); - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_wgmma_wait_group_constant(&mut ctx, entry, 0); - append_return(&mut ctx, entry); - - assert_wgmma_lowering_rejected( - &mut ctx, - module_ptr, - "WGMMA deferred accumulator region supports exactly one commit_group", - ); - Ok(()) -} - -#[test] -fn test_deferred_wgmma_multiple_accumulators_are_rejected() -> Result<(), anyhow::Error> { - let mut ctx = make_test_ctx(); - let (module_ptr, entry, accumulators, desc_a, desc_b, _) = - build_wgmma_pointer_test_kernel(&mut ctx, 2, vec![]); - - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_pointer_wgmma_mma(&mut ctx, entry, accumulators[0], desc_a, desc_b); - append_pointer_wgmma_mma(&mut ctx, entry, accumulators[1], desc_a, desc_b); - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_wgmma_wait_group_constant(&mut ctx, entry, 0); - append_return(&mut ctx, entry); - - assert_wgmma_lowering_rejected( - &mut ctx, - module_ptr, - "WGMMA deferred accumulator region uses more than one accumulator", - ); - Ok(()) -} - -#[test] -fn test_deferred_wgmma_mma_after_commit_is_rejected() -> Result<(), anyhow::Error> { - let mut ctx = make_test_ctx(); - let (module_ptr, entry, accumulators, desc_a, desc_b, _) = - build_wgmma_pointer_test_kernel(&mut ctx, 1, vec![]); - - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_pointer_wgmma_mma(&mut ctx, entry, accumulators[0], desc_a, desc_b); - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_pointer_wgmma_mma(&mut ctx, entry, accumulators[0], desc_a, desc_b); - append_wgmma_wait_group_constant(&mut ctx, entry, 0); - append_return(&mut ctx, entry); - - assert_wgmma_lowering_rejected( - &mut ctx, - module_ptr, - "WGMMA MMA cannot appear after commit_group in a deferred accumulator region", - ); - Ok(()) -} - -#[test] -fn test_deferred_wgmma_branch_is_rejected() -> Result<(), anyhow::Error> { - use pliron::basic_block::BasicBlock; - use pliron::builtin::op_interfaces::OperandSegmentInterface; - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let bool_ty = IntegerType::get(&ctx, 1, Signedness::Signless); - let (module_ptr, entry, accumulators, desc_a, desc_b, trailing) = - build_wgmma_pointer_test_kernel(&mut ctx, 1, vec![bool_ty.into()]); - let condition = trailing[0]; - - let module_region = module_ptr.deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - let function = module_block.deref(&ctx).iter(&ctx).next().unwrap(); - let function_region = function.deref(&ctx).get_region(0); - - let then_block = BasicBlock::new(&mut ctx, None, vec![]); - then_block.insert_at_back(function_region, &ctx); - let else_block = BasicBlock::new(&mut ctx, None, vec![]); - else_block.insert_at_back(function_region, &ctx); - - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_pointer_wgmma_mma(&mut ctx, entry, accumulators[0], desc_a, desc_b); - - let (flat_operands, segment_sizes) = - mir::MirCondBranchOp::compute_segment_sizes(vec![vec![condition], vec![], vec![]]); - let branch = Operation::new( - &mut ctx, - mir::MirCondBranchOp::get_concrete_op_info(), - vec![], - flat_operands, - vec![then_block, else_block], - 0, - ); - Operation::get_op::(branch, &ctx) - .expect("MirCondBranchOp") - .set_operand_segment_sizes(&ctx, segment_sizes); - branch.insert_at_back(entry, &ctx); - - append_return(&mut ctx, then_block); - append_return(&mut ctx, else_block); - - assert_wgmma_lowering_rejected( - &mut ctx, - module_ptr, - "unsupported operation inside WGMMA deferred accumulator region", - ); - Ok(()) -} - -#[test] -fn test_deferred_wgmma_join_is_rejected() -> Result<(), anyhow::Error> { - use pliron::basic_block::BasicBlock; - - let mut ctx = make_test_ctx(); - let (module_ptr, entry, accumulators, desc_a, desc_b, _) = - build_wgmma_pointer_test_kernel(&mut ctx, 1, vec![]); - - let module_region = module_ptr.deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - let function = module_block.deref(&ctx).iter(&ctx).next().unwrap(); - let function_region = function.deref(&ctx).get_region(0); - - let second_predecessor = BasicBlock::new(&mut ctx, None, vec![]); - second_predecessor.insert_at_back(function_region, &ctx); - let join = BasicBlock::new(&mut ctx, None, vec![]); - join.insert_at_back(function_region, &ctx); - - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_pointer_wgmma_mma(&mut ctx, entry, accumulators[0], desc_a, desc_b); - - Operation::new( - &mut ctx, - mir::MirGotoOp::get_concrete_op_info(), - vec![], - vec![], - vec![join], - 0, - ) - .insert_at_back(entry, &ctx); - - Operation::new( - &mut ctx, - mir::MirGotoOp::get_concrete_op_info(), - vec![], - vec![], - vec![join], - 0, - ) - .insert_at_back(second_predecessor, &ctx); - - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(join, &ctx); - append_wgmma_wait_group_constant(&mut ctx, join, 0); - append_return(&mut ctx, join); - - assert_wgmma_lowering_rejected( - &mut ctx, - module_ptr, - "WGMMA deferred accumulator region cannot cross a control-flow join", - ); - Ok(()) -} - -#[test] -fn test_deferred_wgmma_intervening_operation_is_rejected() -> Result<(), anyhow::Error> { - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let (module_ptr, entry, accumulators, desc_a, desc_b, _) = - build_wgmma_pointer_test_kernel(&mut ctx, 1, vec![]); - - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_pointer_wgmma_mma(&mut ctx, entry, accumulators[0], desc_a, desc_b); - - let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); - Operation::new( - &mut ctx, - nvvm::ReadPtxSregTidXOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![], - vec![], - 0, - ) - .insert_at_back(entry, &ctx); - - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_wgmma_wait_group_constant(&mut ctx, entry, 0); - append_return(&mut ctx, entry); - - assert_wgmma_lowering_rejected( - &mut ctx, - module_ptr, - "unsupported operation inside WGMMA deferred accumulator region", - ); - Ok(()) -} - -#[test] -fn test_deferred_wgmma_nested_fence_is_rejected() -> Result<(), anyhow::Error> { - let mut ctx = make_test_ctx(); - let (module_ptr, entry, accumulators, desc_a, desc_b, _) = - build_wgmma_pointer_test_kernel(&mut ctx, 1, vec![]); - - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_pointer_wgmma_mma(&mut ctx, entry, accumulators[0], desc_a, desc_b); - nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_wgmma_wait_group_constant(&mut ctx, entry, 0); - append_return(&mut ctx, entry); - - assert_wgmma_lowering_rejected( - &mut ctx, - module_ptr, - "nested WGMMA fences are not supported in one deferred accumulator region", - ); - Ok(()) -} - -/// Scoped atomic loads and stores must lower to inline PTX, not to -/// `load atomic` / `store atomic`. -/// -/// libNVVM rejects the IR-level form outright ("Atomic loads/stores are not -/// supported"), so lowering to `llvm::AtomicLoadOp` / `llvm::AtomicStoreOp` -/// produced modules that no `--materialize-cubin` build could consume, making -/// every `DeviceAtomic*::load` and `::store` call a build failure. The PTX -/// instructions themselves have existed since sm_70, so the ops lower through -/// inline assembly instead. -/// -/// This asserts the exact template, constraints and asm kind, because all -/// three are load-bearing: the scope qualifier is the whole point of the -/// feature, the constraint register class must match the operand width, and -/// `SideEffect` plus the memory clobber is what stops a publication spin being -/// hoisted out of its loop. -#[test] -fn test_scoped_atomic_load_store_lower_to_inline_ptx() -> Result<(), anyhow::Error> { - use dialect_mir::types::MirPtrType; - use dialect_nvvm::ops::atomic::{ - AtomicOrdering, AtomicScope, NvvmAtomicLoadOp, NvvmAtomicStoreOp, - }; - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = make_test_ctx(); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); - let ptr_ty = MirPtrType::get_generic(&mut ctx, u32_ty.into(), true); - let (module_ptr, entry) = - build_test_kernel(&mut ctx, vec![ptr_ty.into(), u32_ty.into(), u64_ty.into()]); - let address = entry.deref(&ctx).get_argument(0); - let val32 = entry.deref(&ctx).get_argument(1); - let val64 = entry.deref(&ctx).get_argument(2); - - // One per scope, so a regression that hardcodes a scope is caught rather - // than passing on the Device case alone. - NvvmAtomicLoadOp::build( - &mut ctx, - address, - u32_ty.into(), - AtomicOrdering::Relaxed, - AtomicScope::Device, - ) - .get_operation() - .insert_at_back(entry, &ctx); - NvvmAtomicLoadOp::build( - &mut ctx, - address, - u32_ty.into(), - AtomicOrdering::Acquire, - AtomicScope::Block, - ) - .get_operation() - .insert_at_back(entry, &ctx); - NvvmAtomicLoadOp::build( - &mut ctx, - address, - u64_ty.into(), - AtomicOrdering::Relaxed, - AtomicScope::System, - ) - .get_operation() - .insert_at_back(entry, &ctx); - NvvmAtomicStoreOp::build( - &mut ctx, - val32, - address, - AtomicOrdering::Release, - AtomicScope::Device, - ) - .get_operation() - .insert_at_back(entry, &ctx); - NvvmAtomicStoreOp::build( - &mut ctx, - val64, - address, - AtomicOrdering::Relaxed, - AtomicScope::Device, - ) - .get_operation() - .insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let mut lowered = Vec::new(); - let module_region = module_ptr.deref(&ctx).get_region(0); - let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - for op in module_block.deref(&ctx).iter(&ctx) { - let Some(function) = Operation::get_op::(op, &ctx) else { - continue; - }; - if function.get_symbol_name(&ctx).to_string() != "kernel_func" { - continue; - } - let body = function.get_operation().deref(&ctx).get_region(0); - for block in body.deref(&ctx).iter(&ctx) { - for body_op in block.deref(&ctx).iter(&ctx) { - // The IR-level forms are what libNVVM rejects; neither may survive. - assert!( - Operation::get_op::(body_op, &ctx).is_none(), - "scoped atomic load must not lower to `load atomic`: libNVVM rejects it" - ); - assert!( - Operation::get_op::(body_op, &ctx).is_none(), - "scoped atomic store must not lower to `store atomic`: libNVVM rejects it" - ); - let Some(asm) = Operation::get_op::(body_op, &ctx) else { - continue; - }; - let template = asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .unwrap_or_default(); - if !template.starts_with("ld.") && !template.starts_with("st.") { - continue; - } - lowered.push(( - template, - asm.get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .unwrap_or_default(), - llvm::asm_kind(&ctx, &asm), - )); - } - } - } - - // Every access keeps the `~{memory}` clobber, Relaxed included. Without - // it LLVM may move plain loads and stores of the same address across the - // asm, breaking the single-thread coherence Rust still guarantees for - // Relaxed atomics; libcu++ keeps the clobber on its relaxed accesses too. - let expected = [ - ("ld.relaxed.gpu.b32 $0, [$1];", "=r,l,~{memory}"), - ("ld.acquire.cta.b32 $0, [$1];", "=r,l,~{memory}"), - ("ld.relaxed.sys.b64 $0, [$1];", "=l,l,~{memory}"), - ("st.release.gpu.b32 [$0], $1;", "l,r,~{memory}"), - ("st.relaxed.gpu.b64 [$0], $1;", "l,l,~{memory}"), - ]; - assert_eq!( - lowered.len(), - expected.len(), - "expected one inline-asm op per scoped atomic load/store, got {lowered:?}" - ); - for (template, constraints) in expected { - let found = lowered - .iter() - .find(|(t, _, _)| t == template) - .unwrap_or_else(|| panic!("missing lowering for {template}; got {lowered:?}")); - assert_eq!(found.1, constraints, "constraints for {template}"); - // Not Convergent: these are per-thread accesses, not warp-synchronous. - // SideEffect with the memory clobber is what keeps a spin re-reading. - assert_eq!( - found.2, - llvm::AsmKind::SideEffect, - "asm kind for {template}" - ); - } - Ok(()) -} - -/// Orderings a PTX load or store cannot carry must be rejected, not silently -/// weakened. -/// -/// Acquire is load-only, Release is store-only, and AcqRel makes no sense on -/// a single access. Approximating any of them with `relaxed` would be a -/// correctness bug that no runtime test would catch, because the wrong answer -/// only appears under contention on hardware that reorders. -#[test] -fn test_scoped_atomic_rejects_inexpressible_orderings() -> Result<(), anyhow::Error> { - use dialect_mir::types::MirPtrType; - use dialect_nvvm::ops::atomic::{ - AtomicOrdering, AtomicScope, NvvmAtomicLoadOp, NvvmAtomicStoreOp, - }; - use pliron::builtin::types::{IntegerType, Signedness}; - - for (is_load, ordering) in [ - (true, AtomicOrdering::Release), // release is store-only - (true, AtomicOrdering::AcqRel), - (false, AtomicOrdering::Acquire), // acquire is load-only - (false, AtomicOrdering::AcqRel), - ] { - let mut ctx = make_test_ctx(); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let ptr_ty = MirPtrType::get_generic(&mut ctx, u32_ty.into(), true); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![ptr_ty.into(), u32_ty.into()]); - let address = entry.deref(&ctx).get_argument(0); - let val = entry.deref(&ctx).get_argument(1); - - if is_load { - NvvmAtomicLoadOp::build( - &mut ctx, - address, - u32_ty.into(), - ordering.clone(), - AtomicScope::Device, - ) - .get_operation() - .insert_at_back(entry, &ctx); - } else { - NvvmAtomicStoreOp::build( - &mut ctx, - val, - address, - ordering.clone(), - AtomicScope::Device, - ) - .get_operation() - .insert_at_back(entry, &ctx); - } - append_return(&mut ctx, entry); - - let result = mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr); - assert!( - result.is_err(), - "{} with {ordering:?} ordering must be rejected, not approximated", - if is_load { "load" } else { "store" } - ); - } - Ok(()) -} - -/// SeqCst load and store lower to one asm op whose template fuses the -/// `fence.sc` with the access, at every scope. -/// -/// PTX has no sequentially consistent load or store instruction. libcu++ maps -/// a SeqCst load to `fence.sc.{scope}` followed by an acquire load at the -/// same scope, and a SeqCst store to `fence.sc.{scope}` followed by a release -/// store. Fusing the fence into the same template keeps the pair inseparable. -#[test] -fn test_seqcst_atomic_load_store_fuse_fence_into_template() -> Result<(), anyhow::Error> { - use dialect_mir::types::MirPtrType; - use dialect_nvvm::ops::atomic::{ - AtomicOrdering, AtomicScope, NvvmAtomicLoadOp, NvvmAtomicStoreOp, - }; - use pliron::builtin::types::{IntegerType, Signedness}; - - for (scope, ptx) in [ - (AtomicScope::Device, "gpu"), - (AtomicScope::Block, "cta"), - (AtomicScope::System, "sys"), - ] { - let mut ctx = make_test_ctx(); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let ptr_ty = MirPtrType::get_generic(&mut ctx, u32_ty.into(), true); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![ptr_ty.into(), u32_ty.into()]); - let address = entry.deref(&ctx).get_argument(0); - let val = entry.deref(&ctx).get_argument(1); - - NvvmAtomicLoadOp::build( - &mut ctx, - address, - u32_ty.into(), - AtomicOrdering::SeqCst, - scope.clone(), - ) - .get_operation() - .insert_at_back(entry, &ctx); - NvvmAtomicStoreOp::build(&mut ctx, val, address, AtomicOrdering::SeqCst, scope) - .get_operation() - .insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let mut lowered = Vec::new(); - for op in lowered_kernel_body(&ctx, module_ptr) { - let Some(asm) = Operation::get_op::(op, &ctx) else { - continue; - }; - lowered.push(( - asm.get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .unwrap_or_default(), - asm.get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .unwrap_or_default(), - )); - } - let expected = vec![ - ( - format!("fence.sc.{ptx}; ld.acquire.{ptx}.b32 $0, [$1];"), - "=r,l,~{memory}".to_string(), - ), - ( - format!("fence.sc.{ptx}; st.release.{ptx}.b32 [$0], $1;"), - "l,r,~{memory}".to_string(), - ), - ]; - assert_eq!(lowered, expected, "SeqCst lowering at scope {ptx}"); - } - Ok(()) -} - -/// Float atomic loads and stores travel through integer registers with an -/// LLVM-level bitcast on each side. -/// -/// The register classes in the constraints are integer (`h`/`r`/`l`), so -/// handing llc a float-typed asm operand or result is a constraint mismatch. -/// The store bitcasts the float to the same-width integer before the asm; the -/// load returns the integer and bitcasts it back to the float type. The f16 -/// case also covers the `b16`/`h` arm end to end, which is what makes -/// `DeviceAtomicF16::load`/`store` compile at all. -#[test] -fn test_float_atomic_load_store_bitcast_through_integer_registers() -> Result<(), anyhow::Error> { - use dialect_mir::types::{MirFP16Type, MirPtrType}; - use dialect_nvvm::ops::atomic::{ - AtomicOrdering, AtomicScope, NvvmAtomicLoadOp, NvvmAtomicStoreOp, - }; - use llvm_export::types as llvm_types; - use pliron::builtin::types::{FP32Type, IntegerType}; - use pliron::r#type::Typed; - - let is_expected_float = |ctx: &Context, ty: pliron::r#type::TypeHandle, width: u32| -> bool { - let ty_ref = ty.deref(ctx); - match width { - 16 => ty_ref.is::(), - 32 => ty_ref.is::(), - _ => false, - } - }; - let integer_width_of = |ctx: &Context, ty: pliron::r#type::TypeHandle| -> Option { - ty.deref(ctx) - .downcast_ref::() - .map(IntegerType::width) - }; - - for (is_f16, ptx_ty, reg, width) in [(false, "b32", "r", 32u32), (true, "b16", "h", 16u32)] { - let mut ctx = make_test_ctx(); - let elem_ty: pliron::r#type::TypeHandle = if is_f16 { - MirFP16Type::get(&ctx).into() - } else { - FP32Type::get(&ctx).into() - }; - let ptr_ty = MirPtrType::get_generic(&mut ctx, elem_ty, true); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![ptr_ty.into(), elem_ty]); - let address = entry.deref(&ctx).get_argument(0); - let val = entry.deref(&ctx).get_argument(1); - - NvvmAtomicLoadOp::build( - &mut ctx, - address, - elem_ty, - AtomicOrdering::Relaxed, - AtomicScope::Device, - ) - .get_operation() - .insert_at_back(entry, &ctx); - NvvmAtomicStoreOp::build( - &mut ctx, - val, - address, - AtomicOrdering::Relaxed, - AtomicScope::Device, - ) - .get_operation() - .insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let mut ld_asm = None; - let mut st_asm = None; - let mut bitcasts = Vec::new(); - for op in lowered_kernel_body(&ctx, module_ptr) { - if let Some(asm) = Operation::get_op::(op, &ctx) { - let template = asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .unwrap_or_default(); - let constraints = asm - .get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .unwrap_or_default(); - if template.starts_with("ld.") { - ld_asm = Some((op, template, constraints)); - } else if template.starts_with("st.") { - st_asm = Some((op, template, constraints)); - } - } - if Operation::get_op::(op, &ctx).is_some() { - bitcasts.push(op); - } - } - - // Load direction: asm produces the staging integer, then one bitcast - // turns it back into the float. - let (ld_op, ld_template, ld_constraints) = - ld_asm.expect("float atomic load must lower to inline asm"); - assert_eq!(ld_template, format!("ld.relaxed.gpu.{ptx_ty} $0, [$1];")); - assert_eq!(ld_constraints, format!("={reg},l,~{{memory}}")); - let ld_result = ld_op.deref(&ctx).get_result(0); - assert_eq!( - integer_width_of(&ctx, ld_result.get_type(&ctx)), - Some(width), - "load asm must produce the staging integer, not the float" - ); - let load_cast = bitcasts - .iter() - .copied() - .find(|&cast| cast.deref(&ctx).get_operand(0) == ld_result) - .expect("load asm result must be bitcast back to the float type"); - let load_cast_ty = load_cast.deref(&ctx).get_result(0).get_type(&ctx); - assert!( - is_expected_float(&ctx, load_cast_ty, width), - "load bitcast must produce the float type" - ); - - // Store direction: the float is bitcast to the staging integer, and - // that integer is what the asm consumes. - let (st_op, st_template, st_constraints) = - st_asm.expect("float atomic store must lower to inline asm"); - assert_eq!(st_template, format!("st.relaxed.gpu.{ptx_ty} [$0], $1;")); - assert_eq!(st_constraints, format!("l,{reg},~{{memory}}")); - let stored = st_op.deref(&ctx).get_operand(1); - assert_eq!( - integer_width_of(&ctx, stored.get_type(&ctx)), - Some(width), - "store asm must consume the staging integer, not the float" - ); - let store_cast = bitcasts - .iter() - .copied() - .find(|&cast| cast.deref(&ctx).get_result(0) == stored) - .expect("store asm value must come from a float-to-integer bitcast"); - let store_src_ty = store_cast.deref(&ctx).get_operand(0).get_type(&ctx); - assert!( - is_expected_float(&ctx, store_src_ty, width), - "store bitcast must consume the float type" - ); - } - Ok(()) -} - -/// A SeqCst fence must go through the typed NVVM intrinsic, and a weaker one -/// through inline PTX. -/// -/// PTX defines `membar.level` as a synonym for `fence.sc.level`, so -/// `llvm.nvvm.membar.{cta,gl,sys}` is an exact match for SeqCst and is the route -/// the rest of the crate already uses for fences. No intrinsic exists for a -/// weaker fence, and emitting `membar` for AcqRel would silently upgrade the -/// caller's request to sequential consistency, so that case emits -/// `fence.acq_rel.{scope}` directly. -#[test] -fn test_fence_uses_membar_intrinsic_only_for_seqcst() -> Result<(), anyhow::Error> { - use dialect_mir::types::MirPtrType; - use dialect_nvvm::ops::atomic::{AtomicOrdering, AtomicRmwKind, AtomicScope, NvvmAtomicRmwOp}; - use pliron::builtin::types::{IntegerType, Signedness}; - - for (ordering, want_intrinsic, want_asm) in [ - (AtomicOrdering::SeqCst, Some("llvm_nvvm_membar_gl"), None), - (AtomicOrdering::AcqRel, None, Some("fence.acq_rel.gpu;")), - ] { - let mut ctx = make_test_ctx(); - let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); - let ptr_ty = MirPtrType::get_generic(&mut ctx, u32_ty.into(), true); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![ptr_ty.into(), u32_ty.into()]); - let address = entry.deref(&ctx).get_argument(0); - let addend = entry.deref(&ctx).get_argument(1); - - NvvmAtomicRmwOp::build( - &mut ctx, - address, - addend, - u32_ty.into(), - AtomicRmwKind::Add, - ordering.clone(), - AtomicScope::Device, - ) - .get_operation() - .insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let mut saw_intrinsic = false; - let mut saw_asm = None; - for op in lowered_kernel_body(&ctx, module_ptr) { - if let Some(call) = Operation::get_op::(op, &ctx) - && let CallOpCallable::Direct(callee) = call.callee(&ctx) - && callee.to_string().contains("membar") - { - saw_intrinsic = true; - } - if let Some(asm) = Operation::get_op::(op, &ctx) { - let t = asm - .get_attr_inline_asm_template(&ctx) - .map(|v| String::from((*v).clone())) - .unwrap_or_default(); - if t.starts_with("fence.") { - saw_asm = Some(t); - } - } - } - assert_eq!( - saw_intrinsic, - want_intrinsic.is_some(), - "{ordering:?}: membar intrinsic presence" - ); - assert_eq!(saw_asm.as_deref(), want_asm, "{ordering:?}: inline fence"); - // A SeqCst fence must never also emit `membar` as assembly: that would - // mean both routes fired. - if want_intrinsic.is_some() { - assert!(saw_asm.is_none(), "{ordering:?}: emitted both routes"); - } - } - Ok(()) -} - -/// A first-class atomic fence must lower through the same libNVVM-safe routes -/// used by the RMW ordering workaround. -/// -/// `core::sync::atomic::fence` imports with system scope, so the source-level -/// path depends on this exact mapping: Acquire/Release/AcqRel become -/// `fence.acq_rel.sys`, while SeqCst becomes `llvm.nvvm.membar.sys`. -#[test] -fn test_first_class_atomic_fence_lowers_at_system_scope() -> Result<(), anyhow::Error> { - use dialect_nvvm::ops::atomic::{AtomicOrdering, AtomicScope, NvvmAtomicFenceOp}; - - for ordering in [ - AtomicOrdering::Acquire, - AtomicOrdering::Release, - AtomicOrdering::AcqRel, - AtomicOrdering::SeqCst, - ] { - let mut ctx = make_test_ctx(); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![]); - - NvvmAtomicFenceOp::build(&mut ctx, ordering.clone(), AtomicScope::System) - .get_operation() - .insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .map_err(|error| anyhow::anyhow!("{error}"))?; - - let mut membar_calls = 0usize; - let mut fence_asm = Vec::new(); - for op in lowered_kernel_body(&ctx, module_ptr) { - assert!( - Operation::get_op::(op, &ctx).is_none(), - "first-class atomic fence must be fully consumed by lowering" - ); - - if let Some(call) = Operation::get_op::(op, &ctx) - && let CallOpCallable::Direct(callee) = call.callee(&ctx) - && callee.to_string() == "llvm_nvvm_membar_sys" - { - membar_calls += 1; - } - - let Some(asm) = Operation::get_op::(op, &ctx) else { - continue; - }; - let template = asm - .get_attr_inline_asm_template(&ctx) - .map(|value| String::from((*value).clone())) - .unwrap_or_default(); - if template.starts_with("fence.") { - assert_eq!( - asm.get_attr_inline_asm_constraints(&ctx) - .map(|value| String::from((*value).clone())) - .as_deref(), - Some("~{memory}") - ); - assert_eq!(llvm::asm_kind(&ctx, &asm), llvm::AsmKind::SideEffect); - fence_asm.push(template); - } - } - - match ordering { - AtomicOrdering::SeqCst => { - assert_eq!(membar_calls, 1, "SeqCst must use membar.sys"); - assert!( - fence_asm.is_empty(), - "SeqCst must not also emit an inline PTX fence" - ); - } - AtomicOrdering::Acquire | AtomicOrdering::Release | AtomicOrdering::AcqRel => { - assert_eq!(membar_calls, 0, "weak fences must not use membar.sys"); - assert_eq!( - fence_asm, - vec!["fence.acq_rel.sys;".to_owned()], - "{ordering:?} must lower to one system-scope acq_rel fence" - ); - } - AtomicOrdering::Relaxed => unreachable!("Relaxed is not a valid Rust fence ordering"), - } - } - - Ok(()) -} - -#[test] -fn test_first_class_atomic_fence_rejects_relaxed() -> Result<(), anyhow::Error> { - use dialect_nvvm::ops::atomic::{AtomicOrdering, AtomicScope, NvvmAtomicFenceOp}; - - let mut ctx = make_test_ctx(); - let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![]); - NvvmAtomicFenceOp::build(&mut ctx, AtomicOrdering::Relaxed, AtomicScope::System) - .get_operation() - .insert_at_back(entry, &ctx); - append_return(&mut ctx, entry); - - let result = mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr); - assert!( - result.is_err(), - "Relaxed is not a valid atomic fence ordering and must be rejected" - ); - Ok(()) -} - -/// Builds a module holding a single `mir.func` whose body calls the `bswap` -/// placeholder on a `u8` argument, with the call result carrying `result_ty`. -/// Shared by the valid (u8 result) and malformed (tuple result) 8-bit bswap -/// tests below. -fn build_bswap8_call_module( - ctx: &mut Context, - result_ty: pliron::r#type::TypeHandle, -) -> pliron::context::Ptr { - use pliron::basic_block::BasicBlock; - use pliron::builtin::attributes::StringAttr; - use pliron::builtin::attributes::TypeAttr; - use pliron::builtin::types::{FunctionType, IntegerType, Signedness}; - - let u8_ty = IntegerType::get(ctx, 8, Signedness::Unsigned); - - let module = ModuleOp::new(ctx, "m".try_into().unwrap()); - let module_ptr = module.get_operation(); - let module_block = module - .get_operation() - .deref(ctx) - .get_region(0) - .deref(ctx) - .iter(ctx) - .next() - .unwrap(); - - let func_ty = FunctionType::get(ctx, vec![u8_ty.into()], vec![]); - let func_op = Operation::new( - ctx, - mir::MirFuncOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 1, - ); - let func = mir::MirFuncOp::new(ctx, func_op, TypeAttr::new(func_ty.into())); - func.set_symbol_name(ctx, "f".try_into().unwrap()); - { - let region = func.get_operation().deref(ctx).get_region(0); - let block = BasicBlock::new(ctx, None, vec![u8_ty.into()]); - block.insert_at_back(region, ctx); - let arg = block.deref(ctx).get_argument(0); - - let call_ptr = Operation::new( - ctx, - mir::MirCallOp::get_concrete_op_info(), - vec![result_ty], - vec![arg], - vec![], - 0, - ); - let call = mir::MirCallOp::new(call_ptr); - call.set_attr_callee( - ctx, - StringAttr::new(dialect_mir::rust_intrinsics::CALLEE_BSWAP.to_string()), - ); - call_ptr.insert_at_back(block, ctx); - - let ret_op = Operation::new( - ctx, - mir::MirReturnOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 0, - ); - ret_op.insert_at_back(block, ctx); - } - func.get_operation().insert_at_back(module_block, ctx); - - module_ptr -} - -/// An 8-bit `bswap` takes an early bitcast path, since Rust's semantics for a -/// single byte are identity and LLVM has no useful intrinsic for it. That path -/// returned before reaching the cast every other arm goes through, so a result -/// type that was not an 8-bit integer produced `bitcast i8 to `, -/// which LLVM rejects when the module is read back. -/// -/// The known producer of this shape is an importer bug that hands the call -/// the type of the destination *local* rather than of the projected place it -/// writes, so `RET.1 = bswap(x)` on a `(f64, u8)` return arrives here -/// carrying the whole tuple. That bug is the importer's to fix, and whatever -/// the importer does, a malformed producer can still build such a `mir.call`. -/// Lowering cannot repair it, and refusing it is what keeps the emitted -/// module valid. -#[test] -fn bswap8_with_an_aggregate_result_is_refused_rather_than_bitcast() -> Result<(), anyhow::Error> { - use dialect_mir::types::MirTupleType; - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - mir_lower::register(&mut ctx); - - let u8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); - let f64_ty = pliron::builtin::types::FP64Type::get(&ctx); - // The tuple result is the whole point: an 8-bit operand with a result - // the bitcast cannot legally produce. - let tuple_ty = MirTupleType::get(&mut ctx, vec![f64_ty.into(), u8_ty.into()]); - let module_ptr = build_bswap8_call_module(&mut ctx, tuple_ty.into()); - - let result = mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr); - let error = - result.expect_err("an 8-bit bswap with a tuple result must be refused, not lowered"); - assert!( - error - .to_string() - .contains("expected integer type for Rust bit intrinsic"), - "unexpected error: {error}" - ); - Ok(()) -} - -/// The valid counterpart: an 8-bit `bswap` whose result is a plain `u8` must -/// still take the identity-bitcast path and lower cleanly, so the guard above -/// cannot over-reject the case Rust actually produces for `u8::swap_bytes`. -#[test] -fn bswap8_with_a_u8_result_lowers_to_an_identity_bitcast() -> Result<(), anyhow::Error> { - use pliron::builtin::types::{IntegerType, Signedness}; - - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - dialect_nvvm::register(&mut ctx); - mir_lower::register(&mut ctx); - - let u8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); - let module_ptr = build_bswap8_call_module(&mut ctx, u8_ty.into()); - - mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) - .expect("an 8-bit bswap with a u8 result must lower"); - - // The call must have become a bitcast, not an intrinsic call: there is no - // llvm.bswap.i8. - let mut found_bitcast = false; - let module_block = module_ptr - .deref(&ctx) - .get_region(0) - .deref(&ctx) - .iter(&ctx) - .next() - .unwrap(); - for op in module_block.deref(&ctx).iter(&ctx) { - let Some(func_op) = Operation::get_op::(op, &ctx) else { - continue; - }; - let func_region = func_op.get_operation().deref(&ctx).get_region(0); - for func_block in func_region.deref(&ctx).iter(&ctx) { - for body_op in func_block.deref(&ctx).iter(&ctx) { - assert!( - Operation::get_op::(body_op, &ctx).is_none(), - "an 8-bit bswap must not lower to an intrinsic call" - ); - if Operation::get_op::(body_op, &ctx).is_some() { - found_bitcast = true; - } - } - } - } - assert!( - found_bitcast, - "an 8-bit bswap with a u8 result must lower to an identity bitcast" - ); - Ok(()) -} diff --git a/crates/mir-lower/tests/lowering_test/atomics_fences.rs b/crates/mir-lower/tests/lowering_test/atomics_fences.rs new file mode 100644 index 0000000000..367c0d623c --- /dev/null +++ b/crates/mir-lower/tests/lowering_test/atomics_fences.rs @@ -0,0 +1,799 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use dialect_nvvm::ops as nvvm; +use llvm_export::ops as llvm; +use pliron::builtin::op_interfaces::{CallOpCallable, CallOpInterface, SymbolOpInterface}; +use pliron::context::Context; +use pliron::linked_list::ContainsLinkedList; +use pliron::op::Op; +use pliron::operation::Operation; + +use crate::common::{append_return, build_test_kernel, lowered_kernel_body, make_test_ctx}; + +#[test] +fn test_packed_atomic_add_lowers_to_exact_side_effecting_ptx() -> Result<(), anyhow::Error> { + use dialect_mir::types::MirPtrType; + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let ptr_ty = MirPtrType::get_generic(&mut ctx, u32_ty.into(), true); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![ptr_ty.into(), u32_ty.into()]); + let address = entry.deref(&ctx).get_argument(0); + let addend = entry.deref(&ctx).get_argument(1); + + for op_info in [ + nvvm::NvvmAtomAddF16x2Op::get_concrete_op_info(), + nvvm::NvvmAtomAddBf16x2Op::get_concrete_op_info(), + ] { + Operation::new( + &mut ctx, + op_info, + vec![u32_ty.into()], + vec![address, addend], + vec![], + 0, + ) + .insert_at_back(entry, &ctx); + } + for format in [ + nvvm::PackedAtomicFormatAttr::F16x2, + nvvm::PackedAtomicFormatAttr::Bf16x2, + ] { + nvvm::PackedAtomicAddOp::build(&mut ctx, address, addend, format) + .insert_at_back(entry, &ctx); + } + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let expected = [ + "atom.global.add.noftz.f16x2 $0, [$1], $2;", + "atom.global.add.noftz.bf16x2 $0, [$1], $2;", + ]; + let module_region = module_ptr.deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + let mut lowered = Vec::new(); + + for op in module_block.deref(&ctx).iter(&ctx) { + let Some(function) = Operation::get_op::(op, &ctx) else { + continue; + }; + if function.get_symbol_name(&ctx).to_string() != "kernel_func" { + continue; + } + let body = function.get_operation().deref(&ctx).get_region(0); + for block in body.deref(&ctx).iter(&ctx) { + for body_op in block.deref(&ctx).iter(&ctx) { + let Some(asm) = Operation::get_op::(body_op, &ctx) else { + continue; + }; + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .unwrap_or_default(); + assert!( + !template.contains("atom.cas"), + "packed atomic exact-native lowering must not use a CAS loop: {template}" + ); + if !template.starts_with("atom.global.add.noftz.") { + continue; + } + lowered.push(( + template, + asm.get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .unwrap_or_default(), + llvm::asm_kind(&ctx, &asm), + body_op.deref(&ctx).get_num_operands(), + body_op.deref(&ctx).get_num_results(), + )); + } + } + } + + assert_eq!(lowered.len(), expected.len() * 2); + for instruction in expected { + let matches: Vec<_> = lowered + .iter() + .filter(|(template, _, _, _, _)| template == instruction) + .collect(); + assert_eq!( + matches.len(), + 2, + "legacy and generated paths must have exact PTX parity for {instruction}" + ); + for (_, constraints, kind, operands, results) in matches { + assert_eq!(constraints, "=r,l,r,~{memory}"); + assert_eq!(*kind, llvm::AsmKind::SideEffect); + assert_eq!(*operands, 2); + assert_eq!(*results, 1); + } + } + + Ok(()) +} + +#[test] +fn test_generated_packed_atomic_add_libnvvm_route_is_exact() -> Result<(), anyhow::Error> { + use dialect_mir::types::MirPtrType; + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let ptr_ty = MirPtrType::get_generic(&mut ctx, u32_ty.into(), true); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![ptr_ty.into(), u32_ty.into()]); + let address = entry.deref(&ctx).get_argument(0); + let addend = entry.deref(&ctx).get_argument(1); + for format in [ + nvvm::PackedAtomicFormatAttr::F16x2, + nvvm::PackedAtomicFormatAttr::Bf16x2, + ] { + nvvm::PackedAtomicAddOp::build(&mut ctx, address, addend, format) + .insert_at_back(entry, &ctx); + } + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: mir_lower::IntrinsicBackend::LibNvvm, + ..Default::default() + }, + ) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let lowered = lowered_kernel_body(&ctx, module_ptr) + .into_iter() + .filter_map(|op| { + let asm = Operation::get_op::(op, &ctx)?; + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone()))?; + template.starts_with("atom.global.add.noftz.").then(|| { + ( + template, + asm.get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())), + llvm::asm_kind(&ctx, &asm), + ) + }) + }) + .collect::>(); + assert_eq!(lowered.len(), 2); + assert!( + lowered + .iter() + .any(|(template, _, _)| { template == "atom.global.add.noftz.f16x2 $0, [$1], $2;" }) + ); + assert!( + lowered + .iter() + .any(|(template, _, _)| { template == "atom.global.add.noftz.bf16x2 $0, [$1], $2;" }) + ); + for (_, constraints, kind) in lowered { + assert_eq!(constraints.as_deref(), Some("=r,l,r,~{memory}")); + assert_eq!(kind, llvm::AsmKind::SideEffect); + } + Ok(()) +} + +/// Scoped atomic loads and stores must lower to inline PTX, not to +/// `load atomic` / `store atomic`. +/// +/// libNVVM rejects the IR-level form outright ("Atomic loads/stores are not +/// supported"), so lowering to `llvm::AtomicLoadOp` / `llvm::AtomicStoreOp` +/// produced modules that no `--materialize-cubin` build could consume, making +/// every `DeviceAtomic*::load` and `::store` call a build failure. The PTX +/// instructions themselves have existed since sm_70, so the ops lower through +/// inline assembly instead. +/// +/// This asserts the exact template, constraints and asm kind, because all +/// three are load-bearing: the scope qualifier is the whole point of the +/// feature, the constraint register class must match the operand width, and +/// `SideEffect` plus the memory clobber is what stops a publication spin being +/// hoisted out of its loop. +#[test] +fn test_scoped_atomic_load_store_lower_to_inline_ptx() -> Result<(), anyhow::Error> { + use dialect_mir::types::MirPtrType; + use dialect_nvvm::ops::atomic::{ + AtomicOrdering, AtomicScope, NvvmAtomicLoadOp, NvvmAtomicStoreOp, + }; + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let ptr_ty = MirPtrType::get_generic(&mut ctx, u32_ty.into(), true); + let (module_ptr, entry) = + build_test_kernel(&mut ctx, vec![ptr_ty.into(), u32_ty.into(), u64_ty.into()]); + let address = entry.deref(&ctx).get_argument(0); + let val32 = entry.deref(&ctx).get_argument(1); + let val64 = entry.deref(&ctx).get_argument(2); + + // One per scope, so a regression that hardcodes a scope is caught rather + // than passing on the Device case alone. + NvvmAtomicLoadOp::build( + &mut ctx, + address, + u32_ty.into(), + AtomicOrdering::Relaxed, + AtomicScope::Device, + ) + .get_operation() + .insert_at_back(entry, &ctx); + NvvmAtomicLoadOp::build( + &mut ctx, + address, + u32_ty.into(), + AtomicOrdering::Acquire, + AtomicScope::Block, + ) + .get_operation() + .insert_at_back(entry, &ctx); + NvvmAtomicLoadOp::build( + &mut ctx, + address, + u64_ty.into(), + AtomicOrdering::Relaxed, + AtomicScope::System, + ) + .get_operation() + .insert_at_back(entry, &ctx); + NvvmAtomicStoreOp::build( + &mut ctx, + val32, + address, + AtomicOrdering::Release, + AtomicScope::Device, + ) + .get_operation() + .insert_at_back(entry, &ctx); + NvvmAtomicStoreOp::build( + &mut ctx, + val64, + address, + AtomicOrdering::Relaxed, + AtomicScope::Device, + ) + .get_operation() + .insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let mut lowered = Vec::new(); + let module_region = module_ptr.deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + for op in module_block.deref(&ctx).iter(&ctx) { + let Some(function) = Operation::get_op::(op, &ctx) else { + continue; + }; + if function.get_symbol_name(&ctx).to_string() != "kernel_func" { + continue; + } + let body = function.get_operation().deref(&ctx).get_region(0); + for block in body.deref(&ctx).iter(&ctx) { + for body_op in block.deref(&ctx).iter(&ctx) { + // The IR-level forms are what libNVVM rejects; neither may survive. + assert!( + Operation::get_op::(body_op, &ctx).is_none(), + "scoped atomic load must not lower to `load atomic`: libNVVM rejects it" + ); + assert!( + Operation::get_op::(body_op, &ctx).is_none(), + "scoped atomic store must not lower to `store atomic`: libNVVM rejects it" + ); + let Some(asm) = Operation::get_op::(body_op, &ctx) else { + continue; + }; + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .unwrap_or_default(); + if !template.starts_with("ld.") && !template.starts_with("st.") { + continue; + } + lowered.push(( + template, + asm.get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .unwrap_or_default(), + llvm::asm_kind(&ctx, &asm), + )); + } + } + } + + // Every access keeps the `~{memory}` clobber, Relaxed included. Without + // it LLVM may move plain loads and stores of the same address across the + // asm, breaking the single-thread coherence Rust still guarantees for + // Relaxed atomics; libcu++ keeps the clobber on its relaxed accesses too. + let expected = [ + ("ld.relaxed.gpu.b32 $0, [$1];", "=r,l,~{memory}"), + ("ld.acquire.cta.b32 $0, [$1];", "=r,l,~{memory}"), + ("ld.relaxed.sys.b64 $0, [$1];", "=l,l,~{memory}"), + ("st.release.gpu.b32 [$0], $1;", "l,r,~{memory}"), + ("st.relaxed.gpu.b64 [$0], $1;", "l,l,~{memory}"), + ]; + assert_eq!( + lowered.len(), + expected.len(), + "expected one inline-asm op per scoped atomic load/store, got {lowered:?}" + ); + for (template, constraints) in expected { + let found = lowered + .iter() + .find(|(t, _, _)| t == template) + .unwrap_or_else(|| panic!("missing lowering for {template}; got {lowered:?}")); + assert_eq!(found.1, constraints, "constraints for {template}"); + // Not Convergent: these are per-thread accesses, not warp-synchronous. + // SideEffect with the memory clobber is what keeps a spin re-reading. + assert_eq!( + found.2, + llvm::AsmKind::SideEffect, + "asm kind for {template}" + ); + } + Ok(()) +} + +/// Orderings a PTX load or store cannot carry must be rejected, not silently +/// weakened. +/// +/// Acquire is load-only, Release is store-only, and AcqRel makes no sense on +/// a single access. Approximating any of them with `relaxed` would be a +/// correctness bug that no runtime test would catch, because the wrong answer +/// only appears under contention on hardware that reorders. +#[test] +fn test_scoped_atomic_rejects_inexpressible_orderings() -> Result<(), anyhow::Error> { + use dialect_mir::types::MirPtrType; + use dialect_nvvm::ops::atomic::{ + AtomicOrdering, AtomicScope, NvvmAtomicLoadOp, NvvmAtomicStoreOp, + }; + use pliron::builtin::types::{IntegerType, Signedness}; + + for (is_load, ordering) in [ + (true, AtomicOrdering::Release), // release is store-only + (true, AtomicOrdering::AcqRel), + (false, AtomicOrdering::Acquire), // acquire is load-only + (false, AtomicOrdering::AcqRel), + ] { + let mut ctx = make_test_ctx(); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let ptr_ty = MirPtrType::get_generic(&mut ctx, u32_ty.into(), true); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![ptr_ty.into(), u32_ty.into()]); + let address = entry.deref(&ctx).get_argument(0); + let val = entry.deref(&ctx).get_argument(1); + + if is_load { + NvvmAtomicLoadOp::build( + &mut ctx, + address, + u32_ty.into(), + ordering.clone(), + AtomicScope::Device, + ) + .get_operation() + .insert_at_back(entry, &ctx); + } else { + NvvmAtomicStoreOp::build( + &mut ctx, + val, + address, + ordering.clone(), + AtomicScope::Device, + ) + .get_operation() + .insert_at_back(entry, &ctx); + } + append_return(&mut ctx, entry); + + let result = mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr); + assert!( + result.is_err(), + "{} with {ordering:?} ordering must be rejected, not approximated", + if is_load { "load" } else { "store" } + ); + } + Ok(()) +} + +/// SeqCst load and store lower to one asm op whose template fuses the +/// `fence.sc` with the access, at every scope. +/// +/// PTX has no sequentially consistent load or store instruction. libcu++ maps +/// a SeqCst load to `fence.sc.{scope}` followed by an acquire load at the +/// same scope, and a SeqCst store to `fence.sc.{scope}` followed by a release +/// store. Fusing the fence into the same template keeps the pair inseparable. +#[test] +fn test_seqcst_atomic_load_store_fuse_fence_into_template() -> Result<(), anyhow::Error> { + use dialect_mir::types::MirPtrType; + use dialect_nvvm::ops::atomic::{ + AtomicOrdering, AtomicScope, NvvmAtomicLoadOp, NvvmAtomicStoreOp, + }; + use pliron::builtin::types::{IntegerType, Signedness}; + + for (scope, ptx) in [ + (AtomicScope::Device, "gpu"), + (AtomicScope::Block, "cta"), + (AtomicScope::System, "sys"), + ] { + let mut ctx = make_test_ctx(); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let ptr_ty = MirPtrType::get_generic(&mut ctx, u32_ty.into(), true); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![ptr_ty.into(), u32_ty.into()]); + let address = entry.deref(&ctx).get_argument(0); + let val = entry.deref(&ctx).get_argument(1); + + NvvmAtomicLoadOp::build( + &mut ctx, + address, + u32_ty.into(), + AtomicOrdering::SeqCst, + scope.clone(), + ) + .get_operation() + .insert_at_back(entry, &ctx); + NvvmAtomicStoreOp::build(&mut ctx, val, address, AtomicOrdering::SeqCst, scope) + .get_operation() + .insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let mut lowered = Vec::new(); + for op in lowered_kernel_body(&ctx, module_ptr) { + let Some(asm) = Operation::get_op::(op, &ctx) else { + continue; + }; + lowered.push(( + asm.get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .unwrap_or_default(), + asm.get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .unwrap_or_default(), + )); + } + let expected = vec![ + ( + format!("fence.sc.{ptx}; ld.acquire.{ptx}.b32 $0, [$1];"), + "=r,l,~{memory}".to_string(), + ), + ( + format!("fence.sc.{ptx}; st.release.{ptx}.b32 [$0], $1;"), + "l,r,~{memory}".to_string(), + ), + ]; + assert_eq!(lowered, expected, "SeqCst lowering at scope {ptx}"); + } + Ok(()) +} + +/// Float atomic loads and stores travel through integer registers with an +/// LLVM-level bitcast on each side. +/// +/// The register classes in the constraints are integer (`h`/`r`/`l`), so +/// handing llc a float-typed asm operand or result is a constraint mismatch. +/// The store bitcasts the float to the same-width integer before the asm; the +/// load returns the integer and bitcasts it back to the float type. The f16 +/// case also covers the `b16`/`h` arm end to end, which is what makes +/// `DeviceAtomicF16::load`/`store` compile at all. +#[test] +fn test_float_atomic_load_store_bitcast_through_integer_registers() -> Result<(), anyhow::Error> { + use dialect_mir::types::{MirFP16Type, MirPtrType}; + use dialect_nvvm::ops::atomic::{ + AtomicOrdering, AtomicScope, NvvmAtomicLoadOp, NvvmAtomicStoreOp, + }; + use llvm_export::types as llvm_types; + use pliron::builtin::types::{FP32Type, IntegerType}; + use pliron::r#type::Typed; + + let is_expected_float = |ctx: &Context, ty: pliron::r#type::TypeHandle, width: u32| -> bool { + let ty_ref = ty.deref(ctx); + match width { + 16 => ty_ref.is::(), + 32 => ty_ref.is::(), + _ => false, + } + }; + let integer_width_of = |ctx: &Context, ty: pliron::r#type::TypeHandle| -> Option { + ty.deref(ctx) + .downcast_ref::() + .map(IntegerType::width) + }; + + for (is_f16, ptx_ty, reg, width) in [(false, "b32", "r", 32u32), (true, "b16", "h", 16u32)] { + let mut ctx = make_test_ctx(); + let elem_ty: pliron::r#type::TypeHandle = if is_f16 { + MirFP16Type::get(&ctx).into() + } else { + FP32Type::get(&ctx).into() + }; + let ptr_ty = MirPtrType::get_generic(&mut ctx, elem_ty, true); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![ptr_ty.into(), elem_ty]); + let address = entry.deref(&ctx).get_argument(0); + let val = entry.deref(&ctx).get_argument(1); + + NvvmAtomicLoadOp::build( + &mut ctx, + address, + elem_ty, + AtomicOrdering::Relaxed, + AtomicScope::Device, + ) + .get_operation() + .insert_at_back(entry, &ctx); + NvvmAtomicStoreOp::build( + &mut ctx, + val, + address, + AtomicOrdering::Relaxed, + AtomicScope::Device, + ) + .get_operation() + .insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let mut ld_asm = None; + let mut st_asm = None; + let mut bitcasts = Vec::new(); + for op in lowered_kernel_body(&ctx, module_ptr) { + if let Some(asm) = Operation::get_op::(op, &ctx) { + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .unwrap_or_default(); + let constraints = asm + .get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .unwrap_or_default(); + if template.starts_with("ld.") { + ld_asm = Some((op, template, constraints)); + } else if template.starts_with("st.") { + st_asm = Some((op, template, constraints)); + } + } + if Operation::get_op::(op, &ctx).is_some() { + bitcasts.push(op); + } + } + + // Load direction: asm produces the staging integer, then one bitcast + // turns it back into the float. + let (ld_op, ld_template, ld_constraints) = + ld_asm.expect("float atomic load must lower to inline asm"); + assert_eq!(ld_template, format!("ld.relaxed.gpu.{ptx_ty} $0, [$1];")); + assert_eq!(ld_constraints, format!("={reg},l,~{{memory}}")); + let ld_result = ld_op.deref(&ctx).get_result(0); + assert_eq!( + integer_width_of(&ctx, ld_result.get_type(&ctx)), + Some(width), + "load asm must produce the staging integer, not the float" + ); + let load_cast = bitcasts + .iter() + .copied() + .find(|&cast| cast.deref(&ctx).get_operand(0) == ld_result) + .expect("load asm result must be bitcast back to the float type"); + let load_cast_ty = load_cast.deref(&ctx).get_result(0).get_type(&ctx); + assert!( + is_expected_float(&ctx, load_cast_ty, width), + "load bitcast must produce the float type" + ); + + // Store direction: the float is bitcast to the staging integer, and + // that integer is what the asm consumes. + let (st_op, st_template, st_constraints) = + st_asm.expect("float atomic store must lower to inline asm"); + assert_eq!(st_template, format!("st.relaxed.gpu.{ptx_ty} [$0], $1;")); + assert_eq!(st_constraints, format!("l,{reg},~{{memory}}")); + let stored = st_op.deref(&ctx).get_operand(1); + assert_eq!( + integer_width_of(&ctx, stored.get_type(&ctx)), + Some(width), + "store asm must consume the staging integer, not the float" + ); + let store_cast = bitcasts + .iter() + .copied() + .find(|&cast| cast.deref(&ctx).get_result(0) == stored) + .expect("store asm value must come from a float-to-integer bitcast"); + let store_src_ty = store_cast.deref(&ctx).get_operand(0).get_type(&ctx); + assert!( + is_expected_float(&ctx, store_src_ty, width), + "store bitcast must consume the float type" + ); + } + Ok(()) +} + +/// A SeqCst fence must go through the typed NVVM intrinsic, and a weaker one +/// through inline PTX. +/// +/// PTX defines `membar.level` as a synonym for `fence.sc.level`, so +/// `llvm.nvvm.membar.{cta,gl,sys}` is an exact match for SeqCst and is the route +/// the rest of the crate already uses for fences. No intrinsic exists for a +/// weaker fence, and emitting `membar` for AcqRel would silently upgrade the +/// caller's request to sequential consistency, so that case emits +/// `fence.acq_rel.{scope}` directly. +#[test] +fn test_fence_uses_membar_intrinsic_only_for_seqcst() -> Result<(), anyhow::Error> { + use dialect_mir::types::MirPtrType; + use dialect_nvvm::ops::atomic::{AtomicOrdering, AtomicRmwKind, AtomicScope, NvvmAtomicRmwOp}; + use pliron::builtin::types::{IntegerType, Signedness}; + + for (ordering, want_intrinsic, want_asm) in [ + (AtomicOrdering::SeqCst, Some("llvm_nvvm_membar_gl"), None), + (AtomicOrdering::AcqRel, None, Some("fence.acq_rel.gpu;")), + ] { + let mut ctx = make_test_ctx(); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let ptr_ty = MirPtrType::get_generic(&mut ctx, u32_ty.into(), true); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![ptr_ty.into(), u32_ty.into()]); + let address = entry.deref(&ctx).get_argument(0); + let addend = entry.deref(&ctx).get_argument(1); + + NvvmAtomicRmwOp::build( + &mut ctx, + address, + addend, + u32_ty.into(), + AtomicRmwKind::Add, + ordering.clone(), + AtomicScope::Device, + ) + .get_operation() + .insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let mut saw_intrinsic = false; + let mut saw_asm = None; + for op in lowered_kernel_body(&ctx, module_ptr) { + if let Some(call) = Operation::get_op::(op, &ctx) + && let CallOpCallable::Direct(callee) = call.callee(&ctx) + && callee.to_string().contains("membar") + { + saw_intrinsic = true; + } + if let Some(asm) = Operation::get_op::(op, &ctx) { + let t = asm + .get_attr_inline_asm_template(&ctx) + .map(|v| String::from((*v).clone())) + .unwrap_or_default(); + if t.starts_with("fence.") { + saw_asm = Some(t); + } + } + } + assert_eq!( + saw_intrinsic, + want_intrinsic.is_some(), + "{ordering:?}: membar intrinsic presence" + ); + assert_eq!(saw_asm.as_deref(), want_asm, "{ordering:?}: inline fence"); + // A SeqCst fence must never also emit `membar` as assembly: that would + // mean both routes fired. + if want_intrinsic.is_some() { + assert!(saw_asm.is_none(), "{ordering:?}: emitted both routes"); + } + } + Ok(()) +} + +/// A first-class atomic fence must lower through the same libNVVM-safe routes +/// used by the RMW ordering workaround. +/// +/// `core::sync::atomic::fence` imports with system scope, so the source-level +/// path depends on this exact mapping: Acquire/Release/AcqRel become +/// `fence.acq_rel.sys`, while SeqCst becomes `llvm.nvvm.membar.sys`. +#[test] +fn test_first_class_atomic_fence_lowers_at_system_scope() -> Result<(), anyhow::Error> { + use dialect_nvvm::ops::atomic::{AtomicOrdering, AtomicScope, NvvmAtomicFenceOp}; + + for ordering in [ + AtomicOrdering::Acquire, + AtomicOrdering::Release, + AtomicOrdering::AcqRel, + AtomicOrdering::SeqCst, + ] { + let mut ctx = make_test_ctx(); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![]); + + NvvmAtomicFenceOp::build(&mut ctx, ordering.clone(), AtomicScope::System) + .get_operation() + .insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let mut membar_calls = 0usize; + let mut fence_asm = Vec::new(); + for op in lowered_kernel_body(&ctx, module_ptr) { + assert!( + Operation::get_op::(op, &ctx).is_none(), + "first-class atomic fence must be fully consumed by lowering" + ); + + if let Some(call) = Operation::get_op::(op, &ctx) + && let CallOpCallable::Direct(callee) = call.callee(&ctx) + && callee.to_string() == "llvm_nvvm_membar_sys" + { + membar_calls += 1; + } + + let Some(asm) = Operation::get_op::(op, &ctx) else { + continue; + }; + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .unwrap_or_default(); + if template.starts_with("fence.") { + assert_eq!( + asm.get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some("~{memory}") + ); + assert_eq!(llvm::asm_kind(&ctx, &asm), llvm::AsmKind::SideEffect); + fence_asm.push(template); + } + } + + match ordering { + AtomicOrdering::SeqCst => { + assert_eq!(membar_calls, 1, "SeqCst must use membar.sys"); + assert!( + fence_asm.is_empty(), + "SeqCst must not also emit an inline PTX fence" + ); + } + AtomicOrdering::Acquire | AtomicOrdering::Release | AtomicOrdering::AcqRel => { + assert_eq!(membar_calls, 0, "weak fences must not use membar.sys"); + assert_eq!( + fence_asm, + vec!["fence.acq_rel.sys;".to_owned()], + "{ordering:?} must lower to one system-scope acq_rel fence" + ); + } + AtomicOrdering::Relaxed => unreachable!("Relaxed is not a valid Rust fence ordering"), + } + } + + Ok(()) +} + +#[test] +fn test_first_class_atomic_fence_rejects_relaxed() -> Result<(), anyhow::Error> { + use dialect_nvvm::ops::atomic::{AtomicOrdering, AtomicScope, NvvmAtomicFenceOp}; + + let mut ctx = make_test_ctx(); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![]); + NvvmAtomicFenceOp::build(&mut ctx, AtomicOrdering::Relaxed, AtomicScope::System) + .get_operation() + .insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + let result = mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr); + assert!( + result.is_err(), + "Relaxed is not a valid atomic fence ordering and must be rejected" + ); + Ok(()) +} diff --git a/crates/mir-lower/tests/lowering_test/barriers_sync.rs b/crates/mir-lower/tests/lowering_test/barriers_sync.rs new file mode 100644 index 0000000000..e0d797e53d --- /dev/null +++ b/crates/mir-lower/tests/lowering_test/barriers_sync.rs @@ -0,0 +1,743 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use dialect_nvvm::ops as nvvm; +use llvm_export::ops as llvm; +use pliron::builtin::op_interfaces::{CallOpCallable, CallOpInterface, SymbolOpInterface}; +use pliron::builtin::ops::ModuleOp; +use pliron::context::Context; +use pliron::linked_list::ContainsLinkedList; +use pliron::op::Op; +use pliron::operation::Operation; + +use crate::common::{append_return, build_test_kernel, lowered_kernel_body, make_test_ctx}; + +fn lower_basic_mbarrier( + backend: mir_lower::IntrinsicBackend, +) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { + use dialect_mir::types::MirPtrType; + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let bar_ptr_ty = MirPtrType::get_shared(&mut ctx, u64_ty.into(), false); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![bar_ptr_ty.into(), u32_ty.into()]); + let barrier = entry.deref(&ctx).get_argument(0); + let count = entry.deref(&ctx).get_argument(1); + + nvvm::MbarrierInitSharedOp::build(&mut ctx, barrier, count).insert_at_back(entry, &ctx); + let arrive = nvvm::MbarrierArriveSharedOp::build(&mut ctx, barrier); + let token = arrive.deref(&ctx).get_result(0); + arrive.insert_at_back(entry, &ctx); + nvvm::MbarrierTestWaitSharedOp::build(&mut ctx, barrier, token).insert_at_back(entry, &ctx); + nvvm::MbarrierInvalSharedOp::build(&mut ctx, barrier).insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: backend, + ..Default::default() + }, + ) + .map_err(|error| anyhow::anyhow!("{error}"))?; + Ok((ctx, module_ptr)) +} + +#[test] +fn test_generated_basic_mbarrier_uses_shared_lowering_on_both_backends() -> Result<(), anyhow::Error> +{ + let expected_calls = [ + ("llvm_nvvm_mbarrier_init_shared", 2), + ("llvm_nvvm_mbarrier_arrive_shared", 1), + ("llvm_nvvm_mbarrier_inval_shared", 1), + ]; + let init_template = "mbarrier.init.shared.b64 [$0], $1;"; + let arrive_template = "mbarrier.arrive.shared.b64 $0, [$1];"; + let test_wait_template = + "{ .reg .pred %p0; mbarrier.test_wait.shared.b64 %p0, [$1], $2; selp.b32 $0, 1, 0, %p0; }"; + let inval_template = "mbarrier.inval.shared.b64 [$0];"; + + for backend in [ + mir_lower::IntrinsicBackend::LlvmNvptx, + mir_lower::IntrinsicBackend::LibNvvm, + ] { + let (ctx, module_ptr) = lower_basic_mbarrier(backend)?; + let mut call_counts = [0usize; 3]; + let expected_asm = match backend { + mir_lower::IntrinsicBackend::LlvmNvptx => { + vec![(test_wait_template, "=r,l,l,~{memory}", 2)] + } + mir_lower::IntrinsicBackend::LibNvvm => vec![ + (init_template, "l,r,~{memory}", 2), + (arrive_template, "=l,l,~{memory}", 1), + (test_wait_template, "=r,l,l,~{memory}", 2), + (inval_template, "l,~{memory}", 1), + ], + }; + let mut asm_counts = vec![0usize; expected_asm.len()]; + let mut trunc_count = 0; + + for op in lowered_kernel_body(&ctx, module_ptr) { + if let Some(call) = Operation::get_op::(op, &ctx) { + let CallOpCallable::Direct(callee) = call.callee(&ctx) else { + continue; + }; + let callee = callee.to_string(); + let Some(index) = expected_calls + .iter() + .position(|(expected, _)| callee == *expected) + else { + assert_ne!(callee, "llvm_nvvm_mbarrier_test_wait_shared"); + continue; + }; + call_counts[index] += 1; + assert_eq!(op.deref(&ctx).get_num_operands(), expected_calls[index].1); + } + + if let Some(inline_asm) = Operation::get_op::(op, &ctx) { + let template = inline_asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())); + let index = expected_asm + .iter() + .position(|(expected, _, _)| template.as_deref() == Some(*expected)) + .unwrap_or_else(|| panic!("unexpected {backend:?} inline PTX: {template:?}")); + let (_, constraints, operand_count) = expected_asm[index]; + asm_counts[index] += 1; + assert_eq!( + inline_asm + .get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some(constraints) + ); + assert_eq!(llvm::asm_kind(&ctx, &inline_asm), llvm::AsmKind::Convergent); + let asm = op.deref(&ctx); + assert_eq!(asm.get_num_operands(), operand_count); + assert_eq!(asm.get_num_results(), 1); + } + + if Operation::get_op::(op, &ctx).is_some() { + trunc_count += 1; + } + } + + match backend { + mir_lower::IntrinsicBackend::LlvmNvptx => assert_eq!(call_counts, [1; 3]), + mir_lower::IntrinsicBackend::LibNvvm => assert_eq!(call_counts, [0; 3]), + } + assert_eq!(asm_counts, vec![1; expected_asm.len()]); + assert_eq!( + trunc_count, 1, + "test-wait must adapt its i32 predicate to i1" + ); + + let module = Operation::get_op::(module_ptr, &ctx).unwrap(); + let ir = llvm_export::export::export_module_to_string(&ctx, &module) + .map_err(|error| anyhow::anyhow!(error))?; + match backend { + mir_lower::IntrinsicBackend::LlvmNvptx => { + assert!( + ir.contains("call void @llvm.nvvm.mbarrier.init.shared(ptr addrspace(3)"), + "{ir}" + ); + assert!( + ir.contains("call i64 @llvm.nvvm.mbarrier.arrive.shared(ptr addrspace(3)"), + "{ir}" + ); + assert!( + ir.contains("call void @llvm.nvvm.mbarrier.inval.shared(ptr addrspace(3)"), + "{ir}" + ); + assert!(!ir.contains(init_template), "{ir}"); + assert!(!ir.contains(arrive_template), "{ir}"); + assert!(!ir.contains(inval_template), "{ir}"); + } + mir_lower::IntrinsicBackend::LibNvvm => { + for symbol in [ + "@llvm.nvvm.mbarrier.init.shared", + "@llvm.nvvm.mbarrier.arrive.shared", + "@llvm.nvvm.mbarrier.inval.shared", + ] { + assert!( + !ir.contains(symbol), + "libNVVM route retained {symbol}:\n{ir}" + ); + } + for template in [init_template, arrive_template, inval_template] { + assert!(ir.contains(template), "{ir}"); + } + } + } + assert!(ir.contains(test_wait_template), "{ir}"); + assert!(ir.contains("asm sideeffect"), "{ir}"); + assert!(ir.contains("trunc i32") && ir.contains("to i1"), "{ir}"); + assert!(ir.contains("attributes #0 = { convergent }"), "{ir}"); + } + Ok(()) +} + +fn lower_cluster_barriers( + backend: mir_lower::IntrinsicBackend, +) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { + let mut ctx = make_test_ctx(); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![]); + for mode in [ + nvvm::ClusterBarrierModeAttr::Arrive, + nvvm::ClusterBarrierModeAttr::ArriveAligned, + nvvm::ClusterBarrierModeAttr::ArriveRelaxed, + nvvm::ClusterBarrierModeAttr::ArriveRelaxedAligned, + nvvm::ClusterBarrierModeAttr::Wait, + nvvm::ClusterBarrierModeAttr::WaitAligned, + ] { + nvvm::ClusterBarrierOp::build(&mut ctx, mode).insert_at_back(entry, &ctx); + } + append_return(&mut ctx, entry); + mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: backend, + ..Default::default() + }, + ) + .map_err(|error| anyhow::anyhow!("{error}"))?; + Ok((ctx, module_ptr)) +} + +#[test] +fn generated_cluster_barriers_lower_exactly_on_both_backends() -> Result<(), anyhow::Error> { + let recipes = [ + ( + "llvm_nvvm_barrier_cluster_arrive", + "barrier.cluster.arrive;", + ), + ( + "llvm_nvvm_barrier_cluster_arrive_aligned", + "barrier.cluster.arrive.aligned;", + ), + ( + "llvm_nvvm_barrier_cluster_arrive_relaxed", + "barrier.cluster.arrive.relaxed;", + ), + ( + "llvm_nvvm_barrier_cluster_arrive_relaxed_aligned", + "barrier.cluster.arrive.relaxed.aligned;", + ), + ("llvm_nvvm_barrier_cluster_wait", "barrier.cluster.wait;"), + ( + "llvm_nvvm_barrier_cluster_wait_aligned", + "barrier.cluster.wait.aligned;", + ), + ]; + + for backend in [ + mir_lower::IntrinsicBackend::LlvmNvptx, + mir_lower::IntrinsicBackend::LibNvvm, + ] { + let (ctx, module_ptr) = lower_cluster_barriers(backend)?; + let mut calls = [0usize; 6]; + let mut asm = [0usize; 6]; + for op in lowered_kernel_body(&ctx, module_ptr) { + if let Some(call) = Operation::get_op::(op, &ctx) { + let CallOpCallable::Direct(callee) = call.callee(&ctx) else { + continue; + }; + if let Some(index) = recipes + .iter() + .position(|(symbol, _)| callee.to_string() == *symbol) + { + calls[index] += 1; + } + } + if let Some(inline_asm) = Operation::get_op::(op, &ctx) { + let template = inline_asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())); + if let Some(index) = recipes + .iter() + .position(|(_, expected)| template.as_deref() == Some(*expected)) + { + asm[index] += 1; + assert_eq!( + inline_asm + .get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some("~{memory}") + ); + assert_eq!(llvm::asm_kind(&ctx, &inline_asm), llvm::AsmKind::Convergent); + } + } + } + match backend { + mir_lower::IntrinsicBackend::LlvmNvptx => { + assert_eq!(calls, [1; 6]); + assert_eq!(asm, [0; 6]); + } + mir_lower::IntrinsicBackend::LibNvvm => { + assert_eq!(calls, [0; 6]); + assert_eq!(asm, [1; 6]); + } + } + } + Ok(()) +} + +#[test] +fn test_cluster_mbarrier_and_fences_lower_to_exact_inline_ptx() -> Result<(), anyhow::Error> { + use dialect_mir::types::MirPtrType; + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let bar_ptr_ty = MirPtrType::get_shared(&mut ctx, i64_ty.into(), false); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![bar_ptr_ty.into(), i32_ty.into()]); + + let bar_ptr = entry.deref(&ctx).get_argument(0); + let bytes_or_parity = entry.deref(&ctx).get_argument(1); + + let arrive = Operation::new( + &mut ctx, + nvvm::MbarrierArriveExpectTxClusterOp::get_concrete_op_info(), + vec![i64_ty.into()], + vec![bar_ptr, bytes_or_parity], + vec![], + 0, + ); + arrive.insert_at_back(entry, &ctx); + + let try_wait = Operation::new( + &mut ctx, + nvvm::MbarrierTryWaitParityClusterOp::get_concrete_op_info(), + vec![i1_ty.into()], + vec![bar_ptr, bytes_or_parity], + vec![], + 0, + ); + try_wait.insert_at_back(entry, &ctx); + + let mbarrier_fence = Operation::new( + &mut ctx, + nvvm::FenceMbarrierInitReleaseClusterOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + mbarrier_fence.insert_at_back(entry, &ctx); + + let proxy_release_fence = Operation::new( + &mut ctx, + nvvm::FenceProxyAsyncGenericReleaseSharedCtaClusterOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + proxy_release_fence.insert_at_back(entry, &ctx); + + let proxy_acquire_fence = Operation::new( + &mut ctx, + nvvm::FenceProxyAsyncGenericAcquireSharedClusterClusterOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + proxy_acquire_fence.insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let expected = [ + ( + "mbarrier.arrive.expect_tx.relaxed.cluster.shared::cta.b64 $0, [$1], $2;", + "=l,l,r,~{memory}", + ), + ( + "{ .reg .pred %p0; mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 %p0, [$1], $2; selp.b32 $0, 1, 0, %p0; }", + "=r,l,r,~{memory}", + ), + ("fence.mbarrier_init.release.cluster;", "~{memory}"), + ( + "fence.proxy.async::generic.release.sync_restrict::shared::cta.cluster;", + "~{memory}", + ), + ( + "fence.proxy.async::generic.acquire.sync_restrict::shared::cluster.cluster;", + "~{memory}", + ), + ]; + let mut matches = [0usize; 5]; + + let module_region = module_ptr.deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + for op in module_block.deref(&ctx).iter(&ctx) { + let Some(func_op) = Operation::get_op::(op, &ctx) else { + continue; + }; + if func_op.get_symbol_name(&ctx).to_string() != "kernel_func" { + continue; + } + + let func_region = func_op.get_operation().deref(&ctx).get_region(0); + for func_block in func_region.deref(&ctx).iter(&ctx) { + for body_op in func_block.deref(&ctx).iter(&ctx) { + let Some(inline_asm) = Operation::get_op::(body_op, &ctx) else { + continue; + }; + let template = inline_asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())); + let Some(index) = expected.iter().position(|(expected_template, _)| { + template.as_deref() == Some(*expected_template) + }) else { + continue; + }; + + matches[index] += 1; + assert_eq!( + inline_asm + .get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some(expected[index].1) + ); + assert_eq!( + llvm::asm_kind_opt(&ctx, &inline_asm), + Some(llvm::AsmKind::Convergent) + ); + } + } + } + + assert_eq!( + matches, [1; 5], + "each cluster barrier/fence must lower to its exact PTX template once" + ); + Ok(()) +} + +/// The mir-importer encodes `core::sync::atomic::compiler_fence` (issue #781) +/// as an empty, volatile, non-convergent inline-PTX block whose only content +/// is a `~{memory}` clobber. It must lower to a void side-effecting inline asm +/// call with the clobber intact and no instruction text, so no hardware +/// `fence` or `membar` can reach the emitted PTX. +#[test] +fn test_compiler_fence_encoding_lowers_to_empty_sideeffect_asm() -> Result<(), anyhow::Error> { + use llvm_export::types as llvm_types; + use pliron::r#type::Typed; + + let mut ctx = make_test_ctx(); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![]); + + let barrier = nvvm::InlinePtxOp::build(&mut ctx, vec![], vec![], "", "~{memory}", true, false); + barrier.insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{}", e))?; + + let mut found_barrier = false; + for op in lowered_kernel_body(&ctx, module_ptr) { + let Some(inline_asm) = Operation::get_op::(op, &ctx) else { + continue; + }; + assert!( + !found_barrier, + "expected exactly one inline asm op in the lowered kernel" + ); + found_barrier = true; + + let template = inline_asm + .get_attr_inline_asm_template(&ctx) + .map(|s| String::from((*s).clone())); + assert_eq!( + template.as_deref(), + Some(""), + "compiler fence must emit no PTX instruction" + ); + assert_eq!( + inline_asm + .get_attr_inline_asm_constraints(&ctx) + .map(|s| String::from((*s).clone())) + .as_deref(), + Some("~{memory}"), + "compiler fence must keep its memory clobber" + ); + assert!( + llvm::inline_asm_sideeffect(&ctx, inline_asm.get_operation()), + "compiler fence must stay side-effecting so the optimizer cannot drop it" + ); + assert!( + inline_asm + .get_attr_inline_asm_convergent(&ctx) + .is_some_and(|b| !bool::from((*b).clone())), + "compiler fence must not be convergent" + ); + // The zero-result lowering path models the void call as a single + // result of `llvm.void` type, so assert on the type rather than on + // the result count. + let result_ty = inline_asm + .get_operation() + .deref(&ctx) + .get_result(0) + .get_type(&ctx); + assert!( + result_ty + .deref(&ctx) + .downcast_ref::() + .is_some(), + "compiler fence lowers to a void inline asm call" + ); + } + + assert!( + found_barrier, + "expected the compiler-fence inline asm op in the lowered kernel" + ); + Ok(()) +} + +#[test] +fn test_cluster_grid_compatibility_ops_keep_original_lowering() -> Result<(), anyhow::Error> { + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let i32_type = IntegerType::get(&ctx, 32, Signedness::Signless); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![]); + for op_info in [ + nvvm::ReadPtxSregClusterIdxOp::get_concrete_op_info(), + nvvm::ReadPtxSregNclusterIdOp::get_concrete_op_info(), + ] { + Operation::new(&mut ctx, op_info, vec![i32_type.into()], vec![], vec![], 0) + .insert_at_back(entry, &ctx); + } + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let lowered = lowered_kernel_body(&ctx, module_ptr) + .into_iter() + .filter_map(|op| Operation::get_op::(op, &ctx)) + .filter_map(|asm| { + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone()))?; + (template.contains("%clusterid") || template.contains("%nclusterid")) + .then_some((template, asm)) + }) + .collect::>(); + assert_eq!(lowered.len(), 2); + assert!(lowered.iter().any(|(template, _)| { + template + == "{ .reg .u32 %cx, %cy, %cz, %nx, %ny, %nxy, %xy; mov.u32 %cx, %clusterid.x; mov.u32 %cy, %clusterid.y; mov.u32 %cz, %clusterid.z; mov.u32 %nx, %nclusterid.x; mov.u32 %ny, %nclusterid.y; mul.lo.u32 %nxy, %nx, %ny; mad.lo.u32 %xy, %cy, %nx, %cx; mad.lo.u32 $0, %cz, %nxy, %xy; }" + })); + assert!(lowered.iter().any(|(template, _)| { + template + == "{ .reg .u32 %nx, %ny, %nz, %nxy; mov.u32 %nx, %nclusterid.x; mov.u32 %ny, %nclusterid.y; mov.u32 %nz, %nclusterid.z; mul.lo.u32 %nxy, %nx, %ny; mul.lo.u32 $0, %nxy, %nz; }" + })); + for (_, asm) in lowered { + assert_eq!( + asm.get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some("=r") + ); + assert_eq!(llvm::asm_kind(&ctx, &asm), llvm::AsmKind::Convergent); + } + Ok(()) +} + +fn lower_sync_threads( + backend: mir_lower::IntrinsicBackend, +) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { + let mut ctx = make_test_ctx(); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![]); + Operation::new( + &mut ctx, + nvvm::Barrier0Op::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ) + .insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: backend, + ..Default::default() + }, + ) + .map_err(|error| anyhow::anyhow!("{error}"))?; + Ok((ctx, module_ptr)) +} + +#[test] +fn test_sync_threads_llvm_nvptx_uses_typed_intrinsic_with_fixed_zero() -> Result<(), anyhow::Error> +{ + use pliron::builtin::attributes::IntegerAttr; + + let (ctx, module_ptr) = lower_sync_threads(mir_lower::IntrinsicBackend::LlvmNvptx)?; + let body = lowered_kernel_body(&ctx, module_ptr); + let mut found = false; + for op in body { + assert!( + Operation::get_op::(op, &ctx).is_none(), + "LLVM-NVPTX sync_threads must use the typed intrinsic" + ); + let Some(call) = Operation::get_op::(op, &ctx) else { + continue; + }; + let CallOpCallable::Direct(callee) = call.callee(&ctx) else { + continue; + }; + if callee.to_string() != "llvm_nvvm_barrier_cta_sync_aligned_all" { + continue; + } + let call = op.deref(&ctx); + assert_eq!(call.get_num_operands(), 1); + let barrier_id = call.get_operand(0); + let defining_op = barrier_id.defining_op().expect("barrier ID is constant"); + let constant = Operation::get_op::(defining_op, &ctx) + .expect("barrier ID is an LLVM constant"); + let value = constant.get_value(&ctx); + let integer = value + .downcast_ref::() + .expect("barrier ID is an integer"); + assert_eq!(integer.value().bw(), 32); + assert_eq!(integer.value().to_u64(), 0); + found = true; + } + assert!(found, "modern typed CTA barrier call was not emitted"); + + let module = Operation::get_op::(module_ptr, &ctx).unwrap(); + let ir = llvm_export::export::export_module_to_string(&ctx, &module) + .map_err(|error| anyhow::anyhow!(error))?; + assert!(ir.contains("@llvm.nvvm.barrier.cta.sync.aligned.all(i32 0)")); + assert!(!ir.contains("@llvm.nvvm.barrier0")); + Ok(()) +} + +#[test] +fn test_sync_threads_libnvvm_uses_exact_convergent_inline_ptx() -> Result<(), anyhow::Error> { + let (ctx, module_ptr) = lower_sync_threads(mir_lower::IntrinsicBackend::LibNvvm)?; + let body = lowered_kernel_body(&ctx, module_ptr); + let mut found = false; + for op in body { + let Some(inline_asm) = Operation::get_op::(op, &ctx) else { + if let Some(call) = Operation::get_op::(op, &ctx) + && let CallOpCallable::Direct(callee) = call.callee(&ctx) + { + assert_ne!(callee.to_string(), "llvm_nvvm_barrier_cta_sync_aligned_all"); + } + continue; + }; + assert_eq!( + inline_asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some("bar.sync 0;") + ); + assert_eq!( + inline_asm + .get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some("~{memory}") + ); + assert_eq!(llvm::asm_kind(&ctx, &inline_asm), llvm::AsmKind::Convergent); + assert_eq!(op.deref(&ctx).get_num_operands(), 0); + found = true; + } + assert!(found, "exact libNVVM barrier inline PTX was not emitted"); + + let module = Operation::get_op::(module_ptr, &ctx).unwrap(); + let ir = llvm_export::export::export_module_to_string(&ctx, &module) + .map_err(|error| anyhow::anyhow!(error))?; + assert!( + ir.contains("call void asm sideeffect \"bar.sync 0;\", \"~{memory}\"() #0"), + "{ir}" + ); + assert!(ir.contains("attributes #0 = { convergent }"), "{ir}"); + assert!(!ir.contains("@llvm.nvvm.barrier.cta.sync.aligned.all")); + Ok(()) +} + +fn lower_warp_barrier( + backend: mir_lower::IntrinsicBackend, +) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![i32_ty.into()]); + let member_mask = entry.deref(&ctx).get_argument(0); + nvvm::BarWarpSyncOp::build(&mut ctx, member_mask).insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: backend, + ..Default::default() + }, + ) + .map_err(|error| anyhow::anyhow!("{error}"))?; + Ok((ctx, module_ptr)) +} + +#[test] +fn test_warp_barrier_uses_typed_intrinsic_on_both_backends() -> Result<(), anyhow::Error> { + for backend in [ + mir_lower::IntrinsicBackend::LlvmNvptx, + mir_lower::IntrinsicBackend::LibNvvm, + ] { + let (ctx, module_ptr) = lower_warp_barrier(backend)?; + let body = lowered_kernel_body(&ctx, module_ptr); + let mut calls = 0; + for op in body { + assert!( + Operation::get_op::(op, &ctx).is_none(), + "warp barrier must use its typed intrinsic" + ); + let Some(call) = Operation::get_op::(op, &ctx) else { + continue; + }; + let CallOpCallable::Direct(callee) = call.callee(&ctx) else { + continue; + }; + if callee.to_string() == "llvm_nvvm_bar_warp_sync" { + assert_eq!(op.deref(&ctx).get_num_operands(), 1); + assert_eq!(op.deref(&ctx).get_num_results(), 1); + calls += 1; + } + } + assert_eq!(calls, 1, "expected one typed warp-barrier call"); + + let module = Operation::get_op::(module_ptr, &ctx).unwrap(); + let ir = llvm_export::export::export_module_to_string(&ctx, &module) + .map_err(|error| anyhow::anyhow!(error))?; + assert!(ir.contains("@llvm.nvvm.bar.warp.sync(i32"), "{ir}"); + } + Ok(()) +} + +// ============================================================================= +// cp.async lowering tests +// ============================================================================= diff --git a/crates/mir-lower/tests/lowering_test/calls_and_values.rs b/crates/mir-lower/tests/lowering_test/calls_and_values.rs new file mode 100644 index 0000000000..ba83feab64 --- /dev/null +++ b/crates/mir-lower/tests/lowering_test/calls_and_values.rs @@ -0,0 +1,897 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use dialect_mir::ops as mir; +use llvm_export::ops as llvm; +use pliron::builtin::op_interfaces::{CallOpCallable, CallOpInterface, SymbolOpInterface}; +use pliron::builtin::ops::ModuleOp; +use pliron::context::Context; +use pliron::linked_list::ContainsLinkedList; +use pliron::location::{Located, Location}; +use pliron::op::Op; +use pliron::operation::Operation; + +/// Regression cover for explicit MIR address-space coercion before a call. +/// +/// When a caller passes a pointer in one address space to a callee whose +/// declared parameter lives in a different address space (the +/// `*mut SharedArray` / `addrspace(3)` case that surfaces from +/// `block_reduce` and friends), MIR must record the representational cast +/// before the exact call-signature verifier sees the argument. Lowering then +/// emits the corresponding `llvm.addrspacecast`. +/// +/// This test builds two MIR functions in one module: +/// - `callee(p: *mut i32 in addrspace(3))` +/// - `caller(p: *mut i32 in addrspace(0)) { callee(p) }` +/// +/// and asserts the lowered `caller` body contains an `AddrSpaceCastOp`. +#[test] +fn addrspace_coercion_inserts_addrspacecast_at_call_site() -> Result<(), anyhow::Error> { + use dialect_mir::types::MirPtrType; + use llvm_export::ops::AddrSpaceCastOp; + use pliron::basic_block::BasicBlock; + use pliron::builtin::attributes::{StringAttr, TypeAttr}; + use pliron::builtin::types::{FunctionType, IntegerType, Signedness}; + + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + mir_lower::register(&mut ctx); + + let module = ModuleOp::new(&mut ctx, "test_addrspace_coercion".try_into().unwrap()); + let module_ptr = module.get_operation(); + let module_region = module_ptr.deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let shared_ptr_ty = MirPtrType::get_shared(&mut ctx, i32_ty.into(), true); + let generic_ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), true); + let expected_call_location = Location::Named { + name: "source-device-call".to_string(), + child_loc: Box::new(Location::Unknown), + }; + + // Callee: takes a *mut i32 in addrspace(3), returns (). + let callee_func_ty = FunctionType::get(&ctx, vec![shared_ptr_ty.into()], vec![]); + let callee_func_op = Operation::new( + &mut ctx, + mir::MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let callee_func = mir::MirFuncOp::new( + &mut ctx, + callee_func_op, + TypeAttr::new(callee_func_ty.into()), + ); + callee_func.set_symbol_name(&mut ctx, "callee".try_into().unwrap()); + { + let region = callee_func.get_operation().deref(&ctx).get_region(0); + let block = BasicBlock::new(&mut ctx, None, vec![shared_ptr_ty.into()]); + block.insert_at_back(region, &ctx); + + let ret_op = Operation::new( + &mut ctx, + mir::MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + ret_op.insert_at_back(block, &ctx); + } + callee_func + .get_operation() + .insert_at_back(module_block, &ctx); + + // Caller: takes a *mut i32 in addrspace(0), calls `callee` with that + // pointer. The lowerer is responsible for inserting an addrspacecast + // since the callee's declared addrspace differs. + let caller_func_ty = FunctionType::get(&ctx, vec![generic_ptr_ty.into()], vec![]); + let caller_func_op = Operation::new( + &mut ctx, + mir::MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let caller_func = mir::MirFuncOp::new( + &mut ctx, + caller_func_op, + TypeAttr::new(caller_func_ty.into()), + ); + caller_func.set_symbol_name(&mut ctx, "caller".try_into().unwrap()); + { + let region = caller_func.get_operation().deref(&ctx).get_region(0); + let block = BasicBlock::new(&mut ctx, None, vec![generic_ptr_ty.into()]); + block.insert_at_back(region, &ctx); + let arg = block.deref(&ctx).get_argument(0); + let cast_op_ptr = Operation::new( + &mut ctx, + mir::MirCastOp::get_concrete_op_info(), + vec![shared_ptr_ty.into()], + vec![arg], + vec![], + 0, + ); + mir::MirCastOp::new(cast_op_ptr) + .set_attr_cast_kind(&ctx, dialect_mir::attributes::MirCastKindAttr::PtrToPtr); + cast_op_ptr.insert_at_back(block, &ctx); + let coerced_arg = cast_op_ptr.deref(&ctx).get_result(0); + + let call_op_ptr = Operation::new( + &mut ctx, + mir::MirCallOp::get_concrete_op_info(), + vec![], + vec![coerced_arg], + vec![], + 0, + ); + let call_op = mir::MirCallOp::new(call_op_ptr); + call_op.set_attr_callee(&ctx, StringAttr::new("callee".to_string())); + call_op_ptr + .deref_mut(&ctx) + .set_loc(expected_call_location.clone()); + call_op_ptr.insert_at_back(block, &ctx); + + let ret_op = Operation::new( + &mut ctx, + mir::MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + ret_op.insert_at_back(block, &ctx); + } + caller_func + .get_operation() + .insert_at_back(module_block, &ctx); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{}", e))?; + + let mut found_addrspace_cast = false; + let mut found_call = false; + let module_op = module_ptr.deref(&ctx); + let region = module_op.get_region(0); + let block = region.deref(&ctx).iter(&ctx).next().unwrap(); + for op in block.deref(&ctx).iter(&ctx) { + let Some(func_op) = Operation::get_op::(op, &ctx) else { + continue; + }; + if func_op.get_symbol_name(&ctx).to_string() != "caller" { + continue; + } + let func_region = func_op.get_operation().deref(&ctx).get_region(0); + for func_block in func_region.deref(&ctx).iter(&ctx) { + for body_op in func_block.deref(&ctx).iter(&ctx) { + if Operation::get_op::(body_op, &ctx).is_some() { + found_addrspace_cast = true; + } + if let Some(call) = Operation::get_op::(body_op, &ctx) + && matches!( + call.callee(&ctx), + CallOpCallable::Direct(symbol) if symbol.to_string() == "callee" + ) + { + found_call = true; + assert_eq!( + call.get_operation().deref(&ctx).loc(), + expected_call_location + ); + } + } + } + } + + assert!( + found_addrspace_cast, + "caller body must contain llvm.addrspacecast for the addrspace(0) -> (3) coercion at the call site", + ); + assert!(found_call, "caller body must contain the lowered call"); + Ok(()) +} + +/// A zero-sized MIR result is erased from the NVPTX function ABI, but its +/// typed value can remain live inside MIR (for example, when one ZST-returning +/// function returns the result of another). Lowering must keep the void call +/// for side effects and replace only its value result with a typed LLVM undef. +#[test] +fn zst_union_call_result_keeps_void_call_and_replaces_live_uses() -> Result<(), anyhow::Error> { + use dialect_mir::types::{MirTupleType, MirUnionType}; + use pliron::basic_block::BasicBlock; + use pliron::builtin::attributes::{StringAttr, TypeAttr}; + use pliron::builtin::types::FunctionType; + use pliron::r#type::{TypeHandle, Typed}; + + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + mir_lower::register(&mut ctx); + + let unit_ty: TypeHandle = MirTupleType::get(&mut ctx, vec![]).into(); + let union_ty: TypeHandle = MirUnionType::get( + &mut ctx, + "AlignedZeroUnion".into(), + vec!["unit".into()], + vec![unit_ty], + 0, + 16, + ) + .into(); + + let module = ModuleOp::new(&mut ctx, "test_zst_union_call".try_into().unwrap()); + let module_ptr = module.get_operation(); + let module_region = module_ptr.deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + + let callee_ty = FunctionType::get(&ctx, vec![], vec![union_ty]); + let callee_ptr = Operation::new( + &mut ctx, + mir::MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let callee = mir::MirFuncOp::new(&mut ctx, callee_ptr, TypeAttr::new(callee_ty.into())); + callee.set_symbol_name(&mut ctx, "make_zero".try_into().unwrap()); + { + let region = callee.get_operation().deref(&ctx).get_region(0); + let block = BasicBlock::new(&mut ctx, None, vec![]); + block.insert_at_back(region, &ctx); + + let undef = mir::MirUndefOp::new(&mut ctx, union_ty); + undef.get_operation().insert_at_back(block, &ctx); + let value = undef.get_operation().deref(&ctx).get_result(0); + + let ret = Operation::new( + &mut ctx, + mir::MirReturnOp::get_concrete_op_info(), + vec![], + vec![value], + vec![], + 0, + ); + ret.insert_at_back(block, &ctx); + } + callee.get_operation().insert_at_back(module_block, &ctx); + + let caller_ty = FunctionType::get(&ctx, vec![], vec![union_ty]); + let caller_ptr = Operation::new( + &mut ctx, + mir::MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let caller = mir::MirFuncOp::new(&mut ctx, caller_ptr, TypeAttr::new(caller_ty.into())); + caller.set_symbol_name(&mut ctx, "return_called_zero".try_into().unwrap()); + { + let region = caller.get_operation().deref(&ctx).get_region(0); + let block = BasicBlock::new(&mut ctx, None, vec![]); + block.insert_at_back(region, &ctx); + + let call_ptr = Operation::new( + &mut ctx, + mir::MirCallOp::get_concrete_op_info(), + vec![union_ty], + vec![], + vec![], + 0, + ); + let call = mir::MirCallOp::new(call_ptr); + call.set_attr_callee(&ctx, StringAttr::new("make_zero".to_string())); + call_ptr.insert_at_back(block, &ctx); + let value = call_ptr.deref(&ctx).get_result(0); + + let ret = Operation::new( + &mut ctx, + mir::MirReturnOp::get_concrete_op_info(), + vec![], + vec![value], + vec![], + 0, + ); + ret.insert_at_back(block, &ctx); + } + caller.get_operation().insert_at_back(module_block, &ctx); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{e}"))?; + + let mut found_void_call = false; + let mut caller_undefs = 0; + let mut found_void_return = false; + for op in module_block.deref(&ctx).iter(&ctx) { + let Some(func) = Operation::get_op::(op, &ctx) else { + continue; + }; + if func.get_symbol_name(&ctx).to_string() != "return_called_zero" { + continue; + } + let region = func.get_operation().deref(&ctx).get_region(0); + for block in region.deref(&ctx).iter(&ctx) { + for body_op in block.deref(&ctx).iter(&ctx) { + if let Some(call) = Operation::get_op::(body_op, &ctx) + && let CallOpCallable::Direct(callee) = call.callee(&ctx) + && callee.to_string() == "make_zero" + { + let call_op = call.get_operation().deref(&ctx); + assert_eq!(call_op.get_num_results(), 1); + assert!( + call_op + .get_result(0) + .get_type(&ctx) + .deref(&ctx) + .is::(), + "the ZST-returning call must use the void ABI" + ); + found_void_call = true; + } + if Operation::get_op::(body_op, &ctx).is_some() { + caller_undefs += 1; + } + if let Some(ret) = Operation::get_op::(body_op, &ctx) { + assert_eq!( + ret.get_operation().deref(&ctx).get_num_operands(), + 0, + "the caller's ZST return must also use the void ABI" + ); + found_void_return = true; + } + } + } + } + + assert!( + found_void_call, + "the LLVM call must be retained because the callee may have side effects" + ); + assert_eq!( + caller_undefs, 1, + "the live MIR result must be replaced by one typed LLVM undef" + ); + assert!( + found_void_return, + "the caller must retain its return terminator" + ); + Ok(()) +} + +/// Lock the comparison-predicate lowering table to the rustc_codegen_ssa +/// reference (`bin_op_to_fcmp_predicate` / `bin_op_to_icmp_predicate`): +/// +/// | MIR op | float `fcmp` | signed `icmp` | unsigned `icmp` | +/// |----------|-------------------|---------------|-----------------| +/// | `mir.eq` | `oeq` (ordered) | `eq` | `eq` | +/// | `mir.ne` | `une` (UNordered) | `ne` | `ne` | +/// | `mir.lt` | `olt` | `slt` | `ult` | +/// | `mir.le` | `ole` | `sle` | `ule` | +/// | `mir.gt` | `ogt` | `sgt` | `ugt` | +/// | `mir.ge` | `oge` | `sge` | `uge` | +/// +/// `ne` is the one float predicate that must be UNordered: Rust requires +/// `a != b == !(a == b)`, so `x != x` must be true for NaN (issue #123; +/// the ordered `one` folds the canonical NaN check to `false`). +/// +/// The test also locks fastmath flags to *empty* on every lowered `fcmp`: +/// a future `nnan` default would make `fcmp nnan une x, x` poison for NaN +/// and silently re-break NaN detection while the predicate assertion above +/// stays green. +#[test] +fn test_cmp_predicate_lowering() -> Result<(), anyhow::Error> { + use llvm_export::attributes::{FCmpPredicateAttr, FastmathFlagsAttr, ICmpPredicateAttr}; + use llvm_export::op_interfaces::FastMathFlags; + + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + mir_lower::register(&mut ctx); + + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_ptr = module.get_operation(); + + let f32_ty = pliron::builtin::types::FP32Type::get(&ctx); + let i32_signed = pliron::builtin::types::IntegerType::get( + &ctx, + 32, + pliron::builtin::types::Signedness::Signed, + ); + let u32_unsigned = pliron::builtin::types::IntegerType::get( + &ctx, + 32, + pliron::builtin::types::Signedness::Unsigned, + ); + let bool_ty = pliron::builtin::types::IntegerType::get( + &ctx, + 1, + pliron::builtin::types::Signedness::Signless, + ); + + // Args: (f32, f32, i32, u32). The integer args carry pre-conversion + // signedness, which is what selects signed vs unsigned icmp predicates. + let arg_tys: Vec = vec![ + f32_ty.into(), + f32_ty.into(), + i32_signed.into(), + u32_unsigned.into(), + ]; + let func_name = "cmp_func"; + let func_ty = pliron::builtin::types::FunctionType::get(&ctx, arg_tys.clone(), vec![]); + + let func_op_ptr = Operation::new( + &mut ctx, + mir::MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let func_ty_attr = pliron::builtin::attributes::TypeAttr::new(func_ty.into()); + let func = mir::MirFuncOp::new(&mut ctx, func_op_ptr, func_ty_attr); + func.set_symbol_name(&mut ctx, func_name.try_into().unwrap()); + + let region = func.get_operation().deref(&ctx).get_region(0); + let block = { + let b = pliron::basic_block::BasicBlock::new(&mut ctx, None, arg_tys); + b.insert_at_back(region, &ctx); + b + }; + let fa = block.deref(&ctx).get_argument(0); + let fb = block.deref(&ctx).get_argument(1); + let si = block.deref(&ctx).get_argument(2); + let ui = block.deref(&ctx).get_argument(3); + + // One comparison op per table row, in a fixed program order. The raw + // `Operation::new` construction mirrors how the importer builds these + // ops (mir-importer translator/rvalue.rs BinaryOp arm). + let cmp_infos = [ + // Floats: all six predicates. + (mir::MirEqOp::get_concrete_op_info(), fa, fb), + (mir::MirNeOp::get_concrete_op_info(), fa, fb), + (mir::MirLtOp::get_concrete_op_info(), fa, fb), + (mir::MirLeOp::get_concrete_op_info(), fa, fb), + (mir::MirGtOp::get_concrete_op_info(), fa, fb), + (mir::MirGeOp::get_concrete_op_info(), fa, fb), + // Signed integers: eq/ne are sign-agnostic, the rest must be s*. + (mir::MirEqOp::get_concrete_op_info(), si, si), + (mir::MirNeOp::get_concrete_op_info(), si, si), + (mir::MirLtOp::get_concrete_op_info(), si, si), + (mir::MirLeOp::get_concrete_op_info(), si, si), + (mir::MirGtOp::get_concrete_op_info(), si, si), + (mir::MirGeOp::get_concrete_op_info(), si, si), + // Unsigned integers: the relational predicates must be u*. + (mir::MirLtOp::get_concrete_op_info(), ui, ui), + (mir::MirLeOp::get_concrete_op_info(), ui, ui), + (mir::MirGtOp::get_concrete_op_info(), ui, ui), + (mir::MirGeOp::get_concrete_op_info(), ui, ui), + ]; + for (info, lhs, rhs) in cmp_infos { + let op = Operation::new( + &mut ctx, + info, + vec![bool_ty.into()], + vec![lhs, rhs], + vec![], + 0, + ); + op.insert_at_back(block, &ctx); + } + + let ret_op_ptr = Operation::new( + &mut ctx, + mir::MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + ret_op_ptr.insert_at_back(block, &ctx); + + let module_region = module.get_operation().deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + func.get_operation().insert_at_back(module_block, &ctx); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{}", e))?; + + // Collect lowered predicates in program order. + let mut fcmp_preds = Vec::new(); + let mut icmp_preds = Vec::new(); + let module_op = module_ptr.deref(&ctx); + let region = module_op.get_region(0); + let block = region.deref(&ctx).iter(&ctx).next().unwrap(); + for op in block.deref(&ctx).iter(&ctx) { + let Some(func_op) = Operation::get_op::(op, &ctx) else { + continue; + }; + if func_op.get_symbol_name(&ctx).to_string() != func_name { + continue; + } + let func_region = func_op.get_operation().deref(&ctx).get_region(0); + for func_block in func_region.deref(&ctx).iter(&ctx) { + for body_op in func_block.deref(&ctx).iter(&ctx) { + if let Some(fcmp) = Operation::get_op::(body_op, &ctx) { + fcmp_preds.push(fcmp.predicate(&ctx)); + // fcmp carries `contract` (set by add_fastmath_flags) which is a + // no-op for comparisons at the LLVM / PTX level. Critically, nnan + // is NOT set, so NaN checks like `x != x` still evaluate correctly. + let expected: FastmathFlagsAttr = + llvm_export::attributes::FastmathFlags::CONTRACT.into(); + assert_eq!( + fcmp.fast_math_flags(&ctx), + expected, + "fcmp must carry only the contract flag (nnan would poison NaN checks)" + ); + } + if let Some(icmp) = Operation::get_op::(body_op, &ctx) { + icmp_preds.push(icmp.predicate(&ctx)); + } + } + } + } + + assert_eq!( + fcmp_preds, + vec![ + FCmpPredicateAttr::OEQ, + FCmpPredicateAttr::UNE, + FCmpPredicateAttr::OLT, + FCmpPredicateAttr::OLE, + FCmpPredicateAttr::OGT, + FCmpPredicateAttr::OGE, + ], + "float comparison predicates must mirror rustc: ordered except Ne (une)" + ); + assert_eq!( + icmp_preds, + vec![ + ICmpPredicateAttr::EQ, + ICmpPredicateAttr::NE, + ICmpPredicateAttr::SLT, + ICmpPredicateAttr::SLE, + ICmpPredicateAttr::SGT, + ICmpPredicateAttr::SGE, + ICmpPredicateAttr::ULT, + ICmpPredicateAttr::ULE, + ICmpPredicateAttr::UGT, + ICmpPredicateAttr::UGE, + ], + "integer comparison predicates must respect pre-conversion signedness" + ); + Ok(()) +} + +/// Regression cover for PR #141: comparisons whose operand is a bool phi. +/// +/// Bools are signless i1, which `can_convert_type` rejects (signless is +/// already the LLVM form), so DialectConversion records no type history for +/// a bool block argument. `is_signed_int_op` used to error out for such +/// operands ("expected IntegerType or MirPtrType operand in arithmetic op"); +/// it must instead fall back to the live operand type and lower the +/// comparison as unsigned. +/// +/// The function mirrors the MIR of a short-circuit kernel: +/// +/// ```text +/// let p = a || b; // bool phi: merge block argument +/// out = (p == q, p < q); // icmp eq i1 / icmp ult i1 +/// ``` +/// +/// ```text +/// bb0(a: i1, b: i1, q: i1): mir.cond_br a, bb2(a), bb1() +/// bb1(): mir.goto bb2(b) +/// bb2(p: i1): mir.eq p, q ; mir.lt p, q ; mir.return +/// ``` +#[test] +fn test_bool_phi_cmp_lowers_to_unsigned_i1_icmp() -> Result<(), anyhow::Error> { + use llvm_export::attributes::ICmpPredicateAttr; + use pliron::basic_block::BasicBlock; + use pliron::builtin::op_interfaces::OperandSegmentInterface; + use pliron::builtin::types::{FunctionType, IntegerType, Signedness}; + use pliron::r#type::Typed; + + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + mir_lower::register(&mut ctx); + + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_ptr = module.get_operation(); + + let bool_ty = IntegerType::get(&ctx, 1, Signedness::Signless); + let arg_tys: Vec = + vec![bool_ty.into(), bool_ty.into(), bool_ty.into()]; + let func_name = "bool_phi_cmp"; + let func_ty = FunctionType::get(&ctx, arg_tys.clone(), vec![]); + + let func_op_ptr = Operation::new( + &mut ctx, + mir::MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let func_ty_attr = pliron::builtin::attributes::TypeAttr::new(func_ty.into()); + let func = mir::MirFuncOp::new(&mut ctx, func_op_ptr, func_ty_attr); + func.set_symbol_name(&mut ctx, func_name.try_into().unwrap()); + + let region = func.get_operation().deref(&ctx).get_region(0); + + // bb0(a, b, q): the function entry. + let bb0 = BasicBlock::new(&mut ctx, None, arg_tys); + bb0.insert_at_back(region, &ctx); + let a = bb0.deref(&ctx).get_argument(0); + let b = bb0.deref(&ctx).get_argument(1); + let q = bb0.deref(&ctx).get_argument(2); + + // bb1(): the short-circuit "evaluate b" block. + let bb1 = BasicBlock::new(&mut ctx, None, vec![]); + bb1.insert_at_back(region, &ctx); + + // bb2(p): the merge block; `p` is the bool phi. + let bb2 = BasicBlock::new(&mut ctx, None, vec![bool_ty.into()]); + bb2.insert_at_back(region, &ctx); + let p = bb2.deref(&ctx).get_argument(0); + + // bb0: cond_br a, bb2(a), bb1(). On the true edge `a` is true, so + // passing `a` itself is `a || b` without needing a constant. + let (flat_operands, segment_sizes) = + mir::MirCondBranchOp::compute_segment_sizes(vec![vec![a], vec![a], vec![]]); + let cond_br = Operation::new( + &mut ctx, + mir::MirCondBranchOp::get_concrete_op_info(), + vec![], + flat_operands, + vec![bb2, bb1], + 0, + ); + Operation::get_op::(cond_br, &ctx) + .expect("MirCondBranchOp") + .set_operand_segment_sizes(&ctx, segment_sizes); + cond_br.insert_at_back(bb0, &ctx); + + // bb1: goto bb2(b). + let goto = Operation::new( + &mut ctx, + mir::MirGotoOp::get_concrete_op_info(), + vec![], + vec![b], + vec![bb2], + 0, + ); + goto.insert_at_back(bb1, &ctx); + + // bb2: p == q, then p < q. + for info in [ + mir::MirEqOp::get_concrete_op_info(), + mir::MirLtOp::get_concrete_op_info(), + ] { + let cmp = Operation::new(&mut ctx, info, vec![bool_ty.into()], vec![p, q], vec![], 0); + cmp.insert_at_back(bb2, &ctx); + } + let ret_op = Operation::new( + &mut ctx, + mir::MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + ret_op.insert_at_back(bb2, &ctx); + + let module_region = module.get_operation().deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + func.get_operation().insert_at_back(module_block, &ctx); + + // Before the fallback, this failed with "expected IntegerType or + // MirPtrType operand in arithmetic op". + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{}", e))?; + + let mut icmps = Vec::new(); + let module_op = module_ptr.deref(&ctx); + let region = module_op.get_region(0); + let block = region.deref(&ctx).iter(&ctx).next().unwrap(); + for op in block.deref(&ctx).iter(&ctx) { + let Some(func_op) = Operation::get_op::(op, &ctx) else { + continue; + }; + if func_op.get_symbol_name(&ctx).to_string() != func_name { + continue; + } + let func_region = func_op.get_operation().deref(&ctx).get_region(0); + for func_block in func_region.deref(&ctx).iter(&ctx) { + for body_op in func_block.deref(&ctx).iter(&ctx) { + if let Some(icmp) = Operation::get_op::(body_op, &ctx) { + let lhs_ty = body_op.deref(&ctx).get_operand(0).get_type(&ctx); + icmps.push((icmp.predicate(&ctx), lhs_ty)); + } + } + } + } + + let i1: pliron::r#type::TypeHandle = bool_ty.into(); + assert_eq!( + icmps, + vec![(ICmpPredicateAttr::EQ, i1), (ICmpPredicateAttr::ULT, i1),], + "bool-phi comparisons must lower to `icmp eq i1` and `icmp ult i1`" + ); + Ok(()) +} + +// ============================================================================= +// Scalar math lowering tests +// ============================================================================= + +/// Builds a module holding a single `mir.func` whose body calls the `bswap` +/// placeholder on a `u8` argument, with the call result carrying `result_ty`. +/// Shared by the valid (u8 result) and malformed (tuple result) 8-bit bswap +/// tests below. +fn build_bswap8_call_module( + ctx: &mut Context, + result_ty: pliron::r#type::TypeHandle, +) -> pliron::context::Ptr { + use pliron::basic_block::BasicBlock; + use pliron::builtin::attributes::StringAttr; + use pliron::builtin::attributes::TypeAttr; + use pliron::builtin::types::{FunctionType, IntegerType, Signedness}; + + let u8_ty = IntegerType::get(ctx, 8, Signedness::Unsigned); + + let module = ModuleOp::new(ctx, "m".try_into().unwrap()); + let module_ptr = module.get_operation(); + let module_block = module + .get_operation() + .deref(ctx) + .get_region(0) + .deref(ctx) + .iter(ctx) + .next() + .unwrap(); + + let func_ty = FunctionType::get(ctx, vec![u8_ty.into()], vec![]); + let func_op = Operation::new( + ctx, + mir::MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let func = mir::MirFuncOp::new(ctx, func_op, TypeAttr::new(func_ty.into())); + func.set_symbol_name(ctx, "f".try_into().unwrap()); + { + let region = func.get_operation().deref(ctx).get_region(0); + let block = BasicBlock::new(ctx, None, vec![u8_ty.into()]); + block.insert_at_back(region, ctx); + let arg = block.deref(ctx).get_argument(0); + + let call_ptr = Operation::new( + ctx, + mir::MirCallOp::get_concrete_op_info(), + vec![result_ty], + vec![arg], + vec![], + 0, + ); + let call = mir::MirCallOp::new(call_ptr); + call.set_attr_callee( + ctx, + StringAttr::new(dialect_mir::rust_intrinsics::CALLEE_BSWAP.to_string()), + ); + call_ptr.insert_at_back(block, ctx); + + let ret_op = Operation::new( + ctx, + mir::MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + ret_op.insert_at_back(block, ctx); + } + func.get_operation().insert_at_back(module_block, ctx); + + module_ptr +} + +/// An 8-bit `bswap` takes an early bitcast path, since Rust's semantics for a +/// single byte are identity and LLVM has no useful intrinsic for it. That path +/// returned before reaching the cast every other arm goes through, so a result +/// type that was not an 8-bit integer produced `bitcast i8 to `, +/// which LLVM rejects when the module is read back. +/// +/// The known producer of this shape is an importer bug that hands the call +/// the type of the destination *local* rather than of the projected place it +/// writes, so `RET.1 = bswap(x)` on a `(f64, u8)` return arrives here +/// carrying the whole tuple. That bug is the importer's to fix, and whatever +/// the importer does, a malformed producer can still build such a `mir.call`. +/// Lowering cannot repair it, and refusing it is what keeps the emitted +/// module valid. +#[test] +fn bswap8_with_an_aggregate_result_is_refused_rather_than_bitcast() -> Result<(), anyhow::Error> { + use dialect_mir::types::MirTupleType; + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + mir_lower::register(&mut ctx); + + let u8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); + let f64_ty = pliron::builtin::types::FP64Type::get(&ctx); + // The tuple result is the whole point: an 8-bit operand with a result + // the bitcast cannot legally produce. + let tuple_ty = MirTupleType::get(&mut ctx, vec![f64_ty.into(), u8_ty.into()]); + let module_ptr = build_bswap8_call_module(&mut ctx, tuple_ty.into()); + + let result = mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr); + let error = + result.expect_err("an 8-bit bswap with a tuple result must be refused, not lowered"); + assert!( + error + .to_string() + .contains("expected integer type for Rust bit intrinsic"), + "unexpected error: {error}" + ); + Ok(()) +} + +/// The valid counterpart: an 8-bit `bswap` whose result is a plain `u8` must +/// still take the identity-bitcast path and lower cleanly, so the guard above +/// cannot over-reject the case Rust actually produces for `u8::swap_bytes`. +#[test] +fn bswap8_with_a_u8_result_lowers_to_an_identity_bitcast() -> Result<(), anyhow::Error> { + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + mir_lower::register(&mut ctx); + + let u8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); + let module_ptr = build_bswap8_call_module(&mut ctx, u8_ty.into()); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .expect("an 8-bit bswap with a u8 result must lower"); + + // The call must have become a bitcast, not an intrinsic call: there is no + // llvm.bswap.i8. + let mut found_bitcast = false; + let module_block = module_ptr + .deref(&ctx) + .get_region(0) + .deref(&ctx) + .iter(&ctx) + .next() + .unwrap(); + for op in module_block.deref(&ctx).iter(&ctx) { + let Some(func_op) = Operation::get_op::(op, &ctx) else { + continue; + }; + let func_region = func_op.get_operation().deref(&ctx).get_region(0); + for func_block in func_region.deref(&ctx).iter(&ctx) { + for body_op in func_block.deref(&ctx).iter(&ctx) { + assert!( + Operation::get_op::(body_op, &ctx).is_none(), + "an 8-bit bswap must not lower to an intrinsic call" + ); + if Operation::get_op::(body_op, &ctx).is_some() { + found_bitcast = true; + } + } + } + } + assert!( + found_bitcast, + "an 8-bit bswap with a u8 result must lower to an identity bitcast" + ); + Ok(()) +} diff --git a/crates/mir-lower/tests/lowering_test/common.rs b/crates/mir-lower/tests/lowering_test/common.rs new file mode 100644 index 0000000000..0fa5a821d5 --- /dev/null +++ b/crates/mir-lower/tests/lowering_test/common.rs @@ -0,0 +1,101 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use dialect_mir::ops as mir; +use llvm_export::ops as llvm; +use pliron::builtin::op_interfaces::SymbolOpInterface; +use pliron::builtin::ops::ModuleOp; +use pliron::context::Context; +use pliron::linked_list::ContainsLinkedList; +use pliron::op::Op; +use pliron::operation::Operation; + +/// Helper: fresh context with all dialects registered. +pub(super) fn make_test_ctx() -> Context { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + mir_lower::register(&mut ctx); + ctx +} + +/// Helper: build a module + MirFuncOp("kernel_func") with given arg types, +/// returning the module ptr and entry block. +pub(super) fn build_test_kernel( + ctx: &mut Context, + arg_tys: Vec, +) -> ( + pliron::context::Ptr, + pliron::context::Ptr, +) { + use pliron::basic_block::BasicBlock; + use pliron::builtin::attributes::TypeAttr; + use pliron::builtin::types::FunctionType; + + let module = ModuleOp::new(ctx, "test_module".try_into().unwrap()); + let module_ptr = module.get_operation(); + + let func_ty = FunctionType::get(ctx, arg_tys.clone(), vec![]); + let func_op_ptr = Operation::new( + ctx, + mir::MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let func = mir::MirFuncOp::new(ctx, func_op_ptr, TypeAttr::new(func_ty.into())); + func.set_symbol_name(ctx, "kernel_func".try_into().unwrap()); + + let region = func.get_operation().deref(ctx).get_region(0); + let entry = BasicBlock::new(ctx, None, arg_tys); + entry.insert_at_back(region, ctx); + + let module_region = module_ptr.deref(ctx).get_region(0); + let module_block = module_region.deref(ctx).iter(ctx).next().unwrap(); + func.get_operation().insert_at_back(module_block, ctx); + + (module_ptr, entry) +} + +/// Helper: append a mir.return (void) to a block. +pub(super) fn append_return( + ctx: &mut Context, + block: pliron::context::Ptr, +) { + let ret = Operation::new( + ctx, + mir::MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + ret.insert_at_back(block, ctx); +} + +pub(super) fn lowered_kernel_body( + ctx: &Context, + module_ptr: pliron::context::Ptr, +) -> Vec> { + let module_region = module_ptr.deref(ctx).get_region(0); + let module_block = module_region.deref(ctx).iter(ctx).next().unwrap(); + for op in module_block.deref(ctx).iter(ctx) { + let Some(func_op) = Operation::get_op::(op, ctx) else { + continue; + }; + if func_op.get_symbol_name(ctx).to_string() != "kernel_func" { + continue; + } + + let func_region = func_op.get_operation().deref(ctx).get_region(0); + return func_region + .deref(ctx) + .iter(ctx) + .flat_map(|block| block.deref(ctx).iter(ctx)) + .collect(); + } + panic!("lowered kernel function not found") +} diff --git a/crates/mir-lower/tests/lowering_test/cp_async.rs b/crates/mir-lower/tests/lowering_test/cp_async.rs new file mode 100644 index 0000000000..17d2c3c61b --- /dev/null +++ b/crates/mir-lower/tests/lowering_test/cp_async.rs @@ -0,0 +1,727 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use dialect_mir::ops as mir; +use dialect_nvvm::ops as nvvm; +use llvm_export::ops as llvm; +use pliron::builtin::op_interfaces::{CallOpCallable, CallOpInterface, SymbolOpInterface}; +use pliron::builtin::ops::ModuleOp; +use pliron::context::Context; +use pliron::linked_list::ContainsLinkedList; +use pliron::op::Op; +use pliron::operation::Operation; + +use crate::common::{append_return, build_test_kernel, lowered_kernel_body, make_test_ctx}; + +fn lower_all_classic_cp_async( + backend: mir_lower::IntrinsicBackend, +) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { + use dialect_mir::types::MirPtrType; + use pliron::builtin::attributes::IntegerAttr; + use pliron::builtin::types::{IntegerType, Signedness}; + use pliron::utils::apint::APInt; + use std::num::NonZeroUsize; + + let mut ctx = make_test_ctx(); + let u8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let dst_ty = MirPtrType::get_generic(&mut ctx, u32_ty.into(), true); + let src32_ty = MirPtrType::get_generic(&mut ctx, u32_ty.into(), false); + let src8_ty = MirPtrType::get_generic(&mut ctx, u8_ty.into(), false); + let (module_ptr, entry) = build_test_kernel( + &mut ctx, + vec![ + dst_ty.into(), + src32_ty.into(), + src8_ty.into(), + u32_ty.into(), + ], + ); + let dst = entry.deref(&ctx).get_argument(0); + let src32 = entry.deref(&ctx).get_argument(1); + let src8 = entry.deref(&ctx).get_argument(2); + let source_size = entry.deref(&ctx).get_argument(3); + + let zero_op = Operation::new( + &mut ctx, + mir::MirConstantOp::get_concrete_op_info(), + vec![u32_ty.into()], + vec![], + vec![], + 0, + ); + mir::MirConstantOp::new(zero_op).set_attr_value( + &ctx, + IntegerAttr::new(u32_ty, APInt::from_u32(0, NonZeroUsize::new(32).unwrap())), + ); + zero_op.insert_at_back(entry, &ctx); + let zero = zero_op.deref(&ctx).get_result(0); + + for copy in [ + nvvm::CpAsyncCa4Op::build(&mut ctx, dst, src32), + nvvm::CpAsyncCa8Op::build(&mut ctx, dst, src32), + nvvm::CpAsyncCa16Op::build(&mut ctx, dst, src32), + nvvm::CpAsyncCaZfill4Op::build(&mut ctx, dst, src8, source_size), + nvvm::CpAsyncCaZfill8Op::build(&mut ctx, dst, src8, source_size), + nvvm::CpAsyncCaZfill16Op::build(&mut ctx, dst, src8, source_size), + nvvm::CpAsyncCg16Op::build(&mut ctx, dst, src32), + nvvm::CpAsyncCgZfill16Op::build(&mut ctx, dst, src8, source_size), + ] { + copy.insert_at_back(entry, &ctx); + } + for control in [ + nvvm::CpAsyncCommitGroupOp::build(&mut ctx), + nvvm::CpAsyncWaitGroupOp::build(&mut ctx, zero), + nvvm::CpAsyncWaitAllOp::build(&mut ctx), + ] { + control.insert_at_back(entry, &ctx); + } + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: backend, + ..Default::default() + }, + ) + .map_err(|error| anyhow::anyhow!("{error}"))?; + Ok((ctx, module_ptr)) +} + +fn lower_all_cp_async_mbarrier( + backend: mir_lower::IntrinsicBackend, +) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { + use dialect_mir::types::MirPtrType; + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let generic_ty = MirPtrType::get_generic(&mut ctx, u64_ty.into(), true); + let shared_ty = MirPtrType::get_shared(&mut ctx, u64_ty.into(), true); + let (module_ptr, entry) = + build_test_kernel(&mut ctx, vec![generic_ty.into(), shared_ty.into()]); + let generic = entry.deref(&ctx).get_argument(0); + let shared = entry.deref(&ctx).get_argument(1); + + for bridge in [ + nvvm::CpAsyncMbarrierArriveOp::build(&mut ctx, shared), + nvvm::CpAsyncMbarrierArriveSharedOp::build(&mut ctx, generic), + nvvm::CpAsyncMbarrierArriveNoIncOp::build(&mut ctx, shared), + nvvm::CpAsyncMbarrierArriveNoIncSharedOp::build(&mut ctx, generic), + ] { + bridge.insert_at_back(entry, &ctx); + } + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: backend, + ..Default::default() + }, + ) + .map_err(|error| anyhow::anyhow!("{error}"))?; + Ok((ctx, module_ptr)) +} + +#[test] +fn test_generated_cp_async_mbarrier_preserves_backend_and_address_routes() +-> Result<(), anyhow::Error> { + use llvm_export::types::PointerType; + use pliron::r#type::Typed; + + let typed = [ + ("llvm_nvvm_cp_async_mbarrier_arrive", 0), + ("llvm_nvvm_cp_async_mbarrier_arrive_shared", 3), + ("llvm_nvvm_cp_async_mbarrier_arrive_noinc", 0), + ("llvm_nvvm_cp_async_mbarrier_arrive_noinc_shared", 3), + ]; + let templates = [ + ("cp.async.mbarrier.arrive.b64 [$0];", 0), + ("cp.async.mbarrier.arrive.shared.b64 [$0];", 3), + ("cp.async.mbarrier.arrive.noinc.b64 [$0];", 0), + ("cp.async.mbarrier.arrive.noinc.shared.b64 [$0];", 3), + ]; + + for backend in [ + mir_lower::IntrinsicBackend::LlvmNvptx, + mir_lower::IntrinsicBackend::LibNvvm, + ] { + let (ctx, module_ptr) = lower_all_cp_async_mbarrier(backend)?; + let mut call_counts = [0usize; 4]; + let mut asm_counts = [0usize; 4]; + + for op in lowered_kernel_body(&ctx, module_ptr) { + if let Some(call) = Operation::get_op::(op, &ctx) { + let CallOpCallable::Direct(callee) = call.callee(&ctx) else { + continue; + }; + let callee = callee.to_string(); + let Some(index) = typed.iter().position(|(name, _)| callee == *name) else { + continue; + }; + call_counts[index] += 1; + let pointer = op.deref(&ctx).get_operand(0).get_type(&ctx); + assert_eq!( + pointer + .deref(&ctx) + .downcast_ref::() + .unwrap() + .address_space(), + typed[index].1 + ); + } + if let Some(inline_asm) = Operation::get_op::(op, &ctx) { + let template = inline_asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .unwrap(); + let Some(index) = templates + .iter() + .position(|(expected, _)| template == *expected) + else { + continue; + }; + asm_counts[index] += 1; + assert_eq!(llvm::asm_kind(&ctx, &inline_asm), llvm::AsmKind::Convergent); + assert!( + inline_asm + .get_attr_inline_asm_convergent(&ctx) + .is_some_and(|value| bool::from((*value).clone())) + ); + assert_eq!( + inline_asm + .get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some("l,~{memory}") + ); + let pointer = op.deref(&ctx).get_operand(0).get_type(&ctx); + assert_eq!( + pointer + .deref(&ctx) + .downcast_ref::() + .unwrap() + .address_space(), + templates[index].1 + ); + } + } + + let module = Operation::get_op::(module_ptr, &ctx).unwrap(); + let ir = llvm_export::export::export_module_to_string(&ctx, &module) + .map_err(|error| anyhow::anyhow!(error))?; + match backend { + mir_lower::IntrinsicBackend::LlvmNvptx => { + assert_eq!(call_counts, [1; 4]); + assert_eq!(asm_counts, [0; 4]); + assert!( + ir.contains("@llvm.nvvm.cp.async.mbarrier.arrive(ptr"), + "{ir}" + ); + assert!( + ir.contains("@llvm.nvvm.cp.async.mbarrier.arrive.shared(ptr addrspace(3)"), + "{ir}" + ); + assert!( + ir.contains("@llvm.nvvm.cp.async.mbarrier.arrive.noinc(ptr"), + "{ir}" + ); + assert!( + ir.contains( + "@llvm.nvvm.cp.async.mbarrier.arrive.noinc.shared(ptr addrspace(3)" + ), + "{ir}" + ); + assert!(!ir.contains("cp.async.mbarrier.arrive.b64 [$0]"), "{ir}"); + } + mir_lower::IntrinsicBackend::LibNvvm => { + assert_eq!(call_counts, [0; 4]); + assert_eq!(asm_counts, [1; 4]); + assert!(!ir.contains("@llvm.nvvm.cp.async.mbarrier"), "{ir}"); + for (template, _) in templates { + assert!(ir.contains(template), "{ir}"); + } + assert!(ir.contains("asm sideeffect"), "{ir}"); + assert!(ir.contains("convergent"), "{ir}"); + } + } + } + Ok(()) +} + +#[test] +fn test_generated_cp_async_llvm_nvptx_uses_all_typed_intrinsics() -> Result<(), anyhow::Error> { + use llvm_export::types::PointerType; + use pliron::r#type::Typed; + + let (ctx, module_ptr) = lower_all_classic_cp_async(mir_lower::IntrinsicBackend::LlvmNvptx)?; + let mut found = Vec::new(); + for op in lowered_kernel_body(&ctx, module_ptr) { + assert!( + Operation::get_op::(op, &ctx).is_none(), + "LLVM-NVPTX cp.async must use typed intrinsics" + ); + let Some(call) = Operation::get_op::(op, &ctx) else { + continue; + }; + let CallOpCallable::Direct(callee) = call.callee(&ctx) else { + continue; + }; + let callee = callee.to_string(); + if !callee.starts_with("llvm_nvvm_cp_async") { + continue; + } + if callee.contains("shared_global") { + let call = op.deref(&ctx); + let destination_ty = call.get_operand(0).get_type(&ctx); + let source_ty = call.get_operand(1).get_type(&ctx); + assert_eq!( + destination_ty + .deref(&ctx) + .downcast_ref::() + .unwrap() + .address_space(), + 3 + ); + assert_eq!( + source_ty + .deref(&ctx) + .downcast_ref::() + .unwrap() + .address_space(), + 1 + ); + } + found.push(callee); + } + found.sort(); + let mut expected = [ + "llvm_nvvm_cp_async_ca_shared_global_4", + "llvm_nvvm_cp_async_ca_shared_global_4_s", + "llvm_nvvm_cp_async_ca_shared_global_8", + "llvm_nvvm_cp_async_ca_shared_global_8_s", + "llvm_nvvm_cp_async_ca_shared_global_16", + "llvm_nvvm_cp_async_ca_shared_global_16_s", + "llvm_nvvm_cp_async_cg_shared_global_16", + "llvm_nvvm_cp_async_cg_shared_global_16_s", + "llvm_nvvm_cp_async_commit_group", + "llvm_nvvm_cp_async_wait_all", + "llvm_nvvm_cp_async_wait_group", + ] + .map(str::to_owned); + expected.sort(); + assert_eq!(found, expected); + Ok(()) +} + +#[test] +fn test_generated_cp_async_libnvvm_uses_all_exact_inline_ptx() -> Result<(), anyhow::Error> { + let (ctx, module_ptr) = lower_all_classic_cp_async(mir_lower::IntrinsicBackend::LibNvvm)?; + let mut lowered = Vec::new(); + for op in lowered_kernel_body(&ctx, module_ptr) { + let Some(asm) = Operation::get_op::(op, &ctx) else { + continue; + }; + assert_eq!(llvm::asm_kind(&ctx, &asm), llvm::AsmKind::SideEffect); + assert!( + asm.get_attr_inline_asm_convergent(&ctx) + .is_some_and(|value| !bool::from((*value).clone())) + ); + lowered.push(( + asm.get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .unwrap(), + asm.get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .unwrap(), + )); + } + + let expected = [ + ( + "cp.async.ca.shared.global [%smem32], [%gmem64], 4;", + "l,l,~{memory}", + ), + ( + "cp.async.ca.shared.global [%smem32], [%gmem64], 8;", + "l,l,~{memory}", + ), + ( + "cp.async.ca.shared.global [%smem32], [%gmem64], 16;", + "l,l,~{memory}", + ), + ( + "cp.async.ca.shared.global [%smem32], [%gmem64], 4, $2;", + "l,l,r,~{memory}", + ), + ( + "cp.async.ca.shared.global [%smem32], [%gmem64], 8, $2;", + "l,l,r,~{memory}", + ), + ( + "cp.async.ca.shared.global [%smem32], [%gmem64], 16, $2;", + "l,l,r,~{memory}", + ), + ( + "cp.async.cg.shared.global [%smem32], [%gmem64], 16;", + "l,l,~{memory}", + ), + ( + "cp.async.cg.shared.global [%smem32], [%gmem64], 16, $2;", + "l,l,r,~{memory}", + ), + ("cp.async.commit_group;", "~{memory}"), + ("cp.async.wait_group $0;", "n,~{memory}"), + ("cp.async.wait_all;", "~{memory}"), + ]; + for (instruction, constraints) in expected { + assert_eq!( + lowered + .iter() + .filter(|(template, actual_constraints)| { + template.contains(instruction) && actual_constraints == constraints + }) + .count(), + 1, + "missing exact `{instruction}`" + ); + } + assert_eq!(lowered.len(), expected.len()); + Ok(()) +} + +#[test] +fn test_cp_async_ca_4_lowers_to_inline_asm() -> Result<(), anyhow::Error> { + use dialect_mir::types::MirPtrType; + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let dst_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), true); + let src_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), false); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![dst_ty.into(), src_ty.into()]); + + let dst = entry.deref(&ctx).get_argument(0); + let src = entry.deref(&ctx).get_argument(1); + + let op = Operation::new( + &mut ctx, + nvvm::CpAsyncCa4Op::get_concrete_op_info(), + vec![], + vec![dst, src], + vec![], + 0, + ); + op.insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + assert_cp_async_inline_asm_lowering(&mut ctx, module_ptr, 4) +} + +#[test] +fn test_cp_async_ca_8_lowers_to_inline_asm() -> Result<(), anyhow::Error> { + use dialect_mir::types::MirPtrType; + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let dst_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), true); + let src_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), false); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![dst_ty.into(), src_ty.into()]); + + let dst = entry.deref(&ctx).get_argument(0); + let src = entry.deref(&ctx).get_argument(1); + + let op = Operation::new( + &mut ctx, + nvvm::CpAsyncCa8Op::get_concrete_op_info(), + vec![], + vec![dst, src], + vec![], + 0, + ); + op.insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + assert_cp_async_inline_asm_lowering(&mut ctx, module_ptr, 8) +} + +fn assert_cp_async_inline_asm_lowering( + ctx: &mut Context, + module_ptr: pliron::context::Ptr, + copy_size: u32, +) -> Result<(), anyhow::Error> { + use pliron::r#type::Typed; + + mir_lower::lower_mir_to_llvm_with_options( + ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: mir_lower::IntrinsicBackend::LibNvvm, + ..Default::default() + }, + ) + .map_err(|e| anyhow::anyhow!("{e}"))?; + + let expected_template = format!( + "{{ .reg .u64 %smem64; .reg .u32 %smem32; .reg .u64 %gmem64; \ + cvta.to.shared.u64 %smem64, $0; cvt.u32.u64 %smem32, %smem64; \ + cvta.to.global.u64 %gmem64, $1; \ + cp.async.ca.shared.global [%smem32], [%gmem64], {copy_size}; }}" + ); + let mut matches = 0; + let module_region = module_ptr.deref(ctx).get_region(0); + let module_block = module_region.deref(ctx).iter(ctx).next().unwrap(); + + for op in module_block.deref(ctx).iter(ctx) { + let Some(func_op) = Operation::get_op::(op, ctx) else { + continue; + }; + if func_op.get_symbol_name(ctx).to_string() != "kernel_func" { + continue; + } + + let func_region = func_op.get_operation().deref(ctx).get_region(0); + for func_block in func_region.deref(ctx).iter(ctx) { + for body_op in func_block.deref(ctx).iter(ctx) { + let Some(inline_asm) = Operation::get_op::(body_op, ctx) else { + continue; + }; + let template = inline_asm + .get_attr_inline_asm_template(ctx) + .map(|s| String::from((*s).clone())); + if template.as_deref() != Some(expected_template.as_str()) { + continue; + } + + matches += 1; + assert_eq!( + inline_asm + .get_attr_inline_asm_constraints(ctx) + .map(|s| String::from((*s).clone())) + .as_deref(), + Some("l,l,~{memory}") + ); + assert_eq!(llvm::asm_kind(ctx, &inline_asm), llvm::AsmKind::SideEffect); + assert!( + inline_asm + .get_attr_inline_asm_convergent(ctx) + .is_some_and(|value| !bool::from((*value).clone())) + ); + + let operands: Vec<_> = inline_asm.get_operation().deref(ctx).operands().collect(); + assert_eq!(operands.len(), 2); + for operand in operands { + let ty = operand.get_type(ctx); + let ty = ty.deref(ctx); + let ptr_ty = ty + .downcast_ref::() + .expect("cp.async operands must lower to LLVM pointers"); + assert_eq!(ptr_ty.address_space(), 0); + } + } + } + } + + assert_eq!(matches, 1, "missing exact {copy_size}-byte cp.async asm"); + Ok(()) +} + +// ============================================================================= +// cp.async zero-fill lowering tests +// ============================================================================= + +#[test] +fn test_cp_async_ca_zfill_4_lowers_to_inline_asm() -> Result<(), anyhow::Error> { + use dialect_mir::types::MirPtrType; + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signless); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let dst_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), true); + let src_ty = MirPtrType::get_generic(&mut ctx, i8_ty.into(), false); + let (module_ptr, entry) = + build_test_kernel(&mut ctx, vec![dst_ty.into(), src_ty.into(), i32_ty.into()]); + + let dst = entry.deref(&ctx).get_argument(0); + let src = entry.deref(&ctx).get_argument(1); + let src_size = entry.deref(&ctx).get_argument(2); + + let op = Operation::new( + &mut ctx, + nvvm::CpAsyncCaZfill4Op::get_concrete_op_info(), + vec![], + vec![dst, src, src_size], + vec![], + 0, + ); + op.insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + assert_cp_async_zfill_inline_asm_lowering(&mut ctx, module_ptr, 4) +} + +#[test] +fn test_cp_async_ca_zfill_8_lowers_to_inline_asm() -> Result<(), anyhow::Error> { + use dialect_mir::types::MirPtrType; + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signless); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let dst_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), true); + let src_ty = MirPtrType::get_generic(&mut ctx, i8_ty.into(), false); + let (module_ptr, entry) = + build_test_kernel(&mut ctx, vec![dst_ty.into(), src_ty.into(), i32_ty.into()]); + + let dst = entry.deref(&ctx).get_argument(0); + let src = entry.deref(&ctx).get_argument(1); + let src_size = entry.deref(&ctx).get_argument(2); + + let op = Operation::new( + &mut ctx, + nvvm::CpAsyncCaZfill8Op::get_concrete_op_info(), + vec![], + vec![dst, src, src_size], + vec![], + 0, + ); + op.insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + assert_cp_async_zfill_inline_asm_lowering(&mut ctx, module_ptr, 8) +} + +#[test] +fn test_cp_async_ca_zfill_16_lowers_to_inline_asm() -> Result<(), anyhow::Error> { + use dialect_mir::types::MirPtrType; + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signless); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let dst_ty = MirPtrType::get_generic(&mut ctx, i32_ty.into(), true); + let src_ty = MirPtrType::get_generic(&mut ctx, i8_ty.into(), false); + let (module_ptr, entry) = + build_test_kernel(&mut ctx, vec![dst_ty.into(), src_ty.into(), i32_ty.into()]); + + let dst = entry.deref(&ctx).get_argument(0); + let src = entry.deref(&ctx).get_argument(1); + let src_size = entry.deref(&ctx).get_argument(2); + + let op = Operation::new( + &mut ctx, + nvvm::CpAsyncCaZfill16Op::get_concrete_op_info(), + vec![], + vec![dst, src, src_size], + vec![], + 0, + ); + op.insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + assert_cp_async_zfill_inline_asm_lowering(&mut ctx, module_ptr, 16) +} + +fn assert_cp_async_zfill_inline_asm_lowering( + ctx: &mut Context, + module_ptr: pliron::context::Ptr, + copy_size: u32, +) -> Result<(), anyhow::Error> { + use pliron::builtin::types::IntegerType; + use pliron::r#type::Typed; + + mir_lower::lower_mir_to_llvm_with_options( + ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: mir_lower::IntrinsicBackend::LibNvvm, + ..Default::default() + }, + ) + .map_err(|e| anyhow::anyhow!("{e}"))?; + + let expected_template = format!( + "{{ .reg .u64 %smem64; .reg .u32 %smem32; .reg .u64 %gmem64; \ + cvta.to.shared.u64 %smem64, $0; cvt.u32.u64 %smem32, %smem64; \ + cvta.to.global.u64 %gmem64, $1; \ + cp.async.ca.shared.global [%smem32], [%gmem64], {copy_size}, $2; }}" + ); + let mut matches = 0; + let module_region = module_ptr.deref(ctx).get_region(0); + let module_block = module_region.deref(ctx).iter(ctx).next().unwrap(); + + for op in module_block.deref(ctx).iter(ctx) { + let Some(func_op) = Operation::get_op::(op, ctx) else { + continue; + }; + if func_op.get_symbol_name(ctx).to_string() != "kernel_func" { + continue; + } + + let func_region = func_op.get_operation().deref(ctx).get_region(0); + for func_block in func_region.deref(ctx).iter(ctx) { + for body_op in func_block.deref(ctx).iter(ctx) { + let Some(inline_asm) = Operation::get_op::(body_op, ctx) else { + continue; + }; + let template = inline_asm + .get_attr_inline_asm_template(ctx) + .map(|s| String::from((*s).clone())); + if template.as_deref() != Some(expected_template.as_str()) { + continue; + } + + matches += 1; + assert_eq!( + inline_asm + .get_attr_inline_asm_constraints(ctx) + .map(|s| String::from((*s).clone())) + .as_deref(), + Some("l,l,r,~{memory}") + ); + assert_eq!(llvm::asm_kind(ctx, &inline_asm), llvm::AsmKind::SideEffect); + assert!( + inline_asm + .get_attr_inline_asm_convergent(ctx) + .is_some_and(|value| !bool::from((*value).clone())) + ); + + let operands: Vec<_> = inline_asm.get_operation().deref(ctx).operands().collect(); + assert_eq!(operands.len(), 3); + for operand in &operands[..2] { + let ty = operand.get_type(ctx); + let ty = ty.deref(ctx); + let ptr_ty = ty + .downcast_ref::() + .expect("cp.async pointer operands must lower to LLVM pointers"); + assert_eq!(ptr_ty.address_space(), 0); + } + + let src_size_ty = operands[2].get_type(ctx); + let src_size_ty = src_size_ty.deref(ctx); + let src_size_ty = src_size_ty + .downcast_ref::() + .expect("cp.async src_size must lower to an integer"); + assert_eq!(src_size_ty.width(), 32); + } + } + } + + assert_eq!( + matches, 1, + "missing exact {copy_size}-byte zero-fill cp.async asm" + ); + Ok(()) +} + +// ============================================================================= +// Generated packed arithmetic lowering tests +// ============================================================================= diff --git a/crates/mir-lower/tests/lowering_test/inline_ptx.rs b/crates/mir-lower/tests/lowering_test/inline_ptx.rs new file mode 100644 index 0000000000..e7cb6a9235 --- /dev/null +++ b/crates/mir-lower/tests/lowering_test/inline_ptx.rs @@ -0,0 +1,374 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use dialect_nvvm::ops as nvvm; +use llvm_export::ops as llvm; +use pliron::builtin::op_interfaces::SymbolOpInterface; +use pliron::linked_list::ContainsLinkedList; +use pliron::location::{Located, Location}; +use pliron::op::Op; +use pliron::operation::Operation; + +use crate::common::{append_return, build_test_kernel, lowered_kernel_body, make_test_ctx}; + +#[test] +fn test_inline_ptx_op_lowers_to_inline_asm_attrs() -> Result<(), anyhow::Error> { + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![i32_ty.into()]); + let input = entry.deref(&ctx).get_argument(0); + + let inline_ptx = nvvm::InlinePtxOp::build( + &mut ctx, + vec![i32_ty.into()], + vec![input], + "add.u32 $0, $1, $1;", + "=r,r", + true, + true, + ); + inline_ptx.insert_at_back(entry, &ctx); + let register_only_ptx = nvvm::InlinePtxOp::build( + &mut ctx, + vec![i32_ty.into()], + vec![input], + "mul.lo.u32 $0, $1, $1;", + "=r,r", + false, + true, + ); + register_only_ptx.insert_at_back(entry, &ctx); + let may_diverge_ptx = nvvm::InlinePtxOp::build( + &mut ctx, + vec![i32_ty.into()], + vec![input], + "cvt.u32.u32 $0, $1;", + "=r,r", + false, + false, + ); + may_diverge_ptx.insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{}", e))?; + + let mut found_conservative = false; + let mut found_register_only = false; + let mut found_may_diverge = false; + let module_op = module_ptr.deref(&ctx); + let region = module_op.get_region(0); + let block = region.deref(&ctx).iter(&ctx).next().unwrap(); + + for op in block.deref(&ctx).iter(&ctx) { + let Some(func_op) = Operation::get_op::(op, &ctx) else { + continue; + }; + if func_op.get_symbol_name(&ctx).to_string() != "kernel_func" { + continue; + } + let func_region = func_op.get_operation().deref(&ctx).get_region(0); + for func_block in func_region.deref(&ctx).iter(&ctx) { + for body_op in func_block.deref(&ctx).iter(&ctx) { + let Some(inline_asm) = Operation::get_op::(body_op, &ctx) else { + continue; + }; + let template = inline_asm + .get_attr_inline_asm_template(&ctx) + .map(|s| String::from((*s).clone())); + match template.as_deref() { + Some("add.u32 $0, $1, $1;") => { + found_conservative = true; + assert_eq!( + inline_asm + .get_attr_inline_asm_constraints(&ctx) + .map(|s| String::from((*s).clone())) + .as_deref(), + Some("=r,r") + ); + assert!( + inline_asm + .get_attr_inline_asm_convergent(&ctx) + .is_some_and(|b| bool::from((*b).clone())) + ); + assert!(llvm::inline_asm_sideeffect( + &ctx, + inline_asm.get_operation() + )); + } + Some("mul.lo.u32 $0, $1, $1;") => { + found_register_only = true; + assert_eq!( + inline_asm + .get_attr_inline_asm_constraints(&ctx) + .map(|s| String::from((*s).clone())) + .as_deref(), + Some("=r,r") + ); + assert!( + inline_asm + .get_attr_inline_asm_convergent(&ctx) + .is_some_and(|b| bool::from((*b).clone())) + ); + assert!(!llvm::inline_asm_sideeffect( + &ctx, + inline_asm.get_operation() + )); + } + Some("cvt.u32.u32 $0, $1;") => { + found_may_diverge = true; + assert_eq!( + inline_asm + .get_attr_inline_asm_constraints(&ctx) + .map(|s| String::from((*s).clone())) + .as_deref(), + Some("=r,r") + ); + assert!( + inline_asm + .get_attr_inline_asm_convergent(&ctx) + .is_some_and(|b| !bool::from((*b).clone())) + ); + assert!(!llvm::inline_asm_sideeffect( + &ctx, + inline_asm.get_operation() + )); + } + _ => continue, + } + } + } + } + + assert!( + found_conservative, + "Expected conservative inline PTX asm op" + ); + assert!( + found_register_only, + "Expected register-only inline PTX asm op" + ); + assert!(found_may_diverge, "Expected may-diverge inline PTX asm op"); + Ok(()) +} + +#[test] +fn test_multi_result_inline_ptx_lowers_to_struct_asm_and_extractvalues() -> Result<(), anyhow::Error> +{ + use llvm_export::types as llvm_types; + use pliron::builtin::types::{IntegerType, Signedness}; + use pliron::r#type::Typed; + + let mut ctx = make_test_ctx(); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![i32_ty.into()]); + let input = entry.deref(&ctx).get_argument(0); + + let inline_ptx = nvvm::InlinePtxOp::build( + &mut ctx, + vec![i32_ty.into(), i32_ty.into()], + vec![input], + "add.u32 $0, $2, $2; mul.lo.u32 $1, $2, $2;", + "=r,=r,r", + true, + false, + ); + let expected_location = Location::Named { + name: "source-inline-ptx".to_string(), + child_loc: Box::new(Location::Unknown), + }; + inline_ptx + .deref_mut(&ctx) + .set_loc(expected_location.clone()); + inline_ptx.insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{}", e))?; + + let mut asm_result = None; + let mut extract_indices = Vec::new(); + for op in lowered_kernel_body(&ctx, module_ptr) { + if let Some(inline_asm) = Operation::get_op::(op, &ctx) { + assert_eq!( + inline_asm.get_operation().deref(&ctx).loc(), + expected_location + ); + assert_eq!( + inline_asm + .get_attr_inline_asm_template(&ctx) + .map(|s| String::from((*s).clone())) + .as_deref(), + Some("add.u32 $0, $2, $2; mul.lo.u32 $1, $2, $2;") + ); + assert_eq!( + inline_asm + .get_attr_inline_asm_constraints(&ctx) + .map(|s| String::from((*s).clone())) + .as_deref(), + Some("=r,=r,r") + ); + let result = inline_asm.get_operation().deref(&ctx).get_result(0); + let result_ty = result.get_type(&ctx); + let result_ty = result_ty.deref(&ctx); + let struct_ty = result_ty + .downcast_ref::() + .expect("multi-output inline PTX must return an LLVM struct"); + assert_eq!(struct_ty.num_fields(), 2); + for index in 0..2 { + assert_eq!( + struct_ty + .field_type(index) + .deref(&ctx) + .downcast_ref::() + .expect("multi-output inline PTX struct field must stay i32") + .width(), + 32 + ); + } + asm_result = Some(result); + } else if let Some(extract) = Operation::get_op::(op, &ctx) { + assert_eq!(extract.get_operation().deref(&ctx).loc(), expected_location); + let aggregate = extract.get_operation().deref(&ctx).get_operand(0); + assert_eq!( + Some(aggregate), + asm_result, + "extractvalue must consume the struct-returning asm result" + ); + extract_indices.push(extract.indices(&ctx)); + } + } + + assert!( + asm_result.is_some(), + "Expected struct-returning inline PTX asm op" + ); + assert_eq!( + extract_indices, + vec![vec![0], vec![1]], + "each output must be extracted once, in constraint order" + ); + Ok(()) +} + +#[test] +fn test_inline_ptx_supports_thirty_two_tied_f32_results() -> Result<(), anyhow::Error> { + use llvm_export::types as llvm_types; + use pliron::builtin::types::FP32Type; + use pliron::r#type::Typed; + + const ACCUMULATOR_LEN: usize = 32; + + let mut ctx = make_test_ctx(); + let f32_ty = FP32Type::get(&ctx); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![f32_ty.into(); ACCUMULATOR_LEN]); + + let inputs = (0..ACCUMULATOR_LEN) + .map(|index| entry.deref(&ctx).get_argument(index)) + .collect::>(); + + let template = (0..ACCUMULATOR_LEN) + .map(|index| format!("mov.f32 ${index}, ${index};")) + .collect::>() + .join(" "); + + let mut constraints = vec!["=f".to_string(); ACCUMULATOR_LEN]; + constraints.extend((0..ACCUMULATOR_LEN).map(|index| index.to_string())); + let constraints = constraints.join(","); + + let inline_ptx = nvvm::InlinePtxOp::build( + &mut ctx, + vec![f32_ty.into(); ACCUMULATOR_LEN], + inputs, + &template, + &constraints, + true, + true, + ); + inline_ptx.insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let mut asm_result = None; + let mut extract_indices = Vec::new(); + + for op in lowered_kernel_body(&ctx, module_ptr) { + if let Some(inline_asm) = Operation::get_op::(op, &ctx) { + assert_eq!( + inline_asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some(template.as_str()), + ); + + assert_eq!( + inline_asm + .get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some(constraints.as_str()), + ); + + let result = inline_asm.get_operation().deref(&ctx).get_result(0); + let result_ty = result.get_type(&ctx); + let result_ty = result_ty.deref(&ctx); + + let struct_ty = result_ty + .downcast_ref::() + .expect("32-output inline PTX must return an LLVM struct"); + + assert_eq!( + struct_ty.num_fields(), + ACCUMULATOR_LEN, + "inline PTX must return exactly 32 accumulator values", + ); + + for index in 0..ACCUMULATOR_LEN { + assert!( + struct_ty + .field_type(index) + .deref(&ctx) + .downcast_ref::() + .is_some(), + "inline PTX result field {index} must remain f32", + ); + } + + assert!( + asm_result.replace(result).is_none(), + "expected exactly one inline PTX operation", + ); + } else if let Some(extract) = Operation::get_op::(op, &ctx) { + let aggregate = extract.get_operation().deref(&ctx).get_operand(0); + + assert_eq!( + Some(aggregate), + asm_result, + "extractvalue must consume the struct-returning asm result", + ); + + extract_indices.push(extract.indices(&ctx)); + } + } + + assert!( + asm_result.is_some(), + "expected a struct-returning inline PTX asm operation", + ); + + let expected_indices = (0..ACCUMULATOR_LEN) + .map(|index| vec![index as u32]) + .collect::>(); + + assert_eq!( + extract_indices, expected_indices, + "each accumulator output must be extracted exactly once and in constraint order", + ); + + Ok(()) +} diff --git a/crates/mir-lower/tests/lowering_test/main.rs b/crates/mir-lower/tests/lowering_test/main.rs new file mode 100644 index 0000000000..cdc0a69728 --- /dev/null +++ b/crates/mir-lower/tests/lowering_test/main.rs @@ -0,0 +1,21 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +// Tests build kinded fixture types directly; production minting lives in +// mir-importer's facts.rs (see the workspace clippy.toml disallowed-methods). +#![allow(clippy::disallowed_methods)] + +mod atomics_fences; +mod barriers_sync; +mod calls_and_values; +mod common; +mod cp_async; +mod inline_ptx; +mod math_conversions; +mod matrix_memory; +mod mma; +mod sregs_and_warp; +mod wgmma_lowering; +mod wgmma_rejections; diff --git a/crates/mir-lower/tests/lowering_test/math_conversions.rs b/crates/mir-lower/tests/lowering_test/math_conversions.rs new file mode 100644 index 0000000000..7fa48b5396 --- /dev/null +++ b/crates/mir-lower/tests/lowering_test/math_conversions.rs @@ -0,0 +1,1712 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use dialect_mir::ops as mir; +use dialect_nvvm::ops as nvvm; +use llvm_export::ops as llvm; +use pliron::builtin::op_interfaces::{CallOpCallable, CallOpInterface, SymbolOpInterface}; +use pliron::builtin::ops::ModuleOp; +use pliron::context::Context; +use pliron::linked_list::ContainsLinkedList; +use pliron::op::Op; +use pliron::operation::Operation; + +use crate::common::{append_return, build_test_kernel, lowered_kernel_body, make_test_ctx}; + +#[test] +fn test_fast_float_intrinsics_lower_to_explicit_fast_binops() -> Result<(), anyhow::Error> { + use dialect_mir::rust_intrinsics; + use llvm_export::attributes::{FastmathFlags, FastmathFlagsAttr}; + use llvm_export::op_interfaces::FastMathFlags; + use pliron::builtin::attributes::StringAttr; + use pliron::builtin::op_interfaces::CallOpInterface; + use pliron::builtin::types::{FP32Type, FP64Type}; + use pliron::r#type::{TypeHandle, Typed}; + + let mut ctx = make_test_ctx(); + let f32_ty = FP32Type::get(&ctx); + let f64_ty = FP64Type::get(&ctx); + let f32_ty_obj: TypeHandle = f32_ty.into(); + let f64_ty_obj: TypeHandle = f64_ty.into(); + let (module_ptr, entry) = build_test_kernel( + &mut ctx, + vec![f32_ty_obj, f32_ty_obj, f64_ty_obj, f64_ty_obj], + ); + let f32_lhs = entry.deref(&ctx).get_argument(0); + let f32_rhs = entry.deref(&ctx).get_argument(1); + let f64_lhs = entry.deref(&ctx).get_argument(2); + let f64_rhs = entry.deref(&ctx).get_argument(3); + + for (callee, lhs, rhs, result_ty) in [ + ( + rust_intrinsics::CALLEE_FADD_FAST, + f32_lhs, + f32_rhs, + f32_ty_obj, + ), + ( + rust_intrinsics::CALLEE_FSUB_FAST, + f32_lhs, + f32_rhs, + f32_ty_obj, + ), + ( + rust_intrinsics::CALLEE_FMUL_FAST, + f32_lhs, + f32_rhs, + f32_ty_obj, + ), + ( + rust_intrinsics::CALLEE_FDIV_FAST, + f32_lhs, + f32_rhs, + f32_ty_obj, + ), + ( + rust_intrinsics::CALLEE_FREM_FAST, + f32_lhs, + f32_rhs, + f32_ty_obj, + ), + ( + rust_intrinsics::CALLEE_FADD_FAST, + f64_lhs, + f64_rhs, + f64_ty_obj, + ), + ( + rust_intrinsics::CALLEE_FSUB_FAST, + f64_lhs, + f64_rhs, + f64_ty_obj, + ), + ( + rust_intrinsics::CALLEE_FMUL_FAST, + f64_lhs, + f64_rhs, + f64_ty_obj, + ), + ( + rust_intrinsics::CALLEE_FDIV_FAST, + f64_lhs, + f64_rhs, + f64_ty_obj, + ), + ( + rust_intrinsics::CALLEE_FREM_FAST, + f64_lhs, + f64_rhs, + f64_ty_obj, + ), + ] { + let call_ptr = Operation::new( + &mut ctx, + mir::MirCallOp::get_concrete_op_info(), + vec![result_ty], + vec![lhs, rhs], + vec![], + 0, + ); + let call = mir::MirCallOp::new(call_ptr); + call.set_attr_callee(&ctx, StringAttr::new(callee.to_string())); + call_ptr.insert_at_back(entry, &ctx); + } + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{}", e))?; + + let explicit_fast_flags: FastmathFlagsAttr = FastmathFlags::FAST.into(); + assert_ne!( + explicit_fast_flags, + FastmathFlagsAttr::default(), + "FastmathFlagsAttr::default() is empty; f*_fast must use explicit fast flags" + ); + + let mut fadd_counts = [0usize; 2]; + let mut fsub_counts = [0usize; 2]; + let mut fmul_counts = [0usize; 2]; + let mut fdiv_counts = [0usize; 2]; + let mut frem_counts = [0usize; 2]; + + macro_rules! count_fast_binop { + ($body_op:expr, $op_ty:ty, $counts:ident, $name:literal) => { + if let Some(op) = Operation::get_op::<$op_ty>($body_op, &ctx) { + assert_eq!( + op.fast_math_flags(&ctx), + explicit_fast_flags, + concat!($name, " must carry explicit LLVM fast-math flags") + ); + let result_ty = op.get_operation().deref(&ctx).get_result(0).get_type(&ctx); + if result_ty == f32_ty_obj { + $counts[0] += 1; + } else if result_ty == f64_ty_obj { + $counts[1] += 1; + } else { + panic!(concat!($name, " lowered to an unexpected result type")); + } + } + }; + } + + let module_op = module_ptr.deref(&ctx); + let region = module_op.get_region(0); + let block = region.deref(&ctx).iter(&ctx).next().unwrap(); + for op in block.deref(&ctx).iter(&ctx) { + let Some(func_op) = Operation::get_op::(op, &ctx) else { + continue; + }; + if func_op.get_symbol_name(&ctx).to_string() != "kernel_func" { + continue; + } + let func_region = func_op.get_operation().deref(&ctx).get_region(0); + for func_block in func_region.deref(&ctx).iter(&ctx) { + for body_op in func_block.deref(&ctx).iter(&ctx) { + assert!( + Operation::get_op::(body_op, &ctx).is_none(), + "f*_fast placeholder mir.call must not survive MIR lowering" + ); + if let Some(call) = Operation::get_op::(body_op, &ctx) + && let CallOpCallable::Direct(sym) = call.callee(&ctx) + { + let callee = sym.to_string(); + assert!( + !callee.starts_with(rust_intrinsics::PLACEHOLDER_PREFIX), + "lowered LLVM must not call unresolved Rust intrinsic placeholder `{callee}`" + ); + } + count_fast_binop!(body_op, llvm::FAddOp, fadd_counts, "fadd_fast"); + count_fast_binop!(body_op, llvm::FSubOp, fsub_counts, "fsub_fast"); + count_fast_binop!(body_op, llvm::FMulOp, fmul_counts, "fmul_fast"); + count_fast_binop!(body_op, llvm::FDivOp, fdiv_counts, "fdiv_fast"); + count_fast_binop!(body_op, llvm::FRemOp, frem_counts, "frem_fast"); + } + } + } + + assert_eq!(fadd_counts, [1, 1], "fadd_fast must lower for f32 and f64"); + assert_eq!(fsub_counts, [1, 1], "fsub_fast must lower for f32 and f64"); + assert_eq!(fmul_counts, [1, 1], "fmul_fast must lower for f32 and f64"); + assert_eq!(fdiv_counts, [1, 1], "fdiv_fast must lower for f32 and f64"); + assert_eq!(frem_counts, [1, 1], "frem_fast must lower for f32 and f64"); + + Ok(()) +} + +#[test] +fn test_ex2_approx_f16_uses_exact_pure_i16_inline_ptx_on_both_backends() -> Result<(), anyhow::Error> +{ + use pliron::builtin::types::{IntegerType, Signedness}; + use pliron::r#type::Typed; + + for backend in [ + mir_lower::IntrinsicBackend::LlvmNvptx, + mir_lower::IntrinsicBackend::LibNvvm, + ] { + let mut ctx = make_test_ctx(); + let i16_ty = IntegerType::get(&ctx, 16, Signedness::Unsigned); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![i16_ty.into()]); + let operand = entry.deref(&ctx).get_argument(0); + nvvm::ScalarMathOp::build( + &mut ctx, + operand, + nvvm::ScalarMathFormatAttr::F16, + nvvm::ScalarMathOperationAttr::Ex2, + nvvm::ScalarMathPrecisionAttr::Approx, + nvvm::ScalarMathSubnormalAttr::Preserve, + ) + .insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: backend, + ..Default::default() + }, + ) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let inline_asm = lowered_kernel_body(&ctx, module_ptr) + .into_iter() + .filter_map(|op| Operation::get_op::(op, &ctx)) + .collect::>(); + assert_eq!(inline_asm.len(), 1); + let inline_asm = &inline_asm[0]; + assert_eq!( + inline_asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some("ex2.approx.f16 $0, $1;") + ); + assert_eq!( + inline_asm + .get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some("=h,h") + ); + assert_eq!(llvm::asm_kind(&ctx, inline_asm), llvm::AsmKind::Pure); + assert_eq!( + inline_asm + .get_attr_inline_asm_convergent(&ctx) + .map(|value| bool::from((*value).clone())), + Some(false) + ); + + let op = inline_asm.get_operation().deref(&ctx); + assert_eq!(op.get_num_operands(), 1); + assert_eq!(op.get_num_results(), 1); + for ty in [ + op.get_operand(0).get_type(&ctx), + op.get_result(0).get_type(&ctx), + ] { + let ty = ty.deref(&ctx); + let integer = ty + .downcast_ref::() + .expect("f16 values use an i16 transport type"); + assert_eq!(integer.width(), 16); + } + + let module = Operation::get_op::(module_ptr, &ctx).unwrap(); + let ir = llvm_export::export::export_module_to_string(&ctx, &module) + .expect("ex2.approx.f16 module exports to LLVM IR"); + assert!( + ir.contains("call i16 asm \"ex2.approx.f16 $0, $1;\", \"=h,h\"(i16"), + "{ir}" + ); + assert!(!ir.contains("asm sideeffect"), "{ir}"); + assert!(!ir.contains("~{memory}"), "{ir}"); + } + Ok(()) +} + +#[test] +fn test_ex2_approx_f16_lowering_rejects_wrong_type_and_format() { + use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; + + for backend in [ + mir_lower::IntrinsicBackend::LlvmNvptx, + mir_lower::IntrinsicBackend::LibNvvm, + ] { + let mut wrong_type_ctx = make_test_ctx(); + let f32_ty = FP32Type::get(&wrong_type_ctx); + let (module_ptr, entry) = build_test_kernel(&mut wrong_type_ctx, vec![f32_ty.into()]); + let operand = entry.deref(&wrong_type_ctx).get_argument(0); + nvvm::ScalarMathOp::build( + &mut wrong_type_ctx, + operand, + nvvm::ScalarMathFormatAttr::F16, + nvvm::ScalarMathOperationAttr::Ex2, + nvvm::ScalarMathPrecisionAttr::Approx, + nvvm::ScalarMathSubnormalAttr::Preserve, + ) + .insert_at_back(entry, &wrong_type_ctx); + append_return(&mut wrong_type_ctx, entry); + let error = mir_lower::lower_mir_to_llvm_with_options( + &mut wrong_type_ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: backend, + ..Default::default() + }, + ) + .expect_err("an f16 operation with an f32 operand must fail closed") + .to_string(); + assert!( + error.contains("nvvm.scalar_math types do not match its format"), + "{error}" + ); + + let mut wrong_format_ctx = make_test_ctx(); + let i16_ty = IntegerType::get(&wrong_format_ctx, 16, Signedness::Unsigned); + let (module_ptr, entry) = build_test_kernel(&mut wrong_format_ctx, vec![i16_ty.into()]); + let operand = entry.deref(&wrong_format_ctx).get_argument(0); + nvvm::ScalarMathOp::build( + &mut wrong_format_ctx, + operand, + nvvm::ScalarMathFormatAttr::F16, + nvvm::ScalarMathOperationAttr::Ex2, + nvvm::ScalarMathPrecisionAttr::Approx, + nvvm::ScalarMathSubnormalAttr::Ftz, + ) + .insert_at_back(entry, &wrong_format_ctx); + append_return(&mut wrong_format_ctx, entry); + let error = mir_lower::lower_mir_to_llvm_with_options( + &mut wrong_format_ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: backend, + ..Default::default() + }, + ) + .expect_err("unadmitted ex2.approx.ftz.f16 must fail closed") + .to_string(); + assert!( + error.contains("nvvm.scalar_math variant is not admitted"), + "{error}" + ); + } +} + +// ============================================================================= +// Integer dot product (dp4a / dp2a) lowering tests +// ============================================================================= + +const DOT_PRODUCT_TYPED_INTRINSICS: [&str; 4] = [ + "llvm_nvvm_idp4a_s_s", + "llvm_nvvm_idp4a_u_u", + "llvm_nvvm_idp2a_s_s", + "llvm_nvvm_idp2a_u_u", +]; + +const DOT_PRODUCT_PTX: [&str; 4] = [ + "dp4a.s32.s32 $0, $1, $2, $3;", + "dp4a.u32.u32 $0, $1, $2, $3;", + "dp2a.lo.s32.s32 $0, $1, $2, $3;", + "dp2a.lo.u32.u32 $0, $1, $2, $3;", +]; + +fn lower_all_dot_product_forms( + backend: mir_lower::IntrinsicBackend, +) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let (module_ptr, entry) = + build_test_kernel(&mut ctx, vec![i32_ty.into(), i32_ty.into(), i32_ty.into()]); + let operands = (0..3) + .map(|index| entry.deref(&ctx).get_argument(index)) + .collect::>(); + for op_info in [ + nvvm::Dp4aS32Op::get_concrete_op_info(), + nvvm::Dp4aU32Op::get_concrete_op_info(), + nvvm::Dp2aS32Op::get_concrete_op_info(), + nvvm::Dp2aU32Op::get_concrete_op_info(), + ] { + Operation::new( + &mut ctx, + op_info, + vec![i32_ty.into()], + operands.clone(), + vec![], + 0, + ) + .insert_at_back(entry, &ctx); + } + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: backend, + ..Default::default() + }, + ) + .map_err(|error| anyhow::anyhow!("{error}"))?; + Ok((ctx, module_ptr)) +} + +#[test] +fn test_dot_product_llvm_nvptx_uses_typed_intrinsics_and_low_selector() -> Result<(), anyhow::Error> +{ + use pliron::builtin::attributes::IntegerAttr; + + let (ctx, module_ptr) = lower_all_dot_product_forms(mir_lower::IntrinsicBackend::LlvmNvptx)?; + let body = lowered_kernel_body(&ctx, module_ptr); + let mut calls = Vec::new(); + for op in body { + assert!( + Operation::get_op::(op, &ctx).is_none(), + "LLVM-NVPTX dot products must use typed intrinsics" + ); + let Some(call) = Operation::get_op::(op, &ctx) else { + continue; + }; + let CallOpCallable::Direct(callee) = call.callee(&ctx) else { + continue; + }; + let callee = callee.to_string(); + if !callee.starts_with("llvm_nvvm_idp") { + continue; + } + let expected_arity = if callee.contains("idp2a") { 4 } else { 3 }; + assert_eq!(op.deref(&ctx).get_num_operands(), expected_arity); + if expected_arity == 4 { + let selector = op.deref(&ctx).get_operand(2); + let defining_op = selector.defining_op().expect("selector is a constant"); + let constant = Operation::get_op::(defining_op, &ctx) + .expect("selector is an LLVM constant"); + let attribute = constant.get_value(&ctx); + let integer = attribute + .downcast_ref::() + .expect("selector constant is an integer"); + assert_eq!(integer.value().bw(), 1); + assert_eq!(integer.value().to_u64(), 0, "dp2a must select `.lo`"); + } + calls.push(callee); + } + calls.sort(); + let mut expected = DOT_PRODUCT_TYPED_INTRINSICS.map(str::to_owned); + expected.sort(); + assert_eq!(calls, expected); + Ok(()) +} + +#[test] +fn test_dot_product_libnvvm_uses_exact_pure_inline_ptx() -> Result<(), anyhow::Error> { + let (ctx, module_ptr) = lower_all_dot_product_forms(mir_lower::IntrinsicBackend::LibNvvm)?; + let body = lowered_kernel_body(&ctx, module_ptr); + let mut inline_ptx = Vec::new(); + for op in body { + assert!( + Operation::get_op::(op, &ctx).is_none(), + "libNVVM dot products must not use typed intrinsic calls" + ); + let Some(inline_asm) = Operation::get_op::(op, &ctx) else { + continue; + }; + inline_ptx.push( + inline_asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .unwrap_or_default(), + ); + assert_eq!( + inline_asm + .get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some("=r,r,r,r") + ); + assert_eq!(llvm::asm_kind(&ctx, &inline_asm), llvm::AsmKind::Pure); + assert_eq!(op.deref(&ctx).get_num_operands(), 3); + assert_eq!(op.deref(&ctx).get_num_results(), 1); + } + inline_ptx.sort(); + let mut expected = DOT_PRODUCT_PTX.map(str::to_owned); + expected.sort(); + assert_eq!(inline_ptx, expected); + Ok(()) +} + +// ============================================================================= +// Byte permutation lowering tests +// ============================================================================= + +const PRMT_TYPED_INTRINSICS: [(&str, usize); 7] = [ + ("llvm_nvvm_prmt", 3), + ("llvm_nvvm_prmt_f4e", 3), + ("llvm_nvvm_prmt_b4e", 3), + ("llvm_nvvm_prmt_rc8", 2), + ("llvm_nvvm_prmt_ecl", 2), + ("llvm_nvvm_prmt_ecr", 2), + ("llvm_nvvm_prmt_rc16", 2), +]; + +const PRMT_INLINE_PTX: [(&str, &str, usize); 7] = [ + ("prmt.b32 $0, $1, $2, $3;", "=r,r,r,r", 3), + ("prmt.b32.f4e $0, $1, $2, $3;", "=r,r,r,r", 3), + ("prmt.b32.b4e $0, $1, $2, $3;", "=r,r,r,r", 3), + ("prmt.b32.rc8 $0, $1, 0, $2;", "=r,r,r", 2), + ("prmt.b32.ecl $0, $1, 0, $2;", "=r,r,r", 2), + ("prmt.b32.ecr $0, $1, 0, $2;", "=r,r,r", 2), + ("prmt.b32.rc16 $0, $1, 0, $2;", "=r,r,r", 2), +]; + +fn lower_all_prmt_modes( + backend: mir_lower::IntrinsicBackend, +) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let (module_ptr, entry) = + build_test_kernel(&mut ctx, vec![i32_ty.into(), i32_ty.into(), i32_ty.into()]); + let operands = (0..3) + .map(|index| entry.deref(&ctx).get_argument(index)) + .collect::>(); + + for mode in [ + nvvm::PrmtModeAttr::Generic, + nvvm::PrmtModeAttr::F4e, + nvvm::PrmtModeAttr::B4e, + ] { + nvvm::PrmtOp::build(&mut ctx, operands.clone(), mode).insert_at_back(entry, &ctx); + } + for mode in [ + nvvm::PrmtModeAttr::Rc8, + nvvm::PrmtModeAttr::Ecl, + nvvm::PrmtModeAttr::Ecr, + nvvm::PrmtModeAttr::Rc16, + ] { + nvvm::PrmtOp::build(&mut ctx, vec![operands[0], operands[2]], mode) + .insert_at_back(entry, &ctx); + } + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: backend, + ..Default::default() + }, + ) + .map_err(|error| anyhow::anyhow!("{error}"))?; + Ok((ctx, module_ptr)) +} + +#[test] +fn test_prmt_llvm_nvptx_uses_exact_typed_intrinsics() -> Result<(), anyhow::Error> { + let (ctx, module_ptr) = lower_all_prmt_modes(mir_lower::IntrinsicBackend::LlvmNvptx)?; + let body = lowered_kernel_body(&ctx, module_ptr); + let mut calls = Vec::new(); + for op in body { + assert!( + Operation::get_op::(op, &ctx).is_none(), + "LLVM-NVPTX byte permutations must use typed intrinsics" + ); + let Some(call) = Operation::get_op::(op, &ctx) else { + continue; + }; + let CallOpCallable::Direct(callee) = call.callee(&ctx) else { + continue; + }; + let callee = callee.to_string(); + let Some((_, expected_arity)) = PRMT_TYPED_INTRINSICS + .iter() + .find(|(expected, _)| *expected == callee) + else { + continue; + }; + assert_eq!(op.deref(&ctx).get_num_operands(), *expected_arity); + calls.push((callee, *expected_arity)); + } + calls.sort(); + let mut expected = PRMT_TYPED_INTRINSICS.map(|(name, arity)| (name.to_owned(), arity)); + expected.sort(); + assert_eq!(calls, expected); + Ok(()) +} + +#[test] +fn test_prmt_libnvvm_uses_exact_pure_inline_ptx() -> Result<(), anyhow::Error> { + let (ctx, module_ptr) = lower_all_prmt_modes(mir_lower::IntrinsicBackend::LibNvvm)?; + let body = lowered_kernel_body(&ctx, module_ptr); + let mut inline_ptx = Vec::new(); + for op in body { + if let Some(call) = Operation::get_op::(op, &ctx) + && let CallOpCallable::Direct(callee) = call.callee(&ctx) + { + assert!( + !callee.to_string().starts_with("llvm_nvvm_prmt"), + "libNVVM byte permutations must not use typed intrinsic calls" + ); + } + let Some(inline_asm) = Operation::get_op::(op, &ctx) else { + continue; + }; + let template = inline_asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .unwrap_or_default(); + let Some((_, expected_constraints, expected_arity)) = PRMT_INLINE_PTX + .iter() + .find(|(expected, _, _)| *expected == template) + else { + continue; + }; + assert_eq!( + inline_asm + .get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some(*expected_constraints) + ); + assert_eq!(llvm::asm_kind(&ctx, &inline_asm), llvm::AsmKind::Pure); + assert_eq!(op.deref(&ctx).get_num_operands(), *expected_arity); + assert_eq!(op.deref(&ctx).get_num_results(), 1); + inline_ptx.push((template, *expected_constraints, *expected_arity)); + } + inline_ptx.sort(); + let mut expected = PRMT_INLINE_PTX + .map(|(template, constraints, arity)| (template.to_owned(), constraints, arity)); + expected.sort(); + assert_eq!(inline_ptx, expected); + Ok(()) +} + +#[test] +fn test_generated_packed_arithmetic_lowers_to_exact_pure_inline_asm() -> Result<(), anyhow::Error> { + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let (module_ptr, entry) = + build_test_kernel(&mut ctx, vec![i32_ty.into(), i32_ty.into(), i32_ty.into()]); + + type OpInfo = ( + fn(pliron::context::Ptr) -> pliron::op::OpObj, + std::any::TypeId, + ); + let cases: [(OpInfo, usize, &str, &str); 18] = [ + ( + nvvm::FmaBf16x2Op::get_concrete_op_info(), + 3, + "fma.rn.bf16x2 $0, $1, $2, $3;", + "=r,r,r,r", + ), + ( + nvvm::FmaReluBf16x2Op::get_concrete_op_info(), + 3, + "fma.rn.relu.bf16x2 $0, $1, $2, $3;", + "=r,r,r,r", + ), + ( + nvvm::AddBf16x2Op::get_concrete_op_info(), + 2, + "add.rn.bf16x2 $0, $1, $2;", + "=r,r,r", + ), + ( + nvvm::SubBf16x2Op::get_concrete_op_info(), + 2, + "sub.rn.bf16x2 $0, $1, $2;", + "=r,r,r", + ), + ( + nvvm::MulBf16x2Op::get_concrete_op_info(), + 2, + "mul.rn.bf16x2 $0, $1, $2;", + "=r,r,r", + ), + ( + nvvm::MinBf16x2Op::get_concrete_op_info(), + 2, + "min.bf16x2 $0, $1, $2;", + "=r,r,r", + ), + ( + nvvm::MaxBf16x2Op::get_concrete_op_info(), + 2, + "max.bf16x2 $0, $1, $2;", + "=r,r,r", + ), + ( + nvvm::NegBf16x2Op::get_concrete_op_info(), + 1, + "neg.bf16x2 $0, $1;", + "=r,r", + ), + ( + nvvm::AbsBf16x2Op::get_concrete_op_info(), + 1, + "abs.bf16x2 $0, $1;", + "=r,r", + ), + ( + nvvm::FmaF16x2Op::get_concrete_op_info(), + 3, + "fma.rn.f16x2 $0, $1, $2, $3;", + "=r,r,r,r", + ), + ( + nvvm::FmaReluF16x2Op::get_concrete_op_info(), + 3, + "fma.rn.relu.f16x2 $0, $1, $2, $3;", + "=r,r,r,r", + ), + ( + nvvm::AddF16x2Op::get_concrete_op_info(), + 2, + "add.rn.f16x2 $0, $1, $2;", + "=r,r,r", + ), + ( + nvvm::SubF16x2Op::get_concrete_op_info(), + 2, + "sub.rn.f16x2 $0, $1, $2;", + "=r,r,r", + ), + ( + nvvm::MulF16x2Op::get_concrete_op_info(), + 2, + "mul.rn.f16x2 $0, $1, $2;", + "=r,r,r", + ), + ( + nvvm::MinF16x2Op::get_concrete_op_info(), + 2, + "min.f16x2 $0, $1, $2;", + "=r,r,r", + ), + ( + nvvm::MaxF16x2Op::get_concrete_op_info(), + 2, + "max.f16x2 $0, $1, $2;", + "=r,r,r", + ), + ( + nvvm::NegF16x2Op::get_concrete_op_info(), + 1, + "neg.f16x2 $0, $1;", + "=r,r", + ), + ( + nvvm::AbsF16x2Op::get_concrete_op_info(), + 1, + "abs.f16x2 $0, $1;", + "=r,r", + ), + ]; + + let operands = [ + entry.deref(&ctx).get_argument(0), + entry.deref(&ctx).get_argument(1), + entry.deref(&ctx).get_argument(2), + ]; + for &(op_info, operand_count, _, _) in &cases { + let op = Operation::new( + &mut ctx, + op_info, + vec![i32_ty.into()], + operands[..operand_count].to_vec(), + vec![], + 0, + ); + op.insert_at_back(entry, &ctx); + } + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{}", e))?; + + let mut lowered = Vec::new(); + let module_op = module_ptr.deref(&ctx); + let region = module_op.get_region(0); + let block = region.deref(&ctx).iter(&ctx).next().unwrap(); + for op in block.deref(&ctx).iter(&ctx) { + let Some(func_op) = Operation::get_op::(op, &ctx) else { + continue; + }; + if func_op.get_symbol_name(&ctx).to_string() != "kernel_func" { + continue; + } + let func_region = func_op.get_operation().deref(&ctx).get_region(0); + for func_block in func_region.deref(&ctx).iter(&ctx) { + for body_op in func_block.deref(&ctx).iter(&ctx) { + let Some(inline_asm) = Operation::get_op::(body_op, &ctx) else { + continue; + }; + let inline_asm_op = inline_asm.get_operation(); + let operand_count = inline_asm_op.deref(&ctx).operands().count(); + let result_count = inline_asm_op.deref(&ctx).get_num_results(); + lowered.push(( + inline_asm + .get_attr_inline_asm_template(&ctx) + .map(|s| String::from((*s).clone())) + .expect("packed inline asm must have a template"), + inline_asm + .get_attr_inline_asm_constraints(&ctx) + .map(|s| String::from((*s).clone())) + .expect("packed inline asm must have constraints"), + llvm::asm_kind_opt(&ctx, &inline_asm), + inline_asm + .get_attr_inline_asm_convergent(&ctx) + .map(|b| bool::from((*b).clone())), + operand_count, + result_count, + )); + } + } + } + + assert_eq!( + lowered.len(), + cases.len(), + "each packed operation must lower to exactly one inline-asm op" + ); + for &(_, expected_operand_count, expected_template, expected_constraints) in &cases { + let matches: Vec<_> = lowered + .iter() + .filter(|(template, _, _, _, _, _)| template == expected_template) + .collect(); + assert_eq!( + matches.len(), + 1, + "expected one exact `{expected_template}` lowering" + ); + let (_, constraints, kind, convergent, operand_count, result_count) = matches[0]; + assert_eq!(constraints, expected_constraints, "{expected_template}"); + assert_eq!(*kind, Some(llvm::AsmKind::Pure), "{expected_template}"); + assert_eq!(*convergent, Some(false), "{expected_template}"); + assert_eq!( + *operand_count, expected_operand_count, + "{expected_template} input arity" + ); + assert_eq!(*result_count, 1, "{expected_template} result arity"); + } + + Ok(()) +} + +#[test] +fn test_generated_integer_minmax_lowers_to_exact_pure_inline_asm() -> Result<(), anyhow::Error> { + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![i32_ty.into(), i32_ty.into()]); + + type OpInfo = ( + fn(pliron::context::Ptr) -> pliron::op::OpObj, + std::any::TypeId, + ); + let cases: [(OpInfo, &str); 8] = [ + ( + nvvm::MinReluS32Op::get_concrete_op_info(), + "min.relu.s32 $0, $1, $2;", + ), + ( + nvvm::MaxReluS32Op::get_concrete_op_info(), + "max.relu.s32 $0, $1, $2;", + ), + ( + nvvm::MinS16x2Op::get_concrete_op_info(), + "min.s16x2 $0, $1, $2;", + ), + ( + nvvm::MaxS16x2Op::get_concrete_op_info(), + "max.s16x2 $0, $1, $2;", + ), + ( + nvvm::MinU16x2Op::get_concrete_op_info(), + "min.u16x2 $0, $1, $2;", + ), + ( + nvvm::MaxU16x2Op::get_concrete_op_info(), + "max.u16x2 $0, $1, $2;", + ), + ( + nvvm::MinReluS16x2Op::get_concrete_op_info(), + "min.relu.s16x2 $0, $1, $2;", + ), + ( + nvvm::MaxReluS16x2Op::get_concrete_op_info(), + "max.relu.s16x2 $0, $1, $2;", + ), + ]; + + let operands = [ + entry.deref(&ctx).get_argument(0), + entry.deref(&ctx).get_argument(1), + ]; + for &(op_info, _) in &cases { + let op = Operation::new( + &mut ctx, + op_info, + vec![i32_ty.into()], + operands.to_vec(), + vec![], + 0, + ); + op.insert_at_back(entry, &ctx); + } + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{}", e))?; + + let mut lowered = Vec::new(); + let module_op = module_ptr.deref(&ctx); + let region = module_op.get_region(0); + let block = region.deref(&ctx).iter(&ctx).next().unwrap(); + for op in block.deref(&ctx).iter(&ctx) { + let Some(func_op) = Operation::get_op::(op, &ctx) else { + continue; + }; + if func_op.get_symbol_name(&ctx).to_string() != "kernel_func" { + continue; + } + let func_region = func_op.get_operation().deref(&ctx).get_region(0); + for func_block in func_region.deref(&ctx).iter(&ctx) { + for body_op in func_block.deref(&ctx).iter(&ctx) { + let Some(inline_asm) = Operation::get_op::(body_op, &ctx) else { + continue; + }; + let inline_asm_op = inline_asm.get_operation(); + let operand_count = inline_asm_op.deref(&ctx).operands().count(); + let result_count = inline_asm_op.deref(&ctx).get_num_results(); + lowered.push(( + inline_asm + .get_attr_inline_asm_template(&ctx) + .map(|s| String::from((*s).clone())) + .expect("integer min/max inline asm must have a template"), + inline_asm + .get_attr_inline_asm_constraints(&ctx) + .map(|s| String::from((*s).clone())) + .expect("integer min/max inline asm must have constraints"), + llvm::asm_kind_opt(&ctx, &inline_asm), + inline_asm + .get_attr_inline_asm_convergent(&ctx) + .map(|b| bool::from((*b).clone())), + operand_count, + result_count, + )); + } + } + } + + assert_eq!( + lowered.len(), + cases.len(), + "each integer min/max operation must lower to exactly one inline-asm op" + ); + for &(_, expected_template) in &cases { + let matches: Vec<_> = lowered + .iter() + .filter(|(template, _, _, _, _, _)| template == expected_template) + .collect(); + assert_eq!( + matches.len(), + 1, + "expected one exact `{expected_template}` lowering" + ); + let (_, constraints, kind, convergent, operand_count, result_count) = matches[0]; + assert_eq!(constraints, "=r,r,r", "{expected_template}"); + assert_eq!(*kind, Some(llvm::AsmKind::Pure), "{expected_template}"); + assert_eq!(*convergent, Some(false), "{expected_template}"); + assert_eq!(*operand_count, 2, "{expected_template} input arity"); + assert_eq!(*result_count, 1, "{expected_template} result arity"); + } + + Ok(()) +} + +#[test] +fn test_generated_packed_conversions_lower_to_exact_pure_inline_asm() -> Result<(), anyhow::Error> { + use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let f32_ty = FP32Type::get(&ctx); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![f32_ty.into(), f32_ty.into()]); + let low = entry.deref(&ctx).get_argument(0); + let high = entry.deref(&ctx).get_argument(1); + + type OpInfo = ( + fn(pliron::context::Ptr) -> pliron::op::OpObj, + std::any::TypeId, + ); + let cases: [(OpInfo, &str); 6] = [ + ( + nvvm::CvtF32x2Bf16x2Op::get_concrete_op_info(), + "cvt.rn.bf16x2.f32 $0, $2, $1;", + ), + ( + nvvm::CvtF16x2F32Op::get_concrete_op_info(), + "cvt.rn.f16x2.f32 $0, $2, $1;", + ), + ( + nvvm::CvtRzF16x2F32Op::get_concrete_op_info(), + "cvt.rz.f16x2.f32 $0, $2, $1;", + ), + ( + nvvm::CvtRnReluF16x2F32Op::get_concrete_op_info(), + "cvt.rn.relu.f16x2.f32 $0, $2, $1;", + ), + ( + nvvm::CvtRnReluBf16x2F32Op::get_concrete_op_info(), + "cvt.rn.relu.bf16x2.f32 $0, $2, $1;", + ), + ( + nvvm::CvtRzBf16x2F32Op::get_concrete_op_info(), + "cvt.rz.bf16x2.f32 $0, $2, $1;", + ), + ]; + for &(op_info, _) in &cases { + let op = Operation::new( + &mut ctx, + op_info, + vec![i32_ty.into()], + vec![low, high], + vec![], + 0, + ); + op.insert_at_back(entry, &ctx); + } + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{}", e))?; + + let mut lowered = Vec::new(); + let module_region = module_ptr.deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + for op in module_block.deref(&ctx).iter(&ctx) { + let Some(func) = Operation::get_op::(op, &ctx) else { + continue; + }; + if func.get_symbol_name(&ctx).to_string() != "kernel_func" { + continue; + } + let region = func.get_operation().deref(&ctx).get_region(0); + for block in region.deref(&ctx).iter(&ctx) { + for body_op in block.deref(&ctx).iter(&ctx) { + let Some(asm) = Operation::get_op::(body_op, &ctx) else { + continue; + }; + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .expect("packed conversion must have an asm template"); + if !template.starts_with("cvt.") { + continue; + } + lowered.push(( + template, + asm.get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())), + asm.get_attr_inline_asm_convergent(&ctx) + .map(|value| bool::from((*value).clone())), + llvm::asm_kind_opt(&ctx, &asm), + asm.get_operation().deref(&ctx).operands().count(), + asm.get_operation().deref(&ctx).get_num_results(), + )); + } + } + } + + assert_eq!( + lowered.len(), + cases.len(), + "each packed conversion must lower to one inline-asm op" + ); + for &(_, expected_template) in &cases { + let matches: Vec<_> = lowered + .iter() + .filter(|(template, _, _, _, _, _)| template == expected_template) + .collect(); + assert_eq!( + matches.len(), + 1, + "expected one exact `{expected_template}` lowering" + ); + let (_, constraints, convergent, kind, operands, results) = matches[0]; + assert_eq!( + constraints.as_deref(), + Some("=r,f,f"), + "{expected_template}" + ); + assert_eq!(*convergent, Some(false), "{expected_template}"); + assert_eq!(*kind, Some(llvm::AsmKind::Pure), "{expected_template}"); + assert_eq!(*operands, 2, "{expected_template} input arity"); + assert_eq!(*results, 1, "{expected_template} result arity"); + } + + Ok(()) +} + +const FP8_CONVERSION_INTRINSICS: [&str; 4] = [ + "llvm_nvvm_ff_to_e4m3x2_rn", + "llvm_nvvm_ff_to_e4m3x2_rn_relu", + "llvm_nvvm_ff_to_e5m2x2_rn", + "llvm_nvvm_ff_to_e5m2x2_rn_relu", +]; + +const FP8_CONVERSION_PTX: [&str; 4] = [ + "cvt.rn.satfinite.e4m3x2.f32 $0, $2, $1;", + "cvt.rn.satfinite.relu.e4m3x2.f32 $0, $2, $1;", + "cvt.rn.satfinite.e5m2x2.f32 $0, $2, $1;", + "cvt.rn.satfinite.relu.e5m2x2.f32 $0, $2, $1;", +]; + +fn lower_all_fp8_conversions( + backend: mir_lower::IntrinsicBackend, +) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { + use pliron::builtin::types::FP32Type; + + let mut ctx = make_test_ctx(); + let f32_ty = FP32Type::get(&ctx); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![f32_ty.into(), f32_ty.into()]); + let low = entry.deref(&ctx).get_argument(0); + let high = entry.deref(&ctx).get_argument(1); + + nvvm::CvtRnSatfiniteE4m3x2F32Op::build(&mut ctx, low, high).insert_at_back(entry, &ctx); + nvvm::CvtRnSatfiniteReluE4m3x2F32Op::build(&mut ctx, low, high).insert_at_back(entry, &ctx); + nvvm::CvtRnSatfiniteE5m2x2F32Op::build(&mut ctx, low, high).insert_at_back(entry, &ctx); + nvvm::CvtRnSatfiniteReluE5m2x2F32Op::build(&mut ctx, low, high).insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: backend, + ..Default::default() + }, + ) + .map_err(|error| anyhow::anyhow!("{error}"))?; + Ok((ctx, module_ptr)) +} + +#[test] +fn test_fp8_conversions_llvm_nvptx_use_exact_typed_calls() -> Result<(), anyhow::Error> { + use pliron::builtin::types::IntegerType; + use pliron::r#type::Typed; + + let (ctx, module_ptr) = lower_all_fp8_conversions(mir_lower::IntrinsicBackend::LlvmNvptx)?; + let mut calls = Vec::new(); + for op in lowered_kernel_body(&ctx, module_ptr) { + assert!( + Operation::get_op::(op, &ctx).is_none(), + "LLVM-NVPTX FP8 conversions must use typed intrinsics" + ); + let Some(call) = Operation::get_op::(op, &ctx) else { + continue; + }; + let CallOpCallable::Direct(callee) = call.callee(&ctx) else { + continue; + }; + let callee = callee.to_string(); + if !FP8_CONVERSION_INTRINSICS.contains(&callee.as_str()) { + continue; + } + let call_op = call.get_operation(); + assert_eq!(call_op.deref(&ctx).get_num_operands(), 2); + assert_eq!(call_op.deref(&ctx).get_num_results(), 1); + let block = call_op.deref(&ctx).get_parent_block().unwrap(); + assert_eq!( + call_op.deref(&ctx).get_operand(0), + block.deref(&ctx).get_argument(1) + ); + assert_eq!( + call_op.deref(&ctx).get_operand(1), + block.deref(&ctx).get_argument(0) + ); + let result_ty = call_op.deref(&ctx).get_result(0).get_type(&ctx); + assert_eq!( + result_ty + .deref(&ctx) + .downcast_ref::() + .expect("FP8 conversion result is an integer") + .width(), + 16 + ); + calls.push(callee); + } + calls.sort(); + let mut expected = FP8_CONVERSION_INTRINSICS.map(str::to_owned); + expected.sort(); + assert_eq!(calls, expected); + Ok(()) +} + +#[test] +fn test_fp8_conversions_libnvvm_use_exact_pure_inline_ptx() -> Result<(), anyhow::Error> { + use pliron::builtin::types::IntegerType; + use pliron::r#type::Typed; + + let (ctx, module_ptr) = lower_all_fp8_conversions(mir_lower::IntrinsicBackend::LibNvvm)?; + let mut templates = Vec::new(); + for op in lowered_kernel_body(&ctx, module_ptr) { + if let Some(call) = Operation::get_op::(op, &ctx) + && let CallOpCallable::Direct(callee) = call.callee(&ctx) + { + assert!( + !FP8_CONVERSION_INTRINSICS.contains(&callee.as_ref()), + "libNVVM FP8 conversions must not use typed intrinsics" + ); + } + let Some(inline_asm) = Operation::get_op::(op, &ctx) else { + continue; + }; + let template = inline_asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .unwrap_or_default(); + if !FP8_CONVERSION_PTX.contains(&template.as_str()) { + continue; + } + assert_eq!( + inline_asm + .get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some("=h,f,f") + ); + assert_eq!(llvm::asm_kind(&ctx, &inline_asm), llvm::AsmKind::Pure); + assert_eq!( + inline_asm + .get_attr_inline_asm_convergent(&ctx) + .map(|value| bool::from((*value).clone())), + Some(false) + ); + let asm_op = inline_asm.get_operation(); + assert_eq!(asm_op.deref(&ctx).get_num_operands(), 2); + assert_eq!(asm_op.deref(&ctx).get_num_results(), 1); + let result_ty = asm_op.deref(&ctx).get_result(0).get_type(&ctx); + assert_eq!( + result_ty + .deref(&ctx) + .downcast_ref::() + .expect("FP8 conversion result is an integer") + .width(), + 16 + ); + templates.push(template); + } + templates.sort(); + let mut expected = FP8_CONVERSION_PTX.map(str::to_owned); + expected.sort(); + assert_eq!(templates, expected); + Ok(()) +} + +const SCALAR_CONVERSION_INTRINSICS: [&str; 10] = [ + "llvm_nvvm_f2tf32_rna", + "llvm_nvvm_f2tf32_rna_satfinite", + "llvm_nvvm_f2tf32_rn", + "llvm_nvvm_f2tf32_rn_relu", + "llvm_nvvm_f2tf32_rn_satfinite", + "llvm_nvvm_f2tf32_rn_relu_satfinite", + "llvm_nvvm_f2tf32_rz", + "llvm_nvvm_f2tf32_rz_relu", + "llvm_nvvm_f2tf32_rz_satfinite", + "llvm_nvvm_f2tf32_rz_relu_satfinite", +]; + +const SCALAR_CONVERSION_PTX: [&str; 10] = [ + "cvt.rna.tf32.f32 $0, $1;", + "cvt.rna.satfinite.tf32.f32 $0, $1;", + "cvt.rn.tf32.f32 $0, $1;", + "cvt.rn.relu.tf32.f32 $0, $1;", + "cvt.rn.satfinite.tf32.f32 $0, $1;", + "cvt.rn.relu.satfinite.tf32.f32 $0, $1;", + "cvt.rz.tf32.f32 $0, $1;", + "cvt.rz.relu.tf32.f32 $0, $1;", + "cvt.rz.satfinite.tf32.f32 $0, $1;", + "cvt.rz.relu.satfinite.tf32.f32 $0, $1;", +]; + +fn lower_all_scalar_conversions( + backend: mir_lower::IntrinsicBackend, +) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { + use pliron::builtin::types::FP32Type; + + let mut ctx = make_test_ctx(); + let f32_ty = FP32Type::get(&ctx); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![f32_ty.into()]); + let value = entry.deref(&ctx).get_argument(0); + + let variants = [ + ( + nvvm::ScalarConversionRoundingAttr::NearestAway, + nvvm::ScalarConversionSaturationAttr::None, + ), + ( + nvvm::ScalarConversionRoundingAttr::NearestAway, + nvvm::ScalarConversionSaturationAttr::Satfinite, + ), + ( + nvvm::ScalarConversionRoundingAttr::NearestEven, + nvvm::ScalarConversionSaturationAttr::None, + ), + ( + nvvm::ScalarConversionRoundingAttr::NearestEven, + nvvm::ScalarConversionSaturationAttr::Relu, + ), + ( + nvvm::ScalarConversionRoundingAttr::NearestEven, + nvvm::ScalarConversionSaturationAttr::Satfinite, + ), + ( + nvvm::ScalarConversionRoundingAttr::NearestEven, + nvvm::ScalarConversionSaturationAttr::ReluSatfinite, + ), + ( + nvvm::ScalarConversionRoundingAttr::TowardZero, + nvvm::ScalarConversionSaturationAttr::None, + ), + ( + nvvm::ScalarConversionRoundingAttr::TowardZero, + nvvm::ScalarConversionSaturationAttr::Relu, + ), + ( + nvvm::ScalarConversionRoundingAttr::TowardZero, + nvvm::ScalarConversionSaturationAttr::Satfinite, + ), + ( + nvvm::ScalarConversionRoundingAttr::TowardZero, + nvvm::ScalarConversionSaturationAttr::ReluSatfinite, + ), + ]; + for (rounding, saturation) in variants { + nvvm::ScalarConversionOp::build(&mut ctx, value, rounding, saturation) + .insert_at_back(entry, &ctx); + } + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: backend, + ..Default::default() + }, + ) + .map_err(|error| anyhow::anyhow!("{error}"))?; + Ok((ctx, module_ptr)) +} + +#[test] +fn test_scalar_conversions_llvm_nvptx_use_exact_typed_calls() -> Result<(), anyhow::Error> { + use llvm_export::types as llvm_types; + use pliron::builtin::type_interfaces::FunctionTypeInterface; + use pliron::builtin::types::{FP32Type, IntegerType}; + + let (ctx, module_ptr) = lower_all_scalar_conversions(mir_lower::IntrinsicBackend::LlvmNvptx)?; + let mut calls = Vec::new(); + for op in lowered_kernel_body(&ctx, module_ptr) { + assert!( + Operation::get_op::(op, &ctx).is_none(), + "LLVM-NVPTX scalar conversions must use typed intrinsics" + ); + let Some(call) = Operation::get_op::(op, &ctx) else { + continue; + }; + let CallOpCallable::Direct(callee) = call.callee(&ctx) else { + continue; + }; + let callee = callee.to_string(); + if !SCALAR_CONVERSION_INTRINSICS.contains(&callee.as_str()) { + continue; + } + + let call_op = call.get_operation(); + assert_eq!(call_op.deref(&ctx).get_num_operands(), 1, "{callee}"); + assert_eq!(call_op.deref(&ctx).get_num_results(), 1, "{callee}"); + let block = call_op.deref(&ctx).get_parent_block().unwrap(); + assert_eq!( + call_op.deref(&ctx).get_operand(0), + block.deref(&ctx).get_argument(0), + "{callee} source operand" + ); + + let function_ty = call.callee_type(&ctx); + let function_ty = function_ty.deref(&ctx); + let function_ty = function_ty + .downcast_ref::() + .expect("scalar conversion callee has an LLVM function type"); + assert_eq!(function_ty.arg_types().len(), 1, "{callee}"); + assert!( + function_ty.arg_types()[0] + .deref(&ctx) + .downcast_ref::() + .is_some(), + "{callee} source type" + ); + assert_eq!( + function_ty + .result_type() + .deref(&ctx) + .downcast_ref::() + .expect("scalar conversion result is an integer") + .width(), + 32, + "{callee} result type" + ); + calls.push(callee); + } + calls.sort(); + let mut expected = SCALAR_CONVERSION_INTRINSICS.map(str::to_owned); + expected.sort(); + assert_eq!(calls, expected); + + let module_region = module_ptr.deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + let mut declarations: Vec<_> = module_block + .deref(&ctx) + .iter(&ctx) + .filter_map(|op| Operation::get_op::(op, &ctx)) + .map(|func| func.get_symbol_name(&ctx).to_string()) + .filter(|name| SCALAR_CONVERSION_INTRINSICS.contains(&name.as_str())) + .collect(); + declarations.sort(); + assert_eq!(declarations, expected); + Ok(()) +} + +#[test] +fn test_scalar_conversions_libnvvm_use_exact_pure_inline_ptx() -> Result<(), anyhow::Error> { + use pliron::builtin::types::IntegerType; + use pliron::r#type::Typed; + + let (ctx, module_ptr) = lower_all_scalar_conversions(mir_lower::IntrinsicBackend::LibNvvm)?; + let mut templates = Vec::new(); + for op in lowered_kernel_body(&ctx, module_ptr) { + if let Some(call) = Operation::get_op::(op, &ctx) + && let CallOpCallable::Direct(callee) = call.callee(&ctx) + { + assert!( + !SCALAR_CONVERSION_INTRINSICS.contains(&callee.as_ref()), + "libNVVM scalar conversions must not use typed intrinsics" + ); + } + let Some(inline_asm) = Operation::get_op::(op, &ctx) else { + continue; + }; + let template = inline_asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .unwrap_or_default(); + assert!( + SCALAR_CONVERSION_PTX.contains(&template.as_str()), + "unexpected scalar conversion template `{template}`" + ); + assert_eq!( + inline_asm + .get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some("=r,f"), + "{template}" + ); + assert_eq!(llvm::asm_kind(&ctx, &inline_asm), llvm::AsmKind::Pure); + assert_eq!( + inline_asm + .get_attr_inline_asm_convergent(&ctx) + .map(|value| bool::from((*value).clone())), + Some(false), + "{template}" + ); + + let asm_op = inline_asm.get_operation(); + assert_eq!(asm_op.deref(&ctx).get_num_operands(), 1, "{template}"); + assert_eq!(asm_op.deref(&ctx).get_num_results(), 1, "{template}"); + let block = asm_op.deref(&ctx).get_parent_block().unwrap(); + assert_eq!( + asm_op.deref(&ctx).get_operand(0), + block.deref(&ctx).get_argument(0), + "{template} source operand" + ); + assert_eq!( + asm_op + .deref(&ctx) + .get_result(0) + .get_type(&ctx) + .deref(&ctx) + .downcast_ref::() + .expect("scalar conversion result is an integer") + .width(), + 32, + "{template} result type" + ); + templates.push(template); + } + templates.sort(); + let mut expected = SCALAR_CONVERSION_PTX.map(str::to_owned); + expected.sort(); + assert_eq!(templates, expected); + Ok(()) +} + +#[test] +fn test_scalar_conversion_invalid_variant_fails_closed() { + use pliron::builtin::types::FP32Type; + + let mut ctx = make_test_ctx(); + let f32_ty = FP32Type::get(&ctx); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![f32_ty.into()]); + let value = entry.deref(&ctx).get_argument(0); + nvvm::ScalarConversionOp::build( + &mut ctx, + value, + nvvm::ScalarConversionRoundingAttr::NearestAway, + nvvm::ScalarConversionSaturationAttr::Relu, + ) + .insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + let result = mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: mir_lower::IntrinsicBackend::LlvmNvptx, + ..Default::default() + }, + ); + let error = result.expect_err("unadmitted scalar conversion must not lower"); + assert!( + error.to_string().contains("scalar_conversion"), + "unexpected error: {error}" + ); +} + +fn lower_representative_scalar_arithmetic( + backend: mir_lower::IntrinsicBackend, +) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { + use pliron::builtin::types::FP32Type; + + let mut ctx = make_test_ctx(); + let f32_ty = FP32Type::get(&ctx); + let (module_ptr, entry) = + build_test_kernel(&mut ctx, vec![f32_ty.into(), f32_ty.into(), f32_ty.into()]); + let args: Vec<_> = (0..3) + .map(|index| entry.deref(&ctx).get_argument(index)) + .collect(); + + for (operation, saturation, operands) in [ + ( + nvvm::ScalarArithmeticOperationAttr::Add, + nvvm::ScalarArithmeticSaturationAttr::None, + vec![args[0], args[1]], + ), + ( + nvvm::ScalarArithmeticOperationAttr::Add, + nvvm::ScalarArithmeticSaturationAttr::Sat, + vec![args[0], args[1]], + ), + ( + nvvm::ScalarArithmeticOperationAttr::Fma, + nvvm::ScalarArithmeticSaturationAttr::None, + args.clone(), + ), + ( + nvvm::ScalarArithmeticOperationAttr::Fma, + nvvm::ScalarArithmeticSaturationAttr::Sat, + args, + ), + ] { + nvvm::ScalarArithmeticOp::build( + &mut ctx, + operands, + nvvm::ScalarArithmeticFormatAttr::F32, + operation, + nvvm::ScalarArithmeticRoundingAttr::Rn, + nvvm::ScalarArithmeticSubnormalAttr::Preserve, + saturation, + ) + .insert_at_back(entry, &ctx); + } + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: backend, + ..Default::default() + }, + ) + .map_err(|error| anyhow::anyhow!("{error}"))?; + Ok((ctx, module_ptr)) +} + +fn is_representative_scalar_arithmetic_intrinsic(name: &str) -> bool { + name.starts_with("llvm_nvvm_add_rn") || name.starts_with("llvm_nvvm_fma_rn") +} + +#[test] +fn test_scalar_arithmetic_llvm_uses_inline_ptx_only_for_saturation() -> Result<(), anyhow::Error> { + let (ctx, module_ptr) = + lower_representative_scalar_arithmetic(mir_lower::IntrinsicBackend::LlvmNvptx)?; + let mut calls = Vec::new(); + let mut inline_ptx = Vec::new(); + for op in lowered_kernel_body(&ctx, module_ptr) { + if let Some(call) = Operation::get_op::(op, &ctx) + && let CallOpCallable::Direct(callee) = call.callee(&ctx) + && is_representative_scalar_arithmetic_intrinsic(callee.as_ref()) + { + calls.push(callee.to_string()); + } + let Some(inline_asm) = Operation::get_op::(op, &ctx) else { + continue; + }; + let template = inline_asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .unwrap_or_default(); + let constraints = inline_asm + .get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .unwrap_or_default(); + assert_eq!(llvm::asm_kind(&ctx, &inline_asm), llvm::AsmKind::Pure); + assert_eq!( + inline_asm + .get_attr_inline_asm_convergent(&ctx) + .map(|value| bool::from((*value).clone())), + Some(false) + ); + inline_ptx.push((template, constraints)); + } + + calls.sort(); + assert_eq!(calls, ["llvm_nvvm_add_rn_f", "llvm_nvvm_fma_rn_f"]); + inline_ptx.sort(); + assert_eq!( + inline_ptx, + [ + ("add.rn.sat.f32 $0, $1, $2;".into(), "=f,f,f".into()), + ("fma.rn.sat.f32 $0, $1, $2, $3;".into(), "=f,f,f,f".into(),), + ] + ); + + let module_region = module_ptr.deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + let mut declarations: Vec<_> = module_block + .deref(&ctx) + .iter(&ctx) + .filter_map(|op| Operation::get_op::(op, &ctx)) + .map(|func| func.get_symbol_name(&ctx).to_string()) + .filter(|name| is_representative_scalar_arithmetic_intrinsic(name)) + .collect(); + declarations.sort(); + assert_eq!(declarations, calls); + Ok(()) +} + +#[test] +fn test_scalar_arithmetic_libnvvm_uses_exact_inline_ptx() -> Result<(), anyhow::Error> { + let (ctx, module_ptr) = + lower_representative_scalar_arithmetic(mir_lower::IntrinsicBackend::LibNvvm)?; + let mut calls = Vec::new(); + let mut inline_ptx = Vec::new(); + for op in lowered_kernel_body(&ctx, module_ptr) { + if let Some(call) = Operation::get_op::(op, &ctx) + && let CallOpCallable::Direct(callee) = call.callee(&ctx) + && is_representative_scalar_arithmetic_intrinsic(callee.as_ref()) + { + calls.push(callee.to_string()); + } + let Some(inline_asm) = Operation::get_op::(op, &ctx) else { + continue; + }; + let template = inline_asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .unwrap_or_default(); + let constraints = inline_asm + .get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .unwrap_or_default(); + assert_eq!(llvm::asm_kind(&ctx, &inline_asm), llvm::AsmKind::Pure); + inline_ptx.push((template, constraints)); + } + + assert!(calls.is_empty()); + inline_ptx.sort(); + assert_eq!( + inline_ptx, + [ + ("add.rn.f32 $0, $1, $2;".into(), "=f,f,f".into()), + ("add.rn.sat.f32 $0, $1, $2;".into(), "=f,f,f".into()), + ("fma.rn.f32 $0, $1, $2, $3;".into(), "=f,f,f,f".into(),), + ("fma.rn.sat.f32 $0, $1, $2, $3;".into(), "=f,f,f,f".into(),), + ] + ); + Ok(()) +} diff --git a/crates/mir-lower/tests/lowering_test/matrix_memory.rs b/crates/mir-lower/tests/lowering_test/matrix_memory.rs new file mode 100644 index 0000000000..bd05a6da9b --- /dev/null +++ b/crates/mir-lower/tests/lowering_test/matrix_memory.rs @@ -0,0 +1,1214 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use dialect_nvvm::ops as nvvm; +use llvm_export::ops as llvm; +use pliron::builtin::op_interfaces::{CallOpCallable, CallOpInterface, SymbolOpInterface}; +use pliron::builtin::ops::ModuleOp; +use pliron::context::Context; +use pliron::linked_list::ContainsLinkedList; +use pliron::op::Op; +use pliron::operation::Operation; + +use crate::common::{append_return, build_test_kernel, lowered_kernel_body, make_test_ctx}; + +const STMATRIX_TYPED_INTRINSICS: [&str; 4] = [ + "llvm_nvvm_stmatrix_sync_aligned_m8n8_x2_b16_p3", + "llvm_nvvm_stmatrix_sync_aligned_m8n8_x2_trans_b16_p3", + "llvm_nvvm_stmatrix_sync_aligned_m8n8_x4_b16_p3", + "llvm_nvvm_stmatrix_sync_aligned_m8n8_x4_trans_b16_p3", +]; + +const STMATRIX_PTX: [(&str, &str); 4] = [ + ( + "{ .reg .u64 %ptr64; .reg .u32 %ptr32; cvta.to.shared.u64 %ptr64, $0; cvt.u32.u64 %ptr32, %ptr64; stmatrix.sync.aligned.m8n8.x2.shared.b16 [%ptr32], {$1, $2}; }", + "l,r,r,~{memory}", + ), + ( + "{ .reg .u64 %ptr64; .reg .u32 %ptr32; cvta.to.shared.u64 %ptr64, $0; cvt.u32.u64 %ptr32, %ptr64; stmatrix.sync.aligned.m8n8.x2.trans.shared.b16 [%ptr32], {$1, $2}; }", + "l,r,r,~{memory}", + ), + ( + "{ .reg .u64 %ptr64; .reg .u32 %ptr32; cvta.to.shared.u64 %ptr64, $0; cvt.u32.u64 %ptr32, %ptr64; stmatrix.sync.aligned.m8n8.x4.shared.b16 [%ptr32], {$1, $2, $3, $4}; }", + "l,r,r,r,r,~{memory}", + ), + ( + "{ .reg .u64 %ptr64; .reg .u32 %ptr32; cvta.to.shared.u64 %ptr64, $0; cvt.u32.u64 %ptr32, %ptr64; stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 [%ptr32], {$1, $2, $3, $4}; }", + "l,r,r,r,r,~{memory}", + ), +]; + +fn lower_all_stmatrix_forms( + backend: mir_lower::IntrinsicBackend, +) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { + use dialect_mir::types::MirPtrType; + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let i8_ty = IntegerType::get(&ctx, 8, Signedness::Signless); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let ptr_ty = MirPtrType::get_generic(&mut ctx, i8_ty.into(), true); + let (module_ptr, entry) = build_test_kernel( + &mut ctx, + vec![ + ptr_ty.into(), + i32_ty.into(), + i32_ty.into(), + i32_ty.into(), + i32_ty.into(), + ], + ); + let args: Vec<_> = (0..5) + .map(|index| entry.deref(&ctx).get_argument(index)) + .collect(); + + for (op_info, operands) in [ + ( + nvvm::StmatrixM8n8X2Op::get_concrete_op_info(), + args[..3].to_vec(), + ), + ( + nvvm::StmatrixM8n8X2TransOp::get_concrete_op_info(), + args[..3].to_vec(), + ), + (nvvm::StmatrixM8n8X4Op::get_concrete_op_info(), args.clone()), + (nvvm::StmatrixM8n8X4TransOp::get_concrete_op_info(), args), + ] { + Operation::new(&mut ctx, op_info, vec![], operands, vec![], 0).insert_at_back(entry, &ctx); + } + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: backend, + ..Default::default() + }, + ) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + Ok((ctx, module_ptr)) +} + +#[test] +fn test_stmatrix_llvm_nvptx_uses_exact_typed_p3_intrinsics() -> Result<(), anyhow::Error> { + use llvm_export::types as llvm_types; + use pliron::r#type::Typed; + + let (ctx, module_ptr) = lower_all_stmatrix_forms(mir_lower::IntrinsicBackend::LlvmNvptx)?; + let mut callees = Vec::new(); + + for op in lowered_kernel_body(&ctx, module_ptr) { + assert!( + Operation::get_op::(op, &ctx).is_none(), + "LLVM-NVPTX stmatrix lowering must not emit inline PTX" + ); + let Some(call) = Operation::get_op::(op, &ctx) else { + continue; + }; + let CallOpCallable::Direct(callee) = call.callee(&ctx) else { + continue; + }; + let callee = callee.to_string(); + if !STMATRIX_TYPED_INTRINSICS.contains(&callee.as_str()) { + continue; + } + + let call_op = call.get_operation().deref(&ctx); + assert!(matches!(call_op.get_num_operands(), 3 | 5)); + assert_eq!(call_op.get_num_results(), 1); + let pointer_ty = call_op.get_operand(0).get_type(&ctx); + let pointer_ty = pointer_ty.deref(&ctx); + let pointer_ty = pointer_ty + .downcast_ref::() + .expect("stmatrix first argument is a pointer"); + assert_eq!(pointer_ty.address_space(), 3); + callees.push(callee); + } + + callees.sort(); + let mut expected = STMATRIX_TYPED_INTRINSICS.map(str::to_owned); + expected.sort(); + assert_eq!(callees, expected); + + let module = Operation::get_op::(module_ptr, &ctx).unwrap(); + let ir = llvm_export::export::export_module_to_string(&ctx, &module) + .expect("typed stmatrix module exports to LLVM IR"); + for intrinsic in STMATRIX_TYPED_INTRINSICS { + let dotted = intrinsic.replace('_', "."); + assert!( + ir.contains(&format!("@{dotted}(ptr addrspace(3)")), + "missing exact typed stmatrix declaration:\n{ir}" + ); + } + assert!(!ir.contains("asm sideeffect"), "{ir}"); + Ok(()) +} + +#[test] +fn test_stmatrix_libnvvm_uses_exact_convergent_memory_asm() -> Result<(), anyhow::Error> { + let (ctx, module_ptr) = lower_all_stmatrix_forms(mir_lower::IntrinsicBackend::LibNvvm)?; + let mut lowered = Vec::new(); + + for op in lowered_kernel_body(&ctx, module_ptr) { + if let Some(call) = Operation::get_op::(op, &ctx) + && let CallOpCallable::Direct(callee) = call.callee(&ctx) + { + assert!( + !STMATRIX_TYPED_INTRINSICS.contains(&callee.as_ref()), + "libNVVM stmatrix lowering must not emit typed intrinsic calls" + ); + } + let Some(inline_asm) = Operation::get_op::(op, &ctx) else { + continue; + }; + lowered.push(( + inline_asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .unwrap_or_default(), + inline_asm + .get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .unwrap_or_default(), + llvm::asm_kind(&ctx, &inline_asm), + op.deref(&ctx).get_num_operands(), + op.deref(&ctx).get_num_results(), + )); + } + + assert_eq!(lowered.len(), STMATRIX_PTX.len()); + for (template, constraints) in STMATRIX_PTX { + let matches: Vec<_> = lowered + .iter() + .filter(|(actual, _, _, _, _)| actual == template) + .collect(); + assert_eq!(matches.len(), 1, "missing exact PTX {template}"); + let (_, actual_constraints, kind, operands, results) = matches[0]; + assert_eq!(actual_constraints, constraints); + assert_eq!(*kind, llvm::AsmKind::Convergent); + assert!(matches!(*operands, 3 | 5)); + assert_eq!(*results, 1); + } + + let module = Operation::get_op::(module_ptr, &ctx).unwrap(); + let ir = llvm_export::export::export_module_to_string(&ctx, &module) + .expect("inline stmatrix module exports to LLVM IR"); + assert_eq!(ir.matches("asm sideeffect").count(), 4, "{ir}"); + assert_eq!(ir.matches("~{memory}").count(), 4, "{ir}"); + assert!(ir.contains("attributes #0 = { convergent }"), "{ir}"); + assert!(!ir.contains("@llvm.nvvm.stmatrix"), "{ir}"); + Ok(()) +} + +// ============================================================================= +// Warp-level matrix (`movmatrix`) lowering test +// ============================================================================= + +#[test] +fn test_movmatrix_trans_b16_lowers_to_inline_asm() -> Result<(), anyhow::Error> { + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![i32_ty.into()]); + + let a_val = entry.deref(&ctx).get_argument(0); + + let op = Operation::new( + &mut ctx, + nvvm::MovmatrixTransB16Op::get_concrete_op_info(), + vec![i32_ty.into()], + vec![a_val], + vec![], + 0, + ); + op.insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let mut found = 0; + let module_region = module_ptr.deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + for module_op in module_block.deref(&ctx).iter(&ctx) { + let Some(function) = Operation::get_op::(module_op, &ctx) else { + continue; + }; + if function.get_symbol_name(&ctx).to_string() != "kernel_func" { + continue; + } + let body = function.get_operation().deref(&ctx).get_region(0); + for block in body.deref(&ctx).iter(&ctx) { + for body_op in block.deref(&ctx).iter(&ctx) { + let Some(asm) = Operation::get_op::(body_op, &ctx) else { + continue; + }; + found += 1; + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())); + let constraints = asm + .get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())); + assert_eq!( + template.as_deref(), + Some("movmatrix.sync.aligned.m8n8.trans.b16 $0, $1;") + ); + assert_eq!(constraints.as_deref(), Some("=r,r")); + assert_eq!( + llvm::asm_kind_opt(&ctx, &asm), + Some(llvm::AsmKind::Convergent) + ); + assert!( + asm.get_attr_inline_asm_convergent(&ctx) + .is_some_and(|value| bool::from((*value).clone())) + ); + assert!( + !constraints.as_deref().unwrap().contains("memory"), + "register-only movmatrix must not claim a memory clobber" + ); + } + } + } + + assert_eq!(found, 1, "expected one movmatrix inline-asm operation"); + Ok(()) +} + +// ============================================================================= +// ldmatrix lowering tests +// ============================================================================= + +const LDMATRIX_TYPED_INTRINSICS: [&str; 6] = [ + "llvm_nvvm_ldmatrix_sync_aligned_m8n8_x1_b16_p3", + "llvm_nvvm_ldmatrix_sync_aligned_m8n8_x1_trans_b16_p3", + "llvm_nvvm_ldmatrix_sync_aligned_m8n8_x2_b16_p3", + "llvm_nvvm_ldmatrix_sync_aligned_m8n8_x2_trans_b16_p3", + "llvm_nvvm_ldmatrix_sync_aligned_m8n8_x4_b16_p3", + "llvm_nvvm_ldmatrix_sync_aligned_m8n8_x4_trans_b16_p3", +]; + +const LDMATRIX_PTX_TEMPLATES: [&str; 6] = [ + "ldmatrix.sync.aligned.m8n8.x1.shared.b16 {$0}, [$1];", + "ldmatrix.sync.aligned.m8n8.x1.trans.shared.b16 {$0}, [$1];", + "ldmatrix.sync.aligned.m8n8.x2.shared.b16 {$0, $1}, [$2];", + "ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 {$0, $1}, [$2];", + "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {$0, $1, $2, $3}, [$4];", + "ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {$0, $1, $2, $3}, [$4];", +]; + +const LDMATRIX_PTX_CONSTRAINTS: [&str; 6] = [ + "=r,r,~{memory}", + "=r,r,~{memory}", + "=r,=r,r,~{memory}", + "=r,=r,r,~{memory}", + "=r,=r,=r,=r,r,~{memory}", + "=r,=r,=r,=r,r,~{memory}", +]; + +const BLACKWELL_LDMATRIX_CASES: [(&str, &str, &str, usize); 12] = [ + ( + "llvm_nvvm_ldmatrix_sync_aligned_m16n16_x1_trans_b8_p3", + "llvm.nvvm.ldmatrix.sync.aligned.m16n16.x1.trans.b8.p3", + "ldmatrix.sync.aligned.m16n16.x1.trans.shared.b8 {$0, $1}, [$2];", + 2, + ), + ( + "llvm__nvvm_dldmatrix_dsync_daligned_dm16n16_dx1_dtrans_db8x16_db4x16_up64_dp3", + "llvm.nvvm.ldmatrix.sync.aligned.m16n16.x1.trans.b8x16.b4x16_p64.p3", + "ldmatrix.sync.aligned.m16n16.x1.trans.shared.b8x16.b4x16_p64 {$0, $1}, [$2];", + 2, + ), + ( + "llvm__nvvm_dldmatrix_dsync_daligned_dm16n16_dx1_dtrans_db8x16_db6x16_up32_dp3", + "llvm.nvvm.ldmatrix.sync.aligned.m16n16.x1.trans.b8x16.b6x16_p32.p3", + "ldmatrix.sync.aligned.m16n16.x1.trans.shared.b8x16.b6x16_p32 {$0, $1}, [$2];", + 2, + ), + ( + "llvm_nvvm_ldmatrix_sync_aligned_m16n16_x2_trans_b8_p3", + "llvm.nvvm.ldmatrix.sync.aligned.m16n16.x2.trans.b8.p3", + "ldmatrix.sync.aligned.m16n16.x2.trans.shared.b8 {$0, $1, $2, $3}, [$4];", + 4, + ), + ( + "llvm__nvvm_dldmatrix_dsync_daligned_dm16n16_dx2_dtrans_db8x16_db4x16_up64_dp3", + "llvm.nvvm.ldmatrix.sync.aligned.m16n16.x2.trans.b8x16.b4x16_p64.p3", + "ldmatrix.sync.aligned.m16n16.x2.trans.shared.b8x16.b4x16_p64 {$0, $1, $2, $3}, [$4];", + 4, + ), + ( + "llvm__nvvm_dldmatrix_dsync_daligned_dm16n16_dx2_dtrans_db8x16_db6x16_up32_dp3", + "llvm.nvvm.ldmatrix.sync.aligned.m16n16.x2.trans.b8x16.b6x16_p32.p3", + "ldmatrix.sync.aligned.m16n16.x2.trans.shared.b8x16.b6x16_p32 {$0, $1, $2, $3}, [$4];", + 4, + ), + ( + "llvm__nvvm_dldmatrix_dsync_daligned_dm8n16_dx1_db8x16_db4x16_up64_dp3", + "llvm.nvvm.ldmatrix.sync.aligned.m8n16.x1.b8x16.b4x16_p64.p3", + "ldmatrix.sync.aligned.m8n16.x1.shared.b8x16.b4x16_p64 {$0}, [$1];", + 1, + ), + ( + "llvm__nvvm_dldmatrix_dsync_daligned_dm8n16_dx1_db8x16_db6x16_up32_dp3", + "llvm.nvvm.ldmatrix.sync.aligned.m8n16.x1.b8x16.b6x16_p32.p3", + "ldmatrix.sync.aligned.m8n16.x1.shared.b8x16.b6x16_p32 {$0}, [$1];", + 1, + ), + ( + "llvm__nvvm_dldmatrix_dsync_daligned_dm8n16_dx2_db8x16_db4x16_up64_dp3", + "llvm.nvvm.ldmatrix.sync.aligned.m8n16.x2.b8x16.b4x16_p64.p3", + "ldmatrix.sync.aligned.m8n16.x2.shared.b8x16.b4x16_p64 {$0, $1}, [$2];", + 2, + ), + ( + "llvm__nvvm_dldmatrix_dsync_daligned_dm8n16_dx2_db8x16_db6x16_up32_dp3", + "llvm.nvvm.ldmatrix.sync.aligned.m8n16.x2.b8x16.b6x16_p32.p3", + "ldmatrix.sync.aligned.m8n16.x2.shared.b8x16.b6x16_p32 {$0, $1}, [$2];", + 2, + ), + ( + "llvm__nvvm_dldmatrix_dsync_daligned_dm8n16_dx4_db8x16_db4x16_up64_dp3", + "llvm.nvvm.ldmatrix.sync.aligned.m8n16.x4.b8x16.b4x16_p64.p3", + "ldmatrix.sync.aligned.m8n16.x4.shared.b8x16.b4x16_p64 {$0, $1, $2, $3}, [$4];", + 4, + ), + ( + "llvm__nvvm_dldmatrix_dsync_daligned_dm8n16_dx4_db8x16_db6x16_up32_dp3", + "llvm.nvvm.ldmatrix.sync.aligned.m8n16.x4.b8x16.b6x16_p32.p3", + "ldmatrix.sync.aligned.m8n16.x4.shared.b8x16.b6x16_p32 {$0, $1, $2, $3}, [$4];", + 4, + ), +]; + +fn ldmatrix_constraints(register_count: usize) -> &'static str { + match register_count { + 1 => "=r,r,~{memory}", + 2 => "=r,=r,r,~{memory}", + 4 => "=r,=r,=r,=r,r,~{memory}", + _ => unreachable!("closed ldmatrix register count"), + } +} + +fn lower_all_ldmatrix_forms( + address_space: u32, + backend: mir_lower::IntrinsicBackend, + compatibility: bool, +) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { + use dialect_mir::types::MirPtrType; + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let ptr_ty = MirPtrType::get(&mut ctx, u32_ty.into(), true, address_space); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![ptr_ty.into()]); + let pointer = entry.deref(&ctx).get_argument(0); + if compatibility { + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + for (op_info, result_count) in [ + (nvvm::LdmatrixX1Op::get_concrete_op_info(), 1), + (nvvm::LdmatrixX1TransOp::get_concrete_op_info(), 1), + (nvvm::LdmatrixX2Op::get_concrete_op_info(), 2), + (nvvm::LdmatrixX2TransOp::get_concrete_op_info(), 2), + (nvvm::LdmatrixX4Op::get_concrete_op_info(), 4), + (nvvm::LdmatrixX4TransOp::get_concrete_op_info(), 4), + ] { + Operation::new( + &mut ctx, + op_info, + vec![u32_ty.into(); result_count], + vec![pointer], + vec![], + 0, + ) + .insert_at_back(entry, &ctx); + } + } else { + for (multiplicity, layout) in [ + ( + nvvm::LdmatrixMultiplicityAttr::X1, + nvvm::LdmatrixLayoutAttr::Normal, + ), + ( + nvvm::LdmatrixMultiplicityAttr::X1, + nvvm::LdmatrixLayoutAttr::Transposed, + ), + ( + nvvm::LdmatrixMultiplicityAttr::X2, + nvvm::LdmatrixLayoutAttr::Normal, + ), + ( + nvvm::LdmatrixMultiplicityAttr::X2, + nvvm::LdmatrixLayoutAttr::Transposed, + ), + ( + nvvm::LdmatrixMultiplicityAttr::X4, + nvvm::LdmatrixLayoutAttr::Normal, + ), + ( + nvvm::LdmatrixMultiplicityAttr::X4, + nvvm::LdmatrixLayoutAttr::Transposed, + ), + ] { + nvvm::LdmatrixOp::build( + &mut ctx, + pointer, + nvvm::LdmatrixShapeAttr::M8n8, + multiplicity, + layout, + nvvm::LdmatrixElementAttr::B16, + nvvm::LdmatrixStateSpaceAttr::Shared, + ) + .insert_at_back(entry, &ctx); + } + } + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: backend, + ..Default::default() + }, + ) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + Ok((ctx, module_ptr)) +} + +fn lower_all_blackwell_ldmatrix_forms( + address_space: u32, + backend: mir_lower::IntrinsicBackend, +) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { + use dialect_mir::types::MirPtrType; + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let u8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); + let ptr_ty = MirPtrType::get(&mut ctx, u8_ty.into(), true, address_space); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![ptr_ty.into()]); + let pointer = entry.deref(&ctx).get_argument(0); + + for (shape, multiplicity, layout, element) in [ + ( + nvvm::LdmatrixShapeAttr::M16n16, + nvvm::LdmatrixMultiplicityAttr::X1, + nvvm::LdmatrixLayoutAttr::Transposed, + nvvm::LdmatrixElementAttr::B8, + ), + ( + nvvm::LdmatrixShapeAttr::M16n16, + nvvm::LdmatrixMultiplicityAttr::X1, + nvvm::LdmatrixLayoutAttr::Transposed, + nvvm::LdmatrixElementAttr::B8x16B4x16P64, + ), + ( + nvvm::LdmatrixShapeAttr::M16n16, + nvvm::LdmatrixMultiplicityAttr::X1, + nvvm::LdmatrixLayoutAttr::Transposed, + nvvm::LdmatrixElementAttr::B8x16B6x16P32, + ), + ( + nvvm::LdmatrixShapeAttr::M16n16, + nvvm::LdmatrixMultiplicityAttr::X2, + nvvm::LdmatrixLayoutAttr::Transposed, + nvvm::LdmatrixElementAttr::B8, + ), + ( + nvvm::LdmatrixShapeAttr::M16n16, + nvvm::LdmatrixMultiplicityAttr::X2, + nvvm::LdmatrixLayoutAttr::Transposed, + nvvm::LdmatrixElementAttr::B8x16B4x16P64, + ), + ( + nvvm::LdmatrixShapeAttr::M16n16, + nvvm::LdmatrixMultiplicityAttr::X2, + nvvm::LdmatrixLayoutAttr::Transposed, + nvvm::LdmatrixElementAttr::B8x16B6x16P32, + ), + ( + nvvm::LdmatrixShapeAttr::M8n16, + nvvm::LdmatrixMultiplicityAttr::X1, + nvvm::LdmatrixLayoutAttr::Normal, + nvvm::LdmatrixElementAttr::B8x16B4x16P64, + ), + ( + nvvm::LdmatrixShapeAttr::M8n16, + nvvm::LdmatrixMultiplicityAttr::X1, + nvvm::LdmatrixLayoutAttr::Normal, + nvvm::LdmatrixElementAttr::B8x16B6x16P32, + ), + ( + nvvm::LdmatrixShapeAttr::M8n16, + nvvm::LdmatrixMultiplicityAttr::X2, + nvvm::LdmatrixLayoutAttr::Normal, + nvvm::LdmatrixElementAttr::B8x16B4x16P64, + ), + ( + nvvm::LdmatrixShapeAttr::M8n16, + nvvm::LdmatrixMultiplicityAttr::X2, + nvvm::LdmatrixLayoutAttr::Normal, + nvvm::LdmatrixElementAttr::B8x16B6x16P32, + ), + ( + nvvm::LdmatrixShapeAttr::M8n16, + nvvm::LdmatrixMultiplicityAttr::X4, + nvvm::LdmatrixLayoutAttr::Normal, + nvvm::LdmatrixElementAttr::B8x16B4x16P64, + ), + ( + nvvm::LdmatrixShapeAttr::M8n16, + nvvm::LdmatrixMultiplicityAttr::X4, + nvvm::LdmatrixLayoutAttr::Normal, + nvvm::LdmatrixElementAttr::B8x16B6x16P32, + ), + ] { + nvvm::LdmatrixOp::build( + &mut ctx, + pointer, + shape, + multiplicity, + layout, + element, + nvvm::LdmatrixStateSpaceAttr::Shared, + ) + .insert_at_back(entry, &ctx); + } + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: backend, + ..Default::default() + }, + ) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + Ok((ctx, module_ptr)) +} + +fn assert_ldmatrix_producer_result_shape( + ctx: &Context, + op: pliron::context::Ptr, + register_count: usize, +) { + use llvm_export::types as llvm_types; + use pliron::builtin::types::IntegerType; + use pliron::r#type::Typed; + + let ty = op.deref(ctx).get_result(0).get_type(ctx); + let ty = ty.deref(ctx); + if register_count == 1 { + assert_eq!( + ty.downcast_ref::() + .expect("single-register ldmatrix returns i32") + .width(), + 32 + ); + } else { + let ty = ty + .downcast_ref::() + .expect("multi-register ldmatrix returns an LLVM struct"); + assert_eq!(ty.num_fields(), register_count); + for index in 0..register_count { + assert_eq!( + ty.field_type(index) + .deref(ctx) + .downcast_ref::() + .expect("ldmatrix fragment field is i32") + .width(), + 32 + ); + } + } +} + +#[test] +fn test_ldmatrix_llvm_nvptx_uses_exact_typed_p3_intrinsics() -> Result<(), anyhow::Error> { + use llvm_export::types as llvm_types; + use pliron::builtin::type_interfaces::FunctionTypeInterface; + use pliron::builtin::types::IntegerType; + use pliron::r#type::Typed; + + for address_space in [0, 3] { + let (ctx, module_ptr) = + lower_all_ldmatrix_forms(address_space, mir_lower::IntrinsicBackend::LlvmNvptx, false)?; + let body = lowered_kernel_body(&ctx, module_ptr); + let mut callees = Vec::new(); + let mut cast_count = 0; + let mut extract_count = 0; + + for op in body { + if let Some(call) = Operation::get_op::(op, &ctx) + && let CallOpCallable::Direct(callee) = call.callee(&ctx) + { + let callee = callee.to_string(); + let register_count = if callee.contains("_x1_") { + 1 + } else if callee.contains("_x2_") { + 2 + } else { + 4 + }; + let function_ty = call.callee_type(&ctx); + let function_ty = function_ty.deref(&ctx); + let function_ty = function_ty + .downcast_ref::() + .expect("ldmatrix callee has an LLVM function type"); + assert_eq!(function_ty.arg_types().len(), 1); + let argument_ty = function_ty.arg_types()[0].deref(&ctx); + let argument_ty = argument_ty + .downcast_ref::() + .expect("ldmatrix argument is a pointer"); + assert_eq!(argument_ty.address_space(), 3); + + let result_ty = function_ty.result_type(); + let result_ty = result_ty.deref(&ctx); + if register_count == 1 { + let result_ty = result_ty + .downcast_ref::() + .expect("x1 returns i32"); + assert_eq!(result_ty.width(), 32); + } else { + let result_ty = result_ty + .downcast_ref::() + .expect("x2/x4 return an LLVM struct"); + assert_eq!(result_ty.num_fields(), register_count); + for index in 0..result_ty.num_fields() { + let field = result_ty.field_type(index); + let field = field.deref(&ctx); + let field = field + .downcast_ref::() + .expect("fragment register is i32"); + assert_eq!(field.width(), 32); + } + } + + callees.push(callee); + assert_eq!(op.deref(&ctx).get_num_operands(), 1); + assert_eq!(op.deref(&ctx).get_num_results(), 1); + } + if Operation::get_op::(op, &ctx).is_some() { + cast_count += 1; + let cast = op.deref(&ctx); + let source_ty = cast.get_operand(0).get_type(&ctx); + let source_ty = source_ty.deref(&ctx); + let source_ty = source_ty + .downcast_ref::() + .expect("addrspacecast source is a pointer"); + let result_ty = cast.get_result(0).get_type(&ctx); + let result_ty = result_ty.deref(&ctx); + let result_ty = result_ty + .downcast_ref::() + .expect("addrspacecast result is a pointer"); + assert_eq!( + (source_ty.address_space(), result_ty.address_space()), + (0, 3) + ); + } + extract_count += + usize::from(Operation::get_op::(op, &ctx).is_some()); + assert!( + Operation::get_op::(op, &ctx).is_none(), + "LLVM-NVPTX ldmatrix lowering must not emit inline PTX" + ); + assert!( + Operation::get_op::(op, &ctx).is_none(), + "the typed intrinsic consumes the shared pointer directly" + ); + } + + callees.sort(); + let mut expected = LDMATRIX_TYPED_INTRINSICS.map(str::to_owned); + expected.sort(); + assert_eq!(callees, expected); + assert_eq!(cast_count, if address_space == 0 { 6 } else { 0 }); + assert_eq!( + extract_count, 12, + "x2/x4 structs must preserve result order" + ); + + let module = Operation::get_op::(module_ptr, &ctx).unwrap(); + let ir = llvm_export::export::export_module_to_string(&ctx, &module) + .expect("typed ldmatrix module exports to LLVM IR"); + for intrinsic in LDMATRIX_TYPED_INTRINSICS { + let dotted = intrinsic.replace('_', "."); + assert!( + ir.contains(&format!("@{dotted}(ptr addrspace(3)")), + "underscore symbol must export as exact dotted p3 intrinsic:\n{ir}" + ); + } + assert!(!ir.contains("@llvm_nvvm_ldmatrix")); + } + Ok(()) +} + +#[test] +fn test_ldmatrix_libnvvm_uses_exact_convergent_shared_ptx() -> Result<(), anyhow::Error> { + use llvm_export::types as llvm_types; + use pliron::r#type::Typed; + + for address_space in [0, 3] { + let (ctx, module_ptr) = + lower_all_ldmatrix_forms(address_space, mir_lower::IntrinsicBackend::LibNvvm, false)?; + let body = lowered_kernel_body(&ctx, module_ptr); + let mut lowered = Vec::new(); + let mut cast_count = 0; + let mut ptrtoint_count = 0; + + for op in body { + assert!( + Operation::get_op::(op, &ctx).is_none(), + "libNVVM ldmatrix lowering must not emit typed intrinsic calls" + ); + cast_count += + usize::from(Operation::get_op::(op, &ctx).is_some()); + if Operation::get_op::(op, &ctx).is_some() { + ptrtoint_count += 1; + let cast = op.deref(&ctx); + let source_ty = cast.get_operand(0).get_type(&ctx); + let source_ty = source_ty.deref(&ctx); + let source_ty = source_ty + .downcast_ref::() + .expect("ptrtoint source is a pointer"); + assert_eq!(source_ty.address_space(), 3); + } + + let Some(inline_asm) = Operation::get_op::(op, &ctx) else { + continue; + }; + lowered.push(( + inline_asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .unwrap_or_default(), + inline_asm + .get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .unwrap_or_default(), + llvm::asm_kind(&ctx, &inline_asm), + op.deref(&ctx).get_num_operands(), + op.deref(&ctx).get_num_results(), + )); + } + + assert_eq!(lowered.len(), 6); + for (index, expected_template) in LDMATRIX_PTX_TEMPLATES.iter().enumerate() { + let matching: Vec<_> = lowered + .iter() + .filter(|(template, _, _, _, _)| template == expected_template) + .collect(); + assert_eq!(matching.len(), 1, "missing exact PTX `{expected_template}`"); + let (_, constraints, kind, operands, results) = matching[0]; + assert_eq!(constraints, LDMATRIX_PTX_CONSTRAINTS[index]); + assert_eq!(*kind, llvm::AsmKind::Convergent); + assert_eq!(*operands, 1, "inline PTX consumes one i32 shared address"); + assert_eq!(*results, 1, "inline PTX returns one scalar or struct"); + assert!(!expected_template.contains("cvta.to.shared")); + } + assert_eq!(cast_count, if address_space == 0 { 6 } else { 0 }); + assert_eq!(ptrtoint_count, 6); + + let module = Operation::get_op::(module_ptr, &ctx).unwrap(); + let ir = llvm_export::export::export_module_to_string(&ctx, &module) + .expect("inline ldmatrix module exports to LLVM IR"); + assert_eq!(ir.matches("asm sideeffect").count(), 6, "{ir}"); + assert_eq!(ir.matches("~{memory}").count(), 6, "{ir}"); + assert!(ir.contains("attributes #0 = { convergent }"), "{ir}"); + assert!(!ir.contains("@llvm.nvvm.ldmatrix"), "{ir}"); + } + Ok(()) +} + +#[test] +fn test_blackwell_ldmatrix_llvm_uses_all_exact_lossless_p3_intrinsics() -> Result<(), anyhow::Error> +{ + use llvm_export::types as llvm_types; + use pliron::builtin::type_interfaces::FunctionTypeInterface; + + for address_space in [0, 3] { + let (ctx, module_ptr) = lower_all_blackwell_ldmatrix_forms( + address_space, + mir_lower::IntrinsicBackend::LlvmNvptx, + )?; + let mut seen = [0; 12]; + let mut cast_count = 0; + let mut extract_count = 0; + + for op in lowered_kernel_body(&ctx, module_ptr) { + assert!(Operation::get_op::(op, &ctx).is_none()); + assert!(Operation::get_op::(op, &ctx).is_none()); + cast_count += + usize::from(Operation::get_op::(op, &ctx).is_some()); + extract_count += + usize::from(Operation::get_op::(op, &ctx).is_some()); + + let Some(call) = Operation::get_op::(op, &ctx) else { + continue; + }; + let CallOpCallable::Direct(callee) = call.callee(&ctx) else { + panic!("Blackwell ldmatrix intrinsic call must be direct"); + }; + let callee = callee.to_string(); + let index = BLACKWELL_LDMATRIX_CASES + .iter() + .position(|(identifier, _, _, _)| *identifier == callee) + .expect("exact lossless Blackwell ldmatrix identifier"); + seen[index] += 1; + let register_count = BLACKWELL_LDMATRIX_CASES[index].3; + assert_ldmatrix_producer_result_shape(&ctx, op, register_count); + + let function_ty = call.callee_type(&ctx); + let function_ty = function_ty.deref(&ctx); + let function_ty = function_ty + .downcast_ref::() + .expect("Blackwell ldmatrix has an LLVM function type"); + assert_eq!(function_ty.arg_types().len(), 1); + assert_eq!( + function_ty.arg_types()[0] + .deref(&ctx) + .downcast_ref::() + .expect("Blackwell ldmatrix argument is a pointer") + .address_space(), + 3 + ); + } + + assert_eq!(seen, [1; 12]); + assert_eq!(cast_count, if address_space == 0 { 12 } else { 0 }); + assert_eq!(extract_count, 30, "all multi-register results are unpacked"); + + let module = Operation::get_op::(module_ptr, &ctx).unwrap(); + let ir = llvm_export::export::export_module_to_string(&ctx, &module) + .expect("typed Blackwell ldmatrix module exports to LLVM IR"); + for (identifier, symbol, _, _) in BLACKWELL_LDMATRIX_CASES { + assert!( + ir.contains(&format!("@{symbol}(ptr addrspace(3)")), + "missing exact intrinsic symbol {symbol}:\n{ir}" + ); + assert!( + !ir.contains(&format!("@{identifier}(")), + "encoded Rust identifier leaked into LLVM IR: {identifier}" + ); + } + assert!(ir.contains("b4x16_p64.p3"), "literal _p64 was lost: {ir}"); + assert!(ir.contains("b6x16_p32.p3"), "literal _p32 was lost: {ir}"); + } + Ok(()) +} + +#[test] +fn test_blackwell_ldmatrix_libnvvm_uses_all_exact_convergent_templates_without_externs() +-> Result<(), anyhow::Error> { + for address_space in [0, 3] { + let (ctx, module_ptr) = lower_all_blackwell_ldmatrix_forms( + address_space, + mir_lower::IntrinsicBackend::LibNvvm, + )?; + let mut seen = [0; 12]; + let mut cast_count = 0; + let mut ptrtoint_count = 0; + let mut extract_count = 0; + + for op in lowered_kernel_body(&ctx, module_ptr) { + assert!( + Operation::get_op::(op, &ctx).is_none(), + "libNVVM Blackwell ldmatrix must not emit an extern call" + ); + cast_count += + usize::from(Operation::get_op::(op, &ctx).is_some()); + ptrtoint_count += + usize::from(Operation::get_op::(op, &ctx).is_some()); + extract_count += + usize::from(Operation::get_op::(op, &ctx).is_some()); + + let Some(asm) = Operation::get_op::(op, &ctx) else { + continue; + }; + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .unwrap_or_default(); + let index = BLACKWELL_LDMATRIX_CASES + .iter() + .position(|(_, _, expected, _)| *expected == template) + .expect("exact Blackwell ldmatrix inline-PTX template"); + seen[index] += 1; + let register_count = BLACKWELL_LDMATRIX_CASES[index].3; + assert_eq!( + asm.get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some(ldmatrix_constraints(register_count)) + ); + assert_eq!(llvm::asm_kind(&ctx, &asm), llvm::AsmKind::Convergent); + assert_ldmatrix_producer_result_shape(&ctx, op, register_count); + } + + assert_eq!(seen, [1; 12]); + assert_eq!(cast_count, if address_space == 0 { 12 } else { 0 }); + assert_eq!(ptrtoint_count, 12); + assert_eq!(extract_count, 30, "all multi-register results are unpacked"); + + let module = Operation::get_op::(module_ptr, &ctx).unwrap(); + let ir = llvm_export::export::export_module_to_string(&ctx, &module) + .expect("inline Blackwell ldmatrix module exports to LLVM IR"); + assert_eq!(ir.matches("asm sideeffect").count(), 12, "{ir}"); + assert_eq!(ir.matches("~{memory}").count(), 12, "{ir}"); + assert!(!ir.contains("@llvm.nvvm.ldmatrix"), "{ir}"); + assert!(!ir.contains("llvm__nvvm_dldmatrix"), "{ir}"); + } + Ok(()) +} + +#[test] +fn test_blackwell_ldmatrix_rejects_unadmitted_m16n16_x4() { + use dialect_mir::types::MirPtrType; + use pliron::builtin::types::{IntegerType, Signedness}; + + for backend in [ + mir_lower::IntrinsicBackend::LlvmNvptx, + mir_lower::IntrinsicBackend::LibNvvm, + ] { + let mut ctx = make_test_ctx(); + let u8_ty = IntegerType::get(&ctx, 8, Signedness::Unsigned); + let ptr_ty = MirPtrType::get_shared(&mut ctx, u8_ty.into(), false); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![ptr_ty.into()]); + let pointer = entry.deref(&ctx).get_argument(0); + nvvm::LdmatrixOp::build( + &mut ctx, + pointer, + nvvm::LdmatrixShapeAttr::M16n16, + nvvm::LdmatrixMultiplicityAttr::X4, + nvvm::LdmatrixLayoutAttr::Transposed, + nvvm::LdmatrixElementAttr::B8, + nvvm::LdmatrixStateSpaceAttr::Shared, + ) + .insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + let error = mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: backend, + ..Default::default() + }, + ) + .expect_err("m16n16.x4 must fail closed") + .to_string(); + assert!(error.contains("missing or unsupported variant"), "{error}"); + } +} + +#[test] +fn test_classic_ldmatrix_compatibility_ops_keep_exact_lowering() -> Result<(), anyhow::Error> { + use llvm_export::types as llvm_types; + use pliron::builtin::types::IntegerType; + use pliron::r#type::Typed; + + fn assert_result_shape( + ctx: &Context, + op: pliron::context::Ptr, + register_count: usize, + ) { + let operation = op.deref(ctx); + assert_eq!(operation.get_num_operands(), 1); + assert_eq!(operation.get_num_results(), 1); + let result_ty = operation.get_result(0).get_type(ctx); + let result_ty = result_ty.deref(ctx); + if register_count == 1 { + let result_ty = result_ty + .downcast_ref::() + .expect("x1 returns i32"); + assert_eq!(result_ty.width(), 32); + } else { + let result_ty = result_ty + .downcast_ref::() + .expect("x2/x4 return an LLVM struct"); + assert_eq!(result_ty.num_fields(), register_count); + for index in 0..result_ty.num_fields() { + let field = result_ty.field_type(index); + let field = field.deref(ctx); + let field = field + .downcast_ref::() + .expect("fragment register is i32"); + assert_eq!(field.width(), 32); + } + } + } + + let register_counts = [1, 1, 2, 2, 4, 4]; + for backend in [ + mir_lower::IntrinsicBackend::LlvmNvptx, + mir_lower::IntrinsicBackend::LibNvvm, + ] { + let (ctx, module_ptr) = lower_all_ldmatrix_forms(3, backend, true)?; + let mut seen = [0; 6]; + let mut extract_count = 0; + + for op in lowered_kernel_body(&ctx, module_ptr) { + extract_count += + usize::from(Operation::get_op::(op, &ctx).is_some()); + match backend { + mir_lower::IntrinsicBackend::LlvmNvptx => { + assert!(Operation::get_op::(op, &ctx).is_none()); + let Some(call) = Operation::get_op::(op, &ctx) else { + continue; + }; + let CallOpCallable::Direct(callee) = call.callee(&ctx) else { + panic!("ldmatrix intrinsic call must be direct"); + }; + let callee = callee.to_string(); + let index = LDMATRIX_TYPED_INTRINSICS + .iter() + .position(|expected| *expected == callee) + .expect("exact typed ldmatrix intrinsic"); + seen[index] += 1; + assert_result_shape(&ctx, op, register_counts[index]); + } + mir_lower::IntrinsicBackend::LibNvvm => { + assert!(Operation::get_op::(op, &ctx).is_none()); + let Some(inline_asm) = Operation::get_op::(op, &ctx) else { + continue; + }; + let template = inline_asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .unwrap_or_default(); + let index = LDMATRIX_PTX_TEMPLATES + .iter() + .position(|expected| *expected == template) + .expect("exact ldmatrix PTX template"); + seen[index] += 1; + assert_eq!( + inline_asm + .get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some(LDMATRIX_PTX_CONSTRAINTS[index]) + ); + assert_eq!(llvm::asm_kind(&ctx, &inline_asm), llvm::AsmKind::Convergent); + assert_result_shape(&ctx, op, register_counts[index]); + } + } + } + + assert_eq!(seen, [1; 6]); + assert_eq!(extract_count, 12, "x2/x4 results keep their order"); + } + Ok(()) +} + +#[test] +fn test_ldmatrix_rejects_non_shared_pointer_spaces() { + use dialect_mir::types::MirPtrType; + use pliron::builtin::types::{IntegerType, Signedness}; + + for backend in [ + mir_lower::IntrinsicBackend::LlvmNvptx, + mir_lower::IntrinsicBackend::LibNvvm, + ] { + for address_space in [1, 4, 5] { + let mut ctx = make_test_ctx(); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let ptr_ty = MirPtrType::get(&mut ctx, u32_ty.into(), false, address_space); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![ptr_ty.into()]); + let pointer = entry.deref(&ctx).get_argument(0); + nvvm::LdmatrixOp::build( + &mut ctx, + pointer, + nvvm::LdmatrixShapeAttr::M8n8, + nvvm::LdmatrixMultiplicityAttr::X1, + nvvm::LdmatrixLayoutAttr::Normal, + nvvm::LdmatrixElementAttr::B16, + nvvm::LdmatrixStateSpaceAttr::Shared, + ) + .insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + let error = mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: backend, + ..Default::default() + }, + ) + .expect_err("global/constant pointers must fail closed") + .to_string(); + assert!( + error.contains(&format!("not address space {address_space}")), + "{error}" + ); + } + } +} + +#[test] +fn test_ldmatrix_rejects_non_pointer_operand() { + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![i32_ty.into()]); + let not_a_pointer = entry.deref(&ctx).get_argument(0); + nvvm::LdmatrixOp::build( + &mut ctx, + not_a_pointer, + nvvm::LdmatrixShapeAttr::M8n8, + nvvm::LdmatrixMultiplicityAttr::X1, + nvvm::LdmatrixLayoutAttr::Normal, + nvvm::LdmatrixElementAttr::B16, + nvvm::LdmatrixStateSpaceAttr::Shared, + ) + .insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + let error = mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .expect_err("non-pointer ldmatrix input must fail closed") + .to_string(); + assert!(error.contains("operand must be a MIR pointer"), "{error}"); +} + +#[test] +fn test_ldmatrix_rejects_wrong_result_arity() { + use dialect_mir::types::MirPtrType; + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let ptr_ty = MirPtrType::get_shared(&mut ctx, u32_ty.into(), false); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![ptr_ty.into()]); + let pointer = entry.deref(&ctx).get_argument(0); + let op = Operation::new( + &mut ctx, + nvvm::LdmatrixOp::get_concrete_op_info(), + vec![i32_ty.into(), i32_ty.into()], + vec![pointer], + vec![], + 0, + ); + let ldmatrix = nvvm::LdmatrixOp::new(op); + ldmatrix.set_attr_nvvm_ldmatrix_shape(&ctx, nvvm::LdmatrixShapeAttr::M8n8); + ldmatrix.set_attr_nvvm_ldmatrix_multiplicity(&ctx, nvvm::LdmatrixMultiplicityAttr::X1); + ldmatrix.set_attr_nvvm_ldmatrix_layout(&ctx, nvvm::LdmatrixLayoutAttr::Normal); + ldmatrix.set_attr_nvvm_ldmatrix_element(&ctx, nvvm::LdmatrixElementAttr::B16); + ldmatrix.set_attr_nvvm_ldmatrix_state_space(&ctx, nvvm::LdmatrixStateSpaceAttr::Shared); + op.insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + let error = mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .expect_err("x1 must return exactly one register") + .to_string(); + assert!(error.contains("requires 1 u32 results"), "{error}"); +} diff --git a/crates/mir-lower/tests/lowering_test/mma.rs b/crates/mir-lower/tests/lowering_test/mma.rs new file mode 100644 index 0000000000..e2109d53ad --- /dev/null +++ b/crates/mir-lower/tests/lowering_test/mma.rs @@ -0,0 +1,1754 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use dialect_mir::ops as mir; +use dialect_nvvm::ops as nvvm; +use llvm_export::ops as llvm; +use pliron::builtin::op_interfaces::SymbolOpInterface; +use pliron::builtin::ops::ModuleOp; +use pliron::context::Context; +use pliron::linked_list::ContainsLinkedList; +use pliron::op::Op; +use pliron::operation::Operation; + +use crate::common::{append_return, build_test_kernel, lowered_kernel_body, make_test_ctx}; + +#[test] +fn test_generated_register_mma_variants_lower_to_exact_convergent_inline_ptx() +-> Result<(), anyhow::Error> { + use pliron::builtin::types::{FP32Type, FP64Type, IntegerType, Signedness}; + use pliron::r#type::TypeHandle; + + #[derive(Clone, Copy)] + enum Carrier { + F32, + F64, + I32, + U32, + } + + struct Case { + shape: nvvm::RegisterMmaShapeAttr, + operation: nvvm::RegisterMmaOperationAttr, + accumulator: nvvm::RegisterMmaAccumulatorAttr, + a_element: nvvm::RegisterMmaElementAttr, + b_element: nvvm::RegisterMmaElementAttr, + overflow: nvvm::RegisterMmaOverflowAttr, + operands: &'static [Carrier], + results: &'static [Carrier], + template: String, + constraints: &'static str, + } + + let mut cases = vec![ + Case { + shape: nvvm::RegisterMmaShapeAttr::M16n8k16, + operation: nvvm::RegisterMmaOperationAttr::Multiply, + accumulator: nvvm::RegisterMmaAccumulatorAttr::F32, + a_element: nvvm::RegisterMmaElementAttr::Bf16, + b_element: nvvm::RegisterMmaElementAttr::Bf16, + overflow: nvvm::RegisterMmaOverflowAttr::NotApplicable, + operands: &[ + Carrier::F32, + Carrier::F32, + Carrier::F32, + Carrier::F32, + Carrier::U32, + Carrier::U32, + Carrier::U32, + Carrier::U32, + Carrier::U32, + Carrier::U32, + ], + results: &[Carrier::F32; 4], + template: concat!( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 ", + "{$0, $1, $2, $3}, {$8, $9, $10, $11}, ", + "{$12, $13}, {$4, $5, $6, $7};" + ) + .into(), + constraints: "=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r", + }, + Case { + shape: nvvm::RegisterMmaShapeAttr::M16n8k16, + operation: nvvm::RegisterMmaOperationAttr::Multiply, + accumulator: nvvm::RegisterMmaAccumulatorAttr::F32, + a_element: nvvm::RegisterMmaElementAttr::F16, + b_element: nvvm::RegisterMmaElementAttr::F16, + overflow: nvvm::RegisterMmaOverflowAttr::NotApplicable, + operands: &[ + Carrier::F32, + Carrier::F32, + Carrier::F32, + Carrier::F32, + Carrier::U32, + Carrier::U32, + Carrier::U32, + Carrier::U32, + Carrier::U32, + Carrier::U32, + ], + results: &[Carrier::F32; 4], + template: concat!( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 ", + "{$0, $1, $2, $3}, {$8, $9, $10, $11}, ", + "{$12, $13}, {$4, $5, $6, $7};" + ) + .into(), + constraints: "=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r", + }, + Case { + shape: nvvm::RegisterMmaShapeAttr::M16n8k8, + operation: nvvm::RegisterMmaOperationAttr::Multiply, + accumulator: nvvm::RegisterMmaAccumulatorAttr::F32, + a_element: nvvm::RegisterMmaElementAttr::Tf32, + b_element: nvvm::RegisterMmaElementAttr::Tf32, + overflow: nvvm::RegisterMmaOverflowAttr::NotApplicable, + operands: &[ + Carrier::F32, + Carrier::F32, + Carrier::F32, + Carrier::F32, + Carrier::U32, + Carrier::U32, + Carrier::U32, + Carrier::U32, + Carrier::U32, + Carrier::U32, + ], + results: &[Carrier::F32; 4], + template: concat!( + "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 ", + "{$0, $1, $2, $3}, {$8, $9, $10, $11}, ", + "{$12, $13}, {$4, $5, $6, $7};" + ) + .into(), + constraints: "=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r", + }, + Case { + shape: nvvm::RegisterMmaShapeAttr::M8n8k4, + operation: nvvm::RegisterMmaOperationAttr::Multiply, + accumulator: nvvm::RegisterMmaAccumulatorAttr::F64, + a_element: nvvm::RegisterMmaElementAttr::F64, + b_element: nvvm::RegisterMmaElementAttr::F64, + overflow: nvvm::RegisterMmaOverflowAttr::NotApplicable, + operands: &[Carrier::F64; 4], + results: &[Carrier::F64; 2], + template: concat!( + "mma.sync.aligned.m8n8k4.row.col.f64.f64.f64.f64 ", + "{$0, $1}, {$4}, {$5}, {$2, $3};" + ) + .into(), + constraints: "=d,=d,d,d,d,d", + }, + ]; + let c2_a1_b1: &'static [Carrier] = &[Carrier::I32, Carrier::I32, Carrier::U32, Carrier::U32]; + let c4_a2_b1: &'static [Carrier] = &[ + Carrier::I32, + Carrier::I32, + Carrier::I32, + Carrier::I32, + Carrier::U32, + Carrier::U32, + Carrier::U32, + ]; + let c4_a4_b2: &'static [Carrier] = &[ + Carrier::I32, + Carrier::I32, + Carrier::I32, + Carrier::I32, + Carrier::U32, + Carrier::U32, + Carrier::U32, + Carrier::U32, + Carrier::U32, + Carrier::U32, + ]; + let d2_i32: &'static [Carrier] = &[Carrier::I32; 2]; + let d4_i32: &'static [Carrier] = &[Carrier::I32; 4]; + let register_list = |first, count| { + format!( + "{{{}}}", + (first..first + count) + .map(|index| format!("${index}")) + .collect::>() + .join(", ") + ) + }; + + for (shape, shape_attr, operands, results, accumulator_count, a_count, b_count, constraints) in [ + ( + "m8n8k16", + nvvm::RegisterMmaShapeAttr::M8n8k16, + c2_a1_b1, + d2_i32, + 2, + 1, + 1, + "=r,=r,r,r,r,r", + ), + ( + "m16n8k16", + nvvm::RegisterMmaShapeAttr::M16n8k16, + c4_a2_b1, + d4_i32, + 4, + 2, + 1, + "=r,=r,=r,=r,r,r,r,r,r,r,r", + ), + ( + "m16n8k32", + nvvm::RegisterMmaShapeAttr::M16n8k32, + c4_a4_b2, + d4_i32, + 4, + 4, + 2, + "=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,r", + ), + ] { + for (a_name, a_element) in [ + ("s8", nvvm::RegisterMmaElementAttr::S8), + ("u8", nvvm::RegisterMmaElementAttr::U8), + ] { + for (b_name, b_element) in [ + ("s8", nvvm::RegisterMmaElementAttr::S8), + ("u8", nvvm::RegisterMmaElementAttr::U8), + ] { + for (overflow_name, overflow) in [ + ("", nvvm::RegisterMmaOverflowAttr::Wrapping), + (".satfinite", nvvm::RegisterMmaOverflowAttr::Satfinite), + ] { + let result_count = results.len(); + let d = register_list(0, result_count); + let c = register_list(result_count, accumulator_count); + let a = register_list(result_count + accumulator_count, a_count); + let b = register_list(result_count + accumulator_count + a_count, b_count); + cases.push(Case { + shape: shape_attr.clone(), + operation: nvvm::RegisterMmaOperationAttr::Multiply, + accumulator: nvvm::RegisterMmaAccumulatorAttr::S32, + a_element: a_element.clone(), + b_element: b_element.clone(), + overflow, + operands, + results, + template: format!( + "mma.sync.aligned.{shape}.row.col{overflow_name}.s32.{a_name}.{b_name}.s32 {d}, {a}, {b}, {c};" + ), + constraints, + }); + } + } + } + } + + for (shape, shape_attr, operands, results, accumulator_count, a_count, b_count, constraints) in [ + ( + "m8n8k32", + nvvm::RegisterMmaShapeAttr::M8n8k32, + c2_a1_b1, + d2_i32, + 2, + 1, + 1, + "=r,=r,r,r,r,r", + ), + ( + "m16n8k32", + nvvm::RegisterMmaShapeAttr::M16n8k32, + c4_a2_b1, + d4_i32, + 4, + 2, + 1, + "=r,=r,=r,=r,r,r,r,r,r,r,r", + ), + ( + "m16n8k64", + nvvm::RegisterMmaShapeAttr::M16n8k64, + c4_a4_b2, + d4_i32, + 4, + 4, + 2, + "=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,r", + ), + ] { + for (a_name, a_element) in [ + ("s4", nvvm::RegisterMmaElementAttr::S4), + ("u4", nvvm::RegisterMmaElementAttr::U4), + ] { + for (b_name, b_element) in [ + ("s4", nvvm::RegisterMmaElementAttr::S4), + ("u4", nvvm::RegisterMmaElementAttr::U4), + ] { + for (overflow_name, overflow) in [ + ("", nvvm::RegisterMmaOverflowAttr::Wrapping), + (".satfinite", nvvm::RegisterMmaOverflowAttr::Satfinite), + ] { + let result_count = results.len(); + let d = register_list(0, result_count); + let c = register_list(result_count, accumulator_count); + let a = register_list(result_count + accumulator_count, a_count); + let b = register_list(result_count + accumulator_count + a_count, b_count); + cases.push(Case { + shape: shape_attr.clone(), + operation: nvvm::RegisterMmaOperationAttr::Multiply, + accumulator: nvvm::RegisterMmaAccumulatorAttr::S32, + a_element: a_element.clone(), + b_element: b_element.clone(), + overflow, + operands, + results, + template: format!( + "mma.sync.aligned.{shape}.row.col{overflow_name}.s32.{a_name}.{b_name}.s32 {d}, {a}, {b}, {c};" + ), + constraints, + }); + } + } + } + } + + for (shape, shape_attr, operands, results, accumulator_count, a_count, b_count, constraints) in [ + ( + "m8n8k128", + nvvm::RegisterMmaShapeAttr::M8n8k128, + c2_a1_b1, + d2_i32, + 2, + 1, + 1, + "=r,=r,r,r,r,r", + ), + ( + "m16n8k128", + nvvm::RegisterMmaShapeAttr::M16n8k128, + c4_a2_b1, + d4_i32, + 4, + 2, + 1, + "=r,=r,=r,=r,r,r,r,r,r,r,r", + ), + ( + "m16n8k256", + nvvm::RegisterMmaShapeAttr::M16n8k256, + c4_a4_b2, + d4_i32, + 4, + 4, + 2, + "=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,r", + ), + ] { + let result_count = results.len(); + let d = register_list(0, result_count); + let c = register_list(result_count, accumulator_count); + let a = register_list(result_count + accumulator_count, a_count); + let b = register_list(result_count + accumulator_count + a_count, b_count); + for (operation_name, operation) in [ + ("xor", nvvm::RegisterMmaOperationAttr::XorPopc), + ("and", nvvm::RegisterMmaOperationAttr::AndPopc), + ] { + cases.push(Case { + shape: shape_attr.clone(), + operation, + accumulator: nvvm::RegisterMmaAccumulatorAttr::S32, + a_element: nvvm::RegisterMmaElementAttr::B1, + b_element: nvvm::RegisterMmaElementAttr::B1, + overflow: nvvm::RegisterMmaOverflowAttr::Wrapping, + operands, + results, + template: format!( + "mma.sync.aligned.{shape}.row.col.s32.b1.b1.s32.{operation_name}.popc {d}, {a}, {b}, {c};" + ), + constraints, + }); + } + } + assert_eq!(cases.len(), 58); + + let carrier_type = |ctx: &Context, carrier: Carrier| -> TypeHandle { + match carrier { + Carrier::F32 => FP32Type::get(ctx).into(), + Carrier::F64 => FP64Type::get(ctx).into(), + Carrier::I32 => IntegerType::get(ctx, 32, Signedness::Signed).into(), + Carrier::U32 => IntegerType::get(ctx, 32, Signedness::Unsigned).into(), + } + }; + + for backend in [ + mir_lower::IntrinsicBackend::LlvmNvptx, + mir_lower::IntrinsicBackend::LibNvvm, + ] { + for case in &cases { + let mut ctx = make_test_ctx(); + let argument_types = case + .operands + .iter() + .map(|carrier| carrier_type(&ctx, *carrier)) + .collect(); + let result_types = case + .results + .iter() + .map(|carrier| carrier_type(&ctx, *carrier)) + .collect(); + let (module_ptr, entry) = build_test_kernel(&mut ctx, argument_types); + let operands = (0..case.operands.len()) + .map(|index| entry.deref(&ctx).get_argument(index)) + .collect(); + let operation = Operation::new( + &mut ctx, + nvvm::RegisterMmaOp::get_concrete_op_info(), + result_types, + operands, + vec![], + 0, + ); + let mma = nvvm::RegisterMmaOp::new(operation); + mma.set_attr_nvvm_register_mma_shape(&ctx, case.shape.clone()); + let uses_legacy_default = matches!( + (&case.operation, &case.a_element), + ( + nvvm::RegisterMmaOperationAttr::Multiply, + nvvm::RegisterMmaElementAttr::Bf16 + ) + ); + if !uses_legacy_default { + mma.set_attr_nvvm_register_mma_operation(&ctx, case.operation.clone()); + } + mma.set_attr_nvvm_register_mma_accumulator(&ctx, case.accumulator.clone()); + mma.set_attr_nvvm_register_mma_a_element(&ctx, case.a_element.clone()); + mma.set_attr_nvvm_register_mma_b_element(&ctx, case.b_element.clone()); + mma.set_attr_nvvm_register_mma_a_layout(&ctx, nvvm::RegisterMmaLayoutAttr::Row); + mma.set_attr_nvvm_register_mma_b_layout(&ctx, nvvm::RegisterMmaLayoutAttr::Col); + mma.set_attr_nvvm_register_mma_overflow(&ctx, case.overflow.clone()); + operation.insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: backend, + ..Default::default() + }, + ) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let body = lowered_kernel_body(&ctx, module_ptr); + let lowered = body + .iter() + .filter_map(|op| Operation::get_op::(*op, &ctx)) + .collect::>(); + assert_eq!(lowered.len(), 1, "{:?}", backend); + let asm = &lowered[0]; + assert_eq!( + asm.get_attr_inline_asm_template(&ctx) + .as_deref() + .map(|value| String::from(value.clone())), + Some(case.template.clone()) + ); + assert_eq!( + asm.get_attr_inline_asm_constraints(&ctx) + .as_deref() + .map(|value| String::from(value.clone())), + Some(case.constraints.to_string()) + ); + assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); + assert_eq!( + asm.get_operation().deref(&ctx).get_num_operands(), + case.operands.len() + ); + assert_eq!(asm.get_operation().deref(&ctx).get_num_results(), 1); + assert_eq!( + body.iter() + .filter(|op| Operation::get_op::(**op, &ctx).is_some()) + .count(), + case.results.len() + ); + + let module = Operation::get_op::(module_ptr, &ctx).unwrap(); + let ir = llvm_export::export::export_module_to_string(&ctx, &module) + .expect("generated register MMA exports to LLVM IR"); + assert!(ir.contains("asm sideeffect"), "{ir}"); + assert!(ir.contains("{ convergent }"), "{ir}"); + } + } + + Ok(()) +} + +#[test] +fn test_generated_sparse_mma_variants_lower_to_exact_convergent_inline_ptx() +-> Result<(), anyhow::Error> { + use pliron::builtin::attributes::IntegerAttr; + use pliron::builtin::types::{IntegerType, Signedness}; + use pliron::utils::apint::APInt; + use std::num::NonZeroUsize; + + let cases = [ + ( + "s8", + "s8", + "", + nvvm::SparseMmaElementAttr::S8, + nvvm::SparseMmaElementAttr::S8, + nvvm::SparseMmaOverflowAttr::Wrapping, + ), + ( + "s8", + "u8", + "", + nvvm::SparseMmaElementAttr::S8, + nvvm::SparseMmaElementAttr::U8, + nvvm::SparseMmaOverflowAttr::Wrapping, + ), + ( + "u8", + "u8", + "", + nvvm::SparseMmaElementAttr::U8, + nvvm::SparseMmaElementAttr::U8, + nvvm::SparseMmaOverflowAttr::Wrapping, + ), + ( + "u8", + "s8", + "", + nvvm::SparseMmaElementAttr::U8, + nvvm::SparseMmaElementAttr::S8, + nvvm::SparseMmaOverflowAttr::Wrapping, + ), + ( + "s8", + "s8", + ".satfinite", + nvvm::SparseMmaElementAttr::S8, + nvvm::SparseMmaElementAttr::S8, + nvvm::SparseMmaOverflowAttr::Satfinite, + ), + ( + "s8", + "u8", + ".satfinite", + nvvm::SparseMmaElementAttr::S8, + nvvm::SparseMmaElementAttr::U8, + nvvm::SparseMmaOverflowAttr::Satfinite, + ), + ( + "u8", + "u8", + ".satfinite", + nvvm::SparseMmaElementAttr::U8, + nvvm::SparseMmaElementAttr::U8, + nvvm::SparseMmaOverflowAttr::Satfinite, + ), + ( + "u8", + "s8", + ".satfinite", + nvvm::SparseMmaElementAttr::U8, + nvvm::SparseMmaElementAttr::S8, + nvvm::SparseMmaOverflowAttr::Satfinite, + ), + ]; + let metadata_modes = [ + ("sp", nvvm::SparseMmaMetadataAttr::Standard), + ("sp::ordered_metadata", nvvm::SparseMmaMetadataAttr::Ordered), + ]; + + for backend in [ + mir_lower::IntrinsicBackend::LlvmNvptx, + mir_lower::IntrinsicBackend::LibNvvm, + ] { + for (metadata_name, metadata) in &metadata_modes { + for (index, (a_name, b_name, overflow_name, a_element, b_element, overflow)) in + cases.iter().enumerate() + { + let mut ctx = make_test_ctx(); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let argument_types = (0..4) + .map(|_| i32_ty.into()) + .chain((0..5).map(|_| u32_ty.into())) + .collect(); + let (module_ptr, entry) = build_test_kernel(&mut ctx, argument_types); + + let selector_value = (index % 2) as u32; + let selector_op = Operation::new( + &mut ctx, + mir::MirConstantOp::get_concrete_op_info(), + vec![u32_ty.into()], + vec![], + vec![], + 0, + ); + mir::MirConstantOp::new(selector_op).set_attr_value( + &ctx, + IntegerAttr::new( + u32_ty, + APInt::from_u32(selector_value, NonZeroUsize::new(32).unwrap()), + ), + ); + selector_op.insert_at_back(entry, &ctx); + let selector = selector_op.deref(&ctx).get_result(0); + + let operands = (0..9) + .map(|operand| entry.deref(&ctx).get_argument(operand)) + .chain(std::iter::once(selector)) + .collect(); + let operation = Operation::new( + &mut ctx, + nvvm::SparseMmaOp::get_concrete_op_info(), + vec![i32_ty.into(); 4], + operands, + vec![], + 0, + ); + let mma = nvvm::SparseMmaOp::new(operation); + mma.set_attr_nvvm_sparse_mma_shape(&ctx, nvvm::SparseMmaShapeAttr::M16n8k32); + mma.set_attr_nvvm_sparse_mma_accumulator(&ctx, nvvm::SparseMmaAccumulatorAttr::S32); + mma.set_attr_nvvm_sparse_mma_a_element(&ctx, a_element.clone()); + mma.set_attr_nvvm_sparse_mma_b_element(&ctx, b_element.clone()); + mma.set_attr_nvvm_sparse_mma_a_layout(&ctx, nvvm::SparseMmaLayoutAttr::Row); + mma.set_attr_nvvm_sparse_mma_b_layout(&ctx, nvvm::SparseMmaLayoutAttr::Col); + mma.set_attr_nvvm_sparse_mma_overflow(&ctx, overflow.clone()); + mma.set_attr_nvvm_sparse_mma_metadata(&ctx, metadata.clone()); + mma.set_attr_nvvm_sparse_mma_selector( + &ctx, + nvvm::SparseMmaSelectorAttr::ImmediateZeroOrOne, + ); + operation.insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: backend, + ..Default::default() + }, + ) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let body = lowered_kernel_body(&ctx, module_ptr); + let lowered = body + .iter() + .filter_map(|op| Operation::get_op::(*op, &ctx)) + .collect::>(); + assert_eq!(lowered.len(), 1, "{backend:?}"); + let asm = &lowered[0]; + let expected_template = format!( + "mma.{metadata_name}.sync.aligned.m16n8k32.row.col{overflow_name}.s32.{a_name}.{b_name}.s32 {{$0, $1, $2, $3}}, {{$8, $9}}, {{$10, $11}}, {{$4, $5, $6, $7}}, $12, $13;" + ); + assert_eq!( + asm.get_attr_inline_asm_template(&ctx) + .as_deref() + .map(|value| String::from(value.clone())), + Some(expected_template) + ); + assert_eq!( + asm.get_attr_inline_asm_constraints(&ctx) + .as_deref() + .map(|value| String::from(value.clone())), + Some("=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,n".to_string()) + ); + assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); + let asm_operation = asm.get_operation().deref(&ctx); + assert_eq!(asm_operation.get_num_operands(), 10); + assert_eq!(asm_operation.get_num_results(), 1); + let lowered_selector = asm_operation.get_operand(9); + let defining_op = lowered_selector + .defining_op() + .expect("sparse MMA selector remains an LLVM constant"); + let constant = Operation::get_op::(defining_op, &ctx) + .expect("sparse MMA selector remains an LLVM integer constant"); + let attribute = constant.get_value(&ctx); + let integer = attribute + .downcast_ref::() + .expect("sparse MMA selector is an integer"); + assert_eq!(integer.value().bw(), 32); + assert_eq!(integer.value().to_u64(), selector_value as u64); + assert_eq!( + body.iter() + .filter(|op| { + Operation::get_op::(**op, &ctx).is_some() + }) + .count(), + 4 + ); + + let module = Operation::get_op::(module_ptr, &ctx).unwrap(); + let ir = llvm_export::export::export_module_to_string(&ctx, &module) + .expect("generated sparse MMA exports to LLVM IR"); + assert!(ir.contains("asm sideeffect"), "{ir}"); + assert!(ir.contains("{ convergent }"), "{ir}"); + } + } + } + + Ok(()) +} + +#[test] +fn test_generated_sparse_mma_m16n8k64_lowers_to_exact_convergent_inline_ptx() +-> Result<(), anyhow::Error> { + use pliron::builtin::attributes::IntegerAttr; + use pliron::builtin::types::{IntegerType, Signedness}; + use pliron::utils::apint::APInt; + use std::num::NonZeroUsize; + + let cases = [ + ( + "s8", + "u8", + "", + nvvm::SparseMmaElementAttr::S8, + nvvm::SparseMmaElementAttr::U8, + nvvm::SparseMmaOverflowAttr::Wrapping, + ), + ( + "u8", + "s8", + ".satfinite", + nvvm::SparseMmaElementAttr::U8, + nvvm::SparseMmaElementAttr::S8, + nvvm::SparseMmaOverflowAttr::Satfinite, + ), + ]; + let metadata_cases = [ + ("sp", nvvm::SparseMmaMetadataAttr::Standard), + ("sp::ordered_metadata", nvvm::SparseMmaMetadataAttr::Ordered), + ]; + + for backend in [ + mir_lower::IntrinsicBackend::LlvmNvptx, + mir_lower::IntrinsicBackend::LibNvvm, + ] { + for (metadata_name, metadata) in &metadata_cases { + for (a_name, b_name, overflow_name, a_element, b_element, overflow) in &cases { + let mut ctx = make_test_ctx(); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let argument_types = (0..4) + .map(|_| i32_ty.into()) + .chain((0..9).map(|_| u32_ty.into())) + .collect(); + let (module_ptr, entry) = build_test_kernel(&mut ctx, argument_types); + + let selector_op = Operation::new( + &mut ctx, + mir::MirConstantOp::get_concrete_op_info(), + vec![u32_ty.into()], + vec![], + vec![], + 0, + ); + mir::MirConstantOp::new(selector_op).set_attr_value( + &ctx, + IntegerAttr::new(u32_ty, APInt::from_u32(0, NonZeroUsize::new(32).unwrap())), + ); + selector_op.insert_at_back(entry, &ctx); + let selector = selector_op.deref(&ctx).get_result(0); + + let operands = (0..13) + .map(|index| entry.deref(&ctx).get_argument(index)) + .chain(std::iter::once(selector)) + .collect(); + let operation = Operation::new( + &mut ctx, + nvvm::SparseMmaOp::get_concrete_op_info(), + vec![i32_ty.into(); 4], + operands, + vec![], + 0, + ); + let mma = nvvm::SparseMmaOp::new(operation); + mma.set_attr_nvvm_sparse_mma_shape(&ctx, nvvm::SparseMmaShapeAttr::M16n8k64); + mma.set_attr_nvvm_sparse_mma_accumulator(&ctx, nvvm::SparseMmaAccumulatorAttr::S32); + mma.set_attr_nvvm_sparse_mma_a_element(&ctx, a_element.clone()); + mma.set_attr_nvvm_sparse_mma_b_element(&ctx, b_element.clone()); + mma.set_attr_nvvm_sparse_mma_a_layout(&ctx, nvvm::SparseMmaLayoutAttr::Row); + mma.set_attr_nvvm_sparse_mma_b_layout(&ctx, nvvm::SparseMmaLayoutAttr::Col); + mma.set_attr_nvvm_sparse_mma_overflow(&ctx, overflow.clone()); + mma.set_attr_nvvm_sparse_mma_metadata(&ctx, metadata.clone()); + mma.set_attr_nvvm_sparse_mma_selector( + &ctx, + nvvm::SparseMmaSelectorAttr::ImmediateZero, + ); + operation.insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: backend, + ..Default::default() + }, + ) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let body = lowered_kernel_body(&ctx, module_ptr); + let lowered = body + .iter() + .filter_map(|op| Operation::get_op::(*op, &ctx)) + .collect::>(); + assert_eq!(lowered.len(), 1, "{backend:?}"); + let asm = &lowered[0]; + let expected_template = format!( + "mma.{metadata_name}.sync.aligned.m16n8k64.row.col{overflow_name}.s32.{a_name}.{b_name}.s32 {{$0, $1, $2, $3}}, {{$8, $9, $10, $11}}, {{$12, $13, $14, $15}}, {{$4, $5, $6, $7}}, $16, $17;" + ); + assert_eq!( + asm.get_attr_inline_asm_template(&ctx) + .as_deref() + .map(|value| String::from(value.clone())), + Some(expected_template) + ); + assert_eq!( + asm.get_attr_inline_asm_constraints(&ctx) + .as_deref() + .map(|value| String::from(value.clone())), + Some("=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,r,r,r,r,n".to_string()) + ); + assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); + let asm_operation = asm.get_operation().deref(&ctx); + assert_eq!(asm_operation.get_num_operands(), 14); + assert_eq!(asm_operation.get_num_results(), 1); + let lowered_selector = asm_operation.get_operand(13); + let defining_op = lowered_selector + .defining_op() + .expect("sparse MMA selector remains an LLVM constant"); + let constant = Operation::get_op::(defining_op, &ctx) + .expect("sparse MMA selector remains an LLVM integer constant"); + let attribute = constant.get_value(&ctx); + let integer = attribute + .downcast_ref::() + .expect("sparse MMA selector is an integer"); + assert_eq!(integer.value().bw(), 32); + assert_eq!(integer.value().to_u64(), 0); + assert_eq!( + body.iter() + .filter(|op| { + Operation::get_op::(**op, &ctx).is_some() + }) + .count(), + 4 + ); + + let module = Operation::get_op::(module_ptr, &ctx).unwrap(); + let ir = llvm_export::export::export_module_to_string(&ctx, &module) + .expect("generated sparse MMA exports to LLVM IR"); + assert!(ir.contains("asm sideeffect"), "{ir}"); + assert!(ir.contains("{ convergent }"), "{ir}"); + } + } + } + + Ok(()) +} + +#[test] +fn test_generated_sparse_mma_m16n8k64_int4_lowers_both_metadata_modes() -> Result<(), anyhow::Error> +{ + use pliron::builtin::attributes::IntegerAttr; + use pliron::builtin::types::{IntegerType, Signedness}; + use pliron::utils::apint::APInt; + use std::num::NonZeroUsize; + + let cases = [ + ( + "s4", + "s4", + "", + nvvm::SparseMmaElementAttr::S4, + nvvm::SparseMmaElementAttr::S4, + nvvm::SparseMmaOverflowAttr::Wrapping, + ), + ( + "s4", + "u4", + "", + nvvm::SparseMmaElementAttr::S4, + nvvm::SparseMmaElementAttr::U4, + nvvm::SparseMmaOverflowAttr::Wrapping, + ), + ( + "u4", + "u4", + "", + nvvm::SparseMmaElementAttr::U4, + nvvm::SparseMmaElementAttr::U4, + nvvm::SparseMmaOverflowAttr::Wrapping, + ), + ( + "u4", + "s4", + "", + nvvm::SparseMmaElementAttr::U4, + nvvm::SparseMmaElementAttr::S4, + nvvm::SparseMmaOverflowAttr::Wrapping, + ), + ( + "s4", + "s4", + ".satfinite", + nvvm::SparseMmaElementAttr::S4, + nvvm::SparseMmaElementAttr::S4, + nvvm::SparseMmaOverflowAttr::Satfinite, + ), + ( + "s4", + "u4", + ".satfinite", + nvvm::SparseMmaElementAttr::S4, + nvvm::SparseMmaElementAttr::U4, + nvvm::SparseMmaOverflowAttr::Satfinite, + ), + ( + "u4", + "u4", + ".satfinite", + nvvm::SparseMmaElementAttr::U4, + nvvm::SparseMmaElementAttr::U4, + nvvm::SparseMmaOverflowAttr::Satfinite, + ), + ( + "u4", + "s4", + ".satfinite", + nvvm::SparseMmaElementAttr::U4, + nvvm::SparseMmaElementAttr::S4, + nvvm::SparseMmaOverflowAttr::Satfinite, + ), + ]; + + for backend in [ + mir_lower::IntrinsicBackend::LlvmNvptx, + mir_lower::IntrinsicBackend::LibNvvm, + ] { + for (metadata_name, metadata) in [ + ("sp", nvvm::SparseMmaMetadataAttr::Standard), + ("sp::ordered_metadata", nvvm::SparseMmaMetadataAttr::Ordered), + ] { + for (index, (a_name, b_name, overflow_name, a_element, b_element, overflow)) in + cases.iter().enumerate() + { + let mut ctx = make_test_ctx(); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let argument_types = (0..4) + .map(|_| i32_ty.into()) + .chain((0..5).map(|_| u32_ty.into())) + .collect(); + let (module_ptr, entry) = build_test_kernel(&mut ctx, argument_types); + + let selector_value = (index % 2) as u32; + let selector_op = Operation::new( + &mut ctx, + mir::MirConstantOp::get_concrete_op_info(), + vec![u32_ty.into()], + vec![], + vec![], + 0, + ); + mir::MirConstantOp::new(selector_op).set_attr_value( + &ctx, + IntegerAttr::new( + u32_ty, + APInt::from_u32(selector_value, NonZeroUsize::new(32).unwrap()), + ), + ); + selector_op.insert_at_back(entry, &ctx); + let selector = selector_op.deref(&ctx).get_result(0); + + let operands = (0..9) + .map(|operand| entry.deref(&ctx).get_argument(operand)) + .chain(std::iter::once(selector)) + .collect(); + let operation = Operation::new( + &mut ctx, + nvvm::SparseMmaOp::get_concrete_op_info(), + vec![i32_ty.into(); 4], + operands, + vec![], + 0, + ); + let mma = nvvm::SparseMmaOp::new(operation); + mma.set_attr_nvvm_sparse_mma_shape(&ctx, nvvm::SparseMmaShapeAttr::M16n8k64); + mma.set_attr_nvvm_sparse_mma_accumulator(&ctx, nvvm::SparseMmaAccumulatorAttr::S32); + mma.set_attr_nvvm_sparse_mma_a_element(&ctx, a_element.clone()); + mma.set_attr_nvvm_sparse_mma_b_element(&ctx, b_element.clone()); + mma.set_attr_nvvm_sparse_mma_a_layout(&ctx, nvvm::SparseMmaLayoutAttr::Row); + mma.set_attr_nvvm_sparse_mma_b_layout(&ctx, nvvm::SparseMmaLayoutAttr::Col); + mma.set_attr_nvvm_sparse_mma_overflow(&ctx, overflow.clone()); + mma.set_attr_nvvm_sparse_mma_metadata(&ctx, metadata.clone()); + mma.set_attr_nvvm_sparse_mma_selector( + &ctx, + nvvm::SparseMmaSelectorAttr::ImmediateZeroOrOne, + ); + operation.insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: backend, + ..Default::default() + }, + ) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let body = lowered_kernel_body(&ctx, module_ptr); + let lowered = body + .iter() + .filter_map(|op| Operation::get_op::(*op, &ctx)) + .collect::>(); + assert_eq!(lowered.len(), 1, "{backend:?}"); + let asm = &lowered[0]; + let expected_template = format!( + "mma.{metadata_name}.sync.aligned.m16n8k64.row.col{overflow_name}.s32.{a_name}.{b_name}.s32 {{$0, $1, $2, $3}}, {{$8, $9}}, {{$10, $11}}, {{$4, $5, $6, $7}}, $12, $13;" + ); + assert_eq!( + asm.get_attr_inline_asm_template(&ctx) + .as_deref() + .map(|value| String::from(value.clone())), + Some(expected_template) + ); + assert_eq!( + asm.get_attr_inline_asm_constraints(&ctx) + .as_deref() + .map(|value| String::from(value.clone())), + Some("=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,n".to_string()) + ); + assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); + let asm_operation = asm.get_operation().deref(&ctx); + assert_eq!(asm_operation.get_num_operands(), 10); + assert_eq!(asm_operation.get_num_results(), 1); + let lowered_selector = asm_operation.get_operand(9); + let defining_op = lowered_selector + .defining_op() + .expect("sparse MMA selector remains an LLVM constant"); + let constant = Operation::get_op::(defining_op, &ctx) + .expect("sparse MMA selector remains an LLVM integer constant"); + let attribute = constant.get_value(&ctx); + let integer = attribute + .downcast_ref::() + .expect("sparse MMA selector is an integer"); + assert_eq!(integer.value().to_u64(), selector_value as u64); + + let module = Operation::get_op::(module_ptr, &ctx).unwrap(); + let ir = llvm_export::export::export_module_to_string(&ctx, &module) + .expect("generated sparse MMA exports to LLVM IR"); + assert!(ir.contains("asm sideeffect"), "{ir}"); + assert!(ir.contains("{ convergent }"), "{ir}"); + } + } + } + + Ok(()) +} + +#[test] +fn test_generated_sparse_mma_m16n8k128_int4_lowers_both_metadata_modes() -> Result<(), anyhow::Error> +{ + use pliron::builtin::attributes::IntegerAttr; + use pliron::builtin::types::{IntegerType, Signedness}; + use pliron::utils::apint::APInt; + use std::num::NonZeroUsize; + + let cases = [ + ( + "s4", + "s4", + "", + nvvm::SparseMmaElementAttr::S4, + nvvm::SparseMmaElementAttr::S4, + nvvm::SparseMmaOverflowAttr::Wrapping, + ), + ( + "s4", + "u4", + "", + nvvm::SparseMmaElementAttr::S4, + nvvm::SparseMmaElementAttr::U4, + nvvm::SparseMmaOverflowAttr::Wrapping, + ), + ( + "u4", + "u4", + "", + nvvm::SparseMmaElementAttr::U4, + nvvm::SparseMmaElementAttr::U4, + nvvm::SparseMmaOverflowAttr::Wrapping, + ), + ( + "u4", + "s4", + "", + nvvm::SparseMmaElementAttr::U4, + nvvm::SparseMmaElementAttr::S4, + nvvm::SparseMmaOverflowAttr::Wrapping, + ), + ( + "s4", + "s4", + ".satfinite", + nvvm::SparseMmaElementAttr::S4, + nvvm::SparseMmaElementAttr::S4, + nvvm::SparseMmaOverflowAttr::Satfinite, + ), + ( + "s4", + "u4", + ".satfinite", + nvvm::SparseMmaElementAttr::S4, + nvvm::SparseMmaElementAttr::U4, + nvvm::SparseMmaOverflowAttr::Satfinite, + ), + ( + "u4", + "u4", + ".satfinite", + nvvm::SparseMmaElementAttr::U4, + nvvm::SparseMmaElementAttr::U4, + nvvm::SparseMmaOverflowAttr::Satfinite, + ), + ( + "u4", + "s4", + ".satfinite", + nvvm::SparseMmaElementAttr::U4, + nvvm::SparseMmaElementAttr::S4, + nvvm::SparseMmaOverflowAttr::Satfinite, + ), + ]; + + for backend in [ + mir_lower::IntrinsicBackend::LlvmNvptx, + mir_lower::IntrinsicBackend::LibNvvm, + ] { + for (metadata_name, metadata) in [ + ("sp", nvvm::SparseMmaMetadataAttr::Standard), + ("sp::ordered_metadata", nvvm::SparseMmaMetadataAttr::Ordered), + ] { + for (a_name, b_name, overflow_name, a_element, b_element, overflow) in &cases { + let mut ctx = make_test_ctx(); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let argument_types = (0..4) + .map(|_| i32_ty.into()) + .chain((0..9).map(|_| u32_ty.into())) + .collect(); + let (module_ptr, entry) = build_test_kernel(&mut ctx, argument_types); + + let selector_op = Operation::new( + &mut ctx, + mir::MirConstantOp::get_concrete_op_info(), + vec![u32_ty.into()], + vec![], + vec![], + 0, + ); + mir::MirConstantOp::new(selector_op).set_attr_value( + &ctx, + IntegerAttr::new(u32_ty, APInt::from_u32(0, NonZeroUsize::new(32).unwrap())), + ); + selector_op.insert_at_back(entry, &ctx); + let selector = selector_op.deref(&ctx).get_result(0); + + let operands = (0..13) + .map(|operand| entry.deref(&ctx).get_argument(operand)) + .chain(std::iter::once(selector)) + .collect(); + let operation = Operation::new( + &mut ctx, + nvvm::SparseMmaOp::get_concrete_op_info(), + vec![i32_ty.into(); 4], + operands, + vec![], + 0, + ); + let mma = nvvm::SparseMmaOp::new(operation); + mma.set_attr_nvvm_sparse_mma_shape(&ctx, nvvm::SparseMmaShapeAttr::M16n8k128); + mma.set_attr_nvvm_sparse_mma_accumulator(&ctx, nvvm::SparseMmaAccumulatorAttr::S32); + mma.set_attr_nvvm_sparse_mma_a_element(&ctx, a_element.clone()); + mma.set_attr_nvvm_sparse_mma_b_element(&ctx, b_element.clone()); + mma.set_attr_nvvm_sparse_mma_a_layout(&ctx, nvvm::SparseMmaLayoutAttr::Row); + mma.set_attr_nvvm_sparse_mma_b_layout(&ctx, nvvm::SparseMmaLayoutAttr::Col); + mma.set_attr_nvvm_sparse_mma_overflow(&ctx, overflow.clone()); + mma.set_attr_nvvm_sparse_mma_metadata(&ctx, metadata.clone()); + mma.set_attr_nvvm_sparse_mma_selector( + &ctx, + nvvm::SparseMmaSelectorAttr::ImmediateZero, + ); + operation.insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: backend, + ..Default::default() + }, + ) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let body = lowered_kernel_body(&ctx, module_ptr); + let lowered = body + .iter() + .filter_map(|op| Operation::get_op::(*op, &ctx)) + .collect::>(); + assert_eq!(lowered.len(), 1, "{backend:?}"); + let asm = &lowered[0]; + let expected_template = format!( + "mma.{metadata_name}.sync.aligned.m16n8k128.row.col{overflow_name}.s32.{a_name}.{b_name}.s32 {{$0, $1, $2, $3}}, {{$8, $9, $10, $11}}, {{$12, $13, $14, $15}}, {{$4, $5, $6, $7}}, $16, $17;" + ); + assert_eq!( + asm.get_attr_inline_asm_template(&ctx) + .as_deref() + .map(|value| String::from(value.clone())), + Some(expected_template) + ); + assert_eq!( + asm.get_attr_inline_asm_constraints(&ctx) + .as_deref() + .map(|value| String::from(value.clone())), + Some("=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,r,r,r,r,n".to_string()) + ); + assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); + let asm_operation = asm.get_operation().deref(&ctx); + assert_eq!(asm_operation.get_num_operands(), 14); + assert_eq!(asm_operation.get_num_results(), 1); + let lowered_selector = asm_operation.get_operand(13); + let defining_op = lowered_selector + .defining_op() + .expect("sparse MMA selector remains an LLVM constant"); + let constant = Operation::get_op::(defining_op, &ctx) + .expect("sparse MMA selector remains an LLVM integer constant"); + let attribute = constant.get_value(&ctx); + let integer = attribute + .downcast_ref::() + .expect("sparse MMA selector is an integer"); + assert_eq!(integer.value().to_u64(), 0); + + let module = Operation::get_op::(module_ptr, &ctx).unwrap(); + let ir = llvm_export::export::export_module_to_string(&ctx, &module) + .expect("generated sparse MMA exports to LLVM IR"); + assert!(ir.contains("asm sideeffect"), "{ir}"); + assert!(ir.contains("{ convergent }"), "{ir}"); + } + } + } + + Ok(()) +} + +// --------------------------------------------------------------------------- +// mma.sync m16n8k16 bf16 intrinsic lowering test +// --------------------------------------------------------------------------- + +#[test] +fn test_mma_m16n8k16_f32_bf16_lowers_to_inline_asm() -> Result<(), anyhow::Error> { + let mut ctx = make_test_ctx(); + let f32_ty = pliron::builtin::types::FP32Type::get(&ctx); + let i32_ty = pliron::builtin::types::IntegerType::get( + &ctx, + 32, + pliron::builtin::types::Signedness::Signless, + ); + let argument_types = (0..4) + .map(|_| f32_ty.into()) + .chain((0..6).map(|_| i32_ty.into())) + .collect(); + let (module_ptr, entry) = build_test_kernel(&mut ctx, argument_types); + let operands = (0..10) + .map(|index| entry.deref(&ctx).get_argument(index)) + .collect(); + + let op = Operation::new( + &mut ctx, + nvvm::MmaM16N8K16F32Bf16Op::get_concrete_op_info(), + vec![f32_ty.into(); 4], + operands, + vec![], + 0, + ); + op.insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let mut found = 0; + let module_region = module_ptr.deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + for module_op in module_block.deref(&ctx).iter(&ctx) { + let Some(function) = Operation::get_op::(module_op, &ctx) else { + continue; + }; + if function.get_symbol_name(&ctx).to_string() != "kernel_func" { + continue; + } + let body = function.get_operation().deref(&ctx).get_region(0); + for block in body.deref(&ctx).iter(&ctx) { + for body_op in block.deref(&ctx).iter(&ctx) { + let Some(asm) = Operation::get_op::(body_op, &ctx) else { + continue; + }; + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())); + let constraints = asm + .get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())); + if !template.as_deref().is_some_and(|t| { + t.contains("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32") + }) { + continue; + } + found += 1; + let template = template.expect("MMA inline asm must have a template"); + assert_eq!( + template, + concat!( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 ", + "{$0, $1, $2, $3}, ", + "{$8, $9, $10, $11}, ", + "{$12, $13}, ", + "{$4, $5, $6, $7};" + ) + ); + for forbidden in [".reg", "ld.", "st.", "["] { + assert!( + !template.contains(forbidden), + "register-only MMA must not contain {forbidden:?}: {template}" + ); + } + assert_eq!( + constraints.as_deref(), + Some("=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r") + ); + assert_eq!( + llvm::asm_kind_opt(&ctx, &asm), + Some(llvm::AsmKind::Convergent) + ); + assert_eq!( + body_op.deref(&ctx).get_num_operands(), + 10, + "expected C, A, and B scalar register operands" + ); + assert_eq!( + body_op.deref(&ctx).get_num_results(), + 1, + "LLVM inline asm returns the four D registers as one struct" + ); + } + } + } + + assert_eq!(found, 1, "expected one mma.sync inline-asm operation"); + Ok(()) +} + +#[test] +fn test_mma_m16n8k16_f32_f16_lowers_to_inline_asm() -> Result<(), anyhow::Error> { + let mut ctx = make_test_ctx(); + let f32_ty = pliron::builtin::types::FP32Type::get(&ctx); + let i32_ty = pliron::builtin::types::IntegerType::get( + &ctx, + 32, + pliron::builtin::types::Signedness::Signless, + ); + let argument_types = (0..4) + .map(|_| f32_ty.into()) + .chain((0..6).map(|_| i32_ty.into())) + .collect(); + let (module_ptr, entry) = build_test_kernel(&mut ctx, argument_types); + let operands = (0..10) + .map(|index| entry.deref(&ctx).get_argument(index)) + .collect(); + + let op = Operation::new( + &mut ctx, + nvvm::MmaM16N8K16F32F16Op::get_concrete_op_info(), + vec![f32_ty.into(); 4], + operands, + vec![], + 0, + ); + op.insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let mut found = 0; + let module_region = module_ptr.deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + for module_op in module_block.deref(&ctx).iter(&ctx) { + let Some(function) = Operation::get_op::(module_op, &ctx) else { + continue; + }; + if function.get_symbol_name(&ctx).to_string() != "kernel_func" { + continue; + } + let body = function.get_operation().deref(&ctx).get_region(0); + for block in body.deref(&ctx).iter(&ctx) { + for body_op in block.deref(&ctx).iter(&ctx) { + let Some(asm) = Operation::get_op::(body_op, &ctx) else { + continue; + }; + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())); + let constraints = asm + .get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())); + if !template.as_deref().is_some_and(|t| { + t.contains("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32") + }) { + continue; + } + found += 1; + let template = template.expect("MMA inline asm must have a template"); + assert_eq!( + template, + concat!( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 ", + "{$0, $1, $2, $3}, ", + "{$8, $9, $10, $11}, ", + "{$12, $13}, ", + "{$4, $5, $6, $7};" + ) + ); + for forbidden in [".reg", "ld.", "st.", "["] { + assert!( + !template.contains(forbidden), + "register-only MMA must not contain {forbidden:?}: {template}" + ); + } + assert_eq!( + constraints.as_deref(), + Some("=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r") + ); + assert_eq!( + llvm::asm_kind_opt(&ctx, &asm), + Some(llvm::AsmKind::Convergent) + ); + assert_eq!( + body_op.deref(&ctx).get_num_operands(), + 10, + "expected C, A, and B scalar register operands" + ); + assert_eq!( + body_op.deref(&ctx).get_num_results(), + 1, + "LLVM inline asm returns the four D registers as one struct" + ); + } + } + } + + assert_eq!(found, 1, "expected one mma.sync inline-asm operation"); + Ok(()) +} + +#[test] +fn test_mma_m16n8k8_f32_tf32_lowers_to_inline_asm() -> Result<(), anyhow::Error> { + let mut ctx = make_test_ctx(); + let f32_ty = pliron::builtin::types::FP32Type::get(&ctx); + let i32_ty = pliron::builtin::types::IntegerType::get( + &ctx, + 32, + pliron::builtin::types::Signedness::Signless, + ); + let argument_types = (0..4) + .map(|_| f32_ty.into()) + .chain((0..6).map(|_| i32_ty.into())) + .collect(); + let (module_ptr, entry) = build_test_kernel(&mut ctx, argument_types); + let operands = (0..10) + .map(|index| entry.deref(&ctx).get_argument(index)) + .collect(); + + let op = Operation::new( + &mut ctx, + nvvm::MmaM16N8K8F32Tf32Op::get_concrete_op_info(), + vec![f32_ty.into(); 4], + operands, + vec![], + 0, + ); + op.insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let mut found = 0; + let module_region = module_ptr.deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + for module_op in module_block.deref(&ctx).iter(&ctx) { + let Some(function) = Operation::get_op::(module_op, &ctx) else { + continue; + }; + if function.get_symbol_name(&ctx).to_string() != "kernel_func" { + continue; + } + let body = function.get_operation().deref(&ctx).get_region(0); + for block in body.deref(&ctx).iter(&ctx) { + for body_op in block.deref(&ctx).iter(&ctx) { + let Some(asm) = Operation::get_op::(body_op, &ctx) else { + continue; + }; + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())); + let constraints = asm + .get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())); + if !template.as_deref().is_some_and(|t| { + t.contains("mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32") + }) { + continue; + } + found += 1; + let template = template.expect("MMA inline asm must have a template"); + assert_eq!( + template, + concat!( + "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 ", + "{$0, $1, $2, $3}, ", + "{$8, $9, $10, $11}, ", + "{$12, $13}, ", + "{$4, $5, $6, $7};" + ) + ); + for forbidden in [".reg", "ld.", "st.", "["] { + assert!( + !template.contains(forbidden), + "register-only MMA must not contain {forbidden:?}: {template}" + ); + } + assert_eq!( + constraints.as_deref(), + Some("=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r") + ); + assert_eq!( + llvm::asm_kind_opt(&ctx, &asm), + Some(llvm::AsmKind::Convergent) + ); + assert_eq!( + body_op.deref(&ctx).get_num_operands(), + 10, + "expected C, A, and B scalar register operands" + ); + assert_eq!( + body_op.deref(&ctx).get_num_results(), + 1, + "LLVM inline asm returns the four D registers as one struct" + ); + } + } + } + + assert_eq!(found, 1, "expected one mma.sync inline-asm operation"); + Ok(()) +} + +#[test] +fn test_mma_m8n8k4_f64_lowers_to_inline_asm() -> Result<(), anyhow::Error> { + use pliron::builtin::types::FP64Type; + + let mut ctx = make_test_ctx(); + let f64_ty = FP64Type::get(&ctx); + let (module_ptr, entry) = build_test_kernel( + &mut ctx, + vec![f64_ty.into(), f64_ty.into(), f64_ty.into(), f64_ty.into()], + ); + let operands = (0..4) + .map(|index| entry.deref(&ctx).get_argument(index)) + .collect(); + + let op = Operation::new( + &mut ctx, + nvvm::MmaM8N8K4F64Op::get_concrete_op_info(), + vec![f64_ty.into(), f64_ty.into()], + operands, + vec![], + 0, + ); + op.insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let mut found = 0; + let module_region = module_ptr.deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + for module_op in module_block.deref(&ctx).iter(&ctx) { + let Some(function) = Operation::get_op::(module_op, &ctx) else { + continue; + }; + if function.get_symbol_name(&ctx).to_string() != "kernel_func" { + continue; + } + let body = function.get_operation().deref(&ctx).get_region(0); + for block in body.deref(&ctx).iter(&ctx) { + for body_op in block.deref(&ctx).iter(&ctx) { + let Some(asm) = Operation::get_op::(body_op, &ctx) else { + continue; + }; + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())); + let constraints = asm + .get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())); + if !template + .as_deref() + .is_some_and(|t| t.contains("mma.sync.aligned.m8n8k4.row.col.f64.f64.f64.f64")) + { + continue; + } + found += 1; + let template = template.unwrap(); + assert_eq!( + template, + "mma.sync.aligned.m8n8k4.row.col.f64.f64.f64.f64 {$0, $1}, {$4}, {$5}, {$2, $3};" + ); + assert!(!template.contains(".reg")); + assert!(!template.contains("ld.")); + assert!(!template.contains("st.")); + assert_eq!(constraints.as_deref(), Some("=d,=d,d,d,d,d")); + assert_eq!( + llvm::asm_kind_opt(&ctx, &asm), + Some(llvm::AsmKind::Convergent) + ); + assert_eq!( + body_op.deref(&ctx).get_num_operands(), + 4, + "expected four register inputs (c0, c1, a, b)" + ); + assert_eq!( + body_op.deref(&ctx).get_num_results(), + 1, + "LLVM inline asm returns one aggregate containing d0 and d1" + ); + } + } + } + + assert_eq!(found, 1, "expected one mma.sync inline-asm operation"); + Ok(()) +} + +#[test] +fn test_mma_m16n8k32_s32_s8_lowers_to_inline_asm() -> Result<(), anyhow::Error> { + let mut ctx = make_test_ctx(); + let i32_ty = pliron::builtin::types::IntegerType::get( + &ctx, + 32, + pliron::builtin::types::Signedness::Signless, + ); + let argument_types = (0..10).map(|_| i32_ty.into()).collect(); + let (module_ptr, entry) = build_test_kernel(&mut ctx, argument_types); + let operands = (0..10) + .map(|index| entry.deref(&ctx).get_argument(index)) + .collect(); + + let op = Operation::new( + &mut ctx, + nvvm::MmaM16N8K32S32S8Op::get_concrete_op_info(), + vec![i32_ty.into(); 4], + operands, + vec![], + 0, + ); + op.insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let mut found = 0; + let module_region = module_ptr.deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + for module_op in module_block.deref(&ctx).iter(&ctx) { + let Some(function) = Operation::get_op::(module_op, &ctx) else { + continue; + }; + if function.get_symbol_name(&ctx).to_string() != "kernel_func" { + continue; + } + let body = function.get_operation().deref(&ctx).get_region(0); + for block in body.deref(&ctx).iter(&ctx) { + for body_op in block.deref(&ctx).iter(&ctx) { + let Some(asm) = Operation::get_op::(body_op, &ctx) else { + continue; + }; + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())); + let constraints = asm + .get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())); + if !template + .as_deref() + .is_some_and(|t| t.contains("mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32")) + { + continue; + } + found += 1; + let template = template.expect("MMA inline asm must have a template"); + assert_eq!( + template, + concat!( + "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 ", + "{$0, $1, $2, $3}, ", + "{$8, $9, $10, $11}, ", + "{$12, $13}, ", + "{$4, $5, $6, $7};" + ) + ); + for forbidden in [".reg", "ld.", "st.", "["] { + assert!( + !template.contains(forbidden), + "register-only MMA must not contain {forbidden:?}: {template}" + ); + } + assert_eq!( + constraints.as_deref(), + Some("=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,r") + ); + assert_eq!( + llvm::asm_kind_opt(&ctx, &asm), + Some(llvm::AsmKind::Convergent) + ); + assert_eq!( + body_op.deref(&ctx).get_num_operands(), + 10, + "expected C, A, and B scalar register operands" + ); + assert_eq!( + body_op.deref(&ctx).get_num_results(), + 1, + "LLVM inline asm returns the four D registers as one struct" + ); + } + } + } + + assert_eq!(found, 1, "expected one mma.sync inline-asm operation"); + Ok(()) +} diff --git a/crates/mir-lower/tests/lowering_test/sregs_and_warp.rs b/crates/mir-lower/tests/lowering_test/sregs_and_warp.rs new file mode 100644 index 0000000000..1dafe7271a --- /dev/null +++ b/crates/mir-lower/tests/lowering_test/sregs_and_warp.rs @@ -0,0 +1,1618 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use dialect_mir::ops as mir; +use dialect_nvvm::ops as nvvm; +use llvm_export::ops as llvm; +use pliron::builtin::op_interfaces::{ + CallOpCallable, CallOpInterface, OneRegionInterface, SymbolOpInterface, +}; +use pliron::builtin::ops::ModuleOp; +use pliron::context::Context; +use pliron::linked_list::ContainsLinkedList; +use pliron::location::{Located, Location}; +use pliron::op::Op; +use pliron::operation::Operation; + +use crate::common::{append_return, build_test_kernel, lowered_kernel_body, make_test_ctx}; + +#[test] +fn test_standalone_lowering_rejects_builtin_pointer_constant() { + use dialect_mir::types::{MirPointerKind, MirPtrType}; + use pliron::builtin::{ + attributes::IntegerAttr, + ops::ConstantOp, + types::{IntegerType, Signedness}, + }; + use pliron::utils::apint::APInt; + use std::num::NonZeroUsize; + + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + mir_lower::register(&mut ctx); + + let module = ModuleOp::new(&mut ctx, "pointer_constant".try_into().unwrap()); + let block = module + .get_region(&ctx) + .deref(&ctx) + .iter(&ctx) + .next() + .unwrap(); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let pointer_ty = + MirPtrType::get_generic_with_kind(&mut ctx, u32_ty.into(), true, MirPointerKind::UniqueRef); + let value = APInt::from_u64(0, NonZeroUsize::new(32).unwrap()); + let constant = ConstantOp::new(&mut ctx, IntegerAttr::new(u32_ty, value).into()); + let result = constant.get_operation().deref(&ctx).get_result(0); + result.set_type(&ctx, pointer_ty.into()); + constant.get_operation().insert_at_back(block, &ctx); + + let error = mir_lower::lower_mir_to_llvm(&mut ctx, module.get_operation()).unwrap_err(); + assert!( + error + .to_string() + .contains("builtin.constant cannot produce a MIR pointer carrier") + ); +} + +#[test] +fn test_intrinsic_insertion() -> Result<(), anyhow::Error> { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + mir_lower::register(&mut ctx); + + // Create Module + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_ptr = module.get_operation(); + + // Create MirFunc + let func_name = "kernel_func"; + let func_ty = pliron::builtin::types::FunctionType::get(&ctx, vec![], vec![]); + + // Manual construction of MirFuncOp + let func_op_ptr = Operation::new( + &mut ctx, + mir::MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, // 1 region + ); + let func_ty_attr = pliron::builtin::attributes::TypeAttr::new(func_ty.into()); + let func = mir::MirFuncOp::new(&mut ctx, func_op_ptr, func_ty_attr); + func.set_symbol_name(&mut ctx, func_name.try_into().unwrap()); + + // Add body - MirFuncOp has 1 region + let region = func.get_operation().deref(&ctx).get_region(0); + + // Create block if empty (it is empty by default from Operation::new) + let block = { + let b = pliron::basic_block::BasicBlock::new(&mut ctx, None, vec![]); + b.insert_at_back(region, &ctx); + b + }; + + // Add ReadPtxSregTidXOp + let int32_ty = pliron::builtin::types::IntegerType::get( + &ctx, + 32, + pliron::builtin::types::Signedness::Signless, + ); + + let tid_op_ptr = Operation::new( + &mut ctx, + nvvm::ReadPtxSregTidXOp::get_concrete_op_info(), + vec![int32_ty.into()], + vec![], + vec![], + 0, + ); + let tid_op = nvvm::ReadPtxSregTidXOp::new(tid_op_ptr); + let expected_location = Location::Named { + name: "source-tid-x".to_string(), + child_loc: Box::new(Location::Unknown), + }; + tid_op + .get_operation() + .deref_mut(&ctx) + .set_loc(expected_location.clone()); + tid_op.get_operation().insert_at_back(block, &ctx); + + // Add Return + let ret_op_ptr = Operation::new( + &mut ctx, + mir::MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + let ret_op = mir::MirReturnOp::new(ret_op_ptr); + ret_op.get_operation().insert_at_back(block, &ctx); + + // Add Func to Module + let module_region = module.get_operation().deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + func.get_operation().insert_at_back(module_block, &ctx); + + // Run DialectConversion-based lowering + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{}", e))?; + + // Verify result + let mut found_intrinsic = false; + let mut found_intrinsic_call = false; + let mut found_kernel = false; + + let module_op = module_ptr.deref(&ctx); + let region = module_op.get_region(0); + let block = region.deref(&ctx).iter(&ctx).next().unwrap(); + + for op in block.deref(&ctx).iter(&ctx) { + if let Some(func_op) = Operation::get_op::(op, &ctx) { + let name = func_op.get_symbol_name(&ctx).to_string(); + if name == "llvm_nvvm_read_ptx_sreg_tid_x" { + found_intrinsic = true; + // Intrinsic (declaration) should have 0 regions or empty region + let num_regions = func_op.get_operation().deref(&ctx).regions().count(); + if num_regions > 0 { + assert!( + func_op + .get_operation() + .deref(&ctx) + .get_region(0) + .deref(&ctx) + .iter(&ctx) + .next() + .is_none() + ); + } + } else if name == "kernel_func" { + found_kernel = true; + // Kernel should have body (1 region, not empty) + assert!(func_op.get_operation().deref(&ctx).regions().count() > 0); + assert!( + func_op + .get_operation() + .deref(&ctx) + .get_region(0) + .deref(&ctx) + .iter(&ctx) + .next() + .is_some() + ); + let kernel_region = func_op.get_operation().deref(&ctx).get_region(0); + for kernel_block in kernel_region.deref(&ctx).iter(&ctx) { + for body_op in kernel_block.deref(&ctx).iter(&ctx) { + let Some(call) = Operation::get_op::(body_op, &ctx) else { + continue; + }; + if matches!( + call.callee(&ctx), + CallOpCallable::Direct(symbol) + if symbol.to_string() == "llvm_nvvm_read_ptx_sreg_tid_x" + ) { + found_intrinsic_call = true; + assert_eq!(call.get_operation().deref(&ctx).loc(), expected_location); + } + } + } + } + } + } + + assert!(found_intrinsic, "Intrinsic function declaration not found"); + assert!(found_intrinsic_call, "Intrinsic call not found in kernel"); + assert!(found_kernel, "Kernel function not found"); + + Ok(()) +} + +#[test] +fn test_assertfail_orphaned_successor_block_is_removed() -> Result<(), anyhow::Error> { + use dialect_mir::types::MirPtrType; + + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + mir_lower::register(&mut ctx); + + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_ptr = module.get_operation(); + + let func_name = "kernel_func"; + let u8_ty = pliron::builtin::types::IntegerType::get( + &ctx, + 8, + pliron::builtin::types::Signedness::Unsigned, + ); + let ptr_ty = MirPtrType::get_generic(&mut ctx, u8_ty.into(), false); + let u32_ty = pliron::builtin::types::IntegerType::get( + &ctx, + 32, + pliron::builtin::types::Signedness::Unsigned, + ); + let u64_ty = pliron::builtin::types::IntegerType::get( + &ctx, + 64, + pliron::builtin::types::Signedness::Unsigned, + ); + let arg_tys: Vec = vec![ + ptr_ty.into(), + ptr_ty.into(), + u32_ty.into(), + ptr_ty.into(), + u64_ty.into(), + ]; + let func_ty = pliron::builtin::types::FunctionType::get(&ctx, arg_tys.clone(), vec![]); + + let func_op_ptr = Operation::new( + &mut ctx, + mir::MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let func_ty_attr = pliron::builtin::attributes::TypeAttr::new(func_ty.into()); + let func = mir::MirFuncOp::new(&mut ctx, func_op_ptr, func_ty_attr); + func.set_symbol_name(&mut ctx, func_name.try_into().unwrap()); + + let region = func.get_operation().deref(&ctx).get_region(0); + let entry = { + let b = pliron::basic_block::BasicBlock::new(&mut ctx, None, arg_tys); + b.insert_at_back(region, &ctx); + b + }; + // A block reachable only through the assert's success edge, carrying a + // block argument: the CFG a constant-false `gpu_assert!` leaves behind + // after MIR optimization folds the direct edge away. + let tail = { + let b = pliron::basic_block::BasicBlock::new(&mut ctx, None, vec![u32_ty.into()]); + b.insert_at_back(region, &ctx); + b + }; + + let message = entry.deref(&ctx).get_argument(0); + let file = entry.deref(&ctx).get_argument(1); + let line = entry.deref(&ctx).get_argument(2); + let function = entry.deref(&ctx).get_argument(3); + let char_size = entry.deref(&ctx).get_argument(4); + + let assertfail_op = + nvvm::AssertFailOp::build(&mut ctx, message, file, line, function, char_size); + assertfail_op.insert_at_back(entry, &ctx); + + // The success edge forwards `line` into the tail block's argument. + // Lowering erases this branch (the call never returns), orphaning the + // tail block, which then cannot be exported: a block argument needs one + // PHI incoming value per predecessor and there are none left. + let goto_op = Operation::new( + &mut ctx, + mir::MirGotoOp::get_concrete_op_info(), + vec![], + vec![line], + vec![tail], + 0, + ); + goto_op.insert_at_back(entry, &ctx); + + let ret_op = Operation::new( + &mut ctx, + mir::MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + ret_op.insert_at_back(tail, &ctx); + + let module_region = module.get_operation().deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + func.get_operation().insert_at_back(module_block, &ctx); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let mut found_kernel = false; + let module_op = module_ptr.deref(&ctx); + let region = module_op.get_region(0); + let block = region.deref(&ctx).iter(&ctx).next().unwrap(); + for op in block.deref(&ctx).iter(&ctx) { + let Some(func_op) = Operation::get_op::(op, &ctx) else { + continue; + }; + if func_op.get_symbol_name(&ctx).to_string() != func_name { + continue; + } + found_kernel = true; + let func_region = func_op.get_operation().deref(&ctx).get_region(0); + let blocks: Vec<_> = func_region.deref(&ctx).iter(&ctx).collect(); + for block in blocks.iter().skip(1) { + assert!( + !block.preds(&ctx).is_empty(), + "an unreachable block survived lowering; with block arguments \ + it cannot be exported (a PHI needs an incoming value per \ + predecessor and there are none)" + ); + } + } + assert!(found_kernel, "Kernel function not found"); + + let module = Operation::get_op::(module_ptr, &ctx).unwrap(); + let ir = + llvm_export::export::export_module_to_string(&ctx, &module).map_err(anyhow::Error::msg)?; + let lines: Vec<_> = ir.lines().collect(); + let call_line = lines + .iter() + .position(|line| line.contains("call void @__assertfail(")) + .expect("expected exported __assertfail call"); + assert_eq!(lines[call_line + 1].trim(), "unreachable", "{ir}"); + + Ok(()) +} + +#[test] +fn test_globaltimer_lowers_to_intrinsic_call() -> Result<(), anyhow::Error> { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + mir_lower::register(&mut ctx); + + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_ptr = module.get_operation(); + + let func_name = "kernel_func"; + let func_ty = pliron::builtin::types::FunctionType::get(&ctx, vec![], vec![]); + + let func_op_ptr = Operation::new( + &mut ctx, + mir::MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let func_ty_attr = pliron::builtin::attributes::TypeAttr::new(func_ty.into()); + let func = mir::MirFuncOp::new(&mut ctx, func_op_ptr, func_ty_attr); + func.set_symbol_name(&mut ctx, func_name.try_into().unwrap()); + + let region = func.get_operation().deref(&ctx).get_region(0); + let block = { + let b = pliron::basic_block::BasicBlock::new(&mut ctx, None, vec![]); + b.insert_at_back(region, &ctx); + b + }; + + let i64_ty = pliron::builtin::types::IntegerType::get( + &ctx, + 64, + pliron::builtin::types::Signedness::Signless, + ); + let timer_op = Operation::new( + &mut ctx, + nvvm::ReadPtxSregGlobaltimerOp::get_concrete_op_info(), + vec![i64_ty.into()], + vec![], + vec![], + 0, + ); + timer_op.insert_at_back(block, &ctx); + + let ret_op_ptr = Operation::new( + &mut ctx, + mir::MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + let ret_op = mir::MirReturnOp::new(ret_op_ptr); + ret_op.get_operation().insert_at_back(block, &ctx); + + let module_region = module.get_operation().deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + func.get_operation().insert_at_back(module_block, &ctx); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{}", e))?; + + const INTRINSIC: &str = "llvm_nvvm_read_ptx_sreg_globaltimer"; + + let mut found_decl = false; + let mut found_call = false; + let module_op = module_ptr.deref(&ctx); + let region = module_op.get_region(0); + let block = region.deref(&ctx).iter(&ctx).next().unwrap(); + + for op in block.deref(&ctx).iter(&ctx) { + let Some(func_op) = Operation::get_op::(op, &ctx) else { + continue; + }; + let name = func_op.get_symbol_name(&ctx).to_string(); + + if name == INTRINSIC { + // Intrinsic declaration: present with empty body. + found_decl = true; + let num_regions = func_op.get_operation().deref(&ctx).regions().count(); + if num_regions > 0 { + assert!( + func_op + .get_operation() + .deref(&ctx) + .get_region(0) + .deref(&ctx) + .iter(&ctx) + .next() + .is_none(), + "intrinsic declaration must have empty body" + ); + } + } else if name == func_name { + let func_region = func_op.get_operation().deref(&ctx).get_region(0); + for func_block in func_region.deref(&ctx).iter(&ctx) { + for body_op in func_block.deref(&ctx).iter(&ctx) { + if let Some(call) = Operation::get_op::(body_op, &ctx) + && let CallOpCallable::Direct(sym) = call.callee(&ctx) + && sym.to_string() == INTRINSIC + { + found_call = true; + } + assert!( + Operation::get_op::(body_op, &ctx).is_none(), + "globaltimer must not lower to inline asm" + ); + } + } + } + } + + assert!( + found_decl, + "Expected `{INTRINSIC}` declaration in lowered module" + ); + assert!( + found_call, + "Expected call to `{INTRINSIC}` in lowered kernel body" + ); + Ok(()) +} + +#[test] +fn test_assertfail_lowers_to_noreturn_call_and_unreachable() -> Result<(), anyhow::Error> { + use dialect_mir::types::MirPtrType; + + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + mir_lower::register(&mut ctx); + + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_ptr = module.get_operation(); + + let func_name = "kernel_func"; + let u8_ty = pliron::builtin::types::IntegerType::get( + &ctx, + 8, + pliron::builtin::types::Signedness::Unsigned, + ); + let ptr_ty = MirPtrType::get_generic(&mut ctx, u8_ty.into(), false); + let u32_ty = pliron::builtin::types::IntegerType::get( + &ctx, + 32, + pliron::builtin::types::Signedness::Unsigned, + ); + let u64_ty = pliron::builtin::types::IntegerType::get( + &ctx, + 64, + pliron::builtin::types::Signedness::Unsigned, + ); + let arg_tys: Vec = vec![ + ptr_ty.into(), + ptr_ty.into(), + u32_ty.into(), + ptr_ty.into(), + u64_ty.into(), + ]; + let func_ty = pliron::builtin::types::FunctionType::get(&ctx, arg_tys.clone(), vec![]); + + let func_op_ptr = Operation::new( + &mut ctx, + mir::MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let func_ty_attr = pliron::builtin::attributes::TypeAttr::new(func_ty.into()); + let func = mir::MirFuncOp::new(&mut ctx, func_op_ptr, func_ty_attr); + func.set_symbol_name(&mut ctx, func_name.try_into().unwrap()); + + let region = func.get_operation().deref(&ctx).get_region(0); + let block = { + let b = pliron::basic_block::BasicBlock::new(&mut ctx, None, arg_tys); + b.insert_at_back(region, &ctx); + b + }; + + let message = block.deref(&ctx).get_argument(0); + let file = block.deref(&ctx).get_argument(1); + let line = block.deref(&ctx).get_argument(2); + let function = block.deref(&ctx).get_argument(3); + let char_size = block.deref(&ctx).get_argument(4); + + let assertfail_op = + nvvm::AssertFailOp::build(&mut ctx, message, file, line, function, char_size); + assertfail_op.insert_at_back(block, &ctx); + + // This return must be removed by lowering because __assertfail never returns. + let ret_op_ptr = Operation::new( + &mut ctx, + mir::MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + ret_op_ptr.insert_at_back(block, &ctx); + + let module_region = module.get_operation().deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + func.get_operation().insert_at_back(module_block, &ctx); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + const EXTERN: &str = "__assertfail"; + + let mut found_decl = false; + let mut found_call = false; + let mut found_unreachable = false; + + let module_op = module_ptr.deref(&ctx); + let region = module_op.get_region(0); + let block = region.deref(&ctx).iter(&ctx).next().unwrap(); + + for op in block.deref(&ctx).iter(&ctx) { + let Some(func_op) = Operation::get_op::(op, &ctx) else { + continue; + }; + let name = func_op.get_symbol_name(&ctx).to_string(); + + if name == EXTERN { + found_decl = true; + assert!( + llvm::op_noreturn(&ctx, func_op.get_operation()), + "__assertfail declaration must be marked noreturn" + ); + continue; + } + + if name != func_name { + continue; + } + + let func_region = func_op.get_operation().deref(&ctx).get_region(0); + for func_block in func_region.deref(&ctx).iter(&ctx) { + let body_ops: Vec<_> = func_block.deref(&ctx).iter(&ctx).collect(); + + for (index, body_op) in body_ops.iter().copied().enumerate() { + if let Some(call) = Operation::get_op::(body_op, &ctx) + && let CallOpCallable::Direct(sym) = call.callee(&ctx) + && sym.to_string() == EXTERN + { + found_call = true; + + assert!( + llvm::op_noreturn(&ctx, body_op), + "__assertfail call must be marked noreturn" + ); + assert_eq!( + body_op.deref(&ctx).get_num_operands(), + 5, + "__assertfail call must forward all five operands" + ); + assert!( + body_ops.get(index + 1).is_some_and(|next| { + Operation::get_op::(*next, &ctx).is_some() + }), + "llvm.unreachable must immediately follow __assertfail" + ); + assert_eq!( + index + 2, + body_ops.len(), + "no operations may remain after llvm.unreachable" + ); + } + + if Operation::get_op::(body_op, &ctx).is_some() { + found_unreachable = true; + } + + assert!( + Operation::get_op::(body_op, &ctx).is_none(), + "nvvm.assertfail must be fully consumed by the lowering" + ); + } + } + } + + assert!( + found_decl, + "Expected `{EXTERN}` declaration in lowered module" + ); + assert!( + found_call, + "Expected call to `{EXTERN}` in lowered kernel body" + ); + assert!( + found_unreachable, + "__assertfail must terminate the block with llvm.unreachable" + ); + + let module = Operation::get_op::(module_ptr, &ctx).unwrap(); + let ir = + llvm_export::export::export_module_to_string(&ctx, &module).map_err(anyhow::Error::msg)?; + + let declaration = ir + .lines() + .find(|line| line.contains("declare void @__assertfail(")) + .expect("expected exported __assertfail declaration"); + assert!(declaration.contains("noreturn"), "{ir}"); + + let lines: Vec<_> = ir.lines().collect(); + let call_line = lines + .iter() + .position(|line| line.contains("call void @__assertfail(")) + .expect("expected exported __assertfail call"); + assert!(lines[call_line].contains("noreturn"), "{ir}"); + assert_eq!(lines[call_line + 1].trim(), "unreachable", "{ir}"); + + Ok(()) +} + +/// Lower a single zero-operand, i32-result special-register op and assert it +/// emits a declaration of and direct call to `intrinsic` (and no inline asm). +fn assert_sreg_i32_lowers_to_intrinsic( + op_info: ( + fn(pliron::context::Ptr) -> pliron::op::OpObj, + std::any::TypeId, + ), + intrinsic: &str, +) -> Result<(), anyhow::Error> { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + mir_lower::register(&mut ctx); + + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_ptr = module.get_operation(); + + let func_name = "kernel_func"; + let func_ty = pliron::builtin::types::FunctionType::get(&ctx, vec![], vec![]); + + let func_op_ptr = Operation::new( + &mut ctx, + mir::MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let func_ty_attr = pliron::builtin::attributes::TypeAttr::new(func_ty.into()); + let func = mir::MirFuncOp::new(&mut ctx, func_op_ptr, func_ty_attr); + func.set_symbol_name(&mut ctx, func_name.try_into().unwrap()); + + let region = func.get_operation().deref(&ctx).get_region(0); + let block = { + let b = pliron::basic_block::BasicBlock::new(&mut ctx, None, vec![]); + b.insert_at_back(region, &ctx); + b + }; + + let i32_ty = pliron::builtin::types::IntegerType::get( + &ctx, + 32, + pliron::builtin::types::Signedness::Signless, + ); + let sreg_op = Operation::new(&mut ctx, op_info, vec![i32_ty.into()], vec![], vec![], 0); + sreg_op.insert_at_back(block, &ctx); + + let ret_op_ptr = Operation::new( + &mut ctx, + mir::MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + let ret_op = mir::MirReturnOp::new(ret_op_ptr); + ret_op.get_operation().insert_at_back(block, &ctx); + + let module_region = module.get_operation().deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + func.get_operation().insert_at_back(module_block, &ctx); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{}", e))?; + + let mut found_decl = false; + let mut found_call = false; + let module_op = module_ptr.deref(&ctx); + let region = module_op.get_region(0); + let block = region.deref(&ctx).iter(&ctx).next().unwrap(); + + for op in block.deref(&ctx).iter(&ctx) { + let Some(func_op) = Operation::get_op::(op, &ctx) else { + continue; + }; + let name = func_op.get_symbol_name(&ctx).to_string(); + + if name == intrinsic { + found_decl = true; + } else if name == func_name { + let func_region = func_op.get_operation().deref(&ctx).get_region(0); + for func_block in func_region.deref(&ctx).iter(&ctx) { + for body_op in func_block.deref(&ctx).iter(&ctx) { + if let Some(call) = Operation::get_op::(body_op, &ctx) + && let CallOpCallable::Direct(sym) = call.callee(&ctx) + && sym.to_string() == intrinsic + { + found_call = true; + } + assert!( + Operation::get_op::(body_op, &ctx).is_none(), + "{intrinsic} must not lower to inline asm" + ); + } + } + } + } + + assert!( + found_decl, + "Expected `{intrinsic}` declaration in lowered module" + ); + assert!( + found_call, + "Expected call to `{intrinsic}` in lowered kernel body" + ); + Ok(()) +} + +fn assert_sreg_lowers_to_inline_asm( + op_info: ( + fn(pliron::context::Ptr) -> pliron::op::OpObj, + std::any::TypeId, + ), + result_width: u32, + expected_template: &str, + expected_constraints: &str, + expected_kind: llvm::AsmKind, +) -> Result<(), anyhow::Error> { + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let result_ty = IntegerType::get(&ctx, result_width, Signedness::Signless); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![]); + let sreg_op = Operation::new(&mut ctx, op_info, vec![result_ty.into()], vec![], vec![], 0); + sreg_op.insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let mut matches = 0usize; + let module_region = module_ptr.deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + for op in module_block.deref(&ctx).iter(&ctx) { + let Some(func_op) = Operation::get_op::(op, &ctx) else { + continue; + }; + if func_op.get_symbol_name(&ctx).to_string() != "kernel_func" { + continue; + } + + let func_region = func_op.get_operation().deref(&ctx).get_region(0); + for func_block in func_region.deref(&ctx).iter(&ctx) { + for body_op in func_block.deref(&ctx).iter(&ctx) { + let Some(inline_asm) = Operation::get_op::(body_op, &ctx) else { + continue; + }; + let template = inline_asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())); + if template.as_deref() != Some(expected_template) { + continue; + } + + matches += 1; + assert_eq!( + inline_asm + .get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some(expected_constraints) + ); + assert_eq!(llvm::asm_kind(&ctx, &inline_asm), expected_kind); + } + } + } + + assert_eq!(matches, 1, "expected one exact `{expected_template}` read"); + Ok(()) +} + +#[test] +fn test_lanemask_ops_lower_to_sreg_intrinsic_calls() -> Result<(), anyhow::Error> { + // Each lane-position mask op lowers to its matching read-only sreg intrinsic + // (underscores become dots on export: `..._lanemask_lt` -> `...lanemask.lt`). + assert_sreg_i32_lowers_to_intrinsic( + nvvm::ReadPtxSregLanemaskLtOp::get_concrete_op_info(), + "llvm_nvvm_read_ptx_sreg_lanemask_lt", + )?; + assert_sreg_i32_lowers_to_intrinsic( + nvvm::ReadPtxSregLanemaskLeOp::get_concrete_op_info(), + "llvm_nvvm_read_ptx_sreg_lanemask_le", + )?; + assert_sreg_i32_lowers_to_intrinsic( + nvvm::ReadPtxSregLanemaskEqOp::get_concrete_op_info(), + "llvm_nvvm_read_ptx_sreg_lanemask_eq", + )?; + assert_sreg_i32_lowers_to_intrinsic( + nvvm::ReadPtxSregLanemaskGeOp::get_concrete_op_info(), + "llvm_nvvm_read_ptx_sreg_lanemask_ge", + )?; + assert_sreg_i32_lowers_to_intrinsic( + nvvm::ReadPtxSregLanemaskGtOp::get_concrete_op_info(), + "llvm_nvvm_read_ptx_sreg_lanemask_gt", + )?; + Ok(()) +} + +#[test] +fn test_generated_vote_sync_family_lowers_to_exact_typed_intrinsics() -> Result<(), anyhow::Error> { + use pliron::builtin::types::{IntegerType, Signedness}; + use pliron::r#type::Typed; + + let mut ctx = make_test_ctx(); + let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![i32_ty.into(), i1_ty.into()]); + let mask = entry.deref(&ctx).get_argument(0); + let predicate = entry.deref(&ctx).get_argument(1); + + for vote in [ + nvvm::VoteSyncAllOp::build(&mut ctx, mask, predicate), + nvvm::VoteSyncAnyOp::build(&mut ctx, mask, predicate), + nvvm::VoteSyncBallotOp::build(&mut ctx, mask, predicate), + nvvm::VoteSyncUniOp::build(&mut ctx, mask, predicate), + ] { + vote.insert_at_back(entry, &ctx); + } + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let expected = [ + ("llvm_nvvm_vote_all_sync", 1), + ("llvm_nvvm_vote_any_sync", 1), + ("llvm_nvvm_vote_ballot_sync", 32), + ("llvm_nvvm_vote_uni_sync", 1), + ]; + let mut found = Vec::new(); + for op in lowered_kernel_body(&ctx, module_ptr) { + assert!( + Operation::get_op::(op, &ctx).is_none(), + "generated vote.sync operations must use typed LLVM intrinsics" + ); + let Some(call) = Operation::get_op::(op, &ctx) else { + continue; + }; + let CallOpCallable::Direct(callee) = call.callee(&ctx) else { + continue; + }; + let callee = callee.to_string(); + let Some((_, result_width)) = expected.iter().find(|(name, _)| *name == callee) else { + continue; + }; + + let call = call.get_operation().deref(&ctx); + assert_eq!(call.get_num_operands(), 2); + assert_eq!(call.get_num_results(), 1); + + let integer_shape = |value: pliron::value::Value| { + let ty = value.get_type(&ctx); + let ty = ty.deref(&ctx); + let integer = ty + .downcast_ref::() + .expect("vote.sync operands and results are integers"); + (integer.width(), integer.signedness()) + }; + assert_eq!( + [ + integer_shape(call.get_operand(0)), + integer_shape(call.get_operand(1)), + ], + [(32, Signedness::Signless), (1, Signedness::Signless),], + "{callee} must preserve [mask, predicate] operand order" + ); + assert_eq!( + integer_shape(call.get_result(0)), + (*result_width, Signedness::Signless), + "{callee} returned the wrong LLVM integer type" + ); + found.push((callee, *result_width)); + } + + found.sort(); + let mut expected = expected + .map(|(name, width)| (name.to_owned(), width)) + .to_vec(); + expected.sort(); + assert_eq!(found, expected); + Ok(()) +} + +fn lower_generated_active_mask( + backend: mir_lower::IntrinsicBackend, +) -> Result<(Context, pliron::context::Ptr), anyhow::Error> { + let mut ctx = make_test_ctx(); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![]); + nvvm::ActiveMaskOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: backend, + ..Default::default() + }, + ) + .map_err(|error| anyhow::anyhow!("{error}"))?; + Ok((ctx, module_ptr)) +} + +#[test] +fn test_generated_active_mask_llvm_nvptx_uses_typed_intrinsic() -> Result<(), anyhow::Error> { + use pliron::builtin::types::IntegerType; + use pliron::r#type::Typed; + + let (ctx, module_ptr) = lower_generated_active_mask(mir_lower::IntrinsicBackend::LlvmNvptx)?; + let mut found = 0; + for op in lowered_kernel_body(&ctx, module_ptr) { + assert!( + Operation::get_op::(op, &ctx).is_none(), + "LLVM-NVPTX active_mask must use the typed intrinsic" + ); + let Some(call) = Operation::get_op::(op, &ctx) else { + continue; + }; + let CallOpCallable::Direct(callee) = call.callee(&ctx) else { + continue; + }; + if callee.to_string() != "llvm_nvvm_activemask" { + continue; + } + + found += 1; + let call = op.deref(&ctx); + assert_eq!(call.get_num_operands(), 0); + assert_eq!(call.get_num_results(), 1); + let result_ty = call.get_result(0).get_type(&ctx); + let result_ty = result_ty.deref(&ctx); + let result_ty = result_ty + .downcast_ref::() + .expect("active_mask returns an integer"); + assert_eq!(result_ty.width(), 32); + } + assert_eq!(found, 1, "expected one typed active_mask call"); + + let module = Operation::get_op::(module_ptr, &ctx).unwrap(); + let ir = llvm_export::export::export_module_to_string(&ctx, &module) + .map_err(|error| anyhow::anyhow!(error))?; + assert!(ir.contains("call i32 @llvm.nvvm.activemask()"), "{ir}"); + Ok(()) +} + +#[test] +fn test_generated_active_mask_libnvvm_uses_convergent_sideeffect_asm() -> Result<(), anyhow::Error> +{ + use pliron::builtin::types::IntegerType; + use pliron::r#type::Typed; + + let (ctx, module_ptr) = lower_generated_active_mask(mir_lower::IntrinsicBackend::LibNvvm)?; + let mut found = 0; + for op in lowered_kernel_body(&ctx, module_ptr) { + if let Some(call) = Operation::get_op::(op, &ctx) + && let CallOpCallable::Direct(callee) = call.callee(&ctx) + { + assert_ne!(callee.to_string(), "llvm_nvvm_activemask"); + } + let Some(asm) = Operation::get_op::(op, &ctx) else { + continue; + }; + + found += 1; + assert_eq!( + asm.get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some("activemask.b32 $0;") + ); + assert_eq!( + asm.get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some("=r,~{memory}") + ); + assert_eq!(llvm::asm_kind(&ctx, &asm), llvm::AsmKind::Convergent); + assert!( + asm.get_attr_inline_asm_convergent(&ctx) + .is_some_and(|value| bool::from((*value).clone())) + ); + let asm = op.deref(&ctx); + assert_eq!(asm.get_num_operands(), 0); + assert_eq!(asm.get_num_results(), 1); + let result_ty = asm.get_result(0).get_type(&ctx); + let result_ty = result_ty.deref(&ctx); + let result_ty = result_ty + .downcast_ref::() + .expect("active_mask inline asm returns an integer"); + assert_eq!(result_ty.width(), 32); + } + assert_eq!(found, 1, "expected one exact active_mask asm block"); + + let module = Operation::get_op::(module_ptr, &ctx).unwrap(); + let ir = llvm_export::export::export_module_to_string(&ctx, &module) + .map_err(|error| anyhow::anyhow!(error))?; + assert!( + ir.contains("call i32 asm sideeffect \"activemask.b32 $0;\", \"=r,~{memory}\"()"), + "{ir}" + ); + assert!(ir.contains("attributes #0 = { convergent }"), "{ir}"); + Ok(()) +} + +#[test] +fn test_generated_warp_match_family_uses_exact_typed_calls_and_mask_projection() +-> Result<(), anyhow::Error> { + use llvm_export::types::StructType; + use pliron::builtin::types::{IntegerType, Signedness}; + use pliron::r#type::Typed; + + let mut ctx = make_test_ctx(); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); + let (module_ptr, entry) = + build_test_kernel(&mut ctx, vec![i32_ty.into(), i32_ty.into(), i64_ty.into()]); + let mask = entry.deref(&ctx).get_argument(0); + let value32 = entry.deref(&ctx).get_argument(1); + let value64 = entry.deref(&ctx).get_argument(2); + for warp_match in [ + nvvm::MatchAnySyncI32Op::build(&mut ctx, mask, value32), + nvvm::MatchAnySyncI64Op::build(&mut ctx, mask, value64), + nvvm::MatchAllSyncI32Op::build(&mut ctx, mask, value32), + nvvm::MatchAllSyncI64Op::build(&mut ctx, mask, value64), + ] { + warp_match.insert_at_back(entry, &ctx); + } + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: mir_lower::IntrinsicBackend::LlvmNvptx, + ..Default::default() + }, + ) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let expected = [ + ("llvm_nvvm_match_any_sync_i32", 32, false), + ("llvm_nvvm_match_any_sync_i64", 64, false), + ("llvm_nvvm_match_all_sync_i32p", 32, true), + ("llvm_nvvm_match_all_sync_i64p", 64, true), + ]; + let body = lowered_kernel_body(&ctx, module_ptr); + let integer_width = |value: pliron::value::Value| { + let ty = value.get_type(&ctx); + let ty = ty.deref(&ctx); + ty.downcast_ref::() + .expect("warp-match value is an integer") + .width() + }; + let mut found = Vec::new(); + let mut aggregate_results = Vec::new(); + for &op in &body { + assert!( + Operation::get_op::(op, &ctx).is_none(), + "warp match must use typed LLVM intrinsics" + ); + let Some(call) = Operation::get_op::(op, &ctx) else { + continue; + }; + let CallOpCallable::Direct(callee) = call.callee(&ctx) else { + continue; + }; + let callee = callee.to_string(); + let Some((_, value_width, aggregate)) = + expected.iter().find(|(name, _, _)| *name == callee) + else { + continue; + }; + + let call = op.deref(&ctx); + assert_eq!(call.get_num_operands(), 2); + assert_eq!( + [ + integer_width(call.get_operand(0)), + integer_width(call.get_operand(1)), + ], + [32, *value_width], + "{callee} has the wrong typed signature" + ); + assert_eq!(call.get_num_results(), 1); + let result = call.get_result(0); + if *aggregate { + let result_ty = result.get_type(&ctx); + let result_ty = result_ty.deref(&ctx); + let result_ty = result_ty + .downcast_ref::() + .expect("match.all returns an LLVM aggregate"); + assert_eq!(result_ty.num_fields(), 2); + let field_widths = (0..2) + .map(|index| { + let field = result_ty.field_type(index); + let field = field.deref(&ctx); + field + .downcast_ref::() + .expect("match.all aggregate fields are integers") + .width() + }) + .collect::>(); + assert_eq!(field_widths, [32, 1]); + aggregate_results.push((callee.clone(), result)); + } else { + assert_eq!(integer_width(result), 32); + } + found.push(callee); + } + + found.sort(); + let mut expected_calls = expected.map(|(name, _, _)| name.to_owned()); + expected_calls.sort(); + assert_eq!(found, expected_calls); + + let mut projected = Vec::new(); + for &op in &body { + let Some(extract) = Operation::get_op::(op, &ctx) else { + continue; + }; + assert_eq!(extract.indices(&ctx), vec![0]); + let extract = op.deref(&ctx); + assert_eq!(extract.get_num_operands(), 1); + assert_eq!(extract.get_num_results(), 1); + assert_eq!(integer_width(extract.get_result(0)), 32); + let aggregate = extract.get_operand(0); + let callee = aggregate_results + .iter() + .find_map(|(callee, result)| (*result == aggregate).then(|| callee.clone())) + .expect("match.all must extract from its aggregate call result"); + projected.push(callee); + } + projected.sort(); + assert_eq!( + projected, + [ + "llvm_nvvm_match_all_sync_i32p".to_owned(), + "llvm_nvvm_match_all_sync_i64p".to_owned(), + ] + ); + Ok(()) +} + +#[test] +fn test_warpid_ops_preserve_snapshot_semantics() -> Result<(), anyhow::Error> { + assert_sreg_lowers_to_inline_asm( + nvvm::ReadPtxSregWarpIdOp::get_concrete_op_info(), + 32, + "mov.u32 $0, %warpid;", + "=r", + llvm::AsmKind::SideEffect, + )?; + assert_sreg_i32_lowers_to_intrinsic( + nvvm::ReadPtxSregNwarpIdOp::get_concrete_op_info(), + "llvm_nvvm_read_ptx_sreg_nwarpid", + )?; + Ok(()) +} + +#[test] +fn test_smid_ops_preserve_snapshot_semantics() -> Result<(), anyhow::Error> { + assert_sreg_lowers_to_inline_asm( + nvvm::ReadPtxSregSmIdOp::get_concrete_op_info(), + 32, + "mov.u32 $0, %smid;", + "=r", + llvm::AsmKind::SideEffect, + )?; + assert_sreg_i32_lowers_to_intrinsic( + nvvm::ReadPtxSregNsmIdOp::get_concrete_op_info(), + "llvm_nvvm_read_ptx_sreg_nsmid", + )?; + Ok(()) +} + +#[test] +fn test_repeated_location_samples_remain_side_effecting_reads() -> Result<(), anyhow::Error> { + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![]); + + for op_info in [ + nvvm::ReadPtxSregWarpIdOp::get_concrete_op_info(), + nvvm::ReadPtxSregSmIdOp::get_concrete_op_info(), + nvvm::ReadPtxSregWarpIdOp::get_concrete_op_info(), + nvvm::ReadPtxSregSmIdOp::get_concrete_op_info(), + ] { + let op = Operation::new(&mut ctx, op_info, vec![i32_ty.into()], vec![], vec![], 0); + op.insert_at_back(entry, &ctx); + } + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let mut warpid_reads = 0usize; + let mut smid_reads = 0usize; + let module_region = module_ptr.deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + for op in module_block.deref(&ctx).iter(&ctx) { + let Some(func_op) = Operation::get_op::(op, &ctx) else { + continue; + }; + if func_op.get_symbol_name(&ctx).to_string() != "kernel_func" { + continue; + } + let func_region = func_op.get_operation().deref(&ctx).get_region(0); + for func_block in func_region.deref(&ctx).iter(&ctx) { + for body_op in func_block.deref(&ctx).iter(&ctx) { + let Some(inline_asm) = Operation::get_op::(body_op, &ctx) else { + continue; + }; + let template = inline_asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())); + match template.as_deref() { + Some("mov.u32 $0, %warpid;") => warpid_reads += 1, + Some("mov.u32 $0, %smid;") => smid_reads += 1, + _ => continue, + } + assert_eq!( + llvm::asm_kind(&ctx, &inline_asm), + llvm::AsmKind::SideEffect, + "location snapshots must survive LLVM CSE" + ); + } + } + } + + assert_eq!(warpid_reads, 2); + assert_eq!(smid_reads, 2); + Ok(()) +} + +#[test] +fn test_gridid_op_lowers_to_full_width_inline_ptx() -> Result<(), anyhow::Error> { + assert_sreg_lowers_to_inline_asm( + nvvm::ReadPtxSregGridIdOp::get_concrete_op_info(), + 64, + "mov.u64 $0, %gridid;", + "=l", + llvm::AsmKind::Pure, + ) +} + +#[test] +fn test_smem_size_ops_lower_to_portable_inline_ptx() -> Result<(), anyhow::Error> { + assert_sreg_lowers_to_inline_asm( + nvvm::ReadPtxSregDynamicSmemSizeOp::get_concrete_op_info(), + 32, + "mov.u32 $0, %dynamic_smem_size;", + "=r", + llvm::AsmKind::Pure, + )?; + assert_sreg_lowers_to_inline_asm( + nvvm::ReadPtxSregTotalSmemSizeOp::get_concrete_op_info(), + 32, + "mov.u32 $0, %total_smem_size;", + "=r", + llvm::AsmKind::Pure, + )?; + Ok(()) +} + +#[test] +fn generated_threadfences_use_typed_intrinsics_on_both_backends() -> Result<(), anyhow::Error> { + const EXPECTED: [&str; 3] = [ + "llvm_nvvm_membar_cta", + "llvm_nvvm_membar_gl", + "llvm_nvvm_membar_sys", + ]; + + for backend in [ + mir_lower::IntrinsicBackend::LlvmNvptx, + mir_lower::IntrinsicBackend::LibNvvm, + ] { + let mut ctx = make_test_ctx(); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![]); + for op_info in [ + nvvm::ThreadfenceBlockOp::get_concrete_op_info(), + nvvm::ThreadfenceOp::get_concrete_op_info(), + nvvm::ThreadfenceSystemOp::get_concrete_op_info(), + ] { + Operation::new(&mut ctx, op_info, vec![], vec![], vec![], 0) + .insert_at_back(entry, &ctx); + } + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: backend, + ..Default::default() + }, + ) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let mut found = Vec::new(); + for op in lowered_kernel_body(&ctx, module_ptr) { + assert!( + Operation::get_op::(op, &ctx).is_none(), + "threadfences must use their reviewed typed route" + ); + let Some(call) = Operation::get_op::(op, &ctx) else { + continue; + }; + let CallOpCallable::Direct(callee) = call.callee(&ctx) else { + continue; + }; + let callee = callee.to_string(); + if EXPECTED.contains(&callee.as_str()) { + assert_eq!(op.deref(&ctx).get_num_operands(), 0); + found.push(callee); + } + } + found.sort(); + assert_eq!(found, EXPECTED); + + let module = Operation::get_op::(module_ptr, &ctx).unwrap(); + let ir = llvm_export::export::export_module_to_string(&ctx, &module) + .map_err(|error| anyhow::anyhow!(error))?; + for symbol in [ + "llvm.nvvm.membar.cta", + "llvm.nvvm.membar.gl", + "llvm.nvvm.membar.sys", + ] { + assert!(ir.contains(&format!("call void @{symbol}()")), "{ir}"); + } + } + Ok(()) +} + +/// LLVM uses its typed intrinsic. libNVVM uses the reviewed inline-PTX fallback. +#[test] +fn generated_elect_sync_uses_the_selected_backend_route() -> Result<(), anyhow::Error> { + use pliron::builtin::types::{IntegerType, Signedness}; + + for backend in [ + mir_lower::IntrinsicBackend::LlvmNvptx, + mir_lower::IntrinsicBackend::LibNvvm, + ] { + let mut ctx = make_test_ctx(); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); + let (module_ptr, entry) = build_test_kernel(&mut ctx, vec![i32_ty.into()]); + let mask = entry.deref(&ctx).get_argument(0); + Operation::new( + &mut ctx, + nvvm::ElectSyncOp::get_concrete_op_info(), + vec![i32_ty.into(), i1_ty.into()], + vec![mask], + vec![], + 0, + ) + .insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module_ptr, + mir_lower::LoweringOptions { + intrinsic_backend: backend, + ..Default::default() + }, + ) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let mut inline_asm_count = 0; + let mut typed_call_count = 0; + let mut extract_count = 0; + let mut trunc_count = 0; + for body_op in lowered_kernel_body(&ctx, module_ptr) { + if let Some(inline_asm) = Operation::get_op::(body_op, &ctx) { + assert_eq!( + inline_asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some("{ .reg .pred p; elect.sync $0|p, $2; selp.b32 $1, 1, 0, p; }") + ); + assert_eq!( + inline_asm + .get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some("=r,=r,r") + ); + assert!( + inline_asm + .get_attr_inline_asm_convergent(&ctx) + .is_some_and(|value| bool::from((*value).clone())) + ); + inline_asm_count += 1; + } + if let Some(call) = Operation::get_op::(body_op, &ctx) + && let CallOpCallable::Direct(callee) = call.callee(&ctx) + && callee.to_string() == "llvm_nvvm_elect_sync" + { + typed_call_count += 1; + } + if Operation::get_op::(body_op, &ctx).is_some() { + extract_count += 1; + } + if Operation::get_op::(body_op, &ctx).is_some() { + trunc_count += 1; + } + } + + assert_eq!(extract_count, 2); + match backend { + mir_lower::IntrinsicBackend::LlvmNvptx => { + assert_eq!(typed_call_count, 1); + assert_eq!(inline_asm_count, 0); + assert_eq!(trunc_count, 0); + } + mir_lower::IntrinsicBackend::LibNvvm => { + assert_eq!(typed_call_count, 0); + assert_eq!(inline_asm_count, 1); + assert_eq!(trunc_count, 1); + } + } + } + Ok(()) +} + +/// The exact inline-PTX template `convert_shuffle_i64` must emit for `mode`/`clamp`. +/// Mirrors the production `format!` so a drift in either side fails the test. +fn expected_shfl_i64_template(mode: &str, clamp: i32) -> String { + format!( + "{{ .reg .b32 lo; .reg .b32 hi; mov.b64 {{lo, hi}}, $1; \ + shfl.sync.{mode}.b32 lo, lo, $2, {clamp}, $3; \ + shfl.sync.{mode}.b32 hi, hi, $2, {clamp}, $3; \ + mov.b64 $0, {{lo, hi}}; }}" + ) +} + +/// 64-bit warp shuffle has no LLVM intrinsic (`shfl.sync` is 32-bit only), so it +/// lowers to convergent inline PTX that splits the value into two halves and runs +/// two `shfl.sync.*.b32`. Inline asm is opaque to LLVM, so a wrong mnemonic, +/// swapped operand order, wrong clamp, or missing `convergent` would only surface +/// as bad PTX downstream. This pins, for every mode, the exact template (incl. the +/// per-mode clamp: 31 for idx/bfly/down, 0 for up), the `=l,l,r,r` constraints, +/// and the convergent flag. +#[test] +fn test_shuffle_i64_lowers_to_inline_asm() -> Result<(), anyhow::Error> { + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless); + // Kernel args: [mask (i32), value (i64), lane/delta (i32)]. + let (module_ptr, entry) = + build_test_kernel(&mut ctx, vec![i32_ty.into(), i64_ty.into(), i32_ty.into()]); + let mask = entry.deref(&ctx).get_argument(0); + let value = entry.deref(&ctx).get_argument(1); + let lane = entry.deref(&ctx).get_argument(2); + + // One op per mode, all sharing the same [mask, value, lane] operands. + type OpInfo = ( + fn(pliron::context::Ptr) -> pliron::op::OpObj, + std::any::TypeId, + ); + let modes: [(OpInfo, &str, i32); 4] = [ + (nvvm::ShflSyncIdxI64Op::get_concrete_op_info(), "idx", 31), + (nvvm::ShflSyncBflyI64Op::get_concrete_op_info(), "bfly", 31), + (nvvm::ShflSyncDownI64Op::get_concrete_op_info(), "down", 31), + (nvvm::ShflSyncUpI64Op::get_concrete_op_info(), "up", 0), + ]; + for (opid, _, _) in modes { + let op = Operation::new( + &mut ctx, + opid, + vec![i64_ty.into()], + vec![mask, value, lane], + vec![], + 0, + ); + op.insert_at_back(entry, &ctx); + } + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).map_err(|e| anyhow::anyhow!("{}", e))?; + + // Collect every inline-asm template emitted into the kernel body. + let mut templates: Vec = Vec::new(); + let module_op = module_ptr.deref(&ctx); + let region = module_op.get_region(0); + let block = region.deref(&ctx).iter(&ctx).next().unwrap(); + for op in block.deref(&ctx).iter(&ctx) { + let Some(func_op) = Operation::get_op::(op, &ctx) else { + continue; + }; + if func_op.get_symbol_name(&ctx).to_string() != "kernel_func" { + continue; + } + let func_region = func_op.get_operation().deref(&ctx).get_region(0); + for func_block in func_region.deref(&ctx).iter(&ctx) { + for body_op in func_block.deref(&ctx).iter(&ctx) { + let Some(inline_asm) = Operation::get_op::(body_op, &ctx) else { + continue; + }; + assert_eq!( + inline_asm + .get_attr_inline_asm_constraints(&ctx) + .map(|s| String::from((*s).clone())) + .as_deref(), + Some("=l,l,r,r"), + "shfl.b64 constraints must be [out i64, value i64, lane i32, mask i32]" + ); + assert!( + inline_asm + .get_attr_inline_asm_convergent(&ctx) + .is_some_and(|b| bool::from((*b).clone())), + "shfl.b64 inline asm must be convergent" + ); + templates.push( + inline_asm + .get_attr_inline_asm_template(&ctx) + .map(|s| String::from((*s).clone())) + .unwrap_or_default(), + ); + } + } + } + + assert_eq!( + templates.len(), + 4, + "each of the 4 shfl.b64 modes must lower to one inline-asm op" + ); + for (_, mode, clamp) in modes { + let want = expected_shfl_i64_template(mode, clamp); + assert!( + templates.contains(&want), + "missing inline PTX for shfl.sync.{mode}.b32 (clamp {clamp}); got {templates:?}" + ); + } + + Ok(()) +} diff --git a/crates/mir-lower/tests/lowering_test/wgmma_lowering.rs b/crates/mir-lower/tests/lowering_test/wgmma_lowering.rs new file mode 100644 index 0000000000..3f024367a4 --- /dev/null +++ b/crates/mir-lower/tests/lowering_test/wgmma_lowering.rs @@ -0,0 +1,2658 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use dialect_mir::ops as mir; +use dialect_nvvm::ops as nvvm; +use llvm_export::ops as llvm; +use pliron::context::Context; +use pliron::linked_list::ContainsLinkedList; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::r#type::Typed; + +use crate::common::{append_return, build_test_kernel, lowered_kernel_body, make_test_ctx}; + +pub(super) fn build_wgmma_pointer_test_kernel( + ctx: &mut Context, + accumulator_count: usize, + trailing_arg_types: Vec, +) -> ( + pliron::context::Ptr, + pliron::context::Ptr, + Vec, + pliron::value::Value, + pliron::value::Value, + Vec, +) { + use dialect_mir::types::MirPtrType; + use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; + use pliron::r#type::TypeHandle; + + let f32_ty = FP32Type::get(ctx); + let accumulator_ptr_ty: TypeHandle = MirPtrType::get_generic(ctx, f32_ty.into(), true).into(); + let u64_ty: TypeHandle = IntegerType::get(ctx, 64, Signedness::Unsigned).into(); + + let trailing_count = trailing_arg_types.len(); + let mut argument_types = vec![accumulator_ptr_ty; accumulator_count]; + argument_types.push(u64_ty); + argument_types.push(u64_ty); + argument_types.extend(trailing_arg_types); + + let (module_ptr, entry) = build_test_kernel(ctx, argument_types); + + let accumulators = (0..accumulator_count) + .map(|index| entry.deref(ctx).get_argument(index)) + .collect::>(); + let desc_a = entry.deref(ctx).get_argument(accumulator_count); + let desc_b = entry.deref(ctx).get_argument(accumulator_count + 1); + let trailing_arguments = (0..trailing_count) + .map(|index| entry.deref(ctx).get_argument(accumulator_count + 2 + index)) + .collect::>(); + + ( + module_ptr, + entry, + accumulators, + desc_a, + desc_b, + trailing_arguments, + ) +} + +fn build_wgmma_canonical_pointer_test_kernel_with_rows( + ctx: &mut Context, + accumulator_rows: u64, + accumulator_count: usize, + descriptor_count: usize, +) -> ( + pliron::context::Ptr, + pliron::context::Ptr, + Vec, + Vec, +) { + use dialect_mir::types::{MirArrayType, MirPtrType}; + use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; + + let f32_ty = FP32Type::get(ctx); + let row_ty = MirArrayType::get(ctx, f32_ty.into(), 8); + let accumulator_ty = MirArrayType::get(ctx, row_ty.into(), accumulator_rows); + let accumulator_ptr_ty = MirPtrType::get_generic(ctx, accumulator_ty.into(), true); + let u64_ty = IntegerType::get(ctx, 64, Signedness::Unsigned); + + let accumulator_arg_ty: pliron::r#type::TypeHandle = accumulator_ptr_ty.into(); + let u64_arg_ty: pliron::r#type::TypeHandle = u64_ty.into(); + let mut argument_types = vec![accumulator_arg_ty; accumulator_count]; + argument_types.extend(vec![u64_arg_ty; descriptor_count]); + let (module_ptr, entry) = build_test_kernel(ctx, argument_types); + + let accumulators = (0..accumulator_count) + .map(|index| entry.deref(ctx).get_argument(index)) + .collect::>(); + let descriptors = (0..descriptor_count) + .map(|index| entry.deref(ctx).get_argument(accumulator_count + index)) + .collect::>(); + + (module_ptr, entry, accumulators, descriptors) +} + +pub(super) fn build_wgmma_canonical_pointer_test_kernel( + ctx: &mut Context, + accumulator_count: usize, + descriptor_count: usize, +) -> ( + pliron::context::Ptr, + pliron::context::Ptr, + Vec, + Vec, +) { + build_wgmma_canonical_pointer_test_kernel_with_rows(ctx, 4, accumulator_count, descriptor_count) +} + +fn build_wgmma_m64n128_canonical_pointer_test_kernel( + ctx: &mut Context, + accumulator_count: usize, + descriptor_count: usize, +) -> ( + pliron::context::Ptr, + pliron::context::Ptr, + Vec, + Vec, +) { + build_wgmma_canonical_pointer_test_kernel_with_rows(ctx, 8, accumulator_count, descriptor_count) +} + +pub(super) fn append_pointer_wgmma_mma( + ctx: &mut Context, + block: pliron::context::Ptr, + accumulator: pliron::value::Value, + desc_a: pliron::value::Value, + desc_b: pliron::value::Value, +) { + Operation::new( + ctx, + nvvm::WgmmaMmaM64N64K16F32Bf16Op::get_concrete_op_info(), + vec![], + vec![accumulator, desc_a, desc_b], + vec![], + 0, + ) + .insert_at_back(block, ctx); +} + +pub(super) fn append_pointer_wgmma_mma_m64n128( + ctx: &mut Context, + block: pliron::context::Ptr, + accumulator: pliron::value::Value, + desc_a: pliron::value::Value, + desc_b: pliron::value::Value, +) { + Operation::new( + ctx, + nvvm::WgmmaMmaM64N128K16F32Bf16Op::get_concrete_op_info(), + vec![], + vec![accumulator, desc_a, desc_b], + vec![], + 0, + ) + .insert_at_back(block, ctx); +} + +pub(super) fn append_pointer_wgmma_mma_f16( + ctx: &mut Context, + block: pliron::context::Ptr, + accumulator: pliron::value::Value, + desc_a: pliron::value::Value, + desc_b: pliron::value::Value, +) { + Operation::new( + ctx, + nvvm::WgmmaMmaM64N64K16F32F16Op::get_concrete_op_info(), + vec![], + vec![accumulator, desc_a, desc_b], + vec![], + 0, + ) + .insert_at_back(block, ctx); +} + +pub(super) fn append_pointer_wgmma_mma_tf32( + ctx: &mut Context, + block: pliron::context::Ptr, + accumulator: pliron::value::Value, + desc_a: pliron::value::Value, + desc_b: pliron::value::Value, +) { + Operation::new( + ctx, + nvvm::WgmmaMmaM64N64K8F32Tf32Op::get_concrete_op_info(), + vec![], + vec![accumulator, desc_a, desc_b], + vec![], + 0, + ) + .insert_at_back(block, ctx); +} + +pub(super) fn append_wgmma_wait_group_constant( + ctx: &mut Context, + block: pliron::context::Ptr, + pending_groups: i64, +) { + use pliron::builtin::attributes::IntegerAttr; + use pliron::builtin::types::{IntegerType, Signedness}; + use pliron::utils::apint::APInt; + use std::num::NonZeroUsize; + + let u64_ty = IntegerType::get(ctx, 64, Signedness::Unsigned); + let constant_attr = IntegerAttr::new( + u64_ty, + APInt::from_i64(pending_groups, NonZeroUsize::new(64).unwrap()), + ); + let constant = Operation::new( + ctx, + mir::MirConstantOp::get_concrete_op_info(), + vec![u64_ty.into()], + vec![], + vec![], + 0, + ); + mir::MirConstantOp::new(constant).set_attr_value(ctx, constant_attr); + constant.insert_at_back(block, ctx); + + let value = constant.deref(ctx).get_result(0); + nvvm::WgmmaWaitGroupSyncAlignedOp::build(ctx, value).insert_at_back(block, ctx); +} + +pub(super) fn append_mir_unsigned_constant( + ctx: &mut Context, + block: pliron::context::Ptr, + ty: pliron::r#type::TypedHandle, + value: u64, +) -> pliron::value::Value { + use pliron::builtin::attributes::IntegerAttr; + use pliron::utils::apint::APInt; + use std::num::NonZeroUsize; + + let width = usize::try_from(ty.deref(ctx).width()).expect("integer width must fit usize"); + let constant = Operation::new( + ctx, + mir::MirConstantOp::get_concrete_op_info(), + vec![ty.into()], + vec![], + vec![], + 0, + ); + mir::MirConstantOp::new(constant).set_attr_value( + ctx, + IntegerAttr::new( + ty, + APInt::from_u64( + value, + NonZeroUsize::new(width).expect("nonzero integer width"), + ), + ), + ); + constant.insert_at_back(block, ctx); + constant.deref(ctx).get_result(0) +} + +#[test] +fn test_deferred_wgmma_group_lowers_to_one_register_lifetime_scope() -> Result<(), anyhow::Error> { + use dialect_mir::types::MirPtrType; + use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let f32_ty = FP32Type::get(&ctx); + let accumulator_ptr_ty = MirPtrType::get_generic(&mut ctx, f32_ty.into(), true); + let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let (module_ptr, entry) = build_test_kernel( + &mut ctx, + vec![accumulator_ptr_ty.into(), u64_ty.into(), u64_ty.into()], + ); + let accumulator = entry.deref(&ctx).get_argument(0); + let desc_a = entry.deref(&ctx).get_argument(1); + let desc_b = entry.deref(&ctx).get_argument(2); + + nvvm::WgmmaMmaGroupM64N64K16F32Bf16Op::build(&mut ctx, accumulator, vec![desc_a, desc_b]) + .insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let matching = lowered_kernel_body(&ctx, module_ptr) + .into_iter() + .filter_map(|operation| Operation::get_op::(operation, &ctx)) + .filter(|asm| { + asm.get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .is_some_and(|template| { + template.contains("wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16") + }) + }) + .collect::>(); + assert_eq!(matching.len(), 1); + + let asm = &matching[0]; + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .expect("WGMMA template"); + assert_eq!(template.matches("ld.f32 %acc").count(), 32); + assert_eq!(template.matches("st.f32 [$0").count(), 32); + assert_eq!(template.matches("wgmma.mma_async").count(), 1); + let wait = template.find("wgmma.wait_group.sync.aligned 0").unwrap(); + let first_store = template.find("st.f32 [$0").unwrap(); + assert!( + wait < first_store, + "accumulator stores must follow wait_group<0>" + ); + assert_eq!( + asm.get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some("l,l,l,~{memory}") + ); + assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); + Ok(()) +} + +#[test] +fn test_value_form_wgmma_group_lowers_to_tied_register_inline_ptx() -> Result<(), anyhow::Error> { + use llvm_export::types as llvm_types; + use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; + use pliron::r#type::Typed; + + const ACCUMULATOR_LEN: usize = 32; + const DESCRIPTOR_COUNT: usize = 4; + + let mut ctx = make_test_ctx(); + let f32_ty = FP32Type::get(&ctx); + let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + + let argument_types = (0..ACCUMULATOR_LEN) + .map(|_| f32_ty.into()) + .chain((0..DESCRIPTOR_COUNT).map(|_| u64_ty.into())) + .collect::>(); + + let (module_ptr, entry) = build_test_kernel(&mut ctx, argument_types); + + let accumulators = (0..ACCUMULATOR_LEN) + .map(|index| entry.deref(&ctx).get_argument(index)) + .collect::>(); + let descriptors = (0..DESCRIPTOR_COUNT) + .map(|index| entry.deref(&ctx).get_argument(ACCUMULATOR_LEN + index)) + .collect::>(); + + nvvm::WgmmaMmaGroupValuesM64N64K16F32Bf16Op::build(&mut ctx, accumulators, descriptors) + .insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let matching = lowered_kernel_body(&ctx, module_ptr) + .into_iter() + .filter_map(|operation| Operation::get_op::(operation, &ctx)) + .filter(|asm| { + asm.get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .is_some_and(|template| { + template.contains("wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16") + }) + }) + .collect::>(); + assert_eq!(matching.len(), 1); + + let asm = &matching[0]; + let asm_op = asm.get_operation(); + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .expect("value-form WGMMA template"); + + assert_eq!(template.matches("wgmma.fence.sync.aligned").count(), 1); + assert_eq!(template.matches("wgmma.mma_async").count(), 2); + assert_eq!( + template.matches("wgmma.commit_group.sync.aligned").count(), + 1 + ); + assert_eq!( + template.matches("wgmma.wait_group.sync.aligned 0").count(), + 1 + ); + for forbidden in [".reg .f32", "ld.f32", "st.f32"] { + assert!( + !template.contains(forbidden), + "value-form WGMMA must not materialize accumulator memory via {forbidden:?}: {template}" + ); + } + + let accumulator_operands = (0..ACCUMULATOR_LEN) + .map(|index| format!("${index}")) + .collect::>() + .join(", "); + assert!( + template.contains(&format!("{{{accumulator_operands}}}")), + "WGMMA must consume the 32 output operands as one accumulator tuple: {template}" + ); + assert!(template.contains("$64, $65"), "{template}"); + assert!(template.contains("$66, $67"), "{template}"); + + let mut expected_constraints = vec!["=f".to_owned(); ACCUMULATOR_LEN]; + expected_constraints.extend((0..ACCUMULATOR_LEN).map(|index| index.to_string())); + expected_constraints.extend((0..DESCRIPTOR_COUNT).map(|_| "l".to_owned())); + expected_constraints.push("~{memory}".to_owned()); + let expected_constraints = expected_constraints.join(","); + + assert_eq!( + asm.get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some(expected_constraints.as_str()) + ); + assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); + + let asm_ref = asm_op.deref(&ctx); + assert_eq!( + asm_ref.get_num_operands(), + ACCUMULATOR_LEN + DESCRIPTOR_COUNT + ); + assert_eq!( + asm_ref.get_num_results(), + 1, + "LLVM inline asm must return the 32 WGMMA accumulator values as one struct" + ); + + let aggregate = asm_ref.get_result(0); + let aggregate_ty = aggregate.get_type(&ctx); + let aggregate_ty = aggregate_ty.deref(&ctx); + let struct_ty = aggregate_ty + .downcast_ref::() + .expect("value-form WGMMA inline asm must return an LLVM struct"); + assert_eq!(struct_ty.num_fields(), ACCUMULATOR_LEN); + for index in 0..ACCUMULATOR_LEN { + assert!( + struct_ty + .field_type(index) + .deref(&ctx) + .downcast_ref::() + .is_some(), + "WGMMA result field {index} must remain f32" + ); + } + + let extract_indices = lowered_kernel_body(&ctx, module_ptr) + .into_iter() + .filter_map(|operation| Operation::get_op::(operation, &ctx)) + .filter_map(|extract| { + let extract_op = extract.get_operation().deref(&ctx); + (extract_op.get_operand(0) == aggregate).then(|| extract.indices(&ctx)) + }) + .collect::>(); + + let expected_indices = (0..ACCUMULATOR_LEN) + .map(|index| vec![index as u32]) + .collect::>(); + assert_eq!( + extract_indices, expected_indices, + "all 32 value-form WGMMA results must be extracted in constraint order" + ); + + Ok(()) +} + +#[test] +fn test_value_form_m64n128_bf16_wgmma_group_lowers_to_sixty_four_tied_registers() +-> Result<(), anyhow::Error> { + use llvm_export::types as llvm_types; + use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; + use pliron::r#type::Typed; + + const ACCUMULATOR_LEN: usize = 64; + const DESCRIPTOR_COUNT: usize = 4; + + let mut ctx = make_test_ctx(); + let f32_ty = FP32Type::get(&ctx); + let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let argument_types = (0..ACCUMULATOR_LEN) + .map(|_| f32_ty.into()) + .chain((0..DESCRIPTOR_COUNT).map(|_| u64_ty.into())) + .collect::>(); + let (module_ptr, entry) = build_test_kernel(&mut ctx, argument_types); + + let accumulators = (0..ACCUMULATOR_LEN) + .map(|index| entry.deref(&ctx).get_argument(index)) + .collect::>(); + let descriptors = (0..DESCRIPTOR_COUNT) + .map(|index| entry.deref(&ctx).get_argument(ACCUMULATOR_LEN + index)) + .collect::>(); + + nvvm::WgmmaMmaGroupValuesM64N128K16F32Bf16Op::build(&mut ctx, accumulators, descriptors) + .insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let matching = lowered_kernel_body(&ctx, module_ptr) + .into_iter() + .filter_map(|operation| Operation::get_op::(operation, &ctx)) + .filter(|asm| { + asm.get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .is_some_and(|template| template.contains("m64n128k16.f32.bf16.bf16")) + }) + .collect::>(); + assert_eq!(matching.len(), 1); + + let asm = &matching[0]; + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .expect("m64n128 value-form WGMMA template"); + assert_eq!(template.matches("wgmma.mma_async").count(), 2); + assert!(template.contains("m64n128k16.f32.bf16.bf16")); + assert!(template.contains("$128, $129")); + assert!(template.contains("$130, $131")); + assert!(!template.contains("ld.f32")); + assert!(!template.contains("st.f32")); + + let constraints = asm + .get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .expect("m64n128 value-form WGMMA constraints"); + assert_eq!( + constraints + .split(',') + .filter(|value| *value == "=f") + .count(), + 64 + ); + assert_eq!( + constraints.split(',').filter(|value| *value == "l").count(), + 4 + ); + assert!(constraints.ends_with("~{memory}")); + + let aggregate = asm.get_operation().deref(&ctx).get_result(0); + let aggregate_ty = aggregate.get_type(&ctx); + let aggregate_ty = aggregate_ty.deref(&ctx); + let struct_ty = aggregate_ty + .downcast_ref::() + .expect("m64n128 inline asm must return an LLVM struct"); + assert_eq!(struct_ty.num_fields(), ACCUMULATOR_LEN); + assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); + + Ok(()) +} + +#[test] +fn test_value_form_f16_wgmma_group_lowers_to_tied_register_inline_ptx() -> Result<(), anyhow::Error> +{ + use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; + + const ACCUMULATOR_LEN: usize = 32; + const DESCRIPTOR_COUNT: usize = 2; + + let mut ctx = make_test_ctx(); + let f32_ty = FP32Type::get(&ctx); + let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let argument_types = (0..ACCUMULATOR_LEN) + .map(|_| f32_ty.into()) + .chain((0..DESCRIPTOR_COUNT).map(|_| u64_ty.into())) + .collect::>(); + let (module_ptr, entry) = build_test_kernel(&mut ctx, argument_types); + + let accumulators = (0..ACCUMULATOR_LEN) + .map(|index| entry.deref(&ctx).get_argument(index)) + .collect::>(); + let descriptors = (0..DESCRIPTOR_COUNT) + .map(|index| entry.deref(&ctx).get_argument(ACCUMULATOR_LEN + index)) + .collect::>(); + + nvvm::WgmmaMmaGroupValuesM64N64K16F32F16Op::build(&mut ctx, accumulators, descriptors) + .insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let matching = lowered_kernel_body(&ctx, module_ptr) + .into_iter() + .filter_map(|operation| Operation::get_op::(operation, &ctx)) + .filter(|asm| { + asm.get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .is_some_and(|template| template.contains("m64n64k16.f32.f16.f16")) + }) + .collect::>(); + assert_eq!(matching.len(), 1); + + let asm = &matching[0]; + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .expect("F16 value-form WGMMA template"); + assert_eq!(template.matches("wgmma.mma_async").count(), 1); + assert!(template.contains("m64n64k16.f32.f16.f16")); + assert!(!template.contains(".bf16.bf16")); + assert!(template.contains("wgmma.fence.sync.aligned")); + assert!(template.contains("wgmma.commit_group.sync.aligned")); + assert!(template.contains("wgmma.wait_group.sync.aligned 0")); + assert!(!template.contains("ld.f32")); + assert!(!template.contains("st.f32")); + + let constraints = asm + .get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .expect("F16 value-form WGMMA constraints"); + assert_eq!( + constraints + .split(',') + .filter(|value| *value == "=f") + .count(), + 32 + ); + assert_eq!( + constraints.split(',').filter(|value| *value == "l").count(), + 2 + ); + assert!(constraints.ends_with("~{memory}")); + assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); + + Ok(()) +} + +#[test] +fn test_value_form_tf32_wgmma_group_lowers_to_tied_register_inline_ptx() -> Result<(), anyhow::Error> +{ + use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; + + const ACCUMULATOR_LEN: usize = 32; + const DESCRIPTOR_COUNT: usize = 2; + + let mut ctx = make_test_ctx(); + let f32_ty = FP32Type::get(&ctx); + let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let argument_types = (0..ACCUMULATOR_LEN) + .map(|_| f32_ty.into()) + .chain((0..DESCRIPTOR_COUNT).map(|_| u64_ty.into())) + .collect::>(); + let (module_ptr, entry) = build_test_kernel(&mut ctx, argument_types); + + let accumulators = (0..ACCUMULATOR_LEN) + .map(|index| entry.deref(&ctx).get_argument(index)) + .collect::>(); + let descriptors = (0..DESCRIPTOR_COUNT) + .map(|index| entry.deref(&ctx).get_argument(ACCUMULATOR_LEN + index)) + .collect::>(); + + nvvm::WgmmaMmaGroupValuesM64N64K8F32Tf32Op::build(&mut ctx, accumulators, descriptors) + .insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let matching = lowered_kernel_body(&ctx, module_ptr) + .into_iter() + .filter_map(|operation| Operation::get_op::(operation, &ctx)) + .filter(|asm| { + asm.get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .is_some_and(|template| template.contains("m64n64k8.f32.tf32.tf32")) + }) + .collect::>(); + assert_eq!(matching.len(), 1); + + let asm = &matching[0]; + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .expect("TF32 value-form WGMMA template"); + assert_eq!(template.matches("wgmma.mma_async").count(), 1); + assert!(template.contains("m64n64k8.f32.tf32.tf32")); + assert!(!template.contains(".bf16.bf16")); + assert!(!template.contains(".f16.f16")); + assert!(!template.contains("m64n64k16")); + assert!(template.contains("$64, $65, 1, 1, 1;")); + assert!(!template.contains("$64, $65, 1, 1, 1, 0, 0;")); + assert!(template.contains("wgmma.fence.sync.aligned")); + assert!(template.contains("wgmma.commit_group.sync.aligned")); + assert!(template.contains("wgmma.wait_group.sync.aligned 0")); + assert!(!template.contains("ld.f32")); + assert!(!template.contains("st.f32")); + + let constraints = asm + .get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .expect("TF32 value-form WGMMA constraints"); + assert_eq!( + constraints + .split(',') + .filter(|value| *value == "=f") + .count(), + 32 + ); + assert_eq!( + constraints.split(',').filter(|value| *value == "l").count(), + 2 + ); + assert!(constraints.ends_with("~{memory}")); + assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); + + Ok(()) +} + +#[test] +fn test_pointer_form_wgmma_sequence_preserves_deferred_fallback() -> Result<(), anyhow::Error> { + use dialect_mir::types::MirPtrType; + use pliron::builtin::attributes::IntegerAttr; + use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; + use pliron::utils::apint::APInt; + use std::num::NonZeroUsize; + + let mut ctx = make_test_ctx(); + let f32_ty = FP32Type::get(&ctx); + // A direct *mut f32 is intentionally not the public [[f32; 8]; 4] + // accumulator shape. It must therefore retain the deferred pointer path. + let accumulator_ptr_ty = MirPtrType::get_generic(&mut ctx, f32_ty.into(), true); + let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let (module_ptr, entry) = build_test_kernel( + &mut ctx, + vec![accumulator_ptr_ty.into(), u64_ty.into(), u64_ty.into()], + ); + let accumulator = entry.deref(&ctx).get_argument(0); + let desc_a = entry.deref(&ctx).get_argument(1); + let desc_b = entry.deref(&ctx).get_argument(2); + + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + Operation::new( + &mut ctx, + nvvm::WgmmaMmaM64N64K16F32Bf16Op::get_concrete_op_info(), + vec![], + vec![accumulator, desc_a, desc_b], + vec![], + 0, + ) + .insert_at_back(entry, &ctx); + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + + let zero_attr = IntegerAttr::new(u64_ty, APInt::from_i64(0, NonZeroUsize::new(64).unwrap())); + let zero = Operation::new( + &mut ctx, + mir::MirConstantOp::get_concrete_op_info(), + vec![u64_ty.into()], + vec![], + vec![], + 0, + ); + mir::MirConstantOp::new(zero).set_attr_value(&ctx, zero_attr); + zero.insert_at_back(entry, &ctx); + let zero_value = zero.deref(&ctx).get_result(0); + nvvm::WgmmaWaitGroupSyncAlignedOp::build(&mut ctx, zero_value).insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let matching = lowered_kernel_body(&ctx, module_ptr) + .into_iter() + .filter_map(|operation| Operation::get_op::(operation, &ctx)) + .filter(|asm| { + asm.get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .is_some_and(|template| template.contains("wgmma.mma_async")) + }) + .collect::>(); + assert_eq!(matching.len(), 1); + + let asm = &matching[0]; + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .expect("WGMMA template"); + assert_eq!(template.matches("ld.f32 %acc").count(), 32); + assert_eq!(template.matches("st.f32 [$0").count(), 32); + assert_eq!(template.matches("wgmma.mma_async").count(), 1); + assert!(template.contains("wgmma.fence.sync.aligned")); + assert!(template.contains("wgmma.commit_group.sync.aligned")); + assert!(template.contains("wgmma.wait_group.sync.aligned 0")); + assert_eq!( + asm.get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some("l,l,l,~{memory}") + ); + assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); + + Ok(()) +} + +#[test] +fn test_pointer_form_wgmma_region_canonicalizes_reborrow_identity() -> Result<(), anyhow::Error> { + use dialect_mir::attributes::{MirCastKindAttr, MirPointerKindAuthorityAttr}; + use dialect_mir::types::{MirArrayType, MirPointerKind, MirPtrType}; + use pliron::builtin::attributes::IntegerAttr; + use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; + use pliron::utils::apint::APInt; + use std::num::NonZeroUsize; + + let mut ctx = make_test_ctx(); + let f32_ty = FP32Type::get(&ctx); + let row_ty = MirArrayType::get(&mut ctx, f32_ty.into(), 8); + let accumulator_ty = MirArrayType::get(&mut ctx, row_ty.into(), 4); + let erased_ptr_ty = MirPtrType::get_generic(&mut ctx, accumulator_ty.into(), true); + let unique_ptr_ty: pliron::r#type::TypeHandle = MirPtrType::get_generic_with_kind( + &mut ctx, + accumulator_ty.into(), + true, + MirPointerKind::UniqueRef, + ) + .into(); + let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + + let (module_ptr, entry) = build_test_kernel( + &mut ctx, + vec![erased_ptr_ty.into(), u64_ty.into(), u64_ty.into()], + ); + let accumulator = entry.deref(&ctx).get_argument(0); + let desc_a = entry.deref(&ctx).get_argument(1); + let desc_b = entry.deref(&ctx).get_argument(2); + + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + + // The importer emits a fresh kind-only retype for each `Rvalue::Ref`. + // Both typed operands must canonicalize to the same storage identity, while + // the first typed operand remains available to the linear lowering plan. + let first_retype = Operation::new( + &mut ctx, + mir::MirCastOp::get_concrete_op_info(), + vec![unique_ptr_ty], + vec![accumulator], + vec![], + 0, + ); + mir::MirCastOp::new(first_retype).set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + mir::MirCastOp::new(first_retype) + .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); + first_retype.insert_at_back(entry, &ctx); + let first_reborrow = first_retype.deref(&ctx).get_result(0); + + Operation::new( + &mut ctx, + nvvm::WgmmaMmaM64N64K16F32Bf16Op::get_concrete_op_info(), + vec![], + vec![first_reborrow, desc_a, desc_b], + vec![], + 0, + ) + .insert_at_back(entry, &ctx); + + let second_retype = Operation::new( + &mut ctx, + mir::MirCastOp::get_concrete_op_info(), + vec![unique_ptr_ty], + vec![accumulator], + vec![], + 0, + ); + mir::MirCastOp::new(second_retype).set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + mir::MirCastOp::new(second_retype) + .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); + second_retype.insert_at_back(entry, &ctx); + let second_reborrow = second_retype.deref(&ctx).get_result(0); + assert_ne!( + first_reborrow, second_reborrow, + "the regression requires two distinct reborrow SSA values" + ); + let first_reborrow_ty = first_reborrow.get_type(&ctx); + assert_eq!( + first_reborrow_ty + .deref(&ctx) + .downcast_ref::() + .expect("first reborrow must remain a typed MIR pointer") + .pointer_kind(), + MirPointerKind::UniqueRef, + "the regression must provide a typed UniqueRef for the linear plan to retain" + ); + + Operation::new( + &mut ctx, + nvvm::WgmmaMmaM64N64K16F32Bf16Op::get_concrete_op_info(), + vec![], + vec![second_reborrow, desc_a, desc_b], + vec![], + 0, + ) + .insert_at_back(entry, &ctx); + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + + let zero_attr = IntegerAttr::new(u64_ty, APInt::from_i64(0, NonZeroUsize::new(64).unwrap())); + let zero = Operation::new( + &mut ctx, + mir::MirConstantOp::get_concrete_op_info(), + vec![u64_ty.into()], + vec![], + vec![], + 0, + ); + mir::MirConstantOp::new(zero).set_attr_value(&ctx, zero_attr); + zero.insert_at_back(entry, &ctx); + let zero_value = zero.deref(&ctx).get_result(0); + nvvm::WgmmaWaitGroupSyncAlignedOp::build(&mut ctx, zero_value).insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let matching = lowered_kernel_body(&ctx, module_ptr) + .into_iter() + .filter_map(|operation| Operation::get_op::(operation, &ctx)) + .filter(|asm| { + asm.get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .is_some_and(|template| { + template.contains("wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16") + }) + }) + .collect::>(); + assert_eq!( + matching.len(), + 1, + "distinct reborrow SSA values for one accumulator must not break deferred fusion" + ); + let template = matching[0] + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .expect("WGMMA template"); + assert_eq!( + template.matches("wgmma.mma_async").count(), + 2, + "both reborrows must remain in one fused accumulator region" + ); + Ok(()) +} + +#[test] +fn test_pointer_form_wgmma_sequence_uses_value_adapter_before_lowering() -> Result<(), anyhow::Error> +{ + use dialect_mir::types::{MirArrayType, MirPtrType}; + use pliron::builtin::attributes::IntegerAttr; + use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; + use pliron::utils::apint::APInt; + use std::num::NonZeroUsize; + + const ACCUMULATOR_LEN: usize = 32; + const DESCRIPTOR_COUNT: usize = 4; + + let mut ctx = make_test_ctx(); + let f32_ty = FP32Type::get(&ctx); + let row_ty = MirArrayType::get(&mut ctx, f32_ty.into(), 8); + let accumulator_ty = MirArrayType::get(&mut ctx, row_ty.into(), 4); + let accumulator_ptr_ty = MirPtrType::get_generic(&mut ctx, accumulator_ty.into(), true); + let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + + let (module_ptr, entry) = build_test_kernel( + &mut ctx, + vec![ + accumulator_ptr_ty.into(), + u64_ty.into(), + u64_ty.into(), + u64_ty.into(), + u64_ty.into(), + ], + ); + let accumulator = entry.deref(&ctx).get_argument(0); + let descriptors = (1..=DESCRIPTOR_COUNT) + .map(|index| entry.deref(&ctx).get_argument(index)) + .collect::>(); + + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + for pair in descriptors.chunks_exact(2) { + Operation::new( + &mut ctx, + nvvm::WgmmaMmaM64N64K16F32Bf16Op::get_concrete_op_info(), + vec![], + vec![accumulator, pair[0], pair[1]], + vec![], + 0, + ) + .insert_at_back(entry, &ctx); + } + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + + let zero_attr = IntegerAttr::new(u64_ty, APInt::from_i64(0, NonZeroUsize::new(64).unwrap())); + let zero = Operation::new( + &mut ctx, + mir::MirConstantOp::get_concrete_op_info(), + vec![u64_ty.into()], + vec![], + vec![], + 0, + ); + mir::MirConstantOp::new(zero).set_attr_value(&ctx, zero_attr); + zero.insert_at_back(entry, &ctx); + let zero_value = zero.deref(&ctx).get_result(0); + nvvm::WgmmaWaitGroupSyncAlignedOp::build(&mut ctx, zero_value).insert_at_back(entry, &ctx); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let body = lowered_kernel_body(&ctx, module_ptr); + let matching = body + .iter() + .copied() + .filter_map(|operation| { + Operation::get_op::(operation, &ctx) + .map(|inline_asm| (operation, inline_asm)) + }) + .filter(|(_, asm)| { + asm.get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .is_some_and(|template| template.contains("wgmma.mma_async")) + }) + .collect::>(); + assert_eq!(matching.len(), 1); + + let (asm_operation, asm) = &matching[0]; + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .expect("WGMMA template"); + + assert_eq!(template.matches("wgmma.fence.sync.aligned").count(), 1); + assert_eq!(template.matches("wgmma.mma_async").count(), 2); + assert_eq!( + template.matches("wgmma.commit_group.sync.aligned").count(), + 1 + ); + assert_eq!( + template.matches("wgmma.wait_group.sync.aligned 0").count(), + 1 + ); + assert!( + !template.contains(".reg .f32") + && !template.contains("ld.f32") + && !template.contains("st.f32"), + "value-form WGMMA must not materialize accumulator memory inside asm: {template}" + ); + + let mut expected_constraints = vec!["=f".to_owned(); ACCUMULATOR_LEN]; + expected_constraints.extend((0..ACCUMULATOR_LEN).map(|index| index.to_string())); + expected_constraints.extend((0..DESCRIPTOR_COUNT).map(|_| "l".to_owned())); + expected_constraints.push("~{memory}".to_owned()); + let expected_constraints = expected_constraints.join(","); + + assert_eq!( + asm.get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some(expected_constraints.as_str()) + ); + assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); + + let asm_ref = asm_operation.deref(&ctx); + assert_eq!( + asm_ref.get_num_operands(), + ACCUMULATOR_LEN + DESCRIPTOR_COUNT + ); + assert_eq!(asm_ref.get_num_results(), 1); + + let asm_position = body + .iter() + .position(|operation| operation == asm_operation) + .expect("WGMMA asm must be in the lowered kernel body"); + + let load_positions = body + .iter() + .enumerate() + .filter_map(|(index, operation)| { + Operation::get_op::(*operation, &ctx) + .is_some() + .then_some(index) + }) + .collect::>(); + assert_eq!( + load_positions.len(), + ACCUMULATOR_LEN, + "the linear adapter must load each accumulator value exactly once" + ); + assert!( + load_positions.iter().all(|index| *index < asm_position), + "all accumulator loads must happen before the WGMMA region" + ); + + let store_positions = body + .iter() + .enumerate() + .filter_map(|(index, operation)| { + Operation::get_op::(*operation, &ctx) + .is_some() + .then_some(index) + }) + .collect::>(); + assert_eq!( + store_positions.len(), + ACCUMULATOR_LEN, + "the linear adapter must store each accumulator value exactly once" + ); + assert!( + store_positions.iter().all(|index| *index > asm_position), + "all accumulator stores must happen after the final wait" + ); + + assert_eq!( + body.iter() + .filter( + |operation| Operation::get_op::(**operation, &ctx).is_some() + ) + .count(), + ACCUMULATOR_LEN, + "all WGMMA accumulator results must be recovered as scalar SSA values" + ); + + Ok(()) +} + +#[test] +fn test_pointer_form_m64n128_bf16_linear_full_drain_uses_sixty_four_value_adapter() +-> Result<(), anyhow::Error> { + let mut ctx = make_test_ctx(); + let (module_ptr, entry, accumulators, descriptors) = + build_wgmma_m64n128_canonical_pointer_test_kernel(&mut ctx, 1, 4); + let accumulator = accumulators[0]; + + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_pointer_wgmma_mma_m64n128(&mut ctx, entry, accumulator, descriptors[0], descriptors[1]); + append_pointer_wgmma_mma_m64n128(&mut ctx, entry, accumulator, descriptors[2], descriptors[3]); + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_wgmma_wait_group_constant(&mut ctx, entry, 0); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let matching = lowered_kernel_body(&ctx, module_ptr) + .into_iter() + .filter_map(|operation| Operation::get_op::(operation, &ctx)) + .filter(|asm| { + asm.get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .is_some_and(|template| template.contains("m64n128k16.f32.bf16.bf16")) + }) + .collect::>(); + assert_eq!(matching.len(), 1); + + let asm = &matching[0]; + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .expect("m64n128 pointer-form WGMMA template"); + assert_eq!(template.matches("wgmma.mma_async").count(), 2); + assert!(template.contains("m64n128k16.f32.bf16.bf16")); + assert!(!template.contains("ld.f32")); + assert!(!template.contains("st.f32")); + + let constraints = asm + .get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .expect("m64n128 pointer-form WGMMA constraints"); + assert_eq!( + constraints + .split(',') + .filter(|value| *value == "=f") + .count(), + 64 + ); + assert_eq!( + constraints.split(',').filter(|value| *value == "l").count(), + 4 + ); + assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); + + Ok(()) +} + +#[test] +fn test_pointer_form_f16_wgmma_linear_full_drain_uses_value_adapter() -> Result<(), anyhow::Error> { + let mut ctx = make_test_ctx(); + let (module_ptr, entry, accumulators, descriptors) = + build_wgmma_canonical_pointer_test_kernel(&mut ctx, 1, 2); + let accumulator = accumulators[0]; + + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_pointer_wgmma_mma_f16(&mut ctx, entry, accumulator, descriptors[0], descriptors[1]); + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_wgmma_wait_group_constant(&mut ctx, entry, 0); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let matching = lowered_kernel_body(&ctx, module_ptr) + .into_iter() + .filter_map(|operation| Operation::get_op::(operation, &ctx)) + .filter(|asm| { + asm.get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .is_some_and(|template| template.contains("m64n64k16.f32.f16.f16")) + }) + .collect::>(); + assert_eq!(matching.len(), 1); + + let asm = &matching[0]; + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .expect("F16 pointer-form WGMMA template"); + assert_eq!(template.matches("wgmma.fence.sync.aligned").count(), 1); + assert_eq!(template.matches("wgmma.mma_async").count(), 1); + assert_eq!( + template.matches("wgmma.commit_group.sync.aligned").count(), + 1 + ); + assert_eq!( + template.matches("wgmma.wait_group.sync.aligned 0").count(), + 1 + ); + assert!(template.contains("m64n64k16.f32.f16.f16")); + assert!(!template.contains(".bf16.bf16")); + assert!(!template.contains("ld.f32")); + assert!(!template.contains("st.f32")); + assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); + + Ok(()) +} + +#[test] +fn test_pointer_form_tf32_wgmma_linear_full_drain_uses_value_adapter() -> Result<(), anyhow::Error> +{ + let mut ctx = make_test_ctx(); + let (module_ptr, entry, accumulators, descriptors) = + build_wgmma_canonical_pointer_test_kernel(&mut ctx, 1, 2); + let accumulator = accumulators[0]; + + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_pointer_wgmma_mma_tf32(&mut ctx, entry, accumulator, descriptors[0], descriptors[1]); + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_wgmma_wait_group_constant(&mut ctx, entry, 0); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let matching = lowered_kernel_body(&ctx, module_ptr) + .into_iter() + .filter_map(|operation| Operation::get_op::(operation, &ctx)) + .filter(|asm| { + asm.get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .is_some_and(|template| template.contains("m64n64k8.f32.tf32.tf32")) + }) + .collect::>(); + assert_eq!(matching.len(), 1); + + let asm = &matching[0]; + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .expect("TF32 pointer-form WGMMA template"); + assert_eq!(template.matches("wgmma.fence.sync.aligned").count(), 1); + assert_eq!(template.matches("wgmma.mma_async").count(), 1); + assert_eq!( + template.matches("wgmma.commit_group.sync.aligned").count(), + 1 + ); + assert_eq!( + template.matches("wgmma.wait_group.sync.aligned 0").count(), + 1 + ); + assert!(template.contains("m64n64k8.f32.tf32.tf32")); + assert!(!template.contains(".bf16.bf16")); + assert!(!template.contains(".f16.f16")); + assert!(!template.contains("m64n64k16")); + assert!(template.contains("$64, $65, 1, 1, 1;")); + assert!(!template.contains("$64, $65, 1, 1, 1, 0, 0;")); + assert!(!template.contains("ld.f32")); + assert!(!template.contains("st.f32")); + assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); + + Ok(()) +} + +#[test] +fn test_pointer_form_wgmma_partial_wait_pipeline_keeps_multiple_groups_in_flight() +-> Result<(), anyhow::Error> { + const SLOT_COUNT: usize = 2; + const ACCUMULATOR_LEN: usize = 32; + const GROUP_COUNT: usize = 4; + const DESCRIPTOR_COUNT: usize = GROUP_COUNT * 2; + const MAX_PENDING_GROUPS: i64 = 1; + const RESULT_COUNT: usize = SLOT_COUNT * ACCUMULATOR_LEN; + + let mut ctx = make_test_ctx(); + let (module_ptr, entry, accumulators, descriptors) = + build_wgmma_canonical_pointer_test_kernel(&mut ctx, SLOT_COUNT, DESCRIPTOR_COUNT); + + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + for group in 0..GROUP_COUNT { + append_pointer_wgmma_mma( + &mut ctx, + entry, + accumulators[group % SLOT_COUNT], + descriptors[group * 2], + descriptors[group * 2 + 1], + ); + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + if group + 1 >= SLOT_COUNT { + append_wgmma_wait_group_constant(&mut ctx, entry, MAX_PENDING_GROUPS); + } + } + append_wgmma_wait_group_constant(&mut ctx, entry, 0); + append_return(&mut ctx, entry); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let body = lowered_kernel_body(&ctx, module_ptr); + let matching = body + .iter() + .copied() + .filter_map(|operation| { + Operation::get_op::(operation, &ctx) + .map(|inline_asm| (operation, inline_asm)) + }) + .filter(|(_, asm)| { + asm.get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .is_some_and(|template| { + template.contains("wgmma.mma_async") + && template.contains("wgmma.wait_group.sync.aligned 1") + }) + }) + .collect::>(); + assert_eq!( + matching.len(), + 1, + "expected one fused partial-wait WGMMA pipeline" + ); + + let (asm_operation, asm) = &matching[0]; + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .expect("pipeline WGMMA template"); + + assert_eq!(template.matches("wgmma.fence.sync.aligned").count(), 1); + assert_eq!(template.matches("wgmma.mma_async").count(), GROUP_COUNT); + assert_eq!( + template.matches("wgmma.commit_group.sync.aligned").count(), + GROUP_COUNT + ); + assert_eq!( + template.matches("wgmma.wait_group.sync.aligned 1").count(), + GROUP_COUNT - SLOT_COUNT + 1 + ); + assert_eq!( + template.matches("wgmma.wait_group.sync.aligned 0").count(), + 1 + ); + assert!( + template.contains("{$0, $1"), + "slot 0 must use the first accumulator tuple" + ); + assert!( + template.contains("{$32, $33"), + "slot 1 must use the second accumulator tuple" + ); + assert!(template.contains("$128, $129")); + assert!(template.contains("$134, $135")); + assert!( + !template.contains(".reg .f32") + && !template.contains("ld.f32") + && !template.contains("st.f32"), + "pipeline WGMMA must not materialize accumulator memory inside asm: {template}" + ); + + let mut expected_constraints = vec!["=f".to_owned(); RESULT_COUNT]; + expected_constraints.extend((0..RESULT_COUNT).map(|index| index.to_string())); + expected_constraints.extend((0..DESCRIPTOR_COUNT).map(|_| "l".to_owned())); + expected_constraints.push("~{memory}".to_owned()); + let expected_constraints = expected_constraints.join(","); + assert_eq!( + asm.get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some(expected_constraints.as_str()) + ); + assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); + assert_eq!( + asm_operation.deref(&ctx).get_num_operands(), + RESULT_COUNT + DESCRIPTOR_COUNT + ); + assert_eq!(asm_operation.deref(&ctx).get_num_results(), 1); + + let asm_position = body + .iter() + .position(|operation| operation == asm_operation) + .expect("pipeline asm must be in the lowered kernel body"); + let load_positions = body + .iter() + .enumerate() + .filter_map(|(index, operation)| { + Operation::get_op::(*operation, &ctx) + .is_some() + .then_some(index) + }) + .collect::>(); + assert_eq!(load_positions.len(), RESULT_COUNT); + assert!(load_positions.iter().all(|index| *index < asm_position)); + + let store_positions = body + .iter() + .enumerate() + .filter_map(|(index, operation)| { + Operation::get_op::(*operation, &ctx) + .is_some() + .then_some(index) + }) + .collect::>(); + assert_eq!(store_positions.len(), RESULT_COUNT); + assert!(store_positions.iter().all(|index| *index > asm_position)); + assert_eq!( + body.iter() + .filter(|operation| { + Operation::get_op::(**operation, &ctx).is_some() + }) + .count(), + RESULT_COUNT + ); + + Ok(()) +} + +pub(super) fn build_pointer_form_wgmma_counted_pipeline_case( + ctx: &mut Context, + slot_count: usize, + wait_depths: &[i64], + repeat_last_accumulator: bool, +) -> pliron::context::Ptr { + use dialect_mir::types::{MirArrayType, MirPtrType}; + use pliron::basic_block::BasicBlock; + use pliron::builtin::op_interfaces::OperandSegmentInterface; + use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; + + assert_eq!(wait_depths.len(), slot_count); + + let f32_ty = FP32Type::get(ctx); + let row_ty = MirArrayType::get(ctx, f32_ty.into(), 8); + let accumulator_ty = MirArrayType::get(ctx, row_ty.into(), 4); + let accumulator_ptr_ty = MirPtrType::get_generic(ctx, accumulator_ty.into(), true); + let u32_ty = IntegerType::get(ctx, 32, Signedness::Unsigned); + let u64_ty = IntegerType::get(ctx, 64, Signedness::Unsigned); + let i1_ty = IntegerType::get(ctx, 1, Signedness::Signless); + let u64_type: pliron::r#type::TypeHandle = u64_ty.into(); + + let mut argument_types: Vec = + vec![accumulator_ptr_ty.into(); slot_count]; + argument_types.extend(vec![u64_type; slot_count * 2]); + let (module_ptr, preheader) = build_test_kernel(ctx, argument_types); + + let accumulators = (0..slot_count) + .map(|slot| preheader.deref(ctx).get_argument(slot)) + .collect::>(); + let desc_bases = (0..slot_count * 2) + .map(|index| preheader.deref(ctx).get_argument(slot_count + index)) + .collect::>(); + + let module_region = module_ptr.deref(ctx).get_region(0); + let module_block = module_region.deref(ctx).iter(ctx).next().unwrap(); + let function = module_block.deref(ctx).iter(ctx).next().unwrap(); + let function_region = function.deref(ctx).get_region(0); + + let mut header_types: Vec = vec![u32_ty.into()]; + header_types.extend(vec![u64_type; slot_count * 2]); + let header = BasicBlock::new(ctx, None, header_types); + header.insert_at_back(function_region, ctx); + let latch = BasicBlock::new(ctx, None, vec![]); + latch.insert_at_back(function_region, ctx); + let exit = BasicBlock::new(ctx, None, vec![]); + exit.insert_at_back(function_region, ctx); + + nvvm::WgmmaFenceSyncAlignedOp::build(ctx).insert_at_back(preheader, ctx); + let i0 = append_mir_unsigned_constant(ctx, preheader, u32_ty, 0); + let mut initial_values = vec![i0]; + initial_values.extend(desc_bases.iter().copied()); + Operation::new( + ctx, + mir::MirGotoOp::get_concrete_op_info(), + vec![], + initial_values, + vec![header], + 0, + ) + .insert_at_back(preheader, ctx); + + let i = header.deref(ctx).get_argument(0); + let descriptors = (0..slot_count * 2) + .map(|index| header.deref(ctx).get_argument(1 + index)) + .collect::>(); + let bound = append_mir_unsigned_constant(ctx, header, u32_ty, 4); + let lt = Operation::new( + ctx, + mir::MirLtOp::get_concrete_op_info(), + vec![i1_ty.into()], + vec![i, bound], + vec![], + 0, + ); + lt.insert_at_back(header, ctx); + let lt_value = lt.deref(ctx).get_result(0); + let not_lt = Operation::new( + ctx, + mir::MirNotOp::get_concrete_op_info(), + vec![i1_ty.into()], + vec![lt_value], + vec![], + 0, + ); + not_lt.insert_at_back(header, ctx); + let not_lt_value = not_lt.deref(ctx).get_result(0); + let (branch_operands, segment_sizes) = + mir::MirCondBranchOp::compute_segment_sizes(vec![vec![not_lt_value], vec![], vec![]]); + let branch = Operation::new( + ctx, + mir::MirCondBranchOp::get_concrete_op_info(), + vec![], + branch_operands, + vec![exit, latch], + 0, + ); + Operation::get_op::(branch, ctx) + .expect("MirCondBranchOp") + .set_operand_segment_sizes(ctx, segment_sizes); + branch.insert_at_back(header, ctx); + + for slot in 0..slot_count { + let accumulator = if repeat_last_accumulator && slot + 1 == slot_count { + accumulators[0] + } else { + accumulators[slot] + }; + append_pointer_wgmma_mma( + ctx, + latch, + accumulator, + descriptors[slot * 2], + descriptors[slot * 2 + 1], + ); + nvvm::WgmmaCommitGroupSyncAlignedOp::build(ctx).insert_at_back(latch, ctx); + append_wgmma_wait_group_constant(ctx, latch, wait_depths[slot]); + } + + let one = append_mir_unsigned_constant(ctx, latch, u32_ty, 1); + let i_next = Operation::new( + ctx, + mir::MirAddOp::get_concrete_op_info(), + vec![u32_ty.into()], + vec![i, one], + vec![], + 0, + ); + i_next.insert_at_back(latch, ctx); + let i_next = i_next.deref(ctx).get_result(0); + + let mut next_values = vec![i_next]; + for (index, descriptor) in descriptors.iter().copied().enumerate() { + let step = append_mir_unsigned_constant(ctx, latch, u64_ty, 16 * (index as u64 + 1)); + let next = Operation::new( + ctx, + mir::MirAddOp::get_concrete_op_info(), + vec![u64_ty.into()], + vec![descriptor, step], + vec![], + 0, + ); + next.insert_at_back(latch, ctx); + next_values.push(next.deref(ctx).get_result(0)); + } + + Operation::new( + ctx, + mir::MirGotoOp::get_concrete_op_info(), + vec![], + next_values, + vec![header], + 0, + ) + .insert_at_back(latch, ctx); + + append_wgmma_wait_group_constant(ctx, exit, 0); + append_return(ctx, exit); + + module_ptr +} + +#[test] +fn test_pointer_form_wgmma_two_slot_counted_pipeline_stays_register_resident() +-> Result<(), anyhow::Error> { + use dialect_mir::types::{MirArrayType, MirPtrType}; + use pliron::basic_block::BasicBlock; + use pliron::builtin::op_interfaces::OperandSegmentInterface; + use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; + + const SLOT_COUNT: usize = 2; + const ACCUMULATOR_LEN: usize = 32; + const RESULT_COUNT: usize = SLOT_COUNT * ACCUMULATOR_LEN; + const LOOP_CONTROL_COUNT: usize = 9; + const TRIP_COUNT: u64 = 4; + const DESC_A0_STEP: u64 = 16; + const DESC_B0_STEP: u64 = 32; + const DESC_A1_STEP: u64 = 48; + const DESC_B1_STEP: u64 = 64; + + let mut ctx = make_test_ctx(); + let f32_ty = FP32Type::get(&ctx); + let row_ty = MirArrayType::get(&mut ctx, f32_ty.into(), 8); + let accumulator_ty = MirArrayType::get(&mut ctx, row_ty.into(), 4); + let accumulator_ptr_ty = MirPtrType::get_generic(&mut ctx, accumulator_ty.into(), true); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); + + let (module_ptr, preheader) = build_test_kernel( + &mut ctx, + vec![ + accumulator_ptr_ty.into(), + accumulator_ptr_ty.into(), + u64_ty.into(), + u64_ty.into(), + u64_ty.into(), + u64_ty.into(), + ], + ); + let accumulator0 = preheader.deref(&ctx).get_argument(0); + let accumulator1 = preheader.deref(&ctx).get_argument(1); + let desc_a0_base = preheader.deref(&ctx).get_argument(2); + let desc_b0_base = preheader.deref(&ctx).get_argument(3); + let desc_a1_base = preheader.deref(&ctx).get_argument(4); + let desc_b1_base = preheader.deref(&ctx).get_argument(5); + + let module_region = module_ptr.deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + let function = module_block.deref(&ctx).iter(&ctx).next().unwrap(); + let function_region = function.deref(&ctx).get_region(0); + + let header = BasicBlock::new( + &mut ctx, + None, + vec![ + u32_ty.into(), + u64_ty.into(), + u64_ty.into(), + u64_ty.into(), + u64_ty.into(), + ], + ); + header.insert_at_back(function_region, &ctx); + let latch = BasicBlock::new(&mut ctx, None, vec![]); + latch.insert_at_back(function_region, &ctx); + let exit = BasicBlock::new(&mut ctx, None, vec![]); + exit.insert_at_back(function_region, &ctx); + + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(preheader, &ctx); + let i0 = append_mir_unsigned_constant(&mut ctx, preheader, u32_ty, 0); + Operation::new( + &mut ctx, + mir::MirGotoOp::get_concrete_op_info(), + vec![], + vec![i0, desc_a0_base, desc_b0_base, desc_a1_base, desc_b1_base], + vec![header], + 0, + ) + .insert_at_back(preheader, &ctx); + + let i = header.deref(&ctx).get_argument(0); + let desc_a0 = header.deref(&ctx).get_argument(1); + let desc_b0 = header.deref(&ctx).get_argument(2); + let desc_a1 = header.deref(&ctx).get_argument(3); + let desc_b1 = header.deref(&ctx).get_argument(4); + let bound = append_mir_unsigned_constant(&mut ctx, header, u32_ty, TRIP_COUNT); + let lt = Operation::new( + &mut ctx, + mir::MirLtOp::get_concrete_op_info(), + vec![i1_ty.into()], + vec![i, bound], + vec![], + 0, + ); + lt.insert_at_back(header, &ctx); + let lt_value = lt.deref(&ctx).get_result(0); + let not_lt = Operation::new( + &mut ctx, + mir::MirNotOp::get_concrete_op_info(), + vec![i1_ty.into()], + vec![lt_value], + vec![], + 0, + ); + not_lt.insert_at_back(header, &ctx); + let not_lt_value = not_lt.deref(&ctx).get_result(0); + let (branch_operands, segment_sizes) = + mir::MirCondBranchOp::compute_segment_sizes(vec![vec![not_lt_value], vec![], vec![]]); + let branch = Operation::new( + &mut ctx, + mir::MirCondBranchOp::get_concrete_op_info(), + vec![], + branch_operands, + vec![exit, latch], + 0, + ); + Operation::get_op::(branch, &ctx) + .expect("MirCondBranchOp") + .set_operand_segment_sizes(&ctx, segment_sizes); + branch.insert_at_back(header, &ctx); + + append_pointer_wgmma_mma(&mut ctx, latch, accumulator0, desc_a0, desc_b0); + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(latch, &ctx); + append_wgmma_wait_group_constant(&mut ctx, latch, 1); + append_pointer_wgmma_mma(&mut ctx, latch, accumulator1, desc_a1, desc_b1); + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(latch, &ctx); + append_wgmma_wait_group_constant(&mut ctx, latch, 1); + + let one = append_mir_unsigned_constant(&mut ctx, latch, u32_ty, 1); + let i_next = Operation::new( + &mut ctx, + mir::MirAddOp::get_concrete_op_info(), + vec![u32_ty.into()], + vec![i, one], + vec![], + 0, + ); + i_next.insert_at_back(latch, &ctx); + let i_next = i_next.deref(&ctx).get_result(0); + + let desc_a0_step = append_mir_unsigned_constant(&mut ctx, latch, u64_ty, DESC_A0_STEP); + let desc_a0_next = Operation::new( + &mut ctx, + mir::MirAddOp::get_concrete_op_info(), + vec![u64_ty.into()], + vec![desc_a0, desc_a0_step], + vec![], + 0, + ); + desc_a0_next.insert_at_back(latch, &ctx); + let desc_a0_next = desc_a0_next.deref(&ctx).get_result(0); + + let desc_b0_step = append_mir_unsigned_constant(&mut ctx, latch, u64_ty, DESC_B0_STEP); + let desc_b0_next = Operation::new( + &mut ctx, + mir::MirAddOp::get_concrete_op_info(), + vec![u64_ty.into()], + vec![desc_b0, desc_b0_step], + vec![], + 0, + ); + desc_b0_next.insert_at_back(latch, &ctx); + let desc_b0_next = desc_b0_next.deref(&ctx).get_result(0); + + let desc_a1_step = append_mir_unsigned_constant(&mut ctx, latch, u64_ty, DESC_A1_STEP); + let desc_a1_next = Operation::new( + &mut ctx, + mir::MirAddOp::get_concrete_op_info(), + vec![u64_ty.into()], + vec![desc_a1, desc_a1_step], + vec![], + 0, + ); + desc_a1_next.insert_at_back(latch, &ctx); + let desc_a1_next = desc_a1_next.deref(&ctx).get_result(0); + + let desc_b1_step = append_mir_unsigned_constant(&mut ctx, latch, u64_ty, DESC_B1_STEP); + let desc_b1_next = Operation::new( + &mut ctx, + mir::MirAddOp::get_concrete_op_info(), + vec![u64_ty.into()], + vec![desc_b1, desc_b1_step], + vec![], + 0, + ); + desc_b1_next.insert_at_back(latch, &ctx); + let desc_b1_next = desc_b1_next.deref(&ctx).get_result(0); + + Operation::new( + &mut ctx, + mir::MirGotoOp::get_concrete_op_info(), + vec![], + vec![ + i_next, + desc_a0_next, + desc_b0_next, + desc_a1_next, + desc_b1_next, + ], + vec![header], + 0, + ) + .insert_at_back(latch, &ctx); + + append_wgmma_wait_group_constant(&mut ctx, exit, 0); + append_return(&mut ctx, exit); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let body = lowered_kernel_body(&ctx, module_ptr); + let matching = body + .iter() + .copied() + .filter_map(|operation| { + Operation::get_op::(operation, &ctx) + .map(|inline_asm| (operation, inline_asm)) + }) + .filter(|(_, asm)| { + asm.get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .is_some_and(|template| template.contains("L__wgmma_pipeline_loop_${:uid}:")) + }) + .collect::>(); + assert_eq!( + matching.len(), + 1, + "expected one fused two-slot counted WGMMA pipeline" + ); + + let (asm_operation, asm) = &matching[0]; + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .expect("counted-pipeline WGMMA template"); + + assert_eq!(template.matches("wgmma.fence.sync.aligned").count(), 1); + assert_eq!(template.matches("wgmma.mma_async").count(), 2); + assert_eq!( + template.matches("wgmma.commit_group.sync.aligned").count(), + 2 + ); + assert_eq!( + template.matches("wgmma.wait_group.sync.aligned 1").count(), + 2 + ); + assert_eq!( + template.matches("wgmma.wait_group.sync.aligned 0").count(), + 1 + ); + assert!(template.contains("{$0, $1")); + assert!(template.contains("{$32, $33")); + assert!(template.contains("mov.u64 %desc_a0, $128;")); + assert!(template.contains("mov.u64 %desc_b0, $129;")); + assert!(template.contains("mov.u64 %desc_a1, $130;")); + assert!(template.contains("mov.u64 %desc_b1, $131;")); + assert!(template.contains("add.u64 %desc_a0, %desc_a0, $132;")); + assert!(template.contains("add.u64 %desc_b0, %desc_b0, $133;")); + assert!(template.contains("add.u64 %desc_a1, %desc_a1, $134;")); + assert!(template.contains("add.u64 %desc_b1, %desc_b1, $135;")); + assert!(template.contains("mov.u64 %remaining, $136;")); + assert!( + !template.contains(".reg .f32") + && !template.contains("ld.f32") + && !template.contains("st.f32"), + "counted pipeline must keep accumulator memory outside asm: {template}" + ); + + let mut expected_constraints = vec!["=f".to_owned(); RESULT_COUNT]; + expected_constraints.extend((0..RESULT_COUNT).map(|index| index.to_string())); + expected_constraints.extend((0..LOOP_CONTROL_COUNT).map(|_| "l".to_owned())); + expected_constraints.push("~{memory}".to_owned()); + let expected_constraints = expected_constraints.join(","); + assert_eq!( + asm.get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some(expected_constraints.as_str()) + ); + assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); + assert_eq!( + asm_operation.deref(&ctx).get_num_operands(), + RESULT_COUNT + LOOP_CONTROL_COUNT + ); + + let asm_position = body + .iter() + .position(|operation| operation == asm_operation) + .expect("counted-pipeline WGMMA asm must be in the lowered kernel body"); + let load_positions = body + .iter() + .enumerate() + .filter_map(|(index, operation)| { + Operation::get_op::(*operation, &ctx) + .is_some() + .then_some(index) + }) + .collect::>(); + assert_eq!(load_positions.len(), RESULT_COUNT); + assert!(load_positions.iter().all(|index| *index < asm_position)); + + let store_positions = body + .iter() + .enumerate() + .filter_map(|(index, operation)| { + Operation::get_op::(*operation, &ctx) + .is_some() + .then_some(index) + }) + .collect::>(); + assert_eq!(store_positions.len(), RESULT_COUNT); + assert!(store_positions.iter().all(|index| *index > asm_position)); + assert_eq!( + body.iter() + .filter(|operation| { + Operation::get_op::(**operation, &ctx).is_some() + }) + .count(), + RESULT_COUNT + ); + + Ok(()) +} + +#[test] +fn test_pointer_form_wgmma_three_slot_counted_pipeline_stays_register_resident() +-> Result<(), anyhow::Error> { + const RESULT_COUNT: usize = 96; + const LOOP_CONTROL_COUNT: usize = 13; + + let mut ctx = make_test_ctx(); + let module_ptr = build_pointer_form_wgmma_counted_pipeline_case(&mut ctx, 3, &[2, 2, 2], false); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let body = lowered_kernel_body(&ctx, module_ptr); + let matching = body + .iter() + .copied() + .filter_map(|operation| { + Operation::get_op::(operation, &ctx) + .map(|inline_asm| (operation, inline_asm)) + }) + .filter(|(_, asm)| { + asm.get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .is_some_and(|template| template.contains("L__wgmma_pipeline_loop_${:uid}:")) + }) + .collect::>(); + assert_eq!( + matching.len(), + 1, + "expected one fused three-slot counted WGMMA pipeline" + ); + + let (asm_operation, asm) = &matching[0]; + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .expect("three-slot counted-pipeline WGMMA template"); + + assert_eq!(template.matches("wgmma.fence.sync.aligned").count(), 1); + assert_eq!(template.matches("wgmma.mma_async").count(), 3); + assert_eq!( + template.matches("wgmma.commit_group.sync.aligned").count(), + 3 + ); + assert_eq!( + template.matches("wgmma.wait_group.sync.aligned 2").count(), + 3 + ); + assert_eq!( + template.matches("wgmma.wait_group.sync.aligned 0").count(), + 1 + ); + assert!(template.contains("{$0, $1")); + assert!(template.contains("{$32, $33")); + assert!(template.contains("{$64, $65")); + assert!(template.contains("mov.u64 %desc_a0, $192;")); + assert!(template.contains("mov.u64 %desc_b0, $193;")); + assert!(template.contains("mov.u64 %desc_a1, $194;")); + assert!(template.contains("mov.u64 %desc_b1, $195;")); + assert!(template.contains("mov.u64 %desc_a2, $196;")); + assert!(template.contains("mov.u64 %desc_b2, $197;")); + assert!(template.contains("add.u64 %desc_a0, %desc_a0, $198;")); + assert!(template.contains("add.u64 %desc_b0, %desc_b0, $199;")); + assert!(template.contains("add.u64 %desc_a1, %desc_a1, $200;")); + assert!(template.contains("add.u64 %desc_b1, %desc_b1, $201;")); + assert!(template.contains("add.u64 %desc_a2, %desc_a2, $202;")); + assert!(template.contains("add.u64 %desc_b2, %desc_b2, $203;")); + assert!(template.contains("mov.u64 %remaining, $204;")); + assert!( + !template.contains(".reg .f32") + && !template.contains("ld.f32") + && !template.contains("st.f32"), + "three-slot counted pipeline must keep accumulator memory outside asm: {template}" + ); + + let mut expected_constraints = vec!["=f".to_owned(); RESULT_COUNT]; + expected_constraints.extend((0..RESULT_COUNT).map(|index| index.to_string())); + expected_constraints.extend((0..LOOP_CONTROL_COUNT).map(|_| "l".to_owned())); + expected_constraints.push("~{memory}".to_owned()); + let expected_constraints = expected_constraints.join(","); + assert_eq!( + asm.get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some(expected_constraints.as_str()) + ); + assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); + assert_eq!( + asm_operation.deref(&ctx).get_num_operands(), + RESULT_COUNT + LOOP_CONTROL_COUNT + ); + + let asm_position = body + .iter() + .position(|operation| operation == asm_operation) + .expect("three-slot counted-pipeline WGMMA asm must be in the lowered kernel body"); + let load_positions = body + .iter() + .enumerate() + .filter_map(|(index, operation)| { + Operation::get_op::(*operation, &ctx) + .is_some() + .then_some(index) + }) + .collect::>(); + assert_eq!(load_positions.len(), RESULT_COUNT); + assert!(load_positions.iter().all(|index| *index < asm_position)); + + let store_positions = body + .iter() + .enumerate() + .filter_map(|(index, operation)| { + Operation::get_op::(*operation, &ctx) + .is_some() + .then_some(index) + }) + .collect::>(); + assert_eq!(store_positions.len(), RESULT_COUNT); + assert!(store_positions.iter().all(|index| *index > asm_position)); + assert_eq!( + body.iter() + .filter(|operation| { + Operation::get_op::(**operation, &ctx).is_some() + }) + .count(), + RESULT_COUNT + ); + + Ok(()) +} + +#[test] +fn test_pointer_form_wgmma_counted_k_loop_stays_register_resident() -> Result<(), anyhow::Error> { + use dialect_mir::attributes::{MirCastKindAttr, MirPointerKindAuthorityAttr}; + use dialect_mir::types::{MirArrayType, MirPointerKind, MirPtrType}; + use pliron::basic_block::BasicBlock; + use pliron::builtin::op_interfaces::OperandSegmentInterface; + use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; + + const ACCUMULATOR_LEN: usize = 32; + const LOOP_CONTROL_COUNT: usize = 5; + const TRIP_COUNT: u64 = 4; + const DESC_A_STEP: u64 = 16; + const DESC_B_STEP: u64 = 32; + + let mut ctx = make_test_ctx(); + let f32_ty = FP32Type::get(&ctx); + let row_ty = MirArrayType::get(&mut ctx, f32_ty.into(), 8); + let accumulator_ty = MirArrayType::get(&mut ctx, row_ty.into(), 4); + let accumulator_ptr_ty = MirPtrType::get_generic(&mut ctx, accumulator_ty.into(), true); + let unique_accumulator_ptr_ty: pliron::r#type::TypeHandle = MirPtrType::get_generic_with_kind( + &mut ctx, + accumulator_ty.into(), + true, + MirPointerKind::UniqueRef, + ) + .into(); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); + + let (module_ptr, preheader) = build_test_kernel( + &mut ctx, + vec![accumulator_ptr_ty.into(), u64_ty.into(), u64_ty.into()], + ); + let accumulator = preheader.deref(&ctx).get_argument(0); + let desc_a_base = preheader.deref(&ctx).get_argument(1); + let desc_b_base = preheader.deref(&ctx).get_argument(2); + + let module_region = module_ptr.deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + let function = module_block.deref(&ctx).iter(&ctx).next().unwrap(); + let function_region = function.deref(&ctx).get_region(0); + + let header = BasicBlock::new( + &mut ctx, + None, + vec![u32_ty.into(), u64_ty.into(), u64_ty.into()], + ); + header.insert_at_back(function_region, &ctx); + let body = BasicBlock::new(&mut ctx, None, vec![]); + body.insert_at_back(function_region, &ctx); + let latch = BasicBlock::new(&mut ctx, None, vec![]); + latch.insert_at_back(function_region, &ctx); + let exit = BasicBlock::new(&mut ctx, None, vec![]); + exit.insert_at_back(function_region, &ctx); + let wait_block = BasicBlock::new(&mut ctx, None, vec![]); + wait_block.insert_at_back(function_region, &ctx); + + // preheader: fence; i0 = 0; goto header(i0, desc_a_base, desc_b_base) + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(preheader, &ctx); + let i0 = append_mir_unsigned_constant(&mut ctx, preheader, u32_ty, 0); + Operation::new( + &mut ctx, + mir::MirGotoOp::get_concrete_op_info(), + vec![], + vec![i0, desc_a_base, desc_b_base], + vec![header], + 0, + ) + .insert_at_back(preheader, &ctx); + + // header(i, desc_a, desc_b): if !(i < 4) exit else body. + let i = header.deref(&ctx).get_argument(0); + let desc_a = header.deref(&ctx).get_argument(1); + let desc_b = header.deref(&ctx).get_argument(2); + let bound = append_mir_unsigned_constant(&mut ctx, header, u32_ty, TRIP_COUNT); + let lt = Operation::new( + &mut ctx, + mir::MirLtOp::get_concrete_op_info(), + vec![i1_ty.into()], + vec![i, bound], + vec![], + 0, + ); + lt.insert_at_back(header, &ctx); + let lt_value = lt.deref(&ctx).get_result(0); + let not_lt = Operation::new( + &mut ctx, + mir::MirNotOp::get_concrete_op_info(), + vec![i1_ty.into()], + vec![lt_value], + vec![], + 0, + ); + not_lt.insert_at_back(header, &ctx); + let not_lt_value = not_lt.deref(&ctx).get_result(0); + let (branch_operands, segment_sizes) = + mir::MirCondBranchOp::compute_segment_sizes(vec![vec![not_lt_value], vec![], vec![]]); + let branch = Operation::new( + &mut ctx, + mir::MirCondBranchOp::get_concrete_op_info(), + vec![], + branch_operands, + vec![exit, body], + 0, + ); + Operation::get_op::(branch, &ctx) + .expect("MirCondBranchOp") + .set_operand_segment_sizes(&ctx, segment_sizes); + branch.insert_at_back(header, &ctx); + + // body: a fresh Rust `&mut` reborrow and one WGMMA per K iteration. Rust + // calls end a MIR block, so the post-call arithmetic lives in the latch. + // The reborrow is loop-local, but its canonical storage identity is the + // preheader accumulator argument. + let retype = Operation::new( + &mut ctx, + mir::MirCastOp::get_concrete_op_info(), + vec![unique_accumulator_ptr_ty], + vec![accumulator], + vec![], + 0, + ); + mir::MirCastOp::new(retype).set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + mir::MirCastOp::new(retype) + .set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); + retype.insert_at_back(body, &ctx); + let reborrowed_accumulator = retype.deref(&ctx).get_result(0); + append_pointer_wgmma_mma(&mut ctx, body, reborrowed_accumulator, desc_a, desc_b); + Operation::new( + &mut ctx, + mir::MirGotoOp::get_concrete_op_info(), + vec![], + vec![], + vec![latch], + 0, + ) + .insert_at_back(body, &ctx); + + // latch: affine induction and descriptor recurrences, then the back edge. + let one = append_mir_unsigned_constant(&mut ctx, latch, u32_ty, 1); + let i_next = Operation::new( + &mut ctx, + mir::MirAddOp::get_concrete_op_info(), + vec![u32_ty.into()], + vec![i, one], + vec![], + 0, + ); + i_next.insert_at_back(latch, &ctx); + let i_next = i_next.deref(&ctx).get_result(0); + + let desc_a_step = append_mir_unsigned_constant(&mut ctx, latch, u64_ty, DESC_A_STEP); + let desc_a_next = Operation::new( + &mut ctx, + mir::MirAddOp::get_concrete_op_info(), + vec![u64_ty.into()], + vec![desc_a, desc_a_step], + vec![], + 0, + ); + desc_a_next.insert_at_back(latch, &ctx); + let desc_a_next = desc_a_next.deref(&ctx).get_result(0); + + let desc_b_step = append_mir_unsigned_constant(&mut ctx, latch, u64_ty, DESC_B_STEP); + let desc_b_next = Operation::new( + &mut ctx, + mir::MirAddOp::get_concrete_op_info(), + vec![u64_ty.into()], + vec![desc_b, desc_b_step], + vec![], + 0, + ); + desc_b_next.insert_at_back(latch, &ctx); + let desc_b_next = desc_b_next.deref(&ctx).get_result(0); + + Operation::new( + &mut ctx, + mir::MirGotoOp::get_concrete_op_info(), + vec![], + vec![i_next, desc_a_next, desc_b_next], + vec![header], + 0, + ) + .insert_at_back(latch, &ctx); + + // Rust calls split the exit sequence too: commit in the loop exit, then + // wait_group<0> in its unique linear successor. + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(exit, &ctx); + Operation::new( + &mut ctx, + mir::MirGotoOp::get_concrete_op_info(), + vec![], + vec![], + vec![wait_block], + 0, + ) + .insert_at_back(exit, &ctx); + append_wgmma_wait_group_constant(&mut ctx, wait_block, 0); + append_return(&mut ctx, wait_block); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let body = lowered_kernel_body(&ctx, module_ptr); + let matching = body + .iter() + .copied() + .filter_map(|operation| { + Operation::get_op::(operation, &ctx) + .map(|inline_asm| (operation, inline_asm)) + }) + .filter(|(_, asm)| { + asm.get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .is_some_and(|template| template.contains("L__wgmma_loop_${:uid}:")) + }) + .collect::>(); + assert_eq!( + matching.len(), + 1, + "expected one fused counted-loop WGMMA asm" + ); + + let (asm_operation, asm) = &matching[0]; + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .expect("counted-loop WGMMA template"); + + assert_eq!(template.matches("wgmma.fence.sync.aligned").count(), 1); + assert_eq!( + template.matches("wgmma.mma_async").count(), + 1, + "the PTX template contains one MMA instruction controlled by the internal K-loop" + ); + assert_eq!( + template.matches("wgmma.commit_group.sync.aligned").count(), + 1 + ); + assert_eq!( + template.matches("wgmma.wait_group.sync.aligned 0").count(), + 1 + ); + assert!(template.contains("mov.u64 %desc_a, $64;")); + assert!(template.contains("mov.u64 %desc_b, $65;")); + assert!(template.contains("add.u64 %desc_a, %desc_a, $66;")); + assert!(template.contains("add.u64 %desc_b, %desc_b, $67;")); + assert!(template.contains("mov.u64 %remaining, $68;")); + assert!(template.contains("@%loop_more bra.uni L__wgmma_done_${:uid};")); + assert!(template.contains("@%loop_more bra.uni L__wgmma_loop_${:uid};")); + assert!(template.contains("L__wgmma_done_${:uid}:")); + assert!( + !template.contains(".reg .f32") + && !template.contains("ld.f32") + && !template.contains("st.f32"), + "counted-loop WGMMA must keep accumulator memory outside asm: {template}" + ); + + let mut expected_constraints = vec!["=f".to_owned(); ACCUMULATOR_LEN]; + expected_constraints.extend((0..ACCUMULATOR_LEN).map(|index| index.to_string())); + expected_constraints.extend((0..LOOP_CONTROL_COUNT).map(|_| "l".to_owned())); + expected_constraints.push("~{memory}".to_owned()); + let expected_constraints = expected_constraints.join(","); + assert_eq!( + asm.get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some(expected_constraints.as_str()) + ); + assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); + assert_eq!( + asm_operation.deref(&ctx).get_num_operands(), + ACCUMULATOR_LEN + LOOP_CONTROL_COUNT + ); + + let asm_position = body + .iter() + .position(|operation| operation == asm_operation) + .expect("counted-loop WGMMA asm must be in the lowered kernel body"); + + let load_positions = body + .iter() + .enumerate() + .filter_map(|(index, operation)| { + Operation::get_op::(*operation, &ctx) + .is_some() + .then_some(index) + }) + .collect::>(); + assert_eq!( + load_positions.len(), + ACCUMULATOR_LEN, + "K-loop accumulator must be loaded exactly once, not once per iteration" + ); + assert!(load_positions.iter().all(|index| *index < asm_position)); + + let store_positions = body + .iter() + .enumerate() + .filter_map(|(index, operation)| { + Operation::get_op::(*operation, &ctx) + .is_some() + .then_some(index) + }) + .collect::>(); + assert_eq!( + store_positions.len(), + ACCUMULATOR_LEN, + "K-loop accumulator must be stored exactly once after the final wait" + ); + assert!(store_positions.iter().all(|index| *index > asm_position)); + + assert_eq!( + body.iter() + .filter( + |operation| Operation::get_op::(**operation, &ctx).is_some() + ) + .count(), + ACCUMULATOR_LEN + ); + + Ok(()) +} + +#[test] +fn test_f16_wgmma_counted_k_loop_stays_register_resident() -> Result<(), anyhow::Error> { + use dialect_mir::types::{MirArrayType, MirPtrType}; + use pliron::basic_block::BasicBlock; + use pliron::builtin::op_interfaces::OperandSegmentInterface; + use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; + + const ACCUMULATOR_LEN: usize = 32; + const LOOP_CONTROL_COUNT: usize = 5; + const TRIP_COUNT: u64 = 4; + const DESC_A_STEP: u64 = 16; + const DESC_B_STEP: u64 = 32; + + let mut ctx = make_test_ctx(); + let f32_ty = FP32Type::get(&ctx); + let row_ty = MirArrayType::get(&mut ctx, f32_ty.into(), 8); + let accumulator_ty = MirArrayType::get(&mut ctx, row_ty.into(), 4); + let accumulator_ptr_ty = MirPtrType::get_generic(&mut ctx, accumulator_ty.into(), true); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); + + let (module_ptr, preheader) = build_test_kernel( + &mut ctx, + vec![accumulator_ptr_ty.into(), u64_ty.into(), u64_ty.into()], + ); + let accumulator = preheader.deref(&ctx).get_argument(0); + let desc_a_base = preheader.deref(&ctx).get_argument(1); + let desc_b_base = preheader.deref(&ctx).get_argument(2); + + let module_region = module_ptr.deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + let function = module_block.deref(&ctx).iter(&ctx).next().unwrap(); + let function_region = function.deref(&ctx).get_region(0); + + let header = BasicBlock::new( + &mut ctx, + None, + vec![u32_ty.into(), u64_ty.into(), u64_ty.into()], + ); + header.insert_at_back(function_region, &ctx); + let latch = BasicBlock::new(&mut ctx, None, vec![]); + latch.insert_at_back(function_region, &ctx); + let exit = BasicBlock::new(&mut ctx, None, vec![]); + exit.insert_at_back(function_region, &ctx); + + // preheader: fence; i0 = 0; goto header(i0, desc_a_base, desc_b_base) + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(preheader, &ctx); + let i0 = append_mir_unsigned_constant(&mut ctx, preheader, u32_ty, 0); + Operation::new( + &mut ctx, + mir::MirGotoOp::get_concrete_op_info(), + vec![], + vec![i0, desc_a_base, desc_b_base], + vec![header], + 0, + ) + .insert_at_back(preheader, &ctx); + + // header(i, desc_a, desc_b): if !(i < 4) exit else latch. + let i = header.deref(&ctx).get_argument(0); + let desc_a = header.deref(&ctx).get_argument(1); + let desc_b = header.deref(&ctx).get_argument(2); + let bound = append_mir_unsigned_constant(&mut ctx, header, u32_ty, TRIP_COUNT); + let lt = Operation::new( + &mut ctx, + mir::MirLtOp::get_concrete_op_info(), + vec![i1_ty.into()], + vec![i, bound], + vec![], + 0, + ); + lt.insert_at_back(header, &ctx); + let lt_value = lt.deref(&ctx).get_result(0); + let not_lt = Operation::new( + &mut ctx, + mir::MirNotOp::get_concrete_op_info(), + vec![i1_ty.into()], + vec![lt_value], + vec![], + 0, + ); + not_lt.insert_at_back(header, &ctx); + let not_lt_value = not_lt.deref(&ctx).get_result(0); + let (branch_operands, segment_sizes) = + mir::MirCondBranchOp::compute_segment_sizes(vec![vec![not_lt_value], vec![], vec![]]); + let branch = Operation::new( + &mut ctx, + mir::MirCondBranchOp::get_concrete_op_info(), + vec![], + branch_operands, + vec![exit, latch], + 0, + ); + Operation::get_op::(branch, &ctx) + .expect("MirCondBranchOp") + .set_operand_segment_sizes(&ctx, segment_sizes); + branch.insert_at_back(header, &ctx); + + // latch: one WGMMA per K iteration and affine descriptor recurrences. + append_pointer_wgmma_mma_f16(&mut ctx, latch, accumulator, desc_a, desc_b); + + let one = append_mir_unsigned_constant(&mut ctx, latch, u32_ty, 1); + let i_next = Operation::new( + &mut ctx, + mir::MirAddOp::get_concrete_op_info(), + vec![u32_ty.into()], + vec![i, one], + vec![], + 0, + ); + i_next.insert_at_back(latch, &ctx); + let i_next = i_next.deref(&ctx).get_result(0); + + let desc_a_step = append_mir_unsigned_constant(&mut ctx, latch, u64_ty, DESC_A_STEP); + let desc_a_next = Operation::new( + &mut ctx, + mir::MirAddOp::get_concrete_op_info(), + vec![u64_ty.into()], + vec![desc_a, desc_a_step], + vec![], + 0, + ); + desc_a_next.insert_at_back(latch, &ctx); + let desc_a_next = desc_a_next.deref(&ctx).get_result(0); + + let desc_b_step = append_mir_unsigned_constant(&mut ctx, latch, u64_ty, DESC_B_STEP); + let desc_b_next = Operation::new( + &mut ctx, + mir::MirAddOp::get_concrete_op_info(), + vec![u64_ty.into()], + vec![desc_b, desc_b_step], + vec![], + 0, + ); + desc_b_next.insert_at_back(latch, &ctx); + let desc_b_next = desc_b_next.deref(&ctx).get_result(0); + + Operation::new( + &mut ctx, + mir::MirGotoOp::get_concrete_op_info(), + vec![], + vec![i_next, desc_a_next, desc_b_next], + vec![header], + 0, + ) + .insert_at_back(latch, &ctx); + + // exit: the only place where the asynchronous lifetime may become visible. + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(exit, &ctx); + append_wgmma_wait_group_constant(&mut ctx, exit, 0); + append_return(&mut ctx, exit); + + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let body = lowered_kernel_body(&ctx, module_ptr); + let matching = body + .iter() + .copied() + .filter_map(|operation| { + Operation::get_op::(operation, &ctx) + .map(|inline_asm| (operation, inline_asm)) + }) + .filter(|(_, asm)| { + asm.get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .is_some_and(|template| template.contains("L__wgmma_loop_${:uid}:")) + }) + .collect::>(); + assert_eq!( + matching.len(), + 1, + "expected one fused F16 counted-loop WGMMA asm" + ); + + let (asm_operation, asm) = &matching[0]; + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .expect("F16 counted-loop WGMMA template"); + + assert_eq!(template.matches("wgmma.fence.sync.aligned").count(), 1); + assert_eq!( + template + .matches("wgmma.mma_async.sync.aligned.m64n64k16.f32.f16.f16") + .count(), + 1, + "the F16 counted-loop template must contain exactly one F16 MMA" + ); + assert!(!template.contains(".bf16.bf16")); + assert_eq!( + template.matches("wgmma.commit_group.sync.aligned").count(), + 1 + ); + assert_eq!( + template.matches("wgmma.wait_group.sync.aligned 0").count(), + 1 + ); + assert!(template.contains("mov.u64 %desc_a, $64;")); + assert!(template.contains("mov.u64 %desc_b, $65;")); + assert!(template.contains("add.u64 %desc_a, %desc_a, $66;")); + assert!(template.contains("add.u64 %desc_b, %desc_b, $67;")); + assert!(template.contains("mov.u64 %remaining, $68;")); + assert!(template.contains("@%loop_more bra.uni L__wgmma_done_${:uid};")); + assert!(template.contains("@%loop_more bra.uni L__wgmma_loop_${:uid};")); + assert!(template.contains("L__wgmma_done_${:uid}:")); + assert!( + !template.contains(".reg .f32") + && !template.contains("ld.f32") + && !template.contains("st.f32"), + "F16 counted-loop WGMMA must keep accumulator memory outside asm: {template}" + ); + + let mut expected_constraints = vec!["=f".to_owned(); ACCUMULATOR_LEN]; + expected_constraints.extend((0..ACCUMULATOR_LEN).map(|index| index.to_string())); + expected_constraints.extend((0..LOOP_CONTROL_COUNT).map(|_| "l".to_owned())); + expected_constraints.push("~{memory}".to_owned()); + let expected_constraints = expected_constraints.join(","); + assert_eq!( + asm.get_attr_inline_asm_constraints(&ctx) + .map(|value| String::from((*value).clone())) + .as_deref(), + Some(expected_constraints.as_str()) + ); + assert_eq!(llvm::asm_kind(&ctx, asm), llvm::AsmKind::Convergent); + assert_eq!( + asm_operation.deref(&ctx).get_num_operands(), + ACCUMULATOR_LEN + LOOP_CONTROL_COUNT + ); + + let asm_position = body + .iter() + .position(|operation| operation == asm_operation) + .expect("F16 counted-loop WGMMA asm must be in the lowered kernel body"); + + let load_positions = body + .iter() + .enumerate() + .filter_map(|(index, operation)| { + Operation::get_op::(*operation, &ctx) + .is_some() + .then_some(index) + }) + .collect::>(); + assert_eq!( + load_positions.len(), + ACCUMULATOR_LEN, + "F16 K-loop accumulator must be loaded exactly once" + ); + assert!(load_positions.iter().all(|index| *index < asm_position)); + + let store_positions = body + .iter() + .enumerate() + .filter_map(|(index, operation)| { + Operation::get_op::(*operation, &ctx) + .is_some() + .then_some(index) + }) + .collect::>(); + assert_eq!( + store_positions.len(), + ACCUMULATOR_LEN, + "F16 K-loop accumulator must be stored exactly once after the final wait" + ); + assert!(store_positions.iter().all(|index| *index > asm_position)); + + assert_eq!( + body.iter() + .filter( + |operation| Operation::get_op::(**operation, &ctx).is_some() + ) + .count(), + ACCUMULATOR_LEN + ); + + Ok(()) +} diff --git a/crates/mir-lower/tests/lowering_test/wgmma_rejections.rs b/crates/mir-lower/tests/lowering_test/wgmma_rejections.rs new file mode 100644 index 0000000000..1e5af82d93 --- /dev/null +++ b/crates/mir-lower/tests/lowering_test/wgmma_rejections.rs @@ -0,0 +1,747 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use dialect_mir::ops as mir; +use dialect_nvvm::ops as nvvm; +use pliron::context::Context; +use pliron::linked_list::ContainsLinkedList; +use pliron::op::Op; +use pliron::operation::Operation; + +use crate::common::{append_return, build_test_kernel, make_test_ctx}; +use crate::wgmma_lowering::{ + append_mir_unsigned_constant, append_pointer_wgmma_mma, append_pointer_wgmma_mma_f16, + append_pointer_wgmma_mma_m64n128, append_pointer_wgmma_mma_tf32, + append_wgmma_wait_group_constant, build_pointer_form_wgmma_counted_pipeline_case, + build_wgmma_canonical_pointer_test_kernel, build_wgmma_pointer_test_kernel, +}; + +fn assert_wgmma_lowering_rejected( + ctx: &mut Context, + module_ptr: pliron::context::Ptr, + expected_diagnostic: &str, +) { + let error = mir_lower::lower_mir_to_llvm(ctx, module_ptr) + .expect_err("invalid deferred WGMMA sequence must fail closed") + .to_string(); + + assert!( + error.contains(expected_diagnostic), + "expected diagnostic containing `{expected_diagnostic}`, got:\n{error}" + ); +} + +// --------------------------------------------------------------------------- +// mma.sync m16n8k16 f16 intrinsic lowering test +// --------------------------------------------------------------------------- + +#[test] +fn test_pointer_form_wgmma_counted_pipeline_rejects_wait_two_with_two_slots() { + let mut ctx = make_test_ctx(); + let module_ptr = build_pointer_form_wgmma_counted_pipeline_case(&mut ctx, 2, &[2, 2], false); + + assert!( + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).is_err(), + "wait_group<2> must require three counted-pipeline accumulator slots" + ); +} + +#[test] +fn test_pointer_form_wgmma_counted_pipeline_rejects_mixed_partial_waits() { + let mut ctx = make_test_ctx(); + let module_ptr = build_pointer_form_wgmma_counted_pipeline_case(&mut ctx, 3, &[2, 1, 2], false); + + assert!( + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).is_err(), + "all counted-pipeline stages must use the same partial-wait depth" + ); +} + +#[test] +fn test_pointer_form_wgmma_counted_pipeline_rejects_reused_accumulator_slot() { + let mut ctx = make_test_ctx(); + let module_ptr = build_pointer_form_wgmma_counted_pipeline_case(&mut ctx, 3, &[2, 2, 2], true); + + assert!( + mir_lower::lower_mir_to_llvm(&mut ctx, module_ptr).is_err(), + "counted-pipeline accumulator slots must be pairwise distinct" + ); +} + +#[test] +fn test_tf32_wgmma_counted_k_loop_remains_unsupported() -> Result<(), anyhow::Error> { + use dialect_mir::types::{MirArrayType, MirPtrType}; + use pliron::basic_block::BasicBlock; + use pliron::builtin::op_interfaces::OperandSegmentInterface; + use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; + + const TRIP_COUNT: u64 = 4; + const DESC_A_STEP: u64 = 16; + const DESC_B_STEP: u64 = 32; + + let mut ctx = make_test_ctx(); + let f32_ty = FP32Type::get(&ctx); + let row_ty = MirArrayType::get(&mut ctx, f32_ty.into(), 8); + let accumulator_ty = MirArrayType::get(&mut ctx, row_ty.into(), 4); + let accumulator_ptr_ty = MirPtrType::get_generic(&mut ctx, accumulator_ty.into(), true); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); + + let (module_ptr, preheader) = build_test_kernel( + &mut ctx, + vec![accumulator_ptr_ty.into(), u64_ty.into(), u64_ty.into()], + ); + let accumulator = preheader.deref(&ctx).get_argument(0); + let desc_a_base = preheader.deref(&ctx).get_argument(1); + let desc_b_base = preheader.deref(&ctx).get_argument(2); + + let module_region = module_ptr.deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + let function = module_block.deref(&ctx).iter(&ctx).next().unwrap(); + let function_region = function.deref(&ctx).get_region(0); + + let header = BasicBlock::new( + &mut ctx, + None, + vec![u32_ty.into(), u64_ty.into(), u64_ty.into()], + ); + header.insert_at_back(function_region, &ctx); + let latch = BasicBlock::new(&mut ctx, None, vec![]); + latch.insert_at_back(function_region, &ctx); + let exit = BasicBlock::new(&mut ctx, None, vec![]); + exit.insert_at_back(function_region, &ctx); + + // preheader: fence; i0 = 0; goto header(i0, desc_a_base, desc_b_base) + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(preheader, &ctx); + let i0 = append_mir_unsigned_constant(&mut ctx, preheader, u32_ty, 0); + Operation::new( + &mut ctx, + mir::MirGotoOp::get_concrete_op_info(), + vec![], + vec![i0, desc_a_base, desc_b_base], + vec![header], + 0, + ) + .insert_at_back(preheader, &ctx); + + // header(i, desc_a, desc_b): if !(i < 4) exit else latch. + let i = header.deref(&ctx).get_argument(0); + let desc_a = header.deref(&ctx).get_argument(1); + let desc_b = header.deref(&ctx).get_argument(2); + let bound = append_mir_unsigned_constant(&mut ctx, header, u32_ty, TRIP_COUNT); + let lt = Operation::new( + &mut ctx, + mir::MirLtOp::get_concrete_op_info(), + vec![i1_ty.into()], + vec![i, bound], + vec![], + 0, + ); + lt.insert_at_back(header, &ctx); + let lt_value = lt.deref(&ctx).get_result(0); + let not_lt = Operation::new( + &mut ctx, + mir::MirNotOp::get_concrete_op_info(), + vec![i1_ty.into()], + vec![lt_value], + vec![], + 0, + ); + not_lt.insert_at_back(header, &ctx); + let not_lt_value = not_lt.deref(&ctx).get_result(0); + let (branch_operands, segment_sizes) = + mir::MirCondBranchOp::compute_segment_sizes(vec![vec![not_lt_value], vec![], vec![]]); + let branch = Operation::new( + &mut ctx, + mir::MirCondBranchOp::get_concrete_op_info(), + vec![], + branch_operands, + vec![exit, latch], + 0, + ); + Operation::get_op::(branch, &ctx) + .expect("MirCondBranchOp") + .set_operand_segment_sizes(&ctx, segment_sizes); + branch.insert_at_back(header, &ctx); + + // latch: one WGMMA per K iteration and affine descriptor recurrences. + append_pointer_wgmma_mma_tf32(&mut ctx, latch, accumulator, desc_a, desc_b); + + let one = append_mir_unsigned_constant(&mut ctx, latch, u32_ty, 1); + let i_next = Operation::new( + &mut ctx, + mir::MirAddOp::get_concrete_op_info(), + vec![u32_ty.into()], + vec![i, one], + vec![], + 0, + ); + i_next.insert_at_back(latch, &ctx); + let i_next = i_next.deref(&ctx).get_result(0); + + let desc_a_step = append_mir_unsigned_constant(&mut ctx, latch, u64_ty, DESC_A_STEP); + let desc_a_next = Operation::new( + &mut ctx, + mir::MirAddOp::get_concrete_op_info(), + vec![u64_ty.into()], + vec![desc_a, desc_a_step], + vec![], + 0, + ); + desc_a_next.insert_at_back(latch, &ctx); + let desc_a_next = desc_a_next.deref(&ctx).get_result(0); + + let desc_b_step = append_mir_unsigned_constant(&mut ctx, latch, u64_ty, DESC_B_STEP); + let desc_b_next = Operation::new( + &mut ctx, + mir::MirAddOp::get_concrete_op_info(), + vec![u64_ty.into()], + vec![desc_b, desc_b_step], + vec![], + 0, + ); + desc_b_next.insert_at_back(latch, &ctx); + let desc_b_next = desc_b_next.deref(&ctx).get_result(0); + + Operation::new( + &mut ctx, + mir::MirGotoOp::get_concrete_op_info(), + vec![], + vec![i_next, desc_a_next, desc_b_next], + vec![header], + 0, + ) + .insert_at_back(latch, &ctx); + + // exit: the only place where the asynchronous lifetime may become visible. + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(exit, &ctx); + append_wgmma_wait_group_constant(&mut ctx, exit, 0); + append_return(&mut ctx, exit); + + assert_wgmma_lowering_rejected( + &mut ctx, + module_ptr, + "WGMMA MMA reached lowering without deferred accumulator fusion", + ); + Ok(()) +} + +#[test] +fn test_m64n128_wgmma_noncanonical_accumulator_has_no_pointer_fallback() -> Result<(), anyhow::Error> +{ + let mut ctx = make_test_ctx(); + let (module_ptr, entry, accumulators, desc_a, desc_b, _) = + build_wgmma_pointer_test_kernel(&mut ctx, 1, vec![]); + + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_pointer_wgmma_mma_m64n128(&mut ctx, entry, accumulators[0], desc_a, desc_b); + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_wgmma_wait_group_constant(&mut ctx, entry, 0); + append_return(&mut ctx, entry); + + assert_wgmma_lowering_rejected( + &mut ctx, + module_ptr, + "WGMMA linear full-drain lowering for this variant requires a canonical [[f32; 8]; 8] accumulator", + ); + Ok(()) +} + +#[test] +fn test_f16_wgmma_noncanonical_accumulator_has_no_pointer_fallback() -> Result<(), anyhow::Error> { + let mut ctx = make_test_ctx(); + let (module_ptr, entry, accumulators, desc_a, desc_b, _) = + build_wgmma_pointer_test_kernel(&mut ctx, 1, vec![]); + + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_pointer_wgmma_mma_f16(&mut ctx, entry, accumulators[0], desc_a, desc_b); + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_wgmma_wait_group_constant(&mut ctx, entry, 0); + append_return(&mut ctx, entry); + + assert_wgmma_lowering_rejected( + &mut ctx, + module_ptr, + "WGMMA linear full-drain lowering for this variant requires a canonical [[f32; 8]; 4] accumulator", + ); + Ok(()) +} + +#[test] +fn test_tf32_wgmma_noncanonical_accumulator_has_no_pointer_fallback() -> Result<(), anyhow::Error> { + let mut ctx = make_test_ctx(); + let (module_ptr, entry, accumulators, desc_a, desc_b, _) = + build_wgmma_pointer_test_kernel(&mut ctx, 1, vec![]); + + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_pointer_wgmma_mma_tf32(&mut ctx, entry, accumulators[0], desc_a, desc_b); + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_wgmma_wait_group_constant(&mut ctx, entry, 0); + append_return(&mut ctx, entry); + + assert_wgmma_lowering_rejected( + &mut ctx, + module_ptr, + "WGMMA linear full-drain lowering for this variant requires a canonical [[f32; 8]; 4] accumulator", + ); + Ok(()) +} + +#[test] +fn test_linear_wgmma_full_drain_rejects_mixed_bf16_and_f16() -> Result<(), anyhow::Error> { + let mut ctx = make_test_ctx(); + let (module_ptr, entry, accumulators, descriptors) = + build_wgmma_canonical_pointer_test_kernel(&mut ctx, 1, 4); + let accumulator = accumulators[0]; + + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_pointer_wgmma_mma(&mut ctx, entry, accumulator, descriptors[0], descriptors[1]); + append_pointer_wgmma_mma_f16(&mut ctx, entry, accumulator, descriptors[2], descriptors[3]); + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_wgmma_wait_group_constant(&mut ctx, entry, 0); + append_return(&mut ctx, entry); + + assert_wgmma_lowering_rejected( + &mut ctx, + module_ptr, + "one linear WGMMA full-drain region cannot mix MMA variants or shapes", + ); + Ok(()) +} + +#[test] +fn test_linear_wgmma_full_drain_rejects_mixed_bf16_and_tf32() -> Result<(), anyhow::Error> { + let mut ctx = make_test_ctx(); + let (module_ptr, entry, accumulators, descriptors) = + build_wgmma_canonical_pointer_test_kernel(&mut ctx, 1, 4); + let accumulator = accumulators[0]; + + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_pointer_wgmma_mma(&mut ctx, entry, accumulator, descriptors[0], descriptors[1]); + append_pointer_wgmma_mma_tf32(&mut ctx, entry, accumulator, descriptors[2], descriptors[3]); + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_wgmma_wait_group_constant(&mut ctx, entry, 0); + append_return(&mut ctx, entry); + + assert_wgmma_lowering_rejected( + &mut ctx, + module_ptr, + "one linear WGMMA full-drain region cannot mix MMA variants or shapes", + ); + Ok(()) +} + +#[test] +fn test_linear_wgmma_full_drain_rejects_mixed_f16_and_tf32() -> Result<(), anyhow::Error> { + let mut ctx = make_test_ctx(); + let (module_ptr, entry, accumulators, descriptors) = + build_wgmma_canonical_pointer_test_kernel(&mut ctx, 1, 4); + let accumulator = accumulators[0]; + + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_pointer_wgmma_mma_f16(&mut ctx, entry, accumulator, descriptors[0], descriptors[1]); + append_pointer_wgmma_mma_tf32(&mut ctx, entry, accumulator, descriptors[2], descriptors[3]); + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_wgmma_wait_group_constant(&mut ctx, entry, 0); + append_return(&mut ctx, entry); + + assert_wgmma_lowering_rejected( + &mut ctx, + module_ptr, + "one linear WGMMA full-drain region cannot mix MMA variants or shapes", + ); + Ok(()) +} + +#[test] +fn test_f16_wgmma_partial_wait_remains_unsupported() -> Result<(), anyhow::Error> { + let mut ctx = make_test_ctx(); + let (module_ptr, entry, accumulators, descriptors) = + build_wgmma_canonical_pointer_test_kernel(&mut ctx, 1, 2); + + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_pointer_wgmma_mma_f16( + &mut ctx, + entry, + accumulators[0], + descriptors[0], + descriptors[1], + ); + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_wgmma_wait_group_constant(&mut ctx, entry, 1); + append_wgmma_wait_group_constant(&mut ctx, entry, 0); + append_return(&mut ctx, entry); + + assert_wgmma_lowering_rejected( + &mut ctx, + module_ptr, + "WGMMA deferred accumulator lowering requires wait_group<0>", + ); + Ok(()) +} + +#[test] +fn test_tf32_wgmma_partial_wait_remains_unsupported() -> Result<(), anyhow::Error> { + let mut ctx = make_test_ctx(); + let (module_ptr, entry, accumulators, descriptors) = + build_wgmma_canonical_pointer_test_kernel(&mut ctx, 1, 2); + + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_pointer_wgmma_mma_tf32( + &mut ctx, + entry, + accumulators[0], + descriptors[0], + descriptors[1], + ); + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_wgmma_wait_group_constant(&mut ctx, entry, 1); + append_wgmma_wait_group_constant(&mut ctx, entry, 0); + append_return(&mut ctx, entry); + + assert_wgmma_lowering_rejected( + &mut ctx, + module_ptr, + "WGMMA deferred accumulator lowering requires wait_group<0>", + ); + Ok(()) +} + +#[test] +fn test_pointer_form_wgmma_without_fence_is_rejected() -> Result<(), anyhow::Error> { + let mut ctx = make_test_ctx(); + let (module_ptr, entry, accumulators, desc_a, desc_b, _) = + build_wgmma_pointer_test_kernel(&mut ctx, 1, vec![]); + + append_pointer_wgmma_mma(&mut ctx, entry, accumulators[0], desc_a, desc_b); + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_wgmma_wait_group_constant(&mut ctx, entry, 0); + append_return(&mut ctx, entry); + + assert_wgmma_lowering_rejected( + &mut ctx, + module_ptr, + "WGMMA MMA reached lowering without deferred accumulator fusion", + ); + Ok(()) +} + +#[test] +fn test_deferred_wgmma_without_commit_is_rejected() -> Result<(), anyhow::Error> { + let mut ctx = make_test_ctx(); + let (module_ptr, entry, accumulators, desc_a, desc_b, _) = + build_wgmma_pointer_test_kernel(&mut ctx, 1, vec![]); + + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_pointer_wgmma_mma(&mut ctx, entry, accumulators[0], desc_a, desc_b); + append_wgmma_wait_group_constant(&mut ctx, entry, 0); + append_return(&mut ctx, entry); + + assert_wgmma_lowering_rejected( + &mut ctx, + module_ptr, + "WGMMA wait_group requires a preceding commit_group", + ); + Ok(()) +} + +#[test] +fn test_wgmma_partial_wait_without_final_wait_zero_is_rejected() -> Result<(), anyhow::Error> { + let mut ctx = make_test_ctx(); + let (module_ptr, entry, accumulators, desc_a, desc_b, _) = + build_wgmma_pointer_test_kernel(&mut ctx, 1, vec![]); + + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_pointer_wgmma_mma(&mut ctx, entry, accumulators[0], desc_a, desc_b); + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_wgmma_wait_group_constant(&mut ctx, entry, 1); + append_return(&mut ctx, entry); + + assert_wgmma_lowering_rejected( + &mut ctx, + module_ptr, + "WGMMA partial-wait pipeline requires a final wait_group<0>", + ); + Ok(()) +} + +#[test] +fn test_wgmma_partial_wait_pipeline_rejects_unsafe_accumulator_reuse() -> Result<(), anyhow::Error> +{ + let mut ctx = make_test_ctx(); + let (module_ptr, entry, accumulators, descriptors) = + build_wgmma_canonical_pointer_test_kernel(&mut ctx, 1, 4); + + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_pointer_wgmma_mma( + &mut ctx, + entry, + accumulators[0], + descriptors[0], + descriptors[1], + ); + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_pointer_wgmma_mma( + &mut ctx, + entry, + accumulators[0], + descriptors[2], + descriptors[3], + ); + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_wgmma_wait_group_constant(&mut ctx, entry, 1); + append_wgmma_wait_group_constant(&mut ctx, entry, 0); + append_return(&mut ctx, entry); + + assert_wgmma_lowering_rejected( + &mut ctx, + module_ptr, + "WGMMA partial-wait pipeline requires max_pending_groups + 1 distinct accumulator slots", + ); + Ok(()) +} + +#[test] +fn test_wgmma_wait_group_eight_is_rejected() -> Result<(), anyhow::Error> { + let mut ctx = make_test_ctx(); + let (module_ptr, entry, accumulators, desc_a, desc_b, _) = + build_wgmma_pointer_test_kernel(&mut ctx, 1, vec![]); + + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_pointer_wgmma_mma(&mut ctx, entry, accumulators[0], desc_a, desc_b); + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_wgmma_wait_group_constant(&mut ctx, entry, 8); + append_return(&mut ctx, entry); + + assert_wgmma_lowering_rejected( + &mut ctx, + module_ptr, + "WGMMA wait_group immediate must be in 0..=7", + ); + Ok(()) +} + +#[test] +fn test_deferred_wgmma_two_commits_are_rejected() -> Result<(), anyhow::Error> { + let mut ctx = make_test_ctx(); + let (module_ptr, entry, accumulators, desc_a, desc_b, _) = + build_wgmma_pointer_test_kernel(&mut ctx, 1, vec![]); + + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_pointer_wgmma_mma(&mut ctx, entry, accumulators[0], desc_a, desc_b); + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_wgmma_wait_group_constant(&mut ctx, entry, 0); + append_return(&mut ctx, entry); + + assert_wgmma_lowering_rejected( + &mut ctx, + module_ptr, + "WGMMA deferred accumulator region supports exactly one commit_group", + ); + Ok(()) +} + +#[test] +fn test_deferred_wgmma_multiple_accumulators_are_rejected() -> Result<(), anyhow::Error> { + let mut ctx = make_test_ctx(); + let (module_ptr, entry, accumulators, desc_a, desc_b, _) = + build_wgmma_pointer_test_kernel(&mut ctx, 2, vec![]); + + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_pointer_wgmma_mma(&mut ctx, entry, accumulators[0], desc_a, desc_b); + append_pointer_wgmma_mma(&mut ctx, entry, accumulators[1], desc_a, desc_b); + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_wgmma_wait_group_constant(&mut ctx, entry, 0); + append_return(&mut ctx, entry); + + assert_wgmma_lowering_rejected( + &mut ctx, + module_ptr, + "WGMMA deferred accumulator region uses more than one accumulator", + ); + Ok(()) +} + +#[test] +fn test_deferred_wgmma_mma_after_commit_is_rejected() -> Result<(), anyhow::Error> { + let mut ctx = make_test_ctx(); + let (module_ptr, entry, accumulators, desc_a, desc_b, _) = + build_wgmma_pointer_test_kernel(&mut ctx, 1, vec![]); + + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_pointer_wgmma_mma(&mut ctx, entry, accumulators[0], desc_a, desc_b); + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_pointer_wgmma_mma(&mut ctx, entry, accumulators[0], desc_a, desc_b); + append_wgmma_wait_group_constant(&mut ctx, entry, 0); + append_return(&mut ctx, entry); + + assert_wgmma_lowering_rejected( + &mut ctx, + module_ptr, + "WGMMA MMA cannot appear after commit_group in a deferred accumulator region", + ); + Ok(()) +} + +#[test] +fn test_deferred_wgmma_branch_is_rejected() -> Result<(), anyhow::Error> { + use pliron::basic_block::BasicBlock; + use pliron::builtin::op_interfaces::OperandSegmentInterface; + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let bool_ty = IntegerType::get(&ctx, 1, Signedness::Signless); + let (module_ptr, entry, accumulators, desc_a, desc_b, trailing) = + build_wgmma_pointer_test_kernel(&mut ctx, 1, vec![bool_ty.into()]); + let condition = trailing[0]; + + let module_region = module_ptr.deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + let function = module_block.deref(&ctx).iter(&ctx).next().unwrap(); + let function_region = function.deref(&ctx).get_region(0); + + let then_block = BasicBlock::new(&mut ctx, None, vec![]); + then_block.insert_at_back(function_region, &ctx); + let else_block = BasicBlock::new(&mut ctx, None, vec![]); + else_block.insert_at_back(function_region, &ctx); + + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_pointer_wgmma_mma(&mut ctx, entry, accumulators[0], desc_a, desc_b); + + let (flat_operands, segment_sizes) = + mir::MirCondBranchOp::compute_segment_sizes(vec![vec![condition], vec![], vec![]]); + let branch = Operation::new( + &mut ctx, + mir::MirCondBranchOp::get_concrete_op_info(), + vec![], + flat_operands, + vec![then_block, else_block], + 0, + ); + Operation::get_op::(branch, &ctx) + .expect("MirCondBranchOp") + .set_operand_segment_sizes(&ctx, segment_sizes); + branch.insert_at_back(entry, &ctx); + + append_return(&mut ctx, then_block); + append_return(&mut ctx, else_block); + + assert_wgmma_lowering_rejected( + &mut ctx, + module_ptr, + "unsupported operation inside WGMMA deferred accumulator region", + ); + Ok(()) +} + +#[test] +fn test_deferred_wgmma_join_is_rejected() -> Result<(), anyhow::Error> { + use pliron::basic_block::BasicBlock; + + let mut ctx = make_test_ctx(); + let (module_ptr, entry, accumulators, desc_a, desc_b, _) = + build_wgmma_pointer_test_kernel(&mut ctx, 1, vec![]); + + let module_region = module_ptr.deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + let function = module_block.deref(&ctx).iter(&ctx).next().unwrap(); + let function_region = function.deref(&ctx).get_region(0); + + let second_predecessor = BasicBlock::new(&mut ctx, None, vec![]); + second_predecessor.insert_at_back(function_region, &ctx); + let join = BasicBlock::new(&mut ctx, None, vec![]); + join.insert_at_back(function_region, &ctx); + + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_pointer_wgmma_mma(&mut ctx, entry, accumulators[0], desc_a, desc_b); + + Operation::new( + &mut ctx, + mir::MirGotoOp::get_concrete_op_info(), + vec![], + vec![], + vec![join], + 0, + ) + .insert_at_back(entry, &ctx); + + Operation::new( + &mut ctx, + mir::MirGotoOp::get_concrete_op_info(), + vec![], + vec![], + vec![join], + 0, + ) + .insert_at_back(second_predecessor, &ctx); + + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(join, &ctx); + append_wgmma_wait_group_constant(&mut ctx, join, 0); + append_return(&mut ctx, join); + + assert_wgmma_lowering_rejected( + &mut ctx, + module_ptr, + "WGMMA deferred accumulator region cannot cross a control-flow join", + ); + Ok(()) +} + +#[test] +fn test_deferred_wgmma_intervening_operation_is_rejected() -> Result<(), anyhow::Error> { + use pliron::builtin::types::{IntegerType, Signedness}; + + let mut ctx = make_test_ctx(); + let (module_ptr, entry, accumulators, desc_a, desc_b, _) = + build_wgmma_pointer_test_kernel(&mut ctx, 1, vec![]); + + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_pointer_wgmma_mma(&mut ctx, entry, accumulators[0], desc_a, desc_b); + + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + Operation::new( + &mut ctx, + nvvm::ReadPtxSregTidXOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![], + vec![], + 0, + ) + .insert_at_back(entry, &ctx); + + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_wgmma_wait_group_constant(&mut ctx, entry, 0); + append_return(&mut ctx, entry); + + assert_wgmma_lowering_rejected( + &mut ctx, + module_ptr, + "unsupported operation inside WGMMA deferred accumulator region", + ); + Ok(()) +} + +#[test] +fn test_deferred_wgmma_nested_fence_is_rejected() -> Result<(), anyhow::Error> { + let mut ctx = make_test_ctx(); + let (module_ptr, entry, accumulators, desc_a, desc_b, _) = + build_wgmma_pointer_test_kernel(&mut ctx, 1, vec![]); + + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_pointer_wgmma_mma(&mut ctx, entry, accumulators[0], desc_a, desc_b); + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_wgmma_wait_group_constant(&mut ctx, entry, 0); + append_return(&mut ctx, entry); + + assert_wgmma_lowering_rejected( + &mut ctx, + module_ptr, + "nested WGMMA fences are not supported in one deferred accumulator region", + ); + Ok(()) +} From d59199bd5749814ab71ddd16f03c789d4d1ef533 Mon Sep 17 00:00:00 2001 From: nihalpasham Date: Fri, 28 Aug 2026 22:21:07 +0530 Subject: [PATCH 05/18] refactor(mir-lower): split convert/types.rs into convert/types/ modules --- crates/mir-lower/src/convert/types.rs | 5845 ----------------- .../mir-lower/src/convert/types/conversion.rs | 153 + .../src/convert/types/enum_layout.rs | 2014 ++++++ .../mir-lower/src/convert/types/func_abi.rs | 1069 +++ .../src/convert/types/global_layout.rs | 826 +++ crates/mir-lower/src/convert/types/layout.rs | 548 ++ crates/mir-lower/src/convert/types/mod.rs | 145 + .../src/convert/types/pointer_storage.rs | 360 + .../src/convert/types/struct_layout.rs | 619 ++ .../src/convert/types/test_support.rs | 65 + .../src/convert/types/union_storage.rs | 276 + 11 files changed, 6075 insertions(+), 5845 deletions(-) delete mode 100644 crates/mir-lower/src/convert/types.rs create mode 100644 crates/mir-lower/src/convert/types/conversion.rs create mode 100644 crates/mir-lower/src/convert/types/enum_layout.rs create mode 100644 crates/mir-lower/src/convert/types/func_abi.rs create mode 100644 crates/mir-lower/src/convert/types/global_layout.rs create mode 100644 crates/mir-lower/src/convert/types/layout.rs create mode 100644 crates/mir-lower/src/convert/types/mod.rs create mode 100644 crates/mir-lower/src/convert/types/pointer_storage.rs create mode 100644 crates/mir-lower/src/convert/types/struct_layout.rs create mode 100644 crates/mir-lower/src/convert/types/test_support.rs create mode 100644 crates/mir-lower/src/convert/types/union_storage.rs diff --git a/crates/mir-lower/src/convert/types.rs b/crates/mir-lower/src/convert/types.rs deleted file mode 100644 index 79aef815d8..0000000000 --- a/crates/mir-lower/src/convert/types.rs +++ /dev/null @@ -1,5845 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -//! Type conversion from `dialect-mir` types to LLVM dialect types. -//! -//! This module handles the translation of `dialect-mir` type representations -//! to their LLVM dialect equivalents. Type conversion is foundational to -//! the lowering pass—most operation converters depend on it. -//! -//! # Overview -//! -//! `dialect-mir` types are high-level, Rust-like types that preserve semantic -//! information (signedness, slice semantics, etc.). LLVM dialect types are -//! lower-level and match LLVM IR types directly. -//! -//! # Type Mapping Table -//! -//! | `dialect-mir` Type | LLVM dialect Type | Notes | -//! |---------------------------------|-----------------------------------|-----------------------------| -//! | `IntegerType` (signed/unsigned) | `IntegerType` (signless) | Width preserved | -//! | `MirFP16Type` | `HalfType` | Rust `f16` → LLVM `half` | -//! | `FP32Type`, `FP64Type` | Same (builtin) | Pass-through | -//! | `MirPtrType` | `PointerType` | Address space preserved | -//! | `MirSliceType` | `StructType { ptr, i64 }` | Fat pointer | -//! | `MirDisjointSliceType` | `StructType { ptr, i64 }` | Same as slice | -//! | `MirTupleType` | `StructType` | Empty tuple → empty struct | -//! | `MirStructType` | `StructType` | Fields recursively converted| -//! | `MirUnionType` | Aligned shared-storage struct | All fields start at byte zero| -//! | `MirEnumType` | `StructType` (rustc byte layout) | See "Enum Type Representation" | -//! | `ArrayType` | `ArrayType` | Element type converted | -//! | `VectorType` | `VectorType` | Element type converted | -//! -//! # Signedness Handling -//! -//! LLVM IR integers are signless—the signedness is encoded in the operations -//! that use them (e.g., `sdiv` vs `udiv`). During type conversion: -//! -//! - Signed/unsigned MIR integers → signless LLVM integers -//! - The original signedness is preserved in operations (see `arithmetic.rs`) -//! -//! # Address Space Handling -//! -//! GPU memory uses address spaces to distinguish memory types: -//! -//! | Address Space | Memory Type | Usage | -//! |---------------|-------------|---------------------------| -//! | 0 | Generic | Can point to any memory | -//! | 1 | Global | Device memory (VRAM) | -//! | 3 | Shared | Per-block shared memory | -//! | 4 | Constant | Read-only device memory | -//! | 5 | Local | Per-thread stack/spill | -//! -//! Pointer address spaces are preserved through conversion. Slice types use -//! generic address space (0) because they can point to any memory type. -//! -//! # Slice Type Representation -//! -//! Rust slices (`&[T]`) are represented as fat pointers in LLVM: -//! -//! ```text -//! MIR: MirSliceType -//! LLVM: struct { ptr, i64 } ; pointer + length -//! ``` -//! -//! This matches the Rust ABI for slices passed by value. -//! -//! # Enum Type Representation -//! -//! A Rust enum is one tag plus the payload of whichever variant is -//! alive; all variants share the same bytes. We build an LLVM struct -//! that puts the tag and every payload field at the exact byte position -//! rustc chose, inserting `[N x i8]` filler for the gaps: -//! -//! ```text -//! #[repr(u32)] enum E { A(u32), B(f32), C } // rustc: 8 bytes, -//! // tag at 0, payloads at 4 -//! LLVM: { i32, i32 } ; slot 0 = tag, slot 1 = A's payload -//! ; B's f32 also lives at byte 4 but has a -//! ; different type, so it is read/written through -//! ; memory instead of owning a slot -//! ``` -//! -//! Because the bytes match rustc exactly, enum data can cross the -//! host/device boundary safely. The tag slot stores the variant's -//! DECLARED discriminant value (`enum E { A = 7 }` stores 7), not its -//! position. See `build_enum_slot_map` in this module for the full -//! story. -//! -//! # Function Type Conversion -//! -//! Function types undergo ABI transformations: -//! -//! - Slice arguments are flattened to `(ptr, len)` pairs -//! - Struct arguments are flattened to individual fields -//! - Empty tuple return type becomes void -//! -//! This matches the C ABI for GPU kernels. - -use dialect_mir::types::{ - EnumCarrierKind, EnumLayoutKind, MirArrayType, MirDisjointSliceType, MirEnumType, MirFP16Type, - MirPtrType, MirSliceType, MirStructType, MirTupleType, MirUnionType, -}; -use llvm_export::types as llvm_types; -use llvm_export::types::PointerTypeExt; -use pliron::builtin::type_interfaces::{FloatTypeInterface, FunctionTypeInterface}; -use pliron::builtin::types::{FP32Type, FP64Type, FunctionType, IntegerType, Signedness}; -use pliron::context::{Context, Ptr}; -use pliron::operation::Operation; -use pliron::r#type::{TypeHandle, type_cast}; - -use super::enum_payload_storage::{ - MAX_ENUM_PAYLOAD_ARRAY_REWRITE_LEAVES, enum_payload_storage_type, -}; -use super::target_stable_storage::{StorageRewriteOptions, target_stable_storage_type}; -use crate::type_conversion_interface::MirTypeConversion; - -// ============================================================================= -// Kernel-Boundary Detection -// ============================================================================= - -/// Identifier of the attribute that marks a `MirFuncOp` / `llvm.func` as a -/// GPU kernel entry point. -/// -/// Kept as a function (rather than a `const`) because pliron `Identifier` -/// construction needs the `try_into()` fallible path. -fn gpu_kernel_attr() -> pliron::identifier::Identifier { - "gpu_kernel".try_into().expect("static identifier") -} - -/// Returns `true` when `op` carries the `gpu_kernel` attribute. -/// -/// The kernel-entry ABI differs from internal device-function ABI: at -/// kernel boundaries, aggregate parameters (structs, closures) are passed -/// as a single byval value to match what the host pushes via -/// `cuLaunchKernel`. Internal call sites still flatten aggregates the -/// same way they always did. This helper is the single source of truth -/// for that branch and is consumed by both [`convert_function_type`] and -/// the entry-block prologue in `lowering.rs`. -pub fn is_kernel_func(ctx: &Context, op: Ptr) -> bool { - op.deref(ctx) - .attributes - .get::(&gpu_kernel_attr()) - .is_some() -} - -// ============================================================================= -// Zero-Sized Type (ZST) Detection -// ============================================================================= - -/// Check if a type is zero-sized (empty struct). -/// -/// Zero-sized types include: -/// - Empty structs `struct {}` -/// - PhantomData markers (which become empty structs in MIR) -/// - Structs where all fields are themselves zero-sized -/// -/// # Why This Matters -/// -/// LLVM's NVPTX backend doesn't support empty struct types in function -/// signatures. We strip these during type conversion to avoid: -/// `LLVM ERROR: Empty parameter types are not supported` -/// -/// # Background -/// -/// Rust's `#[inline(always)]` attribute is stored in `codegen_fn_attrs`, which -/// is not exposed through the stable_mir API. Since we intercept MIR and generate -/// our own LLVM IR, we don't propagate inline hints. When LLVM decides not to -/// inline a function, the empty struct parameters/returns cause NVPTX to crash. -/// -/// By stripping ZSTs at the LLVM type level, we avoid this issue regardless of -/// inlining decisions. -pub fn is_zero_sized_type(ctx: &Context, ty: TypeHandle) -> bool { - if let Some(array_ty) = ty.deref(ctx).downcast_ref::() { - return array_ty.size() == 0 || is_zero_sized_type(ctx, array_ty.elem_type()); - } - - // Check if LLVM StructType with zero fields - if let Some(struct_ty) = ty.deref(ctx).downcast_ref::() { - let num_fields = struct_ty.num_fields(); - if num_fields == 0 { - return true; - } - // Also check if ALL fields are zero-sized (nested PhantomData) - return struct_ty.fields().all(|f| is_zero_sized_type(ctx, f)); - } - false -} - -// ============================================================================= -// Type Conversion -// ============================================================================= - -/// Convert a `dialect-mir` type to its LLVM dialect equivalent. -/// -/// Dispatches via `MirTypeConversion` type interface — each supported type -/// registers a converter function pointer through `#[type_interface_impl]` -/// in [`super::type_interface_impls`]. -/// -/// The function-pointer indirection avoids a borrow-checker conflict: -/// `type_cast` borrows `ctx` immutably, but conversion needs `&mut ctx`. -/// We extract the `Copy` function pointer, drop the borrow, then call it. -pub fn convert_type(ctx: &mut Context, ty: TypeHandle) -> Result { - // Phase 1: extract a Copy function pointer while ctx is immutably borrowed. - let converter_fn = { - let ty_ref = ty.deref(ctx); - type_cast::(&*ty_ref).map(|conv| conv.converter()) - }; - // Phase 2: borrow dropped — ctx is free for &mut. - if let Some(conv_fn) = converter_fn { - return conv_fn(ty, ctx); - } - - let type_display = ty.deref(ctx).disp(ctx).to_string(); - Err(anyhow::anyhow!( - "Unsupported type conversion: {}\n\ - Supported: integers, fp32, fp64, pointers, slices, tuples, structs, enums, arrays, vectors.", - type_display - )) -} - -/// Return the declaration index and type of the single non-ZST field of a -/// rustc-proven transparent scalar struct. -/// -/// The importer marks the outer ABI from rustc rather than inferring it from -/// source field count. We still validate the MIR shape here so malformed or -/// hand-written dialect input cannot turn an arbitrary aggregate into a scalar -/// ABI value. -fn transparent_scalar_field_with_index( - ctx: &mut Context, - struct_ty: TypeHandle, -) -> Result<(usize, TypeHandle), anyhow::Error> { - let (name, field_types, mem_to_decl, is_transparent_scalar) = { - let ty_ref = struct_ty.deref(ctx); - let s = ty_ref - .downcast_ref::() - .ok_or_else(|| anyhow::anyhow!("transparent scalar ABI requires a MirStructType"))?; - ( - s.name.clone(), - s.field_types.clone(), - s.memory_order(), - s.is_transparent_scalar(), - ) - }; - - if !is_transparent_scalar { - return Err(anyhow::anyhow!( - "struct `{}` is not marked as a transparent scalar", - name - )); - } - - let mut scalar_field = None; - for decl_idx in mem_to_decl { - let field_ty = field_types[decl_idx]; - let converted = convert_type(ctx, field_ty)?; - if is_zero_sized_type(ctx, converted) { - continue; - } - if scalar_field.replace((decl_idx, field_ty)).is_some() { - return Err(anyhow::anyhow!( - "transparent scalar struct `{}` has more than one non-ZST field", - name - )); - } - } - - scalar_field - .ok_or_else(|| anyhow::anyhow!("transparent scalar struct `{}` has no non-ZST field", name)) -} - -/// Return the single non-ZST field of a rustc-proven transparent scalar struct. -pub(crate) fn transparent_scalar_field( - ctx: &mut Context, - struct_ty: TypeHandle, -) -> Result { - Ok(transparent_scalar_field_with_index(ctx, struct_ty)?.1) -} - -/// One aggregate layer traversed when a transparent scalar wrapper crosses an -/// ABI boundary. `field_slot` is the LLVM struct slot containing the next -/// nested wrapper or the final scalar. -#[derive(Clone, Copy, Debug)] -pub(crate) struct TransparentScalarLayer { - pub llvm_struct_ty: TypeHandle, - pub field_slot: u32, -} - -/// Complete ABI projection for a rustc-proven transparent scalar wrapper. -/// -/// `layers` are ordered outermost to innermost. A return lowers by extracting -/// those slots in order; a call result is rebuilt by inserting the scalar into -/// the same layers in reverse order. -#[derive(Clone, Debug)] -pub(crate) struct TransparentScalarAbiInfo { - pub scalar_ty: TypeHandle, - pub layers: Vec, -} - -/// Build the scalar ABI projection, including the exact LLVM slot used at every -/// wrapper layer. -/// -/// Slot indices come from [`build_struct_slot_map`], so ZST markers, explicit -/// padding, and rustc memory order cannot make the return/call paths disagree -/// with ordinary aggregate conversion. -pub(crate) fn transparent_scalar_abi_info( - ctx: &mut Context, - struct_ty: TypeHandle, -) -> Result { - let mut current = struct_ty; - let mut layers = Vec::new(); - - loop { - let (decl_idx, field_ty) = transparent_scalar_field_with_index(ctx, current)?; - let layout = { - let ty_ref = current.deref(ctx); - let s = ty_ref.downcast_ref::().ok_or_else(|| { - anyhow::anyhow!("transparent scalar ABI requires a MirStructType") - })?; - StructLayoutInfo::of_struct(s) - }; - let map = build_struct_slot_map(ctx, &layout)?; - let field_slot = map.decl_to_llvm[decl_idx].ok_or_else(|| { - anyhow::anyhow!( - "transparent scalar field {} unexpectedly lowered as a ZST", - decl_idx - ) - })?; - layers.push(TransparentScalarLayer { - llvm_struct_ty: map.llvm_struct_ty, - field_slot, - }); - - let nested_transparent = { - let field_ref = field_ty.deref(ctx); - field_ref - .downcast_ref::() - .is_some_and(MirStructType::is_transparent_scalar) - }; - if nested_transparent { - current = field_ty; - continue; - } - - let scalar_ty = convert_type(ctx, field_ty)?; - if is_zero_sized_type(ctx, scalar_ty) { - return Err(anyhow::anyhow!( - "transparent scalar ABI resolved to a zero-sized field" - )); - } - return Ok(TransparentScalarAbiInfo { scalar_ty, layers }); - } -} - -/// LLVM ABI type for a rustc-proven transparent scalar struct. -/// -/// Transparent wrappers can nest (`Outer(Inner(u32))`). rustc still reports -/// the outer ADT as one scalar, so recurse through transparent scalar fields -/// until reaching the actual scalar/pointer representation. -pub(crate) fn transparent_scalar_llvm_type( - ctx: &mut Context, - struct_ty: TypeHandle, -) -> Result { - Ok(transparent_scalar_abi_info(ctx, struct_ty)?.scalar_ty) -} - -/// Target-stable ABI projection for packed aggregates containing AS3 leaves -/// across the internal device return boundary. -/// -/// The function body continues to use the semantic packed LLVM struct with -/// shared pointers. Only the physical return value recursively replaces AS3 -/// leaves with generic pointers, whose width is stable across the legacy/PTX -/// and modern NVVM data layouts. -#[derive(Clone, Copy, Debug)] -pub(crate) struct PackedSharedInternalAbiInfo { - pub semantic_ty: TypeHandle, - pub storage_ty: TypeHandle, -} - -/// Maximum number of shared-pointer leaves introduced by fixed-array expansion -/// in one packed-AS3 internal return carrier. -/// -/// Struct/tuple nesting and direct pointer leaves remain proportional to source -/// structure. Arrays can encode an arbitrarily large number of per-element -/// extract/cast/insert sequences compactly, so only array-expanded AS3 leaves -/// count against this code-shape budget. -pub(crate) const MAX_PACKED_SHARED_INTERNAL_ABI_ARRAY_REWRITE_LEAVES: u64 = 16; - -/// Whether a MIR type converts to a zero-sized LLVM type. -/// -/// MIR-side mirror of [`is_zero_sized_type`]: zero-length arrays, arrays of -/// zero-sized elements, and structs/tuples whose fields are all zero-sized -/// (including empty ones such as `PhantomData`) vanish at the LLVM level and -/// must not affect ABI-lane classification. -fn mir_type_is_zero_sized(ctx: &Context, ty: TypeHandle) -> bool { - let ty_ref = ty.deref(ctx); - if let Some(array_ty) = ty_ref.downcast_ref::() { - return array_ty.size() == 0 || mir_type_is_zero_sized(ctx, array_ty.element_type()); - } - if let Some(struct_ty) = ty_ref.downcast_ref::() { - return struct_ty - .field_types - .iter() - .all(|field| mir_type_is_zero_sized(ctx, *field)); - } - if let Some(tuple_ty) = ty_ref.downcast_ref::() { - return tuple_ty - .get_types() - .iter() - .all(|field| mir_type_is_zero_sized(ctx, *field)); - } - false -} - -/// Whether a MIR value shape belongs to the recursive packed-AS3 ABI lane. -/// -/// Structs, tuples, and fixed arrays may nest recursively and may contain any -/// number of scalar pointer leaves. Zero-sized fields are skipped, matching the -/// post-conversion field scan this predicate replaced. The array-specific -/// expansion budget is enforced after conversion from the storage rewrite's -/// exact AS3 leaf count. Vectors and unrelated aggregate kinds remain -/// deliberately fail-closed. -fn packed_shared_internal_abi_mir_shape_is_supported(ctx: &Context, mir_ty: TypeHandle) -> bool { - let children = { - let ty_ref = mir_ty.deref(ctx); - if ty_ref.is::() - || ty_ref.is::() - || ty_ref.is::() - || ty_ref.is::() - || ty_ref.is::() - || ty_ref.is::() - || ty_ref.is::() - { - return true; - } - if let Some(struct_ty) = ty_ref.downcast_ref::() { - Some(struct_ty.field_types.clone()) - } else if let Some(tuple_ty) = ty_ref.downcast_ref::() { - Some(tuple_ty.get_types().to_vec()) - } else { - ty_ref - .downcast_ref::() - .map(|array_ty| vec![array_ty.element_type()]) - } - }; - - children.is_some_and(|children| { - children.into_iter().all(|child| { - mir_type_is_zero_sized(ctx, child) - || packed_shared_internal_abi_mir_shape_is_supported(ctx, child) - }) - }) -} - -/// Recognize a recursive packed-AS3 internal ABI shape. -/// -/// The root must remain a byte-faithful packed struct. Nested structs/tuples, -/// multiple AS3 leaves, and bounded fixed arrays are admitted. Vectors and -/// unrelated aggregate kinds remain out of scope. The target-stable storage -/// utility owns the recursive AS3 -> generic rewrite and this classifier only -/// decides which semantic shapes may use it. -pub(crate) fn packed_shared_internal_abi_info( - ctx: &mut Context, - mir_ty: TypeHandle, -) -> Result, anyhow::Error> { - let layout = { - let ty_ref = mir_ty.deref(ctx); - let Some(struct_ty) = ty_ref.downcast_ref::() else { - return Ok(None); - }; - StructLayoutInfo::of_struct(struct_ty) - }; - - if !packed_shared_internal_abi_mir_shape_is_supported(ctx, mir_ty) { - return Ok(None); - } - - let map = build_struct_slot_map(ctx, &layout)?; - if !map.by_value_layout_faithful { - return Ok(None); - } - let is_packed = map - .llvm_struct_ty - .deref(ctx) - .downcast_ref::() - .is_some_and(|struct_ty| struct_ty.layout() == llvm_types::StructLayout::Packed); - if !is_packed { - return Ok(None); - } - - let rewrite = target_stable_storage_type( - ctx, - map.llvm_struct_ty, - StorageRewriteOptions { - canonicalize_bool: false, - }, - "packed shared internal ABI", - )?; - if rewrite.shared_pointer_leaves == 0 - || rewrite.array_shared_pointer_leaves > MAX_PACKED_SHARED_INTERNAL_ABI_ARRAY_REWRITE_LEAVES - { - return Ok(None); - } - let Some((storage_size, _)) = llvm_type_size_align(ctx, rewrite.ty) else { - return Ok(None); - }; - if layout.total_size > 0 && storage_size != layout.total_size { - return Ok(None); - } - - Ok(Some(PackedSharedInternalAbiInfo { - semantic_ty: map.llvm_struct_ty, - storage_ty: rewrite.ty, - })) -} - -/// Convert a type that crosses a function boundary as one LLVM value. -/// -/// A struct with a natural-layout divergence is legal by value only when -/// [`build_struct_slot_map`] proved that a sequential packed LLVM struct -/// reproduces rustc's offsets and size. This keeps overlapping/union-like -/// legacy struct models fail-closed while allowing real `repr(packed)` values. -/// Packed values containing shared-memory pointers remain target-dependent -/// because AS3 pointer width differs between modern NVVM and PTX/legacy modes. -fn convert_by_value_abi_type( - ctx: &mut Context, - mir_ty: TypeHandle, - role: &str, -) -> Result { - let layout = { - let ty_ref = mir_ty.deref(ctx); - ty_ref - .downcast_ref::() - .map(StructLayoutInfo::of_struct) - }; - let llvm_ty = if let Some(layout) = layout { - let map = build_struct_slot_map(ctx, &layout)?; - if !map.by_value_layout_faithful { - return Err(anyhow::anyhow!( - "{} has a rustc struct layout that cannot be represented by an LLVM struct value", - role - )); - } - map.llvm_struct_ty - } else { - convert_type(ctx, mir_ty)? - }; - - if llvm_packed_struct_contains_pointer_in_address_space( - ctx, - llvm_ty, - llvm_types::address_space::SHARED, - ) { - return Err(anyhow::anyhow!( - "{} contains a packed aggregate with a target-dependent shared-memory pointer", - role - )); - } - - Ok(llvm_ty) -} - -/// Convert a MIR function type to an LLVM function type. -/// -/// This handles the ABI-level transformations required for GPU kernels. -/// The transformations ensure that the generated LLVM IR matches the -/// C ABI expected by the CUDA runtime. -/// -/// # ABI Transformations -/// -/// ## Argument Flattening -/// -/// Aggregate types are flattened to primitive types: -/// -/// ```text -/// MIR: fn kernel(slice: &[f32], point: Point) -/// LLVM: fn internal_fn(ptr: !ptr, len: i64, x: f32, y: f32) -/// ``` -/// -/// | MIR Argument | Internal call ABI | Kernel-entry ABI | -/// |-------------------------|-------------------------|------------------------| -/// | `&[T]` | `(ptr, i64)` | `(ptr, i64)` | -/// | `DisjointSlice` | `(ptr, i64)` | `(ptr, i64)` | -/// | `struct { a: A, b: B }` | `(a: A', b: B')` | one byval `{A', B'}` | -/// | `repr(transparent)` scalar ADT | underlying scalar | underlying scalar | -/// | closure with N captures | N separate field args | one byval struct | -/// | Other | Converted type | Converted type | -/// -/// Slices keep their `(ptr, len)` flattening on both sides because the -/// host-side launch helpers push the pointer and length as two driver -/// args. Structs and closures are unflattened only at kernel boundaries -/// because the host pushes them as a single scalar — see -/// `cuda_host::push_kernel_scalar`. The exception is a rustc-proven -/// `#[repr(transparent)]` `ValueAbi::Scalar` struct: its single non-ZST -/// field is the kernel parameter, matching the source type's transparent ABI. -/// Internal device-side call sites stay flattened: caller and callee are both -/// inside this backend, so the ABI is private and there is no host to disagree -/// with. -/// -/// ## Return Type Handling -/// -/// - Empty tuple `()` becomes `void` -/// - Empty struct `struct {}` becomes `void` -/// - Rustc-proven `repr(transparent)` scalar ADTs return the underlying scalar -/// - Other types are converted normally -/// -/// # Arguments -/// -/// * `ctx` - The pliron context -/// * `func_type` - The MIR function type to convert -/// * `is_kernel_entry` - When `true`, treat aggregate (non-slice) params -/// as single byval values to match the host-side push ABI. When `false`, -/// keep the existing internal device-fn ABI that flattens struct fields -/// into individual scalars. -/// -/// # Returns -/// -/// The equivalent LLVM function type with ABI transformations applied. -/// -/// # Example -/// -/// ```text -/// MIR: fn foo(a: &[f32], b: i32) -> f32 -/// LLVM: fn foo(ptr, i64, i32) -> f32 -/// -/// MIR: fn bar() -> () -/// LLVM: fn bar() -> void -/// ``` -/// -/// # Note -/// -/// At internal device-function boundaries the struct flattening must be -/// reversed in the entry block. At kernel-entry boundaries ordinary structs -/// arrive as one byval aggregate, while rustc-proven transparent scalar -/// wrappers arrive as their single non-ZST scalar field and must be rebuilt. -/// See `lowering.rs::build_entry_prologue` for the reconstruction paths. -pub fn convert_function_type( - ctx: &mut Context, - func_type: pliron::r#type::TypedHandle, - is_kernel_entry: bool, -) -> Result, anyhow::Error> { - // Extract input/output types before mutating context - let (inputs_ptr, results_ptr) = { - let func_ty_ref = func_type.deref(ctx); - let interface = type_cast::(&*func_ty_ref) - .ok_or_else(|| anyhow::anyhow!("Type does not implement FunctionTypeInterface"))?; - (interface.arg_types(), interface.res_types()) - }; - - // Convert inputs, flattening slice/struct types for ABI compatibility. - // Slices flatten on both ABIs; structs flatten only on the internal - // device-fn ABI. - let mut inputs = Vec::new(); - let inputs_vec: Vec<_> = inputs_ptr.to_vec(); - - for t in inputs_vec { - // Determine what kind of flattening this type needs - // Extract all info first, then drop the borrow - enum FlattenKind { - /// `{ ptr, len }`, then one parameter per index-space layout field. - Slice { - space_tys: Vec, - }, - Struct { - field_types: Vec, - mem_to_decl: Vec, - }, - TransparentScalar(TypeHandle), - None, - } - - let flatten_kind = { - let ty_ref = t.deref(ctx); - if ty_ref.is::() { - FlattenKind::Slice { - space_tys: Vec::new(), - } - } else if let Some(slice_ty) = ty_ref.downcast_ref::() { - FlattenKind::Slice { - space_tys: slice_ty.space_tys.clone(), - } - } else if let Some(struct_ty) = ty_ref.downcast_ref::() { - if is_kernel_entry && struct_ty.is_transparent_scalar() { - // rustc proves that this `repr(transparent)` ADT has a - // scalar ABI. Emit exactly its one non-ZST field as the - // kernel parameter instead of an aggregate `.b8[]` param. - FlattenKind::TransparentScalar(t) - } else if is_kernel_entry { - // Ordinary kernel-boundary structs remain intact so the - // host's single `push_kernel_scalar(&closure)` push - // matches a single aggregate .param entry. - FlattenKind::None - } else { - FlattenKind::Struct { - field_types: struct_ty.field_types.clone(), - mem_to_decl: struct_ty.memory_order(), - } - } - } else { - FlattenKind::None - } - }; - - match flatten_kind { - FlattenKind::Slice { space_tys } => { - let ptr_ty = llvm_types::PointerType::get_generic(ctx); - let len_ty = IntegerType::get(ctx, 64, Signedness::Signless); - inputs.push(ptr_ty.into()); - inputs.push(len_ty.into()); - for space_ty in space_tys { - let converted = convert_type(ctx, space_ty)?; - // A zero-sized index-space field contributes no parameter, - // as NVPTX has no empty `.param`. - if !is_zero_sized_type(ctx, converted) { - inputs.push(converted); - } - } - } - FlattenKind::Struct { - field_types, - mem_to_decl, - } => { - // Flatten in MEMORY ORDER to match struct layout - for mem_idx in 0..field_types.len() { - let decl_idx = mem_to_decl[mem_idx]; - let converted = convert_by_value_abi_type( - ctx, - field_types[decl_idx], - "by-value function argument", - )?; - // Skip ZST fields - NVPTX can't handle empty params - if !is_zero_sized_type(ctx, converted) { - inputs.push(converted); - } - } - } - FlattenKind::TransparentScalar(struct_ty) => { - inputs.push(transparent_scalar_llvm_type(ctx, struct_ty)?); - } - FlattenKind::None => { - let converted = convert_by_value_abi_type(ctx, t, "by-value function argument")?; - // Skip ZST args - NVPTX can't handle empty params - if !is_zero_sized_type(ctx, converted) { - inputs.push(converted); - } - } - } - } - - // Convert return type. A rustc-proven transparent scalar wrapper uses - // the underlying scalar at the function ABI boundary, while its body - // continues to use the ordinary converted aggregate representation. - let ret_ty = if results_ptr.is_empty() { - llvm_types::VoidType::get(ctx).into() - } else { - let mir_ret_ty = results_ptr[0]; - let is_transparent_scalar = { - let ty_ref = mir_ret_ty.deref(ctx); - ty_ref - .downcast_ref::() - .is_some_and(MirStructType::is_transparent_scalar) - }; - let ty = if is_transparent_scalar { - transparent_scalar_llvm_type(ctx, mir_ret_ty)? - } else if !is_kernel_entry { - match packed_shared_internal_abi_info(ctx, mir_ret_ty)? { - Some(abi) => abi.storage_ty, - None => convert_by_value_abi_type(ctx, mir_ret_ty, "by-value function return")?, - } - } else { - convert_by_value_abi_type(ctx, mir_ret_ty, "by-value function return")? - }; - // Check if zero-sized (empty struct or struct with only ZST fields). - if is_zero_sized_type(ctx, ty) { - llvm_types::VoidType::get(ctx).into() - } else { - ty - } - }; - - Ok(llvm_types::FuncType::get(ctx, ret_ty, inputs, false)) -} - -// ============================================================================= -// Struct Slot Mapping (single source of truth, issue #128) -// ============================================================================= - -/// Declaration-order layout facts for one MIR aggregate, in the exact form -/// [`build_struct_slot_map`] consumes. -/// -/// Extracting this owned carrier first (and dropping the `Ref` returned by -/// `Ptr::deref`) keeps the borrow checker happy: the slot-map build needs -/// `&mut Context` for type interning. -pub(crate) struct StructLayoutInfo { - /// Field types in declaration order. - pub field_types: Vec, - /// Memory order: `mem_to_decl[mem_idx] = decl_idx`. Always full length - /// (identity when rustc did not reorder). - pub mem_to_decl: Vec, - /// Byte offset of each field in declaration order; empty when rustc - /// layout is unknown. - pub field_offsets: Vec, - /// Total size in bytes including trailing padding; 0 when unknown. - pub total_size: u64, -} - -impl StructLayoutInfo { - /// Layout facts of a `MirStructType`. - pub(crate) fn of_struct(s: &MirStructType) -> Self { - StructLayoutInfo { - field_types: s.field_types.clone(), - mem_to_decl: s.memory_order(), - field_offsets: s.field_offsets().to_vec(), - total_size: s.total_size(), - } - } - - /// Layout facts of a `MirTupleType`. - /// - /// Tuples translated from a rustc type carry rustc's exact layout - /// (offsets, memory order, size), which is consumed here identically to - /// structs, so reordered tuples like `(u32, &T)` lower byte-correctly. - /// Only synthetic layout-less tuples (the unit tuple, hand-built test - /// types) fall back to LLVM natural layout. - pub(crate) fn of_tuple(t: &MirTupleType) -> Self { - StructLayoutInfo { - field_types: t.get_types().to_vec(), - mem_to_decl: t.memory_order(), - field_offsets: t.field_offsets().to_vec(), - total_size: t.total_size(), - } - } -} - -/// One lowered LLVM struct plus the value-level slot mapping into it. -/// -/// [`build_struct_slot_map`] produces the struct type and the index map in -/// the same walk, so every op that indexes into the struct (`insertvalue`, -/// `extractvalue`, GEP, call-boundary flatten/reconstruct) shares the type -/// converter's view of where each field landed. Computing the indices -/// separately is how the issue #128 class of bug (indices that ignore the -/// `[N x i8]` padding slots) happened. -pub(crate) struct StructSlotMap { - /// The final LLVM struct type, including any `[N x i8]` padding slots. - pub llvm_struct_ty: TypeHandle, - /// `decl_to_llvm[decl_idx]` = LLVM slot of that declaration-order field; - /// `None` when the field is zero-sized and was stripped. - pub decl_to_llvm: Vec>, - /// Converted LLVM type of each declaration-order field (ZSTs included). - pub field_llvm_types: Vec, - /// Natural (non-packed) LLVM byte offset of every slot of - /// `llvm_struct_ty`, padding slots included; `None` when some field type - /// cannot be sized ([`llvm_type_size_align`] declined). - pub natural_slot_offsets: Option>, - /// True when rustc's recorded byte layout cannot be represented by a - /// naturally laid-out LLVM struct: some field's natural slot offset differs - /// from rustc's byte offset, or the natural struct size differs from - /// rustc's total size. `repr(packed)` is the canonical case. Address-path - /// consumers retain this signal and the natural offsets so the #859 - /// byte-GEP fallback stays stable for packed field projections. - pub layout_diverges: bool, - /// True when the selected LLVM struct representation reproduces rustc's - /// recorded field offsets and total size for by-value movement. Natural - /// layouts are faithful when they do not diverge; divergent layouts are - /// faithful only when a sequential LLVM packed struct can express them. - /// Overlapping/union-like legacy struct models therefore remain false. - pub by_value_layout_faithful: bool, -} - -/// Whether a `MirStructType` value lowers to a byte-faithful LLVM struct: -/// either its natural LLVM layout already matches rustc's offsets and total -/// size, or the packed-with-explicit-padding representation reproduces them -/// exactly. -/// -/// Public only as a coupling oracle for mir-importer: its constant-promotion -/// gate must never admit a struct layout whose converted storage falls back -/// to a divergent natural layout, because a promoted constant's byte image -/// would then disagree with every typed read through the converted struct. -/// mir-importer asserts agreement against this function in its tests. -pub fn struct_value_lowering_is_byte_faithful( - ctx: &mut Context, - struct_ty: TypeHandle, -) -> Result { - let layout = { - let ty_ref = struct_ty.deref(ctx); - let mir_struct = ty_ref - .downcast_ref::() - .ok_or_else(|| anyhow::anyhow!("expected a MirStructType"))?; - StructLayoutInfo::of_struct(mir_struct) - }; - Ok(build_struct_slot_map(ctx, &layout)?.by_value_layout_faithful) -} - -/// Lower a struct/tuple layout to its LLVM struct type and slot map. -/// -/// When rustc layout is present (`field_offsets` non-empty and -/// `total_size > 0`), fields are placed at their exact byte offsets with -/// explicit `[N x i8]` padding slots in between, plus a trailing pad up to -/// `total_size`. This makes the layout independent of LLVM's datalayout -/// and so ABI-identical to what rustc computed on the host. For -/// `struct Extreme { a: u8, b: i128 }` where rustc puts `b` at offset 0 -/// and `a` at offset 16 with total size 32, we build: -/// -/// ```text -/// { i128, i8, [15 x i8] } ; slots: b = 0, a = 1, pad = 2 -/// ``` -/// -/// Without rustc layout, fields are emitted in memory order with no -/// padding. On both paths zero-sized fields (e.g. `PhantomData`) are -/// stripped, because NVPTX rejects empty types; stripped fields get -/// `None` in `decl_to_llvm`. -/// -/// Malformed layout metadata (a `mem_to_decl` that is not a permutation, -/// or an offsets vector of the wrong length) is rejected loudly: guessing -/// here would scramble every downstream field access. -pub(crate) fn build_struct_slot_map( - ctx: &mut Context, - layout: &StructLayoutInfo, -) -> Result { - let num_fields = layout.field_types.len(); - - if layout.mem_to_decl.len() != num_fields { - return Err(anyhow::anyhow!( - "struct slot map: memory order has {} entries but the struct has {} fields", - layout.mem_to_decl.len(), - num_fields - )); - } - let mut seen = vec![false; num_fields]; - for &decl_idx in &layout.mem_to_decl { - if decl_idx >= num_fields || seen[decl_idx] { - return Err(anyhow::anyhow!( - "struct slot map: memory order {:?} is not a permutation of 0..{}", - layout.mem_to_decl, - num_fields - )); - } - seen[decl_idx] = true; - } - let has_explicit_layout = !layout.field_offsets.is_empty() && layout.total_size > 0; - if has_explicit_layout && layout.field_offsets.len() != num_fields { - return Err(anyhow::anyhow!( - "struct slot map: {} field offsets for {} fields", - layout.field_offsets.len(), - num_fields - )); - } - - // Convert every field up front, in declaration order. - let mut field_llvm_types = Vec::with_capacity(num_fields); - for &field_ty in &layout.field_types { - field_llvm_types.push(convert_type(ctx, field_ty)?); - } - - let mut llvm_fields: Vec = Vec::new(); - let mut decl_to_llvm: Vec> = vec![None; num_fields]; - let mut current_offset: u64 = 0; - - // Place fields in memory order. - for &decl_idx in &layout.mem_to_decl { - let llvm_ty = field_llvm_types[decl_idx]; - - // ZST fields are stripped: no slot, no offset advance (rustc gives - // them size 0). - if is_zero_sized_type(ctx, llvm_ty) { - continue; - } - - if has_explicit_layout { - // Insert padding if needed to reach the rustc field offset. - let target_offset = layout.field_offsets[decl_idx]; - if current_offset < target_offset { - let padding_ty = make_padding_type(ctx, target_offset - current_offset); - llvm_fields.push(padding_ty); - current_offset = target_offset; - } - } - - decl_to_llvm[decl_idx] = Some(llvm_fields.len() as u32); - llvm_fields.push(llvm_ty); - - if has_explicit_layout { - // Offset advance is exact or an error, never guessed. Prefer - // rustc's stored size for the field: nested aggregates carry - // interior/trailing padding the converted type cannot always - // reproduce, and a wrong advance here either forces interior - // padding where rustc has none or overshoots the next field's - // offset. Scalars and other rustc-silent types fall back to the - // LLVM natural size. - let advance = match mir_stored_size(ctx, layout.field_types[decl_idx]) { - Some(size) => size, - None => llvm_type_size_align(ctx, llvm_ty) - .map(|(size, _)| size) - .ok_or_else(|| { - anyhow::anyhow!( - "struct slot map: field {decl_idx} lowers to `{}`, which has no \ - exact size; refusing to guess the next field offset", - llvm_ty.deref(ctx).disp(ctx) - ) - })?, - }; - current_offset += advance; - } - } - - // Add trailing padding to reach total_size. - if has_explicit_layout && current_offset < layout.total_size { - let padding_ty = make_padding_type(ctx, layout.total_size - current_offset); - llvm_fields.push(padding_ty); - } - - // The explicit `[N x i8]` gaps above can only ADD bytes; they cannot make - // LLVM place a slot earlier than its natural alignment allows. So when - // rustc's layout is tighter than natural (repr(packed)), the struct built - // here is a lie at the byte level, and every consumer needs to know. - // Unsizable field types leave no verdict: consumers that need the - // comparison (field_addr with a recorded rustc offset) fail closed at - // their own sites. - let walk = natural_layout_walk(ctx, &llvm_fields); - let natural_slot_offsets = walk.as_ref().map(|(offsets, _, _)| offsets.clone()); - let mut layout_diverges = false; - if has_explicit_layout && let Some((offsets, end, align)) = &walk { - for (decl_idx, slot) in decl_to_llvm.iter().enumerate() { - if let Some(slot) = slot - && offsets[*slot as usize] != layout.field_offsets[decl_idx] - { - layout_diverges = true; - } - } - let natural_size = end.div_ceil(*align) * *align; - if natural_size != layout.total_size { - layout_diverges = true; - } - } - - // A packed LLVM struct is sequential with alignment 1. That faithfully - // represents repr(packed) only when the sequential packed offsets and - // final byte count exactly match rustc's metadata. A divergent layout can - // also be an old union-like/overlapping model; selecting Packed for such a - // shape would merely turn one incorrect sequential layout into another. - let packed_walk = if layout_diverges && has_explicit_layout { - let mut offsets = Vec::with_capacity(llvm_fields.len()); - let mut end = 0u64; - let mut sizeable = true; - for &field in &llvm_fields { - offsets.push(end); - let Some((field_size, _)) = llvm_type_size_align(ctx, field) else { - sizeable = false; - break; - }; - let Some(next) = end.checked_add(field_size) else { - sizeable = false; - break; - }; - end = next; - } - sizeable.then_some((offsets, end)) - } else { - None - }; - let packed_representable = packed_walk.as_ref().is_some_and(|(offsets, end)| { - *end == layout.total_size - && decl_to_llvm - .iter() - .enumerate() - .all(|(decl_idx, slot)| match slot { - Some(slot) => offsets[*slot as usize] == layout.field_offsets[decl_idx], - None => true, - }) - }); - - let struct_layout = if packed_representable { - llvm_types::StructLayout::Packed - } else { - llvm_types::StructLayout::Unpacked - }; - let llvm_struct_ty: TypeHandle = - llvm_types::StructType::get_unnamed(ctx, (llvm_fields, struct_layout)).into(); - let by_value_layout_faithful = !layout_diverges || packed_representable; - - Ok(StructSlotMap { - llvm_struct_ty, - decl_to_llvm, - field_llvm_types, - natural_slot_offsets, - layout_diverges, - by_value_layout_faithful, - }) -} - -/// Create a padding type: `[N x i8]` for N bytes of padding. -fn make_padding_type(ctx: &mut Context, size: u64) -> TypeHandle { - let i8_ty = IntegerType::get(ctx, 8, Signedness::Signless); - llvm_types::ArrayType::get(ctx, i8_ty.into(), size).into() -} - -/// Storage for `size` bytes of enum filler at byte `offset`: bytes the payload -/// occupies that no typed slot claims. -/// -/// `[N x i8]` is byte-exact but costs one leaf *per byte* everywhere a payload -/// value is built, merged, or taken apart. `Option<&[T]>` is the shape that -/// shows it: the niche carrier claims the pointer, the length is left to an -/// 8-byte filler, and building the value then decomposes that length into -/// eight `i8` `insertvalue`s, every block merge carries eight `phi i8`, and -/// reading it back is a chain of `prmt` byte permutes with the aggregate -/// spilled to `.local` in between. Measured on sm_86, it costs the *pointer* -/// too: it rides through the same byte soup, `InferAddressSpaces` can no -/// longer follow it, and the access lands in the generic window (`ld.v2.b64`) -/// instead of `ld.global.v2.b64`. -/// -/// One integer of the same width is the same bytes in one leaf. It is only -/// legal where it moves nothing: -/// -/// - the width is 2, 4 or 8 bytes, so the integer is byte-faithful (a multiple -/// of 8 bits, hence no padding of its own). 16 is deliberately excluded: -/// `i128` is legal LLVM but lowers to register pairs on NVPTX, so widening -/// that far trades one win for another cost and wants its own measurement; -/// - `offset` is a multiple of the width, so LLVM inserts no gap ahead of the -/// field and every later field keeps its byte offset; -/// - the width does not exceed the enum's alignment, or the struct's natural -/// alignment would rise above rustc's. -/// -/// Anything else keeps the byte array. The filler is a *vehicle*, not a claim: -/// nothing reads it field-wise, because a payload that has no typed slot -/// round-trips through memory as a whole aggregate, and both ends of that -/// round trip are byte-exact. `build_enum_slot_map`'s own size and alignment -/// assertions — hard errors, not debug checks — backstop all three conditions. -fn make_enum_filler_type(ctx: &mut Context, offset: u64, size: u64, abi_align: u64) -> TypeHandle { - if matches!(size, 2 | 4 | 8) && offset.is_multiple_of(size) && size <= abi_align.max(1) { - let width = u32::try_from(size * 8).expect("size is at most 8, so width is at most 64"); - return IntegerType::get(ctx, width, Signedness::Signless).into(); - } - make_padding_type(ctx, size) -} - -/// Build byte-exact LLVM storage for a Rust union. -/// -/// A union cannot be represented as an LLVM struct containing every declared -/// field: struct fields are consecutive, while union fields all start at byte -/// zero. We choose one byte-faithful field as the storage view and add explicit -/// tail bytes. A zero-length integer array raises the LLVM type's natural -/// alignment without consuming storage. Pointer-bearing fields are preferred -/// so an ordinary union copy keeps LLVM pointer provenance. -/// -/// NVPTX gives scalar integers natural alignments up to 16 bytes. Stronger -/// Rust alignment is carried explicitly on memory operations because LLVM -/// aggregate types cannot encode over-alignment; the storage type still keeps -/// the union's exact size and therefore its array stride. -pub(crate) fn build_union_storage_type( - ctx: &mut Context, - union_ty: &MirUnionType, -) -> Result { - let size = union_ty.total_size(); - let align = union_ty.abi_align(); - if align == 0 || !align.is_power_of_two() { - return Err(anyhow::anyhow!( - "union `{}` has invalid ABI alignment {}", - union_ty.name(), - align - )); - } - if size > 0 && !size.is_multiple_of(align) { - return Err(anyhow::anyhow!( - "union `{}` size {} is not a multiple of its {}-byte ABI alignment", - union_ty.name(), - size, - align - )); - } - - let mut fields = Vec::with_capacity(union_ty.field_count()); - let mut pointer_carrier: Option = None; - for (index, &field_ty) in union_ty.field_types().iter().enumerate() { - let llvm_field_ty = convert_type(ctx, field_ty)?; - let (field_size, field_align) = - llvm_type_size_align(ctx, llvm_field_ty).ok_or_else(|| { - anyhow::anyhow!( - "union `{}` field {} has unsupported LLVM size/alignment", - union_ty.name(), - index - ) - })?; - if field_size > size { - return Err(anyhow::anyhow!( - "union `{}` field {} lowers to {} bytes but the union is only {} bytes", - union_ty.name(), - index, - field_size, - size - )); - } - if field_align > align { - return Err(anyhow::anyhow!( - "union `{}` field {} lowers with alignment {} but rustc reports union alignment {}", - union_ty.name(), - index, - field_align, - align - )); - } - let contains_pointer = llvm_type_contains_pointer(ctx, llvm_field_ty); - if contains_pointer { - if let Some(first) = pointer_carrier - && first != llvm_field_ty - { - return Err(anyhow::anyhow!( - "union `{}` has pointer-bearing fields with different LLVM representations; preserving provenance for that shape is not yet supported", - union_ty.name() - )); - } - pointer_carrier = Some(llvm_field_ty); - } - fields.push((llvm_field_ty, field_size, field_align, contains_pointer)); - } - - let storage_align = align.min(16); - let anchor_int = IntegerType::get(ctx, (storage_align * 8) as u32, Signedness::Signless); - let anchor: TypeHandle = llvm_types::ArrayType::get(ctx, anchor_int.into(), 0).into(); - let mut storage_fields = vec![anchor]; - if size > 0 { - let representative = fields - .iter() - .filter(|(ty, field_size, _, _)| { - *field_size > 0 - && llvm_type_is_byte_faithful(ctx, *ty) - && (pointer_carrier.is_none() || llvm_type_contains_pointer(ctx, *ty)) - }) - .max_by_key(|(_, field_size, field_align, contains_pointer)| { - (*contains_pointer, *field_align, *field_size) - }); - if let Some(representative) = representative { - storage_fields.push(representative.0); - if representative.1 < size { - storage_fields.push(make_padding_type(ctx, size - representative.1)); - } - } else if pointer_carrier.is_some() { - return Err(anyhow::anyhow!( - "union `{}` has pointer-bearing fields but no byte-faithful pointer carrier; lowering it as raw bytes would discard pointer provenance", - union_ty.name() - )); - } else { - // Pointer-free unions may safely use raw bytes as their SSA - // carrier. Field loads/stores still use their declared types. - storage_fields.push(make_padding_type(ctx, size)); - } - } - let storage: TypeHandle = llvm_types::StructType::get_unnamed( - ctx, - (storage_fields, llvm_types::StructLayout::Unpacked), - ) - .into(); - let (llvm_size, llvm_align) = llvm_type_size_align(ctx, storage).ok_or_else(|| { - anyhow::anyhow!( - "union `{}` storage has unsupported LLVM layout", - union_ty.name() - ) - })?; - if llvm_size != size || llvm_align > align { - return Err(anyhow::anyhow!( - "union `{}` storage lowered to incompatible size/alignment {}/{} but rustc requires {}/{}", - union_ty.name(), - llvm_size, - llvm_align, - size, - align - )); - } - Ok(storage) -} - -fn llvm_type_contains_pointer(ctx: &Context, ty: TypeHandle) -> bool { - let ty_ref = ty.deref(ctx); - if ty_ref.is::() { - return true; - } - if let Some(array) = ty_ref.downcast_ref::() { - return llvm_type_contains_pointer(ctx, array.elem_type()); - } - if let Some(vector) = ty_ref.downcast_ref::() { - return llvm_type_contains_pointer(ctx, vector.elem_type()); - } - if let Some(struct_ty) = ty_ref.downcast_ref::() { - return struct_ty - .fields() - .any(|field| llvm_type_contains_pointer(ctx, field)); - } - false -} - -/// One pointer-valued leaf in an LLVM aggregate's physical storage. -/// -/// Offsets are absolute within the enclosing enum. Keeping the address space -/// in the identity prevents an AS0 pointer carrier from being treated as the -/// same storage as an AS1 payload pointer merely because both are eight bytes. -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)] -struct LlvmPointerStorage { - offset: u64, - size: u64, - address_space: u32, -} - -/// Why a pointer-storage walk refused a type or an overlap. -/// -/// Distinguishing bounded array expansion from genuine provenance loss keeps -/// the user-facing diagnostic specific: an oversized pointer array reports the -/// same "rewrite requires N pointer conversions" contract as the payload -/// storage gate instead of a misleading provenance error. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum PointerOverlapRejection { - /// Expanding the walked type's fixed arrays into per-leaf records would - /// exceed [`MAX_ENUM_PAYLOAD_ARRAY_REWRITE_LEAVES`]; `required` is the - /// total number of array-expanded pointer leaves. - OverArrayLeafBound { required: u64 }, - /// The type holds pointer storage this walk cannot represent, or the - /// overlap would pun a pointer against non-matching bytes. - ProvenanceLoss, -} - -/// Total pointer leaves in `ty`, through arrays and structs, with checked -/// arithmetic. Pointer vectors count zero; every walk that expands leaves -/// fails closed on them separately. -fn count_pointer_leaves(ctx: &Context, ty: TypeHandle) -> Option { - let ty_ref = ty.deref(ctx); - if ty_ref.is::() { - return Some(1); - } - if let Some(array) = ty_ref.downcast_ref::() { - return count_pointer_leaves(ctx, array.elem_type())?.checked_mul(array.size()); - } - if let Some(struct_ty) = ty_ref.downcast_ref::() { - let fields: Vec<_> = struct_ty.fields().collect(); - let mut total = 0u64; - for field in fields { - total = total.checked_add(count_pointer_leaves(ctx, field)?)?; - } - return Some(total); - } - Some(0) -} - -/// Pointer leaves that expanding `ty`'s fixed arrays contributes to a -/// pointer-storage walk. -/// -/// Leaves outside arrays are proportional to the source text and stay -/// unbounded, exactly like struct nesting in `enum_payload_storage_type`. -/// `[&T; N]` expands from three tokens into `N` records, so only -/// array-expanded leaves count against -/// [`MAX_ENUM_PAYLOAD_ARRAY_REWRITE_LEAVES`]. -fn count_array_pointer_leaves(ctx: &Context, ty: TypeHandle) -> Option { - let ty_ref = ty.deref(ctx); - if let Some(array) = ty_ref.downcast_ref::() { - return count_pointer_leaves(ctx, array.elem_type())?.checked_mul(array.size()); - } - if let Some(struct_ty) = ty_ref.downcast_ref::() { - let fields: Vec<_> = struct_ty.fields().collect(); - let mut total = 0u64; - for field in fields { - total = total.checked_add(count_array_pointer_leaves(ctx, field)?)?; - } - return Some(total); - } - Some(0) -} - -/// Record every pointer-valued leaf in `ty` at its natural LLVM byte offset. -/// -/// This is deliberately a physical-layout walk rather than a simple -/// `contains_pointer` predicate. A niche carrier may be one field inside an -/// aggregate payload, for example the pointer at byte 8 in -/// `Option<(usize, &T)>`. In that case the aggregate and the carrier overlap, -/// but they agree exactly about which bytes hold the pointer. -/// -/// Expanding an arbitrary array into one record per pointer would let a valid -/// but enormous type consume unbounded verifier memory, so the walk first -/// counts the array-expanded leaves and refuses anything over -/// [`MAX_ENUM_PAYLOAD_ARRAY_REWRITE_LEAVES`] with the bound-specific -/// rejection before allocating a single record. -fn collect_llvm_pointer_storage( - ctx: &Context, - ty: TypeHandle, - base_offset: u64, - out: &mut Vec, -) -> std::result::Result<(), PointerOverlapRejection> { - fn collect( - ctx: &Context, - ty: TypeHandle, - base_offset: u64, - out: &mut Vec, - ) -> Option<()> { - let ty_ref = ty.deref(ctx); - if let Some(pointer) = ty_ref.downcast_ref::() { - let (size, _) = llvm_type_size_align(ctx, ty)?; - out.push(LlvmPointerStorage { - offset: base_offset, - size, - address_space: pointer.address_space(), - }); - return Some(()); - } - if let Some(array) = ty_ref.downcast_ref::() { - let element_ty = array.elem_type(); - if !llvm_type_contains_pointer(ctx, element_ty) { - return Some(()); - } - let (element_size, _) = llvm_type_size_align(ctx, element_ty)?; - for index in 0..array.size() { - let element_offset = element_size.checked_mul(index)?; - collect( - ctx, - element_ty, - base_offset.checked_add(element_offset)?, - out, - )?; - } - return Some(()); - } - if let Some(vector) = ty_ref.downcast_ref::() { - // Pointer vectors carry vector-specific ABI alignment and cast - // semantics. Keep them fail-closed rather than treating them like - // fixed arrays. - return (!llvm_type_contains_pointer(ctx, vector.elem_type())).then_some(()); - } - if let Some(struct_ty) = ty_ref.downcast_ref::() { - let fields: Vec<_> = struct_ty.fields().collect(); - let packed = struct_ty.layout() == llvm_types::StructLayout::Packed; - let mut end = 0u64; - for field in fields { - let (field_size, field_align) = llvm_type_size_align(ctx, field)?; - let field_offset = if packed { - end - } else { - let field_align = field_align.max(1); - let remainder = end % field_align; - if remainder == 0 { - end - } else { - end.checked_add(field_align - remainder)? - } - }; - collect(ctx, field, base_offset.checked_add(field_offset)?, out)?; - end = field_offset.checked_add(field_size)?; - } - return Some(()); - } - - // All pointer-bearing LLVM types understood by this lowering are - // handled above. Unknown pointer containers must fail closed. - (!llvm_type_contains_pointer(ctx, ty)).then_some(()) - } - - let required = - count_array_pointer_leaves(ctx, ty).ok_or(PointerOverlapRejection::ProvenanceLoss)?; - if required > MAX_ENUM_PAYLOAD_ARRAY_REWRITE_LEAVES { - return Err(PointerOverlapRejection::OverArrayLeafBound { required }); - } - collect(ctx, ty, base_offset, out).ok_or(PointerOverlapRejection::ProvenanceLoss) -} - -/// Analyze how a slotless incoming pointer-bearing field overlaps the enum's -/// already-selected storage, and report how to back every one of its pointers -/// with a real `ptr` slot. -/// -/// Returns `Ok(extra)` when the field is representable without erasing any -/// pointer's provenance: every pointer leaf that coincides with an existing -/// claim reuses that claim, and each remaining ("extra") pointer leaf lands on -/// bytes no other claim covers, so it can be backed by its own fresh `ptr` -/// slot. `extra` lists exactly those leaves for the caller to add as claims. -/// An empty `extra` is the exact-match case (e.g. `Option<(usize, &T)>`, whose -/// single pointer already coincides with the niche carrier); a non-empty -/// `extra` is the multi-pointer payload case (e.g. `Option<(&mut [T], -/// &mut [T])>` from `split_at_mut_checked`, whose second slice pointer needs a -/// slot of its own beside the carrier). -/// -/// Returns [`PointerOverlapRejection::ProvenanceLoss`] when the field cannot -/// be represented without punning a pointer against non-pointer bits: an -/// existing pointer slot the incoming field does not also carry at the same -/// offset/size/address space, or an extra pointer leaf that would overlap an -/// existing (necessarily non-pointer) claim. That genuine pointer/integer -/// union stays fail-closed — there is no single LLVM slot type that is both -/// provenance-carrying and integer-exact. -/// -/// Returns [`PointerOverlapRejection::OverArrayLeafBound`] when a walked type -/// holds more array-expanded pointer leaves than the bounded rewrite limit, -/// so the caller can report the bound instead of a provenance error. -fn analyze_pointer_overlap( - ctx: &Context, - incoming_offset: u64, - incoming_size: u64, - incoming_ty: TypeHandle, - colliding_claims: &[&(u64, u64, TypeHandle)], -) -> std::result::Result, PointerOverlapRejection> { - use PointerOverlapRejection::ProvenanceLoss; - - let incoming_end = incoming_offset - .checked_add(incoming_size) - .ok_or(ProvenanceLoss)?; - - let mut incoming = Vec::new(); - collect_llvm_pointer_storage(ctx, incoming_ty, incoming_offset, &mut incoming)?; - - let mut existing = Vec::new(); - for &&(offset, _size, claim_ty) in colliding_claims { - let mut regions = Vec::new(); - collect_llvm_pointer_storage(ctx, claim_ty, offset, &mut regions)?; - existing.extend(regions.into_iter().filter(|region| { - let Some(region_end) = region.offset.checked_add(region.size) else { - return true; - }; - region.offset < incoming_end && incoming_offset < region_end - })); - } - - // Every existing pointer leaf overlapping this field must be one the field - // also carries at the same offset/size/address space; otherwise a pointer - // slot would be backed by non-matching bytes (an address-space mismatch, or - // a pointer claim where the field has integer bits). - for leaf in &existing { - if !incoming.contains(leaf) { - return Err(ProvenanceLoss); - } - } - - // Each incoming pointer leaf that does not reuse an existing claim needs its - // own slot, and may only get one if it lands on otherwise-unclaimed bytes. A - // leaf overlapping a claim it did not match is a pointer/non-pointer pun and - // stays fail-closed. - let mut extra = Vec::new(); - for leaf in &incoming { - if existing.contains(leaf) { - continue; - } - let leaf_end = leaf.offset.checked_add(leaf.size).ok_or(ProvenanceLoss)?; - let overlaps_claim = colliding_claims.iter().any(|&&(o, s, _)| { - let Some(claim_end) = o.checked_add(s) else { - return true; - }; - o < leaf_end && leaf.offset < claim_end - }); - if overlaps_claim { - return Err(ProvenanceLoss); - } - extra.push(*leaf); - } - - Ok(extra) -} - -pub(crate) fn llvm_type_contains_pointer_in_address_space( - ctx: &Context, - ty: TypeHandle, - address_space: u32, -) -> bool { - let ty_ref = ty.deref(ctx); - if let Some(pointer) = ty_ref.downcast_ref::() { - return pointer.address_space() == address_space; - } - if let Some(array) = ty_ref.downcast_ref::() { - return llvm_type_contains_pointer_in_address_space(ctx, array.elem_type(), address_space); - } - if let Some(vector) = ty_ref.downcast_ref::() { - return llvm_type_contains_pointer_in_address_space(ctx, vector.elem_type(), address_space); - } - if let Some(struct_ty) = ty_ref.downcast_ref::() { - return struct_ty - .fields() - .any(|field| llvm_type_contains_pointer_in_address_space(ctx, field, address_space)); - } - false -} - -/// Whether a physical by-value image contains a packed struct whose bytes -/// include a pointer in `address_space`. -/// -/// Direct pointers in an unpacked aggregate do not trigger this predicate: -/// their ABI can preserve address-space semantics without observing the -/// pointer's raw storage width. A pointer nested anywhere under an LLVM packed -/// struct does trigger it because whole-value construction/load/store/ABI -/// traffic observes that packed physical image. This distinction keeps normal -/// AS3 aggregate handling intact while rejecting the target-dependent packed -/// case (modern NVVM p3:32 versus 64-bit PTX/legacy). -pub(crate) fn llvm_packed_struct_contains_pointer_in_address_space( - ctx: &Context, - ty: TypeHandle, - address_space: u32, -) -> bool { - let ty_ref = ty.deref(ctx); - if let Some(array) = ty_ref.downcast_ref::() { - return llvm_packed_struct_contains_pointer_in_address_space( - ctx, - array.elem_type(), - address_space, - ); - } - if let Some(vector) = ty_ref.downcast_ref::() { - return llvm_packed_struct_contains_pointer_in_address_space( - ctx, - vector.elem_type(), - address_space, - ); - } - if let Some(struct_ty) = ty_ref.downcast_ref::() { - let fields: Vec<_> = struct_ty.fields().collect(); - if struct_ty.layout() == llvm_types::StructLayout::Packed - && fields - .iter() - .copied() - .any(|field| llvm_type_contains_pointer_in_address_space(ctx, field, address_space)) - { - return true; - } - return fields.into_iter().any(|field| { - llvm_packed_struct_contains_pointer_in_address_space(ctx, field, address_space) - }); - } - false -} - -/// Whether this is an aggregate/vector that contains LLVM `i1` storage. -/// -/// Rust `bool` is an SSA `i1`, but its memory representation is one complete -/// byte whose only valid values are 0 and 1. Enum storage never uses `i1` as -/// a physical type: scalar bools claim an explicit i8 byte, and aggregates -/// containing bools claim their byte-faithful twin (see -/// [`llvm_byte_faithful_twin`]), with construction canonicalizing the value -/// before the store. -pub(crate) fn llvm_type_contains_i1(ctx: &Context, ty: TypeHandle) -> bool { - let ty_ref = ty.deref(ctx); - if let Some(integer) = ty_ref.downcast_ref::() { - return integer.width() == 1; - } - if let Some(array) = ty_ref.downcast_ref::() { - return llvm_type_contains_i1(ctx, array.elem_type()); - } - if let Some(vector) = ty_ref.downcast_ref::() { - return llvm_type_contains_i1(ctx, vector.elem_type()); - } - if let Some(struct_ty) = ty_ref.downcast_ref::() { - return struct_ty - .fields() - .any(|field| llvm_type_contains_i1(ctx, field)); - } - false -} - -/// The byte-faithful storage twin of an LLVM type: every `i1` leaf becomes -/// its canonical `i8` memory byte, recursively through structs and arrays. -/// -/// Rust guarantees a bool occupies one full byte holding exactly 0 or 1, so -/// storing the twin (with each bool zero-extended) writes the same bytes the -/// host writes, while re-loading the original type from those canonical -/// bytes remains well-defined. Sizes and alignments are unchanged because an -/// `i1` already occupies one byte of storage. -/// -/// Returns `None` for shapes with a different memory story (`i1` vectors are -/// bit-packed masks) or unknown containers; callers must fail closed. -pub(crate) fn llvm_byte_faithful_twin(ctx: &mut Context, ty: TypeHandle) -> Option { - if !llvm_type_contains_i1(ctx, ty) { - return Some(ty); - } - enum Shape { - Bool, - Array(TypeHandle, u64), - Struct { - fields: Vec, - layout: llvm_types::StructLayout, - }, - Other, - } - let shape = { - let ty_ref = ty.deref(ctx); - if ty_ref - .downcast_ref::() - .is_some_and(|integer| integer.width() == 1) - { - Shape::Bool - } else if let Some(array) = ty_ref.downcast_ref::() { - Shape::Array(array.elem_type(), array.size()) - } else if let Some(struct_ty) = ty_ref.downcast_ref::() { - Shape::Struct { - fields: struct_ty.fields().collect(), - layout: struct_ty.layout(), - } - } else { - Shape::Other - } - }; - match shape { - Shape::Bool => Some(IntegerType::get(ctx, 8, Signedness::Signless).into()), - Shape::Array(elem, count) => { - let twin = llvm_byte_faithful_twin(ctx, elem)?; - Some(llvm_types::ArrayType::get(ctx, twin, count).into()) - } - Shape::Struct { fields, layout } => { - let twins = fields - .into_iter() - .map(|field| llvm_byte_faithful_twin(ctx, field)) - .collect::>>()?; - Some(llvm_types::StructType::get_unnamed(ctx, (twins, layout)).into()) - } - Shape::Other => None, - } -} - -/// Whether a MIR aggregate contains a semantic Rust bool value. -/// -/// This deliberately stops at pointers/slices: a pointee bool does not occupy -/// bytes in the aggregate itself. It also stops at nested enums, whose own -/// slot-map construction is responsible for canonicalizing a top-level bool -/// or rejecting a deeper one. Inspecting the MIR type is necessary because a -/// union may lower to raw i8 storage and thereby hide a bool from the LLVM -/// type-level check above. -fn mir_type_contains_i1(ctx: &Context, ty: TypeHandle) -> bool { - let ty_ref = ty.deref(ctx); - if let Some(integer) = ty_ref.downcast_ref::() { - return integer.width() == 1; - } - if let Some(struct_ty) = ty_ref.downcast_ref::() { - return struct_ty - .field_types() - .iter() - .copied() - .any(|field| mir_type_contains_i1(ctx, field)); - } - if let Some(tuple_ty) = ty_ref.downcast_ref::() { - return tuple_ty - .get_types() - .iter() - .copied() - .any(|field| mir_type_contains_i1(ctx, field)); - } - if let Some(array_ty) = ty_ref.downcast_ref::() { - return mir_type_contains_i1(ctx, array_ty.element_ty); - } - if let Some(union_ty) = ty_ref.downcast_ref::() { - return union_ty - .field_types() - .iter() - .copied() - .any(|field| mir_type_contains_i1(ctx, field)); - } - false -} - -/// Whether loading and storing this LLVM value preserves every byte in its -/// allocation. In particular, `i1` is not byte-faithful: it occupies one -/// addressable byte, but LLVM does not define the upper seven stored bits. -pub(crate) fn llvm_type_is_byte_faithful(ctx: &Context, ty: TypeHandle) -> bool { - let ty_ref = ty.deref(ctx); - if let Some(int_ty) = ty_ref.downcast_ref::() { - return int_ty.width().is_multiple_of(8); - } - if ty_ref.is::() - || ty_ref.is::() - || ty_ref.is::() - || ty_ref.is::() - { - return true; - } - if let Some(array) = ty_ref.downcast_ref::() { - return llvm_type_is_byte_faithful(ctx, array.elem_type()); - } - if let Some(struct_ty) = ty_ref.downcast_ref::() { - let fields: Vec<_> = struct_ty.fields().collect(); - if struct_ty.layout() == llvm_types::StructLayout::Packed { - return fields - .into_iter() - .all(|field| llvm_type_is_byte_faithful(ctx, field)); - } - - let mut end = 0u64; - let mut max_align = 1u64; - for field in fields { - if !llvm_type_is_byte_faithful(ctx, field) { - return false; - } - let Some((field_size, field_align)) = llvm_type_size_align(ctx, field) else { - return false; - }; - let field_align = field_align.max(1); - let aligned_end = end.div_ceil(field_align) * field_align; - if aligned_end != end { - // LLVM would insert bytes that are not represented by an SSA - // field. Loading and re-storing the aggregate would lose them. - return false; - } - end += field_size; - max_align = max_align.max(field_align); - } - // Reject implicit trailing padding for the same reason. Explicit - // `[N x i8]` padding fields keep `end` equal to the allocation size. - return end.div_ceil(max_align) * max_align == end; - } - false -} - -/// Size of a MIR-level type from rustc layout truth, when stored. -/// -/// `MirStructType`, `MirTupleType`, `MirUnionType`, and `MirEnumType` carry -/// `total_size` (interior and trailing padding included) straight from -/// rustc's layout query; arrays of such aggregates multiply it out. Returns -/// `None` when no stored size is available (e.g. niched/single-variant enums -/// store 0) and the caller must fall back to the LLVM-level approximation. -fn mir_stored_size(ctx: &Context, mir_ty: TypeHandle) -> Option { - let ty_ref = mir_ty.deref(ctx); - if let Some(s) = ty_ref.downcast_ref::() { - if s.total_size() > 0 { - return Some(s.total_size()); - } - return None; - } - if let Some(t) = ty_ref.downcast_ref::() { - if t.total_size > 0 { - return Some(t.total_size); - } - return None; - } - if let Some(e) = ty_ref.downcast_ref::() { - if e.total_size() > 0 { - return Some(e.total_size()); - } - return None; - } - if let Some(u) = ty_ref.downcast_ref::() { - return Some(u.total_size()); - } - if let Some(a) = ty_ref.downcast_ref::() { - let elem_ty = a.element_ty; - let size = a.size; - // Checked: alignment claims consume this, and a wrapped product - // would masquerade as a small, trusted stride. - return mir_stored_size(ctx, elem_ty).and_then(|elem_size| elem_size.checked_mul(size)); - } - None -} - -/// Exact byte stride of a MIR array's element, when provable. -/// -/// Element `i` of an array lives at byte `i * stride`, so any alignment -/// claim about an element address is only as strong as the stride it -/// multiplies, and a wrong stride is a miscompile. Scalars have exact -/// sizes; MIR aggregates answer with rustc's stored size (interior and -/// trailing padding included) via [`mir_stored_size`]; arrays of scalars -/// multiply out. Everything else answers `None` so the caller declines -/// instead of guessing. -pub(crate) fn mir_element_stride(ctx: &Context, mir_ty: TypeHandle) -> Option { - if let Some(size) = mir_stored_size(ctx, mir_ty) { - return Some(size); - } - let ty_ref = mir_ty.deref(ctx); - if let Some(int_ty) = ty_ref.downcast_ref::() { - return Some(u64::from(int_ty.width()).div_ceil(8)); - } - // The importer's f16 is MirFP16Type; the HalfType arm covers IR that - // already carries the converted LLVM scalar. - if ty_ref.is::() || ty_ref.is::() { - return Some(2); - } - if ty_ref.is::() { - return Some(4); - } - if ty_ref.is::() { - return Some(8); - } - if let Some(ptr_ty) = ty_ref.downcast_ref::() { - // Generic-space pointers are 64-bit under every data layout the - // exporter can choose; a shared-space pointer is 32-bit under the - // modern NVVM layout alone, so its stored stride is target-dependent. - // Claim the former, decline the latter. - return (ptr_ty.address_space == 0).then_some(8); - } - if let Some(array_ty) = ty_ref.downcast_ref::() { - // Arrays of sized aggregates already answered through - // `mir_stored_size` above; this recursion serves arrays of - // scalars and pointers, and deeper such arrays. - let element_ty = array_ty.element_type(); - let count = array_ty.size(); - drop(ty_ref); - return mir_element_stride(ctx, element_ty)?.checked_mul(count); - } - None -} - -/// Exact ABI alignment carried by a MIR aggregate type, when rustc layout is -/// available. -/// -/// LLVM aggregate types cannot encode a Rust `repr(align(N))` raise. Tuples, -/// structs, enums, and unions therefore carry rustc's alignment explicitly in -/// the MIR dialect. Arrays have the same ABI alignment as their element, so -/// recurse through any number of array layers instead of relying on the -/// converted LLVM element's structural alignment. -pub(crate) fn mir_type_abi_align(ctx: &Context, mir_ty: TypeHandle) -> Option { - let ty_ref = mir_ty.deref(ctx); - if let Some(tuple_ty) = ty_ref.downcast_ref::() { - return Some(tuple_ty.abi_align()).filter(|align| *align > 0); - } - if let Some(struct_ty) = ty_ref.downcast_ref::() { - return Some(struct_ty.abi_align).filter(|align| *align > 0); - } - if let Some(enum_ty) = ty_ref.downcast_ref::() { - return Some(enum_ty.abi_align()).filter(|align| *align > 0); - } - if let Some(union_ty) = ty_ref.downcast_ref::() { - return Some(union_ty.abi_align()).filter(|align| *align > 0); - } - if let Some(array_ty) = ty_ref.downcast_ref::() { - let element_ty = array_ty.element_type(); - return mir_type_abi_align(ctx, element_ty); - } - None -} - -/// LLVM natural-layout `(size, align)` of an exported LLVM type, in bytes. -/// -/// Mirrors LLVM's default data layout for nvptx64 (scalars align to their -/// size, arrays to their element, non-packed structs to their widest field). -/// Computes the real allocation size (interior and trailing padding -/// included), which is what GEP striding and the enum size check below -/// need. Answers `None` rather than guess: sizes here are exact or absent. -pub(crate) fn llvm_type_size_align(ctx: &Context, ty: TypeHandle) -> Option<(u64, u64)> { - let ty_ref = ty.deref(ctx); - - if let Some(int_ty) = ty_ref.downcast_ref::() { - let size = (int_ty.width() as u64).div_ceil(8); - // i8 → 1, i16 → 2, i32 → 4, i64 → 8, i128 → 16. - return Some((size, size.next_power_of_two().min(16))); - } - if ty_ref.is::() { - return Some((2, 2)); - } - if ty_ref.is::() { - return Some((4, 4)); - } - if ty_ref.is::() { - return Some((8, 8)); - } - if ty_ref.is::() { - // Lowering runs before the exporter chooses the minimal, legacy, or - // modern NVPTX data layout. The first two use 64-bit pointers in all - // address spaces; modern NVVM alone uses p3:32. Callers that expose - // physical shared-pointer width must either genericize a direct - // pointer first or reject the target-dependent representation. - return Some((8, 8)); - } - if let Some(arr_ty) = ty_ref.downcast_ref::() { - let (elem_size, elem_align) = llvm_type_size_align(ctx, arr_ty.elem_type())?; - return Some((elem_size.checked_mul(arr_ty.size())?, elem_align.max(1))); - } - if let Some(vector) = ty_ref.downcast_ref::() { - if vector.is_scalable() { - return None; - } - let element_bits = { - let element = vector.elem_type(); - let element_ref = element.deref(ctx); - if let Some(integer) = element_ref.downcast_ref::() { - u64::from(integer.width()) - } else if let Some(float) = type_cast::(&*element_ref) { - u64::try_from(float.get_semantics().bits).ok()? - } else if element_ref.is::() { - 64 - } else { - return None; - } - }; - let total_bits = element_bits.checked_mul(u64::from(vector.num_elements()))?; - let size = total_bits.div_ceil(8); - // Both cuda-oxide NVPTX data layouts explicitly define fixed vector - // ABI alignment only for these widths. Refuse to guess LLVM defaults - // for any other width in physical Rust layout code. - let align = match total_bits { - 16 => 2, - 32 => 4, - 64 => 8, - 128 => 16, - _ => return None, - }; - return Some((size, align)); - } - if let Some(struct_ty) = ty_ref.downcast_ref::() { - let fields: Vec<_> = struct_ty.fields().collect(); - if struct_ty.layout() == llvm_types::StructLayout::Packed { - let mut size = 0u64; - for field in fields { - let (field_size, _) = llvm_type_size_align(ctx, field)?; - size = size.checked_add(field_size)?; - } - return Some((size, 1)); - } - - let (_end, size, align) = natural_struct_layout(ctx, &fields)?; - return Some((size, align)); - } - - None -} - -/// The one walk defining natural (non-packed) LLVM struct placement: each -/// field starts at the running offset rounded up to its own alignment. -/// -/// Returns `(offsets, end, align)`: the byte offset of every field, the -/// unrounded offset just past the last field, and the widest field alignment. -/// [`natural_struct_layout`] and [`StructSlotMap::natural_slot_offsets`] are -/// both views of this walk, so a size question and an offset question can -/// never disagree. -fn natural_layout_walk(ctx: &Context, fields: &[TypeHandle]) -> Option<(Vec, u64, u64)> { - let mut offsets = Vec::with_capacity(fields.len()); - let mut end = 0u64; - let mut align = 1u64; - for field in fields { - let (field_size, field_align) = llvm_type_size_align(ctx, *field)?; - let field_align = field_align.max(1); - end = end.div_ceil(field_align) * field_align; - offsets.push(end); - end = end.checked_add(field_size)?; - align = align.max(field_align); - } - Some((offsets, end, align)) -} - -/// Natural (non-packed) LLVM struct layout over `fields`. -/// -/// Returns `(end, size, align)` where `end` is the unrounded offset just past -/// the last field, `size` is `end` rounded up to the struct alignment (the -/// allocation size LLVM uses for GEP striding), and `align` is the widest -/// field alignment. -pub(crate) fn natural_struct_layout( - ctx: &Context, - fields: &[TypeHandle], -) -> Option<(u64, u64, u64)> { - let (_offsets, end, align) = natural_layout_walk(ctx, fields)?; - let size = end.div_ceil(align) * align; - Some((end, size, align)) -} - -/// The LLVM struct for an enum, plus a map saying where the tag and each -/// payload field ended up. -/// -/// The struct type and the indices into it are produced by one walk in -/// [`build_enum_slot_map`], so they can never disagree. (Computing them -/// separately is how the issue #128 class of bug happened for structs.) -pub(crate) struct EnumSlotMap { - /// The final LLVM struct type, including any `[N x i8]` filler slots. - pub llvm_struct_ty: TypeHandle, - /// Which struct slot holds rustc's physical tag/niche carrier. `None` - /// for `Single` and `Empty` layouts, which have no tag in memory. - pub carrier_slot: Option, - /// Converted physical carrier type (integer or pointer), when present. - pub carrier_llvm_ty: Option, - /// Which struct slot holds each payload field, in the flattened - /// order of `MirEnumType::all_field_types`. `None` means the field - /// has no slot of its own: it is zero-sized, or its bytes are shared - /// with a different-typed field of another variant. Such fields are - /// read and written through memory at `field_offsets` instead. - pub field_slots: Vec>, - /// Byte position of each payload field inside the enum. - pub field_offsets: Vec, - /// Converted LLVM type of each payload field. - pub field_llvm_types: Vec, -} - -/// Build the LLVM struct for an enum, placing everything at the byte -/// positions rustc chose. -/// -/// Why this matters: the host (CPU) lays out enum values with rustc's -/// layout. If the device used different byte positions, every enum -/// passed to a kernel would be read wrong. So the device struct is built -/// to have the same bytes, position for position. -/// -/// The wrinkle is that enum variants SHARE bytes (only one variant is -/// alive at a time), and an LLVM struct cannot say "these two fields -/// overlap". The slot map resolves each field one of three ways: -/// -/// ```text -/// #[repr(u32)] enum E { A(u32), B(f32), C } -/// rustc: 8 bytes, tag at byte 0, A's u32 and B's f32 both at byte 4 -/// -/// LLVM struct: { i32, i32 } -/// | | -/// tag_slot=0 A's payload: own slot (nothing else typed i32 -/// wanted byte 4 first... B did, see below) -/// -/// - own slot: the field's bytes collide with nothing already placed. -/// - shared slot: another variant already placed the SAME type at the -/// SAME position; both map to that slot. (If B were -/// B(u32), A and B would simply share slot 1.) -/// - no slot: the bytes are taken by a different type (B's f32 vs -/// A's u32 here). The field is still at byte 4, just not -/// nameable as a struct field; reads and writes go -/// through memory: spill the value to a stack slot, then -/// use a byte-precise pointer. No slot, but no lie. -/// ``` -/// -/// Gaps between placed fields, and the tail, are covered with explicit -/// `[N x i8]` filler so the struct's size is exactly rustc's no matter -/// what LLVM's own layout rules would have done. -/// -/// Direct and Niche carriers are claimed first, so a semantic field with a -/// different SSA type cannot redefine the same physical bytes. Single and -/// Empty layouts simply have no carrier. Unknown layouts are rejected. -/// -/// If the finished struct's size does not come out equal to rustc's, -/// something is deeply wrong and lowering would miscompile, so that is a -/// hard error rather than a debug assertion. -pub(crate) fn build_enum_slot_map( - ctx: &mut Context, - ty: TypeHandle, -) -> Result { - let ( - name, - discriminant_ty, - all_field_types, - all_field_offsets, - all_field_sizes, - variant_field_counts, - variant_inhabited, - tag_offset, - total_size, - abi_align, - layout_kind, - carrier_kind, - carrier_width, - carrier_address_space, - ) = { - let ty_ref = ty.deref(ctx); - let enum_ty = ty_ref - .downcast_ref::() - .ok_or_else(|| anyhow::anyhow!("build_enum_slot_map: expected MirEnumType"))?; - ( - enum_ty.name().to_string(), - enum_ty.discriminant_ty, - enum_ty.all_field_types.clone(), - enum_ty.all_field_offsets.clone(), - enum_ty.all_field_sizes.clone(), - enum_ty.variant_field_counts.clone(), - enum_ty.variant_inhabited.clone(), - enum_ty.tag_offset(), - enum_ty.total_size(), - enum_ty.abi_align(), - enum_ty.layout_kind, - enum_ty.carrier_kind, - enum_ty.carrier_width, - enum_ty.carrier_address_space, - ) - }; - - if layout_kind == EnumLayoutKind::Unknown { - return Err(anyhow::anyhow!( - "enum slot map: `{}` has unknown physical layout; refusing to guess its bytes", - name - )); - } - if carrier_kind == EnumCarrierKind::Pointer - && carrier_address_space == llvm_types::address_space::SHARED - { - return Err(anyhow::anyhow!( - "enum slot map: `{}` uses a shared-memory pointer carrier whose size is target-mode dependent (64-bit PTX/legacy, 32-bit modern NVVM); refusing target-agnostic enum lowering", - name - )); - } - if carrier_kind == EnumCarrierKind::Integer && !carrier_width.is_multiple_of(8) { - return Err(anyhow::anyhow!( - "enum slot map: `{}` integer carrier width {} is not a whole number of bytes; refusing physical storage with unspecified upper byte bits", - name, - carrier_width - )); - } - - let carrier_ty: Option = match carrier_kind { - EnumCarrierKind::None => None, - EnumCarrierKind::Integer if layout_kind == EnumLayoutKind::Direct => { - let converted = convert_type(ctx, discriminant_ty)?; - let width = converted - .deref(ctx) - .downcast_ref::() - .map(IntegerType::width); - if width != Some(carrier_width) { - return Err(anyhow::anyhow!( - "enum slot map: `{}` direct carrier does not match its declared discriminant type", - name - )); - } - Some(converted) - } - EnumCarrierKind::Integer => { - Some(IntegerType::get(ctx, carrier_width, Signedness::Signless).into()) - } - EnumCarrierKind::Pointer => { - Some(llvm_types::PointerType::get(ctx, carrier_address_space).into()) - } - }; - let mut field_llvm_types = Vec::with_capacity(all_field_types.len()); - for &field_ty in &all_field_types { - field_llvm_types.push(convert_type(ctx, field_ty)?); - } - - if all_field_offsets.len() != all_field_types.len() - || all_field_sizes.len() != all_field_types.len() - { - return Err(anyhow::anyhow!( - "enum slot map: `{}` has {} field offsets for {} fields", - name, - all_field_offsets.len().min(all_field_sizes.len()), - all_field_types.len() - )); - } - - // Phase 1: decide who gets a struct slot. The physical carrier goes - // first so a semantic field can never claim its bytes using a different - // type (e.g. `bool` is i1 semantically but i8 in Option storage). - // claims: (byte position, byte size, converted type), no two overlap. - let mut claims: Vec<(u64, u64, TypeHandle)> = Vec::new(); - let carrier_claim = if let Some(carrier_ty) = carrier_ty { - let (carrier_size, carrier_align) = - llvm_type_size_align(ctx, carrier_ty).ok_or_else(|| { - anyhow::anyhow!( - "enum slot map: `{}` carrier has unsupported LLVM layout", - name - ) - })?; - let expected_size = u64::from(carrier_width).div_ceil(8); - if carrier_size != expected_size - || tag_offset % carrier_align.max(1) != 0 - || tag_offset - .checked_add(carrier_size) - .is_none_or(|end| end > total_size) - { - return Err(anyhow::anyhow!( - "enum slot map: `{}` carrier (size {}, align {}) cannot sit at byte {} of {}", - name, - carrier_size, - carrier_align, - tag_offset, - total_size - )); - } - claims.push((tag_offset, carrier_size, carrier_ty)); - Some(0usize) - } else { - None - }; - - let mut claim_of_field: Vec> = vec![None; field_llvm_types.len()]; - let mut field_is_inhabited = Vec::with_capacity(field_llvm_types.len()); - for (variant, count) in variant_field_counts.iter().enumerate() { - field_is_inhabited.extend(std::iter::repeat_n( - variant_inhabited.get(variant).copied().unwrap_or(0) != 0, - *count as usize, - )); - } - let mut order: Vec = (0..field_llvm_types.len()).collect(); - // At one byte position, prefer a representation that preserves all of - // its stored bits. This makes the result independent of source-variant - // order and prevents an i1/bool view from claiming storage that a later - // i8 view needs to preserve. A scalar i1 always uses an i8 storage claim: - // construction explicitly zero-extends that scalar below. - order.sort_by_key(|&i| { - ( - all_field_offsets[i], - !llvm_type_is_byte_faithful(ctx, field_llvm_types[i]), - i, - ) - }); - for flat in order { - if !field_is_inhabited.get(flat).copied().unwrap_or(false) { - continue; - } - let llvm_ty = field_llvm_types[flat]; - // Enum payload storage uses one target-stable physical view. Shared - // pointer leaves become generic pointers recursively through LLVM - // structs (MIR structs/tuples) and bounded arrays, while bool leaves - // become canonical i8 bytes. Oversized shared-pointer arrays and - // pointer vectors fail closed here. - let storage_ty = enum_payload_storage_type(ctx, llvm_ty).map_err(|error| { - anyhow::anyhow!("enum slot map: `{}` field {}: {error}", name, flat) - })?; - let (size, align) = llvm_type_size_align(ctx, storage_ty).ok_or_else(|| { - anyhow::anyhow!( - "enum slot map: `{}` field {} has unsupported LLVM size/alignment", - name, - flat - ) - })?; - if size != all_field_sizes[flat] { - return Err(anyhow::anyhow!( - "enum slot map: `{}` field {} lowers to {} bytes but rustc says {}", - name, - flat, - size, - all_field_sizes[flat] - )); - } - if size == 0 || is_zero_sized_type(ctx, llvm_ty) { - // ZSTs own no bytes and no slot. - continue; - } - let offset = all_field_offsets[flat]; - if offset.checked_add(size).is_none_or(|end| end > total_size) { - return Err(anyhow::anyhow!( - "enum slot map: `{}` field {} (size {}) at byte {} exceeds total size {}", - name, - flat, - size, - offset, - total_size - )); - } - if !offset.is_multiple_of(align.max(1)) { - return Err(anyhow::anyhow!( - "enum slot map: `{}` field {} requires alignment {} but rustc offset {} is not aligned; packed enum payload access is not yet supported", - name, - flat, - align, - offset - )); - } - - let is_scalar_i1 = llvm_ty - .deref(ctx) - .downcast_ref::() - .is_some_and(|integer| integer.width() == 1); - // A bool the MIR type promises but the LLVM lowering hides (a union - // whose storage is a raw byte blob) cannot be canonicalized here: - // its bytes were written by the union's own stores, not by this - // enum's construction. - if !is_scalar_i1 - && mir_type_contains_i1(ctx, all_field_types[flat]) - && !llvm_type_contains_i1(ctx, llvm_ty) - { - return Err(anyhow::anyhow!( - "enum slot map: `{}` field {} contains bool storage hidden behind a union; cuda-oxide cannot prove those bytes are canonical", - name, - flat - )); - } - - // Rust bool is a semantic i1 but occupies a complete byte in memory. - // Never make i1 the struct's physical storage type: give a standalone - // bool an i8 claim, or reuse the exact i8 carrier/field claim already - // covering that byte. Construction leaves the bool slotless and the - // deferred store below explicitly zero-extends i1 -> i8; extraction - // loads i1 from that byte. - if is_scalar_i1 { - let colliding_claims = claims - .iter() - .filter(|&&(o, s, _)| offset < o + s && o < offset + size) - .collect::>(); - if colliding_claims.is_empty() { - let byte_ty: TypeHandle = IntegerType::get(ctx, 8, Signedness::Signless).into(); - claims.push((offset, 1, byte_ty)); - continue; - } - let has_exact_i8_storage = colliding_claims.iter().all(|&&(o, s, claim_ty)| { - o == offset - && s == 1 - && claim_ty - .deref(ctx) - .downcast_ref::() - .is_some_and(|integer| integer.width() == 8) - }); - if has_exact_i8_storage { - continue; - } - return Err(anyhow::anyhow!( - "enum slot map: `{}` scalar bool field {} overlaps storage other than one exact i8 byte; refusing to expose non-canonical bool bits", - name, - flat - )); - } - - // Aggregates containing bool claim their byte-faithful twin (every - // i1 leaf becomes its canonical i8 memory byte) and stay slotless: - // construction canonicalizes the value and writes it at its byte - // offset; extraction re-loads the original type from the canonical - // bytes, exactly like the scalar-bool path above. - // Bool-bearing aggregates remain slotless so their canonical byte view - // continues through the established spill path. Pointer-only nested - // aggregates may own a slot because construction/extraction now rebuild - // them recursively at the semantic/physical boundary. - let field_gets_slot = !llvm_type_contains_i1(ctx, llvm_ty); - - // Another variant already placed the same storage type at the same - // position? Then both fields can simply use that claim: variants - // share bytes, and here they even agree on the type. - if let Some(ci) = claims - .iter() - .position(|&(o, _, t)| o == offset && t == storage_ty) - { - if field_gets_slot { - claim_of_field[flat] = Some(ci); - } - continue; - } - // The bytes are taken by a different type. Pointer-free values can - // use the memory fallback below. Pointer-bearing values may do so only - // when a physical-layout walk proves that every pointer leaf exactly - // matches an existing pointer slot. This is the common - // `Option<(usize, &T)>` shape: the niche carrier is the tuple's pointer - // field. A pointer/non-pointer overlap, address-space mismatch, or - // additional pointer without a pointer slot still fails closed. - let colliding_claims = claims - .iter() - .filter(|&&(o, s, _)| offset < o + s && o < offset + size) - .collect::>(); - if !colliding_claims.is_empty() { - let has_pointer_overlap = llvm_type_contains_pointer(ctx, storage_ty) - || colliding_claims - .iter() - .any(|&&(_, _, claim_ty)| llvm_type_contains_pointer(ctx, claim_ty)); - // Pointer-bearing overlaps must back every pointer leaf with a real - // `ptr` slot so the memory round-trip preserves provenance. The - // niche carrier backs the leaf that coincides with it; any further - // pointer leaf (the extra slice pointer in `split_at_mut_checked`'s - // `Option<(&mut [T], &mut [T])>`) gets its own fresh `ptr` slot. - // `ProvenanceLoss` means a pointer punned against non-pointer - // bits, which stays fail-closed; `OverArrayLeafBound` reports the - // same bounded rewrite contract as the payload storage gate. - let extra_pointer_claims = if has_pointer_overlap { - match analyze_pointer_overlap(ctx, offset, size, storage_ty, &colliding_claims) { - Ok(extra) => extra, - Err(PointerOverlapRejection::OverArrayLeafBound { required }) => { - return Err(anyhow::anyhow!( - "enum slot map: `{}` field {} overlaps pointer storage whose arrays are not supported above the bounded rewrite limit; rewrite requires {required} pointer conversions, supported bound is {MAX_ENUM_PAYLOAD_ARRAY_REWRITE_LEAVES}", - name, - flat - )); - } - Err(PointerOverlapRejection::ProvenanceLoss) => { - return Err(anyhow::anyhow!( - "enum slot map: `{}` has overlapping pointer and non-identical storage at byte {}; refusing to erase LLVM pointer provenance", - name, - offset - )); - } - } - } else { - Vec::new() - }; - let incoming_is_byte_faithful = llvm_type_is_byte_faithful(ctx, storage_ty); - let claims_are_byte_faithful = colliding_claims - .iter() - .all(|&&(_, _, claim_ty)| llvm_type_is_byte_faithful(ctx, claim_ty)); - if !incoming_is_byte_faithful || !claims_are_byte_faithful { - return Err(anyhow::anyhow!( - "enum slot map: `{}` field {} overlaps non-identical storage but its lowered type is not byte-faithful (for example, it may contain implicit padding); refusing a type-punned store", - name, - flat - )); - } - // Back each extra pointer leaf with its own `ptr` slot (the field - // itself stays slotless and round-trips through memory). Added after - // the byte-faithfulness gate so `colliding_claims`'s borrow of - // `claims` has ended before we push. - for leaf in extra_pointer_claims { - let ptr_ty = llvm_types::PointerType::get(ctx, leaf.address_space).into(); - claims.push((leaf.offset, leaf.size, ptr_ty)); - } - continue; - } - claims.push((offset, size, storage_ty)); - if field_gets_slot { - claim_of_field[flat] = Some(claims.len() - 1); - } - } - - // Phase 2: lay the slots down in byte order, filling every gap (and - // the tail) so the struct's size is exactly rustc's. One integer per gap - // where that is layout-neutral, else [N x i8]; see `make_enum_filler_type`. - let mut emit_order: Vec = (0..claims.len()).collect(); - emit_order.sort_by_key(|&ci| claims[ci].0); - let mut llvm_fields: Vec = Vec::new(); - let mut slot_of_claim: Vec = vec![0; claims.len()]; - let mut current_offset: u64 = 0; - for &ci in &emit_order { - let (offset, size, llvm_ty) = claims[ci]; - if current_offset < offset { - let filler = - make_enum_filler_type(ctx, current_offset, offset - current_offset, abi_align); - llvm_fields.push(filler); - current_offset = offset; - } - slot_of_claim[ci] = llvm_fields.len() as u32; - llvm_fields.push(llvm_ty); - current_offset += size; - } - if current_offset < total_size { - let filler = - make_enum_filler_type(ctx, current_offset, total_size - current_offset, abi_align); - llvm_fields.push(filler); - } - - // Sanity: the struct we just built must be exactly rustc's size. - // Arrays of enums step by this size, so a mismatch means every - // element after the first is read from the wrong place. That is a - // guaranteed miscompile, hence a hard error, not a debug check. - let (_end, natural_size, natural_align) = natural_struct_layout(ctx, &llvm_fields) - .ok_or_else(|| anyhow::anyhow!("enum slot map: `{}` has unsupported LLVM layout", name))?; - if natural_size != total_size { - return Err(anyhow::anyhow!( - "enum slot map: `{}` lowered to {} bytes but rustc says {}", - name, - natural_size, - total_size - )); - } - let required_align = abi_align.max(1); - if natural_align > required_align { - return Err(anyhow::anyhow!( - "enum slot map: `{}` lowered with alignment {} but rustc requires {}; explicit over-aligned enum storage is not yet supported", - name, - natural_align, - abi_align - )); - } - if natural_align < required_align { - // The byte claims alone can under-align the storage, e.g. when the - // only claim is an i8 niche carrier inside a 4-aligned enum. Raise - // the struct's alignment with a zero-length anchor field, the same - // mechanism union storage uses; it occupies no bytes, so every slot - // index simply shifts by one. - let anchor_int = IntegerType::get(ctx, (required_align * 8) as u32, Signedness::Signless); - let anchor: TypeHandle = llvm_types::ArrayType::get(ctx, anchor_int.into(), 0).into(); - llvm_fields.insert(0, anchor); - for slot in &mut slot_of_claim { - *slot += 1; - } - } - - let field_slots = claim_of_field - .into_iter() - .map(|c| c.map(|ci| slot_of_claim[ci])) - .collect(); - Ok(EnumSlotMap { - llvm_struct_ty: llvm_types::StructType::get_unnamed( - ctx, - (llvm_fields, llvm_types::StructLayout::Unpacked), - ) - .into(), - carrier_slot: carrier_claim.map(|claim| slot_of_claim[claim]), - carrier_llvm_ty: carrier_ty, - field_slots, - field_offsets: all_field_offsets, - field_llvm_types, - }) -} - -/// Convert a `MirEnumType` to its LLVM struct representation. -/// -/// Thin wrapper over [`build_enum_slot_map`], which explains the layout. -/// Any op that needs an index into the converted enum must take it from -/// the slot map, never compute it by hand. -pub(crate) fn convert_enum_to_llvm( - ctx: &mut Context, - ty: TypeHandle, -) -> Result { - Ok(build_enum_slot_map(ctx, ty)?.llvm_struct_ty) -} - -/// Return an enum name only when the dialect lacks rustc's physical layout. -/// All importer-produced Direct, Niche, Single, and Empty layouts are -/// byte-faithful; legacy `Unknown` values are rejected everywhere rather than -/// receiving a guessed internal representation. -pub(crate) fn enum_unmodeled_in_memory(ctx: &Context, ty: TypeHandle) -> Option { - let ty_ref = ty.deref(ctx); - let enum_ty = ty_ref.downcast_ref::()?; - (enum_ty.layout_kind == EnumLayoutKind::Unknown).then(|| enum_ty.name().to_string()) -} - -/// Search a kernel parameter's type for an enum the host and device -/// would disagree about (see [`enum_unmodeled_in_memory`]). -/// -/// The search looks everywhere host data can hide: behind pointers, -/// inside slices and arrays, in struct/tuple fields, and in other enums' -/// payloads. It returns the first offending enum's name. -/// -/// Kernel signatures are checked early for a focused diagnostic. Lowering also -/// rejects Unknown layouts for locals, globals, and physical operations, so no -/// guessed representation can escape through an internal-only path. -/// -/// `visited` breaks cycles through recursive types (`TypeHandle` is -/// interned, so equality is identity). -pub(crate) fn find_unmodeled_enum_in_abi( - ctx: &mut Context, - ty: TypeHandle, - visited: &mut Vec, -) -> Result, anyhow::Error> { - if visited.contains(&ty) { - return Ok(None); - } - visited.push(ty); - - if let Some(name) = enum_unmodeled_in_memory(ctx, ty) { - return Ok(Some(name)); - } - - let children: Vec = { - let ty_ref = ty.deref(ctx); - if let Some(p) = ty_ref.downcast_ref::() { - vec![p.pointee] - } else if let Some(s) = ty_ref.downcast_ref::() { - vec![s.element_ty] - } else if let Some(s) = ty_ref.downcast_ref::() { - let mut nested = vec![s.element_ty]; - nested.extend_from_slice(&s.space_tys); - nested - } else if let Some(a) = ty_ref.downcast_ref::() { - vec![a.element_ty] - } else if let Some(s) = ty_ref.downcast_ref::() { - s.field_types.clone() - } else if let Some(u) = ty_ref.downcast_ref::() { - u.field_types.clone() - } else if let Some(t) = ty_ref.downcast_ref::() { - t.get_types().to_vec() - } else if let Some(e) = ty_ref.downcast_ref::() { - e.all_field_types.clone() - } else { - vec![] - } - }; - - for child in children { - if let Some(name) = find_unmodeled_enum_in_abi(ctx, child, visited)? { - return Ok(Some(name)); - } - } - Ok(None) -} - -/// Prove that an initialized Rust global can be accessed through the LLVM -/// semantic type produced by this lowering pipeline. -/// -/// The initializer itself is emitted as exact bytes. That is only half of the -/// contract: later field GEPs and typed loads still use `mir_ty`. If that type -/// places a field at a different byte offset, an exact initializer can still be -/// read incorrectly (or even past the end of the object). Reject every shape -/// for which we cannot prove that the two views agree. -pub(crate) fn validate_initialized_global_layout( - ctx: &mut Context, - mir_ty: TypeHandle, - initializer_size: u64, - initializer_align: u64, -) -> Result<(), anyhow::Error> { - if initializer_align == 0 || !initializer_align.is_power_of_two() { - return Err(anyhow::anyhow!( - "initialized global has invalid rustc allocation alignment {}", - initializer_align - )); - } - - validate_initialized_global_type(ctx, mir_ty, &mut Vec::new())?; - - let llvm_ty = convert_type(ctx, mir_ty)?; - let (llvm_size, llvm_align) = llvm_type_size_align(ctx, llvm_ty) - .ok_or_else(|| anyhow::anyhow!("initialized global has unsupported LLVM size/alignment"))?; - if llvm_size != initializer_size || llvm_align > initializer_align { - return Err(anyhow::anyhow!( - "initialized global type is not byte-compatible with rustc's allocation: the lowered LLVM value has size/alignment {}/{}, but the initializer has size/alignment {}/{}", - llvm_size, - llvm_align, - initializer_size, - initializer_align - )); - } - - Ok(()) -} - -/// Validate an initialized global that carries pointer relocations. -/// -/// Ordinary initialized globals keep their existing conservative semantic -/// layout validator. Relocated globals use a separate segmented physical -/// carrier, so a top-level `repr(packed)` struct may use that relocation path -/// as long as rustc's field ranges are explicit, non-overlapping, and in-bounds. -/// One direct nested struct may use the same relaxation when the top-level -/// struct itself has an ordinary, non-divergent LLVM layout and the child's -/// LLVM packed representation exactly reproduces rustc's offsets and size. A -/// packed top-level struct may not stack the relaxation on a packed child, and -/// deeper packed nesting remains on the ordinary fail-closed path. -pub(crate) fn validate_relocated_initialized_global_layout( - ctx: &mut Context, - mir_ty: TypeHandle, - initializer_size: u64, - initializer_align: u64, -) -> Result<(), anyhow::Error> { - match validate_initialized_global_layout(ctx, mir_ty, initializer_size, initializer_align) { - Ok(()) => Ok(()), - Err(original_error) => { - if initializer_align == 0 || !initializer_align.is_power_of_two() { - return Err(original_error); - } - - let struct_ty = { - let ty_ref = mir_ty.deref(ctx); - ty_ref.downcast_ref::().cloned() - }; - let Some(struct_ty) = struct_ty else { - return Err(original_error); - }; - if !struct_ty.has_explicit_layout() - || struct_ty.total_size() != initializer_size - || struct_ty.abi_align != initializer_align - { - return Err(original_error); - } - - validate_relocated_struct_field_ranges(ctx, &struct_ty)?; - - let top_level_layout = StructLayoutInfo::of_struct(&struct_ty); - let top_level_map = build_struct_slot_map(ctx, &top_level_layout)?; - let top_level_is_ordinary = !top_level_map.layout_diverges - && top_level_map.by_value_layout_faithful - && top_level_map - .llvm_struct_ty - .deref(ctx) - .downcast_ref::() - .is_some_and(|ty| ty.layout() == llvm_types::StructLayout::Unpacked); - - let mut visited = vec![mir_ty]; - for field_ty in struct_ty.field_types.iter().copied() { - if top_level_is_ordinary { - validate_relocated_initialized_global_child(ctx, field_ty, &mut visited)?; - } else { - validate_initialized_global_type(ctx, field_ty, &mut visited)?; - } - } - Ok(()) - } - } -} - -/// Prove that rustc's explicit field ranges for a relocated top-level struct -/// are non-overlapping and fit within the evaluated allocation. -fn validate_relocated_struct_field_ranges( - ctx: &mut Context, - struct_ty: &MirStructType, -) -> Result<(), anyhow::Error> { - let layout = StructLayoutInfo::of_struct(struct_ty); - - // Relocated globals relax only natural LLVM placement. Preserve the - // ordinary struct-layout metadata invariants before accepting rustc's - // explicit byte ranges as the physical initializer contract. - build_struct_slot_map(ctx, &layout)?; - - let mut ranges = Vec::with_capacity(layout.field_types.len()); - for (decl_index, field_ty) in layout.field_types.iter().copied().enumerate() { - let field_size = if let Some(size) = mir_stored_size(ctx, field_ty) { - size - } else { - let llvm_ty = convert_type(ctx, field_ty)?; - llvm_type_size_align(ctx, llvm_ty) - .map(|(size, _)| size) - .ok_or_else(|| { - anyhow::anyhow!( - "relocated initialized struct `{}` field {} has unsupported size", - struct_ty.name(), - decl_index - ) - })? - }; - if field_size == 0 { - continue; - } - - let start = layout.field_offsets[decl_index]; - let end = start.checked_add(field_size).ok_or_else(|| { - anyhow::anyhow!( - "relocated initialized struct `{}` field {} range overflows", - struct_ty.name(), - decl_index - ) - })?; - if end > layout.total_size { - return Err(anyhow::anyhow!( - "relocated initialized struct `{}` field {} occupies bytes {}..{}, but the allocation is only {} bytes", - struct_ty.name(), - decl_index, - start, - end, - layout.total_size - )); - } - ranges.push((start, end, decl_index)); - } - - ranges.sort_by_key(|(start, _, _)| *start); - for pair in ranges.windows(2) { - let (_, previous_end, previous_index) = pair[0]; - let (next_start, _, next_index) = pair[1]; - if next_start < previous_end { - return Err(anyhow::anyhow!( - "relocated initialized struct `{}` fields {} and {} overlap in rustc's byte layout", - struct_ty.name(), - previous_index, - next_index - )); - } - } - - Ok(()) -} - -/// Validate one direct child of an ordinary relocated initialized global. -/// -/// The ordinary initialized-global validator remains the default. Only when it -/// rejects a direct `MirStructType` do we consider the relocation-specific -/// relaxation, and only if `build_struct_slot_map` proves that the divergent -/// rustc layout is represented exactly by an LLVM packed struct. Its own field -/// ranges must still be explicit, non-overlapping, and in-bounds. Children of -/// that packed struct go back through the ordinary validator, deliberately -/// limiting this exception to one nesting level. -fn validate_relocated_initialized_global_child( - ctx: &mut Context, - mir_ty: TypeHandle, - visited: &mut Vec, -) -> Result<(), anyhow::Error> { - if visited.contains(&mir_ty) { - return Ok(()); - } - - let mut ordinary_visited = visited.clone(); - let original_error = match validate_initialized_global_type(ctx, mir_ty, &mut ordinary_visited) - { - Ok(()) => { - *visited = ordinary_visited; - return Ok(()); - } - Err(error) => error, - }; - - let struct_ty = { - let ty_ref = mir_ty.deref(ctx); - ty_ref.downcast_ref::().cloned() - }; - let Some(struct_ty) = struct_ty else { - return Err(original_error); - }; - if !struct_ty.has_explicit_layout() { - return Err(original_error); - } - - let layout = StructLayoutInfo::of_struct(&struct_ty); - let map = build_struct_slot_map(ctx, &layout)?; - let is_exact_packed = map.layout_diverges - && map.by_value_layout_faithful - && map - .llvm_struct_ty - .deref(ctx) - .downcast_ref::() - .is_some_and(|ty| ty.layout() == llvm_types::StructLayout::Packed); - if !is_exact_packed { - return Err(original_error); - } - - validate_relocated_struct_field_ranges(ctx, &struct_ty)?; - - visited.push(mir_ty); - for field_ty in struct_ty.field_types.iter().copied() { - validate_initialized_global_type(ctx, field_ty, visited)?; - } - Ok(()) -} - -fn validate_initialized_global_type( - ctx: &mut Context, - mir_ty: TypeHandle, - visited: &mut Vec, -) -> Result<(), anyhow::Error> { - if visited.contains(&mir_ty) { - return Ok(()); - } - visited.push(mir_ty); - - if let Some(name) = enum_unmodeled_in_memory(ctx, mir_ty) { - return Err(anyhow::anyhow!( - "initialized global contains enum `{}` with unknown physical layout; refusing to guess a byte representation", - name - )); - } - - enum Kind { - Struct(MirStructType), - Tuple(MirTupleType), - Enum(MirEnumType), - Array { element_ty: TypeHandle, size: u64 }, - Leaf, - } - - let kind = { - let ty_ref = mir_ty.deref(ctx); - if let Some(struct_ty) = ty_ref.downcast_ref::() { - Kind::Struct(struct_ty.clone()) - } else if let Some(tuple_ty) = ty_ref.downcast_ref::() { - Kind::Tuple(tuple_ty.clone()) - } else if let Some(enum_ty) = ty_ref.downcast_ref::() { - Kind::Enum(enum_ty.clone()) - } else if let Some(array_ty) = ty_ref.downcast_ref::() { - Kind::Array { - element_ty: array_ty.element_ty, - size: array_ty.size, - } - } else { - Kind::Leaf - } - }; - - match kind { - Kind::Struct(struct_ty) => { - validate_initialized_struct_layout(ctx, mir_ty, &struct_ty)?; - for field_ty in struct_ty.field_types { - validate_initialized_global_type(ctx, field_ty, visited)?; - } - } - Kind::Tuple(tuple_ty) => { - if !tuple_ty.get_types().is_empty() { - // Tuples carry rustc's field offsets exactly like structs; - // prove the lowered aggregate reproduces them byte-for-byte. - let layout = StructLayoutInfo::of_tuple(&tuple_ty); - validate_initialized_aggregate_layout( - ctx, - mir_ty, - "tuple", - "tuple", - &layout, - tuple_ty.abi_align(), - tuple_ty.has_explicit_layout(), - )?; - } - for field_ty in tuple_ty.get_types().iter().copied() { - validate_initialized_global_type(ctx, field_ty, visited)?; - } - } - Kind::Enum(enum_ty) => { - if enum_ty.total_size() > 0 { - let llvm_ty = convert_type(ctx, mir_ty)?; - let (llvm_size, llvm_align) = - llvm_type_size_align(ctx, llvm_ty).ok_or_else(|| { - anyhow::anyhow!( - "initialized enum `{}` has unsupported LLVM size/alignment", - enum_ty.name() - ) - })?; - if llvm_size != enum_ty.total_size() || llvm_align > enum_ty.abi_align() { - return Err(anyhow::anyhow!( - "initialized enum `{}` is not byte-compatible with rustc's layout: the lowered LLVM value has size/alignment {}/{}, but rustc requires {}/{}", - enum_ty.name(), - llvm_size, - llvm_align, - enum_ty.total_size(), - enum_ty.abi_align() - )); - } - } - for field_ty in enum_ty.all_field_types { - validate_initialized_global_type(ctx, field_ty, visited)?; - } - } - Kind::Array { element_ty, size } => { - // A zero-length array contributes no element bytes. Its outer LLVM - // size/alignment is still checked by - // `validate_initialized_global_layout`, but recursively rejecting - // an element layout would incorrectly reject valid promoted - // `&mut [Packed; 0]` constants whose physical allocation is empty. - if size != 0 { - validate_initialized_global_type(ctx, element_ty, visited)?; - } - } - Kind::Leaf => {} - } - - Ok(()) -} - -fn validate_initialized_struct_layout( - ctx: &mut Context, - mir_ty: TypeHandle, - struct_ty: &MirStructType, -) -> Result<(), anyhow::Error> { - let has_layout = struct_ty.has_explicit_layout(); - let layout = StructLayoutInfo::of_struct(struct_ty); - let name = struct_ty.name().to_string(); - validate_initialized_aggregate_layout( - ctx, - mir_ty, - "struct", - &name, - &layout, - struct_ty.abi_align, - has_layout, - ) -} - -/// Shared byte-layout validation for initialized-global structs and tuples. -/// -/// Proves that the lowered LLVM aggregate places every field at exactly the -/// byte offset rustc chose and matches rustc's total size, so a constant -/// initializer written slot-by-slot reproduces the host bytes. -#[allow(clippy::too_many_arguments)] -fn validate_initialized_aggregate_layout( - ctx: &mut Context, - mir_ty: TypeHandle, - kind_noun: &str, - name: &str, - layout: &StructLayoutInfo, - abi_align: u64, - has_explicit_layout: bool, -) -> Result<(), anyhow::Error> { - if layout.total_size == 0 { - let llvm_ty = convert_type(ctx, mir_ty)?; - let (llvm_size, _) = llvm_type_size_align(ctx, llvm_ty).ok_or_else(|| { - anyhow::anyhow!( - "initialized {} `{}` has unsupported LLVM size/alignment", - kind_noun, - name - ) - })?; - if llvm_size == 0 { - return Ok(()); - } - return Err(anyhow::anyhow!( - "initialized {} `{}` has no stored size but lowers to {} bytes", - kind_noun, - name, - llvm_size - )); - } - if !has_explicit_layout { - return Err(anyhow::anyhow!( - "initialized {} `{}` has no rustc field-offset metadata", - kind_noun, - name - )); - } - - let slots = build_struct_slot_map(ctx, layout)?; - let llvm_fields: Vec<_> = slots - .llvm_struct_ty - .deref(ctx) - .downcast_ref::() - .expect("aggregate slot map must produce an LLVM struct") - .fields() - .collect(); - - let mut slot_offsets = Vec::with_capacity(llvm_fields.len()); - let mut current_offset = 0u64; - for llvm_field in &llvm_fields { - let (field_size, field_align) = - llvm_type_size_align(ctx, *llvm_field).ok_or_else(|| { - anyhow::anyhow!( - "initialized {} `{}` field has unsupported LLVM size/alignment", - kind_noun, - name - ) - })?; - current_offset = current_offset.div_ceil(field_align.max(1)) * field_align.max(1); - slot_offsets.push(current_offset); - current_offset += field_size; - } - - for (decl_index, slot) in slots.decl_to_llvm.iter().enumerate() { - let Some(slot) = slot else { - continue; - }; - let actual_offset = slot_offsets[*slot as usize]; - let expected_offset = layout.field_offsets[decl_index]; - if actual_offset != expected_offset { - return Err(anyhow::anyhow!( - "initialized {} `{}` field {} lowers at byte {}, but rustc placed it at byte {}; packed and overlapping field layouts are not yet supported", - kind_noun, - name, - decl_index, - actual_offset, - expected_offset - )); - } - } - - let (llvm_size, llvm_align) = - llvm_type_size_align(ctx, slots.llvm_struct_ty).ok_or_else(|| { - anyhow::anyhow!( - "initialized {} `{}` has unsupported LLVM size/alignment", - kind_noun, - name - ) - })?; - if llvm_size != layout.total_size || llvm_align > abi_align { - return Err(anyhow::anyhow!( - "initialized {} `{}` lowers to size/alignment {}/{}, but rustc requires {}/{}", - kind_noun, - name, - llvm_size, - llvm_align, - layout.total_size, - abi_align - )); - } - - Ok(()) -} - -/// Create the LLVM struct type used for slice representations. -/// -/// Slices are represented as fat pointers: `{ ptr, i64 }` where: -/// - `ptr` is a generic address space (0) pointer to the data -/// - `i64` is the number of elements (not bytes) -/// -/// # Layout -/// -/// ```text -/// struct { -/// ptr: !llvm.ptr, ; offset 0, size 8 -/// len: i64, ; offset 8, size 8 -/// } ; total size: 16 bytes -/// ``` -/// -/// # Address Space -/// -/// The pointer uses generic address space (0) because: -/// - Slices passed to kernels may point to global memory -/// - The kernel doesn't know at compile time which memory space -/// - Generic pointers can be used with any memory type -/// -/// # Usage -/// -/// This type is used for: -/// - `&[T]` slice arguments -/// - `DisjointSlice` (unique-ownership slice) arguments -/// - Any other fat pointer representation -pub(crate) fn make_slice_struct(ctx: &mut Context) -> TypeHandle { - let ptr_ty = llvm_types::PointerType::get_generic(ctx); - let len_ty = IntegerType::get(ctx, 64, Signedness::Signless); - llvm_types::StructType::get_unnamed( - ctx, - ( - vec![ptr_ty.into(), len_ty.into()], - llvm_types::StructLayout::Unpacked, - ), - ) - .into() -} - -/// The fat pointer, followed by an index space's runtime layout fields. -/// -/// With no such fields this is exactly [`make_slice_struct`], so a slice over -/// an index space fixed in its type keeps its two-field representation. -pub(crate) fn make_disjoint_slice_struct( - ctx: &mut Context, - space_tys: &[TypeHandle], -) -> anyhow::Result { - if space_tys.is_empty() { - return Ok(make_slice_struct(ctx)); - } - let ptr_ty = llvm_types::PointerType::get_generic(ctx); - let len_ty = IntegerType::get(ctx, 64, Signedness::Signless); - let mut fields: Vec = vec![ptr_ty.into(), len_ty.into()]; - for space_ty in space_tys { - fields.push(convert_type(ctx, *space_ty)?); - } - Ok( - llvm_types::StructType::get_unnamed(ctx, (fields, llvm_types::StructLayout::Unpacked)) - .into(), - ) -} - -#[cfg(test)] -mod tests { - //! Hardware-free unit tests for [`build_struct_slot_map`]: the slot map - //! and the LLVM struct type are produced by the same walk, so these - //! tests pin down both for the layout shapes from issue #128. - - use super::*; - use dialect_mir::types::{ - EnumEncoding, EnumVariant, MirArrayType, MirEnumType, MirPtrType, MirStructType, - MirTupleType, MirUnionType, StructAbiKind, - }; - - fn make_ctx() -> Context { - let mut ctx = Context::new(); - dialect_mir::register(&mut ctx); - crate::register(&mut ctx); - ctx - } - - /// A MIR-level unsigned integer type (what the importer produces). - fn mir_uint(ctx: &mut Context, width: u32) -> TypeHandle { - IntegerType::get(ctx, width, Signedness::Unsigned).into() - } - - /// A converted (signless) LLVM integer type. - fn llvm_int(ctx: &mut Context, width: u32) -> TypeHandle { - IntegerType::get(ctx, width, Signedness::Signless).into() - } - - /// `[n x i8]` padding type, as `make_padding_type` builds it. - fn pad(ctx: &mut Context, n: u64) -> TypeHandle { - make_padding_type(ctx, n) - } - - /// A zero-sized MIR struct (PhantomData shape). - fn mir_zst(ctx: &mut Context) -> TypeHandle { - MirStructType::get(ctx, "Phantom".into(), vec![], vec![]).into() - } - - fn struct_fields(ctx: &Context, ty: TypeHandle) -> Vec { - ty.deref(ctx) - .downcast_ref::() - .expect("expected an LLVM struct type") - .fields() - .collect() - } - - fn transparent_u32(ctx: &mut Context, name: &str) -> TypeHandle { - let u32_ty = mir_uint(ctx, 32); - MirStructType::get_with_full_layout_and_abi( - ctx, - name.into(), - vec!["value".into()], - vec![u32_ty], - vec![0], - vec![0], - 4, - 4, - StructAbiKind::TransparentScalar, - ) - .into() - } - - #[test] - fn transparent_scalar_return_type_uses_underlying_scalar() { - let mut ctx = make_ctx(); - let wrapper = transparent_u32(&mut ctx, "Scalar"); - let func_ty = FunctionType::get(&ctx, vec![], vec![wrapper]); - - let lowered = convert_function_type(&mut ctx, func_ty, false).unwrap(); - let result_ty = lowered.deref(&ctx).result_type(); - let result_ty_ref = result_ty.deref(&ctx); - let integer = result_ty_ref - .downcast_ref::() - .expect("transparent u32 return must lower to an integer"); - assert_eq!(integer.width(), 32); - } - - #[test] - fn nested_transparent_scalar_abi_records_each_rebuild_layer() { - let mut ctx = make_ctx(); - let inner = transparent_u32(&mut ctx, "Inner"); - let outer: TypeHandle = MirStructType::get_with_full_layout_and_abi( - &mut ctx, - "Outer".into(), - vec!["inner".into()], - vec![inner], - vec![0], - vec![0], - 4, - 4, - StructAbiKind::TransparentScalar, - ) - .into(); - - let info = transparent_scalar_abi_info(&mut ctx, outer).unwrap(); - let scalar_ty_ref = info.scalar_ty.deref(&ctx); - let integer = scalar_ty_ref - .downcast_ref::() - .expect("nested transparent wrapper must resolve to u32"); - assert_eq!(integer.width(), 32); - assert_eq!(info.layers.len(), 2); - assert_eq!(info.layers[0].field_slot, 0); - assert_eq!(info.layers[1].field_slot, 0); - } - - #[test] - fn ordinary_one_field_return_remains_aggregate() { - let mut ctx = make_ctx(); - let u32_ty = mir_uint(&mut ctx, 32); - let ordinary: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "Ordinary".into(), - vec!["value".into()], - vec![u32_ty], - vec![0], - vec![0], - 4, - 4, - ) - .into(); - let func_ty = FunctionType::get(&ctx, vec![], vec![ordinary]); - - let lowered = convert_function_type(&mut ctx, func_ty, false).unwrap(); - assert!( - lowered - .deref(&ctx) - .result_type() - .deref(&ctx) - .is::(), - "ordinary one-field structs must not be scalarized" - ); - } - - #[test] - fn mir_abi_alignment_recurses_through_nested_arrays() { - let mut ctx = make_ctx(); - let byte = mir_uint(&mut ctx, 8); - let marker: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "Align32".into(), - vec![], - vec![], - vec![], - vec![], - 0, - 32, - ) - .into(); - let tuple: TypeHandle = MirTupleType::get_with_layout( - &mut ctx, - vec![marker, byte], - vec![0, 1], - vec![0, 0], - 32, - 32, - ) - .into(); - let inner: TypeHandle = MirArrayType::get(&mut ctx, tuple, 2).into(); - let outer: TypeHandle = MirArrayType::get(&mut ctx, inner, 3).into(); - let plain: TypeHandle = MirArrayType::get(&mut ctx, byte, 4).into(); - - assert_eq!(mir_type_abi_align(&ctx, tuple), Some(32)); - assert_eq!(mir_type_abi_align(&ctx, inner), Some(32)); - assert_eq!(mir_type_abi_align(&ctx, outer), Some(32)); - assert_eq!(mir_type_abi_align(&ctx, plain), None); - } - - #[test] - fn union_storage_has_exact_size_alignment_and_stride() { - let mut ctx = make_ctx(); - let u8_ty = mir_uint(&mut ctx, 8); - let u32_ty = mir_uint(&mut ctx, 32); - let bytes_ty: TypeHandle = MirArrayType::get(&mut ctx, u8_ty, 4).into(); - let union_ty = MirUnionType::get( - &mut ctx, - "Bits".into(), - vec!["word".into(), "bytes".into()], - vec![u32_ty, bytes_ty], - 4, - 4, - ); - let union_data = union_ty.deref(&ctx).clone(); - let storage = build_union_storage_type(&mut ctx, &union_data).unwrap(); - assert_eq!(llvm_type_size_align(&ctx, storage), Some((4, 4))); - - let union_handle: TypeHandle = union_ty.into(); - let array: TypeHandle = MirArrayType::get(&mut ctx, union_handle, 3).into(); - let llvm_array = convert_type(&mut ctx, array).unwrap(); - assert_eq!(llvm_type_size_align(&ctx, llvm_array), Some((12, 4))); - } - - #[test] - fn fixed_vector_layout_uses_packed_bit_width_and_rejects_unknown_widths() { - let mut ctx = make_ctx(); - let i1 = llvm_int(&mut ctx, 1); - let v16i1: TypeHandle = - llvm_types::VectorType::get(&ctx, i1, 16, llvm_types::VectorTypeKind::Fixed).into(); - assert_eq!( - llvm_type_size_align(&ctx, v16i1), - Some((2, 2)), - "<16 x i1> is 16 packed bits, not sixteen bytes" - ); - - let i8 = llvm_int(&mut ctx, 8); - let v3i8: TypeHandle = - llvm_types::VectorType::get(&ctx, i8, 3, llvm_types::VectorTypeKind::Fixed).into(); - assert_eq!( - llvm_type_size_align(&ctx, v3i8), - None, - "the NVPTX data layout does not define a 24-bit vector alignment" - ); - } - - #[test] - fn enum_array_payload_keeps_exact_size_alignment_and_stride() { - let mut ctx = make_ctx(); - let tag = mir_uint(&mut ctx, 8); - let element = mir_uint(&mut ctx, 16); - let payload: TypeHandle = MirArrayType::get(&mut ctx, element, 3).into(); - let enum_ty: TypeHandle = MirEnumType::get_with_layout( - &mut ctx, - "ArrayPayload".into(), - tag, - vec![0, 1], - vec![ - EnumVariant::unit("Empty".into()), - EnumVariant::new_with_layout("Data".into(), vec![payload], vec![2], vec![6]), - ], - 0, - 8, - 2, - ) - .into(); - let map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); - assert_eq!(llvm_type_size_align(&ctx, map.llvm_struct_ty), Some((8, 2))); - - let array: TypeHandle = MirArrayType::get(&mut ctx, enum_ty, 5).into(); - let lowered = convert_type(&mut ctx, array).unwrap(); - assert_eq!(llvm_type_size_align(&ctx, lowered), Some((40, 2))); - } - - #[test] - fn enum_slot_map_rejects_partial_byte_integer_carriers() { - let mut ctx = make_ctx(); - let partial_discriminant = mir_uint(&mut ctx, 7); - let direct: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "PartialDirect".into(), - partial_discriminant, - vec![0, 1], - vec![EnumVariant::unit("A".into()), EnumVariant::unit("B".into())], - EnumEncoding { - tag_offset: 0, - total_size: 1, - abi_align: 1, - layout_kind: EnumLayoutKind::Direct, - carrier_kind: EnumCarrierKind::Integer, - carrier_width: 7, - variant_inhabited: vec![1, 1], - ..EnumEncoding::default() - }, - ) - .into(); - let error = build_enum_slot_map(&mut ctx, direct) - .err() - .expect("partial-byte Direct carrier must fail closed"); - assert!( - error.to_string().contains("whole number of bytes"), - "{error}" - ); - - let logical = mir_uint(&mut ctx, 8); - let payload = mir_uint(&mut ctx, 8); - let niche: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "PartialNiche".into(), - logical, - vec![0, 1], - vec![ - EnumVariant::unit("None".into()), - EnumVariant::new_with_layout("Some".into(), vec![payload], vec![0], vec![1]), - ], - EnumEncoding { - tag_offset: 0, - total_size: 1, - abi_align: 1, - layout_kind: EnumLayoutKind::Niche, - carrier_kind: EnumCarrierKind::Integer, - carrier_width: 7, - untagged_variant: 1, - variant_inhabited: vec![1, 1], - ..EnumEncoding::default() - }, - ) - .into(); - let error = build_enum_slot_map(&mut ctx, niche) - .err() - .expect("partial-byte Niche carrier must fail closed"); - assert!( - error.to_string().contains("whole number of bytes"), - "{error}" - ); - } - - #[test] - fn enum_slot_map_rejects_misaligned_inhabited_payload() { - let mut ctx = make_ctx(); - let tag = mir_uint(&mut ctx, 8); - let word = mir_uint(&mut ctx, 32); - let enum_ty: TypeHandle = MirEnumType::get_with_layout( - &mut ctx, - "PackedPayload".into(), - tag, - vec![0, 1], - vec![ - EnumVariant::unit("Empty".into()), - EnumVariant::new_with_layout("Data".into(), vec![word], vec![1], vec![4]), - ], - 0, - 8, - 4, - ) - .into(); - let error = build_enum_slot_map(&mut ctx, enum_ty) - .err() - .expect("misaligned payload must be rejected"); - assert!(error.to_string().contains("offset 1 is not aligned")); - } - - #[test] - fn enum_slot_map_uses_i8_storage_for_nonoverlapping_scalar_bool() { - let mut ctx = make_ctx(); - let tag = mir_uint(&mut ctx, 32); - let boolean: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Signless).into(); - let enum_ty: TypeHandle = MirEnumType::get_with_layout( - &mut ctx, - "DirectBool".into(), - tag, - vec![0, 1], - vec![ - EnumVariant::new_with_layout("A".into(), vec![boolean], vec![4], vec![1]), - EnumVariant::unit("B".into()), - ], - 0, - 8, - 4, - ) - .into(); - - let map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); - assert_eq!(map.field_slots, vec![None]); - let fields = struct_fields(&ctx, map.llvm_struct_ty); - assert_eq!(fields.len(), 3, "tag, canonical bool byte, tail pad"); - assert_eq!( - fields[1] - .deref(&ctx) - .downcast_ref::() - .map(IntegerType::width), - Some(8), - "Rust bool storage must be an i8 byte, never an LLVM i1 slot" - ); - } - - #[test] - fn enum_slot_map_allows_nonoverlapping_aggregate_padding_without_bool() { - let mut ctx = make_ctx(); - let tag = mir_uint(&mut ctx, 32); - let byte = mir_uint(&mut ctx, 8); - let word = mir_uint(&mut ctx, 32); - let padded: TypeHandle = MirTupleType::get(&mut ctx, vec![byte, word]).into(); - let lowered_padded = convert_type(&mut ctx, padded).unwrap(); - assert!( - !llvm_type_is_byte_faithful(&ctx, lowered_padded), - "the LLVM i8/i32 tuple has harmless implicit padding" - ); - let enum_ty: TypeHandle = MirEnumType::get_with_layout( - &mut ctx, - "PaddedPayload".into(), - tag, - vec![0, 1], - vec![ - EnumVariant::new_with_layout("Data".into(), vec![padded], vec![4], vec![8]), - EnumVariant::unit("Empty".into()), - ], - 0, - 12, - 4, - ) - .into(); - - let map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); - assert_eq!(map.field_slots, vec![Some(1)]); - assert_eq!( - llvm_type_size_align(&ctx, map.llvm_struct_ty), - Some((12, 4)) - ); - } - - #[test] - fn enum_slot_map_rejects_nonoverlapping_nested_bool_storage() { - let mut ctx = make_ctx(); - let tag = mir_uint(&mut ctx, 32); - let boolean: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Signless).into(); - let wrapper: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "BoolWrapper".into(), - vec!["value".into()], - vec![boolean], - vec![0], - vec![0], - 1, - 1, - ) - .into(); - let enum_ty: TypeHandle = MirEnumType::get_with_layout( - &mut ctx, - "DirectBoolWrapper".into(), - tag, - vec![0, 1], - vec![ - EnumVariant::new_with_layout("A".into(), vec![wrapper], vec![4], vec![1]), - EnumVariant::unit("B".into()), - ], - 0, - 8, - 4, - ) - .into(); - - // The wrapper claims its byte-faithful twin ({i8}) and stays - // slotless: construction canonicalizes the bool byte, extraction - // re-loads the original {i1} shape from the canonical byte. - let map = build_enum_slot_map(&mut ctx, enum_ty) - .expect("nested bool storage canonicalizes through its byte-faithful twin"); - assert_eq!(map.field_slots, vec![None]); - assert_eq!(llvm_type_size_align(&ctx, map.llvm_struct_ty), Some((8, 4))); - let struct_fields: Vec<_> = map - .llvm_struct_ty - .deref(&ctx) - .downcast_ref::() - .unwrap() - .fields() - .collect(); - assert!( - struct_fields - .iter() - .all(|field| !llvm_type_contains_i1(&ctx, *field)), - "enum storage must never contain physical i1" - ); - } - - #[test] - fn enum_slot_map_rejects_bool_hidden_by_union_byte_storage() { - let mut ctx = make_ctx(); - let tag = mir_uint(&mut ctx, 32); - let byte = mir_uint(&mut ctx, 8); - let boolean: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Signless).into(); - let union: TypeHandle = MirUnionType::get( - &mut ctx, - "BoolOrByte".into(), - vec!["flag".into(), "byte".into()], - vec![boolean, byte], - 1, - 1, - ) - .into(); - let lowered_union = convert_type(&mut ctx, union).unwrap(); - assert!( - !llvm_type_contains_i1(&ctx, lowered_union), - "the union's raw-byte carrier intentionally hides its semantic bool" - ); - let enum_ty: TypeHandle = MirEnumType::get_with_layout( - &mut ctx, - "DirectUnionBool".into(), - tag, - vec![0, 1], - vec![ - EnumVariant::new_with_layout("A".into(), vec![union], vec![4], vec![1]), - EnumVariant::unit("B".into()), - ], - 0, - 8, - 4, - ) - .into(); - - let error = build_enum_slot_map(&mut ctx, enum_ty) - .err() - .expect("a bool hidden by union byte storage must still fail closed"); - assert!( - error.to_string().contains("hidden behind a union"), - "{error}" - ); - } - - #[test] - fn enum_slot_map_does_not_descend_through_pointer_to_bool() { - let mut ctx = make_ctx(); - let tag = mir_uint(&mut ctx, 32); - let boolean: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Signless).into(); - let pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, boolean, false).into(); - let enum_ty: TypeHandle = MirEnumType::get_with_layout( - &mut ctx, - "PointerToBool".into(), - tag, - vec![0, 1], - vec![ - EnumVariant::new_with_layout("A".into(), vec![pointer], vec![8], vec![8]), - EnumVariant::unit("B".into()), - ], - 0, - 16, - 8, - ) - .into(); - - let map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); - assert!(map.field_slots[0].is_some()); - } - - #[test] - fn enum_slot_map_rejects_pointer_nonpointer_overlap_in_either_variant_order() { - for pointer_first in [false, true] { - let mut ctx = make_ctx(); - let discr = mir_uint(&mut ctx, 32); - let u8_ty = mir_uint(&mut ctx, 8); - let u64_ty = mir_uint(&mut ctx, 64); - let pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, u8_ty, false).into(); - let (first_name, first_ty, second_name, second_ty) = if pointer_first { - ("Ptr", pointer, "Bits", u64_ty) - } else { - ("Bits", u64_ty, "Ptr", pointer) - }; - let enum_ty: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "PointerOrBits".into(), - discr, - vec![0, 1], - vec![ - EnumVariant::new_with_layout( - first_name.into(), - vec![first_ty], - vec![8], - vec![8], - ), - EnumVariant::new_with_layout( - second_name.into(), - vec![second_ty], - vec![8], - vec![8], - ), - ], - EnumEncoding { - tag_offset: 0, - total_size: 16, - abi_align: 8, - layout_kind: EnumLayoutKind::Direct, - carrier_kind: EnumCarrierKind::Integer, - carrier_width: 32, - variant_inhabited: vec![1, 1], - ..EnumEncoding::default() - }, - ) - .into(); - let error = build_enum_slot_map(&mut ctx, enum_ty) - .err() - .expect("pointer overlap must reject"); - assert!( - error.to_string().contains("pointer provenance"), - "pointer_first={pointer_first}: {error}" - ); - } - } - - #[test] - fn enum_slot_map_canonicalizes_overlapping_bool_wrapper_storage() { - let mut ctx = make_ctx(); - let logical = mir_uint(&mut ctx, 8); - let boolean: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Signless).into(); - let wrapper: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "BoolWrapper".into(), - vec!["value".into()], - vec![boolean], - vec![0], - vec![0], - 1, - 1, - ) - .into(); - let enum_ty: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "MaybeBoolWrapper".into(), - logical, - vec![0, 1], - vec![ - EnumVariant::unit("None".into()), - EnumVariant::new_with_layout("Some".into(), vec![wrapper], vec![0], vec![1]), - ], - EnumEncoding { - tag_offset: 0, - total_size: 1, - abi_align: 1, - layout_kind: EnumLayoutKind::Niche, - carrier_kind: EnumCarrierKind::Integer, - carrier_width: 8, - niche_start: 2, - untagged_variant: 1, - variant_inhabited: vec![1, 1], - ..EnumEncoding::default() - }, - ) - .into(); - - // Option: the wrapper's byte-faithful twin ({i8}) - // shares the single canonical byte with the i8 niche carrier. The - // wrapper stays slotless; construction zero-extends its bool. - let map = build_enum_slot_map(&mut ctx, enum_ty) - .expect("a bool wrapper shares its canonical byte with the i8 carrier"); - assert_eq!(map.field_slots, vec![None]); - assert_eq!(llvm_type_size_align(&ctx, map.llvm_struct_ty), Some((1, 1))); - assert!( - !llvm_type_contains_i1(&ctx, map.llvm_struct_ty), - "enum storage must never contain physical i1" - ); - } - - #[test] - fn enum_slot_map_canonicalizes_bool_wrapper_overlap_in_either_variant_order() { - for wrapper_first in [false, true] { - let mut ctx = make_ctx(); - let tag = mir_uint(&mut ctx, 32); - let byte = mir_uint(&mut ctx, 8); - let boolean: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Signless).into(); - let wrapper: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "BoolWrapper".into(), - vec!["value".into()], - vec![boolean], - vec![0], - vec![0], - 1, - 1, - ) - .into(); - let (first, second) = if wrapper_first { - (wrapper, byte) - } else { - (byte, wrapper) - }; - let enum_ty: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "BoolWrapperOrByte".into(), - tag, - vec![0, 1], - vec![ - EnumVariant::new_with_layout("First".into(), vec![first], vec![4], vec![1]), - EnumVariant::new_with_layout("Second".into(), vec![second], vec![4], vec![1]), - ], - EnumEncoding { - tag_offset: 0, - total_size: 8, - abi_align: 4, - layout_kind: EnumLayoutKind::Direct, - carrier_kind: EnumCarrierKind::Integer, - carrier_width: 32, - variant_inhabited: vec![1, 1], - ..EnumEncoding::default() - }, - ) - .into(); - - // The wrapper's canonical twin ({i8}) and the plain u8 variant - // share one byte-faithful byte, independent of declaration - // order; the wrapper stays slotless either way. - let map = build_enum_slot_map(&mut ctx, enum_ty) - .expect("canonical bool bytes may share storage with a u8 variant"); - assert_eq!(llvm_type_size_align(&ctx, map.llvm_struct_ty), Some((8, 4))); - assert!( - !llvm_type_contains_i1(&ctx, map.llvm_struct_ty), - "wrapper_first={wrapper_first}: enum storage must never contain physical i1" - ); - let wrapper_flat = if wrapper_first { 0 } else { 1 }; - assert_eq!(map.field_slots[wrapper_flat], None); - } - } - - #[test] - fn enum_slot_map_backs_pointer_beside_integer_niche_carrier() { - let mut ctx = make_ctx(); - let logical = mir_uint(&mut ctx, 8); - let u32_ty = mir_uint(&mut ctx, 32); - let pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, u32_ty, false).into(); - let payload: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "PointerThenNiche".into(), - vec!["pointer".into(), "niche".into()], - vec![pointer, u32_ty], - vec![0, 1], - vec![0, 8], - 16, - 8, - ) - .into(); - let enum_ty: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "MaybePointerThenNiche".into(), - logical, - vec![0, 1], - vec![ - EnumVariant::unit("None".into()), - EnumVariant::new_with_layout("Some".into(), vec![payload], vec![0], vec![16]), - ], - EnumEncoding { - tag_offset: 8, - total_size: 16, - abi_align: 8, - layout_kind: EnumLayoutKind::Niche, - carrier_kind: EnumCarrierKind::Integer, - carrier_width: 32, - untagged_variant: 1, - variant_inhabited: vec![1, 1], - ..EnumEncoding::default() - }, - ) - .into(); - // The pointer at byte 0 never shares bytes with the integer niche - // carrier at byte 8, so it is representable: it gets its own `ptr` slot - // and the carrier stays an integer slot. `{ptr@0, i32@8, i32 filler}` - // -- the trailing 4 bytes are 4-aligned inside an 8-aligned enum, so - // they lower to one `i32` rather than `[4 x i8]`. - let map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); - assert_eq!(map.field_slots, vec![None]); - let lowered_pointer = convert_type(&mut ctx, pointer).unwrap(); - let carrier: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); - assert_eq!( - struct_fields(&ctx, map.llvm_struct_ty), - vec![lowered_pointer, carrier, llvm_int(&mut ctx, 32)] - ); - assert_eq!( - llvm_type_size_align(&ctx, map.llvm_struct_ty), - Some((16, 8)) - ); - } - - #[test] - fn enum_slot_map_keeps_plain_pointer_niche_in_one_pointer_slot() { - let mut ctx = make_ctx(); - let logical = mir_uint(&mut ctx, 8); - let u32_ty = mir_uint(&mut ctx, 32); - let pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, u32_ty, false).into(); - let enum_ty: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "MaybePointer".into(), - logical, - vec![0, 1], - vec![ - EnumVariant::unit("None".into()), - EnumVariant::new_with_layout("Some".into(), vec![pointer], vec![0], vec![8]), - ], - EnumEncoding { - tag_offset: 0, - total_size: 8, - abi_align: 8, - layout_kind: EnumLayoutKind::Niche, - carrier_kind: EnumCarrierKind::Pointer, - carrier_width: 64, - untagged_variant: 1, - variant_inhabited: vec![1, 1], - ..EnumEncoding::default() - }, - ) - .into(); - let map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); - assert_eq!(map.carrier_slot, Some(0)); - assert_eq!(map.field_slots, vec![Some(0)]); - assert!( - struct_fields(&ctx, map.llvm_struct_ty)[0] - .deref(&ctx) - .is::() - ); - } - - #[test] - fn enum_slot_map_accepts_pointer_first_and_pointer_second_aggregate_niches() { - for pointer_first in [true, false] { - let mut ctx = make_ctx(); - let logical = mir_uint(&mut ctx, 64); - let index = mir_uint(&mut ctx, 64); - let pointee = mir_uint(&mut ctx, 32); - let pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, pointee, false).into(); - let payload_types = if pointer_first { - vec![pointer, index] - } else { - vec![index, pointer] - }; - let payload: TypeHandle = MirTupleType::get(&mut ctx, payload_types).into(); - let tag_offset = if pointer_first { 0 } else { 8 }; - let enum_ty: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "MaybeIndexedRef".into(), - logical, - vec![0, 1], - vec![ - EnumVariant::unit("None".into()), - EnumVariant::new_with_layout("Some".into(), vec![payload], vec![0], vec![16]), - ], - EnumEncoding { - tag_offset, - total_size: 16, - abi_align: 8, - layout_kind: EnumLayoutKind::Niche, - carrier_kind: EnumCarrierKind::Pointer, - carrier_width: 64, - untagged_variant: 1, - variant_inhabited: vec![1, 1], - ..EnumEncoding::default() - }, - ) - .into(); - - let map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); - assert_eq!(map.field_slots, vec![None]); - assert_eq!(map.carrier_slot, Some(if pointer_first { 0 } else { 1 })); - assert_eq!( - llvm_type_size_align(&ctx, map.llvm_struct_ty), - Some((16, 8)) - ); - let lowered_pointer = convert_type(&mut ctx, pointer).unwrap(); - // The 8 bytes the pointer slot does not claim are 8-aligned inside - // an 8-aligned enum, so they lower to one `i64`, not `[8 x i8]`. - let filler = llvm_int(&mut ctx, 64); - let expected = if pointer_first { - vec![lowered_pointer, filler] - } else { - vec![filler, lowered_pointer] - }; - assert_eq!(struct_fields(&ctx, map.llvm_struct_ty), expected); - } - } - - #[test] - fn enum_slot_map_backs_each_aggregate_pointer_leaf_with_its_own_slot() { - let mut ctx = make_ctx(); - let logical = mir_uint(&mut ctx, 64); - let pointee = mir_uint(&mut ctx, 32); - let pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, pointee, false).into(); - let payload: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "TwoRefs".into(), - vec!["first".into(), "second".into()], - vec![pointer, pointer], - vec![0, 1], - vec![0, 8], - 16, - 8, - ) - .into(); - let enum_ty: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "MaybeTwoRefs".into(), - logical, - vec![0, 1], - vec![ - EnumVariant::unit("None".into()), - EnumVariant::new_with_layout("Some".into(), vec![payload], vec![0], vec![16]), - ], - EnumEncoding { - tag_offset: 8, - total_size: 16, - abi_align: 8, - layout_kind: EnumLayoutKind::Niche, - carrier_kind: EnumCarrierKind::Pointer, - carrier_width: 64, - untagged_variant: 1, - variant_inhabited: vec![1, 1], - ..EnumEncoding::default() - }, - ) - .into(); - - // Both pointer leaves are representable without erasing provenance: the - // carrier backs the leaf at byte 8, and the leaf at byte 0 gets its own - // fresh `ptr` slot. The payload stays slotless and round-trips through - // `{ptr@0, ptr@8}`. - let map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); - assert_eq!(map.field_slots, vec![None]); - let lowered_pointer = convert_type(&mut ctx, pointer).unwrap(); - assert_eq!( - struct_fields(&ctx, map.llvm_struct_ty), - vec![lowered_pointer, lowered_pointer] - ); - assert_eq!( - llvm_type_size_align(&ctx, map.llvm_struct_ty), - Some((16, 8)) - ); - } - - #[test] - fn enum_slot_map_backs_split_at_mut_slice_pair() { - // Models `Option<(&mut [u32], &mut [u32])>`, the `split_at_mut_checked` - // return type: two fat slice pointers `{ptr, len}` back to back. The - // None niche lives in the first data pointer; both data pointers must - // keep their own `ptr` slot so provenance survives the memory - // round-trip, with the two `len` fields as raw byte storage. - let mut ctx = make_ctx(); - let logical = mir_uint(&mut ctx, 64); - let len = mir_uint(&mut ctx, 64); - let pointee = mir_uint(&mut ctx, 32); - let pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, pointee, false).into(); - let payload: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "SlicePair".into(), - vec![ - "a_ptr".into(), - "a_len".into(), - "b_ptr".into(), - "b_len".into(), - ], - vec![pointer, len, pointer, len], - vec![0, 1, 2, 3], - vec![0, 8, 16, 24], - 32, - 8, - ) - .into(); - let enum_ty: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "MaybeSlicePair".into(), - logical, - vec![0, 1], - vec![ - EnumVariant::unit("None".into()), - EnumVariant::new_with_layout("Some".into(), vec![payload], vec![0], vec![32]), - ], - EnumEncoding { - tag_offset: 0, - total_size: 32, - abi_align: 8, - layout_kind: EnumLayoutKind::Niche, - carrier_kind: EnumCarrierKind::Pointer, - carrier_width: 64, - untagged_variant: 1, - variant_inhabited: vec![1, 1], - ..EnumEncoding::default() - }, - ) - .into(); - - let map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); - assert_eq!(map.field_slots, vec![None]); - let lowered_pointer = convert_type(&mut ctx, pointer).unwrap(); - // Each slice's length sits in the 8 bytes after its pointer, 8-aligned - // inside an 8-aligned enum, so both lower to one `i64`. This is the - // `split_at_mut_checked` shape, and the whole point of the widening: - // `{ptr, i64, ptr, i64}` moves two lengths as two values, where - // `{ptr, [8 x i8], ptr, [8 x i8]}` moved them as sixteen separate bytes. - let filler = llvm_int(&mut ctx, 64); - assert_eq!( - struct_fields(&ctx, map.llvm_struct_ty), - vec![lowered_pointer, filler, lowered_pointer, filler] - ); - assert_eq!( - llvm_type_size_align(&ctx, map.llvm_struct_ty), - Some((32, 8)) - ); - } - - #[test] - fn enum_slot_map_three_field_tuple_follows_recorded_offsets() { - // rustc lays out `(u32, f32, &T)` with the pointer first in memory: - // ptr @ 0, u32 @ 8, f32 @ 12, size 16. With those offsets recorded on - // the tuple, the pointer leaf coincides exactly with the niche - // carrier at byte 0 and the payload is accepted. - let mut ctx = make_ctx(); - let logical = mir_uint(&mut ctx, 64); - let word = mir_uint(&mut ctx, 32); - let float: TypeHandle = FP32Type::get(&ctx).into(); - let pointee = mir_uint(&mut ctx, 32); - let pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, pointee, false).into(); - let payload: TypeHandle = MirTupleType::get_with_layout( - &mut ctx, - vec![word, float, pointer], - vec![2, 0, 1], - vec![8, 12, 0], - 16, - 8, - ) - .into(); - let enum_ty: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "MaybeMixedTuple".into(), - logical, - vec![0, 1], - vec![ - EnumVariant::unit("None".into()), - EnumVariant::new_with_layout("Some".into(), vec![payload], vec![0], vec![16]), - ], - EnumEncoding { - tag_offset: 0, - total_size: 16, - abi_align: 8, - layout_kind: EnumLayoutKind::Niche, - carrier_kind: EnumCarrierKind::Pointer, - carrier_width: 64, - untagged_variant: 1, - variant_inhabited: vec![1, 1], - ..EnumEncoding::default() - }, - ) - .into(); - - let map = build_enum_slot_map(&mut ctx, enum_ty) - .expect("recorded tuple offsets prove the pointer word matches the carrier"); - assert_eq!(map.carrier_slot, Some(0)); - assert_eq!( - llvm_type_size_align(&ctx, map.llvm_struct_ty), - Some((16, 8)) - ); - - // The same tuple with offsets that put the u32 over the pointer - // carrier (ptr @ 8 instead) must still fail closed: integer bytes - // may not alias pointer storage. - let mismatched_payload: TypeHandle = MirTupleType::get_with_layout( - &mut ctx, - vec![word, float, pointer], - vec![0, 1, 2], - vec![0, 4, 8], - 16, - 8, - ) - .into(); - let enum_ty: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "MaybeMixedTupleShifted".into(), - logical, - vec![0, 1], - vec![ - EnumVariant::unit("None".into()), - EnumVariant::new_with_layout( - "Some".into(), - vec![mismatched_payload], - vec![0], - vec![16], - ), - ], - EnumEncoding { - tag_offset: 0, - total_size: 16, - abi_align: 8, - layout_kind: EnumLayoutKind::Niche, - carrier_kind: EnumCarrierKind::Pointer, - carrier_width: 64, - untagged_variant: 1, - variant_inhabited: vec![1, 1], - ..EnumEncoding::default() - }, - ) - .into(); - let error = build_enum_slot_map(&mut ctx, enum_ty) - .err() - .expect("a non-pointer word over the pointer carrier must fail closed"); - assert!(error.to_string().contains("pointer provenance"), "{error}"); - } - - #[test] - fn enum_slot_map_rejects_shifted_nested_pointer_carrier() { - let mut ctx = make_ctx(); - let logical = mir_uint(&mut ctx, 64); - let index = mir_uint(&mut ctx, 64); - let pointee = mir_uint(&mut ctx, 32); - let pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, pointee, false).into(); - let payload: TypeHandle = MirTupleType::get(&mut ctx, vec![index, pointer]).into(); - let enum_ty: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "ShiftedPointerCarrier".into(), - logical, - vec![0, 1], - vec![ - EnumVariant::unit("None".into()), - EnumVariant::new_with_layout("Some".into(), vec![payload], vec![0], vec![16]), - ], - EnumEncoding { - tag_offset: 0, - total_size: 16, - abi_align: 8, - layout_kind: EnumLayoutKind::Niche, - carrier_kind: EnumCarrierKind::Pointer, - carrier_width: 64, - untagged_variant: 1, - variant_inhabited: vec![1, 1], - ..EnumEncoding::default() - }, - ) - .into(); - - let error = build_enum_slot_map(&mut ctx, enum_ty) - .err() - .expect("pointer leaves at different offsets must not be conflated"); - assert!(error.to_string().contains("pointer provenance"), "{error}"); - } - - #[test] - fn enum_slot_map_rejects_nested_carrier_address_space_mismatch() { - let mut ctx = make_ctx(); - let logical = mir_uint(&mut ctx, 64); - let index = mir_uint(&mut ctx, 64); - let pointee = mir_uint(&mut ctx, 32); - let generic_pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, pointee, false).into(); - let payload: TypeHandle = MirTupleType::get(&mut ctx, vec![generic_pointer, index]).into(); - let enum_ty: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "MismatchedPointerSpace".into(), - logical, - vec![0, 1], - vec![ - EnumVariant::unit("None".into()), - EnumVariant::new_with_layout("Some".into(), vec![payload], vec![0], vec![16]), - ], - EnumEncoding { - tag_offset: 0, - total_size: 16, - abi_align: 8, - layout_kind: EnumLayoutKind::Niche, - carrier_kind: EnumCarrierKind::Pointer, - carrier_width: 64, - carrier_address_space: llvm_types::address_space::GLOBAL, - untagged_variant: 1, - variant_inhabited: vec![1, 1], - ..EnumEncoding::default() - }, - ) - .into(); - - let error = build_enum_slot_map(&mut ctx, enum_ty) - .err() - .expect("equal byte ranges in different address spaces must not alias"); - assert!(error.to_string().contains("pointer provenance"), "{error}"); - } - - #[test] - fn enum_slot_map_genericizes_nonoverlapping_shared_pointer_payload() { - let mut ctx = make_ctx(); - let discr = mir_uint(&mut ctx, 32); - let u32_ty = mir_uint(&mut ctx, 32); - let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, u32_ty, false).into(); - let enum_ty: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "SharedPointerPayload".into(), - discr, - vec![0, 1], - vec![ - EnumVariant::unit("Unit".into()), - EnumVariant::new_with_layout("Ptr".into(), vec![shared], vec![8], vec![8]), - ], - EnumEncoding { - tag_offset: 0, - total_size: 16, - abi_align: 8, - layout_kind: EnumLayoutKind::Direct, - carrier_kind: EnumCarrierKind::Integer, - carrier_width: 32, - variant_inhabited: vec![1, 1], - ..EnumEncoding::default() - }, - ) - .into(); - let map = build_enum_slot_map(&mut ctx, enum_ty) - .expect("a direct shared-pointer payload should use generic physical storage"); - let field_slot = map.field_slots[0].expect("shared pointer field should own a slot"); - let stored_ty = map - .llvm_struct_ty - .deref(&ctx) - .downcast_ref::() - .map(|struct_ty| struct_ty.field_type(field_slot as usize)) - .expect("field slot must exist"); - assert_eq!( - stored_ty - .deref(&ctx) - .downcast_ref::() - .expect("shared pointer storage must remain a pointer") - .address_space(), - llvm_types::address_space::GENERIC, - "enum storage must use a target-stable generic pointer" - ); - assert_eq!( - map.field_llvm_types[0] - .deref(&ctx) - .downcast_ref::() - .expect("semantic field must remain a pointer") - .address_space(), - llvm_types::address_space::SHARED, - "payload operations must retain the semantic shared address space" - ); - } - - #[test] - fn enum_slot_map_genericizes_nested_shared_pointer_payload() { - let mut ctx = make_ctx(); - let discr = mir_uint(&mut ctx, 32); - let u32_ty = mir_uint(&mut ctx, 32); - let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, u32_ty, false).into(); - let wrapper: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "SharedPointerWrapper".into(), - vec!["pointer".into()], - vec![shared], - vec![0], - vec![0], - 8, - 8, - ) - .into(); - let enum_ty: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "NestedSharedPointerPayload".into(), - discr, - vec![0, 1], - vec![ - EnumVariant::unit("Unit".into()), - EnumVariant::new_with_layout("Ptr".into(), vec![wrapper], vec![8], vec![8]), - ], - EnumEncoding { - tag_offset: 0, - total_size: 16, - abi_align: 8, - layout_kind: EnumLayoutKind::Direct, - carrier_kind: EnumCarrierKind::Integer, - carrier_width: 32, - variant_inhabited: vec![1, 1], - ..EnumEncoding::default() - }, - ) - .into(); - - let map = build_enum_slot_map(&mut ctx, enum_ty) - .expect("a struct-nested shared pointer should use generic physical storage"); - let field_slot = map.field_slots[0].expect("nested payload should own a slot"); - let stored_ty = map - .llvm_struct_ty - .deref(&ctx) - .downcast_ref::() - .map(|struct_ty| struct_ty.field_type(field_slot as usize)) - .expect("field slot must exist"); - assert!( - llvm_type_contains_pointer_in_address_space( - &ctx, - stored_ty, - llvm_types::address_space::GENERIC - ), - "physical nested payload must contain a generic pointer" - ); - assert!( - !llvm_type_contains_pointer_in_address_space( - &ctx, - stored_ty, - llvm_types::address_space::SHARED - ), - "physical nested payload must not retain the target-dependent AS3 pointer" - ); - assert!( - llvm_type_contains_pointer_in_address_space( - &ctx, - map.field_llvm_types[0], - llvm_types::address_space::SHARED - ), - "semantic payload type must retain shared address-space semantics" - ); - } - - #[test] - fn enum_slot_map_genericizes_bounded_shared_pointer_array_payload() { - let mut ctx = make_ctx(); - let discr = mir_uint(&mut ctx, 32); - let u32_ty = mir_uint(&mut ctx, 32); - let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, u32_ty, false).into(); - let pointers: TypeHandle = MirArrayType::get(&mut ctx, shared, 2).into(); - let enum_ty: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "SharedPointerArrayPayload".into(), - discr, - vec![0, 1], - vec![ - EnumVariant::unit("Unit".into()), - EnumVariant::new_with_layout("Pointers".into(), vec![pointers], vec![8], vec![16]), - ], - EnumEncoding { - tag_offset: 0, - total_size: 24, - abi_align: 8, - layout_kind: EnumLayoutKind::Direct, - carrier_kind: EnumCarrierKind::Integer, - carrier_width: 32, - variant_inhabited: vec![1, 1], - ..EnumEncoding::default() - }, - ) - .into(); - - let map = build_enum_slot_map(&mut ctx, enum_ty) - .expect("a bounded shared-pointer array should use generic physical storage"); - let field_slot = map.field_slots[0].expect("bounded pointer array should own a slot"); - let stored_ty = map - .llvm_struct_ty - .deref(&ctx) - .downcast_ref::() - .map(|struct_ty| struct_ty.field_type(field_slot as usize)) - .expect("field slot must exist"); - assert!( - llvm_type_contains_pointer_in_address_space( - &ctx, - stored_ty, - llvm_types::address_space::GENERIC - ), - "physical array payload must contain generic pointers" - ); - assert!( - !llvm_type_contains_pointer_in_address_space( - &ctx, - stored_ty, - llvm_types::address_space::SHARED - ), - "physical array payload must not retain target-dependent AS3 pointers" - ); - assert!( - llvm_type_contains_pointer_in_address_space( - &ctx, - map.field_llvm_types[0], - llvm_types::address_space::SHARED - ), - "semantic array payload must retain shared address-space semantics" - ); - } - - #[test] - fn enum_slot_map_rejects_oversized_shared_pointer_array_payload() { - let mut ctx = make_ctx(); - let discr = mir_uint(&mut ctx, 32); - let u32_ty = mir_uint(&mut ctx, 32); - let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, u32_ty, false).into(); - let count = MAX_ENUM_PAYLOAD_ARRAY_REWRITE_LEAVES + 1; - let pointers: TypeHandle = MirArrayType::get(&mut ctx, shared, count).into(); - let payload_size = count * 8; - let enum_ty: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "OversizedSharedPointerArrayPayload".into(), - discr, - vec![0, 1], - vec![ - EnumVariant::unit("Unit".into()), - EnumVariant::new_with_layout( - "Pointers".into(), - vec![pointers], - vec![8], - vec![payload_size], - ), - ], - EnumEncoding { - tag_offset: 0, - total_size: 8 + payload_size, - abi_align: 8, - layout_kind: EnumLayoutKind::Direct, - carrier_kind: EnumCarrierKind::Integer, - carrier_width: 32, - variant_inhabited: vec![1, 1], - ..EnumEncoding::default() - }, - ) - .into(); - - let error = build_enum_slot_map(&mut ctx, enum_ty) - .err() - .expect("an oversized shared-pointer array must remain fail-closed"); - assert!( - error - .to_string() - .contains("arrays containing shared-memory pointers are not supported"), - "{error}" - ); - assert!( - error.to_string().contains(&format!( - "supported bound is {MAX_ENUM_PAYLOAD_ARRAY_REWRITE_LEAVES}" - )), - "{error}" - ); - } - - /// A struct payload holding two shared-pointer arrays whose leaves exceed - /// the bound only in total. The slot map is target-mode agnostic (it is - /// built once, before the exporter picks the legacy 64-bit or modern - /// p3:32 data layout), so this one gate covers both output modes. - fn struct_of_shared_pointer_arrays_over_total_bound(ctx: &mut Context) -> TypeHandle { - let u32_ty = mir_uint(ctx, 32); - let shared: TypeHandle = MirPtrType::get_shared(ctx, u32_ty, false).into(); - let first: TypeHandle = MirArrayType::get(ctx, shared, 9).into(); - let second: TypeHandle = MirArrayType::get(ctx, shared, 8).into(); - MirStructType::get_with_full_layout( - ctx, - "SharedPointerArrayPair".into(), - vec!["first".into(), "second".into()], - vec![first, second], - vec![], - vec![0, 72], - 136, - 8, - ) - .into() - } - - #[test] - fn enum_slot_map_rejects_struct_of_shared_pointer_arrays_over_total_bound() { - // Each array alone (9 and 8 leaves) is within the bound; their total - // of 17 is not. The payload-root gate must reject the direct layout - // with the same bound diagnostic as a single oversized array. - let mut ctx = make_ctx(); - let discr = mir_uint(&mut ctx, 32); - let payload = struct_of_shared_pointer_arrays_over_total_bound(&mut ctx); - let enum_ty: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "StructOfSharedPointerArraysPayload".into(), - discr, - vec![0, 1], - vec![ - EnumVariant::unit("Unit".into()), - EnumVariant::new_with_layout("Pointers".into(), vec![payload], vec![8], vec![136]), - ], - EnumEncoding { - tag_offset: 0, - total_size: 144, - abi_align: 8, - layout_kind: EnumLayoutKind::Direct, - carrier_kind: EnumCarrierKind::Integer, - carrier_width: 32, - variant_inhabited: vec![1, 1], - ..EnumEncoding::default() - }, - ) - .into(); - - let error = build_enum_slot_map(&mut ctx, enum_ty) - .err() - .expect("two in-bound arrays exceeding the bound in total must fail closed"); - assert!( - error - .to_string() - .contains("arrays containing shared-memory pointers are not supported"), - "{error}" - ); - assert!( - error.to_string().contains(&format!( - "rewrite requires 17 pointer conversions, supported bound is {MAX_ENUM_PAYLOAD_ARRAY_REWRITE_LEAVES}" - )), - "{error}" - ); - } - - #[test] - fn enum_slot_map_rejects_struct_of_shared_pointer_arrays_in_niche_layout() { - // The same over-bound payload in a niche layout must report the same - // bound diagnostic as the direct layout, not a pointer-provenance - // error from the overlap walk: the payload-root gate runs first in - // both layouts. - let mut ctx = make_ctx(); - let logical = mir_uint(&mut ctx, 64); - let payload = struct_of_shared_pointer_arrays_over_total_bound(&mut ctx); - let enum_ty: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "NicheStructOfSharedPointerArrays".into(), - logical, - vec![0, 1], - vec![ - EnumVariant::unit("None".into()), - EnumVariant::new_with_layout("Some".into(), vec![payload], vec![0], vec![136]), - ], - EnumEncoding { - tag_offset: 0, - total_size: 136, - abi_align: 8, - layout_kind: EnumLayoutKind::Niche, - carrier_kind: EnumCarrierKind::Pointer, - carrier_width: 64, - untagged_variant: 1, - variant_inhabited: vec![1, 1], - ..EnumEncoding::default() - }, - ) - .into(); - - let error = build_enum_slot_map(&mut ctx, enum_ty) - .err() - .expect("the payload-root bound must also gate niche layouts"); - assert!( - error.to_string().contains(&format!( - "rewrite requires 17 pointer conversions, supported bound is {MAX_ENUM_PAYLOAD_ARRAY_REWRITE_LEAVES}" - )), - "{error}" - ); - assert!( - !error.to_string().contains("pointer provenance"), - "the bound diagnostic must not be masked by the provenance error: {error}" - ); - } - - #[test] - fn enum_slot_map_reports_bound_for_oversized_pointer_array_over_niche_carrier() { - // Generic pointers pass the shared-pointer storage gate untouched, so - // an oversized generic-pointer array reaches the overlap walk. Its - // exhausted expansion budget must surface the bound diagnostic, not - // the unrelated provenance error. - let mut ctx = make_ctx(); - let logical = mir_uint(&mut ctx, 64); - let pointee = mir_uint(&mut ctx, 32); - let pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, pointee, false).into(); - let count = MAX_ENUM_PAYLOAD_ARRAY_REWRITE_LEAVES + 1; - let pointers: TypeHandle = MirArrayType::get(&mut ctx, pointer, count).into(); - let payload_size = count * 8; - let enum_ty: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "NicheOversizedGenericPointerArray".into(), - logical, - vec![0, 1], - vec![ - EnumVariant::unit("None".into()), - EnumVariant::new_with_layout( - "Some".into(), - vec![pointers], - vec![0], - vec![payload_size], - ), - ], - EnumEncoding { - tag_offset: 0, - total_size: payload_size, - abi_align: 8, - layout_kind: EnumLayoutKind::Niche, - carrier_kind: EnumCarrierKind::Pointer, - carrier_width: 64, - untagged_variant: 1, - variant_inhabited: vec![1, 1], - ..EnumEncoding::default() - }, - ) - .into(); - - let error = build_enum_slot_map(&mut ctx, enum_ty) - .err() - .expect("an oversized pointer-array expansion over a niche carrier must fail closed"); - assert!( - error.to_string().contains(&format!( - "rewrite requires {count} pointer conversions, supported bound is {MAX_ENUM_PAYLOAD_ARRAY_REWRITE_LEAVES}" - )), - "{error}" - ); - assert!( - !error.to_string().contains("pointer provenance"), - "the bound diagnostic must not be masked by the provenance error: {error}" - ); - } - - #[test] - fn enum_slot_map_rejects_shared_pointer_vector_payload() { - let mut ctx = make_ctx(); - let tag = mir_uint(&mut ctx, 32); - let shared_pointer: TypeHandle = - llvm_types::PointerType::get(&ctx, llvm_types::address_space::SHARED).into(); - let shared_vector: TypeHandle = - llvm_types::VectorType::get(&ctx, shared_pointer, 2, llvm_types::VectorTypeKind::Fixed) - .into(); - let enum_ty: TypeHandle = MirEnumType::get_with_layout( - &mut ctx, - "SharedPointerVector".into(), - tag, - vec![0, 1], - vec![ - EnumVariant::new_with_layout( - "Data".into(), - vec![shared_vector], - vec![16], - vec![16], - ), - EnumVariant::unit("Empty".into()), - ], - 0, - 32, - 16, - ) - .into(); - - let error = build_enum_slot_map(&mut ctx, enum_ty) - .err() - .expect("a vector of shared pointers must fail closed"); - assert!( - error.to_string().contains("shared-memory pointer"), - "{error}" - ); - } - - #[test] - fn union_storage_prefers_pointer_carrier() { - let mut ctx = make_ctx(); - let u32_ty = mir_uint(&mut ctx, 32); - let u64_ty = mir_uint(&mut ctx, 64); - let ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, u32_ty, false).into(); - let union_ty = MirUnionType::get( - &mut ctx, - "PointerBits".into(), - vec!["ptr".into(), "bits".into()], - vec![ptr_ty, u64_ty], - 8, - 8, - ); - let union_data = union_ty.deref(&ctx).clone(); - let storage = build_union_storage_type(&mut ctx, &union_data).unwrap(); - let fields = struct_fields(&ctx, storage); - assert!(fields[1].deref(&ctx).is::()); - assert_eq!(llvm_type_size_align(&ctx, storage), Some((8, 8))); - } - - #[test] - fn union_storage_rejects_incompatible_pointer_address_spaces() { - let mut ctx = make_ctx(); - let u32_ty = mir_uint(&mut ctx, 32); - let generic: TypeHandle = MirPtrType::get_generic(&mut ctx, u32_ty, false).into(); - let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, u32_ty, false).into(); - let union_ty = MirUnionType::get( - &mut ctx, - "MixedPointers".into(), - vec!["generic".into(), "shared".into()], - vec![generic, shared], - 8, - 8, - ); - let union_data = union_ty.deref(&ctx).clone(); - let err = build_union_storage_type(&mut ctx, &union_data).unwrap_err(); - assert!(err.to_string().contains("different LLVM representations")); - } - - #[test] - fn union_storage_rejects_non_byte_faithful_pointer_carrier() { - let mut ctx = make_ctx(); - let u8_ty = mir_uint(&mut ctx, 8); - let u32_ty = mir_uint(&mut ctx, 32); - let bool_ty = mir_uint(&mut ctx, 1); - let ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, u32_ty, false).into(); - let ptr_bool: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "PtrBool".into(), - vec!["ptr".into(), "flag".into()], - vec![ptr_ty, bool_ty], - vec![0, 1], - vec![0, 8], - 16, - 8, - ) - .into(); - let bytes_ty: TypeHandle = MirArrayType::get(&mut ctx, u8_ty, 16).into(); - let union_ty = MirUnionType::get( - &mut ctx, - "PointerBoolBytes".into(), - vec!["view".into(), "bytes".into()], - vec![ptr_bool, bytes_ty], - 16, - 8, - ); - let union_data = union_ty.deref(&ctx).clone(); - let err = build_union_storage_type(&mut ctx, &union_data).unwrap_err(); - assert!(err.to_string().contains("no byte-faithful pointer carrier")); - } - - #[test] - fn union_storage_preserves_over_aligned_size_and_stride() { - let mut ctx = make_ctx(); - let u32_ty = mir_uint(&mut ctx, 32); - let union_ty = MirUnionType::get( - &mut ctx, - "OverAligned".into(), - vec!["word".into()], - vec![u32_ty], - 32, - 32, - ); - let union_data = union_ty.deref(&ctx).clone(); - let storage = build_union_storage_type(&mut ctx, &union_data).unwrap(); - assert_eq!(llvm_type_size_align(&ctx, storage), Some((32, 16))); - - let union_handle: TypeHandle = union_ty.into(); - let array: TypeHandle = MirArrayType::get(&mut ctx, union_handle, 3).into(); - let llvm_array = convert_type(&mut ctx, array).unwrap(); - assert_eq!(llvm_type_size_align(&ctx, llvm_array), Some((96, 16))); - } - - #[test] - fn slot_map_reorder_only() { - let mut ctx = make_ctx(); - // struct { a: u8, b: u64 }, memory order [b, a], no rustc offsets. - let a = mir_uint(&mut ctx, 8); - let b = mir_uint(&mut ctx, 64); - let layout = StructLayoutInfo { - field_types: vec![a, b], - mem_to_decl: vec![1, 0], - field_offsets: vec![], - total_size: 0, - }; - let map = build_struct_slot_map(&mut ctx, &layout).unwrap(); - - assert_eq!(map.decl_to_llvm, vec![Some(1), Some(0)]); - let i8s = llvm_int(&mut ctx, 8); - let i64s = llvm_int(&mut ctx, 64); - assert_eq!(struct_fields(&ctx, map.llvm_struct_ty), vec![i64s, i8s]); - } - - #[test] - fn slot_map_padding_only() { - let mut ctx = make_ctx(); - // struct { a: u8 @ 0, b: u64 @ 8 }, declaration order == memory - // order, size 16: lowers to { i8, [7 x i8], i64 }. The pad consumes - // slot 1, so b lands at slot 2 (the issue #128 sites used 1). - let a = mir_uint(&mut ctx, 8); - let b = mir_uint(&mut ctx, 64); - let layout = StructLayoutInfo { - field_types: vec![a, b], - mem_to_decl: vec![0, 1], - field_offsets: vec![0, 8], - total_size: 16, - }; - let map = build_struct_slot_map(&mut ctx, &layout).unwrap(); - - assert_eq!(map.decl_to_llvm, vec![Some(0), Some(2)]); - let i8s = llvm_int(&mut ctx, 8); - let i64s = llvm_int(&mut ctx, 64); - let pad7 = pad(&mut ctx, 7); - assert_eq!( - struct_fields(&ctx, map.llvm_struct_ty), - vec![i8s, pad7, i64s] - ); - } - - #[test] - fn slot_map_uses_packed_layout_when_natural_offsets_diverge() { - let mut ctx = make_ctx(); - let tag = mir_uint(&mut ctx, 8); - let value = mir_uint(&mut ctx, 32); - let layout = StructLayoutInfo { - field_types: vec![tag, value], - mem_to_decl: vec![0, 1], - field_offsets: vec![0, 1], - total_size: 5, - }; - - let map = build_struct_slot_map(&mut ctx, &layout).unwrap(); - - assert!(map.layout_diverges); - assert_eq!(map.decl_to_llvm, vec![Some(0), Some(1)]); - let llvm_struct_ty_ref = map.llvm_struct_ty.deref(&ctx); - let struct_ty = llvm_struct_ty_ref - .downcast_ref::() - .expect("slot map must produce an LLVM struct"); - assert_eq!(struct_ty.layout(), llvm_types::StructLayout::Packed); - assert_eq!(llvm_type_size_align(&ctx, map.llvm_struct_ty), Some((5, 1))); - } - - #[test] - fn slot_map_packed_two_keeps_explicit_padding_slot() { - let mut ctx = make_ctx(); - let tag = mir_uint(&mut ctx, 8); - let value = mir_uint(&mut ctx, 32); - let layout = StructLayoutInfo { - field_types: vec![tag, value], - mem_to_decl: vec![0, 1], - field_offsets: vec![0, 2], - total_size: 6, - }; - - let map = build_struct_slot_map(&mut ctx, &layout).unwrap(); - - assert!(map.layout_diverges); - assert_eq!(map.decl_to_llvm, vec![Some(0), Some(2)]); - let i8s = llvm_int(&mut ctx, 8); - let i32s = llvm_int(&mut ctx, 32); - let pad1 = pad(&mut ctx, 1); - assert_eq!( - struct_fields(&ctx, map.llvm_struct_ty), - vec![i8s, pad1, i32s] - ); - let llvm_struct_ty_ref = map.llvm_struct_ty.deref(&ctx); - let struct_ty = llvm_struct_ty_ref - .downcast_ref::() - .expect("slot map must produce an LLVM struct"); - assert_eq!(struct_ty.layout(), llvm_types::StructLayout::Packed); - assert_eq!(llvm_type_size_align(&ctx, map.llvm_struct_ty), Some((6, 1))); - } - - #[test] - fn slot_map_reorder_plus_padding() { - let mut ctx = make_ctx(); - // struct { a: u8 @ 8, b: u64 @ 0 }, memory order [b, a], size 16: - // lowers to { i64, i8, [7 x i8] } with a trailing pad. - let a = mir_uint(&mut ctx, 8); - let b = mir_uint(&mut ctx, 64); - let layout = StructLayoutInfo { - field_types: vec![a, b], - mem_to_decl: vec![1, 0], - field_offsets: vec![8, 0], - total_size: 16, - }; - let map = build_struct_slot_map(&mut ctx, &layout).unwrap(); - - assert_eq!(map.decl_to_llvm, vec![Some(1), Some(0)]); - let i8s = llvm_int(&mut ctx, 8); - let i64s = llvm_int(&mut ctx, 64); - let pad7 = pad(&mut ctx, 7); - assert_eq!( - struct_fields(&ctx, map.llvm_struct_ty), - vec![i64s, i8s, pad7] - ); - } - - #[test] - fn slot_map_zst_interleaving() { - let mut ctx = make_ctx(); - // struct { a: u32 @ 0, z: PhantomData @ 4, b: u32 @ 4 }, size 8. - // The ZST is stripped (no slot, no pad split): { i32, i32 }. - let a = mir_uint(&mut ctx, 32); - let z = mir_zst(&mut ctx); - let b = mir_uint(&mut ctx, 32); - let layout = StructLayoutInfo { - field_types: vec![a, z, b], - mem_to_decl: vec![0, 1, 2], - field_offsets: vec![0, 4, 4], - total_size: 8, - }; - let map = build_struct_slot_map(&mut ctx, &layout).unwrap(); - - assert_eq!(map.decl_to_llvm, vec![Some(0), None, Some(1)]); - let i32s = llvm_int(&mut ctx, 32); - assert_eq!(struct_fields(&ctx, map.llvm_struct_ty), vec![i32s, i32s]); - } - - #[test] - fn slot_map_issue128_arena_shape() { - let mut ctx = make_ctx(); - // The exact shape from issue #128 (examples/struct_field_layout): - // - // enum Layout { Aos, Soa, AoSoA(u32) } // -> { i8, i32 } - // struct Arena { layout: Layout, cap: u32, stride: u32, big: u64 } - // - // rustc layout: layout @ 0 (8 bytes), big @ 8, cap @ 16, - // stride @ 20, size 24. The enum now carries its own explicit - // internal padding, so the OUTER struct needs no extra padding slot: - // - // { { i8, [3 x i8], i32 }, i64, i32, i32 } - // layout=0 big=1 cap=2 stride=3 - let discr = mir_uint(&mut ctx, 8); - let payload = mir_uint(&mut ctx, 32); - let layout_enum: TypeHandle = MirEnumType::get_with_layout( - &mut ctx, - "Layout".into(), - discr, - vec![0, 1, 2], - vec![ - EnumVariant::unit("Aos".into()), - EnumVariant::unit("Soa".into()), - EnumVariant::new_with_layout("AoSoA".into(), vec![payload], vec![4], vec![4]), - ], - 0, - 8, - 4, - ) - .into(); - let cap = mir_uint(&mut ctx, 32); - let stride = mir_uint(&mut ctx, 32); - let big = mir_uint(&mut ctx, 64); - - let layout = StructLayoutInfo { - field_types: vec![layout_enum, cap, stride, big], - mem_to_decl: vec![0, 3, 1, 2], - field_offsets: vec![0, 16, 20, 8], - total_size: 24, - }; - let map = build_struct_slot_map(&mut ctx, &layout).unwrap(); - - assert_eq!( - map.decl_to_llvm, - vec![Some(0), Some(2), Some(3), Some(1)], - "the enum's internal pad must not create an outer struct slot" - ); - - let i8s = llvm_int(&mut ctx, 8); - let i32s = llvm_int(&mut ctx, 32); - let i64s = llvm_int(&mut ctx, 64); - let enum_pad3 = pad(&mut ctx, 3); - let enum_llvm: TypeHandle = llvm_types::StructType::get_unnamed( - &ctx, - ( - vec![i8s, enum_pad3, i32s], - llvm_types::StructLayout::Unpacked, - ), - ) - .into(); - assert_eq!( - struct_fields(&ctx, map.llvm_struct_ty), - vec![enum_llvm, i64s, i32s, i32s] - ); - } - - #[test] - fn slot_map_nested_struct_uses_stored_size() { - let mut ctx = make_ctx(); - // Inner struct whose stored rustc size (16) exceeds the sum of its - // converted LLVM field sizes (i8 + i64 = 9, no offsets stored). - // The outer walk must advance by the stored 16, reaching the next - // field's offset exactly: NO interior pad before it. - let x = mir_uint(&mut ctx, 8); - let y = mir_uint(&mut ctx, 64); - let inner: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "Inner".into(), - vec!["x".into(), "y".into()], - vec![x, y], - vec![], - vec![], - 16, - 0, - ) - .into(); - let c = mir_uint(&mut ctx, 8); - - let layout = StructLayoutInfo { - field_types: vec![inner, c], - mem_to_decl: vec![0, 1], - field_offsets: vec![0, 16], - total_size: 24, - }; - let map = build_struct_slot_map(&mut ctx, &layout).unwrap(); - - // inner = slot 0, c = slot 1 (adjacent), trailing [7 x i8] pad. - assert_eq!(map.decl_to_llvm, vec![Some(0), Some(1)]); - let fields = struct_fields(&ctx, map.llvm_struct_ty); - assert_eq!(fields.len(), 3, "exactly one (trailing) pad slot"); - let pad7 = pad(&mut ctx, 7); - assert_eq!(fields[2], pad7); - } - - #[test] - fn packed_shared_pointer_predicate_ignores_unpacked_direct_pointer() { - let ctx = make_ctx(); - let shared: TypeHandle = - llvm_types::PointerType::get(&ctx, llvm_types::address_space::SHARED).into(); - let unpacked: TypeHandle = llvm_types::StructType::get_unnamed( - &ctx, - (vec![shared], llvm_types::StructLayout::Unpacked), - ) - .into(); - assert!( - !llvm_packed_struct_contains_pointer_in_address_space( - &ctx, - unpacked, - llvm_types::address_space::SHARED, - ), - "a direct AS3 pointer in an unpacked aggregate is not a packed physical-image hazard" - ); - - let packed: TypeHandle = llvm_types::StructType::get_unnamed( - &ctx, - (vec![shared], llvm_types::StructLayout::Packed), - ) - .into(); - assert!( - llvm_packed_struct_contains_pointer_in_address_space( - &ctx, - packed, - llvm_types::address_space::SHARED, - ), - "an AS3 pointer under a packed struct must be rejected by by-value paths" - ); - } - - #[test] - fn packed_shared_internal_abi_genericizes_one_direct_shared_pointer() { - let mut ctx = make_ctx(); - let tag = mir_uint(&mut ctx, 8); - let pointee = mir_uint(&mut ctx, 32); - let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, false).into(); - let packed: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "PackedShared".into(), - vec!["tag".into(), "ptr".into()], - vec![tag, shared], - vec![0, 1], - vec![0, 1], - 9, - 1, - ) - .into(); - - let abi = packed_shared_internal_abi_info(&mut ctx, packed) - .expect("ABI classification must succeed") - .expect("one direct packed AS3 pointer must use the internal carrier"); - assert_eq!(llvm_type_size_align(&ctx, abi.semantic_ty), Some((9, 1))); - assert_eq!(llvm_type_size_align(&ctx, abi.storage_ty), Some((9, 1))); - assert!(llvm_packed_struct_contains_pointer_in_address_space( - &ctx, - abi.semantic_ty, - llvm_types::address_space::SHARED, - )); - assert!(!llvm_type_contains_pointer_in_address_space( - &ctx, - abi.storage_ty, - llvm_types::address_space::SHARED, - )); - assert!(llvm_type_contains_pointer_in_address_space( - &ctx, - abi.storage_ty, - llvm_types::address_space::GENERIC, - )); - } - - #[test] - fn packed_shared_internal_abi_genericizes_multiple_direct_shared_pointers() { - let mut ctx = make_ctx(); - let tag = mir_uint(&mut ctx, 8); - let pointee = mir_uint(&mut ctx, 32); - let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, false).into(); - let packed: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "PackedSharedPair".into(), - vec!["tag".into(), "left".into(), "right".into()], - vec![tag, shared, shared], - vec![0, 1, 2], - vec![0, 1, 9], - 17, - 1, - ) - .into(); - - let abi = packed_shared_internal_abi_info(&mut ctx, packed) - .expect("ABI classification must succeed") - .expect("multiple direct AS3 leaves must use the recursive internal carrier"); - assert_eq!(llvm_type_size_align(&ctx, abi.semantic_ty), Some((17, 1))); - assert_eq!(llvm_type_size_align(&ctx, abi.storage_ty), Some((17, 1))); - assert!(llvm_packed_struct_contains_pointer_in_address_space( - &ctx, - abi.semantic_ty, - llvm_types::address_space::SHARED, - )); - assert!(!llvm_type_contains_pointer_in_address_space( - &ctx, - abi.storage_ty, - llvm_types::address_space::SHARED, - )); - - let fields = struct_fields(&ctx, abi.storage_ty); - assert_eq!(fields.len(), 3); - for field in &fields[1..] { - let field_ref = field.deref(&ctx); - let pointer = field_ref - .downcast_ref::() - .expect("both pointer leaves must remain pointer-typed"); - assert_eq!(pointer.address_space(), llvm_types::address_space::GENERIC); - } - } - - #[test] - fn packed_shared_internal_abi_genericizes_nested_struct_shared_pointers() { - let mut ctx = make_ctx(); - let pointee = mir_uint(&mut ctx, 32); - let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, false).into(); - let inner: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "InnerSharedPair".into(), - vec!["left".into(), "right".into()], - vec![shared, shared], - vec![0, 1], - vec![0, 8], - 16, - 8, - ) - .into(); - let tag = mir_uint(&mut ctx, 8); - let outer: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "PackedNestedShared".into(), - vec!["tag".into(), "inner".into()], - vec![tag, inner], - vec![0, 1], - vec![0, 1], - 17, - 1, - ) - .into(); - - let abi = packed_shared_internal_abi_info(&mut ctx, outer) - .expect("classification must not error") - .expect("nested AS3 leaves must use the recursive internal carrier"); - assert_eq!(llvm_type_size_align(&ctx, abi.semantic_ty), Some((17, 1))); - assert_eq!(llvm_type_size_align(&ctx, abi.storage_ty), Some((17, 1))); - assert!(!llvm_type_contains_pointer_in_address_space( - &ctx, - abi.storage_ty, - llvm_types::address_space::SHARED, - )); - - let outer_fields = struct_fields(&ctx, abi.storage_ty); - let inner_fields = struct_fields(&ctx, outer_fields[1]); - assert_eq!(inner_fields.len(), 2); - for field in inner_fields { - let field_ref = field.deref(&ctx); - let pointer = field_ref - .downcast_ref::() - .expect("nested shared leaves must remain pointer-typed"); - assert_eq!(pointer.address_space(), llvm_types::address_space::GENERIC); - } - } - - #[test] - fn packed_shared_internal_abi_genericizes_nested_tuple_shared_pointer() { - let mut ctx = make_ctx(); - let pointee = mir_uint(&mut ctx, 32); - let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, false).into(); - let word = mir_uint(&mut ctx, 32); - let inner: TypeHandle = MirTupleType::get_with_layout( - &mut ctx, - vec![shared, word], - vec![0, 1], - vec![0, 8], - 16, - 8, - ) - .into(); - let tag = mir_uint(&mut ctx, 8); - let outer: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "PackedNestedTupleShared".into(), - vec!["tag".into(), "inner".into()], - vec![tag, inner], - vec![0, 1], - vec![0, 1], - 17, - 1, - ) - .into(); - - let abi = packed_shared_internal_abi_info(&mut ctx, outer) - .expect("classification must not error") - .expect("an AS3 leaf nested in a tuple must use the recursive carrier"); - let outer_fields = struct_fields(&ctx, abi.storage_ty); - let inner_fields = struct_fields(&ctx, outer_fields[1]); - let inner_field_ref = inner_fields[0].deref(&ctx); - let pointer = inner_field_ref - .downcast_ref::() - .expect("nested tuple leaf must remain pointer-typed"); - assert_eq!(pointer.address_space(), llvm_types::address_space::GENERIC); - assert_eq!(llvm_type_size_align(&ctx, abi.storage_ty), Some((17, 1))); - } - - #[test] - fn packed_shared_internal_abi_skips_zero_sized_fields() { - let mut ctx = make_ctx(); - let tag = mir_uint(&mut ctx, 8); - let byte = mir_uint(&mut ctx, 8); - let marker: TypeHandle = MirArrayType::get(&mut ctx, byte, 0).into(); - let pointee = mir_uint(&mut ctx, 32); - let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, false).into(); - let packed: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "PackedSharedZstMarker".into(), - vec!["tag".into(), "marker".into(), "ptr".into()], - vec![tag, marker, shared], - vec![0, 1, 2], - vec![0, 1, 1], - 9, - 1, - ) - .into(); - - let abi = packed_shared_internal_abi_info(&mut ctx, packed) - .expect("ABI classification must succeed") - .expect("a zero-sized field must not knock the struct out of the ABI lane"); - assert_eq!(llvm_type_size_align(&ctx, abi.semantic_ty), Some((9, 1))); - assert_eq!(llvm_type_size_align(&ctx, abi.storage_ty), Some((9, 1))); - assert!(!llvm_type_contains_pointer_in_address_space( - &ctx, - abi.storage_ty, - llvm_types::address_space::SHARED, - )); - } - - #[test] - fn packed_shared_internal_abi_genericizes_bounded_shared_pointer_array() { - let mut ctx = make_ctx(); - let pointee = mir_uint(&mut ctx, 32); - let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, false).into(); - let shared_array: TypeHandle = MirArrayType::get(&mut ctx, shared, 2).into(); - let tag = mir_uint(&mut ctx, 8); - let outer: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "PackedSharedArray".into(), - vec!["tag".into(), "ptrs".into()], - vec![tag, shared_array], - vec![0, 1], - vec![0, 1], - 17, - 1, - ) - .into(); - - let abi = packed_shared_internal_abi_info(&mut ctx, outer) - .expect("classification must not error") - .expect("a bounded AS3 array must use the target-stable internal carrier"); - assert_eq!(llvm_type_size_align(&ctx, abi.semantic_ty), Some((17, 1))); - assert_eq!(llvm_type_size_align(&ctx, abi.storage_ty), Some((17, 1))); - assert!(!llvm_type_contains_pointer_in_address_space( - &ctx, - abi.storage_ty, - llvm_types::address_space::SHARED, - )); - - let outer_fields = struct_fields(&ctx, abi.storage_ty); - let array_ref = outer_fields[1].deref(&ctx); - let array = array_ref - .downcast_ref::() - .expect("bounded array field must remain an LLVM array"); - assert_eq!(array.size(), 2); - let element = array.elem_type(); - let element_ref = element.deref(&ctx); - let pointer = element_ref - .downcast_ref::() - .expect("array elements must remain pointer-typed"); - assert_eq!(pointer.address_space(), llvm_types::address_space::GENERIC); - } - - #[test] - fn packed_shared_internal_abi_accepts_array_rewrite_at_exact_bound() { - let mut ctx = make_ctx(); - let pointee = mir_uint(&mut ctx, 32); - let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, false).into(); - let count = MAX_PACKED_SHARED_INTERNAL_ABI_ARRAY_REWRITE_LEAVES; - let shared_array: TypeHandle = MirArrayType::get(&mut ctx, shared, count).into(); - let tag = mir_uint(&mut ctx, 8); - let outer: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "PackedSharedArrayAtBound".into(), - vec!["tag".into(), "ptrs".into()], - vec![tag, shared_array], - vec![0, 1], - vec![0, 1], - 1 + 8 * count, - 1, - ) - .into(); - - assert!( - packed_shared_internal_abi_info(&mut ctx, outer) - .expect("classification must not error") - .is_some(), - "the exact bounded-array rewrite limit must remain supported" - ); - } - - #[test] - fn packed_shared_internal_abi_rejects_array_rewrite_above_bound() { - let mut ctx = make_ctx(); - let pointee = mir_uint(&mut ctx, 32); - let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, false).into(); - let count = MAX_PACKED_SHARED_INTERNAL_ABI_ARRAY_REWRITE_LEAVES + 1; - let shared_array: TypeHandle = MirArrayType::get(&mut ctx, shared, count).into(); - let tag = mir_uint(&mut ctx, 8); - let outer: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "PackedSharedArrayAboveBound".into(), - vec!["tag".into(), "ptrs".into()], - vec![tag, shared_array], - vec![0, 1], - vec![0, 1], - 1 + 8 * count, - 1, - ) - .into(); - - assert!( - packed_shared_internal_abi_info(&mut ctx, outer) - .expect("classification must not error") - .is_none(), - "array-expanded AS3 leaves above the explicit budget must fail closed" - ); - } - - #[test] - fn packed_shared_internal_abi_rejects_shared_pointer_vector() { - let mut ctx = make_ctx(); - let tag = mir_uint(&mut ctx, 8); - let shared_pointer: TypeHandle = - llvm_types::PointerType::get(&ctx, llvm_types::address_space::SHARED).into(); - let shared_vector: TypeHandle = - llvm_types::VectorType::get(&ctx, shared_pointer, 2, llvm_types::VectorTypeKind::Fixed) - .into(); - let outer: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "PackedSharedVector".into(), - vec!["tag".into(), "ptrs".into()], - vec![tag, shared_vector], - vec![0, 1], - vec![0, 1], - 17, - 1, - ) - .into(); - - assert!( - packed_shared_internal_abi_info(&mut ctx, outer) - .expect("classification must not error") - .is_none(), - "shared-pointer vectors must remain outside the packed-AS3 internal ABI lane" - ); - } - - #[test] - fn slot_map_rejects_malformed_memory_order() { - let mut ctx = make_ctx(); - let a = mir_uint(&mut ctx, 8); - let b = mir_uint(&mut ctx, 64); - - // Not a permutation: decl index 0 appears twice. - let dup = StructLayoutInfo { - field_types: vec![a, b], - mem_to_decl: vec![0, 0], - field_offsets: vec![], - total_size: 0, - }; - assert!(build_struct_slot_map(&mut ctx, &dup).is_err()); - - // Wrong length. - let short = StructLayoutInfo { - field_types: vec![a, b], - mem_to_decl: vec![0], - field_offsets: vec![], - total_size: 0, - }; - assert!(build_struct_slot_map(&mut ctx, &short).is_err()); - - // Offsets vector length mismatch (with explicit layout engaged). - let bad_offsets = StructLayoutInfo { - field_types: vec![a, b], - mem_to_decl: vec![0, 1], - field_offsets: vec![0], - total_size: 16, - }; - assert!(build_struct_slot_map(&mut ctx, &bad_offsets).is_err()); - } - - #[test] - fn initialized_global_layout_accepts_explicit_overalignment() { - let mut ctx = make_ctx(); - let zst = mir_zst(&mut ctx); - validate_initialized_global_layout(&mut ctx, zst, 0, 1).unwrap(); - - let byte = mir_uint(&mut ctx, 8); - let over_aligned: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "OverAligned".into(), - vec!["byte".into()], - vec![byte], - vec![0], - vec![0], - 16, - 16, - ) - .into(); - - validate_initialized_global_layout(&mut ctx, over_aligned, 16, 16).unwrap(); - } - - #[test] - fn initialized_global_layout_rejects_packed_and_nested_packed_structs() { - let mut ctx = make_ctx(); - let byte = mir_uint(&mut ctx, 8); - let word = mir_uint(&mut ctx, 32); - let packed: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "Packed".into(), - vec!["tag".into(), "word".into()], - vec![byte, word], - vec![0, 1], - vec![0, 1], - 5, - 1, - ) - .into(); - - let empty_packed: TypeHandle = MirArrayType::get(&mut ctx, packed, 0).into(); - validate_initialized_global_layout(&mut ctx, empty_packed, 0, 1) - .expect("[Packed; 0] has no element bytes whose field layout can diverge"); - - let err = validate_initialized_global_layout(&mut ctx, packed, 5, 1).unwrap_err(); - assert!(err.to_string().contains("field 1 lowers at byte 4")); - - // Nesting must not hide the incompatible packed representation. - let wide = mir_uint(&mut ctx, 64); - let outer: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "Outer".into(), - vec!["packed".into(), "wide".into()], - vec![packed, wide], - vec![0, 1], - vec![0, 8], - 16, - 8, - ) - .into(); - let err = validate_initialized_global_layout(&mut ctx, outer, 16, 8).unwrap_err(); - assert!(err.to_string().contains("lowers at byte")); - } - - #[test] - fn relocated_initialized_global_layout_accepts_top_level_packed_struct() { - let mut ctx = make_ctx(); - let byte = mir_uint(&mut ctx, 8); - let target = mir_uint(&mut ctx, 32); - let pointer: TypeHandle = MirPtrType::get_global(&mut ctx, target, false).into(); - let packed: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "PackedRelocation".into(), - vec!["tag".into(), "ptr".into()], - vec![byte, pointer], - vec![0, 1], - vec![0, 1], - 9, - 1, - ) - .into(); - - validate_relocated_initialized_global_layout(&mut ctx, packed, 9, 1) - .expect("relocated top-level packed struct must be accepted"); - assert!(validate_initialized_global_layout(&mut ctx, packed, 9, 1).is_err()); - } - - #[test] - fn relocated_initialized_global_layout_accepts_thin_pointer_union() { - let mut ctx = make_ctx(); - let u32_ty = mir_uint(&mut ctx, 32); - let u8_ty = mir_uint(&mut ctx, 8); - let word_ptr: TypeHandle = MirPtrType::get_generic(&mut ctx, u32_ty, false).into(); - let byte_ptr: TypeHandle = MirPtrType::get_generic(&mut ctx, u8_ty, false).into(); - let union_ty: TypeHandle = MirUnionType::get( - &mut ctx, - "RelocatedPointerUnion".into(), - vec!["word".into(), "byte".into()], - vec![word_ptr, byte_ptr], - 8, - 8, - ) - .into(); - - validate_relocated_initialized_global_layout(&mut ctx, union_ty, 8, 8) - .expect("one pointer-word union must be valid relocated global storage"); - } - - #[test] - fn relocated_initialized_global_layout_rejects_malformed_memory_order() { - let mut ctx = make_ctx(); - let byte = mir_uint(&mut ctx, 8); - let target = mir_uint(&mut ctx, 32); - let pointer: TypeHandle = MirPtrType::get_global(&mut ctx, target, false).into(); - let malformed: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "MalformedPackedRelocation".into(), - vec!["tag".into(), "ptr".into()], - vec![byte, pointer], - vec![0, 0], - vec![0, 1], - 9, - 1, - ) - .into(); - - let error = validate_relocated_initialized_global_layout(&mut ctx, malformed, 9, 1) - .expect_err("malformed memory order must remain unsupported"); - assert!(error.to_string().contains("not a permutation"), "{error}"); - } - - #[test] - fn relocated_initialized_global_layout_accepts_one_direct_nested_packed_struct() { - let mut ctx = make_ctx(); - let byte = mir_uint(&mut ctx, 8); - let word = mir_uint(&mut ctx, 32); - let pointer: TypeHandle = MirPtrType::get_global(&mut ctx, word, false).into(); - let nested_packed: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "NestedPackedRelocation".into(), - vec!["tag".into(), "ptr".into()], - vec![byte, pointer], - vec![0, 1], - vec![0, 1], - 9, - 1, - ) - .into(); - let outer: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "OuterPackedRelocation".into(), - vec!["head".into(), "nested".into()], - vec![word, nested_packed], - vec![0, 1], - vec![0, 4], - 16, - 4, - ) - .into(); - - validate_relocated_initialized_global_layout(&mut ctx, outer, 16, 4) - .expect("one direct nested packed relocation carrier must be accepted"); - assert!(validate_initialized_global_layout(&mut ctx, outer, 16, 4).is_err()); - } - - #[test] - fn relocated_initialized_global_layout_rejects_packed_root_with_packed_child() { - let mut ctx = make_ctx(); - let byte = mir_uint(&mut ctx, 8); - let word = mir_uint(&mut ctx, 32); - let target = mir_uint(&mut ctx, 32); - let pointer: TypeHandle = MirPtrType::get_global(&mut ctx, target, false).into(); - let nested_packed: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "NestedPackedRelocation".into(), - vec!["tag".into(), "ptr".into()], - vec![byte, pointer], - vec![0, 1], - vec![0, 1], - 9, - 1, - ) - .into(); - let packed_root: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "PackedRootRelocation".into(), - vec!["tag".into(), "word".into(), "nested".into()], - vec![byte, word, nested_packed], - vec![0, 1, 2], - vec![0, 1, 5], - 14, - 1, - ) - .into(); - - let error = validate_relocated_initialized_global_layout(&mut ctx, packed_root, 14, 1) - .expect_err("a packed top-level struct must not stack the nested packed relaxation"); - assert!(error.to_string().contains("lowers at byte"), "{error}"); - } - - #[test] - fn relocated_initialized_global_layout_rejects_deeper_packed_nesting() { - let mut ctx = make_ctx(); - let byte = mir_uint(&mut ctx, 8); - let word = mir_uint(&mut ctx, 32); - let target = mir_uint(&mut ctx, 32); - let pointer: TypeHandle = MirPtrType::get_global(&mut ctx, target, false).into(); - let inner_packed: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "InnerPackedRelocation".into(), - vec!["tag".into(), "ptr".into()], - vec![byte, pointer], - vec![0, 1], - vec![0, 1], - 9, - 1, - ) - .into(); - let middle_packed: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "MiddlePackedRelocation".into(), - vec!["tag".into(), "word".into(), "inner".into()], - vec![byte, word, inner_packed], - vec![0, 1, 2], - vec![0, 1, 5], - 14, - 1, - ) - .into(); - let outer: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "OuterDeepPackedRelocation".into(), - vec!["head".into(), "middle".into()], - vec![word, middle_packed], - vec![0, 1], - vec![0, 4], - 20, - 4, - ) - .into(); - - let error = validate_relocated_initialized_global_layout(&mut ctx, outer, 20, 4) - .expect_err("packed relocation relaxation must stop after one nesting level"); - assert!(error.to_string().contains("lowers at byte"), "{error}"); - } - - #[test] - fn relocated_initialized_global_layout_rejects_overlapping_struct_fields() { - let mut ctx = make_ctx(); - let word = mir_uint(&mut ctx, 32); - let overlapping: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "OverlappingRelocation".into(), - vec!["left".into(), "right".into()], - vec![word, word], - vec![0, 1], - vec![0, 0], - 4, - 1, - ) - .into(); - - let error = validate_relocated_initialized_global_layout(&mut ctx, overlapping, 4, 1) - .expect_err("overlapping field ranges must remain unsupported"); - assert!(error.to_string().contains("overlap"), "{error}"); - } - - #[test] - fn initialized_global_layout_rejects_old_union_and_tuple_models() { - let mut ctx = make_ctx(); - let word = mir_uint(&mut ctx, 32); - let union_as_struct: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "UnionBeforeSharedStorageLowering".into(), - vec!["left".into(), "right".into()], - vec![word, word], - vec![0, 1], - vec![0, 0], - 4, - 4, - ) - .into(); - let err = validate_initialized_global_layout(&mut ctx, union_as_struct, 4, 4).unwrap_err(); - assert!(err.to_string().contains("field 1 lowers at byte 4")); - - // A tuple without recorded rustc layout cannot prove its bytes. - let byte = mir_uint(&mut ctx, 8); - let wide = mir_uint(&mut ctx, 64); - let tuple: TypeHandle = MirTupleType::get(&mut ctx, vec![byte, wide]).into(); - let err = validate_initialized_global_layout(&mut ctx, tuple, 16, 8).unwrap_err(); - assert!( - err.to_string() - .contains("no stored size but lowers to 16 bytes"), - "{err}" - ); - - // The same tuple carrying rustc's real (reordered) layout validates: - // memory order is (wide @ 0, byte @ 8), total size 16. - let laid_out_tuple: TypeHandle = MirTupleType::get_with_layout( - &mut ctx, - vec![byte, wide], - vec![1, 0], - vec![8, 0], - 16, - 8, - ) - .into(); - validate_initialized_global_layout(&mut ctx, laid_out_tuple, 16, 8) - .expect("tuples with recorded rustc offsets are provable initialized-global storage"); - } - - #[test] - fn initialized_global_layout_rejects_enum_with_unknown_physical_layout() { - let mut ctx = make_ctx(); - let discr = mir_uint(&mut ctx, 8); - let payload = mir_uint(&mut ctx, 32); - let niche: TypeHandle = MirEnumType::get( - &mut ctx, - "OptionNonZero".into(), - discr, - vec![0, 1], - vec![ - EnumVariant::unit("None".into()), - EnumVariant::new("Some".into(), vec![payload]), - ], - ) - .into(); - - let err = validate_initialized_global_layout(&mut ctx, niche, 4, 4).unwrap_err(); - assert!(err.to_string().contains("unknown physical layout")); - } -} diff --git a/crates/mir-lower/src/convert/types/conversion.rs b/crates/mir-lower/src/convert/types/conversion.rs new file mode 100644 index 0000000000..1df786d65f --- /dev/null +++ b/crates/mir-lower/src/convert/types/conversion.rs @@ -0,0 +1,153 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! `convert_type` dispatch, ZST detection, and the slice fat-pointer builders. + +use llvm_export::types as llvm_types; +use llvm_export::types::PointerTypeExt; +use pliron::builtin::types::{IntegerType, Signedness}; +use pliron::context::Context; +use pliron::r#type::{TypeHandle, type_cast}; + +use crate::type_conversion_interface::MirTypeConversion; + +// ============================================================================= +// Zero-Sized Type (ZST) Detection +// ============================================================================= + +/// Check if a type is zero-sized (empty struct). +/// +/// Zero-sized types include: +/// - Empty structs `struct {}` +/// - PhantomData markers (which become empty structs in MIR) +/// - Structs where all fields are themselves zero-sized +/// +/// # Why This Matters +/// +/// LLVM's NVPTX backend doesn't support empty struct types in function +/// signatures. We strip these during type conversion to avoid: +/// `LLVM ERROR: Empty parameter types are not supported` +/// +/// # Background +/// +/// Rust's `#[inline(always)]` attribute is stored in `codegen_fn_attrs`, which +/// is not exposed through the stable_mir API. Since we intercept MIR and generate +/// our own LLVM IR, we don't propagate inline hints. When LLVM decides not to +/// inline a function, the empty struct parameters/returns cause NVPTX to crash. +/// +/// By stripping ZSTs at the LLVM type level, we avoid this issue regardless of +/// inlining decisions. +pub fn is_zero_sized_type(ctx: &Context, ty: TypeHandle) -> bool { + if let Some(array_ty) = ty.deref(ctx).downcast_ref::() { + return array_ty.size() == 0 || is_zero_sized_type(ctx, array_ty.elem_type()); + } + + // Check if LLVM StructType with zero fields + if let Some(struct_ty) = ty.deref(ctx).downcast_ref::() { + let num_fields = struct_ty.num_fields(); + if num_fields == 0 { + return true; + } + // Also check if ALL fields are zero-sized (nested PhantomData) + return struct_ty.fields().all(|f| is_zero_sized_type(ctx, f)); + } + false +} + +// ============================================================================= +// Type Conversion +// ============================================================================= + +/// Convert a `dialect-mir` type to its LLVM dialect equivalent. +/// +/// Dispatches via `MirTypeConversion` type interface — each supported type +/// registers a converter function pointer through `#[type_interface_impl]` +/// in [`crate::convert::type_interface_impls`]. +/// +/// The function-pointer indirection avoids a borrow-checker conflict: +/// `type_cast` borrows `ctx` immutably, but conversion needs `&mut ctx`. +/// We extract the `Copy` function pointer, drop the borrow, then call it. +pub fn convert_type(ctx: &mut Context, ty: TypeHandle) -> Result { + // Phase 1: extract a Copy function pointer while ctx is immutably borrowed. + let converter_fn = { + let ty_ref = ty.deref(ctx); + type_cast::(&*ty_ref).map(|conv| conv.converter()) + }; + // Phase 2: borrow dropped — ctx is free for &mut. + if let Some(conv_fn) = converter_fn { + return conv_fn(ty, ctx); + } + + let type_display = ty.deref(ctx).disp(ctx).to_string(); + Err(anyhow::anyhow!( + "Unsupported type conversion: {}\n\ + Supported: integers, fp32, fp64, pointers, slices, tuples, structs, enums, arrays, vectors.", + type_display + )) +} + +/// Create the LLVM struct type used for slice representations. +/// +/// Slices are represented as fat pointers: `{ ptr, i64 }` where: +/// - `ptr` is a generic address space (0) pointer to the data +/// - `i64` is the number of elements (not bytes) +/// +/// # Layout +/// +/// ```text +/// struct { +/// ptr: !llvm.ptr, ; offset 0, size 8 +/// len: i64, ; offset 8, size 8 +/// } ; total size: 16 bytes +/// ``` +/// +/// # Address Space +/// +/// The pointer uses generic address space (0) because: +/// - Slices passed to kernels may point to global memory +/// - The kernel doesn't know at compile time which memory space +/// - Generic pointers can be used with any memory type +/// +/// # Usage +/// +/// This type is used for: +/// - `&[T]` slice arguments +/// - `DisjointSlice` (unique-ownership slice) arguments +/// - Any other fat pointer representation +pub(crate) fn make_slice_struct(ctx: &mut Context) -> TypeHandle { + let ptr_ty = llvm_types::PointerType::get_generic(ctx); + let len_ty = IntegerType::get(ctx, 64, Signedness::Signless); + llvm_types::StructType::get_unnamed( + ctx, + ( + vec![ptr_ty.into(), len_ty.into()], + llvm_types::StructLayout::Unpacked, + ), + ) + .into() +} + +/// The fat pointer, followed by an index space's runtime layout fields. +/// +/// With no such fields this is exactly [`make_slice_struct`], so a slice over +/// an index space fixed in its type keeps its two-field representation. +pub(crate) fn make_disjoint_slice_struct( + ctx: &mut Context, + space_tys: &[TypeHandle], +) -> anyhow::Result { + if space_tys.is_empty() { + return Ok(make_slice_struct(ctx)); + } + let ptr_ty = llvm_types::PointerType::get_generic(ctx); + let len_ty = IntegerType::get(ctx, 64, Signedness::Signless); + let mut fields: Vec = vec![ptr_ty.into(), len_ty.into()]; + for space_ty in space_tys { + fields.push(convert_type(ctx, *space_ty)?); + } + Ok( + llvm_types::StructType::get_unnamed(ctx, (fields, llvm_types::StructLayout::Unpacked)) + .into(), + ) +} diff --git a/crates/mir-lower/src/convert/types/enum_layout.rs b/crates/mir-lower/src/convert/types/enum_layout.rs new file mode 100644 index 0000000000..25eb043d10 --- /dev/null +++ b/crates/mir-lower/src/convert/types/enum_layout.rs @@ -0,0 +1,2014 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! The `EnumSlotMap` placement engine and enum type conversion +//! (`build_enum_slot_map`, `convert_enum_to_llvm`, unmodeled-enum probes). + +use dialect_mir::types::{ + EnumCarrierKind, EnumLayoutKind, MirDisjointSliceType, MirEnumType, MirSliceType, + MirStructType, MirTupleType, MirUnionType, +}; +use llvm_export::types as llvm_types; +use pliron::builtin::types::{IntegerType, Signedness}; +use pliron::context::Context; +use pliron::r#type::TypeHandle; + +use super::layout::{make_enum_filler_type, mir_type_contains_i1}; +use super::pointer_storage::{ + PointerOverlapRejection, analyze_pointer_overlap, llvm_type_contains_pointer, +}; +use super::{ + convert_type, is_zero_sized_type, llvm_type_contains_i1, llvm_type_is_byte_faithful, + llvm_type_size_align, natural_struct_layout, +}; +use crate::convert::enum_payload_storage::{ + MAX_ENUM_PAYLOAD_ARRAY_REWRITE_LEAVES, enum_payload_storage_type, +}; + +/// The LLVM struct for an enum, plus a map saying where the tag and each +/// payload field ended up. +/// +/// The struct type and the indices into it are produced by one walk in +/// [`build_enum_slot_map`], so they can never disagree. (Computing them +/// separately is how the issue #128 class of bug happened for structs.) +pub(crate) struct EnumSlotMap { + /// The final LLVM struct type, including any `[N x i8]` filler slots. + pub llvm_struct_ty: TypeHandle, + /// Which struct slot holds rustc's physical tag/niche carrier. `None` + /// for `Single` and `Empty` layouts, which have no tag in memory. + pub carrier_slot: Option, + /// Converted physical carrier type (integer or pointer), when present. + pub carrier_llvm_ty: Option, + /// Which struct slot holds each payload field, in the flattened + /// order of `MirEnumType::all_field_types`. `None` means the field + /// has no slot of its own: it is zero-sized, or its bytes are shared + /// with a different-typed field of another variant. Such fields are + /// read and written through memory at `field_offsets` instead. + pub field_slots: Vec>, + /// Byte position of each payload field inside the enum. + pub field_offsets: Vec, + /// Converted LLVM type of each payload field. + pub field_llvm_types: Vec, +} + +/// Build the LLVM struct for an enum, placing everything at the byte +/// positions rustc chose. +/// +/// Why this matters: the host (CPU) lays out enum values with rustc's +/// layout. If the device used different byte positions, every enum +/// passed to a kernel would be read wrong. So the device struct is built +/// to have the same bytes, position for position. +/// +/// The wrinkle is that enum variants SHARE bytes (only one variant is +/// alive at a time), and an LLVM struct cannot say "these two fields +/// overlap". The slot map resolves each field one of three ways: +/// +/// ```text +/// #[repr(u32)] enum E { A(u32), B(f32), C } +/// rustc: 8 bytes, tag at byte 0, A's u32 and B's f32 both at byte 4 +/// +/// LLVM struct: { i32, i32 } +/// | | +/// tag_slot=0 A's payload: own slot (nothing else typed i32 +/// wanted byte 4 first... B did, see below) +/// +/// - own slot: the field's bytes collide with nothing already placed. +/// - shared slot: another variant already placed the SAME type at the +/// SAME position; both map to that slot. (If B were +/// B(u32), A and B would simply share slot 1.) +/// - no slot: the bytes are taken by a different type (B's f32 vs +/// A's u32 here). The field is still at byte 4, just not +/// nameable as a struct field; reads and writes go +/// through memory: spill the value to a stack slot, then +/// use a byte-precise pointer. No slot, but no lie. +/// ``` +/// +/// Gaps between placed fields, and the tail, are covered with explicit +/// `[N x i8]` filler so the struct's size is exactly rustc's no matter +/// what LLVM's own layout rules would have done. +/// +/// Direct and Niche carriers are claimed first, so a semantic field with a +/// different SSA type cannot redefine the same physical bytes. Single and +/// Empty layouts simply have no carrier. Unknown layouts are rejected. +/// +/// If the finished struct's size does not come out equal to rustc's, +/// something is deeply wrong and lowering would miscompile, so that is a +/// hard error rather than a debug assertion. +pub(crate) fn build_enum_slot_map( + ctx: &mut Context, + ty: TypeHandle, +) -> Result { + let ( + name, + discriminant_ty, + all_field_types, + all_field_offsets, + all_field_sizes, + variant_field_counts, + variant_inhabited, + tag_offset, + total_size, + abi_align, + layout_kind, + carrier_kind, + carrier_width, + carrier_address_space, + ) = { + let ty_ref = ty.deref(ctx); + let enum_ty = ty_ref + .downcast_ref::() + .ok_or_else(|| anyhow::anyhow!("build_enum_slot_map: expected MirEnumType"))?; + ( + enum_ty.name().to_string(), + enum_ty.discriminant_ty, + enum_ty.all_field_types.clone(), + enum_ty.all_field_offsets.clone(), + enum_ty.all_field_sizes.clone(), + enum_ty.variant_field_counts.clone(), + enum_ty.variant_inhabited.clone(), + enum_ty.tag_offset(), + enum_ty.total_size(), + enum_ty.abi_align(), + enum_ty.layout_kind, + enum_ty.carrier_kind, + enum_ty.carrier_width, + enum_ty.carrier_address_space, + ) + }; + + if layout_kind == EnumLayoutKind::Unknown { + return Err(anyhow::anyhow!( + "enum slot map: `{}` has unknown physical layout; refusing to guess its bytes", + name + )); + } + if carrier_kind == EnumCarrierKind::Pointer + && carrier_address_space == llvm_types::address_space::SHARED + { + return Err(anyhow::anyhow!( + "enum slot map: `{}` uses a shared-memory pointer carrier whose size is target-mode dependent (64-bit PTX/legacy, 32-bit modern NVVM); refusing target-agnostic enum lowering", + name + )); + } + if carrier_kind == EnumCarrierKind::Integer && !carrier_width.is_multiple_of(8) { + return Err(anyhow::anyhow!( + "enum slot map: `{}` integer carrier width {} is not a whole number of bytes; refusing physical storage with unspecified upper byte bits", + name, + carrier_width + )); + } + + let carrier_ty: Option = match carrier_kind { + EnumCarrierKind::None => None, + EnumCarrierKind::Integer if layout_kind == EnumLayoutKind::Direct => { + let converted = convert_type(ctx, discriminant_ty)?; + let width = converted + .deref(ctx) + .downcast_ref::() + .map(IntegerType::width); + if width != Some(carrier_width) { + return Err(anyhow::anyhow!( + "enum slot map: `{}` direct carrier does not match its declared discriminant type", + name + )); + } + Some(converted) + } + EnumCarrierKind::Integer => { + Some(IntegerType::get(ctx, carrier_width, Signedness::Signless).into()) + } + EnumCarrierKind::Pointer => { + Some(llvm_types::PointerType::get(ctx, carrier_address_space).into()) + } + }; + let mut field_llvm_types = Vec::with_capacity(all_field_types.len()); + for &field_ty in &all_field_types { + field_llvm_types.push(convert_type(ctx, field_ty)?); + } + + if all_field_offsets.len() != all_field_types.len() + || all_field_sizes.len() != all_field_types.len() + { + return Err(anyhow::anyhow!( + "enum slot map: `{}` has {} field offsets for {} fields", + name, + all_field_offsets.len().min(all_field_sizes.len()), + all_field_types.len() + )); + } + + // Phase 1: decide who gets a struct slot. The physical carrier goes + // first so a semantic field can never claim its bytes using a different + // type (e.g. `bool` is i1 semantically but i8 in Option storage). + // claims: (byte position, byte size, converted type), no two overlap. + let mut claims: Vec<(u64, u64, TypeHandle)> = Vec::new(); + let carrier_claim = if let Some(carrier_ty) = carrier_ty { + let (carrier_size, carrier_align) = + llvm_type_size_align(ctx, carrier_ty).ok_or_else(|| { + anyhow::anyhow!( + "enum slot map: `{}` carrier has unsupported LLVM layout", + name + ) + })?; + let expected_size = u64::from(carrier_width).div_ceil(8); + if carrier_size != expected_size + || tag_offset % carrier_align.max(1) != 0 + || tag_offset + .checked_add(carrier_size) + .is_none_or(|end| end > total_size) + { + return Err(anyhow::anyhow!( + "enum slot map: `{}` carrier (size {}, align {}) cannot sit at byte {} of {}", + name, + carrier_size, + carrier_align, + tag_offset, + total_size + )); + } + claims.push((tag_offset, carrier_size, carrier_ty)); + Some(0usize) + } else { + None + }; + + let mut claim_of_field: Vec> = vec![None; field_llvm_types.len()]; + let mut field_is_inhabited = Vec::with_capacity(field_llvm_types.len()); + for (variant, count) in variant_field_counts.iter().enumerate() { + field_is_inhabited.extend(std::iter::repeat_n( + variant_inhabited.get(variant).copied().unwrap_or(0) != 0, + *count as usize, + )); + } + let mut order: Vec = (0..field_llvm_types.len()).collect(); + // At one byte position, prefer a representation that preserves all of + // its stored bits. This makes the result independent of source-variant + // order and prevents an i1/bool view from claiming storage that a later + // i8 view needs to preserve. A scalar i1 always uses an i8 storage claim: + // construction explicitly zero-extends that scalar below. + order.sort_by_key(|&i| { + ( + all_field_offsets[i], + !llvm_type_is_byte_faithful(ctx, field_llvm_types[i]), + i, + ) + }); + for flat in order { + if !field_is_inhabited.get(flat).copied().unwrap_or(false) { + continue; + } + let llvm_ty = field_llvm_types[flat]; + // Enum payload storage uses one target-stable physical view. Shared + // pointer leaves become generic pointers recursively through LLVM + // structs (MIR structs/tuples) and bounded arrays, while bool leaves + // become canonical i8 bytes. Oversized shared-pointer arrays and + // pointer vectors fail closed here. + let storage_ty = enum_payload_storage_type(ctx, llvm_ty).map_err(|error| { + anyhow::anyhow!("enum slot map: `{}` field {}: {error}", name, flat) + })?; + let (size, align) = llvm_type_size_align(ctx, storage_ty).ok_or_else(|| { + anyhow::anyhow!( + "enum slot map: `{}` field {} has unsupported LLVM size/alignment", + name, + flat + ) + })?; + if size != all_field_sizes[flat] { + return Err(anyhow::anyhow!( + "enum slot map: `{}` field {} lowers to {} bytes but rustc says {}", + name, + flat, + size, + all_field_sizes[flat] + )); + } + if size == 0 || is_zero_sized_type(ctx, llvm_ty) { + // ZSTs own no bytes and no slot. + continue; + } + let offset = all_field_offsets[flat]; + if offset.checked_add(size).is_none_or(|end| end > total_size) { + return Err(anyhow::anyhow!( + "enum slot map: `{}` field {} (size {}) at byte {} exceeds total size {}", + name, + flat, + size, + offset, + total_size + )); + } + if !offset.is_multiple_of(align.max(1)) { + return Err(anyhow::anyhow!( + "enum slot map: `{}` field {} requires alignment {} but rustc offset {} is not aligned; packed enum payload access is not yet supported", + name, + flat, + align, + offset + )); + } + + let is_scalar_i1 = llvm_ty + .deref(ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 1); + // A bool the MIR type promises but the LLVM lowering hides (a union + // whose storage is a raw byte blob) cannot be canonicalized here: + // its bytes were written by the union's own stores, not by this + // enum's construction. + if !is_scalar_i1 + && mir_type_contains_i1(ctx, all_field_types[flat]) + && !llvm_type_contains_i1(ctx, llvm_ty) + { + return Err(anyhow::anyhow!( + "enum slot map: `{}` field {} contains bool storage hidden behind a union; cuda-oxide cannot prove those bytes are canonical", + name, + flat + )); + } + + // Rust bool is a semantic i1 but occupies a complete byte in memory. + // Never make i1 the struct's physical storage type: give a standalone + // bool an i8 claim, or reuse the exact i8 carrier/field claim already + // covering that byte. Construction leaves the bool slotless and the + // deferred store below explicitly zero-extends i1 -> i8; extraction + // loads i1 from that byte. + if is_scalar_i1 { + let colliding_claims = claims + .iter() + .filter(|&&(o, s, _)| offset < o + s && o < offset + size) + .collect::>(); + if colliding_claims.is_empty() { + let byte_ty: TypeHandle = IntegerType::get(ctx, 8, Signedness::Signless).into(); + claims.push((offset, 1, byte_ty)); + continue; + } + let has_exact_i8_storage = colliding_claims.iter().all(|&&(o, s, claim_ty)| { + o == offset + && s == 1 + && claim_ty + .deref(ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 8) + }); + if has_exact_i8_storage { + continue; + } + return Err(anyhow::anyhow!( + "enum slot map: `{}` scalar bool field {} overlaps storage other than one exact i8 byte; refusing to expose non-canonical bool bits", + name, + flat + )); + } + + // Aggregates containing bool claim their byte-faithful twin (every + // i1 leaf becomes its canonical i8 memory byte) and stay slotless: + // construction canonicalizes the value and writes it at its byte + // offset; extraction re-loads the original type from the canonical + // bytes, exactly like the scalar-bool path above. + // Bool-bearing aggregates remain slotless so their canonical byte view + // continues through the established spill path. Pointer-only nested + // aggregates may own a slot because construction/extraction now rebuild + // them recursively at the semantic/physical boundary. + let field_gets_slot = !llvm_type_contains_i1(ctx, llvm_ty); + + // Another variant already placed the same storage type at the same + // position? Then both fields can simply use that claim: variants + // share bytes, and here they even agree on the type. + if let Some(ci) = claims + .iter() + .position(|&(o, _, t)| o == offset && t == storage_ty) + { + if field_gets_slot { + claim_of_field[flat] = Some(ci); + } + continue; + } + // The bytes are taken by a different type. Pointer-free values can + // use the memory fallback below. Pointer-bearing values may do so only + // when a physical-layout walk proves that every pointer leaf exactly + // matches an existing pointer slot. This is the common + // `Option<(usize, &T)>` shape: the niche carrier is the tuple's pointer + // field. A pointer/non-pointer overlap, address-space mismatch, or + // additional pointer without a pointer slot still fails closed. + let colliding_claims = claims + .iter() + .filter(|&&(o, s, _)| offset < o + s && o < offset + size) + .collect::>(); + if !colliding_claims.is_empty() { + let has_pointer_overlap = llvm_type_contains_pointer(ctx, storage_ty) + || colliding_claims + .iter() + .any(|&&(_, _, claim_ty)| llvm_type_contains_pointer(ctx, claim_ty)); + // Pointer-bearing overlaps must back every pointer leaf with a real + // `ptr` slot so the memory round-trip preserves provenance. The + // niche carrier backs the leaf that coincides with it; any further + // pointer leaf (the extra slice pointer in `split_at_mut_checked`'s + // `Option<(&mut [T], &mut [T])>`) gets its own fresh `ptr` slot. + // `ProvenanceLoss` means a pointer punned against non-pointer + // bits, which stays fail-closed; `OverArrayLeafBound` reports the + // same bounded rewrite contract as the payload storage gate. + let extra_pointer_claims = if has_pointer_overlap { + match analyze_pointer_overlap(ctx, offset, size, storage_ty, &colliding_claims) { + Ok(extra) => extra, + Err(PointerOverlapRejection::OverArrayLeafBound { required }) => { + return Err(anyhow::anyhow!( + "enum slot map: `{}` field {} overlaps pointer storage whose arrays are not supported above the bounded rewrite limit; rewrite requires {required} pointer conversions, supported bound is {MAX_ENUM_PAYLOAD_ARRAY_REWRITE_LEAVES}", + name, + flat + )); + } + Err(PointerOverlapRejection::ProvenanceLoss) => { + return Err(anyhow::anyhow!( + "enum slot map: `{}` has overlapping pointer and non-identical storage at byte {}; refusing to erase LLVM pointer provenance", + name, + offset + )); + } + } + } else { + Vec::new() + }; + let incoming_is_byte_faithful = llvm_type_is_byte_faithful(ctx, storage_ty); + let claims_are_byte_faithful = colliding_claims + .iter() + .all(|&&(_, _, claim_ty)| llvm_type_is_byte_faithful(ctx, claim_ty)); + if !incoming_is_byte_faithful || !claims_are_byte_faithful { + return Err(anyhow::anyhow!( + "enum slot map: `{}` field {} overlaps non-identical storage but its lowered type is not byte-faithful (for example, it may contain implicit padding); refusing a type-punned store", + name, + flat + )); + } + // Back each extra pointer leaf with its own `ptr` slot (the field + // itself stays slotless and round-trips through memory). Added after + // the byte-faithfulness gate so `colliding_claims`'s borrow of + // `claims` has ended before we push. + for leaf in extra_pointer_claims { + let ptr_ty = llvm_types::PointerType::get(ctx, leaf.address_space).into(); + claims.push((leaf.offset, leaf.size, ptr_ty)); + } + continue; + } + claims.push((offset, size, storage_ty)); + if field_gets_slot { + claim_of_field[flat] = Some(claims.len() - 1); + } + } + + // Phase 2: lay the slots down in byte order, filling every gap (and + // the tail) so the struct's size is exactly rustc's. One integer per gap + // where that is layout-neutral, else [N x i8]; see `make_enum_filler_type`. + let mut emit_order: Vec = (0..claims.len()).collect(); + emit_order.sort_by_key(|&ci| claims[ci].0); + let mut llvm_fields: Vec = Vec::new(); + let mut slot_of_claim: Vec = vec![0; claims.len()]; + let mut current_offset: u64 = 0; + for &ci in &emit_order { + let (offset, size, llvm_ty) = claims[ci]; + if current_offset < offset { + let filler = + make_enum_filler_type(ctx, current_offset, offset - current_offset, abi_align); + llvm_fields.push(filler); + current_offset = offset; + } + slot_of_claim[ci] = llvm_fields.len() as u32; + llvm_fields.push(llvm_ty); + current_offset += size; + } + if current_offset < total_size { + let filler = + make_enum_filler_type(ctx, current_offset, total_size - current_offset, abi_align); + llvm_fields.push(filler); + } + + // Sanity: the struct we just built must be exactly rustc's size. + // Arrays of enums step by this size, so a mismatch means every + // element after the first is read from the wrong place. That is a + // guaranteed miscompile, hence a hard error, not a debug check. + let (_end, natural_size, natural_align) = natural_struct_layout(ctx, &llvm_fields) + .ok_or_else(|| anyhow::anyhow!("enum slot map: `{}` has unsupported LLVM layout", name))?; + if natural_size != total_size { + return Err(anyhow::anyhow!( + "enum slot map: `{}` lowered to {} bytes but rustc says {}", + name, + natural_size, + total_size + )); + } + let required_align = abi_align.max(1); + if natural_align > required_align { + return Err(anyhow::anyhow!( + "enum slot map: `{}` lowered with alignment {} but rustc requires {}; explicit over-aligned enum storage is not yet supported", + name, + natural_align, + abi_align + )); + } + if natural_align < required_align { + // The byte claims alone can under-align the storage, e.g. when the + // only claim is an i8 niche carrier inside a 4-aligned enum. Raise + // the struct's alignment with a zero-length anchor field, the same + // mechanism union storage uses; it occupies no bytes, so every slot + // index simply shifts by one. + let anchor_int = IntegerType::get(ctx, (required_align * 8) as u32, Signedness::Signless); + let anchor: TypeHandle = llvm_types::ArrayType::get(ctx, anchor_int.into(), 0).into(); + llvm_fields.insert(0, anchor); + for slot in &mut slot_of_claim { + *slot += 1; + } + } + + let field_slots = claim_of_field + .into_iter() + .map(|c| c.map(|ci| slot_of_claim[ci])) + .collect(); + Ok(EnumSlotMap { + llvm_struct_ty: llvm_types::StructType::get_unnamed( + ctx, + (llvm_fields, llvm_types::StructLayout::Unpacked), + ) + .into(), + carrier_slot: carrier_claim.map(|claim| slot_of_claim[claim]), + carrier_llvm_ty: carrier_ty, + field_slots, + field_offsets: all_field_offsets, + field_llvm_types, + }) +} + +/// Convert a `MirEnumType` to its LLVM struct representation. +/// +/// Thin wrapper over [`build_enum_slot_map`], which explains the layout. +/// Any op that needs an index into the converted enum must take it from +/// the slot map, never compute it by hand. +pub(crate) fn convert_enum_to_llvm( + ctx: &mut Context, + ty: TypeHandle, +) -> Result { + Ok(build_enum_slot_map(ctx, ty)?.llvm_struct_ty) +} + +/// Return an enum name only when the dialect lacks rustc's physical layout. +/// All importer-produced Direct, Niche, Single, and Empty layouts are +/// byte-faithful; legacy `Unknown` values are rejected everywhere rather than +/// receiving a guessed internal representation. +pub(crate) fn enum_unmodeled_in_memory(ctx: &Context, ty: TypeHandle) -> Option { + let ty_ref = ty.deref(ctx); + let enum_ty = ty_ref.downcast_ref::()?; + (enum_ty.layout_kind == EnumLayoutKind::Unknown).then(|| enum_ty.name().to_string()) +} + +/// Search a kernel parameter's type for an enum the host and device +/// would disagree about (see [`enum_unmodeled_in_memory`]). +/// +/// The search looks everywhere host data can hide: behind pointers, +/// inside slices and arrays, in struct/tuple fields, and in other enums' +/// payloads. It returns the first offending enum's name. +/// +/// Kernel signatures are checked early for a focused diagnostic. Lowering also +/// rejects Unknown layouts for locals, globals, and physical operations, so no +/// guessed representation can escape through an internal-only path. +/// +/// `visited` breaks cycles through recursive types (`TypeHandle` is +/// interned, so equality is identity). +pub(crate) fn find_unmodeled_enum_in_abi( + ctx: &mut Context, + ty: TypeHandle, + visited: &mut Vec, +) -> Result, anyhow::Error> { + if visited.contains(&ty) { + return Ok(None); + } + visited.push(ty); + + if let Some(name) = enum_unmodeled_in_memory(ctx, ty) { + return Ok(Some(name)); + } + + let children: Vec = { + let ty_ref = ty.deref(ctx); + if let Some(p) = ty_ref.downcast_ref::() { + vec![p.pointee] + } else if let Some(s) = ty_ref.downcast_ref::() { + vec![s.element_ty] + } else if let Some(s) = ty_ref.downcast_ref::() { + let mut nested = vec![s.element_ty]; + nested.extend_from_slice(&s.space_tys); + nested + } else if let Some(a) = ty_ref.downcast_ref::() { + vec![a.element_ty] + } else if let Some(s) = ty_ref.downcast_ref::() { + s.field_types.clone() + } else if let Some(u) = ty_ref.downcast_ref::() { + u.field_types.clone() + } else if let Some(t) = ty_ref.downcast_ref::() { + t.get_types().to_vec() + } else if let Some(e) = ty_ref.downcast_ref::() { + e.all_field_types.clone() + } else { + vec![] + } + }; + + for child in children { + if let Some(name) = find_unmodeled_enum_in_abi(ctx, child, visited)? { + return Ok(Some(name)); + } + } + Ok(None) +} + +#[cfg(test)] +mod tests { + use super::super::llvm_type_contains_pointer_in_address_space; + use super::super::test_support::{llvm_int, make_ctx, mir_uint, struct_fields}; + use super::*; + use dialect_mir::types::{EnumEncoding, EnumVariant, MirArrayType, MirPtrType}; + use pliron::builtin::types::FP32Type; + + #[test] + fn enum_array_payload_keeps_exact_size_alignment_and_stride() { + let mut ctx = make_ctx(); + let tag = mir_uint(&mut ctx, 8); + let element = mir_uint(&mut ctx, 16); + let payload: TypeHandle = MirArrayType::get(&mut ctx, element, 3).into(); + let enum_ty: TypeHandle = MirEnumType::get_with_layout( + &mut ctx, + "ArrayPayload".into(), + tag, + vec![0, 1], + vec![ + EnumVariant::unit("Empty".into()), + EnumVariant::new_with_layout("Data".into(), vec![payload], vec![2], vec![6]), + ], + 0, + 8, + 2, + ) + .into(); + let map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); + assert_eq!(llvm_type_size_align(&ctx, map.llvm_struct_ty), Some((8, 2))); + + let array: TypeHandle = MirArrayType::get(&mut ctx, enum_ty, 5).into(); + let lowered = convert_type(&mut ctx, array).unwrap(); + assert_eq!(llvm_type_size_align(&ctx, lowered), Some((40, 2))); + } + + #[test] + fn enum_slot_map_rejects_partial_byte_integer_carriers() { + let mut ctx = make_ctx(); + let partial_discriminant = mir_uint(&mut ctx, 7); + let direct: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "PartialDirect".into(), + partial_discriminant, + vec![0, 1], + vec![EnumVariant::unit("A".into()), EnumVariant::unit("B".into())], + EnumEncoding { + tag_offset: 0, + total_size: 1, + abi_align: 1, + layout_kind: EnumLayoutKind::Direct, + carrier_kind: EnumCarrierKind::Integer, + carrier_width: 7, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + let error = build_enum_slot_map(&mut ctx, direct) + .err() + .expect("partial-byte Direct carrier must fail closed"); + assert!( + error.to_string().contains("whole number of bytes"), + "{error}" + ); + + let logical = mir_uint(&mut ctx, 8); + let payload = mir_uint(&mut ctx, 8); + let niche: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "PartialNiche".into(), + logical, + vec![0, 1], + vec![ + EnumVariant::unit("None".into()), + EnumVariant::new_with_layout("Some".into(), vec![payload], vec![0], vec![1]), + ], + EnumEncoding { + tag_offset: 0, + total_size: 1, + abi_align: 1, + layout_kind: EnumLayoutKind::Niche, + carrier_kind: EnumCarrierKind::Integer, + carrier_width: 7, + untagged_variant: 1, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + let error = build_enum_slot_map(&mut ctx, niche) + .err() + .expect("partial-byte Niche carrier must fail closed"); + assert!( + error.to_string().contains("whole number of bytes"), + "{error}" + ); + } + + #[test] + fn enum_slot_map_rejects_misaligned_inhabited_payload() { + let mut ctx = make_ctx(); + let tag = mir_uint(&mut ctx, 8); + let word = mir_uint(&mut ctx, 32); + let enum_ty: TypeHandle = MirEnumType::get_with_layout( + &mut ctx, + "PackedPayload".into(), + tag, + vec![0, 1], + vec![ + EnumVariant::unit("Empty".into()), + EnumVariant::new_with_layout("Data".into(), vec![word], vec![1], vec![4]), + ], + 0, + 8, + 4, + ) + .into(); + let error = build_enum_slot_map(&mut ctx, enum_ty) + .err() + .expect("misaligned payload must be rejected"); + assert!(error.to_string().contains("offset 1 is not aligned")); + } + + #[test] + fn enum_slot_map_uses_i8_storage_for_nonoverlapping_scalar_bool() { + let mut ctx = make_ctx(); + let tag = mir_uint(&mut ctx, 32); + let boolean: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Signless).into(); + let enum_ty: TypeHandle = MirEnumType::get_with_layout( + &mut ctx, + "DirectBool".into(), + tag, + vec![0, 1], + vec![ + EnumVariant::new_with_layout("A".into(), vec![boolean], vec![4], vec![1]), + EnumVariant::unit("B".into()), + ], + 0, + 8, + 4, + ) + .into(); + + let map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); + assert_eq!(map.field_slots, vec![None]); + let fields = struct_fields(&ctx, map.llvm_struct_ty); + assert_eq!(fields.len(), 3, "tag, canonical bool byte, tail pad"); + assert_eq!( + fields[1] + .deref(&ctx) + .downcast_ref::() + .map(IntegerType::width), + Some(8), + "Rust bool storage must be an i8 byte, never an LLVM i1 slot" + ); + } + + #[test] + fn enum_slot_map_allows_nonoverlapping_aggregate_padding_without_bool() { + let mut ctx = make_ctx(); + let tag = mir_uint(&mut ctx, 32); + let byte = mir_uint(&mut ctx, 8); + let word = mir_uint(&mut ctx, 32); + let padded: TypeHandle = MirTupleType::get(&mut ctx, vec![byte, word]).into(); + let lowered_padded = convert_type(&mut ctx, padded).unwrap(); + assert!( + !llvm_type_is_byte_faithful(&ctx, lowered_padded), + "the LLVM i8/i32 tuple has harmless implicit padding" + ); + let enum_ty: TypeHandle = MirEnumType::get_with_layout( + &mut ctx, + "PaddedPayload".into(), + tag, + vec![0, 1], + vec![ + EnumVariant::new_with_layout("Data".into(), vec![padded], vec![4], vec![8]), + EnumVariant::unit("Empty".into()), + ], + 0, + 12, + 4, + ) + .into(); + + let map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); + assert_eq!(map.field_slots, vec![Some(1)]); + assert_eq!( + llvm_type_size_align(&ctx, map.llvm_struct_ty), + Some((12, 4)) + ); + } + + #[test] + fn enum_slot_map_rejects_nonoverlapping_nested_bool_storage() { + let mut ctx = make_ctx(); + let tag = mir_uint(&mut ctx, 32); + let boolean: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Signless).into(); + let wrapper: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "BoolWrapper".into(), + vec!["value".into()], + vec![boolean], + vec![0], + vec![0], + 1, + 1, + ) + .into(); + let enum_ty: TypeHandle = MirEnumType::get_with_layout( + &mut ctx, + "DirectBoolWrapper".into(), + tag, + vec![0, 1], + vec![ + EnumVariant::new_with_layout("A".into(), vec![wrapper], vec![4], vec![1]), + EnumVariant::unit("B".into()), + ], + 0, + 8, + 4, + ) + .into(); + + // The wrapper claims its byte-faithful twin ({i8}) and stays + // slotless: construction canonicalizes the bool byte, extraction + // re-loads the original {i1} shape from the canonical byte. + let map = build_enum_slot_map(&mut ctx, enum_ty) + .expect("nested bool storage canonicalizes through its byte-faithful twin"); + assert_eq!(map.field_slots, vec![None]); + assert_eq!(llvm_type_size_align(&ctx, map.llvm_struct_ty), Some((8, 4))); + let struct_fields: Vec<_> = map + .llvm_struct_ty + .deref(&ctx) + .downcast_ref::() + .unwrap() + .fields() + .collect(); + assert!( + struct_fields + .iter() + .all(|field| !llvm_type_contains_i1(&ctx, *field)), + "enum storage must never contain physical i1" + ); + } + + #[test] + fn enum_slot_map_rejects_bool_hidden_by_union_byte_storage() { + let mut ctx = make_ctx(); + let tag = mir_uint(&mut ctx, 32); + let byte = mir_uint(&mut ctx, 8); + let boolean: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Signless).into(); + let union: TypeHandle = MirUnionType::get( + &mut ctx, + "BoolOrByte".into(), + vec!["flag".into(), "byte".into()], + vec![boolean, byte], + 1, + 1, + ) + .into(); + let lowered_union = convert_type(&mut ctx, union).unwrap(); + assert!( + !llvm_type_contains_i1(&ctx, lowered_union), + "the union's raw-byte carrier intentionally hides its semantic bool" + ); + let enum_ty: TypeHandle = MirEnumType::get_with_layout( + &mut ctx, + "DirectUnionBool".into(), + tag, + vec![0, 1], + vec![ + EnumVariant::new_with_layout("A".into(), vec![union], vec![4], vec![1]), + EnumVariant::unit("B".into()), + ], + 0, + 8, + 4, + ) + .into(); + + let error = build_enum_slot_map(&mut ctx, enum_ty) + .err() + .expect("a bool hidden by union byte storage must still fail closed"); + assert!( + error.to_string().contains("hidden behind a union"), + "{error}" + ); + } + + #[test] + fn enum_slot_map_does_not_descend_through_pointer_to_bool() { + let mut ctx = make_ctx(); + let tag = mir_uint(&mut ctx, 32); + let boolean: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Signless).into(); + let pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, boolean, false).into(); + let enum_ty: TypeHandle = MirEnumType::get_with_layout( + &mut ctx, + "PointerToBool".into(), + tag, + vec![0, 1], + vec![ + EnumVariant::new_with_layout("A".into(), vec![pointer], vec![8], vec![8]), + EnumVariant::unit("B".into()), + ], + 0, + 16, + 8, + ) + .into(); + + let map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); + assert!(map.field_slots[0].is_some()); + } + + #[test] + fn enum_slot_map_rejects_pointer_nonpointer_overlap_in_either_variant_order() { + for pointer_first in [false, true] { + let mut ctx = make_ctx(); + let discr = mir_uint(&mut ctx, 32); + let u8_ty = mir_uint(&mut ctx, 8); + let u64_ty = mir_uint(&mut ctx, 64); + let pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, u8_ty, false).into(); + let (first_name, first_ty, second_name, second_ty) = if pointer_first { + ("Ptr", pointer, "Bits", u64_ty) + } else { + ("Bits", u64_ty, "Ptr", pointer) + }; + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "PointerOrBits".into(), + discr, + vec![0, 1], + vec![ + EnumVariant::new_with_layout( + first_name.into(), + vec![first_ty], + vec![8], + vec![8], + ), + EnumVariant::new_with_layout( + second_name.into(), + vec![second_ty], + vec![8], + vec![8], + ), + ], + EnumEncoding { + tag_offset: 0, + total_size: 16, + abi_align: 8, + layout_kind: EnumLayoutKind::Direct, + carrier_kind: EnumCarrierKind::Integer, + carrier_width: 32, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + let error = build_enum_slot_map(&mut ctx, enum_ty) + .err() + .expect("pointer overlap must reject"); + assert!( + error.to_string().contains("pointer provenance"), + "pointer_first={pointer_first}: {error}" + ); + } + } + + #[test] + fn enum_slot_map_canonicalizes_overlapping_bool_wrapper_storage() { + let mut ctx = make_ctx(); + let logical = mir_uint(&mut ctx, 8); + let boolean: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Signless).into(); + let wrapper: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "BoolWrapper".into(), + vec!["value".into()], + vec![boolean], + vec![0], + vec![0], + 1, + 1, + ) + .into(); + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "MaybeBoolWrapper".into(), + logical, + vec![0, 1], + vec![ + EnumVariant::unit("None".into()), + EnumVariant::new_with_layout("Some".into(), vec![wrapper], vec![0], vec![1]), + ], + EnumEncoding { + tag_offset: 0, + total_size: 1, + abi_align: 1, + layout_kind: EnumLayoutKind::Niche, + carrier_kind: EnumCarrierKind::Integer, + carrier_width: 8, + niche_start: 2, + untagged_variant: 1, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + + // Option: the wrapper's byte-faithful twin ({i8}) + // shares the single canonical byte with the i8 niche carrier. The + // wrapper stays slotless; construction zero-extends its bool. + let map = build_enum_slot_map(&mut ctx, enum_ty) + .expect("a bool wrapper shares its canonical byte with the i8 carrier"); + assert_eq!(map.field_slots, vec![None]); + assert_eq!(llvm_type_size_align(&ctx, map.llvm_struct_ty), Some((1, 1))); + assert!( + !llvm_type_contains_i1(&ctx, map.llvm_struct_ty), + "enum storage must never contain physical i1" + ); + } + + #[test] + fn enum_slot_map_canonicalizes_bool_wrapper_overlap_in_either_variant_order() { + for wrapper_first in [false, true] { + let mut ctx = make_ctx(); + let tag = mir_uint(&mut ctx, 32); + let byte = mir_uint(&mut ctx, 8); + let boolean: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Signless).into(); + let wrapper: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "BoolWrapper".into(), + vec!["value".into()], + vec![boolean], + vec![0], + vec![0], + 1, + 1, + ) + .into(); + let (first, second) = if wrapper_first { + (wrapper, byte) + } else { + (byte, wrapper) + }; + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "BoolWrapperOrByte".into(), + tag, + vec![0, 1], + vec![ + EnumVariant::new_with_layout("First".into(), vec![first], vec![4], vec![1]), + EnumVariant::new_with_layout("Second".into(), vec![second], vec![4], vec![1]), + ], + EnumEncoding { + tag_offset: 0, + total_size: 8, + abi_align: 4, + layout_kind: EnumLayoutKind::Direct, + carrier_kind: EnumCarrierKind::Integer, + carrier_width: 32, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + + // The wrapper's canonical twin ({i8}) and the plain u8 variant + // share one byte-faithful byte, independent of declaration + // order; the wrapper stays slotless either way. + let map = build_enum_slot_map(&mut ctx, enum_ty) + .expect("canonical bool bytes may share storage with a u8 variant"); + assert_eq!(llvm_type_size_align(&ctx, map.llvm_struct_ty), Some((8, 4))); + assert!( + !llvm_type_contains_i1(&ctx, map.llvm_struct_ty), + "wrapper_first={wrapper_first}: enum storage must never contain physical i1" + ); + let wrapper_flat = if wrapper_first { 0 } else { 1 }; + assert_eq!(map.field_slots[wrapper_flat], None); + } + } + + #[test] + fn enum_slot_map_backs_pointer_beside_integer_niche_carrier() { + let mut ctx = make_ctx(); + let logical = mir_uint(&mut ctx, 8); + let u32_ty = mir_uint(&mut ctx, 32); + let pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, u32_ty, false).into(); + let payload: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "PointerThenNiche".into(), + vec!["pointer".into(), "niche".into()], + vec![pointer, u32_ty], + vec![0, 1], + vec![0, 8], + 16, + 8, + ) + .into(); + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "MaybePointerThenNiche".into(), + logical, + vec![0, 1], + vec![ + EnumVariant::unit("None".into()), + EnumVariant::new_with_layout("Some".into(), vec![payload], vec![0], vec![16]), + ], + EnumEncoding { + tag_offset: 8, + total_size: 16, + abi_align: 8, + layout_kind: EnumLayoutKind::Niche, + carrier_kind: EnumCarrierKind::Integer, + carrier_width: 32, + untagged_variant: 1, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + // The pointer at byte 0 never shares bytes with the integer niche + // carrier at byte 8, so it is representable: it gets its own `ptr` slot + // and the carrier stays an integer slot. `{ptr@0, i32@8, i32 filler}` + // -- the trailing 4 bytes are 4-aligned inside an 8-aligned enum, so + // they lower to one `i32` rather than `[4 x i8]`. + let map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); + assert_eq!(map.field_slots, vec![None]); + let lowered_pointer = convert_type(&mut ctx, pointer).unwrap(); + let carrier: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); + assert_eq!( + struct_fields(&ctx, map.llvm_struct_ty), + vec![lowered_pointer, carrier, llvm_int(&mut ctx, 32)] + ); + assert_eq!( + llvm_type_size_align(&ctx, map.llvm_struct_ty), + Some((16, 8)) + ); + } + + #[test] + fn enum_slot_map_keeps_plain_pointer_niche_in_one_pointer_slot() { + let mut ctx = make_ctx(); + let logical = mir_uint(&mut ctx, 8); + let u32_ty = mir_uint(&mut ctx, 32); + let pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, u32_ty, false).into(); + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "MaybePointer".into(), + logical, + vec![0, 1], + vec![ + EnumVariant::unit("None".into()), + EnumVariant::new_with_layout("Some".into(), vec![pointer], vec![0], vec![8]), + ], + EnumEncoding { + tag_offset: 0, + total_size: 8, + abi_align: 8, + layout_kind: EnumLayoutKind::Niche, + carrier_kind: EnumCarrierKind::Pointer, + carrier_width: 64, + untagged_variant: 1, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + let map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); + assert_eq!(map.carrier_slot, Some(0)); + assert_eq!(map.field_slots, vec![Some(0)]); + assert!( + struct_fields(&ctx, map.llvm_struct_ty)[0] + .deref(&ctx) + .is::() + ); + } + + #[test] + fn enum_slot_map_accepts_pointer_first_and_pointer_second_aggregate_niches() { + for pointer_first in [true, false] { + let mut ctx = make_ctx(); + let logical = mir_uint(&mut ctx, 64); + let index = mir_uint(&mut ctx, 64); + let pointee = mir_uint(&mut ctx, 32); + let pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, pointee, false).into(); + let payload_types = if pointer_first { + vec![pointer, index] + } else { + vec![index, pointer] + }; + let payload: TypeHandle = MirTupleType::get(&mut ctx, payload_types).into(); + let tag_offset = if pointer_first { 0 } else { 8 }; + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "MaybeIndexedRef".into(), + logical, + vec![0, 1], + vec![ + EnumVariant::unit("None".into()), + EnumVariant::new_with_layout("Some".into(), vec![payload], vec![0], vec![16]), + ], + EnumEncoding { + tag_offset, + total_size: 16, + abi_align: 8, + layout_kind: EnumLayoutKind::Niche, + carrier_kind: EnumCarrierKind::Pointer, + carrier_width: 64, + untagged_variant: 1, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + + let map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); + assert_eq!(map.field_slots, vec![None]); + assert_eq!(map.carrier_slot, Some(if pointer_first { 0 } else { 1 })); + assert_eq!( + llvm_type_size_align(&ctx, map.llvm_struct_ty), + Some((16, 8)) + ); + let lowered_pointer = convert_type(&mut ctx, pointer).unwrap(); + // The 8 bytes the pointer slot does not claim are 8-aligned inside + // an 8-aligned enum, so they lower to one `i64`, not `[8 x i8]`. + let filler = llvm_int(&mut ctx, 64); + let expected = if pointer_first { + vec![lowered_pointer, filler] + } else { + vec![filler, lowered_pointer] + }; + assert_eq!(struct_fields(&ctx, map.llvm_struct_ty), expected); + } + } + + #[test] + fn enum_slot_map_backs_each_aggregate_pointer_leaf_with_its_own_slot() { + let mut ctx = make_ctx(); + let logical = mir_uint(&mut ctx, 64); + let pointee = mir_uint(&mut ctx, 32); + let pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, pointee, false).into(); + let payload: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "TwoRefs".into(), + vec!["first".into(), "second".into()], + vec![pointer, pointer], + vec![0, 1], + vec![0, 8], + 16, + 8, + ) + .into(); + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "MaybeTwoRefs".into(), + logical, + vec![0, 1], + vec![ + EnumVariant::unit("None".into()), + EnumVariant::new_with_layout("Some".into(), vec![payload], vec![0], vec![16]), + ], + EnumEncoding { + tag_offset: 8, + total_size: 16, + abi_align: 8, + layout_kind: EnumLayoutKind::Niche, + carrier_kind: EnumCarrierKind::Pointer, + carrier_width: 64, + untagged_variant: 1, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + + // Both pointer leaves are representable without erasing provenance: the + // carrier backs the leaf at byte 8, and the leaf at byte 0 gets its own + // fresh `ptr` slot. The payload stays slotless and round-trips through + // `{ptr@0, ptr@8}`. + let map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); + assert_eq!(map.field_slots, vec![None]); + let lowered_pointer = convert_type(&mut ctx, pointer).unwrap(); + assert_eq!( + struct_fields(&ctx, map.llvm_struct_ty), + vec![lowered_pointer, lowered_pointer] + ); + assert_eq!( + llvm_type_size_align(&ctx, map.llvm_struct_ty), + Some((16, 8)) + ); + } + + #[test] + fn enum_slot_map_backs_split_at_mut_slice_pair() { + // Models `Option<(&mut [u32], &mut [u32])>`, the `split_at_mut_checked` + // return type: two fat slice pointers `{ptr, len}` back to back. The + // None niche lives in the first data pointer; both data pointers must + // keep their own `ptr` slot so provenance survives the memory + // round-trip, with the two `len` fields as raw byte storage. + let mut ctx = make_ctx(); + let logical = mir_uint(&mut ctx, 64); + let len = mir_uint(&mut ctx, 64); + let pointee = mir_uint(&mut ctx, 32); + let pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, pointee, false).into(); + let payload: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "SlicePair".into(), + vec![ + "a_ptr".into(), + "a_len".into(), + "b_ptr".into(), + "b_len".into(), + ], + vec![pointer, len, pointer, len], + vec![0, 1, 2, 3], + vec![0, 8, 16, 24], + 32, + 8, + ) + .into(); + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "MaybeSlicePair".into(), + logical, + vec![0, 1], + vec![ + EnumVariant::unit("None".into()), + EnumVariant::new_with_layout("Some".into(), vec![payload], vec![0], vec![32]), + ], + EnumEncoding { + tag_offset: 0, + total_size: 32, + abi_align: 8, + layout_kind: EnumLayoutKind::Niche, + carrier_kind: EnumCarrierKind::Pointer, + carrier_width: 64, + untagged_variant: 1, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + + let map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); + assert_eq!(map.field_slots, vec![None]); + let lowered_pointer = convert_type(&mut ctx, pointer).unwrap(); + // Each slice's length sits in the 8 bytes after its pointer, 8-aligned + // inside an 8-aligned enum, so both lower to one `i64`. This is the + // `split_at_mut_checked` shape, and the whole point of the widening: + // `{ptr, i64, ptr, i64}` moves two lengths as two values, where + // `{ptr, [8 x i8], ptr, [8 x i8]}` moved them as sixteen separate bytes. + let filler = llvm_int(&mut ctx, 64); + assert_eq!( + struct_fields(&ctx, map.llvm_struct_ty), + vec![lowered_pointer, filler, lowered_pointer, filler] + ); + assert_eq!( + llvm_type_size_align(&ctx, map.llvm_struct_ty), + Some((32, 8)) + ); + } + + #[test] + fn enum_slot_map_three_field_tuple_follows_recorded_offsets() { + // rustc lays out `(u32, f32, &T)` with the pointer first in memory: + // ptr @ 0, u32 @ 8, f32 @ 12, size 16. With those offsets recorded on + // the tuple, the pointer leaf coincides exactly with the niche + // carrier at byte 0 and the payload is accepted. + let mut ctx = make_ctx(); + let logical = mir_uint(&mut ctx, 64); + let word = mir_uint(&mut ctx, 32); + let float: TypeHandle = FP32Type::get(&ctx).into(); + let pointee = mir_uint(&mut ctx, 32); + let pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, pointee, false).into(); + let payload: TypeHandle = MirTupleType::get_with_layout( + &mut ctx, + vec![word, float, pointer], + vec![2, 0, 1], + vec![8, 12, 0], + 16, + 8, + ) + .into(); + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "MaybeMixedTuple".into(), + logical, + vec![0, 1], + vec![ + EnumVariant::unit("None".into()), + EnumVariant::new_with_layout("Some".into(), vec![payload], vec![0], vec![16]), + ], + EnumEncoding { + tag_offset: 0, + total_size: 16, + abi_align: 8, + layout_kind: EnumLayoutKind::Niche, + carrier_kind: EnumCarrierKind::Pointer, + carrier_width: 64, + untagged_variant: 1, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + + let map = build_enum_slot_map(&mut ctx, enum_ty) + .expect("recorded tuple offsets prove the pointer word matches the carrier"); + assert_eq!(map.carrier_slot, Some(0)); + assert_eq!( + llvm_type_size_align(&ctx, map.llvm_struct_ty), + Some((16, 8)) + ); + + // The same tuple with offsets that put the u32 over the pointer + // carrier (ptr @ 8 instead) must still fail closed: integer bytes + // may not alias pointer storage. + let mismatched_payload: TypeHandle = MirTupleType::get_with_layout( + &mut ctx, + vec![word, float, pointer], + vec![0, 1, 2], + vec![0, 4, 8], + 16, + 8, + ) + .into(); + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "MaybeMixedTupleShifted".into(), + logical, + vec![0, 1], + vec![ + EnumVariant::unit("None".into()), + EnumVariant::new_with_layout( + "Some".into(), + vec![mismatched_payload], + vec![0], + vec![16], + ), + ], + EnumEncoding { + tag_offset: 0, + total_size: 16, + abi_align: 8, + layout_kind: EnumLayoutKind::Niche, + carrier_kind: EnumCarrierKind::Pointer, + carrier_width: 64, + untagged_variant: 1, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + let error = build_enum_slot_map(&mut ctx, enum_ty) + .err() + .expect("a non-pointer word over the pointer carrier must fail closed"); + assert!(error.to_string().contains("pointer provenance"), "{error}"); + } + + #[test] + fn enum_slot_map_rejects_shifted_nested_pointer_carrier() { + let mut ctx = make_ctx(); + let logical = mir_uint(&mut ctx, 64); + let index = mir_uint(&mut ctx, 64); + let pointee = mir_uint(&mut ctx, 32); + let pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, pointee, false).into(); + let payload: TypeHandle = MirTupleType::get(&mut ctx, vec![index, pointer]).into(); + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "ShiftedPointerCarrier".into(), + logical, + vec![0, 1], + vec![ + EnumVariant::unit("None".into()), + EnumVariant::new_with_layout("Some".into(), vec![payload], vec![0], vec![16]), + ], + EnumEncoding { + tag_offset: 0, + total_size: 16, + abi_align: 8, + layout_kind: EnumLayoutKind::Niche, + carrier_kind: EnumCarrierKind::Pointer, + carrier_width: 64, + untagged_variant: 1, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + + let error = build_enum_slot_map(&mut ctx, enum_ty) + .err() + .expect("pointer leaves at different offsets must not be conflated"); + assert!(error.to_string().contains("pointer provenance"), "{error}"); + } + + #[test] + fn enum_slot_map_rejects_nested_carrier_address_space_mismatch() { + let mut ctx = make_ctx(); + let logical = mir_uint(&mut ctx, 64); + let index = mir_uint(&mut ctx, 64); + let pointee = mir_uint(&mut ctx, 32); + let generic_pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, pointee, false).into(); + let payload: TypeHandle = MirTupleType::get(&mut ctx, vec![generic_pointer, index]).into(); + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "MismatchedPointerSpace".into(), + logical, + vec![0, 1], + vec![ + EnumVariant::unit("None".into()), + EnumVariant::new_with_layout("Some".into(), vec![payload], vec![0], vec![16]), + ], + EnumEncoding { + tag_offset: 0, + total_size: 16, + abi_align: 8, + layout_kind: EnumLayoutKind::Niche, + carrier_kind: EnumCarrierKind::Pointer, + carrier_width: 64, + carrier_address_space: llvm_types::address_space::GLOBAL, + untagged_variant: 1, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + + let error = build_enum_slot_map(&mut ctx, enum_ty) + .err() + .expect("equal byte ranges in different address spaces must not alias"); + assert!(error.to_string().contains("pointer provenance"), "{error}"); + } + + #[test] + fn enum_slot_map_genericizes_nonoverlapping_shared_pointer_payload() { + let mut ctx = make_ctx(); + let discr = mir_uint(&mut ctx, 32); + let u32_ty = mir_uint(&mut ctx, 32); + let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, u32_ty, false).into(); + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "SharedPointerPayload".into(), + discr, + vec![0, 1], + vec![ + EnumVariant::unit("Unit".into()), + EnumVariant::new_with_layout("Ptr".into(), vec![shared], vec![8], vec![8]), + ], + EnumEncoding { + tag_offset: 0, + total_size: 16, + abi_align: 8, + layout_kind: EnumLayoutKind::Direct, + carrier_kind: EnumCarrierKind::Integer, + carrier_width: 32, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + let map = build_enum_slot_map(&mut ctx, enum_ty) + .expect("a direct shared-pointer payload should use generic physical storage"); + let field_slot = map.field_slots[0].expect("shared pointer field should own a slot"); + let stored_ty = map + .llvm_struct_ty + .deref(&ctx) + .downcast_ref::() + .map(|struct_ty| struct_ty.field_type(field_slot as usize)) + .expect("field slot must exist"); + assert_eq!( + stored_ty + .deref(&ctx) + .downcast_ref::() + .expect("shared pointer storage must remain a pointer") + .address_space(), + llvm_types::address_space::GENERIC, + "enum storage must use a target-stable generic pointer" + ); + assert_eq!( + map.field_llvm_types[0] + .deref(&ctx) + .downcast_ref::() + .expect("semantic field must remain a pointer") + .address_space(), + llvm_types::address_space::SHARED, + "payload operations must retain the semantic shared address space" + ); + } + + #[test] + fn enum_slot_map_genericizes_nested_shared_pointer_payload() { + let mut ctx = make_ctx(); + let discr = mir_uint(&mut ctx, 32); + let u32_ty = mir_uint(&mut ctx, 32); + let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, u32_ty, false).into(); + let wrapper: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "SharedPointerWrapper".into(), + vec!["pointer".into()], + vec![shared], + vec![0], + vec![0], + 8, + 8, + ) + .into(); + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "NestedSharedPointerPayload".into(), + discr, + vec![0, 1], + vec![ + EnumVariant::unit("Unit".into()), + EnumVariant::new_with_layout("Ptr".into(), vec![wrapper], vec![8], vec![8]), + ], + EnumEncoding { + tag_offset: 0, + total_size: 16, + abi_align: 8, + layout_kind: EnumLayoutKind::Direct, + carrier_kind: EnumCarrierKind::Integer, + carrier_width: 32, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + + let map = build_enum_slot_map(&mut ctx, enum_ty) + .expect("a struct-nested shared pointer should use generic physical storage"); + let field_slot = map.field_slots[0].expect("nested payload should own a slot"); + let stored_ty = map + .llvm_struct_ty + .deref(&ctx) + .downcast_ref::() + .map(|struct_ty| struct_ty.field_type(field_slot as usize)) + .expect("field slot must exist"); + assert!( + llvm_type_contains_pointer_in_address_space( + &ctx, + stored_ty, + llvm_types::address_space::GENERIC + ), + "physical nested payload must contain a generic pointer" + ); + assert!( + !llvm_type_contains_pointer_in_address_space( + &ctx, + stored_ty, + llvm_types::address_space::SHARED + ), + "physical nested payload must not retain the target-dependent AS3 pointer" + ); + assert!( + llvm_type_contains_pointer_in_address_space( + &ctx, + map.field_llvm_types[0], + llvm_types::address_space::SHARED + ), + "semantic payload type must retain shared address-space semantics" + ); + } + + #[test] + fn enum_slot_map_genericizes_bounded_shared_pointer_array_payload() { + let mut ctx = make_ctx(); + let discr = mir_uint(&mut ctx, 32); + let u32_ty = mir_uint(&mut ctx, 32); + let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, u32_ty, false).into(); + let pointers: TypeHandle = MirArrayType::get(&mut ctx, shared, 2).into(); + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "SharedPointerArrayPayload".into(), + discr, + vec![0, 1], + vec![ + EnumVariant::unit("Unit".into()), + EnumVariant::new_with_layout("Pointers".into(), vec![pointers], vec![8], vec![16]), + ], + EnumEncoding { + tag_offset: 0, + total_size: 24, + abi_align: 8, + layout_kind: EnumLayoutKind::Direct, + carrier_kind: EnumCarrierKind::Integer, + carrier_width: 32, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + + let map = build_enum_slot_map(&mut ctx, enum_ty) + .expect("a bounded shared-pointer array should use generic physical storage"); + let field_slot = map.field_slots[0].expect("bounded pointer array should own a slot"); + let stored_ty = map + .llvm_struct_ty + .deref(&ctx) + .downcast_ref::() + .map(|struct_ty| struct_ty.field_type(field_slot as usize)) + .expect("field slot must exist"); + assert!( + llvm_type_contains_pointer_in_address_space( + &ctx, + stored_ty, + llvm_types::address_space::GENERIC + ), + "physical array payload must contain generic pointers" + ); + assert!( + !llvm_type_contains_pointer_in_address_space( + &ctx, + stored_ty, + llvm_types::address_space::SHARED + ), + "physical array payload must not retain target-dependent AS3 pointers" + ); + assert!( + llvm_type_contains_pointer_in_address_space( + &ctx, + map.field_llvm_types[0], + llvm_types::address_space::SHARED + ), + "semantic array payload must retain shared address-space semantics" + ); + } + + #[test] + fn enum_slot_map_rejects_oversized_shared_pointer_array_payload() { + let mut ctx = make_ctx(); + let discr = mir_uint(&mut ctx, 32); + let u32_ty = mir_uint(&mut ctx, 32); + let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, u32_ty, false).into(); + let count = MAX_ENUM_PAYLOAD_ARRAY_REWRITE_LEAVES + 1; + let pointers: TypeHandle = MirArrayType::get(&mut ctx, shared, count).into(); + let payload_size = count * 8; + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "OversizedSharedPointerArrayPayload".into(), + discr, + vec![0, 1], + vec![ + EnumVariant::unit("Unit".into()), + EnumVariant::new_with_layout( + "Pointers".into(), + vec![pointers], + vec![8], + vec![payload_size], + ), + ], + EnumEncoding { + tag_offset: 0, + total_size: 8 + payload_size, + abi_align: 8, + layout_kind: EnumLayoutKind::Direct, + carrier_kind: EnumCarrierKind::Integer, + carrier_width: 32, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + + let error = build_enum_slot_map(&mut ctx, enum_ty) + .err() + .expect("an oversized shared-pointer array must remain fail-closed"); + assert!( + error + .to_string() + .contains("arrays containing shared-memory pointers are not supported"), + "{error}" + ); + assert!( + error.to_string().contains(&format!( + "supported bound is {MAX_ENUM_PAYLOAD_ARRAY_REWRITE_LEAVES}" + )), + "{error}" + ); + } + + /// A struct payload holding two shared-pointer arrays whose leaves exceed + /// the bound only in total. The slot map is target-mode agnostic (it is + /// built once, before the exporter picks the legacy 64-bit or modern + /// p3:32 data layout), so this one gate covers both output modes. + fn struct_of_shared_pointer_arrays_over_total_bound(ctx: &mut Context) -> TypeHandle { + let u32_ty = mir_uint(ctx, 32); + let shared: TypeHandle = MirPtrType::get_shared(ctx, u32_ty, false).into(); + let first: TypeHandle = MirArrayType::get(ctx, shared, 9).into(); + let second: TypeHandle = MirArrayType::get(ctx, shared, 8).into(); + MirStructType::get_with_full_layout( + ctx, + "SharedPointerArrayPair".into(), + vec!["first".into(), "second".into()], + vec![first, second], + vec![], + vec![0, 72], + 136, + 8, + ) + .into() + } + + #[test] + fn enum_slot_map_rejects_struct_of_shared_pointer_arrays_over_total_bound() { + // Each array alone (9 and 8 leaves) is within the bound; their total + // of 17 is not. The payload-root gate must reject the direct layout + // with the same bound diagnostic as a single oversized array. + let mut ctx = make_ctx(); + let discr = mir_uint(&mut ctx, 32); + let payload = struct_of_shared_pointer_arrays_over_total_bound(&mut ctx); + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "StructOfSharedPointerArraysPayload".into(), + discr, + vec![0, 1], + vec![ + EnumVariant::unit("Unit".into()), + EnumVariant::new_with_layout("Pointers".into(), vec![payload], vec![8], vec![136]), + ], + EnumEncoding { + tag_offset: 0, + total_size: 144, + abi_align: 8, + layout_kind: EnumLayoutKind::Direct, + carrier_kind: EnumCarrierKind::Integer, + carrier_width: 32, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + + let error = build_enum_slot_map(&mut ctx, enum_ty) + .err() + .expect("two in-bound arrays exceeding the bound in total must fail closed"); + assert!( + error + .to_string() + .contains("arrays containing shared-memory pointers are not supported"), + "{error}" + ); + assert!( + error.to_string().contains(&format!( + "rewrite requires 17 pointer conversions, supported bound is {MAX_ENUM_PAYLOAD_ARRAY_REWRITE_LEAVES}" + )), + "{error}" + ); + } + + #[test] + fn enum_slot_map_rejects_struct_of_shared_pointer_arrays_in_niche_layout() { + // The same over-bound payload in a niche layout must report the same + // bound diagnostic as the direct layout, not a pointer-provenance + // error from the overlap walk: the payload-root gate runs first in + // both layouts. + let mut ctx = make_ctx(); + let logical = mir_uint(&mut ctx, 64); + let payload = struct_of_shared_pointer_arrays_over_total_bound(&mut ctx); + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "NicheStructOfSharedPointerArrays".into(), + logical, + vec![0, 1], + vec![ + EnumVariant::unit("None".into()), + EnumVariant::new_with_layout("Some".into(), vec![payload], vec![0], vec![136]), + ], + EnumEncoding { + tag_offset: 0, + total_size: 136, + abi_align: 8, + layout_kind: EnumLayoutKind::Niche, + carrier_kind: EnumCarrierKind::Pointer, + carrier_width: 64, + untagged_variant: 1, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + + let error = build_enum_slot_map(&mut ctx, enum_ty) + .err() + .expect("the payload-root bound must also gate niche layouts"); + assert!( + error.to_string().contains(&format!( + "rewrite requires 17 pointer conversions, supported bound is {MAX_ENUM_PAYLOAD_ARRAY_REWRITE_LEAVES}" + )), + "{error}" + ); + assert!( + !error.to_string().contains("pointer provenance"), + "the bound diagnostic must not be masked by the provenance error: {error}" + ); + } + + #[test] + fn enum_slot_map_reports_bound_for_oversized_pointer_array_over_niche_carrier() { + // Generic pointers pass the shared-pointer storage gate untouched, so + // an oversized generic-pointer array reaches the overlap walk. Its + // exhausted expansion budget must surface the bound diagnostic, not + // the unrelated provenance error. + let mut ctx = make_ctx(); + let logical = mir_uint(&mut ctx, 64); + let pointee = mir_uint(&mut ctx, 32); + let pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, pointee, false).into(); + let count = MAX_ENUM_PAYLOAD_ARRAY_REWRITE_LEAVES + 1; + let pointers: TypeHandle = MirArrayType::get(&mut ctx, pointer, count).into(); + let payload_size = count * 8; + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "NicheOversizedGenericPointerArray".into(), + logical, + vec![0, 1], + vec![ + EnumVariant::unit("None".into()), + EnumVariant::new_with_layout( + "Some".into(), + vec![pointers], + vec![0], + vec![payload_size], + ), + ], + EnumEncoding { + tag_offset: 0, + total_size: payload_size, + abi_align: 8, + layout_kind: EnumLayoutKind::Niche, + carrier_kind: EnumCarrierKind::Pointer, + carrier_width: 64, + untagged_variant: 1, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + + let error = build_enum_slot_map(&mut ctx, enum_ty) + .err() + .expect("an oversized pointer-array expansion over a niche carrier must fail closed"); + assert!( + error.to_string().contains(&format!( + "rewrite requires {count} pointer conversions, supported bound is {MAX_ENUM_PAYLOAD_ARRAY_REWRITE_LEAVES}" + )), + "{error}" + ); + assert!( + !error.to_string().contains("pointer provenance"), + "the bound diagnostic must not be masked by the provenance error: {error}" + ); + } + + #[test] + fn enum_slot_map_rejects_shared_pointer_vector_payload() { + let mut ctx = make_ctx(); + let tag = mir_uint(&mut ctx, 32); + let shared_pointer: TypeHandle = + llvm_types::PointerType::get(&ctx, llvm_types::address_space::SHARED).into(); + let shared_vector: TypeHandle = + llvm_types::VectorType::get(&ctx, shared_pointer, 2, llvm_types::VectorTypeKind::Fixed) + .into(); + let enum_ty: TypeHandle = MirEnumType::get_with_layout( + &mut ctx, + "SharedPointerVector".into(), + tag, + vec![0, 1], + vec![ + EnumVariant::new_with_layout( + "Data".into(), + vec![shared_vector], + vec![16], + vec![16], + ), + EnumVariant::unit("Empty".into()), + ], + 0, + 32, + 16, + ) + .into(); + + let error = build_enum_slot_map(&mut ctx, enum_ty) + .err() + .expect("a vector of shared pointers must fail closed"); + assert!( + error.to_string().contains("shared-memory pointer"), + "{error}" + ); + } +} diff --git a/crates/mir-lower/src/convert/types/func_abi.rs b/crates/mir-lower/src/convert/types/func_abi.rs new file mode 100644 index 0000000000..a5f61c909c --- /dev/null +++ b/crates/mir-lower/src/convert/types/func_abi.rs @@ -0,0 +1,1069 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Kernel-boundary detection, transparent-scalar ABI, and function-type +//! conversion (the boundary ABI: slice flattening, struct flattening, +//! packed shared-pointer rewrites). + +use dialect_mir::types::{ + MirArrayType, MirDisjointSliceType, MirFP16Type, MirPtrType, MirSliceType, MirStructType, + MirTupleType, +}; +use llvm_export::types as llvm_types; +use llvm_export::types::PointerTypeExt; +use pliron::builtin::type_interfaces::FunctionTypeInterface; +use pliron::builtin::types::{FP32Type, FP64Type, FunctionType, IntegerType, Signedness}; +use pliron::context::{Context, Ptr}; +use pliron::operation::Operation; +use pliron::r#type::{TypeHandle, type_cast}; + +use super::{ + StructLayoutInfo, build_struct_slot_map, convert_type, is_zero_sized_type, + llvm_packed_struct_contains_pointer_in_address_space, llvm_type_size_align, +}; +use crate::convert::target_stable_storage::{StorageRewriteOptions, target_stable_storage_type}; + +// ============================================================================= +// Kernel-Boundary Detection +// ============================================================================= + +/// Identifier of the attribute that marks a `MirFuncOp` / `llvm.func` as a +/// GPU kernel entry point. +/// +/// Kept as a function (rather than a `const`) because pliron `Identifier` +/// construction needs the `try_into()` fallible path. +fn gpu_kernel_attr() -> pliron::identifier::Identifier { + "gpu_kernel".try_into().expect("static identifier") +} + +/// Returns `true` when `op` carries the `gpu_kernel` attribute. +/// +/// The kernel-entry ABI differs from internal device-function ABI: at +/// kernel boundaries, aggregate parameters (structs, closures) are passed +/// as a single byval value to match what the host pushes via +/// `cuLaunchKernel`. Internal call sites still flatten aggregates the +/// same way they always did. This helper is the single source of truth +/// for that branch and is consumed by both [`convert_function_type`] and +/// the entry-block prologue in `lowering.rs`. +pub fn is_kernel_func(ctx: &Context, op: Ptr) -> bool { + op.deref(ctx) + .attributes + .get::(&gpu_kernel_attr()) + .is_some() +} + +/// Return the declaration index and type of the single non-ZST field of a +/// rustc-proven transparent scalar struct. +/// +/// The importer marks the outer ABI from rustc rather than inferring it from +/// source field count. We still validate the MIR shape here so malformed or +/// hand-written dialect input cannot turn an arbitrary aggregate into a scalar +/// ABI value. +fn transparent_scalar_field_with_index( + ctx: &mut Context, + struct_ty: TypeHandle, +) -> Result<(usize, TypeHandle), anyhow::Error> { + let (name, field_types, mem_to_decl, is_transparent_scalar) = { + let ty_ref = struct_ty.deref(ctx); + let s = ty_ref + .downcast_ref::() + .ok_or_else(|| anyhow::anyhow!("transparent scalar ABI requires a MirStructType"))?; + ( + s.name.clone(), + s.field_types.clone(), + s.memory_order(), + s.is_transparent_scalar(), + ) + }; + + if !is_transparent_scalar { + return Err(anyhow::anyhow!( + "struct `{}` is not marked as a transparent scalar", + name + )); + } + + let mut scalar_field = None; + for decl_idx in mem_to_decl { + let field_ty = field_types[decl_idx]; + let converted = convert_type(ctx, field_ty)?; + if is_zero_sized_type(ctx, converted) { + continue; + } + if scalar_field.replace((decl_idx, field_ty)).is_some() { + return Err(anyhow::anyhow!( + "transparent scalar struct `{}` has more than one non-ZST field", + name + )); + } + } + + scalar_field + .ok_or_else(|| anyhow::anyhow!("transparent scalar struct `{}` has no non-ZST field", name)) +} + +/// Return the single non-ZST field of a rustc-proven transparent scalar struct. +pub(crate) fn transparent_scalar_field( + ctx: &mut Context, + struct_ty: TypeHandle, +) -> Result { + Ok(transparent_scalar_field_with_index(ctx, struct_ty)?.1) +} + +/// One aggregate layer traversed when a transparent scalar wrapper crosses an +/// ABI boundary. `field_slot` is the LLVM struct slot containing the next +/// nested wrapper or the final scalar. +#[derive(Clone, Copy, Debug)] +pub(crate) struct TransparentScalarLayer { + pub llvm_struct_ty: TypeHandle, + pub field_slot: u32, +} + +/// Complete ABI projection for a rustc-proven transparent scalar wrapper. +/// +/// `layers` are ordered outermost to innermost. A return lowers by extracting +/// those slots in order; a call result is rebuilt by inserting the scalar into +/// the same layers in reverse order. +#[derive(Clone, Debug)] +pub(crate) struct TransparentScalarAbiInfo { + pub scalar_ty: TypeHandle, + pub layers: Vec, +} + +/// Build the scalar ABI projection, including the exact LLVM slot used at every +/// wrapper layer. +/// +/// Slot indices come from [`build_struct_slot_map`], so ZST markers, explicit +/// padding, and rustc memory order cannot make the return/call paths disagree +/// with ordinary aggregate conversion. +pub(crate) fn transparent_scalar_abi_info( + ctx: &mut Context, + struct_ty: TypeHandle, +) -> Result { + let mut current = struct_ty; + let mut layers = Vec::new(); + + loop { + let (decl_idx, field_ty) = transparent_scalar_field_with_index(ctx, current)?; + let layout = { + let ty_ref = current.deref(ctx); + let s = ty_ref.downcast_ref::().ok_or_else(|| { + anyhow::anyhow!("transparent scalar ABI requires a MirStructType") + })?; + StructLayoutInfo::of_struct(s) + }; + let map = build_struct_slot_map(ctx, &layout)?; + let field_slot = map.decl_to_llvm[decl_idx].ok_or_else(|| { + anyhow::anyhow!( + "transparent scalar field {} unexpectedly lowered as a ZST", + decl_idx + ) + })?; + layers.push(TransparentScalarLayer { + llvm_struct_ty: map.llvm_struct_ty, + field_slot, + }); + + let nested_transparent = { + let field_ref = field_ty.deref(ctx); + field_ref + .downcast_ref::() + .is_some_and(MirStructType::is_transparent_scalar) + }; + if nested_transparent { + current = field_ty; + continue; + } + + let scalar_ty = convert_type(ctx, field_ty)?; + if is_zero_sized_type(ctx, scalar_ty) { + return Err(anyhow::anyhow!( + "transparent scalar ABI resolved to a zero-sized field" + )); + } + return Ok(TransparentScalarAbiInfo { scalar_ty, layers }); + } +} + +/// LLVM ABI type for a rustc-proven transparent scalar struct. +/// +/// Transparent wrappers can nest (`Outer(Inner(u32))`). rustc still reports +/// the outer ADT as one scalar, so recurse through transparent scalar fields +/// until reaching the actual scalar/pointer representation. +pub(crate) fn transparent_scalar_llvm_type( + ctx: &mut Context, + struct_ty: TypeHandle, +) -> Result { + Ok(transparent_scalar_abi_info(ctx, struct_ty)?.scalar_ty) +} + +/// Target-stable ABI projection for packed aggregates containing AS3 leaves +/// across the internal device return boundary. +/// +/// The function body continues to use the semantic packed LLVM struct with +/// shared pointers. Only the physical return value recursively replaces AS3 +/// leaves with generic pointers, whose width is stable across the legacy/PTX +/// and modern NVVM data layouts. +#[derive(Clone, Copy, Debug)] +pub(crate) struct PackedSharedInternalAbiInfo { + pub semantic_ty: TypeHandle, + pub storage_ty: TypeHandle, +} + +/// Maximum number of shared-pointer leaves introduced by fixed-array expansion +/// in one packed-AS3 internal return carrier. +/// +/// Struct/tuple nesting and direct pointer leaves remain proportional to source +/// structure. Arrays can encode an arbitrarily large number of per-element +/// extract/cast/insert sequences compactly, so only array-expanded AS3 leaves +/// count against this code-shape budget. +pub(crate) const MAX_PACKED_SHARED_INTERNAL_ABI_ARRAY_REWRITE_LEAVES: u64 = 16; + +/// Whether a MIR type converts to a zero-sized LLVM type. +/// +/// MIR-side mirror of [`is_zero_sized_type`]: zero-length arrays, arrays of +/// zero-sized elements, and structs/tuples whose fields are all zero-sized +/// (including empty ones such as `PhantomData`) vanish at the LLVM level and +/// must not affect ABI-lane classification. +fn mir_type_is_zero_sized(ctx: &Context, ty: TypeHandle) -> bool { + let ty_ref = ty.deref(ctx); + if let Some(array_ty) = ty_ref.downcast_ref::() { + return array_ty.size() == 0 || mir_type_is_zero_sized(ctx, array_ty.element_type()); + } + if let Some(struct_ty) = ty_ref.downcast_ref::() { + return struct_ty + .field_types + .iter() + .all(|field| mir_type_is_zero_sized(ctx, *field)); + } + if let Some(tuple_ty) = ty_ref.downcast_ref::() { + return tuple_ty + .get_types() + .iter() + .all(|field| mir_type_is_zero_sized(ctx, *field)); + } + false +} + +/// Whether a MIR value shape belongs to the recursive packed-AS3 ABI lane. +/// +/// Structs, tuples, and fixed arrays may nest recursively and may contain any +/// number of scalar pointer leaves. Zero-sized fields are skipped, matching the +/// post-conversion field scan this predicate replaced. The array-specific +/// expansion budget is enforced after conversion from the storage rewrite's +/// exact AS3 leaf count. Vectors and unrelated aggregate kinds remain +/// deliberately fail-closed. +fn packed_shared_internal_abi_mir_shape_is_supported(ctx: &Context, mir_ty: TypeHandle) -> bool { + let children = { + let ty_ref = mir_ty.deref(ctx); + if ty_ref.is::() + || ty_ref.is::() + || ty_ref.is::() + || ty_ref.is::() + || ty_ref.is::() + || ty_ref.is::() + || ty_ref.is::() + { + return true; + } + if let Some(struct_ty) = ty_ref.downcast_ref::() { + Some(struct_ty.field_types.clone()) + } else if let Some(tuple_ty) = ty_ref.downcast_ref::() { + Some(tuple_ty.get_types().to_vec()) + } else { + ty_ref + .downcast_ref::() + .map(|array_ty| vec![array_ty.element_type()]) + } + }; + + children.is_some_and(|children| { + children.into_iter().all(|child| { + mir_type_is_zero_sized(ctx, child) + || packed_shared_internal_abi_mir_shape_is_supported(ctx, child) + }) + }) +} + +/// Recognize a recursive packed-AS3 internal ABI shape. +/// +/// The root must remain a byte-faithful packed struct. Nested structs/tuples, +/// multiple AS3 leaves, and bounded fixed arrays are admitted. Vectors and +/// unrelated aggregate kinds remain out of scope. The target-stable storage +/// utility owns the recursive AS3 -> generic rewrite and this classifier only +/// decides which semantic shapes may use it. +pub(crate) fn packed_shared_internal_abi_info( + ctx: &mut Context, + mir_ty: TypeHandle, +) -> Result, anyhow::Error> { + let layout = { + let ty_ref = mir_ty.deref(ctx); + let Some(struct_ty) = ty_ref.downcast_ref::() else { + return Ok(None); + }; + StructLayoutInfo::of_struct(struct_ty) + }; + + if !packed_shared_internal_abi_mir_shape_is_supported(ctx, mir_ty) { + return Ok(None); + } + + let map = build_struct_slot_map(ctx, &layout)?; + if !map.by_value_layout_faithful { + return Ok(None); + } + let is_packed = map + .llvm_struct_ty + .deref(ctx) + .downcast_ref::() + .is_some_and(|struct_ty| struct_ty.layout() == llvm_types::StructLayout::Packed); + if !is_packed { + return Ok(None); + } + + let rewrite = target_stable_storage_type( + ctx, + map.llvm_struct_ty, + StorageRewriteOptions { + canonicalize_bool: false, + }, + "packed shared internal ABI", + )?; + if rewrite.shared_pointer_leaves == 0 + || rewrite.array_shared_pointer_leaves > MAX_PACKED_SHARED_INTERNAL_ABI_ARRAY_REWRITE_LEAVES + { + return Ok(None); + } + let Some((storage_size, _)) = llvm_type_size_align(ctx, rewrite.ty) else { + return Ok(None); + }; + if layout.total_size > 0 && storage_size != layout.total_size { + return Ok(None); + } + + Ok(Some(PackedSharedInternalAbiInfo { + semantic_ty: map.llvm_struct_ty, + storage_ty: rewrite.ty, + })) +} + +/// Convert a type that crosses a function boundary as one LLVM value. +/// +/// A struct with a natural-layout divergence is legal by value only when +/// [`build_struct_slot_map`] proved that a sequential packed LLVM struct +/// reproduces rustc's offsets and size. This keeps overlapping/union-like +/// legacy struct models fail-closed while allowing real `repr(packed)` values. +/// Packed values containing shared-memory pointers remain target-dependent +/// because AS3 pointer width differs between modern NVVM and PTX/legacy modes. +fn convert_by_value_abi_type( + ctx: &mut Context, + mir_ty: TypeHandle, + role: &str, +) -> Result { + let layout = { + let ty_ref = mir_ty.deref(ctx); + ty_ref + .downcast_ref::() + .map(StructLayoutInfo::of_struct) + }; + let llvm_ty = if let Some(layout) = layout { + let map = build_struct_slot_map(ctx, &layout)?; + if !map.by_value_layout_faithful { + return Err(anyhow::anyhow!( + "{} has a rustc struct layout that cannot be represented by an LLVM struct value", + role + )); + } + map.llvm_struct_ty + } else { + convert_type(ctx, mir_ty)? + }; + + if llvm_packed_struct_contains_pointer_in_address_space( + ctx, + llvm_ty, + llvm_types::address_space::SHARED, + ) { + return Err(anyhow::anyhow!( + "{} contains a packed aggregate with a target-dependent shared-memory pointer", + role + )); + } + + Ok(llvm_ty) +} + +/// Convert a MIR function type to an LLVM function type. +/// +/// This handles the ABI-level transformations required for GPU kernels. +/// The transformations ensure that the generated LLVM IR matches the +/// C ABI expected by the CUDA runtime. +/// +/// # ABI Transformations +/// +/// ## Argument Flattening +/// +/// Aggregate types are flattened to primitive types: +/// +/// ```text +/// MIR: fn kernel(slice: &[f32], point: Point) +/// LLVM: fn internal_fn(ptr: !ptr, len: i64, x: f32, y: f32) +/// ``` +/// +/// | MIR Argument | Internal call ABI | Kernel-entry ABI | +/// |-------------------------|-------------------------|------------------------| +/// | `&[T]` | `(ptr, i64)` | `(ptr, i64)` | +/// | `DisjointSlice` | `(ptr, i64)` | `(ptr, i64)` | +/// | `struct { a: A, b: B }` | `(a: A', b: B')` | one byval `{A', B'}` | +/// | `repr(transparent)` scalar ADT | underlying scalar | underlying scalar | +/// | closure with N captures | N separate field args | one byval struct | +/// | Other | Converted type | Converted type | +/// +/// Slices keep their `(ptr, len)` flattening on both sides because the +/// host-side launch helpers push the pointer and length as two driver +/// args. Structs and closures are unflattened only at kernel boundaries +/// because the host pushes them as a single scalar — see +/// `cuda_host::push_kernel_scalar`. The exception is a rustc-proven +/// `#[repr(transparent)]` `ValueAbi::Scalar` struct: its single non-ZST +/// field is the kernel parameter, matching the source type's transparent ABI. +/// Internal device-side call sites stay flattened: caller and callee are both +/// inside this backend, so the ABI is private and there is no host to disagree +/// with. +/// +/// ## Return Type Handling +/// +/// - Empty tuple `()` becomes `void` +/// - Empty struct `struct {}` becomes `void` +/// - Rustc-proven `repr(transparent)` scalar ADTs return the underlying scalar +/// - Other types are converted normally +/// +/// # Arguments +/// +/// * `ctx` - The pliron context +/// * `func_type` - The MIR function type to convert +/// * `is_kernel_entry` - When `true`, treat aggregate (non-slice) params +/// as single byval values to match the host-side push ABI. When `false`, +/// keep the existing internal device-fn ABI that flattens struct fields +/// into individual scalars. +/// +/// # Returns +/// +/// The equivalent LLVM function type with ABI transformations applied. +/// +/// # Example +/// +/// ```text +/// MIR: fn foo(a: &[f32], b: i32) -> f32 +/// LLVM: fn foo(ptr, i64, i32) -> f32 +/// +/// MIR: fn bar() -> () +/// LLVM: fn bar() -> void +/// ``` +/// +/// # Note +/// +/// At internal device-function boundaries the struct flattening must be +/// reversed in the entry block. At kernel-entry boundaries ordinary structs +/// arrive as one byval aggregate, while rustc-proven transparent scalar +/// wrappers arrive as their single non-ZST scalar field and must be rebuilt. +/// See `lowering.rs::build_entry_prologue` for the reconstruction paths. +pub fn convert_function_type( + ctx: &mut Context, + func_type: pliron::r#type::TypedHandle, + is_kernel_entry: bool, +) -> Result, anyhow::Error> { + // Extract input/output types before mutating context + let (inputs_ptr, results_ptr) = { + let func_ty_ref = func_type.deref(ctx); + let interface = type_cast::(&*func_ty_ref) + .ok_or_else(|| anyhow::anyhow!("Type does not implement FunctionTypeInterface"))?; + (interface.arg_types(), interface.res_types()) + }; + + // Convert inputs, flattening slice/struct types for ABI compatibility. + // Slices flatten on both ABIs; structs flatten only on the internal + // device-fn ABI. + let mut inputs = Vec::new(); + let inputs_vec: Vec<_> = inputs_ptr.to_vec(); + + for t in inputs_vec { + // Determine what kind of flattening this type needs + // Extract all info first, then drop the borrow + enum FlattenKind { + /// `{ ptr, len }`, then one parameter per index-space layout field. + Slice { + space_tys: Vec, + }, + Struct { + field_types: Vec, + mem_to_decl: Vec, + }, + TransparentScalar(TypeHandle), + None, + } + + let flatten_kind = { + let ty_ref = t.deref(ctx); + if ty_ref.is::() { + FlattenKind::Slice { + space_tys: Vec::new(), + } + } else if let Some(slice_ty) = ty_ref.downcast_ref::() { + FlattenKind::Slice { + space_tys: slice_ty.space_tys.clone(), + } + } else if let Some(struct_ty) = ty_ref.downcast_ref::() { + if is_kernel_entry && struct_ty.is_transparent_scalar() { + // rustc proves that this `repr(transparent)` ADT has a + // scalar ABI. Emit exactly its one non-ZST field as the + // kernel parameter instead of an aggregate `.b8[]` param. + FlattenKind::TransparentScalar(t) + } else if is_kernel_entry { + // Ordinary kernel-boundary structs remain intact so the + // host's single `push_kernel_scalar(&closure)` push + // matches a single aggregate .param entry. + FlattenKind::None + } else { + FlattenKind::Struct { + field_types: struct_ty.field_types.clone(), + mem_to_decl: struct_ty.memory_order(), + } + } + } else { + FlattenKind::None + } + }; + + match flatten_kind { + FlattenKind::Slice { space_tys } => { + let ptr_ty = llvm_types::PointerType::get_generic(ctx); + let len_ty = IntegerType::get(ctx, 64, Signedness::Signless); + inputs.push(ptr_ty.into()); + inputs.push(len_ty.into()); + for space_ty in space_tys { + let converted = convert_type(ctx, space_ty)?; + // A zero-sized index-space field contributes no parameter, + // as NVPTX has no empty `.param`. + if !is_zero_sized_type(ctx, converted) { + inputs.push(converted); + } + } + } + FlattenKind::Struct { + field_types, + mem_to_decl, + } => { + // Flatten in MEMORY ORDER to match struct layout + for mem_idx in 0..field_types.len() { + let decl_idx = mem_to_decl[mem_idx]; + let converted = convert_by_value_abi_type( + ctx, + field_types[decl_idx], + "by-value function argument", + )?; + // Skip ZST fields - NVPTX can't handle empty params + if !is_zero_sized_type(ctx, converted) { + inputs.push(converted); + } + } + } + FlattenKind::TransparentScalar(struct_ty) => { + inputs.push(transparent_scalar_llvm_type(ctx, struct_ty)?); + } + FlattenKind::None => { + let converted = convert_by_value_abi_type(ctx, t, "by-value function argument")?; + // Skip ZST args - NVPTX can't handle empty params + if !is_zero_sized_type(ctx, converted) { + inputs.push(converted); + } + } + } + } + + // Convert return type. A rustc-proven transparent scalar wrapper uses + // the underlying scalar at the function ABI boundary, while its body + // continues to use the ordinary converted aggregate representation. + let ret_ty = if results_ptr.is_empty() { + llvm_types::VoidType::get(ctx).into() + } else { + let mir_ret_ty = results_ptr[0]; + let is_transparent_scalar = { + let ty_ref = mir_ret_ty.deref(ctx); + ty_ref + .downcast_ref::() + .is_some_and(MirStructType::is_transparent_scalar) + }; + let ty = if is_transparent_scalar { + transparent_scalar_llvm_type(ctx, mir_ret_ty)? + } else if !is_kernel_entry { + match packed_shared_internal_abi_info(ctx, mir_ret_ty)? { + Some(abi) => abi.storage_ty, + None => convert_by_value_abi_type(ctx, mir_ret_ty, "by-value function return")?, + } + } else { + convert_by_value_abi_type(ctx, mir_ret_ty, "by-value function return")? + }; + // Check if zero-sized (empty struct or struct with only ZST fields). + if is_zero_sized_type(ctx, ty) { + llvm_types::VoidType::get(ctx).into() + } else { + ty + } + }; + + Ok(llvm_types::FuncType::get(ctx, ret_ty, inputs, false)) +} + +#[cfg(test)] +mod tests { + use super::super::llvm_type_contains_pointer_in_address_space; + use super::super::test_support::{make_ctx, mir_uint, struct_fields, transparent_u32}; + use super::*; + use dialect_mir::types::StructAbiKind; + + #[test] + fn transparent_scalar_return_type_uses_underlying_scalar() { + let mut ctx = make_ctx(); + let wrapper = transparent_u32(&mut ctx, "Scalar"); + let func_ty = FunctionType::get(&ctx, vec![], vec![wrapper]); + + let lowered = convert_function_type(&mut ctx, func_ty, false).unwrap(); + let result_ty = lowered.deref(&ctx).result_type(); + let result_ty_ref = result_ty.deref(&ctx); + let integer = result_ty_ref + .downcast_ref::() + .expect("transparent u32 return must lower to an integer"); + assert_eq!(integer.width(), 32); + } + + #[test] + fn nested_transparent_scalar_abi_records_each_rebuild_layer() { + let mut ctx = make_ctx(); + let inner = transparent_u32(&mut ctx, "Inner"); + let outer: TypeHandle = MirStructType::get_with_full_layout_and_abi( + &mut ctx, + "Outer".into(), + vec!["inner".into()], + vec![inner], + vec![0], + vec![0], + 4, + 4, + StructAbiKind::TransparentScalar, + ) + .into(); + + let info = transparent_scalar_abi_info(&mut ctx, outer).unwrap(); + let scalar_ty_ref = info.scalar_ty.deref(&ctx); + let integer = scalar_ty_ref + .downcast_ref::() + .expect("nested transparent wrapper must resolve to u32"); + assert_eq!(integer.width(), 32); + assert_eq!(info.layers.len(), 2); + assert_eq!(info.layers[0].field_slot, 0); + assert_eq!(info.layers[1].field_slot, 0); + } + + #[test] + fn ordinary_one_field_return_remains_aggregate() { + let mut ctx = make_ctx(); + let u32_ty = mir_uint(&mut ctx, 32); + let ordinary: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "Ordinary".into(), + vec!["value".into()], + vec![u32_ty], + vec![0], + vec![0], + 4, + 4, + ) + .into(); + let func_ty = FunctionType::get(&ctx, vec![], vec![ordinary]); + + let lowered = convert_function_type(&mut ctx, func_ty, false).unwrap(); + assert!( + lowered + .deref(&ctx) + .result_type() + .deref(&ctx) + .is::(), + "ordinary one-field structs must not be scalarized" + ); + } + + #[test] + fn packed_shared_pointer_predicate_ignores_unpacked_direct_pointer() { + let ctx = make_ctx(); + let shared: TypeHandle = + llvm_types::PointerType::get(&ctx, llvm_types::address_space::SHARED).into(); + let unpacked: TypeHandle = llvm_types::StructType::get_unnamed( + &ctx, + (vec![shared], llvm_types::StructLayout::Unpacked), + ) + .into(); + assert!( + !llvm_packed_struct_contains_pointer_in_address_space( + &ctx, + unpacked, + llvm_types::address_space::SHARED, + ), + "a direct AS3 pointer in an unpacked aggregate is not a packed physical-image hazard" + ); + + let packed: TypeHandle = llvm_types::StructType::get_unnamed( + &ctx, + (vec![shared], llvm_types::StructLayout::Packed), + ) + .into(); + assert!( + llvm_packed_struct_contains_pointer_in_address_space( + &ctx, + packed, + llvm_types::address_space::SHARED, + ), + "an AS3 pointer under a packed struct must be rejected by by-value paths" + ); + } + + #[test] + fn packed_shared_internal_abi_genericizes_one_direct_shared_pointer() { + let mut ctx = make_ctx(); + let tag = mir_uint(&mut ctx, 8); + let pointee = mir_uint(&mut ctx, 32); + let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, false).into(); + let packed: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "PackedShared".into(), + vec!["tag".into(), "ptr".into()], + vec![tag, shared], + vec![0, 1], + vec![0, 1], + 9, + 1, + ) + .into(); + + let abi = packed_shared_internal_abi_info(&mut ctx, packed) + .expect("ABI classification must succeed") + .expect("one direct packed AS3 pointer must use the internal carrier"); + assert_eq!(llvm_type_size_align(&ctx, abi.semantic_ty), Some((9, 1))); + assert_eq!(llvm_type_size_align(&ctx, abi.storage_ty), Some((9, 1))); + assert!(llvm_packed_struct_contains_pointer_in_address_space( + &ctx, + abi.semantic_ty, + llvm_types::address_space::SHARED, + )); + assert!(!llvm_type_contains_pointer_in_address_space( + &ctx, + abi.storage_ty, + llvm_types::address_space::SHARED, + )); + assert!(llvm_type_contains_pointer_in_address_space( + &ctx, + abi.storage_ty, + llvm_types::address_space::GENERIC, + )); + } + + #[test] + fn packed_shared_internal_abi_genericizes_multiple_direct_shared_pointers() { + let mut ctx = make_ctx(); + let tag = mir_uint(&mut ctx, 8); + let pointee = mir_uint(&mut ctx, 32); + let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, false).into(); + let packed: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "PackedSharedPair".into(), + vec!["tag".into(), "left".into(), "right".into()], + vec![tag, shared, shared], + vec![0, 1, 2], + vec![0, 1, 9], + 17, + 1, + ) + .into(); + + let abi = packed_shared_internal_abi_info(&mut ctx, packed) + .expect("ABI classification must succeed") + .expect("multiple direct AS3 leaves must use the recursive internal carrier"); + assert_eq!(llvm_type_size_align(&ctx, abi.semantic_ty), Some((17, 1))); + assert_eq!(llvm_type_size_align(&ctx, abi.storage_ty), Some((17, 1))); + assert!(llvm_packed_struct_contains_pointer_in_address_space( + &ctx, + abi.semantic_ty, + llvm_types::address_space::SHARED, + )); + assert!(!llvm_type_contains_pointer_in_address_space( + &ctx, + abi.storage_ty, + llvm_types::address_space::SHARED, + )); + + let fields = struct_fields(&ctx, abi.storage_ty); + assert_eq!(fields.len(), 3); + for field in &fields[1..] { + let field_ref = field.deref(&ctx); + let pointer = field_ref + .downcast_ref::() + .expect("both pointer leaves must remain pointer-typed"); + assert_eq!(pointer.address_space(), llvm_types::address_space::GENERIC); + } + } + + #[test] + fn packed_shared_internal_abi_genericizes_nested_struct_shared_pointers() { + let mut ctx = make_ctx(); + let pointee = mir_uint(&mut ctx, 32); + let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, false).into(); + let inner: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "InnerSharedPair".into(), + vec!["left".into(), "right".into()], + vec![shared, shared], + vec![0, 1], + vec![0, 8], + 16, + 8, + ) + .into(); + let tag = mir_uint(&mut ctx, 8); + let outer: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "PackedNestedShared".into(), + vec!["tag".into(), "inner".into()], + vec![tag, inner], + vec![0, 1], + vec![0, 1], + 17, + 1, + ) + .into(); + + let abi = packed_shared_internal_abi_info(&mut ctx, outer) + .expect("classification must not error") + .expect("nested AS3 leaves must use the recursive internal carrier"); + assert_eq!(llvm_type_size_align(&ctx, abi.semantic_ty), Some((17, 1))); + assert_eq!(llvm_type_size_align(&ctx, abi.storage_ty), Some((17, 1))); + assert!(!llvm_type_contains_pointer_in_address_space( + &ctx, + abi.storage_ty, + llvm_types::address_space::SHARED, + )); + + let outer_fields = struct_fields(&ctx, abi.storage_ty); + let inner_fields = struct_fields(&ctx, outer_fields[1]); + assert_eq!(inner_fields.len(), 2); + for field in inner_fields { + let field_ref = field.deref(&ctx); + let pointer = field_ref + .downcast_ref::() + .expect("nested shared leaves must remain pointer-typed"); + assert_eq!(pointer.address_space(), llvm_types::address_space::GENERIC); + } + } + + #[test] + fn packed_shared_internal_abi_genericizes_nested_tuple_shared_pointer() { + let mut ctx = make_ctx(); + let pointee = mir_uint(&mut ctx, 32); + let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, false).into(); + let word = mir_uint(&mut ctx, 32); + let inner: TypeHandle = MirTupleType::get_with_layout( + &mut ctx, + vec![shared, word], + vec![0, 1], + vec![0, 8], + 16, + 8, + ) + .into(); + let tag = mir_uint(&mut ctx, 8); + let outer: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "PackedNestedTupleShared".into(), + vec!["tag".into(), "inner".into()], + vec![tag, inner], + vec![0, 1], + vec![0, 1], + 17, + 1, + ) + .into(); + + let abi = packed_shared_internal_abi_info(&mut ctx, outer) + .expect("classification must not error") + .expect("an AS3 leaf nested in a tuple must use the recursive carrier"); + let outer_fields = struct_fields(&ctx, abi.storage_ty); + let inner_fields = struct_fields(&ctx, outer_fields[1]); + let inner_field_ref = inner_fields[0].deref(&ctx); + let pointer = inner_field_ref + .downcast_ref::() + .expect("nested tuple leaf must remain pointer-typed"); + assert_eq!(pointer.address_space(), llvm_types::address_space::GENERIC); + assert_eq!(llvm_type_size_align(&ctx, abi.storage_ty), Some((17, 1))); + } + + #[test] + fn packed_shared_internal_abi_skips_zero_sized_fields() { + let mut ctx = make_ctx(); + let tag = mir_uint(&mut ctx, 8); + let byte = mir_uint(&mut ctx, 8); + let marker: TypeHandle = MirArrayType::get(&mut ctx, byte, 0).into(); + let pointee = mir_uint(&mut ctx, 32); + let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, false).into(); + let packed: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "PackedSharedZstMarker".into(), + vec!["tag".into(), "marker".into(), "ptr".into()], + vec![tag, marker, shared], + vec![0, 1, 2], + vec![0, 1, 1], + 9, + 1, + ) + .into(); + + let abi = packed_shared_internal_abi_info(&mut ctx, packed) + .expect("ABI classification must succeed") + .expect("a zero-sized field must not knock the struct out of the ABI lane"); + assert_eq!(llvm_type_size_align(&ctx, abi.semantic_ty), Some((9, 1))); + assert_eq!(llvm_type_size_align(&ctx, abi.storage_ty), Some((9, 1))); + assert!(!llvm_type_contains_pointer_in_address_space( + &ctx, + abi.storage_ty, + llvm_types::address_space::SHARED, + )); + } + + #[test] + fn packed_shared_internal_abi_genericizes_bounded_shared_pointer_array() { + let mut ctx = make_ctx(); + let pointee = mir_uint(&mut ctx, 32); + let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, false).into(); + let shared_array: TypeHandle = MirArrayType::get(&mut ctx, shared, 2).into(); + let tag = mir_uint(&mut ctx, 8); + let outer: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "PackedSharedArray".into(), + vec!["tag".into(), "ptrs".into()], + vec![tag, shared_array], + vec![0, 1], + vec![0, 1], + 17, + 1, + ) + .into(); + + let abi = packed_shared_internal_abi_info(&mut ctx, outer) + .expect("classification must not error") + .expect("a bounded AS3 array must use the target-stable internal carrier"); + assert_eq!(llvm_type_size_align(&ctx, abi.semantic_ty), Some((17, 1))); + assert_eq!(llvm_type_size_align(&ctx, abi.storage_ty), Some((17, 1))); + assert!(!llvm_type_contains_pointer_in_address_space( + &ctx, + abi.storage_ty, + llvm_types::address_space::SHARED, + )); + + let outer_fields = struct_fields(&ctx, abi.storage_ty); + let array_ref = outer_fields[1].deref(&ctx); + let array = array_ref + .downcast_ref::() + .expect("bounded array field must remain an LLVM array"); + assert_eq!(array.size(), 2); + let element = array.elem_type(); + let element_ref = element.deref(&ctx); + let pointer = element_ref + .downcast_ref::() + .expect("array elements must remain pointer-typed"); + assert_eq!(pointer.address_space(), llvm_types::address_space::GENERIC); + } + + #[test] + fn packed_shared_internal_abi_accepts_array_rewrite_at_exact_bound() { + let mut ctx = make_ctx(); + let pointee = mir_uint(&mut ctx, 32); + let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, false).into(); + let count = MAX_PACKED_SHARED_INTERNAL_ABI_ARRAY_REWRITE_LEAVES; + let shared_array: TypeHandle = MirArrayType::get(&mut ctx, shared, count).into(); + let tag = mir_uint(&mut ctx, 8); + let outer: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "PackedSharedArrayAtBound".into(), + vec!["tag".into(), "ptrs".into()], + vec![tag, shared_array], + vec![0, 1], + vec![0, 1], + 1 + 8 * count, + 1, + ) + .into(); + + assert!( + packed_shared_internal_abi_info(&mut ctx, outer) + .expect("classification must not error") + .is_some(), + "the exact bounded-array rewrite limit must remain supported" + ); + } + + #[test] + fn packed_shared_internal_abi_rejects_array_rewrite_above_bound() { + let mut ctx = make_ctx(); + let pointee = mir_uint(&mut ctx, 32); + let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, false).into(); + let count = MAX_PACKED_SHARED_INTERNAL_ABI_ARRAY_REWRITE_LEAVES + 1; + let shared_array: TypeHandle = MirArrayType::get(&mut ctx, shared, count).into(); + let tag = mir_uint(&mut ctx, 8); + let outer: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "PackedSharedArrayAboveBound".into(), + vec!["tag".into(), "ptrs".into()], + vec![tag, shared_array], + vec![0, 1], + vec![0, 1], + 1 + 8 * count, + 1, + ) + .into(); + + assert!( + packed_shared_internal_abi_info(&mut ctx, outer) + .expect("classification must not error") + .is_none(), + "array-expanded AS3 leaves above the explicit budget must fail closed" + ); + } + + #[test] + fn packed_shared_internal_abi_rejects_shared_pointer_vector() { + let mut ctx = make_ctx(); + let tag = mir_uint(&mut ctx, 8); + let shared_pointer: TypeHandle = + llvm_types::PointerType::get(&ctx, llvm_types::address_space::SHARED).into(); + let shared_vector: TypeHandle = + llvm_types::VectorType::get(&ctx, shared_pointer, 2, llvm_types::VectorTypeKind::Fixed) + .into(); + let outer: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "PackedSharedVector".into(), + vec!["tag".into(), "ptrs".into()], + vec![tag, shared_vector], + vec![0, 1], + vec![0, 1], + 17, + 1, + ) + .into(); + + assert!( + packed_shared_internal_abi_info(&mut ctx, outer) + .expect("classification must not error") + .is_none(), + "shared-pointer vectors must remain outside the packed-AS3 internal ABI lane" + ); + } +} diff --git a/crates/mir-lower/src/convert/types/global_layout.rs b/crates/mir-lower/src/convert/types/global_layout.rs new file mode 100644 index 0000000000..89807ca5e1 --- /dev/null +++ b/crates/mir-lower/src/convert/types/global_layout.rs @@ -0,0 +1,826 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Layout validators for initialized (and relocated) device globals. + +use dialect_mir::types::{MirEnumType, MirStructType, MirTupleType}; +use llvm_export::types as llvm_types; +use pliron::context::Context; +use pliron::r#type::TypeHandle; + +use super::layout::mir_stored_size; +use super::{ + StructLayoutInfo, build_struct_slot_map, convert_type, enum_unmodeled_in_memory, + llvm_type_size_align, +}; + +/// Prove that an initialized Rust global can be accessed through the LLVM +/// semantic type produced by this lowering pipeline. +/// +/// The initializer itself is emitted as exact bytes. That is only half of the +/// contract: later field GEPs and typed loads still use `mir_ty`. If that type +/// places a field at a different byte offset, an exact initializer can still be +/// read incorrectly (or even past the end of the object). Reject every shape +/// for which we cannot prove that the two views agree. +pub(crate) fn validate_initialized_global_layout( + ctx: &mut Context, + mir_ty: TypeHandle, + initializer_size: u64, + initializer_align: u64, +) -> Result<(), anyhow::Error> { + if initializer_align == 0 || !initializer_align.is_power_of_two() { + return Err(anyhow::anyhow!( + "initialized global has invalid rustc allocation alignment {}", + initializer_align + )); + } + + validate_initialized_global_type(ctx, mir_ty, &mut Vec::new())?; + + let llvm_ty = convert_type(ctx, mir_ty)?; + let (llvm_size, llvm_align) = llvm_type_size_align(ctx, llvm_ty) + .ok_or_else(|| anyhow::anyhow!("initialized global has unsupported LLVM size/alignment"))?; + if llvm_size != initializer_size || llvm_align > initializer_align { + return Err(anyhow::anyhow!( + "initialized global type is not byte-compatible with rustc's allocation: the lowered LLVM value has size/alignment {}/{}, but the initializer has size/alignment {}/{}", + llvm_size, + llvm_align, + initializer_size, + initializer_align + )); + } + + Ok(()) +} + +/// Validate an initialized global that carries pointer relocations. +/// +/// Ordinary initialized globals keep their existing conservative semantic +/// layout validator. Relocated globals use a separate segmented physical +/// carrier, so a top-level `repr(packed)` struct may use that relocation path +/// as long as rustc's field ranges are explicit, non-overlapping, and in-bounds. +/// One direct nested struct may use the same relaxation when the top-level +/// struct itself has an ordinary, non-divergent LLVM layout and the child's +/// LLVM packed representation exactly reproduces rustc's offsets and size. A +/// packed top-level struct may not stack the relaxation on a packed child, and +/// deeper packed nesting remains on the ordinary fail-closed path. +pub(crate) fn validate_relocated_initialized_global_layout( + ctx: &mut Context, + mir_ty: TypeHandle, + initializer_size: u64, + initializer_align: u64, +) -> Result<(), anyhow::Error> { + match validate_initialized_global_layout(ctx, mir_ty, initializer_size, initializer_align) { + Ok(()) => Ok(()), + Err(original_error) => { + if initializer_align == 0 || !initializer_align.is_power_of_two() { + return Err(original_error); + } + + let struct_ty = { + let ty_ref = mir_ty.deref(ctx); + ty_ref.downcast_ref::().cloned() + }; + let Some(struct_ty) = struct_ty else { + return Err(original_error); + }; + if !struct_ty.has_explicit_layout() + || struct_ty.total_size() != initializer_size + || struct_ty.abi_align != initializer_align + { + return Err(original_error); + } + + validate_relocated_struct_field_ranges(ctx, &struct_ty)?; + + let top_level_layout = StructLayoutInfo::of_struct(&struct_ty); + let top_level_map = build_struct_slot_map(ctx, &top_level_layout)?; + let top_level_is_ordinary = !top_level_map.layout_diverges + && top_level_map.by_value_layout_faithful + && top_level_map + .llvm_struct_ty + .deref(ctx) + .downcast_ref::() + .is_some_and(|ty| ty.layout() == llvm_types::StructLayout::Unpacked); + + let mut visited = vec![mir_ty]; + for field_ty in struct_ty.field_types.iter().copied() { + if top_level_is_ordinary { + validate_relocated_initialized_global_child(ctx, field_ty, &mut visited)?; + } else { + validate_initialized_global_type(ctx, field_ty, &mut visited)?; + } + } + Ok(()) + } + } +} + +/// Prove that rustc's explicit field ranges for a relocated top-level struct +/// are non-overlapping and fit within the evaluated allocation. +fn validate_relocated_struct_field_ranges( + ctx: &mut Context, + struct_ty: &MirStructType, +) -> Result<(), anyhow::Error> { + let layout = StructLayoutInfo::of_struct(struct_ty); + + // Relocated globals relax only natural LLVM placement. Preserve the + // ordinary struct-layout metadata invariants before accepting rustc's + // explicit byte ranges as the physical initializer contract. + build_struct_slot_map(ctx, &layout)?; + + let mut ranges = Vec::with_capacity(layout.field_types.len()); + for (decl_index, field_ty) in layout.field_types.iter().copied().enumerate() { + let field_size = if let Some(size) = mir_stored_size(ctx, field_ty) { + size + } else { + let llvm_ty = convert_type(ctx, field_ty)?; + llvm_type_size_align(ctx, llvm_ty) + .map(|(size, _)| size) + .ok_or_else(|| { + anyhow::anyhow!( + "relocated initialized struct `{}` field {} has unsupported size", + struct_ty.name(), + decl_index + ) + })? + }; + if field_size == 0 { + continue; + } + + let start = layout.field_offsets[decl_index]; + let end = start.checked_add(field_size).ok_or_else(|| { + anyhow::anyhow!( + "relocated initialized struct `{}` field {} range overflows", + struct_ty.name(), + decl_index + ) + })?; + if end > layout.total_size { + return Err(anyhow::anyhow!( + "relocated initialized struct `{}` field {} occupies bytes {}..{}, but the allocation is only {} bytes", + struct_ty.name(), + decl_index, + start, + end, + layout.total_size + )); + } + ranges.push((start, end, decl_index)); + } + + ranges.sort_by_key(|(start, _, _)| *start); + for pair in ranges.windows(2) { + let (_, previous_end, previous_index) = pair[0]; + let (next_start, _, next_index) = pair[1]; + if next_start < previous_end { + return Err(anyhow::anyhow!( + "relocated initialized struct `{}` fields {} and {} overlap in rustc's byte layout", + struct_ty.name(), + previous_index, + next_index + )); + } + } + + Ok(()) +} + +/// Validate one direct child of an ordinary relocated initialized global. +/// +/// The ordinary initialized-global validator remains the default. Only when it +/// rejects a direct `MirStructType` do we consider the relocation-specific +/// relaxation, and only if `build_struct_slot_map` proves that the divergent +/// rustc layout is represented exactly by an LLVM packed struct. Its own field +/// ranges must still be explicit, non-overlapping, and in-bounds. Children of +/// that packed struct go back through the ordinary validator, deliberately +/// limiting this exception to one nesting level. +fn validate_relocated_initialized_global_child( + ctx: &mut Context, + mir_ty: TypeHandle, + visited: &mut Vec, +) -> Result<(), anyhow::Error> { + if visited.contains(&mir_ty) { + return Ok(()); + } + + let mut ordinary_visited = visited.clone(); + let original_error = match validate_initialized_global_type(ctx, mir_ty, &mut ordinary_visited) + { + Ok(()) => { + *visited = ordinary_visited; + return Ok(()); + } + Err(error) => error, + }; + + let struct_ty = { + let ty_ref = mir_ty.deref(ctx); + ty_ref.downcast_ref::().cloned() + }; + let Some(struct_ty) = struct_ty else { + return Err(original_error); + }; + if !struct_ty.has_explicit_layout() { + return Err(original_error); + } + + let layout = StructLayoutInfo::of_struct(&struct_ty); + let map = build_struct_slot_map(ctx, &layout)?; + let is_exact_packed = map.layout_diverges + && map.by_value_layout_faithful + && map + .llvm_struct_ty + .deref(ctx) + .downcast_ref::() + .is_some_and(|ty| ty.layout() == llvm_types::StructLayout::Packed); + if !is_exact_packed { + return Err(original_error); + } + + validate_relocated_struct_field_ranges(ctx, &struct_ty)?; + + visited.push(mir_ty); + for field_ty in struct_ty.field_types.iter().copied() { + validate_initialized_global_type(ctx, field_ty, visited)?; + } + Ok(()) +} + +fn validate_initialized_global_type( + ctx: &mut Context, + mir_ty: TypeHandle, + visited: &mut Vec, +) -> Result<(), anyhow::Error> { + if visited.contains(&mir_ty) { + return Ok(()); + } + visited.push(mir_ty); + + if let Some(name) = enum_unmodeled_in_memory(ctx, mir_ty) { + return Err(anyhow::anyhow!( + "initialized global contains enum `{}` with unknown physical layout; refusing to guess a byte representation", + name + )); + } + + enum Kind { + Struct(MirStructType), + Tuple(MirTupleType), + Enum(MirEnumType), + Array { element_ty: TypeHandle, size: u64 }, + Leaf, + } + + let kind = { + let ty_ref = mir_ty.deref(ctx); + if let Some(struct_ty) = ty_ref.downcast_ref::() { + Kind::Struct(struct_ty.clone()) + } else if let Some(tuple_ty) = ty_ref.downcast_ref::() { + Kind::Tuple(tuple_ty.clone()) + } else if let Some(enum_ty) = ty_ref.downcast_ref::() { + Kind::Enum(enum_ty.clone()) + } else if let Some(array_ty) = ty_ref.downcast_ref::() { + Kind::Array { + element_ty: array_ty.element_ty, + size: array_ty.size, + } + } else { + Kind::Leaf + } + }; + + match kind { + Kind::Struct(struct_ty) => { + validate_initialized_struct_layout(ctx, mir_ty, &struct_ty)?; + for field_ty in struct_ty.field_types { + validate_initialized_global_type(ctx, field_ty, visited)?; + } + } + Kind::Tuple(tuple_ty) => { + if !tuple_ty.get_types().is_empty() { + // Tuples carry rustc's field offsets exactly like structs; + // prove the lowered aggregate reproduces them byte-for-byte. + let layout = StructLayoutInfo::of_tuple(&tuple_ty); + validate_initialized_aggregate_layout( + ctx, + mir_ty, + "tuple", + "tuple", + &layout, + tuple_ty.abi_align(), + tuple_ty.has_explicit_layout(), + )?; + } + for field_ty in tuple_ty.get_types().iter().copied() { + validate_initialized_global_type(ctx, field_ty, visited)?; + } + } + Kind::Enum(enum_ty) => { + if enum_ty.total_size() > 0 { + let llvm_ty = convert_type(ctx, mir_ty)?; + let (llvm_size, llvm_align) = + llvm_type_size_align(ctx, llvm_ty).ok_or_else(|| { + anyhow::anyhow!( + "initialized enum `{}` has unsupported LLVM size/alignment", + enum_ty.name() + ) + })?; + if llvm_size != enum_ty.total_size() || llvm_align > enum_ty.abi_align() { + return Err(anyhow::anyhow!( + "initialized enum `{}` is not byte-compatible with rustc's layout: the lowered LLVM value has size/alignment {}/{}, but rustc requires {}/{}", + enum_ty.name(), + llvm_size, + llvm_align, + enum_ty.total_size(), + enum_ty.abi_align() + )); + } + } + for field_ty in enum_ty.all_field_types { + validate_initialized_global_type(ctx, field_ty, visited)?; + } + } + Kind::Array { element_ty, size } => { + // A zero-length array contributes no element bytes. Its outer LLVM + // size/alignment is still checked by + // `validate_initialized_global_layout`, but recursively rejecting + // an element layout would incorrectly reject valid promoted + // `&mut [Packed; 0]` constants whose physical allocation is empty. + if size != 0 { + validate_initialized_global_type(ctx, element_ty, visited)?; + } + } + Kind::Leaf => {} + } + + Ok(()) +} + +fn validate_initialized_struct_layout( + ctx: &mut Context, + mir_ty: TypeHandle, + struct_ty: &MirStructType, +) -> Result<(), anyhow::Error> { + let has_layout = struct_ty.has_explicit_layout(); + let layout = StructLayoutInfo::of_struct(struct_ty); + let name = struct_ty.name().to_string(); + validate_initialized_aggregate_layout( + ctx, + mir_ty, + "struct", + &name, + &layout, + struct_ty.abi_align, + has_layout, + ) +} + +/// Shared byte-layout validation for initialized-global structs and tuples. +/// +/// Proves that the lowered LLVM aggregate places every field at exactly the +/// byte offset rustc chose and matches rustc's total size, so a constant +/// initializer written slot-by-slot reproduces the host bytes. +#[allow(clippy::too_many_arguments)] +fn validate_initialized_aggregate_layout( + ctx: &mut Context, + mir_ty: TypeHandle, + kind_noun: &str, + name: &str, + layout: &StructLayoutInfo, + abi_align: u64, + has_explicit_layout: bool, +) -> Result<(), anyhow::Error> { + if layout.total_size == 0 { + let llvm_ty = convert_type(ctx, mir_ty)?; + let (llvm_size, _) = llvm_type_size_align(ctx, llvm_ty).ok_or_else(|| { + anyhow::anyhow!( + "initialized {} `{}` has unsupported LLVM size/alignment", + kind_noun, + name + ) + })?; + if llvm_size == 0 { + return Ok(()); + } + return Err(anyhow::anyhow!( + "initialized {} `{}` has no stored size but lowers to {} bytes", + kind_noun, + name, + llvm_size + )); + } + if !has_explicit_layout { + return Err(anyhow::anyhow!( + "initialized {} `{}` has no rustc field-offset metadata", + kind_noun, + name + )); + } + + let slots = build_struct_slot_map(ctx, layout)?; + let llvm_fields: Vec<_> = slots + .llvm_struct_ty + .deref(ctx) + .downcast_ref::() + .expect("aggregate slot map must produce an LLVM struct") + .fields() + .collect(); + + let mut slot_offsets = Vec::with_capacity(llvm_fields.len()); + let mut current_offset = 0u64; + for llvm_field in &llvm_fields { + let (field_size, field_align) = + llvm_type_size_align(ctx, *llvm_field).ok_or_else(|| { + anyhow::anyhow!( + "initialized {} `{}` field has unsupported LLVM size/alignment", + kind_noun, + name + ) + })?; + current_offset = current_offset.div_ceil(field_align.max(1)) * field_align.max(1); + slot_offsets.push(current_offset); + current_offset += field_size; + } + + for (decl_index, slot) in slots.decl_to_llvm.iter().enumerate() { + let Some(slot) = slot else { + continue; + }; + let actual_offset = slot_offsets[*slot as usize]; + let expected_offset = layout.field_offsets[decl_index]; + if actual_offset != expected_offset { + return Err(anyhow::anyhow!( + "initialized {} `{}` field {} lowers at byte {}, but rustc placed it at byte {}; packed and overlapping field layouts are not yet supported", + kind_noun, + name, + decl_index, + actual_offset, + expected_offset + )); + } + } + + let (llvm_size, llvm_align) = + llvm_type_size_align(ctx, slots.llvm_struct_ty).ok_or_else(|| { + anyhow::anyhow!( + "initialized {} `{}` has unsupported LLVM size/alignment", + kind_noun, + name + ) + })?; + if llvm_size != layout.total_size || llvm_align > abi_align { + return Err(anyhow::anyhow!( + "initialized {} `{}` lowers to size/alignment {}/{}, but rustc requires {}/{}", + kind_noun, + name, + llvm_size, + llvm_align, + layout.total_size, + abi_align + )); + } + + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::super::test_support::{make_ctx, mir_uint, mir_zst}; + use super::*; + use dialect_mir::types::{EnumVariant, MirArrayType, MirPtrType, MirUnionType}; + + #[test] + fn initialized_global_layout_accepts_explicit_overalignment() { + let mut ctx = make_ctx(); + let zst = mir_zst(&mut ctx); + validate_initialized_global_layout(&mut ctx, zst, 0, 1).unwrap(); + + let byte = mir_uint(&mut ctx, 8); + let over_aligned: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "OverAligned".into(), + vec!["byte".into()], + vec![byte], + vec![0], + vec![0], + 16, + 16, + ) + .into(); + + validate_initialized_global_layout(&mut ctx, over_aligned, 16, 16).unwrap(); + } + + #[test] + fn initialized_global_layout_rejects_packed_and_nested_packed_structs() { + let mut ctx = make_ctx(); + let byte = mir_uint(&mut ctx, 8); + let word = mir_uint(&mut ctx, 32); + let packed: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "Packed".into(), + vec!["tag".into(), "word".into()], + vec![byte, word], + vec![0, 1], + vec![0, 1], + 5, + 1, + ) + .into(); + + let empty_packed: TypeHandle = MirArrayType::get(&mut ctx, packed, 0).into(); + validate_initialized_global_layout(&mut ctx, empty_packed, 0, 1) + .expect("[Packed; 0] has no element bytes whose field layout can diverge"); + + let err = validate_initialized_global_layout(&mut ctx, packed, 5, 1).unwrap_err(); + assert!(err.to_string().contains("field 1 lowers at byte 4")); + + // Nesting must not hide the incompatible packed representation. + let wide = mir_uint(&mut ctx, 64); + let outer: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "Outer".into(), + vec!["packed".into(), "wide".into()], + vec![packed, wide], + vec![0, 1], + vec![0, 8], + 16, + 8, + ) + .into(); + let err = validate_initialized_global_layout(&mut ctx, outer, 16, 8).unwrap_err(); + assert!(err.to_string().contains("lowers at byte")); + } + + #[test] + fn relocated_initialized_global_layout_accepts_top_level_packed_struct() { + let mut ctx = make_ctx(); + let byte = mir_uint(&mut ctx, 8); + let target = mir_uint(&mut ctx, 32); + let pointer: TypeHandle = MirPtrType::get_global(&mut ctx, target, false).into(); + let packed: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "PackedRelocation".into(), + vec!["tag".into(), "ptr".into()], + vec![byte, pointer], + vec![0, 1], + vec![0, 1], + 9, + 1, + ) + .into(); + + validate_relocated_initialized_global_layout(&mut ctx, packed, 9, 1) + .expect("relocated top-level packed struct must be accepted"); + assert!(validate_initialized_global_layout(&mut ctx, packed, 9, 1).is_err()); + } + + #[test] + fn relocated_initialized_global_layout_accepts_thin_pointer_union() { + let mut ctx = make_ctx(); + let u32_ty = mir_uint(&mut ctx, 32); + let u8_ty = mir_uint(&mut ctx, 8); + let word_ptr: TypeHandle = MirPtrType::get_generic(&mut ctx, u32_ty, false).into(); + let byte_ptr: TypeHandle = MirPtrType::get_generic(&mut ctx, u8_ty, false).into(); + let union_ty: TypeHandle = MirUnionType::get( + &mut ctx, + "RelocatedPointerUnion".into(), + vec!["word".into(), "byte".into()], + vec![word_ptr, byte_ptr], + 8, + 8, + ) + .into(); + + validate_relocated_initialized_global_layout(&mut ctx, union_ty, 8, 8) + .expect("one pointer-word union must be valid relocated global storage"); + } + + #[test] + fn relocated_initialized_global_layout_rejects_malformed_memory_order() { + let mut ctx = make_ctx(); + let byte = mir_uint(&mut ctx, 8); + let target = mir_uint(&mut ctx, 32); + let pointer: TypeHandle = MirPtrType::get_global(&mut ctx, target, false).into(); + let malformed: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "MalformedPackedRelocation".into(), + vec!["tag".into(), "ptr".into()], + vec![byte, pointer], + vec![0, 0], + vec![0, 1], + 9, + 1, + ) + .into(); + + let error = validate_relocated_initialized_global_layout(&mut ctx, malformed, 9, 1) + .expect_err("malformed memory order must remain unsupported"); + assert!(error.to_string().contains("not a permutation"), "{error}"); + } + + #[test] + fn relocated_initialized_global_layout_accepts_one_direct_nested_packed_struct() { + let mut ctx = make_ctx(); + let byte = mir_uint(&mut ctx, 8); + let word = mir_uint(&mut ctx, 32); + let pointer: TypeHandle = MirPtrType::get_global(&mut ctx, word, false).into(); + let nested_packed: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "NestedPackedRelocation".into(), + vec!["tag".into(), "ptr".into()], + vec![byte, pointer], + vec![0, 1], + vec![0, 1], + 9, + 1, + ) + .into(); + let outer: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "OuterPackedRelocation".into(), + vec!["head".into(), "nested".into()], + vec![word, nested_packed], + vec![0, 1], + vec![0, 4], + 16, + 4, + ) + .into(); + + validate_relocated_initialized_global_layout(&mut ctx, outer, 16, 4) + .expect("one direct nested packed relocation carrier must be accepted"); + assert!(validate_initialized_global_layout(&mut ctx, outer, 16, 4).is_err()); + } + + #[test] + fn relocated_initialized_global_layout_rejects_packed_root_with_packed_child() { + let mut ctx = make_ctx(); + let byte = mir_uint(&mut ctx, 8); + let word = mir_uint(&mut ctx, 32); + let target = mir_uint(&mut ctx, 32); + let pointer: TypeHandle = MirPtrType::get_global(&mut ctx, target, false).into(); + let nested_packed: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "NestedPackedRelocation".into(), + vec!["tag".into(), "ptr".into()], + vec![byte, pointer], + vec![0, 1], + vec![0, 1], + 9, + 1, + ) + .into(); + let packed_root: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "PackedRootRelocation".into(), + vec!["tag".into(), "word".into(), "nested".into()], + vec![byte, word, nested_packed], + vec![0, 1, 2], + vec![0, 1, 5], + 14, + 1, + ) + .into(); + + let error = validate_relocated_initialized_global_layout(&mut ctx, packed_root, 14, 1) + .expect_err("a packed top-level struct must not stack the nested packed relaxation"); + assert!(error.to_string().contains("lowers at byte"), "{error}"); + } + + #[test] + fn relocated_initialized_global_layout_rejects_deeper_packed_nesting() { + let mut ctx = make_ctx(); + let byte = mir_uint(&mut ctx, 8); + let word = mir_uint(&mut ctx, 32); + let target = mir_uint(&mut ctx, 32); + let pointer: TypeHandle = MirPtrType::get_global(&mut ctx, target, false).into(); + let inner_packed: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "InnerPackedRelocation".into(), + vec!["tag".into(), "ptr".into()], + vec![byte, pointer], + vec![0, 1], + vec![0, 1], + 9, + 1, + ) + .into(); + let middle_packed: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "MiddlePackedRelocation".into(), + vec!["tag".into(), "word".into(), "inner".into()], + vec![byte, word, inner_packed], + vec![0, 1, 2], + vec![0, 1, 5], + 14, + 1, + ) + .into(); + let outer: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "OuterDeepPackedRelocation".into(), + vec!["head".into(), "middle".into()], + vec![word, middle_packed], + vec![0, 1], + vec![0, 4], + 20, + 4, + ) + .into(); + + let error = validate_relocated_initialized_global_layout(&mut ctx, outer, 20, 4) + .expect_err("packed relocation relaxation must stop after one nesting level"); + assert!(error.to_string().contains("lowers at byte"), "{error}"); + } + + #[test] + fn relocated_initialized_global_layout_rejects_overlapping_struct_fields() { + let mut ctx = make_ctx(); + let word = mir_uint(&mut ctx, 32); + let overlapping: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "OverlappingRelocation".into(), + vec!["left".into(), "right".into()], + vec![word, word], + vec![0, 1], + vec![0, 0], + 4, + 1, + ) + .into(); + + let error = validate_relocated_initialized_global_layout(&mut ctx, overlapping, 4, 1) + .expect_err("overlapping field ranges must remain unsupported"); + assert!(error.to_string().contains("overlap"), "{error}"); + } + + #[test] + fn initialized_global_layout_rejects_old_union_and_tuple_models() { + let mut ctx = make_ctx(); + let word = mir_uint(&mut ctx, 32); + let union_as_struct: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "UnionBeforeSharedStorageLowering".into(), + vec!["left".into(), "right".into()], + vec![word, word], + vec![0, 1], + vec![0, 0], + 4, + 4, + ) + .into(); + let err = validate_initialized_global_layout(&mut ctx, union_as_struct, 4, 4).unwrap_err(); + assert!(err.to_string().contains("field 1 lowers at byte 4")); + + // A tuple without recorded rustc layout cannot prove its bytes. + let byte = mir_uint(&mut ctx, 8); + let wide = mir_uint(&mut ctx, 64); + let tuple: TypeHandle = MirTupleType::get(&mut ctx, vec![byte, wide]).into(); + let err = validate_initialized_global_layout(&mut ctx, tuple, 16, 8).unwrap_err(); + assert!( + err.to_string() + .contains("no stored size but lowers to 16 bytes"), + "{err}" + ); + + // The same tuple carrying rustc's real (reordered) layout validates: + // memory order is (wide @ 0, byte @ 8), total size 16. + let laid_out_tuple: TypeHandle = MirTupleType::get_with_layout( + &mut ctx, + vec![byte, wide], + vec![1, 0], + vec![8, 0], + 16, + 8, + ) + .into(); + validate_initialized_global_layout(&mut ctx, laid_out_tuple, 16, 8) + .expect("tuples with recorded rustc offsets are provable initialized-global storage"); + } + + #[test] + fn initialized_global_layout_rejects_enum_with_unknown_physical_layout() { + let mut ctx = make_ctx(); + let discr = mir_uint(&mut ctx, 8); + let payload = mir_uint(&mut ctx, 32); + let niche: TypeHandle = MirEnumType::get( + &mut ctx, + "OptionNonZero".into(), + discr, + vec![0, 1], + vec![ + EnumVariant::unit("None".into()), + EnumVariant::new("Some".into(), vec![payload]), + ], + ) + .into(); + + let err = validate_initialized_global_layout(&mut ctx, niche, 4, 4).unwrap_err(); + assert!(err.to_string().contains("unknown physical layout")); + } +} diff --git a/crates/mir-lower/src/convert/types/layout.rs b/crates/mir-lower/src/convert/types/layout.rs new file mode 100644 index 0000000000..54083e2424 --- /dev/null +++ b/crates/mir-lower/src/convert/types/layout.rs @@ -0,0 +1,548 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Size/alignment metrics, padding/filler type builders, byte-faithfulness +//! and `i1`-storage predicates, and natural struct layout walks. + +use dialect_mir::types::{ + MirArrayType, MirEnumType, MirFP16Type, MirPtrType, MirStructType, MirTupleType, MirUnionType, +}; +use llvm_export::types as llvm_types; +use pliron::builtin::type_interfaces::FloatTypeInterface; +use pliron::builtin::types::{FP32Type, FP64Type, IntegerType, Signedness}; +use pliron::context::Context; +use pliron::r#type::{TypeHandle, type_cast}; + +/// Create a padding type: `[N x i8]` for N bytes of padding. +pub(super) fn make_padding_type(ctx: &mut Context, size: u64) -> TypeHandle { + let i8_ty = IntegerType::get(ctx, 8, Signedness::Signless); + llvm_types::ArrayType::get(ctx, i8_ty.into(), size).into() +} + +/// Storage for `size` bytes of enum filler at byte `offset`: bytes the payload +/// occupies that no typed slot claims. +/// +/// `[N x i8]` is byte-exact but costs one leaf *per byte* everywhere a payload +/// value is built, merged, or taken apart. `Option<&[T]>` is the shape that +/// shows it: the niche carrier claims the pointer, the length is left to an +/// 8-byte filler, and building the value then decomposes that length into +/// eight `i8` `insertvalue`s, every block merge carries eight `phi i8`, and +/// reading it back is a chain of `prmt` byte permutes with the aggregate +/// spilled to `.local` in between. Measured on sm_86, it costs the *pointer* +/// too: it rides through the same byte soup, `InferAddressSpaces` can no +/// longer follow it, and the access lands in the generic window (`ld.v2.b64`) +/// instead of `ld.global.v2.b64`. +/// +/// One integer of the same width is the same bytes in one leaf. It is only +/// legal where it moves nothing: +/// +/// - the width is 2, 4 or 8 bytes, so the integer is byte-faithful (a multiple +/// of 8 bits, hence no padding of its own). 16 is deliberately excluded: +/// `i128` is legal LLVM but lowers to register pairs on NVPTX, so widening +/// that far trades one win for another cost and wants its own measurement; +/// - `offset` is a multiple of the width, so LLVM inserts no gap ahead of the +/// field and every later field keeps its byte offset; +/// - the width does not exceed the enum's alignment, or the struct's natural +/// alignment would rise above rustc's. +/// +/// Anything else keeps the byte array. The filler is a *vehicle*, not a claim: +/// nothing reads it field-wise, because a payload that has no typed slot +/// round-trips through memory as a whole aggregate, and both ends of that +/// round trip are byte-exact. `build_enum_slot_map`'s own size and alignment +/// assertions — hard errors, not debug checks — backstop all three conditions. +pub(super) fn make_enum_filler_type( + ctx: &mut Context, + offset: u64, + size: u64, + abi_align: u64, +) -> TypeHandle { + if matches!(size, 2 | 4 | 8) && offset.is_multiple_of(size) && size <= abi_align.max(1) { + let width = u32::try_from(size * 8).expect("size is at most 8, so width is at most 64"); + return IntegerType::get(ctx, width, Signedness::Signless).into(); + } + make_padding_type(ctx, size) +} + +/// Whether this is an aggregate/vector that contains LLVM `i1` storage. +/// +/// Rust `bool` is an SSA `i1`, but its memory representation is one complete +/// byte whose only valid values are 0 and 1. Enum storage never uses `i1` as +/// a physical type: scalar bools claim an explicit i8 byte, and aggregates +/// containing bools claim their byte-faithful twin (see +/// [`llvm_byte_faithful_twin`]), with construction canonicalizing the value +/// before the store. +pub(crate) fn llvm_type_contains_i1(ctx: &Context, ty: TypeHandle) -> bool { + let ty_ref = ty.deref(ctx); + if let Some(integer) = ty_ref.downcast_ref::() { + return integer.width() == 1; + } + if let Some(array) = ty_ref.downcast_ref::() { + return llvm_type_contains_i1(ctx, array.elem_type()); + } + if let Some(vector) = ty_ref.downcast_ref::() { + return llvm_type_contains_i1(ctx, vector.elem_type()); + } + if let Some(struct_ty) = ty_ref.downcast_ref::() { + return struct_ty + .fields() + .any(|field| llvm_type_contains_i1(ctx, field)); + } + false +} + +/// The byte-faithful storage twin of an LLVM type: every `i1` leaf becomes +/// its canonical `i8` memory byte, recursively through structs and arrays. +/// +/// Rust guarantees a bool occupies one full byte holding exactly 0 or 1, so +/// storing the twin (with each bool zero-extended) writes the same bytes the +/// host writes, while re-loading the original type from those canonical +/// bytes remains well-defined. Sizes and alignments are unchanged because an +/// `i1` already occupies one byte of storage. +/// +/// Returns `None` for shapes with a different memory story (`i1` vectors are +/// bit-packed masks) or unknown containers; callers must fail closed. +pub(crate) fn llvm_byte_faithful_twin(ctx: &mut Context, ty: TypeHandle) -> Option { + if !llvm_type_contains_i1(ctx, ty) { + return Some(ty); + } + enum Shape { + Bool, + Array(TypeHandle, u64), + Struct { + fields: Vec, + layout: llvm_types::StructLayout, + }, + Other, + } + let shape = { + let ty_ref = ty.deref(ctx); + if ty_ref + .downcast_ref::() + .is_some_and(|integer| integer.width() == 1) + { + Shape::Bool + } else if let Some(array) = ty_ref.downcast_ref::() { + Shape::Array(array.elem_type(), array.size()) + } else if let Some(struct_ty) = ty_ref.downcast_ref::() { + Shape::Struct { + fields: struct_ty.fields().collect(), + layout: struct_ty.layout(), + } + } else { + Shape::Other + } + }; + match shape { + Shape::Bool => Some(IntegerType::get(ctx, 8, Signedness::Signless).into()), + Shape::Array(elem, count) => { + let twin = llvm_byte_faithful_twin(ctx, elem)?; + Some(llvm_types::ArrayType::get(ctx, twin, count).into()) + } + Shape::Struct { fields, layout } => { + let twins = fields + .into_iter() + .map(|field| llvm_byte_faithful_twin(ctx, field)) + .collect::>>()?; + Some(llvm_types::StructType::get_unnamed(ctx, (twins, layout)).into()) + } + Shape::Other => None, + } +} + +/// Whether a MIR aggregate contains a semantic Rust bool value. +/// +/// This deliberately stops at pointers/slices: a pointee bool does not occupy +/// bytes in the aggregate itself. It also stops at nested enums, whose own +/// slot-map construction is responsible for canonicalizing a top-level bool +/// or rejecting a deeper one. Inspecting the MIR type is necessary because a +/// union may lower to raw i8 storage and thereby hide a bool from the LLVM +/// type-level check above. +pub(super) fn mir_type_contains_i1(ctx: &Context, ty: TypeHandle) -> bool { + let ty_ref = ty.deref(ctx); + if let Some(integer) = ty_ref.downcast_ref::() { + return integer.width() == 1; + } + if let Some(struct_ty) = ty_ref.downcast_ref::() { + return struct_ty + .field_types() + .iter() + .copied() + .any(|field| mir_type_contains_i1(ctx, field)); + } + if let Some(tuple_ty) = ty_ref.downcast_ref::() { + return tuple_ty + .get_types() + .iter() + .copied() + .any(|field| mir_type_contains_i1(ctx, field)); + } + if let Some(array_ty) = ty_ref.downcast_ref::() { + return mir_type_contains_i1(ctx, array_ty.element_ty); + } + if let Some(union_ty) = ty_ref.downcast_ref::() { + return union_ty + .field_types() + .iter() + .copied() + .any(|field| mir_type_contains_i1(ctx, field)); + } + false +} + +/// Whether loading and storing this LLVM value preserves every byte in its +/// allocation. In particular, `i1` is not byte-faithful: it occupies one +/// addressable byte, but LLVM does not define the upper seven stored bits. +pub(crate) fn llvm_type_is_byte_faithful(ctx: &Context, ty: TypeHandle) -> bool { + let ty_ref = ty.deref(ctx); + if let Some(int_ty) = ty_ref.downcast_ref::() { + return int_ty.width().is_multiple_of(8); + } + if ty_ref.is::() + || ty_ref.is::() + || ty_ref.is::() + || ty_ref.is::() + { + return true; + } + if let Some(array) = ty_ref.downcast_ref::() { + return llvm_type_is_byte_faithful(ctx, array.elem_type()); + } + if let Some(struct_ty) = ty_ref.downcast_ref::() { + let fields: Vec<_> = struct_ty.fields().collect(); + if struct_ty.layout() == llvm_types::StructLayout::Packed { + return fields + .into_iter() + .all(|field| llvm_type_is_byte_faithful(ctx, field)); + } + + let mut end = 0u64; + let mut max_align = 1u64; + for field in fields { + if !llvm_type_is_byte_faithful(ctx, field) { + return false; + } + let Some((field_size, field_align)) = llvm_type_size_align(ctx, field) else { + return false; + }; + let field_align = field_align.max(1); + let aligned_end = end.div_ceil(field_align) * field_align; + if aligned_end != end { + // LLVM would insert bytes that are not represented by an SSA + // field. Loading and re-storing the aggregate would lose them. + return false; + } + end += field_size; + max_align = max_align.max(field_align); + } + // Reject implicit trailing padding for the same reason. Explicit + // `[N x i8]` padding fields keep `end` equal to the allocation size. + return end.div_ceil(max_align) * max_align == end; + } + false +} + +/// Size of a MIR-level type from rustc layout truth, when stored. +/// +/// `MirStructType`, `MirTupleType`, `MirUnionType`, and `MirEnumType` carry +/// `total_size` (interior and trailing padding included) straight from +/// rustc's layout query; arrays of such aggregates multiply it out. Returns +/// `None` when no stored size is available (e.g. niched/single-variant enums +/// store 0) and the caller must fall back to the LLVM-level approximation. +pub(super) fn mir_stored_size(ctx: &Context, mir_ty: TypeHandle) -> Option { + let ty_ref = mir_ty.deref(ctx); + if let Some(s) = ty_ref.downcast_ref::() { + if s.total_size() > 0 { + return Some(s.total_size()); + } + return None; + } + if let Some(t) = ty_ref.downcast_ref::() { + if t.total_size > 0 { + return Some(t.total_size); + } + return None; + } + if let Some(e) = ty_ref.downcast_ref::() { + if e.total_size() > 0 { + return Some(e.total_size()); + } + return None; + } + if let Some(u) = ty_ref.downcast_ref::() { + return Some(u.total_size()); + } + if let Some(a) = ty_ref.downcast_ref::() { + let elem_ty = a.element_ty; + let size = a.size; + // Checked: alignment claims consume this, and a wrapped product + // would masquerade as a small, trusted stride. + return mir_stored_size(ctx, elem_ty).and_then(|elem_size| elem_size.checked_mul(size)); + } + None +} + +/// Exact byte stride of a MIR array's element, when provable. +/// +/// Element `i` of an array lives at byte `i * stride`, so any alignment +/// claim about an element address is only as strong as the stride it +/// multiplies, and a wrong stride is a miscompile. Scalars have exact +/// sizes; MIR aggregates answer with rustc's stored size (interior and +/// trailing padding included) via [`mir_stored_size`]; arrays of scalars +/// multiply out. Everything else answers `None` so the caller declines +/// instead of guessing. +pub(crate) fn mir_element_stride(ctx: &Context, mir_ty: TypeHandle) -> Option { + if let Some(size) = mir_stored_size(ctx, mir_ty) { + return Some(size); + } + let ty_ref = mir_ty.deref(ctx); + if let Some(int_ty) = ty_ref.downcast_ref::() { + return Some(u64::from(int_ty.width()).div_ceil(8)); + } + // The importer's f16 is MirFP16Type; the HalfType arm covers IR that + // already carries the converted LLVM scalar. + if ty_ref.is::() || ty_ref.is::() { + return Some(2); + } + if ty_ref.is::() { + return Some(4); + } + if ty_ref.is::() { + return Some(8); + } + if let Some(ptr_ty) = ty_ref.downcast_ref::() { + // Generic-space pointers are 64-bit under every data layout the + // exporter can choose; a shared-space pointer is 32-bit under the + // modern NVVM layout alone, so its stored stride is target-dependent. + // Claim the former, decline the latter. + return (ptr_ty.address_space == 0).then_some(8); + } + if let Some(array_ty) = ty_ref.downcast_ref::() { + // Arrays of sized aggregates already answered through + // `mir_stored_size` above; this recursion serves arrays of + // scalars and pointers, and deeper such arrays. + let element_ty = array_ty.element_type(); + let count = array_ty.size(); + drop(ty_ref); + return mir_element_stride(ctx, element_ty)?.checked_mul(count); + } + None +} + +/// Exact ABI alignment carried by a MIR aggregate type, when rustc layout is +/// available. +/// +/// LLVM aggregate types cannot encode a Rust `repr(align(N))` raise. Tuples, +/// structs, enums, and unions therefore carry rustc's alignment explicitly in +/// the MIR dialect. Arrays have the same ABI alignment as their element, so +/// recurse through any number of array layers instead of relying on the +/// converted LLVM element's structural alignment. +pub(crate) fn mir_type_abi_align(ctx: &Context, mir_ty: TypeHandle) -> Option { + let ty_ref = mir_ty.deref(ctx); + if let Some(tuple_ty) = ty_ref.downcast_ref::() { + return Some(tuple_ty.abi_align()).filter(|align| *align > 0); + } + if let Some(struct_ty) = ty_ref.downcast_ref::() { + return Some(struct_ty.abi_align).filter(|align| *align > 0); + } + if let Some(enum_ty) = ty_ref.downcast_ref::() { + return Some(enum_ty.abi_align()).filter(|align| *align > 0); + } + if let Some(union_ty) = ty_ref.downcast_ref::() { + return Some(union_ty.abi_align()).filter(|align| *align > 0); + } + if let Some(array_ty) = ty_ref.downcast_ref::() { + let element_ty = array_ty.element_type(); + return mir_type_abi_align(ctx, element_ty); + } + None +} + +/// LLVM natural-layout `(size, align)` of an exported LLVM type, in bytes. +/// +/// Mirrors LLVM's default data layout for nvptx64 (scalars align to their +/// size, arrays to their element, non-packed structs to their widest field). +/// Computes the real allocation size (interior and trailing padding +/// included), which is what GEP striding and the enum size check below +/// need. Answers `None` rather than guess: sizes here are exact or absent. +pub(crate) fn llvm_type_size_align(ctx: &Context, ty: TypeHandle) -> Option<(u64, u64)> { + let ty_ref = ty.deref(ctx); + + if let Some(int_ty) = ty_ref.downcast_ref::() { + let size = (int_ty.width() as u64).div_ceil(8); + // i8 → 1, i16 → 2, i32 → 4, i64 → 8, i128 → 16. + return Some((size, size.next_power_of_two().min(16))); + } + if ty_ref.is::() { + return Some((2, 2)); + } + if ty_ref.is::() { + return Some((4, 4)); + } + if ty_ref.is::() { + return Some((8, 8)); + } + if ty_ref.is::() { + // Lowering runs before the exporter chooses the minimal, legacy, or + // modern NVPTX data layout. The first two use 64-bit pointers in all + // address spaces; modern NVVM alone uses p3:32. Callers that expose + // physical shared-pointer width must either genericize a direct + // pointer first or reject the target-dependent representation. + return Some((8, 8)); + } + if let Some(arr_ty) = ty_ref.downcast_ref::() { + let (elem_size, elem_align) = llvm_type_size_align(ctx, arr_ty.elem_type())?; + return Some((elem_size.checked_mul(arr_ty.size())?, elem_align.max(1))); + } + if let Some(vector) = ty_ref.downcast_ref::() { + if vector.is_scalable() { + return None; + } + let element_bits = { + let element = vector.elem_type(); + let element_ref = element.deref(ctx); + if let Some(integer) = element_ref.downcast_ref::() { + u64::from(integer.width()) + } else if let Some(float) = type_cast::(&*element_ref) { + u64::try_from(float.get_semantics().bits).ok()? + } else if element_ref.is::() { + 64 + } else { + return None; + } + }; + let total_bits = element_bits.checked_mul(u64::from(vector.num_elements()))?; + let size = total_bits.div_ceil(8); + // Both cuda-oxide NVPTX data layouts explicitly define fixed vector + // ABI alignment only for these widths. Refuse to guess LLVM defaults + // for any other width in physical Rust layout code. + let align = match total_bits { + 16 => 2, + 32 => 4, + 64 => 8, + 128 => 16, + _ => return None, + }; + return Some((size, align)); + } + if let Some(struct_ty) = ty_ref.downcast_ref::() { + let fields: Vec<_> = struct_ty.fields().collect(); + if struct_ty.layout() == llvm_types::StructLayout::Packed { + let mut size = 0u64; + for field in fields { + let (field_size, _) = llvm_type_size_align(ctx, field)?; + size = size.checked_add(field_size)?; + } + return Some((size, 1)); + } + + let (_end, size, align) = natural_struct_layout(ctx, &fields)?; + return Some((size, align)); + } + + None +} + +/// The one walk defining natural (non-packed) LLVM struct placement: each +/// field starts at the running offset rounded up to its own alignment. +/// +/// Returns `(offsets, end, align)`: the byte offset of every field, the +/// unrounded offset just past the last field, and the widest field alignment. +/// [`natural_struct_layout`] and [`StructSlotMap::natural_slot_offsets`] are +/// both views of this walk, so a size question and an offset question can +/// never disagree. +pub(super) fn natural_layout_walk( + ctx: &Context, + fields: &[TypeHandle], +) -> Option<(Vec, u64, u64)> { + let mut offsets = Vec::with_capacity(fields.len()); + let mut end = 0u64; + let mut align = 1u64; + for field in fields { + let (field_size, field_align) = llvm_type_size_align(ctx, *field)?; + let field_align = field_align.max(1); + end = end.div_ceil(field_align) * field_align; + offsets.push(end); + end = end.checked_add(field_size)?; + align = align.max(field_align); + } + Some((offsets, end, align)) +} + +/// Natural (non-packed) LLVM struct layout over `fields`. +/// +/// Returns `(end, size, align)` where `end` is the unrounded offset just past +/// the last field, `size` is `end` rounded up to the struct alignment (the +/// allocation size LLVM uses for GEP striding), and `align` is the widest +/// field alignment. +pub(crate) fn natural_struct_layout( + ctx: &Context, + fields: &[TypeHandle], +) -> Option<(u64, u64, u64)> { + let (_offsets, end, align) = natural_layout_walk(ctx, fields)?; + let size = end.div_ceil(align) * align; + Some((end, size, align)) +} + +#[cfg(test)] +mod tests { + use super::super::test_support::{llvm_int, make_ctx, mir_uint}; + use super::*; + + #[test] + fn mir_abi_alignment_recurses_through_nested_arrays() { + let mut ctx = make_ctx(); + let byte = mir_uint(&mut ctx, 8); + let marker: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "Align32".into(), + vec![], + vec![], + vec![], + vec![], + 0, + 32, + ) + .into(); + let tuple: TypeHandle = MirTupleType::get_with_layout( + &mut ctx, + vec![marker, byte], + vec![0, 1], + vec![0, 0], + 32, + 32, + ) + .into(); + let inner: TypeHandle = MirArrayType::get(&mut ctx, tuple, 2).into(); + let outer: TypeHandle = MirArrayType::get(&mut ctx, inner, 3).into(); + let plain: TypeHandle = MirArrayType::get(&mut ctx, byte, 4).into(); + + assert_eq!(mir_type_abi_align(&ctx, tuple), Some(32)); + assert_eq!(mir_type_abi_align(&ctx, inner), Some(32)); + assert_eq!(mir_type_abi_align(&ctx, outer), Some(32)); + assert_eq!(mir_type_abi_align(&ctx, plain), None); + } + + #[test] + fn fixed_vector_layout_uses_packed_bit_width_and_rejects_unknown_widths() { + let mut ctx = make_ctx(); + let i1 = llvm_int(&mut ctx, 1); + let v16i1: TypeHandle = + llvm_types::VectorType::get(&ctx, i1, 16, llvm_types::VectorTypeKind::Fixed).into(); + assert_eq!( + llvm_type_size_align(&ctx, v16i1), + Some((2, 2)), + "<16 x i1> is 16 packed bits, not sixteen bytes" + ); + + let i8 = llvm_int(&mut ctx, 8); + let v3i8: TypeHandle = + llvm_types::VectorType::get(&ctx, i8, 3, llvm_types::VectorTypeKind::Fixed).into(); + assert_eq!( + llvm_type_size_align(&ctx, v3i8), + None, + "the NVPTX data layout does not define a 24-bit vector alignment" + ); + } +} diff --git a/crates/mir-lower/src/convert/types/mod.rs b/crates/mir-lower/src/convert/types/mod.rs new file mode 100644 index 0000000000..44f3640d20 --- /dev/null +++ b/crates/mir-lower/src/convert/types/mod.rs @@ -0,0 +1,145 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Type conversion from `dialect-mir` types to LLVM dialect types. +//! +//! This module handles the translation of `dialect-mir` type representations +//! to their LLVM dialect equivalents. Type conversion is foundational to +//! the lowering pass—most operation converters depend on it. +//! +//! # Overview +//! +//! `dialect-mir` types are high-level, Rust-like types that preserve semantic +//! information (signedness, slice semantics, etc.). LLVM dialect types are +//! lower-level and match LLVM IR types directly. +//! +//! # Type Mapping Table +//! +//! | `dialect-mir` Type | LLVM dialect Type | Notes | +//! |---------------------------------|-----------------------------------|-----------------------------| +//! | `IntegerType` (signed/unsigned) | `IntegerType` (signless) | Width preserved | +//! | `MirFP16Type` | `HalfType` | Rust `f16` → LLVM `half` | +//! | `FP32Type`, `FP64Type` | Same (builtin) | Pass-through | +//! | `MirPtrType` | `PointerType` | Address space preserved | +//! | `MirSliceType` | `StructType { ptr, i64 }` | Fat pointer | +//! | `MirDisjointSliceType` | `StructType { ptr, i64 }` | Same as slice | +//! | `MirTupleType` | `StructType` | Empty tuple → empty struct | +//! | `MirStructType` | `StructType` | Fields recursively converted| +//! | `MirUnionType` | Aligned shared-storage struct | All fields start at byte zero| +//! | `MirEnumType` | `StructType` (rustc byte layout) | See "Enum Type Representation" | +//! | `ArrayType` | `ArrayType` | Element type converted | +//! | `VectorType` | `VectorType` | Element type converted | +//! +//! # Signedness Handling +//! +//! LLVM IR integers are signless—the signedness is encoded in the operations +//! that use them (e.g., `sdiv` vs `udiv`). During type conversion: +//! +//! - Signed/unsigned MIR integers → signless LLVM integers +//! - The original signedness is preserved in operations (see `arithmetic.rs`) +//! +//! # Address Space Handling +//! +//! GPU memory uses address spaces to distinguish memory types: +//! +//! | Address Space | Memory Type | Usage | +//! |---------------|-------------|---------------------------| +//! | 0 | Generic | Can point to any memory | +//! | 1 | Global | Device memory (VRAM) | +//! | 3 | Shared | Per-block shared memory | +//! | 4 | Constant | Read-only device memory | +//! | 5 | Local | Per-thread stack/spill | +//! +//! Pointer address spaces are preserved through conversion. Slice types use +//! generic address space (0) because they can point to any memory type. +//! +//! # Slice Type Representation +//! +//! Rust slices (`&[T]`) are represented as fat pointers in LLVM: +//! +//! ```text +//! MIR: MirSliceType +//! LLVM: struct { ptr, i64 } ; pointer + length +//! ``` +//! +//! This matches the Rust ABI for slices passed by value. +//! +//! # Enum Type Representation +//! +//! A Rust enum is one tag plus the payload of whichever variant is +//! alive; all variants share the same bytes. We build an LLVM struct +//! that puts the tag and every payload field at the exact byte position +//! rustc chose, inserting `[N x i8]` filler for the gaps: +//! +//! ```text +//! #[repr(u32)] enum E { A(u32), B(f32), C } // rustc: 8 bytes, +//! // tag at 0, payloads at 4 +//! LLVM: { i32, i32 } ; slot 0 = tag, slot 1 = A's payload +//! ; B's f32 also lives at byte 4 but has a +//! ; different type, so it is read/written through +//! ; memory instead of owning a slot +//! ``` +//! +//! Because the bytes match rustc exactly, enum data can cross the +//! host/device boundary safely. The tag slot stores the variant's +//! DECLARED discriminant value (`enum E { A = 7 }` stores 7), not its +//! position. See `build_enum_slot_map` in this module for the full +//! story. +//! +//! # Function Type Conversion +//! +//! Function types undergo ABI transformations: +//! +//! - Slice arguments are flattened to `(ptr, len)` pairs +//! - Struct arguments are flattened to individual fields +//! - Empty tuple return type becomes void +//! +//! This matches the C ABI for GPU kernels. + +mod conversion; +mod enum_layout; +mod func_abi; +mod global_layout; +mod layout; +mod pointer_storage; +mod struct_layout; +#[cfg(test)] +mod test_support; +mod union_storage; + +pub use conversion::{convert_type, is_zero_sized_type}; +pub(crate) use conversion::{make_disjoint_slice_struct, make_slice_struct}; +pub(crate) use enum_layout::{ + EnumSlotMap, build_enum_slot_map, convert_enum_to_llvm, enum_unmodeled_in_memory, + find_unmodeled_enum_in_abi, +}; +pub(crate) use func_abi::{ + TransparentScalarAbiInfo, packed_shared_internal_abi_info, transparent_scalar_abi_info, + transparent_scalar_field, +}; +pub use func_abi::{convert_function_type, is_kernel_func}; +// Surface parity with the pre-split `types.rs`: these four are currently +// consumed only inside `func_abi` itself or under `cfg(test)` (e.g. +// `transparent_scalar_llvm_type` from the `lowering.rs` tests), so the +// re-export trips `unused_imports` in non-test builds. +#[allow(unused_imports)] +pub(crate) use func_abi::{ + MAX_PACKED_SHARED_INTERNAL_ABI_ARRAY_REWRITE_LEAVES, PackedSharedInternalAbiInfo, + TransparentScalarLayer, transparent_scalar_llvm_type, +}; +pub(crate) use global_layout::{ + validate_initialized_global_layout, validate_relocated_initialized_global_layout, +}; +pub(crate) use layout::{ + llvm_byte_faithful_twin, llvm_type_contains_i1, llvm_type_is_byte_faithful, + llvm_type_size_align, mir_element_stride, mir_type_abi_align, natural_struct_layout, +}; +pub(crate) use pointer_storage::{ + llvm_packed_struct_contains_pointer_in_address_space, + llvm_type_contains_pointer_in_address_space, +}; +pub use struct_layout::struct_value_lowering_is_byte_faithful; +pub(crate) use struct_layout::{StructLayoutInfo, StructSlotMap, build_struct_slot_map}; +pub(crate) use union_storage::build_union_storage_type; diff --git a/crates/mir-lower/src/convert/types/pointer_storage.rs b/crates/mir-lower/src/convert/types/pointer_storage.rs new file mode 100644 index 0000000000..ab97df737d --- /dev/null +++ b/crates/mir-lower/src/convert/types/pointer_storage.rs @@ -0,0 +1,360 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Pointer provenance/overlap analysis over lowered LLVM aggregate storage. + +use llvm_export::types as llvm_types; +use pliron::context::Context; +use pliron::r#type::TypeHandle; + +use super::llvm_type_size_align; +use crate::convert::enum_payload_storage::MAX_ENUM_PAYLOAD_ARRAY_REWRITE_LEAVES; + +pub(super) fn llvm_type_contains_pointer(ctx: &Context, ty: TypeHandle) -> bool { + let ty_ref = ty.deref(ctx); + if ty_ref.is::() { + return true; + } + if let Some(array) = ty_ref.downcast_ref::() { + return llvm_type_contains_pointer(ctx, array.elem_type()); + } + if let Some(vector) = ty_ref.downcast_ref::() { + return llvm_type_contains_pointer(ctx, vector.elem_type()); + } + if let Some(struct_ty) = ty_ref.downcast_ref::() { + return struct_ty + .fields() + .any(|field| llvm_type_contains_pointer(ctx, field)); + } + false +} + +/// One pointer-valued leaf in an LLVM aggregate's physical storage. +/// +/// Offsets are absolute within the enclosing enum. Keeping the address space +/// in the identity prevents an AS0 pointer carrier from being treated as the +/// same storage as an AS1 payload pointer merely because both are eight bytes. +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)] +pub(super) struct LlvmPointerStorage { + pub(super) offset: u64, + pub(super) size: u64, + pub(super) address_space: u32, +} + +/// Why a pointer-storage walk refused a type or an overlap. +/// +/// Distinguishing bounded array expansion from genuine provenance loss keeps +/// the user-facing diagnostic specific: an oversized pointer array reports the +/// same "rewrite requires N pointer conversions" contract as the payload +/// storage gate instead of a misleading provenance error. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum PointerOverlapRejection { + /// Expanding the walked type's fixed arrays into per-leaf records would + /// exceed [`MAX_ENUM_PAYLOAD_ARRAY_REWRITE_LEAVES`]; `required` is the + /// total number of array-expanded pointer leaves. + OverArrayLeafBound { required: u64 }, + /// The type holds pointer storage this walk cannot represent, or the + /// overlap would pun a pointer against non-matching bytes. + ProvenanceLoss, +} + +/// Total pointer leaves in `ty`, through arrays and structs, with checked +/// arithmetic. Pointer vectors count zero; every walk that expands leaves +/// fails closed on them separately. +fn count_pointer_leaves(ctx: &Context, ty: TypeHandle) -> Option { + let ty_ref = ty.deref(ctx); + if ty_ref.is::() { + return Some(1); + } + if let Some(array) = ty_ref.downcast_ref::() { + return count_pointer_leaves(ctx, array.elem_type())?.checked_mul(array.size()); + } + if let Some(struct_ty) = ty_ref.downcast_ref::() { + let fields: Vec<_> = struct_ty.fields().collect(); + let mut total = 0u64; + for field in fields { + total = total.checked_add(count_pointer_leaves(ctx, field)?)?; + } + return Some(total); + } + Some(0) +} + +/// Pointer leaves that expanding `ty`'s fixed arrays contributes to a +/// pointer-storage walk. +/// +/// Leaves outside arrays are proportional to the source text and stay +/// unbounded, exactly like struct nesting in `enum_payload_storage_type`. +/// `[&T; N]` expands from three tokens into `N` records, so only +/// array-expanded leaves count against +/// [`MAX_ENUM_PAYLOAD_ARRAY_REWRITE_LEAVES`]. +fn count_array_pointer_leaves(ctx: &Context, ty: TypeHandle) -> Option { + let ty_ref = ty.deref(ctx); + if let Some(array) = ty_ref.downcast_ref::() { + return count_pointer_leaves(ctx, array.elem_type())?.checked_mul(array.size()); + } + if let Some(struct_ty) = ty_ref.downcast_ref::() { + let fields: Vec<_> = struct_ty.fields().collect(); + let mut total = 0u64; + for field in fields { + total = total.checked_add(count_array_pointer_leaves(ctx, field)?)?; + } + return Some(total); + } + Some(0) +} + +/// Record every pointer-valued leaf in `ty` at its natural LLVM byte offset. +/// +/// This is deliberately a physical-layout walk rather than a simple +/// `contains_pointer` predicate. A niche carrier may be one field inside an +/// aggregate payload, for example the pointer at byte 8 in +/// `Option<(usize, &T)>`. In that case the aggregate and the carrier overlap, +/// but they agree exactly about which bytes hold the pointer. +/// +/// Expanding an arbitrary array into one record per pointer would let a valid +/// but enormous type consume unbounded verifier memory, so the walk first +/// counts the array-expanded leaves and refuses anything over +/// [`MAX_ENUM_PAYLOAD_ARRAY_REWRITE_LEAVES`] with the bound-specific +/// rejection before allocating a single record. +fn collect_llvm_pointer_storage( + ctx: &Context, + ty: TypeHandle, + base_offset: u64, + out: &mut Vec, +) -> std::result::Result<(), PointerOverlapRejection> { + fn collect( + ctx: &Context, + ty: TypeHandle, + base_offset: u64, + out: &mut Vec, + ) -> Option<()> { + let ty_ref = ty.deref(ctx); + if let Some(pointer) = ty_ref.downcast_ref::() { + let (size, _) = llvm_type_size_align(ctx, ty)?; + out.push(LlvmPointerStorage { + offset: base_offset, + size, + address_space: pointer.address_space(), + }); + return Some(()); + } + if let Some(array) = ty_ref.downcast_ref::() { + let element_ty = array.elem_type(); + if !llvm_type_contains_pointer(ctx, element_ty) { + return Some(()); + } + let (element_size, _) = llvm_type_size_align(ctx, element_ty)?; + for index in 0..array.size() { + let element_offset = element_size.checked_mul(index)?; + collect( + ctx, + element_ty, + base_offset.checked_add(element_offset)?, + out, + )?; + } + return Some(()); + } + if let Some(vector) = ty_ref.downcast_ref::() { + // Pointer vectors carry vector-specific ABI alignment and cast + // semantics. Keep them fail-closed rather than treating them like + // fixed arrays. + return (!llvm_type_contains_pointer(ctx, vector.elem_type())).then_some(()); + } + if let Some(struct_ty) = ty_ref.downcast_ref::() { + let fields: Vec<_> = struct_ty.fields().collect(); + let packed = struct_ty.layout() == llvm_types::StructLayout::Packed; + let mut end = 0u64; + for field in fields { + let (field_size, field_align) = llvm_type_size_align(ctx, field)?; + let field_offset = if packed { + end + } else { + let field_align = field_align.max(1); + let remainder = end % field_align; + if remainder == 0 { + end + } else { + end.checked_add(field_align - remainder)? + } + }; + collect(ctx, field, base_offset.checked_add(field_offset)?, out)?; + end = field_offset.checked_add(field_size)?; + } + return Some(()); + } + + // All pointer-bearing LLVM types understood by this lowering are + // handled above. Unknown pointer containers must fail closed. + (!llvm_type_contains_pointer(ctx, ty)).then_some(()) + } + + let required = + count_array_pointer_leaves(ctx, ty).ok_or(PointerOverlapRejection::ProvenanceLoss)?; + if required > MAX_ENUM_PAYLOAD_ARRAY_REWRITE_LEAVES { + return Err(PointerOverlapRejection::OverArrayLeafBound { required }); + } + collect(ctx, ty, base_offset, out).ok_or(PointerOverlapRejection::ProvenanceLoss) +} + +/// Analyze how a slotless incoming pointer-bearing field overlaps the enum's +/// already-selected storage, and report how to back every one of its pointers +/// with a real `ptr` slot. +/// +/// Returns `Ok(extra)` when the field is representable without erasing any +/// pointer's provenance: every pointer leaf that coincides with an existing +/// claim reuses that claim, and each remaining ("extra") pointer leaf lands on +/// bytes no other claim covers, so it can be backed by its own fresh `ptr` +/// slot. `extra` lists exactly those leaves for the caller to add as claims. +/// An empty `extra` is the exact-match case (e.g. `Option<(usize, &T)>`, whose +/// single pointer already coincides with the niche carrier); a non-empty +/// `extra` is the multi-pointer payload case (e.g. `Option<(&mut [T], +/// &mut [T])>` from `split_at_mut_checked`, whose second slice pointer needs a +/// slot of its own beside the carrier). +/// +/// Returns [`PointerOverlapRejection::ProvenanceLoss`] when the field cannot +/// be represented without punning a pointer against non-pointer bits: an +/// existing pointer slot the incoming field does not also carry at the same +/// offset/size/address space, or an extra pointer leaf that would overlap an +/// existing (necessarily non-pointer) claim. That genuine pointer/integer +/// union stays fail-closed — there is no single LLVM slot type that is both +/// provenance-carrying and integer-exact. +/// +/// Returns [`PointerOverlapRejection::OverArrayLeafBound`] when a walked type +/// holds more array-expanded pointer leaves than the bounded rewrite limit, +/// so the caller can report the bound instead of a provenance error. +pub(super) fn analyze_pointer_overlap( + ctx: &Context, + incoming_offset: u64, + incoming_size: u64, + incoming_ty: TypeHandle, + colliding_claims: &[&(u64, u64, TypeHandle)], +) -> std::result::Result, PointerOverlapRejection> { + use PointerOverlapRejection::ProvenanceLoss; + + let incoming_end = incoming_offset + .checked_add(incoming_size) + .ok_or(ProvenanceLoss)?; + + let mut incoming = Vec::new(); + collect_llvm_pointer_storage(ctx, incoming_ty, incoming_offset, &mut incoming)?; + + let mut existing = Vec::new(); + for &&(offset, _size, claim_ty) in colliding_claims { + let mut regions = Vec::new(); + collect_llvm_pointer_storage(ctx, claim_ty, offset, &mut regions)?; + existing.extend(regions.into_iter().filter(|region| { + let Some(region_end) = region.offset.checked_add(region.size) else { + return true; + }; + region.offset < incoming_end && incoming_offset < region_end + })); + } + + // Every existing pointer leaf overlapping this field must be one the field + // also carries at the same offset/size/address space; otherwise a pointer + // slot would be backed by non-matching bytes (an address-space mismatch, or + // a pointer claim where the field has integer bits). + for leaf in &existing { + if !incoming.contains(leaf) { + return Err(ProvenanceLoss); + } + } + + // Each incoming pointer leaf that does not reuse an existing claim needs its + // own slot, and may only get one if it lands on otherwise-unclaimed bytes. A + // leaf overlapping a claim it did not match is a pointer/non-pointer pun and + // stays fail-closed. + let mut extra = Vec::new(); + for leaf in &incoming { + if existing.contains(leaf) { + continue; + } + let leaf_end = leaf.offset.checked_add(leaf.size).ok_or(ProvenanceLoss)?; + let overlaps_claim = colliding_claims.iter().any(|&&(o, s, _)| { + let Some(claim_end) = o.checked_add(s) else { + return true; + }; + o < leaf_end && leaf.offset < claim_end + }); + if overlaps_claim { + return Err(ProvenanceLoss); + } + extra.push(*leaf); + } + + Ok(extra) +} + +pub(crate) fn llvm_type_contains_pointer_in_address_space( + ctx: &Context, + ty: TypeHandle, + address_space: u32, +) -> bool { + let ty_ref = ty.deref(ctx); + if let Some(pointer) = ty_ref.downcast_ref::() { + return pointer.address_space() == address_space; + } + if let Some(array) = ty_ref.downcast_ref::() { + return llvm_type_contains_pointer_in_address_space(ctx, array.elem_type(), address_space); + } + if let Some(vector) = ty_ref.downcast_ref::() { + return llvm_type_contains_pointer_in_address_space(ctx, vector.elem_type(), address_space); + } + if let Some(struct_ty) = ty_ref.downcast_ref::() { + return struct_ty + .fields() + .any(|field| llvm_type_contains_pointer_in_address_space(ctx, field, address_space)); + } + false +} + +/// Whether a physical by-value image contains a packed struct whose bytes +/// include a pointer in `address_space`. +/// +/// Direct pointers in an unpacked aggregate do not trigger this predicate: +/// their ABI can preserve address-space semantics without observing the +/// pointer's raw storage width. A pointer nested anywhere under an LLVM packed +/// struct does trigger it because whole-value construction/load/store/ABI +/// traffic observes that packed physical image. This distinction keeps normal +/// AS3 aggregate handling intact while rejecting the target-dependent packed +/// case (modern NVVM p3:32 versus 64-bit PTX/legacy). +pub(crate) fn llvm_packed_struct_contains_pointer_in_address_space( + ctx: &Context, + ty: TypeHandle, + address_space: u32, +) -> bool { + let ty_ref = ty.deref(ctx); + if let Some(array) = ty_ref.downcast_ref::() { + return llvm_packed_struct_contains_pointer_in_address_space( + ctx, + array.elem_type(), + address_space, + ); + } + if let Some(vector) = ty_ref.downcast_ref::() { + return llvm_packed_struct_contains_pointer_in_address_space( + ctx, + vector.elem_type(), + address_space, + ); + } + if let Some(struct_ty) = ty_ref.downcast_ref::() { + let fields: Vec<_> = struct_ty.fields().collect(); + if struct_ty.layout() == llvm_types::StructLayout::Packed + && fields + .iter() + .copied() + .any(|field| llvm_type_contains_pointer_in_address_space(ctx, field, address_space)) + { + return true; + } + return fields.into_iter().any(|field| { + llvm_packed_struct_contains_pointer_in_address_space(ctx, field, address_space) + }); + } + false +} diff --git a/crates/mir-lower/src/convert/types/struct_layout.rs b/crates/mir-lower/src/convert/types/struct_layout.rs new file mode 100644 index 0000000000..d477a60130 --- /dev/null +++ b/crates/mir-lower/src/convert/types/struct_layout.rs @@ -0,0 +1,619 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Struct slot mapping: `StructLayoutInfo`, `StructSlotMap`, and +//! `build_struct_slot_map` (single source of truth, issue #128). + +use dialect_mir::types::{MirStructType, MirTupleType}; +use llvm_export::types as llvm_types; +use pliron::context::Context; +use pliron::r#type::TypeHandle; + +use super::layout::{make_padding_type, mir_stored_size, natural_layout_walk}; +use super::{convert_type, is_zero_sized_type, llvm_type_size_align}; + +// ============================================================================= +// Struct Slot Mapping (single source of truth, issue #128) +// ============================================================================= + +/// Declaration-order layout facts for one MIR aggregate, in the exact form +/// [`build_struct_slot_map`] consumes. +/// +/// Extracting this owned carrier first (and dropping the `Ref` returned by +/// `Ptr::deref`) keeps the borrow checker happy: the slot-map build needs +/// `&mut Context` for type interning. +pub(crate) struct StructLayoutInfo { + /// Field types in declaration order. + pub field_types: Vec, + /// Memory order: `mem_to_decl[mem_idx] = decl_idx`. Always full length + /// (identity when rustc did not reorder). + pub mem_to_decl: Vec, + /// Byte offset of each field in declaration order; empty when rustc + /// layout is unknown. + pub field_offsets: Vec, + /// Total size in bytes including trailing padding; 0 when unknown. + pub total_size: u64, +} + +impl StructLayoutInfo { + /// Layout facts of a `MirStructType`. + pub(crate) fn of_struct(s: &MirStructType) -> Self { + StructLayoutInfo { + field_types: s.field_types.clone(), + mem_to_decl: s.memory_order(), + field_offsets: s.field_offsets().to_vec(), + total_size: s.total_size(), + } + } + + /// Layout facts of a `MirTupleType`. + /// + /// Tuples translated from a rustc type carry rustc's exact layout + /// (offsets, memory order, size), which is consumed here identically to + /// structs, so reordered tuples like `(u32, &T)` lower byte-correctly. + /// Only synthetic layout-less tuples (the unit tuple, hand-built test + /// types) fall back to LLVM natural layout. + pub(crate) fn of_tuple(t: &MirTupleType) -> Self { + StructLayoutInfo { + field_types: t.get_types().to_vec(), + mem_to_decl: t.memory_order(), + field_offsets: t.field_offsets().to_vec(), + total_size: t.total_size(), + } + } +} + +/// One lowered LLVM struct plus the value-level slot mapping into it. +/// +/// [`build_struct_slot_map`] produces the struct type and the index map in +/// the same walk, so every op that indexes into the struct (`insertvalue`, +/// `extractvalue`, GEP, call-boundary flatten/reconstruct) shares the type +/// converter's view of where each field landed. Computing the indices +/// separately is how the issue #128 class of bug (indices that ignore the +/// `[N x i8]` padding slots) happened. +pub(crate) struct StructSlotMap { + /// The final LLVM struct type, including any `[N x i8]` padding slots. + pub llvm_struct_ty: TypeHandle, + /// `decl_to_llvm[decl_idx]` = LLVM slot of that declaration-order field; + /// `None` when the field is zero-sized and was stripped. + pub decl_to_llvm: Vec>, + /// Converted LLVM type of each declaration-order field (ZSTs included). + pub field_llvm_types: Vec, + /// Natural (non-packed) LLVM byte offset of every slot of + /// `llvm_struct_ty`, padding slots included; `None` when some field type + /// cannot be sized ([`llvm_type_size_align`] declined). + pub natural_slot_offsets: Option>, + /// True when rustc's recorded byte layout cannot be represented by a + /// naturally laid-out LLVM struct: some field's natural slot offset differs + /// from rustc's byte offset, or the natural struct size differs from + /// rustc's total size. `repr(packed)` is the canonical case. Address-path + /// consumers retain this signal and the natural offsets so the #859 + /// byte-GEP fallback stays stable for packed field projections. + pub layout_diverges: bool, + /// True when the selected LLVM struct representation reproduces rustc's + /// recorded field offsets and total size for by-value movement. Natural + /// layouts are faithful when they do not diverge; divergent layouts are + /// faithful only when a sequential LLVM packed struct can express them. + /// Overlapping/union-like legacy struct models therefore remain false. + pub by_value_layout_faithful: bool, +} + +/// Whether a `MirStructType` value lowers to a byte-faithful LLVM struct: +/// either its natural LLVM layout already matches rustc's offsets and total +/// size, or the packed-with-explicit-padding representation reproduces them +/// exactly. +/// +/// Public only as a coupling oracle for mir-importer: its constant-promotion +/// gate must never admit a struct layout whose converted storage falls back +/// to a divergent natural layout, because a promoted constant's byte image +/// would then disagree with every typed read through the converted struct. +/// mir-importer asserts agreement against this function in its tests. +pub fn struct_value_lowering_is_byte_faithful( + ctx: &mut Context, + struct_ty: TypeHandle, +) -> Result { + let layout = { + let ty_ref = struct_ty.deref(ctx); + let mir_struct = ty_ref + .downcast_ref::() + .ok_or_else(|| anyhow::anyhow!("expected a MirStructType"))?; + StructLayoutInfo::of_struct(mir_struct) + }; + Ok(build_struct_slot_map(ctx, &layout)?.by_value_layout_faithful) +} + +/// Lower a struct/tuple layout to its LLVM struct type and slot map. +/// +/// When rustc layout is present (`field_offsets` non-empty and +/// `total_size > 0`), fields are placed at their exact byte offsets with +/// explicit `[N x i8]` padding slots in between, plus a trailing pad up to +/// `total_size`. This makes the layout independent of LLVM's datalayout +/// and so ABI-identical to what rustc computed on the host. For +/// `struct Extreme { a: u8, b: i128 }` where rustc puts `b` at offset 0 +/// and `a` at offset 16 with total size 32, we build: +/// +/// ```text +/// { i128, i8, [15 x i8] } ; slots: b = 0, a = 1, pad = 2 +/// ``` +/// +/// Without rustc layout, fields are emitted in memory order with no +/// padding. On both paths zero-sized fields (e.g. `PhantomData`) are +/// stripped, because NVPTX rejects empty types; stripped fields get +/// `None` in `decl_to_llvm`. +/// +/// Malformed layout metadata (a `mem_to_decl` that is not a permutation, +/// or an offsets vector of the wrong length) is rejected loudly: guessing +/// here would scramble every downstream field access. +pub(crate) fn build_struct_slot_map( + ctx: &mut Context, + layout: &StructLayoutInfo, +) -> Result { + let num_fields = layout.field_types.len(); + + if layout.mem_to_decl.len() != num_fields { + return Err(anyhow::anyhow!( + "struct slot map: memory order has {} entries but the struct has {} fields", + layout.mem_to_decl.len(), + num_fields + )); + } + let mut seen = vec![false; num_fields]; + for &decl_idx in &layout.mem_to_decl { + if decl_idx >= num_fields || seen[decl_idx] { + return Err(anyhow::anyhow!( + "struct slot map: memory order {:?} is not a permutation of 0..{}", + layout.mem_to_decl, + num_fields + )); + } + seen[decl_idx] = true; + } + let has_explicit_layout = !layout.field_offsets.is_empty() && layout.total_size > 0; + if has_explicit_layout && layout.field_offsets.len() != num_fields { + return Err(anyhow::anyhow!( + "struct slot map: {} field offsets for {} fields", + layout.field_offsets.len(), + num_fields + )); + } + + // Convert every field up front, in declaration order. + let mut field_llvm_types = Vec::with_capacity(num_fields); + for &field_ty in &layout.field_types { + field_llvm_types.push(convert_type(ctx, field_ty)?); + } + + let mut llvm_fields: Vec = Vec::new(); + let mut decl_to_llvm: Vec> = vec![None; num_fields]; + let mut current_offset: u64 = 0; + + // Place fields in memory order. + for &decl_idx in &layout.mem_to_decl { + let llvm_ty = field_llvm_types[decl_idx]; + + // ZST fields are stripped: no slot, no offset advance (rustc gives + // them size 0). + if is_zero_sized_type(ctx, llvm_ty) { + continue; + } + + if has_explicit_layout { + // Insert padding if needed to reach the rustc field offset. + let target_offset = layout.field_offsets[decl_idx]; + if current_offset < target_offset { + let padding_ty = make_padding_type(ctx, target_offset - current_offset); + llvm_fields.push(padding_ty); + current_offset = target_offset; + } + } + + decl_to_llvm[decl_idx] = Some(llvm_fields.len() as u32); + llvm_fields.push(llvm_ty); + + if has_explicit_layout { + // Offset advance is exact or an error, never guessed. Prefer + // rustc's stored size for the field: nested aggregates carry + // interior/trailing padding the converted type cannot always + // reproduce, and a wrong advance here either forces interior + // padding where rustc has none or overshoots the next field's + // offset. Scalars and other rustc-silent types fall back to the + // LLVM natural size. + let advance = match mir_stored_size(ctx, layout.field_types[decl_idx]) { + Some(size) => size, + None => llvm_type_size_align(ctx, llvm_ty) + .map(|(size, _)| size) + .ok_or_else(|| { + anyhow::anyhow!( + "struct slot map: field {decl_idx} lowers to `{}`, which has no \ + exact size; refusing to guess the next field offset", + llvm_ty.deref(ctx).disp(ctx) + ) + })?, + }; + current_offset += advance; + } + } + + // Add trailing padding to reach total_size. + if has_explicit_layout && current_offset < layout.total_size { + let padding_ty = make_padding_type(ctx, layout.total_size - current_offset); + llvm_fields.push(padding_ty); + } + + // The explicit `[N x i8]` gaps above can only ADD bytes; they cannot make + // LLVM place a slot earlier than its natural alignment allows. So when + // rustc's layout is tighter than natural (repr(packed)), the struct built + // here is a lie at the byte level, and every consumer needs to know. + // Unsizable field types leave no verdict: consumers that need the + // comparison (field_addr with a recorded rustc offset) fail closed at + // their own sites. + let walk = natural_layout_walk(ctx, &llvm_fields); + let natural_slot_offsets = walk.as_ref().map(|(offsets, _, _)| offsets.clone()); + let mut layout_diverges = false; + if has_explicit_layout && let Some((offsets, end, align)) = &walk { + for (decl_idx, slot) in decl_to_llvm.iter().enumerate() { + if let Some(slot) = slot + && offsets[*slot as usize] != layout.field_offsets[decl_idx] + { + layout_diverges = true; + } + } + let natural_size = end.div_ceil(*align) * *align; + if natural_size != layout.total_size { + layout_diverges = true; + } + } + + // A packed LLVM struct is sequential with alignment 1. That faithfully + // represents repr(packed) only when the sequential packed offsets and + // final byte count exactly match rustc's metadata. A divergent layout can + // also be an old union-like/overlapping model; selecting Packed for such a + // shape would merely turn one incorrect sequential layout into another. + let packed_walk = if layout_diverges && has_explicit_layout { + let mut offsets = Vec::with_capacity(llvm_fields.len()); + let mut end = 0u64; + let mut sizeable = true; + for &field in &llvm_fields { + offsets.push(end); + let Some((field_size, _)) = llvm_type_size_align(ctx, field) else { + sizeable = false; + break; + }; + let Some(next) = end.checked_add(field_size) else { + sizeable = false; + break; + }; + end = next; + } + sizeable.then_some((offsets, end)) + } else { + None + }; + let packed_representable = packed_walk.as_ref().is_some_and(|(offsets, end)| { + *end == layout.total_size + && decl_to_llvm + .iter() + .enumerate() + .all(|(decl_idx, slot)| match slot { + Some(slot) => offsets[*slot as usize] == layout.field_offsets[decl_idx], + None => true, + }) + }); + + let struct_layout = if packed_representable { + llvm_types::StructLayout::Packed + } else { + llvm_types::StructLayout::Unpacked + }; + let llvm_struct_ty: TypeHandle = + llvm_types::StructType::get_unnamed(ctx, (llvm_fields, struct_layout)).into(); + let by_value_layout_faithful = !layout_diverges || packed_representable; + + Ok(StructSlotMap { + llvm_struct_ty, + decl_to_llvm, + field_llvm_types, + natural_slot_offsets, + layout_diverges, + by_value_layout_faithful, + }) +} + +#[cfg(test)] +mod tests { + //! Hardware-free unit tests for [`build_struct_slot_map`]: the slot map + //! and the LLVM struct type are produced by the same walk, so these + //! tests pin down both for the layout shapes from issue #128. + + use super::super::test_support::{llvm_int, make_ctx, mir_uint, mir_zst, pad, struct_fields}; + use super::*; + use dialect_mir::types::{EnumVariant, MirEnumType}; + + #[test] + fn slot_map_reorder_only() { + let mut ctx = make_ctx(); + // struct { a: u8, b: u64 }, memory order [b, a], no rustc offsets. + let a = mir_uint(&mut ctx, 8); + let b = mir_uint(&mut ctx, 64); + let layout = StructLayoutInfo { + field_types: vec![a, b], + mem_to_decl: vec![1, 0], + field_offsets: vec![], + total_size: 0, + }; + let map = build_struct_slot_map(&mut ctx, &layout).unwrap(); + + assert_eq!(map.decl_to_llvm, vec![Some(1), Some(0)]); + let i8s = llvm_int(&mut ctx, 8); + let i64s = llvm_int(&mut ctx, 64); + assert_eq!(struct_fields(&ctx, map.llvm_struct_ty), vec![i64s, i8s]); + } + + #[test] + fn slot_map_padding_only() { + let mut ctx = make_ctx(); + // struct { a: u8 @ 0, b: u64 @ 8 }, declaration order == memory + // order, size 16: lowers to { i8, [7 x i8], i64 }. The pad consumes + // slot 1, so b lands at slot 2 (the issue #128 sites used 1). + let a = mir_uint(&mut ctx, 8); + let b = mir_uint(&mut ctx, 64); + let layout = StructLayoutInfo { + field_types: vec![a, b], + mem_to_decl: vec![0, 1], + field_offsets: vec![0, 8], + total_size: 16, + }; + let map = build_struct_slot_map(&mut ctx, &layout).unwrap(); + + assert_eq!(map.decl_to_llvm, vec![Some(0), Some(2)]); + let i8s = llvm_int(&mut ctx, 8); + let i64s = llvm_int(&mut ctx, 64); + let pad7 = pad(&mut ctx, 7); + assert_eq!( + struct_fields(&ctx, map.llvm_struct_ty), + vec![i8s, pad7, i64s] + ); + } + + #[test] + fn slot_map_uses_packed_layout_when_natural_offsets_diverge() { + let mut ctx = make_ctx(); + let tag = mir_uint(&mut ctx, 8); + let value = mir_uint(&mut ctx, 32); + let layout = StructLayoutInfo { + field_types: vec![tag, value], + mem_to_decl: vec![0, 1], + field_offsets: vec![0, 1], + total_size: 5, + }; + + let map = build_struct_slot_map(&mut ctx, &layout).unwrap(); + + assert!(map.layout_diverges); + assert_eq!(map.decl_to_llvm, vec![Some(0), Some(1)]); + let llvm_struct_ty_ref = map.llvm_struct_ty.deref(&ctx); + let struct_ty = llvm_struct_ty_ref + .downcast_ref::() + .expect("slot map must produce an LLVM struct"); + assert_eq!(struct_ty.layout(), llvm_types::StructLayout::Packed); + assert_eq!(llvm_type_size_align(&ctx, map.llvm_struct_ty), Some((5, 1))); + } + + #[test] + fn slot_map_packed_two_keeps_explicit_padding_slot() { + let mut ctx = make_ctx(); + let tag = mir_uint(&mut ctx, 8); + let value = mir_uint(&mut ctx, 32); + let layout = StructLayoutInfo { + field_types: vec![tag, value], + mem_to_decl: vec![0, 1], + field_offsets: vec![0, 2], + total_size: 6, + }; + + let map = build_struct_slot_map(&mut ctx, &layout).unwrap(); + + assert!(map.layout_diverges); + assert_eq!(map.decl_to_llvm, vec![Some(0), Some(2)]); + let i8s = llvm_int(&mut ctx, 8); + let i32s = llvm_int(&mut ctx, 32); + let pad1 = pad(&mut ctx, 1); + assert_eq!( + struct_fields(&ctx, map.llvm_struct_ty), + vec![i8s, pad1, i32s] + ); + let llvm_struct_ty_ref = map.llvm_struct_ty.deref(&ctx); + let struct_ty = llvm_struct_ty_ref + .downcast_ref::() + .expect("slot map must produce an LLVM struct"); + assert_eq!(struct_ty.layout(), llvm_types::StructLayout::Packed); + assert_eq!(llvm_type_size_align(&ctx, map.llvm_struct_ty), Some((6, 1))); + } + + #[test] + fn slot_map_reorder_plus_padding() { + let mut ctx = make_ctx(); + // struct { a: u8 @ 8, b: u64 @ 0 }, memory order [b, a], size 16: + // lowers to { i64, i8, [7 x i8] } with a trailing pad. + let a = mir_uint(&mut ctx, 8); + let b = mir_uint(&mut ctx, 64); + let layout = StructLayoutInfo { + field_types: vec![a, b], + mem_to_decl: vec![1, 0], + field_offsets: vec![8, 0], + total_size: 16, + }; + let map = build_struct_slot_map(&mut ctx, &layout).unwrap(); + + assert_eq!(map.decl_to_llvm, vec![Some(1), Some(0)]); + let i8s = llvm_int(&mut ctx, 8); + let i64s = llvm_int(&mut ctx, 64); + let pad7 = pad(&mut ctx, 7); + assert_eq!( + struct_fields(&ctx, map.llvm_struct_ty), + vec![i64s, i8s, pad7] + ); + } + + #[test] + fn slot_map_zst_interleaving() { + let mut ctx = make_ctx(); + // struct { a: u32 @ 0, z: PhantomData @ 4, b: u32 @ 4 }, size 8. + // The ZST is stripped (no slot, no pad split): { i32, i32 }. + let a = mir_uint(&mut ctx, 32); + let z = mir_zst(&mut ctx); + let b = mir_uint(&mut ctx, 32); + let layout = StructLayoutInfo { + field_types: vec![a, z, b], + mem_to_decl: vec![0, 1, 2], + field_offsets: vec![0, 4, 4], + total_size: 8, + }; + let map = build_struct_slot_map(&mut ctx, &layout).unwrap(); + + assert_eq!(map.decl_to_llvm, vec![Some(0), None, Some(1)]); + let i32s = llvm_int(&mut ctx, 32); + assert_eq!(struct_fields(&ctx, map.llvm_struct_ty), vec![i32s, i32s]); + } + + #[test] + fn slot_map_issue128_arena_shape() { + let mut ctx = make_ctx(); + // The exact shape from issue #128 (examples/struct_field_layout): + // + // enum Layout { Aos, Soa, AoSoA(u32) } // -> { i8, i32 } + // struct Arena { layout: Layout, cap: u32, stride: u32, big: u64 } + // + // rustc layout: layout @ 0 (8 bytes), big @ 8, cap @ 16, + // stride @ 20, size 24. The enum now carries its own explicit + // internal padding, so the OUTER struct needs no extra padding slot: + // + // { { i8, [3 x i8], i32 }, i64, i32, i32 } + // layout=0 big=1 cap=2 stride=3 + let discr = mir_uint(&mut ctx, 8); + let payload = mir_uint(&mut ctx, 32); + let layout_enum: TypeHandle = MirEnumType::get_with_layout( + &mut ctx, + "Layout".into(), + discr, + vec![0, 1, 2], + vec![ + EnumVariant::unit("Aos".into()), + EnumVariant::unit("Soa".into()), + EnumVariant::new_with_layout("AoSoA".into(), vec![payload], vec![4], vec![4]), + ], + 0, + 8, + 4, + ) + .into(); + let cap = mir_uint(&mut ctx, 32); + let stride = mir_uint(&mut ctx, 32); + let big = mir_uint(&mut ctx, 64); + + let layout = StructLayoutInfo { + field_types: vec![layout_enum, cap, stride, big], + mem_to_decl: vec![0, 3, 1, 2], + field_offsets: vec![0, 16, 20, 8], + total_size: 24, + }; + let map = build_struct_slot_map(&mut ctx, &layout).unwrap(); + + assert_eq!( + map.decl_to_llvm, + vec![Some(0), Some(2), Some(3), Some(1)], + "the enum's internal pad must not create an outer struct slot" + ); + + let i8s = llvm_int(&mut ctx, 8); + let i32s = llvm_int(&mut ctx, 32); + let i64s = llvm_int(&mut ctx, 64); + let enum_pad3 = pad(&mut ctx, 3); + let enum_llvm: TypeHandle = llvm_types::StructType::get_unnamed( + &ctx, + ( + vec![i8s, enum_pad3, i32s], + llvm_types::StructLayout::Unpacked, + ), + ) + .into(); + assert_eq!( + struct_fields(&ctx, map.llvm_struct_ty), + vec![enum_llvm, i64s, i32s, i32s] + ); + } + + #[test] + fn slot_map_nested_struct_uses_stored_size() { + let mut ctx = make_ctx(); + // Inner struct whose stored rustc size (16) exceeds the sum of its + // converted LLVM field sizes (i8 + i64 = 9, no offsets stored). + // The outer walk must advance by the stored 16, reaching the next + // field's offset exactly: NO interior pad before it. + let x = mir_uint(&mut ctx, 8); + let y = mir_uint(&mut ctx, 64); + let inner: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "Inner".into(), + vec!["x".into(), "y".into()], + vec![x, y], + vec![], + vec![], + 16, + 0, + ) + .into(); + let c = mir_uint(&mut ctx, 8); + + let layout = StructLayoutInfo { + field_types: vec![inner, c], + mem_to_decl: vec![0, 1], + field_offsets: vec![0, 16], + total_size: 24, + }; + let map = build_struct_slot_map(&mut ctx, &layout).unwrap(); + + // inner = slot 0, c = slot 1 (adjacent), trailing [7 x i8] pad. + assert_eq!(map.decl_to_llvm, vec![Some(0), Some(1)]); + let fields = struct_fields(&ctx, map.llvm_struct_ty); + assert_eq!(fields.len(), 3, "exactly one (trailing) pad slot"); + let pad7 = pad(&mut ctx, 7); + assert_eq!(fields[2], pad7); + } + + #[test] + fn slot_map_rejects_malformed_memory_order() { + let mut ctx = make_ctx(); + let a = mir_uint(&mut ctx, 8); + let b = mir_uint(&mut ctx, 64); + + // Not a permutation: decl index 0 appears twice. + let dup = StructLayoutInfo { + field_types: vec![a, b], + mem_to_decl: vec![0, 0], + field_offsets: vec![], + total_size: 0, + }; + assert!(build_struct_slot_map(&mut ctx, &dup).is_err()); + + // Wrong length. + let short = StructLayoutInfo { + field_types: vec![a, b], + mem_to_decl: vec![0], + field_offsets: vec![], + total_size: 0, + }; + assert!(build_struct_slot_map(&mut ctx, &short).is_err()); + + // Offsets vector length mismatch (with explicit layout engaged). + let bad_offsets = StructLayoutInfo { + field_types: vec![a, b], + mem_to_decl: vec![0, 1], + field_offsets: vec![0], + total_size: 16, + }; + assert!(build_struct_slot_map(&mut ctx, &bad_offsets).is_err()); + } +} diff --git a/crates/mir-lower/src/convert/types/test_support.rs b/crates/mir-lower/src/convert/types/test_support.rs new file mode 100644 index 0000000000..5e47bfcb1d --- /dev/null +++ b/crates/mir-lower/src/convert/types/test_support.rs @@ -0,0 +1,65 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Shared fixtures for the `types` submodule unit tests. + +use dialect_mir::types::{MirStructType, StructAbiKind}; +use llvm_export::types as llvm_types; +use pliron::builtin::types::{IntegerType, Signedness}; +use pliron::context::Context; +use pliron::r#type::TypeHandle; + +use super::layout::make_padding_type; + +pub(super) fn make_ctx() -> Context { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + ctx +} + +/// A MIR-level unsigned integer type (what the importer produces). +pub(super) fn mir_uint(ctx: &mut Context, width: u32) -> TypeHandle { + IntegerType::get(ctx, width, Signedness::Unsigned).into() +} + +/// A converted (signless) LLVM integer type. +pub(super) fn llvm_int(ctx: &mut Context, width: u32) -> TypeHandle { + IntegerType::get(ctx, width, Signedness::Signless).into() +} + +/// `[n x i8]` padding type, as `make_padding_type` builds it. +pub(super) fn pad(ctx: &mut Context, n: u64) -> TypeHandle { + make_padding_type(ctx, n) +} + +/// A zero-sized MIR struct (PhantomData shape). +pub(super) fn mir_zst(ctx: &mut Context) -> TypeHandle { + MirStructType::get(ctx, "Phantom".into(), vec![], vec![]).into() +} + +pub(super) fn struct_fields(ctx: &Context, ty: TypeHandle) -> Vec { + ty.deref(ctx) + .downcast_ref::() + .expect("expected an LLVM struct type") + .fields() + .collect() +} + +pub(super) fn transparent_u32(ctx: &mut Context, name: &str) -> TypeHandle { + let u32_ty = mir_uint(ctx, 32); + MirStructType::get_with_full_layout_and_abi( + ctx, + name.into(), + vec!["value".into()], + vec![u32_ty], + vec![0], + vec![0], + 4, + 4, + StructAbiKind::TransparentScalar, + ) + .into() +} diff --git a/crates/mir-lower/src/convert/types/union_storage.rs b/crates/mir-lower/src/convert/types/union_storage.rs new file mode 100644 index 0000000000..aa821690ff --- /dev/null +++ b/crates/mir-lower/src/convert/types/union_storage.rs @@ -0,0 +1,276 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Byte-exact LLVM storage for Rust unions. + +use dialect_mir::types::MirUnionType; +use llvm_export::types as llvm_types; +use pliron::builtin::types::{IntegerType, Signedness}; +use pliron::context::Context; +use pliron::r#type::TypeHandle; + +use super::layout::make_padding_type; +use super::pointer_storage::llvm_type_contains_pointer; +use super::{convert_type, llvm_type_is_byte_faithful, llvm_type_size_align}; + +/// Build byte-exact LLVM storage for a Rust union. +/// +/// A union cannot be represented as an LLVM struct containing every declared +/// field: struct fields are consecutive, while union fields all start at byte +/// zero. We choose one byte-faithful field as the storage view and add explicit +/// tail bytes. A zero-length integer array raises the LLVM type's natural +/// alignment without consuming storage. Pointer-bearing fields are preferred +/// so an ordinary union copy keeps LLVM pointer provenance. +/// +/// NVPTX gives scalar integers natural alignments up to 16 bytes. Stronger +/// Rust alignment is carried explicitly on memory operations because LLVM +/// aggregate types cannot encode over-alignment; the storage type still keeps +/// the union's exact size and therefore its array stride. +pub(crate) fn build_union_storage_type( + ctx: &mut Context, + union_ty: &MirUnionType, +) -> Result { + let size = union_ty.total_size(); + let align = union_ty.abi_align(); + if align == 0 || !align.is_power_of_two() { + return Err(anyhow::anyhow!( + "union `{}` has invalid ABI alignment {}", + union_ty.name(), + align + )); + } + if size > 0 && !size.is_multiple_of(align) { + return Err(anyhow::anyhow!( + "union `{}` size {} is not a multiple of its {}-byte ABI alignment", + union_ty.name(), + size, + align + )); + } + + let mut fields = Vec::with_capacity(union_ty.field_count()); + let mut pointer_carrier: Option = None; + for (index, &field_ty) in union_ty.field_types().iter().enumerate() { + let llvm_field_ty = convert_type(ctx, field_ty)?; + let (field_size, field_align) = + llvm_type_size_align(ctx, llvm_field_ty).ok_or_else(|| { + anyhow::anyhow!( + "union `{}` field {} has unsupported LLVM size/alignment", + union_ty.name(), + index + ) + })?; + if field_size > size { + return Err(anyhow::anyhow!( + "union `{}` field {} lowers to {} bytes but the union is only {} bytes", + union_ty.name(), + index, + field_size, + size + )); + } + if field_align > align { + return Err(anyhow::anyhow!( + "union `{}` field {} lowers with alignment {} but rustc reports union alignment {}", + union_ty.name(), + index, + field_align, + align + )); + } + let contains_pointer = llvm_type_contains_pointer(ctx, llvm_field_ty); + if contains_pointer { + if let Some(first) = pointer_carrier + && first != llvm_field_ty + { + return Err(anyhow::anyhow!( + "union `{}` has pointer-bearing fields with different LLVM representations; preserving provenance for that shape is not yet supported", + union_ty.name() + )); + } + pointer_carrier = Some(llvm_field_ty); + } + fields.push((llvm_field_ty, field_size, field_align, contains_pointer)); + } + + let storage_align = align.min(16); + let anchor_int = IntegerType::get(ctx, (storage_align * 8) as u32, Signedness::Signless); + let anchor: TypeHandle = llvm_types::ArrayType::get(ctx, anchor_int.into(), 0).into(); + let mut storage_fields = vec![anchor]; + if size > 0 { + let representative = fields + .iter() + .filter(|(ty, field_size, _, _)| { + *field_size > 0 + && llvm_type_is_byte_faithful(ctx, *ty) + && (pointer_carrier.is_none() || llvm_type_contains_pointer(ctx, *ty)) + }) + .max_by_key(|(_, field_size, field_align, contains_pointer)| { + (*contains_pointer, *field_align, *field_size) + }); + if let Some(representative) = representative { + storage_fields.push(representative.0); + if representative.1 < size { + storage_fields.push(make_padding_type(ctx, size - representative.1)); + } + } else if pointer_carrier.is_some() { + return Err(anyhow::anyhow!( + "union `{}` has pointer-bearing fields but no byte-faithful pointer carrier; lowering it as raw bytes would discard pointer provenance", + union_ty.name() + )); + } else { + // Pointer-free unions may safely use raw bytes as their SSA + // carrier. Field loads/stores still use their declared types. + storage_fields.push(make_padding_type(ctx, size)); + } + } + let storage: TypeHandle = llvm_types::StructType::get_unnamed( + ctx, + (storage_fields, llvm_types::StructLayout::Unpacked), + ) + .into(); + let (llvm_size, llvm_align) = llvm_type_size_align(ctx, storage).ok_or_else(|| { + anyhow::anyhow!( + "union `{}` storage has unsupported LLVM layout", + union_ty.name() + ) + })?; + if llvm_size != size || llvm_align > align { + return Err(anyhow::anyhow!( + "union `{}` storage lowered to incompatible size/alignment {}/{} but rustc requires {}/{}", + union_ty.name(), + llvm_size, + llvm_align, + size, + align + )); + } + Ok(storage) +} + +#[cfg(test)] +mod tests { + use super::super::test_support::{make_ctx, mir_uint, struct_fields}; + use super::*; + use dialect_mir::types::{MirArrayType, MirPtrType, MirStructType}; + + #[test] + fn union_storage_has_exact_size_alignment_and_stride() { + let mut ctx = make_ctx(); + let u8_ty = mir_uint(&mut ctx, 8); + let u32_ty = mir_uint(&mut ctx, 32); + let bytes_ty: TypeHandle = MirArrayType::get(&mut ctx, u8_ty, 4).into(); + let union_ty = MirUnionType::get( + &mut ctx, + "Bits".into(), + vec!["word".into(), "bytes".into()], + vec![u32_ty, bytes_ty], + 4, + 4, + ); + let union_data = union_ty.deref(&ctx).clone(); + let storage = build_union_storage_type(&mut ctx, &union_data).unwrap(); + assert_eq!(llvm_type_size_align(&ctx, storage), Some((4, 4))); + + let union_handle: TypeHandle = union_ty.into(); + let array: TypeHandle = MirArrayType::get(&mut ctx, union_handle, 3).into(); + let llvm_array = convert_type(&mut ctx, array).unwrap(); + assert_eq!(llvm_type_size_align(&ctx, llvm_array), Some((12, 4))); + } + + #[test] + fn union_storage_prefers_pointer_carrier() { + let mut ctx = make_ctx(); + let u32_ty = mir_uint(&mut ctx, 32); + let u64_ty = mir_uint(&mut ctx, 64); + let ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, u32_ty, false).into(); + let union_ty = MirUnionType::get( + &mut ctx, + "PointerBits".into(), + vec!["ptr".into(), "bits".into()], + vec![ptr_ty, u64_ty], + 8, + 8, + ); + let union_data = union_ty.deref(&ctx).clone(); + let storage = build_union_storage_type(&mut ctx, &union_data).unwrap(); + let fields = struct_fields(&ctx, storage); + assert!(fields[1].deref(&ctx).is::()); + assert_eq!(llvm_type_size_align(&ctx, storage), Some((8, 8))); + } + + #[test] + fn union_storage_rejects_incompatible_pointer_address_spaces() { + let mut ctx = make_ctx(); + let u32_ty = mir_uint(&mut ctx, 32); + let generic: TypeHandle = MirPtrType::get_generic(&mut ctx, u32_ty, false).into(); + let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, u32_ty, false).into(); + let union_ty = MirUnionType::get( + &mut ctx, + "MixedPointers".into(), + vec!["generic".into(), "shared".into()], + vec![generic, shared], + 8, + 8, + ); + let union_data = union_ty.deref(&ctx).clone(); + let err = build_union_storage_type(&mut ctx, &union_data).unwrap_err(); + assert!(err.to_string().contains("different LLVM representations")); + } + + #[test] + fn union_storage_rejects_non_byte_faithful_pointer_carrier() { + let mut ctx = make_ctx(); + let u8_ty = mir_uint(&mut ctx, 8); + let u32_ty = mir_uint(&mut ctx, 32); + let bool_ty = mir_uint(&mut ctx, 1); + let ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, u32_ty, false).into(); + let ptr_bool: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "PtrBool".into(), + vec!["ptr".into(), "flag".into()], + vec![ptr_ty, bool_ty], + vec![0, 1], + vec![0, 8], + 16, + 8, + ) + .into(); + let bytes_ty: TypeHandle = MirArrayType::get(&mut ctx, u8_ty, 16).into(); + let union_ty = MirUnionType::get( + &mut ctx, + "PointerBoolBytes".into(), + vec!["view".into(), "bytes".into()], + vec![ptr_bool, bytes_ty], + 16, + 8, + ); + let union_data = union_ty.deref(&ctx).clone(); + let err = build_union_storage_type(&mut ctx, &union_data).unwrap_err(); + assert!(err.to_string().contains("no byte-faithful pointer carrier")); + } + + #[test] + fn union_storage_preserves_over_aligned_size_and_stride() { + let mut ctx = make_ctx(); + let u32_ty = mir_uint(&mut ctx, 32); + let union_ty = MirUnionType::get( + &mut ctx, + "OverAligned".into(), + vec!["word".into()], + vec![u32_ty], + 32, + 32, + ); + let union_data = union_ty.deref(&ctx).clone(); + let storage = build_union_storage_type(&mut ctx, &union_data).unwrap(); + assert_eq!(llvm_type_size_align(&ctx, storage), Some((32, 16))); + + let union_handle: TypeHandle = union_ty.into(); + let array: TypeHandle = MirArrayType::get(&mut ctx, union_handle, 3).into(); + let llvm_array = convert_type(&mut ctx, array).unwrap(); + assert_eq!(llvm_type_size_align(&ctx, llvm_array), Some((96, 16))); + } +} From d7fea9ff95f655a87107de6fcbdb5e313aeb68f0 Mon Sep 17 00:00:00 2001 From: nihalpasham Date: Fri, 28 Aug 2026 22:16:58 +0530 Subject: [PATCH 06/18] refactor(mir-lower): split ops/aggregate.rs into ops/aggregate/ modules --- crates/mir-lower/src/convert/ops/aggregate.rs | 4967 ----------------- .../src/convert/ops/aggregate/addressing.rs | 1185 ++++ .../convert/ops/aggregate/array_extract.rs | 351 ++ .../src/convert/ops/aggregate/common.rs | 99 + .../src/convert/ops/aggregate/construct.rs | 1001 ++++ .../src/convert/ops/aggregate/enum_layout.rs | 220 + .../src/convert/ops/aggregate/enums.rs | 1704 ++++++ .../src/convert/ops/aggregate/fields.rs | 489 ++ .../src/convert/ops/aggregate/mod.rs | 59 + .../src/convert/ops/aggregate/test_support.rs | 109 + 10 files changed, 5217 insertions(+), 4967 deletions(-) delete mode 100644 crates/mir-lower/src/convert/ops/aggregate.rs create mode 100644 crates/mir-lower/src/convert/ops/aggregate/addressing.rs create mode 100644 crates/mir-lower/src/convert/ops/aggregate/array_extract.rs create mode 100644 crates/mir-lower/src/convert/ops/aggregate/common.rs create mode 100644 crates/mir-lower/src/convert/ops/aggregate/construct.rs create mode 100644 crates/mir-lower/src/convert/ops/aggregate/enum_layout.rs create mode 100644 crates/mir-lower/src/convert/ops/aggregate/enums.rs create mode 100644 crates/mir-lower/src/convert/ops/aggregate/fields.rs create mode 100644 crates/mir-lower/src/convert/ops/aggregate/mod.rs create mode 100644 crates/mir-lower/src/convert/ops/aggregate/test_support.rs diff --git a/crates/mir-lower/src/convert/ops/aggregate.rs b/crates/mir-lower/src/convert/ops/aggregate.rs deleted file mode 100644 index e60d3e5214..0000000000 --- a/crates/mir-lower/src/convert/ops/aggregate.rs +++ /dev/null @@ -1,4967 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -//! Aggregate operation conversion: `dialect-mir` → LLVM dialect. -//! -//! Converts `dialect-mir` aggregate operations (structs, tuples, enums) to -//! their LLVM dialect equivalents. -//! -//! # Operations -//! -//! | MIR Operation | LLVM Operation(s) | Description | -//! |--------------------------|--------------------------------------|------------------------| -//! | `mir.extract_field` | `llvm.extractvalue` | Get struct/tuple field | -//! | `mir.insert_field` | `llvm.insertvalue` | Set struct/tuple field | -//! | `mir.construct_struct` | `llvm.undef` + `llvm.insertvalue` | Build struct | -//! | `mir.construct_tuple` | `llvm.undef` + `llvm.insertvalue` | Build tuple | -//! | `mir.construct_slice` | `llvm.undef` + `llvm.insertvalue` | Build slice fat ptr | -//! | `mir.construct_enum` | `llvm.undef` + `llvm.insertvalue` | Build enum | -//! | `mir.get_discriminant` | `llvm.extractvalue` | Get enum tag | -//! | `mir.set_discriminant` | `llvm.getelementptr` + `llvm.store` | Write enum tag | -//! | `mir.enum_payload` | `llvm.extractvalue` | Get enum payload | -//! -//! # Enum Representation -//! -//! Enums use rustc's physical layout, not a cuda-oxide-only tagged struct. -//! `build_enum_slot_map` places the direct tag or niche carrier and every -//! payload at rustc's byte offsets, reuses identical overlapping storage, and -//! routes differently typed overlaps through byte-addressed memory. `Single` -//! and `Empty` layouts have no carrier at all. -//! -//! A direct tag holds the variant's DECLARED discriminant value, not its -//! position. For `core::cmp::Ordering`, `Less` therefore stores -1 (the i8 bit -//! pattern 255), `Equal` stores 0, and `Greater` stores 1. A niche layout -//! instead uses rustc's wrapping `niche_start + variant_offset` encoding and -//! introduces no extra tag. - -use crate::convert::enum_payload_storage::{coerce_enum_payload_value, enum_payload_storage_type}; -use crate::convert::types::{ - EnumSlotMap, StructLayoutInfo, StructSlotMap, build_enum_slot_map, build_struct_slot_map, - build_union_storage_type, convert_type, is_zero_sized_type, llvm_byte_faithful_twin, - llvm_packed_struct_contains_pointer_in_address_space, llvm_type_contains_i1, - llvm_type_size_align, make_slice_struct, mir_element_stride, mir_type_abi_align, - packed_shared_internal_abi_info, -}; -use dialect_mir::ops::{ - MirConstantOp, MirConstructEnumOp, MirEnumPayloadOp, MirExtractFieldOp, MirFieldAddrOp, - MirInsertFieldOp, MirSetDiscriminantOp, -}; -use dialect_mir::types::{ - EnumCarrierKind, EnumLayoutKind, MirArrayType, MirDisjointSliceType, MirEnumType, MirPtrType, - MirSliceType, MirStructType, MirTupleType, MirUnionType, -}; -use llvm_export::attributes::{ICmpPredicateAttr, IntegerOverflowFlagsAttr}; -use llvm_export::op_interfaces::{ - CastOpInterface, CastOpWithNNegInterface, IntBinArithOpWithOverflowFlag, -}; -use llvm_export::ops as llvm; -use llvm_export::types as llvm_types; -use pliron::builtin::types::{IntegerType, Signedness}; -use pliron::context::{Context, Ptr}; -use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; -use pliron::irbuild::inserter::Inserter; -use pliron::irbuild::rewriter::Rewriter; -use pliron::location::Located; -use pliron::op::Op; -use pliron::operation::Operation; -use pliron::result::Result; -use pliron::r#type::{TypeHandle, Typed}; -use pliron::utils::apint::APInt; -use pliron::value::Value; -use std::num::NonZeroUsize; - -fn anyhow_to_pliron(e: anyhow::Error) -> pliron::result::Error { - pliron::input_error_noloc!("{e}") -} - -/// How the MIR-level field indices of an aggregate operand map onto the -/// lowered LLVM aggregate. -enum AggregateSlots { - /// Lowered from a `MirStructType`/`MirTupleType`: use the slot map the - /// type converter built (accounts for reordering, `[N x i8]` padding - /// slots and stripped ZST fields). - Mapped(StructSlotMap), - /// The MIR index is already the final LLVM index. Sound only for - /// aggregates whose lowered layout is index-preserving by construction: - /// arrays and slice fat pointers (`{ ptr, i64 }`). - Identity, -} - -/// Resolve how field indices of `aggregate` map onto its lowered type. -/// -/// Recover-or-error (issue #128): when the operand has no recorded -/// `MirStructType`/`MirTupleType` conversion history, identity indexing is -/// only sound for aggregates the converter lowers without reordering, -/// padding, or ZST stripping: arrays and slice fat pointers. Anything -/// else is a lowering bug; guessing identity there silently reads or -/// writes the wrong field, so we error out loudly instead. -fn resolve_aggregate_slots( - ctx: &mut Context, - operands_info: &OperandsInfo, - aggregate: Value, -) -> Result { - let layout = operands_info - .lookup_most_recent_of_type::(ctx, aggregate) - .map(|struct_ref| StructLayoutInfo::of_struct(&struct_ref)) - .or_else(|| { - operands_info - .lookup_most_recent_of_type::(ctx, aggregate) - .map(|tuple_ref| StructLayoutInfo::of_tuple(&tuple_ref)) - }); - - if let Some(layout) = layout { - let map = build_struct_slot_map(ctx, &layout).map_err(anyhow_to_pliron)?; - return Ok(AggregateSlots::Mapped(map)); - } - - // Arrays keep their element indices: `[N x T]` has no reorder, no - // padding, no ZST stripping. - let is_array_history = operands_info - .lookup_most_recent_of_type::(ctx, aggregate) - .is_some(); - // Slices lower to the `{ ptr, i64 }` fat pointer, where index 0 = ptr - // and index 1 = len by construction. - let is_slice_history = operands_info - .lookup_most_recent_of_type::(ctx, aggregate) - .is_some() - || operands_info - .lookup_most_recent_of_type::(ctx, aggregate) - .is_some(); - if is_array_history || is_slice_history { - return Ok(AggregateSlots::Identity); - } - - // No conversion history at all (e.g. a slice reconstructed in the entry - // prologue, which is born as an LLVM struct). Identity is still fine if - // the current type is a slice fat pointer or an LLVM array. A disjoint - // slice with a runtime row width is the same case one word longer: every - // runtime index space today carries a single `u32` row width, and the struct the type - // converter builds for it (`make_disjoint_slice_struct`) is - // index-preserving by construction, exactly like the two-field fat - // pointer. Both shape checks share the two-field caveat that an unnamed - // user struct with the identical lowered layout would be accepted too; - // that has been the accepted trade-off for `{ ptr, i64 }` since issue - // #128, and a struct needing a slot map still errors below whenever its - // lowered shape differs (padding slots, reordering, other field types). - let aggregate_ty = aggregate.get_type(ctx); - let slice_struct_ty = make_slice_struct(ctx); - let row_width_slice_struct_ty = { - // The one runtime layout current `SpaceLayout` impls produce: a - // single u32 row width. Built through the real constructor so this - // check cannot drift from the lowered shape. A future space with a - // different `Data` shape is not listed here and thus fails closed - // into the refuse-to-guess error until this site learns it. - let width_ty: TypeHandle = IntegerType::get(ctx, 32, Signedness::Unsigned).into(); - crate::convert::types::make_disjoint_slice_struct(ctx, &[width_ty]) - .map_err(anyhow_to_pliron)? - }; - let is_llvm_array = aggregate_ty - .deref(ctx) - .is::(); - if aggregate_ty == slice_struct_ty || aggregate_ty == row_width_slice_struct_ty || is_llvm_array - { - return Ok(AggregateSlots::Identity); - } - - let ty_disp = aggregate_ty.deref(ctx).disp(ctx).to_string(); - pliron::input_err_noloc!( - "Cannot map field indices for aggregate of type {ty_disp}: no struct/tuple \ - conversion history was recorded for this operand, and identity indexing is \ - only sound for arrays and slice fat pointers (with or without the runtime \ - row-width word). Refusing to guess a field mapping (issue #128)." - ) -} - -/// Convert `mir.extract_field` to `llvm.extractvalue`. -/// -/// Handles scalar-lowered newtype case: if the operand is a scalar (e.g., `ThreadIndex`), -/// no extraction is needed. -/// -/// The declaration-order field index is mapped to the LLVM slot via -/// [`resolve_aggregate_slots`], which shares the type converter's view of -/// the struct (reorder, `[N x i8]` padding slots, stripped ZSTs). If -/// extracting a ZST field, we return undef of its (empty) type. -pub(crate) fn convert_extract_field( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - operands_info: &OperandsInfo, -) -> Result<()> { - let aggregate = op.deref(ctx).get_operand(0); - - let extract_op = MirExtractFieldOp::new(op); - let decl_index = match extract_op.get_attr_index(ctx) { - Some(attr) => attr.0 as usize, - None => return pliron::input_err_noloc!("Missing index attribute on extract_field"), - }; - - if operands_info - .lookup_most_recent_of_type::(ctx, aggregate) - .is_some() - { - return convert_extract_union_field( - ctx, - rewriter, - op, - aggregate, - decl_index, - operands_info, - ); - } - - let is_scalar = aggregate - .get_type(ctx) - .deref(ctx) - .downcast_ref::() - .is_some(); - - if is_scalar { - rewriter.replace_operation_with_values(ctx, op, vec![aggregate]); - return Ok(()); - } - - let llvm_index = match resolve_aggregate_slots(ctx, operands_info, aggregate)? { - AggregateSlots::Mapped(map) => match map.decl_to_llvm.get(decl_index) { - Some(Some(slot)) => *slot, - Some(None) => { - // ZST field: stripped from the LLVM struct, so there is - // nothing to extract. Materialize undef of its empty type. - let zst_ty = map.field_llvm_types[decl_index]; - let undef_op = llvm::UndefOp::new(ctx, zst_ty); - rewriter.insert_operation(ctx, undef_op.get_operation()); - rewriter.replace_operation(ctx, op, undef_op.get_operation()); - return Ok(()); - } - None => { - return pliron::input_err_noloc!( - "extract_field index {} out of bounds for aggregate with {} fields", - decl_index, - map.decl_to_llvm.len() - ); - } - }, - AggregateSlots::Identity => decl_index as u32, - }; - - let llvm_extract = llvm::ExtractValueOp::new(ctx, aggregate, vec![llvm_index])?; - rewriter.insert_operation(ctx, llvm_extract.get_operation()); - rewriter.replace_operation(ctx, op, llvm_extract.get_operation()); - - Ok(()) -} - -/// Convert `mir.insert_field` to `llvm.insertvalue`. -/// -/// Operands: `[aggregate, new_value]` -/// Returns a new aggregate with the field at `insert_index` replaced. -/// -/// The declaration-order field index is mapped to the LLVM slot via -/// [`resolve_aggregate_slots`] (arrays keep their element index). If -/// inserting a ZST field, we return the original aggregate unchanged. -pub(crate) fn convert_insert_field( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - operands_info: &OperandsInfo, -) -> Result<()> { - let aggregate = op.deref(ctx).get_operand(0); - let new_value = op.deref(ctx).get_operand(1); - - let insert_op = MirInsertFieldOp::new(op); - let decl_index = match insert_op.get_attr_insert_index(ctx) { - Some(attr) => attr.0 as usize, - None => return pliron::input_err_noloc!("Missing insert_index attribute on insert_field"), - }; - - if operands_info - .lookup_most_recent_of_type::(ctx, aggregate) - .is_some() - { - return convert_insert_union_field( - ctx, - rewriter, - op, - aggregate, - new_value, - decl_index, - operands_info, - ); - } - - let llvm_index = match resolve_aggregate_slots(ctx, operands_info, aggregate)? { - AggregateSlots::Mapped(map) => match map.decl_to_llvm.get(decl_index) { - Some(Some(slot)) => *slot, - Some(None) => { - // ZST field: stripped from the LLVM struct, so inserting - // into it is a no-op. Forward the aggregate unchanged. - rewriter.replace_operation_with_values(ctx, op, vec![aggregate]); - return Ok(()); - } - None => { - return pliron::input_err_noloc!( - "insert_field index {} out of bounds for aggregate with {} fields", - decl_index, - map.decl_to_llvm.len() - ); - } - }, - AggregateSlots::Identity => decl_index as u32, - }; - - let llvm_insert = llvm::InsertValueOp::new(ctx, aggregate, new_value, vec![llvm_index]); - rewriter.insert_operation(ctx, llvm_insert.get_operation()); - rewriter.replace_operation(ctx, op, llvm_insert.get_operation()); - - Ok(()) -} - -fn union_type_of_operand( - ctx: &Context, - operands_info: &OperandsInfo, - value: Value, -) -> Result { - operands_info - .lookup_most_recent_of_type::(ctx, value) - .map(|union_ty| union_ty.clone()) - .ok_or_else(|| { - pliron::create_error!( - pliron::location::Location::Unknown, - pliron::result::ErrorKind::VerificationFailed, - pliron::result::StringError( - "Expected MirUnionType conversion history for union value".to_string() - ) - ) - }) -} - -/// Read one typed view of a union's shared bytes. -fn convert_extract_union_field( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - union_value: Value, - field_index: usize, - operands_info: &OperandsInfo, -) -> Result<()> { - let union_ty = union_type_of_operand(ctx, operands_info, union_value)?; - let Some(field_mir_ty) = union_ty.get_field_type(field_index) else { - return pliron::input_err_noloc!( - "union field index {} is out of bounds for `{}`", - field_index, - union_ty.name() - ); - }; - let field_llvm_ty = convert_type(ctx, field_mir_ty).map_err(anyhow_to_pliron)?; - if is_zero_sized_type(ctx, field_llvm_ty) { - let undef = llvm::UndefOp::new(ctx, field_llvm_ty); - rewriter.insert_operation(ctx, undef.get_operation()); - rewriter.replace_operation(ctx, op, undef.get_operation()); - return Ok(()); - } - - let storage_ty = build_union_storage_type(ctx, &union_ty).map_err(anyhow_to_pliron)?; - let ptr = spill_enum_value(ctx, rewriter, union_value, storage_ty, union_ty.abi_align()); - let load = llvm::LoadOp::new(ctx, ptr, field_llvm_ty); - llvm_export::ops::set_op_alignment(ctx, load.get_operation(), union_ty.abi_align() as u32); - rewriter.insert_operation(ctx, load.get_operation()); - rewriter.replace_operation(ctx, op, load.get_operation()); - Ok(()) -} - -/// Write one typed view at byte zero while preserving the rest of the union. -fn convert_insert_union_field( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - union_value: Value, - new_value: Value, - field_index: usize, - operands_info: &OperandsInfo, -) -> Result<()> { - let union_ty = union_type_of_operand(ctx, operands_info, union_value)?; - let Some(field_mir_ty) = union_ty.get_field_type(field_index) else { - return pliron::input_err_noloc!( - "union field index {} is out of bounds for `{}`", - field_index, - union_ty.name() - ); - }; - let field_llvm_ty = convert_type(ctx, field_mir_ty).map_err(anyhow_to_pliron)?; - if is_zero_sized_type(ctx, field_llvm_ty) { - rewriter.replace_operation_with_values(ctx, op, vec![union_value]); - return Ok(()); - } - - let storage_ty = build_union_storage_type(ctx, &union_ty).map_err(anyhow_to_pliron)?; - let ptr = spill_enum_value(ctx, rewriter, union_value, storage_ty, union_ty.abi_align()); - let store = llvm::StoreOp::new(ctx, new_value, ptr); - llvm_export::ops::set_op_alignment(ctx, store.get_operation(), union_ty.abi_align() as u32); - rewriter.insert_operation(ctx, store.get_operation()); - - let load = llvm::LoadOp::new(ctx, ptr, storage_ty); - llvm_export::ops::set_op_alignment(ctx, load.get_operation(), union_ty.abi_align() as u32); - rewriter.insert_operation(ctx, load.get_operation()); - rewriter.replace_operation(ctx, op, load.get_operation()); - Ok(()) -} - -/// Convert `mir.construct_struct` to a chain of `llvm.insertvalue` operations. -/// -/// Builds a struct by: -/// 1. Creating an `undef` value of the lowered struct type -/// 2. Inserting each operand at the LLVM slot its field landed in -/// -/// Operand order matches field order in the struct type (declaration order). -/// The LLVM struct type and the slot of each field both come from -/// [`build_struct_slot_map`], so the insert indices skip `[N x i8]` padding -/// slots exactly the way the type converter laid them out. ZST fields -/// (e.g. PhantomData) have no slot and are skipped. -pub(crate) fn convert_construct_struct( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - _operands_info: &OperandsInfo, -) -> Result<()> { - let (result_ty, operands) = { - let mir_op = op.deref(ctx); - let result_ty = mir_op.get_result(0).get_type(ctx); - let operands: Vec<_> = mir_op.operands().collect(); - (result_ty, operands) - }; - - let layout = { - let ty_ref = result_ty.deref(ctx); - match ty_ref.downcast_ref::() { - Some(s) => StructLayoutInfo::of_struct(s), - None => { - return pliron::input_err_noloc!( - "MirConstructStructOp result type must be MirStructType" - ); - } - } - }; - - if operands.len() != layout.field_types.len() { - return pliron::input_err_noloc!( - "construct_struct has {} operands for a struct with {} fields", - operands.len(), - layout.field_types.len() - ); - } - - let map = build_struct_slot_map(ctx, &layout).map_err(anyhow_to_pliron)?; - - // A divergent rustc layout is constructible by value only when the slot - // map proved that a sequential LLVM packed struct reproduces every byte. - // Overlapping/union-like legacy struct models remain unrepresentable. - if !map.by_value_layout_faithful { - return pliron::input_err_noloc!( - "constructing a struct whose rustc layout cannot be represented by an LLVM \ - struct value is not supported; keep the value behind a pointer and access \ - fields through their byte-accurate address path" - ); - } - - // A packed struct containing AS3 remains target-dependent as a physical - // memory image. The internal-call ABI exception is safe to construct in SSA - // because its return boundary recursively rebuilds every supported AS3 leaf - // into a target-stable generic-pointer carrier before the aggregate crosses - // the function ABI. - if llvm_packed_struct_contains_pointer_in_address_space( - ctx, - map.llvm_struct_ty, - llvm_types::address_space::SHARED, - ) && packed_shared_internal_abi_info(ctx, result_ty) - .map_err(anyhow_to_pliron)? - .is_none() - { - return pliron::input_err_noloc!( - "constructing a packed aggregate containing a shared-memory pointer by value is \ - target-mode dependent and is not yet supported" - ); - } - - let undef_op = llvm::UndefOp::new(ctx, map.llvm_struct_ty); - rewriter.insert_operation(ctx, undef_op.get_operation()); - let mut current_struct = undef_op.get_operation().deref(ctx).get_result(0); - - let mut last_insert: Option> = None; - // Walk in memory order so the insertvalue chain ascends slot indices. - for &decl_idx in &layout.mem_to_decl { - let Some(slot) = map.decl_to_llvm[decl_idx] else { - continue; // ZST field: no slot in the LLVM struct. - }; - - let insert_op = - llvm::InsertValueOp::new(ctx, current_struct, operands[decl_idx], vec![slot]); - rewriter.insert_operation(ctx, insert_op.get_operation()); - current_struct = insert_op.get_operation().deref(ctx).get_result(0); - last_insert = Some(insert_op.get_operation()); - } - - match last_insert { - Some(last_op) => rewriter.replace_operation(ctx, op, last_op), - None => rewriter.replace_operation(ctx, op, undef_op.get_operation()), - } - - Ok(()) -} - -/// Convert `mir.construct_tuple` to a chain of `llvm.insertvalue` operations. -/// -/// Tuples are represented as LLVM structs. Same construction pattern as -/// structs, and like structs the element slots come from -/// [`build_struct_slot_map`] (identity order, no padding; ZST elements are -/// stripped and skipped). -pub(crate) fn convert_construct_tuple( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - _operands_info: &OperandsInfo, -) -> Result<()> { - let (result_ty, operands) = { - let mir_op = op.deref(ctx); - let result_ty = mir_op.get_result(0).get_type(ctx); - let operands: Vec<_> = mir_op.operands().collect(); - (result_ty, operands) - }; - - let layout = { - let ty_ref = result_ty.deref(ctx); - match ty_ref.downcast_ref::() { - Some(t) => StructLayoutInfo::of_tuple(t), - None => { - return pliron::input_err_noloc!( - "MirConstructTupleOp result type must be MirTupleType" - ); - } - } - }; - - if operands.len() != layout.field_types.len() { - return pliron::input_err_noloc!( - "construct_tuple has {} operands for a tuple with {} elements", - operands.len(), - layout.field_types.len() - ); - } - - let map = build_struct_slot_map(ctx, &layout).map_err(anyhow_to_pliron)?; - - let undef_op = llvm::UndefOp::new(ctx, map.llvm_struct_ty); - rewriter.insert_operation(ctx, undef_op.get_operation()); - let mut current_tuple = undef_op.get_operation().deref(ctx).get_result(0); - - let mut last_insert: Option> = None; - for (mir_idx, operand) in operands.iter().enumerate() { - let Some(slot) = map.decl_to_llvm[mir_idx] else { - continue; // ZST element: no slot in the LLVM struct. - }; - - let insert_op = llvm::InsertValueOp::new(ctx, current_tuple, *operand, vec![slot]); - rewriter.insert_operation(ctx, insert_op.get_operation()); - current_tuple = insert_op.get_operation().deref(ctx).get_result(0); - last_insert = Some(insert_op.get_operation()); - } - - match last_insert { - Some(last_op) => rewriter.replace_operation(ctx, op, last_op), - None => rewriter.replace_operation(ctx, op, undef_op.get_operation()), - } - - Ok(()) -} - -/// Convert `mir.construct_slice` to `llvm.undef` + two `llvm.insertvalue`s. -/// -/// `MirSliceType` lowers to the `{ ptr, i64 }` fat-pointer struct, where -/// field 0 is the data pointer and field 1 is the element count by -/// construction (the same layout the entry prologue's `reconstruct_slice` -/// and the Unsize cast path build). `MirSliceType::kind` is intentionally -/// semantic-only: all reference/raw-pointer kinds share this physical layout -/// and no LLVM alias metadata is inferred here. -pub(crate) fn convert_construct_slice( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - _operands_info: &OperandsInfo, -) -> Result<()> { - let (result_ty, data_val, len_val) = { - let mir_op = op.deref(ctx); - ( - mir_op.get_result(0).get_type(ctx), - mir_op.get_operand(0), - mir_op.get_operand(1), - ) - }; - - if !result_ty.deref(ctx).is::() { - return pliron::input_err_noloc!("MirConstructSliceOp result type must be MirSliceType"); - } - - let slice_struct_ty = make_slice_struct(ctx); - - let undef_op = llvm::UndefOp::new(ctx, slice_struct_ty); - rewriter.insert_operation(ctx, undef_op.get_operation()); - let undef_val = undef_op.get_operation().deref(ctx).get_result(0); - - let insert_ptr = llvm::InsertValueOp::new(ctx, undef_val, data_val, vec![0]); - rewriter.insert_operation(ctx, insert_ptr.get_operation()); - let with_ptr = insert_ptr.get_operation().deref(ctx).get_result(0); - - let insert_len = llvm::InsertValueOp::new(ctx, with_ptr, len_val, vec![1]); - rewriter.insert_operation(ctx, insert_len.get_operation()); - - rewriter.replace_operation(ctx, op, insert_len.get_operation()); - - Ok(()) -} - -/// Convert `mir.construct_disjoint_slice` to a chain of `llvm.insertvalue` -/// operations. -/// -/// The same shape as [`convert_construct_slice`], one insert longer per -/// runtime layout word the index space carries. The struct type comes from -/// [`make_disjoint_slice_struct`], the same constructor the type converter -/// uses, so the field order here cannot drift from the lowered layout that -/// [`resolve_aggregate_slots`] indexes identically. -pub(crate) fn convert_construct_disjoint_slice( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - _operands_info: &OperandsInfo, -) -> Result<()> { - let (result_ty, operands) = { - let mir_op = op.deref(ctx); - let operands: Vec<_> = mir_op.operands().collect(); - (mir_op.get_result(0).get_type(ctx), operands) - }; - - let space_tys = { - let ty_ref = result_ty.deref(ctx); - match ty_ref.downcast_ref::() { - Some(slice_ty) => slice_ty.space_types().to_vec(), - None => { - return pliron::input_err_noloc!( - "MirConstructDisjointSliceOp result type must be MirDisjointSliceType" - ); - } - } - }; - - // The verifier already pins the operand count to the result type. Check - // it here too: this pass runs on operations the verifier accepted, and a - // short operand list would otherwise index past the end below. - let expected_operands = 2 + space_tys.len(); - if operands.len() != expected_operands { - return pliron::input_err_noloc!( - "MirConstructDisjointSliceOp expects {expected_operands} operands, got {}", - operands.len() - ); - } - - let slice_struct_ty = crate::convert::types::make_disjoint_slice_struct(ctx, &space_tys) - .map_err(anyhow_to_pliron)?; - - let undef_op = llvm::UndefOp::new(ctx, slice_struct_ty); - rewriter.insert_operation(ctx, undef_op.get_operation()); - let mut current = undef_op.get_operation().deref(ctx).get_result(0); - let mut last_insert = undef_op.get_operation(); - - for (slot, operand) in operands.into_iter().enumerate() { - let insert = llvm::InsertValueOp::new(ctx, current, operand, vec![slot as u32]); - rewriter.insert_operation(ctx, insert.get_operation()); - current = insert.get_operation().deref(ctx).get_result(0); - last_insert = insert.get_operation(); - } - - rewriter.replace_operation(ctx, op, last_insert); - - Ok(()) -} - -/// Convert `mir.construct_array` to a chain of `llvm.insertvalue` operations. -/// -/// Arrays are represented as LLVM arrays. Same construction pattern as structs: -/// 1. Create `undef` of the array type -/// 2. Insert each element at its index -pub(crate) fn convert_construct_array( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - _operands_info: &OperandsInfo, -) -> Result<()> { - let (result_ty, operands) = { - let mir_op = op.deref(ctx); - let result_ty = mir_op.get_result(0).get_type(ctx); - let operands: Vec<_> = mir_op.operands().collect(); - (result_ty, operands) - }; - - let (element_ty, array_size) = { - let ty_ref = result_ty.deref(ctx); - match ty_ref.downcast_ref::() { - Some(a) => (a.element_type(), a.size()), - None => { - return pliron::input_err_noloc!( - "MirConstructArrayOp result type must be MirArrayType" - ); - } - } - }; - - let llvm_element_ty = convert_type(ctx, element_ty).map_err(anyhow_to_pliron)?; - let llvm_array_ty = llvm_export::types::ArrayType::get(ctx, llvm_element_ty, array_size); - - let undef_op = llvm::UndefOp::new(ctx, llvm_array_ty.into()); - rewriter.insert_operation(ctx, undef_op.get_operation()); - let mut current_array = undef_op.get_operation().deref(ctx).get_result(0); - - let mut last_insert: Option> = None; - for (i, operand) in operands.iter().enumerate() { - let insert_op = llvm::InsertValueOp::new(ctx, current_array, *operand, vec![i as u32]); - rewriter.insert_operation(ctx, insert_op.get_operation()); - current_array = insert_op.get_operation().deref(ctx).get_result(0); - last_insert = Some(insert_op.get_operation()); - } - - match last_insert { - Some(last_op) => rewriter.replace_operation(ctx, op, last_op), - None => rewriter.replace_operation(ctx, op, undef_op.get_operation()), - } - - Ok(()) -} -/// Convert `mir.extract_array_element` to LLVM operations. -/// -/// A runtime index normally requires materializing the array in memory because -/// LLVM `extractvalue` accepts only constant indices. When the index is proven -/// to be `urem value, C`, however, it is in `0..C`. For small `C` within the -/// array bounds, emit one constant `extractvalue` per candidate and select the -/// runtime result in SSA. This avoids the temporary alloca that otherwise -/// becomes NVPTX local memory. -/// -/// Unbounded, oversized, or otherwise unsupported indices retain the existing -/// alloca+store+GEP+load fallback. -pub(crate) fn convert_extract_array_element( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - operands_info: &OperandsInfo, -) -> Result<()> { - // One shared cap for the candidate chain, so the mir-transforms - // canonicalization and this lowering fast path cannot drift apart. - use dialect_mir::ops::MAX_SCALARIZED_CANDIDATES; - - fn integer_constant_u64(ctx: &Context, value: Value) -> Option { - let defining_op = value.defining_op()?; - let constant = Operation::get_op::(defining_op, ctx)?; - let attribute = constant.get_value(ctx); - let integer = attribute.downcast_ref::()?; - let integer_value = integer.value(); - // `APInt::to_u64` truncates wider values, so a >64-bit constant could - // be misread as a small in-range divisor. Fail closed on such widths. - (integer_value.bw() <= 64).then(|| integer_value.to_u64()) - } - - fn bounded_urem_candidate_count(ctx: &Context, index: Value, array_size: u64) -> Option { - let defining_op = index.defining_op()?; - Operation::get_op::(defining_op, ctx)?; - - let divisor = defining_op.deref(ctx).get_operand(1); - let candidate_count = integer_constant_u64(ctx, divisor)?; - (candidate_count > 0 - && candidate_count <= array_size - && candidate_count <= MAX_SCALARIZED_CANDIDATES) - .then_some(candidate_count) - } - - fn integer_constant_like( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - reference: Value, - value: u64, - ) -> Result { - let reference_ty = reference.get_type(ctx); - let width = reference_ty - .deref(ctx) - .downcast_ref::() - .ok_or_else(|| { - pliron::input_error_noloc!( - "mir.extract_array_element index must lower to an integer" - ) - })? - .width(); - - let integer_ty = IntegerType::get(ctx, width, Signedness::Signless); - let attribute = pliron::builtin::attributes::IntegerAttr::new( - integer_ty, - APInt::from_u64( - value, - NonZeroUsize::new(width as usize).expect("integer width is nonzero"), - ), - ); - let constant = llvm::ConstantOp::new(ctx, attribute.into()); - rewriter.insert_operation(ctx, constant.get_operation()); - Ok(constant.get_operation().deref(ctx).get_result(0)) - } - - let array_val = op.deref(ctx).get_operand(0); - let index_val = op.deref(ctx).get_operand(1); - - let (element_ty, array_size) = { - match operands_info.lookup_most_recent_of_type::(ctx, array_val) { - Some(r) => (r.element_type(), r.size()), - None => return pliron::input_err_noloc!("Expected MirArrayType"), - } - }; - - if let Some(candidate_count) = bounded_urem_candidate_count(ctx, index_val, array_size) { - let mut candidates = Vec::with_capacity(candidate_count as usize); - for candidate_index in 0..candidate_count { - let extract = llvm::ExtractValueOp::new(ctx, array_val, vec![candidate_index as u32])?; - rewriter.insert_operation(ctx, extract.get_operation()); - candidates.push(extract.get_operation().deref(ctx).get_result(0)); - } - - let mut selected = *candidates - .last() - .expect("candidate count is proven nonzero"); - for candidate_index in (0..candidates.len().saturating_sub(1)).rev() { - let candidate_constant = - integer_constant_like(ctx, rewriter, index_val, candidate_index as u64)?; - let compare = - llvm::ICmpOp::new(ctx, ICmpPredicateAttr::EQ, index_val, candidate_constant); - rewriter.insert_operation(ctx, compare.get_operation()); - let condition = compare.get_operation().deref(ctx).get_result(0); - - let select = llvm::SelectOp::new(ctx, condition, candidates[candidate_index], selected); - rewriter.insert_operation(ctx, select.get_operation()); - selected = select.get_operation().deref(ctx).get_result(0); - } - - rewriter.replace_operation_with_values(ctx, op, vec![selected]); - return Ok(()); - } - - let llvm_element_ty = convert_type(ctx, element_ty).map_err(anyhow_to_pliron)?; - let llvm_array_ty = llvm_export::types::ArrayType::get(ctx, llvm_element_ty, array_size); - let abi_align = mir_type_abi_align(ctx, element_ty); - - let i64_ty = IntegerType::get(ctx, 64, Signedness::Signless); - let one_val = { - let one_apint = APInt::from_i64(1, NonZeroUsize::new(64).unwrap()); - let one_attr = pliron::builtin::attributes::IntegerAttr::new(i64_ty, one_apint); - let const_op = llvm::ConstantOp::new(ctx, one_attr.into()); - rewriter.insert_operation(ctx, const_op.get_operation()); - const_op.get_operation().deref(ctx).get_result(0) - }; - - let alloca_op = llvm::AllocaOp::new(ctx, llvm_array_ty.into(), one_val); - rewriter.insert_operation(ctx, alloca_op.get_operation()); - if let Some(align) = abi_align { - llvm_export::ops::set_op_alignment(ctx, alloca_op.get_operation(), align as u32); - } - let array_ptr = alloca_op.get_operation().deref(ctx).get_result(0); - - let store_op = llvm::StoreOp::new(ctx, array_val, array_ptr); - rewriter.insert_operation(ctx, store_op.get_operation()); - if let Some(align) = abi_align { - llvm_export::ops::set_op_alignment(ctx, store_op.get_operation(), align as u32); - } - - use llvm_export::ops::GepIndex; - let gep_indices = vec![GepIndex::Constant(0), GepIndex::Value(index_val)]; - let gep_op = llvm::GetElementPtrOp::new(ctx, array_ptr, gep_indices, llvm_array_ty.into()); - rewriter.insert_operation(ctx, gep_op.get_operation()); - let element_ptr = gep_op.get_operation().deref(ctx).get_result(0); - - let load_op = llvm::LoadOp::new(ctx, element_ptr, llvm_element_ty); - rewriter.insert_operation(ctx, load_op.get_operation()); - if let Some(align) = abi_align { - llvm_export::ops::set_op_alignment(ctx, load_op.get_operation(), align as u32); - } - rewriter.replace_operation(ctx, op, load_op.get_operation()); - - Ok(()) -} - -/// Copy an enum value into a fresh stack slot and return the pointer. -/// -/// This is how we reach a payload field that has no struct slot of its -/// own (its bytes are shared with a different-typed field of another -/// variant): once the value sits in memory, a byte-precise pointer can -/// read or write any part of it, no struct field needed. -/// -/// The slot is marked with the enum's real (rustc) alignment. The struct -/// type alone can look under-aligned: `{ i8, [7 x i8] }` says "align 1" -/// to LLVM, while Rust may require 8. -/// -/// The alloca lands at the use site, same as -/// [`convert_extract_array_element`]; the standard `opt -O2` run (SROA) -/// removes it again. Hoisting these into the function's entry block is a -/// known follow-up for the unoptimized (`CUDA_OXIDE_NO_OPT=1`) path. -fn spill_enum_value( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - enum_val: Value, - llvm_struct_ty: TypeHandle, - abi_align: u64, -) -> Value { - let i64_ty = IntegerType::get(ctx, 64, Signedness::Signless); - let one_apint = APInt::from_i64(1, NonZeroUsize::new(64).unwrap()); - let one_attr = pliron::builtin::attributes::IntegerAttr::new(i64_ty, one_apint); - let one_const = llvm::ConstantOp::new(ctx, one_attr.into()); - rewriter.insert_operation(ctx, one_const.get_operation()); - let one_val = one_const.get_operation().deref(ctx).get_result(0); - - let alloca_op = llvm::AllocaOp::new(ctx, llvm_struct_ty, one_val); - rewriter.insert_operation(ctx, alloca_op.get_operation()); - if abi_align > 0 { - llvm_export::ops::set_op_alignment(ctx, alloca_op.get_operation(), abi_align as u32); - } - let slot_ptr = alloca_op.get_operation().deref(ctx).get_result(0); - - let store_op = llvm::StoreOp::new(ctx, enum_val, slot_ptr); - rewriter.insert_operation(ctx, store_op.get_operation()); - if abi_align > 0 { - llvm_export::ops::set_op_alignment(ctx, store_op.get_operation(), abi_align as u32); - } - slot_ptr -} - -/// Pointer to `base + offset` bytes (`getelementptr i8, ptr base, offset`). -/// -/// Used whenever rustc's physical byte offset cannot be represented faithfully -/// by a typed LLVM aggregate GEP, as with overlapping enum payloads or packed -/// struct fields. -fn byte_offset_gep( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - base: Value, - offset: u64, -) -> Value { - use llvm_export::ops::GepIndex; - let i8_ty: TypeHandle = IntegerType::get(ctx, 8, Signedness::Signless).into(); - let offset = emit_integer_constant(ctx, rewriter, 64, u128::from(offset)); - let gep_op = llvm::GetElementPtrOp::new(ctx, base, vec![GepIndex::Value(offset)], i8_ty); - rewriter.insert_operation(ctx, gep_op.get_operation()); - gep_op.get_operation().deref(ctx).get_result(0) -} - -/// The physical carrier write required to select one source variant. -/// `None` is a real semantic result for the untagged niche variant and for -/// rustc's single inhabited variant; it must never be replaced by a guessed -/// memory write. -fn enum_carrier_bits_for_variant( - enum_ty: &MirEnumType, - variant: usize, -) -> std::result::Result, String> { - if variant >= enum_ty.variant_count() || enum_ty.variant_is_inhabited(variant) != Some(true) { - return Err(format!( - "cannot select uninhabited or missing variant {} of '{}'", - variant, - enum_ty.name() - )); - } - - match enum_ty.layout_kind { - EnumLayoutKind::Direct => enum_ty - .variant_discriminants - .get(variant) - .copied() - .map(|bits| Some(u128::from(bits))) - .ok_or_else(|| format!("variant {} has no declared discriminant", variant)), - EnumLayoutKind::Niche => { - // This check deliberately comes before the encoded range check: - // rustc permits the untagged variant index to lie inside that - // range. Its range position is a dead niche value; selecting the - // actual untagged variant remains a no-op. - if variant == enum_ty.untagged_variant as usize { - return Ok(None); - } - if !(enum_ty.niche_variant_start as usize..=enum_ty.niche_variant_end as usize) - .contains(&variant) - { - return Err(format!( - "inhabited variant {} is not representable by niche layout of '{}'", - variant, - enum_ty.name() - )); - } - let offset = (variant as u128) - u128::from(enum_ty.niche_variant_start); - let mut bits = enum_ty.niche_start().wrapping_add(offset); - if enum_ty.carrier_width < 128 { - bits &= (1u128 << enum_ty.carrier_width) - 1; - } - Ok(Some(bits)) - } - EnumLayoutKind::Single if variant == enum_ty.single_variant as usize => Ok(None), - EnumLayoutKind::Single => Err(format!( - "variant {} is not the single inhabited variant of '{}'", - variant, - enum_ty.name() - )), - EnumLayoutKind::Empty => Err(format!("enum '{}' is uninhabited", enum_ty.name())), - EnumLayoutKind::Unknown => Err(format!( - "enum '{}' has unknown physical layout", - enum_ty.name() - )), - } -} - -fn emit_integer_constant( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - width: u32, - bits: u128, -) -> Value { - let ty = IntegerType::get(ctx, width, Signedness::Signless); - let attr = pliron::builtin::attributes::IntegerAttr::new( - ty, - APInt::from_u128(bits, NonZeroUsize::new(width as usize).unwrap()), - ); - let op = llvm::ConstantOp::new(ctx, attr.into()); - rewriter.insert_operation(ctx, op.get_operation()); - op.get_operation().deref(ctx).get_result(0) -} - -fn emit_carrier_constant( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - enum_ty: &MirEnumType, - carrier_ty: TypeHandle, - bits: u128, -) -> Result { - let integer = emit_integer_constant(ctx, rewriter, enum_ty.carrier_width, bits); - match enum_ty.carrier_kind { - EnumCarrierKind::Integer => Ok(integer), - EnumCarrierKind::Pointer => { - let cast = llvm::IntToPtrOp::new(ctx, integer, carrier_ty); - rewriter.insert_operation(ctx, cast.get_operation()); - Ok(cast.get_operation().deref(ctx).get_result(0)) - } - _ => pliron::input_err_noloc!("enum carrier constant requested without a carrier"), - } -} - -/// Convert a value to its byte-faithful storage twin: every `i1` leaf is -/// zero-extended to its canonical `i8` memory byte, recursively through -/// structs and arrays. Values without `i1` storage pass through unchanged. -/// -/// This is the value-level half of [`llvm_byte_faithful_twin`]: the enum -/// slot map claims twin-typed storage for bool-bearing payloads, and this -/// produces the twin-typed value the store into that storage needs. -fn canonicalize_bool_value_bytes( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - value: Value, -) -> Result { - let ty = value.get_type(ctx); - if !llvm_type_contains_i1(ctx, ty) { - return Ok(value); - } - let is_scalar_i1 = ty - .deref(ctx) - .downcast_ref::() - .is_some_and(|integer| integer.width() == 1); - if is_scalar_i1 { - let byte_ty: TypeHandle = IntegerType::get(ctx, 8, Signedness::Signless).into(); - let zext = llvm::ZExtOp::new_with_nneg(ctx, value, byte_ty, false); - rewriter.insert_operation(ctx, zext.get_operation()); - return Ok(zext.get_operation().deref(ctx).get_result(0)); - } - let Some(twin) = llvm_byte_faithful_twin(ctx, ty) else { - return pliron::input_err_noloc!( - "enum construction: bool storage in this value's shape cannot be canonicalized" - ); - }; - let element_count = { - let ty_ref = ty.deref(ctx); - if let Some(struct_ty) = ty_ref.downcast_ref::() { - struct_ty.fields().count() as u64 - } else if let Some(array_ty) = ty_ref.downcast_ref::() { - array_ty.size() - } else { - return pliron::input_err_noloc!( - "enum construction: unexpected container for bool storage canonicalization" - ); - } - }; - let undef_op = llvm::UndefOp::new(ctx, twin); - rewriter.insert_operation(ctx, undef_op.get_operation()); - let mut current = undef_op.get_operation().deref(ctx).get_result(0); - for index in 0..element_count { - let extract_op = llvm::ExtractValueOp::new(ctx, value, vec![index as u32])?; - rewriter.insert_operation(ctx, extract_op.get_operation()); - let element = extract_op.get_operation().deref(ctx).get_result(0); - let converted = canonicalize_bool_value_bytes(ctx, rewriter, element)?; - let insert_op = llvm::InsertValueOp::new(ctx, current, converted, vec![index as u32]); - rewriter.insert_operation(ctx, insert_op.get_operation()); - current = insert_op.get_operation().deref(ctx).get_result(0); - } - Ok(current) -} - -/// Adapt a semantic enum payload value to or from its physical storage type. -/// -/// Shared pointer leaves are converted through CUDA generic space recursively -/// through struct/tuple payloads. Bool leaves are canonicalized separately on -/// construction and narrowed recursively on extraction. -fn coerce_enum_payload_storage( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - value: Value, - target_ty: TypeHandle, -) -> Result { - coerce_enum_payload_value(ctx, rewriter, value, target_ty) -} - -/// Convert `mir.construct_enum` (e.g. `E::A(x)`) to LLVM operations. -/// -/// Builds the enum value slot by slot, taking every index from -/// [`build_enum_slot_map`] (indexes are never computed by hand here): -/// -/// 1. Put the variant's declared discriminant VALUE into the tag slot. -/// 2. `insertvalue` each payload field that owns a struct slot. -/// 3. If some field has no slot (its bytes are shared with a -/// different-typed field of another variant), finish through memory: -/// copy the value to a stack slot, store that field at its byte -/// position, and load the completed enum back. -pub(crate) fn convert_construct_enum( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - _operands_info: &OperandsInfo, -) -> Result<()> { - let (result_ty, operands, variant_index) = { - let mir_op = op.deref(ctx); - let result_ty = mir_op.get_result(0).get_type(ctx); - let operands: Vec<_> = mir_op.operands().collect(); - - let enum_op = MirConstructEnumOp::new(op); - let variant_index = enum_op - .get_attr_construct_enum_variant_index(ctx) - .map(|attr| attr.0 as usize) - .unwrap_or(0); - - (result_ty, operands, variant_index) - }; - - let enum_ty: MirEnumType = { - let ty_ref = result_ty.deref(ctx); - match ty_ref.downcast_ref::() { - Some(e) => e.clone(), - None => { - return pliron::input_err_noloc!( - "MirConstructEnumOp result type must be MirEnumType" - ); - } - } - }; - - // Build the value as the SAME struct type the type converter - // produces everywhere else (block args, loads, allocas, ...). Taking - // both the type and the indices from one slot map is what keeps them - // in agreement. Filler slots are simply never written. - let slot_map = build_enum_slot_map(ctx, result_ty).map_err(anyhow_to_pliron)?; - let llvm_struct_ty = slot_map.llvm_struct_ty; - - let undef_op = llvm::UndefOp::new(ctx, llvm_struct_ty); - rewriter.insert_operation(ctx, undef_op.get_operation()); - let mut current_struct = undef_op.get_operation().deref(ctx).get_result(0); - let mut last_op = undef_op.get_operation(); - - if let Some(bits) = enum_carrier_bits_for_variant(&enum_ty, variant_index) - .map_err(|error| pliron::input_error_noloc!("MirConstructEnumOp: {error}"))? - { - let carrier_slot = slot_map.carrier_slot.ok_or_else(|| { - pliron::input_error_noloc!("MirConstructEnumOp requires a physical carrier slot") - })?; - let carrier_ty = slot_map.carrier_llvm_ty.ok_or_else(|| { - pliron::input_error_noloc!("MirConstructEnumOp requires a physical carrier type") - })?; - let carrier = emit_carrier_constant(ctx, rewriter, &enum_ty, carrier_ty, bits)?; - let insert = llvm::InsertValueOp::new(ctx, current_struct, carrier, vec![carrier_slot]); - rewriter.insert_operation(ctx, insert.get_operation()); - current_struct = insert.get_operation().deref(ctx).get_result(0); - last_op = insert.get_operation(); - } - - let field_base: usize = enum_ty - .variant_field_counts - .iter() - .take(variant_index) - .map(|&c| c as usize) - .sum(); - - // Insert every payload field that owns a struct slot; remember the - // slotless ones for the memory pass below. - let mut deferred: Vec<(usize, Value)> = Vec::new(); - for (i, operand) in operands.into_iter().enumerate() { - let flat = field_base + i; - let Some(slot) = slot_map.field_slots.get(flat) else { - return pliron::input_err_noloc!( - "MirConstructEnumOp field {} of variant {} is out of range for the enum's {} fields", - i, - variant_index, - slot_map.field_slots.len() - ); - }; - match slot { - Some(slot) => { - let storage_ty = { - let storage_type = llvm_struct_ty.deref(ctx); - let storage_struct = storage_type - .downcast_ref::() - .ok_or_else(|| { - pliron::input_error_noloc!( - "MirConstructEnumOp physical storage is not an LLVM struct" - ) - })?; - let storage_slot = *slot as usize; - if storage_slot >= storage_struct.num_fields() { - return pliron::input_err_noloc!( - "MirConstructEnumOp physical field slot {} is out of range", - slot - ); - } - storage_struct.field_type(storage_slot) - }; - let stored_operand = - coerce_enum_payload_storage(ctx, rewriter, operand, storage_ty)?; - let insert_op = - llvm::InsertValueOp::new(ctx, current_struct, stored_operand, vec![*slot]); - rewriter.insert_operation(ctx, insert_op.get_operation()); - current_struct = insert_op.get_operation().deref(ctx).get_result(0); - last_op = insert_op.get_operation(); - } - None => { - // Zero-sized fields own no bytes; nothing to write. - if is_zero_sized_type(ctx, slot_map.field_llvm_types[flat]) { - continue; - } - deferred.push((flat, operand)); - } - } - } - - if deferred.is_empty() { - rewriter.replace_operation(ctx, op, last_op); - return Ok(()); - } - - // Slotless fields: copy the half-built value to the stack, write - // each remaining payload at its byte position, and load the finished - // enum back as the result. - let abi_align = enum_ty.abi_align(); - let slot_ptr = spill_enum_value(ctx, rewriter, current_struct, llvm_struct_ty, abi_align); - for (flat, operand) in deferred { - let field_ptr = byte_offset_gep(ctx, rewriter, slot_ptr, slot_map.field_offsets[flat]); - // `bool` is an LLVM i1 as a value but occupies one full byte in - // Rust memory. Enum storage claims the byte-faithful twin of every - // bool-bearing payload (scalar i8 byte, or an aggregate with each - // i1 leaf widened to i8), so canonicalize the stored value to that - // twin: every physical bool byte becomes an unambiguous 0 or 1. - let semantic_ty = slot_map.field_llvm_types[flat]; - let storage_ty = enum_payload_storage_type(ctx, semantic_ty).map_err(anyhow_to_pliron)?; - let canonical_operand = canonicalize_bool_value_bytes(ctx, rewriter, operand)?; - let stored_operand = - coerce_enum_payload_storage(ctx, rewriter, canonical_operand, storage_ty)?; - let store_op = llvm::StoreOp::new(ctx, stored_operand, field_ptr); - rewriter.insert_operation(ctx, store_op.get_operation()); - } - let load_op = llvm::LoadOp::new(ctx, slot_ptr, llvm_struct_ty); - rewriter.insert_operation(ctx, load_op.get_operation()); - if abi_align > 0 { - llvm_export::ops::set_op_alignment(ctx, load_op.get_operation(), abi_align as u32); - } - rewriter.replace_operation(ctx, op, load_op.get_operation()); - - Ok(()) -} - -/// Get the slot map for an enum operand. -/// -/// By the time an op is converted, its operand's type has already been -/// rewritten to the LLVM struct, so we look up the ORIGINAL `MirEnumType` -/// the framework recorded for it and rebuild the map from that. Also -/// returns the enum's rustc alignment, which spill slots need. -fn enum_slot_map_of_operand( - ctx: &mut Context, - operands_info: &OperandsInfo, - enum_val: Value, -) -> Result<(EnumSlotMap, u64)> { - // Clone the type data out so the `Ref` borrow of `ctx` ends before - // re-interning (types are hash-consed: registering an equal instance - // returns the existing pointer). - let enum_ty: MirEnumType = { - match operands_info.lookup_most_recent_of_type::(ctx, enum_val) { - Some(r) => r.clone(), - None => { - return pliron::input_err_noloc!("Expected MirEnumType for enum value access"); - } - } - }; - let abi_align = enum_ty.abi_align(); - let mir_ty: TypeHandle = pliron::r#type::Type::instantiate(enum_ty, ctx).into(); - let map = build_enum_slot_map(ctx, mir_ty).map_err(anyhow_to_pliron)?; - Ok((map, abi_align)) -} - -/// Convert `mir.set_discriminant` to the one physical carrier write rustc's -/// layout requires. Direct layouts write the declared discriminant; Niche -/// layouts write wrapping `niche_start + range_offset`. Selecting the -/// untagged Niche variant or an inhabited Single variant is a no-op. -/// -/// The enum layout comes from the op's own `set_discriminant_enum_ty` -/// attribute, stamped at build time. Operand type history is not usable -/// here: a kind-only `mir.cast` lowers to a plain value forwarding, history -/// does not follow that edge, and a stale hit would write the tag at the -/// wrong offset. -pub(crate) fn convert_set_discriminant( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - _operands_info: &OperandsInfo, -) -> Result<()> { - let enum_ptr = op.deref(ctx).get_operand(0); - let target = MirSetDiscriminantOp::new(op) - .get_attr_set_discriminant_variant_index(ctx) - .map(|attr| attr.0 as usize) - .ok_or_else(|| { - pliron::input_error_noloc!( - "MirSetDiscriminantOp missing set_discriminant_variant_index" - ) - })?; - - let enum_ty: MirEnumType = { - let stamped_ty = MirSetDiscriminantOp::new(op) - .get_attr_set_discriminant_enum_ty(ctx) - .map(|attr| attr.get_type(ctx)) - .ok_or_else(|| { - pliron::input_error_noloc!( - "mir.set_discriminant missing enum type attribute; \ - discriminant write has no fact to derive from" - ) - })?; - match stamped_ty.deref(ctx).downcast_ref::() { - Some(et) => et.clone(), - None => { - return pliron::input_err_noloc!( - "mir.set_discriminant enum type attribute must be an enum type" - ); - } - } - }; - - let Some(bits) = enum_carrier_bits_for_variant(&enum_ty, target) - .map_err(|error| pliron::input_error_noloc!("MirSetDiscriminantOp: {error}"))? - else { - rewriter.erase_operation(ctx, op); - return Ok(()); - }; - - let tag_offset = enum_ty.tag_offset(); - let mir_ty: TypeHandle = pliron::r#type::Type::instantiate(enum_ty.clone(), ctx).into(); - let slot_map = build_enum_slot_map(ctx, mir_ty).map_err(anyhow_to_pliron)?; - let carrier_ty = slot_map.carrier_llvm_ty.ok_or_else(|| { - pliron::input_error_noloc!("MirSetDiscriminantOp physical write has no carrier type") - })?; - let carrier = emit_carrier_constant(ctx, rewriter, &enum_ty, carrier_ty, bits)?; - let carrier_ptr = byte_offset_gep(ctx, rewriter, enum_ptr, tag_offset); - let store_op = llvm::StoreOp::new(ctx, carrier, carrier_ptr); - rewriter.insert_operation(ctx, store_op.get_operation()); - - rewriter.erase_operation(ctx, op); - Ok(()) -} - -/// Convert `mir.get_discriminant` (reading which variant is alive) to -/// `llvm.extractvalue`. -/// -/// Direct layouts read the tag from the slot map's carrier slot. Niche -/// layouts decode rustc's wrapping carrier range; Single layouts materialize -/// their one logical discriminant as a constant. No slot number is assumed. -pub(crate) fn convert_get_discriminant( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - operands_info: &OperandsInfo, -) -> Result<()> { - let enum_val = match op.deref(ctx).operands().next() { - Some(v) => v, - None => return pliron::input_err_noloc!("MirGetDiscriminantOp requires an operand"), - }; - - let enum_ty: MirEnumType = operands_info - .lookup_most_recent_of_type::(ctx, enum_val) - .map(|ty| ty.clone()) - .ok_or_else(|| pliron::input_error_noloc!("Expected MirEnumType for discriminant read"))?; - let mir_ty: TypeHandle = pliron::r#type::Type::instantiate(enum_ty.clone(), ctx).into(); - let slot_map = build_enum_slot_map(ctx, mir_ty).map_err(anyhow_to_pliron)?; - let logical_ty = convert_type(ctx, enum_ty.discriminant_ty).map_err(anyhow_to_pliron)?; - let logical_width = logical_ty - .deref(ctx) - .downcast_ref::() - .map(IntegerType::width) - .ok_or_else(|| { - pliron::input_error_noloc!("MirGetDiscriminantOp logical result must be integer") - })?; - - let result = match enum_ty.layout_kind { - EnumLayoutKind::Direct => { - let slot = slot_map.carrier_slot.ok_or_else(|| { - pliron::input_error_noloc!("Direct enum has no physical carrier slot") - })?; - let extract = llvm::ExtractValueOp::new(ctx, enum_val, vec![slot])?; - rewriter.insert_operation(ctx, extract.get_operation()); - extract.get_operation().deref(ctx).get_result(0) - } - EnumLayoutKind::Single => { - if enum_ty.variant_is_inhabited(enum_ty.single_variant as usize) != Some(true) { - return pliron::input_err_noloc!( - "Cannot read discriminant of an uninhabited single-variant enum" - ); - } - let value = *enum_ty - .variant_discriminants - .get(enum_ty.single_variant as usize) - .ok_or_else(|| { - pliron::input_error_noloc!("Single enum has no declared discriminant") - })?; - emit_integer_constant(ctx, rewriter, logical_width, u128::from(value)) - } - EnumLayoutKind::Niche => { - let slot = slot_map.carrier_slot.ok_or_else(|| { - pliron::input_error_noloc!("Niche enum has no physical carrier slot") - })?; - let extract = llvm::ExtractValueOp::new(ctx, enum_val, vec![slot])?; - rewriter.insert_operation(ctx, extract.get_operation()); - let carrier = extract.get_operation().deref(ctx).get_result(0); - let carrier_int_ty: TypeHandle = - IntegerType::get(ctx, enum_ty.carrier_width, Signedness::Signless).into(); - let carrier_int = if enum_ty.carrier_kind == EnumCarrierKind::Pointer { - let cast = llvm::PtrToIntOp::new(ctx, carrier, carrier_int_ty); - rewriter.insert_operation(ctx, cast.get_operation()); - cast.get_operation().deref(ctx).get_result(0) - } else { - carrier - }; - - let niche_start = - emit_integer_constant(ctx, rewriter, enum_ty.carrier_width, enum_ty.niche_start()); - let relative = llvm::SubOp::new_with_overflow_flag( - ctx, - carrier_int, - niche_start, - IntegerOverflowFlagsAttr::default(), - ) - .get_operation(); - rewriter.insert_operation(ctx, relative); - let relative_val = relative.deref(ctx).get_result(0); - let max = emit_integer_constant( - ctx, - rewriter, - enum_ty.carrier_width, - u128::from(enum_ty.niche_variant_end - enum_ty.niche_variant_start), - ); - let in_range = - llvm::ICmpOp::new(ctx, ICmpPredicateAttr::ULE, relative_val, max).get_operation(); - rewriter.insert_operation(ctx, in_range); - // The range test is carrier-width wrapping arithmetic, exactly - // like rustc. Variant indices themselves live at the logical - // discriminant width: the start index may be larger than the - // carrier can represent (e.g. variants 298..=299 in an i8 niche). - let logical_relative = match enum_ty.carrier_width.cmp(&logical_width) { - std::cmp::Ordering::Equal => relative_val, - std::cmp::Ordering::Greater => { - let cast = llvm::TruncOp::new(ctx, relative_val, logical_ty).get_operation(); - rewriter.insert_operation(ctx, cast); - cast.deref(ctx).get_result(0) - } - std::cmp::Ordering::Less => { - // LLVM's zext op requires its explicit `nneg` flag even - // when the flag is false. Niche decoding is ordinary - // unsigned extension, so it must never claim nneg. - let cast = llvm::ZExtOp::new_with_nneg(ctx, relative_val, logical_ty, false) - .get_operation(); - rewriter.insert_operation(ctx, cast); - cast.deref(ctx).get_result(0) - } - }; - let niche_base = emit_integer_constant( - ctx, - rewriter, - logical_width, - u128::from(enum_ty.niche_variant_start), - ); - let niche_variant = llvm::AddOp::new_with_overflow_flag( - ctx, - logical_relative, - niche_base, - IntegerOverflowFlagsAttr::default(), - ) - .get_operation(); - rewriter.insert_operation(ctx, niche_variant); - let untagged = emit_integer_constant( - ctx, - rewriter, - logical_width, - u128::from(enum_ty.untagged_variant), - ); - let in_range_value = in_range.deref(ctx).get_result(0); - let niche_variant_value = niche_variant.deref(ctx).get_result(0); - let select = llvm::SelectOp::new(ctx, in_range_value, niche_variant_value, untagged) - .get_operation(); - rewriter.insert_operation(ctx, select); - select.deref(ctx).get_result(0) - } - EnumLayoutKind::Empty => { - return pliron::input_err_noloc!("Cannot read discriminant of uninhabited enum"); - } - _ => { - return pliron::input_err_noloc!( - "Cannot read discriminant of enum with unknown physical layout" - ); - } - }; - - rewriter.replace_operation_with_values(ctx, op, vec![result]); - - Ok(()) -} - -/// Convert `mir.enum_payload` (reading a variant's field, e.g. the `x` -/// in `E::A(x) => x`) to a payload-field read. -/// -/// Three cases, decided by the [`EnumSlotMap`]: -/// -/// - The field owns a struct slot: a plain `llvm.extractvalue`. -/// - The field has no slot (its bytes are shared with a different-typed -/// field of another variant): go through memory. Copy the enum to a -/// stack slot, point at the field's byte position, and load it with -/// its own type. Same trick as [`convert_extract_array_element`], and -/// it avoids LLVM `bitcast` entirely. -/// - The field is zero-sized: there is nothing to read; produce `undef`. -pub(crate) fn convert_enum_payload( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - operands_info: &OperandsInfo, -) -> Result<()> { - let enum_val = match op.deref(ctx).operands().next() { - Some(v) => v, - None => return pliron::input_err_noloc!("MirEnumPayloadOp requires an operand"), - }; - - let payload_op = MirEnumPayloadOp::new(op); - let variant_index = payload_op - .get_attr_payload_variant_index(ctx) - .map(|attr| attr.0 as usize) - .unwrap_or(0); - let field_index = payload_op - .get_attr_payload_field_index(ctx) - .map(|attr| attr.0 as usize) - .unwrap_or(0); - - let variant_field_counts = { - match operands_info.lookup_most_recent_of_type::(ctx, enum_val) { - Some(r) => r.variant_field_counts.clone(), - None => { - return pliron::input_err_noloc!( - "Expected MirEnumType for enum payload extraction" - ); - } - } - }; - let (slot_map, abi_align) = enum_slot_map_of_operand(ctx, operands_info, enum_val)?; - - let field_base: usize = variant_field_counts - .iter() - .take(variant_index) - .map(|&c| c as usize) - .sum(); - let flat = field_base + field_index; - let Some(slot) = slot_map.field_slots.get(flat).copied() else { - return pliron::input_err_noloc!( - "MirEnumPayloadOp field {} of variant {} is out of range for the enum's {} fields", - field_index, - variant_index, - slot_map.field_slots.len() - ); - }; - - match slot { - Some(slot) => { - let extract_op = llvm::ExtractValueOp::new(ctx, enum_val, vec![slot])?; - rewriter.insert_operation(ctx, extract_op.get_operation()); - let extracted = extract_op.get_operation().deref(ctx).get_result(0); - let semantic_value = coerce_enum_payload_storage( - ctx, - rewriter, - extracted, - slot_map.field_llvm_types[flat], - )?; - rewriter.replace_operation_with_values(ctx, op, vec![semantic_value]); - } - None if is_zero_sized_type(ctx, slot_map.field_llvm_types[flat]) => { - let undef_op = llvm::UndefOp::new(ctx, slot_map.field_llvm_types[flat]); - rewriter.insert_operation(ctx, undef_op.get_operation()); - rewriter.replace_operation(ctx, op, undef_op.get_operation()); - } - None => { - let slot_ptr = - spill_enum_value(ctx, rewriter, enum_val, slot_map.llvm_struct_ty, abi_align); - let field_ptr = byte_offset_gep(ctx, rewriter, slot_ptr, slot_map.field_offsets[flat]); - let semantic_ty = slot_map.field_llvm_types[flat]; - let storage_ty = - enum_payload_storage_type(ctx, semantic_ty).map_err(anyhow_to_pliron)?; - let load_op = llvm::LoadOp::new(ctx, field_ptr, storage_ty); - rewriter.insert_operation(ctx, load_op.get_operation()); - let stored = load_op.get_operation().deref(ctx).get_result(0); - let semantic = coerce_enum_payload_storage(ctx, rewriter, stored, semantic_ty)?; - rewriter.replace_operation_with_values(ctx, op, vec![semantic]); - } - } - - Ok(()) -} - -// ============================================================================ -// MirFieldAddrOp Conversion -// ============================================================================ - -/// Convert `mir.field_addr` to `llvm.getelementptr`. -/// -/// Computes the address of a struct field using GEP. This is needed when -/// Rust code takes `&mut self.field` — we need the ADDRESS of the field, -/// not a COPY of its value. -/// -/// The GEP field index and the struct type it indexes into both come from -/// [`build_struct_slot_map`], so the index accounts for reordering, -/// `[N x i8]` padding slots and stripped ZSTs (ZST-ness is decided on the -/// converted LLVM field type, like the value-level sites). -/// -/// Taking the address of a ZST field emits a distinct zero-offset `i8` GEP -/// off the base pointer (a ZST field lives at byte 0 of the struct), the -/// same idiom the union branch uses. It must NOT forward the base SSA value -/// itself: a 1:1 `replace_operation_with_values` that aliases the op's -/// result to an already-existing value records the result's pointee type -/// (the ZST field's own zero-field type) onto that value's conversion type -/// history, so a sibling `field_addr` on the same base would later resolve -/// the base's pointee to the ZST type instead of the struct and fail with -/// "field_addr index N out of bounds for struct with 0 fields". The -/// distinct GEP result leaves the base pointer's recorded pointee intact -/// and carries the field's pointee type on its own result, which is what -/// nested projections through the field address look up. -pub(crate) fn convert_field_addr( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - operands_info: &OperandsInfo, -) -> Result<()> { - let ptr_operand = op.deref(ctx).get_operand(0); - - let field_addr_op = MirFieldAddrOp::new(op); - let field_index = match field_addr_op.get_attr_field_index(ctx) { - Some(attr) => attr.0 as usize, - None => return pliron::input_err_noloc!("MirFieldAddrOp missing field_index attribute"), - }; - - let mir_ptr_pointee = - match operands_info.lookup_most_recent_of_type::(ctx, ptr_operand) { - Some(r) => r.pointee, - None => { - return pliron::input_err_noloc!("MirFieldAddrOp operand must be pointer type"); - } - }; - - let union_field_count = mir_ptr_pointee - .deref(ctx) - .downcast_ref::() - .map(MirUnionType::field_count); - if let Some(field_count) = union_field_count { - if field_index >= field_count { - return pliron::input_err_noloc!( - "field_addr index {} out of bounds for union with {} fields", - field_index, - field_count - ); - } - // Every union field begins at byte zero. Emit an explicit zero-offset - // GEP instead of forwarding the base SSA value directly: the distinct - // result keeps dialect conversion's pointer-type history unambiguous - // for repeated field accesses and for `union.struct_field.inner`. - use llvm_export::ops::GepIndex; - let i8_ty: TypeHandle = IntegerType::get(ctx, 8, Signedness::Signless).into(); - let gep = llvm::GetElementPtrOp::new(ctx, ptr_operand, vec![GepIndex::Constant(0)], i8_ty); - rewriter.insert_operation(ctx, gep.get_operation()); - rewriter.replace_operation(ctx, op, gep.get_operation()); - return Ok(()); - } - - // An enum payload field, addressed by its position in the flattened - // `all_field_types`. Variants share bytes, so the slot map resolves each - // field one of two ways: its own LLVM slot, or, when its bytes are already - // held by a differently typed field of another variant, a byte offset into - // the enum. Both give the address of the same storage, which is what makes - // a write through the returned pointer land in the enum rather than a copy. - let enum_name = mir_ptr_pointee - .deref(ctx) - .downcast_ref::() - .map(|enum_ty| enum_ty.name().to_string()); - if let Some(enum_name) = enum_name { - let map = build_enum_slot_map(ctx, mir_ptr_pointee).map_err(anyhow_to_pliron)?; - let Some(slot_entry) = map.field_slots.get(field_index).copied() else { - return pliron::input_err_noloc!( - "field_addr index {} out of bounds for enum with {} payload fields", - field_index, - map.field_slots.len() - ); - }; - - // Enum payload bytes hold one CANONICAL storage type that can differ - // from the field's semantic type: a bool is physically an i8 byte and - // a shared-memory pointer is stored as a generic pointer, with the - // value paths (construct/extract) coercing exactly at that boundary. - // The address computed here ESCAPES this site: the loads and stores - // made through it happen at arbitrary other sites and are typed with - // the SEMANTIC type, so no storage coercion can be attached at - // address-formation time. Handing the address out anyway would let an - // i1 store leave the byte's upper seven bits undefined for the i8 and - // niche-tag readers, or write a shared-pointer representation into - // bytes every other reader interprets (and, on modern NVVM, sizes) as - // a generic pointer. Fail closed instead, the same way the slot map - // already rejects bool bytes hidden behind unions. - let semantic_ty = map.field_llvm_types[field_index]; - let storage_ty = enum_payload_storage_type(ctx, semantic_ty).map_err(anyhow_to_pliron)?; - if storage_ty != semantic_ty { - return pliron::input_err_noloc!( - "field_addr: cannot hand out the in-place address of payload field {} of enum `{}`: its bytes use canonical storage type {} while its semantic type is {}, and loads or stores made through an escaped payload address are typed with the semantic type, which the canonical bytes cannot honor; shared reads of such payloads are compiled through a value copy automatically, and a write that stays inside its function is compiled by rebuilding the enum around the new payload; a borrow that escapes into a call keeps no such rewrite and is refused here", - field_index, - enum_name, - storage_ty.deref(ctx).disp(ctx), - semantic_ty.deref(ctx).disp(ctx) - ); - } - - use llvm_export::ops::GepIndex; - let gep = match slot_entry { - Some(slot) => llvm::GetElementPtrOp::new( - ctx, - ptr_operand, - vec![GepIndex::Constant(0), GepIndex::Constant(slot)], - map.llvm_struct_ty, - ), - None => { - // No slot of its own: address the bytes directly, off the - // ORIGINAL enum pointer, so a write through the result lands - // in the enum. A zero-sized field lands at its offset like - // any other and simply spans nothing. The offset is always - // present: the slot map builds `field_offsets` and - // `field_slots` from one walk over the same fields, and the - // bounds check above already validated the index. - let offset = map.field_offsets[field_index]; - let offset = u32::try_from(offset).map_err(|_| { - pliron::input_error_noloc!( - "field_addr: payload byte offset {} of enum `{}` exceeds u32", - offset, - enum_name - ) - })?; - let i8_ty: TypeHandle = IntegerType::get(ctx, 8, Signedness::Signless).into(); - llvm::GetElementPtrOp::new( - ctx, - ptr_operand, - vec![GepIndex::Constant(offset)], - i8_ty, - ) - } - }; - rewriter.insert_operation(ctx, gep.get_operation()); - rewriter.replace_operation(ctx, op, gep.get_operation()); - return Ok(()); - } - - // Carried alongside the layout so the field address can record what it - // proves about its own alignment; see `stamp_field_address_alignment`. - let aggregate_abi_align; - let layout = { - let pointee_ref = mir_ptr_pointee.deref(ctx); - if let Some(struct_ty) = pointee_ref.downcast_ref::() { - aggregate_abi_align = struct_ty.abi_align; - StructLayoutInfo::of_struct(struct_ty) - } else if let Some(tuple_ty) = pointee_ref.downcast_ref::() { - aggregate_abi_align = tuple_ty.abi_align(); - StructLayoutInfo::of_tuple(tuple_ty) - } else { - return pliron::input_err_noloc!( - "MirFieldAddrOp pointer must point to a struct, tuple, union or enum type, got {}", - mir_ptr_pointee.deref(ctx).disp(ctx) - ); - } - }; - - let map = build_struct_slot_map(ctx, &layout).map_err(anyhow_to_pliron)?; - - let slot = match map.decl_to_llvm.get(field_index) { - Some(Some(slot)) => *slot, - Some(None) => { - // ZST field: it has no storage; the struct address stands in for - // the field address. Emit an explicit zero-offset GEP instead of - // forwarding the base SSA value directly (mirrors the union branch - // above): the distinct result keeps dialect conversion's pointer-type - // history unambiguous for repeated field accesses off the same base - // pointer. Forwarding the base value here type-puns it to the field's - // pointee, corrupting the base pointer's recorded type so a sibling - // field_addr on the same base later resolves to the wrong (0-field) - // pointee -- the "field_addr index N out of bounds for struct with 0 - // fields" failure on a struct with >= 2 ZST fields (e.g. a closure - // that captures other ZST closures, like iter map_fold). - use llvm_export::ops::GepIndex; - let i8_ty: TypeHandle = IntegerType::get(ctx, 8, Signedness::Signless).into(); - let gep = - llvm::GetElementPtrOp::new(ctx, ptr_operand, vec![GepIndex::Constant(0)], i8_ty); - rewriter.insert_operation(ctx, gep.get_operation()); - rewriter.replace_operation(ctx, op, gep.get_operation()); - return Ok(()); - } - None => { - return pliron::input_err_noloc!( - "field_addr index {} out of bounds for struct with {} fields", - field_index, - map.decl_to_llvm.len() - ); - } - }; - - let rustc_offset = layout.field_offsets.get(field_index).copied(); - - // Preserve the #859 address-path contract independently of the value - // representation. `build_struct_slot_map` retains the offsets the same - // fields would have under natural LLVM layout; when those differ from - // rustc's recorded offsets, keep using the original aggregate pointer plus - // a byte GEP. The semantic value type may now be an LLVM packed struct, but - // changing this established field-address path is outside this change. - if let Some(expected_offset) = rustc_offset { - let actual_offset = map - .natural_slot_offsets - .as_ref() - .and_then(|offsets| offsets.get(slot as usize).copied()) - .ok_or_else(|| { - pliron::input_error_noloc!( - "field_addr: cannot determine LLVM byte offset of slot {} for field {}", - slot, - field_index - ) - })?; - if map.layout_diverges && actual_offset != expected_offset { - let field_ptr = byte_offset_gep(ctx, rewriter, ptr_operand, expected_offset); - let gep = field_ptr - .defining_op() - .expect("byte_offset_gep always returns a GEP result"); - stamp_field_address_alignment(ctx, gep, aggregate_abi_align, Some(expected_offset)); - rewriter.replace_operation(ctx, op, gep); - return Ok(()); - } - } - - use llvm_export::ops::GepIndex; - let gep_indices = vec![GepIndex::Constant(0), GepIndex::Constant(slot)]; - - let gep_op = llvm::GetElementPtrOp::new(ctx, ptr_operand, gep_indices, map.llvm_struct_ty); - rewriter.insert_operation(ctx, gep_op.get_operation()); - stamp_field_address_alignment( - ctx, - gep_op.get_operation(), - aggregate_abi_align, - rustc_offset, - ); - rewriter.replace_operation(ctx, op, gep_op.get_operation()); - - Ok(()) -} - -/// Record on a field address what its alignment provably is. -/// -/// A load reads its alignment from its own result type, and a scalar records -/// none -- so `p.x` on an `#[repr(C, align(8))]` struct would export with LLVM's -/// default `align 4`, dropping the guarantee rustc gave the aggregate. That is -/// what stops LoadStoreVectorizer from fusing two adjacent field reads into one -/// wide access; a whole-element copy of the same struct already vectorizes, -/// because there the load's result type *is* the aggregate. -/// -/// The aggregate is aligned to `abi_align`, so its field at byte `offset` is -/// aligned to `gcd(abi_align, offset)`, and to `abi_align` itself at offset -/// zero. Both numbers are rustc's own layout, so this claims nothing the source -/// did not already guarantee. -/// -/// Stamps nothing when the aggregate records no alignment (`abi_align == 0`, -/// which also stands for "layout unknown") or the field has no recorded offset. -/// A wrong alignment here is a miscompile, so every uncertain case declines and -/// leaves the previous, weaker-but-sound behaviour. -fn stamp_field_address_alignment( - ctx: &mut Context, - gep: Ptr, - abi_align: u64, - field_offset: Option, -) { - const fn gcd(a: u64, b: u64) -> u64 { - if b == 0 { a } else { gcd(b, a % b) } - } - - if abi_align == 0 { - return; - } - let Some(offset) = field_offset else { - return; - }; - let provable = if offset == 0 { - abi_align - } else { - gcd(abi_align, offset) - }; - // Every rustc layout has a power-of-two `abi_align`, but dialect-mir only - // verifier-enforces that for unions and enums; a malformed hand-built - // struct or tuple layout could reach here with e.g. 12, and a - // non-power-of-two `align N` is invalid LLVM IR that llc rejects. - if !provable.is_power_of_two() { - return; - } - if let Ok(align) = u32::try_from(provable) { - llvm_export::ops::set_address_alignment(ctx, gep, align); - } -} - -// ============================================================================ -// MirArrayElementAddrOp Conversion -// ============================================================================ - -/// Convert `mir.array_element_addr` to `llvm.getelementptr`. -/// -/// ```text -/// &arr[i] ──► getelementptr T, ptr %arr_ptr, i64 %i -/// ``` -/// -/// - The element type `T` comes from the op's OWN result type -/// (`mir.ptr`), which the dialect verifier ties to the operand -/// array's element type. Same address as the old two-index -/// `[N x T]` GEP, minus the dead `[N x T]` base. -/// - Operand type history is not usable here: a kind-only `mir.cast` -/// lowers to a plain value forwarding, history does not follow that -/// edge, and a stale hit would stride by the wrong element size. -pub(crate) fn convert_array_element_addr( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - _operands_info: &OperandsInfo, -) -> Result<()> { - let loc = op.deref(ctx).loc(); - let arr_ptr = op.deref(ctx).get_operand(0); - let index = op.deref(ctx).get_operand(1); - - let result_ty = op.deref(ctx).get_result(0).get_type(ctx); - let element_ty = result_ty - .deref(ctx) - .downcast_ref::() - .map(|mir_ptr| mir_ptr.pointee) - .ok_or_else(|| { - pliron::input_error!( - loc.clone(), - "mir.array_element_addr result must be a MIR pointer type; \ - element sizing has no fact to derive from" - ) - })?; - - let llvm_element_ty = convert_type(ctx, element_ty).map_err(anyhow_to_pliron)?; - - // The typed GEP strides by the LLVM allocation size of the element type. - // Keep the rustc-stride comparison as a fail-closed backstop for any - // element representation whose LLVM allocation size differs. - { - let rustc_stride = mir_element_stride(ctx, element_ty); - let llvm_size = llvm_type_size_align(ctx, llvm_element_ty).map(|(size, _)| size); - if let (Some(stride), Some(llvm_size)) = (rustc_stride, llvm_size) - && stride != llvm_size - { - return pliron::input_err_noloc!( - "addressing elements of an array whose LLVM allocation size differs from \ - rustc's stored stride is not supported: rustc strides by {} bytes but \ - the LLVM element type occupies {}", - stride, - llvm_size - ); - } - } - - use llvm_export::ops::GepIndex; - let gep_indices = vec![GepIndex::Value(index)]; - - let element_align = element_address_provable_alignment(ctx, arr_ptr, element_ty, index); - - let gep_op = llvm::GetElementPtrOp::new(ctx, arr_ptr, gep_indices, llvm_element_ty); - rewriter.insert_operation(ctx, gep_op.get_operation()); - if let Some(align) = element_align { - llvm_export::ops::set_address_alignment(ctx, gep_op.get_operation(), align); - } - rewriter.replace_operation(ctx, op, gep_op.get_operation()); - - Ok(()) -} - -/// Alignment an array-element address provably has, in bytes. -/// -/// The field path records this already ([`stamp_field_address_alignment`]), but -/// an element address recorded nothing, so `lanes[0] + lanes[1]` on an -/// `#[repr(C, align(8))]` `[f32; 2]` still exported two `align 4` loads and -/// LoadStoreVectorizer refused to fuse them. Reading the same element into a -/// local first happens to work, because SROA then loads the whole aggregate and -/// the alignment comes from its type — so the cost depended on whether the -/// source copied the element or read through it. -/// -/// Two facts combine. The element's `abi_align` equals the array's (a Rust -/// array aligns to its element), and whatever the *base pointer* already -/// proved is stronger when the array is itself a field of an over-aligned -/// aggregate: `&table[i].lanes` carries the outer struct's alignment, which -/// the element type alone does not know. Element `i` then sits at byte -/// `i * stride`, so it is aligned to `gcd(base, i * stride)`, and to `base` -/// itself at index zero. A runtime index can land on any element, so it gets -/// `gcd(base, stride)`, which every element satisfies. -/// -/// Answers `None`, leaving the previous behaviour, when neither the base nor -/// the element type records an alignment, or the stride is unknown. As on the -/// field path, a wrong answer here is a miscompile, so every uncertain case -/// declines rather than guesses. That is why the stride comes from -/// [`mir_element_stride`], which is exact or `None`: an LLVM-level size -/// approximation would guess 8 for a MIR aggregate element it does not -/// model, and e.g. `[[f32; 3]; 4]` under an align-8 base would then claim -/// align 8 on element addresses that are only 4-aligned. -fn element_address_provable_alignment( - ctx: &Context, - arr_ptr: Value, - element_ty: TypeHandle, - index: Value, -) -> Option { - const fn gcd(a: u64, b: u64) -> u64 { - if b == 0 { a } else { gcd(b, a % b) } - } - - // What the base address already proved, else what the element type - // states (identical to the array's own alignment). - let base_align = arr_ptr - .defining_op() - .and_then(|def| llvm_export::ops::address_alignment(ctx, def)) - .map(u64::from) - .or_else(|| mir_type_abi_align(ctx, element_ty))?; - if base_align == 0 { - return None; - } - - let stride = mir_element_stride(ctx, element_ty)?; - if stride == 0 { - return None; - } - - let provable = match constant_index_value(ctx, index) { - Some(0) => base_align, - Some(i) => gcd(base_align, i.checked_mul(stride)?), - // Any element is reachable, so claim only what every stride preserves. - None => gcd(base_align, stride), - }; - - // Same guard as the field path: dialect-mir does not verifier-enforce a - // power-of-two `abi_align` for every aggregate, and a non-power-of-two - // `align N` is invalid LLVM IR that llc rejects. - if !provable.is_power_of_two() { - return None; - } - u32::try_from(provable).ok() -} - -/// The constant an index operand holds, if it is one. -/// -/// `APInt::to_u64` truncates wider values, so a >64-bit constant could be -/// misread as a small offset multiplier. Fail closed on such widths, as -/// `integer_constant_u64` in the extract-element fast path does. -fn constant_index_value(ctx: &Context, index: Value) -> Option { - let defining_op = index.defining_op()?; - if let Some(constant) = Operation::get_op::(defining_op, ctx) { - let value = constant.get_attr_value(ctx)?.value(); - return (value.bw() <= 64).then(|| value.to_u64()); - } - let constant = Operation::get_op::(defining_op, ctx)?; - let attribute = constant.get_value(ctx); - let integer = attribute.downcast_ref::()?; - let value = integer.value(); - (value.bw() <= 64).then(|| value.to_u64()) -} - -#[cfg(test)] -// Tests build kinded fixture types directly; production minting lives in mir-importer's facts.rs. -#[allow(clippy::disallowed_methods)] -mod tests { - use super::*; - use crate::convert::ops::test_util::*; - use dialect_mir::attributes::{FieldIndexAttr, MirCastKindAttr, VariantIndexAttr}; - use dialect_mir::ops as mir; - use dialect_mir::types::{ - EnumEncoding, EnumVariant, MirArrayType, MirPointerKind, MirPtrType, MirSliceType, - MirStructType, MirTupleType, - }; - use llvm_export::types as llvm_types; - use pliron::builtin::attributes::IntegerAttr; - use pliron::common_traits::Verify; - - fn insert_indices(ctx: &Context, inserts: &[llvm::InsertValueOp]) -> Vec> { - inserts.iter().map(|op| op.indices(ctx)).collect() - } - - fn empty_struct_ty(ctx: &mut Context, name: &str) -> TypeHandle { - MirStructType::get(ctx, name.to_string(), vec![], vec![]).into() - } - - fn padded_struct_with_zst_ty(ctx: &mut Context) -> (TypeHandle, TypeHandle) { - let i8_ty: TypeHandle = IntegerType::get(ctx, 8, Signedness::Signless).into(); - let i64_ty: TypeHandle = IntegerType::get(ctx, 64, Signedness::Signless).into(); - let zst_ty = empty_struct_ty(ctx, "Marker"); - - let struct_ty = MirStructType::get_with_full_layout( - ctx, - "Padded".to_string(), - vec!["a".to_string(), "marker".to_string(), "b".to_string()], - vec![i8_ty, zst_ty, i64_ty], - vec![0, 1, 2], - vec![0, 1, 8], - 16, - 8, - ); - - (struct_ty.into(), zst_ty) - } - - fn append_empty_struct_value( - ctx: &mut Context, - block: Ptr, - zst_ty: TypeHandle, - ) -> Value { - let op = Operation::new( - ctx, - mir::MirConstructStructOp::get_concrete_op_info(), - vec![zst_ty], - vec![], - vec![], - 0, - ); - op.insert_at_back(block, ctx); - op.deref(ctx).get_result(0) - } - - fn byte_gep_constant_offset(ctx: &Context, gep: &llvm::GetElementPtrOp) -> Option { - use llvm_export::ops::GepIndex; - - let indices = gep.indices(ctx); - let [GepIndex::Value(offset)] = indices.as_slice() else { - return None; - }; - let defining_op = offset.defining_op()?; - let constant = Operation::get_op::(defining_op, ctx)?; - let attribute = constant.get_value(ctx); - attribute - .downcast_ref::() - .map(|integer| integer.value().to_u64()) - } - - fn assert_byte_addressed_field( - ctx: &Context, - module: Ptr, - expected_offset: u64, - expected_alignment: u32, - ) { - let body = kernel_blocks(ctx, module); - let geps = find_all::(ctx, &body); - assert_eq!(geps.len(), 1, "one field_addr must lower to one GEP"); - - let gep = &geps[0]; - assert_eq!( - gep.src_elem_type(ctx), - IntegerType::get(ctx, 8, Signedness::Signless).into(), - "a layout-mismatched field must be addressed in byte units" - ); - assert_eq!( - byte_gep_constant_offset(ctx, gep), - Some(expected_offset), - "the byte GEP must use rustc's exact field offset" - ); - assert_eq!( - llvm_export::ops::address_alignment(ctx, gep.get_operation()), - Some(expected_alignment), - "the byte GEP must retain the alignment proved by rustc's aggregate layout" - ); - } - - /// `mir.construct_slice` lowers to the canonical fat-pointer value: - /// `undef { ptr, i64 }`, then insert data pointer at slot 0 and length at slot 1. - #[test] - fn construct_slice_lowers_to_ptr_len_insert_values() { - let mut ctx = make_ctx(); - - let i8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); - let usize_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); - let ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, i8_ty, false).into(); - let slice_ty: TypeHandle = MirSliceType::get(&mut ctx, i8_ty).into(); - - let (module_ptr, block) = build_kernel(&mut ctx, vec![ptr_ty, usize_ty], vec![]); - let data_ptr = block.deref(&ctx).get_argument(0); - let len = block.deref(&ctx).get_argument(1); - - let op = Operation::new( - &mut ctx, - mir::MirConstructSliceOp::get_concrete_op_info(), - vec![slice_ty], - vec![data_ptr, len], - vec![], - 0, - ); - op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - let inserts = find_all::(&ctx, &body); - - assert_eq!( - insert_indices(&ctx, &inserts), - vec![vec![0], vec![1]], - "slice construction must insert data pointer at slot 0 and length at slot 1" - ); - let first_insert = inserts[0].get_operation(); - let second_insert = inserts[1].get_operation(); - assert_eq!( - first_insert.deref(&ctx).get_operand(1), - data_ptr, - "slice slot 0 must receive the original data pointer" - ); - assert_eq!( - second_insert.deref(&ctx).get_operand(0), - first_insert.deref(&ctx).get_result(0), - "the length insertion must consume the aggregate produced by the pointer insertion" - ); - assert_eq!( - second_insert.deref(&ctx).get_operand(1), - len, - "slice slot 1 must receive the original length" - ); - assert!( - inserts.iter().all(|insert| insert.verify(&ctx).is_ok()), - "both slice insertions must satisfy LLVM dialect verification" - ); - assert_eq!( - count_ops::(&ctx, &body), - 1, - "slice construction should start from one undef aggregate" - ); - } - - /// `mir.construct_disjoint_slice` lowers to the same insert chain as the - /// fat pointer for an index space with no runtime layout. - #[test] - fn construct_disjoint_slice_lowers_to_ptr_len_insert_values() { - let mut ctx = make_ctx(); - - let f32_ty: TypeHandle = pliron::builtin::types::FP32Type::get(&ctx).into(); - let usize_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); - let ptr_ty: TypeHandle = - MirPtrType::get_generic_with_kind(&mut ctx, f32_ty, true, MirPointerKind::RawMut) - .into(); - let slice_ty: TypeHandle = MirDisjointSliceType::get(&mut ctx, f32_ty).into(); - - let (module_ptr, block) = build_kernel(&mut ctx, vec![ptr_ty, usize_ty], vec![]); - let data_ptr = block.deref(&ctx).get_argument(0); - let len = block.deref(&ctx).get_argument(1); - - let op = Operation::new( - &mut ctx, - mir::MirConstructDisjointSliceOp::get_concrete_op_info(), - vec![slice_ty], - vec![data_ptr, len], - vec![], - 0, - ); - op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - let inserts = find_all::(&ctx, &body); - assert_eq!( - insert_indices(&ctx, &inserts), - vec![vec![0], vec![1]], - "a space-free disjoint slice must insert pointer at slot 0 and length at slot 1" - ); - assert_eq!( - inserts[0].get_operation().deref(&ctx).get_operand(1), - data_ptr, - "slot 0 must receive the original data pointer" - ); - assert_eq!( - inserts[1].get_operation().deref(&ctx).get_operand(1), - len, - "slot 1 must receive the original length" - ); - assert_eq!( - count_ops::(&ctx, &body), - 1, - "construction should start from one undef aggregate" - ); - } - - /// The runtime row width is the third operand and must land in slot 2. - /// Writing it into the length slot, or dropping it, gives a slice whose - /// row width reads back as something else at every access site. - #[test] - fn construct_disjoint_slice_places_the_row_width_after_ptr_and_len() { - let mut ctx = make_ctx(); - - let f32_ty: TypeHandle = pliron::builtin::types::FP32Type::get(&ctx).into(); - let usize_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); - let width_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); - let ptr_ty: TypeHandle = - MirPtrType::get_generic_with_kind(&mut ctx, f32_ty, true, MirPointerKind::RawMut) - .into(); - let slice_ty: TypeHandle = - MirDisjointSliceType::get_with_space(&mut ctx, f32_ty, vec![width_ty]).into(); - - let (module_ptr, block) = build_kernel(&mut ctx, vec![ptr_ty, usize_ty, width_ty], vec![]); - let data_ptr = block.deref(&ctx).get_argument(0); - let len = block.deref(&ctx).get_argument(1); - let width = block.deref(&ctx).get_argument(2); - - let op = Operation::new( - &mut ctx, - mir::MirConstructDisjointSliceOp::get_concrete_op_info(), - vec![slice_ty], - vec![data_ptr, len, width], - vec![], - 0, - ); - op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - let inserts = find_all::(&ctx, &body); - assert_eq!( - insert_indices(&ctx, &inserts), - vec![vec![0], vec![1], vec![2]], - "the row width must occupy slot 2, after the pointer and the length" - ); - assert_eq!( - inserts[2].get_operation().deref(&ctx).get_operand(1), - width, - "slot 2 must receive the row width operand, not the length" - ); - assert_eq!( - inserts[1].get_operation().deref(&ctx).get_operand(1), - len, - "slot 1 must still receive the length" - ); - } - - /// Explicit rustc layout must be respected: field `b` is at byte offset 8, - /// so the lowered LLVM struct has a padding slot between `a` and `b`. - /// The ZST marker field is stripped and receives no insert_value. - #[test] - fn construct_struct_uses_layout_slots_and_skips_zst() { - let mut ctx = make_ctx(); - - let i8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Signless).into(); - let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signless).into(); - let (struct_ty, zst_ty) = padded_struct_with_zst_ty(&mut ctx); - - let (module_ptr, block) = build_kernel(&mut ctx, vec![i8_ty, i64_ty], vec![]); - let a = block.deref(&ctx).get_argument(0); - let b = block.deref(&ctx).get_argument(1); - let marker = append_empty_struct_value(&mut ctx, block, zst_ty); - - let op = Operation::new( - &mut ctx, - mir::MirConstructStructOp::get_concrete_op_info(), - vec![struct_ty], - vec![a, marker, b], - vec![], - 0, - ); - op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - let inserts = find_all::(&ctx, &body); - - assert_eq!( - insert_indices(&ctx, &inserts), - vec![vec![0], vec![2]], - "non-ZST fields must be inserted at their layout slots, skipping padding and ZSTs" - ); - let first_insert = inserts[0].get_operation(); - let second_insert = inserts[1].get_operation(); - assert_eq!( - first_insert.deref(&ctx).get_operand(1), - a, - "struct slot 0 must receive field `a`" - ); - assert_eq!( - second_insert.deref(&ctx).get_operand(0), - first_insert.deref(&ctx).get_result(0), - "field `b` must be inserted into the aggregate containing field `a`" - ); - assert_eq!( - second_insert.deref(&ctx).get_operand(1), - b, - "struct slot 2 must receive field `b`" - ); - assert!( - inserts.iter().all(|insert| insert.verify(&ctx).is_ok()), - "both struct insertions must satisfy LLVM dialect verification" - ); - } - - /// Extracting a ZST field must not emit `extract_value`: the field has no - /// storage in the lowered LLVM struct, so lowering materializes an undef - /// zero-sized value instead. - #[test] - fn extract_zst_field_lowers_to_undef_without_extract_value() { - let mut ctx = make_ctx(); - - let i8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Signless).into(); - let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signless).into(); - let (struct_ty, zst_ty) = padded_struct_with_zst_ty(&mut ctx); - - let (module_ptr, block) = build_kernel(&mut ctx, vec![i8_ty, i64_ty], vec![]); - let a = block.deref(&ctx).get_argument(0); - let b = block.deref(&ctx).get_argument(1); - let marker = append_empty_struct_value(&mut ctx, block, zst_ty); - - let construct = Operation::new( - &mut ctx, - mir::MirConstructStructOp::get_concrete_op_info(), - vec![struct_ty], - vec![a, marker, b], - vec![], - 0, - ); - construct.insert_at_back(block, &ctx); - let aggregate = construct.deref(&ctx).get_result(0); - - let extract = Operation::new( - &mut ctx, - MirExtractFieldOp::get_concrete_op_info(), - vec![zst_ty], - vec![aggregate], - vec![], - 0, - ); - MirExtractFieldOp::new(extract).set_attr_index(&ctx, FieldIndexAttr(1)); - extract.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - - assert_eq!( - count_ops::(&ctx, &body), - 0, - "extracting a stripped ZST field must not emit llvm.extractvalue" - ); - - let zst_un_defs = find_all::(&ctx, &body) - .into_iter() - .filter(|op| { - let result_ty = op.get_operation().deref(&ctx).get_result(0).get_type(&ctx); - is_zero_sized_type(&ctx, result_ty) - }) - .count(); - - assert_eq!( - zst_un_defs, 2, - "one undef should build the ZST value and one should materialize the extracted ZST" - ); - } - - /// Sibling ZST field addresses off one base pointer must each lower to - /// their own zero-offset `i8` GEP, never to the base SSA value itself. - /// Forwarding the base (the pre-fix behavior) recorded the first field's - /// zero-field pointee type onto the base pointer's conversion history, so - /// the second `field_addr` resolved the base's pointee to "struct with 0 - /// fields" and lowering aborted. This is the `iter().map(..).sum()` shape: - /// a composed closure holding two captureless (zero-sized) closures. - #[test] - fn sibling_zst_field_addrs_lower_to_distinct_zero_offset_geps() { - use llvm_export::ops::GepIndex; - - let mut ctx = make_ctx(); - - let i8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Signless).into(); - let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signless).into(); - let zst_f = empty_struct_ty(&mut ctx, "MapClosure"); - let zst_g = empty_struct_ty(&mut ctx, "SumClosure"); - - // struct Composed { f: MapClosure (ZST), g: SumClosure (ZST), acc: i64 } - let struct_ty: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "Composed".to_string(), - vec!["f".to_string(), "g".to_string(), "acc".to_string()], - vec![zst_f, zst_g, i64_ty], - vec![0, 1, 2], - vec![0, 0, 0], - 8, - 8, - ) - .into(); - - let base_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, struct_ty, true).into(); - // Each field_addr result records the field's own pointee type; for the - // ZST fields that zero-field pointee is exactly what poisoned the base - // pointer's history when the result aliased the base. - let f_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, zst_f, true).into(); - let g_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, zst_g, true).into(); - let acc_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, i64_ty, true).into(); - - let (module_ptr, block) = build_kernel(&mut ctx, vec![base_ptr_ty], vec![]); - let base = block.deref(&ctx).get_argument(0); - - for (field_index, result_ty) in [(0u32, f_ptr_ty), (1, g_ptr_ty), (2, acc_ptr_ty)] { - let op = MirFieldAddrOp::build(&mut ctx, base, result_ty, field_index) - .expect("field_addr build"); - op.insert_at_back(block, &ctx); - } - append_mir_return(&mut ctx, block, vec![]); - - // Pre-fix this failed with "field_addr index 1 out of bounds for - // struct with 0 fields" on the second (sibling) ZST field_addr. - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - let geps = find_all::(&ctx, &body); - assert_eq!(geps.len(), 3, "each field_addr must lower to its own GEP"); - assert!( - geps.iter().all(|gep| gep.verify(&ctx).is_ok()), - "every field-address GEP must satisfy LLVM dialect verification" - ); - - let zst_geps: Vec<_> = geps - .iter() - .filter(|gep| { - gep.src_elem_type(&ctx) == i8_ty - && matches!(gep.indices(&ctx).as_slice(), [GepIndex::Constant(0)]) - }) - .collect(); - assert_eq!( - zst_geps.len(), - 2, - "both ZST field addresses must lower to zero-offset i8 GEPs" - ); - - let gep_base = |gep: &llvm::GetElementPtrOp| gep.get_operation().deref(&ctx).get_operand(0); - let gep_result = - |gep: &llvm::GetElementPtrOp| gep.get_operation().deref(&ctx).get_result(0); - assert_eq!( - gep_base(zst_geps[0]), - gep_base(zst_geps[1]), - "both ZST field addresses must be taken off the same base pointer" - ); - assert_ne!( - gep_result(zst_geps[0]), - gep_base(zst_geps[0]), - "a ZST field address must be a value distinct from the base pointer" - ); - assert_ne!( - gep_result(zst_geps[1]), - gep_base(zst_geps[1]), - "a ZST field address must be a value distinct from the base pointer" - ); - assert_ne!( - gep_result(zst_geps[0]), - gep_result(zst_geps[1]), - "each sibling ZST field address must get its own GEP result" - ); - - // The non-ZST sibling still resolves the base's pointee to the struct - // (slot 0 holds `acc`): the base pointer's type history stayed intact. - let acc_geps: Vec<_> = geps - .iter() - .filter(|gep| { - matches!( - gep.indices(&ctx).as_slice(), - [GepIndex::Constant(0), GepIndex::Constant(0)] - ) - }) - .collect(); - assert_eq!( - acc_geps.len(), - 1, - "the non-ZST field address must index the struct's layout slot" - ); - } - - /// `mir.field_addr` on a TUPLE pointee (the `#693` shape: `let (a, b) = - /// TABLE[i];`) must resolve the GEP index through the tuple's memory-order - /// layout, exactly like the struct path above, not through the - /// declaration index directly. - /// - /// `(u8, u32)` is rustc's own layout for this pair: the `u32` field is - /// placed FIRST in memory for alignment, so declaration field 0 (`u8`, - /// `.0`) lives at memory slot 1 and declaration field 1 (`u32`, `.1`) - /// lives at memory slot 0. A GEP index equal to the declaration index - /// would silently address the WRONG field's bytes. - #[test] - fn field_addr_tuple_pointee_resolves_memory_order_gep_index() { - use llvm_export::ops::GepIndex; - - let mut ctx = make_ctx(); - - let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); - let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); - - // decl order [u8, u32]; memory order [u32, u8] (mem_to_decl = [1, 0]). - let tuple_ty: TypeHandle = MirTupleType::get_with_layout( - &mut ctx, - vec![u8_ty, u32_ty], - vec![1, 0], - vec![4, 0], - 8, - 4, - ) - .into(); - - let tuple_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, tuple_ty, false).into(); - let u8_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, u8_ty, false).into(); - let u32_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, u32_ty, false).into(); - - let (module_ptr, block) = build_kernel(&mut ctx, vec![tuple_ptr_ty], vec![]); - let base = block.deref(&ctx).get_argument(0); - - // Declaration field 0 (`.0`, u8) first, then declaration field 1 - // (`.1`, u32) -- source order, not memory order. - for (field_index, result_ty) in [(0u32, u8_ptr_ty), (1, u32_ptr_ty)] { - let op = MirFieldAddrOp::build(&mut ctx, base, result_ty, field_index) - .expect("field_addr build"); - op.insert_at_back(block, &ctx); - } - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - let geps = find_all::(&ctx, &body); - assert_eq!(geps.len(), 2, "each field_addr must lower to its own GEP"); - assert!( - geps.iter().all(|gep| gep.verify(&ctx).is_ok()), - "every field-address GEP must satisfy LLVM dialect verification" - ); - - // `.0` (u8, declared first) must land at MEMORY slot 1. - let field0_geps: Vec<_> = geps - .iter() - .filter(|gep| { - matches!( - gep.indices(&ctx).as_slice(), - [GepIndex::Constant(0), GepIndex::Constant(1)] - ) - }) - .collect(); - assert_eq!( - field0_geps.len(), - 1, - "declaration field 0 (u8) must resolve to its memory slot 1, not slot 0" - ); - - // `.1` (u32, declared second) must land at MEMORY slot 0. - let field1_geps: Vec<_> = geps - .iter() - .filter(|gep| { - matches!( - gep.indices(&ctx).as_slice(), - [GepIndex::Constant(0), GepIndex::Constant(0)] - ) - }) - .collect(); - assert_eq!( - field1_geps.len(), - 1, - "declaration field 1 (u32) must resolve to its memory slot 0, not slot 1" - ); - } - - /// A `#[repr(C, packed)]`-style layout can place a naturally aligned - /// scalar at an offset LLVM's ordinary struct layout cannot express. The - /// field address must therefore use rustc's byte offset rather than a - /// typed struct GEP that would silently land at byte 4. - #[test] - fn packed_field_addr_uses_rustc_byte_offset_and_alignment_one() { - let mut ctx = make_ctx(); - - let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); - let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); - let packed_ty: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "Packed".into(), - vec!["tag".into(), "value".into()], - vec![u8_ty, u32_ty], - vec![0, 1], - vec![0, 1], - 5, - 1, - ) - .into(); - - let base_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, packed_ty, false).into(); - let field_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, u32_ty, false).into(); - let (module, block) = build_kernel(&mut ctx, vec![base_ptr_ty], vec![]); - let base = block.deref(&ctx).get_argument(0); - - let field_addr = - MirFieldAddrOp::build(&mut ctx, base, field_ptr_ty, 1).expect("field_addr build"); - field_addr.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); - assert_byte_addressed_field(&ctx, module, 1, 1); - } - - /// `repr(packed(2))` is not equivalent to byte alignment: the same u32 - /// field is allowed to sit at byte 2 and the address still proves align 2. - #[test] - fn packed_two_field_addr_uses_rustc_byte_offset_and_alignment_two() { - let mut ctx = make_ctx(); - - let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); - let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); - let packed_ty: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "Packed2".into(), - vec!["tag".into(), "value".into()], - vec![u8_ty, u32_ty], - vec![0, 1], - vec![0, 2], - 6, - 2, - ) - .into(); - - let base_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, packed_ty, false).into(); - let field_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, u32_ty, false).into(); - let (module, block) = build_kernel(&mut ctx, vec![base_ptr_ty], vec![]); - let base = block.deref(&ctx).get_argument(0); - - let field_addr = - MirFieldAddrOp::build(&mut ctx, base, field_ptr_ty, 1).expect("field_addr build"); - field_addr.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); - assert_byte_addressed_field(&ctx, module, 2, 2); - } - - /// The decision must compare physical offsets, not merely the selected - /// field's own alignment. In `{ u8, u32, u8 }` the final u8 is naturally - /// byte-aligned, but LLVM has already shifted it because the packed u32 - /// before it was placed at byte 4 instead of byte 1. - #[test] - fn packed_trailing_byte_field_uses_accumulated_rustc_offset() { - let mut ctx = make_ctx(); - - let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); - let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); - let packed_ty: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "PackedTrailing".into(), - vec!["head".into(), "value".into(), "tail".into()], - vec![u8_ty, u32_ty, u8_ty], - vec![0, 1, 2], - vec![0, 1, 5], - 6, - 1, - ) - .into(); - - let base_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, packed_ty, false).into(); - let field_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, u8_ty, false).into(); - let (module, block) = build_kernel(&mut ctx, vec![base_ptr_ty], vec![]); - let base = block.deref(&ctx).get_argument(0); - - let field_addr = - MirFieldAddrOp::build(&mut ctx, base, field_ptr_ty, 2).expect("field_addr build"); - field_addr.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); - assert_byte_addressed_field(&ctx, module, 5, 1); - } - - /// A by-value packed struct is represented by an LLVM packed struct, so - /// construction can use the ordinary insertvalue chain without changing - /// rustc's field offsets or total byte size. - #[test] - fn packed_struct_construction_by_value_uses_packed_layout() { - use dialect_mir::ops::MirConstructStructOp; - - let mut ctx = make_ctx(); - - let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); - let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); - let packed_ty: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "Packed".into(), - vec!["tag".into(), "value".into()], - vec![u8_ty, u32_ty], - vec![0, 1], - vec![0, 1], - 5, - 1, - ) - .into(); - - let (module, block) = build_kernel(&mut ctx, vec![u8_ty, u32_ty], vec![]); - let tag = block.deref(&ctx).get_argument(0); - let value = block.deref(&ctx).get_argument(1); - - let construct = Operation::new( - &mut ctx, - MirConstructStructOp::get_concrete_op_info(), - vec![packed_ty], - vec![tag, value], - vec![], - 0, - ); - construct.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module) - .expect("constructing a pointer-free packed struct by value must lower"); - - let body = kernel_blocks(&ctx, module); - let inserts = find_all::(&ctx, &body); - assert_eq!( - insert_indices(&ctx, &inserts), - vec![vec![0], vec![1]], - "packed construction must keep declaration fields in their packed LLVM slots" - ); - let result_ty = inserts - .last() - .expect("packed construction must emit insertvalue") - .get_operation() - .deref(&ctx) - .get_result(0) - .get_type(&ctx); - let result_ty_ref = result_ty.deref(&ctx); - let struct_ty = result_ty_ref - .downcast_ref::() - .expect("packed construction result must be an LLVM struct"); - assert_eq!(struct_ty.layout(), llvm_types::StructLayout::Packed); - assert_eq!(llvm_type_size_align(&ctx, result_ty), Some((5, 1))); - } - - #[test] - fn packed_struct_construction_with_one_direct_shared_pointer_stays_semantic_in_ssa() { - use dialect_mir::ops::MirConstructStructOp; - - let mut ctx = make_ctx(); - let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); - let pointee: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); - let shared_ty: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, false).into(); - let packed_ty: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "PackedShared".into(), - vec!["tag".into(), "ptr".into()], - vec![u8_ty, shared_ty], - vec![0, 1], - vec![0, 1], - 9, - 1, - ) - .into(); - - let (module, block) = build_kernel(&mut ctx, vec![u8_ty, shared_ty], vec![]); - let tag = block.deref(&ctx).get_argument(0); - let ptr = block.deref(&ctx).get_argument(1); - let construct = Operation::new( - &mut ctx, - MirConstructStructOp::get_concrete_op_info(), - vec![packed_ty], - vec![tag, ptr], - vec![], - 0, - ); - construct.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module) - .expect("one direct AS3 pointer in a packed SSA aggregate must lower"); - - let body = kernel_blocks(&ctx, module); - let inserts = find_all::(&ctx, &body); - assert_eq!( - insert_indices(&ctx, &inserts), - vec![vec![0], vec![1]], - "packed AS3 construction must preserve the semantic packed field slots" - ); - - let result_ty = inserts - .last() - .expect("packed AS3 construction must emit insertvalue") - .get_operation() - .deref(&ctx) - .get_result(0) - .get_type(&ctx); - let result_ty_ref = result_ty.deref(&ctx); - let struct_ty = result_ty_ref - .downcast_ref::() - .expect("packed AS3 construction result must be an LLVM struct"); - assert_eq!(struct_ty.layout(), llvm_types::StructLayout::Packed); - assert_eq!(llvm_type_size_align(&ctx, result_ty), Some((9, 1))); - - let pointer_ty = struct_ty.field_type(1); - let pointer_ty_ref = pointer_ty.deref(&ctx); - let pointer = pointer_ty_ref - .downcast_ref::() - .expect("second packed field must remain a pointer"); - assert_eq!( - pointer.address_space(), - llvm_types::address_space::SHARED, - "construction must keep the semantic pointer in AS3" - ); - assert_eq!( - count_ops::(&ctx, &body), - 0, - "construction itself must not genericize the shared pointer" - ); - } - - #[test] - fn packed_struct_construction_with_multiple_direct_shared_pointers_stays_semantic_in_ssa() { - use dialect_mir::ops::MirConstructStructOp; - - let mut ctx = make_ctx(); - let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); - let pointee: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); - let shared_ty: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, false).into(); - let packed_ty: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "PackedSharedPair".into(), - vec!["tag".into(), "left".into(), "right".into()], - vec![u8_ty, shared_ty, shared_ty], - vec![0, 1, 2], - vec![0, 1, 9], - 17, - 1, - ) - .into(); - - let (module, block) = build_kernel(&mut ctx, vec![u8_ty, shared_ty, shared_ty], vec![]); - let tag = block.deref(&ctx).get_argument(0); - let left = block.deref(&ctx).get_argument(1); - let right = block.deref(&ctx).get_argument(2); - let construct = Operation::new( - &mut ctx, - MirConstructStructOp::get_concrete_op_info(), - vec![packed_ty], - vec![tag, left, right], - vec![], - 0, - ); - construct.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module) - .expect("multiple direct AS3 leaves in a packed SSA aggregate must lower"); - - let body = kernel_blocks(&ctx, module); - let inserts = find_all::(&ctx, &body); - assert_eq!( - insert_indices(&ctx, &inserts), - vec![vec![0], vec![1], vec![2]], - "all direct shared leaves must remain in their semantic packed slots" - ); - - let result_ty = inserts - .last() - .expect("packed AS3 construction must emit insertvalue") - .get_operation() - .deref(&ctx) - .get_result(0) - .get_type(&ctx); - let result_ty_ref = result_ty.deref(&ctx); - let struct_ty = result_ty_ref - .downcast_ref::() - .expect("packed AS3 construction result must be an LLVM struct"); - assert_eq!(struct_ty.layout(), llvm_types::StructLayout::Packed); - assert_eq!(llvm_type_size_align(&ctx, result_ty), Some((17, 1))); - for slot in [1usize, 2] { - let pointer_ty = struct_ty.field_type(slot); - let pointer_ty_ref = pointer_ty.deref(&ctx); - let pointer = pointer_ty_ref - .downcast_ref::() - .expect("direct shared leaves must remain pointer-typed"); - assert_eq!(pointer.address_space(), llvm_types::address_space::SHARED); - } - assert_eq!( - count_ops::(&ctx, &body), - 0, - "semantic construction must not genericize any direct shared leaf" - ); - } - - #[test] - fn packed_struct_construction_with_nested_shared_pointers_stays_semantic_in_ssa() { - use dialect_mir::ops::MirConstructStructOp; - - let mut ctx = make_ctx(); - let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); - let pointee: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); - let shared_ty: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, false).into(); - let inner_ty: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "SharedPair".into(), - vec!["left".into(), "right".into()], - vec![shared_ty, shared_ty], - vec![0, 1], - vec![0, 8], - 16, - 8, - ) - .into(); - let packed_ty: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "PackedNestedShared".into(), - vec!["tag".into(), "pair".into()], - vec![u8_ty, inner_ty], - vec![0, 1], - vec![0, 1], - 17, - 1, - ) - .into(); - - let (module, block) = build_kernel(&mut ctx, vec![u8_ty, shared_ty, shared_ty], vec![]); - let tag = block.deref(&ctx).get_argument(0); - let left = block.deref(&ctx).get_argument(1); - let right = block.deref(&ctx).get_argument(2); - - let inner = Operation::new( - &mut ctx, - MirConstructStructOp::get_concrete_op_info(), - vec![inner_ty], - vec![left, right], - vec![], - 0, - ); - inner.insert_at_back(block, &ctx); - let inner_value = inner.deref(&ctx).get_result(0); - - let outer = Operation::new( - &mut ctx, - MirConstructStructOp::get_concrete_op_info(), - vec![packed_ty], - vec![tag, inner_value], - vec![], - 0, - ); - outer.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module) - .expect("nested AS3 leaves in a packed SSA aggregate must lower"); - - let body = kernel_blocks(&ctx, module); - assert_eq!( - count_ops::(&ctx, &body), - 0, - "semantic construction must keep nested shared leaves in AS3" - ); - - let inserts = find_all::(&ctx, &body); - let packed_result = inserts - .iter() - .filter_map(|insert| { - let result_ty = insert - .get_operation() - .deref(&ctx) - .get_result(0) - .get_type(&ctx); - let is_packed = result_ty - .deref(&ctx) - .downcast_ref::() - .is_some_and(|ty| ty.layout() == llvm_types::StructLayout::Packed); - is_packed.then_some(result_ty) - }) - .last() - .expect("outer packed construction must produce a packed LLVM struct"); - - assert_eq!(llvm_type_size_align(&ctx, packed_result), Some((17, 1))); - assert!( - llvm_packed_struct_contains_pointer_in_address_space( - &ctx, - packed_result, - llvm_types::address_space::SHARED, - ), - "the semantic packed value must retain its recursively nested AS3 leaves" - ); - } - - #[test] - fn packed_struct_construction_with_bounded_shared_pointer_array_stays_semantic_in_ssa() { - use dialect_mir::ops::{MirConstructArrayOp, MirConstructStructOp}; - - let mut ctx = make_ctx(); - let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); - let pointee: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); - let shared_ty: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, false).into(); - let array_ty: TypeHandle = MirArrayType::get(&mut ctx, shared_ty, 2).into(); - let packed_ty: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "PackedSharedArray".into(), - vec!["tag".into(), "ptrs".into()], - vec![u8_ty, array_ty], - vec![0, 1], - vec![0, 1], - 17, - 1, - ) - .into(); - - let (module, block) = build_kernel(&mut ctx, vec![u8_ty, shared_ty, shared_ty], vec![]); - let tag = block.deref(&ctx).get_argument(0); - let left = block.deref(&ctx).get_argument(1); - let right = block.deref(&ctx).get_argument(2); - - let array = Operation::new( - &mut ctx, - MirConstructArrayOp::get_concrete_op_info(), - vec![array_ty], - vec![left, right], - vec![], - 0, - ); - array.insert_at_back(block, &ctx); - let array_value = array.deref(&ctx).get_result(0); - - let outer = Operation::new( - &mut ctx, - MirConstructStructOp::get_concrete_op_info(), - vec![packed_ty], - vec![tag, array_value], - vec![], - 0, - ); - outer.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module) - .expect("a bounded AS3 array in a packed SSA aggregate must lower"); - - let body = kernel_blocks(&ctx, module); - assert_eq!( - count_ops::(&ctx, &body), - 0, - "semantic construction must keep array elements in AS3" - ); - - let inserts = find_all::(&ctx, &body); - let packed_result = inserts - .iter() - .filter_map(|insert| { - let result_ty = insert - .get_operation() - .deref(&ctx) - .get_result(0) - .get_type(&ctx); - let is_packed = result_ty - .deref(&ctx) - .downcast_ref::() - .is_some_and(|ty| ty.layout() == llvm_types::StructLayout::Packed); - is_packed.then_some(result_ty) - }) - .last() - .expect("outer packed array construction must produce a packed LLVM struct"); - assert_eq!(llvm_type_size_align(&ctx, packed_result), Some((17, 1))); - assert!( - llvm_packed_struct_contains_pointer_in_address_space( - &ctx, - packed_result, - llvm_types::address_space::SHARED, - ), - "the semantic packed value must retain AS3 array elements" - ); - } - - /// Packed LLVM element types carry the same allocation size as rustc, so - /// typed array GEPs now stride by the correct packed byte count. - #[test] - fn packed_array_element_addressing_uses_packed_stride() { - use dialect_mir::ops::MirArrayElementAddrOp; - - let mut ctx = make_ctx(); - - let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); - let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); - let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); - let packed_ty: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "Packed".into(), - vec!["tag".into(), "value".into()], - vec![u8_ty, u32_ty], - vec![0, 1], - vec![0, 1], - 5, - 1, - ) - .into(); - let array_ty: TypeHandle = MirArrayType::get(&mut ctx, packed_ty, 4).into(); - let base_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, array_ty, false).into(); - let elem_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, packed_ty, false).into(); - - let (module, block) = build_kernel(&mut ctx, vec![base_ptr_ty, u64_ty], vec![]); - let base = block.deref(&ctx).get_argument(0); - let index = block.deref(&ctx).get_argument(1); - - let elem_addr = Operation::new( - &mut ctx, - MirArrayElementAddrOp::get_concrete_op_info(), - vec![elem_ptr_ty], - vec![base, index], - vec![], - 0, - ); - elem_addr.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module) - .expect("packed element addressing must use the packed LLVM stride"); - let body = kernel_blocks(&ctx, module); - assert_eq!( - count_ops::(&ctx, &body), - 1, - "packed array element addressing should lower to one typed GEP" - ); - } - - /// Enum construction must store the declared discriminant value, not the - /// variant index. This locks the `Ordering::Less = -1` style case as the - /// i8 bit-pattern `255`. - #[test] - fn construct_enum_uses_declared_discriminant_not_variant_index() { - let mut ctx = make_ctx(); - - let discr_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Signed).into(); - let enum_ty: TypeHandle = MirEnumType::get_with_layout( - &mut ctx, - "OrderingLike".to_string(), - discr_ty, - vec![255, 0, 1], - vec![ - EnumVariant::unit("Less".to_string()), - EnumVariant::unit("Equal".to_string()), - EnumVariant::unit("Greater".to_string()), - ], - 0, - 1, - 1, - ) - .into(); - - let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); - - let op = Operation::new( - &mut ctx, - MirConstructEnumOp::get_concrete_op_info(), - vec![enum_ty], - vec![], - vec![], - 0, - ); - MirConstructEnumOp::new(op) - .set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(0)); - op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - let inserts = find_all::(&ctx, &body); - - assert_eq!( - insert_indices(&ctx, &inserts), - vec![vec![0]], - "unit enum construction should insert only the discriminant tag" - ); - let tag_insert = &inserts[0]; - assert!( - tag_insert.verify(&ctx).is_ok(), - "the enum tag insertion must satisfy LLVM dialect verification" - ); - let tag = tag_insert.get_operation().deref(&ctx).get_operand(1); - let tag_def = tag - .defining_op() - .expect("the inserted enum tag must have a defining operation"); - let tag_constant = Operation::get_op::(tag_def, &ctx) - .expect("the inserted enum tag must be defined by llvm.constant"); - let tag_attr = tag_constant.get_value(&ctx); - let tag_integer = tag_attr - .downcast_ref::() - .expect("the inserted enum tag must be an integer constant"); - assert_eq!(tag_integer.value().bw(), 8, "the enum tag must be 8-bit"); - assert_eq!( - tag_integer.value().to_u64(), - 255, - "Less must lower to its declared i8 bit-pattern 255, not variant index 0" - ); - } - - #[test] - fn nested_pointer_niche_tuple_construct_extract_and_discriminant_lower() { - let mut ctx = make_ctx(); - let logical: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signed).into(); - let index: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); - let pointee: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); - let pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, pointee, false).into(); - let tuple_ty: TypeHandle = MirTupleType::get(&mut ctx, vec![index, pointer]).into(); - let enum_ty: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "Option".into(), - logical, - vec![0, 1], - vec![ - EnumVariant::unit("None".into()), - EnumVariant::new_with_layout("Some".into(), vec![tuple_ty], vec![0], vec![16]), - ], - EnumEncoding { - tag_offset: 8, - total_size: 16, - abi_align: 8, - layout_kind: EnumLayoutKind::Niche, - carrier_kind: EnumCarrierKind::Pointer, - carrier_width: 64, - untagged_variant: 1, - variant_inhabited: vec![1, 1], - ..EnumEncoding::default() - }, - ) - .into(); - let slot_map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); - assert_eq!(slot_map.carrier_slot, Some(1)); - assert_eq!(slot_map.field_slots, vec![None]); - let lowered_tuple = convert_type(&mut ctx, tuple_ty).unwrap(); - - let (module, block) = build_kernel(&mut ctx, vec![index, pointer], vec![]); - let index_value = block.deref(&ctx).get_argument(0); - let pointer_value = block.deref(&ctx).get_argument(1); - let tuple = Operation::new( - &mut ctx, - mir::MirConstructTupleOp::get_concrete_op_info(), - vec![tuple_ty], - vec![index_value, pointer_value], - vec![], - 0, - ); - tuple.insert_at_back(block, &ctx); - let tuple_value = tuple.deref(&ctx).get_result(0); - - let construct = Operation::new( - &mut ctx, - mir::MirConstructEnumOp::get_concrete_op_info(), - vec![enum_ty], - vec![tuple_value], - vec![], - 0, - ); - mir::MirConstructEnumOp::new(construct) - .set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(1)); - construct.insert_at_back(block, &ctx); - let enum_value = construct.deref(&ctx).get_result(0); - - let payload = Operation::new( - &mut ctx, - mir::MirEnumPayloadOp::get_concrete_op_info(), - vec![tuple_ty], - vec![enum_value], - vec![], - 0, - ); - mir::MirEnumPayloadOp::new(payload) - .set_attr_payload_variant_index(&ctx, VariantIndexAttr(1)); - mir::MirEnumPayloadOp::new(payload).set_attr_payload_field_index(&ctx, FieldIndexAttr(0)); - payload.insert_at_back(block, &ctx); - - let discriminant = Operation::new( - &mut ctx, - mir::MirGetDiscriminantOp::get_concrete_op_info(), - vec![logical], - vec![enum_value], - vec![], - 0, - ); - discriminant.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); - let body = kernel_blocks(&ctx, module); - assert_eq!( - count_ops::(&ctx, &body), - 0, - "constructing the untagged Some payload must not recreate its pointer from bits" - ); - assert_eq!( - count_ops::(&ctx, &body), - 1, - "reading the pointer niche should inspect the carrier exactly once" - ); - assert_eq!( - count_ops::(&ctx, &body), - 3, - "construction and extraction should each spill the enum, plus one tuple payload store" - ); - assert_eq!( - count_ops::(&ctx, &body), - 2, - "construction should reload the enum and extraction should load the tuple payload" - ); - // What this test is about is that the payload moves as one unit, never - // field by field. Assert that property directly instead of counting - // whole-aggregate accesses. - // - // Counting was the fragile form. The enum's physical storage here is - // `{i64, ptr}` -- the *same interned type* as the lowered payload tuple, - // since the niche carrier claims the pointer at byte 8 and the 8 bytes - // below it become one `i64` filler. So a count of `lowered_tuple`-typed - // accesses cannot separate the payload write from the enum spill, and - // the expected numbers move whenever that coincidence appears or - // disappears -- which has nothing to do with the property under test. - // - // A lowering that decomposed the payload is recognisable by what it - // emits instead: traffic in the tuple's *field* types. Look for that, - // and the assertion holds whatever the enum storage type happens to be. - let field_tys: Vec = lowered_tuple - .deref(&ctx) - .downcast_ref::() - .expect("the lowered tuple is an LLVM struct") - .fields() - .collect(); - - let store_tys: Vec = find_all::(&ctx, &body) - .iter() - .map(|store| store.get_operand_value(&ctx).get_type(&ctx)) - .collect(); - let load_tys: Vec = find_all::(&ctx, &body) - .iter() - .map(|load| { - load.get_operation() - .deref(&ctx) - .get_result(0) - .get_type(&ctx) - }) - .collect(); - - let describe = |tys: &[TypeHandle]| -> Vec { - tys.iter() - .filter(|ty| field_tys.contains(ty)) - .map(|ty| ty.deref(&ctx).disp(&ctx).to_string()) - .collect() - }; - assert!( - describe(&store_tys).is_empty(), - "the payload must be stored whole, but these field-typed stores appear: {:?}", - describe(&store_tys) - ); - assert!( - describe(&load_tys).is_empty(), - "the payload must be read whole, but these field-typed loads appear: {:?}", - describe(&load_tys) - ); - - // And it must actually be moved: without this the checks above would - // also pass a lowering that emitted no payload traffic at all. - assert!( - store_tys.contains(&lowered_tuple), - "at least one store must move the complete {{i64, ptr}} payload" - ); - assert!( - load_tys.contains(&lowered_tuple), - "at least one load must read the complete {{i64, ptr}} payload" - ); - } - - /// SetDiscriminant must use the slot map instead of assuming that the tag - /// is field zero, and its GEP must retain the source pointer's GPU address - /// space. This shape puts an i64 payload first and an i8 tag above the - /// u32 range, proving the byte GEP does not truncate offsets to u32. - #[test] - fn set_discriminant_uses_tag_slot_and_preserves_shared_address_space() { - use llvm_export::ops::GepIndex; - use llvm_export::types::{PointerType, address_space}; - - let mut ctx = make_ctx(); - let discr_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); - let payload_a: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); - let payload_b: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); - let tag_offset = u64::from(u32::MAX) + 1; - let enum_ty: TypeHandle = MirEnumType::get_with_layout( - &mut ctx, - "TagAfterPayload".to_string(), - discr_ty, - vec![3, 7], - vec![ - EnumVariant::new_with_layout("A".to_string(), vec![payload_a], vec![0], vec![8]), - EnumVariant::new_with_layout("B".to_string(), vec![payload_b], vec![0], vec![8]), - ], - tag_offset, - tag_offset + 8, - 8, - ) - .into(); - let ptr_ty: TypeHandle = MirPtrType::get_shared(&mut ctx, enum_ty, true).into(); - - let (module_ptr, block) = build_kernel(&mut ctx, vec![ptr_ty], vec![]); - let enum_ptr = block.deref(&ctx).get_argument(0); - let set = Operation::new( - &mut ctx, - mir::MirSetDiscriminantOp::get_concrete_op_info(), - vec![], - vec![enum_ptr], - vec![], - 0, - ); - let set_op = mir::MirSetDiscriminantOp::new(set); - set_op.set_attr_set_discriminant_variant_index(&ctx, VariantIndexAttr(1)); - set_op.set_attr_set_discriminant_enum_ty( - &ctx, - pliron::builtin::attributes::TypeAttr::new(enum_ty), - ); - set.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - assert_eq!(count_ops::(&ctx, &body), 0); - assert_eq!(count_ops::(&ctx, &body), 1); - assert_eq!(count_ops::(&ctx, &body), 1); - - let gep = find_first::(&ctx, &body).unwrap(); - let indices = gep.indices(&ctx); - assert!(matches!(indices.as_slice(), [GepIndex::Value(_)])); - let GepIndex::Value(offset) = indices[0] else { - unreachable!() - }; - let offset_def = offset.defining_op().expect("byte offset must be constant"); - let offset_constant = Operation::get_op::(offset_def, &ctx) - .expect("byte offset must be an LLVM constant"); - let offset_attr = offset_constant.get_value(&ctx); - assert_eq!( - offset_attr - .downcast_ref::() - .expect("byte offset must be integer") - .value() - .to_u64(), - tag_offset, - "the write must use rustc's absolute carrier byte offset" - ); - let gep_result_ty = gep.get_operation().deref(&ctx).get_result(0).get_type(&ctx); - assert_eq!( - gep_result_ty - .deref(&ctx) - .downcast_ref::() - .expect("GEP result must be a pointer") - .address_space(), - address_space::SHARED, - "tag GEP must preserve shared address space" - ); - - let store = find_first::(&ctx, &body).unwrap(); - let stored_ty = store.get_operand_value(&ctx).get_type(&ctx); - assert_eq!( - stored_ty - .deref(&ctx) - .downcast_ref::() - .expect("stored tag must be an integer") - .width(), - 8 - ); - assert_eq!( - store.get_operand_address(&ctx), - gep.get_operation().deref(&ctx).get_result(0), - "the store must use the tag GEP result" - ); - } - - fn unit_niche_enum( - ctx: &mut Context, - carrier: (EnumCarrierKind, u32, u32), - niche_start: u128, - niche_range: std::ops::RangeInclusive, - untagged_variant: u32, - inhabited: Vec, - ) -> TypeHandle { - let (carrier_kind, carrier_width, carrier_address_space) = carrier; - let logical_width = if inhabited.len() > u8::MAX as usize { - 16 - } else { - 8 - }; - let logical_ty: TypeHandle = - IntegerType::get(ctx, logical_width, Signedness::Unsigned).into(); - let variants = (0..inhabited.len()) - .map(|index| EnumVariant::unit(format!("V{index}"))) - .collect::>(); - let discriminants = (0..inhabited.len() as u64).collect(); - let carrier_size = u64::from(carrier_width).div_ceil(8); - let carrier_align = carrier_size.next_power_of_two().min(16); - MirEnumType::get_with_encoding( - ctx, - "UnitNiche".into(), - logical_ty, - discriminants, - variants, - EnumEncoding { - tag_offset: 0, - total_size: carrier_size, - abi_align: carrier_align, - layout_kind: EnumLayoutKind::Niche, - carrier_kind, - carrier_width, - carrier_address_space, - niche_start, - niche_variant_start: *niche_range.start(), - niche_variant_end: *niche_range.end(), - untagged_variant, - variant_inhabited: inhabited, - ..EnumEncoding::default() - }, - ) - .into() - } - - fn over_aligned_tuple_ty(ctx: &mut Context) -> TypeHandle { - let byte: TypeHandle = IntegerType::get(ctx, 8, Signedness::Unsigned).into(); - let marker: TypeHandle = MirStructType::get_with_full_layout( - ctx, - "Align32".into(), - vec![], - vec![], - vec![], - vec![], - 0, - 32, - ) - .into(); - MirTupleType::get_with_layout(ctx, vec![marker, byte], vec![0, 1], vec![0, 0], 32, 32) - .into() - } - - fn lower_array_extract_case( - ctx: &mut Context, - array_size: u64, - divisor: Option, - ) -> Ptr { - let element_type: TypeHandle = IntegerType::get(ctx, 32, Signedness::Unsigned).into(); - let index_type = IntegerType::get(ctx, 64, Signedness::Unsigned); - let index_handle: TypeHandle = index_type.into(); - let array_type: TypeHandle = MirArrayType::get(ctx, element_type, array_size).into(); - - let (module, block) = build_kernel(ctx, vec![index_handle], vec![element_type]); - let raw_index = block.deref(ctx).get_argument(0); - - let undef = mir::MirUndefOp::new(ctx, array_type); - undef.get_operation().insert_at_back(block, ctx); - let array = undef.get_operation().deref(ctx).get_result(0); - - let index = if let Some(divisor) = divisor { - let constant = Operation::new( - ctx, - mir::MirConstantOp::get_concrete_op_info(), - vec![index_handle], - vec![], - vec![], - 0, - ); - mir::MirConstantOp::new(constant).set_attr_value( - ctx, - IntegerAttr::new( - index_type, - APInt::from_u64(divisor, NonZeroUsize::new(64).unwrap()), - ), - ); - constant.insert_at_back(block, ctx); - let divisor_value = constant.deref(ctx).get_result(0); - - let rem = Operation::new( - ctx, - mir::MirRemOp::get_concrete_op_info(), - vec![index_handle], - vec![raw_index, divisor_value], - vec![], - 0, - ); - rem.insert_at_back(block, ctx); - rem.deref(ctx).get_result(0) - } else { - raw_index - }; - - let extract = Operation::new( - ctx, - mir::MirExtractArrayElementOp::get_concrete_op_info(), - vec![element_type], - vec![array, index], - vec![], - 0, - ); - extract.insert_at_back(block, ctx); - let result = extract.deref(ctx).get_result(0); - append_mir_return(ctx, block, vec![result]); - - crate::lower_mir_to_llvm(ctx, module).expect("lowering failed"); - module - } - - fn assert_array_extract_memory_fallback(ctx: &Context, module: Ptr) { - let body = kernel_blocks(ctx, module); - assert_eq!(count_ops::(ctx, &body), 1); - assert_eq!(count_ops::(ctx, &body), 1); - assert_eq!(count_ops::(ctx, &body), 1); - assert_eq!(count_ops::(ctx, &body), 1); - assert_eq!(count_ops::(ctx, &body), 0); - } - - #[test] - fn bounded_urem_array_extract_stays_in_ssa() { - let mut ctx = make_ctx(); - let module = lower_array_extract_case(&mut ctx, 3, Some(3)); - let body = kernel_blocks(&ctx, module); - - assert_eq!(count_ops::(&ctx, &body), 3); - assert_eq!(count_ops::(&ctx, &body), 2); - assert_eq!(count_ops::(&ctx, &body), 2); - assert_eq!(count_ops::(&ctx, &body), 0); - assert_eq!(count_ops::(&ctx, &body), 0); - assert_eq!(count_ops::(&ctx, &body), 0); - assert_eq!(count_ops::(&ctx, &body), 0); - } - - #[test] - fn unbounded_array_extract_keeps_memory_fallback() { - let mut ctx = make_ctx(); - let module = lower_array_extract_case(&mut ctx, 3, None); - assert_array_extract_memory_fallback(&ctx, module); - } - - #[test] - fn oversized_urem_array_extract_keeps_memory_fallback() { - let mut ctx = make_ctx(); - let module = lower_array_extract_case(&mut ctx, 17, Some(17)); - assert_array_extract_memory_fallback(&ctx, module); - } - - #[test] - fn dynamic_array_extract_preserves_recursive_element_alignment() { - let mut ctx = make_ctx(); - let tuple_ty = over_aligned_tuple_ty(&mut ctx); - let inner: TypeHandle = MirArrayType::get(&mut ctx, tuple_ty, 2).into(); - let outer: TypeHandle = MirArrayType::get(&mut ctx, inner, 3).into(); - let index_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); - let (module_ptr, block) = build_kernel(&mut ctx, vec![index_ty], vec![]); - let index = block.deref(&ctx).get_argument(0); - - let undef = mir::MirUndefOp::new(&mut ctx, outer); - undef.get_operation().insert_at_back(block, &ctx); - let array = undef.get_operation().deref(&ctx).get_result(0); - let extract = Operation::new( - &mut ctx, - mir::MirExtractArrayElementOp::get_concrete_op_info(), - vec![inner], - vec![array, index], - vec![], - 0, - ); - extract.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - let alloca = find_first::(&ctx, &body).expect("expected llvm.alloca"); - let store = find_first::(&ctx, &body).expect("expected llvm.store"); - let load = find_first::(&ctx, &body).expect("expected llvm.load"); - for memory_op in [ - alloca.get_operation(), - store.get_operation(), - load.get_operation(), - ] { - assert_eq!(llvm_export::ops::op_alignment(&ctx, memory_op), Some(32)); - } - } - - #[test] - fn niche_encoding_handles_untagged_inside_range_and_u128_wrap() { - let mut ctx = make_ctx(); - let inside = unit_niche_enum( - &mut ctx, - (EnumCarrierKind::Integer, 8, 0), - 42, - 0..=1, - 1, - vec![1, 1], - ); - let inside_ref = inside.deref(&ctx); - let inside_enum = inside_ref.downcast_ref::().unwrap(); - assert_eq!(enum_carrier_bits_for_variant(inside_enum, 0), Ok(Some(42))); - assert_eq!( - enum_carrier_bits_for_variant(inside_enum, 1), - Ok(None), - "the untagged variant is a no-op even when its index is in the niche range" - ); - drop(inside_ref); - - let wrapping = unit_niche_enum( - &mut ctx, - (EnumCarrierKind::Integer, 128, 0), - u128::MAX, - 0..=1, - 0, - vec![1, 1], - ); - let wrapping_ref = wrapping.deref(&ctx); - let wrapping_enum = wrapping_ref.downcast_ref::().unwrap(); - assert_eq!( - enum_carrier_bits_for_variant(wrapping_enum, 1), - Ok(Some(0)), - "niche arithmetic must wrap across the full u128 carrier" - ); - } - - #[test] - fn set_discriminant_niche_writes_only_tagged_variant_carrier() { - for (target, expected_stores) in [(0, 1), (1, 0)] { - let mut ctx = make_ctx(); - let enum_ty = unit_niche_enum( - &mut ctx, - (EnumCarrierKind::Integer, 8, 0), - 0, - 0..=0, - 1, - vec![1, 1], - ); - let ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, enum_ty, true).into(); - let (module, block) = build_kernel(&mut ctx, vec![ptr_ty], vec![]); - let ptr = block.deref(&ctx).get_argument(0); - let set = Operation::new( - &mut ctx, - mir::MirSetDiscriminantOp::get_concrete_op_info(), - vec![], - vec![ptr], - vec![], - 0, - ); - let set_op = mir::MirSetDiscriminantOp::new(set); - set_op.set_attr_set_discriminant_variant_index(&ctx, VariantIndexAttr(target)); - set_op.set_attr_set_discriminant_enum_ty( - &ctx, - pliron::builtin::attributes::TypeAttr::new(enum_ty), - ); - set.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); - let body = kernel_blocks(&ctx, module); - assert_eq!(count_ops::(&ctx, &body), expected_stores); - assert_eq!(count_ops::(&ctx, &body), 0); - } - } - - #[test] - fn shared_pointer_niche_carrier_rejects_target_dependent_width() { - let mut ctx = make_ctx(); - let enum_ty = unit_niche_enum( - &mut ctx, - (EnumCarrierKind::Pointer, 64, 3), - 0, - 0..=0, - 1, - vec![1, 1], - ); - let error = build_enum_slot_map(&mut ctx, enum_ty) - .err() - .expect("shared pointer carrier must reject"); - assert!( - error.to_string().contains("target-mode dependent"), - "{error}" - ); - } - - #[test] - fn shared_pointer_niche_payload_round_trips_through_generic_storage() { - let mut ctx = make_ctx(); - let logical: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signed).into(); - let pointee: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); - let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, true).into(); - let enum_ty: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "OptionShared".into(), - logical, - vec![0, 1], - vec![ - EnumVariant::unit("None".into()), - EnumVariant::new_with_layout("Some".into(), vec![shared], vec![0], vec![8]), - ], - EnumEncoding { - tag_offset: 0, - total_size: 8, - abi_align: 8, - layout_kind: EnumLayoutKind::Niche, - carrier_kind: EnumCarrierKind::Pointer, - carrier_width: 64, - carrier_address_space: llvm_types::address_space::GENERIC, - niche_start: 0, - niche_variant_start: 0, - niche_variant_end: 0, - untagged_variant: 1, - variant_inhabited: vec![1, 1], - ..EnumEncoding::default() - }, - ) - .into(); - - let (module, block) = build_kernel(&mut ctx, vec![shared], vec![shared]); - let pointer = block.deref(&ctx).get_argument(0); - let construct = Operation::new( - &mut ctx, - MirConstructEnumOp::get_concrete_op_info(), - vec![enum_ty], - vec![pointer], - vec![], - 0, - ); - MirConstructEnumOp::new(construct) - .set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(1)); - construct.insert_at_back(block, &ctx); - let option = construct.deref(&ctx).get_result(0); - - let discriminant = Operation::new( - &mut ctx, - mir::MirGetDiscriminantOp::get_concrete_op_info(), - vec![logical], - vec![option], - vec![], - 0, - ); - discriminant.insert_at_back(block, &ctx); - - let payload = Operation::new( - &mut ctx, - MirEnumPayloadOp::get_concrete_op_info(), - vec![shared], - vec![option], - vec![], - 0, - ); - let payload_op = MirEnumPayloadOp::new(payload); - payload_op.set_attr_payload_variant_index(&ctx, VariantIndexAttr(1)); - payload_op.set_attr_payload_field_index(&ctx, FieldIndexAttr(0)); - payload.insert_at_back(block, &ctx); - let result = payload.deref(&ctx).get_result(0); - append_mir_return(&mut ctx, block, vec![result]); - - crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); - let body = kernel_blocks(&ctx, module); - assert_eq!( - count_ops::(&ctx, &body), - 2, - "construction must genericize the pointer and extraction must restore shared space" - ); - assert_eq!( - count_ops::(&ctx, &body), - 1, - "niche discrimination must inspect the generic pointer carrier" - ); - assert_eq!(count_ops::(&ctx, &body), 0); - assert_eq!(count_ops::(&ctx, &body), 0); - } - - #[test] - fn nested_shared_pointer_payload_round_trips_through_recursive_storage() { - let mut ctx = make_ctx(); - let logical: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); - let pointee: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); - let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, true).into(); - let inner: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "SharedPointerInner".into(), - vec!["pointer".into(), "cookie".into()], - vec![shared, logical], - vec![0, 1], - vec![0, 8], - 16, - 8, - ) - .into(); - let outer: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "SharedPointerOuter".into(), - vec!["inner".into(), "guard".into()], - vec![inner, logical], - vec![0, 1], - vec![0, 16], - 24, - 8, - ) - .into(); - let enum_ty: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "NestedSharedPointer".into(), - logical, - vec![0, 1], - vec![ - EnumVariant::unit("None".into()), - EnumVariant::new_with_layout("Some".into(), vec![outer], vec![8], vec![24]), - ], - EnumEncoding { - tag_offset: 0, - total_size: 32, - abi_align: 8, - layout_kind: EnumLayoutKind::Direct, - carrier_kind: EnumCarrierKind::Integer, - carrier_width: 32, - variant_inhabited: vec![1, 1], - ..EnumEncoding::default() - }, - ) - .into(); - - let (module, block) = build_kernel(&mut ctx, vec![outer], vec![outer]); - let wrapper = block.deref(&ctx).get_argument(0); - let construct = Operation::new( - &mut ctx, - MirConstructEnumOp::get_concrete_op_info(), - vec![enum_ty], - vec![wrapper], - vec![], - 0, - ); - MirConstructEnumOp::new(construct) - .set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(1)); - construct.insert_at_back(block, &ctx); - let option = construct.deref(&ctx).get_result(0); - - let payload = Operation::new( - &mut ctx, - MirEnumPayloadOp::get_concrete_op_info(), - vec![outer], - vec![option], - vec![], - 0, - ); - let payload_op = MirEnumPayloadOp::new(payload); - payload_op.set_attr_payload_variant_index(&ctx, VariantIndexAttr(1)); - payload_op.set_attr_payload_field_index(&ctx, FieldIndexAttr(0)); - payload.insert_at_back(block, &ctx); - let result = payload.deref(&ctx).get_result(0); - append_mir_return(&mut ctx, block, vec![result]); - - crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); - let body = kernel_blocks(&ctx, module); - assert_eq!( - count_ops::(&ctx, &body), - 2, - "construction and extraction must cast the nested shared pointer leaf" - ); - assert_eq!(count_ops::(&ctx, &body), 0); - assert_eq!(count_ops::(&ctx, &body), 0); - } - - #[test] - fn bounded_shared_pointer_array_niche_payload_round_trips_through_recursive_storage() { - let mut ctx = make_ctx(); - let logical: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signed).into(); - let pointee: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); - let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, true).into(); - let pointers: TypeHandle = MirArrayType::get(&mut ctx, shared, 2).into(); - let enum_ty: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "OptionSharedPointerArray".into(), - logical, - vec![0, 1], - vec![ - EnumVariant::unit("None".into()), - EnumVariant::new_with_layout("Some".into(), vec![pointers], vec![0], vec![16]), - ], - EnumEncoding { - tag_offset: 0, - total_size: 16, - abi_align: 8, - layout_kind: EnumLayoutKind::Niche, - carrier_kind: EnumCarrierKind::Pointer, - carrier_width: 64, - carrier_address_space: llvm_types::address_space::GENERIC, - niche_start: 0, - niche_variant_start: 0, - niche_variant_end: 0, - untagged_variant: 1, - variant_inhabited: vec![1, 1], - ..EnumEncoding::default() - }, - ) - .into(); - - let (module, block) = build_kernel(&mut ctx, vec![pointers], vec![pointers]); - let array = block.deref(&ctx).get_argument(0); - let construct = Operation::new( - &mut ctx, - MirConstructEnumOp::get_concrete_op_info(), - vec![enum_ty], - vec![array], - vec![], - 0, - ); - MirConstructEnumOp::new(construct) - .set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(1)); - construct.insert_at_back(block, &ctx); - let option = construct.deref(&ctx).get_result(0); - - let discriminant = Operation::new( - &mut ctx, - mir::MirGetDiscriminantOp::get_concrete_op_info(), - vec![logical], - vec![option], - vec![], - 0, - ); - discriminant.insert_at_back(block, &ctx); - - let payload = Operation::new( - &mut ctx, - MirEnumPayloadOp::get_concrete_op_info(), - vec![pointers], - vec![option], - vec![], - 0, - ); - let payload_op = MirEnumPayloadOp::new(payload); - payload_op.set_attr_payload_variant_index(&ctx, VariantIndexAttr(1)); - payload_op.set_attr_payload_field_index(&ctx, FieldIndexAttr(0)); - payload.insert_at_back(block, &ctx); - let result = payload.deref(&ctx).get_result(0); - append_mir_return(&mut ctx, block, vec![result]); - - crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); - let body = kernel_blocks(&ctx, module); - assert_eq!( - count_ops::(&ctx, &body), - 4, - "construction and extraction must cast both shared-pointer array elements" - ); - assert_eq!( - count_ops::(&ctx, &body), - 1, - "niche discrimination must inspect the generic first-pointer carrier" - ); - assert_eq!(count_ops::(&ctx, &body), 0); - assert_eq!(count_ops::(&ctx, &body), 0); - } - - #[test] - fn option_bool_uses_i8_carrier_and_spills_i1_payload() { - let mut ctx = make_ctx(); - let logical: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); - let bool_ty: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Signless).into(); - let enum_ty: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "OptionBool".into(), - logical, - vec![0, 1], - vec![ - EnumVariant::unit("None".into()), - EnumVariant::new_with_layout("Some".into(), vec![bool_ty], vec![0], vec![1]), - ], - EnumEncoding { - tag_offset: 0, - total_size: 1, - abi_align: 1, - layout_kind: EnumLayoutKind::Niche, - carrier_kind: EnumCarrierKind::Integer, - carrier_width: 8, - niche_start: 2, - untagged_variant: 1, - variant_inhabited: vec![1, 1], - ..EnumEncoding::default() - }, - ) - .into(); - let slot_map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); - assert_eq!( - slot_map - .carrier_llvm_ty - .unwrap() - .deref(&ctx) - .downcast_ref::() - .unwrap() - .width(), - 8, - "Option's memory carrier is i8, not bool's semantic i1" - ); - assert_eq!(slot_map.field_slots, vec![None]); - - let (module, block) = build_kernel(&mut ctx, vec![bool_ty], vec![enum_ty]); - let payload = block.deref(&ctx).get_argument(0); - let construct = Operation::new( - &mut ctx, - mir::MirConstructEnumOp::get_concrete_op_info(), - vec![enum_ty], - vec![payload], - vec![], - 0, - ); - mir::MirConstructEnumOp::new(construct) - .set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(1)); - construct.insert_at_back(block, &ctx); - let result = construct.deref(&ctx).get_result(0); - append_mir_return(&mut ctx, block, vec![result]); - - crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); - let body = kernel_blocks(&ctx, module); - assert_eq!(count_ops::(&ctx, &body), 0); - assert!(find_all::(&ctx, &body).iter().any(|zext| { - zext.get_operation() - .deref(&ctx) - .get_operand(0) - .get_type(&ctx) - .deref(&ctx) - .downcast_ref::() - .is_some_and(|integer| integer.width() == 1) - && zext - .get_operation() - .deref(&ctx) - .get_result(0) - .get_type(&ctx) - .deref(&ctx) - .downcast_ref::() - .is_some_and(|integer| integer.width() == 8) - })); - assert!(find_all::(&ctx, &body).iter().any(|store| { - store - .get_operand_value(&ctx) - .get_type(&ctx) - .deref(&ctx) - .downcast_ref::() - .is_some_and(|integer| integer.width() == 8) - })); - } - - #[test] - fn direct_bool_uses_i8_storage_and_spills_i1_payload() { - let mut ctx = make_ctx(); - let tag: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); - let bool_ty: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Signless).into(); - let enum_ty: TypeHandle = MirEnumType::get_with_layout( - &mut ctx, - "DirectBool".into(), - tag, - vec![0, 1], - vec![ - EnumVariant::new_with_layout("A".into(), vec![bool_ty], vec![4], vec![1]), - EnumVariant::unit("B".into()), - ], - 0, - 8, - 4, - ) - .into(); - let slot_map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); - assert_eq!(slot_map.field_slots, vec![None]); - let storage_fields = slot_map - .llvm_struct_ty - .deref(&ctx) - .downcast_ref::() - .expect("enum storage must be an LLVM struct") - .fields() - .collect::>(); - assert_eq!( - storage_fields[1] - .deref(&ctx) - .downcast_ref::() - .map(IntegerType::width), - Some(8), - "the standalone Rust bool byte must use physical i8 storage" - ); - - let (module, block) = build_kernel(&mut ctx, vec![bool_ty], vec![enum_ty]); - let payload = block.deref(&ctx).get_argument(0); - let construct = Operation::new( - &mut ctx, - mir::MirConstructEnumOp::get_concrete_op_info(), - vec![enum_ty], - vec![payload], - vec![], - 0, - ); - mir::MirConstructEnumOp::new(construct) - .set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(0)); - construct.insert_at_back(block, &ctx); - let result = construct.deref(&ctx).get_result(0); - append_mir_return(&mut ctx, block, vec![result]); - - crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); - let body = kernel_blocks(&ctx, module); - assert_eq!( - count_ops::(&ctx, &body), - 1, - "only the direct tag should be inserted as an SSA struct field" - ); - assert!(find_all::(&ctx, &body).iter().any(|zext| { - zext.get_operation() - .deref(&ctx) - .get_operand(0) - .get_type(&ctx) - .deref(&ctx) - .downcast_ref::() - .is_some_and(|integer| integer.width() == 1) - && zext - .get_operation() - .deref(&ctx) - .get_result(0) - .get_type(&ctx) - .deref(&ctx) - .downcast_ref::() - .is_some_and(|integer| integer.width() == 8) - })); - assert!(find_all::(&ctx, &body).iter().any(|store| { - store - .get_operand_value(&ctx) - .get_type(&ctx) - .deref(&ctx) - .downcast_ref::() - .is_some_and(|integer| integer.width() == 8) - })); - } - - /// Taking the in-place address of a bool payload must fail loudly: the - /// payload's canonical storage is an i8 byte, and a semantic i1 store - /// through the escaped address would leave that byte's upper seven bits - /// undefined for every i8 reader (including a niche tag sharing them). - #[test] - fn bool_payload_field_addr_fails_closed() { - let mut ctx = make_ctx(); - let tag: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); - let bool_ty: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Signless).into(); - let enum_ty: TypeHandle = MirEnumType::get_with_layout( - &mut ctx, - "DirectBool".into(), - tag, - vec![0, 1], - vec![ - EnumVariant::new_with_layout("A".into(), vec![bool_ty], vec![4], vec![1]), - EnumVariant::unit("B".into()), - ], - 0, - 8, - 4, - ) - .into(); - let enum_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, enum_ty, true).into(); - let bool_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, bool_ty, true).into(); - - let (module, block) = build_kernel(&mut ctx, vec![enum_ptr_ty], vec![]); - let base = block.deref(&ctx).get_argument(0); - let op = MirFieldAddrOp::build(&mut ctx, base, bool_ptr_ty, 0).expect("field_addr build"); - op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - let err = crate::lower_mir_to_llvm(&mut ctx, module) - .expect_err("addressing a bool enum payload in place must fail to lower"); - assert!( - err.err.to_string().contains("canonical storage type"), - "unexpected error: {}", - err.err - ); - } - - /// Same gate for the slot arm: a shared-memory pointer payload is stored - /// as a GENERIC pointer (its slot is typed `ptr`), so an in-place address - /// would let a semantic `ptr addrspace(3)` store write a representation - /// every other reader interprets as a generic pointer. - #[test] - fn shared_pointer_payload_field_addr_fails_closed() { - let mut ctx = make_ctx(); - let logical: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signed).into(); - let pointee: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); - let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, true).into(); - let enum_ty: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "OptionShared".into(), - logical, - vec![0, 1], - vec![ - EnumVariant::unit("None".into()), - EnumVariant::new_with_layout("Some".into(), vec![shared], vec![0], vec![8]), - ], - EnumEncoding { - tag_offset: 0, - total_size: 8, - abi_align: 8, - layout_kind: EnumLayoutKind::Niche, - carrier_kind: EnumCarrierKind::Pointer, - carrier_width: 64, - carrier_address_space: llvm_types::address_space::GENERIC, - niche_start: 0, - niche_variant_start: 0, - niche_variant_end: 0, - untagged_variant: 1, - variant_inhabited: vec![1, 1], - ..EnumEncoding::default() - }, - ) - .into(); - let enum_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, enum_ty, true).into(); - let payload_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, shared, true).into(); - - let (module, block) = build_kernel(&mut ctx, vec![enum_ptr_ty], vec![]); - let base = block.deref(&ctx).get_argument(0); - let op = - MirFieldAddrOp::build(&mut ctx, base, payload_ptr_ty, 0).expect("field_addr build"); - op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - let err = crate::lower_mir_to_llvm(&mut ctx, module) - .expect_err("addressing a shared-pointer enum payload in place must fail to lower"); - assert!( - err.err.to_string().contains("canonical storage type"), - "unexpected error: {}", - err.err - ); - } - - /// A payload with no slot of its own is addressed at its byte offset off - /// the ORIGINAL enum pointer: no stack spill is introduced, so a write - /// through the result lands in the enum rather than in a copy. - #[test] - fn slotless_payload_field_addr_geps_original_storage_at_byte_offset() { - use llvm_export::ops::GepIndex; - use pliron::builtin::types::FP32Type; - - let mut ctx = make_ctx(); - let tag: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); - let f32_ty: TypeHandle = FP32Type::get(&ctx).into(); - let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); - let enum_ty: TypeHandle = MirEnumType::get_with_layout( - &mut ctx, - "Either".into(), - tag, - vec![0, 1], - vec![ - EnumVariant::new_with_layout("Real".into(), vec![f32_ty], vec![4], vec![4]), - EnumVariant::new_with_layout("Bits".into(), vec![u32_ty], vec![4], vec![4]), - ], - 0, - 8, - 4, - ) - .into(); - let slot_map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); - assert_eq!( - slot_map.field_slots, - vec![Some(1), None], - "Real's f32 claims byte 4 first; Bits shares those bytes slotless" - ); - - let enum_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, enum_ty, true).into(); - let u32_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, u32_ty, true).into(); - let (module, block) = build_kernel(&mut ctx, vec![enum_ptr_ty], vec![]); - let base = block.deref(&ctx).get_argument(0); - let op = MirFieldAddrOp::build(&mut ctx, base, u32_ptr_ty, 1).expect("field_addr build"); - op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); - let body = kernel_blocks(&ctx, module); - assert_eq!( - count_ops::(&ctx, &body), - 0, - "an in-place payload address must not spill the enum to a stack copy" - ); - let geps = find_all::(&ctx, &body); - assert_eq!(geps.len(), 1, "one field_addr lowers to one GEP"); - let gep = &geps[0]; - // The MIR entry block and its arguments are the ORIGINALS (moved by - // `inline_region`), so the enum pointer argument keeps its identity - // through lowering and the GEP must be based directly on it. - assert_eq!( - gep.get_operation().deref(&ctx).get_operand(0), - base, - "the byte-offset GEP must address the ORIGINAL enum storage" - ); - assert!( - matches!(gep.indices(&ctx).as_slice(), [GepIndex::Constant(4)]), - "the slotless payload must be addressed at its rustc byte offset" - ); - assert_eq!( - gep.src_elem_type(&ctx), - IntegerType::get(&ctx, 8, Signedness::Signless).into(), - "byte addressing must step in i8 units" - ); - } - - /// The capability split, read side: a payload whose storage IS its - /// semantic type keeps the address path even for a SHARED borrow. The - /// borrow lowers to one GEP into the ORIGINAL enum storage and the read - /// to one load through it: no stack spill, no value copy. (Non-canonical - /// payloads never get here for shared borrows; the importer punts them - /// to the value-copy fallback before an address is formed.) - #[test] - fn canonical_payload_shared_read_loads_through_gep_without_copy() { - use llvm_export::ops::GepIndex; - use pliron::builtin::types::FP32Type; - - let mut ctx = make_ctx(); - let tag: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); - let f32_ty: TypeHandle = FP32Type::get(&ctx).into(); - let enum_ty: TypeHandle = MirEnumType::get_with_layout( - &mut ctx, - "Slot".into(), - tag, - vec![0, 1], - vec![ - EnumVariant::new_with_layout("Occupied".into(), vec![f32_ty], vec![4], vec![4]), - EnumVariant::unit("Empty".into()), - ], - 0, - 8, - 4, - ) - .into(); - let slot_map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); - assert_eq!( - slot_map.field_slots, - vec![Some(1)], - "an f32 payload owns its LLVM slot; storage equals semantic type" - ); - - // Immutable pointer types model the shared borrow. - let enum_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, enum_ty, false).into(); - let f32_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, f32_ty, false).into(); - let (module, block) = build_kernel(&mut ctx, vec![enum_ptr_ty], vec![f32_ty]); - let base = block.deref(&ctx).get_argument(0); - let addr = MirFieldAddrOp::build(&mut ctx, base, f32_ptr_ty, 0).expect("field_addr build"); - addr.insert_at_back(block, &ctx); - let payload_ptr = addr.deref(&ctx).get_result(0); - - let load = Operation::new( - &mut ctx, - mir::MirLoadOp::get_concrete_op_info(), - vec![f32_ty], - vec![payload_ptr], - vec![], - 0, - ); - load.insert_at_back(block, &ctx); - let loaded = load.deref(&ctx).get_result(0); - append_mir_return(&mut ctx, block, vec![loaded]); - - crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); - let body = kernel_blocks(&ctx, module); - assert_eq!( - count_ops::(&ctx, &body), - 0, - "a canonical payload read must not spill or copy the enum" - ); - let geps = find_all::(&ctx, &body); - assert_eq!(geps.len(), 1, "one field_addr lowers to one GEP"); - let gep = &geps[0]; - assert_eq!( - gep.get_operation().deref(&ctx).get_operand(0), - base, - "the GEP must address the ORIGINAL enum storage" - ); - assert!( - matches!( - gep.indices(&ctx).as_slice(), - [GepIndex::Constant(0), GepIndex::Constant(1)] - ), - "an own-slot payload is addressed through its struct slot" - ); - let loads = find_all::(&ctx, &body); - assert_eq!( - loads.len(), - 1, - "the read is a single load through the payload address" - ); - let gep_result = gep.get_operation().deref(&ctx).get_result(0); - assert_eq!( - loads[0].get_operation().deref(&ctx).get_operand(0), - gep_result, - "the load must go through the GEP result, not a copy" - ); - assert_eq!( - loads[0] - .get_operation() - .deref(&ctx) - .get_result(0) - .get_type(&ctx), - f32_ty, - "the load reads the payload at its semantic type" - ); - } - - #[test] - fn later_field_niche_set_writes_exact_carrier_offset() { - use llvm_export::ops::GepIndex; - - let mut ctx = make_ctx(); - let logical: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); - let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); - let wrapper: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "Wrapper".into(), - vec!["pad".into(), "nz".into()], - vec![u32_ty, u32_ty], - vec![0, 1], - vec![0, 4], - 8, - 4, - ) - .into(); - let enum_ty: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "MaybeWrapper".into(), - logical, - vec![0, 1], - vec![ - EnumVariant::unit("None".into()), - EnumVariant::new_with_layout("Some".into(), vec![wrapper], vec![0], vec![8]), - ], - EnumEncoding { - tag_offset: 4, - total_size: 8, - abi_align: 4, - layout_kind: EnumLayoutKind::Niche, - carrier_kind: EnumCarrierKind::Integer, - carrier_width: 32, - untagged_variant: 1, - variant_inhabited: vec![1, 1], - ..EnumEncoding::default() - }, - ) - .into(); - let ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, enum_ty, true).into(); - let (module, block) = build_kernel(&mut ctx, vec![ptr_ty], vec![]); - let ptr = block.deref(&ctx).get_argument(0); - let set = Operation::new( - &mut ctx, - mir::MirSetDiscriminantOp::get_concrete_op_info(), - vec![], - vec![ptr], - vec![], - 0, - ); - let set_op = mir::MirSetDiscriminantOp::new(set); - set_op.set_attr_set_discriminant_variant_index(&ctx, VariantIndexAttr(0)); - set_op.set_attr_set_discriminant_enum_ty( - &ctx, - pliron::builtin::attributes::TypeAttr::new(enum_ty), - ); - set.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); - let body = kernel_blocks(&ctx, module); - let gep = find_first::(&ctx, &body).unwrap(); - let indices = gep.indices(&ctx); - let [GepIndex::Value(offset)] = indices.as_slice() else { - panic!("carrier access must use a byte-offset SSA value"); - }; - let constant = - Operation::get_op::(offset.defining_op().unwrap(), &ctx).unwrap(); - assert_eq!( - constant - .get_value(&ctx) - .downcast_ref::() - .unwrap() - .value() - .to_u64(), - 4 - ); - } - - #[test] - fn get_niche_discriminant_adds_large_range_start_at_logical_width() { - let mut ctx = make_ctx(); - let logical_ty: TypeHandle = IntegerType::get(&ctx, 16, Signedness::Unsigned).into(); - let carrier_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); - let mut variants = (0..300) - .map(|index| EnumVariant::unit(format!("V{index}"))) - .collect::>(); - // A valid niche layout stores the carrier inside the untagged - // variant's payload. Keep the large logical-variant test realistic - // instead of relying on a carrier with no backing field. - variants[0] = EnumVariant::new_with_layout("V0".into(), vec![carrier_ty], vec![0], vec![1]); - let enum_ty: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "LargeNiche".into(), - logical_ty, - (0..300).collect(), - variants, - EnumEncoding { - total_size: 1, - abi_align: 1, - layout_kind: EnumLayoutKind::Niche, - carrier_kind: EnumCarrierKind::Integer, - carrier_width: 8, - niche_variant_start: 298, - niche_variant_end: 299, - untagged_variant: 0, - variant_inhabited: { - let mut inhabited = vec![0; 300]; - inhabited[0] = 1; - inhabited[298] = 1; - inhabited[299] = 1; - inhabited - }, - ..EnumEncoding::default() - }, - ) - .into(); - let (module, block) = build_kernel(&mut ctx, vec![enum_ty], vec![logical_ty]); - let value = block.deref(&ctx).get_argument(0); - let get = Operation::new( - &mut ctx, - mir::MirGetDiscriminantOp::get_concrete_op_info(), - vec![logical_ty], - vec![value], - vec![], - 0, - ); - get.insert_at_back(block, &ctx); - let result = get.deref(&ctx).get_result(0); - append_mir_return(&mut ctx, block, vec![result]); - - crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); - let body = kernel_blocks(&ctx, module); - assert_eq!(count_ops::(&ctx, &body), 1); - let add = find_first::(&ctx, &body).expect("logical variant add"); - for operand in add.get_operation().deref(&ctx).operands() { - assert_eq!( - operand - .get_type(&ctx) - .deref(&ctx) - .downcast_ref::() - .unwrap() - .width(), - 16, - "variant-index arithmetic must occur at logical width" - ); - } - } - - #[test] - fn direct_negative_discriminant_read_remains_signed_for_widening() { - let mut ctx = make_ctx(); - let i8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Signed).into(); - let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signed).into(); - let enum_ty: TypeHandle = MirEnumType::get_with_layout( - &mut ctx, - "Negative".into(), - i8_ty, - vec![255, 0], - vec![EnumVariant::unit("N".into()), EnumVariant::unit("Z".into())], - 0, - 1, - 1, - ) - .into(); - let (module, block) = build_kernel(&mut ctx, vec![], vec![i32_ty]); - let construct = Operation::new( - &mut ctx, - mir::MirConstructEnumOp::get_concrete_op_info(), - vec![enum_ty], - vec![], - vec![], - 0, - ); - mir::MirConstructEnumOp::new(construct) - .set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(0)); - construct.insert_at_back(block, &ctx); - let enum_value = construct.deref(&ctx).get_result(0); - let get = Operation::new( - &mut ctx, - mir::MirGetDiscriminantOp::get_concrete_op_info(), - vec![i8_ty], - vec![enum_value], - vec![], - 0, - ); - get.insert_at_back(block, &ctx); - let discriminant = get.deref(&ctx).get_result(0); - let cast = Operation::new( - &mut ctx, - mir::MirCastOp::get_concrete_op_info(), - vec![i32_ty], - vec![discriminant], - vec![], - 0, - ); - mir::MirCastOp::new(cast).set_attr_cast_kind(&ctx, MirCastKindAttr::IntToInt); - cast.insert_at_back(block, &ctx); - let widened = cast.deref(&ctx).get_result(0); - append_mir_return(&mut ctx, block, vec![widened]); - - crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); - let body = kernel_blocks(&ctx, module); - assert_eq!(count_ops::(&ctx, &body), 1); - assert_eq!(count_ops::(&ctx, &body), 0); - } - - #[test] - fn single_layout_preserves_large_and_negative_declared_discriminants() { - for (width, signedness, bits, widened_width, expects_sext) in [ - (16, Signedness::Unsigned, 1_000, 16, false), - (8, Signedness::Signed, 251, 32, true), - ] { - let mut ctx = make_ctx(); - let logical: TypeHandle = IntegerType::get(&ctx, width, signedness).into(); - let destination: TypeHandle = IntegerType::get(&ctx, widened_width, signedness).into(); - let enum_ty: TypeHandle = MirEnumType::get_with_encoding( - &mut ctx, - "Single".into(), - logical, - vec![bits], - vec![EnumVariant::unit("Only".into())], - EnumEncoding { - tag_offset: 0, - total_size: 0, - abi_align: 1, - layout_kind: EnumLayoutKind::Single, - variant_inhabited: vec![1], - ..EnumEncoding::default() - }, - ) - .into(); - let (module, block) = build_kernel(&mut ctx, vec![], vec![destination]); - let construct = Operation::new( - &mut ctx, - mir::MirConstructEnumOp::get_concrete_op_info(), - vec![enum_ty], - vec![], - vec![], - 0, - ); - mir::MirConstructEnumOp::new(construct) - .set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(0)); - construct.insert_at_back(block, &ctx); - let enum_value = construct.deref(&ctx).get_result(0); - let get = Operation::new( - &mut ctx, - mir::MirGetDiscriminantOp::get_concrete_op_info(), - vec![logical], - vec![enum_value], - vec![], - 0, - ); - get.insert_at_back(block, &ctx); - let discr = get.deref(&ctx).get_result(0); - let returned = if widened_width == width { - discr - } else { - let cast = Operation::new( - &mut ctx, - mir::MirCastOp::get_concrete_op_info(), - vec![destination], - vec![discr], - vec![], - 0, - ); - mir::MirCastOp::new(cast).set_attr_cast_kind(&ctx, MirCastKindAttr::IntToInt); - cast.insert_at_back(block, &ctx); - cast.deref(&ctx).get_result(0) - }; - append_mir_return(&mut ctx, block, vec![returned]); - - crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); - let body = kernel_blocks(&ctx, module); - let found_bits = find_all::(&ctx, &body) - .iter() - .any(|constant| { - constant - .get_value(&ctx) - .downcast_ref::() - .is_some_and(|value| value.value().to_u64() == bits) - }); - assert!( - found_bits, - "single discriminant {bits} must not be truncated" - ); - assert_eq!(count_ops::(&ctx, &body) == 1, expects_sext); - } - } - - /// The row-width arm of [`resolve_aggregate_slots`]' no-history fallback - /// is defensive symmetry: no current lowering path produces a runtime-width - /// slice value with an empty conversion history, so no end-to-end pipeline - /// test can reach it. This exercises the fallback directly: a value born as - /// the lowered `{ ptr, i64, i32 }` row-width shape (a block argument, which - /// carries no history) must resolve to identity indexing exactly like the - /// two-field `{ ptr, i64 }` fat pointer, and any other unrecognized - /// no-history shape must keep failing closed into the refuse-to-guess - /// error (issue #128). - #[test] - fn no_history_fallback_resolves_row_width_slice_shape_and_stays_closed_otherwise() { - let mut ctx = make_ctx(); - - // The exact struct the type converter lowers a runtime-width - // disjoint slice to, built through the same constructor the fallback - // uses so the test cannot drift from the lowered shape. - let width_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); - let row_width_ty = crate::convert::types::make_disjoint_slice_struct(&mut ctx, &[width_ty]) - .expect("building the row-width slice struct must succeed"); - - // A three-field control shape that is NOT the row-width slice: i64 - // where the u32 row-width word belongs. - use llvm_export::types::PointerTypeExt; - let ptr_ty: TypeHandle = llvm_types::PointerType::get_generic(&mut ctx).into(); - let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signless).into(); - let foreign_ty: TypeHandle = llvm_types::StructType::get_unnamed( - &ctx, - ( - vec![ptr_ty, i64_ty, i64_ty], - llvm_types::StructLayout::Unpacked, - ), - ) - .into(); - - let (_module, block) = build_kernel(&mut ctx, vec![row_width_ty, foreign_ty], vec![]); - let row_width_value = block.deref(&ctx).get_argument(0); - let foreign_value = block.deref(&ctx).get_argument(1); - - // Block arguments have no recorded conversion history at all. - let no_history = OperandsInfo::default(); - - let resolved = resolve_aggregate_slots(&mut ctx, &no_history, row_width_value) - .expect("the row-width slice shape with no history must resolve"); - assert!( - matches!(resolved, AggregateSlots::Identity), - "the row-width {{ ptr, i64, u32 }} shape is index-preserving and must map identically" - ); - - let refused = resolve_aggregate_slots(&mut ctx, &no_history, foreign_value); - assert!( - refused.is_err(), - "a no-history shape that is not a slice fat pointer (with or without \ - the row-width word) must fail closed" - ); - } -} diff --git a/crates/mir-lower/src/convert/ops/aggregate/addressing.rs b/crates/mir-lower/src/convert/ops/aggregate/addressing.rs new file mode 100644 index 0000000000..e112f93681 --- /dev/null +++ b/crates/mir-lower/src/convert/ops/aggregate/addressing.rs @@ -0,0 +1,1185 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use super::common::{anyhow_to_pliron, byte_offset_gep}; +use crate::convert::enum_payload_storage::enum_payload_storage_type; +use crate::convert::types::{ + StructLayoutInfo, build_enum_slot_map, build_struct_slot_map, convert_type, + llvm_type_size_align, mir_element_stride, mir_type_abi_align, +}; +use dialect_mir::ops::{MirConstantOp, MirFieldAddrOp}; +use dialect_mir::types::{MirEnumType, MirPtrType, MirStructType, MirTupleType, MirUnionType}; +use llvm_export::ops as llvm; +use pliron::builtin::types::{IntegerType, Signedness}; +use pliron::context::{Context, Ptr}; +use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; +use pliron::irbuild::inserter::Inserter; +use pliron::irbuild::rewriter::Rewriter; +use pliron::location::Located; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::result::Result; +use pliron::r#type::{TypeHandle, Typed}; +use pliron::value::Value; + +// ============================================================================ +// MirFieldAddrOp Conversion +// ============================================================================ + +/// Convert `mir.field_addr` to `llvm.getelementptr`. +/// +/// Computes the address of a struct field using GEP. This is needed when +/// Rust code takes `&mut self.field` — we need the ADDRESS of the field, +/// not a COPY of its value. +/// +/// The GEP field index and the struct type it indexes into both come from +/// [`build_struct_slot_map`], so the index accounts for reordering, +/// `[N x i8]` padding slots and stripped ZSTs (ZST-ness is decided on the +/// converted LLVM field type, like the value-level sites). +/// +/// Taking the address of a ZST field emits a distinct zero-offset `i8` GEP +/// off the base pointer (a ZST field lives at byte 0 of the struct), the +/// same idiom the union branch uses. It must NOT forward the base SSA value +/// itself: a 1:1 `replace_operation_with_values` that aliases the op's +/// result to an already-existing value records the result's pointee type +/// (the ZST field's own zero-field type) onto that value's conversion type +/// history, so a sibling `field_addr` on the same base would later resolve +/// the base's pointee to the ZST type instead of the struct and fail with +/// "field_addr index N out of bounds for struct with 0 fields". The +/// distinct GEP result leaves the base pointer's recorded pointee intact +/// and carries the field's pointee type on its own result, which is what +/// nested projections through the field address look up. +pub(crate) fn convert_field_addr( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + operands_info: &OperandsInfo, +) -> Result<()> { + let ptr_operand = op.deref(ctx).get_operand(0); + + let field_addr_op = MirFieldAddrOp::new(op); + let field_index = match field_addr_op.get_attr_field_index(ctx) { + Some(attr) => attr.0 as usize, + None => return pliron::input_err_noloc!("MirFieldAddrOp missing field_index attribute"), + }; + + let mir_ptr_pointee = + match operands_info.lookup_most_recent_of_type::(ctx, ptr_operand) { + Some(r) => r.pointee, + None => { + return pliron::input_err_noloc!("MirFieldAddrOp operand must be pointer type"); + } + }; + + let union_field_count = mir_ptr_pointee + .deref(ctx) + .downcast_ref::() + .map(MirUnionType::field_count); + if let Some(field_count) = union_field_count { + if field_index >= field_count { + return pliron::input_err_noloc!( + "field_addr index {} out of bounds for union with {} fields", + field_index, + field_count + ); + } + // Every union field begins at byte zero. Emit an explicit zero-offset + // GEP instead of forwarding the base SSA value directly: the distinct + // result keeps dialect conversion's pointer-type history unambiguous + // for repeated field accesses and for `union.struct_field.inner`. + use llvm_export::ops::GepIndex; + let i8_ty: TypeHandle = IntegerType::get(ctx, 8, Signedness::Signless).into(); + let gep = llvm::GetElementPtrOp::new(ctx, ptr_operand, vec![GepIndex::Constant(0)], i8_ty); + rewriter.insert_operation(ctx, gep.get_operation()); + rewriter.replace_operation(ctx, op, gep.get_operation()); + return Ok(()); + } + + // An enum payload field, addressed by its position in the flattened + // `all_field_types`. Variants share bytes, so the slot map resolves each + // field one of two ways: its own LLVM slot, or, when its bytes are already + // held by a differently typed field of another variant, a byte offset into + // the enum. Both give the address of the same storage, which is what makes + // a write through the returned pointer land in the enum rather than a copy. + let enum_name = mir_ptr_pointee + .deref(ctx) + .downcast_ref::() + .map(|enum_ty| enum_ty.name().to_string()); + if let Some(enum_name) = enum_name { + let map = build_enum_slot_map(ctx, mir_ptr_pointee).map_err(anyhow_to_pliron)?; + let Some(slot_entry) = map.field_slots.get(field_index).copied() else { + return pliron::input_err_noloc!( + "field_addr index {} out of bounds for enum with {} payload fields", + field_index, + map.field_slots.len() + ); + }; + + // Enum payload bytes hold one CANONICAL storage type that can differ + // from the field's semantic type: a bool is physically an i8 byte and + // a shared-memory pointer is stored as a generic pointer, with the + // value paths (construct/extract) coercing exactly at that boundary. + // The address computed here ESCAPES this site: the loads and stores + // made through it happen at arbitrary other sites and are typed with + // the SEMANTIC type, so no storage coercion can be attached at + // address-formation time. Handing the address out anyway would let an + // i1 store leave the byte's upper seven bits undefined for the i8 and + // niche-tag readers, or write a shared-pointer representation into + // bytes every other reader interprets (and, on modern NVVM, sizes) as + // a generic pointer. Fail closed instead, the same way the slot map + // already rejects bool bytes hidden behind unions. + let semantic_ty = map.field_llvm_types[field_index]; + let storage_ty = enum_payload_storage_type(ctx, semantic_ty).map_err(anyhow_to_pliron)?; + if storage_ty != semantic_ty { + return pliron::input_err_noloc!( + "field_addr: cannot hand out the in-place address of payload field {} of enum `{}`: its bytes use canonical storage type {} while its semantic type is {}, and loads or stores made through an escaped payload address are typed with the semantic type, which the canonical bytes cannot honor; shared reads of such payloads are compiled through a value copy automatically, and a write that stays inside its function is compiled by rebuilding the enum around the new payload; a borrow that escapes into a call keeps no such rewrite and is refused here", + field_index, + enum_name, + storage_ty.deref(ctx).disp(ctx), + semantic_ty.deref(ctx).disp(ctx) + ); + } + + use llvm_export::ops::GepIndex; + let gep = match slot_entry { + Some(slot) => llvm::GetElementPtrOp::new( + ctx, + ptr_operand, + vec![GepIndex::Constant(0), GepIndex::Constant(slot)], + map.llvm_struct_ty, + ), + None => { + // No slot of its own: address the bytes directly, off the + // ORIGINAL enum pointer, so a write through the result lands + // in the enum. A zero-sized field lands at its offset like + // any other and simply spans nothing. The offset is always + // present: the slot map builds `field_offsets` and + // `field_slots` from one walk over the same fields, and the + // bounds check above already validated the index. + let offset = map.field_offsets[field_index]; + let offset = u32::try_from(offset).map_err(|_| { + pliron::input_error_noloc!( + "field_addr: payload byte offset {} of enum `{}` exceeds u32", + offset, + enum_name + ) + })?; + let i8_ty: TypeHandle = IntegerType::get(ctx, 8, Signedness::Signless).into(); + llvm::GetElementPtrOp::new( + ctx, + ptr_operand, + vec![GepIndex::Constant(offset)], + i8_ty, + ) + } + }; + rewriter.insert_operation(ctx, gep.get_operation()); + rewriter.replace_operation(ctx, op, gep.get_operation()); + return Ok(()); + } + + // Carried alongside the layout so the field address can record what it + // proves about its own alignment; see `stamp_field_address_alignment`. + let aggregate_abi_align; + let layout = { + let pointee_ref = mir_ptr_pointee.deref(ctx); + if let Some(struct_ty) = pointee_ref.downcast_ref::() { + aggregate_abi_align = struct_ty.abi_align; + StructLayoutInfo::of_struct(struct_ty) + } else if let Some(tuple_ty) = pointee_ref.downcast_ref::() { + aggregate_abi_align = tuple_ty.abi_align(); + StructLayoutInfo::of_tuple(tuple_ty) + } else { + return pliron::input_err_noloc!( + "MirFieldAddrOp pointer must point to a struct, tuple, union or enum type, got {}", + mir_ptr_pointee.deref(ctx).disp(ctx) + ); + } + }; + + let map = build_struct_slot_map(ctx, &layout).map_err(anyhow_to_pliron)?; + + let slot = match map.decl_to_llvm.get(field_index) { + Some(Some(slot)) => *slot, + Some(None) => { + // ZST field: it has no storage; the struct address stands in for + // the field address. Emit an explicit zero-offset GEP instead of + // forwarding the base SSA value directly (mirrors the union branch + // above): the distinct result keeps dialect conversion's pointer-type + // history unambiguous for repeated field accesses off the same base + // pointer. Forwarding the base value here type-puns it to the field's + // pointee, corrupting the base pointer's recorded type so a sibling + // field_addr on the same base later resolves to the wrong (0-field) + // pointee -- the "field_addr index N out of bounds for struct with 0 + // fields" failure on a struct with >= 2 ZST fields (e.g. a closure + // that captures other ZST closures, like iter map_fold). + use llvm_export::ops::GepIndex; + let i8_ty: TypeHandle = IntegerType::get(ctx, 8, Signedness::Signless).into(); + let gep = + llvm::GetElementPtrOp::new(ctx, ptr_operand, vec![GepIndex::Constant(0)], i8_ty); + rewriter.insert_operation(ctx, gep.get_operation()); + rewriter.replace_operation(ctx, op, gep.get_operation()); + return Ok(()); + } + None => { + return pliron::input_err_noloc!( + "field_addr index {} out of bounds for struct with {} fields", + field_index, + map.decl_to_llvm.len() + ); + } + }; + + let rustc_offset = layout.field_offsets.get(field_index).copied(); + + // Preserve the #859 address-path contract independently of the value + // representation. `build_struct_slot_map` retains the offsets the same + // fields would have under natural LLVM layout; when those differ from + // rustc's recorded offsets, keep using the original aggregate pointer plus + // a byte GEP. The semantic value type may now be an LLVM packed struct, but + // changing this established field-address path is outside this change. + if let Some(expected_offset) = rustc_offset { + let actual_offset = map + .natural_slot_offsets + .as_ref() + .and_then(|offsets| offsets.get(slot as usize).copied()) + .ok_or_else(|| { + pliron::input_error_noloc!( + "field_addr: cannot determine LLVM byte offset of slot {} for field {}", + slot, + field_index + ) + })?; + if map.layout_diverges && actual_offset != expected_offset { + let field_ptr = byte_offset_gep(ctx, rewriter, ptr_operand, expected_offset); + let gep = field_ptr + .defining_op() + .expect("byte_offset_gep always returns a GEP result"); + stamp_field_address_alignment(ctx, gep, aggregate_abi_align, Some(expected_offset)); + rewriter.replace_operation(ctx, op, gep); + return Ok(()); + } + } + + use llvm_export::ops::GepIndex; + let gep_indices = vec![GepIndex::Constant(0), GepIndex::Constant(slot)]; + + let gep_op = llvm::GetElementPtrOp::new(ctx, ptr_operand, gep_indices, map.llvm_struct_ty); + rewriter.insert_operation(ctx, gep_op.get_operation()); + stamp_field_address_alignment( + ctx, + gep_op.get_operation(), + aggregate_abi_align, + rustc_offset, + ); + rewriter.replace_operation(ctx, op, gep_op.get_operation()); + + Ok(()) +} + +/// Record on a field address what its alignment provably is. +/// +/// A load reads its alignment from its own result type, and a scalar records +/// none -- so `p.x` on an `#[repr(C, align(8))]` struct would export with LLVM's +/// default `align 4`, dropping the guarantee rustc gave the aggregate. That is +/// what stops LoadStoreVectorizer from fusing two adjacent field reads into one +/// wide access; a whole-element copy of the same struct already vectorizes, +/// because there the load's result type *is* the aggregate. +/// +/// The aggregate is aligned to `abi_align`, so its field at byte `offset` is +/// aligned to `gcd(abi_align, offset)`, and to `abi_align` itself at offset +/// zero. Both numbers are rustc's own layout, so this claims nothing the source +/// did not already guarantee. +/// +/// Stamps nothing when the aggregate records no alignment (`abi_align == 0`, +/// which also stands for "layout unknown") or the field has no recorded offset. +/// A wrong alignment here is a miscompile, so every uncertain case declines and +/// leaves the previous, weaker-but-sound behaviour. +fn stamp_field_address_alignment( + ctx: &mut Context, + gep: Ptr, + abi_align: u64, + field_offset: Option, +) { + const fn gcd(a: u64, b: u64) -> u64 { + if b == 0 { a } else { gcd(b, a % b) } + } + + if abi_align == 0 { + return; + } + let Some(offset) = field_offset else { + return; + }; + let provable = if offset == 0 { + abi_align + } else { + gcd(abi_align, offset) + }; + // Every rustc layout has a power-of-two `abi_align`, but dialect-mir only + // verifier-enforces that for unions and enums; a malformed hand-built + // struct or tuple layout could reach here with e.g. 12, and a + // non-power-of-two `align N` is invalid LLVM IR that llc rejects. + if !provable.is_power_of_two() { + return; + } + if let Ok(align) = u32::try_from(provable) { + llvm_export::ops::set_address_alignment(ctx, gep, align); + } +} + +// ============================================================================ +// MirArrayElementAddrOp Conversion +// ============================================================================ + +/// Convert `mir.array_element_addr` to `llvm.getelementptr`. +/// +/// ```text +/// &arr[i] ──► getelementptr T, ptr %arr_ptr, i64 %i +/// ``` +/// +/// - The element type `T` comes from the op's OWN result type +/// (`mir.ptr`), which the dialect verifier ties to the operand +/// array's element type. Same address as the old two-index +/// `[N x T]` GEP, minus the dead `[N x T]` base. +/// - Operand type history is not usable here: a kind-only `mir.cast` +/// lowers to a plain value forwarding, history does not follow that +/// edge, and a stale hit would stride by the wrong element size. +pub(crate) fn convert_array_element_addr( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + _operands_info: &OperandsInfo, +) -> Result<()> { + let loc = op.deref(ctx).loc(); + let arr_ptr = op.deref(ctx).get_operand(0); + let index = op.deref(ctx).get_operand(1); + + let result_ty = op.deref(ctx).get_result(0).get_type(ctx); + let element_ty = result_ty + .deref(ctx) + .downcast_ref::() + .map(|mir_ptr| mir_ptr.pointee) + .ok_or_else(|| { + pliron::input_error!( + loc.clone(), + "mir.array_element_addr result must be a MIR pointer type; \ + element sizing has no fact to derive from" + ) + })?; + + let llvm_element_ty = convert_type(ctx, element_ty).map_err(anyhow_to_pliron)?; + + // The typed GEP strides by the LLVM allocation size of the element type. + // Keep the rustc-stride comparison as a fail-closed backstop for any + // element representation whose LLVM allocation size differs. + { + let rustc_stride = mir_element_stride(ctx, element_ty); + let llvm_size = llvm_type_size_align(ctx, llvm_element_ty).map(|(size, _)| size); + if let (Some(stride), Some(llvm_size)) = (rustc_stride, llvm_size) + && stride != llvm_size + { + return pliron::input_err_noloc!( + "addressing elements of an array whose LLVM allocation size differs from \ + rustc's stored stride is not supported: rustc strides by {} bytes but \ + the LLVM element type occupies {}", + stride, + llvm_size + ); + } + } + + use llvm_export::ops::GepIndex; + let gep_indices = vec![GepIndex::Value(index)]; + + let element_align = element_address_provable_alignment(ctx, arr_ptr, element_ty, index); + + let gep_op = llvm::GetElementPtrOp::new(ctx, arr_ptr, gep_indices, llvm_element_ty); + rewriter.insert_operation(ctx, gep_op.get_operation()); + if let Some(align) = element_align { + llvm_export::ops::set_address_alignment(ctx, gep_op.get_operation(), align); + } + rewriter.replace_operation(ctx, op, gep_op.get_operation()); + + Ok(()) +} + +/// Alignment an array-element address provably has, in bytes. +/// +/// The field path records this already ([`stamp_field_address_alignment`]), but +/// an element address recorded nothing, so `lanes[0] + lanes[1]` on an +/// `#[repr(C, align(8))]` `[f32; 2]` still exported two `align 4` loads and +/// LoadStoreVectorizer refused to fuse them. Reading the same element into a +/// local first happens to work, because SROA then loads the whole aggregate and +/// the alignment comes from its type — so the cost depended on whether the +/// source copied the element or read through it. +/// +/// Two facts combine. The element's `abi_align` equals the array's (a Rust +/// array aligns to its element), and whatever the *base pointer* already +/// proved is stronger when the array is itself a field of an over-aligned +/// aggregate: `&table[i].lanes` carries the outer struct's alignment, which +/// the element type alone does not know. Element `i` then sits at byte +/// `i * stride`, so it is aligned to `gcd(base, i * stride)`, and to `base` +/// itself at index zero. A runtime index can land on any element, so it gets +/// `gcd(base, stride)`, which every element satisfies. +/// +/// Answers `None`, leaving the previous behaviour, when neither the base nor +/// the element type records an alignment, or the stride is unknown. As on the +/// field path, a wrong answer here is a miscompile, so every uncertain case +/// declines rather than guesses. That is why the stride comes from +/// [`mir_element_stride`], which is exact or `None`: an LLVM-level size +/// approximation would guess 8 for a MIR aggregate element it does not +/// model, and e.g. `[[f32; 3]; 4]` under an align-8 base would then claim +/// align 8 on element addresses that are only 4-aligned. +fn element_address_provable_alignment( + ctx: &Context, + arr_ptr: Value, + element_ty: TypeHandle, + index: Value, +) -> Option { + const fn gcd(a: u64, b: u64) -> u64 { + if b == 0 { a } else { gcd(b, a % b) } + } + + // What the base address already proved, else what the element type + // states (identical to the array's own alignment). + let base_align = arr_ptr + .defining_op() + .and_then(|def| llvm_export::ops::address_alignment(ctx, def)) + .map(u64::from) + .or_else(|| mir_type_abi_align(ctx, element_ty))?; + if base_align == 0 { + return None; + } + + let stride = mir_element_stride(ctx, element_ty)?; + if stride == 0 { + return None; + } + + let provable = match constant_index_value(ctx, index) { + Some(0) => base_align, + Some(i) => gcd(base_align, i.checked_mul(stride)?), + // Any element is reachable, so claim only what every stride preserves. + None => gcd(base_align, stride), + }; + + // Same guard as the field path: dialect-mir does not verifier-enforce a + // power-of-two `abi_align` for every aggregate, and a non-power-of-two + // `align N` is invalid LLVM IR that llc rejects. + if !provable.is_power_of_two() { + return None; + } + u32::try_from(provable).ok() +} + +/// The constant an index operand holds, if it is one. +/// +/// `APInt::to_u64` truncates wider values, so a >64-bit constant could be +/// misread as a small offset multiplier. Fail closed on such widths, as +/// `integer_constant_u64` in the extract-element fast path does. +fn constant_index_value(ctx: &Context, index: Value) -> Option { + let defining_op = index.defining_op()?; + if let Some(constant) = Operation::get_op::(defining_op, ctx) { + let value = constant.get_attr_value(ctx)?.value(); + return (value.bw() <= 64).then(|| value.to_u64()); + } + let constant = Operation::get_op::(defining_op, ctx)?; + let attribute = constant.get_value(ctx); + let integer = attribute.downcast_ref::()?; + let value = integer.value(); + (value.bw() <= 64).then(|| value.to_u64()) +} + +#[cfg(test)] +// Tests build kinded fixture types directly; production minting lives in mir-importer's facts.rs. +#[allow(clippy::disallowed_methods)] +mod tests { + use super::*; + use crate::convert::ops::test_util::*; + + use dialect_mir::ops as mir; + use dialect_mir::types::{ + EnumCarrierKind, EnumEncoding, EnumLayoutKind, EnumVariant, MirArrayType, MirPtrType, + MirStructType, MirTupleType, + }; + use llvm_export::types as llvm_types; + use pliron::builtin::attributes::IntegerAttr; + use pliron::common_traits::Verify; + + use super::super::test_support::*; + + fn byte_gep_constant_offset(ctx: &Context, gep: &llvm::GetElementPtrOp) -> Option { + use llvm_export::ops::GepIndex; + + let indices = gep.indices(ctx); + let [GepIndex::Value(offset)] = indices.as_slice() else { + return None; + }; + let defining_op = offset.defining_op()?; + let constant = Operation::get_op::(defining_op, ctx)?; + let attribute = constant.get_value(ctx); + attribute + .downcast_ref::() + .map(|integer| integer.value().to_u64()) + } + + fn assert_byte_addressed_field( + ctx: &Context, + module: Ptr, + expected_offset: u64, + expected_alignment: u32, + ) { + let body = kernel_blocks(ctx, module); + let geps = find_all::(ctx, &body); + assert_eq!(geps.len(), 1, "one field_addr must lower to one GEP"); + + let gep = &geps[0]; + assert_eq!( + gep.src_elem_type(ctx), + IntegerType::get(ctx, 8, Signedness::Signless).into(), + "a layout-mismatched field must be addressed in byte units" + ); + assert_eq!( + byte_gep_constant_offset(ctx, gep), + Some(expected_offset), + "the byte GEP must use rustc's exact field offset" + ); + assert_eq!( + llvm_export::ops::address_alignment(ctx, gep.get_operation()), + Some(expected_alignment), + "the byte GEP must retain the alignment proved by rustc's aggregate layout" + ); + } + + /// Sibling ZST field addresses off one base pointer must each lower to + /// their own zero-offset `i8` GEP, never to the base SSA value itself. + /// Forwarding the base (the pre-fix behavior) recorded the first field's + /// zero-field pointee type onto the base pointer's conversion history, so + /// the second `field_addr` resolved the base's pointee to "struct with 0 + /// fields" and lowering aborted. This is the `iter().map(..).sum()` shape: + /// a composed closure holding two captureless (zero-sized) closures. + #[test] + fn sibling_zst_field_addrs_lower_to_distinct_zero_offset_geps() { + use llvm_export::ops::GepIndex; + + let mut ctx = make_ctx(); + + let i8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Signless).into(); + let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signless).into(); + let zst_f = empty_struct_ty(&mut ctx, "MapClosure"); + let zst_g = empty_struct_ty(&mut ctx, "SumClosure"); + + // struct Composed { f: MapClosure (ZST), g: SumClosure (ZST), acc: i64 } + let struct_ty: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "Composed".to_string(), + vec!["f".to_string(), "g".to_string(), "acc".to_string()], + vec![zst_f, zst_g, i64_ty], + vec![0, 1, 2], + vec![0, 0, 0], + 8, + 8, + ) + .into(); + + let base_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, struct_ty, true).into(); + // Each field_addr result records the field's own pointee type; for the + // ZST fields that zero-field pointee is exactly what poisoned the base + // pointer's history when the result aliased the base. + let f_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, zst_f, true).into(); + let g_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, zst_g, true).into(); + let acc_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, i64_ty, true).into(); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![base_ptr_ty], vec![]); + let base = block.deref(&ctx).get_argument(0); + + for (field_index, result_ty) in [(0u32, f_ptr_ty), (1, g_ptr_ty), (2, acc_ptr_ty)] { + let op = MirFieldAddrOp::build(&mut ctx, base, result_ty, field_index) + .expect("field_addr build"); + op.insert_at_back(block, &ctx); + } + append_mir_return(&mut ctx, block, vec![]); + + // Pre-fix this failed with "field_addr index 1 out of bounds for + // struct with 0 fields" on the second (sibling) ZST field_addr. + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + let geps = find_all::(&ctx, &body); + assert_eq!(geps.len(), 3, "each field_addr must lower to its own GEP"); + assert!( + geps.iter().all(|gep| gep.verify(&ctx).is_ok()), + "every field-address GEP must satisfy LLVM dialect verification" + ); + + let zst_geps: Vec<_> = geps + .iter() + .filter(|gep| { + gep.src_elem_type(&ctx) == i8_ty + && matches!(gep.indices(&ctx).as_slice(), [GepIndex::Constant(0)]) + }) + .collect(); + assert_eq!( + zst_geps.len(), + 2, + "both ZST field addresses must lower to zero-offset i8 GEPs" + ); + + let gep_base = |gep: &llvm::GetElementPtrOp| gep.get_operation().deref(&ctx).get_operand(0); + let gep_result = + |gep: &llvm::GetElementPtrOp| gep.get_operation().deref(&ctx).get_result(0); + assert_eq!( + gep_base(zst_geps[0]), + gep_base(zst_geps[1]), + "both ZST field addresses must be taken off the same base pointer" + ); + assert_ne!( + gep_result(zst_geps[0]), + gep_base(zst_geps[0]), + "a ZST field address must be a value distinct from the base pointer" + ); + assert_ne!( + gep_result(zst_geps[1]), + gep_base(zst_geps[1]), + "a ZST field address must be a value distinct from the base pointer" + ); + assert_ne!( + gep_result(zst_geps[0]), + gep_result(zst_geps[1]), + "each sibling ZST field address must get its own GEP result" + ); + + // The non-ZST sibling still resolves the base's pointee to the struct + // (slot 0 holds `acc`): the base pointer's type history stayed intact. + let acc_geps: Vec<_> = geps + .iter() + .filter(|gep| { + matches!( + gep.indices(&ctx).as_slice(), + [GepIndex::Constant(0), GepIndex::Constant(0)] + ) + }) + .collect(); + assert_eq!( + acc_geps.len(), + 1, + "the non-ZST field address must index the struct's layout slot" + ); + } + + /// `mir.field_addr` on a TUPLE pointee (the `#693` shape: `let (a, b) = + /// TABLE[i];`) must resolve the GEP index through the tuple's memory-order + /// layout, exactly like the struct path above, not through the + /// declaration index directly. + /// + /// `(u8, u32)` is rustc's own layout for this pair: the `u32` field is + /// placed FIRST in memory for alignment, so declaration field 0 (`u8`, + /// `.0`) lives at memory slot 1 and declaration field 1 (`u32`, `.1`) + /// lives at memory slot 0. A GEP index equal to the declaration index + /// would silently address the WRONG field's bytes. + #[test] + fn field_addr_tuple_pointee_resolves_memory_order_gep_index() { + use llvm_export::ops::GepIndex; + + let mut ctx = make_ctx(); + + let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + + // decl order [u8, u32]; memory order [u32, u8] (mem_to_decl = [1, 0]). + let tuple_ty: TypeHandle = MirTupleType::get_with_layout( + &mut ctx, + vec![u8_ty, u32_ty], + vec![1, 0], + vec![4, 0], + 8, + 4, + ) + .into(); + + let tuple_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, tuple_ty, false).into(); + let u8_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, u8_ty, false).into(); + let u32_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, u32_ty, false).into(); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![tuple_ptr_ty], vec![]); + let base = block.deref(&ctx).get_argument(0); + + // Declaration field 0 (`.0`, u8) first, then declaration field 1 + // (`.1`, u32) -- source order, not memory order. + for (field_index, result_ty) in [(0u32, u8_ptr_ty), (1, u32_ptr_ty)] { + let op = MirFieldAddrOp::build(&mut ctx, base, result_ty, field_index) + .expect("field_addr build"); + op.insert_at_back(block, &ctx); + } + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + let geps = find_all::(&ctx, &body); + assert_eq!(geps.len(), 2, "each field_addr must lower to its own GEP"); + assert!( + geps.iter().all(|gep| gep.verify(&ctx).is_ok()), + "every field-address GEP must satisfy LLVM dialect verification" + ); + + // `.0` (u8, declared first) must land at MEMORY slot 1. + let field0_geps: Vec<_> = geps + .iter() + .filter(|gep| { + matches!( + gep.indices(&ctx).as_slice(), + [GepIndex::Constant(0), GepIndex::Constant(1)] + ) + }) + .collect(); + assert_eq!( + field0_geps.len(), + 1, + "declaration field 0 (u8) must resolve to its memory slot 1, not slot 0" + ); + + // `.1` (u32, declared second) must land at MEMORY slot 0. + let field1_geps: Vec<_> = geps + .iter() + .filter(|gep| { + matches!( + gep.indices(&ctx).as_slice(), + [GepIndex::Constant(0), GepIndex::Constant(0)] + ) + }) + .collect(); + assert_eq!( + field1_geps.len(), + 1, + "declaration field 1 (u32) must resolve to its memory slot 0, not slot 1" + ); + } + + /// A `#[repr(C, packed)]`-style layout can place a naturally aligned + /// scalar at an offset LLVM's ordinary struct layout cannot express. The + /// field address must therefore use rustc's byte offset rather than a + /// typed struct GEP that would silently land at byte 4. + #[test] + fn packed_field_addr_uses_rustc_byte_offset_and_alignment_one() { + let mut ctx = make_ctx(); + + let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let packed_ty: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "Packed".into(), + vec!["tag".into(), "value".into()], + vec![u8_ty, u32_ty], + vec![0, 1], + vec![0, 1], + 5, + 1, + ) + .into(); + + let base_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, packed_ty, false).into(); + let field_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, u32_ty, false).into(); + let (module, block) = build_kernel(&mut ctx, vec![base_ptr_ty], vec![]); + let base = block.deref(&ctx).get_argument(0); + + let field_addr = + MirFieldAddrOp::build(&mut ctx, base, field_ptr_ty, 1).expect("field_addr build"); + field_addr.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); + assert_byte_addressed_field(&ctx, module, 1, 1); + } + + /// `repr(packed(2))` is not equivalent to byte alignment: the same u32 + /// field is allowed to sit at byte 2 and the address still proves align 2. + #[test] + fn packed_two_field_addr_uses_rustc_byte_offset_and_alignment_two() { + let mut ctx = make_ctx(); + + let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let packed_ty: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "Packed2".into(), + vec!["tag".into(), "value".into()], + vec![u8_ty, u32_ty], + vec![0, 1], + vec![0, 2], + 6, + 2, + ) + .into(); + + let base_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, packed_ty, false).into(); + let field_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, u32_ty, false).into(); + let (module, block) = build_kernel(&mut ctx, vec![base_ptr_ty], vec![]); + let base = block.deref(&ctx).get_argument(0); + + let field_addr = + MirFieldAddrOp::build(&mut ctx, base, field_ptr_ty, 1).expect("field_addr build"); + field_addr.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); + assert_byte_addressed_field(&ctx, module, 2, 2); + } + + /// The decision must compare physical offsets, not merely the selected + /// field's own alignment. In `{ u8, u32, u8 }` the final u8 is naturally + /// byte-aligned, but LLVM has already shifted it because the packed u32 + /// before it was placed at byte 4 instead of byte 1. + #[test] + fn packed_trailing_byte_field_uses_accumulated_rustc_offset() { + let mut ctx = make_ctx(); + + let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let packed_ty: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "PackedTrailing".into(), + vec!["head".into(), "value".into(), "tail".into()], + vec![u8_ty, u32_ty, u8_ty], + vec![0, 1, 2], + vec![0, 1, 5], + 6, + 1, + ) + .into(); + + let base_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, packed_ty, false).into(); + let field_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, u8_ty, false).into(); + let (module, block) = build_kernel(&mut ctx, vec![base_ptr_ty], vec![]); + let base = block.deref(&ctx).get_argument(0); + + let field_addr = + MirFieldAddrOp::build(&mut ctx, base, field_ptr_ty, 2).expect("field_addr build"); + field_addr.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); + assert_byte_addressed_field(&ctx, module, 5, 1); + } + + /// Packed LLVM element types carry the same allocation size as rustc, so + /// typed array GEPs now stride by the correct packed byte count. + #[test] + fn packed_array_element_addressing_uses_packed_stride() { + use dialect_mir::ops::MirArrayElementAddrOp; + + let mut ctx = make_ctx(); + + let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let packed_ty: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "Packed".into(), + vec!["tag".into(), "value".into()], + vec![u8_ty, u32_ty], + vec![0, 1], + vec![0, 1], + 5, + 1, + ) + .into(); + let array_ty: TypeHandle = MirArrayType::get(&mut ctx, packed_ty, 4).into(); + let base_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, array_ty, false).into(); + let elem_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, packed_ty, false).into(); + + let (module, block) = build_kernel(&mut ctx, vec![base_ptr_ty, u64_ty], vec![]); + let base = block.deref(&ctx).get_argument(0); + let index = block.deref(&ctx).get_argument(1); + + let elem_addr = Operation::new( + &mut ctx, + MirArrayElementAddrOp::get_concrete_op_info(), + vec![elem_ptr_ty], + vec![base, index], + vec![], + 0, + ); + elem_addr.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module) + .expect("packed element addressing must use the packed LLVM stride"); + let body = kernel_blocks(&ctx, module); + assert_eq!( + count_ops::(&ctx, &body), + 1, + "packed array element addressing should lower to one typed GEP" + ); + } + + /// Taking the in-place address of a bool payload must fail loudly: the + /// payload's canonical storage is an i8 byte, and a semantic i1 store + /// through the escaped address would leave that byte's upper seven bits + /// undefined for every i8 reader (including a niche tag sharing them). + #[test] + fn bool_payload_field_addr_fails_closed() { + let mut ctx = make_ctx(); + let tag: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let bool_ty: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Signless).into(); + let enum_ty: TypeHandle = MirEnumType::get_with_layout( + &mut ctx, + "DirectBool".into(), + tag, + vec![0, 1], + vec![ + EnumVariant::new_with_layout("A".into(), vec![bool_ty], vec![4], vec![1]), + EnumVariant::unit("B".into()), + ], + 0, + 8, + 4, + ) + .into(); + let enum_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, enum_ty, true).into(); + let bool_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, bool_ty, true).into(); + + let (module, block) = build_kernel(&mut ctx, vec![enum_ptr_ty], vec![]); + let base = block.deref(&ctx).get_argument(0); + let op = MirFieldAddrOp::build(&mut ctx, base, bool_ptr_ty, 0).expect("field_addr build"); + op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + let err = crate::lower_mir_to_llvm(&mut ctx, module) + .expect_err("addressing a bool enum payload in place must fail to lower"); + assert!( + err.err.to_string().contains("canonical storage type"), + "unexpected error: {}", + err.err + ); + } + + /// Same gate for the slot arm: a shared-memory pointer payload is stored + /// as a GENERIC pointer (its slot is typed `ptr`), so an in-place address + /// would let a semantic `ptr addrspace(3)` store write a representation + /// every other reader interprets as a generic pointer. + #[test] + fn shared_pointer_payload_field_addr_fails_closed() { + let mut ctx = make_ctx(); + let logical: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signed).into(); + let pointee: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, true).into(); + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "OptionShared".into(), + logical, + vec![0, 1], + vec![ + EnumVariant::unit("None".into()), + EnumVariant::new_with_layout("Some".into(), vec![shared], vec![0], vec![8]), + ], + EnumEncoding { + tag_offset: 0, + total_size: 8, + abi_align: 8, + layout_kind: EnumLayoutKind::Niche, + carrier_kind: EnumCarrierKind::Pointer, + carrier_width: 64, + carrier_address_space: llvm_types::address_space::GENERIC, + niche_start: 0, + niche_variant_start: 0, + niche_variant_end: 0, + untagged_variant: 1, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + let enum_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, enum_ty, true).into(); + let payload_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, shared, true).into(); + + let (module, block) = build_kernel(&mut ctx, vec![enum_ptr_ty], vec![]); + let base = block.deref(&ctx).get_argument(0); + let op = + MirFieldAddrOp::build(&mut ctx, base, payload_ptr_ty, 0).expect("field_addr build"); + op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + let err = crate::lower_mir_to_llvm(&mut ctx, module) + .expect_err("addressing a shared-pointer enum payload in place must fail to lower"); + assert!( + err.err.to_string().contains("canonical storage type"), + "unexpected error: {}", + err.err + ); + } + + /// A payload with no slot of its own is addressed at its byte offset off + /// the ORIGINAL enum pointer: no stack spill is introduced, so a write + /// through the result lands in the enum rather than in a copy. + #[test] + fn slotless_payload_field_addr_geps_original_storage_at_byte_offset() { + use llvm_export::ops::GepIndex; + use pliron::builtin::types::FP32Type; + + let mut ctx = make_ctx(); + let tag: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let f32_ty: TypeHandle = FP32Type::get(&ctx).into(); + let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let enum_ty: TypeHandle = MirEnumType::get_with_layout( + &mut ctx, + "Either".into(), + tag, + vec![0, 1], + vec![ + EnumVariant::new_with_layout("Real".into(), vec![f32_ty], vec![4], vec![4]), + EnumVariant::new_with_layout("Bits".into(), vec![u32_ty], vec![4], vec![4]), + ], + 0, + 8, + 4, + ) + .into(); + let slot_map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); + assert_eq!( + slot_map.field_slots, + vec![Some(1), None], + "Real's f32 claims byte 4 first; Bits shares those bytes slotless" + ); + + let enum_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, enum_ty, true).into(); + let u32_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, u32_ty, true).into(); + let (module, block) = build_kernel(&mut ctx, vec![enum_ptr_ty], vec![]); + let base = block.deref(&ctx).get_argument(0); + let op = MirFieldAddrOp::build(&mut ctx, base, u32_ptr_ty, 1).expect("field_addr build"); + op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); + let body = kernel_blocks(&ctx, module); + assert_eq!( + count_ops::(&ctx, &body), + 0, + "an in-place payload address must not spill the enum to a stack copy" + ); + let geps = find_all::(&ctx, &body); + assert_eq!(geps.len(), 1, "one field_addr lowers to one GEP"); + let gep = &geps[0]; + // The MIR entry block and its arguments are the ORIGINALS (moved by + // `inline_region`), so the enum pointer argument keeps its identity + // through lowering and the GEP must be based directly on it. + assert_eq!( + gep.get_operation().deref(&ctx).get_operand(0), + base, + "the byte-offset GEP must address the ORIGINAL enum storage" + ); + assert!( + matches!(gep.indices(&ctx).as_slice(), [GepIndex::Constant(4)]), + "the slotless payload must be addressed at its rustc byte offset" + ); + assert_eq!( + gep.src_elem_type(&ctx), + IntegerType::get(&ctx, 8, Signedness::Signless).into(), + "byte addressing must step in i8 units" + ); + } + + /// The capability split, read side: a payload whose storage IS its + /// semantic type keeps the address path even for a SHARED borrow. The + /// borrow lowers to one GEP into the ORIGINAL enum storage and the read + /// to one load through it: no stack spill, no value copy. (Non-canonical + /// payloads never get here for shared borrows; the importer punts them + /// to the value-copy fallback before an address is formed.) + #[test] + fn canonical_payload_shared_read_loads_through_gep_without_copy() { + use llvm_export::ops::GepIndex; + use pliron::builtin::types::FP32Type; + + let mut ctx = make_ctx(); + let tag: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let f32_ty: TypeHandle = FP32Type::get(&ctx).into(); + let enum_ty: TypeHandle = MirEnumType::get_with_layout( + &mut ctx, + "Slot".into(), + tag, + vec![0, 1], + vec![ + EnumVariant::new_with_layout("Occupied".into(), vec![f32_ty], vec![4], vec![4]), + EnumVariant::unit("Empty".into()), + ], + 0, + 8, + 4, + ) + .into(); + let slot_map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); + assert_eq!( + slot_map.field_slots, + vec![Some(1)], + "an f32 payload owns its LLVM slot; storage equals semantic type" + ); + + // Immutable pointer types model the shared borrow. + let enum_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, enum_ty, false).into(); + let f32_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, f32_ty, false).into(); + let (module, block) = build_kernel(&mut ctx, vec![enum_ptr_ty], vec![f32_ty]); + let base = block.deref(&ctx).get_argument(0); + let addr = MirFieldAddrOp::build(&mut ctx, base, f32_ptr_ty, 0).expect("field_addr build"); + addr.insert_at_back(block, &ctx); + let payload_ptr = addr.deref(&ctx).get_result(0); + + let load = Operation::new( + &mut ctx, + mir::MirLoadOp::get_concrete_op_info(), + vec![f32_ty], + vec![payload_ptr], + vec![], + 0, + ); + load.insert_at_back(block, &ctx); + let loaded = load.deref(&ctx).get_result(0); + append_mir_return(&mut ctx, block, vec![loaded]); + + crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); + let body = kernel_blocks(&ctx, module); + assert_eq!( + count_ops::(&ctx, &body), + 0, + "a canonical payload read must not spill or copy the enum" + ); + let geps = find_all::(&ctx, &body); + assert_eq!(geps.len(), 1, "one field_addr lowers to one GEP"); + let gep = &geps[0]; + assert_eq!( + gep.get_operation().deref(&ctx).get_operand(0), + base, + "the GEP must address the ORIGINAL enum storage" + ); + assert!( + matches!( + gep.indices(&ctx).as_slice(), + [GepIndex::Constant(0), GepIndex::Constant(1)] + ), + "an own-slot payload is addressed through its struct slot" + ); + let loads = find_all::(&ctx, &body); + assert_eq!( + loads.len(), + 1, + "the read is a single load through the payload address" + ); + let gep_result = gep.get_operation().deref(&ctx).get_result(0); + assert_eq!( + loads[0].get_operation().deref(&ctx).get_operand(0), + gep_result, + "the load must go through the GEP result, not a copy" + ); + assert_eq!( + loads[0] + .get_operation() + .deref(&ctx) + .get_result(0) + .get_type(&ctx), + f32_ty, + "the load reads the payload at its semantic type" + ); + } +} diff --git a/crates/mir-lower/src/convert/ops/aggregate/array_extract.rs b/crates/mir-lower/src/convert/ops/aggregate/array_extract.rs new file mode 100644 index 0000000000..b4d0a36ba7 --- /dev/null +++ b/crates/mir-lower/src/convert/ops/aggregate/array_extract.rs @@ -0,0 +1,351 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use super::common::anyhow_to_pliron; +use crate::convert::types::{convert_type, mir_type_abi_align}; +use dialect_mir::types::MirArrayType; +use llvm_export::attributes::ICmpPredicateAttr; +use llvm_export::ops as llvm; +use pliron::builtin::types::{IntegerType, Signedness}; +use pliron::context::{Context, Ptr}; +use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; +use pliron::irbuild::inserter::Inserter; +use pliron::irbuild::rewriter::Rewriter; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::result::Result; +use pliron::r#type::Typed; +use pliron::utils::apint::APInt; +use pliron::value::Value; +use std::num::NonZeroUsize; + +/// Convert `mir.extract_array_element` to LLVM operations. +/// +/// A runtime index normally requires materializing the array in memory because +/// LLVM `extractvalue` accepts only constant indices. When the index is proven +/// to be `urem value, C`, however, it is in `0..C`. For small `C` within the +/// array bounds, emit one constant `extractvalue` per candidate and select the +/// runtime result in SSA. This avoids the temporary alloca that otherwise +/// becomes NVPTX local memory. +/// +/// Unbounded, oversized, or otherwise unsupported indices retain the existing +/// alloca+store+GEP+load fallback. +pub(crate) fn convert_extract_array_element( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + operands_info: &OperandsInfo, +) -> Result<()> { + // One shared cap for the candidate chain, so the mir-transforms + // canonicalization and this lowering fast path cannot drift apart. + use dialect_mir::ops::MAX_SCALARIZED_CANDIDATES; + + fn integer_constant_u64(ctx: &Context, value: Value) -> Option { + let defining_op = value.defining_op()?; + let constant = Operation::get_op::(defining_op, ctx)?; + let attribute = constant.get_value(ctx); + let integer = attribute.downcast_ref::()?; + let integer_value = integer.value(); + // `APInt::to_u64` truncates wider values, so a >64-bit constant could + // be misread as a small in-range divisor. Fail closed on such widths. + (integer_value.bw() <= 64).then(|| integer_value.to_u64()) + } + + fn bounded_urem_candidate_count(ctx: &Context, index: Value, array_size: u64) -> Option { + let defining_op = index.defining_op()?; + Operation::get_op::(defining_op, ctx)?; + + let divisor = defining_op.deref(ctx).get_operand(1); + let candidate_count = integer_constant_u64(ctx, divisor)?; + (candidate_count > 0 + && candidate_count <= array_size + && candidate_count <= MAX_SCALARIZED_CANDIDATES) + .then_some(candidate_count) + } + + fn integer_constant_like( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + reference: Value, + value: u64, + ) -> Result { + let reference_ty = reference.get_type(ctx); + let width = reference_ty + .deref(ctx) + .downcast_ref::() + .ok_or_else(|| { + pliron::input_error_noloc!( + "mir.extract_array_element index must lower to an integer" + ) + })? + .width(); + + let integer_ty = IntegerType::get(ctx, width, Signedness::Signless); + let attribute = pliron::builtin::attributes::IntegerAttr::new( + integer_ty, + APInt::from_u64( + value, + NonZeroUsize::new(width as usize).expect("integer width is nonzero"), + ), + ); + let constant = llvm::ConstantOp::new(ctx, attribute.into()); + rewriter.insert_operation(ctx, constant.get_operation()); + Ok(constant.get_operation().deref(ctx).get_result(0)) + } + + let array_val = op.deref(ctx).get_operand(0); + let index_val = op.deref(ctx).get_operand(1); + + let (element_ty, array_size) = { + match operands_info.lookup_most_recent_of_type::(ctx, array_val) { + Some(r) => (r.element_type(), r.size()), + None => return pliron::input_err_noloc!("Expected MirArrayType"), + } + }; + + if let Some(candidate_count) = bounded_urem_candidate_count(ctx, index_val, array_size) { + let mut candidates = Vec::with_capacity(candidate_count as usize); + for candidate_index in 0..candidate_count { + let extract = llvm::ExtractValueOp::new(ctx, array_val, vec![candidate_index as u32])?; + rewriter.insert_operation(ctx, extract.get_operation()); + candidates.push(extract.get_operation().deref(ctx).get_result(0)); + } + + let mut selected = *candidates + .last() + .expect("candidate count is proven nonzero"); + for candidate_index in (0..candidates.len().saturating_sub(1)).rev() { + let candidate_constant = + integer_constant_like(ctx, rewriter, index_val, candidate_index as u64)?; + let compare = + llvm::ICmpOp::new(ctx, ICmpPredicateAttr::EQ, index_val, candidate_constant); + rewriter.insert_operation(ctx, compare.get_operation()); + let condition = compare.get_operation().deref(ctx).get_result(0); + + let select = llvm::SelectOp::new(ctx, condition, candidates[candidate_index], selected); + rewriter.insert_operation(ctx, select.get_operation()); + selected = select.get_operation().deref(ctx).get_result(0); + } + + rewriter.replace_operation_with_values(ctx, op, vec![selected]); + return Ok(()); + } + + let llvm_element_ty = convert_type(ctx, element_ty).map_err(anyhow_to_pliron)?; + let llvm_array_ty = llvm_export::types::ArrayType::get(ctx, llvm_element_ty, array_size); + let abi_align = mir_type_abi_align(ctx, element_ty); + + let i64_ty = IntegerType::get(ctx, 64, Signedness::Signless); + let one_val = { + let one_apint = APInt::from_i64(1, NonZeroUsize::new(64).unwrap()); + let one_attr = pliron::builtin::attributes::IntegerAttr::new(i64_ty, one_apint); + let const_op = llvm::ConstantOp::new(ctx, one_attr.into()); + rewriter.insert_operation(ctx, const_op.get_operation()); + const_op.get_operation().deref(ctx).get_result(0) + }; + + let alloca_op = llvm::AllocaOp::new(ctx, llvm_array_ty.into(), one_val); + rewriter.insert_operation(ctx, alloca_op.get_operation()); + if let Some(align) = abi_align { + llvm_export::ops::set_op_alignment(ctx, alloca_op.get_operation(), align as u32); + } + let array_ptr = alloca_op.get_operation().deref(ctx).get_result(0); + + let store_op = llvm::StoreOp::new(ctx, array_val, array_ptr); + rewriter.insert_operation(ctx, store_op.get_operation()); + if let Some(align) = abi_align { + llvm_export::ops::set_op_alignment(ctx, store_op.get_operation(), align as u32); + } + + use llvm_export::ops::GepIndex; + let gep_indices = vec![GepIndex::Constant(0), GepIndex::Value(index_val)]; + let gep_op = llvm::GetElementPtrOp::new(ctx, array_ptr, gep_indices, llvm_array_ty.into()); + rewriter.insert_operation(ctx, gep_op.get_operation()); + let element_ptr = gep_op.get_operation().deref(ctx).get_result(0); + + let load_op = llvm::LoadOp::new(ctx, element_ptr, llvm_element_ty); + rewriter.insert_operation(ctx, load_op.get_operation()); + if let Some(align) = abi_align { + llvm_export::ops::set_op_alignment(ctx, load_op.get_operation(), align as u32); + } + rewriter.replace_operation(ctx, op, load_op.get_operation()); + + Ok(()) +} + +#[cfg(test)] +// Tests build kinded fixture types directly; production minting lives in mir-importer's facts.rs. +#[allow(clippy::disallowed_methods)] +mod tests { + use super::*; + use crate::convert::ops::test_util::*; + + use dialect_mir::ops as mir; + use dialect_mir::types::{MirArrayType, MirStructType, MirTupleType}; + + use pliron::builtin::attributes::IntegerAttr; + use pliron::builtin::types::{IntegerType, Signedness}; + use pliron::r#type::TypeHandle; + + fn over_aligned_tuple_ty(ctx: &mut Context) -> TypeHandle { + let byte: TypeHandle = IntegerType::get(ctx, 8, Signedness::Unsigned).into(); + let marker: TypeHandle = MirStructType::get_with_full_layout( + ctx, + "Align32".into(), + vec![], + vec![], + vec![], + vec![], + 0, + 32, + ) + .into(); + MirTupleType::get_with_layout(ctx, vec![marker, byte], vec![0, 1], vec![0, 0], 32, 32) + .into() + } + + fn lower_array_extract_case( + ctx: &mut Context, + array_size: u64, + divisor: Option, + ) -> Ptr { + let element_type: TypeHandle = IntegerType::get(ctx, 32, Signedness::Unsigned).into(); + let index_type = IntegerType::get(ctx, 64, Signedness::Unsigned); + let index_handle: TypeHandle = index_type.into(); + let array_type: TypeHandle = MirArrayType::get(ctx, element_type, array_size).into(); + + let (module, block) = build_kernel(ctx, vec![index_handle], vec![element_type]); + let raw_index = block.deref(ctx).get_argument(0); + + let undef = mir::MirUndefOp::new(ctx, array_type); + undef.get_operation().insert_at_back(block, ctx); + let array = undef.get_operation().deref(ctx).get_result(0); + + let index = if let Some(divisor) = divisor { + let constant = Operation::new( + ctx, + mir::MirConstantOp::get_concrete_op_info(), + vec![index_handle], + vec![], + vec![], + 0, + ); + mir::MirConstantOp::new(constant).set_attr_value( + ctx, + IntegerAttr::new( + index_type, + APInt::from_u64(divisor, NonZeroUsize::new(64).unwrap()), + ), + ); + constant.insert_at_back(block, ctx); + let divisor_value = constant.deref(ctx).get_result(0); + + let rem = Operation::new( + ctx, + mir::MirRemOp::get_concrete_op_info(), + vec![index_handle], + vec![raw_index, divisor_value], + vec![], + 0, + ); + rem.insert_at_back(block, ctx); + rem.deref(ctx).get_result(0) + } else { + raw_index + }; + + let extract = Operation::new( + ctx, + mir::MirExtractArrayElementOp::get_concrete_op_info(), + vec![element_type], + vec![array, index], + vec![], + 0, + ); + extract.insert_at_back(block, ctx); + let result = extract.deref(ctx).get_result(0); + append_mir_return(ctx, block, vec![result]); + + crate::lower_mir_to_llvm(ctx, module).expect("lowering failed"); + module + } + + fn assert_array_extract_memory_fallback(ctx: &Context, module: Ptr) { + let body = kernel_blocks(ctx, module); + assert_eq!(count_ops::(ctx, &body), 1); + assert_eq!(count_ops::(ctx, &body), 1); + assert_eq!(count_ops::(ctx, &body), 1); + assert_eq!(count_ops::(ctx, &body), 1); + assert_eq!(count_ops::(ctx, &body), 0); + } + + #[test] + fn bounded_urem_array_extract_stays_in_ssa() { + let mut ctx = make_ctx(); + let module = lower_array_extract_case(&mut ctx, 3, Some(3)); + let body = kernel_blocks(&ctx, module); + + assert_eq!(count_ops::(&ctx, &body), 3); + assert_eq!(count_ops::(&ctx, &body), 2); + assert_eq!(count_ops::(&ctx, &body), 2); + assert_eq!(count_ops::(&ctx, &body), 0); + assert_eq!(count_ops::(&ctx, &body), 0); + assert_eq!(count_ops::(&ctx, &body), 0); + assert_eq!(count_ops::(&ctx, &body), 0); + } + + #[test] + fn unbounded_array_extract_keeps_memory_fallback() { + let mut ctx = make_ctx(); + let module = lower_array_extract_case(&mut ctx, 3, None); + assert_array_extract_memory_fallback(&ctx, module); + } + + #[test] + fn oversized_urem_array_extract_keeps_memory_fallback() { + let mut ctx = make_ctx(); + let module = lower_array_extract_case(&mut ctx, 17, Some(17)); + assert_array_extract_memory_fallback(&ctx, module); + } + + #[test] + fn dynamic_array_extract_preserves_recursive_element_alignment() { + let mut ctx = make_ctx(); + let tuple_ty = over_aligned_tuple_ty(&mut ctx); + let inner: TypeHandle = MirArrayType::get(&mut ctx, tuple_ty, 2).into(); + let outer: TypeHandle = MirArrayType::get(&mut ctx, inner, 3).into(); + let index_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![index_ty], vec![]); + let index = block.deref(&ctx).get_argument(0); + + let undef = mir::MirUndefOp::new(&mut ctx, outer); + undef.get_operation().insert_at_back(block, &ctx); + let array = undef.get_operation().deref(&ctx).get_result(0); + let extract = Operation::new( + &mut ctx, + mir::MirExtractArrayElementOp::get_concrete_op_info(), + vec![inner], + vec![array, index], + vec![], + 0, + ); + extract.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + let alloca = find_first::(&ctx, &body).expect("expected llvm.alloca"); + let store = find_first::(&ctx, &body).expect("expected llvm.store"); + let load = find_first::(&ctx, &body).expect("expected llvm.load"); + for memory_op in [ + alloca.get_operation(), + store.get_operation(), + load.get_operation(), + ] { + assert_eq!(llvm_export::ops::op_alignment(&ctx, memory_op), Some(32)); + } + } +} diff --git a/crates/mir-lower/src/convert/ops/aggregate/common.rs b/crates/mir-lower/src/convert/ops/aggregate/common.rs new file mode 100644 index 0000000000..5ac33287ee --- /dev/null +++ b/crates/mir-lower/src/convert/ops/aggregate/common.rs @@ -0,0 +1,99 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use llvm_export::ops as llvm; +use pliron::builtin::types::{IntegerType, Signedness}; +use pliron::context::Context; +use pliron::irbuild::dialect_conversion::DialectConversionRewriter; +use pliron::irbuild::inserter::Inserter; +use pliron::op::Op; +use pliron::r#type::TypeHandle; +use pliron::utils::apint::APInt; +use pliron::value::Value; +use std::num::NonZeroUsize; + +pub(super) fn anyhow_to_pliron(e: anyhow::Error) -> pliron::result::Error { + pliron::input_error_noloc!("{e}") +} + +/// Copy an enum value into a fresh stack slot and return the pointer. +/// +/// This is how we reach a payload field that has no struct slot of its +/// own (its bytes are shared with a different-typed field of another +/// variant): once the value sits in memory, a byte-precise pointer can +/// read or write any part of it, no struct field needed. +/// +/// The slot is marked with the enum's real (rustc) alignment. The struct +/// type alone can look under-aligned: `{ i8, [7 x i8] }` says "align 1" +/// to LLVM, while Rust may require 8. +/// +/// The alloca lands at the use site, same as +/// [`convert_extract_array_element`](super::array_extract::convert_extract_array_element); +/// the standard `opt -O2` run (SROA) +/// removes it again. Hoisting these into the function's entry block is a +/// known follow-up for the unoptimized (`CUDA_OXIDE_NO_OPT=1`) path. +pub(super) fn spill_enum_value( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + enum_val: Value, + llvm_struct_ty: TypeHandle, + abi_align: u64, +) -> Value { + let i64_ty = IntegerType::get(ctx, 64, Signedness::Signless); + let one_apint = APInt::from_i64(1, NonZeroUsize::new(64).unwrap()); + let one_attr = pliron::builtin::attributes::IntegerAttr::new(i64_ty, one_apint); + let one_const = llvm::ConstantOp::new(ctx, one_attr.into()); + rewriter.insert_operation(ctx, one_const.get_operation()); + let one_val = one_const.get_operation().deref(ctx).get_result(0); + + let alloca_op = llvm::AllocaOp::new(ctx, llvm_struct_ty, one_val); + rewriter.insert_operation(ctx, alloca_op.get_operation()); + if abi_align > 0 { + llvm_export::ops::set_op_alignment(ctx, alloca_op.get_operation(), abi_align as u32); + } + let slot_ptr = alloca_op.get_operation().deref(ctx).get_result(0); + + let store_op = llvm::StoreOp::new(ctx, enum_val, slot_ptr); + rewriter.insert_operation(ctx, store_op.get_operation()); + if abi_align > 0 { + llvm_export::ops::set_op_alignment(ctx, store_op.get_operation(), abi_align as u32); + } + slot_ptr +} + +/// Pointer to `base + offset` bytes (`getelementptr i8, ptr base, offset`). +/// +/// Used whenever rustc's physical byte offset cannot be represented faithfully +/// by a typed LLVM aggregate GEP, as with overlapping enum payloads or packed +/// struct fields. +pub(super) fn byte_offset_gep( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + base: Value, + offset: u64, +) -> Value { + use llvm_export::ops::GepIndex; + let i8_ty: TypeHandle = IntegerType::get(ctx, 8, Signedness::Signless).into(); + let offset = emit_integer_constant(ctx, rewriter, 64, u128::from(offset)); + let gep_op = llvm::GetElementPtrOp::new(ctx, base, vec![GepIndex::Value(offset)], i8_ty); + rewriter.insert_operation(ctx, gep_op.get_operation()); + gep_op.get_operation().deref(ctx).get_result(0) +} + +pub(super) fn emit_integer_constant( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + width: u32, + bits: u128, +) -> Value { + let ty = IntegerType::get(ctx, width, Signedness::Signless); + let attr = pliron::builtin::attributes::IntegerAttr::new( + ty, + APInt::from_u128(bits, NonZeroUsize::new(width as usize).unwrap()), + ); + let op = llvm::ConstantOp::new(ctx, attr.into()); + rewriter.insert_operation(ctx, op.get_operation()); + op.get_operation().deref(ctx).get_result(0) +} diff --git a/crates/mir-lower/src/convert/ops/aggregate/construct.rs b/crates/mir-lower/src/convert/ops/aggregate/construct.rs new file mode 100644 index 0000000000..cc377cff49 --- /dev/null +++ b/crates/mir-lower/src/convert/ops/aggregate/construct.rs @@ -0,0 +1,1001 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use super::common::anyhow_to_pliron; +use crate::convert::types::{ + StructLayoutInfo, build_struct_slot_map, convert_type, + llvm_packed_struct_contains_pointer_in_address_space, make_slice_struct, + packed_shared_internal_abi_info, +}; +use dialect_mir::types::{ + MirArrayType, MirDisjointSliceType, MirSliceType, MirStructType, MirTupleType, +}; +use llvm_export::ops as llvm; +use llvm_export::types as llvm_types; +use pliron::context::{Context, Ptr}; +use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; +use pliron::irbuild::inserter::Inserter; +use pliron::irbuild::rewriter::Rewriter; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::result::Result; +use pliron::r#type::Typed; + +/// Convert `mir.construct_struct` to a chain of `llvm.insertvalue` operations. +/// +/// Builds a struct by: +/// 1. Creating an `undef` value of the lowered struct type +/// 2. Inserting each operand at the LLVM slot its field landed in +/// +/// Operand order matches field order in the struct type (declaration order). +/// The LLVM struct type and the slot of each field both come from +/// [`build_struct_slot_map`], so the insert indices skip `[N x i8]` padding +/// slots exactly the way the type converter laid them out. ZST fields +/// (e.g. PhantomData) have no slot and are skipped. +pub(crate) fn convert_construct_struct( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + _operands_info: &OperandsInfo, +) -> Result<()> { + let (result_ty, operands) = { + let mir_op = op.deref(ctx); + let result_ty = mir_op.get_result(0).get_type(ctx); + let operands: Vec<_> = mir_op.operands().collect(); + (result_ty, operands) + }; + + let layout = { + let ty_ref = result_ty.deref(ctx); + match ty_ref.downcast_ref::() { + Some(s) => StructLayoutInfo::of_struct(s), + None => { + return pliron::input_err_noloc!( + "MirConstructStructOp result type must be MirStructType" + ); + } + } + }; + + if operands.len() != layout.field_types.len() { + return pliron::input_err_noloc!( + "construct_struct has {} operands for a struct with {} fields", + operands.len(), + layout.field_types.len() + ); + } + + let map = build_struct_slot_map(ctx, &layout).map_err(anyhow_to_pliron)?; + + // A divergent rustc layout is constructible by value only when the slot + // map proved that a sequential LLVM packed struct reproduces every byte. + // Overlapping/union-like legacy struct models remain unrepresentable. + if !map.by_value_layout_faithful { + return pliron::input_err_noloc!( + "constructing a struct whose rustc layout cannot be represented by an LLVM \ + struct value is not supported; keep the value behind a pointer and access \ + fields through their byte-accurate address path" + ); + } + + // A packed struct containing AS3 remains target-dependent as a physical + // memory image. The internal-call ABI exception is safe to construct in SSA + // because its return boundary recursively rebuilds every supported AS3 leaf + // into a target-stable generic-pointer carrier before the aggregate crosses + // the function ABI. + if llvm_packed_struct_contains_pointer_in_address_space( + ctx, + map.llvm_struct_ty, + llvm_types::address_space::SHARED, + ) && packed_shared_internal_abi_info(ctx, result_ty) + .map_err(anyhow_to_pliron)? + .is_none() + { + return pliron::input_err_noloc!( + "constructing a packed aggregate containing a shared-memory pointer by value is \ + target-mode dependent and is not yet supported" + ); + } + + let undef_op = llvm::UndefOp::new(ctx, map.llvm_struct_ty); + rewriter.insert_operation(ctx, undef_op.get_operation()); + let mut current_struct = undef_op.get_operation().deref(ctx).get_result(0); + + let mut last_insert: Option> = None; + // Walk in memory order so the insertvalue chain ascends slot indices. + for &decl_idx in &layout.mem_to_decl { + let Some(slot) = map.decl_to_llvm[decl_idx] else { + continue; // ZST field: no slot in the LLVM struct. + }; + + let insert_op = + llvm::InsertValueOp::new(ctx, current_struct, operands[decl_idx], vec![slot]); + rewriter.insert_operation(ctx, insert_op.get_operation()); + current_struct = insert_op.get_operation().deref(ctx).get_result(0); + last_insert = Some(insert_op.get_operation()); + } + + match last_insert { + Some(last_op) => rewriter.replace_operation(ctx, op, last_op), + None => rewriter.replace_operation(ctx, op, undef_op.get_operation()), + } + + Ok(()) +} + +/// Convert `mir.construct_tuple` to a chain of `llvm.insertvalue` operations. +/// +/// Tuples are represented as LLVM structs. Same construction pattern as +/// structs, and like structs the element slots come from +/// [`build_struct_slot_map`] (identity order, no padding; ZST elements are +/// stripped and skipped). +pub(crate) fn convert_construct_tuple( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + _operands_info: &OperandsInfo, +) -> Result<()> { + let (result_ty, operands) = { + let mir_op = op.deref(ctx); + let result_ty = mir_op.get_result(0).get_type(ctx); + let operands: Vec<_> = mir_op.operands().collect(); + (result_ty, operands) + }; + + let layout = { + let ty_ref = result_ty.deref(ctx); + match ty_ref.downcast_ref::() { + Some(t) => StructLayoutInfo::of_tuple(t), + None => { + return pliron::input_err_noloc!( + "MirConstructTupleOp result type must be MirTupleType" + ); + } + } + }; + + if operands.len() != layout.field_types.len() { + return pliron::input_err_noloc!( + "construct_tuple has {} operands for a tuple with {} elements", + operands.len(), + layout.field_types.len() + ); + } + + let map = build_struct_slot_map(ctx, &layout).map_err(anyhow_to_pliron)?; + + let undef_op = llvm::UndefOp::new(ctx, map.llvm_struct_ty); + rewriter.insert_operation(ctx, undef_op.get_operation()); + let mut current_tuple = undef_op.get_operation().deref(ctx).get_result(0); + + let mut last_insert: Option> = None; + for (mir_idx, operand) in operands.iter().enumerate() { + let Some(slot) = map.decl_to_llvm[mir_idx] else { + continue; // ZST element: no slot in the LLVM struct. + }; + + let insert_op = llvm::InsertValueOp::new(ctx, current_tuple, *operand, vec![slot]); + rewriter.insert_operation(ctx, insert_op.get_operation()); + current_tuple = insert_op.get_operation().deref(ctx).get_result(0); + last_insert = Some(insert_op.get_operation()); + } + + match last_insert { + Some(last_op) => rewriter.replace_operation(ctx, op, last_op), + None => rewriter.replace_operation(ctx, op, undef_op.get_operation()), + } + + Ok(()) +} + +/// Convert `mir.construct_slice` to `llvm.undef` + two `llvm.insertvalue`s. +/// +/// `MirSliceType` lowers to the `{ ptr, i64 }` fat-pointer struct, where +/// field 0 is the data pointer and field 1 is the element count by +/// construction (the same layout the entry prologue's `reconstruct_slice` +/// and the Unsize cast path build). `MirSliceType::kind` is intentionally +/// semantic-only: all reference/raw-pointer kinds share this physical layout +/// and no LLVM alias metadata is inferred here. +pub(crate) fn convert_construct_slice( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + _operands_info: &OperandsInfo, +) -> Result<()> { + let (result_ty, data_val, len_val) = { + let mir_op = op.deref(ctx); + ( + mir_op.get_result(0).get_type(ctx), + mir_op.get_operand(0), + mir_op.get_operand(1), + ) + }; + + if !result_ty.deref(ctx).is::() { + return pliron::input_err_noloc!("MirConstructSliceOp result type must be MirSliceType"); + } + + let slice_struct_ty = make_slice_struct(ctx); + + let undef_op = llvm::UndefOp::new(ctx, slice_struct_ty); + rewriter.insert_operation(ctx, undef_op.get_operation()); + let undef_val = undef_op.get_operation().deref(ctx).get_result(0); + + let insert_ptr = llvm::InsertValueOp::new(ctx, undef_val, data_val, vec![0]); + rewriter.insert_operation(ctx, insert_ptr.get_operation()); + let with_ptr = insert_ptr.get_operation().deref(ctx).get_result(0); + + let insert_len = llvm::InsertValueOp::new(ctx, with_ptr, len_val, vec![1]); + rewriter.insert_operation(ctx, insert_len.get_operation()); + + rewriter.replace_operation(ctx, op, insert_len.get_operation()); + + Ok(()) +} + +/// Convert `mir.construct_disjoint_slice` to a chain of `llvm.insertvalue` +/// operations. +/// +/// The same shape as [`convert_construct_slice`], one insert longer per +/// runtime layout word the index space carries. The struct type comes from +/// [`make_disjoint_slice_struct`], the same constructor the type converter +/// uses, so the field order here cannot drift from the lowered layout that +/// [`resolve_aggregate_slots`](super::fields::resolve_aggregate_slots) indexes identically. +pub(crate) fn convert_construct_disjoint_slice( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + _operands_info: &OperandsInfo, +) -> Result<()> { + let (result_ty, operands) = { + let mir_op = op.deref(ctx); + let operands: Vec<_> = mir_op.operands().collect(); + (mir_op.get_result(0).get_type(ctx), operands) + }; + + let space_tys = { + let ty_ref = result_ty.deref(ctx); + match ty_ref.downcast_ref::() { + Some(slice_ty) => slice_ty.space_types().to_vec(), + None => { + return pliron::input_err_noloc!( + "MirConstructDisjointSliceOp result type must be MirDisjointSliceType" + ); + } + } + }; + + // The verifier already pins the operand count to the result type. Check + // it here too: this pass runs on operations the verifier accepted, and a + // short operand list would otherwise index past the end below. + let expected_operands = 2 + space_tys.len(); + if operands.len() != expected_operands { + return pliron::input_err_noloc!( + "MirConstructDisjointSliceOp expects {expected_operands} operands, got {}", + operands.len() + ); + } + + let slice_struct_ty = crate::convert::types::make_disjoint_slice_struct(ctx, &space_tys) + .map_err(anyhow_to_pliron)?; + + let undef_op = llvm::UndefOp::new(ctx, slice_struct_ty); + rewriter.insert_operation(ctx, undef_op.get_operation()); + let mut current = undef_op.get_operation().deref(ctx).get_result(0); + let mut last_insert = undef_op.get_operation(); + + for (slot, operand) in operands.into_iter().enumerate() { + let insert = llvm::InsertValueOp::new(ctx, current, operand, vec![slot as u32]); + rewriter.insert_operation(ctx, insert.get_operation()); + current = insert.get_operation().deref(ctx).get_result(0); + last_insert = insert.get_operation(); + } + + rewriter.replace_operation(ctx, op, last_insert); + + Ok(()) +} + +/// Convert `mir.construct_array` to a chain of `llvm.insertvalue` operations. +/// +/// Arrays are represented as LLVM arrays. Same construction pattern as structs: +/// 1. Create `undef` of the array type +/// 2. Insert each element at its index +pub(crate) fn convert_construct_array( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + _operands_info: &OperandsInfo, +) -> Result<()> { + let (result_ty, operands) = { + let mir_op = op.deref(ctx); + let result_ty = mir_op.get_result(0).get_type(ctx); + let operands: Vec<_> = mir_op.operands().collect(); + (result_ty, operands) + }; + + let (element_ty, array_size) = { + let ty_ref = result_ty.deref(ctx); + match ty_ref.downcast_ref::() { + Some(a) => (a.element_type(), a.size()), + None => { + return pliron::input_err_noloc!( + "MirConstructArrayOp result type must be MirArrayType" + ); + } + } + }; + + let llvm_element_ty = convert_type(ctx, element_ty).map_err(anyhow_to_pliron)?; + let llvm_array_ty = llvm_export::types::ArrayType::get(ctx, llvm_element_ty, array_size); + + let undef_op = llvm::UndefOp::new(ctx, llvm_array_ty.into()); + rewriter.insert_operation(ctx, undef_op.get_operation()); + let mut current_array = undef_op.get_operation().deref(ctx).get_result(0); + + let mut last_insert: Option> = None; + for (i, operand) in operands.iter().enumerate() { + let insert_op = llvm::InsertValueOp::new(ctx, current_array, *operand, vec![i as u32]); + rewriter.insert_operation(ctx, insert_op.get_operation()); + current_array = insert_op.get_operation().deref(ctx).get_result(0); + last_insert = Some(insert_op.get_operation()); + } + + match last_insert { + Some(last_op) => rewriter.replace_operation(ctx, op, last_op), + None => rewriter.replace_operation(ctx, op, undef_op.get_operation()), + } + + Ok(()) +} + +#[cfg(test)] +// Tests build kinded fixture types directly; production minting lives in mir-importer's facts.rs. +#[allow(clippy::disallowed_methods)] +mod tests { + use super::*; + use crate::convert::ops::test_util::*; + + use dialect_mir::ops as mir; + use dialect_mir::types::{ + MirArrayType, MirPointerKind, MirPtrType, MirSliceType, MirStructType, + }; + use llvm_export::types as llvm_types; + + use pliron::common_traits::Verify; + + use crate::convert::types::llvm_type_size_align; + use pliron::builtin::types::{IntegerType, Signedness}; + use pliron::r#type::TypeHandle; + + use super::super::test_support::*; + + /// `mir.construct_slice` lowers to the canonical fat-pointer value: + /// `undef { ptr, i64 }`, then insert data pointer at slot 0 and length at slot 1. + #[test] + fn construct_slice_lowers_to_ptr_len_insert_values() { + let mut ctx = make_ctx(); + + let i8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let usize_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, i8_ty, false).into(); + let slice_ty: TypeHandle = MirSliceType::get(&mut ctx, i8_ty).into(); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![ptr_ty, usize_ty], vec![]); + let data_ptr = block.deref(&ctx).get_argument(0); + let len = block.deref(&ctx).get_argument(1); + + let op = Operation::new( + &mut ctx, + mir::MirConstructSliceOp::get_concrete_op_info(), + vec![slice_ty], + vec![data_ptr, len], + vec![], + 0, + ); + op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + let inserts = find_all::(&ctx, &body); + + assert_eq!( + insert_indices(&ctx, &inserts), + vec![vec![0], vec![1]], + "slice construction must insert data pointer at slot 0 and length at slot 1" + ); + let first_insert = inserts[0].get_operation(); + let second_insert = inserts[1].get_operation(); + assert_eq!( + first_insert.deref(&ctx).get_operand(1), + data_ptr, + "slice slot 0 must receive the original data pointer" + ); + assert_eq!( + second_insert.deref(&ctx).get_operand(0), + first_insert.deref(&ctx).get_result(0), + "the length insertion must consume the aggregate produced by the pointer insertion" + ); + assert_eq!( + second_insert.deref(&ctx).get_operand(1), + len, + "slice slot 1 must receive the original length" + ); + assert!( + inserts.iter().all(|insert| insert.verify(&ctx).is_ok()), + "both slice insertions must satisfy LLVM dialect verification" + ); + assert_eq!( + count_ops::(&ctx, &body), + 1, + "slice construction should start from one undef aggregate" + ); + } + + /// `mir.construct_disjoint_slice` lowers to the same insert chain as the + /// fat pointer for an index space with no runtime layout. + #[test] + fn construct_disjoint_slice_lowers_to_ptr_len_insert_values() { + let mut ctx = make_ctx(); + + let f32_ty: TypeHandle = pliron::builtin::types::FP32Type::get(&ctx).into(); + let usize_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let ptr_ty: TypeHandle = + MirPtrType::get_generic_with_kind(&mut ctx, f32_ty, true, MirPointerKind::RawMut) + .into(); + let slice_ty: TypeHandle = MirDisjointSliceType::get(&mut ctx, f32_ty).into(); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![ptr_ty, usize_ty], vec![]); + let data_ptr = block.deref(&ctx).get_argument(0); + let len = block.deref(&ctx).get_argument(1); + + let op = Operation::new( + &mut ctx, + mir::MirConstructDisjointSliceOp::get_concrete_op_info(), + vec![slice_ty], + vec![data_ptr, len], + vec![], + 0, + ); + op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + let inserts = find_all::(&ctx, &body); + assert_eq!( + insert_indices(&ctx, &inserts), + vec![vec![0], vec![1]], + "a space-free disjoint slice must insert pointer at slot 0 and length at slot 1" + ); + assert_eq!( + inserts[0].get_operation().deref(&ctx).get_operand(1), + data_ptr, + "slot 0 must receive the original data pointer" + ); + assert_eq!( + inserts[1].get_operation().deref(&ctx).get_operand(1), + len, + "slot 1 must receive the original length" + ); + assert_eq!( + count_ops::(&ctx, &body), + 1, + "construction should start from one undef aggregate" + ); + } + + /// The runtime row width is the third operand and must land in slot 2. + /// Writing it into the length slot, or dropping it, gives a slice whose + /// row width reads back as something else at every access site. + #[test] + fn construct_disjoint_slice_places_the_row_width_after_ptr_and_len() { + let mut ctx = make_ctx(); + + let f32_ty: TypeHandle = pliron::builtin::types::FP32Type::get(&ctx).into(); + let usize_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let width_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let ptr_ty: TypeHandle = + MirPtrType::get_generic_with_kind(&mut ctx, f32_ty, true, MirPointerKind::RawMut) + .into(); + let slice_ty: TypeHandle = + MirDisjointSliceType::get_with_space(&mut ctx, f32_ty, vec![width_ty]).into(); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![ptr_ty, usize_ty, width_ty], vec![]); + let data_ptr = block.deref(&ctx).get_argument(0); + let len = block.deref(&ctx).get_argument(1); + let width = block.deref(&ctx).get_argument(2); + + let op = Operation::new( + &mut ctx, + mir::MirConstructDisjointSliceOp::get_concrete_op_info(), + vec![slice_ty], + vec![data_ptr, len, width], + vec![], + 0, + ); + op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + let inserts = find_all::(&ctx, &body); + assert_eq!( + insert_indices(&ctx, &inserts), + vec![vec![0], vec![1], vec![2]], + "the row width must occupy slot 2, after the pointer and the length" + ); + assert_eq!( + inserts[2].get_operation().deref(&ctx).get_operand(1), + width, + "slot 2 must receive the row width operand, not the length" + ); + assert_eq!( + inserts[1].get_operation().deref(&ctx).get_operand(1), + len, + "slot 1 must still receive the length" + ); + } + + /// Explicit rustc layout must be respected: field `b` is at byte offset 8, + /// so the lowered LLVM struct has a padding slot between `a` and `b`. + /// The ZST marker field is stripped and receives no insert_value. + #[test] + fn construct_struct_uses_layout_slots_and_skips_zst() { + let mut ctx = make_ctx(); + + let i8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Signless).into(); + let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signless).into(); + let (struct_ty, zst_ty) = padded_struct_with_zst_ty(&mut ctx); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![i8_ty, i64_ty], vec![]); + let a = block.deref(&ctx).get_argument(0); + let b = block.deref(&ctx).get_argument(1); + let marker = append_empty_struct_value(&mut ctx, block, zst_ty); + + let op = Operation::new( + &mut ctx, + mir::MirConstructStructOp::get_concrete_op_info(), + vec![struct_ty], + vec![a, marker, b], + vec![], + 0, + ); + op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + let inserts = find_all::(&ctx, &body); + + assert_eq!( + insert_indices(&ctx, &inserts), + vec![vec![0], vec![2]], + "non-ZST fields must be inserted at their layout slots, skipping padding and ZSTs" + ); + let first_insert = inserts[0].get_operation(); + let second_insert = inserts[1].get_operation(); + assert_eq!( + first_insert.deref(&ctx).get_operand(1), + a, + "struct slot 0 must receive field `a`" + ); + assert_eq!( + second_insert.deref(&ctx).get_operand(0), + first_insert.deref(&ctx).get_result(0), + "field `b` must be inserted into the aggregate containing field `a`" + ); + assert_eq!( + second_insert.deref(&ctx).get_operand(1), + b, + "struct slot 2 must receive field `b`" + ); + assert!( + inserts.iter().all(|insert| insert.verify(&ctx).is_ok()), + "both struct insertions must satisfy LLVM dialect verification" + ); + } + + /// A by-value packed struct is represented by an LLVM packed struct, so + /// construction can use the ordinary insertvalue chain without changing + /// rustc's field offsets or total byte size. + #[test] + fn packed_struct_construction_by_value_uses_packed_layout() { + use dialect_mir::ops::MirConstructStructOp; + + let mut ctx = make_ctx(); + + let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let packed_ty: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "Packed".into(), + vec!["tag".into(), "value".into()], + vec![u8_ty, u32_ty], + vec![0, 1], + vec![0, 1], + 5, + 1, + ) + .into(); + + let (module, block) = build_kernel(&mut ctx, vec![u8_ty, u32_ty], vec![]); + let tag = block.deref(&ctx).get_argument(0); + let value = block.deref(&ctx).get_argument(1); + + let construct = Operation::new( + &mut ctx, + MirConstructStructOp::get_concrete_op_info(), + vec![packed_ty], + vec![tag, value], + vec![], + 0, + ); + construct.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module) + .expect("constructing a pointer-free packed struct by value must lower"); + + let body = kernel_blocks(&ctx, module); + let inserts = find_all::(&ctx, &body); + assert_eq!( + insert_indices(&ctx, &inserts), + vec![vec![0], vec![1]], + "packed construction must keep declaration fields in their packed LLVM slots" + ); + let result_ty = inserts + .last() + .expect("packed construction must emit insertvalue") + .get_operation() + .deref(&ctx) + .get_result(0) + .get_type(&ctx); + let result_ty_ref = result_ty.deref(&ctx); + let struct_ty = result_ty_ref + .downcast_ref::() + .expect("packed construction result must be an LLVM struct"); + assert_eq!(struct_ty.layout(), llvm_types::StructLayout::Packed); + assert_eq!(llvm_type_size_align(&ctx, result_ty), Some((5, 1))); + } + + #[test] + fn packed_struct_construction_with_one_direct_shared_pointer_stays_semantic_in_ssa() { + use dialect_mir::ops::MirConstructStructOp; + + let mut ctx = make_ctx(); + let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let pointee: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let shared_ty: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, false).into(); + let packed_ty: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "PackedShared".into(), + vec!["tag".into(), "ptr".into()], + vec![u8_ty, shared_ty], + vec![0, 1], + vec![0, 1], + 9, + 1, + ) + .into(); + + let (module, block) = build_kernel(&mut ctx, vec![u8_ty, shared_ty], vec![]); + let tag = block.deref(&ctx).get_argument(0); + let ptr = block.deref(&ctx).get_argument(1); + let construct = Operation::new( + &mut ctx, + MirConstructStructOp::get_concrete_op_info(), + vec![packed_ty], + vec![tag, ptr], + vec![], + 0, + ); + construct.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module) + .expect("one direct AS3 pointer in a packed SSA aggregate must lower"); + + let body = kernel_blocks(&ctx, module); + let inserts = find_all::(&ctx, &body); + assert_eq!( + insert_indices(&ctx, &inserts), + vec![vec![0], vec![1]], + "packed AS3 construction must preserve the semantic packed field slots" + ); + + let result_ty = inserts + .last() + .expect("packed AS3 construction must emit insertvalue") + .get_operation() + .deref(&ctx) + .get_result(0) + .get_type(&ctx); + let result_ty_ref = result_ty.deref(&ctx); + let struct_ty = result_ty_ref + .downcast_ref::() + .expect("packed AS3 construction result must be an LLVM struct"); + assert_eq!(struct_ty.layout(), llvm_types::StructLayout::Packed); + assert_eq!(llvm_type_size_align(&ctx, result_ty), Some((9, 1))); + + let pointer_ty = struct_ty.field_type(1); + let pointer_ty_ref = pointer_ty.deref(&ctx); + let pointer = pointer_ty_ref + .downcast_ref::() + .expect("second packed field must remain a pointer"); + assert_eq!( + pointer.address_space(), + llvm_types::address_space::SHARED, + "construction must keep the semantic pointer in AS3" + ); + assert_eq!( + count_ops::(&ctx, &body), + 0, + "construction itself must not genericize the shared pointer" + ); + } + + #[test] + fn packed_struct_construction_with_multiple_direct_shared_pointers_stays_semantic_in_ssa() { + use dialect_mir::ops::MirConstructStructOp; + + let mut ctx = make_ctx(); + let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let pointee: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let shared_ty: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, false).into(); + let packed_ty: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "PackedSharedPair".into(), + vec!["tag".into(), "left".into(), "right".into()], + vec![u8_ty, shared_ty, shared_ty], + vec![0, 1, 2], + vec![0, 1, 9], + 17, + 1, + ) + .into(); + + let (module, block) = build_kernel(&mut ctx, vec![u8_ty, shared_ty, shared_ty], vec![]); + let tag = block.deref(&ctx).get_argument(0); + let left = block.deref(&ctx).get_argument(1); + let right = block.deref(&ctx).get_argument(2); + let construct = Operation::new( + &mut ctx, + MirConstructStructOp::get_concrete_op_info(), + vec![packed_ty], + vec![tag, left, right], + vec![], + 0, + ); + construct.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module) + .expect("multiple direct AS3 leaves in a packed SSA aggregate must lower"); + + let body = kernel_blocks(&ctx, module); + let inserts = find_all::(&ctx, &body); + assert_eq!( + insert_indices(&ctx, &inserts), + vec![vec![0], vec![1], vec![2]], + "all direct shared leaves must remain in their semantic packed slots" + ); + + let result_ty = inserts + .last() + .expect("packed AS3 construction must emit insertvalue") + .get_operation() + .deref(&ctx) + .get_result(0) + .get_type(&ctx); + let result_ty_ref = result_ty.deref(&ctx); + let struct_ty = result_ty_ref + .downcast_ref::() + .expect("packed AS3 construction result must be an LLVM struct"); + assert_eq!(struct_ty.layout(), llvm_types::StructLayout::Packed); + assert_eq!(llvm_type_size_align(&ctx, result_ty), Some((17, 1))); + for slot in [1usize, 2] { + let pointer_ty = struct_ty.field_type(slot); + let pointer_ty_ref = pointer_ty.deref(&ctx); + let pointer = pointer_ty_ref + .downcast_ref::() + .expect("direct shared leaves must remain pointer-typed"); + assert_eq!(pointer.address_space(), llvm_types::address_space::SHARED); + } + assert_eq!( + count_ops::(&ctx, &body), + 0, + "semantic construction must not genericize any direct shared leaf" + ); + } + + #[test] + fn packed_struct_construction_with_nested_shared_pointers_stays_semantic_in_ssa() { + use dialect_mir::ops::MirConstructStructOp; + + let mut ctx = make_ctx(); + let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let pointee: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let shared_ty: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, false).into(); + let inner_ty: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "SharedPair".into(), + vec!["left".into(), "right".into()], + vec![shared_ty, shared_ty], + vec![0, 1], + vec![0, 8], + 16, + 8, + ) + .into(); + let packed_ty: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "PackedNestedShared".into(), + vec!["tag".into(), "pair".into()], + vec![u8_ty, inner_ty], + vec![0, 1], + vec![0, 1], + 17, + 1, + ) + .into(); + + let (module, block) = build_kernel(&mut ctx, vec![u8_ty, shared_ty, shared_ty], vec![]); + let tag = block.deref(&ctx).get_argument(0); + let left = block.deref(&ctx).get_argument(1); + let right = block.deref(&ctx).get_argument(2); + + let inner = Operation::new( + &mut ctx, + MirConstructStructOp::get_concrete_op_info(), + vec![inner_ty], + vec![left, right], + vec![], + 0, + ); + inner.insert_at_back(block, &ctx); + let inner_value = inner.deref(&ctx).get_result(0); + + let outer = Operation::new( + &mut ctx, + MirConstructStructOp::get_concrete_op_info(), + vec![packed_ty], + vec![tag, inner_value], + vec![], + 0, + ); + outer.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module) + .expect("nested AS3 leaves in a packed SSA aggregate must lower"); + + let body = kernel_blocks(&ctx, module); + assert_eq!( + count_ops::(&ctx, &body), + 0, + "semantic construction must keep nested shared leaves in AS3" + ); + + let inserts = find_all::(&ctx, &body); + let packed_result = inserts + .iter() + .filter_map(|insert| { + let result_ty = insert + .get_operation() + .deref(&ctx) + .get_result(0) + .get_type(&ctx); + let is_packed = result_ty + .deref(&ctx) + .downcast_ref::() + .is_some_and(|ty| ty.layout() == llvm_types::StructLayout::Packed); + is_packed.then_some(result_ty) + }) + .last() + .expect("outer packed construction must produce a packed LLVM struct"); + + assert_eq!(llvm_type_size_align(&ctx, packed_result), Some((17, 1))); + assert!( + llvm_packed_struct_contains_pointer_in_address_space( + &ctx, + packed_result, + llvm_types::address_space::SHARED, + ), + "the semantic packed value must retain its recursively nested AS3 leaves" + ); + } + + #[test] + fn packed_struct_construction_with_bounded_shared_pointer_array_stays_semantic_in_ssa() { + use dialect_mir::ops::{MirConstructArrayOp, MirConstructStructOp}; + + let mut ctx = make_ctx(); + let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let pointee: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let shared_ty: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, false).into(); + let array_ty: TypeHandle = MirArrayType::get(&mut ctx, shared_ty, 2).into(); + let packed_ty: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "PackedSharedArray".into(), + vec!["tag".into(), "ptrs".into()], + vec![u8_ty, array_ty], + vec![0, 1], + vec![0, 1], + 17, + 1, + ) + .into(); + + let (module, block) = build_kernel(&mut ctx, vec![u8_ty, shared_ty, shared_ty], vec![]); + let tag = block.deref(&ctx).get_argument(0); + let left = block.deref(&ctx).get_argument(1); + let right = block.deref(&ctx).get_argument(2); + + let array = Operation::new( + &mut ctx, + MirConstructArrayOp::get_concrete_op_info(), + vec![array_ty], + vec![left, right], + vec![], + 0, + ); + array.insert_at_back(block, &ctx); + let array_value = array.deref(&ctx).get_result(0); + + let outer = Operation::new( + &mut ctx, + MirConstructStructOp::get_concrete_op_info(), + vec![packed_ty], + vec![tag, array_value], + vec![], + 0, + ); + outer.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module) + .expect("a bounded AS3 array in a packed SSA aggregate must lower"); + + let body = kernel_blocks(&ctx, module); + assert_eq!( + count_ops::(&ctx, &body), + 0, + "semantic construction must keep array elements in AS3" + ); + + let inserts = find_all::(&ctx, &body); + let packed_result = inserts + .iter() + .filter_map(|insert| { + let result_ty = insert + .get_operation() + .deref(&ctx) + .get_result(0) + .get_type(&ctx); + let is_packed = result_ty + .deref(&ctx) + .downcast_ref::() + .is_some_and(|ty| ty.layout() == llvm_types::StructLayout::Packed); + is_packed.then_some(result_ty) + }) + .last() + .expect("outer packed array construction must produce a packed LLVM struct"); + assert_eq!(llvm_type_size_align(&ctx, packed_result), Some((17, 1))); + assert!( + llvm_packed_struct_contains_pointer_in_address_space( + &ctx, + packed_result, + llvm_types::address_space::SHARED, + ), + "the semantic packed value must retain AS3 array elements" + ); + } +} diff --git a/crates/mir-lower/src/convert/ops/aggregate/enum_layout.rs b/crates/mir-lower/src/convert/ops/aggregate/enum_layout.rs new file mode 100644 index 0000000000..179ce032b6 --- /dev/null +++ b/crates/mir-lower/src/convert/ops/aggregate/enum_layout.rs @@ -0,0 +1,220 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use super::common::emit_integer_constant; +use crate::convert::enum_payload_storage::coerce_enum_payload_value; +use crate::convert::types::{llvm_byte_faithful_twin, llvm_type_contains_i1}; +use dialect_mir::types::{EnumCarrierKind, EnumLayoutKind, MirEnumType}; +use llvm_export::op_interfaces::{CastOpInterface, CastOpWithNNegInterface}; +use llvm_export::ops as llvm; +use llvm_export::types as llvm_types; +use pliron::builtin::types::{IntegerType, Signedness}; +use pliron::context::Context; +use pliron::irbuild::dialect_conversion::DialectConversionRewriter; +use pliron::irbuild::inserter::Inserter; +use pliron::op::Op; +use pliron::result::Result; +use pliron::r#type::{TypeHandle, Typed}; +use pliron::value::Value; + +/// The physical carrier write required to select one source variant. +/// `None` is a real semantic result for the untagged niche variant and for +/// rustc's single inhabited variant; it must never be replaced by a guessed +/// memory write. +pub(super) fn enum_carrier_bits_for_variant( + enum_ty: &MirEnumType, + variant: usize, +) -> std::result::Result, String> { + if variant >= enum_ty.variant_count() || enum_ty.variant_is_inhabited(variant) != Some(true) { + return Err(format!( + "cannot select uninhabited or missing variant {} of '{}'", + variant, + enum_ty.name() + )); + } + + match enum_ty.layout_kind { + EnumLayoutKind::Direct => enum_ty + .variant_discriminants + .get(variant) + .copied() + .map(|bits| Some(u128::from(bits))) + .ok_or_else(|| format!("variant {} has no declared discriminant", variant)), + EnumLayoutKind::Niche => { + // This check deliberately comes before the encoded range check: + // rustc permits the untagged variant index to lie inside that + // range. Its range position is a dead niche value; selecting the + // actual untagged variant remains a no-op. + if variant == enum_ty.untagged_variant as usize { + return Ok(None); + } + if !(enum_ty.niche_variant_start as usize..=enum_ty.niche_variant_end as usize) + .contains(&variant) + { + return Err(format!( + "inhabited variant {} is not representable by niche layout of '{}'", + variant, + enum_ty.name() + )); + } + let offset = (variant as u128) - u128::from(enum_ty.niche_variant_start); + let mut bits = enum_ty.niche_start().wrapping_add(offset); + if enum_ty.carrier_width < 128 { + bits &= (1u128 << enum_ty.carrier_width) - 1; + } + Ok(Some(bits)) + } + EnumLayoutKind::Single if variant == enum_ty.single_variant as usize => Ok(None), + EnumLayoutKind::Single => Err(format!( + "variant {} is not the single inhabited variant of '{}'", + variant, + enum_ty.name() + )), + EnumLayoutKind::Empty => Err(format!("enum '{}' is uninhabited", enum_ty.name())), + EnumLayoutKind::Unknown => Err(format!( + "enum '{}' has unknown physical layout", + enum_ty.name() + )), + } +} + +pub(super) fn emit_carrier_constant( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + enum_ty: &MirEnumType, + carrier_ty: TypeHandle, + bits: u128, +) -> Result { + let integer = emit_integer_constant(ctx, rewriter, enum_ty.carrier_width, bits); + match enum_ty.carrier_kind { + EnumCarrierKind::Integer => Ok(integer), + EnumCarrierKind::Pointer => { + let cast = llvm::IntToPtrOp::new(ctx, integer, carrier_ty); + rewriter.insert_operation(ctx, cast.get_operation()); + Ok(cast.get_operation().deref(ctx).get_result(0)) + } + _ => pliron::input_err_noloc!("enum carrier constant requested without a carrier"), + } +} + +/// Convert a value to its byte-faithful storage twin: every `i1` leaf is +/// zero-extended to its canonical `i8` memory byte, recursively through +/// structs and arrays. Values without `i1` storage pass through unchanged. +/// +/// This is the value-level half of [`llvm_byte_faithful_twin`]: the enum +/// slot map claims twin-typed storage for bool-bearing payloads, and this +/// produces the twin-typed value the store into that storage needs. +pub(super) fn canonicalize_bool_value_bytes( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + value: Value, +) -> Result { + let ty = value.get_type(ctx); + if !llvm_type_contains_i1(ctx, ty) { + return Ok(value); + } + let is_scalar_i1 = ty + .deref(ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 1); + if is_scalar_i1 { + let byte_ty: TypeHandle = IntegerType::get(ctx, 8, Signedness::Signless).into(); + let zext = llvm::ZExtOp::new_with_nneg(ctx, value, byte_ty, false); + rewriter.insert_operation(ctx, zext.get_operation()); + return Ok(zext.get_operation().deref(ctx).get_result(0)); + } + let Some(twin) = llvm_byte_faithful_twin(ctx, ty) else { + return pliron::input_err_noloc!( + "enum construction: bool storage in this value's shape cannot be canonicalized" + ); + }; + let element_count = { + let ty_ref = ty.deref(ctx); + if let Some(struct_ty) = ty_ref.downcast_ref::() { + struct_ty.fields().count() as u64 + } else if let Some(array_ty) = ty_ref.downcast_ref::() { + array_ty.size() + } else { + return pliron::input_err_noloc!( + "enum construction: unexpected container for bool storage canonicalization" + ); + } + }; + let undef_op = llvm::UndefOp::new(ctx, twin); + rewriter.insert_operation(ctx, undef_op.get_operation()); + let mut current = undef_op.get_operation().deref(ctx).get_result(0); + for index in 0..element_count { + let extract_op = llvm::ExtractValueOp::new(ctx, value, vec![index as u32])?; + rewriter.insert_operation(ctx, extract_op.get_operation()); + let element = extract_op.get_operation().deref(ctx).get_result(0); + let converted = canonicalize_bool_value_bytes(ctx, rewriter, element)?; + let insert_op = llvm::InsertValueOp::new(ctx, current, converted, vec![index as u32]); + rewriter.insert_operation(ctx, insert_op.get_operation()); + current = insert_op.get_operation().deref(ctx).get_result(0); + } + Ok(current) +} + +/// Adapt a semantic enum payload value to or from its physical storage type. +/// +/// Shared pointer leaves are converted through CUDA generic space recursively +/// through struct/tuple payloads. Bool leaves are canonicalized separately on +/// construction and narrowed recursively on extraction. +pub(super) fn coerce_enum_payload_storage( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + value: Value, + target_ty: TypeHandle, +) -> Result { + coerce_enum_payload_value(ctx, rewriter, value, target_ty) +} + +#[cfg(test)] +// Tests build kinded fixture types directly; production minting lives in mir-importer's facts.rs. +#[allow(clippy::disallowed_methods)] +mod tests { + use super::*; + use crate::convert::ops::test_util::*; + + use super::super::test_support::*; + + #[test] + fn niche_encoding_handles_untagged_inside_range_and_u128_wrap() { + let mut ctx = make_ctx(); + let inside = unit_niche_enum( + &mut ctx, + (EnumCarrierKind::Integer, 8, 0), + 42, + 0..=1, + 1, + vec![1, 1], + ); + let inside_ref = inside.deref(&ctx); + let inside_enum = inside_ref.downcast_ref::().unwrap(); + assert_eq!(enum_carrier_bits_for_variant(inside_enum, 0), Ok(Some(42))); + assert_eq!( + enum_carrier_bits_for_variant(inside_enum, 1), + Ok(None), + "the untagged variant is a no-op even when its index is in the niche range" + ); + drop(inside_ref); + + let wrapping = unit_niche_enum( + &mut ctx, + (EnumCarrierKind::Integer, 128, 0), + u128::MAX, + 0..=1, + 0, + vec![1, 1], + ); + let wrapping_ref = wrapping.deref(&ctx); + let wrapping_enum = wrapping_ref.downcast_ref::().unwrap(); + assert_eq!( + enum_carrier_bits_for_variant(wrapping_enum, 1), + Ok(Some(0)), + "niche arithmetic must wrap across the full u128 carrier" + ); + } +} diff --git a/crates/mir-lower/src/convert/ops/aggregate/enums.rs b/crates/mir-lower/src/convert/ops/aggregate/enums.rs new file mode 100644 index 0000000000..99a33c080e --- /dev/null +++ b/crates/mir-lower/src/convert/ops/aggregate/enums.rs @@ -0,0 +1,1704 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use super::common::{anyhow_to_pliron, byte_offset_gep, emit_integer_constant, spill_enum_value}; +use super::enum_layout::{ + canonicalize_bool_value_bytes, coerce_enum_payload_storage, emit_carrier_constant, + enum_carrier_bits_for_variant, +}; +use crate::convert::enum_payload_storage::enum_payload_storage_type; +use crate::convert::types::{EnumSlotMap, build_enum_slot_map, convert_type, is_zero_sized_type}; +use dialect_mir::ops::{MirConstructEnumOp, MirEnumPayloadOp, MirSetDiscriminantOp}; +use dialect_mir::types::{EnumCarrierKind, EnumLayoutKind, MirEnumType}; +use llvm_export::attributes::{ICmpPredicateAttr, IntegerOverflowFlagsAttr}; +use llvm_export::op_interfaces::{ + CastOpInterface, CastOpWithNNegInterface, IntBinArithOpWithOverflowFlag, +}; +use llvm_export::ops as llvm; +use llvm_export::types as llvm_types; +use pliron::builtin::types::{IntegerType, Signedness}; +use pliron::context::{Context, Ptr}; +use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; +use pliron::irbuild::inserter::Inserter; +use pliron::irbuild::rewriter::Rewriter; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::result::Result; +use pliron::r#type::{TypeHandle, Typed}; +use pliron::value::Value; + +/// Convert `mir.construct_enum` (e.g. `E::A(x)`) to LLVM operations. +/// +/// Builds the enum value slot by slot, taking every index from +/// [`build_enum_slot_map`] (indexes are never computed by hand here): +/// +/// 1. Put the variant's declared discriminant VALUE into the tag slot. +/// 2. `insertvalue` each payload field that owns a struct slot. +/// 3. If some field has no slot (its bytes are shared with a +/// different-typed field of another variant), finish through memory: +/// copy the value to a stack slot, store that field at its byte +/// position, and load the completed enum back. +pub(crate) fn convert_construct_enum( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + _operands_info: &OperandsInfo, +) -> Result<()> { + let (result_ty, operands, variant_index) = { + let mir_op = op.deref(ctx); + let result_ty = mir_op.get_result(0).get_type(ctx); + let operands: Vec<_> = mir_op.operands().collect(); + + let enum_op = MirConstructEnumOp::new(op); + let variant_index = enum_op + .get_attr_construct_enum_variant_index(ctx) + .map(|attr| attr.0 as usize) + .unwrap_or(0); + + (result_ty, operands, variant_index) + }; + + let enum_ty: MirEnumType = { + let ty_ref = result_ty.deref(ctx); + match ty_ref.downcast_ref::() { + Some(e) => e.clone(), + None => { + return pliron::input_err_noloc!( + "MirConstructEnumOp result type must be MirEnumType" + ); + } + } + }; + + // Build the value as the SAME struct type the type converter + // produces everywhere else (block args, loads, allocas, ...). Taking + // both the type and the indices from one slot map is what keeps them + // in agreement. Filler slots are simply never written. + let slot_map = build_enum_slot_map(ctx, result_ty).map_err(anyhow_to_pliron)?; + let llvm_struct_ty = slot_map.llvm_struct_ty; + + let undef_op = llvm::UndefOp::new(ctx, llvm_struct_ty); + rewriter.insert_operation(ctx, undef_op.get_operation()); + let mut current_struct = undef_op.get_operation().deref(ctx).get_result(0); + let mut last_op = undef_op.get_operation(); + + if let Some(bits) = enum_carrier_bits_for_variant(&enum_ty, variant_index) + .map_err(|error| pliron::input_error_noloc!("MirConstructEnumOp: {error}"))? + { + let carrier_slot = slot_map.carrier_slot.ok_or_else(|| { + pliron::input_error_noloc!("MirConstructEnumOp requires a physical carrier slot") + })?; + let carrier_ty = slot_map.carrier_llvm_ty.ok_or_else(|| { + pliron::input_error_noloc!("MirConstructEnumOp requires a physical carrier type") + })?; + let carrier = emit_carrier_constant(ctx, rewriter, &enum_ty, carrier_ty, bits)?; + let insert = llvm::InsertValueOp::new(ctx, current_struct, carrier, vec![carrier_slot]); + rewriter.insert_operation(ctx, insert.get_operation()); + current_struct = insert.get_operation().deref(ctx).get_result(0); + last_op = insert.get_operation(); + } + + let field_base: usize = enum_ty + .variant_field_counts + .iter() + .take(variant_index) + .map(|&c| c as usize) + .sum(); + + // Insert every payload field that owns a struct slot; remember the + // slotless ones for the memory pass below. + let mut deferred: Vec<(usize, Value)> = Vec::new(); + for (i, operand) in operands.into_iter().enumerate() { + let flat = field_base + i; + let Some(slot) = slot_map.field_slots.get(flat) else { + return pliron::input_err_noloc!( + "MirConstructEnumOp field {} of variant {} is out of range for the enum's {} fields", + i, + variant_index, + slot_map.field_slots.len() + ); + }; + match slot { + Some(slot) => { + let storage_ty = { + let storage_type = llvm_struct_ty.deref(ctx); + let storage_struct = storage_type + .downcast_ref::() + .ok_or_else(|| { + pliron::input_error_noloc!( + "MirConstructEnumOp physical storage is not an LLVM struct" + ) + })?; + let storage_slot = *slot as usize; + if storage_slot >= storage_struct.num_fields() { + return pliron::input_err_noloc!( + "MirConstructEnumOp physical field slot {} is out of range", + slot + ); + } + storage_struct.field_type(storage_slot) + }; + let stored_operand = + coerce_enum_payload_storage(ctx, rewriter, operand, storage_ty)?; + let insert_op = + llvm::InsertValueOp::new(ctx, current_struct, stored_operand, vec![*slot]); + rewriter.insert_operation(ctx, insert_op.get_operation()); + current_struct = insert_op.get_operation().deref(ctx).get_result(0); + last_op = insert_op.get_operation(); + } + None => { + // Zero-sized fields own no bytes; nothing to write. + if is_zero_sized_type(ctx, slot_map.field_llvm_types[flat]) { + continue; + } + deferred.push((flat, operand)); + } + } + } + + if deferred.is_empty() { + rewriter.replace_operation(ctx, op, last_op); + return Ok(()); + } + + // Slotless fields: copy the half-built value to the stack, write + // each remaining payload at its byte position, and load the finished + // enum back as the result. + let abi_align = enum_ty.abi_align(); + let slot_ptr = spill_enum_value(ctx, rewriter, current_struct, llvm_struct_ty, abi_align); + for (flat, operand) in deferred { + let field_ptr = byte_offset_gep(ctx, rewriter, slot_ptr, slot_map.field_offsets[flat]); + // `bool` is an LLVM i1 as a value but occupies one full byte in + // Rust memory. Enum storage claims the byte-faithful twin of every + // bool-bearing payload (scalar i8 byte, or an aggregate with each + // i1 leaf widened to i8), so canonicalize the stored value to that + // twin: every physical bool byte becomes an unambiguous 0 or 1. + let semantic_ty = slot_map.field_llvm_types[flat]; + let storage_ty = enum_payload_storage_type(ctx, semantic_ty).map_err(anyhow_to_pliron)?; + let canonical_operand = canonicalize_bool_value_bytes(ctx, rewriter, operand)?; + let stored_operand = + coerce_enum_payload_storage(ctx, rewriter, canonical_operand, storage_ty)?; + let store_op = llvm::StoreOp::new(ctx, stored_operand, field_ptr); + rewriter.insert_operation(ctx, store_op.get_operation()); + } + let load_op = llvm::LoadOp::new(ctx, slot_ptr, llvm_struct_ty); + rewriter.insert_operation(ctx, load_op.get_operation()); + if abi_align > 0 { + llvm_export::ops::set_op_alignment(ctx, load_op.get_operation(), abi_align as u32); + } + rewriter.replace_operation(ctx, op, load_op.get_operation()); + + Ok(()) +} + +/// Get the slot map for an enum operand. +/// +/// By the time an op is converted, its operand's type has already been +/// rewritten to the LLVM struct, so we look up the ORIGINAL `MirEnumType` +/// the framework recorded for it and rebuild the map from that. Also +/// returns the enum's rustc alignment, which spill slots need. +fn enum_slot_map_of_operand( + ctx: &mut Context, + operands_info: &OperandsInfo, + enum_val: Value, +) -> Result<(EnumSlotMap, u64)> { + // Clone the type data out so the `Ref` borrow of `ctx` ends before + // re-interning (types are hash-consed: registering an equal instance + // returns the existing pointer). + let enum_ty: MirEnumType = { + match operands_info.lookup_most_recent_of_type::(ctx, enum_val) { + Some(r) => r.clone(), + None => { + return pliron::input_err_noloc!("Expected MirEnumType for enum value access"); + } + } + }; + let abi_align = enum_ty.abi_align(); + let mir_ty: TypeHandle = pliron::r#type::Type::instantiate(enum_ty, ctx).into(); + let map = build_enum_slot_map(ctx, mir_ty).map_err(anyhow_to_pliron)?; + Ok((map, abi_align)) +} + +/// Convert `mir.set_discriminant` to the one physical carrier write rustc's +/// layout requires. Direct layouts write the declared discriminant; Niche +/// layouts write wrapping `niche_start + range_offset`. Selecting the +/// untagged Niche variant or an inhabited Single variant is a no-op. +/// +/// The enum layout comes from the op's own `set_discriminant_enum_ty` +/// attribute, stamped at build time. Operand type history is not usable +/// here: a kind-only `mir.cast` lowers to a plain value forwarding, history +/// does not follow that edge, and a stale hit would write the tag at the +/// wrong offset. +pub(crate) fn convert_set_discriminant( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + _operands_info: &OperandsInfo, +) -> Result<()> { + let enum_ptr = op.deref(ctx).get_operand(0); + let target = MirSetDiscriminantOp::new(op) + .get_attr_set_discriminant_variant_index(ctx) + .map(|attr| attr.0 as usize) + .ok_or_else(|| { + pliron::input_error_noloc!( + "MirSetDiscriminantOp missing set_discriminant_variant_index" + ) + })?; + + let enum_ty: MirEnumType = { + let stamped_ty = MirSetDiscriminantOp::new(op) + .get_attr_set_discriminant_enum_ty(ctx) + .map(|attr| attr.get_type(ctx)) + .ok_or_else(|| { + pliron::input_error_noloc!( + "mir.set_discriminant missing enum type attribute; \ + discriminant write has no fact to derive from" + ) + })?; + match stamped_ty.deref(ctx).downcast_ref::() { + Some(et) => et.clone(), + None => { + return pliron::input_err_noloc!( + "mir.set_discriminant enum type attribute must be an enum type" + ); + } + } + }; + + let Some(bits) = enum_carrier_bits_for_variant(&enum_ty, target) + .map_err(|error| pliron::input_error_noloc!("MirSetDiscriminantOp: {error}"))? + else { + rewriter.erase_operation(ctx, op); + return Ok(()); + }; + + let tag_offset = enum_ty.tag_offset(); + let mir_ty: TypeHandle = pliron::r#type::Type::instantiate(enum_ty.clone(), ctx).into(); + let slot_map = build_enum_slot_map(ctx, mir_ty).map_err(anyhow_to_pliron)?; + let carrier_ty = slot_map.carrier_llvm_ty.ok_or_else(|| { + pliron::input_error_noloc!("MirSetDiscriminantOp physical write has no carrier type") + })?; + let carrier = emit_carrier_constant(ctx, rewriter, &enum_ty, carrier_ty, bits)?; + let carrier_ptr = byte_offset_gep(ctx, rewriter, enum_ptr, tag_offset); + let store_op = llvm::StoreOp::new(ctx, carrier, carrier_ptr); + rewriter.insert_operation(ctx, store_op.get_operation()); + + rewriter.erase_operation(ctx, op); + Ok(()) +} + +/// Convert `mir.get_discriminant` (reading which variant is alive) to +/// `llvm.extractvalue`. +/// +/// Direct layouts read the tag from the slot map's carrier slot. Niche +/// layouts decode rustc's wrapping carrier range; Single layouts materialize +/// their one logical discriminant as a constant. No slot number is assumed. +pub(crate) fn convert_get_discriminant( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + operands_info: &OperandsInfo, +) -> Result<()> { + let enum_val = match op.deref(ctx).operands().next() { + Some(v) => v, + None => return pliron::input_err_noloc!("MirGetDiscriminantOp requires an operand"), + }; + + let enum_ty: MirEnumType = operands_info + .lookup_most_recent_of_type::(ctx, enum_val) + .map(|ty| ty.clone()) + .ok_or_else(|| pliron::input_error_noloc!("Expected MirEnumType for discriminant read"))?; + let mir_ty: TypeHandle = pliron::r#type::Type::instantiate(enum_ty.clone(), ctx).into(); + let slot_map = build_enum_slot_map(ctx, mir_ty).map_err(anyhow_to_pliron)?; + let logical_ty = convert_type(ctx, enum_ty.discriminant_ty).map_err(anyhow_to_pliron)?; + let logical_width = logical_ty + .deref(ctx) + .downcast_ref::() + .map(IntegerType::width) + .ok_or_else(|| { + pliron::input_error_noloc!("MirGetDiscriminantOp logical result must be integer") + })?; + + let result = match enum_ty.layout_kind { + EnumLayoutKind::Direct => { + let slot = slot_map.carrier_slot.ok_or_else(|| { + pliron::input_error_noloc!("Direct enum has no physical carrier slot") + })?; + let extract = llvm::ExtractValueOp::new(ctx, enum_val, vec![slot])?; + rewriter.insert_operation(ctx, extract.get_operation()); + extract.get_operation().deref(ctx).get_result(0) + } + EnumLayoutKind::Single => { + if enum_ty.variant_is_inhabited(enum_ty.single_variant as usize) != Some(true) { + return pliron::input_err_noloc!( + "Cannot read discriminant of an uninhabited single-variant enum" + ); + } + let value = *enum_ty + .variant_discriminants + .get(enum_ty.single_variant as usize) + .ok_or_else(|| { + pliron::input_error_noloc!("Single enum has no declared discriminant") + })?; + emit_integer_constant(ctx, rewriter, logical_width, u128::from(value)) + } + EnumLayoutKind::Niche => { + let slot = slot_map.carrier_slot.ok_or_else(|| { + pliron::input_error_noloc!("Niche enum has no physical carrier slot") + })?; + let extract = llvm::ExtractValueOp::new(ctx, enum_val, vec![slot])?; + rewriter.insert_operation(ctx, extract.get_operation()); + let carrier = extract.get_operation().deref(ctx).get_result(0); + let carrier_int_ty: TypeHandle = + IntegerType::get(ctx, enum_ty.carrier_width, Signedness::Signless).into(); + let carrier_int = if enum_ty.carrier_kind == EnumCarrierKind::Pointer { + let cast = llvm::PtrToIntOp::new(ctx, carrier, carrier_int_ty); + rewriter.insert_operation(ctx, cast.get_operation()); + cast.get_operation().deref(ctx).get_result(0) + } else { + carrier + }; + + let niche_start = + emit_integer_constant(ctx, rewriter, enum_ty.carrier_width, enum_ty.niche_start()); + let relative = llvm::SubOp::new_with_overflow_flag( + ctx, + carrier_int, + niche_start, + IntegerOverflowFlagsAttr::default(), + ) + .get_operation(); + rewriter.insert_operation(ctx, relative); + let relative_val = relative.deref(ctx).get_result(0); + let max = emit_integer_constant( + ctx, + rewriter, + enum_ty.carrier_width, + u128::from(enum_ty.niche_variant_end - enum_ty.niche_variant_start), + ); + let in_range = + llvm::ICmpOp::new(ctx, ICmpPredicateAttr::ULE, relative_val, max).get_operation(); + rewriter.insert_operation(ctx, in_range); + // The range test is carrier-width wrapping arithmetic, exactly + // like rustc. Variant indices themselves live at the logical + // discriminant width: the start index may be larger than the + // carrier can represent (e.g. variants 298..=299 in an i8 niche). + let logical_relative = match enum_ty.carrier_width.cmp(&logical_width) { + std::cmp::Ordering::Equal => relative_val, + std::cmp::Ordering::Greater => { + let cast = llvm::TruncOp::new(ctx, relative_val, logical_ty).get_operation(); + rewriter.insert_operation(ctx, cast); + cast.deref(ctx).get_result(0) + } + std::cmp::Ordering::Less => { + // LLVM's zext op requires its explicit `nneg` flag even + // when the flag is false. Niche decoding is ordinary + // unsigned extension, so it must never claim nneg. + let cast = llvm::ZExtOp::new_with_nneg(ctx, relative_val, logical_ty, false) + .get_operation(); + rewriter.insert_operation(ctx, cast); + cast.deref(ctx).get_result(0) + } + }; + let niche_base = emit_integer_constant( + ctx, + rewriter, + logical_width, + u128::from(enum_ty.niche_variant_start), + ); + let niche_variant = llvm::AddOp::new_with_overflow_flag( + ctx, + logical_relative, + niche_base, + IntegerOverflowFlagsAttr::default(), + ) + .get_operation(); + rewriter.insert_operation(ctx, niche_variant); + let untagged = emit_integer_constant( + ctx, + rewriter, + logical_width, + u128::from(enum_ty.untagged_variant), + ); + let in_range_value = in_range.deref(ctx).get_result(0); + let niche_variant_value = niche_variant.deref(ctx).get_result(0); + let select = llvm::SelectOp::new(ctx, in_range_value, niche_variant_value, untagged) + .get_operation(); + rewriter.insert_operation(ctx, select); + select.deref(ctx).get_result(0) + } + EnumLayoutKind::Empty => { + return pliron::input_err_noloc!("Cannot read discriminant of uninhabited enum"); + } + _ => { + return pliron::input_err_noloc!( + "Cannot read discriminant of enum with unknown physical layout" + ); + } + }; + + rewriter.replace_operation_with_values(ctx, op, vec![result]); + + Ok(()) +} + +/// Convert `mir.enum_payload` (reading a variant's field, e.g. the `x` +/// in `E::A(x) => x`) to a payload-field read. +/// +/// Three cases, decided by the [`EnumSlotMap`]: +/// +/// - The field owns a struct slot: a plain `llvm.extractvalue`. +/// - The field has no slot (its bytes are shared with a different-typed +/// field of another variant): go through memory. Copy the enum to a +/// stack slot, point at the field's byte position, and load it with +/// its own type. Same trick as +/// [`convert_extract_array_element`](super::array_extract::convert_extract_array_element), +/// and +/// it avoids LLVM `bitcast` entirely. +/// - The field is zero-sized: there is nothing to read; produce `undef`. +pub(crate) fn convert_enum_payload( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + operands_info: &OperandsInfo, +) -> Result<()> { + let enum_val = match op.deref(ctx).operands().next() { + Some(v) => v, + None => return pliron::input_err_noloc!("MirEnumPayloadOp requires an operand"), + }; + + let payload_op = MirEnumPayloadOp::new(op); + let variant_index = payload_op + .get_attr_payload_variant_index(ctx) + .map(|attr| attr.0 as usize) + .unwrap_or(0); + let field_index = payload_op + .get_attr_payload_field_index(ctx) + .map(|attr| attr.0 as usize) + .unwrap_or(0); + + let variant_field_counts = { + match operands_info.lookup_most_recent_of_type::(ctx, enum_val) { + Some(r) => r.variant_field_counts.clone(), + None => { + return pliron::input_err_noloc!( + "Expected MirEnumType for enum payload extraction" + ); + } + } + }; + let (slot_map, abi_align) = enum_slot_map_of_operand(ctx, operands_info, enum_val)?; + + let field_base: usize = variant_field_counts + .iter() + .take(variant_index) + .map(|&c| c as usize) + .sum(); + let flat = field_base + field_index; + let Some(slot) = slot_map.field_slots.get(flat).copied() else { + return pliron::input_err_noloc!( + "MirEnumPayloadOp field {} of variant {} is out of range for the enum's {} fields", + field_index, + variant_index, + slot_map.field_slots.len() + ); + }; + + match slot { + Some(slot) => { + let extract_op = llvm::ExtractValueOp::new(ctx, enum_val, vec![slot])?; + rewriter.insert_operation(ctx, extract_op.get_operation()); + let extracted = extract_op.get_operation().deref(ctx).get_result(0); + let semantic_value = coerce_enum_payload_storage( + ctx, + rewriter, + extracted, + slot_map.field_llvm_types[flat], + )?; + rewriter.replace_operation_with_values(ctx, op, vec![semantic_value]); + } + None if is_zero_sized_type(ctx, slot_map.field_llvm_types[flat]) => { + let undef_op = llvm::UndefOp::new(ctx, slot_map.field_llvm_types[flat]); + rewriter.insert_operation(ctx, undef_op.get_operation()); + rewriter.replace_operation(ctx, op, undef_op.get_operation()); + } + None => { + let slot_ptr = + spill_enum_value(ctx, rewriter, enum_val, slot_map.llvm_struct_ty, abi_align); + let field_ptr = byte_offset_gep(ctx, rewriter, slot_ptr, slot_map.field_offsets[flat]); + let semantic_ty = slot_map.field_llvm_types[flat]; + let storage_ty = + enum_payload_storage_type(ctx, semantic_ty).map_err(anyhow_to_pliron)?; + let load_op = llvm::LoadOp::new(ctx, field_ptr, storage_ty); + rewriter.insert_operation(ctx, load_op.get_operation()); + let stored = load_op.get_operation().deref(ctx).get_result(0); + let semantic = coerce_enum_payload_storage(ctx, rewriter, stored, semantic_ty)?; + rewriter.replace_operation_with_values(ctx, op, vec![semantic]); + } + } + + Ok(()) +} + +#[cfg(test)] +// Tests build kinded fixture types directly; production minting lives in mir-importer's facts.rs. +#[allow(clippy::disallowed_methods)] +mod tests { + use super::*; + use crate::convert::ops::test_util::*; + use dialect_mir::attributes::{FieldIndexAttr, MirCastKindAttr, VariantIndexAttr}; + use dialect_mir::ops as mir; + use dialect_mir::types::{ + EnumEncoding, EnumVariant, MirArrayType, MirPtrType, MirStructType, MirTupleType, + }; + use llvm_export::types as llvm_types; + use pliron::builtin::attributes::IntegerAttr; + use pliron::common_traits::Verify; + + use super::super::test_support::*; + + /// Enum construction must store the declared discriminant value, not the + /// variant index. This locks the `Ordering::Less = -1` style case as the + /// i8 bit-pattern `255`. + #[test] + fn construct_enum_uses_declared_discriminant_not_variant_index() { + let mut ctx = make_ctx(); + + let discr_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Signed).into(); + let enum_ty: TypeHandle = MirEnumType::get_with_layout( + &mut ctx, + "OrderingLike".to_string(), + discr_ty, + vec![255, 0, 1], + vec![ + EnumVariant::unit("Less".to_string()), + EnumVariant::unit("Equal".to_string()), + EnumVariant::unit("Greater".to_string()), + ], + 0, + 1, + 1, + ) + .into(); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + + let op = Operation::new( + &mut ctx, + MirConstructEnumOp::get_concrete_op_info(), + vec![enum_ty], + vec![], + vec![], + 0, + ); + MirConstructEnumOp::new(op) + .set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(0)); + op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + let inserts = find_all::(&ctx, &body); + + assert_eq!( + insert_indices(&ctx, &inserts), + vec![vec![0]], + "unit enum construction should insert only the discriminant tag" + ); + let tag_insert = &inserts[0]; + assert!( + tag_insert.verify(&ctx).is_ok(), + "the enum tag insertion must satisfy LLVM dialect verification" + ); + let tag = tag_insert.get_operation().deref(&ctx).get_operand(1); + let tag_def = tag + .defining_op() + .expect("the inserted enum tag must have a defining operation"); + let tag_constant = Operation::get_op::(tag_def, &ctx) + .expect("the inserted enum tag must be defined by llvm.constant"); + let tag_attr = tag_constant.get_value(&ctx); + let tag_integer = tag_attr + .downcast_ref::() + .expect("the inserted enum tag must be an integer constant"); + assert_eq!(tag_integer.value().bw(), 8, "the enum tag must be 8-bit"); + assert_eq!( + tag_integer.value().to_u64(), + 255, + "Less must lower to its declared i8 bit-pattern 255, not variant index 0" + ); + } + + #[test] + fn nested_pointer_niche_tuple_construct_extract_and_discriminant_lower() { + let mut ctx = make_ctx(); + let logical: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signed).into(); + let index: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let pointee: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, pointee, false).into(); + let tuple_ty: TypeHandle = MirTupleType::get(&mut ctx, vec![index, pointer]).into(); + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "Option".into(), + logical, + vec![0, 1], + vec![ + EnumVariant::unit("None".into()), + EnumVariant::new_with_layout("Some".into(), vec![tuple_ty], vec![0], vec![16]), + ], + EnumEncoding { + tag_offset: 8, + total_size: 16, + abi_align: 8, + layout_kind: EnumLayoutKind::Niche, + carrier_kind: EnumCarrierKind::Pointer, + carrier_width: 64, + untagged_variant: 1, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + let slot_map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); + assert_eq!(slot_map.carrier_slot, Some(1)); + assert_eq!(slot_map.field_slots, vec![None]); + let lowered_tuple = convert_type(&mut ctx, tuple_ty).unwrap(); + + let (module, block) = build_kernel(&mut ctx, vec![index, pointer], vec![]); + let index_value = block.deref(&ctx).get_argument(0); + let pointer_value = block.deref(&ctx).get_argument(1); + let tuple = Operation::new( + &mut ctx, + mir::MirConstructTupleOp::get_concrete_op_info(), + vec![tuple_ty], + vec![index_value, pointer_value], + vec![], + 0, + ); + tuple.insert_at_back(block, &ctx); + let tuple_value = tuple.deref(&ctx).get_result(0); + + let construct = Operation::new( + &mut ctx, + mir::MirConstructEnumOp::get_concrete_op_info(), + vec![enum_ty], + vec![tuple_value], + vec![], + 0, + ); + mir::MirConstructEnumOp::new(construct) + .set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(1)); + construct.insert_at_back(block, &ctx); + let enum_value = construct.deref(&ctx).get_result(0); + + let payload = Operation::new( + &mut ctx, + mir::MirEnumPayloadOp::get_concrete_op_info(), + vec![tuple_ty], + vec![enum_value], + vec![], + 0, + ); + mir::MirEnumPayloadOp::new(payload) + .set_attr_payload_variant_index(&ctx, VariantIndexAttr(1)); + mir::MirEnumPayloadOp::new(payload).set_attr_payload_field_index(&ctx, FieldIndexAttr(0)); + payload.insert_at_back(block, &ctx); + + let discriminant = Operation::new( + &mut ctx, + mir::MirGetDiscriminantOp::get_concrete_op_info(), + vec![logical], + vec![enum_value], + vec![], + 0, + ); + discriminant.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); + let body = kernel_blocks(&ctx, module); + assert_eq!( + count_ops::(&ctx, &body), + 0, + "constructing the untagged Some payload must not recreate its pointer from bits" + ); + assert_eq!( + count_ops::(&ctx, &body), + 1, + "reading the pointer niche should inspect the carrier exactly once" + ); + assert_eq!( + count_ops::(&ctx, &body), + 3, + "construction and extraction should each spill the enum, plus one tuple payload store" + ); + assert_eq!( + count_ops::(&ctx, &body), + 2, + "construction should reload the enum and extraction should load the tuple payload" + ); + // What this test is about is that the payload moves as one unit, never + // field by field. Assert that property directly instead of counting + // whole-aggregate accesses. + // + // Counting was the fragile form. The enum's physical storage here is + // `{i64, ptr}` -- the *same interned type* as the lowered payload tuple, + // since the niche carrier claims the pointer at byte 8 and the 8 bytes + // below it become one `i64` filler. So a count of `lowered_tuple`-typed + // accesses cannot separate the payload write from the enum spill, and + // the expected numbers move whenever that coincidence appears or + // disappears -- which has nothing to do with the property under test. + // + // A lowering that decomposed the payload is recognisable by what it + // emits instead: traffic in the tuple's *field* types. Look for that, + // and the assertion holds whatever the enum storage type happens to be. + let field_tys: Vec = lowered_tuple + .deref(&ctx) + .downcast_ref::() + .expect("the lowered tuple is an LLVM struct") + .fields() + .collect(); + + let store_tys: Vec = find_all::(&ctx, &body) + .iter() + .map(|store| store.get_operand_value(&ctx).get_type(&ctx)) + .collect(); + let load_tys: Vec = find_all::(&ctx, &body) + .iter() + .map(|load| { + load.get_operation() + .deref(&ctx) + .get_result(0) + .get_type(&ctx) + }) + .collect(); + + let describe = |tys: &[TypeHandle]| -> Vec { + tys.iter() + .filter(|ty| field_tys.contains(ty)) + .map(|ty| ty.deref(&ctx).disp(&ctx).to_string()) + .collect() + }; + assert!( + describe(&store_tys).is_empty(), + "the payload must be stored whole, but these field-typed stores appear: {:?}", + describe(&store_tys) + ); + assert!( + describe(&load_tys).is_empty(), + "the payload must be read whole, but these field-typed loads appear: {:?}", + describe(&load_tys) + ); + + // And it must actually be moved: without this the checks above would + // also pass a lowering that emitted no payload traffic at all. + assert!( + store_tys.contains(&lowered_tuple), + "at least one store must move the complete {{i64, ptr}} payload" + ); + assert!( + load_tys.contains(&lowered_tuple), + "at least one load must read the complete {{i64, ptr}} payload" + ); + } + + /// SetDiscriminant must use the slot map instead of assuming that the tag + /// is field zero, and its GEP must retain the source pointer's GPU address + /// space. This shape puts an i64 payload first and an i8 tag above the + /// u32 range, proving the byte GEP does not truncate offsets to u32. + #[test] + fn set_discriminant_uses_tag_slot_and_preserves_shared_address_space() { + use llvm_export::ops::GepIndex; + use llvm_export::types::{PointerType, address_space}; + + let mut ctx = make_ctx(); + let discr_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let payload_a: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let payload_b: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let tag_offset = u64::from(u32::MAX) + 1; + let enum_ty: TypeHandle = MirEnumType::get_with_layout( + &mut ctx, + "TagAfterPayload".to_string(), + discr_ty, + vec![3, 7], + vec![ + EnumVariant::new_with_layout("A".to_string(), vec![payload_a], vec![0], vec![8]), + EnumVariant::new_with_layout("B".to_string(), vec![payload_b], vec![0], vec![8]), + ], + tag_offset, + tag_offset + 8, + 8, + ) + .into(); + let ptr_ty: TypeHandle = MirPtrType::get_shared(&mut ctx, enum_ty, true).into(); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![ptr_ty], vec![]); + let enum_ptr = block.deref(&ctx).get_argument(0); + let set = Operation::new( + &mut ctx, + mir::MirSetDiscriminantOp::get_concrete_op_info(), + vec![], + vec![enum_ptr], + vec![], + 0, + ); + let set_op = mir::MirSetDiscriminantOp::new(set); + set_op.set_attr_set_discriminant_variant_index(&ctx, VariantIndexAttr(1)); + set_op.set_attr_set_discriminant_enum_ty( + &ctx, + pliron::builtin::attributes::TypeAttr::new(enum_ty), + ); + set.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + assert_eq!(count_ops::(&ctx, &body), 0); + assert_eq!(count_ops::(&ctx, &body), 1); + assert_eq!(count_ops::(&ctx, &body), 1); + + let gep = find_first::(&ctx, &body).unwrap(); + let indices = gep.indices(&ctx); + assert!(matches!(indices.as_slice(), [GepIndex::Value(_)])); + let GepIndex::Value(offset) = indices[0] else { + unreachable!() + }; + let offset_def = offset.defining_op().expect("byte offset must be constant"); + let offset_constant = Operation::get_op::(offset_def, &ctx) + .expect("byte offset must be an LLVM constant"); + let offset_attr = offset_constant.get_value(&ctx); + assert_eq!( + offset_attr + .downcast_ref::() + .expect("byte offset must be integer") + .value() + .to_u64(), + tag_offset, + "the write must use rustc's absolute carrier byte offset" + ); + let gep_result_ty = gep.get_operation().deref(&ctx).get_result(0).get_type(&ctx); + assert_eq!( + gep_result_ty + .deref(&ctx) + .downcast_ref::() + .expect("GEP result must be a pointer") + .address_space(), + address_space::SHARED, + "tag GEP must preserve shared address space" + ); + + let store = find_first::(&ctx, &body).unwrap(); + let stored_ty = store.get_operand_value(&ctx).get_type(&ctx); + assert_eq!( + stored_ty + .deref(&ctx) + .downcast_ref::() + .expect("stored tag must be an integer") + .width(), + 8 + ); + assert_eq!( + store.get_operand_address(&ctx), + gep.get_operation().deref(&ctx).get_result(0), + "the store must use the tag GEP result" + ); + } + + #[test] + fn set_discriminant_niche_writes_only_tagged_variant_carrier() { + for (target, expected_stores) in [(0, 1), (1, 0)] { + let mut ctx = make_ctx(); + let enum_ty = unit_niche_enum( + &mut ctx, + (EnumCarrierKind::Integer, 8, 0), + 0, + 0..=0, + 1, + vec![1, 1], + ); + let ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, enum_ty, true).into(); + let (module, block) = build_kernel(&mut ctx, vec![ptr_ty], vec![]); + let ptr = block.deref(&ctx).get_argument(0); + let set = Operation::new( + &mut ctx, + mir::MirSetDiscriminantOp::get_concrete_op_info(), + vec![], + vec![ptr], + vec![], + 0, + ); + let set_op = mir::MirSetDiscriminantOp::new(set); + set_op.set_attr_set_discriminant_variant_index(&ctx, VariantIndexAttr(target)); + set_op.set_attr_set_discriminant_enum_ty( + &ctx, + pliron::builtin::attributes::TypeAttr::new(enum_ty), + ); + set.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); + let body = kernel_blocks(&ctx, module); + assert_eq!(count_ops::(&ctx, &body), expected_stores); + assert_eq!(count_ops::(&ctx, &body), 0); + } + } + + #[test] + fn shared_pointer_niche_carrier_rejects_target_dependent_width() { + let mut ctx = make_ctx(); + let enum_ty = unit_niche_enum( + &mut ctx, + (EnumCarrierKind::Pointer, 64, 3), + 0, + 0..=0, + 1, + vec![1, 1], + ); + let error = build_enum_slot_map(&mut ctx, enum_ty) + .err() + .expect("shared pointer carrier must reject"); + assert!( + error.to_string().contains("target-mode dependent"), + "{error}" + ); + } + + #[test] + fn shared_pointer_niche_payload_round_trips_through_generic_storage() { + let mut ctx = make_ctx(); + let logical: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signed).into(); + let pointee: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, true).into(); + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "OptionShared".into(), + logical, + vec![0, 1], + vec![ + EnumVariant::unit("None".into()), + EnumVariant::new_with_layout("Some".into(), vec![shared], vec![0], vec![8]), + ], + EnumEncoding { + tag_offset: 0, + total_size: 8, + abi_align: 8, + layout_kind: EnumLayoutKind::Niche, + carrier_kind: EnumCarrierKind::Pointer, + carrier_width: 64, + carrier_address_space: llvm_types::address_space::GENERIC, + niche_start: 0, + niche_variant_start: 0, + niche_variant_end: 0, + untagged_variant: 1, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + + let (module, block) = build_kernel(&mut ctx, vec![shared], vec![shared]); + let pointer = block.deref(&ctx).get_argument(0); + let construct = Operation::new( + &mut ctx, + MirConstructEnumOp::get_concrete_op_info(), + vec![enum_ty], + vec![pointer], + vec![], + 0, + ); + MirConstructEnumOp::new(construct) + .set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(1)); + construct.insert_at_back(block, &ctx); + let option = construct.deref(&ctx).get_result(0); + + let discriminant = Operation::new( + &mut ctx, + mir::MirGetDiscriminantOp::get_concrete_op_info(), + vec![logical], + vec![option], + vec![], + 0, + ); + discriminant.insert_at_back(block, &ctx); + + let payload = Operation::new( + &mut ctx, + MirEnumPayloadOp::get_concrete_op_info(), + vec![shared], + vec![option], + vec![], + 0, + ); + let payload_op = MirEnumPayloadOp::new(payload); + payload_op.set_attr_payload_variant_index(&ctx, VariantIndexAttr(1)); + payload_op.set_attr_payload_field_index(&ctx, FieldIndexAttr(0)); + payload.insert_at_back(block, &ctx); + let result = payload.deref(&ctx).get_result(0); + append_mir_return(&mut ctx, block, vec![result]); + + crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); + let body = kernel_blocks(&ctx, module); + assert_eq!( + count_ops::(&ctx, &body), + 2, + "construction must genericize the pointer and extraction must restore shared space" + ); + assert_eq!( + count_ops::(&ctx, &body), + 1, + "niche discrimination must inspect the generic pointer carrier" + ); + assert_eq!(count_ops::(&ctx, &body), 0); + assert_eq!(count_ops::(&ctx, &body), 0); + } + + #[test] + fn nested_shared_pointer_payload_round_trips_through_recursive_storage() { + let mut ctx = make_ctx(); + let logical: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let pointee: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, true).into(); + let inner: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "SharedPointerInner".into(), + vec!["pointer".into(), "cookie".into()], + vec![shared, logical], + vec![0, 1], + vec![0, 8], + 16, + 8, + ) + .into(); + let outer: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "SharedPointerOuter".into(), + vec!["inner".into(), "guard".into()], + vec![inner, logical], + vec![0, 1], + vec![0, 16], + 24, + 8, + ) + .into(); + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "NestedSharedPointer".into(), + logical, + vec![0, 1], + vec![ + EnumVariant::unit("None".into()), + EnumVariant::new_with_layout("Some".into(), vec![outer], vec![8], vec![24]), + ], + EnumEncoding { + tag_offset: 0, + total_size: 32, + abi_align: 8, + layout_kind: EnumLayoutKind::Direct, + carrier_kind: EnumCarrierKind::Integer, + carrier_width: 32, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + + let (module, block) = build_kernel(&mut ctx, vec![outer], vec![outer]); + let wrapper = block.deref(&ctx).get_argument(0); + let construct = Operation::new( + &mut ctx, + MirConstructEnumOp::get_concrete_op_info(), + vec![enum_ty], + vec![wrapper], + vec![], + 0, + ); + MirConstructEnumOp::new(construct) + .set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(1)); + construct.insert_at_back(block, &ctx); + let option = construct.deref(&ctx).get_result(0); + + let payload = Operation::new( + &mut ctx, + MirEnumPayloadOp::get_concrete_op_info(), + vec![outer], + vec![option], + vec![], + 0, + ); + let payload_op = MirEnumPayloadOp::new(payload); + payload_op.set_attr_payload_variant_index(&ctx, VariantIndexAttr(1)); + payload_op.set_attr_payload_field_index(&ctx, FieldIndexAttr(0)); + payload.insert_at_back(block, &ctx); + let result = payload.deref(&ctx).get_result(0); + append_mir_return(&mut ctx, block, vec![result]); + + crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); + let body = kernel_blocks(&ctx, module); + assert_eq!( + count_ops::(&ctx, &body), + 2, + "construction and extraction must cast the nested shared pointer leaf" + ); + assert_eq!(count_ops::(&ctx, &body), 0); + assert_eq!(count_ops::(&ctx, &body), 0); + } + + #[test] + fn bounded_shared_pointer_array_niche_payload_round_trips_through_recursive_storage() { + let mut ctx = make_ctx(); + let logical: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signed).into(); + let pointee: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, true).into(); + let pointers: TypeHandle = MirArrayType::get(&mut ctx, shared, 2).into(); + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "OptionSharedPointerArray".into(), + logical, + vec![0, 1], + vec![ + EnumVariant::unit("None".into()), + EnumVariant::new_with_layout("Some".into(), vec![pointers], vec![0], vec![16]), + ], + EnumEncoding { + tag_offset: 0, + total_size: 16, + abi_align: 8, + layout_kind: EnumLayoutKind::Niche, + carrier_kind: EnumCarrierKind::Pointer, + carrier_width: 64, + carrier_address_space: llvm_types::address_space::GENERIC, + niche_start: 0, + niche_variant_start: 0, + niche_variant_end: 0, + untagged_variant: 1, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + + let (module, block) = build_kernel(&mut ctx, vec![pointers], vec![pointers]); + let array = block.deref(&ctx).get_argument(0); + let construct = Operation::new( + &mut ctx, + MirConstructEnumOp::get_concrete_op_info(), + vec![enum_ty], + vec![array], + vec![], + 0, + ); + MirConstructEnumOp::new(construct) + .set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(1)); + construct.insert_at_back(block, &ctx); + let option = construct.deref(&ctx).get_result(0); + + let discriminant = Operation::new( + &mut ctx, + mir::MirGetDiscriminantOp::get_concrete_op_info(), + vec![logical], + vec![option], + vec![], + 0, + ); + discriminant.insert_at_back(block, &ctx); + + let payload = Operation::new( + &mut ctx, + MirEnumPayloadOp::get_concrete_op_info(), + vec![pointers], + vec![option], + vec![], + 0, + ); + let payload_op = MirEnumPayloadOp::new(payload); + payload_op.set_attr_payload_variant_index(&ctx, VariantIndexAttr(1)); + payload_op.set_attr_payload_field_index(&ctx, FieldIndexAttr(0)); + payload.insert_at_back(block, &ctx); + let result = payload.deref(&ctx).get_result(0); + append_mir_return(&mut ctx, block, vec![result]); + + crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); + let body = kernel_blocks(&ctx, module); + assert_eq!( + count_ops::(&ctx, &body), + 4, + "construction and extraction must cast both shared-pointer array elements" + ); + assert_eq!( + count_ops::(&ctx, &body), + 1, + "niche discrimination must inspect the generic first-pointer carrier" + ); + assert_eq!(count_ops::(&ctx, &body), 0); + assert_eq!(count_ops::(&ctx, &body), 0); + } + + #[test] + fn option_bool_uses_i8_carrier_and_spills_i1_payload() { + let mut ctx = make_ctx(); + let logical: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let bool_ty: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Signless).into(); + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "OptionBool".into(), + logical, + vec![0, 1], + vec![ + EnumVariant::unit("None".into()), + EnumVariant::new_with_layout("Some".into(), vec![bool_ty], vec![0], vec![1]), + ], + EnumEncoding { + tag_offset: 0, + total_size: 1, + abi_align: 1, + layout_kind: EnumLayoutKind::Niche, + carrier_kind: EnumCarrierKind::Integer, + carrier_width: 8, + niche_start: 2, + untagged_variant: 1, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + let slot_map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); + assert_eq!( + slot_map + .carrier_llvm_ty + .unwrap() + .deref(&ctx) + .downcast_ref::() + .unwrap() + .width(), + 8, + "Option's memory carrier is i8, not bool's semantic i1" + ); + assert_eq!(slot_map.field_slots, vec![None]); + + let (module, block) = build_kernel(&mut ctx, vec![bool_ty], vec![enum_ty]); + let payload = block.deref(&ctx).get_argument(0); + let construct = Operation::new( + &mut ctx, + mir::MirConstructEnumOp::get_concrete_op_info(), + vec![enum_ty], + vec![payload], + vec![], + 0, + ); + mir::MirConstructEnumOp::new(construct) + .set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(1)); + construct.insert_at_back(block, &ctx); + let result = construct.deref(&ctx).get_result(0); + append_mir_return(&mut ctx, block, vec![result]); + + crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); + let body = kernel_blocks(&ctx, module); + assert_eq!(count_ops::(&ctx, &body), 0); + assert!(find_all::(&ctx, &body).iter().any(|zext| { + zext.get_operation() + .deref(&ctx) + .get_operand(0) + .get_type(&ctx) + .deref(&ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 1) + && zext + .get_operation() + .deref(&ctx) + .get_result(0) + .get_type(&ctx) + .deref(&ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 8) + })); + assert!(find_all::(&ctx, &body).iter().any(|store| { + store + .get_operand_value(&ctx) + .get_type(&ctx) + .deref(&ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 8) + })); + } + + #[test] + fn direct_bool_uses_i8_storage_and_spills_i1_payload() { + let mut ctx = make_ctx(); + let tag: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let bool_ty: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Signless).into(); + let enum_ty: TypeHandle = MirEnumType::get_with_layout( + &mut ctx, + "DirectBool".into(), + tag, + vec![0, 1], + vec![ + EnumVariant::new_with_layout("A".into(), vec![bool_ty], vec![4], vec![1]), + EnumVariant::unit("B".into()), + ], + 0, + 8, + 4, + ) + .into(); + let slot_map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); + assert_eq!(slot_map.field_slots, vec![None]); + let storage_fields = slot_map + .llvm_struct_ty + .deref(&ctx) + .downcast_ref::() + .expect("enum storage must be an LLVM struct") + .fields() + .collect::>(); + assert_eq!( + storage_fields[1] + .deref(&ctx) + .downcast_ref::() + .map(IntegerType::width), + Some(8), + "the standalone Rust bool byte must use physical i8 storage" + ); + + let (module, block) = build_kernel(&mut ctx, vec![bool_ty], vec![enum_ty]); + let payload = block.deref(&ctx).get_argument(0); + let construct = Operation::new( + &mut ctx, + mir::MirConstructEnumOp::get_concrete_op_info(), + vec![enum_ty], + vec![payload], + vec![], + 0, + ); + mir::MirConstructEnumOp::new(construct) + .set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(0)); + construct.insert_at_back(block, &ctx); + let result = construct.deref(&ctx).get_result(0); + append_mir_return(&mut ctx, block, vec![result]); + + crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); + let body = kernel_blocks(&ctx, module); + assert_eq!( + count_ops::(&ctx, &body), + 1, + "only the direct tag should be inserted as an SSA struct field" + ); + assert!(find_all::(&ctx, &body).iter().any(|zext| { + zext.get_operation() + .deref(&ctx) + .get_operand(0) + .get_type(&ctx) + .deref(&ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 1) + && zext + .get_operation() + .deref(&ctx) + .get_result(0) + .get_type(&ctx) + .deref(&ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 8) + })); + assert!(find_all::(&ctx, &body).iter().any(|store| { + store + .get_operand_value(&ctx) + .get_type(&ctx) + .deref(&ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 8) + })); + } + + #[test] + fn later_field_niche_set_writes_exact_carrier_offset() { + use llvm_export::ops::GepIndex; + + let mut ctx = make_ctx(); + let logical: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let wrapper: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "Wrapper".into(), + vec!["pad".into(), "nz".into()], + vec![u32_ty, u32_ty], + vec![0, 1], + vec![0, 4], + 8, + 4, + ) + .into(); + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "MaybeWrapper".into(), + logical, + vec![0, 1], + vec![ + EnumVariant::unit("None".into()), + EnumVariant::new_with_layout("Some".into(), vec![wrapper], vec![0], vec![8]), + ], + EnumEncoding { + tag_offset: 4, + total_size: 8, + abi_align: 4, + layout_kind: EnumLayoutKind::Niche, + carrier_kind: EnumCarrierKind::Integer, + carrier_width: 32, + untagged_variant: 1, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + let ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, enum_ty, true).into(); + let (module, block) = build_kernel(&mut ctx, vec![ptr_ty], vec![]); + let ptr = block.deref(&ctx).get_argument(0); + let set = Operation::new( + &mut ctx, + mir::MirSetDiscriminantOp::get_concrete_op_info(), + vec![], + vec![ptr], + vec![], + 0, + ); + let set_op = mir::MirSetDiscriminantOp::new(set); + set_op.set_attr_set_discriminant_variant_index(&ctx, VariantIndexAttr(0)); + set_op.set_attr_set_discriminant_enum_ty( + &ctx, + pliron::builtin::attributes::TypeAttr::new(enum_ty), + ); + set.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); + let body = kernel_blocks(&ctx, module); + let gep = find_first::(&ctx, &body).unwrap(); + let indices = gep.indices(&ctx); + let [GepIndex::Value(offset)] = indices.as_slice() else { + panic!("carrier access must use a byte-offset SSA value"); + }; + let constant = + Operation::get_op::(offset.defining_op().unwrap(), &ctx).unwrap(); + assert_eq!( + constant + .get_value(&ctx) + .downcast_ref::() + .unwrap() + .value() + .to_u64(), + 4 + ); + } + + #[test] + fn get_niche_discriminant_adds_large_range_start_at_logical_width() { + let mut ctx = make_ctx(); + let logical_ty: TypeHandle = IntegerType::get(&ctx, 16, Signedness::Unsigned).into(); + let carrier_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let mut variants = (0..300) + .map(|index| EnumVariant::unit(format!("V{index}"))) + .collect::>(); + // A valid niche layout stores the carrier inside the untagged + // variant's payload. Keep the large logical-variant test realistic + // instead of relying on a carrier with no backing field. + variants[0] = EnumVariant::new_with_layout("V0".into(), vec![carrier_ty], vec![0], vec![1]); + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "LargeNiche".into(), + logical_ty, + (0..300).collect(), + variants, + EnumEncoding { + total_size: 1, + abi_align: 1, + layout_kind: EnumLayoutKind::Niche, + carrier_kind: EnumCarrierKind::Integer, + carrier_width: 8, + niche_variant_start: 298, + niche_variant_end: 299, + untagged_variant: 0, + variant_inhabited: { + let mut inhabited = vec![0; 300]; + inhabited[0] = 1; + inhabited[298] = 1; + inhabited[299] = 1; + inhabited + }, + ..EnumEncoding::default() + }, + ) + .into(); + let (module, block) = build_kernel(&mut ctx, vec![enum_ty], vec![logical_ty]); + let value = block.deref(&ctx).get_argument(0); + let get = Operation::new( + &mut ctx, + mir::MirGetDiscriminantOp::get_concrete_op_info(), + vec![logical_ty], + vec![value], + vec![], + 0, + ); + get.insert_at_back(block, &ctx); + let result = get.deref(&ctx).get_result(0); + append_mir_return(&mut ctx, block, vec![result]); + + crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); + let body = kernel_blocks(&ctx, module); + assert_eq!(count_ops::(&ctx, &body), 1); + let add = find_first::(&ctx, &body).expect("logical variant add"); + for operand in add.get_operation().deref(&ctx).operands() { + assert_eq!( + operand + .get_type(&ctx) + .deref(&ctx) + .downcast_ref::() + .unwrap() + .width(), + 16, + "variant-index arithmetic must occur at logical width" + ); + } + } + + #[test] + fn direct_negative_discriminant_read_remains_signed_for_widening() { + let mut ctx = make_ctx(); + let i8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Signed).into(); + let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signed).into(); + let enum_ty: TypeHandle = MirEnumType::get_with_layout( + &mut ctx, + "Negative".into(), + i8_ty, + vec![255, 0], + vec![EnumVariant::unit("N".into()), EnumVariant::unit("Z".into())], + 0, + 1, + 1, + ) + .into(); + let (module, block) = build_kernel(&mut ctx, vec![], vec![i32_ty]); + let construct = Operation::new( + &mut ctx, + mir::MirConstructEnumOp::get_concrete_op_info(), + vec![enum_ty], + vec![], + vec![], + 0, + ); + mir::MirConstructEnumOp::new(construct) + .set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(0)); + construct.insert_at_back(block, &ctx); + let enum_value = construct.deref(&ctx).get_result(0); + let get = Operation::new( + &mut ctx, + mir::MirGetDiscriminantOp::get_concrete_op_info(), + vec![i8_ty], + vec![enum_value], + vec![], + 0, + ); + get.insert_at_back(block, &ctx); + let discriminant = get.deref(&ctx).get_result(0); + let cast = Operation::new( + &mut ctx, + mir::MirCastOp::get_concrete_op_info(), + vec![i32_ty], + vec![discriminant], + vec![], + 0, + ); + mir::MirCastOp::new(cast).set_attr_cast_kind(&ctx, MirCastKindAttr::IntToInt); + cast.insert_at_back(block, &ctx); + let widened = cast.deref(&ctx).get_result(0); + append_mir_return(&mut ctx, block, vec![widened]); + + crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); + let body = kernel_blocks(&ctx, module); + assert_eq!(count_ops::(&ctx, &body), 1); + assert_eq!(count_ops::(&ctx, &body), 0); + } + + #[test] + fn single_layout_preserves_large_and_negative_declared_discriminants() { + for (width, signedness, bits, widened_width, expects_sext) in [ + (16, Signedness::Unsigned, 1_000, 16, false), + (8, Signedness::Signed, 251, 32, true), + ] { + let mut ctx = make_ctx(); + let logical: TypeHandle = IntegerType::get(&ctx, width, signedness).into(); + let destination: TypeHandle = IntegerType::get(&ctx, widened_width, signedness).into(); + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "Single".into(), + logical, + vec![bits], + vec![EnumVariant::unit("Only".into())], + EnumEncoding { + tag_offset: 0, + total_size: 0, + abi_align: 1, + layout_kind: EnumLayoutKind::Single, + variant_inhabited: vec![1], + ..EnumEncoding::default() + }, + ) + .into(); + let (module, block) = build_kernel(&mut ctx, vec![], vec![destination]); + let construct = Operation::new( + &mut ctx, + mir::MirConstructEnumOp::get_concrete_op_info(), + vec![enum_ty], + vec![], + vec![], + 0, + ); + mir::MirConstructEnumOp::new(construct) + .set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(0)); + construct.insert_at_back(block, &ctx); + let enum_value = construct.deref(&ctx).get_result(0); + let get = Operation::new( + &mut ctx, + mir::MirGetDiscriminantOp::get_concrete_op_info(), + vec![logical], + vec![enum_value], + vec![], + 0, + ); + get.insert_at_back(block, &ctx); + let discr = get.deref(&ctx).get_result(0); + let returned = if widened_width == width { + discr + } else { + let cast = Operation::new( + &mut ctx, + mir::MirCastOp::get_concrete_op_info(), + vec![destination], + vec![discr], + vec![], + 0, + ); + mir::MirCastOp::new(cast).set_attr_cast_kind(&ctx, MirCastKindAttr::IntToInt); + cast.insert_at_back(block, &ctx); + cast.deref(&ctx).get_result(0) + }; + append_mir_return(&mut ctx, block, vec![returned]); + + crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); + let body = kernel_blocks(&ctx, module); + let found_bits = find_all::(&ctx, &body) + .iter() + .any(|constant| { + constant + .get_value(&ctx) + .downcast_ref::() + .is_some_and(|value| value.value().to_u64() == bits) + }); + assert!( + found_bits, + "single discriminant {bits} must not be truncated" + ); + assert_eq!(count_ops::(&ctx, &body) == 1, expects_sext); + } + } +} diff --git a/crates/mir-lower/src/convert/ops/aggregate/fields.rs b/crates/mir-lower/src/convert/ops/aggregate/fields.rs new file mode 100644 index 0000000000..41ca9756cf --- /dev/null +++ b/crates/mir-lower/src/convert/ops/aggregate/fields.rs @@ -0,0 +1,489 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use super::common::{anyhow_to_pliron, spill_enum_value}; +use crate::convert::types::{ + StructLayoutInfo, StructSlotMap, build_struct_slot_map, build_union_storage_type, convert_type, + is_zero_sized_type, make_slice_struct, +}; +use dialect_mir::ops::{MirExtractFieldOp, MirInsertFieldOp}; +use dialect_mir::types::{ + MirArrayType, MirDisjointSliceType, MirSliceType, MirStructType, MirTupleType, MirUnionType, +}; +use llvm_export::ops as llvm; +use pliron::builtin::types::{IntegerType, Signedness}; +use pliron::context::{Context, Ptr}; +use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; +use pliron::irbuild::inserter::Inserter; +use pliron::irbuild::rewriter::Rewriter; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::result::Result; +use pliron::r#type::{TypeHandle, Typed}; +use pliron::value::Value; + +/// How the MIR-level field indices of an aggregate operand map onto the +/// lowered LLVM aggregate. +pub(super) enum AggregateSlots { + /// Lowered from a `MirStructType`/`MirTupleType`: use the slot map the + /// type converter built (accounts for reordering, `[N x i8]` padding + /// slots and stripped ZST fields). + Mapped(StructSlotMap), + /// The MIR index is already the final LLVM index. Sound only for + /// aggregates whose lowered layout is index-preserving by construction: + /// arrays and slice fat pointers (`{ ptr, i64 }`). + Identity, +} + +/// Resolve how field indices of `aggregate` map onto its lowered type. +/// +/// Recover-or-error (issue #128): when the operand has no recorded +/// `MirStructType`/`MirTupleType` conversion history, identity indexing is +/// only sound for aggregates the converter lowers without reordering, +/// padding, or ZST stripping: arrays and slice fat pointers. Anything +/// else is a lowering bug; guessing identity there silently reads or +/// writes the wrong field, so we error out loudly instead. +pub(super) fn resolve_aggregate_slots( + ctx: &mut Context, + operands_info: &OperandsInfo, + aggregate: Value, +) -> Result { + let layout = operands_info + .lookup_most_recent_of_type::(ctx, aggregate) + .map(|struct_ref| StructLayoutInfo::of_struct(&struct_ref)) + .or_else(|| { + operands_info + .lookup_most_recent_of_type::(ctx, aggregate) + .map(|tuple_ref| StructLayoutInfo::of_tuple(&tuple_ref)) + }); + + if let Some(layout) = layout { + let map = build_struct_slot_map(ctx, &layout).map_err(anyhow_to_pliron)?; + return Ok(AggregateSlots::Mapped(map)); + } + + // Arrays keep their element indices: `[N x T]` has no reorder, no + // padding, no ZST stripping. + let is_array_history = operands_info + .lookup_most_recent_of_type::(ctx, aggregate) + .is_some(); + // Slices lower to the `{ ptr, i64 }` fat pointer, where index 0 = ptr + // and index 1 = len by construction. + let is_slice_history = operands_info + .lookup_most_recent_of_type::(ctx, aggregate) + .is_some() + || operands_info + .lookup_most_recent_of_type::(ctx, aggregate) + .is_some(); + if is_array_history || is_slice_history { + return Ok(AggregateSlots::Identity); + } + + // No conversion history at all (e.g. a slice reconstructed in the entry + // prologue, which is born as an LLVM struct). Identity is still fine if + // the current type is a slice fat pointer or an LLVM array. A disjoint + // slice with a runtime row width is the same case one word longer: every + // runtime index space today carries a single `u32` row width, and the struct the type + // converter builds for it (`make_disjoint_slice_struct`) is + // index-preserving by construction, exactly like the two-field fat + // pointer. Both shape checks share the two-field caveat that an unnamed + // user struct with the identical lowered layout would be accepted too; + // that has been the accepted trade-off for `{ ptr, i64 }` since issue + // #128, and a struct needing a slot map still errors below whenever its + // lowered shape differs (padding slots, reordering, other field types). + let aggregate_ty = aggregate.get_type(ctx); + let slice_struct_ty = make_slice_struct(ctx); + let row_width_slice_struct_ty = { + // The one runtime layout current `SpaceLayout` impls produce: a + // single u32 row width. Built through the real constructor so this + // check cannot drift from the lowered shape. A future space with a + // different `Data` shape is not listed here and thus fails closed + // into the refuse-to-guess error until this site learns it. + let width_ty: TypeHandle = IntegerType::get(ctx, 32, Signedness::Unsigned).into(); + crate::convert::types::make_disjoint_slice_struct(ctx, &[width_ty]) + .map_err(anyhow_to_pliron)? + }; + let is_llvm_array = aggregate_ty + .deref(ctx) + .is::(); + if aggregate_ty == slice_struct_ty || aggregate_ty == row_width_slice_struct_ty || is_llvm_array + { + return Ok(AggregateSlots::Identity); + } + + let ty_disp = aggregate_ty.deref(ctx).disp(ctx).to_string(); + pliron::input_err_noloc!( + "Cannot map field indices for aggregate of type {ty_disp}: no struct/tuple \ + conversion history was recorded for this operand, and identity indexing is \ + only sound for arrays and slice fat pointers (with or without the runtime \ + row-width word). Refusing to guess a field mapping (issue #128)." + ) +} + +/// Convert `mir.extract_field` to `llvm.extractvalue`. +/// +/// Handles scalar-lowered newtype case: if the operand is a scalar (e.g., `ThreadIndex`), +/// no extraction is needed. +/// +/// The declaration-order field index is mapped to the LLVM slot via +/// [`resolve_aggregate_slots`], which shares the type converter's view of +/// the struct (reorder, `[N x i8]` padding slots, stripped ZSTs). If +/// extracting a ZST field, we return undef of its (empty) type. +pub(crate) fn convert_extract_field( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + operands_info: &OperandsInfo, +) -> Result<()> { + let aggregate = op.deref(ctx).get_operand(0); + + let extract_op = MirExtractFieldOp::new(op); + let decl_index = match extract_op.get_attr_index(ctx) { + Some(attr) => attr.0 as usize, + None => return pliron::input_err_noloc!("Missing index attribute on extract_field"), + }; + + if operands_info + .lookup_most_recent_of_type::(ctx, aggregate) + .is_some() + { + return convert_extract_union_field( + ctx, + rewriter, + op, + aggregate, + decl_index, + operands_info, + ); + } + + let is_scalar = aggregate + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_some(); + + if is_scalar { + rewriter.replace_operation_with_values(ctx, op, vec![aggregate]); + return Ok(()); + } + + let llvm_index = match resolve_aggregate_slots(ctx, operands_info, aggregate)? { + AggregateSlots::Mapped(map) => match map.decl_to_llvm.get(decl_index) { + Some(Some(slot)) => *slot, + Some(None) => { + // ZST field: stripped from the LLVM struct, so there is + // nothing to extract. Materialize undef of its empty type. + let zst_ty = map.field_llvm_types[decl_index]; + let undef_op = llvm::UndefOp::new(ctx, zst_ty); + rewriter.insert_operation(ctx, undef_op.get_operation()); + rewriter.replace_operation(ctx, op, undef_op.get_operation()); + return Ok(()); + } + None => { + return pliron::input_err_noloc!( + "extract_field index {} out of bounds for aggregate with {} fields", + decl_index, + map.decl_to_llvm.len() + ); + } + }, + AggregateSlots::Identity => decl_index as u32, + }; + + let llvm_extract = llvm::ExtractValueOp::new(ctx, aggregate, vec![llvm_index])?; + rewriter.insert_operation(ctx, llvm_extract.get_operation()); + rewriter.replace_operation(ctx, op, llvm_extract.get_operation()); + + Ok(()) +} + +/// Convert `mir.insert_field` to `llvm.insertvalue`. +/// +/// Operands: `[aggregate, new_value]` +/// Returns a new aggregate with the field at `insert_index` replaced. +/// +/// The declaration-order field index is mapped to the LLVM slot via +/// [`resolve_aggregate_slots`] (arrays keep their element index). If +/// inserting a ZST field, we return the original aggregate unchanged. +pub(crate) fn convert_insert_field( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + operands_info: &OperandsInfo, +) -> Result<()> { + let aggregate = op.deref(ctx).get_operand(0); + let new_value = op.deref(ctx).get_operand(1); + + let insert_op = MirInsertFieldOp::new(op); + let decl_index = match insert_op.get_attr_insert_index(ctx) { + Some(attr) => attr.0 as usize, + None => return pliron::input_err_noloc!("Missing insert_index attribute on insert_field"), + }; + + if operands_info + .lookup_most_recent_of_type::(ctx, aggregate) + .is_some() + { + return convert_insert_union_field( + ctx, + rewriter, + op, + aggregate, + new_value, + decl_index, + operands_info, + ); + } + + let llvm_index = match resolve_aggregate_slots(ctx, operands_info, aggregate)? { + AggregateSlots::Mapped(map) => match map.decl_to_llvm.get(decl_index) { + Some(Some(slot)) => *slot, + Some(None) => { + // ZST field: stripped from the LLVM struct, so inserting + // into it is a no-op. Forward the aggregate unchanged. + rewriter.replace_operation_with_values(ctx, op, vec![aggregate]); + return Ok(()); + } + None => { + return pliron::input_err_noloc!( + "insert_field index {} out of bounds for aggregate with {} fields", + decl_index, + map.decl_to_llvm.len() + ); + } + }, + AggregateSlots::Identity => decl_index as u32, + }; + + let llvm_insert = llvm::InsertValueOp::new(ctx, aggregate, new_value, vec![llvm_index]); + rewriter.insert_operation(ctx, llvm_insert.get_operation()); + rewriter.replace_operation(ctx, op, llvm_insert.get_operation()); + + Ok(()) +} + +fn union_type_of_operand( + ctx: &Context, + operands_info: &OperandsInfo, + value: Value, +) -> Result { + operands_info + .lookup_most_recent_of_type::(ctx, value) + .map(|union_ty| union_ty.clone()) + .ok_or_else(|| { + pliron::create_error!( + pliron::location::Location::Unknown, + pliron::result::ErrorKind::VerificationFailed, + pliron::result::StringError( + "Expected MirUnionType conversion history for union value".to_string() + ) + ) + }) +} + +/// Read one typed view of a union's shared bytes. +fn convert_extract_union_field( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + union_value: Value, + field_index: usize, + operands_info: &OperandsInfo, +) -> Result<()> { + let union_ty = union_type_of_operand(ctx, operands_info, union_value)?; + let Some(field_mir_ty) = union_ty.get_field_type(field_index) else { + return pliron::input_err_noloc!( + "union field index {} is out of bounds for `{}`", + field_index, + union_ty.name() + ); + }; + let field_llvm_ty = convert_type(ctx, field_mir_ty).map_err(anyhow_to_pliron)?; + if is_zero_sized_type(ctx, field_llvm_ty) { + let undef = llvm::UndefOp::new(ctx, field_llvm_ty); + rewriter.insert_operation(ctx, undef.get_operation()); + rewriter.replace_operation(ctx, op, undef.get_operation()); + return Ok(()); + } + + let storage_ty = build_union_storage_type(ctx, &union_ty).map_err(anyhow_to_pliron)?; + let ptr = spill_enum_value(ctx, rewriter, union_value, storage_ty, union_ty.abi_align()); + let load = llvm::LoadOp::new(ctx, ptr, field_llvm_ty); + llvm_export::ops::set_op_alignment(ctx, load.get_operation(), union_ty.abi_align() as u32); + rewriter.insert_operation(ctx, load.get_operation()); + rewriter.replace_operation(ctx, op, load.get_operation()); + Ok(()) +} + +/// Write one typed view at byte zero while preserving the rest of the union. +fn convert_insert_union_field( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + union_value: Value, + new_value: Value, + field_index: usize, + operands_info: &OperandsInfo, +) -> Result<()> { + let union_ty = union_type_of_operand(ctx, operands_info, union_value)?; + let Some(field_mir_ty) = union_ty.get_field_type(field_index) else { + return pliron::input_err_noloc!( + "union field index {} is out of bounds for `{}`", + field_index, + union_ty.name() + ); + }; + let field_llvm_ty = convert_type(ctx, field_mir_ty).map_err(anyhow_to_pliron)?; + if is_zero_sized_type(ctx, field_llvm_ty) { + rewriter.replace_operation_with_values(ctx, op, vec![union_value]); + return Ok(()); + } + + let storage_ty = build_union_storage_type(ctx, &union_ty).map_err(anyhow_to_pliron)?; + let ptr = spill_enum_value(ctx, rewriter, union_value, storage_ty, union_ty.abi_align()); + let store = llvm::StoreOp::new(ctx, new_value, ptr); + llvm_export::ops::set_op_alignment(ctx, store.get_operation(), union_ty.abi_align() as u32); + rewriter.insert_operation(ctx, store.get_operation()); + + let load = llvm::LoadOp::new(ctx, ptr, storage_ty); + llvm_export::ops::set_op_alignment(ctx, load.get_operation(), union_ty.abi_align() as u32); + rewriter.insert_operation(ctx, load.get_operation()); + rewriter.replace_operation(ctx, op, load.get_operation()); + Ok(()) +} + +#[cfg(test)] +// Tests build kinded fixture types directly; production minting lives in mir-importer's facts.rs. +#[allow(clippy::disallowed_methods)] +mod tests { + use super::*; + use crate::convert::ops::test_util::*; + use dialect_mir::attributes::FieldIndexAttr; + use dialect_mir::ops as mir; + + use llvm_export::types as llvm_types; + + use super::super::test_support::*; + + /// Extracting a ZST field must not emit `extract_value`: the field has no + /// storage in the lowered LLVM struct, so lowering materializes an undef + /// zero-sized value instead. + #[test] + fn extract_zst_field_lowers_to_undef_without_extract_value() { + let mut ctx = make_ctx(); + + let i8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Signless).into(); + let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signless).into(); + let (struct_ty, zst_ty) = padded_struct_with_zst_ty(&mut ctx); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![i8_ty, i64_ty], vec![]); + let a = block.deref(&ctx).get_argument(0); + let b = block.deref(&ctx).get_argument(1); + let marker = append_empty_struct_value(&mut ctx, block, zst_ty); + + let construct = Operation::new( + &mut ctx, + mir::MirConstructStructOp::get_concrete_op_info(), + vec![struct_ty], + vec![a, marker, b], + vec![], + 0, + ); + construct.insert_at_back(block, &ctx); + let aggregate = construct.deref(&ctx).get_result(0); + + let extract = Operation::new( + &mut ctx, + MirExtractFieldOp::get_concrete_op_info(), + vec![zst_ty], + vec![aggregate], + vec![], + 0, + ); + MirExtractFieldOp::new(extract).set_attr_index(&ctx, FieldIndexAttr(1)); + extract.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + + assert_eq!( + count_ops::(&ctx, &body), + 0, + "extracting a stripped ZST field must not emit llvm.extractvalue" + ); + + let zst_un_defs = find_all::(&ctx, &body) + .into_iter() + .filter(|op| { + let result_ty = op.get_operation().deref(&ctx).get_result(0).get_type(&ctx); + is_zero_sized_type(&ctx, result_ty) + }) + .count(); + + assert_eq!( + zst_un_defs, 2, + "one undef should build the ZST value and one should materialize the extracted ZST" + ); + } + + /// The row-width arm of [`resolve_aggregate_slots`]' no-history fallback + /// is defensive symmetry: no current lowering path produces a runtime-width + /// slice value with an empty conversion history, so no end-to-end pipeline + /// test can reach it. This exercises the fallback directly: a value born as + /// the lowered `{ ptr, i64, i32 }` row-width shape (a block argument, which + /// carries no history) must resolve to identity indexing exactly like the + /// two-field `{ ptr, i64 }` fat pointer, and any other unrecognized + /// no-history shape must keep failing closed into the refuse-to-guess + /// error (issue #128). + #[test] + fn no_history_fallback_resolves_row_width_slice_shape_and_stays_closed_otherwise() { + let mut ctx = make_ctx(); + + // The exact struct the type converter lowers a runtime-width + // disjoint slice to, built through the same constructor the fallback + // uses so the test cannot drift from the lowered shape. + let width_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let row_width_ty = crate::convert::types::make_disjoint_slice_struct(&mut ctx, &[width_ty]) + .expect("building the row-width slice struct must succeed"); + + // A three-field control shape that is NOT the row-width slice: i64 + // where the u32 row-width word belongs. + use llvm_export::types::PointerTypeExt; + let ptr_ty: TypeHandle = llvm_types::PointerType::get_generic(&mut ctx).into(); + let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signless).into(); + let foreign_ty: TypeHandle = llvm_types::StructType::get_unnamed( + &ctx, + ( + vec![ptr_ty, i64_ty, i64_ty], + llvm_types::StructLayout::Unpacked, + ), + ) + .into(); + + let (_module, block) = build_kernel(&mut ctx, vec![row_width_ty, foreign_ty], vec![]); + let row_width_value = block.deref(&ctx).get_argument(0); + let foreign_value = block.deref(&ctx).get_argument(1); + + // Block arguments have no recorded conversion history at all. + let no_history = OperandsInfo::default(); + + let resolved = resolve_aggregate_slots(&mut ctx, &no_history, row_width_value) + .expect("the row-width slice shape with no history must resolve"); + assert!( + matches!(resolved, AggregateSlots::Identity), + "the row-width {{ ptr, i64, u32 }} shape is index-preserving and must map identically" + ); + + let refused = resolve_aggregate_slots(&mut ctx, &no_history, foreign_value); + assert!( + refused.is_err(), + "a no-history shape that is not a slice fat pointer (with or without \ + the row-width word) must fail closed" + ); + } +} diff --git a/crates/mir-lower/src/convert/ops/aggregate/mod.rs b/crates/mir-lower/src/convert/ops/aggregate/mod.rs new file mode 100644 index 0000000000..37a574d53c --- /dev/null +++ b/crates/mir-lower/src/convert/ops/aggregate/mod.rs @@ -0,0 +1,59 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Aggregate operation conversion: `dialect-mir` → LLVM dialect. +//! +//! Converts `dialect-mir` aggregate operations (structs, tuples, enums) to +//! their LLVM dialect equivalents. +//! +//! # Operations +//! +//! | MIR Operation | LLVM Operation(s) | Description | +//! |--------------------------|--------------------------------------|------------------------| +//! | `mir.extract_field` | `llvm.extractvalue` | Get struct/tuple field | +//! | `mir.insert_field` | `llvm.insertvalue` | Set struct/tuple field | +//! | `mir.construct_struct` | `llvm.undef` + `llvm.insertvalue` | Build struct | +//! | `mir.construct_tuple` | `llvm.undef` + `llvm.insertvalue` | Build tuple | +//! | `mir.construct_slice` | `llvm.undef` + `llvm.insertvalue` | Build slice fat ptr | +//! | `mir.construct_enum` | `llvm.undef` + `llvm.insertvalue` | Build enum | +//! | `mir.get_discriminant` | `llvm.extractvalue` | Get enum tag | +//! | `mir.set_discriminant` | `llvm.getelementptr` + `llvm.store` | Write enum tag | +//! | `mir.enum_payload` | `llvm.extractvalue` | Get enum payload | +//! +//! # Enum Representation +//! +//! Enums use rustc's physical layout, not a cuda-oxide-only tagged struct. +//! `build_enum_slot_map` places the direct tag or niche carrier and every +//! payload at rustc's byte offsets, reuses identical overlapping storage, and +//! routes differently typed overlaps through byte-addressed memory. `Single` +//! and `Empty` layouts have no carrier at all. +//! +//! A direct tag holds the variant's DECLARED discriminant value, not its +//! position. For `core::cmp::Ordering`, `Less` therefore stores -1 (the i8 bit +//! pattern 255), `Equal` stores 0, and `Greater` stores 1. A niche layout +//! instead uses rustc's wrapping `niche_start + variant_offset` encoding and +//! introduces no extra tag. + +mod addressing; +mod array_extract; +mod common; +mod construct; +mod enum_layout; +mod enums; +mod fields; +#[cfg(test)] +mod test_support; + +pub(crate) use addressing::{convert_array_element_addr, convert_field_addr}; +pub(crate) use array_extract::convert_extract_array_element; +pub(crate) use construct::{ + convert_construct_array, convert_construct_disjoint_slice, convert_construct_slice, + convert_construct_struct, convert_construct_tuple, +}; +pub(crate) use enums::{ + convert_construct_enum, convert_enum_payload, convert_get_discriminant, + convert_set_discriminant, +}; +pub(crate) use fields::{convert_extract_field, convert_insert_field}; diff --git a/crates/mir-lower/src/convert/ops/aggregate/test_support.rs b/crates/mir-lower/src/convert/ops/aggregate/test_support.rs new file mode 100644 index 0000000000..22d3a22d4a --- /dev/null +++ b/crates/mir-lower/src/convert/ops/aggregate/test_support.rs @@ -0,0 +1,109 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +// Tests build kinded fixture types directly; production minting lives in mir-importer's facts.rs. +#![allow(clippy::disallowed_methods)] + +use dialect_mir::ops as mir; +use dialect_mir::types::{ + EnumCarrierKind, EnumEncoding, EnumLayoutKind, EnumVariant, MirEnumType, MirStructType, +}; +use llvm_export::ops as llvm; +use pliron::builtin::types::{IntegerType, Signedness}; +use pliron::context::{Context, Ptr}; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::r#type::TypeHandle; +use pliron::value::Value; + +pub(super) fn insert_indices(ctx: &Context, inserts: &[llvm::InsertValueOp]) -> Vec> { + inserts.iter().map(|op| op.indices(ctx)).collect() +} + +pub(super) fn empty_struct_ty(ctx: &mut Context, name: &str) -> TypeHandle { + MirStructType::get(ctx, name.to_string(), vec![], vec![]).into() +} + +pub(super) fn padded_struct_with_zst_ty(ctx: &mut Context) -> (TypeHandle, TypeHandle) { + let i8_ty: TypeHandle = IntegerType::get(ctx, 8, Signedness::Signless).into(); + let i64_ty: TypeHandle = IntegerType::get(ctx, 64, Signedness::Signless).into(); + let zst_ty = empty_struct_ty(ctx, "Marker"); + + let struct_ty = MirStructType::get_with_full_layout( + ctx, + "Padded".to_string(), + vec!["a".to_string(), "marker".to_string(), "b".to_string()], + vec![i8_ty, zst_ty, i64_ty], + vec![0, 1, 2], + vec![0, 1, 8], + 16, + 8, + ); + + (struct_ty.into(), zst_ty) +} + +pub(super) fn append_empty_struct_value( + ctx: &mut Context, + block: Ptr, + zst_ty: TypeHandle, +) -> Value { + let op = Operation::new( + ctx, + mir::MirConstructStructOp::get_concrete_op_info(), + vec![zst_ty], + vec![], + vec![], + 0, + ); + op.insert_at_back(block, ctx); + op.deref(ctx).get_result(0) +} + +pub(super) fn unit_niche_enum( + ctx: &mut Context, + carrier: (EnumCarrierKind, u32, u32), + niche_start: u128, + niche_range: std::ops::RangeInclusive, + untagged_variant: u32, + inhabited: Vec, +) -> TypeHandle { + let (carrier_kind, carrier_width, carrier_address_space) = carrier; + let logical_width = if inhabited.len() > u8::MAX as usize { + 16 + } else { + 8 + }; + let logical_ty: TypeHandle = IntegerType::get(ctx, logical_width, Signedness::Unsigned).into(); + let variants = (0..inhabited.len()) + .map(|index| EnumVariant::unit(format!("V{index}"))) + .collect::>(); + let discriminants = (0..inhabited.len() as u64).collect(); + let carrier_size = u64::from(carrier_width).div_ceil(8); + let carrier_align = carrier_size.next_power_of_two().min(16); + MirEnumType::get_with_encoding( + ctx, + "UnitNiche".into(), + logical_ty, + discriminants, + variants, + EnumEncoding { + tag_offset: 0, + total_size: carrier_size, + abi_align: carrier_align, + layout_kind: EnumLayoutKind::Niche, + carrier_kind, + carrier_width, + carrier_address_space, + niche_start, + niche_variant_start: *niche_range.start(), + niche_variant_end: *niche_range.end(), + untagged_variant, + variant_inhabited: inhabited, + ..EnumEncoding::default() + }, + ) + .into() +} From 7557be25e3891a3ae763987362e414d0c5c2f356 Mon Sep 17 00:00:00 2001 From: nihalpasham Date: Fri, 28 Aug 2026 22:35:42 +0530 Subject: [PATCH 07/18] refactor(mir-lower): split ops/memory.rs into ops/memory/ modules --- crates/mir-lower/src/convert/ops/memory.rs | 4754 ----------------- .../src/convert/ops/memory/access.rs | 274 + .../src/convert/ops/memory/common.rs | 158 + .../mir-lower/src/convert/ops/memory/debug.rs | 92 + .../src/convert/ops/memory/device_global.rs | 410 ++ .../src/convert/ops/memory/extern_shared.rs | 209 + .../mir-lower/src/convert/ops/memory/mod.rs | 69 + .../src/convert/ops/memory/shared.rs | 342 ++ .../src/convert/ops/memory/tests/access.rs | 473 ++ .../src/convert/ops/memory/tests/alignment.rs | 806 +++ .../ops/memory/tests/debug_provenance.rs | 381 ++ .../convert/ops/memory/tests/device_global.rs | 555 ++ .../src/convert/ops/memory/tests/enums.rs | 322 ++ .../src/convert/ops/memory/tests/mod.rs | 83 + .../src/convert/ops/memory/tests/shared.rs | 560 ++ .../src/convert/ops/memory/transfer.rs | 237 + 16 files changed, 4971 insertions(+), 4754 deletions(-) delete mode 100644 crates/mir-lower/src/convert/ops/memory.rs create mode 100644 crates/mir-lower/src/convert/ops/memory/access.rs create mode 100644 crates/mir-lower/src/convert/ops/memory/common.rs create mode 100644 crates/mir-lower/src/convert/ops/memory/debug.rs create mode 100644 crates/mir-lower/src/convert/ops/memory/device_global.rs create mode 100644 crates/mir-lower/src/convert/ops/memory/extern_shared.rs create mode 100644 crates/mir-lower/src/convert/ops/memory/mod.rs create mode 100644 crates/mir-lower/src/convert/ops/memory/shared.rs create mode 100644 crates/mir-lower/src/convert/ops/memory/tests/access.rs create mode 100644 crates/mir-lower/src/convert/ops/memory/tests/alignment.rs create mode 100644 crates/mir-lower/src/convert/ops/memory/tests/debug_provenance.rs create mode 100644 crates/mir-lower/src/convert/ops/memory/tests/device_global.rs create mode 100644 crates/mir-lower/src/convert/ops/memory/tests/enums.rs create mode 100644 crates/mir-lower/src/convert/ops/memory/tests/mod.rs create mode 100644 crates/mir-lower/src/convert/ops/memory/tests/shared.rs create mode 100644 crates/mir-lower/src/convert/ops/memory/transfer.rs diff --git a/crates/mir-lower/src/convert/ops/memory.rs b/crates/mir-lower/src/convert/ops/memory.rs deleted file mode 100644 index 232b516d5e..0000000000 --- a/crates/mir-lower/src/convert/ops/memory.rs +++ /dev/null @@ -1,4754 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -//! Memory operation conversion: `dialect-mir` → LLVM dialect. -//! -//! Converts `dialect-mir` memory operations to their LLVM dialect equivalents. -//! -//! # Operations -//! -//! | MIR Operation | LLVM Operation(s) | Description | -//! |----------------------|-----------------------------------|------------------------------| -//! | `mir.load` | `llvm.load` | Load from pointer | -//! | `mir.store` | `llvm.store` | Store to pointer | -//! | `mir.ref` | `llvm.alloca` + `llvm.store` | Materialize aggregate in mem | -//! | `mir.ptr_offset` | `llvm.getelementptr` | Pointer arithmetic | -//! | `mir.shared_alloc` | `llvm.global` + `llvm.addressof` | Static shared memory | -//! | `mir.extern_shared` | `llvm.global` + `llvm.addressof` | Dynamic shared memory | -//! -//! # Shared Memory -//! -//! ## Static Shared Memory (`SharedArray`) -//! -//! Each static shared memory allocation gets a unique global symbol (`__shared_mem_N`). -//! Multiple allocations in the same or different kernels each have their own symbol -//! with their own size and alignment. -//! -//! ## Dynamic Shared Memory (`DynamicSharedArray`) -//! -//! Dynamic shared memory uses a symbol for each function that owns an access -//! (`__dynamic_smem_{function_name}`). -//! Key characteristics: -//! -//! - **Per-owner symbols**: Each function containing an access gets an extern symbol -//! - **Pre-computed alignment**: A pre-pass combines the owner's body alignment with -//! the strongest launch-contract marker that can reach it -//! - **Single runtime pool per launch**: The symbols refer to dynamic shared memory -//! sized by `shared_mem_bytes` at launch -//! -//! ### PTX Output Example -//! -//! ```ptx -//! ; Kernel with 128-byte aligned dynamic shared memory -//! .extern .shared .align 128 .b8 __dynamic_smem_my_kernel[]; -//! -//! ; Another kernel with 16-byte aligned (default) -//! .extern .shared .align 16 .b8 __dynamic_smem_other_kernel[]; -//! ``` - -use crate::context::{ - DeviceGlobalDeclaration, DeviceGlobalRecord, DeviceGlobalsMap, DynamicSmemAlignmentMap, - SharedGlobalDeclaration, SharedGlobalKind, SharedGlobalRecord, SharedGlobalsMap, -}; -use crate::convert::types::{ - StructLayoutInfo, build_struct_slot_map, convert_type, - llvm_packed_struct_contains_pointer_in_address_space, llvm_type_size_align, mir_element_stride, - mir_type_abi_align, validate_initialized_global_layout, - validate_relocated_initialized_global_layout, -}; -use crate::helpers; -use dialect_mir::types::{MirPtrType, MirStructType}; -use llvm_export::attributes::IntegerOverflowFlagsAttr; -use llvm_export::op_interfaces::IntBinArithOpWithOverflowFlag; -use llvm_export::ops as llvm; -use llvm_export::ops::GlobalOpExt; -use llvm_export::types::{ArrayType, FuncType, StructLayout, StructType, VoidType}; -use pliron::attribute::AttrObj; -use pliron::builtin::attributes::IntegerAttr; -use pliron::builtin::op_interfaces::CallOpCallable; -use pliron::builtin::op_interfaces::SymbolOpInterface; -use pliron::builtin::types::{IntegerType, Signedness}; -use pliron::context::{Context, Ptr}; -use pliron::identifier::Identifier; -use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; -use pliron::irbuild::inserter::Inserter; -use pliron::irbuild::rewriter::Rewriter; -use pliron::linked_list::ContainsLinkedList; -use pliron::location::Located; -use pliron::op::Op; -use pliron::operation::Operation; -use pliron::result::Result; -use pliron::r#type::{TypeHandle, Typed}; -use pliron::utils::apint::APInt; -use pliron::value::Value; - -fn anyhow_to_pliron(e: anyhow::Error) -> pliron::result::Error { - pliron::create_error!( - pliron::location::Location::Unknown, - pliron::result::ErrorKind::VerificationFailed, - pliron::result::StringError(e.to_string()) - ) -} - -/// Convert `mir.store` to `llvm.store`. -/// -/// Operand order: `[ptr, value]` - stores `value` to address `ptr`. -/// No result is produced (store is a side effect). -pub(crate) fn convert_store( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - operands_info: &OperandsInfo, -) -> Result<()> { - let operands: Vec<_> = op.deref(ctx).operands().collect(); - - let (ptr, val) = match operands.as_slice() { - [ptr, val] => (*ptr, *val), - _ => { - return pliron::input_err_noloc!("Store operation requires exactly 2 operands"); - } - }; - - // Packed whole-value stores are byte-faithful now that divergent rustc - // layouts lower to LLVM packed structs. Keep the target-dependent AS3 case - // fail-closed because its physical pointer width is selected only later. - fail_on_target_dependent_packed_aggregate( - ctx, - value_mir_type(ctx, operands_info, val), - "storing", - )?; - - let llvm_store = llvm::StoreOp::new(ctx, val, ptr); - if dialect_mir::ops::MirStoreOp::new(op).is_volatile(ctx) { - llvm_export::ops::set_op_volatile(ctx, llvm_store.get_operation(), true); - } - // The stored value's own type answers first, as it did before. A scalar - // records none, though, so fall back to whatever the address itself proved - // when it was computed -- for a field projection that is the aggregate's - // `abi_align` narrowed to the field's offset, which is otherwise lost here - // and costs the pair its vectorization. This mirrors `convert_load`, which - // consults the same record for the same reason. When both answer, the - // weaker wins: a field of a packed aggregate can place an abi-aligned type - // at a byte-aligned address, and the address's proved alignment is the - // ceiling of what the store may claim. - let abi = value_abi_align(ctx, operands_info, val); - let proved = pointer_proved_alignment(ctx, ptr); - let align = match (abi, proved) { - (Some(abi), Some(proved)) => Some(abi.min(proved)), - (abi, proved) => abi.or(proved), - }; - if let Some(align) = align { - llvm_export::ops::set_op_alignment(ctx, llvm_store.get_operation(), align as u32); - } - crate::convert::preserve_location(ctx, op, llvm_store.get_operation()); - rewriter.insert_operation(ctx, llvm_store.get_operation()); - rewriter.erase_operation(ctx, op); - Ok(()) -} - -/// Recover the `repr(align(N))` ABI alignment of `value` at conversion time. -/// -/// The alignment lives on MIR aggregate types (`abi_align`); the converted -/// LLVM struct types cannot express over-alignment. The conversion driver may -/// already have converted the value's type (block arguments are converted -/// before any rewrite runs; replaced op results carry the new type), but it -/// records every such change. So check the current type first, then walk the -/// value's conversion history, newest first, for a MIR type that records an -/// alignment. -fn value_abi_align(ctx: &Context, operands_info: &OperandsInfo, value: Value) -> Option { - mir_type_abi_align(ctx, value.get_type(ctx)).or_else(|| { - operands_info - .lookup_operand_history(value) - .iter() - .rev() - .find_map(|ty| mir_type_abi_align(ctx, *ty)) - }) -} - -/// The MIR type of `value`: its current type when that is still a MIR type, -/// else the newest MIR type in its conversion history (same walk as -/// [`value_abi_align`], for the same reason). -fn value_mir_type(ctx: &Context, operands_info: &OperandsInfo, value: Value) -> TypeHandle { - let current = value.get_type(ctx); - if current.deref(ctx).is::() { - return current; - } - operands_info - .lookup_operand_history(value) - .iter() - .rev() - .copied() - .find(|ty| ty.deref(ctx).is::()) - .unwrap_or(current) -} - -/// Refuse only packed by-value images whose physical bytes depend on the -/// selected NVPTX mode. -/// -/// Shared-memory pointers are 32-bit in modern NVVM p3:32 but 64-bit in the -/// PTX/legacy layouts. Lowering runs before that target mode is selected, so a -/// packed aggregate containing AS3 cannot safely be loaded or stored as one -/// physical value. Pointer-free packed aggregates and packed aggregates whose -/// pointers use target-stable address spaces are allowed. -fn fail_on_target_dependent_packed_aggregate( - ctx: &mut Context, - ty: TypeHandle, - verb: &str, -) -> Result<()> { - let llvm_ty = { - let layout = { - let ty_ref = ty.deref(ctx); - ty_ref - .downcast_ref::() - .map(StructLayoutInfo::of_struct) - }; - if let Some(layout) = layout { - let map = build_struct_slot_map(ctx, &layout).map_err(anyhow_to_pliron)?; - if !map.by_value_layout_faithful { - return pliron::input_err_noloc!( - "{} a struct whose rustc layout cannot be represented by an LLVM \ - struct value is not supported", - verb - ); - } - map.llvm_struct_ty - } else { - convert_type(ctx, ty).map_err(anyhow_to_pliron)? - } - }; - if llvm_packed_struct_contains_pointer_in_address_space( - ctx, - llvm_ty, - llvm_export::types::address_space::SHARED, - ) { - return pliron::input_err_noloc!( - "{} a packed aggregate containing a shared-memory pointer by value is \ - target-mode dependent and is not yet supported", - verb - ); - } - Ok(()) -} - -fn copy_debug_local_variable(ctx: &mut Context, mir_op: Ptr, llvm_op: Ptr) { - if let Some(info) = llvm_export::ops::debug_local_variable(ctx, mir_op) { - llvm_export::ops::set_debug_local_variable(ctx, llvm_op, info); - } - let projected = llvm_export::ops::debug_projected_variables(ctx, mir_op); - if !projected.is_empty() { - llvm_export::ops::set_debug_projected_variables(ctx, llvm_op, &projected); - } - let fragments = llvm_export::ops::debug_fragment_variables(ctx, mir_op); - if !fragments.is_empty() { - llvm_export::ops::set_debug_fragment_variables(ctx, llvm_op, &fragments); - } - if let Some(scope) = llvm_export::ops::debug_local_source_scope(ctx, mir_op) { - llvm_export::ops::set_debug_local_source_scope(ctx, llvm_op, scope); - } - if let Some((file, pos)) = llvm_export::ops::debug_local_declaration_location(ctx, mir_op) { - llvm_export::ops::set_debug_local_declaration_location( - ctx, llvm_op, file, pos.line, pos.column, - ); - } - if let Some(expression) = llvm_export::ops::debug_value_expression(ctx, mir_op) { - llvm_export::ops::set_debug_value_expression(ctx, llvm_op, &expression); - } -} - -/// Carry Rust-local provenance from the MIR alloca to the LLVM alloca. -/// -/// This is compiler-only metadata stored in the in-memory dialect operation; -/// it is not emitted as LLVM metadata because an unrecognized metadata node is -/// not a preservation contract for middle-end passes. The textual exporter -/// instead consumes the attribute to choose a stable SSA name before `opt`. -fn copy_local_memory_provenance( - ctx: &mut Context, - mir_op: Ptr, - llvm_op: Ptr, -) { - if let Some(provenance) = llvm_export::ops::local_memory_provenance(ctx, mir_op) { - llvm_export::ops::set_local_memory_provenance(ctx, llvm_op, provenance); - } -} - -/// Convert `mir.memcpy` to the matching `llvm.memcpy.p.p.i`. -/// -/// MIR's count is measured in pointee elements, while LLVM's memcpy intrinsic -/// expects bytes. The pre-conversion destination pointer type still carries the -/// MIR pointee, so use it to scale the count before emitting the call. -/// -/// The intrinsic name is an overload: LLVM encodes each pointer's address -/// space and the length width into it, and its verifier rejects a call whose -/// argument types disagree with the name. Today every pointer reaching a -/// `copy_nonoverlapping` is a Rust raw pointer, which cuda-oxide normalizes to -/// the generic address space (an `addrspacecast` is inserted when the raw -/// pointer is formed), so the operands are always `p0` and `i64`. We still -/// derive the suffix from the real operand types rather than hardcoding -/// `p0.p0.i64`: it matches how every other overloaded intrinsic is named here -/// (`ctpop`, `fptosi.sat`, ...), and it keeps this lowering correct if raw -/// pointers ever start carrying a non-generic address space. -pub(crate) fn convert_memcpy( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - operands_info: &OperandsInfo, -) -> Result<()> { - convert_mem_transfer(ctx, rewriter, op, operands_info, "memcpy") -} - -/// Convert `mir.memmove` to the matching `llvm.memmove.p.p.i`. -/// -/// Identical to [`convert_memcpy`] except it emits the overlap-safe -/// `llvm.memmove` intrinsic. `mir.memmove` backs `core::intrinsics::copy` -/// (`ptr::copy`); `mir.memcpy` backs the non-overlapping variant. -pub(crate) fn convert_memmove( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - operands_info: &OperandsInfo, -) -> Result<()> { - convert_mem_transfer(ctx, rewriter, op, operands_info, "memmove") -} - -/// Shared lowering for `mir.memcpy` / `mir.memmove`. `intrinsic_base` selects -/// the LLVM intrinsic family ("memcpy" or "memmove"); both share the same -/// `(dst, src, len_bytes, isvolatile)` signature and element->byte count scaling. -/// -/// The element type comes from the op's own `elem_type` attribute, stamped -/// at build time from dst's pointer type. Operand type history is not -/// usable here: a kind-only `mir.cast` lowers to a plain value forwarding, -/// history does not follow that edge, and a stale hit would scale the byte -/// count by the wrong element size. -fn convert_mem_transfer( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - _operands_info: &OperandsInfo, - intrinsic_base: &str, -) -> Result<()> { - let operands: Vec<_> = op.deref(ctx).operands().collect(); - let (dst, src, count) = match operands.as_slice() { - [dst, src, count] => (*dst, *src, *count), - _ => { - return pliron::input_err_noloc!( - "{intrinsic_base} operation requires exactly 3 operands" - ); - } - }; - - let pointee = { - let op_ref = op.deref(ctx); - op_ref - .attributes - .get::( - &format!("{intrinsic_base}_elem_type").try_into().unwrap(), - ) - .map(|attr| attr.get_type(ctx)) - .ok_or_else(|| { - pliron::create_error!( - op_ref.loc(), - pliron::result::ErrorKind::VerificationFailed, - pliron::result::StringError(format!( - "mir.{intrinsic_base} missing its elem_type attribute; \ - byte count has no fact to derive from" - )) - ) - })? - }; - // Byte-count policy: exact or error, never guessed. rustc's stride of the - // stamped MIR elem type is the primary fact; the converted LLVM type's - // natural layout is the fallback. - let elem_size = match mir_element_stride(ctx, pointee) { - Some(stride) => stride, - None => { - let elem_ty = convert_type(ctx, pointee).map_err(anyhow_to_pliron)?; - match llvm_type_size_align(ctx, elem_ty) { - Some((size, _)) => size, - None => { - let type_display = pointee.deref(ctx).disp(ctx).to_string(); - return Err(pliron::create_error!( - op.deref(ctx).loc(), - pliron::result::ErrorKind::VerificationFailed, - pliron::result::StringError(format!( - "mir.{intrinsic_base} element type `{type_display}` has no exact \ - byte size; refusing to guess the copy length" - )) - )); - } - } - } - }; - - let bytes = if elem_size == 1 { - count - } else { - let count_ty = count.get_type(ctx); - let bits = count_ty - .deref(ctx) - .downcast_ref::() - .map(|ty| ty.width()) - .unwrap_or(64); - let count_int_ty = IntegerType::get(ctx, bits, Signedness::Signless); - let size_attr: AttrObj = IntegerAttr::new( - count_int_ty, - APInt::from_u64( - elem_size, - std::num::NonZeroUsize::new(bits as usize).unwrap(), - ), - ) - .into(); - let size_const = llvm::ConstantOp::new(ctx, size_attr); - let size_val = size_const.get_operation().deref(ctx).get_result(0); - rewriter.insert_operation(ctx, size_const.get_operation()); - - let flags = IntegerOverflowFlagsAttr::default(); - let mul = llvm::MulOp::new_with_overflow_flag(ctx, count, size_val, flags); - rewriter.insert_operation(ctx, mul.get_operation()); - mul.get_operation().deref(ctx).get_result(0) - }; - - let i1_ty = IntegerType::get(ctx, 1, Signedness::Signless); - let false_attr: AttrObj = IntegerAttr::new( - i1_ty, - APInt::from_u64(0, std::num::NonZeroUsize::new(1).unwrap()), - ) - .into(); - let volatile = llvm::ConstantOp::new(ctx, false_attr); - rewriter.insert_operation(ctx, volatile.get_operation()); - let volatile_val = volatile.get_operation().deref(ctx).get_result(0); - - let void_ty = VoidType::get(ctx); - let func_ty = FuncType::get( - ctx, - void_ty.into(), - vec![ - dst.get_type(ctx), - src.get_type(ctx), - bytes.get_type(ctx), - volatile_val.get_type(ctx), - ], - false, - ); - let parent_block = op.deref(ctx).get_parent_block().ok_or_else(|| { - pliron::create_error!( - op.deref(ctx).loc(), - pliron::result::ErrorKind::VerificationFailed, - pliron::result::StringError(format!("{intrinsic_base} operation has no parent block")) - ) - })?; - // Derive the overload suffix from the real (already type-converted) - // operands so the name can never disagree with the argument types. - let dst_ty = dst.get_type(ctx); - let dst_as = dst_ty - .deref(ctx) - .downcast_ref::() - .map(|pt| pt.address_space()) - .unwrap_or(0); - let src_ty = src.get_type(ctx); - let src_as = src_ty - .deref(ctx) - .downcast_ref::() - .map(|pt| pt.address_space()) - .unwrap_or(0); - let len_bits = bytes - .get_type(ctx) - .deref(ctx) - .downcast_ref::() - .map(|t| t.width()) - .unwrap_or(64); - let intrinsic_name = format!("llvm_{intrinsic_base}_p{dst_as}_p{src_as}_i{len_bits}"); - helpers::ensure_intrinsic_declared(ctx, parent_block, &intrinsic_name, func_ty) - .map_err(anyhow_to_pliron)?; - - let callee: Identifier = intrinsic_name.as_str().try_into().map_err(|e| { - pliron::create_error!( - op.deref(ctx).loc(), - pliron::result::ErrorKind::VerificationFailed, - pliron::result::StringError(format!("Invalid memcpy intrinsic name: {e:?}")) - ) - })?; - let call = llvm::CallOp::new( - ctx, - CallOpCallable::Direct(callee), - func_ty, - vec![dst, src, bytes, volatile_val], - ); - crate::convert::preserve_location(ctx, op, call.get_operation()); - rewriter.insert_operation(ctx, call.get_operation()); - rewriter.erase_operation(ctx, op); - Ok(()) -} - -/// Alignment the op that computed `ptr` recorded about the address it produced. -/// -/// Field addresses stamp this during their own conversion, where the aggregate's -/// `abi_align` is still in hand; by the time a load that consumes the address is -/// converted, `mir.field_addr` is already a GEP and the MIR aggregate type is no -/// longer reachable from the load's operands. -/// -/// Known fidelity limit: the claim rides the address value, not the place -/// access that justified it. `&raw const (*p).x` through a merely 4-aligned -/// `p: *const Aligned8` is legal Rust so long as the place itself is never -/// accessed, yet the later read through the raw pointer consumes this same -/// stamped GEP and inherits the aggregate's alignment, an over-claim for that -/// pointer. rustc keeps the two apart by threading alignment through place -/// evaluation and dropping it at the `&raw` boundary; dialect-mir erases that -/// distinction before lowering runs, so full fidelity would require stamping -/// loads during MIR translation instead. Accepted for now: the pattern needs -/// an actually under-aligned pointer to an over-aligned aggregate, projected -/// via `&raw` inside a kernel, and any direct access through such a place is -/// UB to begin with. -fn pointer_proved_alignment(ctx: &Context, ptr: Value) -> Option { - let defining_op = ptr.defining_op()?; - llvm_export::ops::address_alignment(ctx, defining_op).map(u64::from) -} - -/// Convert `mir.load` to `llvm.load`. -/// -/// Takes a single pointer operand and returns the loaded value. -/// The result type is derived from the MIR operation's result type. -pub(crate) fn convert_load( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - _operands_info: &OperandsInfo, -) -> Result<()> { - let ptr = op.deref(ctx).get_operand(0); - let result_ty = op.deref(ctx).get_result(0).get_type(ctx); - - // Packed whole-value loads are byte-faithful now that divergent rustc - // layouts lower to LLVM packed structs. Keep only the target-dependent AS3 - // physical-image case fail-closed. - fail_on_target_dependent_packed_aggregate(ctx, result_ty, "loading")?; - - let llvm_ty = convert_type(ctx, result_ty).map_err(anyhow_to_pliron)?; - - let llvm_load = llvm::LoadOp::new(ctx, ptr, llvm_ty); - if dialect_mir::ops::MirLoadOp::new(op).is_volatile(ctx) { - llvm_export::ops::set_op_volatile(ctx, llvm_load.get_operation(), true); - } - // The loaded value's ABI alignment comes from this op's own result type, - // which is still the MIR type: result types are only converted by the - // op's own rewrite. A scalar records none, so fall back to whatever the - // address itself proved when it was computed -- for a field projection - // that is the aggregate's `abi_align` narrowed to the field's offset, - // which is otherwise lost here and costs the pair its vectorization. - // When both answer, the weaker wins: a field of a packed aggregate can - // place an abi-aligned type at a byte-aligned address, and the address's - // proved alignment is the ceiling of what the load may claim. - let abi = mir_type_abi_align(ctx, result_ty); - let proved = pointer_proved_alignment(ctx, ptr); - let align = match (abi, proved) { - (Some(abi), Some(proved)) => Some(abi.min(proved)), - (abi, proved) => abi.or(proved), - }; - if let Some(align) = align { - llvm_export::ops::set_op_alignment(ctx, llvm_load.get_operation(), align as u32); - } - rewriter.insert_operation(ctx, llvm_load.get_operation()); - rewriter.replace_operation(ctx, op, llvm_load.get_operation()); - - Ok(()) -} - -/// Convert `mir.dbg_value` to the LLVM-export debug marker. -/// -/// The op is still debug-only after lowering. The textual LLVM exporter later -/// prints it as an `llvm.dbg.value` intrinsic call. -pub(crate) fn convert_dbg_value( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - _operands_info: &OperandsInfo, -) -> Result<()> { - let value = op.deref(ctx).get_operand(0); - let loc = op.deref(ctx).loc().clone(); - let llvm_dbg_value = llvm::DebugValueOp::new(ctx, value); - llvm_dbg_value.get_operation().deref_mut(ctx).set_loc(loc); - copy_debug_local_variable(ctx, op, llvm_dbg_value.get_operation()); - rewriter.insert_operation(ctx, llvm_dbg_value.get_operation()); - rewriter.erase_operation(ctx, op); - Ok(()) -} - -/// Convert `mir.dbg_value_list` to the LLVM-export multi-value debug marker. -/// -/// The ordered operands become a `DIArgList` during textual export. The typed -/// location recipe is carried as generic metadata and copied unchanged here. -pub(crate) fn convert_dbg_value_list( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - _operands_info: &OperandsInfo, -) -> Result<()> { - let values: Vec<_> = op.deref(ctx).operands().collect(); - if values.len() < 2 { - return pliron::input_err_noloc!("mir.dbg_value_list requires at least two operands"); - } - if llvm_export::ops::debug_value_expression(ctx, op).is_none() { - return pliron::input_err_noloc!("mir.dbg_value_list is missing its debug expression"); - } - - let loc = op.deref(ctx).loc().clone(); - let llvm_dbg_value = llvm::DebugValueListOp::new(ctx, values); - llvm_dbg_value.get_operation().deref_mut(ctx).set_loc(loc); - copy_debug_local_variable(ctx, op, llvm_dbg_value.get_operation()); - rewriter.insert_operation(ctx, llvm_dbg_value.get_operation()); - rewriter.erase_operation(ctx, op); - Ok(()) -} - -/// Convert `mir.alloca` to `llvm.alloca`. -/// -/// `mir.alloca` carries its element type on the result pointer's pointee, and -/// emits a single-element stack slot of that type. We therefore convert the -/// pointee to an LLVM type and emit `llvm.alloca , i32 1`. -/// -/// No value is stored into the slot; that is the caller's job via subsequent -/// `mir.store` / `llvm.store` ops. This matches the mem2reg-ready translator -/// model where every local is backed by one alloca in the entry block and -/// defs/uses go through `store`/`load` rather than SSA values. -pub(crate) fn convert_alloca( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - _operands_info: &OperandsInfo, -) -> Result<()> { - let result_ty = op.deref(ctx).get_result(0).get_type(ctx); - let mir_pointee = { - let ty_ref = result_ty.deref(ctx); - let mir_ptr = ty_ref.downcast_ref::().ok_or_else(|| { - anyhow_to_pliron(anyhow::anyhow!( - "MirAllocaOp result must be MirPtrType (enforced by verifier)" - )) - })?; - mir_ptr.pointee - }; - let llvm_pointee = convert_type(ctx, mir_pointee).map_err(anyhow_to_pliron)?; - - let i32_ty = IntegerType::get(ctx, 32, Signedness::Signless); - let one_apint = - pliron::utils::apint::APInt::from_i64(1, std::num::NonZeroUsize::new(32).unwrap()); - let one_attr = pliron::builtin::attributes::IntegerAttr::new(i32_ty, one_apint); - let one_const = llvm::ConstantOp::new(ctx, one_attr.into()); - rewriter.insert_operation(ctx, one_const.get_operation()); - let one_val = one_const.get_operation().deref(ctx).get_result(0); - - let alloca = llvm::AllocaOp::new(ctx, llvm_pointee, one_val); - // The allocated type's ABI alignment comes from this op's own result - // pointee, which is still the MIR type at rewrite time. - if let Some(align) = mir_type_abi_align(ctx, mir_pointee) { - llvm_export::ops::set_op_alignment(ctx, alloca.get_operation(), align as u32); - } - copy_debug_local_variable(ctx, op, alloca.get_operation()); - copy_local_memory_provenance(ctx, op, alloca.get_operation()); - rewriter.insert_operation(ctx, alloca.get_operation()); - rewriter.replace_operation(ctx, op, alloca.get_operation()); - - Ok(()) -} - -/// Convert `mir.ref` — materialize the operand in stack memory via alloca+store. -/// -/// `mir.ref` creates a pointer to an SSA value. In SSA form, values don't have -/// addresses, so we must place the value in memory to obtain a pointer. -/// This applies to all types: scalars (e.g. `&factor` where factor is `u32`), -/// aggregates (e.g. `&closure_env`), and pointers (e.g. `&&T`). -pub(crate) fn convert_ref( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - operands_info: &OperandsInfo, -) -> Result<()> { - let operand = op.deref(ctx).get_operand(0); - let operand_ty = operand.get_type(ctx); - let abi_align = value_abi_align(ctx, operands_info, operand); - - let i32_ty = IntegerType::get(ctx, 32, Signedness::Signless); - let one_apint = - pliron::utils::apint::APInt::from_i64(1, std::num::NonZeroUsize::new(32).unwrap()); - let one_attr = pliron::builtin::attributes::IntegerAttr::new(i32_ty, one_apint); - let one_const = llvm::ConstantOp::new(ctx, one_attr.into()); - rewriter.insert_operation(ctx, one_const.get_operation()); - let one_val = one_const.get_operation().deref(ctx).get_result(0); - - let alloca = llvm::AllocaOp::new(ctx, operand_ty, one_val); - // Honour the referent's repr(align(N)) ABI alignment. Without this, the - // synthesised alloca would be under-aligned relative to any loads/stores - // that claim the struct's true alignment. - if let Some(align) = abi_align { - llvm_export::ops::set_op_alignment(ctx, alloca.get_operation(), align as u32); - } - rewriter.insert_operation(ctx, alloca.get_operation()); - let alloca_ptr = alloca.get_operation().deref(ctx).get_result(0); - - let store = llvm::StoreOp::new(ctx, operand, alloca_ptr); - if let Some(align) = abi_align { - llvm_export::ops::set_op_alignment(ctx, store.get_operation(), align as u32); - } - rewriter.insert_operation(ctx, store.get_operation()); - - rewriter.replace_operation_with_values(ctx, op, vec![alloca_ptr]); - - Ok(()) -} - -/// Convert `mir.ptr_offset` to `llvm.getelementptr`. -/// -/// Operands: `[ptr, offset]` where offset is an integer index. -/// Element sizing comes from the op's own result type, which is still the -/// MIR pointer type when this converter runs. The operand's recorded type -/// history is not usable here: a kind-only `mir.cast` lowers to a plain -/// value forwarding, and the history does not follow that replacement -/// edge, so a history miss would silently misscale the offset. -pub(crate) fn convert_ptr_offset( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - _operands_info: &OperandsInfo, -) -> Result<()> { - let loc = op.deref(ctx).loc(); - let operands: Vec<_> = op.deref(ctx).operands().collect(); - - let (ptr, offset) = match operands.as_slice() { - [ptr, offset] => (*ptr, *offset), - _ => return pliron::input_err!(loc, "PtrOffset requires exactly 2 operands"), - }; - - let result_ty = op.deref(ctx).get_result(0).get_type(ctx); - let pointee = result_ty - .deref(ctx) - .downcast_ref::() - .map(|mir_ptr| mir_ptr.pointee) - .ok_or_else(|| { - pliron::input_error!( - loc.clone(), - "mir.ptr_offset result must be a MIR pointer type; \ - element sizing has no fact to derive from" - ) - })?; - let elem_ty = convert_type(ctx, pointee).map_err(anyhow_to_pliron)?; - - let llvm_gep = llvm::GetElementPtrOp::new( - ctx, - ptr, - vec![llvm_export::ops::GepIndex::Value(offset)], - elem_ty, - ); - let inbounds = dialect_mir::ops::MirPtrOffsetOp::new(op).is_inbounds(ctx); - llvm::set_gep_inbounds(ctx, llvm_gep.get_operation(), inbounds); - rewriter.insert_operation(ctx, llvm_gep.get_operation()); - rewriter.replace_operation(ctx, op, llvm_gep.get_operation()); - - Ok(()) -} - -/// Convert `mir.shared_alloc` to LLVM global variable in shared address space. -/// -/// GPU shared memory is represented as a global variable with address space 3. -/// Uses `shared_globals` to deduplicate: multiple allocations with the same -/// `alloc_key` share the same global. -/// -/// Called directly from `MirToLlvmConversionDriver::rewrite` (not through -/// op_cast dispatch) because it needs the cross-function `SharedGlobalsMap`. -pub fn convert_shared_alloc_dc( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - _operands_info: &OperandsInfo, - shared_globals: &mut SharedGlobalsMap, - next_shared_mem_index: &mut usize, -) -> Result<()> { - use pliron::builtin::attributes::{IntegerAttr, TypeAttr}; - - let (alloc_key, source_name, mir_elem_type, size, alignment, debug_info, debug_owner_function) = { - let shared_alloc_op = dialect_mir::ops::MirSharedAllocOp::new(op); - let op_ref = op.deref(ctx); - - let alloc_key: Option = shared_alloc_op - .get_attr_alloc_key(ctx) - .map(|s| String::from((*s).clone())); - - // Optional and diagnostic: the Rust path of the originating `static`, - // carried through so the emitted global can name its source. - let source_name: Option = shared_alloc_op - .get_attr_source_name(ctx) - .map(|s| String::from((*s).clone())); - - let elem_type_attr = op_ref - .attributes - .get::(&"elem_type".try_into().unwrap()) - .ok_or_else(|| { - anyhow_to_pliron(anyhow::anyhow!( - "MirSharedAllocOp missing elem_type TypeAttr attribute" - )) - })?; - let mir_elem_type = elem_type_attr.get_type(ctx); - - let size_attr = op_ref - .attributes - .get::(&"size".try_into().unwrap()) - .ok_or_else(|| { - anyhow_to_pliron(anyhow::anyhow!( - "MirSharedAllocOp missing size IntegerAttr attribute" - )) - })?; - let size = size_attr.value().to_u64(); - - let alignment = shared_alloc_op.get_alignment_value(ctx).unwrap_or(0); - let debug_info = llvm::debug_global_variable(ctx, op); - let debug_owner_function = llvm::debug_global_owner_function(ctx, op); - - ( - alloc_key, - source_name, - mir_elem_type, - size, - alignment, - debug_info, - debug_owner_function, - ) - }; - - // A shared variable DIE must describe the physical allocation. Gate the - // optional debug identity against the lowered element layout before it - // can reach the declaration record or the emitted global. - let llvm_elem_type = convert_type(ctx, mir_elem_type).map_err(anyhow_to_pliron)?; - let debug_info = debug_info - .filter(|info| shared_debug_type_matches_physical(ctx, info, llvm_elem_type, size)); - let debug_owner_function = debug_info.as_ref().and(debug_owner_function); - - let declaration = SharedGlobalDeclaration { - kind: SharedGlobalKind::Static, - mir_elem_type, - size, - alignment, - debug_info, - debug_owner_function, - }; - - // A key names one physical allocation, not merely a preferred symbol. - // Reuse it only when the complete storage declaration agrees; otherwise - // the later address would silently inherit the first allocation's type, - // extent, or alignment. - let global_name = if let Some(key) = alloc_key.as_ref() - && let Some(existing) = shared_globals.get(key) - { - if !existing.declaration.same_storage(&declaration) { - return Err(anyhow_to_pliron(anyhow::anyhow!( - "duplicate shared alloc_key {:?} has an incompatible declaration", - key - ))); - } - let symbol = existing.symbol.clone(); - let debug_identity_matches = existing.declaration.debug_info == declaration.debug_info - && existing.declaration.debug_owner_function == declaration.debug_owner_function; - if !debug_identity_matches { - // One physical allocation reached under two debug identities, for - // example a function-local static whose owning function was - // materialized (or inlined) into two kernels. Debug metadata must - // never fail a build that a release build accepts, and DWARF - // cannot truthfully scope one AS3 object to two subprograms, so - // fail open: drop the attachment from the materialized global and - // demote the cached record to metadata-free. - strip_shared_debug_identity(ctx, op, &symbol)?; - if let Some(record) = shared_globals.get_mut(key) { - record.declaration.debug_info = None; - record.declaration.debug_owner_function = None; - } - } - symbol - } else { - create_shared_global( - ctx, - op, - shared_globals, - next_shared_mem_index, - SharedAllocSpec { - mir_elem_type, - size, - alignment, - alloc_key, - source_name: source_name.as_deref(), - debug_info: declaration.debug_info.as_ref(), - debug_owner_function: declaration.debug_owner_function.as_deref(), - }, - )? - }; - - let address_of_op = llvm::AddressOfOp::new(ctx, global_name, 3); - rewriter.insert_operation(ctx, address_of_op.get_operation()); - rewriter.replace_operation(ctx, op, address_of_op.get_operation()); - - Ok(()) -} - -/// Everything `create_shared_global` needs about one `mir.shared_alloc`. -/// -/// Mirrors [`DeviceGlobalSpec`] for the shared-memory path. -struct SharedAllocSpec<'a> { - mir_elem_type: TypeHandle, - size: u64, - alignment: u64, - alloc_key: Option, - source_name: Option<&'a str>, - debug_info: Option<&'a llvm::DebugGlobalVariableInfo>, - debug_owner_function: Option<&'a str>, -} - -/// Create a shared memory global variable in the module. -/// -/// Creates an LLVM global variable with: -/// - Array type: `[size x element_type]` -/// - Address space 3 (shared memory) -/// - Optional alignment -/// - Unique generated name (`__shared_mem_N`) -/// -/// The global is inserted at the front of the module block. When -/// `spec.alloc_key` is `Some`, the key is moved into `shared_globals` so that -/// later allocations with the same key reuse this global only after the caller -/// has checked that their complete declarations agree. -/// -/// `spec.source_name`, when present, is the Rust path of the `static` this -/// allocation came from. The generated symbol stays anonymous; the name is -/// recorded as an attribute on the global so the exporter can render it -/// beside the definition. Only the allocation that *creates* the global -/// contributes a name — a later allocation with the same `alloc_key` hits -/// the cache and never reaches this function — which is consistent because -/// the key and the name are both derived from the same constant. -/// -/// `next_shared_mem_index` is scoped to one `MirToLlvmConversionDriver` -/// instance (one module), not a process-global counter: `N` is a function of -/// this module's own MIR walk order, not of how many other modules have -/// lowered a shared allocation earlier in the process (#706). -fn create_shared_global( - ctx: &mut Context, - op: Ptr, - shared_globals: &mut SharedGlobalsMap, - next_shared_mem_index: &mut usize, - spec: SharedAllocSpec<'_>, -) -> Result { - let llvm_elem_type = convert_type(ctx, spec.mir_elem_type).map_err(anyhow_to_pliron)?; - let array_type = ArrayType::get(ctx, llvm_elem_type, spec.size); - - let counter = *next_shared_mem_index; - *next_shared_mem_index += 1; - let name: pliron::identifier::Identifier = - format!("__shared_mem_{counter}").try_into().unwrap(); - - let global_op = if spec.alignment > 0 { - llvm::GlobalOp::new_with_alignment(ctx, name.clone(), array_type.into(), spec.alignment) - } else { - llvm::GlobalOp::new(ctx, name.clone(), array_type.into()) - }; - global_op.set_address_space(ctx, llvm_export::types::address_space::SHARED); - if let Some(source_name) = spec.source_name { - use llvm_export::ops::GlobalOpExt; - global_op.set_shared_source_name(ctx, source_name); - } - if let Some(info) = spec.debug_info { - llvm::set_debug_global_variable(ctx, global_op.get_operation(), info); - } - if let Some(owner) = spec.debug_owner_function { - llvm::set_debug_global_owner_function(ctx, global_op.get_operation(), owner); - } - - let parent_block = op - .deref(ctx) - .get_parent_block() - .ok_or_else(|| anyhow_to_pliron(anyhow::anyhow!("Op has no parent block")))?; - let module_op = helpers::get_module_from_block(ctx, parent_block).map_err(anyhow_to_pliron)?; - let region = module_op.deref(ctx).get_region(0); - let module_block = region - .deref(ctx) - .iter(ctx) - .next() - .ok_or_else(|| anyhow_to_pliron(anyhow::anyhow!("Module is empty")))?; - - global_op.get_operation().insert_at_front(module_block, ctx); - - if let Some(key) = spec.alloc_key { - shared_globals.insert( - key, - SharedGlobalRecord { - symbol: name.clone(), - declaration: SharedGlobalDeclaration { - kind: SharedGlobalKind::Static, - mir_elem_type: spec.mir_elem_type, - size: spec.size, - alignment: spec.alignment, - debug_info: spec.debug_info.cloned(), - debug_owner_function: spec.debug_owner_function.map(str::to_owned), - }, - }, - ); - } - - Ok(name) -} - -/// A shared variable DIE must describe the compiler-materialized allocation, -/// never the marker ZST or an element graph whose pointer base types are not -/// representable yet. The frontend checks semantic element layout; this final -/// boundary verifies total physical bits/count against the LLVM backing. -fn shared_debug_type_matches_physical( - ctx: &Context, - info: &llvm::DebugGlobalVariableInfo, - llvm_elem_type: TypeHandle, - count: u64, -) -> bool { - let Some((element_size, _)) = llvm_type_size_align(ctx, llvm_elem_type) else { - return false; - }; - let Some(element_bytes) = element_size.checked_mul(count) else { - return false; - }; - let Some(physical_bits) = element_bytes.checked_mul(8) else { - return false; - }; - if info.ty.size_bits() != physical_bits { - return false; - } - match &info.ty { - llvm::DebugLocalTypeKind::Array { - element, - count: debug_count, - .. - } => *debug_count == count && element.size_bits() == element_size.saturating_mul(8), - // `Barrier` retains its semantic repr(C) struct while the backing is - // one i64. Other marker types are never admitted by the frontend. - llvm::DebugLocalTypeKind::Struct { .. } => count == 1, - _ => false, - } -} - -/// Detach the debug identity from an already-materialized shared global. -/// -/// The fail-open path for divergent debug identities on one `alloc_key`: -/// the storage stays shared and valid, only the optional DWARF attachment is -/// dropped so it cannot misattribute the allocation to the wrong static or -/// owner function. -fn strip_shared_debug_identity( - ctx: &mut Context, - op: Ptr, - name: &pliron::identifier::Identifier, -) -> Result<()> { - let parent_block = op - .deref(ctx) - .get_parent_block() - .ok_or_else(|| anyhow_to_pliron(anyhow::anyhow!("Op has no parent block")))?; - let module_op = helpers::get_module_from_block(ctx, parent_block).map_err(anyhow_to_pliron)?; - let module_block = module_op - .deref(ctx) - .get_region(0) - .deref(ctx) - .iter(ctx) - .next() - .ok_or_else(|| anyhow_to_pliron(anyhow::anyhow!("Module is empty")))?; - let existing = module_block - .deref(ctx) - .iter(ctx) - .filter_map(|candidate| Operation::get_op::(candidate, ctx)) - .find(|global| global.get_symbol_name(ctx) == *name) - .ok_or_else(|| { - anyhow_to_pliron(anyhow::anyhow!( - "shared allocation cache refers to missing LLVM global `@{name}`" - )) - })?; - llvm::clear_debug_global_identity(ctx, existing.get_operation()); - Ok(()) -} - -/// Convert `mir.global_alloc` to an LLVM global in CUDA global memory. -/// -/// Ordinary Rust `static` / `static mut` values have grid scope and -/// application lifetime, so they live in address space 1 rather than the -/// per-block shared-memory address space. -pub fn convert_global_alloc_dc( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - _operands_info: &OperandsInfo, - device_globals: &mut DeviceGlobalsMap, - next_device_global_index: &mut usize, -) -> Result<()> { - use pliron::builtin::attributes::{StringAttr, TypeAttr}; - - let ( - global_key, - mir_global_type, - alignment, - addr_space, - initializer_hex, - initializer_relocations, - immutable, - debug_info, - ) = { - let global_op = dialect_mir::ops::MirGlobalAllocOp::new(op); - let op_ref = op.deref(ctx); - - let global_key_attr = op_ref - .attributes - .get::(&"global_key".try_into().unwrap()) - .ok_or_else(|| { - anyhow_to_pliron(anyhow::anyhow!( - "MirGlobalAllocOp missing global_key StringAttr attribute" - )) - })?; - let global_key = String::from((*global_key_attr).clone()); - - let global_type_attr = op_ref - .attributes - .get::(&"global_type".try_into().unwrap()) - .ok_or_else(|| { - anyhow_to_pliron(anyhow::anyhow!( - "MirGlobalAllocOp missing global_type TypeAttr attribute" - )) - })?; - let mir_global_type = global_type_attr.get_type(ctx); - - let alignment = global_op.get_alignment_value(ctx).unwrap_or(0); - let initializer_hex = op_ref - .attributes - .get::(&"global_initializer_hex".try_into().unwrap()) - .map(|attr| String::from((*attr).clone())); - let initializer_relocations = op_ref - .attributes - .get::(&"global_initializer_relocations".try_into().unwrap()) - .map(|attr| String::from((*attr).clone())); - let debug_info = llvm::debug_global_variable(ctx, op); - - // Read the address space the op's result already carries — set by - // mir-importer based on the static's type (`ConstantMemory` → 4, - // ordinary → 1). The dialect verifier accepts both. - let res_ty = op_ref.get_result(0).get_type(ctx); - let addr_space = res_ty - .deref(ctx) - .downcast_ref::() - .map(|p| { - if p.address_space == dialect_mir::types::address_space::CONSTANT { - llvm_export::types::address_space::CONSTANT - } else { - llvm_export::types::address_space::GLOBAL - } - }) - .ok_or_else(|| { - anyhow_to_pliron(anyhow::anyhow!( - "MirGlobalAllocOp result is not a MirPtrType" - )) - })?; - - ( - global_key, - mir_global_type, - alignment, - addr_space, - initializer_hex, - initializer_relocations, - global_op.is_immutable(ctx), - debug_info, - ) - }; - - let declaration = DeviceGlobalDeclaration { - mir_type: mir_global_type, - alignment, - addr_space, - initializer_hex: initializer_hex.clone(), - initializer_relocations: initializer_relocations.clone(), - debug_info: debug_info.clone(), - immutable, - }; - - let global_name = if let Some(existing) = device_globals.get(&global_key) { - if existing.declaration != declaration { - return Err(anyhow_to_pliron(anyhow::anyhow!( - "duplicate global_key {:?} has an incompatible declaration", - global_key - ))); - } - existing.symbol.clone() - } else { - create_device_global( - ctx, - op, - device_globals, - next_device_global_index, - DeviceGlobalSpec { - key: &global_key, - mir_type: mir_global_type, - alignment, - addr_space, - initializer_hex: initializer_hex.as_deref(), - initializer_relocations: initializer_relocations.as_deref(), - debug_info: debug_info.as_ref(), - immutable, - }, - )? - }; - - let address_of_op = llvm::AddressOfOp::new(ctx, global_name, addr_space); - rewriter.insert_operation(ctx, address_of_op.get_operation()); - rewriter.replace_operation(ctx, op, address_of_op.get_operation()); - - Ok(()) -} - -struct DeviceGlobalSpec<'a> { - key: &'a str, - mir_type: TypeHandle, - alignment: u64, - addr_space: u32, - initializer_hex: Option<&'a str>, - initializer_relocations: Option<&'a str>, - debug_info: Option<&'a llvm::DebugGlobalVariableInfo>, - /// Nothing writes this storage, so it exports as LLVM `constant`. Set only - /// for the compiler's own promoted constants; see `MirGlobalAllocOp`. - immutable: bool, -} - -/// `next_device_global_index` is scoped to one `MirToLlvmConversionDriver` -/// instance (one module), not a process-global counter: `N` is a function of -/// this module's own MIR walk order, not of how many other modules have -/// lowered a device global earlier in the process (#706). -fn create_device_global( - ctx: &mut Context, - op: Ptr, - device_globals: &mut DeviceGlobalsMap, - next_device_global_index: &mut usize, - spec: DeviceGlobalSpec<'_>, -) -> Result { - // An explicit initializer is already the evaluated Rust allocation image. - // Pointer-free data stays `[N x i8]`. Initializers with relocations use a - // segmented LLVM struct whose literal spans remain byte arrays and whose - // pointer slots become pointer-width integers. This preserves both exact - // bytes and linker-visible pointer provenance. - let semantic_llvm_type = convert_type(ctx, spec.mir_type).map_err(anyhow_to_pliron)?; - let (llvm_global_type, alignment) = if let Some(initializer_hex) = spec.initializer_hex { - let byte_count = initializer_hex_byte_count(initializer_hex).map_err(anyhow_to_pliron)?; - if spec.alignment == 0 { - return Err(anyhow_to_pliron(anyhow::anyhow!( - "device global initializer is missing its evaluated Rust allocation alignment" - ))); - } - let storage_type = if let Some(encoded) = spec.initializer_relocations { - validate_relocated_initialized_global_layout( - ctx, - spec.mir_type, - byte_count, - spec.alignment, - ) - .map_err(anyhow_to_pliron)?; - relocated_initializer_storage_type(ctx, byte_count, spec.alignment, encoded) - .map_err(anyhow_to_pliron)? - } else { - validate_initialized_global_layout(ctx, spec.mir_type, byte_count, spec.alignment) - .map_err(anyhow_to_pliron)?; - let i8_ty = IntegerType::get(ctx, 8, Signedness::Signless); - ArrayType::get(ctx, i8_ty.into(), byte_count).into() - }; - (storage_type, spec.alignment) - } else { - if spec.initializer_relocations.is_some() { - return Err(anyhow_to_pliron(anyhow::anyhow!( - "device global carries relocation metadata without initializer bytes" - ))); - } - (semantic_llvm_type, spec.alignment) - }; - - // Constant-memory globals reuse the Rust-side mangled name so host code can - // resolve them by name via `cuModuleGetGlobal`. Ordinary device globals - // are private to the kernel and get a counter-based unique name. - let name: pliron::identifier::Identifier = - if spec.addr_space == llvm_export::types::address_space::CONSTANT { - let symbol = dialect_mir::ops::rust_static_symbol_from_global_key(spec.key) - .ok_or_else(|| { - anyhow_to_pliron(anyhow::anyhow!( - "constant global_key {:?} is not a tagged Rust static identity", - spec.key - )) - })?; - symbol.try_into().map_err(|e| { - anyhow_to_pliron(anyhow::anyhow!( - "constant Rust static symbol {:?} is not a valid identifier: {e:?}", - symbol - )) - })? - } else { - let counter = *next_device_global_index; - *next_device_global_index += 1; - format!("__device_global_{counter}").try_into().unwrap() - }; - - let global_op = if alignment > 0 { - llvm::GlobalOp::new_with_alignment(ctx, name.clone(), llvm_global_type, alignment) - } else { - llvm::GlobalOp::new(ctx, name.clone(), llvm_global_type) - }; - global_op.set_address_space(ctx, spec.addr_space); - global_op.set_source_global_key(ctx, spec.key); - if spec.addr_space == llvm_export::types::address_space::GLOBAL - && let Some(info) = spec.debug_info - { - llvm::set_debug_global_variable(ctx, global_op.get_operation(), info); - } - if let Some(initializer_hex) = spec.initializer_hex { - global_op.set_initializer_hex(ctx, initializer_hex); - } - if let Some(initializer_relocations) = spec.initializer_relocations { - global_op.set_initializer_relocations(ctx, initializer_relocations); - } - if spec.immutable { - global_op.mark_immutable(ctx); - } - - let parent_block = op - .deref(ctx) - .get_parent_block() - .ok_or_else(|| anyhow_to_pliron(anyhow::anyhow!("Op has no parent block")))?; - let module_op = helpers::get_module_from_block(ctx, parent_block).map_err(anyhow_to_pliron)?; - let region = module_op.deref(ctx).get_region(0); - let module_block = region - .deref(ctx) - .iter(ctx) - .next() - .ok_or_else(|| anyhow_to_pliron(anyhow::anyhow!("Module is empty")))?; - - global_op.get_operation().insert_at_front(module_block, ctx); - device_globals.insert( - spec.key.to_string(), - DeviceGlobalRecord { - symbol: name.clone(), - declaration: DeviceGlobalDeclaration { - mir_type: spec.mir_type, - alignment: spec.alignment, - addr_space: spec.addr_space, - initializer_hex: spec.initializer_hex.map(str::to_owned), - initializer_relocations: spec.initializer_relocations.map(str::to_owned), - debug_info: spec.debug_info.cloned(), - immutable: spec.immutable, - }, - }, - ); - - Ok(name) -} - -fn relocated_initializer_storage_type( - ctx: &mut Context, - byte_count: u64, - allocation_alignment: u64, - encoded: &str, -) -> std::result::Result { - let mut relocations = - llvm::decode_global_initializer_relocations(encoded).map_err(anyhow::Error::msg)?; - if relocations.is_empty() { - anyhow::bail!("device global relocation metadata contains no relocations"); - } - relocations.sort_by_key(|relocation| relocation.source_offset); - - let mut cursor = 0u64; - let mut fields = Vec::with_capacity(relocations.len() * 2 + 1); - let mut requires_packed_storage = false; - let i8_ty = IntegerType::get(ctx, 8, Signedness::Signless); - - for (index, relocation) in relocations.iter().enumerate() { - if relocation.width_bytes != 8 { - anyhow::bail!( - "device global relocation {index} uses unsupported {}-byte pointer storage; CUDA global/constant pointers require 8 bytes", - relocation.width_bytes - ); - } - if !matches!(relocation.target_address_space, 1 | 4) { - anyhow::bail!( - "device global relocation {index} targets unsupported CUDA address space {}", - relocation.target_address_space - ); - } - if relocation.target_key.is_empty() { - anyhow::bail!("device global relocation {index} has an empty target key"); - } - - let width = u64::from(relocation.width_bytes); - requires_packed_storage |= - allocation_alignment < width || !relocation.source_offset.is_multiple_of(width); - if relocation.source_offset < cursor { - anyhow::bail!( - "device global relocation {index} overlaps the previous relocation or literal span" - ); - } - let end = relocation - .source_offset - .checked_add(width) - .ok_or_else(|| anyhow::anyhow!("device global relocation {index} offset overflows"))?; - if end > byte_count { - anyhow::bail!( - "device global relocation {index} occupies bytes {}..{} but the initializer is only {} bytes", - relocation.source_offset, - end, - byte_count - ); - } - - if relocation.source_offset > cursor { - fields - .push(ArrayType::get(ctx, i8_ty.into(), relocation.source_offset - cursor).into()); - } - fields.push(IntegerType::get(ctx, relocation.width_bytes * 8, Signedness::Signless).into()); - cursor = end; - } - - if cursor < byte_count { - fields.push(ArrayType::get(ctx, i8_ty.into(), byte_count - cursor).into()); - } - - let layout = if requires_packed_storage { - StructLayout::Packed - } else { - StructLayout::Unpacked - }; - let storage: TypeHandle = StructType::get_unnamed(ctx, (fields, layout)).into(); - // Exact or error, never guessed: the storage type is built from i8 arrays - // and integer slots, so its natural size is always computable, and it must - // land exactly on rustc's byte count or the relocation offsets are wrong. - let Some((lowered_size, _)) = llvm_type_size_align(ctx, storage) else { - anyhow::bail!( - "relocated device global storage `{}` has no exact size", - storage.deref(ctx).disp(ctx) - ); - }; - if lowered_size != byte_count { - anyhow::bail!( - "relocated device global storage lowers to {} bytes, but rustc evaluated {} bytes", - lowered_size, - byte_count - ); - } - Ok(storage) -} - -fn initializer_hex_byte_count(hex: &str) -> std::result::Result { - if !hex.len().is_multiple_of(2) { - anyhow::bail!("device global initializer has an odd-length hex byte string"); - } - if let Some(invalid) = hex.bytes().find(|byte| !byte.is_ascii_hexdigit()) { - anyhow::bail!( - "device global initializer contains invalid hex digit {:?}", - invalid as char - ); - } - u64::try_from(hex.len() / 2) - .map_err(|_| anyhow::anyhow!("device global initializer is too large for LLVM")) -} - -/// Convert `mir.extern_shared` to LLVM extern global variable in shared address space. -/// -/// Dynamic (extern) shared memory is represented as an external global variable -/// with address space 3 and zero-length array type `[0 x i8]`. The actual size -/// is determined at kernel launch via `LaunchConfig::shared_mem_bytes`. -/// -/// # Per-Owner Symbols -/// -/// Each function that owns an access gets a dynamic shared-memory symbol -/// (`__dynamic_smem_{function_name}`). -/// -/// # Alignment -/// -/// The alignment is pre-computed during the lowering pre-pass. It is the -/// maximum of the owner's body requirements and every launch-contract marker -/// that can reach it. -/// -/// # Byte Offset -/// -/// - `DynamicSharedArray::get()` / `get_raw()`: offset = 0, returns base pointer -/// - `DynamicSharedArray::offset(N)`: offset = N bytes, returns base + N -/// -/// Called directly from `MirToLlvmConversionDriver::rewrite` (not through -/// op_cast dispatch) because it needs cross-function state maps. -pub fn convert_extern_shared_dc( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - _operands_info: &OperandsInfo, - shared_globals: &mut SharedGlobalsMap, - dynamic_smem_alignments: &mut DynamicSmemAlignmentMap, -) -> Result<()> { - let (byte_offset, alignment) = { - let extern_shared_op = dialect_mir::ops::MirExternSharedOp::new(op); - let byte_offset = extern_shared_op.get_byte_offset_value(ctx); - let alignment = extern_shared_op.get_alignment_value(ctx); - (byte_offset, alignment) - }; - - let func_name: String = { - let parent_block = op - .deref(ctx) - .get_parent_block() - .ok_or_else(|| anyhow_to_pliron(anyhow::anyhow!("Op has no parent block")))?; - let func_op_ptr = helpers::get_parent_func(ctx, parent_block).map_err(anyhow_to_pliron)?; - let llvm_func = Operation::get_op::(func_op_ptr, ctx) - .ok_or_else(|| anyhow_to_pliron(anyhow::anyhow!("Parent op is not an llvm.func")))?; - llvm_func.get_symbol_name(ctx).to_string() - }; - - let global_name = get_or_create_extern_shared_global( - ctx, - op, - &func_name, - shared_globals, - dynamic_smem_alignments, - alignment, - )?; - - let address_of_op = llvm::AddressOfOp::new(ctx, global_name, 3); - rewriter.insert_operation(ctx, address_of_op.get_operation()); - - let base_ptr = address_of_op.get_operation().deref(ctx).get_result(0); - - if byte_offset > 0 { - let i64_ty = IntegerType::get(ctx, 64, Signedness::Signless); - let offset_attr = pliron::builtin::attributes::IntegerAttr::new( - i64_ty, - pliron::utils::apint::APInt::from_u64( - byte_offset, - std::num::NonZeroUsize::new(64).unwrap(), - ), - ); - let offset_const = llvm::ConstantOp::new(ctx, offset_attr.into()); - rewriter.insert_operation(ctx, offset_const.get_operation()); - let offset_value = offset_const.get_operation().deref(ctx).get_result(0); - - let i8_ty = IntegerType::get(ctx, 8, Signedness::Signless); - let gep_op = llvm::GetElementPtrOp::new( - ctx, - base_ptr, - vec![llvm_export::ops::GepIndex::Value(offset_value)], - i8_ty.into(), - ); - rewriter.insert_operation(ctx, gep_op.get_operation()); - rewriter.replace_operation(ctx, op, gep_op.get_operation()); - } else { - rewriter.replace_operation(ctx, op, address_of_op.get_operation()); - } - - Ok(()) -} - -/// Get or create the extern shared memory global for an owning function. -/// -/// Creates an LLVM global variable with: -/// - Zero-length array type: `[0 x i8]` -/// - External linkage (no initializer) -/// - Address space 3 (shared memory) -/// - Pre-computed body and calling-kernel contract alignment -/// -/// Each owning function gets its own dynamic shared memory symbol. Uses -/// `shared_globals` for deduplication (only one global per function). -fn get_or_create_extern_shared_global( - ctx: &mut Context, - op: Ptr, - func_name: &str, - shared_globals: &mut SharedGlobalsMap, - dynamic_smem_alignments: &mut DynamicSmemAlignmentMap, - _requested_alignment: u64, -) -> Result { - let (symbol_name, max_alignment) = dynamic_smem_alignments.get(func_name).cloned().ok_or_else( - || { - anyhow_to_pliron(anyhow::anyhow!( - "Internal error: dynamic shared memory alignment not pre-computed for function '{}'. \ - This should have been done in compute_max_dynamic_smem_alignment.", - func_name - )) - }, - )?; - - let i8_ty = IntegerType::get(ctx, 8, Signedness::Signless); - let declaration = SharedGlobalDeclaration { - kind: SharedGlobalKind::DynamicExtern, - mir_elem_type: i8_ty.into(), - size: 0, - alignment: max_alignment, - // The dynamic pool is per-launch storage with no originating static; - // it never carries a source-level debug identity. - debug_info: None, - debug_owner_function: None, - }; - let global_created_key = format!("__dynamic_smem_global_created_{}", func_name); - if let Some(existing) = shared_globals.get(&global_created_key) { - if existing.declaration != declaration || existing.symbol != symbol_name { - return Err(anyhow_to_pliron(anyhow::anyhow!( - "dynamic shared-memory key {:?} has an incompatible declaration", - global_created_key - ))); - } - return Ok(existing.symbol.clone()); - } - - let array_type = ArrayType::get(ctx, i8_ty.into(), 0); - - let global_op = llvm::GlobalOp::new_with_alignment( - ctx, - symbol_name.clone(), - array_type.into(), - max_alignment, - ); - global_op.set_address_space(ctx, llvm_export::types::address_space::SHARED); - - { - use llvm_export::attributes::LinkageAttr; - global_op.set_attr_llvm_global_linkage(ctx, LinkageAttr::ExternalLinkage); - } - - let parent_block = op - .deref(ctx) - .get_parent_block() - .ok_or_else(|| anyhow_to_pliron(anyhow::anyhow!("Op has no parent block")))?; - let module_op = helpers::get_module_from_block(ctx, parent_block).map_err(anyhow_to_pliron)?; - let region = module_op.deref(ctx).get_region(0); - let module_block = region - .deref(ctx) - .iter(ctx) - .next() - .ok_or_else(|| anyhow_to_pliron(anyhow::anyhow!("Module is empty")))?; - - global_op.get_operation().insert_at_front(module_block, ctx); - - shared_globals.insert( - global_created_key, - SharedGlobalRecord { - symbol: symbol_name.clone(), - declaration, - }, - ); - - Ok(symbol_name) -} - -#[cfg(test)] -// Tests build kinded fixture types directly; production minting lives in mir-importer's facts.rs. -#[allow(clippy::disallowed_methods)] -mod tests { - //! End-to-end lowering tests for `dialect-mir` memory ops. - //! - //! The `convert_*` functions in this file take a live - //! `DialectConversionRewriter`, which is owned by pliron's conversion - //! driver and not constructible standalone. So each test builds a small - //! MIR module, runs the full `lower_mir_to_llvm` pass on it, and asserts - //! the lowered module contains the expected `dialect-llvm` shape — same - //! pattern as `tests/lowering_test.rs`. - - use super::*; - use crate::convert::ops::test_util::*; - use dialect_mir::attributes::MirPointerKindAuthorityAttr; - use dialect_mir::ops as mir; - use dialect_mir::types::{ - MirArrayType, MirPointerKind, MirPtrType, MirStructType, MirTupleType, MirUnionType, - }; - use llvm_export::op_interfaces::PointerTypeResult; - use llvm_export::ops as llvm; - use llvm_export::types::{PointerType, StructType, address_space as llvm_addr}; - use pliron::basic_block::BasicBlock; - use pliron::builtin::attributes::{StringAttr, TypeAttr}; - use pliron::builtin::op_interfaces::SymbolOpInterface; - use pliron::builtin::types::{IntegerType, Signedness}; - use pliron::context::Context; - use pliron::linked_list::ContainsLinkedList; - use pliron::location::{Location, Source}; - use pliron::op::Op; - use pliron::operation::Operation; - use std::path::PathBuf; - - fn ptr_addrspace(ctx: &Context, ty: TypeHandle) -> u32 { - ty.deref(ctx) - .downcast_ref::() - .expect("expected llvm.PointerType") - .address_space() - } - - fn src_location(ctx: &mut Context, file: &str, line: i32, column: i32) -> Location { - Location::SrcPos { - src: Source::new_from_file(ctx, PathBuf::from(file)), - pos: combine::stream::position::SourcePosition { line, column }, - } - } - - fn over_aligned_tuple_ty(ctx: &mut Context) -> TypeHandle { - let byte: TypeHandle = IntegerType::get(ctx, 8, Signedness::Unsigned).into(); - let marker: TypeHandle = MirStructType::get_with_full_layout( - ctx, - "Align32".into(), - vec![], - vec![], - vec![], - vec![], - 0, - 32, - ) - .into(); - MirTupleType::get_with_layout(ctx, vec![marker, byte], vec![0, 1], vec![0, 0], 32, 32) - .into() - } - - #[test] - fn convert_alloca_lowers_to_llvm_alloca() { - let mut ctx = make_ctx(); - let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); - let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty, true); - - let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); - - let alloca_op = Operation::new( - &mut ctx, - mir::MirAllocaOp::get_concrete_op_info(), - vec![mir_ptr_ty.into()], - vec![], - vec![], - 0, - ); - alloca_op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - assert_eq!( - count_ops::(&ctx, &body), - 1, - "expected exactly one llvm.alloca" - ); - let alloca = find_first::(&ctx, &body).unwrap(); - // Element type should round-trip through convert_type as i32. - let elem_ty = alloca.result_pointee_type(&ctx); - assert!(elem_ty.deref(&ctx).is::()); - } - - #[test] - fn convert_alloca_preserves_nested_array_element_alignment() { - let mut ctx = make_ctx(); - let tuple_ty = over_aligned_tuple_ty(&mut ctx); - let inner: TypeHandle = MirArrayType::get(&mut ctx, tuple_ty, 2).into(); - let outer: TypeHandle = MirArrayType::get(&mut ctx, inner, 3).into(); - let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, outer, true); - let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); - - let alloca_op = Operation::new( - &mut ctx, - mir::MirAllocaOp::get_concrete_op_info(), - vec![mir_ptr_ty.into()], - vec![], - vec![], - 0, - ); - alloca_op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - let alloca = find_first::(&ctx, &body).expect("expected llvm.alloca"); - assert_eq!( - llvm_export::ops::op_alignment(&ctx, alloca.get_operation()), - Some(32) - ); - } - - #[test] - fn convert_alloca_preserves_debug_local_metadata() { - let mut ctx = make_ctx(); - let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); - let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty, true); - - let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); - - let alloca_op = Operation::new( - &mut ctx, - mir::MirAllocaOp::get_concrete_op_info(), - vec![mir_ptr_ty.into()], - vec![], - vec![], - 0, - ); - llvm::set_debug_local_variable( - &mut ctx, - alloca_op, - llvm::DebugLocalVariableInfo { - name: "x".to_string(), - argument_index: Some(1), - ty: llvm::DebugLocalTypeKind::Basic { - name: "i32".to_string(), - size_bits: 32, - encoding: "DW_ATE_signed", - }, - }, - ); - alloca_op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - let alloca = find_first::(&ctx, &body).unwrap(); - let info = llvm::debug_local_variable(&ctx, alloca.get_operation()) - .expect("debug local metadata should survive lowering"); - - assert_eq!(info.name, "x"); - assert_eq!(info.argument_index, Some(1)); - assert_eq!( - info.ty, - llvm::DebugLocalTypeKind::Basic { - name: "i32".to_string(), - size_bits: 32, - encoding: "DW_ATE_signed", - } - ); - } - - #[test] - fn convert_store_lowers_to_llvm_store() { - let mut ctx = make_ctx(); - let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); - let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty, true); - - // Kernel takes (ptr, val) so we can store one into the other. - let (module_ptr, block) = build_kernel(&mut ctx, vec![mir_ptr_ty.into(), i32_ty], vec![]); - let ptr_val = block.deref(&ctx).get_argument(0); - let val = block.deref(&ctx).get_argument(1); - - let store_op = Operation::new( - &mut ctx, - mir::MirStoreOp::get_concrete_op_info(), - vec![], - vec![ptr_val, val], - vec![], - 0, - ); - store_op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - assert_eq!( - count_ops::(&ctx, &body), - 1, - "expected one llvm.store" - ); - // The original mir.store must be gone. - assert_eq!(count_ops::(&ctx, &body), 0); - - // convert_store swaps operand order: mir.store is [ptr, value] but - // llvm.store takes (value, ptr). Verify that mapping survived. - let store = find_first::(&ctx, &body).unwrap(); - let addr_ty = store.get_operand_address(&ctx).get_type(&ctx); - assert!(addr_ty.deref(&ctx).is::(), "operand 1 is ptr"); - } - - #[test] - fn convert_store_preserves_volatile() { - let mut ctx = make_ctx(); - let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); - let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty, true); - - let (module_ptr, block) = build_kernel(&mut ctx, vec![mir_ptr_ty.into(), i32_ty], vec![]); - let ptr_val = block.deref(&ctx).get_argument(0); - let val = block.deref(&ctx).get_argument(1); - - let store_op = Operation::new( - &mut ctx, - mir::MirStoreOp::get_concrete_op_info(), - vec![], - vec![ptr_val, val], - vec![], - 0, - ); - mir::MirStoreOp::new(store_op).set_volatile(&mut ctx, true); - store_op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - let store = find_first::(&ctx, &body).unwrap(); - assert!( - llvm_export::ops::op_volatile(&ctx, store.get_operation()), - "volatile mir.store must lower to a volatile llvm.store" - ); - } - - #[test] - fn convert_load_lowers_to_llvm_load() { - let mut ctx = make_ctx(); - let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); - let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty, false); - - let (module_ptr, block) = build_kernel(&mut ctx, vec![mir_ptr_ty.into()], vec![]); - let ptr_val = block.deref(&ctx).get_argument(0); - - let load_op = Operation::new( - &mut ctx, - mir::MirLoadOp::get_concrete_op_info(), - vec![i32_ty], - vec![ptr_val], - vec![], - 0, - ); - load_op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - assert_eq!(count_ops::(&ctx, &body), 1); - assert_eq!(count_ops::(&ctx, &body), 0); - } - - /// Lower `mir.load (mir.field_addr %p, field_index)` for a struct of - /// signless integer fields with the given layout and report the alignment - /// stamped on the resulting `llvm.load`. `None` means no stamp survived - /// and the exporter's natural-alignment default applies. - fn lowered_field_load_alignment( - field_bit_widths: Vec, - field_offsets: Vec, - total_size: u64, - abi_align: u64, - field_index: u32, - ) -> Option { - let mut ctx = make_ctx(); - let field_types: Vec = field_bit_widths - .iter() - .map(|w| IntegerType::get(&ctx, *w, Signedness::Signless).into()) - .collect(); - let field_names = (0..field_types.len()).map(|i| format!("f{i}")).collect(); - let struct_ty: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "FieldLoadAlign".into(), - field_names, - field_types.clone(), - vec![], - field_offsets, - total_size, - abi_align, - ) - .into(); - let struct_ptr_ty = MirPtrType::get_generic(&mut ctx, struct_ty, false); - let field_ty = field_types[field_index as usize]; - let field_ptr_ty = MirPtrType::get_generic(&mut ctx, field_ty, false); - - let (module_ptr, block) = build_kernel(&mut ctx, vec![struct_ptr_ty.into()], vec![]); - let struct_ptr_val = block.deref(&ctx).get_argument(0); - - let field_addr_op = - mir::MirFieldAddrOp::build(&mut ctx, struct_ptr_val, field_ptr_ty.into(), field_index) - .expect("field_addr build"); - field_addr_op.insert_at_back(block, &ctx); - let field_ptr_val = field_addr_op.deref(&ctx).get_result(0); - - let load_op = Operation::new( - &mut ctx, - mir::MirLoadOp::get_concrete_op_info(), - vec![field_ty], - vec![field_ptr_val], - vec![], - 0, - ); - load_op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - let load = find_first::(&ctx, &body).expect("expected one llvm.load"); - llvm_export::ops::op_alignment(&ctx, load.get_operation()) - } - - /// Field 0 of an over-aligned struct sits at the aggregate's own - /// alignment, which the field's scalar result type cannot state on its - /// own. This is what lets LoadStoreVectorizer fuse the adjacent pair. - #[test] - fn convert_load_inherits_overaligned_field_alignment_at_offset_zero() { - // #[repr(C, align(8))] struct { a: i32, b: i32 } - assert_eq!( - lowered_field_load_alignment(vec![32, 32], vec![0, 4], 8, 8, 0), - Some(8) - ); - } - - /// A field at a nonzero offset proves `gcd(abi_align, offset)`: an i32 at - /// offset 8 of an align-16 struct proves 8, beating its natural 4. - #[test] - fn convert_load_narrows_field_alignment_to_gcd_of_align_and_offset() { - // #[repr(C, align(16))] struct { a: i64, b: i32 } - assert_eq!( - lowered_field_load_alignment(vec![64, 32], vec![0, 8], 16, 16, 1), - Some(8) - ); - } - - /// Whole-value loads of pointer-free packed structs use the packed LLVM - /// representation, preserving rustc's byte size and field offsets. - #[test] - fn packed_struct_whole_value_load_uses_packed_layout() { - let mut ctx = make_ctx(); - let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); - let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); - let packed_ty: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "Packed".into(), - vec!["tag".into(), "value".into()], - vec![u8_ty, u32_ty], - vec![0, 1], - vec![0, 1], - 5, - 1, - ) - .into(); - let ptr_ty = MirPtrType::get_generic(&mut ctx, packed_ty, false); - - let (module_ptr, block) = build_kernel(&mut ctx, vec![ptr_ty.into()], vec![]); - let ptr_val = block.deref(&ctx).get_argument(0); - - let load_op = Operation::new( - &mut ctx, - mir::MirLoadOp::get_concrete_op_info(), - vec![packed_ty], - vec![ptr_val], - vec![], - 0, - ); - load_op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr) - .expect("whole-value load of a pointer-free packed struct must lower"); - - let body = kernel_blocks(&ctx, module_ptr); - let load = find_first::(&ctx, &body).expect("expected packed llvm.load"); - let result_ty = load - .get_operation() - .deref(&ctx) - .get_result(0) - .get_type(&ctx); - let result_ty_ref = result_ty.deref(&ctx); - let struct_ty = result_ty_ref - .downcast_ref::() - .expect("packed load result must be an LLVM struct"); - assert_eq!(struct_ty.layout(), StructLayout::Packed); - assert_eq!( - crate::convert::types::llvm_type_size_align(&ctx, result_ty), - Some((5, 1)) - ); - assert_eq!( - llvm_export::ops::op_alignment(&ctx, load.get_operation()), - Some(1) - ); - } - - #[test] - fn packed_struct_whole_value_load_with_shared_pointer_fails_closed() { - let mut ctx = make_ctx(); - let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); - let pointee: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); - let shared_ty: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, false).into(); - let packed_ty: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "PackedShared".into(), - vec!["tag".into(), "ptr".into()], - vec![u8_ty, shared_ty], - vec![0, 1], - vec![0, 1], - 9, - 1, - ) - .into(); - let ptr_ty = MirPtrType::get_generic(&mut ctx, packed_ty, false); - - let (module_ptr, block) = build_kernel(&mut ctx, vec![ptr_ty.into()], vec![]); - let ptr_val = block.deref(&ctx).get_argument(0); - let load_op = Operation::new( - &mut ctx, - mir::MirLoadOp::get_concrete_op_info(), - vec![packed_ty], - vec![ptr_val], - vec![], - 0, - ); - load_op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - let err = crate::lower_mir_to_llvm(&mut ctx, module_ptr) - .expect_err("packed whole-value load containing AS3 must remain fail-closed"); - assert!( - format!("{err:?}").contains("target-mode dependent"), - "the refusal must identify the target-dependent packed AS3 image: {err:?}" - ); - } - - /// Whole-value stores use the same packed representation as construction - /// and loads, while preserving the MIR aggregate's proved ABI alignment. - #[test] - fn packed_struct_whole_value_store_uses_packed_layout() { - let mut ctx = make_ctx(); - let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); - let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); - let packed_ty: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "Packed".into(), - vec!["tag".into(), "value".into()], - vec![u8_ty, u32_ty], - vec![0, 1], - vec![0, 1], - 5, - 1, - ) - .into(); - let ptr_ty = MirPtrType::get_generic(&mut ctx, packed_ty, false); - - let (module_ptr, block) = build_kernel(&mut ctx, vec![ptr_ty.into(), packed_ty], vec![]); - let ptr_val = block.deref(&ctx).get_argument(0); - let val = block.deref(&ctx).get_argument(1); - - let store_op = Operation::new( - &mut ctx, - mir::MirStoreOp::get_concrete_op_info(), - vec![], - vec![ptr_val, val], - vec![], - 0, - ); - store_op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr) - .expect("whole-value store of a pointer-free packed struct must lower"); - - let body = kernel_blocks(&ctx, module_ptr); - let store = find_first::(&ctx, &body).expect("expected packed llvm.store"); - let value_ty = store - .get_operation() - .deref(&ctx) - .get_operand(0) - .get_type(&ctx); - let value_ty_ref = value_ty.deref(&ctx); - let struct_ty = value_ty_ref - .downcast_ref::() - .expect("packed store value must be an LLVM struct"); - assert_eq!(struct_ty.layout(), StructLayout::Packed); - assert_eq!( - crate::convert::types::llvm_type_size_align(&ctx, value_ty), - Some((5, 1)) - ); - assert_eq!( - llvm_export::ops::op_alignment(&ctx, store.get_operation()), - Some(1) - ); - } - - /// A naturally aligned inner struct sitting at a packed byte offset: the - /// field address only proves align 1, and the load must claim that over - /// the inner type's recorded abi alignment. Claiming the abi would stamp - /// `align 4` on a 1-aligned address, which llc may honor with a wider - /// access than the bytes allow. - #[test] - fn convert_load_claims_address_alignment_over_abi_at_packed_offsets() { - let mut ctx = make_ctx(); - let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); - let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); - let inner_ty: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "Inner".into(), - vec!["v".into()], - vec![u32_ty], - vec![0], - vec![0], - 4, - 4, - ) - .into(); - let outer_ty: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "PackedOuter".into(), - vec!["tag".into(), "inner".into()], - vec![u8_ty, inner_ty], - vec![0, 1], - vec![0, 1], - 5, - 1, - ) - .into(); - let outer_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, outer_ty, false).into(); - let inner_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, inner_ty, false).into(); - - let (module_ptr, block) = build_kernel(&mut ctx, vec![outer_ptr_ty], vec![]); - let base = block.deref(&ctx).get_argument(0); - - let field_addr_op = - mir::MirFieldAddrOp::build(&mut ctx, base, inner_ptr_ty, 1).expect("field_addr build"); - field_addr_op.insert_at_back(block, &ctx); - let field_ptr_val = field_addr_op.deref(&ctx).get_result(0); - - let load_op = Operation::new( - &mut ctx, - mir::MirLoadOp::get_concrete_op_info(), - vec![inner_ty], - vec![field_ptr_val], - vec![], - 0, - ); - load_op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - let load = find_first::(&ctx, &body).expect("expected one llvm.load"); - assert_eq!( - llvm_export::ops::op_alignment(&ctx, load.get_operation()), - Some(1), - "the load must claim the address's proved alignment, not the inner abi" - ); - } - - /// A struct with no extra alignment proves nothing beyond the scalar's - /// natural alignment: the stamp equals the exporter's default 4, so the - /// emitted access is unchanged. - #[test] - fn convert_load_keeps_natural_alignment_without_overalignment() { - // struct { a: i32, b: i32 } with rustc's natural abi_align 4 - assert_eq!( - lowered_field_load_alignment(vec![32, 32], vec![0, 4], 8, 4, 1), - Some(4) - ); - } - - /// dialect-mir only verifier-enforces power-of-two alignment for unions - /// and enums. A malformed hand-built struct layout must decline the stamp - /// rather than emit a non-power-of-two `align N` that llc rejects. - #[test] - fn convert_load_declines_non_power_of_two_field_alignment() { - assert_eq!( - lowered_field_load_alignment(vec![32], vec![0], 12, 12, 0), - None - ); - } - - /// Lower `mir.store %v, (mir.field_addr %p, field_index)` for a struct of - /// signless integer fields with the given layout and report the alignment - /// stamped on the resulting `llvm.store`. `None` means no stamp survived - /// and the exporter's natural-alignment default applies. - fn lowered_field_store_alignment( - field_bit_widths: Vec, - field_offsets: Vec, - total_size: u64, - abi_align: u64, - field_index: u32, - ) -> Option { - let mut ctx = make_ctx(); - let field_types: Vec = field_bit_widths - .iter() - .map(|w| IntegerType::get(&ctx, *w, Signedness::Signless).into()) - .collect(); - let field_names = (0..field_types.len()).map(|i| format!("f{i}")).collect(); - let struct_ty: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "FieldStoreAlign".into(), - field_names, - field_types.clone(), - vec![], - field_offsets, - total_size, - abi_align, - ) - .into(); - let struct_ptr_ty = MirPtrType::get_generic(&mut ctx, struct_ty, true); - let field_ty = field_types[field_index as usize]; - let field_ptr_ty = MirPtrType::get_generic(&mut ctx, field_ty, true); - - // The stored value arrives as a kernel argument of the field's own - // scalar type, so `value_abi_align` reports nothing about it and the - // address's stamp is the only alignment left -- the case this covers. - let (module_ptr, block) = - build_kernel(&mut ctx, vec![struct_ptr_ty.into(), field_ty], vec![]); - let struct_ptr_val = block.deref(&ctx).get_argument(0); - let val = block.deref(&ctx).get_argument(1); - - let field_addr_op = - mir::MirFieldAddrOp::build(&mut ctx, struct_ptr_val, field_ptr_ty.into(), field_index) - .expect("field_addr build"); - field_addr_op.insert_at_back(block, &ctx); - let field_ptr_val = field_addr_op.deref(&ctx).get_result(0); - - let store_op = Operation::new( - &mut ctx, - mir::MirStoreOp::get_concrete_op_info(), - vec![], - vec![field_ptr_val, val], - vec![], - 0, - ); - store_op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - let store = find_first::(&ctx, &body).expect("expected one llvm.store"); - llvm_export::ops::op_alignment(&ctx, store.get_operation()) - } - - /// Field 0 of an over-aligned struct sits at the aggregate's own alignment, - /// which the stored scalar's type cannot state. This is what lets - /// LoadStoreVectorizer fuse the adjacent pair into one wide store. - #[test] - fn convert_store_inherits_overaligned_field_alignment_at_offset_zero() { - // #[repr(C, align(8))] struct { a: i32, b: i32 } - assert_eq!( - lowered_field_store_alignment(vec![32, 32], vec![0, 4], 8, 8, 0), - Some(8) - ); - } - - /// A field at a nonzero offset proves `gcd(abi_align, offset)`: an i32 at - /// offset 8 of an align-16 struct proves 8, beating its natural 4. - #[test] - fn convert_store_narrows_field_alignment_to_gcd_of_align_and_offset() { - // #[repr(C, align(16))] struct { a: i64, b: i32 } - assert_eq!( - lowered_field_store_alignment(vec![64, 32], vec![0, 8], 16, 16, 1), - Some(8) - ); - } - - /// A struct with no extra alignment proves nothing beyond the scalar's - /// natural alignment, so the emitted store is unchanged. Widening here - /// would claim an alignment the source never guaranteed. - #[test] - fn convert_store_keeps_natural_alignment_without_overalignment() { - // struct { a: i32, b: i32 } with rustc's natural abi_align 4 - assert_eq!( - lowered_field_store_alignment(vec![32, 32], vec![0, 4], 8, 4, 1), - Some(4) - ); - } - - /// A malformed hand-built layout must decline the stamp rather than emit a - /// non-power-of-two `align N` that llc rejects. Same guard the load path - /// has, and it matters more here: an over-aligned store instruction on an - /// under-aligned address is undefined, not merely slow. - #[test] - fn convert_store_declines_non_power_of_two_field_alignment() { - assert_eq!( - lowered_field_store_alignment(vec![32], vec![0], 12, 12, 0), - None - ); - } - - /// Lower `load (&arr[index])` where `arr` is the array field of an - /// over-aligned struct, and report the alignment the load ends up with. - /// - /// `index` of `Some(i)` builds a constant index, `None` a runtime one. - fn lowered_element_load_alignment( - element_bits: u32, - element_count: u64, - struct_abi_align: u64, - index: Option, - ) -> Option { - use pliron::builtin::attributes::IntegerAttr; - use std::num::NonZeroUsize; - - let mut ctx = make_ctx(); - let element_ty: TypeHandle = - IntegerType::get(&ctx, element_bits, Signedness::Signless).into(); - let array_ty: TypeHandle = - dialect_mir::types::MirArrayType::get(&mut ctx, element_ty, element_count).into(); - let elem_bytes = u64::from(element_bits) / 8; - let struct_ty: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "ElementLoadAlign".into(), - vec!["lanes".into()], - vec![array_ty], - vec![], - vec![0], - elem_bytes * element_count, - struct_abi_align, - ) - .into(); - let struct_ptr_ty = MirPtrType::get_generic(&mut ctx, struct_ty, false); - let array_ptr_ty = MirPtrType::get_generic(&mut ctx, array_ty, false); - let element_ptr_ty = MirPtrType::get_generic(&mut ctx, element_ty, false); - let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signed).into(); - - let (module_ptr, block) = - build_kernel(&mut ctx, vec![struct_ptr_ty.into(), i64_ty], vec![]); - let struct_ptr_val = block.deref(&ctx).get_argument(0); - - // &s.lanes -- carries the struct's alignment onto the array address. - let field_addr_op = - mir::MirFieldAddrOp::build(&mut ctx, struct_ptr_val, array_ptr_ty.into(), 0) - .expect("field_addr build"); - field_addr_op.insert_at_back(block, &ctx); - let array_ptr_val = field_addr_op.deref(&ctx).get_result(0); - - let index_val = match index { - Some(i) => { - let constant = Operation::new( - &mut ctx, - mir::MirConstantOp::get_concrete_op_info(), - vec![i64_ty], - vec![], - vec![], - 0, - ); - mir::MirConstantOp::new(constant).set_attr_value( - &ctx, - IntegerAttr::new( - IntegerType::get(&ctx, 64, Signedness::Signed), - APInt::from_u64(i, NonZeroUsize::new(64).unwrap()), - ), - ); - constant.insert_at_back(block, &ctx); - constant.deref(&ctx).get_result(0) - } - None => block.deref(&ctx).get_argument(1), - }; - - let elem_addr_op = Operation::new( - &mut ctx, - mir::MirArrayElementAddrOp::get_concrete_op_info(), - vec![element_ptr_ty.into()], - vec![array_ptr_val, index_val], - vec![], - 0, - ); - elem_addr_op.insert_at_back(block, &ctx); - let elem_ptr_val = elem_addr_op.deref(&ctx).get_result(0); - - let load_op = Operation::new( - &mut ctx, - mir::MirLoadOp::get_concrete_op_info(), - vec![element_ty], - vec![elem_ptr_val], - vec![], - 0, - ); - load_op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - let load = find_first::(&ctx, &body).expect("expected one llvm.load"); - llvm_export::ops::op_alignment(&ctx, load.get_operation()) - } - - /// Element 0 inherits the whole alignment the base address proved, which is - /// what lets the adjacent pair fuse into one wide load. - #[test] - fn convert_load_inherits_base_alignment_at_element_zero() { - // &(#[repr(C, align(8))] struct { lanes: [i32; 2] }).lanes[0] - assert_eq!(lowered_element_load_alignment(32, 2, 8, Some(0)), Some(8)); - } - - /// A nonzero constant index proves `gcd(base, i * stride)`: element 1 of an - /// align-8 `[i32; 2]` sits at byte 4, so it proves 4, not 8. - #[test] - fn convert_load_narrows_element_alignment_to_gcd_with_offset() { - assert_eq!(lowered_element_load_alignment(32, 2, 8, Some(1)), Some(4)); - } - - /// A runtime index can land on any element, so only what every stride - /// preserves may be claimed -- `gcd(base, stride)`, never the base itself. - #[test] - fn convert_load_claims_only_stride_alignment_for_a_runtime_index() { - assert_eq!(lowered_element_load_alignment(32, 2, 8, None), Some(4)); - } - - /// A base with no extra alignment proves nothing beyond the element's own - /// natural alignment, so the emitted access is unchanged. - #[test] - fn convert_load_keeps_natural_element_alignment_without_overalignment() { - assert_eq!(lowered_element_load_alignment(32, 2, 4, Some(0)), Some(4)); - } - - /// Like [`lowered_element_load_alignment`], but for an array whose element - /// is an aggregate built by `element_ty_of` (which also reports the - /// element's stored size in bytes). The claim on the element address must - /// then come from the element's *exact* stride — rustc's stored size, - /// padding included — not from any LLVM-level approximation. - /// - /// With `load_first_scalar` the element must itself be an array and the - /// access becomes `s.lanes[index][0]`, mirroring the nested-read chain - /// where the outer stamp is inherited by the inner index-0 address and - /// ends up on a scalar load that LoadStoreVectorizer trusts. Without it, - /// the element itself is loaded. - /// - /// Reports `(element address stamp, final load alignment)`. - fn lowered_aggregate_element_alignments( - element_ty_of: impl FnOnce(&mut Context) -> (TypeHandle, u64), - element_count: u64, - struct_abi_align: u64, - index: Option, - load_first_scalar: bool, - ) -> (Option, Option) { - use pliron::builtin::attributes::IntegerAttr; - use std::num::NonZeroUsize; - - let mut ctx = make_ctx(); - let (element_ty, elem_bytes) = element_ty_of(&mut ctx); - let inner_scalar_ty = load_first_scalar.then(|| { - let element_ref = element_ty.deref(&ctx); - element_ref - .downcast_ref::() - .expect("load_first_scalar needs an array element") - .element_type() - }); - let array_ty: TypeHandle = MirArrayType::get(&mut ctx, element_ty, element_count).into(); - let struct_ty: TypeHandle = MirStructType::get_with_full_layout( - &mut ctx, - "AggregateElementAlign".into(), - vec!["lanes".into()], - vec![array_ty], - vec![], - vec![0], - elem_bytes * element_count, - struct_abi_align, - ) - .into(); - let struct_ptr_ty = MirPtrType::get_generic(&mut ctx, struct_ty, false); - let array_ptr_ty = MirPtrType::get_generic(&mut ctx, array_ty, false); - let element_ptr_ty = MirPtrType::get_generic(&mut ctx, element_ty, false); - let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signed).into(); - - let (module_ptr, block) = - build_kernel(&mut ctx, vec![struct_ptr_ty.into(), i64_ty], vec![]); - let struct_ptr_val = block.deref(&ctx).get_argument(0); - - // &s.lanes -- carries the struct's alignment onto the array address. - let field_addr_op = - mir::MirFieldAddrOp::build(&mut ctx, struct_ptr_val, array_ptr_ty.into(), 0) - .expect("field_addr build"); - field_addr_op.insert_at_back(block, &ctx); - let array_ptr_val = field_addr_op.deref(&ctx).get_result(0); - - let constant_index = |ctx: &mut Context, i: u64| { - let constant = Operation::new( - ctx, - mir::MirConstantOp::get_concrete_op_info(), - vec![i64_ty], - vec![], - vec![], - 0, - ); - mir::MirConstantOp::new(constant).set_attr_value( - ctx, - IntegerAttr::new( - IntegerType::get(ctx, 64, Signedness::Signed), - APInt::from_u64(i, NonZeroUsize::new(64).unwrap()), - ), - ); - constant.insert_at_back(block, ctx); - constant.deref(ctx).get_result(0) - }; - - let index_val = match index { - Some(i) => constant_index(&mut ctx, i), - None => block.deref(&ctx).get_argument(1), - }; - - let elem_addr_op = Operation::new( - &mut ctx, - mir::MirArrayElementAddrOp::get_concrete_op_info(), - vec![element_ptr_ty.into()], - vec![array_ptr_val, index_val], - vec![], - 0, - ); - elem_addr_op.insert_at_back(block, &ctx); - let elem_ptr_val = elem_addr_op.deref(&ctx).get_result(0); - - let (loaded_ty, loaded_ptr_val) = match inner_scalar_ty { - Some(scalar_ty) => { - let zero_val = constant_index(&mut ctx, 0); - let scalar_ptr_ty = MirPtrType::get_generic(&mut ctx, scalar_ty, false); - let inner_addr_op = Operation::new( - &mut ctx, - mir::MirArrayElementAddrOp::get_concrete_op_info(), - vec![scalar_ptr_ty.into()], - vec![elem_ptr_val, zero_val], - vec![], - 0, - ); - inner_addr_op.insert_at_back(block, &ctx); - (scalar_ty, inner_addr_op.deref(&ctx).get_result(0)) - } - None => (element_ty, elem_ptr_val), - }; - - let load_op = Operation::new( - &mut ctx, - mir::MirLoadOp::get_concrete_op_info(), - vec![loaded_ty], - vec![loaded_ptr_val], - vec![], - 0, - ); - load_op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - // GEP order follows source order: field address, then the element - // address under test (then the inner index-0 address when nested). - let geps = find_all::(&ctx, &body); - assert_eq!(geps.len(), if load_first_scalar { 3 } else { 2 }); - let element_gep_align = llvm_export::ops::address_alignment(&ctx, geps[1].get_operation()); - let load = find_first::(&ctx, &body).expect("expected one llvm.load"); - let load_align = llvm_export::ops::op_alignment(&ctx, load.get_operation()); - (element_gep_align, load_align) - } - - /// `[f32; 3]` element under an align-8 base: stride is 12, so a runtime - /// index proves `gcd(8, 12) = 4` — and the inner index-0 scalar read - /// inherits exactly that. Guards against sizing the element through an - /// LLVM-level approximation, whose guessed stride of 8 would stamp - /// align 8 onto addresses that are only 4-aligned (a miscompile once - /// LoadStoreVectorizer trusts it). - #[test] - fn convert_load_claims_exact_aggregate_stride_for_nested_array_elements() { - use pliron::builtin::types::FP32Type; - // &(#[repr(C, align(8))] struct { lanes: [[f32; 3]; 4] }).lanes[i][0] - let nested_f32x3 = |ctx: &mut Context| { - let f32_ty: TypeHandle = FP32Type::get(ctx).into(); - (MirArrayType::get(ctx, f32_ty, 3).into(), 12) - }; - assert_eq!( - lowered_aggregate_element_alignments(nested_f32x3, 4, 8, None, true), - (Some(4), Some(4)) - ); - } - - /// A constant index into the same nested array uses the exact byte - /// offset: element 1 sits at byte 12 (`gcd(8, 12) = 4`), element 2 at - /// byte 24 (`gcd(8, 24) = 8`, the full base alignment again). - #[test] - fn convert_load_narrows_nested_element_alignment_by_exact_byte_offset() { - use pliron::builtin::types::FP32Type; - let nested_f32x3 = |ctx: &mut Context| { - let f32_ty: TypeHandle = FP32Type::get(ctx).into(); - (MirArrayType::get(ctx, f32_ty, 3).into(), 12) - }; - assert_eq!( - lowered_aggregate_element_alignments(nested_f32x3, 4, 8, Some(1), true), - (Some(4), Some(4)) - ); - assert_eq!( - lowered_aggregate_element_alignments(nested_f32x3, 4, 8, Some(2), true), - (Some(8), Some(8)) - ); - } - - /// A tuple element's stride comes from rustc's recorded `total_size` - /// (trailing padding included): `(f32, f32, f32)` stores 12 bytes, so an - /// align-8 base proves only 4 on a runtime element address. - #[test] - fn convert_array_element_addr_takes_tuple_stride_from_recorded_layout() { - use pliron::builtin::types::FP32Type; - let f32x3_tuple = |ctx: &mut Context| { - let f32_ty: TypeHandle = FP32Type::get(ctx).into(); - let tuple_ty: TypeHandle = - MirTupleType::get_with_layout(ctx, vec![f32_ty; 3], vec![], vec![0, 4, 8], 12, 4) - .into(); - (tuple_ty, 12) - }; - let (element_gep_align, _load_align) = - lowered_aggregate_element_alignments(f32x3_tuple, 4, 8, None, false); - assert_eq!(element_gep_align, Some(4)); - } - - /// `f16` arrives from the importer as `MirFP16Type`, not the converted - /// LLVM `half`, and its stride is exactly 2: an align-8 base proves 2 on - /// a runtime element address and `gcd(8, 4) = 4` at element 2. Guards - /// the arm the importer actually exercises — the old sizing guessed 8 - /// for this type too, the same over-claim as the aggregate cases. - #[test] - fn convert_load_claims_exact_f16_element_stride() { - let f16_scalar = |ctx: &mut Context| { - let f16_ty: TypeHandle = dialect_mir::types::MirFP16Type::get(ctx).into(); - (f16_ty, 2) - }; - assert_eq!( - lowered_aggregate_element_alignments(f16_scalar, 4, 8, None, false), - (Some(2), Some(2)) - ); - assert_eq!( - lowered_aggregate_element_alignments(f16_scalar, 4, 8, Some(2), false), - (Some(4), Some(4)) - ); - } - - /// An element whose stored size is unknown (a struct built without rustc - /// layout) must not have its stride guessed: the element address claims - /// nothing and the load keeps the previous, weaker-but-sound behaviour. - #[test] - fn convert_array_element_addr_declines_unknown_element_stride() { - use pliron::builtin::types::FP32Type; - let opaque_struct = |ctx: &mut Context| { - let f32_ty: TypeHandle = FP32Type::get(ctx).into(); - let struct_ty: TypeHandle = - MirStructType::get(ctx, "OpaqueElement".into(), vec!["x".into()], vec![f32_ty]) - .into(); - (struct_ty, 4) - }; - assert_eq!( - lowered_aggregate_element_alignments(opaque_struct, 4, 8, None, false), - (None, None) - ); - } - - #[test] - fn convert_dbg_value_lowers_to_llvm_dbg_value() { - let mut ctx = make_ctx(); - let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); - - let (module_ptr, block) = build_kernel(&mut ctx, vec![i32_ty], vec![]); - let value = block.deref(&ctx).get_argument(0); - - let dbg_op = mir::MirDbgValueOp::new(&mut ctx, value); - let dbg_loc = pliron::location::Location::Named { - name: "current value location".to_string(), - child_loc: Box::new(pliron::location::Location::Unknown), - }; - dbg_op - .get_operation() - .deref_mut(&ctx) - .set_loc(dbg_loc.clone()); - llvm::set_debug_local_variable( - &mut ctx, - dbg_op.get_operation(), - llvm::DebugLocalVariableInfo { - name: "x".to_string(), - argument_index: None, - ty: llvm::DebugLocalTypeKind::Basic { - name: "i32".to_string(), - size_bits: 32, - encoding: "DW_ATE_signed", - }, - }, - ); - llvm::set_debug_local_source_scope(&mut ctx, dbg_op.get_operation(), 42); - llvm::set_debug_fragment_variables( - &mut ctx, - dbg_op.get_operation(), - &[llvm::DebugFragmentVariableInfo { - variable: llvm::DebugLocalVariableInfo { - name: "pair".to_string(), - argument_index: None, - ty: llvm::DebugLocalTypeKind::Array { - name: "[u32; 2]".to_string(), - size_bits: 64, - element: Box::new(llvm::DebugLocalTypeKind::Basic { - name: "u32".to_string(), - size_bits: 32, - encoding: "DW_ATE_unsigned", - }), - count: 2, - }, - }, - fragment: llvm::DebugFragment { - offset_bits: 32, - size_bits: 32, - }, - source_scope: Some(42), - declaration: Some(llvm::DebugSourcePosition { - file: PathBuf::from("decl.rs"), - line: 7, - column: 3, - }), - }], - ); - llvm::set_debug_local_declaration_location( - &mut ctx, - dbg_op.get_operation(), - PathBuf::from("decl.rs"), - 7, - 3, - ); - dbg_op.get_operation().insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - assert_eq!(count_ops::(&ctx, &body), 0); - let dbg_value = find_first::(&ctx, &body) - .expect("expected lowered llvm.dbg_value marker"); - assert_eq!( - dbg_value.get_operation().deref(&ctx).loc(), - dbg_loc, - "dbg.value lowering should keep the current-value source location" - ); - let info = llvm::debug_local_variable(&ctx, dbg_value.get_operation()) - .expect("debug local metadata should survive dbg_value lowering"); - - assert_eq!(info.name, "x"); - assert_eq!( - llvm::debug_local_source_scope(&ctx, dbg_value.get_operation()), - Some(42), - "dbg.value lowering should keep the MIR source-scope owner" - ); - let (decl_file, decl_pos) = - llvm::debug_local_declaration_location(&ctx, dbg_value.get_operation()) - .expect("declaration location should survive dbg_value lowering"); - assert_eq!(decl_file, PathBuf::from("decl.rs")); - assert_eq!(decl_pos.line, 7); - assert_eq!(decl_pos.column, 3); - assert_eq!( - info.ty, - llvm::DebugLocalTypeKind::Basic { - name: "i32".to_string(), - size_bits: 32, - encoding: "DW_ATE_signed", - } - ); - let fragments = llvm::debug_fragment_variables(&ctx, dbg_value.get_operation()); - assert_eq!(fragments.len(), 1); - assert_eq!(fragments[0].fragment.offset_bits, 32); - assert_eq!(fragments[0].fragment.size_bits, 32); - assert_eq!(fragments[0].variable.name, "pair"); - } - - #[test] - fn convert_dbg_value_list_preserves_operands_and_expression() { - let mut ctx = make_ctx(); - let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); - let ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty, false); - let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signless).into(); - - let (module_ptr, block) = build_kernel(&mut ctx, vec![ptr_ty.into(), i64_ty], vec![]); - let base = block.deref(&ctx).get_argument(0); - let index = block.deref(&ctx).get_argument(1); - - let dbg_op = mir::MirDbgValueListOp::new(&mut ctx, vec![base, index]); - llvm::set_debug_local_variable( - &mut ctx, - dbg_op.get_operation(), - llvm::DebugLocalVariableInfo { - name: "item".to_string(), - argument_index: None, - ty: llvm::DebugLocalTypeKind::Basic { - name: "u32".to_string(), - size_bits: 32, - encoding: "DW_ATE_unsigned", - }, - }, - ); - let expression = llvm::DebugValueExpression::new(vec![ - llvm::DebugValueExpressionOp::Arg(0), - llvm::DebugValueExpressionOp::Arg(1), - llvm::DebugValueExpressionOp::ConstU(4), - llvm::DebugValueExpressionOp::Mul, - llvm::DebugValueExpressionOp::Plus, - llvm::DebugValueExpressionOp::Deref, - ]); - llvm::set_debug_value_expression(&mut ctx, dbg_op.get_operation(), &expression); - dbg_op.get_operation().insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - assert_eq!(count_ops::(&ctx, &body), 0); - let dbg_value = find_first::(&ctx, &body) - .expect("expected lowered llvm.dbg_value_list marker"); - assert_eq!(dbg_value.values(&ctx), vec![base, index]); - assert_eq!( - llvm::debug_value_expression(&ctx, dbg_value.get_operation()), - Some(expression) - ); - } - - #[test] - fn convert_alloca_preserves_local_memory_provenance() { - let mut ctx = make_ctx(); - let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); - let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty, true); - let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); - - let alloca_op = Operation::new( - &mut ctx, - mir::MirAllocaOp::get_concrete_op_info(), - vec![mir_ptr_ty.into()], - vec![], - vec![], - 0, - ); - let provenance = llvm_export::ops::LocalMemoryProvenanceAttr { - local_index: 3, - size_bytes: 16, - binding_name: "scratch".into(), - type_name: "[u32; 4]".into(), - }; - llvm_export::ops::set_local_memory_provenance(&mut ctx, alloca_op, provenance.clone()); - alloca_op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - let alloca = find_first::(&ctx, &body).unwrap(); - let copied = - llvm_export::ops::local_memory_provenance(&ctx, alloca.get_operation()).unwrap(); - assert_eq!(copied, provenance); - } - - #[test] - fn mem2reg_salvages_tagged_alloca_into_mir_dbg_value() { - let mut ctx = make_ctx(); - let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); - let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty, true); - - let (module_ptr, block) = build_kernel(&mut ctx, vec![i32_ty], vec![i32_ty]); - let arg = block.deref(&ctx).get_argument(0); - - let alloca_op = Operation::new( - &mut ctx, - mir::MirAllocaOp::get_concrete_op_info(), - vec![mir_ptr_ty.into()], - vec![], - vec![], - 0, - ); - let decl_loc = src_location(&mut ctx, "kernel.rs", 12, 9); - alloca_op.deref_mut(&ctx).set_loc(decl_loc.clone()); - llvm::set_debug_local_variable( - &mut ctx, - alloca_op, - llvm::DebugLocalVariableInfo { - name: "x".to_string(), - argument_index: Some(1), - ty: llvm::DebugLocalTypeKind::Basic { - name: "i32".to_string(), - size_bits: 32, - encoding: "DW_ATE_signed", - }, - }, - ); - llvm::set_debug_local_source_scope(&mut ctx, alloca_op, 9); - alloca_op.insert_at_back(block, &ctx); - let slot = alloca_op.deref(&ctx).get_result(0); - - let store_op = Operation::new( - &mut ctx, - mir::MirStoreOp::get_concrete_op_info(), - vec![], - vec![slot, arg], - vec![], - 0, - ); - store_op.insert_at_back(block, &ctx); - - let load_op = Operation::new( - &mut ctx, - mir::MirLoadOp::get_concrete_op_info(), - vec![i32_ty], - vec![slot], - vec![], - 0, - ); - load_op.insert_at_back(block, &ctx); - let loaded = load_op.deref(&ctx).get_result(0); - append_mir_return(&mut ctx, block, vec![loaded]); - - let mut analyses = pliron::pass::AnalysisManager::default(); - pliron::opts::mem2reg::mem2reg(module_ptr, &mut ctx, &mut analyses) - .expect("mem2reg should promote the local slot"); - - let blocks = vec![block]; - assert_eq!(count_ops::(&ctx, &blocks), 0); - assert_eq!(count_ops::(&ctx, &blocks), 0); - assert_eq!(count_ops::(&ctx, &blocks), 0); - - let dbg_values = find_all::(&ctx, &blocks); - assert!( - !dbg_values.is_empty(), - "mem2reg should leave value-based debug records for promoted locals" - ); - let info = llvm::debug_local_variable(&ctx, dbg_values[0].get_operation()) - .expect("mir.dbg_value should carry the promoted local metadata"); - assert_eq!(info.name, "x"); - assert_eq!(info.argument_index, Some(1)); - assert_eq!( - llvm::debug_local_source_scope(&ctx, dbg_values[0].get_operation()), - Some(9), - "mem2reg salvage should keep the local's MIR source-scope owner" - ); - assert_eq!( - dbg_values[0].get_operation().deref(&ctx).loc(), - decl_loc, - "debug records for source-less promoted ops should fall back to the local declaration" - ); - } - - #[test] - fn mem2reg_salvages_fragment_only_alloca_into_mir_dbg_value() { - let mut ctx = make_ctx(); - let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); - let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty, true); - - let (module_ptr, block) = build_kernel(&mut ctx, vec![i32_ty], vec![i32_ty]); - let arg = block.deref(&ctx).get_argument(0); - - let alloca_op = Operation::new( - &mut ctx, - mir::MirAllocaOp::get_concrete_op_info(), - vec![mir_ptr_ty.into()], - vec![], - vec![], - 0, - ); - let alloca_loc = src_location(&mut ctx, "kernel.rs", 20, 9); - alloca_op.deref_mut(&ctx).set_loc(alloca_loc); - llvm::set_debug_fragment_variables( - &mut ctx, - alloca_op, - &[llvm::DebugFragmentVariableInfo { - variable: llvm::DebugLocalVariableInfo { - name: "pair".to_string(), - argument_index: None, - ty: llvm::DebugLocalTypeKind::Array { - name: "[u32; 2]".to_string(), - size_bits: 64, - element: Box::new(llvm::DebugLocalTypeKind::Basic { - name: "u32".to_string(), - size_bits: 32, - encoding: "DW_ATE_unsigned", - }), - count: 2, - }, - }, - fragment: llvm::DebugFragment { - offset_bits: 32, - size_bits: 32, - }, - source_scope: Some(9), - declaration: Some(llvm::DebugSourcePosition { - file: PathBuf::from("kernel.rs"), - line: 20, - column: 9, - }), - }], - ); - alloca_op.insert_at_back(block, &ctx); - let slot = alloca_op.deref(&ctx).get_result(0); - - let store_op = Operation::new( - &mut ctx, - mir::MirStoreOp::get_concrete_op_info(), - vec![], - vec![slot, arg], - vec![], - 0, - ); - store_op.insert_at_back(block, &ctx); - - let load_op = Operation::new( - &mut ctx, - mir::MirLoadOp::get_concrete_op_info(), - vec![i32_ty], - vec![slot], - vec![], - 0, - ); - load_op.insert_at_back(block, &ctx); - let loaded = load_op.deref(&ctx).get_result(0); - append_mir_return(&mut ctx, block, vec![loaded]); - - let mut analyses = pliron::pass::AnalysisManager::default(); - pliron::opts::mem2reg::mem2reg(module_ptr, &mut ctx, &mut analyses) - .expect("mem2reg should promote fragment storage"); - - let dbg_values = find_all::(&ctx, &[block]); - assert!( - !dbg_values.is_empty(), - "fragment-only storage should still produce mir.dbg_value salvage" - ); - let fragments = llvm::debug_fragment_variables(&ctx, dbg_values[0].get_operation()); - assert_eq!(fragments.len(), 1); - assert_eq!(fragments[0].variable.name, "pair"); - assert_eq!(fragments[0].fragment.offset_bits, 32); - assert_eq!(fragments[0].fragment.size_bits, 32); - } - - #[test] - fn convert_load_preserves_volatile() { - let mut ctx = make_ctx(); - let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); - let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty, false); - - let (module_ptr, block) = build_kernel(&mut ctx, vec![mir_ptr_ty.into()], vec![]); - let ptr_val = block.deref(&ctx).get_argument(0); - - let load_op = Operation::new( - &mut ctx, - mir::MirLoadOp::get_concrete_op_info(), - vec![i32_ty], - vec![ptr_val], - vec![], - 0, - ); - mir::MirLoadOp::new(load_op).set_volatile(&mut ctx, true); - load_op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - let load = find_first::(&ctx, &body).unwrap(); - assert!( - llvm_export::ops::op_volatile(&ctx, load.get_operation()), - "volatile mir.load must lower to a volatile llvm.load" - ); - } - - #[test] - fn convert_ref_lowers_to_alloca_then_store() { - let mut ctx = make_ctx(); - let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); - let mir_ptr_ty = - MirPtrType::get_generic_with_kind(&mut ctx, i32_ty, false, MirPointerKind::SharedRef); - - // Take a u32 by value, build `&x`. - let (module_ptr, block) = build_kernel(&mut ctx, vec![i32_ty], vec![]); - let arg = block.deref(&ctx).get_argument(0); - - let ref_op_ptr = Operation::new( - &mut ctx, - mir::MirRefOp::get_concrete_op_info(), - vec![mir_ptr_ty.into()], - vec![arg], - vec![], - 0, - ); - let ref_op = mir::MirRefOp::new(ref_op_ptr); - ref_op.set_mutable(&mut ctx, false); - ref_op.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); - ref_op_ptr.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - assert_eq!( - count_ops::(&ctx, &body), - 1, - "ref must materialize via alloca" - ); - assert_eq!( - count_ops::(&ctx, &body), - 1, - "ref must store the value into the alloca" - ); - assert_eq!(count_ops::(&ctx, &body), 0); - } - - #[test] - fn convert_ref_preserves_tuple_alignment_on_alloca_and_store() { - let mut ctx = make_ctx(); - let tuple_ty = over_aligned_tuple_ty(&mut ctx); - let mir_ptr_ty = - MirPtrType::get_generic_with_kind(&mut ctx, tuple_ty, false, MirPointerKind::SharedRef); - let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); - - let undef = mir::MirUndefOp::new(&mut ctx, tuple_ty); - undef.get_operation().insert_at_back(block, &ctx); - let value = undef.get_operation().deref(&ctx).get_result(0); - let ref_op = Operation::new( - &mut ctx, - mir::MirRefOp::get_concrete_op_info(), - vec![mir_ptr_ty.into()], - vec![value], - vec![], - 0, - ); - let mir_ref = mir::MirRefOp::new(ref_op); - mir_ref.set_mutable(&mut ctx, false); - mir_ref.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); - ref_op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - let alloca = find_first::(&ctx, &body).expect("expected llvm.alloca"); - let store = find_first::(&ctx, &body).expect("expected llvm.store"); - assert_eq!( - llvm_export::ops::op_alignment(&ctx, alloca.get_operation()), - Some(32) - ); - assert_eq!( - llvm_export::ops::op_alignment(&ctx, store.get_operation()), - Some(32) - ); - } - - #[test] - fn convert_ref_preserves_over_aligned_union_array_layout_and_alignment() { - for abi_align in [32, 64] { - let mut ctx = make_ctx(); - let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); - let union_ty: TypeHandle = MirUnionType::get( - &mut ctx, - format!("Align{abi_align}Union"), - vec!["word".into()], - vec![u32_ty], - abi_align, - abi_align, - ) - .into(); - let array_ty: TypeHandle = MirArrayType::get(&mut ctx, union_ty, 3).into(); - let mir_ptr_ty = MirPtrType::get_generic_with_kind( - &mut ctx, - array_ty, - false, - MirPointerKind::SharedRef, - ); - let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); - - let undef = mir::MirUndefOp::new(&mut ctx, array_ty); - undef.get_operation().insert_at_back(block, &ctx); - let value = undef.get_operation().deref(&ctx).get_result(0); - let ref_op = Operation::new( - &mut ctx, - mir::MirRefOp::get_concrete_op_info(), - vec![mir_ptr_ty.into()], - vec![value], - vec![], - 0, - ); - let mir_ref = mir::MirRefOp::new(ref_op); - mir_ref.set_mutable(&mut ctx, false); - mir_ref.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); - ref_op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - let alloca = find_first::(&ctx, &body).expect("expected llvm.alloca"); - let store = find_first::(&ctx, &body).expect("expected llvm.store"); - let llvm_array_ty = alloca.result_pointee_type(&ctx); - let llvm_array_data = llvm_array_ty.deref(&ctx); - let llvm_array = llvm_array_data - .downcast_ref::() - .expect("over-aligned union array must remain an LLVM array"); - - assert_eq!(llvm_array.size(), 3); - assert_eq!( - crate::convert::types::llvm_type_size_align(&ctx, llvm_array.elem_type()), - Some((abi_align, 16)) - ); - assert_eq!( - crate::convert::types::llvm_type_size_align(&ctx, llvm_array_ty), - Some((abi_align * 3, 16)) - ); - assert_eq!( - llvm_export::ops::op_alignment(&ctx, alloca.get_operation()), - Some(abi_align as u32) - ); - assert_eq!( - llvm_export::ops::op_alignment(&ctx, store.get_operation()), - Some(abi_align as u32) - ); - } - } - - #[test] - fn convert_ptr_offset_lowers_to_gep_with_pointee_elem_type() { - let mut ctx = make_ctx(); - let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); - let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signless).into(); - let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty, true); - - let (module_ptr, block) = build_kernel(&mut ctx, vec![mir_ptr_ty.into(), i64_ty], vec![]); - let ptr_val = block.deref(&ctx).get_argument(0); - let off_val = block.deref(&ctx).get_argument(1); - - let off_op = Operation::new( - &mut ctx, - mir::MirPtrOffsetOp::get_concrete_op_info(), - vec![mir_ptr_ty.into()], - vec![ptr_val, off_val], - vec![], - 0, - ); - off_op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - let gep = find_first::(&ctx, &body).expect("expected one llvm.gep"); - // Element type must come from the MirPtrType pointee (i32), not the - // i8 fallback used when no operand-type info is available. - let elem_ty = gep.src_elem_type(&ctx); - let elem_ty_ref = elem_ty.deref(&ctx); - let int_ty = elem_ty_ref - .downcast_ref::() - .expect("gep src_elem_type should be IntegerType"); - assert_eq!(int_ty.width(), 32, "gep elem type must be i32 (pointee)"); - assert!( - llvm::gep_inbounds(&ctx, gep.get_operation()), - "ordinary pointer offsets retain the in-bounds contract" - ); - } - - #[test] - fn convert_wrapping_ptr_offset_lowers_to_non_inbounds_gep() { - let mut ctx = make_ctx(); - let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); - let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signed).into(); - let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty, false); - - let (module_ptr, block) = build_kernel(&mut ctx, vec![mir_ptr_ty.into(), i64_ty], vec![]); - let ptr_val = block.deref(&ctx).get_argument(0); - let off_val = block.deref(&ctx).get_argument(1); - - let off_op = Operation::new( - &mut ctx, - mir::MirPtrOffsetOp::get_concrete_op_info(), - vec![mir_ptr_ty.into()], - vec![ptr_val, off_val], - vec![], - 0, - ); - mir::MirPtrOffsetOp::new(off_op).set_inbounds(&mut ctx, false); - off_op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let body = kernel_blocks(&ctx, module_ptr); - let gep = find_first::(&ctx, &body).expect("expected one llvm.gep"); - assert!( - !llvm::gep_inbounds(&ctx, gep.get_operation()), - "wrapping pointer offsets must not promise in-bounds arithmetic" - ); - } - - // ========================================================================= - // Enum layout: converted width per shape + divergent-enum rejection - // ========================================================================= - - use dialect_mir::types::{EnumVariant, MirEnumType}; - - /// Build a Direct-tag `MirEnumType` the way the importer does: - /// unsigned tag of `tag_bits`, plus rustc's `total_size`/`abi_align`. - fn make_enum_ty( - ctx: &mut Context, - name: &str, - tag_bits: u32, - variants: Vec, - total_size: u64, - abi_align: u64, - ) -> TypeHandle { - let tag_ty: TypeHandle = IntegerType::get(ctx, tag_bits, Signedness::Unsigned).into(); - // Sequential 0..n discriminants: these layout tests only exercise - // size/width, not value mapping. - let discriminants: Vec = (0..variants.len() as u64).collect(); - MirEnumType::get_with_layout( - ctx, - name.to_string(), - tag_ty, - discriminants, - variants, - 0, // tag at byte 0, like every shape these tests exercise - total_size, - abi_align, - ) - .into() - } - - fn unit_variants(n: usize) -> Vec { - (0..n).map(|i| EnumVariant::unit(format!("V{i}"))).collect() - } - - /// Converted enum allocation size must equal rustc's `total_size` for - /// every memory-faithful tag shape: that size is what GEP strides by. - #[test] - fn enum_conversion_strides_by_rustc_size() { - use crate::convert::types::llvm_type_size_align; - - let mut ctx = make_ctx(); - - // #[repr(u32)] fieldless (issue #118 shape): {i32}, 4 bytes. - let repr_u32 = make_enum_ty(&mut ctx, "ReprU32", 32, unit_variants(4), 4, 4); - let conv = convert_type(&mut ctx, repr_u32).unwrap(); - assert_eq!( - llvm_type_size_align(&ctx, conv), - Some((4, 4)), - "repr(u32) tag" - ); - - // #[repr(usize)] fieldless: {i64}, 8 bytes. - let repr_usize = make_enum_ty(&mut ctx, "ReprUsize", 64, unit_variants(4), 8, 8); - let conv = convert_type(&mut ctx, repr_usize).unwrap(); - assert_eq!( - llvm_type_size_align(&ctx, conv), - Some((8, 8)), - "repr(usize) tag" - ); - - // repr(align(8)) with an i8 tag: the byte claims alone would give - // alignment 1, so the slot map raises the storage alignment with a - // zero-length anchor field. Size and alignment must both match - // rustc, or arrays and by-value ABI uses would be unsound. - let padded = make_enum_ty(&mut ctx, "Padded", 8, unit_variants(2), 8, 8); - let conv = convert_type(&mut ctx, padded).unwrap(); - assert_eq!( - llvm_type_size_align(&ctx, conv), - Some((8, 8)), - "repr(align(8)) i8 tag" - ); - - // u8 tag + i64 payload, rustc size 16: the slot map places the - // payload at its rustc byte offset 8 behind an explicit - // [7 x i8] filler, making the layout datalayout-independent. - let i64_payload: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); - let payload = make_enum_ty( - &mut ctx, - "OnePayload", - 8, - vec![ - EnumVariant::new_with_layout("A".to_string(), vec![i64_payload], vec![8], vec![8]), - EnumVariant::unit("B".to_string()), - ], - 16, - 8, - ); - let conv = convert_type(&mut ctx, payload).unwrap(); - let (size, _align) = llvm_type_size_align(&ctx, conv).unwrap(); - assert_eq!(size, 16, "natural layout matches rustc size, no pad"); - let conv_ref = conv.deref(&ctx); - let struct_ty = conv_ref - .downcast_ref::() - .expect("converted enum is a struct"); - assert_eq!( - struct_ty.fields().count(), - 3, - "{{tag, [7 x i8] filler, payload}}: explicit filler to byte 8" - ); - } - - /// Multi-payload enum: variants overlap in Rust, and identical - /// (offset, converted type) payloads share one typed slot, so the - /// converted struct is byte-identical to rustc's layout AND every - /// access stays pure SSA (no spill). - #[test] - fn multi_payload_enum_shares_payload_slot() { - use crate::convert::types::{build_enum_slot_map, llvm_type_size_align}; - - let mut ctx = make_ctx(); - let e = make_multi_payload_enum_ty(&mut ctx); - let map = build_enum_slot_map(&mut ctx, e).unwrap(); - assert_eq!(map.carrier_slot, Some(0)); - assert_eq!( - map.field_slots, - vec![Some(1), Some(1)], - "A.0 and B.0 overlap at byte 4 with the same type: one shared slot" - ); - assert_eq!( - llvm_type_size_align(&ctx, map.llvm_struct_ty), - Some((8, 4)), - "byte-identical to rustc's 8-byte layout, not the 12-byte concat" - ); - } - - /// rustc may place the tag AFTER payload bytes; the slot map must - /// follow the recorded tag_offset, never assume slot 0. - #[test] - fn enum_slot_map_tag_not_at_zero() { - use crate::convert::types::{build_enum_slot_map, llvm_type_size_align}; - - let mut ctx = make_ctx(); - let u64_a: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); - let u64_b: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); - let tag_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); - // enum F { A(u64), B(u64) }: payloads share byte 0, tag at byte 8. - let ty: TypeHandle = MirEnumType::get_with_layout( - &mut ctx, - "TagAtEight".to_string(), - tag_ty, - vec![0, 1], - vec![ - EnumVariant::new_with_layout("A".to_string(), vec![u64_a], vec![0], vec![8]), - EnumVariant::new_with_layout("B".to_string(), vec![u64_b], vec![0], vec![8]), - ], - 8, - 16, - 8, - ) - .into(); - let map = build_enum_slot_map(&mut ctx, ty).unwrap(); - assert_eq!( - map.field_slots, - vec![Some(0), Some(0)], - "payloads share the first slot" - ); - assert_eq!( - map.carrier_slot, - Some(1), - "tag claims its own slot at byte 8" - ); - let (size, _align) = llvm_type_size_align(&ctx, map.llvm_struct_ty).unwrap(); - assert_eq!(size, 16, "{{ i64, i8, [7 x i8] }}"); - } - - /// Multi-payload enum whose variants overlap in Rust: both payloads use - /// bytes 4..8 after the direct `u32` tag, so the complete value is eight - /// bytes. Mimics `#[repr(u32)] enum E { A(u32), B(u32) }`. - fn make_multi_payload_enum_ty(ctx: &mut Context) -> TypeHandle { - let i32_a: TypeHandle = IntegerType::get(ctx, 32, Signedness::Unsigned).into(); - let i32_b: TypeHandle = IntegerType::get(ctx, 32, Signedness::Unsigned).into(); - make_enum_ty( - ctx, - "MultiPayload", - 32, - vec![ - EnumVariant::new_with_layout("A".to_string(), vec![i32_a], vec![4], vec![4]), - EnumVariant::new_with_layout("B".to_string(), vec![i32_b], vec![4], vec![4]), - ], - 8, - 4, - ) - } - - /// Device-local GEP + load must use the same exact eight-byte rustc layout - /// as every other enum path. This protects pointer stride in issue #131's - /// in-kernel `[E; 4]` arrays; the sibling kernel-parameter test proves the - /// same representation is also valid at the host/device boundary. - #[test] - fn device_local_multi_payload_enum_gep_and_load_lower() { - let mut ctx = make_ctx(); - let enum_ty = make_multi_payload_enum_ty(&mut ctx); - let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signless).into(); - let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, enum_ty, true); - - let (module_ptr, block) = build_kernel(&mut ctx, vec![mir_ptr_ty.into(), i64_ty], vec![]); - let ptr_val = block.deref(&ctx).get_argument(0); - let off_val = block.deref(&ctx).get_argument(1); - - let off_op = Operation::new( - &mut ctx, - mir::MirPtrOffsetOp::get_concrete_op_info(), - vec![mir_ptr_ty.into()], - vec![ptr_val, off_val], - vec![], - 0, - ); - off_op.insert_at_back(block, &ctx); - let elem_ptr = off_op.deref(&ctx).get_result(0); - - let load_op = Operation::new( - &mut ctx, - mir::MirLoadOp::get_concrete_op_info(), - vec![enum_ty], - vec![elem_ptr], - vec![], - 0, - ); - load_op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr) - .expect("device-local rustc-layout enum GEP + load must lower"); - } - - /// A kernel parameter may carry this enum because the device slot map is - /// byte-identical to rustc's overlapped host layout. - #[test] - fn kernel_param_accepts_multi_payload_enum() { - use pliron::builtin::attributes::StringAttr; - - let mut ctx = make_ctx(); - let enum_ty = make_multi_payload_enum_ty(&mut ctx); - let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, enum_ty, false); - - let (module_ptr, block) = build_kernel(&mut ctx, vec![mir_ptr_ty.into()], vec![]); - append_mir_return(&mut ctx, block, vec![]); - - // Mark the function as a GPU kernel the way the importer does. - { - let module_block = module_ptr - .deref(&ctx) - .get_region(0) - .deref(&ctx) - .iter(&ctx) - .next() - .unwrap(); - let func_op = module_block.deref(&ctx).iter(&ctx).next().unwrap(); - let kernel_attr = StringAttr::new("true".to_string()); - let key: pliron::identifier::Identifier = "gpu_kernel".try_into().unwrap(); - func_op.deref_mut(&ctx).attributes.set(key, kernel_attr); - } - - // The slot map lowers MultiPayload byte-identically to rustc's - // layout ({ i32, i32 }, 8 bytes), so the kernel ABI accepts it. - crate::lower_mir_to_llvm(&mut ctx, module_ptr) - .expect("multi-payload enum kernel param must lower"); - } - - /// A legacy enum without physical rustc metadata must be rejected at the - /// kernel boundary. Importer-produced niche layouts are now known and are - /// accepted; this fixture deliberately constructs `Unknown` metadata. - #[test] - fn kernel_param_rejects_enum_with_unknown_layout() { - use pliron::builtin::attributes::StringAttr; - - let mut ctx = make_ctx(); - // MirEnumType::get with no layout is the legacy Unknown form. - let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); - let pointee = MirPtrType::get_generic(&mut ctx, i32_ty, false); - let tag_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); - let niched: TypeHandle = MirEnumType::get( - &mut ctx, - "Option".to_string(), - tag_ty, - vec![0, 1], - vec![ - EnumVariant::unit("None".to_string()), - EnumVariant::new("Some".to_string(), vec![pointee.into()]), - ], - ) - .into(); - let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, niched, false); - - let (module_ptr, block) = build_kernel(&mut ctx, vec![mir_ptr_ty.into()], vec![]); - append_mir_return(&mut ctx, block, vec![]); - - { - let module_block = module_ptr - .deref(&ctx) - .get_region(0) - .deref(&ctx) - .iter(&ctx) - .next() - .unwrap(); - let func_op = module_block.deref(&ctx).iter(&ctx).next().unwrap(); - let kernel_attr = StringAttr::new("true".to_string()); - let key: pliron::identifier::Identifier = "gpu_kernel".try_into().unwrap(); - func_op.deref_mut(&ctx).attributes.set(key, kernel_attr); - } - - let err = crate::lower_mir_to_llvm(&mut ctx, module_ptr) - .expect_err("unknown-layout enum kernel param must be rejected"); - let msg = format!("{err}"); - assert!( - msg.contains("Option") && msg.contains("kernel boundary"), - "error must name the enum and the kernel boundary, got: {msg}" - ); - assert!(msg.contains("unknown physical rustc layout")); - } - - /// Build a `mir.shared_alloc` returning `MirPtrType` of - /// length `size`, with the given alloc_key, and append it to `block`. - fn append_shared_alloc( - ctx: &mut Context, - block: Ptr, - alloc_key: &str, - size: u64, - ) -> Ptr { - append_shared_alloc_named(ctx, block, alloc_key, size, None) - } - - /// As [`append_shared_alloc`], additionally carrying the Rust path of the - /// `static` the allocation came from. - fn append_shared_alloc_named( - ctx: &mut Context, - block: Ptr, - alloc_key: &str, - size: u64, - source_name: Option<&str>, - ) -> Ptr { - let i32_ty: TypeHandle = IntegerType::get(ctx, 32, Signedness::Signless).into(); - append_shared_alloc_typed(ctx, block, alloc_key, i32_ty, size, source_name, 0) - } - - #[allow(clippy::too_many_arguments)] - fn append_shared_alloc_typed( - ctx: &mut Context, - block: Ptr, - alloc_key: &str, - element_type: TypeHandle, - size: u64, - source_name: Option<&str>, - alignment: u64, - ) -> Ptr { - use pliron::builtin::attributes::IntegerAttr; - use pliron::utils::apint::APInt; - - let i32_ty = element_type; - let result_ty = MirPtrType::get_shared(ctx, i32_ty, true); - let op = Operation::new( - ctx, - mir::MirSharedAllocOp::get_concrete_op_info(), - vec![result_ty.into()], - vec![], - vec![], - 0, - ); - let alloc = mir::MirSharedAllocOp::new(op); - alloc.set_attr_elem_type(ctx, TypeAttr::new(i32_ty)); - let size_attr = IntegerAttr::new( - IntegerType::get(ctx, 64, Signedness::Signless), - APInt::from_u64(size, std::num::NonZeroUsize::new(64).unwrap()), - ); - alloc.set_attr_size(ctx, size_attr); - alloc.set_attr_alloc_key(ctx, StringAttr::new(alloc_key.to_string())); - if let Some(source_name) = source_name { - alloc.set_attr_source_name(ctx, StringAttr::new(source_name.to_string())); - } - if alignment != 0 { - alloc.set_alignment_value(ctx, alignment); - } - op.insert_at_back(block, ctx); - op - } - - fn shared_array_debug_info(count: u64) -> llvm::DebugGlobalVariableInfo { - llvm::DebugGlobalVariableInfo { - name: "TILE".to_string(), - namespace: vec!["fixture".to_string(), "kernel".to_string()], - ty: llvm::DebugLocalTypeKind::Array { - name: format!("[i32; {count}]"), - size_bits: count * 32, - element: Box::new(llvm::DebugLocalTypeKind::Basic { - name: "i32".to_string(), - size_bits: 32, - encoding: "DW_ATE_signed", - }), - count, - }, - declaration: llvm::DebugSourcePosition { - file: PathBuf::from("/tmp/shared.rs"), - line: 7, - column: 5, - }, - is_local_to_unit: true, - is_function_local: true, - } - } - - fn shared_static_key(symbol: &str) -> String { - dialect_mir::ops::encode_rust_static_global_key(symbol) - } - - #[test] - fn convert_shared_alloc_creates_global_in_addrspace_3() { - let mut ctx = make_ctx(); - let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); - append_shared_alloc(&mut ctx, block, "k1", 64); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - // Global lives at module level; addressof lives in the function body. - let top = module_top_block(&ctx, module_ptr); - let global = top - .deref(&ctx) - .iter(&ctx) - .find_map(|op| Operation::get_op::(op, &ctx)) - .expect("expected an llvm.global for the shared allocation"); - assert_eq!( - global.address_space(&ctx), - llvm_addr::SHARED, - "shared_alloc global must live in addrspace 3" - ); - assert!( - global - .get_symbol_name(&ctx) - .to_string() - .starts_with("__shared_mem_"), - "shared global should have __shared_mem_ prefix" - ); - - let body = kernel_blocks(&ctx, module_ptr); - let addrof = - find_first::(&ctx, &body).expect("expected an llvm.addressof"); - assert_eq!( - ptr_addrspace( - &ctx, - addrof - .get_operation() - .deref(&ctx) - .get_result(0) - .get_type(&ctx) - ), - llvm_addr::SHARED, - ); - } - - #[test] - fn convert_shared_alloc_deduplicates_by_alloc_key() { - let mut ctx = make_ctx(); - let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); - // Two allocations sharing the same alloc_key — they must collapse to - // a single underlying global (this is what enables a single `static` - // referenced from multiple sites to land in one shared region). - append_shared_alloc(&mut ctx, block, "same-key", 64); - append_shared_alloc(&mut ctx, block, "same-key", 64); - // A third with a different key must NOT dedupe with them. - append_shared_alloc(&mut ctx, block, "other-key", 32); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let top = module_top_block(&ctx, module_ptr); - let shared_globals = top - .deref(&ctx) - .iter(&ctx) - .filter_map(|op| Operation::get_op::(op, &ctx)) - .filter(|g| g.address_space(&ctx) == llvm_addr::SHARED) - .count(); - assert_eq!( - shared_globals, 2, - "two distinct alloc_keys must produce two globals" - ); - - // Each of the three mir.shared_alloc ops becomes one addressof. - let body = kernel_blocks(&ctx, module_ptr); - assert_eq!(count_ops::(&ctx, &body), 3); - } - - #[test] - fn convert_shared_alloc_rejects_conflicting_key_declarations() { - let mut ctx = make_ctx(); - let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); - append_shared_alloc(&mut ctx, block, "conflicting-key", 1); - let differently_aligned = append_shared_alloc(&mut ctx, block, "conflicting-key", 1); - mir::MirSharedAllocOp::new(differently_aligned).set_alignment_value(&mut ctx, 16); - append_mir_return(&mut ctx, block, vec![]); - - let error = crate::lower_mir_to_llvm(&mut ctx, module_ptr) - .expect_err("one shared alloc_key must not name differently aligned storage"); - assert!( - error.to_string().contains("incompatible declaration"), - "{error}" - ); - } - - #[test] - fn shared_alloc_repeated_debug_identity_is_attached_once() { - let mut ctx = make_ctx(); - let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); - let key = shared_static_key("_Rfixture4TILE"); - for _ in 0..2 { - let op = - append_shared_alloc_named(&mut ctx, block, &key, 32, Some("fixture::kernel::TILE")); - let info = shared_array_debug_info(32); - llvm::set_debug_global_variable(&mut ctx, op, &info); - llvm::set_debug_global_owner_function(&mut ctx, op, "fixture_kernel"); - } - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - let top = module_top_block(&ctx, module_ptr); - let globals: Vec<_> = top - .deref(&ctx) - .iter(&ctx) - .filter_map(|op| Operation::get_op::(op, &ctx)) - .filter(|global| global.address_space(&ctx) == llvm_addr::SHARED) - .collect(); - assert_eq!( - globals.len(), - 1, - "repeated references must share one global" - ); - assert_eq!( - llvm::debug_global_variable(&ctx, globals[0].get_operation()), - Some(shared_array_debug_info(32)) - ); - assert_eq!( - llvm::debug_global_owner_function(&ctx, globals[0].get_operation()).as_deref(), - Some("fixture_kernel") - ); - } - - /// One function-local shared static materialized in two owning functions - /// (for example a `#[device]` helper inlined into two kernels) is a valid - /// release build, so debug metadata must not turn it into an error. DWARF - /// cannot truthfully scope one AS3 object to two subprograms either, so - /// the divergent identity fails open: the storage is still shared and the - /// debug attachment is dropped. - #[test] - fn shared_alloc_owner_conflict_drops_debug_metadata() { - let mut ctx = make_ctx(); - let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); - let key = shared_static_key("_Rfixture4TILE"); - for owner in ["fixture_kernel", "other_kernel"] { - let op = - append_shared_alloc_named(&mut ctx, block, &key, 32, Some("fixture::kernel::TILE")); - let info = shared_array_debug_info(32); - llvm::set_debug_global_variable(&mut ctx, op, &info); - llvm::set_debug_global_owner_function(&mut ctx, op, owner); - } - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr) - .expect("divergent debug identity must not fail a build release accepts"); - let top = module_top_block(&ctx, module_ptr); - let globals: Vec<_> = top - .deref(&ctx) - .iter(&ctx) - .filter_map(|op| Operation::get_op::(op, &ctx)) - .filter(|global| global.address_space(&ctx) == llvm_addr::SHARED) - .collect(); - assert_eq!( - globals.len(), - 1, - "the physical storage must still be shared" - ); - assert!(llvm::debug_global_variable(&ctx, globals[0].get_operation()).is_none()); - assert!(llvm::debug_global_owner_function(&ctx, globals[0].get_operation()).is_none()); - } - - #[test] - fn dynamic_extern_then_static_shared_rejects_reserved_key_collision() { - use pliron::builtin::attributes::IntegerAttr; - use pliron::utils::apint::APInt; - - let mut ctx = make_ctx(); - let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); - let i8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Signless).into(); - let shared_i8 = MirPtrType::get_shared(&mut ctx, i8_ty, true); - - // Lower the dynamic declaration first. Its internal cache key is - // deliberately then reused by a static allocation with the same - // physical type, size, and alignment. Those declarations must still - // be rejected because one is extern storage and one is fixed storage. - let dynamic = Operation::new( - &mut ctx, - mir::MirExternSharedOp::get_concrete_op_info(), - vec![shared_i8.into()], - vec![], - vec![], - 0, - ); - mir::MirExternSharedOp::new(dynamic).set_alignment_value(&mut ctx, 128); - dynamic.insert_at_back(block, &ctx); - - let fixed = Operation::new( - &mut ctx, - mir::MirSharedAllocOp::get_concrete_op_info(), - vec![shared_i8.into()], - vec![], - vec![], - 0, - ); - let fixed_alloc = mir::MirSharedAllocOp::new(fixed); - fixed_alloc.set_attr_elem_type(&ctx, TypeAttr::new(i8_ty)); - fixed_alloc.set_attr_size( - &ctx, - IntegerAttr::new( - IntegerType::get(&ctx, 64, Signedness::Unsigned), - APInt::from_u64(0, std::num::NonZeroUsize::new(64).unwrap()), - ), - ); - fixed_alloc.set_attr_alloc_key( - &ctx, - StringAttr::new("__dynamic_smem_global_created_kernel_func".to_string()), - ); - fixed_alloc.set_alignment_value(&mut ctx, 128); - fixed.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - let error = crate::lower_mir_to_llvm(&mut ctx, module_ptr) - .expect_err("a static allocation must not reuse a dynamic extern declaration"); - let message = error.to_string(); - assert!(message.contains("incompatible declaration"), "{message}"); - assert!( - message.contains("duplicate shared alloc_key"), - "the dynamic declaration must be observed first, got: {message}" - ); - } - - #[test] - fn shared_alloc_mismatched_debug_type_fails_closed_without_failing_lowering() { - let mut ctx = make_ctx(); - let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); - let key = shared_static_key("_Rfixture4TILE"); - let op = - append_shared_alloc_named(&mut ctx, block, &key, 32, Some("fixture::kernel::TILE")); - llvm::set_debug_global_variable(&mut ctx, op, &shared_array_debug_info(31)); - llvm::set_debug_global_owner_function(&mut ctx, op, "fixture_kernel"); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr) - .expect("bad optional metadata must not fail compilation"); - let top = module_top_block(&ctx, module_ptr); - let global = top - .deref(&ctx) - .iter(&ctx) - .find_map(|candidate| Operation::get_op::(candidate, &ctx)) - .expect("shared global"); - assert!(llvm::debug_global_variable(&ctx, global.get_operation()).is_none()); - assert!(llvm::debug_global_owner_function(&ctx, global.get_operation()).is_none()); - } - - #[test] - fn barrier_semantic_struct_matches_single_i64_shared_backing() { - let mut ctx = make_ctx(); - let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); - let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); - let key = shared_static_key("_Rfixture7BARRIER"); - let op = append_shared_alloc_typed( - &mut ctx, - block, - &key, - i64_ty, - 1, - Some("fixture::kernel::BARRIER"), - 8, - ); - let info = llvm::DebugGlobalVariableInfo { - name: "BARRIER".to_string(), - namespace: vec!["fixture".to_string(), "kernel".to_string()], - ty: llvm::DebugLocalTypeKind::Struct { - name: "Barrier".to_string(), - size_bits: 64, - members: vec![llvm::DebugTypeMember { - name: "_state".to_string(), - offset_bits: 0, - ty: llvm::DebugLocalTypeKind::Basic { - name: "u64".to_string(), - size_bits: 64, - encoding: "DW_ATE_unsigned", - }, - }], - }, - declaration: llvm::DebugSourcePosition { - file: PathBuf::from("/tmp/barrier.rs"), - line: 8, - column: 9, - }, - is_local_to_unit: true, - is_function_local: true, - }; - llvm::set_debug_global_variable(&mut ctx, op, &info); - llvm::set_debug_global_owner_function(&mut ctx, op, "fixture_kernel"); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("Barrier lowering succeeds"); - let top = module_top_block(&ctx, module_ptr); - let global = top - .deref(&ctx) - .iter(&ctx) - .find_map(|candidate| Operation::get_op::(candidate, &ctx)) - .expect("Barrier shared global"); - assert_eq!(global.get_alignment(&ctx), Some(8)); - assert_eq!( - llvm::debug_global_variable(&ctx, global.get_operation()), - Some(info) - ); - } - - /// Collect `(symbol, source_name)` for every shared global in the module. - fn shared_global_source_names( - ctx: &Context, - module_ptr: Ptr, - ) -> Vec<(String, Option)> { - use llvm_export::ops::GlobalOpExt; - - let top = module_top_block(ctx, module_ptr); - let mut named: Vec<_> = top - .deref(ctx) - .iter(ctx) - .filter_map(|op| Operation::get_op::(op, ctx)) - .filter(|g| g.address_space(ctx) == llvm_addr::SHARED) - .map(|g| { - ( - g.get_symbol_name(ctx).to_string(), - g.shared_source_name(ctx), - ) - }) - .collect(); - // Globals are inserted at the front of the module block, so iteration - // order is the reverse of creation order. Sort for a stable assertion. - named.sort(); - named - } - - #[test] - fn shared_alloc_source_name_reaches_the_generated_global() { - let mut ctx = make_ctx(); - let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); - append_shared_alloc_named(&mut ctx, block, "k1", 64, Some("my_kernel::TILE")); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let named = shared_global_source_names(&ctx, module_ptr); - assert_eq!(named.len(), 1, "expected exactly one shared global"); - let (symbol, source_name) = &named[0]; - // The symbol itself must stay anonymous: the whole point of the - // sidecar attribute is that it does not perturb the emitted name. - assert!( - symbol.starts_with("__shared_mem_"), - "the generated symbol must not be renamed, got `{symbol}`" - ); - assert_eq!(source_name.as_deref(), Some("my_kernel::TILE")); - } - - #[test] - fn shared_alloc_without_source_name_leaves_the_global_unlabelled() { - let mut ctx = make_ctx(); - let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); - append_shared_alloc(&mut ctx, block, "k1", 64); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let named = shared_global_source_names(&ctx, module_ptr); - assert_eq!(named.len(), 1); - assert_eq!(named[0].1, None, "an unnamed allocation must stay unnamed"); - } - - #[test] - fn shared_alloc_source_names_are_per_global_not_shared_across_them() { - let mut ctx = make_ctx(); - let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); - // Two references to one static dedupe onto a single global, and a - // second static gets its own. Each global must carry its own name — - // the failure this guards is one name leaking onto every allocation. - append_shared_alloc_named(&mut ctx, block, "tile", 64, Some("my_kernel::TILE")); - append_shared_alloc_named(&mut ctx, block, "tile", 64, Some("my_kernel::TILE")); - append_shared_alloc_named(&mut ctx, block, "scratch", 32, Some("my_kernel::SCRATCH")); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let names: Vec<_> = shared_global_source_names(&ctx, module_ptr) - .into_iter() - .map(|(_, source_name)| source_name) - .collect(); - assert_eq!(names.len(), 2, "the shared alloc_key must still dedupe"); - let mut names: Vec<_> = names.into_iter().map(|n| n.expect("named")).collect(); - names.sort(); - assert_eq!(names, vec!["my_kernel::SCRATCH", "my_kernel::TILE"]); - } - - #[test] - fn shared_alloc_source_name_reaches_the_exported_llvm_ir() { - // The end the feature exists for: a consumer holding only the emitted - // artifact can tell which Rust `static` a `__shared_mem_N` block is. - let mut ctx = make_ctx(); - let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); - append_shared_alloc_named(&mut ctx, block, "tile", 64, Some("my_kernel::TILE")); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let module = Operation::get_op::(module_ptr, &ctx) - .expect("lowered top-level op is a module"); - let ir = llvm_export::export::export_module_to_string(&ctx, &module).expect("export"); - - let comment_index = ir - .find("; shared source: my_kernel::TILE") - .unwrap_or_else(|| panic!("exported IR must name the shared source:\n{ir}")); - let definition_index = ir - .find("__shared_mem_") - .expect("exported IR must declare the shared global"); - assert!( - comment_index < definition_index, - "the source comment must precede the global it describes:\n{ir}" - ); - } - - /// A `__shared_mem_N` or `__device_global_N` index must depend only on - /// the module being lowered, not on how many allocations any OTHER - /// module has already lowered in this process (#706). Before the fix, - /// each `N` came from a `static AtomicUsize` shared across every call in - /// the process, so lowering the second of these two modules would have - /// produced `__shared_mem_1` and `__device_global_1`, not the `_0` names. - #[test] - fn shared_and_device_global_indices_are_per_module_not_process_global() { - for _ in 0..2 { - let mut ctx = make_ctx(); - let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); - append_shared_alloc(&mut ctx, block, "k", 64); - append_global_alloc(&mut ctx, block, "ordinary_static", false); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let top = module_top_block(&ctx, module_ptr); - let names: Vec = top - .deref(&ctx) - .iter(&ctx) - .filter_map(|op| Operation::get_op::(op, &ctx)) - .map(|g| g.get_symbol_name(&ctx).to_string()) - .collect(); - assert!( - names.iter().any(|n| n == "__shared_mem_0"), - "a module with exactly one shared allocation must always name it \ - __shared_mem_0, regardless of how many other modules already lowered \ - one in this process (got {names:?})" - ); - assert!( - names.iter().any(|n| n == "__device_global_0"), - "a module with exactly one ordinary device global must always name it \ - __device_global_0, regardless of how many other modules already \ - lowered one in this process (got {names:?})" - ); - } - } - - fn append_global_alloc( - ctx: &mut Context, - block: Ptr, - global_key: &str, - constant: bool, - ) -> Ptr { - let i32_ty: TypeHandle = IntegerType::get(ctx, 32, Signedness::Signless).into(); - append_global_alloc_typed(ctx, block, global_key, constant, !constant, i32_ty) - } - - fn append_global_alloc_typed( - ctx: &mut Context, - block: Ptr, - global_key: &str, - constant: bool, - is_mutable: bool, - global_ty: TypeHandle, - ) -> Ptr { - let result_ty = if constant { - MirPtrType::get_constant(ctx, global_ty, is_mutable) - } else { - MirPtrType::get_global(ctx, global_ty, is_mutable) - }; - let op = Operation::new( - ctx, - mir::MirGlobalAllocOp::get_concrete_op_info(), - vec![result_ty.into()], - vec![], - vec![], - 0, - ); - let alloc = mir::MirGlobalAllocOp::new(op); - alloc.set_attr_global_type(ctx, TypeAttr::new(global_ty)); - alloc.set_attr_global_key(ctx, StringAttr::new(global_key.to_string())); - op.insert_at_back(block, ctx); - op - } - - #[test] - fn convert_global_alloc_deduplicates_matching_storage_across_pointer_mutability() { - let mut ctx = make_ctx(); - let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); - let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); - append_global_alloc_typed(&mut ctx, block, "same-global", false, false, i32_ty); - append_global_alloc_typed(&mut ctx, block, "same-global", false, true, i32_ty); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr) - .expect("pointer-carrier mutability does not change physical global storage"); - - let top = module_top_block(&ctx, module_ptr); - let globals = top - .deref(&ctx) - .iter(&ctx) - .filter_map(|op| Operation::get_op::(op, &ctx)) - .filter(|global| global.source_global_key(&ctx).as_deref() == Some("same-global")) - .count(); - assert_eq!(globals, 1); - assert_eq!( - count_ops::(&ctx, &kernel_blocks(&ctx, module_ptr)), - 2 - ); - } - - #[test] - fn convert_global_alloc_rejects_conflicting_key_declarations() { - let mut ctx = make_ctx(); - let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); - let byte_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Signless).into(); - let word_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); - let empty_words: TypeHandle = MirArrayType::get(&mut ctx, word_ty, 0).into(); - - let byte = - append_global_alloc_typed(&mut ctx, block, "conflicting-global", false, false, byte_ty); - let byte = mir::MirGlobalAllocOp::new(byte); - byte.set_alignment_value(&mut ctx, 1); - byte.mark_immutable(&mut ctx); - - let words = append_global_alloc_typed( - &mut ctx, - block, - "conflicting-global", - false, - false, - empty_words, - ); - let words = mir::MirGlobalAllocOp::new(words); - words.set_alignment_value(&mut ctx, 4); - words.mark_immutable(&mut ctx); - append_mir_return(&mut ctx, block, vec![]); - - let error = crate::lower_mir_to_llvm(&mut ctx, module_ptr) - .expect_err("one global_key must not name incompatible type/alignment storage"); - assert!( - error.to_string().contains("incompatible declaration"), - "{error}" - ); - } - - #[test] - fn device_global_declaration_identity_covers_every_storage_property() { - let ctx = make_ctx(); - let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); - let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signless).into(); - let base = DeviceGlobalDeclaration { - mir_type: i32_ty, - alignment: 4, - addr_space: llvm_addr::GLOBAL, - initializer_hex: Some("00000000".to_string()), - initializer_relocations: Some("reloc-a".to_string()), - debug_info: None, - immutable: true, - }; - - let mut changed = base.clone(); - changed.mir_type = i64_ty; - assert!(base != changed); - changed = base.clone(); - changed.alignment = 8; - assert!(base != changed); - changed = base.clone(); - changed.addr_space = llvm_addr::CONSTANT; - assert!(base != changed); - changed = base.clone(); - changed.initializer_hex = Some("01000000".to_string()); - assert!(base != changed); - changed = base.clone(); - changed.initializer_relocations = Some("reloc-b".to_string()); - assert!(base != changed); - changed = base.clone(); - changed.debug_info = Some(llvm::DebugGlobalVariableInfo { - name: "COUNTER".to_string(), - namespace: vec!["my_crate".to_string()], - ty: llvm::DebugLocalTypeKind::Basic { - name: "u32".to_string(), - size_bits: 32, - encoding: "DW_ATE_unsigned", - }, - declaration: llvm::DebugSourcePosition { - file: PathBuf::from("/tmp/global.rs"), - line: 7, - column: 1, - }, - is_local_to_unit: true, - is_function_local: false, - }); - assert!(base != changed); - changed = base.clone(); - changed.immutable = false; - assert!(base != changed); - } - - #[test] - fn convert_global_alloc_places_in_global_or_constant_addrspace() { - let mut ctx = make_ctx(); - let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); - append_global_alloc(&mut ctx, block, "ordinary_static", false); - let constant_key = dialect_mir::ops::encode_rust_static_global_key("_ZN7my_mod3KEYE"); - append_global_alloc(&mut ctx, block, &constant_key, true); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let top = module_top_block(&ctx, module_ptr); - let globals: Vec<_> = top - .deref(&ctx) - .iter(&ctx) - .filter_map(|op| Operation::get_op::(op, &ctx)) - .collect(); - let global_addr_global = globals - .iter() - .find(|g| g.address_space(&ctx) == llvm_addr::GLOBAL) - .expect("expected one global in addrspace(1)"); - let global_addr_const = globals - .iter() - .find(|g| g.address_space(&ctx) == llvm_addr::CONSTANT) - .expect("expected one global in addrspace(4)"); - - // Constant-memory globals reuse the Rust mangled name so host code can - // resolve them by name via `cuModuleGetGlobal`; ordinary globals get - // a counter-suffixed `__device_global_N`. - assert_eq!( - global_addr_const.get_symbol_name(&ctx).to_string(), - "_ZN7my_mod3KEYE", - "constant globals must keep the raw rustc symbol payload as their name" - ); - assert!( - global_addr_global - .get_symbol_name(&ctx) - .to_string() - .starts_with("__device_global_"), - "ordinary device globals get the __device_global_ prefix" - ); - } - - #[test] - fn ordinary_global_debug_info_survives_lowering_but_constant_memory_stays_out_of_scope() { - let mut ctx = make_ctx(); - let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); - let ordinary = append_global_alloc(&mut ctx, block, "crate::GLOBAL_COUNTER", false); - let constant_key = dialect_mir::ops::encode_rust_static_global_key("_ZN5crate5COEFFE"); - let constant = append_global_alloc(&mut ctx, block, &constant_key, true); - let info = llvm::DebugGlobalVariableInfo { - name: "GLOBAL_COUNTER".to_string(), - namespace: vec!["crate".to_string(), "state".to_string()], - ty: llvm::DebugLocalTypeKind::Basic { - name: "u32".to_string(), - size_bits: 32, - encoding: "DW_ATE_unsigned", - }, - declaration: llvm::DebugSourcePosition { - file: PathBuf::from("/tmp/global.rs"), - line: 7, - column: 1, - }, - is_local_to_unit: true, - is_function_local: false, - }; - llvm::set_debug_global_variable(&mut ctx, ordinary, &info); - // Deliberately tag AS4 too: the AS1-only implementation must not - // accidentally expand its behavior merely because the generic debug - // carrier can be attached to any op. - llvm::set_debug_global_variable(&mut ctx, constant, &info); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let top = module_top_block(&ctx, module_ptr); - let globals: Vec<_> = top - .deref(&ctx) - .iter(&ctx) - .filter_map(|op| Operation::get_op::(op, &ctx)) - .collect(); - let ordinary = globals - .iter() - .find(|global| global.address_space(&ctx) == llvm_addr::GLOBAL) - .expect("ordinary global"); - let constant = globals - .iter() - .find(|global| global.address_space(&ctx) == llvm_addr::CONSTANT) - .expect("constant global"); - - assert_eq!( - llvm::debug_global_variable(&ctx, ordinary.get_operation()), - Some(info), - "source identity and semantic type must survive MIR-to-LLVM lowering" - ); - assert!( - llvm::debug_global_variable(&ctx, constant.get_operation()).is_none(), - "AS4 debug metadata is a separate feature and must not leak into this AS1 change" - ); - } - - #[test] - fn immutable_marking_survives_lowering_and_is_not_assumed() { - let mut ctx = make_ctx(); - let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); - - // Two ordinary addrspace(1) globals, distinguished by their source key. - // Only the promoted one claims immutability; the plain static must not - // acquire it, or the exporter would write `constant` for storage the - // host can still overwrite by symbol. - let promoted = append_global_alloc(&mut ctx, block, "promoted_table", false); - mir::MirGlobalAllocOp::new(promoted).mark_immutable(&mut ctx); - append_global_alloc(&mut ctx, block, "plain_static", false); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let top = module_top_block(&ctx, module_ptr); - let globals: Vec = top - .deref(&ctx) - .iter(&ctx) - .filter_map(|op| Operation::get_op::(op, &ctx)) - .collect(); - let by_key = |key: &str| -> llvm::GlobalOp { - *globals - .iter() - .find(|g| g.source_global_key(&ctx).as_deref() == Some(key)) - .unwrap_or_else(|| panic!("no lowered global carries source key {key}")) - }; - - assert!( - by_key("promoted_table").is_immutable(&ctx), - "a global marked immutable in MIR must stay immutable through lowering" - ); - assert!( - !by_key("plain_static").is_immutable(&ctx), - "lowering must not infer immutability; only the promoted-constant \ - sites may claim it" - ); - } - - #[test] - fn convert_global_alloc_rejects_conflicting_immutability() { - // A plain static and a promoted constant that end up sharing one key - // differ only in the immutable flag; that alone must fail closed - // instead of silently reusing the first allocation's storage. - let mut ctx = make_ctx(); - let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); - append_global_alloc(&mut ctx, block, "collision", false); - let promoted = append_global_alloc(&mut ctx, block, "collision", false); - mir::MirGlobalAllocOp::new(promoted).mark_immutable(&mut ctx); - append_mir_return(&mut ctx, block, vec![]); - - let error = crate::lower_mir_to_llvm(&mut ctx, module_ptr) - .expect_err("a shared string with different origins must fail closed"); - assert!( - error.to_string().contains("incompatible declaration"), - "{error}" - ); - } - - #[test] - fn initialized_global_lowers_to_byte_storage() { - let mut ctx = make_ctx(); - let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); - let op = append_global_alloc(&mut ctx, block, "nan_payload", false); - let alloc = mir::MirGlobalAllocOp::new(op); - alloc.set_alignment_value(&mut ctx, 4); - let initializer_key: Identifier = "global_initializer_hex".try_into().unwrap(); - op.deref_mut(&ctx) - .attributes - .set(initializer_key, StringAttr::new("3412c07f".to_string())); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let top = module_top_block(&ctx, module_ptr); - let global = top - .deref(&ctx) - .iter(&ctx) - .find_map(|op| Operation::get_op::(op, &ctx)) - .expect("expected lowered device global"); - let global_ty = global.get_type(&ctx); - let global_ty_ref = global_ty.deref(&ctx); - let array_ty = global_ty_ref - .downcast_ref::() - .expect("initialized global must use byte-array storage"); - assert_eq!(array_ty.size(), 4); - let elem_ty = array_ty.elem_type(); - let elem_ty_ref = elem_ty.deref(&ctx); - let elem = elem_ty_ref - .downcast_ref::() - .expect("byte-array element must be an integer"); - assert_eq!(elem.width(), 8); - assert_eq!(global.get_alignment(&ctx), Some(4)); - assert_eq!(global.initializer_hex(&ctx).as_deref(), Some("3412c07f")); - } - - #[test] - fn relocated_global_lowers_to_segmented_storage() { - let mut ctx = make_ctx(); - let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); - let word_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); - let mir_global_ty: TypeHandle = MirArrayType::get(&mut ctx, word_ty, 3).into(); - let result_ty = MirPtrType::get_global(&mut ctx, mir_global_ty, false); - let op = Operation::new( - &mut ctx, - mir::MirGlobalAllocOp::get_concrete_op_info(), - vec![result_ty.into()], - vec![], - vec![], - 0, - ); - let alloc = mir::MirGlobalAllocOp::new(op); - alloc.set_attr_global_type(&ctx, TypeAttr::new(mir_global_ty)); - alloc.set_attr_global_key(&ctx, StringAttr::new("reference_table".to_string())); - alloc.set_alignment_value(&mut ctx, 8); - op.deref_mut(&ctx).attributes.set( - "global_initializer_hex".try_into().unwrap(), - StringAttr::new("000000000000000000000000000000000000000000000000".to_string()), - ); - let encoded = - llvm::encode_global_initializer_relocations(&[llvm::GlobalInitializerRelocation { - source_offset: 8, - width_bytes: 8, - target_address_space: llvm_addr::GLOBAL, - target_addend: 4, - target_key: "target_static".to_string(), - }]); - op.deref_mut(&ctx).attributes.set( - "global_initializer_relocations".try_into().unwrap(), - StringAttr::new(encoded.clone()), - ); - op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); - - let top = module_top_block(&ctx, module_ptr); - let global = top - .deref(&ctx) - .iter(&ctx) - .find_map(|op| Operation::get_op::(op, &ctx)) - .expect("expected lowered device global"); - let global_ty = global.get_type(&ctx); - let global_ty_ref = global_ty.deref(&ctx); - let storage = global_ty_ref - .downcast_ref::() - .expect("relocated initializer must use segmented struct storage"); - assert_eq!(storage.num_fields(), 3); - assert_eq!( - global.source_global_key(&ctx).as_deref(), - Some("reference_table") - ); - assert_eq!( - global.initializer_relocations(&ctx).as_deref(), - Some(encoded.as_str()) - ); - } - - #[test] - fn relocated_global_uses_packed_storage_for_unaligned_pointer_slot() { - let mut ctx = make_ctx(); - let encoded = - llvm::encode_global_initializer_relocations(&[llvm::GlobalInitializerRelocation { - source_offset: 1, - width_bytes: 8, - target_address_space: llvm_addr::GLOBAL, - target_addend: 0, - target_key: "target_static".to_string(), - }]); - - let storage = relocated_initializer_storage_type(&mut ctx, 9, 1, &encoded) - .expect("unaligned relocation should use packed storage"); - let storage_ref = storage.deref(&ctx); - let struct_ty = storage_ref - .downcast_ref::() - .expect("relocated initializer must use struct storage"); - assert_eq!(struct_ty.layout(), StructLayout::Packed); - assert_eq!(struct_ty.num_fields(), 2); - assert_eq!( - crate::convert::types::llvm_type_size_align(&ctx, storage), - Some((9, 1)) - ); - - let fields: Vec<_> = struct_ty.fields().collect(); - let literal_ref = fields[0].deref(&ctx); - let literal = literal_ref - .downcast_ref::() - .expect("leading literal span must be a byte array"); - assert_eq!(literal.size(), 1); - let pointer_ref = fields[1].deref(&ctx); - let pointer = pointer_ref - .downcast_ref::() - .expect("relocation slot must be an integer carrier"); - assert_eq!(pointer.width(), 64); - } - - #[test] - fn relocated_global_uses_packed_storage_for_underaligned_allocation() { - let mut ctx = make_ctx(); - let encoded = - llvm::encode_global_initializer_relocations(&[llvm::GlobalInitializerRelocation { - source_offset: 0, - width_bytes: 8, - target_address_space: llvm_addr::GLOBAL, - target_addend: 0, - target_key: "target_static".to_string(), - }]); - - let storage = relocated_initializer_storage_type(&mut ctx, 8, 1, &encoded) - .expect("underaligned allocation should use packed storage"); - let storage_ref = storage.deref(&ctx); - let struct_ty = storage_ref - .downcast_ref::() - .expect("relocated initializer must use struct storage"); - assert_eq!(struct_ty.layout(), StructLayout::Packed); - assert_eq!( - crate::convert::types::llvm_type_size_align(&ctx, storage), - Some((8, 1)) - ); - } - - #[test] - fn relocated_global_keeps_naturally_aligned_storage_unpacked() { - let mut ctx = make_ctx(); - let encoded = - llvm::encode_global_initializer_relocations(&[llvm::GlobalInitializerRelocation { - source_offset: 8, - width_bytes: 8, - target_address_space: llvm_addr::GLOBAL, - target_addend: 0, - target_key: "target_static".to_string(), - }]); - - let storage = relocated_initializer_storage_type(&mut ctx, 16, 8, &encoded) - .expect("aligned relocation should keep ordinary storage"); - let storage_ref = storage.deref(&ctx); - let struct_ty = storage_ref - .downcast_ref::() - .expect("relocated initializer must use struct storage"); - assert_eq!(struct_ty.layout(), StructLayout::Unpacked); - } - - #[test] - fn relocated_global_rejects_overlapping_slots() { - let mut ctx = make_ctx(); - let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); - let word_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); - let mir_global_ty: TypeHandle = MirArrayType::get(&mut ctx, word_ty, 2).into(); - let result_ty = MirPtrType::get_global(&mut ctx, mir_global_ty, false); - let op = Operation::new( - &mut ctx, - mir::MirGlobalAllocOp::get_concrete_op_info(), - vec![result_ty.into()], - vec![], - vec![], - 0, - ); - let alloc = mir::MirGlobalAllocOp::new(op); - alloc.set_attr_global_type(&ctx, TypeAttr::new(mir_global_ty)); - alloc.set_attr_global_key(&ctx, StringAttr::new("overlap".to_string())); - alloc.set_alignment_value(&mut ctx, 8); - op.deref_mut(&ctx).attributes.set( - "global_initializer_hex".try_into().unwrap(), - StringAttr::new("00000000000000000000000000000000".to_string()), - ); - let encoded = llvm::encode_global_initializer_relocations(&[ - llvm::GlobalInitializerRelocation { - source_offset: 0, - width_bytes: 8, - target_address_space: llvm_addr::GLOBAL, - target_addend: 0, - target_key: "a".to_string(), - }, - llvm::GlobalInitializerRelocation { - source_offset: 0, - width_bytes: 8, - target_address_space: llvm_addr::GLOBAL, - target_addend: 0, - target_key: "b".to_string(), - }, - ]); - op.deref_mut(&ctx).attributes.set( - "global_initializer_relocations".try_into().unwrap(), - StringAttr::new(encoded), - ); - op.insert_at_back(block, &ctx); - append_mir_return(&mut ctx, block, vec![]); - - let error = crate::lower_mir_to_llvm(&mut ctx, module_ptr) - .expect_err("overlapping relocations must fail closed"); - assert!(error.to_string().contains("overlaps"), "{error}"); - } -} diff --git a/crates/mir-lower/src/convert/ops/memory/access.rs b/crates/mir-lower/src/convert/ops/memory/access.rs new file mode 100644 index 0000000000..933b249874 --- /dev/null +++ b/crates/mir-lower/src/convert/ops/memory/access.rs @@ -0,0 +1,274 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Conversions for `mir.store`, `mir.load`, `mir.alloca`, `mir.ref`, and `mir.ptr_offset`. + +use super::common::{ + anyhow_to_pliron, copy_local_memory_provenance, fail_on_target_dependent_packed_aggregate, + pointer_proved_alignment, value_abi_align, value_mir_type, +}; +use super::debug::copy_debug_local_variable; +use crate::convert::types::{convert_type, mir_type_abi_align}; +use dialect_mir::types::MirPtrType; +use llvm_export::ops as llvm; +use pliron::builtin::types::{IntegerType, Signedness}; +use pliron::context::{Context, Ptr}; +use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; +use pliron::irbuild::inserter::Inserter; +use pliron::irbuild::rewriter::Rewriter; +use pliron::location::Located; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::result::Result; +use pliron::r#type::Typed; + +/// Convert `mir.store` to `llvm.store`. +/// +/// Operand order: `[ptr, value]` - stores `value` to address `ptr`. +/// No result is produced (store is a side effect). +pub(crate) fn convert_store( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + operands_info: &OperandsInfo, +) -> Result<()> { + let operands: Vec<_> = op.deref(ctx).operands().collect(); + + let (ptr, val) = match operands.as_slice() { + [ptr, val] => (*ptr, *val), + _ => { + return pliron::input_err_noloc!("Store operation requires exactly 2 operands"); + } + }; + + // Packed whole-value stores are byte-faithful now that divergent rustc + // layouts lower to LLVM packed structs. Keep the target-dependent AS3 case + // fail-closed because its physical pointer width is selected only later. + fail_on_target_dependent_packed_aggregate( + ctx, + value_mir_type(ctx, operands_info, val), + "storing", + )?; + + let llvm_store = llvm::StoreOp::new(ctx, val, ptr); + if dialect_mir::ops::MirStoreOp::new(op).is_volatile(ctx) { + llvm_export::ops::set_op_volatile(ctx, llvm_store.get_operation(), true); + } + // The stored value's own type answers first, as it did before. A scalar + // records none, though, so fall back to whatever the address itself proved + // when it was computed -- for a field projection that is the aggregate's + // `abi_align` narrowed to the field's offset, which is otherwise lost here + // and costs the pair its vectorization. This mirrors `convert_load`, which + // consults the same record for the same reason. When both answer, the + // weaker wins: a field of a packed aggregate can place an abi-aligned type + // at a byte-aligned address, and the address's proved alignment is the + // ceiling of what the store may claim. + let abi = value_abi_align(ctx, operands_info, val); + let proved = pointer_proved_alignment(ctx, ptr); + let align = match (abi, proved) { + (Some(abi), Some(proved)) => Some(abi.min(proved)), + (abi, proved) => abi.or(proved), + }; + if let Some(align) = align { + llvm_export::ops::set_op_alignment(ctx, llvm_store.get_operation(), align as u32); + } + crate::convert::preserve_location(ctx, op, llvm_store.get_operation()); + rewriter.insert_operation(ctx, llvm_store.get_operation()); + rewriter.erase_operation(ctx, op); + Ok(()) +} + +/// Convert `mir.load` to `llvm.load`. +/// +/// Takes a single pointer operand and returns the loaded value. +/// The result type is derived from the MIR operation's result type. +pub(crate) fn convert_load( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + _operands_info: &OperandsInfo, +) -> Result<()> { + let ptr = op.deref(ctx).get_operand(0); + let result_ty = op.deref(ctx).get_result(0).get_type(ctx); + + // Packed whole-value loads are byte-faithful now that divergent rustc + // layouts lower to LLVM packed structs. Keep only the target-dependent AS3 + // physical-image case fail-closed. + fail_on_target_dependent_packed_aggregate(ctx, result_ty, "loading")?; + + let llvm_ty = convert_type(ctx, result_ty).map_err(anyhow_to_pliron)?; + + let llvm_load = llvm::LoadOp::new(ctx, ptr, llvm_ty); + if dialect_mir::ops::MirLoadOp::new(op).is_volatile(ctx) { + llvm_export::ops::set_op_volatile(ctx, llvm_load.get_operation(), true); + } + // The loaded value's ABI alignment comes from this op's own result type, + // which is still the MIR type: result types are only converted by the + // op's own rewrite. A scalar records none, so fall back to whatever the + // address itself proved when it was computed -- for a field projection + // that is the aggregate's `abi_align` narrowed to the field's offset, + // which is otherwise lost here and costs the pair its vectorization. + // When both answer, the weaker wins: a field of a packed aggregate can + // place an abi-aligned type at a byte-aligned address, and the address's + // proved alignment is the ceiling of what the load may claim. + let abi = mir_type_abi_align(ctx, result_ty); + let proved = pointer_proved_alignment(ctx, ptr); + let align = match (abi, proved) { + (Some(abi), Some(proved)) => Some(abi.min(proved)), + (abi, proved) => abi.or(proved), + }; + if let Some(align) = align { + llvm_export::ops::set_op_alignment(ctx, llvm_load.get_operation(), align as u32); + } + rewriter.insert_operation(ctx, llvm_load.get_operation()); + rewriter.replace_operation(ctx, op, llvm_load.get_operation()); + + Ok(()) +} + +/// Convert `mir.alloca` to `llvm.alloca`. +/// +/// `mir.alloca` carries its element type on the result pointer's pointee, and +/// emits a single-element stack slot of that type. We therefore convert the +/// pointee to an LLVM type and emit `llvm.alloca , i32 1`. +/// +/// No value is stored into the slot; that is the caller's job via subsequent +/// `mir.store` / `llvm.store` ops. This matches the mem2reg-ready translator +/// model where every local is backed by one alloca in the entry block and +/// defs/uses go through `store`/`load` rather than SSA values. +pub(crate) fn convert_alloca( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + _operands_info: &OperandsInfo, +) -> Result<()> { + let result_ty = op.deref(ctx).get_result(0).get_type(ctx); + let mir_pointee = { + let ty_ref = result_ty.deref(ctx); + let mir_ptr = ty_ref.downcast_ref::().ok_or_else(|| { + anyhow_to_pliron(anyhow::anyhow!( + "MirAllocaOp result must be MirPtrType (enforced by verifier)" + )) + })?; + mir_ptr.pointee + }; + let llvm_pointee = convert_type(ctx, mir_pointee).map_err(anyhow_to_pliron)?; + + let i32_ty = IntegerType::get(ctx, 32, Signedness::Signless); + let one_apint = + pliron::utils::apint::APInt::from_i64(1, std::num::NonZeroUsize::new(32).unwrap()); + let one_attr = pliron::builtin::attributes::IntegerAttr::new(i32_ty, one_apint); + let one_const = llvm::ConstantOp::new(ctx, one_attr.into()); + rewriter.insert_operation(ctx, one_const.get_operation()); + let one_val = one_const.get_operation().deref(ctx).get_result(0); + + let alloca = llvm::AllocaOp::new(ctx, llvm_pointee, one_val); + // The allocated type's ABI alignment comes from this op's own result + // pointee, which is still the MIR type at rewrite time. + if let Some(align) = mir_type_abi_align(ctx, mir_pointee) { + llvm_export::ops::set_op_alignment(ctx, alloca.get_operation(), align as u32); + } + copy_debug_local_variable(ctx, op, alloca.get_operation()); + copy_local_memory_provenance(ctx, op, alloca.get_operation()); + rewriter.insert_operation(ctx, alloca.get_operation()); + rewriter.replace_operation(ctx, op, alloca.get_operation()); + + Ok(()) +} + +/// Convert `mir.ref` — materialize the operand in stack memory via alloca+store. +/// +/// `mir.ref` creates a pointer to an SSA value. In SSA form, values don't have +/// addresses, so we must place the value in memory to obtain a pointer. +/// This applies to all types: scalars (e.g. `&factor` where factor is `u32`), +/// aggregates (e.g. `&closure_env`), and pointers (e.g. `&&T`). +pub(crate) fn convert_ref( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + operands_info: &OperandsInfo, +) -> Result<()> { + let operand = op.deref(ctx).get_operand(0); + let operand_ty = operand.get_type(ctx); + let abi_align = value_abi_align(ctx, operands_info, operand); + + let i32_ty = IntegerType::get(ctx, 32, Signedness::Signless); + let one_apint = + pliron::utils::apint::APInt::from_i64(1, std::num::NonZeroUsize::new(32).unwrap()); + let one_attr = pliron::builtin::attributes::IntegerAttr::new(i32_ty, one_apint); + let one_const = llvm::ConstantOp::new(ctx, one_attr.into()); + rewriter.insert_operation(ctx, one_const.get_operation()); + let one_val = one_const.get_operation().deref(ctx).get_result(0); + + let alloca = llvm::AllocaOp::new(ctx, operand_ty, one_val); + // Honour the referent's repr(align(N)) ABI alignment. Without this, the + // synthesised alloca would be under-aligned relative to any loads/stores + // that claim the struct's true alignment. + if let Some(align) = abi_align { + llvm_export::ops::set_op_alignment(ctx, alloca.get_operation(), align as u32); + } + rewriter.insert_operation(ctx, alloca.get_operation()); + let alloca_ptr = alloca.get_operation().deref(ctx).get_result(0); + + let store = llvm::StoreOp::new(ctx, operand, alloca_ptr); + if let Some(align) = abi_align { + llvm_export::ops::set_op_alignment(ctx, store.get_operation(), align as u32); + } + rewriter.insert_operation(ctx, store.get_operation()); + + rewriter.replace_operation_with_values(ctx, op, vec![alloca_ptr]); + + Ok(()) +} + +/// Convert `mir.ptr_offset` to `llvm.getelementptr`. +/// +/// Operands: `[ptr, offset]` where offset is an integer index. +/// Element sizing comes from the op's own result type, which is still the +/// MIR pointer type when this converter runs. The operand's recorded type +/// history is not usable here: a kind-only `mir.cast` lowers to a plain +/// value forwarding, and the history does not follow that replacement +/// edge, so a history miss would silently misscale the offset. +pub(crate) fn convert_ptr_offset( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + _operands_info: &OperandsInfo, +) -> Result<()> { + let loc = op.deref(ctx).loc(); + let operands: Vec<_> = op.deref(ctx).operands().collect(); + + let (ptr, offset) = match operands.as_slice() { + [ptr, offset] => (*ptr, *offset), + _ => return pliron::input_err!(loc, "PtrOffset requires exactly 2 operands"), + }; + + let result_ty = op.deref(ctx).get_result(0).get_type(ctx); + let pointee = result_ty + .deref(ctx) + .downcast_ref::() + .map(|mir_ptr| mir_ptr.pointee) + .ok_or_else(|| { + pliron::input_error!( + loc.clone(), + "mir.ptr_offset result must be a MIR pointer type; \ + element sizing has no fact to derive from" + ) + })?; + let elem_ty = convert_type(ctx, pointee).map_err(anyhow_to_pliron)?; + + let llvm_gep = llvm::GetElementPtrOp::new( + ctx, + ptr, + vec![llvm_export::ops::GepIndex::Value(offset)], + elem_ty, + ); + let inbounds = dialect_mir::ops::MirPtrOffsetOp::new(op).is_inbounds(ctx); + llvm::set_gep_inbounds(ctx, llvm_gep.get_operation(), inbounds); + rewriter.insert_operation(ctx, llvm_gep.get_operation()); + rewriter.replace_operation(ctx, op, llvm_gep.get_operation()); + + Ok(()) +} diff --git a/crates/mir-lower/src/convert/ops/memory/common.rs b/crates/mir-lower/src/convert/ops/memory/common.rs new file mode 100644 index 0000000000..14bba3d328 --- /dev/null +++ b/crates/mir-lower/src/convert/ops/memory/common.rs @@ -0,0 +1,158 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Helpers shared by the memory-op conversions in this module. + +use crate::convert::types::{ + StructLayoutInfo, build_struct_slot_map, convert_type, + llvm_packed_struct_contains_pointer_in_address_space, mir_type_abi_align, +}; +use dialect_mir::types::MirStructType; +use pliron::context::{Context, Ptr}; +use pliron::irbuild::dialect_conversion::OperandsInfo; +use pliron::operation::Operation; +use pliron::result::Result; +use pliron::r#type::{TypeHandle, Typed}; +use pliron::value::Value; + +pub(super) fn anyhow_to_pliron(e: anyhow::Error) -> pliron::result::Error { + pliron::create_error!( + pliron::location::Location::Unknown, + pliron::result::ErrorKind::VerificationFailed, + pliron::result::StringError(e.to_string()) + ) +} + +/// Recover the `repr(align(N))` ABI alignment of `value` at conversion time. +/// +/// The alignment lives on MIR aggregate types (`abi_align`); the converted +/// LLVM struct types cannot express over-alignment. The conversion driver may +/// already have converted the value's type (block arguments are converted +/// before any rewrite runs; replaced op results carry the new type), but it +/// records every such change. So check the current type first, then walk the +/// value's conversion history, newest first, for a MIR type that records an +/// alignment. +pub(super) fn value_abi_align( + ctx: &Context, + operands_info: &OperandsInfo, + value: Value, +) -> Option { + mir_type_abi_align(ctx, value.get_type(ctx)).or_else(|| { + operands_info + .lookup_operand_history(value) + .iter() + .rev() + .find_map(|ty| mir_type_abi_align(ctx, *ty)) + }) +} + +/// The MIR type of `value`: its current type when that is still a MIR type, +/// else the newest MIR type in its conversion history (same walk as +/// [`value_abi_align`], for the same reason). +pub(super) fn value_mir_type( + ctx: &Context, + operands_info: &OperandsInfo, + value: Value, +) -> TypeHandle { + let current = value.get_type(ctx); + if current.deref(ctx).is::() { + return current; + } + operands_info + .lookup_operand_history(value) + .iter() + .rev() + .copied() + .find(|ty| ty.deref(ctx).is::()) + .unwrap_or(current) +} + +/// Refuse only packed by-value images whose physical bytes depend on the +/// selected NVPTX mode. +/// +/// Shared-memory pointers are 32-bit in modern NVVM p3:32 but 64-bit in the +/// PTX/legacy layouts. Lowering runs before that target mode is selected, so a +/// packed aggregate containing AS3 cannot safely be loaded or stored as one +/// physical value. Pointer-free packed aggregates and packed aggregates whose +/// pointers use target-stable address spaces are allowed. +pub(super) fn fail_on_target_dependent_packed_aggregate( + ctx: &mut Context, + ty: TypeHandle, + verb: &str, +) -> Result<()> { + let llvm_ty = { + let layout = { + let ty_ref = ty.deref(ctx); + ty_ref + .downcast_ref::() + .map(StructLayoutInfo::of_struct) + }; + if let Some(layout) = layout { + let map = build_struct_slot_map(ctx, &layout).map_err(anyhow_to_pliron)?; + if !map.by_value_layout_faithful { + return pliron::input_err_noloc!( + "{} a struct whose rustc layout cannot be represented by an LLVM \ + struct value is not supported", + verb + ); + } + map.llvm_struct_ty + } else { + convert_type(ctx, ty).map_err(anyhow_to_pliron)? + } + }; + if llvm_packed_struct_contains_pointer_in_address_space( + ctx, + llvm_ty, + llvm_export::types::address_space::SHARED, + ) { + return pliron::input_err_noloc!( + "{} a packed aggregate containing a shared-memory pointer by value is \ + target-mode dependent and is not yet supported", + verb + ); + } + Ok(()) +} + +/// Carry Rust-local provenance from the MIR alloca to the LLVM alloca. +/// +/// This is compiler-only metadata stored in the in-memory dialect operation; +/// it is not emitted as LLVM metadata because an unrecognized metadata node is +/// not a preservation contract for middle-end passes. The textual exporter +/// instead consumes the attribute to choose a stable SSA name before `opt`. +pub(super) fn copy_local_memory_provenance( + ctx: &mut Context, + mir_op: Ptr, + llvm_op: Ptr, +) { + if let Some(provenance) = llvm_export::ops::local_memory_provenance(ctx, mir_op) { + llvm_export::ops::set_local_memory_provenance(ctx, llvm_op, provenance); + } +} + +/// Alignment the op that computed `ptr` recorded about the address it produced. +/// +/// Field addresses stamp this during their own conversion, where the aggregate's +/// `abi_align` is still in hand; by the time a load that consumes the address is +/// converted, `mir.field_addr` is already a GEP and the MIR aggregate type is no +/// longer reachable from the load's operands. +/// +/// Known fidelity limit: the claim rides the address value, not the place +/// access that justified it. `&raw const (*p).x` through a merely 4-aligned +/// `p: *const Aligned8` is legal Rust so long as the place itself is never +/// accessed, yet the later read through the raw pointer consumes this same +/// stamped GEP and inherits the aggregate's alignment, an over-claim for that +/// pointer. rustc keeps the two apart by threading alignment through place +/// evaluation and dropping it at the `&raw` boundary; dialect-mir erases that +/// distinction before lowering runs, so full fidelity would require stamping +/// loads during MIR translation instead. Accepted for now: the pattern needs +/// an actually under-aligned pointer to an over-aligned aggregate, projected +/// via `&raw` inside a kernel, and any direct access through such a place is +/// UB to begin with. +pub(super) fn pointer_proved_alignment(ctx: &Context, ptr: Value) -> Option { + let defining_op = ptr.defining_op()?; + llvm_export::ops::address_alignment(ctx, defining_op).map(u64::from) +} diff --git a/crates/mir-lower/src/convert/ops/memory/debug.rs b/crates/mir-lower/src/convert/ops/memory/debug.rs new file mode 100644 index 0000000000..91e9e95b67 --- /dev/null +++ b/crates/mir-lower/src/convert/ops/memory/debug.rs @@ -0,0 +1,92 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Conversions for the debug-value markers and debug-local metadata copying. + +use llvm_export::ops as llvm; +use pliron::context::{Context, Ptr}; +use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; +use pliron::irbuild::inserter::Inserter; +use pliron::irbuild::rewriter::Rewriter; +use pliron::location::Located; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::result::Result; + +pub(super) fn copy_debug_local_variable( + ctx: &mut Context, + mir_op: Ptr, + llvm_op: Ptr, +) { + if let Some(info) = llvm_export::ops::debug_local_variable(ctx, mir_op) { + llvm_export::ops::set_debug_local_variable(ctx, llvm_op, info); + } + let projected = llvm_export::ops::debug_projected_variables(ctx, mir_op); + if !projected.is_empty() { + llvm_export::ops::set_debug_projected_variables(ctx, llvm_op, &projected); + } + let fragments = llvm_export::ops::debug_fragment_variables(ctx, mir_op); + if !fragments.is_empty() { + llvm_export::ops::set_debug_fragment_variables(ctx, llvm_op, &fragments); + } + if let Some(scope) = llvm_export::ops::debug_local_source_scope(ctx, mir_op) { + llvm_export::ops::set_debug_local_source_scope(ctx, llvm_op, scope); + } + if let Some((file, pos)) = llvm_export::ops::debug_local_declaration_location(ctx, mir_op) { + llvm_export::ops::set_debug_local_declaration_location( + ctx, llvm_op, file, pos.line, pos.column, + ); + } + if let Some(expression) = llvm_export::ops::debug_value_expression(ctx, mir_op) { + llvm_export::ops::set_debug_value_expression(ctx, llvm_op, &expression); + } +} + +/// Convert `mir.dbg_value` to the LLVM-export debug marker. +/// +/// The op is still debug-only after lowering. The textual LLVM exporter later +/// prints it as an `llvm.dbg.value` intrinsic call. +pub(crate) fn convert_dbg_value( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + _operands_info: &OperandsInfo, +) -> Result<()> { + let value = op.deref(ctx).get_operand(0); + let loc = op.deref(ctx).loc().clone(); + let llvm_dbg_value = llvm::DebugValueOp::new(ctx, value); + llvm_dbg_value.get_operation().deref_mut(ctx).set_loc(loc); + copy_debug_local_variable(ctx, op, llvm_dbg_value.get_operation()); + rewriter.insert_operation(ctx, llvm_dbg_value.get_operation()); + rewriter.erase_operation(ctx, op); + Ok(()) +} + +/// Convert `mir.dbg_value_list` to the LLVM-export multi-value debug marker. +/// +/// The ordered operands become a `DIArgList` during textual export. The typed +/// location recipe is carried as generic metadata and copied unchanged here. +pub(crate) fn convert_dbg_value_list( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + _operands_info: &OperandsInfo, +) -> Result<()> { + let values: Vec<_> = op.deref(ctx).operands().collect(); + if values.len() < 2 { + return pliron::input_err_noloc!("mir.dbg_value_list requires at least two operands"); + } + if llvm_export::ops::debug_value_expression(ctx, op).is_none() { + return pliron::input_err_noloc!("mir.dbg_value_list is missing its debug expression"); + } + + let loc = op.deref(ctx).loc().clone(); + let llvm_dbg_value = llvm::DebugValueListOp::new(ctx, values); + llvm_dbg_value.get_operation().deref_mut(ctx).set_loc(loc); + copy_debug_local_variable(ctx, op, llvm_dbg_value.get_operation()); + rewriter.insert_operation(ctx, llvm_dbg_value.get_operation()); + rewriter.erase_operation(ctx, op); + Ok(()) +} diff --git a/crates/mir-lower/src/convert/ops/memory/device_global.rs b/crates/mir-lower/src/convert/ops/memory/device_global.rs new file mode 100644 index 0000000000..dd591c88a1 --- /dev/null +++ b/crates/mir-lower/src/convert/ops/memory/device_global.rs @@ -0,0 +1,410 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Conversion of `mir.global_alloc` to device globals, including relocated initializers. + +use super::common::anyhow_to_pliron; +use crate::context::{DeviceGlobalDeclaration, DeviceGlobalRecord, DeviceGlobalsMap}; +use crate::convert::types::{ + convert_type, llvm_type_size_align, validate_initialized_global_layout, + validate_relocated_initialized_global_layout, +}; +use crate::helpers; +use llvm_export::ops as llvm; +use llvm_export::ops::GlobalOpExt; +use llvm_export::types::{ArrayType, StructLayout, StructType}; +use pliron::builtin::types::{IntegerType, Signedness}; +use pliron::context::{Context, Ptr}; +use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; +use pliron::irbuild::inserter::Inserter; +use pliron::irbuild::rewriter::Rewriter; +use pliron::linked_list::ContainsLinkedList; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::result::Result; +use pliron::r#type::{TypeHandle, Typed}; + +/// Convert `mir.global_alloc` to an LLVM global in CUDA global memory. +/// +/// Ordinary Rust `static` / `static mut` values have grid scope and +/// application lifetime, so they live in address space 1 rather than the +/// per-block shared-memory address space. +pub fn convert_global_alloc_dc( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + _operands_info: &OperandsInfo, + device_globals: &mut DeviceGlobalsMap, + next_device_global_index: &mut usize, +) -> Result<()> { + use pliron::builtin::attributes::{StringAttr, TypeAttr}; + + let ( + global_key, + mir_global_type, + alignment, + addr_space, + initializer_hex, + initializer_relocations, + immutable, + debug_info, + ) = { + let global_op = dialect_mir::ops::MirGlobalAllocOp::new(op); + let op_ref = op.deref(ctx); + + let global_key_attr = op_ref + .attributes + .get::(&"global_key".try_into().unwrap()) + .ok_or_else(|| { + anyhow_to_pliron(anyhow::anyhow!( + "MirGlobalAllocOp missing global_key StringAttr attribute" + )) + })?; + let global_key = String::from((*global_key_attr).clone()); + + let global_type_attr = op_ref + .attributes + .get::(&"global_type".try_into().unwrap()) + .ok_or_else(|| { + anyhow_to_pliron(anyhow::anyhow!( + "MirGlobalAllocOp missing global_type TypeAttr attribute" + )) + })?; + let mir_global_type = global_type_attr.get_type(ctx); + + let alignment = global_op.get_alignment_value(ctx).unwrap_or(0); + let initializer_hex = op_ref + .attributes + .get::(&"global_initializer_hex".try_into().unwrap()) + .map(|attr| String::from((*attr).clone())); + let initializer_relocations = op_ref + .attributes + .get::(&"global_initializer_relocations".try_into().unwrap()) + .map(|attr| String::from((*attr).clone())); + let debug_info = llvm::debug_global_variable(ctx, op); + + // Read the address space the op's result already carries — set by + // mir-importer based on the static's type (`ConstantMemory` → 4, + // ordinary → 1). The dialect verifier accepts both. + let res_ty = op_ref.get_result(0).get_type(ctx); + let addr_space = res_ty + .deref(ctx) + .downcast_ref::() + .map(|p| { + if p.address_space == dialect_mir::types::address_space::CONSTANT { + llvm_export::types::address_space::CONSTANT + } else { + llvm_export::types::address_space::GLOBAL + } + }) + .ok_or_else(|| { + anyhow_to_pliron(anyhow::anyhow!( + "MirGlobalAllocOp result is not a MirPtrType" + )) + })?; + + ( + global_key, + mir_global_type, + alignment, + addr_space, + initializer_hex, + initializer_relocations, + global_op.is_immutable(ctx), + debug_info, + ) + }; + + let declaration = DeviceGlobalDeclaration { + mir_type: mir_global_type, + alignment, + addr_space, + initializer_hex: initializer_hex.clone(), + initializer_relocations: initializer_relocations.clone(), + debug_info: debug_info.clone(), + immutable, + }; + + let global_name = if let Some(existing) = device_globals.get(&global_key) { + if existing.declaration != declaration { + return Err(anyhow_to_pliron(anyhow::anyhow!( + "duplicate global_key {:?} has an incompatible declaration", + global_key + ))); + } + existing.symbol.clone() + } else { + create_device_global( + ctx, + op, + device_globals, + next_device_global_index, + DeviceGlobalSpec { + key: &global_key, + mir_type: mir_global_type, + alignment, + addr_space, + initializer_hex: initializer_hex.as_deref(), + initializer_relocations: initializer_relocations.as_deref(), + debug_info: debug_info.as_ref(), + immutable, + }, + )? + }; + + let address_of_op = llvm::AddressOfOp::new(ctx, global_name, addr_space); + rewriter.insert_operation(ctx, address_of_op.get_operation()); + rewriter.replace_operation(ctx, op, address_of_op.get_operation()); + + Ok(()) +} + +struct DeviceGlobalSpec<'a> { + key: &'a str, + mir_type: TypeHandle, + alignment: u64, + addr_space: u32, + initializer_hex: Option<&'a str>, + initializer_relocations: Option<&'a str>, + debug_info: Option<&'a llvm::DebugGlobalVariableInfo>, + /// Nothing writes this storage, so it exports as LLVM `constant`. Set only + /// for the compiler's own promoted constants; see `MirGlobalAllocOp`. + immutable: bool, +} + +/// `next_device_global_index` is scoped to one `MirToLlvmConversionDriver` +/// instance (one module), not a process-global counter: `N` is a function of +/// this module's own MIR walk order, not of how many other modules have +/// lowered a device global earlier in the process (#706). +fn create_device_global( + ctx: &mut Context, + op: Ptr, + device_globals: &mut DeviceGlobalsMap, + next_device_global_index: &mut usize, + spec: DeviceGlobalSpec<'_>, +) -> Result { + // An explicit initializer is already the evaluated Rust allocation image. + // Pointer-free data stays `[N x i8]`. Initializers with relocations use a + // segmented LLVM struct whose literal spans remain byte arrays and whose + // pointer slots become pointer-width integers. This preserves both exact + // bytes and linker-visible pointer provenance. + let semantic_llvm_type = convert_type(ctx, spec.mir_type).map_err(anyhow_to_pliron)?; + let (llvm_global_type, alignment) = if let Some(initializer_hex) = spec.initializer_hex { + let byte_count = initializer_hex_byte_count(initializer_hex).map_err(anyhow_to_pliron)?; + if spec.alignment == 0 { + return Err(anyhow_to_pliron(anyhow::anyhow!( + "device global initializer is missing its evaluated Rust allocation alignment" + ))); + } + let storage_type = if let Some(encoded) = spec.initializer_relocations { + validate_relocated_initialized_global_layout( + ctx, + spec.mir_type, + byte_count, + spec.alignment, + ) + .map_err(anyhow_to_pliron)?; + relocated_initializer_storage_type(ctx, byte_count, spec.alignment, encoded) + .map_err(anyhow_to_pliron)? + } else { + validate_initialized_global_layout(ctx, spec.mir_type, byte_count, spec.alignment) + .map_err(anyhow_to_pliron)?; + let i8_ty = IntegerType::get(ctx, 8, Signedness::Signless); + ArrayType::get(ctx, i8_ty.into(), byte_count).into() + }; + (storage_type, spec.alignment) + } else { + if spec.initializer_relocations.is_some() { + return Err(anyhow_to_pliron(anyhow::anyhow!( + "device global carries relocation metadata without initializer bytes" + ))); + } + (semantic_llvm_type, spec.alignment) + }; + + // Constant-memory globals reuse the Rust-side mangled name so host code can + // resolve them by name via `cuModuleGetGlobal`. Ordinary device globals + // are private to the kernel and get a counter-based unique name. + let name: pliron::identifier::Identifier = + if spec.addr_space == llvm_export::types::address_space::CONSTANT { + let symbol = dialect_mir::ops::rust_static_symbol_from_global_key(spec.key) + .ok_or_else(|| { + anyhow_to_pliron(anyhow::anyhow!( + "constant global_key {:?} is not a tagged Rust static identity", + spec.key + )) + })?; + symbol.try_into().map_err(|e| { + anyhow_to_pliron(anyhow::anyhow!( + "constant Rust static symbol {:?} is not a valid identifier: {e:?}", + symbol + )) + })? + } else { + let counter = *next_device_global_index; + *next_device_global_index += 1; + format!("__device_global_{counter}").try_into().unwrap() + }; + + let global_op = if alignment > 0 { + llvm::GlobalOp::new_with_alignment(ctx, name.clone(), llvm_global_type, alignment) + } else { + llvm::GlobalOp::new(ctx, name.clone(), llvm_global_type) + }; + global_op.set_address_space(ctx, spec.addr_space); + global_op.set_source_global_key(ctx, spec.key); + if spec.addr_space == llvm_export::types::address_space::GLOBAL + && let Some(info) = spec.debug_info + { + llvm::set_debug_global_variable(ctx, global_op.get_operation(), info); + } + if let Some(initializer_hex) = spec.initializer_hex { + global_op.set_initializer_hex(ctx, initializer_hex); + } + if let Some(initializer_relocations) = spec.initializer_relocations { + global_op.set_initializer_relocations(ctx, initializer_relocations); + } + if spec.immutable { + global_op.mark_immutable(ctx); + } + + let parent_block = op + .deref(ctx) + .get_parent_block() + .ok_or_else(|| anyhow_to_pliron(anyhow::anyhow!("Op has no parent block")))?; + let module_op = helpers::get_module_from_block(ctx, parent_block).map_err(anyhow_to_pliron)?; + let region = module_op.deref(ctx).get_region(0); + let module_block = region + .deref(ctx) + .iter(ctx) + .next() + .ok_or_else(|| anyhow_to_pliron(anyhow::anyhow!("Module is empty")))?; + + global_op.get_operation().insert_at_front(module_block, ctx); + device_globals.insert( + spec.key.to_string(), + DeviceGlobalRecord { + symbol: name.clone(), + declaration: DeviceGlobalDeclaration { + mir_type: spec.mir_type, + alignment: spec.alignment, + addr_space: spec.addr_space, + initializer_hex: spec.initializer_hex.map(str::to_owned), + initializer_relocations: spec.initializer_relocations.map(str::to_owned), + debug_info: spec.debug_info.cloned(), + immutable: spec.immutable, + }, + }, + ); + + Ok(name) +} + +pub(super) fn relocated_initializer_storage_type( + ctx: &mut Context, + byte_count: u64, + allocation_alignment: u64, + encoded: &str, +) -> std::result::Result { + let mut relocations = + llvm::decode_global_initializer_relocations(encoded).map_err(anyhow::Error::msg)?; + if relocations.is_empty() { + anyhow::bail!("device global relocation metadata contains no relocations"); + } + relocations.sort_by_key(|relocation| relocation.source_offset); + + let mut cursor = 0u64; + let mut fields = Vec::with_capacity(relocations.len() * 2 + 1); + let mut requires_packed_storage = false; + let i8_ty = IntegerType::get(ctx, 8, Signedness::Signless); + + for (index, relocation) in relocations.iter().enumerate() { + if relocation.width_bytes != 8 { + anyhow::bail!( + "device global relocation {index} uses unsupported {}-byte pointer storage; CUDA global/constant pointers require 8 bytes", + relocation.width_bytes + ); + } + if !matches!(relocation.target_address_space, 1 | 4) { + anyhow::bail!( + "device global relocation {index} targets unsupported CUDA address space {}", + relocation.target_address_space + ); + } + if relocation.target_key.is_empty() { + anyhow::bail!("device global relocation {index} has an empty target key"); + } + + let width = u64::from(relocation.width_bytes); + requires_packed_storage |= + allocation_alignment < width || !relocation.source_offset.is_multiple_of(width); + if relocation.source_offset < cursor { + anyhow::bail!( + "device global relocation {index} overlaps the previous relocation or literal span" + ); + } + let end = relocation + .source_offset + .checked_add(width) + .ok_or_else(|| anyhow::anyhow!("device global relocation {index} offset overflows"))?; + if end > byte_count { + anyhow::bail!( + "device global relocation {index} occupies bytes {}..{} but the initializer is only {} bytes", + relocation.source_offset, + end, + byte_count + ); + } + + if relocation.source_offset > cursor { + fields + .push(ArrayType::get(ctx, i8_ty.into(), relocation.source_offset - cursor).into()); + } + fields.push(IntegerType::get(ctx, relocation.width_bytes * 8, Signedness::Signless).into()); + cursor = end; + } + + if cursor < byte_count { + fields.push(ArrayType::get(ctx, i8_ty.into(), byte_count - cursor).into()); + } + + let layout = if requires_packed_storage { + StructLayout::Packed + } else { + StructLayout::Unpacked + }; + let storage: TypeHandle = StructType::get_unnamed(ctx, (fields, layout)).into(); + // Exact or error, never guessed: the storage type is built from i8 arrays + // and integer slots, so its natural size is always computable, and it must + // land exactly on rustc's byte count or the relocation offsets are wrong. + let Some((lowered_size, _)) = llvm_type_size_align(ctx, storage) else { + anyhow::bail!( + "relocated device global storage `{}` has no exact size", + storage.deref(ctx).disp(ctx) + ); + }; + if lowered_size != byte_count { + anyhow::bail!( + "relocated device global storage lowers to {} bytes, but rustc evaluated {} bytes", + lowered_size, + byte_count + ); + } + Ok(storage) +} + +fn initializer_hex_byte_count(hex: &str) -> std::result::Result { + if !hex.len().is_multiple_of(2) { + anyhow::bail!("device global initializer has an odd-length hex byte string"); + } + if let Some(invalid) = hex.bytes().find(|byte| !byte.is_ascii_hexdigit()) { + anyhow::bail!( + "device global initializer contains invalid hex digit {:?}", + invalid as char + ); + } + u64::try_from(hex.len() / 2) + .map_err(|_| anyhow::anyhow!("device global initializer is too large for LLVM")) +} diff --git a/crates/mir-lower/src/convert/ops/memory/extern_shared.rs b/crates/mir-lower/src/convert/ops/memory/extern_shared.rs new file mode 100644 index 0000000000..4db8ff6aa6 --- /dev/null +++ b/crates/mir-lower/src/convert/ops/memory/extern_shared.rs @@ -0,0 +1,209 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Conversion of `mir.extern_shared` to dynamic shared-memory extern globals. + +use super::common::anyhow_to_pliron; +use crate::context::{ + DynamicSmemAlignmentMap, SharedGlobalDeclaration, SharedGlobalKind, SharedGlobalRecord, + SharedGlobalsMap, +}; +use crate::helpers; +use llvm_export::ops as llvm; +use llvm_export::ops::GlobalOpExt; +use llvm_export::types::ArrayType; +use pliron::builtin::op_interfaces::SymbolOpInterface; +use pliron::builtin::types::{IntegerType, Signedness}; +use pliron::context::{Context, Ptr}; +use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; +use pliron::irbuild::inserter::Inserter; +use pliron::irbuild::rewriter::Rewriter; +use pliron::linked_list::ContainsLinkedList; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::result::Result; + +/// Convert `mir.extern_shared` to LLVM extern global variable in shared address space. +/// +/// Dynamic (extern) shared memory is represented as an external global variable +/// with address space 3 and zero-length array type `[0 x i8]`. The actual size +/// is determined at kernel launch via `LaunchConfig::shared_mem_bytes`. +/// +/// # Per-Owner Symbols +/// +/// Each function that owns an access gets a dynamic shared-memory symbol +/// (`__dynamic_smem_{function_name}`). +/// +/// # Alignment +/// +/// The alignment is pre-computed during the lowering pre-pass. It is the +/// maximum of the owner's body requirements and every launch-contract marker +/// that can reach it. +/// +/// # Byte Offset +/// +/// - `DynamicSharedArray::get()` / `get_raw()`: offset = 0, returns base pointer +/// - `DynamicSharedArray::offset(N)`: offset = N bytes, returns base + N +/// +/// Called directly from `MirToLlvmConversionDriver::rewrite` (not through +/// op_cast dispatch) because it needs cross-function state maps. +pub fn convert_extern_shared_dc( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + _operands_info: &OperandsInfo, + shared_globals: &mut SharedGlobalsMap, + dynamic_smem_alignments: &mut DynamicSmemAlignmentMap, +) -> Result<()> { + let (byte_offset, alignment) = { + let extern_shared_op = dialect_mir::ops::MirExternSharedOp::new(op); + let byte_offset = extern_shared_op.get_byte_offset_value(ctx); + let alignment = extern_shared_op.get_alignment_value(ctx); + (byte_offset, alignment) + }; + + let func_name: String = { + let parent_block = op + .deref(ctx) + .get_parent_block() + .ok_or_else(|| anyhow_to_pliron(anyhow::anyhow!("Op has no parent block")))?; + let func_op_ptr = helpers::get_parent_func(ctx, parent_block).map_err(anyhow_to_pliron)?; + let llvm_func = Operation::get_op::(func_op_ptr, ctx) + .ok_or_else(|| anyhow_to_pliron(anyhow::anyhow!("Parent op is not an llvm.func")))?; + llvm_func.get_symbol_name(ctx).to_string() + }; + + let global_name = get_or_create_extern_shared_global( + ctx, + op, + &func_name, + shared_globals, + dynamic_smem_alignments, + alignment, + )?; + + let address_of_op = llvm::AddressOfOp::new(ctx, global_name, 3); + rewriter.insert_operation(ctx, address_of_op.get_operation()); + + let base_ptr = address_of_op.get_operation().deref(ctx).get_result(0); + + if byte_offset > 0 { + let i64_ty = IntegerType::get(ctx, 64, Signedness::Signless); + let offset_attr = pliron::builtin::attributes::IntegerAttr::new( + i64_ty, + pliron::utils::apint::APInt::from_u64( + byte_offset, + std::num::NonZeroUsize::new(64).unwrap(), + ), + ); + let offset_const = llvm::ConstantOp::new(ctx, offset_attr.into()); + rewriter.insert_operation(ctx, offset_const.get_operation()); + let offset_value = offset_const.get_operation().deref(ctx).get_result(0); + + let i8_ty = IntegerType::get(ctx, 8, Signedness::Signless); + let gep_op = llvm::GetElementPtrOp::new( + ctx, + base_ptr, + vec![llvm_export::ops::GepIndex::Value(offset_value)], + i8_ty.into(), + ); + rewriter.insert_operation(ctx, gep_op.get_operation()); + rewriter.replace_operation(ctx, op, gep_op.get_operation()); + } else { + rewriter.replace_operation(ctx, op, address_of_op.get_operation()); + } + + Ok(()) +} + +/// Get or create the extern shared memory global for an owning function. +/// +/// Creates an LLVM global variable with: +/// - Zero-length array type: `[0 x i8]` +/// - External linkage (no initializer) +/// - Address space 3 (shared memory) +/// - Pre-computed body and calling-kernel contract alignment +/// +/// Each owning function gets its own dynamic shared memory symbol. Uses +/// `shared_globals` for deduplication (only one global per function). +fn get_or_create_extern_shared_global( + ctx: &mut Context, + op: Ptr, + func_name: &str, + shared_globals: &mut SharedGlobalsMap, + dynamic_smem_alignments: &mut DynamicSmemAlignmentMap, + _requested_alignment: u64, +) -> Result { + let (symbol_name, max_alignment) = dynamic_smem_alignments.get(func_name).cloned().ok_or_else( + || { + anyhow_to_pliron(anyhow::anyhow!( + "Internal error: dynamic shared memory alignment not pre-computed for function '{}'. \ + This should have been done in compute_max_dynamic_smem_alignment.", + func_name + )) + }, + )?; + + let i8_ty = IntegerType::get(ctx, 8, Signedness::Signless); + let declaration = SharedGlobalDeclaration { + kind: SharedGlobalKind::DynamicExtern, + mir_elem_type: i8_ty.into(), + size: 0, + alignment: max_alignment, + // The dynamic pool is per-launch storage with no originating static; + // it never carries a source-level debug identity. + debug_info: None, + debug_owner_function: None, + }; + let global_created_key = format!("__dynamic_smem_global_created_{}", func_name); + if let Some(existing) = shared_globals.get(&global_created_key) { + if existing.declaration != declaration || existing.symbol != symbol_name { + return Err(anyhow_to_pliron(anyhow::anyhow!( + "dynamic shared-memory key {:?} has an incompatible declaration", + global_created_key + ))); + } + return Ok(existing.symbol.clone()); + } + + let array_type = ArrayType::get(ctx, i8_ty.into(), 0); + + let global_op = llvm::GlobalOp::new_with_alignment( + ctx, + symbol_name.clone(), + array_type.into(), + max_alignment, + ); + global_op.set_address_space(ctx, llvm_export::types::address_space::SHARED); + + { + use llvm_export::attributes::LinkageAttr; + global_op.set_attr_llvm_global_linkage(ctx, LinkageAttr::ExternalLinkage); + } + + let parent_block = op + .deref(ctx) + .get_parent_block() + .ok_or_else(|| anyhow_to_pliron(anyhow::anyhow!("Op has no parent block")))?; + let module_op = helpers::get_module_from_block(ctx, parent_block).map_err(anyhow_to_pliron)?; + let region = module_op.deref(ctx).get_region(0); + let module_block = region + .deref(ctx) + .iter(ctx) + .next() + .ok_or_else(|| anyhow_to_pliron(anyhow::anyhow!("Module is empty")))?; + + global_op.get_operation().insert_at_front(module_block, ctx); + + shared_globals.insert( + global_created_key, + SharedGlobalRecord { + symbol: symbol_name.clone(), + declaration, + }, + ); + + Ok(symbol_name) +} diff --git a/crates/mir-lower/src/convert/ops/memory/mod.rs b/crates/mir-lower/src/convert/ops/memory/mod.rs new file mode 100644 index 0000000000..99f085cce3 --- /dev/null +++ b/crates/mir-lower/src/convert/ops/memory/mod.rs @@ -0,0 +1,69 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Memory operation conversion: `dialect-mir` → LLVM dialect. +//! +//! Converts `dialect-mir` memory operations to their LLVM dialect equivalents. +//! +//! # Operations +//! +//! | MIR Operation | LLVM Operation(s) | Description | +//! |----------------------|-----------------------------------|------------------------------| +//! | `mir.load` | `llvm.load` | Load from pointer | +//! | `mir.store` | `llvm.store` | Store to pointer | +//! | `mir.ref` | `llvm.alloca` + `llvm.store` | Materialize aggregate in mem | +//! | `mir.ptr_offset` | `llvm.getelementptr` | Pointer arithmetic | +//! | `mir.shared_alloc` | `llvm.global` + `llvm.addressof` | Static shared memory | +//! | `mir.extern_shared` | `llvm.global` + `llvm.addressof` | Dynamic shared memory | +//! +//! # Shared Memory +//! +//! ## Static Shared Memory (`SharedArray`) +//! +//! Each static shared memory allocation gets a unique global symbol (`__shared_mem_N`). +//! Multiple allocations in the same or different kernels each have their own symbol +//! with their own size and alignment. +//! +//! ## Dynamic Shared Memory (`DynamicSharedArray`) +//! +//! Dynamic shared memory uses a symbol for each function that owns an access +//! (`__dynamic_smem_{function_name}`). +//! Key characteristics: +//! +//! - **Per-owner symbols**: Each function containing an access gets an extern symbol +//! - **Pre-computed alignment**: A pre-pass combines the owner's body alignment with +//! the strongest launch-contract marker that can reach it +//! - **Single runtime pool per launch**: The symbols refer to dynamic shared memory +//! sized by `shared_mem_bytes` at launch +//! +//! ### PTX Output Example +//! +//! ```ptx +//! ; Kernel with 128-byte aligned dynamic shared memory +//! .extern .shared .align 128 .b8 __dynamic_smem_my_kernel[]; +//! +//! ; Another kernel with 16-byte aligned (default) +//! .extern .shared .align 16 .b8 __dynamic_smem_other_kernel[]; +//! ``` + +mod access; +mod common; +mod debug; +mod device_global; +mod extern_shared; +mod shared; +mod transfer; + +pub(crate) use access::{ + convert_alloca, convert_load, convert_ptr_offset, convert_ref, convert_store, +}; +pub(crate) use debug::{convert_dbg_value, convert_dbg_value_list}; +pub use device_global::convert_global_alloc_dc; +pub use extern_shared::convert_extern_shared_dc; +pub use shared::convert_shared_alloc_dc; +pub(crate) use transfer::{convert_memcpy, convert_memmove}; + +#[cfg(test)] +mod tests; diff --git a/crates/mir-lower/src/convert/ops/memory/shared.rs b/crates/mir-lower/src/convert/ops/memory/shared.rs new file mode 100644 index 0000000000..05d63840da --- /dev/null +++ b/crates/mir-lower/src/convert/ops/memory/shared.rs @@ -0,0 +1,342 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Conversion of `mir.shared_alloc` to static shared-memory globals. + +use super::common::anyhow_to_pliron; +use crate::context::{ + SharedGlobalDeclaration, SharedGlobalKind, SharedGlobalRecord, SharedGlobalsMap, +}; +use crate::convert::types::{convert_type, llvm_type_size_align}; +use crate::helpers; +use llvm_export::ops as llvm; +use llvm_export::ops::GlobalOpExt; +use llvm_export::types::ArrayType; +use pliron::builtin::op_interfaces::SymbolOpInterface; +use pliron::context::{Context, Ptr}; +use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; +use pliron::irbuild::inserter::Inserter; +use pliron::irbuild::rewriter::Rewriter; +use pliron::linked_list::ContainsLinkedList; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::result::Result; +use pliron::r#type::{TypeHandle, Typed}; + +/// Convert `mir.shared_alloc` to LLVM global variable in shared address space. +/// +/// GPU shared memory is represented as a global variable with address space 3. +/// Uses `shared_globals` to deduplicate: multiple allocations with the same +/// `alloc_key` share the same global. +/// +/// Called directly from `MirToLlvmConversionDriver::rewrite` (not through +/// op_cast dispatch) because it needs the cross-function `SharedGlobalsMap`. +pub fn convert_shared_alloc_dc( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + _operands_info: &OperandsInfo, + shared_globals: &mut SharedGlobalsMap, + next_shared_mem_index: &mut usize, +) -> Result<()> { + use pliron::builtin::attributes::{IntegerAttr, TypeAttr}; + + let (alloc_key, source_name, mir_elem_type, size, alignment, debug_info, debug_owner_function) = { + let shared_alloc_op = dialect_mir::ops::MirSharedAllocOp::new(op); + let op_ref = op.deref(ctx); + + let alloc_key: Option = shared_alloc_op + .get_attr_alloc_key(ctx) + .map(|s| String::from((*s).clone())); + + // Optional and diagnostic: the Rust path of the originating `static`, + // carried through so the emitted global can name its source. + let source_name: Option = shared_alloc_op + .get_attr_source_name(ctx) + .map(|s| String::from((*s).clone())); + + let elem_type_attr = op_ref + .attributes + .get::(&"elem_type".try_into().unwrap()) + .ok_or_else(|| { + anyhow_to_pliron(anyhow::anyhow!( + "MirSharedAllocOp missing elem_type TypeAttr attribute" + )) + })?; + let mir_elem_type = elem_type_attr.get_type(ctx); + + let size_attr = op_ref + .attributes + .get::(&"size".try_into().unwrap()) + .ok_or_else(|| { + anyhow_to_pliron(anyhow::anyhow!( + "MirSharedAllocOp missing size IntegerAttr attribute" + )) + })?; + let size = size_attr.value().to_u64(); + + let alignment = shared_alloc_op.get_alignment_value(ctx).unwrap_or(0); + let debug_info = llvm::debug_global_variable(ctx, op); + let debug_owner_function = llvm::debug_global_owner_function(ctx, op); + + ( + alloc_key, + source_name, + mir_elem_type, + size, + alignment, + debug_info, + debug_owner_function, + ) + }; + + // A shared variable DIE must describe the physical allocation. Gate the + // optional debug identity against the lowered element layout before it + // can reach the declaration record or the emitted global. + let llvm_elem_type = convert_type(ctx, mir_elem_type).map_err(anyhow_to_pliron)?; + let debug_info = debug_info + .filter(|info| shared_debug_type_matches_physical(ctx, info, llvm_elem_type, size)); + let debug_owner_function = debug_info.as_ref().and(debug_owner_function); + + let declaration = SharedGlobalDeclaration { + kind: SharedGlobalKind::Static, + mir_elem_type, + size, + alignment, + debug_info, + debug_owner_function, + }; + + // A key names one physical allocation, not merely a preferred symbol. + // Reuse it only when the complete storage declaration agrees; otherwise + // the later address would silently inherit the first allocation's type, + // extent, or alignment. + let global_name = if let Some(key) = alloc_key.as_ref() + && let Some(existing) = shared_globals.get(key) + { + if !existing.declaration.same_storage(&declaration) { + return Err(anyhow_to_pliron(anyhow::anyhow!( + "duplicate shared alloc_key {:?} has an incompatible declaration", + key + ))); + } + let symbol = existing.symbol.clone(); + let debug_identity_matches = existing.declaration.debug_info == declaration.debug_info + && existing.declaration.debug_owner_function == declaration.debug_owner_function; + if !debug_identity_matches { + // One physical allocation reached under two debug identities, for + // example a function-local static whose owning function was + // materialized (or inlined) into two kernels. Debug metadata must + // never fail a build that a release build accepts, and DWARF + // cannot truthfully scope one AS3 object to two subprograms, so + // fail open: drop the attachment from the materialized global and + // demote the cached record to metadata-free. + strip_shared_debug_identity(ctx, op, &symbol)?; + if let Some(record) = shared_globals.get_mut(key) { + record.declaration.debug_info = None; + record.declaration.debug_owner_function = None; + } + } + symbol + } else { + create_shared_global( + ctx, + op, + shared_globals, + next_shared_mem_index, + SharedAllocSpec { + mir_elem_type, + size, + alignment, + alloc_key, + source_name: source_name.as_deref(), + debug_info: declaration.debug_info.as_ref(), + debug_owner_function: declaration.debug_owner_function.as_deref(), + }, + )? + }; + + let address_of_op = llvm::AddressOfOp::new(ctx, global_name, 3); + rewriter.insert_operation(ctx, address_of_op.get_operation()); + rewriter.replace_operation(ctx, op, address_of_op.get_operation()); + + Ok(()) +} + +/// Everything `create_shared_global` needs about one `mir.shared_alloc`. +/// +/// Mirrors [`DeviceGlobalSpec`] for the shared-memory path. +struct SharedAllocSpec<'a> { + mir_elem_type: TypeHandle, + size: u64, + alignment: u64, + alloc_key: Option, + source_name: Option<&'a str>, + debug_info: Option<&'a llvm::DebugGlobalVariableInfo>, + debug_owner_function: Option<&'a str>, +} + +/// Create a shared memory global variable in the module. +/// +/// Creates an LLVM global variable with: +/// - Array type: `[size x element_type]` +/// - Address space 3 (shared memory) +/// - Optional alignment +/// - Unique generated name (`__shared_mem_N`) +/// +/// The global is inserted at the front of the module block. When +/// `spec.alloc_key` is `Some`, the key is moved into `shared_globals` so that +/// later allocations with the same key reuse this global only after the caller +/// has checked that their complete declarations agree. +/// +/// `spec.source_name`, when present, is the Rust path of the `static` this +/// allocation came from. The generated symbol stays anonymous; the name is +/// recorded as an attribute on the global so the exporter can render it +/// beside the definition. Only the allocation that *creates* the global +/// contributes a name — a later allocation with the same `alloc_key` hits +/// the cache and never reaches this function — which is consistent because +/// the key and the name are both derived from the same constant. +/// +/// `next_shared_mem_index` is scoped to one `MirToLlvmConversionDriver` +/// instance (one module), not a process-global counter: `N` is a function of +/// this module's own MIR walk order, not of how many other modules have +/// lowered a shared allocation earlier in the process (#706). +fn create_shared_global( + ctx: &mut Context, + op: Ptr, + shared_globals: &mut SharedGlobalsMap, + next_shared_mem_index: &mut usize, + spec: SharedAllocSpec<'_>, +) -> Result { + let llvm_elem_type = convert_type(ctx, spec.mir_elem_type).map_err(anyhow_to_pliron)?; + let array_type = ArrayType::get(ctx, llvm_elem_type, spec.size); + + let counter = *next_shared_mem_index; + *next_shared_mem_index += 1; + let name: pliron::identifier::Identifier = + format!("__shared_mem_{counter}").try_into().unwrap(); + + let global_op = if spec.alignment > 0 { + llvm::GlobalOp::new_with_alignment(ctx, name.clone(), array_type.into(), spec.alignment) + } else { + llvm::GlobalOp::new(ctx, name.clone(), array_type.into()) + }; + global_op.set_address_space(ctx, llvm_export::types::address_space::SHARED); + if let Some(source_name) = spec.source_name { + use llvm_export::ops::GlobalOpExt; + global_op.set_shared_source_name(ctx, source_name); + } + if let Some(info) = spec.debug_info { + llvm::set_debug_global_variable(ctx, global_op.get_operation(), info); + } + if let Some(owner) = spec.debug_owner_function { + llvm::set_debug_global_owner_function(ctx, global_op.get_operation(), owner); + } + + let parent_block = op + .deref(ctx) + .get_parent_block() + .ok_or_else(|| anyhow_to_pliron(anyhow::anyhow!("Op has no parent block")))?; + let module_op = helpers::get_module_from_block(ctx, parent_block).map_err(anyhow_to_pliron)?; + let region = module_op.deref(ctx).get_region(0); + let module_block = region + .deref(ctx) + .iter(ctx) + .next() + .ok_or_else(|| anyhow_to_pliron(anyhow::anyhow!("Module is empty")))?; + + global_op.get_operation().insert_at_front(module_block, ctx); + + if let Some(key) = spec.alloc_key { + shared_globals.insert( + key, + SharedGlobalRecord { + symbol: name.clone(), + declaration: SharedGlobalDeclaration { + kind: SharedGlobalKind::Static, + mir_elem_type: spec.mir_elem_type, + size: spec.size, + alignment: spec.alignment, + debug_info: spec.debug_info.cloned(), + debug_owner_function: spec.debug_owner_function.map(str::to_owned), + }, + }, + ); + } + + Ok(name) +} + +/// A shared variable DIE must describe the compiler-materialized allocation, +/// never the marker ZST or an element graph whose pointer base types are not +/// representable yet. The frontend checks semantic element layout; this final +/// boundary verifies total physical bits/count against the LLVM backing. +fn shared_debug_type_matches_physical( + ctx: &Context, + info: &llvm::DebugGlobalVariableInfo, + llvm_elem_type: TypeHandle, + count: u64, +) -> bool { + let Some((element_size, _)) = llvm_type_size_align(ctx, llvm_elem_type) else { + return false; + }; + let Some(element_bytes) = element_size.checked_mul(count) else { + return false; + }; + let Some(physical_bits) = element_bytes.checked_mul(8) else { + return false; + }; + if info.ty.size_bits() != physical_bits { + return false; + } + match &info.ty { + llvm::DebugLocalTypeKind::Array { + element, + count: debug_count, + .. + } => *debug_count == count && element.size_bits() == element_size.saturating_mul(8), + // `Barrier` retains its semantic repr(C) struct while the backing is + // one i64. Other marker types are never admitted by the frontend. + llvm::DebugLocalTypeKind::Struct { .. } => count == 1, + _ => false, + } +} + +/// Detach the debug identity from an already-materialized shared global. +/// +/// The fail-open path for divergent debug identities on one `alloc_key`: +/// the storage stays shared and valid, only the optional DWARF attachment is +/// dropped so it cannot misattribute the allocation to the wrong static or +/// owner function. +fn strip_shared_debug_identity( + ctx: &mut Context, + op: Ptr, + name: &pliron::identifier::Identifier, +) -> Result<()> { + let parent_block = op + .deref(ctx) + .get_parent_block() + .ok_or_else(|| anyhow_to_pliron(anyhow::anyhow!("Op has no parent block")))?; + let module_op = helpers::get_module_from_block(ctx, parent_block).map_err(anyhow_to_pliron)?; + let module_block = module_op + .deref(ctx) + .get_region(0) + .deref(ctx) + .iter(ctx) + .next() + .ok_or_else(|| anyhow_to_pliron(anyhow::anyhow!("Module is empty")))?; + let existing = module_block + .deref(ctx) + .iter(ctx) + .filter_map(|candidate| Operation::get_op::(candidate, ctx)) + .find(|global| global.get_symbol_name(ctx) == *name) + .ok_or_else(|| { + anyhow_to_pliron(anyhow::anyhow!( + "shared allocation cache refers to missing LLVM global `@{name}`" + )) + })?; + llvm::clear_debug_global_identity(ctx, existing.get_operation()); + Ok(()) +} diff --git a/crates/mir-lower/src/convert/ops/memory/tests/access.rs b/crates/mir-lower/src/convert/ops/memory/tests/access.rs new file mode 100644 index 0000000000..954c1afc27 --- /dev/null +++ b/crates/mir-lower/src/convert/ops/memory/tests/access.rs @@ -0,0 +1,473 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#![allow(clippy::disallowed_methods)] + +use super::*; + +#[test] +fn convert_alloca_lowers_to_llvm_alloca() { + let mut ctx = make_ctx(); + let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); + let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty, true); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + + let alloca_op = Operation::new( + &mut ctx, + mir::MirAllocaOp::get_concrete_op_info(), + vec![mir_ptr_ty.into()], + vec![], + vec![], + 0, + ); + alloca_op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + assert_eq!( + count_ops::(&ctx, &body), + 1, + "expected exactly one llvm.alloca" + ); + let alloca = find_first::(&ctx, &body).unwrap(); + // Element type should round-trip through convert_type as i32. + let elem_ty = alloca.result_pointee_type(&ctx); + assert!(elem_ty.deref(&ctx).is::()); +} + +#[test] +fn convert_alloca_preserves_nested_array_element_alignment() { + let mut ctx = make_ctx(); + let tuple_ty = over_aligned_tuple_ty(&mut ctx); + let inner: TypeHandle = MirArrayType::get(&mut ctx, tuple_ty, 2).into(); + let outer: TypeHandle = MirArrayType::get(&mut ctx, inner, 3).into(); + let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, outer, true); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + + let alloca_op = Operation::new( + &mut ctx, + mir::MirAllocaOp::get_concrete_op_info(), + vec![mir_ptr_ty.into()], + vec![], + vec![], + 0, + ); + alloca_op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + let alloca = find_first::(&ctx, &body).expect("expected llvm.alloca"); + assert_eq!( + llvm_export::ops::op_alignment(&ctx, alloca.get_operation()), + Some(32) + ); +} + +#[test] +fn convert_alloca_preserves_debug_local_metadata() { + let mut ctx = make_ctx(); + let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); + let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty, true); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + + let alloca_op = Operation::new( + &mut ctx, + mir::MirAllocaOp::get_concrete_op_info(), + vec![mir_ptr_ty.into()], + vec![], + vec![], + 0, + ); + llvm::set_debug_local_variable( + &mut ctx, + alloca_op, + llvm::DebugLocalVariableInfo { + name: "x".to_string(), + argument_index: Some(1), + ty: llvm::DebugLocalTypeKind::Basic { + name: "i32".to_string(), + size_bits: 32, + encoding: "DW_ATE_signed", + }, + }, + ); + alloca_op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + let alloca = find_first::(&ctx, &body).unwrap(); + let info = llvm::debug_local_variable(&ctx, alloca.get_operation()) + .expect("debug local metadata should survive lowering"); + + assert_eq!(info.name, "x"); + assert_eq!(info.argument_index, Some(1)); + assert_eq!( + info.ty, + llvm::DebugLocalTypeKind::Basic { + name: "i32".to_string(), + size_bits: 32, + encoding: "DW_ATE_signed", + } + ); +} + +#[test] +fn convert_store_lowers_to_llvm_store() { + let mut ctx = make_ctx(); + let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); + let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty, true); + + // Kernel takes (ptr, val) so we can store one into the other. + let (module_ptr, block) = build_kernel(&mut ctx, vec![mir_ptr_ty.into(), i32_ty], vec![]); + let ptr_val = block.deref(&ctx).get_argument(0); + let val = block.deref(&ctx).get_argument(1); + + let store_op = Operation::new( + &mut ctx, + mir::MirStoreOp::get_concrete_op_info(), + vec![], + vec![ptr_val, val], + vec![], + 0, + ); + store_op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + assert_eq!( + count_ops::(&ctx, &body), + 1, + "expected one llvm.store" + ); + // The original mir.store must be gone. + assert_eq!(count_ops::(&ctx, &body), 0); + + // convert_store swaps operand order: mir.store is [ptr, value] but + // llvm.store takes (value, ptr). Verify that mapping survived. + let store = find_first::(&ctx, &body).unwrap(); + let addr_ty = store.get_operand_address(&ctx).get_type(&ctx); + assert!(addr_ty.deref(&ctx).is::(), "operand 1 is ptr"); +} + +#[test] +fn convert_store_preserves_volatile() { + let mut ctx = make_ctx(); + let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); + let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty, true); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![mir_ptr_ty.into(), i32_ty], vec![]); + let ptr_val = block.deref(&ctx).get_argument(0); + let val = block.deref(&ctx).get_argument(1); + + let store_op = Operation::new( + &mut ctx, + mir::MirStoreOp::get_concrete_op_info(), + vec![], + vec![ptr_val, val], + vec![], + 0, + ); + mir::MirStoreOp::new(store_op).set_volatile(&mut ctx, true); + store_op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + let store = find_first::(&ctx, &body).unwrap(); + assert!( + llvm_export::ops::op_volatile(&ctx, store.get_operation()), + "volatile mir.store must lower to a volatile llvm.store" + ); +} + +#[test] +fn convert_load_lowers_to_llvm_load() { + let mut ctx = make_ctx(); + let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); + let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty, false); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![mir_ptr_ty.into()], vec![]); + let ptr_val = block.deref(&ctx).get_argument(0); + + let load_op = Operation::new( + &mut ctx, + mir::MirLoadOp::get_concrete_op_info(), + vec![i32_ty], + vec![ptr_val], + vec![], + 0, + ); + load_op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + assert_eq!(count_ops::(&ctx, &body), 1); + assert_eq!(count_ops::(&ctx, &body), 0); +} + +#[test] +fn convert_load_preserves_volatile() { + let mut ctx = make_ctx(); + let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); + let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty, false); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![mir_ptr_ty.into()], vec![]); + let ptr_val = block.deref(&ctx).get_argument(0); + + let load_op = Operation::new( + &mut ctx, + mir::MirLoadOp::get_concrete_op_info(), + vec![i32_ty], + vec![ptr_val], + vec![], + 0, + ); + mir::MirLoadOp::new(load_op).set_volatile(&mut ctx, true); + load_op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + let load = find_first::(&ctx, &body).unwrap(); + assert!( + llvm_export::ops::op_volatile(&ctx, load.get_operation()), + "volatile mir.load must lower to a volatile llvm.load" + ); +} + +#[test] +fn convert_ref_lowers_to_alloca_then_store() { + let mut ctx = make_ctx(); + let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); + let mir_ptr_ty = + MirPtrType::get_generic_with_kind(&mut ctx, i32_ty, false, MirPointerKind::SharedRef); + + // Take a u32 by value, build `&x`. + let (module_ptr, block) = build_kernel(&mut ctx, vec![i32_ty], vec![]); + let arg = block.deref(&ctx).get_argument(0); + + let ref_op_ptr = Operation::new( + &mut ctx, + mir::MirRefOp::get_concrete_op_info(), + vec![mir_ptr_ty.into()], + vec![arg], + vec![], + 0, + ); + let ref_op = mir::MirRefOp::new(ref_op_ptr); + ref_op.set_mutable(&mut ctx, false); + ref_op.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); + ref_op_ptr.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + assert_eq!( + count_ops::(&ctx, &body), + 1, + "ref must materialize via alloca" + ); + assert_eq!( + count_ops::(&ctx, &body), + 1, + "ref must store the value into the alloca" + ); + assert_eq!(count_ops::(&ctx, &body), 0); +} + +#[test] +fn convert_ref_preserves_tuple_alignment_on_alloca_and_store() { + let mut ctx = make_ctx(); + let tuple_ty = over_aligned_tuple_ty(&mut ctx); + let mir_ptr_ty = + MirPtrType::get_generic_with_kind(&mut ctx, tuple_ty, false, MirPointerKind::SharedRef); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + + let undef = mir::MirUndefOp::new(&mut ctx, tuple_ty); + undef.get_operation().insert_at_back(block, &ctx); + let value = undef.get_operation().deref(&ctx).get_result(0); + let ref_op = Operation::new( + &mut ctx, + mir::MirRefOp::get_concrete_op_info(), + vec![mir_ptr_ty.into()], + vec![value], + vec![], + 0, + ); + let mir_ref = mir::MirRefOp::new(ref_op); + mir_ref.set_mutable(&mut ctx, false); + mir_ref.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); + ref_op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + let alloca = find_first::(&ctx, &body).expect("expected llvm.alloca"); + let store = find_first::(&ctx, &body).expect("expected llvm.store"); + assert_eq!( + llvm_export::ops::op_alignment(&ctx, alloca.get_operation()), + Some(32) + ); + assert_eq!( + llvm_export::ops::op_alignment(&ctx, store.get_operation()), + Some(32) + ); +} + +#[test] +fn convert_ref_preserves_over_aligned_union_array_layout_and_alignment() { + for abi_align in [32, 64] { + let mut ctx = make_ctx(); + let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let union_ty: TypeHandle = MirUnionType::get( + &mut ctx, + format!("Align{abi_align}Union"), + vec!["word".into()], + vec![u32_ty], + abi_align, + abi_align, + ) + .into(); + let array_ty: TypeHandle = MirArrayType::get(&mut ctx, union_ty, 3).into(); + let mir_ptr_ty = + MirPtrType::get_generic_with_kind(&mut ctx, array_ty, false, MirPointerKind::SharedRef); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + + let undef = mir::MirUndefOp::new(&mut ctx, array_ty); + undef.get_operation().insert_at_back(block, &ctx); + let value = undef.get_operation().deref(&ctx).get_result(0); + let ref_op = Operation::new( + &mut ctx, + mir::MirRefOp::get_concrete_op_info(), + vec![mir_ptr_ty.into()], + vec![value], + vec![], + 0, + ); + let mir_ref = mir::MirRefOp::new(ref_op); + mir_ref.set_mutable(&mut ctx, false); + mir_ref.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); + ref_op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + let alloca = find_first::(&ctx, &body).expect("expected llvm.alloca"); + let store = find_first::(&ctx, &body).expect("expected llvm.store"); + let llvm_array_ty = alloca.result_pointee_type(&ctx); + let llvm_array_data = llvm_array_ty.deref(&ctx); + let llvm_array = llvm_array_data + .downcast_ref::() + .expect("over-aligned union array must remain an LLVM array"); + + assert_eq!(llvm_array.size(), 3); + assert_eq!( + crate::convert::types::llvm_type_size_align(&ctx, llvm_array.elem_type()), + Some((abi_align, 16)) + ); + assert_eq!( + crate::convert::types::llvm_type_size_align(&ctx, llvm_array_ty), + Some((abi_align * 3, 16)) + ); + assert_eq!( + llvm_export::ops::op_alignment(&ctx, alloca.get_operation()), + Some(abi_align as u32) + ); + assert_eq!( + llvm_export::ops::op_alignment(&ctx, store.get_operation()), + Some(abi_align as u32) + ); + } +} + +#[test] +fn convert_ptr_offset_lowers_to_gep_with_pointee_elem_type() { + let mut ctx = make_ctx(); + let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); + let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signless).into(); + let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty, true); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![mir_ptr_ty.into(), i64_ty], vec![]); + let ptr_val = block.deref(&ctx).get_argument(0); + let off_val = block.deref(&ctx).get_argument(1); + + let off_op = Operation::new( + &mut ctx, + mir::MirPtrOffsetOp::get_concrete_op_info(), + vec![mir_ptr_ty.into()], + vec![ptr_val, off_val], + vec![], + 0, + ); + off_op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + let gep = find_first::(&ctx, &body).expect("expected one llvm.gep"); + // Element type must come from the MirPtrType pointee (i32), not the + // i8 fallback used when no operand-type info is available. + let elem_ty = gep.src_elem_type(&ctx); + let elem_ty_ref = elem_ty.deref(&ctx); + let int_ty = elem_ty_ref + .downcast_ref::() + .expect("gep src_elem_type should be IntegerType"); + assert_eq!(int_ty.width(), 32, "gep elem type must be i32 (pointee)"); + assert!( + llvm::gep_inbounds(&ctx, gep.get_operation()), + "ordinary pointer offsets retain the in-bounds contract" + ); +} + +#[test] +fn convert_wrapping_ptr_offset_lowers_to_non_inbounds_gep() { + let mut ctx = make_ctx(); + let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); + let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signed).into(); + let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty, false); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![mir_ptr_ty.into(), i64_ty], vec![]); + let ptr_val = block.deref(&ctx).get_argument(0); + let off_val = block.deref(&ctx).get_argument(1); + + let off_op = Operation::new( + &mut ctx, + mir::MirPtrOffsetOp::get_concrete_op_info(), + vec![mir_ptr_ty.into()], + vec![ptr_val, off_val], + vec![], + 0, + ); + mir::MirPtrOffsetOp::new(off_op).set_inbounds(&mut ctx, false); + off_op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + let gep = find_first::(&ctx, &body).expect("expected one llvm.gep"); + assert!( + !llvm::gep_inbounds(&ctx, gep.get_operation()), + "wrapping pointer offsets must not promise in-bounds arithmetic" + ); +} diff --git a/crates/mir-lower/src/convert/ops/memory/tests/alignment.rs b/crates/mir-lower/src/convert/ops/memory/tests/alignment.rs new file mode 100644 index 0000000000..aae268bb08 --- /dev/null +++ b/crates/mir-lower/src/convert/ops/memory/tests/alignment.rs @@ -0,0 +1,806 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#![allow(clippy::disallowed_methods)] + +use super::*; + +/// Lower `mir.load (mir.field_addr %p, field_index)` for a struct of +/// signless integer fields with the given layout and report the alignment +/// stamped on the resulting `llvm.load`. `None` means no stamp survived +/// and the exporter's natural-alignment default applies. +fn lowered_field_load_alignment( + field_bit_widths: Vec, + field_offsets: Vec, + total_size: u64, + abi_align: u64, + field_index: u32, +) -> Option { + let mut ctx = make_ctx(); + let field_types: Vec = field_bit_widths + .iter() + .map(|w| IntegerType::get(&ctx, *w, Signedness::Signless).into()) + .collect(); + let field_names = (0..field_types.len()).map(|i| format!("f{i}")).collect(); + let struct_ty: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "FieldLoadAlign".into(), + field_names, + field_types.clone(), + vec![], + field_offsets, + total_size, + abi_align, + ) + .into(); + let struct_ptr_ty = MirPtrType::get_generic(&mut ctx, struct_ty, false); + let field_ty = field_types[field_index as usize]; + let field_ptr_ty = MirPtrType::get_generic(&mut ctx, field_ty, false); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![struct_ptr_ty.into()], vec![]); + let struct_ptr_val = block.deref(&ctx).get_argument(0); + + let field_addr_op = + mir::MirFieldAddrOp::build(&mut ctx, struct_ptr_val, field_ptr_ty.into(), field_index) + .expect("field_addr build"); + field_addr_op.insert_at_back(block, &ctx); + let field_ptr_val = field_addr_op.deref(&ctx).get_result(0); + + let load_op = Operation::new( + &mut ctx, + mir::MirLoadOp::get_concrete_op_info(), + vec![field_ty], + vec![field_ptr_val], + vec![], + 0, + ); + load_op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + let load = find_first::(&ctx, &body).expect("expected one llvm.load"); + llvm_export::ops::op_alignment(&ctx, load.get_operation()) +} + +/// Field 0 of an over-aligned struct sits at the aggregate's own +/// alignment, which the field's scalar result type cannot state on its +/// own. This is what lets LoadStoreVectorizer fuse the adjacent pair. +#[test] +fn convert_load_inherits_overaligned_field_alignment_at_offset_zero() { + // #[repr(C, align(8))] struct { a: i32, b: i32 } + assert_eq!( + lowered_field_load_alignment(vec![32, 32], vec![0, 4], 8, 8, 0), + Some(8) + ); +} + +/// A field at a nonzero offset proves `gcd(abi_align, offset)`: an i32 at +/// offset 8 of an align-16 struct proves 8, beating its natural 4. +#[test] +fn convert_load_narrows_field_alignment_to_gcd_of_align_and_offset() { + // #[repr(C, align(16))] struct { a: i64, b: i32 } + assert_eq!( + lowered_field_load_alignment(vec![64, 32], vec![0, 8], 16, 16, 1), + Some(8) + ); +} + +/// Whole-value loads of pointer-free packed structs use the packed LLVM +/// representation, preserving rustc's byte size and field offsets. +#[test] +fn packed_struct_whole_value_load_uses_packed_layout() { + let mut ctx = make_ctx(); + let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let packed_ty: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "Packed".into(), + vec!["tag".into(), "value".into()], + vec![u8_ty, u32_ty], + vec![0, 1], + vec![0, 1], + 5, + 1, + ) + .into(); + let ptr_ty = MirPtrType::get_generic(&mut ctx, packed_ty, false); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![ptr_ty.into()], vec![]); + let ptr_val = block.deref(&ctx).get_argument(0); + + let load_op = Operation::new( + &mut ctx, + mir::MirLoadOp::get_concrete_op_info(), + vec![packed_ty], + vec![ptr_val], + vec![], + 0, + ); + load_op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr) + .expect("whole-value load of a pointer-free packed struct must lower"); + + let body = kernel_blocks(&ctx, module_ptr); + let load = find_first::(&ctx, &body).expect("expected packed llvm.load"); + let result_ty = load + .get_operation() + .deref(&ctx) + .get_result(0) + .get_type(&ctx); + let result_ty_ref = result_ty.deref(&ctx); + let struct_ty = result_ty_ref + .downcast_ref::() + .expect("packed load result must be an LLVM struct"); + assert_eq!(struct_ty.layout(), StructLayout::Packed); + assert_eq!( + crate::convert::types::llvm_type_size_align(&ctx, result_ty), + Some((5, 1)) + ); + assert_eq!( + llvm_export::ops::op_alignment(&ctx, load.get_operation()), + Some(1) + ); +} + +#[test] +fn packed_struct_whole_value_load_with_shared_pointer_fails_closed() { + let mut ctx = make_ctx(); + let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let pointee: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let shared_ty: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, false).into(); + let packed_ty: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "PackedShared".into(), + vec!["tag".into(), "ptr".into()], + vec![u8_ty, shared_ty], + vec![0, 1], + vec![0, 1], + 9, + 1, + ) + .into(); + let ptr_ty = MirPtrType::get_generic(&mut ctx, packed_ty, false); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![ptr_ty.into()], vec![]); + let ptr_val = block.deref(&ctx).get_argument(0); + let load_op = Operation::new( + &mut ctx, + mir::MirLoadOp::get_concrete_op_info(), + vec![packed_ty], + vec![ptr_val], + vec![], + 0, + ); + load_op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + let err = crate::lower_mir_to_llvm(&mut ctx, module_ptr) + .expect_err("packed whole-value load containing AS3 must remain fail-closed"); + assert!( + format!("{err:?}").contains("target-mode dependent"), + "the refusal must identify the target-dependent packed AS3 image: {err:?}" + ); +} + +/// Whole-value stores use the same packed representation as construction +/// and loads, while preserving the MIR aggregate's proved ABI alignment. +#[test] +fn packed_struct_whole_value_store_uses_packed_layout() { + let mut ctx = make_ctx(); + let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let packed_ty: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "Packed".into(), + vec!["tag".into(), "value".into()], + vec![u8_ty, u32_ty], + vec![0, 1], + vec![0, 1], + 5, + 1, + ) + .into(); + let ptr_ty = MirPtrType::get_generic(&mut ctx, packed_ty, false); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![ptr_ty.into(), packed_ty], vec![]); + let ptr_val = block.deref(&ctx).get_argument(0); + let val = block.deref(&ctx).get_argument(1); + + let store_op = Operation::new( + &mut ctx, + mir::MirStoreOp::get_concrete_op_info(), + vec![], + vec![ptr_val, val], + vec![], + 0, + ); + store_op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr) + .expect("whole-value store of a pointer-free packed struct must lower"); + + let body = kernel_blocks(&ctx, module_ptr); + let store = find_first::(&ctx, &body).expect("expected packed llvm.store"); + let value_ty = store + .get_operation() + .deref(&ctx) + .get_operand(0) + .get_type(&ctx); + let value_ty_ref = value_ty.deref(&ctx); + let struct_ty = value_ty_ref + .downcast_ref::() + .expect("packed store value must be an LLVM struct"); + assert_eq!(struct_ty.layout(), StructLayout::Packed); + assert_eq!( + crate::convert::types::llvm_type_size_align(&ctx, value_ty), + Some((5, 1)) + ); + assert_eq!( + llvm_export::ops::op_alignment(&ctx, store.get_operation()), + Some(1) + ); +} + +/// A naturally aligned inner struct sitting at a packed byte offset: the +/// field address only proves align 1, and the load must claim that over +/// the inner type's recorded abi alignment. Claiming the abi would stamp +/// `align 4` on a 1-aligned address, which llc may honor with a wider +/// access than the bytes allow. +#[test] +fn convert_load_claims_address_alignment_over_abi_at_packed_offsets() { + let mut ctx = make_ctx(); + let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let inner_ty: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "Inner".into(), + vec!["v".into()], + vec![u32_ty], + vec![0], + vec![0], + 4, + 4, + ) + .into(); + let outer_ty: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "PackedOuter".into(), + vec!["tag".into(), "inner".into()], + vec![u8_ty, inner_ty], + vec![0, 1], + vec![0, 1], + 5, + 1, + ) + .into(); + let outer_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, outer_ty, false).into(); + let inner_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, inner_ty, false).into(); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![outer_ptr_ty], vec![]); + let base = block.deref(&ctx).get_argument(0); + + let field_addr_op = + mir::MirFieldAddrOp::build(&mut ctx, base, inner_ptr_ty, 1).expect("field_addr build"); + field_addr_op.insert_at_back(block, &ctx); + let field_ptr_val = field_addr_op.deref(&ctx).get_result(0); + + let load_op = Operation::new( + &mut ctx, + mir::MirLoadOp::get_concrete_op_info(), + vec![inner_ty], + vec![field_ptr_val], + vec![], + 0, + ); + load_op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + let load = find_first::(&ctx, &body).expect("expected one llvm.load"); + assert_eq!( + llvm_export::ops::op_alignment(&ctx, load.get_operation()), + Some(1), + "the load must claim the address's proved alignment, not the inner abi" + ); +} + +/// A struct with no extra alignment proves nothing beyond the scalar's +/// natural alignment: the stamp equals the exporter's default 4, so the +/// emitted access is unchanged. +#[test] +fn convert_load_keeps_natural_alignment_without_overalignment() { + // struct { a: i32, b: i32 } with rustc's natural abi_align 4 + assert_eq!( + lowered_field_load_alignment(vec![32, 32], vec![0, 4], 8, 4, 1), + Some(4) + ); +} + +/// dialect-mir only verifier-enforces power-of-two alignment for unions +/// and enums. A malformed hand-built struct layout must decline the stamp +/// rather than emit a non-power-of-two `align N` that llc rejects. +#[test] +fn convert_load_declines_non_power_of_two_field_alignment() { + assert_eq!( + lowered_field_load_alignment(vec![32], vec![0], 12, 12, 0), + None + ); +} + +/// Lower `mir.store %v, (mir.field_addr %p, field_index)` for a struct of +/// signless integer fields with the given layout and report the alignment +/// stamped on the resulting `llvm.store`. `None` means no stamp survived +/// and the exporter's natural-alignment default applies. +fn lowered_field_store_alignment( + field_bit_widths: Vec, + field_offsets: Vec, + total_size: u64, + abi_align: u64, + field_index: u32, +) -> Option { + let mut ctx = make_ctx(); + let field_types: Vec = field_bit_widths + .iter() + .map(|w| IntegerType::get(&ctx, *w, Signedness::Signless).into()) + .collect(); + let field_names = (0..field_types.len()).map(|i| format!("f{i}")).collect(); + let struct_ty: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "FieldStoreAlign".into(), + field_names, + field_types.clone(), + vec![], + field_offsets, + total_size, + abi_align, + ) + .into(); + let struct_ptr_ty = MirPtrType::get_generic(&mut ctx, struct_ty, true); + let field_ty = field_types[field_index as usize]; + let field_ptr_ty = MirPtrType::get_generic(&mut ctx, field_ty, true); + + // The stored value arrives as a kernel argument of the field's own + // scalar type, so `value_abi_align` reports nothing about it and the + // address's stamp is the only alignment left -- the case this covers. + let (module_ptr, block) = build_kernel(&mut ctx, vec![struct_ptr_ty.into(), field_ty], vec![]); + let struct_ptr_val = block.deref(&ctx).get_argument(0); + let val = block.deref(&ctx).get_argument(1); + + let field_addr_op = + mir::MirFieldAddrOp::build(&mut ctx, struct_ptr_val, field_ptr_ty.into(), field_index) + .expect("field_addr build"); + field_addr_op.insert_at_back(block, &ctx); + let field_ptr_val = field_addr_op.deref(&ctx).get_result(0); + + let store_op = Operation::new( + &mut ctx, + mir::MirStoreOp::get_concrete_op_info(), + vec![], + vec![field_ptr_val, val], + vec![], + 0, + ); + store_op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + let store = find_first::(&ctx, &body).expect("expected one llvm.store"); + llvm_export::ops::op_alignment(&ctx, store.get_operation()) +} + +/// Field 0 of an over-aligned struct sits at the aggregate's own alignment, +/// which the stored scalar's type cannot state. This is what lets +/// LoadStoreVectorizer fuse the adjacent pair into one wide store. +#[test] +fn convert_store_inherits_overaligned_field_alignment_at_offset_zero() { + // #[repr(C, align(8))] struct { a: i32, b: i32 } + assert_eq!( + lowered_field_store_alignment(vec![32, 32], vec![0, 4], 8, 8, 0), + Some(8) + ); +} + +/// A field at a nonzero offset proves `gcd(abi_align, offset)`: an i32 at +/// offset 8 of an align-16 struct proves 8, beating its natural 4. +#[test] +fn convert_store_narrows_field_alignment_to_gcd_of_align_and_offset() { + // #[repr(C, align(16))] struct { a: i64, b: i32 } + assert_eq!( + lowered_field_store_alignment(vec![64, 32], vec![0, 8], 16, 16, 1), + Some(8) + ); +} + +/// A struct with no extra alignment proves nothing beyond the scalar's +/// natural alignment, so the emitted store is unchanged. Widening here +/// would claim an alignment the source never guaranteed. +#[test] +fn convert_store_keeps_natural_alignment_without_overalignment() { + // struct { a: i32, b: i32 } with rustc's natural abi_align 4 + assert_eq!( + lowered_field_store_alignment(vec![32, 32], vec![0, 4], 8, 4, 1), + Some(4) + ); +} + +/// A malformed hand-built layout must decline the stamp rather than emit a +/// non-power-of-two `align N` that llc rejects. Same guard the load path +/// has, and it matters more here: an over-aligned store instruction on an +/// under-aligned address is undefined, not merely slow. +#[test] +fn convert_store_declines_non_power_of_two_field_alignment() { + assert_eq!( + lowered_field_store_alignment(vec![32], vec![0], 12, 12, 0), + None + ); +} + +/// Lower `load (&arr[index])` where `arr` is the array field of an +/// over-aligned struct, and report the alignment the load ends up with. +/// +/// `index` of `Some(i)` builds a constant index, `None` a runtime one. +fn lowered_element_load_alignment( + element_bits: u32, + element_count: u64, + struct_abi_align: u64, + index: Option, +) -> Option { + use pliron::builtin::attributes::IntegerAttr; + use std::num::NonZeroUsize; + + let mut ctx = make_ctx(); + let element_ty: TypeHandle = IntegerType::get(&ctx, element_bits, Signedness::Signless).into(); + let array_ty: TypeHandle = + dialect_mir::types::MirArrayType::get(&mut ctx, element_ty, element_count).into(); + let elem_bytes = u64::from(element_bits) / 8; + let struct_ty: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "ElementLoadAlign".into(), + vec!["lanes".into()], + vec![array_ty], + vec![], + vec![0], + elem_bytes * element_count, + struct_abi_align, + ) + .into(); + let struct_ptr_ty = MirPtrType::get_generic(&mut ctx, struct_ty, false); + let array_ptr_ty = MirPtrType::get_generic(&mut ctx, array_ty, false); + let element_ptr_ty = MirPtrType::get_generic(&mut ctx, element_ty, false); + let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signed).into(); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![struct_ptr_ty.into(), i64_ty], vec![]); + let struct_ptr_val = block.deref(&ctx).get_argument(0); + + // &s.lanes -- carries the struct's alignment onto the array address. + let field_addr_op = + mir::MirFieldAddrOp::build(&mut ctx, struct_ptr_val, array_ptr_ty.into(), 0) + .expect("field_addr build"); + field_addr_op.insert_at_back(block, &ctx); + let array_ptr_val = field_addr_op.deref(&ctx).get_result(0); + + let index_val = match index { + Some(i) => { + let constant = Operation::new( + &mut ctx, + mir::MirConstantOp::get_concrete_op_info(), + vec![i64_ty], + vec![], + vec![], + 0, + ); + mir::MirConstantOp::new(constant).set_attr_value( + &ctx, + IntegerAttr::new( + IntegerType::get(&ctx, 64, Signedness::Signed), + APInt::from_u64(i, NonZeroUsize::new(64).unwrap()), + ), + ); + constant.insert_at_back(block, &ctx); + constant.deref(&ctx).get_result(0) + } + None => block.deref(&ctx).get_argument(1), + }; + + let elem_addr_op = Operation::new( + &mut ctx, + mir::MirArrayElementAddrOp::get_concrete_op_info(), + vec![element_ptr_ty.into()], + vec![array_ptr_val, index_val], + vec![], + 0, + ); + elem_addr_op.insert_at_back(block, &ctx); + let elem_ptr_val = elem_addr_op.deref(&ctx).get_result(0); + + let load_op = Operation::new( + &mut ctx, + mir::MirLoadOp::get_concrete_op_info(), + vec![element_ty], + vec![elem_ptr_val], + vec![], + 0, + ); + load_op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + let load = find_first::(&ctx, &body).expect("expected one llvm.load"); + llvm_export::ops::op_alignment(&ctx, load.get_operation()) +} + +/// Element 0 inherits the whole alignment the base address proved, which is +/// what lets the adjacent pair fuse into one wide load. +#[test] +fn convert_load_inherits_base_alignment_at_element_zero() { + // &(#[repr(C, align(8))] struct { lanes: [i32; 2] }).lanes[0] + assert_eq!(lowered_element_load_alignment(32, 2, 8, Some(0)), Some(8)); +} + +/// A nonzero constant index proves `gcd(base, i * stride)`: element 1 of an +/// align-8 `[i32; 2]` sits at byte 4, so it proves 4, not 8. +#[test] +fn convert_load_narrows_element_alignment_to_gcd_with_offset() { + assert_eq!(lowered_element_load_alignment(32, 2, 8, Some(1)), Some(4)); +} + +/// A runtime index can land on any element, so only what every stride +/// preserves may be claimed -- `gcd(base, stride)`, never the base itself. +#[test] +fn convert_load_claims_only_stride_alignment_for_a_runtime_index() { + assert_eq!(lowered_element_load_alignment(32, 2, 8, None), Some(4)); +} + +/// A base with no extra alignment proves nothing beyond the element's own +/// natural alignment, so the emitted access is unchanged. +#[test] +fn convert_load_keeps_natural_element_alignment_without_overalignment() { + assert_eq!(lowered_element_load_alignment(32, 2, 4, Some(0)), Some(4)); +} + +/// Like [`lowered_element_load_alignment`], but for an array whose element +/// is an aggregate built by `element_ty_of` (which also reports the +/// element's stored size in bytes). The claim on the element address must +/// then come from the element's *exact* stride — rustc's stored size, +/// padding included — not from any LLVM-level approximation. +/// +/// With `load_first_scalar` the element must itself be an array and the +/// access becomes `s.lanes[index][0]`, mirroring the nested-read chain +/// where the outer stamp is inherited by the inner index-0 address and +/// ends up on a scalar load that LoadStoreVectorizer trusts. Without it, +/// the element itself is loaded. +/// +/// Reports `(element address stamp, final load alignment)`. +fn lowered_aggregate_element_alignments( + element_ty_of: impl FnOnce(&mut Context) -> (TypeHandle, u64), + element_count: u64, + struct_abi_align: u64, + index: Option, + load_first_scalar: bool, +) -> (Option, Option) { + use pliron::builtin::attributes::IntegerAttr; + use std::num::NonZeroUsize; + + let mut ctx = make_ctx(); + let (element_ty, elem_bytes) = element_ty_of(&mut ctx); + let inner_scalar_ty = load_first_scalar.then(|| { + let element_ref = element_ty.deref(&ctx); + element_ref + .downcast_ref::() + .expect("load_first_scalar needs an array element") + .element_type() + }); + let array_ty: TypeHandle = MirArrayType::get(&mut ctx, element_ty, element_count).into(); + let struct_ty: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "AggregateElementAlign".into(), + vec!["lanes".into()], + vec![array_ty], + vec![], + vec![0], + elem_bytes * element_count, + struct_abi_align, + ) + .into(); + let struct_ptr_ty = MirPtrType::get_generic(&mut ctx, struct_ty, false); + let array_ptr_ty = MirPtrType::get_generic(&mut ctx, array_ty, false); + let element_ptr_ty = MirPtrType::get_generic(&mut ctx, element_ty, false); + let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signed).into(); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![struct_ptr_ty.into(), i64_ty], vec![]); + let struct_ptr_val = block.deref(&ctx).get_argument(0); + + // &s.lanes -- carries the struct's alignment onto the array address. + let field_addr_op = + mir::MirFieldAddrOp::build(&mut ctx, struct_ptr_val, array_ptr_ty.into(), 0) + .expect("field_addr build"); + field_addr_op.insert_at_back(block, &ctx); + let array_ptr_val = field_addr_op.deref(&ctx).get_result(0); + + let constant_index = |ctx: &mut Context, i: u64| { + let constant = Operation::new( + ctx, + mir::MirConstantOp::get_concrete_op_info(), + vec![i64_ty], + vec![], + vec![], + 0, + ); + mir::MirConstantOp::new(constant).set_attr_value( + ctx, + IntegerAttr::new( + IntegerType::get(ctx, 64, Signedness::Signed), + APInt::from_u64(i, NonZeroUsize::new(64).unwrap()), + ), + ); + constant.insert_at_back(block, ctx); + constant.deref(ctx).get_result(0) + }; + + let index_val = match index { + Some(i) => constant_index(&mut ctx, i), + None => block.deref(&ctx).get_argument(1), + }; + + let elem_addr_op = Operation::new( + &mut ctx, + mir::MirArrayElementAddrOp::get_concrete_op_info(), + vec![element_ptr_ty.into()], + vec![array_ptr_val, index_val], + vec![], + 0, + ); + elem_addr_op.insert_at_back(block, &ctx); + let elem_ptr_val = elem_addr_op.deref(&ctx).get_result(0); + + let (loaded_ty, loaded_ptr_val) = match inner_scalar_ty { + Some(scalar_ty) => { + let zero_val = constant_index(&mut ctx, 0); + let scalar_ptr_ty = MirPtrType::get_generic(&mut ctx, scalar_ty, false); + let inner_addr_op = Operation::new( + &mut ctx, + mir::MirArrayElementAddrOp::get_concrete_op_info(), + vec![scalar_ptr_ty.into()], + vec![elem_ptr_val, zero_val], + vec![], + 0, + ); + inner_addr_op.insert_at_back(block, &ctx); + (scalar_ty, inner_addr_op.deref(&ctx).get_result(0)) + } + None => (element_ty, elem_ptr_val), + }; + + let load_op = Operation::new( + &mut ctx, + mir::MirLoadOp::get_concrete_op_info(), + vec![loaded_ty], + vec![loaded_ptr_val], + vec![], + 0, + ); + load_op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + // GEP order follows source order: field address, then the element + // address under test (then the inner index-0 address when nested). + let geps = find_all::(&ctx, &body); + assert_eq!(geps.len(), if load_first_scalar { 3 } else { 2 }); + let element_gep_align = llvm_export::ops::address_alignment(&ctx, geps[1].get_operation()); + let load = find_first::(&ctx, &body).expect("expected one llvm.load"); + let load_align = llvm_export::ops::op_alignment(&ctx, load.get_operation()); + (element_gep_align, load_align) +} + +/// `[f32; 3]` element under an align-8 base: stride is 12, so a runtime +/// index proves `gcd(8, 12) = 4` — and the inner index-0 scalar read +/// inherits exactly that. Guards against sizing the element through an +/// LLVM-level approximation, whose guessed stride of 8 would stamp +/// align 8 onto addresses that are only 4-aligned (a miscompile once +/// LoadStoreVectorizer trusts it). +#[test] +fn convert_load_claims_exact_aggregate_stride_for_nested_array_elements() { + use pliron::builtin::types::FP32Type; + // &(#[repr(C, align(8))] struct { lanes: [[f32; 3]; 4] }).lanes[i][0] + let nested_f32x3 = |ctx: &mut Context| { + let f32_ty: TypeHandle = FP32Type::get(ctx).into(); + (MirArrayType::get(ctx, f32_ty, 3).into(), 12) + }; + assert_eq!( + lowered_aggregate_element_alignments(nested_f32x3, 4, 8, None, true), + (Some(4), Some(4)) + ); +} + +/// A constant index into the same nested array uses the exact byte +/// offset: element 1 sits at byte 12 (`gcd(8, 12) = 4`), element 2 at +/// byte 24 (`gcd(8, 24) = 8`, the full base alignment again). +#[test] +fn convert_load_narrows_nested_element_alignment_by_exact_byte_offset() { + use pliron::builtin::types::FP32Type; + let nested_f32x3 = |ctx: &mut Context| { + let f32_ty: TypeHandle = FP32Type::get(ctx).into(); + (MirArrayType::get(ctx, f32_ty, 3).into(), 12) + }; + assert_eq!( + lowered_aggregate_element_alignments(nested_f32x3, 4, 8, Some(1), true), + (Some(4), Some(4)) + ); + assert_eq!( + lowered_aggregate_element_alignments(nested_f32x3, 4, 8, Some(2), true), + (Some(8), Some(8)) + ); +} + +/// A tuple element's stride comes from rustc's recorded `total_size` +/// (trailing padding included): `(f32, f32, f32)` stores 12 bytes, so an +/// align-8 base proves only 4 on a runtime element address. +#[test] +fn convert_array_element_addr_takes_tuple_stride_from_recorded_layout() { + use pliron::builtin::types::FP32Type; + let f32x3_tuple = |ctx: &mut Context| { + let f32_ty: TypeHandle = FP32Type::get(ctx).into(); + let tuple_ty: TypeHandle = + MirTupleType::get_with_layout(ctx, vec![f32_ty; 3], vec![], vec![0, 4, 8], 12, 4) + .into(); + (tuple_ty, 12) + }; + let (element_gep_align, _load_align) = + lowered_aggregate_element_alignments(f32x3_tuple, 4, 8, None, false); + assert_eq!(element_gep_align, Some(4)); +} + +/// `f16` arrives from the importer as `MirFP16Type`, not the converted +/// LLVM `half`, and its stride is exactly 2: an align-8 base proves 2 on +/// a runtime element address and `gcd(8, 4) = 4` at element 2. Guards +/// the arm the importer actually exercises — the old sizing guessed 8 +/// for this type too, the same over-claim as the aggregate cases. +#[test] +fn convert_load_claims_exact_f16_element_stride() { + let f16_scalar = |ctx: &mut Context| { + let f16_ty: TypeHandle = dialect_mir::types::MirFP16Type::get(ctx).into(); + (f16_ty, 2) + }; + assert_eq!( + lowered_aggregate_element_alignments(f16_scalar, 4, 8, None, false), + (Some(2), Some(2)) + ); + assert_eq!( + lowered_aggregate_element_alignments(f16_scalar, 4, 8, Some(2), false), + (Some(4), Some(4)) + ); +} + +/// An element whose stored size is unknown (a struct built without rustc +/// layout) must not have its stride guessed: the element address claims +/// nothing and the load keeps the previous, weaker-but-sound behaviour. +#[test] +fn convert_array_element_addr_declines_unknown_element_stride() { + use pliron::builtin::types::FP32Type; + let opaque_struct = |ctx: &mut Context| { + let f32_ty: TypeHandle = FP32Type::get(ctx).into(); + let struct_ty: TypeHandle = + MirStructType::get(ctx, "OpaqueElement".into(), vec!["x".into()], vec![f32_ty]).into(); + (struct_ty, 4) + }; + assert_eq!( + lowered_aggregate_element_alignments(opaque_struct, 4, 8, None, false), + (None, None) + ); +} diff --git a/crates/mir-lower/src/convert/ops/memory/tests/debug_provenance.rs b/crates/mir-lower/src/convert/ops/memory/tests/debug_provenance.rs new file mode 100644 index 0000000000..5b34793741 --- /dev/null +++ b/crates/mir-lower/src/convert/ops/memory/tests/debug_provenance.rs @@ -0,0 +1,381 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#![allow(clippy::disallowed_methods)] + +use super::*; + +#[test] +fn convert_dbg_value_lowers_to_llvm_dbg_value() { + let mut ctx = make_ctx(); + let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![i32_ty], vec![]); + let value = block.deref(&ctx).get_argument(0); + + let dbg_op = mir::MirDbgValueOp::new(&mut ctx, value); + let dbg_loc = pliron::location::Location::Named { + name: "current value location".to_string(), + child_loc: Box::new(pliron::location::Location::Unknown), + }; + dbg_op + .get_operation() + .deref_mut(&ctx) + .set_loc(dbg_loc.clone()); + llvm::set_debug_local_variable( + &mut ctx, + dbg_op.get_operation(), + llvm::DebugLocalVariableInfo { + name: "x".to_string(), + argument_index: None, + ty: llvm::DebugLocalTypeKind::Basic { + name: "i32".to_string(), + size_bits: 32, + encoding: "DW_ATE_signed", + }, + }, + ); + llvm::set_debug_local_source_scope(&mut ctx, dbg_op.get_operation(), 42); + llvm::set_debug_fragment_variables( + &mut ctx, + dbg_op.get_operation(), + &[llvm::DebugFragmentVariableInfo { + variable: llvm::DebugLocalVariableInfo { + name: "pair".to_string(), + argument_index: None, + ty: llvm::DebugLocalTypeKind::Array { + name: "[u32; 2]".to_string(), + size_bits: 64, + element: Box::new(llvm::DebugLocalTypeKind::Basic { + name: "u32".to_string(), + size_bits: 32, + encoding: "DW_ATE_unsigned", + }), + count: 2, + }, + }, + fragment: llvm::DebugFragment { + offset_bits: 32, + size_bits: 32, + }, + source_scope: Some(42), + declaration: Some(llvm::DebugSourcePosition { + file: PathBuf::from("decl.rs"), + line: 7, + column: 3, + }), + }], + ); + llvm::set_debug_local_declaration_location( + &mut ctx, + dbg_op.get_operation(), + PathBuf::from("decl.rs"), + 7, + 3, + ); + dbg_op.get_operation().insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + assert_eq!(count_ops::(&ctx, &body), 0); + let dbg_value = find_first::(&ctx, &body) + .expect("expected lowered llvm.dbg_value marker"); + assert_eq!( + dbg_value.get_operation().deref(&ctx).loc(), + dbg_loc, + "dbg.value lowering should keep the current-value source location" + ); + let info = llvm::debug_local_variable(&ctx, dbg_value.get_operation()) + .expect("debug local metadata should survive dbg_value lowering"); + + assert_eq!(info.name, "x"); + assert_eq!( + llvm::debug_local_source_scope(&ctx, dbg_value.get_operation()), + Some(42), + "dbg.value lowering should keep the MIR source-scope owner" + ); + let (decl_file, decl_pos) = + llvm::debug_local_declaration_location(&ctx, dbg_value.get_operation()) + .expect("declaration location should survive dbg_value lowering"); + assert_eq!(decl_file, PathBuf::from("decl.rs")); + assert_eq!(decl_pos.line, 7); + assert_eq!(decl_pos.column, 3); + assert_eq!( + info.ty, + llvm::DebugLocalTypeKind::Basic { + name: "i32".to_string(), + size_bits: 32, + encoding: "DW_ATE_signed", + } + ); + let fragments = llvm::debug_fragment_variables(&ctx, dbg_value.get_operation()); + assert_eq!(fragments.len(), 1); + assert_eq!(fragments[0].fragment.offset_bits, 32); + assert_eq!(fragments[0].fragment.size_bits, 32); + assert_eq!(fragments[0].variable.name, "pair"); +} + +#[test] +fn convert_dbg_value_list_preserves_operands_and_expression() { + let mut ctx = make_ctx(); + let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); + let ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty, false); + let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signless).into(); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![ptr_ty.into(), i64_ty], vec![]); + let base = block.deref(&ctx).get_argument(0); + let index = block.deref(&ctx).get_argument(1); + + let dbg_op = mir::MirDbgValueListOp::new(&mut ctx, vec![base, index]); + llvm::set_debug_local_variable( + &mut ctx, + dbg_op.get_operation(), + llvm::DebugLocalVariableInfo { + name: "item".to_string(), + argument_index: None, + ty: llvm::DebugLocalTypeKind::Basic { + name: "u32".to_string(), + size_bits: 32, + encoding: "DW_ATE_unsigned", + }, + }, + ); + let expression = llvm::DebugValueExpression::new(vec![ + llvm::DebugValueExpressionOp::Arg(0), + llvm::DebugValueExpressionOp::Arg(1), + llvm::DebugValueExpressionOp::ConstU(4), + llvm::DebugValueExpressionOp::Mul, + llvm::DebugValueExpressionOp::Plus, + llvm::DebugValueExpressionOp::Deref, + ]); + llvm::set_debug_value_expression(&mut ctx, dbg_op.get_operation(), &expression); + dbg_op.get_operation().insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + assert_eq!(count_ops::(&ctx, &body), 0); + let dbg_value = find_first::(&ctx, &body) + .expect("expected lowered llvm.dbg_value_list marker"); + assert_eq!(dbg_value.values(&ctx), vec![base, index]); + assert_eq!( + llvm::debug_value_expression(&ctx, dbg_value.get_operation()), + Some(expression) + ); +} + +#[test] +fn convert_alloca_preserves_local_memory_provenance() { + let mut ctx = make_ctx(); + let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); + let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty, true); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + + let alloca_op = Operation::new( + &mut ctx, + mir::MirAllocaOp::get_concrete_op_info(), + vec![mir_ptr_ty.into()], + vec![], + vec![], + 0, + ); + let provenance = llvm_export::ops::LocalMemoryProvenanceAttr { + local_index: 3, + size_bytes: 16, + binding_name: "scratch".into(), + type_name: "[u32; 4]".into(), + }; + llvm_export::ops::set_local_memory_provenance(&mut ctx, alloca_op, provenance.clone()); + alloca_op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + let alloca = find_first::(&ctx, &body).unwrap(); + let copied = llvm_export::ops::local_memory_provenance(&ctx, alloca.get_operation()).unwrap(); + assert_eq!(copied, provenance); +} + +#[test] +fn mem2reg_salvages_tagged_alloca_into_mir_dbg_value() { + let mut ctx = make_ctx(); + let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); + let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty, true); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![i32_ty], vec![i32_ty]); + let arg = block.deref(&ctx).get_argument(0); + + let alloca_op = Operation::new( + &mut ctx, + mir::MirAllocaOp::get_concrete_op_info(), + vec![mir_ptr_ty.into()], + vec![], + vec![], + 0, + ); + let decl_loc = src_location(&mut ctx, "kernel.rs", 12, 9); + alloca_op.deref_mut(&ctx).set_loc(decl_loc.clone()); + llvm::set_debug_local_variable( + &mut ctx, + alloca_op, + llvm::DebugLocalVariableInfo { + name: "x".to_string(), + argument_index: Some(1), + ty: llvm::DebugLocalTypeKind::Basic { + name: "i32".to_string(), + size_bits: 32, + encoding: "DW_ATE_signed", + }, + }, + ); + llvm::set_debug_local_source_scope(&mut ctx, alloca_op, 9); + alloca_op.insert_at_back(block, &ctx); + let slot = alloca_op.deref(&ctx).get_result(0); + + let store_op = Operation::new( + &mut ctx, + mir::MirStoreOp::get_concrete_op_info(), + vec![], + vec![slot, arg], + vec![], + 0, + ); + store_op.insert_at_back(block, &ctx); + + let load_op = Operation::new( + &mut ctx, + mir::MirLoadOp::get_concrete_op_info(), + vec![i32_ty], + vec![slot], + vec![], + 0, + ); + load_op.insert_at_back(block, &ctx); + let loaded = load_op.deref(&ctx).get_result(0); + append_mir_return(&mut ctx, block, vec![loaded]); + + let mut analyses = pliron::pass::AnalysisManager::default(); + pliron::opts::mem2reg::mem2reg(module_ptr, &mut ctx, &mut analyses) + .expect("mem2reg should promote the local slot"); + + let blocks = vec![block]; + assert_eq!(count_ops::(&ctx, &blocks), 0); + assert_eq!(count_ops::(&ctx, &blocks), 0); + assert_eq!(count_ops::(&ctx, &blocks), 0); + + let dbg_values = find_all::(&ctx, &blocks); + assert!( + !dbg_values.is_empty(), + "mem2reg should leave value-based debug records for promoted locals" + ); + let info = llvm::debug_local_variable(&ctx, dbg_values[0].get_operation()) + .expect("mir.dbg_value should carry the promoted local metadata"); + assert_eq!(info.name, "x"); + assert_eq!(info.argument_index, Some(1)); + assert_eq!( + llvm::debug_local_source_scope(&ctx, dbg_values[0].get_operation()), + Some(9), + "mem2reg salvage should keep the local's MIR source-scope owner" + ); + assert_eq!( + dbg_values[0].get_operation().deref(&ctx).loc(), + decl_loc, + "debug records for source-less promoted ops should fall back to the local declaration" + ); +} + +#[test] +fn mem2reg_salvages_fragment_only_alloca_into_mir_dbg_value() { + let mut ctx = make_ctx(); + let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); + let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, i32_ty, true); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![i32_ty], vec![i32_ty]); + let arg = block.deref(&ctx).get_argument(0); + + let alloca_op = Operation::new( + &mut ctx, + mir::MirAllocaOp::get_concrete_op_info(), + vec![mir_ptr_ty.into()], + vec![], + vec![], + 0, + ); + let alloca_loc = src_location(&mut ctx, "kernel.rs", 20, 9); + alloca_op.deref_mut(&ctx).set_loc(alloca_loc); + llvm::set_debug_fragment_variables( + &mut ctx, + alloca_op, + &[llvm::DebugFragmentVariableInfo { + variable: llvm::DebugLocalVariableInfo { + name: "pair".to_string(), + argument_index: None, + ty: llvm::DebugLocalTypeKind::Array { + name: "[u32; 2]".to_string(), + size_bits: 64, + element: Box::new(llvm::DebugLocalTypeKind::Basic { + name: "u32".to_string(), + size_bits: 32, + encoding: "DW_ATE_unsigned", + }), + count: 2, + }, + }, + fragment: llvm::DebugFragment { + offset_bits: 32, + size_bits: 32, + }, + source_scope: Some(9), + declaration: Some(llvm::DebugSourcePosition { + file: PathBuf::from("kernel.rs"), + line: 20, + column: 9, + }), + }], + ); + alloca_op.insert_at_back(block, &ctx); + let slot = alloca_op.deref(&ctx).get_result(0); + + let store_op = Operation::new( + &mut ctx, + mir::MirStoreOp::get_concrete_op_info(), + vec![], + vec![slot, arg], + vec![], + 0, + ); + store_op.insert_at_back(block, &ctx); + + let load_op = Operation::new( + &mut ctx, + mir::MirLoadOp::get_concrete_op_info(), + vec![i32_ty], + vec![slot], + vec![], + 0, + ); + load_op.insert_at_back(block, &ctx); + let loaded = load_op.deref(&ctx).get_result(0); + append_mir_return(&mut ctx, block, vec![loaded]); + + let mut analyses = pliron::pass::AnalysisManager::default(); + pliron::opts::mem2reg::mem2reg(module_ptr, &mut ctx, &mut analyses) + .expect("mem2reg should promote fragment storage"); + + let dbg_values = find_all::(&ctx, &[block]); + assert!( + !dbg_values.is_empty(), + "fragment-only storage should still produce mir.dbg_value salvage" + ); + let fragments = llvm::debug_fragment_variables(&ctx, dbg_values[0].get_operation()); + assert_eq!(fragments.len(), 1); + assert_eq!(fragments[0].variable.name, "pair"); + assert_eq!(fragments[0].fragment.offset_bits, 32); + assert_eq!(fragments[0].fragment.size_bits, 32); +} diff --git a/crates/mir-lower/src/convert/ops/memory/tests/device_global.rs b/crates/mir-lower/src/convert/ops/memory/tests/device_global.rs new file mode 100644 index 0000000000..0d1a236fc6 --- /dev/null +++ b/crates/mir-lower/src/convert/ops/memory/tests/device_global.rs @@ -0,0 +1,555 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#![allow(clippy::disallowed_methods)] + +use super::*; + +pub(super) fn append_global_alloc( + ctx: &mut Context, + block: Ptr, + global_key: &str, + constant: bool, +) -> Ptr { + let i32_ty: TypeHandle = IntegerType::get(ctx, 32, Signedness::Signless).into(); + append_global_alloc_typed(ctx, block, global_key, constant, !constant, i32_ty) +} + +fn append_global_alloc_typed( + ctx: &mut Context, + block: Ptr, + global_key: &str, + constant: bool, + is_mutable: bool, + global_ty: TypeHandle, +) -> Ptr { + let result_ty = if constant { + MirPtrType::get_constant(ctx, global_ty, is_mutable) + } else { + MirPtrType::get_global(ctx, global_ty, is_mutable) + }; + let op = Operation::new( + ctx, + mir::MirGlobalAllocOp::get_concrete_op_info(), + vec![result_ty.into()], + vec![], + vec![], + 0, + ); + let alloc = mir::MirGlobalAllocOp::new(op); + alloc.set_attr_global_type(ctx, TypeAttr::new(global_ty)); + alloc.set_attr_global_key(ctx, StringAttr::new(global_key.to_string())); + op.insert_at_back(block, ctx); + op +} + +#[test] +fn convert_global_alloc_deduplicates_matching_storage_across_pointer_mutability() { + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); + append_global_alloc_typed(&mut ctx, block, "same-global", false, false, i32_ty); + append_global_alloc_typed(&mut ctx, block, "same-global", false, true, i32_ty); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr) + .expect("pointer-carrier mutability does not change physical global storage"); + + let top = module_top_block(&ctx, module_ptr); + let globals = top + .deref(&ctx) + .iter(&ctx) + .filter_map(|op| Operation::get_op::(op, &ctx)) + .filter(|global| global.source_global_key(&ctx).as_deref() == Some("same-global")) + .count(); + assert_eq!(globals, 1); + assert_eq!( + count_ops::(&ctx, &kernel_blocks(&ctx, module_ptr)), + 2 + ); +} + +#[test] +fn convert_global_alloc_rejects_conflicting_key_declarations() { + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + let byte_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Signless).into(); + let word_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); + let empty_words: TypeHandle = MirArrayType::get(&mut ctx, word_ty, 0).into(); + + let byte = + append_global_alloc_typed(&mut ctx, block, "conflicting-global", false, false, byte_ty); + let byte = mir::MirGlobalAllocOp::new(byte); + byte.set_alignment_value(&mut ctx, 1); + byte.mark_immutable(&mut ctx); + + let words = append_global_alloc_typed( + &mut ctx, + block, + "conflicting-global", + false, + false, + empty_words, + ); + let words = mir::MirGlobalAllocOp::new(words); + words.set_alignment_value(&mut ctx, 4); + words.mark_immutable(&mut ctx); + append_mir_return(&mut ctx, block, vec![]); + + let error = crate::lower_mir_to_llvm(&mut ctx, module_ptr) + .expect_err("one global_key must not name incompatible type/alignment storage"); + assert!( + error.to_string().contains("incompatible declaration"), + "{error}" + ); +} + +#[test] +fn device_global_declaration_identity_covers_every_storage_property() { + let ctx = make_ctx(); + let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); + let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signless).into(); + let base = DeviceGlobalDeclaration { + mir_type: i32_ty, + alignment: 4, + addr_space: llvm_addr::GLOBAL, + initializer_hex: Some("00000000".to_string()), + initializer_relocations: Some("reloc-a".to_string()), + debug_info: None, + immutable: true, + }; + + let mut changed = base.clone(); + changed.mir_type = i64_ty; + assert!(base != changed); + changed = base.clone(); + changed.alignment = 8; + assert!(base != changed); + changed = base.clone(); + changed.addr_space = llvm_addr::CONSTANT; + assert!(base != changed); + changed = base.clone(); + changed.initializer_hex = Some("01000000".to_string()); + assert!(base != changed); + changed = base.clone(); + changed.initializer_relocations = Some("reloc-b".to_string()); + assert!(base != changed); + changed = base.clone(); + changed.debug_info = Some(llvm::DebugGlobalVariableInfo { + name: "COUNTER".to_string(), + namespace: vec!["my_crate".to_string()], + ty: llvm::DebugLocalTypeKind::Basic { + name: "u32".to_string(), + size_bits: 32, + encoding: "DW_ATE_unsigned", + }, + declaration: llvm::DebugSourcePosition { + file: PathBuf::from("/tmp/global.rs"), + line: 7, + column: 1, + }, + is_local_to_unit: true, + is_function_local: false, + }); + assert!(base != changed); + changed = base.clone(); + changed.immutable = false; + assert!(base != changed); +} + +#[test] +fn convert_global_alloc_places_in_global_or_constant_addrspace() { + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + append_global_alloc(&mut ctx, block, "ordinary_static", false); + let constant_key = dialect_mir::ops::encode_rust_static_global_key("_ZN7my_mod3KEYE"); + append_global_alloc(&mut ctx, block, &constant_key, true); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let top = module_top_block(&ctx, module_ptr); + let globals: Vec<_> = top + .deref(&ctx) + .iter(&ctx) + .filter_map(|op| Operation::get_op::(op, &ctx)) + .collect(); + let global_addr_global = globals + .iter() + .find(|g| g.address_space(&ctx) == llvm_addr::GLOBAL) + .expect("expected one global in addrspace(1)"); + let global_addr_const = globals + .iter() + .find(|g| g.address_space(&ctx) == llvm_addr::CONSTANT) + .expect("expected one global in addrspace(4)"); + + // Constant-memory globals reuse the Rust mangled name so host code can + // resolve them by name via `cuModuleGetGlobal`; ordinary globals get + // a counter-suffixed `__device_global_N`. + assert_eq!( + global_addr_const.get_symbol_name(&ctx).to_string(), + "_ZN7my_mod3KEYE", + "constant globals must keep the raw rustc symbol payload as their name" + ); + assert!( + global_addr_global + .get_symbol_name(&ctx) + .to_string() + .starts_with("__device_global_"), + "ordinary device globals get the __device_global_ prefix" + ); +} + +#[test] +fn ordinary_global_debug_info_survives_lowering_but_constant_memory_stays_out_of_scope() { + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + let ordinary = append_global_alloc(&mut ctx, block, "crate::GLOBAL_COUNTER", false); + let constant_key = dialect_mir::ops::encode_rust_static_global_key("_ZN5crate5COEFFE"); + let constant = append_global_alloc(&mut ctx, block, &constant_key, true); + let info = llvm::DebugGlobalVariableInfo { + name: "GLOBAL_COUNTER".to_string(), + namespace: vec!["crate".to_string(), "state".to_string()], + ty: llvm::DebugLocalTypeKind::Basic { + name: "u32".to_string(), + size_bits: 32, + encoding: "DW_ATE_unsigned", + }, + declaration: llvm::DebugSourcePosition { + file: PathBuf::from("/tmp/global.rs"), + line: 7, + column: 1, + }, + is_local_to_unit: true, + is_function_local: false, + }; + llvm::set_debug_global_variable(&mut ctx, ordinary, &info); + // Deliberately tag AS4 too: the AS1-only implementation must not + // accidentally expand its behavior merely because the generic debug + // carrier can be attached to any op. + llvm::set_debug_global_variable(&mut ctx, constant, &info); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let top = module_top_block(&ctx, module_ptr); + let globals: Vec<_> = top + .deref(&ctx) + .iter(&ctx) + .filter_map(|op| Operation::get_op::(op, &ctx)) + .collect(); + let ordinary = globals + .iter() + .find(|global| global.address_space(&ctx) == llvm_addr::GLOBAL) + .expect("ordinary global"); + let constant = globals + .iter() + .find(|global| global.address_space(&ctx) == llvm_addr::CONSTANT) + .expect("constant global"); + + assert_eq!( + llvm::debug_global_variable(&ctx, ordinary.get_operation()), + Some(info), + "source identity and semantic type must survive MIR-to-LLVM lowering" + ); + assert!( + llvm::debug_global_variable(&ctx, constant.get_operation()).is_none(), + "AS4 debug metadata is a separate feature and must not leak into this AS1 change" + ); +} + +#[test] +fn immutable_marking_survives_lowering_and_is_not_assumed() { + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + + // Two ordinary addrspace(1) globals, distinguished by their source key. + // Only the promoted one claims immutability; the plain static must not + // acquire it, or the exporter would write `constant` for storage the + // host can still overwrite by symbol. + let promoted = append_global_alloc(&mut ctx, block, "promoted_table", false); + mir::MirGlobalAllocOp::new(promoted).mark_immutable(&mut ctx); + append_global_alloc(&mut ctx, block, "plain_static", false); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let top = module_top_block(&ctx, module_ptr); + let globals: Vec = top + .deref(&ctx) + .iter(&ctx) + .filter_map(|op| Operation::get_op::(op, &ctx)) + .collect(); + let by_key = |key: &str| -> llvm::GlobalOp { + *globals + .iter() + .find(|g| g.source_global_key(&ctx).as_deref() == Some(key)) + .unwrap_or_else(|| panic!("no lowered global carries source key {key}")) + }; + + assert!( + by_key("promoted_table").is_immutable(&ctx), + "a global marked immutable in MIR must stay immutable through lowering" + ); + assert!( + !by_key("plain_static").is_immutable(&ctx), + "lowering must not infer immutability; only the promoted-constant \ + sites may claim it" + ); +} + +#[test] +fn convert_global_alloc_rejects_conflicting_immutability() { + // A plain static and a promoted constant that end up sharing one key + // differ only in the immutable flag; that alone must fail closed + // instead of silently reusing the first allocation's storage. + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + append_global_alloc(&mut ctx, block, "collision", false); + let promoted = append_global_alloc(&mut ctx, block, "collision", false); + mir::MirGlobalAllocOp::new(promoted).mark_immutable(&mut ctx); + append_mir_return(&mut ctx, block, vec![]); + + let error = crate::lower_mir_to_llvm(&mut ctx, module_ptr) + .expect_err("a shared string with different origins must fail closed"); + assert!( + error.to_string().contains("incompatible declaration"), + "{error}" + ); +} + +#[test] +fn initialized_global_lowers_to_byte_storage() { + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + let op = append_global_alloc(&mut ctx, block, "nan_payload", false); + let alloc = mir::MirGlobalAllocOp::new(op); + alloc.set_alignment_value(&mut ctx, 4); + let initializer_key: Identifier = "global_initializer_hex".try_into().unwrap(); + op.deref_mut(&ctx) + .attributes + .set(initializer_key, StringAttr::new("3412c07f".to_string())); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let top = module_top_block(&ctx, module_ptr); + let global = top + .deref(&ctx) + .iter(&ctx) + .find_map(|op| Operation::get_op::(op, &ctx)) + .expect("expected lowered device global"); + let global_ty = global.get_type(&ctx); + let global_ty_ref = global_ty.deref(&ctx); + let array_ty = global_ty_ref + .downcast_ref::() + .expect("initialized global must use byte-array storage"); + assert_eq!(array_ty.size(), 4); + let elem_ty = array_ty.elem_type(); + let elem_ty_ref = elem_ty.deref(&ctx); + let elem = elem_ty_ref + .downcast_ref::() + .expect("byte-array element must be an integer"); + assert_eq!(elem.width(), 8); + assert_eq!(global.get_alignment(&ctx), Some(4)); + assert_eq!(global.initializer_hex(&ctx).as_deref(), Some("3412c07f")); +} + +#[test] +fn relocated_global_lowers_to_segmented_storage() { + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + let word_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let mir_global_ty: TypeHandle = MirArrayType::get(&mut ctx, word_ty, 3).into(); + let result_ty = MirPtrType::get_global(&mut ctx, mir_global_ty, false); + let op = Operation::new( + &mut ctx, + mir::MirGlobalAllocOp::get_concrete_op_info(), + vec![result_ty.into()], + vec![], + vec![], + 0, + ); + let alloc = mir::MirGlobalAllocOp::new(op); + alloc.set_attr_global_type(&ctx, TypeAttr::new(mir_global_ty)); + alloc.set_attr_global_key(&ctx, StringAttr::new("reference_table".to_string())); + alloc.set_alignment_value(&mut ctx, 8); + op.deref_mut(&ctx).attributes.set( + "global_initializer_hex".try_into().unwrap(), + StringAttr::new("000000000000000000000000000000000000000000000000".to_string()), + ); + let encoded = + llvm::encode_global_initializer_relocations(&[llvm::GlobalInitializerRelocation { + source_offset: 8, + width_bytes: 8, + target_address_space: llvm_addr::GLOBAL, + target_addend: 4, + target_key: "target_static".to_string(), + }]); + op.deref_mut(&ctx).attributes.set( + "global_initializer_relocations".try_into().unwrap(), + StringAttr::new(encoded.clone()), + ); + op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let top = module_top_block(&ctx, module_ptr); + let global = top + .deref(&ctx) + .iter(&ctx) + .find_map(|op| Operation::get_op::(op, &ctx)) + .expect("expected lowered device global"); + let global_ty = global.get_type(&ctx); + let global_ty_ref = global_ty.deref(&ctx); + let storage = global_ty_ref + .downcast_ref::() + .expect("relocated initializer must use segmented struct storage"); + assert_eq!(storage.num_fields(), 3); + assert_eq!( + global.source_global_key(&ctx).as_deref(), + Some("reference_table") + ); + assert_eq!( + global.initializer_relocations(&ctx).as_deref(), + Some(encoded.as_str()) + ); +} + +#[test] +fn relocated_global_uses_packed_storage_for_unaligned_pointer_slot() { + let mut ctx = make_ctx(); + let encoded = + llvm::encode_global_initializer_relocations(&[llvm::GlobalInitializerRelocation { + source_offset: 1, + width_bytes: 8, + target_address_space: llvm_addr::GLOBAL, + target_addend: 0, + target_key: "target_static".to_string(), + }]); + + let storage = relocated_initializer_storage_type(&mut ctx, 9, 1, &encoded) + .expect("unaligned relocation should use packed storage"); + let storage_ref = storage.deref(&ctx); + let struct_ty = storage_ref + .downcast_ref::() + .expect("relocated initializer must use struct storage"); + assert_eq!(struct_ty.layout(), StructLayout::Packed); + assert_eq!(struct_ty.num_fields(), 2); + assert_eq!( + crate::convert::types::llvm_type_size_align(&ctx, storage), + Some((9, 1)) + ); + + let fields: Vec<_> = struct_ty.fields().collect(); + let literal_ref = fields[0].deref(&ctx); + let literal = literal_ref + .downcast_ref::() + .expect("leading literal span must be a byte array"); + assert_eq!(literal.size(), 1); + let pointer_ref = fields[1].deref(&ctx); + let pointer = pointer_ref + .downcast_ref::() + .expect("relocation slot must be an integer carrier"); + assert_eq!(pointer.width(), 64); +} + +#[test] +fn relocated_global_uses_packed_storage_for_underaligned_allocation() { + let mut ctx = make_ctx(); + let encoded = + llvm::encode_global_initializer_relocations(&[llvm::GlobalInitializerRelocation { + source_offset: 0, + width_bytes: 8, + target_address_space: llvm_addr::GLOBAL, + target_addend: 0, + target_key: "target_static".to_string(), + }]); + + let storage = relocated_initializer_storage_type(&mut ctx, 8, 1, &encoded) + .expect("underaligned allocation should use packed storage"); + let storage_ref = storage.deref(&ctx); + let struct_ty = storage_ref + .downcast_ref::() + .expect("relocated initializer must use struct storage"); + assert_eq!(struct_ty.layout(), StructLayout::Packed); + assert_eq!( + crate::convert::types::llvm_type_size_align(&ctx, storage), + Some((8, 1)) + ); +} + +#[test] +fn relocated_global_keeps_naturally_aligned_storage_unpacked() { + let mut ctx = make_ctx(); + let encoded = + llvm::encode_global_initializer_relocations(&[llvm::GlobalInitializerRelocation { + source_offset: 8, + width_bytes: 8, + target_address_space: llvm_addr::GLOBAL, + target_addend: 0, + target_key: "target_static".to_string(), + }]); + + let storage = relocated_initializer_storage_type(&mut ctx, 16, 8, &encoded) + .expect("aligned relocation should keep ordinary storage"); + let storage_ref = storage.deref(&ctx); + let struct_ty = storage_ref + .downcast_ref::() + .expect("relocated initializer must use struct storage"); + assert_eq!(struct_ty.layout(), StructLayout::Unpacked); +} + +#[test] +fn relocated_global_rejects_overlapping_slots() { + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + let word_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let mir_global_ty: TypeHandle = MirArrayType::get(&mut ctx, word_ty, 2).into(); + let result_ty = MirPtrType::get_global(&mut ctx, mir_global_ty, false); + let op = Operation::new( + &mut ctx, + mir::MirGlobalAllocOp::get_concrete_op_info(), + vec![result_ty.into()], + vec![], + vec![], + 0, + ); + let alloc = mir::MirGlobalAllocOp::new(op); + alloc.set_attr_global_type(&ctx, TypeAttr::new(mir_global_ty)); + alloc.set_attr_global_key(&ctx, StringAttr::new("overlap".to_string())); + alloc.set_alignment_value(&mut ctx, 8); + op.deref_mut(&ctx).attributes.set( + "global_initializer_hex".try_into().unwrap(), + StringAttr::new("00000000000000000000000000000000".to_string()), + ); + let encoded = llvm::encode_global_initializer_relocations(&[ + llvm::GlobalInitializerRelocation { + source_offset: 0, + width_bytes: 8, + target_address_space: llvm_addr::GLOBAL, + target_addend: 0, + target_key: "a".to_string(), + }, + llvm::GlobalInitializerRelocation { + source_offset: 0, + width_bytes: 8, + target_address_space: llvm_addr::GLOBAL, + target_addend: 0, + target_key: "b".to_string(), + }, + ]); + op.deref_mut(&ctx).attributes.set( + "global_initializer_relocations".try_into().unwrap(), + StringAttr::new(encoded), + ); + op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + let error = crate::lower_mir_to_llvm(&mut ctx, module_ptr) + .expect_err("overlapping relocations must fail closed"); + assert!(error.to_string().contains("overlaps"), "{error}"); +} diff --git a/crates/mir-lower/src/convert/ops/memory/tests/enums.rs b/crates/mir-lower/src/convert/ops/memory/tests/enums.rs new file mode 100644 index 0000000000..5658a39cc1 --- /dev/null +++ b/crates/mir-lower/src/convert/ops/memory/tests/enums.rs @@ -0,0 +1,322 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#![allow(clippy::disallowed_methods)] + +use super::*; + +// ========================================================================= +// Enum layout: converted width per shape + divergent-enum rejection +// ========================================================================= + +use dialect_mir::types::{EnumVariant, MirEnumType}; + +/// Build a Direct-tag `MirEnumType` the way the importer does: +/// unsigned tag of `tag_bits`, plus rustc's `total_size`/`abi_align`. +fn make_enum_ty( + ctx: &mut Context, + name: &str, + tag_bits: u32, + variants: Vec, + total_size: u64, + abi_align: u64, +) -> TypeHandle { + let tag_ty: TypeHandle = IntegerType::get(ctx, tag_bits, Signedness::Unsigned).into(); + // Sequential 0..n discriminants: these layout tests only exercise + // size/width, not value mapping. + let discriminants: Vec = (0..variants.len() as u64).collect(); + MirEnumType::get_with_layout( + ctx, + name.to_string(), + tag_ty, + discriminants, + variants, + 0, // tag at byte 0, like every shape these tests exercise + total_size, + abi_align, + ) + .into() +} + +fn unit_variants(n: usize) -> Vec { + (0..n).map(|i| EnumVariant::unit(format!("V{i}"))).collect() +} + +/// Converted enum allocation size must equal rustc's `total_size` for +/// every memory-faithful tag shape: that size is what GEP strides by. +#[test] +fn enum_conversion_strides_by_rustc_size() { + use crate::convert::types::llvm_type_size_align; + + let mut ctx = make_ctx(); + + // #[repr(u32)] fieldless (issue #118 shape): {i32}, 4 bytes. + let repr_u32 = make_enum_ty(&mut ctx, "ReprU32", 32, unit_variants(4), 4, 4); + let conv = convert_type(&mut ctx, repr_u32).unwrap(); + assert_eq!( + llvm_type_size_align(&ctx, conv), + Some((4, 4)), + "repr(u32) tag" + ); + + // #[repr(usize)] fieldless: {i64}, 8 bytes. + let repr_usize = make_enum_ty(&mut ctx, "ReprUsize", 64, unit_variants(4), 8, 8); + let conv = convert_type(&mut ctx, repr_usize).unwrap(); + assert_eq!( + llvm_type_size_align(&ctx, conv), + Some((8, 8)), + "repr(usize) tag" + ); + + // repr(align(8)) with an i8 tag: the byte claims alone would give + // alignment 1, so the slot map raises the storage alignment with a + // zero-length anchor field. Size and alignment must both match + // rustc, or arrays and by-value ABI uses would be unsound. + let padded = make_enum_ty(&mut ctx, "Padded", 8, unit_variants(2), 8, 8); + let conv = convert_type(&mut ctx, padded).unwrap(); + assert_eq!( + llvm_type_size_align(&ctx, conv), + Some((8, 8)), + "repr(align(8)) i8 tag" + ); + + // u8 tag + i64 payload, rustc size 16: the slot map places the + // payload at its rustc byte offset 8 behind an explicit + // [7 x i8] filler, making the layout datalayout-independent. + let i64_payload: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let payload = make_enum_ty( + &mut ctx, + "OnePayload", + 8, + vec![ + EnumVariant::new_with_layout("A".to_string(), vec![i64_payload], vec![8], vec![8]), + EnumVariant::unit("B".to_string()), + ], + 16, + 8, + ); + let conv = convert_type(&mut ctx, payload).unwrap(); + let (size, _align) = llvm_type_size_align(&ctx, conv).unwrap(); + assert_eq!(size, 16, "natural layout matches rustc size, no pad"); + let conv_ref = conv.deref(&ctx); + let struct_ty = conv_ref + .downcast_ref::() + .expect("converted enum is a struct"); + assert_eq!( + struct_ty.fields().count(), + 3, + "{{tag, [7 x i8] filler, payload}}: explicit filler to byte 8" + ); +} + +/// Multi-payload enum: variants overlap in Rust, and identical +/// (offset, converted type) payloads share one typed slot, so the +/// converted struct is byte-identical to rustc's layout AND every +/// access stays pure SSA (no spill). +#[test] +fn multi_payload_enum_shares_payload_slot() { + use crate::convert::types::{build_enum_slot_map, llvm_type_size_align}; + + let mut ctx = make_ctx(); + let e = make_multi_payload_enum_ty(&mut ctx); + let map = build_enum_slot_map(&mut ctx, e).unwrap(); + assert_eq!(map.carrier_slot, Some(0)); + assert_eq!( + map.field_slots, + vec![Some(1), Some(1)], + "A.0 and B.0 overlap at byte 4 with the same type: one shared slot" + ); + assert_eq!( + llvm_type_size_align(&ctx, map.llvm_struct_ty), + Some((8, 4)), + "byte-identical to rustc's 8-byte layout, not the 12-byte concat" + ); +} + +/// rustc may place the tag AFTER payload bytes; the slot map must +/// follow the recorded tag_offset, never assume slot 0. +#[test] +fn enum_slot_map_tag_not_at_zero() { + use crate::convert::types::{build_enum_slot_map, llvm_type_size_align}; + + let mut ctx = make_ctx(); + let u64_a: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let u64_b: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let tag_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + // enum F { A(u64), B(u64) }: payloads share byte 0, tag at byte 8. + let ty: TypeHandle = MirEnumType::get_with_layout( + &mut ctx, + "TagAtEight".to_string(), + tag_ty, + vec![0, 1], + vec![ + EnumVariant::new_with_layout("A".to_string(), vec![u64_a], vec![0], vec![8]), + EnumVariant::new_with_layout("B".to_string(), vec![u64_b], vec![0], vec![8]), + ], + 8, + 16, + 8, + ) + .into(); + let map = build_enum_slot_map(&mut ctx, ty).unwrap(); + assert_eq!( + map.field_slots, + vec![Some(0), Some(0)], + "payloads share the first slot" + ); + assert_eq!( + map.carrier_slot, + Some(1), + "tag claims its own slot at byte 8" + ); + let (size, _align) = llvm_type_size_align(&ctx, map.llvm_struct_ty).unwrap(); + assert_eq!(size, 16, "{{ i64, i8, [7 x i8] }}"); +} + +/// Multi-payload enum whose variants overlap in Rust: both payloads use +/// bytes 4..8 after the direct `u32` tag, so the complete value is eight +/// bytes. Mimics `#[repr(u32)] enum E { A(u32), B(u32) }`. +fn make_multi_payload_enum_ty(ctx: &mut Context) -> TypeHandle { + let i32_a: TypeHandle = IntegerType::get(ctx, 32, Signedness::Unsigned).into(); + let i32_b: TypeHandle = IntegerType::get(ctx, 32, Signedness::Unsigned).into(); + make_enum_ty( + ctx, + "MultiPayload", + 32, + vec![ + EnumVariant::new_with_layout("A".to_string(), vec![i32_a], vec![4], vec![4]), + EnumVariant::new_with_layout("B".to_string(), vec![i32_b], vec![4], vec![4]), + ], + 8, + 4, + ) +} + +/// Device-local GEP + load must use the same exact eight-byte rustc layout +/// as every other enum path. This protects pointer stride in issue #131's +/// in-kernel `[E; 4]` arrays; the sibling kernel-parameter test proves the +/// same representation is also valid at the host/device boundary. +#[test] +fn device_local_multi_payload_enum_gep_and_load_lower() { + let mut ctx = make_ctx(); + let enum_ty = make_multi_payload_enum_ty(&mut ctx); + let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signless).into(); + let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, enum_ty, true); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![mir_ptr_ty.into(), i64_ty], vec![]); + let ptr_val = block.deref(&ctx).get_argument(0); + let off_val = block.deref(&ctx).get_argument(1); + + let off_op = Operation::new( + &mut ctx, + mir::MirPtrOffsetOp::get_concrete_op_info(), + vec![mir_ptr_ty.into()], + vec![ptr_val, off_val], + vec![], + 0, + ); + off_op.insert_at_back(block, &ctx); + let elem_ptr = off_op.deref(&ctx).get_result(0); + + let load_op = Operation::new( + &mut ctx, + mir::MirLoadOp::get_concrete_op_info(), + vec![enum_ty], + vec![elem_ptr], + vec![], + 0, + ); + load_op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr) + .expect("device-local rustc-layout enum GEP + load must lower"); +} + +/// A kernel parameter may carry this enum because the device slot map is +/// byte-identical to rustc's overlapped host layout. +#[test] +fn kernel_param_accepts_multi_payload_enum() { + use pliron::builtin::attributes::StringAttr; + + let mut ctx = make_ctx(); + let enum_ty = make_multi_payload_enum_ty(&mut ctx); + let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, enum_ty, false); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![mir_ptr_ty.into()], vec![]); + append_mir_return(&mut ctx, block, vec![]); + + // Mark the function as a GPU kernel the way the importer does. + { + let module_block = module_ptr + .deref(&ctx) + .get_region(0) + .deref(&ctx) + .iter(&ctx) + .next() + .unwrap(); + let func_op = module_block.deref(&ctx).iter(&ctx).next().unwrap(); + let kernel_attr = StringAttr::new("true".to_string()); + let key: pliron::identifier::Identifier = "gpu_kernel".try_into().unwrap(); + func_op.deref_mut(&ctx).attributes.set(key, kernel_attr); + } + + // The slot map lowers MultiPayload byte-identically to rustc's + // layout ({ i32, i32 }, 8 bytes), so the kernel ABI accepts it. + crate::lower_mir_to_llvm(&mut ctx, module_ptr) + .expect("multi-payload enum kernel param must lower"); +} + +/// A legacy enum without physical rustc metadata must be rejected at the +/// kernel boundary. Importer-produced niche layouts are now known and are +/// accepted; this fixture deliberately constructs `Unknown` metadata. +#[test] +fn kernel_param_rejects_enum_with_unknown_layout() { + use pliron::builtin::attributes::StringAttr; + + let mut ctx = make_ctx(); + // MirEnumType::get with no layout is the legacy Unknown form. + let i32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); + let pointee = MirPtrType::get_generic(&mut ctx, i32_ty, false); + let tag_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let niched: TypeHandle = MirEnumType::get( + &mut ctx, + "Option".to_string(), + tag_ty, + vec![0, 1], + vec![ + EnumVariant::unit("None".to_string()), + EnumVariant::new("Some".to_string(), vec![pointee.into()]), + ], + ) + .into(); + let mir_ptr_ty = MirPtrType::get_generic(&mut ctx, niched, false); + + let (module_ptr, block) = build_kernel(&mut ctx, vec![mir_ptr_ty.into()], vec![]); + append_mir_return(&mut ctx, block, vec![]); + + { + let module_block = module_ptr + .deref(&ctx) + .get_region(0) + .deref(&ctx) + .iter(&ctx) + .next() + .unwrap(); + let func_op = module_block.deref(&ctx).iter(&ctx).next().unwrap(); + let kernel_attr = StringAttr::new("true".to_string()); + let key: pliron::identifier::Identifier = "gpu_kernel".try_into().unwrap(); + func_op.deref_mut(&ctx).attributes.set(key, kernel_attr); + } + + let err = crate::lower_mir_to_llvm(&mut ctx, module_ptr) + .expect_err("unknown-layout enum kernel param must be rejected"); + let msg = format!("{err}"); + assert!( + msg.contains("Option") && msg.contains("kernel boundary"), + "error must name the enum and the kernel boundary, got: {msg}" + ); + assert!(msg.contains("unknown physical rustc layout")); +} diff --git a/crates/mir-lower/src/convert/ops/memory/tests/mod.rs b/crates/mir-lower/src/convert/ops/memory/tests/mod.rs new file mode 100644 index 0000000000..616fe704f7 --- /dev/null +++ b/crates/mir-lower/src/convert/ops/memory/tests/mod.rs @@ -0,0 +1,83 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! End-to-end lowering tests for `dialect-mir` memory ops. +//! +//! The `convert_*` functions in this file take a live +//! `DialectConversionRewriter`, which is owned by pliron's conversion +//! driver and not constructible standalone. So each test builds a small +//! MIR module, runs the full `lower_mir_to_llvm` pass on it, and asserts +//! the lowered module contains the expected `dialect-llvm` shape — same +//! pattern as `tests/lowering_test.rs`. + +// Tests build kinded fixture types directly; production minting lives in mir-importer's facts.rs. +#![allow(clippy::disallowed_methods)] + +mod access; +mod alignment; +mod debug_provenance; +mod device_global; +mod enums; +mod shared; + +use crate::context::DeviceGlobalDeclaration; +use crate::convert::ops::test_util::*; +use crate::convert::types::convert_type; +use dialect_mir::attributes::MirPointerKindAuthorityAttr; +use dialect_mir::ops as mir; +use dialect_mir::types::{ + MirArrayType, MirPointerKind, MirPtrType, MirStructType, MirTupleType, MirUnionType, +}; +use llvm_export::op_interfaces::PointerTypeResult; +use llvm_export::ops as llvm; +use llvm_export::ops::GlobalOpExt; +use llvm_export::types::{ + ArrayType, PointerType, StructLayout, StructType, address_space as llvm_addr, +}; +use pliron::basic_block::BasicBlock; +use pliron::builtin::attributes::{StringAttr, TypeAttr}; +use pliron::builtin::op_interfaces::SymbolOpInterface; +use pliron::builtin::types::{IntegerType, Signedness}; +use pliron::context::{Context, Ptr}; +use pliron::identifier::Identifier; +use pliron::linked_list::ContainsLinkedList; +use pliron::location::{Located, Location, Source}; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::r#type::{TypeHandle, Typed}; +use pliron::utils::apint::APInt; +use std::path::PathBuf; + +use super::device_global::relocated_initializer_storage_type; + +fn ptr_addrspace(ctx: &Context, ty: TypeHandle) -> u32 { + ty.deref(ctx) + .downcast_ref::() + .expect("expected llvm.PointerType") + .address_space() +} + +fn src_location(ctx: &mut Context, file: &str, line: i32, column: i32) -> Location { + Location::SrcPos { + src: Source::new_from_file(ctx, PathBuf::from(file)), + pos: combine::stream::position::SourcePosition { line, column }, + } +} + +fn over_aligned_tuple_ty(ctx: &mut Context) -> TypeHandle { + let byte: TypeHandle = IntegerType::get(ctx, 8, Signedness::Unsigned).into(); + let marker: TypeHandle = MirStructType::get_with_full_layout( + ctx, + "Align32".into(), + vec![], + vec![], + vec![], + vec![], + 0, + 32, + ) + .into(); + MirTupleType::get_with_layout(ctx, vec![marker, byte], vec![0, 1], vec![0, 0], 32, 32).into() +} diff --git a/crates/mir-lower/src/convert/ops/memory/tests/shared.rs b/crates/mir-lower/src/convert/ops/memory/tests/shared.rs new file mode 100644 index 0000000000..9cb4fa0cef --- /dev/null +++ b/crates/mir-lower/src/convert/ops/memory/tests/shared.rs @@ -0,0 +1,560 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#![allow(clippy::disallowed_methods)] + +use super::device_global::append_global_alloc; +use super::*; + +/// Build a `mir.shared_alloc` returning `MirPtrType` of +/// length `size`, with the given alloc_key, and append it to `block`. +fn append_shared_alloc( + ctx: &mut Context, + block: Ptr, + alloc_key: &str, + size: u64, +) -> Ptr { + append_shared_alloc_named(ctx, block, alloc_key, size, None) +} + +/// As [`append_shared_alloc`], additionally carrying the Rust path of the +/// `static` the allocation came from. +fn append_shared_alloc_named( + ctx: &mut Context, + block: Ptr, + alloc_key: &str, + size: u64, + source_name: Option<&str>, +) -> Ptr { + let i32_ty: TypeHandle = IntegerType::get(ctx, 32, Signedness::Signless).into(); + append_shared_alloc_typed(ctx, block, alloc_key, i32_ty, size, source_name, 0) +} + +#[allow(clippy::too_many_arguments)] +fn append_shared_alloc_typed( + ctx: &mut Context, + block: Ptr, + alloc_key: &str, + element_type: TypeHandle, + size: u64, + source_name: Option<&str>, + alignment: u64, +) -> Ptr { + use pliron::builtin::attributes::IntegerAttr; + use pliron::utils::apint::APInt; + + let i32_ty = element_type; + let result_ty = MirPtrType::get_shared(ctx, i32_ty, true); + let op = Operation::new( + ctx, + mir::MirSharedAllocOp::get_concrete_op_info(), + vec![result_ty.into()], + vec![], + vec![], + 0, + ); + let alloc = mir::MirSharedAllocOp::new(op); + alloc.set_attr_elem_type(ctx, TypeAttr::new(i32_ty)); + let size_attr = IntegerAttr::new( + IntegerType::get(ctx, 64, Signedness::Signless), + APInt::from_u64(size, std::num::NonZeroUsize::new(64).unwrap()), + ); + alloc.set_attr_size(ctx, size_attr); + alloc.set_attr_alloc_key(ctx, StringAttr::new(alloc_key.to_string())); + if let Some(source_name) = source_name { + alloc.set_attr_source_name(ctx, StringAttr::new(source_name.to_string())); + } + if alignment != 0 { + alloc.set_alignment_value(ctx, alignment); + } + op.insert_at_back(block, ctx); + op +} + +fn shared_array_debug_info(count: u64) -> llvm::DebugGlobalVariableInfo { + llvm::DebugGlobalVariableInfo { + name: "TILE".to_string(), + namespace: vec!["fixture".to_string(), "kernel".to_string()], + ty: llvm::DebugLocalTypeKind::Array { + name: format!("[i32; {count}]"), + size_bits: count * 32, + element: Box::new(llvm::DebugLocalTypeKind::Basic { + name: "i32".to_string(), + size_bits: 32, + encoding: "DW_ATE_signed", + }), + count, + }, + declaration: llvm::DebugSourcePosition { + file: PathBuf::from("/tmp/shared.rs"), + line: 7, + column: 5, + }, + is_local_to_unit: true, + is_function_local: true, + } +} + +fn shared_static_key(symbol: &str) -> String { + dialect_mir::ops::encode_rust_static_global_key(symbol) +} + +#[test] +fn convert_shared_alloc_creates_global_in_addrspace_3() { + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + append_shared_alloc(&mut ctx, block, "k1", 64); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + // Global lives at module level; addressof lives in the function body. + let top = module_top_block(&ctx, module_ptr); + let global = top + .deref(&ctx) + .iter(&ctx) + .find_map(|op| Operation::get_op::(op, &ctx)) + .expect("expected an llvm.global for the shared allocation"); + assert_eq!( + global.address_space(&ctx), + llvm_addr::SHARED, + "shared_alloc global must live in addrspace 3" + ); + assert!( + global + .get_symbol_name(&ctx) + .to_string() + .starts_with("__shared_mem_"), + "shared global should have __shared_mem_ prefix" + ); + + let body = kernel_blocks(&ctx, module_ptr); + let addrof = find_first::(&ctx, &body).expect("expected an llvm.addressof"); + assert_eq!( + ptr_addrspace( + &ctx, + addrof + .get_operation() + .deref(&ctx) + .get_result(0) + .get_type(&ctx) + ), + llvm_addr::SHARED, + ); +} + +#[test] +fn convert_shared_alloc_deduplicates_by_alloc_key() { + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + // Two allocations sharing the same alloc_key — they must collapse to + // a single underlying global (this is what enables a single `static` + // referenced from multiple sites to land in one shared region). + append_shared_alloc(&mut ctx, block, "same-key", 64); + append_shared_alloc(&mut ctx, block, "same-key", 64); + // A third with a different key must NOT dedupe with them. + append_shared_alloc(&mut ctx, block, "other-key", 32); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let top = module_top_block(&ctx, module_ptr); + let shared_globals = top + .deref(&ctx) + .iter(&ctx) + .filter_map(|op| Operation::get_op::(op, &ctx)) + .filter(|g| g.address_space(&ctx) == llvm_addr::SHARED) + .count(); + assert_eq!( + shared_globals, 2, + "two distinct alloc_keys must produce two globals" + ); + + // Each of the three mir.shared_alloc ops becomes one addressof. + let body = kernel_blocks(&ctx, module_ptr); + assert_eq!(count_ops::(&ctx, &body), 3); +} + +#[test] +fn convert_shared_alloc_rejects_conflicting_key_declarations() { + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + append_shared_alloc(&mut ctx, block, "conflicting-key", 1); + let differently_aligned = append_shared_alloc(&mut ctx, block, "conflicting-key", 1); + mir::MirSharedAllocOp::new(differently_aligned).set_alignment_value(&mut ctx, 16); + append_mir_return(&mut ctx, block, vec![]); + + let error = crate::lower_mir_to_llvm(&mut ctx, module_ptr) + .expect_err("one shared alloc_key must not name differently aligned storage"); + assert!( + error.to_string().contains("incompatible declaration"), + "{error}" + ); +} + +#[test] +fn shared_alloc_repeated_debug_identity_is_attached_once() { + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + let key = shared_static_key("_Rfixture4TILE"); + for _ in 0..2 { + let op = + append_shared_alloc_named(&mut ctx, block, &key, 32, Some("fixture::kernel::TILE")); + let info = shared_array_debug_info(32); + llvm::set_debug_global_variable(&mut ctx, op, &info); + llvm::set_debug_global_owner_function(&mut ctx, op, "fixture_kernel"); + } + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + let top = module_top_block(&ctx, module_ptr); + let globals: Vec<_> = top + .deref(&ctx) + .iter(&ctx) + .filter_map(|op| Operation::get_op::(op, &ctx)) + .filter(|global| global.address_space(&ctx) == llvm_addr::SHARED) + .collect(); + assert_eq!( + globals.len(), + 1, + "repeated references must share one global" + ); + assert_eq!( + llvm::debug_global_variable(&ctx, globals[0].get_operation()), + Some(shared_array_debug_info(32)) + ); + assert_eq!( + llvm::debug_global_owner_function(&ctx, globals[0].get_operation()).as_deref(), + Some("fixture_kernel") + ); +} + +/// One function-local shared static materialized in two owning functions +/// (for example a `#[device]` helper inlined into two kernels) is a valid +/// release build, so debug metadata must not turn it into an error. DWARF +/// cannot truthfully scope one AS3 object to two subprograms either, so +/// the divergent identity fails open: the storage is still shared and the +/// debug attachment is dropped. +#[test] +fn shared_alloc_owner_conflict_drops_debug_metadata() { + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + let key = shared_static_key("_Rfixture4TILE"); + for owner in ["fixture_kernel", "other_kernel"] { + let op = + append_shared_alloc_named(&mut ctx, block, &key, 32, Some("fixture::kernel::TILE")); + let info = shared_array_debug_info(32); + llvm::set_debug_global_variable(&mut ctx, op, &info); + llvm::set_debug_global_owner_function(&mut ctx, op, owner); + } + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr) + .expect("divergent debug identity must not fail a build release accepts"); + let top = module_top_block(&ctx, module_ptr); + let globals: Vec<_> = top + .deref(&ctx) + .iter(&ctx) + .filter_map(|op| Operation::get_op::(op, &ctx)) + .filter(|global| global.address_space(&ctx) == llvm_addr::SHARED) + .collect(); + assert_eq!( + globals.len(), + 1, + "the physical storage must still be shared" + ); + assert!(llvm::debug_global_variable(&ctx, globals[0].get_operation()).is_none()); + assert!(llvm::debug_global_owner_function(&ctx, globals[0].get_operation()).is_none()); +} + +#[test] +fn dynamic_extern_then_static_shared_rejects_reserved_key_collision() { + use pliron::builtin::attributes::IntegerAttr; + use pliron::utils::apint::APInt; + + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + let i8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Signless).into(); + let shared_i8 = MirPtrType::get_shared(&mut ctx, i8_ty, true); + + // Lower the dynamic declaration first. Its internal cache key is + // deliberately then reused by a static allocation with the same + // physical type, size, and alignment. Those declarations must still + // be rejected because one is extern storage and one is fixed storage. + let dynamic = Operation::new( + &mut ctx, + mir::MirExternSharedOp::get_concrete_op_info(), + vec![shared_i8.into()], + vec![], + vec![], + 0, + ); + mir::MirExternSharedOp::new(dynamic).set_alignment_value(&mut ctx, 128); + dynamic.insert_at_back(block, &ctx); + + let fixed = Operation::new( + &mut ctx, + mir::MirSharedAllocOp::get_concrete_op_info(), + vec![shared_i8.into()], + vec![], + vec![], + 0, + ); + let fixed_alloc = mir::MirSharedAllocOp::new(fixed); + fixed_alloc.set_attr_elem_type(&ctx, TypeAttr::new(i8_ty)); + fixed_alloc.set_attr_size( + &ctx, + IntegerAttr::new( + IntegerType::get(&ctx, 64, Signedness::Unsigned), + APInt::from_u64(0, std::num::NonZeroUsize::new(64).unwrap()), + ), + ); + fixed_alloc.set_attr_alloc_key( + &ctx, + StringAttr::new("__dynamic_smem_global_created_kernel_func".to_string()), + ); + fixed_alloc.set_alignment_value(&mut ctx, 128); + fixed.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + let error = crate::lower_mir_to_llvm(&mut ctx, module_ptr) + .expect_err("a static allocation must not reuse a dynamic extern declaration"); + let message = error.to_string(); + assert!(message.contains("incompatible declaration"), "{message}"); + assert!( + message.contains("duplicate shared alloc_key"), + "the dynamic declaration must be observed first, got: {message}" + ); +} + +#[test] +fn shared_alloc_mismatched_debug_type_fails_closed_without_failing_lowering() { + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + let key = shared_static_key("_Rfixture4TILE"); + let op = append_shared_alloc_named(&mut ctx, block, &key, 32, Some("fixture::kernel::TILE")); + llvm::set_debug_global_variable(&mut ctx, op, &shared_array_debug_info(31)); + llvm::set_debug_global_owner_function(&mut ctx, op, "fixture_kernel"); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr) + .expect("bad optional metadata must not fail compilation"); + let top = module_top_block(&ctx, module_ptr); + let global = top + .deref(&ctx) + .iter(&ctx) + .find_map(|candidate| Operation::get_op::(candidate, &ctx)) + .expect("shared global"); + assert!(llvm::debug_global_variable(&ctx, global.get_operation()).is_none()); + assert!(llvm::debug_global_owner_function(&ctx, global.get_operation()).is_none()); +} + +#[test] +fn barrier_semantic_struct_matches_single_i64_shared_backing() { + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let key = shared_static_key("_Rfixture7BARRIER"); + let op = append_shared_alloc_typed( + &mut ctx, + block, + &key, + i64_ty, + 1, + Some("fixture::kernel::BARRIER"), + 8, + ); + let info = llvm::DebugGlobalVariableInfo { + name: "BARRIER".to_string(), + namespace: vec!["fixture".to_string(), "kernel".to_string()], + ty: llvm::DebugLocalTypeKind::Struct { + name: "Barrier".to_string(), + size_bits: 64, + members: vec![llvm::DebugTypeMember { + name: "_state".to_string(), + offset_bits: 0, + ty: llvm::DebugLocalTypeKind::Basic { + name: "u64".to_string(), + size_bits: 64, + encoding: "DW_ATE_unsigned", + }, + }], + }, + declaration: llvm::DebugSourcePosition { + file: PathBuf::from("/tmp/barrier.rs"), + line: 8, + column: 9, + }, + is_local_to_unit: true, + is_function_local: true, + }; + llvm::set_debug_global_variable(&mut ctx, op, &info); + llvm::set_debug_global_owner_function(&mut ctx, op, "fixture_kernel"); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("Barrier lowering succeeds"); + let top = module_top_block(&ctx, module_ptr); + let global = top + .deref(&ctx) + .iter(&ctx) + .find_map(|candidate| Operation::get_op::(candidate, &ctx)) + .expect("Barrier shared global"); + assert_eq!(global.get_alignment(&ctx), Some(8)); + assert_eq!( + llvm::debug_global_variable(&ctx, global.get_operation()), + Some(info) + ); +} + +/// Collect `(symbol, source_name)` for every shared global in the module. +fn shared_global_source_names( + ctx: &Context, + module_ptr: Ptr, +) -> Vec<(String, Option)> { + use llvm_export::ops::GlobalOpExt; + + let top = module_top_block(ctx, module_ptr); + let mut named: Vec<_> = top + .deref(ctx) + .iter(ctx) + .filter_map(|op| Operation::get_op::(op, ctx)) + .filter(|g| g.address_space(ctx) == llvm_addr::SHARED) + .map(|g| { + ( + g.get_symbol_name(ctx).to_string(), + g.shared_source_name(ctx), + ) + }) + .collect(); + // Globals are inserted at the front of the module block, so iteration + // order is the reverse of creation order. Sort for a stable assertion. + named.sort(); + named +} + +#[test] +fn shared_alloc_source_name_reaches_the_generated_global() { + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + append_shared_alloc_named(&mut ctx, block, "k1", 64, Some("my_kernel::TILE")); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let named = shared_global_source_names(&ctx, module_ptr); + assert_eq!(named.len(), 1, "expected exactly one shared global"); + let (symbol, source_name) = &named[0]; + // The symbol itself must stay anonymous: the whole point of the + // sidecar attribute is that it does not perturb the emitted name. + assert!( + symbol.starts_with("__shared_mem_"), + "the generated symbol must not be renamed, got `{symbol}`" + ); + assert_eq!(source_name.as_deref(), Some("my_kernel::TILE")); +} + +#[test] +fn shared_alloc_without_source_name_leaves_the_global_unlabelled() { + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + append_shared_alloc(&mut ctx, block, "k1", 64); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let named = shared_global_source_names(&ctx, module_ptr); + assert_eq!(named.len(), 1); + assert_eq!(named[0].1, None, "an unnamed allocation must stay unnamed"); +} + +#[test] +fn shared_alloc_source_names_are_per_global_not_shared_across_them() { + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + // Two references to one static dedupe onto a single global, and a + // second static gets its own. Each global must carry its own name — + // the failure this guards is one name leaking onto every allocation. + append_shared_alloc_named(&mut ctx, block, "tile", 64, Some("my_kernel::TILE")); + append_shared_alloc_named(&mut ctx, block, "tile", 64, Some("my_kernel::TILE")); + append_shared_alloc_named(&mut ctx, block, "scratch", 32, Some("my_kernel::SCRATCH")); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let names: Vec<_> = shared_global_source_names(&ctx, module_ptr) + .into_iter() + .map(|(_, source_name)| source_name) + .collect(); + assert_eq!(names.len(), 2, "the shared alloc_key must still dedupe"); + let mut names: Vec<_> = names.into_iter().map(|n| n.expect("named")).collect(); + names.sort(); + assert_eq!(names, vec!["my_kernel::SCRATCH", "my_kernel::TILE"]); +} + +#[test] +fn shared_alloc_source_name_reaches_the_exported_llvm_ir() { + // The end the feature exists for: a consumer holding only the emitted + // artifact can tell which Rust `static` a `__shared_mem_N` block is. + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + append_shared_alloc_named(&mut ctx, block, "tile", 64, Some("my_kernel::TILE")); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let module = Operation::get_op::(module_ptr, &ctx) + .expect("lowered top-level op is a module"); + let ir = llvm_export::export::export_module_to_string(&ctx, &module).expect("export"); + + let comment_index = ir + .find("; shared source: my_kernel::TILE") + .unwrap_or_else(|| panic!("exported IR must name the shared source:\n{ir}")); + let definition_index = ir + .find("__shared_mem_") + .expect("exported IR must declare the shared global"); + assert!( + comment_index < definition_index, + "the source comment must precede the global it describes:\n{ir}" + ); +} + +/// A `__shared_mem_N` or `__device_global_N` index must depend only on +/// the module being lowered, not on how many allocations any OTHER +/// module has already lowered in this process (#706). Before the fix, +/// each `N` came from a `static AtomicUsize` shared across every call in +/// the process, so lowering the second of these two modules would have +/// produced `__shared_mem_1` and `__device_global_1`, not the `_0` names. +#[test] +fn shared_and_device_global_indices_are_per_module_not_process_global() { + for _ in 0..2 { + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + append_shared_alloc(&mut ctx, block, "k", 64); + append_global_alloc(&mut ctx, block, "ordinary_static", false); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let top = module_top_block(&ctx, module_ptr); + let names: Vec = top + .deref(&ctx) + .iter(&ctx) + .filter_map(|op| Operation::get_op::(op, &ctx)) + .map(|g| g.get_symbol_name(&ctx).to_string()) + .collect(); + assert!( + names.iter().any(|n| n == "__shared_mem_0"), + "a module with exactly one shared allocation must always name it \ + __shared_mem_0, regardless of how many other modules already lowered \ + one in this process (got {names:?})" + ); + assert!( + names.iter().any(|n| n == "__device_global_0"), + "a module with exactly one ordinary device global must always name it \ + __device_global_0, regardless of how many other modules already \ + lowered one in this process (got {names:?})" + ); + } +} diff --git a/crates/mir-lower/src/convert/ops/memory/transfer.rs b/crates/mir-lower/src/convert/ops/memory/transfer.rs new file mode 100644 index 0000000000..f2d5367b07 --- /dev/null +++ b/crates/mir-lower/src/convert/ops/memory/transfer.rs @@ -0,0 +1,237 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Conversions for `mir.memcpy` / `mir.memmove` and their shared lowering body. + +use super::common::anyhow_to_pliron; +use crate::convert::types::{convert_type, llvm_type_size_align, mir_element_stride}; +use crate::helpers; +use llvm_export::attributes::IntegerOverflowFlagsAttr; +use llvm_export::op_interfaces::IntBinArithOpWithOverflowFlag; +use llvm_export::ops as llvm; +use llvm_export::types::{FuncType, VoidType}; +use pliron::attribute::AttrObj; +use pliron::builtin::attributes::IntegerAttr; +use pliron::builtin::op_interfaces::CallOpCallable; +use pliron::builtin::types::{IntegerType, Signedness}; +use pliron::context::{Context, Ptr}; +use pliron::identifier::Identifier; +use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; +use pliron::irbuild::inserter::Inserter; +use pliron::irbuild::rewriter::Rewriter; +use pliron::location::Located; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::result::Result; +use pliron::r#type::Typed; +use pliron::utils::apint::APInt; + +/// Convert `mir.memcpy` to the matching `llvm.memcpy.p.p.i`. +/// +/// MIR's count is measured in pointee elements, while LLVM's memcpy intrinsic +/// expects bytes. The pre-conversion destination pointer type still carries the +/// MIR pointee, so use it to scale the count before emitting the call. +/// +/// The intrinsic name is an overload: LLVM encodes each pointer's address +/// space and the length width into it, and its verifier rejects a call whose +/// argument types disagree with the name. Today every pointer reaching a +/// `copy_nonoverlapping` is a Rust raw pointer, which cuda-oxide normalizes to +/// the generic address space (an `addrspacecast` is inserted when the raw +/// pointer is formed), so the operands are always `p0` and `i64`. We still +/// derive the suffix from the real operand types rather than hardcoding +/// `p0.p0.i64`: it matches how every other overloaded intrinsic is named here +/// (`ctpop`, `fptosi.sat`, ...), and it keeps this lowering correct if raw +/// pointers ever start carrying a non-generic address space. +pub(crate) fn convert_memcpy( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + operands_info: &OperandsInfo, +) -> Result<()> { + convert_mem_transfer(ctx, rewriter, op, operands_info, "memcpy") +} + +/// Convert `mir.memmove` to the matching `llvm.memmove.p.p.i`. +/// +/// Identical to [`convert_memcpy`] except it emits the overlap-safe +/// `llvm.memmove` intrinsic. `mir.memmove` backs `core::intrinsics::copy` +/// (`ptr::copy`); `mir.memcpy` backs the non-overlapping variant. +pub(crate) fn convert_memmove( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + operands_info: &OperandsInfo, +) -> Result<()> { + convert_mem_transfer(ctx, rewriter, op, operands_info, "memmove") +} + +/// Shared lowering for `mir.memcpy` / `mir.memmove`. `intrinsic_base` selects +/// the LLVM intrinsic family ("memcpy" or "memmove"); both share the same +/// `(dst, src, len_bytes, isvolatile)` signature and element->byte count scaling. +/// +/// The element type comes from the op's own `elem_type` attribute, stamped +/// at build time from dst's pointer type. Operand type history is not +/// usable here: a kind-only `mir.cast` lowers to a plain value forwarding, +/// history does not follow that edge, and a stale hit would scale the byte +/// count by the wrong element size. +fn convert_mem_transfer( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + _operands_info: &OperandsInfo, + intrinsic_base: &str, +) -> Result<()> { + let operands: Vec<_> = op.deref(ctx).operands().collect(); + let (dst, src, count) = match operands.as_slice() { + [dst, src, count] => (*dst, *src, *count), + _ => { + return pliron::input_err_noloc!( + "{intrinsic_base} operation requires exactly 3 operands" + ); + } + }; + + let pointee = { + let op_ref = op.deref(ctx); + op_ref + .attributes + .get::( + &format!("{intrinsic_base}_elem_type").try_into().unwrap(), + ) + .map(|attr| attr.get_type(ctx)) + .ok_or_else(|| { + pliron::create_error!( + op_ref.loc(), + pliron::result::ErrorKind::VerificationFailed, + pliron::result::StringError(format!( + "mir.{intrinsic_base} missing its elem_type attribute; \ + byte count has no fact to derive from" + )) + ) + })? + }; + // Byte-count policy: exact or error, never guessed. rustc's stride of the + // stamped MIR elem type is the primary fact; the converted LLVM type's + // natural layout is the fallback. + let elem_size = match mir_element_stride(ctx, pointee) { + Some(stride) => stride, + None => { + let elem_ty = convert_type(ctx, pointee).map_err(anyhow_to_pliron)?; + match llvm_type_size_align(ctx, elem_ty) { + Some((size, _)) => size, + None => { + let type_display = pointee.deref(ctx).disp(ctx).to_string(); + return Err(pliron::create_error!( + op.deref(ctx).loc(), + pliron::result::ErrorKind::VerificationFailed, + pliron::result::StringError(format!( + "mir.{intrinsic_base} element type `{type_display}` has no exact \ + byte size; refusing to guess the copy length" + )) + )); + } + } + } + }; + + let bytes = if elem_size == 1 { + count + } else { + let count_ty = count.get_type(ctx); + let bits = count_ty + .deref(ctx) + .downcast_ref::() + .map(|ty| ty.width()) + .unwrap_or(64); + let count_int_ty = IntegerType::get(ctx, bits, Signedness::Signless); + let size_attr: AttrObj = IntegerAttr::new( + count_int_ty, + APInt::from_u64( + elem_size, + std::num::NonZeroUsize::new(bits as usize).unwrap(), + ), + ) + .into(); + let size_const = llvm::ConstantOp::new(ctx, size_attr); + let size_val = size_const.get_operation().deref(ctx).get_result(0); + rewriter.insert_operation(ctx, size_const.get_operation()); + + let flags = IntegerOverflowFlagsAttr::default(); + let mul = llvm::MulOp::new_with_overflow_flag(ctx, count, size_val, flags); + rewriter.insert_operation(ctx, mul.get_operation()); + mul.get_operation().deref(ctx).get_result(0) + }; + + let i1_ty = IntegerType::get(ctx, 1, Signedness::Signless); + let false_attr: AttrObj = IntegerAttr::new( + i1_ty, + APInt::from_u64(0, std::num::NonZeroUsize::new(1).unwrap()), + ) + .into(); + let volatile = llvm::ConstantOp::new(ctx, false_attr); + rewriter.insert_operation(ctx, volatile.get_operation()); + let volatile_val = volatile.get_operation().deref(ctx).get_result(0); + + let void_ty = VoidType::get(ctx); + let func_ty = FuncType::get( + ctx, + void_ty.into(), + vec![ + dst.get_type(ctx), + src.get_type(ctx), + bytes.get_type(ctx), + volatile_val.get_type(ctx), + ], + false, + ); + let parent_block = op.deref(ctx).get_parent_block().ok_or_else(|| { + pliron::create_error!( + op.deref(ctx).loc(), + pliron::result::ErrorKind::VerificationFailed, + pliron::result::StringError(format!("{intrinsic_base} operation has no parent block")) + ) + })?; + // Derive the overload suffix from the real (already type-converted) + // operands so the name can never disagree with the argument types. + let dst_ty = dst.get_type(ctx); + let dst_as = dst_ty + .deref(ctx) + .downcast_ref::() + .map(|pt| pt.address_space()) + .unwrap_or(0); + let src_ty = src.get_type(ctx); + let src_as = src_ty + .deref(ctx) + .downcast_ref::() + .map(|pt| pt.address_space()) + .unwrap_or(0); + let len_bits = bytes + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .map(|t| t.width()) + .unwrap_or(64); + let intrinsic_name = format!("llvm_{intrinsic_base}_p{dst_as}_p{src_as}_i{len_bits}"); + helpers::ensure_intrinsic_declared(ctx, parent_block, &intrinsic_name, func_ty) + .map_err(anyhow_to_pliron)?; + + let callee: Identifier = intrinsic_name.as_str().try_into().map_err(|e| { + pliron::create_error!( + op.deref(ctx).loc(), + pliron::result::ErrorKind::VerificationFailed, + pliron::result::StringError(format!("Invalid memcpy intrinsic name: {e:?}")) + ) + })?; + let call = llvm::CallOp::new( + ctx, + CallOpCallable::Direct(callee), + func_ty, + vec![dst, src, bytes, volatile_val], + ); + crate::convert::preserve_location(ctx, op, call.get_operation()); + rewriter.insert_operation(ctx, call.get_operation()); + rewriter.erase_operation(ctx, op); + Ok(()) +} From 9862c4a738eb0fea1dec41243ffa07ace170a663 Mon Sep 17 00:00:00 2001 From: nihalpasham Date: Fri, 28 Aug 2026 22:12:42 +0530 Subject: [PATCH 08/18] refactor(intrinsics-gen): split model.rs into model/ modules --- crates/cuda-intrinsics-gen/src/model.rs | 4730 ----------------- .../src/model/abi_ledger.rs | 33 + .../cuda-intrinsics-gen/src/model/catalog.rs | 458 ++ .../src/model/contracts/control.rs | 266 + .../src/model/contracts/matrix.rs | 605 +++ .../src/model/contracts/memory.rs | 464 ++ .../src/model/contracts/mod.rs | 20 + .../src/model/contracts/scalar.rs | 511 ++ .../src/model/contracts/tcgen05.rs | 398 ++ .../src/model/contracts/tma.rs | 243 + .../src/model/contracts/warp.rs | 503 ++ crates/cuda-intrinsics-gen/src/model/core.rs | 92 + .../cuda-intrinsics-gen/src/model/evidence.rs | 208 + .../cuda-intrinsics-gen/src/model/imported.rs | 195 + crates/cuda-intrinsics-gen/src/model/mod.rs | 20 + .../cuda-intrinsics-gen/src/model/overlay.rs | 915 ++++ 16 files changed, 4931 insertions(+), 4730 deletions(-) delete mode 100644 crates/cuda-intrinsics-gen/src/model.rs create mode 100644 crates/cuda-intrinsics-gen/src/model/abi_ledger.rs create mode 100644 crates/cuda-intrinsics-gen/src/model/catalog.rs create mode 100644 crates/cuda-intrinsics-gen/src/model/contracts/control.rs create mode 100644 crates/cuda-intrinsics-gen/src/model/contracts/matrix.rs create mode 100644 crates/cuda-intrinsics-gen/src/model/contracts/memory.rs create mode 100644 crates/cuda-intrinsics-gen/src/model/contracts/mod.rs create mode 100644 crates/cuda-intrinsics-gen/src/model/contracts/scalar.rs create mode 100644 crates/cuda-intrinsics-gen/src/model/contracts/tcgen05.rs create mode 100644 crates/cuda-intrinsics-gen/src/model/contracts/tma.rs create mode 100644 crates/cuda-intrinsics-gen/src/model/contracts/warp.rs create mode 100644 crates/cuda-intrinsics-gen/src/model/core.rs create mode 100644 crates/cuda-intrinsics-gen/src/model/evidence.rs create mode 100644 crates/cuda-intrinsics-gen/src/model/imported.rs create mode 100644 crates/cuda-intrinsics-gen/src/model/mod.rs create mode 100644 crates/cuda-intrinsics-gen/src/model/overlay.rs diff --git a/crates/cuda-intrinsics-gen/src/model.rs b/crates/cuda-intrinsics-gen/src/model.rs deleted file mode 100644 index 2a3392ae2e..0000000000 --- a/crates/cuda-intrinsics-gen/src/model.rs +++ /dev/null @@ -1,4730 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -use crate::ptx::InstructionPattern; -use serde::{Deserialize, Serialize}; -use std::{fmt, str::FromStr}; - -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct UpstreamLock { - pub schema: u32, - pub llvm: LockedLlvm, - pub llvm_tblgen: LockedTool, - #[serde(default)] - pub comparison_tools: Vec, - pub dumps: LockedDumps, -} - -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct LockedLlvm { - pub repository: String, - pub revision: String, - pub provenance: String, - pub public_output_allowed: bool, -} - -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct LockedTool { - pub name: String, - pub version_line: String, - pub sha256: String, - #[serde(default)] - pub enforce_sha256: bool, - pub provenance: String, - pub built_from_llvm_revision: Option, -} - -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct LockedDumps { - pub intrinsics_sha256: String, - pub nvptx_sha256: String, - pub normalized_imported_sha256: String, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ImportedFile { - pub schema: u32, - pub source: ImportedSource, - pub intrinsics: Vec, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ImportedSource { - pub llvm_repository: String, - pub llvm_revision: String, - pub llvm_tblgen_version: String, - pub llvm_tblgen_source_revision: String, - pub intrinsics_json_sha256: String, - pub nvptx_json_sha256: String, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ImportedIntrinsic { - pub source_record: String, - pub llvm_name: String, - pub arguments: Vec, - pub results: Vec, - pub classes: Vec, - pub properties: Vec, - pub selections: Vec, -} - -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ImportedSelection { - pub source_record: String, - pub asm: String, - pub predicates: Vec, - #[serde( - default, - skip_serializing_if = "ImportedSelectionConstraints::is_empty" - )] - pub constraints: ImportedSelectionConstraints, -} - -/// Normalized constraints attached to an NVPTX instruction-selection record. -/// -/// TableGen represents address-space-specific patterns through anonymous -/// `PatFrag` records and can bind intrinsic arguments to integer literals. -/// Keeping those facts separate from the assembly spelling lets policy select -/// an exact lowering without parsing PTX text. -#[derive(Debug, Clone, Default, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ImportedSelectionConstraints { - #[serde(default, skip_serializing_if = "Option::is_none")] - pub address_space: Option, - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub immediate_bindings: Vec, -} - -/// One integer literal fixed by an NVPTX instruction-selection pattern. -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ImportedImmediateBinding { - pub argument_index: usize, - pub value: i64, -} - -impl ImportedSelectionConstraints { - pub fn is_empty(&self) -> bool { - self.address_space.is_none() && self.immediate_bindings.is_empty() - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ImportedAddressSpace { - Generic, - Shared, -} - -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct OverlayFile { - pub schema: u32, - pub catalog_version: String, - pub intrinsic_abi: u32, - pub backend_profile: String, - #[serde(default)] - pub shards: Vec, - #[serde(rename = "intrinsic")] - #[serde(default)] - pub intrinsics: Vec, -} - -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct OverlayShardFile { - pub schema: u32, - pub family: String, - #[serde(rename = "intrinsic")] - #[serde(default)] - pub intrinsics: Vec, - #[serde(default)] - pub register_mma_int4: Option, - #[serde(default)] - pub register_mma_int8: Option, - #[serde(default)] - pub register_mma_b1: Option, - #[serde(default)] - pub register_mma_f8f6f4_f32: Option, - #[serde(default)] - pub register_mma_f8f6f4_f16: Option, - #[serde(default)] - pub register_mma_mxf8f6f4_f32: Option, - #[serde(default)] - pub register_mma_fp8: Option, - #[serde(default)] - pub register_mma_ampere_float: Option, - #[serde(default, alias = "sparse_mma_int8")] - pub sparse_mma_integer: Option, - #[serde(default)] - pub sparse_mma_f8f6f4_f32: Option, - #[serde(default)] - pub sparse_mma_f8f6f4_f16: Option, - #[serde(default)] - pub prmt: Option, - #[serde(default)] - pub packed_conversion_fp8: Option, - #[serde(default)] - pub packed_conversion_fp8_f16x2: Option, - #[serde(default)] - pub scalar_conversion: Option, - #[serde(default)] - pub scalar_arithmetic: Option, - #[serde(default)] - pub extended_minmax: Option, - #[serde(default)] - pub cluster_sreg: Option, - #[serde(default)] - pub cluster_barrier: Option, - #[serde(default)] - pub mbarrier_extended: Option, - #[serde(default)] - pub special_registers: Option, - #[serde(default)] - pub debug_control: Option, - #[serde(default)] - pub threadfence: Option, - #[serde(default)] - pub cluster_memory: Option, - #[serde(default)] - pub stmatrix: Option, - #[serde(default)] - pub clc: Option, - #[serde(default)] - pub wgmma_controls: Option, - #[serde(default)] - pub tma: Option, - #[serde(default)] - pub tcgen05: Option, - #[serde(default)] - pub scalar_math: Option, -} - -/// Compact admission for unary scalar floating-point math operations. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ScalarMathAdmission { - pub llvm_evidence_profile: String, - pub libnvvm_evidence_profile: String, - pub runtime_validation: RuntimeValidation, - #[serde(rename = "variant")] - pub variants: Vec, -} - -/// One reviewed scalar math variant. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ScalarMathAdmissionVariant { - pub abi_id: String, - #[serde(default)] - pub libnvvm_evidence_profile: Option, - pub format: ScalarMathFormat, - pub operation: ScalarMathOperation, - pub precision: ScalarMathPrecision, - pub subnormal: ScalarMathSubnormal, -} - -/// Compact admission for the four existing `stmatrix` stores. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct StmatrixAdmission { - pub llvm_evidence_profile: String, - pub libnvvm_evidence_profile: String, - pub runtime_validation: RuntimeValidation, - #[serde(rename = "variant")] - pub variants: Vec, -} - -/// One reviewed `stmatrix` multiplicity and layout. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct StmatrixAdmissionVariant { - pub abi_id: String, - pub multiplicity: StmatrixMultiplicity, - pub layout: StmatrixLayout, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum StmatrixMultiplicity { - X2, - X4, -} - -impl StmatrixMultiplicity { - pub const fn register_count(self) -> usize { - match self { - Self::X2 => 2, - Self::X4 => 4, - } - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum StmatrixLayout { - Normal, - Transposed, -} - -/// Compact admission for the remaining handwritten mbarrier operations. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct MbarrierExtendedAdmission { - pub llvm_evidence_profile: String, - pub libnvvm_evidence_profile: String, - pub runtime_validation: RuntimeValidation, - #[serde(rename = "variant")] - pub variants: Vec, -} - -/// One reviewed extended-mbarrier operation and its reserved ABI ID. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct MbarrierExtendedAdmissionVariant { - pub abi_id: String, - pub operation: MbarrierExtendedOperation, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum MbarrierExtendedOperation { - ArriveExpectTxCta, - ArriveExpectTxCluster, - ArriveRemoteCluster, - TryWaitTokenCta, - TryWaitParityCta, - TryWaitParityCluster, - FenceProxyAsyncSharedCta, - FenceMbarrierInitReleaseCluster, - FenceProxyAsyncGenericReleaseSharedCtaCluster, - FenceProxyAsyncGenericAcquireSharedClusterCluster, - Nanosleep, -} - -/// Compact admission for Hopper cluster special registers. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ClusterSregAdmission { - pub axes: Vec, - pub xyz_product_count: usize, - pub record_count: usize, -} - -/// Compact admission for the six cluster-barrier instructions. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ClusterBarrierAdmission { - pub llvm_evidence_profile: String, - pub libnvvm_evidence_profile: String, - pub runtime_validation: RuntimeValidation, - #[serde(rename = "variant")] - pub variants: Vec, -} - -/// One reviewed cluster-barrier spelling. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ClusterBarrierAdmissionVariant { - pub abi_id: String, - pub mode: ClusterBarrierMode, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ClusterBarrierMode { - Arrive, - ArriveAligned, - ArriveRelaxed, - ArriveRelaxedAligned, - Wait, - WaitAligned, -} - -/// Closed semantics for one cluster-barrier spelling. -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ClusterBarrier { - pub mode: ClusterBarrierMode, - pub ordering: ClusterBarrierOrdering, - pub aligned: bool, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ClusterBarrierOrdering { - Release, - Relaxed, - Acquire, -} - -/// Compact admission for cluster address mapping and remote shared reads. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ClusterMemoryAdmission { - pub llvm_evidence_profile: String, - pub libnvvm_evidence_profile: String, - pub runtime_validation: RuntimeValidation, - #[serde(rename = "variant")] - pub variants: Vec, -} - -/// One reviewed cluster-memory operation and its reserved ABI ID. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ClusterMemoryAdmissionVariant { - pub abi_id: String, - pub operation: ClusterMemoryOperation, -} - -/// Compact admission for the three WGMMA control instructions. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct WgmmaControlAdmission { - pub llvm_evidence_profile: String, - pub libnvvm_evidence_profile: String, - pub runtime_validation: RuntimeValidation, - #[serde(rename = "variant")] - pub variants: Vec, -} - -/// One reviewed WGMMA control operation. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct WgmmaControlAdmissionVariant { - pub abi_id: String, - pub mode: WgmmaControlMode, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ClusterMemoryOperation { - MapSharedRank, - ReadU32, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum WgmmaControlMode { - Fence, - CommitGroup, - WaitGroup, -} - -/// Closed semantics for one WGMMA control operation. -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct WgmmaControl { - pub mode: WgmmaControlMode, - pub adapter: WgmmaControlAdapter, - pub participation: WgmmaControlParticipation, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum WgmmaControlAdapter { - NoArguments, - ConstGenericU32ToI64Immediate, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum WgmmaControlParticipation { - WarpgroupAllThreadsSameInstruction, -} - -/// Compact admission for the reviewed non-launch special-register reads. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct SpecialRegisterAdmission { - pub llvm_evidence_profile: String, - pub libnvvm_evidence_profile: String, - pub runtime_validation: RuntimeValidation, - pub registers: Vec, - pub product_count: usize, -} - -/// Compact admission for PTX debug controls. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct DebugControlAdmission { - pub llvm_evidence_profile: String, - pub libnvvm_evidence_profile: String, - pub runtime_validation: RuntimeValidation, - pub operations: Vec, - /// Filled only when this pending shard is aggregated. - #[serde(default)] - pub abi_ids: Vec, -} - -/// Compact admission for the three CUDA thread fences. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ThreadfenceAdmission { - pub llvm_evidence_profile: String, - pub libnvvm_evidence_profile: String, - pub runtime_validation: RuntimeValidation, - #[serde(rename = "variant")] - pub variants: Vec, -} - -/// One reviewed thread-fence scope. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ThreadfenceAdmissionVariant { - pub abi_id: String, - pub scope: ThreadfenceScope, -} - -/// Scope encoded by a PTX `membar` instruction. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ThreadfenceScope { - Cta, - Device, - System, -} - -/// Compact admission for Cluster Launch Control. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ClcAdmission { - pub llvm_evidence_profile: String, - pub libnvvm_evidence_profile: String, - pub runtime_validation: RuntimeValidation, - #[serde(rename = "variant")] - pub variants: Vec, -} - -/// One reviewed Cluster Launch Control operation. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ClcAdmissionVariant { - pub abi_id: String, - pub operation: ClcOperation, -} - -/// Compact admission for the existing TMA copy and group operations. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct TmaAdmission { - pub llvm_evidence_profile: String, - pub libnvvm_evidence_profile: String, - #[serde(default)] - pub reduce_llvm_evidence_profile: Option, - #[serde(default)] - pub reduce_libnvvm_evidence_profile: Option, - pub runtime_validation: RuntimeValidation, - #[serde(rename = "variant")] - pub variants: Vec, - #[serde(rename = "reduce_variant", default)] - pub reduce_variants: Vec, -} - -/// One reviewed TMA operation and its reserved ABI ID. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct TmaAdmissionVariant { - pub abi_id: String, - pub operation: TmaOperation, -} - -/// One reviewed tensor-reduction operation and its reserved ABI ID. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct TmaReductionAdmissionVariant { - pub abi_id: String, - pub operation: TmaReductionOperation, - pub load_mode: TmaReductionLoadMode, - pub dimensions: u8, -} - -/// Compact admission for the existing Tensor Core Generation 5 operations. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct Tcgen05Admission { - pub llvm_evidence_profile: String, - pub libnvvm_evidence_profile: String, - #[serde(default)] - pub cp_llvm_evidence_profile: Option, - #[serde(default)] - pub cp_libnvvm_evidence_profile: Option, - #[serde(default)] - pub ld_llvm_evidence_profile: Option, - #[serde(default)] - pub ld_libnvvm_evidence_profile: Option, - #[serde(default)] - pub st_llvm_evidence_profile: Option, - #[serde(default)] - pub st_libnvvm_evidence_profile: Option, - #[serde(default)] - pub offset_llvm_evidence_profile: Option, - #[serde(default)] - pub offset_libnvvm_evidence_profile: Option, - #[serde(default)] - pub control_llvm_evidence_profile: Option, - #[serde(default)] - pub control_libnvvm_evidence_profile: Option, - #[serde(default)] - pub mma_llvm_evidence_profile: Option, - #[serde(default)] - pub mma_libnvvm_evidence_profile: Option, - #[serde(rename = "mma_llvm_target_contract", default)] - pub mma_llvm_target_contracts: Vec, - #[serde(rename = "mma_libnvvm_target_contract", default)] - pub mma_libnvvm_target_contracts: Vec, - pub runtime_validation: RuntimeValidation, - #[serde(rename = "variant")] - pub variants: Vec, - #[serde(rename = "cp_variant", default)] - pub cp_variants: Vec, - #[serde(rename = "ld_variant", default)] - pub ld_variants: Vec, - #[serde(rename = "st_variant", default)] - pub st_variants: Vec, - #[serde(rename = "ld_offset_variant", default)] - pub ld_offset_variants: Vec, - #[serde(rename = "st_offset_variant", default)] - pub st_offset_variants: Vec, - #[serde(rename = "mma_variant", default)] - pub mma_variants: Vec, -} - -/// One reviewed tcgen05 operation and its reserved ABI ID. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct Tcgen05AdmissionVariant { - pub abi_id: String, - pub operation: Tcgen05Operation, -} - -/// One reviewed tcgen05 copy member and CTA group. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct Tcgen05CpAdmissionVariant { - pub abi_id: String, - pub member: Tcgen05CpMember, - pub group: Tcgen05CpGroup, -} - -/// One reviewed tcgen05 load shape, repetition, and packing mode. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct Tcgen05LdAdmissionVariant { - pub abi_id: String, - pub shape: Tcgen05LdShape, - pub multiplicity: Tcgen05LdMultiplicity, - pub pack16: bool, -} - -/// One reviewed tcgen05 store shape, repetition, and unpacking mode. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct Tcgen05StAdmissionVariant { - pub abi_id: String, - pub shape: Tcgen05LdShape, - pub multiplicity: Tcgen05LdMultiplicity, - pub unpack16: bool, -} - -/// One reviewed tcgen05 MMA source form or compatibility alias. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct Tcgen05MmaAdmissionVariant { - pub abi_id: String, - pub form: Tcgen05MmaForm, - #[serde(default)] - pub alias: Option, -} - -/// Compact admission for the closed `prmt` family. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct PrmtAdmission { - pub llvm_evidence_profile: String, - pub libnvvm_evidence_profile: String, - pub runtime_validation: RuntimeValidation, - #[serde(rename = "variant")] - pub variants: Vec, -} - -/// One reviewed member of the `prmt` family. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct PrmtAdmissionVariant { - pub abi_id: String, - pub mode: PrmtMode, -} - -/// Compact admission for the closed scalar-f32 to packed-FP8 family. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct PackedConversionFp8Admission { - pub llvm_evidence_profile: String, - pub libnvvm_evidence_profile: String, - pub runtime_validation: RuntimeValidation, - pub destination_formats: Vec, - pub saturations: Vec, - pub product_count: usize, -} - -/// Compact admission for packed FP8 conversions whose other side is `f16x2`. -/// -/// Covers both directions: packing `f16x2` down to `e4m3x2`/`e5m2x2`, and -/// unpacking those back to `f16x2`. Both are single-operand conversions, unlike -/// the scalar-`f32` pair admitted by [`PackedConversionFp8Admission`]. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct PackedConversionFp8F16x2Admission { - pub llvm_evidence_profile: String, - pub libnvvm_evidence_profile: String, - pub runtime_validation: RuntimeValidation, - pub fp8_formats: Vec, - pub directions: Vec, - pub relu_variants: bool, - pub product_count: usize, -} - -/// The FP8 side of an `f16x2` conversion pair. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PackedConversionFp8Format { - E4m3x2, - E5m2x2, -} - -/// Which way an FP8/`f16x2` conversion runs. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PackedConversionFp8Direction { - /// `f16x2` narrowed to packed FP8, always saturating to finite. - Pack, - /// Packed FP8 widened back to `f16x2`, which is always exact. - Unpack, -} - -/// Compact admission for scalar F32-to-TF32 conversions. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ScalarConversionAdmission { - pub llvm_evidence_profile: String, - pub libnvvm_evidence_profile: String, - pub runtime_validation: RuntimeValidation, - #[serde(rename = "variant")] - pub variants: Vec, -} - -/// One reviewed scalar conversion variant. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ScalarConversionAdmissionVariant { - pub abi_id: String, - pub rounding: ScalarConversionRounding, - pub saturation: ScalarConversionSaturation, -} - -/// Compact admission for scalar floating-point arithmetic. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ScalarArithmeticAdmission { - pub llvm_evidence_profile: String, - pub libnvvm_evidence_profile: String, - pub runtime_validation: RuntimeValidation, - #[serde(rename = "variant")] - pub variants: Vec, -} - -/// One reviewed scalar arithmetic variant. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ScalarArithmeticAdmissionVariant { - pub abi_id: String, - pub format: ScalarArithmeticFormat, - pub operation: ScalarArithmeticOperation, - pub rounding: ScalarArithmeticRounding, - pub subnormal: ScalarArithmeticSubnormal, - pub saturation: ScalarArithmeticSaturation, -} - -/// Compact admission for the exact extended floating-point min/max family. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ExtendedMinMaxAdmission { - pub llvm_evidence_profile: String, - pub libnvvm_evidence_profile: String, - pub runtime_validation: RuntimeValidation, - #[serde(rename = "variant")] - pub variants: Vec, -} - -/// One reviewed extended min/max variant. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ExtendedMinMaxAdmissionVariant { - pub abi_id: String, - pub format: ExtendedMinMaxFormat, - pub operation: ExtendedMinMaxOperation, - pub subnormal: ExtendedMinMaxSubnormal, - pub nan: ExtendedMinMaxNan, - pub xorsign_abs: bool, -} - -/// Compact admission for a closed dense integer register-MMA family. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct RegisterMmaIntegerAdmission { - pub llvm_evidence_profile: String, - pub libnvvm_evidence_profile: String, - pub runtime_validation: RuntimeValidation, - #[serde(rename = "variant")] - pub variants: Vec, -} - -/// One reviewed member of a dense integer register-MMA family. -#[derive(Debug, Clone, Copy, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct RegisterMmaIntegerVariant { - pub shape: RegisterMmaShape, - pub a_element: RegisterMmaElement, - pub b_element: RegisterMmaElement, - pub overflow: RegisterMmaOverflow, -} - -/// Compact admission for the closed dense binary register-MMA family. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct RegisterMmaBinaryAdmission { - pub llvm_evidence_profile: String, - pub libnvvm_evidence_profile: String, - pub runtime_validation: RuntimeValidation, - #[serde(rename = "variant")] - pub variants: Vec, -} - -/// One reviewed member of the dense binary register-MMA family. -#[derive(Debug, Clone, Copy, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct RegisterMmaBinaryVariant { - pub shape: RegisterMmaShape, - pub operation: RegisterMmaOperation, -} - -/// Compact admission for one dense Blackwell `kind::f8f6f4` matrix. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct RegisterMmaF8F6F4Admission { - pub llvm_evidence_profile: String, - pub libnvvm_evidence_profile: String, - pub runtime_validation: RuntimeValidation, - /// Legacy shard metadata retained only so older overlays continue to parse. - /// ABI identity is bound from the append-only ledger by catalog ID. - #[serde(default, rename = "first_abi_id")] - pub _legacy_first_abi_id: Option, - pub a_elements: Vec, - pub b_elements: Vec, - pub product_count: usize, - pub targets: Vec, -} - -/// Compact admission for the standard FP8 register-MMA family. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct RegisterMmaFp8Admission { - pub llvm_evidence_profile: String, - pub libnvvm_evidence_profile: String, - pub runtime_validation: RuntimeValidation, - /// Legacy shard metadata retained only so older overlays continue to parse. - /// ABI identity is bound from the append-only ledger by catalog ID. - #[serde(default, rename = "first_abi_id")] - pub _legacy_first_abi_id: Option, - pub shapes: Vec, - pub accumulators: Vec, - pub a_elements: Vec, - pub b_elements: Vec, - pub product_count: usize, -} - -/// Compact admission for the reviewed Ampere floating-point MMA forms. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct RegisterMmaAmpereFloatAdmission { - pub llvm_evidence_profile: String, - pub libnvvm_evidence_profile: String, - pub runtime_validation: RuntimeValidation, - /// Legacy shard metadata retained only so older overlays continue to parse. - /// ABI identity is bound from the append-only ledger by catalog ID. - #[serde(default, rename = "first_abi_id")] - pub _legacy_first_abi_id: Option, - pub product_count: usize, - #[serde(rename = "variant")] - pub variants: Vec, -} - -/// One reviewed Ampere floating-point MMA form. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct RegisterMmaAmpereFloatVariant { - pub shape: RegisterMmaShape, - pub accumulator: RegisterMmaAccumulator, - pub element: RegisterMmaElement, -} - -/// Compact admission for a sparse integer register-MMA family. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct SparseMmaIntegerAdmission { - pub llvm_evidence_profile: String, - pub libnvvm_evidence_profile: String, - pub runtime_validation: RuntimeValidation, - pub metadata: SparseMmaMetadata, - #[serde(rename = "variant")] - pub variants: Vec, -} - -/// One reviewed member of a sparse integer register-MMA family. -#[derive(Debug, Clone, Copy, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct SparseMmaIntegerVariant { - pub shape: SparseMmaShape, - pub a_element: SparseMmaElement, - pub b_element: SparseMmaElement, - pub overflow: SparseMmaOverflow, -} - -/// Compact admission for ordered sparse `kind::f8f6f4` F32 MMA. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct SparseMmaF8F6F4Admission { - pub llvm_evidence_profile: String, - pub libnvvm_evidence_profile: String, - pub runtime_validation: RuntimeValidation, - pub a_elements: Vec, - pub b_elements: Vec, - pub product_count: usize, -} - -/// Compact admission for ordered sparse `kind::f8f6f4` packed-F16 MMA. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct SparseMmaF8F6F4F16Admission { - pub llvm_evidence_profile: String, - pub libnvvm_evidence_profile: String, - pub runtime_validation: RuntimeValidation, - /// Legacy shard metadata retained only so older overlays continue to parse. - /// ABI identity is bound from the append-only ledger by catalog ID. - #[serde(default, rename = "first_abi_id")] - pub _legacy_first_abi_id: Option, - pub a_elements: Vec, - pub b_elements: Vec, - pub product_count: usize, -} - -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct OverlayIntrinsic { - pub id: String, - pub abi_id: String, - pub operation_key: String, - pub family: String, - /// Imported LLVM records use the legacy `source_record` field below. - /// PTX-native records must instead carry an explicit tagged source. - #[serde(default)] - pub source: Option, - #[serde(default)] - pub source_record: Option, - pub rust_module: String, - pub rust_name: String, - #[serde(default)] - pub rust_arguments: Vec, - pub rust_result: String, - pub safe: bool, - #[serde(default)] - pub must_use: bool, - pub safe_allowlist_reason: Option, - pub public_rust_path: String, - #[serde(default)] - pub compatibility_rust_paths: Vec, - pub dialect_op_type: String, - pub dialect_op_name: String, - #[serde(default)] - pub dialect_operands: Vec, - #[serde(default)] - pub dialect_results: Vec, - #[serde(default)] - pub llvm_symbol: Option, - #[serde(default)] - pub resolved_llvm_symbol: Option, - #[serde(default)] - pub llvm_arguments: Vec, - #[serde(default)] - pub llvm_results: Vec, - pub pure: bool, - pub memory: String, - pub convergent: bool, - pub execution_scope: String, - pub minimum_ptx: String, - #[serde(default)] - pub minimum_sm: Option, - pub ptx_result: String, - pub targets: String, - pub ptx_isa_version: String, - pub ptx_isa_section: String, - pub ptx_isa_url: String, - pub lowering: String, - #[serde(default)] - pub backend_lowerings: Vec, - #[serde(default)] - pub packed_atomic: Option, - #[serde(default)] - pub redux: Option, - #[serde(default)] - pub vote: Option, - #[serde(default)] - pub active_mask: Option, - #[serde(default)] - pub warp_match: Option, - #[serde(default)] - pub warp_barrier: Option, - #[serde(default)] - pub warp_shuffle: Option, - #[serde(default)] - pub dot_product: Option, - #[serde(default)] - pub packed_alu: Option, - #[serde(default)] - pub integer_minmax: Option, - #[serde(default)] - pub packed_conversion: Option, - #[serde(default)] - pub scalar_conversion: Option, - #[serde(default)] - pub scalar_arithmetic: Option, - #[serde(default)] - pub scalar_math: Option, - #[serde(default)] - pub extended_minmax: Option, - #[serde(default)] - pub cp_async_copy: Option, - #[serde(default)] - pub cp_async_control: Option, - #[serde(default)] - pub cp_async_mbarrier: Option, - #[serde(default)] - pub mbarrier_basic: Option, - #[serde(default)] - pub movmatrix: Option, - #[serde(default)] - pub mbarrier_extended: Option, - #[serde(default)] - pub register_mma: Option, - #[serde(default)] - pub sparse_mma: Option, - #[serde(default)] - pub prmt: Option, - #[serde(default)] - pub cluster_barrier: Option, - #[serde(default)] - pub wgmma_control: Option, - #[serde(default)] - pub special_register: Option, - #[serde(default)] - pub debug_control: Option, - #[serde(default)] - pub cluster_memory: Option, - #[serde(default)] - pub clc: Option, - #[serde(default)] - pub tma: Option, - #[serde(default)] - pub tcgen05: Option, - #[serde(default)] - pub ldmatrix_variant: Option, - #[serde(default)] - pub ldmatrix_safety: Option, - #[serde(default)] - pub ldmatrix_adapter: Option, - #[serde(default)] - pub selected_address_space: Option, - pub expected_ptx: InstructionPattern, - pub summary: String, -} - -/// Closed semantic contract for byte permutation. -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct Prmt { - pub mode: PrmtMode, - pub adapter: PrmtAdapter, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PrmtMode { - Generic, - F4e, - B4e, - Rc8, - Ecl, - Ecr, - Rc16, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PrmtAdapter { - DirectThreeOperands, - InsertZeroSecondSource, -} - -/// Closed semantic and lowering contract for one special-register read. -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct SpecialRegister { - pub register: SpecialRegisterKind, - pub observation: SpecialRegisterObservation, - pub result_width: SpecialRegisterWidth, - pub ptx_type: SpecialRegisterPtxType, - pub output_constraint: SpecialRegisterOutputConstraint, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub llvm_exclusion: Option, -} - -/// Closed semantic contract for PTX debug controls. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct DebugControl { - pub operation: DebugControlOperation, - pub adapter: DebugControlAdapter, - pub runtime_validation: RuntimeValidation, -} - -/// Closed contract for cluster address mapping and remote shared reads. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ClusterMemory { - pub operation: ClusterMemoryOperation, - pub adapter: ClusterMemoryAdapter, - pub source_contract: ClusterMemorySourceContract, - pub runtime_validation: RuntimeValidation, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ClusterMemoryAdapter { - GenericConstAndMutPointerRankToSamePointer, - ConstU32PointerRankToU32, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ClusterMemorySourceContract { - LlvmMapaSharedClusterAs7IdentityInlinePtx, - PtxNativeMapaThenWeakClusterLoad, -} - -/// Closed semantic contract for Cluster Launch Control. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct Clc { - pub operation: ClcOperation, - pub adapter: ClcAdapter, - pub runtime_validation: RuntimeValidation, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ClcOperation { - TryCancel, - TryCancelMulticast, - QueryIsCanceled, - QueryGetFirstCtaidX, - QueryGetFirstCtaidY, - QueryGetFirstCtaidZ, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ClcAdapter { - GenericPointersToShared, - PairU64ToI128BoolToU32, - PairU64ToI128U32, -} - -/// Closed semantic contract for a TMA operation. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct Tma { - pub operation: TmaOperation, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub reduction: Option, - pub adapter: TmaAdapter, - pub runtime_validation: RuntimeValidation, -} - -impl Tma { - pub const fn dimensions(&self) -> Option { - match &self.reduction { - Some(reduction) => Some(reduction.dimensions as usize), - None => self.operation.dimensions(), - } - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum TmaOperation { - G2sTile1d, - G2sTile2d, - G2sTile2dMulticast, - G2sTile2dMulticastCg2, - G2sTile3d, - G2sTile4d, - G2sTile5d, - S2gTile1d, - S2gTile2d, - S2gTile3d, - S2gTile4d, - S2gTile5d, - Reduce, - CommitGroup, - WaitGroup, - WaitGroupRead, - PrefetchTensorMap, - PrefetchTile1d, - PrefetchTile2d, - PrefetchTile3d, - PrefetchTile4d, - PrefetchTile5d, - #[serde(rename = "prefetch_tile_gather4_2d")] - PrefetchTileGather4TwoDimensional, - PrefetchTile1dCacheHint, - PrefetchTile2dCacheHint, - PrefetchTile3dCacheHint, - PrefetchTile4dCacheHint, - PrefetchTile5dCacheHint, - #[serde(rename = "prefetch_tile_gather4_2d_cache_hint")] - PrefetchTileGather4TwoDimensionalCacheHint, - ReplaceBoxDim, - ReplaceElementStride, - ReplaceElementType, - ReplaceFillMode, - ReplaceGlobalAddress, - ReplaceGlobalDim, - ReplaceGlobalStride, - ReplaceInterleaveLayout, - ReplaceRank, - ReplaceSwizzleAtomicity, - ReplaceSwizzleMode, - FenceProxyTensorMapAcquireCluster, - FenceProxyTensorMapAcquireCta, - FenceProxyTensorMapAcquireGpu, - FenceProxyTensorMapAcquireSystem, - FenceProxyTensorMapReleaseCluster, - FenceProxyTensorMapReleaseCta, - FenceProxyTensorMapReleaseGpu, - FenceProxyTensorMapReleaseSystem, -} - -impl TmaOperation { - pub const fn dimensions(self) -> Option { - match self { - Self::G2sTile1d | Self::S2gTile1d => Some(1), - Self::G2sTile2d - | Self::G2sTile2dMulticast - | Self::G2sTile2dMulticastCg2 - | Self::S2gTile2d => Some(2), - Self::G2sTile3d | Self::S2gTile3d => Some(3), - Self::G2sTile4d | Self::S2gTile4d => Some(4), - Self::G2sTile5d | Self::S2gTile5d => Some(5), - Self::Reduce - | Self::CommitGroup - | Self::WaitGroup - | Self::WaitGroupRead - | Self::PrefetchTensorMap - | Self::PrefetchTile1d - | Self::PrefetchTile2d - | Self::PrefetchTile3d - | Self::PrefetchTile4d - | Self::PrefetchTile5d - | Self::PrefetchTileGather4TwoDimensional - | Self::PrefetchTile1dCacheHint - | Self::PrefetchTile2dCacheHint - | Self::PrefetchTile3dCacheHint - | Self::PrefetchTile4dCacheHint - | Self::PrefetchTile5dCacheHint - | Self::PrefetchTileGather4TwoDimensionalCacheHint - | Self::ReplaceBoxDim - | Self::ReplaceElementStride - | Self::ReplaceElementType - | Self::ReplaceFillMode - | Self::ReplaceGlobalAddress - | Self::ReplaceGlobalDim - | Self::ReplaceGlobalStride - | Self::ReplaceInterleaveLayout - | Self::ReplaceRank - | Self::ReplaceSwizzleAtomicity - | Self::ReplaceSwizzleMode - | Self::FenceProxyTensorMapAcquireCluster - | Self::FenceProxyTensorMapAcquireCta - | Self::FenceProxyTensorMapAcquireGpu - | Self::FenceProxyTensorMapAcquireSystem - | Self::FenceProxyTensorMapReleaseCluster - | Self::FenceProxyTensorMapReleaseCta - | Self::FenceProxyTensorMapReleaseGpu - | Self::FenceProxyTensorMapReleaseSystem => None, - } - } - - pub const fn prefetch_coordinate_count(self) -> Option { - match self { - Self::PrefetchTile1d | Self::PrefetchTile1dCacheHint => Some(1), - Self::PrefetchTile2d | Self::PrefetchTile2dCacheHint => Some(2), - Self::PrefetchTile3d | Self::PrefetchTile3dCacheHint => Some(3), - Self::PrefetchTile4d | Self::PrefetchTile4dCacheHint => Some(4), - Self::PrefetchTile5d - | Self::PrefetchTile5dCacheHint - | Self::PrefetchTileGather4TwoDimensional - | Self::PrefetchTileGather4TwoDimensionalCacheHint => Some(5), - _ => None, - } - } - - pub const fn uses_prefetch_cache_hint(self) -> bool { - matches!( - self, - Self::PrefetchTile1dCacheHint - | Self::PrefetchTile2dCacheHint - | Self::PrefetchTile3dCacheHint - | Self::PrefetchTile4dCacheHint - | Self::PrefetchTile5dCacheHint - | Self::PrefetchTileGather4TwoDimensionalCacheHint - ) - } -} - -/// Closed identity for one TMA tensor-reduction operation. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct TmaReduction { - pub operation: TmaReductionOperation, - pub load_mode: TmaReductionLoadMode, - pub dimensions: u8, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum TmaReductionOperation { - Add, - And, - Dec, - Inc, - Max, - Min, - Or, - Xor, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum TmaReductionLoadMode { - Tile, - Im2col, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum TmaAdapter { - G2sPointersCoordinatesBarrierInjectDefaults, - G2sPointersCoordinatesBarrierMaskInjectDefaults, - S2gPointersCoordinatesInjectDefaults, - ReductionPointersCoordinatesInjectDefaults, - NoOperands, - CompileTimeConstantMaxPending, - DescriptorPointer, - DescriptorCoordinatesInjectDefaults, - DescriptorCoordinatesCacheHintInjectFlag, - DescriptorAndAddressPointers, - DescriptorOrdinalAndU32, - DescriptorOrdinalAndU64, - DescriptorAndImmediateU32, - DescriptorAndRuntimeU32, - DescriptorPointerInjectBytes, -} - -/// Closed semantic contract for one tcgen05 operation. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct Tcgen05 { - pub operation: Tcgen05Operation, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub cp: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub ld: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub st: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub mma: Option, - pub adapter: Tcgen05Adapter, - pub source_contract: Tcgen05SourceContract, - pub runtime_validation: RuntimeValidation, -} - -/// Closed identity and selector contract for one tcgen05 MMA API. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct Tcgen05Mma { - pub form: Tcgen05MmaForm, - pub selector_layout: Tcgen05MmaSelectorLayout, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub fixed_selectors: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub alias: Option, - pub llvm_target: CatalogTargetRequirement, - pub libnvvm_target: CatalogTargetRequirement, -} - -/// The 14 LLVM source forms covered by the first tcgen05 MMA batch. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum Tcgen05MmaForm { - Shared, - Tensor, - TensorAshift, - SpShared, - SpTensor, - SpTensorAshift, - WsShared, - WsSharedZeroColMask, - WsSpShared, - WsSpSharedZeroColMask, - WsSpTensor, - WsSpTensorZeroColMask, - WsTensor, - WsTensorZeroColMask, -} - -/// Immediate arguments that select one imported tcgen05 MMA spelling. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(tag = "kind", rename_all = "snake_case")] -pub enum Tcgen05MmaSelectorLayout { - Base { - kind_argument: u8, - cta_group_argument: u8, - collector_a_argument: u8, - collector_a_upper_exclusive: u8, - }, - WarpSpecialized { - kind_argument: u8, - b_buffer_argument: u8, - b_usage_argument: u8, - }, -} - -/// A fixed warp-specialized selector tuple used by compatibility aliases. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct Tcgen05MmaFixedSelectors { - pub kind: Tcgen05MmaKind, - pub b_buffer: u8, - pub b_usage: Tcgen05MmaBUsage, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum Tcgen05MmaKind { - F16, - Tf32, - F8f6f4, - I8, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum Tcgen05MmaBUsage { - Discard, - LastUse, - Fill, - Use, -} - -impl Tcgen05MmaBUsage { - pub const fn selector_value(self) -> u8 { - match self { - Self::Discard => 0, - Self::LastUse => 1, - Self::Fill => 2, - Self::Use => 3, - } - } -} - -/// Public names proposed by PR #346 for the generic f8f6f4 carrier. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum Tcgen05MmaAlias { - E4m3, - E5m2, - E2m3, - E3m2, - E2m1, -} - -/// Closed identity for one tcgen05 shared-to-tensor-memory copy. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct Tcgen05Cp { - pub member: Tcgen05CpMember, - pub group: Tcgen05CpGroup, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -pub enum Tcgen05CpMember { - #[serde(rename = "128x128b_b4x16_p64")] - M128x128bB4x16P64, - #[serde(rename = "128x128b_b6x16_p32")] - M128x128bB6x16P32, - #[serde(rename = "128x128b")] - M128x128b, - #[serde(rename = "128x256b_b4x16_p64")] - M128x256bB4x16P64, - #[serde(rename = "128x256b_b6x16_p32")] - M128x256bB6x16P32, - #[serde(rename = "32x128b_warpx4_b4x16_p64")] - M32x128bWarpx4B4x16P64, - #[serde(rename = "32x128b_warpx4_b6x16_p32")] - M32x128bWarpx4B6x16P32, - #[serde(rename = "32x128b_warpx4")] - M32x128bWarpx4, - #[serde(rename = "4x256b_b4x16_p64")] - M4x256bB4x16P64, - #[serde(rename = "4x256b_b6x16_p32")] - M4x256bB6x16P32, - #[serde(rename = "4x256b")] - M4x256b, - #[serde(rename = "64x128b_warpx2_01_23_b4x16_p64")] - M64x128bWarpx2Pair0123B4x16P64, - #[serde(rename = "64x128b_warpx2_01_23_b6x16_p32")] - M64x128bWarpx2Pair0123B6x16P32, - #[serde(rename = "64x128b_warpx2_01_23")] - M64x128bWarpx2Pair0123, - #[serde(rename = "64x128b_warpx2_02_13_b4x16_p64")] - M64x128bWarpx2Pair0213B4x16P64, - #[serde(rename = "64x128b_warpx2_02_13_b6x16_p32")] - M64x128bWarpx2Pair0213B6x16P32, - #[serde(rename = "64x128b_warpx2_02_13")] - M64x128bWarpx2Pair0213, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum Tcgen05CpGroup { - Cg1, - Cg2, -} - -/// Closed identity for one tcgen05 tensor-memory load. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct Tcgen05Ld { - pub shape: Tcgen05LdShape, - pub multiplicity: Tcgen05LdMultiplicity, - pub pack16: bool, -} - -/// Closed identity for one tcgen05 tensor-memory store. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct Tcgen05St { - pub shape: Tcgen05LdShape, - pub multiplicity: Tcgen05LdMultiplicity, - pub unpack16: bool, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -pub enum Tcgen05LdShape { - #[serde(rename = "16x32bx2")] - M16x32bx2, - #[serde(rename = "16x64b")] - M16x64b, - #[serde(rename = "16x128b")] - M16x128b, - #[serde(rename = "16x256b")] - M16x256b, - #[serde(rename = "32x32b")] - M32x32b, -} - -impl Tcgen05LdShape { - pub const fn register_multiplier(self) -> usize { - match self { - Self::M16x32bx2 | Self::M16x64b | Self::M32x32b => 1, - Self::M16x128b => 2, - Self::M16x256b => 4, - } - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum Tcgen05LdMultiplicity { - X1, - X2, - X4, - X8, - X16, - X32, - X64, - X128, -} - -impl Tcgen05LdMultiplicity { - pub const fn count(self) -> usize { - match self { - Self::X1 => 1, - Self::X2 => 2, - Self::X4 => 4, - Self::X8 => 8, - Self::X16 => 16, - Self::X32 => 32, - Self::X64 => 64, - Self::X128 => 128, - } - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum Tcgen05Operation { - Alloc, - Dealloc, - RelinquishAllocPermit, - FenceBeforeThreadSync, - FenceAfterThreadSync, - Commit, - CommitSharedCluster, - MmaWsF16, - MmaF16, - MmaWsBf16, - MmaWsTf32, - CpSmemToTmem, - Ld16x256bX8Pure, - Ld16x256bPure, - LoadWait, - StoreWait, - AllocCg2, - DeallocCg2, - RelinquishAllocPermitCg2, - MmaF16Cg2, - CommitCg2, - CommitSharedClusterCg2, - CommitMulticastCg2, - CpSmemToTmemCg2, - Ld, - St, - CommitMulticast, - ShiftDown, - ShiftDownCg2, - Mma, -} - -impl Tcgen05Operation { - pub const fn execution_scope(self) -> &'static str { - match self { - Self::Alloc - | Self::Dealloc - | Self::RelinquishAllocPermit - | Self::Ld16x256bX8Pure - | Self::Ld16x256bPure - | Self::LoadWait - | Self::StoreWait - | Self::AllocCg2 - | Self::DeallocCg2 - | Self::RelinquishAllocPermitCg2 - | Self::Ld - | Self::St => "warp", - Self::FenceBeforeThreadSync - | Self::FenceAfterThreadSync - | Self::Commit - | Self::CommitSharedCluster - | Self::MmaWsF16 - | Self::MmaF16 - | Self::MmaWsBf16 - | Self::MmaWsTf32 - | Self::CpSmemToTmem - | Self::MmaF16Cg2 - | Self::CommitCg2 - | Self::CommitSharedClusterCg2 - | Self::CommitMulticastCg2 - | Self::CpSmemToTmemCg2 - | Self::CommitMulticast - | Self::ShiftDown - | Self::ShiftDownCg2 - | Self::Mma => "thread", - } - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum Tcgen05Adapter { - SharedPointerColumnsToVoid, - TmemAddressColumnsToVoid, - NoOperands, - BarrierPointerToVoid, - MmaWsDropLegacyADescriptor, - MmaInjectZeroDisableLanes, - TmemDescriptorToVoid, - TmemToF32x32, - TmemToF32x4, - BarrierPointerMaskToVoid, - TmemInjectPack16ToU32Registers, - TmemU32RegistersInjectUnpack16ToVoid, - TmemHalfSplitOffsetInjectPack16ToU32Registers, - TmemHalfSplitOffsetU32RegistersInjectUnpack16ToVoid, - TmemAddressToVoid, - MmaDirectSelectors, - MmaWsFixedSelectorsDropLegacyADescriptor, -} - -/// Relationship between the public operation and LLVM's NVPTX selection. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum Tcgen05SourceContract { - ExactTablegenSelection, - TablegenSelectionChangesPtx, - LlvmCustomLoweringWithoutSelection, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum SpecialRegisterKind { - Clock, - Clock64, - Globaltimer, - Envreg1, - Envreg2, - Smid, - Nsmid, - Gridid, - Warpid, - Nwarpid, - DynamicSmemSize, - TotalSmemSize, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum DebugControlOperation { - Trap, - Breakpoint, - Pmevent, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum SpecialRegisterObservation { - StablePure, - VolatileObservation, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum SpecialRegisterWidth { - B32, - B64, -} - -impl SpecialRegisterWidth { - pub const fn bits(self) -> u32 { - match self { - Self::B32 => 32, - Self::B64 => 64, - } - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum SpecialRegisterPtxType { - B32, - U32, - U64, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum SpecialRegisterOutputConstraint { - Register32, - Register64, -} - -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct SpecialRegisterLlvmExclusion { - pub source_record: String, - pub llvm_symbol: String, - pub imported_result_width: SpecialRegisterWidth, - pub reason: SpecialRegisterLlvmExclusionReason, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum SpecialRegisterLlvmExclusionReason { - ResultWidthMismatch, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum DebugControlAdapter { - Direct, - ConstGenericToImmediateU32, -} - -/// Backend-specific lowering selected by reviewed evidence. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct OverlayBackendLowering { - pub backend: IntrinsicBackend, - pub mechanism: BackendLoweringMechanism, - pub evidence_profile: String, - /// Optional exact target alternatives for this backend route. - #[serde(default)] - pub targets: Option, - /// Optional backend-profile floor. When absent, the intrinsic's native - /// target requirement is used. - #[serde(default)] - pub minimum_ptx: Option, - #[serde(default)] - pub minimum_sm: Option, -} - -/// Provenance for a generated intrinsic. PTX-native operations deliberately -/// have no invented LLVM TableGen record or LLVM intrinsic symbol. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)] -pub enum IntrinsicSource { - LlvmImported { source_record: String }, - PtxNative { instruction: String }, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum IntrinsicBackend { - LlvmNvptx, - LibNvvm, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum BackendLoweringMechanism { - TypedNvvm, - InlinePtx, -} - -/// Closed identity for the small execution-control families that share a -/// result-less MIR/dialect representation. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] -pub enum ExecutionControlOperation { - BarrierCtaSync, - BarrierCtaSyncAligned, - BarrierCtaArrive, - BarrierCtaArriveAligned, - GridDependencyLaunchDependents, - GridDependencyWait, - SetMaxNRegInc, - SetMaxNRegDec, -} - -impl ExecutionControlOperation { - pub fn from_catalog_id(id: &str) -> Option { - Some(match id { - "barrier_cta_sync" => Self::BarrierCtaSync, - "barrier_cta_sync_aligned" => Self::BarrierCtaSyncAligned, - "barrier_cta_arrive" => Self::BarrierCtaArrive, - "barrier_cta_arrive_aligned" => Self::BarrierCtaArriveAligned, - "grid_dependency_launch_dependents" => Self::GridDependencyLaunchDependents, - "grid_dependency_wait" => Self::GridDependencyWait, - "setmaxnreg_inc" => Self::SetMaxNRegInc, - "setmaxnreg_dec" => Self::SetMaxNRegDec, - _ => return None, - }) - } - - pub const fn family(self) -> &'static str { - match self { - Self::BarrierCtaSync - | Self::BarrierCtaSyncAligned - | Self::BarrierCtaArrive - | Self::BarrierCtaArriveAligned => "counted_barrier", - Self::GridDependencyLaunchDependents | Self::GridDependencyWait => "grid_dependency", - Self::SetMaxNRegInc | Self::SetMaxNRegDec => "register_control", - } - } - - pub const fn operand_count(self) -> usize { - match self { - Self::BarrierCtaSync - | Self::BarrierCtaSyncAligned - | Self::BarrierCtaArrive - | Self::BarrierCtaArriveAligned => 2, - Self::GridDependencyLaunchDependents | Self::GridDependencyWait => 0, - Self::SetMaxNRegInc | Self::SetMaxNRegDec => 1, - } - } - - pub const fn requires_immediate_operands(self) -> bool { - matches!(self, Self::SetMaxNRegInc | Self::SetMaxNRegDec) - } -} - -/// Closed semantic identity for the generated `ldmatrix` family. -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct LdmatrixVariant { - pub shape: LdmatrixShape, - pub multiplicity: LdmatrixMultiplicity, - pub layout: LdmatrixLayout, - pub element: LdmatrixElement, - pub state_space: LdmatrixStateSpace, -} - -impl LdmatrixVariant { - pub const fn register_count(&self) -> usize { - let matrices = self.multiplicity.register_count(); - match self.shape { - LdmatrixShape::M8n8 | LdmatrixShape::M8n16 => matrices, - LdmatrixShape::M16n16 => matrices * 2, - } - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum LdmatrixShape { - M8n8, - M8n16, - M16n16, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum LdmatrixMultiplicity { - X1, - X2, - X4, -} - -impl LdmatrixMultiplicity { - pub const fn register_count(self) -> usize { - match self { - Self::X1 => 1, - Self::X2 => 2, - Self::X4 => 4, - } - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum LdmatrixLayout { - Normal, - Transposed, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum LdmatrixElement { - B16, - B8, - B8x16B4x16P64, - B8x16B6x16P32, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum LdmatrixStateSpace { - Shared, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct LdmatrixSafety { - pub participation: LdmatrixParticipation, - pub address_contract: LdmatrixAddressContract, - pub memory_order: LdmatrixMemoryOrder, - pub runtime_validation: RuntimeValidation, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum LdmatrixParticipation { - AllWarpLanesSameInstruction, - AllWarpLanesSameInstructionAndQualifiersNoExitedLanes, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum LdmatrixAddressContract { - WarpLaneAddressesMappedByMultiplicitySixteenByteAlignedSixteenBytesReadableWithSm75Replication, - WarpLaneAddressesMappedByMultiplicitySixteenByteAlignedSixteenBytesReadable, - WarpLaneAddressesMappedByMultiplicitySixteenByteAlignedThirtyTwoBytesReadable, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum LdmatrixMemoryOrder { - Weak, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum RuntimeValidation { - Unexecuted, - Executed, -} - -/// Closed contract for the in-register warp matrix transpose. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct Movmatrix { - pub participation: MovmatrixParticipation, - pub adapter: MovmatrixAdapter, - pub runtime_validation: RuntimeValidation, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum MovmatrixParticipation { - AllWarpLanesSameInstructionNoExitedLanes, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum MovmatrixAdapter { - PackedB16x2U32ToPackedB16x2U32, -} - -/// Closed contract for register-only warp-level `mma.sync` operations. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct RegisterMma { - pub shape: RegisterMmaShape, - pub operation: RegisterMmaOperation, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub kind: Option, - pub accumulator: RegisterMmaAccumulator, - pub a_element: RegisterMmaElement, - pub b_element: RegisterMmaElement, - pub a_layout: RegisterMmaLayout, - pub b_layout: RegisterMmaLayout, - pub overflow: RegisterMmaOverflow, - pub participation: RegisterMmaParticipation, - pub adapter: RegisterMmaAdapter, - pub compatibility_source: RegisterMmaCompatibilitySource, - pub runtime_validation: RuntimeValidation, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum RegisterMmaKind { - Standard, - F8f6f4, - Mxf8f6f4, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum RegisterMmaShape { - M8n8k4, - M8n8k16, - M8n8k32, - M16n8k4, - M16n8k8, - M16n8k16, - M16n8k32, - M16n8k64, - M8n8k128, - M16n8k128, - M16n8k256, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum RegisterMmaOperation { - Multiply, - AndPopc, - XorPopc, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum RegisterMmaAccumulator { - F16, - F32, - F64, - S32, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum RegisterMmaElement { - B1, - Bf16, - E2m1, - E2m3, - E3m2, - E4m3, - E5m2, - F16, - Tf32, - F64, - S4, - U4, - S8, - U8, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum RegisterMmaLayout { - Row, - Col, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum RegisterMmaOverflow { - NotApplicable, - Wrapping, - Satfinite, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum RegisterMmaParticipation { - AllWarpLanesSameInstructionAndQualifiersNoExitedLanes, -} - -/// Rust `C, A, B` fragment shape used by the importer. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum RegisterMmaAdapter { - C2U32A2U32B1U32ToD2U32, - C2U32A4U32B2U32ToD2U32, - C4F32A2U32B1U32ToD4F32, - C4F32A4U32B2U32ToD4F32, - C2F64A1F64B1F64ToD2F64, - C2I32A1U32B1U32ToD2I32, - C4I32A4U32B2U32ToD4I32, - C4I32A2U32B1U32ToD4I32, - C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32, -} - -/// Where the stable `cuda_device::wmma` callable is defined. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum RegisterMmaCompatibilitySource { - ExistingStub, - GeneratedStub, -} - -/// Closed semantic contract for register-only sparse `mma.sp` operations. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct SparseMma { - pub shape: SparseMmaShape, - pub accumulator: SparseMmaAccumulator, - pub a_element: SparseMmaElement, - pub b_element: SparseMmaElement, - pub a_layout: SparseMmaLayout, - pub b_layout: SparseMmaLayout, - pub overflow: SparseMmaOverflow, - pub metadata: SparseMmaMetadata, - pub selector: SparseMmaSelector, - pub participation: SparseMmaParticipation, - pub adapter: SparseMmaAdapter, - pub llvm_adapter: SparseMmaLlvmAdapter, - pub compatibility_source: SparseMmaCompatibilitySource, - pub runtime_validation: RuntimeValidation, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum SparseMmaShape { - M16n8k32, - M16n8k64, - M16n8k128, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum SparseMmaAccumulator { - F16, - F32, - S32, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum SparseMmaElement { - E2m1, - E2m3, - E3m2, - E4m3, - E5m2, - S4, - U4, - S8, - U8, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum SparseMmaLayout { - Row, - Col, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum SparseMmaOverflow { - NotApplicable, - Wrapping, - Satfinite, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum SparseMmaMetadata { - Standard, - Ordered, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum SparseMmaSelector { - ImmediateZeroOrOne, - ImmediateZero, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum SparseMmaParticipation { - AllWarpLanesSameInstructionAndQualifiersNoExitedLanes, -} - -/// Rust `C, A, B, metadata, selector` shape used by the importer. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum SparseMmaAdapter { - C2U32A4U32B4U32MetadataU32SelectorU32ToD2U32, - C4F32A4U32B4U32MetadataU32SelectorU32ToD4F32, - C4I32A2U32B2U32MetadataU32SelectorU32ToD4I32, - C4I32A4U32B4U32MetadataU32SelectorU32ToD4I32, -} - -/// LLVM `A, B, C, metadata, selector` shape. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum SparseMmaLlvmAdapter { - A4I32B4I32C2V2F16MetadataI32SelectorI32ToD2V2F16, - A4I32B4I32C4F32MetadataI32SelectorI32ToD4F32, - A2I32B2I32C4I32MetadataI32SelectorI32ToD4I32, - A4I32B4I32C4I32MetadataI32SelectorI32ToD4I32, -} - -/// Where the stable `cuda_device::wmma` callable is defined. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum SparseMmaCompatibilitySource { - GeneratedStub, -} - -/// Closed semantic contract for the generated packed global atomic-add -/// family. These fields are intentionally enums rather than free-form strings: -/// accepting an unreviewed state space, scope, or floating-point mode must -/// require a generator change. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct PackedAtomic { - pub format: PackedAtomicFormat, - /// PTX ISA hardware floor, kept separate from cuda-oxide's admitted floor - /// and from backend-profile floors. - pub native_minimum_sm: u16, - pub operation: PackedAtomicOperation, - pub state_space: PackedAtomicStateSpace, - pub ordering: PackedAtomicOrdering, - pub scope: PackedAtomicScope, - pub rounding: PackedAtomicRounding, - pub subnormal: PackedAtomicSubnormal, - pub atomicity: PackedAtomicAtomicity, - pub pointer_contract: PackedAtomicPointerContract, - pub access_contract: PackedAtomicAccessContract, - pub scope_contract: PackedAtomicScopeContract, - pub codegen_contract: PackedAtomicCodegenContract, - pub return_contract: PackedAtomicReturnContract, - pub adapter: PackedAtomicAdapter, - pub runtime_validation: RuntimeValidation, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PackedAtomicFormat { - F16x2, - Bf16x2, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PackedAtomicOperation { - Add, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PackedAtomicStateSpace { - Global, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PackedAtomicOrdering { - Relaxed, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PackedAtomicScope { - Gpu, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PackedAtomicRounding { - NearestEven, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PackedAtomicSubnormal { - Preserve, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PackedAtomicAtomicity { - PerElement, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PackedAtomicPointerContract { - MutableGlobalU32Aligned4, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PackedAtomicAccessContract { - NoMixedWholeWordOrNonAtomicAccess, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PackedAtomicScopeContract { - RacingAtomicsMutuallyInclusive, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PackedAtomicCodegenContract { - ExactNativeInstruction, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PackedAtomicReturnContract { - OldValuesPerElementMayBeNoncoherent, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PackedAtomicAdapter { - OldPackedU32, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum LdmatrixAdapter { - SingleResultDirect, - MultipleResultsToArray, -} - -/// Closed semantic and lowering contract for the generated `redux.sync` -/// family. -/// -/// The Rust and NVVM dialect APIs intentionally put the participation mask -/// first, while LLVM's NVVM intrinsic puts the lane value first. Keeping that -/// adapter typed prevents a generic direct-call renderer from silently -/// swapping the collective's source and member mask. -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct Redux { - pub operation: ReduxOperation, - pub participation: ReduxParticipation, - pub adapter: ReduxAdapter, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ReduxOperation { - Add, - Umin, - Min, - Umax, - Max, - And, - Or, - Xor, - Fmin, - FminNan, - FminAbs, - FminAbsNan, - Fmax, - FmaxNan, - FmaxAbs, - FmaxAbsNan, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ReduxParticipation { - ExecutingLaneNamedAllNamedLanesSameInstructionAndMask, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ReduxAdapter { - MaskValueToSourceMemberMask, -} - -/// Closed semantic and lowering contract for the generated `vote.sync` -/// family. -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct Vote { - pub mode: VoteMode, - pub participation: VoteParticipation, - pub legacy_pre_sm70: PreSm70MemberMaskRule, - pub adapter: VoteAdapter, - pub mask_encoding: MaskEncoding, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum VoteMode { - All, - Any, - Ballot, - Uni, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum VoteParticipation { - ExecutingLaneNamedAllNamedLanesSameInstructionAndMask, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum VoteAdapter { - DirectMaskPredicate, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum MaskEncoding { - RegisterOrImmediate, -} - -/// Closed semantic and lowering contract for `activemask`. -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ActiveMask { - pub observation: ActiveMaskObservation, - pub adapter: ActiveMaskAdapter, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ActiveMaskObservation { - ExecutingLanesAtInstruction, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ActiveMaskAdapter { - DirectZeroOperandMask, -} - -/// Closed semantic and lowering contract for `match.sync`. -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct WarpMatch { - pub mode: WarpMatchMode, - pub value_width: WarpMatchValueWidth, - pub participation: WarpMatchParticipation, - pub adapter: WarpMatchAdapter, - pub value_encoding: MatchOperandEncoding, - pub mask_encoding: MatchOperandEncoding, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum WarpMatchMode { - Any, - All, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum WarpMatchValueWidth { - B32, - B64, -} - -impl WarpMatchValueWidth { - pub const fn bits(self) -> u32 { - match self { - Self::B32 => 32, - Self::B64 => 64, - } - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum WarpMatchParticipation { - ExecutingLaneNamedAllNamedLanesSameInstructionAndMask, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum WarpMatchAdapter { - DirectMask, - ProjectMaskDiscardPredicate, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum MatchOperandEncoding { - RegisterOrImmediate, -} - -/// Closed semantic and lowering contract for `bar.warp.sync`. -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct WarpBarrier { - pub participation: WarpBarrierParticipation, - pub legacy_pre_sm70: PreSm70MemberMaskRule, - pub adapter: WarpBarrierAdapter, - pub mask_encoding: WarpBarrierMaskEncoding, - pub memory_ordering: WarpBarrierMemoryOrdering, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum WarpBarrierParticipation { - ExecutingLaneNamedAllNamedLanesSameInstructionAndMask, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PreSm70MemberMaskRule { - AllNamedLanesConvergedAndOnlyNamedLanesActive, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum WarpBarrierAdapter { - DirectMemberMask, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum WarpBarrierMaskEncoding { - RegisterOrImmediate, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum WarpBarrierMemoryOrdering { - ParticipatingLanes, -} - -/// Closed semantic and lowering contract for `shfl.sync`. -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct WarpShuffle { - pub mode: WarpShuffleMode, - pub value_kind: WarpShuffleValueKind, - pub participation: WarpShuffleParticipation, - pub legacy_pre_sm70: PreSm70MemberMaskRule, - pub source_lane: WarpShuffleSourceLane, - pub adapter: WarpShuffleAdapter, - pub clamp: u32, - pub lane_encoding: WarpShuffleOperandEncoding, - pub mask_encoding: WarpShuffleOperandEncoding, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum WarpShuffleMode { - Idx, - Bfly, - Down, - Up, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum WarpShuffleValueKind { - I32, - F32, - I64, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum WarpShuffleParticipation { - ExecutingLaneNamedAllNamedLanesSameInstructionAndMask, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum WarpShuffleSourceLane { - InRangeSourceActiveAndNamedOutOfRangeCopiesSelf, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum WarpShuffleAdapter { - MaskValueLaneOrDeltaInsertClamp, - /// Split i64 into low/high b32 halves, shuffle both in one convergent - /// side-effecting block, then reassemble the original bit layout. - MaskValueLaneOrDeltaSplitI64LowHighB32InsertClampReassemble, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum WarpShuffleOperandEncoding { - RegisterOrImmediate, - RegisterOnly, -} - -/// Closed identity and source adapter for generated packed integer dot products. -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct DotProduct { - pub operation: DotProductOperation, - pub signedness: DotProductSignedness, - pub adapter: DotProductAdapter, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum DotProductOperation { - Dp2a, - Dp4a, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum DotProductSignedness { - Signed, - Unsigned, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum DotProductAdapter { - DirectThreeOperands, - InsertLowHalfFalse, -} - -/// Closed identity and carrier contract for packed floating-point ALU ops. -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct PackedAlu { - pub format: PackedAluFormat, - /// Hardware floor of the native PTX instruction, independent of the - /// target floor admitted by cuda-oxide. - pub native_minimum_sm: u16, - pub operation: PackedAluOperation, - pub adapter: PackedAluAdapter, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PackedAluFormat { - Bf16x2, - F16x2, - F32x2, -} - -/// Closed identity and carrier contract for extended integer min/max ops. -/// -/// Covers the PTX ISA 8.0 integer min/max extensions: the `.relu` saturation -/// qualifier on `s32` and `s16x2`, plus the packed `s16x2`/`u16x2` forms. -/// These are the DPX-adjacent shapes ptxas fuses into `VIMNMX`-family SASS. -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct IntegerMinMax { - pub format: IntegerMinMaxFormat, - /// Hardware floor of the native PTX instruction, independent of the - /// target floor admitted by cuda-oxide. - pub native_minimum_sm: u16, - pub operation: IntegerMinMaxOperation, - pub relu: bool, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum IntegerMinMaxFormat { - S32, - S16x2, - U16x2, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum IntegerMinMaxOperation { - Min, - Max, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PackedAluOperation { - Add, - AddFtz, - Sub, - SubFtz, - Mul, - MulFtz, - Fma, - FmaFtz, - FmaSat, - FmaFtzSat, - FmaRelu, - FmaFtzRelu, - Min, - Max, - Neg, - Abs, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PackedAluAdapter { - DirectPackedU32, - DirectPackedU64, -} - -/// Closed contract for scalar floating-point arithmetic. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ScalarArithmetic { - pub format: ScalarArithmeticFormat, - pub operation: ScalarArithmeticOperation, - pub rounding: ScalarArithmeticRounding, - pub subnormal: ScalarArithmeticSubnormal, - pub saturation: ScalarArithmeticSaturation, - pub runtime_validation: RuntimeValidation, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ScalarArithmeticFormat { - F32, - F64, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ScalarArithmeticOperation { - Mul, - Div, - Fma, - Add, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ScalarArithmeticRounding { - Rn, - Rz, - Rm, - Rp, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ScalarArithmeticSubnormal { - Preserve, - Ftz, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ScalarArithmeticSaturation { - None, - Sat, -} - -/// Closed contract for unary scalar floating-point math operations. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ScalarMath { - pub format: ScalarMathFormat, - pub operation: ScalarMathOperation, - pub precision: ScalarMathPrecision, - pub subnormal: ScalarMathSubnormal, - pub runtime_validation: RuntimeValidation, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ScalarMathFormat { - F16, - F32, - F64, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ScalarMathOperation { - Sin, - Cos, - /// Routes through inline PTX: LLVM 22's tblgen models ex2 as the - /// overloaded `int_nvvm_ex2_approx{,_ftz}` records (anyfloat, no DAG - /// selection pattern), so the evidence contract cannot admit a typed - /// call. The legacy `llvm.nvvm.ex2.approx.f`/`.ftz.f` names still select - /// directly on both llc 21 and 22, so this is promotable to a typed - /// route once the import resolves the overloaded family. - Ex2, - Lg2, - Rcp, - Rsqrt, - Sqrt, - /// Routes through inline PTX: LLVM 22.1.2's tblgen export carries no - /// record for tanh at all (llc 22 selects `llvm.nvvm.tanh.approx.f32` - /// via NVVMIntrinsic-class matching, and llc 21 miscompiles it into an - /// extern funcall), so this is the family's only PTX-native source. - /// Hardware floor is sm_75 (PTX ISA 7.0); the family contract gates it - /// at the attested sm_80 evidence floor like every other variant. - Tanh, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ScalarMathPrecision { - Approx, - Rn, - Rz, - Rm, - Rp, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ScalarMathSubnormal { - Preserve, - Ftz, -} - -/// Closed identity and carrier contract for extended floating-point min/max. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ExtendedMinMax { - pub format: ExtendedMinMaxFormat, - pub operation: ExtendedMinMaxOperation, - pub subnormal: ExtendedMinMaxSubnormal, - pub nan: ExtendedMinMaxNan, - pub xorsign_abs: bool, - pub adapter: ExtendedMinMaxAdapter, - pub runtime_validation: RuntimeValidation, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ExtendedMinMaxFormat { - F32, - F16, - Bf16, - F16x2, - Bf16x2, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ExtendedMinMaxOperation { - Min, - Max, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ExtendedMinMaxSubnormal { - Preserve, - Ftz, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ExtendedMinMaxNan { - Number, - Nan, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ExtendedMinMaxAdapter { - DirectF32, - /// A single 16-bit float carried as its `u16` bit pattern, matching how - /// `packed_conversion` already carries `e4m3x2` and `e5m2x2` operands. - DirectHalfU16, - DirectPackedU32, -} - -/// Closed contract for converting between scalar and packed values. -/// -/// The source may be a scalar pair (`f32x2`, two operands) or an already-packed -/// 16-bit or 32-bit register (`f16x2`, `e4m3x2`, `e5m2x2`, one operand), so the -/// operand arity follows [`PackedConversionSourceFormat`] rather than being -/// fixed at two. -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct PackedConversion { - pub source_format: PackedConversionSourceFormat, - pub destination_format: PackedConversionDestinationFormat, - pub rounding: PackedConversionRounding, - pub saturation: PackedConversionSaturation, - pub adapter: PackedConversionAdapter, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PackedConversionSourceFormat { - E4m3x2, - E5m2x2, - F16x2, - F32x2, -} - -impl PackedConversionSourceFormat { - /// Number of source operands the conversion takes. - /// - /// `f32x2` names two scalar `f32` operands; every packed source arrives in a - /// single register. - pub fn operand_count(self) -> usize { - match self { - Self::F32x2 => 2, - Self::E4m3x2 | Self::E5m2x2 | Self::F16x2 => 1, - } - } - - /// PTX source-type token, used as the trailing `cvt` modifier. - pub fn ptx_token(self) -> &'static str { - match self { - Self::E4m3x2 => "e4m3x2", - Self::E5m2x2 => "e5m2x2", - Self::F16x2 => "f16x2", - Self::F32x2 => "f32", - } - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PackedConversionDestinationFormat { - Bf16x2, - E4m3x2, - E5m2x2, - F16x2, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PackedConversionRounding { - NearestEven, - TowardZero, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PackedConversionSaturation { - None, - Relu, - Satfinite, - SatfiniteRelu, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PackedConversionAdapter { - /// Forward the single packed source operand unchanged. - /// - /// PTX orders `cvt` operands as `d, a`, so a one-operand conversion needs no - /// reordering to keep the Rust argument order. - Identity, - /// Swap the two scalar operands. - /// - /// PTX writes the second source operand into the low half, so the operands - /// are reversed to keep the first Rust argument in the low half. - ReverseHighLowOperands, -} - -/// Closed contract for one scalar floating-point conversion. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct ScalarConversion { - pub source_format: ScalarConversionSourceFormat, - pub destination_format: ScalarConversionDestinationFormat, - pub rounding: ScalarConversionRounding, - pub saturation: ScalarConversionSaturation, - pub result_representation: ScalarConversionResultRepresentation, - pub adapter: ScalarConversionAdapter, - pub runtime_validation: RuntimeValidation, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ScalarConversionSourceFormat { - F32, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ScalarConversionDestinationFormat { - Tf32, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ScalarConversionRounding { - NearestAway, - NearestEven, - TowardZero, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ScalarConversionSaturation { - None, - Relu, - Satfinite, - ReluSatfinite, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ScalarConversionResultRepresentation { - RawU32Bits, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ScalarConversionAdapter { - DirectF32ToRawU32Bits, -} - -/// Closed contract for classic global-to-shared `cp.async` copies. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct CpAsyncCopy { - pub cache_policy: CpAsyncCachePolicy, - pub copy_size: CpAsyncCopySize, - pub source_size: CpAsyncSourceSize, - pub adapter: CpAsyncAdapter, - pub runtime_validation: RuntimeValidation, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum CpAsyncCachePolicy { - Ca, - Cg, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum CpAsyncCopySize { - B4, - B8, - B16, -} - -impl CpAsyncCopySize { - pub const fn bytes(self) -> u32 { - match self { - Self::B4 => 4, - Self::B8 => 8, - Self::B16 => 16, - } - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum CpAsyncSourceSize { - Full, - Runtime, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum CpAsyncAdapter { - DirectPointers, - DirectPointersAndSourceSize, -} - -/// Closed contract for classic `cp.async` group controls. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct CpAsyncControl { - pub operation: CpAsyncControlOperation, - pub adapter: CpAsyncControlAdapter, - pub runtime_validation: RuntimeValidation, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum CpAsyncControlOperation { - CommitGroup, - WaitAll, - WaitGroup, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum CpAsyncControlAdapter { - NoOperands, - CompileTimeConstantMaxPending, -} - -/// Closed contract for associating classic `cp.async` completion with an mbarrier. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct CpAsyncMbarrier { - pub operation: CpAsyncMbarrierOperation, - pub state_space: CpAsyncMbarrierStateSpace, - pub adapter: CpAsyncMbarrierAdapter, - pub runtime_validation: RuntimeValidation, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum CpAsyncMbarrierOperation { - Arrive, - ArriveNoInc, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum CpAsyncMbarrierStateSpace { - Generic, - Shared, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum CpAsyncMbarrierAdapter { - PointerToVoid, -} - -/// Closed contract for the basic shared-memory mbarrier lifecycle. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct MbarrierBasic { - pub operation: MbarrierBasicOperation, - pub state_space: MbarrierStateSpace, - pub adapter: MbarrierBasicAdapter, - pub runtime_validation: RuntimeValidation, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum MbarrierBasicOperation { - Init, - Arrive, - TestWait, - Inval, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum MbarrierStateSpace { - Shared, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum MbarrierBasicAdapter { - #[serde(rename = "pointer_count_to_void")] - InitPointerCountToVoid, - #[serde(rename = "pointer_to_token")] - ArrivePointerToToken, - #[serde(rename = "pointer_token_to_predicate")] - TestWaitPointerTokenToPredicate, - #[serde(rename = "pointer_to_void")] - InvalPointerToVoid, -} - -/// Closed contract for the remaining handwritten mbarrier operations. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct MbarrierExtended { - pub operation: MbarrierExtendedOperation, - pub adapter: MbarrierExtendedAdapter, - pub source_contract: MbarrierExtendedSourceContract, - pub runtime_validation: RuntimeValidation, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum MbarrierExtendedAdapter { - PointerTxCountBytesToTokenDroppingTxCount, - RawClusterAddressToVoid, - PointerTokenToPredicate, - PointerParityToPredicate, - ZeroOperandsToVoid, - NanosecondsToVoid, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum MbarrierExtendedSourceContract { - LlvmImported, - PtxNativeRawClusterAddress, -} - -#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct AbiLedgerFile { - pub schema: u32, - pub intrinsic_abi: u32, - #[serde(rename = "entry")] - pub entries: Vec, -} - -#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct AbiLedgerEntry { - pub abi_id: String, - pub status: String, - pub catalog_id: String, - pub operation_key: String, - pub raw_rust_signature: AbiRawRustSignature, -} - -#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct AbiRawRustSignature { - pub safe: bool, - pub arguments: Vec, - pub result: String, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct EvidenceFile { - pub schema: u32, - pub backend_profile: String, - #[serde(default)] - pub backend_kind: Option, - pub llvm_revision: String, - pub backend_version: String, - pub backend_sha256: String, - #[serde(default)] - pub artifact_path: Option, - #[serde(default)] - pub build_id_prefix: Option, - #[serde(default)] - pub nvvm_ir_version: Option, - #[serde(default)] - pub debug_ir_version: Option, - pub records: Vec, -} - -/// Schema-6 evidence before matrix expansion. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct EvidenceFileV6 { - pub schema: u32, - pub backend_profile: String, - #[serde(default)] - pub backend_kind: Option, - pub llvm_revision: String, - pub backend_version: String, - pub backend_sha256: String, - #[serde(default)] - pub artifact_path: Option, - #[serde(default)] - pub build_id_prefix: Option, - #[serde(default)] - pub nvvm_ir_version: Option, - #[serde(default)] - pub debug_ir_version: Option, - #[serde(default)] - pub defaults: EvidenceRecordDefaults, - #[serde(default)] - pub fixtures: Vec, - #[serde(default)] - pub matrices: Vec, - #[serde(default)] - pub records: Vec, -} - -/// Facts shared by every record in one evidence matrix. -#[derive(Debug, Clone, Default, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct EvidenceRecordDefaults { - #[serde(default)] - pub resolved_llvm_symbol: Option, - #[serde(default)] - pub llvm_arguments: Option>, - #[serde(default)] - pub llvm_results: Option>, - #[serde(default)] - pub concrete_llvm_arguments: Option>, - #[serde(default)] - pub concrete_llvm_results: Option>, - #[serde(default)] - pub target_triple: Option, - #[serde(default)] - pub gpu_target: Option, - #[serde(default)] - pub ptx_feature: Option, - #[serde(default)] - pub status: Option, - #[serde(default)] - pub stages: Vec, - #[serde(default)] - pub declaration_attributes_canonicalized: Option, - #[serde(default)] - pub runtime_validation: Option, -} - -/// One shared fixture and the number of records it covers. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct EvidenceFixture { - pub id: String, - pub coverage_count: usize, - pub stages: Vec, -} - -/// One Cartesian evidence matrix. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct EvidenceMatrix { - pub axes: Vec, - pub product_count: usize, - #[serde(default)] - pub fixtures: Vec, - pub template: EvidenceMatrixTemplate, -} - -/// One named matrix axis. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct EvidenceMatrixAxis { - pub name: String, - pub values: Vec, -} - -/// Identity and matrix-specific facts for one record template. -#[derive(Debug, Clone, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct EvidenceMatrixTemplate { - pub id: String, - #[serde(default)] - pub source: Option, - #[serde(default)] - pub source_record: Option, - #[serde(deserialize_with = "deserialize_required_optional_string")] - pub llvm_symbol: Option, - pub expected_ptx: InstructionPattern, - #[serde(default)] - pub facts: EvidenceRecordDefaults, -} - -fn deserialize_required_optional_string<'de, D>(deserializer: D) -> Result, D::Error> -where - D: serde::Deserializer<'de>, -{ - Option::::deserialize(deserializer) -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct EvidenceRecord { - pub id: String, - #[serde(default)] - pub source: Option, - #[serde(default)] - pub source_record: Option, - #[serde(default)] - pub llvm_symbol: Option, - #[serde(default)] - pub resolved_llvm_symbol: Option, - #[serde(default)] - pub llvm_arguments: Vec, - #[serde(default)] - pub llvm_results: Vec, - #[serde(default)] - pub concrete_llvm_arguments: Vec, - #[serde(default)] - pub concrete_llvm_results: Vec, - pub target_triple: String, - pub gpu_target: String, - pub ptx_feature: String, - pub status: String, - #[serde(default)] - pub stages: Vec, - #[serde(default)] - pub declaration_attributes_canonicalized: Option, - #[serde(default)] - pub runtime_validation: Option, - pub expected_ptx: InstructionPattern, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct EvidenceStage { - pub targets: Vec, - pub representation: String, - pub stage: EvidenceStageKind, - #[serde(default)] - pub mechanism: Option, - pub outcome: String, - pub detail: String, - #[serde(default)] - pub artifact_kind: Option, - #[serde(default)] - pub tool_path: Option, - #[serde(default)] - pub tool_version: Option, - #[serde(default)] - pub tool_sha256: Option, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum EvidenceArtifactKind { - Cubin, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum EvidenceStageKind { - DeclarationCanonicalization, - BackendCodegen, - DeviceLink, - PtxAssembly, - Runtime, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct CatalogFile { - pub schema: u32, - pub catalog_version: String, - pub intrinsic_abi: u32, - pub generator_version: String, - pub source: CatalogSource, - pub inputs: CatalogInputs, - pub intrinsics: Vec, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct CatalogSource { - pub llvm_repository: String, - pub llvm_revision: String, - pub llvm_tblgen_version: String, - pub llvm_tblgen_source_revision: String, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct CatalogInputs { - pub imported_sha256: String, - pub overlay_sha256: String, - pub abi_ledger_sha256: String, - pub evidence_sha256: Vec, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct CatalogIntrinsic { - pub id: String, - pub operation_key: String, - pub family: String, - pub source: IntrinsicSource, - pub selections: Vec, - pub rust: CatalogRust, - pub dialect: CatalogDialect, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub llvm: Option, - pub semantics: CatalogSemantics, - pub target: CatalogTarget, - pub backend: CatalogBackend, - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub backend_lowerings: Vec, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub packed_atomic: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub redux: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub vote: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub active_mask: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub warp_match: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub warp_barrier: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub warp_shuffle: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub dot_product: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub packed_alu: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub integer_minmax: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub packed_conversion: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub scalar_conversion: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub scalar_arithmetic: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub scalar_math: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub extended_minmax: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub cp_async_copy: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub cp_async_control: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub cp_async_mbarrier: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub mbarrier_basic: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub movmatrix: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub mbarrier_extended: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub register_mma: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub sparse_mma: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub prmt: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub cluster_barrier: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub wgmma_control: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub special_register: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub debug_control: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub cluster_memory: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub clc: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub tma: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub tcgen05: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub ldmatrix: Option, - pub lowering: String, - pub expected_ptx: InstructionPattern, - pub summary: String, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct CatalogSelection { - pub source_record: String, - pub asm: String, - pub predicates: Vec, - #[serde( - default, - skip_serializing_if = "ImportedSelectionConstraints::is_empty" - )] - pub constraints: ImportedSelectionConstraints, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct CatalogRust { - pub abi_id: String, - pub module: String, - pub name: String, - pub arguments: Vec, - pub result: String, - pub safe: bool, - pub must_use: bool, - pub safe_allowlist_reason: Option, - pub canonical_path: String, - pub public_path: String, - pub compatibility_paths: Vec, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct CatalogDialect { - pub op_type: String, - pub op_name: String, - pub operands: Vec, - pub results: Vec, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct CatalogLlvm { - pub symbol: String, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub resolved_symbol: Option, - pub arguments: Vec, - pub results: Vec, - pub properties: Vec, - pub result_facts: CatalogLlvmResultFacts, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct CatalogLlvmResultFacts { - pub no_undef: bool, - pub range: Option, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct CatalogHalfOpenRange { - pub lower: String, - pub upper_exclusive: String, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct CatalogSemantics { - pub pure: bool, - pub memory: String, - pub convergent: bool, - pub execution_scope: String, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct CatalogTarget { - pub minimum_ptx: PtxVersion, - pub hardware: CatalogHardwareTarget, - pub ptx_result: String, - pub targets: String, - pub ptx_isa_version: String, - pub ptx_isa_section: String, - pub ptx_isa_url: String, -} - -/// A PTX ISA version encoded as `major * 10 + minor`. -/// -/// PTX currently uses one decimal minor digit. The resolver validates that -/// shape before constructing this value, so generated consumers compare a -/// number rather than reparsing policy text. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] -pub struct PtxVersion(u16); - -impl PtxVersion { - pub const fn encoded(self) -> u16 { - self.0 - } - - pub const fn major(self) -> u16 { - self.0 / 10 - } - - pub const fn minor(self) -> u16 { - self.0 % 10 - } -} - -impl FromStr for PtxVersion { - type Err = String; - - fn from_str(value: &str) -> Result { - let (major, minor) = value - .split_once('.') - .ok_or_else(|| "expected major.minor".to_owned())?; - if major.is_empty() - || !major.bytes().all(|byte| byte.is_ascii_digit()) - || minor.len() != 1 - || !minor.bytes().all(|byte| byte.is_ascii_digit()) - { - return Err("expected numeric major.minor with one minor digit".to_owned()); - } - let major: u16 = major.parse().map_err(|_| "major version is too large")?; - let minor: u16 = minor.parse().unwrap(); - if format!("{major}.{minor}") != value { - return Err("version is not in canonical major.minor form".to_owned()); - } - let encoded = major - .checked_mul(10) - .and_then(|value| value.checked_add(minor)) - .ok_or_else(|| "version is too large".to_owned())?; - Ok(Self(encoded)) - } -} - -impl Serialize for PtxVersion { - fn serialize(&self, serializer: S) -> Result - where - S: serde::Serializer, - { - serializer.serialize_str(&self.to_string()) - } -} - -impl<'de> Deserialize<'de> for PtxVersion { - fn deserialize(deserializer: D) -> Result - where - D: serde::Deserializer<'de>, - { - let value = String::deserialize(deserializer)?; - value.parse().map_err(serde::de::Error::custom) - } -} - -impl fmt::Display for PtxVersion { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(formatter, "{}.{}", self.major(), self.minor()) - } -} - -/// Reviewed hardware availability for an intrinsic. -/// -/// Exact `a` and `f` targets stay distinct from monotonic minimums. A target -/// matrix also keeps each hardware target paired with its PTX floor. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(tag = "kind", rename_all = "snake_case")] -pub enum CatalogHardwareTarget { - All, - AnyOf { - alternatives: Vec, - }, - /// Closed selector contracts with their exact PTX and hardware pairs. - TargetMatrix { - contracts: Vec, - }, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(tag = "kind", rename_all = "snake_case")] -pub enum CatalogHardwareAlternative { - MinimumSm { sm: u16 }, - ExactArchitecture { sm: u16 }, - FamilyTarget { sm: u16 }, -} - -/// One PTX floor paired with one hardware alternative. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct CatalogTargetAlternative { - pub minimum_ptx: PtxVersion, - pub hardware: CatalogHardwareAlternative, -} - -/// One selector tuple and its exact PTX and hardware pairs. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct CatalogTargetContract { - pub selectors: Vec, - pub alternatives: Vec, -} - -/// One field/value pair that selects a target contract. -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct TargetSelectorBinding { - pub name: String, - pub value: String, -} - -/// One selector-specific target contract from admission policy. -#[allow(dead_code)] -#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct TargetContract { - #[serde(default)] - pub selectors: Vec, - pub alternatives: Vec, -} - -/// One target spelling and PTX floor. -#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct TargetContractAlternative { - pub target: String, - pub minimum_ptx: String, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct CatalogBackend { - pub profile: String, - pub version: String, - pub sha256: String, - pub status: String, - pub target_triple: String, - pub gpu_target: String, - pub ptx_feature: String, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct CatalogBackendLowering { - pub backend: IntrinsicBackend, - pub mechanism: BackendLoweringMechanism, - pub evidence_profile: String, - pub target: CatalogTargetRequirement, - pub version: String, - pub sha256: String, - pub artifact_path: Option, - pub build_id_prefix: Option, - pub status: String, - pub stages: Vec, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct CatalogTargetRequirement { - pub minimum_ptx: PtxVersion, - pub hardware: CatalogHardwareTarget, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct CatalogLdmatrix { - pub variant: LdmatrixVariant, - pub safety: LdmatrixSafety, - pub adapter: LdmatrixAdapter, - pub selected_address_space: ImportedAddressSpace, -} - -impl CatalogIntrinsic { - pub fn scalar_width(&self) -> Option { - match self.rust.result.as_str() { - "u32" => Some(32), - "u64" => Some(64), - _ => None, - } - } - - pub fn llvm_identifier(&self) -> String { - llvm_symbol_to_identifier(&self.llvm.as_ref().expect("LLVM-backed intrinsic").symbol) - } - - pub fn resolved_llvm_identifier(&self) -> String { - let llvm = self.llvm.as_ref().expect("LLVM-backed intrinsic"); - llvm_symbol_to_identifier(llvm.resolved_symbol.as_deref().unwrap_or(&llvm.symbol)) - } -} - -fn llvm_symbol_to_identifier(symbol: &str) -> String { - if !symbol.contains('_') { - return symbol.replace('.', "_"); - } - - let suffix = symbol - .strip_prefix("llvm.") - .expect("LLVM intrinsic symbol must start with llvm."); - let mut output = String::from("llvm__"); - for ch in suffix.chars() { - match ch { - '.' => output.push_str("_d"), - '_' => output.push_str("_u"), - ch if ch.is_ascii_alphanumeric() => output.push(ch), - _ => panic!("LLVM intrinsic symbol contains an unsupported character"), - } - } - output -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn llvm_identifier_encoding_preserves_literal_underscores() { - assert_eq!( - llvm_symbol_to_identifier("llvm.nvvm.read.ptx.sreg.tid.x"), - "llvm_nvvm_read_ptx_sreg_tid_x" - ); - assert_eq!( - llvm_symbol_to_identifier("llvm.nvvm.wgmma.wait_group.sync.aligned"), - "llvm__nvvm_dwgmma_dwait_ugroup_dsync_daligned" - ); - assert_eq!( - llvm_symbol_to_identifier( - "llvm.nvvm.ldmatrix.sync.aligned.m16n16.x1.trans.b8x16.b4x16_p64.p3" - ), - "llvm__nvvm_dldmatrix_dsync_daligned_dm16n16_dx1_dtrans_db8x16_db4x16_up64_dp3" - ); - assert_eq!( - llvm_symbol_to_identifier( - "llvm.nvvm.ldmatrix.sync.aligned.m8n16.x4.b8x16.b6x16_p32.p3" - ), - "llvm__nvvm_dldmatrix_dsync_daligned_dm8n16_dx4_db8x16_db6x16_up32_dp3" - ); - } - - #[test] - fn ldmatrix_register_count_includes_the_shape_width() { - let variant = |shape, multiplicity| LdmatrixVariant { - shape, - multiplicity, - layout: LdmatrixLayout::Normal, - element: LdmatrixElement::B8, - state_space: LdmatrixStateSpace::Shared, - }; - - assert_eq!( - variant(LdmatrixShape::M8n16, LdmatrixMultiplicity::X4).register_count(), - 4 - ); - assert_eq!( - variant(LdmatrixShape::M16n16, LdmatrixMultiplicity::X1).register_count(), - 2 - ); - assert_eq!( - variant(LdmatrixShape::M16n16, LdmatrixMultiplicity::X2).register_count(), - 4 - ); - } - - #[test] - fn blackwell_ldmatrix_address_contracts_keep_readable_widths_distinct() { - assert_eq!( - serde_json::from_str::( - r#""warp_lane_addresses_mapped_by_multiplicity_sixteen_byte_aligned_sixteen_bytes_readable""# - ) - .unwrap(), - LdmatrixAddressContract::WarpLaneAddressesMappedByMultiplicitySixteenByteAlignedSixteenBytesReadable - ); - assert_eq!( - serde_json::from_str::( - r#""warp_lane_addresses_mapped_by_multiplicity_sixteen_byte_aligned_thirty_two_bytes_readable""# - ) - .unwrap(), - LdmatrixAddressContract::WarpLaneAddressesMappedByMultiplicitySixteenByteAlignedThirtyTwoBytesReadable - ); - } - - #[test] - fn wgmma_control_contract_is_closed() { - let parsed: WgmmaControl = serde_json::from_str( - r#"{ - "mode": "wait_group", - "adapter": "const_generic_u32_to_i64_immediate", - "participation": "warpgroup_all_threads_same_instruction" - }"#, - ) - .unwrap(); - assert_eq!(parsed.mode, WgmmaControlMode::WaitGroup); - assert_eq!( - parsed.adapter, - WgmmaControlAdapter::ConstGenericU32ToI64Immediate - ); - assert_eq!( - parsed.participation, - WgmmaControlParticipation::WarpgroupAllThreadsSameInstruction - ); - - let open_ended = r#"{ - "mode": "wait_group", - "adapter": "const_generic_u32_to_i64_immediate", - "participation": "warpgroup_all_threads_same_instruction", - "custom_ptx": "wgmma.wait_group.sync.aligned 0;" - }"#; - assert!(serde_json::from_str::(open_ended).is_err()); - } - - #[test] - fn locked_tool_rejects_misspelled_security_field() { - let input = r#" -name = "llvm-tblgen" -version_line = "LLVM version test" -sha256 = "abc" -enforce_sha25 = true -provenance = "test" -"#; - let error = toml::from_str::(input).unwrap_err(); - assert!(error.to_string().contains("enforce_sha25")); - } - - #[test] - fn imported_selection_rejects_misspelled_constraint() { - let input = r#"{ - "source_record": "selection", - "asm": "op;", - "predicates": [], - "constraints": { "adress_space": "shared" } - }"#; - let error = serde_json::from_str::(input).unwrap_err(); - assert!(error.to_string().contains("adress_space")); - } - - #[test] - fn imported_selection_preserves_immediate_binding() { - let input = r#"{ - "source_record": "DOT2_lo_ss", - "asm": "dp2a.lo.s32.s32 $dst, $a, $b, $c;", - "predicates": ["hasDotInstructions"], - "constraints": { - "immediate_bindings": [ - { "argument_index": 2, "value": 0 } - ] - } - }"#; - let selection = serde_json::from_str::(input).unwrap(); - assert_eq!( - selection.constraints.immediate_bindings, - [ImportedImmediateBinding { - argument_index: 2, - value: 0, - }] - ); - assert!(!selection.constraints.is_empty()); - } - - #[test] - fn imported_immediate_binding_rejects_misspelled_index() { - let input = r#"{ - "source_record": "DOT2_lo_ss", - "asm": "dp2a.lo.s32.s32 $dst, $a, $b, $c;", - "predicates": [], - "constraints": { - "immediate_bindings": [ - { "argument_indx": 2, "value": 0 } - ] - } - }"#; - let error = serde_json::from_str::(input).unwrap_err(); - assert!(error.to_string().contains("argument_indx")); - } - - #[test] - fn redux_contract_rejects_unknown_operand_adapter() { - let input = r#" -operation = "add" -participation = "executing_lane_named_all_named_lanes_same_instruction_and_mask" -adapter = "mask_value_direct" -"#; - let error = toml::from_str::(input).unwrap_err(); - assert!(error.to_string().contains("mask_value_direct")); - } - - #[test] - fn packed_alu_contract_rejects_unknown_format_operation_and_adapter() { - let valid = r#" -format = "bf16x2" -native_minimum_sm = 80 -operation = "fma" -adapter = "direct_packed_u32" -"#; - toml::from_str::(valid).unwrap(); - for invalid in [ - valid.replace("format = \"bf16x2\"", "format = \"bf16\""), - valid.replace("native_minimum_sm = 80\n", ""), - valid.replace("native_minimum_sm = 80", "native_minimum_sm = \"80\""), - valid.replace("operation = \"fma\"", "operation = \"mad\""), - valid.replace( - "adapter = \"direct_packed_u32\"", - "adapter = \"bitcast_any\"", - ), - format!("{valid}unreviewed = true\n"), - ] { - assert!(toml::from_str::(&invalid).is_err(), "{invalid}"); - } - } - - #[test] - fn packed_conversion_contract_rejects_open_ended_policy() { - let valid = r#" -source_format = "f32x2" -destination_format = "bf16x2" -rounding = "nearest_even" -saturation = "none" -adapter = "reverse_high_low_operands" -"#; - toml::from_str::(valid).unwrap(); - for invalid in [ - valid.replace("source_format = \"f32x2\"", "source_format = \"f64x2\""), - valid.replace( - "destination_format = \"bf16x2\"", - "destination_format = \"f8x2\"", - ), - valid.replace("rounding = \"nearest_even\"", "rounding = \"zero\""), - valid.replace("saturation = \"none\"", "saturation = \"finite\""), - valid.replace( - "adapter = \"reverse_high_low_operands\"", - "adapter = \"direct\"", - ), - format!("{valid}unreviewed = true\n"), - ] { - assert!( - toml::from_str::(&invalid).is_err(), - "{invalid}" - ); - } - } - - #[test] - fn vote_contract_rejects_unknown_modes_and_mask_encodings() { - let valid = r#" -mode = "all" -participation = "executing_lane_named_all_named_lanes_same_instruction_and_mask" -legacy_pre_sm70 = "all_named_lanes_converged_and_only_named_lanes_active" -adapter = "direct_mask_predicate" -mask_encoding = "register_or_immediate" -"#; - toml::from_str::(valid).unwrap(); - - for invalid in [ - valid.replace("mode = \"all\"", "mode = \"match\""), - valid.replace( - "mask_encoding = \"register_or_immediate\"", - "mask_encoding = \"any_operand\"", - ), - valid.replace( - "legacy_pre_sm70 = \"all_named_lanes_converged_and_only_named_lanes_active\"", - "legacy_pre_sm70 = \"independent_threads\"", - ), - format!("{valid}unreviewed = true\n"), - ] { - assert!(toml::from_str::(&invalid).is_err(), "{invalid}"); - } - } - - #[test] - fn warp_shuffle_contract_rejects_open_ended_policy() { - let valid = r#" -mode = "idx" -value_kind = "i32" -participation = "executing_lane_named_all_named_lanes_same_instruction_and_mask" -legacy_pre_sm70 = "all_named_lanes_converged_and_only_named_lanes_active" -source_lane = "in_range_source_active_and_named_out_of_range_copies_self" -adapter = "mask_value_lane_or_delta_insert_clamp" -clamp = 31 -lane_encoding = "register_or_immediate" -mask_encoding = "register_or_immediate" -"#; - toml::from_str::(valid).unwrap(); - - for invalid in [ - valid.replace("mode = \"idx\"", "mode = \"rotate\""), - valid.replace("value_kind = \"i32\"", "value_kind = \"b32\""), - valid.replace( - "source_lane = \"in_range_source_active_and_named_out_of_range_copies_self\"", - "source_lane = \"unchecked\"", - ), - valid.replace( - "lane_encoding = \"register_or_immediate\"", - "lane_encoding = \"anything\"", - ), - format!("{valid}unreviewed = true\n"), - ] { - assert!( - toml::from_str::(&invalid).is_err(), - "{invalid}" - ); - } - - let i64 = r#" -mode = "down" -value_kind = "i64" -participation = "executing_lane_named_all_named_lanes_same_instruction_and_mask" -legacy_pre_sm70 = "all_named_lanes_converged_and_only_named_lanes_active" -source_lane = "in_range_source_active_and_named_out_of_range_copies_self" -adapter = "mask_value_lane_or_delta_split_i64_low_high_b32_insert_clamp_reassemble" -clamp = 31 -lane_encoding = "register_only" -mask_encoding = "register_only" -"#; - let parsed = toml::from_str::(i64).unwrap(); - assert_eq!(parsed.value_kind, WarpShuffleValueKind::I64); - assert_eq!( - parsed.adapter, - WarpShuffleAdapter::MaskValueLaneOrDeltaSplitI64LowHighB32InsertClampReassemble - ); - assert_eq!( - parsed.lane_encoding, - WarpShuffleOperandEncoding::RegisterOnly - ); - - for invalid in [ - i64.replace("value_kind = \"i64\"", "value_kind = \"u64\""), - i64.replace( - "adapter = \"mask_value_lane_or_delta_split_i64_low_high_b32_insert_clamp_reassemble\"", - "adapter = \"split_any_width\"", - ), - i64.replace( - "mask_encoding = \"register_only\"", - "mask_encoding = \"any_operand\"", - ), - ] { - assert!( - toml::from_str::(&invalid).is_err(), - "{invalid}" - ); - } - } - - #[test] - fn warp_match_contract_rejects_open_ended_adapters_and_encodings() { - let valid = r#" -mode = "all" -value_width = "b64" -participation = "executing_lane_named_all_named_lanes_same_instruction_and_mask" -adapter = "project_mask_discard_predicate" -value_encoding = "register_or_immediate" -mask_encoding = "register_or_immediate" -"#; - toml::from_str::(valid).unwrap(); - - for invalid in [ - valid.replace("mode = \"all\"", "mode = \"equal\""), - valid.replace( - "adapter = \"project_mask_discard_predicate\"", - "adapter = \"first_result\"", - ), - valid.replace( - "value_encoding = \"register_or_immediate\"", - "value_encoding = \"anything\"", - ), - format!("{valid}unreviewed = true\n"), - ] { - assert!(toml::from_str::(&invalid).is_err(), "{invalid}"); - } - } - - #[test] - fn warp_barrier_contract_rejects_open_ended_policy() { - let valid = r#" -participation = "executing_lane_named_all_named_lanes_same_instruction_and_mask" -legacy_pre_sm70 = "all_named_lanes_converged_and_only_named_lanes_active" -adapter = "direct_member_mask" -mask_encoding = "register_or_immediate" -memory_ordering = "participating_lanes" -"#; - toml::from_str::(valid).unwrap(); - - for invalid in [ - valid.replace("adapter = \"direct_member_mask\"", "adapter = \"direct\""), - valid.replace( - "legacy_pre_sm70 = \"all_named_lanes_converged_and_only_named_lanes_active\"", - "legacy_pre_sm70 = \"independent_threads\"", - ), - valid.replace( - "mask_encoding = \"register_or_immediate\"", - "mask_encoding = \"any_operand\"", - ), - valid.replace( - "memory_ordering = \"participating_lanes\"", - "memory_ordering = \"none\"", - ), - format!("{valid}unreviewed = true\n"), - ] { - assert!( - toml::from_str::(&invalid).is_err(), - "{invalid}" - ); - } - } - - #[test] - fn mbarrier_basic_contract_rejects_open_ended_policy() { - let valid = r#" -operation = "test_wait" -state_space = "shared" -adapter = "pointer_token_to_predicate" -runtime_validation = "unexecuted" -"#; - let parsed = toml::from_str::(valid).unwrap(); - assert_eq!(parsed.operation, MbarrierBasicOperation::TestWait); - assert_eq!( - parsed.adapter, - MbarrierBasicAdapter::TestWaitPointerTokenToPredicate - ); - - for invalid in [ - valid.replace("operation = \"test_wait\"", "operation = \"wait\""), - valid.replace("state_space = \"shared\"", "state_space = \"global\""), - valid.replace( - "adapter = \"pointer_token_to_predicate\"", - "adapter = \"direct\"", - ), - format!("{valid}unreviewed = true\n"), - ] { - assert!( - toml::from_str::(&invalid).is_err(), - "{invalid}" - ); - } - } - - #[test] - fn movmatrix_contract_rejects_open_ended_policy() { - let valid = r#" -participation = "all_warp_lanes_same_instruction_no_exited_lanes" -adapter = "packed_b16x2_u32_to_packed_b16x2_u32" -runtime_validation = "unexecuted" -"#; - let parsed = toml::from_str::(valid).unwrap(); - assert_eq!( - parsed.participation, - MovmatrixParticipation::AllWarpLanesSameInstructionNoExitedLanes - ); - - for invalid in [ - valid.replace( - "all_warp_lanes_same_instruction_no_exited_lanes", - "participating_lanes", - ), - valid.replace("packed_b16x2_u32_to_packed_b16x2_u32", "direct_u32"), - format!("{valid}unreviewed = true\n"), - ] { - assert!(toml::from_str::(&invalid).is_err(), "{invalid}"); - } - } - - #[test] - fn mbarrier_extended_contract_rejects_open_ended_policy() { - let valid = r#" -operation = "arrive_expect_tx_cta" -adapter = "pointer_tx_count_bytes_to_token_dropping_tx_count" -source_contract = "llvm_imported" -runtime_validation = "unexecuted" -"#; - let parsed = toml::from_str::(valid).unwrap(); - assert_eq!( - parsed.adapter, - MbarrierExtendedAdapter::PointerTxCountBytesToTokenDroppingTxCount - ); - assert_eq!( - parsed.source_contract, - MbarrierExtendedSourceContract::LlvmImported - ); - - for invalid in [ - valid.replace("llvm_imported", "auto"), - valid.replace( - "pointer_tx_count_bytes_to_token_dropping_tx_count", - "direct", - ), - format!("{valid}unreviewed = true\n"), - ] { - assert!( - toml::from_str::(&invalid).is_err(), - "{invalid}" - ); - } - } - - #[test] - fn register_mma_contract_parses_unsigned_k16_generated_stub() { - let valid = r#" -shape = "m16n8k16" -operation = "multiply" -accumulator = "s32" -a_element = "s8" -b_element = "u8" -a_layout = "row" -b_layout = "col" -overflow = "satfinite" -participation = "all_warp_lanes_same_instruction_and_qualifiers_no_exited_lanes" -adapter = "c4_i32_a2_u32_b1_u32_to_d4_i32" -compatibility_source = "generated_stub" -runtime_validation = "unexecuted" -"#; - let parsed = toml::from_str::(valid).unwrap(); - assert_eq!(parsed.b_element, RegisterMmaElement::U8); - assert_eq!(parsed.adapter, RegisterMmaAdapter::C4I32A2U32B1U32ToD4I32); - assert_eq!( - parsed.compatibility_source, - RegisterMmaCompatibilitySource::GeneratedStub - ); - - for invalid in [ - valid.replace("b_element = \"u8\"", "b_element = \"i8\""), - valid.replace( - "adapter = \"c4_i32_a2_u32_b1_u32_to_d4_i32\"", - "adapter = \"direct\"", - ), - valid.replace( - "compatibility_source = \"generated_stub\"", - "compatibility_source = \"automatic\"", - ), - format!("{valid}unreviewed = true\n"), - ] { - assert!( - toml::from_str::(&invalid).is_err(), - "{invalid}" - ); - } - } - - #[test] - fn sparse_mma_contract_closes_the_selector_and_metadata_modes() { - let valid = r#" -shape = "m16n8k32" -accumulator = "s32" -a_element = "s8" -b_element = "u8" -a_layout = "row" -b_layout = "col" -overflow = "satfinite" -metadata = "standard" -selector = "immediate_zero_or_one" -participation = "all_warp_lanes_same_instruction_and_qualifiers_no_exited_lanes" -adapter = "c4_i32_a2_u32_b2_u32_metadata_u32_selector_u32_to_d4_i32" -llvm_adapter = "a2_i32_b2_i32_c4_i32_metadata_i32_selector_i32_to_d4_i32" -compatibility_source = "generated_stub" -runtime_validation = "unexecuted" -"#; - let parsed = toml::from_str::(valid).unwrap(); - assert_eq!(parsed.metadata, SparseMmaMetadata::Standard); - assert_eq!(parsed.selector, SparseMmaSelector::ImmediateZeroOrOne); - - let ordered = valid.replace("metadata = \"standard\"", "metadata = \"ordered\""); - assert_eq!( - toml::from_str::(&ordered).unwrap().metadata, - SparseMmaMetadata::Ordered - ); - - let k64 = ordered - .replace("shape = \"m16n8k32\"", "shape = \"m16n8k64\"") - .replace( - "selector = \"immediate_zero_or_one\"", - "selector = \"immediate_zero\"", - ) - .replace( - "adapter = \"c4_i32_a2_u32_b2_u32_metadata_u32_selector_u32_to_d4_i32\"", - "adapter = \"c4_i32_a4_u32_b4_u32_metadata_u32_selector_u32_to_d4_i32\"", - ) - .replace( - "llvm_adapter = \"a2_i32_b2_i32_c4_i32_metadata_i32_selector_i32_to_d4_i32\"", - "llvm_adapter = \"a4_i32_b4_i32_c4_i32_metadata_i32_selector_i32_to_d4_i32\"", - ); - let parsed_k64 = toml::from_str::(&k64).unwrap(); - assert_eq!(parsed_k64.shape, SparseMmaShape::M16n8k64); - assert_eq!(parsed_k64.selector, SparseMmaSelector::ImmediateZero); - assert_eq!( - parsed_k64.adapter, - SparseMmaAdapter::C4I32A4U32B4U32MetadataU32SelectorU32ToD4I32 - ); - assert_eq!( - parsed_k64.llvm_adapter, - SparseMmaLlvmAdapter::A4I32B4I32C4I32MetadataI32SelectorI32ToD4I32 - ); - - let int4 = ordered - .replace("shape = \"m16n8k32\"", "shape = \"m16n8k64\"") - .replace("a_element = \"s8\"", "a_element = \"s4\"") - .replace("b_element = \"u8\"", "b_element = \"u4\""); - let parsed_int4 = toml::from_str::(&int4).unwrap(); - assert_eq!(parsed_int4.a_element, SparseMmaElement::S4); - assert_eq!(parsed_int4.b_element, SparseMmaElement::U4); - - let k128_int4 = int4 - .replace("shape = \"m16n8k64\"", "shape = \"m16n8k128\"") - .replace( - "selector = \"immediate_zero_or_one\"", - "selector = \"immediate_zero\"", - ) - .replace( - "adapter = \"c4_i32_a2_u32_b2_u32_metadata_u32_selector_u32_to_d4_i32\"", - "adapter = \"c4_i32_a4_u32_b4_u32_metadata_u32_selector_u32_to_d4_i32\"", - ) - .replace( - "llvm_adapter = \"a2_i32_b2_i32_c4_i32_metadata_i32_selector_i32_to_d4_i32\"", - "llvm_adapter = \"a4_i32_b4_i32_c4_i32_metadata_i32_selector_i32_to_d4_i32\"", - ); - let parsed_k128_int4 = toml::from_str::(&k128_int4).unwrap(); - assert_eq!(parsed_k128_int4.shape, SparseMmaShape::M16n8k128); - assert_eq!(parsed_k128_int4.selector, SparseMmaSelector::ImmediateZero); - - for invalid in [ - valid.replace( - "selector = \"immediate_zero_or_one\"", - "selector = \"runtime\"", - ), - valid.replace("metadata = \"standard\"", "metadata = \"unreviewed\""), - valid.replace( - "llvm_adapter = \"a2_i32_b2_i32_c4_i32_metadata_i32_selector_i32_to_d4_i32\"", - "llvm_adapter = \"c_then_a_then_b\"", - ), - valid.replace("shape = \"m16n8k32\"", "shape = \"m16n8k256\""), - format!("{valid}unreviewed = true\n"), - ] { - assert!(toml::from_str::(&invalid).is_err(), "{invalid}"); - } - } - - #[test] - fn sparse_mma_admission_accepts_the_canonical_name_and_legacy_alias() { - let canonical = r#" -schema = 25 -family = "sparse_mma" - -[sparse_mma_integer] -llvm_evidence_profile = "llvm" -libnvvm_evidence_profile = "libnvvm" -runtime_validation = "unexecuted" -metadata = "ordered" - -[[sparse_mma_integer.variant]] -shape = "m16n8k64" -a_element = "s4" -b_element = "u4" -overflow = "wrapping" -"#; - let parsed = toml::from_str::(canonical).unwrap(); - assert_eq!( - parsed.sparse_mma_integer.unwrap().variants[0].b_element, - SparseMmaElement::U4 - ); - - let legacy = canonical.replace("sparse_mma_integer", "sparse_mma_int8"); - assert!( - toml::from_str::(&legacy) - .unwrap() - .sparse_mma_integer - .is_some() - ); - } - - #[test] - fn cp_async_mbarrier_contract_rejects_open_ended_policy() { - let valid = r#" -operation = "arrive_no_inc" -state_space = "shared" -adapter = "pointer_to_void" -runtime_validation = "unexecuted" -"#; - let parsed = toml::from_str::(valid).unwrap(); - assert_eq!(parsed.operation, CpAsyncMbarrierOperation::ArriveNoInc); - assert_eq!(parsed.state_space, CpAsyncMbarrierStateSpace::Shared); - - for invalid in [ - valid.replace("operation = \"arrive_no_inc\"", "operation = \"wait\""), - valid.replace("state_space = \"shared\"", "state_space = \"global\""), - valid.replace("adapter = \"pointer_to_void\"", "adapter = \"direct\""), - format!("{valid}unreviewed = true\n"), - ] { - assert!( - toml::from_str::(&invalid).is_err(), - "{invalid}" - ); - } - } - - #[test] - fn cluster_barrier_contract_rejects_open_ended_semantics() { - let valid = r#" -mode = "arrive_relaxed_aligned" -ordering = "relaxed" -aligned = true -"#; - let parsed = toml::from_str::(valid).unwrap(); - assert_eq!(parsed.mode, ClusterBarrierMode::ArriveRelaxedAligned); - assert_eq!(parsed.ordering, ClusterBarrierOrdering::Relaxed); - assert!(parsed.aligned); - - for invalid in [ - valid.replace("ordering = \"relaxed\"", "ordering = \"unordered\""), - valid.replace("aligned = true", "aligned = \"sometimes\""), - valid.replace( - "mode = \"arrive_relaxed_aligned\"", - "mode = \"arrive_release_aligned\"", - ), - format!("{valid}unreviewed = true\n"), - ] { - assert!( - toml::from_str::(&invalid).is_err(), - "{invalid}" - ); - } - } - - #[test] - fn debug_control_contract_rejects_open_ended_policy() { - let valid = r#" -operation = "pmevent" -adapter = "const_generic_to_immediate_u32" -runtime_validation = "unexecuted" -"#; - let parsed = toml::from_str::(valid).unwrap(); - assert_eq!(parsed.operation, DebugControlOperation::Pmevent); - assert_eq!( - parsed.adapter, - DebugControlAdapter::ConstGenericToImmediateU32 - ); - - for invalid in [ - valid.replace("operation = \"pmevent\"", "operation = \"profiler\""), - valid.replace( - "adapter = \"const_generic_to_immediate_u32\"", - "adapter = \"runtime_u32\"", - ), - valid.replace( - "runtime_validation = \"unexecuted\"", - "runtime_validation = \"assumed\"", - ), - format!("{valid}unreviewed = true\n"), - ] { - assert!( - toml::from_str::(&invalid).is_err(), - "{invalid}" - ); - } - } - - #[test] - fn cluster_memory_contract_rejects_open_ended_policy() { - let valid = r#" -operation = "map_shared_rank" -adapter = "generic_const_and_mut_pointer_rank_to_same_pointer" -source_contract = "llvm_mapa_shared_cluster_as7_identity_inline_ptx" -runtime_validation = "unexecuted" -"#; - let parsed = toml::from_str::(valid).unwrap(); - assert_eq!(parsed.operation, ClusterMemoryOperation::MapSharedRank); - assert_eq!( - parsed.source_contract, - ClusterMemorySourceContract::LlvmMapaSharedClusterAs7IdentityInlinePtx - ); - - for invalid in [ - valid.replace("map_shared_rank", "map_generic_rank"), - valid.replace( - "generic_const_and_mut_pointer_rank_to_same_pointer", - "direct", - ), - valid.replace( - "llvm_mapa_shared_cluster_as7_identity_inline_ptx", - "llvm_typed_as3", - ), - format!("{valid}unreviewed = true\n"), - ] { - assert!( - toml::from_str::(&invalid).is_err(), - "{invalid}" - ); - } - } - - #[test] - fn clc_contract_rejects_open_ended_policy() { - let valid = r#" -operation = "query_is_canceled" -adapter = "pair_u64_to_i128_bool_to_u32" -runtime_validation = "unexecuted" -"#; - let parsed = toml::from_str::(valid).unwrap(); - assert_eq!(parsed.operation, ClcOperation::QueryIsCanceled); - assert_eq!(parsed.adapter, ClcAdapter::PairU64ToI128BoolToU32); - - for invalid in [ - valid.replace("operation = \"query_is_canceled\"", "operation = \"query\""), - valid.replace( - "adapter = \"pair_u64_to_i128_bool_to_u32\"", - "adapter = \"pair_u64_to_i128\"", - ), - valid.replace( - "runtime_validation = \"unexecuted\"", - "runtime_validation = \"assumed\"", - ), - format!("{valid}unreviewed = true\n"), - ] { - assert!(toml::from_str::(&invalid).is_err(), "{invalid}"); - } - } - - #[test] - fn tma_contract_rejects_open_ended_policy() { - let valid = r#" -operation = "g2s_tile2d_multicast" -adapter = "g2s_pointers_coordinates_barrier_mask_inject_defaults" -runtime_validation = "unexecuted" -"#; - let parsed = toml::from_str::(valid).unwrap(); - assert_eq!(parsed.operation, TmaOperation::G2sTile2dMulticast); - assert_eq!( - parsed.adapter, - TmaAdapter::G2sPointersCoordinatesBarrierMaskInjectDefaults - ); - - for invalid in [ - valid.replace("g2s_tile2d_multicast", "g2s_multicast"), - valid.replace( - "g2s_pointers_coordinates_barrier_mask_inject_defaults", - "direct", - ), - valid.replace( - "runtime_validation = \"unexecuted\"", - "runtime_validation = \"assumed\"", - ), - format!("{valid}unreviewed = true\n"), - ] { - assert!(toml::from_str::(&invalid).is_err(), "{invalid}"); - } - } - - #[test] - fn tcgen05_execution_scope_is_closed() { - for operation in [ - Tcgen05Operation::Alloc, - Tcgen05Operation::Dealloc, - Tcgen05Operation::RelinquishAllocPermit, - Tcgen05Operation::AllocCg2, - Tcgen05Operation::DeallocCg2, - Tcgen05Operation::RelinquishAllocPermitCg2, - Tcgen05Operation::Ld16x256bX8Pure, - Tcgen05Operation::Ld16x256bPure, - Tcgen05Operation::LoadWait, - Tcgen05Operation::StoreWait, - Tcgen05Operation::Ld, - Tcgen05Operation::St, - ] { - assert_eq!(operation.execution_scope(), "warp"); - } - for operation in [ - Tcgen05Operation::FenceBeforeThreadSync, - Tcgen05Operation::FenceAfterThreadSync, - Tcgen05Operation::Commit, - Tcgen05Operation::CommitSharedCluster, - Tcgen05Operation::MmaWsF16, - Tcgen05Operation::MmaWsBf16, - Tcgen05Operation::MmaWsTf32, - Tcgen05Operation::MmaF16, - Tcgen05Operation::CpSmemToTmem, - Tcgen05Operation::MmaF16Cg2, - Tcgen05Operation::CommitCg2, - Tcgen05Operation::CommitSharedClusterCg2, - Tcgen05Operation::CommitMulticastCg2, - Tcgen05Operation::CpSmemToTmemCg2, - Tcgen05Operation::CommitMulticast, - Tcgen05Operation::ShiftDown, - Tcgen05Operation::ShiftDownCg2, - ] { - assert_eq!(operation.execution_scope(), "thread"); - } - } -} diff --git a/crates/cuda-intrinsics-gen/src/model/abi_ledger.rs b/crates/cuda-intrinsics-gen/src/model/abi_ledger.rs new file mode 100644 index 0000000000..9e9e496ad7 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/model/abi_ledger.rs @@ -0,0 +1,33 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use serde::Deserialize; + +#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct AbiLedgerFile { + pub schema: u32, + pub intrinsic_abi: u32, + #[serde(rename = "entry")] + pub entries: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct AbiLedgerEntry { + pub abi_id: String, + pub status: String, + pub catalog_id: String, + pub operation_key: String, + pub raw_rust_signature: AbiRawRustSignature, +} + +#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct AbiRawRustSignature { + pub safe: bool, + pub arguments: Vec, + pub result: String, +} diff --git a/crates/cuda-intrinsics-gen/src/model/catalog.rs b/crates/cuda-intrinsics-gen/src/model/catalog.rs new file mode 100644 index 0000000000..064d40446c --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/model/catalog.rs @@ -0,0 +1,458 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use super::contracts::{ + ActiveMask, Clc, ClusterBarrier, ClusterMemory, CpAsyncControl, CpAsyncCopy, CpAsyncMbarrier, + DebugControl, DotProduct, ExtendedMinMax, IntegerMinMax, LdmatrixAdapter, LdmatrixSafety, + LdmatrixVariant, MbarrierBasic, MbarrierExtended, Movmatrix, PackedAlu, PackedAtomic, + PackedConversion, Prmt, Redux, RegisterMma, ScalarArithmetic, ScalarConversion, ScalarMath, + SparseMma, SpecialRegister, Tcgen05, Tma, Vote, WarpBarrier, WarpMatch, WarpShuffle, + WgmmaControl, +}; +use super::core::{BackendLoweringMechanism, IntrinsicBackend, IntrinsicSource}; +use super::evidence::EvidenceStage; +use super::imported::{ImportedAddressSpace, ImportedSelectionConstraints}; +use crate::ptx::InstructionPattern; +use serde::{Deserialize, Serialize}; +use std::{fmt, str::FromStr}; + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct CatalogFile { + pub schema: u32, + pub catalog_version: String, + pub intrinsic_abi: u32, + pub generator_version: String, + pub source: CatalogSource, + pub inputs: CatalogInputs, + pub intrinsics: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct CatalogSource { + pub llvm_repository: String, + pub llvm_revision: String, + pub llvm_tblgen_version: String, + pub llvm_tblgen_source_revision: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct CatalogInputs { + pub imported_sha256: String, + pub overlay_sha256: String, + pub abi_ledger_sha256: String, + pub evidence_sha256: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct CatalogIntrinsic { + pub id: String, + pub operation_key: String, + pub family: String, + pub source: IntrinsicSource, + pub selections: Vec, + pub rust: CatalogRust, + pub dialect: CatalogDialect, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub llvm: Option, + pub semantics: CatalogSemantics, + pub target: CatalogTarget, + pub backend: CatalogBackend, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub backend_lowerings: Vec, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub packed_atomic: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub redux: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub vote: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub active_mask: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub warp_match: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub warp_barrier: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub warp_shuffle: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub dot_product: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub packed_alu: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub integer_minmax: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub packed_conversion: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub scalar_conversion: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub scalar_arithmetic: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub scalar_math: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub extended_minmax: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cp_async_copy: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cp_async_control: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cp_async_mbarrier: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub mbarrier_basic: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub movmatrix: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub mbarrier_extended: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub register_mma: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub sparse_mma: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub prmt: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cluster_barrier: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub wgmma_control: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub special_register: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub debug_control: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cluster_memory: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub clc: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tma: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tcgen05: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub ldmatrix: Option, + pub lowering: String, + pub expected_ptx: InstructionPattern, + pub summary: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct CatalogSelection { + pub source_record: String, + pub asm: String, + pub predicates: Vec, + #[serde( + default, + skip_serializing_if = "ImportedSelectionConstraints::is_empty" + )] + pub constraints: ImportedSelectionConstraints, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct CatalogRust { + pub abi_id: String, + pub module: String, + pub name: String, + pub arguments: Vec, + pub result: String, + pub safe: bool, + pub must_use: bool, + pub safe_allowlist_reason: Option, + pub canonical_path: String, + pub public_path: String, + pub compatibility_paths: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct CatalogDialect { + pub op_type: String, + pub op_name: String, + pub operands: Vec, + pub results: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct CatalogLlvm { + pub symbol: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub resolved_symbol: Option, + pub arguments: Vec, + pub results: Vec, + pub properties: Vec, + pub result_facts: CatalogLlvmResultFacts, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct CatalogLlvmResultFacts { + pub no_undef: bool, + pub range: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct CatalogHalfOpenRange { + pub lower: String, + pub upper_exclusive: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct CatalogSemantics { + pub pure: bool, + pub memory: String, + pub convergent: bool, + pub execution_scope: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct CatalogTarget { + pub minimum_ptx: PtxVersion, + pub hardware: CatalogHardwareTarget, + pub ptx_result: String, + pub targets: String, + pub ptx_isa_version: String, + pub ptx_isa_section: String, + pub ptx_isa_url: String, +} + +/// A PTX ISA version encoded as `major * 10 + minor`. +/// +/// PTX currently uses one decimal minor digit. The resolver validates that +/// shape before constructing this value, so generated consumers compare a +/// number rather than reparsing policy text. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub struct PtxVersion(u16); + +impl PtxVersion { + pub const fn encoded(self) -> u16 { + self.0 + } + + pub const fn major(self) -> u16 { + self.0 / 10 + } + + pub const fn minor(self) -> u16 { + self.0 % 10 + } +} + +impl FromStr for PtxVersion { + type Err = String; + + fn from_str(value: &str) -> Result { + let (major, minor) = value + .split_once('.') + .ok_or_else(|| "expected major.minor".to_owned())?; + if major.is_empty() + || !major.bytes().all(|byte| byte.is_ascii_digit()) + || minor.len() != 1 + || !minor.bytes().all(|byte| byte.is_ascii_digit()) + { + return Err("expected numeric major.minor with one minor digit".to_owned()); + } + let major: u16 = major.parse().map_err(|_| "major version is too large")?; + let minor: u16 = minor.parse().unwrap(); + if format!("{major}.{minor}") != value { + return Err("version is not in canonical major.minor form".to_owned()); + } + let encoded = major + .checked_mul(10) + .and_then(|value| value.checked_add(minor)) + .ok_or_else(|| "version is too large".to_owned())?; + Ok(Self(encoded)) + } +} + +impl Serialize for PtxVersion { + fn serialize(&self, serializer: S) -> Result + where + S: serde::Serializer, + { + serializer.serialize_str(&self.to_string()) + } +} + +impl<'de> Deserialize<'de> for PtxVersion { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + let value = String::deserialize(deserializer)?; + value.parse().map_err(serde::de::Error::custom) + } +} + +impl fmt::Display for PtxVersion { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(formatter, "{}.{}", self.major(), self.minor()) + } +} + +/// Reviewed hardware availability for an intrinsic. +/// +/// Exact `a` and `f` targets stay distinct from monotonic minimums. A target +/// matrix also keeps each hardware target paired with its PTX floor. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "kind", rename_all = "snake_case")] +pub enum CatalogHardwareTarget { + All, + AnyOf { + alternatives: Vec, + }, + /// Closed selector contracts with their exact PTX and hardware pairs. + TargetMatrix { + contracts: Vec, + }, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "kind", rename_all = "snake_case")] +pub enum CatalogHardwareAlternative { + MinimumSm { sm: u16 }, + ExactArchitecture { sm: u16 }, + FamilyTarget { sm: u16 }, +} + +/// One PTX floor paired with one hardware alternative. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct CatalogTargetAlternative { + pub minimum_ptx: PtxVersion, + pub hardware: CatalogHardwareAlternative, +} + +/// One selector tuple and its exact PTX and hardware pairs. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct CatalogTargetContract { + pub selectors: Vec, + pub alternatives: Vec, +} + +/// One field/value pair that selects a target contract. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct TargetSelectorBinding { + pub name: String, + pub value: String, +} + +/// One selector-specific target contract from admission policy. +#[allow(dead_code)] +#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct TargetContract { + #[serde(default)] + pub selectors: Vec, + pub alternatives: Vec, +} + +/// One target spelling and PTX floor. +#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct TargetContractAlternative { + pub target: String, + pub minimum_ptx: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct CatalogBackend { + pub profile: String, + pub version: String, + pub sha256: String, + pub status: String, + pub target_triple: String, + pub gpu_target: String, + pub ptx_feature: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct CatalogBackendLowering { + pub backend: IntrinsicBackend, + pub mechanism: BackendLoweringMechanism, + pub evidence_profile: String, + pub target: CatalogTargetRequirement, + pub version: String, + pub sha256: String, + pub artifact_path: Option, + pub build_id_prefix: Option, + pub status: String, + pub stages: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct CatalogTargetRequirement { + pub minimum_ptx: PtxVersion, + pub hardware: CatalogHardwareTarget, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct CatalogLdmatrix { + pub variant: LdmatrixVariant, + pub safety: LdmatrixSafety, + pub adapter: LdmatrixAdapter, + pub selected_address_space: ImportedAddressSpace, +} + +impl CatalogIntrinsic { + pub fn scalar_width(&self) -> Option { + match self.rust.result.as_str() { + "u32" => Some(32), + "u64" => Some(64), + _ => None, + } + } + + pub fn llvm_identifier(&self) -> String { + llvm_symbol_to_identifier(&self.llvm.as_ref().expect("LLVM-backed intrinsic").symbol) + } + + pub fn resolved_llvm_identifier(&self) -> String { + let llvm = self.llvm.as_ref().expect("LLVM-backed intrinsic"); + llvm_symbol_to_identifier(llvm.resolved_symbol.as_deref().unwrap_or(&llvm.symbol)) + } +} + +fn llvm_symbol_to_identifier(symbol: &str) -> String { + if !symbol.contains('_') { + return symbol.replace('.', "_"); + } + + let suffix = symbol + .strip_prefix("llvm.") + .expect("LLVM intrinsic symbol must start with llvm."); + let mut output = String::from("llvm__"); + for ch in suffix.chars() { + match ch { + '.' => output.push_str("_d"), + '_' => output.push_str("_u"), + ch if ch.is_ascii_alphanumeric() => output.push(ch), + _ => panic!("LLVM intrinsic symbol contains an unsupported character"), + } + } + output +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn llvm_identifier_encoding_preserves_literal_underscores() { + assert_eq!( + llvm_symbol_to_identifier("llvm.nvvm.read.ptx.sreg.tid.x"), + "llvm_nvvm_read_ptx_sreg_tid_x" + ); + assert_eq!( + llvm_symbol_to_identifier("llvm.nvvm.wgmma.wait_group.sync.aligned"), + "llvm__nvvm_dwgmma_dwait_ugroup_dsync_daligned" + ); + assert_eq!( + llvm_symbol_to_identifier( + "llvm.nvvm.ldmatrix.sync.aligned.m16n16.x1.trans.b8x16.b4x16_p64.p3" + ), + "llvm__nvvm_dldmatrix_dsync_daligned_dm16n16_dx1_dtrans_db8x16_db4x16_up64_dp3" + ); + assert_eq!( + llvm_symbol_to_identifier( + "llvm.nvvm.ldmatrix.sync.aligned.m8n16.x4.b8x16.b6x16_p32.p3" + ), + "llvm__nvvm_dldmatrix_dsync_daligned_dm8n16_dx4_db8x16_db6x16_up32_dp3" + ); + } +} diff --git a/crates/cuda-intrinsics-gen/src/model/contracts/control.rs b/crates/cuda-intrinsics-gen/src/model/contracts/control.rs new file mode 100644 index 0000000000..10225ce0ba --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/model/contracts/control.rs @@ -0,0 +1,266 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use super::super::core::RuntimeValidation; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum WgmmaControlMode { + Fence, + CommitGroup, + WaitGroup, +} + +/// Closed semantics for one WGMMA control operation. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct WgmmaControl { + pub mode: WgmmaControlMode, + pub adapter: WgmmaControlAdapter, + pub participation: WgmmaControlParticipation, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum WgmmaControlAdapter { + NoArguments, + ConstGenericU32ToI64Immediate, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum WgmmaControlParticipation { + WarpgroupAllThreadsSameInstruction, +} + +/// Closed semantic and lowering contract for one special-register read. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct SpecialRegister { + pub register: SpecialRegisterKind, + pub observation: SpecialRegisterObservation, + pub result_width: SpecialRegisterWidth, + pub ptx_type: SpecialRegisterPtxType, + pub output_constraint: SpecialRegisterOutputConstraint, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub llvm_exclusion: Option, +} + +/// Closed semantic contract for PTX debug controls. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct DebugControl { + pub operation: DebugControlOperation, + pub adapter: DebugControlAdapter, + pub runtime_validation: RuntimeValidation, +} + +/// Closed semantic contract for Cluster Launch Control. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Clc { + pub operation: ClcOperation, + pub adapter: ClcAdapter, + pub runtime_validation: RuntimeValidation, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ClcOperation { + TryCancel, + TryCancelMulticast, + QueryIsCanceled, + QueryGetFirstCtaidX, + QueryGetFirstCtaidY, + QueryGetFirstCtaidZ, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ClcAdapter { + GenericPointersToShared, + PairU64ToI128BoolToU32, + PairU64ToI128U32, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SpecialRegisterKind { + Clock, + Clock64, + Globaltimer, + Envreg1, + Envreg2, + Smid, + Nsmid, + Gridid, + Warpid, + Nwarpid, + DynamicSmemSize, + TotalSmemSize, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum DebugControlOperation { + Trap, + Breakpoint, + Pmevent, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SpecialRegisterObservation { + StablePure, + VolatileObservation, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SpecialRegisterWidth { + B32, + B64, +} + +impl SpecialRegisterWidth { + pub const fn bits(self) -> u32 { + match self { + Self::B32 => 32, + Self::B64 => 64, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SpecialRegisterPtxType { + B32, + U32, + U64, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SpecialRegisterOutputConstraint { + Register32, + Register64, +} + +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct SpecialRegisterLlvmExclusion { + pub source_record: String, + pub llvm_symbol: String, + pub imported_result_width: SpecialRegisterWidth, + pub reason: SpecialRegisterLlvmExclusionReason, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SpecialRegisterLlvmExclusionReason { + ResultWidthMismatch, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum DebugControlAdapter { + Direct, + ConstGenericToImmediateU32, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn wgmma_control_contract_is_closed() { + let parsed: WgmmaControl = serde_json::from_str( + r#"{ + "mode": "wait_group", + "adapter": "const_generic_u32_to_i64_immediate", + "participation": "warpgroup_all_threads_same_instruction" + }"#, + ) + .unwrap(); + assert_eq!(parsed.mode, WgmmaControlMode::WaitGroup); + assert_eq!( + parsed.adapter, + WgmmaControlAdapter::ConstGenericU32ToI64Immediate + ); + assert_eq!( + parsed.participation, + WgmmaControlParticipation::WarpgroupAllThreadsSameInstruction + ); + + let open_ended = r#"{ + "mode": "wait_group", + "adapter": "const_generic_u32_to_i64_immediate", + "participation": "warpgroup_all_threads_same_instruction", + "custom_ptx": "wgmma.wait_group.sync.aligned 0;" + }"#; + assert!(serde_json::from_str::(open_ended).is_err()); + } + + #[test] + fn debug_control_contract_rejects_open_ended_policy() { + let valid = r#" +operation = "pmevent" +adapter = "const_generic_to_immediate_u32" +runtime_validation = "unexecuted" +"#; + let parsed = toml::from_str::(valid).unwrap(); + assert_eq!(parsed.operation, DebugControlOperation::Pmevent); + assert_eq!( + parsed.adapter, + DebugControlAdapter::ConstGenericToImmediateU32 + ); + + for invalid in [ + valid.replace("operation = \"pmevent\"", "operation = \"profiler\""), + valid.replace( + "adapter = \"const_generic_to_immediate_u32\"", + "adapter = \"runtime_u32\"", + ), + valid.replace( + "runtime_validation = \"unexecuted\"", + "runtime_validation = \"assumed\"", + ), + format!("{valid}unreviewed = true\n"), + ] { + assert!( + toml::from_str::(&invalid).is_err(), + "{invalid}" + ); + } + } + + #[test] + fn clc_contract_rejects_open_ended_policy() { + let valid = r#" +operation = "query_is_canceled" +adapter = "pair_u64_to_i128_bool_to_u32" +runtime_validation = "unexecuted" +"#; + let parsed = toml::from_str::(valid).unwrap(); + assert_eq!(parsed.operation, ClcOperation::QueryIsCanceled); + assert_eq!(parsed.adapter, ClcAdapter::PairU64ToI128BoolToU32); + + for invalid in [ + valid.replace("operation = \"query_is_canceled\"", "operation = \"query\""), + valid.replace( + "adapter = \"pair_u64_to_i128_bool_to_u32\"", + "adapter = \"pair_u64_to_i128\"", + ), + valid.replace( + "runtime_validation = \"unexecuted\"", + "runtime_validation = \"assumed\"", + ), + format!("{valid}unreviewed = true\n"), + ] { + assert!(toml::from_str::(&invalid).is_err(), "{invalid}"); + } + } +} diff --git a/crates/cuda-intrinsics-gen/src/model/contracts/matrix.rs b/crates/cuda-intrinsics-gen/src/model/contracts/matrix.rs new file mode 100644 index 0000000000..d62f5562d1 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/model/contracts/matrix.rs @@ -0,0 +1,605 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use super::super::core::RuntimeValidation; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum StmatrixMultiplicity { + X2, + X4, +} + +impl StmatrixMultiplicity { + pub const fn register_count(self) -> usize { + match self { + Self::X2 => 2, + Self::X4 => 4, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum StmatrixLayout { + Normal, + Transposed, +} + +/// Closed semantic identity for the generated `ldmatrix` family. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct LdmatrixVariant { + pub shape: LdmatrixShape, + pub multiplicity: LdmatrixMultiplicity, + pub layout: LdmatrixLayout, + pub element: LdmatrixElement, + pub state_space: LdmatrixStateSpace, +} + +impl LdmatrixVariant { + pub const fn register_count(&self) -> usize { + let matrices = self.multiplicity.register_count(); + match self.shape { + LdmatrixShape::M8n8 | LdmatrixShape::M8n16 => matrices, + LdmatrixShape::M16n16 => matrices * 2, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum LdmatrixShape { + M8n8, + M8n16, + M16n16, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum LdmatrixMultiplicity { + X1, + X2, + X4, +} + +impl LdmatrixMultiplicity { + pub const fn register_count(self) -> usize { + match self { + Self::X1 => 1, + Self::X2 => 2, + Self::X4 => 4, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum LdmatrixLayout { + Normal, + Transposed, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum LdmatrixElement { + B16, + B8, + B8x16B4x16P64, + B8x16B6x16P32, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum LdmatrixStateSpace { + Shared, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct LdmatrixSafety { + pub participation: LdmatrixParticipation, + pub address_contract: LdmatrixAddressContract, + pub memory_order: LdmatrixMemoryOrder, + pub runtime_validation: RuntimeValidation, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum LdmatrixParticipation { + AllWarpLanesSameInstruction, + AllWarpLanesSameInstructionAndQualifiersNoExitedLanes, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum LdmatrixAddressContract { + WarpLaneAddressesMappedByMultiplicitySixteenByteAlignedSixteenBytesReadableWithSm75Replication, + WarpLaneAddressesMappedByMultiplicitySixteenByteAlignedSixteenBytesReadable, + WarpLaneAddressesMappedByMultiplicitySixteenByteAlignedThirtyTwoBytesReadable, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum LdmatrixMemoryOrder { + Weak, +} + +/// Closed contract for the in-register warp matrix transpose. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Movmatrix { + pub participation: MovmatrixParticipation, + pub adapter: MovmatrixAdapter, + pub runtime_validation: RuntimeValidation, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum MovmatrixParticipation { + AllWarpLanesSameInstructionNoExitedLanes, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum MovmatrixAdapter { + PackedB16x2U32ToPackedB16x2U32, +} + +/// Closed contract for register-only warp-level `mma.sync` operations. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct RegisterMma { + pub shape: RegisterMmaShape, + pub operation: RegisterMmaOperation, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub kind: Option, + pub accumulator: RegisterMmaAccumulator, + pub a_element: RegisterMmaElement, + pub b_element: RegisterMmaElement, + pub a_layout: RegisterMmaLayout, + pub b_layout: RegisterMmaLayout, + pub overflow: RegisterMmaOverflow, + pub participation: RegisterMmaParticipation, + pub adapter: RegisterMmaAdapter, + pub compatibility_source: RegisterMmaCompatibilitySource, + pub runtime_validation: RuntimeValidation, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum RegisterMmaKind { + Standard, + F8f6f4, + Mxf8f6f4, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum RegisterMmaShape { + M8n8k4, + M8n8k16, + M8n8k32, + M16n8k4, + M16n8k8, + M16n8k16, + M16n8k32, + M16n8k64, + M8n8k128, + M16n8k128, + M16n8k256, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum RegisterMmaOperation { + Multiply, + AndPopc, + XorPopc, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum RegisterMmaAccumulator { + F16, + F32, + F64, + S32, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum RegisterMmaElement { + B1, + Bf16, + E2m1, + E2m3, + E3m2, + E4m3, + E5m2, + F16, + Tf32, + F64, + S4, + U4, + S8, + U8, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum RegisterMmaLayout { + Row, + Col, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum RegisterMmaOverflow { + NotApplicable, + Wrapping, + Satfinite, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum RegisterMmaParticipation { + AllWarpLanesSameInstructionAndQualifiersNoExitedLanes, +} + +/// Rust `C, A, B` fragment shape used by the importer. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum RegisterMmaAdapter { + C2U32A2U32B1U32ToD2U32, + C2U32A4U32B2U32ToD2U32, + C4F32A2U32B1U32ToD4F32, + C4F32A4U32B2U32ToD4F32, + C2F64A1F64B1F64ToD2F64, + C2I32A1U32B1U32ToD2I32, + C4I32A4U32B2U32ToD4I32, + C4I32A2U32B1U32ToD4I32, + C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32, +} + +/// Where the stable `cuda_device::wmma` callable is defined. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum RegisterMmaCompatibilitySource { + ExistingStub, + GeneratedStub, +} + +/// Closed semantic contract for register-only sparse `mma.sp` operations. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct SparseMma { + pub shape: SparseMmaShape, + pub accumulator: SparseMmaAccumulator, + pub a_element: SparseMmaElement, + pub b_element: SparseMmaElement, + pub a_layout: SparseMmaLayout, + pub b_layout: SparseMmaLayout, + pub overflow: SparseMmaOverflow, + pub metadata: SparseMmaMetadata, + pub selector: SparseMmaSelector, + pub participation: SparseMmaParticipation, + pub adapter: SparseMmaAdapter, + pub llvm_adapter: SparseMmaLlvmAdapter, + pub compatibility_source: SparseMmaCompatibilitySource, + pub runtime_validation: RuntimeValidation, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SparseMmaShape { + M16n8k32, + M16n8k64, + M16n8k128, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SparseMmaAccumulator { + F16, + F32, + S32, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SparseMmaElement { + E2m1, + E2m3, + E3m2, + E4m3, + E5m2, + S4, + U4, + S8, + U8, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SparseMmaLayout { + Row, + Col, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SparseMmaOverflow { + NotApplicable, + Wrapping, + Satfinite, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SparseMmaMetadata { + Standard, + Ordered, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SparseMmaSelector { + ImmediateZeroOrOne, + ImmediateZero, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SparseMmaParticipation { + AllWarpLanesSameInstructionAndQualifiersNoExitedLanes, +} + +/// Rust `C, A, B, metadata, selector` shape used by the importer. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SparseMmaAdapter { + C2U32A4U32B4U32MetadataU32SelectorU32ToD2U32, + C4F32A4U32B4U32MetadataU32SelectorU32ToD4F32, + C4I32A2U32B2U32MetadataU32SelectorU32ToD4I32, + C4I32A4U32B4U32MetadataU32SelectorU32ToD4I32, +} + +/// LLVM `A, B, C, metadata, selector` shape. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SparseMmaLlvmAdapter { + A4I32B4I32C2V2F16MetadataI32SelectorI32ToD2V2F16, + A4I32B4I32C4F32MetadataI32SelectorI32ToD4F32, + A2I32B2I32C4I32MetadataI32SelectorI32ToD4I32, + A4I32B4I32C4I32MetadataI32SelectorI32ToD4I32, +} + +/// Where the stable `cuda_device::wmma` callable is defined. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SparseMmaCompatibilitySource { + GeneratedStub, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum LdmatrixAdapter { + SingleResultDirect, + MultipleResultsToArray, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn ldmatrix_register_count_includes_the_shape_width() { + let variant = |shape, multiplicity| LdmatrixVariant { + shape, + multiplicity, + layout: LdmatrixLayout::Normal, + element: LdmatrixElement::B8, + state_space: LdmatrixStateSpace::Shared, + }; + + assert_eq!( + variant(LdmatrixShape::M8n16, LdmatrixMultiplicity::X4).register_count(), + 4 + ); + assert_eq!( + variant(LdmatrixShape::M16n16, LdmatrixMultiplicity::X1).register_count(), + 2 + ); + assert_eq!( + variant(LdmatrixShape::M16n16, LdmatrixMultiplicity::X2).register_count(), + 4 + ); + } + + #[test] + fn blackwell_ldmatrix_address_contracts_keep_readable_widths_distinct() { + assert_eq!( + serde_json::from_str::( + r#""warp_lane_addresses_mapped_by_multiplicity_sixteen_byte_aligned_sixteen_bytes_readable""# + ) + .unwrap(), + LdmatrixAddressContract::WarpLaneAddressesMappedByMultiplicitySixteenByteAlignedSixteenBytesReadable + ); + assert_eq!( + serde_json::from_str::( + r#""warp_lane_addresses_mapped_by_multiplicity_sixteen_byte_aligned_thirty_two_bytes_readable""# + ) + .unwrap(), + LdmatrixAddressContract::WarpLaneAddressesMappedByMultiplicitySixteenByteAlignedThirtyTwoBytesReadable + ); + } + + #[test] + fn movmatrix_contract_rejects_open_ended_policy() { + let valid = r#" +participation = "all_warp_lanes_same_instruction_no_exited_lanes" +adapter = "packed_b16x2_u32_to_packed_b16x2_u32" +runtime_validation = "unexecuted" +"#; + let parsed = toml::from_str::(valid).unwrap(); + assert_eq!( + parsed.participation, + MovmatrixParticipation::AllWarpLanesSameInstructionNoExitedLanes + ); + + for invalid in [ + valid.replace( + "all_warp_lanes_same_instruction_no_exited_lanes", + "participating_lanes", + ), + valid.replace("packed_b16x2_u32_to_packed_b16x2_u32", "direct_u32"), + format!("{valid}unreviewed = true\n"), + ] { + assert!(toml::from_str::(&invalid).is_err(), "{invalid}"); + } + } + + #[test] + fn register_mma_contract_parses_unsigned_k16_generated_stub() { + let valid = r#" +shape = "m16n8k16" +operation = "multiply" +accumulator = "s32" +a_element = "s8" +b_element = "u8" +a_layout = "row" +b_layout = "col" +overflow = "satfinite" +participation = "all_warp_lanes_same_instruction_and_qualifiers_no_exited_lanes" +adapter = "c4_i32_a2_u32_b1_u32_to_d4_i32" +compatibility_source = "generated_stub" +runtime_validation = "unexecuted" +"#; + let parsed = toml::from_str::(valid).unwrap(); + assert_eq!(parsed.b_element, RegisterMmaElement::U8); + assert_eq!(parsed.adapter, RegisterMmaAdapter::C4I32A2U32B1U32ToD4I32); + assert_eq!( + parsed.compatibility_source, + RegisterMmaCompatibilitySource::GeneratedStub + ); + + for invalid in [ + valid.replace("b_element = \"u8\"", "b_element = \"i8\""), + valid.replace( + "adapter = \"c4_i32_a2_u32_b1_u32_to_d4_i32\"", + "adapter = \"direct\"", + ), + valid.replace( + "compatibility_source = \"generated_stub\"", + "compatibility_source = \"automatic\"", + ), + format!("{valid}unreviewed = true\n"), + ] { + assert!( + toml::from_str::(&invalid).is_err(), + "{invalid}" + ); + } + } + + #[test] + fn sparse_mma_contract_closes_the_selector_and_metadata_modes() { + let valid = r#" +shape = "m16n8k32" +accumulator = "s32" +a_element = "s8" +b_element = "u8" +a_layout = "row" +b_layout = "col" +overflow = "satfinite" +metadata = "standard" +selector = "immediate_zero_or_one" +participation = "all_warp_lanes_same_instruction_and_qualifiers_no_exited_lanes" +adapter = "c4_i32_a2_u32_b2_u32_metadata_u32_selector_u32_to_d4_i32" +llvm_adapter = "a2_i32_b2_i32_c4_i32_metadata_i32_selector_i32_to_d4_i32" +compatibility_source = "generated_stub" +runtime_validation = "unexecuted" +"#; + let parsed = toml::from_str::(valid).unwrap(); + assert_eq!(parsed.metadata, SparseMmaMetadata::Standard); + assert_eq!(parsed.selector, SparseMmaSelector::ImmediateZeroOrOne); + + let ordered = valid.replace("metadata = \"standard\"", "metadata = \"ordered\""); + assert_eq!( + toml::from_str::(&ordered).unwrap().metadata, + SparseMmaMetadata::Ordered + ); + + let k64 = ordered + .replace("shape = \"m16n8k32\"", "shape = \"m16n8k64\"") + .replace( + "selector = \"immediate_zero_or_one\"", + "selector = \"immediate_zero\"", + ) + .replace( + "adapter = \"c4_i32_a2_u32_b2_u32_metadata_u32_selector_u32_to_d4_i32\"", + "adapter = \"c4_i32_a4_u32_b4_u32_metadata_u32_selector_u32_to_d4_i32\"", + ) + .replace( + "llvm_adapter = \"a2_i32_b2_i32_c4_i32_metadata_i32_selector_i32_to_d4_i32\"", + "llvm_adapter = \"a4_i32_b4_i32_c4_i32_metadata_i32_selector_i32_to_d4_i32\"", + ); + let parsed_k64 = toml::from_str::(&k64).unwrap(); + assert_eq!(parsed_k64.shape, SparseMmaShape::M16n8k64); + assert_eq!(parsed_k64.selector, SparseMmaSelector::ImmediateZero); + assert_eq!( + parsed_k64.adapter, + SparseMmaAdapter::C4I32A4U32B4U32MetadataU32SelectorU32ToD4I32 + ); + assert_eq!( + parsed_k64.llvm_adapter, + SparseMmaLlvmAdapter::A4I32B4I32C4I32MetadataI32SelectorI32ToD4I32 + ); + + let int4 = ordered + .replace("shape = \"m16n8k32\"", "shape = \"m16n8k64\"") + .replace("a_element = \"s8\"", "a_element = \"s4\"") + .replace("b_element = \"u8\"", "b_element = \"u4\""); + let parsed_int4 = toml::from_str::(&int4).unwrap(); + assert_eq!(parsed_int4.a_element, SparseMmaElement::S4); + assert_eq!(parsed_int4.b_element, SparseMmaElement::U4); + + let k128_int4 = int4 + .replace("shape = \"m16n8k64\"", "shape = \"m16n8k128\"") + .replace( + "selector = \"immediate_zero_or_one\"", + "selector = \"immediate_zero\"", + ) + .replace( + "adapter = \"c4_i32_a2_u32_b2_u32_metadata_u32_selector_u32_to_d4_i32\"", + "adapter = \"c4_i32_a4_u32_b4_u32_metadata_u32_selector_u32_to_d4_i32\"", + ) + .replace( + "llvm_adapter = \"a2_i32_b2_i32_c4_i32_metadata_i32_selector_i32_to_d4_i32\"", + "llvm_adapter = \"a4_i32_b4_i32_c4_i32_metadata_i32_selector_i32_to_d4_i32\"", + ); + let parsed_k128_int4 = toml::from_str::(&k128_int4).unwrap(); + assert_eq!(parsed_k128_int4.shape, SparseMmaShape::M16n8k128); + assert_eq!(parsed_k128_int4.selector, SparseMmaSelector::ImmediateZero); + + for invalid in [ + valid.replace( + "selector = \"immediate_zero_or_one\"", + "selector = \"runtime\"", + ), + valid.replace("metadata = \"standard\"", "metadata = \"unreviewed\""), + valid.replace( + "llvm_adapter = \"a2_i32_b2_i32_c4_i32_metadata_i32_selector_i32_to_d4_i32\"", + "llvm_adapter = \"c_then_a_then_b\"", + ), + valid.replace("shape = \"m16n8k32\"", "shape = \"m16n8k256\""), + format!("{valid}unreviewed = true\n"), + ] { + assert!(toml::from_str::(&invalid).is_err(), "{invalid}"); + } + } +} diff --git a/crates/cuda-intrinsics-gen/src/model/contracts/memory.rs b/crates/cuda-intrinsics-gen/src/model/contracts/memory.rs new file mode 100644 index 0000000000..be03a8862d --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/model/contracts/memory.rs @@ -0,0 +1,464 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use super::super::core::RuntimeValidation; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum MbarrierExtendedOperation { + ArriveExpectTxCta, + ArriveExpectTxCluster, + ArriveRemoteCluster, + TryWaitTokenCta, + TryWaitParityCta, + TryWaitParityCluster, + FenceProxyAsyncSharedCta, + FenceMbarrierInitReleaseCluster, + FenceProxyAsyncGenericReleaseSharedCtaCluster, + FenceProxyAsyncGenericAcquireSharedClusterCluster, + Nanosleep, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ClusterMemoryOperation { + MapSharedRank, + ReadU32, +} + +/// Closed contract for cluster address mapping and remote shared reads. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ClusterMemory { + pub operation: ClusterMemoryOperation, + pub adapter: ClusterMemoryAdapter, + pub source_contract: ClusterMemorySourceContract, + pub runtime_validation: RuntimeValidation, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ClusterMemoryAdapter { + GenericConstAndMutPointerRankToSamePointer, + ConstU32PointerRankToU32, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ClusterMemorySourceContract { + LlvmMapaSharedClusterAs7IdentityInlinePtx, + PtxNativeMapaThenWeakClusterLoad, +} + +/// Closed semantic contract for the generated packed global atomic-add +/// family. These fields are intentionally enums rather than free-form strings: +/// accepting an unreviewed state space, scope, or floating-point mode must +/// require a generator change. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct PackedAtomic { + pub format: PackedAtomicFormat, + /// PTX ISA hardware floor, kept separate from cuda-oxide's admitted floor + /// and from backend-profile floors. + pub native_minimum_sm: u16, + pub operation: PackedAtomicOperation, + pub state_space: PackedAtomicStateSpace, + pub ordering: PackedAtomicOrdering, + pub scope: PackedAtomicScope, + pub rounding: PackedAtomicRounding, + pub subnormal: PackedAtomicSubnormal, + pub atomicity: PackedAtomicAtomicity, + pub pointer_contract: PackedAtomicPointerContract, + pub access_contract: PackedAtomicAccessContract, + pub scope_contract: PackedAtomicScopeContract, + pub codegen_contract: PackedAtomicCodegenContract, + pub return_contract: PackedAtomicReturnContract, + pub adapter: PackedAtomicAdapter, + pub runtime_validation: RuntimeValidation, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PackedAtomicFormat { + F16x2, + Bf16x2, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PackedAtomicOperation { + Add, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PackedAtomicStateSpace { + Global, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PackedAtomicOrdering { + Relaxed, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PackedAtomicScope { + Gpu, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PackedAtomicRounding { + NearestEven, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PackedAtomicSubnormal { + Preserve, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PackedAtomicAtomicity { + PerElement, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PackedAtomicPointerContract { + MutableGlobalU32Aligned4, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PackedAtomicAccessContract { + NoMixedWholeWordOrNonAtomicAccess, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PackedAtomicScopeContract { + RacingAtomicsMutuallyInclusive, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PackedAtomicCodegenContract { + ExactNativeInstruction, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PackedAtomicReturnContract { + OldValuesPerElementMayBeNoncoherent, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PackedAtomicAdapter { + OldPackedU32, +} + +/// Closed contract for classic global-to-shared `cp.async` copies. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct CpAsyncCopy { + pub cache_policy: CpAsyncCachePolicy, + pub copy_size: CpAsyncCopySize, + pub source_size: CpAsyncSourceSize, + pub adapter: CpAsyncAdapter, + pub runtime_validation: RuntimeValidation, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum CpAsyncCachePolicy { + Ca, + Cg, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum CpAsyncCopySize { + B4, + B8, + B16, +} + +impl CpAsyncCopySize { + pub const fn bytes(self) -> u32 { + match self { + Self::B4 => 4, + Self::B8 => 8, + Self::B16 => 16, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum CpAsyncSourceSize { + Full, + Runtime, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum CpAsyncAdapter { + DirectPointers, + DirectPointersAndSourceSize, +} + +/// Closed contract for classic `cp.async` group controls. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct CpAsyncControl { + pub operation: CpAsyncControlOperation, + pub adapter: CpAsyncControlAdapter, + pub runtime_validation: RuntimeValidation, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum CpAsyncControlOperation { + CommitGroup, + WaitAll, + WaitGroup, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum CpAsyncControlAdapter { + NoOperands, + CompileTimeConstantMaxPending, +} + +/// Closed contract for associating classic `cp.async` completion with an mbarrier. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct CpAsyncMbarrier { + pub operation: CpAsyncMbarrierOperation, + pub state_space: CpAsyncMbarrierStateSpace, + pub adapter: CpAsyncMbarrierAdapter, + pub runtime_validation: RuntimeValidation, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum CpAsyncMbarrierOperation { + Arrive, + ArriveNoInc, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum CpAsyncMbarrierStateSpace { + Generic, + Shared, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum CpAsyncMbarrierAdapter { + PointerToVoid, +} + +/// Closed contract for the basic shared-memory mbarrier lifecycle. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct MbarrierBasic { + pub operation: MbarrierBasicOperation, + pub state_space: MbarrierStateSpace, + pub adapter: MbarrierBasicAdapter, + pub runtime_validation: RuntimeValidation, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum MbarrierBasicOperation { + Init, + Arrive, + TestWait, + Inval, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum MbarrierStateSpace { + Shared, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum MbarrierBasicAdapter { + #[serde(rename = "pointer_count_to_void")] + InitPointerCountToVoid, + #[serde(rename = "pointer_to_token")] + ArrivePointerToToken, + #[serde(rename = "pointer_token_to_predicate")] + TestWaitPointerTokenToPredicate, + #[serde(rename = "pointer_to_void")] + InvalPointerToVoid, +} + +/// Closed contract for the remaining handwritten mbarrier operations. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct MbarrierExtended { + pub operation: MbarrierExtendedOperation, + pub adapter: MbarrierExtendedAdapter, + pub source_contract: MbarrierExtendedSourceContract, + pub runtime_validation: RuntimeValidation, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum MbarrierExtendedAdapter { + PointerTxCountBytesToTokenDroppingTxCount, + RawClusterAddressToVoid, + PointerTokenToPredicate, + PointerParityToPredicate, + ZeroOperandsToVoid, + NanosecondsToVoid, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum MbarrierExtendedSourceContract { + LlvmImported, + PtxNativeRawClusterAddress, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn mbarrier_basic_contract_rejects_open_ended_policy() { + let valid = r#" +operation = "test_wait" +state_space = "shared" +adapter = "pointer_token_to_predicate" +runtime_validation = "unexecuted" +"#; + let parsed = toml::from_str::(valid).unwrap(); + assert_eq!(parsed.operation, MbarrierBasicOperation::TestWait); + assert_eq!( + parsed.adapter, + MbarrierBasicAdapter::TestWaitPointerTokenToPredicate + ); + + for invalid in [ + valid.replace("operation = \"test_wait\"", "operation = \"wait\""), + valid.replace("state_space = \"shared\"", "state_space = \"global\""), + valid.replace( + "adapter = \"pointer_token_to_predicate\"", + "adapter = \"direct\"", + ), + format!("{valid}unreviewed = true\n"), + ] { + assert!( + toml::from_str::(&invalid).is_err(), + "{invalid}" + ); + } + } + + #[test] + fn mbarrier_extended_contract_rejects_open_ended_policy() { + let valid = r#" +operation = "arrive_expect_tx_cta" +adapter = "pointer_tx_count_bytes_to_token_dropping_tx_count" +source_contract = "llvm_imported" +runtime_validation = "unexecuted" +"#; + let parsed = toml::from_str::(valid).unwrap(); + assert_eq!( + parsed.adapter, + MbarrierExtendedAdapter::PointerTxCountBytesToTokenDroppingTxCount + ); + assert_eq!( + parsed.source_contract, + MbarrierExtendedSourceContract::LlvmImported + ); + + for invalid in [ + valid.replace("llvm_imported", "auto"), + valid.replace( + "pointer_tx_count_bytes_to_token_dropping_tx_count", + "direct", + ), + format!("{valid}unreviewed = true\n"), + ] { + assert!( + toml::from_str::(&invalid).is_err(), + "{invalid}" + ); + } + } + + #[test] + fn cp_async_mbarrier_contract_rejects_open_ended_policy() { + let valid = r#" +operation = "arrive_no_inc" +state_space = "shared" +adapter = "pointer_to_void" +runtime_validation = "unexecuted" +"#; + let parsed = toml::from_str::(valid).unwrap(); + assert_eq!(parsed.operation, CpAsyncMbarrierOperation::ArriveNoInc); + assert_eq!(parsed.state_space, CpAsyncMbarrierStateSpace::Shared); + + for invalid in [ + valid.replace("operation = \"arrive_no_inc\"", "operation = \"wait\""), + valid.replace("state_space = \"shared\"", "state_space = \"global\""), + valid.replace("adapter = \"pointer_to_void\"", "adapter = \"direct\""), + format!("{valid}unreviewed = true\n"), + ] { + assert!( + toml::from_str::(&invalid).is_err(), + "{invalid}" + ); + } + } + + #[test] + fn cluster_memory_contract_rejects_open_ended_policy() { + let valid = r#" +operation = "map_shared_rank" +adapter = "generic_const_and_mut_pointer_rank_to_same_pointer" +source_contract = "llvm_mapa_shared_cluster_as7_identity_inline_ptx" +runtime_validation = "unexecuted" +"#; + let parsed = toml::from_str::(valid).unwrap(); + assert_eq!(parsed.operation, ClusterMemoryOperation::MapSharedRank); + assert_eq!( + parsed.source_contract, + ClusterMemorySourceContract::LlvmMapaSharedClusterAs7IdentityInlinePtx + ); + + for invalid in [ + valid.replace("map_shared_rank", "map_generic_rank"), + valid.replace( + "generic_const_and_mut_pointer_rank_to_same_pointer", + "direct", + ), + valid.replace( + "llvm_mapa_shared_cluster_as7_identity_inline_ptx", + "llvm_typed_as3", + ), + format!("{valid}unreviewed = true\n"), + ] { + assert!( + toml::from_str::(&invalid).is_err(), + "{invalid}" + ); + } + } +} diff --git a/crates/cuda-intrinsics-gen/src/model/contracts/mod.rs b/crates/cuda-intrinsics-gen/src/model/contracts/mod.rs new file mode 100644 index 0000000000..b92f736bda --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/model/contracts/mod.rs @@ -0,0 +1,20 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +mod control; +mod matrix; +mod memory; +mod scalar; +mod tcgen05; +mod tma; +mod warp; + +pub use self::control::*; +pub use self::matrix::*; +pub use self::memory::*; +pub use self::scalar::*; +pub use self::tcgen05::*; +pub use self::tma::*; +pub use self::warp::*; diff --git a/crates/cuda-intrinsics-gen/src/model/contracts/scalar.rs b/crates/cuda-intrinsics-gen/src/model/contracts/scalar.rs new file mode 100644 index 0000000000..257a1a85e5 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/model/contracts/scalar.rs @@ -0,0 +1,511 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use super::super::core::RuntimeValidation; +use serde::{Deserialize, Serialize}; + +/// Closed semantic contract for byte permutation. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Prmt { + pub mode: PrmtMode, + pub adapter: PrmtAdapter, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PrmtMode { + Generic, + F4e, + B4e, + Rc8, + Ecl, + Ecr, + Rc16, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PrmtAdapter { + DirectThreeOperands, + InsertZeroSecondSource, +} + +/// Closed identity and source adapter for generated packed integer dot products. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct DotProduct { + pub operation: DotProductOperation, + pub signedness: DotProductSignedness, + pub adapter: DotProductAdapter, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum DotProductOperation { + Dp2a, + Dp4a, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum DotProductSignedness { + Signed, + Unsigned, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum DotProductAdapter { + DirectThreeOperands, + InsertLowHalfFalse, +} + +/// Closed identity and carrier contract for packed floating-point ALU ops. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct PackedAlu { + pub format: PackedAluFormat, + /// Hardware floor of the native PTX instruction, independent of the + /// target floor admitted by cuda-oxide. + pub native_minimum_sm: u16, + pub operation: PackedAluOperation, + pub adapter: PackedAluAdapter, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PackedAluFormat { + Bf16x2, + F16x2, + F32x2, +} + +/// Closed identity and carrier contract for extended integer min/max ops. +/// +/// Covers the PTX ISA 8.0 integer min/max extensions: the `.relu` saturation +/// qualifier on `s32` and `s16x2`, plus the packed `s16x2`/`u16x2` forms. +/// These are the DPX-adjacent shapes ptxas fuses into `VIMNMX`-family SASS. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct IntegerMinMax { + pub format: IntegerMinMaxFormat, + /// Hardware floor of the native PTX instruction, independent of the + /// target floor admitted by cuda-oxide. + pub native_minimum_sm: u16, + pub operation: IntegerMinMaxOperation, + pub relu: bool, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum IntegerMinMaxFormat { + S32, + S16x2, + U16x2, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum IntegerMinMaxOperation { + Min, + Max, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PackedAluOperation { + Add, + AddFtz, + Sub, + SubFtz, + Mul, + MulFtz, + Fma, + FmaFtz, + FmaSat, + FmaFtzSat, + FmaRelu, + FmaFtzRelu, + Min, + Max, + Neg, + Abs, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PackedAluAdapter { + DirectPackedU32, + DirectPackedU64, +} + +/// Closed contract for scalar floating-point arithmetic. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ScalarArithmetic { + pub format: ScalarArithmeticFormat, + pub operation: ScalarArithmeticOperation, + pub rounding: ScalarArithmeticRounding, + pub subnormal: ScalarArithmeticSubnormal, + pub saturation: ScalarArithmeticSaturation, + pub runtime_validation: RuntimeValidation, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ScalarArithmeticFormat { + F32, + F64, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ScalarArithmeticOperation { + Mul, + Div, + Fma, + Add, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ScalarArithmeticRounding { + Rn, + Rz, + Rm, + Rp, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ScalarArithmeticSubnormal { + Preserve, + Ftz, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ScalarArithmeticSaturation { + None, + Sat, +} + +/// Closed contract for unary scalar floating-point math operations. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ScalarMath { + pub format: ScalarMathFormat, + pub operation: ScalarMathOperation, + pub precision: ScalarMathPrecision, + pub subnormal: ScalarMathSubnormal, + pub runtime_validation: RuntimeValidation, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ScalarMathFormat { + F16, + F32, + F64, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ScalarMathOperation { + Sin, + Cos, + /// Routes through inline PTX: LLVM 22's tblgen models ex2 as the + /// overloaded `int_nvvm_ex2_approx{,_ftz}` records (anyfloat, no DAG + /// selection pattern), so the evidence contract cannot admit a typed + /// call. The legacy `llvm.nvvm.ex2.approx.f`/`.ftz.f` names still select + /// directly on both llc 21 and 22, so this is promotable to a typed + /// route once the import resolves the overloaded family. + Ex2, + Lg2, + Rcp, + Rsqrt, + Sqrt, + /// Routes through inline PTX: LLVM 22.1.2's tblgen export carries no + /// record for tanh at all (llc 22 selects `llvm.nvvm.tanh.approx.f32` + /// via NVVMIntrinsic-class matching, and llc 21 miscompiles it into an + /// extern funcall), so this is the family's only PTX-native source. + /// Hardware floor is sm_75 (PTX ISA 7.0); the family contract gates it + /// at the attested sm_80 evidence floor like every other variant. + Tanh, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ScalarMathPrecision { + Approx, + Rn, + Rz, + Rm, + Rp, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ScalarMathSubnormal { + Preserve, + Ftz, +} + +/// Closed identity and carrier contract for extended floating-point min/max. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ExtendedMinMax { + pub format: ExtendedMinMaxFormat, + pub operation: ExtendedMinMaxOperation, + pub subnormal: ExtendedMinMaxSubnormal, + pub nan: ExtendedMinMaxNan, + pub xorsign_abs: bool, + pub adapter: ExtendedMinMaxAdapter, + pub runtime_validation: RuntimeValidation, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ExtendedMinMaxFormat { + F32, + F16, + Bf16, + F16x2, + Bf16x2, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ExtendedMinMaxOperation { + Min, + Max, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ExtendedMinMaxSubnormal { + Preserve, + Ftz, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ExtendedMinMaxNan { + Number, + Nan, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ExtendedMinMaxAdapter { + DirectF32, + /// A single 16-bit float carried as its `u16` bit pattern, matching how + /// `packed_conversion` already carries `e4m3x2` and `e5m2x2` operands. + DirectHalfU16, + DirectPackedU32, +} + +/// Closed contract for converting between scalar and packed values. +/// +/// The source may be a scalar pair (`f32x2`, two operands) or an already-packed +/// 16-bit or 32-bit register (`f16x2`, `e4m3x2`, `e5m2x2`, one operand), so the +/// operand arity follows [`PackedConversionSourceFormat`] rather than being +/// fixed at two. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct PackedConversion { + pub source_format: PackedConversionSourceFormat, + pub destination_format: PackedConversionDestinationFormat, + pub rounding: PackedConversionRounding, + pub saturation: PackedConversionSaturation, + pub adapter: PackedConversionAdapter, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PackedConversionSourceFormat { + E4m3x2, + E5m2x2, + F16x2, + F32x2, +} + +impl PackedConversionSourceFormat { + /// Number of source operands the conversion takes. + /// + /// `f32x2` names two scalar `f32` operands; every packed source arrives in a + /// single register. + pub fn operand_count(self) -> usize { + match self { + Self::F32x2 => 2, + Self::E4m3x2 | Self::E5m2x2 | Self::F16x2 => 1, + } + } + + /// PTX source-type token, used as the trailing `cvt` modifier. + pub fn ptx_token(self) -> &'static str { + match self { + Self::E4m3x2 => "e4m3x2", + Self::E5m2x2 => "e5m2x2", + Self::F16x2 => "f16x2", + Self::F32x2 => "f32", + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PackedConversionDestinationFormat { + Bf16x2, + E4m3x2, + E5m2x2, + F16x2, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PackedConversionRounding { + NearestEven, + TowardZero, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PackedConversionSaturation { + None, + Relu, + Satfinite, + SatfiniteRelu, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PackedConversionAdapter { + /// Forward the single packed source operand unchanged. + /// + /// PTX orders `cvt` operands as `d, a`, so a one-operand conversion needs no + /// reordering to keep the Rust argument order. + Identity, + /// Swap the two scalar operands. + /// + /// PTX writes the second source operand into the low half, so the operands + /// are reversed to keep the first Rust argument in the low half. + ReverseHighLowOperands, +} + +/// Closed contract for one scalar floating-point conversion. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ScalarConversion { + pub source_format: ScalarConversionSourceFormat, + pub destination_format: ScalarConversionDestinationFormat, + pub rounding: ScalarConversionRounding, + pub saturation: ScalarConversionSaturation, + pub result_representation: ScalarConversionResultRepresentation, + pub adapter: ScalarConversionAdapter, + pub runtime_validation: RuntimeValidation, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ScalarConversionSourceFormat { + F32, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ScalarConversionDestinationFormat { + Tf32, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ScalarConversionRounding { + NearestAway, + NearestEven, + TowardZero, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ScalarConversionSaturation { + None, + Relu, + Satfinite, + ReluSatfinite, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ScalarConversionResultRepresentation { + RawU32Bits, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ScalarConversionAdapter { + DirectF32ToRawU32Bits, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn packed_alu_contract_rejects_unknown_format_operation_and_adapter() { + let valid = r#" +format = "bf16x2" +native_minimum_sm = 80 +operation = "fma" +adapter = "direct_packed_u32" +"#; + toml::from_str::(valid).unwrap(); + for invalid in [ + valid.replace("format = \"bf16x2\"", "format = \"bf16\""), + valid.replace("native_minimum_sm = 80\n", ""), + valid.replace("native_minimum_sm = 80", "native_minimum_sm = \"80\""), + valid.replace("operation = \"fma\"", "operation = \"mad\""), + valid.replace( + "adapter = \"direct_packed_u32\"", + "adapter = \"bitcast_any\"", + ), + format!("{valid}unreviewed = true\n"), + ] { + assert!(toml::from_str::(&invalid).is_err(), "{invalid}"); + } + } + + #[test] + fn packed_conversion_contract_rejects_open_ended_policy() { + let valid = r#" +source_format = "f32x2" +destination_format = "bf16x2" +rounding = "nearest_even" +saturation = "none" +adapter = "reverse_high_low_operands" +"#; + toml::from_str::(valid).unwrap(); + for invalid in [ + valid.replace("source_format = \"f32x2\"", "source_format = \"f64x2\""), + valid.replace( + "destination_format = \"bf16x2\"", + "destination_format = \"f8x2\"", + ), + valid.replace("rounding = \"nearest_even\"", "rounding = \"zero\""), + valid.replace("saturation = \"none\"", "saturation = \"finite\""), + valid.replace( + "adapter = \"reverse_high_low_operands\"", + "adapter = \"direct\"", + ), + format!("{valid}unreviewed = true\n"), + ] { + assert!( + toml::from_str::(&invalid).is_err(), + "{invalid}" + ); + } + } +} diff --git a/crates/cuda-intrinsics-gen/src/model/contracts/tcgen05.rs b/crates/cuda-intrinsics-gen/src/model/contracts/tcgen05.rs new file mode 100644 index 0000000000..cbbba12ed2 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/model/contracts/tcgen05.rs @@ -0,0 +1,398 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use super::super::catalog::CatalogTargetRequirement; +use super::super::core::RuntimeValidation; +use serde::{Deserialize, Serialize}; + +/// Closed semantic contract for one tcgen05 operation. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Tcgen05 { + pub operation: Tcgen05Operation, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cp: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub ld: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub st: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub mma: Option, + pub adapter: Tcgen05Adapter, + pub source_contract: Tcgen05SourceContract, + pub runtime_validation: RuntimeValidation, +} + +/// Closed identity and selector contract for one tcgen05 MMA API. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Tcgen05Mma { + pub form: Tcgen05MmaForm, + pub selector_layout: Tcgen05MmaSelectorLayout, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub fixed_selectors: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub alias: Option, + pub llvm_target: CatalogTargetRequirement, + pub libnvvm_target: CatalogTargetRequirement, +} + +/// The 14 LLVM source forms covered by the first tcgen05 MMA batch. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Tcgen05MmaForm { + Shared, + Tensor, + TensorAshift, + SpShared, + SpTensor, + SpTensorAshift, + WsShared, + WsSharedZeroColMask, + WsSpShared, + WsSpSharedZeroColMask, + WsSpTensor, + WsSpTensorZeroColMask, + WsTensor, + WsTensorZeroColMask, +} + +/// Immediate arguments that select one imported tcgen05 MMA spelling. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(tag = "kind", rename_all = "snake_case")] +pub enum Tcgen05MmaSelectorLayout { + Base { + kind_argument: u8, + cta_group_argument: u8, + collector_a_argument: u8, + collector_a_upper_exclusive: u8, + }, + WarpSpecialized { + kind_argument: u8, + b_buffer_argument: u8, + b_usage_argument: u8, + }, +} + +/// A fixed warp-specialized selector tuple used by compatibility aliases. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Tcgen05MmaFixedSelectors { + pub kind: Tcgen05MmaKind, + pub b_buffer: u8, + pub b_usage: Tcgen05MmaBUsage, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Tcgen05MmaKind { + F16, + Tf32, + F8f6f4, + I8, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Tcgen05MmaBUsage { + Discard, + LastUse, + Fill, + Use, +} + +impl Tcgen05MmaBUsage { + pub const fn selector_value(self) -> u8 { + match self { + Self::Discard => 0, + Self::LastUse => 1, + Self::Fill => 2, + Self::Use => 3, + } + } +} + +/// Public names proposed by PR #346 for the generic f8f6f4 carrier. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Tcgen05MmaAlias { + E4m3, + E5m2, + E2m3, + E3m2, + E2m1, +} + +/// Closed identity for one tcgen05 shared-to-tensor-memory copy. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Tcgen05Cp { + pub member: Tcgen05CpMember, + pub group: Tcgen05CpGroup, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +pub enum Tcgen05CpMember { + #[serde(rename = "128x128b_b4x16_p64")] + M128x128bB4x16P64, + #[serde(rename = "128x128b_b6x16_p32")] + M128x128bB6x16P32, + #[serde(rename = "128x128b")] + M128x128b, + #[serde(rename = "128x256b_b4x16_p64")] + M128x256bB4x16P64, + #[serde(rename = "128x256b_b6x16_p32")] + M128x256bB6x16P32, + #[serde(rename = "32x128b_warpx4_b4x16_p64")] + M32x128bWarpx4B4x16P64, + #[serde(rename = "32x128b_warpx4_b6x16_p32")] + M32x128bWarpx4B6x16P32, + #[serde(rename = "32x128b_warpx4")] + M32x128bWarpx4, + #[serde(rename = "4x256b_b4x16_p64")] + M4x256bB4x16P64, + #[serde(rename = "4x256b_b6x16_p32")] + M4x256bB6x16P32, + #[serde(rename = "4x256b")] + M4x256b, + #[serde(rename = "64x128b_warpx2_01_23_b4x16_p64")] + M64x128bWarpx2Pair0123B4x16P64, + #[serde(rename = "64x128b_warpx2_01_23_b6x16_p32")] + M64x128bWarpx2Pair0123B6x16P32, + #[serde(rename = "64x128b_warpx2_01_23")] + M64x128bWarpx2Pair0123, + #[serde(rename = "64x128b_warpx2_02_13_b4x16_p64")] + M64x128bWarpx2Pair0213B4x16P64, + #[serde(rename = "64x128b_warpx2_02_13_b6x16_p32")] + M64x128bWarpx2Pair0213B6x16P32, + #[serde(rename = "64x128b_warpx2_02_13")] + M64x128bWarpx2Pair0213, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Tcgen05CpGroup { + Cg1, + Cg2, +} + +/// Closed identity for one tcgen05 tensor-memory load. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Tcgen05Ld { + pub shape: Tcgen05LdShape, + pub multiplicity: Tcgen05LdMultiplicity, + pub pack16: bool, +} + +/// Closed identity for one tcgen05 tensor-memory store. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Tcgen05St { + pub shape: Tcgen05LdShape, + pub multiplicity: Tcgen05LdMultiplicity, + pub unpack16: bool, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +pub enum Tcgen05LdShape { + #[serde(rename = "16x32bx2")] + M16x32bx2, + #[serde(rename = "16x64b")] + M16x64b, + #[serde(rename = "16x128b")] + M16x128b, + #[serde(rename = "16x256b")] + M16x256b, + #[serde(rename = "32x32b")] + M32x32b, +} + +impl Tcgen05LdShape { + pub const fn register_multiplier(self) -> usize { + match self { + Self::M16x32bx2 | Self::M16x64b | Self::M32x32b => 1, + Self::M16x128b => 2, + Self::M16x256b => 4, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Tcgen05LdMultiplicity { + X1, + X2, + X4, + X8, + X16, + X32, + X64, + X128, +} + +impl Tcgen05LdMultiplicity { + pub const fn count(self) -> usize { + match self { + Self::X1 => 1, + Self::X2 => 2, + Self::X4 => 4, + Self::X8 => 8, + Self::X16 => 16, + Self::X32 => 32, + Self::X64 => 64, + Self::X128 => 128, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Tcgen05Operation { + Alloc, + Dealloc, + RelinquishAllocPermit, + FenceBeforeThreadSync, + FenceAfterThreadSync, + Commit, + CommitSharedCluster, + MmaWsF16, + MmaF16, + MmaWsBf16, + MmaWsTf32, + CpSmemToTmem, + Ld16x256bX8Pure, + Ld16x256bPure, + LoadWait, + StoreWait, + AllocCg2, + DeallocCg2, + RelinquishAllocPermitCg2, + MmaF16Cg2, + CommitCg2, + CommitSharedClusterCg2, + CommitMulticastCg2, + CpSmemToTmemCg2, + Ld, + St, + CommitMulticast, + ShiftDown, + ShiftDownCg2, + Mma, +} + +impl Tcgen05Operation { + pub const fn execution_scope(self) -> &'static str { + match self { + Self::Alloc + | Self::Dealloc + | Self::RelinquishAllocPermit + | Self::Ld16x256bX8Pure + | Self::Ld16x256bPure + | Self::LoadWait + | Self::StoreWait + | Self::AllocCg2 + | Self::DeallocCg2 + | Self::RelinquishAllocPermitCg2 + | Self::Ld + | Self::St => "warp", + Self::FenceBeforeThreadSync + | Self::FenceAfterThreadSync + | Self::Commit + | Self::CommitSharedCluster + | Self::MmaWsF16 + | Self::MmaF16 + | Self::MmaWsBf16 + | Self::MmaWsTf32 + | Self::CpSmemToTmem + | Self::MmaF16Cg2 + | Self::CommitCg2 + | Self::CommitSharedClusterCg2 + | Self::CommitMulticastCg2 + | Self::CpSmemToTmemCg2 + | Self::CommitMulticast + | Self::ShiftDown + | Self::ShiftDownCg2 + | Self::Mma => "thread", + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Tcgen05Adapter { + SharedPointerColumnsToVoid, + TmemAddressColumnsToVoid, + NoOperands, + BarrierPointerToVoid, + MmaWsDropLegacyADescriptor, + MmaInjectZeroDisableLanes, + TmemDescriptorToVoid, + TmemToF32x32, + TmemToF32x4, + BarrierPointerMaskToVoid, + TmemInjectPack16ToU32Registers, + TmemU32RegistersInjectUnpack16ToVoid, + TmemHalfSplitOffsetInjectPack16ToU32Registers, + TmemHalfSplitOffsetU32RegistersInjectUnpack16ToVoid, + TmemAddressToVoid, + MmaDirectSelectors, + MmaWsFixedSelectorsDropLegacyADescriptor, +} + +/// Relationship between the public operation and LLVM's NVPTX selection. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Tcgen05SourceContract { + ExactTablegenSelection, + TablegenSelectionChangesPtx, + LlvmCustomLoweringWithoutSelection, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn tcgen05_execution_scope_is_closed() { + for operation in [ + Tcgen05Operation::Alloc, + Tcgen05Operation::Dealloc, + Tcgen05Operation::RelinquishAllocPermit, + Tcgen05Operation::AllocCg2, + Tcgen05Operation::DeallocCg2, + Tcgen05Operation::RelinquishAllocPermitCg2, + Tcgen05Operation::Ld16x256bX8Pure, + Tcgen05Operation::Ld16x256bPure, + Tcgen05Operation::LoadWait, + Tcgen05Operation::StoreWait, + Tcgen05Operation::Ld, + Tcgen05Operation::St, + ] { + assert_eq!(operation.execution_scope(), "warp"); + } + for operation in [ + Tcgen05Operation::FenceBeforeThreadSync, + Tcgen05Operation::FenceAfterThreadSync, + Tcgen05Operation::Commit, + Tcgen05Operation::CommitSharedCluster, + Tcgen05Operation::MmaWsF16, + Tcgen05Operation::MmaWsBf16, + Tcgen05Operation::MmaWsTf32, + Tcgen05Operation::MmaF16, + Tcgen05Operation::CpSmemToTmem, + Tcgen05Operation::MmaF16Cg2, + Tcgen05Operation::CommitCg2, + Tcgen05Operation::CommitSharedClusterCg2, + Tcgen05Operation::CommitMulticastCg2, + Tcgen05Operation::CpSmemToTmemCg2, + Tcgen05Operation::CommitMulticast, + Tcgen05Operation::ShiftDown, + Tcgen05Operation::ShiftDownCg2, + ] { + assert_eq!(operation.execution_scope(), "thread"); + } + } +} diff --git a/crates/cuda-intrinsics-gen/src/model/contracts/tma.rs b/crates/cuda-intrinsics-gen/src/model/contracts/tma.rs new file mode 100644 index 0000000000..d4690333ec --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/model/contracts/tma.rs @@ -0,0 +1,243 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use super::super::core::RuntimeValidation; +use serde::{Deserialize, Serialize}; + +/// Closed semantic contract for a TMA operation. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Tma { + pub operation: TmaOperation, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub reduction: Option, + pub adapter: TmaAdapter, + pub runtime_validation: RuntimeValidation, +} + +impl Tma { + pub const fn dimensions(&self) -> Option { + match &self.reduction { + Some(reduction) => Some(reduction.dimensions as usize), + None => self.operation.dimensions(), + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum TmaOperation { + G2sTile1d, + G2sTile2d, + G2sTile2dMulticast, + G2sTile2dMulticastCg2, + G2sTile3d, + G2sTile4d, + G2sTile5d, + S2gTile1d, + S2gTile2d, + S2gTile3d, + S2gTile4d, + S2gTile5d, + Reduce, + CommitGroup, + WaitGroup, + WaitGroupRead, + PrefetchTensorMap, + PrefetchTile1d, + PrefetchTile2d, + PrefetchTile3d, + PrefetchTile4d, + PrefetchTile5d, + #[serde(rename = "prefetch_tile_gather4_2d")] + PrefetchTileGather4TwoDimensional, + PrefetchTile1dCacheHint, + PrefetchTile2dCacheHint, + PrefetchTile3dCacheHint, + PrefetchTile4dCacheHint, + PrefetchTile5dCacheHint, + #[serde(rename = "prefetch_tile_gather4_2d_cache_hint")] + PrefetchTileGather4TwoDimensionalCacheHint, + ReplaceBoxDim, + ReplaceElementStride, + ReplaceElementType, + ReplaceFillMode, + ReplaceGlobalAddress, + ReplaceGlobalDim, + ReplaceGlobalStride, + ReplaceInterleaveLayout, + ReplaceRank, + ReplaceSwizzleAtomicity, + ReplaceSwizzleMode, + FenceProxyTensorMapAcquireCluster, + FenceProxyTensorMapAcquireCta, + FenceProxyTensorMapAcquireGpu, + FenceProxyTensorMapAcquireSystem, + FenceProxyTensorMapReleaseCluster, + FenceProxyTensorMapReleaseCta, + FenceProxyTensorMapReleaseGpu, + FenceProxyTensorMapReleaseSystem, +} + +impl TmaOperation { + pub const fn dimensions(self) -> Option { + match self { + Self::G2sTile1d | Self::S2gTile1d => Some(1), + Self::G2sTile2d + | Self::G2sTile2dMulticast + | Self::G2sTile2dMulticastCg2 + | Self::S2gTile2d => Some(2), + Self::G2sTile3d | Self::S2gTile3d => Some(3), + Self::G2sTile4d | Self::S2gTile4d => Some(4), + Self::G2sTile5d | Self::S2gTile5d => Some(5), + Self::Reduce + | Self::CommitGroup + | Self::WaitGroup + | Self::WaitGroupRead + | Self::PrefetchTensorMap + | Self::PrefetchTile1d + | Self::PrefetchTile2d + | Self::PrefetchTile3d + | Self::PrefetchTile4d + | Self::PrefetchTile5d + | Self::PrefetchTileGather4TwoDimensional + | Self::PrefetchTile1dCacheHint + | Self::PrefetchTile2dCacheHint + | Self::PrefetchTile3dCacheHint + | Self::PrefetchTile4dCacheHint + | Self::PrefetchTile5dCacheHint + | Self::PrefetchTileGather4TwoDimensionalCacheHint + | Self::ReplaceBoxDim + | Self::ReplaceElementStride + | Self::ReplaceElementType + | Self::ReplaceFillMode + | Self::ReplaceGlobalAddress + | Self::ReplaceGlobalDim + | Self::ReplaceGlobalStride + | Self::ReplaceInterleaveLayout + | Self::ReplaceRank + | Self::ReplaceSwizzleAtomicity + | Self::ReplaceSwizzleMode + | Self::FenceProxyTensorMapAcquireCluster + | Self::FenceProxyTensorMapAcquireCta + | Self::FenceProxyTensorMapAcquireGpu + | Self::FenceProxyTensorMapAcquireSystem + | Self::FenceProxyTensorMapReleaseCluster + | Self::FenceProxyTensorMapReleaseCta + | Self::FenceProxyTensorMapReleaseGpu + | Self::FenceProxyTensorMapReleaseSystem => None, + } + } + + pub const fn prefetch_coordinate_count(self) -> Option { + match self { + Self::PrefetchTile1d | Self::PrefetchTile1dCacheHint => Some(1), + Self::PrefetchTile2d | Self::PrefetchTile2dCacheHint => Some(2), + Self::PrefetchTile3d | Self::PrefetchTile3dCacheHint => Some(3), + Self::PrefetchTile4d | Self::PrefetchTile4dCacheHint => Some(4), + Self::PrefetchTile5d + | Self::PrefetchTile5dCacheHint + | Self::PrefetchTileGather4TwoDimensional + | Self::PrefetchTileGather4TwoDimensionalCacheHint => Some(5), + _ => None, + } + } + + pub const fn uses_prefetch_cache_hint(self) -> bool { + matches!( + self, + Self::PrefetchTile1dCacheHint + | Self::PrefetchTile2dCacheHint + | Self::PrefetchTile3dCacheHint + | Self::PrefetchTile4dCacheHint + | Self::PrefetchTile5dCacheHint + | Self::PrefetchTileGather4TwoDimensionalCacheHint + ) + } +} + +/// Closed identity for one TMA tensor-reduction operation. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct TmaReduction { + pub operation: TmaReductionOperation, + pub load_mode: TmaReductionLoadMode, + pub dimensions: u8, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum TmaReductionOperation { + Add, + And, + Dec, + Inc, + Max, + Min, + Or, + Xor, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum TmaReductionLoadMode { + Tile, + Im2col, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum TmaAdapter { + G2sPointersCoordinatesBarrierInjectDefaults, + G2sPointersCoordinatesBarrierMaskInjectDefaults, + S2gPointersCoordinatesInjectDefaults, + ReductionPointersCoordinatesInjectDefaults, + NoOperands, + CompileTimeConstantMaxPending, + DescriptorPointer, + DescriptorCoordinatesInjectDefaults, + DescriptorCoordinatesCacheHintInjectFlag, + DescriptorAndAddressPointers, + DescriptorOrdinalAndU32, + DescriptorOrdinalAndU64, + DescriptorAndImmediateU32, + DescriptorAndRuntimeU32, + DescriptorPointerInjectBytes, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn tma_contract_rejects_open_ended_policy() { + let valid = r#" +operation = "g2s_tile2d_multicast" +adapter = "g2s_pointers_coordinates_barrier_mask_inject_defaults" +runtime_validation = "unexecuted" +"#; + let parsed = toml::from_str::(valid).unwrap(); + assert_eq!(parsed.operation, TmaOperation::G2sTile2dMulticast); + assert_eq!( + parsed.adapter, + TmaAdapter::G2sPointersCoordinatesBarrierMaskInjectDefaults + ); + + for invalid in [ + valid.replace("g2s_tile2d_multicast", "g2s_multicast"), + valid.replace( + "g2s_pointers_coordinates_barrier_mask_inject_defaults", + "direct", + ), + valid.replace( + "runtime_validation = \"unexecuted\"", + "runtime_validation = \"assumed\"", + ), + format!("{valid}unreviewed = true\n"), + ] { + assert!(toml::from_str::(&invalid).is_err(), "{invalid}"); + } + } +} diff --git a/crates/cuda-intrinsics-gen/src/model/contracts/warp.rs b/crates/cuda-intrinsics-gen/src/model/contracts/warp.rs new file mode 100644 index 0000000000..c94c87e1d8 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/model/contracts/warp.rs @@ -0,0 +1,503 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ClusterBarrierMode { + Arrive, + ArriveAligned, + ArriveRelaxed, + ArriveRelaxedAligned, + Wait, + WaitAligned, +} + +/// Closed semantics for one cluster-barrier spelling. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ClusterBarrier { + pub mode: ClusterBarrierMode, + pub ordering: ClusterBarrierOrdering, + pub aligned: bool, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ClusterBarrierOrdering { + Release, + Relaxed, + Acquire, +} + +/// Closed semantic and lowering contract for the generated `redux.sync` +/// family. +/// +/// The Rust and NVVM dialect APIs intentionally put the participation mask +/// first, while LLVM's NVVM intrinsic puts the lane value first. Keeping that +/// adapter typed prevents a generic direct-call renderer from silently +/// swapping the collective's source and member mask. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Redux { + pub operation: ReduxOperation, + pub participation: ReduxParticipation, + pub adapter: ReduxAdapter, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ReduxOperation { + Add, + Umin, + Min, + Umax, + Max, + And, + Or, + Xor, + Fmin, + FminNan, + FminAbs, + FminAbsNan, + Fmax, + FmaxNan, + FmaxAbs, + FmaxAbsNan, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ReduxParticipation { + ExecutingLaneNamedAllNamedLanesSameInstructionAndMask, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ReduxAdapter { + MaskValueToSourceMemberMask, +} + +/// Closed semantic and lowering contract for the generated `vote.sync` +/// family. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Vote { + pub mode: VoteMode, + pub participation: VoteParticipation, + pub legacy_pre_sm70: PreSm70MemberMaskRule, + pub adapter: VoteAdapter, + pub mask_encoding: MaskEncoding, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum VoteMode { + All, + Any, + Ballot, + Uni, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum VoteParticipation { + ExecutingLaneNamedAllNamedLanesSameInstructionAndMask, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum VoteAdapter { + DirectMaskPredicate, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum MaskEncoding { + RegisterOrImmediate, +} + +/// Closed semantic and lowering contract for `activemask`. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ActiveMask { + pub observation: ActiveMaskObservation, + pub adapter: ActiveMaskAdapter, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ActiveMaskObservation { + ExecutingLanesAtInstruction, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ActiveMaskAdapter { + DirectZeroOperandMask, +} + +/// Closed semantic and lowering contract for `match.sync`. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct WarpMatch { + pub mode: WarpMatchMode, + pub value_width: WarpMatchValueWidth, + pub participation: WarpMatchParticipation, + pub adapter: WarpMatchAdapter, + pub value_encoding: MatchOperandEncoding, + pub mask_encoding: MatchOperandEncoding, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum WarpMatchMode { + Any, + All, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum WarpMatchValueWidth { + B32, + B64, +} + +impl WarpMatchValueWidth { + pub const fn bits(self) -> u32 { + match self { + Self::B32 => 32, + Self::B64 => 64, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum WarpMatchParticipation { + ExecutingLaneNamedAllNamedLanesSameInstructionAndMask, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum WarpMatchAdapter { + DirectMask, + ProjectMaskDiscardPredicate, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum MatchOperandEncoding { + RegisterOrImmediate, +} + +/// Closed semantic and lowering contract for `bar.warp.sync`. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct WarpBarrier { + pub participation: WarpBarrierParticipation, + pub legacy_pre_sm70: PreSm70MemberMaskRule, + pub adapter: WarpBarrierAdapter, + pub mask_encoding: WarpBarrierMaskEncoding, + pub memory_ordering: WarpBarrierMemoryOrdering, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum WarpBarrierParticipation { + ExecutingLaneNamedAllNamedLanesSameInstructionAndMask, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PreSm70MemberMaskRule { + AllNamedLanesConvergedAndOnlyNamedLanesActive, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum WarpBarrierAdapter { + DirectMemberMask, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum WarpBarrierMaskEncoding { + RegisterOrImmediate, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum WarpBarrierMemoryOrdering { + ParticipatingLanes, +} + +/// Closed semantic and lowering contract for `shfl.sync`. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct WarpShuffle { + pub mode: WarpShuffleMode, + pub value_kind: WarpShuffleValueKind, + pub participation: WarpShuffleParticipation, + pub legacy_pre_sm70: PreSm70MemberMaskRule, + pub source_lane: WarpShuffleSourceLane, + pub adapter: WarpShuffleAdapter, + pub clamp: u32, + pub lane_encoding: WarpShuffleOperandEncoding, + pub mask_encoding: WarpShuffleOperandEncoding, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum WarpShuffleMode { + Idx, + Bfly, + Down, + Up, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum WarpShuffleValueKind { + I32, + F32, + I64, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum WarpShuffleParticipation { + ExecutingLaneNamedAllNamedLanesSameInstructionAndMask, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum WarpShuffleSourceLane { + InRangeSourceActiveAndNamedOutOfRangeCopiesSelf, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum WarpShuffleAdapter { + MaskValueLaneOrDeltaInsertClamp, + /// Split i64 into low/high b32 halves, shuffle both in one convergent + /// side-effecting block, then reassemble the original bit layout. + MaskValueLaneOrDeltaSplitI64LowHighB32InsertClampReassemble, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum WarpShuffleOperandEncoding { + RegisterOrImmediate, + RegisterOnly, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn redux_contract_rejects_unknown_operand_adapter() { + let input = r#" +operation = "add" +participation = "executing_lane_named_all_named_lanes_same_instruction_and_mask" +adapter = "mask_value_direct" +"#; + let error = toml::from_str::(input).unwrap_err(); + assert!(error.to_string().contains("mask_value_direct")); + } + + #[test] + fn vote_contract_rejects_unknown_modes_and_mask_encodings() { + let valid = r#" +mode = "all" +participation = "executing_lane_named_all_named_lanes_same_instruction_and_mask" +legacy_pre_sm70 = "all_named_lanes_converged_and_only_named_lanes_active" +adapter = "direct_mask_predicate" +mask_encoding = "register_or_immediate" +"#; + toml::from_str::(valid).unwrap(); + + for invalid in [ + valid.replace("mode = \"all\"", "mode = \"match\""), + valid.replace( + "mask_encoding = \"register_or_immediate\"", + "mask_encoding = \"any_operand\"", + ), + valid.replace( + "legacy_pre_sm70 = \"all_named_lanes_converged_and_only_named_lanes_active\"", + "legacy_pre_sm70 = \"independent_threads\"", + ), + format!("{valid}unreviewed = true\n"), + ] { + assert!(toml::from_str::(&invalid).is_err(), "{invalid}"); + } + } + + #[test] + fn warp_shuffle_contract_rejects_open_ended_policy() { + let valid = r#" +mode = "idx" +value_kind = "i32" +participation = "executing_lane_named_all_named_lanes_same_instruction_and_mask" +legacy_pre_sm70 = "all_named_lanes_converged_and_only_named_lanes_active" +source_lane = "in_range_source_active_and_named_out_of_range_copies_self" +adapter = "mask_value_lane_or_delta_insert_clamp" +clamp = 31 +lane_encoding = "register_or_immediate" +mask_encoding = "register_or_immediate" +"#; + toml::from_str::(valid).unwrap(); + + for invalid in [ + valid.replace("mode = \"idx\"", "mode = \"rotate\""), + valid.replace("value_kind = \"i32\"", "value_kind = \"b32\""), + valid.replace( + "source_lane = \"in_range_source_active_and_named_out_of_range_copies_self\"", + "source_lane = \"unchecked\"", + ), + valid.replace( + "lane_encoding = \"register_or_immediate\"", + "lane_encoding = \"anything\"", + ), + format!("{valid}unreviewed = true\n"), + ] { + assert!( + toml::from_str::(&invalid).is_err(), + "{invalid}" + ); + } + + let i64 = r#" +mode = "down" +value_kind = "i64" +participation = "executing_lane_named_all_named_lanes_same_instruction_and_mask" +legacy_pre_sm70 = "all_named_lanes_converged_and_only_named_lanes_active" +source_lane = "in_range_source_active_and_named_out_of_range_copies_self" +adapter = "mask_value_lane_or_delta_split_i64_low_high_b32_insert_clamp_reassemble" +clamp = 31 +lane_encoding = "register_only" +mask_encoding = "register_only" +"#; + let parsed = toml::from_str::(i64).unwrap(); + assert_eq!(parsed.value_kind, WarpShuffleValueKind::I64); + assert_eq!( + parsed.adapter, + WarpShuffleAdapter::MaskValueLaneOrDeltaSplitI64LowHighB32InsertClampReassemble + ); + assert_eq!( + parsed.lane_encoding, + WarpShuffleOperandEncoding::RegisterOnly + ); + + for invalid in [ + i64.replace("value_kind = \"i64\"", "value_kind = \"u64\""), + i64.replace( + "adapter = \"mask_value_lane_or_delta_split_i64_low_high_b32_insert_clamp_reassemble\"", + "adapter = \"split_any_width\"", + ), + i64.replace( + "mask_encoding = \"register_only\"", + "mask_encoding = \"any_operand\"", + ), + ] { + assert!( + toml::from_str::(&invalid).is_err(), + "{invalid}" + ); + } + } + + #[test] + fn warp_match_contract_rejects_open_ended_adapters_and_encodings() { + let valid = r#" +mode = "all" +value_width = "b64" +participation = "executing_lane_named_all_named_lanes_same_instruction_and_mask" +adapter = "project_mask_discard_predicate" +value_encoding = "register_or_immediate" +mask_encoding = "register_or_immediate" +"#; + toml::from_str::(valid).unwrap(); + + for invalid in [ + valid.replace("mode = \"all\"", "mode = \"equal\""), + valid.replace( + "adapter = \"project_mask_discard_predicate\"", + "adapter = \"first_result\"", + ), + valid.replace( + "value_encoding = \"register_or_immediate\"", + "value_encoding = \"anything\"", + ), + format!("{valid}unreviewed = true\n"), + ] { + assert!(toml::from_str::(&invalid).is_err(), "{invalid}"); + } + } + + #[test] + fn warp_barrier_contract_rejects_open_ended_policy() { + let valid = r#" +participation = "executing_lane_named_all_named_lanes_same_instruction_and_mask" +legacy_pre_sm70 = "all_named_lanes_converged_and_only_named_lanes_active" +adapter = "direct_member_mask" +mask_encoding = "register_or_immediate" +memory_ordering = "participating_lanes" +"#; + toml::from_str::(valid).unwrap(); + + for invalid in [ + valid.replace("adapter = \"direct_member_mask\"", "adapter = \"direct\""), + valid.replace( + "legacy_pre_sm70 = \"all_named_lanes_converged_and_only_named_lanes_active\"", + "legacy_pre_sm70 = \"independent_threads\"", + ), + valid.replace( + "mask_encoding = \"register_or_immediate\"", + "mask_encoding = \"any_operand\"", + ), + valid.replace( + "memory_ordering = \"participating_lanes\"", + "memory_ordering = \"none\"", + ), + format!("{valid}unreviewed = true\n"), + ] { + assert!( + toml::from_str::(&invalid).is_err(), + "{invalid}" + ); + } + } + + #[test] + fn cluster_barrier_contract_rejects_open_ended_semantics() { + let valid = r#" +mode = "arrive_relaxed_aligned" +ordering = "relaxed" +aligned = true +"#; + let parsed = toml::from_str::(valid).unwrap(); + assert_eq!(parsed.mode, ClusterBarrierMode::ArriveRelaxedAligned); + assert_eq!(parsed.ordering, ClusterBarrierOrdering::Relaxed); + assert!(parsed.aligned); + + for invalid in [ + valid.replace("ordering = \"relaxed\"", "ordering = \"unordered\""), + valid.replace("aligned = true", "aligned = \"sometimes\""), + valid.replace( + "mode = \"arrive_relaxed_aligned\"", + "mode = \"arrive_release_aligned\"", + ), + format!("{valid}unreviewed = true\n"), + ] { + assert!( + toml::from_str::(&invalid).is_err(), + "{invalid}" + ); + } + } +} diff --git a/crates/cuda-intrinsics-gen/src/model/core.rs b/crates/cuda-intrinsics-gen/src/model/core.rs new file mode 100644 index 0000000000..f594dc884c --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/model/core.rs @@ -0,0 +1,92 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use serde::{Deserialize, Serialize}; + +/// Provenance for a generated intrinsic. PTX-native operations deliberately +/// have no invented LLVM TableGen record or LLVM intrinsic symbol. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)] +pub enum IntrinsicSource { + LlvmImported { source_record: String }, + PtxNative { instruction: String }, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum IntrinsicBackend { + LlvmNvptx, + LibNvvm, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum BackendLoweringMechanism { + TypedNvvm, + InlinePtx, +} + +/// Closed identity for the small execution-control families that share a +/// result-less MIR/dialect representation. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub enum ExecutionControlOperation { + BarrierCtaSync, + BarrierCtaSyncAligned, + BarrierCtaArrive, + BarrierCtaArriveAligned, + GridDependencyLaunchDependents, + GridDependencyWait, + SetMaxNRegInc, + SetMaxNRegDec, +} + +impl ExecutionControlOperation { + pub fn from_catalog_id(id: &str) -> Option { + Some(match id { + "barrier_cta_sync" => Self::BarrierCtaSync, + "barrier_cta_sync_aligned" => Self::BarrierCtaSyncAligned, + "barrier_cta_arrive" => Self::BarrierCtaArrive, + "barrier_cta_arrive_aligned" => Self::BarrierCtaArriveAligned, + "grid_dependency_launch_dependents" => Self::GridDependencyLaunchDependents, + "grid_dependency_wait" => Self::GridDependencyWait, + "setmaxnreg_inc" => Self::SetMaxNRegInc, + "setmaxnreg_dec" => Self::SetMaxNRegDec, + _ => return None, + }) + } + + pub const fn family(self) -> &'static str { + match self { + Self::BarrierCtaSync + | Self::BarrierCtaSyncAligned + | Self::BarrierCtaArrive + | Self::BarrierCtaArriveAligned => "counted_barrier", + Self::GridDependencyLaunchDependents | Self::GridDependencyWait => "grid_dependency", + Self::SetMaxNRegInc | Self::SetMaxNRegDec => "register_control", + } + } + + pub const fn operand_count(self) -> usize { + match self { + Self::BarrierCtaSync + | Self::BarrierCtaSyncAligned + | Self::BarrierCtaArrive + | Self::BarrierCtaArriveAligned => 2, + Self::GridDependencyLaunchDependents | Self::GridDependencyWait => 0, + Self::SetMaxNRegInc | Self::SetMaxNRegDec => 1, + } + } + + pub const fn requires_immediate_operands(self) -> bool { + matches!(self, Self::SetMaxNRegInc | Self::SetMaxNRegDec) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum RuntimeValidation { + Unexecuted, + Executed, +} diff --git a/crates/cuda-intrinsics-gen/src/model/evidence.rs b/crates/cuda-intrinsics-gen/src/model/evidence.rs new file mode 100644 index 0000000000..2800f68c28 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/model/evidence.rs @@ -0,0 +1,208 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use super::core::{BackendLoweringMechanism, IntrinsicBackend, IntrinsicSource, RuntimeValidation}; +use crate::ptx::InstructionPattern; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct EvidenceFile { + pub schema: u32, + pub backend_profile: String, + #[serde(default)] + pub backend_kind: Option, + pub llvm_revision: String, + pub backend_version: String, + pub backend_sha256: String, + #[serde(default)] + pub artifact_path: Option, + #[serde(default)] + pub build_id_prefix: Option, + #[serde(default)] + pub nvvm_ir_version: Option, + #[serde(default)] + pub debug_ir_version: Option, + pub records: Vec, +} + +/// Schema-6 evidence before matrix expansion. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct EvidenceFileV6 { + pub schema: u32, + pub backend_profile: String, + #[serde(default)] + pub backend_kind: Option, + pub llvm_revision: String, + pub backend_version: String, + pub backend_sha256: String, + #[serde(default)] + pub artifact_path: Option, + #[serde(default)] + pub build_id_prefix: Option, + #[serde(default)] + pub nvvm_ir_version: Option, + #[serde(default)] + pub debug_ir_version: Option, + #[serde(default)] + pub defaults: EvidenceRecordDefaults, + #[serde(default)] + pub fixtures: Vec, + #[serde(default)] + pub matrices: Vec, + #[serde(default)] + pub records: Vec, +} + +/// Facts shared by every record in one evidence matrix. +#[derive(Debug, Clone, Default, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct EvidenceRecordDefaults { + #[serde(default)] + pub resolved_llvm_symbol: Option, + #[serde(default)] + pub llvm_arguments: Option>, + #[serde(default)] + pub llvm_results: Option>, + #[serde(default)] + pub concrete_llvm_arguments: Option>, + #[serde(default)] + pub concrete_llvm_results: Option>, + #[serde(default)] + pub target_triple: Option, + #[serde(default)] + pub gpu_target: Option, + #[serde(default)] + pub ptx_feature: Option, + #[serde(default)] + pub status: Option, + #[serde(default)] + pub stages: Vec, + #[serde(default)] + pub declaration_attributes_canonicalized: Option, + #[serde(default)] + pub runtime_validation: Option, +} + +/// One shared fixture and the number of records it covers. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct EvidenceFixture { + pub id: String, + pub coverage_count: usize, + pub stages: Vec, +} + +/// One Cartesian evidence matrix. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct EvidenceMatrix { + pub axes: Vec, + pub product_count: usize, + #[serde(default)] + pub fixtures: Vec, + pub template: EvidenceMatrixTemplate, +} + +/// One named matrix axis. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct EvidenceMatrixAxis { + pub name: String, + pub values: Vec, +} + +/// Identity and matrix-specific facts for one record template. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct EvidenceMatrixTemplate { + pub id: String, + #[serde(default)] + pub source: Option, + #[serde(default)] + pub source_record: Option, + #[serde(deserialize_with = "deserialize_required_optional_string")] + pub llvm_symbol: Option, + pub expected_ptx: InstructionPattern, + #[serde(default)] + pub facts: EvidenceRecordDefaults, +} + +fn deserialize_required_optional_string<'de, D>(deserializer: D) -> Result, D::Error> +where + D: serde::Deserializer<'de>, +{ + Option::::deserialize(deserializer) +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct EvidenceRecord { + pub id: String, + #[serde(default)] + pub source: Option, + #[serde(default)] + pub source_record: Option, + #[serde(default)] + pub llvm_symbol: Option, + #[serde(default)] + pub resolved_llvm_symbol: Option, + #[serde(default)] + pub llvm_arguments: Vec, + #[serde(default)] + pub llvm_results: Vec, + #[serde(default)] + pub concrete_llvm_arguments: Vec, + #[serde(default)] + pub concrete_llvm_results: Vec, + pub target_triple: String, + pub gpu_target: String, + pub ptx_feature: String, + pub status: String, + #[serde(default)] + pub stages: Vec, + #[serde(default)] + pub declaration_attributes_canonicalized: Option, + #[serde(default)] + pub runtime_validation: Option, + pub expected_ptx: InstructionPattern, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct EvidenceStage { + pub targets: Vec, + pub representation: String, + pub stage: EvidenceStageKind, + #[serde(default)] + pub mechanism: Option, + pub outcome: String, + pub detail: String, + #[serde(default)] + pub artifact_kind: Option, + #[serde(default)] + pub tool_path: Option, + #[serde(default)] + pub tool_version: Option, + #[serde(default)] + pub tool_sha256: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum EvidenceArtifactKind { + Cubin, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum EvidenceStageKind { + DeclarationCanonicalization, + BackendCodegen, + DeviceLink, + PtxAssembly, + Runtime, +} diff --git a/crates/cuda-intrinsics-gen/src/model/imported.rs b/crates/cuda-intrinsics-gen/src/model/imported.rs new file mode 100644 index 0000000000..7b5f796120 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/model/imported.rs @@ -0,0 +1,195 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct UpstreamLock { + pub schema: u32, + pub llvm: LockedLlvm, + pub llvm_tblgen: LockedTool, + #[serde(default)] + pub comparison_tools: Vec, + pub dumps: LockedDumps, +} + +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct LockedLlvm { + pub repository: String, + pub revision: String, + pub provenance: String, + pub public_output_allowed: bool, +} + +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct LockedTool { + pub name: String, + pub version_line: String, + pub sha256: String, + #[serde(default)] + pub enforce_sha256: bool, + pub provenance: String, + pub built_from_llvm_revision: Option, +} + +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct LockedDumps { + pub intrinsics_sha256: String, + pub nvptx_sha256: String, + pub normalized_imported_sha256: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ImportedFile { + pub schema: u32, + pub source: ImportedSource, + pub intrinsics: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ImportedSource { + pub llvm_repository: String, + pub llvm_revision: String, + pub llvm_tblgen_version: String, + pub llvm_tblgen_source_revision: String, + pub intrinsics_json_sha256: String, + pub nvptx_json_sha256: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ImportedIntrinsic { + pub source_record: String, + pub llvm_name: String, + pub arguments: Vec, + pub results: Vec, + pub classes: Vec, + pub properties: Vec, + pub selections: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ImportedSelection { + pub source_record: String, + pub asm: String, + pub predicates: Vec, + #[serde( + default, + skip_serializing_if = "ImportedSelectionConstraints::is_empty" + )] + pub constraints: ImportedSelectionConstraints, +} + +/// Normalized constraints attached to an NVPTX instruction-selection record. +/// +/// TableGen represents address-space-specific patterns through anonymous +/// `PatFrag` records and can bind intrinsic arguments to integer literals. +/// Keeping those facts separate from the assembly spelling lets policy select +/// an exact lowering without parsing PTX text. +#[derive(Debug, Clone, Default, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ImportedSelectionConstraints { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub address_space: Option, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub immediate_bindings: Vec, +} + +/// One integer literal fixed by an NVPTX instruction-selection pattern. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ImportedImmediateBinding { + pub argument_index: usize, + pub value: i64, +} + +impl ImportedSelectionConstraints { + pub fn is_empty(&self) -> bool { + self.address_space.is_none() && self.immediate_bindings.is_empty() + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ImportedAddressSpace { + Generic, + Shared, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn locked_tool_rejects_misspelled_security_field() { + let input = r#" +name = "llvm-tblgen" +version_line = "LLVM version test" +sha256 = "abc" +enforce_sha25 = true +provenance = "test" +"#; + let error = toml::from_str::(input).unwrap_err(); + assert!(error.to_string().contains("enforce_sha25")); + } + + #[test] + fn imported_selection_rejects_misspelled_constraint() { + let input = r#"{ + "source_record": "selection", + "asm": "op;", + "predicates": [], + "constraints": { "adress_space": "shared" } + }"#; + let error = serde_json::from_str::(input).unwrap_err(); + assert!(error.to_string().contains("adress_space")); + } + + #[test] + fn imported_selection_preserves_immediate_binding() { + let input = r#"{ + "source_record": "DOT2_lo_ss", + "asm": "dp2a.lo.s32.s32 $dst, $a, $b, $c;", + "predicates": ["hasDotInstructions"], + "constraints": { + "immediate_bindings": [ + { "argument_index": 2, "value": 0 } + ] + } + }"#; + let selection = serde_json::from_str::(input).unwrap(); + assert_eq!( + selection.constraints.immediate_bindings, + [ImportedImmediateBinding { + argument_index: 2, + value: 0, + }] + ); + assert!(!selection.constraints.is_empty()); + } + + #[test] + fn imported_immediate_binding_rejects_misspelled_index() { + let input = r#"{ + "source_record": "DOT2_lo_ss", + "asm": "dp2a.lo.s32.s32 $dst, $a, $b, $c;", + "predicates": [], + "constraints": { + "immediate_bindings": [ + { "argument_indx": 2, "value": 0 } + ] + } + }"#; + let error = serde_json::from_str::(input).unwrap_err(); + assert!(error.to_string().contains("argument_indx")); + } +} diff --git a/crates/cuda-intrinsics-gen/src/model/mod.rs b/crates/cuda-intrinsics-gen/src/model/mod.rs new file mode 100644 index 0000000000..10713b4d27 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/model/mod.rs @@ -0,0 +1,20 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +mod abi_ledger; +mod catalog; +mod contracts; +mod core; +mod evidence; +mod imported; +mod overlay; + +pub use self::abi_ledger::*; +pub use self::catalog::*; +pub use self::contracts::*; +pub use self::core::*; +pub use self::evidence::*; +pub use self::imported::*; +pub use self::overlay::*; diff --git a/crates/cuda-intrinsics-gen/src/model/overlay.rs b/crates/cuda-intrinsics-gen/src/model/overlay.rs new file mode 100644 index 0000000000..60b1b6e351 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/model/overlay.rs @@ -0,0 +1,915 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use super::catalog::TargetContract; +use super::contracts::{ + ActiveMask, Clc, ClcOperation, ClusterBarrier, ClusterBarrierMode, ClusterMemory, + ClusterMemoryOperation, CpAsyncControl, CpAsyncCopy, CpAsyncMbarrier, DebugControl, + DebugControlOperation, DotProduct, ExtendedMinMax, ExtendedMinMaxFormat, ExtendedMinMaxNan, + ExtendedMinMaxOperation, ExtendedMinMaxSubnormal, IntegerMinMax, LdmatrixAdapter, + LdmatrixSafety, LdmatrixVariant, MbarrierBasic, MbarrierExtended, MbarrierExtendedOperation, + Movmatrix, PackedAlu, PackedAtomic, PackedConversion, PackedConversionDestinationFormat, + PackedConversionSaturation, Prmt, PrmtMode, Redux, RegisterMma, RegisterMmaAccumulator, + RegisterMmaElement, RegisterMmaOperation, RegisterMmaOverflow, RegisterMmaShape, + ScalarArithmetic, ScalarArithmeticFormat, ScalarArithmeticOperation, ScalarArithmeticRounding, + ScalarArithmeticSaturation, ScalarArithmeticSubnormal, ScalarConversion, + ScalarConversionRounding, ScalarConversionSaturation, ScalarMath, ScalarMathFormat, + ScalarMathOperation, ScalarMathPrecision, ScalarMathSubnormal, SparseMma, SparseMmaElement, + SparseMmaMetadata, SparseMmaOverflow, SparseMmaShape, SpecialRegister, SpecialRegisterKind, + StmatrixLayout, StmatrixMultiplicity, Tcgen05, Tcgen05CpGroup, Tcgen05CpMember, + Tcgen05LdMultiplicity, Tcgen05LdShape, Tcgen05MmaAlias, Tcgen05MmaForm, Tcgen05Operation, Tma, + TmaOperation, TmaReductionLoadMode, TmaReductionOperation, Vote, WarpBarrier, WarpMatch, + WarpShuffle, WgmmaControl, WgmmaControlMode, +}; +use super::core::{BackendLoweringMechanism, IntrinsicBackend, IntrinsicSource, RuntimeValidation}; +use super::imported::ImportedAddressSpace; +use crate::ptx::InstructionPattern; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct OverlayFile { + pub schema: u32, + pub catalog_version: String, + pub intrinsic_abi: u32, + pub backend_profile: String, + #[serde(default)] + pub shards: Vec, + #[serde(rename = "intrinsic")] + #[serde(default)] + pub intrinsics: Vec, +} + +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct OverlayShardFile { + pub schema: u32, + pub family: String, + #[serde(rename = "intrinsic")] + #[serde(default)] + pub intrinsics: Vec, + #[serde(default)] + pub register_mma_int4: Option, + #[serde(default)] + pub register_mma_int8: Option, + #[serde(default)] + pub register_mma_b1: Option, + #[serde(default)] + pub register_mma_f8f6f4_f32: Option, + #[serde(default)] + pub register_mma_f8f6f4_f16: Option, + #[serde(default)] + pub register_mma_mxf8f6f4_f32: Option, + #[serde(default)] + pub register_mma_fp8: Option, + #[serde(default)] + pub register_mma_ampere_float: Option, + #[serde(default, alias = "sparse_mma_int8")] + pub sparse_mma_integer: Option, + #[serde(default)] + pub sparse_mma_f8f6f4_f32: Option, + #[serde(default)] + pub sparse_mma_f8f6f4_f16: Option, + #[serde(default)] + pub prmt: Option, + #[serde(default)] + pub packed_conversion_fp8: Option, + #[serde(default)] + pub packed_conversion_fp8_f16x2: Option, + #[serde(default)] + pub scalar_conversion: Option, + #[serde(default)] + pub scalar_arithmetic: Option, + #[serde(default)] + pub extended_minmax: Option, + #[serde(default)] + pub cluster_sreg: Option, + #[serde(default)] + pub cluster_barrier: Option, + #[serde(default)] + pub mbarrier_extended: Option, + #[serde(default)] + pub special_registers: Option, + #[serde(default)] + pub debug_control: Option, + #[serde(default)] + pub threadfence: Option, + #[serde(default)] + pub cluster_memory: Option, + #[serde(default)] + pub stmatrix: Option, + #[serde(default)] + pub clc: Option, + #[serde(default)] + pub wgmma_controls: Option, + #[serde(default)] + pub tma: Option, + #[serde(default)] + pub tcgen05: Option, + #[serde(default)] + pub scalar_math: Option, +} + +/// Compact admission for unary scalar floating-point math operations. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ScalarMathAdmission { + pub llvm_evidence_profile: String, + pub libnvvm_evidence_profile: String, + pub runtime_validation: RuntimeValidation, + #[serde(rename = "variant")] + pub variants: Vec, +} + +/// One reviewed scalar math variant. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ScalarMathAdmissionVariant { + pub abi_id: String, + #[serde(default)] + pub libnvvm_evidence_profile: Option, + pub format: ScalarMathFormat, + pub operation: ScalarMathOperation, + pub precision: ScalarMathPrecision, + pub subnormal: ScalarMathSubnormal, +} + +/// Compact admission for the four existing `stmatrix` stores. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct StmatrixAdmission { + pub llvm_evidence_profile: String, + pub libnvvm_evidence_profile: String, + pub runtime_validation: RuntimeValidation, + #[serde(rename = "variant")] + pub variants: Vec, +} + +/// One reviewed `stmatrix` multiplicity and layout. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct StmatrixAdmissionVariant { + pub abi_id: String, + pub multiplicity: StmatrixMultiplicity, + pub layout: StmatrixLayout, +} + +/// Compact admission for the remaining handwritten mbarrier operations. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct MbarrierExtendedAdmission { + pub llvm_evidence_profile: String, + pub libnvvm_evidence_profile: String, + pub runtime_validation: RuntimeValidation, + #[serde(rename = "variant")] + pub variants: Vec, +} + +/// One reviewed extended-mbarrier operation and its reserved ABI ID. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct MbarrierExtendedAdmissionVariant { + pub abi_id: String, + pub operation: MbarrierExtendedOperation, +} + +/// Compact admission for Hopper cluster special registers. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ClusterSregAdmission { + pub axes: Vec, + pub xyz_product_count: usize, + pub record_count: usize, +} + +/// Compact admission for the six cluster-barrier instructions. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ClusterBarrierAdmission { + pub llvm_evidence_profile: String, + pub libnvvm_evidence_profile: String, + pub runtime_validation: RuntimeValidation, + #[serde(rename = "variant")] + pub variants: Vec, +} + +/// One reviewed cluster-barrier spelling. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ClusterBarrierAdmissionVariant { + pub abi_id: String, + pub mode: ClusterBarrierMode, +} + +/// Compact admission for cluster address mapping and remote shared reads. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ClusterMemoryAdmission { + pub llvm_evidence_profile: String, + pub libnvvm_evidence_profile: String, + pub runtime_validation: RuntimeValidation, + #[serde(rename = "variant")] + pub variants: Vec, +} + +/// One reviewed cluster-memory operation and its reserved ABI ID. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ClusterMemoryAdmissionVariant { + pub abi_id: String, + pub operation: ClusterMemoryOperation, +} + +/// Compact admission for the three WGMMA control instructions. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct WgmmaControlAdmission { + pub llvm_evidence_profile: String, + pub libnvvm_evidence_profile: String, + pub runtime_validation: RuntimeValidation, + #[serde(rename = "variant")] + pub variants: Vec, +} + +/// One reviewed WGMMA control operation. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct WgmmaControlAdmissionVariant { + pub abi_id: String, + pub mode: WgmmaControlMode, +} + +/// Compact admission for the reviewed non-launch special-register reads. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct SpecialRegisterAdmission { + pub llvm_evidence_profile: String, + pub libnvvm_evidence_profile: String, + pub runtime_validation: RuntimeValidation, + pub registers: Vec, + pub product_count: usize, +} + +/// Compact admission for PTX debug controls. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct DebugControlAdmission { + pub llvm_evidence_profile: String, + pub libnvvm_evidence_profile: String, + pub runtime_validation: RuntimeValidation, + pub operations: Vec, + /// Filled only when this pending shard is aggregated. + #[serde(default)] + pub abi_ids: Vec, +} + +/// Compact admission for the three CUDA thread fences. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ThreadfenceAdmission { + pub llvm_evidence_profile: String, + pub libnvvm_evidence_profile: String, + pub runtime_validation: RuntimeValidation, + #[serde(rename = "variant")] + pub variants: Vec, +} + +/// One reviewed thread-fence scope. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ThreadfenceAdmissionVariant { + pub abi_id: String, + pub scope: ThreadfenceScope, +} + +/// Scope encoded by a PTX `membar` instruction. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ThreadfenceScope { + Cta, + Device, + System, +} + +/// Compact admission for Cluster Launch Control. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ClcAdmission { + pub llvm_evidence_profile: String, + pub libnvvm_evidence_profile: String, + pub runtime_validation: RuntimeValidation, + #[serde(rename = "variant")] + pub variants: Vec, +} + +/// One reviewed Cluster Launch Control operation. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ClcAdmissionVariant { + pub abi_id: String, + pub operation: ClcOperation, +} + +/// Compact admission for the existing TMA copy and group operations. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct TmaAdmission { + pub llvm_evidence_profile: String, + pub libnvvm_evidence_profile: String, + #[serde(default)] + pub reduce_llvm_evidence_profile: Option, + #[serde(default)] + pub reduce_libnvvm_evidence_profile: Option, + pub runtime_validation: RuntimeValidation, + #[serde(rename = "variant")] + pub variants: Vec, + #[serde(rename = "reduce_variant", default)] + pub reduce_variants: Vec, +} + +/// One reviewed TMA operation and its reserved ABI ID. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct TmaAdmissionVariant { + pub abi_id: String, + pub operation: TmaOperation, +} + +/// One reviewed tensor-reduction operation and its reserved ABI ID. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct TmaReductionAdmissionVariant { + pub abi_id: String, + pub operation: TmaReductionOperation, + pub load_mode: TmaReductionLoadMode, + pub dimensions: u8, +} + +/// Compact admission for the existing Tensor Core Generation 5 operations. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Tcgen05Admission { + pub llvm_evidence_profile: String, + pub libnvvm_evidence_profile: String, + #[serde(default)] + pub cp_llvm_evidence_profile: Option, + #[serde(default)] + pub cp_libnvvm_evidence_profile: Option, + #[serde(default)] + pub ld_llvm_evidence_profile: Option, + #[serde(default)] + pub ld_libnvvm_evidence_profile: Option, + #[serde(default)] + pub st_llvm_evidence_profile: Option, + #[serde(default)] + pub st_libnvvm_evidence_profile: Option, + #[serde(default)] + pub offset_llvm_evidence_profile: Option, + #[serde(default)] + pub offset_libnvvm_evidence_profile: Option, + #[serde(default)] + pub control_llvm_evidence_profile: Option, + #[serde(default)] + pub control_libnvvm_evidence_profile: Option, + #[serde(default)] + pub mma_llvm_evidence_profile: Option, + #[serde(default)] + pub mma_libnvvm_evidence_profile: Option, + #[serde(rename = "mma_llvm_target_contract", default)] + pub mma_llvm_target_contracts: Vec, + #[serde(rename = "mma_libnvvm_target_contract", default)] + pub mma_libnvvm_target_contracts: Vec, + pub runtime_validation: RuntimeValidation, + #[serde(rename = "variant")] + pub variants: Vec, + #[serde(rename = "cp_variant", default)] + pub cp_variants: Vec, + #[serde(rename = "ld_variant", default)] + pub ld_variants: Vec, + #[serde(rename = "st_variant", default)] + pub st_variants: Vec, + #[serde(rename = "ld_offset_variant", default)] + pub ld_offset_variants: Vec, + #[serde(rename = "st_offset_variant", default)] + pub st_offset_variants: Vec, + #[serde(rename = "mma_variant", default)] + pub mma_variants: Vec, +} + +/// One reviewed tcgen05 operation and its reserved ABI ID. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Tcgen05AdmissionVariant { + pub abi_id: String, + pub operation: Tcgen05Operation, +} + +/// One reviewed tcgen05 copy member and CTA group. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Tcgen05CpAdmissionVariant { + pub abi_id: String, + pub member: Tcgen05CpMember, + pub group: Tcgen05CpGroup, +} + +/// One reviewed tcgen05 load shape, repetition, and packing mode. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Tcgen05LdAdmissionVariant { + pub abi_id: String, + pub shape: Tcgen05LdShape, + pub multiplicity: Tcgen05LdMultiplicity, + pub pack16: bool, +} + +/// One reviewed tcgen05 store shape, repetition, and unpacking mode. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Tcgen05StAdmissionVariant { + pub abi_id: String, + pub shape: Tcgen05LdShape, + pub multiplicity: Tcgen05LdMultiplicity, + pub unpack16: bool, +} + +/// One reviewed tcgen05 MMA source form or compatibility alias. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Tcgen05MmaAdmissionVariant { + pub abi_id: String, + pub form: Tcgen05MmaForm, + #[serde(default)] + pub alias: Option, +} + +/// Compact admission for the closed `prmt` family. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct PrmtAdmission { + pub llvm_evidence_profile: String, + pub libnvvm_evidence_profile: String, + pub runtime_validation: RuntimeValidation, + #[serde(rename = "variant")] + pub variants: Vec, +} + +/// One reviewed member of the `prmt` family. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct PrmtAdmissionVariant { + pub abi_id: String, + pub mode: PrmtMode, +} + +/// Compact admission for the closed scalar-f32 to packed-FP8 family. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct PackedConversionFp8Admission { + pub llvm_evidence_profile: String, + pub libnvvm_evidence_profile: String, + pub runtime_validation: RuntimeValidation, + pub destination_formats: Vec, + pub saturations: Vec, + pub product_count: usize, +} + +/// Compact admission for packed FP8 conversions whose other side is `f16x2`. +/// +/// Covers both directions: packing `f16x2` down to `e4m3x2`/`e5m2x2`, and +/// unpacking those back to `f16x2`. Both are single-operand conversions, unlike +/// the scalar-`f32` pair admitted by [`PackedConversionFp8Admission`]. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct PackedConversionFp8F16x2Admission { + pub llvm_evidence_profile: String, + pub libnvvm_evidence_profile: String, + pub runtime_validation: RuntimeValidation, + pub fp8_formats: Vec, + pub directions: Vec, + pub relu_variants: bool, + pub product_count: usize, +} + +/// The FP8 side of an `f16x2` conversion pair. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PackedConversionFp8Format { + E4m3x2, + E5m2x2, +} + +/// Which way an FP8/`f16x2` conversion runs. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PackedConversionFp8Direction { + /// `f16x2` narrowed to packed FP8, always saturating to finite. + Pack, + /// Packed FP8 widened back to `f16x2`, which is always exact. + Unpack, +} + +/// Compact admission for scalar F32-to-TF32 conversions. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ScalarConversionAdmission { + pub llvm_evidence_profile: String, + pub libnvvm_evidence_profile: String, + pub runtime_validation: RuntimeValidation, + #[serde(rename = "variant")] + pub variants: Vec, +} + +/// One reviewed scalar conversion variant. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ScalarConversionAdmissionVariant { + pub abi_id: String, + pub rounding: ScalarConversionRounding, + pub saturation: ScalarConversionSaturation, +} + +/// Compact admission for scalar floating-point arithmetic. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ScalarArithmeticAdmission { + pub llvm_evidence_profile: String, + pub libnvvm_evidence_profile: String, + pub runtime_validation: RuntimeValidation, + #[serde(rename = "variant")] + pub variants: Vec, +} + +/// One reviewed scalar arithmetic variant. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ScalarArithmeticAdmissionVariant { + pub abi_id: String, + pub format: ScalarArithmeticFormat, + pub operation: ScalarArithmeticOperation, + pub rounding: ScalarArithmeticRounding, + pub subnormal: ScalarArithmeticSubnormal, + pub saturation: ScalarArithmeticSaturation, +} + +/// Compact admission for the exact extended floating-point min/max family. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ExtendedMinMaxAdmission { + pub llvm_evidence_profile: String, + pub libnvvm_evidence_profile: String, + pub runtime_validation: RuntimeValidation, + #[serde(rename = "variant")] + pub variants: Vec, +} + +/// One reviewed extended min/max variant. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ExtendedMinMaxAdmissionVariant { + pub abi_id: String, + pub format: ExtendedMinMaxFormat, + pub operation: ExtendedMinMaxOperation, + pub subnormal: ExtendedMinMaxSubnormal, + pub nan: ExtendedMinMaxNan, + pub xorsign_abs: bool, +} + +/// Compact admission for a closed dense integer register-MMA family. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct RegisterMmaIntegerAdmission { + pub llvm_evidence_profile: String, + pub libnvvm_evidence_profile: String, + pub runtime_validation: RuntimeValidation, + #[serde(rename = "variant")] + pub variants: Vec, +} + +/// One reviewed member of a dense integer register-MMA family. +#[derive(Debug, Clone, Copy, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct RegisterMmaIntegerVariant { + pub shape: RegisterMmaShape, + pub a_element: RegisterMmaElement, + pub b_element: RegisterMmaElement, + pub overflow: RegisterMmaOverflow, +} + +/// Compact admission for the closed dense binary register-MMA family. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct RegisterMmaBinaryAdmission { + pub llvm_evidence_profile: String, + pub libnvvm_evidence_profile: String, + pub runtime_validation: RuntimeValidation, + #[serde(rename = "variant")] + pub variants: Vec, +} + +/// One reviewed member of the dense binary register-MMA family. +#[derive(Debug, Clone, Copy, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct RegisterMmaBinaryVariant { + pub shape: RegisterMmaShape, + pub operation: RegisterMmaOperation, +} + +/// Compact admission for one dense Blackwell `kind::f8f6f4` matrix. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct RegisterMmaF8F6F4Admission { + pub llvm_evidence_profile: String, + pub libnvvm_evidence_profile: String, + pub runtime_validation: RuntimeValidation, + /// Legacy shard metadata retained only so older overlays continue to parse. + /// ABI identity is bound from the append-only ledger by catalog ID. + #[serde(default, rename = "first_abi_id")] + pub _legacy_first_abi_id: Option, + pub a_elements: Vec, + pub b_elements: Vec, + pub product_count: usize, + pub targets: Vec, +} + +/// Compact admission for the standard FP8 register-MMA family. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct RegisterMmaFp8Admission { + pub llvm_evidence_profile: String, + pub libnvvm_evidence_profile: String, + pub runtime_validation: RuntimeValidation, + /// Legacy shard metadata retained only so older overlays continue to parse. + /// ABI identity is bound from the append-only ledger by catalog ID. + #[serde(default, rename = "first_abi_id")] + pub _legacy_first_abi_id: Option, + pub shapes: Vec, + pub accumulators: Vec, + pub a_elements: Vec, + pub b_elements: Vec, + pub product_count: usize, +} + +/// Compact admission for the reviewed Ampere floating-point MMA forms. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct RegisterMmaAmpereFloatAdmission { + pub llvm_evidence_profile: String, + pub libnvvm_evidence_profile: String, + pub runtime_validation: RuntimeValidation, + /// Legacy shard metadata retained only so older overlays continue to parse. + /// ABI identity is bound from the append-only ledger by catalog ID. + #[serde(default, rename = "first_abi_id")] + pub _legacy_first_abi_id: Option, + pub product_count: usize, + #[serde(rename = "variant")] + pub variants: Vec, +} + +/// One reviewed Ampere floating-point MMA form. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct RegisterMmaAmpereFloatVariant { + pub shape: RegisterMmaShape, + pub accumulator: RegisterMmaAccumulator, + pub element: RegisterMmaElement, +} + +/// Compact admission for a sparse integer register-MMA family. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct SparseMmaIntegerAdmission { + pub llvm_evidence_profile: String, + pub libnvvm_evidence_profile: String, + pub runtime_validation: RuntimeValidation, + pub metadata: SparseMmaMetadata, + #[serde(rename = "variant")] + pub variants: Vec, +} + +/// One reviewed member of a sparse integer register-MMA family. +#[derive(Debug, Clone, Copy, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct SparseMmaIntegerVariant { + pub shape: SparseMmaShape, + pub a_element: SparseMmaElement, + pub b_element: SparseMmaElement, + pub overflow: SparseMmaOverflow, +} + +/// Compact admission for ordered sparse `kind::f8f6f4` F32 MMA. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct SparseMmaF8F6F4Admission { + pub llvm_evidence_profile: String, + pub libnvvm_evidence_profile: String, + pub runtime_validation: RuntimeValidation, + pub a_elements: Vec, + pub b_elements: Vec, + pub product_count: usize, +} + +/// Compact admission for ordered sparse `kind::f8f6f4` packed-F16 MMA. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct SparseMmaF8F6F4F16Admission { + pub llvm_evidence_profile: String, + pub libnvvm_evidence_profile: String, + pub runtime_validation: RuntimeValidation, + /// Legacy shard metadata retained only so older overlays continue to parse. + /// ABI identity is bound from the append-only ledger by catalog ID. + #[serde(default, rename = "first_abi_id")] + pub _legacy_first_abi_id: Option, + pub a_elements: Vec, + pub b_elements: Vec, + pub product_count: usize, +} + +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct OverlayIntrinsic { + pub id: String, + pub abi_id: String, + pub operation_key: String, + pub family: String, + /// Imported LLVM records use the legacy `source_record` field below. + /// PTX-native records must instead carry an explicit tagged source. + #[serde(default)] + pub source: Option, + #[serde(default)] + pub source_record: Option, + pub rust_module: String, + pub rust_name: String, + #[serde(default)] + pub rust_arguments: Vec, + pub rust_result: String, + pub safe: bool, + #[serde(default)] + pub must_use: bool, + pub safe_allowlist_reason: Option, + pub public_rust_path: String, + #[serde(default)] + pub compatibility_rust_paths: Vec, + pub dialect_op_type: String, + pub dialect_op_name: String, + #[serde(default)] + pub dialect_operands: Vec, + #[serde(default)] + pub dialect_results: Vec, + #[serde(default)] + pub llvm_symbol: Option, + #[serde(default)] + pub resolved_llvm_symbol: Option, + #[serde(default)] + pub llvm_arguments: Vec, + #[serde(default)] + pub llvm_results: Vec, + pub pure: bool, + pub memory: String, + pub convergent: bool, + pub execution_scope: String, + pub minimum_ptx: String, + #[serde(default)] + pub minimum_sm: Option, + pub ptx_result: String, + pub targets: String, + pub ptx_isa_version: String, + pub ptx_isa_section: String, + pub ptx_isa_url: String, + pub lowering: String, + #[serde(default)] + pub backend_lowerings: Vec, + #[serde(default)] + pub packed_atomic: Option, + #[serde(default)] + pub redux: Option, + #[serde(default)] + pub vote: Option, + #[serde(default)] + pub active_mask: Option, + #[serde(default)] + pub warp_match: Option, + #[serde(default)] + pub warp_barrier: Option, + #[serde(default)] + pub warp_shuffle: Option, + #[serde(default)] + pub dot_product: Option, + #[serde(default)] + pub packed_alu: Option, + #[serde(default)] + pub integer_minmax: Option, + #[serde(default)] + pub packed_conversion: Option, + #[serde(default)] + pub scalar_conversion: Option, + #[serde(default)] + pub scalar_arithmetic: Option, + #[serde(default)] + pub scalar_math: Option, + #[serde(default)] + pub extended_minmax: Option, + #[serde(default)] + pub cp_async_copy: Option, + #[serde(default)] + pub cp_async_control: Option, + #[serde(default)] + pub cp_async_mbarrier: Option, + #[serde(default)] + pub mbarrier_basic: Option, + #[serde(default)] + pub movmatrix: Option, + #[serde(default)] + pub mbarrier_extended: Option, + #[serde(default)] + pub register_mma: Option, + #[serde(default)] + pub sparse_mma: Option, + #[serde(default)] + pub prmt: Option, + #[serde(default)] + pub cluster_barrier: Option, + #[serde(default)] + pub wgmma_control: Option, + #[serde(default)] + pub special_register: Option, + #[serde(default)] + pub debug_control: Option, + #[serde(default)] + pub cluster_memory: Option, + #[serde(default)] + pub clc: Option, + #[serde(default)] + pub tma: Option, + #[serde(default)] + pub tcgen05: Option, + #[serde(default)] + pub ldmatrix_variant: Option, + #[serde(default)] + pub ldmatrix_safety: Option, + #[serde(default)] + pub ldmatrix_adapter: Option, + #[serde(default)] + pub selected_address_space: Option, + pub expected_ptx: InstructionPattern, + pub summary: String, +} + +/// Backend-specific lowering selected by reviewed evidence. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct OverlayBackendLowering { + pub backend: IntrinsicBackend, + pub mechanism: BackendLoweringMechanism, + pub evidence_profile: String, + /// Optional exact target alternatives for this backend route. + #[serde(default)] + pub targets: Option, + /// Optional backend-profile floor. When absent, the intrinsic's native + /// target requirement is used. + #[serde(default)] + pub minimum_ptx: Option, + #[serde(default)] + pub minimum_sm: Option, +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::model::SparseMmaElement; + + #[test] + fn sparse_mma_admission_accepts_the_canonical_name_and_legacy_alias() { + let canonical = r#" +schema = 25 +family = "sparse_mma" + +[sparse_mma_integer] +llvm_evidence_profile = "llvm" +libnvvm_evidence_profile = "libnvvm" +runtime_validation = "unexecuted" +metadata = "ordered" + +[[sparse_mma_integer.variant]] +shape = "m16n8k64" +a_element = "s4" +b_element = "u4" +overflow = "wrapping" +"#; + let parsed = toml::from_str::(canonical).unwrap(); + assert_eq!( + parsed.sparse_mma_integer.unwrap().variants[0].b_element, + SparseMmaElement::U4 + ); + + let legacy = canonical.replace("sparse_mma_integer", "sparse_mma_int8"); + assert!( + toml::from_str::(&legacy) + .unwrap() + .sparse_mma_integer + .is_some() + ); + } +} From a79f41f32fca5288190730f60623dd27c8f997c9 Mon Sep 17 00:00:00 2001 From: nihalpasham Date: Fri, 28 Aug 2026 22:26:44 +0530 Subject: [PATCH 09/18] refactor(intrinsics-gen): split render.rs into render/ modules --- crates/cuda-intrinsics-gen/src/render.rs | 25281 ---------------- .../src/render/collector_targets.rs | 931 + .../cuda-intrinsics-gen/src/render/common.rs | 202 + .../cuda-intrinsics-gen/src/render/compat.rs | 2214 ++ .../src/render/dialect/cluster_tensor.rs | 903 + .../src/render/dialect/mma.rs | 932 + .../src/render/dialect/mod.rs | 212 + .../src/render/dialect/scalar_packed.rs | 1190 + .../src/render/dialect/warp_sync.rs | 1132 + .../src/render/families/misc.rs | 198 + .../src/render/families/mma.rs | 566 + .../src/render/families/mod.rs | 301 + .../src/render/families/packed.rs | 449 + .../src/render/families/scalar.rs | 319 + .../src/render/families/tcgen05.rs | 1398 + .../src/render/importer.rs | 2999 ++ .../src/render/lowering.rs | 2152 ++ crates/cuda-intrinsics-gen/src/render/mod.rs | 465 + .../cuda-intrinsics-gen/src/render/probes.rs | 2050 ++ .../cuda-intrinsics-gen/src/render/raw_abi.rs | 580 + .../src/render/reference.rs | 733 + .../src/render/tests/mma.rs | 1160 + .../src/render/tests/mod.rs | 241 + .../src/render/tests/scalar_misc.rs | 1094 + .../src/render/tests/tensor.rs | 1133 + .../src/render/tests/warp_memory.rs | 964 + .../src/render/validate.rs | 1369 + 27 files changed, 25887 insertions(+), 25281 deletions(-) delete mode 100644 crates/cuda-intrinsics-gen/src/render.rs create mode 100644 crates/cuda-intrinsics-gen/src/render/collector_targets.rs create mode 100644 crates/cuda-intrinsics-gen/src/render/common.rs create mode 100644 crates/cuda-intrinsics-gen/src/render/compat.rs create mode 100644 crates/cuda-intrinsics-gen/src/render/dialect/cluster_tensor.rs create mode 100644 crates/cuda-intrinsics-gen/src/render/dialect/mma.rs create mode 100644 crates/cuda-intrinsics-gen/src/render/dialect/mod.rs create mode 100644 crates/cuda-intrinsics-gen/src/render/dialect/scalar_packed.rs create mode 100644 crates/cuda-intrinsics-gen/src/render/dialect/warp_sync.rs create mode 100644 crates/cuda-intrinsics-gen/src/render/families/misc.rs create mode 100644 crates/cuda-intrinsics-gen/src/render/families/mma.rs create mode 100644 crates/cuda-intrinsics-gen/src/render/families/mod.rs create mode 100644 crates/cuda-intrinsics-gen/src/render/families/packed.rs create mode 100644 crates/cuda-intrinsics-gen/src/render/families/scalar.rs create mode 100644 crates/cuda-intrinsics-gen/src/render/families/tcgen05.rs create mode 100644 crates/cuda-intrinsics-gen/src/render/importer.rs create mode 100644 crates/cuda-intrinsics-gen/src/render/lowering.rs create mode 100644 crates/cuda-intrinsics-gen/src/render/mod.rs create mode 100644 crates/cuda-intrinsics-gen/src/render/probes.rs create mode 100644 crates/cuda-intrinsics-gen/src/render/raw_abi.rs create mode 100644 crates/cuda-intrinsics-gen/src/render/reference.rs create mode 100644 crates/cuda-intrinsics-gen/src/render/tests/mma.rs create mode 100644 crates/cuda-intrinsics-gen/src/render/tests/mod.rs create mode 100644 crates/cuda-intrinsics-gen/src/render/tests/scalar_misc.rs create mode 100644 crates/cuda-intrinsics-gen/src/render/tests/tensor.rs create mode 100644 crates/cuda-intrinsics-gen/src/render/tests/warp_memory.rs create mode 100644 crates/cuda-intrinsics-gen/src/render/validate.rs diff --git a/crates/cuda-intrinsics-gen/src/render.rs b/crates/cuda-intrinsics-gen/src/render.rs deleted file mode 100644 index 7ea55ce743..0000000000 --- a/crates/cuda-intrinsics-gen/src/render.rs +++ /dev/null @@ -1,25281 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -use crate::model::{ - ActiveMaskAdapter, BackendLoweringMechanism, CatalogFile, CatalogHardwareAlternative, - CatalogHardwareTarget, CatalogIntrinsic, CatalogLlvm, CatalogSelection, CatalogTargetContract, - ClcAdapter, ClcOperation, ClusterBarrierMode, ClusterBarrierOrdering, ClusterMemoryAdapter, - ClusterMemoryOperation, ClusterMemorySourceContract, CpAsyncCachePolicy, - CpAsyncControlOperation, CpAsyncMbarrierAdapter, CpAsyncMbarrierOperation, - CpAsyncMbarrierStateSpace, CpAsyncSourceSize, DebugControlAdapter, DebugControlOperation, - DotProductAdapter, DotProductOperation, DotProductSignedness, EvidenceArtifactKind, - EvidenceStageKind, ExecutionControlOperation, ExtendedMinMaxAdapter, ExtendedMinMaxFormat, - ExtendedMinMaxNan, ExtendedMinMaxOperation, ExtendedMinMaxSubnormal, ImportedAddressSpace, - IntegerMinMaxFormat, IntegerMinMaxOperation, IntrinsicBackend, IntrinsicSource, - LdmatrixElement, LdmatrixLayout, LdmatrixMultiplicity, LdmatrixParticipation, LdmatrixShape, - LdmatrixStateSpace, MbarrierBasicAdapter, MbarrierBasicOperation, MbarrierExtendedAdapter, - MbarrierExtendedOperation, MbarrierExtendedSourceContract, MbarrierStateSpace, - PackedAluAdapter, PackedAluFormat, PackedAluOperation, PackedAtomicFormat, - PackedConversionAdapter, PackedConversionDestinationFormat, PackedConversionRounding, - PackedConversionSaturation, PackedConversionSourceFormat, PrmtAdapter, PrmtMode, ReduxAdapter, - RegisterMmaAccumulator, RegisterMmaAdapter, RegisterMmaCompatibilitySource, RegisterMmaElement, - RegisterMmaKind, RegisterMmaLayout, RegisterMmaOperation, RegisterMmaOverflow, - RegisterMmaShape, RuntimeValidation, ScalarArithmeticFormat, ScalarArithmeticOperation, - ScalarArithmeticRounding, ScalarArithmeticSaturation, ScalarArithmeticSubnormal, - ScalarConversionRounding, ScalarConversionSaturation, ScalarMathFormat, ScalarMathOperation, - ScalarMathPrecision, ScalarMathSubnormal, SparseMma, SparseMmaAccumulator, SparseMmaAdapter, - SparseMmaCompatibilitySource, SparseMmaElement, SparseMmaLayout, SparseMmaMetadata, - SparseMmaOverflow, SparseMmaSelector, SparseMmaShape, SpecialRegisterObservation, - SpecialRegisterOutputConstraint, SpecialRegisterPtxType, StmatrixLayout, StmatrixMultiplicity, - Tcgen05Adapter, Tcgen05CpGroup, Tcgen05CpMember, Tcgen05LdMultiplicity, Tcgen05LdShape, - Tcgen05Mma, Tcgen05MmaAlias, Tcgen05MmaBUsage, Tcgen05MmaForm, Tcgen05MmaKind, - Tcgen05MmaSelectorLayout, Tcgen05Operation, Tcgen05SourceContract, TmaAdapter, TmaOperation, - TmaReductionLoadMode, TmaReductionOperation, VoteAdapter, VoteMode, WarpBarrierAdapter, - WarpMatchAdapter, WarpMatchMode, WarpShuffleAdapter, WarpShuffleMode, - WarpShuffleOperandEncoding, WarpShuffleValueKind, WgmmaControlAdapter, WgmmaControlMode, - WgmmaControlParticipation, -}; -use anyhow::{Result, ensure}; -use std::collections::{BTreeMap, BTreeSet}; -use std::fmt::Write as _; -use std::path::PathBuf; - -pub fn all_outputs( - catalog: &CatalogFile, - catalog_json: String, - catalog_sha256: &str, -) -> Result> { - validate_renderable(catalog)?; - let mut outputs = BTreeMap::new(); - outputs.insert("intrinsics/catalog.json".into(), catalog_json); - outputs.insert( - "crates/cuda-intrinsics/src/generated/mod.rs".into(), - render_raw_mod(catalog, catalog_sha256), - ); - outputs.insert( - format!( - "crates/cuda-intrinsics/src/generated/abi_v{}.rs", - catalog.intrinsic_abi - ) - .into(), - render_raw_abi(catalog, catalog_sha256)?, - ); - outputs.insert( - "crates/cuda-device/src/generated/register_mma.rs".into(), - render_compat_register_mma(catalog, catalog_sha256), - ); - outputs.insert( - "crates/cuda-device/src/generated/ldmatrix.rs".into(), - render_compat_ldmatrix(catalog, catalog_sha256), - ); - outputs.insert( - "crates/cuda-device/src/generated/sparse_mma.rs".into(), - render_compat_sparse_mma(catalog, catalog_sha256), - ); - outputs.insert( - "crates/cuda-device/src/generated/sreg.rs".into(), - render_compat_sreg(catalog, catalog_sha256), - ); - for module in ["debug", "grid", "shared", "warp"] { - outputs.insert( - format!("crates/cuda-device/src/generated/{module}_sreg.rs").into(), - render_compat_special_register_module(catalog, catalog_sha256, module), - ); - } - outputs.insert( - "crates/cuda-device/src/generated/cluster_sreg.rs".into(), - render_compat_cluster_sreg(catalog, catalog_sha256), - ); - if sync_intrinsics(catalog).any(|record| threadfence_ptx_level(record).is_some()) { - outputs.insert( - "crates/cuda-device/src/generated/fence.rs".into(), - render_compat_fence(catalog, catalog_sha256), - ); - } - outputs.insert( - "crates/cuda-device/src/generated/dotprod.rs".into(), - render_compat_dotprod(catalog, catalog_sha256), - ); - outputs.insert( - "crates/cuda-device/src/generated/prmt.rs".into(), - render_compat_prmt(catalog, catalog_sha256), - ); - outputs.insert( - "crates/cuda-device/src/generated/cluster_barrier.rs".into(), - render_compat_cluster_barrier(catalog, catalog_sha256), - ); - if cluster_memory(catalog).next().is_some() { - outputs.insert( - "crates/cuda-device/src/generated/cluster_memory.rs".into(), - render_compat_cluster_memory(catalog, catalog_sha256), - ); - } - if stmatrices(catalog).next().is_some() { - outputs.insert( - "crates/cuda-device/src/generated/stmatrix.rs".into(), - render_compat_stmatrix(catalog, catalog_sha256), - ); - } - if movmatrix(catalog).next().is_some() { - outputs.insert( - "crates/cuda-device/src/generated/movmatrix.rs".into(), - render_compat_movmatrix(catalog, catalog_sha256), - ); - } - if wgmma_controls(catalog).next().is_some() { - outputs.insert( - "crates/cuda-device/src/generated/wgmma_control.rs".into(), - render_compat_wgmma_control(catalog, catalog_sha256), - ); - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/wgmma_control.rs".into(), - render_dialect_wgmma_control(catalog, catalog_sha256), - ); - } - outputs.insert( - "crates/cuda-device/src/generated/atomic.rs".into(), - render_compat_packed_atomic(catalog, catalog_sha256), - ); - outputs.insert( - "crates/cuda-device/src/generated/async_copy.rs".into(), - render_compat_cp_async_copy(catalog, catalog_sha256), - ); - outputs.insert( - "crates/cuda-device/src/generated/mbarrier_basic.rs".into(), - render_compat_mbarrier_basic(catalog, catalog_sha256), - ); - if mbarrier_extended(catalog).next().is_some() { - outputs.insert( - "crates/cuda-device/src/generated/mbarrier_extended.rs".into(), - render_compat_mbarrier_extended(catalog, catalog_sha256), - ); - } - outputs.insert( - "crates/cuda-device/src/generated/bf16x2.rs".into(), - render_compat_packed_alu(catalog, catalog_sha256, PackedAluFormat::Bf16x2), - ); - for (module, format) in [ - ("f16", ExtendedMinMaxFormat::F16), - ("bf16", ExtendedMinMaxFormat::Bf16), - ] { - if extended_minmax(catalog).any(|record| extended_minmax_contract(record).format == format) - { - outputs.insert( - format!("crates/cuda-device/src/generated/{module}.rs").into(), - render_compat_scalar_minmax(catalog, catalog_sha256, module), - ); - } - } - outputs.insert( - "crates/cuda-device/src/generated/f16x2.rs".into(), - render_compat_packed_alu(catalog, catalog_sha256, PackedAluFormat::F16x2), - ); - for module in ["i16x2", "int"] { - if integer_minmaxes(catalog).any(|record| record.rust.module == module) { - outputs.insert( - format!("crates/cuda-device/src/generated/{module}.rs").into(), - render_compat_integer_minmax(catalog, catalog_sha256, module), - ); - } - } - if integer_minmaxes(catalog).next().is_some() { - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/integer_minmax.rs".into(), - render_dialect_integer_minmax(catalog, catalog_sha256), - ); - } - outputs.insert( - "crates/cuda-device/src/generated/f32x2.rs".into(), - render_compat_packed_alu(catalog, catalog_sha256, PackedAluFormat::F32x2), - ); - outputs.insert( - "crates/cuda-device/src/generated/convert.rs".into(), - render_compat_packed_conversion( - catalog, - catalog_sha256, - "cuda_device::convert::", - "convert", - ("lo", "hi"), - ), - ); - if let Some((path, source)) = render_compat_float_output(catalog, catalog_sha256) { - outputs.insert(path, source); - } - if extended_minmax(catalog).next().is_some() { - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/extended_minmax.rs".into(), - render_dialect_extended_minmax(catalog, catalog_sha256), - ); - } - if debug_controls(catalog).next().is_some() { - outputs.insert( - "crates/cuda-device/src/generated/debug_control.rs".into(), - render_compat_debug_control(catalog, catalog_sha256), - ); - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/debug_control.rs".into(), - render_dialect_debug_control(catalog, catalog_sha256), - ); - } - if clc_intrinsics(catalog).next().is_some() { - outputs.insert( - "crates/cuda-device/src/generated/clc.rs".into(), - render_compat_clc(catalog, catalog_sha256), - ); - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/clc.rs".into(), - render_dialect_clc(catalog, catalog_sha256), - ); - } - if tma_intrinsics(catalog).next().is_some() { - outputs.insert( - "crates/cuda-device/src/generated/tma.rs".into(), - render_compat_tma(catalog, catalog_sha256), - ); - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/tma.rs".into(), - render_dialect_tma(catalog, catalog_sha256), - ); - } - if execution_controls(catalog).next().is_some() { - outputs.insert( - "crates/cuda-device/src/generated/counted_barrier.rs".into(), - render_compat_counted_barrier(catalog, catalog_sha256), - ); - outputs.insert( - "crates/cuda-device/src/generated/grid_dependency.rs".into(), - render_compat_grid_dependency(catalog, catalog_sha256), - ); - outputs.insert( - "crates/cuda-device/src/generated/register_control.rs".into(), - render_compat_register_control(catalog, catalog_sha256), - ); - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/execution_control.rs".into(), - render_dialect_execution_control(catalog, catalog_sha256), - ); - } - if tcgen05_intrinsics(catalog).next().is_some() { - outputs.insert( - "crates/cuda-device/src/generated/tcgen05.rs".into(), - render_compat_tcgen05(catalog, catalog_sha256), - ); - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/tcgen05.rs".into(), - render_dialect_tcgen05(catalog, catalog_sha256), - ); - } - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/mod.rs".into(), - render_dialect_mod(catalog, catalog_sha256), - ); - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/sreg.rs".into(), - render_dialect_sreg(catalog, catalog_sha256), - ); - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/dotprod.rs".into(), - render_dialect_dotprod(catalog, catalog_sha256), - ); - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/ldmatrix.rs".into(), - render_dialect_ldmatrix(catalog, catalog_sha256), - ); - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/register_mma.rs".into(), - render_dialect_register_mma(catalog, catalog_sha256), - ); - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/sparse_mma.rs".into(), - render_dialect_sparse_mma(catalog, catalog_sha256), - ); - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/packed_atomic.rs".into(), - render_dialect_packed_atomic(catalog, catalog_sha256), - ); - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/packed_alu.rs".into(), - render_dialect_packed_alu(catalog, catalog_sha256), - ); - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/packed_conversion.rs".into(), - render_dialect_packed_conversion(catalog, catalog_sha256), - ); - if scalar_conversions(catalog).next().is_some() { - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/scalar_conversion.rs".into(), - render_dialect_scalar_conversion(catalog, catalog_sha256), - ); - } - if scalar_arithmetics(catalog).next().is_some() { - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/scalar_arithmetic.rs".into(), - render_dialect_scalar_arithmetic(catalog, catalog_sha256), - ); - } - if scalar_maths(catalog).next().is_some() { - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/scalar_math.rs".into(), - render_dialect_scalar_math(catalog, catalog_sha256), - ); - } - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/prmt.rs".into(), - render_dialect_prmt(catalog, catalog_sha256), - ); - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/cluster_barrier.rs".into(), - render_dialect_cluster_barrier(catalog, catalog_sha256), - ); - if cluster_memory(catalog).next().is_some() { - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/cluster_memory.rs".into(), - render_dialect_cluster_memory(catalog, catalog_sha256), - ); - } - if stmatrices(catalog).next().is_some() { - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/stmatrix.rs".into(), - render_dialect_stmatrix(catalog, catalog_sha256), - ); - } - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/redux.rs".into(), - render_dialect_redux(catalog, catalog_sha256), - ); - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/sync.rs".into(), - render_dialect_sync(catalog, catalog_sha256), - ); - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/vote.rs".into(), - render_dialect_vote(catalog, catalog_sha256), - ); - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/active_mask.rs".into(), - render_dialect_active_mask(catalog, catalog_sha256), - ); - if elect_intrinsics(catalog).next().is_some() { - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/elect.rs".into(), - render_dialect_elect(catalog, catalog_sha256), - ); - } - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/cp_async.rs".into(), - render_dialect_cp_async_copy(catalog, catalog_sha256), - ); - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/mbarrier_basic.rs".into(), - render_dialect_mbarrier_basic(catalog, catalog_sha256), - ); - if movmatrix(catalog).next().is_some() { - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/movmatrix.rs".into(), - render_dialect_movmatrix(catalog, catalog_sha256), - ); - } - if mbarrier_extended(catalog).next().is_some() { - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/mbarrier_extended.rs".into(), - render_dialect_mbarrier_extended(catalog, catalog_sha256), - ); - } - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/warp_match.rs".into(), - render_dialect_warp_match(catalog, catalog_sha256), - ); - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/warp_barrier.rs".into(), - render_dialect_warp_barrier(catalog, catalog_sha256), - ); - outputs.insert( - "crates/dialect-nvvm/src/ops/generated/warp_shuffle.rs".into(), - render_dialect_warp_shuffle(catalog, catalog_sha256), - ); - outputs.insert( - "crates/mir-importer/src/translator/terminator/intrinsics/generated.rs".into(), - render_importer(catalog, catalog_sha256), - ); - outputs.insert( - "crates/mir-lower/src/convert/generated_intrinsics.rs".into(), - render_lowering(catalog, catalog_sha256), - ); - outputs.insert( - "crates/rustc-codegen-cuda/src/generated_intrinsics.rs".into(), - render_collector(catalog, catalog_sha256), - ); - outputs.insert( - "crates/cuda-oxide-codegen/src/generated_intrinsic_targets.rs".into(), - render_targets(catalog, catalog_sha256), - ); - for record in &catalog.intrinsics { - if record.special_register.is_some() || record.family == "elect" { - for backend in [IntrinsicBackend::LlvmNvptx, IntrinsicBackend::LibNvvm] { - outputs.insert( - format!( - "intrinsics/probes/{}.{}.ll", - record.id, - backend_label(backend) - ) - .into(), - if record.family == "elect" { - render_elect_probe(catalog, record, catalog_sha256, backend) - } else { - render_special_register_probe(catalog, record, catalog_sha256, backend) - }, - ); - } - } else { - outputs.insert( - format!("intrinsics/probes/{}.ll", record.id).into(), - render_probe(catalog, record, catalog_sha256), - ); - } - } - outputs.insert( - "intrinsics/generated-reference.md".into(), - render_reference(catalog, catalog_sha256), - ); - Ok(outputs) -} - -fn validate_renderable(catalog: &CatalogFile) -> Result<()> { - ensure!( - !catalog.intrinsics.is_empty(), - "catalog contains no intrinsics" - ); - for record in &catalog.intrinsics { - match record.family.as_str() { - "sreg" => { - if let Some(special) = &record.special_register { - ensure!( - matches!(record.rust.module.as_str(), "debug" | "grid" | "thread" | "warp" | "shared") - && record.rust.arguments.is_empty() - && record.dialect.operands.is_empty() - && record.dialect.results - == [format!("i{}", special.result_width.bits())] - && record.scalar_width() == Some(special.result_width.bits()) - && record.lowering == "generated_special_register" - && record.backend_lowerings.len() == 2 - && record.semantics.pure - == (special.observation - == SpecialRegisterObservation::StablePure), - "{} is outside the closed generated special-register recipe", - record.id - ); - } else { - ensure!( - record.rust.module == "sreg" - && record.rust.arguments.is_empty() - && llvm(record).arguments.is_empty() - && record.lowering == "direct_nvvm" - && record.scalar_width().is_some(), - "{} is outside the zero-operand scalar direct-NVVM sreg recipe", - record.id - ); - } - } - "ldmatrix" => ensure!( - record.rust.module == "matrix" - && record.lowering == "generated_ldmatrix" - && record.ldmatrix.as_ref().is_some_and(|ldmatrix| { - record.rust.arguments - == [if ldmatrix.variant.element == LdmatrixElement::B16 { - "*const u32" - } else { - "*const u8" - }] - }), - "{} is outside the generated ldmatrix recipe", - record.id - ), - "stmatrix" => ensure!( - stmatrix_variant(record).is_some() - && record.rust.module == "matrix" - && !record.rust.safe - && !record.rust.must_use - && record.rust.result == "()" - && record.dialect.results.is_empty() - && record.llvm.as_ref().is_some_and(|llvm| llvm.results.is_empty()) - && record.semantics.memory == "write" - && record.semantics.convergent - && record.lowering == "generated_stmatrix", - "{} is outside the closed generated stmatrix recipe", - record.id - ), - "register_mma" => ensure!( - record.rust.module == "matrix" - && record.rust.arguments.len() == 3 + register_mma_extra_operand_count(record) - && !record.rust.safe - && record.rust.must_use - && record.semantics.memory == "none" - && record.semantics.convergent - && record.dialect.op_type == "RegisterMmaOp" - && record.dialect.op_name == "nvvm.register_mma" - && record.lowering == "generated_register_mma" - && record.register_mma.is_some(), - "{} is outside the closed generated register-MMA recipe", - record.id - ), - "movmatrix" => ensure!( - record.rust.module == "matrix" - && record.rust.arguments == ["u32"] - && record.rust.result == "u32" - && !record.rust.safe - && record.rust.must_use - && record.dialect.op_type == "MovmatrixTransB16Op" - && record.dialect.op_name == "nvvm.movmatrix_trans_b16" - && record.dialect.operands == ["i32"] - && record.dialect.results == ["i32"] - && !record.semantics.pure - && record.semantics.memory == "inaccessible_read_write" - && record.semantics.convergent - && record.lowering == "generated_movmatrix_inline_ptx" - && record.movmatrix.is_some(), - "{} is outside the generated movmatrix recipe", - record.id - ), - "sparse_mma" => { - ensure!( - record.sparse_mma.is_some(), - "{} is outside the closed generated sparse-MMA recipe", - record.id - ); - let (c_count, a_count, b_count, d_count) = - sparse_mma_fragment_counts(record); - let accumulator = match record.sparse_mma.as_ref().unwrap().accumulator { - SparseMmaAccumulator::F16 => "u32", - SparseMmaAccumulator::F32 => "f32", - SparseMmaAccumulator::S32 => "i32", - }; - let arguments = [ - format!("[{accumulator}; {c_count}]"), - format!("[u32; {a_count}]"), - format!("[u32; {b_count}]"), - "u32".to_owned(), - "u32".to_owned(), - ]; - ensure!( - record.rust.module == "matrix" - && record.rust.arguments == arguments - && record.rust.result == format!("[{accumulator}; {d_count}]") - && !record.rust.safe - && record.rust.must_use - && record.semantics.memory == "none" - && record.semantics.convergent - && record.dialect.op_type == "SparseMmaOp" - && record.dialect.op_name == "nvvm.sparse_mma" - && record.lowering == "generated_sparse_mma" - && record.sparse_mma.is_some(), - "{} is outside the closed generated sparse-MMA recipe", - record.id - ) - } - "packed_atomic" => ensure!( - record.rust.module == "atomic" - && record.rust.arguments == ["*mut u32", "u32"] - && record.rust.result == "u32" - && record.rust.must_use - && record.llvm.is_none() - && record.lowering == "generated_packed_atomic_inline_ptx" - && record.packed_atomic.is_some(), - "{} is outside the closed generated packed-atomic recipe", - record.id - ), - "redux" => ensure!( - record.rust.module == "warp" - && matches!(record.rust.arguments.as_slice(), [mask, value] - if mask == "u32" && value == &record.rust.result) - && matches!(record.rust.result.as_str(), "u32" | "i32" | "f32") - && !record.rust.safe - && matches!(record.dialect.results.as_slice(), [result] - if matches!(result.as_str(), "i32" | "f32")) - && matches!(record.dialect.results.as_slice(), [result] - if record.llvm.as_ref().is_some_and(|llvm| { - llvm.arguments == [result.as_str(), "i32"] - && llvm.results == [result.as_str()] - }) && record.dialect.operands == ["i32", result.as_str()]) - && record.lowering == "generated_redux" - && record.redux.is_some(), - "{} is outside the closed generated redux recipe", - record.id - ), - "dotprod" => ensure!( - record.rust.module == "dotprod" - && matches!(record.rust.arguments.as_slice(), [a, b, c] - if a == "u32" && b == "u32" && matches!(c.as_str(), "u32" | "i32")) - && record.rust.result == *record.rust.arguments.last().unwrap() - && record.rust.safe - && !record.rust.must_use - && record.llvm.as_ref().is_some_and(|llvm| { - llvm.results == ["i32"] - && (matches!(llvm.arguments.as_slice(), [a, b, c] - if a == "i32" && b == "i32" && c == "i32") - || matches!(llvm.arguments.as_slice(), [a, b, selector, c] - if a == "i32" && b == "i32" && selector == "i1" && c == "i32")) - }) - && record.dialect.operands == ["i32", "i32", "i32"] - && record.dialect.results == ["i32"] - && record.lowering == "generated_dotprod" - && record.dot_product.is_some(), - "{} is outside the closed generated dot-product recipe", - record.id - ), - "packed_alu" => ensure!( - record.packed_alu.as_ref().is_some_and(|packed| { - let (module, must_use, rust_type, dialect_type, adapter) = - packed_alu_format_shape(packed.format); - record.rust.module == module - && record.rust.must_use == must_use - && packed.adapter == adapter - && record - .rust - .arguments - .iter() - .all(|argument| argument == rust_type) - && record.rust.result == rust_type - && record - .dialect - .operands - .iter() - .all(|operand| operand == dialect_type) - && record.dialect.results == [dialect_type] - }) - && (1..=3).contains(&record.rust.arguments.len()) - && record.rust.safe - && record.dialect.operands.len() == record.rust.arguments.len() - && record.lowering == "generated_packed_alu_inline_ptx", - "{} is outside the closed generated packed-ALU recipe", - record.id - ), - "packed_conversion" => ensure!( - record.rust.module == "convert" - && record.rust.arguments == packed_conversion_rust_arguments(record) - && record.rust.result == packed_conversion_rust_type(record) - && record.rust.safe - && !record.rust.must_use - && record.dialect.operands == packed_conversion_dialect_operands(record) - && record.dialect.results == [packed_conversion_dialect_type(record)] - && record.lowering == packed_conversion_lowering_name(record) - && record - .packed_conversion - .as_ref() - .is_some_and(packed_conversion_is_closed_recipe), - "{} is outside the closed generated packed-conversion recipe", - record.id - ), - "scalar_conversion" => ensure!( - record.rust.module == "convert" - && record.rust.arguments == ["f32"] - && record.rust.result == "u32" - && record.rust.safe - && record.rust.must_use - && record.dialect.op_type == "ScalarConversionOp" - && record.dialect.op_name == "nvvm.scalar_conversion" - && record.dialect.operands == ["f32"] - && record.dialect.results == ["i32"] - && record.lowering == "generated_scalar_conversion" - && record.scalar_conversion.is_some(), - "{} is outside the closed generated scalar-conversion recipe", - record.id - ), - "scalar_arithmetic" => ensure!( - record.scalar_arithmetic.as_ref().is_some_and(|arithmetic| { - let ty = match arithmetic.format { - ScalarArithmeticFormat::F32 => "f32", - ScalarArithmeticFormat::F64 => "f64", - }; - let arity = match arithmetic.operation { - ScalarArithmeticOperation::Mul - | ScalarArithmeticOperation::Div - | ScalarArithmeticOperation::Add => 2, - ScalarArithmeticOperation::Fma => 3, - }; - record.rust.arguments == vec![ty; arity] - && record.rust.result == ty - && record.dialect.operands == vec![ty; arity] - && record.dialect.results == [ty] - && record.semantics.pure - == (arithmetic.operation != ScalarArithmeticOperation::Div) - }) - && record.rust.module == "float" - && record.rust.safe - && record.rust.must_use - && record.dialect.op_type == "ScalarArithmeticOp" - && record.dialect.op_name == "nvvm.scalar_arithmetic" - && record.semantics.memory == "none" - && !record.semantics.convergent - && record.lowering == "generated_scalar_arithmetic", - "{} is outside the closed generated scalar-arithmetic recipe", - record.id - ), - "scalar_math" => ensure!( - record.scalar_math.as_ref().is_some_and(|math| { - let (rust_ty, dialect_ty) = match math.format { - ScalarMathFormat::F16 => ("u16", "i16"), - ScalarMathFormat::F32 => ("f32", "f32"), - ScalarMathFormat::F64 => ("f64", "f64"), - }; - record.rust.arguments == [rust_ty] - && record.rust.result == rust_ty - && record.dialect.operands == [dialect_ty] - && record.dialect.results == [dialect_ty] - && record.semantics.pure - }) - && record.rust.module == "float" - && record.rust.safe - && record.rust.must_use - && record.dialect.op_type == "ScalarMathOp" - && record.dialect.op_name == "nvvm.scalar_math" - && record.semantics.memory == "none" - && !record.semantics.convergent - && record.lowering == "generated_scalar_math", - "{} is outside the closed generated scalar-math recipe", - record.id - ), - "extended_minmax" => ensure!( - record.extended_minmax.as_ref().is_some_and(|minmax| { - let (module, rust_type, dialect_type, adapter) = match minmax.format { - ExtendedMinMaxFormat::F32 => ( - "float", - "f32", - "f32", - ExtendedMinMaxAdapter::DirectF32, - ), - ExtendedMinMaxFormat::F16 => ( - "f16", - "u16", - "i16", - ExtendedMinMaxAdapter::DirectHalfU16, - ), - ExtendedMinMaxFormat::Bf16 => ( - "bf16", - "u16", - "i16", - ExtendedMinMaxAdapter::DirectHalfU16, - ), - ExtendedMinMaxFormat::F16x2 => ( - "f16x2", - "u32", - "i32", - ExtendedMinMaxAdapter::DirectPackedU32, - ), - ExtendedMinMaxFormat::Bf16x2 => ( - "bf16x2", - "u32", - "i32", - ExtendedMinMaxAdapter::DirectPackedU32, - ), - }; - record.rust.module == module - && record.rust.arguments == [rust_type, rust_type] - && record.rust.result == rust_type - && record.dialect.operands == [dialect_type, dialect_type] - && record.dialect.results == [dialect_type] - && minmax.adapter == adapter - }) - && record.rust.safe - && record.rust.must_use - && record.dialect.op_type == "ExtendedMinMaxOp" - && record.dialect.op_name == "nvvm.extended_minmax" - && record.semantics.pure - && record.semantics.memory == "none" - && !record.semantics.convergent - && record.lowering == "generated_extended_minmax", - "{} is outside the closed generated extended-minmax recipe", - record.id - ), - "prmt" => ensure!( - record.rust.module == "prmt" - && matches!(record.rust.arguments.len(), 2 | 3) - && record.rust.arguments.iter().all(|argument| argument == "u32") - && record.rust.result == "u32" - && record.rust.safe - && record.rust.must_use - && record.dialect.op_type == "PrmtOp" - && record.dialect.op_name == "nvvm.prmt" - && record.dialect.operands.len() == record.rust.arguments.len() - && record.dialect.operands.iter().all(|operand| operand == "i32") - && record.dialect.results == ["i32"] - && record.lowering == "generated_prmt" - && record.prmt.is_some(), - "{} is outside the closed generated prmt recipe", - record.id - ), - "cluster_barrier" => ensure!( - matches!( - record.id.as_str(), - "barrier_cluster_arrive" - | "barrier_cluster_arrive_aligned" - | "barrier_cluster_arrive_relaxed" - | "barrier_cluster_arrive_relaxed_aligned" - | "barrier_cluster_wait" - | "barrier_cluster_wait_aligned" - ) - && record.rust.module == "cluster" - && record.rust.arguments.is_empty() - && record.rust.result == "()" - && !record.rust.safe - && !record.rust.must_use - && record.dialect.op_type == "ClusterBarrierOp" - && record.dialect.op_name == "nvvm.cluster_barrier" - && record.dialect.operands.is_empty() - && record.dialect.results.is_empty() - && record.llvm.as_ref().is_some_and(|llvm| { - llvm.arguments.is_empty() && llvm.results.is_empty() - }) - && record.cluster_barrier.as_ref().is_some_and(|barrier| { - matches!( - (barrier.mode, barrier.ordering, barrier.aligned), - ( - ClusterBarrierMode::Arrive, - ClusterBarrierOrdering::Release, - false - ) | ( - ClusterBarrierMode::ArriveAligned, - ClusterBarrierOrdering::Release, - true - ) | ( - ClusterBarrierMode::ArriveRelaxed, - ClusterBarrierOrdering::Relaxed, - false - ) | ( - ClusterBarrierMode::ArriveRelaxedAligned, - ClusterBarrierOrdering::Relaxed, - true - ) | ( - ClusterBarrierMode::Wait, - ClusterBarrierOrdering::Acquire, - false - ) | ( - ClusterBarrierMode::WaitAligned, - ClusterBarrierOrdering::Acquire, - true - ) - ) - }) - && record.lowering == "generated_cluster_barrier", - "{} is outside the closed generated cluster-barrier recipe", - record.id - ), - "cluster_memory" => ensure!( - record.rust.module == "cluster" - && !record.rust.safe - && record.rust.must_use - && record.semantics.convergent - && record.semantics.execution_scope == "cluster" - && record.target.minimum_ptx.to_string() == "7.8" - && matches!( - &record.target.hardware, - CatalogHardwareTarget::AnyOf { alternatives } - if alternatives.as_slice() - == [CatalogHardwareAlternative::MinimumSm { sm: 90 }] - ) - && record.lowering == "generated_cluster_memory_inline_ptx" - && record.backend_lowerings.len() == 2 - && record.backend_lowerings.iter().all(|lowering| { - lowering.mechanism == BackendLoweringMechanism::InlinePtx - && lowering.target.minimum_ptx.to_string() == "7.8" - && matches!( - &lowering.target.hardware, - CatalogHardwareTarget::AnyOf { alternatives } - if alternatives.as_slice() - == [CatalogHardwareAlternative::MinimumSm { sm: 90 }] - ) - }) - && record.cluster_memory.as_ref().is_some_and(|cluster| { - cluster.runtime_validation == RuntimeValidation::Unexecuted - && match (cluster.operation, cluster.adapter, cluster.source_contract) { - ( - ClusterMemoryOperation::MapSharedRank, - ClusterMemoryAdapter::GenericConstAndMutPointerRankToSamePointer, - ClusterMemorySourceContract::LlvmMapaSharedClusterAs7IdentityInlinePtx, - ) => { - record.id == "map_shared_rank" - && record.rust.arguments == ["*const u8", "u32"] - && record.rust.result == "*const u8" - && record.rust.compatibility_paths - == [ - "cuda_device::cluster::map_shared_rank", - "cuda_device::cluster::map_shared_rank_mut", - ] - && record.dialect.op_type == "MapaSharedClusterOp" - && record.dialect.op_name == "nvvm.mapa_shared_cluster" - && record.dialect.operands == ["ptr", "i32"] - && record.dialect.results == ["ptr"] - && record.semantics.memory == "none" - && matches!( - record.source, - IntrinsicSource::LlvmImported { ref source_record } - if source_record == "int_nvvm_mapa_shared_cluster" - ) - && record.llvm.as_ref().is_some_and(|llvm| { - llvm.symbol == "llvm.nvvm.mapa.shared.cluster" - && llvm.arguments == ["shared_ptr", "i32"] - && llvm.results == ["shared_cluster_ptr"] - && llvm.properties - == [ - "IntrNoMem", - "IntrSpeculatable", - "NoCapture", - ] - }) - && record - .selections - .iter() - .map(|selection| selection.source_record.as_str()) - .collect::>() - == BTreeSet::from([ - "mapa_shared_cluster_64", - "mapa_shared_cluster_64i", - ]) - } - ( - ClusterMemoryOperation::ReadU32, - ClusterMemoryAdapter::ConstU32PointerRankToU32, - ClusterMemorySourceContract::PtxNativeMapaThenWeakClusterLoad, - ) => { - record.id == "dsmem_read_u32" - && record.rust.arguments == ["*const u32", "u32"] - && record.rust.result == "u32" - && record.rust.compatibility_paths - == ["cuda_device::cluster::dsmem_read_u32"] - && record.dialect.op_type == "DsmemReadU32Op" - && record.dialect.op_name == "nvvm.dsmem_read_u32" - && record.dialect.operands == ["ptr", "i32"] - && record.dialect.results == ["i32"] - && record.semantics.memory == "read" - && record.llvm.is_none() - && record.selections.is_empty() - && matches!( - record.source, - IntrinsicSource::PtxNative { ref instruction } - if instruction - == "mapa.shared::cluster.u64 + ld.shared::cluster.u32" - ) - } - _ => false, - } - }), - "{} is outside the closed generated cluster-memory recipe", - record.id - ), - "debug_control" => ensure!( - record.rust.module == "debug" - && record.rust.safe - && !record.rust.must_use - && record.dialect.operands.is_empty() - && record.dialect.results.is_empty() - && record.llvm.is_none() - && record.lowering == "generated_debug_control" - && record.debug_control.as_ref().is_some_and(|debug| { - debug.runtime_validation == RuntimeValidation::Unexecuted - && match debug.operation { - DebugControlOperation::Trap => { - debug.adapter == DebugControlAdapter::Direct - && record.rust.arguments.is_empty() - && record.rust.result == "!" - && record.dialect.op_type == "TrapOp" - && record.dialect.op_name == "nvvm.trap" - } - DebugControlOperation::Breakpoint => { - debug.adapter == DebugControlAdapter::Direct - && record.rust.arguments.is_empty() - && record.rust.result == "()" - && record.dialect.op_type == "BreakpointOp" - && record.dialect.op_name == "nvvm.brkpt" - } - DebugControlOperation::Pmevent => { - debug.adapter - == DebugControlAdapter::ConstGenericToImmediateU32 - && record.rust.arguments == ["u32"] - && record.rust.result == "()" - && record.dialect.op_type == "PmEventOp" - && record.dialect.op_name == "nvvm.pmevent" - } - } - }), - "{} is outside the closed generated debug-control recipe", - record.id - ), - "clc" => ensure!( - record.rust.module == "clc" - && !record.rust.safe - && !record.rust.must_use - && record.lowering == "generated_clc" - && record.llvm.is_some() - && record.clc.as_ref().is_some_and(|clc| { - clc.runtime_validation == RuntimeValidation::Unexecuted - && match (clc.operation, clc.adapter) { - ( - ClcOperation::TryCancel - | ClcOperation::TryCancelMulticast, - ClcAdapter::GenericPointersToShared, - ) => { - record.rust.arguments == ["*mut u8", "*mut u64"] - && record.rust.result == "()" - && record.dialect.operands == ["ptr", "ptr"] - && record.dialect.results.is_empty() - && record.llvm.as_ref().is_some_and(|llvm| { - llvm.arguments == ["shared_ptr", "shared_ptr"] - && llvm.results.is_empty() - }) - } - ( - ClcOperation::QueryIsCanceled, - ClcAdapter::PairU64ToI128BoolToU32, - ) => { - record.rust.arguments == ["u64", "u64"] - && record.rust.result == "u32" - && record.dialect.operands == ["i64", "i64"] - && record.dialect.results == ["i32"] - && record.llvm.as_ref().is_some_and(|llvm| { - llvm.arguments == ["i128"] - && llvm.results == ["i1"] - }) - } - ( - ClcOperation::QueryGetFirstCtaidX - | ClcOperation::QueryGetFirstCtaidY - | ClcOperation::QueryGetFirstCtaidZ, - ClcAdapter::PairU64ToI128U32, - ) => { - record.rust.arguments == ["u64", "u64"] - && record.rust.result == "u32" - && record.dialect.operands == ["i64", "i64"] - && record.dialect.results == ["i32"] - && record.llvm.as_ref().is_some_and(|llvm| { - llvm.arguments == ["i128"] - && llvm.results == ["i32"] - }) - } - _ => false, - } - }), - "{} is outside the closed generated CLC recipe", - record.id - ), - "wgmma_control" => ensure!( - record.rust.module == "wgmma" - && record.rust.result == "()" - && !record.rust.safe - && !record.rust.must_use - && record.dialect.results.is_empty() - && record.llvm.as_ref().is_some_and(|llvm| llvm.results.is_empty()) - && !record.semantics.pure - && record.semantics.memory == "read_write" - && record.semantics.convergent - && record.semantics.execution_scope == "warpgroup" - && record.target.minimum_ptx.encoded() == 80 - && record.target.hardware - == CatalogHardwareTarget::AnyOf { - alternatives: vec![CatalogHardwareAlternative::ExactArchitecture { - sm: 90, - }], - } - && record.backend_lowerings.len() == 2 - && record.backend_lowerings.iter().any(|lowering| { - lowering.backend == IntrinsicBackend::LlvmNvptx - && lowering.mechanism == BackendLoweringMechanism::TypedNvvm - }) - && record.backend_lowerings.iter().any(|lowering| { - lowering.backend == IntrinsicBackend::LibNvvm - && lowering.mechanism == BackendLoweringMechanism::InlinePtx - }) - && record.lowering == "generated_wgmma_control" - && record.wgmma_control.as_ref().is_some_and(|control| { - control.participation - == WgmmaControlParticipation::WarpgroupAllThreadsSameInstruction - && match control.mode { - WgmmaControlMode::Fence => { - control.adapter == WgmmaControlAdapter::NoArguments - && record.id == "wgmma_fence" - && record.rust.arguments.is_empty() - && record.dialect.op_type == "WgmmaFenceSyncAlignedOp" - && record.dialect.op_name - == "nvvm.wgmma_fence_sync_aligned" - && record.dialect.operands.is_empty() - && record.llvm.as_ref().is_some_and(|llvm| { - llvm.arguments.is_empty() - }) - } - WgmmaControlMode::CommitGroup => { - control.adapter == WgmmaControlAdapter::NoArguments - && record.id == "wgmma_commit_group" - && record.rust.arguments.is_empty() - && record.dialect.op_type - == "WgmmaCommitGroupSyncAlignedOp" - && record.dialect.op_name - == "nvvm.wgmma_commit_group_sync_aligned" - && record.dialect.operands.is_empty() - && record.llvm.as_ref().is_some_and(|llvm| { - llvm.arguments.is_empty() - }) - } - WgmmaControlMode::WaitGroup => { - control.adapter - == WgmmaControlAdapter::ConstGenericU32ToI64Immediate - && record.id == "wgmma_wait_group" - && record.rust.arguments == ["u64"] - && record.rust.compatibility_paths - == ["cuda_device::wgmma::__wgmma_wait_group"] - && record.dialect.op_type - == "WgmmaWaitGroupSyncAlignedOp" - && record.dialect.op_name - == "nvvm.wgmma_wait_group_sync_aligned" - && record.dialect.operands == ["i64"] - && record.llvm.as_ref().is_some_and(|llvm| { - llvm.arguments == ["i64"] - }) - } - } - }), - "{} is outside the closed generated WGMMA-control recipe", - record.id - ), - "tma" => ensure!( - record.rust.module == "tma" - && record.rust.result == "()" - && !record.rust.must_use - && record.dialect.results.is_empty() - && record.llvm.as_ref().is_some_and(|llvm| llvm.results.is_empty()) - && record.lowering == "generated_tma" - && record.tma.as_ref().is_some_and(|tma| { - tma.runtime_validation == RuntimeValidation::Unexecuted - && match (tma.operation, tma.adapter) { - ( - TmaOperation::G2sTile1d - | TmaOperation::G2sTile2d - | TmaOperation::G2sTile3d - | TmaOperation::G2sTile4d - | TmaOperation::G2sTile5d, - TmaAdapter::G2sPointersCoordinatesBarrierInjectDefaults, - ) => { - !record.rust.safe - && record.semantics.convergent - && record.semantics.memory == "read_write" - } - ( - TmaOperation::G2sTile2dMulticast - | TmaOperation::G2sTile2dMulticastCg2, - TmaAdapter::G2sPointersCoordinatesBarrierMaskInjectDefaults, - ) => { - !record.rust.safe - && record.semantics.convergent - && record.semantics.memory == "read_write" - } - ( - TmaOperation::S2gTile1d - | TmaOperation::S2gTile2d - | TmaOperation::S2gTile3d - | TmaOperation::S2gTile4d - | TmaOperation::S2gTile5d, - TmaAdapter::S2gPointersCoordinatesInjectDefaults, - ) => { - !record.rust.safe - && record.semantics.convergent - && record.semantics.memory == "read_write" - } - ( - TmaOperation::Reduce, - TmaAdapter::ReductionPointersCoordinatesInjectDefaults, - ) => { - !record.rust.safe - && record.semantics.convergent - && record.semantics.memory == "read_write" - && tma.reduction.is_some() - } - (TmaOperation::CommitGroup, TmaAdapter::NoOperands) => { - record.rust.safe - && !record.semantics.convergent - && record.semantics.memory == "read_write" - } - ( - TmaOperation::WaitGroup | TmaOperation::WaitGroupRead, - TmaAdapter::CompileTimeConstantMaxPending, - ) => { - record.rust.safe - && !record.semantics.convergent - && record.semantics.memory == "read_write" - } - ( - TmaOperation::PrefetchTensorMap, - TmaAdapter::DescriptorPointer, - ) => { - !record.rust.safe - && !record.semantics.convergent - && record.semantics.memory == "read" - } - ( - TmaOperation::PrefetchTile1d - | TmaOperation::PrefetchTile2d - | TmaOperation::PrefetchTile3d - | TmaOperation::PrefetchTile4d - | TmaOperation::PrefetchTile5d - | TmaOperation::PrefetchTileGather4TwoDimensional, - TmaAdapter::DescriptorCoordinatesInjectDefaults, - ) => { - !record.rust.safe - && record.semantics.convergent - && record.semantics.memory == "read" - } - ( - TmaOperation::PrefetchTile1dCacheHint - | TmaOperation::PrefetchTile2dCacheHint - | TmaOperation::PrefetchTile3dCacheHint - | TmaOperation::PrefetchTile4dCacheHint - | TmaOperation::PrefetchTile5dCacheHint - | TmaOperation::PrefetchTileGather4TwoDimensionalCacheHint, - TmaAdapter::DescriptorCoordinatesCacheHintInjectFlag, - ) => { - !record.rust.safe - && record.semantics.convergent - && record.semantics.memory == "read" - } - ( - TmaOperation::ReplaceGlobalAddress, - TmaAdapter::DescriptorAndAddressPointers, - ) - | ( - TmaOperation::ReplaceBoxDim - | TmaOperation::ReplaceElementStride - | TmaOperation::ReplaceGlobalDim, - TmaAdapter::DescriptorOrdinalAndU32, - ) - | ( - TmaOperation::ReplaceGlobalStride, - TmaAdapter::DescriptorOrdinalAndU64, - ) - | ( - TmaOperation::ReplaceElementType - | TmaOperation::ReplaceFillMode - | TmaOperation::ReplaceInterleaveLayout - | TmaOperation::ReplaceSwizzleAtomicity - | TmaOperation::ReplaceSwizzleMode, - TmaAdapter::DescriptorAndImmediateU32, - ) - | ( - TmaOperation::ReplaceRank, - TmaAdapter::DescriptorAndRuntimeU32, - ) => { - !record.rust.safe - && !record.semantics.convergent - && record.semantics.memory == "write" - } - ( - TmaOperation::FenceProxyTensorMapAcquireCluster - | TmaOperation::FenceProxyTensorMapAcquireCta - | TmaOperation::FenceProxyTensorMapAcquireGpu - | TmaOperation::FenceProxyTensorMapAcquireSystem, - TmaAdapter::DescriptorPointerInjectBytes, - ) => { - !record.rust.safe - && !record.semantics.convergent - && record.semantics.memory == "read_write" - } - ( - TmaOperation::FenceProxyTensorMapReleaseCluster - | TmaOperation::FenceProxyTensorMapReleaseCta - | TmaOperation::FenceProxyTensorMapReleaseGpu - | TmaOperation::FenceProxyTensorMapReleaseSystem, - TmaAdapter::NoOperands, - ) => { - record.rust.safe - && !record.semantics.convergent - && record.semantics.memory == "read_write" - } - _ => false, - } - }), - "{} is outside the closed generated TMA recipe", - record.id - ), - "counted_barrier" | "grid_dependency" | "register_control" => ensure!( - ExecutionControlOperation::from_catalog_id(&record.id).is_some_and(|operation| { - record.family == operation.family() - && record.rust.result == "()" - && !record.rust.safe - && !record.rust.must_use - && record.dialect.results.is_empty() - && record - .llvm - .as_ref() - .is_some_and(|llvm| llvm.results.is_empty()) - && record.lowering == "generated_execution_control" - && match operation { - ExecutionControlOperation::BarrierCtaSync - | ExecutionControlOperation::BarrierCtaSyncAligned - | ExecutionControlOperation::BarrierCtaArrive - | ExecutionControlOperation::BarrierCtaArriveAligned => { - record.rust.module == "barrier" - && record.rust.arguments == ["u32", "u32"] - && record.dialect.operands == ["i32", "i32"] - && record.llvm.as_ref().is_some_and(|llvm| { - llvm.arguments == ["i32", "i32"] - }) - } - ExecutionControlOperation::GridDependencyLaunchDependents - | ExecutionControlOperation::GridDependencyWait => { - record.rust.module == "grid" - && record.rust.arguments.is_empty() - && record.dialect.operands.is_empty() - && record - .llvm - .as_ref() - .is_some_and(|llvm| llvm.arguments.is_empty()) - } - ExecutionControlOperation::SetMaxNRegInc - | ExecutionControlOperation::SetMaxNRegDec => { - record.rust.module == "thread" - && record.rust.arguments == ["u32"] - && record.dialect.operands.is_empty() - && record - .llvm - .as_ref() - .is_some_and(|llvm| llvm.arguments == ["i32"]) - } - } - }), - "{} is outside the closed generated execution-control recipe", - record.id - ), - "tcgen05" => ensure!( - tcgen05_render_contract(record), - "{} is outside the closed generated tcgen05 recipe", - record.id - ), - "cp_async_copy" => ensure!( - record.rust.module == "async_copy" - && record.rust.result == "()" - && !record.rust.safe - && !record.rust.must_use - && record.llvm.as_ref().is_some_and(|llvm| { - matches!(llvm.arguments.as_slice(), [dst, src] - if dst == "shared_ptr" && src == "global_ptr") - || matches!(llvm.arguments.as_slice(), [dst, src, size] - if dst == "shared_ptr" && src == "global_ptr" && size == "i32") - }) - && record.llvm.as_ref().is_some_and(|llvm| llvm.results.is_empty()) - && record.dialect.results.is_empty() - && record.lowering == "generated_cp_async_copy" - && record.cp_async_copy.is_some(), - "{} is outside the closed generated cp.async copy recipe", - record.id - ), - "cp_async_control" => ensure!( - record.rust.module == "async_copy" - && record.rust.result == "()" - && !record.rust.safe - && !record.rust.must_use - && record.llvm.as_ref().is_some_and(|llvm| { - llvm.results.is_empty() - && (llvm.arguments.is_empty() || llvm.arguments == ["i32"]) - }) - && record.dialect.results.is_empty() - && record.lowering == "generated_cp_async_control" - && record.cp_async_control.is_some(), - "{} is outside the closed generated cp.async control recipe", - record.id - ), - "cp_async_mbarrier" => ensure!( - record.rust.module == "async_copy" - && record.rust.arguments == ["*mut u64"] - && record.rust.result == "()" - && !record.rust.safe - && !record.rust.must_use - && record.dialect.operands == ["ptr"] - && record.dialect.results.is_empty() - && record.lowering == "generated_cp_async_mbarrier" - && record.cp_async_mbarrier.as_ref().is_some_and(|bridge| { - bridge.adapter == CpAsyncMbarrierAdapter::PointerToVoid - && record.llvm.as_ref().is_some_and(|llvm| { - llvm.results.is_empty() - && match bridge.state_space { - CpAsyncMbarrierStateSpace::Generic => { - llvm.arguments == ["ptr"] - } - CpAsyncMbarrierStateSpace::Shared => { - llvm.arguments == ["shared_ptr"] - } - } - }) - }), - "{} is outside the closed generated cp.async mbarrier recipe", - record.id - ), - "mbarrier_basic" => ensure!( - record.rust.module == "barrier" - && !record.rust.safe - && record.lowering == "generated_mbarrier_basic" - && record.mbarrier_basic.as_ref().is_some_and(|mbarrier| { - mbarrier.state_space == MbarrierStateSpace::Shared - && match (mbarrier.operation, mbarrier.adapter) { - ( - MbarrierBasicOperation::Init, - MbarrierBasicAdapter::InitPointerCountToVoid, - ) => { - record.rust.arguments == ["*mut u64", "u32"] - && record.rust.result == "()" - && !record.rust.must_use - && record.llvm.as_ref().is_some_and(|llvm| { - llvm.arguments == ["shared_ptr", "i32"] - && llvm.results.is_empty() - }) - && record.dialect.operands == ["ptr", "i32"] - && record.dialect.results.is_empty() - } - ( - MbarrierBasicOperation::Arrive, - MbarrierBasicAdapter::ArrivePointerToToken, - ) => { - record.rust.arguments == ["*const u64"] - && record.rust.result == "u64" - && record.rust.must_use - && record.llvm.as_ref().is_some_and(|llvm| { - llvm.arguments == ["shared_ptr"] - && llvm.results == ["i64"] - }) - && record.dialect.operands == ["ptr"] - && record.dialect.results == ["i64"] - } - ( - MbarrierBasicOperation::TestWait, - MbarrierBasicAdapter::TestWaitPointerTokenToPredicate, - ) => { - record.rust.arguments == ["*const u64", "u64"] - && record.rust.result == "bool" - && record.rust.must_use - && record.llvm.as_ref().is_some_and(|llvm| { - llvm.arguments == ["shared_ptr", "i64"] - && llvm.results == ["i1"] - }) - && record.dialect.operands == ["ptr", "i64"] - && record.dialect.results == ["i1"] - } - ( - MbarrierBasicOperation::Inval, - MbarrierBasicAdapter::InvalPointerToVoid, - ) => { - record.rust.arguments == ["*mut u64"] - && record.rust.result == "()" - && !record.rust.must_use - && record.llvm.as_ref().is_some_and(|llvm| { - llvm.arguments == ["shared_ptr"] - && llvm.results.is_empty() - }) - && record.dialect.operands == ["ptr"] - && record.dialect.results.is_empty() - } - _ => false, - } - }), - "{} is outside the closed generated basic mbarrier recipe", - record.id - ), - "mbarrier_extended" => ensure!( - record.rust.module == "barrier" - && !record.rust.safe - && record.semantics.memory == "read_write" - && record.semantics.convergent - && record.lowering == "generated_mbarrier_extended_inline_ptx" - && record.backend_lowerings.iter().all(|lowering| { - lowering.mechanism == BackendLoweringMechanism::InlinePtx - }) - && record.mbarrier_extended.as_ref().is_some_and(|mbarrier| { - let source_matches = match mbarrier.source_contract { - MbarrierExtendedSourceContract::LlvmImported => record.llvm.is_some(), - MbarrierExtendedSourceContract::PtxNativeRawClusterAddress => { - record.llvm.is_none() - && matches!(record.source, IntrinsicSource::PtxNative { .. }) - } - }; - source_matches - && match (mbarrier.operation, mbarrier.adapter) { - ( - MbarrierExtendedOperation::ArriveExpectTxCta, - MbarrierExtendedAdapter::PointerTxCountBytesToTokenDroppingTxCount, - ) => record.id == "mbarrier_arrive_expect_tx", - ( - MbarrierExtendedOperation::ArriveExpectTxCluster, - MbarrierExtendedAdapter::PointerTxCountBytesToTokenDroppingTxCount, - ) => record.id == "mbarrier_arrive_expect_tx_cluster", - ( - MbarrierExtendedOperation::ArriveRemoteCluster, - MbarrierExtendedAdapter::RawClusterAddressToVoid, - ) => record.id == "mbarrier_arrive_cluster", - ( - MbarrierExtendedOperation::TryWaitTokenCta, - MbarrierExtendedAdapter::PointerTokenToPredicate, - ) => record.id == "mbarrier_try_wait", - ( - MbarrierExtendedOperation::TryWaitParityCta, - MbarrierExtendedAdapter::PointerParityToPredicate, - ) => record.id == "mbarrier_try_wait_parity", - ( - MbarrierExtendedOperation::TryWaitParityCluster, - MbarrierExtendedAdapter::PointerParityToPredicate, - ) => record.id == "mbarrier_try_wait_parity_cluster", - ( - MbarrierExtendedOperation::FenceProxyAsyncSharedCta, - MbarrierExtendedAdapter::ZeroOperandsToVoid, - ) => record.id == "fence_proxy_async_shared_cta", - ( - MbarrierExtendedOperation::FenceMbarrierInitReleaseCluster, - MbarrierExtendedAdapter::ZeroOperandsToVoid, - ) => record.id == "fence_mbarrier_init_release_cluster", - ( - MbarrierExtendedOperation::FenceProxyAsyncGenericReleaseSharedCtaCluster, - MbarrierExtendedAdapter::ZeroOperandsToVoid, - ) => record.id - == "fence_proxy_async_generic_release_shared_cta_cluster", - ( - MbarrierExtendedOperation::FenceProxyAsyncGenericAcquireSharedClusterCluster, - MbarrierExtendedAdapter::ZeroOperandsToVoid, - ) => record.id - == "fence_proxy_async_generic_acquire_shared_cluster_cluster", - ( - MbarrierExtendedOperation::Nanosleep, - MbarrierExtendedAdapter::NanosecondsToVoid, - ) => record.id == "nanosleep", - _ => false, - } - }), - "{} is outside the closed generated extended-mbarrier recipe", - record.id - ), - "sync" => { - if record.id == "sync_threads" { - ensure!( - record.rust.module == "thread" - && record.rust.name == "sync_threads" - && record.rust.arguments.is_empty() - && record.rust.result == "()" - && !record.rust.safe - && !record.rust.must_use - && record.llvm.as_ref().is_some_and(|llvm| { - llvm.symbol == "llvm.nvvm.barrier.cta.sync.aligned.all" - && llvm.arguments == ["i32"] - && llvm.results.is_empty() - }) - && record.dialect.op_type == "Barrier0Op" - && record.dialect.op_name == "nvvm.barrier0" - && record.dialect.operands.is_empty() - && record.dialect.results.is_empty() - && record.lowering == "generated_sync_threads", - "{} is outside the fixed-zero generated sync_threads recipe", - record.id - ); - } else { - ensure!( - threadfence_ptx_level(record).is_some() - && record.rust.module == "fence" - && record.rust.name == record.id - && record.rust.arguments.is_empty() - && record.rust.result == "()" - && record.rust.safe - && !record.rust.must_use - && record.llvm.as_ref().is_some_and(|llvm| { - llvm.arguments.is_empty() && llvm.results.is_empty() - }) - && record.dialect.operands.is_empty() - && record.dialect.results.is_empty() - && record.lowering == "direct_nvvm" - && record.semantics.memory == "read_write" - && !record.semantics.pure - && !record.semantics.convergent, - "{} is outside the closed generated thread-fence recipe", - record.id - ); - } - } - "vote" => ensure!( - record.rust.module == "warp" - && record.rust.arguments == ["u32", "bool"] - && matches!(record.rust.result.as_str(), "bool" | "u32") - && !record.rust.safe - && record.rust.must_use - && record.llvm.as_ref().is_some_and(|llvm| { - llvm.arguments == ["i32", "i1"] - && matches!(llvm.results.as_slice(), [result] - if result == "i1" || result == "i32") - }) - && record.dialect.operands == ["i32", "i1"] - && matches!(record.dialect.results.as_slice(), [result] - if result == "i1" || result == "i32") - && record.lowering == "generated_vote" - && record.vote.is_some(), - "{} is outside the closed generated vote.sync recipe", - record.id - ), - "active_mask" => ensure!( - record.id == "active_mask" - && record.rust.module == "warp" - && record.rust.arguments.is_empty() - && record.rust.result == "u32" - && record.rust.safe - && record.rust.must_use - && record.llvm.as_ref().is_some_and(|llvm| { - llvm.arguments.is_empty() && llvm.results == ["i32"] - }) - && record.dialect.operands.is_empty() - && record.dialect.results == ["i32"] - && record.lowering == "generated_active_mask" - && record.active_mask.is_some(), - "{} is outside the closed generated active-mask recipe", - record.id - ), - "warp_match" => ensure!( - record.rust.module == "warp" - && matches!(record.rust.arguments.as_slice(), [mask, value] - if mask == "u32" && matches!(value.as_str(), "u32" | "u64")) - && record.rust.result == "u32" - && !record.rust.safe - && record.rust.must_use - && record.llvm.as_ref().is_some_and(|llvm| { - matches!(llvm.arguments.as_slice(), [mask, value] - if mask == "i32" && matches!(value.as_str(), "i32" | "i64")) - && (matches!(llvm.results.as_slice(), [mask] if mask == "i32") - || matches!(llvm.results.as_slice(), [mask, predicate] - if mask == "i32" && predicate == "i1")) - }) - && matches!(record.dialect.operands.as_slice(), [mask, value] - if mask == "i32" && matches!(value.as_str(), "i32" | "i64")) - && record.dialect.results == ["i32"] - && record.lowering == "generated_warp_match" - && record.warp_match.is_some(), - "{} is outside the closed generated match.sync recipe", - record.id - ), - "elect" => ensure!( - record.id == "elect_sync" - && record.rust.module == "warp" - && record.rust.name == "elect_sync" - && record.rust.arguments == ["u32"] - && record.rust.result == "(u32, bool)" - && !record.rust.safe - && record.rust.must_use - && record.llvm.as_ref().is_some_and(|llvm| { - llvm.symbol == "llvm.nvvm.elect.sync" - && llvm.arguments == ["i32"] - && llvm.results == ["i32", "i1"] - }) - && record.dialect.op_type == "ElectSyncOp" - && record.dialect.op_name == "nvvm.elect_sync" - && record.dialect.operands == ["i32"] - && record.dialect.results == ["i32", "i1"] - && record.lowering == "generated_elect", - "{} is outside the closed generated elect.sync recipe", - record.id - ), - "warp_barrier" => ensure!( - record.id == "sync_mask" - && record.rust.module == "warp" - && record.rust.name == "sync_mask" - && record.rust.arguments == ["u32"] - && record.rust.result == "()" - && !record.rust.safe - && !record.rust.must_use - && record.llvm.as_ref().is_some_and(|llvm| { - llvm.symbol == "llvm.nvvm.bar.warp.sync" - && llvm.arguments == ["i32"] - && llvm.results.is_empty() - }) - && record.dialect.op_type == "BarWarpSyncOp" - && record.dialect.op_name == "nvvm.bar_warp_sync" - && record.dialect.operands == ["i32"] - && record.dialect.results.is_empty() - && record.lowering == "generated_warp_barrier" - && record.warp_barrier.as_ref().is_some_and(|barrier| { - barrier.adapter == WarpBarrierAdapter::DirectMemberMask - }), - "{} is outside the closed generated bar.warp.sync recipe", - record.id - ), - "warp_shuffle" => ensure!( - record.rust.module == "warp" - && matches!(record.rust.arguments.as_slice(), [mask, value, lane] - if mask == "u32" - && lane == "u32" - && matches!(value.as_str(), "u32" | "f32" | "u64") - && value == &record.rust.result) - && !record.rust.safe - && record.rust.must_use - && matches!(record.dialect.operands.as_slice(), [mask, value, lane] - if mask == "i32" - && lane == "i32" - && matches!(value.as_str(), "i32" | "f32" | "i64")) - && matches!(record.dialect.results.as_slice(), [result] - if result == &record.dialect.operands[1]) - && record.warp_shuffle.as_ref().is_some_and(|shuffle| { - match shuffle.value_kind { - WarpShuffleValueKind::I32 | WarpShuffleValueKind::F32 => { - let (value_ty, rust_ty) = match shuffle.value_kind { - WarpShuffleValueKind::I32 => ("i32", "u32"), - WarpShuffleValueKind::F32 => ("f32", "f32"), - WarpShuffleValueKind::I64 => unreachable!(), - }; - shuffle.adapter - == WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp - && shuffle.lane_encoding - == WarpShuffleOperandEncoding::RegisterOrImmediate - && shuffle.mask_encoding - == WarpShuffleOperandEncoding::RegisterOrImmediate - && record.lowering == "generated_warp_shuffle" - && record.rust.result == rust_ty - && record.dialect.operands[1] == value_ty - && record.llvm.as_ref().is_some_and(|llvm| { - matches!(llvm.arguments.as_slice(), [mask, value, lane, clamp] - if mask == "i32" - && value == value_ty - && lane == "i32" - && clamp == "i32") - && llvm.results == [value_ty] - }) - } - WarpShuffleValueKind::I64 => { - shuffle.adapter - == WarpShuffleAdapter::MaskValueLaneOrDeltaSplitI64LowHighB32InsertClampReassemble - && shuffle.lane_encoding - == WarpShuffleOperandEncoding::RegisterOnly - && shuffle.mask_encoding - == WarpShuffleOperandEncoding::RegisterOnly - && record.rust.result == "u64" - && record.dialect.operands[1] == "i64" - && record.lowering - == "generated_warp_shuffle_i64_inline_ptx" - && record.llvm.is_none() - && matches!(record.source, IntrinsicSource::PtxNative { .. }) - } - } - }), - "{} is outside the closed generated shfl.sync recipe", - record.id - ), - "integer_minmax" => ensure!( - record.integer_minmax.as_ref().is_some_and(|minmax| { - let module = match minmax.format { - IntegerMinMaxFormat::S32 => "int", - IntegerMinMaxFormat::S16x2 | IntegerMinMaxFormat::U16x2 => "i16x2", - }; - record.rust.module == module - }) - && record.rust.arguments.len() == 2 - && record - .rust - .arguments - .iter() - .all(|argument| argument == &record.rust.result) - && matches!(record.rust.result.as_str(), "i32" | "u32") - && record.rust.safe - && record.rust.must_use - && record.dialect.operands == ["i32", "i32"] - && record.dialect.results == ["i32"] - && record.lowering == "generated_integer_minmax_inline_ptx", - "{} is outside the closed generated integer-min/max recipe", - record.id - ), - family => ensure!(false, "{} has unrenderable family {family}", record.id), - }; - } - Ok(()) -} - -fn rust_header(catalog: &CatalogFile, catalog_sha256: &str) -> String { - format!( - "/*\n * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: Apache-2.0\n */\n\n// @generated by cuda-intrinsics-gen {}. DO NOT EDIT.\n// Catalog schema/version: {}/{}; intrinsic ABI: v{}; catalog SHA-256: {}.\n// LLVM source: {}.\n\n", - catalog.generator_version, - catalog.schema, - catalog.catalog_version, - catalog.intrinsic_abi, - catalog_sha256, - catalog.source.llvm_revision - ) -} - -fn llvm_header(catalog: &CatalogFile, catalog_sha256: &str) -> String { - format!( - "; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n; SPDX-License-Identifier: Apache-2.0\n; @generated by cuda-intrinsics-gen {}. DO NOT EDIT.\n; Catalog schema/version: {}/{}; intrinsic ABI: v{}; catalog SHA-256: {}.\n; LLVM source: {}.\n\n", - catalog.generator_version, - catalog.schema, - catalog.catalog_version, - catalog.intrinsic_abi, - catalog_sha256, - catalog.source.llvm_revision - ) -} - -fn markdown_header(catalog: &CatalogFile, catalog_sha256: &str) -> String { - format!( - "\n\n", - catalog.generator_version, - catalog.schema, - catalog.catalog_version, - catalog.intrinsic_abi, - catalog_sha256, - catalog.source.llvm_revision - ) -} - -fn sregs(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "sreg") -} - -fn ldmatrix(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "ldmatrix") -} - -fn is_blackwell_ldmatrix(record: &CatalogIntrinsic) -> bool { - record - .ldmatrix - .as_ref() - .is_some_and(|ldmatrix| ldmatrix.variant.shape != LdmatrixShape::M8n8) -} - -const BLACKWELL_LDMATRIX_EFFECTIVE_FLOORS: &str = "sm_100a PTX 8.6, sm_100f PTX 8.8, sm_103a PTX 8.8, sm_103f PTX 8.8, sm_110a PTX 9.0, sm_110f PTX 9.0, sm_120a PTX 8.7, sm_120f PTX 8.8, sm_121a PTX 8.8, sm_121f PTX 8.8"; - -fn ldmatrix_compat_op(record: &CatalogIntrinsic) -> Option<(&'static str, &'static str)> { - match record.id.as_str() { - "ldmatrix_m8n8_x1_b16" => Some(("LdmatrixX1Op", "nvvm.ldmatrix_x1")), - "ldmatrix_m8n8_x1_trans_b16" => Some(("LdmatrixX1TransOp", "nvvm.ldmatrix_x1_trans")), - "ldmatrix_m8n8_x2_b16" => Some(("LdmatrixX2Op", "nvvm.ldmatrix_x2")), - "ldmatrix_m8n8_x2_trans_b16" => Some(("LdmatrixX2TransOp", "nvvm.ldmatrix_x2_trans")), - "ldmatrix_m8n8_x4_b16" => Some(("LdmatrixX4Op", "nvvm.ldmatrix_x4")), - "ldmatrix_m8n8_x4_trans_b16" => Some(("LdmatrixX4TransOp", "nvvm.ldmatrix_x4_trans")), - _ => None, - } -} - -fn stmatrices(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "stmatrix") -} - -fn stmatrix_variant(record: &CatalogIntrinsic) -> Option<(StmatrixMultiplicity, StmatrixLayout)> { - match record.id.as_str() { - "stmatrix_m8n8_x2_b16" => Some((StmatrixMultiplicity::X2, StmatrixLayout::Normal)), - "stmatrix_m8n8_x2_trans_b16" => { - Some((StmatrixMultiplicity::X2, StmatrixLayout::Transposed)) - } - "stmatrix_m8n8_x4_b16" => Some((StmatrixMultiplicity::X4, StmatrixLayout::Normal)), - "stmatrix_m8n8_x4_trans_b16" => { - Some((StmatrixMultiplicity::X4, StmatrixLayout::Transposed)) - } - _ => None, - } -} - -fn stmatrix_compatibility_name(record: &CatalogIntrinsic) -> &'static str { - match record.id.as_str() { - "stmatrix_m8n8_x2_b16" => "stmatrix_m8n8_x2", - "stmatrix_m8n8_x2_trans_b16" => "stmatrix_m8n8_x2_trans", - "stmatrix_m8n8_x4_b16" => "stmatrix_m8n8_x4", - "stmatrix_m8n8_x4_trans_b16" => "stmatrix_m8n8_x4_trans", - _ => panic!("unknown stmatrix record {}", record.id), - } -} - -fn register_mmas(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "register_mma") -} - -fn sparse_mmas(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "sparse_mma") -} - -fn packed_atomics(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "packed_atomic") -} - -fn redux(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "redux") -} - -fn dot_products(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "dotprod") -} - -fn packed_alus(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "packed_alu") -} - -fn packed_alu_format_shape( - format: PackedAluFormat, -) -> ( - &'static str, - bool, - &'static str, - &'static str, - PackedAluAdapter, -) { - match format { - PackedAluFormat::Bf16x2 => ( - "bf16x2", - false, - "u32", - "i32", - PackedAluAdapter::DirectPackedU32, - ), - PackedAluFormat::F16x2 => ( - "f16x2", - true, - "u32", - "i32", - PackedAluAdapter::DirectPackedU32, - ), - PackedAluFormat::F32x2 => ( - "f32x2", - true, - "u64", - "i64", - PackedAluAdapter::DirectPackedU64, - ), - } -} - -fn packed_alu_width(record: &CatalogIntrinsic) -> u32 { - match record - .packed_alu - .as_ref() - .expect("packed-ALU record") - .adapter - { - PackedAluAdapter::DirectPackedU32 => 32, - PackedAluAdapter::DirectPackedU64 => 64, - } -} - -fn packed_alu_register_constraint(record: &CatalogIntrinsic) -> &'static str { - match packed_alu_width(record) { - 32 => "r", - 64 => "l", - _ => unreachable!("closed packed-ALU carrier width"), - } -} - -fn packed_conversions(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "packed_conversion") -} - -fn scalar_conversions(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "scalar_conversion") -} - -fn scalar_conversion_rounding_attr(record: &CatalogIntrinsic) -> &'static str { - match record - .scalar_conversion - .as_ref() - .expect("scalar-conversion contract") - .rounding - { - ScalarConversionRounding::NearestAway => "ScalarConversionRoundingAttr::NearestAway", - ScalarConversionRounding::NearestEven => "ScalarConversionRoundingAttr::NearestEven", - ScalarConversionRounding::TowardZero => "ScalarConversionRoundingAttr::TowardZero", - } -} - -fn scalar_conversion_saturation_attr(record: &CatalogIntrinsic) -> &'static str { - match record - .scalar_conversion - .as_ref() - .expect("scalar-conversion contract") - .saturation - { - ScalarConversionSaturation::None => "ScalarConversionSaturationAttr::None", - ScalarConversionSaturation::Relu => "ScalarConversionSaturationAttr::Relu", - ScalarConversionSaturation::Satfinite => "ScalarConversionSaturationAttr::Satfinite", - ScalarConversionSaturation::ReluSatfinite => { - "ScalarConversionSaturationAttr::ReluSatfinite" - } - } -} - -fn scalar_conversion_ptx_mnemonic(record: &CatalogIntrinsic) -> String { - format!( - "{}.{}", - record.expected_ptx.mnemonic, - record.expected_ptx.modifiers.join(".") - ) -} - -fn scalar_arithmetics(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "scalar_arithmetic") -} - -fn extended_minmax(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "extended_minmax") -} - -fn extended_minmax_contract(record: &CatalogIntrinsic) -> &crate::model::ExtendedMinMax { - record - .extended_minmax - .as_ref() - .expect("extended-minmax contract") -} - -fn extended_minmax_format_attr(record: &CatalogIntrinsic) -> &'static str { - match extended_minmax_contract(record).format { - ExtendedMinMaxFormat::F32 => "ExtendedMinMaxFormatAttr::F32", - ExtendedMinMaxFormat::F16 => "ExtendedMinMaxFormatAttr::F16", - ExtendedMinMaxFormat::Bf16 => "ExtendedMinMaxFormatAttr::Bf16", - ExtendedMinMaxFormat::F16x2 => "ExtendedMinMaxFormatAttr::F16x2", - ExtendedMinMaxFormat::Bf16x2 => "ExtendedMinMaxFormatAttr::Bf16x2", - } -} - -fn extended_minmax_operation_attr(record: &CatalogIntrinsic) -> &'static str { - match extended_minmax_contract(record).operation { - ExtendedMinMaxOperation::Min => "ExtendedMinMaxOperationAttr::Min", - ExtendedMinMaxOperation::Max => "ExtendedMinMaxOperationAttr::Max", - } -} - -fn extended_minmax_subnormal_attr(record: &CatalogIntrinsic) -> &'static str { - match extended_minmax_contract(record).subnormal { - ExtendedMinMaxSubnormal::Preserve => "ExtendedMinMaxSubnormalAttr::Preserve", - ExtendedMinMaxSubnormal::Ftz => "ExtendedMinMaxSubnormalAttr::Ftz", - } -} - -fn extended_minmax_nan_attr(record: &CatalogIntrinsic) -> &'static str { - match extended_minmax_contract(record).nan { - ExtendedMinMaxNan::Number => "ExtendedMinMaxNanAttr::Number", - ExtendedMinMaxNan::Nan => "ExtendedMinMaxNanAttr::Nan", - } -} - -fn extended_minmax_xorsign_abs_attr(record: &CatalogIntrinsic) -> &'static str { - if extended_minmax_contract(record).xorsign_abs { - "ExtendedMinMaxXorSignAbsAttr::Enabled" - } else { - "ExtendedMinMaxXorSignAbsAttr::Disabled" - } -} - -fn extended_minmax_carrier(record: &CatalogIntrinsic) -> &'static str { - match extended_minmax_contract(record).format { - ExtendedMinMaxFormat::F32 => "MinMaxCarrier::F32", - ExtendedMinMaxFormat::F16 | ExtendedMinMaxFormat::Bf16 => "MinMaxCarrier::Half16", - ExtendedMinMaxFormat::F16x2 | ExtendedMinMaxFormat::Bf16x2 => "MinMaxCarrier::PackedU32", - } -} - -fn extended_minmax_rust_type(record: &CatalogIntrinsic) -> &'static str { - match extended_minmax_contract(record).format { - ExtendedMinMaxFormat::F32 => "f32", - ExtendedMinMaxFormat::F16 | ExtendedMinMaxFormat::Bf16 => "u16", - ExtendedMinMaxFormat::F16x2 | ExtendedMinMaxFormat::Bf16x2 => "u32", - } -} - -fn extended_minmax_ptx_mnemonic(record: &CatalogIntrinsic) -> String { - format!( - "{}.{}", - record.expected_ptx.mnemonic, - record.expected_ptx.modifiers.join(".") - ) -} - -fn scalar_arithmetic_contract(record: &CatalogIntrinsic) -> &crate::model::ScalarArithmetic { - record - .scalar_arithmetic - .as_ref() - .expect("scalar-arithmetic contract") -} - -fn scalar_arithmetic_format_attr(record: &CatalogIntrinsic) -> &'static str { - match scalar_arithmetic_contract(record).format { - ScalarArithmeticFormat::F32 => "ScalarArithmeticFormatAttr::F32", - ScalarArithmeticFormat::F64 => "ScalarArithmeticFormatAttr::F64", - } -} - -fn scalar_arithmetic_operation_attr(record: &CatalogIntrinsic) -> &'static str { - match scalar_arithmetic_contract(record).operation { - ScalarArithmeticOperation::Mul => "ScalarArithmeticOperationAttr::Mul", - ScalarArithmeticOperation::Div => "ScalarArithmeticOperationAttr::Div", - ScalarArithmeticOperation::Fma => "ScalarArithmeticOperationAttr::Fma", - ScalarArithmeticOperation::Add => "ScalarArithmeticOperationAttr::Add", - } -} - -fn scalar_arithmetic_rounding_attr(record: &CatalogIntrinsic) -> &'static str { - match scalar_arithmetic_contract(record).rounding { - ScalarArithmeticRounding::Rn => "ScalarArithmeticRoundingAttr::Rn", - ScalarArithmeticRounding::Rz => "ScalarArithmeticRoundingAttr::Rz", - ScalarArithmeticRounding::Rm => "ScalarArithmeticRoundingAttr::Rm", - ScalarArithmeticRounding::Rp => "ScalarArithmeticRoundingAttr::Rp", - } -} - -fn scalar_arithmetic_subnormal_attr(record: &CatalogIntrinsic) -> &'static str { - match scalar_arithmetic_contract(record).subnormal { - ScalarArithmeticSubnormal::Preserve => "ScalarArithmeticSubnormalAttr::Preserve", - ScalarArithmeticSubnormal::Ftz => "ScalarArithmeticSubnormalAttr::Ftz", - } -} - -fn scalar_arithmetic_saturation_attr(record: &CatalogIntrinsic) -> &'static str { - match scalar_arithmetic_contract(record).saturation { - ScalarArithmeticSaturation::None => "ScalarArithmeticSaturationAttr::None", - ScalarArithmeticSaturation::Sat => "ScalarArithmeticSaturationAttr::Sat", - } -} - -fn scalar_arithmetic_arity(record: &CatalogIntrinsic) -> usize { - match scalar_arithmetic_contract(record).operation { - ScalarArithmeticOperation::Mul - | ScalarArithmeticOperation::Div - | ScalarArithmeticOperation::Add => 2, - ScalarArithmeticOperation::Fma => 3, - } -} - -fn scalar_arithmetic_rust_type(record: &CatalogIntrinsic) -> &'static str { - match scalar_arithmetic_contract(record).format { - ScalarArithmeticFormat::F32 => "f32", - ScalarArithmeticFormat::F64 => "f64", - } -} - -fn scalar_arithmetic_llvm_type(record: &CatalogIntrinsic) -> &'static str { - match scalar_arithmetic_contract(record).format { - ScalarArithmeticFormat::F32 => "float", - ScalarArithmeticFormat::F64 => "double", - } -} - -fn scalar_arithmetic_ptx_mnemonic(record: &CatalogIntrinsic) -> String { - format!( - "{}.{}", - record.expected_ptx.mnemonic, - record.expected_ptx.modifiers.join(".") - ) -} - -fn scalar_arithmetic_llvm_mechanism(record: &CatalogIntrinsic) -> BackendLoweringMechanism { - record - .backend_lowerings - .iter() - .find(|lowering| lowering.backend == IntrinsicBackend::LlvmNvptx) - .expect("scalar arithmetic has an LLVM-NVPTX route") - .mechanism -} - -fn scalar_maths(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "scalar_math") -} - -fn scalar_math_contract(record: &CatalogIntrinsic) -> &crate::model::ScalarMath { - record.scalar_math.as_ref().expect("scalar-math contract") -} - -fn scalar_math_format_attr(record: &CatalogIntrinsic) -> &'static str { - match scalar_math_contract(record).format { - ScalarMathFormat::F16 => "ScalarMathFormatAttr::F16", - ScalarMathFormat::F32 => "ScalarMathFormatAttr::F32", - ScalarMathFormat::F64 => "ScalarMathFormatAttr::F64", - } -} - -fn scalar_math_operation_attr(record: &CatalogIntrinsic) -> &'static str { - match scalar_math_contract(record).operation { - ScalarMathOperation::Sin => "ScalarMathOperationAttr::Sin", - ScalarMathOperation::Cos => "ScalarMathOperationAttr::Cos", - ScalarMathOperation::Ex2 => "ScalarMathOperationAttr::Ex2", - ScalarMathOperation::Lg2 => "ScalarMathOperationAttr::Lg2", - ScalarMathOperation::Rcp => "ScalarMathOperationAttr::Rcp", - ScalarMathOperation::Rsqrt => "ScalarMathOperationAttr::Rsqrt", - ScalarMathOperation::Sqrt => "ScalarMathOperationAttr::Sqrt", - ScalarMathOperation::Tanh => "ScalarMathOperationAttr::Tanh", - } -} - -fn scalar_math_precision_attr(record: &CatalogIntrinsic) -> &'static str { - match scalar_math_contract(record).precision { - ScalarMathPrecision::Approx => "ScalarMathPrecisionAttr::Approx", - ScalarMathPrecision::Rn => "ScalarMathPrecisionAttr::Rn", - ScalarMathPrecision::Rz => "ScalarMathPrecisionAttr::Rz", - ScalarMathPrecision::Rm => "ScalarMathPrecisionAttr::Rm", - ScalarMathPrecision::Rp => "ScalarMathPrecisionAttr::Rp", - } -} - -fn scalar_math_subnormal_attr(record: &CatalogIntrinsic) -> &'static str { - match scalar_math_contract(record).subnormal { - ScalarMathSubnormal::Preserve => "ScalarMathSubnormalAttr::Preserve", - ScalarMathSubnormal::Ftz => "ScalarMathSubnormalAttr::Ftz", - } -} - -fn scalar_math_llvm_type(record: &CatalogIntrinsic) -> &'static str { - match scalar_math_contract(record).format { - ScalarMathFormat::F16 => "i16", - ScalarMathFormat::F32 => "float", - ScalarMathFormat::F64 => "double", - } -} - -fn scalar_math_ptx_mnemonic(record: &CatalogIntrinsic) -> String { - format!( - "{}.{}", - record.expected_ptx.mnemonic, - record.expected_ptx.modifiers.join(".") - ) -} - -fn scalar_math_llvm_mechanism(record: &CatalogIntrinsic) -> BackendLoweringMechanism { - record - .backend_lowerings - .iter() - .find(|lowering| lowering.backend == IntrinsicBackend::LlvmNvptx) - .expect("scalar math has an LLVM-NVPTX route") - .mechanism -} - -fn prmts(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "prmt") -} - -fn cluster_barriers(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "cluster_barrier") -} - -fn cluster_barrier_attr(record: &CatalogIntrinsic) -> &'static str { - match record - .cluster_barrier - .as_ref() - .expect("cluster-barrier contract") - .mode - { - ClusterBarrierMode::Arrive => "ClusterBarrierModeAttr::Arrive", - ClusterBarrierMode::ArriveAligned => "ClusterBarrierModeAttr::ArriveAligned", - ClusterBarrierMode::ArriveRelaxed => "ClusterBarrierModeAttr::ArriveRelaxed", - ClusterBarrierMode::ArriveRelaxedAligned => "ClusterBarrierModeAttr::ArriveRelaxedAligned", - ClusterBarrierMode::Wait => "ClusterBarrierModeAttr::Wait", - ClusterBarrierMode::WaitAligned => "ClusterBarrierModeAttr::WaitAligned", - } -} - -fn cluster_barrier_template(record: &CatalogIntrinsic) -> String { - format!( - "{}.{};", - record.expected_ptx.mnemonic, - record.expected_ptx.modifiers.join(".") - ) -} - -fn debug_controls(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "debug_control") -} - -fn cluster_memory(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "cluster_memory") -} - -fn clc_intrinsics(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "clc") -} - -#[derive(Clone, Copy)] -enum ClcSafetyArgNames { - RawAbi, - Compatibility, -} - -fn clc_safety_lines(operation: ClcOperation, names: ClcSafetyArgNames) -> Vec { - let (response, mbar, resp_lo, resp_hi) = match names { - ClcSafetyArgNames::RawAbi => ("`_arg0`", "`_arg1`", "`_arg0`", "`_arg1`"), - ClcSafetyArgNames::Compatibility => ("`response`", "`mbar`", "`resp_lo`", "`resp_hi`"), - }; - match operation { - ClcOperation::TryCancel => vec![ - format!( - "{response} must designate a writable, 16-byte-aligned 16-byte region in `shared::cta` memory, and {mbar} must designate a valid, aligned shared-memory mbarrier." - ), - format!( - "The calling CTA must issue `arrive_expect_tx` for {mbar} with 16 bytes before this request." - ), - "This CTA must not have observed a prior CLC try-cancel request fail.".into(), - ], - ClcOperation::TryCancelMulticast => vec![ - format!( - "{response} must designate a writable, 16-byte-aligned 16-byte region in `shared::cta` memory, and {mbar} must designate a valid, aligned shared-memory mbarrier." - ), - format!( - "Each CTA must issue `arrive_expect_tx` for its own corresponding mbarrier with 16 bytes before this request." - ), - "This CTA must not have observed a prior CLC try-cancel request fail.".into(), - "The response is written to every CTA's shared-memory window and signals every CTA's mbarrier.".into(), - "Every CTA in the cluster must still be active.".into(), - ], - ClcOperation::QueryIsCanceled => vec![format!( - "{resp_lo} and {resp_hi} must be the two halves of an opaque response observed complete from a prior CLC try-cancel request." - )], - ClcOperation::QueryGetFirstCtaidX - | ClcOperation::QueryGetFirstCtaidY - | ClcOperation::QueryGetFirstCtaidZ => vec![ - format!( - "{resp_lo} and {resp_hi} must be the two halves of an opaque response observed complete from a prior CLC try-cancel request." - ), - "The result is meaningful only when `clc_query_is_canceled` returned one for this response.".into(), - ], - } -} - -fn render_clc_safety_lines(output: &mut String, operation: ClcOperation, names: ClcSafetyArgNames) { - for line in clc_safety_lines(operation, names) { - writeln!(output, "/// {line}").unwrap(); - } -} - -fn tma_intrinsics(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "tma") -} - -fn execution_controls(catalog: &CatalogFile) -> impl Iterator { - catalog.intrinsics.iter().filter(|record| { - matches!( - record.family.as_str(), - "counted_barrier" | "grid_dependency" | "register_control" - ) - }) -} - -fn execution_control_family<'a>( - catalog: &'a CatalogFile, - family: &'a str, -) -> impl Iterator { - execution_controls(catalog).filter(move |record| record.family == family) -} - -fn wgmma_controls(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "wgmma_control") -} - -fn wgmma_control(catalog: &CatalogFile, mode: WgmmaControlMode) -> &CatalogIntrinsic { - wgmma_controls(catalog) - .find(|record| { - record - .wgmma_control - .as_ref() - .is_some_and(|control| control.mode == mode) - }) - .expect("complete WGMMA-control family") -} - -fn wgmma_control_template(record: &CatalogIntrinsic) -> String { - let operand = if record - .wgmma_control - .as_ref() - .is_some_and(|control| control.mode == WgmmaControlMode::WaitGroup) - { - " $0" - } else { - "" - }; - format!( - "{}.{}{operand};", - record.expected_ptx.mnemonic, - record.expected_ptx.modifiers.join(".") - ) -} - -fn tcgen05_intrinsics(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "tcgen05") -} - -fn tcgen05_mma_intrinsics(catalog: &CatalogFile) -> impl Iterator { - tcgen05_intrinsics(catalog).filter(|record| { - record - .tcgen05 - .as_ref() - .is_some_and(|tcgen05| tcgen05.mma.is_some()) - }) -} - -fn tcgen05_non_mma_intrinsics(catalog: &CatalogFile) -> impl Iterator { - tcgen05_intrinsics(catalog).filter(|record| { - record - .tcgen05 - .as_ref() - .is_some_and(|tcgen05| tcgen05.mma.is_none()) - }) -} - -fn tcgen05_mma_form_name(form: Tcgen05MmaForm) -> &'static str { - match form { - Tcgen05MmaForm::Shared => "Shared", - Tcgen05MmaForm::Tensor => "Tensor", - Tcgen05MmaForm::TensorAshift => "TensorAshift", - Tcgen05MmaForm::SpShared => "SpShared", - Tcgen05MmaForm::SpTensor => "SpTensor", - Tcgen05MmaForm::SpTensorAshift => "SpTensorAshift", - Tcgen05MmaForm::WsShared => "WsShared", - Tcgen05MmaForm::WsSharedZeroColMask => "WsSharedZeroColMask", - Tcgen05MmaForm::WsSpShared => "WsSpShared", - Tcgen05MmaForm::WsSpSharedZeroColMask => "WsSpSharedZeroColMask", - Tcgen05MmaForm::WsSpTensor => "WsSpTensor", - Tcgen05MmaForm::WsSpTensorZeroColMask => "WsSpTensorZeroColMask", - Tcgen05MmaForm::WsTensor => "WsTensor", - Tcgen05MmaForm::WsTensorZeroColMask => "WsTensorZeroColMask", - } -} - -fn tcgen05_mma_form_attr(form: Tcgen05MmaForm) -> String { - format!("Tcgen05MmaFormAttr::{}", tcgen05_mma_form_name(form)) -} - -fn tcgen05_mma_kind_attr(kind: Tcgen05MmaKind) -> &'static str { - match kind { - Tcgen05MmaKind::F16 => "Tcgen05MmaKindAttr::F16", - Tcgen05MmaKind::Tf32 => "Tcgen05MmaKindAttr::Tf32", - Tcgen05MmaKind::F8f6f4 => "Tcgen05MmaKindAttr::F8f6f4", - Tcgen05MmaKind::I8 => "Tcgen05MmaKindAttr::I8", - } -} - -fn tcgen05_mma_b_usage_attr(usage: Tcgen05MmaBUsage) -> &'static str { - match usage { - Tcgen05MmaBUsage::Discard => "Tcgen05MmaBUsageAttr::Discard", - Tcgen05MmaBUsage::Fill => "Tcgen05MmaBUsageAttr::Fill", - Tcgen05MmaBUsage::Use => "Tcgen05MmaBUsageAttr::Use", - Tcgen05MmaBUsage::LastUse => "Tcgen05MmaBUsageAttr::LastUse", - } -} - -fn tcgen05_mma_runtime_parameters(mma: &Tcgen05Mma) -> Vec<(&'static str, &'static str)> { - if mma.alias.is_some() && mma.form == Tcgen05MmaForm::WsTensor { - return vec![ - ("d_tmem", "u32"), - ("a_tmem", "u32"), - ("legacy_a_desc", "u64"), - ("b_desc", "u64"), - ("idesc", "u32"), - ("enable_d", "bool"), - ]; - } - let mut parameters = match mma.form { - Tcgen05MmaForm::Shared - | Tcgen05MmaForm::SpShared - | Tcgen05MmaForm::WsShared - | Tcgen05MmaForm::WsSharedZeroColMask - | Tcgen05MmaForm::WsSpShared - | Tcgen05MmaForm::WsSpSharedZeroColMask => vec![ - ("d_tmem", "u32"), - ("a_desc", "u64"), - ("b_desc", "u64"), - ("idesc", "u32"), - ("enable_d", "bool"), - ], - Tcgen05MmaForm::Tensor - | Tcgen05MmaForm::TensorAshift - | Tcgen05MmaForm::SpTensor - | Tcgen05MmaForm::SpTensorAshift - | Tcgen05MmaForm::WsSpTensor - | Tcgen05MmaForm::WsSpTensorZeroColMask - | Tcgen05MmaForm::WsTensor - | Tcgen05MmaForm::WsTensorZeroColMask => vec![ - ("d_tmem", "u32"), - ("a_tmem", "u32"), - ("b_desc", "u64"), - ("idesc", "u32"), - ("enable_d", "bool"), - ], - }; - if matches!( - mma.form, - Tcgen05MmaForm::SpShared - | Tcgen05MmaForm::SpTensor - | Tcgen05MmaForm::SpTensorAshift - | Tcgen05MmaForm::WsSpShared - | Tcgen05MmaForm::WsSpSharedZeroColMask - | Tcgen05MmaForm::WsSpTensor - | Tcgen05MmaForm::WsSpTensorZeroColMask - ) { - parameters.push(("metadata_tmem", "u32")); - } - if matches!( - mma.form, - Tcgen05MmaForm::WsSharedZeroColMask - | Tcgen05MmaForm::WsSpSharedZeroColMask - | Tcgen05MmaForm::WsSpTensorZeroColMask - | Tcgen05MmaForm::WsTensorZeroColMask - ) { - parameters.push(("zero_column_mask", "u64")); - } - parameters -} - -fn tcgen05_mma_selector_parameters( - layout: Tcgen05MmaSelectorLayout, -) -> [(&'static str, &'static str); 3] { - match layout { - Tcgen05MmaSelectorLayout::Base { .. } => [ - ("KIND", "kind"), - ("CTA_GROUP", "cta_group"), - ("COLLECTOR_A", "collector_a"), - ], - Tcgen05MmaSelectorLayout::WarpSpecialized { .. } => [ - ("KIND", "kind"), - ("B_BUFFER", "b_buffer"), - ("B_USAGE", "b_usage"), - ], - } -} - -fn tcgen05_mma_is_ws(form: Tcgen05MmaForm) -> bool { - matches!( - form, - Tcgen05MmaForm::WsShared - | Tcgen05MmaForm::WsSharedZeroColMask - | Tcgen05MmaForm::WsSpShared - | Tcgen05MmaForm::WsSpSharedZeroColMask - | Tcgen05MmaForm::WsSpTensor - | Tcgen05MmaForm::WsSpTensorZeroColMask - | Tcgen05MmaForm::WsTensor - | Tcgen05MmaForm::WsTensorZeroColMask - ) -} - -fn tcgen05_mma_is_sparse(form: Tcgen05MmaForm) -> bool { - matches!( - form, - Tcgen05MmaForm::SpShared - | Tcgen05MmaForm::SpTensor - | Tcgen05MmaForm::SpTensorAshift - | Tcgen05MmaForm::WsSpShared - | Tcgen05MmaForm::WsSpSharedZeroColMask - | Tcgen05MmaForm::WsSpTensor - | Tcgen05MmaForm::WsSpTensorZeroColMask - ) -} - -fn tcgen05_mma_is_tensor_a(form: Tcgen05MmaForm) -> bool { - matches!( - form, - Tcgen05MmaForm::Tensor - | Tcgen05MmaForm::TensorAshift - | Tcgen05MmaForm::SpTensor - | Tcgen05MmaForm::SpTensorAshift - | Tcgen05MmaForm::WsSpTensor - | Tcgen05MmaForm::WsSpTensorZeroColMask - | Tcgen05MmaForm::WsTensor - | Tcgen05MmaForm::WsTensorZeroColMask - ) -} - -fn tcgen05_mma_is_ashift(form: Tcgen05MmaForm) -> bool { - matches!( - form, - Tcgen05MmaForm::TensorAshift | Tcgen05MmaForm::SpTensorAshift - ) -} - -fn tcgen05_mma_has_zero_col_mask(form: Tcgen05MmaForm) -> bool { - matches!( - form, - Tcgen05MmaForm::WsSharedZeroColMask - | Tcgen05MmaForm::WsSpSharedZeroColMask - | Tcgen05MmaForm::WsSpTensorZeroColMask - | Tcgen05MmaForm::WsTensorZeroColMask - ) -} - -fn tcgen05_mma_kind_name(kind: Tcgen05MmaKind) -> &'static str { - match kind { - Tcgen05MmaKind::F16 => "f16", - Tcgen05MmaKind::Tf32 => "tf32", - Tcgen05MmaKind::F8f6f4 => "f8f6f4", - Tcgen05MmaKind::I8 => "i8", - } -} - -fn tcgen05_mma_b_usage_name(usage: Tcgen05MmaBUsage) -> &'static str { - match usage { - Tcgen05MmaBUsage::Discard => "discard", - Tcgen05MmaBUsage::LastUse => "lastuse", - Tcgen05MmaBUsage::Fill => "fill", - Tcgen05MmaBUsage::Use => "use", - } -} - -fn tcgen05_mma_inline_asm( - form: Tcgen05MmaForm, - kind: Tcgen05MmaKind, - cta_group: u8, - collector_a: Option<&str>, - b_buffer: Option, - b_usage: Option, -) -> (String, String) { - let mut instruction = "tcgen05.mma".to_owned(); - if tcgen05_mma_is_ws(form) { - instruction.push_str(".ws"); - } - if tcgen05_mma_is_sparse(form) { - instruction.push_str(".sp"); - } - write!( - instruction, - ".cta_group::{cta_group}.kind::{}", - tcgen05_mma_kind_name(kind) - ) - .unwrap(); - if tcgen05_mma_is_ws(form) { - write!( - instruction, - ".collector::b{}::{}", - b_buffer.expect("warp-specialized B buffer"), - tcgen05_mma_b_usage_name(b_usage.expect("warp-specialized B usage")) - ) - .unwrap(); - } else { - write!( - instruction, - ".collector::a::{}", - collector_a.expect("base collector A usage") - ) - .unwrap(); - if tcgen05_mma_is_ashift(form) { - instruction.push_str(".ashift"); - } - } - - let a = if tcgen05_mma_is_tensor_a(form) { - "[$1]" - } else { - "$1" - }; - write!(instruction, " [$0], {a}, $2").unwrap(); - if tcgen05_mma_is_sparse(form) { - instruction.push_str(", [$5]"); - } - instruction.push_str(", $3, %enable_pred"); - if tcgen05_mma_has_zero_col_mask(form) { - instruction.push_str(if tcgen05_mma_is_sparse(form) { - ", $6" - } else { - ", $5" - }); - } - instruction.push(';'); - - let template = - format!("{{ .reg .pred %enable_pred; setp.ne.s32 %enable_pred, $4, 0; {instruction} }}"); - let mut constraints = vec![ - "r", - if tcgen05_mma_is_tensor_a(form) { - "r" - } else { - "l" - }, - "l", - "r", - "r", - ]; - if tcgen05_mma_is_sparse(form) { - constraints.push("r"); - } - if tcgen05_mma_has_zero_col_mask(form) { - constraints.push("l"); - } - constraints.push("~{memory}"); - (template, constraints.join(",")) -} - -fn tcgen05_participation_doc(operation: Tcgen05Operation) -> Option<&'static str> { - match operation { - Tcgen05Operation::AllocCg2 - | Tcgen05Operation::DeallocCg2 - | Tcgen05Operation::RelinquishAllocPermitCg2 => Some( - "One full warp in each peer CTA must execute this instruction; lanes within each warp must execute uniformly with the same operands.", - ), - Tcgen05Operation::Alloc - | Tcgen05Operation::Dealloc - | Tcgen05Operation::RelinquishAllocPermit - | Tcgen05Operation::Ld16x256bX8Pure - | Tcgen05Operation::Ld16x256bPure - | Tcgen05Operation::LoadWait - | Tcgen05Operation::StoreWait => { - Some("One full warp must execute this instruction uniformly with the same operands.") - } - Tcgen05Operation::Commit - | Tcgen05Operation::CommitSharedCluster - | Tcgen05Operation::CommitMulticast - | Tcgen05Operation::ShiftDown => Some("One thread in the CTA issues this instruction."), - Tcgen05Operation::CommitCg2 - | Tcgen05Operation::CommitSharedClusterCg2 - | Tcgen05Operation::CommitMulticastCg2 - | Tcgen05Operation::ShiftDownCg2 => Some( - "One thread in the CTA pair issues this instruction; the peer CTA must be active and must not have exited.", - ), - _ => None, - } -} - -fn tcgen05_is_commit(operation: Tcgen05Operation) -> bool { - matches!( - operation, - Tcgen05Operation::Commit - | Tcgen05Operation::CommitSharedCluster - | Tcgen05Operation::CommitMulticast - | Tcgen05Operation::CommitCg2 - | Tcgen05Operation::CommitSharedClusterCg2 - | Tcgen05Operation::CommitMulticastCg2 - ) -} - -fn tcgen05_is_multicast_commit(operation: Tcgen05Operation) -> bool { - matches!( - operation, - Tcgen05Operation::CommitMulticast | Tcgen05Operation::CommitMulticastCg2 - ) -} - -fn tcgen05_is_shift(operation: Tcgen05Operation) -> bool { - matches!( - operation, - Tcgen05Operation::ShiftDown | Tcgen05Operation::ShiftDownCg2 - ) -} - -fn tcgen05_cp_ptx_suffix(member: Tcgen05CpMember) -> &'static str { - use Tcgen05CpMember::*; - match member { - M128x128bB4x16P64 => "128x128b.b8x16.b4x16_p64", - M128x128bB6x16P32 => "128x128b.b8x16.b6x16_p32", - M128x128b => "128x128b", - M128x256bB4x16P64 => "128x256b.b8x16.b4x16_p64", - M128x256bB6x16P32 => "128x256b.b8x16.b6x16_p32", - M32x128bWarpx4B4x16P64 => "32x128b.warpx4.b8x16.b4x16_p64", - M32x128bWarpx4B6x16P32 => "32x128b.warpx4.b8x16.b6x16_p32", - M32x128bWarpx4 => "32x128b.warpx4", - M4x256bB4x16P64 => "4x256b.b8x16.b4x16_p64", - M4x256bB6x16P32 => "4x256b.b8x16.b6x16_p32", - M4x256b => "4x256b", - M64x128bWarpx2Pair0123B4x16P64 => "64x128b.warpx2::01_23.b8x16.b4x16_p64", - M64x128bWarpx2Pair0123B6x16P32 => "64x128b.warpx2::01_23.b8x16.b6x16_p32", - M64x128bWarpx2Pair0123 => "64x128b.warpx2::01_23", - M64x128bWarpx2Pair0213B4x16P64 => "64x128b.warpx2::02_13.b8x16.b4x16_p64", - M64x128bWarpx2Pair0213B6x16P32 => "64x128b.warpx2::02_13.b8x16.b6x16_p32", - M64x128bWarpx2Pair0213 => "64x128b.warpx2::02_13", - } -} - -fn tcgen05_ld_shape_label(shape: Tcgen05LdShape) -> &'static str { - match shape { - Tcgen05LdShape::M16x32bx2 => "16x32bx2", - Tcgen05LdShape::M16x64b => "16x64b", - Tcgen05LdShape::M16x128b => "16x128b", - Tcgen05LdShape::M16x256b => "16x256b", - Tcgen05LdShape::M32x32b => "32x32b", - } -} - -fn tcgen05_ld_multiplicity_label(multiplicity: Tcgen05LdMultiplicity) -> &'static str { - match multiplicity { - Tcgen05LdMultiplicity::X1 => "x1", - Tcgen05LdMultiplicity::X2 => "x2", - Tcgen05LdMultiplicity::X4 => "x4", - Tcgen05LdMultiplicity::X8 => "x8", - Tcgen05LdMultiplicity::X16 => "x16", - Tcgen05LdMultiplicity::X32 => "x32", - Tcgen05LdMultiplicity::X64 => "x64", - Tcgen05LdMultiplicity::X128 => "x128", - } -} - -fn tcgen05_ld_register_count(record: &CatalogIntrinsic) -> usize { - let ld = record - .tcgen05 - .as_ref() - .and_then(|tcgen05| tcgen05.ld) - .expect("generated tcgen05 load identity"); - ld.shape.register_multiplier() * ld.multiplicity.count() -} - -fn tcgen05_st_register_count(record: &CatalogIntrinsic) -> usize { - let st = record - .tcgen05 - .as_ref() - .and_then(|tcgen05| tcgen05.st) - .expect("generated tcgen05 store identity"); - st.shape.register_multiplier() * st.multiplicity.count() -} - -fn tcgen05_mma_target_matrix_is_closed( - target: &crate::model::CatalogTargetRequirement, - backend: IntrinsicBackend, - fixed_kind: Option, -) -> bool { - let CatalogHardwareTarget::TargetMatrix { contracts } = &target.hardware else { - return false; - }; - let expected_kinds: Vec<&str> = fixed_kind.map_or_else( - || vec!["f16", "tf32", "f8f6f4", "i8"], - |kind| vec![tcgen05_mma_kind_name(kind)], - ); - if target.minimum_ptx.to_string() != "8.6" || contracts.len() != expected_kinds.len() { - return false; - } - let common = match backend { - IntrinsicBackend::LlvmNvptx => vec![ - ( - CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, - "8.6", - ), - (CatalogHardwareAlternative::FamilyTarget { sm: 100 }, "8.8"), - ( - CatalogHardwareAlternative::ExactArchitecture { sm: 101 }, - "8.6", - ), - (CatalogHardwareAlternative::FamilyTarget { sm: 101 }, "8.8"), - ( - CatalogHardwareAlternative::ExactArchitecture { sm: 103 }, - "8.8", - ), - (CatalogHardwareAlternative::FamilyTarget { sm: 103 }, "8.8"), - ( - CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, - "9.0", - ), - (CatalogHardwareAlternative::FamilyTarget { sm: 110 }, "9.0"), - ], - IntrinsicBackend::LibNvvm => vec![ - ( - CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, - "8.6", - ), - (CatalogHardwareAlternative::FamilyTarget { sm: 100 }, "8.8"), - ( - CatalogHardwareAlternative::ExactArchitecture { sm: 103 }, - "8.8", - ), - (CatalogHardwareAlternative::FamilyTarget { sm: 103 }, "8.8"), - ( - CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, - "9.0", - ), - (CatalogHardwareAlternative::FamilyTarget { sm: 110 }, "9.0"), - ], - }; - let i8 = match backend { - IntrinsicBackend::LlvmNvptx => vec![ - ( - CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, - "8.6", - ), - ( - CatalogHardwareAlternative::ExactArchitecture { sm: 101 }, - "8.6", - ), - ( - CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, - "9.0", - ), - ], - IntrinsicBackend::LibNvvm => vec![ - ( - CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, - "8.6", - ), - ( - CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, - "9.0", - ), - ], - }; - expected_kinds.into_iter().all(|kind| { - let Some(contract) = contracts.iter().find(|contract| { - contract.selectors.len() == 1 - && contract.selectors[0].name == "kind" - && contract.selectors[0].value == kind - }) else { - return false; - }; - let expected = if kind == "i8" { &i8 } else { &common }; - contract.alternatives.len() == expected.len() - && contract.alternatives.iter().zip(expected).all( - |(actual, (hardware, minimum_ptx))| { - actual.hardware == *hardware && actual.minimum_ptx.to_string() == *minimum_ptx - }, - ) - }) -} - -fn tcgen05_mma_render_contract( - record: &CatalogIntrinsic, - tcgen05: &crate::model::Tcgen05, - mma: &Tcgen05Mma, - llvm_route: &crate::model::CatalogBackendLowering, - libnvvm_route: &crate::model::CatalogBackendLowering, -) -> bool { - if tcgen05.operation != Tcgen05Operation::Mma - || tcgen05.cp.is_some() - || tcgen05.ld.is_some() - || tcgen05.st.is_some() - || tcgen05.source_contract != Tcgen05SourceContract::TablegenSelectionChangesPtx - || tcgen05.runtime_validation != RuntimeValidation::Unexecuted - || record.rust.module != "tcgen05" - || record.rust.safe - || record.rust.must_use - || record.rust.result != "()" - || record.dialect.op_type != "Tcgen05MmaOp" - || record.dialect.op_name != "nvvm.tcgen05_mma" - || !record.dialect.results.is_empty() - || record.lowering != "generated_tcgen05_mma" - || record.semantics.pure - || record.semantics.memory != "read_write" - || !record.semantics.convergent - || record.semantics.execution_scope != "thread" - || record.target.minimum_ptx.to_string() != "8.6" - || record.target.hardware != mma.llvm_target.hardware - || record.target.targets != "sm_100a|sm_101a|sm_103a|sm_110a" - || llvm_route.backend != IntrinsicBackend::LlvmNvptx - || llvm_route.mechanism != BackendLoweringMechanism::InlinePtx - || llvm_route.target != mma.llvm_target - || !tcgen05_mma_target_matrix_is_closed( - &llvm_route.target, - IntrinsicBackend::LlvmNvptx, - mma.fixed_selectors.map(|fixed| fixed.kind), - ) - || libnvvm_route.backend != IntrinsicBackend::LibNvvm - || libnvvm_route.mechanism != BackendLoweringMechanism::InlinePtx - || libnvvm_route.target != mma.libnvvm_target - || !tcgen05_mma_target_matrix_is_closed( - &libnvvm_route.target, - IntrinsicBackend::LibNvvm, - mma.fixed_selectors.map(|fixed| fixed.kind), - ) - { - return false; - } - - let mut expected_operands = vec![ - "i32".to_owned(), - if tcgen05_mma_is_tensor_a(mma.form) { - "i32".to_owned() - } else { - "i64".to_owned() - }, - "i64".to_owned(), - "i32".to_owned(), - "i1".to_owned(), - ]; - if tcgen05_mma_is_sparse(mma.form) { - expected_operands.push("i32".into()); - } - if tcgen05_mma_has_zero_col_mask(mma.form) { - expected_operands.push("i64".into()); - } - - let llvm = llvm(record); - let representative_kind = mma - .fixed_selectors - .map_or(Tcgen05MmaKind::F16, |fixed| fixed.kind); - let representative_buffer = mma.fixed_selectors.map_or(0, |fixed| fixed.b_buffer); - let representative_usage = mma - .fixed_selectors - .map_or(Tcgen05MmaBUsage::Discard, |fixed| fixed.b_usage); - let representative = if tcgen05_mma_is_ws(mma.form) { - tcgen05_mma_inline_asm( - mma.form, - representative_kind, - 1, - None, - Some(representative_buffer), - Some(representative_usage), - ) - } else { - tcgen05_mma_inline_asm( - mma.form, - representative_kind, - 1, - Some("discard"), - None, - None, - ) - }; - let instruction = representative - .0 - .find("tcgen05.mma") - .map_or("", |start| &representative.0[start..]); - let head = instruction.split_whitespace().next().unwrap_or_default(); - let mut components = head.split('.'); - let expected_mnemonic = components.next().unwrap_or_default(); - let expected_modifiers = components.map(str::to_owned).collect::>(); - let mut expected_ptx_operands = vec![crate::ptx::OperandPattern::Address]; - expected_ptx_operands.push(if tcgen05_mma_is_tensor_a(mma.form) { - crate::ptx::OperandPattern::Address - } else { - crate::ptx::OperandPattern::Register - }); - expected_ptx_operands.push(crate::ptx::OperandPattern::Register); - if tcgen05_mma_is_sparse(mma.form) { - expected_ptx_operands.push(crate::ptx::OperandPattern::Address); - } - expected_ptx_operands.extend([ - crate::ptx::OperandPattern::Register, - crate::ptx::OperandPattern::Exact { - value: "%enable_pred".into(), - }, - ]); - if tcgen05_mma_has_zero_col_mask(mma.form) { - expected_ptx_operands.push(crate::ptx::OperandPattern::Register); - } - if record.dialect.operands != expected_operands - || !llvm.results.is_empty() - || record.expected_ptx.mnemonic != expected_mnemonic - || record.expected_ptx.modifiers != expected_modifiers - || record.expected_ptx.operands != expected_ptx_operands - { - return false; - } - - match (tcgen05.adapter, mma.alias, mma.fixed_selectors) { - (Tcgen05Adapter::MmaDirectSelectors, None, None) => { - let selector_indices = match mma.selector_layout { - Tcgen05MmaSelectorLayout::Base { - kind_argument, - cta_group_argument, - collector_a_argument, - collector_a_upper_exclusive, - } => { - if collector_a_upper_exclusive - != if tcgen05_mma_is_ashift(mma.form) { - 2 - } else { - 4 - } - { - return false; - } - [kind_argument, cta_group_argument, collector_a_argument] - } - Tcgen05MmaSelectorLayout::WarpSpecialized { - kind_argument, - b_buffer_argument, - b_usage_argument, - } => [kind_argument, b_buffer_argument, b_usage_argument], - }; - let first_selector = expected_operands.len() as u8; - record.rust.arguments.len() == expected_operands.len() + 3 - && record - .rust - .arguments - .ends_with(&["u32".into(), "u32".into(), "u32".into()]) - && llvm.arguments.len() == expected_operands.len() + 3 - && selector_indices == [first_selector, first_selector + 1, first_selector + 2] - && tcgen05_mma_is_ws(mma.form) - == matches!( - mma.selector_layout, - Tcgen05MmaSelectorLayout::WarpSpecialized { .. } - ) - } - ( - Tcgen05Adapter::MmaDirectSelectors, - Some( - Tcgen05MmaAlias::E4m3 - | Tcgen05MmaAlias::E5m2 - | Tcgen05MmaAlias::E2m3 - | Tcgen05MmaAlias::E3m2 - | Tcgen05MmaAlias::E2m1, - ), - Some(fixed), - ) => { - mma.form == Tcgen05MmaForm::Shared - && fixed.kind == Tcgen05MmaKind::F8f6f4 - && fixed.b_buffer == 0 - && fixed.b_usage == Tcgen05MmaBUsage::Discard - && matches!( - mma.selector_layout, - Tcgen05MmaSelectorLayout::Base { - kind_argument: 5, - cta_group_argument: 6, - collector_a_argument: 7, - collector_a_upper_exclusive: 4, - } - ) - && record.rust.arguments == ["u32", "u64", "u64", "u32", "bool"] - && record.dialect.operands == ["i32", "i64", "i64", "i32", "i1"] - && llvm.arguments.len() == 8 - } - ( - Tcgen05Adapter::MmaWsFixedSelectorsDropLegacyADescriptor, - Some( - Tcgen05MmaAlias::E4m3 - | Tcgen05MmaAlias::E5m2 - | Tcgen05MmaAlias::E2m3 - | Tcgen05MmaAlias::E3m2 - | Tcgen05MmaAlias::E2m1, - ), - Some(fixed), - ) => { - mma.form == Tcgen05MmaForm::WsTensor - && fixed.kind == Tcgen05MmaKind::F8f6f4 - && fixed.b_buffer == 0 - && fixed.b_usage == Tcgen05MmaBUsage::Discard - && record.rust.arguments == ["u32", "u32", "u64", "u64", "u32", "bool"] - && record.dialect.operands == ["i32", "i32", "i64", "i32", "i1"] - && llvm.arguments.len() == 8 - } - _ => false, - } -} - -fn tcgen05_render_contract(record: &CatalogIntrinsic) -> bool { - let Some(tcgen05) = &record.tcgen05 else { - return false; - }; - let [llvm_route, libnvvm_route] = record.backend_lowerings.as_slice() else { - return false; - }; - if let Some(mma) = &tcgen05.mma { - return tcgen05_mma_render_contract(record, tcgen05, mma, llvm_route, libnvvm_route); - } - let llvm_hardware = CatalogHardwareTarget::AnyOf { - alternatives: vec![ - CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 101 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 103 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, - ], - }; - let libnvvm_hardware = CatalogHardwareTarget::AnyOf { - alternatives: vec![ - CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 103 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, - ], - }; - if record.rust.module != "tcgen05" - || record.rust.must_use != (tcgen05.operation == Tcgen05Operation::Ld) - || record.lowering != "generated_tcgen05" - || !record.semantics.convergent - || record.semantics.execution_scope != tcgen05.operation.execution_scope() - || record.target.minimum_ptx.to_string() != "8.6" - || record.target.targets != "sm_100a|sm_101a|sm_103a|sm_110a" - || llvm_route.backend != IntrinsicBackend::LlvmNvptx - || llvm_route.mechanism != BackendLoweringMechanism::InlinePtx - || llvm_route.target.minimum_ptx.to_string() != "8.6" - || llvm_route.target.hardware != llvm_hardware - || libnvvm_route.backend != IntrinsicBackend::LibNvvm - || libnvvm_route.mechanism != BackendLoweringMechanism::InlinePtx - || libnvvm_route.target.minimum_ptx.to_string() != "8.6" - || libnvvm_route.target.hardware != libnvvm_hardware - || tcgen05.runtime_validation != RuntimeValidation::Unexecuted - { - return false; - } - match tcgen05.cp { - Some(cp) => { - let expected_operation = match cp.group { - Tcgen05CpGroup::Cg1 => Tcgen05Operation::CpSmemToTmem, - Tcgen05CpGroup::Cg2 => Tcgen05Operation::CpSmemToTmemCg2, - }; - let group = match cp.group { - Tcgen05CpGroup::Cg1 => "cta_group::1", - Tcgen05CpGroup::Cg2 => "cta_group::2", - }; - let modifiers = std::iter::once("cp") - .chain(std::iter::once(group)) - .chain(tcgen05_cp_ptx_suffix(cp.member).split('.')) - .collect::>(); - if tcgen05.operation != expected_operation - || tcgen05.ld.is_some() - || tcgen05.st.is_some() - || record.semantics.pure - || record.semantics.memory != "read_write" - || record.expected_ptx.mnemonic != "tcgen05" - || record.expected_ptx.modifiers != modifiers - || record.expected_ptx.operands - != [ - crate::ptx::OperandPattern::Address, - crate::ptx::OperandPattern::Register, - ] - { - return false; - } - } - None if !matches!( - tcgen05.operation, - Tcgen05Operation::CpSmemToTmem | Tcgen05Operation::CpSmemToTmemCg2 - ) => {} - None if matches!( - record.id.as_str(), - "tcgen05_cp_smem_to_tmem" | "tcgen05_cp_smem_to_tmem_cg2" - ) => {} - None => return false, - } - match tcgen05.ld { - Some(ld) => { - let count = tcgen05_ld_register_count(record); - let mut modifiers: Vec = vec![ - "ld".into(), - "sync".into(), - "aligned".into(), - tcgen05_ld_shape_label(ld.shape).into(), - tcgen05_ld_multiplicity_label(ld.multiplicity).into(), - ]; - if ld.pack16 { - modifiers.push("pack::16b".into()); - } - modifiers.push("b32".into()); - let destination = crate::ptx::OperandPattern::RegisterList { length: count }; - let mut operands = vec![destination, crate::ptx::OperandPattern::Address]; - if ld.shape == Tcgen05LdShape::M16x32bx2 { - operands.push(crate::ptx::OperandPattern::Immediate); - } - if tcgen05.operation != Tcgen05Operation::Ld - || tcgen05.cp.is_some() - || tcgen05.st.is_some() - || record.expected_ptx.mnemonic != "tcgen05" - || record.expected_ptx.modifiers != modifiers - || record.expected_ptx.operands != operands - { - return false; - } - } - None if tcgen05.operation != Tcgen05Operation::Ld => {} - None => return false, - } - match tcgen05.st { - Some(st) => { - let count = tcgen05_st_register_count(record); - let mut modifiers: Vec = vec![ - "st".into(), - "sync".into(), - "aligned".into(), - tcgen05_ld_shape_label(st.shape).into(), - tcgen05_ld_multiplicity_label(st.multiplicity).into(), - ]; - if st.unpack16 { - modifiers.push("unpack::16b".into()); - } - modifiers.push("b32".into()); - let data = crate::ptx::OperandPattern::RegisterList { length: count }; - let mut operands = vec![crate::ptx::OperandPattern::Address]; - if st.shape == Tcgen05LdShape::M16x32bx2 { - operands.push(crate::ptx::OperandPattern::Immediate); - } - operands.push(data); - if tcgen05.operation != Tcgen05Operation::St - || tcgen05.cp.is_some() - || tcgen05.ld.is_some() - || record.expected_ptx.mnemonic != "tcgen05" - || record.expected_ptx.modifiers != modifiers - || record.expected_ptx.operands != operands - { - return false; - } - } - None if tcgen05.operation != Tcgen05Operation::St => {} - None => return false, - } - let llvm = llvm(record); - match (tcgen05.operation, tcgen05.adapter) { - ( - Tcgen05Operation::Alloc | Tcgen05Operation::AllocCg2, - Tcgen05Adapter::SharedPointerColumnsToVoid, - ) => { - !record.rust.safe - && record.rust.arguments == ["*mut u32", "u32"] - && record.rust.result == "()" - && record.dialect.operands == ["ptr", "i32"] - && record.dialect.results.is_empty() - && llvm.arguments == ["shared_ptr", "i32"] - && llvm.results.is_empty() - && tcgen05.source_contract == Tcgen05SourceContract::ExactTablegenSelection - } - ( - Tcgen05Operation::Dealloc | Tcgen05Operation::DeallocCg2, - Tcgen05Adapter::TmemAddressColumnsToVoid, - ) => { - !record.rust.safe - && record.rust.arguments == ["u32", "u32"] - && record.rust.result == "()" - && record.dialect.operands == ["i32", "i32"] - && llvm.arguments == ["tmem_ptr", "i32"] - && tcgen05.source_contract == Tcgen05SourceContract::ExactTablegenSelection - } - ( - Tcgen05Operation::RelinquishAllocPermit - | Tcgen05Operation::RelinquishAllocPermitCg2 - | Tcgen05Operation::FenceBeforeThreadSync - | Tcgen05Operation::FenceAfterThreadSync - | Tcgen05Operation::LoadWait - | Tcgen05Operation::StoreWait, - Tcgen05Adapter::NoOperands, - ) => { - record.rust.safe - && record.rust.arguments.is_empty() - && record.rust.result == "()" - && record.dialect.operands.is_empty() - && llvm.arguments.is_empty() - && tcgen05.source_contract == Tcgen05SourceContract::ExactTablegenSelection - } - ( - Tcgen05Operation::Commit - | Tcgen05Operation::CommitCg2 - | Tcgen05Operation::CommitSharedCluster - | Tcgen05Operation::CommitSharedClusterCg2, - Tcgen05Adapter::BarrierPointerToVoid, - ) => { - !record.rust.safe - && record.rust.arguments == ["*mut u64"] - && record.rust.result == "()" - && record.dialect.operands == ["ptr"] - && llvm.arguments.len() == 1 - && if matches!( - tcgen05.operation, - Tcgen05Operation::Commit | Tcgen05Operation::CommitCg2 - ) { - tcgen05.source_contract == Tcgen05SourceContract::TablegenSelectionChangesPtx - && llvm.arguments == ["ptr"] - } else { - tcgen05.source_contract == Tcgen05SourceContract::ExactTablegenSelection - && llvm.arguments == ["shared_ptr"] - } - } - ( - Tcgen05Operation::MmaWsF16 | Tcgen05Operation::MmaWsBf16 | Tcgen05Operation::MmaWsTf32, - Tcgen05Adapter::MmaWsDropLegacyADescriptor, - ) => { - !record.rust.safe - && record.rust.arguments == ["u32", "u32", "u64", "u64", "u32", "bool"] - && record.rust.result == "()" - && record.dialect.operands == ["i32", "i32", "i64", "i64", "i32", "i1"] - && llvm.arguments - == [ - "tmem_ptr", "tmem_ptr", "i64", "i32", "i1", "i32", "i32", "i32", - ] - && tcgen05.source_contract == Tcgen05SourceContract::TablegenSelectionChangesPtx - && (tcgen05.operation != Tcgen05Operation::MmaWsBf16 - || (record.expected_ptx.modifiers.contains(&"kind::f16".into()) - && !record.expected_ptx.modifiers.contains(&"kind::bf16".into()))) - } - ( - Tcgen05Operation::MmaF16 | Tcgen05Operation::MmaF16Cg2, - Tcgen05Adapter::MmaInjectZeroDisableLanes, - ) => { - !record.rust.safe - && record.rust.arguments == ["u32", "u64", "u64", "u32", "bool"] - && record.rust.result == "()" - && record.dialect.operands == ["i32", "i64", "i64", "i32", "i1"] - && llvm.arguments.len() == 8 - && tcgen05.source_contract - == Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection - } - ( - Tcgen05Operation::CpSmemToTmem | Tcgen05Operation::CpSmemToTmemCg2, - Tcgen05Adapter::TmemDescriptorToVoid, - ) => { - !record.rust.safe - && record.rust.arguments == ["u32", "u64"] - && record.rust.result == "()" - && record.dialect.operands == ["i32", "i64"] - && llvm.arguments == ["tmem_ptr", "i64"] - && tcgen05.source_contract == Tcgen05SourceContract::ExactTablegenSelection - } - (Tcgen05Operation::Ld16x256bX8Pure, Tcgen05Adapter::TmemToF32x32) => { - !record.rust.safe - && record.rust.arguments == ["u32"] - && record.rust.result == "[f32; 32]" - && record.dialect.operands == ["i32"] - && record.dialect.results == vec!["f32"; 32] - && llvm.arguments == ["tmem_ptr", "i1"] - && llvm.results == ["v32i32"] - && tcgen05.source_contract - == Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection - } - (Tcgen05Operation::Ld16x256bPure, Tcgen05Adapter::TmemToF32x4) => { - !record.rust.safe - && record.rust.arguments == ["u32"] - && record.rust.result == "[f32; 4]" - && record.dialect.operands == ["i32"] - && record.dialect.results == ["f32", "f32", "f32", "f32"] - && llvm.arguments == ["tmem_ptr", "i1"] - && llvm.results == ["v4i32"] - && tcgen05.source_contract - == Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection - } - ( - Tcgen05Operation::CommitMulticast | Tcgen05Operation::CommitMulticastCg2, - Tcgen05Adapter::BarrierPointerMaskToVoid, - ) => { - !record.rust.safe - && record.rust.arguments == ["*mut u64", "u16"] - && record.rust.result == "()" - && record.dialect.operands == ["ptr", "i16"] - && llvm.arguments == ["shared_ptr", "i16"] - && tcgen05.source_contract == Tcgen05SourceContract::ExactTablegenSelection - } - ( - Tcgen05Operation::ShiftDown | Tcgen05Operation::ShiftDownCg2, - Tcgen05Adapter::TmemAddressToVoid, - ) => { - !record.rust.safe - && record.rust.arguments == ["u32"] - && record.rust.result == "()" - && record.dialect.operands == ["i32"] - && record.dialect.results.is_empty() - && llvm.arguments == ["tmem_ptr"] - && llvm.results.is_empty() - && tcgen05.source_contract == Tcgen05SourceContract::ExactTablegenSelection - } - ( - Tcgen05Operation::Ld, - Tcgen05Adapter::TmemInjectPack16ToU32Registers - | Tcgen05Adapter::TmemHalfSplitOffsetInjectPack16ToU32Registers, - ) => { - let count = tcgen05_ld_register_count(record); - let has_half_split_offset = - tcgen05.adapter == Tcgen05Adapter::TmemHalfSplitOffsetInjectPack16ToU32Registers; - let rust_result = if count == 1 { - "u32".into() - } else { - format!("[u32; {count}]") - }; - let llvm_result = if count == 1 { - "i32".into() - } else { - format!("v{count}i32") - }; - !record.rust.safe - && record.rust.arguments - == if has_half_split_offset { - vec!["u32", "i64"] - } else { - vec!["u32"] - } - && record.rust.result == rust_result - && record.dialect.operands - == if has_half_split_offset { - vec!["i32", "i64"] - } else { - vec!["i32"] - } - && record.dialect.results == vec!["i32"; count] - && llvm.arguments - == if has_half_split_offset { - vec!["tmem_ptr", "i64", "i1"] - } else { - vec!["tmem_ptr", "i1"] - } - && llvm.results == [llvm_result] - && !record.semantics.pure - && record.semantics.memory == "read" - && tcgen05.source_contract - == Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection - } - ( - Tcgen05Operation::St, - Tcgen05Adapter::TmemU32RegistersInjectUnpack16ToVoid - | Tcgen05Adapter::TmemHalfSplitOffsetU32RegistersInjectUnpack16ToVoid, - ) => { - let count = tcgen05_st_register_count(record); - let has_half_split_offset = tcgen05.adapter - == Tcgen05Adapter::TmemHalfSplitOffsetU32RegistersInjectUnpack16ToVoid; - let rust_data = if count == 1 { - "u32".into() - } else { - format!("[u32; {count}]") - }; - let llvm_data = if count == 1 { - "i32".into() - } else { - format!("v{count}i32") - }; - !record.rust.safe - && record.rust.arguments - == if has_half_split_offset { - vec!["u32", "i64", rust_data.as_str()] - } else { - vec!["u32", rust_data.as_str()] - } - && record.rust.result == "()" - && record.dialect.operands - == std::iter::once("i32") - .chain(has_half_split_offset.then_some("i64")) - .chain(std::iter::repeat_n("i32", count)) - .collect::>() - && record.dialect.results.is_empty() - && llvm.arguments - == if has_half_split_offset { - vec!["tmem_ptr", "i64", llvm_data.as_str(), "i1"] - } else { - vec!["tmem_ptr", llvm_data.as_str(), "i1"] - } - && llvm.results.is_empty() - && !record.semantics.pure - && record.semantics.memory == "write" - && tcgen05.source_contract - == Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection - } - _ => false, - } -} - -fn cp_async_copies(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "cp_async_copy") -} - -fn cp_async_controls(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "cp_async_control") -} - -fn cp_async_mbarriers(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "cp_async_mbarrier") -} - -fn mbarrier_basics(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "mbarrier_basic") -} - -fn movmatrix(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "movmatrix") -} - -fn movmatrix_template(record: &CatalogIntrinsic) -> String { - format!( - "{}.{} $0, $1;", - record.expected_ptx.mnemonic, - record.expected_ptx.modifiers.join(".") - ) -} - -fn mbarrier_extended(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "mbarrier_extended") -} - -fn sync_intrinsics(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "sync") -} - -fn threadfence_ptx_level(record: &CatalogIntrinsic) -> Option<&'static str> { - match record.id.as_str() { - "threadfence_block" => Some("cta"), - "threadfence" => Some("gl"), - "threadfence_system" => Some("sys"), - _ => None, - } -} - -fn vote_intrinsics(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "vote") -} - -fn active_masks(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "active_mask") -} - -fn elect_intrinsics(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "elect") -} - -fn warp_matches(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "warp_match") -} - -fn warp_barriers(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "warp_barrier") -} - -fn warp_shuffles(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "warp_shuffle") -} - -fn dot_product_ptx(record: &CatalogIntrinsic) -> &'static str { - let dot = record.dot_product.as_ref().expect("dot-product record"); - match (dot.operation, dot.signedness, dot.adapter) { - ( - DotProductOperation::Dp4a, - DotProductSignedness::Signed, - DotProductAdapter::DirectThreeOperands, - ) => "dp4a.s32.s32 $0, $1, $2, $3;", - ( - DotProductOperation::Dp4a, - DotProductSignedness::Unsigned, - DotProductAdapter::DirectThreeOperands, - ) => "dp4a.u32.u32 $0, $1, $2, $3;", - ( - DotProductOperation::Dp2a, - DotProductSignedness::Signed, - DotProductAdapter::InsertLowHalfFalse, - ) => "dp2a.lo.s32.s32 $0, $1, $2, $3;", - ( - DotProductOperation::Dp2a, - DotProductSignedness::Unsigned, - DotProductAdapter::InsertLowHalfFalse, - ) => "dp2a.lo.u32.u32 $0, $1, $2, $3;", - combination => panic!("unsupported generated dot-product recipe {combination:?}"), - } -} - -fn integer_minmaxes(catalog: &CatalogFile) -> impl Iterator { - catalog - .intrinsics - .iter() - .filter(|record| record.family == "integer_minmax") -} - -fn integer_minmax_ptx_mnemonic(record: &CatalogIntrinsic) -> &'static str { - let minmax = record - .integer_minmax - .as_ref() - .expect("integer-min/max record"); - match (minmax.format, minmax.operation, minmax.relu) { - (IntegerMinMaxFormat::S32, IntegerMinMaxOperation::Min, true) => "min.relu.s32", - (IntegerMinMaxFormat::S32, IntegerMinMaxOperation::Max, true) => "max.relu.s32", - (IntegerMinMaxFormat::S16x2, IntegerMinMaxOperation::Min, false) => "min.s16x2", - (IntegerMinMaxFormat::S16x2, IntegerMinMaxOperation::Max, false) => "max.s16x2", - (IntegerMinMaxFormat::U16x2, IntegerMinMaxOperation::Min, false) => "min.u16x2", - (IntegerMinMaxFormat::U16x2, IntegerMinMaxOperation::Max, false) => "max.u16x2", - (IntegerMinMaxFormat::S16x2, IntegerMinMaxOperation::Min, true) => "min.relu.s16x2", - (IntegerMinMaxFormat::S16x2, IntegerMinMaxOperation::Max, true) => "max.relu.s16x2", - // `integer_minmax_recipe` rejects those combinations before a record - // can exist. - (IntegerMinMaxFormat::S32, _, false) | (IntegerMinMaxFormat::U16x2, _, true) => { - panic!("{} is outside the closed integer-min/max recipe", record.id) - } - } -} - -fn packed_alu_ptx_mnemonic(record: &CatalogIntrinsic) -> &'static str { - let packed = record.packed_alu.as_ref().expect("packed-ALU record"); - match (packed.format, packed.operation, packed.adapter) { - (PackedAluFormat::Bf16x2, PackedAluOperation::Add, PackedAluAdapter::DirectPackedU32) => { - "add.rn.bf16x2" - } - (PackedAluFormat::Bf16x2, PackedAluOperation::Sub, PackedAluAdapter::DirectPackedU32) => { - "sub.rn.bf16x2" - } - (PackedAluFormat::Bf16x2, PackedAluOperation::Mul, PackedAluAdapter::DirectPackedU32) => { - "mul.rn.bf16x2" - } - (PackedAluFormat::Bf16x2, PackedAluOperation::Fma, PackedAluAdapter::DirectPackedU32) => { - "fma.rn.bf16x2" - } - ( - PackedAluFormat::Bf16x2, - PackedAluOperation::FmaRelu, - PackedAluAdapter::DirectPackedU32, - ) => "fma.rn.relu.bf16x2", - (PackedAluFormat::Bf16x2, PackedAluOperation::Min, PackedAluAdapter::DirectPackedU32) => { - "min.bf16x2" - } - (PackedAluFormat::Bf16x2, PackedAluOperation::Max, PackedAluAdapter::DirectPackedU32) => { - "max.bf16x2" - } - (PackedAluFormat::Bf16x2, PackedAluOperation::Neg, PackedAluAdapter::DirectPackedU32) => { - "neg.bf16x2" - } - (PackedAluFormat::Bf16x2, PackedAluOperation::Abs, PackedAluAdapter::DirectPackedU32) => { - "abs.bf16x2" - } - (PackedAluFormat::F16x2, PackedAluOperation::Add, PackedAluAdapter::DirectPackedU32) => { - "add.rn.f16x2" - } - (PackedAluFormat::F16x2, PackedAluOperation::Sub, PackedAluAdapter::DirectPackedU32) => { - "sub.rn.f16x2" - } - (PackedAluFormat::F16x2, PackedAluOperation::Mul, PackedAluAdapter::DirectPackedU32) => { - "mul.rn.f16x2" - } - (PackedAluFormat::F16x2, PackedAluOperation::Fma, PackedAluAdapter::DirectPackedU32) => { - "fma.rn.f16x2" - } - ( - PackedAluFormat::F16x2, - PackedAluOperation::FmaRelu, - PackedAluAdapter::DirectPackedU32, - ) => "fma.rn.relu.f16x2", - (PackedAluFormat::F16x2, PackedAluOperation::FmaFtz, PackedAluAdapter::DirectPackedU32) => { - "fma.rn.ftz.f16x2" - } - (PackedAluFormat::F16x2, PackedAluOperation::FmaSat, PackedAluAdapter::DirectPackedU32) => { - "fma.rn.sat.f16x2" - } - ( - PackedAluFormat::F16x2, - PackedAluOperation::FmaFtzSat, - PackedAluAdapter::DirectPackedU32, - ) => "fma.rn.ftz.sat.f16x2", - ( - PackedAluFormat::F16x2, - PackedAluOperation::FmaFtzRelu, - PackedAluAdapter::DirectPackedU32, - ) => "fma.rn.ftz.relu.f16x2", - (PackedAluFormat::F16x2, PackedAluOperation::Min, PackedAluAdapter::DirectPackedU32) => { - "min.f16x2" - } - (PackedAluFormat::F16x2, PackedAluOperation::Max, PackedAluAdapter::DirectPackedU32) => { - "max.f16x2" - } - (PackedAluFormat::F16x2, PackedAluOperation::Neg, PackedAluAdapter::DirectPackedU32) => { - "neg.f16x2" - } - (PackedAluFormat::F16x2, PackedAluOperation::Abs, PackedAluAdapter::DirectPackedU32) => { - "abs.f16x2" - } - (PackedAluFormat::F32x2, PackedAluOperation::Add, PackedAluAdapter::DirectPackedU64) => { - "add.rn.f32x2" - } - (PackedAluFormat::F32x2, PackedAluOperation::AddFtz, PackedAluAdapter::DirectPackedU64) => { - "add.rn.ftz.f32x2" - } - (PackedAluFormat::F32x2, PackedAluOperation::Sub, PackedAluAdapter::DirectPackedU64) => { - "sub.rn.f32x2" - } - (PackedAluFormat::F32x2, PackedAluOperation::SubFtz, PackedAluAdapter::DirectPackedU64) => { - "sub.rn.ftz.f32x2" - } - (PackedAluFormat::F32x2, PackedAluOperation::Mul, PackedAluAdapter::DirectPackedU64) => { - "mul.rn.f32x2" - } - (PackedAluFormat::F32x2, PackedAluOperation::MulFtz, PackedAluAdapter::DirectPackedU64) => { - "mul.rn.ftz.f32x2" - } - (PackedAluFormat::F32x2, PackedAluOperation::Fma, PackedAluAdapter::DirectPackedU64) => { - "fma.rn.f32x2" - } - (PackedAluFormat::F32x2, PackedAluOperation::FmaFtz, PackedAluAdapter::DirectPackedU64) => { - "fma.rn.ftz.f32x2" - } - // bf16x2 has no NVPTX selection pattern for the ftz and sat fma forms; - // `packed_alu_recipe` rejects those pairs before a record can exist. - combination => panic!("unsupported generated packed-ALU recipe {combination:?}"), - } -} - -fn packed_conversion_destination(record: &CatalogIntrinsic) -> &'static str { - match record - .packed_conversion - .as_ref() - .expect("packed-conversion record") - .destination_format - { - PackedConversionDestinationFormat::Bf16x2 => "bf16x2", - PackedConversionDestinationFormat::E4m3x2 => "e4m3x2", - PackedConversionDestinationFormat::E5m2x2 => "e5m2x2", - PackedConversionDestinationFormat::F16x2 => "f16x2", - } -} - -fn packed_conversion_element(record: &CatalogIntrinsic) -> &'static str { - match record - .packed_conversion - .as_ref() - .expect("packed-conversion record") - .destination_format - { - PackedConversionDestinationFormat::Bf16x2 => "bf16", - PackedConversionDestinationFormat::E4m3x2 => "e4m3", - PackedConversionDestinationFormat::E5m2x2 => "e5m2", - PackedConversionDestinationFormat::F16x2 => "f16", - } -} - -/// Whether the conversion is one of the closed, reviewed packed-conversion -/// recipes: scalar `f32` pairs narrowed to a packed format, `f16x2` narrowed to -/// packed FP8, or packed FP8 widened back to `f16x2`. -fn packed_conversion_is_closed_recipe(conversion: &crate::model::PackedConversion) -> bool { - use PackedConversionDestinationFormat as Dst; - use PackedConversionRounding as Round; - use PackedConversionSaturation as Sat; - use PackedConversionSourceFormat as Src; - - let adapter_matches = conversion.adapter - == match conversion.source_format { - Src::F32x2 => PackedConversionAdapter::ReverseHighLowOperands, - Src::E4m3x2 | Src::E5m2x2 | Src::F16x2 => PackedConversionAdapter::Identity, - }; - - adapter_matches - && matches!( - ( - conversion.source_format, - conversion.destination_format, - conversion.rounding, - conversion.saturation, - ), - (Src::F32x2, Dst::Bf16x2, Round::NearestEven, Sat::None) - | (Src::F32x2, Dst::Bf16x2, Round::NearestEven, Sat::Relu) - | (Src::F32x2, Dst::Bf16x2, Round::TowardZero, Sat::None) - | (Src::F32x2, Dst::F16x2, Round::NearestEven, Sat::None) - | (Src::F32x2, Dst::F16x2, Round::NearestEven, Sat::Relu) - | (Src::F32x2, Dst::F16x2, Round::TowardZero, Sat::None) - | (Src::F32x2, Dst::E4m3x2, Round::NearestEven, Sat::Satfinite) - | ( - Src::F32x2, - Dst::E4m3x2, - Round::NearestEven, - Sat::SatfiniteRelu - ) - | (Src::F32x2, Dst::E5m2x2, Round::NearestEven, Sat::Satfinite) - | ( - Src::F32x2, - Dst::E5m2x2, - Round::NearestEven, - Sat::SatfiniteRelu - ) - | (Src::F16x2, Dst::E4m3x2, Round::NearestEven, Sat::Satfinite) - | ( - Src::F16x2, - Dst::E4m3x2, - Round::NearestEven, - Sat::SatfiniteRelu - ) - | (Src::F16x2, Dst::E5m2x2, Round::NearestEven, Sat::Satfinite) - | ( - Src::F16x2, - Dst::E5m2x2, - Round::NearestEven, - Sat::SatfiniteRelu - ) - | (Src::E4m3x2, Dst::F16x2, Round::NearestEven, Sat::None) - | (Src::E4m3x2, Dst::F16x2, Round::NearestEven, Sat::Relu) - | (Src::E5m2x2, Dst::F16x2, Round::NearestEven, Sat::None) - | (Src::E5m2x2, Dst::F16x2, Round::NearestEven, Sat::Relu) - ) -} - -fn packed_conversion_source(record: &CatalogIntrinsic) -> PackedConversionSourceFormat { - record - .packed_conversion - .as_ref() - .expect("packed-conversion record") - .source_format -} - -/// Register width of a single packed source operand. -fn packed_conversion_source_width(record: &CatalogIntrinsic) -> u32 { - match packed_conversion_source(record) { - PackedConversionSourceFormat::F16x2 => 32, - PackedConversionSourceFormat::E4m3x2 | PackedConversionSourceFormat::E5m2x2 => 16, - PackedConversionSourceFormat::F32x2 => { - unreachable!("f32x2 conversions do not have a single packed source") - } - } -} - -/// Rust argument types, one per source operand. -fn packed_conversion_rust_arguments(record: &CatalogIntrinsic) -> Vec<&'static str> { - match packed_conversion_source(record) { - PackedConversionSourceFormat::F32x2 => vec!["f32", "f32"], - PackedConversionSourceFormat::F16x2 => vec!["u32"], - PackedConversionSourceFormat::E4m3x2 | PackedConversionSourceFormat::E5m2x2 => vec!["u16"], - } -} - -/// Dialect operand types, one per source operand. -fn packed_conversion_dialect_operands(record: &CatalogIntrinsic) -> Vec<&'static str> { - match packed_conversion_source(record) { - PackedConversionSourceFormat::F32x2 => vec!["f32", "f32"], - PackedConversionSourceFormat::F16x2 => vec!["i32"], - PackedConversionSourceFormat::E4m3x2 | PackedConversionSourceFormat::E5m2x2 => vec!["i16"], - } -} - -fn packed_conversion_result_width(record: &CatalogIntrinsic) -> u32 { - match record - .packed_conversion - .as_ref() - .expect("packed-conversion record") - .destination_format - { - PackedConversionDestinationFormat::Bf16x2 | PackedConversionDestinationFormat::F16x2 => 32, - PackedConversionDestinationFormat::E4m3x2 | PackedConversionDestinationFormat::E5m2x2 => 16, - } -} - -fn packed_conversion_rust_type(record: &CatalogIntrinsic) -> &'static str { - match packed_conversion_result_width(record) { - 16 => "u16", - 32 => "u32", - _ => unreachable!("closed packed-conversion result width"), - } -} - -fn packed_conversion_dialect_type(record: &CatalogIntrinsic) -> &'static str { - match packed_conversion_result_width(record) { - 16 => "i16", - 32 => "i32", - _ => unreachable!("closed packed-conversion result width"), - } -} - -/// Inline-asm constraint string: one result register, then one per source -/// operand. `h` is a 16-bit register, `r` a 32-bit one, and `f` an f32. -fn packed_conversion_constraint(record: &CatalogIntrinsic) -> &'static str { - match ( - packed_conversion_result_width(record), - packed_conversion_source(record), - ) { - (16, PackedConversionSourceFormat::F32x2) => "=h,f,f", - (32, PackedConversionSourceFormat::F32x2) => "=r,f,f", - (16, PackedConversionSourceFormat::F16x2) => "=h,r", - (32, PackedConversionSourceFormat::E4m3x2 | PackedConversionSourceFormat::E5m2x2) => "=r,h", - _ => unreachable!("closed packed-conversion result width and source format"), - } -} - -/// Whether the conversion lowers through a typed NVVM intrinsic call rather -/// than inline PTX. Only the scalar-`f32` pair narrowed to FP8 does; see the -/// matching predicate in `resolve.rs` for why. -fn packed_conversion_uses_typed_nvvm(record: &CatalogIntrinsic) -> bool { - let conversion = record - .packed_conversion - .as_ref() - .expect("packed-conversion record"); - conversion.source_format == PackedConversionSourceFormat::F32x2 - && matches!( - conversion.destination_format, - PackedConversionDestinationFormat::E4m3x2 | PackedConversionDestinationFormat::E5m2x2 - ) -} - -fn packed_conversion_lowering_name(record: &CatalogIntrinsic) -> &'static str { - if packed_conversion_uses_typed_nvvm(record) { - "generated_packed_conversion_backend" - } else { - "generated_packed_conversion_inline_ptx" - } -} - -fn packed_conversion_typed_llvm_name(record: &CatalogIntrinsic) -> Option { - let route = record - .backend_lowerings - .iter() - .find(|lowering| lowering.backend == IntrinsicBackend::LlvmNvptx) - .expect("packed conversion has an LLVM-NVPTX route"); - match route.mechanism { - BackendLoweringMechanism::TypedNvvm => Some(record.llvm_identifier()), - BackendLoweringMechanism::InlinePtx => None, - } -} - -fn packed_conversion_ptx_mnemonic(record: &CatalogIntrinsic) -> String { - let conversion = record - .packed_conversion - .as_ref() - .expect("packed-conversion record"); - debug_assert_eq!( - conversion.adapter, - match conversion.source_format { - PackedConversionSourceFormat::F32x2 => PackedConversionAdapter::ReverseHighLowOperands, - PackedConversionSourceFormat::E4m3x2 - | PackedConversionSourceFormat::E5m2x2 - | PackedConversionSourceFormat::F16x2 => PackedConversionAdapter::Identity, - } - ); - let rounding = match conversion.rounding { - PackedConversionRounding::NearestEven => "rn", - PackedConversionRounding::TowardZero => "rz", - }; - let saturation = match conversion.saturation { - PackedConversionSaturation::None => "", - PackedConversionSaturation::Relu => ".relu", - PackedConversionSaturation::Satfinite => ".satfinite", - PackedConversionSaturation::SatfiniteRelu => ".satfinite.relu", - }; - format!( - "cvt.{rounding}{saturation}.{}.{}", - packed_conversion_destination(record), - conversion.source_format.ptx_token() - ) -} - -fn llvm(record: &CatalogIntrinsic) -> &CatalogLlvm { - record.llvm.as_ref().expect("LLVM-imported intrinsic") -} - -fn source_label(record: &CatalogIntrinsic) -> String { - match &record.source { - IntrinsicSource::LlvmImported { source_record } => { - format!("LLVM `{}` from `{source_record}`", llvm(record).symbol) - } - IntrinsicSource::PtxNative { instruction } => { - format!("PTX-native `{instruction}`") - } - } -} - -fn ldmatrix_attr_variants( - record: &CatalogIntrinsic, -) -> ( - &'static str, - &'static str, - &'static str, - &'static str, - &'static str, -) { - let variant = &record.ldmatrix.as_ref().expect("ldmatrix record").variant; - let shape = match variant.shape { - LdmatrixShape::M8n8 => "LdmatrixShapeAttr::M8n8", - LdmatrixShape::M8n16 => "LdmatrixShapeAttr::M8n16", - LdmatrixShape::M16n16 => "LdmatrixShapeAttr::M16n16", - }; - let multiplicity = match variant.multiplicity { - LdmatrixMultiplicity::X1 => "LdmatrixMultiplicityAttr::X1", - LdmatrixMultiplicity::X2 => "LdmatrixMultiplicityAttr::X2", - LdmatrixMultiplicity::X4 => "LdmatrixMultiplicityAttr::X4", - }; - let layout = match variant.layout { - LdmatrixLayout::Normal => "LdmatrixLayoutAttr::Normal", - LdmatrixLayout::Transposed => "LdmatrixLayoutAttr::Transposed", - }; - let element = match variant.element { - LdmatrixElement::B16 => "LdmatrixElementAttr::B16", - LdmatrixElement::B8 => "LdmatrixElementAttr::B8", - LdmatrixElement::B8x16B4x16P64 => "LdmatrixElementAttr::B8x16B4x16P64", - LdmatrixElement::B8x16B6x16P32 => "LdmatrixElementAttr::B8x16B6x16P32", - }; - let state_space = match variant.state_space { - LdmatrixStateSpace::Shared => "LdmatrixStateSpaceAttr::Shared", - }; - (shape, multiplicity, layout, element, state_space) -} - -fn register_mma_effective_kind(record: &CatalogIntrinsic) -> RegisterMmaKind { - let mma = record.register_mma.as_ref().expect("register-MMA record"); - mma.kind.unwrap_or_else(|| { - if record - .llvm - .as_ref() - .is_some_and(|llvm| llvm.symbol.contains(".kind.f8f6f4.")) - { - RegisterMmaKind::F8f6f4 - } else { - RegisterMmaKind::Standard - } - }) -} - -fn register_mma_attr_variants( - record: &CatalogIntrinsic, -) -> ( - &'static str, - &'static str, - &'static str, - &'static str, - &'static str, - &'static str, - &'static str, - &'static str, - &'static str, -) { - let mma = record.register_mma.as_ref().expect("register-MMA record"); - let shape = match mma.shape { - RegisterMmaShape::M8n8k4 => "RegisterMmaShapeAttr::M8n8k4", - RegisterMmaShape::M8n8k16 => "RegisterMmaShapeAttr::M8n8k16", - RegisterMmaShape::M8n8k32 => "RegisterMmaShapeAttr::M8n8k32", - RegisterMmaShape::M8n8k128 => "RegisterMmaShapeAttr::M8n8k128", - RegisterMmaShape::M16n8k4 => "RegisterMmaShapeAttr::M16n8k4", - RegisterMmaShape::M16n8k8 => "RegisterMmaShapeAttr::M16n8k8", - RegisterMmaShape::M16n8k16 => "RegisterMmaShapeAttr::M16n8k16", - RegisterMmaShape::M16n8k32 => "RegisterMmaShapeAttr::M16n8k32", - RegisterMmaShape::M16n8k64 => "RegisterMmaShapeAttr::M16n8k64", - RegisterMmaShape::M16n8k128 => "RegisterMmaShapeAttr::M16n8k128", - RegisterMmaShape::M16n8k256 => "RegisterMmaShapeAttr::M16n8k256", - }; - let operation = match mma.operation { - RegisterMmaOperation::Multiply => "RegisterMmaOperationAttr::Multiply", - RegisterMmaOperation::AndPopc => "RegisterMmaOperationAttr::AndPopc", - RegisterMmaOperation::XorPopc => "RegisterMmaOperationAttr::XorPopc", - }; - let kind = match register_mma_effective_kind(record) { - RegisterMmaKind::Standard => "RegisterMmaKindAttr::Standard", - RegisterMmaKind::F8f6f4 => "RegisterMmaKindAttr::F8f6f4", - RegisterMmaKind::Mxf8f6f4 => "RegisterMmaKindAttr::Mxf8f6f4", - }; - let accumulator = match mma.accumulator { - RegisterMmaAccumulator::F16 => "RegisterMmaAccumulatorAttr::F16", - RegisterMmaAccumulator::F32 => "RegisterMmaAccumulatorAttr::F32", - RegisterMmaAccumulator::F64 => "RegisterMmaAccumulatorAttr::F64", - RegisterMmaAccumulator::S32 => "RegisterMmaAccumulatorAttr::S32", - }; - let element = |element| match element { - RegisterMmaElement::Bf16 => "RegisterMmaElementAttr::Bf16", - RegisterMmaElement::E2m1 => "RegisterMmaElementAttr::E2m1", - RegisterMmaElement::E2m3 => "RegisterMmaElementAttr::E2m3", - RegisterMmaElement::E3m2 => "RegisterMmaElementAttr::E3m2", - RegisterMmaElement::E4m3 => "RegisterMmaElementAttr::E4m3", - RegisterMmaElement::E5m2 => "RegisterMmaElementAttr::E5m2", - RegisterMmaElement::F16 => "RegisterMmaElementAttr::F16", - RegisterMmaElement::Tf32 => "RegisterMmaElementAttr::Tf32", - RegisterMmaElement::F64 => "RegisterMmaElementAttr::F64", - RegisterMmaElement::B1 => "RegisterMmaElementAttr::B1", - RegisterMmaElement::S4 => "RegisterMmaElementAttr::S4", - RegisterMmaElement::U4 => "RegisterMmaElementAttr::U4", - RegisterMmaElement::S8 => "RegisterMmaElementAttr::S8", - RegisterMmaElement::U8 => "RegisterMmaElementAttr::U8", - }; - let layout = |layout| match layout { - RegisterMmaLayout::Row => "RegisterMmaLayoutAttr::Row", - RegisterMmaLayout::Col => "RegisterMmaLayoutAttr::Col", - }; - let overflow = match mma.overflow { - RegisterMmaOverflow::NotApplicable => "RegisterMmaOverflowAttr::NotApplicable", - RegisterMmaOverflow::Wrapping => "RegisterMmaOverflowAttr::Wrapping", - RegisterMmaOverflow::Satfinite => "RegisterMmaOverflowAttr::Satfinite", - }; - ( - shape, - operation, - kind, - accumulator, - element(mma.a_element), - element(mma.b_element), - layout(mma.a_layout), - layout(mma.b_layout), - overflow, - ) -} - -fn register_mma_fragment_counts(record: &CatalogIntrinsic) -> (usize, usize, usize, usize) { - match record.register_mma.as_ref().unwrap().adapter { - RegisterMmaAdapter::C2U32A2U32B1U32ToD2U32 => (2, 2, 1, 2), - RegisterMmaAdapter::C2U32A4U32B2U32ToD2U32 => (2, 4, 2, 2), - RegisterMmaAdapter::C4F32A2U32B1U32ToD4F32 => (4, 2, 1, 4), - RegisterMmaAdapter::C4F32A4U32B2U32ToD4F32 | RegisterMmaAdapter::C4I32A4U32B2U32ToD4I32 => { - (4, 4, 2, 4) - } - RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 => (4, 4, 2, 4), - RegisterMmaAdapter::C2F64A1F64B1F64ToD2F64 | RegisterMmaAdapter::C2I32A1U32B1U32ToD2I32 => { - (2, 1, 1, 2) - } - RegisterMmaAdapter::C4I32A2U32B1U32ToD4I32 => (4, 2, 1, 4), - } -} - -fn register_mma_extra_operand_count(record: &CatalogIntrinsic) -> usize { - usize::from( - record.register_mma.as_ref().unwrap().adapter - == RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32, - ) * 6 -} - -fn expected_ptx_head(record: &CatalogIntrinsic) -> String { - let mut head = record.expected_ptx.mnemonic.clone(); - for modifier in &record.expected_ptx.modifiers { - write!(head, ".{modifier}").unwrap(); - } - head -} - -fn register_list(first: usize, count: usize) -> String { - format!( - "{{{}}}", - (first..first + count) - .map(|index| format!("${index}")) - .collect::>() - .join(", ") - ) -} - -fn register_mma_template(record: &CatalogIntrinsic) -> String { - let (c_count, a_count, b_count, d_count) = register_mma_fragment_counts(record); - let d = register_list(0, d_count); - let c = register_list(d_count, c_count); - let a = register_list(d_count + c_count, a_count); - let b = register_list(d_count + c_count + a_count, b_count); - if record.register_mma.as_ref().unwrap().adapter - == RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 - { - let scale_a = d_count + c_count + a_count + b_count; - let byte_id_a = scale_a + 1; - let thread_id_a = scale_a + 2; - let scale_b = scale_a + 3; - let byte_id_b = scale_a + 4; - let thread_id_b = scale_a + 5; - format!( - "{} {d}, {a}, {b}, {c}, ${scale_a}, {{${byte_id_a}, ${thread_id_a}}}, ${scale_b}, {{${byte_id_b}, ${thread_id_b}}};", - expected_ptx_head(record) - ) - } else { - format!("{} {d}, {a}, {b}, {c};", expected_ptx_head(record)) - } -} - -fn register_mma_constraints(record: &CatalogIntrinsic) -> String { - let mma = record.register_mma.as_ref().unwrap(); - let (c_count, a_count, b_count, d_count) = register_mma_fragment_counts(record); - let (output, c, packed) = match mma.adapter { - RegisterMmaAdapter::C2U32A2U32B1U32ToD2U32 | RegisterMmaAdapter::C2U32A4U32B2U32ToD2U32 => { - ("=r", "r", "r") - } - RegisterMmaAdapter::C4F32A2U32B1U32ToD4F32 - | RegisterMmaAdapter::C4F32A4U32B2U32ToD4F32 - | RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 => ("=f", "f", "r"), - RegisterMmaAdapter::C2F64A1F64B1F64ToD2F64 => ("=d", "d", "d"), - RegisterMmaAdapter::C2I32A1U32B1U32ToD2I32 - | RegisterMmaAdapter::C4I32A4U32B2U32ToD4I32 - | RegisterMmaAdapter::C4I32A2U32B1U32ToD4I32 => ("=r", "r", "r"), - }; - let mut constraints = std::iter::repeat_n(output, d_count) - .chain(std::iter::repeat_n(c, c_count)) - .chain(std::iter::repeat_n(packed, a_count + b_count)) - .collect::>(); - if mma.adapter == RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 { - constraints.extend(["r", "h", "h", "r", "h", "h"]); - } - constraints.join(",") -} - -fn register_mma_result_variant(record: &CatalogIntrinsic) -> &'static str { - match record.register_mma.as_ref().unwrap().adapter { - RegisterMmaAdapter::C2U32A2U32B1U32ToD2U32 | RegisterMmaAdapter::C2U32A4U32B2U32ToD2U32 => { - "GeneratedMmaResultType::I32" - } - RegisterMmaAdapter::C4F32A2U32B1U32ToD4F32 | RegisterMmaAdapter::C4F32A4U32B2U32ToD4F32 => { - "GeneratedMmaResultType::F32" - } - RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 => { - "GeneratedMmaResultType::F32" - } - RegisterMmaAdapter::C2F64A1F64B1F64ToD2F64 => "GeneratedMmaResultType::F64", - RegisterMmaAdapter::C2I32A1U32B1U32ToD2I32 - | RegisterMmaAdapter::C4I32A4U32B2U32ToD4I32 - | RegisterMmaAdapter::C4I32A2U32B1U32ToD4I32 => "GeneratedMmaResultType::I32", - } -} - -fn sparse_mma_fragment_counts(record: &CatalogIntrinsic) -> (usize, usize, usize, usize) { - match record.sparse_mma.as_ref().unwrap().adapter { - SparseMmaAdapter::C2U32A4U32B4U32MetadataU32SelectorU32ToD2U32 => (2, 4, 4, 2), - SparseMmaAdapter::C4F32A4U32B4U32MetadataU32SelectorU32ToD4F32 - | SparseMmaAdapter::C4I32A4U32B4U32MetadataU32SelectorU32ToD4I32 => (4, 4, 4, 4), - SparseMmaAdapter::C4I32A2U32B2U32MetadataU32SelectorU32ToD4I32 => (4, 2, 2, 4), - } -} - -fn sparse_mma_selector_values(record: &CatalogIntrinsic) -> &'static [u32] { - match record.sparse_mma.as_ref().unwrap().selector { - SparseMmaSelector::ImmediateZeroOrOne => &[0, 1], - SparseMmaSelector::ImmediateZero => &[0], - } -} - -fn sparse_mma_selector_description(record: &CatalogIntrinsic) -> &'static str { - match record.sparse_mma.as_ref().unwrap().selector { - SparseMmaSelector::ImmediateZeroOrOne => "the compile-time constant `0` or `1`", - SparseMmaSelector::ImmediateZero => "the compile-time constant `0`", - } -} - -fn sparse_mma_selector_error(record: &CatalogIntrinsic) -> &'static str { - match record.sparse_mma.as_ref().unwrap().selector { - SparseMmaSelector::ImmediateZeroOrOne => { - "sparse MMA selector must be the compile-time constant 0 or 1" - } - SparseMmaSelector::ImmediateZero => { - "sparse MMA selector must be the compile-time constant 0" - } - } -} - -const SPARSE_MMA_STANDARD_METADATA_RULE: &str = "Every 4-bit metadata group must encode two distinct 2-bit indices; `0x0`, `0x5`, `0xa`, and `0xf` are undefined behavior."; -const SPARSE_MMA_ORDERED_METADATA_RULE: &str = "Every 4-bit metadata group must be `0x4`, `0x8`, `0x9`, `0xc`, `0xd`, or `0xe`; any other value is undefined behavior."; - -fn sparse_mma_metadata_rule(mma: &SparseMma) -> &'static str { - match mma.metadata { - SparseMmaMetadata::Standard => SPARSE_MMA_STANDARD_METADATA_RULE, - SparseMmaMetadata::Ordered => SPARSE_MMA_ORDERED_METADATA_RULE, - } -} - -fn sparse_mma_import_adapter(record: &CatalogIntrinsic) -> &'static str { - match record.sparse_mma.as_ref().unwrap().adapter { - SparseMmaAdapter::C2U32A4U32B4U32MetadataU32SelectorU32ToD2U32 => { - "GeneratedMmaImportAdapter::C2U32A4U32B4U32ToD2U32" - } - SparseMmaAdapter::C4F32A4U32B4U32MetadataU32SelectorU32ToD4F32 => { - "GeneratedMmaImportAdapter::C4F32A4U32B4U32ToD4F32" - } - SparseMmaAdapter::C4I32A2U32B2U32MetadataU32SelectorU32ToD4I32 => { - "GeneratedMmaImportAdapter::C4I32A2U32B2U32ToD4I32" - } - SparseMmaAdapter::C4I32A4U32B4U32MetadataU32SelectorU32ToD4I32 => { - "GeneratedMmaImportAdapter::C4I32A4U32B4U32ToD4I32" - } - } -} - -fn sparse_mma_attr_variants( - record: &CatalogIntrinsic, -) -> ( - &'static str, - &'static str, - &'static str, - &'static str, - &'static str, - &'static str, - &'static str, - &'static str, - &'static str, -) { - let mma = record.sparse_mma.as_ref().expect("sparse-MMA record"); - let shape = match mma.shape { - SparseMmaShape::M16n8k32 => "SparseMmaShapeAttr::M16n8k32", - SparseMmaShape::M16n8k64 => "SparseMmaShapeAttr::M16n8k64", - SparseMmaShape::M16n8k128 => "SparseMmaShapeAttr::M16n8k128", - }; - let accumulator = match mma.accumulator { - SparseMmaAccumulator::F16 => "SparseMmaAccumulatorAttr::F16", - SparseMmaAccumulator::F32 => "SparseMmaAccumulatorAttr::F32", - SparseMmaAccumulator::S32 => "SparseMmaAccumulatorAttr::S32", - }; - let element = |element| match element { - SparseMmaElement::E2m1 => "SparseMmaElementAttr::E2m1", - SparseMmaElement::E2m3 => "SparseMmaElementAttr::E2m3", - SparseMmaElement::E3m2 => "SparseMmaElementAttr::E3m2", - SparseMmaElement::E4m3 => "SparseMmaElementAttr::E4m3", - SparseMmaElement::E5m2 => "SparseMmaElementAttr::E5m2", - SparseMmaElement::S4 => "SparseMmaElementAttr::S4", - SparseMmaElement::U4 => "SparseMmaElementAttr::U4", - SparseMmaElement::S8 => "SparseMmaElementAttr::S8", - SparseMmaElement::U8 => "SparseMmaElementAttr::U8", - }; - let layout = |layout| match layout { - SparseMmaLayout::Row => "SparseMmaLayoutAttr::Row", - SparseMmaLayout::Col => "SparseMmaLayoutAttr::Col", - }; - let overflow = match mma.overflow { - SparseMmaOverflow::NotApplicable => "SparseMmaOverflowAttr::NotApplicable", - SparseMmaOverflow::Wrapping => "SparseMmaOverflowAttr::Wrapping", - SparseMmaOverflow::Satfinite => "SparseMmaOverflowAttr::Satfinite", - }; - let metadata = match mma.metadata { - SparseMmaMetadata::Standard => "SparseMmaMetadataAttr::Standard", - SparseMmaMetadata::Ordered => "SparseMmaMetadataAttr::Ordered", - }; - let selector = match mma.selector { - SparseMmaSelector::ImmediateZeroOrOne => "SparseMmaSelectorAttr::ImmediateZeroOrOne", - SparseMmaSelector::ImmediateZero => "SparseMmaSelectorAttr::ImmediateZero", - }; - ( - shape, - accumulator, - element(mma.a_element), - element(mma.b_element), - layout(mma.a_layout), - layout(mma.b_layout), - overflow, - metadata, - selector, - ) -} - -fn sparse_mma_ptx_head(record: &CatalogIntrinsic) -> String { - let mut head = record.expected_ptx.mnemonic.clone(); - for modifier in &record.expected_ptx.modifiers { - write!(head, ".{modifier}").unwrap(); - } - head -} - -fn sparse_mma_template(record: &CatalogIntrinsic) -> String { - let (c_count, a_count, b_count, d_count) = sparse_mma_fragment_counts(record); - let d = register_list(0, d_count); - let c = register_list(d_count, c_count); - let a = register_list(d_count + c_count, a_count); - let b = register_list(d_count + c_count + a_count, b_count); - let metadata = d_count + c_count + a_count + b_count; - let selector = metadata + 1; - format!( - "{} {d}, {a}, {b}, {c}, ${metadata}, ${selector};", - sparse_mma_ptx_head(record) - ) -} - -fn sparse_mma_constraints(record: &CatalogIntrinsic) -> String { - let mma = record.sparse_mma.as_ref().unwrap(); - let (c_count, a_count, b_count, d_count) = sparse_mma_fragment_counts(record); - let (output, accumulator) = match mma.accumulator { - SparseMmaAccumulator::F16 => ("=r", "r"), - SparseMmaAccumulator::F32 => ("=f", "f"), - SparseMmaAccumulator::S32 => ("=r", "r"), - }; - std::iter::repeat_n(output, d_count) - .chain(std::iter::repeat_n(accumulator, c_count)) - .chain(std::iter::repeat_n("r", a_count + b_count + 1)) - .chain(std::iter::once("n")) - .collect::>() - .join(",") -} - -fn sparse_mma_result_variant(record: &CatalogIntrinsic) -> &'static str { - match record.sparse_mma.as_ref().unwrap().accumulator { - SparseMmaAccumulator::F16 => "GeneratedMmaResultType::I32", - SparseMmaAccumulator::F32 => "GeneratedMmaResultType::F32", - SparseMmaAccumulator::S32 => "GeneratedMmaResultType::I32", - } -} - -fn sparse_mma_carriers(record: &CatalogIntrinsic) -> (String, String) { - let mma = record.sparse_mma.as_ref().unwrap(); - let (c_count, a_count, b_count, d_count) = sparse_mma_fragment_counts(record); - let accumulator = match mma.accumulator { - SparseMmaAccumulator::F16 => "MmaCarrier::U32", - SparseMmaAccumulator::F32 => "MmaCarrier::F32", - SparseMmaAccumulator::S32 => "MmaCarrier::I32", - }; - let carrier_slice = |carriers: Vec<&str>| format!("&[{}]", carriers.join(", ")); - let operands = carrier_slice( - std::iter::repeat_n(accumulator, c_count) - .chain(std::iter::repeat_n( - "MmaCarrier::U32", - a_count + b_count + 2, - )) - .collect(), - ); - let results = carrier_slice(std::iter::repeat_n(accumulator, d_count).collect()); - (operands, results) -} - -fn intrinsic_marker(catalog: &CatalogFile, record: &CatalogIntrinsic) -> String { - format!("v{}:{}", catalog.intrinsic_abi, record.rust.abi_id) -} - -fn hardware_target_label(target: &CatalogHardwareTarget) -> String { - match target { - CatalogHardwareTarget::All => "all".to_owned(), - CatalogHardwareTarget::AnyOf { alternatives } => alternatives - .iter() - .map(hardware_alternative_label) - .collect::>() - .join(" or "), - CatalogHardwareTarget::TargetMatrix { contracts } => contracts - .iter() - .map(|contract| { - let selectors = contract - .selectors - .iter() - .map(|selector| format!("{}={}", selector.name, selector.value)) - .collect::>() - .join(", "); - let alternatives = contract - .alternatives - .iter() - .map(|alternative| { - format!( - "{} at PTX {}", - hardware_alternative_label(&alternative.hardware), - alternative.minimum_ptx - ) - }) - .collect::>() - .join(" or "); - if selectors.is_empty() { - alternatives - } else { - format!("{alternatives} for {selectors}") - } - }) - .collect::>() - .join("; "), - } -} - -fn hardware_alternative_label(alternative: &CatalogHardwareAlternative) -> String { - match alternative { - CatalogHardwareAlternative::MinimumSm { sm } => format!("sm_{sm}+"), - CatalogHardwareAlternative::ExactArchitecture { sm } => format!("sm_{sm}a"), - CatalogHardwareAlternative::FamilyTarget { sm } => format!("sm_{sm}f"), - } -} - -fn generated_hardware_target(target: &CatalogHardwareTarget) -> String { - match target { - CatalogHardwareTarget::All => "GeneratedHardwareTarget::All".to_owned(), - CatalogHardwareTarget::AnyOf { alternatives } => { - let alternatives = alternatives - .iter() - .map(generated_hardware_alternative) - .collect::>() - .join(", "); - format!("GeneratedHardwareTarget::AnyOf(&[{alternatives}])") - } - CatalogHardwareTarget::TargetMatrix { contracts } => { - let contracts = contracts - .iter() - .map(generated_target_contract) - .collect::>() - .join(", "); - format!("GeneratedHardwareTarget::TargetMatrix {{ contracts: &[{contracts}] }}") - } - } -} - -fn generated_target_contract(contract: &CatalogTargetContract) -> String { - let selectors = contract - .selectors - .iter() - .map(|selector| { - format!( - "GeneratedTargetSelectorBinding {{ name: {:?}, value: {:?} }}", - selector.name, selector.value - ) - }) - .collect::>() - .join(", "); - let alternatives = contract - .alternatives - .iter() - .map(|alternative| { - format!( - "GeneratedTargetAlternative {{ minimum_ptx: GeneratedPtxVersion::from_encoded({}), hardware: {} }}", - alternative.minimum_ptx.encoded(), - generated_hardware_alternative(&alternative.hardware) - ) - }) - .collect::>() - .join(", "); - format!( - "GeneratedTargetContract {{ selectors: &[{selectors}], alternatives: &[{alternatives}] }}" - ) -} - -fn generated_hardware_alternative(alternative: &CatalogHardwareAlternative) -> String { - match alternative { - CatalogHardwareAlternative::MinimumSm { sm } => { - format!("GeneratedHardwareAlternative::MinimumSm({sm})") - } - CatalogHardwareAlternative::ExactArchitecture { sm } => { - format!("GeneratedHardwareAlternative::ExactArchitecture({sm})") - } - CatalogHardwareAlternative::FamilyTarget { sm } => { - format!("GeneratedHardwareAlternative::FamilyTarget({sm})") - } - } -} - -fn backend_label(backend: IntrinsicBackend) -> &'static str { - match backend { - IntrinsicBackend::LlvmNvptx => "llvm_nvptx", - IntrinsicBackend::LibNvvm => "lib_nvvm", - } -} - -fn lowering_mechanism_label(mechanism: BackendLoweringMechanism) -> &'static str { - match mechanism { - BackendLoweringMechanism::TypedNvvm => "typed_nvvm", - BackendLoweringMechanism::InlinePtx => "inline_ptx", - } -} - -fn evidence_stage_label(stage: EvidenceStageKind) -> &'static str { - match stage { - EvidenceStageKind::DeclarationCanonicalization => "declaration canonicalization", - EvidenceStageKind::BackendCodegen => "backend codegen", - EvidenceStageKind::DeviceLink => "device link", - EvidenceStageKind::PtxAssembly => "PTX assembly", - EvidenceStageKind::Runtime => "runtime", - } -} - -fn render_raw_mod(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - let abi_module = format!("__cuda_oxide_intrinsic_abi_v{}", catalog.intrinsic_abi); - writeln!( - output, - "#[doc(hidden)]\n#[path = \"abi_v{}.rs\"]\npub mod {abi_module};\n", - catalog.intrinsic_abi - ) - .unwrap(); - for module in modules(catalog) { - writeln!( - output, - "/// Generated `{module}` intrinsic source API.\npub mod {module} {{" - ) - .unwrap(); - for record in catalog - .intrinsics - .iter() - .filter(|record| record.rust.module == module) - { - writeln!( - output, - " pub use crate::{abi_module}::{} as {};", - record.rust.abi_id, record.rust.name - ) - .unwrap(); - } - output.push_str("}\n\n"); - } - output.push_str("#[cfg(test)]\n#[allow(clippy::type_complexity)]\nmod tests {\n"); - for record in &catalog.intrinsics { - let arguments = record.rust.arguments.join(", "); - writeln!( - output, - " #[test]\n fn public_{}_reexports_abi_{}() {{", - record.rust.name, record.rust.abi_id - ) - .unwrap(); - writeln!( - output, - " let public: {}fn({}) -> {} = super::{}::{};", - if record.rust.safe { "" } else { "unsafe " }, - arguments, - record.rust.result, - record.rust.module, - record.rust.name - ) - .unwrap(); - writeln!( - output, - " let canonical: {}fn({}) -> {} = super::{abi_module}::{};", - if record.rust.safe { "" } else { "unsafe " }, - arguments, - record.rust.result, - record.rust.abi_id - ) - .unwrap(); - output.push_str(" assert_eq!(public as usize, canonical as usize);\n }\n"); - } - output.push_str("}\n"); - output -} - -fn render_raw_abi(catalog: &CatalogFile, hash: &str) -> Result { - let mut output = rust_header(catalog, hash); - output.push_str("//! Raw ABI functions recognized by cuda-oxide.\n\n"); - for record in &catalog.intrinsics { - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!(output, "///").unwrap(); - writeln!( - output, - "/// Catalog ID: `{}`. Source: {}; expects PTX `{}`.", - record.id, - source_label(record), - record.expected_ptx - ) - .unwrap(); - if is_blackwell_ldmatrix(record) { - writeln!( - output, - "/// Available on `{}` targets. Instruction floor PTX {}; the selected target may require a newer PTX version.", - hardware_target_label(&record.target.hardware), - record.target.minimum_ptx - ) - .unwrap(); - } else { - writeln!( - output, - "/// Available on `{}` targets from PTX {}.", - hardware_target_label(&record.target.hardware), - record.target.minimum_ptx - ) - .unwrap(); - } - if let Some(operation) = record.tcgen05.as_ref().map(|tcgen05| tcgen05.operation) - && let Some(participation) = tcgen05_participation_doc(operation) - { - writeln!(output, "/// {participation}").unwrap(); - } - if record.tcgen05.is_some() { - output.push_str( - "/// All tcgen05 operations in the kernel must use the same CTA-group mode.\n", - ); - } - if let Some(reason) = &record.rust.safe_allowlist_reason { - writeln!(output, "/// Safe because {reason}").unwrap(); - } - if !record.rust.safe { - output.push_str("///\n/// # Safety\n"); - if let Some(ldmatrix) = &record.ldmatrix { - let variant = &ldmatrix.variant; - let contributing_lanes = variant.multiplicity.register_count() * 8; - let readable_bytes = if variant.shape == LdmatrixShape::M16n16 { - 32 - } else { - 16 - }; - let participation = match ldmatrix.safety.participation { - LdmatrixParticipation::AllWarpLanesSameInstruction => { - "All 32 warp lanes must execute the same instruction." - } - LdmatrixParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes => { - "All 32 warp lanes must execute the same instruction and qualifiers; no lane may have exited." - } - }; - writeln!( - output, - "/// {participation} PTX maps {contributing_lanes} lane-provided row addresses for this {:?} variant; addresses may alias, but each used address must be 16-byte aligned and have {readable_bytes} readable shared-memory bytes.", - variant.multiplicity - ) - .unwrap(); - if matches!( - ldmatrix.safety.address_contract, - crate::model::LdmatrixAddressContract::WarpLaneAddressesMappedByMultiplicitySixteenByteAlignedSixteenBytesReadableWithSm75Replication - ) && contributing_lanes < 32 - { - output.push_str("/// For portable sm_75 behavior, otherwise-unused lanes must also carry valid addresses replicated from the contributing lanes.\n"); - } - output.push_str( - "/// This weak memory operation does not replace a required barrier or fence.\n", - ); - } else if let Some((multiplicity, _)) = stmatrix_variant(record) { - let address_lanes = multiplicity.register_count() * 8; - writeln!( - output, - "/// All 32 warp lanes must execute the same instruction. The first {address_lanes} lanes must provide valid, 16-byte-aligned shared-memory row addresses." - ) - .unwrap(); - output.push_str( - "/// Register operands must contain the packed b16 fragments required by the selected layout.\n/// This weak memory operation does not replace a required barrier or fence.\n", - ); - } else if let Some(mma) = &record.sparse_mma { - let (c_count, a_count, b_count, _) = sparse_mma_fragment_counts(record); - output.push_str( - "/// All 32 warp lanes must execute the same sparse MMA instruction with the same qualifiers, and no lane may have exited.\n", - ); - writeln!( - output, - "/// `_arg0`, `_arg1`, and `_arg2` hold this lane's {c_count}-register C, {a_count}-register A, and {b_count}-register B fragments; `_arg3` holds its sparse metadata." - ) - .unwrap(); - writeln!( - output, - "/// `_arg4` must be {}.", - sparse_mma_selector_description(record) - ) - .unwrap(); - output.push_str("/// The operation is register-only and is not a memory fence.\n"); - writeln!(output, "/// {}", sparse_mma_metadata_rule(mma)).unwrap(); - writeln!( - output, - "/// See the [PTX sparse MMA fragment layouts]({}).", - record.target.ptx_isa_url - ) - .unwrap(); - match mma.overflow { - SparseMmaOverflow::NotApplicable => {} - SparseMmaOverflow::Wrapping => { - output.push_str( - "/// Signed accumulator overflow wraps because this form omits `.satfinite`.\n", - ); - } - SparseMmaOverflow::Satfinite => output.push_str( - "/// Signed accumulator overflow clamps to the finite `i32` range.\n", - ), - } - } else if let Some(mma) = &record.register_mma { - output.push_str( - "/// All 32 warp lanes must execute the same `mma.sync.aligned` instruction with the same qualifiers, and no lane may have exited.\n\ - /// `_arg0`, `_arg1`, and `_arg2` must contain this lane's C, A, and B fragments in the documented PTX layout.\n\ - /// The operation is register-only and is not a memory fence.\n", - ); - if mma.adapter == RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 - { - output.push_str( - "/// `_arg3` and `_arg6` contain this lane's packed A and B scale data. `_arg4`/`_arg5` and `_arg7`/`_arg8` are the corresponding byte/thread selectors.\n\ - /// For `scale_vec::1X`, A and B byte selectors must be in `0..=3`, the A thread selector in `0..=1`, and the B thread selector in `0..=3`; other values make the PTX operation undefined.\n", - ); - } - writeln!( - output, - "/// See the [PTX MMA fragment layouts]({}).", - record.target.ptx_isa_url - ) - .unwrap(); - if mma.overflow == RegisterMmaOverflow::Wrapping { - output.push_str("/// Signed accumulator overflow wraps because this form omits `.satfinite`.\n"); - } else if mma.overflow == RegisterMmaOverflow::Satfinite { - output.push_str( - "/// Signed accumulator overflow clamps to the finite `i32` range.\n", - ); - } - } else if record.movmatrix.is_some() { - output.push_str( - "/// All 32 warp lanes must execute the same instruction, and no lane may have exited.\n\ - /// `_arg0` must contain this lane's two packed b16 input elements.\n\ - /// The operation is register-only and is not a memory fence.\n", - ); - } else if record.redux.is_some() { - output.push_str( - "/// The executing lane must be named in `mask`. Every non-exited lane named in `mask` must execute the same `redux.sync` operation with the same qualifiers and mask.\n\ - /// The instruction waits for those lanes; violating this participation contract makes the PTX operation undefined.\n", - ); - } else if record.warp_barrier.is_some() { - output.push_str( - "/// The executing lane must be named in `mask`. Every non-exited lane named in `mask` must execute the same `bar.warp.sync` operation with the same mask.\n\ - /// On `sm_6x` and earlier, all lanes named in `mask` must execute the barrier in convergence, and no lane outside `mask` may be active when it executes.\n\ - /// The barrier orders memory accesses among participating lanes; violating the participation contract makes the PTX operation undefined.\n", - ); - } else if record.warp_shuffle.is_some() { - output.push_str( - "/// The executing lane must be named in `mask`. Every non-exited lane named in `mask` must execute the same `shfl.sync` operation with the same qualifiers and mask.\n\ - /// On `sm_6x` and earlier, all lanes named in `mask` must execute in convergence, and no lane outside `mask` may be active.\n\ - /// If the computed source lane is in range, it must be active and named in `mask`; otherwise the result is undefined. If PTX marks the computed source out of range, the calling lane's input is copied.\n", - ); - if record - .warp_shuffle - .as_ref() - .is_some_and(|shuffle| shuffle.value_kind == WarpShuffleValueKind::I64) - { - output.push_str( - "/// The 64-bit value is moved by two `b32` shuffles in one convergent block.\n", - ); - } - } else if record.warp_match.is_some() { - output.push_str( - "/// The executing lane must be named in `mask`. Every non-exited lane named in `mask` must execute the same `match.sync` operation with the same qualifiers and mask.\n\ - /// Violating this participation contract makes the PTX operation undefined.\n", - ); - } else if record.family == "elect" { - output.push_str( - "/// The executing lane must be named in `mask`. Every non-exited lane named in `mask` must execute the same `elect.sync` operation with the same mask.\n\ - /// Violating this participation contract makes the PTX operation undefined.\n", - ); - } else if record.vote.is_some() { - output.push_str( - "/// The executing lane must be named in `mask`. Every non-exited lane named in `mask` must execute the same `vote.sync` operation with the same qualifiers and mask.\n\ - /// On `sm_6x` and earlier, all lanes named in `mask` must execute in convergence, and no lane outside `mask` may be active.\n\ - /// Violating this participation contract makes the PTX operation undefined.\n", - ); - } else if record.family == "sync" { - output.push_str( - "/// Every active thread in the CTA must reach the same barrier. Calling it from divergent control flow can deadlock the CTA.\n", - ); - } else if record.family == "cluster_barrier" { - let barrier = record - .cluster_barrier - .as_ref() - .expect("cluster-barrier contract"); - output.push_str( - "/// Each non-exited cluster thread must arrive exactly once before the barrier completes, then execute the matching wait.\n", - ); - if barrier.aligned { - output.push_str( - "/// Every non-exited thread in the warp must execute this aligned instruction with identical control flow.\n", - ); - } - if barrier.ordering == ClusterBarrierOrdering::Relaxed { - output.push_str( - "/// This relaxed arrival does not publish earlier memory accesses; add the required cluster-scope fence before it.\n", - ); - } - } else if let Some(cluster) = &record.cluster_memory { - match cluster.operation { - ClusterMemoryOperation::MapSharedRank => output.push_str( - "/// `_arg0` must point into CTA shared memory, and `_arg1` must name a rank in the same live cluster.\n\ - /// The result is a cluster-shared pointer in address space 7. Dereferencing it performs a remote DSMEM access; ordinary loads and stores compile to `ld.shared::cluster` and `st.shared::cluster`. The target CTA must remain live and synchronization must make the access race-free.\n", - ), - ClusterMemoryOperation::ReadU32 => output.push_str( - "/// `_arg0` must point to an aligned readable `u32` in CTA shared memory, and `_arg1` must name a rank in the same live cluster.\n\ - /// The target CTA must publish the value with the required cluster synchronization before this weak load.\n", - ), - } - } else if let Some(control) = &record.wgmma_control { - output.push_str( - "/// All four warps in the warpgroup must execute this instruction in convergence with the same control flow.\n", - ); - match control.mode { - WgmmaControlMode::Fence => output.push_str( - "/// Issue it after register writes and before the WGMMA operations that consume those registers.\n", - ), - WgmmaControlMode::CommitGroup => output.push_str( - "/// Issue it only after the warpgroup has submitted the WGMMA operations to commit.\n", - ), - WgmmaControlMode::WaitGroup => output.push_str( - "/// `_arg0` must be a compile-time constant shared by the whole warpgroup.\n", - ), - } - } else if let Some(tma) = &record.tma { - match tma.adapter { - TmaAdapter::G2sPointersCoordinatesBarrierInjectDefaults - | TmaAdapter::G2sPointersCoordinatesBarrierMaskInjectDefaults => output - .push_str( - "/// The destination and barrier must be valid shared-memory objects, and the tensor map must be a live descriptor for this dimensionality.\n\ - /// Keep every object alive until the asynchronous copy completes. Only the designated issuing thread may start this transfer.\n", - ), - TmaAdapter::S2gPointersCoordinatesInjectDefaults => output.push_str( - "/// The source must name a live shared-memory tile, and the tensor map must be a live descriptor for this dimensionality.\n\ - /// Keep both objects alive until the committed bulk-copy group completes.\n", - ), - TmaAdapter::ReductionPointersCoordinatesInjectDefaults => output.push_str( - "/// The source must name a live shared-memory tile, and the tensor map must describe a compatible global tensor destination for this dimensionality.\n\ - /// Keep both objects alive until the committed asynchronous reduction completes.\n", - ), - TmaAdapter::DescriptorPointer - | TmaAdapter::DescriptorCoordinatesInjectDefaults - | TmaAdapter::DescriptorCoordinatesCacheHintInjectFlag => output.push_str( - "/// The tensor-map pointer must name a live, correctly encoded descriptor for this operation.\n", - ), - TmaAdapter::DescriptorAndAddressPointers - | TmaAdapter::DescriptorOrdinalAndU32 - | TmaAdapter::DescriptorOrdinalAndU64 - | TmaAdapter::DescriptorAndImmediateU32 - | TmaAdapter::DescriptorAndRuntimeU32 => output.push_str( - "/// The tensor-map pointer must name a writable, 128-byte descriptor in global memory, and every replacement value must satisfy the PTX field contract.\n", - ), - TmaAdapter::DescriptorPointerInjectBytes => output.push_str( - "/// The tensor-map pointer must name a live 128-byte descriptor covered by the matching generic-proxy release fence.\n", - ), - TmaAdapter::NoOperands | TmaAdapter::CompileTimeConstantMaxPending => {} - } - } else if let Some(operation) = ExecutionControlOperation::from_catalog_id(&record.id) { - match operation { - ExecutionControlOperation::BarrierCtaSync - | ExecutionControlOperation::BarrierCtaSyncAligned - | ExecutionControlOperation::BarrierCtaArrive - | ExecutionControlOperation::BarrierCtaArriveAligned => output.push_str( - "/// `_arg0` must identify a CTA barrier and `_arg1` must be the compatible expected thread count used by every participant.\n", - ), - ExecutionControlOperation::GridDependencyLaunchDependents - | ExecutionControlOperation::GridDependencyWait => output.push_str( - "/// The kernel launch must participate in a valid programmatic dependent-launch protocol.\n", - ), - ExecutionControlOperation::SetMaxNRegInc - | ExecutionControlOperation::SetMaxNRegDec => output.push_str( - "/// `_arg0` must be a compile-time multiple of eight in `24..=256`, and every thread in the warpgroup must execute the same operation and count.\n", - ), - } - } else if let Some(bridge) = &record.cp_async_mbarrier { - output.push_str( - "/// `_arg0` must point to a live, initialized, eight-byte-aligned mbarrier object in shared memory.\n\ - /// The issuing thread must have prior `cp.async` operations, and the object must remain valid until they complete.\n", - ); - match bridge.operation { - CpAsyncMbarrierOperation::Arrive => output.push_str( - "/// This instruction increments the pending count before scheduling the asynchronous arrival; that increment must not exceed the barrier's pending-count limit.\n", - ), - CpAsyncMbarrierOperation::ArriveNoInc => output.push_str( - "/// The barrier's initial pending count must already include the asynchronous arrival because this form does not increment it.\n", - ), - } - } else if let Some(mbarrier) = &record.mbarrier_basic { - output.push_str( - "/// `_arg0` must point to a live, eight-byte-aligned mbarrier object in shared memory.\n", - ); - match mbarrier.operation { - MbarrierBasicOperation::Init => output.push_str( - "/// Exactly one thread may initialize the object. `_arg1` must be in `1..=0xFFFFF` and include every arrival in the phase.\n\ - /// The object must be uninitialized or invalidated, with no concurrent barrier operation.\n", - ), - MbarrierBasicOperation::Arrive => output.push_str( - "/// The object must be initialized, and this arrival must be included in the current phase's expected count.\n\ - /// Use the returned token only with the same object and phase.\n", - ), - MbarrierBasicOperation::TestWait => output.push_str( - "/// The object must be initialized. `_arg1` must be a token returned for the same object and phase.\n", - ), - MbarrierBasicOperation::Inval => output.push_str( - "/// The object must be initialized and no thread or asynchronous operation may still use it.\n\ - /// Exactly one thread may invalidate the object.\n", - ), - } - } else if let Some(mbarrier) = &record.mbarrier_extended { - output.push_str( - "/// This is a convergent, side-effecting operation with a compiler memory clobber.\n", - ); - match mbarrier.adapter { - MbarrierExtendedAdapter::PointerTxCountBytesToTokenDroppingTxCount => { - output.push_str( - "/// `_arg0` must point to a live, initialized, eight-byte-aligned mbarrier in shared memory. `_arg2` must match the asynchronous transaction bytes.\n\ - /// `_arg1` is retained for compatibility and is not a PTX operand.\n", - ); - } - MbarrierExtendedAdapter::RawClusterAddressToVoid => output.push_str( - "/// `_arg0` must be a live remote shared-memory barrier address obtained through cluster address mapping.\n", - ), - MbarrierExtendedAdapter::PointerTokenToPredicate => output.push_str( - "/// `_arg0` must point to a live initialized barrier. `_arg1` must be a token for that barrier and phase.\n", - ), - MbarrierExtendedAdapter::PointerParityToPredicate => output.push_str( - "/// `_arg0` must point to a live initialized barrier. `_arg1` must be the expected phase parity.\n", - ), - MbarrierExtendedAdapter::ZeroOperandsToVoid => output.push_str( - "/// The call site must satisfy the scope and proxy-ordering contract named by the operation.\n", - ), - MbarrierExtendedAdapter::NanosecondsToVoid => output.push_str( - "/// `_arg0` is an approximate suspension duration; hardware may resume the thread earlier.\n", - ), - } - } else if let Some(copy) = &record.cp_async_copy { - let bytes = copy.copy_size.bytes(); - writeln!( - output, - "/// `_arg0` must point to {bytes} writable bytes in shared memory and be aligned to {bytes} bytes." - ) - .unwrap(); - if copy.source_size == CpAsyncSourceSize::Runtime { - writeln!( - output, - "/// `_arg2` must be at most {bytes}; `_arg1` must point to that many readable bytes in global memory and be aligned to {bytes} bytes." - ) - .unwrap(); - } else { - writeln!( - output, - "/// `_arg1` must point to {bytes} readable bytes in global memory and be aligned to {bytes} bytes." - ) - .unwrap(); - } - output.push_str( - "/// Both ranges must remain valid, the source must remain unchanged, and the destination must not be accessed until this copy completes.\n\ - /// The issuing thread must use a matching `cp.async` completion operation. Synchronize threads after completion before another thread accesses the destination.\n\ - /// User-authored completion assembly must include a compiler memory clobber.\n", - ); - } else if let Some(control) = &record.cp_async_control { - if control.operation == CpAsyncControlOperation::WaitGroup { - output.push_str( - "/// `_arg0` must be a compile-time constant. Access only destinations whose copy groups this wait completes.\n", - ); - } else if control.operation == CpAsyncControlOperation::WaitAll { - output.push_str( - "/// This waits only for copies issued by the executing thread. Synchronize threads before another thread accesses a completed destination.\n", - ); - } else { - output.push_str( - "/// This commits only copies issued by the executing thread and does not wait for completion.\n", - ); - } - } else if let Some(tcgen05) = &record.tcgen05 { - if tcgen05_is_commit(tcgen05.operation) { - if tcgen05_is_multicast_commit(tcgen05.operation) { - output.push_str( - "/// `_arg0` must point to a live initialized cluster-shared mbarrier. `_arg1` must select valid CTA ranks in its cluster.\n", - ); - } else { - output.push_str( - "/// `_arg0` must point to a live initialized mbarrier valid for this CTA-group mode.\n", - ); - } - output.push_str( - "/// The same thread that issued the tracked asynchronous tcgen05 operations must issue this commit.\n", - ); - } else if tcgen05_is_shift(tcgen05.operation) { - output.push_str( - "/// `_arg0` must name a live tensor-memory allocation, and its lane component must be a multiple of 32.\n\ - /// Completion must be tracked by a matching commit from that same thread and observed through the selected mbarrier before relying on shifted data.\n", - ); - } else if tcgen05.operation == Tcgen05Operation::Ld { - output.push_str( - "/// `_arg0` must name a live tensor-memory allocation covering the selected tile.\n\ - /// All active warp lanes must execute convergently with the same address.\n\ - /// Complete the matching tensor-memory load wait before consuming the returned registers.\n", - ); - if tcgen05 - .ld - .is_some_and(|ld| ld.shape == Tcgen05LdShape::M16x32bx2) - { - output.push_str( - "/// `_arg1` must be a compile-time constant; inline PTX encodes its low 32 bits.\n", - ); - } - } else if tcgen05.operation == Tcgen05Operation::St { - output.push_str( - "/// `_arg0` must name a live tensor-memory allocation covering the selected tile.\n\ - /// All active warp lanes must execute convergently with the same address.\n\ - /// Complete the matching tensor-memory store wait before relying on completion or reusing the affected storage.\n", - ); - if tcgen05 - .st - .is_some_and(|st| st.shape == Tcgen05LdShape::M16x32bx2) - { - output.push_str( - "/// `_arg1` must be a compile-time constant; inline PTX encodes its low 32 bits.\n", - ); - } - } else { - output.push_str( - "/// The caller must satisfy the tcgen05 address, lifetime, and participation rules.\n", - ); - } - } else if record.packed_atomic.is_some() { - output.push_str( - "/// `addr` must designate four writable bytes in global memory and be naturally aligned to four bytes.\n\ - /// Do not overlap this operation with a whole-word atomic or any non-atomic access to either 16-bit lane.\n\ - /// Racing atomics must use scopes that include each other; this relaxed GPU-scope operation is not atomic with host/system access.\n", - ); - } else if let Some(clc) = &record.clc { - render_clc_safety_lines(&mut output, clc.operation, ClcSafetyArgNames::RawAbi); - } else { - anyhow::bail!( - "{} ({}) is unsafe but has no dedicated family safety renderer", - record.id, - record.family - ); - } - } - if record.rust.arguments.len() > 7 { - output.push_str("#[allow(clippy::too_many_arguments)]\n"); - } - if record.rust.must_use { - output.push_str("#[must_use]\n"); - } - output.push_str("#[inline(never)]\n"); - let safety = if record.rust.safe { "" } else { "unsafe " }; - let arguments = record - .rust - .arguments - .iter() - .enumerate() - .map(|(index, ty)| format!("_arg{index}: {ty}")) - .collect::>() - .join(", "); - writeln!( - output, - "pub {safety}fn {}({arguments}) -> {} {{", - record.rust.abi_id, record.rust.result, - ) - .unwrap(); - writeln!( - output, - " unreachable!(\"generated CUDA intrinsic `{}` executed outside device compilation\")", - record.rust.canonical_path - ) - .unwrap(); - output.push_str("}\n\n"); - } - Ok(output) -} - -fn render_compat_register_mma(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str("// Included inside `cuda_device::wmma` to keep public paths stable.\n\n"); - for record in register_mmas(catalog).filter(|record| { - record.register_mma.as_ref().is_some_and(|mma| { - mma.compatibility_source == RegisterMmaCompatibilitySource::GeneratedStub - }) - }) { - let mma = record - .register_mma - .as_ref() - .expect("register-MMA semantics"); - let [path] = record.rust.compatibility_paths.as_slice() else { - panic!("generated register-MMA API requires one compatibility path"); - }; - assert_eq!(path, &format!("cuda_device::wmma::{}", record.rust.name)); - assert!(!record.rust.safe); - assert!(record.rust.must_use); - let argument_names: &[&str] = - if mma.adapter == RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 { - &[ - "c", - "a", - "b", - "scale_a", - "byte_id_a", - "thread_id_a", - "scale_b", - "byte_id_b", - "thread_id_b", - ] - } else { - &["c", "a", "b"] - }; - assert_eq!(record.rust.arguments.len(), argument_names.len()); - - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!( - output, - "/// Lowers to `{}`. C, A, and B are this lane's fragments in PTX register order.", - expected_ptx_head(record) - ) - .unwrap(); - writeln!( - output, - "/// Requires PTX {} and `{}`.", - record.target.minimum_ptx, - hardware_target_label(&record.target.hardware) - ) - .unwrap(); - match mma.overflow { - RegisterMmaOverflow::Wrapping => { - output.push_str("/// Signed accumulator overflow wraps.\n"); - } - RegisterMmaOverflow::Satfinite => { - output.push_str( - "/// Signed accumulator overflow clamps to the finite `i32` range.\n", - ); - } - RegisterMmaOverflow::NotApplicable => {} - } - output.push_str("///\n/// # Safety\n"); - output.push_str( - "/// All 32 lanes must execute the same instruction with the same qualifiers, and no lane may have exited.\n", - ); - output.push_str( - "/// `c`, `a`, and `b` must contain this lane's fragments in the documented PTX layout.\n", - ); - if mma.adapter == RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 { - output.push_str( - "/// `scale_a` and `scale_b` contain this lane's packed scale data.\n\ - /// For `scale_vec::1X`, `byte_id_a` and `byte_id_b` must be in `0..=3`, `thread_id_a` in `0..=1`, and `thread_id_b` in `0..=3`; other values make the PTX operation undefined.\n", - ); - } - writeln!( - output, - "/// See the [PTX MMA fragment layouts]({}).", - record.target.ptx_isa_url - ) - .unwrap(); - output.push_str("/// This register-only operation is not a memory fence.\n"); - if record.rust.arguments.len() > 7 { - output.push_str("#[allow(clippy::too_many_arguments)]\n"); - } - output.push_str("#[must_use]\n#[inline(never)]\n"); - let arguments = argument_names - .iter() - .copied() - .zip(&record.rust.arguments) - .map(|(name, ty)| format!("{name}: {ty}")) - .collect::>() - .join(", "); - writeln!( - output, - "pub unsafe fn {}({arguments}) -> {} {{", - record.rust.name, record.rust.result - ) - .unwrap(); - writeln!(output, " let _ = ({});", argument_names.join(", ")).unwrap(); - writeln!( - output, - " unreachable!(\"{} called outside CUDA kernel context\")", - record.rust.name - ) - .unwrap(); - output.push_str("}\n\n"); - } - output -} - -fn render_compat_ldmatrix(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str("// Included inside `cuda_device::wmma` to keep public paths stable.\n\n"); - for record in ldmatrix(catalog).filter(|record| { - is_blackwell_ldmatrix(record) && !record.rust.compatibility_paths.is_empty() - }) { - let [path] = record.rust.compatibility_paths.as_slice() else { - panic!("generated Blackwell ldmatrix API requires one compatibility path"); - }; - assert_eq!(path, &format!("cuda_device::wmma::{}", record.rust.name)); - assert_eq!(record.rust.arguments, ["*const u8"]); - assert!(!record.rust.safe); - assert!(record.rust.must_use); - - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!(output, "/// Lowers to `{}`.", expected_ptx_head(record)).unwrap(); - writeln!( - output, - "/// Requires PTX {} and `{}`.", - record.target.minimum_ptx, - hardware_target_label(&record.target.hardware) - ) - .unwrap(); - output.push_str( - "///\n/// # Safety\n\ - /// All 32 lanes must execute the same instruction and qualifiers, with no exited lanes.\n\ - /// `smem_ptr` must satisfy the lane-address mapping, alignment, and readable-byte contract documented by PTX.\n\ - #[must_use]\n#[inline(never)]\n", - ); - writeln!( - output, - "pub unsafe fn {}(smem_ptr: *const u8) -> {} {{", - record.rust.name, record.rust.result - ) - .unwrap(); - output.push_str(" let _ = smem_ptr;\n"); - writeln!( - output, - " unreachable!(\"{} called outside CUDA kernel context\")", - record.rust.name - ) - .unwrap(); - output.push_str("}\n\n"); - } - output -} - -fn render_compat_sparse_mma(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str("// Included inside `cuda_device::wmma` to keep public paths stable.\n\n"); - for record in sparse_mmas(catalog).filter(|record| { - record.sparse_mma.as_ref().is_some_and(|mma| { - mma.compatibility_source == SparseMmaCompatibilitySource::GeneratedStub - }) - }) { - let mma = record.sparse_mma.as_ref().expect("sparse-MMA semantics"); - let [path] = record.rust.compatibility_paths.as_slice() else { - panic!("generated sparse-MMA API requires one compatibility path"); - }; - assert_eq!(path, &format!("cuda_device::wmma::{}", record.rust.name)); - assert!(matches!( - mma.adapter, - SparseMmaAdapter::C2U32A4U32B4U32MetadataU32SelectorU32ToD2U32 - | SparseMmaAdapter::C4F32A4U32B4U32MetadataU32SelectorU32ToD4F32 - | SparseMmaAdapter::C4I32A2U32B2U32MetadataU32SelectorU32ToD4I32 - | SparseMmaAdapter::C4I32A4U32B4U32MetadataU32SelectorU32ToD4I32 - )); - - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!( - output, - "/// Lowers to `{}`. C, A, B, and metadata are this lane's PTX fragments.", - sparse_mma_ptx_head(record) - ) - .unwrap(); - writeln!( - output, - "/// Requires PTX {} and `{}`.", - record.target.minimum_ptx, - hardware_target_label(&record.target.hardware) - ) - .unwrap(); - match mma.overflow { - SparseMmaOverflow::NotApplicable => {} - SparseMmaOverflow::Wrapping => { - output.push_str("/// Signed accumulator overflow wraps.\n"); - } - SparseMmaOverflow::Satfinite => output - .push_str("/// Signed accumulator overflow clamps to the finite `i32` range.\n"), - } - output.push_str("///\n/// # Safety\n"); - output.push_str( - "/// All 32 lanes must execute the same instruction with the same qualifiers, and no lane may have exited.\n", - ); - let (c_count, a_count, b_count, _) = sparse_mma_fragment_counts(record); - writeln!( - output, - "/// `c`, `a`, and `b` must contain this lane's {c_count}-, {a_count}-, and {b_count}-register fragments; `metadata` must contain its sparse metadata." - ) - .unwrap(); - writeln!(output, "/// {}", sparse_mma_metadata_rule(mma)).unwrap(); - writeln!( - output, - "/// `selector` must be {}.", - sparse_mma_selector_description(record) - ) - .unwrap(); - writeln!( - output, - "/// See the [PTX sparse MMA fragment layouts]({}).", - record.target.ptx_isa_url - ) - .unwrap(); - output.push_str("/// This register-only operation is not a memory fence.\n"); - output.push_str("#[must_use]\n#[inline(never)]\n"); - let arguments = ["c", "a", "b", "metadata", "selector"] - .into_iter() - .zip(&record.rust.arguments) - .map(|(name, ty)| format!("{name}: {ty}")) - .collect::>() - .join(", "); - writeln!( - output, - "pub unsafe fn {}({arguments}) -> {} {{", - record.rust.name, record.rust.result - ) - .unwrap(); - output.push_str(" let _ = (c, a, b, metadata, selector);\n"); - writeln!( - output, - " unreachable!(\"{} called outside CUDA kernel context\")", - record.rust.name - ) - .unwrap(); - output.push_str("}\n\n"); - } - output -} - -fn render_compat_sreg(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - "// This file is included lexically inside `cuda_device::thread` so existing\n// DefPaths remain stable during the generated-intrinsics migration.\n\n", - ); - for record in sregs(catalog) { - let Some(path) = record - .rust - .compatibility_paths - .iter() - .find(|path| path.starts_with("cuda_device::thread::")) - else { - continue; - }; - let name = path.rsplit("::").next().unwrap(); - writeln!( - output, - "/// Compatibility spelling for `{}`.", - record.rust.public_path - ) - .unwrap(); - writeln!( - output, - "/// The compiler replaces this call with {}.", - source_label(record) - ) - .unwrap(); - output.push_str("#[allow(non_snake_case)]\n#[inline(never)]\n"); - let safety = if record.rust.safe { "" } else { "unsafe " }; - writeln!( - output, - "pub {safety}fn {name}() -> {} {{", - record.rust.result - ) - .unwrap(); - writeln!( - output, - " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" - ) - .unwrap(); - output.push_str("}\n\n"); - } - output -} - -fn render_compat_cluster_sreg(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - "// This file is included inside `cuda_device::cluster`.\n// Its private leaves let public helpers compose generated reads.\n\n", - ); - for record in sregs(catalog) { - let Some(path) = record - .rust - .compatibility_paths - .iter() - .find(|path| path.starts_with("cuda_device::cluster::__")) - else { - continue; - }; - let name = path.rsplit("::").next().unwrap(); - writeln!( - output, - "/// Private compatibility spelling for `{}`.", - record.rust.public_path - ) - .unwrap(); - writeln!( - output, - "/// The compiler replaces this call with `{}`.", - llvm(record).symbol - ) - .unwrap(); - output.push_str("#[allow(non_snake_case)]\n#[inline(never)]\n"); - writeln!( - output, - "pub(crate) fn {name}() -> {} {{", - record.rust.result - ) - .unwrap(); - writeln!( - output, - " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" - ) - .unwrap(); - output.push_str("}\n\n"); - } - output -} - -fn render_compat_special_register_module( - catalog: &CatalogFile, - hash: &str, - module: &str, -) -> String { - let mut output = rust_header(catalog, hash); - writeln!( - output, - "// Included inside `cuda_device::{module}` to keep existing paths stable.\n" - ) - .unwrap(); - let prefix = format!("cuda_device::{module}::"); - for record in sregs(catalog).filter(|record| record.special_register.is_some()) { - let Some(path) = record - .rust - .compatibility_paths - .iter() - .find(|path| path.starts_with(&prefix)) - else { - continue; - }; - let name = path.rsplit("::").next().unwrap(); - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!(output, "/// Generated from {}.", source_label(record)).unwrap(); - output.push_str("#[inline(never)]\n"); - let safety = if record.rust.safe { "" } else { "unsafe " }; - writeln!( - output, - "pub {safety}fn {name}() -> {} {{", - record.rust.result - ) - .unwrap(); - writeln!( - output, - " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" - ) - .unwrap(); - output.push_str("}\n\n"); - } - output -} - -fn render_compat_fence(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - "// This file is included inside `cuda_device::fence` so existing paths stay stable.\n\n", - ); - for record in sync_intrinsics(catalog).filter(|record| threadfence_ptx_level(record).is_some()) - { - let Some(path) = record - .rust - .compatibility_paths - .iter() - .find(|path| path.starts_with("cuda_device::fence::")) - else { - continue; - }; - let name = path.rsplit("::").next().unwrap(); - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!( - output, - "/// The compiler replaces this call with `{}`.", - llvm(record).symbol - ) - .unwrap(); - output.push_str("#[inline(never)]\n"); - writeln!(output, "pub fn {name}() {{").unwrap(); - writeln!( - output, - " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" - ) - .unwrap(); - output.push_str("}\n\n"); - } - output -} - -fn render_compat_dotprod(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str("// Included inside `cuda_device::dotprod` to keep existing paths stable.\n\n"); - for record in dot_products(catalog) { - let path = record - .rust - .compatibility_paths - .iter() - .find(|path| path.starts_with("cuda_device::dotprod::")) - .expect("dot-product compatibility path"); - let arguments = record - .rust - .arguments - .iter() - .enumerate() - .map(|(index, ty)| format!("arg{index}: {ty}")) - .collect::>() - .join(", "); - let values = (0..record.rust.arguments.len()) - .map(|index| format!("arg{index}")) - .collect::>() - .join(", "); - writeln!(output, "/// {}", record.summary).unwrap(); - output.push_str("#[inline(never)]\n"); - writeln!( - output, - "pub fn {}({arguments}) -> {} {{", - record.rust.name, record.rust.result - ) - .unwrap(); - if record.rust.arguments.len() == 1 { - writeln!(output, " let _ = {values};").unwrap(); - } else { - writeln!(output, " let _ = ({values});").unwrap(); - } - writeln!( - output, - " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" - ) - .unwrap(); - output.push_str("}\n\n"); - } - output -} - -fn render_compat_prmt(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str("// Included inside `cuda_device::prmt` to keep the public paths stable.\n\n"); - for record in prmts(catalog) { - let path = record - .rust - .compatibility_paths - .iter() - .find(|path| path.starts_with("cuda_device::prmt::")) - .expect("prmt compatibility path"); - let arguments = record - .rust - .arguments - .iter() - .enumerate() - .map(|(index, ty)| format!("arg{index}: {ty}")) - .collect::>() - .join(", "); - let values = (0..record.rust.arguments.len()) - .map(|index| format!("arg{index}")) - .collect::>() - .join(", "); - writeln!(output, "/// {}", record.summary).unwrap(); - output.push_str("#[must_use]\n#[inline(never)]\n"); - writeln!(output, "pub fn {}({arguments}) -> u32 {{", record.rust.name).unwrap(); - writeln!(output, " let _ = ({values});").unwrap(); - writeln!( - output, - " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" - ) - .unwrap(); - output.push_str("}\n\n"); - } - output -} - -fn render_compat_cluster_barrier(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str("// Included inside `cuda_device::cluster` to keep public paths stable.\n\n"); - for record in cluster_barriers(catalog) { - let barrier = record - .cluster_barrier - .as_ref() - .expect("cluster-barrier contract"); - let path = record - .rust - .compatibility_paths - .iter() - .find(|path| path.starts_with("cuda_device::cluster::")) - .expect("cluster-barrier compatibility path"); - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!(output, "/// Lowers to `{}`.", record.expected_ptx).unwrap(); - output.push_str("///\n/// # Safety\n"); - output.push_str( - "/// Each non-exited cluster thread must arrive exactly once before completion, then execute the matching wait.\n", - ); - if barrier.aligned { - output.push_str( - "/// Every non-exited thread in the warp must execute this aligned instruction with identical control flow.\n", - ); - } - if barrier.ordering == ClusterBarrierOrdering::Relaxed { - output.push_str( - "/// This relaxed arrival does not publish earlier memory accesses; use the required cluster-scope fence first.\n", - ); - } - output.push_str("#[inline(never)]\n"); - writeln!(output, "pub unsafe fn {}() {{", record.rust.name).unwrap(); - writeln!(output, " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")").unwrap(); - output.push_str("}\n\n"); - } - output -} - -fn render_compat_cluster_memory(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str("// Included inside `cuda_device::cluster` to keep its public API stable.\n\n"); - for record in cluster_memory(catalog) { - let cluster = record - .cluster_memory - .as_ref() - .expect("cluster-memory contract"); - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!(output, "/// Lowers to `{}`.", record.expected_ptx).unwrap(); - output.push_str("///\n/// # Safety\n"); - match cluster.operation { - ClusterMemoryOperation::MapSharedRank => { - output.push_str( - "/// `local_ptr` must point into CTA shared memory, and `target_rank` must name a rank in the same live cluster.\n\ - /// The result is a cluster-shared pointer in address space 7. Dereferencing it performs a remote DSMEM access; the target CTA must remain live and synchronization must make the access race-free.\n", - ); - output.push_str("#[must_use]\n#[inline(never)]\n"); - output.push_str( - "pub unsafe fn map_shared_rank(local_ptr: *const T, target_rank: u32) -> *const T {\n\ - let _ = (local_ptr, target_rank);\n\ - unreachable!(\"map_shared_rank called outside CUDA kernel context\")\n\ - }\n\n", - ); - output.push_str( - "/// Maps a mutable CTA-shared address to another cluster rank.\n\ - ///\n\ - /// # Safety\n\ - /// The mapping requirements above apply. The target CTA must remain live, and remote writes must be synchronized and race-free.\n\ - #[must_use]\n#[inline(never)]\n\ - pub unsafe fn map_shared_rank_mut(local_ptr: *mut T, target_rank: u32) -> *mut T {\n\ - let _ = (local_ptr, target_rank);\n\ - unreachable!(\"map_shared_rank_mut called outside CUDA kernel context\")\n\ - }\n\n", - ); - } - ClusterMemoryOperation::ReadU32 => { - output.push_str( - "/// `local_ptr` must identify an aligned readable `u32` in CTA shared memory. `target_rank` must name a rank in the same live cluster.\n\ - /// The target CTA must publish the value before this weak cluster load.\n", - ); - output.push_str("#[must_use]\n#[inline(never)]\n"); - output.push_str( - "pub unsafe fn dsmem_read_u32(local_ptr: *const u32, target_rank: u32) -> u32 {\n\ - let _ = (local_ptr, target_rank);\n\ - unreachable!(\"dsmem_read_u32 called outside CUDA kernel context\")\n\ - }\n\n", - ); - } - } - } - output -} - -fn render_compat_movmatrix(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str("// Included inside `cuda_device::wmma` to keep its public API stable.\n\n"); - for record in movmatrix(catalog) { - let path = record - .rust - .compatibility_paths - .iter() - .find(|path| path.starts_with("cuda_device::wmma::")) - .expect("movmatrix compatibility path"); - assert_eq!(path, &format!("cuda_device::wmma::{}", record.rust.name)); - writeln!(output, "/// {}", record.summary).unwrap(); - output.push_str( - "///\n\ - /// Each lane supplies two packed b16 values. The warp collectively transposes the 8x8 tile.\n\ - /// No addressable memory is touched, but the result depends on every lane's\n\ - /// input, so the call is modelled as reading and writing inaccessible state:\n\ - /// two calls with equal operands are not interchangeable.\n\ - ///\n\ - /// # Safety\n\ - /// All 32 warp lanes must execute the same call, and no lane may have exited.\n", - ); - writeln!( - output, - "/// Requires PTX {} and `sm_75+`.", - record.target.minimum_ptx - ) - .unwrap(); - output.push_str("#[must_use]\n#[inline(never)]\n"); - writeln!( - output, - "pub unsafe fn {}(value: u32) -> u32 {{", - record.rust.name - ) - .unwrap(); - output.push_str(" let _ = value;\n"); - writeln!( - output, - " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" - ) - .unwrap(); - output.push_str("}\n\n"); - } - output -} - -fn render_compat_debug_control(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str("// Included inside `cuda_device::debug` to keep its public API stable.\n\n"); - for operation in [ - DebugControlOperation::Trap, - DebugControlOperation::Breakpoint, - DebugControlOperation::Pmevent, - ] { - let record = debug_controls(catalog) - .find(|record| { - record - .debug_control - .as_ref() - .is_some_and(|debug| debug.operation == operation) - }) - .expect("complete debug-control family"); - match operation { - DebugControlOperation::Trap => { - writeln!(output, "/// {}", record.summary).unwrap(); - output.push_str("#[inline(never)]\npub fn trap() -> ! {\n"); - output.push_str( - " unreachable!(\"trap called outside CUDA kernel context\")\n}\n\n", - ); - } - DebugControlOperation::Breakpoint => { - writeln!(output, "/// {}", record.summary).unwrap(); - output.push_str("#[inline(never)]\npub fn breakpoint() {\n"); - output.push_str( - " unreachable!(\"breakpoint called outside CUDA kernel context\")\n}\n\n", - ); - } - DebugControlOperation::Pmevent => { - output.push_str( - r#"/// Triggers performance monitor event `N`, which must be in `0..=15`. -#[inline(always)] -pub fn prof_trigger() { - __prof_trigger(N); -} - -#[doc(hidden)] -#[inline(never)] -pub(crate) fn __prof_trigger(_event_id: u32) { - unreachable!("prof_trigger called outside CUDA kernel context") -} - -"#, - ); - } - } - } - output -} - -fn render_compat_stmatrix(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str("// Included inside `cuda_device::tcgen05` to keep public paths stable.\n\n"); - for record in stmatrices(catalog) { - let (multiplicity, _) = stmatrix_variant(record).expect("stmatrix variant"); - let count = multiplicity.register_count(); - let name = stmatrix_compatibility_name(record); - let path = record - .rust - .compatibility_paths - .first() - .expect("stmatrix compatibility path"); - let arguments = std::iter::once("smem_ptr: *mut u8".to_owned()) - .chain((0..count).map(|index| format!("r{index}: u32"))) - .collect::>() - .join(", "); - let values = std::iter::once("smem_ptr".to_owned()) - .chain((0..count).map(|index| format!("r{index}"))) - .collect::>() - .join(", "); - writeln!(output, "/// {}", record.summary).unwrap(); - output.push_str("///\n/// # Safety\n"); - output.push_str( - "/// All warp lanes must participate. The used row addresses must be valid, shared, and 16-byte aligned.\n", - ); - output.push_str("#[inline(never)]\n"); - writeln!(output, "pub unsafe fn {name}({arguments}) {{").unwrap(); - writeln!(output, " let _ = ({values});").unwrap(); - writeln!( - output, - " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" - ) - .unwrap(); - output.push_str("}\n\n"); - } - output -} - -fn render_compat_clc(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str("// Included inside `cuda_device::clc` to keep its public API stable.\n\n"); - for record in clc_intrinsics(catalog) { - let operation = record.clc.as_ref().expect("CLC contract").operation; - writeln!(output, "/// {}", record.summary).unwrap(); - output.push_str("///\n/// # Safety\n"); - render_clc_safety_lines(&mut output, operation, ClcSafetyArgNames::Compatibility); - match operation { - ClcOperation::TryCancel | ClcOperation::TryCancelMulticast => { - output.push_str("#[inline(never)]\n"); - writeln!( - output, - "pub unsafe fn {}(response: *mut u8, mbar: *mut Barrier) {{", - record.rust.name - ) - .unwrap(); - output.push_str(" let _ = (response, mbar);\n"); - } - ClcOperation::QueryIsCanceled - | ClcOperation::QueryGetFirstCtaidX - | ClcOperation::QueryGetFirstCtaidY - | ClcOperation::QueryGetFirstCtaidZ => { - output.push_str("#[inline(never)]\n"); - writeln!( - output, - "pub unsafe fn {}(resp_lo: u64, resp_hi: u64) -> u32 {{", - record.rust.name - ) - .unwrap(); - output.push_str(" let _ = (resp_lo, resp_hi);\n"); - } - } - writeln!( - output, - " unreachable!(\"{} called outside CUDA kernel context\")", - record.rust.name - ) - .unwrap(); - output.push_str("}\n\n"); - } - output -} - -fn render_compat_counted_barrier(catalog: &CatalogFile, hash: &str) -> String { - assert_eq!( - execution_control_family(catalog, "counted_barrier").count(), - 4 - ); - let mut output = rust_header(catalog, hash); - output.push_str("// Included inside `cuda_device::barrier` to keep its public API stable.\n\n"); - for record in execution_control_family(catalog, "counted_barrier") { - writeln!(output, "/// {}", record.summary).unwrap(); - output.push_str( - "///\n/// # Safety\n/// Every participating thread must use a compatible barrier ID and expected thread count.\n#[inline(never)]\n", - ); - writeln!( - output, - "pub unsafe fn {}(barrier_id: u32, thread_count: u32) {{", - record.rust.name - ) - .unwrap(); - output.push_str(" let _ = (barrier_id, thread_count);\n"); - writeln!( - output, - " unreachable!(\"{} called outside CUDA kernel context\")", - record.rust.name - ) - .unwrap(); - output.push_str("}\n\n"); - } - output -} - -fn render_compat_grid_dependency(catalog: &CatalogFile, hash: &str) -> String { - assert_eq!( - execution_control_family(catalog, "grid_dependency").count(), - 2 - ); - let mut output = rust_header(catalog, hash); - output.push_str( - "// Included inside `cuda_device::grid` to keep its public API stable.\n\npub mod dependency {\n", - ); - for (id, name) in [ - ("grid_dependency_launch_dependents", "trigger_dependents"), - ("grid_dependency_wait", "wait"), - ] { - let record = execution_control_family(catalog, "grid_dependency") - .find(|record| record.id == id) - .expect("complete grid-dependency family"); - writeln!(output, " /// {}", record.summary).unwrap(); - output.push_str( - " ///\n /// # Safety\n /// The kernel launch must participate in a valid programmatic dependent-launch protocol.\n #[inline(never)]\n", - ); - writeln!(output, " pub unsafe fn {name}() {{").unwrap(); - writeln!( - output, - " unreachable!(\"{name} called outside CUDA kernel context\")" - ) - .unwrap(); - output.push_str(" }\n\n"); - } - output.push_str("}\n"); - output -} - -fn render_compat_register_control(catalog: &CatalogFile, hash: &str) -> String { - assert_eq!( - execution_control_family(catalog, "register_control").count(), - 2 - ); - let mut output = rust_header(catalog, hash); - output.push_str("// Included inside `cuda_device::thread` to keep its public API stable.\n\n"); - for (id, public_name, hidden_name) in [ - ("setmaxnreg_inc", "setmaxnreg_inc", "__setmaxnreg_inc"), - ("setmaxnreg_dec", "setmaxnreg_dec", "__setmaxnreg_dec"), - ] { - let record = execution_control_family(catalog, "register_control") - .find(|record| record.id == id) - .expect("complete register-control family"); - writeln!(output, "/// {}", record.summary).unwrap(); - output.push_str( - "///\n/// # Safety\n/// Every thread in the warpgroup must execute the same operation with the same count.\n#[inline(always)]\n", - ); - writeln!(output, "pub unsafe fn {public_name}() {{").unwrap(); - writeln!(output, " unsafe {{ {hidden_name}(N) }}").unwrap(); - output.push_str("}\n\n#[doc(hidden)]\n#[inline(never)]\n"); - writeln!( - output, - "pub(crate) unsafe fn {hidden_name}(_register_count: u32) {{" - ) - .unwrap(); - writeln!( - output, - " unreachable!(\"{public_name} called outside CUDA kernel context\")" - ) - .unwrap(); - output.push_str("}\n\n"); - } - output -} - -fn render_compat_tma(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str("// Included inside `cuda_device::tma` to keep its public API stable.\n\n"); - for record in tma_intrinsics(catalog) { - let tma = record.tma.as_ref().expect("TMA contract"); - let operation = tma.operation; - writeln!(output, "/// {}", record.summary).unwrap(); - let dimensions = tma.dimensions(); - let is_g2s = matches!( - operation, - TmaOperation::G2sTile1d - | TmaOperation::G2sTile2d - | TmaOperation::G2sTile2dMulticast - | TmaOperation::G2sTile2dMulticastCg2 - | TmaOperation::G2sTile3d - | TmaOperation::G2sTile4d - | TmaOperation::G2sTile5d - ); - let is_s2g = matches!( - operation, - TmaOperation::S2gTile1d - | TmaOperation::S2gTile2d - | TmaOperation::S2gTile3d - | TmaOperation::S2gTile4d - | TmaOperation::S2gTile5d - ); - let is_reduction = operation == TmaOperation::Reduce; - if !record.rust.safe { - output.push_str("///\n/// # Safety\n"); - if is_g2s { - output.push_str( - "/// `dst`, `tensor_map`, and `barrier` must remain valid until the copy completes.\n", - ); - } else if is_s2g { - output.push_str( - "/// `src` and `tensor_map` must remain valid until the committed copy group completes.\n", - ); - } else if is_reduction { - output.push_str( - "/// `src` must name a live shared-memory tile and `tensor_map` must describe a compatible global destination until the committed reduction completes.\n", - ); - } else if matches!( - tma.adapter, - TmaAdapter::DescriptorAndAddressPointers - | TmaAdapter::DescriptorOrdinalAndU32 - | TmaAdapter::DescriptorOrdinalAndU64 - | TmaAdapter::DescriptorAndImmediateU32 - | TmaAdapter::DescriptorAndRuntimeU32 - ) { - output.push_str( - "/// `tensor_map` must point to a writable, 128-byte tensor-map descriptor in global memory.\n", - ); - } else { - output.push_str("/// `tensor_map` must point to a live tensor-map descriptor.\n"); - } - } - output.push_str("#[inline(never)]\n"); - if is_g2s { - let dimensions = dimensions.unwrap(); - let mut arguments = vec![ - "dst: *mut u8".to_owned(), - "tensor_map: *const TmaDescriptor".to_owned(), - ]; - let mut values = vec!["dst".to_owned(), "tensor_map".to_owned()]; - for index in 0..dimensions { - arguments.push(format!("coord{index}: i32")); - values.push(format!("coord{index}")); - } - arguments.push("barrier: *mut Barrier".into()); - values.push("barrier".into()); - if matches!( - operation, - TmaOperation::G2sTile2dMulticast | TmaOperation::G2sTile2dMulticastCg2 - ) { - arguments.push("cta_mask: u16".into()); - values.push("cta_mask".into()); - } - if arguments.len() > 7 { - output.push_str("#[allow(clippy::too_many_arguments)]\n"); - } - writeln!( - output, - "pub unsafe fn {}({}) {{", - record.rust.name, - arguments.join(", ") - ) - .unwrap(); - writeln!(output, " let _ = ({});", values.join(", ")).unwrap(); - } else if is_s2g || is_reduction { - let dimensions = dimensions.unwrap(); - let mut arguments = vec![ - "src: *const u8".to_owned(), - "tensor_map: *const TmaDescriptor".to_owned(), - ]; - let mut values = vec!["src".to_owned(), "tensor_map".to_owned()]; - for index in 0..dimensions { - arguments.push(format!("coord{index}: i32")); - values.push(format!("coord{index}")); - } - writeln!( - output, - "pub unsafe fn {}({}) {{", - record.rust.name, - arguments.join(", ") - ) - .unwrap(); - writeln!(output, " let _ = ({});", values.join(", ")).unwrap(); - } else if operation == TmaOperation::CommitGroup { - writeln!(output, "pub fn {}() {{", record.rust.name).unwrap(); - } else if matches!( - operation, - TmaOperation::WaitGroup | TmaOperation::WaitGroupRead - ) { - writeln!(output, "pub fn {}(n: u32) {{", record.rust.name).unwrap(); - output.push_str(" let _ = n;\n"); - } else if operation == TmaOperation::PrefetchTensorMap { - writeln!( - output, - "pub unsafe fn {}(tensor_map: *const TmaDescriptor) {{", - record.rust.name - ) - .unwrap(); - output.push_str(" let _ = tensor_map;\n"); - } else if let Some(coordinate_count) = operation.prefetch_coordinate_count() { - let mut arguments = vec!["tensor_map: *const TmaDescriptor".to_owned()]; - let mut values = vec!["tensor_map".to_owned()]; - for index in 0..coordinate_count { - arguments.push(format!("coord{index}: i32")); - values.push(format!("coord{index}")); - } - if operation.uses_prefetch_cache_hint() { - arguments.push("cache_hint: u64".into()); - values.push("cache_hint".into()); - } - writeln!( - output, - "pub unsafe fn {}({}) {{", - record.rust.name, - arguments.join(", ") - ) - .unwrap(); - writeln!(output, " let _ = ({});", values.join(", ")).unwrap(); - } else { - match tma.adapter { - TmaAdapter::DescriptorAndAddressPointers => { - writeln!( - output, - "pub unsafe fn {}(tensor_map: *mut TmaDescriptor, new_address: *const u8) {{", - record.rust.name - ) - .unwrap(); - output.push_str(" let _ = (tensor_map, new_address);\n"); - } - TmaAdapter::DescriptorOrdinalAndU32 => { - writeln!( - output, - "pub unsafe fn {}(tensor_map: *mut TmaDescriptor, ordinal: u32, new_value: u32) {{", - record.rust.name - ) - .unwrap(); - output.push_str(" let _ = (tensor_map, ordinal, new_value);\n"); - } - TmaAdapter::DescriptorOrdinalAndU64 => { - writeln!( - output, - "pub unsafe fn {}(tensor_map: *mut TmaDescriptor, ordinal: u32, new_value: u64) {{", - record.rust.name - ) - .unwrap(); - output.push_str(" let _ = (tensor_map, ordinal, new_value);\n"); - } - TmaAdapter::DescriptorAndImmediateU32 | TmaAdapter::DescriptorAndRuntimeU32 => { - writeln!( - output, - "pub unsafe fn {}(tensor_map: *mut TmaDescriptor, new_value: u32) {{", - record.rust.name - ) - .unwrap(); - output.push_str(" let _ = (tensor_map, new_value);\n"); - } - TmaAdapter::DescriptorPointerInjectBytes => { - writeln!( - output, - "pub unsafe fn {}(tensor_map: *const TmaDescriptor) {{", - record.rust.name - ) - .unwrap(); - output.push_str(" let _ = tensor_map;\n"); - } - TmaAdapter::NoOperands => { - writeln!(output, "pub fn {}() {{", record.rust.name).unwrap(); - } - _ => unreachable!("TMA compatibility operation category was matched"), - } - } - writeln!( - output, - " unreachable!(\"{} called outside CUDA kernel context\")", - record.rust.name - ) - .unwrap(); - output.push_str("}\n\n"); - } - output.push_str( - r#"/// Replace a global tensor-map address and publish it to the tensor-map proxy. -/// -/// # Safety -/// `tensor_map` must point to a writable, 128-byte tensor-map descriptor in -/// global memory, and `new_address` must remain valid for every later TMA use. -#[inline(always)] -pub unsafe fn replace_tma_global_address( - tensor_map: *mut TmaDescriptor, - new_address: *const u8, -) { - unsafe { tensormap_replace_global_address(tensor_map, new_address) }; - fence_proxy_tensormap_generic_release_gpu(); - unsafe { fence_proxy_tensormap_generic_acquire_gpu(tensor_map) }; -} -"#, - ); - output -} - -fn render_compat_wgmma_control(catalog: &CatalogFile, hash: &str) -> String { - assert_eq!(wgmma_controls(catalog).count(), 3); - let mut output = rust_header(catalog, hash); - output.push_str("// Included inside `cuda_device::wgmma` to keep its public API stable.\n\n"); - - for mode in [WgmmaControlMode::Fence, WgmmaControlMode::CommitGroup] { - let record = wgmma_control(catalog, mode); - writeln!(output, "/// {}", record.summary).unwrap(); - output.push_str("#[inline(never)]\n"); - writeln!(output, "pub fn {}() {{", record.rust.name).unwrap(); - writeln!( - output, - " unreachable!({:?})", - format!("{} called outside CUDA kernel context", record.rust.name) - ) - .unwrap(); - output.push_str("}\n\n"); - } - - let wait = wgmma_control(catalog, WgmmaControlMode::WaitGroup); - writeln!(output, "/// {}", wait.summary).unwrap(); - output.push_str( - r#"#[inline(always)] -pub fn wgmma_wait_group() { - __wgmma_wait_group(N as u64); -} - -#[doc(hidden)] -#[inline(never)] -pub(crate) fn __wgmma_wait_group(_max_pending: u64) { - unreachable!("wgmma_wait_group called outside CUDA kernel context") -} -"#, - ); - output -} - -fn render_compat_tcgen05_mma_record(output: &mut String, record: &CatalogIntrinsic) { - let mma = record - .tcgen05 - .as_ref() - .and_then(|tcgen05| tcgen05.mma.as_ref()) - .expect("tcgen05 MMA record"); - let parameters = tcgen05_mma_runtime_parameters(mma); - let signature = parameters - .iter() - .map(|(name, ty)| format!("{name}: {ty}")) - .collect::>() - .join(", "); - let values = parameters - .iter() - .map(|(name, _)| *name) - .collect::>() - .join(", "); - - writeln!(output, "/// {}", record.summary).unwrap(); - output.push_str("/// All tcgen05 operations in the kernel must use the same CTA-group mode.\n"); - if mma.alias.is_some() && mma.form == Tcgen05MmaForm::WsTensor { - output.push_str("/// This uses kind f8f6f4 and collector b0::discard.\n"); - output.push_str("/// `legacy_a_desc` is kept for compatibility; tensor A uses `a_tmem`.\n"); - } else if mma.alias.is_some() { - output.push_str("/// This uses kind f8f6f4, CTA group 1, and collector a::discard.\n"); - } else { - match mma.selector_layout { - Tcgen05MmaSelectorLayout::Base { .. } => { - output.push_str("/// `KIND` is 0=f16, 1=tf32, 2=f8f6f4, or 3=i8.\n"); - output.push_str( - "/// `CTA_GROUP` is 1 or 2. `COLLECTOR_A` is 0=discard, 1=lastuse, 2=fill, or 3=use.\n", - ); - } - Tcgen05MmaSelectorLayout::WarpSpecialized { .. } => { - output.push_str("/// `KIND` is 0=f16, 1=tf32, 2=f8f6f4, or 3=i8.\n"); - output.push_str( - "/// `B_BUFFER` is 0 through 3. `B_USAGE` is 0=discard, 1=lastuse, 2=fill, or 3=use.\n", - ); - } - } - } - output.push_str("///\n/// # Safety\n"); - output - .push_str("/// Tensor-memory addresses and descriptors must be valid for this MMA form.\n"); - if mma.alias.is_some() { - output.push_str("#[inline(never)]\n"); - output.push_str("#[allow(clippy::too_many_arguments)]\n"); - writeln!(output, "pub unsafe fn {}({signature}) {{", record.rust.name).unwrap(); - writeln!(output, " let _ = ({values});").unwrap(); - writeln!( - output, - " unreachable!(\"{} called outside CUDA kernel context\")", - record.rust.name - ) - .unwrap(); - output.push_str("}\n\n"); - return; - } - - let selectors = tcgen05_mma_selector_parameters(mma.selector_layout); - let const_parameters = selectors - .iter() - .map(|(name, _)| format!("const {name}: u32")) - .collect::>() - .join(", "); - let selector_values = selectors - .iter() - .map(|(name, _)| *name) - .collect::>() - .join(", "); - output.push_str("#[inline(always)]\n"); - if parameters.len() > 5 { - output.push_str("#[allow(clippy::too_many_arguments)]\n"); - } - writeln!( - output, - "pub unsafe fn {}<{const_parameters}>({signature}) {{", - record.rust.name - ) - .unwrap(); - writeln!( - output, - " unsafe {{ __{}({values}, {selector_values}) }}", - record.rust.name - ) - .unwrap(); - output.push_str("}\n\n#[doc(hidden)]\n#[inline(never)]\n"); - output.push_str("#[allow(clippy::too_many_arguments)]\n"); - let hidden_parameters = parameters - .iter() - .map(|(name, ty)| format!("_{name}: {ty}")) - .chain(selectors.iter().map(|(_, name)| format!("_{name}: u32"))) - .collect::>() - .join(", "); - writeln!( - output, - "pub(crate) unsafe fn __{}({hidden_parameters}) {{", - record.rust.name - ) - .unwrap(); - writeln!( - output, - " unreachable!(\"{} called outside CUDA kernel context\")", - record.rust.name - ) - .unwrap(); - output.push_str("}\n\n"); -} - -fn render_compat_tcgen05(catalog: &CatalogFile, hash: &str) -> String { - assert_eq!(tcgen05_intrinsics(catalog).count(), 233); - let mut output = rust_header(catalog, hash); - output.push_str("// Included inside `cuda_device::tcgen05` to keep its public API stable.\n\n"); - for record in tcgen05_intrinsics(catalog) { - let tcgen05 = record.tcgen05.as_ref().unwrap(); - if tcgen05.mma.is_some() { - render_compat_tcgen05_mma_record(&mut output, record); - continue; - } - let operation = tcgen05.operation; - let has_half_split_offset = tcgen05 - .ld - .is_some_and(|ld| ld.shape == Tcgen05LdShape::M16x32bx2) - || tcgen05 - .st - .is_some_and(|st| st.shape == Tcgen05LdShape::M16x32bx2); - if has_half_split_offset { - let (result, data) = if operation == Tcgen05Operation::Ld { - let count = tcgen05_ld_register_count(record); - ( - if count == 1 { - "u32".into() - } else { - format!("CuSimd") - }, - None, - ) - } else { - let count = tcgen05_st_register_count(record); - ( - "()".into(), - Some(if count == 1 { - "u32".into() - } else { - format!("CuSimd") - }), - ) - }; - writeln!(output, "/// {}", record.summary).unwrap(); - output.push_str( - "/// All tcgen05 operations in the kernel must use the same CTA-group mode.\n", - ); - output.push_str("///\n/// # Safety\n"); - output.push_str( - "/// The tensor-memory address must be live, and the warp must execute this call uniformly.\n", - ); - output.push_str("#[inline(always)]\n"); - if let Some(data) = &data { - writeln!( - output, - "pub unsafe fn {}(tmem_addr: u32, data: {data}) {{", - record.rust.name - ) - .unwrap(); - writeln!( - output, - " unsafe {{ __{}(tmem_addr, HALF_SPLIT_OFFSET as i64, data) }}", - record.rust.name - ) - .unwrap(); - } else { - writeln!( - output, - "pub unsafe fn {}(tmem_addr: u32) -> {result} {{", - record.rust.name - ) - .unwrap(); - writeln!( - output, - " unsafe {{ __{}(tmem_addr, HALF_SPLIT_OFFSET as i64) }}", - record.rust.name - ) - .unwrap(); - } - output.push_str("}\n\n#[doc(hidden)]\n#[inline(never)]\n"); - if let Some(data) = data { - writeln!( - output, - "pub(crate) unsafe fn __{}(_tmem_addr: u32, _half_split_offset: i64, _data: {data}) {{", - record.rust.name - ) - .unwrap(); - } else { - writeln!( - output, - "pub(crate) unsafe fn __{}(_tmem_addr: u32, _half_split_offset: i64) -> {result} {{", - record.rust.name - ) - .unwrap(); - } - writeln!( - output, - " unreachable!(\"{} called outside CUDA kernel context\")", - record.rust.name - ) - .unwrap(); - output.push_str("}\n\n"); - continue; - } - let (arguments, values): (String, String) = if operation == Tcgen05Operation::St { - let count = tcgen05_st_register_count(record); - let data = if count == 1 { - "u32".into() - } else { - format!("CuSimd") - }; - ( - format!("tmem_addr: u32, data: {data}"), - "tmem_addr, data".into(), - ) - } else { - let (arguments, values) = match operation { - Tcgen05Operation::Alloc | Tcgen05Operation::AllocCg2 => { - ("dst_smem: *mut u32, n_cols: u32", "dst_smem, n_cols") - } - Tcgen05Operation::Dealloc | Tcgen05Operation::DeallocCg2 => { - ("tmem_addr: u32, n_cols: u32", "tmem_addr, n_cols") - } - Tcgen05Operation::RelinquishAllocPermit - | Tcgen05Operation::FenceBeforeThreadSync - | Tcgen05Operation::FenceAfterThreadSync - | Tcgen05Operation::LoadWait - | Tcgen05Operation::StoreWait - | Tcgen05Operation::RelinquishAllocPermitCg2 => ("", ""), - Tcgen05Operation::Commit - | Tcgen05Operation::CommitSharedCluster - | Tcgen05Operation::CommitCg2 - | Tcgen05Operation::CommitSharedClusterCg2 => ("mbar: *mut u64", "mbar"), - Tcgen05Operation::MmaWsF16 - | Tcgen05Operation::MmaWsBf16 - | Tcgen05Operation::MmaWsTf32 => ( - "d_tmem: u32, a_tmem: u32, a_desc: u64, b_desc: u64, idesc: u32, enable_d: bool", - "d_tmem, a_tmem, a_desc, b_desc, idesc, enable_d", - ), - Tcgen05Operation::MmaF16 | Tcgen05Operation::MmaF16Cg2 => ( - "d_tmem: u32, a_desc: u64, b_desc: u64, idesc: u32, enable_d: bool", - "d_tmem, a_desc, b_desc, idesc, enable_d", - ), - Tcgen05Operation::CpSmemToTmem | Tcgen05Operation::CpSmemToTmemCg2 => { - ("tmem_addr: u32, smem_desc: u64", "tmem_addr, smem_desc") - } - Tcgen05Operation::Ld16x256bX8Pure - | Tcgen05Operation::Ld16x256bPure - | Tcgen05Operation::Ld => ("tmem_addr: u32", "tmem_addr"), - Tcgen05Operation::CommitMulticast | Tcgen05Operation::CommitMulticastCg2 => { - ("mbar: *mut u64, cta_mask: u16", "mbar, cta_mask") - } - Tcgen05Operation::ShiftDown | Tcgen05Operation::ShiftDownCg2 => { - ("tmem_addr: u32", "tmem_addr") - } - Tcgen05Operation::St => unreachable!("store handled above"), - Tcgen05Operation::Mma => unreachable!("generic MMA handled above"), - }; - (arguments.into(), values.into()) - }; - let result: String = match operation { - Tcgen05Operation::Ld16x256bX8Pure => "TmemF32x32".into(), - Tcgen05Operation::Ld16x256bPure => "TmemF32x4".into(), - Tcgen05Operation::Ld => { - let count = tcgen05_ld_register_count(record); - if count == 1 { - "u32".into() - } else { - format!("CuSimd") - } - } - Tcgen05Operation::St => "()".into(), - _ => "()".into(), - }; - writeln!(output, "/// {}", record.summary).unwrap(); - if let Some(participation) = tcgen05_participation_doc(operation) { - writeln!(output, "/// {participation}").unwrap(); - } - output.push_str( - "/// All tcgen05 operations in the kernel must use the same CTA-group mode.\n", - ); - if !record.rust.safe { - output.push_str("///\n/// # Safety\n"); - if operation == Tcgen05Operation::St { - output.push_str( - "/// The tensor-memory address must remain live and cover the selected tile.\n\ - /// All active warp lanes must execute convergently with the same address. Complete the store wait before relying on completion.\n", - ); - } else if operation == Tcgen05Operation::Ld { - output.push_str( - "/// The tensor-memory address must remain live and cover the selected tile.\n\ - /// All active warp lanes must execute convergently with the same address. Complete the load wait before using the result.\n", - ); - } else if tcgen05_is_commit(operation) { - if tcgen05_is_multicast_commit(operation) { - output.push_str( - "/// `mbar` must point to a live initialized cluster-shared mbarrier. `cta_mask` must select valid CTA ranks in its cluster.\n", - ); - } else { - output.push_str( - "/// `mbar` must point to a live initialized mbarrier valid for this CTA-group mode.\n", - ); - } - output.push_str( - "/// The same thread that issued the tracked asynchronous tcgen05 operations must issue this commit.\n", - ); - } else if tcgen05_is_shift(operation) { - output.push_str( - "/// `tmem_addr` must name a live tensor-memory allocation, and its lane component must be a multiple of 32.\n\ - /// Completion must be tracked by a matching commit from that same thread and observed through the selected mbarrier before relying on shifted data.\n", - ); - } else { - output.push_str("/// The caller must satisfy the tcgen05 address, lifetime, and participation rules.\n"); - } - } - output.push_str("#[inline(never)]\n"); - if record.rust.arguments.len() > 5 { - output.push_str("#[allow(clippy::too_many_arguments)]\n"); - } - let safety = if record.rust.safe { "" } else { "unsafe " }; - writeln!( - output, - "pub {safety}fn {}({arguments}) -> {result} {{", - record.rust.name - ) - .unwrap(); - if !values.is_empty() { - if values.contains(',') { - writeln!(output, " let _ = ({values});").unwrap(); - } else { - writeln!(output, " let _ = {values};").unwrap(); - } - } - writeln!( - output, - " unreachable!(\"{} called outside CUDA kernel context\")", - record.rust.name - ) - .unwrap(); - output.push_str("}\n\n"); - } - output -} - -fn render_compat_packed_atomic(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str("// Included inside `cuda_device::atomic` to keep existing paths stable.\n\n"); - for record in packed_atomics(catalog) { - let path = record - .rust - .compatibility_paths - .iter() - .find(|path| path.starts_with("cuda_device::atomic::")) - .expect("packed-atomic compatibility path"); - assert_eq!(path, &format!("cuda_device::atomic::{}", record.rust.name)); - let packed = record - .packed_atomic - .as_ref() - .expect("packed-atomic semantics"); - let lane_type = match packed.format { - PackedAtomicFormat::F16x2 => "f16", - PackedAtomicFormat::Bf16x2 => "bf16", - }; - let minimum_sm = match &record.target.hardware { - CatalogHardwareTarget::AnyOf { alternatives } => match alternatives.as_slice() { - [CatalogHardwareAlternative::MinimumSm { sm }] => *sm, - _ => panic!("packed-atomic compatibility API requires one minimum SM"), - }, - _ => panic!("packed-atomic compatibility API requires one minimum SM"), - }; - assert!(!record.rust.safe); - assert!(record.rust.must_use); - assert_eq!(record.rust.arguments, ["*mut u32", "u32"]); - assert_eq!(record.rust.result, "u32"); - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!( - output, - "/// `val` and the result pack two {lane_type} lanes into `u32`, low lane first." - ) - .unwrap(); - output.push_str( - "/// The lanes are atomic independently and may not form one old 32-bit snapshot.\n", - ); - output.push_str("/// This is a relaxed GPU-scope operation. Each lane rounds to nearest-even and preserves subnormals.\n"); - writeln!( - output, - "/// Requires PTX {} and `sm_{minimum_sm}+`.", - record.target.minimum_ptx - ) - .unwrap(); - output.push_str("///\n/// # Safety\n"); - output.push_str( - "/// `addr` must point to four writable, four-byte-aligned bytes in global memory.\n", - ); - output.push_str("/// Do not overlap this operation with a whole-word atomic or non-atomic lane access.\n"); - output.push_str("/// Racing atomics must use mutually inclusive scopes; host/system access is not included.\n"); - output.push_str("#[must_use]\n#[inline(never)]\n"); - writeln!( - output, - "pub unsafe fn {}(addr: *mut u32, val: u32) -> u32 {{", - record.rust.name - ) - .unwrap(); - output.push_str(" let _ = (addr, val);\n"); - writeln!( - output, - " unreachable!(\"{} called outside CUDA kernel context\")", - record.rust.name - ) - .unwrap(); - output.push_str("}\n\n"); - } - output -} - -fn render_compat_cp_async_copy(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - "// Included inside `cuda_device::async_copy` to keep existing paths stable.\n\n", - ); - for record in cp_async_copies(catalog) { - let copy = record.cp_async_copy.as_ref().expect("cp.async semantics"); - let bytes = copy.copy_size.bytes(); - let cache = match copy.cache_policy { - CpAsyncCachePolicy::Ca => "cache-all", - CpAsyncCachePolicy::Cg => "cache-global", - }; - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!( - output, - "/// Uses the {cache} policy and copies {bytes} bytes." - ) - .unwrap(); - if copy.source_size == CpAsyncSourceSize::Runtime { - writeln!( - output, - "/// Bytes after `src_size` are filled with zero; `src_size` must be at most {bytes}." - ) - .unwrap(); - } - output.push_str("///\n/// # Safety\n"); - writeln!( - output, - "/// `shared_dst` must point to {bytes} writable bytes in shared memory and be aligned to {bytes} bytes." - ) - .unwrap(); - if copy.source_size == CpAsyncSourceSize::Runtime { - writeln!( - output, - "/// `global_src` must point to `src_size` readable bytes in global memory and be aligned to {bytes} bytes." - ) - .unwrap(); - } else { - writeln!( - output, - "/// `global_src` must point to {bytes} readable bytes in global memory and be aligned to {bytes} bytes." - ) - .unwrap(); - } - output.push_str( - "/// Both ranges must remain valid, the source must remain unchanged, and the destination must not be accessed until this copy completes.\n\ - /// The issuing thread must complete the copy. Synchronize threads afterward before another thread accesses the destination.\n\ - /// User-authored completion assembly must include a compiler memory clobber.\n", - ); - output.push_str("#[inline(never)]\n"); - if copy.source_size == CpAsyncSourceSize::Runtime { - writeln!( - output, - "pub unsafe fn {}(_shared_dst: *mut u32, _global_src: *const u8, _src_size: u32) {{", - record.rust.name - ) - .unwrap(); - } else { - writeln!( - output, - "pub unsafe fn {}(_shared_dst: *mut u32, _global_src: *const u32) {{", - record.rust.name - ) - .unwrap(); - } - writeln!( - output, - " unreachable!(\"{} called outside CUDA kernel context\")", - record.rust.name - ) - .unwrap(); - output.push_str("}\n\n"); - } - for record in cp_async_controls(catalog) { - let control = record - .cp_async_control - .as_ref() - .expect("cp.async control semantics"); - writeln!(output, "/// {}", record.summary).unwrap(); - output.push_str("///\n/// # Safety\n"); - match control.operation { - CpAsyncControlOperation::CommitGroup => output.push_str( - "/// This commits only copies issued by the executing thread and does not wait for completion.\n", - ), - CpAsyncControlOperation::WaitAll => output.push_str( - "/// This waits only for copies issued by this thread. Synchronize threads before another thread accesses a completed destination.\n", - ), - CpAsyncControlOperation::WaitGroup => output.push_str( - "/// `max_pending` must be a compile-time constant. Access only destinations whose copy groups this wait completes.\n", - ), - } - output.push_str("#[inline(never)]\n"); - if control.operation == CpAsyncControlOperation::WaitGroup { - writeln!( - output, - "pub unsafe fn {}(_max_pending: u32) {{", - record.rust.name - ) - .unwrap(); - } else { - writeln!(output, "pub unsafe fn {}() {{", record.rust.name).unwrap(); - } - writeln!( - output, - " unreachable!(\"{} called outside CUDA kernel context\")", - record.rust.name - ) - .unwrap(); - output.push_str("}\n\n"); - } - for record in cp_async_mbarriers(catalog) { - let bridge = record - .cp_async_mbarrier - .as_ref() - .expect("cp.async mbarrier semantics"); - let path = record - .rust - .compatibility_paths - .iter() - .find(|path| path.starts_with("cuda_device::async_copy::")) - .expect("cp.async mbarrier compatibility path"); - assert_eq!( - path, - &format!("cuda_device::async_copy::{}", record.rust.name) - ); - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!(output, "/// Lowers to `{}`.", record.expected_ptx).unwrap(); - output.push_str("///\n/// # Safety\n"); - output.push_str( - "/// `barrier` must point to a live, initialized, eight-byte-aligned mbarrier object in shared memory.\n\ - /// This thread must have issued the `cp.async` operations being associated with it.\n\ - /// The object must remain valid until those operations complete.\n", - ); - match bridge.operation { - CpAsyncMbarrierOperation::Arrive => output.push_str( - "/// This form increments the pending count before scheduling the asynchronous arrival.\n\ - /// That increment must not exceed the barrier's pending-count limit.\n", - ), - CpAsyncMbarrierOperation::ArriveNoInc => output.push_str( - "/// The initial pending count must already include the asynchronous arrival.\n", - ), - } - output.push_str("#[inline(never)]\n"); - writeln!( - output, - "pub unsafe fn {}(_barrier: *mut crate::barrier::Barrier) {{", - record.rust.name - ) - .unwrap(); - writeln!( - output, - " unreachable!(\"{} called outside CUDA kernel context\")", - record.rust.name - ) - .unwrap(); - output.push_str("}\n\n"); - } - output -} - -fn render_compat_mbarrier_basic(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str("// Included inside `cuda_device::barrier` to keep existing paths stable.\n\n"); - for record in mbarrier_basics(catalog) { - let mbarrier = record - .mbarrier_basic - .as_ref() - .expect("basic mbarrier semantics"); - let path = record - .rust - .compatibility_paths - .iter() - .find(|path| path.starts_with("cuda_device::barrier::")) - .expect("basic mbarrier compatibility path"); - assert_eq!(path, &format!("cuda_device::barrier::{}", record.rust.name)); - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!(output, "/// Lowers to `{}`.", record.expected_ptx).unwrap(); - output.push_str("///\n/// # Safety\n"); - output.push_str( - "/// `bar` must point to a live, eight-byte-aligned `Barrier` in shared memory.\n", - ); - match mbarrier.operation { - MbarrierBasicOperation::Init => output.push_str( - "/// Exactly one thread may initialize it. `expected_count` must be in `1..=0xFFFFF` and include every arrival in the phase.\n\ - /// The barrier must be uninitialized or invalidated, with no concurrent barrier operation.\n", - ), - MbarrierBasicOperation::Arrive => output.push_str( - "/// The barrier must be initialized, and this arrival must be included in the current phase's expected count.\n\ - /// Use the returned token only with this barrier and phase.\n", - ), - MbarrierBasicOperation::TestWait => output.push_str( - "/// The barrier must be initialized. `token` must come from this barrier and phase.\n", - ), - MbarrierBasicOperation::Inval => output.push_str( - "/// The barrier must be initialized and unused by all threads and asynchronous operations.\n\ - /// Exactly one thread may invalidate it.\n", - ), - } - output.push_str("#[inline(never)]\n"); - match mbarrier.operation { - MbarrierBasicOperation::Init => { - writeln!( - output, - "pub unsafe fn {}(bar: *mut Barrier, expected_count: u32) {{", - record.rust.name - ) - .unwrap(); - output.push_str(" let _ = (bar, expected_count);\n"); - } - MbarrierBasicOperation::Arrive => { - writeln!( - output, - "pub unsafe fn {}(bar: *const Barrier) -> u64 {{", - record.rust.name - ) - .unwrap(); - output.push_str(" let _ = bar;\n"); - } - MbarrierBasicOperation::TestWait => { - writeln!( - output, - "pub unsafe fn {}(bar: *const Barrier, token: u64) -> bool {{", - record.rust.name - ) - .unwrap(); - output.push_str(" let _ = (bar, token);\n"); - } - MbarrierBasicOperation::Inval => { - writeln!( - output, - "pub unsafe fn {}(bar: *mut Barrier) {{", - record.rust.name - ) - .unwrap(); - output.push_str(" let _ = bar;\n"); - } - } - writeln!( - output, - " unreachable!(\"{} called outside CUDA kernel context\")", - record.rust.name - ) - .unwrap(); - output.push_str("}\n\n"); - } - output -} - -fn render_compat_mbarrier_extended(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str("// Included inside `cuda_device::barrier` to keep existing paths stable.\n\n"); - for record in mbarrier_extended(catalog) { - let contract = record - .mbarrier_extended - .as_ref() - .expect("extended-mbarrier contract"); - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!(output, "/// Lowers to `{}`.", record.expected_ptx).unwrap(); - output.push_str("///\n/// # Safety\n"); - output.push_str( - "/// The caller must satisfy the barrier, scope, and memory-ordering contract.\n", - ); - output.push_str("#[inline(never)]\n"); - match contract.operation { - MbarrierExtendedOperation::ArriveExpectTxCta - | MbarrierExtendedOperation::ArriveExpectTxCluster => { - writeln!( - output, - "pub unsafe fn {}(bar: *const Barrier, _tx_count: u32, bytes: u32) -> u64 {{", - record.rust.name - ) - .unwrap(); - output.push_str(" let _ = (bar, bytes);\n"); - } - MbarrierExtendedOperation::ArriveRemoteCluster => { - writeln!( - output, - "pub unsafe fn {}(remote_bar_addr: u64) {{", - record.rust.name - ) - .unwrap(); - output.push_str(" let _ = remote_bar_addr;\n"); - } - MbarrierExtendedOperation::TryWaitTokenCta => { - writeln!( - output, - "pub unsafe fn {}(bar: *const Barrier, token: u64) -> bool {{", - record.rust.name - ) - .unwrap(); - output.push_str(" let _ = (bar, token);\n"); - } - MbarrierExtendedOperation::TryWaitParityCta - | MbarrierExtendedOperation::TryWaitParityCluster => { - writeln!( - output, - "pub unsafe fn {}(bar: *const Barrier, parity: u32) -> bool {{", - record.rust.name - ) - .unwrap(); - output.push_str(" let _ = (bar, parity);\n"); - } - MbarrierExtendedOperation::FenceProxyAsyncSharedCta - | MbarrierExtendedOperation::FenceMbarrierInitReleaseCluster - | MbarrierExtendedOperation::FenceProxyAsyncGenericReleaseSharedCtaCluster - | MbarrierExtendedOperation::FenceProxyAsyncGenericAcquireSharedClusterCluster => { - writeln!(output, "pub unsafe fn {}() {{", record.rust.name).unwrap(); - } - MbarrierExtendedOperation::Nanosleep => { - writeln!(output, "pub unsafe fn {}(ns: u32) {{", record.rust.name).unwrap(); - output.push_str(" let _ = ns;\n"); - } - } - writeln!( - output, - " unreachable!(\"{} called outside CUDA kernel context\")", - record.rust.name - ) - .unwrap(); - output.push_str("}\n\n"); - } - output -} - -fn render_compat_integer_minmax(catalog: &CatalogFile, hash: &str, module: &str) -> String { - let mut output = rust_header(catalog, hash); - writeln!( - output, - "// Included inside `cuda_device::{module}` to keep existing paths stable.\n" - ) - .unwrap(); - for record in integer_minmaxes(catalog).filter(|record| record.rust.module == module) { - let path = record - .rust - .compatibility_paths - .iter() - .find(|path| path.starts_with(&format!("cuda_device::{module}::"))) - .expect("integer-min/max compatibility path"); - let scalar = &record.rust.result; - writeln!(output, "/// {}", record.summary).unwrap(); - output.push_str("#[must_use]\n#[inline(never)]\n"); - writeln!( - output, - "pub fn {}(arg0: {scalar}, arg1: {scalar}) -> {scalar} {{", - record.rust.name - ) - .unwrap(); - output.push_str(" let _ = (arg0, arg1);\n"); - writeln!( - output, - " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" - ) - .unwrap(); - output.push_str("}\n\n"); - } - output -} - -fn render_dialect_integer_minmax(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - "//! Structural operations for generated extended integer min/max.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_i32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32)\n}\n\n", - ); - for record in integer_minmaxes(catalog) { - let signedness = if record.rust.result == "i32" { - "Signed" - } else { - "Unsigned" - }; - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!( - output, - "///\n/// Lowers to `{}`.", - integer_minmax_ptx_mnemonic(record) - ) - .unwrap(); - writeln!( - output, - "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<2>, NResultsInterface<1>],\n)]", - record.dialect.op_name - ) - .unwrap(); - writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); - writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); - writeln!( - output, - " pub fn new(op: Ptr) -> Self {{\n Self {{ op }}\n }}\n\n pub fn build(ctx: &mut Context, arg0: Value, arg1: Value) -> Ptr {{\n let result_ty = IntegerType::get(ctx, 32, Signedness::{signedness});\n Operation::new(\n ctx,\n Self::get_concrete_op_info(),\n vec![result_ty.into()],\n vec![arg0, arg1],\n vec![],\n 0,\n )\n }}\n}}" - ) - .unwrap(); - writeln!(output, "\nimpl Verify for {} {{", record.dialect.op_type).unwrap(); - writeln!( - output, - " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n let op = self.get_operation().deref(ctx);\n if op.get_num_operands() != 2 || op.get_num_results() != 1 {{\n return verify_err!(op.loc(), {:?});\n }}\n if !(0..2).all(|index| is_i32(ctx, op.get_operand(index).get_type(ctx)))\n || !is_i32(ctx, op.get_result(0).get_type(ctx))\n {{\n return verify_err!(op.loc(), {:?});\n }}\n Ok(())\n }}\n}}\n", - format!( - "{} requires exactly 2 operands and one result", - record.dialect.op_name - ), - format!( - "{} operands and result must be 32-bit integers", - record.dialect.op_name - ), - ) - .unwrap(); - } - output.push_str("\npub(super) fn register(ctx: &mut Context) {\n"); - for record in integer_minmaxes(catalog) { - writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); - } - output.push_str("}\n"); - output -} - -fn render_compat_packed_alu(catalog: &CatalogFile, hash: &str, format: PackedAluFormat) -> String { - let mut output = rust_header(catalog, hash); - let (module, _, _, _, _) = packed_alu_format_shape(format); - writeln!( - output, - "// Included inside `cuda_device::{module}` to keep existing paths stable.\n" - ) - .unwrap(); - for record in packed_alus(catalog).filter(|record| { - record - .packed_alu - .as_ref() - .is_some_and(|packed| packed.format == format) - }) { - let path = record - .rust - .compatibility_paths - .iter() - .find(|path| path.starts_with(&format!("cuda_device::{module}::"))) - .expect("packed-ALU compatibility path"); - let arguments = record - .rust - .arguments - .iter() - .enumerate() - .map(|(index, ty)| format!("arg{index}: {ty}")) - .collect::>() - .join(", "); - let values = (0..record.rust.arguments.len()) - .map(|index| format!("arg{index}")) - .collect::>() - .join(", "); - writeln!(output, "/// {}", record.summary).unwrap(); - if record.rust.must_use { - output.push_str("#[must_use]\n"); - } - output.push_str("#[inline(never)]\n"); - writeln!( - output, - "pub fn {}({arguments}) -> {} {{", - record.rust.name, record.rust.result - ) - .unwrap(); - if record.rust.arguments.len() == 1 { - writeln!(output, " let _ = {values};").unwrap(); - } else { - writeln!(output, " let _ = ({values});").unwrap(); - } - writeln!( - output, - " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" - ) - .unwrap(); - output.push_str("}\n\n"); - } - render_compat_extended_minmax_into(&mut output, catalog, module); - output -} - -fn render_compat_scalar_minmax(catalog: &CatalogFile, hash: &str, module: &str) -> String { - let mut output = rust_header(catalog, hash); - writeln!( - output, - "// Included inside `cuda_device::{module}` to keep existing paths stable.\n" - ) - .unwrap(); - render_compat_extended_minmax_into(&mut output, catalog, module); - output -} - -fn render_compat_extended_minmax_into(output: &mut String, catalog: &CatalogFile, module: &str) { - for record in extended_minmax(catalog).filter(|record| record.rust.module == module) { - let path = record - .rust - .compatibility_paths - .iter() - .find(|path| path.starts_with(&format!("cuda_device::{module}::"))) - .expect("extended-minmax compatibility path"); - let ty = extended_minmax_rust_type(record); - writeln!(output, "/// {}", record.summary).unwrap(); - output.push_str("#[must_use]\n#[inline(never)]\n"); - writeln!( - output, - "pub fn {}(a: {ty}, b: {ty}) -> {ty} {{", - record.rust.name - ) - .unwrap(); - output.push_str(" let _ = (a, b);\n"); - writeln!( - output, - " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" - ) - .unwrap(); - output.push_str("}\n\n"); - } -} - -fn render_compat_packed_conversion( - catalog: &CatalogFile, - hash: &str, - path_prefix: &str, - containing_module: &str, - argument_names: (&str, &str), -) -> String { - let mut output = rust_header(catalog, hash); - writeln!( - output, - "// Included inside `cuda_device::{containing_module}`.\n" - ) - .unwrap(); - for record in packed_conversions(catalog).filter(|record| { - record - .rust - .compatibility_paths - .iter() - .any(|path| path.starts_with(path_prefix)) - }) { - let path = record - .rust - .compatibility_paths - .iter() - .find(|path| path.starts_with(path_prefix)) - .expect("packed-conversion compatibility path"); - // A packed source arrives in one register, so the parameter list - // follows the record's own argument types rather than a fixed f32 pair. - let parameter_names: Vec<&str> = match packed_conversion_source(record) { - PackedConversionSourceFormat::F32x2 => vec![argument_names.0, argument_names.1], - PackedConversionSourceFormat::E4m3x2 - | PackedConversionSourceFormat::E5m2x2 - | PackedConversionSourceFormat::F16x2 => vec!["packed"], - }; - let parameter_types = packed_conversion_rust_arguments(record); - let parameters = parameter_names - .iter() - .zip(¶meter_types) - .map(|(name, ty)| format!("{name}: {ty}")) - .collect::>() - .join(", "); - let discards = parameter_names.join(", "); - writeln!(output, "/// {}", record.summary).unwrap(); - output.push_str("#[inline(never)]\n"); - writeln!( - output, - "pub fn {}({parameters}) -> {} {{", - path.rsplit("::") - .next() - .expect("packed-conversion compatibility function name"), - record.rust.result, - ) - .unwrap(); - if parameter_names.len() == 1 { - writeln!(output, " let _ = {discards};").unwrap(); - } else { - writeln!(output, " let _ = ({discards});").unwrap(); - } - writeln!( - output, - " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" - ) - .unwrap(); - output.push_str("}\n\n"); - } - if containing_module == "convert" { - for record in scalar_conversions(catalog) { - let path = record - .rust - .compatibility_paths - .iter() - .find(|path| path.starts_with("cuda_device::convert::")) - .expect("scalar-conversion compatibility path"); - writeln!(output, "/// {}", record.summary).unwrap(); - output.push_str("#[must_use]\n#[inline(never)]\n"); - writeln!(output, "pub fn {}(value: f32) -> u32 {{", record.rust.name).unwrap(); - output.push_str(" let _ = value;\n"); - writeln!( - output, - " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" - ) - .unwrap(); - output.push_str("}\n\n"); - } - } - output -} - -fn render_compat_float(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str("// Included inside `cuda_device::float`.\n\n"); - for record in scalar_arithmetics(catalog) { - let path = record - .rust - .compatibility_paths - .iter() - .find(|path| path.starts_with("cuda_device::float::")) - .expect("scalar-arithmetic compatibility path"); - let ty = scalar_arithmetic_rust_type(record); - let arity = scalar_arithmetic_arity(record); - let parameters = (0..arity) - .map(|index| format!("arg{index}: {ty}")) - .collect::>() - .join(", "); - let arguments = (0..arity) - .map(|index| format!("arg{index}")) - .collect::>() - .join(", "); - writeln!(output, "/// {}", record.summary).unwrap(); - output.push_str("#[must_use]\n#[inline(never)]\n"); - writeln!( - output, - "pub fn {}({parameters}) -> {ty} {{", - record.rust.name - ) - .unwrap(); - writeln!(output, " let _ = ({arguments});").unwrap(); - writeln!( - output, - " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" - ) - .unwrap(); - output.push_str("}\n\n"); - } - for record in scalar_maths(catalog) { - let path = record - .rust - .compatibility_paths - .iter() - .find(|path| path.starts_with("cuda_device::float::")) - .expect("scalar-math compatibility path"); - let ty = match scalar_math_contract(record).format { - ScalarMathFormat::F16 => "u16", - ScalarMathFormat::F32 => "f32", - ScalarMathFormat::F64 => "f64", - }; - writeln!(output, "/// {}", record.summary).unwrap(); - output.push_str("#[must_use]\n#[inline(never)]\n"); - writeln!(output, "pub fn {}(arg0: {ty}) -> {ty} {{", record.rust.name).unwrap(); - writeln!(output, " let _ = arg0;").unwrap(); - writeln!( - output, - " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" - ) - .unwrap(); - output.push_str("}\n\n"); - } - render_compat_extended_minmax_into(&mut output, catalog, "float"); - output -} - -fn render_compat_float_output(catalog: &CatalogFile, hash: &str) -> Option<(PathBuf, String)> { - (scalar_arithmetics(catalog).next().is_some() - || extended_minmax(catalog).next().is_some() - || scalar_maths(catalog).next().is_some()) - .then(|| { - ( - "crates/cuda-device/src/generated/float.rs".into(), - render_compat_float(catalog, hash), - ) - }) -} - -fn render_dialect_mod(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - "mod active_mask;\nmod cluster_barrier;\nmod cp_async;\nmod dotprod;\nmod ldmatrix;\nmod mbarrier_basic;\nmod movmatrix;\nmod packed_alu;\nmod packed_atomic;\nmod packed_conversion;\nmod prmt;\nmod redux;\nmod register_mma;\nmod sparse_mma;\nmod sreg;\nmod sync;\nmod vote;\nmod warp_barrier;\nmod warp_match;\nmod warp_shuffle;\n\npub use active_mask::*;\npub use cluster_barrier::*;\npub use cp_async::*;\npub use dotprod::*;\npub use ldmatrix::*;\npub use mbarrier_basic::*;\npub use movmatrix::*;\npub use packed_alu::*;\npub use packed_atomic::*;\npub use packed_conversion::*;\npub use prmt::*;\npub use redux::*;\npub use register_mma::*;\npub use sparse_mma::*;\npub use sreg::*;\npub use sync::*;\npub use vote::*;\npub use warp_barrier::*;\npub use warp_match::*;\npub use warp_shuffle::*;\n\nuse pliron::context::Context;\n\npub(super) fn register(ctx: &mut Context) {\n active_mask::register(ctx);\n cluster_barrier::register(ctx);\n cp_async::register(ctx);\n dotprod::register(ctx);\n ldmatrix::register(ctx);\n mbarrier_basic::register(ctx);\n movmatrix::register(ctx);\n packed_alu::register(ctx);\n packed_atomic::register(ctx);\n packed_conversion::register(ctx);\n prmt::register(ctx);\n redux::register(ctx);\n register_mma::register(ctx);\n sparse_mma::register(ctx);\n sreg::register(ctx);\n sync::register(ctx);\n vote::register(ctx);\n warp_barrier::register(ctx);\n warp_match::register(ctx);\n warp_shuffle::register(ctx);\n}\n", - ); - if scalar_conversions(catalog).next().is_some() { - output = output - .replace("mod redux;", "mod redux;\nmod scalar_conversion;") - .replace( - "pub use redux::*;", - "pub use redux::*;\npub use scalar_conversion::*;", - ) - .replace( - " redux::register(ctx);", - " redux::register(ctx);\n scalar_conversion::register(ctx);", - ); - } - if scalar_arithmetics(catalog).next().is_some() { - output = output - .replace("mod redux;", "mod redux;\nmod scalar_arithmetic;") - .replace( - "pub use redux::*;", - "pub use redux::*;\npub use scalar_arithmetic::*;", - ) - .replace( - " redux::register(ctx);", - " redux::register(ctx);\n scalar_arithmetic::register(ctx);", - ); - } - if scalar_maths(catalog).next().is_some() { - output = output - .replace("mod redux;", "mod redux;\nmod scalar_math;") - .replace( - "pub use redux::*;", - "pub use redux::*;\npub use scalar_math::*;", - ) - .replace( - " redux::register(ctx);", - " redux::register(ctx);\n scalar_math::register(ctx);", - ); - } - if integer_minmaxes(catalog).next().is_some() { - output = output - .replace("mod ldmatrix;", "mod integer_minmax;\nmod ldmatrix;") - .replace( - "pub use ldmatrix::*;", - "pub use integer_minmax::*;\npub use ldmatrix::*;", - ) - .replace( - " ldmatrix::register(ctx);", - " integer_minmax::register(ctx);\n ldmatrix::register(ctx);", - ); - } - if extended_minmax(catalog).next().is_some() { - output = output - .replace("mod dotprod;", "mod dotprod;\nmod extended_minmax;") - .replace( - "pub use dotprod::*;", - "pub use dotprod::*;\npub use extended_minmax::*;", - ) - .replace( - " dotprod::register(ctx);", - " dotprod::register(ctx);\n extended_minmax::register(ctx);", - ); - } - if debug_controls(catalog).next().is_some() { - output = output - .replace("mod dotprod;", "mod debug_control;\nmod dotprod;") - .replace( - "pub use dotprod::*;", - "pub use debug_control::*;\npub use dotprod::*;", - ) - .replace( - " dotprod::register(ctx);", - " debug_control::register(ctx);\n dotprod::register(ctx);", - ); - } - if cluster_memory(catalog).next().is_some() { - output = output - .replace( - "mod cluster_barrier;", - "mod cluster_barrier;\nmod cluster_memory;", - ) - .replace( - "pub use cluster_barrier::*;", - "pub use cluster_barrier::*;\npub use cluster_memory::*;", - ) - .replace( - " cluster_barrier::register(ctx);", - " cluster_barrier::register(ctx);\n cluster_memory::register(ctx);", - ); - } - if stmatrices(catalog).next().is_some() { - output = output - .replace("mod sync;", "mod stmatrix;\nmod sync;") - .replace("pub use sync::*;", "pub use stmatrix::*;\npub use sync::*;") - .replace( - " sync::register(ctx);", - " stmatrix::register(ctx);\n sync::register(ctx);", - ); - } - if movmatrix(catalog).next().is_none() { - output = output.replace("mod movmatrix;\n", ""); - output = output.replace("pub use movmatrix::*;\n", ""); - output = output.replace(" movmatrix::register(ctx);\n", ""); - } - if clc_intrinsics(catalog).next().is_some() { - output = output - .replace("mod cluster_barrier;", "mod clc;\nmod cluster_barrier;") - .replace( - "pub use cluster_barrier::*;", - "pub use clc::*;\npub use cluster_barrier::*;", - ) - .replace( - " cluster_barrier::register(ctx);", - " clc::register(ctx);\n cluster_barrier::register(ctx);", - ); - } - if elect_intrinsics(catalog).next().is_some() { - output = output - .replace("mod dotprod;", "mod dotprod;\nmod elect;") - .replace( - "pub use dotprod::*;", - "pub use dotprod::*;\npub use elect::*;", - ) - .replace( - " dotprod::register(ctx);", - " dotprod::register(ctx);\n elect::register(ctx);", - ); - } - if mbarrier_extended(catalog).next().is_some() { - output = output.replace( - "mod mbarrier_basic;\n", - "mod mbarrier_basic;\nmod mbarrier_extended;\n", - ); - output = output.replace( - "pub use mbarrier_basic::*;\n", - "pub use mbarrier_basic::*;\npub use mbarrier_extended::*;\n", - ); - output = output.replace( - " mbarrier_basic::register(ctx);\n", - " mbarrier_basic::register(ctx);\n mbarrier_extended::register(ctx);\n", - ); - } - if wgmma_controls(catalog).next().is_some() { - output = output - .replace("mod warp_shuffle;", "mod warp_shuffle;\nmod wgmma_control;") - .replace( - "pub use warp_shuffle::*;", - "pub use warp_shuffle::*;\npub use wgmma_control::*;", - ) - .replace( - " warp_shuffle::register(ctx);", - " warp_shuffle::register(ctx);\n wgmma_control::register(ctx);", - ); - } - if execution_controls(catalog).next().is_some() { - output = output - .replace("mod dotprod;", "mod dotprod;\nmod execution_control;") - .replace( - "pub use dotprod::*;", - "pub use dotprod::*;\npub use execution_control::*;", - ) - .replace( - " dotprod::register(ctx);", - " dotprod::register(ctx);\n execution_control::register(ctx);", - ); - } - if tma_intrinsics(catalog).next().is_some() { - output = output - .replace("mod sync;", "mod sync;\nmod tma;") - .replace("pub use sync::*;", "pub use sync::*;\npub use tma::*;") - .replace( - " sync::register(ctx);", - " sync::register(ctx);\n tma::register(ctx);", - ); - } - if tcgen05_intrinsics(catalog).next().is_some() { - output = output - .replace("mod sync;", "mod sync;\nmod tcgen05;") - .replace("pub use sync::*;", "pub use sync::*;\npub use tcgen05::*;") - .replace( - " sync::register(ctx);", - " sync::register(ctx);\n tcgen05::register(ctx);", - ); - } - output -} - -fn render_dialect_elect(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - "//! Generated warp leader-election operation.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::IntegerType,\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::{TypeHandle, Typed},\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\n", - ); - for record in elect_intrinsics(catalog) { - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!( - output, - "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<1>, NResultsInterface<2>],\n)]", - record.dialect.op_name - ) - .unwrap(); - writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); - writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); - output.push_str(" pub fn new(op: Ptr) -> Self { Self { op } }\n}\n\n"); - writeln!(output, "impl Verify for {} {{", record.dialect.op_type).unwrap(); - output.push_str( - " fn verify(&self, ctx: &Context) -> Result<(), Error> {\n let op = self.get_operation().deref(ctx);\n if op.get_num_operands() != 1 || op.get_num_results() != 2 {\n return verify_err!(op.loc(), \"nvvm.elect_sync requires one i32 operand and results [i32, i1]\");\n }\n if !is_integer_width(ctx, op.get_operand(0).get_type(ctx), 32)\n || !is_integer_width(ctx, op.get_result(0).get_type(ctx), 32)\n || !is_integer_width(ctx, op.get_result(1).get_type(ctx), 1)\n {\n return verify_err!(op.loc(), \"nvvm.elect_sync requires one i32 operand and results [i32, i1]\");\n }\n Ok(())\n }\n}\n\n", - ); - } - output.push_str( - "fn is_integer_width(ctx: &Context, ty: TypeHandle, width: u32) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == width)\n}\n\npub(super) fn register(ctx: &mut Context) {\n", - ); - for record in elect_intrinsics(catalog) { - writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); - } - output.push_str("}\n"); - output -} - -fn render_dialect_sreg(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - "//! Structural NVVM operations for generated special-register reads.\n\nuse pliron::{\n builtin::op_interfaces::{NOpdsInterface, NResultsInterface},\n builtin::types::IntegerType,\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\n", - ); - for record in sregs(catalog) { - let width = record.scalar_width().unwrap(); - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!( - output, - "/// Catalog ID `{}`; {} returns one `i{width}` result.", - record.id, - source_label(record) - ) - .unwrap(); - writeln!( - output, - "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<0>, NResultsInterface<1>],\n)]", - record.dialect.op_name - ) - .unwrap(); - writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); - writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); - output.push_str(" pub fn new(op: Ptr) -> Self {\n"); - writeln!(output, " Self {{ op }}").unwrap(); - output.push_str(" }\n}\n"); - writeln!(output, "\nimpl Verify for {} {{", record.dialect.op_type).unwrap(); - output.push_str(" fn verify(&self, ctx: &Context) -> Result<(), Error> {\n"); - writeln!( - output, - " verify_scalar_result(ctx, self.get_operation(), {:?}, {width})", - record.dialect.op_name - ) - .unwrap(); - output.push_str(" }\n}\n\n"); - } - output.push_str( - "fn verify_scalar_result(\n ctx: &Context,\n op: Ptr,\n name: &str,\n width: u32,\n) -> Result<(), Error> {\n let op = op.deref(ctx);\n let ty = op.get_result(0).get_type(ctx);\n let ty_object = ty.deref(ctx);\n let Some(integer) = ty_object.downcast_ref::() else {\n return verify_err!(op.loc(), \"{} result must be an integer\", name);\n };\n if integer.width() != width {\n return verify_err!(\n op.loc(),\n \"{} result must be a {}-bit integer\",\n name,\n width\n );\n }\n Ok(())\n}\n\npub(super) fn register(ctx: &mut Context) {\n", - ); - for record in sregs(catalog) { - writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); - } - output.push_str("}\n"); - output -} - -fn render_dialect_sync(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - "//! Structural operation for generated CTA synchronization.\n\nuse pliron::{\n builtin::op_interfaces::{NOpdsInterface, NResultsInterface},\n context::{Context, Ptr},\n op::Op,\n operation::Operation,\n};\nuse pliron_derive::pliron_op;\n\n", - ); - for record in sync_intrinsics(catalog) { - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!( - output, - "#[pliron_op(\n name = {:?},\n format,\n verifier = \"succ\",\n interfaces = [NOpdsInterface<0>, NResultsInterface<0>],\n)]", - record.dialect.op_name - ) - .unwrap(); - writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); - writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); - output.push_str( - " pub fn new(op: Ptr) -> Self {\n Self { op }\n }\n}\n\n", - ); - } - output.push_str("pub(super) fn register(ctx: &mut Context) {\n"); - for record in sync_intrinsics(catalog) { - writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); - } - output.push_str("}\n"); - output -} - -fn render_dialect_vote(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - "//! Structural operations for the generated `vote.sync` family.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_integer_width(ctx: &Context, ty: pliron::r#type::TypeHandle, width: u32) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == width)\n}\n\n", - ); - for record in vote_intrinsics(catalog) { - let result_width = match record.vote.as_ref().unwrap().mode { - VoteMode::All | VoteMode::Any | VoteMode::Uni => 1, - VoteMode::Ballot => 32, - }; - let result_signedness = if result_width == 32 { - "Unsigned" - } else { - "Signless" - }; - writeln!(output, "/// {}", record.summary).unwrap(); - output.push_str( - "///\n/// Operands are `[member_mask, predicate]`. The generated verifier keeps\n/// the mask, predicate, and result types exact.\n", - ); - writeln!( - output, - "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<2>, NResultsInterface<1>],\n)]", - record.dialect.op_name - ) - .unwrap(); - writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); - writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); - writeln!( - output, - " pub fn new(op: Ptr) -> Self {{\n Self {{ op }}\n }}\n\n pub fn build(ctx: &mut Context, member_mask: Value, predicate: Value) -> Ptr {{\n let result_ty = IntegerType::get(ctx, {result_width}, Signedness::{result_signedness});\n Operation::new(\n ctx,\n Self::get_concrete_op_info(),\n vec![result_ty.into()],\n vec![member_mask, predicate],\n vec![],\n 0,\n )\n }}\n}}" - ) - .unwrap(); - writeln!(output, "\nimpl Verify for {} {{", record.dialect.op_type).unwrap(); - writeln!( - output, - " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n let op = self.get_operation().deref(ctx);\n if op.get_num_operands() != 2 || op.get_num_results() != 1 {{\n return verify_err!(op.loc(), {:?});\n }}\n if !is_integer_width(ctx, op.get_operand(0).get_type(ctx), 32)\n || !is_integer_width(ctx, op.get_operand(1).get_type(ctx), 1)\n || !is_integer_width(ctx, op.get_result(0).get_type(ctx), {result_width})\n {{\n return verify_err!(op.loc(), {:?});\n }}\n Ok(())\n }}\n}}\n", - format!( - "{} requires exactly two operands [member_mask, predicate] and one result", - record.dialect.op_name - ), - format!( - "{} requires i32 member mask, i1 predicate, and i{result_width} result", - record.dialect.op_name - ), - ) - .unwrap(); - } - output.push_str("\npub(super) fn register(ctx: &mut Context) {\n"); - for record in vote_intrinsics(catalog) { - writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); - } - output.push_str("}\n"); - output -} - -fn render_dialect_active_mask(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - "//! Structural operation for the generated active warp mask.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\n", - ); - for record in active_masks(catalog) { - debug_assert_eq!( - record.active_mask.as_ref().unwrap().adapter, - ActiveMaskAdapter::DirectZeroOperandMask - ); - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!( - output, - "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<0>, NResultsInterface<1>],\n)]", - record.dialect.op_name - ) - .unwrap(); - writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); - writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); - output.push_str( - " pub fn new(op: Ptr) -> Self {\n Self { op }\n }\n\n pub fn build(ctx: &mut Context) -> Ptr {\n let result_ty = IntegerType::get(ctx, 32, Signedness::Unsigned);\n Operation::new(\n ctx,\n Self::get_concrete_op_info(),\n vec![result_ty.into()],\n vec![],\n vec![],\n 0,\n )\n }\n}\n", - ); - writeln!(output, "\nimpl Verify for {} {{", record.dialect.op_type).unwrap(); - writeln!( - output, - " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n let op = self.get_operation().deref(ctx);\n let valid = op.get_num_operands() == 0\n && op.get_num_results() == 1\n && op.get_result(0).get_type(ctx).deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32);\n if !valid {{\n return verify_err!(op.loc(), {:?});\n }}\n Ok(())\n }}\n}}\n", - format!("{} requires no operands and one i32 result", record.dialect.op_name) - ) - .unwrap(); - } - output.push_str("\npub(super) fn register(ctx: &mut Context) {\n"); - for record in active_masks(catalog) { - writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); - } - output.push_str("}\n"); - output -} - -fn render_dialect_cp_async_copy(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - r#"//! Structural operations for classic `cp.async` instructions. - -use dialect_mir::{ops::MirConstantOp, types::{MirPtrType, address_space}}; -use pliron::{ - builtin::{ - op_interfaces::{NOpdsInterface, NResultsInterface}, - ops::ConstantOp, - types::{IntegerType, Signedness}, - }, - 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; - -fn verify_pointer( - ctx: &Context, - op: &Operation, - value: Value, - role: &str, - allowed_address_spaces: &[u32], -) -> Result<(), Error> { - let ty = value.get_type(ctx); - let ty = ty.deref(ctx); - let Some(pointer) = ty.downcast_ref::() else { - return verify_err!(op.loc(), "{role} must be a MIR pointer"); - }; - if !allowed_address_spaces.contains(&pointer.address_space) { - return verify_err!(op.loc(), "{role} has the wrong address space"); - } - Ok(()) -} - -fn verify_cp_async_copy( - ctx: &Context, - operation: Ptr, - name: &str, - has_source_size: bool, -) -> Result<(), Error> { - let op = operation.deref(ctx); - let expected_operands = if has_source_size { 3 } else { 2 }; - if op.get_num_operands() != expected_operands || op.get_num_results() != 0 { - return verify_err!(op.loc(), "{name} has the wrong operand or result count"); - } - verify_pointer( - ctx, - &op, - op.get_operand(0), - "shared destination", - &[address_space::GENERIC, address_space::SHARED], - )?; - verify_pointer( - ctx, - &op, - op.get_operand(1), - "global source", - &[address_space::GENERIC, address_space::GLOBAL], - )?; - if has_source_size { - let ty = op.get_operand(2).get_type(ctx); - let ty = ty.deref(ctx); - let Some(integer) = ty.downcast_ref::() else { - return verify_err!(op.loc(), "source size must be u32"); - }; - if integer.width() != 32 || integer.signedness() != Signedness::Unsigned { - return verify_err!(op.loc(), "source size must be u32"); - } - } - Ok(()) -} - -fn verify_cp_async_control( - ctx: &Context, - operation: Ptr, - name: &str, - has_immediate: bool, -) -> Result<(), Error> { - let op = operation.deref(ctx); - let expected_operands = usize::from(has_immediate); - if op.get_num_operands() != expected_operands || op.get_num_results() != 0 { - return verify_err!(op.loc(), "{name} has the wrong operand or result count"); - } - if has_immediate { - let value = op.get_operand(0); - let ty = value.get_type(ctx); - let ty = ty.deref(ctx); - let Some(integer) = ty.downcast_ref::() else { - return verify_err!(op.loc(), "maximum pending group count must be u32"); - }; - if integer.width() != 32 || integer.signedness() != Signedness::Unsigned { - return verify_err!(op.loc(), "maximum pending group count must be u32"); - } - let Some(defining_op) = value.defining_op() else { - return verify_err!(op.loc(), "maximum pending group count must be a compile-time constant"); - }; - if Operation::get_op::(defining_op, ctx).is_none() - && Operation::get_op::(defining_op, ctx).is_none() - { - return verify_err!(op.loc(), "maximum pending group count must be a compile-time constant"); - } - } - Ok(()) -} - -fn verify_cp_async_mbarrier( - ctx: &Context, - operation: Ptr, - name: &str, -) -> Result<(), Error> { - let op = operation.deref(ctx); - if op.get_num_operands() != 1 || op.get_num_results() != 0 { - return verify_err!(op.loc(), "{name} has the wrong operand or result count"); - } - let ty = op.get_operand(0).get_type(ctx); - let ty = ty.deref(ctx); - let Some(pointer) = ty.downcast_ref::() else { - return verify_err!(op.loc(), "mbarrier address must be a MIR pointer"); - }; - let pointee = pointer.pointee.deref(ctx); - let valid_pointee = pointee - .downcast_ref::() - .is_some_and(|integer| { - integer.width() == 64 && integer.signedness() == Signedness::Unsigned - }); - if !pointer.is_mutable - || !valid_pointee - || !matches!( - pointer.address_space, - address_space::GENERIC | address_space::SHARED - ) - { - return verify_err!( - op.loc(), - "mbarrier address must be a mutable generic/shared pointer to u64" - ); - } - Ok(()) -} - -"#, - ); - for record in cp_async_copies(catalog) { - let copy = record.cp_async_copy.as_ref().unwrap(); - let dynamic = copy.source_size == CpAsyncSourceSize::Runtime; - let operand_count = if dynamic { 3 } else { 2 }; - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!(output, "/// Lowers to `{}`.", record.expected_ptx).unwrap(); - writeln!( - output, - "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<{operand_count}>, NResultsInterface<0>],\n)]\npub struct {};", - record.dialect.op_name, record.dialect.op_type - ) - .unwrap(); - writeln!(output, "impl {} {{", record.dialect.op_type).unwrap(); - output.push_str(" pub fn new(op: Ptr) -> Self { Self { op } }\n\n"); - if dynamic { - output.push_str( - " pub fn build(ctx: &mut Context, shared_dst: Value, global_src: Value, source_size: Value) -> Ptr {\n Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![shared_dst, global_src, source_size], vec![], 0)\n }\n", - ); - } else { - output.push_str( - " pub fn build(ctx: &mut Context, shared_dst: Value, global_src: Value) -> Ptr {\n Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![shared_dst, global_src], vec![], 0)\n }\n", - ); - } - output.push_str("}\n\n"); - writeln!(output, "impl Verify for {} {{", record.dialect.op_type).unwrap(); - writeln!( - output, - " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n verify_cp_async_copy(ctx, self.get_operation(), {:?}, {dynamic})\n }}\n}}\n", - record.dialect.op_name - ) - .unwrap(); - } - for record in cp_async_controls(catalog) { - let control = record.cp_async_control.as_ref().unwrap(); - let has_immediate = control.operation == CpAsyncControlOperation::WaitGroup; - let operand_count = usize::from(has_immediate); - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!(output, "/// Lowers to `{}`.", record.expected_ptx).unwrap(); - writeln!( - output, - "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<{operand_count}>, NResultsInterface<0>],\n)]\npub struct {};", - record.dialect.op_name, record.dialect.op_type - ) - .unwrap(); - writeln!(output, "impl {} {{", record.dialect.op_type).unwrap(); - output.push_str(" pub fn new(op: Ptr) -> Self { Self { op } }\n\n"); - if has_immediate { - output.push_str( - " pub fn build(ctx: &mut Context, max_pending: Value) -> Ptr {\n Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![max_pending], vec![], 0)\n }\n", - ); - } else { - output.push_str( - " pub fn build(ctx: &mut Context) -> Ptr {\n Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![], vec![], 0)\n }\n", - ); - } - output.push_str("}\n\n"); - writeln!(output, "impl Verify for {} {{", record.dialect.op_type).unwrap(); - writeln!( - output, - " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n verify_cp_async_control(ctx, self.get_operation(), {:?}, {has_immediate})\n }}\n}}\n", - record.dialect.op_name - ) - .unwrap(); - } - for record in cp_async_mbarriers(catalog) { - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!(output, "/// Lowers to `{}`.", record.expected_ptx).unwrap(); - writeln!( - output, - "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<1>, NResultsInterface<0>],\n)]\npub struct {};", - record.dialect.op_name, record.dialect.op_type - ) - .unwrap(); - writeln!(output, "impl {} {{", record.dialect.op_type).unwrap(); - output.push_str( - " pub fn new(op: Ptr) -> Self { Self { op } }\n\n pub fn build(ctx: &mut Context, barrier: Value) -> Ptr {\n Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![barrier], vec![], 0)\n }\n}\n\n", - ); - writeln!(output, "impl Verify for {} {{", record.dialect.op_type).unwrap(); - writeln!( - output, - " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n verify_cp_async_mbarrier(ctx, self.get_operation(), {:?})\n }}\n}}\n", - record.dialect.op_name - ) - .unwrap(); - } - output.push_str("pub(super) fn register(ctx: &mut Context) {\n"); - for record in cp_async_copies(catalog) { - writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); - } - for record in cp_async_controls(catalog) { - writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); - } - for record in cp_async_mbarriers(catalog) { - writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); - } - output.push_str("}\n"); - output -} - -fn render_dialect_mbarrier_basic(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - r#"//! Structural operations for the basic shared-memory mbarrier lifecycle. - -use dialect_mir::types::{MirPtrType, address_space}; -use pliron::{ - builtin::{ - op_interfaces::{NOpdsInterface, NResultsInterface}, - types::{IntegerType, Signedness}, - }, - 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; - -#[derive(Clone, Copy)] -enum MbarrierBasicShape { - Init, - Arrive, - TestWait, - Inval, -} - -fn is_integer( - ctx: &Context, - value: Value, - width: u32, - signedness: Signedness, -) -> bool { - value - .get_type(ctx) - .deref(ctx) - .downcast_ref::() - .is_some_and(|integer| { - integer.width() == width && integer.signedness() == signedness - }) -} - -fn verify_barrier_pointer(ctx: &Context, op: &Operation, value: Value) -> Result<(), Error> { - let ty = value.get_type(ctx); - let ty = ty.deref(ctx); - let Some(pointer) = ty.downcast_ref::() else { - return verify_err!(op.loc(), "mbarrier address must be a MIR pointer"); - }; - if !matches!(pointer.address_space, address_space::GENERIC | address_space::SHARED) { - return verify_err!(op.loc(), "mbarrier address must be generic or shared"); - } - Ok(()) -} - -fn verify_mbarrier_basic( - ctx: &Context, - operation: Ptr, - name: &str, - shape: MbarrierBasicShape, -) -> Result<(), Error> { - let op = operation.deref(ctx); - let (operands, results) = match shape { - MbarrierBasicShape::Init | MbarrierBasicShape::TestWait => (2, usize::from(matches!(shape, MbarrierBasicShape::TestWait))), - MbarrierBasicShape::Arrive => (1, 1), - MbarrierBasicShape::Inval => (1, 0), - }; - if op.get_num_operands() != operands || op.get_num_results() != results { - return verify_err!(op.loc(), "{name} has the wrong operand or result count"); - } - verify_barrier_pointer(ctx, &op, op.get_operand(0))?; - match shape { - MbarrierBasicShape::Init => { - if !is_integer(ctx, op.get_operand(1), 32, Signedness::Unsigned) { - return verify_err!(op.loc(), "mbarrier expected count must be u32"); - } - } - MbarrierBasicShape::Arrive => { - if !is_integer(ctx, op.get_result(0), 64, Signedness::Unsigned) { - return verify_err!(op.loc(), "mbarrier arrival token must be u64"); - } - } - MbarrierBasicShape::TestWait => { - if !is_integer(ctx, op.get_operand(1), 64, Signedness::Unsigned) - || !is_integer(ctx, op.get_result(0), 1, Signedness::Signless) - { - return verify_err!(op.loc(), "mbarrier test-wait requires a u64 token and i1 result"); - } - } - MbarrierBasicShape::Inval => {} - } - Ok(()) -} - -"#, - ); - for record in mbarrier_basics(catalog) { - let mbarrier = record.mbarrier_basic.as_ref().unwrap(); - debug_assert_eq!(mbarrier.state_space, MbarrierStateSpace::Shared); - let (operand_count, result_count, shape, build) = match mbarrier.operation { - MbarrierBasicOperation::Init => { - debug_assert_eq!( - mbarrier.adapter, - MbarrierBasicAdapter::InitPointerCountToVoid - ); - ( - 2, - 0, - "Init", - " pub fn build(ctx: &mut Context, barrier: Value, expected_count: Value) -> Ptr {\n Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![barrier, expected_count], vec![], 0)\n }\n", - ) - } - MbarrierBasicOperation::Arrive => { - debug_assert_eq!(mbarrier.adapter, MbarrierBasicAdapter::ArrivePointerToToken); - ( - 1, - 1, - "Arrive", - " pub fn build(ctx: &mut Context, barrier: 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], vec![], 0)\n }\n", - ) - } - MbarrierBasicOperation::TestWait => { - debug_assert_eq!( - mbarrier.adapter, - MbarrierBasicAdapter::TestWaitPointerTokenToPredicate - ); - ( - 2, - 1, - "TestWait", - " pub fn build(ctx: &mut Context, barrier: Value, token: Value) -> Ptr {\n let predicate_ty = IntegerType::get(ctx, 1, Signedness::Signless);\n Operation::new(ctx, Self::get_concrete_op_info(), vec![predicate_ty.into()], vec![barrier, token], vec![], 0)\n }\n", - ) - } - MbarrierBasicOperation::Inval => { - debug_assert_eq!(mbarrier.adapter, MbarrierBasicAdapter::InvalPointerToVoid); - ( - 1, - 0, - "Inval", - " pub fn build(ctx: &mut Context, barrier: Value) -> Ptr {\n Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![barrier], vec![], 0)\n }\n", - ) - } - }; - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!(output, "/// Lowers to `{}`.", record.expected_ptx).unwrap(); - writeln!( - output, - "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<{operand_count}>, NResultsInterface<{result_count}>],\n)]\npub struct {};", - record.dialect.op_name, record.dialect.op_type - ) - .unwrap(); - writeln!(output, "impl {} {{", record.dialect.op_type).unwrap(); - output.push_str(" pub fn new(op: Ptr) -> Self { Self { op } }\n\n"); - output.push_str(build); - output.push_str("}\n\n"); - writeln!(output, "impl Verify for {} {{", record.dialect.op_type).unwrap(); - writeln!( - output, - " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n verify_mbarrier_basic(ctx, self.get_operation(), {:?}, MbarrierBasicShape::{shape})\n }}\n}}\n", - record.dialect.op_name - ) - .unwrap(); - } - output.push_str("pub(super) fn register(ctx: &mut Context) {\n"); - for record in mbarrier_basics(catalog) { - writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); - } - output.push_str("}\n"); - output -} - -fn render_dialect_movmatrix(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - if movmatrix(catalog).next().is_none() { - output.push_str( - "//! Structural operation for generated in-register matrix transpose.\n\nuse pliron::context::Context;\n\npub(super) fn register(_ctx: &mut Context) {}\n", - ); - return output; - } - let record = movmatrix(catalog).next().unwrap(); - output.push_str( - "//! Structural operation for generated in-register matrix transpose.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_i32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32)\n}\n\n", - ); - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!( - output, - "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<1>, NResultsInterface<1>],\n)]\npub struct {};\n", - record.dialect.op_name, record.dialect.op_type - ) - .unwrap(); - writeln!( - output, - "impl {} {{\n pub fn new(op: Ptr) -> Self {{\n Self {{ op }}\n }}\n\n pub fn build(ctx: &mut Context, value: Value) -> Ptr {{\n let result_ty = IntegerType::get(ctx, 32, Signedness::Unsigned);\n Operation::new(\n ctx,\n Self::get_concrete_op_info(),\n vec![result_ty.into()],\n vec![value],\n vec![],\n 0,\n )\n }}\n}}\n", - record.dialect.op_type - ) - .unwrap(); - writeln!( - output, - "impl Verify for {} {{\n fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n let op = self.get_operation().deref(ctx);\n if op.get_num_operands() != 1 || op.get_num_results() != 1 {{\n return verify_err!(op.loc(), {:?});\n }}\n if !is_i32(ctx, op.get_operand(0).get_type(ctx))\n || !is_i32(ctx, op.get_result(0).get_type(ctx))\n {{\n return verify_err!(op.loc(), {:?});\n }}\n Ok(())\n }}\n}}\n\npub(super) fn register(ctx: &mut Context) {{\n {}::register(ctx);\n}}\n", - record.dialect.op_type, - format!("{} requires one operand and one result", record.dialect.op_name), - format!( - "{} operand and result must be 32-bit integers", - record.dialect.op_name - ), - record.dialect.op_type, - ) - .unwrap(); - output -} - -fn render_dialect_mbarrier_extended(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - r#"//! Structural operations for extended mbarrier and async-proxy instructions. - -use dialect_mir::types::{MirPtrType, address_space}; -use pliron::{ - builtin::{ - op_interfaces::{NOpdsInterface, NResultsInterface}, - types::{IntegerType, Signedness}, - }, - 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; - -#[derive(Clone, Copy)] -enum MbarrierExtendedShape { - PointerU32ToU64, - RawU64ToVoid, - PointerU64ToI1, - PointerU32ToI1, - ZeroToVoid, - U32ToVoid, -} - -fn is_integer(ctx: &Context, value: Value, width: u32, signedness: Signedness) -> bool { - value - .get_type(ctx) - .deref(ctx) - .downcast_ref::() - .is_some_and(|integer| { - integer.width() == width && integer.signedness() == signedness - }) -} - -fn verify_barrier_pointer(ctx: &Context, op: &Operation, value: Value) -> Result<(), Error> { - let ty = value.get_type(ctx); - let ty = ty.deref(ctx); - let Some(pointer) = ty.downcast_ref::() else { - return verify_err!(op.loc(), "mbarrier address must be a MIR pointer"); - }; - if !matches!(pointer.address_space, address_space::GENERIC | address_space::SHARED) { - return verify_err!(op.loc(), "mbarrier address must be generic or shared"); - } - Ok(()) -} - -fn verify_mbarrier_extended( - ctx: &Context, - operation: Ptr, - name: &str, - shape: MbarrierExtendedShape, -) -> Result<(), Error> { - let op = operation.deref(ctx); - let (operands, results) = match shape { - MbarrierExtendedShape::PointerU32ToU64 - | MbarrierExtendedShape::PointerU64ToI1 - | MbarrierExtendedShape::PointerU32ToI1 => (2, 1), - MbarrierExtendedShape::RawU64ToVoid | MbarrierExtendedShape::U32ToVoid => (1, 0), - MbarrierExtendedShape::ZeroToVoid => (0, 0), - }; - if op.get_num_operands() != operands || op.get_num_results() != results { - return verify_err!(op.loc(), "{name} has the wrong operand or result count"); - } - match shape { - MbarrierExtendedShape::PointerU32ToU64 => { - verify_barrier_pointer(ctx, &op, op.get_operand(0))?; - if !is_integer(ctx, op.get_operand(1), 32, Signedness::Unsigned) - || !is_integer(ctx, op.get_result(0), 64, Signedness::Unsigned) - { - return verify_err!(op.loc(), "mbarrier arrival requires u32 bytes and a u64 token"); - } - } - MbarrierExtendedShape::PointerU64ToI1 => { - verify_barrier_pointer(ctx, &op, op.get_operand(0))?; - if !is_integer(ctx, op.get_operand(1), 64, Signedness::Unsigned) - || !is_integer(ctx, op.get_result(0), 1, Signedness::Signless) - { - return verify_err!(op.loc(), "mbarrier wait requires a u64 token and i1 result"); - } - } - MbarrierExtendedShape::PointerU32ToI1 => { - verify_barrier_pointer(ctx, &op, op.get_operand(0))?; - if !is_integer(ctx, op.get_operand(1), 32, Signedness::Unsigned) - || !is_integer(ctx, op.get_result(0), 1, Signedness::Signless) - { - return verify_err!(op.loc(), "mbarrier wait requires u32 parity and i1 result"); - } - } - MbarrierExtendedShape::RawU64ToVoid => { - if !is_integer(ctx, op.get_operand(0), 64, Signedness::Unsigned) { - return verify_err!(op.loc(), "remote mbarrier address must be u64"); - } - } - MbarrierExtendedShape::U32ToVoid => { - if !is_integer(ctx, op.get_operand(0), 32, Signedness::Unsigned) { - return verify_err!(op.loc(), "nanosleep duration must be u32"); - } - } - MbarrierExtendedShape::ZeroToVoid => {} - } - Ok(()) -} - -"#, - ); - for record in mbarrier_extended(catalog) { - let contract = record.mbarrier_extended.as_ref().unwrap(); - let (operand_count, result_count, shape, build) = match contract.operation { - MbarrierExtendedOperation::ArriveExpectTxCta - | MbarrierExtendedOperation::ArriveExpectTxCluster => ( - 2, - 1, - "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 => ( - 1, - 0, - "RawU64ToVoid", - " pub fn build(ctx: &mut Context, address: Value) -> Ptr {\n Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![address], vec![], 0)\n }\n", - ), - MbarrierExtendedOperation::TryWaitTokenCta => ( - 2, - 1, - "PointerU64ToI1", - " 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 => ( - 2, - 1, - "PointerU32ToI1", - " pub fn build(ctx: &mut Context, barrier: Value, parity: 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, parity], vec![], 0)\n }\n", - ), - MbarrierExtendedOperation::FenceProxyAsyncSharedCta - | MbarrierExtendedOperation::FenceMbarrierInitReleaseCluster - | MbarrierExtendedOperation::FenceProxyAsyncGenericReleaseSharedCtaCluster - | MbarrierExtendedOperation::FenceProxyAsyncGenericAcquireSharedClusterCluster => ( - 0, - 0, - "ZeroToVoid", - " pub fn build(ctx: &mut Context) -> Ptr {\n Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![], vec![], 0)\n }\n", - ), - MbarrierExtendedOperation::Nanosleep => ( - 1, - 0, - "U32ToVoid", - " pub fn build(ctx: &mut Context, ns: Value) -> Ptr {\n Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![ns], vec![], 0)\n }\n", - ), - }; - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!(output, "/// Lowers to `{}`.", record.expected_ptx).unwrap(); - writeln!( - output, - "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<{operand_count}>, NResultsInterface<{result_count}>],\n)]\npub struct {};", - record.dialect.op_name, record.dialect.op_type - ) - .unwrap(); - writeln!(output, "impl {} {{", record.dialect.op_type).unwrap(); - output.push_str(" pub fn new(op: Ptr) -> Self { Self { op } }\n\n"); - output.push_str(build); - output.push_str("}\n\n"); - writeln!(output, "impl Verify for {} {{", record.dialect.op_type).unwrap(); - writeln!( - output, - " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n verify_mbarrier_extended(ctx, self.get_operation(), {:?}, MbarrierExtendedShape::{shape})\n }}\n}}\n", - record.dialect.op_name - ) - .unwrap(); - } - output.push_str("pub(super) fn register(ctx: &mut Context) {\n"); - for record in mbarrier_extended(catalog) { - writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); - } - output.push_str("}\n"); - output -} - -fn render_dialect_warp_match(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - "//! Structural operations for the generated `match.sync` family.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_integer_width(ctx: &Context, ty: pliron::r#type::TypeHandle, width: u32) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == width)\n}\n\n", - ); - for record in warp_matches(catalog) { - let warp_match = record.warp_match.as_ref().unwrap(); - let value_width = warp_match.value_width.bits(); - writeln!(output, "/// {}", record.summary).unwrap(); - output.push_str( - "///\n/// Operands are `[member_mask, value]`; the result is a 32-bit lane mask.\n", - ); - if warp_match.mode == WarpMatchMode::All { - output.push_str( - "/// LLVM also returns a predicate, which the established API discards.\n", - ); - } - writeln!( - output, - "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<2>, NResultsInterface<1>],\n)]", - record.dialect.op_name - ) - .unwrap(); - writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); - writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); - output.push_str( - " pub fn new(op: Ptr) -> Self {\n Self { op }\n }\n\n pub fn build(ctx: &mut Context, member_mask: Value, value: Value) -> Ptr {\n let result_ty = IntegerType::get(ctx, 32, Signedness::Unsigned);\n Operation::new(\n ctx,\n Self::get_concrete_op_info(),\n vec![result_ty.into()],\n vec![member_mask, value],\n vec![],\n 0,\n )\n }\n}\n", - ); - writeln!(output, "\nimpl Verify for {} {{", record.dialect.op_type).unwrap(); - writeln!( - output, - " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n let op = self.get_operation().deref(ctx);\n if op.get_num_operands() != 2 || op.get_num_results() != 1 {{\n return verify_err!(op.loc(), {:?});\n }}\n if !is_integer_width(ctx, op.get_operand(0).get_type(ctx), 32)\n || !is_integer_width(ctx, op.get_operand(1).get_type(ctx), {value_width})\n || !is_integer_width(ctx, op.get_result(0).get_type(ctx), 32)\n {{\n return verify_err!(op.loc(), {:?});\n }}\n Ok(())\n }}\n}}\n", - format!( - "{} requires exactly [member_mask, value] and one mask result", - record.dialect.op_name - ), - format!( - "{} requires i32 member mask, i{value_width} value, and i32 result", - record.dialect.op_name - ), - ) - .unwrap(); - } - output.push_str("\npub(super) fn register(ctx: &mut Context) {\n"); - for record in warp_matches(catalog) { - writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); - } - output.push_str("}\n"); - output -} - -fn render_dialect_warp_barrier(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - "//! Structural operation for generated warp synchronization.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::IntegerType,\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_i32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32)\n}\n\n", - ); - for record in warp_barriers(catalog) { - debug_assert_eq!( - record.warp_barrier.as_ref().unwrap().adapter, - WarpBarrierAdapter::DirectMemberMask - ); - writeln!(output, "/// {}", record.summary).unwrap(); - output.push_str("///\n/// The operand is the 32-bit warp participation mask.\n"); - writeln!( - output, - "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<1>, NResultsInterface<0>],\n)]", - record.dialect.op_name - ) - .unwrap(); - writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); - writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); - output.push_str( - " pub fn new(op: Ptr) -> Self {\n Self { op }\n }\n\n pub fn build(ctx: &mut Context, member_mask: Value) -> Ptr {\n Operation::new(\n ctx,\n Self::get_concrete_op_info(),\n vec![],\n vec![member_mask],\n vec![],\n 0,\n )\n }\n}\n", - ); - writeln!(output, "\nimpl Verify for {} {{", record.dialect.op_type).unwrap(); - writeln!( - output, - " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n let op = self.get_operation().deref(ctx);\n if op.get_num_operands() != 1 || op.get_num_results() != 0 {{\n return verify_err!(op.loc(), {:?});\n }}\n if !is_i32(ctx, op.get_operand(0).get_type(ctx)) {{\n return verify_err!(op.loc(), {:?});\n }}\n Ok(())\n }}\n}}\n", - format!( - "{} requires exactly one member-mask operand and no results", - record.dialect.op_name - ), - format!("{} member mask must be i32", record.dialect.op_name), - ) - .unwrap(); - } - output.push_str("\npub(super) fn register(ctx: &mut Context) {\n"); - for record in warp_barriers(catalog) { - writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); - } - output.push_str("}\n"); - output -} - -fn render_dialect_warp_shuffle(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - "//! Structural operations for the generated `shfl.sync` family.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{FP32Type, IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_i32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32)\n}\n\nfn is_i64(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 64)\n}\n\nfn is_f32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx).downcast_ref::().is_some()\n}\n\n", - ); - for record in warp_shuffles(catalog) { - let shuffle = record.warp_shuffle.as_ref().unwrap(); - debug_assert!(matches!( - (shuffle.value_kind, shuffle.adapter), - ( - WarpShuffleValueKind::I32 | WarpShuffleValueKind::F32, - WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp - ) | ( - WarpShuffleValueKind::I64, - WarpShuffleAdapter::MaskValueLaneOrDeltaSplitI64LowHighB32InsertClampReassemble - ) - )); - let (result_builder, value_check, value_label) = match shuffle.value_kind { - WarpShuffleValueKind::I32 => ( - "IntegerType::get(ctx, 32, Signedness::Unsigned).into()", - "is_i32", - "i32", - ), - WarpShuffleValueKind::F32 => ("FP32Type::get(ctx).into()", "is_f32", "f32"), - WarpShuffleValueKind::I64 => ( - "IntegerType::get(ctx, 64, Signedness::Unsigned).into()", - "is_i64", - "i64", - ), - }; - writeln!(output, "/// {}", record.summary).unwrap(); - output.push_str("///\n/// Operands are `[member_mask, value, lane_or_delta]`.\n"); - if shuffle.value_kind == WarpShuffleValueKind::I64 { - output.push_str( - "/// Lowering splits the value into two `b32` halves, shuffles both, and reassembles it.\n", - ); - } else { - output.push_str( - "/// Generated lowering inserts the fixed clamp required by the selected shuffle mode.\n", - ); - } - writeln!( - output, - "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<3>, NResultsInterface<1>],\n)]", - record.dialect.op_name - ) - .unwrap(); - writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); - writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); - writeln!( - output, - " pub fn new(op: Ptr) -> Self {{\n Self {{ op }}\n }}\n\n pub fn build(\n ctx: &mut Context,\n member_mask: Value,\n value: Value,\n lane_or_delta: Value,\n ) -> Ptr {{\n let result_ty: pliron::r#type::TypeHandle = {result_builder};\n Operation::new(\n ctx,\n Self::get_concrete_op_info(),\n vec![result_ty],\n vec![member_mask, value, lane_or_delta],\n vec![],\n 0,\n )\n }}\n}}" - ) - .unwrap(); - writeln!(output, "\nimpl Verify for {} {{", record.dialect.op_type).unwrap(); - writeln!( - output, - " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n let op = self.get_operation().deref(ctx);\n if op.get_num_operands() != 3 || op.get_num_results() != 1 {{\n return verify_err!(op.loc(), {:?});\n }}\n if !is_i32(ctx, op.get_operand(0).get_type(ctx))\n || !{value_check}(ctx, op.get_operand(1).get_type(ctx))\n || !is_i32(ctx, op.get_operand(2).get_type(ctx))\n || !{value_check}(ctx, op.get_result(0).get_type(ctx))\n {{\n return verify_err!(op.loc(), {:?});\n }}\n Ok(())\n }}\n}}\n", - format!( - "{} requires exactly [member_mask, value, lane_or_delta] and one result", - record.dialect.op_name - ), - format!( - "{} requires i32 mask/lane and {value_label} value/result", - record.dialect.op_name - ), - ) - .unwrap(); - } - output.push_str("\npub(super) fn register(ctx: &mut Context) {\n"); - for record in warp_shuffles(catalog) { - writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); - } - output.push_str("}\n"); - output -} - -fn render_dialect_dotprod(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - "//! Structural operations for generated packed integer dot products.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_i32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32)\n}\n\n", - ); - for record in dot_products(catalog) { - let signedness = match record.rust.result.as_str() { - "i32" => "Signed", - "u32" => "Unsigned", - result => panic!("unsupported dot-product result {result}"), - }; - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!( - output, - "/// Lowers to `{}`.", - dot_product_ptx(record).replace("$0, $1, $2, $3", "%d, %a, %b, %c") - ) - .unwrap(); - writeln!( - output, - "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<3>, NResultsInterface<1>],\n)]", - record.dialect.op_name - ) - .unwrap(); - writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); - writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); - writeln!( - output, - " pub fn new(op: Ptr) -> Self {{\n Self {{ op }}\n }}\n\n pub fn build(ctx: &mut Context, a: Value, b: Value, c: Value) -> Ptr {{\n let result_ty = IntegerType::get(ctx, 32, Signedness::{signedness});\n Operation::new(\n ctx,\n Self::get_concrete_op_info(),\n vec![result_ty.into()],\n vec![a, b, c],\n vec![],\n 0,\n )\n }}\n}}" - ) - .unwrap(); - writeln!(output, "\nimpl Verify for {} {{", record.dialect.op_type).unwrap(); - writeln!( - output, - " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n let op = self.get_operation().deref(ctx);\n if op.get_num_operands() != 3 || op.get_num_results() != 1 {{\n return verify_err!(op.loc(), {:?});\n }}\n if !(0..3).all(|index| is_i32(ctx, op.get_operand(index).get_type(ctx)))\n || !is_i32(ctx, op.get_result(0).get_type(ctx))\n {{\n return verify_err!(op.loc(), {:?});\n }}\n Ok(())\n }}\n}}\n", - format!( - "{} requires exactly three operands and one result", - record.dialect.op_name - ), - format!( - "{} operands and result must be 32-bit integers", - record.dialect.op_name - ), - ) - .unwrap(); - } - output.push_str("\npub(super) fn register(ctx: &mut Context) {\n"); - for record in dot_products(catalog) { - writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); - } - output.push_str("}\n"); - output -} - -fn render_dialect_redux(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - let has_f32 = redux(catalog).any(|record| record.dialect.results == ["f32"]); - if has_f32 { - output.push_str( - "//! Structural NVVM operations for the closed generated `redux.sync` family.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{FP32Type, IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_i32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32)\n}\n\nfn is_f32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx).downcast_ref::().is_some()\n}\n\n", - ); - } else { - output.push_str( - "//! Structural NVVM operations for the closed generated `redux.sync` family.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_i32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32)\n}\n\n", - ); - } - for record in redux(catalog) { - writeln!(output, "/// {}", record.summary).unwrap(); - output.push_str( - "///\n/// Dialect operands are `[member_mask, value]`; generated lowering adapts\n/// them to LLVM's `(value, member_mask)` signature.\n", - ); - writeln!( - output, - "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<2>, NResultsInterface<1>],\n)]", - record.dialect.op_name - ) - .unwrap(); - writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); - writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); - let (result_type, value_predicate, type_error) = match record.rust.result.as_str() { - "u32" => ( - "IntegerType::get(ctx, 32, Signedness::Unsigned)", - "is_i32", - format!( - "{} member mask, value, and result must be 32-bit integers", - record.dialect.op_name - ), - ), - "i32" => ( - "IntegerType::get(ctx, 32, Signedness::Signed)", - "is_i32", - format!( - "{} member mask, value, and result must be 32-bit integers", - record.dialect.op_name - ), - ), - "f32" => ( - "FP32Type::get(ctx)", - "is_f32", - format!( - "{} member mask must be a 32-bit integer and value and result must be f32", - record.dialect.op_name - ), - ), - result => panic!("unsupported redux result {result}"), - }; - writeln!( - output, - " pub fn new(op: Ptr) -> Self {{\n Self {{ op }}\n }}\n\n pub fn build(ctx: &mut Context, member_mask: Value, value: Value) -> Ptr {{\n let result_ty = {result_type};\n Operation::new(\n ctx,\n Self::get_concrete_op_info(),\n vec![result_ty.into()],\n vec![member_mask, value],\n vec![],\n 0,\n )\n }}\n}}" - ) - .unwrap(); - writeln!(output, "\nimpl Verify for {} {{", record.dialect.op_type).unwrap(); - writeln!( - output, - " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n let op = self.get_operation().deref(ctx);\n if op.get_num_operands() != 2 || op.get_num_results() != 1 {{\n return verify_err!(op.loc(), {:?});\n }}\n if !is_i32(ctx, op.get_operand(0).get_type(ctx))\n || !{value_predicate}(ctx, op.get_operand(1).get_type(ctx))\n || !{value_predicate}(ctx, op.get_result(0).get_type(ctx))\n {{\n return verify_err!(op.loc(), {:?});\n }}\n Ok(())\n }}\n}}\n", - format!( - "{} requires exactly two operands [member_mask, value] and one result", - record.dialect.op_name - ), - type_error, - ) - .unwrap(); - } - output.push_str("\npub(super) fn register(ctx: &mut Context) {\n"); - for record in redux(catalog) { - writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); - } - output.push_str("}\n"); - output -} - -fn render_dialect_ldmatrix(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - r##"//! Structural operation and compatibility carriers for generated `ldmatrix`. - -use dialect_mir::types::{address_space, MirPtrType}; -use pliron::{ - attribute::Attribute, - builtin::{ - op_interfaces::{NOpdsInterface, NResultsInterface}, - types::{IntegerType, Signedness}, - }, - 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_attr, pliron_op}; - -#[pliron_attr(name = "nvvm.ldmatrix_shape", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum LdmatrixShapeAttr { - M8n8, - M8n16, - M16n16, -} - -impl LdmatrixShapeAttr { - pub const fn register_count(&self, multiplicity: &LdmatrixMultiplicityAttr) -> usize { - let matrices = multiplicity.register_count(); - match self { - Self::M8n8 | Self::M8n16 => matrices, - Self::M16n16 => matrices * 2, - } - } -} - -#[pliron_attr(name = "nvvm.ldmatrix_multiplicity", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum LdmatrixMultiplicityAttr { - X1, - X2, - X4, -} - -impl LdmatrixMultiplicityAttr { - pub const fn register_count(&self) -> usize { - match self { - Self::X1 => 1, - Self::X2 => 2, - Self::X4 => 4, - } - } -} - -#[pliron_attr(name = "nvvm.ldmatrix_layout", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum LdmatrixLayoutAttr { - Normal, - Transposed, -} - -#[pliron_attr(name = "nvvm.ldmatrix_element", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum LdmatrixElementAttr { - B16, - B8, - B8x16B4x16P64, - B8x16B6x16P32, -} - -#[pliron_attr(name = "nvvm.ldmatrix_state_space", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum LdmatrixStateSpaceAttr { - Shared, -} - -/// A warp-cooperative matrix load whose exact variant is carried by attributes. -#[pliron_op( - name = "nvvm.ldmatrix", - format, - interfaces = [NOpdsInterface<1>], - attributes = ( - nvvm_ldmatrix_shape: LdmatrixShapeAttr, - nvvm_ldmatrix_multiplicity: LdmatrixMultiplicityAttr, - nvvm_ldmatrix_layout: LdmatrixLayoutAttr, - nvvm_ldmatrix_element: LdmatrixElementAttr, - nvvm_ldmatrix_state_space: LdmatrixStateSpaceAttr - ) -)] -pub struct LdmatrixOp; - -impl LdmatrixOp { - pub fn new(op: Ptr) -> Self { - Self { op } - } - - pub fn build( - ctx: &mut Context, - address: Value, - shape: LdmatrixShapeAttr, - multiplicity: LdmatrixMultiplicityAttr, - layout: LdmatrixLayoutAttr, - element: LdmatrixElementAttr, - state_space: LdmatrixStateSpaceAttr, - ) -> Ptr { - let register_count = shape.register_count(&multiplicity); - let u32_ty = IntegerType::get(ctx, 32, Signedness::Unsigned); - let op = Operation::new( - ctx, - Self::get_concrete_op_info(), - vec![u32_ty.into(); register_count], - vec![address], - vec![], - 0, - ); - let this = Self { op }; - this.set_attr_nvvm_ldmatrix_shape(ctx, shape); - this.set_attr_nvvm_ldmatrix_multiplicity(ctx, multiplicity); - this.set_attr_nvvm_ldmatrix_layout(ctx, layout); - this.set_attr_nvvm_ldmatrix_element(ctx, element); - this.set_attr_nvvm_ldmatrix_state_space(ctx, state_space); - this.get_operation() - } -} - -impl Verify for LdmatrixOp { - fn verify(&self, ctx: &Context) -> Result<(), Error> { - let op = self.get_operation().deref(ctx); - let operands: Vec<_> = op.operands().collect(); - if operands.len() != 1 { - return verify_err!(op.loc(), "nvvm.ldmatrix requires exactly one pointer operand"); - } - let pointer = operands[0].get_type(ctx); - let pointer_object = pointer.deref(ctx); - let Some(pointer) = pointer_object.downcast_ref::() else { - return verify_err!(op.loc(), "nvvm.ldmatrix operand must be a MIR pointer"); - }; - if !matches!(pointer.address_space, address_space::GENERIC | address_space::SHARED) { - return verify_err!( - op.loc(), - "nvvm.ldmatrix pointer must be generic (p0) or shared (p3), not address space {}", - pointer.address_space - ); - } - let pointee = pointer.pointee.deref(ctx); - let Some(pointee) = pointee.downcast_ref::() else { - return verify_err!(op.loc(), "nvvm.ldmatrix pointer must point to an unsigned integer"); - }; - - let Some(shape) = self.get_attr_nvvm_ldmatrix_shape(ctx) else { - return verify_err!(op.loc(), "nvvm.ldmatrix requires a shape attribute"); - }; - let Some(multiplicity) = self.get_attr_nvvm_ldmatrix_multiplicity(ctx) else { - return verify_err!(op.loc(), "nvvm.ldmatrix requires a multiplicity attribute"); - }; - let layout = self.get_attr_nvvm_ldmatrix_layout(ctx); - let element = self.get_attr_nvvm_ldmatrix_element(ctx); - let expected_pointee_width = if element.as_deref() == Some(&LdmatrixElementAttr::B16) { - 32 - } else { - 8 - }; - if pointee.width() != expected_pointee_width - || pointee.signedness() != Signedness::Unsigned - { - return verify_err!( - op.loc(), - "nvvm.ldmatrix pointer element width does not match its matrix element format" - ); - } - let supported = matches!( - (&*shape, layout.as_deref(), element.as_deref()), - (LdmatrixShapeAttr::M8n8, Some(_), Some(LdmatrixElementAttr::B16)) - | - ( - LdmatrixShapeAttr::M8n16, - Some(LdmatrixLayoutAttr::Normal), - Some( - LdmatrixElementAttr::B8x16B4x16P64 - | LdmatrixElementAttr::B8x16B6x16P32, - ), - ) - | - ( - LdmatrixShapeAttr::M16n16, - Some(LdmatrixLayoutAttr::Transposed), - Some( - LdmatrixElementAttr::B8 - | LdmatrixElementAttr::B8x16B4x16P64 - | LdmatrixElementAttr::B8x16B6x16P32, - ), - ) - ); - if !supported - || (*shape == LdmatrixShapeAttr::M16n16 - && *multiplicity == LdmatrixMultiplicityAttr::X4) - || self.get_attr_nvvm_ldmatrix_state_space(ctx).as_deref() - != Some(&LdmatrixStateSpaceAttr::Shared) - { - return verify_err!(op.loc(), "nvvm.ldmatrix has a missing or unsupported variant attribute"); - } - - let register_count = shape.register_count(&multiplicity); - if op.get_num_results() != register_count { - return verify_err!( - op.loc(), - "nvvm.ldmatrix {:?} requires {} u32 results", - multiplicity, - register_count - ); - } - for index in 0..register_count { - let ty = op.get_result(index).get_type(ctx); - let ty_object = ty.deref(ctx); - let Some(integer) = ty_object.downcast_ref::() else { - return verify_err!(op.loc(), "nvvm.ldmatrix result {} must be u32", index); - }; - if integer.width() != 32 || integer.signedness() != Signedness::Unsigned { - return verify_err!(op.loc(), "nvvm.ldmatrix result {} must be u32", index); - } - } - Ok(()) - } -} - -pub(super) fn register(ctx: &mut Context) { - LdmatrixShapeAttr::register(ctx); - LdmatrixMultiplicityAttr::register(ctx); - LdmatrixLayoutAttr::register(ctx); - LdmatrixElementAttr::register(ctx); - LdmatrixStateSpaceAttr::register(ctx); - LdmatrixOp::register(ctx); -} -"##, - ); - - let compatibility = ldmatrix(catalog) - .filter_map(|record| ldmatrix_compat_op(record).map(|compat| (record, compat))) - .collect::>(); - assert_eq!(compatibility.len(), 6); - - let mut definitions = String::from( - r#" -fn verify_compat_ldmatrix( - ctx: &Context, - op: Ptr, - op_name: &str, - result_count: usize, -) -> Result<(), Error> { - let op = op.deref(ctx); - let operands: Vec<_> = op.operands().collect(); - if operands.len() != 1 { - return verify_err!(op.loc(), "{} requires one shared-memory pointer", op_name); - } - if operands[0] - .get_type(ctx) - .deref(ctx) - .downcast_ref::() - .is_none() - { - return verify_err!(op.loc(), "{} operand 0 must be a MIR pointer", op_name); - } - if op.get_num_results() != result_count { - return verify_err!( - op.loc(), - "{} requires {} register results", - op_name, - result_count - ); - } - for index in 0..result_count { - let ty = op.get_result(index).get_type(ctx); - let ty = ty.deref(ctx); - let Some(integer) = ty.downcast_ref::() else { - return verify_err!(op.loc(), "{} result {} must be an integer", op_name, index); - }; - if integer.width() != 32 { - return verify_err!(op.loc(), "{} result {} must be 32 bits", op_name, index); - } - } - Ok(()) -} - -"#, - ); - for (record, (op_type, op_name)) in &compatibility { - let result_count = record - .ldmatrix - .as_ref() - .expect("ldmatrix compatibility record") - .variant - .multiplicity - .register_count(); - writeln!( - definitions, - "/// Compatibility carrier for `{op_name}`.\n#[pliron_op(\n name = {op_name:?},\n format,\n interfaces = [NOpdsInterface<1>, NResultsInterface<{result_count}>],\n)]\npub struct {op_type};\n\nimpl {op_type} {{\n pub fn new(op: Ptr) -> Self {{\n Self {{ op }}\n }}\n}}\n\nimpl Verify for {op_type} {{\n fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n verify_compat_ldmatrix(ctx, self.get_operation(), {op_name:?}, {result_count})\n }}\n}}\n" - ) - .unwrap(); - } - let register_start = output - .find("pub(super) fn register(ctx: &mut Context) {") - .expect("ldmatrix register function"); - output.insert_str(register_start, &definitions); - let register_anchor = " LdmatrixOp::register(ctx);\n"; - let mut registrations = String::from(register_anchor); - for (_, (op_type, _)) in compatibility { - writeln!(registrations, " {op_type}::register(ctx);").unwrap(); - } - output = output.replacen(register_anchor, ®istrations, 1); - output -} - -fn render_dialect_stmatrix(catalog: &CatalogFile, hash: &str) -> String { - assert_eq!(stmatrices(catalog).count(), 4); - let mut output = rust_header(catalog, hash); - output.push_str( - r#"//! Structural operations for the four generated `stmatrix` stores. - -use dialect_mir::types::MirPtrType; -use pliron::{ - builtin::{ - op_interfaces::{NOpdsInterface, NResultsInterface}, - types::IntegerType, - }, - common_traits::Verify, - context::{Context, Ptr}, - location::Located, - op::Op, - operation::Operation, - result::Error, - r#type::Typed, - verify_err, -}; -use pliron_derive::pliron_op; - -fn verify_stmatrix_operands( - ctx: &Context, - op: Ptr, - op_name: &str, - register_count: usize, -) -> Result<(), Error> { - let op = op.deref(ctx); - let operands: Vec<_> = op.operands().collect(); - if operands.len() != register_count + 1 { - return verify_err!( - op.loc(), - "{} requires one pointer and {} register operands", - op_name, - register_count - ); - } - if operands[0] - .get_type(ctx) - .deref(ctx) - .downcast_ref::() - .is_none() - { - return verify_err!(op.loc(), "{} operand 0 must be a MIR pointer", op_name); - } - for (index, register) in operands.iter().enumerate().skip(1) { - let ty = register.get_type(ctx); - let ty = ty.deref(ctx); - let Some(integer) = ty.downcast_ref::() else { - return verify_err!( - op.loc(), - "{} register operand {} must be an integer", - op_name, - index - 1 - ); - }; - if integer.width() != 32 { - return verify_err!( - op.loc(), - "{} register operand {} must be 32 bits", - op_name, - index - 1 - ); - } - } - Ok(()) -} - -"#, - ); - for record in stmatrices(catalog) { - let (multiplicity, _) = stmatrix_variant(record).expect("stmatrix variant"); - let count = multiplicity.register_count(); - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!( - output, - "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<{}>, NResultsInterface<0>],\n)]", - record.dialect.op_name, - count + 1 - ) - .unwrap(); - writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); - writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); - output.push_str( - " pub fn new(op: Ptr) -> Self {\n Self { op }\n }\n}\n", - ); - writeln!(output, "\nimpl Verify for {} {{", record.dialect.op_type).unwrap(); - writeln!( - output, - " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n verify_stmatrix_operands(ctx, self.get_operation(), {:?}, {count})\n }}\n}}\n", - record.dialect.op_name - ) - .unwrap(); - } - output.push_str("pub(super) fn register(ctx: &mut Context) {\n"); - for record in stmatrices(catalog) { - writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); - } - output.push_str("}\n"); - output -} - -fn register_mma_carriers(record: &CatalogIntrinsic) -> (&'static str, &'static str) { - match record.register_mma.as_ref().unwrap().adapter { - RegisterMmaAdapter::C2U32A2U32B1U32ToD2U32 => ( - "&[MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32]", - "&[MmaCarrier::U32, MmaCarrier::U32]", - ), - RegisterMmaAdapter::C2U32A4U32B2U32ToD2U32 => ( - "&[MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32]", - "&[MmaCarrier::U32, MmaCarrier::U32]", - ), - RegisterMmaAdapter::C4F32A2U32B1U32ToD4F32 => ( - "&[MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32]", - "&[MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32]", - ), - RegisterMmaAdapter::C4F32A4U32B2U32ToD4F32 => ( - "&[MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32]", - "&[MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32]", - ), - RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 => ( - "&[MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U16, MmaCarrier::U16, MmaCarrier::U32, MmaCarrier::U16, MmaCarrier::U16]", - "&[MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32]", - ), - RegisterMmaAdapter::C2F64A1F64B1F64ToD2F64 => ( - "&[MmaCarrier::F64, MmaCarrier::F64, MmaCarrier::F64, MmaCarrier::F64]", - "&[MmaCarrier::F64, MmaCarrier::F64]", - ), - RegisterMmaAdapter::C2I32A1U32B1U32ToD2I32 => ( - "&[MmaCarrier::I32, MmaCarrier::I32, MmaCarrier::U32, MmaCarrier::U32]", - "&[MmaCarrier::I32, MmaCarrier::I32]", - ), - RegisterMmaAdapter::C4I32A4U32B2U32ToD4I32 => ( - "&[MmaCarrier::I32, MmaCarrier::I32, MmaCarrier::I32, MmaCarrier::I32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32]", - "&[MmaCarrier::I32, MmaCarrier::I32, MmaCarrier::I32, MmaCarrier::I32]", - ), - RegisterMmaAdapter::C4I32A2U32B1U32ToD4I32 => ( - "&[MmaCarrier::I32, MmaCarrier::I32, MmaCarrier::I32, MmaCarrier::I32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32]", - "&[MmaCarrier::I32, MmaCarrier::I32, MmaCarrier::I32, MmaCarrier::I32]", - ), - } -} - -fn register_mma_compat_op_type(record: &CatalogIntrinsic) -> Option<&'static str> { - match record.id.as_str() { - "mma_m16n8k16_f32_bf16" => Some("MmaM16N8K16F32Bf16Op"), - "mma_m16n8k16_f32_f16" => Some("MmaM16N8K16F32F16Op"), - "mma_m16n8k8_f32_tf32" => Some("MmaM16N8K8F32Tf32Op"), - "mma_m16n8k32_s32_s8" => Some("MmaM16N8K32S32S8Op"), - "mma_m8n8k4_f64" => Some("MmaM8N8K4F64Op"), - _ => None, - } -} - -fn render_dialect_register_mma(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - r#"//! One structural operation for generated register-only `mma.sync` variants. - -use pliron::{ - attribute::Attribute, - builtin::{ - op_interfaces::{NOpdsInterface, NResultsInterface}, - types::{FP32Type, FP64Type, IntegerType, Signedness}, - }, - common_traits::Verify, - context::{Context, Ptr}, - location::Located, - op::Op, - operation::Operation, - result::Error, - r#type::{TypeHandle, Typed}, - verify_err, -}; -use pliron_derive::{pliron_attr, pliron_op}; - -#[pliron_attr(name = "nvvm.register_mma_shape", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum RegisterMmaShapeAttr { M8n8k4, M8n8k16, M8n8k32, M8n8k128, M16n8k4, M16n8k8, M16n8k16, M16n8k32, M16n8k64, M16n8k128, M16n8k256 } - -#[pliron_attr(name = "nvvm.register_mma_operation", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum RegisterMmaOperationAttr { Multiply, AndPopc, XorPopc } - -#[pliron_attr(name = "nvvm.register_mma_kind", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum RegisterMmaKindAttr { Standard, F8f6f4, Mxf8f6f4 } - -#[pliron_attr(name = "nvvm.register_mma_accumulator", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum RegisterMmaAccumulatorAttr { F16, F32, F64, S32 } - -#[pliron_attr(name = "nvvm.register_mma_element", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum RegisterMmaElementAttr { Bf16, E2m1, E2m3, E3m2, E4m3, E5m2, F16, Tf32, F64, B1, S4, U4, S8, U8 } - -#[pliron_attr(name = "nvvm.register_mma_layout", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum RegisterMmaLayoutAttr { Row, Col } - -#[pliron_attr(name = "nvvm.register_mma_overflow", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum RegisterMmaOverflowAttr { NotApplicable, Wrapping, Satfinite } - -#[pliron_op( - name = "nvvm.register_mma", - format, - attributes = ( - nvvm_register_mma_shape: RegisterMmaShapeAttr, - nvvm_register_mma_operation: RegisterMmaOperationAttr, - nvvm_register_mma_kind: RegisterMmaKindAttr, - nvvm_register_mma_accumulator: RegisterMmaAccumulatorAttr, - nvvm_register_mma_a_element: RegisterMmaElementAttr, - nvvm_register_mma_b_element: RegisterMmaElementAttr, - nvvm_register_mma_a_layout: RegisterMmaLayoutAttr, - nvvm_register_mma_b_layout: RegisterMmaLayoutAttr, - nvvm_register_mma_overflow: RegisterMmaOverflowAttr - ) -)] -pub struct RegisterMmaOp; - -impl RegisterMmaOp { - pub fn new(op: Ptr) -> Self { Self { op } } - - /// Defaults older multiply-only IR to multiply. - pub fn operation_or_multiply(&self, ctx: &Context) -> RegisterMmaOperationAttr { - self.get_attr_nvvm_register_mma_operation(ctx) - .as_deref() - .cloned() - .unwrap_or(RegisterMmaOperationAttr::Multiply) - } - - /// Infers the kind used by older generated IR. - pub fn kind_or_inferred(&self, ctx: &Context) -> RegisterMmaKindAttr { - if let Some(kind) = self.get_attr_nvvm_register_mma_kind(ctx).as_deref() { - return kind.clone(); - } - let low_format = |element: Option<&RegisterMmaElementAttr>| { - matches!( - element, - Some( - RegisterMmaElementAttr::E2m1 - | RegisterMmaElementAttr::E2m3 - | RegisterMmaElementAttr::E3m2 - | RegisterMmaElementAttr::E4m3 - | RegisterMmaElementAttr::E5m2 - ) - ) - }; - let old_f8f6f4 = self.get_attr_nvvm_register_mma_shape(ctx).as_deref() - == Some(&RegisterMmaShapeAttr::M16n8k32) - && low_format(self.get_attr_nvvm_register_mma_a_element(ctx).as_deref()) - && low_format(self.get_attr_nvvm_register_mma_b_element(ctx).as_deref()); - if old_f8f6f4 { - RegisterMmaKindAttr::F8f6f4 - } else { - RegisterMmaKindAttr::Standard - } - } -} - -#[derive(Clone, Copy)] -enum MmaCarrier { F32, F64, I32, U16, U32 } - -fn is_carrier(ctx: &Context, ty: TypeHandle, carrier: MmaCarrier) -> bool { - match carrier { - MmaCarrier::F32 => ty.deref(ctx).downcast_ref::().is_some(), - MmaCarrier::F64 => ty.deref(ctx).downcast_ref::().is_some(), - MmaCarrier::I32 | MmaCarrier::U16 | MmaCarrier::U32 => { - let expected = if matches!(carrier, MmaCarrier::I32) { - Signedness::Signed - } else { - Signedness::Unsigned - }; - ty.deref(ctx) - .downcast_ref::() - .is_some_and(|integer| { - integer.width() - == if matches!(carrier, MmaCarrier::U16) { 16 } else { 32 } - && integer.signedness() == expected - }) - } - } -} - -impl Verify for RegisterMmaOp { - fn verify(&self, ctx: &Context) -> Result<(), Error> { - let op = self.get_operation().deref(ctx); - let operation = self.operation_or_multiply(ctx); - let kind = self.kind_or_inferred(ctx); - let recipe: (&[MmaCarrier], &[MmaCarrier]) = match ( - self.get_attr_nvvm_register_mma_shape(ctx).as_deref(), - operation, - kind, - self.get_attr_nvvm_register_mma_accumulator(ctx).as_deref(), - self.get_attr_nvvm_register_mma_a_element(ctx).as_deref(), - self.get_attr_nvvm_register_mma_b_element(ctx).as_deref(), - self.get_attr_nvvm_register_mma_a_layout(ctx).as_deref(), - self.get_attr_nvvm_register_mma_b_layout(ctx).as_deref(), - self.get_attr_nvvm_register_mma_overflow(ctx).as_deref(), - ) { -"#, - ); - for record in register_mmas(catalog) { - let ( - shape, - operation, - kind, - accumulator, - a_element, - b_element, - a_layout, - b_layout, - overflow, - ) = register_mma_attr_variants(record); - let (operands, results) = register_mma_carriers(record); - writeln!( - output, - " (Some(&{shape}), {operation}, {kind}, Some(&{accumulator}), Some(&{a_element}), Some(&{b_element}), Some(&{a_layout}), Some(&{b_layout}), Some(&{overflow})) => ({operands}, {results})," - ) - .unwrap(); - } - output.push_str( - r#" _ => return verify_err!(op.loc(), "nvvm.register_mma has a missing or unsupported variant"), - }; - let operands: Vec<_> = op.operands().collect(); - if operands.len() != recipe.0.len() || op.get_num_results() != recipe.1.len() { - return verify_err!(op.loc(), "nvvm.register_mma has the wrong register count"); - } - for (index, (operand, carrier)) in operands.iter().zip(recipe.0).enumerate() { - if !is_carrier(ctx, operand.get_type(ctx), *carrier) { - return verify_err!(op.loc(), "nvvm.register_mma operand {} has the wrong carrier type", index); - } - } - for (index, carrier) in recipe.1.iter().enumerate() { - if !is_carrier(ctx, op.get_result(index).get_type(ctx), *carrier) { - return verify_err!(op.loc(), "nvvm.register_mma result {} has the wrong carrier type", index); - } - } - Ok(()) - } -} - -"#, - ); - output.push_str( - r#" -fn is_compat_carrier(ctx: &Context, ty: TypeHandle, carrier: MmaCarrier) -> bool { - match carrier { - MmaCarrier::F32 => ty.deref(ctx).downcast_ref::().is_some(), - MmaCarrier::F64 => ty.deref(ctx).downcast_ref::().is_some(), - MmaCarrier::I32 | MmaCarrier::U16 | MmaCarrier::U32 => { - let expected_width = if matches!(carrier, MmaCarrier::U16) { - 16 - } else { - 32 - }; - ty.deref(ctx) - .downcast_ref::() - .is_some_and(|integer| integer.width() == expected_width) - } - } -} - -fn verify_compat_register_mma( - ctx: &Context, - op_ptr: Ptr, - op_name: &str, - operand_carriers: &[MmaCarrier], - result_carriers: &[MmaCarrier], -) -> Result<(), Error> { - let op = op_ptr.deref(ctx); - let operands: Vec<_> = op.operands().collect(); - if operands.len() != operand_carriers.len() || op.get_num_results() != result_carriers.len() { - return verify_err!(op.loc(), "{} has the wrong register count", op_name); - } - for (index, (operand, carrier)) in operands.iter().zip(operand_carriers).enumerate() { - if !is_compat_carrier(ctx, operand.get_type(ctx), *carrier) { - return verify_err!(op.loc(), "{} operand {} has the wrong carrier type", op_name, index); - } - } - for (index, carrier) in result_carriers.iter().enumerate() { - if !is_compat_carrier(ctx, op.get_result(index).get_type(ctx), *carrier) { - return verify_err!(op.loc(), "{} result {} has the wrong carrier type", op_name, index); - } - } - Ok(()) -} - -"#, - ); - for record in - register_mmas(catalog).filter(|record| register_mma_compat_op_type(record).is_some()) - { - let op_type = register_mma_compat_op_type(record).unwrap(); - let op_name = format!("nvvm.{}", record.id); - let operand_count = record.dialect.operands.len(); - let result_count = record.dialect.results.len(); - let (operands, results) = register_mma_carriers(record); - writeln!(output, "/// Compatibility carrier for `{}`.", record.id).unwrap(); - writeln!( - output, - "#[pliron_op(\n name = {op_name:?},\n format,\n interfaces = [NOpdsInterface<{operand_count}>, NResultsInterface<{result_count}>],\n)]\npub struct {op_type};\n" - ) - .unwrap(); - writeln!( - output, - "impl {op_type} {{\n pub fn new(op: Ptr) -> Self {{ Self {{ op }} }}\n}}\n" - ) - .unwrap(); - writeln!( - output, - "impl Verify for {op_type} {{\n fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n verify_compat_register_mma(ctx, self.get_operation(), {op_name:?}, {operands}, {results})\n }}\n}}\n" - ) - .unwrap(); - } - output.push_str( - "pub(super) fn register(ctx: &mut Context) {\n RegisterMmaShapeAttr::register(ctx);\n RegisterMmaOperationAttr::register(ctx);\n RegisterMmaKindAttr::register(ctx);\n RegisterMmaAccumulatorAttr::register(ctx);\n RegisterMmaElementAttr::register(ctx);\n RegisterMmaLayoutAttr::register(ctx);\n RegisterMmaOverflowAttr::register(ctx);\n RegisterMmaOp::register(ctx);\n", - ); - for record in - register_mmas(catalog).filter(|record| register_mma_compat_op_type(record).is_some()) - { - writeln!( - output, - " {}::register(ctx);", - register_mma_compat_op_type(record).unwrap() - ) - .unwrap(); - } - output.push_str("}\n"); - output -} - -fn render_dialect_sparse_mma(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - r#"//! One structural operation for generated sparse MMA variants. - -use dialect_mir::ops::MirConstantOp; -use pliron::{ - attribute::Attribute, - builtin::{attributes::IntegerAttr, ops::ConstantOp, types::{FP32Type, IntegerType, Signedness}}, - 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_attr, pliron_op}; - -#[pliron_attr(name = "nvvm.sparse_mma_shape", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum SparseMmaShapeAttr { M16n8k32, M16n8k64, M16n8k128 } - -#[pliron_attr(name = "nvvm.sparse_mma_accumulator", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum SparseMmaAccumulatorAttr { F16, F32, S32 } - -#[pliron_attr(name = "nvvm.sparse_mma_element", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum SparseMmaElementAttr { E2m1, E2m3, E3m2, E4m3, E5m2, S4, U4, S8, U8 } - -#[pliron_attr(name = "nvvm.sparse_mma_layout", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum SparseMmaLayoutAttr { Row, Col } - -#[pliron_attr(name = "nvvm.sparse_mma_overflow", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum SparseMmaOverflowAttr { NotApplicable, Wrapping, Satfinite } - -#[pliron_attr(name = "nvvm.sparse_mma_metadata", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum SparseMmaMetadataAttr { Standard, Ordered } - -#[pliron_attr(name = "nvvm.sparse_mma_selector", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum SparseMmaSelectorAttr { ImmediateZeroOrOne, ImmediateZero } - -#[pliron_op( - name = "nvvm.sparse_mma", - format, - attributes = ( - nvvm_sparse_mma_shape: SparseMmaShapeAttr, - nvvm_sparse_mma_accumulator: SparseMmaAccumulatorAttr, - nvvm_sparse_mma_a_element: SparseMmaElementAttr, - nvvm_sparse_mma_b_element: SparseMmaElementAttr, - nvvm_sparse_mma_a_layout: SparseMmaLayoutAttr, - nvvm_sparse_mma_b_layout: SparseMmaLayoutAttr, - nvvm_sparse_mma_overflow: SparseMmaOverflowAttr, - nvvm_sparse_mma_metadata: SparseMmaMetadataAttr, - nvvm_sparse_mma_selector: SparseMmaSelectorAttr - ) -)] -pub struct SparseMmaOp; - -impl SparseMmaOp { - pub fn new(op: Ptr) -> Self { Self { op } } -} - -#[derive(Clone, Copy)] -enum MmaCarrier { F32, I32, U32 } - -fn is_carrier(ctx: &Context, ty: TypeHandle, carrier: MmaCarrier) -> bool { - match carrier { - MmaCarrier::F32 => ty.deref(ctx).downcast_ref::().is_some(), - MmaCarrier::I32 | MmaCarrier::U32 => { - let expected = if matches!(carrier, MmaCarrier::I32) { - Signedness::Signed - } else { - Signedness::Unsigned - }; - ty.deref(ctx) - .downcast_ref::() - .is_some_and(|integer| integer.width() == 32 && integer.signedness() == expected) - } - } -} - -fn constant_u32(ctx: &Context, value: Value) -> Option { - let defining_op = value.defining_op()?; - if let Some(constant) = Operation::get_op::(defining_op, ctx) { - return constant.get_attr_value(ctx).map(|value| value.value().to_u64()); - } - let constant = Operation::get_op::(defining_op, ctx)?; - constant - .get_value(ctx) - .downcast_ref::() - .map(|value| value.value().to_u64()) -} - -impl Verify for SparseMmaOp { - fn verify(&self, ctx: &Context) -> Result<(), Error> { - let op = self.get_operation().deref(ctx); - let (expected_operands, expected_results, selector_upper_exclusive): (&[MmaCarrier], &[MmaCarrier], usize) = match ( - self.get_attr_nvvm_sparse_mma_shape(ctx).as_deref(), - self.get_attr_nvvm_sparse_mma_accumulator(ctx).as_deref(), - self.get_attr_nvvm_sparse_mma_a_element(ctx).as_deref(), - self.get_attr_nvvm_sparse_mma_b_element(ctx).as_deref(), - self.get_attr_nvvm_sparse_mma_a_layout(ctx).as_deref(), - self.get_attr_nvvm_sparse_mma_b_layout(ctx).as_deref(), - self.get_attr_nvvm_sparse_mma_overflow(ctx).as_deref(), - self.get_attr_nvvm_sparse_mma_metadata(ctx).as_deref(), - self.get_attr_nvvm_sparse_mma_selector(ctx).as_deref(), - ) { -"#, - ); - for record in sparse_mmas(catalog) { - let ( - shape, - accumulator, - a_element, - b_element, - a_layout, - b_layout, - overflow, - metadata, - selector, - ) = sparse_mma_attr_variants(record); - let (expected_operands, expected_results) = sparse_mma_carriers(record); - let selector_upper_exclusive = sparse_mma_selector_values(record).len(); - writeln!( - output, - " (Some(&{shape}), Some(&{accumulator}), Some(&{a_element}), Some(&{b_element}), Some(&{a_layout}), Some(&{b_layout}), Some(&{overflow}), Some(&{metadata}), Some(&{selector})) => ({expected_operands}, {expected_results}, {selector_upper_exclusive})," - ) - .unwrap(); - } - output.push_str( - r#" _ => return verify_err!(op.loc(), "nvvm.sparse_mma has a missing or unsupported variant"), - }; - let operands: Vec<_> = op.operands().collect(); - if operands.len() != expected_operands.len() || op.get_num_results() != expected_results.len() { - return verify_err!(op.loc(), "nvvm.sparse_mma has the wrong register count"); - } - for (index, (operand, carrier)) in operands.iter().zip(expected_operands).enumerate() { - if !is_carrier(ctx, operand.get_type(ctx), *carrier) { - return verify_err!(op.loc(), "nvvm.sparse_mma operand {} has the wrong carrier type", index); - } - } - for (index, carrier) in expected_results.iter().enumerate() { - if !is_carrier(ctx, op.get_result(index).get_type(ctx), *carrier) { - return verify_err!(op.loc(), "nvvm.sparse_mma result {} has the wrong carrier type", index); - } - } - if constant_u32(ctx, operands[expected_operands.len() - 1]) - .is_none_or(|selector| selector >= selector_upper_exclusive as u64) - { - return verify_err!(op.loc(), "nvvm.sparse_mma selector must be a compile-time constant in 0..{}", selector_upper_exclusive); - } - Ok(()) - } -} - -pub(super) fn register(ctx: &mut Context) { - SparseMmaShapeAttr::register(ctx); - SparseMmaAccumulatorAttr::register(ctx); - SparseMmaElementAttr::register(ctx); - SparseMmaLayoutAttr::register(ctx); - SparseMmaOverflowAttr::register(ctx); - SparseMmaMetadataAttr::register(ctx); - SparseMmaSelectorAttr::register(ctx); - SparseMmaOp::register(ctx); -} -"#, - ); - output -} - -fn render_dialect_packed_atomic(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - r##"//! One structural operation for the closed generated packed-atomic family. - -use dialect_mir::types::{address_space, MirPtrType}; -use pliron::{ - attribute::Attribute, - builtin::{ - op_interfaces::{NOpdsInterface, NResultsInterface}, - types::{IntegerType, Signedness}, - }, - 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_attr, pliron_op}; - -#[pliron_attr(name = "nvvm.packed_atomic_format", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum PackedAtomicFormatAttr { F16x2, Bf16x2 } - -#[pliron_attr(name = "nvvm.packed_atomic_state_space", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum PackedAtomicStateSpaceAttr { Global } - -#[pliron_attr(name = "nvvm.packed_atomic_ordering", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum PackedAtomicOrderingAttr { Relaxed } - -#[pliron_attr(name = "nvvm.packed_atomic_scope", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum PackedAtomicScopeAttr { Gpu } - -#[pliron_attr(name = "nvvm.packed_atomic_rounding", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum PackedAtomicRoundingAttr { Rn } - -#[pliron_attr(name = "nvvm.packed_atomic_subnormal", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum PackedAtomicSubnormalAttr { NoFtz } - -#[pliron_attr(name = "nvvm.packed_atomic_atomicity", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum PackedAtomicAtomicityAttr { PerElement } - -/// Packed global atomic add with exact format and semantic attributes. -#[pliron_op( - name = "nvvm.packed_atomic_add", - format, - interfaces = [NOpdsInterface<2>, NResultsInterface<1>], - attributes = ( - nvvm_packed_atomic_format: PackedAtomicFormatAttr, - nvvm_packed_atomic_state_space: PackedAtomicStateSpaceAttr, - nvvm_packed_atomic_ordering: PackedAtomicOrderingAttr, - nvvm_packed_atomic_scope: PackedAtomicScopeAttr, - nvvm_packed_atomic_rounding: PackedAtomicRoundingAttr, - nvvm_packed_atomic_subnormal: PackedAtomicSubnormalAttr, - nvvm_packed_atomic_atomicity: PackedAtomicAtomicityAttr - ) -)] -pub struct PackedAtomicAddOp; - -impl PackedAtomicAddOp { - pub fn new(op: Ptr) -> Self { Self { op } } - - pub fn build( - ctx: &mut Context, - address: Value, - addend: Value, - format: PackedAtomicFormatAttr, - ) -> Ptr { - let u32_ty = IntegerType::get(ctx, 32, Signedness::Unsigned); - let op = Operation::new( - ctx, - Self::get_concrete_op_info(), - vec![u32_ty.into()], - vec![address, addend], - vec![], - 0, - ); - let this = Self { op }; - this.set_attr_nvvm_packed_atomic_format(ctx, format); - this.set_attr_nvvm_packed_atomic_state_space(ctx, PackedAtomicStateSpaceAttr::Global); - this.set_attr_nvvm_packed_atomic_ordering(ctx, PackedAtomicOrderingAttr::Relaxed); - this.set_attr_nvvm_packed_atomic_scope(ctx, PackedAtomicScopeAttr::Gpu); - this.set_attr_nvvm_packed_atomic_rounding(ctx, PackedAtomicRoundingAttr::Rn); - this.set_attr_nvvm_packed_atomic_subnormal(ctx, PackedAtomicSubnormalAttr::NoFtz); - this.set_attr_nvvm_packed_atomic_atomicity(ctx, PackedAtomicAtomicityAttr::PerElement); - this.get_operation() - } -} - -fn is_u32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool { - ty.deref(ctx).downcast_ref::().is_some_and(|integer| { - integer.width() == 32 && integer.signedness() == Signedness::Unsigned - }) -} - -fn verify_packed_atomic_signature( - ctx: &Context, - op_ptr: Ptr, - op_name: &str, -) -> Result<(), Error> { - let op = op_ptr.deref(ctx); - if op.get_num_operands() != 2 || op.get_num_results() != 1 { - return verify_err!(op.loc(), "{} requires exactly two operands and one result", op_name); - } - let pointer_ty = op.get_operand(0).get_type(ctx); - let pointer_object = pointer_ty.deref(ctx); - let Some(pointer) = pointer_object.downcast_ref::() else { - return verify_err!(op.loc(), "{} address must be a MIR pointer", op_name); - }; - if !pointer.is_mutable() - || !matches!(pointer.address_space(), address_space::GENERIC | address_space::GLOBAL) - || !is_u32(ctx, pointer.pointee) - { - return verify_err!(op.loc(), "{} address must be a mutable generic/global pointer to u32", op_name); - } - if !is_u32(ctx, op.get_operand(1).get_type(ctx)) || !is_u32(ctx, op.get_result(0).get_type(ctx)) { - return verify_err!(op.loc(), "{} addend and result must be u32", op_name); - } - Ok(()) -} - -impl Verify for PackedAtomicAddOp { - fn verify(&self, ctx: &Context) -> Result<(), Error> { - verify_packed_atomic_signature(ctx, self.get_operation(), "nvvm.packed_atomic_add")?; - let op = self.get_operation().deref(ctx); - if self.get_attr_nvvm_packed_atomic_format(ctx).is_none() - || self.get_attr_nvvm_packed_atomic_state_space(ctx).as_deref() != Some(&PackedAtomicStateSpaceAttr::Global) - || self.get_attr_nvvm_packed_atomic_ordering(ctx).as_deref() != Some(&PackedAtomicOrderingAttr::Relaxed) - || self.get_attr_nvvm_packed_atomic_scope(ctx).as_deref() != Some(&PackedAtomicScopeAttr::Gpu) - || self.get_attr_nvvm_packed_atomic_rounding(ctx).as_deref() != Some(&PackedAtomicRoundingAttr::Rn) - || self.get_attr_nvvm_packed_atomic_subnormal(ctx).as_deref() != Some(&PackedAtomicSubnormalAttr::NoFtz) - || self.get_attr_nvvm_packed_atomic_atomicity(ctx).as_deref() != Some(&PackedAtomicAtomicityAttr::PerElement) - { - return verify_err!(op.loc(), "nvvm.packed_atomic_add has a missing or unsupported semantic attribute"); - } - Ok(()) - } -} - -/// Compatibility operation for the existing `nvvm.atom_add_f16x2` carrier. -#[pliron_op( - name = "nvvm.atom_add_f16x2", - format, - interfaces = [NOpdsInterface<2>, NResultsInterface<1>], -)] -pub struct NvvmAtomAddF16x2Op; - -impl NvvmAtomAddF16x2Op { - pub fn new(op: Ptr) -> Self { Self { op } } -} - -impl Verify for NvvmAtomAddF16x2Op { - fn verify(&self, ctx: &Context) -> Result<(), Error> { - verify_packed_atomic_signature(ctx, self.get_operation(), "nvvm.atom_add_f16x2") - } -} - -/// Compatibility operation for the existing `nvvm.atom_add_bf16x2` carrier. -#[pliron_op( - name = "nvvm.atom_add_bf16x2", - format, - interfaces = [NOpdsInterface<2>, NResultsInterface<1>], -)] -pub struct NvvmAtomAddBf16x2Op; - -impl NvvmAtomAddBf16x2Op { - pub fn new(op: Ptr) -> Self { Self { op } } -} - -impl Verify for NvvmAtomAddBf16x2Op { - fn verify(&self, ctx: &Context) -> Result<(), Error> { - verify_packed_atomic_signature(ctx, self.get_operation(), "nvvm.atom_add_bf16x2") - } -} - -pub(super) fn register(ctx: &mut Context) { - PackedAtomicFormatAttr::register(ctx); - PackedAtomicStateSpaceAttr::register(ctx); - PackedAtomicOrderingAttr::register(ctx); - PackedAtomicScopeAttr::register(ctx); - PackedAtomicRoundingAttr::register(ctx); - PackedAtomicSubnormalAttr::register(ctx); - PackedAtomicAtomicityAttr::register(ctx); - PackedAtomicAddOp::register(ctx); - NvvmAtomAddF16x2Op::register(ctx); - NvvmAtomAddBf16x2Op::register(ctx); -} -"##, - ); - output -} - -fn render_dialect_packed_alu(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - "//! Structural operations for generated packed floating-point arithmetic.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_integer_width(\n ctx: &Context,\n ty: pliron::r#type::TypeHandle,\n width: u32,\n) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == width)\n}\n\n", - ); - for record in packed_alus(catalog) { - let arity = record.rust.arguments.len(); - let width = packed_alu_width(record); - let parameters = (0..arity) - .map(|index| format!("arg{index}: Value")) - .collect::>() - .join(", "); - let operands = (0..arity) - .map(|index| format!("arg{index}")) - .collect::>() - .join(", "); - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!( - output, - "///\n/// Lowers to `{}`.", - packed_alu_ptx_mnemonic(record) - ) - .unwrap(); - writeln!( - output, - "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<{arity}>, NResultsInterface<1>],\n)]", - record.dialect.op_name - ) - .unwrap(); - writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); - writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); - writeln!( - output, - " pub fn new(op: Ptr) -> Self {{\n Self {{ op }}\n }}\n\n pub fn build(ctx: &mut Context, {parameters}) -> Ptr {{\n let result_ty = IntegerType::get(ctx, {width}, Signedness::Unsigned);\n Operation::new(\n ctx,\n Self::get_concrete_op_info(),\n vec![result_ty.into()],\n vec![{operands}],\n vec![],\n 0,\n )\n }}\n}}" - ) - .unwrap(); - writeln!(output, "\nimpl Verify for {} {{", record.dialect.op_type).unwrap(); - writeln!( - output, - " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n let op = self.get_operation().deref(ctx);\n if op.get_num_operands() != {arity} || op.get_num_results() != 1 {{\n return verify_err!(op.loc(), {:?});\n }}\n if !(0..{arity}).all(|index| is_integer_width(ctx, op.get_operand(index).get_type(ctx), {width}))\n || !is_integer_width(ctx, op.get_result(0).get_type(ctx), {width})\n {{\n return verify_err!(op.loc(), {:?});\n }}\n Ok(())\n }}\n}}\n", - format!( - "{} requires exactly {arity} operands and one result", - record.dialect.op_name - ), - format!( - "{} operands and result must be {width}-bit integers", - record.dialect.op_name, - ), - ) - .unwrap(); - } - output.push_str("\npub(super) fn register(ctx: &mut Context) {\n"); - for record in packed_alus(catalog) { - writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); - } - output.push_str("}\n"); - output -} - -fn render_dialect_packed_conversion(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - "//! Structural operation for generated packed conversion.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{FP32Type, IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_f32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx).downcast_ref::().is_some()\n}\n\nfn is_integer_width(\n ctx: &Context,\n ty: pliron::r#type::TypeHandle,\n width: u32,\n) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == width)\n}\n\n", - ); - for record in packed_conversions(catalog) { - let result_width = packed_conversion_result_width(record); - writeln!(output, "/// {}", record.summary).unwrap(); - match packed_conversion_source(record) { - PackedConversionSourceFormat::F32x2 => { - writeln!( - output, - "///\n/// The first input becomes the low {} lane; the second becomes the high lane.", - packed_conversion_element(record) - ) - .unwrap(); - writeln!( - output, - "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<2>, NResultsInterface<1>],\n)]", - record.dialect.op_name - ) - .unwrap(); - writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); - writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); - writeln!( - output, - " pub fn new(op: Ptr) -> Self {{\n Self {{ op }}\n }}\n\n pub fn build(ctx: &mut Context, low: Value, high: Value) -> Ptr {{\n let result_ty = IntegerType::get(ctx, {result_width}, Signedness::Unsigned);\n Operation::new(\n ctx,\n Self::get_concrete_op_info(),\n vec![result_ty.into()],\n vec![low, high],\n vec![],\n 0,\n )\n }}\n}}", - ) - .unwrap(); - writeln!(output, "\nimpl Verify for {} {{", record.dialect.op_type).unwrap(); - writeln!( - output, - " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n let op = self.get_operation().deref(ctx);\n if op.get_num_operands() != 2 || op.get_num_results() != 1 {{\n return verify_err!(op.loc(), {:?});\n }}\n if !is_f32(ctx, op.get_operand(0).get_type(ctx))\n || !is_f32(ctx, op.get_operand(1).get_type(ctx))\n || !is_integer_width(ctx, op.get_result(0).get_type(ctx), {})\n {{\n return verify_err!(op.loc(), {:?});\n }}\n Ok(())\n }}\n}}\n", - format!("{} requires two operands and one result", record.dialect.op_name), - result_width, - format!( - "{} requires f32 operands and one {result_width}-bit integer result", - record.dialect.op_name, - ), - ) - .unwrap(); - } - PackedConversionSourceFormat::E4m3x2 - | PackedConversionSourceFormat::E5m2x2 - | PackedConversionSourceFormat::F16x2 => { - let source_width = packed_conversion_source_width(record); - writeln!( - output, - "///\n/// The single input carries both packed lanes, and lane order is preserved." - ) - .unwrap(); - writeln!( - output, - "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<1>, NResultsInterface<1>],\n)]", - record.dialect.op_name - ) - .unwrap(); - writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); - writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); - writeln!( - output, - " pub fn new(op: Ptr) -> Self {{\n Self {{ op }}\n }}\n\n pub fn build(ctx: &mut Context, packed: Value) -> Ptr {{\n let result_ty = IntegerType::get(ctx, {result_width}, Signedness::Unsigned);\n Operation::new(\n ctx,\n Self::get_concrete_op_info(),\n vec![result_ty.into()],\n vec![packed],\n vec![],\n 0,\n )\n }}\n}}", - ) - .unwrap(); - writeln!(output, "\nimpl Verify for {} {{", record.dialect.op_type).unwrap(); - writeln!( - output, - " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n let op = self.get_operation().deref(ctx);\n if op.get_num_operands() != 1 || op.get_num_results() != 1 {{\n return verify_err!(op.loc(), {:?});\n }}\n if !is_integer_width(ctx, op.get_operand(0).get_type(ctx), {source_width})\n || !is_integer_width(ctx, op.get_result(0).get_type(ctx), {result_width})\n {{\n return verify_err!(op.loc(), {:?});\n }}\n Ok(())\n }}\n}}\n", - format!("{} requires one operand and one result", record.dialect.op_name), - format!( - "{} requires one {source_width}-bit integer operand and one {result_width}-bit integer result", - record.dialect.op_name, - ), - ) - .unwrap(); - } - } - } - output.push_str("\npub(super) fn register(ctx: &mut Context) {\n"); - for record in packed_conversions(catalog) { - writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); - } - output.push_str("}\n"); - output -} - -fn render_dialect_scalar_conversion(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - r##"//! One structural operation for generated `f32` to TF32 conversions. - -use pliron::{ - attribute::Attribute, - builtin::{ - op_interfaces::{NOpdsInterface, NResultsInterface}, - types::{FP32Type, IntegerType, Signedness}, - }, - 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_attr, pliron_op}; - -#[pliron_attr(name = "nvvm.scalar_conversion_rounding", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum ScalarConversionRoundingAttr { - NearestAway, - NearestEven, - TowardZero, -} - -#[pliron_attr(name = "nvvm.scalar_conversion_saturation", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum ScalarConversionSaturationAttr { - None, - Relu, - Satfinite, - ReluSatfinite, -} - -/// Converts one `f32` value and returns the raw TF32 bits. -#[pliron_op( - name = "nvvm.scalar_conversion", - format, - interfaces = [NOpdsInterface<1>, NResultsInterface<1>], - attributes = ( - nvvm_scalar_conversion_rounding: ScalarConversionRoundingAttr, - nvvm_scalar_conversion_saturation: ScalarConversionSaturationAttr - ) -)] -pub struct ScalarConversionOp; - -impl ScalarConversionOp { - pub fn new(op: Ptr) -> Self { - Self { op } - } - - pub fn build( - ctx: &mut Context, - value: Value, - rounding: ScalarConversionRoundingAttr, - saturation: ScalarConversionSaturationAttr, - ) -> Ptr { - let result_ty = IntegerType::get(ctx, 32, Signedness::Unsigned); - let op = Operation::new( - ctx, - Self::get_concrete_op_info(), - vec![result_ty.into()], - vec![value], - vec![], - 0, - ); - let this = Self { op }; - this.set_attr_nvvm_scalar_conversion_rounding(ctx, rounding); - this.set_attr_nvvm_scalar_conversion_saturation(ctx, saturation); - this.get_operation() - } -} - -impl Verify for ScalarConversionOp { - fn verify(&self, ctx: &Context) -> Result<(), Error> { - let op = self.get_operation().deref(ctx); - let Some(rounding) = self.get_attr_nvvm_scalar_conversion_rounding(ctx) else { - return verify_err!(op.loc(), "nvvm.scalar_conversion requires rounding"); - }; - let Some(saturation) = self.get_attr_nvvm_scalar_conversion_saturation(ctx) else { - return verify_err!(op.loc(), "nvvm.scalar_conversion requires saturation"); - }; - let admitted = matches!( - (&*rounding, &*saturation), -"##, - ); - for (index, record) in scalar_conversions(catalog).enumerate() { - if index != 0 { - output.push_str(" |\n"); - } - writeln!( - output, - " ({}, {})", - scalar_conversion_rounding_attr(record), - scalar_conversion_saturation_attr(record), - ) - .unwrap(); - } - output.push_str( - r##" ); - if !admitted { - return verify_err!(op.loc(), "nvvm.scalar_conversion variant is not admitted"); - } - if op - .get_operand(0) - .get_type(ctx) - .deref(ctx) - .downcast_ref::() - .is_none() - { - return verify_err!(op.loc(), "nvvm.scalar_conversion operand must be f32"); - } - if op - .get_result(0) - .get_type(ctx) - .deref(ctx) - .downcast_ref::() - .is_none_or(|integer| integer.width() != 32) - { - return verify_err!(op.loc(), "nvvm.scalar_conversion result must be a 32-bit integer"); - } - Ok(()) - } -} - -pub(super) fn register(ctx: &mut Context) { - ScalarConversionRoundingAttr::register(ctx); - ScalarConversionSaturationAttr::register(ctx); - ScalarConversionOp::register(ctx); -} -"##, - ); - output -} - -fn render_dialect_scalar_arithmetic(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - r##"//! One structural operation for generated scalar arithmetic. - -use pliron::{ - attribute::Attribute, - builtin::{op_interfaces::NResultsInterface, types::{FP32Type, FP64Type}}, - 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_attr, pliron_op}; - -#[pliron_attr(name = "nvvm.scalar_arithmetic_format", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum ScalarArithmeticFormatAttr { F32, F64 } - -#[pliron_attr(name = "nvvm.scalar_arithmetic_operation", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum ScalarArithmeticOperationAttr { Mul, Div, Fma, Add } - -#[pliron_attr(name = "nvvm.scalar_arithmetic_rounding", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum ScalarArithmeticRoundingAttr { Rn, Rz, Rm, Rp } - -#[pliron_attr(name = "nvvm.scalar_arithmetic_subnormal", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum ScalarArithmeticSubnormalAttr { Preserve, Ftz } - -#[pliron_attr(name = "nvvm.scalar_arithmetic_saturation", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum ScalarArithmeticSaturationAttr { None, Sat } - -/// Scalar arithmetic with explicit PTX modifiers. -#[pliron_op( - name = "nvvm.scalar_arithmetic", - format, - interfaces = [NResultsInterface<1>], - attributes = ( - nvvm_scalar_arithmetic_format: ScalarArithmeticFormatAttr, - nvvm_scalar_arithmetic_operation: ScalarArithmeticOperationAttr, - nvvm_scalar_arithmetic_rounding: ScalarArithmeticRoundingAttr, - nvvm_scalar_arithmetic_subnormal: ScalarArithmeticSubnormalAttr, - nvvm_scalar_arithmetic_saturation: ScalarArithmeticSaturationAttr - ) -)] -pub struct ScalarArithmeticOp; - -impl ScalarArithmeticOp { - pub fn new(op: Ptr) -> Self { Self { op } } - - pub fn build( - ctx: &mut Context, - operands: Vec, - format: ScalarArithmeticFormatAttr, - operation: ScalarArithmeticOperationAttr, - rounding: ScalarArithmeticRoundingAttr, - subnormal: ScalarArithmeticSubnormalAttr, - saturation: ScalarArithmeticSaturationAttr, - ) -> Ptr { - let result_ty = match &format { - ScalarArithmeticFormatAttr::F32 => FP32Type::get(ctx).into(), - ScalarArithmeticFormatAttr::F64 => FP64Type::get(ctx).into(), - }; - let op = Operation::new( - ctx, - Self::get_concrete_op_info(), - vec![result_ty], - operands, - vec![], - 0, - ); - let this = Self { op }; - this.set_attr_nvvm_scalar_arithmetic_format(ctx, format); - this.set_attr_nvvm_scalar_arithmetic_operation(ctx, operation); - this.set_attr_nvvm_scalar_arithmetic_rounding(ctx, rounding); - this.set_attr_nvvm_scalar_arithmetic_subnormal(ctx, subnormal); - this.set_attr_nvvm_scalar_arithmetic_saturation(ctx, saturation); - this.get_operation() - } -} - -impl Verify for ScalarArithmeticOp { - fn verify(&self, ctx: &Context) -> Result<(), Error> { - let op = self.get_operation().deref(ctx); - let Some(format) = self.get_attr_nvvm_scalar_arithmetic_format(ctx) else { - return verify_err!(op.loc(), "nvvm.scalar_arithmetic requires a format"); - }; - let Some(operation) = self.get_attr_nvvm_scalar_arithmetic_operation(ctx) else { - return verify_err!(op.loc(), "nvvm.scalar_arithmetic requires an operation"); - }; - let Some(rounding) = self.get_attr_nvvm_scalar_arithmetic_rounding(ctx) else { - return verify_err!(op.loc(), "nvvm.scalar_arithmetic requires rounding"); - }; - let Some(subnormal) = self.get_attr_nvvm_scalar_arithmetic_subnormal(ctx) else { - return verify_err!(op.loc(), "nvvm.scalar_arithmetic requires a subnormal mode"); - }; - let Some(saturation) = self.get_attr_nvvm_scalar_arithmetic_saturation(ctx) else { - return verify_err!(op.loc(), "nvvm.scalar_arithmetic requires saturation"); - }; - let admitted = matches!( - (&*format, &*operation, &*rounding, &*subnormal, &*saturation), -"##, - ); - for (index, record) in scalar_arithmetics(catalog).enumerate() { - if index != 0 { - output.push_str(" |\n"); - } - writeln!( - output, - " ({}, {}, {}, {}, {})", - scalar_arithmetic_format_attr(record), - scalar_arithmetic_operation_attr(record), - scalar_arithmetic_rounding_attr(record), - scalar_arithmetic_subnormal_attr(record), - scalar_arithmetic_saturation_attr(record), - ) - .unwrap(); - } - output.push_str( - r##" ); - if !admitted { - return verify_err!(op.loc(), "nvvm.scalar_arithmetic variant is not admitted"); - } - let expected_operands = match &*operation { - ScalarArithmeticOperationAttr::Mul - | ScalarArithmeticOperationAttr::Div - | ScalarArithmeticOperationAttr::Add => 2, - ScalarArithmeticOperationAttr::Fma => 3, - }; - if op.get_num_operands() != expected_operands || op.get_num_results() != 1 { - return verify_err!( - op.loc(), - "nvvm.scalar_arithmetic requires {} operands and one result", - expected_operands, - ); - } - let type_matches = |ty: pliron::r#type::TypeHandle| match &*format { - ScalarArithmeticFormatAttr::F32 => ty.deref(ctx).downcast_ref::().is_some(), - ScalarArithmeticFormatAttr::F64 => ty.deref(ctx).downcast_ref::().is_some(), - }; - if !(0..expected_operands).all(|index| type_matches(op.get_operand(index).get_type(ctx))) - || !type_matches(op.get_result(0).get_type(ctx)) - { - return verify_err!(op.loc(), "nvvm.scalar_arithmetic types do not match its format"); - } - Ok(()) - } -} - -pub(super) fn register(ctx: &mut Context) { - ScalarArithmeticFormatAttr::register(ctx); - ScalarArithmeticOperationAttr::register(ctx); - ScalarArithmeticRoundingAttr::register(ctx); - ScalarArithmeticSubnormalAttr::register(ctx); - ScalarArithmeticSaturationAttr::register(ctx); - ScalarArithmeticOp::register(ctx); -} -"##, - ); - output -} - -fn render_dialect_scalar_math(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - r##"//! One structural operation for generated scalar math. - -use pliron::{ - attribute::Attribute, - builtin::{op_interfaces::{NOpdsInterface, NResultsInterface}, types::{FP32Type, FP64Type, IntegerType, Signedness}}, - 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_attr, pliron_op}; - -#[pliron_attr(name = "nvvm.scalar_math_format", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum ScalarMathFormatAttr { F16, F32, F64 } - -#[pliron_attr(name = "nvvm.scalar_math_operation", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum ScalarMathOperationAttr { Sin, Cos, Ex2, Lg2, Rcp, Rsqrt, Sqrt, Tanh } - -#[pliron_attr(name = "nvvm.scalar_math_precision", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum ScalarMathPrecisionAttr { Approx, Rn, Rz, Rm, Rp } - -#[pliron_attr(name = "nvvm.scalar_math_subnormal", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum ScalarMathSubnormalAttr { Preserve, Ftz } - -/// Scalar math with exact PTX modifiers. -#[pliron_op( - name = "nvvm.scalar_math", - format, - interfaces = [NOpdsInterface<1>, NResultsInterface<1>], - attributes = ( - nvvm_scalar_math_format: ScalarMathFormatAttr, - nvvm_scalar_math_operation: ScalarMathOperationAttr, - nvvm_scalar_math_precision: ScalarMathPrecisionAttr, - nvvm_scalar_math_subnormal: ScalarMathSubnormalAttr - ) -)] -pub struct ScalarMathOp; - -impl ScalarMathOp { - pub fn new(op: Ptr) -> Self { Self { op } } - - pub fn build( - ctx: &mut Context, - operand: Value, - format: ScalarMathFormatAttr, - operation: ScalarMathOperationAttr, - precision: ScalarMathPrecisionAttr, - subnormal: ScalarMathSubnormalAttr, - ) -> Ptr { - let result_ty = match &format { - ScalarMathFormatAttr::F16 => IntegerType::get(ctx, 16, Signedness::Unsigned).into(), - ScalarMathFormatAttr::F32 => FP32Type::get(ctx).into(), - ScalarMathFormatAttr::F64 => FP64Type::get(ctx).into(), - }; - let op = Operation::new( - ctx, - Self::get_concrete_op_info(), - vec![result_ty], - vec![operand], - vec![], - 0, - ); - let this = Self { op }; - this.set_attr_nvvm_scalar_math_format(ctx, format); - this.set_attr_nvvm_scalar_math_operation(ctx, operation); - this.set_attr_nvvm_scalar_math_precision(ctx, precision); - this.set_attr_nvvm_scalar_math_subnormal(ctx, subnormal); - this.get_operation() - } -} - -impl Verify for ScalarMathOp { - fn verify(&self, ctx: &Context) -> Result<(), Error> { - let op = self.get_operation().deref(ctx); - let Some(format) = self.get_attr_nvvm_scalar_math_format(ctx) else { - return verify_err!(op.loc(), "nvvm.scalar_math requires a format"); - }; - let Some(operation) = self.get_attr_nvvm_scalar_math_operation(ctx) else { - return verify_err!(op.loc(), "nvvm.scalar_math requires an operation"); - }; - let Some(precision) = self.get_attr_nvvm_scalar_math_precision(ctx) else { - return verify_err!(op.loc(), "nvvm.scalar_math requires a precision"); - }; - let Some(subnormal) = self.get_attr_nvvm_scalar_math_subnormal(ctx) else { - return verify_err!(op.loc(), "nvvm.scalar_math requires a subnormal mode"); - }; - let admitted = matches!( - (&*format, &*operation, &*precision, &*subnormal), -"##, - ); - for (index, record) in scalar_maths(catalog).enumerate() { - if index != 0 { - output.push_str(" |\n"); - } - writeln!( - output, - " ({}, {}, {}, {})", - scalar_math_format_attr(record), - scalar_math_operation_attr(record), - scalar_math_precision_attr(record), - scalar_math_subnormal_attr(record), - ) - .unwrap(); - } - output.push_str( - r##" ); - if !admitted { - return verify_err!(op.loc(), "nvvm.scalar_math variant is not admitted"); - } - if op.get_num_operands() != 1 || op.get_num_results() != 1 { - return verify_err!( - op.loc(), - "nvvm.scalar_math requires one operand and one result", - ); - } - let type_matches = |ty: pliron::r#type::TypeHandle| match &*format { - ScalarMathFormatAttr::F16 => ty.deref(ctx).downcast_ref::().is_some_and(|integer| integer.width() == 16), - ScalarMathFormatAttr::F32 => ty.deref(ctx).downcast_ref::().is_some(), - ScalarMathFormatAttr::F64 => ty.deref(ctx).downcast_ref::().is_some(), - }; - if !type_matches(op.get_operand(0).get_type(ctx)) - || !type_matches(op.get_result(0).get_type(ctx)) - { - return verify_err!(op.loc(), "nvvm.scalar_math types do not match its format"); - } - Ok(()) - } -} - -pub(super) fn register(ctx: &mut Context) { - ScalarMathFormatAttr::register(ctx); - ScalarMathOperationAttr::register(ctx); - ScalarMathPrecisionAttr::register(ctx); - ScalarMathSubnormalAttr::register(ctx); - ScalarMathOp::register(ctx); -} -"##, - ); - output -} - -fn render_dialect_extended_minmax(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - r##"//! One structural operation for extended floating-point min/max. - -use pliron::{ - attribute::Attribute, - builtin::{ - op_interfaces::{NOpdsInterface, NResultsInterface}, - types::{FP32Type, IntegerType, Signedness}, - }, - 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_attr, pliron_op}; - -#[pliron_attr(name = "nvvm.extended_minmax_format", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum ExtendedMinMaxFormatAttr { F32, F16, Bf16, F16x2, Bf16x2 } - -#[pliron_attr(name = "nvvm.extended_minmax_operation", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum ExtendedMinMaxOperationAttr { Min, Max } - -#[pliron_attr(name = "nvvm.extended_minmax_subnormal", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum ExtendedMinMaxSubnormalAttr { Preserve, Ftz } - -#[pliron_attr(name = "nvvm.extended_minmax_nan", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum ExtendedMinMaxNanAttr { Number, Nan } - -#[pliron_attr(name = "nvvm.extended_minmax_xorsign_abs", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum ExtendedMinMaxXorSignAbsAttr { Disabled, Enabled } - -/// Extended min/max with exact PTX modifiers. -#[pliron_op( - name = "nvvm.extended_minmax", - format, - interfaces = [NOpdsInterface<2>, NResultsInterface<1>], - attributes = ( - nvvm_extended_minmax_format: ExtendedMinMaxFormatAttr, - nvvm_extended_minmax_operation: ExtendedMinMaxOperationAttr, - nvvm_extended_minmax_subnormal: ExtendedMinMaxSubnormalAttr, - nvvm_extended_minmax_nan: ExtendedMinMaxNanAttr, - nvvm_extended_minmax_xorsign_abs: ExtendedMinMaxXorSignAbsAttr - ) -)] -pub struct ExtendedMinMaxOp; - -impl ExtendedMinMaxOp { - pub fn new(op: Ptr) -> Self { Self { op } } - - #[allow(clippy::too_many_arguments)] - pub fn build( - ctx: &mut Context, - a: Value, - b: Value, - format: ExtendedMinMaxFormatAttr, - operation: ExtendedMinMaxOperationAttr, - subnormal: ExtendedMinMaxSubnormalAttr, - nan: ExtendedMinMaxNanAttr, - xorsign_abs: ExtendedMinMaxXorSignAbsAttr, - ) -> Ptr { - let result_ty = match &format { - ExtendedMinMaxFormatAttr::F32 => FP32Type::get(ctx).into(), - ExtendedMinMaxFormatAttr::F16 | ExtendedMinMaxFormatAttr::Bf16 => { - IntegerType::get(ctx, 16, Signedness::Unsigned).into() - } - ExtendedMinMaxFormatAttr::F16x2 | ExtendedMinMaxFormatAttr::Bf16x2 => { - IntegerType::get(ctx, 32, Signedness::Unsigned).into() - } - }; - let op = Operation::new( - ctx, - Self::get_concrete_op_info(), - vec![result_ty], - vec![a, b], - vec![], - 0, - ); - let this = Self { op }; - this.set_attr_nvvm_extended_minmax_format(ctx, format); - this.set_attr_nvvm_extended_minmax_operation(ctx, operation); - this.set_attr_nvvm_extended_minmax_subnormal(ctx, subnormal); - this.set_attr_nvvm_extended_minmax_nan(ctx, nan); - this.set_attr_nvvm_extended_minmax_xorsign_abs(ctx, xorsign_abs); - this.get_operation() - } -} - -impl Verify for ExtendedMinMaxOp { - fn verify(&self, ctx: &Context) -> Result<(), Error> { - let op = self.get_operation().deref(ctx); - let Some(format) = self.get_attr_nvvm_extended_minmax_format(ctx) else { - return verify_err!(op.loc(), "nvvm.extended_minmax requires a format"); - }; - let Some(operation) = self.get_attr_nvvm_extended_minmax_operation(ctx) else { - return verify_err!(op.loc(), "nvvm.extended_minmax requires an operation"); - }; - let Some(subnormal) = self.get_attr_nvvm_extended_minmax_subnormal(ctx) else { - return verify_err!(op.loc(), "nvvm.extended_minmax requires a subnormal mode"); - }; - let Some(nan) = self.get_attr_nvvm_extended_minmax_nan(ctx) else { - return verify_err!(op.loc(), "nvvm.extended_minmax requires a NaN mode"); - }; - let Some(xorsign_abs) = self.get_attr_nvvm_extended_minmax_xorsign_abs(ctx) else { - return verify_err!(op.loc(), "nvvm.extended_minmax requires an xorsign.abs mode"); - }; - let admitted = matches!( - (&*format, &*operation, &*subnormal, &*nan, &*xorsign_abs), -"##, - ); - for (index, record) in extended_minmax(catalog).enumerate() { - if index != 0 { - output.push_str(" |\n"); - } - writeln!( - output, - " ({}, {}, {}, {}, {})", - extended_minmax_format_attr(record), - extended_minmax_operation_attr(record), - extended_minmax_subnormal_attr(record), - extended_minmax_nan_attr(record), - extended_minmax_xorsign_abs_attr(record), - ) - .unwrap(); - } - output.push_str( - r##" ); - if !admitted { - return verify_err!(op.loc(), "nvvm.extended_minmax variant is not admitted"); - } - let type_matches = |ty: pliron::r#type::TypeHandle| match &*format { - ExtendedMinMaxFormatAttr::F32 => ty.deref(ctx).downcast_ref::().is_some(), - ExtendedMinMaxFormatAttr::F16 | ExtendedMinMaxFormatAttr::Bf16 => ty - .deref(ctx) - .downcast_ref::() - .is_some_and(|integer| integer.width() == 16), - ExtendedMinMaxFormatAttr::F16x2 | ExtendedMinMaxFormatAttr::Bf16x2 => ty - .deref(ctx) - .downcast_ref::() - .is_some_and(|integer| integer.width() == 32), - }; - if !type_matches(op.get_operand(0).get_type(ctx)) - || !type_matches(op.get_operand(1).get_type(ctx)) - || !type_matches(op.get_result(0).get_type(ctx)) - { - return verify_err!(op.loc(), "nvvm.extended_minmax types do not match its format"); - } - Ok(()) - } -} - -pub(super) fn register(ctx: &mut Context) { - ExtendedMinMaxFormatAttr::register(ctx); - ExtendedMinMaxOperationAttr::register(ctx); - ExtendedMinMaxSubnormalAttr::register(ctx); - ExtendedMinMaxNanAttr::register(ctx); - ExtendedMinMaxXorSignAbsAttr::register(ctx); - ExtendedMinMaxOp::register(ctx); -} -"##, - ); - output -} - -fn render_dialect_prmt(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - r##"//! One structural operation for the closed generated `prmt` family. - -use pliron::{ - attribute::Attribute, - builtin::{ - op_interfaces::NResultsInterface, - types::{IntegerType, Signedness}, - }, - 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_attr, pliron_op}; - -#[pliron_attr(name = "nvvm.prmt_mode", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum PrmtModeAttr { - Generic, - F4e, - B4e, - Rc8, - Ecl, - Ecr, - Rc16, -} - -/// Byte permutation whose exact mode is carried by an attribute. -#[pliron_op( - name = "nvvm.prmt", - format, - interfaces = [NResultsInterface<1>], - attributes = (nvvm_prmt_mode: PrmtModeAttr) -)] -pub struct PrmtOp; - -impl PrmtOp { - pub fn new(op: Ptr) -> Self { - Self { op } - } - - pub fn build(ctx: &mut Context, operands: Vec, mode: PrmtModeAttr) -> Ptr { - let result_ty = IntegerType::get(ctx, 32, Signedness::Unsigned); - let op = Operation::new( - ctx, - Self::get_concrete_op_info(), - vec![result_ty.into()], - operands, - vec![], - 0, - ); - let this = Self { op }; - this.set_attr_nvvm_prmt_mode(ctx, mode); - this.get_operation() - } -} - -impl Verify for PrmtOp { - fn verify(&self, ctx: &Context) -> Result<(), Error> { - let op = self.get_operation().deref(ctx); - let Some(mode) = self.get_attr_nvvm_prmt_mode(ctx) else { - return verify_err!(op.loc(), "nvvm.prmt requires a mode attribute"); - }; - let expected_operands = match &*mode { - PrmtModeAttr::Generic | PrmtModeAttr::F4e | PrmtModeAttr::B4e => 3, - PrmtModeAttr::Rc8 | PrmtModeAttr::Ecl | PrmtModeAttr::Ecr | PrmtModeAttr::Rc16 => 2, - }; - if op.get_num_operands() != expected_operands || op.get_num_results() != 1 { - return verify_err!( - op.loc(), - "nvvm.prmt {:?} requires {} operands and one result", - mode, - expected_operands - ); - } - for index in 0..expected_operands { - let ty = op.get_operand(index).get_type(ctx); - if ty - .deref(ctx) - .downcast_ref::() - .is_none_or(|integer| integer.width() != 32) - { - return verify_err!(op.loc(), "nvvm.prmt operand {} must be a 32-bit integer", index); - } - } - let result_ty = op.get_result(0).get_type(ctx); - if result_ty - .deref(ctx) - .downcast_ref::() - .is_none_or(|integer| integer.width() != 32) - { - return verify_err!(op.loc(), "nvvm.prmt result must be a 32-bit integer"); - } - Ok(()) - } -} - -pub(super) fn register(ctx: &mut Context) { - PrmtModeAttr::register(ctx); - PrmtOp::register(ctx); -} -"##, - ); - output -} - -fn render_dialect_cluster_barrier(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - if cluster_barriers(catalog).next().is_none() { - return output; - } - output.push_str( - r##"//! Generated operations for the closed cluster-barrier family. - -use pliron::{ - attribute::Attribute, - builtin::op_interfaces::{NOpdsInterface, NResultsInterface}, - common_traits::Verify, - context::{Context, Ptr}, - location::Located, - op::Op, - operation::Operation, - result::Error, - verify_err, -}; -use pliron_derive::{pliron_attr, pliron_op}; - -#[pliron_attr(name = "nvvm.cluster_barrier_mode", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum ClusterBarrierModeAttr { - Arrive, - ArriveAligned, - ArriveRelaxed, - ArriveRelaxedAligned, - Wait, - WaitAligned, -} - -/// Cluster synchronization whose exact spelling is carried by an attribute. -#[pliron_op( - name = "nvvm.cluster_barrier", - format, - interfaces = [NOpdsInterface<0>, NResultsInterface<0>], - attributes = (nvvm_cluster_barrier_mode: ClusterBarrierModeAttr) -)] -pub struct ClusterBarrierOp; - -impl ClusterBarrierOp { - pub fn new(op: Ptr) -> Self { - Self { op } - } - - pub fn build(ctx: &mut Context, mode: ClusterBarrierModeAttr) -> Ptr { - let op = Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![], vec![], 0); - let this = Self { op }; - this.set_attr_nvvm_cluster_barrier_mode(ctx, mode); - this.get_operation() - } -} - -impl Verify for ClusterBarrierOp { - fn verify(&self, ctx: &Context) -> Result<(), Error> { - let op = self.get_operation().deref(ctx); - if self.get_attr_nvvm_cluster_barrier_mode(ctx).is_none() { - return verify_err!(op.loc(), "nvvm.cluster_barrier requires a mode attribute"); - } - if op.get_num_operands() != 0 || op.get_num_results() != 0 { - return verify_err!(op.loc(), "nvvm.cluster_barrier takes no operands or results"); - } - Ok(()) - } -} - -/// Compatibility operation for a complete aligned cluster synchronization. -#[pliron_op( - name = "nvvm.cluster_sync", - format, - verifier = "succ", - interfaces = [NOpdsInterface<0>, NResultsInterface<0>], -)] -pub struct ClusterSyncOp; - -impl ClusterSyncOp { - pub fn new(op: Ptr) -> Self { - Self { op } - } -} - -pub(super) fn register(ctx: &mut Context) { - ClusterBarrierModeAttr::register(ctx); - ClusterBarrierOp::register(ctx); - ClusterSyncOp::register(ctx); -} -"##, - ); - output -} - -fn render_dialect_cluster_memory(catalog: &CatalogFile, hash: &str) -> String { - assert_eq!(cluster_memory(catalog).count(), 2); - let mut output = rust_header(catalog, hash); - output.push_str( - r#"//! Structural operations for cluster address mapping and remote shared reads. - -use dialect_mir::types::{MirPointerKind, MirPtrType, address_space}; -use pliron::{ - builtin::{ - op_interfaces::{NOpdsInterface, NResultsInterface}, - types::{IntegerType, Signedness}, - }, - 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; - -fn verify_pointer_rank( - ctx: &Context, - operation: Ptr, - name: &str, -) -> Result<(), Error> { - let op = operation.deref(ctx); - if op.get_num_operands() != 2 || op.get_num_results() != 1 { - return verify_err!(op.loc(), "{name} requires two operands and one result"); - } - if op - .get_operand(0) - .get_type(ctx) - .deref(ctx) - .downcast_ref::() - .is_none() - { - return verify_err!(op.loc(), "{name} source must be a MIR pointer"); - } - let rank_ty = op.get_operand(1).get_type(ctx); - let rank_ty = rank_ty.deref(ctx); - if !rank_ty - .downcast_ref::() - .is_some_and(|ty| ty.width() == 32 && ty.signedness() == Signedness::Unsigned) - { - return verify_err!(op.loc(), "{name} rank must be u32"); - } - Ok(()) -} - -#[pliron_op( - name = "nvvm.mapa_shared_cluster", - format, - interfaces = [NOpdsInterface<2>, NResultsInterface<1>], -)] -pub struct MapaSharedClusterOp; - -impl MapaSharedClusterOp { - pub fn new(op: Ptr) -> Self { Self { op } } - - // Kind-preserving addrspace retype: copies the source pointer's kind - // verbatim; concrete kinds are minted only in mir-importer's facts.rs. - #[allow(clippy::disallowed_methods)] - pub fn build(ctx: &mut Context, source: Value, rank: Value) -> Ptr { - let source_ty = source.get_type(ctx); - let (pointee, is_mutable, pointer_kind) = { - let source_ty_obj = source_ty.deref(ctx); - let source_ptr = source_ty_obj - .downcast_ref::() - .expect("nvvm.mapa_shared_cluster source must be a MIR pointer"); - ( - source_ptr.pointee, - source_ptr.is_mutable(), - source_ptr.pointer_kind(), - ) - }; - let result_ty = MirPtrType::get_with_kind( - ctx, - pointee, - is_mutable, - address_space::CLUSTER_SHARED, - pointer_kind, - ); - Operation::new( - ctx, - Self::get_concrete_op_info(), - vec![result_ty.into()], - vec![source, rank], - vec![], - 0, - ) - } -} - -impl Verify for MapaSharedClusterOp { - fn verify(&self, ctx: &Context) -> Result<(), Error> { - let operation = self.get_operation(); - verify_pointer_rank(ctx, operation, "nvvm.mapa_shared_cluster")?; - let op = operation.deref(ctx); - let source_ty = op.get_operand(0).get_type(ctx); - let result_ty = op.get_result(0).get_type(ctx); - let source_ty_obj = source_ty.deref(ctx); - let result_ty_obj = result_ty.deref(ctx); - let source_ptr = source_ty_obj.downcast_ref::().unwrap(); - let Some(result_ptr) = result_ty_obj.downcast_ref::() else { - return verify_err!( - op.loc(), - "nvvm.mapa_shared_cluster result must be a MIR pointer" - ); - }; - if !matches!( - source_ptr.pointer_kind(), - MirPointerKind::RawConst | MirPointerKind::RawMut - ) - || result_ptr.pointee != source_ptr.pointee - || result_ptr.is_mutable() != source_ptr.is_mutable() - || !result_ptr.is_cluster_shared() - || result_ptr.pointer_kind() != source_ptr.pointer_kind() - { - return verify_err!( - op.loc(), - "nvvm.mapa_shared_cluster requires a raw source pointer and must preserve its pointee, mutability, and raw kind while returning addrspace(7)" - ); - } - Ok(()) - } -} - -#[pliron_op( - name = "nvvm.dsmem_read_u32", - format, - interfaces = [NOpdsInterface<2>, NResultsInterface<1>], -)] -pub struct DsmemReadU32Op; - -impl DsmemReadU32Op { - pub fn new(op: Ptr) -> Self { Self { op } } - - pub fn build(ctx: &mut Context, source: Value, rank: Value) -> Ptr { - let result_ty = IntegerType::get(ctx, 32, Signedness::Unsigned); - Operation::new( - ctx, - Self::get_concrete_op_info(), - vec![result_ty.to_handle()], - vec![source, rank], - vec![], - 0, - ) - } -} - -impl Verify for DsmemReadU32Op { - fn verify(&self, ctx: &Context) -> Result<(), Error> { - let operation = self.get_operation(); - verify_pointer_rank(ctx, operation, "nvvm.dsmem_read_u32")?; - let op = operation.deref(ctx); - let result_ty = op.get_result(0).get_type(ctx); - let result_ty = result_ty.deref(ctx); - if !result_ty - .downcast_ref::() - .is_some_and(|ty| ty.width() == 32 && ty.signedness() == Signedness::Unsigned) - { - return verify_err!(op.loc(), "nvvm.dsmem_read_u32 result must be u32"); - } - Ok(()) - } -} - -pub(super) fn register(ctx: &mut Context) { - MapaSharedClusterOp::register(ctx); - DsmemReadU32Op::register(ctx); -} -"#, - ); - output -} - -fn render_dialect_wgmma_control(catalog: &CatalogFile, hash: &str) -> String { - assert_eq!(wgmma_controls(catalog).count(), 3); - let mut output = rust_header(catalog, hash); - output.push_str( - r##"//! Structural operations for generated WGMMA controls. - -use pliron::{ - builtin::{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; - -#[pliron_op( - name = "nvvm.wgmma_fence_sync_aligned", - format, - verifier = "succ", - interfaces = [NOpdsInterface<0>, NResultsInterface<0>], -)] -pub struct WgmmaFenceSyncAlignedOp; - -impl WgmmaFenceSyncAlignedOp { - 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) - } -} - -#[pliron_op( - name = "nvvm.wgmma_commit_group_sync_aligned", - format, - verifier = "succ", - interfaces = [NOpdsInterface<0>, NResultsInterface<0>], -)] -pub struct WgmmaCommitGroupSyncAlignedOp; - -impl WgmmaCommitGroupSyncAlignedOp { - 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) - } -} - -#[pliron_op( - name = "nvvm.wgmma_wait_group_sync_aligned", - format, - interfaces = [NOpdsInterface<1>, NResultsInterface<0>], -)] -pub struct WgmmaWaitGroupSyncAlignedOp; - -impl WgmmaWaitGroupSyncAlignedOp { - pub fn new(op: Ptr) -> Self { Self { op } } - - pub fn build(ctx: &mut Context, max_pending: Value) -> Ptr { - Operation::new( - ctx, - Self::get_concrete_op_info(), - vec![], - vec![max_pending], - vec![], - 0, - ) - } -} - -impl Verify for WgmmaWaitGroupSyncAlignedOp { - 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(), "nvvm.wgmma_wait_group_sync_aligned takes one operand and no results"); - } - if !is_i64(ctx, op.get_operand(0).get_type(ctx)) { - return verify_err!(op.loc(), "nvvm.wgmma_wait_group_sync_aligned operand must be i64"); - } - Ok(()) - } -} - -fn is_i64(ctx: &Context, ty: TypeHandle) -> bool { - ty.deref(ctx) - .downcast_ref::() - .is_some_and(|integer| integer.width() == 64) -} - -pub(super) fn register(ctx: &mut Context) { - WgmmaFenceSyncAlignedOp::register(ctx); - WgmmaCommitGroupSyncAlignedOp::register(ctx); - WgmmaWaitGroupSyncAlignedOp::register(ctx); -} -"##, - ); - output -} - -fn render_dialect_debug_control(catalog: &CatalogFile, hash: &str) -> String { - assert_eq!(debug_controls(catalog).count(), 3); - let mut output = rust_header(catalog, hash); - output.push_str( - "//! Structural operations for generated PTX debug controls.\n\n\ - use pliron::{\n\ - builtin::{\n\ - attributes::IntegerAttr,\n\ - op_interfaces::{NOpdsInterface, NResultsInterface},\n\ - types::{IntegerType, Signedness},\n\ - },\n\ - common_traits::Verify,\n\ - context::{Context, Ptr},\n\ - identifier::Identifier,\n\ - location::Located,\n\ - op::Op,\n\ - operation::Operation,\n\ - result::Error,\n\ - verify_err,\n\ - };\n\ - use pliron::utils::apint::APInt;\n\ - use pliron_derive::pliron_op;\n\ - use std::num::NonZeroUsize;\n\n\ - #[pliron_op(\n\ - name = \"nvvm.trap\",\n\ - format,\n\ - verifier = \"succ\",\n\ - interfaces = [NOpdsInterface<0>, NResultsInterface<0>],\n\ - )]\n\ - pub struct TrapOp;\n\n\ - impl TrapOp {\n\ - pub fn new(op: Ptr) -> Self { Self { op } }\n\ - pub fn build(ctx: &mut Context) -> Ptr {\n\ - Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![], vec![], 0)\n\ - }\n\ - }\n\n\ - #[pliron_op(\n\ - name = \"nvvm.brkpt\",\n\ - format,\n\ - verifier = \"succ\",\n\ - interfaces = [NOpdsInterface<0>, NResultsInterface<0>],\n\ - )]\n\ - pub struct BreakpointOp;\n\n\ - impl BreakpointOp {\n\ - pub fn new(op: Ptr) -> Self { Self { op } }\n\ - pub fn build(ctx: &mut Context) -> Ptr {\n\ - Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![], vec![], 0)\n\ - }\n\ - }\n\n\ - #[pliron_op(\n\ - name = \"nvvm.pmevent\",\n\ - format,\n\ - interfaces = [NOpdsInterface<0>, NResultsInterface<0>],\n\ - )]\n\ - pub struct PmEventOp;\n\n\ - impl PmEventOp {\n\ - pub fn new(op: Ptr) -> Self { Self { op } }\n\n\ - pub fn build(ctx: &mut Context, event_id: u32) -> Ptr {\n\ - let op = Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![], vec![], 0);\n\ - let ty = IntegerType::get(ctx, 32, Signedness::Unsigned);\n\ - let value = APInt::from_u64(event_id.into(), NonZeroUsize::new(32).unwrap());\n\ - op.deref_mut(ctx).attributes.set(\n\ - Identifier::try_from(\"event_id\").unwrap(),\n\ - IntegerAttr::new(ty, value),\n\ - );\n\ - op\n\ - }\n\n\ - pub fn new_with_event_id(ctx: &mut Context, event_id: u32) -> Ptr {\n\ - Self::build(ctx, event_id)\n\ - }\n\n\ - pub fn event_id(&self, ctx: &Context) -> Option {\n\ - let key = Identifier::try_from(\"event_id\").unwrap();\n\ - let operation = self.get_operation().deref(ctx);\n\ - let attribute: &IntegerAttr = operation.attributes.get(&key)?;\n\ - let ty_handle = attribute.get_type();\n\ - let ty = ty_handle.deref(ctx);\n\ - if ty.width() != 32 || ty.signedness() != Signedness::Unsigned {\n\ - return None;\n\ - }\n\ - u32::try_from(attribute.value().to_u64()).ok().filter(|value| *value <= 15)\n\ - }\n\ -\n\ - pub fn get_event_id(&self, ctx: &Context) -> Option {\n\ - self.event_id(ctx)\n\ - }\n\ - }\n\n\ - impl Verify for PmEventOp {\n\ - fn verify(&self, ctx: &Context) -> Result<(), Error> {\n\ - let op = self.get_operation().deref(ctx);\n\ - if self.event_id(ctx).is_none() {\n\ - return verify_err!(op.loc(), \"nvvm.pmevent requires a u32 event ID in 0..=15\");\n\ - }\n\ - Ok(())\n\ - }\n\ - }\n\n\ - pub(super) fn register(ctx: &mut Context) {\n\ - TrapOp::register(ctx);\n\ - BreakpointOp::register(ctx);\n\ - PmEventOp::register(ctx);\n\ - }\n", - ); - output -} - -fn render_dialect_clc(catalog: &CatalogFile, hash: &str) -> String { - assert_eq!(clc_intrinsics(catalog).count(), 6); - let mut output = rust_header(catalog, hash); - output.push_str( - "//! Generated Cluster Launch Control operations.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::IntegerType,\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::{TypeHandle, Typed},\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\n", - ); - for record in clc_intrinsics(catalog) { - let result_count = record.dialect.results.len(); - let verifier = if result_count == 0 { - " verifier = \"succ\",\n" - } else { - "" - }; - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!( - output, - "#[pliron_op(\n name = {:?},\n format,\n{verifier} interfaces = [NOpdsInterface<2>, NResultsInterface<{result_count}>],\n)]", - record.dialect.op_name, - ) - .unwrap(); - writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); - writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); - output.push_str(" pub fn new(op: Ptr) -> Self { Self { op } }\n}\n\n"); - if result_count == 1 { - writeln!(output, "impl Verify for {} {{", record.dialect.op_type).unwrap(); - writeln!( - output, - " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n verify_query_shape(ctx, self.get_operation(), {:?})\n }}\n}}\n", - record.dialect.op_name, - ) - .unwrap(); - } - } - output.push_str( - "fn verify_query_shape(\n ctx: &Context,\n operation: Ptr,\n name: &str,\n) -> Result<(), Error> {\n let op = operation.deref(ctx);\n let valid = op.get_num_operands() == 2\n && op.get_num_results() == 1\n && is_integer_width(ctx, op.get_operand(0).get_type(ctx), 64)\n && is_integer_width(ctx, op.get_operand(1).get_type(ctx), 64)\n && is_integer_width(ctx, op.get_result(0).get_type(ctx), 32);\n if !valid {\n return verify_err!(\n op.loc(),\n \"{} requires two i64 operands and one i32 result\",\n name\n );\n }\n Ok(())\n}\n\nfn is_integer_width(ctx: &Context, ty: TypeHandle, width: u32) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == width)\n}\n\npub(super) fn register(ctx: &mut Context) {\n", - ); - for record in clc_intrinsics(catalog) { - writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); - } - output.push_str("}\n"); - output -} - -fn render_dialect_execution_control(catalog: &CatalogFile, hash: &str) -> String { - assert_eq!(execution_controls(catalog).count(), 8); - let mut output = rust_header(catalog, hash); - output.push_str( - "//! Generated counted-barrier, grid-dependency, and register-control operations.\n\nuse pliron::{\n builtin::{attributes::IntegerAttr, op_interfaces::{NOpdsInterface, NResultsInterface}, types::{IntegerType, Signedness}},\n common_traits::Verify,\n context::{Context, Ptr},\n identifier::Identifier,\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n verify_err,\n};\nuse pliron::utils::apint::APInt;\nuse pliron_derive::pliron_op;\nuse std::num::NonZeroUsize;\n\n", - ); - for record in execution_controls(catalog).filter(|record| record.family != "register_control") { - let operand_count = record.dialect.operands.len(); - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!( - output, - "#[pliron_op(\n name = {:?},\n format,\n verifier = \"succ\",\n interfaces = [NOpdsInterface<{operand_count}>, NResultsInterface<0>],\n)]", - record.dialect.op_name - ) - .unwrap(); - writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); - writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); - output.push_str(" pub fn new(op: Ptr) -> Self { Self { op } }\n}\n\n"); - } - for record in execution_control_family(catalog, "register_control") { - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!( - output, - "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<0>, NResultsInterface<0>],\n)]", - record.dialect.op_name - ) - .unwrap(); - writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); - writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); - output.push_str( - " pub fn new(op: Ptr) -> Self { Self { op } }\n\n pub fn build(ctx: &mut Context, register_count: u32) -> Ptr {\n let op = Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![], vec![], 0);\n let ty = IntegerType::get(ctx, 32, Signedness::Unsigned);\n let value = APInt::from_u64(register_count.into(), NonZeroUsize::new(32).unwrap());\n op.deref_mut(ctx).attributes.set(\n Identifier::try_from(\"register_count\").unwrap(),\n IntegerAttr::new(ty, value),\n );\n op\n }\n\n pub fn register_count(&self, ctx: &Context) -> Option {\n let key = Identifier::try_from(\"register_count\").unwrap();\n let operation = self.get_operation().deref(ctx);\n let attribute: &IntegerAttr = operation.attributes.get(&key)?;\n let ty_handle = attribute.get_type();\n let ty = ty_handle.deref(ctx);\n if ty.width() != 32 || ty.signedness() != Signedness::Unsigned {\n return None;\n }\n u32::try_from(attribute.value().to_u64()).ok().filter(|value| (24..=256).contains(value) && value % 8 == 0)\n }\n}\n\n", - ); - writeln!(output, "impl Verify for {} {{", record.dialect.op_type).unwrap(); - output.push_str( - " fn verify(&self, ctx: &Context) -> Result<(), Error> {\n let op = self.get_operation().deref(ctx);\n if self.register_count(ctx).is_none() {\n return verify_err!(op.loc(), \"setmaxnreg requires an immediate register count in 24..=256 divisible by 8\");\n }\n Ok(())\n }\n}\n\n", - ); - } - output.push_str("pub(super) fn register(ctx: &mut Context) {\n"); - for record in execution_controls(catalog) { - writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); - } - output.push_str("}\n"); - output -} - -fn render_dialect_tma(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - "//! Generated Tensor Memory Accelerator operations.\n\nuse pliron::{\n builtin::op_interfaces::{NOpdsInterface, NResultsInterface},\n context::{Context, Ptr},\n op::Op,\n operation::Operation,\n};\nuse pliron_derive::pliron_op;\n\n", - ); - for record in tma_intrinsics(catalog) { - let operand_count = record.dialect.operands.len(); - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!( - output, - "#[pliron_op(\n name = {:?},\n format,\n verifier = \"succ\",\n interfaces = [NOpdsInterface<{operand_count}>, NResultsInterface<0>],\n)]", - record.dialect.op_name - ) - .unwrap(); - writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); - writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); - output.push_str(" pub fn new(op: Ptr) -> Self { Self { op } }\n}\n\n"); - } - output.push_str("pub(super) fn register(ctx: &mut Context) {\n"); - for record in tma_intrinsics(catalog) { - writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); - } - output.push_str("}\n"); - output -} - -fn tcgen05_carrier_name(ty: &str) -> &'static str { - match ty { - "ptr" => "Ptr", - "i1" => "I1", - "i16" => "I16", - "i32" => "I32", - "i64" => "I64", - "f32" => "F32", - _ => panic!("unsupported tcgen05 dialect carrier {ty}"), - } -} - -fn render_tcgen05_carrier_runs(types: &[String]) -> String { - let mut runs = Vec::<(&str, usize)>::new(); - for ty in types { - let carrier = tcgen05_carrier_name(ty); - if let Some((previous, count)) = runs.last_mut() - && *previous == carrier - { - *count += 1; - } else { - runs.push((carrier, 1)); - } - } - - let mut output = String::from("&["); - for (index, (carrier, count)) in runs.into_iter().enumerate() { - if index != 0 { - output.push_str(", "); - } - write!(output, "(Tcgen05Carrier::{carrier}, {count})").unwrap(); - } - output.push(']'); - output -} - -fn render_dialect_tcgen05(catalog: &CatalogFile, hash: &str) -> String { - assert_eq!(tcgen05_intrinsics(catalog).count(), 233); - let mut output = rust_header(catalog, hash); - output.push_str( - r#"//! Generated Tensor Core Generation 5 operations. - -use dialect_mir::{ops::MirConstantOp, types::MirPtrType}; -use pliron::{ - attribute::Attribute, - builtin::{ - op_interfaces::{NOpdsInterface, NResultsInterface}, - ops::ConstantOp, - types::{FP32Type, IntegerType}, - }, - common_traits::Verify, - context::{Context, Ptr}, - location::Located, - op::Op, - operation::Operation, - result::Error, - r#type::{TypeHandle, Typed}, - verify_err, -}; -use pliron_derive::{pliron_attr, pliron_op}; - -#[derive(Clone, Copy)] -enum Tcgen05Carrier { - Ptr, - I1, - I16, - I32, - I64, - F32, -} - -impl Tcgen05Carrier { - fn matches(self, ctx: &Context, ty: TypeHandle) -> bool { - match self { - Self::Ptr => ty.deref(ctx).downcast_ref::().is_some(), - Self::I1 => is_integer_width(ctx, ty, 1), - Self::I16 => is_integer_width(ctx, ty, 16), - Self::I32 => is_integer_width(ctx, ty, 32), - Self::I64 => is_integer_width(ctx, ty, 64), - Self::F32 => ty.deref(ctx).downcast_ref::().is_some(), - } - } - - fn name(self) -> &'static str { - match self { - Self::Ptr => "MIR pointer", - Self::I1 => "i1", - Self::I16 => "i16", - Self::I32 => "i32", - Self::I64 => "i64", - Self::F32 => "f32", - } - } -} - -fn is_integer_width(ctx: &Context, ty: TypeHandle, width: u32) -> bool { - ty.deref(ctx) - .downcast_ref::() - .is_some_and(|integer| integer.width() == width) -} - -fn expected_carrier_count(runs: &[(Tcgen05Carrier, usize)]) -> usize { - runs.iter().map(|(_, count)| *count).sum() -} - -fn verify_tcgen05_signature( - ctx: &Context, - operation: Ptr, - name: &str, - operand_runs: &[(Tcgen05Carrier, usize)], - result_runs: &[(Tcgen05Carrier, usize)], - constant_operand: Option, -) -> Result<(), Error> { - let op = operation.deref(ctx); - let expected_operands = expected_carrier_count(operand_runs); - let expected_results = expected_carrier_count(result_runs); - if op.get_num_operands() != expected_operands || op.get_num_results() != expected_results { - return verify_err!( - op.loc(), - "{name} requires {expected_operands} operands and {expected_results} results" - ); - } - - let mut index = 0; - for (carrier, count) in operand_runs { - for _ in 0..*count { - if !carrier.matches(ctx, op.get_operand(index).get_type(ctx)) { - return verify_err!( - op.loc(), - "{name} operand {index} must be {}", - carrier.name() - ); - } - index += 1; - } - } - - index = 0; - for (carrier, count) in result_runs { - for _ in 0..*count { - if !carrier.matches(ctx, op.get_result(index).get_type(ctx)) { - return verify_err!( - op.loc(), - "{name} result {index} must be {}", - carrier.name() - ); - } - index += 1; - } - } - - if let Some(index) = constant_operand { - let value = op.get_operand(index); - let is_constant = value.defining_op().is_some_and(|defining_op| { - Operation::get_op::(defining_op, ctx).is_some() - || Operation::get_op::(defining_op, ctx).is_some() - }); - if !is_constant { - return verify_err!( - op.loc(), - "{name} operand {index} must be a compile-time constant" - ); - } - } - - Ok(()) -} - -"#, - ); - if tcgen05_mma_intrinsics(catalog).next().is_some() { - output.push_str( - r#"#[pliron_attr(name = "nvvm.tcgen05_mma_form", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum Tcgen05MmaFormAttr { - Shared, Tensor, TensorAshift, SpShared, SpTensor, SpTensorAshift, - WsShared, WsSharedZeroColMask, WsSpShared, WsSpSharedZeroColMask, - WsSpTensor, WsSpTensorZeroColMask, WsTensor, WsTensorZeroColMask, -} - -#[pliron_attr(name = "nvvm.tcgen05_mma_kind", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum Tcgen05MmaKindAttr { F16, Tf32, F8f6f4, I8 } - -#[pliron_attr(name = "nvvm.tcgen05_mma_cta_group", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum Tcgen05MmaCtaGroupAttr { Cg1, Cg2 } - -#[pliron_attr(name = "nvvm.tcgen05_mma_collector_a", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum Tcgen05MmaCollectorAAttr { Discard, LastUse, Fill, Use } - -#[pliron_attr(name = "nvvm.tcgen05_mma_b_buffer", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum Tcgen05MmaBBufferAttr { B0, B1, B2, B3 } - -#[pliron_attr(name = "nvvm.tcgen05_mma_b_usage", format, verifier = "succ")] -#[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum Tcgen05MmaBUsageAttr { Discard, LastUse, Fill, Use } - -#[pliron_op( - name = "nvvm.tcgen05_mma", - format, - attributes = ( - nvvm_tcgen05_mma_form: Tcgen05MmaFormAttr, - nvvm_tcgen05_mma_kind: Tcgen05MmaKindAttr, - nvvm_tcgen05_mma_cta_group: Tcgen05MmaCtaGroupAttr, - nvvm_tcgen05_mma_collector_a: Tcgen05MmaCollectorAAttr, - nvvm_tcgen05_mma_b_buffer: Tcgen05MmaBBufferAttr, - nvvm_tcgen05_mma_b_usage: Tcgen05MmaBUsageAttr - ) -)] -pub struct Tcgen05MmaOp; - -impl Tcgen05MmaOp { - pub fn new(op: Ptr) -> Self { Self { op } } -} - -impl Verify for Tcgen05MmaOp { - fn verify(&self, ctx: &Context) -> Result<(), Error> { - use Tcgen05Carrier::{I1, I32, I64}; - let op = self.get_operation(); - let form = self.get_attr_nvvm_tcgen05_mma_form(ctx); - let Some(form) = form.as_deref() else { - return verify_err!(op.deref(ctx).loc(), "nvvm.tcgen05_mma requires a form"); - }; - if self.get_attr_nvvm_tcgen05_mma_kind(ctx).is_none() { - return verify_err!(op.deref(ctx).loc(), "nvvm.tcgen05_mma requires a kind"); - } - let (operands, base): (&[(Tcgen05Carrier, usize)], bool) = match form { - Tcgen05MmaFormAttr::Shared | Tcgen05MmaFormAttr::WsShared => - (&[(I32, 1), (I64, 2), (I32, 1), (I1, 1)], matches!(form, Tcgen05MmaFormAttr::Shared)), - Tcgen05MmaFormAttr::Tensor | Tcgen05MmaFormAttr::TensorAshift | Tcgen05MmaFormAttr::WsTensor => - (&[(I32, 2), (I64, 1), (I32, 1), (I1, 1)], !matches!(form, Tcgen05MmaFormAttr::WsTensor)), - Tcgen05MmaFormAttr::SpShared | Tcgen05MmaFormAttr::WsSpShared => - (&[(I32, 1), (I64, 2), (I32, 1), (I1, 1), (I32, 1)], matches!(form, Tcgen05MmaFormAttr::SpShared)), - Tcgen05MmaFormAttr::SpTensor | Tcgen05MmaFormAttr::SpTensorAshift | Tcgen05MmaFormAttr::WsSpTensor => - (&[(I32, 2), (I64, 1), (I32, 1), (I1, 1), (I32, 1)], !matches!(form, Tcgen05MmaFormAttr::WsSpTensor)), - Tcgen05MmaFormAttr::WsSharedZeroColMask => - (&[(I32, 1), (I64, 2), (I32, 1), (I1, 1), (I64, 1)], false), - Tcgen05MmaFormAttr::WsSpSharedZeroColMask => - (&[(I32, 1), (I64, 2), (I32, 1), (I1, 1), (I32, 1), (I64, 1)], false), - Tcgen05MmaFormAttr::WsSpTensorZeroColMask => - (&[(I32, 2), (I64, 1), (I32, 1), (I1, 1), (I32, 1), (I64, 1)], false), - Tcgen05MmaFormAttr::WsTensorZeroColMask => - (&[(I32, 2), (I64, 1), (I32, 1), (I1, 1), (I64, 1)], false), - }; - verify_tcgen05_signature(ctx, op, "nvvm.tcgen05_mma", operands, &[], None)?; - - let cta_group = self.get_attr_nvvm_tcgen05_mma_cta_group(ctx); - let collector_a = self.get_attr_nvvm_tcgen05_mma_collector_a(ctx); - let b_buffer = self.get_attr_nvvm_tcgen05_mma_b_buffer(ctx); - let b_usage = self.get_attr_nvvm_tcgen05_mma_b_usage(ctx); - if base { - if cta_group.is_none() || collector_a.is_none() || b_buffer.is_some() || b_usage.is_some() { - return verify_err!(op.deref(ctx).loc(), "base tcgen05 MMA requires CTA-group and collector-A selectors"); - } - if matches!(form, Tcgen05MmaFormAttr::TensorAshift | Tcgen05MmaFormAttr::SpTensorAshift) - && !matches!(collector_a.as_deref(), Some(Tcgen05MmaCollectorAAttr::Discard | Tcgen05MmaCollectorAAttr::LastUse)) - { - return verify_err!(op.deref(ctx).loc(), "ashift tcgen05 MMA only supports discard or last-use collector-A"); - } - } else if cta_group.is_some() || collector_a.is_some() || b_buffer.is_none() || b_usage.is_none() { - return verify_err!(op.deref(ctx).loc(), "warp-specialized tcgen05 MMA requires B-buffer and B-usage selectors"); - } - Ok(()) - } -} - -"#, - ); - } - for record in tcgen05_non_mma_intrinsics(catalog) { - let operand_count = record.dialect.operands.len(); - let result_count = record.dialect.results.len(); - let operand_runs = render_tcgen05_carrier_runs(&record.dialect.operands); - let result_runs = render_tcgen05_carrier_runs(&record.dialect.results); - let constant_operand = record - .tcgen05 - .as_ref() - .is_some_and(|tcgen05| { - matches!( - tcgen05.adapter, - Tcgen05Adapter::TmemHalfSplitOffsetInjectPack16ToU32Registers - | Tcgen05Adapter::TmemHalfSplitOffsetU32RegistersInjectUnpack16ToVoid - ) - }) - .then_some(1); - writeln!(output, "/// {}", record.summary).unwrap(); - writeln!( - output, - "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<{operand_count}>, NResultsInterface<{result_count}>],\n)]", - record.dialect.op_name - ) - .unwrap(); - writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); - writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); - output.push_str(" pub fn new(op: Ptr) -> Self { Self { op } }\n}\n\n"); - writeln!(output, "impl Verify for {} {{", record.dialect.op_type).unwrap(); - output.push_str(" fn verify(&self, ctx: &Context) -> Result<(), Error> {\n"); - output.push_str(" verify_tcgen05_signature(\n"); - output.push_str(" ctx,\n"); - output.push_str(" self.get_operation(),\n"); - writeln!(output, " {:?},", record.dialect.op_name).unwrap(); - writeln!(output, " {operand_runs},").unwrap(); - writeln!(output, " {result_runs},").unwrap(); - writeln!(output, " {constant_operand:?},").unwrap(); - output.push_str(" )\n"); - output.push_str(" }\n"); - output.push_str("}\n\n"); - } - output.push_str("pub(super) fn register(ctx: &mut Context) {\n"); - if tcgen05_mma_intrinsics(catalog).next().is_some() { - output.push_str( - " Tcgen05MmaFormAttr::register(ctx);\n Tcgen05MmaKindAttr::register(ctx);\n Tcgen05MmaCtaGroupAttr::register(ctx);\n Tcgen05MmaCollectorAAttr::register(ctx);\n Tcgen05MmaBBufferAttr::register(ctx);\n Tcgen05MmaBUsageAttr::register(ctx);\n Tcgen05MmaOp::register(ctx);\n", - ); - } - for record in tcgen05_non_mma_intrinsics(catalog) { - writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); - } - output.push_str("}\n"); - output -} - -fn render_importer_pure_value_dispatch( - output: &mut String, - catalog: &CatalogFile, - record: &CatalogIntrinsic, -) { - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(output, &path_refs); - output.push_str(" => {\n"); - writeln!( - output, - " require_arity(name, args.len(), {}, &loc)?;", - record.rust.arguments.len() - ) - .unwrap(); - for index in 0..record.rust.arguments.len() { - let previous = if index == 0 { "prev_op" } else { "last_op" }; - writeln!( - output, - " let (arg{index}, last_op) = rvalue::translate_operand(\n ctx, body, &args[{index}], value_map, block_ptr, {previous}, loc.clone(),\n )?;" - ) - .unwrap(); - } - let arguments = (0..record.rust.arguments.len()) - .map(|index| format!("arg{index}")) - .collect::>() - .join(", "); - writeln!( - output, - " let intrinsic = {}::build(ctx, {arguments});", - record.dialect.op_type - ) - .unwrap(); - output.push_str(" intrinsic.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, intrinsic, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, intrinsic, block_ptr, last_op);\n let result = intrinsic.deref(ctx).get_result(0);\n", - ); - writeln!( - output, - " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, intrinsic, value_map, block_map, loc,\n {:?},\n )?))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n"); -} - -fn render_importer_scalar_conversion_dispatch( - output: &mut String, - catalog: &CatalogFile, - record: &CatalogIntrinsic, -) { - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(output, &path_refs); - output.push_str(" => {\n"); - output.push_str( - " require_arity(name, args.len(), 1, &loc)?;\n\ - let (value, last_op) = rvalue::translate_operand(\n\ - ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n\ - )?;\n", - ); - writeln!( - output, - " let intrinsic = ScalarConversionOp::build(ctx, value, {}, {});", - scalar_conversion_rounding_attr(record), - scalar_conversion_saturation_attr(record), - ) - .unwrap(); - output.push_str(" intrinsic.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, intrinsic, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, intrinsic, block_ptr, last_op);\n\ - let result = intrinsic.deref(ctx).get_result(0);\n", - ); - writeln!( - output, - " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, intrinsic, value_map, block_map, loc,\n {:?},\n )?))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n"); -} - -fn render_importer_scalar_arithmetic_dispatch( - output: &mut String, - catalog: &CatalogFile, - record: &CatalogIntrinsic, -) { - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(output, &path_refs); - output.push_str(" => {\n"); - let arity = scalar_arithmetic_arity(record); - writeln!( - output, - " require_arity(name, args.len(), {arity}, &loc)?;" - ) - .unwrap(); - for index in 0..arity { - let previous = if index == 0 { "prev_op" } else { "last_op" }; - writeln!( - output, - " let (arg{index}, last_op) = rvalue::translate_operand(\n ctx, body, &args[{index}], value_map, block_ptr, {previous}, loc.clone(),\n )?;" - ) - .unwrap(); - } - let arguments = (0..arity) - .map(|index| format!("arg{index}")) - .collect::>() - .join(", "); - writeln!( - output, - " let intrinsic = ScalarArithmeticOp::build(ctx, vec![{arguments}], {}, {}, {}, {}, {});", - scalar_arithmetic_format_attr(record), - scalar_arithmetic_operation_attr(record), - scalar_arithmetic_rounding_attr(record), - scalar_arithmetic_subnormal_attr(record), - scalar_arithmetic_saturation_attr(record), - ) - .unwrap(); - output.push_str(" intrinsic.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, intrinsic, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, intrinsic, block_ptr, last_op);\n\ - let result = intrinsic.deref(ctx).get_result(0);\n", - ); - writeln!( - output, - " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, intrinsic, value_map, block_map, loc,\n {:?},\n )?))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n"); -} - -fn render_importer_scalar_math_dispatch( - output: &mut String, - catalog: &CatalogFile, - record: &CatalogIntrinsic, -) { - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(output, &path_refs); - output.push_str(" => {\n"); - writeln!( - output, - " require_arity(name, args.len(), 1, &loc)?;" - ) - .unwrap(); - writeln!( - output, - " let (arg0, last_op) = rvalue::translate_operand(\n ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n )?;" - ) - .unwrap(); - writeln!( - output, - " let intrinsic = ScalarMathOp::build(ctx, arg0, {}, {}, {}, {});", - scalar_math_format_attr(record), - scalar_math_operation_attr(record), - scalar_math_precision_attr(record), - scalar_math_subnormal_attr(record), - ) - .unwrap(); - output.push_str(" intrinsic.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, intrinsic, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, intrinsic, block_ptr, last_op);\n\ - let result = intrinsic.deref(ctx).get_result(0);\n", - ); - writeln!( - output, - " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, intrinsic, value_map, block_map, loc,\n {:?},\n )?))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n"); -} - -fn render_importer_extended_minmax_dispatch( - output: &mut String, - catalog: &CatalogFile, - record: &CatalogIntrinsic, -) { - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(output, &path_refs); - output.push_str(" => {\n"); - output.push_str( - " require_arity(name, args.len(), 2, &loc)?;\n\ - let (a, last_op) = rvalue::translate_operand(\n\ - ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n\ - )?;\n\ - let (b, last_op) = rvalue::translate_operand(\n\ - ctx, body, &args[1], value_map, block_ptr, last_op, loc.clone(),\n\ - )?;\n", - ); - writeln!( - output, - " let intrinsic = ExtendedMinMaxOp::build(ctx, a, b, {}, {}, {}, {}, {});", - extended_minmax_format_attr(record), - extended_minmax_operation_attr(record), - extended_minmax_subnormal_attr(record), - extended_minmax_nan_attr(record), - extended_minmax_xorsign_abs_attr(record), - ) - .unwrap(); - output.push_str(" intrinsic.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, intrinsic, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, intrinsic, block_ptr, last_op);\n\ - let result = intrinsic.deref(ctx).get_result(0);\n", - ); - writeln!( - output, - " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, intrinsic, value_map, block_map, loc,\n {:?},\n )?))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n"); -} - -fn render_importer_elect_dispatch( - output: &mut String, - catalog: &CatalogFile, - record: &CatalogIntrinsic, -) { - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(output, &path_refs); - output.push_str(" => {\n"); - output.push_str( - " require_arity(name, args.len(), 1, &loc)?;\n\ - \n\ - let tuple_ty = crate::translator::types::translate_destination_type(\n\ - ctx, body, destination, &loc,\n\ - )?;\n\ - let (leader_ty, elected_ty) = {\n\ - let ty = tuple_ty.deref(ctx);\n\ - match ty.downcast_ref::() {\n\ - Some(tuple) if tuple.get_types().len() == 2 => {\n\ - (tuple.get_types()[0], tuple.get_types()[1])\n\ - }\n\ - _ => {\n\ - return input_err!(\n\ - loc.clone(),\n\ - TranslationErr::unsupported(\n\ - \"warp::elect_sync destination must be a (u32, bool) tuple\"\n\ - .to_owned()\n\ - )\n\ - );\n\ - }\n\ - }\n\ - };\n\ - let (mask, last_op) = rvalue::translate_operand(\n\ - ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n\ - )?;\n", - ); - writeln!( - output, - " let elect = Operation::new(\n ctx, {}::get_concrete_op_info(), vec![leader_ty, elected_ty],\n vec![mask], vec![], 0,\n );", - record.dialect.op_type - ) - .unwrap(); - output.push_str(" elect.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, elect, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, elect, block_ptr, last_op);\n\ - let leader = elect.deref(ctx).get_result(0);\n\ - let elected = elect.deref(ctx).get_result(1);\n\ - let tuple = Operation::new(\n\ - ctx, MirConstructTupleOp::get_concrete_op_info(), vec![tuple_ty],\n\ - vec![leader, elected],\n\ - vec![], 0,\n\ - );\n\ - tuple.deref_mut(ctx).set_loc(loc.clone());\n\ - tuple.insert_after(ctx, elect);\n\ - let result = tuple.deref(ctx).get_result(0);\n", - ); - writeln!( - output, - " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, tuple, value_map, block_map, loc,\n {:?},\n )?))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n"); -} - -fn render_importer_tcgen05_mma_dispatch( - output: &mut String, - catalog: &CatalogFile, - record: &CatalogIntrinsic, -) { - let mma = record - .tcgen05 - .as_ref() - .and_then(|tcgen05| tcgen05.mma.as_ref()) - .expect("tcgen05 MMA record"); - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(output, &path_refs); - output.push_str(" => {\n"); - writeln!( - output, - " require_arity(name, args.len(), {}, &loc)?;", - record.rust.arguments.len() - ) - .unwrap(); - - let runtime_indices = if mma.alias.is_some() && mma.form == Tcgen05MmaForm::WsTensor { - vec![0, 1, 3, 4, 5] - } else if mma.alias.is_some() { - (0..record.rust.arguments.len()).collect() - } else { - let selector_indices = match mma.selector_layout { - Tcgen05MmaSelectorLayout::Base { - kind_argument, - cta_group_argument, - collector_a_argument, - .. - } => [ - kind_argument as usize, - cta_group_argument as usize, - collector_a_argument as usize, - ], - Tcgen05MmaSelectorLayout::WarpSpecialized { - kind_argument, - b_buffer_argument, - b_usage_argument, - } => [ - kind_argument as usize, - b_buffer_argument as usize, - b_usage_argument as usize, - ], - }; - (0..record.rust.arguments.len()) - .filter(|index| !selector_indices.contains(index)) - .collect() - }; - writeln!( - output, - " let runtime_indices: &[usize] = &{runtime_indices:?};" - ) - .unwrap(); - output.push_str( - " let mut last_op = prev_op;\n let mut operands = Vec::with_capacity(runtime_indices.len());\n for &index in runtime_indices {\n let (value, translated) = rvalue::translate_operand(\n ctx, body, &args[index], value_map, block_ptr, last_op, loc.clone(),\n )?;\n last_op = translated;\n operands.push(value);\n }\n", - ); - - if let Some(fixed) = mma.fixed_selectors { - writeln!( - output, - " let kind = {};", - tcgen05_mma_kind_attr(fixed.kind) - ) - .unwrap(); - match mma.selector_layout { - Tcgen05MmaSelectorLayout::Base { .. } => { - output.push_str( - " let cta_group = Tcgen05MmaCtaGroupAttr::Cg1;\n let collector_a = Tcgen05MmaCollectorAAttr::Discard;\n", - ); - } - Tcgen05MmaSelectorLayout::WarpSpecialized { .. } => { - let b_buffer = match fixed.b_buffer { - 0 => "Tcgen05MmaBBufferAttr::B0", - 1 => "Tcgen05MmaBBufferAttr::B1", - 2 => "Tcgen05MmaBBufferAttr::B2", - 3 => "Tcgen05MmaBBufferAttr::B3", - _ => panic!("admitted tcgen05 MMA B buffer"), - }; - writeln!(output, " let b_buffer = {b_buffer};").unwrap(); - writeln!( - output, - " let b_usage = {};", - tcgen05_mma_b_usage_attr(fixed.b_usage) - ) - .unwrap(); - } - } - } else { - let (kind_argument, second_argument, third_argument) = match mma.selector_layout { - Tcgen05MmaSelectorLayout::Base { - kind_argument, - cta_group_argument, - collector_a_argument, - .. - } => (kind_argument, cta_group_argument, collector_a_argument), - Tcgen05MmaSelectorLayout::WarpSpecialized { - kind_argument, - b_buffer_argument, - b_usage_argument, - } => (kind_argument, b_buffer_argument, b_usage_argument), - }; - writeln!( - output, - " let kind = match generated_tcgen05_mma_selector(args.get({kind_argument}usize)) {{\n Some(0) => Tcgen05MmaKindAttr::F16,\n Some(1) => Tcgen05MmaKindAttr::Tf32,\n Some(2) => Tcgen05MmaKindAttr::F8f6f4,\n Some(3) => Tcgen05MmaKindAttr::I8,\n _ => return input_err!(loc.clone(), TranslationErr::unsupported(\"tcgen05 MMA kind must be a compile-time u32 constant in 0..=3\".to_owned())),\n }};" - ) - .unwrap(); - match mma.selector_layout { - Tcgen05MmaSelectorLayout::Base { - collector_a_upper_exclusive, - .. - } => { - writeln!( - output, - " let cta_group = match generated_tcgen05_mma_selector(args.get({second_argument}usize)) {{\n Some(1) => Tcgen05MmaCtaGroupAttr::Cg1,\n Some(2) => Tcgen05MmaCtaGroupAttr::Cg2,\n _ => {{\n return input_err!(\n loc.clone(),\n TranslationErr::unsupported(\n \"tcgen05 MMA CTA group must be the compile-time u32 constant 1 or 2\"\n .to_owned()\n )\n );\n }}\n }};" - ) - .unwrap(); - let collector_arms = if collector_a_upper_exclusive == 2 { - " Some(0) => Tcgen05MmaCollectorAAttr::Discard,\n Some(1) => Tcgen05MmaCollectorAAttr::LastUse,\n" - } else { - " Some(0) => Tcgen05MmaCollectorAAttr::Discard,\n Some(1) => Tcgen05MmaCollectorAAttr::LastUse,\n Some(2) => Tcgen05MmaCollectorAAttr::Fill,\n Some(3) => Tcgen05MmaCollectorAAttr::Use,\n" - }; - writeln!( - output, - " let collector_a = match generated_tcgen05_mma_selector(args.get({third_argument}usize)) {{\n{collector_arms} _ => return input_err!(loc.clone(), TranslationErr::unsupported(\"tcgen05 MMA collector-A selector is outside its closed range\".to_owned())),\n }};" - ) - .unwrap(); - } - Tcgen05MmaSelectorLayout::WarpSpecialized { .. } => { - writeln!( - output, - " let b_buffer = match generated_tcgen05_mma_selector(args.get({second_argument}usize)) {{\n Some(0) => Tcgen05MmaBBufferAttr::B0,\n Some(1) => Tcgen05MmaBBufferAttr::B1,\n Some(2) => Tcgen05MmaBBufferAttr::B2,\n Some(3) => Tcgen05MmaBBufferAttr::B3,\n _ => return input_err!(loc.clone(), TranslationErr::unsupported(\"tcgen05 MMA B buffer must be a compile-time u32 constant in 0..=3\".to_owned())),\n }};" - ) - .unwrap(); - writeln!( - output, - " let b_usage = match generated_tcgen05_mma_selector(args.get({third_argument}usize)) {{\n Some(0) => Tcgen05MmaBUsageAttr::Discard,\n Some(1) => Tcgen05MmaBUsageAttr::LastUse,\n Some(2) => Tcgen05MmaBUsageAttr::Fill,\n Some(3) => Tcgen05MmaBUsageAttr::Use,\n _ => return input_err!(loc.clone(), TranslationErr::unsupported(\"tcgen05 MMA B usage must be a compile-time u32 constant in 0..=3\".to_owned())),\n }};" - ) - .unwrap(); - } - } - } - - output.push_str( - " let intrinsic = Operation::new(ctx, Tcgen05MmaOp::get_concrete_op_info(), vec![], operands, vec![], 0);\n intrinsic.deref_mut(ctx).set_loc(loc.clone());\n let mma_op = Tcgen05MmaOp::new(intrinsic);\n", - ); - writeln!( - output, - " mma_op.set_attr_nvvm_tcgen05_mma_form(ctx, {});", - tcgen05_mma_form_attr(mma.form) - ) - .unwrap(); - output.push_str(" mma_op.set_attr_nvvm_tcgen05_mma_kind(ctx, kind);\n"); - match mma.selector_layout { - Tcgen05MmaSelectorLayout::Base { .. } => output.push_str( - " mma_op.set_attr_nvvm_tcgen05_mma_cta_group(ctx, cta_group);\n mma_op.set_attr_nvvm_tcgen05_mma_collector_a(ctx, collector_a);\n", - ), - Tcgen05MmaSelectorLayout::WarpSpecialized { .. } => output.push_str( - " mma_op.set_attr_nvvm_tcgen05_mma_b_buffer(ctx, b_buffer);\n mma_op.set_attr_nvvm_tcgen05_mma_b_usage(ctx, b_usage);\n", - ), - } - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, intrinsic, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, intrinsic, block_ptr, last_op);\n if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, intrinsic, block_map, loc)))\n } else {\n", - ); - writeln!( - output, - " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n }\n"); -} - -fn render_importer(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - "//! Generated raw/compatibility path dispatch for CUDA intrinsics.\n\nuse crate::error::{TranslationErr, TranslationResult};\nuse crate::translator::{rvalue, terminator::helpers};\nuse crate::translator::values::ValueMap;\nuse dialect_nvvm::ops::{", - ); - for (index, record) in sregs(catalog).enumerate() { - if index != 0 { - output.push_str(", "); - } - output.push_str(&record.dialect.op_type); - } - if ldmatrix(catalog).next().is_some() { - if sregs(catalog).next().is_some() { - output.push_str(", "); - } - output.push_str("LdmatrixElementAttr, LdmatrixLayoutAttr, LdmatrixMultiplicityAttr, LdmatrixOp, LdmatrixShapeAttr, LdmatrixStateSpaceAttr"); - } - if let Some(record) = movmatrix(catalog).next() { - output.push_str(", "); - output.push_str(&record.dialect.op_type); - } - if register_mmas(catalog).next().is_some() { - output.push_str(", RegisterMmaAccumulatorAttr, RegisterMmaElementAttr, RegisterMmaKindAttr, RegisterMmaLayoutAttr, RegisterMmaOp, RegisterMmaOperationAttr, RegisterMmaOverflowAttr, RegisterMmaShapeAttr"); - } - if sparse_mmas(catalog).next().is_some() { - output.push_str(", SparseMmaAccumulatorAttr, SparseMmaElementAttr, SparseMmaLayoutAttr, SparseMmaMetadataAttr, SparseMmaOp, SparseMmaOverflowAttr, SparseMmaSelectorAttr, SparseMmaShapeAttr"); - } - if prmts(catalog).next().is_some() { - output.push_str(", PrmtModeAttr, PrmtOp"); - } - if scalar_conversions(catalog).next().is_some() { - output.push_str( - ", ScalarConversionOp, ScalarConversionRoundingAttr, ScalarConversionSaturationAttr", - ); - } - if scalar_arithmetics(catalog).next().is_some() { - output.push_str( - ", ScalarArithmeticFormatAttr, ScalarArithmeticOp, ScalarArithmeticOperationAttr, ScalarArithmeticRoundingAttr, ScalarArithmeticSaturationAttr, ScalarArithmeticSubnormalAttr", - ); - } - if scalar_maths(catalog).next().is_some() { - output.push_str( - ", ScalarMathFormatAttr, ScalarMathOp, ScalarMathOperationAttr, ScalarMathPrecisionAttr, ScalarMathSubnormalAttr", - ); - } - if extended_minmax(catalog).next().is_some() { - output.push_str( - ", ExtendedMinMaxFormatAttr, ExtendedMinMaxNanAttr, ExtendedMinMaxOp, ExtendedMinMaxOperationAttr, ExtendedMinMaxSubnormalAttr, ExtendedMinMaxXorSignAbsAttr", - ); - } - if cluster_barriers(catalog).next().is_some() { - output.push_str(", ClusterBarrierModeAttr, ClusterBarrierOp"); - } - if debug_controls(catalog).next().is_some() { - output.push_str(", BreakpointOp, PmEventOp, TrapOp"); - } - for record in cluster_memory(catalog) { - output.push_str(", "); - output.push_str(&record.dialect.op_type); - } - for record in clc_intrinsics(catalog) { - output.push_str(", "); - output.push_str(&record.dialect.op_type); - } - for record in execution_controls(catalog) { - output.push_str(", "); - output.push_str(&record.dialect.op_type); - } - if wgmma_controls(catalog).next().is_some() { - output.push_str( - ", WgmmaCommitGroupSyncAlignedOp, WgmmaFenceSyncAlignedOp, WgmmaWaitGroupSyncAlignedOp", - ); - } - if packed_atomics(catalog).next().is_some() { - if sregs(catalog).next().is_some() || ldmatrix(catalog).next().is_some() { - output.push_str(", "); - } - output.push_str("PackedAtomicAddOp, PackedAtomicFormatAttr"); - } - if redux(catalog).next().is_some() { - if sregs(catalog).next().is_some() - || ldmatrix(catalog).next().is_some() - || packed_atomics(catalog).next().is_some() - { - output.push_str(", "); - } - for (index, record) in redux(catalog).enumerate() { - if index != 0 { - output.push_str(", "); - } - output.push_str(&record.dialect.op_type); - } - } - if vote_intrinsics(catalog).next().is_some() { - if sregs(catalog).next().is_some() - || ldmatrix(catalog).next().is_some() - || packed_atomics(catalog).next().is_some() - || redux(catalog).next().is_some() - { - output.push_str(", "); - } - for (index, record) in vote_intrinsics(catalog).enumerate() { - if index != 0 { - output.push_str(", "); - } - output.push_str(&record.dialect.op_type); - } - } - if active_masks(catalog).next().is_some() { - output.push_str(", "); - for (index, record) in active_masks(catalog).enumerate() { - if index != 0 { - output.push_str(", "); - } - output.push_str(&record.dialect.op_type); - } - } - if warp_matches(catalog).next().is_some() { - output.push_str(", "); - for (index, record) in warp_matches(catalog).enumerate() { - if index != 0 { - output.push_str(", "); - } - output.push_str(&record.dialect.op_type); - } - } - for record in elect_intrinsics(catalog) { - output.push_str(", "); - output.push_str(&record.dialect.op_type); - } - if warp_barriers(catalog).next().is_some() { - output.push_str(", "); - for (index, record) in warp_barriers(catalog).enumerate() { - if index != 0 { - output.push_str(", "); - } - output.push_str(&record.dialect.op_type); - } - } - if warp_shuffles(catalog).next().is_some() { - output.push_str(", "); - for (index, record) in warp_shuffles(catalog).enumerate() { - if index != 0 { - output.push_str(", "); - } - output.push_str(&record.dialect.op_type); - } - } - if packed_alus(catalog).next().is_some() { - if sregs(catalog).next().is_some() - || ldmatrix(catalog).next().is_some() - || packed_atomics(catalog).next().is_some() - || redux(catalog).next().is_some() - || vote_intrinsics(catalog).next().is_some() - || active_masks(catalog).next().is_some() - || warp_matches(catalog).next().is_some() - || warp_barriers(catalog).next().is_some() - || warp_shuffles(catalog).next().is_some() - { - output.push_str(", "); - } - for (index, record) in packed_alus(catalog).enumerate() { - if index != 0 { - output.push_str(", "); - } - output.push_str(&record.dialect.op_type); - } - } - if integer_minmaxes(catalog).next().is_some() { - // packed_alu records always exist, so a separator is always needed. - output.push_str(", "); - for (index, record) in integer_minmaxes(catalog).enumerate() { - if index != 0 { - output.push_str(", "); - } - output.push_str(&record.dialect.op_type); - } - } - if packed_conversions(catalog).next().is_some() { - if sregs(catalog).next().is_some() - || ldmatrix(catalog).next().is_some() - || packed_atomics(catalog).next().is_some() - || redux(catalog).next().is_some() - || vote_intrinsics(catalog).next().is_some() - || active_masks(catalog).next().is_some() - || warp_matches(catalog).next().is_some() - || warp_barriers(catalog).next().is_some() - || warp_shuffles(catalog).next().is_some() - || packed_alus(catalog).next().is_some() - { - output.push_str(", "); - } - for (index, record) in packed_conversions(catalog).enumerate() { - if index != 0 { - output.push_str(", "); - } - output.push_str(&record.dialect.op_type); - } - } - if cp_async_copies(catalog).next().is_some() { - output.push_str(", "); - for (index, record) in cp_async_copies(catalog).enumerate() { - if index != 0 { - output.push_str(", "); - } - output.push_str(&record.dialect.op_type); - } - } - if cp_async_controls(catalog).next().is_some() { - output.push_str(", "); - for (index, record) in cp_async_controls(catalog).enumerate() { - if index != 0 { - output.push_str(", "); - } - output.push_str(&record.dialect.op_type); - } - } - if cp_async_mbarriers(catalog).next().is_some() { - output.push_str(", "); - for (index, record) in cp_async_mbarriers(catalog).enumerate() { - if index != 0 { - output.push_str(", "); - } - output.push_str(&record.dialect.op_type); - } - } - for record in mbarrier_basics(catalog) { - output.push_str(", "); - output.push_str(&record.dialect.op_type); - } - for record in mbarrier_extended(catalog) { - output.push_str(", "); - output.push_str(&record.dialect.op_type); - } - if dot_products(catalog).next().is_some() { - if sregs(catalog).next().is_some() - || ldmatrix(catalog).next().is_some() - || packed_atomics(catalog).next().is_some() - || redux(catalog).next().is_some() - || vote_intrinsics(catalog).next().is_some() - || packed_alus(catalog).next().is_some() - || packed_conversions(catalog).next().is_some() - { - output.push_str(", "); - } - for (index, record) in dot_products(catalog).enumerate() { - if index != 0 { - output.push_str(", "); - } - output.push_str(&record.dialect.op_type); - } - } - if sync_intrinsics(catalog).next().is_some() { - if sregs(catalog).next().is_some() - || ldmatrix(catalog).next().is_some() - || packed_atomics(catalog).next().is_some() - || redux(catalog).next().is_some() - || vote_intrinsics(catalog).next().is_some() - || dot_products(catalog).next().is_some() - || packed_alus(catalog).next().is_some() - || packed_conversions(catalog).next().is_some() - { - output.push_str(", "); - } - for (index, record) in sync_intrinsics(catalog).enumerate() { - if index != 0 { - output.push_str(", "); - } - output.push_str(&record.dialect.op_type); - } - } - for record in stmatrices(catalog) { - output.push_str(", "); - output.push_str(&record.dialect.op_type); - } - for record in tcgen05_non_mma_intrinsics(catalog) { - output.push_str(", "); - output.push_str(&record.dialect.op_type); - } - if tcgen05_mma_intrinsics(catalog).next().is_some() { - output.push_str(", Tcgen05MmaBBufferAttr, Tcgen05MmaBUsageAttr, Tcgen05MmaCollectorAAttr, Tcgen05MmaCtaGroupAttr, Tcgen05MmaFormAttr, Tcgen05MmaKindAttr, Tcgen05MmaOp"); - } - output.push_str( - "};\nuse pliron::basic_block::BasicBlock;\nuse pliron::context::{Context, Ptr};\nuse pliron::input_err;\nuse pliron::location::{Located, Location};\nuse pliron::op::Op;\nuse pliron::operation::Operation;\nuse rustc_public::{CrateDef, mir, ty::FnDef};\n\n", - ); - if debug_controls(catalog).next().is_some() - || execution_control_family(catalog, "register_control") - .next() - .is_some() - { - output.push_str("use dialect_mir::ops::{MirConstantOp, MirUnreachableOp};\n\n"); - } else if wgmma_controls(catalog).next().is_some() { - output.push_str("use dialect_mir::ops::MirConstantOp;\n\n"); - } - if register_mmas(catalog).next().is_some() - || sparse_mmas(catalog).next().is_some() - || tcgen05_intrinsics(catalog).next().is_some() - { - output.push_str( - "use dialect_mir::{attributes::FieldIndexAttr, ops::{MirConstructArrayOp, MirConstructStructOp, MirExtractFieldOp}, types::{MirArrayType, MirStructType}};\nuse pliron::{builtin::types::{FP32Type, FP64Type, IntegerType, Signedness}, r#type::{TypeHandle, Typed}, value::Value};\n\n", - ); - } - if tcgen05_intrinsics(catalog).next().is_some() { - output = output.replace( - "use crate::translator::{rvalue, terminator::helpers};", - "use crate::translator::{rvalue, terminator::helpers, types};", - ); - } - if elect_intrinsics(catalog).next().is_some() { - output.push_str("use dialect_mir::{ops::MirConstructTupleOp, types::MirTupleType};\n\n"); - } - writeln!( - output, - "pub const GENERATED_INTRINSIC_ABI: u32 = {};", - catalog.intrinsic_abi - ) - .unwrap(); - writeln!( - output, - "pub const GENERATED_INTRINSIC_ABI_NAMESPACE: &str = {:?};\n", - format!("__cuda_oxide_intrinsic_abi_v{}", catalog.intrinsic_abi) - ) - .unwrap(); - output.push_str( - "#[derive(Debug, Clone, PartialEq, Eq)]\npub enum RawIntrinsicIdentity {\n NotRawCrate,\n Known(String),\n UnsupportedAbi(String),\n UnknownId(String),\n}\n\n/// Classify the unrewritten raw intrinsic DefPath. rustc's ordinary name\n/// printer may prefer a public re-export path, which is source API rather than ABI.\npub fn classify_raw_intrinsic(fn_def: FnDef) -> RawIntrinsicIdentity {\n let crate_name = fn_def.krate().name.to_string();\n if !matches!(crate_name.as_str(), \"cuda_intrinsics\" | \"cuda-intrinsics\") {\n return RawIntrinsicIdentity::NotRawCrate;\n }\n\n let mut segments = Vec::new();\n let mut current = Some(fn_def.def_id());\n while let Some(def_id) = current {\n let printed = def_id.name();\n let segment = printed.as_str().rsplit(\"::\").next().unwrap_or_default();\n if segment != crate_name {\n segments.push(segment.to_owned());\n }\n current = def_id.parent();\n }\n segments.reverse();\n classify_raw_intrinsic_path(&crate_name, format!(\"{crate_name}::{}\", segments.join(\"::\")))\n}\n\nfn classify_raw_intrinsic_path(crate_name: &str, path: String) -> RawIntrinsicIdentity {\n if !matches!(crate_name, \"cuda_intrinsics\" | \"cuda-intrinsics\") {\n return RawIntrinsicIdentity::NotRawCrate;\n }\n if is_raw_generated_intrinsic_path(&path) {\n return RawIntrinsicIdentity::Known(path);\n }\n let namespace = path.split(\"::\").nth(1).unwrap_or_default();\n if namespace.starts_with(\"__cuda_oxide_intrinsic_abi_v\")\n && namespace != GENERATED_INTRINSIC_ABI_NAMESPACE\n {\n RawIntrinsicIdentity::UnsupportedAbi(path)\n } else {\n RawIntrinsicIdentity::UnknownId(path)\n }\n}\n\npub fn require_supported_raw_intrinsic(\n fn_def: FnDef,\n loc: &Location,\n) -> TranslationResult> {\n match classify_raw_intrinsic(fn_def) {\n RawIntrinsicIdentity::NotRawCrate => Ok(None),\n RawIntrinsicIdentity::Known(path) => Ok(Some(path)),\n RawIntrinsicIdentity::UnsupportedAbi(path) => input_err!(\n loc.clone(),\n TranslationErr::unsupported(format!(\n \"cuda-intrinsics ABI mismatch: `{path}` uses an unsupported intrinsic ABI; this compiler supports ABI v{GENERATED_INTRINSIC_ABI}\"\n ))\n ),\n RawIntrinsicIdentity::UnknownId(path) => input_err!(\n loc.clone(),\n TranslationErr::unsupported(format!(\n \"cuda-intrinsics ABI mismatch: `{path}` is not a known intrinsic ID in ABI v{GENERATED_INTRINSIC_ABI}\"\n ))\n ),\n }\n}\n\n", - ); - output - .push_str("pub fn is_generated_intrinsic_path(name: &str) -> bool {\n matches!(name,\n"); - render_compiler_path_patterns(&mut output, catalog, " "); - output.push_str(" )\n}\n\npub fn is_raw_generated_intrinsic_path(name: &str) -> bool {\n matches!(name,\n"); - let raw_paths: Vec<_> = catalog - .intrinsics - .iter() - .map(|record| record.rust.canonical_path.as_str()) - .collect(); - render_string_patterns(&mut output, &raw_paths, " "); - // Test-only. The dispatch arms rendered below carry each marker as a - // literal, and the op-name-keyed lookup the compiler actually calls lives - // in `cuda-oxide-codegen`; nothing outside this file's own generated tests - // reads this path-keyed table. Gate it rather than leave a 2000-line match - // sitting unreferenced in the compiler. - output.push_str(" )\n}\n\n#[cfg(test)]\npub fn generated_intrinsic_marker(name: &str) -> Option<&'static str> {\n match name {\n"); - for record in &catalog.intrinsics { - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(&mut output, &path_refs); - writeln!(output, " => Some({:?}),", intrinsic_marker(catalog, record)).unwrap(); - } - output.push_str(" _ => None,\n }\n}\n\n"); - output.push_str( - "#[allow(clippy::too_many_arguments)]\npub fn try_dispatch_generated_intrinsic(\n ctx: &mut Context,\n body: &mir::Body,\n name: &str,\n args: &[mir::Operand],\n destination: &mir::Place,\n target: &Option,\n block_ptr: Ptr,\n prev_op: Option>,\n value_map: &mut ValueMap,\n block_map: &[Ptr],\n loc: Location,\n) -> TranslationResult>> {\n match name {\n", - ); - for record in sregs(catalog) { - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(&mut output, &path_refs); - output.push_str(" => {\n"); - writeln!( - output, - " require_arity(name, args.len(), 0, &loc)?;" - ) - .unwrap(); - let helper = if record.scalar_width() == Some(64) { - "emit_generated_nvvm_intrinsic_u64" - } else { - "emit_generated_nvvm_intrinsic" - }; - writeln!(output, " Ok(Some(helpers::{helper}(").unwrap(); - writeln!( - output, - " ctx, {}::get_concrete_op_info(), {:?}, destination, target, block_ptr,", - record.dialect.op_type, - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " prev_op, value_map, block_map, loc,\n )?))\n }\n", - ); - } - for record in active_masks(catalog) { - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(&mut output, &path_refs); - output.push_str(" => {\n"); - output.push_str(" require_arity(name, args.len(), 0, &loc)?;\n"); - writeln!( - output, - " let active_mask = {}::build(ctx);", - record.dialect.op_type - ) - .unwrap(); - output.push_str(" active_mask.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, active_mask, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, active_mask, block_ptr, prev_op);\n let result = active_mask.deref(ctx).get_result(0);\n", - ); - writeln!( - output, - " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, active_mask, value_map, block_map, loc,\n {:?},\n )?))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n"); - } - for record in ldmatrix(catalog) { - let (shape, multiplicity, layout, element, state_space) = ldmatrix_attr_variants(record); - let register_count = record.ldmatrix.as_ref().unwrap().variant.register_count(); - output.push_str(" "); - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - render_inline_patterns(&mut output, &path_refs); - output.push_str(" => {\n"); - output.push_str(" require_arity(name, args.len(), 1, &loc)?;\n"); - output.push_str( - " let (address, last_op) = rvalue::translate_operand(\n ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n )?;\n", - ); - writeln!( - output, - " let load = LdmatrixOp::build(ctx, address, {shape}, {multiplicity}, {layout}, {element}, {state_space});" - ) - .unwrap(); - output.push_str(" load.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, load, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str(" helpers::insert_op(ctx, load, block_ptr, last_op);\n"); - if register_count == 1 { - output.push_str(" let value = load.deref(ctx).get_result(0);\n"); - output.push_str(" let last_op = load;\n"); - } else { - writeln!( - output, - " let (value, last_op) = helpers::bundle_generated_u32_results_as_array(ctx, load, {register_count}, loc.clone());" - ) - .unwrap(); - } - writeln!( - output, - " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, value, target, block_ptr, last_op, value_map, block_map, loc,\n {:?},\n )?))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n"); - } - for record in stmatrices(catalog) { - let (multiplicity, _) = stmatrix_variant(record).expect("stmatrix variant"); - let arity = multiplicity.register_count() + 1; - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(&mut output, &path_refs); - output.push_str(" => {\n"); - writeln!( - output, - " require_arity(name, args.len(), {arity}, &loc)?;" - ) - .unwrap(); - output.push_str( - " let mut last_op = prev_op;\n let mut operands = Vec::with_capacity(args.len());\n for arg in args {\n let (value, translated) = rvalue::translate_operand(\n ctx, body, arg, value_map, block_ptr, last_op, loc.clone(),\n )?;\n last_op = translated;\n operands.push(value);\n }\n", - ); - writeln!( - output, - " let store = Operation::new(ctx, {}::get_concrete_op_info(), vec![], operands, vec![], 0);", - record.dialect.op_type - ) - .unwrap(); - output.push_str(" store.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, store, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, store, block_ptr, last_op);\n if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, store, block_map, loc)))\n } else {\n", - ); - writeln!( - output, - " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", - format!( - "{} call without target block", - stmatrix_compatibility_name(record) - ) - ) - .unwrap(); - output.push_str(" }\n }\n"); - } - for record in register_mmas(catalog) { - let adapter = match record.register_mma.as_ref().unwrap().adapter { - RegisterMmaAdapter::C2U32A2U32B1U32ToD2U32 => { - "GeneratedMmaImportAdapter::C2U32A2U32B1U32ToD2U32" - } - RegisterMmaAdapter::C2U32A4U32B2U32ToD2U32 => { - "GeneratedMmaImportAdapter::C2U32A4U32B2U32ToD2U32" - } - RegisterMmaAdapter::C4F32A2U32B1U32ToD4F32 => { - "GeneratedMmaImportAdapter::C4F32A2U32B1U32ToD4F32" - } - RegisterMmaAdapter::C4F32A4U32B2U32ToD4F32 => { - "GeneratedMmaImportAdapter::C4F32A4U32B2U32ToD4F32" - } - RegisterMmaAdapter::C2F64A1F64B1F64ToD2F64 => { - "GeneratedMmaImportAdapter::C2F64A1F64B1F64ToD2F64" - } - RegisterMmaAdapter::C2I32A1U32B1U32ToD2I32 => { - "GeneratedMmaImportAdapter::C2I32A1U32B1U32ToD2I32" - } - RegisterMmaAdapter::C4I32A4U32B2U32ToD4I32 => { - "GeneratedMmaImportAdapter::C4I32A4U32B2U32ToD4I32" - } - RegisterMmaAdapter::C4I32A2U32B1U32ToD4I32 => { - "GeneratedMmaImportAdapter::C4I32A2U32B1U32ToD4I32" - } - RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 => { - "GeneratedMmaImportAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32" - } - }; - let ( - shape, - operation, - kind, - accumulator, - a_element, - b_element, - a_layout, - b_layout, - overflow, - ) = register_mma_attr_variants(record); - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(&mut output, &path_refs); - output.push_str(" => {\n"); - let arity = record.rust.arguments.len(); - writeln!( - output, - " require_arity(name, args.len(), {arity}, &loc)?;" - ) - .unwrap(); - writeln!( - output, - " let (operands, last_op, result_ty, result_count) = import_generated_mma_operands(ctx, body, args, block_ptr, prev_op, value_map, loc.clone(), {adapter})?;" - ) - .unwrap(); - output.push_str( - " let mma = Operation::new(ctx, RegisterMmaOp::get_concrete_op_info(), vec![result_ty; result_count], operands, vec![], 0);\n mma.deref_mut(ctx).set_loc(loc.clone());\n let mma = RegisterMmaOp::new(mma);\n", - ); - writeln!( - output, - " mma.set_attr_nvvm_register_mma_shape(ctx, {shape});\n mma.set_attr_nvvm_register_mma_operation(ctx, {operation});\n mma.set_attr_nvvm_register_mma_kind(ctx, {kind});\n mma.set_attr_nvvm_register_mma_accumulator(ctx, {accumulator});\n mma.set_attr_nvvm_register_mma_a_element(ctx, {a_element});\n mma.set_attr_nvvm_register_mma_b_element(ctx, {b_element});\n mma.set_attr_nvvm_register_mma_a_layout(ctx, {a_layout});\n mma.set_attr_nvvm_register_mma_b_layout(ctx, {b_layout});\n mma.set_attr_nvvm_register_mma_overflow(ctx, {overflow});" - ) - .unwrap(); - output.push_str(" let mma = mma.get_operation();\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, mma, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, mma, block_ptr, Some(last_op));\n let (result, last_op) = bundle_generated_mma_results(ctx, mma, result_ty, result_count, loc.clone());\n", - ); - writeln!( - output, - " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, last_op, value_map, block_map, loc,\n {:?},\n )?))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n"); - } - for record in sparse_mmas(catalog) { - let ( - shape, - accumulator, - a_element, - b_element, - a_layout, - b_layout, - overflow, - metadata, - selector, - ) = sparse_mma_attr_variants(record); - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(&mut output, &path_refs); - output.push_str(" => {\n"); - output.push_str(" require_arity(name, args.len(), 5, &loc)?;\n"); - writeln!( - output, - " if !matches!(&args[4], mir::Operand::Constant(_)) {{\n return input_err!(\n loc,\n TranslationErr::unsupported(\n {:?}.to_owned()\n )\n );\n }}", - sparse_mma_selector_error(record) - ) - .unwrap(); - writeln!( - output, - " let (mut operands, last_op, result_ty, result_count) = import_generated_mma_operands(ctx, body, args, block_ptr, prev_op, value_map, loc.clone(), {})?;", - sparse_mma_import_adapter(record) - ) - .unwrap(); - output.push_str( - " let (metadata, last_op) = rvalue::translate_operand(\n ctx, body, &args[3], value_map, block_ptr, Some(last_op), loc.clone(),\n )?;\n operands.push(metadata);\n let (selector_value, last_op) = rvalue::translate_operand(\n ctx, body, &args[4], value_map, block_ptr, last_op, loc.clone(),\n )?;\n operands.push(selector_value);\n", - ); - output.push_str( - " let mma = Operation::new(ctx, SparseMmaOp::get_concrete_op_info(), vec![result_ty; result_count], operands, vec![], 0);\n mma.deref_mut(ctx).set_loc(loc.clone());\n let mma = SparseMmaOp::new(mma);\n", - ); - writeln!( - output, - " mma.set_attr_nvvm_sparse_mma_shape(ctx, {shape});\n mma.set_attr_nvvm_sparse_mma_accumulator(ctx, {accumulator});\n mma.set_attr_nvvm_sparse_mma_a_element(ctx, {a_element});\n mma.set_attr_nvvm_sparse_mma_b_element(ctx, {b_element});\n mma.set_attr_nvvm_sparse_mma_a_layout(ctx, {a_layout});\n mma.set_attr_nvvm_sparse_mma_b_layout(ctx, {b_layout});\n mma.set_attr_nvvm_sparse_mma_overflow(ctx, {overflow});\n mma.set_attr_nvvm_sparse_mma_metadata(ctx, {metadata});\n mma.set_attr_nvvm_sparse_mma_selector(ctx, {selector});" - ) - .unwrap(); - output.push_str(" let mma = mma.get_operation();\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, mma, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, mma, block_ptr, last_op);\n let (result, last_op) = bundle_generated_mma_results(ctx, mma, result_ty, result_count, loc.clone());\n", - ); - writeln!( - output, - " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, last_op, value_map, block_map, loc,\n {:?},\n )?))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n"); - } - for record in packed_atomics(catalog) { - let format = match record.packed_atomic.as_ref().unwrap().format { - PackedAtomicFormat::F16x2 => "PackedAtomicFormatAttr::F16x2", - PackedAtomicFormat::Bf16x2 => "PackedAtomicFormatAttr::Bf16x2", - }; - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(&mut output, &path_refs); - output.push_str(" => {\n"); - output.push_str(" require_arity(name, args.len(), 2, &loc)?;\n"); - output.push_str( - " let (address, last_op) = rvalue::translate_operand(\n ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n )?;\n", - ); - output.push_str( - " let (addend, last_op) = rvalue::translate_operand(\n ctx, body, &args[1], value_map, block_ptr, last_op, loc.clone(),\n )?;\n", - ); - writeln!( - output, - " let atom = PackedAtomicAddOp::build(ctx, address, addend, {format});" - ) - .unwrap(); - output.push_str(" atom.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, atom, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str(" helpers::insert_op(ctx, atom, block_ptr, last_op);\n"); - output.push_str(" let value = atom.deref(ctx).get_result(0);\n"); - writeln!( - output, - " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, value, target, block_ptr, atom, value_map, block_map, loc,\n {:?},\n )?))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n"); - } - for record in redux(catalog) { - debug_assert_eq!( - record.redux.as_ref().unwrap().adapter, - ReduxAdapter::MaskValueToSourceMemberMask - ); - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(&mut output, &path_refs); - output.push_str(" => {\n"); - output.push_str(" require_arity(name, args.len(), 2, &loc)?;\n"); - output.push_str( - " let (member_mask, last_op) = rvalue::translate_operand(\n ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n )?;\n", - ); - output.push_str( - " let (value, last_op) = rvalue::translate_operand(\n ctx, body, &args[1], value_map, block_ptr, last_op, loc.clone(),\n )?;\n", - ); - writeln!( - output, - " let reduction = {}::build(ctx, member_mask, value);", - record.dialect.op_type - ) - .unwrap(); - output.push_str(" reduction.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, reduction, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, reduction, block_ptr, last_op);\n let result = reduction.deref(ctx).get_result(0);\n", - ); - writeln!( - output, - " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, reduction, value_map, block_map, loc,\n {:?},\n )?))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n"); - } - for record in vote_intrinsics(catalog) { - debug_assert_eq!( - record.vote.as_ref().unwrap().adapter, - VoteAdapter::DirectMaskPredicate - ); - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(&mut output, &path_refs); - output.push_str(" => {\n"); - output.push_str(" require_arity(name, args.len(), 2, &loc)?;\n"); - output.push_str( - " let (member_mask, last_op) = rvalue::translate_operand(\n ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n )?;\n", - ); - output.push_str( - " let (predicate, last_op) = rvalue::translate_operand(\n ctx, body, &args[1], value_map, block_ptr, last_op, loc.clone(),\n )?;\n", - ); - writeln!( - output, - " let vote = {}::build(ctx, member_mask, predicate);", - record.dialect.op_type - ) - .unwrap(); - output.push_str(" vote.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, vote, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, vote, block_ptr, last_op);\n let result = vote.deref(ctx).get_result(0);\n", - ); - writeln!( - output, - " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, vote, value_map, block_map, loc,\n {:?},\n )?))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n"); - } - for record in warp_matches(catalog) { - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(&mut output, &path_refs); - output.push_str(" => {\n"); - output.push_str(" require_arity(name, args.len(), 2, &loc)?;\n"); - output.push_str( - " let (member_mask, last_op) = rvalue::translate_operand(\n ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n )?;\n", - ); - output.push_str( - " let (value, last_op) = rvalue::translate_operand(\n ctx, body, &args[1], value_map, block_ptr, last_op, loc.clone(),\n )?;\n", - ); - writeln!( - output, - " let warp_match = {}::build(ctx, member_mask, value);", - record.dialect.op_type - ) - .unwrap(); - output.push_str(" warp_match.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, warp_match, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, warp_match, block_ptr, last_op);\n let result = warp_match.deref(ctx).get_result(0);\n", - ); - writeln!( - output, - " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, warp_match, value_map, block_map, loc,\n {:?},\n )?))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n"); - } - for record in elect_intrinsics(catalog) { - render_importer_elect_dispatch(&mut output, catalog, record); - } - for record in warp_barriers(catalog) { - debug_assert_eq!( - record.warp_barrier.as_ref().unwrap().adapter, - WarpBarrierAdapter::DirectMemberMask - ); - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(&mut output, &path_refs); - output.push_str(" => {\n"); - output.push_str(" require_arity(name, args.len(), 1, &loc)?;\n"); - output.push_str( - " let (member_mask, last_op) = rvalue::translate_operand(\n ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n )?;\n", - ); - writeln!( - output, - " let barrier = {}::build(ctx, member_mask);", - record.dialect.op_type - ) - .unwrap(); - output.push_str(" barrier.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, barrier, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, barrier, block_ptr, last_op);\n if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, barrier, block_map, loc)))\n } else {\n", - ); - writeln!( - output, - " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n }\n"); - } - for record in warp_shuffles(catalog) { - debug_assert!(matches!( - record.warp_shuffle.as_ref().unwrap().adapter, - WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp - | WarpShuffleAdapter::MaskValueLaneOrDeltaSplitI64LowHighB32InsertClampReassemble - )); - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(&mut output, &path_refs); - output.push_str(" => {\n"); - output.push_str(" require_arity(name, args.len(), 3, &loc)?;\n"); - output.push_str( - " let (member_mask, last_op) = rvalue::translate_operand(\n ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n )?;\n", - ); - output.push_str( - " let (value, last_op) = rvalue::translate_operand(\n ctx, body, &args[1], value_map, block_ptr, last_op, loc.clone(),\n )?;\n", - ); - output.push_str( - " let (lane_or_delta, last_op) = rvalue::translate_operand(\n ctx, body, &args[2], value_map, block_ptr, last_op, loc.clone(),\n )?;\n", - ); - writeln!( - output, - " let shuffle = {}::build(ctx, member_mask, value, lane_or_delta);", - record.dialect.op_type - ) - .unwrap(); - output.push_str(" shuffle.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, shuffle, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, shuffle, block_ptr, last_op);\n let result = shuffle.deref(ctx).get_result(0);\n", - ); - writeln!( - output, - " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, shuffle, value_map, block_map, loc,\n {:?},\n )?))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n"); - } - for record in dot_products(catalog) { - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(&mut output, &path_refs); - output.push_str(" => {\n"); - output.push_str(" require_arity(name, args.len(), 3, &loc)?;\n"); - output.push_str( - " let (a, last_op) = rvalue::translate_operand(\n ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n )?;\n", - ); - output.push_str( - " let (b, last_op) = rvalue::translate_operand(\n ctx, body, &args[1], value_map, block_ptr, last_op, loc.clone(),\n )?;\n", - ); - output.push_str( - " let (c, last_op) = rvalue::translate_operand(\n ctx, body, &args[2], value_map, block_ptr, last_op, loc.clone(),\n )?;\n", - ); - writeln!( - output, - " let dot = {}::build(ctx, a, b, c);", - record.dialect.op_type - ) - .unwrap(); - output.push_str(" dot.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, dot, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, dot, block_ptr, last_op);\n let result = dot.deref(ctx).get_result(0);\n", - ); - writeln!( - output, - " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, dot, value_map, block_map, loc,\n {:?},\n )?))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n"); - } - for record in packed_alus(catalog) { - render_importer_pure_value_dispatch(&mut output, catalog, record); - } - for record in integer_minmaxes(catalog) { - render_importer_pure_value_dispatch(&mut output, catalog, record); - } - for record in packed_conversions(catalog) { - render_importer_pure_value_dispatch(&mut output, catalog, record); - } - for record in scalar_conversions(catalog) { - render_importer_scalar_conversion_dispatch(&mut output, catalog, record); - } - for record in scalar_arithmetics(catalog) { - render_importer_scalar_arithmetic_dispatch(&mut output, catalog, record); - } - for record in scalar_maths(catalog) { - render_importer_scalar_math_dispatch(&mut output, catalog, record); - } - for record in extended_minmax(catalog) { - render_importer_extended_minmax_dispatch(&mut output, catalog, record); - } - for record in movmatrix(catalog) { - render_importer_pure_value_dispatch(&mut output, catalog, record); - } - for record in prmts(catalog) { - let mode = match record.prmt.as_ref().unwrap().mode { - PrmtMode::Generic => "PrmtModeAttr::Generic", - PrmtMode::F4e => "PrmtModeAttr::F4e", - PrmtMode::B4e => "PrmtModeAttr::B4e", - PrmtMode::Rc8 => "PrmtModeAttr::Rc8", - PrmtMode::Ecl => "PrmtModeAttr::Ecl", - PrmtMode::Ecr => "PrmtModeAttr::Ecr", - PrmtMode::Rc16 => "PrmtModeAttr::Rc16", - }; - let arity = record.rust.arguments.len(); - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(&mut output, &path_refs); - output.push_str(" => {\n"); - writeln!( - output, - " require_arity(name, args.len(), {arity}, &loc)?;" - ) - .unwrap(); - for index in 0..arity { - let previous = if index == 0 { "prev_op" } else { "last_op" }; - writeln!( - output, - " let (arg{index}, last_op) = rvalue::translate_operand(\n ctx, body, &args[{index}], value_map, block_ptr, {previous}, loc.clone(),\n )?;" - ) - .unwrap(); - } - let arguments = (0..arity) - .map(|index| format!("arg{index}")) - .collect::>() - .join(", "); - writeln!( - output, - " let prmt = PrmtOp::build(ctx, vec![{arguments}], {mode});" - ) - .unwrap(); - output.push_str(" prmt.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, prmt, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, prmt, block_ptr, last_op);\n let result = prmt.deref(ctx).get_result(0);\n", - ); - writeln!( - output, - " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, prmt, value_map, block_map, loc,\n {:?},\n )?))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n"); - } - if cluster_barriers(catalog).next().is_some() { - let arrive = cluster_barriers(catalog) - .find(|record| { - record - .cluster_barrier - .as_ref() - .is_some_and(|barrier| barrier.mode == ClusterBarrierMode::ArriveAligned) - }) - .expect("aligned cluster arrival"); - let wait = cluster_barriers(catalog) - .find(|record| { - record - .cluster_barrier - .as_ref() - .is_some_and(|barrier| barrier.mode == ClusterBarrierMode::WaitAligned) - }) - .expect("aligned cluster wait"); - output.push_str( - " \"cuda_device::cluster::cluster_sync\" => {\n require_arity(name, args.len(), 0, &loc)?;\n let arrive = ClusterBarrierOp::build(ctx, ClusterBarrierModeAttr::ArriveAligned);\n arrive.deref_mut(ctx).set_loc(loc.clone());\n", - ); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, arrive, {:?});", - intrinsic_marker(catalog, arrive) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, arrive, block_ptr, prev_op);\n let wait = ClusterBarrierOp::build(ctx, ClusterBarrierModeAttr::WaitAligned);\n wait.deref_mut(ctx).set_loc(loc.clone());\n", - ); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, wait, {:?});", - intrinsic_marker(catalog, wait) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, wait, block_ptr, Some(arrive));\n if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, wait, block_map, loc)))\n } else {\n input_err!(loc, TranslationErr::unsupported(\"cluster_sync call without target block\".to_owned()))\n }\n }\n", - ); - } - for record in cluster_barriers(catalog) { - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(&mut output, &path_refs); - output.push_str(" => {\n require_arity(name, args.len(), 0, &loc)?;\n"); - writeln!( - output, - " let barrier = ClusterBarrierOp::build(ctx, {});", - cluster_barrier_attr(record) - ) - .unwrap(); - output.push_str(" barrier.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, barrier, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, barrier, block_ptr, prev_op);\n if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, barrier, block_map, loc)))\n } else {\n", - ); - writeln!( - output, - " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n }\n"); - } - for record in cluster_memory(catalog) { - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(&mut output, &path_refs); - output.push_str(" => {\n"); - output.push_str( - " require_arity(name, args.len(), 2, &loc)?;\n\ - let (source, last_op) = rvalue::translate_operand(\n\ - ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n\ - )?;\n\ - let (rank, last_op) = rvalue::translate_operand(\n\ - ctx, body, &args[1], value_map, block_ptr, last_op, loc.clone(),\n\ - )?;\n", - ); - writeln!( - output, - " let cluster = {}::build(ctx, source, rank);", - record.dialect.op_type - ) - .unwrap(); - output.push_str(" cluster.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, cluster, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, cluster, block_ptr, last_op);\n\ - let result = cluster.deref(ctx).get_result(0);\n", - ); - writeln!( - output, - " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, cluster, value_map, block_map, loc,\n {:?},\n )?))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n"); - } - for record in debug_controls(catalog) { - let operation = record.debug_control.as_ref().unwrap().operation; - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(&mut output, &path_refs); - output.push_str(" => {\n"); - match operation { - DebugControlOperation::Trap => { - output.push_str(" require_arity(name, args.len(), 0, &loc)?;\n"); - output.push_str(" let debug = TrapOp::build(ctx);\n"); - output.push_str(" debug.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, debug, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, debug, block_ptr, prev_op);\n\ - let unreachable = Operation::new(\n\ - ctx, MirUnreachableOp::get_concrete_op_info(), vec![], vec![], vec![], 0,\n\ - );\n\ - unreachable.deref_mut(ctx).set_loc(loc);\n\ - helpers::insert_op(ctx, unreachable, block_ptr, Some(debug));\n\ - Ok(Some(unreachable))\n", - ); - } - DebugControlOperation::Breakpoint => { - output.push_str(" require_arity(name, args.len(), 0, &loc)?;\n"); - output.push_str(" let debug = BreakpointOp::build(ctx);\n"); - output.push_str(" debug.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, debug, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, debug, block_ptr, prev_op);\n\ - if let Some(target_idx) = target {\n\ - Ok(Some(helpers::emit_goto(ctx, *target_idx, debug, block_map, loc)))\n\ - } else {\n", - ); - writeln!( - output, - " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n"); - } - DebugControlOperation::Pmevent => { - output.push_str( - " require_arity(name, args.len(), 1, &loc)?;\n\ - if !matches!(&args[0], mir::Operand::Constant(_)) {\n\ - return input_err!(\n\ - loc,\n\ - TranslationErr::unsupported(\n\ - \"prof_trigger requires a compile-time constant event ID in 0..=15\".to_owned()\n\ - )\n\ - );\n\ - }\n\ - let (event_id_value, last_op) = rvalue::translate_operand(\n\ - ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n\ - )?;\n\ - let event_id = event_id_value\n\ - .defining_op()\n\ - .and_then(|defining_op| Operation::get_op::(defining_op, ctx))\n\ - .and_then(|constant| constant.get_attr_value(ctx))\n\ - .map(|value| value.value().to_u64())\n\ - .and_then(|value| u32::try_from(value).ok())\n\ - .filter(|value| *value <= 15);\n\ - let Some(event_id) = event_id else {\n\ - return input_err!(\n\ - loc,\n\ - TranslationErr::unsupported(\n\ - \"prof_trigger requires a compile-time constant event ID in 0..=15\".to_owned()\n\ - )\n\ - );\n\ - };\n\ - let debug = PmEventOp::build(ctx, event_id);\n\ - debug.deref_mut(ctx).set_loc(loc.clone());\n", - ); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, debug, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, debug, block_ptr, last_op);\n\ - if let Some(target_idx) = target {\n\ - Ok(Some(helpers::emit_goto(ctx, *target_idx, debug, block_map, loc)))\n\ - } else {\n", - ); - writeln!( - output, - " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n"); - } - } - output.push_str(" }\n"); - } - for record in clc_intrinsics(catalog) { - let query = !matches!( - record.clc.as_ref().unwrap().operation, - ClcOperation::TryCancel | ClcOperation::TryCancelMulticast - ); - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(&mut output, &path_refs); - output.push_str(" => {\n require_arity(name, args.len(), 2, &loc)?;\n"); - output.push_str( - " let (arg0, last_op) = rvalue::translate_operand(\n ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n )?;\n let (arg1, last_op) = rvalue::translate_operand(\n ctx, body, &args[1], value_map, block_ptr, last_op, loc.clone(),\n )?;\n", - ); - let result_types = if query { - "vec![pliron::builtin::types::IntegerType::get(ctx, 32, pliron::builtin::types::Signedness::Unsigned).into()]" - } else { - "vec![]" - }; - writeln!( - output, - " let intrinsic = Operation::new(ctx, {}::get_concrete_op_info(), {result_types}, vec![arg0, arg1], vec![], 0);", - record.dialect.op_type - ) - .unwrap(); - output.push_str(" intrinsic.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, intrinsic, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str(" helpers::insert_op(ctx, intrinsic, block_ptr, last_op);\n"); - if query { - output.push_str(" let result = intrinsic.deref(ctx).get_result(0);\n"); - writeln!( - output, - " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, intrinsic, value_map, block_map, loc,\n {:?},\n )?))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - } else { - output.push_str(" if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, intrinsic, block_map, loc)))\n } else {\n"); - writeln!( - output, - " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n"); - } - output.push_str(" }\n"); - } - for record in wgmma_controls(catalog) { - let control = record.wgmma_control.as_ref().unwrap(); - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(&mut output, &path_refs); - output.push_str(" => {\n"); - match control.mode { - WgmmaControlMode::Fence | WgmmaControlMode::CommitGroup => { - output.push_str(" require_arity(name, args.len(), 0, &loc)?;\n"); - writeln!( - output, - " let control = {}::build(ctx);", - record.dialect.op_type - ) - .unwrap(); - output.push_str(" control.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, control, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, control, block_ptr, prev_op);\n", - ); - } - WgmmaControlMode::WaitGroup => { - output.push_str( - " require_arity(name, args.len(), 1, &loc)?;\n\ - if !matches!(&args[0], mir::Operand::Constant(_)) {\n\ - return input_err!(\n\ - loc,\n\ - TranslationErr::unsupported(\n\ - \"wgmma_wait_group requires a compile-time constant\".to_owned()\n\ - )\n\ - );\n\ - }\n\ - let (max_pending, last_op) = rvalue::translate_operand(\n\ - ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n\ - )?;\n\ - if max_pending\n\ - .defining_op()\n\ - .and_then(|op| Operation::get_op::(op, ctx))\n\ - .is_none()\n\ - {\n\ - return input_err!(\n\ - loc,\n\ - TranslationErr::unsupported(\n\ - \"wgmma_wait_group requires a compile-time constant\".to_owned()\n\ - )\n\ - );\n\ - }\n\ - let control = WgmmaWaitGroupSyncAlignedOp::build(ctx, max_pending);\n\ - control.deref_mut(ctx).set_loc(loc.clone());\n", - ); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, control, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, control, block_ptr, last_op);\n", - ); - } - } - output.push_str( - " if let Some(target_idx) = target {\n\ - Ok(Some(helpers::emit_goto(ctx, *target_idx, control, block_map, loc)))\n\ - } else {\n", - ); - writeln!( - output, - " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n }\n"); - } - for record in cp_async_copies(catalog) { - let copy = record.cp_async_copy.as_ref().unwrap(); - let dynamic = copy.source_size == CpAsyncSourceSize::Runtime; - let arity = if dynamic { 3 } else { 2 }; - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(&mut output, &path_refs); - output.push_str(" => {\n"); - writeln!( - output, - " require_arity(name, args.len(), {arity}, &loc)?;" - ) - .unwrap(); - output.push_str( - " let (shared_dst, last_op) = rvalue::translate_operand(\n ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n )?;\n", - ); - output.push_str( - " let (global_src, last_op) = rvalue::translate_operand(\n ctx, body, &args[1], value_map, block_ptr, last_op, loc.clone(),\n )?;\n", - ); - if dynamic { - output.push_str( - " let (source_size, last_op) = rvalue::translate_operand(\n ctx, body, &args[2], value_map, block_ptr, last_op, loc.clone(),\n )?;\n", - ); - writeln!( - output, - " let copy = {}::build(ctx, shared_dst, global_src, source_size);", - record.dialect.op_type - ) - .unwrap(); - } else { - writeln!( - output, - " let copy = {}::build(ctx, shared_dst, global_src);", - record.dialect.op_type - ) - .unwrap(); - } - output.push_str(" copy.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, copy, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, copy, block_ptr, last_op);\n if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, copy, block_map, loc)))\n } else {\n", - ); - writeln!( - output, - " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n }\n"); - } - for record in cp_async_controls(catalog) { - let control = record.cp_async_control.as_ref().unwrap(); - let has_immediate = control.operation == CpAsyncControlOperation::WaitGroup; - let arity = usize::from(has_immediate); - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(&mut output, &path_refs); - output.push_str(" => {\n"); - writeln!( - output, - " require_arity(name, args.len(), {arity}, &loc)?;" - ) - .unwrap(); - if has_immediate { - output.push_str( - " if !matches!(&args[0], mir::Operand::Constant(_)) {\n return input_err!(\n loc,\n TranslationErr::unsupported(\n \"cp_async_wait_group requires a compile-time constant max_pending value\".to_owned()\n )\n );\n }\n", - ); - output.push_str( - " let (max_pending, last_op) = rvalue::translate_operand(\n ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n )?;\n", - ); - writeln!( - output, - " let control = {}::build(ctx, max_pending);", - record.dialect.op_type - ) - .unwrap(); - } else { - output.push_str(" let last_op = prev_op;\n"); - writeln!( - output, - " let control = {}::build(ctx);", - record.dialect.op_type - ) - .unwrap(); - } - output.push_str(" control.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, control, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, control, block_ptr, last_op);\n if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, control, block_map, loc)))\n } else {\n", - ); - writeln!( - output, - " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n }\n"); - } - for record in cp_async_mbarriers(catalog) { - debug_assert_eq!( - record.cp_async_mbarrier.as_ref().unwrap().adapter, - CpAsyncMbarrierAdapter::PointerToVoid - ); - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(&mut output, &path_refs); - output.push_str(" => {\n"); - output.push_str(" require_arity(name, args.len(), 1, &loc)?;\n"); - output.push_str( - " let (barrier, last_op) = rvalue::translate_operand(\n ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n )?;\n", - ); - writeln!( - output, - " let bridge = {}::build(ctx, barrier);", - record.dialect.op_type - ) - .unwrap(); - output.push_str(" bridge.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, bridge, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, bridge, block_ptr, last_op);\n if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, bridge, block_map, loc)))\n } else {\n", - ); - writeln!( - output, - " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n }\n"); - } - for record in mbarrier_basics(catalog) { - let mbarrier = record.mbarrier_basic.as_ref().unwrap(); - let (argument_names, returns_value) = match mbarrier.operation { - MbarrierBasicOperation::Init => { - debug_assert_eq!( - mbarrier.adapter, - MbarrierBasicAdapter::InitPointerCountToVoid - ); - (&["barrier", "expected_count"][..], false) - } - MbarrierBasicOperation::Arrive => { - debug_assert_eq!(mbarrier.adapter, MbarrierBasicAdapter::ArrivePointerToToken); - (&["barrier"][..], true) - } - MbarrierBasicOperation::TestWait => { - debug_assert_eq!( - mbarrier.adapter, - MbarrierBasicAdapter::TestWaitPointerTokenToPredicate - ); - (&["barrier", "token"][..], true) - } - MbarrierBasicOperation::Inval => { - debug_assert_eq!(mbarrier.adapter, MbarrierBasicAdapter::InvalPointerToVoid); - (&["barrier"][..], false) - } - }; - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(&mut output, &path_refs); - output.push_str(" => {\n"); - writeln!( - output, - " require_arity(name, args.len(), {}, &loc)?;", - argument_names.len() - ) - .unwrap(); - for (index, argument_name) in argument_names.iter().enumerate() { - let previous = if index == 0 { "prev_op" } else { "last_op" }; - writeln!( - output, - " let ({argument_name}, last_op) = rvalue::translate_operand(\n ctx, body, &args[{index}], value_map, block_ptr, {previous}, loc.clone(),\n )?;" - ) - .unwrap(); - } - writeln!( - output, - " let mbarrier = {}::build(ctx, {});", - record.dialect.op_type, - argument_names.join(", ") - ) - .unwrap(); - output.push_str(" mbarrier.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, mbarrier, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str(" helpers::insert_op(ctx, mbarrier, block_ptr, last_op);\n"); - if returns_value { - output.push_str(" let result = mbarrier.deref(ctx).get_result(0);\n"); - writeln!( - output, - " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, mbarrier, value_map, block_map, loc,\n {:?},\n )?))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - } else { - output.push_str( - " if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, mbarrier, block_map, loc)))\n } else {\n", - ); - writeln!( - output, - " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n"); - } - output.push_str(" }\n"); - } - for record in mbarrier_extended(catalog) { - let contract = record.mbarrier_extended.as_ref().unwrap(); - let (arguments, returns_value): (Vec<(usize, &str)>, bool) = match contract.adapter { - MbarrierExtendedAdapter::PointerTxCountBytesToTokenDroppingTxCount => { - (vec![(0, "barrier"), (2, "bytes")], true) - } - MbarrierExtendedAdapter::RawClusterAddressToVoid => (vec![(0, "address")], false), - MbarrierExtendedAdapter::PointerTokenToPredicate => { - (vec![(0, "barrier"), (1, "token")], true) - } - MbarrierExtendedAdapter::PointerParityToPredicate => { - (vec![(0, "barrier"), (1, "parity")], true) - } - MbarrierExtendedAdapter::ZeroOperandsToVoid => (vec![], false), - MbarrierExtendedAdapter::NanosecondsToVoid => (vec![(0, "ns")], false), - }; - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(&mut output, &path_refs); - output.push_str(" => {\n"); - writeln!( - output, - " require_arity(name, args.len(), {}, &loc)?;", - record.rust.arguments.len() - ) - .unwrap(); - if arguments.is_empty() { - output.push_str(" let last_op = prev_op;\n"); - } - for (position, (argument_index, argument_name)) in arguments.iter().enumerate() { - let previous = if position == 0 { "prev_op" } else { "last_op" }; - writeln!( - output, - " let ({argument_name}, last_op) = rvalue::translate_operand(\n ctx, body, &args[{argument_index}], value_map, block_ptr, {previous}, loc.clone(),\n )?;" - ) - .unwrap(); - } - let build_arguments = arguments - .iter() - .map(|(_, name)| *name) - .collect::>() - .join(", "); - writeln!( - output, - " let extended = {}::build(ctx{}{});", - record.dialect.op_type, - if build_arguments.is_empty() { "" } else { ", " }, - build_arguments, - ) - .unwrap(); - output.push_str(" extended.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, extended, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str(" helpers::insert_op(ctx, extended, block_ptr, last_op);\n"); - if returns_value { - output.push_str(" let result = extended.deref(ctx).get_result(0);\n"); - writeln!( - output, - " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, extended, value_map, block_map, loc,\n {:?},\n )?))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - } else { - output.push_str( - " if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, extended, block_map, loc)))\n } else {\n", - ); - writeln!( - output, - " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n"); - } - output.push_str(" }\n"); - } - for record in sync_intrinsics(catalog) { - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(&mut output, &path_refs); - output.push_str(" => {\n"); - output.push_str(" require_arity(name, args.len(), 0, &loc)?;\n"); - writeln!( - output, - " let barrier = Operation::new(ctx, {}::get_concrete_op_info(), vec![], vec![], vec![], 0);", - record.dialect.op_type - ) - .unwrap(); - output.push_str(" barrier.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, barrier, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, barrier, block_ptr, prev_op);\n if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, barrier, block_map, loc)))\n } else {\n", - ); - writeln!( - output, - " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n }\n"); - } - for record in execution_controls(catalog) { - let operation = ExecutionControlOperation::from_catalog_id(&record.id) - .expect("closed execution-control record"); - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(&mut output, &path_refs); - output.push_str(" => {\n"); - writeln!( - output, - " require_arity(name, args.len(), {}, &loc)?;", - operation.operand_count() - ) - .unwrap(); - if operation.requires_immediate_operands() { - output.push_str( - " if !matches!(&args[0], mir::Operand::Constant(_)) {\n return input_err!(\n loc,\n TranslationErr::unsupported(\n \"setmaxnreg requires a compile-time register count in 24..=256 divisible by 8\".to_owned()\n )\n );\n }\n let (register_count_value, last_op) = rvalue::translate_operand(\n ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n )?;\n let register_count = register_count_value\n .defining_op()\n .and_then(|defining_op| Operation::get_op::(defining_op, ctx))\n .and_then(|constant| constant.get_attr_value(ctx))\n .map(|value| value.value().to_u64())\n .and_then(|value| u32::try_from(value).ok())\n .filter(|value| (24..=256).contains(value) && value % 8 == 0);\n let Some(register_count) = register_count else {\n return input_err!(\n loc,\n TranslationErr::unsupported(\n \"setmaxnreg requires a compile-time register count in 24..=256 divisible by 8\".to_owned()\n )\n );\n };\n", - ); - writeln!( - output, - " let control = {}::build(ctx, register_count);", - record.dialect.op_type - ) - .unwrap(); - } else if operation.operand_count() == 0 { - output.push_str(" let last_op = prev_op;\n"); - writeln!( - output, - " let control = Operation::new(ctx, {}::get_concrete_op_info(), vec![], vec![], vec![], 0);", - record.dialect.op_type - ) - .unwrap(); - } else { - output.push_str( - " let mut last_op = prev_op;\n let mut operands = Vec::with_capacity(args.len());\n for arg in args {\n let (value, translated) = rvalue::translate_operand(\n ctx, body, arg, value_map, block_ptr, last_op, loc.clone(),\n )?;\n last_op = translated;\n operands.push(value);\n }\n", - ); - writeln!( - output, - " let control = Operation::new(ctx, {}::get_concrete_op_info(), vec![], operands, vec![], 0);", - record.dialect.op_type - ) - .unwrap(); - } - output.push_str(" control.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, control, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, control, block_ptr, last_op);\n if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, control, block_map, loc)))\n } else {\n", - ); - writeln!( - output, - " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n }\n"); - } - for record in tma_intrinsics(catalog) { - let operation = record.tma.as_ref().unwrap().operation; - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(&mut output, &path_refs); - output.push_str(" => {\n"); - if matches!( - operation, - TmaOperation::WaitGroup | TmaOperation::WaitGroupRead - ) { - output.push_str( - " require_arity(name, args.len(), 1, &loc)?;\n if !matches!(args.first(), Some(mir::Operand::Constant(_))) {\n return input_err!(\n loc,\n TranslationErr::unsupported(\n \"TMA wait-group count must be a compile-time constant\".to_owned()\n )\n );\n }\n", - ); - } - let marker = intrinsic_marker(catalog, record); - match operation { - TmaOperation::G2sTile1d - | TmaOperation::G2sTile2d - | TmaOperation::G2sTile3d - | TmaOperation::G2sTile4d - | TmaOperation::G2sTile5d => { - writeln!( - output, - " Ok(Some(super::tma::emit_tma_g2s(\n ctx, body, args, target, block_ptr, prev_op, value_map, block_map, loc, {}, {marker:?},\n )?))", - operation.dimensions().unwrap() - ) - .unwrap(); - } - TmaOperation::G2sTile2dMulticast => { - writeln!( - output, - " Ok(Some(super::tma::emit_tma_g2s_multicast(\n ctx, body, args, target, block_ptr, prev_op, value_map, block_map, loc, {marker:?},\n )?))" - ) - .unwrap(); - } - TmaOperation::G2sTile2dMulticastCg2 => { - writeln!( - output, - " Ok(Some(super::tma::emit_tma_g2s_multicast_cg2(\n ctx, body, args, target, block_ptr, prev_op, value_map, block_map, loc, {marker:?},\n )?))" - ) - .unwrap(); - } - TmaOperation::S2gTile1d - | TmaOperation::S2gTile2d - | TmaOperation::S2gTile3d - | TmaOperation::S2gTile4d - | TmaOperation::S2gTile5d => { - writeln!( - output, - " Ok(Some(super::tma::emit_tma_s2g(\n ctx, body, args, target, block_ptr, prev_op, value_map, block_map, loc, {}, {marker:?},\n )?))", - operation.dimensions().unwrap() - ) - .unwrap(); - } - TmaOperation::CommitGroup => { - writeln!( - output, - " Ok(Some(super::tma::emit_tma_commit_group(\n ctx, args, target, block_ptr, prev_op, block_map, loc, {marker:?},\n )?))" - ) - .unwrap(); - } - TmaOperation::WaitGroup | TmaOperation::WaitGroupRead => { - writeln!( - output, - " Ok(Some(super::tma::emit_tma_wait_group(\n ctx, body, args, target, block_ptr, prev_op, value_map, block_map, loc, {}, {marker:?},\n )?))", - operation == TmaOperation::WaitGroupRead - ) - .unwrap(); - } - _ => { - let arity = record.dialect.operands.len(); - writeln!( - output, - " require_arity(name, args.len(), {arity}, &loc)?;" - ) - .unwrap(); - if matches!( - record.tma.as_ref().unwrap().adapter, - TmaAdapter::DescriptorOrdinalAndU32 - | TmaAdapter::DescriptorOrdinalAndU64 - | TmaAdapter::DescriptorAndImmediateU32 - ) { - output.push_str( - " if !matches!(args.get(1), Some(mir::Operand::Constant(_))) {\n return input_err!(\n loc,\n TranslationErr::unsupported(\n \"tensor-map replacement selector must be a compile-time constant\".to_owned()\n )\n );\n }\n", - ); - } - output.push_str( - " let mut last_op = prev_op;\n let mut operands = Vec::with_capacity(args.len());\n for arg in args {\n let (value, translated) = rvalue::translate_operand(\n ctx, body, arg, value_map, block_ptr, last_op, loc.clone(),\n )?;\n last_op = translated;\n operands.push(value);\n }\n", - ); - writeln!( - output, - " let intrinsic = Operation::new(ctx, dialect_nvvm::ops::{}::get_concrete_op_info(), vec![], operands, vec![], 0);", - record.dialect.op_type - ) - .unwrap(); - output.push_str(" intrinsic.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, intrinsic, {marker:?});" - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, intrinsic, block_ptr, last_op);\n if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, intrinsic, block_map, loc)))\n } else {\n", - ); - writeln!( - output, - " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n"); - } - } - output.push_str(" }\n"); - } - for record in tcgen05_intrinsics(catalog) { - let tcgen05 = record.tcgen05.as_ref().unwrap(); - if tcgen05.mma.is_some() { - render_importer_tcgen05_mma_dispatch(&mut output, catalog, record); - continue; - } - let operation = tcgen05.operation; - let has_half_split_offset = tcgen05 - .ld - .is_some_and(|ld| ld.shape == Tcgen05LdShape::M16x32bx2) - || tcgen05 - .st - .is_some_and(|st| st.shape == Tcgen05LdShape::M16x32bx2); - let mut path_refs = vec![record.rust.canonical_path.as_str()]; - path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); - output.push_str(" "); - render_inline_patterns(&mut output, &path_refs); - output.push_str(" => {\n"); - if operation == Tcgen05Operation::St { - let count = tcgen05_st_register_count(record); - writeln!( - output, - " require_arity(name, args.len(), {}, &loc)?;", - if has_half_split_offset { 3 } else { 2 } - ) - .unwrap(); - if has_half_split_offset { - output.push_str( - " if !matches!(args.get(1), Some(mir::Operand::Constant(_))) {\n return input_err!(\n loc,\n TranslationErr::unsupported(\n \"tcgen05 16x32bx2 half-split offset must be a compile-time constant\".to_owned()\n )\n );\n }\n", - ); - } - writeln!( - output, - " let (operands, last_op) = import_generated_tcgen05_store_operands(\n ctx, body, args, {count}, {has_half_split_offset}, block_ptr, prev_op, value_map, loc.clone(),\n )?;" - ) - .unwrap(); - writeln!( - output, - " let intrinsic = Operation::new(ctx, {}::get_concrete_op_info(), vec![], operands, vec![], 0);", - record.dialect.op_type - ) - .unwrap(); - output.push_str(" intrinsic.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, intrinsic, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, intrinsic, block_ptr, last_op);\n if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, intrinsic, block_map, loc)))\n } else {\n", - ); - writeln!( - output, - " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n }\n"); - continue; - } - let arity = record.dialect.operands.len(); - writeln!( - output, - " require_arity(name, args.len(), {arity}, &loc)?;" - ) - .unwrap(); - if has_half_split_offset { - output.push_str( - " if !matches!(args.get(1), Some(mir::Operand::Constant(_))) {\n return input_err!(\n loc,\n TranslationErr::unsupported(\n \"tcgen05 16x32bx2 half-split offset must be a compile-time constant\".to_owned()\n )\n );\n }\n", - ); - } - output.push_str( - " let mut last_op = prev_op;\n let mut operands = Vec::with_capacity(args.len());\n for arg in args {\n let (value, translated) = rvalue::translate_operand(\n ctx, body, arg, value_map, block_ptr, last_op, loc.clone(),\n )?;\n last_op = translated;\n operands.push(value);\n }\n", - ); - if has_half_split_offset { - output.push_str( - " if operands.get(1).and_then(|value| value.defining_op()).and_then(|op| Operation::get_op::(op, ctx)).is_none() {\n return input_err!(\n loc,\n TranslationErr::unsupported(\n \"tcgen05 16x32bx2 half-split offset must lower to a constant\".to_owned()\n )\n );\n }\n", - ); - } - let load = match operation { - Tcgen05Operation::Ld16x256bX8Pure => Some((32, "FP32Type::get(ctx).into()")), - Tcgen05Operation::Ld16x256bPure => Some((4, "FP32Type::get(ctx).into()")), - Tcgen05Operation::Ld => Some(( - tcgen05_ld_register_count(record), - "IntegerType::get(ctx, 32, Signedness::Unsigned).into()", - )), - _ => None, - }; - if let Some((count, result_ty)) = load { - writeln!( - output, - " let result_ty: TypeHandle = {result_ty};" - ) - .unwrap(); - writeln!( - output, - " let result_types = (0..{count}).map(|_| result_ty).collect();" - ) - .unwrap(); - writeln!( - output, - " let intrinsic = Operation::new(ctx, {}::get_concrete_op_info(), result_types, operands, vec![], 0);", - record.dialect.op_type - ) - .unwrap(); - output.push_str(" intrinsic.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, intrinsic, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output - .push_str(" helpers::insert_op(ctx, intrinsic, block_ptr, last_op);\n"); - if count == 1 { - output.push_str( - " let result = intrinsic.deref(ctx).get_result(0);\n let value = intrinsic;\n", - ); - } else { - writeln!( - output, - " let results: Vec = (0..{count}).map(|index| intrinsic.deref(ctx).get_result(index)).collect();" - ) - .unwrap(); - writeln!( - output, - " let array_ty = MirArrayType::get(ctx, result_ty, {count});" - ) - .unwrap(); - output.push_str( - " let array = Operation::new(ctx, MirConstructArrayOp::get_concrete_op_info(), vec![array_ty.into()], results, vec![], 0);\n array.deref_mut(ctx).set_loc(loc.clone());\n array.insert_after(ctx, intrinsic);\n let destination_rust_ty = match destination.ty(body.locals()) {\n Ok(ty) => ty,\n Err(error) => {\n return input_err!(\n loc,\n TranslationErr::unsupported(format!(\n \"failed to resolve destination type for intrinsic result: {error:?}\"\n ))\n );\n }\n };\n let destination_ty = types::translate_type(ctx, &destination_rust_ty)?;\n let array_result = array.deref(ctx).get_result(0);\n let (result, value) = if destination_ty == array_result.get_type(ctx) {\n (array_result, array)\n } else {\n let value = Operation::new(ctx, MirConstructStructOp::get_concrete_op_info(), vec![destination_ty], vec![array_result], vec![], 0);\n value.deref_mut(ctx).set_loc(loc.clone());\n value.insert_after(ctx, array);\n (value.deref(ctx).get_result(0), value)\n };\n helpers::set_compiler_result_bundle_marker(ctx, value);\n", - ); - } - writeln!( - output, - " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, value, value_map, block_map, loc,\n {:?},\n )?))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - } else { - writeln!( - output, - " let intrinsic = Operation::new(ctx, {}::get_concrete_op_info(), vec![], operands, vec![], 0);", - record.dialect.op_type - ) - .unwrap(); - output.push_str(" intrinsic.deref_mut(ctx).set_loc(loc.clone());\n"); - writeln!( - output, - " helpers::set_generated_intrinsic_marker(ctx, intrinsic, {:?});", - intrinsic_marker(catalog, record) - ) - .unwrap(); - output.push_str( - " helpers::insert_op(ctx, intrinsic, block_ptr, last_op);\n if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, intrinsic, block_map, loc)))\n } else {\n", - ); - writeln!( - output, - " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", - format!("{} call without target block", record.rust.name) - ) - .unwrap(); - output.push_str(" }\n"); - } - output.push_str(" }\n"); - } - output.push_str(" _ => Ok(None),\n }\n}\n\n"); - output.push_str( - "fn require_arity(\n name: &str,\n actual: usize,\n expected: usize,\n loc: &Location,\n) -> TranslationResult<()> {\n if actual != expected {\n return input_err!(\n loc.clone(),\n TranslationErr::unsupported(format!(\n \"generated intrinsic `{name}` expects {expected} arguments, got {actual}\"\n ))\n );\n }\n Ok(())\n}\n", - ); - if tcgen05_mma_intrinsics(catalog).next().is_some() { - output.push_str( - r#" - -fn generated_tcgen05_mma_selector(operand: Option<&mir::Operand>) -> Option { - use rustc_public::ty::{ConstantKind, TyConstKind}; - let mir::Operand::Constant(constant) = operand? else { return None; }; - let value: u128 = match constant.const_.kind() { - ConstantKind::Allocated(alloc) => alloc.read_uint().ok()?, - ConstantKind::Ty(value) => match value.kind() { - TyConstKind::Value(_, alloc) => alloc.read_uint().ok()?, - _ => value.eval_target_usize().ok()? as u128, - }, - ConstantKind::Unevaluated(_) | ConstantKind::Param(_) => return None, - ConstantKind::ZeroSized => return None, - }; - u32::try_from(value).ok() -} -"#, - ); - } - if tcgen05_intrinsics(catalog) - .any(|record| record.tcgen05.as_ref().unwrap().operation == Tcgen05Operation::St) - { - output.push_str( - r#" - -#[allow(clippy::too_many_arguments)] -fn import_generated_tcgen05_store_operands( - ctx: &mut Context, - body: &mir::Body, - args: &[mir::Operand], - expected_len: usize, - has_half_split_offset: bool, - block_ptr: Ptr, - prev_op: Option>, - value_map: &mut ValueMap, - loc: Location, -) -> TranslationResult<(Vec, Option>)> { - let expected_arity = if has_half_split_offset { 3 } else { 2 }; - if args.len() != expected_arity || !(1..=128).contains(&expected_len) { - return input_err!( - loc, - TranslationErr::unsupported( - "generated tcgen05 store has an invalid address, offset, or register count".to_owned() - ) - ); - } - let u32_ty: TypeHandle = IntegerType::get(ctx, 32, Signedness::Unsigned).into(); - let (address, last_op) = rvalue::translate_operand( - ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(), - )?; - let (half_split_offset, last_op) = if has_half_split_offset { - let (offset, translated) = rvalue::translate_operand( - ctx, body, &args[1], value_map, block_ptr, last_op, loc.clone(), - )?; - if offset - .defining_op() - .and_then(|op| Operation::get_op::(op, ctx)) - .is_none() - { - return input_err!( - loc, - TranslationErr::unsupported( - "tcgen05 16x32bx2 half-split offset must lower to a constant".to_owned() - ) - ); - } - (Some(offset), translated) - } else { - (None, last_op) - }; - let data_index = if has_half_split_offset { 2 } else { 1 }; - let (data, mut last_op) = rvalue::translate_operand( - ctx, body, &args[data_index], value_map, block_ptr, last_op, loc.clone(), - )?; - if expected_len == 1 { - if data.get_type(ctx) != u32_ty { - return input_err!( - loc, - TranslationErr::unsupported( - "generated tcgen05 store data must be a u32 register".to_owned() - ) - ); - } - let mut operands = vec![address]; - operands.extend(half_split_offset); - operands.push(data); - return Ok((operands, last_op)); - } - - let data_ty = data.get_type(ctx); - let direct_array = data_ty - .deref(ctx) - .downcast_ref::() - .is_some_and(|array| { - array.size() == expected_len as u64 && array.element_type() == u32_ty - }); - let array = if direct_array { - data - } else { - let wrapped_array_ty = data_ty - .deref(ctx) - .downcast_ref::() - .filter(|structure| structure.field_count() == 1) - .and_then(|structure| structure.get_field_type(0)); - let valid_wrapper = wrapped_array_ty.is_some_and(|field_ty| { - field_ty - .deref(ctx) - .downcast_ref::() - .is_some_and(|array| { - array.size() == expected_len as u64 && array.element_type() == u32_ty - }) - }); - if !valid_wrapper { - return input_err!( - loc, - TranslationErr::unsupported(format!( - "generated tcgen05 store data must contain {expected_len} u32 registers" - )) - ); - } - let array_ty = wrapped_array_ty.expect("validated tcgen05 store wrapper"); - let extract = Operation::new( - ctx, - MirExtractFieldOp::get_concrete_op_info(), - vec![array_ty], - vec![data], - vec![], - 0, - ); - extract.deref_mut(ctx).set_loc(loc.clone()); - let extract = MirExtractFieldOp::new(extract); - extract.set_attr_index(ctx, FieldIndexAttr(0)); - helpers::insert_op(ctx, extract.get_operation(), block_ptr, last_op); - last_op = Some(extract.get_operation()); - extract.get_operation().deref(ctx).get_result(0) - }; - - let mut operands = Vec::with_capacity(expected_len + 1 + usize::from(has_half_split_offset)); - operands.push(address); - operands.extend(half_split_offset); - for index in 0..expected_len { - let extract = Operation::new( - ctx, - MirExtractFieldOp::get_concrete_op_info(), - vec![u32_ty], - vec![array], - vec![], - 0, - ); - extract.deref_mut(ctx).set_loc(loc.clone()); - let extract = MirExtractFieldOp::new(extract); - extract.set_attr_index(ctx, FieldIndexAttr(index as u32)); - helpers::insert_op(ctx, extract.get_operation(), block_ptr, last_op); - last_op = Some(extract.get_operation()); - operands.push(extract.get_operation().deref(ctx).get_result(0)); - } - Ok((operands, last_op)) -} -"#, - ); - } - if register_mmas(catalog).next().is_some() || sparse_mmas(catalog).next().is_some() { - output.push_str( - r#" - -#[derive(Clone, Copy)] -enum GeneratedMmaImportAdapter { - C2U32A2U32B1U32ToD2U32, - C2U32A4U32B2U32ToD2U32, - C2U32A4U32B4U32ToD2U32, - C4F32A2U32B1U32ToD4F32, - C4F32A4U32B2U32ToD4F32, - C4F32A4U32B4U32ToD4F32, - C2F64A1F64B1F64ToD2F64, - C2I32A1U32B1U32ToD2I32, - C4I32A4U32B2U32ToD4I32, - C4I32A4U32B4U32ToD4I32, - C4I32A2U32B1U32ToD4I32, - C4I32A2U32B2U32ToD4I32, - C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32, -} - -#[allow(clippy::too_many_arguments)] -fn extract_generated_mma_array( - ctx: &mut Context, - array: Value, - expected_element_ty: TypeHandle, - expected_len: usize, - block_ptr: Ptr, - mut last_op: Option>, - loc: Location, -) -> TranslationResult<(Vec, Ptr)> { - let valid = array - .get_type(ctx) - .deref(ctx) - .downcast_ref::() - .is_some_and(|array_ty| { - array_ty.size() == expected_len as u64 - && array_ty.element_type() == expected_element_ty - }); - if !valid { - return input_err!( - loc, - TranslationErr::unsupported(format!( - "generated MMA fragment must be an array of {expected_len} scalar registers" - )) - ); - } - let mut registers = Vec::with_capacity(expected_len); - for index in 0..expected_len { - let extract = Operation::new( - ctx, - MirExtractFieldOp::get_concrete_op_info(), - vec![expected_element_ty], - vec![array], - vec![], - 0, - ); - extract.deref_mut(ctx).set_loc(loc.clone()); - let extract = MirExtractFieldOp::new(extract); - extract.set_attr_index(ctx, FieldIndexAttr(index as u32)); - helpers::insert_op(ctx, extract.get_operation(), block_ptr, last_op); - last_op = Some(extract.get_operation()); - registers.push(extract.get_operation().deref(ctx).get_result(0)); - } - Ok((registers, last_op.expect("non-empty generated MMA fragment"))) -} - -#[allow(clippy::too_many_arguments)] -fn import_generated_mma_operands( - ctx: &mut Context, - body: &mir::Body, - args: &[mir::Operand], - block_ptr: Ptr, - prev_op: Option>, - value_map: &mut ValueMap, - loc: Location, - adapter: GeneratedMmaImportAdapter, -) -> TranslationResult<(Vec, Ptr, TypeHandle, usize)> { - let f32_ty: TypeHandle = FP32Type::get(ctx).into(); - let f64_ty: TypeHandle = FP64Type::get(ctx).into(); - let i32_ty: TypeHandle = IntegerType::get(ctx, 32, Signedness::Signed).into(); - let u32_ty: TypeHandle = IntegerType::get(ctx, 32, Signedness::Unsigned).into(); - let (c_ty, c_count, a_ty, a_count, a_array, b_ty, b_count, b_array, result_ty, result_count) = - match adapter { - GeneratedMmaImportAdapter::C2U32A2U32B1U32ToD2U32 => - (u32_ty, 2, u32_ty, 2, true, u32_ty, 1, false, u32_ty, 2), - GeneratedMmaImportAdapter::C2U32A4U32B2U32ToD2U32 => - (u32_ty, 2, u32_ty, 4, true, u32_ty, 2, true, u32_ty, 2), - GeneratedMmaImportAdapter::C2U32A4U32B4U32ToD2U32 => - (u32_ty, 2, u32_ty, 4, true, u32_ty, 4, true, u32_ty, 2), - GeneratedMmaImportAdapter::C4F32A2U32B1U32ToD4F32 => - (f32_ty, 4, u32_ty, 2, true, u32_ty, 1, false, f32_ty, 4), - GeneratedMmaImportAdapter::C4F32A4U32B2U32ToD4F32 => - (f32_ty, 4, u32_ty, 4, true, u32_ty, 2, true, f32_ty, 4), - GeneratedMmaImportAdapter::C4F32A4U32B4U32ToD4F32 => - (f32_ty, 4, u32_ty, 4, true, u32_ty, 4, true, f32_ty, 4), - GeneratedMmaImportAdapter::C2F64A1F64B1F64ToD2F64 => - (f64_ty, 2, f64_ty, 1, false, f64_ty, 1, false, f64_ty, 2), - GeneratedMmaImportAdapter::C2I32A1U32B1U32ToD2I32 => - (i32_ty, 2, u32_ty, 1, false, u32_ty, 1, false, i32_ty, 2), - GeneratedMmaImportAdapter::C4I32A4U32B2U32ToD4I32 => - (i32_ty, 4, u32_ty, 4, true, u32_ty, 2, true, i32_ty, 4), - GeneratedMmaImportAdapter::C4I32A4U32B4U32ToD4I32 => - (i32_ty, 4, u32_ty, 4, true, u32_ty, 4, true, i32_ty, 4), - GeneratedMmaImportAdapter::C4I32A2U32B1U32ToD4I32 => - (i32_ty, 4, u32_ty, 2, true, u32_ty, 1, false, i32_ty, 4), - GeneratedMmaImportAdapter::C4I32A2U32B2U32ToD4I32 => - (i32_ty, 4, u32_ty, 2, true, u32_ty, 2, true, i32_ty, 4), - GeneratedMmaImportAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 => - (f32_ty, 4, u32_ty, 4, true, u32_ty, 2, true, f32_ty, 4), - }; - let (c_array, last_op) = rvalue::translate_operand( - ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(), - )?; - let (mut operands, last_op) = extract_generated_mma_array( - ctx, c_array, c_ty, c_count, block_ptr, last_op, loc.clone(), - )?; - let (a_value, last_after_a) = rvalue::translate_operand( - ctx, body, &args[1], value_map, block_ptr, Some(last_op), loc.clone(), - )?; - let (a_registers, last_op) = if a_array { - extract_generated_mma_array( - ctx, a_value, a_ty, a_count, block_ptr, last_after_a, loc.clone(), - )? - } else { - (vec![a_value], last_after_a.expect("generated MMA A translation keeps predecessor")) - }; - operands.extend(a_registers); - let (b_value, last_after_b) = rvalue::translate_operand( - ctx, body, &args[2], value_map, block_ptr, Some(last_op), loc.clone(), - )?; - let (b_registers, last_op) = if b_array { - extract_generated_mma_array( - ctx, b_value, b_ty, b_count, block_ptr, last_after_b, loc.clone(), - )? - } else { - (vec![b_value], last_after_b.expect("generated MMA B translation keeps predecessor")) - }; - operands.extend(b_registers); - let mut last_op = last_op; - if matches!( - adapter, - GeneratedMmaImportAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 - ) { - for arg in &args[3..] { - let (value, translated) = rvalue::translate_operand( - ctx, - body, - arg, - value_map, - block_ptr, - Some(last_op), - loc.clone(), - )?; - last_op = translated.expect("generated block-scale MMA operand translation keeps predecessor"); - operands.push(value); - } - } - Ok((operands, last_op, result_ty, result_count)) -} - -fn bundle_generated_mma_results( - ctx: &mut Context, - mma: Ptr, - result_ty: TypeHandle, - result_count: usize, - loc: Location, -) -> (Value, Ptr) { - let results = (0..result_count) - .map(|index| mma.deref(ctx).get_result(index)) - .collect(); - let array_ty = MirArrayType::get(ctx, result_ty, result_count as u64); - let array = Operation::new( - ctx, - MirConstructArrayOp::get_concrete_op_info(), - vec![array_ty.into()], - results, - vec![], - 0, - ); - array.deref_mut(ctx).set_loc(loc); - array.insert_after(ctx, mma); - (array.deref(ctx).get_result(0), array) -} -"#, - ); - } - output.push_str("\n#[cfg(test)]\nmod tests {\n use super::*;\n\n"); - for record in &catalog.intrinsics { - writeln!( - output, - " #[test]\n fn {}_uses_only_canonical_or_compatibility_defpaths() {{", - record.id - ) - .unwrap(); - writeln!( - output, - " assert!(is_generated_intrinsic_path({:?}));", - record.rust.canonical_path - ) - .unwrap(); - writeln!( - output, - " assert!(is_raw_generated_intrinsic_path({:?}));", - record.rust.canonical_path - ) - .unwrap(); - writeln!( - output, - " assert_eq!(generated_intrinsic_marker({:?}), Some({:?}));", - record.rust.canonical_path, - intrinsic_marker(catalog, record) - ) - .unwrap(); - writeln!( - output, - " assert!(matches!(classify_raw_intrinsic_path(\"cuda_intrinsics\", {:?}.into()), RawIntrinsicIdentity::Known(_)));", - record.rust.canonical_path - ) - .unwrap(); - writeln!( - output, - " assert!(!is_generated_intrinsic_path({:?}));", - record.rust.public_path - ) - .unwrap(); - for compatibility_path in &record.rust.compatibility_paths { - writeln!( - output, - " assert!(is_generated_intrinsic_path({compatibility_path:?}));\n assert!(!is_raw_generated_intrinsic_path({compatibility_path:?}));\n assert_eq!(generated_intrinsic_marker({compatibility_path:?}), Some({:?}));", - intrinsic_marker(catalog, record) - ) - .unwrap(); - } - writeln!( - output, - " assert!(!is_generated_intrinsic_path(\"cuda_intrinsics::__cuda_oxide_intrinsic_abi_v{}::{}\"));", - catalog.intrinsic_abi + 1, - record.rust.abi_id - ) - .unwrap(); - output.push_str(" }\n"); - } - output.push_str( - "\n #[test]\n fn raw_intrinsic_identity_classification_fails_closed() {\n assert_eq!(\n classify_raw_intrinsic_path(\"serde\", \"serde::helper\".into()),\n RawIntrinsicIdentity::NotRawCrate\n );\n assert!(matches!(\n classify_raw_intrinsic_path(\"cuda_intrinsics\", \"cuda_intrinsics::__cuda_oxide_intrinsic_abi_v2::i0001\".into()),\n RawIntrinsicIdentity::UnsupportedAbi(_)\n ));\n assert!(matches!(\n classify_raw_intrinsic_path(\"cuda_intrinsics\", \"cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i9999\".into()),\n RawIntrinsicIdentity::UnknownId(_)\n ));\n assert!(matches!(\n classify_raw_intrinsic_path(\"cuda_intrinsics\", \"cuda_intrinsics::helper\".into()),\n RawIntrinsicIdentity::UnknownId(_)\n ));\n }\n", - ); - output.push_str("}\n"); - output -} - -fn special_register_ptx_type(record: &CatalogIntrinsic) -> &'static str { - match record - .special_register - .as_ref() - .expect("special-register record") - .ptx_type - { - SpecialRegisterPtxType::B32 => "b32", - SpecialRegisterPtxType::U32 => "u32", - SpecialRegisterPtxType::U64 => "u64", - } -} - -fn special_register_output_constraint(record: &CatalogIntrinsic) -> &'static str { - match record - .special_register - .as_ref() - .expect("special-register record") - .output_constraint - { - SpecialRegisterOutputConstraint::Register32 => "=r", - SpecialRegisterOutputConstraint::Register64 => "=l", - } -} - -fn special_register_inline_template(record: &CatalogIntrinsic) -> String { - let register = match record.expected_ptx.operands.as_slice() { - [_, crate::ptx::OperandPattern::Exact { value }] => value, - _ => unreachable!("closed special-register PTX operands"), - }; - format!("mov.{} $0, {register};", special_register_ptx_type(record)) -} - -fn special_register_asm_kind(record: &CatalogIntrinsic) -> &'static str { - match record - .special_register - .as_ref() - .expect("special-register record") - .observation - { - SpecialRegisterObservation::StablePure => "AsmKind::Pure", - SpecialRegisterObservation::VolatileObservation => "AsmKind::SideEffect", - } -} - -fn special_register_backend_mechanism( - record: &CatalogIntrinsic, - backend: IntrinsicBackend, -) -> BackendLoweringMechanism { - record - .backend_lowerings - .iter() - .find(|route| route.backend == backend) - .expect("special-register backend route") - .mechanism -} - -fn tcgen05_inline_asm(record: &CatalogIntrinsic) -> (String, String, Option) { - let operation = record.tcgen05.as_ref().expect("tcgen05 contract").operation; - match operation { - Tcgen05Operation::Alloc | Tcgen05Operation::AllocCg2 => { - let group = if operation == Tcgen05Operation::AllocCg2 { 2 } else { 1 }; - ( - format!("{{ .reg .u64 %shared64; .reg .u32 %shared32; cvta.to.shared.u64 %shared64, $0; cvt.u32.u64 %shared32, %shared64; tcgen05.alloc.cta_group::{group}.sync.aligned.shared::cta.b32 [%shared32], $1; }}"), - "l,r,~{memory}".into(), - None, - ) - } - Tcgen05Operation::Dealloc | Tcgen05Operation::DeallocCg2 => { - let group = if operation == Tcgen05Operation::DeallocCg2 { 2 } else { 1 }; - ( - format!("tcgen05.dealloc.cta_group::{group}.sync.aligned.b32 $0, $1;"), - "r,r,~{memory}".into(), - None, - ) - } - Tcgen05Operation::RelinquishAllocPermit - | Tcgen05Operation::RelinquishAllocPermitCg2 => { - let group = if operation == Tcgen05Operation::RelinquishAllocPermitCg2 { - 2 - } else { - 1 - }; - ( - format!("tcgen05.relinquish_alloc_permit.cta_group::{group}.sync.aligned;"), - "~{memory}".into(), - None, - ) - } - Tcgen05Operation::FenceBeforeThreadSync => ( - "tcgen05.fence::before_thread_sync;".into(), - "~{memory}".into(), - None, - ), - Tcgen05Operation::FenceAfterThreadSync => ( - "tcgen05.fence::after_thread_sync;".into(), - "~{memory}".into(), - None, - ), - Tcgen05Operation::Commit - | Tcgen05Operation::CommitSharedCluster - | Tcgen05Operation::CommitCg2 - | Tcgen05Operation::CommitSharedClusterCg2 => { - let group = if matches!( - operation, - Tcgen05Operation::CommitCg2 | Tcgen05Operation::CommitSharedClusterCg2 - ) { - 2 - } else { - 1 - }; - let shared = if matches!( - operation, - Tcgen05Operation::CommitSharedCluster - | Tcgen05Operation::CommitSharedClusterCg2 - ) { - ".shared::cluster" - } else { - "" - }; - ( - format!("tcgen05.commit.cta_group::{group}.mbarrier::arrive::one{shared}.b64 [$0];"), - "r,~{memory}".into(), - None, - ) - } - Tcgen05Operation::MmaWsF16 | Tcgen05Operation::MmaWsBf16 => ( - "{ .reg .pred %enable_pred; setp.ne.s32 %enable_pred, $5, 0; tcgen05.mma.ws.cta_group::1.kind::f16 [$0], [$1], $3, $4, %enable_pred; }".into(), - "r,r,l,l,r,r,~{memory}".into(), - None, - ), - Tcgen05Operation::MmaWsTf32 => ( - "{ .reg .pred %enable_pred; setp.ne.s32 %enable_pred, $5, 0; tcgen05.mma.ws.cta_group::1.kind::tf32 [$0], [$1], $3, $4, %enable_pred; }".into(), - "r,r,l,l,r,r,~{memory}".into(), - None, - ), - Tcgen05Operation::MmaF16 | Tcgen05Operation::MmaF16Cg2 => { - let group = if operation == Tcgen05Operation::MmaF16Cg2 { 2 } else { 1 }; - let zeros = std::iter::repeat_n("%z", if group == 2 { 8 } else { 4 }) - .collect::>() - .join(", "); - ( - format!("{{ .reg .pred %enable_pred; setp.ne.s32 %enable_pred, $4, 0; .reg .u32 %z; mov.u32 %z, 0; tcgen05.mma.cta_group::{group}.kind::f16 [$0], $1, $2, $3, {{{zeros}}}, %enable_pred; }}"), - "r,l,l,r,r,~{memory}".into(), - None, - ) - } - Tcgen05Operation::CpSmemToTmem | Tcgen05Operation::CpSmemToTmemCg2 => { - ( - format!( - "{}.{} [$0], $1;", - record.expected_ptx.mnemonic, - record.expected_ptx.modifiers.join(".") - ), - "r,l,~{memory}".into(), - None, - ) - } - Tcgen05Operation::Ld16x256bX8Pure | Tcgen05Operation::Ld16x256bPure => { - let count = if operation == Tcgen05Operation::Ld16x256bX8Pure { - 32 - } else { - 4 - }; - let registers = (0..count) - .map(|index| format!("${index}")) - .collect::>() - .join(","); - let constraints = std::iter::repeat_n("=f", count) - .chain(std::iter::once("r")) - .collect::>() - .join(","); - ( - format!( - "tcgen05.ld.sync.aligned.16x256b.{}.b32 {{{registers}}}, [${count}];", - if count == 32 { "x8" } else { "x1" } - ), - constraints, - Some(count), - ) - } - Tcgen05Operation::Ld => { - let count = tcgen05_ld_register_count(record); - let has_half_split_offset = record - .tcgen05 - .as_ref() - .and_then(|tcgen05| tcgen05.ld) - .is_some_and(|ld| ld.shape == Tcgen05LdShape::M16x32bx2); - let registers = (0..count) - .map(|index| format!("${index}")) - .collect::>() - .join(","); - let constraints = std::iter::repeat_n("=r", count) - .chain(if has_half_split_offset { - vec!["r", "n", "~{memory}"] - } else { - vec!["r", "~{memory}"] - }) - .collect::>() - .join(","); - ( - format!( - "{}.{} {{{registers}}}, [${count}]{};", - record.expected_ptx.mnemonic, - record.expected_ptx.modifiers.join("."), - if has_half_split_offset { - format!(", ${}", count + 1) - } else { - String::new() - } - ), - constraints, - Some(count), - ) - } - Tcgen05Operation::St => { - let count = tcgen05_st_register_count(record); - let has_half_split_offset = record - .tcgen05 - .as_ref() - .and_then(|tcgen05| tcgen05.st) - .is_some_and(|st| st.shape == Tcgen05LdShape::M16x32bx2); - let first_data = if has_half_split_offset { 2 } else { 1 }; - let registers = (first_data..first_data + count) - .map(|index| format!("${index}")) - .collect::>() - .join(","); - let constraints = std::iter::once("r") - .chain(has_half_split_offset.then_some("n")) - .chain(std::iter::repeat_n("r", count)) - .chain(std::iter::once("~{memory}")) - .collect::>() - .join(","); - ( - format!( - "{}.{} [$0], {}{{{registers}}};", - record.expected_ptx.mnemonic, - record.expected_ptx.modifiers.join("."), - if has_half_split_offset { - "$1, " - } else { - "" - } - ), - constraints, - None, - ) - } - Tcgen05Operation::LoadWait => ( - "tcgen05.wait::ld.sync.aligned;".into(), - "~{memory}".into(), - None, - ), - Tcgen05Operation::StoreWait => ( - "tcgen05.wait::st.sync.aligned;".into(), - "~{memory}".into(), - None, - ), - Tcgen05Operation::CommitMulticast | Tcgen05Operation::CommitMulticastCg2 => { - let group = if operation == Tcgen05Operation::CommitMulticastCg2 { - 2 - } else { - 1 - }; - ( - format!("tcgen05.commit.cta_group::{group}.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 [$0], $1;"), - "r,h,~{memory}".into(), - None, - ) - } - Tcgen05Operation::ShiftDown | Tcgen05Operation::ShiftDownCg2 => { - let group = if operation == Tcgen05Operation::ShiftDownCg2 { - 2 - } else { - 1 - }; - ( - format!("tcgen05.shift.cta_group::{group}.down [$0];"), - "r,~{memory}".into(), - None, - ) - } - Tcgen05Operation::Mma => unreachable!("generic MMA uses attribute-driven lowering"), - } -} - -fn render_lowering(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - output.push_str( - "//! Generated conversion interfaces for admitted CUDA intrinsic families.\n\nuse crate::conversion_interface::MirToLlvmConversion;\nuse crate::convert::intrinsics::{atomic::convert_packed_atom_add, basic::convert_sreg_read_inline, common::{call_intrinsic, create_i32_const, inline_asm_convergent}, cp_async::{convert_generated_cp_async_control, convert_generated_cp_async_copy, convert_generated_cp_async_mbarrier}, dotprod::convert_generated_dot_product, ldmatrix::convert_generated_ldmatrix, mbarrier::{convert_arrive, convert_init, convert_inval, convert_test_wait}, packed::{convert_generated_packed_alu, convert_generated_packed_f32x2, convert_generated_packed_unary}, prmt::convert_generated_prmt, warp::{convert_active_mask, convert_bar_warp_sync, convert_match_all, convert_match_any, convert_redux, convert_shuffle_f32, convert_shuffle_i32, convert_shuffle_i64, convert_vote}, wmma::{convert_generated_register_mma, convert_generated_sparse_mma, GeneratedMmaResultType}};\nuse crate::{context, IntrinsicBackend};\nuse dialect_nvvm::ops::{", - ); - if debug_controls(catalog).next().is_some() { - output = output.replace( - "inline_asm_convergent}", - "inline_asm_convergent, inline_asm_sideeffect}", - ); - } - if integer_minmaxes(catalog).next().is_some() { - output = output.replace( - "ldmatrix::convert_generated_ldmatrix, ", - "integer_minmax::convert_generated_integer_minmax, ldmatrix::convert_generated_ldmatrix, ", - ); - } - if stmatrices(catalog).next().is_some() { - output = output.replace( - "common::{call_intrinsic,", - "common::{call_intrinsic, cast_to_shared_addrspace,", - ); - } - if clc_intrinsics(catalog).next().is_some() { - output = output.replace( - "cp_async::{", - "clc::{convert_generated_clc_query, convert_generated_clc_try_cancel}, cp_async::{", - ); - } - if tma_intrinsics(catalog).next().is_some() { - output = output.replace( - "prmt::convert_generated_prmt, warp::{", - "prmt::convert_generated_prmt, tma::{convert_control, convert_g2s, convert_g2s_multicast_cg2, convert_prefetch_tensormap, convert_prefetch_tile, convert_reduce_s2g, convert_s2g, convert_tensormap_fence, convert_tensormap_replace, PrefetchTileConfig, ReduceConfig}, warp::{", - ); - } - if execution_controls(catalog).next().is_some() { - output = output.replace( - "dotprod::convert_generated_dot_product, ", - "dotprod::convert_generated_dot_product, execution_control::{convert_counted_barrier, convert_grid_dependency, convert_setmaxnreg}, ", - ); - } - if scalar_conversions(catalog).next().is_some() { - output = output.replace( - "prmt::convert_generated_prmt, ", - "prmt::convert_generated_prmt, scalar_conversion::convert_generated_scalar_conversion, ", - ); - } - if scalar_arithmetics(catalog).next().is_some() { - output = output.replace( - "prmt::convert_generated_prmt, ", - "prmt::convert_generated_prmt, scalar_arithmetic::convert_generated_scalar_arithmetic, ", - ); - } - if scalar_maths(catalog).next().is_some() { - output = output.replace( - "prmt::convert_generated_prmt, ", - "prmt::convert_generated_prmt, scalar_math::convert_generated_scalar_math, ", - ); - } - if extended_minmax(catalog).next().is_some() { - output = output.replace( - "dotprod::convert_generated_dot_product, ", - "dotprod::convert_generated_dot_product, extended_minmax::{MinMaxCarrier, convert_generated_extended_minmax}, ", - ); - } - if elect_intrinsics(catalog).next().is_some() { - output = output.replace( - "warp::{convert_active_mask,", - "warp::{convert_active_mask, convert_elect_sync_inline, convert_elect_sync_typed,", - ); - } - for (index, record) in sregs(catalog).enumerate() { - if index != 0 { - output.push_str(", "); - } - output.push_str(&record.dialect.op_type); - } - if ldmatrix(catalog).next().is_some() { - if sregs(catalog).next().is_some() { - output.push_str(", "); - } - output.push_str("LdmatrixElementAttr, LdmatrixLayoutAttr, LdmatrixMultiplicityAttr, LdmatrixOp, LdmatrixShapeAttr, LdmatrixStateSpaceAttr"); - for record in ldmatrix(catalog) { - if let Some((op_type, _)) = ldmatrix_compat_op(record) { - output.push_str(", "); - output.push_str(op_type); - } - } - } - if let Some(record) = movmatrix(catalog).next() { - output.push_str(", "); - output.push_str(&record.dialect.op_type); - } - if register_mmas(catalog).next().is_some() { - for record in - register_mmas(catalog).filter(|record| register_mma_compat_op_type(record).is_some()) - { - output.push_str(", "); - output.push_str(register_mma_compat_op_type(record).unwrap()); - } - output.push_str(", RegisterMmaAccumulatorAttr, RegisterMmaElementAttr, RegisterMmaKindAttr, RegisterMmaLayoutAttr, RegisterMmaOp, RegisterMmaOperationAttr, RegisterMmaOverflowAttr, RegisterMmaShapeAttr"); - } - if sparse_mmas(catalog).next().is_some() { - output.push_str(", SparseMmaAccumulatorAttr, SparseMmaElementAttr, SparseMmaLayoutAttr, SparseMmaMetadataAttr, SparseMmaOp, SparseMmaOverflowAttr, SparseMmaSelectorAttr, SparseMmaShapeAttr"); - } - if prmts(catalog).next().is_some() { - output.push_str(", PrmtModeAttr, PrmtOp"); - } - if scalar_conversions(catalog).next().is_some() { - output.push_str( - ", ScalarConversionOp, ScalarConversionRoundingAttr, ScalarConversionSaturationAttr", - ); - } - if scalar_arithmetics(catalog).next().is_some() { - output.push_str( - ", ScalarArithmeticFormatAttr, ScalarArithmeticOp, ScalarArithmeticOperationAttr, ScalarArithmeticRoundingAttr, ScalarArithmeticSaturationAttr, ScalarArithmeticSubnormalAttr", - ); - } - if scalar_maths(catalog).next().is_some() { - output.push_str( - ", ScalarMathFormatAttr, ScalarMathOp, ScalarMathOperationAttr, ScalarMathPrecisionAttr, ScalarMathSubnormalAttr", - ); - } - if extended_minmax(catalog).next().is_some() { - output.push_str( - ", ExtendedMinMaxFormatAttr, ExtendedMinMaxNanAttr, ExtendedMinMaxOp, ExtendedMinMaxOperationAttr, ExtendedMinMaxSubnormalAttr, ExtendedMinMaxXorSignAbsAttr", - ); - } - if cluster_barriers(catalog).next().is_some() { - output.push_str(", ClusterBarrierModeAttr, ClusterBarrierOp, ClusterSyncOp"); - } - if debug_controls(catalog).next().is_some() { - output.push_str(", BreakpointOp, PmEventOp, TrapOp"); - } - for record in clc_intrinsics(catalog) { - output.push_str(", "); - output.push_str(&record.dialect.op_type); - } - if wgmma_controls(catalog).next().is_some() { - output.push_str( - ", WgmmaCommitGroupSyncAlignedOp, WgmmaFenceSyncAlignedOp, WgmmaWaitGroupSyncAlignedOp", - ); - } - for record in tcgen05_non_mma_intrinsics(catalog) { - output.push_str(", "); - output.push_str(&record.dialect.op_type); - } - if tcgen05_mma_intrinsics(catalog).next().is_some() { - output.push_str(", Tcgen05MmaBBufferAttr, Tcgen05MmaBUsageAttr, Tcgen05MmaCollectorAAttr, Tcgen05MmaCtaGroupAttr, Tcgen05MmaFormAttr, Tcgen05MmaKindAttr, Tcgen05MmaOp"); - } - if packed_atomics(catalog).next().is_some() { - if sregs(catalog).next().is_some() || ldmatrix(catalog).next().is_some() { - output.push_str(", "); - } - output.push_str("NvvmAtomAddBf16x2Op, NvvmAtomAddF16x2Op, PackedAtomicAddOp, PackedAtomicAtomicityAttr, PackedAtomicFormatAttr, PackedAtomicOrderingAttr, PackedAtomicRoundingAttr, PackedAtomicScopeAttr, PackedAtomicStateSpaceAttr, PackedAtomicSubnormalAttr"); - } - if redux(catalog).next().is_some() { - if sregs(catalog).next().is_some() - || ldmatrix(catalog).next().is_some() - || packed_atomics(catalog).next().is_some() - { - output.push_str(", "); - } - for (index, record) in redux(catalog).enumerate() { - if index != 0 { - output.push_str(", "); - } - output.push_str(&record.dialect.op_type); - } - } - if vote_intrinsics(catalog).next().is_some() { - if sregs(catalog).next().is_some() - || ldmatrix(catalog).next().is_some() - || packed_atomics(catalog).next().is_some() - || redux(catalog).next().is_some() - { - output.push_str(", "); - } - for (index, record) in vote_intrinsics(catalog).enumerate() { - if index != 0 { - output.push_str(", "); - } - output.push_str(&record.dialect.op_type); - } - } - if active_masks(catalog).next().is_some() { - output.push_str(", "); - for (index, record) in active_masks(catalog).enumerate() { - if index != 0 { - output.push_str(", "); - } - output.push_str(&record.dialect.op_type); - } - } - if warp_matches(catalog).next().is_some() { - output.push_str(", "); - for (index, record) in warp_matches(catalog).enumerate() { - if index != 0 { - output.push_str(", "); - } - output.push_str(&record.dialect.op_type); - } - } - for record in elect_intrinsics(catalog) { - output.push_str(", "); - output.push_str(&record.dialect.op_type); - } - if warp_barriers(catalog).next().is_some() { - output.push_str(", "); - for (index, record) in warp_barriers(catalog).enumerate() { - if index != 0 { - output.push_str(", "); - } - output.push_str(&record.dialect.op_type); - } - } - if dot_products(catalog).next().is_some() { - if sregs(catalog).next().is_some() - || ldmatrix(catalog).next().is_some() - || packed_atomics(catalog).next().is_some() - || redux(catalog).next().is_some() - || vote_intrinsics(catalog).next().is_some() - { - output.push_str(", "); - } - for (index, record) in dot_products(catalog).enumerate() { - if index != 0 { - output.push_str(", "); - } - output.push_str(&record.dialect.op_type); - } - } - if sync_intrinsics(catalog).next().is_some() { - if sregs(catalog).next().is_some() - || ldmatrix(catalog).next().is_some() - || packed_atomics(catalog).next().is_some() - || redux(catalog).next().is_some() - || vote_intrinsics(catalog).next().is_some() - || dot_products(catalog).next().is_some() - { - output.push_str(", "); - } - for (index, record) in sync_intrinsics(catalog).enumerate() { - if index != 0 { - output.push_str(", "); - } - output.push_str(&record.dialect.op_type); - } - } - if warp_shuffles(catalog).next().is_some() { - output.push_str(", "); - for (index, record) in warp_shuffles(catalog).enumerate() { - if index != 0 { - output.push_str(", "); - } - output.push_str(&record.dialect.op_type); - } - } - if packed_alus(catalog).next().is_some() { - if catalog - .intrinsics - .iter() - .any(|record| record.family != "packed_alu") - { - output.push_str(", "); - } - for (index, record) in packed_alus(catalog).enumerate() { - if index != 0 { - output.push_str(", "); - } - output.push_str(&record.dialect.op_type); - } - } - if integer_minmaxes(catalog).next().is_some() { - // packed_alu records always exist, so a separator is always needed. - output.push_str(", "); - for (index, record) in integer_minmaxes(catalog).enumerate() { - if index != 0 { - output.push_str(", "); - } - output.push_str(&record.dialect.op_type); - } - } - if packed_conversions(catalog).next().is_some() { - if catalog - .intrinsics - .iter() - .any(|record| record.family != "packed_conversion") - { - output.push_str(", "); - } - for (index, record) in packed_conversions(catalog).enumerate() { - if index != 0 { - output.push_str(", "); - } - output.push_str(&record.dialect.op_type); - } - } - for record in cp_async_copies(catalog) - .chain(cp_async_controls(catalog)) - .chain(cp_async_mbarriers(catalog)) - { - output.push_str(", "); - output.push_str(&record.dialect.op_type); - } - for record in mbarrier_basics(catalog) { - output.push_str(", "); - output.push_str(&record.dialect.op_type); - } - for record in cluster_memory(catalog) { - output.push_str(", "); - output.push_str(&record.dialect.op_type); - } - for record in stmatrices(catalog) { - output.push_str(", "); - output.push_str(&record.dialect.op_type); - } - for record in mbarrier_extended(catalog) { - output.push_str(", "); - output.push_str(&record.dialect.op_type); - } - for record in tma_intrinsics(catalog) { - output.push_str(", "); - output.push_str(&record.dialect.op_type); - } - for record in execution_controls(catalog) { - output.push_str(", "); - output.push_str(&record.dialect.op_type); - } - output.push_str( - "};\nuse llvm_export::{ops::AsmKind, types as llvm_types};\nuse pliron::{\n builtin::types::{IntegerType, Signedness},\n context::{Context, Ptr},\n derive::op_interface_impl,\n irbuild::{\n dialect_conversion::{DialectConversionRewriter, OperandsInfo},\n rewriter::Rewriter,\n },\n op::Op,\n operation::Operation,\n result::Result,\n};\n\n", - ); - if cluster_memory(catalog).next().is_some() { - if !output.contains("cast_to_shared_addrspace") { - output = output.replace( - "common::{call_intrinsic, ", - "common::{call_intrinsic, cast_to_shared_addrspace, ", - ); - } - output = output - .replace( - "use llvm_export::{ops::AsmKind, types as llvm_types};", - "use llvm_export::{op_interfaces::CastOpInterface, ops as llvm_ops, ops::AsmKind, types as llvm_types};", - ) - .replace( - "dialect_conversion::{DialectConversionRewriter, OperandsInfo},\n rewriter::Rewriter,", - "dialect_conversion::{DialectConversionRewriter, OperandsInfo},\n inserter::Inserter,\n rewriter::Rewriter,", - ); - } - if tcgen05_intrinsics(catalog).next().is_some() { - output = output.replace( - "use llvm_export::{ops::AsmKind, types as llvm_types};", - "use llvm_export::{ops as llvm_ops, ops::AsmKind, types as llvm_types};", - ); - output = output.replace( - "builtin::types::{IntegerType, Signedness}", - "builtin::types::{FP32Type, IntegerType, Signedness}", - ); - output = output.replace( - "dialect_conversion::{DialectConversionRewriter, OperandsInfo},\n rewriter::Rewriter,", - "dialect_conversion::{DialectConversionRewriter, OperandsInfo},\n inserter::Inserter,\n rewriter::Rewriter,", - ); - } - output.push_str("use pliron::location::Located;\n\n"); - if mbarrier_extended(catalog).next().is_some() { - if !output.contains("cast_to_shared_addrspace") { - output = output.replace( - "common::{call_intrinsic,", - "common::{call_intrinsic, cast_to_shared_addrspace,", - ); - } - output = if output.contains("inline_asm_sideeffect}") { - output.replace( - "inline_asm_sideeffect}", - "inline_asm_sideeffect, trunc_to_i1}", - ) - } else { - output.replace( - "inline_asm_convergent}", - "inline_asm_convergent, trunc_to_i1}", - ) - }; - output.push_str("use pliron::value::DefiningEntity;\n\n"); - } - output.push_str( - "fn convert_zero_operand_scalar_direct(\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n op: Ptr,\n width: u32,\n intrinsic_name: &str,\n) -> Result<()> {\n let result_ty = IntegerType::get(ctx, width, Signedness::Signless);\n let function_ty = llvm_types::FuncType::get(ctx, result_ty.into(), vec![], false);\n let call = call_intrinsic(ctx, rewriter, op, intrinsic_name, function_ty, vec![])?;\n rewriter.replace_operation(ctx, op, call);\n Ok(())\n}\n\n", - ); - if tcgen05_intrinsics(catalog).next().is_some() { - output.push_str( - r#"fn convert_generated_tcgen05_void( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - arity: usize, - template: &str, - constraints: &str, -) -> Result<()> { - let operands: Vec<_> = op.deref(ctx).operands().collect(); - if operands.len() != arity { - return pliron::input_err_noloc!( - "generated tcgen05 operation expects {arity} operands, got {}", - operands.len() - ); - } - let void_ty = llvm_types::VoidType::get(ctx); - inline_asm_convergent( - ctx, - rewriter, - op, - void_ty.into(), - operands, - template, - constraints, - ); - rewriter.erase_operation(ctx, op); - Ok(()) -} - -#[allow(clippy::too_many_arguments)] -fn convert_generated_tcgen05_load( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - arity: usize, - count: usize, - integer_results: bool, - template: &str, - constraints: &str, -) -> Result<()> { - let operands: Vec<_> = op.deref(ctx).operands().collect(); - if operands.len() != arity || !(1..=2).contains(&arity) || !(1..=128).contains(&count) || op.deref(ctx).get_num_results() != count { - return pliron::input_err_noloc!( - "generated tcgen05 load requires 1..=2 operands and 1..=128 results" - ); - } - let scalar_ty: pliron::r#type::TypeHandle = if integer_results { - IntegerType::get(ctx, 32, Signedness::Signless).into() - } else { - FP32Type::get(ctx).into() - }; - if count == 1 { - let inline_asm = inline_asm_convergent( - ctx, - rewriter, - op, - scalar_ty, - operands, - template, - constraints, - ); - let result = inline_asm.deref(ctx).get_result(0); - rewriter.replace_operation_with_values(ctx, op, vec![result]); - return Ok(()); - } - let result_ty = llvm_types::StructType::get_unnamed( - ctx, - ( - (0..count).map(|_| scalar_ty).collect(), - llvm_types::StructLayout::Unpacked, - ), - ); - let inline_asm = inline_asm_convergent( - ctx, - rewriter, - op, - result_ty.into(), - operands, - template, - constraints, - ); - let result = inline_asm.deref(ctx).get_result(0); - let mut values = Vec::with_capacity(count); - for index in 0..count as u32 { - let extract = llvm_ops::ExtractValueOp::new(ctx, result, vec![index]) - .map_err(|error| pliron::input_error_noloc!("{}", error))?; - rewriter.insert_operation(ctx, extract.get_operation()); - values.push(extract.get_operation().deref(ctx).get_result(0)); - } - rewriter.replace_operation_with_values(ctx, op, values); - Ok(()) -} - -"#, - ); - } - if stmatrices(catalog).next().is_some() { - output.push_str( - r#"fn convert_generated_stmatrix( - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - op: Ptr, - register_count: usize, - transposed: bool, - typed_intrinsic_name: &str, -) -> Result<()> { - let operands: Vec<_> = op.deref(ctx).operands().collect(); - if operands.len() != register_count + 1 || !matches!(register_count, 2 | 4) { - return pliron::input_err_noloc!( - "stmatrix requires one pointer and two or four registers" - ); - } - let void_ty = llvm_types::VoidType::get(ctx); - match context::lowering_options(ctx).intrinsic_backend { - IntrinsicBackend::LlvmNvptx => { - let shared = cast_to_shared_addrspace(ctx, rewriter, operands[0]); - let shared_ty = llvm_types::PointerType::get( - ctx, - llvm_types::address_space::SHARED, - ); - let i32_ty = IntegerType::get(ctx, 32, Signedness::Signless); - let argument_types = std::iter::once(shared_ty.into()) - .chain(std::iter::repeat_n(i32_ty.into(), register_count)) - .collect(); - let function_ty = llvm_types::FuncType::get( - ctx, - void_ty.into(), - argument_types, - false, - ); - let arguments = std::iter::once(shared) - .chain(operands.iter().skip(1).copied()) - .collect(); - call_intrinsic( - ctx, - rewriter, - op, - typed_intrinsic_name, - function_ty, - arguments, - )?; - } - IntrinsicBackend::LibNvvm => { - let registers = (1..=register_count) - .map(|index| format!("${index}")) - .collect::>() - .join(", "); - let trans = if transposed { ".trans" } else { "" }; - let template = format!( - "{{ .reg .u64 %ptr64; .reg .u32 %ptr32; cvta.to.shared.u64 %ptr64, $0; cvt.u32.u64 %ptr32, %ptr64; stmatrix.sync.aligned.m8n8.x{register_count}{trans}.shared.b16 [%ptr32], {{{registers}}}; }}" - ); - let constraints = std::iter::once("l") - .chain(std::iter::repeat_n("r", register_count)) - .chain(std::iter::once("~{memory}")) - .collect::>() - .join(","); - inline_asm_convergent( - ctx, - rewriter, - op, - void_ty.into(), - operands, - &template, - &constraints, - ); - } - } - rewriter.erase_operation(ctx, op); - Ok(()) -} - -"#, - ); - } - for record in sregs(catalog) { - writeln!( - output, - "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", - record.dialect.op_type - ) - .unwrap(); - output.push_str( - " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n", - ); - if record.special_register.is_none() { - writeln!( - output, - " convert_zero_operand_scalar_direct(ctx, rewriter, self.get_operation(), {}, {:?})", - record.scalar_width().unwrap(), - record.llvm_identifier() - ) - .unwrap(); - } else { - output.push_str(" match context::lowering_options(ctx).intrinsic_backend {\n"); - for (backend, backend_variant) in [ - (IntrinsicBackend::LlvmNvptx, "LlvmNvptx"), - (IntrinsicBackend::LibNvvm, "LibNvvm"), - ] { - writeln!( - output, - " IntrinsicBackend::{backend_variant} => {{" - ) - .unwrap(); - match special_register_backend_mechanism(record, backend) { - BackendLoweringMechanism::TypedNvvm => writeln!( - output, - " convert_zero_operand_scalar_direct(ctx, rewriter, self.get_operation(), {}, {:?})", - record.scalar_width().unwrap(), - record.llvm_identifier() - ) - .unwrap(), - BackendLoweringMechanism::InlinePtx => writeln!( - output, - " convert_sreg_read_inline(ctx, rewriter, self.get_operation(), {}, {:?}, {:?}, {})", - record.scalar_width().unwrap(), - special_register_inline_template(record), - special_register_output_constraint(record), - special_register_asm_kind(record), - ) - .unwrap(), - } - output.push_str(" }\n"); - } - output.push_str(" }\n"); - } - output.push_str(" }\n}\n\n"); - } - for record in active_masks(catalog) { - debug_assert_eq!( - record.active_mask.as_ref().unwrap().adapter, - ActiveMaskAdapter::DirectZeroOperandMask - ); - writeln!( - output, - "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", - record.dialect.op_type - ) - .unwrap(); - output.push_str( - " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n operands_info: &OperandsInfo,\n ) -> Result<()> {\n let op = self.get_operation();\n match context::lowering_options(ctx).intrinsic_backend {\n IntrinsicBackend::LlvmNvptx => {\n convert_active_mask(ctx, rewriter, op, operands_info)\n }\n IntrinsicBackend::LibNvvm => {\n let i32_ty = IntegerType::get(ctx, 32, Signedness::Signless);\n let inline_asm = inline_asm_convergent(\n ctx,\n rewriter,\n op,\n i32_ty.into(),\n vec![],\n \"activemask.b32 $0;\",\n \"=r,~{memory}\",\n );\n rewriter.replace_operation(ctx, op, inline_asm);\n Ok(())\n }\n }\n }\n}\n\n", - ); - } - if ldmatrix(catalog).next().is_some() { - output.push_str( - "#[op_interface_impl]\nimpl MirToLlvmConversion for LdmatrixOp {\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let recipe = {\n let shape = self.get_attr_nvvm_ldmatrix_shape(ctx);\n let multiplicity = self.get_attr_nvvm_ldmatrix_multiplicity(ctx);\n let layout = self.get_attr_nvvm_ldmatrix_layout(ctx);\n let element = self.get_attr_nvvm_ldmatrix_element(ctx);\n let state_space = self.get_attr_nvvm_ldmatrix_state_space(ctx);\n match (shape.as_deref(), multiplicity.as_deref(), layout.as_deref(), element.as_deref(), state_space.as_deref()) {\n", - ); - for record in ldmatrix(catalog) { - let (shape, multiplicity, layout, element, state_space) = - ldmatrix_attr_variants(record); - let variant = &record.ldmatrix.as_ref().unwrap().variant; - let register_count = variant.register_count(); - let instruction_head = expected_ptx_head(record); - let intrinsic_name = record.resolved_llvm_identifier(); - writeln!( - output, - " (Some(&{shape}), Some(&{multiplicity}), Some(&{layout}), Some(&{element}), Some(&{state_space})) => ({register_count}, {instruction_head:?}, {intrinsic_name:?})," - ) - .unwrap(); - } - output.push_str( - " _ => return pliron::input_err!(\n self.get_operation().deref(ctx).loc(),\n \"nvvm.ldmatrix variant has no generated lowering recipe\",\n ),\n }\n };\n convert_generated_ldmatrix(ctx, rewriter, self.get_operation(), recipe.0, recipe.1, recipe.2)\n }\n}\n\n", - ); - for record in ldmatrix(catalog) { - let Some((op_type, _)) = ldmatrix_compat_op(record) else { - continue; - }; - let variant = &record - .ldmatrix - .as_ref() - .expect("ldmatrix compatibility record") - .variant; - let register_count = variant.register_count(); - let instruction_head = expected_ptx_head(record); - let intrinsic_name = record.resolved_llvm_identifier(); - writeln!( - output, - "#[op_interface_impl]\nimpl MirToLlvmConversion for {op_type} {{\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {{\n convert_generated_ldmatrix(\n ctx, rewriter, self.get_operation(), {register_count}, {instruction_head:?}, {intrinsic_name:?},\n )\n }}\n}}\n" - ) - .unwrap(); - } - } - for record in stmatrices(catalog) { - let (multiplicity, layout) = stmatrix_variant(record).expect("stmatrix variant"); - let count = multiplicity.register_count(); - let transposed = layout == StmatrixLayout::Transposed; - let intrinsic_name = record.resolved_llvm_identifier(); - writeln!( - output, - "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", - record.dialect.op_type - ) - .unwrap(); - writeln!( - output, - " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {{\n convert_generated_stmatrix(\n ctx, rewriter, self.get_operation(), {count}, {transposed}, {intrinsic_name:?},\n )\n }}\n}}\n" - ) - .unwrap(); - } - if let Some(record) = movmatrix(catalog).next() { - writeln!( - output, - "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {{\n let op = self.get_operation();\n let operands: Vec<_> = op.deref(ctx).operands().collect();\n if operands.len() != 1 {{\n return pliron::input_err_noloc!(\n \"movmatrix_trans_b16 requires 1 operand, got {{}}\",\n operands.len()\n );\n }}\n let result_ty = IntegerType::get(ctx, 32, Signedness::Signless);\n let asm = inline_asm_convergent(\n ctx, rewriter, op, result_ty.into(), operands, {:?}, \"=r,r\",\n );\n rewriter.replace_operation(ctx, op, asm);\n Ok(())\n }}\n}}\n", - record.dialect.op_type, - movmatrix_template(record), - ) - .unwrap(); - } - if register_mmas(catalog).next().is_some() { - output.push_str( - "#[op_interface_impl]\nimpl MirToLlvmConversion for RegisterMmaOp {\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let operation = self.operation_or_multiply(ctx);\n let kind = self.kind_or_inferred(ctx);\n let recipe = match (\n self.get_attr_nvvm_register_mma_shape(ctx).as_deref(),\n operation,\n kind,\n self.get_attr_nvvm_register_mma_accumulator(ctx).as_deref(),\n self.get_attr_nvvm_register_mma_a_element(ctx).as_deref(),\n self.get_attr_nvvm_register_mma_b_element(ctx).as_deref(),\n self.get_attr_nvvm_register_mma_a_layout(ctx).as_deref(),\n self.get_attr_nvvm_register_mma_b_layout(ctx).as_deref(),\n self.get_attr_nvvm_register_mma_overflow(ctx).as_deref(),\n ) {\n", - ); - for record in register_mmas(catalog) { - let ( - shape, - operation, - kind, - accumulator, - a_element, - b_element, - a_layout, - b_layout, - overflow, - ) = register_mma_attr_variants(record); - let (c_count, a_count, b_count, result_count) = register_mma_fragment_counts(record); - let expected_operands = - c_count + a_count + b_count + register_mma_extra_operand_count(record); - writeln!( - output, - " (Some(&{shape}), {operation}, {kind}, Some(&{accumulator}), Some(&{a_element}), Some(&{b_element}), Some(&{a_layout}), Some(&{b_layout}), Some(&{overflow})) => ({}, {result_count}, {expected_operands}, {:?}, {:?}),", - register_mma_result_variant(record), - register_mma_template(record), - register_mma_constraints(record), - ) - .unwrap(); - } - output.push_str( - " _ => return pliron::input_err!(\n self.get_operation().deref(ctx).loc(),\n \"nvvm.register_mma variant has no generated lowering recipe\",\n ),\n };\n convert_generated_register_mma(\n ctx, rewriter, self.get_operation(), recipe.0, recipe.1, recipe.2, recipe.3, recipe.4,\n )\n }\n}\n\n", - ); - for record in - register_mmas(catalog).filter(|record| register_mma_compat_op_type(record).is_some()) - { - let op_type = register_mma_compat_op_type(record).unwrap(); - let (c_count, a_count, b_count, result_count) = register_mma_fragment_counts(record); - let expected_operands = c_count + a_count + b_count; - writeln!( - output, - "#[op_interface_impl]\nimpl MirToLlvmConversion for {op_type} {{\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {{\n convert_generated_register_mma(\n ctx, rewriter, self.get_operation(), {}, {result_count}, {expected_operands}, {:?}, {:?},\n )\n }}\n}}\n", - register_mma_result_variant(record), - register_mma_template(record), - register_mma_constraints(record), - ) - .unwrap(); - } - } - if sparse_mmas(catalog).next().is_some() { - output.push_str( - "#[op_interface_impl]\nimpl MirToLlvmConversion for SparseMmaOp {\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let recipe = match (\n self.get_attr_nvvm_sparse_mma_shape(ctx).as_deref(),\n self.get_attr_nvvm_sparse_mma_accumulator(ctx).as_deref(),\n self.get_attr_nvvm_sparse_mma_a_element(ctx).as_deref(),\n self.get_attr_nvvm_sparse_mma_b_element(ctx).as_deref(),\n self.get_attr_nvvm_sparse_mma_a_layout(ctx).as_deref(),\n self.get_attr_nvvm_sparse_mma_b_layout(ctx).as_deref(),\n self.get_attr_nvvm_sparse_mma_overflow(ctx).as_deref(),\n self.get_attr_nvvm_sparse_mma_metadata(ctx).as_deref(),\n self.get_attr_nvvm_sparse_mma_selector(ctx).as_deref(),\n ) {\n", - ); - for record in sparse_mmas(catalog) { - let ( - shape, - accumulator, - a_element, - b_element, - a_layout, - b_layout, - overflow, - metadata, - selector, - ) = sparse_mma_attr_variants(record); - let (c_count, a_count, b_count, result_count) = sparse_mma_fragment_counts(record); - let expected_operands = c_count + a_count + b_count + 2; - writeln!( - output, - " (Some(&{shape}), Some(&{accumulator}), Some(&{a_element}), Some(&{b_element}), Some(&{a_layout}), Some(&{b_layout}), Some(&{overflow}), Some(&{metadata}), Some(&{selector})) => ({}, {result_count}, {expected_operands}, {:?}, {:?}),", - sparse_mma_result_variant(record), - sparse_mma_template(record), - sparse_mma_constraints(record), - ) - .unwrap(); - } - output.push_str( - " _ => return pliron::input_err!(\n self.get_operation().deref(ctx).loc(),\n \"nvvm.sparse_mma variant has no generated lowering recipe\",\n ),\n };\n convert_generated_sparse_mma(\n ctx, rewriter, self.get_operation(), recipe.0, recipe.1, recipe.2, recipe.3, recipe.4,\n )\n }\n}\n\n", - ); - } - if prmts(catalog).next().is_some() { - output.push_str( - "#[op_interface_impl]\nimpl MirToLlvmConversion for PrmtOp {\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let recipe = match self.get_attr_nvvm_prmt_mode(ctx).as_deref() {\n", - ); - for record in prmts(catalog) { - let prmt = record.prmt.as_ref().unwrap(); - let mode = match prmt.mode { - PrmtMode::Generic => "PrmtModeAttr::Generic", - PrmtMode::F4e => "PrmtModeAttr::F4e", - PrmtMode::B4e => "PrmtModeAttr::B4e", - PrmtMode::Rc8 => "PrmtModeAttr::Rc8", - PrmtMode::Ecl => "PrmtModeAttr::Ecl", - PrmtMode::Ecr => "PrmtModeAttr::Ecr", - PrmtMode::Rc16 => "PrmtModeAttr::Rc16", - }; - let modifier = match prmt.mode { - PrmtMode::Generic => "", - PrmtMode::F4e => ".f4e", - PrmtMode::B4e => ".b4e", - PrmtMode::Rc8 => ".rc8", - PrmtMode::Ecl => ".ecl", - PrmtMode::Ecr => ".ecr", - PrmtMode::Rc16 => ".rc16", - }; - let template = match prmt.adapter { - PrmtAdapter::DirectThreeOperands => { - format!("prmt.b32{modifier} $0, $1, $2, $3;") - } - PrmtAdapter::InsertZeroSecondSource => { - format!("prmt.b32{modifier} $0, $1, 0, $2;") - } - }; - writeln!( - output, - " Some(&{mode}) => ({:?}, {:?}),", - record.llvm_identifier(), - template - ) - .unwrap(); - } - output.push_str( - " _ => return pliron::input_err!(\n self.get_operation().deref(ctx).loc(),\n \"nvvm.prmt mode has no generated lowering recipe\",\n ),\n };\n convert_generated_prmt(ctx, rewriter, self.get_operation(), recipe.0, recipe.1)\n }\n}\n\n", - ); - } - if scalar_conversions(catalog).next().is_some() { - output.push_str( - "#[op_interface_impl]\nimpl MirToLlvmConversion for ScalarConversionOp {\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let recipe = match (\n self.get_attr_nvvm_scalar_conversion_rounding(ctx).as_deref(),\n self.get_attr_nvvm_scalar_conversion_saturation(ctx).as_deref(),\n ) {\n", - ); - for record in scalar_conversions(catalog) { - writeln!( - output, - " (Some(&{}), Some(&{})) => ({:?}, {:?}),", - scalar_conversion_rounding_attr(record), - scalar_conversion_saturation_attr(record), - record.llvm_identifier(), - scalar_conversion_ptx_mnemonic(record), - ) - .unwrap(); - } - output.push_str( - " _ => return pliron::input_err!(\n self.get_operation().deref(ctx).loc(),\n \"nvvm.scalar_conversion variant has no generated lowering recipe\",\n ),\n };\n convert_generated_scalar_conversion(\n ctx, rewriter, self.get_operation(), recipe.0, recipe.1,\n )\n }\n}\n\n", - ); - } - if scalar_arithmetics(catalog).next().is_some() { - output.push_str( - "#[op_interface_impl]\nimpl MirToLlvmConversion for ScalarArithmeticOp {\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let recipe = match (\n self.get_attr_nvvm_scalar_arithmetic_format(ctx).as_deref(),\n self.get_attr_nvvm_scalar_arithmetic_operation(ctx).as_deref(),\n self.get_attr_nvvm_scalar_arithmetic_rounding(ctx).as_deref(),\n self.get_attr_nvvm_scalar_arithmetic_subnormal(ctx).as_deref(),\n self.get_attr_nvvm_scalar_arithmetic_saturation(ctx).as_deref(),\n ) {\n", - ); - for record in scalar_arithmetics(catalog) { - writeln!( - output, - " (Some(&{}), Some(&{}), Some(&{}), Some(&{}), Some(&{})) => ({:?}, {:?}, {}, {}),", - scalar_arithmetic_format_attr(record), - scalar_arithmetic_operation_attr(record), - scalar_arithmetic_rounding_attr(record), - scalar_arithmetic_subnormal_attr(record), - scalar_arithmetic_saturation_attr(record), - record.llvm_identifier(), - scalar_arithmetic_ptx_mnemonic(record), - scalar_arithmetic_contract(record).format == ScalarArithmeticFormat::F64, - scalar_arithmetic_llvm_mechanism(record) - == BackendLoweringMechanism::InlinePtx, - ) - .unwrap(); - } - output.push_str( - " _ => return pliron::input_err!(\n self.get_operation().deref(ctx).loc(),\n \"nvvm.scalar_arithmetic variant has no generated lowering recipe\",\n ),\n };\n convert_generated_scalar_arithmetic(\n ctx, rewriter, self.get_operation(), recipe.0, recipe.1, recipe.2, recipe.3,\n )\n }\n}\n\n", - ); - } - if scalar_maths(catalog).next().is_some() { - output.push_str( - "#[op_interface_impl]\nimpl MirToLlvmConversion for ScalarMathOp {\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let recipe = match (\n self.get_attr_nvvm_scalar_math_format(ctx).as_deref(),\n self.get_attr_nvvm_scalar_math_operation(ctx).as_deref(),\n self.get_attr_nvvm_scalar_math_precision(ctx).as_deref(),\n self.get_attr_nvvm_scalar_math_subnormal(ctx).as_deref(),\n ) {\n", - ); - for record in scalar_maths(catalog) { - // PTX-native records carry no LLVM symbol. The identifier is - // only consumed by the typed route, which such records never - // take (their mechanism is always inline PTX), so render an - // empty name; the lowering helper rejects it defensively. - let intrinsic_name = if record.llvm.is_some() { - record.llvm_identifier() - } else { - String::new() - }; - writeln!( - output, - " (Some(&{}), Some(&{}), Some(&{}), Some(&{})) => ({:?}, {:?}, {}, {}, {}),", - scalar_math_format_attr(record), - scalar_math_operation_attr(record), - scalar_math_precision_attr(record), - scalar_math_subnormal_attr(record), - intrinsic_name, - scalar_math_ptx_mnemonic(record), - scalar_math_contract(record).format == ScalarMathFormat::F16, - scalar_math_contract(record).format == ScalarMathFormat::F64, - scalar_math_llvm_mechanism(record) == BackendLoweringMechanism::InlinePtx, - ) - .unwrap(); - } - output.push_str( - " _ => return pliron::input_err!(\n self.get_operation().deref(ctx).loc(),\n \"nvvm.scalar_math variant has no generated lowering recipe\",\n ),\n };\n convert_generated_scalar_math(\n ctx, rewriter, self.get_operation(), recipe.0, recipe.1, recipe.2, recipe.3, recipe.4,\n )\n }\n}\n\n", - ); - } - if extended_minmax(catalog).next().is_some() { - output.push_str( - "#[op_interface_impl]\nimpl MirToLlvmConversion for ExtendedMinMaxOp {\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let recipe = match (\n self.get_attr_nvvm_extended_minmax_format(ctx).as_deref(),\n self.get_attr_nvvm_extended_minmax_operation(ctx).as_deref(),\n self.get_attr_nvvm_extended_minmax_subnormal(ctx).as_deref(),\n self.get_attr_nvvm_extended_minmax_nan(ctx).as_deref(),\n self.get_attr_nvvm_extended_minmax_xorsign_abs(ctx).as_deref(),\n ) {\n", - ); - for record in extended_minmax(catalog) { - writeln!( - output, - " (Some(&{}), Some(&{}), Some(&{}), Some(&{}), Some(&{})) => ({:?}, {}),", - extended_minmax_format_attr(record), - extended_minmax_operation_attr(record), - extended_minmax_subnormal_attr(record), - extended_minmax_nan_attr(record), - extended_minmax_xorsign_abs_attr(record), - extended_minmax_ptx_mnemonic(record), - extended_minmax_carrier(record), - ) - .unwrap(); - } - output.push_str( - " _ => return pliron::input_err!(\n self.get_operation().deref(ctx).loc(),\n \"nvvm.extended_minmax variant has no generated lowering recipe\",\n ),\n };\n convert_generated_extended_minmax(\n ctx, rewriter, self.get_operation(), recipe.0, recipe.1,\n )\n }\n}\n\n", - ); - } - if cluster_barriers(catalog).next().is_some() { - output.push_str( - "#[op_interface_impl]\nimpl MirToLlvmConversion for ClusterBarrierOp {\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let recipe = match self.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() {\n", - ); - for record in cluster_barriers(catalog) { - writeln!( - output, - " Some(&{}) => ({:?}, {:?}),", - cluster_barrier_attr(record), - record.llvm_identifier(), - cluster_barrier_template(record) - ) - .unwrap(); - } - output.push_str( - " _ => return pliron::input_err!(\n self.get_operation().deref(ctx).loc(),\n \"nvvm.cluster_barrier mode has no generated lowering recipe\",\n ),\n };\n let op = self.get_operation();\n let void_ty = llvm_types::VoidType::get(ctx);\n match context::lowering_options(ctx).intrinsic_backend {\n IntrinsicBackend::LlvmNvptx => {\n let function_ty = llvm_types::FuncType::get(ctx, void_ty.into(), vec![], false);\n call_intrinsic(ctx, rewriter, op, recipe.0, function_ty, vec![])?;\n }\n IntrinsicBackend::LibNvvm => {\n inline_asm_convergent(ctx, rewriter, op, void_ty.into(), vec![], recipe.1, \"~{memory}\");\n }\n }\n rewriter.erase_operation(ctx, op);\n Ok(())\n }\n}\n\n", - ); - - let arrive = cluster_barriers(catalog) - .find(|record| { - record - .cluster_barrier - .as_ref() - .is_some_and(|barrier| barrier.mode == ClusterBarrierMode::ArriveAligned) - }) - .expect("aligned cluster arrival"); - let wait = cluster_barriers(catalog) - .find(|record| { - record - .cluster_barrier - .as_ref() - .is_some_and(|barrier| barrier.mode == ClusterBarrierMode::WaitAligned) - }) - .expect("aligned cluster wait"); - output.push_str( - "#[op_interface_impl]\nimpl MirToLlvmConversion for ClusterSyncOp {\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let op = self.get_operation();\n let void_ty = llvm_types::VoidType::get(ctx);\n match context::lowering_options(ctx).intrinsic_backend {\n IntrinsicBackend::LlvmNvptx => {\n let function_ty = llvm_types::FuncType::get(ctx, void_ty.into(), vec![], false);\n", - ); - writeln!( - output, - " call_intrinsic(ctx, rewriter, op, {:?}, function_ty, vec![])?;", - arrive.llvm_identifier() - ) - .unwrap(); - writeln!( - output, - " call_intrinsic(ctx, rewriter, op, {:?}, function_ty, vec![])?;", - wait.llvm_identifier() - ) - .unwrap(); - output.push_str(" }\n IntrinsicBackend::LibNvvm => {\n"); - writeln!( - output, - " inline_asm_convergent(ctx, rewriter, op, void_ty.into(), vec![], {:?}, \"~{{memory}}\");", - format!( - "{} {}", - cluster_barrier_template(arrive), - cluster_barrier_template(wait) - ) - ) - .unwrap(); - output.push_str( - " }\n }\n rewriter.erase_operation(ctx, op);\n Ok(())\n }\n}\n\n", - ); - } - if wgmma_controls(catalog).next().is_some() { - output.push_str( - "fn convert_generated_wgmma_control(\n\ - ctx: &mut Context,\n\ - rewriter: &mut DialectConversionRewriter,\n\ - op: Ptr,\n\ - intrinsic_name: &str,\n\ - ptx: &str,\n\ - has_immediate: bool,\n\ - ) -> Result<()> {\n\ - let operands: Vec<_> = op.deref(ctx).operands().collect();\n\ - if operands.len() != usize::from(has_immediate) {\n\ - return pliron::input_err!(\n\ - op.deref(ctx).loc(),\n\ - \"generated WGMMA control has the wrong operand count\",\n\ - );\n\ - }\n\ - let void_ty = llvm_types::VoidType::get(ctx);\n\ - match context::lowering_options(ctx).intrinsic_backend {\n\ - IntrinsicBackend::LlvmNvptx => {\n\ - let arguments = if has_immediate {\n\ - vec![IntegerType::get(ctx, 64, Signedness::Signless).into()]\n\ - } else {\n\ - vec![]\n\ - };\n\ - let function_ty = llvm_types::FuncType::get(\n\ - ctx, void_ty.into(), arguments, false,\n\ - );\n\ - call_intrinsic(\n\ - ctx, rewriter, op, intrinsic_name, function_ty, operands,\n\ - )?;\n\ - }\n\ - IntrinsicBackend::LibNvvm => {\n\ - let constraints = if has_immediate {\n\ - \"n,~{memory}\"\n\ - } else {\n\ - \"~{memory}\"\n\ - };\n\ - inline_asm_convergent(\n\ - ctx, rewriter, op, void_ty.into(), operands, ptx, constraints,\n\ - );\n\ - }\n\ - }\n\ - rewriter.erase_operation(ctx, op);\n\ - Ok(())\n\ - }\n\n", - ); - for record in wgmma_controls(catalog) { - let has_immediate = record - .wgmma_control - .as_ref() - .is_some_and(|control| control.mode == WgmmaControlMode::WaitGroup); - writeln!( - output, - "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {{\n convert_generated_wgmma_control(\n ctx, rewriter, self.get_operation(), {:?}, {:?}, {has_immediate},\n )\n }}\n}}\n", - record.dialect.op_type, - record.llvm_identifier(), - wgmma_control_template(record), - ) - .unwrap(); - } - } - if packed_atomics(catalog).next().is_some() { - output.push_str( - "#[op_interface_impl]\nimpl MirToLlvmConversion for PackedAtomicAddOp {\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let format = self.get_attr_nvvm_packed_atomic_format(ctx);\n let state_space = self.get_attr_nvvm_packed_atomic_state_space(ctx);\n let ordering = self.get_attr_nvvm_packed_atomic_ordering(ctx);\n let scope = self.get_attr_nvvm_packed_atomic_scope(ctx);\n let rounding = self.get_attr_nvvm_packed_atomic_rounding(ctx);\n let subnormal = self.get_attr_nvvm_packed_atomic_subnormal(ctx);\n let atomicity = self.get_attr_nvvm_packed_atomic_atomicity(ctx);\n let ptx_type = match (format.as_deref(), state_space.as_deref(), ordering.as_deref(), scope.as_deref(), rounding.as_deref(), subnormal.as_deref(), atomicity.as_deref()) {\n (Some(&PackedAtomicFormatAttr::F16x2), Some(&PackedAtomicStateSpaceAttr::Global), Some(&PackedAtomicOrderingAttr::Relaxed), Some(&PackedAtomicScopeAttr::Gpu), Some(&PackedAtomicRoundingAttr::Rn), Some(&PackedAtomicSubnormalAttr::NoFtz), Some(&PackedAtomicAtomicityAttr::PerElement)) => \"f16x2\",\n (Some(&PackedAtomicFormatAttr::Bf16x2), Some(&PackedAtomicStateSpaceAttr::Global), Some(&PackedAtomicOrderingAttr::Relaxed), Some(&PackedAtomicScopeAttr::Gpu), Some(&PackedAtomicRoundingAttr::Rn), Some(&PackedAtomicSubnormalAttr::NoFtz), Some(&PackedAtomicAtomicityAttr::PerElement)) => \"bf16x2\",\n _ => return pliron::input_err!(\n self.get_operation().deref(ctx).loc(),\n \"nvvm.packed_atomic_add attributes have no generated lowering recipe\",\n ),\n };\n convert_packed_atom_add(ctx, rewriter, self.get_operation(), ptx_type)\n }\n}\n\n", - ); - output = output.replace( - " convert_packed_atom_add(ctx, rewriter, self.get_operation(), ptx_type)", - " drop((format, state_space, ordering, scope, rounding, subnormal, atomicity));\n convert_packed_atom_add(ctx, rewriter, self.get_operation(), ptx_type)", - ); - output.push_str( - "#[op_interface_impl]\nimpl MirToLlvmConversion for NvvmAtomAddF16x2Op {\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n convert_packed_atom_add(ctx, rewriter, self.get_operation(), \"f16x2\")\n }\n}\n\n#[op_interface_impl]\nimpl MirToLlvmConversion for NvvmAtomAddBf16x2Op {\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n convert_packed_atom_add(ctx, rewriter, self.get_operation(), \"bf16x2\")\n }\n}\n\n", - ); - } - for record in redux(catalog) { - debug_assert_eq!( - record.redux.as_ref().unwrap().adapter, - ReduxAdapter::MaskValueToSourceMemberMask - ); - writeln!( - output, - "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", - record.dialect.op_type - ) - .unwrap(); - output.push_str( - " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n operands_info: &OperandsInfo,\n ) -> Result<()> {\n", - ); - writeln!( - output, - " convert_redux(ctx, rewriter, self.get_operation(), operands_info, {:?})", - record.llvm_identifier() - ) - .unwrap(); - output.push_str(" }\n}\n\n"); - } - for record in vote_intrinsics(catalog) { - debug_assert_eq!( - record.vote.as_ref().unwrap().adapter, - VoteAdapter::DirectMaskPredicate - ); - writeln!( - output, - "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", - record.dialect.op_type - ) - .unwrap(); - output.push_str( - " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n operands_info: &OperandsInfo,\n ) -> Result<()> {\n", - ); - writeln!( - output, - " convert_vote(ctx, rewriter, self.get_operation(), operands_info, {:?})", - record.llvm_identifier() - ) - .unwrap(); - output.push_str(" }\n}\n\n"); - } - for record in warp_matches(catalog) { - let warp_match = record.warp_match.as_ref().unwrap(); - let helper = match warp_match.adapter { - WarpMatchAdapter::DirectMask => "convert_match_any", - WarpMatchAdapter::ProjectMaskDiscardPredicate => "convert_match_all", - }; - let value_width = warp_match.value_width.bits(); - writeln!( - output, - "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", - record.dialect.op_type - ) - .unwrap(); - output.push_str( - " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n operands_info: &OperandsInfo,\n ) -> Result<()> {\n", - ); - writeln!( - output, - " let value_ty = IntegerType::get(ctx, {value_width}, Signedness::Signless);" - ) - .unwrap(); - writeln!( - output, - " {helper}(ctx, rewriter, self.get_operation(), operands_info, {:?}, value_ty.into())", - record.llvm_identifier() - ) - .unwrap(); - output.push_str(" }\n}\n\n"); - } - for record in elect_intrinsics(catalog) { - writeln!( - output, - "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", - record.dialect.op_type - ) - .unwrap(); - output.push_str( - " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n operands_info: &OperandsInfo,\n ) -> Result<()> {\n match context::lowering_options(ctx).intrinsic_backend {\n", - ); - for (backend, variant) in [ - (IntrinsicBackend::LlvmNvptx, "LlvmNvptx"), - (IntrinsicBackend::LibNvvm, "LibNvvm"), - ] { - let mechanism = record - .backend_lowerings - .iter() - .find(|route| route.backend == backend) - .expect("elect backend route") - .mechanism; - writeln!(output, " IntrinsicBackend::{variant} => {{").unwrap(); - match mechanism { - BackendLoweringMechanism::TypedNvvm => writeln!( - output, - " convert_elect_sync_typed(ctx, rewriter, self.get_operation(), operands_info, {:?})", - record.llvm_identifier() - ) - .unwrap(), - BackendLoweringMechanism::InlinePtx => output.push_str( - " convert_elect_sync_inline(ctx, rewriter, self.get_operation(), operands_info)\n", - ), - } - output.push_str(" }\n"); - } - output.push_str(" }\n }\n}\n\n"); - } - for record in warp_barriers(catalog) { - debug_assert_eq!( - record.warp_barrier.as_ref().unwrap().adapter, - WarpBarrierAdapter::DirectMemberMask - ); - writeln!( - output, - "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", - record.dialect.op_type - ) - .unwrap(); - output.push_str( - " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n operands_info: &OperandsInfo,\n ) -> Result<()> {\n convert_bar_warp_sync(ctx, rewriter, self.get_operation(), operands_info)\n }\n}\n\n", - ); - } - for record in warp_shuffles(catalog) { - let shuffle = record.warp_shuffle.as_ref().unwrap(); - writeln!( - output, - "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", - record.dialect.op_type - ) - .unwrap(); - output.push_str( - " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n operands_info: &OperandsInfo,\n ) -> Result<()> {\n", - ); - match shuffle.value_kind { - WarpShuffleValueKind::I32 | WarpShuffleValueKind::F32 => { - debug_assert_eq!( - shuffle.adapter, - WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp - ); - let helper = match shuffle.value_kind { - WarpShuffleValueKind::I32 => "convert_shuffle_i32", - WarpShuffleValueKind::F32 => "convert_shuffle_f32", - WarpShuffleValueKind::I64 => unreachable!(), - }; - writeln!( - output, - " {helper}(ctx, rewriter, self.get_operation(), operands_info, {:?}, {})", - record.llvm_identifier(), - shuffle.clamp, - ) - .unwrap(); - } - WarpShuffleValueKind::I64 => { - debug_assert_eq!( - shuffle.adapter, - WarpShuffleAdapter::MaskValueLaneOrDeltaSplitI64LowHighB32InsertClampReassemble - ); - let mode = match shuffle.mode { - WarpShuffleMode::Idx => "idx", - WarpShuffleMode::Bfly => "bfly", - WarpShuffleMode::Down => "down", - WarpShuffleMode::Up => "up", - }; - writeln!( - output, - " convert_shuffle_i64(ctx, rewriter, self.get_operation(), operands_info, {mode:?}, {})", - shuffle.clamp, - ) - .unwrap(); - } - } - output.push_str(" }\n}\n\n"); - } - for record in packed_alus(catalog) { - let width = packed_alu_width(record); - writeln!( - output, - "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", - record.dialect.op_type - ) - .unwrap(); - output.push_str( - " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n", - ); - writeln!( - output, - " convert_generated_packed_alu(ctx, rewriter, self.get_operation(), {:?}, {width})", - packed_alu_ptx_mnemonic(record), - ) - .unwrap(); - output.push_str(" }\n}\n\n"); - } - for record in integer_minmaxes(catalog) { - writeln!( - output, - "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", - record.dialect.op_type - ) - .unwrap(); - output.push_str( - " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n", - ); - writeln!( - output, - " convert_generated_integer_minmax(ctx, rewriter, self.get_operation(), {:?})", - integer_minmax_ptx_mnemonic(record) - ) - .unwrap(); - output.push_str(" }\n}\n\n"); - } - for record in packed_conversions(catalog) { - let conversion = record - .packed_conversion - .as_ref() - .expect("packed-conversion record"); - writeln!( - output, - "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", - record.dialect.op_type - ) - .unwrap(); - output.push_str( - " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n", - ); - match conversion.source_format { - PackedConversionSourceFormat::F32x2 => { - debug_assert_eq!( - conversion.adapter, - PackedConversionAdapter::ReverseHighLowOperands - ); - let typed_intrinsic = packed_conversion_typed_llvm_name(record) - .map(|name| format!("Some({name:?})")) - .unwrap_or_else(|| "None".into()); - writeln!( - output, - " convert_generated_packed_f32x2(ctx, rewriter, self.get_operation(), {typed_intrinsic}, {:?}, {})", - packed_conversion_ptx_mnemonic(record), - packed_conversion_result_width(record), - ) - .unwrap(); - } - PackedConversionSourceFormat::E4m3x2 - | PackedConversionSourceFormat::E5m2x2 - | PackedConversionSourceFormat::F16x2 => { - debug_assert_eq!(conversion.adapter, PackedConversionAdapter::Identity); - writeln!( - output, - " convert_generated_packed_unary(ctx, rewriter, self.get_operation(), {:?}, {}, {})", - packed_conversion_ptx_mnemonic(record), - packed_conversion_result_width(record), - packed_conversion_source_width(record), - ) - .unwrap(); - } - } - output.push_str(" }\n}\n\n"); - } - for record in cp_async_copies(catalog) { - let copy = record.cp_async_copy.as_ref().unwrap(); - let cache_policy = match copy.cache_policy { - CpAsyncCachePolicy::Ca => "ca", - CpAsyncCachePolicy::Cg => "cg", - }; - let has_source_size = copy.source_size == CpAsyncSourceSize::Runtime; - writeln!( - output, - "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", - record.dialect.op_type - ) - .unwrap(); - output.push_str( - " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n", - ); - writeln!( - output, - " convert_generated_cp_async_copy(ctx, rewriter, self.get_operation(), {cache_policy:?}, {}, {has_source_size}, {:?})", - copy.copy_size.bytes(), - record.llvm_identifier(), - ) - .unwrap(); - output.push_str(" }\n}\n\n"); - } - for record in cp_async_controls(catalog) { - let operation = match record.cp_async_control.as_ref().unwrap().operation { - CpAsyncControlOperation::CommitGroup => "commit_group", - CpAsyncControlOperation::WaitAll => "wait_all", - CpAsyncControlOperation::WaitGroup => "wait_group", - }; - writeln!( - output, - "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", - record.dialect.op_type - ) - .unwrap(); - output.push_str( - " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n", - ); - writeln!( - output, - " convert_generated_cp_async_control(ctx, rewriter, self.get_operation(), {operation:?}, {:?})", - record.llvm_identifier(), - ) - .unwrap(); - output.push_str(" }\n}\n\n"); - } - for record in cp_async_mbarriers(catalog) { - let bridge = record.cp_async_mbarrier.as_ref().unwrap(); - let operation = match bridge.operation { - CpAsyncMbarrierOperation::Arrive => "arrive", - CpAsyncMbarrierOperation::ArriveNoInc => "arrive_no_inc", - }; - let state_space = match bridge.state_space { - CpAsyncMbarrierStateSpace::Generic => "generic", - CpAsyncMbarrierStateSpace::Shared => "shared", - }; - writeln!( - output, - "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", - record.dialect.op_type - ) - .unwrap(); - output.push_str( - " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n", - ); - writeln!( - output, - " convert_generated_cp_async_mbarrier(ctx, rewriter, self.get_operation(), {operation:?}, {state_space:?}, {:?})", - record.llvm_identifier(), - ) - .unwrap(); - output.push_str(" }\n}\n\n"); - } - for record in mbarrier_basics(catalog) { - let mbarrier = record.mbarrier_basic.as_ref().unwrap(); - let helper = match (mbarrier.operation, mbarrier.adapter) { - (MbarrierBasicOperation::Init, MbarrierBasicAdapter::InitPointerCountToVoid) => { - "convert_init" - } - (MbarrierBasicOperation::Arrive, MbarrierBasicAdapter::ArrivePointerToToken) => { - "convert_arrive" - } - ( - MbarrierBasicOperation::TestWait, - MbarrierBasicAdapter::TestWaitPointerTokenToPredicate, - ) => "convert_test_wait", - (MbarrierBasicOperation::Inval, MbarrierBasicAdapter::InvalPointerToVoid) => { - "convert_inval" - } - _ => unreachable!("resolver admitted an invalid basic mbarrier adapter"), - }; - writeln!( - output, - "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", - record.dialect.op_type - ) - .unwrap(); - output.push_str( - " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n operands_info: &OperandsInfo,\n ) -> Result<()> {\n", - ); - writeln!( - output, - " {helper}(ctx, rewriter, self.get_operation(), operands_info)" - ) - .unwrap(); - output.push_str(" }\n}\n\n"); - } - for record in cluster_memory(catalog) { - let cluster = record.cluster_memory.as_ref().unwrap(); - let (template, constraints) = - crate::resolve::cluster_memory_inline_recipe(cluster.operation); - writeln!( - output, - "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", - record.dialect.op_type - ) - .unwrap(); - output.push_str( - " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let op = self.get_operation();\n let operands: Vec<_> = op.deref(ctx).operands().collect();\n", - ); - writeln!( - output, - " if operands.len() != 2 {{\n return pliron::input_err_noloc!({:?}, operands.len());\n }}", - format!("{} requires 2 operands, got {{}}", record.rust.name) - ) - .unwrap(); - output.push_str( - " let shared_pointer = cast_to_shared_addrspace(ctx, rewriter, operands[0]);\n let rank = operands[1];\n", - ); - match cluster.operation { - ClusterMemoryOperation::MapSharedRank => { - writeln!( - output, - " let i64_ty = IntegerType::get(ctx, 64, Signedness::Signless);\n let asm = inline_asm_convergent(\n ctx, rewriter, op, i64_ty.into(), vec![shared_pointer, rank], {template:?}, {constraints:?},\n );\n let mapped = asm.deref(ctx).get_result(0);\n let cluster_shared_pointer_ty = llvm_types::PointerType::get(ctx, 7);\n let int_to_ptr = llvm_ops::IntToPtrOp::new(ctx, mapped, cluster_shared_pointer_ty.into());\n rewriter.insert_operation(ctx, int_to_ptr.get_operation());\n rewriter.replace_operation(ctx, op, int_to_ptr.get_operation());\n Ok(())" - ) - .unwrap(); - } - ClusterMemoryOperation::ReadU32 => { - writeln!( - output, - " let i32_ty = IntegerType::get(ctx, 32, Signedness::Signless);\n let asm = inline_asm_convergent(\n ctx, rewriter, op, i32_ty.into(), vec![shared_pointer, rank], {template:?}, {constraints:?},\n );\n rewriter.replace_operation(ctx, op, asm);\n Ok(())" - ) - .unwrap(); - } - } - output.push_str(" }\n}\n\n"); - } - for record in mbarrier_extended(catalog) { - let contract = record.mbarrier_extended.as_ref().unwrap(); - let (template, constraints) = - crate::resolve::mbarrier_extended_inline_recipe(contract.operation); - writeln!( - output, - "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", - record.dialect.op_type - ) - .unwrap(); - let mutability = if matches!( - contract.adapter, - MbarrierExtendedAdapter::PointerTxCountBytesToTokenDroppingTxCount - | MbarrierExtendedAdapter::PointerTokenToPredicate - | MbarrierExtendedAdapter::PointerParityToPredicate - ) { - "mut " - } else { - "" - }; - output.push_str( - " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let op = self.get_operation();\n", - ); - writeln!( - output, - " let {mutability}operands: Vec<_> = op.deref(ctx).operands().collect();" - ) - .unwrap(); - let operand_count = record.dialect.operands.len(); - let operand_count_check = if operand_count == 0 { - "!operands.is_empty()".to_owned() - } else { - format!("operands.len() != {operand_count}") - }; - writeln!( - output, - " if {operand_count_check} {{\n return pliron::input_err_noloc!({:?}, operands.len());\n }}", - format!( - "{} requires {} operands, got {{}}", - record.rust.name, - operand_count - ) - ) - .unwrap(); - if matches!( - contract.adapter, - MbarrierExtendedAdapter::PointerTxCountBytesToTokenDroppingTxCount - | MbarrierExtendedAdapter::PointerTokenToPredicate - | MbarrierExtendedAdapter::PointerParityToPredicate - ) { - output.push_str( - " operands[0] = cast_to_shared_addrspace(ctx, rewriter, operands[0]);\n", - ); - } - match contract.adapter { - MbarrierExtendedAdapter::PointerTxCountBytesToTokenDroppingTxCount => { - writeln!( - output, - " let result_ty = IntegerType::get(ctx, 64, Signedness::Signless);\n let asm = inline_asm_convergent(ctx, rewriter, op, result_ty.into(), operands, {template:?}, {constraints:?});\n rewriter.replace_operation(ctx, op, asm);\n Ok(())" - ) - .unwrap(); - } - MbarrierExtendedAdapter::PointerTokenToPredicate - | MbarrierExtendedAdapter::PointerParityToPredicate => { - writeln!( - output, - " let result_ty = IntegerType::get(ctx, 32, Signedness::Signless);\n let asm = inline_asm_convergent(ctx, rewriter, op, result_ty.into(), operands, {template:?}, {constraints:?});\n let asm_result = asm.deref(ctx).get_result(0);\n let result = trunc_to_i1(ctx, rewriter, asm_result);\n let DefiningEntity::Op(result_op) = result.defining_entity() else {{ unreachable!() }};\n rewriter.replace_operation(ctx, op, result_op);\n Ok(())" - ) - .unwrap(); - } - MbarrierExtendedAdapter::RawClusterAddressToVoid - | MbarrierExtendedAdapter::ZeroOperandsToVoid - | MbarrierExtendedAdapter::NanosecondsToVoid => { - writeln!( - output, - " let void_ty = llvm_types::VoidType::get(ctx);\n inline_asm_convergent(ctx, rewriter, op, void_ty.into(), operands, {template:?}, {constraints:?});\n rewriter.erase_operation(ctx, op);\n Ok(())" - ) - .unwrap(); - } - } - output.push_str(" }\n}\n\n"); - } - for record in dot_products(catalog) { - let insert_low_half_selector = matches!( - record.dot_product.as_ref().unwrap().adapter, - DotProductAdapter::InsertLowHalfFalse - ); - writeln!( - output, - "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", - record.dialect.op_type - ) - .unwrap(); - output.push_str( - " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n", - ); - writeln!( - output, - " convert_generated_dot_product(ctx, rewriter, self.get_operation(), {:?}, {:?}, {insert_low_half_selector})", - record.llvm_identifier(), - dot_product_ptx(record), - ) - .unwrap(); - output.push_str(" }\n}\n\n"); - } - for record in clc_intrinsics(catalog) { - writeln!( - output, - "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", - record.dialect.op_type - ) - .unwrap(); - output.push_str( - " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n operands_info: &OperandsInfo,\n ) -> Result<()> {\n", - ); - match record.clc.as_ref().unwrap().operation { - ClcOperation::TryCancel | ClcOperation::TryCancelMulticast => { - writeln!( - output, - " convert_generated_clc_try_cancel(ctx, rewriter, self.get_operation(), operands_info, {:?})", - record.llvm_identifier() - ) - .unwrap(); - } - ClcOperation::QueryIsCanceled - | ClcOperation::QueryGetFirstCtaidX - | ClcOperation::QueryGetFirstCtaidY - | ClcOperation::QueryGetFirstCtaidZ => { - writeln!( - output, - " convert_generated_clc_query(ctx, rewriter, self.get_operation(), operands_info, {:?}, {})", - record.llvm_identifier(), - record.clc.as_ref().unwrap().operation == ClcOperation::QueryIsCanceled - ) - .unwrap(); - } - } - output.push_str(" }\n}\n\n"); - } - for record in execution_controls(catalog) { - let operation = ExecutionControlOperation::from_catalog_id(&record.id) - .expect("closed execution-control record"); - writeln!( - output, - "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", - record.dialect.op_type - ) - .unwrap(); - output.push_str( - " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n operands_info: &OperandsInfo,\n ) -> Result<()> {\n", - ); - match operation { - ExecutionControlOperation::BarrierCtaSync - | ExecutionControlOperation::BarrierCtaSyncAligned - | ExecutionControlOperation::BarrierCtaArrive - | ExecutionControlOperation::BarrierCtaArriveAligned => { - let template = format!( - "{}.{} $0, $1;", - record.expected_ptx.mnemonic, - record.expected_ptx.modifiers.join(".") - ); - writeln!( - output, - " convert_counted_barrier(ctx, rewriter, self.get_operation(), operands_info, {:?}, {template:?})", - record.llvm_identifier() - ) - .unwrap(); - } - ExecutionControlOperation::GridDependencyLaunchDependents - | ExecutionControlOperation::GridDependencyWait => { - let template = format!( - "{}.{};", - record.expected_ptx.mnemonic, - record.expected_ptx.modifiers.join(".") - ); - writeln!( - output, - " convert_grid_dependency(ctx, rewriter, self.get_operation(), operands_info, {:?}, {template:?})", - record.llvm_identifier() - ) - .unwrap(); - } - ExecutionControlOperation::SetMaxNRegInc | ExecutionControlOperation::SetMaxNRegDec => { - let direction = match operation { - ExecutionControlOperation::SetMaxNRegInc => "inc", - ExecutionControlOperation::SetMaxNRegDec => "dec", - _ => unreachable!("setmaxnreg operation was matched"), - }; - output.push_str( - " let register_count = self.register_count(ctx).ok_or_else(|| pliron::input_error_noloc!(\"setmaxnreg register-count attribute is invalid\"))?;\n", - ); - writeln!( - output, - " convert_setmaxnreg(ctx, rewriter, self.get_operation(), operands_info, register_count, {:?}, {direction:?})", - record.llvm_identifier() - ) - .unwrap(); - } - } - output.push_str(" }\n}\n\n"); - } - for record in tma_intrinsics(catalog) { - let operation = record.tma.as_ref().unwrap().operation; - writeln!( - output, - "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", - record.dialect.op_type - ) - .unwrap(); - output.push_str( - " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n operands_info: &OperandsInfo,\n ) -> Result<()> {\n", - ); - match operation { - TmaOperation::G2sTile1d - | TmaOperation::G2sTile2d - | TmaOperation::G2sTile3d - | TmaOperation::G2sTile4d - | TmaOperation::G2sTile5d => { - writeln!( - output, - " convert_g2s(ctx, rewriter, self.get_operation(), operands_info, {}, false)", - operation.dimensions().unwrap() - ) - .unwrap(); - } - TmaOperation::G2sTile2dMulticast => output.push_str( - " convert_g2s(ctx, rewriter, self.get_operation(), operands_info, 2, true)\n", - ), - TmaOperation::G2sTile2dMulticastCg2 => output.push_str( - " convert_g2s_multicast_cg2(ctx, rewriter, self.get_operation(), operands_info)\n", - ), - TmaOperation::S2gTile1d - | TmaOperation::S2gTile2d - | TmaOperation::S2gTile3d - | TmaOperation::S2gTile4d - | TmaOperation::S2gTile5d => { - writeln!( - output, - " convert_s2g(ctx, rewriter, self.get_operation(), operands_info, {})", - operation.dimensions().unwrap() - ) - .unwrap(); - } - TmaOperation::Reduce => { - let reduction = record - .tma - .as_ref() - .and_then(|tma| tma.reduction.as_ref()) - .expect("TMA reduction contract"); - let reduction_name = match reduction.operation { - TmaReductionOperation::Add => "add", - TmaReductionOperation::And => "and", - TmaReductionOperation::Dec => "dec", - TmaReductionOperation::Inc => "inc", - TmaReductionOperation::Max => "max", - TmaReductionOperation::Min => "min", - TmaReductionOperation::Or => "or", - TmaReductionOperation::Xor => "xor", - }; - let load_mode = match reduction.load_mode { - TmaReductionLoadMode::Tile => "tile", - TmaReductionLoadMode::Im2col => "im2col", - }; - writeln!( - output, - " convert_reduce_s2g(ctx, rewriter, self.get_operation(), operands_info, ReduceConfig::new({}, {reduction_name:?}, {load_mode:?}, {:?}))", - reduction.dimensions, - record.resolved_llvm_identifier() - ) - .unwrap(); - } - TmaOperation::CommitGroup | TmaOperation::WaitGroup | TmaOperation::WaitGroupRead => { - let operation_name = match operation { - TmaOperation::CommitGroup => "commit_group", - TmaOperation::WaitGroup => "wait_group", - TmaOperation::WaitGroupRead => "wait_group_read", - _ => unreachable!("TMA control operation was matched"), - }; - writeln!( - output, - " convert_control(ctx, rewriter, self.get_operation(), operands_info, {operation_name:?}, {:?})", - record.resolved_llvm_identifier() - ) - .unwrap(); - } - TmaOperation::PrefetchTensorMap => { - writeln!( - output, - " convert_prefetch_tensormap(ctx, rewriter, self.get_operation(), operands_info, {:?})", - record.resolved_llvm_identifier() - ) - .unwrap(); - } - TmaOperation::PrefetchTile1d - | TmaOperation::PrefetchTile2d - | TmaOperation::PrefetchTile3d - | TmaOperation::PrefetchTile4d - | TmaOperation::PrefetchTile5d - | TmaOperation::PrefetchTileGather4TwoDimensional - | TmaOperation::PrefetchTile1dCacheHint - | TmaOperation::PrefetchTile2dCacheHint - | TmaOperation::PrefetchTile3dCacheHint - | TmaOperation::PrefetchTile4dCacheHint - | TmaOperation::PrefetchTile5dCacheHint - | TmaOperation::PrefetchTileGather4TwoDimensionalCacheHint => { - writeln!( - output, - " convert_prefetch_tile(ctx, rewriter, self.get_operation(), operands_info, PrefetchTileConfig::new({}, {}, {}, {:?}))", - operation.prefetch_coordinate_count().unwrap(), - matches!( - operation, - TmaOperation::PrefetchTileGather4TwoDimensional - | TmaOperation::PrefetchTileGather4TwoDimensionalCacheHint - ), - operation.uses_prefetch_cache_hint(), - record.resolved_llvm_identifier() - ) - .unwrap(); - } - TmaOperation::ReplaceBoxDim - | TmaOperation::ReplaceElementStride - | TmaOperation::ReplaceElementType - | TmaOperation::ReplaceFillMode - | TmaOperation::ReplaceGlobalAddress - | TmaOperation::ReplaceGlobalDim - | TmaOperation::ReplaceGlobalStride - | TmaOperation::ReplaceInterleaveLayout - | TmaOperation::ReplaceRank - | TmaOperation::ReplaceSwizzleAtomicity - | TmaOperation::ReplaceSwizzleMode => { - let (field, value_kind, ordinal, immediate) = match operation { - TmaOperation::ReplaceBoxDim => ("box_dim", "u32", true, false), - TmaOperation::ReplaceElementStride => { - ("element_stride", "u32", true, false) - } - TmaOperation::ReplaceElementType => ("elemtype", "u32", false, true), - TmaOperation::ReplaceFillMode => ("fill_mode", "u32", false, true), - TmaOperation::ReplaceGlobalAddress => { - ("global_address", "address", false, false) - } - TmaOperation::ReplaceGlobalDim => ("global_dim", "u32", true, false), - TmaOperation::ReplaceGlobalStride => { - ("global_stride", "u64", true, false) - } - TmaOperation::ReplaceInterleaveLayout => { - ("interleave_layout", "u32", false, true) - } - TmaOperation::ReplaceRank => ("rank", "u32", false, false), - TmaOperation::ReplaceSwizzleAtomicity => { - ("swizzle_atomicity", "u32", false, true) - } - TmaOperation::ReplaceSwizzleMode => { - ("swizzle_mode", "u32", false, true) - } - _ => unreachable!("TMA tensor-map replace operation was matched"), - }; - writeln!( - output, - " convert_tensormap_replace(ctx, rewriter, self.get_operation(), operands_info, {:?}, {field:?}, {value_kind:?}, {ordinal}, {immediate})", - record.resolved_llvm_identifier() - ) - .unwrap(); - } - TmaOperation::FenceProxyTensorMapAcquireCluster - | TmaOperation::FenceProxyTensorMapAcquireCta - | TmaOperation::FenceProxyTensorMapAcquireGpu - | TmaOperation::FenceProxyTensorMapAcquireSystem - | TmaOperation::FenceProxyTensorMapReleaseCluster - | TmaOperation::FenceProxyTensorMapReleaseCta - | TmaOperation::FenceProxyTensorMapReleaseGpu - | TmaOperation::FenceProxyTensorMapReleaseSystem => { - let (acquire, scope) = match operation { - TmaOperation::FenceProxyTensorMapAcquireCluster => (true, "cluster"), - TmaOperation::FenceProxyTensorMapAcquireCta => (true, "cta"), - TmaOperation::FenceProxyTensorMapAcquireGpu => (true, "gpu"), - TmaOperation::FenceProxyTensorMapAcquireSystem => (true, "sys"), - TmaOperation::FenceProxyTensorMapReleaseCluster => (false, "cluster"), - TmaOperation::FenceProxyTensorMapReleaseCta => (false, "cta"), - TmaOperation::FenceProxyTensorMapReleaseGpu => (false, "gpu"), - TmaOperation::FenceProxyTensorMapReleaseSystem => (false, "sys"), - _ => unreachable!("TMA tensor-map fence operation was matched"), - }; - writeln!( - output, - " convert_tensormap_fence(ctx, rewriter, self.get_operation(), operands_info, {:?}, {acquire}, {scope:?})", - record.resolved_llvm_identifier() - ) - .unwrap(); - } - } - output.push_str(" }\n}\n\n"); - } - for record in tcgen05_non_mma_intrinsics(catalog) { - let operation = record.tcgen05.as_ref().unwrap().operation; - let (template, constraints, result_count) = tcgen05_inline_asm(record); - writeln!( - output, - "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", - record.dialect.op_type - ) - .unwrap(); - output.push_str( - " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n", - ); - if let Some(count) = result_count { - let integer_results = operation == Tcgen05Operation::Ld; - writeln!( - output, - " convert_generated_tcgen05_load(ctx, rewriter, self.get_operation(), {}, {count}, {integer_results}, {template:?}, {constraints:?})", - record.dialect.operands.len() - ) - .unwrap(); - } else { - writeln!( - output, - " convert_generated_tcgen05_void(ctx, rewriter, self.get_operation(), {}, {template:?}, {constraints:?})", - record.dialect.operands.len() - ) - .unwrap(); - } - output.push_str(" }\n}\n\n"); - } - if tcgen05_mma_intrinsics(catalog).next().is_some() { - output.push_str( - r#"#[op_interface_impl] -impl MirToLlvmConversion for Tcgen05MmaOp { - fn convert( - &self, - ctx: &mut Context, - rewriter: &mut DialectConversionRewriter, - _operands_info: &OperandsInfo, - ) -> Result<()> { - let op = self.get_operation(); - let form = self - .get_attr_nvvm_tcgen05_mma_form(ctx) - .as_deref() - .cloned(); - let kind = self - .get_attr_nvvm_tcgen05_mma_kind(ctx) - .as_deref() - .cloned(); - let cta_group = self - .get_attr_nvvm_tcgen05_mma_cta_group(ctx) - .as_deref() - .cloned(); - let collector_a = self - .get_attr_nvvm_tcgen05_mma_collector_a(ctx) - .as_deref() - .cloned(); - let b_buffer = self - .get_attr_nvvm_tcgen05_mma_b_buffer(ctx) - .as_deref() - .cloned(); - let b_usage = self - .get_attr_nvvm_tcgen05_mma_b_usage(ctx) - .as_deref() - .cloned(); - let kind = match kind { - Some(Tcgen05MmaKindAttr::F16) => "f16", - Some(Tcgen05MmaKindAttr::Tf32) => "tf32", - Some(Tcgen05MmaKindAttr::F8f6f4) => "f8f6f4", - Some(Tcgen05MmaKindAttr::I8) => "i8", - None => return pliron::input_err!( - op.deref(ctx).loc(), - "nvvm.tcgen05_mma requires a kind", - ), - }; - let (prefix, a, metadata, zero_mask, constraints, arity, base, ashift) = - match form { - Some(Tcgen05MmaFormAttr::Shared) => - ("tcgen05.mma", "$1", None, None, "r,l,l,r,r,~{memory}", 5, true, false), - Some(Tcgen05MmaFormAttr::Tensor) => - ("tcgen05.mma", "[$1]", None, None, "r,r,l,r,r,~{memory}", 5, true, false), - Some(Tcgen05MmaFormAttr::TensorAshift) => - ("tcgen05.mma", "[$1]", None, None, "r,r,l,r,r,~{memory}", 5, true, true), - Some(Tcgen05MmaFormAttr::SpShared) => - ("tcgen05.mma.sp", "$1", Some("$5"), None, "r,l,l,r,r,r,~{memory}", 6, true, false), - Some(Tcgen05MmaFormAttr::SpTensor) => - ("tcgen05.mma.sp", "[$1]", Some("$5"), None, "r,r,l,r,r,r,~{memory}", 6, true, false), - Some(Tcgen05MmaFormAttr::SpTensorAshift) => - ("tcgen05.mma.sp", "[$1]", Some("$5"), None, "r,r,l,r,r,r,~{memory}", 6, true, true), - Some(Tcgen05MmaFormAttr::WsShared) => - ("tcgen05.mma.ws", "$1", None, None, "r,l,l,r,r,~{memory}", 5, false, false), - Some(Tcgen05MmaFormAttr::WsSharedZeroColMask) => - ("tcgen05.mma.ws", "$1", None, Some("$5"), "r,l,l,r,r,l,~{memory}", 6, false, false), - Some(Tcgen05MmaFormAttr::WsSpShared) => - ("tcgen05.mma.ws.sp", "$1", Some("$5"), None, "r,l,l,r,r,r,~{memory}", 6, false, false), - Some(Tcgen05MmaFormAttr::WsSpSharedZeroColMask) => - ("tcgen05.mma.ws.sp", "$1", Some("$5"), Some("$6"), "r,l,l,r,r,r,l,~{memory}", 7, false, false), - Some(Tcgen05MmaFormAttr::WsSpTensor) => - ("tcgen05.mma.ws.sp", "[$1]", Some("$5"), None, "r,r,l,r,r,r,~{memory}", 6, false, false), - Some(Tcgen05MmaFormAttr::WsSpTensorZeroColMask) => - ("tcgen05.mma.ws.sp", "[$1]", Some("$5"), Some("$6"), "r,r,l,r,r,r,l,~{memory}", 7, false, false), - Some(Tcgen05MmaFormAttr::WsTensor) => - ("tcgen05.mma.ws", "[$1]", None, None, "r,r,l,r,r,~{memory}", 5, false, false), - Some(Tcgen05MmaFormAttr::WsTensorZeroColMask) => - ("tcgen05.mma.ws", "[$1]", None, Some("$5"), "r,r,l,r,r,l,~{memory}", 6, false, false), - None => return pliron::input_err!( - op.deref(ctx).loc(), - "nvvm.tcgen05_mma requires a form", - ), - }; - - let selector = if base { - let group = match cta_group { - Some(Tcgen05MmaCtaGroupAttr::Cg1) => 1, - Some(Tcgen05MmaCtaGroupAttr::Cg2) => 2, - None => return pliron::input_err!( - op.deref(ctx).loc(), - "base tcgen05 MMA requires a CTA-group selector", - ), - }; - let usage = match collector_a { - Some(Tcgen05MmaCollectorAAttr::Discard) => "discard", - Some(Tcgen05MmaCollectorAAttr::LastUse) => "lastuse", - Some(Tcgen05MmaCollectorAAttr::Fill) if !ashift => "fill", - Some(Tcgen05MmaCollectorAAttr::Use) if !ashift => "use", - _ => return pliron::input_err!( - op.deref(ctx).loc(), - "tcgen05 MMA collector-A selector has no generated recipe", - ), - }; - format!( - ".cta_group::{group}.kind::{kind}.collector::a::{usage}{}", - if ashift { ".ashift" } else { "" }, - ) - } else { - let buffer = match b_buffer { - Some(Tcgen05MmaBBufferAttr::B0) => 0, - Some(Tcgen05MmaBBufferAttr::B1) => 1, - Some(Tcgen05MmaBBufferAttr::B2) => 2, - Some(Tcgen05MmaBBufferAttr::B3) => 3, - None => return pliron::input_err!( - op.deref(ctx).loc(), - "warp-specialized tcgen05 MMA requires a B-buffer selector", - ), - }; - let usage = match b_usage { - Some(Tcgen05MmaBUsageAttr::Discard) => "discard", - Some(Tcgen05MmaBUsageAttr::LastUse) => "lastuse", - Some(Tcgen05MmaBUsageAttr::Fill) => "fill", - Some(Tcgen05MmaBUsageAttr::Use) => "use", - None => return pliron::input_err!( - op.deref(ctx).loc(), - "warp-specialized tcgen05 MMA requires a B-usage selector", - ), - }; - format!( - ".cta_group::1.kind::{kind}.collector::b{buffer}::{usage}" - ) - }; - - let mut operands = format!("[$0], {a}, $2"); - if let Some(metadata) = metadata { - operands.push_str(&format!(", [{metadata}]")); - } - operands.push_str(", $3, %enable_pred"); - if let Some(zero_mask) = zero_mask { - operands.push_str(&format!(", {zero_mask}")); - } - let template = format!( - "{{ .reg .pred %enable_pred; setp.ne.s32 %enable_pred, $4, 0; {prefix}{selector} {operands}; }}" - ); - convert_generated_tcgen05_void( - ctx, rewriter, op, arity, &template, constraints, - ) - } -} - -"#, - ); - } - for record in debug_controls(catalog) { - writeln!( - output, - "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", - record.dialect.op_type - ) - .unwrap(); - output.push_str( - " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let op = self.get_operation();\n let void_ty = llvm_types::VoidType::get(ctx);\n", - ); - match record.debug_control.as_ref().unwrap().operation { - DebugControlOperation::Trap => output.push_str( - " inline_asm_sideeffect(ctx, rewriter, op, void_ty.into(), vec![], \"trap;\", \"\");\n", - ), - DebugControlOperation::Breakpoint => output.push_str( - " inline_asm_sideeffect(ctx, rewriter, op, void_ty.into(), vec![], \"brkpt;\", \"\");\n", - ), - DebugControlOperation::Pmevent => output.push_str( - " let Some(event_id) = self.event_id(ctx) else {\n\ - return pliron::input_err!(\n\ - op.deref(ctx).loc(),\n\ - \"nvvm.pmevent requires a u32 event ID in 0..=15\",\n\ - );\n\ - };\n\ - let template = format!(\"pmevent {event_id};\");\n\ - inline_asm_sideeffect(ctx, rewriter, op, void_ty.into(), vec![], &template, \"\");\n", - ), - } - output.push_str(" rewriter.erase_operation(ctx, op);\n Ok(())\n }\n}\n\n"); - } - for record in sync_intrinsics(catalog) { - writeln!( - output, - "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", - record.dialect.op_type - ) - .unwrap(); - if record.id == "sync_threads" { - output.push_str( - " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let op = self.get_operation();\n let void_ty = llvm_types::VoidType::get(ctx);\n match context::lowering_options(ctx).intrinsic_backend {\n IntrinsicBackend::LlvmNvptx => {\n let i32_ty = IntegerType::get(ctx, 32, Signedness::Signless);\n let barrier_id = create_i32_const(ctx, rewriter, 0);\n let function_ty = llvm_types::FuncType::get(ctx, void_ty.into(), vec![i32_ty.into()], false);\n", - ); - writeln!( - output, - " call_intrinsic(ctx, rewriter, op, {:?}, function_ty, vec![barrier_id])?;", - record.llvm_identifier() - ) - .unwrap(); - output.push_str( - " }\n IntrinsicBackend::LibNvvm => {\n inline_asm_convergent(ctx, rewriter, op, void_ty.into(), vec![], \"bar.sync 0;\", \"~{memory}\");\n }\n }\n rewriter.erase_operation(ctx, op);\n Ok(())\n }\n}\n\n", - ); - } else { - debug_assert!(threadfence_ptx_level(record).is_some()); - output.push_str( - " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let op = self.get_operation();\n let void_ty = llvm_types::VoidType::get(ctx);\n let function_ty = llvm_types::FuncType::get(ctx, void_ty.into(), vec![], false);\n", - ); - writeln!( - output, - " call_intrinsic(ctx, rewriter, op, {:?}, function_ty, vec![])?;", - record.llvm_identifier() - ) - .unwrap(); - output.push_str( - " rewriter.erase_operation(ctx, op);\n Ok(())\n }\n}\n\n", - ); - } - } - output -} - -fn render_collector(catalog: &CatalogFile, hash: &str) -> String { - let mut output = rust_header(catalog, hash); - let abi_namespace = format!("__cuda_oxide_intrinsic_abi_v{}", catalog.intrinsic_abi); - output.push_str("//! Generated collector predicates for intrinsic placeholder functions.\n\n"); - writeln!( - output, - "pub(crate) const GENERATED_INTRINSIC_ABI: u32 = {};", - catalog.intrinsic_abi - ) - .unwrap(); - writeln!( - output, - "pub(crate) const GENERATED_INTRINSIC_ABI_NAMESPACE: &str = {abi_namespace:?};" - ) - .unwrap(); - output.push_str( - "pub(crate) const GENERATED_INTRINSIC_CRATES: &[&str] = &[\n \"cuda_intrinsics\",\n \"cuda-intrinsics\",\n];\n\npub(crate) const GENERATED_INTRINSIC_CANONICAL_PATHS: &[&str] = &[\n", - ); - for record in &catalog.intrinsics { - writeln!(output, " {:?},", record.rust.canonical_path).unwrap(); - } - output.push_str("];\n\npub(crate) const GENERATED_INTRINSIC_PUBLIC_PATHS: &[&str] = &[\n"); - for record in &catalog.intrinsics { - writeln!(output, " {:?},", record.rust.public_path).unwrap(); - } - output.push_str("];\n\npub(crate) fn is_generated_intrinsic_crate(crate_name: &str) -> bool {\n GENERATED_INTRINSIC_CRATES.contains(&crate_name)\n}\n\npub(crate) fn is_generated_intrinsic_canonical_path(path: &str) -> bool {\n matches!(path,\n"); - let canonical_paths: Vec<_> = catalog - .intrinsics - .iter() - .map(|record| record.rust.canonical_path.as_str()) - .collect(); - render_string_patterns(&mut output, &canonical_paths, " "); - output.push_str(" )\n}\n\npub(crate) fn is_generated_intrinsic_compatibility_path(path: &str) -> bool {\n"); - let compatibility_paths: Vec<_> = catalog - .intrinsics - .iter() - .flat_map(|record| record.rust.compatibility_paths.iter().map(String::as_str)) - .collect(); - if compatibility_paths.is_empty() { - output.push_str(" let _ = path;\n false\n"); - } else { - output.push_str(" matches!(path,\n"); - render_string_patterns(&mut output, &compatibility_paths, " "); - output.push_str(" )\n"); - } - output.push_str( - "}\n\npub(crate) fn is_generated_intrinsic_placeholder(crate_name: &str, path: &str) -> bool {\n if is_generated_intrinsic_crate(crate_name) {\n is_generated_intrinsic_canonical_path(path)\n } else if matches!(crate_name, \"cuda_device\" | \"cuda-device\") {\n is_generated_intrinsic_compatibility_path(path)\n } else {\n false\n }\n}\n", - ); - output -} - -fn generated_selection_constraints(selection: &CatalogSelection) -> String { - let address_space = match selection.constraints.address_space { - None => "None", - Some(ImportedAddressSpace::Generic) => "Some(GeneratedSelectionAddressSpace::Generic)", - Some(ImportedAddressSpace::Shared) => "Some(GeneratedSelectionAddressSpace::Shared)", - }; - let immediate_bindings = selection - .constraints - .immediate_bindings - .iter() - .map(|binding| { - format!( - "GeneratedImmediateBinding {{ argument_index: {}, value: {} }}", - binding.argument_index, binding.value - ) - }) - .collect::>() - .join(", "); - format!( - "GeneratedSelectionConstraints {{ address_space: {address_space}, immediate_bindings: &[{immediate_bindings}] }}" - ) -} - -fn generated_selection_alternatives(selections: &[CatalogSelection]) -> String { - let mut output = String::from("&["); - for selection in selections { - write!( - output, - "GeneratedSelectionAlternative {{ source_record: {:?}, asm: {:?}, predicates: &{:?}, constraints: {} }},", - selection.source_record, - selection.asm, - selection.predicates, - generated_selection_constraints(selection), - ) - .unwrap(); - } - output.push(']'); - output -} - -fn generated_intrinsic_variant(record: &CatalogIntrinsic) -> String { - if let Some(mma) = record - .tcgen05 - .as_ref() - .and_then(|tcgen05| tcgen05.mma.as_ref()) - { - return format!( - "GeneratedIntrinsicVariant::Tcgen05Mma {{ form: GeneratedTcgen05MmaForm::{}, target_selector: GeneratedTcgen05MmaTargetSelector::Kind, compatibility_alias: {} }}", - tcgen05_mma_form_name(mma.form), - mma.alias.is_some(), - ); - } - if let Some(control) = &record.wgmma_control { - let mode = match control.mode { - WgmmaControlMode::Fence => "GeneratedWgmmaControlMode::Fence", - WgmmaControlMode::CommitGroup => "GeneratedWgmmaControlMode::CommitGroup", - WgmmaControlMode::WaitGroup => "GeneratedWgmmaControlMode::WaitGroup", - }; - return format!("GeneratedIntrinsicVariant::WgmmaControl {{ mode: {mode} }}"); - } - if let Some(barrier) = &record.cluster_barrier { - let mode = match barrier.mode { - ClusterBarrierMode::Arrive => "GeneratedClusterBarrierMode::Arrive", - ClusterBarrierMode::ArriveAligned => "GeneratedClusterBarrierMode::ArriveAligned", - ClusterBarrierMode::ArriveRelaxed => "GeneratedClusterBarrierMode::ArriveRelaxed", - ClusterBarrierMode::ArriveRelaxedAligned => { - "GeneratedClusterBarrierMode::ArriveRelaxedAligned" - } - ClusterBarrierMode::Wait => "GeneratedClusterBarrierMode::Wait", - ClusterBarrierMode::WaitAligned => "GeneratedClusterBarrierMode::WaitAligned", - }; - return format!("GeneratedIntrinsicVariant::ClusterBarrier {{ mode: {mode} }}"); - } - if let Some(prmt) = &record.prmt { - let mode = match prmt.mode { - PrmtMode::Generic => "GeneratedPrmtMode::Generic", - PrmtMode::F4e => "GeneratedPrmtMode::F4e", - PrmtMode::B4e => "GeneratedPrmtMode::B4e", - PrmtMode::Rc8 => "GeneratedPrmtMode::Rc8", - PrmtMode::Ecl => "GeneratedPrmtMode::Ecl", - PrmtMode::Ecr => "GeneratedPrmtMode::Ecr", - PrmtMode::Rc16 => "GeneratedPrmtMode::Rc16", - }; - return format!("GeneratedIntrinsicVariant::Prmt {{ mode: {mode} }}"); - } - if let Some(packed) = &record.packed_atomic { - let format = match packed.format { - PackedAtomicFormat::F16x2 => "GeneratedPackedAtomicFormat::F16x2", - PackedAtomicFormat::Bf16x2 => "GeneratedPackedAtomicFormat::Bf16x2", - }; - return format!("GeneratedIntrinsicVariant::PackedAtomic {{ format: {format} }}"); - } - if let Some(mma) = &record.register_mma { - let shape = match mma.shape { - RegisterMmaShape::M8n8k4 => "GeneratedRegisterMmaShape::M8n8k4", - RegisterMmaShape::M8n8k16 => "GeneratedRegisterMmaShape::M8n8k16", - RegisterMmaShape::M8n8k32 => "GeneratedRegisterMmaShape::M8n8k32", - RegisterMmaShape::M8n8k128 => "GeneratedRegisterMmaShape::M8n8k128", - RegisterMmaShape::M16n8k4 => "GeneratedRegisterMmaShape::M16n8k4", - RegisterMmaShape::M16n8k8 => "GeneratedRegisterMmaShape::M16n8k8", - RegisterMmaShape::M16n8k16 => "GeneratedRegisterMmaShape::M16n8k16", - RegisterMmaShape::M16n8k32 => "GeneratedRegisterMmaShape::M16n8k32", - RegisterMmaShape::M16n8k64 => "GeneratedRegisterMmaShape::M16n8k64", - RegisterMmaShape::M16n8k128 => "GeneratedRegisterMmaShape::M16n8k128", - RegisterMmaShape::M16n8k256 => "GeneratedRegisterMmaShape::M16n8k256", - }; - let operation = match mma.operation { - RegisterMmaOperation::Multiply => "GeneratedRegisterMmaOperation::Multiply", - RegisterMmaOperation::AndPopc => "GeneratedRegisterMmaOperation::AndPopc", - RegisterMmaOperation::XorPopc => "GeneratedRegisterMmaOperation::XorPopc", - }; - let kind = match register_mma_effective_kind(record) { - RegisterMmaKind::Standard => "GeneratedRegisterMmaKind::Standard", - RegisterMmaKind::F8f6f4 => "GeneratedRegisterMmaKind::F8f6f4", - RegisterMmaKind::Mxf8f6f4 => "GeneratedRegisterMmaKind::Mxf8f6f4", - }; - let accumulator = match mma.accumulator { - RegisterMmaAccumulator::F16 => "GeneratedRegisterMmaAccumulator::F16", - RegisterMmaAccumulator::F32 => "GeneratedRegisterMmaAccumulator::F32", - RegisterMmaAccumulator::F64 => "GeneratedRegisterMmaAccumulator::F64", - RegisterMmaAccumulator::S32 => "GeneratedRegisterMmaAccumulator::S32", - }; - let element = |element| match element { - RegisterMmaElement::Bf16 => "GeneratedRegisterMmaElement::Bf16", - RegisterMmaElement::E2m1 => "GeneratedRegisterMmaElement::E2m1", - RegisterMmaElement::E2m3 => "GeneratedRegisterMmaElement::E2m3", - RegisterMmaElement::E3m2 => "GeneratedRegisterMmaElement::E3m2", - RegisterMmaElement::E4m3 => "GeneratedRegisterMmaElement::E4m3", - RegisterMmaElement::E5m2 => "GeneratedRegisterMmaElement::E5m2", - RegisterMmaElement::F16 => "GeneratedRegisterMmaElement::F16", - RegisterMmaElement::Tf32 => "GeneratedRegisterMmaElement::Tf32", - RegisterMmaElement::F64 => "GeneratedRegisterMmaElement::F64", - RegisterMmaElement::B1 => "GeneratedRegisterMmaElement::B1", - RegisterMmaElement::S4 => "GeneratedRegisterMmaElement::S4", - RegisterMmaElement::U4 => "GeneratedRegisterMmaElement::U4", - RegisterMmaElement::S8 => "GeneratedRegisterMmaElement::S8", - RegisterMmaElement::U8 => "GeneratedRegisterMmaElement::U8", - }; - let layout = |layout| match layout { - RegisterMmaLayout::Row => "GeneratedRegisterMmaLayout::Row", - RegisterMmaLayout::Col => "GeneratedRegisterMmaLayout::Col", - }; - let overflow = match mma.overflow { - RegisterMmaOverflow::NotApplicable => "GeneratedRegisterMmaOverflow::NotApplicable", - RegisterMmaOverflow::Wrapping => "GeneratedRegisterMmaOverflow::Wrapping", - RegisterMmaOverflow::Satfinite => "GeneratedRegisterMmaOverflow::Satfinite", - }; - return format!( - "GeneratedIntrinsicVariant::RegisterMma {{ shape: {shape}, operation: {operation}, kind: {kind}, accumulator: {accumulator}, a_element: {}, b_element: {}, a_layout: {}, b_layout: {}, overflow: {overflow} }}", - element(mma.a_element), - element(mma.b_element), - layout(mma.a_layout), - layout(mma.b_layout), - ); - } - if let Some(mma) = &record.sparse_mma { - let shape = match mma.shape { - SparseMmaShape::M16n8k32 => "GeneratedSparseMmaShape::M16n8k32", - SparseMmaShape::M16n8k64 => "GeneratedSparseMmaShape::M16n8k64", - SparseMmaShape::M16n8k128 => "GeneratedSparseMmaShape::M16n8k128", - }; - let accumulator = match mma.accumulator { - SparseMmaAccumulator::F16 => "GeneratedSparseMmaAccumulator::F16", - SparseMmaAccumulator::F32 => "GeneratedSparseMmaAccumulator::F32", - SparseMmaAccumulator::S32 => "GeneratedSparseMmaAccumulator::S32", - }; - let element = |element| match element { - SparseMmaElement::E2m1 => "GeneratedSparseMmaElement::E2m1", - SparseMmaElement::E2m3 => "GeneratedSparseMmaElement::E2m3", - SparseMmaElement::E3m2 => "GeneratedSparseMmaElement::E3m2", - SparseMmaElement::E4m3 => "GeneratedSparseMmaElement::E4m3", - SparseMmaElement::E5m2 => "GeneratedSparseMmaElement::E5m2", - SparseMmaElement::S4 => "GeneratedSparseMmaElement::S4", - SparseMmaElement::U4 => "GeneratedSparseMmaElement::U4", - SparseMmaElement::S8 => "GeneratedSparseMmaElement::S8", - SparseMmaElement::U8 => "GeneratedSparseMmaElement::U8", - }; - let layout = |layout| match layout { - SparseMmaLayout::Row => "GeneratedSparseMmaLayout::Row", - SparseMmaLayout::Col => "GeneratedSparseMmaLayout::Col", - }; - let overflow = match mma.overflow { - SparseMmaOverflow::NotApplicable => "GeneratedSparseMmaOverflow::NotApplicable", - SparseMmaOverflow::Wrapping => "GeneratedSparseMmaOverflow::Wrapping", - SparseMmaOverflow::Satfinite => "GeneratedSparseMmaOverflow::Satfinite", - }; - let metadata = match mma.metadata { - SparseMmaMetadata::Standard => "GeneratedSparseMmaMetadata::Standard", - SparseMmaMetadata::Ordered => "GeneratedSparseMmaMetadata::Ordered", - }; - let selector = match mma.selector { - SparseMmaSelector::ImmediateZeroOrOne => { - "GeneratedSparseMmaSelector::ImmediateZeroOrOne" - } - SparseMmaSelector::ImmediateZero => "GeneratedSparseMmaSelector::ImmediateZero", - }; - return format!( - "GeneratedIntrinsicVariant::SparseMma {{ shape: {shape}, accumulator: {accumulator}, a_element: {}, b_element: {}, a_layout: {}, b_layout: {}, overflow: {overflow}, metadata: {metadata}, selector: {selector} }}", - element(mma.a_element), - element(mma.b_element), - layout(mma.a_layout), - layout(mma.b_layout), - ); - } - if let Some(arithmetic) = &record.scalar_arithmetic { - let format = match arithmetic.format { - ScalarArithmeticFormat::F32 => "GeneratedScalarArithmeticFormat::F32", - ScalarArithmeticFormat::F64 => "GeneratedScalarArithmeticFormat::F64", - }; - let operation = match arithmetic.operation { - ScalarArithmeticOperation::Mul => "GeneratedScalarArithmeticOperation::Mul", - ScalarArithmeticOperation::Div => "GeneratedScalarArithmeticOperation::Div", - ScalarArithmeticOperation::Fma => "GeneratedScalarArithmeticOperation::Fma", - ScalarArithmeticOperation::Add => "GeneratedScalarArithmeticOperation::Add", - }; - let rounding = match arithmetic.rounding { - ScalarArithmeticRounding::Rn => "GeneratedScalarArithmeticRounding::Rn", - ScalarArithmeticRounding::Rz => "GeneratedScalarArithmeticRounding::Rz", - ScalarArithmeticRounding::Rm => "GeneratedScalarArithmeticRounding::Rm", - ScalarArithmeticRounding::Rp => "GeneratedScalarArithmeticRounding::Rp", - }; - let subnormal = match arithmetic.subnormal { - ScalarArithmeticSubnormal::Preserve => "GeneratedScalarArithmeticSubnormal::Preserve", - ScalarArithmeticSubnormal::Ftz => "GeneratedScalarArithmeticSubnormal::Ftz", - }; - let saturation = match arithmetic.saturation { - ScalarArithmeticSaturation::None => "GeneratedScalarArithmeticSaturation::None", - ScalarArithmeticSaturation::Sat => "GeneratedScalarArithmeticSaturation::Sat", - }; - return format!( - "GeneratedIntrinsicVariant::ScalarArithmetic {{ format: {format}, operation: {operation}, rounding: {rounding}, subnormal: {subnormal}, saturation: {saturation} }}" - ); - } - if let Some(minmax) = &record.extended_minmax { - let format = match minmax.format { - ExtendedMinMaxFormat::F32 => "GeneratedExtendedMinMaxFormat::F32", - ExtendedMinMaxFormat::F16 => "GeneratedExtendedMinMaxFormat::F16", - ExtendedMinMaxFormat::Bf16 => "GeneratedExtendedMinMaxFormat::Bf16", - ExtendedMinMaxFormat::F16x2 => "GeneratedExtendedMinMaxFormat::F16x2", - ExtendedMinMaxFormat::Bf16x2 => "GeneratedExtendedMinMaxFormat::Bf16x2", - }; - let operation = match minmax.operation { - ExtendedMinMaxOperation::Min => "GeneratedExtendedMinMaxOperation::Min", - ExtendedMinMaxOperation::Max => "GeneratedExtendedMinMaxOperation::Max", - }; - let subnormal = match minmax.subnormal { - ExtendedMinMaxSubnormal::Preserve => "GeneratedExtendedMinMaxSubnormal::Preserve", - ExtendedMinMaxSubnormal::Ftz => "GeneratedExtendedMinMaxSubnormal::Ftz", - }; - let nan = match minmax.nan { - ExtendedMinMaxNan::Number => "GeneratedExtendedMinMaxNan::Number", - ExtendedMinMaxNan::Nan => "GeneratedExtendedMinMaxNan::Nan", - }; - return format!( - "GeneratedIntrinsicVariant::ExtendedMinMax {{ format: {format}, operation: {operation}, subnormal: {subnormal}, nan: {nan}, xorsign_abs: {} }}", - minmax.xorsign_abs - ); - } - if let Some(conversion) = &record.scalar_conversion { - let rounding = match conversion.rounding { - ScalarConversionRounding::NearestAway => { - "GeneratedScalarConversionRounding::NearestAway" - } - ScalarConversionRounding::NearestEven => { - "GeneratedScalarConversionRounding::NearestEven" - } - ScalarConversionRounding::TowardZero => "GeneratedScalarConversionRounding::TowardZero", - }; - let saturation = match conversion.saturation { - ScalarConversionSaturation::None => "GeneratedScalarConversionSaturation::None", - ScalarConversionSaturation::Relu => "GeneratedScalarConversionSaturation::Relu", - ScalarConversionSaturation::Satfinite => { - "GeneratedScalarConversionSaturation::Satfinite" - } - ScalarConversionSaturation::ReluSatfinite => { - "GeneratedScalarConversionSaturation::ReluSatfinite" - } - }; - return format!( - "GeneratedIntrinsicVariant::ScalarConversion {{ rounding: {rounding}, saturation: {saturation} }}" - ); - } - let Some(ldmatrix) = &record.ldmatrix else { - return "GeneratedIntrinsicVariant::Scalar".to_owned(); - }; - let variant = &ldmatrix.variant; - let shape = match variant.shape { - LdmatrixShape::M8n8 => "GeneratedLdmatrixShape::M8n8", - LdmatrixShape::M8n16 => "GeneratedLdmatrixShape::M8n16", - LdmatrixShape::M16n16 => "GeneratedLdmatrixShape::M16n16", - }; - let multiplicity = match variant.multiplicity { - LdmatrixMultiplicity::X1 => "GeneratedLdmatrixMultiplicity::X1", - LdmatrixMultiplicity::X2 => "GeneratedLdmatrixMultiplicity::X2", - LdmatrixMultiplicity::X4 => "GeneratedLdmatrixMultiplicity::X4", - }; - let layout = match variant.layout { - LdmatrixLayout::Normal => "GeneratedLdmatrixLayout::Normal", - LdmatrixLayout::Transposed => "GeneratedLdmatrixLayout::Transposed", - }; - let element = match variant.element { - LdmatrixElement::B16 => "GeneratedLdmatrixElement::B16", - LdmatrixElement::B8 => "GeneratedLdmatrixElement::B8", - LdmatrixElement::B8x16B4x16P64 => "GeneratedLdmatrixElement::B8x16B4x16P64", - LdmatrixElement::B8x16B6x16P32 => "GeneratedLdmatrixElement::B8x16B6x16P32", - }; - format!( - "GeneratedIntrinsicVariant::Ldmatrix {{ shape: {shape}, multiplicity: {multiplicity}, layout: {layout}, element: {element} }}" - ) -} - -fn generated_backend_requirements(record: &CatalogIntrinsic) -> String { - let requirements = record - .backend_lowerings - .iter() - .map(|lowering| { - let backend = match lowering.backend { - IntrinsicBackend::LlvmNvptx => "GeneratedIntrinsicBackend::LlvmNvptx", - IntrinsicBackend::LibNvvm => "GeneratedIntrinsicBackend::LibNvvm", - }; - format!( - "GeneratedBackendRequirement {{ backend: {backend}, requirement: GeneratedTargetRequirement {{ minimum_ptx: GeneratedPtxVersion::from_encoded({}), hardware: {} }} }}", - lowering.target.minimum_ptx.encoded(), - generated_hardware_target(&lowering.target.hardware), - ) - }) - .collect::>() - .join(", "); - format!("&[{requirements}]") -} - -fn replace_exact_render_fragment(output: &mut String, fragment: &str, replacement: &str) { - assert_eq!( - output.matches(fragment).count(), - 1, - "render fragment must occur exactly once" - ); - let start = output.find(fragment).expect("checked render fragment"); - output.replace_range(start..start + fragment.len(), replacement); -} - -fn render_targets(catalog: &CatalogFile, hash: &str) -> String { - let mut records = catalog.intrinsics.iter().collect::>(); - records.sort_by(|left, right| left.rust.abi_id.cmp(&right.rust.abi_id)); - let mut output = rust_header(catalog, hash); - output.push_str( - "//! Generated target requirements and separately imported LLVM/selection facts.\n\npub const GENERATED_INTRINSIC_MARKER_ATTR: &str = \"cuda_oxide_intrinsic_marker\";\n\n#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]\npub struct GeneratedPtxVersion(u16);\nimpl GeneratedPtxVersion {\n pub const fn from_encoded(encoded: u16) -> Self { Self(encoded) }\n pub const fn encoded(self) -> u16 { self.0 }\n pub const fn major(self) -> u16 { self.0 / 10 }\n pub const fn minor(self) -> u16 { self.0 % 10 }\n}\n\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedHardwareAlternative { MinimumSm(u16), ExactArchitecture(u16), FamilyTarget(u16) }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedTargetSelectorBinding { pub name: &'static str, pub value: &'static str }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedTargetAlternative { pub minimum_ptx: GeneratedPtxVersion, pub hardware: GeneratedHardwareAlternative }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedHardwareTarget { All, AnyOf(&'static [GeneratedHardwareAlternative]), TargetMatrix { selectors: &'static [GeneratedTargetSelectorBinding], alternatives: &'static [GeneratedTargetAlternative] } }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedTargetRequirement { pub minimum_ptx: GeneratedPtxVersion, pub hardware: GeneratedHardwareTarget }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedIntrinsicBackend { LlvmNvptx, LibNvvm }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedBackendRequirement { pub backend: GeneratedIntrinsicBackend, pub requirement: GeneratedTargetRequirement }\n\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedSelectionAddressSpace { Generic, Shared }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedImmediateBinding { pub argument_index: usize, pub value: i64 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedSelectionConstraints { pub address_space: Option, pub immediate_bindings: &'static [GeneratedImmediateBinding] }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedSelectionAlternative { pub source_record: &'static str, pub asm: &'static str, pub predicates: &'static [&'static str], pub constraints: GeneratedSelectionConstraints }\n\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedIntrinsicRange { pub lower: &'static str, pub upper_exclusive: &'static str }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedLlvmFacts { pub properties: &'static [&'static str], pub result_no_undef: bool, pub result_range: Option }\n\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedLdmatrixShape { M8n8 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedLdmatrixMultiplicity { X1, X2, X4 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedLdmatrixLayout { Normal, Transposed }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedPackedAtomicFormat { F16x2, Bf16x2 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedRegisterMmaShape { M8n8k4, M16n8k8, M16n8k16, M16n8k32 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedRegisterMmaAccumulator { F32, F64, S32 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedRegisterMmaElement { Bf16, F16, Tf32, F64, S8, U8 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedRegisterMmaLayout { Row, Col }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedRegisterMmaOverflow { NotApplicable, Wrapping, Satfinite }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedIntrinsicVariant {\n Scalar,\n Ldmatrix { shape: GeneratedLdmatrixShape, multiplicity: GeneratedLdmatrixMultiplicity, layout: GeneratedLdmatrixLayout },\n PackedAtomic { format: GeneratedPackedAtomicFormat },\n RegisterMma { shape: GeneratedRegisterMmaShape, accumulator: GeneratedRegisterMmaAccumulator, a_element: GeneratedRegisterMmaElement, b_element: GeneratedRegisterMmaElement, a_layout: GeneratedRegisterMmaLayout, b_layout: GeneratedRegisterMmaLayout, overflow: GeneratedRegisterMmaOverflow },\n}\n\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedIntrinsicTarget {\n pub marker: &'static str,\n pub id: &'static str,\n pub abi_id: &'static str,\n pub dialect_op: &'static str,\n pub variant: GeneratedIntrinsicVariant,\n pub requirement: GeneratedTargetRequirement,\n pub backend_requirements: &'static [GeneratedBackendRequirement],\n pub selections: &'static [GeneratedSelectionAlternative],\n pub llvm: Option,\n}\n\nimpl GeneratedIntrinsicTarget {\n pub fn requirement_for_backend(&self, backend: GeneratedIntrinsicBackend) -> GeneratedTargetRequirement {\n self.backend_requirements.iter().find(|entry| entry.backend == backend).map(|entry| entry.requirement).unwrap_or(self.requirement)\n }\n}\n\npub const GENERATED_INTRINSIC_TARGETS: &[GeneratedIntrinsicTarget] = &[\n", - ); - replace_exact_render_fragment( - &mut output, - "#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedTargetAlternative { pub minimum_ptx: GeneratedPtxVersion, pub hardware: GeneratedHardwareAlternative }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedHardwareTarget { All, AnyOf(&'static [GeneratedHardwareAlternative]), TargetMatrix { selectors: &'static [GeneratedTargetSelectorBinding], alternatives: &'static [GeneratedTargetAlternative] } }", - "#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedTargetAlternative { pub minimum_ptx: GeneratedPtxVersion, pub hardware: GeneratedHardwareAlternative }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedTargetContract { pub selectors: &'static [GeneratedTargetSelectorBinding], pub alternatives: &'static [GeneratedTargetAlternative] }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedHardwareTarget { All, AnyOf(&'static [GeneratedHardwareAlternative]), TargetMatrix { contracts: &'static [GeneratedTargetContract] } }", - ); - replace_exact_render_fragment( - &mut output, - "impl GeneratedIntrinsicTarget {\n pub fn requirement_for_backend(&self, backend: GeneratedIntrinsicBackend) -> GeneratedTargetRequirement {\n self.backend_requirements.iter().find(|entry| entry.backend == backend).map(|entry| entry.requirement).unwrap_or(self.requirement)\n }\n}", - r#"impl GeneratedHardwareTarget { - pub fn contract_for_selector( - self, - name: &str, - value: &str, - ) -> Option<&'static GeneratedTargetContract> { - let Self::TargetMatrix { contracts } = self else { return None; }; - let mut matching = contracts.iter().filter(|contract| { - matches!(contract.selectors, [binding] if binding.name == name && binding.value == value) - }); - let contract = matching.next()?; - if matching.next().is_some() { None } else { Some(contract) } - } -} - -impl GeneratedIntrinsicTarget { - pub fn requirement_for_backend(&self, backend: GeneratedIntrinsicBackend) -> GeneratedTargetRequirement { - self.backend_requirements.iter().find(|entry| entry.backend == backend).map(|entry| entry.requirement).unwrap_or(self.requirement) - } - - pub fn target_contract_for_backend_selector( - &self, - backend: GeneratedIntrinsicBackend, - name: &str, - value: &str, - ) -> Option<&'static GeneratedTargetContract> { - self.requirement_for_backend(backend).hardware.contract_for_selector(name, value) - } -}"#, - ); - replace_exact_render_fragment( - &mut output, - "GeneratedLdmatrixShape { M8n8 }", - "GeneratedLdmatrixShape { M8n8, M8n16, M16n16 }", - ); - replace_exact_render_fragment( - &mut output, - "pub enum GeneratedPackedAtomicFormat { F16x2, Bf16x2 }", - "pub enum GeneratedLdmatrixElement { B16, B8, B8x16B4x16P64, B8x16B6x16P32 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedPackedAtomicFormat { F16x2, Bf16x2 }", - ); - replace_exact_render_fragment( - &mut output, - "Ldmatrix { shape: GeneratedLdmatrixShape, multiplicity: GeneratedLdmatrixMultiplicity, layout: GeneratedLdmatrixLayout }", - "Ldmatrix { shape: GeneratedLdmatrixShape, multiplicity: GeneratedLdmatrixMultiplicity, layout: GeneratedLdmatrixLayout, element: GeneratedLdmatrixElement }", - ); - replace_exact_render_fragment( - &mut output, - "GeneratedRegisterMmaShape { M8n8k4, M16n8k8, M16n8k16, M16n8k32 }", - "GeneratedRegisterMmaShape { M8n8k4, M8n8k16, M8n8k32, M8n8k128, M16n8k4, M16n8k8, M16n8k16, M16n8k32, M16n8k64, M16n8k128, M16n8k256 }", - ); - replace_exact_render_fragment( - &mut output, - "GeneratedRegisterMmaElement { Bf16, F16, Tf32, F64, S8, U8 }", - "GeneratedRegisterMmaElement { Bf16, E2m1, E2m3, E3m2, E4m3, E5m2, F16, Tf32, F64, B1, S4, U4, S8, U8 }", - ); - replace_exact_render_fragment( - &mut output, - "pub enum GeneratedRegisterMmaAccumulator { F32, F64, S32 }", - "pub enum GeneratedRegisterMmaOperation { Multiply, AndPopc, XorPopc }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedRegisterMmaKind { Standard, F8f6f4, Mxf8f6f4 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedRegisterMmaAccumulator { F16, F32, F64, S32 }", - ); - replace_exact_render_fragment( - &mut output, - "RegisterMma { shape: GeneratedRegisterMmaShape, accumulator: GeneratedRegisterMmaAccumulator,", - "RegisterMma { shape: GeneratedRegisterMmaShape, operation: GeneratedRegisterMmaOperation, kind: GeneratedRegisterMmaKind, accumulator: GeneratedRegisterMmaAccumulator,", - ); - replace_exact_render_fragment( - &mut output, - "#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedIntrinsicVariant {", - "#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedSparseMmaShape { M16n8k32, M16n8k64, M16n8k128 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedSparseMmaAccumulator { F16, F32, S32 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedSparseMmaElement { E2m1, E2m3, E3m2, E4m3, E5m2, S4, U4, S8, U8 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedSparseMmaLayout { Row, Col }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedSparseMmaOverflow { NotApplicable, Wrapping, Satfinite }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedSparseMmaMetadata { Standard, Ordered }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedSparseMmaSelector { ImmediateZeroOrOne, ImmediateZero }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedPrmtMode { Generic, F4e, B4e, Rc8, Ecl, Ecr, Rc16 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedClusterBarrierMode { Arrive, ArriveAligned, ArriveRelaxed, ArriveRelaxedAligned, Wait, WaitAligned }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedScalarConversionRounding { NearestAway, NearestEven, TowardZero }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedScalarConversionSaturation { None, Relu, Satfinite, ReluSatfinite }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedIntrinsicVariant {", - ); - replace_exact_render_fragment( - &mut output, - " RegisterMma { shape: GeneratedRegisterMmaShape, operation: GeneratedRegisterMmaOperation, kind: GeneratedRegisterMmaKind, accumulator: GeneratedRegisterMmaAccumulator, a_element: GeneratedRegisterMmaElement, b_element: GeneratedRegisterMmaElement, a_layout: GeneratedRegisterMmaLayout, b_layout: GeneratedRegisterMmaLayout, overflow: GeneratedRegisterMmaOverflow },\n}", - " RegisterMma { shape: GeneratedRegisterMmaShape, operation: GeneratedRegisterMmaOperation, kind: GeneratedRegisterMmaKind, accumulator: GeneratedRegisterMmaAccumulator, a_element: GeneratedRegisterMmaElement, b_element: GeneratedRegisterMmaElement, a_layout: GeneratedRegisterMmaLayout, b_layout: GeneratedRegisterMmaLayout, overflow: GeneratedRegisterMmaOverflow },\n SparseMma { shape: GeneratedSparseMmaShape, accumulator: GeneratedSparseMmaAccumulator, a_element: GeneratedSparseMmaElement, b_element: GeneratedSparseMmaElement, a_layout: GeneratedSparseMmaLayout, b_layout: GeneratedSparseMmaLayout, overflow: GeneratedSparseMmaOverflow, metadata: GeneratedSparseMmaMetadata, selector: GeneratedSparseMmaSelector },\n Prmt { mode: GeneratedPrmtMode },\n ClusterBarrier { mode: GeneratedClusterBarrierMode },\n ScalarConversion { rounding: GeneratedScalarConversionRounding, saturation: GeneratedScalarConversionSaturation },\n}", - ); - if wgmma_controls(catalog).next().is_some() { - replace_exact_render_fragment( - &mut output, - "#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedIntrinsicVariant {", - "#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedWgmmaControlMode { Fence, CommitGroup, WaitGroup }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedIntrinsicVariant {", - ); - replace_exact_render_fragment( - &mut output, - " ClusterBarrier { mode: GeneratedClusterBarrierMode },\n ScalarConversion { rounding: GeneratedScalarConversionRounding, saturation: GeneratedScalarConversionSaturation },\n}", - " ClusterBarrier { mode: GeneratedClusterBarrierMode },\n ScalarConversion { rounding: GeneratedScalarConversionRounding, saturation: GeneratedScalarConversionSaturation },\n WgmmaControl { mode: GeneratedWgmmaControlMode },\n}", - ); - } - if scalar_arithmetics(catalog).next().is_some() { - replace_exact_render_fragment( - &mut output, - "#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedIntrinsicVariant {", - "#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedScalarArithmeticFormat { F32, F64 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedScalarArithmeticOperation { Mul, Div, Fma, Add }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedScalarArithmeticRounding { Rn, Rz, Rm, Rp }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedScalarArithmeticSubnormal { Preserve, Ftz }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedScalarArithmeticSaturation { None, Sat }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedIntrinsicVariant {", - ); - if wgmma_controls(catalog).next().is_some() { - replace_exact_render_fragment( - &mut output, - " WgmmaControl { mode: GeneratedWgmmaControlMode },\n}", - " WgmmaControl { mode: GeneratedWgmmaControlMode },\n ScalarArithmetic { format: GeneratedScalarArithmeticFormat, operation: GeneratedScalarArithmeticOperation, rounding: GeneratedScalarArithmeticRounding, subnormal: GeneratedScalarArithmeticSubnormal, saturation: GeneratedScalarArithmeticSaturation },\n}", - ); - } else { - replace_exact_render_fragment( - &mut output, - " ScalarConversion { rounding: GeneratedScalarConversionRounding, saturation: GeneratedScalarConversionSaturation },\n}", - " ScalarConversion { rounding: GeneratedScalarConversionRounding, saturation: GeneratedScalarConversionSaturation },\n ScalarArithmetic { format: GeneratedScalarArithmeticFormat, operation: GeneratedScalarArithmeticOperation, rounding: GeneratedScalarArithmeticRounding, subnormal: GeneratedScalarArithmeticSubnormal, saturation: GeneratedScalarArithmeticSaturation },\n}", - ); - } - } - if extended_minmax(catalog).next().is_some() { - replace_exact_render_fragment( - &mut output, - "#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedIntrinsicVariant {", - "#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedExtendedMinMaxFormat { F32, F16, Bf16, F16x2, Bf16x2 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedExtendedMinMaxOperation { Min, Max }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedExtendedMinMaxSubnormal { Preserve, Ftz }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedExtendedMinMaxNan { Number, Nan }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedIntrinsicVariant {", - ); - replace_exact_render_fragment( - &mut output, - "}\n\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedIntrinsicTarget", - " ExtendedMinMax { format: GeneratedExtendedMinMaxFormat, operation: GeneratedExtendedMinMaxOperation, subnormal: GeneratedExtendedMinMaxSubnormal, nan: GeneratedExtendedMinMaxNan, xorsign_abs: bool },\n}\n\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedIntrinsicTarget", - ); - } - if tcgen05_mma_intrinsics(catalog).next().is_some() { - replace_exact_render_fragment( - &mut output, - "#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedIntrinsicVariant {", - "#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedTcgen05MmaForm { Shared, Tensor, TensorAshift, SpShared, SpTensor, SpTensorAshift, WsShared, WsSharedZeroColMask, WsSpShared, WsSpSharedZeroColMask, WsSpTensor, WsSpTensorZeroColMask, WsTensor, WsTensorZeroColMask }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedTcgen05MmaTargetSelector { Kind }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedIntrinsicVariant {", - ); - replace_exact_render_fragment( - &mut output, - "}\n\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedIntrinsicTarget", - " Tcgen05Mma { form: GeneratedTcgen05MmaForm, target_selector: GeneratedTcgen05MmaTargetSelector, compatibility_alias: bool },\n}\n\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedIntrinsicTarget", - ); - } - for record in records.iter().copied() { - let llvm_facts = match &record.llvm { - Some(llvm) => { - let result_range = match &llvm.result_facts.range { - Some(range) => format!( - "Some(GeneratedIntrinsicRange {{ lower: {:?}, upper_exclusive: {:?} }})", - range.lower, range.upper_exclusive - ), - None => "None".to_owned(), - }; - format!( - "Some(GeneratedLlvmFacts {{ properties: &{:?}, result_no_undef: {}, result_range: {} }})", - llvm.properties, llvm.result_facts.no_undef, result_range - ) - } - None => "None".to_owned(), - }; - writeln!( - output, - " GeneratedIntrinsicTarget {{ marker: {:?}, id: {:?}, abi_id: {:?}, dialect_op: {:?}, variant: {}, requirement: GeneratedTargetRequirement {{ minimum_ptx: GeneratedPtxVersion::from_encoded({}), hardware: {} }}, backend_requirements: {}, selections: {}, llvm: {} }},", - intrinsic_marker(catalog, record), - record.id, - record.rust.abi_id, - record.dialect.op_name, - generated_intrinsic_variant(record), - record.target.minimum_ptx.encoded(), - generated_hardware_target(&record.target.hardware), - generated_backend_requirements(record), - generated_selection_alternatives(&record.selections), - llvm_facts, - ) - .unwrap(); - } - output.push_str( - "];\n\npub fn generated_intrinsic_target_by_marker(marker: &str) -> Option<&'static GeneratedIntrinsicTarget> {\n GENERATED_INTRINSIC_TARGETS.iter().find(|target| target.marker == marker)\n}\n\npub fn generated_intrinsic_targets_by_op_name(op_name: &str) -> impl Iterator + '_ {\n GENERATED_INTRINSIC_TARGETS.iter().filter(move |target| target.dialect_op == op_name)\n}\n\npub fn generated_intrinsic_target_by_op_name(op_name: &str) -> Option<&'static GeneratedIntrinsicTarget> {\n generated_intrinsic_targets_by_op_name(op_name).next()\n}\n\npub fn generated_intrinsic_target(op_name: &str) -> Option<&'static GeneratedIntrinsicTarget> {\n generated_intrinsic_target_by_op_name(op_name)\n}\n\npub fn generated_intrinsic_operation_matches(ctx: &Context, target: &GeneratedIntrinsicTarget, operation: Ptr) -> bool {\n match target.variant {\n GeneratedIntrinsicVariant::Scalar => true,\n GeneratedIntrinsicVariant::Ldmatrix { shape, multiplicity, layout } => {\n let Some(op) = Operation::get_op::(operation, ctx) else { return false; };\n let shape_matches = matches!(shape, GeneratedLdmatrixShape::M8n8) && op.get_attr_nvvm_ldmatrix_shape(ctx).as_deref() == Some(&LdmatrixShapeAttr::M8n8);\n let multiplicity_matches = match multiplicity {\n GeneratedLdmatrixMultiplicity::X1 => op.get_attr_nvvm_ldmatrix_multiplicity(ctx).as_deref() == Some(&LdmatrixMultiplicityAttr::X1),\n GeneratedLdmatrixMultiplicity::X2 => op.get_attr_nvvm_ldmatrix_multiplicity(ctx).as_deref() == Some(&LdmatrixMultiplicityAttr::X2),\n GeneratedLdmatrixMultiplicity::X4 => op.get_attr_nvvm_ldmatrix_multiplicity(ctx).as_deref() == Some(&LdmatrixMultiplicityAttr::X4),\n };\n let layout_matches = match layout {\n GeneratedLdmatrixLayout::Normal => op.get_attr_nvvm_ldmatrix_layout(ctx).as_deref() == Some(&LdmatrixLayoutAttr::Normal),\n GeneratedLdmatrixLayout::Transposed => op.get_attr_nvvm_ldmatrix_layout(ctx).as_deref() == Some(&LdmatrixLayoutAttr::Transposed),\n };\n shape_matches && multiplicity_matches && layout_matches\n && op.get_attr_nvvm_ldmatrix_element(ctx).as_deref() == Some(&LdmatrixElementAttr::B16)\n && op.get_attr_nvvm_ldmatrix_state_space(ctx).as_deref() == Some(&LdmatrixStateSpaceAttr::Shared)\n }\n GeneratedIntrinsicVariant::PackedAtomic { format } => {\n let Some(op) = Operation::get_op::(operation, ctx) else { return false; };\n let format_matches = match format {\n GeneratedPackedAtomicFormat::F16x2 => op.get_attr_nvvm_packed_atomic_format(ctx).as_deref() == Some(&PackedAtomicFormatAttr::F16x2),\n GeneratedPackedAtomicFormat::Bf16x2 => op.get_attr_nvvm_packed_atomic_format(ctx).as_deref() == Some(&PackedAtomicFormatAttr::Bf16x2),\n };\n format_matches\n && op.get_attr_nvvm_packed_atomic_state_space(ctx).as_deref() == Some(&PackedAtomicStateSpaceAttr::Global)\n && op.get_attr_nvvm_packed_atomic_ordering(ctx).as_deref() == Some(&PackedAtomicOrderingAttr::Relaxed)\n && op.get_attr_nvvm_packed_atomic_scope(ctx).as_deref() == Some(&PackedAtomicScopeAttr::Gpu)\n && op.get_attr_nvvm_packed_atomic_rounding(ctx).as_deref() == Some(&PackedAtomicRoundingAttr::Rn)\n && op.get_attr_nvvm_packed_atomic_subnormal(ctx).as_deref() == Some(&PackedAtomicSubnormalAttr::NoFtz)\n && op.get_attr_nvvm_packed_atomic_atomicity(ctx).as_deref() == Some(&PackedAtomicAtomicityAttr::PerElement)\n }\n GeneratedIntrinsicVariant::RegisterMma { shape, accumulator, a_element, b_element, a_layout, b_layout, overflow } => {\n let Some(op) = Operation::get_op::(operation, ctx) else { return false; };\n let shape_matches = match shape {\n GeneratedRegisterMmaShape::M8n8k4 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M8n8k4),\n GeneratedRegisterMmaShape::M16n8k8 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M16n8k8),\n GeneratedRegisterMmaShape::M16n8k16 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M16n8k16),\n GeneratedRegisterMmaShape::M16n8k32 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M16n8k32),\n };\n let accumulator_matches = match accumulator {\n GeneratedRegisterMmaAccumulator::F32 => op.get_attr_nvvm_register_mma_accumulator(ctx).as_deref() == Some(&RegisterMmaAccumulatorAttr::F32),\n GeneratedRegisterMmaAccumulator::F64 => op.get_attr_nvvm_register_mma_accumulator(ctx).as_deref() == Some(&RegisterMmaAccumulatorAttr::F64),\n GeneratedRegisterMmaAccumulator::S32 => op.get_attr_nvvm_register_mma_accumulator(ctx).as_deref() == Some(&RegisterMmaAccumulatorAttr::S32),\n };\n let element_matches = |expected, actual: Option<&RegisterMmaElementAttr>| match expected {\n GeneratedRegisterMmaElement::Bf16 => actual == Some(&RegisterMmaElementAttr::Bf16),\n GeneratedRegisterMmaElement::F16 => actual == Some(&RegisterMmaElementAttr::F16),\n GeneratedRegisterMmaElement::Tf32 => actual == Some(&RegisterMmaElementAttr::Tf32),\n GeneratedRegisterMmaElement::F64 => actual == Some(&RegisterMmaElementAttr::F64),\n GeneratedRegisterMmaElement::S8 => actual == Some(&RegisterMmaElementAttr::S8),\n GeneratedRegisterMmaElement::U8 => actual == Some(&RegisterMmaElementAttr::U8),\n };\n let layout_matches = |expected, actual: Option<&RegisterMmaLayoutAttr>| match expected {\n GeneratedRegisterMmaLayout::Row => actual == Some(&RegisterMmaLayoutAttr::Row),\n GeneratedRegisterMmaLayout::Col => actual == Some(&RegisterMmaLayoutAttr::Col),\n };\n let overflow_matches = match overflow {\n GeneratedRegisterMmaOverflow::NotApplicable => op.get_attr_nvvm_register_mma_overflow(ctx).as_deref() == Some(&RegisterMmaOverflowAttr::NotApplicable),\n GeneratedRegisterMmaOverflow::Wrapping => op.get_attr_nvvm_register_mma_overflow(ctx).as_deref() == Some(&RegisterMmaOverflowAttr::Wrapping),\n GeneratedRegisterMmaOverflow::Satfinite => op.get_attr_nvvm_register_mma_overflow(ctx).as_deref() == Some(&RegisterMmaOverflowAttr::Satfinite),\n };\n shape_matches && accumulator_matches\n && element_matches(a_element, op.get_attr_nvvm_register_mma_a_element(ctx).as_deref())\n && element_matches(b_element, op.get_attr_nvvm_register_mma_b_element(ctx).as_deref())\n && layout_matches(a_layout, op.get_attr_nvvm_register_mma_a_layout(ctx).as_deref())\n && layout_matches(b_layout, op.get_attr_nvvm_register_mma_b_layout(ctx).as_deref())\n && overflow_matches\n }\n }\n}\n", - ); - if tcgen05_mma_intrinsics(catalog).next().is_some() { - replace_exact_render_fragment( - &mut output, - "pub fn generated_intrinsic_target(op_name: &str) -> Option<&'static GeneratedIntrinsicTarget> {\n generated_intrinsic_target_by_op_name(op_name)\n}\n\npub fn generated_intrinsic_operation_matches", - "pub fn generated_intrinsic_target(op_name: &str) -> Option<&'static GeneratedIntrinsicTarget> {\n generated_intrinsic_target_by_op_name(op_name)\n}\n\npub fn generated_intrinsic_target_is_direct_dialect_candidate(target: &GeneratedIntrinsicTarget) -> bool {\n !matches!(target.variant, GeneratedIntrinsicVariant::Tcgen05Mma { compatibility_alias: true, .. })\n}\n\npub fn generated_intrinsic_operation_matches", - ); - replace_exact_render_fragment( - &mut output, - " GeneratedIntrinsicVariant::Scalar => true,", - r#" GeneratedIntrinsicVariant::Scalar => true, - GeneratedIntrinsicVariant::Tcgen05Mma { form, target_selector, compatibility_alias } => { - let Some(op) = Operation::get_op::(operation, ctx) else { return false; }; - let form_matches = match form { - GeneratedTcgen05MmaForm::Shared => op.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() == Some(&Tcgen05MmaFormAttr::Shared), - GeneratedTcgen05MmaForm::Tensor => op.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() == Some(&Tcgen05MmaFormAttr::Tensor), - GeneratedTcgen05MmaForm::TensorAshift => op.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() == Some(&Tcgen05MmaFormAttr::TensorAshift), - GeneratedTcgen05MmaForm::SpShared => op.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() == Some(&Tcgen05MmaFormAttr::SpShared), - GeneratedTcgen05MmaForm::SpTensor => op.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() == Some(&Tcgen05MmaFormAttr::SpTensor), - GeneratedTcgen05MmaForm::SpTensorAshift => op.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() == Some(&Tcgen05MmaFormAttr::SpTensorAshift), - GeneratedTcgen05MmaForm::WsShared => op.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() == Some(&Tcgen05MmaFormAttr::WsShared), - GeneratedTcgen05MmaForm::WsSharedZeroColMask => op.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() == Some(&Tcgen05MmaFormAttr::WsSharedZeroColMask), - GeneratedTcgen05MmaForm::WsSpShared => op.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() == Some(&Tcgen05MmaFormAttr::WsSpShared), - GeneratedTcgen05MmaForm::WsSpSharedZeroColMask => op.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() == Some(&Tcgen05MmaFormAttr::WsSpSharedZeroColMask), - GeneratedTcgen05MmaForm::WsSpTensor => op.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() == Some(&Tcgen05MmaFormAttr::WsSpTensor), - GeneratedTcgen05MmaForm::WsSpTensorZeroColMask => op.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() == Some(&Tcgen05MmaFormAttr::WsSpTensorZeroColMask), - GeneratedTcgen05MmaForm::WsTensor => op.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() == Some(&Tcgen05MmaFormAttr::WsTensor), - GeneratedTcgen05MmaForm::WsTensorZeroColMask => op.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() == Some(&Tcgen05MmaFormAttr::WsTensorZeroColMask), - }; - let selector_matches = matches!(target_selector, GeneratedTcgen05MmaTargetSelector::Kind) - && op.get_attr_nvvm_tcgen05_mma_kind(ctx).is_some(); - let alias_matches = !compatibility_alias || ( - op.get_attr_nvvm_tcgen05_mma_kind(ctx).as_deref() == Some(&Tcgen05MmaKindAttr::F8f6f4) - && match form { - GeneratedTcgen05MmaForm::Shared => { - op.get_attr_nvvm_tcgen05_mma_cta_group(ctx).as_deref() == Some(&Tcgen05MmaCtaGroupAttr::Cg1) - && op.get_attr_nvvm_tcgen05_mma_collector_a(ctx).as_deref() == Some(&Tcgen05MmaCollectorAAttr::Discard) - } - GeneratedTcgen05MmaForm::WsTensor => { - op.get_attr_nvvm_tcgen05_mma_b_buffer(ctx).as_deref() == Some(&Tcgen05MmaBBufferAttr::B0) - && op.get_attr_nvvm_tcgen05_mma_b_usage(ctx).as_deref() == Some(&Tcgen05MmaBUsageAttr::Discard) - } - _ => false, - } - ); - form_matches && selector_matches && alias_matches - }"#, - ); - } - replace_exact_render_fragment( - &mut output, - "GeneratedIntrinsicVariant::Ldmatrix { shape, multiplicity, layout } => {", - "GeneratedIntrinsicVariant::Ldmatrix { shape, multiplicity, layout, element } => {", - ); - replace_exact_render_fragment( - &mut output, - " let shape_matches = matches!(shape, GeneratedLdmatrixShape::M8n8) && op.get_attr_nvvm_ldmatrix_shape(ctx).as_deref() == Some(&LdmatrixShapeAttr::M8n8);", - " let shape_matches = match shape {\n GeneratedLdmatrixShape::M8n8 => op.get_attr_nvvm_ldmatrix_shape(ctx).as_deref() == Some(&LdmatrixShapeAttr::M8n8),\n GeneratedLdmatrixShape::M8n16 => op.get_attr_nvvm_ldmatrix_shape(ctx).as_deref() == Some(&LdmatrixShapeAttr::M8n16),\n GeneratedLdmatrixShape::M16n16 => op.get_attr_nvvm_ldmatrix_shape(ctx).as_deref() == Some(&LdmatrixShapeAttr::M16n16),\n };", - ); - replace_exact_render_fragment( - &mut output, - " shape_matches && multiplicity_matches && layout_matches\n && op.get_attr_nvvm_ldmatrix_element(ctx).as_deref() == Some(&LdmatrixElementAttr::B16)", - " let element_matches = match element {\n GeneratedLdmatrixElement::B16 => op.get_attr_nvvm_ldmatrix_element(ctx).as_deref() == Some(&LdmatrixElementAttr::B16),\n GeneratedLdmatrixElement::B8 => op.get_attr_nvvm_ldmatrix_element(ctx).as_deref() == Some(&LdmatrixElementAttr::B8),\n GeneratedLdmatrixElement::B8x16B4x16P64 => op.get_attr_nvvm_ldmatrix_element(ctx).as_deref() == Some(&LdmatrixElementAttr::B8x16B4x16P64),\n GeneratedLdmatrixElement::B8x16B6x16P32 => op.get_attr_nvvm_ldmatrix_element(ctx).as_deref() == Some(&LdmatrixElementAttr::B8x16B6x16P32),\n };\n shape_matches && multiplicity_matches && layout_matches && element_matches", - ); - let mut compatibility_aliases = String::from( - "fn generated_intrinsic_compatibility_marker_by_op_name(op_name: &str) -> Option<&'static str> {\n match op_name {\n", - ); - let mut compatibility_count = 0; - for record in ldmatrix(catalog) { - let Some((_, op_name)) = ldmatrix_compat_op(record) else { - continue; - }; - compatibility_count += 1; - writeln!( - compatibility_aliases, - " {op_name:?} => Some({:?}),", - intrinsic_marker(catalog, record) - ) - .unwrap(); - } - assert_eq!(compatibility_count, 6); - compatibility_aliases.push_str(" _ => None,\n }\n}\n\n"); - replace_exact_render_fragment( - &mut output, - "pub fn generated_intrinsic_target_by_marker(marker: &str)", - &format!( - "{compatibility_aliases}pub fn generated_intrinsic_target_by_marker(marker: &str)" - ), - ); - replace_exact_render_fragment( - &mut output, - " GENERATED_INTRINSIC_TARGETS.iter().filter(move |target| target.dialect_op == op_name)\n", - " let compatibility_marker = generated_intrinsic_compatibility_marker_by_op_name(op_name);\n GENERATED_INTRINSIC_TARGETS.iter().filter(move |target| {\n target.dialect_op == op_name || compatibility_marker == Some(target.marker)\n })\n", - ); - replace_exact_render_fragment( - &mut output, - " GeneratedIntrinsicVariant::Ldmatrix { shape, multiplicity, layout, element } => {\n let Some(op) = Operation::get_op::(operation, ctx) else { return false; };", - " GeneratedIntrinsicVariant::Ldmatrix { shape, multiplicity, layout, element } => {\n let operation_name = Operation::get_opid(operation, ctx).to_string();\n if let Some(marker) = generated_intrinsic_compatibility_marker_by_op_name(&operation_name) {\n return marker == target.marker;\n }\n let Some(op) = Operation::get_op::(operation, ctx) else { return false; };", - ); - replace_exact_render_fragment( - &mut output, - "GeneratedIntrinsicVariant::RegisterMma { shape, accumulator,", - "GeneratedIntrinsicVariant::RegisterMma { shape, operation: mma_operation, kind, accumulator,", - ); - replace_exact_render_fragment( - &mut output, - " };\n let accumulator_matches = match accumulator {", - " };\n let operation_matches = match mma_operation {\n GeneratedRegisterMmaOperation::Multiply => op.operation_or_multiply(ctx) == RegisterMmaOperationAttr::Multiply,\n GeneratedRegisterMmaOperation::AndPopc => op.operation_or_multiply(ctx) == RegisterMmaOperationAttr::AndPopc,\n GeneratedRegisterMmaOperation::XorPopc => op.operation_or_multiply(ctx) == RegisterMmaOperationAttr::XorPopc,\n };\n let kind_matches = match kind {\n GeneratedRegisterMmaKind::Standard => op.kind_or_inferred(ctx) == RegisterMmaKindAttr::Standard,\n GeneratedRegisterMmaKind::F8f6f4 => op.kind_or_inferred(ctx) == RegisterMmaKindAttr::F8f6f4,\n GeneratedRegisterMmaKind::Mxf8f6f4 => op.kind_or_inferred(ctx) == RegisterMmaKindAttr::Mxf8f6f4,\n };\n let accumulator_matches = match accumulator {", - ); - replace_exact_render_fragment( - &mut output, - " let accumulator_matches = match accumulator {\n GeneratedRegisterMmaAccumulator::F32", - " let accumulator_matches = match accumulator {\n GeneratedRegisterMmaAccumulator::F16 => op.get_attr_nvvm_register_mma_accumulator(ctx).as_deref() == Some(&RegisterMmaAccumulatorAttr::F16),\n GeneratedRegisterMmaAccumulator::F32", - ); - replace_exact_render_fragment( - &mut output, - " shape_matches && accumulator_matches", - " shape_matches && operation_matches && kind_matches && accumulator_matches", - ); - replace_exact_render_fragment( - &mut output, - "GeneratedRegisterMmaShape::M8n8k4 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M8n8k4),\n GeneratedRegisterMmaShape::M16n8k8", - "GeneratedRegisterMmaShape::M8n8k4 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M8n8k4),\n GeneratedRegisterMmaShape::M8n8k16 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M8n8k16),\n GeneratedRegisterMmaShape::M8n8k32 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M8n8k32),\n GeneratedRegisterMmaShape::M8n8k128 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M8n8k128),\n GeneratedRegisterMmaShape::M16n8k4 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M16n8k4),\n GeneratedRegisterMmaShape::M16n8k8", - ); - replace_exact_render_fragment( - &mut output, - "GeneratedRegisterMmaShape::M16n8k32 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M16n8k32),\n };", - "GeneratedRegisterMmaShape::M16n8k32 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M16n8k32),\n GeneratedRegisterMmaShape::M16n8k64 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M16n8k64),\n GeneratedRegisterMmaShape::M16n8k128 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M16n8k128),\n GeneratedRegisterMmaShape::M16n8k256 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M16n8k256),\n };", - ); - replace_exact_render_fragment( - &mut output, - "GeneratedRegisterMmaElement::F64 => actual == Some(&RegisterMmaElementAttr::F64),\n GeneratedRegisterMmaElement::S8", - "GeneratedRegisterMmaElement::F64 => actual == Some(&RegisterMmaElementAttr::F64),\n GeneratedRegisterMmaElement::E2m1 => actual == Some(&RegisterMmaElementAttr::E2m1),\n GeneratedRegisterMmaElement::E2m3 => actual == Some(&RegisterMmaElementAttr::E2m3),\n GeneratedRegisterMmaElement::E3m2 => actual == Some(&RegisterMmaElementAttr::E3m2),\n GeneratedRegisterMmaElement::E4m3 => actual == Some(&RegisterMmaElementAttr::E4m3),\n GeneratedRegisterMmaElement::E5m2 => actual == Some(&RegisterMmaElementAttr::E5m2),\n GeneratedRegisterMmaElement::B1 => actual == Some(&RegisterMmaElementAttr::B1),\n GeneratedRegisterMmaElement::S4 => actual == Some(&RegisterMmaElementAttr::S4),\n GeneratedRegisterMmaElement::U4 => actual == Some(&RegisterMmaElementAttr::U4),\n GeneratedRegisterMmaElement::S8", - ); - replace_exact_render_fragment( - &mut output, - " && overflow_matches\n }\n }\n}\n", - " && overflow_matches\n }\n GeneratedIntrinsicVariant::SparseMma { shape, accumulator, a_element, b_element, a_layout, b_layout, overflow, metadata, selector } => {\n let Some(op) = Operation::get_op::(operation, ctx) else { return false; };\n let element_matches = |expected, actual: Option<&SparseMmaElementAttr>| match expected {\n GeneratedSparseMmaElement::E2m1 => actual == Some(&SparseMmaElementAttr::E2m1),\n GeneratedSparseMmaElement::E2m3 => actual == Some(&SparseMmaElementAttr::E2m3),\n GeneratedSparseMmaElement::E3m2 => actual == Some(&SparseMmaElementAttr::E3m2),\n GeneratedSparseMmaElement::E4m3 => actual == Some(&SparseMmaElementAttr::E4m3),\n GeneratedSparseMmaElement::E5m2 => actual == Some(&SparseMmaElementAttr::E5m2),\n GeneratedSparseMmaElement::S4 => actual == Some(&SparseMmaElementAttr::S4),\n GeneratedSparseMmaElement::U4 => actual == Some(&SparseMmaElementAttr::U4),\n GeneratedSparseMmaElement::S8 => actual == Some(&SparseMmaElementAttr::S8),\n GeneratedSparseMmaElement::U8 => actual == Some(&SparseMmaElementAttr::U8),\n };\n let layout_matches = |expected, actual: Option<&SparseMmaLayoutAttr>| match expected {\n GeneratedSparseMmaLayout::Row => actual == Some(&SparseMmaLayoutAttr::Row),\n GeneratedSparseMmaLayout::Col => actual == Some(&SparseMmaLayoutAttr::Col),\n };\n let overflow_matches = match overflow {\n GeneratedSparseMmaOverflow::NotApplicable => op.get_attr_nvvm_sparse_mma_overflow(ctx).as_deref() == Some(&SparseMmaOverflowAttr::NotApplicable),\n GeneratedSparseMmaOverflow::Wrapping => op.get_attr_nvvm_sparse_mma_overflow(ctx).as_deref() == Some(&SparseMmaOverflowAttr::Wrapping),\n GeneratedSparseMmaOverflow::Satfinite => op.get_attr_nvvm_sparse_mma_overflow(ctx).as_deref() == Some(&SparseMmaOverflowAttr::Satfinite),\n };\n let metadata_matches = match metadata {\n GeneratedSparseMmaMetadata::Standard => op.get_attr_nvvm_sparse_mma_metadata(ctx).as_deref() == Some(&SparseMmaMetadataAttr::Standard),\n GeneratedSparseMmaMetadata::Ordered => op.get_attr_nvvm_sparse_mma_metadata(ctx).as_deref() == Some(&SparseMmaMetadataAttr::Ordered),\n };\n matches!(shape, GeneratedSparseMmaShape::M16n8k32)\n && op.get_attr_nvvm_sparse_mma_shape(ctx).as_deref() == Some(&SparseMmaShapeAttr::M16n8k32)\n && matches!(accumulator, GeneratedSparseMmaAccumulator::S32)\n && op.get_attr_nvvm_sparse_mma_accumulator(ctx).as_deref() == Some(&SparseMmaAccumulatorAttr::S32)\n && element_matches(a_element, op.get_attr_nvvm_sparse_mma_a_element(ctx).as_deref())\n && element_matches(b_element, op.get_attr_nvvm_sparse_mma_b_element(ctx).as_deref())\n && layout_matches(a_layout, op.get_attr_nvvm_sparse_mma_a_layout(ctx).as_deref())\n && layout_matches(b_layout, op.get_attr_nvvm_sparse_mma_b_layout(ctx).as_deref())\n && overflow_matches\n && metadata_matches\n && matches!(selector, GeneratedSparseMmaSelector::ImmediateZeroOrOne)\n && op.get_attr_nvvm_sparse_mma_selector(ctx).as_deref() == Some(&SparseMmaSelectorAttr::ImmediateZeroOrOne)\n }\n }\n}\n", - ); - replace_exact_render_fragment( - &mut output, - " matches!(shape, GeneratedSparseMmaShape::M16n8k32)\n && op.get_attr_nvvm_sparse_mma_shape(ctx).as_deref() == Some(&SparseMmaShapeAttr::M16n8k32)\n && matches!(accumulator, GeneratedSparseMmaAccumulator::S32)\n && op.get_attr_nvvm_sparse_mma_accumulator(ctx).as_deref() == Some(&SparseMmaAccumulatorAttr::S32)", - " let shape_matches = match shape {\n GeneratedSparseMmaShape::M16n8k32 => op.get_attr_nvvm_sparse_mma_shape(ctx).as_deref() == Some(&SparseMmaShapeAttr::M16n8k32),\n GeneratedSparseMmaShape::M16n8k64 => op.get_attr_nvvm_sparse_mma_shape(ctx).as_deref() == Some(&SparseMmaShapeAttr::M16n8k64),\n GeneratedSparseMmaShape::M16n8k128 => op.get_attr_nvvm_sparse_mma_shape(ctx).as_deref() == Some(&SparseMmaShapeAttr::M16n8k128),\n };\n let accumulator_matches = match accumulator {\n GeneratedSparseMmaAccumulator::F16 => op.get_attr_nvvm_sparse_mma_accumulator(ctx).as_deref() == Some(&SparseMmaAccumulatorAttr::F16),\n GeneratedSparseMmaAccumulator::F32 => op.get_attr_nvvm_sparse_mma_accumulator(ctx).as_deref() == Some(&SparseMmaAccumulatorAttr::F32),\n GeneratedSparseMmaAccumulator::S32 => op.get_attr_nvvm_sparse_mma_accumulator(ctx).as_deref() == Some(&SparseMmaAccumulatorAttr::S32),\n };\n let selector_matches = match selector {\n GeneratedSparseMmaSelector::ImmediateZeroOrOne => op.get_attr_nvvm_sparse_mma_selector(ctx).as_deref() == Some(&SparseMmaSelectorAttr::ImmediateZeroOrOne),\n GeneratedSparseMmaSelector::ImmediateZero => op.get_attr_nvvm_sparse_mma_selector(ctx).as_deref() == Some(&SparseMmaSelectorAttr::ImmediateZero),\n };\n shape_matches\n && accumulator_matches", - ); - replace_exact_render_fragment( - &mut output, - " && metadata_matches\n && matches!(selector, GeneratedSparseMmaSelector::ImmediateZeroOrOne)\n && op.get_attr_nvvm_sparse_mma_selector(ctx).as_deref() == Some(&SparseMmaSelectorAttr::ImmediateZeroOrOne)", - " && metadata_matches\n && selector_matches", - ); - replace_exact_render_fragment( - &mut output, - " }\n }\n}\n", - " }\n GeneratedIntrinsicVariant::Prmt { mode } => {\n let Some(op) = Operation::get_op::(operation, ctx) else { return false; };\n match mode {\n GeneratedPrmtMode::Generic => op.get_attr_nvvm_prmt_mode(ctx).as_deref() == Some(&PrmtModeAttr::Generic),\n GeneratedPrmtMode::F4e => op.get_attr_nvvm_prmt_mode(ctx).as_deref() == Some(&PrmtModeAttr::F4e),\n GeneratedPrmtMode::B4e => op.get_attr_nvvm_prmt_mode(ctx).as_deref() == Some(&PrmtModeAttr::B4e),\n GeneratedPrmtMode::Rc8 => op.get_attr_nvvm_prmt_mode(ctx).as_deref() == Some(&PrmtModeAttr::Rc8),\n GeneratedPrmtMode::Ecl => op.get_attr_nvvm_prmt_mode(ctx).as_deref() == Some(&PrmtModeAttr::Ecl),\n GeneratedPrmtMode::Ecr => op.get_attr_nvvm_prmt_mode(ctx).as_deref() == Some(&PrmtModeAttr::Ecr),\n GeneratedPrmtMode::Rc16 => op.get_attr_nvvm_prmt_mode(ctx).as_deref() == Some(&PrmtModeAttr::Rc16),\n }\n }\n }\n}\n", - ); - replace_exact_render_fragment( - &mut output, - " }\n }\n}\n", - " }\n GeneratedIntrinsicVariant::ClusterBarrier { mode } => {\n let Some(op) = Operation::get_op::(operation, ctx) else { return false; };\n match mode {\n GeneratedClusterBarrierMode::Arrive => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::Arrive),\n GeneratedClusterBarrierMode::ArriveAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveAligned),\n GeneratedClusterBarrierMode::ArriveRelaxed => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveRelaxed),\n GeneratedClusterBarrierMode::ArriveRelaxedAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveRelaxedAligned),\n GeneratedClusterBarrierMode::Wait => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::Wait),\n GeneratedClusterBarrierMode::WaitAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::WaitAligned),\n }\n }\n }\n}\n", - ); - if wgmma_controls(catalog).next().is_some() { - replace_exact_render_fragment( - &mut output, - " GeneratedIntrinsicVariant::ClusterBarrier { mode } => {\n let Some(op) = Operation::get_op::(operation, ctx) else { return false; };\n match mode {\n GeneratedClusterBarrierMode::Arrive => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::Arrive),\n GeneratedClusterBarrierMode::ArriveAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveAligned),\n GeneratedClusterBarrierMode::ArriveRelaxed => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveRelaxed),\n GeneratedClusterBarrierMode::ArriveRelaxedAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveRelaxedAligned),\n GeneratedClusterBarrierMode::Wait => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::Wait),\n GeneratedClusterBarrierMode::WaitAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::WaitAligned),\n }\n }\n }\n}\n", - " GeneratedIntrinsicVariant::ClusterBarrier { mode } => {\n let Some(op) = Operation::get_op::(operation, ctx) else { return false; };\n match mode {\n GeneratedClusterBarrierMode::Arrive => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::Arrive),\n GeneratedClusterBarrierMode::ArriveAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveAligned),\n GeneratedClusterBarrierMode::ArriveRelaxed => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveRelaxed),\n GeneratedClusterBarrierMode::ArriveRelaxedAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveRelaxedAligned),\n GeneratedClusterBarrierMode::Wait => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::Wait),\n GeneratedClusterBarrierMode::WaitAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::WaitAligned),\n }\n }\n GeneratedIntrinsicVariant::WgmmaControl { mode } => match mode {\n GeneratedWgmmaControlMode::Fence => Operation::get_op::(operation, ctx).is_some(),\n GeneratedWgmmaControlMode::CommitGroup => Operation::get_op::(operation, ctx).is_some(),\n GeneratedWgmmaControlMode::WaitGroup => Operation::get_op::(operation, ctx).is_some(),\n },\n }\n}\n", - ); - } - if scalar_conversions(catalog).next().is_some() { - let scalar_match = " GeneratedIntrinsicVariant::ScalarConversion { rounding, saturation } => {\n let Some(op) = Operation::get_op::(operation, ctx) else { return false; };\n let rounding_matches = match rounding {\n GeneratedScalarConversionRounding::NearestAway => op.get_attr_nvvm_scalar_conversion_rounding(ctx).as_deref() == Some(&ScalarConversionRoundingAttr::NearestAway),\n GeneratedScalarConversionRounding::NearestEven => op.get_attr_nvvm_scalar_conversion_rounding(ctx).as_deref() == Some(&ScalarConversionRoundingAttr::NearestEven),\n GeneratedScalarConversionRounding::TowardZero => op.get_attr_nvvm_scalar_conversion_rounding(ctx).as_deref() == Some(&ScalarConversionRoundingAttr::TowardZero),\n };\n let saturation_matches = match saturation {\n GeneratedScalarConversionSaturation::None => op.get_attr_nvvm_scalar_conversion_saturation(ctx).as_deref() == Some(&ScalarConversionSaturationAttr::None),\n GeneratedScalarConversionSaturation::Relu => op.get_attr_nvvm_scalar_conversion_saturation(ctx).as_deref() == Some(&ScalarConversionSaturationAttr::Relu),\n GeneratedScalarConversionSaturation::Satfinite => op.get_attr_nvvm_scalar_conversion_saturation(ctx).as_deref() == Some(&ScalarConversionSaturationAttr::Satfinite),\n GeneratedScalarConversionSaturation::ReluSatfinite => op.get_attr_nvvm_scalar_conversion_saturation(ctx).as_deref() == Some(&ScalarConversionSaturationAttr::ReluSatfinite),\n };\n rounding_matches && saturation_matches\n }\n"; - if wgmma_controls(catalog).next().is_some() { - replace_exact_render_fragment( - &mut output, - " GeneratedIntrinsicVariant::WgmmaControl { mode } => match mode {\n GeneratedWgmmaControlMode::Fence => Operation::get_op::(operation, ctx).is_some(),\n GeneratedWgmmaControlMode::CommitGroup => Operation::get_op::(operation, ctx).is_some(),\n GeneratedWgmmaControlMode::WaitGroup => Operation::get_op::(operation, ctx).is_some(),\n },\n }\n}\n", - &format!( - " GeneratedIntrinsicVariant::WgmmaControl {{ mode }} => match mode {{\n GeneratedWgmmaControlMode::Fence => Operation::get_op::(operation, ctx).is_some(),\n GeneratedWgmmaControlMode::CommitGroup => Operation::get_op::(operation, ctx).is_some(),\n GeneratedWgmmaControlMode::WaitGroup => Operation::get_op::(operation, ctx).is_some(),\n }},\n{scalar_match} }}\n}}\n" - ), - ); - } else { - replace_exact_render_fragment( - &mut output, - " GeneratedIntrinsicVariant::ClusterBarrier { mode } => {\n let Some(op) = Operation::get_op::(operation, ctx) else { return false; };\n match mode {\n GeneratedClusterBarrierMode::Arrive => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::Arrive),\n GeneratedClusterBarrierMode::ArriveAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveAligned),\n GeneratedClusterBarrierMode::ArriveRelaxed => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveRelaxed),\n GeneratedClusterBarrierMode::ArriveRelaxedAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveRelaxedAligned),\n GeneratedClusterBarrierMode::Wait => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::Wait),\n GeneratedClusterBarrierMode::WaitAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::WaitAligned),\n }\n }\n }\n}\n", - &format!( - " GeneratedIntrinsicVariant::ClusterBarrier {{ mode }} => {{\n let Some(op) = Operation::get_op::(operation, ctx) else {{ return false; }};\n match mode {{\n GeneratedClusterBarrierMode::Arrive => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::Arrive),\n GeneratedClusterBarrierMode::ArriveAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveAligned),\n GeneratedClusterBarrierMode::ArriveRelaxed => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveRelaxed),\n GeneratedClusterBarrierMode::ArriveRelaxedAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveRelaxedAligned),\n GeneratedClusterBarrierMode::Wait => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::Wait),\n GeneratedClusterBarrierMode::WaitAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::WaitAligned),\n }}\n }}\n{scalar_match} }}\n}}\n" - ), - ); - } - } - if scalar_arithmetics(catalog).next().is_some() { - let arithmetic_match = r#" GeneratedIntrinsicVariant::ScalarArithmetic { format, operation: arithmetic_operation, rounding, subnormal, saturation } => { - let Some(op) = Operation::get_op::(operation, ctx) else { return false; }; - let format_matches = match format { - GeneratedScalarArithmeticFormat::F32 => op.get_attr_nvvm_scalar_arithmetic_format(ctx).as_deref() == Some(&ScalarArithmeticFormatAttr::F32), - GeneratedScalarArithmeticFormat::F64 => op.get_attr_nvvm_scalar_arithmetic_format(ctx).as_deref() == Some(&ScalarArithmeticFormatAttr::F64), - }; - let operation_matches = match arithmetic_operation { - GeneratedScalarArithmeticOperation::Mul => op.get_attr_nvvm_scalar_arithmetic_operation(ctx).as_deref() == Some(&ScalarArithmeticOperationAttr::Mul), - GeneratedScalarArithmeticOperation::Div => op.get_attr_nvvm_scalar_arithmetic_operation(ctx).as_deref() == Some(&ScalarArithmeticOperationAttr::Div), - GeneratedScalarArithmeticOperation::Fma => op.get_attr_nvvm_scalar_arithmetic_operation(ctx).as_deref() == Some(&ScalarArithmeticOperationAttr::Fma), - GeneratedScalarArithmeticOperation::Add => op.get_attr_nvvm_scalar_arithmetic_operation(ctx).as_deref() == Some(&ScalarArithmeticOperationAttr::Add), - }; - let rounding_matches = match rounding { - GeneratedScalarArithmeticRounding::Rn => op.get_attr_nvvm_scalar_arithmetic_rounding(ctx).as_deref() == Some(&ScalarArithmeticRoundingAttr::Rn), - GeneratedScalarArithmeticRounding::Rz => op.get_attr_nvvm_scalar_arithmetic_rounding(ctx).as_deref() == Some(&ScalarArithmeticRoundingAttr::Rz), - GeneratedScalarArithmeticRounding::Rm => op.get_attr_nvvm_scalar_arithmetic_rounding(ctx).as_deref() == Some(&ScalarArithmeticRoundingAttr::Rm), - GeneratedScalarArithmeticRounding::Rp => op.get_attr_nvvm_scalar_arithmetic_rounding(ctx).as_deref() == Some(&ScalarArithmeticRoundingAttr::Rp), - }; - let subnormal_matches = match subnormal { - GeneratedScalarArithmeticSubnormal::Preserve => op.get_attr_nvvm_scalar_arithmetic_subnormal(ctx).as_deref() == Some(&ScalarArithmeticSubnormalAttr::Preserve), - GeneratedScalarArithmeticSubnormal::Ftz => op.get_attr_nvvm_scalar_arithmetic_subnormal(ctx).as_deref() == Some(&ScalarArithmeticSubnormalAttr::Ftz), - }; - let saturation_matches = match saturation { - GeneratedScalarArithmeticSaturation::None => op.get_attr_nvvm_scalar_arithmetic_saturation(ctx).as_deref() == Some(&ScalarArithmeticSaturationAttr::None), - GeneratedScalarArithmeticSaturation::Sat => op.get_attr_nvvm_scalar_arithmetic_saturation(ctx).as_deref() == Some(&ScalarArithmeticSaturationAttr::Sat), - }; - format_matches && operation_matches && rounding_matches && subnormal_matches && saturation_matches - } -"#; - replace_exact_render_fragment( - &mut output, - " rounding_matches && saturation_matches\n }\n }\n}\n", - &format!( - " rounding_matches && saturation_matches\n }}\n{arithmetic_match} }}\n}}\n" - ), - ); - } - if extended_minmax(catalog).next().is_some() { - let minmax_match = r#" GeneratedIntrinsicVariant::ExtendedMinMax { format, operation: minmax_operation, subnormal, nan, xorsign_abs } => { - let Some(op) = Operation::get_op::(operation, ctx) else { return false; }; - let format_matches = match format { - GeneratedExtendedMinMaxFormat::F32 => op.get_attr_nvvm_extended_minmax_format(ctx).as_deref() == Some(&ExtendedMinMaxFormatAttr::F32), - GeneratedExtendedMinMaxFormat::F16 => op.get_attr_nvvm_extended_minmax_format(ctx).as_deref() == Some(&ExtendedMinMaxFormatAttr::F16), - GeneratedExtendedMinMaxFormat::Bf16 => op.get_attr_nvvm_extended_minmax_format(ctx).as_deref() == Some(&ExtendedMinMaxFormatAttr::Bf16), - GeneratedExtendedMinMaxFormat::F16x2 => op.get_attr_nvvm_extended_minmax_format(ctx).as_deref() == Some(&ExtendedMinMaxFormatAttr::F16x2), - GeneratedExtendedMinMaxFormat::Bf16x2 => op.get_attr_nvvm_extended_minmax_format(ctx).as_deref() == Some(&ExtendedMinMaxFormatAttr::Bf16x2), - }; - let operation_matches = match minmax_operation { - GeneratedExtendedMinMaxOperation::Min => op.get_attr_nvvm_extended_minmax_operation(ctx).as_deref() == Some(&ExtendedMinMaxOperationAttr::Min), - GeneratedExtendedMinMaxOperation::Max => op.get_attr_nvvm_extended_minmax_operation(ctx).as_deref() == Some(&ExtendedMinMaxOperationAttr::Max), - }; - let subnormal_matches = match subnormal { - GeneratedExtendedMinMaxSubnormal::Preserve => op.get_attr_nvvm_extended_minmax_subnormal(ctx).as_deref() == Some(&ExtendedMinMaxSubnormalAttr::Preserve), - GeneratedExtendedMinMaxSubnormal::Ftz => op.get_attr_nvvm_extended_minmax_subnormal(ctx).as_deref() == Some(&ExtendedMinMaxSubnormalAttr::Ftz), - }; - let nan_matches = match nan { - GeneratedExtendedMinMaxNan::Number => op.get_attr_nvvm_extended_minmax_nan(ctx).as_deref() == Some(&ExtendedMinMaxNanAttr::Number), - GeneratedExtendedMinMaxNan::Nan => op.get_attr_nvvm_extended_minmax_nan(ctx).as_deref() == Some(&ExtendedMinMaxNanAttr::Nan), - }; - let xorsign_abs_matches = match xorsign_abs { - false => op.get_attr_nvvm_extended_minmax_xorsign_abs(ctx).as_deref() == Some(&ExtendedMinMaxXorSignAbsAttr::Disabled), - true => op.get_attr_nvvm_extended_minmax_xorsign_abs(ctx).as_deref() == Some(&ExtendedMinMaxXorSignAbsAttr::Enabled), - }; - format_matches && operation_matches && subnormal_matches && nan_matches && xorsign_abs_matches - } -"#; - let match_end = output - .rfind(" }\n}\n") - .expect("generated target matcher terminator"); - output.insert_str(match_end, minmax_match); - } - output.push_str( - "\nuse dialect_nvvm::ops::{ClusterBarrierModeAttr, ClusterBarrierOp, LdmatrixElementAttr, LdmatrixLayoutAttr, LdmatrixMultiplicityAttr, LdmatrixOp, LdmatrixShapeAttr, LdmatrixStateSpaceAttr, PackedAtomicAddOp, PackedAtomicAtomicityAttr, PackedAtomicFormatAttr, PackedAtomicOrderingAttr, PackedAtomicRoundingAttr, PackedAtomicScopeAttr, PackedAtomicStateSpaceAttr, PackedAtomicSubnormalAttr, PrmtModeAttr, PrmtOp, RegisterMmaAccumulatorAttr, RegisterMmaElementAttr, RegisterMmaKindAttr, RegisterMmaLayoutAttr, RegisterMmaOp, RegisterMmaOperationAttr, RegisterMmaOverflowAttr, RegisterMmaShapeAttr, SparseMmaAccumulatorAttr, SparseMmaElementAttr, SparseMmaLayoutAttr, SparseMmaMetadataAttr, SparseMmaOp, SparseMmaOverflowAttr, SparseMmaSelectorAttr, SparseMmaShapeAttr};\nuse pliron::{context::{Context, Ptr}, operation::Operation};\n", - ); - if wgmma_controls(catalog).next().is_some() { - replace_exact_render_fragment( - &mut output, - "SparseMmaSelectorAttr, SparseMmaShapeAttr};", - "SparseMmaSelectorAttr, SparseMmaShapeAttr, WgmmaCommitGroupSyncAlignedOp, WgmmaFenceSyncAlignedOp, WgmmaWaitGroupSyncAlignedOp};", - ); - } - if scalar_conversions(catalog).next().is_some() { - replace_exact_render_fragment( - &mut output, - "SparseMmaSelectorAttr, SparseMmaShapeAttr", - "ScalarConversionOp, ScalarConversionRoundingAttr, ScalarConversionSaturationAttr, SparseMmaSelectorAttr, SparseMmaShapeAttr", - ); - } - if scalar_arithmetics(catalog).next().is_some() { - replace_exact_render_fragment( - &mut output, - "SparseMmaSelectorAttr, SparseMmaShapeAttr", - "ScalarArithmeticFormatAttr, ScalarArithmeticOp, ScalarArithmeticOperationAttr, ScalarArithmeticRoundingAttr, ScalarArithmeticSaturationAttr, ScalarArithmeticSubnormalAttr, SparseMmaSelectorAttr, SparseMmaShapeAttr", - ); - } - if extended_minmax(catalog).next().is_some() { - replace_exact_render_fragment( - &mut output, - "SparseMmaSelectorAttr, SparseMmaShapeAttr", - "ExtendedMinMaxFormatAttr, ExtendedMinMaxNanAttr, ExtendedMinMaxOp, ExtendedMinMaxOperationAttr, ExtendedMinMaxSubnormalAttr, ExtendedMinMaxXorSignAbsAttr, SparseMmaSelectorAttr, SparseMmaShapeAttr", - ); - } - if tcgen05_mma_intrinsics(catalog).next().is_some() { - replace_exact_render_fragment( - &mut output, - "SparseMmaSelectorAttr, SparseMmaShapeAttr", - "SparseMmaSelectorAttr, SparseMmaShapeAttr, Tcgen05MmaBBufferAttr, Tcgen05MmaBUsageAttr, Tcgen05MmaCollectorAAttr, Tcgen05MmaCtaGroupAttr, Tcgen05MmaFormAttr, Tcgen05MmaKindAttr, Tcgen05MmaOp", - ); - } - output.push_str( - "\n#[cfg(test)]\nmod tests {\n use super::*;\n use std::collections::BTreeSet;\n\n #[test]\n fn generated_target_table_is_unique_and_lookup_is_complete() {\n let mut ids = BTreeSet::new();\n let mut markers = BTreeSet::new();\n let mut previous_abi_id = None;\n for target in GENERATED_INTRINSIC_TARGETS {\n if let Some(previous) = previous_abi_id {\n assert!(previous < target.abi_id, \"generated ABI IDs are not strictly increasing: {previous} then {}\", target.abi_id);\n }\n previous_abi_id = Some(target.abi_id);\n assert!(ids.insert(target.id), \"duplicate generated intrinsic ID {}\", target.id);\n assert!(markers.insert(target.marker), \"duplicate generated marker {}\", target.marker);\n assert_eq!(generated_intrinsic_target_by_marker(target.marker), Some(target));\n assert!(generated_intrinsic_targets_by_op_name(target.dialect_op).any(|candidate| candidate == target));\n }\n", - ); - for record in records { - writeln!( - output, - " let target = generated_intrinsic_target_by_marker({:?}).unwrap();", - intrinsic_marker(catalog, record) - ) - .unwrap(); - writeln!( - output, - " assert_eq!(target.id, {:?});\n assert_eq!(target.abi_id, {:?});\n assert_eq!(target.dialect_op, {:?});\n assert_eq!(target.variant, {});\n assert_eq!(target.requirement.minimum_ptx.encoded(), {});\n assert_eq!(target.requirement.hardware, const {{ {} }});\n assert_eq!(target.backend_requirements, const {{ {} }});\n assert_eq!(target.selections, {});", - record.id, - record.rust.abi_id, - record.dialect.op_name, - generated_intrinsic_variant(record), - record.target.minimum_ptx.encoded(), - generated_hardware_target(&record.target.hardware), - generated_backend_requirements(record), - generated_selection_alternatives(&record.selections), - ) - .unwrap(); - match &record.llvm { - Some(llvm) => { - writeln!( - output, - " assert_eq!(target.llvm.unwrap().properties, &{:?} as &[&str]);", - llvm.properties - ) - .unwrap(); - } - None => output.push_str(" assert!(target.llvm.is_none());\n"), - } - } - output.push_str(" assert!(generated_intrinsic_target_by_marker(\"v1:i9999\").is_none());\n }\n}\n"); - output -} - -fn render_special_register_probe( - catalog: &CatalogFile, - record: &CatalogIntrinsic, - hash: &str, - backend: IntrinsicBackend, -) -> String { - let mut output = llvm_header(catalog, hash); - output.push_str("target triple = \"nvptx64-nvidia-cuda\"\n\n"); - let width = record.scalar_width().expect("special-register width"); - let function_name = format!("probe_{}_{}", record.id, backend_label(backend)); - match special_register_backend_mechanism(record, backend) { - BackendLoweringMechanism::TypedNvvm => { - let symbol = &record - .llvm - .as_ref() - .expect("typed special-register route") - .symbol; - writeln!(output, "declare i{width} @{symbol}()\n").unwrap(); - writeln!(output, "define i{width} @{function_name}() {{").unwrap(); - writeln!(output, " %result = call i{width} @{symbol}()").unwrap(); - } - BackendLoweringMechanism::InlinePtx => { - writeln!(output, "define i{width} @{function_name}() {{").unwrap(); - let sideeffect = if record.special_register.as_ref().is_some_and(|special| { - special.observation == SpecialRegisterObservation::VolatileObservation - }) { - " sideeffect" - } else { - "" - }; - writeln!( - output, - " %result = call i{width} asm{sideeffect} {:?}, {:?}()", - special_register_inline_template(record), - special_register_output_constraint(record) - ) - .unwrap(); - } - } - writeln!(output, " ret i{width} %result\n}}").unwrap(); - output -} - -fn render_tma_probe(catalog: &CatalogFile, record: &CatalogIntrinsic, hash: &str) -> String { - let tma = record.tma.as_ref().expect("TMA contract"); - let operation = tma.operation; - let llvm = llvm(record); - let symbol = llvm.resolved_symbol.as_ref().unwrap_or(&llvm.symbol); - let mut output = llvm_header(catalog, hash); - output.push_str("target triple = \"nvptx64-nvidia-cuda\"\n\n"); - if let Some(dimensions) = tma.dimensions() { - let is_g2s = matches!( - operation, - TmaOperation::G2sTile1d - | TmaOperation::G2sTile2d - | TmaOperation::G2sTile2dMulticast - | TmaOperation::G2sTile2dMulticastCg2 - | TmaOperation::G2sTile3d - | TmaOperation::G2sTile4d - | TmaOperation::G2sTile5d - ); - let coordinates = std::iter::repeat_n("i32", dimensions) - .collect::>() - .join(", "); - let coordinate_parameters = (0..dimensions) - .map(|index| format!("i32 %coord{index}")) - .collect::>() - .join(", "); - let coordinate_arguments = (0..dimensions) - .map(|index| format!("i32 %coord{index}")) - .collect::>() - .join(", "); - if is_g2s { - let multicast = matches!( - operation, - TmaOperation::G2sTile2dMulticast | TmaOperation::G2sTile2dMulticastCg2 - ); - let declaration_coordinates = if coordinates.is_empty() { - String::new() - } else { - format!(", {coordinates}") - }; - writeln!( - output, - "declare void @{symbol}(ptr addrspace(7), ptr addrspace(3), ptr{declaration_coordinates}, i16, i64, i1, i1, i32) #0\n" - ) - .unwrap(); - let parameters = if coordinate_parameters.is_empty() { - String::new() - } else { - format!(", {coordinate_parameters}") - }; - let mask_parameter = if multicast { ", i16 %cta_mask" } else { "" }; - writeln!( - output, - "define void @probe_{}(ptr %dst_generic, ptr %barrier_generic, ptr %tensor_map{parameters}{mask_parameter}) #0 {{", - record.id - ) - .unwrap(); - output.push_str( - " %dst = addrspacecast ptr %dst_generic to ptr addrspace(7)\n %barrier = addrspacecast ptr %barrier_generic to ptr addrspace(3)\n", - ); - let arguments = if coordinate_arguments.is_empty() { - String::new() - } else { - format!(", {coordinate_arguments}") - }; - let mask = if multicast { "%cta_mask" } else { "0" }; - let group = if operation == TmaOperation::G2sTile2dMulticastCg2 { - 2 - } else { - 0 - }; - writeln!( - output, - " call void @{symbol}(ptr addrspace(7) %dst, ptr addrspace(3) %barrier, ptr %tensor_map{arguments}, i16 {mask}, i64 0, i1 {multicast}, i1 false, i32 {group}) #0" - ) - .unwrap(); - } else { - let declaration_coordinates = if coordinates.is_empty() { - String::new() - } else { - format!(", {coordinates}") - }; - writeln!( - output, - "declare void @{symbol}(ptr addrspace(3), ptr{declaration_coordinates}, i64, i1) #0\n" - ) - .unwrap(); - let parameters = if coordinate_parameters.is_empty() { - String::new() - } else { - format!(", {coordinate_parameters}") - }; - writeln!( - output, - "define void @probe_{}(ptr %src_generic, ptr %tensor_map{parameters}) #0 {{", - record.id - ) - .unwrap(); - output.push_str(" %src = addrspacecast ptr %src_generic to ptr addrspace(3)\n"); - let arguments = if coordinate_arguments.is_empty() { - String::new() - } else { - format!(", {coordinate_arguments}") - }; - writeln!( - output, - " call void @{symbol}(ptr addrspace(3) %src, ptr %tensor_map{arguments}, i64 0, i1 false) #0" - ) - .unwrap(); - } - output.push_str(" ret void\n}\n\nattributes #0 = { convergent }\n"); - } else if operation == TmaOperation::PrefetchTensorMap { - writeln!(output, "declare void @{symbol}(ptr)\n").unwrap(); - writeln!( - output, - "define void @probe_{}(ptr %tensor_map) {{", - record.id - ) - .unwrap(); - writeln!(output, " call void @{symbol}(ptr %tensor_map)").unwrap(); - output.push_str(" ret void\n}\n"); - } else if let Some(coordinate_count) = operation.prefetch_coordinate_count() { - let coordinates = std::iter::repeat_n("i32", coordinate_count) - .collect::>() - .join(", "); - let declaration_coordinates = if coordinates.is_empty() { - String::new() - } else { - format!(", {coordinates}") - }; - let coordinate_parameters = (0..coordinate_count) - .map(|index| format!("i32 %coord{index}")) - .collect::>() - .join(", "); - let parameters = if coordinate_parameters.is_empty() { - String::new() - } else { - format!(", {coordinate_parameters}") - }; - let cache_hint_parameter = if operation.uses_prefetch_cache_hint() { - ", i64 %cache_hint" - } else { - "" - }; - let coordinate_arguments = (0..coordinate_count) - .map(|index| format!("i32 %coord{index}")) - .collect::>() - .join(", "); - let arguments = if coordinate_arguments.is_empty() { - String::new() - } else { - format!(", {coordinate_arguments}") - }; - writeln!( - output, - "declare void @{symbol}(ptr{declaration_coordinates}, i64, i1) #0\n" - ) - .unwrap(); - writeln!( - output, - "define void @probe_{}(ptr %tensor_map{parameters}{cache_hint_parameter}) #0 {{", - record.id, - ) - .unwrap(); - let cache_hint = if operation.uses_prefetch_cache_hint() { - "%cache_hint" - } else { - "0" - }; - let use_cache_hint = operation.uses_prefetch_cache_hint(); - writeln!( - output, - " call void @{symbol}(ptr %tensor_map{arguments}, i64 {cache_hint}, i1 {use_cache_hint}) #0" - ) - .unwrap(); - output.push_str(" ret void\n}\n\nattributes #0 = { convergent }\n"); - } else if matches!( - operation, - TmaOperation::ReplaceBoxDim - | TmaOperation::ReplaceElementStride - | TmaOperation::ReplaceElementType - | TmaOperation::ReplaceFillMode - | TmaOperation::ReplaceGlobalAddress - | TmaOperation::ReplaceGlobalDim - | TmaOperation::ReplaceGlobalStride - | TmaOperation::ReplaceInterleaveLayout - | TmaOperation::ReplaceRank - | TmaOperation::ReplaceSwizzleAtomicity - | TmaOperation::ReplaceSwizzleMode - ) { - let (field, width, ordinal, value_ty, immediate, address) = match operation { - TmaOperation::ReplaceBoxDim => ("box_dim", "b32", true, "i32", false, false), - TmaOperation::ReplaceElementStride => { - ("element_stride", "b32", true, "i32", false, false) - } - TmaOperation::ReplaceElementType => ("elemtype", "b32", false, "i32", true, false), - TmaOperation::ReplaceFillMode => ("fill_mode", "b32", false, "i32", true, false), - TmaOperation::ReplaceGlobalAddress => { - ("global_address", "b64", false, "i64", false, true) - } - TmaOperation::ReplaceGlobalDim => ("global_dim", "b32", true, "i32", false, false), - TmaOperation::ReplaceGlobalStride => { - ("global_stride", "b64", true, "i64", false, false) - } - TmaOperation::ReplaceInterleaveLayout => { - ("interleave_layout", "b32", false, "i32", true, false) - } - TmaOperation::ReplaceRank => ("rank", "b32", false, "i32", false, false), - TmaOperation::ReplaceSwizzleAtomicity => { - ("swizzle_atomicity", "b32", false, "i32", true, false) - } - TmaOperation::ReplaceSwizzleMode => ("swizzle_mode", "b32", false, "i32", true, false), - _ => unreachable!("TMA tensor-map replace operation was matched"), - }; - let value_parameter = if immediate { - String::new() - } else if address { - ", ptr %new_address".into() - } else { - format!(", {value_ty} %new_value") - }; - writeln!( - output, - "define void @probe_{}(ptr %tensor_map{value_parameter}) {{", - record.id - ) - .unwrap(); - let mut arguments = vec!["ptr %tensor_map".to_owned()]; - let mut constraints = vec!["l"]; - if ordinal { - arguments.push("i32 0".into()); - constraints.push("n"); - } - if immediate { - arguments.push(format!("{value_ty} 0")); - constraints.push("n"); - } else if address { - arguments.push("ptr %new_address".into()); - constraints.push("l"); - } else { - arguments.push(format!("{value_ty} %new_value")); - constraints.push(if value_ty == "i64" { "l" } else { "r" }); - } - constraints.push("~{memory}"); - let template = if ordinal { - format!("tensormap.replace.tile.{field}.global.b1024.{width} [$0], $1, $2;") - } else { - format!("tensormap.replace.tile.{field}.global.b1024.{width} [$0], $1;") - }; - writeln!( - output, - " call void asm sideeffect {template:?}, {:?}({})", - constraints.join(","), - arguments.join(", ") - ) - .unwrap(); - output.push_str(" ret void\n}\n"); - } else if matches!( - operation, - TmaOperation::FenceProxyTensorMapAcquireCluster - | TmaOperation::FenceProxyTensorMapAcquireCta - | TmaOperation::FenceProxyTensorMapAcquireGpu - | TmaOperation::FenceProxyTensorMapAcquireSystem - ) { - writeln!(output, "declare void @{symbol}(ptr, i32)\n").unwrap(); - writeln!( - output, - "define void @probe_{}(ptr %tensor_map) {{", - record.id - ) - .unwrap(); - writeln!(output, " call void @{symbol}(ptr %tensor_map, i32 128)").unwrap(); - output.push_str(" ret void\n}\n"); - } else if operation == TmaOperation::CommitGroup - || matches!( - operation, - TmaOperation::FenceProxyTensorMapReleaseCluster - | TmaOperation::FenceProxyTensorMapReleaseCta - | TmaOperation::FenceProxyTensorMapReleaseGpu - | TmaOperation::FenceProxyTensorMapReleaseSystem - ) - { - writeln!(output, "declare void @{symbol}()\n").unwrap(); - writeln!(output, "define void @probe_{}() {{", record.id).unwrap(); - writeln!(output, " call void @{symbol}()").unwrap(); - output.push_str(" ret void\n}\n"); - } else { - writeln!(output, "declare void @{symbol}(i32)\n").unwrap(); - writeln!(output, "define void @probe_{}() {{", record.id).unwrap(); - writeln!(output, " call void @{symbol}(i32 0)").unwrap(); - output.push_str(" ret void\n}\n"); - } - output -} - -fn render_execution_control_probe( - catalog: &CatalogFile, - record: &CatalogIntrinsic, - hash: &str, -) -> String { - let operation = ExecutionControlOperation::from_catalog_id(&record.id) - .expect("closed execution-control operation"); - let symbol = &llvm(record).symbol; - let mut output = llvm_header(catalog, hash); - output.push_str("target triple = \"nvptx64-nvidia-cuda\"\n\n"); - match operation { - ExecutionControlOperation::BarrierCtaSync - | ExecutionControlOperation::BarrierCtaSyncAligned - | ExecutionControlOperation::BarrierCtaArrive - | ExecutionControlOperation::BarrierCtaArriveAligned => { - writeln!(output, "declare void @{symbol}(i32, i32)\n").unwrap(); - for (suffix, parameters, barrier_id, thread_count) in [ - ( - "rr", - "i32 %barrier_id, i32 %thread_count", - "%barrier_id", - "%thread_count", - ), - ("ri", "i32 %barrier_id", "%barrier_id", "32"), - ("ir", "i32 %thread_count", "1", "%thread_count"), - ("ii", "", "1", "32"), - ] { - writeln!( - output, - "define void @probe_{}_{suffix}({parameters}) #0 {{", - record.id - ) - .unwrap(); - writeln!( - output, - " call void @{symbol}(i32 {barrier_id}, i32 {thread_count})" - ) - .unwrap(); - output.push_str(" ret void\n}\n\n"); - } - output.push_str("attributes #0 = { convergent }\n"); - } - ExecutionControlOperation::GridDependencyLaunchDependents - | ExecutionControlOperation::GridDependencyWait => { - writeln!(output, "declare void @{symbol}()\n").unwrap(); - writeln!(output, "define void @probe_{}() {{", record.id).unwrap(); - writeln!(output, " call void @{symbol}()").unwrap(); - output.push_str(" ret void\n}\n"); - } - ExecutionControlOperation::SetMaxNRegInc | ExecutionControlOperation::SetMaxNRegDec => { - writeln!(output, "declare void @{symbol}(i32)\n").unwrap(); - writeln!(output, "define void @probe_{}() #0 {{", record.id).unwrap(); - writeln!(output, " call void @{symbol}(i32 64)").unwrap(); - output.push_str(" ret void\n}\n\nattributes #0 = { convergent }\n"); - } - } - output -} - -fn render_tcgen05_mma_probe( - catalog: &CatalogFile, - record: &CatalogIntrinsic, - hash: &str, - mma: &Tcgen05Mma, -) -> String { - let (kind, b_buffer, b_usage) = mma.fixed_selectors.map_or( - (Tcgen05MmaKind::F16, 0, Tcgen05MmaBUsage::Discard), - |fixed| (fixed.kind, fixed.b_buffer, fixed.b_usage), - ); - let (template, constraints) = if tcgen05_mma_is_ws(mma.form) { - tcgen05_mma_inline_asm(mma.form, kind, 1, None, Some(b_buffer), Some(b_usage)) - } else { - tcgen05_mma_inline_asm(mma.form, kind, 1, Some("discard"), None, None) - }; - let parameters = tcgen05_mma_runtime_parameters(mma); - let llvm_parameter = |(name, ty): &(&str, &str)| { - let ty = match *ty { - "u32" => "i32", - "u64" => "i64", - "bool" => "i1", - _ => unreachable!("closed tcgen05 MMA probe type"), - }; - format!("{ty} %{name}") - }; - let signature = parameters - .iter() - .map(llvm_parameter) - .collect::>() - .join(", "); - let argument_indices: Vec = - if mma.alias.is_some() && mma.form == Tcgen05MmaForm::WsTensor { - vec![0, 1, 3, 4, 5] - } else { - (0..parameters.len()).collect() - }; - let arguments = argument_indices - .iter() - .map(|&index| llvm_parameter(¶meters[index])) - .collect::>() - .join(", "); - - let mut output = llvm_header(catalog, hash); - output.push_str("target triple = \"nvptx64-nvidia-cuda\"\n\n"); - writeln!( - output, - "; LLVM source: {}; route: inline PTX; source contract: {:?}", - llvm(record).symbol, - record.tcgen05.as_ref().unwrap().source_contract - ) - .unwrap(); - writeln!( - output, - "define void @probe_{}({signature}) #0 {{", - record.id - ) - .unwrap(); - writeln!( - output, - " call void asm sideeffect {template:?}, {constraints:?}({arguments}) #0" - ) - .unwrap(); - output.push_str(" ret void\n}\n\nattributes #0 = { convergent }\n"); - output -} - -fn render_tcgen05_probe(catalog: &CatalogFile, record: &CatalogIntrinsic, hash: &str) -> String { - let tcgen05 = record.tcgen05.as_ref().expect("tcgen05 contract"); - if let Some(mma) = &tcgen05.mma { - return render_tcgen05_mma_probe(catalog, record, hash, mma); - } - let operation = tcgen05.operation; - let mut output = llvm_header(catalog, hash); - output.push_str("target triple = \"nvptx64-nvidia-cuda\"\n\n"); - writeln!( - output, - "; LLVM source: {}; route: inline PTX; source contract: {:?}", - llvm(record).symbol, - record.tcgen05.as_ref().unwrap().source_contract - ) - .unwrap(); - let (parameters, result_ty, template, constraints, arguments) = match operation { - Tcgen05Operation::Alloc | Tcgen05Operation::AllocCg2 => { - let group = if operation == Tcgen05Operation::AllocCg2 { - "2" - } else { - "1" - }; - ( - "ptr %dst, i32 %ncols".into(), - "void".into(), - format!( - "{{ .reg .u64 %shared64; .reg .u32 %shared32; cvta.to.shared.u64 %shared64, $0; cvt.u32.u64 %shared32, %shared64; tcgen05.alloc.cta_group::{group}.sync.aligned.shared::cta.b32 [%shared32], $1; }}" - ), - "l,r,~{memory}".into(), - "ptr %dst, i32 %ncols".into(), - ) - } - Tcgen05Operation::Dealloc | Tcgen05Operation::DeallocCg2 => { - let group = if operation == Tcgen05Operation::DeallocCg2 { - "2" - } else { - "1" - }; - ( - "i32 %tmem, i32 %ncols".into(), - "void".into(), - format!("tcgen05.dealloc.cta_group::{group}.sync.aligned.b32 $0, $1;"), - "r,r,~{memory}".into(), - "i32 %tmem, i32 %ncols".into(), - ) - } - Tcgen05Operation::RelinquishAllocPermit | Tcgen05Operation::RelinquishAllocPermitCg2 => { - let group = if operation == Tcgen05Operation::RelinquishAllocPermitCg2 { - "2" - } else { - "1" - }; - ( - String::new(), - "void".into(), - format!("tcgen05.relinquish_alloc_permit.cta_group::{group}.sync.aligned;"), - "~{memory}".into(), - String::new(), - ) - } - Tcgen05Operation::FenceBeforeThreadSync | Tcgen05Operation::FenceAfterThreadSync => { - let modifier = if operation == Tcgen05Operation::FenceBeforeThreadSync { - "before_thread_sync" - } else { - "after_thread_sync" - }; - ( - String::new(), - "void".into(), - format!("tcgen05.fence::{modifier};"), - "~{memory}".into(), - String::new(), - ) - } - Tcgen05Operation::Commit - | Tcgen05Operation::CommitSharedCluster - | Tcgen05Operation::CommitCg2 - | Tcgen05Operation::CommitSharedClusterCg2 => { - let group = if matches!( - operation, - Tcgen05Operation::CommitCg2 | Tcgen05Operation::CommitSharedClusterCg2 - ) { - "2" - } else { - "1" - }; - let shared = matches!( - operation, - Tcgen05Operation::CommitSharedCluster | Tcgen05Operation::CommitSharedClusterCg2 - ); - ( - "i32 %mbar".into(), - "void".into(), - format!( - "tcgen05.commit.cta_group::{group}.mbarrier::arrive::one{}.b64 [$0];", - if shared { ".shared::cluster" } else { "" } - ), - "r,~{memory}".into(), - "i32 %mbar".into(), - ) - } - Tcgen05Operation::MmaWsF16 | Tcgen05Operation::MmaWsBf16 | Tcgen05Operation::MmaWsTf32 => { - let kind = if operation == Tcgen05Operation::MmaWsTf32 { - "tf32" - } else { - "f16" - }; - ( - "i32 %d, i32 %a, i64 %legacy_a_desc, i64 %b, i32 %idesc, i1 %enable".into(), - "void".into(), - format!( - "{{ .reg .pred %enable_pred; setp.ne.s32 %enable_pred, $5, 0; tcgen05.mma.ws.cta_group::1.kind::{kind} [$0], [$1], $3, $4, %enable_pred; }}" - ), - "r,r,l,l,r,r,~{memory}".into(), - "i32 %d, i32 %a, i64 %legacy_a_desc, i64 %b, i32 %idesc, i1 %enable".into(), - ) - } - Tcgen05Operation::MmaF16 | Tcgen05Operation::MmaF16Cg2 => { - let group = if operation == Tcgen05Operation::MmaF16Cg2 { - "2" - } else { - "1" - }; - let zeros = if operation == Tcgen05Operation::MmaF16Cg2 { - "%z, %z, %z, %z, %z, %z, %z, %z" - } else { - "%z, %z, %z, %z" - }; - ( - "i32 %d, i64 %a, i64 %b, i32 %idesc, i1 %enable".into(), - "void".into(), - format!( - "{{ .reg .pred %enable_pred; setp.ne.s32 %enable_pred, $4, 0; .reg .u32 %z; mov.u32 %z, 0; tcgen05.mma.cta_group::{group}.kind::f16 [$0], $1, $2, $3, {{{zeros}}}, %enable_pred; }}" - ), - "r,l,l,r,r,~{memory}".into(), - "i32 %d, i64 %a, i64 %b, i32 %idesc, i1 %enable".into(), - ) - } - Tcgen05Operation::CpSmemToTmem | Tcgen05Operation::CpSmemToTmemCg2 => ( - "i32 %tmem, i64 %desc".into(), - "void".into(), - format!( - "{}.{} [$0], $1;", - record.expected_ptx.mnemonic, - record.expected_ptx.modifiers.join(".") - ), - "r,l,~{memory}".into(), - "i32 %tmem, i64 %desc".into(), - ), - Tcgen05Operation::Ld16x256bX8Pure | Tcgen05Operation::Ld16x256bPure => { - let count = if operation == Tcgen05Operation::Ld16x256bX8Pure { - 32 - } else { - 4 - }; - let registers = (0..count) - .map(|index| format!("${index}")) - .collect::>() - .join(","); - let constraints = std::iter::repeat_n("=f", count) - .chain(std::iter::once("r")) - .collect::>() - .join(","); - let fields = std::iter::repeat_n("float", count) - .collect::>() - .join(", "); - ( - "i32 %tmem".into(), - format!("{{ {fields} }}"), - format!( - "tcgen05.ld.sync.aligned.16x256b.{}.b32 {{{registers}}}, [${count}];", - if count == 32 { "x8" } else { "x1" } - ), - constraints, - "i32 %tmem".into(), - ) - } - Tcgen05Operation::St => { - let count = tcgen05_st_register_count(record); - let has_half_split_offset = record - .tcgen05 - .as_ref() - .and_then(|tcgen05| tcgen05.st) - .is_some_and(|st| st.shape == Tcgen05LdShape::M16x32bx2); - let parameters = std::iter::once("i32 %tmem".into()) - .chain((0..count).map(|index| format!("i32 %d{index}"))) - .collect::>() - .join(", "); - let arguments = std::iter::once("i32 %tmem".into()) - .chain(has_half_split_offset.then(|| "i64 16".into())) - .chain((0..count).map(|index| format!("i32 %d{index}"))) - .collect::>() - .join(", "); - let (template, constraints, _) = tcgen05_inline_asm(record); - (parameters, "void".into(), template, constraints, arguments) - } - Tcgen05Operation::Ld => { - let count = tcgen05_ld_register_count(record); - let has_half_split_offset = record - .tcgen05 - .as_ref() - .and_then(|tcgen05| tcgen05.ld) - .is_some_and(|ld| ld.shape == Tcgen05LdShape::M16x32bx2); - let fields = std::iter::repeat_n("i32", count) - .collect::>() - .join(", "); - let (template, constraints, _) = tcgen05_inline_asm(record); - ( - "i32 %tmem".into(), - if count == 1 { - "i32".into() - } else { - format!("{{ {fields} }}") - }, - template, - constraints, - if has_half_split_offset { - "i32 %tmem, i64 16".into() - } else { - "i32 %tmem".into() - }, - ) - } - Tcgen05Operation::LoadWait | Tcgen05Operation::StoreWait => { - let kind = if operation == Tcgen05Operation::LoadWait { - "ld" - } else { - "st" - }; - ( - String::new(), - "void".into(), - format!("tcgen05.wait::{kind}.sync.aligned;"), - "~{memory}".into(), - String::new(), - ) - } - Tcgen05Operation::CommitMulticast | Tcgen05Operation::CommitMulticastCg2 => { - let group = if operation == Tcgen05Operation::CommitMulticastCg2 { - 2 - } else { - 1 - }; - ( - "i32 %mbar, i16 %mask".into(), - "void".into(), - format!( - "tcgen05.commit.cta_group::{group}.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 [$0], $1;" - ), - "r,h,~{memory}".into(), - "i32 %mbar, i16 %mask".into(), - ) - } - Tcgen05Operation::ShiftDown | Tcgen05Operation::ShiftDownCg2 => { - let group = if operation == Tcgen05Operation::ShiftDownCg2 { - 2 - } else { - 1 - }; - ( - "i32 %tmem".into(), - "void".into(), - format!("tcgen05.shift.cta_group::{group}.down [$0];"), - "r,~{memory}".into(), - "i32 %tmem".into(), - ) - } - Tcgen05Operation::Mma => unreachable!("generic MMA probe handled above"), - }; - let (lowering_template, lowering_constraints, lowering_results) = tcgen05_inline_asm(record); - assert_eq!(template, lowering_template); - assert_eq!(constraints, lowering_constraints); - assert_eq!(result_ty != "void", lowering_results.is_some()); - writeln!( - output, - "define void @probe_{}({parameters}) #0 {{", - record.id - ) - .unwrap(); - if result_ty == "void" { - writeln!( - output, - " call void asm sideeffect {template:?}, {constraints:?}({arguments}) #0" - ) - .unwrap(); - } else { - writeln!( - output, - " %result = call {result_ty} asm sideeffect {template:?}, {constraints:?}({arguments}) #0" - ) - .unwrap(); - } - output.push_str(" ret void\n}\n\nattributes #0 = { convergent }\n"); - output -} - -fn render_elect_probe( - catalog: &CatalogFile, - record: &CatalogIntrinsic, - hash: &str, - backend: IntrinsicBackend, -) -> String { - let mut output = llvm_header(catalog, hash); - output.push_str("target triple = \"nvptx64-nvidia-cuda\"\n\n"); - let mechanism = record - .backend_lowerings - .iter() - .find(|route| route.backend == backend) - .expect("elect backend route") - .mechanism; - match mechanism { - BackendLoweringMechanism::TypedNvvm => { - writeln!( - output, - "declare {{ i32, i1 }} @{}(i32) #0\n", - llvm(record).symbol - ) - .unwrap(); - for (suffix, parameters, mask) in - [("register", "i32 %mask", "%mask"), ("immediate", "", "-1")] - { - writeln!( - output, - "define {{ i32, i1 }} @probe_{}_{}_{suffix}({parameters}) #0 {{", - record.id, - backend_label(backend) - ) - .unwrap(); - writeln!( - output, - " %result = call {{ i32, i1 }} @{}(i32 {mask}) #0", - llvm(record).symbol - ) - .unwrap(); - output.push_str(" ret { i32, i1 } %result\n}\n"); - } - } - BackendLoweringMechanism::InlinePtx => { - for (suffix, parameters, mask) in [ - ("register", "i32 %mask", "%mask"), - ("constant_mask", "", "-1"), - ] { - writeln!( - output, - "define {{ i32, i1 }} @probe_{}_{}_{suffix}({parameters}) #0 {{", - record.id, - backend_label(backend) - ) - .unwrap(); - writeln!( - output, - " %raw = call {{ i32, i32 }} asm sideeffect \"{{ .reg .pred p; elect.sync $0|p, $2; selp.b32 $1, 1, 0, p; }}\", \"=r,=r,r\"(i32 {mask}) #0" - ) - .unwrap(); - output.push_str( - " %leader = extractvalue { i32, i32 } %raw, 0\n\ - \x20 %predicate.i32 = extractvalue { i32, i32 } %raw, 1\n\ - \x20 %predicate = trunc i32 %predicate.i32 to i1\n\ - \x20 %result.0 = insertvalue { i32, i1 } poison, i32 %leader, 0\n\ - \x20 %result.1 = insertvalue { i32, i1 } %result.0, i1 %predicate, 1\n\ - \x20 ret { i32, i1 } %result.1\n}\n", - ); - } - } - } - output.push_str("\nattributes #0 = { convergent }\n"); - output -} - -pub(crate) fn render_probe(catalog: &CatalogFile, record: &CatalogIntrinsic, hash: &str) -> String { - if record.special_register.is_some() { - return render_special_register_probe(catalog, record, hash, IntrinsicBackend::LlvmNvptx); - } - if record.tma.is_some() { - return render_tma_probe(catalog, record, hash); - } - if ExecutionControlOperation::from_catalog_id(&record.id).is_some() { - return render_execution_control_probe(catalog, record, hash); - } - if record.tcgen05.is_some() { - return render_tcgen05_probe(catalog, record, hash); - } - if record.family == "elect" { - return render_elect_probe(catalog, record, hash, IntrinsicBackend::LlvmNvptx); - } - let mut output = llvm_header(catalog, hash); - output.push_str("target triple = \"nvptx64-nvidia-cuda\"\n\n"); - if let Some(cluster) = &record.cluster_memory { - let (template, constraints) = - crate::resolve::cluster_memory_inline_recipe(cluster.operation); - match cluster.operation { - ClusterMemoryOperation::MapSharedRank => { - writeln!( - output, - "define ptr addrspace(7) @probe_{}(ptr %source_generic, i32 %rank) #0 {{", - record.id - ) - .unwrap(); - output.push_str( - " %source = addrspacecast ptr %source_generic to ptr addrspace(3)\n", - ); - writeln!( - output, - " %mapped_integer = call i64 asm sideeffect {template:?}, {constraints:?}(ptr addrspace(3) %source, i32 %rank) #0" - ) - .unwrap(); - output.push_str( - " %mapped = inttoptr i64 %mapped_integer to ptr addrspace(7)\n ret ptr addrspace(7) %mapped\n}\n\nattributes #0 = { convergent }\n", - ); - } - ClusterMemoryOperation::ReadU32 => { - writeln!( - output, - "define i32 @probe_{}(ptr %source_generic, i32 %rank) #0 {{", - record.id - ) - .unwrap(); - output.push_str( - " %source = addrspacecast ptr %source_generic to ptr addrspace(3)\n", - ); - writeln!( - output, - " %value = call i32 asm sideeffect {template:?}, {constraints:?}(ptr addrspace(3) %source, i32 %rank) #0" - ) - .unwrap(); - output.push_str(" ret i32 %value\n}\n\nattributes #0 = { convergent }\n"); - } - } - } else if let Some(clc) = &record.clc { - match clc.adapter { - ClcAdapter::GenericPointersToShared => { - writeln!( - output, - "declare void @{}(ptr addrspace(3), ptr addrspace(3))\n", - llvm(record).symbol - ) - .unwrap(); - writeln!( - output, - "define void @probe_{}(ptr %response_generic, ptr %mbarrier_generic) {{", - record.id - ) - .unwrap(); - output.push_str( - " %response = addrspacecast ptr %response_generic to ptr addrspace(3)\n\ - \x20 %mbarrier = addrspacecast ptr %mbarrier_generic to ptr addrspace(3)\n", - ); - writeln!( - output, - " call void @{}(ptr addrspace(3) %response, ptr addrspace(3) %mbarrier)", - llvm(record).symbol - ) - .unwrap(); - output.push_str(" ret void\n}\n"); - } - ClcAdapter::PairU64ToI128BoolToU32 | ClcAdapter::PairU64ToI128U32 => { - let llvm_result = if clc.adapter == ClcAdapter::PairU64ToI128BoolToU32 { - "i1" - } else { - "i32" - }; - writeln!( - output, - "declare {llvm_result} @{}(i128)\n", - llvm(record).symbol - ) - .unwrap(); - writeln!( - output, - "define i32 @probe_{}(i64 %response_low, i64 %response_high) {{", - record.id - ) - .unwrap(); - output.push_str( - " %response_low_i128 = zext i64 %response_low to i128\n\ - \x20 %response_high_i128 = zext i64 %response_high to i128\n\ - \x20 %response_high_shifted = shl i128 %response_high_i128, 64\n\ - \x20 %response = or i128 %response_low_i128, %response_high_shifted\n", - ); - writeln!( - output, - " %raw_result = call {llvm_result} @{}(i128 %response)", - llvm(record).symbol - ) - .unwrap(); - if clc.adapter == ClcAdapter::PairU64ToI128BoolToU32 { - output.push_str( - " %result = zext i1 %raw_result to i32\n\ - \x20 ret i32 %result\n}\n", - ); - } else { - output.push_str(" ret i32 %raw_result\n}\n"); - } - } - } - } else if let Some(bridge) = &record.cp_async_mbarrier { - let shared = bridge.state_space == CpAsyncMbarrierStateSpace::Shared; - let pointer_type = if shared { "ptr addrspace(3)" } else { "ptr" }; - writeln!( - output, - "declare void @{}({pointer_type})\n", - llvm(record).symbol - ) - .unwrap(); - writeln!( - output, - "define void @probe_{}(ptr %barrier_generic) {{", - record.id - ) - .unwrap(); - if shared { - output - .push_str(" %barrier = addrspacecast ptr %barrier_generic to ptr addrspace(3)\n"); - writeln!( - output, - " call void @{}(ptr addrspace(3) %barrier)", - llvm(record).symbol - ) - .unwrap(); - } else { - writeln!( - output, - " call void @{}(ptr %barrier_generic)", - llvm(record).symbol - ) - .unwrap(); - } - output.push_str(" ret void\n}\n"); - } else if let Some(mbarrier) = &record.mbarrier_basic { - match mbarrier.operation { - MbarrierBasicOperation::Init => { - writeln!( - output, - "declare void @{}(ptr addrspace(3), i32)\n", - llvm(record).symbol - ) - .unwrap(); - writeln!( - output, - "define void @probe_{}(ptr %barrier_generic, i32 %expected_count) {{", - record.id - ) - .unwrap(); - output.push_str( - " %barrier = addrspacecast ptr %barrier_generic to ptr addrspace(3)\n", - ); - writeln!( - output, - " call void @{}(ptr addrspace(3) %barrier, i32 %expected_count)", - llvm(record).symbol - ) - .unwrap(); - output.push_str(" ret void\n}\n"); - } - MbarrierBasicOperation::Arrive => { - writeln!( - output, - "declare i64 @{}(ptr addrspace(3))\n", - llvm(record).symbol - ) - .unwrap(); - writeln!( - output, - "define i64 @probe_{}(ptr %barrier_generic) {{", - record.id - ) - .unwrap(); - output.push_str( - " %barrier = addrspacecast ptr %barrier_generic to ptr addrspace(3)\n", - ); - writeln!( - output, - " %token = call i64 @{}(ptr addrspace(3) %barrier)", - llvm(record).symbol - ) - .unwrap(); - output.push_str(" ret i64 %token\n}\n"); - } - MbarrierBasicOperation::TestWait => { - writeln!( - output, - "define i1 @probe_{}(ptr %barrier_generic, i64 %token) #0 {{", - record.id - ) - .unwrap(); - output.push_str( - " %barrier = addrspacecast ptr %barrier_generic to ptr addrspace(3)\n\ - \x20 %result_i32 = call i32 asm sideeffect \"{ .reg .pred %p0; mbarrier.test_wait.shared.b64 %p0, [$1], $2; selp.b32 $0, 1, 0, %p0; }\", \"=r,l,l,~{memory}\"(ptr addrspace(3) %barrier, i64 %token) #0\n\ - \x20 %result = trunc i32 %result_i32 to i1\n\ - \x20 ret i1 %result\n\ - }\n\n\ - attributes #0 = { convergent }\n", - ); - } - MbarrierBasicOperation::Inval => { - writeln!( - output, - "declare void @{}(ptr addrspace(3))\n", - llvm(record).symbol - ) - .unwrap(); - writeln!( - output, - "define void @probe_{}(ptr %barrier_generic) {{", - record.id - ) - .unwrap(); - output.push_str( - " %barrier = addrspacecast ptr %barrier_generic to ptr addrspace(3)\n", - ); - writeln!( - output, - " call void @{}(ptr addrspace(3) %barrier)", - llvm(record).symbol - ) - .unwrap(); - output.push_str(" ret void\n}\n"); - } - } - } else if record.movmatrix.is_some() { - writeln!(output, "define i32 @probe_{}(i32 %value) #0 {{", record.id).unwrap(); - writeln!( - output, - " %result = call i32 asm {:?}, \"=r,r\"(i32 %value) #0", - movmatrix_template(record), - ) - .unwrap(); - output.push_str(" ret i32 %result\n}\n\nattributes #0 = { convergent }\n"); - } else if let Some(mbarrier) = &record.mbarrier_extended { - let (template, constraints) = - crate::resolve::mbarrier_extended_inline_recipe(mbarrier.operation); - match mbarrier.adapter { - MbarrierExtendedAdapter::PointerTxCountBytesToTokenDroppingTxCount => { - writeln!( - output, - "define i64 @probe_{}(ptr %barrier_generic, i32 %bytes) #0 {{", - record.id - ) - .unwrap(); - output.push_str( - " %barrier = addrspacecast ptr %barrier_generic to ptr addrspace(3)\n", - ); - writeln!( - output, - " %token = call i64 asm sideeffect {template:?}, {constraints:?}(ptr addrspace(3) %barrier, i32 %bytes) #0" - ) - .unwrap(); - output.push_str(" ret i64 %token\n}\n\nattributes #0 = { convergent }\n"); - } - MbarrierExtendedAdapter::PointerTokenToPredicate => { - writeln!( - output, - "define i1 @probe_{}(ptr %barrier_generic, i64 %token) #0 {{", - record.id - ) - .unwrap(); - output.push_str( - " %barrier = addrspacecast ptr %barrier_generic to ptr addrspace(3)\n", - ); - writeln!( - output, - " %result_i32 = call i32 asm sideeffect {template:?}, {constraints:?}(ptr addrspace(3) %barrier, i64 %token) #0" - ) - .unwrap(); - output.push_str( - " %result = trunc i32 %result_i32 to i1\n ret i1 %result\n}\n\nattributes #0 = { convergent }\n", - ); - } - MbarrierExtendedAdapter::PointerParityToPredicate => { - writeln!( - output, - "define i1 @probe_{}(ptr %barrier_generic, i32 %parity) #0 {{", - record.id - ) - .unwrap(); - output.push_str( - " %barrier = addrspacecast ptr %barrier_generic to ptr addrspace(3)\n", - ); - writeln!( - output, - " %result_i32 = call i32 asm sideeffect {template:?}, {constraints:?}(ptr addrspace(3) %barrier, i32 %parity) #0" - ) - .unwrap(); - output.push_str( - " %result = trunc i32 %result_i32 to i1\n ret i1 %result\n}\n\nattributes #0 = { convergent }\n", - ); - } - MbarrierExtendedAdapter::RawClusterAddressToVoid => { - writeln!( - output, - "define void @probe_{}(i64 %address) #0 {{", - record.id - ) - .unwrap(); - writeln!( - output, - " call void asm sideeffect {template:?}, {constraints:?}(i64 %address) #0" - ) - .unwrap(); - output.push_str(" ret void\n}\n\nattributes #0 = { convergent }\n"); - } - MbarrierExtendedAdapter::ZeroOperandsToVoid => { - writeln!(output, "define void @probe_{}() #0 {{", record.id).unwrap(); - writeln!( - output, - " call void asm sideeffect {template:?}, {constraints:?}() #0" - ) - .unwrap(); - output.push_str(" ret void\n}\n\nattributes #0 = { convergent }\n"); - } - MbarrierExtendedAdapter::NanosecondsToVoid => { - writeln!(output, "define void @probe_{}(i32 %ns) #0 {{", record.id).unwrap(); - writeln!( - output, - " call void asm sideeffect {template:?}, {constraints:?}(i32 %ns) #0" - ) - .unwrap(); - output.push_str(" ret void\n}\n\nattributes #0 = { convergent }\n"); - } - } - } else if record.integer_minmax.is_some() { - let parameters = "i32 %arg0, i32 %arg1"; - writeln!(output, "define i32 @probe_{}({parameters}) {{", record.id).unwrap(); - writeln!( - output, - " %result = call i32 asm \"{} $0, $1, $2;\", \"=r,r,r\"(i32 %arg0, i32 %arg1)", - integer_minmax_ptx_mnemonic(record) - ) - .unwrap(); - output.push_str(" ret i32 %result\n}\n"); - } else if record.packed_alu.is_some() { - let arity = record.rust.arguments.len(); - let width = packed_alu_width(record); - let llvm_type = format!("i{width}"); - let register_constraint = packed_alu_register_constraint(record); - let parameters = (0..arity) - .map(|index| format!("{llvm_type} %arg{index}")) - .collect::>() - .join(", "); - let arguments = (0..arity) - .map(|index| format!("{llvm_type} %arg{index}")) - .collect::>() - .join(", "); - let operands = (0..=arity) - .map(|index| format!("${index}")) - .collect::>() - .join(", "); - let constraints = std::iter::once(format!("={register_constraint}")) - .chain(std::iter::repeat_n(register_constraint.to_owned(), arity)) - .collect::>() - .join(","); - writeln!( - output, - "define {llvm_type} @probe_{}({parameters}) {{", - record.id - ) - .unwrap(); - writeln!( - output, - " %result = call {llvm_type} asm \"{} {operands};\", \"{constraints}\"({arguments})", - packed_alu_ptx_mnemonic(record) - ) - .unwrap(); - writeln!(output, " ret {llvm_type} %result\n}}").unwrap(); - } else if record.packed_conversion.is_some() { - let result_ty = packed_conversion_dialect_type(record); - if packed_conversion_typed_llvm_name(record).is_some() { - let llvm = llvm(record); - let symbol = llvm.resolved_symbol.as_deref().unwrap_or(&llvm.symbol); - writeln!(output, "declare {result_ty} @{symbol}(float, float)\n").unwrap(); - writeln!( - output, - "define {result_ty} @probe_{}(float %low, float %high) {{", - record.id, - ) - .unwrap(); - writeln!( - output, - " %result = call {result_ty} @{symbol}(float %high, float %low)" - ) - .unwrap(); - } else if packed_conversion_source(record) == PackedConversionSourceFormat::F32x2 { - writeln!( - output, - "define {result_ty} @probe_{}(float %low, float %high) {{", - record.id, - ) - .unwrap(); - writeln!( - output, - " %result = call {result_ty} asm \"{} $0, $2, $1;\", \"{}\"(float %low, float %high)", - packed_conversion_ptx_mnemonic(record), - packed_conversion_constraint(record), - ) - .unwrap(); - } else { - // A packed source arrives in one integer register, so the probe - // takes a single operand and needs no reordering. - let source_ty = format!("i{}", packed_conversion_source_width(record)); - writeln!( - output, - "define {result_ty} @probe_{}({source_ty} %packed) {{", - record.id, - ) - .unwrap(); - writeln!( - output, - " %result = call {result_ty} asm \"{} $0, $1;\", \"{}\"({source_ty} %packed)", - packed_conversion_ptx_mnemonic(record), - packed_conversion_constraint(record), - ) - .unwrap(); - } - writeln!(output, " ret {result_ty} %result\n}}").unwrap(); - } else if let Some(packed) = &record.packed_atomic { - let format = match packed.format { - PackedAtomicFormat::F16x2 => "f16x2", - PackedAtomicFormat::Bf16x2 => "bf16x2", - }; - writeln!( - output, - "define i32 @probe_{}(ptr %address, i32 %addend) {{", - record.id - ) - .unwrap(); - writeln!( - output, - " %old = call i32 asm sideeffect \"atom.global.add.noftz.{format} $0, [$1], $2;\", \"=r,l,r,~{{memory}}\"(ptr %address, i32 %addend)" - ) - .unwrap(); - output.push_str(" ret i32 %old\n}\n"); - } else if let Some(copy) = &record.cp_async_copy { - let dynamic_source_size = copy.source_size == CpAsyncSourceSize::Runtime; - let declaration_arguments = if dynamic_source_size { - "ptr addrspace(3), ptr addrspace(1), i32" - } else { - "ptr addrspace(3), ptr addrspace(1)" - }; - writeln!( - output, - "declare void @{}({declaration_arguments})", - llvm(record).symbol - ) - .unwrap(); - output.push('\n'); - if dynamic_source_size { - writeln!( - output, - "define void @probe_{}_register(ptr %shared_generic, ptr %global_generic, i32 %source_size) {{", - record.id - ) - .unwrap(); - output.push_str( - " %shared = addrspacecast ptr %shared_generic to ptr addrspace(3)\n %global = addrspacecast ptr %global_generic to ptr addrspace(1)\n", - ); - writeln!( - output, - " call void @{}(ptr addrspace(3) %shared, ptr addrspace(1) %global, i32 %source_size)", - llvm(record).symbol - ) - .unwrap(); - output.push_str(" ret void\n}\n"); - writeln!( - output, - "define void @probe_{}_immediate(ptr %shared_generic, ptr %global_generic) {{", - record.id - ) - .unwrap(); - output.push_str( - " %shared = addrspacecast ptr %shared_generic to ptr addrspace(3)\n %global = addrspacecast ptr %global_generic to ptr addrspace(1)\n", - ); - writeln!( - output, - " call void @{}(ptr addrspace(3) %shared, ptr addrspace(1) %global, i32 3)", - llvm(record).symbol - ) - .unwrap(); - output.push_str(" ret void\n}\n"); - } else { - writeln!( - output, - "define void @probe_{}(ptr %shared_generic, ptr %global_generic) {{", - record.id - ) - .unwrap(); - output.push_str( - " %shared = addrspacecast ptr %shared_generic to ptr addrspace(3)\n %global = addrspacecast ptr %global_generic to ptr addrspace(1)\n", - ); - writeln!( - output, - " call void @{}(ptr addrspace(3) %shared, ptr addrspace(1) %global)", - llvm(record).symbol - ) - .unwrap(); - output.push_str(" ret void\n}\n"); - } - } else if let Some(control) = &record.cp_async_control { - let has_immediate = control.operation == CpAsyncControlOperation::WaitGroup; - let declaration_arguments = if has_immediate { "i32" } else { "" }; - writeln!( - output, - "declare void @{}({declaration_arguments})", - llvm(record).symbol - ) - .unwrap(); - output.push('\n'); - writeln!(output, "define void @probe_{}() {{", record.id).unwrap(); - if has_immediate { - writeln!(output, " call void @{}(i32 3)", llvm(record).symbol).unwrap(); - } else { - writeln!(output, " call void @{}()", llvm(record).symbol).unwrap(); - } - output.push_str(" ret void\n}\n"); - } else if let Some(control) = &record.wgmma_control { - if control.mode == WgmmaControlMode::WaitGroup { - writeln!( - output, - "declare void @{}(i64 immarg) #0\n", - llvm(record).symbol - ) - .unwrap(); - } else { - writeln!(output, "declare void @{}() #0\n", llvm(record).symbol).unwrap(); - } - writeln!(output, "define void @probe_{}() #0 {{", record.id).unwrap(); - if control.mode == WgmmaControlMode::WaitGroup { - writeln!(output, " call void @{}(i64 0) #0", llvm(record).symbol,).unwrap(); - } else { - writeln!(output, " call void @{}() #0", llvm(record).symbol,).unwrap(); - } - output.push_str(" ret void\n}\n\nattributes #0 = { convergent }\n"); - } else if record.cluster_barrier.is_some() { - writeln!(output, "declare void @{}() #0", llvm(record).symbol).unwrap(); - output.push('\n'); - writeln!(output, "define void @probe_{}() #0 {{", record.id).unwrap(); - writeln!(output, " call void @{}() #0", llvm(record).symbol).unwrap(); - output.push_str(" ret void\n}\n\nattributes #0 = { convergent }\n"); - } else if record.family == "sync" { - let arguments = if record.id == "sync_threads" { - "i32" - } else { - debug_assert!(threadfence_ptx_level(record).is_some()); - "" - }; - writeln!(output, "declare void @{}({arguments})", llvm(record).symbol).unwrap(); - output.push('\n'); - writeln!(output, "define void @probe_{}() {{", record.id).unwrap(); - let call_arguments = if record.id == "sync_threads" { - "i32 0" - } else { - "" - }; - writeln!( - output, - " call void @{}({call_arguments})", - llvm(record).symbol - ) - .unwrap(); - output.push_str(" ret void\n}\n"); - } else if let Some(warp_barrier) = &record.warp_barrier { - debug_assert_eq!(warp_barrier.adapter, WarpBarrierAdapter::DirectMemberMask); - writeln!(output, "declare void @{}(i32)", llvm(record).symbol).unwrap(); - output.push('\n'); - writeln!( - output, - "define void @probe_{}(i32 %member_mask) {{", - record.id - ) - .unwrap(); - writeln!( - output, - " call void @{}(i32 %member_mask)", - llvm(record).symbol - ) - .unwrap(); - output.push_str(" ret void\n}\n"); - writeln!(output, "define void @probe_{}_immediate() {{", record.id).unwrap(); - writeln!(output, " call void @{}(i32 -1)", llvm(record).symbol).unwrap(); - output.push_str(" ret void\n}\n"); - } else if let Some(dot) = &record.dot_product { - match dot.adapter { - DotProductAdapter::DirectThreeOperands => { - writeln!( - output, - "declare i32 @{}(i32, i32, i32)", - llvm(record).symbol - ) - .unwrap(); - output.push('\n'); - writeln!( - output, - "define i32 @probe_{}(i32 %a, i32 %b, i32 %c) {{", - record.id - ) - .unwrap(); - writeln!( - output, - " %result = call i32 @{}(i32 %a, i32 %b, i32 %c)", - llvm(record).symbol - ) - .unwrap(); - } - DotProductAdapter::InsertLowHalfFalse => { - writeln!( - output, - "declare i32 @{}(i32, i32, i1, i32)", - llvm(record).symbol - ) - .unwrap(); - output.push('\n'); - writeln!( - output, - "define i32 @probe_{}(i32 %a, i32 %b, i32 %c) {{", - record.id - ) - .unwrap(); - writeln!( - output, - " %result = call i32 @{}(i32 %a, i32 %b, i1 false, i32 %c)", - llvm(record).symbol - ) - .unwrap(); - } - } - output.push_str(" ret i32 %result\n}\n"); - } else if let Some(vote) = &record.vote { - debug_assert_eq!(vote.adapter, VoteAdapter::DirectMaskPredicate); - let result_ty = match vote.mode { - VoteMode::All | VoteMode::Any | VoteMode::Uni => "i1", - VoteMode::Ballot => "i32", - }; - writeln!( - output, - "declare {result_ty} @{}(i32, i1)", - llvm(record).symbol - ) - .unwrap(); - output.push('\n'); - writeln!( - output, - "define {result_ty} @probe_{}(i32 %member_mask, i1 %predicate) {{", - record.id - ) - .unwrap(); - writeln!( - output, - " %result = call {result_ty} @{}(i32 %member_mask, i1 %predicate)", - llvm(record).symbol - ) - .unwrap(); - writeln!(output, " ret {result_ty} %result\n}}").unwrap(); - writeln!( - output, - "define {result_ty} @probe_{}_immediate(i1 %predicate) {{", - record.id - ) - .unwrap(); - writeln!( - output, - " %result = call {result_ty} @{}(i32 -1, i1 %predicate)", - llvm(record).symbol - ) - .unwrap(); - writeln!(output, " ret {result_ty} %result\n}}").unwrap(); - } else if let Some(warp_match) = &record.warp_match { - let value_ty = format!("i{}", warp_match.value_width.bits()); - let result_ty = match warp_match.mode { - WarpMatchMode::Any => "i32".to_owned(), - WarpMatchMode::All => "{ i32, i1 }".to_owned(), - }; - writeln!( - output, - "declare {result_ty} @{}(i32, {value_ty})", - llvm(record).symbol - ) - .unwrap(); - output.push('\n'); - let forms = [ - ( - "rr", - "i32 %member_mask, ", - "i32 %member_mask", - format!("{value_ty} %value"), - ), - ("ri", "", "i32 -1", format!("{value_ty} %value")), - ( - "ir", - "i32 %member_mask", - "i32 %member_mask", - format!("{value_ty} 7"), - ), - ("ii", "", "i32 -1", format!("{value_ty} 7")), - ]; - for (suffix, first_parameter, mask, value) in forms { - let parameters = match suffix { - "rr" => format!("{first_parameter}{value_ty} %value"), - "ri" => format!("{value_ty} %value"), - "ir" => first_parameter.to_owned(), - "ii" => String::new(), - _ => unreachable!(), - }; - writeln!( - output, - "define {result_ty} @probe_{}_{suffix}({parameters}) {{", - record.id - ) - .unwrap(); - writeln!( - output, - " %result = call {result_ty} @{}({mask}, {value})", - llvm(record).symbol - ) - .unwrap(); - writeln!(output, " ret {result_ty} %result\n}}").unwrap(); - } - } else if let Some(warp_shuffle) = &record.warp_shuffle { - if warp_shuffle.value_kind == WarpShuffleValueKind::I64 { - debug_assert_eq!( - warp_shuffle.adapter, - WarpShuffleAdapter::MaskValueLaneOrDeltaSplitI64LowHighB32InsertClampReassemble - ); - let mode = match warp_shuffle.mode { - WarpShuffleMode::Idx => "idx", - WarpShuffleMode::Bfly => "bfly", - WarpShuffleMode::Down => "down", - WarpShuffleMode::Up => "up", - }; - let asm = format!( - "{{ .reg .b32 lo; .reg .b32 hi; mov.b64 {{lo, hi}}, $1; shfl.sync.{mode}.b32 lo, lo, $2, {}, $3; shfl.sync.{mode}.b32 hi, hi, $2, {}, $3; mov.b64 $0, {{lo, hi}}; }}", - warp_shuffle.clamp, warp_shuffle.clamp - ); - writeln!( - output, - "define i64 @probe_{}(i32 %member_mask, i64 %value, i32 %lane) #0 {{", - record.id - ) - .unwrap(); - writeln!( - output, - " %result = call i64 asm sideeffect {asm:?}, \"=l,l,r,r\"(i64 %value, i32 %lane, i32 %member_mask) #0" - ) - .unwrap(); - output.push_str(" ret i64 %result\n}\n\nattributes #0 = { convergent }\n"); - } else { - debug_assert_eq!( - warp_shuffle.adapter, - WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp - ); - let value_ty = match warp_shuffle.value_kind { - WarpShuffleValueKind::I32 => "i32", - WarpShuffleValueKind::F32 => "float", - WarpShuffleValueKind::I64 => unreachable!(), - }; - writeln!( - output, - "declare {value_ty} @{}(i32, {value_ty}, i32, i32)", - llvm(record).symbol - ) - .unwrap(); - output.push('\n'); - let forms = [ - ( - "rr", - format!("i32 %member_mask, {value_ty} %value, i32 %lane"), - "i32 %member_mask", - "i32 %lane", - ), - ( - "ri", - format!("{value_ty} %value, i32 %lane"), - "i32 -1", - "i32 %lane", - ), - ( - "ir", - format!("i32 %member_mask, {value_ty} %value"), - "i32 %member_mask", - "i32 1", - ), - ("ii", format!("{value_ty} %value"), "i32 -1", "i32 1"), - ]; - for (suffix, parameters, member_mask, lane) in forms { - writeln!( - output, - "define {value_ty} @probe_{}_{suffix}({parameters}) {{", - record.id - ) - .unwrap(); - writeln!( - output, - " %result = call {value_ty} @{}({member_mask}, {value_ty} %value, {lane}, i32 {})", - llvm(record).symbol, - warp_shuffle.clamp, - ) - .unwrap(); - writeln!(output, " ret {value_ty} %result\n}}").unwrap(); - } - } - } else if let Some(redux) = &record.redux { - debug_assert_eq!(redux.adapter, ReduxAdapter::MaskValueToSourceMemberMask); - let value_type = match llvm(record).results[0].as_str() { - "i32" => "i32", - "f32" => "float", - other => panic!("unsupported redux value type {other}"), - }; - writeln!( - output, - "declare {value_type} @{}({value_type}, i32)", - llvm(record).symbol - ) - .unwrap(); - output.push('\n'); - writeln!( - output, - "define {value_type} @probe_{}(i32 %member_mask, {value_type} %value) {{", - record.id - ) - .unwrap(); - writeln!( - output, - " %result = call {value_type} @{}({value_type} %value, i32 %member_mask)", - llvm(record).symbol - ) - .unwrap(); - writeln!(output, " ret {value_type} %result\n}}").unwrap(); - } else if let Some(sparse_mma) = &record.sparse_mma { - let accumulator = match sparse_mma.accumulator { - SparseMmaAccumulator::F16 => "i32", - SparseMmaAccumulator::F32 => "float", - SparseMmaAccumulator::S32 => "i32", - }; - let (c_count, a_count, b_count, d_count) = sparse_mma_fragment_counts(record); - let result = format!( - "{{ {} }}", - std::iter::repeat_n(accumulator, d_count) - .collect::>() - .join(", ") - ); - let parameters = (0..c_count) - .map(|index| format!("{accumulator} %c{index}")) - .chain((0..a_count).map(|index| format!("i32 %a{index}"))) - .chain((0..b_count).map(|index| format!("i32 %b{index}"))) - .chain(std::iter::once("i32 %metadata".to_owned())) - .collect::>() - .join(", "); - let arguments = (0..c_count) - .map(|index| format!("{accumulator} %c{index}")) - .chain((0..a_count).map(|index| format!("i32 %a{index}"))) - .chain((0..b_count).map(|index| format!("i32 %b{index}"))) - .chain(std::iter::once("i32 %metadata".to_owned())) - .collect::>() - .join(", "); - for selector in sparse_mma_selector_values(record) { - writeln!( - output, - "define {result} @probe_{}_selector_{selector}({parameters}) #0 {{", - record.id - ) - .unwrap(); - writeln!( - output, - " %result = call {result} asm sideeffect {:?}, {:?}({arguments}, i32 {selector}) #0", - sparse_mma_template(record), - sparse_mma_constraints(record), - ) - .unwrap(); - writeln!(output, " ret {result} %result\n}}").unwrap(); - } - output.push_str("\nattributes #0 = { convergent }\n"); - } else if let Some(mma) = &record.register_mma { - let (c_count, a_count, b_count, d_count) = register_mma_fragment_counts(record); - let (c_type, packed_type, result_type) = match mma.adapter { - RegisterMmaAdapter::C2U32A2U32B1U32ToD2U32 - | RegisterMmaAdapter::C2U32A4U32B2U32ToD2U32 => ("i32", "i32", "i32"), - RegisterMmaAdapter::C4F32A2U32B1U32ToD4F32 - | RegisterMmaAdapter::C4F32A4U32B2U32ToD4F32 - | RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 => { - ("float", "i32", "float") - } - RegisterMmaAdapter::C2F64A1F64B1F64ToD2F64 => ("double", "double", "double"), - RegisterMmaAdapter::C2I32A1U32B1U32ToD2I32 - | RegisterMmaAdapter::C4I32A4U32B2U32ToD4I32 - | RegisterMmaAdapter::C4I32A2U32B1U32ToD4I32 => ("i32", "i32", "i32"), - }; - let mut parameter_values = (0..c_count) - .map(|index| format!("{c_type} %c{index}")) - .chain((0..a_count).map(|index| format!("{packed_type} %a{index}"))) - .chain((0..b_count).map(|index| format!("{packed_type} %b{index}"))) - .collect::>(); - let mut argument_values = (0..c_count) - .map(|index| format!("{c_type} %c{index}")) - .chain((0..a_count).map(|index| format!("{packed_type} %a{index}"))) - .chain((0..b_count).map(|index| format!("{packed_type} %b{index}"))) - .collect::>(); - if mma.adapter == RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 { - let block_scale_parameters = [ - "i32 %scale_a", - "i16 %byte_id_a", - "i16 %thread_id_a", - "i32 %scale_b", - "i16 %byte_id_b", - "i16 %thread_id_b", - ]; - parameter_values.extend(block_scale_parameters.map(str::to_owned)); - argument_values.extend(block_scale_parameters.map(str::to_owned)); - } - let parameters = parameter_values.join(", "); - let arguments = argument_values.join(", "); - let result = format!( - "{{ {} }}", - std::iter::repeat_n(result_type, d_count) - .collect::>() - .join(", ") - ); - writeln!( - output, - "define {result} @probe_{}({parameters}) #0 {{", - record.id - ) - .unwrap(); - writeln!( - output, - " %result = call {result} asm sideeffect {:?}, {:?}({arguments}) #0", - register_mma_template(record), - register_mma_constraints(record), - ) - .unwrap(); - output.push_str(" ret "); - output.push_str(&result); - output.push_str(" %result\n}\n\nattributes #0 = { convergent }\n"); - } else if let Some(debug) = &record.debug_control { - let template = match debug.operation { - DebugControlOperation::Trap => "trap;", - DebugControlOperation::Breakpoint => "brkpt;", - DebugControlOperation::Pmevent => "pmevent 15;", - }; - writeln!(output, "define void @probe_{}() {{", record.id).unwrap(); - writeln!(output, " call void asm sideeffect {:?}, \"\"()", template).unwrap(); - output.push_str(" ret void\n}\n"); - } else if record.family == "stmatrix" { - let (multiplicity, _) = stmatrix_variant(record).expect("stmatrix variant"); - let count = multiplicity.register_count(); - let symbol = llvm(record) - .resolved_symbol - .as_ref() - .expect("stmatrix resolved symbol"); - let declaration = std::iter::once("ptr addrspace(3)".to_owned()) - .chain(std::iter::repeat_n("i32".to_owned(), count)) - .collect::>() - .join(", "); - let parameters = std::iter::once("ptr %generic".to_owned()) - .chain((0..count).map(|index| format!("i32 %r{index}"))) - .collect::>() - .join(", "); - let arguments = std::iter::once("ptr addrspace(3) %shared".to_owned()) - .chain((0..count).map(|index| format!("i32 %r{index}"))) - .collect::>() - .join(", "); - writeln!(output, "declare void @{symbol}({declaration})\n").unwrap(); - writeln!(output, "define void @probe_{}({parameters}) {{", record.id).unwrap(); - output.push_str(" %shared = addrspacecast ptr %generic to ptr addrspace(3)\n"); - writeln!(output, " call void @{symbol}({arguments})").unwrap(); - output.push_str(" ret void\n}\n"); - } else if record.prmt.is_some() { - let arity = record.rust.arguments.len(); - let declaration = std::iter::repeat_n("i32", arity) - .collect::>() - .join(", "); - let parameters = (0..arity) - .map(|index| format!("i32 %arg{index}")) - .collect::>() - .join(", "); - let arguments = (0..arity) - .map(|index| format!("i32 %arg{index}")) - .collect::>() - .join(", "); - writeln!( - output, - "declare i32 @{}({declaration})\n", - llvm(record).symbol - ) - .unwrap(); - writeln!(output, "define i32 @probe_{}({parameters}) {{", record.id).unwrap(); - writeln!( - output, - " %result = call i32 @{}({arguments})", - llvm(record).symbol - ) - .unwrap(); - output.push_str(" ret i32 %result\n}\n"); - } else if let Some(ldmatrix) = &record.ldmatrix { - let register_count = ldmatrix.variant.register_count(); - let result_ty = if register_count == 1 { - "i32".to_owned() - } else { - format!( - "{{ {} }}", - std::iter::repeat_n("i32", register_count) - .collect::>() - .join(", ") - ) - }; - let symbol = record - .llvm - .as_ref() - .expect("ldmatrix LLVM facts") - .resolved_symbol - .as_ref() - .expect("ldmatrix resolved symbol"); - writeln!(output, "declare {result_ty} @{symbol}(ptr addrspace(3))").unwrap(); - output.push('\n'); - writeln!( - output, - "define {result_ty} @probe_{}(ptr %generic) {{", - record.id - ) - .unwrap(); - output.push_str(" %shared = addrspacecast ptr %generic to ptr addrspace(3)\n"); - writeln!( - output, - " %result = call {result_ty} @{symbol}(ptr addrspace(3) %shared)" - ) - .unwrap(); - writeln!(output, " ret {result_ty} %result\n}}\n").unwrap(); - } else if record.scalar_arithmetic.is_some() { - let ty = scalar_arithmetic_llvm_type(record); - let arity = scalar_arithmetic_arity(record); - let parameters = (0..arity) - .map(|index| format!("{ty} %arg{index}")) - .collect::>() - .join(", "); - let arguments = (0..arity) - .map(|index| format!("{ty} %arg{index}")) - .collect::>() - .join(", "); - if scalar_arithmetic_llvm_mechanism(record) == BackendLoweringMechanism::TypedNvvm { - let declaration = std::iter::repeat_n(ty, arity) - .collect::>() - .join(", "); - writeln!( - output, - "declare {ty} @{}({declaration})\n", - llvm(record).symbol - ) - .unwrap(); - } - writeln!(output, "define {ty} @probe_{}({parameters}) {{", record.id).unwrap(); - match scalar_arithmetic_llvm_mechanism(record) { - BackendLoweringMechanism::TypedNvvm => { - writeln!( - output, - " %result = call {ty} @{}({arguments})", - llvm(record).symbol - ) - .unwrap(); - } - BackendLoweringMechanism::InlinePtx => { - let register = if ty == "double" { "d" } else { "f" }; - let constraints = std::iter::once(format!("={register}")) - .chain(std::iter::repeat_n(register.to_owned(), arity)) - .collect::>() - .join(","); - let operands = (0..=arity) - .map(|index| format!("${index}")) - .collect::>() - .join(", "); - writeln!( - output, - " %result = call {ty} asm {:?}, {:?}({arguments})", - format!("{} {operands};", scalar_arithmetic_ptx_mnemonic(record)), - constraints, - ) - .unwrap(); - } - } - writeln!(output, " ret {ty} %result\n}}").unwrap(); - } else if record.scalar_math.is_some() { - let ty = scalar_math_llvm_type(record); - let parameters = format!("{ty} %arg0"); - let arguments = format!("{ty} %arg0"); - if scalar_math_llvm_mechanism(record) == BackendLoweringMechanism::TypedNvvm { - writeln!(output, "declare {ty} @{}({ty})\n", llvm(record).symbol).unwrap(); - } - writeln!(output, "define {ty} @probe_{}({parameters}) {{", record.id).unwrap(); - match scalar_math_llvm_mechanism(record) { - BackendLoweringMechanism::TypedNvvm => { - writeln!( - output, - " %result = call {ty} @{}({arguments})", - llvm(record).symbol - ) - .unwrap(); - } - BackendLoweringMechanism::InlinePtx => { - let register = match ty { - "double" => "d", - "i16" => "h", - _ => "f", - }; - writeln!( - output, - " %result = call {ty} asm {:?}, {:?}({arguments})", - format!("{} $0, $1;", scalar_math_ptx_mnemonic(record)), - format!("={register},{register}"), - ) - .unwrap(); - } - } - writeln!(output, " ret {ty} %result\n}}").unwrap(); - } else if record.extended_minmax.is_some() { - let (ty, register) = match extended_minmax_contract(record).format { - ExtendedMinMaxFormat::F32 => ("float", "f"), - ExtendedMinMaxFormat::F16 | ExtendedMinMaxFormat::Bf16 => ("i16", "h"), - ExtendedMinMaxFormat::F16x2 | ExtendedMinMaxFormat::Bf16x2 => ("i32", "r"), - }; - writeln!( - output, - "define {ty} @probe_{}({ty} %a, {ty} %b) {{", - record.id - ) - .unwrap(); - writeln!( - output, - " %result = call {ty} asm {:?}, {:?}({ty} %a, {ty} %b)", - format!("{} $0, $1, $2;", extended_minmax_ptx_mnemonic(record)), - format!("={register},{register},{register}"), - ) - .unwrap(); - writeln!(output, " ret {ty} %result\n}}").unwrap(); - } else if record.scalar_conversion.is_some() { - writeln!(output, "declare i32 @{}(float)", llvm(record).symbol).unwrap(); - output.push('\n'); - writeln!(output, "define i32 @probe_{}(float %value) {{", record.id).unwrap(); - writeln!( - output, - " %result = call i32 @{}(float %value)", - llvm(record).symbol - ) - .unwrap(); - output.push_str(" ret i32 %result\n}\n"); - } else if let Some(width) = record.scalar_width() { - writeln!(output, "declare i{width} @{}()", llvm(record).symbol).unwrap(); - output.push('\n'); - writeln!(output, "define i{width} @probe_{}() {{", record.id).unwrap(); - writeln!( - output, - " %result = call i{width} @{}()", - llvm(record).symbol - ) - .unwrap(); - writeln!(output, " ret i{width} %result\n}}\n").unwrap(); - } else { - unreachable!("generated intrinsic has no probe renderer: {}", record.id); - } - output -} - -/// Escape a value for a GFM table cell. -/// -/// A `|` splits the row before inline parsing runs, so it splits the cell even -/// inside a code span. GFM's own rule is that a literal pipe must be written -/// `\|` wherever it appears. Three catalog entries carry a literal pipe today: -/// the `elect.sync` and `match.all.sync` PTX patterns write their two -/// destinations as ``. Their rows rendered with a seventh -/// cell, which dropped the backend-evidence column. -fn escape_table_cell(value: impl std::fmt::Display) -> String { - value.to_string().replace('|', "\\|") -} - -fn render_reference(catalog: &CatalogFile, hash: &str) -> String { - let mut output = markdown_header(catalog, hash); - output.push_str( - "# Generated Intrinsic Reference\n\nThis table is generated from the resolved catalog. Evidence stages below distinguish backend code generation, terminal validation, and GPU runtime execution; no runtime claim is made unless an executed stage is recorded.\n\n| Rust function | CUDA operation | Source | Effects | Availability | Backend evidence |\n|:--|:--|:--|:--|:--|:--|\n", - ); - for record in &catalog.intrinsics { - let safety = if record.rust.safe { "safe" } else { "unsafe" }; - let availability = if is_blackwell_ldmatrix(record) { - format!( - "Instruction floor PTX {} on {}; effective LLVM target floors: {BLACKWELL_LDMATRIX_EFFECTIVE_FLOORS}", - record.target.minimum_ptx, - hardware_target_label(&record.target.hardware) - ) - } else { - format!( - "PTX {} on {}", - record.target.minimum_ptx, - hardware_target_label(&record.target.hardware) - ) - }; - writeln!( - output, - "| `{}` | `{}` | {} | {safety}; scope {}; {}; memory {}; convergent {} | {availability} ([PTX ISA {}, {}]({})) | {} on `{}`; expects `{}` (`{}`, SHA-256 `{}`) |", - escape_table_cell(&record.rust.public_path), - escape_table_cell(&record.dialect.op_name), - escape_table_cell(source_label(record)), - escape_table_cell(&record.semantics.execution_scope), - if record.semantics.pure { "pure" } else { "impure" }, - escape_table_cell(&record.semantics.memory), - record.semantics.convergent, - escape_table_cell(&record.target.ptx_isa_version), - escape_table_cell(&record.target.ptx_isa_section), - escape_table_cell(&record.target.ptx_isa_url), - escape_table_cell(&record.backend.status), - escape_table_cell(&record.backend.profile), - escape_table_cell(&record.expected_ptx), - escape_table_cell(&record.backend.version), - escape_table_cell(&record.backend.sha256), - ) - .unwrap(); - } - output.push_str("\n## Compiler identity and compatibility paths\n\n"); - for record in &catalog.intrinsics { - writeln!(output, "- `{}` (`{}`)", record.id, record.rust.abi_id).unwrap(); - writeln!( - output, - " - canonical compiler path: `{}`", - record.rust.canonical_path - ) - .unwrap(); - writeln!( - output, - " - public source path: `{}`", - record.rust.public_path - ) - .unwrap(); - for path in &record.rust.compatibility_paths { - writeln!(output, " - compatibility compiler path: `{path}`").unwrap(); - } - } - output.push_str("\n## Register-MMA contracts\n\n"); - for record in register_mmas(catalog) { - let mma = record.register_mma.as_ref().unwrap(); - let operation = match mma.operation { - RegisterMmaOperation::Multiply => "multiply-accumulate", - RegisterMmaOperation::AndPopc => "AND, population count, and accumulate", - RegisterMmaOperation::XorPopc => "XOR, population count, and accumulate", - }; - let overflow = match mma.overflow { - RegisterMmaOverflow::NotApplicable => "not applicable", - RegisterMmaOverflow::Wrapping => "wrapping", - RegisterMmaOverflow::Satfinite => "finite saturation", - }; - let runtime = match mma.runtime_validation { - RuntimeValidation::Unexecuted => "not executed on a GPU", - RuntimeValidation::Executed => "executed on a GPU", - }; - if mma.adapter == RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 { - writeln!( - output, - "- `{}` takes C, A, B, scale-A data/selectors, and scale-B data/selectors in PTX operand order, performs {operation}, and lowers to one convergent, register-only `{}` instruction. For `scale_vec::1X`, byte selectors must be in `0..=3`, the A thread selector in `0..=1`, and the B thread selector in `0..=3`. Every non-exited warp lane must execute the same instruction and qualifiers. Integer overflow is {overflow}; runtime validation is {runtime}.", - record.id, - expected_ptx_head(record), - ) - .unwrap(); - } else { - writeln!( - output, - "- `{}` takes fragments in C, A, B order, performs {operation}, and lowers to one convergent, register-only `{}` instruction. Every non-exited warp lane must execute the same instruction and qualifiers. Integer overflow is {overflow}; runtime validation is {runtime}.", - record.id, - expected_ptx_head(record), - ) - .unwrap(); - } - } - output.push_str("\n## Sparse-MMA contracts\n\n"); - for record in sparse_mmas(catalog) { - let mma = record.sparse_mma.as_ref().unwrap(); - let overflow = match mma.overflow { - SparseMmaOverflow::NotApplicable => "Overflow mode is not applicable.", - SparseMmaOverflow::Wrapping => "Integer overflow wraps.", - SparseMmaOverflow::Satfinite => "Integer overflow uses finite saturation.", - }; - let runtime = match mma.runtime_validation { - RuntimeValidation::Unexecuted => "not executed on a GPU", - RuntimeValidation::Executed => "executed on a GPU", - }; - let metadata = sparse_mma_metadata_rule(mma); - writeln!( - output, - "- `{}` takes fragments in C, A, B, metadata, selector order and lowers to one convergent, register-only `{}` instruction. Its LLVM source record uses A, B, C, metadata, selector order. The selector must be {}. Every non-exited warp lane must execute the same instruction and qualifiers. {metadata} {overflow} Runtime validation is {runtime}.", - record.id, - sparse_mma_ptx_head(record), - sparse_mma_selector_description(record), - ) - .unwrap(); - } - output.push_str("\n## Packed-atomic contracts\n\n"); - for record in packed_atomics(catalog) { - let packed = record.packed_atomic.as_ref().unwrap(); - let format = match packed.format { - PackedAtomicFormat::F16x2 => "f16x2", - PackedAtomicFormat::Bf16x2 => "bf16x2", - }; - writeln!( - output, - "- `{}` (`{format}`): the native PTX instruction starts at `sm_{}`; cuda-oxide admits it from {}. Omitted `.sem` / `.scope` mean relaxed GPU scope. Each 16-bit element rounds to nearest-even and `.noftz` preserves subnormal inputs and results. The elements are atomic independently, so the returned `u32` contains old per-element values that may not form one coherent pair. The pointer must address four writable, four-byte-aligned global bytes; do not mix whole-word or non-atomic overlapping access, and every racing atomic must use a mutually inclusive scope.", - record.id, - packed.native_minimum_sm, - hardware_target_label(&record.target.hardware), - ) - .unwrap(); - } - output.push_str("\n## Redux contracts\n\n"); - for record in redux(catalog) { - writeln!( - output, - "- `{}`: raw and dialect operands are `[member_mask, value]`, adapted to LLVM `(value, membermask)`. The executing lane must be named in the mask, and every non-exited named lane must execute the same instruction with the same qualifiers and mask.", - record.id, - ) - .unwrap(); - } - output.push_str("\n## Packed-ALU contracts\n\n"); - for record in packed_alus(catalog) { - let packed = record.packed_alu.as_ref().unwrap(); - let (format, _, carrier, _, _) = packed_alu_format_shape(packed.format); - let backend_floors = record - .backend_lowerings - .iter() - .map(|lowering| { - let backend = match lowering.backend { - IntrinsicBackend::LlvmNvptx => "LLVM-NVPTX", - IntrinsicBackend::LibNvvm => "libNVVM", - }; - format!( - "{backend} PTX {} on {}", - lowering.target.minimum_ptx, - hardware_target_label(&lowering.target.hardware), - ) - }) - .collect::>() - .join("; "); - writeln!( - output, - "- `{}` carries one packed `{format}` value in a `{carrier}` and lowers to one pure `{}` instruction. The native instruction starts at PTX {} / `sm_{}`; cuda-oxide admits it from {}. Backend profile floors: {backend_floors}.", - record.id, - packed_alu_ptx_mnemonic(record), - record.target.minimum_ptx, - packed.native_minimum_sm, - hardware_target_label(&record.target.hardware), - ) - .unwrap(); - } - output.push_str("\n## Packed-conversion contract\n\n"); - for record in packed_conversions(catalog) { - let conversion = record.packed_conversion.as_ref().unwrap(); - let rounding = match conversion.rounding { - PackedConversionRounding::NearestEven => "nearest-even", - PackedConversionRounding::TowardZero => "toward-zero", - }; - let saturation = match conversion.saturation { - PackedConversionSaturation::None => "without saturation", - PackedConversionSaturation::Relu => "with ReLU", - PackedConversionSaturation::Satfinite => "with finite saturation", - PackedConversionSaturation::SatfiniteRelu => "with finite saturation and ReLU", - }; - if packed_conversion_typed_llvm_name(record).is_some() { - writeln!( - output, - "- `{}` converts two `f32` inputs to packed `{}` using {rounding} rounding {saturation}. LLVM-NVPTX uses typed `{}` with `[high, low]` inputs; libNVVM uses pure `{}` inline PTX. The public first input becomes the low lane and the second becomes the high lane.", - record.id, - packed_conversion_destination(record), - llvm(record).symbol, - packed_conversion_ptx_mnemonic(record), - ) - .unwrap(); - } else { - writeln!( - output, - "- `{}` converts two `f32` inputs to packed `{}` using {rounding} rounding {saturation}. It lowers to pure `{}` inline PTX. The first input becomes the low lane and the second becomes the high lane, so PTX prints the inputs in reverse order.", - record.id, - packed_conversion_destination(record), - packed_conversion_ptx_mnemonic(record), - ) - .unwrap(); - } - } - output.push_str("\n## Scalar TF32-conversion contract\n\n"); - for record in scalar_conversions(catalog) { - writeln!( - output, - "- `{}` converts one `f32` value to raw TF32 bits with `{}`. LLVM-NVPTX uses the typed intrinsic; libNVVM uses pure inline PTX so the native target floor is preserved.", - record.id, - scalar_conversion_ptx_mnemonic(record), - ) - .unwrap(); - } - output.push_str("\n## Explicit-rounding scalar-arithmetic contracts\n\n"); - for record in scalar_arithmetics(catalog) { - let llvm_route = match scalar_arithmetic_llvm_mechanism(record) { - BackendLoweringMechanism::TypedNvvm => "the typed intrinsic", - BackendLoweringMechanism::InlinePtx => "pure inline PTX", - }; - writeln!( - output, - "- `{}` performs `{}`. LLVM-NVPTX uses {llvm_route}; libNVVM uses pure inline PTX.", - record.id, - scalar_arithmetic_ptx_mnemonic(record), - ) - .unwrap(); - } - output.push_str("\n## Warp vote contracts\n\n"); - for record in vote_intrinsics(catalog) { - writeln!( - output, - "- `{}` keeps raw and dialect operands in `[member_mask, predicate]` order. The executing lane must be named in the mask, and every non-exited named lane must execute the same `vote.sync` instruction with the same qualifiers and mask. On `sm_6x` and earlier, all named lanes must execute in convergence and no unnamed lane may be active. Both immediate and register member masks are admitted.", - record.id, - ) - .unwrap(); - } - output.push_str("\n## Active-mask contract\n\n"); - for record in active_masks(catalog) { - writeln!( - output, - "- `{}` observes the lanes active at the instruction. LLVM uses the typed intrinsic; libNVVM uses reviewed convergent, side-effecting inline PTX because that backend does not select the intrinsic.", - record.id, - ) - .unwrap(); - } - output.push_str("\n## Warp-match contracts\n\n"); - for record in warp_matches(catalog) { - let adapter = match record.warp_match.as_ref().unwrap().adapter { - WarpMatchAdapter::DirectMask => "returns LLVM's mask directly", - WarpMatchAdapter::ProjectMaskDiscardPredicate => { - "projects field 0 from LLVM's `{i32, i1}` result" - } - }; - writeln!( - output, - "- `{}` keeps operands in `[member_mask, value]` order and {adapter}. The executing lane must be named in the mask, and every non-exited named lane must execute the same `match.sync` operation with the same qualifiers and mask. All register/immediate value and mask forms are admitted.", - record.id, - ) - .unwrap(); - } - output.push_str("\n## Warp-barrier contract\n\n"); - for record in warp_barriers(catalog) { - writeln!( - output, - "- `{}` passes the 32-bit member mask directly to the typed LLVM intrinsic on both backends. The executing lane must be named in the mask, and every non-exited named lane must execute the same `bar.warp.sync` operation with the same mask. On `sm_6x` and earlier, all named lanes must execute the barrier in convergence, and no unnamed lane may be active when it executes. The barrier orders memory accesses among participating lanes. Both immediate and register masks are admitted.", - record.id, - ) - .unwrap(); - } - output.push_str("\n## Warp-shuffle contracts\n\n"); - for record in warp_shuffles(catalog) { - let shuffle = record.warp_shuffle.as_ref().unwrap(); - let lowering = match shuffle.value_kind { - WarpShuffleValueKind::I32 | WarpShuffleValueKind::F32 => { - "Register and immediate lane/mask forms are admitted." - } - WarpShuffleValueKind::I64 => { - "One convergent, side-effecting inline-PTX block splits `i64`, performs low then high `b32` shuffles with the same register lane/mask, and reassembles it." - } - }; - writeln!( - output, - "- `{}` keeps raw and dialect operands in `[member_mask, value, lane_or_delta]` order and inserts clamp `{}` during lowering. The executing lane must be named in the mask, and every non-exited named lane must execute the same `shfl.sync` operation with the same qualifiers and mask. On `sm_6x` and earlier, all named lanes must execute in convergence and no unnamed lane may be active. A computed in-range source must be active and named; if PTX marks it out of range, the calling lane's input is copied. {lowering}", - record.id, - shuffle.clamp, - ) - .unwrap(); - } - output.push_str("\n## Cluster-barrier contracts\n\n"); - for record in cluster_barriers(catalog) { - let barrier = record.cluster_barrier.as_ref().unwrap(); - let ordering = match barrier.ordering { - ClusterBarrierOrdering::Release => "release", - ClusterBarrierOrdering::Relaxed => "relaxed", - ClusterBarrierOrdering::Acquire => "acquire", - }; - let alignment = if barrier.aligned { - "Every non-exited warp thread must execute the same aligned instruction in identical control flow." - } else { - "The instruction is not aligned." - }; - writeln!( - output, - "- `{}` has {ordering} ordering. Each non-exited cluster thread must arrive exactly once before completion, then execute the matching wait. {alignment}", - record.id, - ) - .unwrap(); - } - if cluster_barriers(catalog).next().is_some() { - output.push_str( - "- `cuda_device::cluster::cluster_sync` is the generated compatibility operation: aligned arrive followed by aligned wait.\n", - ); - } - if cluster_memory(catalog).next().is_some() { - output.push_str("\n## Cluster-memory contracts\n\n"); - for record in cluster_memory(catalog) { - match record.cluster_memory.as_ref().unwrap().source_contract { - ClusterMemorySourceContract::LlvmMapaSharedClusterAs7IdentityInlinePtx => { - writeln!( - output, - "- `{}` preserves LLVM 22's address-space-7 result. Both backends use exact convergent `mapa.shared::cluster.u64` inline PTX, and ordinary Rust dereference of the mapped pointer lowers through cluster shared memory.", - record.id - ) - .unwrap(); - } - ClusterMemorySourceContract::PtxNativeMapaThenWeakClusterLoad => { - writeln!( - output, - "- `{}` has no one-to-one LLVM intrinsic. Both backends use one convergent `mapa.shared::cluster.u64` plus `ld.shared::cluster.u32` block with a compiler memory clobber.", - record.id - ) - .unwrap(); - } - } - } - } - if movmatrix(catalog).next().is_some() { - output.push_str("\n## movmatrix contracts\n\n"); - for record in movmatrix(catalog) { - writeln!( - output, - "- `{}` transposes one packed b16 fragment in registers. All warp lanes must execute the same instruction, and no lane may have exited.", - record.id, - ) - .unwrap(); - } - } - if mbarrier_extended(catalog).next().is_some() { - output.push_str("\n## Extended mbarrier contracts\n\n"); - for record in mbarrier_extended(catalog) { - let contract = record.mbarrier_extended.as_ref().unwrap(); - let source = match contract.source_contract { - MbarrierExtendedSourceContract::LlvmImported => "LLVM 22 TableGen", - MbarrierExtendedSourceContract::PtxNativeRawClusterAddress => { - "PTX-native raw-address contract" - } - }; - writeln!( - output, - "- `{}` keeps its established CUDA-device signature and exact convergent inline-PTX memory clobber. Its source identity is {source}.", - record.id, - ) - .unwrap(); - } - } - if wgmma_controls(catalog).next().is_some() { - output.push_str("\n## WGMMA-control contracts\n\n"); - for record in wgmma_controls(catalog) { - let detail = match record.wgmma_control.as_ref().unwrap().mode { - WgmmaControlMode::Fence => { - "It orders register writes before later WGMMA operations." - } - WgmmaControlMode::CommitGroup => { - "It commits the warpgroup's prior uncommitted WGMMA operations." - } - WgmmaControlMode::WaitGroup => { - "Its `u64` argument must be a compile-time constant." - } - }; - writeln!( - output, - "- `{}` must be executed by all four warps in the warpgroup with the same control flow. {detail}", - record.id, - ) - .unwrap(); - } - } - if sync_intrinsics(catalog).any(|record| threadfence_ptx_level(record).is_some()) { - output.push_str("\n## Synchronization contracts\n\n"); - } else { - output.push_str("\n## CTA synchronization contracts\n\n"); - } - for record in sync_intrinsics(catalog) { - if record.id == "sync_threads" { - writeln!( - output, - "- `{}` inserts the fixed barrier ID `0`. Every active CTA thread must reach the same barrier; divergent use can deadlock the CTA.", - record.id, - ) - .unwrap(); - } else { - let level = threadfence_ptx_level(record) - .expect("resolver admitted an unknown generated synchronization record"); - writeln!( - output, - "- `{}` emits `membar.{level}` through the reviewed typed NVVM route on both backends.", - record.id, - ) - .unwrap(); - } - } - output.push_str("\n## cp.async mbarrier contracts\n\n"); - for record in cp_async_mbarriers(catalog) { - let bridge = record.cp_async_mbarrier.as_ref().unwrap(); - let address = match bridge.state_space { - CpAsyncMbarrierStateSpace::Generic => "generic", - CpAsyncMbarrierStateSpace::Shared => "explicit shared", - }; - let counting = match bridge.operation { - CpAsyncMbarrierOperation::Arrive => { - "increments the pending count before the later asynchronous decrement" - } - CpAsyncMbarrierOperation::ArriveNoInc => { - "does not increment the pending count, so initialization must already include the later asynchronous decrement" - } - }; - writeln!( - output, - "- `{}` uses {address} addressing and {counting}. It associates only the executing thread's prior `cp.async` operations with a live, eight-byte-aligned shared-memory barrier. LLVM-NVPTX uses the typed intrinsic; libNVVM uses reviewed convergent, side-effecting inline PTX with a memory clobber.", - record.id, - ) - .unwrap(); - } - output.push_str("\n## Basic mbarrier contracts\n\n"); - for record in mbarrier_basics(catalog) { - let mbarrier = record.mbarrier_basic.as_ref().unwrap(); - let contract = match mbarrier.operation { - MbarrierBasicOperation::Init => { - "initializes one eight-byte-aligned shared-memory object exactly once with an expected count in `1..=0xFFFFF`" - } - MbarrierBasicOperation::Arrive => { - "records one expected arrival and returns a token for the same barrier phase" - } - MbarrierBasicOperation::TestWait => { - "tests one token from the same barrier phase with convergent, side-effecting inline PTX" - } - MbarrierBasicOperation::Inval => { - "invalidates an initialized barrier only after all users are finished" - } - }; - let mechanism = if mbarrier.operation == MbarrierBasicOperation::TestWait { - "inline PTX on both backends" - } else { - "the typed NVVM intrinsic with LLVM-NVPTX and inline PTX with libNVVM" - }; - writeln!( - output, - "- `{}` {contract}. It lowers through {mechanism}; generic MIR pointers are normalized to shared address space during lowering.", - record.id, - ) - .unwrap(); - } - output.push_str("\n## Imported LLVM properties and result facts\n\n"); - for record in &catalog.intrinsics { - let Some(llvm) = &record.llvm else { - writeln!( - output, - "- `{}`: PTX-native source; no LLVM record, symbol, properties, or selection facts.", - record.id - ) - .unwrap(); - continue; - }; - let properties = record - .llvm - .as_ref() - .unwrap() - .properties - .iter() - .map(|property| format!("`{property}`")) - .collect::>() - .join(", "); - let range = record - .llvm - .as_ref() - .unwrap() - .result_facts - .range - .as_ref() - .map(|range| format!("[{}, {})", range.lower, range.upper_exclusive)) - .unwrap_or_else(|| "none".to_owned()); - let selection_records = record - .selections - .iter() - .map(|selection| format!("`{}`", selection.source_record)) - .collect::>() - .into_iter() - .collect::>() - .join(", "); - let selection_predicates = record - .selections - .iter() - .flat_map(|selection| selection.predicates.iter()) - .map(|predicate| format!("`{predicate}`")) - .collect::>() - .into_iter() - .collect::>() - .join(", "); - writeln!( - output, - "- `{}`: properties {}; result `NoUndef` {}; half-open result range `{}`; selection records {}; selection predicates {}", - record.id, - properties, - llvm.result_facts.no_undef, - range, - selection_records, - if selection_predicates.is_empty() { - "none" - } else { - &selection_predicates - } - ) - .unwrap(); - } - output.push_str("\n## Backend-specific lowering evidence\n\n"); - for record in &catalog.intrinsics { - if record.backend_lowerings.is_empty() { - continue; - } - let runtime = record - .ldmatrix - .as_ref() - .map(|record| format!("{:?}", record.safety.runtime_validation).to_lowercase()) - .or_else(|| { - record - .packed_atomic - .as_ref() - .map(|record| format!("{:?}", record.runtime_validation).to_lowercase()) - }) - .or_else(|| { - record - .mbarrier_basic - .as_ref() - .map(|record| format!("{:?}", record.runtime_validation).to_lowercase()) - }) - .or_else(|| { - record - .movmatrix - .as_ref() - .map(|record| format!("{:?}", record.runtime_validation).to_lowercase()) - }) - .or_else(|| { - record - .mbarrier_extended - .as_ref() - .map(|record| format!("{:?}", record.runtime_validation).to_lowercase()) - }) - .or_else(|| { - record - .cp_async_mbarrier - .as_ref() - .map(|record| format!("{:?}", record.runtime_validation).to_lowercase()) - }) - .or_else(|| { - record - .register_mma - .as_ref() - .map(|record| format!("{:?}", record.runtime_validation).to_lowercase()) - }) - .or_else(|| { - record - .sparse_mma - .as_ref() - .map(|record| format!("{:?}", record.runtime_validation).to_lowercase()) - }) - .or_else(|| { - record - .debug_control - .as_ref() - .map(|record| format!("{:?}", record.runtime_validation).to_lowercase()) - }) - .or_else(|| { - record - .cluster_memory - .as_ref() - .map(|record| format!("{:?}", record.runtime_validation).to_lowercase()) - }) - .or_else(|| { - record - .tcgen05 - .as_ref() - .map(|record| format!("{:?}", record.runtime_validation).to_lowercase()) - }) - .or_else(|| (record.family == "stmatrix").then(|| "unexecuted".to_owned())) - .unwrap_or_else(|| "not recorded".to_owned()); - writeln!(output, "- `{}`: runtime `{runtime}`", record.id).unwrap(); - for lowering in &record.backend_lowerings { - writeln!( - output, - " - `{}` uses `{}` from profile `{}` at PTX {} / {}: status `{}` (`{}`, SHA-256 `{}`)", - backend_label(lowering.backend), - lowering_mechanism_label(lowering.mechanism), - lowering.evidence_profile, - lowering.target.minimum_ptx, - hardware_target_label(&lowering.target.hardware), - lowering.status, - lowering.version, - lowering.sha256, - ) - .unwrap(); - for stage in &lowering.stages { - let tool = match ( - stage.tool_path.as_deref(), - stage.tool_version.as_deref(), - stage.tool_sha256.as_deref(), - ) { - (Some(path), Some(version), Some(sha256)) => { - format!(" Tool `{path}` reports `{version}` (SHA-256 `{sha256}`).") - } - _ => String::new(), - }; - let artifact = match stage.artifact_kind { - Some(EvidenceArtifactKind::Cubin) => " Artifact `cubin`.", - None => "", - }; - writeln!( - output, - " - {} on `{}`: `{}` — {}{}{}", - evidence_stage_label(stage.stage), - stage.targets.join(", "), - stage.outcome, - stage.detail, - tool, - artifact, - ) - .unwrap(); - } - } - } - output -} - -fn render_compiler_path_patterns(output: &mut String, catalog: &CatalogFile, indent: &str) { - let paths: Vec<_> = catalog - .intrinsics - .iter() - .flat_map(|record| { - std::iter::once(record.rust.canonical_path.as_str()) - .chain(record.rust.compatibility_paths.iter().map(String::as_str)) - }) - .collect(); - render_string_patterns(output, &paths, indent); -} - -fn render_string_patterns(output: &mut String, values: &[&str], indent: &str) { - for (index, value) in values.iter().enumerate() { - if index == 0 { - writeln!(output, "{indent}{value:?}").unwrap(); - } else { - writeln!(output, "{indent}| {value:?}").unwrap(); - } - } -} - -fn render_inline_patterns(output: &mut String, values: &[&str]) { - for (index, value) in values.iter().enumerate() { - if index != 0 { - output.push_str(" | "); - } - write!(output, "{value:?}").unwrap(); - } -} - -#[allow(dead_code)] -fn modules(catalog: &CatalogFile) -> BTreeSet<&str> { - catalog - .intrinsics - .iter() - .map(|record| record.rust.module.as_str()) - .collect() -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::model::{ - CatalogTargetAlternative, ImportedSelectionConstraints, TargetSelectorBinding, - }; - use crate::ptx::{InstructionPattern, OperandPattern}; - use crate::util::read_json; - use std::path::Path; - - #[test] - fn paired_target_matrix_renders_selectors_and_ptx_floors() { - let target = CatalogHardwareTarget::TargetMatrix { - contracts: vec![CatalogTargetContract { - selectors: vec![TargetSelectorBinding { - name: "kind".into(), - value: "i8".into(), - }], - alternatives: vec![ - CatalogTargetAlternative { - minimum_ptx: "8.8".parse().unwrap(), - hardware: CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, - }, - CatalogTargetAlternative { - minimum_ptx: "9.0".parse().unwrap(), - hardware: CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, - }, - ], - }], - }; - - assert_eq!( - hardware_target_label(&target), - "sm_100a at PTX 8.8 or sm_110a at PTX 9.0 for kind=i8" - ); - assert_eq!( - generated_hardware_target(&target), - "GeneratedHardwareTarget::TargetMatrix { contracts: &[GeneratedTargetContract { selectors: &[GeneratedTargetSelectorBinding { name: \"kind\", value: \"i8\" }], alternatives: &[GeneratedTargetAlternative { minimum_ptx: GeneratedPtxVersion::from_encoded(88), hardware: GeneratedHardwareAlternative::ExactArchitecture(100) }, GeneratedTargetAlternative { minimum_ptx: GeneratedPtxVersion::from_encoded(90), hardware: GeneratedHardwareAlternative::ExactArchitecture(110) }] }] }" - ); - } - - #[test] - fn tcgen05_mma_inline_asm_closes_selectors_and_operand_order() { - for (kind, spelling) in [ - (Tcgen05MmaKind::F16, "kind::f16"), - (Tcgen05MmaKind::Tf32, "kind::tf32"), - (Tcgen05MmaKind::F8f6f4, "kind::f8f6f4"), - (Tcgen05MmaKind::I8, "kind::i8"), - ] { - let (template, _) = tcgen05_mma_inline_asm( - Tcgen05MmaForm::WsTensor, - kind, - 1, - None, - Some(0), - Some(Tcgen05MmaBUsage::Discard), - ); - assert!(template.contains(spelling)); - } - for (usage, spelling) in [ - (Tcgen05MmaBUsage::Discard, "collector::b0::discard"), - (Tcgen05MmaBUsage::LastUse, "collector::b0::lastuse"), - (Tcgen05MmaBUsage::Fill, "collector::b0::fill"), - (Tcgen05MmaBUsage::Use, "collector::b0::use"), - ] { - let (template, _) = tcgen05_mma_inline_asm( - Tcgen05MmaForm::WsTensor, - Tcgen05MmaKind::F16, - 1, - None, - Some(0), - Some(usage), - ); - assert!(template.contains(spelling)); - } - - for (form, operands, constraints) in [ - ( - Tcgen05MmaForm::WsTensor, - "[$0], [$1], $2, $3, %enable_pred;", - "r,r,l,r,r,~{memory}", - ), - ( - Tcgen05MmaForm::WsSpTensor, - "[$0], [$1], $2, [$5], $3, %enable_pred;", - "r,r,l,r,r,r,~{memory}", - ), - ( - Tcgen05MmaForm::WsTensorZeroColMask, - "[$0], [$1], $2, $3, %enable_pred, $5;", - "r,r,l,r,r,l,~{memory}", - ), - ( - Tcgen05MmaForm::WsSpTensorZeroColMask, - "[$0], [$1], $2, [$5], $3, %enable_pred, $6;", - "r,r,l,r,r,r,l,~{memory}", - ), - ] { - let (template, actual_constraints) = tcgen05_mma_inline_asm( - form, - Tcgen05MmaKind::F16, - 1, - None, - Some(3), - Some(Tcgen05MmaBUsage::LastUse), - ); - assert!(template.contains(operands)); - assert_eq!(actual_constraints, constraints); - } - } - - fn catalog_with_debug_controls() -> CatalogFile { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - assert_eq!(debug_controls(&catalog).count(), 3); - catalog - } - - fn catalog_with_stmatrix() -> CatalogFile { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - assert_eq!(stmatrices(&catalog).count(), 4); - catalog - } - - fn catalog_with_clc() -> CatalogFile { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::test_catalog_with_clc(&repo_root).unwrap(); - assert_eq!(clc_intrinsics(&catalog).count(), 6); - catalog - } - - fn catalog_with_tma() -> CatalogFile { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::test_catalog_with_tma(&repo_root).unwrap(); - assert_eq!(tma_intrinsics(&catalog).count(), 111); - catalog - } - - fn catalog_with_wgmma_controls() -> CatalogFile { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let mut catalog = crate::resolve::resolve(&repo_root).unwrap(); - if wgmma_controls(&catalog).count() == 3 { - return catalog; - } - let template = catalog - .intrinsics - .iter() - .find(|record| record.cluster_barrier.is_some()) - .expect("cluster-barrier template") - .clone(); - - let recipes = [ - ( - WgmmaControlMode::Fence, - "i0317", - "wgmma_fence", - "WgmmaFenceSyncAlignedOp", - "nvvm.wgmma_fence_sync_aligned", - "llvm.nvvm.wgmma.fence.sync.aligned", - "WGMMA_FENCE_SYNC_ALIGNED", - "fence.sync.aligned", - WgmmaControlAdapter::NoArguments, - "cuda_device::wgmma::wgmma_fence", - ), - ( - WgmmaControlMode::CommitGroup, - "i0318", - "wgmma_commit_group", - "WgmmaCommitGroupSyncAlignedOp", - "nvvm.wgmma_commit_group_sync_aligned", - "llvm.nvvm.wgmma.commit_group.sync.aligned", - "WGMMA_COMMIT_GROUP_SYNC_ALIGNED", - "commit_group.sync.aligned", - WgmmaControlAdapter::NoArguments, - "cuda_device::wgmma::wgmma_commit_group", - ), - ( - WgmmaControlMode::WaitGroup, - "i0319", - "wgmma_wait_group", - "WgmmaWaitGroupSyncAlignedOp", - "nvvm.wgmma_wait_group_sync_aligned", - "llvm.nvvm.wgmma.wait_group.sync.aligned", - "WGMMA_WAIT_GROUP_SYNC_ALIGNED", - "wait_group.sync.aligned", - WgmmaControlAdapter::ConstGenericU32ToI64Immediate, - "cuda_device::wgmma::__wgmma_wait_group", - ), - ]; - - for ( - mode, - abi_id, - id, - op_type, - op_name, - llvm_symbol, - selection_record, - suffix, - adapter, - compatibility_path, - ) in recipes - { - let wait = mode == WgmmaControlMode::WaitGroup; - let mut record = template.clone(); - record.id = id.into(); - record.operation_key = format!("wgmma.control.{suffix}"); - record.family = "wgmma_control".into(); - record.source = crate::model::IntrinsicSource::LlvmImported { - source_record: format!("int_nvvm_wgmma_{}", suffix.replace('.', "_")), - }; - record.selections = vec![CatalogSelection { - source_record: selection_record.into(), - asm: format!("wgmma.{suffix};"), - predicates: vec!["Subtarget->getPTXVersion() >= 80".into(), "hasSM90a".into()], - constraints: ImportedSelectionConstraints::default(), - }]; - record.rust.abi_id = abi_id.into(); - record.rust.module = "wgmma".into(); - record.rust.name = id.into(); - record.rust.arguments = if wait { vec!["u64".into()] } else { vec![] }; - record.rust.result = "()".into(); - record.rust.safe = false; - record.rust.must_use = false; - record.rust.safe_allowlist_reason = None; - record.rust.canonical_path = - format!("cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::{abi_id}"); - record.rust.public_path = format!("cuda_intrinsics::wgmma::{id}"); - record.rust.compatibility_paths = vec![compatibility_path.into()]; - record.dialect.op_type = op_type.into(); - record.dialect.op_name = op_name.into(); - record.dialect.operands = if wait { vec!["i64".into()] } else { vec![] }; - record.dialect.results.clear(); - let llvm = record.llvm.as_mut().unwrap(); - llvm.symbol = llvm_symbol.into(); - llvm.resolved_symbol = None; - llvm.arguments = if wait { vec!["i64".into()] } else { vec![] }; - llvm.results.clear(); - llvm.properties = if wait { - vec!["ImmArg".into(), "IntrConvergent".into()] - } else { - vec!["IntrConvergent".into()] - }; - record.semantics.pure = false; - record.semantics.memory = "read_write".into(); - record.semantics.convergent = true; - record.semantics.execution_scope = "warpgroup".into(); - record.target.minimum_ptx = "8.0".parse().unwrap(); - record.target.hardware = CatalogHardwareTarget::AnyOf { - alternatives: vec![CatalogHardwareAlternative::ExactArchitecture { sm: 90 }], - }; - record.target.ptx_result = "()".into(); - record.target.targets = "sm_90a".into(); - record.target.ptx_isa_version = "9.3".into(); - record.target.ptx_isa_section = "WGMMA control".into(); - record.target.ptx_isa_url = "https://docs.nvidia.com/cuda/parallel-thread-execution/#asynchronous-warpgroup-level-matrix-instructions".into(); - for lowering in &mut record.backend_lowerings { - lowering.target.minimum_ptx = "8.0".parse().unwrap(); - lowering.target.hardware = record.target.hardware.clone(); - lowering.mechanism = match lowering.backend { - IntrinsicBackend::LlvmNvptx => BackendLoweringMechanism::TypedNvvm, - IntrinsicBackend::LibNvvm => BackendLoweringMechanism::InlinePtx, - }; - } - record.packed_atomic = None; - record.redux = None; - record.vote = None; - record.active_mask = None; - record.warp_match = None; - record.warp_barrier = None; - record.warp_shuffle = None; - record.dot_product = None; - record.packed_alu = None; - record.packed_conversion = None; - record.cp_async_copy = None; - record.cp_async_control = None; - record.cp_async_mbarrier = None; - record.mbarrier_basic = None; - record.register_mma = None; - record.sparse_mma = None; - record.prmt = None; - record.cluster_barrier = None; - record.wgmma_control = Some(crate::model::WgmmaControl { - mode, - adapter, - participation: WgmmaControlParticipation::WarpgroupAllThreadsSameInstruction, - }); - record.special_register = None; - record.ldmatrix = None; - record.lowering = "generated_wgmma_control".into(); - record.expected_ptx = InstructionPattern { - mnemonic: "wgmma".into(), - modifiers: suffix.split('.').map(str::to_owned).collect(), - operands: if wait { - vec![OperandPattern::Immediate] - } else { - vec![] - }, - }; - record.summary = format!("Generated {id} test record."); - catalog.intrinsics.push(record); - } - catalog - } - - fn raw_abi_item<'a>(raw: &'a str, id: &str) -> &'a str { - let marker = format!("Catalog ID: `{id}`"); - let marker_offset = raw.find(&marker).unwrap(); - let start = raw[..marker_offset] - .rfind("\n\n") - .map_or(0, |offset| offset + 2); - let end = raw[marker_offset..] - .find("\n}\n") - .map_or(raw.len(), |offset| marker_offset + offset + 3); - &raw[start..end] - } - - fn raw_abi_safety_block<'a>(raw: &'a str, id: &str) -> &'a str { - let item = raw_abi_item(raw, id); - let marker = "/// # Safety\n"; - let start = item.find(marker).unwrap() + marker.len(); - let end = item[start..].find("#[").unwrap() + start; - &item[start..end] - } - - #[test] - fn unsafe_safety_docs_never_reference_missing_addr() { - let catalog = catalog_with_clc(); - let raw = render_raw_abi(&catalog, "test-hash").unwrap(); - let offenders = catalog - .intrinsics - .iter() - .filter(|record| !record.rust.safe) - .filter(|record| raw_abi_item(&raw, &record.id).contains("/// `addr`")) - .filter(|record| { - !record - .rust - .arguments - .iter() - .any(|argument| argument.starts_with('*')) - }) - .map(|record| record.id.as_str()) - .collect::>(); - assert!( - offenders.is_empty(), - "unsafe safety docs reference `addr` without a pointer argument: {}", - offenders.join(", ") - ); - } - - #[test] - fn unsafe_safety_blocks_are_exclusive_to_one_family() { - let catalog = catalog_with_clc(); - let raw = render_raw_abi(&catalog, "test-hash").unwrap(); - let mut families_by_block = BTreeMap::<&str, BTreeSet<&str>>::new(); - for record in catalog.intrinsics.iter().filter(|record| !record.rust.safe) { - families_by_block - .entry(raw_abi_safety_block(&raw, &record.id)) - .or_default() - .insert(&record.family); - } - let collisions = families_by_block - .values() - .filter(|families| families.len() != 1) - .collect::>(); - assert!( - collisions.is_empty(), - "unsafe safety blocks shared across families: {collisions:?}; render a legitimate shared contract through a shared arm" - ); - } - - #[test] - fn wgmma_controls_render_closed_compatibility_and_backend_routes() { - let catalog = catalog_with_wgmma_controls(); - validate_renderable(&catalog).unwrap(); - assert_eq!(wgmma_controls(&catalog).count(), 3); - - let compatibility = render_compat_wgmma_control(&catalog, "test-hash"); - assert!(compatibility.contains("pub fn wgmma_fence()")); - assert!(compatibility.contains("pub fn wgmma_commit_group()")); - assert!(compatibility.contains("pub fn wgmma_wait_group()")); - assert!(compatibility.contains("__wgmma_wait_group(N as u64);")); - assert!(compatibility.contains("pub(crate) fn __wgmma_wait_group(_max_pending: u64)")); - - let raw = render_raw_abi(&catalog, "test-hash").unwrap(); - assert!(raw.contains("pub unsafe fn i0317()")); - assert!(raw.contains("pub unsafe fn i0318()")); - assert!(raw.contains("pub unsafe fn i0319(_arg0: u64)")); - let raw_mod = render_raw_mod(&catalog, "test-hash"); - assert!( - raw_mod.contains("pub use crate::__cuda_oxide_intrinsic_abi_v1::i0317 as wgmma_fence;") - ); - assert!( - raw_mod.contains( - "pub use crate::__cuda_oxide_intrinsic_abi_v1::i0319 as wgmma_wait_group;" - ) - ); - - let dialect = render_dialect_wgmma_control(&catalog, "test-hash"); - for op in [ - "WgmmaFenceSyncAlignedOp", - "WgmmaCommitGroupSyncAlignedOp", - "WgmmaWaitGroupSyncAlignedOp", - ] { - assert!(dialect.contains(&format!("pub struct {op};"))); - assert!(dialect.contains(&format!("{op}::register(ctx);"))); - } - assert!(dialect.contains("operand must be i64")); - - let importer = render_importer(&catalog, "test-hash"); - assert!(importer.contains("cuda_device::wgmma::wgmma_fence")); - assert!(importer.contains("cuda_device::wgmma::wgmma_commit_group")); - assert!(importer.contains("cuda_device::wgmma::__wgmma_wait_group")); - assert!(importer.contains("wgmma_wait_group requires a compile-time constant")); - assert!(importer.contains("WgmmaWaitGroupSyncAlignedOp::build(ctx, max_pending)")); - - let lowering = render_lowering(&catalog, "test-hash"); - assert!(lowering.contains("llvm_nvvm_wgmma_fence_sync_aligned")); - assert!(lowering.contains("llvm__nvvm_dwgmma_dcommit_ugroup_dsync_daligned")); - assert!(lowering.contains("llvm__nvvm_dwgmma_dwait_ugroup_dsync_daligned")); - assert!(lowering.contains("IntrinsicBackend::LlvmNvptx")); - assert!(lowering.contains("IntrinsicBackend::LibNvvm")); - assert!(lowering.contains("wgmma.fence.sync.aligned;")); - assert!(lowering.contains("wgmma.commit_group.sync.aligned;")); - assert!(lowering.contains("wgmma.wait_group.sync.aligned $0;")); - assert!(lowering.contains("\"n,~{memory}\"")); - - let targets = render_targets(&catalog, "test-hash"); - assert!( - targets - .contains("pub enum GeneratedWgmmaControlMode { Fence, CommitGroup, WaitGroup }") - ); - assert!(targets.contains("GeneratedIntrinsicVariant::WgmmaControl")); - assert!(targets.contains("WgmmaFenceSyncAlignedOp>(operation, ctx).is_some()")); - - for record in wgmma_controls(&catalog) { - let probe = render_probe(&catalog, record, "test-hash"); - assert!(probe.contains(&format!("@{}", llvm(record).symbol))); - assert!(!probe.contains("asm sideeffect")); - assert!(probe.contains("attributes #0 = { convergent }")); - } - - let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); - assert!(outputs.contains_key(&PathBuf::from( - "crates/cuda-device/src/generated/wgmma_control.rs" - ))); - assert!(outputs.contains_key(&PathBuf::from( - "crates/dialect-nvvm/src/ops/generated/wgmma_control.rs" - ))); - let reference = render_reference(&catalog, "test-hash"); - assert!(reference.contains("## WGMMA-control contracts")); - } - - fn catalog_with_tcgen05() -> CatalogFile { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::test_catalog_with_tcgen05(&repo_root).unwrap(); - assert_eq!(tcgen05_intrinsics(&catalog).count(), 233); - catalog - } - - #[test] - fn selection_alternatives_keep_predicates_and_constraints_grouped() { - let selections = vec![ - CatalogSelection { - source_record: "SELECT_A".into(), - asm: "op.a $dst;".into(), - predicates: vec!["HasA".into(), "HasCommon".into()], - constraints: ImportedSelectionConstraints { - address_space: Some(ImportedAddressSpace::Generic), - immediate_bindings: vec![], - }, - }, - CatalogSelection { - source_record: "SELECT_B".into(), - asm: "op.b $dst;".into(), - predicates: vec!["HasB".into()], - constraints: ImportedSelectionConstraints { - address_space: Some(ImportedAddressSpace::Shared), - immediate_bindings: vec![], - }, - }, - ]; - - let rendered = generated_selection_alternatives(&selections); - assert_eq!(rendered.matches("GeneratedSelectionAlternative").count(), 2); - assert!(rendered.contains( - "source_record: \"SELECT_A\", asm: \"op.a $dst;\", predicates: &[\"HasA\", \"HasCommon\"], constraints: GeneratedSelectionConstraints { address_space: Some(GeneratedSelectionAddressSpace::Generic), immediate_bindings: &[] }" - )); - assert!(rendered.contains( - "source_record: \"SELECT_B\", asm: \"op.b $dst;\", predicates: &[\"HasB\"], constraints: GeneratedSelectionConstraints { address_space: Some(GeneratedSelectionAddressSpace::Shared), immediate_bindings: &[] }" - )); - } - - #[test] - fn stmatrix_rendering_preserves_all_four_public_and_backend_contracts() { - let catalog = catalog_with_stmatrix(); - validate_renderable(&catalog).unwrap(); - assert_eq!(stmatrices(&catalog).count(), 4); - - let compatibility = render_compat_stmatrix(&catalog, "test-hash"); - for signature in [ - "pub unsafe fn stmatrix_m8n8_x2(smem_ptr: *mut u8, r0: u32, r1: u32)", - "pub unsafe fn stmatrix_m8n8_x2_trans(smem_ptr: *mut u8, r0: u32, r1: u32)", - "pub unsafe fn stmatrix_m8n8_x4(smem_ptr: *mut u8, r0: u32, r1: u32, r2: u32, r3: u32)", - "pub unsafe fn stmatrix_m8n8_x4_trans(smem_ptr: *mut u8, r0: u32, r1: u32, r2: u32, r3: u32)", - ] { - assert!(compatibility.contains(signature)); - } - assert!(compatibility.contains("All warp lanes must participate")); - - let dialect_mod = render_dialect_mod(&catalog, "test-hash"); - assert!(dialect_mod.contains("mod stmatrix;")); - assert!(dialect_mod.contains("stmatrix::register(ctx)")); - let dialect = render_dialect_stmatrix(&catalog, "test-hash"); - for (op_type, op_name, operands) in [ - ("StmatrixM8n8X2Op", "nvvm.stmatrix_m8n8_x2", 3), - ("StmatrixM8n8X2TransOp", "nvvm.stmatrix_m8n8_x2_trans", 3), - ("StmatrixM8n8X4Op", "nvvm.stmatrix_m8n8_x4", 5), - ("StmatrixM8n8X4TransOp", "nvvm.stmatrix_m8n8_x4_trans", 5), - ] { - assert!(dialect.contains(&format!("pub struct {op_type};"))); - assert!(dialect.contains(&format!("name = \"{op_name}\""))); - assert!(dialect.contains(&format!("NOpdsInterface<{operands}>"))); - assert!(dialect.contains(&format!("{op_type}::register(ctx)"))); - } - - let importer = render_importer(&catalog, "test-hash"); - for (path, op_type, marker) in [ - ( - "cuda_device::tcgen05::stmatrix_m8n8_x2", - "StmatrixM8n8X2Op", - "v1:i0301", - ), - ( - "cuda_device::tcgen05::stmatrix_m8n8_x2_trans", - "StmatrixM8n8X2TransOp", - "v1:i0302", - ), - ( - "cuda_device::tcgen05::stmatrix_m8n8_x4", - "StmatrixM8n8X4Op", - "v1:i0303", - ), - ( - "cuda_device::tcgen05::stmatrix_m8n8_x4_trans", - "StmatrixM8n8X4TransOp", - "v1:i0304", - ), - ] { - assert!(importer.contains(path)); - assert!(importer.contains(&format!("{op_type}::get_concrete_op_info()"))); - assert!(importer.contains(&format!( - "set_generated_intrinsic_marker(ctx, store, \"{marker}\")" - ))); - } - - let lowering = render_lowering(&catalog, "test-hash"); - assert_eq!( - lowering - .matches("impl MirToLlvmConversion for StmatrixM8n8") - .count(), - 4 - ); - assert!(lowering.contains("IntrinsicBackend::LlvmNvptx")); - assert!(lowering.contains("IntrinsicBackend::LibNvvm")); - assert!(lowering.contains("cast_to_shared_addrspace")); - assert!(lowering.contains("llvm_nvvm_stmatrix_sync_aligned_m8n8_x2_b16_p3")); - assert!(lowering.contains("llvm_nvvm_stmatrix_sync_aligned_m8n8_x4_trans_b16_p3")); - assert!( - lowering.contains("stmatrix.sync.aligned.m8n8.x{register_count}{trans}.shared.b16") - ); - assert!(lowering.contains("std::iter::once(\"~{memory}\")")); - - let x2 = stmatrices(&catalog) - .find(|record| record.id == "stmatrix_m8n8_x2_b16") - .unwrap(); - let probe = render_probe(&catalog, x2, "test-hash"); - assert!(probe.contains( - "declare void @llvm.nvvm.stmatrix.sync.aligned.m8n8.x2.b16.p3(ptr addrspace(3), i32, i32)" - )); - assert!(probe.contains("%shared = addrspacecast ptr %generic to ptr addrspace(3)")); - assert!(probe.contains( - "call void @llvm.nvvm.stmatrix.sync.aligned.m8n8.x2.b16.p3(ptr addrspace(3) %shared, i32 %r0, i32 %r1)" - )); - - let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); - assert!(outputs.contains_key(&PathBuf::from( - "crates/cuda-device/src/generated/stmatrix.rs" - ))); - assert!(outputs.contains_key(&PathBuf::from( - "crates/dialect-nvvm/src/ops/generated/stmatrix.rs" - ))); - } - - #[test] - fn cluster_composites_use_generated_private_sreg_leaves() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - validate_renderable(&catalog).unwrap(); - - let leaves = render_compat_cluster_sreg(&catalog, "test-hash"); - assert_eq!(leaves.matches("pub(crate) fn __cluster_").count(), 6); - for name in [ - "__cluster_idxX", - "__cluster_idxY", - "__cluster_idxZ", - "__cluster_grid_dimX", - "__cluster_grid_dimY", - "__cluster_grid_dimZ", - ] { - assert!(leaves.contains(&format!("pub(crate) fn {name}() -> u32"))); - } - assert!(!leaves.contains("pub fn cluster_ctaidX")); - assert!(!leaves.contains("pub fn cluster_nctaidX")); - - let importer = render_importer(&catalog, "test-hash"); - assert!(importer.contains("cuda_device::cluster::__cluster_idxX")); - assert!(importer.contains("cuda_device::cluster::__cluster_grid_dimZ")); - - let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); - assert!(outputs.contains_key(&PathBuf::from( - "crates/cuda-device/src/generated/cluster_sreg.rs" - ))); - } - - #[test] - fn packed_alu_and_conversion_render_exact_pure_inline_ptx_adapters() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - validate_renderable(&catalog).unwrap(); - assert_eq!(packed_alus(&catalog).count(), 30); - assert_eq!(packed_conversions(&catalog).count(), 18); - - let dialect = render_dialect_packed_alu(&catalog, "test-hash"); - for op in [ - "FmaBf16x2Op", - "FmaReluBf16x2Op", - "AddBf16x2Op", - "SubBf16x2Op", - "MulBf16x2Op", - "MinBf16x2Op", - "MaxBf16x2Op", - "NegBf16x2Op", - "AbsBf16x2Op", - "FmaF16x2Op", - "FmaReluF16x2Op", - "AddF16x2Op", - "SubF16x2Op", - "MulF16x2Op", - "MinF16x2Op", - "MaxF16x2Op", - "NegF16x2Op", - "AbsF16x2Op", - "AddF32x2Op", - "AddFtzF32x2Op", - "SubF32x2Op", - "SubFtzF32x2Op", - "MulF32x2Op", - "MulFtzF32x2Op", - "FmaF32x2Op", - "FmaFtzF32x2Op", - ] { - assert!(dialect.contains(&format!("pub struct {op}"))); - assert!(dialect.contains(&format!("{op}::register(ctx)"))); - } - let conversion_dialect = render_dialect_packed_conversion(&catalog, "test-hash"); - for op in [ - "CvtF32x2Bf16x2Op", - "CvtF16x2F32Op", - "CvtRzF16x2F32Op", - "CvtRnReluF16x2F32Op", - "CvtRnReluBf16x2F32Op", - "CvtRzBf16x2F32Op", - "CvtRnSatfiniteE4m3x2F32Op", - "CvtRnSatfiniteReluE4m3x2F32Op", - "CvtRnSatfiniteE5m2x2F32Op", - "CvtRnSatfiniteReluE5m2x2F32Op", - ] { - assert!(conversion_dialect.contains(&format!("pub struct {op}"))); - assert!(conversion_dialect.contains(&format!("{op}::register(ctx)"))); - } - assert!(conversion_dialect.contains("low f16 lane")); - assert!(conversion_dialect.contains("low bf16 lane")); - assert!(conversion_dialect.contains("low e4m3 lane")); - assert!(conversion_dialect.contains("low e5m2 lane")); - assert!(conversion_dialect.contains("vec![low, high]")); - - let importer = render_importer(&catalog, "test-hash"); - assert!(importer.contains("cuda_device::bf16x2::fma_bf16x2")); - assert!(importer.contains("cuda_device::f16x2::fma_f16x2")); - assert!(importer.contains("cuda_device::convert::cvt_bf16x2_f32")); - assert!(importer.contains("cuda_device::convert::cvt_f16x2_f32")); - assert!(importer.contains("cuda_device::convert::cvt_rz_bf16x2_f32")); - assert!(importer.contains("cuda_device::convert::cvt_rn_satfinite_e4m3x2_f32")); - assert!(importer.contains("cuda_device::convert::cvt_rn_satfinite_relu_e5m2x2_f32")); - assert!(importer.contains("FmaBf16x2Op::build(ctx, arg0, arg1, arg2)")); - assert!(importer.contains("CvtF32x2Bf16x2Op::build(ctx, arg0, arg1)")); - - let lowering = render_lowering(&catalog, "test-hash"); - for mnemonic in [ - "fma.rn.bf16x2", - "fma.rn.relu.bf16x2", - "add.rn.bf16x2", - "sub.rn.bf16x2", - "mul.rn.bf16x2", - "min.bf16x2", - "max.bf16x2", - "neg.bf16x2", - "abs.bf16x2", - "fma.rn.f16x2", - "fma.rn.relu.f16x2", - "add.rn.f16x2", - "sub.rn.f16x2", - "mul.rn.f16x2", - "min.f16x2", - "max.f16x2", - "neg.f16x2", - "abs.f16x2", - ] { - assert!(lowering.contains(&format!( - "convert_generated_packed_alu(ctx, rewriter, self.get_operation(), \"{mnemonic}\", 32)" - ))); - } - for mnemonic in [ - "add.rn.f32x2", - "add.rn.ftz.f32x2", - "sub.rn.f32x2", - "sub.rn.ftz.f32x2", - "mul.rn.f32x2", - "mul.rn.ftz.f32x2", - "fma.rn.f32x2", - "fma.rn.ftz.f32x2", - ] { - assert!(lowering.contains(&format!( - "convert_generated_packed_alu(ctx, rewriter, self.get_operation(), \"{mnemonic}\", 64)" - ))); - } - for mnemonic in [ - "cvt.rn.bf16x2.f32", - "cvt.rn.f16x2.f32", - "cvt.rz.f16x2.f32", - "cvt.rn.relu.f16x2.f32", - "cvt.rn.relu.bf16x2.f32", - "cvt.rz.bf16x2.f32", - ] { - assert!(lowering.contains(&format!( - "convert_generated_packed_f32x2(ctx, rewriter, self.get_operation(), None, \"{mnemonic}\", 32)" - ))); - } - for (intrinsic, mnemonic) in [ - ("llvm_nvvm_ff_to_e4m3x2_rn", "cvt.rn.satfinite.e4m3x2.f32"), - ( - "llvm_nvvm_ff_to_e4m3x2_rn_relu", - "cvt.rn.satfinite.relu.e4m3x2.f32", - ), - ("llvm_nvvm_ff_to_e5m2x2_rn", "cvt.rn.satfinite.e5m2x2.f32"), - ( - "llvm_nvvm_ff_to_e5m2x2_rn_relu", - "cvt.rn.satfinite.relu.e5m2x2.f32", - ), - ] { - assert!(lowering.contains(&format!( - "convert_generated_packed_f32x2(ctx, rewriter, self.get_operation(), Some(\"{intrinsic}\"), \"{mnemonic}\", 16)" - ))); - } - - for record in packed_alus(&catalog) { - let probe = render_probe(&catalog, record, "test-hash"); - let register = packed_alu_register_constraint(record); - let constraints = std::iter::once(format!("={register}")) - .chain(std::iter::repeat_n( - register.to_owned(), - record.rust.arguments.len(), - )) - .collect::>() - .join(","); - assert!(probe.contains(&format!("\", \"{constraints}\""))); - assert!(!probe.contains("asm sideeffect")); - assert!(!probe.contains("~{memory}")); - } - let mut typed_probe_count = 0; - for conversion in packed_conversions(&catalog) { - let probe = render_probe(&catalog, conversion, "test-hash"); - if packed_conversion_typed_llvm_name(conversion).is_some() { - typed_probe_count += 1; - let llvm = llvm(conversion); - let symbol = llvm.resolved_symbol.as_deref().unwrap_or(&llvm.symbol); - assert!(probe.contains(&format!("declare i16 @{symbol}(float, float)"))); - assert!(probe.contains(&format!("call i16 @{symbol}(float %high, float %low)"))); - assert!(!probe.contains(" asm ")); - } else if packed_conversion_source(conversion) == PackedConversionSourceFormat::F32x2 { - assert!(probe.contains(&format!( - "asm \"{} $0, $2, $1;\", \"{}\"(float %low, float %high)", - packed_conversion_ptx_mnemonic(conversion), - packed_conversion_constraint(conversion), - ))); - assert!(!probe.contains("declare ")); - } else { - // One packed source operand, so no high/low reordering. - let source_ty = format!("i{}", packed_conversion_source_width(conversion)); - assert!(probe.contains(&format!( - "asm \"{} $0, $1;\", \"{}\"({source_ty} %packed)", - packed_conversion_ptx_mnemonic(conversion), - packed_conversion_constraint(conversion), - ))); - assert!(!probe.contains("declare ")); - } - assert!(!probe.contains("asm sideeffect")); - } - assert_eq!(typed_probe_count, 4); - - let raw = render_raw_abi(&catalog, "test-hash").unwrap(); - assert!(raw.contains("pub fn i0062(_arg0: u32, _arg1: u32, _arg2: u32) -> u32")); - assert!(raw.contains("pub fn i0071(_arg0: f32, _arg1: f32) -> u32")); - assert!(raw.contains("pub fn i0072(_arg0: u32, _arg1: u32, _arg2: u32) -> u32")); - assert!(raw.contains("pub fn i0085(_arg0: f32, _arg1: f32) -> u32")); - assert!(raw.contains("pub fn i0259(_arg0: f32, _arg1: f32) -> u16")); - assert!(raw.contains("pub fn i0262(_arg0: f32, _arg1: f32) -> u16")); - assert!(raw.contains("pub fn i0995(_arg0: u64, _arg1: u64) -> u64")); - assert!(raw.contains("pub fn i1002(_arg0: u64, _arg1: u64, _arg2: u64) -> u64")); - assert!(!raw.contains("#[must_use]\n#[inline(never)]\npub fn i0062")); - let f16_raw = raw.find("pub fn i0072").unwrap(); - assert!(raw[..f16_raw].ends_with("#[must_use]\n#[inline(never)]\n")); - - let compatibility = - render_compat_packed_alu(&catalog, "test-hash", PackedAluFormat::Bf16x2); - assert!(compatibility.contains("pub fn fma_bf16x2(arg0: u32, arg1: u32, arg2: u32)")); - assert!(!compatibility.contains("fma_f16x2")); - assert!(compatibility.contains("let _ = arg0;")); - assert!(!compatibility.contains("let _ = (arg0);")); - let compatibility = render_compat_packed_alu(&catalog, "test-hash", PackedAluFormat::F16x2); - let f16_compat = compatibility.find("pub fn fma_f16x2").unwrap(); - assert!(compatibility[..f16_compat].ends_with("#[must_use]\n#[inline(never)]\n")); - assert!(!compatibility.contains("fma_bf16x2")); - let compatibility = render_compat_packed_alu(&catalog, "test-hash", PackedAluFormat::F32x2); - assert!(compatibility.contains("pub fn add_f32x2(arg0: u64, arg1: u64) -> u64")); - assert!( - compatibility.contains("pub fn fma_ftz_f32x2(arg0: u64, arg1: u64, arg2: u64) -> u64") - ); - assert!(!compatibility.contains("fma_f16x2")); - let reference = render_reference(&catalog, "test-hash"); - assert!(reference.contains("`fma_f16x2` carries one packed `f16x2` value in a `u32`")); - assert!(reference.contains( - "native instruction starts at PTX 4.2 / `sm_53`; cuda-oxide admits it from sm_70+" - )); - assert!(reference.contains("LLVM-NVPTX PTX 6.0 on sm_70+")); - assert!(reference.contains("libNVVM PTX 4.2 on sm_75+")); - assert!(reference.contains( - "`cvt_rn_relu_f16x2_f32` converts two `f32` inputs to packed `f16x2` using nearest-even rounding with ReLU" - )); - assert!(reference.contains("pure `cvt.rz.bf16x2.f32` inline PTX")); - assert!(reference.contains( - "LLVM-NVPTX uses typed `llvm.nvvm.ff.to.e4m3x2.rn` with `[high, low]` inputs; libNVVM uses pure `cvt.rn.satfinite.e4m3x2.f32` inline PTX" - )); - let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); - assert!(outputs.contains_key(&PathBuf::from("crates/cuda-device/src/generated/bf16x2.rs"))); - assert!(outputs.contains_key(&PathBuf::from("crates/cuda-device/src/generated/f16x2.rs"))); - assert!(outputs.contains_key(&PathBuf::from("crates/cuda-device/src/generated/f32x2.rs"))); - assert!(outputs.contains_key(&PathBuf::from( - "crates/cuda-device/src/generated/convert.rs" - ))); - let compatibility = render_compat_packed_conversion( - &catalog, - "test-hash", - "cuda_device::convert::", - "convert", - ("lo", "hi"), - ); - assert!(compatibility.contains("pub fn cvt_bf16x2_f32(lo: f32, hi: f32) -> u32")); - assert!(compatibility.contains("pub fn cvt_f16x2_f32(lo: f32, hi: f32) -> u32")); - assert!(compatibility.contains("pub fn cvt_rz_bf16x2_f32(lo: f32, hi: f32) -> u32")); - assert!( - compatibility.contains("pub fn cvt_rn_satfinite_e4m3x2_f32(lo: f32, hi: f32) -> u16") - ); - assert!( - compatibility - .contains("pub fn cvt_rn_satfinite_relu_e5m2x2_f32(lo: f32, hi: f32) -> u16") - ); - assert!(!compatibility.contains("pub fn cvt_f32x2_bf16x2(")); - assert!(!outputs.contains_key(&PathBuf::from( - "crates/cuda-device/src/generated/tcgen05_conversion.rs" - ))); - } - - #[test] - fn scalar_tf32_conversions_render_one_exact_attribute_carrier() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - validate_renderable(&catalog).unwrap(); - let records = scalar_conversions(&catalog).collect::>(); - assert_eq!(records.len(), 10); - - for record in &records { - assert_eq!(record.rust.arguments, ["f32"]); - assert_eq!(record.rust.result, "u32"); - assert!(record.backend_lowerings.iter().any(|lowering| { - lowering.backend == IntrinsicBackend::LlvmNvptx - && lowering.mechanism == BackendLoweringMechanism::TypedNvvm - })); - assert!(record.backend_lowerings.iter().any(|lowering| { - lowering.backend == IntrinsicBackend::LibNvvm - && lowering.mechanism == BackendLoweringMechanism::InlinePtx - })); - } - - let dialect = render_dialect_scalar_conversion(&catalog, "test-hash"); - assert_eq!(dialect.matches("pub struct ScalarConversionOp;").count(), 1); - assert!(dialect.contains("pub enum ScalarConversionRoundingAttr")); - assert!(dialect.contains("pub enum ScalarConversionSaturationAttr")); - assert_eq!( - dialect - .matches("ScalarConversionOp::register(ctx);") - .count(), - 1 - ); - for record in &records { - assert!(dialect.contains(scalar_conversion_rounding_attr(record))); - assert!(dialect.contains(scalar_conversion_saturation_attr(record))); - } - - let importer = render_importer(&catalog, "test-hash"); - assert_eq!( - importer - .matches("let intrinsic = ScalarConversionOp::build") - .count(), - 10 - ); - let lowering = render_lowering(&catalog, "test-hash"); - assert_eq!( - lowering - .matches("impl MirToLlvmConversion for ScalarConversionOp") - .count(), - 1 - ); - for record in &records { - assert!(lowering.contains(&record.llvm_identifier())); - assert!(lowering.contains(&scalar_conversion_ptx_mnemonic(record))); - let probe = render_probe(&catalog, record, "test-hash"); - assert!(probe.contains(&format!("declare i32 @{}(float)", llvm(record).symbol))); - } - - let targets = render_targets(&catalog, "test-hash"); - let production_targets = targets.split("\n#[cfg(test)]").next().unwrap(); - assert_eq!( - production_targets - .matches("GeneratedIntrinsicVariant::ScalarConversion {") - .count(), - 11 - ); - assert!(targets.contains("Operation::get_op::")); - assert!(targets.contains("GeneratedScalarConversionRounding::NearestAway")); - assert!(targets.contains("GeneratedScalarConversionSaturation::ReluSatfinite")); - - let compatibility = render_compat_packed_conversion( - &catalog, - "test-hash", - "cuda_device::convert::", - "convert", - ("lo", "hi"), - ); - assert!(compatibility.contains("pub fn cvt_rna_tf32_f32(value: f32) -> u32")); - assert!(compatibility.contains("pub fn cvt_rz_relu_satfinite_tf32_f32(value: f32) -> u32")); - let raw = render_raw_abi(&catalog, "test-hash").unwrap(); - assert!(raw.contains("pub fn i0368(_arg0: f32) -> u32")); - assert!(raw.contains("pub fn i0377(_arg0: f32) -> u32")); - } - - #[test] - fn ldmatrix_family_lowering_uses_one_attribute_dispatch_impl() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - - let dialect = render_dialect_ldmatrix(&catalog, "test-hash"); - let rendered = render_lowering(&catalog, "test-hash"); - assert_eq!( - rendered - .matches("impl MirToLlvmConversion for LdmatrixOp") - .count(), - 1 - ); - assert!(rendered.contains("LdmatrixMultiplicityAttr::X4")); - assert!(rendered.contains("LdmatrixMultiplicityAttr::X2")); - assert!(rendered.contains("LdmatrixMultiplicityAttr::X1")); - assert!(rendered.contains("LdmatrixLayoutAttr::Transposed")); - assert!(rendered.contains("llvm_nvvm_ldmatrix_sync_aligned_m8n8_x2_trans_b16_p3")); - for (op_type, op_name, count, instruction_head, intrinsic) in [ - ( - "LdmatrixX1Op", - "nvvm.ldmatrix_x1", - 1, - "ldmatrix.sync.aligned.m8n8.x1.shared.b16", - "llvm_nvvm_ldmatrix_sync_aligned_m8n8_x1_b16_p3", - ), - ( - "LdmatrixX1TransOp", - "nvvm.ldmatrix_x1_trans", - 1, - "ldmatrix.sync.aligned.m8n8.x1.trans.shared.b16", - "llvm_nvvm_ldmatrix_sync_aligned_m8n8_x1_trans_b16_p3", - ), - ( - "LdmatrixX2Op", - "nvvm.ldmatrix_x2", - 2, - "ldmatrix.sync.aligned.m8n8.x2.shared.b16", - "llvm_nvvm_ldmatrix_sync_aligned_m8n8_x2_b16_p3", - ), - ( - "LdmatrixX2TransOp", - "nvvm.ldmatrix_x2_trans", - 2, - "ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16", - "llvm_nvvm_ldmatrix_sync_aligned_m8n8_x2_trans_b16_p3", - ), - ( - "LdmatrixX4Op", - "nvvm.ldmatrix_x4", - 4, - "ldmatrix.sync.aligned.m8n8.x4.shared.b16", - "llvm_nvvm_ldmatrix_sync_aligned_m8n8_x4_b16_p3", - ), - ( - "LdmatrixX4TransOp", - "nvvm.ldmatrix_x4_trans", - 4, - "ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16", - "llvm_nvvm_ldmatrix_sync_aligned_m8n8_x4_trans_b16_p3", - ), - ] { - assert!(dialect.contains(&format!("pub struct {op_type};"))); - assert!(dialect.contains(&format!("name = {op_name:?}"))); - assert!(dialect.contains(&format!("NResultsInterface<{count}>"))); - assert!(dialect.contains(&format!("{op_type}::register(ctx);"))); - assert_eq!( - rendered - .matches(&format!("impl MirToLlvmConversion for {op_type}")) - .count(), - 1 - ); - assert!(rendered.contains(&format!( - "self.get_operation(), {count}, {instruction_head:?}, {intrinsic:?}" - ))); - } - - let importer = render_importer(&catalog, "test-hash"); - assert!(!importer.contains("LdmatrixX1Op")); - assert!(!importer.contains("LdmatrixX4TransOp")); - let targets = render_targets(&catalog, "test-hash"); - for op_name in [ - "nvvm.ldmatrix_x1", - "nvvm.ldmatrix_x1_trans", - "nvvm.ldmatrix_x2", - "nvvm.ldmatrix_x2_trans", - "nvvm.ldmatrix_x4", - "nvvm.ldmatrix_x4_trans", - ] { - assert!(targets.contains(&format!("{op_name:?} => Some("))); - } - assert_eq!(ldmatrix(&catalog).count(), 18); - assert!(dialect.contains("LdmatrixShapeAttr::M16n16")); - assert!(dialect.contains("LdmatrixShapeAttr::M8n16")); - assert!(dialect.contains("LdmatrixElementAttr::B8x16B4x16P64")); - assert!(rendered.contains( - "llvm__nvvm_dldmatrix_dsync_daligned_dm16n16_dx1_dtrans_db8x16_db4x16_up64_dp3" - )); - assert!(rendered.contains("ldmatrix.sync.aligned.m16n16.x2.trans.shared.b8x16.b6x16_p32")); - - let raw = render_raw_abi(&catalog, "test-hash").unwrap(); - assert!(raw.contains( - "Instruction floor PTX 8.6; the selected target may require a newer PTX version." - )); - assert!(raw.contains("no lane may have exited")); - assert!(raw.contains("have 32 readable shared-memory bytes")); - assert!(raw.contains("have 16 readable shared-memory bytes")); - - let reference = render_reference(&catalog, "test-hash"); - assert!(reference.contains(BLACKWELL_LDMATRIX_EFFECTIVE_FLOORS)); - } - - #[test] - fn ldmatrix_x1_probe_keeps_its_pointer_operand() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - let record = catalog - .intrinsics - .iter() - .find(|record| record.id == "ldmatrix_m8n8_x1_b16") - .unwrap(); - - let rendered = render_probe(&catalog, record, "test-hash"); - assert!(rendered.contains( - "declare i32 @llvm.nvvm.ldmatrix.sync.aligned.m8n8.x1.b16.p3(ptr addrspace(3))" - )); - assert!(rendered.contains("define i32 @probe_ldmatrix_m8n8_x1_b16(ptr %generic)")); - assert!(!rendered.contains("@llvm.nvvm.ldmatrix.sync.aligned.m8n8.x1.b16.p3()")); - } - - #[test] - fn packed_atomic_family_uses_one_attribute_dispatch_impl() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - let rendered = render_lowering(&catalog, "test-hash"); - assert_eq!( - rendered - .matches("impl MirToLlvmConversion for PackedAtomicAddOp") - .count(), - 1 - ); - assert_eq!( - rendered.matches("PackedAtomicFormatAttr::F16x2)").count(), - 1 - ); - assert_eq!( - rendered.matches("PackedAtomicFormatAttr::Bf16x2)").count(), - 1 - ); - for (op_type, ptx_type) in [ - ("NvvmAtomAddF16x2Op", "f16x2"), - ("NvvmAtomAddBf16x2Op", "bf16x2"), - ] { - assert_eq!( - rendered - .matches(&format!("impl MirToLlvmConversion for {op_type}")) - .count(), - 1 - ); - assert!(rendered.contains(&format!( - "convert_packed_atom_add(ctx, rewriter, self.get_operation(), \"{ptx_type}\")" - ))); - } - - let dialect = render_dialect_packed_atomic(&catalog, "test-hash"); - for (op_type, op_name) in [ - ("NvvmAtomAddF16x2Op", "nvvm.atom_add_f16x2"), - ("NvvmAtomAddBf16x2Op", "nvvm.atom_add_bf16x2"), - ] { - assert!(dialect.contains(&format!("pub struct {op_type};"))); - assert!(dialect.contains(&format!("name = \"{op_name}\""))); - assert!(dialect.contains(&format!("{op_type}::register(ctx);"))); - } - } - - #[test] - fn packed_atomic_raw_abi_is_unsafe_and_must_use() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - let rendered = render_raw_abi(&catalog, "test-hash").unwrap(); - for abi_id in ["i0014", "i0015"] { - let signature = format!("pub unsafe fn {abi_id}(_arg0: *mut u32, _arg1: u32) -> u32"); - let index = rendered.find(&signature).unwrap(); - assert!(rendered[..index].ends_with("#[must_use]\n#[inline(never)]\n")); - } - } - - #[test] - fn packed_atomic_compatibility_preserves_paths_signature_and_safety() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - let rendered = render_compat_packed_atomic(&catalog, "test-hash"); - - for name in ["atom_add_f16x2", "atom_add_bf16x2"] { - let signature = format!("pub unsafe fn {name}(addr: *mut u32, val: u32) -> u32"); - let index = rendered.find(&signature).unwrap(); - assert!(rendered[..index].ends_with("#[must_use]\n#[inline(never)]\n")); - assert!(rendered.contains(&format!( - "unreachable!(\"{name} called outside CUDA kernel context\")" - ))); - } - assert!(rendered.contains("relaxed GPU-scope operation")); - assert!(rendered.contains("low lane first")); - assert!(rendered.contains("may not form one old 32-bit snapshot")); - assert!(rendered.contains("Requires PTX 6.2 and `sm_70+`")); - assert!(rendered.contains("Requires PTX 7.8 and `sm_90+`")); - assert!(rendered.contains("four writable, four-byte-aligned bytes in global memory")); - assert!(rendered.contains("whole-word atomic or non-atomic lane access")); - assert!(rendered.contains("Racing atomics must use mutually inclusive scopes")); - - let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); - assert_eq!( - outputs.get(&PathBuf::from("crates/cuda-device/src/generated/atomic.rs")), - Some(&rendered) - ); - } - - #[test] - fn cp_async_rendering_preserves_compatibility_dispatch_and_backend_routes() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - validate_renderable(&catalog).unwrap(); - assert_eq!(cp_async_copies(&catalog).count(), 8); - assert_eq!(cp_async_controls(&catalog).count(), 3); - assert_eq!(cp_async_mbarriers(&catalog).count(), 4); - - let compatibility = render_compat_cp_async_copy(&catalog, "test-hash"); - for signature in [ - "pub unsafe fn cp_async_ca_4(_shared_dst: *mut u32, _global_src: *const u32)", - "pub unsafe fn cp_async_ca_8(_shared_dst: *mut u32, _global_src: *const u32)", - "pub unsafe fn cp_async_ca_16(_shared_dst: *mut u32, _global_src: *const u32)", - "pub unsafe fn cp_async_ca_zfill_4(_shared_dst: *mut u32, _global_src: *const u8, _src_size: u32)", - "pub unsafe fn cp_async_ca_zfill_8(_shared_dst: *mut u32, _global_src: *const u8, _src_size: u32)", - "pub unsafe fn cp_async_ca_zfill_16(_shared_dst: *mut u32, _global_src: *const u8, _src_size: u32)", - "pub unsafe fn cp_async_cg_16(_shared_dst: *mut u32, _global_src: *const u32)", - "pub unsafe fn cp_async_cg_zfill_16(_shared_dst: *mut u32, _global_src: *const u8, _src_size: u32)", - "pub unsafe fn cp_async_commit_group()", - "pub unsafe fn cp_async_wait_all()", - "pub unsafe fn cp_async_wait_group(_max_pending: u32)", - "pub unsafe fn cp_async_mbarrier_arrive(_barrier: *mut crate::barrier::Barrier)", - "pub unsafe fn cp_async_mbarrier_arrive_shared(_barrier: *mut crate::barrier::Barrier)", - "pub unsafe fn cp_async_mbarrier_arrive_noinc(_barrier: *mut crate::barrier::Barrier)", - "pub unsafe fn cp_async_mbarrier_arrive_noinc_shared(_barrier: *mut crate::barrier::Barrier)", - ] { - assert!(compatibility.contains(signature)); - } - - let dialect = render_dialect_cp_async_copy(&catalog, "test-hash"); - let importer = render_importer(&catalog, "test-hash"); - let lowering = render_lowering(&catalog, "test-hash"); - let targets = render_targets(&catalog, "test-hash"); - for record in cp_async_copies(&catalog) - .chain(cp_async_controls(&catalog)) - .chain(cp_async_mbarriers(&catalog)) - { - assert!(dialect.contains(&format!("pub struct {}", record.dialect.op_type))); - assert!(dialect.contains(&format!("{}::register(ctx)", record.dialect.op_type))); - assert!(importer.contains(&record.rust.canonical_path)); - for path in &record.rust.compatibility_paths { - assert!(importer.contains(path)); - } - assert!(importer.contains(&intrinsic_marker(&catalog, record))); - assert!(lowering.contains(&format!( - "impl MirToLlvmConversion for {}", - record.dialect.op_type - ))); - assert!(lowering.contains(&record.llvm_identifier())); - assert!( - record - .backend_lowerings - .iter() - .any(|entry| entry.backend == IntrinsicBackend::LlvmNvptx) - ); - assert!( - record - .backend_lowerings - .iter() - .any(|entry| entry.backend == IntrinsicBackend::LibNvvm) - ); - assert!(targets.contains(&format!("id: {:?}", record.id))); - } - assert_eq!(dialect.matches("::register(ctx);").count(), 15); - assert!(lowering.contains("convert_generated_cp_async_copy")); - assert!(lowering.contains("convert_generated_cp_async_control")); - assert!(lowering.contains("convert_generated_cp_async_mbarrier")); - } - - #[test] - fn cp_async_mbarrier_rendering_preserves_counting_and_state_space_routes() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - let records: Vec<_> = cp_async_mbarriers(&catalog).collect(); - assert_eq!(records.len(), 4); - - let raw = render_raw_abi(&catalog, "test-hash").unwrap(); - assert!(raw.contains("pub unsafe fn i0101(_arg0: *mut u64)")); - assert!(raw.contains("live, initialized, eight-byte-aligned mbarrier")); - assert!(raw.contains("initial pending count must already include")); - - let compatibility = render_compat_cp_async_copy(&catalog, "test-hash"); - assert_eq!( - compatibility - .matches("That increment must not exceed the barrier's pending-count limit.") - .count(), - 2 - ); - - let dialect = render_dialect_cp_async_copy(&catalog, "test-hash"); - assert!(dialect.contains("mutable generic/shared pointer to u64")); - let importer = render_importer(&catalog, "test-hash"); - let lowering = render_lowering(&catalog, "test-hash"); - for record in &records { - assert!(importer.contains(&format!( - "let bridge = {}::build(ctx, barrier)", - record.dialect.op_type - ))); - assert!(lowering.contains(&record.llvm_identifier())); - let routes: BTreeSet<_> = record - .backend_lowerings - .iter() - .map(|route| (route.backend, route.mechanism)) - .collect(); - assert_eq!( - routes, - BTreeSet::from([ - ( - IntrinsicBackend::LlvmNvptx, - BackendLoweringMechanism::TypedNvvm, - ), - ( - IntrinsicBackend::LibNvvm, - BackendLoweringMechanism::InlinePtx, - ), - ]) - ); - } - assert!(lowering.contains("\"arrive\", \"generic\"")); - assert!(lowering.contains("\"arrive\", \"shared\"")); - assert!(lowering.contains("\"arrive_no_inc\", \"generic\"")); - assert!(lowering.contains("\"arrive_no_inc\", \"shared\"")); - - let generic = records - .iter() - .find(|record| record.id == "cp_async_mbarrier_arrive") - .unwrap(); - let generic_probe = render_probe(&catalog, generic, "test-hash"); - assert!(generic_probe.contains("declare void @llvm.nvvm.cp.async.mbarrier.arrive(ptr)")); - assert!(!generic_probe.contains("addrspacecast")); - let shared = records - .iter() - .find(|record| record.id == "cp_async_mbarrier_arrive_shared") - .unwrap(); - let shared_probe = render_probe(&catalog, shared, "test-hash"); - assert!(shared_probe.contains("ptr addrspace(3)")); - assert!(shared_probe.contains("addrspacecast ptr %barrier_generic")); - - let reference = render_reference(&catalog, "test-hash"); - assert!(reference.contains("## cp.async mbarrier contracts")); - assert!(reference.contains("cp_async_mbarrier_arrive_noinc`: runtime `unexecuted`")); - - let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); - assert!(outputs.contains_key(&PathBuf::from( - "crates/cuda-device/src/generated/async_copy.rs" - ))); - assert!(outputs.contains_key(&PathBuf::from( - "crates/dialect-nvvm/src/ops/generated/cp_async.rs" - ))); - assert!(!outputs.contains_key(&PathBuf::from( - "crates/dialect-nvvm/src/ops/generated/cp_async_mbarrier.rs" - ))); - } - - #[test] - fn redux_rendering_preserves_mask_first_api_and_source_first_llvm_order() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - assert_eq!(catalog.schema, crate::resolve::CATALOG_SCHEMA); - validate_renderable(&catalog).unwrap(); - - assert_eq!(redux(&catalog).count(), 16); - let record = redux(&catalog).next().unwrap(); - assert_eq!( - record.redux.as_ref().unwrap().adapter, - ReduxAdapter::MaskValueToSourceMemberMask - ); - - let dialect = render_dialect_redux(&catalog, "test-hash"); - assert!(dialect.contains("name = \"nvvm.redux_sync_add\"")); - assert!(dialect.contains("name = \"nvvm.redux_sync_min\"")); - assert!(dialect.contains("vec![member_mask, value]")); - assert!(dialect.contains("ReduxSyncAddOp::register(ctx)")); - assert!(dialect.contains("Signedness::Signed")); - - assert!(dialect.contains("name = \"nvvm.redux_sync_fmin\"")); - assert!(dialect.contains("types::{FP32Type, IntegerType, Signedness}")); - assert!(dialect.contains("let result_ty = FP32Type::get(ctx);")); - assert!(dialect.contains("is_f32(ctx, op.get_operand(1).get_type(ctx))")); - assert!(dialect.contains("is_f32(ctx, op.get_result(0).get_type(ctx))")); - - let importer = render_importer(&catalog, "test-hash"); - assert!(importer.contains("cuda_device::warp::redux_sync_add")); - assert!(importer.contains("let (member_mask, last_op)")); - assert!( - importer.contains("let reduction = ReduxSyncAddOp::build(ctx, member_mask, value)") - ); - assert!(importer.contains("set_generated_intrinsic_marker(ctx, reduction, \"v1:i0017\")")); - - let lowering = render_lowering(&catalog, "test-hash"); - assert!(lowering.contains("impl MirToLlvmConversion for ReduxSyncAddOp")); - assert!( - lowering.contains("convert_redux(ctx, rewriter, self.get_operation(), operands_info") - ); - assert!(lowering.contains("\"llvm_nvvm_redux_sync_add\"")); - - let probe = render_probe(&catalog, record, "test-hash"); - assert!(probe.contains("define i32 @probe_redux_sync_add(i32 %member_mask, i32 %value)")); - assert!(probe.contains("call i32 @llvm.nvvm.redux.sync.add(i32 %value, i32 %member_mask)")); - - let f32_record = redux(&catalog) - .find(|record| record.id == "redux_sync_min_f32") - .unwrap(); - let f32_probe = render_probe(&catalog, f32_record, "test-hash"); - assert!( - f32_probe - .contains("define float @probe_redux_sync_min_f32(i32 %member_mask, float %value)") - ); - assert!( - f32_probe - .contains("call float @llvm.nvvm.redux.sync.fmin(float %value, i32 %member_mask)") - ); - - let raw = render_raw_abi(&catalog, "test-hash").unwrap(); - let signature = "pub unsafe fn i0017(_arg0: u32, _arg1: u32) -> u32"; - let index = raw.find(signature).unwrap(); - assert!(raw[..index].ends_with("#[must_use]\n#[inline(never)]\n")); - assert!(raw.contains("pub unsafe fn i0019(_arg0: u32, _arg1: i32) -> i32")); - assert!(raw.contains("pub unsafe fn i0024(_arg0: u32, _arg1: u32) -> u32")); - assert!(raw.contains("pub unsafe fn i1003(_arg0: u32, _arg1: f32) -> f32")); - assert!(raw.contains("The executing lane must be named in `mask`")); - } - - #[test] - fn vote_rendering_keeps_types_selection_pairs_and_raw_only_uni() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - validate_renderable(&catalog).unwrap(); - assert_eq!(vote_intrinsics(&catalog).count(), 4); - - let dialect = render_dialect_vote(&catalog, "test-hash"); - for op in [ - "VoteSyncAllOp", - "VoteSyncAnyOp", - "VoteSyncBallotOp", - "VoteSyncUniOp", - ] { - assert!(dialect.contains(&format!("pub struct {op}"))); - assert!(dialect.contains(&format!("{op}::register(ctx)"))); - } - assert!(dialect.contains("requires i32 member mask, i1 predicate, and i1 result")); - assert!(dialect.contains("requires i32 member mask, i1 predicate, and i32 result")); - - let importer = render_importer(&catalog, "test-hash"); - assert!(importer.contains("cuda_device::warp::all_sync")); - assert!(!importer.contains("cuda_device::warp::uni_sync")); - assert!(importer.contains("let vote = VoteSyncUniOp::build(ctx, member_mask, predicate)")); - assert!(importer.contains("set_generated_intrinsic_marker(ctx, vote, \"v1:i0043\")")); - - let lowering = render_lowering(&catalog, "test-hash"); - assert!(lowering.contains("impl MirToLlvmConversion for VoteSyncUniOp")); - assert!(lowering.contains("\"llvm_nvvm_vote_uni_sync\"")); - assert!( - lowering.contains("convert_vote(ctx, rewriter, self.get_operation(), operands_info") - ); - - let record = vote_intrinsics(&catalog) - .find(|record| record.id == "ballot_sync") - .unwrap(); - assert_eq!(record.selections.len(), 2); - let probe = render_probe(&catalog, record, "test-hash"); - assert!(probe.contains("define i32 @probe_ballot_sync(i32 %member_mask, i1 %predicate)")); - assert!(probe.contains("define i32 @probe_ballot_sync_immediate(i1 %predicate)")); - assert!(probe.contains("i32 -1, i1 %predicate")); - - let raw = render_raw_abi(&catalog, "test-hash").unwrap(); - for abi_id in ["i0040", "i0041", "i0042", "i0043"] { - assert!(raw.contains(&format!("pub unsafe fn {abi_id}"))); - } - assert!(raw.contains("Every non-exited lane named in `mask`")); - } - - #[test] - fn active_mask_and_warp_match_rendering_preserves_backend_and_abi_contracts() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - validate_renderable(&catalog).unwrap(); - assert_eq!(active_masks(&catalog).count(), 1); - assert_eq!(warp_matches(&catalog).count(), 4); - - let dialect = render_dialect_active_mask(&catalog, "test-hash"); - assert!(dialect.contains("pub struct ActiveMaskOp")); - assert!(dialect.contains("NOpdsInterface<0>, NResultsInterface<1>")); - - let match_dialect = render_dialect_warp_match(&catalog, "test-hash"); - for op in [ - "MatchAnySyncI32Op", - "MatchAnySyncI64Op", - "MatchAllSyncI32Op", - "MatchAllSyncI64Op", - ] { - assert!(match_dialect.contains(&format!("pub struct {op}"))); - assert!(match_dialect.contains(&format!("{op}::register(ctx)"))); - } - - let lowering = render_lowering(&catalog, "test-hash"); - assert!(lowering.contains("IntrinsicBackend::LlvmNvptx => {")); - assert!(lowering.contains("convert_active_mask(ctx, rewriter, op, operands_info)")); - assert!(lowering.contains("IntrinsicBackend::LibNvvm =>")); - assert!(lowering.contains("\"activemask.b32 $0;\"")); - assert!(lowering.contains("\"=r,~{memory}\"")); - assert!(lowering.contains("convert_match_any(")); - assert!(lowering.contains("convert_match_all(")); - assert!(lowering.contains("\"llvm_nvvm_match_all_sync_i64p\"")); - - let match_all = warp_matches(&catalog) - .find(|record| record.id == "match_all_sync") - .unwrap(); - let probe = render_probe(&catalog, match_all, "test-hash"); - for suffix in ["rr", "ri", "ir", "ii"] { - assert!(probe.contains(&format!("@probe_match_all_sync_{suffix}"))); - } - assert!(probe.contains("declare { i32, i1 } @llvm.nvvm.match.all.sync.i32p")); - - let raw = render_raw_abi(&catalog, "test-hash").unwrap(); - assert!(raw.contains("pub fn i0044() -> u32")); - assert!(!raw.contains("pub unsafe fn i0044() -> u32")); - assert!(raw.contains("pub unsafe fn i0045(_arg0: u32, _arg1: u32) -> u32")); - assert!(raw.contains("pub unsafe fn i0046(_arg0: u32, _arg1: u64) -> u32")); - assert!(raw.contains("pub unsafe fn i0047(_arg0: u32, _arg1: u32) -> u32")); - assert!(raw.contains("pub unsafe fn i0048(_arg0: u32, _arg1: u64) -> u32")); - } - - #[test] - fn warp_barrier_rendering_preserves_mask_and_void_contracts() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - validate_renderable(&catalog).unwrap(); - assert_eq!(warp_barriers(&catalog).count(), 1); - - let record = warp_barriers(&catalog).next().unwrap(); - assert_eq!(record.id, "sync_mask"); - assert_eq!( - record.warp_barrier.as_ref().unwrap().adapter, - WarpBarrierAdapter::DirectMemberMask - ); - - let dialect_mod = render_dialect_mod(&catalog, "test-hash"); - assert!(dialect_mod.contains("mod warp_barrier;")); - assert!(dialect_mod.contains("warp_barrier::register(ctx)")); - - let dialect = render_dialect_warp_barrier(&catalog, "test-hash"); - assert!(dialect.contains("pub struct BarWarpSyncOp")); - assert!(dialect.contains("NOpdsInterface<1>, NResultsInterface<0>")); - assert!(dialect.contains("vec![member_mask]")); - assert!(dialect.contains("op.get_num_operands() != 1 || op.get_num_results() != 0")); - assert!(dialect.contains("if !is_i32(ctx, op.get_operand(0).get_type(ctx))")); - assert!(dialect.contains("requires exactly one member-mask operand and no results")); - assert!(dialect.contains("member mask must be i32")); - assert!(dialect.contains("BarWarpSyncOp::register(ctx)")); - - let importer = render_importer(&catalog, "test-hash"); - assert!(importer.contains("cuda_device::warp::sync_mask")); - assert!(importer.contains("let barrier = BarWarpSyncOp::build(ctx, member_mask)")); - assert!(importer.contains("set_generated_intrinsic_marker(ctx, barrier, \"v1:i0049\")")); - assert!(importer.contains("helpers::emit_goto(ctx, *target_idx, barrier")); - - let lowering = render_lowering(&catalog, "test-hash"); - assert!(lowering.contains("impl MirToLlvmConversion for BarWarpSyncOp")); - assert!( - lowering.contains( - "convert_bar_warp_sync(ctx, rewriter, self.get_operation(), operands_info)" - ) - ); - - let probe = render_probe(&catalog, record, "test-hash"); - assert!(probe.contains("declare void @llvm.nvvm.bar.warp.sync(i32)")); - assert!(probe.contains("define void @probe_sync_mask(i32 %member_mask)")); - assert!(probe.contains("call void @llvm.nvvm.bar.warp.sync(i32 %member_mask)")); - assert!(probe.contains("define void @probe_sync_mask_immediate()")); - assert!(probe.contains("call void @llvm.nvvm.bar.warp.sync(i32 -1)")); - - let raw = render_raw_abi(&catalog, "test-hash").unwrap(); - assert!(raw.contains("pub unsafe fn i0049(_arg0: u32) -> ()")); - assert!(raw.contains("On `sm_6x` and earlier")); - assert!(raw.contains("no lane outside `mask` may be active")); - assert!(raw.contains("The barrier orders memory accesses among participating lanes")); - - let reference = render_reference(&catalog, "test-hash"); - assert!(reference.contains("## Warp-barrier contract")); - assert!(reference.contains("no unnamed lane may be active")); - assert!(reference.contains("Both immediate and register masks are admitted")); - } - - #[test] - fn warp_shuffle_rendering_owns_all_i32_f32_and_i64_modes() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - validate_renderable(&catalog).unwrap(); - assert_eq!(warp_shuffles(&catalog).count(), 12); - - let dialect_mod = render_dialect_mod(&catalog, "test-hash"); - assert!(dialect_mod.contains("mod warp_shuffle;")); - assert!(dialect_mod.contains("warp_shuffle::register(ctx)")); - - let dialect = render_dialect_warp_shuffle(&catalog, "test-hash"); - for op in [ - "ShflSyncIdxI32Op", - "ShflSyncBflyI32Op", - "ShflSyncDownI32Op", - "ShflSyncUpI32Op", - "ShflSyncIdxF32Op", - "ShflSyncBflyF32Op", - "ShflSyncDownF32Op", - "ShflSyncUpF32Op", - "ShflSyncIdxI64Op", - "ShflSyncBflyI64Op", - "ShflSyncDownI64Op", - "ShflSyncUpI64Op", - ] { - assert!(dialect.contains(&format!("pub struct {op}"))); - assert!(dialect.contains(&format!("{op}::register(ctx)"))); - } - assert!(dialect.contains("vec![member_mask, value, lane_or_delta]")); - assert!(dialect.contains("requires i32 mask/lane and i32 value/result")); - assert!(dialect.contains("requires i32 mask/lane and f32 value/result")); - assert!(dialect.contains("requires i32 mask/lane and i64 value/result")); - assert!(dialect.contains("fn is_i64")); - assert!(dialect.contains("integer.width() == 64")); - assert!(dialect.contains("IntegerType::get(ctx, 64, Signedness::Unsigned)")); - - let importer = render_importer(&catalog, "test-hash"); - assert!(importer.contains("cuda_device::warp::shuffle_sync")); - assert!(importer.contains("cuda_device::warp::shuffle_up_f32_sync")); - assert!(importer.contains( - "let shuffle = ShflSyncIdxI32Op::build(ctx, member_mask, value, lane_or_delta)" - )); - assert!(importer.contains("set_generated_intrinsic_marker(ctx, shuffle, \"v1:i0050\")")); - for (name, op, marker) in [ - ("shuffle_u64_sync", "ShflSyncIdxI64Op", "v1:i0058"), - ("shuffle_xor_u64_sync", "ShflSyncBflyI64Op", "v1:i0059"), - ("shuffle_down_u64_sync", "ShflSyncDownI64Op", "v1:i0060"), - ("shuffle_up_u64_sync", "ShflSyncUpI64Op", "v1:i0061"), - ] { - assert!(importer.contains(&format!( - "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::{}", - marker.strip_prefix("v1:").unwrap() - ))); - assert!(importer.contains(&format!("cuda_device::warp::{name}"))); - assert!(importer.contains(&format!( - "let shuffle = {op}::build(ctx, member_mask, value, lane_or_delta)" - ))); - assert!(importer.contains(&format!( - "set_generated_intrinsic_marker(ctx, shuffle, {marker:?})" - ))); - } - - let lowering = render_lowering(&catalog, "test-hash"); - assert!(lowering.contains("impl MirToLlvmConversion for ShflSyncIdxI32Op")); - assert!(lowering.contains("convert_shuffle_i32(ctx, rewriter")); - assert!(lowering.contains("\"llvm_nvvm_shfl_sync_idx_i32\", 31)")); - assert!(lowering.contains("impl MirToLlvmConversion for ShflSyncUpF32Op")); - assert!(lowering.contains("convert_shuffle_f32(ctx, rewriter")); - assert!(lowering.contains("\"llvm_nvvm_shfl_sync_up_f32\", 0)")); - - for (op, mode, clamp) in [ - ("ShflSyncIdxI64Op", "idx", 31), - ("ShflSyncBflyI64Op", "bfly", 31), - ("ShflSyncDownI64Op", "down", 31), - ("ShflSyncUpI64Op", "up", 0), - ] { - assert!(lowering.contains(&format!("impl MirToLlvmConversion for {op}"))); - assert!(lowering.contains(&format!( - "convert_shuffle_i64(ctx, rewriter, self.get_operation(), operands_info, {mode:?}, {clamp})" - ))); - } - assert!(!lowering.contains("llvm_nvvm_shfl_sync_idx_i64")); - assert!(!lowering.contains("llvm_nvvm_shfl_sync_bfly_i64")); - assert!(!lowering.contains("llvm_nvvm_shfl_sync_down_i64")); - assert!(!lowering.contains("llvm_nvvm_shfl_sync_up_i64")); - - for record in warp_shuffles(&catalog) { - let probe = render_probe(&catalog, record, "test-hash"); - let shuffle = record.warp_shuffle.as_ref().unwrap(); - if shuffle.value_kind == WarpShuffleValueKind::I64 { - let mode = match shuffle.mode { - WarpShuffleMode::Idx => "idx", - WarpShuffleMode::Bfly => "bfly", - WarpShuffleMode::Down => "down", - WarpShuffleMode::Up => "up", - }; - let asm = format!( - "{{ .reg .b32 lo; .reg .b32 hi; mov.b64 {{lo, hi}}, $1; shfl.sync.{mode}.b32 lo, lo, $2, {}, $3; shfl.sync.{mode}.b32 hi, hi, $2, {}, $3; mov.b64 $0, {{lo, hi}}; }}", - shuffle.clamp, shuffle.clamp - ); - assert!(probe.contains(&format!( - "define i64 @probe_{}(i32 %member_mask, i64 %value, i32 %lane) #0", - record.id - ))); - assert!(probe.contains(&format!( - "call i64 asm sideeffect {asm:?}, \"=l,l,r,r\"(i64 %value, i32 %lane, i32 %member_mask) #0" - ))); - assert_eq!(probe.matches("asm sideeffect").count(), 1); - assert_eq!(probe.matches("attributes #0 = { convergent }").count(), 1); - assert!(!probe.contains("declare i64 @llvm.nvvm.shfl")); - for suffix in ["rr", "ri", "ir", "ii"] { - assert!(!probe.contains(&format!("@probe_{}_{suffix}", record.id))); - } - } else { - for suffix in ["rr", "ri", "ir", "ii"] { - assert!(probe.contains(&format!("@probe_{}_{suffix}", record.id))); - } - assert!(probe.contains(&format!(", i32 {})", shuffle.clamp))); - } - } - - let raw = render_raw_abi(&catalog, "test-hash").unwrap(); - assert!(raw.contains("pub unsafe fn i0050(_arg0: u32, _arg1: u32, _arg2: u32) -> u32")); - assert!(raw.contains("pub unsafe fn i0057(_arg0: u32, _arg1: f32, _arg2: u32) -> f32")); - for abi_id in ["i0058", "i0059", "i0060", "i0061"] { - assert!(raw.contains(&format!( - "pub unsafe fn {abi_id}(_arg0: u32, _arg1: u64, _arg2: u32) -> u64" - ))); - } - assert!(raw.contains("If the computed source lane is in range")); - assert!(raw.contains( - "If PTX marks the computed source out of range, the calling lane's input is copied" - )); - assert!(raw.contains("two `b32` shuffles in one convergent block")); - - let reference = render_reference(&catalog, "test-hash"); - assert!(reference.contains("## Warp-shuffle contracts")); - assert!(reference.contains("inserts clamp `31` during lowering")); - assert!(reference.contains("inserts clamp `0` during lowering")); - assert!(reference.contains("One convergent, side-effecting inline-PTX block splits `i64`")); - assert!(reference.contains("PTX-native source; no LLVM record")); - - let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); - assert!(outputs.contains_key(&PathBuf::from( - "crates/dialect-nvvm/src/ops/generated/warp_shuffle.rs" - ))); - } - - #[test] - fn dot_product_rendering_preserves_stable_paths_and_low_selector() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - validate_renderable(&catalog).unwrap(); - assert_eq!(dot_products(&catalog).count(), 4); - - let compatibility = render_compat_dotprod(&catalog, "test-hash"); - for name in ["dp4a_s32", "dp4a_u32", "dp2a_s32", "dp2a_u32"] { - assert!(compatibility.contains(&format!("pub fn {name}("))); - } - - let dialect = render_dialect_dotprod(&catalog, "test-hash"); - assert!(dialect.contains("pub struct Dp4aS32Op")); - assert!(dialect.contains("pub struct Dp2aU32Op")); - assert!(dialect.contains("NOpdsInterface<3>")); - - let importer = render_importer(&catalog, "test-hash"); - assert!(importer.contains("cuda_device::dotprod::dp2a_s32")); - assert!(importer.contains("let dot = Dp2aS32Op::build(ctx, a, b, c)")); - assert!(importer.contains("set_generated_intrinsic_marker(ctx, dot, \"v1:i0032\")")); - - let lowering = render_lowering(&catalog, "test-hash"); - assert!(lowering.contains("impl MirToLlvmConversion for Dp4aS32Op")); - assert!(lowering.contains("\"llvm_nvvm_idp2a_s_s\"")); - assert!(lowering.contains("\"dp2a.lo.s32.s32 $0, $1, $2, $3;\"")); - assert!(lowering.contains("\"dp2a.lo.s32.s32 $0, $1, $2, $3;\", true)")); - - let low = dot_products(&catalog) - .find(|record| record.id == "dp2a_s32") - .unwrap(); - let probe = render_probe(&catalog, low, "test-hash"); - assert!(probe.contains("call i32 @llvm.nvvm.idp2a.s.s(i32 %a, i32 %b, i1 false, i32 %c)")); - assert!(!probe.contains("i1 true")); - - let target = render_targets(&catalog, "test-hash"); - assert!(target.contains("GeneratedImmediateBinding { argument_index: 2, value: 0 }")); - assert!(!target.contains("GeneratedImmediateBinding { argument_index: 2, value: -1 }")); - } - - #[test] - fn sync_rendering_keeps_barrier_and_threadfence_contracts_explicit() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - validate_renderable(&catalog).unwrap(); - let record = sync_intrinsics(&catalog).next().unwrap(); - assert_eq!(sync_intrinsics(&catalog).count(), 4); - - let raw = render_raw_abi(&catalog, "test-hash").unwrap(); - assert!(raw.contains("pub unsafe fn i0034() -> ()")); - assert!(raw.contains("Every active thread in the CTA must reach the same barrier")); - - let dialect = render_dialect_sync(&catalog, "test-hash"); - assert!(dialect.contains("pub struct Barrier0Op")); - assert!(dialect.contains("NOpdsInterface<0>, NResultsInterface<0>")); - assert!(dialect.contains("Barrier0Op::register(ctx)")); - - let importer = render_importer(&catalog, "test-hash"); - assert!(importer.contains("cuda_device::thread::sync_threads")); - assert!(importer.contains("cuda_device::sync_threads")); - assert!(importer.contains("Barrier0Op::get_concrete_op_info()")); - assert!(importer.contains("set_generated_intrinsic_marker(ctx, barrier, \"v1:i0034\")")); - assert!(importer.contains("helpers::emit_goto(ctx, *target_idx, barrier")); - - let lowering = render_lowering(&catalog, "test-hash"); - assert!(lowering.contains("impl MirToLlvmConversion for Barrier0Op")); - assert!(lowering.contains("create_i32_const(ctx, rewriter, 0)")); - assert!(lowering.contains("\"llvm_nvvm_barrier_cta_sync_aligned_all\"")); - assert!(lowering.contains("IntrinsicBackend::LlvmNvptx")); - assert!(lowering.contains("IntrinsicBackend::LibNvvm")); - assert!(lowering.contains("\"bar.sync 0;\", \"~{memory}\"")); - - let probe = render_probe(&catalog, record, "test-hash"); - assert!(probe.contains("declare void @llvm.nvvm.barrier.cta.sync.aligned.all(i32)")); - assert!(probe.contains("call void @llvm.nvvm.barrier.cta.sync.aligned.all(i32 0)")); - - let targets = render_targets(&catalog, "test-hash"); - assert!(targets.contains("id: \"sync_threads\", abi_id: \"i0034\"")); - assert!(targets.contains("source_record: \"BARRIER_CTA_SYNC_ALIGNED_ALL_i\"")); - assert!(targets.contains( - "backend: GeneratedIntrinsicBackend::LlvmNvptx, requirement: GeneratedTargetRequirement { minimum_ptx: GeneratedPtxVersion::from_encoded(32), hardware: GeneratedHardwareTarget::AnyOf(&[GeneratedHardwareAlternative::MinimumSm(20)]) }" - )); - assert!(targets.contains( - "backend: GeneratedIntrinsicBackend::LibNvvm, requirement: GeneratedTargetRequirement { minimum_ptx: GeneratedPtxVersion::from_encoded(10), hardware: GeneratedHardwareTarget::AnyOf(&[GeneratedHardwareAlternative::MinimumSm(75)]) }" - )); - assert!( - !record - .selections - .iter() - .any(|selection| selection.source_record.ends_with("_r")) - ); - - let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); - assert!(outputs.contains_key(&PathBuf::from( - "crates/dialect-nvvm/src/ops/generated/sync.rs" - ))); - assert!(outputs.contains_key(&PathBuf::from("crates/cuda-device/src/generated/fence.rs"))); - - let compatibility = render_compat_fence(&catalog, "test-hash"); - for (id, abi_id, op_type, op_name, llvm_identifier, ptx) in [ - ( - "threadfence_block", - "i0298", - "ThreadfenceBlockOp", - "nvvm.threadfence_block", - "llvm_nvvm_membar_cta", - "membar.cta;", - ), - ( - "threadfence", - "i0299", - "ThreadfenceOp", - "nvvm.threadfence", - "llvm_nvvm_membar_gl", - "membar.gl;", - ), - ( - "threadfence_system", - "i0300", - "ThreadfenceSystemOp", - "nvvm.threadfence_system", - "llvm_nvvm_membar_sys", - "membar.sys;", - ), - ] { - let record = sync_intrinsics(&catalog) - .find(|record| record.id == id) - .unwrap(); - assert_eq!(record.rust.abi_id, abi_id); - assert!(compatibility.contains(&format!("pub fn {id}()"))); - assert!(dialect.contains(&format!("pub struct {op_type}"))); - assert!(dialect.contains(&format!("name = \"{op_name}\""))); - assert!(dialect.contains(&format!("{op_type}::register(ctx)"))); - assert!(importer.contains(&format!("cuda_device::fence::{id}"))); - assert!(importer.contains(&format!("cuda_device::{id}"))); - assert!(importer.contains(&format!( - "set_generated_intrinsic_marker(ctx, barrier, \"v1:{abi_id}\")" - ))); - assert!(lowering.contains(&format!("impl MirToLlvmConversion for {op_type}"))); - assert!(lowering.contains(&format!("\"{llvm_identifier}\""))); - - let probe = render_probe(&catalog, record, "test-hash"); - assert!(probe.contains(&format!("declare void @{}()", llvm(record).symbol))); - assert!(probe.contains(&format!("call void @{}()", llvm(record).symbol))); - assert!(targets.contains(&format!("id: \"{id}\", abi_id: \"{abi_id}\""))); - assert!(targets.contains(&format!("asm: \"{ptx}\""))); - } - for ptx in ["membar.cta;", "membar.gl;", "membar.sys;"] { - assert!(!lowering.contains(&format!("\"{ptx}\""))); - } - } - - #[test] - fn basic_mbarrier_rendering_preserves_existing_paths_shapes_and_routes() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - validate_renderable(&catalog).unwrap(); - assert_eq!(mbarrier_basics(&catalog).count(), 4); - - let compatibility = render_compat_mbarrier_basic(&catalog, "test-hash"); - for signature in [ - "pub unsafe fn mbarrier_init(bar: *mut Barrier, expected_count: u32)", - "pub unsafe fn mbarrier_arrive(bar: *const Barrier) -> u64", - "pub unsafe fn mbarrier_test_wait(bar: *const Barrier, token: u64) -> bool", - "pub unsafe fn mbarrier_inval(bar: *mut Barrier)", - ] { - assert!(compatibility.contains(signature)); - } - assert!(compatibility.contains("eight-byte-aligned `Barrier` in shared memory")); - assert!(compatibility.contains("expected_count` must be in `1..=0xFFFFF`")); - let arrive = compatibility.find("pub unsafe fn mbarrier_arrive").unwrap(); - assert!(compatibility[..arrive].ends_with("#[inline(never)]\n")); - assert!(!compatibility.contains("#[must_use]")); - - let dialect_mod = render_dialect_mod(&catalog, "test-hash"); - assert!(dialect_mod.contains("mod mbarrier_basic;")); - assert!(dialect_mod.contains("mbarrier_basic::register(ctx)")); - let dialect = render_dialect_mbarrier_basic(&catalog, "test-hash"); - assert!(dialect.contains("address_space::GENERIC | address_space::SHARED")); - assert!(dialect.contains("mbarrier expected count must be u32")); - assert!(dialect.contains("mbarrier arrival token must be u64")); - assert!(dialect.contains("mbarrier test-wait requires a u64 token and i1 result")); - - let importer = render_importer(&catalog, "test-hash"); - let lowering = render_lowering(&catalog, "test-hash"); - let targets = render_targets(&catalog, "test-hash"); - let raw = render_raw_abi(&catalog, "test-hash").unwrap(); - for record in mbarrier_basics(&catalog) { - assert!(dialect.contains(&format!("pub struct {}", record.dialect.op_type))); - assert!(dialect.contains(&format!("{}::register(ctx)", record.dialect.op_type))); - assert!(importer.contains(&record.rust.canonical_path)); - assert!(importer.contains(&format!( - "set_generated_intrinsic_marker(ctx, mbarrier, {:?})", - intrinsic_marker(&catalog, record) - ))); - assert!(lowering.contains(&format!( - "impl MirToLlvmConversion for {}", - record.dialect.op_type - ))); - assert!(targets.contains(&format!("id: {:?}", record.id))); - assert!(raw.contains(&format!("pub unsafe fn {}", record.rust.abi_id))); - - let probe = render_probe(&catalog, record, "test-hash"); - assert!( - probe.contains("%barrier = addrspacecast ptr %barrier_generic to ptr addrspace(3)") - ); - let mbarrier = record.mbarrier_basic.as_ref().unwrap(); - let llvm_route = record - .backend_lowerings - .iter() - .find(|lowering| lowering.backend == IntrinsicBackend::LlvmNvptx) - .unwrap(); - let libnvvm_route = record - .backend_lowerings - .iter() - .find(|lowering| lowering.backend == IntrinsicBackend::LibNvvm) - .unwrap(); - let expected_llvm_mechanism = match mbarrier.operation { - MbarrierBasicOperation::TestWait => BackendLoweringMechanism::InlinePtx, - _ => BackendLoweringMechanism::TypedNvvm, - }; - assert_eq!(llvm_route.mechanism, expected_llvm_mechanism); - assert_eq!(libnvvm_route.mechanism, BackendLoweringMechanism::InlinePtx); - match mbarrier.operation { - MbarrierBasicOperation::Init => { - assert!( - dialect - .contains("MbarrierInitSharedOp::build(ctx, barrier, expected_count)") - || importer.contains( - "MbarrierInitSharedOp::build(ctx, barrier, expected_count)" - ) - ); - assert!(lowering.contains( - "convert_init(ctx, rewriter, self.get_operation(), operands_info)" - )); - assert!(probe.contains(&format!( - "declare void @{}(ptr addrspace(3), i32)", - llvm(record).symbol - ))); - } - MbarrierBasicOperation::Arrive => { - assert!(importer.contains("MbarrierArriveSharedOp::build(ctx, barrier)")); - assert!(lowering.contains( - "convert_arrive(ctx, rewriter, self.get_operation(), operands_info)" - )); - assert!(probe.contains(&format!( - "declare i64 @{}(ptr addrspace(3))", - llvm(record).symbol - ))); - } - MbarrierBasicOperation::TestWait => { - assert!( - importer.contains("MbarrierTestWaitSharedOp::build(ctx, barrier, token)") - ); - assert!(lowering.contains( - "convert_test_wait(ctx, rewriter, self.get_operation(), operands_info)" - )); - assert!(probe.contains("mbarrier.test_wait.shared.b64")); - assert!(probe.contains("asm sideeffect")); - assert!(probe.contains("attributes #0 = { convergent }")); - assert!(!probe.contains(&format!("declare i1 @{}", llvm(record).symbol))); - } - MbarrierBasicOperation::Inval => { - assert!(importer.contains("MbarrierInvalSharedOp::build(ctx, barrier)")); - assert!(lowering.contains( - "convert_inval(ctx, rewriter, self.get_operation(), operands_info)" - )); - assert!(probe.contains(&format!( - "declare void @{}(ptr addrspace(3))", - llvm(record).symbol - ))); - } - } - } - - let reference = render_reference(&catalog, "test-hash"); - assert!(reference.contains("## Basic mbarrier contracts")); - assert!(reference.contains("inline PTX on both backends")); - assert!(reference.contains("typed NVVM intrinsic with LLVM-NVPTX")); - assert!(reference.contains("inline PTX with libNVVM")); - - let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); - assert!(outputs.contains_key(&PathBuf::from( - "crates/cuda-device/src/generated/mbarrier_basic.rs" - ))); - assert!(outputs.contains_key(&PathBuf::from( - "crates/dialect-nvvm/src/ops/generated/mbarrier_basic.rs" - ))); - } - - #[test] - fn movmatrix_rendering_owns_dialect_import_and_lowering() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let mut catalog = crate::resolve::resolve(&repo_root).unwrap(); - let mut record = catalog - .intrinsics - .iter() - .find(|record| record.id == "packed_atomic_add_f16x2") - .unwrap() - .clone(); - record.id = "movmatrix_trans_b16".into(); - record.operation_key = "movmatrix.m8n8.trans.b16".into(); - record.family = "movmatrix".into(); - record.rust.abi_id = "i0305".into(); - record.rust.module = "matrix".into(); - record.rust.name = "movmatrix_trans_b16".into(); - record.rust.arguments = vec!["u32".into()]; - record.rust.result = "u32".into(); - record.rust.safe = false; - record.rust.must_use = true; - record.rust.canonical_path = "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0305".into(); - record.rust.public_path = "cuda_intrinsics::matrix::movmatrix_trans_b16".into(); - record.rust.compatibility_paths = vec!["cuda_device::wmma::movmatrix_trans_b16".into()]; - record.dialect.op_type = "MovmatrixTransB16Op".into(); - record.dialect.op_name = "nvvm.movmatrix_trans_b16".into(); - record.dialect.operands = vec!["i32".into()]; - record.dialect.results = vec!["i32".into()]; - record.semantics.pure = false; - record.semantics.memory = "inaccessible_read_write".into(); - record.semantics.convergent = true; - record.semantics.execution_scope = "warp".into(); - record.packed_atomic = None; - record.movmatrix = Some(crate::model::Movmatrix { - participation: - crate::model::MovmatrixParticipation::AllWarpLanesSameInstructionNoExitedLanes, - adapter: crate::model::MovmatrixAdapter::PackedB16x2U32ToPackedB16x2U32, - runtime_validation: RuntimeValidation::Unexecuted, - }); - record.lowering = "generated_movmatrix_inline_ptx".into(); - record.expected_ptx = crate::ptx::InstructionPattern::new( - "movmatrix", - &["sync", "aligned", "m8n8", "trans", "b16"], - vec![ - crate::ptx::OperandPattern::Register, - crate::ptx::OperandPattern::Register, - ], - ); - record.summary = "Transposes one packed b16 matrix fragment across a warp.".into(); - catalog.intrinsics.push(record); - - validate_renderable(&catalog).unwrap(); - let dialect_mod = render_dialect_mod(&catalog, "test-hash"); - assert!(dialect_mod.contains("mod movmatrix;")); - assert!(dialect_mod.contains("movmatrix::register(ctx);")); - - let dialect = render_dialect_movmatrix(&catalog, "test-hash"); - assert!(dialect.contains("pub struct MovmatrixTransB16Op;")); - assert!(dialect.contains("name = \"nvvm.movmatrix_trans_b16\"")); - assert!(dialect.contains("MovmatrixTransB16Op::register(ctx);")); - - let compatibility = render_compat_movmatrix(&catalog, "test-hash"); - assert!(compatibility.contains("pub unsafe fn movmatrix_trans_b16(value: u32) -> u32")); - assert!(compatibility.contains("All 32 warp lanes must execute the same call")); - - let importer = render_importer(&catalog, "test-hash"); - assert!(importer.contains("cuda_device::wmma::movmatrix_trans_b16")); - assert!(importer.contains("MovmatrixTransB16Op::build(ctx, arg0)")); - - let lowering = render_lowering(&catalog, "test-hash"); - assert!(lowering.contains("impl MirToLlvmConversion for MovmatrixTransB16Op")); - assert!(lowering.contains("movmatrix.sync.aligned.m8n8.trans.b16 $0, $1;")); - assert!(lowering.contains("\"=r,r\"")); - assert!(lowering.contains("inline_asm_convergent")); - - let probe_record = movmatrix(&catalog).next().unwrap(); - let probe = render_probe(&catalog, probe_record, "test-hash"); - assert!(probe.contains("call i32 asm \"movmatrix.sync.aligned.m8n8.trans.b16 $0, $1;\"")); - assert!(probe.contains("attributes #0 = { convergent }")); - - let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); - assert!(outputs.contains_key(&PathBuf::from( - "crates/cuda-device/src/generated/movmatrix.rs" - ))); - } - - #[test] - fn extended_mbarrier_rendering_preserves_all_manual_contracts() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - validate_renderable(&catalog).unwrap(); - let records = mbarrier_extended(&catalog).collect::>(); - assert_eq!(records.len(), 11); - - 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_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", - "pub unsafe fn fence_proxy_async_shared_cta()", - "pub unsafe fn fence_mbarrier_init_release_cluster()", - "pub unsafe fn fence_proxy_async_generic_release_shared_cta_cluster()", - "pub unsafe fn fence_proxy_async_generic_acquire_shared_cluster_cluster()", - "pub unsafe fn nanosleep(ns: u32)", - ] { - assert!(compatibility.contains(signature), "missing {signature}"); - } - assert!(!compatibility.contains("#[must_use]")); - - let dialect_mod = render_dialect_mod(&catalog, "test-hash"); - assert!(dialect_mod.contains("mod mbarrier_extended;")); - assert!(dialect_mod.contains("pub use mbarrier_extended::*;")); - assert!(dialect_mod.contains("mbarrier_extended::register(ctx);")); - - let dialect = render_dialect_mbarrier_extended(&catalog, "test-hash"); - let importer = render_importer(&catalog, "test-hash"); - let lowering = render_lowering(&catalog, "test-hash"); - let raw = render_raw_abi(&catalog, "test-hash").unwrap(); - for record in &records { - assert!(dialect.contains(&format!("pub struct {}", record.dialect.op_type))); - assert!(dialect.contains(&format!("{}::register(ctx);", record.dialect.op_type))); - assert!(importer.contains(&record.rust.canonical_path)); - assert!(importer.contains(&record.rust.compatibility_paths[0])); - assert!(lowering.contains(&format!( - "impl MirToLlvmConversion for {}", - record.dialect.op_type - ))); - assert!(raw.contains(&format!("pub unsafe fn {}", record.rust.abi_id))); - for route in &record.backend_lowerings { - assert_eq!(route.mechanism, BackendLoweringMechanism::InlinePtx); - } - let (template, constraints) = crate::resolve::mbarrier_extended_inline_recipe( - record.mbarrier_extended.as_ref().unwrap().operation, - ); - assert!(lowering.contains(&format!("{template:?}, {constraints:?}"))); - let probe = render_probe(&catalog, record, "test-hash"); - assert!(probe.contains("asm sideeffect")); - assert!(probe.contains(template)); - assert!(probe.contains("~{memory}")); - assert!(probe.contains("attributes #0 = { convergent }")); - } - - assert!(dialect.contains("address_space::GENERIC | address_space::SHARED")); - assert!(dialect.contains("mbarrier arrival requires u32 bytes and a u64 token")); - assert!(dialect.contains("mbarrier wait requires a u64 token and i1 result")); - assert!(dialect.contains("mbarrier wait requires u32 parity and i1 result")); - assert!(dialect.contains("remote mbarrier address must be u64")); - assert!(dialect.contains("nanosleep duration must be u32")); - - assert!(importer.contains("MbarrierArriveExpectTxSharedOp::build(ctx, barrier, bytes)")); - assert!(importer.contains("MbarrierArriveExpectTxClusterOp::build(ctx, barrier, bytes)")); - assert!( - !importer - .contains("MbarrierArriveExpectTxSharedOp::build(ctx, barrier, tx_count, bytes)") - ); - assert!(importer.contains("MbarrierArriveClusterOp::build(ctx, address)")); - 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("FenceProxyAsyncSharedCtaOp::build(ctx)")); - assert!(importer.contains("FenceMbarrierInitReleaseClusterOp::build(ctx)")); - assert!(importer.contains("NanosleepOp::build(ctx, ns)")); - - assert!(lowering.contains("cast_to_shared_addrspace")); - assert!(lowering.contains("trunc_to_i1")); - assert!(lowering.contains("DefiningEntity::Op(result_op)")); - assert!(lowering.contains("rewriter.erase_operation(ctx, op)")); - assert!(lowering.contains("mbarrier.arrive.release.cluster.shared::cluster.b64 _, [$0];")); - let remote = records - .iter() - .find(|record| record.id == "mbarrier_arrive_cluster") - .unwrap(); - assert!(remote.llvm.is_none()); - assert!(matches!(remote.source, IntrinsicSource::PtxNative { .. })); - assert!(!render_probe(&catalog, remote, "test-hash").contains("declare")); - - let reference = render_reference(&catalog, "test-hash"); - assert!(reference.contains("## Extended mbarrier contracts")); - assert!(reference.contains("LLVM 22 TableGen")); - assert!(reference.contains("PTX-native raw-address contract")); - - let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); - assert!(outputs.contains_key(&PathBuf::from( - "crates/cuda-device/src/generated/mbarrier_extended.rs" - ))); - assert!(outputs.contains_key(&PathBuf::from( - "crates/dialect-nvvm/src/ops/generated/mbarrier_extended.rs" - ))); - } - - #[test] - 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(), 1016); - let records: Vec<_> = register_mmas(&catalog).collect(); - assert_eq!(records.len(), 154); - let generated_records = records - .iter() - .copied() - .filter(|record| { - record.register_mma.as_ref().unwrap().compatibility_source - == RegisterMmaCompatibilitySource::GeneratedStub - }) - .collect::>(); - let existing_records = records - .iter() - .copied() - .filter(|record| { - record.register_mma.as_ref().unwrap().compatibility_source - == RegisterMmaCompatibilitySource::ExistingStub - }) - .collect::>(); - assert_eq!(generated_records.len(), 149); - assert_eq!(existing_records.len(), 5); - - let raw = render_raw_abi(&catalog, "test-hash").unwrap(); - for record in &records { - assert!(raw.contains(&format!("pub unsafe fn {}(", record.rust.abi_id))); - } - assert!(raw.contains("no lane may have exited")); - assert!(raw.contains("Signed accumulator overflow wraps")); - assert!(raw.contains("Signed accumulator overflow clamps")); - - let compatibility = render_compat_register_mma(&catalog, "test-hash"); - assert_eq!(compatibility.matches("pub unsafe fn ").count(), 149); - for record in generated_records { - let argument_names: &[&str] = if record.register_mma.as_ref().unwrap().adapter - == RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 - { - &[ - "c", - "a", - "b", - "scale_a", - "byte_id_a", - "thread_id_a", - "scale_b", - "byte_id_b", - "thread_id_b", - ] - } else { - &["c", "a", "b"] - }; - let arguments = argument_names - .iter() - .zip(&record.rust.arguments) - .map(|(name, ty)| format!("{name}: {ty}")) - .collect::>() - .join(", "); - assert!(compatibility.contains(&format!( - "pub unsafe fn {}({arguments}) -> {} {{", - record.rust.name, record.rust.result - ))); - } - for record in existing_records { - assert!(!compatibility.contains(&format!("pub unsafe fn {}(", record.rust.name))); - } - - let dialect = render_dialect_register_mma(&catalog, "test-hash"); - assert_eq!(dialect.matches("pub struct RegisterMmaOp").count(), 1); - assert!(dialect.contains("let recipe: (&[MmaCarrier], &[MmaCarrier])")); - assert!(dialect.contains("RegisterMmaShapeAttr::M8n8k4")); - assert!(dialect.contains("RegisterMmaShapeAttr::M8n8k16")); - assert!(dialect.contains("RegisterMmaShapeAttr::M8n8k32")); - assert!(dialect.contains("RegisterMmaShapeAttr::M8n8k128")); - assert!(dialect.contains("RegisterMmaShapeAttr::M16n8k4")); - assert!(dialect.contains("RegisterMmaShapeAttr::M16n8k32")); - assert!(dialect.contains("RegisterMmaShapeAttr::M16n8k64")); - assert!(dialect.contains("RegisterMmaShapeAttr::M16n8k128")); - assert!(dialect.contains("RegisterMmaShapeAttr::M16n8k256")); - assert!(dialect.contains("RegisterMmaOperationAttr::Multiply")); - assert!(dialect.contains("RegisterMmaOperationAttr::AndPopc")); - assert!(dialect.contains("RegisterMmaOperationAttr::XorPopc")); - assert!(dialect.contains("pub enum RegisterMmaKindAttr { Standard, F8f6f4, Mxf8f6f4 }")); - assert!(dialect.contains("kind_or_inferred")); - assert!(dialect.contains("RegisterMmaAccumulatorAttr::F16")); - assert!(dialect.contains("operation_or_multiply")); - assert!(dialect.contains("MmaCarrier::I32 | MmaCarrier::U16 | MmaCarrier::U32")); - assert!(dialect.contains("if matches!(carrier, MmaCarrier::U16)")); - assert!(dialect.contains("RegisterMmaElementAttr::B1")); - assert!(dialect.contains("RegisterMmaElementAttr::E2m1")); - assert!(dialect.contains("RegisterMmaElementAttr::E5m2")); - assert!(dialect.contains("RegisterMmaElementAttr::S4")); - assert!(dialect.contains("RegisterMmaElementAttr::U4")); - assert!(dialect.contains("RegisterMmaElementAttr::U8")); - assert!(dialect.contains("RegisterMmaOverflowAttr::Wrapping")); - assert!(dialect.contains("RegisterMmaOverflowAttr::Satfinite")); - assert!(dialect.contains( - "&[MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32]" - )); - for (op_type, op_name) in [ - ("MmaM16N8K16F32Bf16Op", "nvvm.mma_m16n8k16_f32_bf16"), - ("MmaM16N8K16F32F16Op", "nvvm.mma_m16n8k16_f32_f16"), - ("MmaM16N8K8F32Tf32Op", "nvvm.mma_m16n8k8_f32_tf32"), - ("MmaM16N8K32S32S8Op", "nvvm.mma_m16n8k32_s32_s8"), - ("MmaM8N8K4F64Op", "nvvm.mma_m8n8k4_f64"), - ] { - assert!(dialect.contains(&format!("pub struct {op_type};"))); - assert!(dialect.contains(&format!("name = {op_name:?}"))); - assert!(dialect.contains(&format!("{op_type}::register(ctx);"))); - } - - let importer = render_importer(&catalog, "test-hash"); - assert!(importer.contains("enum GeneratedMmaImportAdapter")); - assert!(importer.contains("C2U32A4U32B2U32ToD2U32")); - assert!(importer.contains("C2U32A2U32B1U32ToD2U32")); - assert!(importer.contains("C4F32A2U32B1U32ToD4F32")); - assert!(importer.contains("C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32")); - assert!( - importer.contains("ctx, b_value, b_ty, b_count, block_ptr, last_after_b, loc.clone()") - ); - assert!(importer.contains("C4I32A2U32B1U32ToD4I32")); - assert!(importer.contains("C2I32A1U32B1U32ToD2I32")); - assert!(importer.contains("(i32_ty, 2, u32_ty, 1, false, u32_ty, 1, false, i32_ty, 2)")); - assert!(importer.contains("(i32_ty, 4, u32_ty, 2, true, u32_ty, 1, false, i32_ty, 4)")); - assert!(importer.contains("(u32_ty, 2, u32_ty, 4, true, u32_ty, 2, true, u32_ty, 2)")); - assert!(importer.contains("import_generated_mma_operands")); - assert!(importer.contains("bundle_generated_mma_results")); - assert!(importer.contains("set_attr_nvvm_register_mma_operation")); - assert!(importer.contains("set_attr_nvvm_register_mma_kind")); - for record in &records { - assert!(importer.contains(&record.rust.canonical_path)); - assert!(importer.contains(&record.rust.compatibility_paths[0])); - assert!(importer.contains(&format!( - "set_generated_intrinsic_marker(ctx, mma, {:?})", - intrinsic_marker(&catalog, record) - ))); - } - - let lowering = render_lowering(&catalog, "test-hash"); - assert_eq!( - lowering - .matches("impl MirToLlvmConversion for RegisterMmaOp") - .count(), - 1 - ); - assert!(lowering.contains("convert_generated_register_mma")); - assert!(lowering.contains("operation_or_multiply")); - assert!(lowering.contains("kind_or_inferred")); - assert!(lowering.contains("mma.sync.aligned.m16n8k16.row.col.f16.e4m3.e4m3.f16")); - assert!(lowering.contains("mma.sync.aligned.m16n8k4.row.col.f32.tf32.tf32.f32")); - assert!(lowering.contains("mma.sync.aligned.m16n8k8.row.col.f32.bf16.bf16.f32")); - assert!(lowering.contains("mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16")); - assert!(lowering.contains("mma.sync.aligned.m16n8k32.row.col.f32.e5m2.e5m2.f32")); - assert!(lowering.contains("=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r")); - assert!(lowering.contains("=d,=d,d,d,d,d")); - assert!(lowering.contains("=r,=r,=r,=r,r,r,r,r,r,r,r")); - assert!(lowering.contains("=r,=r,r,r,r,r")); - assert!(lowering.contains("=r,=r,r,r,r,r,r,r,r,r")); - assert!(lowering.contains("mma.sync.aligned.m8n8k16.row.col.s32.s8.u8.s32")); - assert!(lowering.contains("mma.sync.aligned.m8n8k16.row.col.satfinite.s32.u8.s8.s32")); - assert!(lowering.contains("mma.sync.aligned.m8n8k32.row.col.s32.s4.u4.s32")); - assert!(lowering.contains("mma.sync.aligned.m8n8k32.row.col.satfinite.s32.u4.s4.s32")); - assert!(lowering.contains("mma.sync.aligned.m16n8k16.row.col.s32.s8.u8.s32")); - assert!(lowering.contains("mma.sync.aligned.m16n8k32.row.col.satfinite.s32.u8.s8.s32")); - assert!(lowering.contains("mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32")); - assert!(lowering.contains("mma.sync.aligned.m16n8k32.row.col.s32.s4.u4.s32")); - assert!(lowering.contains("mma.sync.aligned.m16n8k32.row.col.satfinite.s32.u4.s4.s32")); - assert!(lowering.contains("mma.sync.aligned.m16n8k64.row.col.s32.s4.u4.s32")); - assert!(lowering.contains("mma.sync.aligned.m16n8k64.row.col.satfinite.s32.u4.s4.s32")); - assert!(lowering.contains("mma.sync.aligned.m8n8k128.row.col.s32.b1.b1.s32.xor.popc")); - assert!(lowering.contains("mma.sync.aligned.m8n8k128.row.col.s32.b1.b1.s32.and.popc")); - assert!(lowering.contains("mma.sync.aligned.m16n8k128.row.col.s32.b1.b1.s32.xor.popc")); - assert!(lowering.contains("mma.sync.aligned.m16n8k128.row.col.s32.b1.b1.s32.and.popc")); - assert!(lowering.contains("mma.sync.aligned.m16n8k256.row.col.s32.b1.b1.s32.xor.popc")); - assert!(lowering.contains("mma.sync.aligned.m16n8k256.row.col.s32.b1.b1.s32.and.popc")); - for op_type in [ - "MmaM16N8K16F32Bf16Op", - "MmaM16N8K16F32F16Op", - "MmaM16N8K8F32Tf32Op", - "MmaM16N8K32S32S8Op", - "MmaM8N8K4F64Op", - ] { - assert_eq!( - lowering - .matches(&format!("impl MirToLlvmConversion for {op_type}")) - .count(), - 1 - ); - } - assert!(lowering.contains( - r#"(GeneratedMmaResultType::I32, 4, 7, "mma.sync.aligned.m16n8k32.row.col.s32.s4.u4.s32 {$0, $1, $2, $3}, {$8, $9}, {$10}, {$4, $5, $6, $7};", "=r,=r,=r,=r,r,r,r,r,r,r,r")"# - )); - assert!(lowering.contains( - r#"(GeneratedMmaResultType::I32, 4, 10, "mma.sync.aligned.m16n8k64.row.col.satfinite.s32.u4.s4.s32 {$0, $1, $2, $3}, {$8, $9, $10, $11}, {$12, $13}, {$4, $5, $6, $7};", "=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,r")"# - )); - assert!(lowering.contains( - r#"(GeneratedMmaResultType::I32, 2, 4, "mma.sync.aligned.m8n8k128.row.col.s32.b1.b1.s32.xor.popc {$0, $1}, {$4}, {$5}, {$2, $3};", "=r,=r,r,r,r,r")"# - )); - assert!(lowering.contains( - r#"(GeneratedMmaResultType::I32, 4, 7, "mma.sync.aligned.m16n8k128.row.col.s32.b1.b1.s32.and.popc {$0, $1, $2, $3}, {$8, $9}, {$10}, {$4, $5, $6, $7};", "=r,=r,=r,=r,r,r,r,r,r,r,r")"# - )); - assert!(lowering.contains( - r#"(GeneratedMmaResultType::I32, 4, 10, "mma.sync.aligned.m16n8k256.row.col.s32.b1.b1.s32.xor.popc {$0, $1, $2, $3}, {$8, $9, $10, $11}, {$12, $13}, {$4, $5, $6, $7};", "=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,r")"# - )); - assert!(lowering.contains( - r#"(GeneratedMmaResultType::F32, 4, 10, "mma.sync.aligned.m16n8k32.row.col.kind::f8f6f4.f32.e2m1.e2m1.f32 {$0, $1, $2, $3}, {$8, $9, $10, $11}, {$12, $13}, {$4, $5, $6, $7};", "=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r")"# - )); - assert!(lowering.contains( - r#"(GeneratedMmaResultType::F32, 4, 10, "mma.sync.aligned.m16n8k32.row.col.kind::f8f6f4.f32.e5m2.e5m2.f32 {$0, $1, $2, $3}, {$8, $9, $10, $11}, {$12, $13}, {$4, $5, $6, $7};", "=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r")"# - )); - assert!(lowering.contains( - r#"(GeneratedMmaResultType::I32, 2, 8, "mma.sync.aligned.m16n8k32.row.col.kind::f8f6f4.f16.e2m1.e2m1.f16 {$0, $1}, {$4, $5, $6, $7}, {$8, $9}, {$2, $3};", "=r,=r,r,r,r,r,r,r,r,r")"# - )); - assert!(lowering.contains( - r#"(GeneratedMmaResultType::I32, 2, 8, "mma.sync.aligned.m16n8k32.row.col.kind::f8f6f4.f16.e5m2.e5m2.f16 {$0, $1}, {$4, $5, $6, $7}, {$8, $9}, {$2, $3};", "=r,=r,r,r,r,r,r,r,r,r")"# - )); - assert!(lowering.contains( - r#"(GeneratedMmaResultType::F32, 4, 16, "mma.sync.aligned.m16n8k32.row.col.kind::mxf8f6f4.block_scale.f32.e2m1.e2m1.f32.ue8m0 {$0, $1, $2, $3}, {$8, $9, $10, $11}, {$12, $13}, {$4, $5, $6, $7}, $14, {$15, $16}, $17, {$18, $19};", "=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r,r,h,h,r,h,h")"# - )); - - let targets = render_targets(&catalog, "test-hash"); - assert!(targets.contains("GeneratedIntrinsicVariant::RegisterMma")); - assert!(targets.contains("GeneratedRegisterMmaShape::M8n8k4")); - assert!(targets.contains("GeneratedRegisterMmaShape::M16n8k4")); - assert!(targets.contains( - "pub enum GeneratedRegisterMmaShape { M8n8k4, M8n8k16, M8n8k32, M8n8k128, M16n8k4, M16n8k8, M16n8k16, M16n8k32, M16n8k64, M16n8k128, M16n8k256 }" - )); - assert!(targets.contains( - "GeneratedRegisterMmaShape::M8n8k16 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M8n8k16)" - )); - assert!(targets.contains( - "GeneratedRegisterMmaShape::M8n8k32 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M8n8k32)" - )); - assert!(targets.contains( - "GeneratedRegisterMmaShape::M8n8k128 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M8n8k128)" - )); - assert!(targets.contains("GeneratedRegisterMmaShape::M16n8k32")); - assert!(targets.contains( - "GeneratedRegisterMmaShape::M16n8k64 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M16n8k64)" - )); - assert!(targets.contains( - "GeneratedRegisterMmaShape::M16n8k128 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M16n8k128)" - )); - assert!(targets.contains( - "GeneratedRegisterMmaShape::M16n8k256 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M16n8k256)" - )); - assert!(targets.contains("GeneratedRegisterMmaOperation::Multiply")); - assert!(targets.contains("GeneratedRegisterMmaOperation::AndPopc")); - assert!(targets.contains("GeneratedRegisterMmaOperation::XorPopc")); - assert!( - targets.contains("pub enum GeneratedRegisterMmaKind { Standard, F8f6f4, Mxf8f6f4 }") - ); - assert!(targets.contains("kind: GeneratedRegisterMmaKind::Standard")); - assert!(targets.contains("kind: GeneratedRegisterMmaKind::F8f6f4")); - assert!(targets.contains("kind: GeneratedRegisterMmaKind::Mxf8f6f4")); - assert!(targets.contains("kind_or_inferred")); - assert!(targets.contains("RegisterMmaKindAttr::Standard")); - assert!(targets.contains("RegisterMmaKindAttr::F8f6f4")); - assert!(targets.contains("RegisterMmaKindAttr::Mxf8f6f4")); - assert!(targets.contains("GeneratedRegisterMmaAccumulator::F16")); - assert!(targets.contains("operation: GeneratedRegisterMmaOperation::AndPopc")); - assert!(targets.contains("operation: GeneratedRegisterMmaOperation::XorPopc")); - assert!(targets.contains("operation: mma_operation")); - assert!(targets.contains("operation_or_multiply")); - assert!(targets.contains("GeneratedRegisterMmaElement::B1")); - assert!(targets.contains("GeneratedRegisterMmaElement::E2m1")); - assert!(targets.contains("GeneratedRegisterMmaElement::E5m2")); - assert!(targets.contains("GeneratedRegisterMmaElement::S4")); - assert!(targets.contains("GeneratedRegisterMmaElement::U4")); - assert!(targets.contains("GeneratedRegisterMmaElement::U8")); - assert!(targets.contains( - "a_element: GeneratedRegisterMmaElement::S4, b_element: GeneratedRegisterMmaElement::U4" - )); - assert!(targets.contains( - "a_element: GeneratedRegisterMmaElement::U4, b_element: GeneratedRegisterMmaElement::S4" - )); - assert!(targets.contains( - "a_element: GeneratedRegisterMmaElement::S8, b_element: GeneratedRegisterMmaElement::U8" - )); - assert!(targets.contains( - "a_element: GeneratedRegisterMmaElement::U8, b_element: GeneratedRegisterMmaElement::S8" - )); - assert!(targets.contains("overflow: GeneratedRegisterMmaOverflow::Satfinite")); - assert!(targets.contains("get_attr_nvvm_register_mma_overflow")); - assert!(targets.contains("minimum_ptx: GeneratedPtxVersion::from_encoded(87)")); - assert!(targets.contains("minimum_ptx: GeneratedPtxVersion::from_encoded(84)")); - assert!(targets.contains("GeneratedHardwareAlternative::ExactArchitecture(120)")); - assert!(targets.contains("GeneratedHardwareAlternative::FamilyTarget(121)")); - - let first_dense = records - .iter() - .find(|record| record.id == "mma_m16n8k32_f32_e2m1_e2m1") - .unwrap(); - let last_dense = records - .iter() - .find(|record| record.id == "mma_m16n8k32_f32_e5m2_e5m2") - .unwrap(); - assert_eq!(first_dense.rust.abi_id, "i0454"); - assert_eq!(last_dense.rust.abi_id, "i0478"); - assert_eq!( - register_mma_constraints(first_dense), - "=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r" - ); - let first_dense_f16 = records - .iter() - .find(|record| record.id == "mma_m16n8k32_f16_e2m1_e2m1") - .unwrap(); - let last_dense_f16 = records - .iter() - .find(|record| record.id == "mma_m16n8k32_f16_e5m2_e5m2") - .unwrap(); - assert_eq!(first_dense_f16.rust.abi_id, "i0479"); - assert_eq!(last_dense_f16.rust.abi_id, "i0503"); - assert_eq!( - register_mma_constraints(first_dense_f16), - "=r,=r,r,r,r,r,r,r,r,r" - ); - assert!( - render_probe(&catalog, first_dense_f16, "test-hash").contains("define { i32, i32 }") - ); - - let first_mxf8f6f4 = records - .iter() - .find(|record| record.id == "mma_m16n8k32_mxf8f6f4_f32_e2m1_e2m1") - .unwrap(); - let last_mxf8f6f4 = records - .iter() - .find(|record| record.id == "mma_m16n8k32_mxf8f6f4_f32_e5m2_e5m2") - .unwrap(); - assert_eq!(first_mxf8f6f4.rust.abi_id, "i0858"); - assert_eq!(last_mxf8f6f4.rust.abi_id, "i0882"); - assert_eq!( - register_mma_constraints(first_mxf8f6f4), - "=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r,r,h,h,r,h,h" - ); - assert!( - render_probe(&catalog, first_mxf8f6f4, "test-hash") - .contains("kind::mxf8f6f4.block_scale") - ); - let first_mxf8f6f4_compatibility = format!("pub unsafe fn {}(", first_mxf8f6f4.rust.name); - let first_mxf8f6f4_compatibility = compatibility - .find(&first_mxf8f6f4_compatibility) - .expect("generated MXF8F6F4 compatibility wrapper"); - assert!( - compatibility[..first_mxf8f6f4_compatibility] - .ends_with("#[allow(clippy::too_many_arguments)]\n#[must_use]\n#[inline(never)]\n") - ); - - for record in &records { - let probe = render_probe(&catalog, record, "test-hash"); - assert!(probe.contains("asm sideeffect")); - assert!(probe.contains("attributes #0 = { convergent }")); - assert!(probe.contains(®ister_mma_template(record))); - assert!(probe.contains(®ister_mma_constraints(record))); - } - - let reference = render_reference(&catalog, "test-hash"); - assert!(reference.contains("## Register-MMA contracts")); - assert!(reference.contains("performs XOR, population count, and accumulate")); - assert!(reference.contains("performs AND, population count, and accumulate")); - assert!(reference.contains("runtime validation is not executed on a GPU")); - for record in &records { - assert!(reference.contains(&format!("- `{}`: runtime `unexecuted`", record.id))); - } - - let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); - assert!(outputs.contains_key(&PathBuf::from( - "crates/cuda-device/src/generated/register_mma.rs" - ))); - assert!(outputs.contains_key(&PathBuf::from( - "crates/dialect-nvvm/src/ops/generated/register_mma.rs" - ))); - } - - #[test] - fn register_mma_kind_is_explicit_and_legacy_ir_stays_compatible() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let mut catalog: CatalogFile = - read_json(&repo_root.join("intrinsics/catalog.json")).unwrap(); - let old = catalog - .intrinsics - .iter() - .find(|record| record.id == "mma_m16n8k32_f16_e4m3_e4m3") - .unwrap() - .clone(); - assert_eq!(old.register_mma.as_ref().unwrap().kind, None); - assert_eq!( - register_mma_attr_variants(&old).2, - "RegisterMmaKindAttr::F8f6f4" - ); - assert!( - generated_intrinsic_variant(&old).contains("kind: GeneratedRegisterMmaKind::F8f6f4") - ); - - let mut standard = old.clone(); - standard.id = "mma_m16n8k32_fp8_f16_e4m3_e4m3".into(); - standard.operation_key = - "matrix.mma.m16n8k32.row.col.standard_fp8.f16.e4m3.e4m3.f16".into(); - standard.rust.abi_id = "i9999".into(); - standard.rust.name = standard.id.clone(); - standard.rust.canonical_path = format!("cuda_intrinsics::matrix::{}", standard.id); - standard.rust.public_path = standard.rust.canonical_path.clone(); - standard.rust.compatibility_paths = vec![format!("cuda_device::wmma::{}", standard.id)]; - standard.llvm.as_mut().unwrap().symbol = - "llvm.nvvm.mma.m16n8k32.row.col.f16.e4m3.e4m3.f16".into(); - standard.register_mma.as_mut().unwrap().kind = Some(RegisterMmaKind::Standard); - standard - .expected_ptx - .modifiers - .retain(|modifier| modifier != "kind::f8f6f4"); - assert_eq!( - register_mma_attr_variants(&standard).2, - "RegisterMmaKindAttr::Standard" - ); - assert!( - generated_intrinsic_variant(&standard) - .contains("kind: GeneratedRegisterMmaKind::Standard") - ); - catalog.intrinsics.push(standard); - - let dialect = render_dialect_register_mma(&catalog, "test-hash"); - assert!(dialect.contains("pub fn kind_or_inferred")); - assert!( - dialect - .contains("let old_f8f6f4 = self.get_attr_nvvm_register_mma_shape(ctx).as_deref()") - ); - assert!(dialect.contains("== Some(&RegisterMmaShapeAttr::M16n8k32)")); - assert!( - dialect.contains( - "&& low_format(self.get_attr_nvvm_register_mma_a_element(ctx).as_deref())" - ) - ); - assert!( - dialect.contains( - "&& low_format(self.get_attr_nvvm_register_mma_b_element(ctx).as_deref())" - ) - ); - assert!(dialect.contains("RegisterMmaKindAttr::F8f6f4")); - assert!(dialect.contains("RegisterMmaKindAttr::Standard")); - - let importer = render_importer(&catalog, "test-hash"); - assert!( - importer.contains("set_attr_nvvm_register_mma_kind(ctx, RegisterMmaKindAttr::F8f6f4)") - ); - assert!( - importer - .contains("set_attr_nvvm_register_mma_kind(ctx, RegisterMmaKindAttr::Standard)") - ); - - let lowering = render_lowering(&catalog, "test-hash"); - assert!(lowering.contains("let kind = self.kind_or_inferred(ctx)")); - assert!(lowering.contains("kind::f8f6f4.f16.e4m3.e4m3.f16")); - assert!(lowering.contains("row.col.f16.e4m3.e4m3.f16")); - - let targets = render_targets(&catalog, "test-hash"); - assert!(targets.contains("kind: GeneratedRegisterMmaKind::F8f6f4")); - assert!(targets.contains("kind: GeneratedRegisterMmaKind::Standard")); - assert!(targets.contains( - "GeneratedRegisterMmaKind::Standard => op.kind_or_inferred(ctx) == RegisterMmaKindAttr::Standard" - )); - assert!(targets.contains( - "GeneratedRegisterMmaKind::F8f6f4 => op.kind_or_inferred(ctx) == RegisterMmaKindAttr::F8f6f4" - )); - assert!(targets.contains("kind_matches && accumulator_matches")); - } - - #[test] - fn sparse_mma_rendering_enforces_selector_and_keeps_family_distinct() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - validate_renderable(&catalog).unwrap(); - let records = sparse_mmas(&catalog).collect::>(); - let k32 = records - .iter() - .filter(|record| record.sparse_mma.as_ref().unwrap().shape == SparseMmaShape::M16n8k32) - .count(); - let k64 = records - .iter() - .filter(|record| record.sparse_mma.as_ref().unwrap().shape == SparseMmaShape::M16n8k64) - .count(); - let k128 = records - .iter() - .filter(|record| record.sparse_mma.as_ref().unwrap().shape == SparseMmaShape::M16n8k128) - .count(); - assert_eq!((records.len(), k32, k64, k128), (114, 16, 82, 16)); - let standard_k64 = records - .iter() - .copied() - .find(|record| record.id == "mma_sp_m16n8k64_s32_s8") - .unwrap(); - assert_eq!( - standard_k64.sparse_mma.as_ref().unwrap().adapter, - SparseMmaAdapter::C4I32A4U32B4U32MetadataU32SelectorU32ToD4I32 - ); - assert_eq!( - standard_k64.sparse_mma.as_ref().unwrap().llvm_adapter, - crate::model::SparseMmaLlvmAdapter::A4I32B4I32C4I32MetadataI32SelectorI32ToD4I32 - ); - assert_eq!( - sparse_mma_template(standard_k64), - "mma.sp.sync.aligned.m16n8k64.row.col.s32.s8.s8.s32 {$0, $1, $2, $3}, {$8, $9, $10, $11}, {$12, $13, $14, $15}, {$4, $5, $6, $7}, $16, $17;" - ); - assert_eq!( - sparse_mma_constraints(standard_k64), - "=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,r,r,r,r,n" - ); - let standard_int4_k64 = records - .iter() - .copied() - .find(|record| record.id == "mma_sp_m16n8k64_s32_s4") - .unwrap(); - assert_eq!( - sparse_mma_template(standard_int4_k64), - "mma.sp.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 {$0, $1, $2, $3}, {$8, $9}, {$10, $11}, {$4, $5, $6, $7}, $12, $13;" - ); - assert_eq!( - sparse_mma_constraints(standard_int4_k64), - "=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,n" - ); - let ordered_int4_k64 = records - .iter() - .copied() - .find(|record| record.id == "mma_sp_ordered_metadata_m16n8k64_s32_s4") - .unwrap(); - assert_eq!( - ordered_int4_k64.sparse_mma.as_ref().unwrap().adapter, - SparseMmaAdapter::C4I32A2U32B2U32MetadataU32SelectorU32ToD4I32 - ); - assert_eq!( - ordered_int4_k64.sparse_mma.as_ref().unwrap().llvm_adapter, - crate::model::SparseMmaLlvmAdapter::A2I32B2I32C4I32MetadataI32SelectorI32ToD4I32 - ); - assert_eq!( - sparse_mma_template(ordered_int4_k64), - "mma.sp::ordered_metadata.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 {$0, $1, $2, $3}, {$8, $9}, {$10, $11}, {$4, $5, $6, $7}, $12, $13;" - ); - assert_eq!( - sparse_mma_constraints(ordered_int4_k64), - "=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,n" - ); - let ordered_int4_k128 = records - .iter() - .copied() - .find(|record| record.id == "mma_sp_ordered_metadata_m16n8k128_s32_s4") - .unwrap(); - assert_eq!( - ordered_int4_k128.sparse_mma.as_ref().unwrap().adapter, - SparseMmaAdapter::C4I32A4U32B4U32MetadataU32SelectorU32ToD4I32 - ); - assert_eq!( - ordered_int4_k128.sparse_mma.as_ref().unwrap().llvm_adapter, - crate::model::SparseMmaLlvmAdapter::A4I32B4I32C4I32MetadataI32SelectorI32ToD4I32 - ); - assert_eq!( - sparse_mma_template(ordered_int4_k128), - "mma.sp::ordered_metadata.sync.aligned.m16n8k128.row.col.s32.s4.s4.s32 {$0, $1, $2, $3}, {$8, $9, $10, $11}, {$12, $13, $14, $15}, {$4, $5, $6, $7}, $16, $17;" - ); - assert_eq!( - sparse_mma_constraints(ordered_int4_k128), - "=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,r,r,r,r,n" - ); - let standard_int4_k128 = records - .iter() - .copied() - .find(|record| record.id == "mma_sp_m16n8k128_s32_s4") - .unwrap(); - assert_eq!( - sparse_mma_template(standard_int4_k128), - "mma.sp.sync.aligned.m16n8k128.row.col.s32.s4.s4.s32 {$0, $1, $2, $3}, {$8, $9, $10, $11}, {$12, $13, $14, $15}, {$4, $5, $6, $7}, $16, $17;" - ); - assert_eq!( - sparse_mma_constraints(standard_int4_k128), - "=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,r,r,r,r,n" - ); - let ordered_f8f6f4 = records - .iter() - .copied() - .find(|record| { - record.id == "mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m1_e2m1_f32" - }) - .unwrap(); - assert_eq!( - ordered_f8f6f4.sparse_mma.as_ref().unwrap().adapter, - SparseMmaAdapter::C4F32A4U32B4U32MetadataU32SelectorU32ToD4F32 - ); - assert_eq!( - sparse_mma_template(ordered_f8f6f4), - "mma.sp::ordered_metadata.sync.aligned.m16n8k64.row.col.kind::f8f6f4.f32.e2m1.e2m1.f32 {$0, $1, $2, $3}, {$8, $9, $10, $11}, {$12, $13, $14, $15}, {$4, $5, $6, $7}, $16, $17;" - ); - assert_eq!( - sparse_mma_constraints(ordered_f8f6f4), - "=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r,r,r,r,n" - ); - assert_eq!( - sparse_mma_carriers(ordered_f8f6f4), - ( - "&[MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32]".into(), - "&[MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32]".into(), - ) - ); - let ordered_f8f6f4_f16 = records - .iter() - .copied() - .find(|record| { - record.id == "mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e2m1_f16" - }) - .unwrap(); - assert_eq!( - ordered_f8f6f4_f16.sparse_mma.as_ref().unwrap().adapter, - SparseMmaAdapter::C2U32A4U32B4U32MetadataU32SelectorU32ToD2U32 - ); - assert_eq!(sparse_mma_fragment_counts(ordered_f8f6f4_f16), (2, 4, 4, 2)); - assert_eq!( - sparse_mma_template(ordered_f8f6f4_f16), - "mma.sp::ordered_metadata.sync.aligned.m16n8k64.row.col.kind::f8f6f4.f16.e2m1.e2m1.f16 {$0, $1}, {$4, $5, $6, $7}, {$8, $9, $10, $11}, {$2, $3}, $12, $13;" - ); - assert_eq!( - sparse_mma_constraints(ordered_f8f6f4_f16), - "=r,=r,r,r,r,r,r,r,r,r,r,r,r,n" - ); - assert_eq!( - sparse_mma_carriers(ordered_f8f6f4_f16), - ( - "&[MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32]".into(), - "&[MmaCarrier::U32, MmaCarrier::U32]".into(), - ) - ); - assert_eq!(sparse_mma_selector_values(ordered_f8f6f4_f16), &[0]); - - let raw = render_raw_abi(&catalog, "test-hash").unwrap(); - assert!(raw.contains("must be the compile-time constant `0` or `1`")); - assert!(raw.contains("must be the compile-time constant `0`")); - for record in &records { - let (c_count, a_count, b_count, d_count) = sparse_mma_fragment_counts(record); - let scalar = match record.sparse_mma.as_ref().unwrap().accumulator { - SparseMmaAccumulator::F16 => "u32", - SparseMmaAccumulator::F32 => "f32", - SparseMmaAccumulator::S32 => "i32", - }; - assert!(raw.contains(&format!( - "pub unsafe fn {}(_arg0: [{scalar}; {c_count}], _arg1: [u32; {a_count}], _arg2: [u32; {b_count}], _arg3: u32, _arg4: u32) -> [{scalar}; {d_count}]", - record.rust.abi_id, - ))); - } - - let compatibility = render_compat_sparse_mma(&catalog, "test-hash"); - assert_eq!(compatibility.matches("pub unsafe fn ").count(), 114); - assert!( - compatibility - .contains("c: [i32; 4], a: [u32; 2], b: [u32; 2], metadata: u32, selector: u32") - ); - assert!( - compatibility - .contains("c: [i32; 4], a: [u32; 4], b: [u32; 4], metadata: u32, selector: u32") - ); - assert!( - compatibility - .contains("c: [f32; 4], a: [u32; 4], b: [u32; 4], metadata: u32, selector: u32") - ); - assert!( - compatibility - .contains("c: [u32; 2], a: [u32; 4], b: [u32; 4], metadata: u32, selector: u32") - ); - - let dialect = render_dialect_sparse_mma(&catalog, "test-hash"); - assert_eq!(dialect.matches("pub struct SparseMmaOp").count(), 1); - assert!(dialect.contains("MmaCarrier::F32")); - assert!(dialect.contains("MmaCarrier::I32")); - assert!(dialect.contains("MmaCarrier::U32")); - assert!(dialect.contains("operands.len() != expected_operands")); - assert!(dialect.contains("SparseMmaShapeAttr { M16n8k32, M16n8k64, M16n8k128 }")); - assert!(dialect.contains("SparseMmaAccumulatorAttr { F16, F32, S32 }")); - assert!(dialect.contains("SparseMmaSelectorAttr { ImmediateZeroOrOne, ImmediateZero }")); - assert!(dialect.contains("SparseMmaElementAttr::E2m1")); - assert!(dialect.contains("SparseMmaElementAttr::E2m3")); - assert!(dialect.contains("SparseMmaElementAttr::E3m2")); - assert!(dialect.contains("SparseMmaElementAttr::E4m3")); - assert!(dialect.contains("SparseMmaElementAttr::E5m2")); - assert!(dialect.contains("SparseMmaElementAttr::S4")); - assert!(dialect.contains("SparseMmaElementAttr::U4")); - assert!(dialect.contains("SparseMmaElementAttr::S8")); - assert!(dialect.contains("SparseMmaElementAttr::U8")); - assert!(dialect.contains("SparseMmaOverflowAttr::Wrapping")); - assert!(dialect.contains("SparseMmaOverflowAttr::Satfinite")); - assert!(dialect.contains("SparseMmaOverflowAttr::NotApplicable")); - assert!(dialect.contains("SparseMmaMetadataAttr { Standard, Ordered }")); - - let importer = render_importer(&catalog, "test-hash"); - assert!(importer.contains("sparse MMA selector must be the compile-time constant 0 or 1")); - assert!(importer.contains("sparse MMA selector must be the compile-time constant 0")); - assert!(importer.contains("GeneratedMmaImportAdapter::C4I32A2U32B2U32ToD4I32")); - assert!(importer.contains("GeneratedMmaImportAdapter::C4I32A4U32B4U32ToD4I32")); - assert!(importer.contains("GeneratedMmaImportAdapter::C4F32A4U32B4U32ToD4F32")); - assert!(importer.contains("GeneratedMmaImportAdapter::C2U32A4U32B4U32ToD2U32")); - assert!(importer.contains("(u32_ty, 2, u32_ty, 4, true, u32_ty, 4, true, u32_ty, 2)")); - assert!(importer.contains("let (c_array, last_op) = rvalue::translate_operand(")); - assert!(importer.contains("ctx, body, &args[0], value_map")); - assert!(importer.contains("let (a_value, last_after_a) = rvalue::translate_operand(")); - assert!(importer.contains("ctx, body, &args[1], value_map")); - assert!(importer.contains("let (b_value, last_after_b) = rvalue::translate_operand(")); - assert!(importer.contains("ctx, body, &args[2], value_map")); - assert!(importer.contains("operands.push(metadata)")); - assert!(importer.contains("operands.push(selector_value)")); - assert!(importer.contains("SparseMmaOp::get_concrete_op_info")); - - let lowering = render_lowering(&catalog, "test-hash"); - assert_eq!( - lowering - .matches("impl MirToLlvmConversion for SparseMmaOp") - .count(), - 1 - ); - assert!(lowering.contains("convert_generated_sparse_mma")); - for record in &records { - assert!(lowering.contains(&sparse_mma_template(record))); - assert!(lowering.contains(&sparse_mma_constraints(record))); - } - - let targets = render_targets(&catalog, "test-hash"); - assert!(targets.contains("GeneratedIntrinsicVariant::SparseMma")); - assert!(targets.contains("Operation::get_op::")); - assert!(targets.contains("get_attr_nvvm_sparse_mma_selector")); - assert!(targets.contains("GeneratedSparseMmaSelector::ImmediateZeroOrOne")); - assert!(targets.contains("GeneratedSparseMmaSelector::ImmediateZero")); - assert!(targets.contains("GeneratedSparseMmaShape::M16n8k64")); - assert!(targets.contains("SparseMmaShapeAttr::M16n8k64")); - assert!(targets.contains("GeneratedSparseMmaShape::M16n8k128")); - assert!(targets.contains("SparseMmaShapeAttr::M16n8k128")); - assert!(targets.contains("GeneratedSparseMmaElement::S4")); - assert!(targets.contains("GeneratedSparseMmaElement::U4")); - assert!(targets.contains("GeneratedSparseMmaElement::E2m1")); - assert!(targets.contains("GeneratedSparseMmaElement::E5m2")); - assert!(targets.contains("SparseMmaElementAttr::S4")); - assert!(targets.contains("SparseMmaElementAttr::U4")); - assert!(targets.contains("GeneratedSparseMmaMetadata::Standard")); - assert!(targets.contains("SparseMmaMetadataAttr::Standard")); - assert!(targets.contains("GeneratedSparseMmaMetadata::Ordered")); - assert!(targets.contains("SparseMmaMetadataAttr::Ordered")); - assert!(targets.contains("GeneratedSparseMmaAccumulator::F16")); - assert!(targets.contains("GeneratedSparseMmaAccumulator::F32")); - assert!(targets.contains("GeneratedSparseMmaOverflow::NotApplicable")); - assert!(targets.contains("GeneratedHardwareAlternative::ExactArchitecture(120)")); - - assert_eq!(raw.matches(SPARSE_MMA_STANDARD_METADATA_RULE).count(), 32); - assert_eq!(raw.matches(SPARSE_MMA_ORDERED_METADATA_RULE).count(), 82); - assert_eq!( - compatibility - .matches(SPARSE_MMA_STANDARD_METADATA_RULE) - .count(), - 32 - ); - assert_eq!( - compatibility - .matches(SPARSE_MMA_ORDERED_METADATA_RULE) - .count(), - 82 - ); - assert!( - lowering - .contains("mma.sp::ordered_metadata.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32") - ); - assert!( - lowering - .contains("mma.sp::ordered_metadata.sync.aligned.m16n8k64.row.col.s32.s8.s8.s32") - ); - assert!( - lowering - .contains("mma.sp::ordered_metadata.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32") - ); - assert!( - lowering - .contains("mma.sp::ordered_metadata.sync.aligned.m16n8k128.row.col.s32.s4.s4.s32") - ); - assert!(lowering.contains("mma.sp.sync.aligned.m16n8k64.row.col.s32.s8.s8.s32")); - assert!(lowering.contains("mma.sp.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32")); - assert!(lowering.contains("mma.sp.sync.aligned.m16n8k128.row.col.s32.s4.s4.s32")); - assert!(lowering.contains( - "mma.sp::ordered_metadata.sync.aligned.m16n8k64.row.col.kind::f8f6f4.f32.e2m1.e2m1.f32" - )); - assert!(lowering.contains( - r#"(GeneratedMmaResultType::F32, 4, 14, "mma.sp::ordered_metadata.sync.aligned.m16n8k64.row.col.kind::f8f6f4.f32.e2m1.e2m1.f32 {$0, $1, $2, $3}, {$8, $9, $10, $11}, {$12, $13, $14, $15}, {$4, $5, $6, $7}, $16, $17;", "=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r,r,r,r,n")"# - )); - assert!(lowering.contains( - r#"(GeneratedMmaResultType::I32, 2, 12, "mma.sp::ordered_metadata.sync.aligned.m16n8k64.row.col.kind::f8f6f4.f16.e2m1.e2m1.f16 {$0, $1}, {$4, $5, $6, $7}, {$8, $9, $10, $11}, {$2, $3}, $12, $13;", "=r,=r,r,r,r,r,r,r,r,r,r,r,r,n")"# - )); - let f16_probe = render_probe(&catalog, ordered_f8f6f4_f16, "test-hash"); - assert!(f16_probe.contains( - "define { i32, i32 } @probe_mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e2m1_f16_selector_0" - )); - assert!(!f16_probe.contains("_selector_1")); - - for record in &records { - let probe = render_probe(&catalog, record, "test-hash"); - assert!(probe.contains(&format!("probe_{}_selector_0", record.id))); - let selectors = sparse_mma_selector_values(record); - assert_eq!( - probe.contains(&format!("probe_{}_selector_1", record.id)), - selectors.contains(&1) - ); - assert_eq!(probe.matches("asm sideeffect").count(), selectors.len()); - assert!(probe.contains(&sparse_mma_template(record))); - assert!(probe.contains(&sparse_mma_constraints(record))); - assert!(probe.contains("attributes #0 = { convergent }")); - } - - let reference = render_reference(&catalog, "test-hash"); - assert!(reference.contains("## Sparse-MMA contracts")); - assert!(reference.contains("C, A, B, metadata, selector order")); - assert!(reference.contains("LLVM source record uses A, B, C, metadata, selector order")); - assert_eq!( - reference.matches(SPARSE_MMA_STANDARD_METADATA_RULE).count(), - 32 - ); - assert_eq!( - reference.matches(SPARSE_MMA_ORDERED_METADATA_RULE).count(), - 82 - ); - let sparse_reference = reference - .split("## Sparse-MMA contracts") - .nth(1) - .unwrap() - .split("## Packed-atomic contracts") - .next() - .unwrap(); - assert!(sparse_reference.contains("Overflow mode is not applicable.")); - assert!(!sparse_reference.contains("Integer overflow is not applicable")); - for record in &records { - assert!(reference.contains(&format!("- `{}`: runtime `unexecuted`", record.id))); - } - - let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); - assert!(outputs.contains_key(&PathBuf::from( - "crates/cuda-device/src/generated/sparse_mma.rs" - ))); - assert!(outputs.contains_key(&PathBuf::from( - "crates/dialect-nvvm/src/ops/generated/sparse_mma.rs" - ))); - } - - #[test] - fn prmt_rendering_keeps_modes_and_zero_source_exact() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - validate_renderable(&catalog).unwrap(); - assert_eq!(prmts(&catalog).count(), 7); - - let dialect = render_dialect_prmt(&catalog, "test-hash"); - assert_eq!(dialect.matches("pub struct PrmtOp").count(), 1); - for mode in ["Generic", "F4e", "B4e", "Rc8", "Ecl", "Ecr", "Rc16"] { - assert!(dialect.contains(mode)); - } - - let lowering = render_lowering(&catalog, "test-hash"); - for mode in ["rc8", "ecl", "ecr", "rc16"] { - assert!(lowering.contains(&format!("prmt.b32.{mode} $0, $1, 0, $2;"))); - } - assert!(lowering.contains("prmt.b32 $0, $1, $2, $3;")); - assert!(lowering.contains("prmt.b32.f4e $0, $1, $2, $3;")); - assert!(lowering.contains("prmt.b32.b4e $0, $1, $2, $3;")); - assert!(lowering.contains("convert_generated_prmt")); - - let targets = render_targets(&catalog, "test-hash"); - assert!(targets.contains("GeneratedIntrinsicVariant::Prmt")); - assert!(targets.contains("Operation::get_op::")); - assert!(targets.contains("GeneratedPrmtMode::Rc16")); - - let compatibility = render_compat_prmt(&catalog, "test-hash"); - assert_eq!(compatibility.matches("pub fn prmt").count(), 7); - } - - #[test] - fn selected_special_register_probes_use_the_llvm_route() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - - let clock = catalog - .intrinsics - .iter() - .find(|record| record.id == "clock") - .unwrap(); - let typed = render_probe(&catalog, clock, "test-hash"); - assert!(typed.contains("declare i32 @llvm.nvvm.read.ptx.sreg.clock()")); - assert!(typed.contains("call i32 @llvm.nvvm.read.ptx.sreg.clock()")); - - let gridid = catalog - .intrinsics - .iter() - .find(|record| record.id == "gridid") - .unwrap(); - let inline = render_probe(&catalog, gridid, "test-hash"); - assert!(inline.contains("define i64 @probe_gridid_llvm_nvptx()")); - assert!(inline.contains("call i64 asm \"mov.u64 $0, %gridid;\", \"=l\"()")); - } - - #[test] - fn cluster_barrier_rendering_owns_compatibility_sync_and_both_backends() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - assert_eq!(cluster_barriers(&catalog).count(), 6); - - let dialect = render_dialect_cluster_barrier(&catalog, "test-hash"); - assert_eq!(dialect.matches("pub struct ClusterBarrierOp").count(), 1); - assert_eq!(dialect.matches("pub struct ClusterSyncOp").count(), 1); - assert!(dialect.contains("name = \"nvvm.cluster_sync\"")); - assert!(dialect.contains("ClusterSyncOp::register(ctx);")); - - let importer = render_importer(&catalog, "test-hash"); - assert!(importer.contains("\"cuda_device::cluster::cluster_sync\" =>")); - assert!(importer.contains("ClusterBarrierModeAttr::ArriveAligned")); - assert!(importer.contains("ClusterBarrierModeAttr::WaitAligned")); - assert!(!importer.contains("ClusterSyncOp")); - - let lowering = render_lowering(&catalog, "test-hash"); - assert_eq!( - lowering - .matches("impl MirToLlvmConversion for ClusterSyncOp") - .count(), - 1 - ); - assert!(lowering.contains("llvm_nvvm_barrier_cluster_arrive_aligned")); - assert!(lowering.contains("llvm_nvvm_barrier_cluster_wait_aligned")); - assert!(lowering.contains("barrier.cluster.arrive.aligned; barrier.cluster.wait.aligned;")); - - let reference = render_reference(&catalog, "test-hash"); - assert!(reference.contains("## Cluster-barrier contracts")); - assert!(reference.contains( - "`cuda_device::cluster::cluster_sync` is the generated compatibility operation" - )); - - let targets = render_targets(&catalog, "test-hash"); - assert!(targets.contains( - "GeneratedIntrinsicVariant::ClusterBarrier { mode: GeneratedClusterBarrierMode::ArriveRelaxedAligned }" - )); - assert!(targets.contains("Operation::get_op::")); - assert!(targets.contains("ClusterBarrierModeAttr::WaitAligned")); - } - - #[test] - fn cluster_memory_rendering_preserves_cluster_shared_addrspace_and_composite_load() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - validate_renderable(&catalog).unwrap(); - let records = cluster_memory(&catalog).collect::>(); - assert_eq!(records.len(), 2); - - let compatibility = render_compat_cluster_memory(&catalog, "test-hash"); - for signature in [ - "pub unsafe fn map_shared_rank(local_ptr: *const T, target_rank: u32) -> *const T", - "pub unsafe fn map_shared_rank_mut(local_ptr: *mut T, target_rank: u32) -> *mut T", - "pub unsafe fn dsmem_read_u32(local_ptr: *const u32, target_rank: u32) -> u32", - ] { - assert!(compatibility.contains(signature), "missing {signature}"); - } - assert!(compatibility.contains("cluster-shared pointer in address space 7")); - - let dialect_mod = render_dialect_mod(&catalog, "test-hash"); - assert!(dialect_mod.contains("mod cluster_memory;")); - assert!(dialect_mod.contains("pub use cluster_memory::*;")); - assert!(dialect_mod.contains("cluster_memory::register(ctx);")); - - let dialect = render_dialect_cluster_memory(&catalog, "test-hash"); - assert!(dialect.contains("pub struct MapaSharedClusterOp")); - assert!(dialect.contains("pub struct DsmemReadU32Op")); - assert!(dialect.contains("source_ptr.pointer_kind()")); - assert!(dialect.contains("address_space::CLUSTER_SHARED")); - assert!(dialect.contains("MirPointerKind::RawConst | MirPointerKind::RawMut")); - assert!(dialect.contains("result_ptr.pointer_kind() != source_ptr.pointer_kind()")); - assert!(dialect.contains( - "nvvm.mapa_shared_cluster requires a raw source pointer and must preserve its pointee, mutability, and raw kind while returning addrspace(7)" - )); - assert!(dialect.contains("nvvm.dsmem_read_u32 result must be u32")); - assert!(dialect.contains("MapaSharedClusterOp::register(ctx);")); - assert!(dialect.contains("DsmemReadU32Op::register(ctx);")); - - let importer = render_importer(&catalog, "test-hash"); - for path in [ - "cuda_device::cluster::map_shared_rank", - "cuda_device::cluster::map_shared_rank_mut", - "cuda_device::cluster::dsmem_read_u32", - ] { - assert!(importer.contains(path)); - } - assert!(importer.contains("MapaSharedClusterOp::build(ctx, source, rank)")); - assert!(importer.contains("DsmemReadU32Op::build(ctx, source, rank)")); - - let lowering = render_lowering(&catalog, "test-hash"); - assert!(lowering.contains("cast_to_shared_addrspace")); - assert!(lowering.contains("llvm_ops::IntToPtrOp::new")); - assert!(lowering.contains("llvm_types::PointerType::get(ctx, 7)")); - assert!(lowering.contains("mapa.shared::cluster.u64 $0, $1, $2;")); - assert!(lowering.contains( - "{ .reg .u64 %mapped; mapa.shared::cluster.u64 %mapped, $1, $2; ld.shared::cluster.u32 $0, [%mapped]; }" - )); - assert!(lowering.contains("\"=l,l,r\"")); - assert!(lowering.contains("\"=r,l,r,~{memory}\"")); - - let map = records - .iter() - .find(|record| record.id == "map_shared_rank") - .unwrap(); - assert_eq!(map.llvm.as_ref().unwrap().results, ["shared_cluster_ptr"]); - let map_probe = render_probe(&catalog, map, "test-hash"); - assert!(map_probe.contains("define ptr addrspace(7)")); - assert!(map_probe.contains("inttoptr i64")); - assert!(map_probe.contains("asm sideeffect")); - assert!(!map_probe.contains("~{memory}")); - assert!(map_probe.contains("attributes #0 = { convergent }")); - - let read = records - .iter() - .find(|record| record.id == "dsmem_read_u32") - .unwrap(); - assert!(read.llvm.is_none()); - let read_probe = render_probe(&catalog, read, "test-hash"); - assert!(read_probe.contains("ld.shared::cluster.u32")); - assert!(read_probe.contains("~{memory}")); - assert!(read_probe.contains("attributes #0 = { convergent }")); - - let raw = render_raw_abi(&catalog, "test-hash").unwrap(); - assert!(raw.contains("pub unsafe fn i0320(_arg0: *const u8, _arg1: u32) -> *const u8")); - assert!(raw.contains("pub unsafe fn i0321(_arg0: *const u32, _arg1: u32) -> u32")); - assert!(raw.contains( - "ordinary loads and stores compile to `ld.shared::cluster` and `st.shared::cluster`" - )); - - let reference = render_reference(&catalog, "test-hash"); - assert!(reference.contains("## Cluster-memory contracts")); - assert!(reference.contains("preserves LLVM 22's address-space-7 result")); - assert!(reference.contains("has no one-to-one LLVM intrinsic")); - - let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); - assert!(outputs.contains_key(&PathBuf::from( - "crates/cuda-device/src/generated/cluster_memory.rs" - ))); - assert!(outputs.contains_key(&PathBuf::from( - "crates/dialect-nvvm/src/ops/generated/cluster_memory.rs" - ))); - - let mut wrong_as = catalog.clone(); - wrong_as - .intrinsics - .iter_mut() - .find(|record| record.id == "map_shared_rank") - .unwrap() - .llvm - .as_mut() - .unwrap() - .results = vec!["shared_ptr".into()]; - assert!(validate_renderable(&wrong_as).is_err()); - - let mut wrong_source = catalog; - wrong_source - .intrinsics - .iter_mut() - .find(|record| record.id == "dsmem_read_u32") - .unwrap() - .source = IntrinsicSource::LlvmImported { - source_record: "invented".into(), - }; - assert!(validate_renderable(&wrong_source).is_err()); - } - - #[test] - fn lowering_imports_cast_to_shared_addrspace_exactly_once_across_families() { - // The stmatrix, cluster_memory, and mbarrier_extended families all - // insert the `cast_to_shared_addrspace` helper into the same - // `common::{call_intrinsic, ...}` use group. An unguarded second - // insert duplicates the name inside one use group, which is a hard - // rustc error (E0252) in the raw generator output. The rustfmt pass - // run over generated files silently dedupes the duplicate, so assert - // on the raw (pre-rustfmt) output to catch a regressed guard loudly. - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - assert!(stmatrices(&catalog).next().is_some()); - assert!(cluster_memory(&catalog).next().is_some()); - assert!(mbarrier_extended(&catalog).next().is_some()); - - let import_count = |lowering: &str| { - // Everything before the first rendered helper function is the - // module header plus the use block; helper call sites in the - // function bodies below must not count. - let bodies_start = lowering - .find("fn convert_zero_operand_scalar_direct") - .expect("lowering output must contain the first helper function"); - lowering[..bodies_start] - .matches("cast_to_shared_addrspace") - .count() - }; - - assert_eq!( - import_count(&render_lowering(&catalog, "test-hash")), - 1, - "cast_to_shared_addrspace must be imported exactly once in the \ - raw lowering output when multiple families need it" - ); - - // The guard must not suppress the import either: with stmatrix - // absent, cluster_memory still needs the helper and must insert it. - let mut without_stmatrix = catalog; - without_stmatrix - .intrinsics - .retain(|record| record.family != "stmatrix"); - assert!(stmatrices(&without_stmatrix).next().is_none()); - assert_eq!( - import_count(&render_lowering(&without_stmatrix, "test-hash")), - 1, - "cast_to_shared_addrspace must still be imported when only \ - cluster_memory and mbarrier_extended need it" - ); - } - - #[test] - fn clc_rendering_preserves_api_and_uses_typed_llvm_routes() { - let catalog = catalog_with_clc(); - validate_renderable(&catalog).unwrap(); - - let compatibility = render_compat_clc(&catalog, "test-hash"); - assert!( - compatibility - .contains("pub unsafe fn clc_try_cancel(response: *mut u8, mbar: *mut Barrier)") - ); - assert!(compatibility.contains( - "pub unsafe fn clc_try_cancel_multicast(response: *mut u8, mbar: *mut Barrier)" - )); - assert!( - compatibility - .contains("pub unsafe fn clc_query_is_canceled(resp_lo: u64, resp_hi: u64) -> u32") - ); - - let dialect = render_dialect_clc(&catalog, "test-hash"); - for op in [ - "ClcTryCancelOp", - "ClcTryCancelMulticastOp", - "ClcQueryIsCanceledOp", - "ClcQueryGetFirstCtaidXOp", - "ClcQueryGetFirstCtaidYOp", - "ClcQueryGetFirstCtaidZOp", - ] { - assert_eq!(dialect.matches(&format!("pub struct {op}")).count(), 1); - assert!(dialect.contains(&format!("{op}::register(ctx)"))); - } - assert_eq!(dialect.matches("NResultsInterface<0>").count(), 2); - assert_eq!(dialect.matches("NResultsInterface<1>").count(), 4); - assert_eq!(dialect.matches("impl Verify for ClcQuery").count(), 4); - assert!(dialect.contains("is_integer_width(ctx, op.get_operand(0).get_type(ctx), 64)")); - assert!(dialect.contains("is_integer_width(ctx, op.get_result(0).get_type(ctx), 32)")); - - let importer = render_importer(&catalog, "test-hash"); - assert!(importer.contains("cuda_device::clc::clc_try_cancel")); - assert!(importer.contains("ClcTryCancelOp::get_concrete_op_info()")); - assert!(importer.contains("ClcQueryIsCanceledOp::get_concrete_op_info()")); - assert!(importer.contains("helpers::emit_store_result_and_goto")); - - let lowering = render_lowering(&catalog, "test-hash"); - assert!( - lowering - .contains("clc::{convert_generated_clc_query, convert_generated_clc_try_cancel}") - ); - assert!(lowering.contains( - "convert_generated_clc_try_cancel(ctx, rewriter, self.get_operation(), operands_info, \"llvm__nvvm_dclusterlaunchcontrol_dtry_ucancel_dasync_dshared\")" - )); - assert!(lowering.contains( - "convert_generated_clc_query(ctx, rewriter, self.get_operation(), operands_info, \"llvm__nvvm_dclusterlaunchcontrol_dquery_ucancel_dis_ucanceled\", true)" - )); - - let request = clc_intrinsics(&catalog) - .find(|record| record.id == "clc_try_cancel") - .unwrap(); - let request_probe = render_probe(&catalog, request, "test-hash"); - assert_eq!(request_probe.matches("addrspacecast ptr").count(), 2); - assert!(request_probe.contains( - "call void @llvm.nvvm.clusterlaunchcontrol.try_cancel.async.shared(ptr addrspace(3) %response, ptr addrspace(3) %mbarrier)" - )); - - let query = clc_intrinsics(&catalog) - .find(|record| record.id == "clc_query_is_canceled") - .unwrap(); - let query_probe = render_probe(&catalog, query, "test-hash"); - assert!(query_probe.contains("%response_high_shifted = shl i128 %response_high_i128, 64")); - assert!( - query_probe.contains("%response = or i128 %response_low_i128, %response_high_shifted") - ); - assert!(query_probe.contains("%result = zext i1 %raw_result to i32")); - - let targets = render_targets(&catalog, "test-hash"); - assert!(targets.contains( - "GeneratedHardwareTarget::AnyOf(&[GeneratedHardwareAlternative::ExactArchitecture(100), GeneratedHardwareAlternative::ExactArchitecture(101), GeneratedHardwareAlternative::ExactArchitecture(103), GeneratedHardwareAlternative::ExactArchitecture(110), GeneratedHardwareAlternative::ExactArchitecture(120), GeneratedHardwareAlternative::ExactArchitecture(121)])" - )); - - let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); - assert!(outputs.contains_key(&PathBuf::from("crates/cuda-device/src/generated/clc.rs"))); - assert!(outputs.contains_key(&PathBuf::from( - "crates/dialect-nvvm/src/ops/generated/clc.rs" - ))); - - let mut wrong_adapter = catalog; - wrong_adapter - .intrinsics - .iter_mut() - .find(|record| record.id == "clc_query_is_canceled") - .unwrap() - .clc - .as_mut() - .unwrap() - .adapter = ClcAdapter::PairU64ToI128U32; - assert!(validate_renderable(&wrong_adapter).is_err()); - } - - #[test] - fn elect_dialect_verifies_its_scalar_shape() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - let dialect = render_dialect_elect(&catalog, "test-hash"); - - assert!(dialect.contains("impl Verify for ElectSyncOp")); - assert!(dialect.contains("is_integer_width(ctx, op.get_operand(0).get_type(ctx), 32)")); - assert!(dialect.contains("is_integer_width(ctx, op.get_result(0).get_type(ctx), 32)")); - assert!(dialect.contains("is_integer_width(ctx, op.get_result(1).get_type(ctx), 1)")); - assert!(!dialect.contains("verifier = \"succ\"")); - } - - #[test] - fn tma_rendering_preserves_api_and_injects_backend_defaults() { - let catalog = catalog_with_tma(); - validate_renderable(&catalog).unwrap(); - - let compatibility = render_compat_tma(&catalog, "test-hash"); - assert!(compatibility.contains( - "pub unsafe fn cp_async_bulk_tensor_1d_g2s(dst: *mut u8, tensor_map: *const TmaDescriptor, coord0: i32, barrier: *mut Barrier)" - )); - assert!(compatibility.contains( - "pub unsafe fn cp_async_bulk_tensor_2d_g2s_multicast(dst: *mut u8, tensor_map: *const TmaDescriptor, coord0: i32, coord1: i32, barrier: *mut Barrier, cta_mask: u16)" - )); - assert!(compatibility.contains( - "pub unsafe fn cp_async_bulk_tensor_5d_s2g(src: *const u8, tensor_map: *const TmaDescriptor, coord0: i32, coord1: i32, coord2: i32, coord3: i32, coord4: i32)" - )); - assert!(compatibility.contains("pub fn cp_async_bulk_commit_group()")); - assert!(compatibility.contains("pub fn cp_async_bulk_wait_group(n: u32)")); - assert!(compatibility.contains( - "pub unsafe fn cp_async_bulk_tensor_reduce_add_tile_2d(src: *const u8, tensor_map: *const TmaDescriptor, coord0: i32, coord1: i32)" - )); - assert!(compatibility.contains( - "pub unsafe fn cp_async_bulk_tensor_reduce_xor_im2col_5d(src: *const u8, tensor_map: *const TmaDescriptor, coord0: i32, coord1: i32, coord2: i32, coord3: i32, coord4: i32)" - )); - for dimensions in 1..=5 { - assert!(compatibility.contains(&format!( - "pub unsafe fn cp_async_bulk_prefetch_tensor_{dimensions}d_l2(" - ))); - assert!(compatibility.contains(&format!( - "pub unsafe fn cp_async_bulk_prefetch_tensor_{dimensions}d_l2_cache_hint(" - ))); - } - assert!( - compatibility.contains("pub unsafe fn cp_async_bulk_prefetch_tensor_gather4_2d_l2(") - ); - assert!( - compatibility - .contains("pub unsafe fn cp_async_bulk_prefetch_tensor_gather4_2d_l2_cache_hint(") - ); - assert!(compatibility.contains("pub unsafe fn tensormap_replace_swizzle_atomicity(")); - assert!( - compatibility.contains("pub unsafe fn fence_proxy_tensormap_generic_acquire_system(") - ); - assert!(compatibility.contains("pub fn fence_proxy_tensormap_generic_release_system()")); - - let dialect = render_dialect_tma(&catalog, "test-hash"); - assert_eq!(dialect.matches("pub struct ").count(), 111); - assert_eq!(dialect.matches("NResultsInterface<0>").count(), 111); - assert!(dialect.contains("NOpdsInterface<10>")); - assert!(dialect.contains("CpAsyncBulkWaitGroupReadOp::register(ctx)")); - assert!( - dialect.contains("CpAsyncBulkPrefetchTensorGather4TwoDimensionalL2Op::register(ctx)") - ); - assert!(dialect.contains( - "CpAsyncBulkPrefetchTensorGather4TwoDimensionalL2CacheHintOp::register(ctx)" - )); - assert!(dialect.contains("ReplaceTensorMapSwizzleAtomicityOp::register(ctx)")); - assert!(dialect.contains("FenceProxyTensorMapGenericReleaseSystemOp::register(ctx)")); - - let dialect_mod = render_dialect_mod(&catalog, "test-hash"); - assert!(dialect_mod.contains("mod tma;")); - assert!(dialect_mod.contains("pub use tma::*;")); - assert!(dialect_mod.contains("tma::register(ctx);")); - - let importer = render_importer(&catalog, "test-hash"); - assert!(importer.contains("cuda_device::tma::cp_async_bulk_tensor_1d_g2s")); - assert!(importer.contains("super::tma::emit_tma_g2s(")); - assert!(importer.contains("super::tma::emit_tma_g2s_multicast_cg2(")); - assert!(importer.contains("super::tma::emit_tma_s2g(")); - assert!(importer.contains("TMA wait-group count must be a compile-time constant")); - assert!(importer.contains("require_arity(name, args.len(), 1, &loc)?;")); - assert!(importer.contains("args.first(), Some(mir::Operand::Constant(_))")); - assert!(importer.contains("\"v1:i0328\"")); - assert!(importer.contains("cp_async_bulk_tensor_reduce_add_tile_2d")); - assert!(importer.contains( - "dialect_nvvm::ops::CpAsyncBulkTensorReduceAddTile2dOp::get_concrete_op_info()" - )); - - let lowering = render_lowering(&catalog, "test-hash"); - assert!(lowering.contains("CpAsyncBulkTensorG2sTile1dOp")); - assert!(lowering.contains( - "tma::{convert_control, convert_g2s, convert_g2s_multicast_cg2, convert_prefetch_tensormap, convert_prefetch_tile, convert_reduce_s2g, convert_s2g, convert_tensormap_fence, convert_tensormap_replace, PrefetchTileConfig, ReduceConfig}" - )); - assert!( - lowering.contains( - "convert_g2s(ctx, rewriter, self.get_operation(), operands_info, 5, false)" - ) - ); - assert!( - lowering.contains("convert_s2g(ctx, rewriter, self.get_operation(), operands_info, 5)") - ); - assert!(lowering.contains( - "convert_control(ctx, rewriter, self.get_operation(), operands_info, \"commit_group\", \"llvm_nvvm_cp_async_bulk_commit_group\")" - )); - assert!(lowering.contains( - "convert_reduce_s2g(ctx, rewriter, self.get_operation(), operands_info, ReduceConfig::new(2, \"add\", \"tile\", \"llvm_nvvm_cp_async_bulk_tensor_reduce_add_tile_2d\"))" - )); - assert!(lowering.contains( - "convert_prefetch_tile(ctx, rewriter, self.get_operation(), operands_info, PrefetchTileConfig::new(5, false, false, \"llvm_nvvm_cp_async_bulk_tensor_prefetch_tile_5d\"))" - )); - assert!(lowering.contains( - "convert_tensormap_replace(ctx, rewriter, self.get_operation(), operands_info, \"llvm_nvvm_tensormap_replace_global_address\", \"global_address\", \"address\", false, false)" - )); - let release_system = tma_intrinsics(&catalog) - .find(|record| record.id == "fence_proxy_tensormap_generic_release_system") - .unwrap(); - assert!(lowering.contains(&format!( - "convert_tensormap_fence(ctx, rewriter, self.get_operation(), operands_info, {:?}, false, \"sys\")", - release_system.resolved_llvm_identifier() - ))); - - let g2s = tma_intrinsics(&catalog) - .find(|record| record.id == "cp_async_bulk_tensor_2d_g2s_multicast_cg2") - .unwrap(); - let g2s_probe = render_probe(&catalog, g2s, "test-hash"); - assert!(g2s_probe.contains("i16 %cta_mask, i64 0, i1 true, i1 false, i32 2")); - assert!(g2s_probe.contains("addrspacecast ptr %dst_generic to ptr addrspace(7)")); - - let s2g = tma_intrinsics(&catalog) - .find(|record| record.id == "cp_async_bulk_tensor_2d_s2g") - .unwrap(); - let s2g_probe = render_probe(&catalog, s2g, "test-hash"); - assert!(s2g_probe.contains("ptr %tensor_map, i32 %coord0, i32 %coord1, i64 0, i1 false")); - - let reduce = tma_intrinsics(&catalog) - .find(|record| record.id == "cp_async_bulk_tensor_reduce_add_tile_2d") - .unwrap(); - let reduce_probe = render_probe(&catalog, reduce, "test-hash"); - assert!(reduce_probe.contains( - "declare void @llvm.nvvm.cp.async.bulk.tensor.reduce.add.tile.2d(ptr addrspace(3), ptr, i32, i32, i64, i1)" - )); - assert!(reduce_probe.contains( - "call void @llvm.nvvm.cp.async.bulk.tensor.reduce.add.tile.2d(ptr addrspace(3) %src, ptr %tensor_map, i32 %coord0, i32 %coord1, i64 0, i1 false)" - )); - - for stem in [ - "cp_async_bulk_prefetch_tensor_1d_l2", - "cp_async_bulk_prefetch_tensor_2d_l2", - "cp_async_bulk_prefetch_tensor_3d_l2", - "cp_async_bulk_prefetch_tensor_4d_l2", - "cp_async_bulk_prefetch_tensor_5d_l2", - "cp_async_bulk_prefetch_tensor_gather4_2d_l2", - ] { - let prefetch_plain = tma_intrinsics(&catalog) - .find(|record| record.id == stem) - .unwrap(); - let prefetch_plain_probe = render_probe(&catalog, prefetch_plain, "test-hash"); - assert!(prefetch_plain_probe.contains("i64 0, i1 false")); - assert!(!prefetch_plain_probe.contains("%cache_hint")); - - let cache_hint_id = format!("{stem}_cache_hint"); - let prefetch_cache_hint = tma_intrinsics(&catalog) - .find(|record| record.id == cache_hint_id) - .unwrap(); - let prefetch_cache_hint_probe = - render_probe(&catalog, prefetch_cache_hint, "test-hash"); - assert!(prefetch_cache_hint_probe.contains("i64 %cache_hint")); - assert!(prefetch_cache_hint_probe.contains("i64 %cache_hint, i1 true")); - } - - let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); - assert!(outputs.contains_key(&PathBuf::from("crates/cuda-device/src/generated/tma.rs"))); - assert!(outputs.contains_key(&PathBuf::from( - "crates/dialect-nvvm/src/ops/generated/tma.rs" - ))); - - let mut wrong_adapter = catalog; - wrong_adapter - .intrinsics - .iter_mut() - .find(|record| record.id == "cp_async_bulk_tensor_1d_g2s") - .unwrap() - .tma - .as_mut() - .unwrap() - .adapter = TmaAdapter::NoOperands; - assert!(validate_renderable(&wrong_adapter).is_err()); - } - - #[test] - fn tcgen05_rendering_preserves_api_and_inline_ptx_routes() { - let catalog = catalog_with_tcgen05(); - validate_renderable(&catalog).unwrap(); - - let compatibility = render_compat_tcgen05(&catalog, "test-hash"); - assert!( - compatibility - .contains("pub unsafe fn tcgen05_alloc(dst_smem: *mut u32, n_cols: u32) -> ()") - ); - assert!(compatibility.contains( - "pub unsafe fn tcgen05_mma_ws_f16(d_tmem: u32, a_tmem: u32, a_desc: u64, b_desc: u64, idesc: u32, enable_d: bool) -> ()" - )); - for alias in ["e4m3", "e5m2", "e2m3", "e3m2", "e2m1"] { - assert!(compatibility.contains(&format!( - "pub unsafe fn tcgen05_mma_{alias}(d_tmem: u32, a_desc: u64, b_desc: u64, idesc: u32, enable_d: bool)" - ))); - } - assert!( - compatibility - .contains("pub unsafe fn tcgen05_ld_16x256b_x8_pure(tmem_addr: u32) -> TmemF32x32") - ); - assert!(compatibility.contains("pub fn tcgen05_load_wait() -> ()")); - assert!(compatibility.contains("pub fn tcgen05_relinquish_alloc_permit() -> ()")); - assert!(compatibility.contains("pub fn tcgen05_relinquish_alloc_permit_cg2() -> ()")); - assert!(compatibility.contains( - "One full warp must execute this instruction uniformly with the same operands." - )); - assert!(compatibility.contains( - "One full warp in each peer CTA must execute this instruction; lanes within each warp must execute uniformly with the same operands." - )); - assert!(compatibility.contains( - "pub unsafe fn tcgen05_commit_multicast_cg2(mbar: *mut u64, cta_mask: u16) -> ()" - )); - assert!(compatibility.contains( - "pub unsafe fn tcgen05_cp_128x128b_b4x16_p64(tmem_addr: u32, smem_desc: u64) -> ()" - )); - assert!(compatibility.contains( - "pub unsafe fn tcgen05_cp_64x128b_warpx2_02_13_cg2(tmem_addr: u32, smem_desc: u64) -> ()" - )); - assert!( - compatibility.contains("pub unsafe fn tcgen05_ld_16x64b_x1_raw(tmem_addr: u32) -> u32") - ); - assert!( - compatibility.contains( - "pub unsafe fn tcgen05_ld_16x64b_x2_raw(tmem_addr: u32) -> CuSimd" - ) - ); - assert!(compatibility.contains( - "pub unsafe fn tcgen05_ld_32x32b_x128_pack16(tmem_addr: u32) -> CuSimd" - )); - assert!( - compatibility.contains( - "pub unsafe fn tcgen05_st_16x64b_x1_raw(tmem_addr: u32, data: u32) -> ()" - ) - ); - assert!(compatibility.contains( - "pub unsafe fn tcgen05_st_16x64b_x2_raw(tmem_addr: u32, data: CuSimd) -> ()" - )); - assert!(compatibility.contains( - "pub unsafe fn tcgen05_st_32x32b_x128_unpack16(tmem_addr: u32, data: CuSimd) -> ()" - )); - assert!(compatibility.contains( - "pub unsafe fn tcgen05_ld_16x32bx2_x1_raw(tmem_addr: u32) -> u32" - )); - assert!(compatibility.contains( - "pub(crate) unsafe fn __tcgen05_ld_16x32bx2_x1_raw(_tmem_addr: u32, _half_split_offset: i64) -> u32" - )); - assert!(compatibility.contains( - "unsafe { __tcgen05_ld_16x32bx2_x1_raw(tmem_addr, HALF_SPLIT_OFFSET as i64) }" - )); - assert!(compatibility.contains( - "pub unsafe fn tcgen05_st_16x32bx2_x128_unpack16(tmem_addr: u32, data: CuSimd)" - )); - assert!(compatibility.contains( - "pub(crate) unsafe fn __tcgen05_st_16x32bx2_x128_unpack16(_tmem_addr: u32, _half_split_offset: i64, _data: CuSimd)" - )); - assert!(compatibility.contains( - "pub unsafe fn tcgen05_commit_multicast(mbar: *mut u64, cta_mask: u16) -> ()" - )); - assert!(compatibility.contains("pub unsafe fn tcgen05_shift_down(tmem_addr: u32) -> ()")); - assert!( - compatibility.contains("pub unsafe fn tcgen05_shift_down_cg2(tmem_addr: u32) -> ()") - ); - assert_eq!( - compatibility - .matches("/// One thread in the CTA issues this instruction.\n") - .count(), - 4 - ); - assert_eq!( - compatibility - .matches( - "/// One thread in the CTA pair issues this instruction; the peer CTA must be active and must not have exited.\n", - ) - .count(), - 4 - ); - assert_eq!( - compatibility - .matches( - "/// The same thread that issued the tracked asynchronous tcgen05 operations must issue this commit.\n", - ) - .count(), - 6 - ); - assert_eq!( - compatibility - .matches( - "/// Completion must be tracked by a matching commit from that same thread and observed through the selected mbarrier before relying on shifted data.\n", - ) - .count(), - 2 - ); - assert_eq!( - compatibility - .matches( - "/// All tcgen05 operations in the kernel must use the same CTA-group mode.\n", - ) - .count(), - 233 - ); - assert!(compatibility.contains("/// `KIND` is 0=f16, 1=tf32, 2=f8f6f4, or 3=i8.")); - assert!(compatibility.contains( - "/// `CTA_GROUP` is 1 or 2. `COLLECTOR_A` is 0=discard, 1=lastuse, 2=fill, or 3=use." - )); - assert!(compatibility.contains( - "/// `B_BUFFER` is 0 through 3. `B_USAGE` is 0=discard, 1=lastuse, 2=fill, or 3=use." - )); - assert_eq!( - compatibility - .matches( - "/// `legacy_a_desc` is kept for compatibility; tensor A uses `a_tmem`.\n", - ) - .count(), - 5 - ); - assert_eq!( - compatibility - .matches("/// This uses kind f8f6f4, CTA group 1, and collector a::discard.\n",) - .count(), - 5 - ); - assert!(!compatibility.contains("tcgen05_mma_ws_f16_with_collector")); - - let raw = render_raw_abi(&catalog, "test-hash").unwrap(); - assert!(raw.contains("pub unsafe fn i0578(_arg0: u32, _arg1: u64) -> ()")); - assert!(raw.contains("pub unsafe fn i0611(_arg0: u32, _arg1: u64) -> ()")); - assert!(raw.contains("pub unsafe fn i0612(_arg0: u32) -> u32")); - assert!(raw.contains("pub unsafe fn i0614(_arg0: u32) -> [u32; 2]")); - assert!(raw.contains("pub unsafe fn i0669(_arg0: u32) -> [u32; 128]")); - assert!(raw.contains("pub unsafe fn i0670(_arg0: u32, _arg1: u32) -> ()")); - assert!(raw.contains("pub unsafe fn i0672(_arg0: u32, _arg1: [u32; 2]) -> ()")); - assert!(raw.contains("pub unsafe fn i0727(_arg0: u32, _arg1: [u32; 128]) -> ()")); - assert!(raw.contains("pub unsafe fn i0728(_arg0: u32, _arg1: i64) -> u32")); - assert!(raw.contains("pub unsafe fn i0743(_arg0: u32, _arg1: i64) -> [u32; 128]")); - assert!( - raw.contains("pub unsafe fn i0759(_arg0: u32, _arg1: i64, _arg2: [u32; 128]) -> ()") - ); - assert!(raw.contains("pub unsafe fn i0760(_arg0: *mut u64, _arg1: u16) -> ()")); - assert!(raw.contains("pub unsafe fn i0761(_arg0: u32) -> ()")); - assert!(raw.contains("pub unsafe fn i0762(_arg0: u32) -> ()")); - let mma = raw.find("pub unsafe fn i0763").unwrap(); - assert!(raw[..mma].ends_with("#[allow(clippy::too_many_arguments)]\n#[inline(never)]\n")); - assert!(raw.contains("lane component must be a multiple of 32")); - assert!(raw.contains("`_arg0` must name a live tensor-memory allocation")); - assert_eq!( - raw.matches("/// One thread in the CTA issues this instruction.\n") - .count(), - 4 - ); - assert_eq!( - raw.matches( - "/// One thread in the CTA pair issues this instruction; the peer CTA must be active and must not have exited.\n", - ) - .count(), - 4 - ); - assert_eq!( - raw.matches( - "/// The same thread that issued the tracked asynchronous tcgen05 operations must issue this commit.\n", - ) - .count(), - 6 - ); - assert_eq!( - raw.matches( - "/// Completion must be tracked by a matching commit from that same thread and observed through the selected mbarrier before relying on shifted data.\n", - ) - .count(), - 2 - ); - assert_eq!( - raw.matches( - "/// All tcgen05 operations in the kernel must use the same CTA-group mode.\n", - ) - .count(), - 233 - ); - assert!(raw.contains("pub fn i0345() -> ()")); - assert!(raw.contains("pub fn i0357() -> ()")); - assert!(raw.contains("pub fn i0358() -> ()")); - assert!(raw.contains("pub fn i0361() -> ()")); - assert!(raw.contains("Complete the matching tensor-memory load wait")); - assert!(raw.contains("Complete the matching tensor-memory store wait")); - let raw_mod = render_raw_mod(&catalog, "test-hash"); - assert!(raw_mod.contains( - "pub use crate::__cuda_oxide_intrinsic_abi_v1::i0578 as tcgen05_cp_128x128b_b4x16_p64;" - )); - assert!(raw_mod.contains( - "pub use crate::__cuda_oxide_intrinsic_abi_v1::i0611 as tcgen05_cp_64x128b_warpx2_02_13_cg2;" - )); - assert!(raw_mod.contains( - "pub use crate::__cuda_oxide_intrinsic_abi_v1::i0727 as tcgen05_st_32x32b_x128_unpack16;" - )); - assert!(raw_mod.contains( - "pub use crate::__cuda_oxide_intrinsic_abi_v1::i0728 as tcgen05_ld_16x32bx2_x1_raw;" - )); - assert!(raw_mod.contains( - "pub use crate::__cuda_oxide_intrinsic_abi_v1::i0759 as tcgen05_st_16x32bx2_x128_unpack16;" - )); - assert!(raw_mod.contains( - "pub use crate::__cuda_oxide_intrinsic_abi_v1::i0760 as tcgen05_commit_multicast;" - )); - assert!(raw_mod.contains( - "pub use crate::__cuda_oxide_intrinsic_abi_v1::i0761 as tcgen05_shift_down;" - )); - assert!(raw_mod.contains( - "pub use crate::__cuda_oxide_intrinsic_abi_v1::i0762 as tcgen05_shift_down_cg2;" - )); - - let dialect = render_dialect_tcgen05(&catalog, "test-hash"); - assert_eq!(dialect.matches("pub struct Tcgen05").count(), 210); - assert_eq!(dialect.matches("impl Verify for Tcgen05").count(), 210); - assert_eq!(dialect.matches("::register(ctx)").count(), 216); - assert_eq!(dialect.matches(" Some(1),").count(), 32); - assert_eq!(dialect.matches("verifier = \"succ\"").count(), 6); - assert!(dialect.contains("Operation::get_op::")); - assert!(dialect.contains("Operation::get_op::")); - assert!(dialect.contains("&[(Tcgen05Carrier::I32, 1), (Tcgen05Carrier::I64, 1)],")); - assert!(dialect.contains( - "&[(Tcgen05Carrier::I32, 1), (Tcgen05Carrier::I64, 1), (Tcgen05Carrier::I32, 128)]," - )); - assert!(dialect.contains("pub struct Tcgen05Ld16x64bX1RawOp")); - assert!(dialect.contains("pub struct Tcgen05Ld32x32bX128Pack16Op")); - assert!(dialect.contains("pub struct Tcgen05St16x64bX1RawOp")); - assert!(dialect.contains("pub struct Tcgen05St32x32bX128Unpack16Op")); - assert!(dialect.contains("pub struct Tcgen05Ld16x32bx2X1RawOp")); - assert!(dialect.contains("pub struct Tcgen05St16x32bx2X128Unpack16Op")); - assert!(dialect.contains("pub struct Tcgen05CommitMulticastOp")); - assert!(dialect.contains("pub struct Tcgen05ShiftDownOp")); - assert!(dialect.contains("pub struct Tcgen05ShiftDownCg2Op")); - assert!(dialect.contains("NResultsInterface<128>")); - assert!(dialect.contains("NResultsInterface<32>")); - assert!(dialect.contains("NResultsInterface<4>")); - assert!(dialect.contains("NOpdsInterface<6>")); - - let dialect_mod = render_dialect_mod(&catalog, "test-hash"); - assert!(dialect_mod.contains("mod tcgen05;")); - assert!(dialect_mod.contains("pub use tcgen05::*;")); - assert!(dialect_mod.contains("tcgen05::register(ctx);")); - - let importer = render_importer(&catalog, "test-hash"); - assert!(importer.contains("cuda_device::tcgen05::tcgen05_alloc")); - assert!(importer.contains("Tcgen05AllocOp::get_concrete_op_info()")); - assert!(importer.contains("Tcgen05Ld16x256bX8PureOp::get_concrete_op_info()")); - assert!(importer.contains("MirConstructStructOp::get_concrete_op_info()")); - assert!(importer.contains("Tcgen05Ld16x64bX1RawOp::get_concrete_op_info()")); - assert!(importer.contains("Tcgen05Ld32x32bX128Pack16Op::get_concrete_op_info()")); - assert!(importer.contains("Tcgen05St16x64bX1RawOp::get_concrete_op_info()")); - assert!(importer.contains("Tcgen05St32x32bX128Unpack16Op::get_concrete_op_info()")); - assert!(importer.contains("import_generated_tcgen05_store_operands(")); - assert!(importer.contains("downcast_ref::()")); - assert!( - importer - .contains("generated tcgen05 store data must contain {expected_len} u32 registers") - ); - assert!(importer.contains("destination_ty == array_result.get_type(ctx)")); - assert!(!importer.contains("super::tcgen05::emit_")); - assert!(importer.contains("\"v1:i0343\"")); - assert!(importer.contains("\"v1:i0366\"")); - assert!(importer.contains("cuda_device::tcgen05::tcgen05_cp_128x128b_b4x16_p64")); - assert!(importer.contains("Tcgen05Cp128x128bB4x16P64Op::get_concrete_op_info()")); - assert!(importer.contains("\"v1:i0578\"")); - assert!(importer.contains("\"v1:i0611\"")); - assert!(importer.contains("\"v1:i0670\"")); - assert!(importer.contains("\"v1:i0727\"")); - assert!(importer.contains("cuda_device::tcgen05::__tcgen05_ld_16x32bx2_x1_raw")); - assert!(importer.contains("cuda_device::tcgen05::__tcgen05_st_16x32bx2_x128_unpack16")); - assert!( - importer.contains("tcgen05 16x32bx2 half-split offset must be a compile-time constant") - ); - assert!(importer.contains("half-split offset must lower to a constant")); - assert!(importer.contains("args, 128, true, block_ptr")); - assert!(importer.contains("\"v1:i0728\"")); - assert!(importer.contains("\"v1:i0759\"")); - assert!(importer.contains("\"v1:i0760\"")); - assert!(importer.contains("\"v1:i0761\"")); - assert!(importer.contains("\"v1:i0762\"")); - - let lowering = render_lowering(&catalog, "test-hash"); - assert!(!lowering.contains("use crate::convert::intrinsics::tcgen05;")); - assert!(lowering.contains("impl MirToLlvmConversion for Tcgen05AllocOp")); - assert!( - lowering.contains( - "tcgen05.mma.ws.cta_group::1.kind::f16 [$0], [$1], $3, $4, %enable_pred;" - ) - ); - assert!(!lowering.contains(".kind::bf16")); - assert!(lowering.contains("convert_generated_tcgen05_load(")); - assert!(lowering.contains("tcgen05.ld.sync.aligned.16x64b.x1.b32 {$0}, [$1];")); - assert!(lowering.contains("tcgen05.ld.sync.aligned.16x64b.x1.pack::16b.b32 {$0}, [$1];")); - assert!(lowering.contains("\"=r,r,~{memory}\"")); - assert!(lowering.contains( - "convert_generated_tcgen05_load(ctx, rewriter, self.get_operation(), 1, 1, true" - )); - assert!(lowering.contains("convert_generated_tcgen05_void(")); - assert!(lowering.contains("tcgen05.st.sync.aligned.16x64b.x1.b32 [$0], {$1};")); - assert!( - lowering.contains("tcgen05.st.sync.aligned.32x32b.x128.unpack::16b.b32 [$0], {$1,$2") - ); - assert!(lowering.contains("\"r,r,~{memory}\"")); - assert!(lowering.contains("tcgen05.ld.sync.aligned.16x32bx2.x1.b32 {$0}, [$1], $2;")); - assert!(lowering.contains("tcgen05.st.sync.aligned.16x32bx2.x1.b32 [$0], $1, {$2};")); - assert!(lowering.contains("\"=r,r,n,~{memory}\"")); - assert!(lowering.contains("\"r,n,r,~{memory}\"")); - assert!(lowering.contains( - "convert_generated_tcgen05_load(ctx, rewriter, self.get_operation(), 2, 1, true" - )); - assert!(lowering.contains("builtin::types::{FP32Type, IntegerType, Signedness}")); - assert!(lowering.contains("ops as llvm_ops")); - assert!(lowering.contains("inserter::Inserter")); - assert!(lowering.contains("tcgen05.cp.cta_group::1.128x128b.b8x16.b4x16_p64 [$0], $1;")); - assert!(lowering.contains("tcgen05.cp.cta_group::1.32x128b.warpx4 [$0], $1;")); - assert!( - lowering.contains( - "tcgen05.cp.cta_group::2.64x128b.warpx2::01_23.b8x16.b6x16_p32 [$0], $1;" - ) - ); - assert!(lowering.contains("tcgen05.cp.cta_group::1.64x128b.warpx2::02_13 [$0], $1;")); - assert!(lowering.contains("\"r,l,~{memory}\"")); - assert!(lowering.contains( - "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 [$0], $1;" - )); - assert!(lowering.contains("tcgen05.shift.cta_group::1.down [$0];")); - assert!(lowering.contains("tcgen05.shift.cta_group::2.down [$0];")); - assert!(lowering.contains("\"r,h,~{memory}\"")); - - let bf16 = tcgen05_intrinsics(&catalog) - .find(|record| record.id == "tcgen05_mma_ws_bf16") - .unwrap(); - let bf16_probe = render_probe(&catalog, bf16, "test-hash"); - assert!(bf16_probe.contains("route: inline PTX")); - assert!(bf16_probe.contains(".kind::f16 [$0]")); - assert!(!bf16_probe.contains(".kind::bf16")); - - let base_commit = tcgen05_intrinsics(&catalog) - .find(|record| record.id == "tcgen05_commit") - .unwrap(); - let base_commit_probe = render_probe(&catalog, base_commit, "test-hash"); - assert!(base_commit_probe.contains("mbarrier::arrive::one.b64")); - assert!(!base_commit_probe.contains("one.shared::cluster.b64")); - - let shared_commit = tcgen05_intrinsics(&catalog) - .find(|record| record.id == "tcgen05_commit_shared_cluster") - .unwrap(); - let shared_commit_probe = render_probe(&catalog, shared_commit, "test-hash"); - assert!(shared_commit_probe.contains("one.shared::cluster.b64")); - - for (id, spelling) in [ - ( - "tcgen05_commit_multicast", - "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 [$0], $1;", - ), - ( - "tcgen05_shift_down", - "tcgen05.shift.cta_group::1.down [$0];", - ), - ( - "tcgen05_shift_down_cg2", - "tcgen05.shift.cta_group::2.down [$0];", - ), - ] { - let record = tcgen05_intrinsics(&catalog) - .find(|record| record.id == id) - .unwrap(); - let probe = render_probe(&catalog, record, "test-hash"); - assert!(probe.contains(spelling)); - assert!(probe.contains("asm sideeffect")); - assert!(probe.contains("attributes #0 = { convergent }")); - } - - for (id, spelling) in [ - ( - "tcgen05_cp_128x128b_b4x16_p64", - "tcgen05.cp.cta_group::1.128x128b.b8x16.b4x16_p64", - ), - ( - "tcgen05_cp_32x128b_warpx4", - "tcgen05.cp.cta_group::1.32x128b.warpx4", - ), - ( - "tcgen05_cp_64x128b_warpx2_01_23_b6x16_p32_cg2", - "tcgen05.cp.cta_group::2.64x128b.warpx2::01_23.b8x16.b6x16_p32", - ), - ( - "tcgen05_cp_64x128b_warpx2_02_13", - "tcgen05.cp.cta_group::1.64x128b.warpx2::02_13", - ), - ] { - let record = tcgen05_intrinsics(&catalog) - .find(|record| record.id == id) - .unwrap(); - let probe = render_probe(&catalog, record, "test-hash"); - assert!(probe.contains(spelling)); - assert!(probe.contains("asm sideeffect")); - assert!(probe.contains("\"r,l,~{memory}\"")); - assert!(probe.contains("attributes #0 = { convergent }")); - } - - let scalar_load = tcgen05_intrinsics(&catalog) - .find(|record| record.id == "tcgen05_ld_16x64b_x1_raw") - .unwrap(); - let scalar_load_probe = render_probe(&catalog, scalar_load, "test-hash"); - assert!(scalar_load_probe.contains("%result = call i32 asm sideeffect")); - assert!(scalar_load_probe.contains("tcgen05.ld.sync.aligned.16x64b.x1.b32")); - assert!(scalar_load_probe.contains("\"=r,r,~{memory}\"")); - - let packed_load = tcgen05_intrinsics(&catalog) - .find(|record| record.id == "tcgen05_ld_32x32b_x128_pack16") - .unwrap(); - let packed_load_probe = render_probe(&catalog, packed_load, "test-hash"); - assert!(packed_load_probe.contains("tcgen05.ld.sync.aligned.32x32b.x128.pack::16b.b32")); - assert!(packed_load_probe.contains("{ i32, i32")); - - let scalar_store = tcgen05_intrinsics(&catalog) - .find(|record| record.id == "tcgen05_st_16x64b_x1_raw") - .unwrap(); - let scalar_store_probe = render_probe(&catalog, scalar_store, "test-hash"); - assert!(scalar_store_probe.contains("call void asm sideeffect")); - assert!(scalar_store_probe.contains("tcgen05.st.sync.aligned.16x64b.x1.b32")); - assert!(scalar_store_probe.contains("\"r,r,~{memory}\"")); - - let unpacked_store = tcgen05_intrinsics(&catalog) - .find(|record| record.id == "tcgen05_st_32x32b_x128_unpack16") - .unwrap(); - let unpacked_store_probe = render_probe(&catalog, unpacked_store, "test-hash"); - assert!( - unpacked_store_probe.contains("tcgen05.st.sync.aligned.32x32b.x128.unpack::16b.b32") - ); - assert!(unpacked_store_probe.contains("i32 %d127")); - - let offset_load = tcgen05_intrinsics(&catalog) - .find(|record| record.id == "tcgen05_ld_16x32bx2_x1_pack16") - .unwrap(); - let offset_load_probe = render_probe(&catalog, offset_load, "test-hash"); - assert!(offset_load_probe.contains("i64 16")); - assert!(offset_load_probe.contains("[$1], $2;")); - assert!(offset_load_probe.contains("\"=r,r,n,~{memory}\"")); - - let offset_store = tcgen05_intrinsics(&catalog) - .find(|record| record.id == "tcgen05_st_16x32bx2_x1_unpack16") - .unwrap(); - let offset_store_probe = render_probe(&catalog, offset_store, "test-hash"); - assert!(offset_store_probe.contains("i64 16")); - assert!(offset_store_probe.contains("[$0], $1, {$2};")); - assert!(offset_store_probe.contains("\"r,n,r,~{memory}\"")); - - let target = render_targets(&catalog, "test-hash"); - assert!(target.contains("\"v1:i0343\"")); - assert!(target.contains("\"v1:i0578\"")); - assert!(target.contains("\"v1:i0611\"")); - assert!(target.contains("\"v1:i0727\"")); - assert!(target.contains("\"v1:i0728\"")); - assert!(target.contains("\"v1:i0759\"")); - assert!(target.contains( - "Tcgen05MmaBUsageAttr, Tcgen05MmaCollectorAAttr, Tcgen05MmaCtaGroupAttr, Tcgen05MmaFormAttr" - )); - assert!(target.contains( - "GeneratedHardwareTarget::AnyOf(&[GeneratedHardwareAlternative::ExactArchitecture(100), GeneratedHardwareAlternative::ExactArchitecture(101), GeneratedHardwareAlternative::ExactArchitecture(103), GeneratedHardwareAlternative::ExactArchitecture(110)])" - )); - assert!(target.contains( - "GeneratedHardwareTarget::AnyOf(&[GeneratedHardwareAlternative::ExactArchitecture(100), GeneratedHardwareAlternative::ExactArchitecture(103), GeneratedHardwareAlternative::ExactArchitecture(110)])" - )); - - let reference = render_reference(&catalog, "test-hash"); - assert!(reference.contains("- `tcgen05_st_16x64b_x1_raw`: runtime `unexecuted`")); - assert!(!reference.contains("- `tcgen05_st_16x64b_x1_raw`: runtime `not recorded`")); - - let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); - assert!(outputs.contains_key(&PathBuf::from( - "crates/cuda-device/src/generated/tcgen05.rs" - ))); - assert!(outputs.contains_key(&PathBuf::from( - "crates/dialect-nvvm/src/ops/generated/tcgen05.rs" - ))); - - let mut wrong_adapter = catalog.clone(); - wrong_adapter - .intrinsics - .iter_mut() - .find(|record| record.id == "tcgen05_alloc") - .unwrap() - .tcgen05 - .as_mut() - .unwrap() - .adapter = Tcgen05Adapter::NoOperands; - assert!(validate_renderable(&wrong_adapter).is_err()); - - let mut wrong_libnvvm_target = catalog.clone(); - wrong_libnvvm_target - .intrinsics - .iter_mut() - .find(|record| record.id == "tcgen05_cp_128x128b_b4x16_p64") - .unwrap() - .backend_lowerings[1] - .target - .hardware = CatalogHardwareTarget::AnyOf { - alternatives: vec![ - CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 101 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 103 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, - ], - }; - assert!(validate_renderable(&wrong_libnvvm_target).is_err()); - - let mut wrong_load_selector = catalog.clone(); - wrong_load_selector - .intrinsics - .iter_mut() - .find(|record| record.id == "tcgen05_ld_16x64b_x1_pack16") - .unwrap() - .expected_ptx - .modifiers - .remove(5); - assert!(validate_renderable(&wrong_load_selector).is_err()); - - let mut wrong_store_selector = catalog.clone(); - wrong_store_selector - .intrinsics - .iter_mut() - .find(|record| record.id == "tcgen05_st_16x64b_x1_unpack16") - .unwrap() - .expected_ptx - .modifiers - .remove(5); - assert!(validate_renderable(&wrong_store_selector).is_err()); - - let mut wrong_offset_scope = catalog.clone(); - wrong_offset_scope - .intrinsics - .iter_mut() - .find(|record| record.id == "tcgen05_ld_16x32bx2_x1_raw") - .unwrap() - .semantics - .execution_scope = "thread".into(); - assert!(validate_renderable(&wrong_offset_scope).is_err()); - - let mut wrong_copy_spelling = catalog; - wrong_copy_spelling - .intrinsics - .iter_mut() - .find(|record| record.id == "tcgen05_cp_128x128b_b4x16_p64") - .unwrap() - .expected_ptx - .modifiers - .remove(3); - assert!(validate_renderable(&wrong_copy_spelling).is_err()); - } - - #[test] - fn debug_control_rendering_preserves_api_immediates_and_side_effects() { - let catalog = catalog_with_debug_controls(); - validate_renderable(&catalog).unwrap(); - assert_eq!(debug_controls(&catalog).count(), 3); - - let compatibility = render_compat_debug_control(&catalog, "test-hash"); - assert!(compatibility.contains("pub fn trap() -> !")); - assert!(compatibility.contains("pub fn breakpoint()")); - assert!(compatibility.contains("pub fn prof_trigger()")); - assert!(compatibility.contains("pub(crate) fn __prof_trigger(_event_id: u32)")); - assert!(compatibility.contains("__prof_trigger(N);")); - - let dialect = render_dialect_debug_control(&catalog, "test-hash"); - for op in ["TrapOp", "BreakpointOp", "PmEventOp"] { - assert_eq!(dialect.matches(&format!("pub struct {op}")).count(), 1); - assert!(dialect.contains(&format!("{op}::register(ctx)"))); - } - assert!(dialect.contains("event_id: u32")); - assert!(dialect.contains("pub fn new_with_event_id")); - assert!(dialect.contains("pub fn get_event_id")); - assert!(dialect.contains("filter(|value| *value <= 15)")); - assert!(dialect.contains("requires a u32 event ID in 0..=15")); - - let importer = render_importer(&catalog, "test-hash"); - assert!(importer.contains("cuda_device::debug::__prof_trigger")); - assert!(importer.contains("mir::Operand::Constant")); - assert!(importer.contains("u32::try_from(value)")); - assert!(importer.contains("filter(|value| *value <= 15)")); - assert!(importer.contains("PmEventOp::build(ctx, event_id)")); - assert!(importer.contains("MirUnreachableOp::get_concrete_op_info()")); - assert!( - importer.contains("prof_trigger requires a compile-time constant event ID in 0..=15") - ); - - let lowering = render_lowering(&catalog, "test-hash"); - for op in ["TrapOp", "BreakpointOp", "PmEventOp"] { - assert!(lowering.contains(&format!("impl MirToLlvmConversion for {op}"))); - } - assert!(lowering.contains( - "inline_asm_sideeffect(ctx, rewriter, op, void_ty.into(), vec![], \"trap;\", \"\")" - )); - assert!(lowering.contains( - "inline_asm_sideeffect(ctx, rewriter, op, void_ty.into(), vec![], \"brkpt;\", \"\")" - )); - assert!(lowering.contains("let template = format!(\"pmevent {event_id};\")")); - - for record in debug_controls(&catalog) { - let probe = render_probe(&catalog, record, "test-hash"); - assert!(probe.contains("call void asm sideeffect")); - assert!(!probe.contains("attributes #0 = { convergent }")); - } - - let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); - assert!(outputs.contains_key(&PathBuf::from( - "crates/cuda-device/src/generated/debug_control.rs" - ))); - assert!(outputs.contains_key(&PathBuf::from( - "crates/dialect-nvvm/src/ops/generated/debug_control.rs" - ))); - - let mut wrong_adapter = catalog.clone(); - wrong_adapter - .intrinsics - .iter_mut() - .find(|record| record.id == "pmevent") - .unwrap() - .debug_control - .as_mut() - .unwrap() - .adapter = DebugControlAdapter::Direct; - assert!(validate_renderable(&wrong_adapter).is_err()); - - let mut wrong_runtime = catalog; - wrong_runtime - .intrinsics - .iter_mut() - .find(|record| record.id == "trap") - .unwrap() - .debug_control - .as_mut() - .unwrap() - .runtime_validation = RuntimeValidation::Executed; - assert!(validate_renderable(&wrong_runtime).is_err()); - } - - #[test] - fn scalar_arithmetic_rendering_uses_one_closed_carrier() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - validate_renderable(&catalog).unwrap(); - let records = scalar_arithmetics(&catalog).collect::>(); - assert_eq!(records.len(), 64); - assert!(records.iter().all(|record| { - record.semantics.pure - == (record.scalar_arithmetic.as_ref().unwrap().operation - != ScalarArithmeticOperation::Div) - })); - - let compatibility = render_compat_float(&catalog, "test-hash"); - assert!(compatibility.contains("pub fn mul_rn_f64(arg0: f64, arg1: f64) -> f64")); - assert!( - compatibility - .contains("pub fn fma_rp_ftz_sat_f32(arg0: f32, arg1: f32, arg2: f32) -> f32") - ); - assert!(compatibility.contains("pub fn add_rp_ftz_sat_f32(arg0: f32, arg1: f32) -> f32")); - assert_eq!(compatibility.matches("#[must_use]").count(), 113); - - let dialect = render_dialect_scalar_arithmetic(&catalog, "test-hash"); - assert_eq!(dialect.matches("pub struct ScalarArithmeticOp").count(), 1); - for attr in [ - "ScalarArithmeticFormatAttr", - "ScalarArithmeticOperationAttr", - "ScalarArithmeticRoundingAttr", - "ScalarArithmeticSubnormalAttr", - "ScalarArithmeticSaturationAttr", - ] { - assert!(dialect.contains(attr)); - assert!(dialect.contains(&format!("{attr}::register(ctx)"))); - } - assert!(dialect.contains("variant is not admitted")); - assert!(dialect.contains("ScalarArithmeticOperationAttr::Fma => 3")); - - let importer = render_importer(&catalog, "test-hash"); - assert_eq!( - importer.matches("ScalarArithmeticOp::build(ctx").count(), - 64 - ); - assert!(importer.contains("cuda_device::float::fma_rp_ftz_sat_f32")); - assert!(importer.contains("cuda_device::float::add_rp_ftz_sat_f32")); - - let lowering = render_lowering(&catalog, "test-hash"); - assert_eq!( - lowering - .matches("impl MirToLlvmConversion for ScalarArithmeticOp") - .count(), - 1 - ); - assert!(lowering.contains("llvm_nvvm_mul_rn_d")); - assert!(lowering.contains("fma.rp.ftz.sat.f32")); - assert!(lowering.contains("llvm_nvvm_add_rn_d")); - assert!(lowering.contains("add.rp.sat.ftz.f32")); - assert!(lowering.contains("recipe.0, recipe.1, recipe.2, recipe.3")); - - let targets = render_targets(&catalog, "test-hash"); - assert!(targets.contains("GeneratedIntrinsicVariant::ScalarArithmetic")); - assert!(targets.contains("GeneratedScalarArithmeticFormat::F64")); - assert!(targets.contains("GeneratedScalarArithmeticSaturation::Sat")); - assert!(targets.contains("Operation::get_op::")); - - let f64_mul = records - .iter() - .copied() - .find(|record| record.id == "mul_rn_f64") - .unwrap(); - let f64_probe = render_probe(&catalog, f64_mul, "test-hash"); - assert!(f64_probe.contains("declare double @llvm.nvvm.mul.rn.d(double, double)")); - let f32_sat = records - .iter() - .copied() - .find(|record| record.id == "fma_rp_ftz_sat_f32") - .unwrap(); - let f32_probe = render_probe(&catalog, f32_sat, "test-hash"); - assert!( - f32_probe - .contains("call float asm \"fma.rp.ftz.sat.f32 $0, $1, $2, $3;\", \"=f,f,f,f\"") - ); - assert!(!f32_probe.contains("@llvm.nvvm.fma.rp.ftz.sat.f")); - - let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); - assert!(outputs.contains_key(&PathBuf::from("crates/cuda-device/src/generated/float.rs"))); - assert!(outputs.contains_key(&PathBuf::from( - "crates/dialect-nvvm/src/ops/generated/scalar_arithmetic.rs" - ))); - } - - #[test] - fn extended_minmax_rendering_is_closed_pure_and_exact() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - validate_renderable(&catalog).unwrap(); - let records = extended_minmax(&catalog).collect::>(); - assert_eq!(records.len(), 52); - assert_eq!( - records - .iter() - .filter(|record| record.target.minimum_ptx.encoded() == 70) - .count(), - 20 - ); - assert_eq!( - records - .iter() - .filter(|record| record.target.minimum_ptx.encoded() == 72) - .count(), - 32 - ); - for id in ["min_f16x2", "max_f16x2", "min_bf16x2", "max_bf16x2"] { - assert_eq!( - catalog - .intrinsics - .iter() - .find(|record| record.id == id) - .unwrap() - .family, - "packed_alu" - ); - assert!(records.iter().all(|record| record.id != id)); - } - - let f16 = render_compat_packed_alu(&catalog, "test-hash", PackedAluFormat::F16x2); - assert!(f16.contains("pub fn min_ftz_nan_xorsign_abs_f16x2(a: u32, b: u32) -> u32")); - let bf16 = render_compat_packed_alu(&catalog, "test-hash", PackedAluFormat::Bf16x2); - assert!(bf16.contains("pub fn max_nan_xorsign_abs_bf16x2(a: u32, b: u32) -> u32")); - let float = render_compat_float(&catalog, "test-hash"); - assert!(float.contains("pub fn min_ftz_nan_xorsign_abs_f32(a: f32, b: f32) -> f32")); - - let dialect = render_dialect_extended_minmax(&catalog, "test-hash"); - assert_eq!(dialect.matches("pub struct ExtendedMinMaxOp").count(), 1); - assert_eq!( - dialect - .matches(" (ExtendedMinMaxFormatAttr::") - .count(), - 52 - ); - for attr in [ - "ExtendedMinMaxFormatAttr", - "ExtendedMinMaxOperationAttr", - "ExtendedMinMaxSubnormalAttr", - "ExtendedMinMaxNanAttr", - "ExtendedMinMaxXorSignAbsAttr", - ] { - assert!(dialect.contains(attr)); - assert!(dialect.contains(&format!("{attr}::register(ctx)"))); - } - assert!(dialect.contains("variant is not admitted")); - - let importer = render_importer(&catalog, "test-hash"); - assert_eq!(importer.matches("ExtendedMinMaxOp::build(ctx").count(), 52); - assert!(importer.contains("cuda_device::float::min_xorsign_abs_f32")); - assert!(importer.contains("cuda_device::bf16x2::max_nan_bf16x2")); - - let lowering = render_lowering(&catalog, "test-hash"); - assert_eq!( - lowering - .matches("impl MirToLlvmConversion for ExtendedMinMaxOp") - .count(), - 1 - ); - assert!(lowering.contains("convert_generated_extended_minmax")); - assert!(lowering.contains("min.ftz.NaN.xorsign.abs.f32")); - assert!(lowering.contains("max.NaN.xorsign.abs.bf16x2")); - - let targets = render_targets(&catalog, "test-hash"); - assert!(targets.contains("GeneratedIntrinsicVariant::ExtendedMinMax")); - assert!(targets.contains("GeneratedPtxVersion::from_encoded(72)")); - assert!(targets.contains("Operation::get_op::")); - assert!(targets.contains("ExtendedMinMaxXorSignAbsAttr::Enabled")); - - let packed_record = records - .iter() - .copied() - .find(|record| record.id == "min_ftz_nan_f16x2") - .unwrap(); - let packed_probe = render_probe(&catalog, packed_record, "test-hash"); - assert!( - packed_probe.contains("call i32 asm \"min.ftz.NaN.f16x2 $0, $1, $2;\", \"=r,r,r\"") - ); - assert!(!packed_probe.contains("sideeffect")); - assert!(!packed_probe.contains("convergent")); - let f32_record = records - .iter() - .copied() - .find(|record| record.id == "max_ftz_nan_xorsign_abs_f32") - .unwrap(); - let f32_probe = render_probe(&catalog, f32_record, "test-hash"); - assert!( - f32_probe - .contains("call float asm \"max.ftz.NaN.xorsign.abs.f32 $0, $1, $2;\", \"=f,f,f\"") - ); - // Scalar 16-bit forms ride in `h` registers, not the packed `r` pair. - let half_record = records - .iter() - .copied() - .find(|record| record.id == "min_ftz_nan_f16") - .unwrap(); - let half_probe = render_probe(&catalog, half_record, "test-hash"); - assert!(half_probe.contains("call i16 asm \"min.ftz.NaN.f16 $0, $1, $2;\", \"=h,h,h\"")); - let bf16_record = records - .iter() - .copied() - .find(|record| record.id == "max_nan_xorsign_abs_bf16") - .unwrap(); - let bf16_probe = render_probe(&catalog, bf16_record, "test-hash"); - assert!( - bf16_probe - .contains("call i16 asm \"max.NaN.xorsign.abs.bf16 $0, $1, $2;\", \"=h,h,h\"") - ); - assert!(lowering.contains("MinMaxCarrier::Half16")); - - let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); - assert!(outputs.contains_key(&PathBuf::from( - "crates/dialect-nvvm/src/ops/generated/extended_minmax.rs" - ))); - for module in ["float", "f16", "bf16", "f16x2", "bf16x2"] { - assert!(outputs.contains_key(&PathBuf::from(format!( - "crates/cuda-device/src/generated/{module}.rs" - )))); - } - - let mut wrong_adapter = catalog; - wrong_adapter - .intrinsics - .iter_mut() - .find(|record| record.id == "min_xorsign_abs_f32") - .unwrap() - .extended_minmax - .as_mut() - .unwrap() - .adapter = ExtendedMinMaxAdapter::DirectPackedU32; - assert!(validate_renderable(&wrong_adapter).is_err()); - } - - #[test] - fn extended_minmax_alone_renders_float_wrappers() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let mut catalog = crate::resolve::resolve(&repo_root).unwrap(); - catalog - .intrinsics - .retain(|record| record.family == "extended_minmax"); - - let (path, float) = render_compat_float_output(&catalog, "test-hash") - .expect("extended min/max alone must render float wrappers"); - assert_eq!( - path, - PathBuf::from("crates/cuda-device/src/generated/float.rs") - ); - assert!(float.contains("pub fn min_xorsign_abs_f32(a: f32, b: f32) -> f32")); - assert!(!float.contains("pub fn add_rn_f32")); - } - - /// Every reference row must carry exactly the six cells of its header. - /// - /// A `|` in cell data splits the row before inline parsing, so it splits - /// the cell even inside a code span. Three catalog entries write a PTX - /// destination pair as ``, and their rows rendered with - /// a seventh cell, which drops the backend-evidence column. Counting - /// unescaped pipes per row is the direct check. - #[test] - fn reference_rows_have_one_cell_per_header_column() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = crate::resolve::resolve(&repo_root).unwrap(); - let reference = render_reference(&catalog, "test-hash"); - - let mut rows = 0usize; - let mut carried_a_pipe = 0usize; - for line in reference - .lines() - .filter(|line| line.starts_with("| `cuda_intrinsics::")) - { - rows += 1; - if line.contains("\\|") { - carried_a_pipe += 1; - } - // Delimiters only: drop every escaped pipe first, then a six-column - // row leaves seven. - let delimiters = line.replace("\\|", "").matches('|').count(); - assert_eq!( - delimiters, 7, - "row has {delimiters} cell delimiters, expected 7: {line}" - ); - } - - assert!( - rows > 900, - "expected the full reference, rendered {rows} rows" - ); - assert!( - carried_a_pipe >= 3, - "expected the elect.sync and match.all.sync patterns to need escaping, \ - found {carried_a_pipe} rows with an escaped pipe" - ); - } -} diff --git a/crates/cuda-intrinsics-gen/src/render/collector_targets.rs b/crates/cuda-intrinsics-gen/src/render/collector_targets.rs new file mode 100644 index 0000000000..68fc4ff2a6 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/render/collector_targets.rs @@ -0,0 +1,931 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + CatalogFile, CatalogIntrinsic, CatalogSelection, ClusterBarrierMode, ExtendedMinMaxFormat, + ExtendedMinMaxNan, ExtendedMinMaxOperation, ExtendedMinMaxSubnormal, ImportedAddressSpace, + IntrinsicBackend, LdmatrixElement, LdmatrixLayout, LdmatrixMultiplicity, LdmatrixShape, + PackedAtomicFormat, PrmtMode, RegisterMmaAccumulator, RegisterMmaElement, RegisterMmaKind, + RegisterMmaLayout, RegisterMmaOperation, RegisterMmaOverflow, RegisterMmaShape, + ScalarArithmeticFormat, ScalarArithmeticOperation, ScalarArithmeticRounding, + ScalarArithmeticSaturation, ScalarArithmeticSubnormal, ScalarConversionRounding, + ScalarConversionSaturation, SparseMmaAccumulator, SparseMmaElement, SparseMmaLayout, + SparseMmaMetadata, SparseMmaOverflow, SparseMmaSelector, SparseMmaShape, WgmmaControlMode, +}; +use crate::render::common::{generated_hardware_target, intrinsic_marker, rust_header}; +use crate::render::families::{ + extended_minmax, ldmatrix, ldmatrix_compat_op, register_mma_effective_kind, scalar_arithmetics, + scalar_conversions, tcgen05_mma_form_name, tcgen05_mma_intrinsics, wgmma_controls, +}; +use crate::render::reference::render_string_patterns; +use std::fmt::Write as _; + +pub(super) fn render_collector(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + let abi_namespace = format!("__cuda_oxide_intrinsic_abi_v{}", catalog.intrinsic_abi); + output.push_str("//! Generated collector predicates for intrinsic placeholder functions.\n\n"); + writeln!( + output, + "pub(crate) const GENERATED_INTRINSIC_ABI: u32 = {};", + catalog.intrinsic_abi + ) + .unwrap(); + writeln!( + output, + "pub(crate) const GENERATED_INTRINSIC_ABI_NAMESPACE: &str = {abi_namespace:?};" + ) + .unwrap(); + output.push_str( + "pub(crate) const GENERATED_INTRINSIC_CRATES: &[&str] = &[\n \"cuda_intrinsics\",\n \"cuda-intrinsics\",\n];\n\npub(crate) const GENERATED_INTRINSIC_CANONICAL_PATHS: &[&str] = &[\n", + ); + for record in &catalog.intrinsics { + writeln!(output, " {:?},", record.rust.canonical_path).unwrap(); + } + output.push_str("];\n\npub(crate) const GENERATED_INTRINSIC_PUBLIC_PATHS: &[&str] = &[\n"); + for record in &catalog.intrinsics { + writeln!(output, " {:?},", record.rust.public_path).unwrap(); + } + output.push_str("];\n\npub(crate) fn is_generated_intrinsic_crate(crate_name: &str) -> bool {\n GENERATED_INTRINSIC_CRATES.contains(&crate_name)\n}\n\npub(crate) fn is_generated_intrinsic_canonical_path(path: &str) -> bool {\n matches!(path,\n"); + let canonical_paths: Vec<_> = catalog + .intrinsics + .iter() + .map(|record| record.rust.canonical_path.as_str()) + .collect(); + render_string_patterns(&mut output, &canonical_paths, " "); + output.push_str(" )\n}\n\npub(crate) fn is_generated_intrinsic_compatibility_path(path: &str) -> bool {\n"); + let compatibility_paths: Vec<_> = catalog + .intrinsics + .iter() + .flat_map(|record| record.rust.compatibility_paths.iter().map(String::as_str)) + .collect(); + if compatibility_paths.is_empty() { + output.push_str(" let _ = path;\n false\n"); + } else { + output.push_str(" matches!(path,\n"); + render_string_patterns(&mut output, &compatibility_paths, " "); + output.push_str(" )\n"); + } + output.push_str( + "}\n\npub(crate) fn is_generated_intrinsic_placeholder(crate_name: &str, path: &str) -> bool {\n if is_generated_intrinsic_crate(crate_name) {\n is_generated_intrinsic_canonical_path(path)\n } else if matches!(crate_name, \"cuda_device\" | \"cuda-device\") {\n is_generated_intrinsic_compatibility_path(path)\n } else {\n false\n }\n}\n", + ); + output +} + +fn generated_selection_constraints(selection: &CatalogSelection) -> String { + let address_space = match selection.constraints.address_space { + None => "None", + Some(ImportedAddressSpace::Generic) => "Some(GeneratedSelectionAddressSpace::Generic)", + Some(ImportedAddressSpace::Shared) => "Some(GeneratedSelectionAddressSpace::Shared)", + }; + let immediate_bindings = selection + .constraints + .immediate_bindings + .iter() + .map(|binding| { + format!( + "GeneratedImmediateBinding {{ argument_index: {}, value: {} }}", + binding.argument_index, binding.value + ) + }) + .collect::>() + .join(", "); + format!( + "GeneratedSelectionConstraints {{ address_space: {address_space}, immediate_bindings: &[{immediate_bindings}] }}" + ) +} + +pub(super) fn generated_selection_alternatives(selections: &[CatalogSelection]) -> String { + let mut output = String::from("&["); + for selection in selections { + write!( + output, + "GeneratedSelectionAlternative {{ source_record: {:?}, asm: {:?}, predicates: &{:?}, constraints: {} }},", + selection.source_record, + selection.asm, + selection.predicates, + generated_selection_constraints(selection), + ) + .unwrap(); + } + output.push(']'); + output +} + +pub(super) fn generated_intrinsic_variant(record: &CatalogIntrinsic) -> String { + if let Some(mma) = record + .tcgen05 + .as_ref() + .and_then(|tcgen05| tcgen05.mma.as_ref()) + { + return format!( + "GeneratedIntrinsicVariant::Tcgen05Mma {{ form: GeneratedTcgen05MmaForm::{}, target_selector: GeneratedTcgen05MmaTargetSelector::Kind, compatibility_alias: {} }}", + tcgen05_mma_form_name(mma.form), + mma.alias.is_some(), + ); + } + if let Some(control) = &record.wgmma_control { + let mode = match control.mode { + WgmmaControlMode::Fence => "GeneratedWgmmaControlMode::Fence", + WgmmaControlMode::CommitGroup => "GeneratedWgmmaControlMode::CommitGroup", + WgmmaControlMode::WaitGroup => "GeneratedWgmmaControlMode::WaitGroup", + }; + return format!("GeneratedIntrinsicVariant::WgmmaControl {{ mode: {mode} }}"); + } + if let Some(barrier) = &record.cluster_barrier { + let mode = match barrier.mode { + ClusterBarrierMode::Arrive => "GeneratedClusterBarrierMode::Arrive", + ClusterBarrierMode::ArriveAligned => "GeneratedClusterBarrierMode::ArriveAligned", + ClusterBarrierMode::ArriveRelaxed => "GeneratedClusterBarrierMode::ArriveRelaxed", + ClusterBarrierMode::ArriveRelaxedAligned => { + "GeneratedClusterBarrierMode::ArriveRelaxedAligned" + } + ClusterBarrierMode::Wait => "GeneratedClusterBarrierMode::Wait", + ClusterBarrierMode::WaitAligned => "GeneratedClusterBarrierMode::WaitAligned", + }; + return format!("GeneratedIntrinsicVariant::ClusterBarrier {{ mode: {mode} }}"); + } + if let Some(prmt) = &record.prmt { + let mode = match prmt.mode { + PrmtMode::Generic => "GeneratedPrmtMode::Generic", + PrmtMode::F4e => "GeneratedPrmtMode::F4e", + PrmtMode::B4e => "GeneratedPrmtMode::B4e", + PrmtMode::Rc8 => "GeneratedPrmtMode::Rc8", + PrmtMode::Ecl => "GeneratedPrmtMode::Ecl", + PrmtMode::Ecr => "GeneratedPrmtMode::Ecr", + PrmtMode::Rc16 => "GeneratedPrmtMode::Rc16", + }; + return format!("GeneratedIntrinsicVariant::Prmt {{ mode: {mode} }}"); + } + if let Some(packed) = &record.packed_atomic { + let format = match packed.format { + PackedAtomicFormat::F16x2 => "GeneratedPackedAtomicFormat::F16x2", + PackedAtomicFormat::Bf16x2 => "GeneratedPackedAtomicFormat::Bf16x2", + }; + return format!("GeneratedIntrinsicVariant::PackedAtomic {{ format: {format} }}"); + } + if let Some(mma) = &record.register_mma { + let shape = match mma.shape { + RegisterMmaShape::M8n8k4 => "GeneratedRegisterMmaShape::M8n8k4", + RegisterMmaShape::M8n8k16 => "GeneratedRegisterMmaShape::M8n8k16", + RegisterMmaShape::M8n8k32 => "GeneratedRegisterMmaShape::M8n8k32", + RegisterMmaShape::M8n8k128 => "GeneratedRegisterMmaShape::M8n8k128", + RegisterMmaShape::M16n8k4 => "GeneratedRegisterMmaShape::M16n8k4", + RegisterMmaShape::M16n8k8 => "GeneratedRegisterMmaShape::M16n8k8", + RegisterMmaShape::M16n8k16 => "GeneratedRegisterMmaShape::M16n8k16", + RegisterMmaShape::M16n8k32 => "GeneratedRegisterMmaShape::M16n8k32", + RegisterMmaShape::M16n8k64 => "GeneratedRegisterMmaShape::M16n8k64", + RegisterMmaShape::M16n8k128 => "GeneratedRegisterMmaShape::M16n8k128", + RegisterMmaShape::M16n8k256 => "GeneratedRegisterMmaShape::M16n8k256", + }; + let operation = match mma.operation { + RegisterMmaOperation::Multiply => "GeneratedRegisterMmaOperation::Multiply", + RegisterMmaOperation::AndPopc => "GeneratedRegisterMmaOperation::AndPopc", + RegisterMmaOperation::XorPopc => "GeneratedRegisterMmaOperation::XorPopc", + }; + let kind = match register_mma_effective_kind(record) { + RegisterMmaKind::Standard => "GeneratedRegisterMmaKind::Standard", + RegisterMmaKind::F8f6f4 => "GeneratedRegisterMmaKind::F8f6f4", + RegisterMmaKind::Mxf8f6f4 => "GeneratedRegisterMmaKind::Mxf8f6f4", + }; + let accumulator = match mma.accumulator { + RegisterMmaAccumulator::F16 => "GeneratedRegisterMmaAccumulator::F16", + RegisterMmaAccumulator::F32 => "GeneratedRegisterMmaAccumulator::F32", + RegisterMmaAccumulator::F64 => "GeneratedRegisterMmaAccumulator::F64", + RegisterMmaAccumulator::S32 => "GeneratedRegisterMmaAccumulator::S32", + }; + let element = |element| match element { + RegisterMmaElement::Bf16 => "GeneratedRegisterMmaElement::Bf16", + RegisterMmaElement::E2m1 => "GeneratedRegisterMmaElement::E2m1", + RegisterMmaElement::E2m3 => "GeneratedRegisterMmaElement::E2m3", + RegisterMmaElement::E3m2 => "GeneratedRegisterMmaElement::E3m2", + RegisterMmaElement::E4m3 => "GeneratedRegisterMmaElement::E4m3", + RegisterMmaElement::E5m2 => "GeneratedRegisterMmaElement::E5m2", + RegisterMmaElement::F16 => "GeneratedRegisterMmaElement::F16", + RegisterMmaElement::Tf32 => "GeneratedRegisterMmaElement::Tf32", + RegisterMmaElement::F64 => "GeneratedRegisterMmaElement::F64", + RegisterMmaElement::B1 => "GeneratedRegisterMmaElement::B1", + RegisterMmaElement::S4 => "GeneratedRegisterMmaElement::S4", + RegisterMmaElement::U4 => "GeneratedRegisterMmaElement::U4", + RegisterMmaElement::S8 => "GeneratedRegisterMmaElement::S8", + RegisterMmaElement::U8 => "GeneratedRegisterMmaElement::U8", + }; + let layout = |layout| match layout { + RegisterMmaLayout::Row => "GeneratedRegisterMmaLayout::Row", + RegisterMmaLayout::Col => "GeneratedRegisterMmaLayout::Col", + }; + let overflow = match mma.overflow { + RegisterMmaOverflow::NotApplicable => "GeneratedRegisterMmaOverflow::NotApplicable", + RegisterMmaOverflow::Wrapping => "GeneratedRegisterMmaOverflow::Wrapping", + RegisterMmaOverflow::Satfinite => "GeneratedRegisterMmaOverflow::Satfinite", + }; + return format!( + "GeneratedIntrinsicVariant::RegisterMma {{ shape: {shape}, operation: {operation}, kind: {kind}, accumulator: {accumulator}, a_element: {}, b_element: {}, a_layout: {}, b_layout: {}, overflow: {overflow} }}", + element(mma.a_element), + element(mma.b_element), + layout(mma.a_layout), + layout(mma.b_layout), + ); + } + if let Some(mma) = &record.sparse_mma { + let shape = match mma.shape { + SparseMmaShape::M16n8k32 => "GeneratedSparseMmaShape::M16n8k32", + SparseMmaShape::M16n8k64 => "GeneratedSparseMmaShape::M16n8k64", + SparseMmaShape::M16n8k128 => "GeneratedSparseMmaShape::M16n8k128", + }; + let accumulator = match mma.accumulator { + SparseMmaAccumulator::F16 => "GeneratedSparseMmaAccumulator::F16", + SparseMmaAccumulator::F32 => "GeneratedSparseMmaAccumulator::F32", + SparseMmaAccumulator::S32 => "GeneratedSparseMmaAccumulator::S32", + }; + let element = |element| match element { + SparseMmaElement::E2m1 => "GeneratedSparseMmaElement::E2m1", + SparseMmaElement::E2m3 => "GeneratedSparseMmaElement::E2m3", + SparseMmaElement::E3m2 => "GeneratedSparseMmaElement::E3m2", + SparseMmaElement::E4m3 => "GeneratedSparseMmaElement::E4m3", + SparseMmaElement::E5m2 => "GeneratedSparseMmaElement::E5m2", + SparseMmaElement::S4 => "GeneratedSparseMmaElement::S4", + SparseMmaElement::U4 => "GeneratedSparseMmaElement::U4", + SparseMmaElement::S8 => "GeneratedSparseMmaElement::S8", + SparseMmaElement::U8 => "GeneratedSparseMmaElement::U8", + }; + let layout = |layout| match layout { + SparseMmaLayout::Row => "GeneratedSparseMmaLayout::Row", + SparseMmaLayout::Col => "GeneratedSparseMmaLayout::Col", + }; + let overflow = match mma.overflow { + SparseMmaOverflow::NotApplicable => "GeneratedSparseMmaOverflow::NotApplicable", + SparseMmaOverflow::Wrapping => "GeneratedSparseMmaOverflow::Wrapping", + SparseMmaOverflow::Satfinite => "GeneratedSparseMmaOverflow::Satfinite", + }; + let metadata = match mma.metadata { + SparseMmaMetadata::Standard => "GeneratedSparseMmaMetadata::Standard", + SparseMmaMetadata::Ordered => "GeneratedSparseMmaMetadata::Ordered", + }; + let selector = match mma.selector { + SparseMmaSelector::ImmediateZeroOrOne => { + "GeneratedSparseMmaSelector::ImmediateZeroOrOne" + } + SparseMmaSelector::ImmediateZero => "GeneratedSparseMmaSelector::ImmediateZero", + }; + return format!( + "GeneratedIntrinsicVariant::SparseMma {{ shape: {shape}, accumulator: {accumulator}, a_element: {}, b_element: {}, a_layout: {}, b_layout: {}, overflow: {overflow}, metadata: {metadata}, selector: {selector} }}", + element(mma.a_element), + element(mma.b_element), + layout(mma.a_layout), + layout(mma.b_layout), + ); + } + if let Some(arithmetic) = &record.scalar_arithmetic { + let format = match arithmetic.format { + ScalarArithmeticFormat::F32 => "GeneratedScalarArithmeticFormat::F32", + ScalarArithmeticFormat::F64 => "GeneratedScalarArithmeticFormat::F64", + }; + let operation = match arithmetic.operation { + ScalarArithmeticOperation::Mul => "GeneratedScalarArithmeticOperation::Mul", + ScalarArithmeticOperation::Div => "GeneratedScalarArithmeticOperation::Div", + ScalarArithmeticOperation::Fma => "GeneratedScalarArithmeticOperation::Fma", + ScalarArithmeticOperation::Add => "GeneratedScalarArithmeticOperation::Add", + }; + let rounding = match arithmetic.rounding { + ScalarArithmeticRounding::Rn => "GeneratedScalarArithmeticRounding::Rn", + ScalarArithmeticRounding::Rz => "GeneratedScalarArithmeticRounding::Rz", + ScalarArithmeticRounding::Rm => "GeneratedScalarArithmeticRounding::Rm", + ScalarArithmeticRounding::Rp => "GeneratedScalarArithmeticRounding::Rp", + }; + let subnormal = match arithmetic.subnormal { + ScalarArithmeticSubnormal::Preserve => "GeneratedScalarArithmeticSubnormal::Preserve", + ScalarArithmeticSubnormal::Ftz => "GeneratedScalarArithmeticSubnormal::Ftz", + }; + let saturation = match arithmetic.saturation { + ScalarArithmeticSaturation::None => "GeneratedScalarArithmeticSaturation::None", + ScalarArithmeticSaturation::Sat => "GeneratedScalarArithmeticSaturation::Sat", + }; + return format!( + "GeneratedIntrinsicVariant::ScalarArithmetic {{ format: {format}, operation: {operation}, rounding: {rounding}, subnormal: {subnormal}, saturation: {saturation} }}" + ); + } + if let Some(minmax) = &record.extended_minmax { + let format = match minmax.format { + ExtendedMinMaxFormat::F32 => "GeneratedExtendedMinMaxFormat::F32", + ExtendedMinMaxFormat::F16 => "GeneratedExtendedMinMaxFormat::F16", + ExtendedMinMaxFormat::Bf16 => "GeneratedExtendedMinMaxFormat::Bf16", + ExtendedMinMaxFormat::F16x2 => "GeneratedExtendedMinMaxFormat::F16x2", + ExtendedMinMaxFormat::Bf16x2 => "GeneratedExtendedMinMaxFormat::Bf16x2", + }; + let operation = match minmax.operation { + ExtendedMinMaxOperation::Min => "GeneratedExtendedMinMaxOperation::Min", + ExtendedMinMaxOperation::Max => "GeneratedExtendedMinMaxOperation::Max", + }; + let subnormal = match minmax.subnormal { + ExtendedMinMaxSubnormal::Preserve => "GeneratedExtendedMinMaxSubnormal::Preserve", + ExtendedMinMaxSubnormal::Ftz => "GeneratedExtendedMinMaxSubnormal::Ftz", + }; + let nan = match minmax.nan { + ExtendedMinMaxNan::Number => "GeneratedExtendedMinMaxNan::Number", + ExtendedMinMaxNan::Nan => "GeneratedExtendedMinMaxNan::Nan", + }; + return format!( + "GeneratedIntrinsicVariant::ExtendedMinMax {{ format: {format}, operation: {operation}, subnormal: {subnormal}, nan: {nan}, xorsign_abs: {} }}", + minmax.xorsign_abs + ); + } + if let Some(conversion) = &record.scalar_conversion { + let rounding = match conversion.rounding { + ScalarConversionRounding::NearestAway => { + "GeneratedScalarConversionRounding::NearestAway" + } + ScalarConversionRounding::NearestEven => { + "GeneratedScalarConversionRounding::NearestEven" + } + ScalarConversionRounding::TowardZero => "GeneratedScalarConversionRounding::TowardZero", + }; + let saturation = match conversion.saturation { + ScalarConversionSaturation::None => "GeneratedScalarConversionSaturation::None", + ScalarConversionSaturation::Relu => "GeneratedScalarConversionSaturation::Relu", + ScalarConversionSaturation::Satfinite => { + "GeneratedScalarConversionSaturation::Satfinite" + } + ScalarConversionSaturation::ReluSatfinite => { + "GeneratedScalarConversionSaturation::ReluSatfinite" + } + }; + return format!( + "GeneratedIntrinsicVariant::ScalarConversion {{ rounding: {rounding}, saturation: {saturation} }}" + ); + } + let Some(ldmatrix) = &record.ldmatrix else { + return "GeneratedIntrinsicVariant::Scalar".to_owned(); + }; + let variant = &ldmatrix.variant; + let shape = match variant.shape { + LdmatrixShape::M8n8 => "GeneratedLdmatrixShape::M8n8", + LdmatrixShape::M8n16 => "GeneratedLdmatrixShape::M8n16", + LdmatrixShape::M16n16 => "GeneratedLdmatrixShape::M16n16", + }; + let multiplicity = match variant.multiplicity { + LdmatrixMultiplicity::X1 => "GeneratedLdmatrixMultiplicity::X1", + LdmatrixMultiplicity::X2 => "GeneratedLdmatrixMultiplicity::X2", + LdmatrixMultiplicity::X4 => "GeneratedLdmatrixMultiplicity::X4", + }; + let layout = match variant.layout { + LdmatrixLayout::Normal => "GeneratedLdmatrixLayout::Normal", + LdmatrixLayout::Transposed => "GeneratedLdmatrixLayout::Transposed", + }; + let element = match variant.element { + LdmatrixElement::B16 => "GeneratedLdmatrixElement::B16", + LdmatrixElement::B8 => "GeneratedLdmatrixElement::B8", + LdmatrixElement::B8x16B4x16P64 => "GeneratedLdmatrixElement::B8x16B4x16P64", + LdmatrixElement::B8x16B6x16P32 => "GeneratedLdmatrixElement::B8x16B6x16P32", + }; + format!( + "GeneratedIntrinsicVariant::Ldmatrix {{ shape: {shape}, multiplicity: {multiplicity}, layout: {layout}, element: {element} }}" + ) +} + +fn generated_backend_requirements(record: &CatalogIntrinsic) -> String { + let requirements = record + .backend_lowerings + .iter() + .map(|lowering| { + let backend = match lowering.backend { + IntrinsicBackend::LlvmNvptx => "GeneratedIntrinsicBackend::LlvmNvptx", + IntrinsicBackend::LibNvvm => "GeneratedIntrinsicBackend::LibNvvm", + }; + format!( + "GeneratedBackendRequirement {{ backend: {backend}, requirement: GeneratedTargetRequirement {{ minimum_ptx: GeneratedPtxVersion::from_encoded({}), hardware: {} }} }}", + lowering.target.minimum_ptx.encoded(), + generated_hardware_target(&lowering.target.hardware), + ) + }) + .collect::>() + .join(", "); + format!("&[{requirements}]") +} + +fn replace_exact_render_fragment(output: &mut String, fragment: &str, replacement: &str) { + assert_eq!( + output.matches(fragment).count(), + 1, + "render fragment must occur exactly once" + ); + let start = output.find(fragment).expect("checked render fragment"); + output.replace_range(start..start + fragment.len(), replacement); +} + +pub(super) fn render_targets(catalog: &CatalogFile, hash: &str) -> String { + let mut records = catalog.intrinsics.iter().collect::>(); + records.sort_by(|left, right| left.rust.abi_id.cmp(&right.rust.abi_id)); + let mut output = rust_header(catalog, hash); + output.push_str( + "//! Generated target requirements and separately imported LLVM/selection facts.\n\npub const GENERATED_INTRINSIC_MARKER_ATTR: &str = \"cuda_oxide_intrinsic_marker\";\n\n#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]\npub struct GeneratedPtxVersion(u16);\nimpl GeneratedPtxVersion {\n pub const fn from_encoded(encoded: u16) -> Self { Self(encoded) }\n pub const fn encoded(self) -> u16 { self.0 }\n pub const fn major(self) -> u16 { self.0 / 10 }\n pub const fn minor(self) -> u16 { self.0 % 10 }\n}\n\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedHardwareAlternative { MinimumSm(u16), ExactArchitecture(u16), FamilyTarget(u16) }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedTargetSelectorBinding { pub name: &'static str, pub value: &'static str }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedTargetAlternative { pub minimum_ptx: GeneratedPtxVersion, pub hardware: GeneratedHardwareAlternative }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedHardwareTarget { All, AnyOf(&'static [GeneratedHardwareAlternative]), TargetMatrix { selectors: &'static [GeneratedTargetSelectorBinding], alternatives: &'static [GeneratedTargetAlternative] } }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedTargetRequirement { pub minimum_ptx: GeneratedPtxVersion, pub hardware: GeneratedHardwareTarget }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedIntrinsicBackend { LlvmNvptx, LibNvvm }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedBackendRequirement { pub backend: GeneratedIntrinsicBackend, pub requirement: GeneratedTargetRequirement }\n\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedSelectionAddressSpace { Generic, Shared }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedImmediateBinding { pub argument_index: usize, pub value: i64 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedSelectionConstraints { pub address_space: Option, pub immediate_bindings: &'static [GeneratedImmediateBinding] }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedSelectionAlternative { pub source_record: &'static str, pub asm: &'static str, pub predicates: &'static [&'static str], pub constraints: GeneratedSelectionConstraints }\n\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedIntrinsicRange { pub lower: &'static str, pub upper_exclusive: &'static str }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedLlvmFacts { pub properties: &'static [&'static str], pub result_no_undef: bool, pub result_range: Option }\n\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedLdmatrixShape { M8n8 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedLdmatrixMultiplicity { X1, X2, X4 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedLdmatrixLayout { Normal, Transposed }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedPackedAtomicFormat { F16x2, Bf16x2 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedRegisterMmaShape { M8n8k4, M16n8k8, M16n8k16, M16n8k32 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedRegisterMmaAccumulator { F32, F64, S32 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedRegisterMmaElement { Bf16, F16, Tf32, F64, S8, U8 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedRegisterMmaLayout { Row, Col }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedRegisterMmaOverflow { NotApplicable, Wrapping, Satfinite }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedIntrinsicVariant {\n Scalar,\n Ldmatrix { shape: GeneratedLdmatrixShape, multiplicity: GeneratedLdmatrixMultiplicity, layout: GeneratedLdmatrixLayout },\n PackedAtomic { format: GeneratedPackedAtomicFormat },\n RegisterMma { shape: GeneratedRegisterMmaShape, accumulator: GeneratedRegisterMmaAccumulator, a_element: GeneratedRegisterMmaElement, b_element: GeneratedRegisterMmaElement, a_layout: GeneratedRegisterMmaLayout, b_layout: GeneratedRegisterMmaLayout, overflow: GeneratedRegisterMmaOverflow },\n}\n\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedIntrinsicTarget {\n pub marker: &'static str,\n pub id: &'static str,\n pub abi_id: &'static str,\n pub dialect_op: &'static str,\n pub variant: GeneratedIntrinsicVariant,\n pub requirement: GeneratedTargetRequirement,\n pub backend_requirements: &'static [GeneratedBackendRequirement],\n pub selections: &'static [GeneratedSelectionAlternative],\n pub llvm: Option,\n}\n\nimpl GeneratedIntrinsicTarget {\n pub fn requirement_for_backend(&self, backend: GeneratedIntrinsicBackend) -> GeneratedTargetRequirement {\n self.backend_requirements.iter().find(|entry| entry.backend == backend).map(|entry| entry.requirement).unwrap_or(self.requirement)\n }\n}\n\npub const GENERATED_INTRINSIC_TARGETS: &[GeneratedIntrinsicTarget] = &[\n", + ); + replace_exact_render_fragment( + &mut output, + "#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedTargetAlternative { pub minimum_ptx: GeneratedPtxVersion, pub hardware: GeneratedHardwareAlternative }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedHardwareTarget { All, AnyOf(&'static [GeneratedHardwareAlternative]), TargetMatrix { selectors: &'static [GeneratedTargetSelectorBinding], alternatives: &'static [GeneratedTargetAlternative] } }", + "#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedTargetAlternative { pub minimum_ptx: GeneratedPtxVersion, pub hardware: GeneratedHardwareAlternative }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedTargetContract { pub selectors: &'static [GeneratedTargetSelectorBinding], pub alternatives: &'static [GeneratedTargetAlternative] }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedHardwareTarget { All, AnyOf(&'static [GeneratedHardwareAlternative]), TargetMatrix { contracts: &'static [GeneratedTargetContract] } }", + ); + replace_exact_render_fragment( + &mut output, + "impl GeneratedIntrinsicTarget {\n pub fn requirement_for_backend(&self, backend: GeneratedIntrinsicBackend) -> GeneratedTargetRequirement {\n self.backend_requirements.iter().find(|entry| entry.backend == backend).map(|entry| entry.requirement).unwrap_or(self.requirement)\n }\n}", + r#"impl GeneratedHardwareTarget { + pub fn contract_for_selector( + self, + name: &str, + value: &str, + ) -> Option<&'static GeneratedTargetContract> { + let Self::TargetMatrix { contracts } = self else { return None; }; + let mut matching = contracts.iter().filter(|contract| { + matches!(contract.selectors, [binding] if binding.name == name && binding.value == value) + }); + let contract = matching.next()?; + if matching.next().is_some() { None } else { Some(contract) } + } +} + +impl GeneratedIntrinsicTarget { + pub fn requirement_for_backend(&self, backend: GeneratedIntrinsicBackend) -> GeneratedTargetRequirement { + self.backend_requirements.iter().find(|entry| entry.backend == backend).map(|entry| entry.requirement).unwrap_or(self.requirement) + } + + pub fn target_contract_for_backend_selector( + &self, + backend: GeneratedIntrinsicBackend, + name: &str, + value: &str, + ) -> Option<&'static GeneratedTargetContract> { + self.requirement_for_backend(backend).hardware.contract_for_selector(name, value) + } +}"#, + ); + replace_exact_render_fragment( + &mut output, + "GeneratedLdmatrixShape { M8n8 }", + "GeneratedLdmatrixShape { M8n8, M8n16, M16n16 }", + ); + replace_exact_render_fragment( + &mut output, + "pub enum GeneratedPackedAtomicFormat { F16x2, Bf16x2 }", + "pub enum GeneratedLdmatrixElement { B16, B8, B8x16B4x16P64, B8x16B6x16P32 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedPackedAtomicFormat { F16x2, Bf16x2 }", + ); + replace_exact_render_fragment( + &mut output, + "Ldmatrix { shape: GeneratedLdmatrixShape, multiplicity: GeneratedLdmatrixMultiplicity, layout: GeneratedLdmatrixLayout }", + "Ldmatrix { shape: GeneratedLdmatrixShape, multiplicity: GeneratedLdmatrixMultiplicity, layout: GeneratedLdmatrixLayout, element: GeneratedLdmatrixElement }", + ); + replace_exact_render_fragment( + &mut output, + "GeneratedRegisterMmaShape { M8n8k4, M16n8k8, M16n8k16, M16n8k32 }", + "GeneratedRegisterMmaShape { M8n8k4, M8n8k16, M8n8k32, M8n8k128, M16n8k4, M16n8k8, M16n8k16, M16n8k32, M16n8k64, M16n8k128, M16n8k256 }", + ); + replace_exact_render_fragment( + &mut output, + "GeneratedRegisterMmaElement { Bf16, F16, Tf32, F64, S8, U8 }", + "GeneratedRegisterMmaElement { Bf16, E2m1, E2m3, E3m2, E4m3, E5m2, F16, Tf32, F64, B1, S4, U4, S8, U8 }", + ); + replace_exact_render_fragment( + &mut output, + "pub enum GeneratedRegisterMmaAccumulator { F32, F64, S32 }", + "pub enum GeneratedRegisterMmaOperation { Multiply, AndPopc, XorPopc }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedRegisterMmaKind { Standard, F8f6f4, Mxf8f6f4 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedRegisterMmaAccumulator { F16, F32, F64, S32 }", + ); + replace_exact_render_fragment( + &mut output, + "RegisterMma { shape: GeneratedRegisterMmaShape, accumulator: GeneratedRegisterMmaAccumulator,", + "RegisterMma { shape: GeneratedRegisterMmaShape, operation: GeneratedRegisterMmaOperation, kind: GeneratedRegisterMmaKind, accumulator: GeneratedRegisterMmaAccumulator,", + ); + replace_exact_render_fragment( + &mut output, + "#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedIntrinsicVariant {", + "#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedSparseMmaShape { M16n8k32, M16n8k64, M16n8k128 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedSparseMmaAccumulator { F16, F32, S32 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedSparseMmaElement { E2m1, E2m3, E3m2, E4m3, E5m2, S4, U4, S8, U8 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedSparseMmaLayout { Row, Col }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedSparseMmaOverflow { NotApplicable, Wrapping, Satfinite }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedSparseMmaMetadata { Standard, Ordered }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedSparseMmaSelector { ImmediateZeroOrOne, ImmediateZero }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedPrmtMode { Generic, F4e, B4e, Rc8, Ecl, Ecr, Rc16 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedClusterBarrierMode { Arrive, ArriveAligned, ArriveRelaxed, ArriveRelaxedAligned, Wait, WaitAligned }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedScalarConversionRounding { NearestAway, NearestEven, TowardZero }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedScalarConversionSaturation { None, Relu, Satfinite, ReluSatfinite }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedIntrinsicVariant {", + ); + replace_exact_render_fragment( + &mut output, + " RegisterMma { shape: GeneratedRegisterMmaShape, operation: GeneratedRegisterMmaOperation, kind: GeneratedRegisterMmaKind, accumulator: GeneratedRegisterMmaAccumulator, a_element: GeneratedRegisterMmaElement, b_element: GeneratedRegisterMmaElement, a_layout: GeneratedRegisterMmaLayout, b_layout: GeneratedRegisterMmaLayout, overflow: GeneratedRegisterMmaOverflow },\n}", + " RegisterMma { shape: GeneratedRegisterMmaShape, operation: GeneratedRegisterMmaOperation, kind: GeneratedRegisterMmaKind, accumulator: GeneratedRegisterMmaAccumulator, a_element: GeneratedRegisterMmaElement, b_element: GeneratedRegisterMmaElement, a_layout: GeneratedRegisterMmaLayout, b_layout: GeneratedRegisterMmaLayout, overflow: GeneratedRegisterMmaOverflow },\n SparseMma { shape: GeneratedSparseMmaShape, accumulator: GeneratedSparseMmaAccumulator, a_element: GeneratedSparseMmaElement, b_element: GeneratedSparseMmaElement, a_layout: GeneratedSparseMmaLayout, b_layout: GeneratedSparseMmaLayout, overflow: GeneratedSparseMmaOverflow, metadata: GeneratedSparseMmaMetadata, selector: GeneratedSparseMmaSelector },\n Prmt { mode: GeneratedPrmtMode },\n ClusterBarrier { mode: GeneratedClusterBarrierMode },\n ScalarConversion { rounding: GeneratedScalarConversionRounding, saturation: GeneratedScalarConversionSaturation },\n}", + ); + if wgmma_controls(catalog).next().is_some() { + replace_exact_render_fragment( + &mut output, + "#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedIntrinsicVariant {", + "#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedWgmmaControlMode { Fence, CommitGroup, WaitGroup }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedIntrinsicVariant {", + ); + replace_exact_render_fragment( + &mut output, + " ClusterBarrier { mode: GeneratedClusterBarrierMode },\n ScalarConversion { rounding: GeneratedScalarConversionRounding, saturation: GeneratedScalarConversionSaturation },\n}", + " ClusterBarrier { mode: GeneratedClusterBarrierMode },\n ScalarConversion { rounding: GeneratedScalarConversionRounding, saturation: GeneratedScalarConversionSaturation },\n WgmmaControl { mode: GeneratedWgmmaControlMode },\n}", + ); + } + if scalar_arithmetics(catalog).next().is_some() { + replace_exact_render_fragment( + &mut output, + "#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedIntrinsicVariant {", + "#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedScalarArithmeticFormat { F32, F64 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedScalarArithmeticOperation { Mul, Div, Fma, Add }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedScalarArithmeticRounding { Rn, Rz, Rm, Rp }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedScalarArithmeticSubnormal { Preserve, Ftz }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedScalarArithmeticSaturation { None, Sat }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedIntrinsicVariant {", + ); + if wgmma_controls(catalog).next().is_some() { + replace_exact_render_fragment( + &mut output, + " WgmmaControl { mode: GeneratedWgmmaControlMode },\n}", + " WgmmaControl { mode: GeneratedWgmmaControlMode },\n ScalarArithmetic { format: GeneratedScalarArithmeticFormat, operation: GeneratedScalarArithmeticOperation, rounding: GeneratedScalarArithmeticRounding, subnormal: GeneratedScalarArithmeticSubnormal, saturation: GeneratedScalarArithmeticSaturation },\n}", + ); + } else { + replace_exact_render_fragment( + &mut output, + " ScalarConversion { rounding: GeneratedScalarConversionRounding, saturation: GeneratedScalarConversionSaturation },\n}", + " ScalarConversion { rounding: GeneratedScalarConversionRounding, saturation: GeneratedScalarConversionSaturation },\n ScalarArithmetic { format: GeneratedScalarArithmeticFormat, operation: GeneratedScalarArithmeticOperation, rounding: GeneratedScalarArithmeticRounding, subnormal: GeneratedScalarArithmeticSubnormal, saturation: GeneratedScalarArithmeticSaturation },\n}", + ); + } + } + if extended_minmax(catalog).next().is_some() { + replace_exact_render_fragment( + &mut output, + "#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedIntrinsicVariant {", + "#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedExtendedMinMaxFormat { F32, F16, Bf16, F16x2, Bf16x2 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedExtendedMinMaxOperation { Min, Max }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedExtendedMinMaxSubnormal { Preserve, Ftz }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedExtendedMinMaxNan { Number, Nan }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedIntrinsicVariant {", + ); + replace_exact_render_fragment( + &mut output, + "}\n\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedIntrinsicTarget", + " ExtendedMinMax { format: GeneratedExtendedMinMaxFormat, operation: GeneratedExtendedMinMaxOperation, subnormal: GeneratedExtendedMinMaxSubnormal, nan: GeneratedExtendedMinMaxNan, xorsign_abs: bool },\n}\n\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedIntrinsicTarget", + ); + } + if tcgen05_mma_intrinsics(catalog).next().is_some() { + replace_exact_render_fragment( + &mut output, + "#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedIntrinsicVariant {", + "#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedTcgen05MmaForm { Shared, Tensor, TensorAshift, SpShared, SpTensor, SpTensorAshift, WsShared, WsSharedZeroColMask, WsSpShared, WsSpSharedZeroColMask, WsSpTensor, WsSpTensorZeroColMask, WsTensor, WsTensorZeroColMask }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedTcgen05MmaTargetSelector { Kind }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedIntrinsicVariant {", + ); + replace_exact_render_fragment( + &mut output, + "}\n\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedIntrinsicTarget", + " Tcgen05Mma { form: GeneratedTcgen05MmaForm, target_selector: GeneratedTcgen05MmaTargetSelector, compatibility_alias: bool },\n}\n\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub struct GeneratedIntrinsicTarget", + ); + } + for record in records.iter().copied() { + let llvm_facts = match &record.llvm { + Some(llvm) => { + let result_range = match &llvm.result_facts.range { + Some(range) => format!( + "Some(GeneratedIntrinsicRange {{ lower: {:?}, upper_exclusive: {:?} }})", + range.lower, range.upper_exclusive + ), + None => "None".to_owned(), + }; + format!( + "Some(GeneratedLlvmFacts {{ properties: &{:?}, result_no_undef: {}, result_range: {} }})", + llvm.properties, llvm.result_facts.no_undef, result_range + ) + } + None => "None".to_owned(), + }; + writeln!( + output, + " GeneratedIntrinsicTarget {{ marker: {:?}, id: {:?}, abi_id: {:?}, dialect_op: {:?}, variant: {}, requirement: GeneratedTargetRequirement {{ minimum_ptx: GeneratedPtxVersion::from_encoded({}), hardware: {} }}, backend_requirements: {}, selections: {}, llvm: {} }},", + intrinsic_marker(catalog, record), + record.id, + record.rust.abi_id, + record.dialect.op_name, + generated_intrinsic_variant(record), + record.target.minimum_ptx.encoded(), + generated_hardware_target(&record.target.hardware), + generated_backend_requirements(record), + generated_selection_alternatives(&record.selections), + llvm_facts, + ) + .unwrap(); + } + output.push_str( + "];\n\npub fn generated_intrinsic_target_by_marker(marker: &str) -> Option<&'static GeneratedIntrinsicTarget> {\n GENERATED_INTRINSIC_TARGETS.iter().find(|target| target.marker == marker)\n}\n\npub fn generated_intrinsic_targets_by_op_name(op_name: &str) -> impl Iterator + '_ {\n GENERATED_INTRINSIC_TARGETS.iter().filter(move |target| target.dialect_op == op_name)\n}\n\npub fn generated_intrinsic_target_by_op_name(op_name: &str) -> Option<&'static GeneratedIntrinsicTarget> {\n generated_intrinsic_targets_by_op_name(op_name).next()\n}\n\npub fn generated_intrinsic_target(op_name: &str) -> Option<&'static GeneratedIntrinsicTarget> {\n generated_intrinsic_target_by_op_name(op_name)\n}\n\npub fn generated_intrinsic_operation_matches(ctx: &Context, target: &GeneratedIntrinsicTarget, operation: Ptr) -> bool {\n match target.variant {\n GeneratedIntrinsicVariant::Scalar => true,\n GeneratedIntrinsicVariant::Ldmatrix { shape, multiplicity, layout } => {\n let Some(op) = Operation::get_op::(operation, ctx) else { return false; };\n let shape_matches = matches!(shape, GeneratedLdmatrixShape::M8n8) && op.get_attr_nvvm_ldmatrix_shape(ctx).as_deref() == Some(&LdmatrixShapeAttr::M8n8);\n let multiplicity_matches = match multiplicity {\n GeneratedLdmatrixMultiplicity::X1 => op.get_attr_nvvm_ldmatrix_multiplicity(ctx).as_deref() == Some(&LdmatrixMultiplicityAttr::X1),\n GeneratedLdmatrixMultiplicity::X2 => op.get_attr_nvvm_ldmatrix_multiplicity(ctx).as_deref() == Some(&LdmatrixMultiplicityAttr::X2),\n GeneratedLdmatrixMultiplicity::X4 => op.get_attr_nvvm_ldmatrix_multiplicity(ctx).as_deref() == Some(&LdmatrixMultiplicityAttr::X4),\n };\n let layout_matches = match layout {\n GeneratedLdmatrixLayout::Normal => op.get_attr_nvvm_ldmatrix_layout(ctx).as_deref() == Some(&LdmatrixLayoutAttr::Normal),\n GeneratedLdmatrixLayout::Transposed => op.get_attr_nvvm_ldmatrix_layout(ctx).as_deref() == Some(&LdmatrixLayoutAttr::Transposed),\n };\n shape_matches && multiplicity_matches && layout_matches\n && op.get_attr_nvvm_ldmatrix_element(ctx).as_deref() == Some(&LdmatrixElementAttr::B16)\n && op.get_attr_nvvm_ldmatrix_state_space(ctx).as_deref() == Some(&LdmatrixStateSpaceAttr::Shared)\n }\n GeneratedIntrinsicVariant::PackedAtomic { format } => {\n let Some(op) = Operation::get_op::(operation, ctx) else { return false; };\n let format_matches = match format {\n GeneratedPackedAtomicFormat::F16x2 => op.get_attr_nvvm_packed_atomic_format(ctx).as_deref() == Some(&PackedAtomicFormatAttr::F16x2),\n GeneratedPackedAtomicFormat::Bf16x2 => op.get_attr_nvvm_packed_atomic_format(ctx).as_deref() == Some(&PackedAtomicFormatAttr::Bf16x2),\n };\n format_matches\n && op.get_attr_nvvm_packed_atomic_state_space(ctx).as_deref() == Some(&PackedAtomicStateSpaceAttr::Global)\n && op.get_attr_nvvm_packed_atomic_ordering(ctx).as_deref() == Some(&PackedAtomicOrderingAttr::Relaxed)\n && op.get_attr_nvvm_packed_atomic_scope(ctx).as_deref() == Some(&PackedAtomicScopeAttr::Gpu)\n && op.get_attr_nvvm_packed_atomic_rounding(ctx).as_deref() == Some(&PackedAtomicRoundingAttr::Rn)\n && op.get_attr_nvvm_packed_atomic_subnormal(ctx).as_deref() == Some(&PackedAtomicSubnormalAttr::NoFtz)\n && op.get_attr_nvvm_packed_atomic_atomicity(ctx).as_deref() == Some(&PackedAtomicAtomicityAttr::PerElement)\n }\n GeneratedIntrinsicVariant::RegisterMma { shape, accumulator, a_element, b_element, a_layout, b_layout, overflow } => {\n let Some(op) = Operation::get_op::(operation, ctx) else { return false; };\n let shape_matches = match shape {\n GeneratedRegisterMmaShape::M8n8k4 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M8n8k4),\n GeneratedRegisterMmaShape::M16n8k8 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M16n8k8),\n GeneratedRegisterMmaShape::M16n8k16 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M16n8k16),\n GeneratedRegisterMmaShape::M16n8k32 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M16n8k32),\n };\n let accumulator_matches = match accumulator {\n GeneratedRegisterMmaAccumulator::F32 => op.get_attr_nvvm_register_mma_accumulator(ctx).as_deref() == Some(&RegisterMmaAccumulatorAttr::F32),\n GeneratedRegisterMmaAccumulator::F64 => op.get_attr_nvvm_register_mma_accumulator(ctx).as_deref() == Some(&RegisterMmaAccumulatorAttr::F64),\n GeneratedRegisterMmaAccumulator::S32 => op.get_attr_nvvm_register_mma_accumulator(ctx).as_deref() == Some(&RegisterMmaAccumulatorAttr::S32),\n };\n let element_matches = |expected, actual: Option<&RegisterMmaElementAttr>| match expected {\n GeneratedRegisterMmaElement::Bf16 => actual == Some(&RegisterMmaElementAttr::Bf16),\n GeneratedRegisterMmaElement::F16 => actual == Some(&RegisterMmaElementAttr::F16),\n GeneratedRegisterMmaElement::Tf32 => actual == Some(&RegisterMmaElementAttr::Tf32),\n GeneratedRegisterMmaElement::F64 => actual == Some(&RegisterMmaElementAttr::F64),\n GeneratedRegisterMmaElement::S8 => actual == Some(&RegisterMmaElementAttr::S8),\n GeneratedRegisterMmaElement::U8 => actual == Some(&RegisterMmaElementAttr::U8),\n };\n let layout_matches = |expected, actual: Option<&RegisterMmaLayoutAttr>| match expected {\n GeneratedRegisterMmaLayout::Row => actual == Some(&RegisterMmaLayoutAttr::Row),\n GeneratedRegisterMmaLayout::Col => actual == Some(&RegisterMmaLayoutAttr::Col),\n };\n let overflow_matches = match overflow {\n GeneratedRegisterMmaOverflow::NotApplicable => op.get_attr_nvvm_register_mma_overflow(ctx).as_deref() == Some(&RegisterMmaOverflowAttr::NotApplicable),\n GeneratedRegisterMmaOverflow::Wrapping => op.get_attr_nvvm_register_mma_overflow(ctx).as_deref() == Some(&RegisterMmaOverflowAttr::Wrapping),\n GeneratedRegisterMmaOverflow::Satfinite => op.get_attr_nvvm_register_mma_overflow(ctx).as_deref() == Some(&RegisterMmaOverflowAttr::Satfinite),\n };\n shape_matches && accumulator_matches\n && element_matches(a_element, op.get_attr_nvvm_register_mma_a_element(ctx).as_deref())\n && element_matches(b_element, op.get_attr_nvvm_register_mma_b_element(ctx).as_deref())\n && layout_matches(a_layout, op.get_attr_nvvm_register_mma_a_layout(ctx).as_deref())\n && layout_matches(b_layout, op.get_attr_nvvm_register_mma_b_layout(ctx).as_deref())\n && overflow_matches\n }\n }\n}\n", + ); + if tcgen05_mma_intrinsics(catalog).next().is_some() { + replace_exact_render_fragment( + &mut output, + "pub fn generated_intrinsic_target(op_name: &str) -> Option<&'static GeneratedIntrinsicTarget> {\n generated_intrinsic_target_by_op_name(op_name)\n}\n\npub fn generated_intrinsic_operation_matches", + "pub fn generated_intrinsic_target(op_name: &str) -> Option<&'static GeneratedIntrinsicTarget> {\n generated_intrinsic_target_by_op_name(op_name)\n}\n\npub fn generated_intrinsic_target_is_direct_dialect_candidate(target: &GeneratedIntrinsicTarget) -> bool {\n !matches!(target.variant, GeneratedIntrinsicVariant::Tcgen05Mma { compatibility_alias: true, .. })\n}\n\npub fn generated_intrinsic_operation_matches", + ); + replace_exact_render_fragment( + &mut output, + " GeneratedIntrinsicVariant::Scalar => true,", + r#" GeneratedIntrinsicVariant::Scalar => true, + GeneratedIntrinsicVariant::Tcgen05Mma { form, target_selector, compatibility_alias } => { + let Some(op) = Operation::get_op::(operation, ctx) else { return false; }; + let form_matches = match form { + GeneratedTcgen05MmaForm::Shared => op.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() == Some(&Tcgen05MmaFormAttr::Shared), + GeneratedTcgen05MmaForm::Tensor => op.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() == Some(&Tcgen05MmaFormAttr::Tensor), + GeneratedTcgen05MmaForm::TensorAshift => op.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() == Some(&Tcgen05MmaFormAttr::TensorAshift), + GeneratedTcgen05MmaForm::SpShared => op.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() == Some(&Tcgen05MmaFormAttr::SpShared), + GeneratedTcgen05MmaForm::SpTensor => op.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() == Some(&Tcgen05MmaFormAttr::SpTensor), + GeneratedTcgen05MmaForm::SpTensorAshift => op.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() == Some(&Tcgen05MmaFormAttr::SpTensorAshift), + GeneratedTcgen05MmaForm::WsShared => op.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() == Some(&Tcgen05MmaFormAttr::WsShared), + GeneratedTcgen05MmaForm::WsSharedZeroColMask => op.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() == Some(&Tcgen05MmaFormAttr::WsSharedZeroColMask), + GeneratedTcgen05MmaForm::WsSpShared => op.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() == Some(&Tcgen05MmaFormAttr::WsSpShared), + GeneratedTcgen05MmaForm::WsSpSharedZeroColMask => op.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() == Some(&Tcgen05MmaFormAttr::WsSpSharedZeroColMask), + GeneratedTcgen05MmaForm::WsSpTensor => op.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() == Some(&Tcgen05MmaFormAttr::WsSpTensor), + GeneratedTcgen05MmaForm::WsSpTensorZeroColMask => op.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() == Some(&Tcgen05MmaFormAttr::WsSpTensorZeroColMask), + GeneratedTcgen05MmaForm::WsTensor => op.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() == Some(&Tcgen05MmaFormAttr::WsTensor), + GeneratedTcgen05MmaForm::WsTensorZeroColMask => op.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() == Some(&Tcgen05MmaFormAttr::WsTensorZeroColMask), + }; + let selector_matches = matches!(target_selector, GeneratedTcgen05MmaTargetSelector::Kind) + && op.get_attr_nvvm_tcgen05_mma_kind(ctx).is_some(); + let alias_matches = !compatibility_alias || ( + op.get_attr_nvvm_tcgen05_mma_kind(ctx).as_deref() == Some(&Tcgen05MmaKindAttr::F8f6f4) + && match form { + GeneratedTcgen05MmaForm::Shared => { + op.get_attr_nvvm_tcgen05_mma_cta_group(ctx).as_deref() == Some(&Tcgen05MmaCtaGroupAttr::Cg1) + && op.get_attr_nvvm_tcgen05_mma_collector_a(ctx).as_deref() == Some(&Tcgen05MmaCollectorAAttr::Discard) + } + GeneratedTcgen05MmaForm::WsTensor => { + op.get_attr_nvvm_tcgen05_mma_b_buffer(ctx).as_deref() == Some(&Tcgen05MmaBBufferAttr::B0) + && op.get_attr_nvvm_tcgen05_mma_b_usage(ctx).as_deref() == Some(&Tcgen05MmaBUsageAttr::Discard) + } + _ => false, + } + ); + form_matches && selector_matches && alias_matches + }"#, + ); + } + replace_exact_render_fragment( + &mut output, + "GeneratedIntrinsicVariant::Ldmatrix { shape, multiplicity, layout } => {", + "GeneratedIntrinsicVariant::Ldmatrix { shape, multiplicity, layout, element } => {", + ); + replace_exact_render_fragment( + &mut output, + " let shape_matches = matches!(shape, GeneratedLdmatrixShape::M8n8) && op.get_attr_nvvm_ldmatrix_shape(ctx).as_deref() == Some(&LdmatrixShapeAttr::M8n8);", + " let shape_matches = match shape {\n GeneratedLdmatrixShape::M8n8 => op.get_attr_nvvm_ldmatrix_shape(ctx).as_deref() == Some(&LdmatrixShapeAttr::M8n8),\n GeneratedLdmatrixShape::M8n16 => op.get_attr_nvvm_ldmatrix_shape(ctx).as_deref() == Some(&LdmatrixShapeAttr::M8n16),\n GeneratedLdmatrixShape::M16n16 => op.get_attr_nvvm_ldmatrix_shape(ctx).as_deref() == Some(&LdmatrixShapeAttr::M16n16),\n };", + ); + replace_exact_render_fragment( + &mut output, + " shape_matches && multiplicity_matches && layout_matches\n && op.get_attr_nvvm_ldmatrix_element(ctx).as_deref() == Some(&LdmatrixElementAttr::B16)", + " let element_matches = match element {\n GeneratedLdmatrixElement::B16 => op.get_attr_nvvm_ldmatrix_element(ctx).as_deref() == Some(&LdmatrixElementAttr::B16),\n GeneratedLdmatrixElement::B8 => op.get_attr_nvvm_ldmatrix_element(ctx).as_deref() == Some(&LdmatrixElementAttr::B8),\n GeneratedLdmatrixElement::B8x16B4x16P64 => op.get_attr_nvvm_ldmatrix_element(ctx).as_deref() == Some(&LdmatrixElementAttr::B8x16B4x16P64),\n GeneratedLdmatrixElement::B8x16B6x16P32 => op.get_attr_nvvm_ldmatrix_element(ctx).as_deref() == Some(&LdmatrixElementAttr::B8x16B6x16P32),\n };\n shape_matches && multiplicity_matches && layout_matches && element_matches", + ); + let mut compatibility_aliases = String::from( + "fn generated_intrinsic_compatibility_marker_by_op_name(op_name: &str) -> Option<&'static str> {\n match op_name {\n", + ); + let mut compatibility_count = 0; + for record in ldmatrix(catalog) { + let Some((_, op_name)) = ldmatrix_compat_op(record) else { + continue; + }; + compatibility_count += 1; + writeln!( + compatibility_aliases, + " {op_name:?} => Some({:?}),", + intrinsic_marker(catalog, record) + ) + .unwrap(); + } + assert_eq!(compatibility_count, 6); + compatibility_aliases.push_str(" _ => None,\n }\n}\n\n"); + replace_exact_render_fragment( + &mut output, + "pub fn generated_intrinsic_target_by_marker(marker: &str)", + &format!( + "{compatibility_aliases}pub fn generated_intrinsic_target_by_marker(marker: &str)" + ), + ); + replace_exact_render_fragment( + &mut output, + " GENERATED_INTRINSIC_TARGETS.iter().filter(move |target| target.dialect_op == op_name)\n", + " let compatibility_marker = generated_intrinsic_compatibility_marker_by_op_name(op_name);\n GENERATED_INTRINSIC_TARGETS.iter().filter(move |target| {\n target.dialect_op == op_name || compatibility_marker == Some(target.marker)\n })\n", + ); + replace_exact_render_fragment( + &mut output, + " GeneratedIntrinsicVariant::Ldmatrix { shape, multiplicity, layout, element } => {\n let Some(op) = Operation::get_op::(operation, ctx) else { return false; };", + " GeneratedIntrinsicVariant::Ldmatrix { shape, multiplicity, layout, element } => {\n let operation_name = Operation::get_opid(operation, ctx).to_string();\n if let Some(marker) = generated_intrinsic_compatibility_marker_by_op_name(&operation_name) {\n return marker == target.marker;\n }\n let Some(op) = Operation::get_op::(operation, ctx) else { return false; };", + ); + replace_exact_render_fragment( + &mut output, + "GeneratedIntrinsicVariant::RegisterMma { shape, accumulator,", + "GeneratedIntrinsicVariant::RegisterMma { shape, operation: mma_operation, kind, accumulator,", + ); + replace_exact_render_fragment( + &mut output, + " };\n let accumulator_matches = match accumulator {", + " };\n let operation_matches = match mma_operation {\n GeneratedRegisterMmaOperation::Multiply => op.operation_or_multiply(ctx) == RegisterMmaOperationAttr::Multiply,\n GeneratedRegisterMmaOperation::AndPopc => op.operation_or_multiply(ctx) == RegisterMmaOperationAttr::AndPopc,\n GeneratedRegisterMmaOperation::XorPopc => op.operation_or_multiply(ctx) == RegisterMmaOperationAttr::XorPopc,\n };\n let kind_matches = match kind {\n GeneratedRegisterMmaKind::Standard => op.kind_or_inferred(ctx) == RegisterMmaKindAttr::Standard,\n GeneratedRegisterMmaKind::F8f6f4 => op.kind_or_inferred(ctx) == RegisterMmaKindAttr::F8f6f4,\n GeneratedRegisterMmaKind::Mxf8f6f4 => op.kind_or_inferred(ctx) == RegisterMmaKindAttr::Mxf8f6f4,\n };\n let accumulator_matches = match accumulator {", + ); + replace_exact_render_fragment( + &mut output, + " let accumulator_matches = match accumulator {\n GeneratedRegisterMmaAccumulator::F32", + " let accumulator_matches = match accumulator {\n GeneratedRegisterMmaAccumulator::F16 => op.get_attr_nvvm_register_mma_accumulator(ctx).as_deref() == Some(&RegisterMmaAccumulatorAttr::F16),\n GeneratedRegisterMmaAccumulator::F32", + ); + replace_exact_render_fragment( + &mut output, + " shape_matches && accumulator_matches", + " shape_matches && operation_matches && kind_matches && accumulator_matches", + ); + replace_exact_render_fragment( + &mut output, + "GeneratedRegisterMmaShape::M8n8k4 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M8n8k4),\n GeneratedRegisterMmaShape::M16n8k8", + "GeneratedRegisterMmaShape::M8n8k4 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M8n8k4),\n GeneratedRegisterMmaShape::M8n8k16 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M8n8k16),\n GeneratedRegisterMmaShape::M8n8k32 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M8n8k32),\n GeneratedRegisterMmaShape::M8n8k128 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M8n8k128),\n GeneratedRegisterMmaShape::M16n8k4 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M16n8k4),\n GeneratedRegisterMmaShape::M16n8k8", + ); + replace_exact_render_fragment( + &mut output, + "GeneratedRegisterMmaShape::M16n8k32 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M16n8k32),\n };", + "GeneratedRegisterMmaShape::M16n8k32 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M16n8k32),\n GeneratedRegisterMmaShape::M16n8k64 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M16n8k64),\n GeneratedRegisterMmaShape::M16n8k128 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M16n8k128),\n GeneratedRegisterMmaShape::M16n8k256 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M16n8k256),\n };", + ); + replace_exact_render_fragment( + &mut output, + "GeneratedRegisterMmaElement::F64 => actual == Some(&RegisterMmaElementAttr::F64),\n GeneratedRegisterMmaElement::S8", + "GeneratedRegisterMmaElement::F64 => actual == Some(&RegisterMmaElementAttr::F64),\n GeneratedRegisterMmaElement::E2m1 => actual == Some(&RegisterMmaElementAttr::E2m1),\n GeneratedRegisterMmaElement::E2m3 => actual == Some(&RegisterMmaElementAttr::E2m3),\n GeneratedRegisterMmaElement::E3m2 => actual == Some(&RegisterMmaElementAttr::E3m2),\n GeneratedRegisterMmaElement::E4m3 => actual == Some(&RegisterMmaElementAttr::E4m3),\n GeneratedRegisterMmaElement::E5m2 => actual == Some(&RegisterMmaElementAttr::E5m2),\n GeneratedRegisterMmaElement::B1 => actual == Some(&RegisterMmaElementAttr::B1),\n GeneratedRegisterMmaElement::S4 => actual == Some(&RegisterMmaElementAttr::S4),\n GeneratedRegisterMmaElement::U4 => actual == Some(&RegisterMmaElementAttr::U4),\n GeneratedRegisterMmaElement::S8", + ); + replace_exact_render_fragment( + &mut output, + " && overflow_matches\n }\n }\n}\n", + " && overflow_matches\n }\n GeneratedIntrinsicVariant::SparseMma { shape, accumulator, a_element, b_element, a_layout, b_layout, overflow, metadata, selector } => {\n let Some(op) = Operation::get_op::(operation, ctx) else { return false; };\n let element_matches = |expected, actual: Option<&SparseMmaElementAttr>| match expected {\n GeneratedSparseMmaElement::E2m1 => actual == Some(&SparseMmaElementAttr::E2m1),\n GeneratedSparseMmaElement::E2m3 => actual == Some(&SparseMmaElementAttr::E2m3),\n GeneratedSparseMmaElement::E3m2 => actual == Some(&SparseMmaElementAttr::E3m2),\n GeneratedSparseMmaElement::E4m3 => actual == Some(&SparseMmaElementAttr::E4m3),\n GeneratedSparseMmaElement::E5m2 => actual == Some(&SparseMmaElementAttr::E5m2),\n GeneratedSparseMmaElement::S4 => actual == Some(&SparseMmaElementAttr::S4),\n GeneratedSparseMmaElement::U4 => actual == Some(&SparseMmaElementAttr::U4),\n GeneratedSparseMmaElement::S8 => actual == Some(&SparseMmaElementAttr::S8),\n GeneratedSparseMmaElement::U8 => actual == Some(&SparseMmaElementAttr::U8),\n };\n let layout_matches = |expected, actual: Option<&SparseMmaLayoutAttr>| match expected {\n GeneratedSparseMmaLayout::Row => actual == Some(&SparseMmaLayoutAttr::Row),\n GeneratedSparseMmaLayout::Col => actual == Some(&SparseMmaLayoutAttr::Col),\n };\n let overflow_matches = match overflow {\n GeneratedSparseMmaOverflow::NotApplicable => op.get_attr_nvvm_sparse_mma_overflow(ctx).as_deref() == Some(&SparseMmaOverflowAttr::NotApplicable),\n GeneratedSparseMmaOverflow::Wrapping => op.get_attr_nvvm_sparse_mma_overflow(ctx).as_deref() == Some(&SparseMmaOverflowAttr::Wrapping),\n GeneratedSparseMmaOverflow::Satfinite => op.get_attr_nvvm_sparse_mma_overflow(ctx).as_deref() == Some(&SparseMmaOverflowAttr::Satfinite),\n };\n let metadata_matches = match metadata {\n GeneratedSparseMmaMetadata::Standard => op.get_attr_nvvm_sparse_mma_metadata(ctx).as_deref() == Some(&SparseMmaMetadataAttr::Standard),\n GeneratedSparseMmaMetadata::Ordered => op.get_attr_nvvm_sparse_mma_metadata(ctx).as_deref() == Some(&SparseMmaMetadataAttr::Ordered),\n };\n matches!(shape, GeneratedSparseMmaShape::M16n8k32)\n && op.get_attr_nvvm_sparse_mma_shape(ctx).as_deref() == Some(&SparseMmaShapeAttr::M16n8k32)\n && matches!(accumulator, GeneratedSparseMmaAccumulator::S32)\n && op.get_attr_nvvm_sparse_mma_accumulator(ctx).as_deref() == Some(&SparseMmaAccumulatorAttr::S32)\n && element_matches(a_element, op.get_attr_nvvm_sparse_mma_a_element(ctx).as_deref())\n && element_matches(b_element, op.get_attr_nvvm_sparse_mma_b_element(ctx).as_deref())\n && layout_matches(a_layout, op.get_attr_nvvm_sparse_mma_a_layout(ctx).as_deref())\n && layout_matches(b_layout, op.get_attr_nvvm_sparse_mma_b_layout(ctx).as_deref())\n && overflow_matches\n && metadata_matches\n && matches!(selector, GeneratedSparseMmaSelector::ImmediateZeroOrOne)\n && op.get_attr_nvvm_sparse_mma_selector(ctx).as_deref() == Some(&SparseMmaSelectorAttr::ImmediateZeroOrOne)\n }\n }\n}\n", + ); + replace_exact_render_fragment( + &mut output, + " matches!(shape, GeneratedSparseMmaShape::M16n8k32)\n && op.get_attr_nvvm_sparse_mma_shape(ctx).as_deref() == Some(&SparseMmaShapeAttr::M16n8k32)\n && matches!(accumulator, GeneratedSparseMmaAccumulator::S32)\n && op.get_attr_nvvm_sparse_mma_accumulator(ctx).as_deref() == Some(&SparseMmaAccumulatorAttr::S32)", + " let shape_matches = match shape {\n GeneratedSparseMmaShape::M16n8k32 => op.get_attr_nvvm_sparse_mma_shape(ctx).as_deref() == Some(&SparseMmaShapeAttr::M16n8k32),\n GeneratedSparseMmaShape::M16n8k64 => op.get_attr_nvvm_sparse_mma_shape(ctx).as_deref() == Some(&SparseMmaShapeAttr::M16n8k64),\n GeneratedSparseMmaShape::M16n8k128 => op.get_attr_nvvm_sparse_mma_shape(ctx).as_deref() == Some(&SparseMmaShapeAttr::M16n8k128),\n };\n let accumulator_matches = match accumulator {\n GeneratedSparseMmaAccumulator::F16 => op.get_attr_nvvm_sparse_mma_accumulator(ctx).as_deref() == Some(&SparseMmaAccumulatorAttr::F16),\n GeneratedSparseMmaAccumulator::F32 => op.get_attr_nvvm_sparse_mma_accumulator(ctx).as_deref() == Some(&SparseMmaAccumulatorAttr::F32),\n GeneratedSparseMmaAccumulator::S32 => op.get_attr_nvvm_sparse_mma_accumulator(ctx).as_deref() == Some(&SparseMmaAccumulatorAttr::S32),\n };\n let selector_matches = match selector {\n GeneratedSparseMmaSelector::ImmediateZeroOrOne => op.get_attr_nvvm_sparse_mma_selector(ctx).as_deref() == Some(&SparseMmaSelectorAttr::ImmediateZeroOrOne),\n GeneratedSparseMmaSelector::ImmediateZero => op.get_attr_nvvm_sparse_mma_selector(ctx).as_deref() == Some(&SparseMmaSelectorAttr::ImmediateZero),\n };\n shape_matches\n && accumulator_matches", + ); + replace_exact_render_fragment( + &mut output, + " && metadata_matches\n && matches!(selector, GeneratedSparseMmaSelector::ImmediateZeroOrOne)\n && op.get_attr_nvvm_sparse_mma_selector(ctx).as_deref() == Some(&SparseMmaSelectorAttr::ImmediateZeroOrOne)", + " && metadata_matches\n && selector_matches", + ); + replace_exact_render_fragment( + &mut output, + " }\n }\n}\n", + " }\n GeneratedIntrinsicVariant::Prmt { mode } => {\n let Some(op) = Operation::get_op::(operation, ctx) else { return false; };\n match mode {\n GeneratedPrmtMode::Generic => op.get_attr_nvvm_prmt_mode(ctx).as_deref() == Some(&PrmtModeAttr::Generic),\n GeneratedPrmtMode::F4e => op.get_attr_nvvm_prmt_mode(ctx).as_deref() == Some(&PrmtModeAttr::F4e),\n GeneratedPrmtMode::B4e => op.get_attr_nvvm_prmt_mode(ctx).as_deref() == Some(&PrmtModeAttr::B4e),\n GeneratedPrmtMode::Rc8 => op.get_attr_nvvm_prmt_mode(ctx).as_deref() == Some(&PrmtModeAttr::Rc8),\n GeneratedPrmtMode::Ecl => op.get_attr_nvvm_prmt_mode(ctx).as_deref() == Some(&PrmtModeAttr::Ecl),\n GeneratedPrmtMode::Ecr => op.get_attr_nvvm_prmt_mode(ctx).as_deref() == Some(&PrmtModeAttr::Ecr),\n GeneratedPrmtMode::Rc16 => op.get_attr_nvvm_prmt_mode(ctx).as_deref() == Some(&PrmtModeAttr::Rc16),\n }\n }\n }\n}\n", + ); + replace_exact_render_fragment( + &mut output, + " }\n }\n}\n", + " }\n GeneratedIntrinsicVariant::ClusterBarrier { mode } => {\n let Some(op) = Operation::get_op::(operation, ctx) else { return false; };\n match mode {\n GeneratedClusterBarrierMode::Arrive => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::Arrive),\n GeneratedClusterBarrierMode::ArriveAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveAligned),\n GeneratedClusterBarrierMode::ArriveRelaxed => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveRelaxed),\n GeneratedClusterBarrierMode::ArriveRelaxedAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveRelaxedAligned),\n GeneratedClusterBarrierMode::Wait => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::Wait),\n GeneratedClusterBarrierMode::WaitAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::WaitAligned),\n }\n }\n }\n}\n", + ); + if wgmma_controls(catalog).next().is_some() { + replace_exact_render_fragment( + &mut output, + " GeneratedIntrinsicVariant::ClusterBarrier { mode } => {\n let Some(op) = Operation::get_op::(operation, ctx) else { return false; };\n match mode {\n GeneratedClusterBarrierMode::Arrive => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::Arrive),\n GeneratedClusterBarrierMode::ArriveAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveAligned),\n GeneratedClusterBarrierMode::ArriveRelaxed => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveRelaxed),\n GeneratedClusterBarrierMode::ArriveRelaxedAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveRelaxedAligned),\n GeneratedClusterBarrierMode::Wait => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::Wait),\n GeneratedClusterBarrierMode::WaitAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::WaitAligned),\n }\n }\n }\n}\n", + " GeneratedIntrinsicVariant::ClusterBarrier { mode } => {\n let Some(op) = Operation::get_op::(operation, ctx) else { return false; };\n match mode {\n GeneratedClusterBarrierMode::Arrive => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::Arrive),\n GeneratedClusterBarrierMode::ArriveAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveAligned),\n GeneratedClusterBarrierMode::ArriveRelaxed => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveRelaxed),\n GeneratedClusterBarrierMode::ArriveRelaxedAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveRelaxedAligned),\n GeneratedClusterBarrierMode::Wait => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::Wait),\n GeneratedClusterBarrierMode::WaitAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::WaitAligned),\n }\n }\n GeneratedIntrinsicVariant::WgmmaControl { mode } => match mode {\n GeneratedWgmmaControlMode::Fence => Operation::get_op::(operation, ctx).is_some(),\n GeneratedWgmmaControlMode::CommitGroup => Operation::get_op::(operation, ctx).is_some(),\n GeneratedWgmmaControlMode::WaitGroup => Operation::get_op::(operation, ctx).is_some(),\n },\n }\n}\n", + ); + } + if scalar_conversions(catalog).next().is_some() { + let scalar_match = " GeneratedIntrinsicVariant::ScalarConversion { rounding, saturation } => {\n let Some(op) = Operation::get_op::(operation, ctx) else { return false; };\n let rounding_matches = match rounding {\n GeneratedScalarConversionRounding::NearestAway => op.get_attr_nvvm_scalar_conversion_rounding(ctx).as_deref() == Some(&ScalarConversionRoundingAttr::NearestAway),\n GeneratedScalarConversionRounding::NearestEven => op.get_attr_nvvm_scalar_conversion_rounding(ctx).as_deref() == Some(&ScalarConversionRoundingAttr::NearestEven),\n GeneratedScalarConversionRounding::TowardZero => op.get_attr_nvvm_scalar_conversion_rounding(ctx).as_deref() == Some(&ScalarConversionRoundingAttr::TowardZero),\n };\n let saturation_matches = match saturation {\n GeneratedScalarConversionSaturation::None => op.get_attr_nvvm_scalar_conversion_saturation(ctx).as_deref() == Some(&ScalarConversionSaturationAttr::None),\n GeneratedScalarConversionSaturation::Relu => op.get_attr_nvvm_scalar_conversion_saturation(ctx).as_deref() == Some(&ScalarConversionSaturationAttr::Relu),\n GeneratedScalarConversionSaturation::Satfinite => op.get_attr_nvvm_scalar_conversion_saturation(ctx).as_deref() == Some(&ScalarConversionSaturationAttr::Satfinite),\n GeneratedScalarConversionSaturation::ReluSatfinite => op.get_attr_nvvm_scalar_conversion_saturation(ctx).as_deref() == Some(&ScalarConversionSaturationAttr::ReluSatfinite),\n };\n rounding_matches && saturation_matches\n }\n"; + if wgmma_controls(catalog).next().is_some() { + replace_exact_render_fragment( + &mut output, + " GeneratedIntrinsicVariant::WgmmaControl { mode } => match mode {\n GeneratedWgmmaControlMode::Fence => Operation::get_op::(operation, ctx).is_some(),\n GeneratedWgmmaControlMode::CommitGroup => Operation::get_op::(operation, ctx).is_some(),\n GeneratedWgmmaControlMode::WaitGroup => Operation::get_op::(operation, ctx).is_some(),\n },\n }\n}\n", + &format!( + " GeneratedIntrinsicVariant::WgmmaControl {{ mode }} => match mode {{\n GeneratedWgmmaControlMode::Fence => Operation::get_op::(operation, ctx).is_some(),\n GeneratedWgmmaControlMode::CommitGroup => Operation::get_op::(operation, ctx).is_some(),\n GeneratedWgmmaControlMode::WaitGroup => Operation::get_op::(operation, ctx).is_some(),\n }},\n{scalar_match} }}\n}}\n" + ), + ); + } else { + replace_exact_render_fragment( + &mut output, + " GeneratedIntrinsicVariant::ClusterBarrier { mode } => {\n let Some(op) = Operation::get_op::(operation, ctx) else { return false; };\n match mode {\n GeneratedClusterBarrierMode::Arrive => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::Arrive),\n GeneratedClusterBarrierMode::ArriveAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveAligned),\n GeneratedClusterBarrierMode::ArriveRelaxed => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveRelaxed),\n GeneratedClusterBarrierMode::ArriveRelaxedAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveRelaxedAligned),\n GeneratedClusterBarrierMode::Wait => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::Wait),\n GeneratedClusterBarrierMode::WaitAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::WaitAligned),\n }\n }\n }\n}\n", + &format!( + " GeneratedIntrinsicVariant::ClusterBarrier {{ mode }} => {{\n let Some(op) = Operation::get_op::(operation, ctx) else {{ return false; }};\n match mode {{\n GeneratedClusterBarrierMode::Arrive => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::Arrive),\n GeneratedClusterBarrierMode::ArriveAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveAligned),\n GeneratedClusterBarrierMode::ArriveRelaxed => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveRelaxed),\n GeneratedClusterBarrierMode::ArriveRelaxedAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::ArriveRelaxedAligned),\n GeneratedClusterBarrierMode::Wait => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::Wait),\n GeneratedClusterBarrierMode::WaitAligned => op.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() == Some(&ClusterBarrierModeAttr::WaitAligned),\n }}\n }}\n{scalar_match} }}\n}}\n" + ), + ); + } + } + if scalar_arithmetics(catalog).next().is_some() { + let arithmetic_match = r#" GeneratedIntrinsicVariant::ScalarArithmetic { format, operation: arithmetic_operation, rounding, subnormal, saturation } => { + let Some(op) = Operation::get_op::(operation, ctx) else { return false; }; + let format_matches = match format { + GeneratedScalarArithmeticFormat::F32 => op.get_attr_nvvm_scalar_arithmetic_format(ctx).as_deref() == Some(&ScalarArithmeticFormatAttr::F32), + GeneratedScalarArithmeticFormat::F64 => op.get_attr_nvvm_scalar_arithmetic_format(ctx).as_deref() == Some(&ScalarArithmeticFormatAttr::F64), + }; + let operation_matches = match arithmetic_operation { + GeneratedScalarArithmeticOperation::Mul => op.get_attr_nvvm_scalar_arithmetic_operation(ctx).as_deref() == Some(&ScalarArithmeticOperationAttr::Mul), + GeneratedScalarArithmeticOperation::Div => op.get_attr_nvvm_scalar_arithmetic_operation(ctx).as_deref() == Some(&ScalarArithmeticOperationAttr::Div), + GeneratedScalarArithmeticOperation::Fma => op.get_attr_nvvm_scalar_arithmetic_operation(ctx).as_deref() == Some(&ScalarArithmeticOperationAttr::Fma), + GeneratedScalarArithmeticOperation::Add => op.get_attr_nvvm_scalar_arithmetic_operation(ctx).as_deref() == Some(&ScalarArithmeticOperationAttr::Add), + }; + let rounding_matches = match rounding { + GeneratedScalarArithmeticRounding::Rn => op.get_attr_nvvm_scalar_arithmetic_rounding(ctx).as_deref() == Some(&ScalarArithmeticRoundingAttr::Rn), + GeneratedScalarArithmeticRounding::Rz => op.get_attr_nvvm_scalar_arithmetic_rounding(ctx).as_deref() == Some(&ScalarArithmeticRoundingAttr::Rz), + GeneratedScalarArithmeticRounding::Rm => op.get_attr_nvvm_scalar_arithmetic_rounding(ctx).as_deref() == Some(&ScalarArithmeticRoundingAttr::Rm), + GeneratedScalarArithmeticRounding::Rp => op.get_attr_nvvm_scalar_arithmetic_rounding(ctx).as_deref() == Some(&ScalarArithmeticRoundingAttr::Rp), + }; + let subnormal_matches = match subnormal { + GeneratedScalarArithmeticSubnormal::Preserve => op.get_attr_nvvm_scalar_arithmetic_subnormal(ctx).as_deref() == Some(&ScalarArithmeticSubnormalAttr::Preserve), + GeneratedScalarArithmeticSubnormal::Ftz => op.get_attr_nvvm_scalar_arithmetic_subnormal(ctx).as_deref() == Some(&ScalarArithmeticSubnormalAttr::Ftz), + }; + let saturation_matches = match saturation { + GeneratedScalarArithmeticSaturation::None => op.get_attr_nvvm_scalar_arithmetic_saturation(ctx).as_deref() == Some(&ScalarArithmeticSaturationAttr::None), + GeneratedScalarArithmeticSaturation::Sat => op.get_attr_nvvm_scalar_arithmetic_saturation(ctx).as_deref() == Some(&ScalarArithmeticSaturationAttr::Sat), + }; + format_matches && operation_matches && rounding_matches && subnormal_matches && saturation_matches + } +"#; + replace_exact_render_fragment( + &mut output, + " rounding_matches && saturation_matches\n }\n }\n}\n", + &format!( + " rounding_matches && saturation_matches\n }}\n{arithmetic_match} }}\n}}\n" + ), + ); + } + if extended_minmax(catalog).next().is_some() { + let minmax_match = r#" GeneratedIntrinsicVariant::ExtendedMinMax { format, operation: minmax_operation, subnormal, nan, xorsign_abs } => { + let Some(op) = Operation::get_op::(operation, ctx) else { return false; }; + let format_matches = match format { + GeneratedExtendedMinMaxFormat::F32 => op.get_attr_nvvm_extended_minmax_format(ctx).as_deref() == Some(&ExtendedMinMaxFormatAttr::F32), + GeneratedExtendedMinMaxFormat::F16 => op.get_attr_nvvm_extended_minmax_format(ctx).as_deref() == Some(&ExtendedMinMaxFormatAttr::F16), + GeneratedExtendedMinMaxFormat::Bf16 => op.get_attr_nvvm_extended_minmax_format(ctx).as_deref() == Some(&ExtendedMinMaxFormatAttr::Bf16), + GeneratedExtendedMinMaxFormat::F16x2 => op.get_attr_nvvm_extended_minmax_format(ctx).as_deref() == Some(&ExtendedMinMaxFormatAttr::F16x2), + GeneratedExtendedMinMaxFormat::Bf16x2 => op.get_attr_nvvm_extended_minmax_format(ctx).as_deref() == Some(&ExtendedMinMaxFormatAttr::Bf16x2), + }; + let operation_matches = match minmax_operation { + GeneratedExtendedMinMaxOperation::Min => op.get_attr_nvvm_extended_minmax_operation(ctx).as_deref() == Some(&ExtendedMinMaxOperationAttr::Min), + GeneratedExtendedMinMaxOperation::Max => op.get_attr_nvvm_extended_minmax_operation(ctx).as_deref() == Some(&ExtendedMinMaxOperationAttr::Max), + }; + let subnormal_matches = match subnormal { + GeneratedExtendedMinMaxSubnormal::Preserve => op.get_attr_nvvm_extended_minmax_subnormal(ctx).as_deref() == Some(&ExtendedMinMaxSubnormalAttr::Preserve), + GeneratedExtendedMinMaxSubnormal::Ftz => op.get_attr_nvvm_extended_minmax_subnormal(ctx).as_deref() == Some(&ExtendedMinMaxSubnormalAttr::Ftz), + }; + let nan_matches = match nan { + GeneratedExtendedMinMaxNan::Number => op.get_attr_nvvm_extended_minmax_nan(ctx).as_deref() == Some(&ExtendedMinMaxNanAttr::Number), + GeneratedExtendedMinMaxNan::Nan => op.get_attr_nvvm_extended_minmax_nan(ctx).as_deref() == Some(&ExtendedMinMaxNanAttr::Nan), + }; + let xorsign_abs_matches = match xorsign_abs { + false => op.get_attr_nvvm_extended_minmax_xorsign_abs(ctx).as_deref() == Some(&ExtendedMinMaxXorSignAbsAttr::Disabled), + true => op.get_attr_nvvm_extended_minmax_xorsign_abs(ctx).as_deref() == Some(&ExtendedMinMaxXorSignAbsAttr::Enabled), + }; + format_matches && operation_matches && subnormal_matches && nan_matches && xorsign_abs_matches + } +"#; + let match_end = output + .rfind(" }\n}\n") + .expect("generated target matcher terminator"); + output.insert_str(match_end, minmax_match); + } + output.push_str( + "\nuse dialect_nvvm::ops::{ClusterBarrierModeAttr, ClusterBarrierOp, LdmatrixElementAttr, LdmatrixLayoutAttr, LdmatrixMultiplicityAttr, LdmatrixOp, LdmatrixShapeAttr, LdmatrixStateSpaceAttr, PackedAtomicAddOp, PackedAtomicAtomicityAttr, PackedAtomicFormatAttr, PackedAtomicOrderingAttr, PackedAtomicRoundingAttr, PackedAtomicScopeAttr, PackedAtomicStateSpaceAttr, PackedAtomicSubnormalAttr, PrmtModeAttr, PrmtOp, RegisterMmaAccumulatorAttr, RegisterMmaElementAttr, RegisterMmaKindAttr, RegisterMmaLayoutAttr, RegisterMmaOp, RegisterMmaOperationAttr, RegisterMmaOverflowAttr, RegisterMmaShapeAttr, SparseMmaAccumulatorAttr, SparseMmaElementAttr, SparseMmaLayoutAttr, SparseMmaMetadataAttr, SparseMmaOp, SparseMmaOverflowAttr, SparseMmaSelectorAttr, SparseMmaShapeAttr};\nuse pliron::{context::{Context, Ptr}, operation::Operation};\n", + ); + if wgmma_controls(catalog).next().is_some() { + replace_exact_render_fragment( + &mut output, + "SparseMmaSelectorAttr, SparseMmaShapeAttr};", + "SparseMmaSelectorAttr, SparseMmaShapeAttr, WgmmaCommitGroupSyncAlignedOp, WgmmaFenceSyncAlignedOp, WgmmaWaitGroupSyncAlignedOp};", + ); + } + if scalar_conversions(catalog).next().is_some() { + replace_exact_render_fragment( + &mut output, + "SparseMmaSelectorAttr, SparseMmaShapeAttr", + "ScalarConversionOp, ScalarConversionRoundingAttr, ScalarConversionSaturationAttr, SparseMmaSelectorAttr, SparseMmaShapeAttr", + ); + } + if scalar_arithmetics(catalog).next().is_some() { + replace_exact_render_fragment( + &mut output, + "SparseMmaSelectorAttr, SparseMmaShapeAttr", + "ScalarArithmeticFormatAttr, ScalarArithmeticOp, ScalarArithmeticOperationAttr, ScalarArithmeticRoundingAttr, ScalarArithmeticSaturationAttr, ScalarArithmeticSubnormalAttr, SparseMmaSelectorAttr, SparseMmaShapeAttr", + ); + } + if extended_minmax(catalog).next().is_some() { + replace_exact_render_fragment( + &mut output, + "SparseMmaSelectorAttr, SparseMmaShapeAttr", + "ExtendedMinMaxFormatAttr, ExtendedMinMaxNanAttr, ExtendedMinMaxOp, ExtendedMinMaxOperationAttr, ExtendedMinMaxSubnormalAttr, ExtendedMinMaxXorSignAbsAttr, SparseMmaSelectorAttr, SparseMmaShapeAttr", + ); + } + if tcgen05_mma_intrinsics(catalog).next().is_some() { + replace_exact_render_fragment( + &mut output, + "SparseMmaSelectorAttr, SparseMmaShapeAttr", + "SparseMmaSelectorAttr, SparseMmaShapeAttr, Tcgen05MmaBBufferAttr, Tcgen05MmaBUsageAttr, Tcgen05MmaCollectorAAttr, Tcgen05MmaCtaGroupAttr, Tcgen05MmaFormAttr, Tcgen05MmaKindAttr, Tcgen05MmaOp", + ); + } + output.push_str( + "\n#[cfg(test)]\nmod tests {\n use super::*;\n use std::collections::BTreeSet;\n\n #[test]\n fn generated_target_table_is_unique_and_lookup_is_complete() {\n let mut ids = BTreeSet::new();\n let mut markers = BTreeSet::new();\n let mut previous_abi_id = None;\n for target in GENERATED_INTRINSIC_TARGETS {\n if let Some(previous) = previous_abi_id {\n assert!(previous < target.abi_id, \"generated ABI IDs are not strictly increasing: {previous} then {}\", target.abi_id);\n }\n previous_abi_id = Some(target.abi_id);\n assert!(ids.insert(target.id), \"duplicate generated intrinsic ID {}\", target.id);\n assert!(markers.insert(target.marker), \"duplicate generated marker {}\", target.marker);\n assert_eq!(generated_intrinsic_target_by_marker(target.marker), Some(target));\n assert!(generated_intrinsic_targets_by_op_name(target.dialect_op).any(|candidate| candidate == target));\n }\n", + ); + for record in records { + writeln!( + output, + " let target = generated_intrinsic_target_by_marker({:?}).unwrap();", + intrinsic_marker(catalog, record) + ) + .unwrap(); + writeln!( + output, + " assert_eq!(target.id, {:?});\n assert_eq!(target.abi_id, {:?});\n assert_eq!(target.dialect_op, {:?});\n assert_eq!(target.variant, {});\n assert_eq!(target.requirement.minimum_ptx.encoded(), {});\n assert_eq!(target.requirement.hardware, const {{ {} }});\n assert_eq!(target.backend_requirements, const {{ {} }});\n assert_eq!(target.selections, {});", + record.id, + record.rust.abi_id, + record.dialect.op_name, + generated_intrinsic_variant(record), + record.target.minimum_ptx.encoded(), + generated_hardware_target(&record.target.hardware), + generated_backend_requirements(record), + generated_selection_alternatives(&record.selections), + ) + .unwrap(); + match &record.llvm { + Some(llvm) => { + writeln!( + output, + " assert_eq!(target.llvm.unwrap().properties, &{:?} as &[&str]);", + llvm.properties + ) + .unwrap(); + } + None => output.push_str(" assert!(target.llvm.is_none());\n"), + } + } + output.push_str(" assert!(generated_intrinsic_target_by_marker(\"v1:i9999\").is_none());\n }\n}\n"); + output +} diff --git a/crates/cuda-intrinsics-gen/src/render/common.rs b/crates/cuda-intrinsics-gen/src/render/common.rs new file mode 100644 index 0000000000..607da9aba2 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/render/common.rs @@ -0,0 +1,202 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, CatalogFile, CatalogHardwareAlternative, CatalogHardwareTarget, + CatalogIntrinsic, CatalogLlvm, CatalogTargetContract, EvidenceStageKind, IntrinsicBackend, + IntrinsicSource, +}; + +pub(super) fn rust_header(catalog: &CatalogFile, catalog_sha256: &str) -> String { + format!( + "/*\n * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: Apache-2.0\n */\n\n// @generated by cuda-intrinsics-gen {}. DO NOT EDIT.\n// Catalog schema/version: {}/{}; intrinsic ABI: v{}; catalog SHA-256: {}.\n// LLVM source: {}.\n\n", + catalog.generator_version, + catalog.schema, + catalog.catalog_version, + catalog.intrinsic_abi, + catalog_sha256, + catalog.source.llvm_revision + ) +} + +pub(super) fn llvm_header(catalog: &CatalogFile, catalog_sha256: &str) -> String { + format!( + "; SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n; SPDX-License-Identifier: Apache-2.0\n; @generated by cuda-intrinsics-gen {}. DO NOT EDIT.\n; Catalog schema/version: {}/{}; intrinsic ABI: v{}; catalog SHA-256: {}.\n; LLVM source: {}.\n\n", + catalog.generator_version, + catalog.schema, + catalog.catalog_version, + catalog.intrinsic_abi, + catalog_sha256, + catalog.source.llvm_revision + ) +} + +pub(super) fn markdown_header(catalog: &CatalogFile, catalog_sha256: &str) -> String { + format!( + "\n\n", + catalog.generator_version, + catalog.schema, + catalog.catalog_version, + catalog.intrinsic_abi, + catalog_sha256, + catalog.source.llvm_revision + ) +} + +pub(super) fn llvm(record: &CatalogIntrinsic) -> &CatalogLlvm { + record.llvm.as_ref().expect("LLVM-imported intrinsic") +} + +pub(super) fn source_label(record: &CatalogIntrinsic) -> String { + match &record.source { + IntrinsicSource::LlvmImported { source_record } => { + format!("LLVM `{}` from `{source_record}`", llvm(record).symbol) + } + IntrinsicSource::PtxNative { instruction } => { + format!("PTX-native `{instruction}`") + } + } +} + +pub(super) fn intrinsic_marker(catalog: &CatalogFile, record: &CatalogIntrinsic) -> String { + format!("v{}:{}", catalog.intrinsic_abi, record.rust.abi_id) +} + +pub(super) fn hardware_target_label(target: &CatalogHardwareTarget) -> String { + match target { + CatalogHardwareTarget::All => "all".to_owned(), + CatalogHardwareTarget::AnyOf { alternatives } => alternatives + .iter() + .map(hardware_alternative_label) + .collect::>() + .join(" or "), + CatalogHardwareTarget::TargetMatrix { contracts } => contracts + .iter() + .map(|contract| { + let selectors = contract + .selectors + .iter() + .map(|selector| format!("{}={}", selector.name, selector.value)) + .collect::>() + .join(", "); + let alternatives = contract + .alternatives + .iter() + .map(|alternative| { + format!( + "{} at PTX {}", + hardware_alternative_label(&alternative.hardware), + alternative.minimum_ptx + ) + }) + .collect::>() + .join(" or "); + if selectors.is_empty() { + alternatives + } else { + format!("{alternatives} for {selectors}") + } + }) + .collect::>() + .join("; "), + } +} + +fn hardware_alternative_label(alternative: &CatalogHardwareAlternative) -> String { + match alternative { + CatalogHardwareAlternative::MinimumSm { sm } => format!("sm_{sm}+"), + CatalogHardwareAlternative::ExactArchitecture { sm } => format!("sm_{sm}a"), + CatalogHardwareAlternative::FamilyTarget { sm } => format!("sm_{sm}f"), + } +} + +pub(super) fn generated_hardware_target(target: &CatalogHardwareTarget) -> String { + match target { + CatalogHardwareTarget::All => "GeneratedHardwareTarget::All".to_owned(), + CatalogHardwareTarget::AnyOf { alternatives } => { + let alternatives = alternatives + .iter() + .map(generated_hardware_alternative) + .collect::>() + .join(", "); + format!("GeneratedHardwareTarget::AnyOf(&[{alternatives}])") + } + CatalogHardwareTarget::TargetMatrix { contracts } => { + let contracts = contracts + .iter() + .map(generated_target_contract) + .collect::>() + .join(", "); + format!("GeneratedHardwareTarget::TargetMatrix {{ contracts: &[{contracts}] }}") + } + } +} + +fn generated_target_contract(contract: &CatalogTargetContract) -> String { + let selectors = contract + .selectors + .iter() + .map(|selector| { + format!( + "GeneratedTargetSelectorBinding {{ name: {:?}, value: {:?} }}", + selector.name, selector.value + ) + }) + .collect::>() + .join(", "); + let alternatives = contract + .alternatives + .iter() + .map(|alternative| { + format!( + "GeneratedTargetAlternative {{ minimum_ptx: GeneratedPtxVersion::from_encoded({}), hardware: {} }}", + alternative.minimum_ptx.encoded(), + generated_hardware_alternative(&alternative.hardware) + ) + }) + .collect::>() + .join(", "); + format!( + "GeneratedTargetContract {{ selectors: &[{selectors}], alternatives: &[{alternatives}] }}" + ) +} + +fn generated_hardware_alternative(alternative: &CatalogHardwareAlternative) -> String { + match alternative { + CatalogHardwareAlternative::MinimumSm { sm } => { + format!("GeneratedHardwareAlternative::MinimumSm({sm})") + } + CatalogHardwareAlternative::ExactArchitecture { sm } => { + format!("GeneratedHardwareAlternative::ExactArchitecture({sm})") + } + CatalogHardwareAlternative::FamilyTarget { sm } => { + format!("GeneratedHardwareAlternative::FamilyTarget({sm})") + } + } +} + +pub(super) fn backend_label(backend: IntrinsicBackend) -> &'static str { + match backend { + IntrinsicBackend::LlvmNvptx => "llvm_nvptx", + IntrinsicBackend::LibNvvm => "lib_nvvm", + } +} + +pub(super) fn lowering_mechanism_label(mechanism: BackendLoweringMechanism) -> &'static str { + match mechanism { + BackendLoweringMechanism::TypedNvvm => "typed_nvvm", + BackendLoweringMechanism::InlinePtx => "inline_ptx", + } +} + +pub(super) fn evidence_stage_label(stage: EvidenceStageKind) -> &'static str { + match stage { + EvidenceStageKind::DeclarationCanonicalization => "declaration canonicalization", + EvidenceStageKind::BackendCodegen => "backend codegen", + EvidenceStageKind::DeviceLink => "device link", + EvidenceStageKind::PtxAssembly => "PTX assembly", + EvidenceStageKind::Runtime => "runtime", + } +} diff --git a/crates/cuda-intrinsics-gen/src/render/compat.rs b/crates/cuda-intrinsics-gen/src/render/compat.rs new file mode 100644 index 0000000000..979f9445fa --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/render/compat.rs @@ -0,0 +1,2214 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + CatalogFile, CatalogHardwareAlternative, CatalogHardwareTarget, CatalogIntrinsic, ClcOperation, + ClusterBarrierOrdering, ClusterMemoryOperation, CpAsyncCachePolicy, CpAsyncControlOperation, + CpAsyncMbarrierOperation, CpAsyncSourceSize, DebugControlOperation, MbarrierBasicOperation, + MbarrierExtendedOperation, PackedAluFormat, PackedAtomicFormat, PackedConversionSourceFormat, + RegisterMmaAdapter, RegisterMmaCompatibilitySource, RegisterMmaOverflow, ScalarMathFormat, + SparseMmaAdapter, SparseMmaCompatibilitySource, SparseMmaOverflow, Tcgen05LdShape, + Tcgen05MmaForm, Tcgen05MmaSelectorLayout, Tcgen05Operation, TmaAdapter, TmaOperation, + WgmmaControlMode, +}; +use crate::render::common::{hardware_target_label, llvm, rust_header, source_label}; +use crate::render::families::{ + ClcSafetyArgNames, clc_intrinsics, cluster_barriers, cluster_memory, cp_async_controls, + cp_async_copies, cp_async_mbarriers, debug_controls, dot_products, execution_control_family, + expected_ptx_head, extended_minmax, extended_minmax_rust_type, integer_minmaxes, + is_blackwell_ldmatrix, ldmatrix, mbarrier_basics, mbarrier_extended, movmatrix, + packed_alu_format_shape, packed_alus, packed_atomics, packed_conversion_rust_arguments, + packed_conversion_source, packed_conversions, prmts, register_mmas, render_clc_safety_lines, + scalar_arithmetic_arity, scalar_arithmetic_rust_type, scalar_arithmetics, scalar_conversions, + scalar_math_contract, scalar_maths, sparse_mma_fragment_counts, sparse_mma_metadata_rule, + sparse_mma_ptx_head, sparse_mma_selector_description, sparse_mmas, sregs, stmatrices, + stmatrix_compatibility_name, stmatrix_variant, sync_intrinsics, tcgen05_intrinsics, + tcgen05_is_commit, tcgen05_is_multicast_commit, tcgen05_is_shift, tcgen05_ld_register_count, + tcgen05_mma_runtime_parameters, tcgen05_mma_selector_parameters, tcgen05_participation_doc, + tcgen05_st_register_count, threadfence_ptx_level, tma_intrinsics, wgmma_control, + wgmma_controls, +}; +use std::fmt::Write as _; +use std::path::PathBuf; + +pub(super) fn render_compat_register_mma(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str("// Included inside `cuda_device::wmma` to keep public paths stable.\n\n"); + for record in register_mmas(catalog).filter(|record| { + record.register_mma.as_ref().is_some_and(|mma| { + mma.compatibility_source == RegisterMmaCompatibilitySource::GeneratedStub + }) + }) { + let mma = record + .register_mma + .as_ref() + .expect("register-MMA semantics"); + let [path] = record.rust.compatibility_paths.as_slice() else { + panic!("generated register-MMA API requires one compatibility path"); + }; + assert_eq!(path, &format!("cuda_device::wmma::{}", record.rust.name)); + assert!(!record.rust.safe); + assert!(record.rust.must_use); + let argument_names: &[&str] = + if mma.adapter == RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 { + &[ + "c", + "a", + "b", + "scale_a", + "byte_id_a", + "thread_id_a", + "scale_b", + "byte_id_b", + "thread_id_b", + ] + } else { + &["c", "a", "b"] + }; + assert_eq!(record.rust.arguments.len(), argument_names.len()); + + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!( + output, + "/// Lowers to `{}`. C, A, and B are this lane's fragments in PTX register order.", + expected_ptx_head(record) + ) + .unwrap(); + writeln!( + output, + "/// Requires PTX {} and `{}`.", + record.target.minimum_ptx, + hardware_target_label(&record.target.hardware) + ) + .unwrap(); + match mma.overflow { + RegisterMmaOverflow::Wrapping => { + output.push_str("/// Signed accumulator overflow wraps.\n"); + } + RegisterMmaOverflow::Satfinite => { + output.push_str( + "/// Signed accumulator overflow clamps to the finite `i32` range.\n", + ); + } + RegisterMmaOverflow::NotApplicable => {} + } + output.push_str("///\n/// # Safety\n"); + output.push_str( + "/// All 32 lanes must execute the same instruction with the same qualifiers, and no lane may have exited.\n", + ); + output.push_str( + "/// `c`, `a`, and `b` must contain this lane's fragments in the documented PTX layout.\n", + ); + if mma.adapter == RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 { + output.push_str( + "/// `scale_a` and `scale_b` contain this lane's packed scale data.\n\ + /// For `scale_vec::1X`, `byte_id_a` and `byte_id_b` must be in `0..=3`, `thread_id_a` in `0..=1`, and `thread_id_b` in `0..=3`; other values make the PTX operation undefined.\n", + ); + } + writeln!( + output, + "/// See the [PTX MMA fragment layouts]({}).", + record.target.ptx_isa_url + ) + .unwrap(); + output.push_str("/// This register-only operation is not a memory fence.\n"); + if record.rust.arguments.len() > 7 { + output.push_str("#[allow(clippy::too_many_arguments)]\n"); + } + output.push_str("#[must_use]\n#[inline(never)]\n"); + let arguments = argument_names + .iter() + .copied() + .zip(&record.rust.arguments) + .map(|(name, ty)| format!("{name}: {ty}")) + .collect::>() + .join(", "); + writeln!( + output, + "pub unsafe fn {}({arguments}) -> {} {{", + record.rust.name, record.rust.result + ) + .unwrap(); + writeln!(output, " let _ = ({});", argument_names.join(", ")).unwrap(); + writeln!( + output, + " unreachable!(\"{} called outside CUDA kernel context\")", + record.rust.name + ) + .unwrap(); + output.push_str("}\n\n"); + } + output +} + +pub(super) fn render_compat_ldmatrix(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str("// Included inside `cuda_device::wmma` to keep public paths stable.\n\n"); + for record in ldmatrix(catalog).filter(|record| { + is_blackwell_ldmatrix(record) && !record.rust.compatibility_paths.is_empty() + }) { + let [path] = record.rust.compatibility_paths.as_slice() else { + panic!("generated Blackwell ldmatrix API requires one compatibility path"); + }; + assert_eq!(path, &format!("cuda_device::wmma::{}", record.rust.name)); + assert_eq!(record.rust.arguments, ["*const u8"]); + assert!(!record.rust.safe); + assert!(record.rust.must_use); + + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!(output, "/// Lowers to `{}`.", expected_ptx_head(record)).unwrap(); + writeln!( + output, + "/// Requires PTX {} and `{}`.", + record.target.minimum_ptx, + hardware_target_label(&record.target.hardware) + ) + .unwrap(); + output.push_str( + "///\n/// # Safety\n\ + /// All 32 lanes must execute the same instruction and qualifiers, with no exited lanes.\n\ + /// `smem_ptr` must satisfy the lane-address mapping, alignment, and readable-byte contract documented by PTX.\n\ + #[must_use]\n#[inline(never)]\n", + ); + writeln!( + output, + "pub unsafe fn {}(smem_ptr: *const u8) -> {} {{", + record.rust.name, record.rust.result + ) + .unwrap(); + output.push_str(" let _ = smem_ptr;\n"); + writeln!( + output, + " unreachable!(\"{} called outside CUDA kernel context\")", + record.rust.name + ) + .unwrap(); + output.push_str("}\n\n"); + } + output +} + +pub(super) fn render_compat_sparse_mma(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str("// Included inside `cuda_device::wmma` to keep public paths stable.\n\n"); + for record in sparse_mmas(catalog).filter(|record| { + record.sparse_mma.as_ref().is_some_and(|mma| { + mma.compatibility_source == SparseMmaCompatibilitySource::GeneratedStub + }) + }) { + let mma = record.sparse_mma.as_ref().expect("sparse-MMA semantics"); + let [path] = record.rust.compatibility_paths.as_slice() else { + panic!("generated sparse-MMA API requires one compatibility path"); + }; + assert_eq!(path, &format!("cuda_device::wmma::{}", record.rust.name)); + assert!(matches!( + mma.adapter, + SparseMmaAdapter::C2U32A4U32B4U32MetadataU32SelectorU32ToD2U32 + | SparseMmaAdapter::C4F32A4U32B4U32MetadataU32SelectorU32ToD4F32 + | SparseMmaAdapter::C4I32A2U32B2U32MetadataU32SelectorU32ToD4I32 + | SparseMmaAdapter::C4I32A4U32B4U32MetadataU32SelectorU32ToD4I32 + )); + + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!( + output, + "/// Lowers to `{}`. C, A, B, and metadata are this lane's PTX fragments.", + sparse_mma_ptx_head(record) + ) + .unwrap(); + writeln!( + output, + "/// Requires PTX {} and `{}`.", + record.target.minimum_ptx, + hardware_target_label(&record.target.hardware) + ) + .unwrap(); + match mma.overflow { + SparseMmaOverflow::NotApplicable => {} + SparseMmaOverflow::Wrapping => { + output.push_str("/// Signed accumulator overflow wraps.\n"); + } + SparseMmaOverflow::Satfinite => output + .push_str("/// Signed accumulator overflow clamps to the finite `i32` range.\n"), + } + output.push_str("///\n/// # Safety\n"); + output.push_str( + "/// All 32 lanes must execute the same instruction with the same qualifiers, and no lane may have exited.\n", + ); + let (c_count, a_count, b_count, _) = sparse_mma_fragment_counts(record); + writeln!( + output, + "/// `c`, `a`, and `b` must contain this lane's {c_count}-, {a_count}-, and {b_count}-register fragments; `metadata` must contain its sparse metadata." + ) + .unwrap(); + writeln!(output, "/// {}", sparse_mma_metadata_rule(mma)).unwrap(); + writeln!( + output, + "/// `selector` must be {}.", + sparse_mma_selector_description(record) + ) + .unwrap(); + writeln!( + output, + "/// See the [PTX sparse MMA fragment layouts]({}).", + record.target.ptx_isa_url + ) + .unwrap(); + output.push_str("/// This register-only operation is not a memory fence.\n"); + output.push_str("#[must_use]\n#[inline(never)]\n"); + let arguments = ["c", "a", "b", "metadata", "selector"] + .into_iter() + .zip(&record.rust.arguments) + .map(|(name, ty)| format!("{name}: {ty}")) + .collect::>() + .join(", "); + writeln!( + output, + "pub unsafe fn {}({arguments}) -> {} {{", + record.rust.name, record.rust.result + ) + .unwrap(); + output.push_str(" let _ = (c, a, b, metadata, selector);\n"); + writeln!( + output, + " unreachable!(\"{} called outside CUDA kernel context\")", + record.rust.name + ) + .unwrap(); + output.push_str("}\n\n"); + } + output +} + +pub(super) fn render_compat_sreg(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + "// This file is included lexically inside `cuda_device::thread` so existing\n// DefPaths remain stable during the generated-intrinsics migration.\n\n", + ); + for record in sregs(catalog) { + let Some(path) = record + .rust + .compatibility_paths + .iter() + .find(|path| path.starts_with("cuda_device::thread::")) + else { + continue; + }; + let name = path.rsplit("::").next().unwrap(); + writeln!( + output, + "/// Compatibility spelling for `{}`.", + record.rust.public_path + ) + .unwrap(); + writeln!( + output, + "/// The compiler replaces this call with {}.", + source_label(record) + ) + .unwrap(); + output.push_str("#[allow(non_snake_case)]\n#[inline(never)]\n"); + let safety = if record.rust.safe { "" } else { "unsafe " }; + writeln!( + output, + "pub {safety}fn {name}() -> {} {{", + record.rust.result + ) + .unwrap(); + writeln!( + output, + " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" + ) + .unwrap(); + output.push_str("}\n\n"); + } + output +} + +pub(super) fn render_compat_cluster_sreg(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + "// This file is included inside `cuda_device::cluster`.\n// Its private leaves let public helpers compose generated reads.\n\n", + ); + for record in sregs(catalog) { + let Some(path) = record + .rust + .compatibility_paths + .iter() + .find(|path| path.starts_with("cuda_device::cluster::__")) + else { + continue; + }; + let name = path.rsplit("::").next().unwrap(); + writeln!( + output, + "/// Private compatibility spelling for `{}`.", + record.rust.public_path + ) + .unwrap(); + writeln!( + output, + "/// The compiler replaces this call with `{}`.", + llvm(record).symbol + ) + .unwrap(); + output.push_str("#[allow(non_snake_case)]\n#[inline(never)]\n"); + writeln!( + output, + "pub(crate) fn {name}() -> {} {{", + record.rust.result + ) + .unwrap(); + writeln!( + output, + " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" + ) + .unwrap(); + output.push_str("}\n\n"); + } + output +} + +pub(super) fn render_compat_special_register_module( + catalog: &CatalogFile, + hash: &str, + module: &str, +) -> String { + let mut output = rust_header(catalog, hash); + writeln!( + output, + "// Included inside `cuda_device::{module}` to keep existing paths stable.\n" + ) + .unwrap(); + let prefix = format!("cuda_device::{module}::"); + for record in sregs(catalog).filter(|record| record.special_register.is_some()) { + let Some(path) = record + .rust + .compatibility_paths + .iter() + .find(|path| path.starts_with(&prefix)) + else { + continue; + }; + let name = path.rsplit("::").next().unwrap(); + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!(output, "/// Generated from {}.", source_label(record)).unwrap(); + output.push_str("#[inline(never)]\n"); + let safety = if record.rust.safe { "" } else { "unsafe " }; + writeln!( + output, + "pub {safety}fn {name}() -> {} {{", + record.rust.result + ) + .unwrap(); + writeln!( + output, + " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" + ) + .unwrap(); + output.push_str("}\n\n"); + } + output +} + +pub(super) fn render_compat_fence(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + "// This file is included inside `cuda_device::fence` so existing paths stay stable.\n\n", + ); + for record in sync_intrinsics(catalog).filter(|record| threadfence_ptx_level(record).is_some()) + { + let Some(path) = record + .rust + .compatibility_paths + .iter() + .find(|path| path.starts_with("cuda_device::fence::")) + else { + continue; + }; + let name = path.rsplit("::").next().unwrap(); + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!( + output, + "/// The compiler replaces this call with `{}`.", + llvm(record).symbol + ) + .unwrap(); + output.push_str("#[inline(never)]\n"); + writeln!(output, "pub fn {name}() {{").unwrap(); + writeln!( + output, + " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" + ) + .unwrap(); + output.push_str("}\n\n"); + } + output +} + +pub(super) fn render_compat_dotprod(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str("// Included inside `cuda_device::dotprod` to keep existing paths stable.\n\n"); + for record in dot_products(catalog) { + let path = record + .rust + .compatibility_paths + .iter() + .find(|path| path.starts_with("cuda_device::dotprod::")) + .expect("dot-product compatibility path"); + let arguments = record + .rust + .arguments + .iter() + .enumerate() + .map(|(index, ty)| format!("arg{index}: {ty}")) + .collect::>() + .join(", "); + let values = (0..record.rust.arguments.len()) + .map(|index| format!("arg{index}")) + .collect::>() + .join(", "); + writeln!(output, "/// {}", record.summary).unwrap(); + output.push_str("#[inline(never)]\n"); + writeln!( + output, + "pub fn {}({arguments}) -> {} {{", + record.rust.name, record.rust.result + ) + .unwrap(); + if record.rust.arguments.len() == 1 { + writeln!(output, " let _ = {values};").unwrap(); + } else { + writeln!(output, " let _ = ({values});").unwrap(); + } + writeln!( + output, + " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" + ) + .unwrap(); + output.push_str("}\n\n"); + } + output +} + +pub(super) fn render_compat_prmt(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str("// Included inside `cuda_device::prmt` to keep the public paths stable.\n\n"); + for record in prmts(catalog) { + let path = record + .rust + .compatibility_paths + .iter() + .find(|path| path.starts_with("cuda_device::prmt::")) + .expect("prmt compatibility path"); + let arguments = record + .rust + .arguments + .iter() + .enumerate() + .map(|(index, ty)| format!("arg{index}: {ty}")) + .collect::>() + .join(", "); + let values = (0..record.rust.arguments.len()) + .map(|index| format!("arg{index}")) + .collect::>() + .join(", "); + writeln!(output, "/// {}", record.summary).unwrap(); + output.push_str("#[must_use]\n#[inline(never)]\n"); + writeln!(output, "pub fn {}({arguments}) -> u32 {{", record.rust.name).unwrap(); + writeln!(output, " let _ = ({values});").unwrap(); + writeln!( + output, + " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" + ) + .unwrap(); + output.push_str("}\n\n"); + } + output +} + +pub(super) fn render_compat_cluster_barrier(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str("// Included inside `cuda_device::cluster` to keep public paths stable.\n\n"); + for record in cluster_barriers(catalog) { + let barrier = record + .cluster_barrier + .as_ref() + .expect("cluster-barrier contract"); + let path = record + .rust + .compatibility_paths + .iter() + .find(|path| path.starts_with("cuda_device::cluster::")) + .expect("cluster-barrier compatibility path"); + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!(output, "/// Lowers to `{}`.", record.expected_ptx).unwrap(); + output.push_str("///\n/// # Safety\n"); + output.push_str( + "/// Each non-exited cluster thread must arrive exactly once before completion, then execute the matching wait.\n", + ); + if barrier.aligned { + output.push_str( + "/// Every non-exited thread in the warp must execute this aligned instruction with identical control flow.\n", + ); + } + if barrier.ordering == ClusterBarrierOrdering::Relaxed { + output.push_str( + "/// This relaxed arrival does not publish earlier memory accesses; use the required cluster-scope fence first.\n", + ); + } + output.push_str("#[inline(never)]\n"); + writeln!(output, "pub unsafe fn {}() {{", record.rust.name).unwrap(); + writeln!(output, " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")").unwrap(); + output.push_str("}\n\n"); + } + output +} + +pub(super) fn render_compat_cluster_memory(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str("// Included inside `cuda_device::cluster` to keep its public API stable.\n\n"); + for record in cluster_memory(catalog) { + let cluster = record + .cluster_memory + .as_ref() + .expect("cluster-memory contract"); + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!(output, "/// Lowers to `{}`.", record.expected_ptx).unwrap(); + output.push_str("///\n/// # Safety\n"); + match cluster.operation { + ClusterMemoryOperation::MapSharedRank => { + output.push_str( + "/// `local_ptr` must point into CTA shared memory, and `target_rank` must name a rank in the same live cluster.\n\ + /// The result is a cluster-shared pointer in address space 7. Dereferencing it performs a remote DSMEM access; the target CTA must remain live and synchronization must make the access race-free.\n", + ); + output.push_str("#[must_use]\n#[inline(never)]\n"); + output.push_str( + "pub unsafe fn map_shared_rank(local_ptr: *const T, target_rank: u32) -> *const T {\n\ + let _ = (local_ptr, target_rank);\n\ + unreachable!(\"map_shared_rank called outside CUDA kernel context\")\n\ + }\n\n", + ); + output.push_str( + "/// Maps a mutable CTA-shared address to another cluster rank.\n\ + ///\n\ + /// # Safety\n\ + /// The mapping requirements above apply. The target CTA must remain live, and remote writes must be synchronized and race-free.\n\ + #[must_use]\n#[inline(never)]\n\ + pub unsafe fn map_shared_rank_mut(local_ptr: *mut T, target_rank: u32) -> *mut T {\n\ + let _ = (local_ptr, target_rank);\n\ + unreachable!(\"map_shared_rank_mut called outside CUDA kernel context\")\n\ + }\n\n", + ); + } + ClusterMemoryOperation::ReadU32 => { + output.push_str( + "/// `local_ptr` must identify an aligned readable `u32` in CTA shared memory. `target_rank` must name a rank in the same live cluster.\n\ + /// The target CTA must publish the value before this weak cluster load.\n", + ); + output.push_str("#[must_use]\n#[inline(never)]\n"); + output.push_str( + "pub unsafe fn dsmem_read_u32(local_ptr: *const u32, target_rank: u32) -> u32 {\n\ + let _ = (local_ptr, target_rank);\n\ + unreachable!(\"dsmem_read_u32 called outside CUDA kernel context\")\n\ + }\n\n", + ); + } + } + } + output +} + +pub(super) fn render_compat_movmatrix(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str("// Included inside `cuda_device::wmma` to keep its public API stable.\n\n"); + for record in movmatrix(catalog) { + let path = record + .rust + .compatibility_paths + .iter() + .find(|path| path.starts_with("cuda_device::wmma::")) + .expect("movmatrix compatibility path"); + assert_eq!(path, &format!("cuda_device::wmma::{}", record.rust.name)); + writeln!(output, "/// {}", record.summary).unwrap(); + output.push_str( + "///\n\ + /// Each lane supplies two packed b16 values. The warp collectively transposes the 8x8 tile.\n\ + /// No addressable memory is touched, but the result depends on every lane's\n\ + /// input, so the call is modelled as reading and writing inaccessible state:\n\ + /// two calls with equal operands are not interchangeable.\n\ + ///\n\ + /// # Safety\n\ + /// All 32 warp lanes must execute the same call, and no lane may have exited.\n", + ); + writeln!( + output, + "/// Requires PTX {} and `sm_75+`.", + record.target.minimum_ptx + ) + .unwrap(); + output.push_str("#[must_use]\n#[inline(never)]\n"); + writeln!( + output, + "pub unsafe fn {}(value: u32) -> u32 {{", + record.rust.name + ) + .unwrap(); + output.push_str(" let _ = value;\n"); + writeln!( + output, + " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" + ) + .unwrap(); + output.push_str("}\n\n"); + } + output +} + +pub(super) fn render_compat_debug_control(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str("// Included inside `cuda_device::debug` to keep its public API stable.\n\n"); + for operation in [ + DebugControlOperation::Trap, + DebugControlOperation::Breakpoint, + DebugControlOperation::Pmevent, + ] { + let record = debug_controls(catalog) + .find(|record| { + record + .debug_control + .as_ref() + .is_some_and(|debug| debug.operation == operation) + }) + .expect("complete debug-control family"); + match operation { + DebugControlOperation::Trap => { + writeln!(output, "/// {}", record.summary).unwrap(); + output.push_str("#[inline(never)]\npub fn trap() -> ! {\n"); + output.push_str( + " unreachable!(\"trap called outside CUDA kernel context\")\n}\n\n", + ); + } + DebugControlOperation::Breakpoint => { + writeln!(output, "/// {}", record.summary).unwrap(); + output.push_str("#[inline(never)]\npub fn breakpoint() {\n"); + output.push_str( + " unreachable!(\"breakpoint called outside CUDA kernel context\")\n}\n\n", + ); + } + DebugControlOperation::Pmevent => { + output.push_str( + r#"/// Triggers performance monitor event `N`, which must be in `0..=15`. +#[inline(always)] +pub fn prof_trigger() { + __prof_trigger(N); +} + +#[doc(hidden)] +#[inline(never)] +pub(crate) fn __prof_trigger(_event_id: u32) { + unreachable!("prof_trigger called outside CUDA kernel context") +} + +"#, + ); + } + } + } + output +} + +pub(super) fn render_compat_stmatrix(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str("// Included inside `cuda_device::tcgen05` to keep public paths stable.\n\n"); + for record in stmatrices(catalog) { + let (multiplicity, _) = stmatrix_variant(record).expect("stmatrix variant"); + let count = multiplicity.register_count(); + let name = stmatrix_compatibility_name(record); + let path = record + .rust + .compatibility_paths + .first() + .expect("stmatrix compatibility path"); + let arguments = std::iter::once("smem_ptr: *mut u8".to_owned()) + .chain((0..count).map(|index| format!("r{index}: u32"))) + .collect::>() + .join(", "); + let values = std::iter::once("smem_ptr".to_owned()) + .chain((0..count).map(|index| format!("r{index}"))) + .collect::>() + .join(", "); + writeln!(output, "/// {}", record.summary).unwrap(); + output.push_str("///\n/// # Safety\n"); + output.push_str( + "/// All warp lanes must participate. The used row addresses must be valid, shared, and 16-byte aligned.\n", + ); + output.push_str("#[inline(never)]\n"); + writeln!(output, "pub unsafe fn {name}({arguments}) {{").unwrap(); + writeln!(output, " let _ = ({values});").unwrap(); + writeln!( + output, + " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" + ) + .unwrap(); + output.push_str("}\n\n"); + } + output +} + +pub(super) fn render_compat_clc(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str("// Included inside `cuda_device::clc` to keep its public API stable.\n\n"); + for record in clc_intrinsics(catalog) { + let operation = record.clc.as_ref().expect("CLC contract").operation; + writeln!(output, "/// {}", record.summary).unwrap(); + output.push_str("///\n/// # Safety\n"); + render_clc_safety_lines(&mut output, operation, ClcSafetyArgNames::Compatibility); + match operation { + ClcOperation::TryCancel | ClcOperation::TryCancelMulticast => { + output.push_str("#[inline(never)]\n"); + writeln!( + output, + "pub unsafe fn {}(response: *mut u8, mbar: *mut Barrier) {{", + record.rust.name + ) + .unwrap(); + output.push_str(" let _ = (response, mbar);\n"); + } + ClcOperation::QueryIsCanceled + | ClcOperation::QueryGetFirstCtaidX + | ClcOperation::QueryGetFirstCtaidY + | ClcOperation::QueryGetFirstCtaidZ => { + output.push_str("#[inline(never)]\n"); + writeln!( + output, + "pub unsafe fn {}(resp_lo: u64, resp_hi: u64) -> u32 {{", + record.rust.name + ) + .unwrap(); + output.push_str(" let _ = (resp_lo, resp_hi);\n"); + } + } + writeln!( + output, + " unreachable!(\"{} called outside CUDA kernel context\")", + record.rust.name + ) + .unwrap(); + output.push_str("}\n\n"); + } + output +} + +pub(super) fn render_compat_counted_barrier(catalog: &CatalogFile, hash: &str) -> String { + assert_eq!( + execution_control_family(catalog, "counted_barrier").count(), + 4 + ); + let mut output = rust_header(catalog, hash); + output.push_str("// Included inside `cuda_device::barrier` to keep its public API stable.\n\n"); + for record in execution_control_family(catalog, "counted_barrier") { + writeln!(output, "/// {}", record.summary).unwrap(); + output.push_str( + "///\n/// # Safety\n/// Every participating thread must use a compatible barrier ID and expected thread count.\n#[inline(never)]\n", + ); + writeln!( + output, + "pub unsafe fn {}(barrier_id: u32, thread_count: u32) {{", + record.rust.name + ) + .unwrap(); + output.push_str(" let _ = (barrier_id, thread_count);\n"); + writeln!( + output, + " unreachable!(\"{} called outside CUDA kernel context\")", + record.rust.name + ) + .unwrap(); + output.push_str("}\n\n"); + } + output +} + +pub(super) fn render_compat_grid_dependency(catalog: &CatalogFile, hash: &str) -> String { + assert_eq!( + execution_control_family(catalog, "grid_dependency").count(), + 2 + ); + let mut output = rust_header(catalog, hash); + output.push_str( + "// Included inside `cuda_device::grid` to keep its public API stable.\n\npub mod dependency {\n", + ); + for (id, name) in [ + ("grid_dependency_launch_dependents", "trigger_dependents"), + ("grid_dependency_wait", "wait"), + ] { + let record = execution_control_family(catalog, "grid_dependency") + .find(|record| record.id == id) + .expect("complete grid-dependency family"); + writeln!(output, " /// {}", record.summary).unwrap(); + output.push_str( + " ///\n /// # Safety\n /// The kernel launch must participate in a valid programmatic dependent-launch protocol.\n #[inline(never)]\n", + ); + writeln!(output, " pub unsafe fn {name}() {{").unwrap(); + writeln!( + output, + " unreachable!(\"{name} called outside CUDA kernel context\")" + ) + .unwrap(); + output.push_str(" }\n\n"); + } + output.push_str("}\n"); + output +} + +pub(super) fn render_compat_register_control(catalog: &CatalogFile, hash: &str) -> String { + assert_eq!( + execution_control_family(catalog, "register_control").count(), + 2 + ); + let mut output = rust_header(catalog, hash); + output.push_str("// Included inside `cuda_device::thread` to keep its public API stable.\n\n"); + for (id, public_name, hidden_name) in [ + ("setmaxnreg_inc", "setmaxnreg_inc", "__setmaxnreg_inc"), + ("setmaxnreg_dec", "setmaxnreg_dec", "__setmaxnreg_dec"), + ] { + let record = execution_control_family(catalog, "register_control") + .find(|record| record.id == id) + .expect("complete register-control family"); + writeln!(output, "/// {}", record.summary).unwrap(); + output.push_str( + "///\n/// # Safety\n/// Every thread in the warpgroup must execute the same operation with the same count.\n#[inline(always)]\n", + ); + writeln!(output, "pub unsafe fn {public_name}() {{").unwrap(); + writeln!(output, " unsafe {{ {hidden_name}(N) }}").unwrap(); + output.push_str("}\n\n#[doc(hidden)]\n#[inline(never)]\n"); + writeln!( + output, + "pub(crate) unsafe fn {hidden_name}(_register_count: u32) {{" + ) + .unwrap(); + writeln!( + output, + " unreachable!(\"{public_name} called outside CUDA kernel context\")" + ) + .unwrap(); + output.push_str("}\n\n"); + } + output +} + +pub(super) fn render_compat_tma(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str("// Included inside `cuda_device::tma` to keep its public API stable.\n\n"); + for record in tma_intrinsics(catalog) { + let tma = record.tma.as_ref().expect("TMA contract"); + let operation = tma.operation; + writeln!(output, "/// {}", record.summary).unwrap(); + let dimensions = tma.dimensions(); + let is_g2s = matches!( + operation, + TmaOperation::G2sTile1d + | TmaOperation::G2sTile2d + | TmaOperation::G2sTile2dMulticast + | TmaOperation::G2sTile2dMulticastCg2 + | TmaOperation::G2sTile3d + | TmaOperation::G2sTile4d + | TmaOperation::G2sTile5d + ); + let is_s2g = matches!( + operation, + TmaOperation::S2gTile1d + | TmaOperation::S2gTile2d + | TmaOperation::S2gTile3d + | TmaOperation::S2gTile4d + | TmaOperation::S2gTile5d + ); + let is_reduction = operation == TmaOperation::Reduce; + if !record.rust.safe { + output.push_str("///\n/// # Safety\n"); + if is_g2s { + output.push_str( + "/// `dst`, `tensor_map`, and `barrier` must remain valid until the copy completes.\n", + ); + } else if is_s2g { + output.push_str( + "/// `src` and `tensor_map` must remain valid until the committed copy group completes.\n", + ); + } else if is_reduction { + output.push_str( + "/// `src` must name a live shared-memory tile and `tensor_map` must describe a compatible global destination until the committed reduction completes.\n", + ); + } else if matches!( + tma.adapter, + TmaAdapter::DescriptorAndAddressPointers + | TmaAdapter::DescriptorOrdinalAndU32 + | TmaAdapter::DescriptorOrdinalAndU64 + | TmaAdapter::DescriptorAndImmediateU32 + | TmaAdapter::DescriptorAndRuntimeU32 + ) { + output.push_str( + "/// `tensor_map` must point to a writable, 128-byte tensor-map descriptor in global memory.\n", + ); + } else { + output.push_str("/// `tensor_map` must point to a live tensor-map descriptor.\n"); + } + } + output.push_str("#[inline(never)]\n"); + if is_g2s { + let dimensions = dimensions.unwrap(); + let mut arguments = vec![ + "dst: *mut u8".to_owned(), + "tensor_map: *const TmaDescriptor".to_owned(), + ]; + let mut values = vec!["dst".to_owned(), "tensor_map".to_owned()]; + for index in 0..dimensions { + arguments.push(format!("coord{index}: i32")); + values.push(format!("coord{index}")); + } + arguments.push("barrier: *mut Barrier".into()); + values.push("barrier".into()); + if matches!( + operation, + TmaOperation::G2sTile2dMulticast | TmaOperation::G2sTile2dMulticastCg2 + ) { + arguments.push("cta_mask: u16".into()); + values.push("cta_mask".into()); + } + if arguments.len() > 7 { + output.push_str("#[allow(clippy::too_many_arguments)]\n"); + } + writeln!( + output, + "pub unsafe fn {}({}) {{", + record.rust.name, + arguments.join(", ") + ) + .unwrap(); + writeln!(output, " let _ = ({});", values.join(", ")).unwrap(); + } else if is_s2g || is_reduction { + let dimensions = dimensions.unwrap(); + let mut arguments = vec![ + "src: *const u8".to_owned(), + "tensor_map: *const TmaDescriptor".to_owned(), + ]; + let mut values = vec!["src".to_owned(), "tensor_map".to_owned()]; + for index in 0..dimensions { + arguments.push(format!("coord{index}: i32")); + values.push(format!("coord{index}")); + } + writeln!( + output, + "pub unsafe fn {}({}) {{", + record.rust.name, + arguments.join(", ") + ) + .unwrap(); + writeln!(output, " let _ = ({});", values.join(", ")).unwrap(); + } else if operation == TmaOperation::CommitGroup { + writeln!(output, "pub fn {}() {{", record.rust.name).unwrap(); + } else if matches!( + operation, + TmaOperation::WaitGroup | TmaOperation::WaitGroupRead + ) { + writeln!(output, "pub fn {}(n: u32) {{", record.rust.name).unwrap(); + output.push_str(" let _ = n;\n"); + } else if operation == TmaOperation::PrefetchTensorMap { + writeln!( + output, + "pub unsafe fn {}(tensor_map: *const TmaDescriptor) {{", + record.rust.name + ) + .unwrap(); + output.push_str(" let _ = tensor_map;\n"); + } else if let Some(coordinate_count) = operation.prefetch_coordinate_count() { + let mut arguments = vec!["tensor_map: *const TmaDescriptor".to_owned()]; + let mut values = vec!["tensor_map".to_owned()]; + for index in 0..coordinate_count { + arguments.push(format!("coord{index}: i32")); + values.push(format!("coord{index}")); + } + if operation.uses_prefetch_cache_hint() { + arguments.push("cache_hint: u64".into()); + values.push("cache_hint".into()); + } + writeln!( + output, + "pub unsafe fn {}({}) {{", + record.rust.name, + arguments.join(", ") + ) + .unwrap(); + writeln!(output, " let _ = ({});", values.join(", ")).unwrap(); + } else { + match tma.adapter { + TmaAdapter::DescriptorAndAddressPointers => { + writeln!( + output, + "pub unsafe fn {}(tensor_map: *mut TmaDescriptor, new_address: *const u8) {{", + record.rust.name + ) + .unwrap(); + output.push_str(" let _ = (tensor_map, new_address);\n"); + } + TmaAdapter::DescriptorOrdinalAndU32 => { + writeln!( + output, + "pub unsafe fn {}(tensor_map: *mut TmaDescriptor, ordinal: u32, new_value: u32) {{", + record.rust.name + ) + .unwrap(); + output.push_str(" let _ = (tensor_map, ordinal, new_value);\n"); + } + TmaAdapter::DescriptorOrdinalAndU64 => { + writeln!( + output, + "pub unsafe fn {}(tensor_map: *mut TmaDescriptor, ordinal: u32, new_value: u64) {{", + record.rust.name + ) + .unwrap(); + output.push_str(" let _ = (tensor_map, ordinal, new_value);\n"); + } + TmaAdapter::DescriptorAndImmediateU32 | TmaAdapter::DescriptorAndRuntimeU32 => { + writeln!( + output, + "pub unsafe fn {}(tensor_map: *mut TmaDescriptor, new_value: u32) {{", + record.rust.name + ) + .unwrap(); + output.push_str(" let _ = (tensor_map, new_value);\n"); + } + TmaAdapter::DescriptorPointerInjectBytes => { + writeln!( + output, + "pub unsafe fn {}(tensor_map: *const TmaDescriptor) {{", + record.rust.name + ) + .unwrap(); + output.push_str(" let _ = tensor_map;\n"); + } + TmaAdapter::NoOperands => { + writeln!(output, "pub fn {}() {{", record.rust.name).unwrap(); + } + _ => unreachable!("TMA compatibility operation category was matched"), + } + } + writeln!( + output, + " unreachable!(\"{} called outside CUDA kernel context\")", + record.rust.name + ) + .unwrap(); + output.push_str("}\n\n"); + } + output.push_str( + r#"/// Replace a global tensor-map address and publish it to the tensor-map proxy. +/// +/// # Safety +/// `tensor_map` must point to a writable, 128-byte tensor-map descriptor in +/// global memory, and `new_address` must remain valid for every later TMA use. +#[inline(always)] +pub unsafe fn replace_tma_global_address( + tensor_map: *mut TmaDescriptor, + new_address: *const u8, +) { + unsafe { tensormap_replace_global_address(tensor_map, new_address) }; + fence_proxy_tensormap_generic_release_gpu(); + unsafe { fence_proxy_tensormap_generic_acquire_gpu(tensor_map) }; +} +"#, + ); + output +} + +pub(super) fn render_compat_wgmma_control(catalog: &CatalogFile, hash: &str) -> String { + assert_eq!(wgmma_controls(catalog).count(), 3); + let mut output = rust_header(catalog, hash); + output.push_str("// Included inside `cuda_device::wgmma` to keep its public API stable.\n\n"); + + for mode in [WgmmaControlMode::Fence, WgmmaControlMode::CommitGroup] { + let record = wgmma_control(catalog, mode); + writeln!(output, "/// {}", record.summary).unwrap(); + output.push_str("#[inline(never)]\n"); + writeln!(output, "pub fn {}() {{", record.rust.name).unwrap(); + writeln!( + output, + " unreachable!({:?})", + format!("{} called outside CUDA kernel context", record.rust.name) + ) + .unwrap(); + output.push_str("}\n\n"); + } + + let wait = wgmma_control(catalog, WgmmaControlMode::WaitGroup); + writeln!(output, "/// {}", wait.summary).unwrap(); + output.push_str( + r#"#[inline(always)] +pub fn wgmma_wait_group() { + __wgmma_wait_group(N as u64); +} + +#[doc(hidden)] +#[inline(never)] +pub(crate) fn __wgmma_wait_group(_max_pending: u64) { + unreachable!("wgmma_wait_group called outside CUDA kernel context") +} +"#, + ); + output +} + +fn render_compat_tcgen05_mma_record(output: &mut String, record: &CatalogIntrinsic) { + let mma = record + .tcgen05 + .as_ref() + .and_then(|tcgen05| tcgen05.mma.as_ref()) + .expect("tcgen05 MMA record"); + let parameters = tcgen05_mma_runtime_parameters(mma); + let signature = parameters + .iter() + .map(|(name, ty)| format!("{name}: {ty}")) + .collect::>() + .join(", "); + let values = parameters + .iter() + .map(|(name, _)| *name) + .collect::>() + .join(", "); + + writeln!(output, "/// {}", record.summary).unwrap(); + output.push_str("/// All tcgen05 operations in the kernel must use the same CTA-group mode.\n"); + if mma.alias.is_some() && mma.form == Tcgen05MmaForm::WsTensor { + output.push_str("/// This uses kind f8f6f4 and collector b0::discard.\n"); + output.push_str("/// `legacy_a_desc` is kept for compatibility; tensor A uses `a_tmem`.\n"); + } else if mma.alias.is_some() { + output.push_str("/// This uses kind f8f6f4, CTA group 1, and collector a::discard.\n"); + } else { + match mma.selector_layout { + Tcgen05MmaSelectorLayout::Base { .. } => { + output.push_str("/// `KIND` is 0=f16, 1=tf32, 2=f8f6f4, or 3=i8.\n"); + output.push_str( + "/// `CTA_GROUP` is 1 or 2. `COLLECTOR_A` is 0=discard, 1=lastuse, 2=fill, or 3=use.\n", + ); + } + Tcgen05MmaSelectorLayout::WarpSpecialized { .. } => { + output.push_str("/// `KIND` is 0=f16, 1=tf32, 2=f8f6f4, or 3=i8.\n"); + output.push_str( + "/// `B_BUFFER` is 0 through 3. `B_USAGE` is 0=discard, 1=lastuse, 2=fill, or 3=use.\n", + ); + } + } + } + output.push_str("///\n/// # Safety\n"); + output + .push_str("/// Tensor-memory addresses and descriptors must be valid for this MMA form.\n"); + if mma.alias.is_some() { + output.push_str("#[inline(never)]\n"); + output.push_str("#[allow(clippy::too_many_arguments)]\n"); + writeln!(output, "pub unsafe fn {}({signature}) {{", record.rust.name).unwrap(); + writeln!(output, " let _ = ({values});").unwrap(); + writeln!( + output, + " unreachable!(\"{} called outside CUDA kernel context\")", + record.rust.name + ) + .unwrap(); + output.push_str("}\n\n"); + return; + } + + let selectors = tcgen05_mma_selector_parameters(mma.selector_layout); + let const_parameters = selectors + .iter() + .map(|(name, _)| format!("const {name}: u32")) + .collect::>() + .join(", "); + let selector_values = selectors + .iter() + .map(|(name, _)| *name) + .collect::>() + .join(", "); + output.push_str("#[inline(always)]\n"); + if parameters.len() > 5 { + output.push_str("#[allow(clippy::too_many_arguments)]\n"); + } + writeln!( + output, + "pub unsafe fn {}<{const_parameters}>({signature}) {{", + record.rust.name + ) + .unwrap(); + writeln!( + output, + " unsafe {{ __{}({values}, {selector_values}) }}", + record.rust.name + ) + .unwrap(); + output.push_str("}\n\n#[doc(hidden)]\n#[inline(never)]\n"); + output.push_str("#[allow(clippy::too_many_arguments)]\n"); + let hidden_parameters = parameters + .iter() + .map(|(name, ty)| format!("_{name}: {ty}")) + .chain(selectors.iter().map(|(_, name)| format!("_{name}: u32"))) + .collect::>() + .join(", "); + writeln!( + output, + "pub(crate) unsafe fn __{}({hidden_parameters}) {{", + record.rust.name + ) + .unwrap(); + writeln!( + output, + " unreachable!(\"{} called outside CUDA kernel context\")", + record.rust.name + ) + .unwrap(); + output.push_str("}\n\n"); +} + +pub(super) fn render_compat_tcgen05(catalog: &CatalogFile, hash: &str) -> String { + assert_eq!(tcgen05_intrinsics(catalog).count(), 233); + let mut output = rust_header(catalog, hash); + output.push_str("// Included inside `cuda_device::tcgen05` to keep its public API stable.\n\n"); + for record in tcgen05_intrinsics(catalog) { + let tcgen05 = record.tcgen05.as_ref().unwrap(); + if tcgen05.mma.is_some() { + render_compat_tcgen05_mma_record(&mut output, record); + continue; + } + let operation = tcgen05.operation; + let has_half_split_offset = tcgen05 + .ld + .is_some_and(|ld| ld.shape == Tcgen05LdShape::M16x32bx2) + || tcgen05 + .st + .is_some_and(|st| st.shape == Tcgen05LdShape::M16x32bx2); + if has_half_split_offset { + let (result, data) = if operation == Tcgen05Operation::Ld { + let count = tcgen05_ld_register_count(record); + ( + if count == 1 { + "u32".into() + } else { + format!("CuSimd") + }, + None, + ) + } else { + let count = tcgen05_st_register_count(record); + ( + "()".into(), + Some(if count == 1 { + "u32".into() + } else { + format!("CuSimd") + }), + ) + }; + writeln!(output, "/// {}", record.summary).unwrap(); + output.push_str( + "/// All tcgen05 operations in the kernel must use the same CTA-group mode.\n", + ); + output.push_str("///\n/// # Safety\n"); + output.push_str( + "/// The tensor-memory address must be live, and the warp must execute this call uniformly.\n", + ); + output.push_str("#[inline(always)]\n"); + if let Some(data) = &data { + writeln!( + output, + "pub unsafe fn {}(tmem_addr: u32, data: {data}) {{", + record.rust.name + ) + .unwrap(); + writeln!( + output, + " unsafe {{ __{}(tmem_addr, HALF_SPLIT_OFFSET as i64, data) }}", + record.rust.name + ) + .unwrap(); + } else { + writeln!( + output, + "pub unsafe fn {}(tmem_addr: u32) -> {result} {{", + record.rust.name + ) + .unwrap(); + writeln!( + output, + " unsafe {{ __{}(tmem_addr, HALF_SPLIT_OFFSET as i64) }}", + record.rust.name + ) + .unwrap(); + } + output.push_str("}\n\n#[doc(hidden)]\n#[inline(never)]\n"); + if let Some(data) = data { + writeln!( + output, + "pub(crate) unsafe fn __{}(_tmem_addr: u32, _half_split_offset: i64, _data: {data}) {{", + record.rust.name + ) + .unwrap(); + } else { + writeln!( + output, + "pub(crate) unsafe fn __{}(_tmem_addr: u32, _half_split_offset: i64) -> {result} {{", + record.rust.name + ) + .unwrap(); + } + writeln!( + output, + " unreachable!(\"{} called outside CUDA kernel context\")", + record.rust.name + ) + .unwrap(); + output.push_str("}\n\n"); + continue; + } + let (arguments, values): (String, String) = if operation == Tcgen05Operation::St { + let count = tcgen05_st_register_count(record); + let data = if count == 1 { + "u32".into() + } else { + format!("CuSimd") + }; + ( + format!("tmem_addr: u32, data: {data}"), + "tmem_addr, data".into(), + ) + } else { + let (arguments, values) = match operation { + Tcgen05Operation::Alloc | Tcgen05Operation::AllocCg2 => { + ("dst_smem: *mut u32, n_cols: u32", "dst_smem, n_cols") + } + Tcgen05Operation::Dealloc | Tcgen05Operation::DeallocCg2 => { + ("tmem_addr: u32, n_cols: u32", "tmem_addr, n_cols") + } + Tcgen05Operation::RelinquishAllocPermit + | Tcgen05Operation::FenceBeforeThreadSync + | Tcgen05Operation::FenceAfterThreadSync + | Tcgen05Operation::LoadWait + | Tcgen05Operation::StoreWait + | Tcgen05Operation::RelinquishAllocPermitCg2 => ("", ""), + Tcgen05Operation::Commit + | Tcgen05Operation::CommitSharedCluster + | Tcgen05Operation::CommitCg2 + | Tcgen05Operation::CommitSharedClusterCg2 => ("mbar: *mut u64", "mbar"), + Tcgen05Operation::MmaWsF16 + | Tcgen05Operation::MmaWsBf16 + | Tcgen05Operation::MmaWsTf32 => ( + "d_tmem: u32, a_tmem: u32, a_desc: u64, b_desc: u64, idesc: u32, enable_d: bool", + "d_tmem, a_tmem, a_desc, b_desc, idesc, enable_d", + ), + Tcgen05Operation::MmaF16 | Tcgen05Operation::MmaF16Cg2 => ( + "d_tmem: u32, a_desc: u64, b_desc: u64, idesc: u32, enable_d: bool", + "d_tmem, a_desc, b_desc, idesc, enable_d", + ), + Tcgen05Operation::CpSmemToTmem | Tcgen05Operation::CpSmemToTmemCg2 => { + ("tmem_addr: u32, smem_desc: u64", "tmem_addr, smem_desc") + } + Tcgen05Operation::Ld16x256bX8Pure + | Tcgen05Operation::Ld16x256bPure + | Tcgen05Operation::Ld => ("tmem_addr: u32", "tmem_addr"), + Tcgen05Operation::CommitMulticast | Tcgen05Operation::CommitMulticastCg2 => { + ("mbar: *mut u64, cta_mask: u16", "mbar, cta_mask") + } + Tcgen05Operation::ShiftDown | Tcgen05Operation::ShiftDownCg2 => { + ("tmem_addr: u32", "tmem_addr") + } + Tcgen05Operation::St => unreachable!("store handled above"), + Tcgen05Operation::Mma => unreachable!("generic MMA handled above"), + }; + (arguments.into(), values.into()) + }; + let result: String = match operation { + Tcgen05Operation::Ld16x256bX8Pure => "TmemF32x32".into(), + Tcgen05Operation::Ld16x256bPure => "TmemF32x4".into(), + Tcgen05Operation::Ld => { + let count = tcgen05_ld_register_count(record); + if count == 1 { + "u32".into() + } else { + format!("CuSimd") + } + } + Tcgen05Operation::St => "()".into(), + _ => "()".into(), + }; + writeln!(output, "/// {}", record.summary).unwrap(); + if let Some(participation) = tcgen05_participation_doc(operation) { + writeln!(output, "/// {participation}").unwrap(); + } + output.push_str( + "/// All tcgen05 operations in the kernel must use the same CTA-group mode.\n", + ); + if !record.rust.safe { + output.push_str("///\n/// # Safety\n"); + if operation == Tcgen05Operation::St { + output.push_str( + "/// The tensor-memory address must remain live and cover the selected tile.\n\ + /// All active warp lanes must execute convergently with the same address. Complete the store wait before relying on completion.\n", + ); + } else if operation == Tcgen05Operation::Ld { + output.push_str( + "/// The tensor-memory address must remain live and cover the selected tile.\n\ + /// All active warp lanes must execute convergently with the same address. Complete the load wait before using the result.\n", + ); + } else if tcgen05_is_commit(operation) { + if tcgen05_is_multicast_commit(operation) { + output.push_str( + "/// `mbar` must point to a live initialized cluster-shared mbarrier. `cta_mask` must select valid CTA ranks in its cluster.\n", + ); + } else { + output.push_str( + "/// `mbar` must point to a live initialized mbarrier valid for this CTA-group mode.\n", + ); + } + output.push_str( + "/// The same thread that issued the tracked asynchronous tcgen05 operations must issue this commit.\n", + ); + } else if tcgen05_is_shift(operation) { + output.push_str( + "/// `tmem_addr` must name a live tensor-memory allocation, and its lane component must be a multiple of 32.\n\ + /// Completion must be tracked by a matching commit from that same thread and observed through the selected mbarrier before relying on shifted data.\n", + ); + } else { + output.push_str("/// The caller must satisfy the tcgen05 address, lifetime, and participation rules.\n"); + } + } + output.push_str("#[inline(never)]\n"); + if record.rust.arguments.len() > 5 { + output.push_str("#[allow(clippy::too_many_arguments)]\n"); + } + let safety = if record.rust.safe { "" } else { "unsafe " }; + writeln!( + output, + "pub {safety}fn {}({arguments}) -> {result} {{", + record.rust.name + ) + .unwrap(); + if !values.is_empty() { + if values.contains(',') { + writeln!(output, " let _ = ({values});").unwrap(); + } else { + writeln!(output, " let _ = {values};").unwrap(); + } + } + writeln!( + output, + " unreachable!(\"{} called outside CUDA kernel context\")", + record.rust.name + ) + .unwrap(); + output.push_str("}\n\n"); + } + output +} + +pub(super) fn render_compat_packed_atomic(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str("// Included inside `cuda_device::atomic` to keep existing paths stable.\n\n"); + for record in packed_atomics(catalog) { + let path = record + .rust + .compatibility_paths + .iter() + .find(|path| path.starts_with("cuda_device::atomic::")) + .expect("packed-atomic compatibility path"); + assert_eq!(path, &format!("cuda_device::atomic::{}", record.rust.name)); + let packed = record + .packed_atomic + .as_ref() + .expect("packed-atomic semantics"); + let lane_type = match packed.format { + PackedAtomicFormat::F16x2 => "f16", + PackedAtomicFormat::Bf16x2 => "bf16", + }; + let minimum_sm = match &record.target.hardware { + CatalogHardwareTarget::AnyOf { alternatives } => match alternatives.as_slice() { + [CatalogHardwareAlternative::MinimumSm { sm }] => *sm, + _ => panic!("packed-atomic compatibility API requires one minimum SM"), + }, + _ => panic!("packed-atomic compatibility API requires one minimum SM"), + }; + assert!(!record.rust.safe); + assert!(record.rust.must_use); + assert_eq!(record.rust.arguments, ["*mut u32", "u32"]); + assert_eq!(record.rust.result, "u32"); + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!( + output, + "/// `val` and the result pack two {lane_type} lanes into `u32`, low lane first." + ) + .unwrap(); + output.push_str( + "/// The lanes are atomic independently and may not form one old 32-bit snapshot.\n", + ); + output.push_str("/// This is a relaxed GPU-scope operation. Each lane rounds to nearest-even and preserves subnormals.\n"); + writeln!( + output, + "/// Requires PTX {} and `sm_{minimum_sm}+`.", + record.target.minimum_ptx + ) + .unwrap(); + output.push_str("///\n/// # Safety\n"); + output.push_str( + "/// `addr` must point to four writable, four-byte-aligned bytes in global memory.\n", + ); + output.push_str("/// Do not overlap this operation with a whole-word atomic or non-atomic lane access.\n"); + output.push_str("/// Racing atomics must use mutually inclusive scopes; host/system access is not included.\n"); + output.push_str("#[must_use]\n#[inline(never)]\n"); + writeln!( + output, + "pub unsafe fn {}(addr: *mut u32, val: u32) -> u32 {{", + record.rust.name + ) + .unwrap(); + output.push_str(" let _ = (addr, val);\n"); + writeln!( + output, + " unreachable!(\"{} called outside CUDA kernel context\")", + record.rust.name + ) + .unwrap(); + output.push_str("}\n\n"); + } + output +} + +pub(super) fn render_compat_cp_async_copy(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + "// Included inside `cuda_device::async_copy` to keep existing paths stable.\n\n", + ); + for record in cp_async_copies(catalog) { + let copy = record.cp_async_copy.as_ref().expect("cp.async semantics"); + let bytes = copy.copy_size.bytes(); + let cache = match copy.cache_policy { + CpAsyncCachePolicy::Ca => "cache-all", + CpAsyncCachePolicy::Cg => "cache-global", + }; + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!( + output, + "/// Uses the {cache} policy and copies {bytes} bytes." + ) + .unwrap(); + if copy.source_size == CpAsyncSourceSize::Runtime { + writeln!( + output, + "/// Bytes after `src_size` are filled with zero; `src_size` must be at most {bytes}." + ) + .unwrap(); + } + output.push_str("///\n/// # Safety\n"); + writeln!( + output, + "/// `shared_dst` must point to {bytes} writable bytes in shared memory and be aligned to {bytes} bytes." + ) + .unwrap(); + if copy.source_size == CpAsyncSourceSize::Runtime { + writeln!( + output, + "/// `global_src` must point to `src_size` readable bytes in global memory and be aligned to {bytes} bytes." + ) + .unwrap(); + } else { + writeln!( + output, + "/// `global_src` must point to {bytes} readable bytes in global memory and be aligned to {bytes} bytes." + ) + .unwrap(); + } + output.push_str( + "/// Both ranges must remain valid, the source must remain unchanged, and the destination must not be accessed until this copy completes.\n\ + /// The issuing thread must complete the copy. Synchronize threads afterward before another thread accesses the destination.\n\ + /// User-authored completion assembly must include a compiler memory clobber.\n", + ); + output.push_str("#[inline(never)]\n"); + if copy.source_size == CpAsyncSourceSize::Runtime { + writeln!( + output, + "pub unsafe fn {}(_shared_dst: *mut u32, _global_src: *const u8, _src_size: u32) {{", + record.rust.name + ) + .unwrap(); + } else { + writeln!( + output, + "pub unsafe fn {}(_shared_dst: *mut u32, _global_src: *const u32) {{", + record.rust.name + ) + .unwrap(); + } + writeln!( + output, + " unreachable!(\"{} called outside CUDA kernel context\")", + record.rust.name + ) + .unwrap(); + output.push_str("}\n\n"); + } + for record in cp_async_controls(catalog) { + let control = record + .cp_async_control + .as_ref() + .expect("cp.async control semantics"); + writeln!(output, "/// {}", record.summary).unwrap(); + output.push_str("///\n/// # Safety\n"); + match control.operation { + CpAsyncControlOperation::CommitGroup => output.push_str( + "/// This commits only copies issued by the executing thread and does not wait for completion.\n", + ), + CpAsyncControlOperation::WaitAll => output.push_str( + "/// This waits only for copies issued by this thread. Synchronize threads before another thread accesses a completed destination.\n", + ), + CpAsyncControlOperation::WaitGroup => output.push_str( + "/// `max_pending` must be a compile-time constant. Access only destinations whose copy groups this wait completes.\n", + ), + } + output.push_str("#[inline(never)]\n"); + if control.operation == CpAsyncControlOperation::WaitGroup { + writeln!( + output, + "pub unsafe fn {}(_max_pending: u32) {{", + record.rust.name + ) + .unwrap(); + } else { + writeln!(output, "pub unsafe fn {}() {{", record.rust.name).unwrap(); + } + writeln!( + output, + " unreachable!(\"{} called outside CUDA kernel context\")", + record.rust.name + ) + .unwrap(); + output.push_str("}\n\n"); + } + for record in cp_async_mbarriers(catalog) { + let bridge = record + .cp_async_mbarrier + .as_ref() + .expect("cp.async mbarrier semantics"); + let path = record + .rust + .compatibility_paths + .iter() + .find(|path| path.starts_with("cuda_device::async_copy::")) + .expect("cp.async mbarrier compatibility path"); + assert_eq!( + path, + &format!("cuda_device::async_copy::{}", record.rust.name) + ); + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!(output, "/// Lowers to `{}`.", record.expected_ptx).unwrap(); + output.push_str("///\n/// # Safety\n"); + output.push_str( + "/// `barrier` must point to a live, initialized, eight-byte-aligned mbarrier object in shared memory.\n\ + /// This thread must have issued the `cp.async` operations being associated with it.\n\ + /// The object must remain valid until those operations complete.\n", + ); + match bridge.operation { + CpAsyncMbarrierOperation::Arrive => output.push_str( + "/// This form increments the pending count before scheduling the asynchronous arrival.\n\ + /// That increment must not exceed the barrier's pending-count limit.\n", + ), + CpAsyncMbarrierOperation::ArriveNoInc => output.push_str( + "/// The initial pending count must already include the asynchronous arrival.\n", + ), + } + output.push_str("#[inline(never)]\n"); + writeln!( + output, + "pub unsafe fn {}(_barrier: *mut crate::barrier::Barrier) {{", + record.rust.name + ) + .unwrap(); + writeln!( + output, + " unreachable!(\"{} called outside CUDA kernel context\")", + record.rust.name + ) + .unwrap(); + output.push_str("}\n\n"); + } + output +} + +pub(super) fn render_compat_mbarrier_basic(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str("// Included inside `cuda_device::barrier` to keep existing paths stable.\n\n"); + for record in mbarrier_basics(catalog) { + let mbarrier = record + .mbarrier_basic + .as_ref() + .expect("basic mbarrier semantics"); + let path = record + .rust + .compatibility_paths + .iter() + .find(|path| path.starts_with("cuda_device::barrier::")) + .expect("basic mbarrier compatibility path"); + assert_eq!(path, &format!("cuda_device::barrier::{}", record.rust.name)); + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!(output, "/// Lowers to `{}`.", record.expected_ptx).unwrap(); + output.push_str("///\n/// # Safety\n"); + output.push_str( + "/// `bar` must point to a live, eight-byte-aligned `Barrier` in shared memory.\n", + ); + match mbarrier.operation { + MbarrierBasicOperation::Init => output.push_str( + "/// Exactly one thread may initialize it. `expected_count` must be in `1..=0xFFFFF` and include every arrival in the phase.\n\ + /// The barrier must be uninitialized or invalidated, with no concurrent barrier operation.\n", + ), + MbarrierBasicOperation::Arrive => output.push_str( + "/// The barrier must be initialized, and this arrival must be included in the current phase's expected count.\n\ + /// Use the returned token only with this barrier and phase.\n", + ), + MbarrierBasicOperation::TestWait => output.push_str( + "/// The barrier must be initialized. `token` must come from this barrier and phase.\n", + ), + MbarrierBasicOperation::Inval => output.push_str( + "/// The barrier must be initialized and unused by all threads and asynchronous operations.\n\ + /// Exactly one thread may invalidate it.\n", + ), + } + output.push_str("#[inline(never)]\n"); + match mbarrier.operation { + MbarrierBasicOperation::Init => { + writeln!( + output, + "pub unsafe fn {}(bar: *mut Barrier, expected_count: u32) {{", + record.rust.name + ) + .unwrap(); + output.push_str(" let _ = (bar, expected_count);\n"); + } + MbarrierBasicOperation::Arrive => { + writeln!( + output, + "pub unsafe fn {}(bar: *const Barrier) -> u64 {{", + record.rust.name + ) + .unwrap(); + output.push_str(" let _ = bar;\n"); + } + MbarrierBasicOperation::TestWait => { + writeln!( + output, + "pub unsafe fn {}(bar: *const Barrier, token: u64) -> bool {{", + record.rust.name + ) + .unwrap(); + output.push_str(" let _ = (bar, token);\n"); + } + MbarrierBasicOperation::Inval => { + writeln!( + output, + "pub unsafe fn {}(bar: *mut Barrier) {{", + record.rust.name + ) + .unwrap(); + output.push_str(" let _ = bar;\n"); + } + } + writeln!( + output, + " unreachable!(\"{} called outside CUDA kernel context\")", + record.rust.name + ) + .unwrap(); + output.push_str("}\n\n"); + } + output +} + +pub(super) fn render_compat_mbarrier_extended(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str("// Included inside `cuda_device::barrier` to keep existing paths stable.\n\n"); + for record in mbarrier_extended(catalog) { + let contract = record + .mbarrier_extended + .as_ref() + .expect("extended-mbarrier contract"); + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!(output, "/// Lowers to `{}`.", record.expected_ptx).unwrap(); + output.push_str("///\n/// # Safety\n"); + output.push_str( + "/// The caller must satisfy the barrier, scope, and memory-ordering contract.\n", + ); + output.push_str("#[inline(never)]\n"); + match contract.operation { + MbarrierExtendedOperation::ArriveExpectTxCta + | MbarrierExtendedOperation::ArriveExpectTxCluster => { + writeln!( + output, + "pub unsafe fn {}(bar: *const Barrier, _tx_count: u32, bytes: u32) -> u64 {{", + record.rust.name + ) + .unwrap(); + output.push_str(" let _ = (bar, bytes);\n"); + } + MbarrierExtendedOperation::ArriveRemoteCluster => { + writeln!( + output, + "pub unsafe fn {}(remote_bar_addr: u64) {{", + record.rust.name + ) + .unwrap(); + output.push_str(" let _ = remote_bar_addr;\n"); + } + MbarrierExtendedOperation::TryWaitTokenCta => { + writeln!( + output, + "pub unsafe fn {}(bar: *const Barrier, token: u64) -> bool {{", + record.rust.name + ) + .unwrap(); + output.push_str(" let _ = (bar, token);\n"); + } + MbarrierExtendedOperation::TryWaitParityCta + | MbarrierExtendedOperation::TryWaitParityCluster => { + writeln!( + output, + "pub unsafe fn {}(bar: *const Barrier, parity: u32) -> bool {{", + record.rust.name + ) + .unwrap(); + output.push_str(" let _ = (bar, parity);\n"); + } + MbarrierExtendedOperation::FenceProxyAsyncSharedCta + | MbarrierExtendedOperation::FenceMbarrierInitReleaseCluster + | MbarrierExtendedOperation::FenceProxyAsyncGenericReleaseSharedCtaCluster + | MbarrierExtendedOperation::FenceProxyAsyncGenericAcquireSharedClusterCluster => { + writeln!(output, "pub unsafe fn {}() {{", record.rust.name).unwrap(); + } + MbarrierExtendedOperation::Nanosleep => { + writeln!(output, "pub unsafe fn {}(ns: u32) {{", record.rust.name).unwrap(); + output.push_str(" let _ = ns;\n"); + } + } + writeln!( + output, + " unreachable!(\"{} called outside CUDA kernel context\")", + record.rust.name + ) + .unwrap(); + output.push_str("}\n\n"); + } + output +} + +pub(super) fn render_compat_integer_minmax( + catalog: &CatalogFile, + hash: &str, + module: &str, +) -> String { + let mut output = rust_header(catalog, hash); + writeln!( + output, + "// Included inside `cuda_device::{module}` to keep existing paths stable.\n" + ) + .unwrap(); + for record in integer_minmaxes(catalog).filter(|record| record.rust.module == module) { + let path = record + .rust + .compatibility_paths + .iter() + .find(|path| path.starts_with(&format!("cuda_device::{module}::"))) + .expect("integer-min/max compatibility path"); + let scalar = &record.rust.result; + writeln!(output, "/// {}", record.summary).unwrap(); + output.push_str("#[must_use]\n#[inline(never)]\n"); + writeln!( + output, + "pub fn {}(arg0: {scalar}, arg1: {scalar}) -> {scalar} {{", + record.rust.name + ) + .unwrap(); + output.push_str(" let _ = (arg0, arg1);\n"); + writeln!( + output, + " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" + ) + .unwrap(); + output.push_str("}\n\n"); + } + output +} + +pub(super) fn render_compat_packed_alu( + catalog: &CatalogFile, + hash: &str, + format: PackedAluFormat, +) -> String { + let mut output = rust_header(catalog, hash); + let (module, _, _, _, _) = packed_alu_format_shape(format); + writeln!( + output, + "// Included inside `cuda_device::{module}` to keep existing paths stable.\n" + ) + .unwrap(); + for record in packed_alus(catalog).filter(|record| { + record + .packed_alu + .as_ref() + .is_some_and(|packed| packed.format == format) + }) { + let path = record + .rust + .compatibility_paths + .iter() + .find(|path| path.starts_with(&format!("cuda_device::{module}::"))) + .expect("packed-ALU compatibility path"); + let arguments = record + .rust + .arguments + .iter() + .enumerate() + .map(|(index, ty)| format!("arg{index}: {ty}")) + .collect::>() + .join(", "); + let values = (0..record.rust.arguments.len()) + .map(|index| format!("arg{index}")) + .collect::>() + .join(", "); + writeln!(output, "/// {}", record.summary).unwrap(); + if record.rust.must_use { + output.push_str("#[must_use]\n"); + } + output.push_str("#[inline(never)]\n"); + writeln!( + output, + "pub fn {}({arguments}) -> {} {{", + record.rust.name, record.rust.result + ) + .unwrap(); + if record.rust.arguments.len() == 1 { + writeln!(output, " let _ = {values};").unwrap(); + } else { + writeln!(output, " let _ = ({values});").unwrap(); + } + writeln!( + output, + " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" + ) + .unwrap(); + output.push_str("}\n\n"); + } + render_compat_extended_minmax_into(&mut output, catalog, module); + output +} + +pub(super) fn render_compat_scalar_minmax( + catalog: &CatalogFile, + hash: &str, + module: &str, +) -> String { + let mut output = rust_header(catalog, hash); + writeln!( + output, + "// Included inside `cuda_device::{module}` to keep existing paths stable.\n" + ) + .unwrap(); + render_compat_extended_minmax_into(&mut output, catalog, module); + output +} + +fn render_compat_extended_minmax_into(output: &mut String, catalog: &CatalogFile, module: &str) { + for record in extended_minmax(catalog).filter(|record| record.rust.module == module) { + let path = record + .rust + .compatibility_paths + .iter() + .find(|path| path.starts_with(&format!("cuda_device::{module}::"))) + .expect("extended-minmax compatibility path"); + let ty = extended_minmax_rust_type(record); + writeln!(output, "/// {}", record.summary).unwrap(); + output.push_str("#[must_use]\n#[inline(never)]\n"); + writeln!( + output, + "pub fn {}(a: {ty}, b: {ty}) -> {ty} {{", + record.rust.name + ) + .unwrap(); + output.push_str(" let _ = (a, b);\n"); + writeln!( + output, + " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" + ) + .unwrap(); + output.push_str("}\n\n"); + } +} + +pub(super) fn render_compat_packed_conversion( + catalog: &CatalogFile, + hash: &str, + path_prefix: &str, + containing_module: &str, + argument_names: (&str, &str), +) -> String { + let mut output = rust_header(catalog, hash); + writeln!( + output, + "// Included inside `cuda_device::{containing_module}`.\n" + ) + .unwrap(); + for record in packed_conversions(catalog).filter(|record| { + record + .rust + .compatibility_paths + .iter() + .any(|path| path.starts_with(path_prefix)) + }) { + let path = record + .rust + .compatibility_paths + .iter() + .find(|path| path.starts_with(path_prefix)) + .expect("packed-conversion compatibility path"); + // A packed source arrives in one register, so the parameter list + // follows the record's own argument types rather than a fixed f32 pair. + let parameter_names: Vec<&str> = match packed_conversion_source(record) { + PackedConversionSourceFormat::F32x2 => vec![argument_names.0, argument_names.1], + PackedConversionSourceFormat::E4m3x2 + | PackedConversionSourceFormat::E5m2x2 + | PackedConversionSourceFormat::F16x2 => vec!["packed"], + }; + let parameter_types = packed_conversion_rust_arguments(record); + let parameters = parameter_names + .iter() + .zip(¶meter_types) + .map(|(name, ty)| format!("{name}: {ty}")) + .collect::>() + .join(", "); + let discards = parameter_names.join(", "); + writeln!(output, "/// {}", record.summary).unwrap(); + output.push_str("#[inline(never)]\n"); + writeln!( + output, + "pub fn {}({parameters}) -> {} {{", + path.rsplit("::") + .next() + .expect("packed-conversion compatibility function name"), + record.rust.result, + ) + .unwrap(); + if parameter_names.len() == 1 { + writeln!(output, " let _ = {discards};").unwrap(); + } else { + writeln!(output, " let _ = ({discards});").unwrap(); + } + writeln!( + output, + " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" + ) + .unwrap(); + output.push_str("}\n\n"); + } + if containing_module == "convert" { + for record in scalar_conversions(catalog) { + let path = record + .rust + .compatibility_paths + .iter() + .find(|path| path.starts_with("cuda_device::convert::")) + .expect("scalar-conversion compatibility path"); + writeln!(output, "/// {}", record.summary).unwrap(); + output.push_str("#[must_use]\n#[inline(never)]\n"); + writeln!(output, "pub fn {}(value: f32) -> u32 {{", record.rust.name).unwrap(); + output.push_str(" let _ = value;\n"); + writeln!( + output, + " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" + ) + .unwrap(); + output.push_str("}\n\n"); + } + } + output +} + +pub(super) fn render_compat_float(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str("// Included inside `cuda_device::float`.\n\n"); + for record in scalar_arithmetics(catalog) { + let path = record + .rust + .compatibility_paths + .iter() + .find(|path| path.starts_with("cuda_device::float::")) + .expect("scalar-arithmetic compatibility path"); + let ty = scalar_arithmetic_rust_type(record); + let arity = scalar_arithmetic_arity(record); + let parameters = (0..arity) + .map(|index| format!("arg{index}: {ty}")) + .collect::>() + .join(", "); + let arguments = (0..arity) + .map(|index| format!("arg{index}")) + .collect::>() + .join(", "); + writeln!(output, "/// {}", record.summary).unwrap(); + output.push_str("#[must_use]\n#[inline(never)]\n"); + writeln!( + output, + "pub fn {}({parameters}) -> {ty} {{", + record.rust.name + ) + .unwrap(); + writeln!(output, " let _ = ({arguments});").unwrap(); + writeln!( + output, + " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" + ) + .unwrap(); + output.push_str("}\n\n"); + } + for record in scalar_maths(catalog) { + let path = record + .rust + .compatibility_paths + .iter() + .find(|path| path.starts_with("cuda_device::float::")) + .expect("scalar-math compatibility path"); + let ty = match scalar_math_contract(record).format { + ScalarMathFormat::F16 => "u16", + ScalarMathFormat::F32 => "f32", + ScalarMathFormat::F64 => "f64", + }; + writeln!(output, "/// {}", record.summary).unwrap(); + output.push_str("#[must_use]\n#[inline(never)]\n"); + writeln!(output, "pub fn {}(arg0: {ty}) -> {ty} {{", record.rust.name).unwrap(); + writeln!(output, " let _ = arg0;").unwrap(); + writeln!( + output, + " unreachable!(\"generated CUDA intrinsic `{path}` executed outside device compilation\")" + ) + .unwrap(); + output.push_str("}\n\n"); + } + render_compat_extended_minmax_into(&mut output, catalog, "float"); + output +} + +pub(super) fn render_compat_float_output( + catalog: &CatalogFile, + hash: &str, +) -> Option<(PathBuf, String)> { + (scalar_arithmetics(catalog).next().is_some() + || extended_minmax(catalog).next().is_some() + || scalar_maths(catalog).next().is_some()) + .then(|| { + ( + "crates/cuda-device/src/generated/float.rs".into(), + render_compat_float(catalog, hash), + ) + }) +} diff --git a/crates/cuda-intrinsics-gen/src/render/dialect/cluster_tensor.rs b/crates/cuda-intrinsics-gen/src/render/dialect/cluster_tensor.rs new file mode 100644 index 0000000000..0a66d7d299 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/render/dialect/cluster_tensor.rs @@ -0,0 +1,903 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{CatalogFile, Tcgen05Adapter}; +use crate::render::common::rust_header; +use crate::render::families::{ + clc_intrinsics, cluster_barriers, cluster_memory, debug_controls, execution_control_family, + execution_controls, render_tcgen05_carrier_runs, tcgen05_intrinsics, tcgen05_mma_intrinsics, + tcgen05_non_mma_intrinsics, tma_intrinsics, wgmma_controls, +}; +use std::fmt::Write as _; + +pub(in crate::render) fn render_dialect_cluster_barrier( + catalog: &CatalogFile, + hash: &str, +) -> String { + let mut output = rust_header(catalog, hash); + if cluster_barriers(catalog).next().is_none() { + return output; + } + output.push_str( + r##"//! Generated operations for the closed cluster-barrier family. + +use pliron::{ + attribute::Attribute, + builtin::op_interfaces::{NOpdsInterface, NResultsInterface}, + common_traits::Verify, + context::{Context, Ptr}, + location::Located, + op::Op, + operation::Operation, + result::Error, + verify_err, +}; +use pliron_derive::{pliron_attr, pliron_op}; + +#[pliron_attr(name = "nvvm.cluster_barrier_mode", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum ClusterBarrierModeAttr { + Arrive, + ArriveAligned, + ArriveRelaxed, + ArriveRelaxedAligned, + Wait, + WaitAligned, +} + +/// Cluster synchronization whose exact spelling is carried by an attribute. +#[pliron_op( + name = "nvvm.cluster_barrier", + format, + interfaces = [NOpdsInterface<0>, NResultsInterface<0>], + attributes = (nvvm_cluster_barrier_mode: ClusterBarrierModeAttr) +)] +pub struct ClusterBarrierOp; + +impl ClusterBarrierOp { + pub fn new(op: Ptr) -> Self { + Self { op } + } + + pub fn build(ctx: &mut Context, mode: ClusterBarrierModeAttr) -> Ptr { + let op = Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![], vec![], 0); + let this = Self { op }; + this.set_attr_nvvm_cluster_barrier_mode(ctx, mode); + this.get_operation() + } +} + +impl Verify for ClusterBarrierOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if self.get_attr_nvvm_cluster_barrier_mode(ctx).is_none() { + return verify_err!(op.loc(), "nvvm.cluster_barrier requires a mode attribute"); + } + if op.get_num_operands() != 0 || op.get_num_results() != 0 { + return verify_err!(op.loc(), "nvvm.cluster_barrier takes no operands or results"); + } + Ok(()) + } +} + +/// Compatibility operation for a complete aligned cluster synchronization. +#[pliron_op( + name = "nvvm.cluster_sync", + format, + verifier = "succ", + interfaces = [NOpdsInterface<0>, NResultsInterface<0>], +)] +pub struct ClusterSyncOp; + +impl ClusterSyncOp { + pub fn new(op: Ptr) -> Self { + Self { op } + } +} + +pub(super) fn register(ctx: &mut Context) { + ClusterBarrierModeAttr::register(ctx); + ClusterBarrierOp::register(ctx); + ClusterSyncOp::register(ctx); +} +"##, + ); + output +} + +pub(in crate::render) fn render_dialect_cluster_memory( + catalog: &CatalogFile, + hash: &str, +) -> String { + assert_eq!(cluster_memory(catalog).count(), 2); + let mut output = rust_header(catalog, hash); + output.push_str( + r#"//! Structural operations for cluster address mapping and remote shared reads. + +use dialect_mir::types::{MirPointerKind, MirPtrType, address_space}; +use pliron::{ + builtin::{ + op_interfaces::{NOpdsInterface, NResultsInterface}, + types::{IntegerType, Signedness}, + }, + 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; + +fn verify_pointer_rank( + ctx: &Context, + operation: Ptr, + name: &str, +) -> Result<(), Error> { + let op = operation.deref(ctx); + if op.get_num_operands() != 2 || op.get_num_results() != 1 { + return verify_err!(op.loc(), "{name} requires two operands and one result"); + } + if op + .get_operand(0) + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_none() + { + return verify_err!(op.loc(), "{name} source must be a MIR pointer"); + } + let rank_ty = op.get_operand(1).get_type(ctx); + let rank_ty = rank_ty.deref(ctx); + if !rank_ty + .downcast_ref::() + .is_some_and(|ty| ty.width() == 32 && ty.signedness() == Signedness::Unsigned) + { + return verify_err!(op.loc(), "{name} rank must be u32"); + } + Ok(()) +} + +#[pliron_op( + name = "nvvm.mapa_shared_cluster", + format, + interfaces = [NOpdsInterface<2>, NResultsInterface<1>], +)] +pub struct MapaSharedClusterOp; + +impl MapaSharedClusterOp { + pub fn new(op: Ptr) -> Self { Self { op } } + + // Kind-preserving addrspace retype: copies the source pointer's kind + // verbatim; concrete kinds are minted only in mir-importer's facts.rs. + #[allow(clippy::disallowed_methods)] + pub fn build(ctx: &mut Context, source: Value, rank: Value) -> Ptr { + let source_ty = source.get_type(ctx); + let (pointee, is_mutable, pointer_kind) = { + let source_ty_obj = source_ty.deref(ctx); + let source_ptr = source_ty_obj + .downcast_ref::() + .expect("nvvm.mapa_shared_cluster source must be a MIR pointer"); + ( + source_ptr.pointee, + source_ptr.is_mutable(), + source_ptr.pointer_kind(), + ) + }; + let result_ty = MirPtrType::get_with_kind( + ctx, + pointee, + is_mutable, + address_space::CLUSTER_SHARED, + pointer_kind, + ); + Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![result_ty.into()], + vec![source, rank], + vec![], + 0, + ) + } +} + +impl Verify for MapaSharedClusterOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let operation = self.get_operation(); + verify_pointer_rank(ctx, operation, "nvvm.mapa_shared_cluster")?; + let op = operation.deref(ctx); + let source_ty = op.get_operand(0).get_type(ctx); + let result_ty = op.get_result(0).get_type(ctx); + let source_ty_obj = source_ty.deref(ctx); + let result_ty_obj = result_ty.deref(ctx); + let source_ptr = source_ty_obj.downcast_ref::().unwrap(); + let Some(result_ptr) = result_ty_obj.downcast_ref::() else { + return verify_err!( + op.loc(), + "nvvm.mapa_shared_cluster result must be a MIR pointer" + ); + }; + if !matches!( + source_ptr.pointer_kind(), + MirPointerKind::RawConst | MirPointerKind::RawMut + ) + || result_ptr.pointee != source_ptr.pointee + || result_ptr.is_mutable() != source_ptr.is_mutable() + || !result_ptr.is_cluster_shared() + || result_ptr.pointer_kind() != source_ptr.pointer_kind() + { + return verify_err!( + op.loc(), + "nvvm.mapa_shared_cluster requires a raw source pointer and must preserve its pointee, mutability, and raw kind while returning addrspace(7)" + ); + } + Ok(()) + } +} + +#[pliron_op( + name = "nvvm.dsmem_read_u32", + format, + interfaces = [NOpdsInterface<2>, NResultsInterface<1>], +)] +pub struct DsmemReadU32Op; + +impl DsmemReadU32Op { + pub fn new(op: Ptr) -> Self { Self { op } } + + pub fn build(ctx: &mut Context, source: Value, rank: Value) -> Ptr { + let result_ty = IntegerType::get(ctx, 32, Signedness::Unsigned); + Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![result_ty.to_handle()], + vec![source, rank], + vec![], + 0, + ) + } +} + +impl Verify for DsmemReadU32Op { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let operation = self.get_operation(); + verify_pointer_rank(ctx, operation, "nvvm.dsmem_read_u32")?; + let op = operation.deref(ctx); + let result_ty = op.get_result(0).get_type(ctx); + let result_ty = result_ty.deref(ctx); + if !result_ty + .downcast_ref::() + .is_some_and(|ty| ty.width() == 32 && ty.signedness() == Signedness::Unsigned) + { + return verify_err!(op.loc(), "nvvm.dsmem_read_u32 result must be u32"); + } + Ok(()) + } +} + +pub(super) fn register(ctx: &mut Context) { + MapaSharedClusterOp::register(ctx); + DsmemReadU32Op::register(ctx); +} +"#, + ); + output +} + +pub(in crate::render) fn render_dialect_wgmma_control(catalog: &CatalogFile, hash: &str) -> String { + assert_eq!(wgmma_controls(catalog).count(), 3); + let mut output = rust_header(catalog, hash); + output.push_str( + r##"//! Structural operations for generated WGMMA controls. + +use pliron::{ + builtin::{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; + +#[pliron_op( + name = "nvvm.wgmma_fence_sync_aligned", + format, + verifier = "succ", + interfaces = [NOpdsInterface<0>, NResultsInterface<0>], +)] +pub struct WgmmaFenceSyncAlignedOp; + +impl WgmmaFenceSyncAlignedOp { + 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) + } +} + +#[pliron_op( + name = "nvvm.wgmma_commit_group_sync_aligned", + format, + verifier = "succ", + interfaces = [NOpdsInterface<0>, NResultsInterface<0>], +)] +pub struct WgmmaCommitGroupSyncAlignedOp; + +impl WgmmaCommitGroupSyncAlignedOp { + 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) + } +} + +#[pliron_op( + name = "nvvm.wgmma_wait_group_sync_aligned", + format, + interfaces = [NOpdsInterface<1>, NResultsInterface<0>], +)] +pub struct WgmmaWaitGroupSyncAlignedOp; + +impl WgmmaWaitGroupSyncAlignedOp { + pub fn new(op: Ptr) -> Self { Self { op } } + + pub fn build(ctx: &mut Context, max_pending: Value) -> Ptr { + Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![], + vec![max_pending], + vec![], + 0, + ) + } +} + +impl Verify for WgmmaWaitGroupSyncAlignedOp { + 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(), "nvvm.wgmma_wait_group_sync_aligned takes one operand and no results"); + } + if !is_i64(ctx, op.get_operand(0).get_type(ctx)) { + return verify_err!(op.loc(), "nvvm.wgmma_wait_group_sync_aligned operand must be i64"); + } + Ok(()) + } +} + +fn is_i64(ctx: &Context, ty: TypeHandle) -> bool { + ty.deref(ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 64) +} + +pub(super) fn register(ctx: &mut Context) { + WgmmaFenceSyncAlignedOp::register(ctx); + WgmmaCommitGroupSyncAlignedOp::register(ctx); + WgmmaWaitGroupSyncAlignedOp::register(ctx); +} +"##, + ); + output +} + +pub(in crate::render) fn render_dialect_debug_control(catalog: &CatalogFile, hash: &str) -> String { + assert_eq!(debug_controls(catalog).count(), 3); + let mut output = rust_header(catalog, hash); + output.push_str( + "//! Structural operations for generated PTX debug controls.\n\n\ + use pliron::{\n\ + builtin::{\n\ + attributes::IntegerAttr,\n\ + op_interfaces::{NOpdsInterface, NResultsInterface},\n\ + types::{IntegerType, Signedness},\n\ + },\n\ + common_traits::Verify,\n\ + context::{Context, Ptr},\n\ + identifier::Identifier,\n\ + location::Located,\n\ + op::Op,\n\ + operation::Operation,\n\ + result::Error,\n\ + verify_err,\n\ + };\n\ + use pliron::utils::apint::APInt;\n\ + use pliron_derive::pliron_op;\n\ + use std::num::NonZeroUsize;\n\n\ + #[pliron_op(\n\ + name = \"nvvm.trap\",\n\ + format,\n\ + verifier = \"succ\",\n\ + interfaces = [NOpdsInterface<0>, NResultsInterface<0>],\n\ + )]\n\ + pub struct TrapOp;\n\n\ + impl TrapOp {\n\ + pub fn new(op: Ptr) -> Self { Self { op } }\n\ + pub fn build(ctx: &mut Context) -> Ptr {\n\ + Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![], vec![], 0)\n\ + }\n\ + }\n\n\ + #[pliron_op(\n\ + name = \"nvvm.brkpt\",\n\ + format,\n\ + verifier = \"succ\",\n\ + interfaces = [NOpdsInterface<0>, NResultsInterface<0>],\n\ + )]\n\ + pub struct BreakpointOp;\n\n\ + impl BreakpointOp {\n\ + pub fn new(op: Ptr) -> Self { Self { op } }\n\ + pub fn build(ctx: &mut Context) -> Ptr {\n\ + Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![], vec![], 0)\n\ + }\n\ + }\n\n\ + #[pliron_op(\n\ + name = \"nvvm.pmevent\",\n\ + format,\n\ + interfaces = [NOpdsInterface<0>, NResultsInterface<0>],\n\ + )]\n\ + pub struct PmEventOp;\n\n\ + impl PmEventOp {\n\ + pub fn new(op: Ptr) -> Self { Self { op } }\n\n\ + pub fn build(ctx: &mut Context, event_id: u32) -> Ptr {\n\ + let op = Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![], vec![], 0);\n\ + let ty = IntegerType::get(ctx, 32, Signedness::Unsigned);\n\ + let value = APInt::from_u64(event_id.into(), NonZeroUsize::new(32).unwrap());\n\ + op.deref_mut(ctx).attributes.set(\n\ + Identifier::try_from(\"event_id\").unwrap(),\n\ + IntegerAttr::new(ty, value),\n\ + );\n\ + op\n\ + }\n\n\ + pub fn new_with_event_id(ctx: &mut Context, event_id: u32) -> Ptr {\n\ + Self::build(ctx, event_id)\n\ + }\n\n\ + pub fn event_id(&self, ctx: &Context) -> Option {\n\ + let key = Identifier::try_from(\"event_id\").unwrap();\n\ + let operation = self.get_operation().deref(ctx);\n\ + let attribute: &IntegerAttr = operation.attributes.get(&key)?;\n\ + let ty_handle = attribute.get_type();\n\ + let ty = ty_handle.deref(ctx);\n\ + if ty.width() != 32 || ty.signedness() != Signedness::Unsigned {\n\ + return None;\n\ + }\n\ + u32::try_from(attribute.value().to_u64()).ok().filter(|value| *value <= 15)\n\ + }\n\ +\n\ + pub fn get_event_id(&self, ctx: &Context) -> Option {\n\ + self.event_id(ctx)\n\ + }\n\ + }\n\n\ + impl Verify for PmEventOp {\n\ + fn verify(&self, ctx: &Context) -> Result<(), Error> {\n\ + let op = self.get_operation().deref(ctx);\n\ + if self.event_id(ctx).is_none() {\n\ + return verify_err!(op.loc(), \"nvvm.pmevent requires a u32 event ID in 0..=15\");\n\ + }\n\ + Ok(())\n\ + }\n\ + }\n\n\ + pub(super) fn register(ctx: &mut Context) {\n\ + TrapOp::register(ctx);\n\ + BreakpointOp::register(ctx);\n\ + PmEventOp::register(ctx);\n\ + }\n", + ); + output +} + +pub(in crate::render) fn render_dialect_clc(catalog: &CatalogFile, hash: &str) -> String { + assert_eq!(clc_intrinsics(catalog).count(), 6); + let mut output = rust_header(catalog, hash); + output.push_str( + "//! Generated Cluster Launch Control operations.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::IntegerType,\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::{TypeHandle, Typed},\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\n", + ); + for record in clc_intrinsics(catalog) { + let result_count = record.dialect.results.len(); + let verifier = if result_count == 0 { + " verifier = \"succ\",\n" + } else { + "" + }; + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!( + output, + "#[pliron_op(\n name = {:?},\n format,\n{verifier} interfaces = [NOpdsInterface<2>, NResultsInterface<{result_count}>],\n)]", + record.dialect.op_name, + ) + .unwrap(); + writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); + writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); + output.push_str(" pub fn new(op: Ptr) -> Self { Self { op } }\n}\n\n"); + if result_count == 1 { + writeln!(output, "impl Verify for {} {{", record.dialect.op_type).unwrap(); + writeln!( + output, + " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n verify_query_shape(ctx, self.get_operation(), {:?})\n }}\n}}\n", + record.dialect.op_name, + ) + .unwrap(); + } + } + output.push_str( + "fn verify_query_shape(\n ctx: &Context,\n operation: Ptr,\n name: &str,\n) -> Result<(), Error> {\n let op = operation.deref(ctx);\n let valid = op.get_num_operands() == 2\n && op.get_num_results() == 1\n && is_integer_width(ctx, op.get_operand(0).get_type(ctx), 64)\n && is_integer_width(ctx, op.get_operand(1).get_type(ctx), 64)\n && is_integer_width(ctx, op.get_result(0).get_type(ctx), 32);\n if !valid {\n return verify_err!(\n op.loc(),\n \"{} requires two i64 operands and one i32 result\",\n name\n );\n }\n Ok(())\n}\n\nfn is_integer_width(ctx: &Context, ty: TypeHandle, width: u32) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == width)\n}\n\npub(super) fn register(ctx: &mut Context) {\n", + ); + for record in clc_intrinsics(catalog) { + writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); + } + output.push_str("}\n"); + output +} + +pub(in crate::render) fn render_dialect_execution_control( + catalog: &CatalogFile, + hash: &str, +) -> String { + assert_eq!(execution_controls(catalog).count(), 8); + let mut output = rust_header(catalog, hash); + output.push_str( + "//! Generated counted-barrier, grid-dependency, and register-control operations.\n\nuse pliron::{\n builtin::{attributes::IntegerAttr, op_interfaces::{NOpdsInterface, NResultsInterface}, types::{IntegerType, Signedness}},\n common_traits::Verify,\n context::{Context, Ptr},\n identifier::Identifier,\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n verify_err,\n};\nuse pliron::utils::apint::APInt;\nuse pliron_derive::pliron_op;\nuse std::num::NonZeroUsize;\n\n", + ); + for record in execution_controls(catalog).filter(|record| record.family != "register_control") { + let operand_count = record.dialect.operands.len(); + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!( + output, + "#[pliron_op(\n name = {:?},\n format,\n verifier = \"succ\",\n interfaces = [NOpdsInterface<{operand_count}>, NResultsInterface<0>],\n)]", + record.dialect.op_name + ) + .unwrap(); + writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); + writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); + output.push_str(" pub fn new(op: Ptr) -> Self { Self { op } }\n}\n\n"); + } + for record in execution_control_family(catalog, "register_control") { + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!( + output, + "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<0>, NResultsInterface<0>],\n)]", + record.dialect.op_name + ) + .unwrap(); + writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); + writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); + output.push_str( + " pub fn new(op: Ptr) -> Self { Self { op } }\n\n pub fn build(ctx: &mut Context, register_count: u32) -> Ptr {\n let op = Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![], vec![], 0);\n let ty = IntegerType::get(ctx, 32, Signedness::Unsigned);\n let value = APInt::from_u64(register_count.into(), NonZeroUsize::new(32).unwrap());\n op.deref_mut(ctx).attributes.set(\n Identifier::try_from(\"register_count\").unwrap(),\n IntegerAttr::new(ty, value),\n );\n op\n }\n\n pub fn register_count(&self, ctx: &Context) -> Option {\n let key = Identifier::try_from(\"register_count\").unwrap();\n let operation = self.get_operation().deref(ctx);\n let attribute: &IntegerAttr = operation.attributes.get(&key)?;\n let ty_handle = attribute.get_type();\n let ty = ty_handle.deref(ctx);\n if ty.width() != 32 || ty.signedness() != Signedness::Unsigned {\n return None;\n }\n u32::try_from(attribute.value().to_u64()).ok().filter(|value| (24..=256).contains(value) && value % 8 == 0)\n }\n}\n\n", + ); + writeln!(output, "impl Verify for {} {{", record.dialect.op_type).unwrap(); + output.push_str( + " fn verify(&self, ctx: &Context) -> Result<(), Error> {\n let op = self.get_operation().deref(ctx);\n if self.register_count(ctx).is_none() {\n return verify_err!(op.loc(), \"setmaxnreg requires an immediate register count in 24..=256 divisible by 8\");\n }\n Ok(())\n }\n}\n\n", + ); + } + output.push_str("pub(super) fn register(ctx: &mut Context) {\n"); + for record in execution_controls(catalog) { + writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); + } + output.push_str("}\n"); + output +} + +pub(in crate::render) fn render_dialect_tma(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + "//! Generated Tensor Memory Accelerator operations.\n\nuse pliron::{\n builtin::op_interfaces::{NOpdsInterface, NResultsInterface},\n context::{Context, Ptr},\n op::Op,\n operation::Operation,\n};\nuse pliron_derive::pliron_op;\n\n", + ); + for record in tma_intrinsics(catalog) { + let operand_count = record.dialect.operands.len(); + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!( + output, + "#[pliron_op(\n name = {:?},\n format,\n verifier = \"succ\",\n interfaces = [NOpdsInterface<{operand_count}>, NResultsInterface<0>],\n)]", + record.dialect.op_name + ) + .unwrap(); + writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); + writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); + output.push_str(" pub fn new(op: Ptr) -> Self { Self { op } }\n}\n\n"); + } + output.push_str("pub(super) fn register(ctx: &mut Context) {\n"); + for record in tma_intrinsics(catalog) { + writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); + } + output.push_str("}\n"); + output +} + +pub(in crate::render) fn render_dialect_tcgen05(catalog: &CatalogFile, hash: &str) -> String { + assert_eq!(tcgen05_intrinsics(catalog).count(), 233); + let mut output = rust_header(catalog, hash); + output.push_str( + r#"//! Generated Tensor Core Generation 5 operations. + +use dialect_mir::{ops::MirConstantOp, types::MirPtrType}; +use pliron::{ + attribute::Attribute, + builtin::{ + op_interfaces::{NOpdsInterface, NResultsInterface}, + ops::ConstantOp, + types::{FP32Type, IntegerType}, + }, + common_traits::Verify, + context::{Context, Ptr}, + location::Located, + op::Op, + operation::Operation, + result::Error, + r#type::{TypeHandle, Typed}, + verify_err, +}; +use pliron_derive::{pliron_attr, pliron_op}; + +#[derive(Clone, Copy)] +enum Tcgen05Carrier { + Ptr, + I1, + I16, + I32, + I64, + F32, +} + +impl Tcgen05Carrier { + fn matches(self, ctx: &Context, ty: TypeHandle) -> bool { + match self { + Self::Ptr => ty.deref(ctx).downcast_ref::().is_some(), + Self::I1 => is_integer_width(ctx, ty, 1), + Self::I16 => is_integer_width(ctx, ty, 16), + Self::I32 => is_integer_width(ctx, ty, 32), + Self::I64 => is_integer_width(ctx, ty, 64), + Self::F32 => ty.deref(ctx).downcast_ref::().is_some(), + } + } + + fn name(self) -> &'static str { + match self { + Self::Ptr => "MIR pointer", + Self::I1 => "i1", + Self::I16 => "i16", + Self::I32 => "i32", + Self::I64 => "i64", + Self::F32 => "f32", + } + } +} + +fn is_integer_width(ctx: &Context, ty: TypeHandle, width: u32) -> bool { + ty.deref(ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == width) +} + +fn expected_carrier_count(runs: &[(Tcgen05Carrier, usize)]) -> usize { + runs.iter().map(|(_, count)| *count).sum() +} + +fn verify_tcgen05_signature( + ctx: &Context, + operation: Ptr, + name: &str, + operand_runs: &[(Tcgen05Carrier, usize)], + result_runs: &[(Tcgen05Carrier, usize)], + constant_operand: Option, +) -> Result<(), Error> { + let op = operation.deref(ctx); + let expected_operands = expected_carrier_count(operand_runs); + let expected_results = expected_carrier_count(result_runs); + if op.get_num_operands() != expected_operands || op.get_num_results() != expected_results { + return verify_err!( + op.loc(), + "{name} requires {expected_operands} operands and {expected_results} results" + ); + } + + let mut index = 0; + for (carrier, count) in operand_runs { + for _ in 0..*count { + if !carrier.matches(ctx, op.get_operand(index).get_type(ctx)) { + return verify_err!( + op.loc(), + "{name} operand {index} must be {}", + carrier.name() + ); + } + index += 1; + } + } + + index = 0; + for (carrier, count) in result_runs { + for _ in 0..*count { + if !carrier.matches(ctx, op.get_result(index).get_type(ctx)) { + return verify_err!( + op.loc(), + "{name} result {index} must be {}", + carrier.name() + ); + } + index += 1; + } + } + + if let Some(index) = constant_operand { + let value = op.get_operand(index); + let is_constant = value.defining_op().is_some_and(|defining_op| { + Operation::get_op::(defining_op, ctx).is_some() + || Operation::get_op::(defining_op, ctx).is_some() + }); + if !is_constant { + return verify_err!( + op.loc(), + "{name} operand {index} must be a compile-time constant" + ); + } + } + + Ok(()) +} + +"#, + ); + if tcgen05_mma_intrinsics(catalog).next().is_some() { + output.push_str( + r#"#[pliron_attr(name = "nvvm.tcgen05_mma_form", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum Tcgen05MmaFormAttr { + Shared, Tensor, TensorAshift, SpShared, SpTensor, SpTensorAshift, + WsShared, WsSharedZeroColMask, WsSpShared, WsSpSharedZeroColMask, + WsSpTensor, WsSpTensorZeroColMask, WsTensor, WsTensorZeroColMask, +} + +#[pliron_attr(name = "nvvm.tcgen05_mma_kind", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum Tcgen05MmaKindAttr { F16, Tf32, F8f6f4, I8 } + +#[pliron_attr(name = "nvvm.tcgen05_mma_cta_group", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum Tcgen05MmaCtaGroupAttr { Cg1, Cg2 } + +#[pliron_attr(name = "nvvm.tcgen05_mma_collector_a", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum Tcgen05MmaCollectorAAttr { Discard, LastUse, Fill, Use } + +#[pliron_attr(name = "nvvm.tcgen05_mma_b_buffer", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum Tcgen05MmaBBufferAttr { B0, B1, B2, B3 } + +#[pliron_attr(name = "nvvm.tcgen05_mma_b_usage", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum Tcgen05MmaBUsageAttr { Discard, LastUse, Fill, Use } + +#[pliron_op( + name = "nvvm.tcgen05_mma", + format, + attributes = ( + nvvm_tcgen05_mma_form: Tcgen05MmaFormAttr, + nvvm_tcgen05_mma_kind: Tcgen05MmaKindAttr, + nvvm_tcgen05_mma_cta_group: Tcgen05MmaCtaGroupAttr, + nvvm_tcgen05_mma_collector_a: Tcgen05MmaCollectorAAttr, + nvvm_tcgen05_mma_b_buffer: Tcgen05MmaBBufferAttr, + nvvm_tcgen05_mma_b_usage: Tcgen05MmaBUsageAttr + ) +)] +pub struct Tcgen05MmaOp; + +impl Tcgen05MmaOp { + pub fn new(op: Ptr) -> Self { Self { op } } +} + +impl Verify for Tcgen05MmaOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + use Tcgen05Carrier::{I1, I32, I64}; + let op = self.get_operation(); + let form = self.get_attr_nvvm_tcgen05_mma_form(ctx); + let Some(form) = form.as_deref() else { + return verify_err!(op.deref(ctx).loc(), "nvvm.tcgen05_mma requires a form"); + }; + if self.get_attr_nvvm_tcgen05_mma_kind(ctx).is_none() { + return verify_err!(op.deref(ctx).loc(), "nvvm.tcgen05_mma requires a kind"); + } + let (operands, base): (&[(Tcgen05Carrier, usize)], bool) = match form { + Tcgen05MmaFormAttr::Shared | Tcgen05MmaFormAttr::WsShared => + (&[(I32, 1), (I64, 2), (I32, 1), (I1, 1)], matches!(form, Tcgen05MmaFormAttr::Shared)), + Tcgen05MmaFormAttr::Tensor | Tcgen05MmaFormAttr::TensorAshift | Tcgen05MmaFormAttr::WsTensor => + (&[(I32, 2), (I64, 1), (I32, 1), (I1, 1)], !matches!(form, Tcgen05MmaFormAttr::WsTensor)), + Tcgen05MmaFormAttr::SpShared | Tcgen05MmaFormAttr::WsSpShared => + (&[(I32, 1), (I64, 2), (I32, 1), (I1, 1), (I32, 1)], matches!(form, Tcgen05MmaFormAttr::SpShared)), + Tcgen05MmaFormAttr::SpTensor | Tcgen05MmaFormAttr::SpTensorAshift | Tcgen05MmaFormAttr::WsSpTensor => + (&[(I32, 2), (I64, 1), (I32, 1), (I1, 1), (I32, 1)], !matches!(form, Tcgen05MmaFormAttr::WsSpTensor)), + Tcgen05MmaFormAttr::WsSharedZeroColMask => + (&[(I32, 1), (I64, 2), (I32, 1), (I1, 1), (I64, 1)], false), + Tcgen05MmaFormAttr::WsSpSharedZeroColMask => + (&[(I32, 1), (I64, 2), (I32, 1), (I1, 1), (I32, 1), (I64, 1)], false), + Tcgen05MmaFormAttr::WsSpTensorZeroColMask => + (&[(I32, 2), (I64, 1), (I32, 1), (I1, 1), (I32, 1), (I64, 1)], false), + Tcgen05MmaFormAttr::WsTensorZeroColMask => + (&[(I32, 2), (I64, 1), (I32, 1), (I1, 1), (I64, 1)], false), + }; + verify_tcgen05_signature(ctx, op, "nvvm.tcgen05_mma", operands, &[], None)?; + + let cta_group = self.get_attr_nvvm_tcgen05_mma_cta_group(ctx); + let collector_a = self.get_attr_nvvm_tcgen05_mma_collector_a(ctx); + let b_buffer = self.get_attr_nvvm_tcgen05_mma_b_buffer(ctx); + let b_usage = self.get_attr_nvvm_tcgen05_mma_b_usage(ctx); + if base { + if cta_group.is_none() || collector_a.is_none() || b_buffer.is_some() || b_usage.is_some() { + return verify_err!(op.deref(ctx).loc(), "base tcgen05 MMA requires CTA-group and collector-A selectors"); + } + if matches!(form, Tcgen05MmaFormAttr::TensorAshift | Tcgen05MmaFormAttr::SpTensorAshift) + && !matches!(collector_a.as_deref(), Some(Tcgen05MmaCollectorAAttr::Discard | Tcgen05MmaCollectorAAttr::LastUse)) + { + return verify_err!(op.deref(ctx).loc(), "ashift tcgen05 MMA only supports discard or last-use collector-A"); + } + } else if cta_group.is_some() || collector_a.is_some() || b_buffer.is_none() || b_usage.is_none() { + return verify_err!(op.deref(ctx).loc(), "warp-specialized tcgen05 MMA requires B-buffer and B-usage selectors"); + } + Ok(()) + } +} + +"#, + ); + } + for record in tcgen05_non_mma_intrinsics(catalog) { + let operand_count = record.dialect.operands.len(); + let result_count = record.dialect.results.len(); + let operand_runs = render_tcgen05_carrier_runs(&record.dialect.operands); + let result_runs = render_tcgen05_carrier_runs(&record.dialect.results); + let constant_operand = record + .tcgen05 + .as_ref() + .is_some_and(|tcgen05| { + matches!( + tcgen05.adapter, + Tcgen05Adapter::TmemHalfSplitOffsetInjectPack16ToU32Registers + | Tcgen05Adapter::TmemHalfSplitOffsetU32RegistersInjectUnpack16ToVoid + ) + }) + .then_some(1); + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!( + output, + "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<{operand_count}>, NResultsInterface<{result_count}>],\n)]", + record.dialect.op_name + ) + .unwrap(); + writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); + writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); + output.push_str(" pub fn new(op: Ptr) -> Self { Self { op } }\n}\n\n"); + writeln!(output, "impl Verify for {} {{", record.dialect.op_type).unwrap(); + output.push_str(" fn verify(&self, ctx: &Context) -> Result<(), Error> {\n"); + output.push_str(" verify_tcgen05_signature(\n"); + output.push_str(" ctx,\n"); + output.push_str(" self.get_operation(),\n"); + writeln!(output, " {:?},", record.dialect.op_name).unwrap(); + writeln!(output, " {operand_runs},").unwrap(); + writeln!(output, " {result_runs},").unwrap(); + writeln!(output, " {constant_operand:?},").unwrap(); + output.push_str(" )\n"); + output.push_str(" }\n"); + output.push_str("}\n\n"); + } + output.push_str("pub(super) fn register(ctx: &mut Context) {\n"); + if tcgen05_mma_intrinsics(catalog).next().is_some() { + output.push_str( + " Tcgen05MmaFormAttr::register(ctx);\n Tcgen05MmaKindAttr::register(ctx);\n Tcgen05MmaCtaGroupAttr::register(ctx);\n Tcgen05MmaCollectorAAttr::register(ctx);\n Tcgen05MmaBBufferAttr::register(ctx);\n Tcgen05MmaBUsageAttr::register(ctx);\n Tcgen05MmaOp::register(ctx);\n", + ); + } + for record in tcgen05_non_mma_intrinsics(catalog) { + writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); + } + output.push_str("}\n"); + output +} diff --git a/crates/cuda-intrinsics-gen/src/render/dialect/mma.rs b/crates/cuda-intrinsics-gen/src/render/dialect/mma.rs new file mode 100644 index 0000000000..bd793ae84d --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/render/dialect/mma.rs @@ -0,0 +1,932 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::CatalogFile; +use crate::render::common::rust_header; +use crate::render::families::{ + ldmatrix, ldmatrix_compat_op, movmatrix, register_mma_attr_variants, register_mma_carriers, + register_mma_compat_op_type, register_mmas, sparse_mma_attr_variants, sparse_mma_carriers, + sparse_mma_selector_values, sparse_mmas, stmatrices, stmatrix_variant, +}; +use std::fmt::Write as _; + +pub(in crate::render) fn render_dialect_movmatrix(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + if movmatrix(catalog).next().is_none() { + output.push_str( + "//! Structural operation for generated in-register matrix transpose.\n\nuse pliron::context::Context;\n\npub(super) fn register(_ctx: &mut Context) {}\n", + ); + return output; + } + let record = movmatrix(catalog).next().unwrap(); + output.push_str( + "//! Structural operation for generated in-register matrix transpose.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_i32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32)\n}\n\n", + ); + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!( + output, + "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<1>, NResultsInterface<1>],\n)]\npub struct {};\n", + record.dialect.op_name, record.dialect.op_type + ) + .unwrap(); + writeln!( + output, + "impl {} {{\n pub fn new(op: Ptr) -> Self {{\n Self {{ op }}\n }}\n\n pub fn build(ctx: &mut Context, value: Value) -> Ptr {{\n let result_ty = IntegerType::get(ctx, 32, Signedness::Unsigned);\n Operation::new(\n ctx,\n Self::get_concrete_op_info(),\n vec![result_ty.into()],\n vec![value],\n vec![],\n 0,\n )\n }}\n}}\n", + record.dialect.op_type + ) + .unwrap(); + writeln!( + output, + "impl Verify for {} {{\n fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n let op = self.get_operation().deref(ctx);\n if op.get_num_operands() != 1 || op.get_num_results() != 1 {{\n return verify_err!(op.loc(), {:?});\n }}\n if !is_i32(ctx, op.get_operand(0).get_type(ctx))\n || !is_i32(ctx, op.get_result(0).get_type(ctx))\n {{\n return verify_err!(op.loc(), {:?});\n }}\n Ok(())\n }}\n}}\n\npub(super) fn register(ctx: &mut Context) {{\n {}::register(ctx);\n}}\n", + record.dialect.op_type, + format!("{} requires one operand and one result", record.dialect.op_name), + format!( + "{} operand and result must be 32-bit integers", + record.dialect.op_name + ), + record.dialect.op_type, + ) + .unwrap(); + output +} + +pub(in crate::render) fn render_dialect_ldmatrix(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + r##"//! Structural operation and compatibility carriers for generated `ldmatrix`. + +use dialect_mir::types::{address_space, MirPtrType}; +use pliron::{ + attribute::Attribute, + builtin::{ + op_interfaces::{NOpdsInterface, NResultsInterface}, + types::{IntegerType, Signedness}, + }, + 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_attr, pliron_op}; + +#[pliron_attr(name = "nvvm.ldmatrix_shape", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum LdmatrixShapeAttr { + M8n8, + M8n16, + M16n16, +} + +impl LdmatrixShapeAttr { + pub const fn register_count(&self, multiplicity: &LdmatrixMultiplicityAttr) -> usize { + let matrices = multiplicity.register_count(); + match self { + Self::M8n8 | Self::M8n16 => matrices, + Self::M16n16 => matrices * 2, + } + } +} + +#[pliron_attr(name = "nvvm.ldmatrix_multiplicity", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum LdmatrixMultiplicityAttr { + X1, + X2, + X4, +} + +impl LdmatrixMultiplicityAttr { + pub const fn register_count(&self) -> usize { + match self { + Self::X1 => 1, + Self::X2 => 2, + Self::X4 => 4, + } + } +} + +#[pliron_attr(name = "nvvm.ldmatrix_layout", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum LdmatrixLayoutAttr { + Normal, + Transposed, +} + +#[pliron_attr(name = "nvvm.ldmatrix_element", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum LdmatrixElementAttr { + B16, + B8, + B8x16B4x16P64, + B8x16B6x16P32, +} + +#[pliron_attr(name = "nvvm.ldmatrix_state_space", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum LdmatrixStateSpaceAttr { + Shared, +} + +/// A warp-cooperative matrix load whose exact variant is carried by attributes. +#[pliron_op( + name = "nvvm.ldmatrix", + format, + interfaces = [NOpdsInterface<1>], + attributes = ( + nvvm_ldmatrix_shape: LdmatrixShapeAttr, + nvvm_ldmatrix_multiplicity: LdmatrixMultiplicityAttr, + nvvm_ldmatrix_layout: LdmatrixLayoutAttr, + nvvm_ldmatrix_element: LdmatrixElementAttr, + nvvm_ldmatrix_state_space: LdmatrixStateSpaceAttr + ) +)] +pub struct LdmatrixOp; + +impl LdmatrixOp { + pub fn new(op: Ptr) -> Self { + Self { op } + } + + pub fn build( + ctx: &mut Context, + address: Value, + shape: LdmatrixShapeAttr, + multiplicity: LdmatrixMultiplicityAttr, + layout: LdmatrixLayoutAttr, + element: LdmatrixElementAttr, + state_space: LdmatrixStateSpaceAttr, + ) -> Ptr { + let register_count = shape.register_count(&multiplicity); + let u32_ty = IntegerType::get(ctx, 32, Signedness::Unsigned); + let op = Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![u32_ty.into(); register_count], + vec![address], + vec![], + 0, + ); + let this = Self { op }; + this.set_attr_nvvm_ldmatrix_shape(ctx, shape); + this.set_attr_nvvm_ldmatrix_multiplicity(ctx, multiplicity); + this.set_attr_nvvm_ldmatrix_layout(ctx, layout); + this.set_attr_nvvm_ldmatrix_element(ctx, element); + this.set_attr_nvvm_ldmatrix_state_space(ctx, state_space); + this.get_operation() + } +} + +impl Verify for LdmatrixOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + let operands: Vec<_> = op.operands().collect(); + if operands.len() != 1 { + return verify_err!(op.loc(), "nvvm.ldmatrix requires exactly one pointer operand"); + } + let pointer = operands[0].get_type(ctx); + let pointer_object = pointer.deref(ctx); + let Some(pointer) = pointer_object.downcast_ref::() else { + return verify_err!(op.loc(), "nvvm.ldmatrix operand must be a MIR pointer"); + }; + if !matches!(pointer.address_space, address_space::GENERIC | address_space::SHARED) { + return verify_err!( + op.loc(), + "nvvm.ldmatrix pointer must be generic (p0) or shared (p3), not address space {}", + pointer.address_space + ); + } + let pointee = pointer.pointee.deref(ctx); + let Some(pointee) = pointee.downcast_ref::() else { + return verify_err!(op.loc(), "nvvm.ldmatrix pointer must point to an unsigned integer"); + }; + + let Some(shape) = self.get_attr_nvvm_ldmatrix_shape(ctx) else { + return verify_err!(op.loc(), "nvvm.ldmatrix requires a shape attribute"); + }; + let Some(multiplicity) = self.get_attr_nvvm_ldmatrix_multiplicity(ctx) else { + return verify_err!(op.loc(), "nvvm.ldmatrix requires a multiplicity attribute"); + }; + let layout = self.get_attr_nvvm_ldmatrix_layout(ctx); + let element = self.get_attr_nvvm_ldmatrix_element(ctx); + let expected_pointee_width = if element.as_deref() == Some(&LdmatrixElementAttr::B16) { + 32 + } else { + 8 + }; + if pointee.width() != expected_pointee_width + || pointee.signedness() != Signedness::Unsigned + { + return verify_err!( + op.loc(), + "nvvm.ldmatrix pointer element width does not match its matrix element format" + ); + } + let supported = matches!( + (&*shape, layout.as_deref(), element.as_deref()), + (LdmatrixShapeAttr::M8n8, Some(_), Some(LdmatrixElementAttr::B16)) + | + ( + LdmatrixShapeAttr::M8n16, + Some(LdmatrixLayoutAttr::Normal), + Some( + LdmatrixElementAttr::B8x16B4x16P64 + | LdmatrixElementAttr::B8x16B6x16P32, + ), + ) + | + ( + LdmatrixShapeAttr::M16n16, + Some(LdmatrixLayoutAttr::Transposed), + Some( + LdmatrixElementAttr::B8 + | LdmatrixElementAttr::B8x16B4x16P64 + | LdmatrixElementAttr::B8x16B6x16P32, + ), + ) + ); + if !supported + || (*shape == LdmatrixShapeAttr::M16n16 + && *multiplicity == LdmatrixMultiplicityAttr::X4) + || self.get_attr_nvvm_ldmatrix_state_space(ctx).as_deref() + != Some(&LdmatrixStateSpaceAttr::Shared) + { + return verify_err!(op.loc(), "nvvm.ldmatrix has a missing or unsupported variant attribute"); + } + + let register_count = shape.register_count(&multiplicity); + if op.get_num_results() != register_count { + return verify_err!( + op.loc(), + "nvvm.ldmatrix {:?} requires {} u32 results", + multiplicity, + register_count + ); + } + for index in 0..register_count { + let ty = op.get_result(index).get_type(ctx); + let ty_object = ty.deref(ctx); + let Some(integer) = ty_object.downcast_ref::() else { + return verify_err!(op.loc(), "nvvm.ldmatrix result {} must be u32", index); + }; + if integer.width() != 32 || integer.signedness() != Signedness::Unsigned { + return verify_err!(op.loc(), "nvvm.ldmatrix result {} must be u32", index); + } + } + Ok(()) + } +} + +pub(super) fn register(ctx: &mut Context) { + LdmatrixShapeAttr::register(ctx); + LdmatrixMultiplicityAttr::register(ctx); + LdmatrixLayoutAttr::register(ctx); + LdmatrixElementAttr::register(ctx); + LdmatrixStateSpaceAttr::register(ctx); + LdmatrixOp::register(ctx); +} +"##, + ); + + let compatibility = ldmatrix(catalog) + .filter_map(|record| ldmatrix_compat_op(record).map(|compat| (record, compat))) + .collect::>(); + assert_eq!(compatibility.len(), 6); + + let mut definitions = String::from( + r#" +fn verify_compat_ldmatrix( + ctx: &Context, + op: Ptr, + op_name: &str, + result_count: usize, +) -> Result<(), Error> { + let op = op.deref(ctx); + let operands: Vec<_> = op.operands().collect(); + if operands.len() != 1 { + return verify_err!(op.loc(), "{} requires one shared-memory pointer", op_name); + } + if operands[0] + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_none() + { + return verify_err!(op.loc(), "{} operand 0 must be a MIR pointer", op_name); + } + if op.get_num_results() != result_count { + return verify_err!( + op.loc(), + "{} requires {} register results", + op_name, + result_count + ); + } + for index in 0..result_count { + let ty = op.get_result(index).get_type(ctx); + let ty = ty.deref(ctx); + let Some(integer) = ty.downcast_ref::() else { + return verify_err!(op.loc(), "{} result {} must be an integer", op_name, index); + }; + if integer.width() != 32 { + return verify_err!(op.loc(), "{} result {} must be 32 bits", op_name, index); + } + } + Ok(()) +} + +"#, + ); + for (record, (op_type, op_name)) in &compatibility { + let result_count = record + .ldmatrix + .as_ref() + .expect("ldmatrix compatibility record") + .variant + .multiplicity + .register_count(); + writeln!( + definitions, + "/// Compatibility carrier for `{op_name}`.\n#[pliron_op(\n name = {op_name:?},\n format,\n interfaces = [NOpdsInterface<1>, NResultsInterface<{result_count}>],\n)]\npub struct {op_type};\n\nimpl {op_type} {{\n pub fn new(op: Ptr) -> Self {{\n Self {{ op }}\n }}\n}}\n\nimpl Verify for {op_type} {{\n fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n verify_compat_ldmatrix(ctx, self.get_operation(), {op_name:?}, {result_count})\n }}\n}}\n" + ) + .unwrap(); + } + let register_start = output + .find("pub(super) fn register(ctx: &mut Context) {") + .expect("ldmatrix register function"); + output.insert_str(register_start, &definitions); + let register_anchor = " LdmatrixOp::register(ctx);\n"; + let mut registrations = String::from(register_anchor); + for (_, (op_type, _)) in compatibility { + writeln!(registrations, " {op_type}::register(ctx);").unwrap(); + } + output = output.replacen(register_anchor, ®istrations, 1); + output +} + +pub(in crate::render) fn render_dialect_stmatrix(catalog: &CatalogFile, hash: &str) -> String { + assert_eq!(stmatrices(catalog).count(), 4); + let mut output = rust_header(catalog, hash); + output.push_str( + r#"//! Structural operations for the four generated `stmatrix` stores. + +use dialect_mir::types::MirPtrType; +use pliron::{ + builtin::{ + op_interfaces::{NOpdsInterface, NResultsInterface}, + types::IntegerType, + }, + common_traits::Verify, + context::{Context, Ptr}, + location::Located, + op::Op, + operation::Operation, + result::Error, + r#type::Typed, + verify_err, +}; +use pliron_derive::pliron_op; + +fn verify_stmatrix_operands( + ctx: &Context, + op: Ptr, + op_name: &str, + register_count: usize, +) -> Result<(), Error> { + let op = op.deref(ctx); + let operands: Vec<_> = op.operands().collect(); + if operands.len() != register_count + 1 { + return verify_err!( + op.loc(), + "{} requires one pointer and {} register operands", + op_name, + register_count + ); + } + if operands[0] + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_none() + { + return verify_err!(op.loc(), "{} operand 0 must be a MIR pointer", op_name); + } + for (index, register) in operands.iter().enumerate().skip(1) { + let ty = register.get_type(ctx); + let ty = ty.deref(ctx); + let Some(integer) = ty.downcast_ref::() else { + return verify_err!( + op.loc(), + "{} register operand {} must be an integer", + op_name, + index - 1 + ); + }; + if integer.width() != 32 { + return verify_err!( + op.loc(), + "{} register operand {} must be 32 bits", + op_name, + index - 1 + ); + } + } + Ok(()) +} + +"#, + ); + for record in stmatrices(catalog) { + let (multiplicity, _) = stmatrix_variant(record).expect("stmatrix variant"); + let count = multiplicity.register_count(); + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!( + output, + "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<{}>, NResultsInterface<0>],\n)]", + record.dialect.op_name, + count + 1 + ) + .unwrap(); + writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); + writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); + output.push_str( + " pub fn new(op: Ptr) -> Self {\n Self { op }\n }\n}\n", + ); + writeln!(output, "\nimpl Verify for {} {{", record.dialect.op_type).unwrap(); + writeln!( + output, + " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n verify_stmatrix_operands(ctx, self.get_operation(), {:?}, {count})\n }}\n}}\n", + record.dialect.op_name + ) + .unwrap(); + } + output.push_str("pub(super) fn register(ctx: &mut Context) {\n"); + for record in stmatrices(catalog) { + writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); + } + output.push_str("}\n"); + output +} + +pub(in crate::render) fn render_dialect_register_mma(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + r#"//! One structural operation for generated register-only `mma.sync` variants. + +use pliron::{ + attribute::Attribute, + builtin::{ + op_interfaces::{NOpdsInterface, NResultsInterface}, + types::{FP32Type, FP64Type, IntegerType, Signedness}, + }, + common_traits::Verify, + context::{Context, Ptr}, + location::Located, + op::Op, + operation::Operation, + result::Error, + r#type::{TypeHandle, Typed}, + verify_err, +}; +use pliron_derive::{pliron_attr, pliron_op}; + +#[pliron_attr(name = "nvvm.register_mma_shape", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum RegisterMmaShapeAttr { M8n8k4, M8n8k16, M8n8k32, M8n8k128, M16n8k4, M16n8k8, M16n8k16, M16n8k32, M16n8k64, M16n8k128, M16n8k256 } + +#[pliron_attr(name = "nvvm.register_mma_operation", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum RegisterMmaOperationAttr { Multiply, AndPopc, XorPopc } + +#[pliron_attr(name = "nvvm.register_mma_kind", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum RegisterMmaKindAttr { Standard, F8f6f4, Mxf8f6f4 } + +#[pliron_attr(name = "nvvm.register_mma_accumulator", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum RegisterMmaAccumulatorAttr { F16, F32, F64, S32 } + +#[pliron_attr(name = "nvvm.register_mma_element", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum RegisterMmaElementAttr { Bf16, E2m1, E2m3, E3m2, E4m3, E5m2, F16, Tf32, F64, B1, S4, U4, S8, U8 } + +#[pliron_attr(name = "nvvm.register_mma_layout", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum RegisterMmaLayoutAttr { Row, Col } + +#[pliron_attr(name = "nvvm.register_mma_overflow", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum RegisterMmaOverflowAttr { NotApplicable, Wrapping, Satfinite } + +#[pliron_op( + name = "nvvm.register_mma", + format, + attributes = ( + nvvm_register_mma_shape: RegisterMmaShapeAttr, + nvvm_register_mma_operation: RegisterMmaOperationAttr, + nvvm_register_mma_kind: RegisterMmaKindAttr, + nvvm_register_mma_accumulator: RegisterMmaAccumulatorAttr, + nvvm_register_mma_a_element: RegisterMmaElementAttr, + nvvm_register_mma_b_element: RegisterMmaElementAttr, + nvvm_register_mma_a_layout: RegisterMmaLayoutAttr, + nvvm_register_mma_b_layout: RegisterMmaLayoutAttr, + nvvm_register_mma_overflow: RegisterMmaOverflowAttr + ) +)] +pub struct RegisterMmaOp; + +impl RegisterMmaOp { + pub fn new(op: Ptr) -> Self { Self { op } } + + /// Defaults older multiply-only IR to multiply. + pub fn operation_or_multiply(&self, ctx: &Context) -> RegisterMmaOperationAttr { + self.get_attr_nvvm_register_mma_operation(ctx) + .as_deref() + .cloned() + .unwrap_or(RegisterMmaOperationAttr::Multiply) + } + + /// Infers the kind used by older generated IR. + pub fn kind_or_inferred(&self, ctx: &Context) -> RegisterMmaKindAttr { + if let Some(kind) = self.get_attr_nvvm_register_mma_kind(ctx).as_deref() { + return kind.clone(); + } + let low_format = |element: Option<&RegisterMmaElementAttr>| { + matches!( + element, + Some( + RegisterMmaElementAttr::E2m1 + | RegisterMmaElementAttr::E2m3 + | RegisterMmaElementAttr::E3m2 + | RegisterMmaElementAttr::E4m3 + | RegisterMmaElementAttr::E5m2 + ) + ) + }; + let old_f8f6f4 = self.get_attr_nvvm_register_mma_shape(ctx).as_deref() + == Some(&RegisterMmaShapeAttr::M16n8k32) + && low_format(self.get_attr_nvvm_register_mma_a_element(ctx).as_deref()) + && low_format(self.get_attr_nvvm_register_mma_b_element(ctx).as_deref()); + if old_f8f6f4 { + RegisterMmaKindAttr::F8f6f4 + } else { + RegisterMmaKindAttr::Standard + } + } +} + +#[derive(Clone, Copy)] +enum MmaCarrier { F32, F64, I32, U16, U32 } + +fn is_carrier(ctx: &Context, ty: TypeHandle, carrier: MmaCarrier) -> bool { + match carrier { + MmaCarrier::F32 => ty.deref(ctx).downcast_ref::().is_some(), + MmaCarrier::F64 => ty.deref(ctx).downcast_ref::().is_some(), + MmaCarrier::I32 | MmaCarrier::U16 | MmaCarrier::U32 => { + let expected = if matches!(carrier, MmaCarrier::I32) { + Signedness::Signed + } else { + Signedness::Unsigned + }; + ty.deref(ctx) + .downcast_ref::() + .is_some_and(|integer| { + integer.width() + == if matches!(carrier, MmaCarrier::U16) { 16 } else { 32 } + && integer.signedness() == expected + }) + } + } +} + +impl Verify for RegisterMmaOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + let operation = self.operation_or_multiply(ctx); + let kind = self.kind_or_inferred(ctx); + let recipe: (&[MmaCarrier], &[MmaCarrier]) = match ( + self.get_attr_nvvm_register_mma_shape(ctx).as_deref(), + operation, + kind, + self.get_attr_nvvm_register_mma_accumulator(ctx).as_deref(), + self.get_attr_nvvm_register_mma_a_element(ctx).as_deref(), + self.get_attr_nvvm_register_mma_b_element(ctx).as_deref(), + self.get_attr_nvvm_register_mma_a_layout(ctx).as_deref(), + self.get_attr_nvvm_register_mma_b_layout(ctx).as_deref(), + self.get_attr_nvvm_register_mma_overflow(ctx).as_deref(), + ) { +"#, + ); + for record in register_mmas(catalog) { + let ( + shape, + operation, + kind, + accumulator, + a_element, + b_element, + a_layout, + b_layout, + overflow, + ) = register_mma_attr_variants(record); + let (operands, results) = register_mma_carriers(record); + writeln!( + output, + " (Some(&{shape}), {operation}, {kind}, Some(&{accumulator}), Some(&{a_element}), Some(&{b_element}), Some(&{a_layout}), Some(&{b_layout}), Some(&{overflow})) => ({operands}, {results})," + ) + .unwrap(); + } + output.push_str( + r#" _ => return verify_err!(op.loc(), "nvvm.register_mma has a missing or unsupported variant"), + }; + let operands: Vec<_> = op.operands().collect(); + if operands.len() != recipe.0.len() || op.get_num_results() != recipe.1.len() { + return verify_err!(op.loc(), "nvvm.register_mma has the wrong register count"); + } + for (index, (operand, carrier)) in operands.iter().zip(recipe.0).enumerate() { + if !is_carrier(ctx, operand.get_type(ctx), *carrier) { + return verify_err!(op.loc(), "nvvm.register_mma operand {} has the wrong carrier type", index); + } + } + for (index, carrier) in recipe.1.iter().enumerate() { + if !is_carrier(ctx, op.get_result(index).get_type(ctx), *carrier) { + return verify_err!(op.loc(), "nvvm.register_mma result {} has the wrong carrier type", index); + } + } + Ok(()) + } +} + +"#, + ); + output.push_str( + r#" +fn is_compat_carrier(ctx: &Context, ty: TypeHandle, carrier: MmaCarrier) -> bool { + match carrier { + MmaCarrier::F32 => ty.deref(ctx).downcast_ref::().is_some(), + MmaCarrier::F64 => ty.deref(ctx).downcast_ref::().is_some(), + MmaCarrier::I32 | MmaCarrier::U16 | MmaCarrier::U32 => { + let expected_width = if matches!(carrier, MmaCarrier::U16) { + 16 + } else { + 32 + }; + ty.deref(ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == expected_width) + } + } +} + +fn verify_compat_register_mma( + ctx: &Context, + op_ptr: Ptr, + op_name: &str, + operand_carriers: &[MmaCarrier], + result_carriers: &[MmaCarrier], +) -> Result<(), Error> { + let op = op_ptr.deref(ctx); + let operands: Vec<_> = op.operands().collect(); + if operands.len() != operand_carriers.len() || op.get_num_results() != result_carriers.len() { + return verify_err!(op.loc(), "{} has the wrong register count", op_name); + } + for (index, (operand, carrier)) in operands.iter().zip(operand_carriers).enumerate() { + if !is_compat_carrier(ctx, operand.get_type(ctx), *carrier) { + return verify_err!(op.loc(), "{} operand {} has the wrong carrier type", op_name, index); + } + } + for (index, carrier) in result_carriers.iter().enumerate() { + if !is_compat_carrier(ctx, op.get_result(index).get_type(ctx), *carrier) { + return verify_err!(op.loc(), "{} result {} has the wrong carrier type", op_name, index); + } + } + Ok(()) +} + +"#, + ); + for record in + register_mmas(catalog).filter(|record| register_mma_compat_op_type(record).is_some()) + { + let op_type = register_mma_compat_op_type(record).unwrap(); + let op_name = format!("nvvm.{}", record.id); + let operand_count = record.dialect.operands.len(); + let result_count = record.dialect.results.len(); + let (operands, results) = register_mma_carriers(record); + writeln!(output, "/// Compatibility carrier for `{}`.", record.id).unwrap(); + writeln!( + output, + "#[pliron_op(\n name = {op_name:?},\n format,\n interfaces = [NOpdsInterface<{operand_count}>, NResultsInterface<{result_count}>],\n)]\npub struct {op_type};\n" + ) + .unwrap(); + writeln!( + output, + "impl {op_type} {{\n pub fn new(op: Ptr) -> Self {{ Self {{ op }} }}\n}}\n" + ) + .unwrap(); + writeln!( + output, + "impl Verify for {op_type} {{\n fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n verify_compat_register_mma(ctx, self.get_operation(), {op_name:?}, {operands}, {results})\n }}\n}}\n" + ) + .unwrap(); + } + output.push_str( + "pub(super) fn register(ctx: &mut Context) {\n RegisterMmaShapeAttr::register(ctx);\n RegisterMmaOperationAttr::register(ctx);\n RegisterMmaKindAttr::register(ctx);\n RegisterMmaAccumulatorAttr::register(ctx);\n RegisterMmaElementAttr::register(ctx);\n RegisterMmaLayoutAttr::register(ctx);\n RegisterMmaOverflowAttr::register(ctx);\n RegisterMmaOp::register(ctx);\n", + ); + for record in + register_mmas(catalog).filter(|record| register_mma_compat_op_type(record).is_some()) + { + writeln!( + output, + " {}::register(ctx);", + register_mma_compat_op_type(record).unwrap() + ) + .unwrap(); + } + output.push_str("}\n"); + output +} + +pub(in crate::render) fn render_dialect_sparse_mma(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + r#"//! One structural operation for generated sparse MMA variants. + +use dialect_mir::ops::MirConstantOp; +use pliron::{ + attribute::Attribute, + builtin::{attributes::IntegerAttr, ops::ConstantOp, types::{FP32Type, IntegerType, Signedness}}, + 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_attr, pliron_op}; + +#[pliron_attr(name = "nvvm.sparse_mma_shape", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum SparseMmaShapeAttr { M16n8k32, M16n8k64, M16n8k128 } + +#[pliron_attr(name = "nvvm.sparse_mma_accumulator", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum SparseMmaAccumulatorAttr { F16, F32, S32 } + +#[pliron_attr(name = "nvvm.sparse_mma_element", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum SparseMmaElementAttr { E2m1, E2m3, E3m2, E4m3, E5m2, S4, U4, S8, U8 } + +#[pliron_attr(name = "nvvm.sparse_mma_layout", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum SparseMmaLayoutAttr { Row, Col } + +#[pliron_attr(name = "nvvm.sparse_mma_overflow", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum SparseMmaOverflowAttr { NotApplicable, Wrapping, Satfinite } + +#[pliron_attr(name = "nvvm.sparse_mma_metadata", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum SparseMmaMetadataAttr { Standard, Ordered } + +#[pliron_attr(name = "nvvm.sparse_mma_selector", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum SparseMmaSelectorAttr { ImmediateZeroOrOne, ImmediateZero } + +#[pliron_op( + name = "nvvm.sparse_mma", + format, + attributes = ( + nvvm_sparse_mma_shape: SparseMmaShapeAttr, + nvvm_sparse_mma_accumulator: SparseMmaAccumulatorAttr, + nvvm_sparse_mma_a_element: SparseMmaElementAttr, + nvvm_sparse_mma_b_element: SparseMmaElementAttr, + nvvm_sparse_mma_a_layout: SparseMmaLayoutAttr, + nvvm_sparse_mma_b_layout: SparseMmaLayoutAttr, + nvvm_sparse_mma_overflow: SparseMmaOverflowAttr, + nvvm_sparse_mma_metadata: SparseMmaMetadataAttr, + nvvm_sparse_mma_selector: SparseMmaSelectorAttr + ) +)] +pub struct SparseMmaOp; + +impl SparseMmaOp { + pub fn new(op: Ptr) -> Self { Self { op } } +} + +#[derive(Clone, Copy)] +enum MmaCarrier { F32, I32, U32 } + +fn is_carrier(ctx: &Context, ty: TypeHandle, carrier: MmaCarrier) -> bool { + match carrier { + MmaCarrier::F32 => ty.deref(ctx).downcast_ref::().is_some(), + MmaCarrier::I32 | MmaCarrier::U32 => { + let expected = if matches!(carrier, MmaCarrier::I32) { + Signedness::Signed + } else { + Signedness::Unsigned + }; + ty.deref(ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 32 && integer.signedness() == expected) + } + } +} + +fn constant_u32(ctx: &Context, value: Value) -> Option { + let defining_op = value.defining_op()?; + if let Some(constant) = Operation::get_op::(defining_op, ctx) { + return constant.get_attr_value(ctx).map(|value| value.value().to_u64()); + } + let constant = Operation::get_op::(defining_op, ctx)?; + constant + .get_value(ctx) + .downcast_ref::() + .map(|value| value.value().to_u64()) +} + +impl Verify for SparseMmaOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + let (expected_operands, expected_results, selector_upper_exclusive): (&[MmaCarrier], &[MmaCarrier], usize) = match ( + self.get_attr_nvvm_sparse_mma_shape(ctx).as_deref(), + self.get_attr_nvvm_sparse_mma_accumulator(ctx).as_deref(), + self.get_attr_nvvm_sparse_mma_a_element(ctx).as_deref(), + self.get_attr_nvvm_sparse_mma_b_element(ctx).as_deref(), + self.get_attr_nvvm_sparse_mma_a_layout(ctx).as_deref(), + self.get_attr_nvvm_sparse_mma_b_layout(ctx).as_deref(), + self.get_attr_nvvm_sparse_mma_overflow(ctx).as_deref(), + self.get_attr_nvvm_sparse_mma_metadata(ctx).as_deref(), + self.get_attr_nvvm_sparse_mma_selector(ctx).as_deref(), + ) { +"#, + ); + for record in sparse_mmas(catalog) { + let ( + shape, + accumulator, + a_element, + b_element, + a_layout, + b_layout, + overflow, + metadata, + selector, + ) = sparse_mma_attr_variants(record); + let (expected_operands, expected_results) = sparse_mma_carriers(record); + let selector_upper_exclusive = sparse_mma_selector_values(record).len(); + writeln!( + output, + " (Some(&{shape}), Some(&{accumulator}), Some(&{a_element}), Some(&{b_element}), Some(&{a_layout}), Some(&{b_layout}), Some(&{overflow}), Some(&{metadata}), Some(&{selector})) => ({expected_operands}, {expected_results}, {selector_upper_exclusive})," + ) + .unwrap(); + } + output.push_str( + r#" _ => return verify_err!(op.loc(), "nvvm.sparse_mma has a missing or unsupported variant"), + }; + let operands: Vec<_> = op.operands().collect(); + if operands.len() != expected_operands.len() || op.get_num_results() != expected_results.len() { + return verify_err!(op.loc(), "nvvm.sparse_mma has the wrong register count"); + } + for (index, (operand, carrier)) in operands.iter().zip(expected_operands).enumerate() { + if !is_carrier(ctx, operand.get_type(ctx), *carrier) { + return verify_err!(op.loc(), "nvvm.sparse_mma operand {} has the wrong carrier type", index); + } + } + for (index, carrier) in expected_results.iter().enumerate() { + if !is_carrier(ctx, op.get_result(index).get_type(ctx), *carrier) { + return verify_err!(op.loc(), "nvvm.sparse_mma result {} has the wrong carrier type", index); + } + } + if constant_u32(ctx, operands[expected_operands.len() - 1]) + .is_none_or(|selector| selector >= selector_upper_exclusive as u64) + { + return verify_err!(op.loc(), "nvvm.sparse_mma selector must be a compile-time constant in 0..{}", selector_upper_exclusive); + } + Ok(()) + } +} + +pub(super) fn register(ctx: &mut Context) { + SparseMmaShapeAttr::register(ctx); + SparseMmaAccumulatorAttr::register(ctx); + SparseMmaElementAttr::register(ctx); + SparseMmaLayoutAttr::register(ctx); + SparseMmaOverflowAttr::register(ctx); + SparseMmaMetadataAttr::register(ctx); + SparseMmaSelectorAttr::register(ctx); + SparseMmaOp::register(ctx); +} +"#, + ); + output +} diff --git a/crates/cuda-intrinsics-gen/src/render/dialect/mod.rs b/crates/cuda-intrinsics-gen/src/render/dialect/mod.rs new file mode 100644 index 0000000000..95f0b41e50 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/render/dialect/mod.rs @@ -0,0 +1,212 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::CatalogFile; +use crate::render::common::rust_header; +use crate::render::families::{ + clc_intrinsics, cluster_memory, debug_controls, elect_intrinsics, execution_controls, + extended_minmax, integer_minmaxes, mbarrier_extended, movmatrix, scalar_arithmetics, + scalar_conversions, scalar_maths, stmatrices, tcgen05_intrinsics, tma_intrinsics, + wgmma_controls, +}; + +mod cluster_tensor; +mod mma; +mod scalar_packed; +mod warp_sync; + +pub(super) use cluster_tensor::*; +pub(super) use mma::*; +pub(super) use scalar_packed::*; +pub(super) use warp_sync::*; + +pub(super) fn render_dialect_mod(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + "mod active_mask;\nmod cluster_barrier;\nmod cp_async;\nmod dotprod;\nmod ldmatrix;\nmod mbarrier_basic;\nmod movmatrix;\nmod packed_alu;\nmod packed_atomic;\nmod packed_conversion;\nmod prmt;\nmod redux;\nmod register_mma;\nmod sparse_mma;\nmod sreg;\nmod sync;\nmod vote;\nmod warp_barrier;\nmod warp_match;\nmod warp_shuffle;\n\npub use active_mask::*;\npub use cluster_barrier::*;\npub use cp_async::*;\npub use dotprod::*;\npub use ldmatrix::*;\npub use mbarrier_basic::*;\npub use movmatrix::*;\npub use packed_alu::*;\npub use packed_atomic::*;\npub use packed_conversion::*;\npub use prmt::*;\npub use redux::*;\npub use register_mma::*;\npub use sparse_mma::*;\npub use sreg::*;\npub use sync::*;\npub use vote::*;\npub use warp_barrier::*;\npub use warp_match::*;\npub use warp_shuffle::*;\n\nuse pliron::context::Context;\n\npub(super) fn register(ctx: &mut Context) {\n active_mask::register(ctx);\n cluster_barrier::register(ctx);\n cp_async::register(ctx);\n dotprod::register(ctx);\n ldmatrix::register(ctx);\n mbarrier_basic::register(ctx);\n movmatrix::register(ctx);\n packed_alu::register(ctx);\n packed_atomic::register(ctx);\n packed_conversion::register(ctx);\n prmt::register(ctx);\n redux::register(ctx);\n register_mma::register(ctx);\n sparse_mma::register(ctx);\n sreg::register(ctx);\n sync::register(ctx);\n vote::register(ctx);\n warp_barrier::register(ctx);\n warp_match::register(ctx);\n warp_shuffle::register(ctx);\n}\n", + ); + if scalar_conversions(catalog).next().is_some() { + output = output + .replace("mod redux;", "mod redux;\nmod scalar_conversion;") + .replace( + "pub use redux::*;", + "pub use redux::*;\npub use scalar_conversion::*;", + ) + .replace( + " redux::register(ctx);", + " redux::register(ctx);\n scalar_conversion::register(ctx);", + ); + } + if scalar_arithmetics(catalog).next().is_some() { + output = output + .replace("mod redux;", "mod redux;\nmod scalar_arithmetic;") + .replace( + "pub use redux::*;", + "pub use redux::*;\npub use scalar_arithmetic::*;", + ) + .replace( + " redux::register(ctx);", + " redux::register(ctx);\n scalar_arithmetic::register(ctx);", + ); + } + if scalar_maths(catalog).next().is_some() { + output = output + .replace("mod redux;", "mod redux;\nmod scalar_math;") + .replace( + "pub use redux::*;", + "pub use redux::*;\npub use scalar_math::*;", + ) + .replace( + " redux::register(ctx);", + " redux::register(ctx);\n scalar_math::register(ctx);", + ); + } + if integer_minmaxes(catalog).next().is_some() { + output = output + .replace("mod ldmatrix;", "mod integer_minmax;\nmod ldmatrix;") + .replace( + "pub use ldmatrix::*;", + "pub use integer_minmax::*;\npub use ldmatrix::*;", + ) + .replace( + " ldmatrix::register(ctx);", + " integer_minmax::register(ctx);\n ldmatrix::register(ctx);", + ); + } + if extended_minmax(catalog).next().is_some() { + output = output + .replace("mod dotprod;", "mod dotprod;\nmod extended_minmax;") + .replace( + "pub use dotprod::*;", + "pub use dotprod::*;\npub use extended_minmax::*;", + ) + .replace( + " dotprod::register(ctx);", + " dotprod::register(ctx);\n extended_minmax::register(ctx);", + ); + } + if debug_controls(catalog).next().is_some() { + output = output + .replace("mod dotprod;", "mod debug_control;\nmod dotprod;") + .replace( + "pub use dotprod::*;", + "pub use debug_control::*;\npub use dotprod::*;", + ) + .replace( + " dotprod::register(ctx);", + " debug_control::register(ctx);\n dotprod::register(ctx);", + ); + } + if cluster_memory(catalog).next().is_some() { + output = output + .replace( + "mod cluster_barrier;", + "mod cluster_barrier;\nmod cluster_memory;", + ) + .replace( + "pub use cluster_barrier::*;", + "pub use cluster_barrier::*;\npub use cluster_memory::*;", + ) + .replace( + " cluster_barrier::register(ctx);", + " cluster_barrier::register(ctx);\n cluster_memory::register(ctx);", + ); + } + if stmatrices(catalog).next().is_some() { + output = output + .replace("mod sync;", "mod stmatrix;\nmod sync;") + .replace("pub use sync::*;", "pub use stmatrix::*;\npub use sync::*;") + .replace( + " sync::register(ctx);", + " stmatrix::register(ctx);\n sync::register(ctx);", + ); + } + if movmatrix(catalog).next().is_none() { + output = output.replace("mod movmatrix;\n", ""); + output = output.replace("pub use movmatrix::*;\n", ""); + output = output.replace(" movmatrix::register(ctx);\n", ""); + } + if clc_intrinsics(catalog).next().is_some() { + output = output + .replace("mod cluster_barrier;", "mod clc;\nmod cluster_barrier;") + .replace( + "pub use cluster_barrier::*;", + "pub use clc::*;\npub use cluster_barrier::*;", + ) + .replace( + " cluster_barrier::register(ctx);", + " clc::register(ctx);\n cluster_barrier::register(ctx);", + ); + } + if elect_intrinsics(catalog).next().is_some() { + output = output + .replace("mod dotprod;", "mod dotprod;\nmod elect;") + .replace( + "pub use dotprod::*;", + "pub use dotprod::*;\npub use elect::*;", + ) + .replace( + " dotprod::register(ctx);", + " dotprod::register(ctx);\n elect::register(ctx);", + ); + } + if mbarrier_extended(catalog).next().is_some() { + output = output.replace( + "mod mbarrier_basic;\n", + "mod mbarrier_basic;\nmod mbarrier_extended;\n", + ); + output = output.replace( + "pub use mbarrier_basic::*;\n", + "pub use mbarrier_basic::*;\npub use mbarrier_extended::*;\n", + ); + output = output.replace( + " mbarrier_basic::register(ctx);\n", + " mbarrier_basic::register(ctx);\n mbarrier_extended::register(ctx);\n", + ); + } + if wgmma_controls(catalog).next().is_some() { + output = output + .replace("mod warp_shuffle;", "mod warp_shuffle;\nmod wgmma_control;") + .replace( + "pub use warp_shuffle::*;", + "pub use warp_shuffle::*;\npub use wgmma_control::*;", + ) + .replace( + " warp_shuffle::register(ctx);", + " warp_shuffle::register(ctx);\n wgmma_control::register(ctx);", + ); + } + if execution_controls(catalog).next().is_some() { + output = output + .replace("mod dotprod;", "mod dotprod;\nmod execution_control;") + .replace( + "pub use dotprod::*;", + "pub use dotprod::*;\npub use execution_control::*;", + ) + .replace( + " dotprod::register(ctx);", + " dotprod::register(ctx);\n execution_control::register(ctx);", + ); + } + if tma_intrinsics(catalog).next().is_some() { + output = output + .replace("mod sync;", "mod sync;\nmod tma;") + .replace("pub use sync::*;", "pub use sync::*;\npub use tma::*;") + .replace( + " sync::register(ctx);", + " sync::register(ctx);\n tma::register(ctx);", + ); + } + if tcgen05_intrinsics(catalog).next().is_some() { + output = output + .replace("mod sync;", "mod sync;\nmod tcgen05;") + .replace("pub use sync::*;", "pub use sync::*;\npub use tcgen05::*;") + .replace( + " sync::register(ctx);", + " sync::register(ctx);\n tcgen05::register(ctx);", + ); + } + output +} diff --git a/crates/cuda-intrinsics-gen/src/render/dialect/scalar_packed.rs b/crates/cuda-intrinsics-gen/src/render/dialect/scalar_packed.rs new file mode 100644 index 0000000000..045f7c5c83 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/render/dialect/scalar_packed.rs @@ -0,0 +1,1190 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{CatalogFile, PackedConversionSourceFormat}; +use crate::render::common::rust_header; +use crate::render::families::{ + extended_minmax, extended_minmax_format_attr, extended_minmax_nan_attr, + extended_minmax_operation_attr, extended_minmax_subnormal_attr, + extended_minmax_xorsign_abs_attr, integer_minmax_ptx_mnemonic, integer_minmaxes, + packed_alu_ptx_mnemonic, packed_alu_width, packed_alus, packed_conversion_element, + packed_conversion_result_width, packed_conversion_source, packed_conversion_source_width, + packed_conversions, scalar_arithmetic_format_attr, scalar_arithmetic_operation_attr, + scalar_arithmetic_rounding_attr, scalar_arithmetic_saturation_attr, + scalar_arithmetic_subnormal_attr, scalar_arithmetics, scalar_conversion_rounding_attr, + scalar_conversion_saturation_attr, scalar_conversions, scalar_math_format_attr, + scalar_math_operation_attr, scalar_math_precision_attr, scalar_math_subnormal_attr, + scalar_maths, +}; +use std::fmt::Write as _; + +pub(in crate::render) fn render_dialect_integer_minmax( + catalog: &CatalogFile, + hash: &str, +) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + "//! Structural operations for generated extended integer min/max.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_i32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32)\n}\n\n", + ); + for record in integer_minmaxes(catalog) { + let signedness = if record.rust.result == "i32" { + "Signed" + } else { + "Unsigned" + }; + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!( + output, + "///\n/// Lowers to `{}`.", + integer_minmax_ptx_mnemonic(record) + ) + .unwrap(); + writeln!( + output, + "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<2>, NResultsInterface<1>],\n)]", + record.dialect.op_name + ) + .unwrap(); + writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); + writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); + writeln!( + output, + " pub fn new(op: Ptr) -> Self {{\n Self {{ op }}\n }}\n\n pub fn build(ctx: &mut Context, arg0: Value, arg1: Value) -> Ptr {{\n let result_ty = IntegerType::get(ctx, 32, Signedness::{signedness});\n Operation::new(\n ctx,\n Self::get_concrete_op_info(),\n vec![result_ty.into()],\n vec![arg0, arg1],\n vec![],\n 0,\n )\n }}\n}}" + ) + .unwrap(); + writeln!(output, "\nimpl Verify for {} {{", record.dialect.op_type).unwrap(); + writeln!( + output, + " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n let op = self.get_operation().deref(ctx);\n if op.get_num_operands() != 2 || op.get_num_results() != 1 {{\n return verify_err!(op.loc(), {:?});\n }}\n if !(0..2).all(|index| is_i32(ctx, op.get_operand(index).get_type(ctx)))\n || !is_i32(ctx, op.get_result(0).get_type(ctx))\n {{\n return verify_err!(op.loc(), {:?});\n }}\n Ok(())\n }}\n}}\n", + format!( + "{} requires exactly 2 operands and one result", + record.dialect.op_name + ), + format!( + "{} operands and result must be 32-bit integers", + record.dialect.op_name + ), + ) + .unwrap(); + } + output.push_str("\npub(super) fn register(ctx: &mut Context) {\n"); + for record in integer_minmaxes(catalog) { + writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); + } + output.push_str("}\n"); + output +} + +pub(in crate::render) fn render_dialect_packed_atomic(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + r##"//! One structural operation for the closed generated packed-atomic family. + +use dialect_mir::types::{address_space, MirPtrType}; +use pliron::{ + attribute::Attribute, + builtin::{ + op_interfaces::{NOpdsInterface, NResultsInterface}, + types::{IntegerType, Signedness}, + }, + 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_attr, pliron_op}; + +#[pliron_attr(name = "nvvm.packed_atomic_format", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum PackedAtomicFormatAttr { F16x2, Bf16x2 } + +#[pliron_attr(name = "nvvm.packed_atomic_state_space", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum PackedAtomicStateSpaceAttr { Global } + +#[pliron_attr(name = "nvvm.packed_atomic_ordering", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum PackedAtomicOrderingAttr { Relaxed } + +#[pliron_attr(name = "nvvm.packed_atomic_scope", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum PackedAtomicScopeAttr { Gpu } + +#[pliron_attr(name = "nvvm.packed_atomic_rounding", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum PackedAtomicRoundingAttr { Rn } + +#[pliron_attr(name = "nvvm.packed_atomic_subnormal", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum PackedAtomicSubnormalAttr { NoFtz } + +#[pliron_attr(name = "nvvm.packed_atomic_atomicity", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum PackedAtomicAtomicityAttr { PerElement } + +/// Packed global atomic add with exact format and semantic attributes. +#[pliron_op( + name = "nvvm.packed_atomic_add", + format, + interfaces = [NOpdsInterface<2>, NResultsInterface<1>], + attributes = ( + nvvm_packed_atomic_format: PackedAtomicFormatAttr, + nvvm_packed_atomic_state_space: PackedAtomicStateSpaceAttr, + nvvm_packed_atomic_ordering: PackedAtomicOrderingAttr, + nvvm_packed_atomic_scope: PackedAtomicScopeAttr, + nvvm_packed_atomic_rounding: PackedAtomicRoundingAttr, + nvvm_packed_atomic_subnormal: PackedAtomicSubnormalAttr, + nvvm_packed_atomic_atomicity: PackedAtomicAtomicityAttr + ) +)] +pub struct PackedAtomicAddOp; + +impl PackedAtomicAddOp { + pub fn new(op: Ptr) -> Self { Self { op } } + + pub fn build( + ctx: &mut Context, + address: Value, + addend: Value, + format: PackedAtomicFormatAttr, + ) -> Ptr { + let u32_ty = IntegerType::get(ctx, 32, Signedness::Unsigned); + let op = Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![u32_ty.into()], + vec![address, addend], + vec![], + 0, + ); + let this = Self { op }; + this.set_attr_nvvm_packed_atomic_format(ctx, format); + this.set_attr_nvvm_packed_atomic_state_space(ctx, PackedAtomicStateSpaceAttr::Global); + this.set_attr_nvvm_packed_atomic_ordering(ctx, PackedAtomicOrderingAttr::Relaxed); + this.set_attr_nvvm_packed_atomic_scope(ctx, PackedAtomicScopeAttr::Gpu); + this.set_attr_nvvm_packed_atomic_rounding(ctx, PackedAtomicRoundingAttr::Rn); + this.set_attr_nvvm_packed_atomic_subnormal(ctx, PackedAtomicSubnormalAttr::NoFtz); + this.set_attr_nvvm_packed_atomic_atomicity(ctx, PackedAtomicAtomicityAttr::PerElement); + this.get_operation() + } +} + +fn is_u32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool { + ty.deref(ctx).downcast_ref::().is_some_and(|integer| { + integer.width() == 32 && integer.signedness() == Signedness::Unsigned + }) +} + +fn verify_packed_atomic_signature( + ctx: &Context, + op_ptr: Ptr, + op_name: &str, +) -> Result<(), Error> { + let op = op_ptr.deref(ctx); + if op.get_num_operands() != 2 || op.get_num_results() != 1 { + return verify_err!(op.loc(), "{} requires exactly two operands and one result", op_name); + } + let pointer_ty = op.get_operand(0).get_type(ctx); + let pointer_object = pointer_ty.deref(ctx); + let Some(pointer) = pointer_object.downcast_ref::() else { + return verify_err!(op.loc(), "{} address must be a MIR pointer", op_name); + }; + if !pointer.is_mutable() + || !matches!(pointer.address_space(), address_space::GENERIC | address_space::GLOBAL) + || !is_u32(ctx, pointer.pointee) + { + return verify_err!(op.loc(), "{} address must be a mutable generic/global pointer to u32", op_name); + } + if !is_u32(ctx, op.get_operand(1).get_type(ctx)) || !is_u32(ctx, op.get_result(0).get_type(ctx)) { + return verify_err!(op.loc(), "{} addend and result must be u32", op_name); + } + Ok(()) +} + +impl Verify for PackedAtomicAddOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + verify_packed_atomic_signature(ctx, self.get_operation(), "nvvm.packed_atomic_add")?; + let op = self.get_operation().deref(ctx); + if self.get_attr_nvvm_packed_atomic_format(ctx).is_none() + || self.get_attr_nvvm_packed_atomic_state_space(ctx).as_deref() != Some(&PackedAtomicStateSpaceAttr::Global) + || self.get_attr_nvvm_packed_atomic_ordering(ctx).as_deref() != Some(&PackedAtomicOrderingAttr::Relaxed) + || self.get_attr_nvvm_packed_atomic_scope(ctx).as_deref() != Some(&PackedAtomicScopeAttr::Gpu) + || self.get_attr_nvvm_packed_atomic_rounding(ctx).as_deref() != Some(&PackedAtomicRoundingAttr::Rn) + || self.get_attr_nvvm_packed_atomic_subnormal(ctx).as_deref() != Some(&PackedAtomicSubnormalAttr::NoFtz) + || self.get_attr_nvvm_packed_atomic_atomicity(ctx).as_deref() != Some(&PackedAtomicAtomicityAttr::PerElement) + { + return verify_err!(op.loc(), "nvvm.packed_atomic_add has a missing or unsupported semantic attribute"); + } + Ok(()) + } +} + +/// Compatibility operation for the existing `nvvm.atom_add_f16x2` carrier. +#[pliron_op( + name = "nvvm.atom_add_f16x2", + format, + interfaces = [NOpdsInterface<2>, NResultsInterface<1>], +)] +pub struct NvvmAtomAddF16x2Op; + +impl NvvmAtomAddF16x2Op { + pub fn new(op: Ptr) -> Self { Self { op } } +} + +impl Verify for NvvmAtomAddF16x2Op { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + verify_packed_atomic_signature(ctx, self.get_operation(), "nvvm.atom_add_f16x2") + } +} + +/// Compatibility operation for the existing `nvvm.atom_add_bf16x2` carrier. +#[pliron_op( + name = "nvvm.atom_add_bf16x2", + format, + interfaces = [NOpdsInterface<2>, NResultsInterface<1>], +)] +pub struct NvvmAtomAddBf16x2Op; + +impl NvvmAtomAddBf16x2Op { + pub fn new(op: Ptr) -> Self { Self { op } } +} + +impl Verify for NvvmAtomAddBf16x2Op { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + verify_packed_atomic_signature(ctx, self.get_operation(), "nvvm.atom_add_bf16x2") + } +} + +pub(super) fn register(ctx: &mut Context) { + PackedAtomicFormatAttr::register(ctx); + PackedAtomicStateSpaceAttr::register(ctx); + PackedAtomicOrderingAttr::register(ctx); + PackedAtomicScopeAttr::register(ctx); + PackedAtomicRoundingAttr::register(ctx); + PackedAtomicSubnormalAttr::register(ctx); + PackedAtomicAtomicityAttr::register(ctx); + PackedAtomicAddOp::register(ctx); + NvvmAtomAddF16x2Op::register(ctx); + NvvmAtomAddBf16x2Op::register(ctx); +} +"##, + ); + output +} + +pub(in crate::render) fn render_dialect_packed_alu(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + "//! Structural operations for generated packed floating-point arithmetic.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_integer_width(\n ctx: &Context,\n ty: pliron::r#type::TypeHandle,\n width: u32,\n) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == width)\n}\n\n", + ); + for record in packed_alus(catalog) { + let arity = record.rust.arguments.len(); + let width = packed_alu_width(record); + let parameters = (0..arity) + .map(|index| format!("arg{index}: Value")) + .collect::>() + .join(", "); + let operands = (0..arity) + .map(|index| format!("arg{index}")) + .collect::>() + .join(", "); + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!( + output, + "///\n/// Lowers to `{}`.", + packed_alu_ptx_mnemonic(record) + ) + .unwrap(); + writeln!( + output, + "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<{arity}>, NResultsInterface<1>],\n)]", + record.dialect.op_name + ) + .unwrap(); + writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); + writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); + writeln!( + output, + " pub fn new(op: Ptr) -> Self {{\n Self {{ op }}\n }}\n\n pub fn build(ctx: &mut Context, {parameters}) -> Ptr {{\n let result_ty = IntegerType::get(ctx, {width}, Signedness::Unsigned);\n Operation::new(\n ctx,\n Self::get_concrete_op_info(),\n vec![result_ty.into()],\n vec![{operands}],\n vec![],\n 0,\n )\n }}\n}}" + ) + .unwrap(); + writeln!(output, "\nimpl Verify for {} {{", record.dialect.op_type).unwrap(); + writeln!( + output, + " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n let op = self.get_operation().deref(ctx);\n if op.get_num_operands() != {arity} || op.get_num_results() != 1 {{\n return verify_err!(op.loc(), {:?});\n }}\n if !(0..{arity}).all(|index| is_integer_width(ctx, op.get_operand(index).get_type(ctx), {width}))\n || !is_integer_width(ctx, op.get_result(0).get_type(ctx), {width})\n {{\n return verify_err!(op.loc(), {:?});\n }}\n Ok(())\n }}\n}}\n", + format!( + "{} requires exactly {arity} operands and one result", + record.dialect.op_name + ), + format!( + "{} operands and result must be {width}-bit integers", + record.dialect.op_name, + ), + ) + .unwrap(); + } + output.push_str("\npub(super) fn register(ctx: &mut Context) {\n"); + for record in packed_alus(catalog) { + writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); + } + output.push_str("}\n"); + output +} + +pub(in crate::render) fn render_dialect_packed_conversion( + catalog: &CatalogFile, + hash: &str, +) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + "//! Structural operation for generated packed conversion.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{FP32Type, IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_f32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx).downcast_ref::().is_some()\n}\n\nfn is_integer_width(\n ctx: &Context,\n ty: pliron::r#type::TypeHandle,\n width: u32,\n) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == width)\n}\n\n", + ); + for record in packed_conversions(catalog) { + let result_width = packed_conversion_result_width(record); + writeln!(output, "/// {}", record.summary).unwrap(); + match packed_conversion_source(record) { + PackedConversionSourceFormat::F32x2 => { + writeln!( + output, + "///\n/// The first input becomes the low {} lane; the second becomes the high lane.", + packed_conversion_element(record) + ) + .unwrap(); + writeln!( + output, + "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<2>, NResultsInterface<1>],\n)]", + record.dialect.op_name + ) + .unwrap(); + writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); + writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); + writeln!( + output, + " pub fn new(op: Ptr) -> Self {{\n Self {{ op }}\n }}\n\n pub fn build(ctx: &mut Context, low: Value, high: Value) -> Ptr {{\n let result_ty = IntegerType::get(ctx, {result_width}, Signedness::Unsigned);\n Operation::new(\n ctx,\n Self::get_concrete_op_info(),\n vec![result_ty.into()],\n vec![low, high],\n vec![],\n 0,\n )\n }}\n}}", + ) + .unwrap(); + writeln!(output, "\nimpl Verify for {} {{", record.dialect.op_type).unwrap(); + writeln!( + output, + " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n let op = self.get_operation().deref(ctx);\n if op.get_num_operands() != 2 || op.get_num_results() != 1 {{\n return verify_err!(op.loc(), {:?});\n }}\n if !is_f32(ctx, op.get_operand(0).get_type(ctx))\n || !is_f32(ctx, op.get_operand(1).get_type(ctx))\n || !is_integer_width(ctx, op.get_result(0).get_type(ctx), {})\n {{\n return verify_err!(op.loc(), {:?});\n }}\n Ok(())\n }}\n}}\n", + format!("{} requires two operands and one result", record.dialect.op_name), + result_width, + format!( + "{} requires f32 operands and one {result_width}-bit integer result", + record.dialect.op_name, + ), + ) + .unwrap(); + } + PackedConversionSourceFormat::E4m3x2 + | PackedConversionSourceFormat::E5m2x2 + | PackedConversionSourceFormat::F16x2 => { + let source_width = packed_conversion_source_width(record); + writeln!( + output, + "///\n/// The single input carries both packed lanes, and lane order is preserved." + ) + .unwrap(); + writeln!( + output, + "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<1>, NResultsInterface<1>],\n)]", + record.dialect.op_name + ) + .unwrap(); + writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); + writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); + writeln!( + output, + " pub fn new(op: Ptr) -> Self {{\n Self {{ op }}\n }}\n\n pub fn build(ctx: &mut Context, packed: Value) -> Ptr {{\n let result_ty = IntegerType::get(ctx, {result_width}, Signedness::Unsigned);\n Operation::new(\n ctx,\n Self::get_concrete_op_info(),\n vec![result_ty.into()],\n vec![packed],\n vec![],\n 0,\n )\n }}\n}}", + ) + .unwrap(); + writeln!(output, "\nimpl Verify for {} {{", record.dialect.op_type).unwrap(); + writeln!( + output, + " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n let op = self.get_operation().deref(ctx);\n if op.get_num_operands() != 1 || op.get_num_results() != 1 {{\n return verify_err!(op.loc(), {:?});\n }}\n if !is_integer_width(ctx, op.get_operand(0).get_type(ctx), {source_width})\n || !is_integer_width(ctx, op.get_result(0).get_type(ctx), {result_width})\n {{\n return verify_err!(op.loc(), {:?});\n }}\n Ok(())\n }}\n}}\n", + format!("{} requires one operand and one result", record.dialect.op_name), + format!( + "{} requires one {source_width}-bit integer operand and one {result_width}-bit integer result", + record.dialect.op_name, + ), + ) + .unwrap(); + } + } + } + output.push_str("\npub(super) fn register(ctx: &mut Context) {\n"); + for record in packed_conversions(catalog) { + writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); + } + output.push_str("}\n"); + output +} + +pub(in crate::render) fn render_dialect_scalar_conversion( + catalog: &CatalogFile, + hash: &str, +) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + r##"//! One structural operation for generated `f32` to TF32 conversions. + +use pliron::{ + attribute::Attribute, + builtin::{ + op_interfaces::{NOpdsInterface, NResultsInterface}, + types::{FP32Type, IntegerType, Signedness}, + }, + 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_attr, pliron_op}; + +#[pliron_attr(name = "nvvm.scalar_conversion_rounding", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum ScalarConversionRoundingAttr { + NearestAway, + NearestEven, + TowardZero, +} + +#[pliron_attr(name = "nvvm.scalar_conversion_saturation", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum ScalarConversionSaturationAttr { + None, + Relu, + Satfinite, + ReluSatfinite, +} + +/// Converts one `f32` value and returns the raw TF32 bits. +#[pliron_op( + name = "nvvm.scalar_conversion", + format, + interfaces = [NOpdsInterface<1>, NResultsInterface<1>], + attributes = ( + nvvm_scalar_conversion_rounding: ScalarConversionRoundingAttr, + nvvm_scalar_conversion_saturation: ScalarConversionSaturationAttr + ) +)] +pub struct ScalarConversionOp; + +impl ScalarConversionOp { + pub fn new(op: Ptr) -> Self { + Self { op } + } + + pub fn build( + ctx: &mut Context, + value: Value, + rounding: ScalarConversionRoundingAttr, + saturation: ScalarConversionSaturationAttr, + ) -> Ptr { + let result_ty = IntegerType::get(ctx, 32, Signedness::Unsigned); + let op = Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![result_ty.into()], + vec![value], + vec![], + 0, + ); + let this = Self { op }; + this.set_attr_nvvm_scalar_conversion_rounding(ctx, rounding); + this.set_attr_nvvm_scalar_conversion_saturation(ctx, saturation); + this.get_operation() + } +} + +impl Verify for ScalarConversionOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + let Some(rounding) = self.get_attr_nvvm_scalar_conversion_rounding(ctx) else { + return verify_err!(op.loc(), "nvvm.scalar_conversion requires rounding"); + }; + let Some(saturation) = self.get_attr_nvvm_scalar_conversion_saturation(ctx) else { + return verify_err!(op.loc(), "nvvm.scalar_conversion requires saturation"); + }; + let admitted = matches!( + (&*rounding, &*saturation), +"##, + ); + for (index, record) in scalar_conversions(catalog).enumerate() { + if index != 0 { + output.push_str(" |\n"); + } + writeln!( + output, + " ({}, {})", + scalar_conversion_rounding_attr(record), + scalar_conversion_saturation_attr(record), + ) + .unwrap(); + } + output.push_str( + r##" ); + if !admitted { + return verify_err!(op.loc(), "nvvm.scalar_conversion variant is not admitted"); + } + if op + .get_operand(0) + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_none() + { + return verify_err!(op.loc(), "nvvm.scalar_conversion operand must be f32"); + } + if op + .get_result(0) + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_none_or(|integer| integer.width() != 32) + { + return verify_err!(op.loc(), "nvvm.scalar_conversion result must be a 32-bit integer"); + } + Ok(()) + } +} + +pub(super) fn register(ctx: &mut Context) { + ScalarConversionRoundingAttr::register(ctx); + ScalarConversionSaturationAttr::register(ctx); + ScalarConversionOp::register(ctx); +} +"##, + ); + output +} + +pub(in crate::render) fn render_dialect_scalar_arithmetic( + catalog: &CatalogFile, + hash: &str, +) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + r##"//! One structural operation for generated scalar arithmetic. + +use pliron::{ + attribute::Attribute, + builtin::{op_interfaces::NResultsInterface, types::{FP32Type, FP64Type}}, + 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_attr, pliron_op}; + +#[pliron_attr(name = "nvvm.scalar_arithmetic_format", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum ScalarArithmeticFormatAttr { F32, F64 } + +#[pliron_attr(name = "nvvm.scalar_arithmetic_operation", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum ScalarArithmeticOperationAttr { Mul, Div, Fma, Add } + +#[pliron_attr(name = "nvvm.scalar_arithmetic_rounding", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum ScalarArithmeticRoundingAttr { Rn, Rz, Rm, Rp } + +#[pliron_attr(name = "nvvm.scalar_arithmetic_subnormal", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum ScalarArithmeticSubnormalAttr { Preserve, Ftz } + +#[pliron_attr(name = "nvvm.scalar_arithmetic_saturation", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum ScalarArithmeticSaturationAttr { None, Sat } + +/// Scalar arithmetic with explicit PTX modifiers. +#[pliron_op( + name = "nvvm.scalar_arithmetic", + format, + interfaces = [NResultsInterface<1>], + attributes = ( + nvvm_scalar_arithmetic_format: ScalarArithmeticFormatAttr, + nvvm_scalar_arithmetic_operation: ScalarArithmeticOperationAttr, + nvvm_scalar_arithmetic_rounding: ScalarArithmeticRoundingAttr, + nvvm_scalar_arithmetic_subnormal: ScalarArithmeticSubnormalAttr, + nvvm_scalar_arithmetic_saturation: ScalarArithmeticSaturationAttr + ) +)] +pub struct ScalarArithmeticOp; + +impl ScalarArithmeticOp { + pub fn new(op: Ptr) -> Self { Self { op } } + + pub fn build( + ctx: &mut Context, + operands: Vec, + format: ScalarArithmeticFormatAttr, + operation: ScalarArithmeticOperationAttr, + rounding: ScalarArithmeticRoundingAttr, + subnormal: ScalarArithmeticSubnormalAttr, + saturation: ScalarArithmeticSaturationAttr, + ) -> Ptr { + let result_ty = match &format { + ScalarArithmeticFormatAttr::F32 => FP32Type::get(ctx).into(), + ScalarArithmeticFormatAttr::F64 => FP64Type::get(ctx).into(), + }; + let op = Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![result_ty], + operands, + vec![], + 0, + ); + let this = Self { op }; + this.set_attr_nvvm_scalar_arithmetic_format(ctx, format); + this.set_attr_nvvm_scalar_arithmetic_operation(ctx, operation); + this.set_attr_nvvm_scalar_arithmetic_rounding(ctx, rounding); + this.set_attr_nvvm_scalar_arithmetic_subnormal(ctx, subnormal); + this.set_attr_nvvm_scalar_arithmetic_saturation(ctx, saturation); + this.get_operation() + } +} + +impl Verify for ScalarArithmeticOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + let Some(format) = self.get_attr_nvvm_scalar_arithmetic_format(ctx) else { + return verify_err!(op.loc(), "nvvm.scalar_arithmetic requires a format"); + }; + let Some(operation) = self.get_attr_nvvm_scalar_arithmetic_operation(ctx) else { + return verify_err!(op.loc(), "nvvm.scalar_arithmetic requires an operation"); + }; + let Some(rounding) = self.get_attr_nvvm_scalar_arithmetic_rounding(ctx) else { + return verify_err!(op.loc(), "nvvm.scalar_arithmetic requires rounding"); + }; + let Some(subnormal) = self.get_attr_nvvm_scalar_arithmetic_subnormal(ctx) else { + return verify_err!(op.loc(), "nvvm.scalar_arithmetic requires a subnormal mode"); + }; + let Some(saturation) = self.get_attr_nvvm_scalar_arithmetic_saturation(ctx) else { + return verify_err!(op.loc(), "nvvm.scalar_arithmetic requires saturation"); + }; + let admitted = matches!( + (&*format, &*operation, &*rounding, &*subnormal, &*saturation), +"##, + ); + for (index, record) in scalar_arithmetics(catalog).enumerate() { + if index != 0 { + output.push_str(" |\n"); + } + writeln!( + output, + " ({}, {}, {}, {}, {})", + scalar_arithmetic_format_attr(record), + scalar_arithmetic_operation_attr(record), + scalar_arithmetic_rounding_attr(record), + scalar_arithmetic_subnormal_attr(record), + scalar_arithmetic_saturation_attr(record), + ) + .unwrap(); + } + output.push_str( + r##" ); + if !admitted { + return verify_err!(op.loc(), "nvvm.scalar_arithmetic variant is not admitted"); + } + let expected_operands = match &*operation { + ScalarArithmeticOperationAttr::Mul + | ScalarArithmeticOperationAttr::Div + | ScalarArithmeticOperationAttr::Add => 2, + ScalarArithmeticOperationAttr::Fma => 3, + }; + if op.get_num_operands() != expected_operands || op.get_num_results() != 1 { + return verify_err!( + op.loc(), + "nvvm.scalar_arithmetic requires {} operands and one result", + expected_operands, + ); + } + let type_matches = |ty: pliron::r#type::TypeHandle| match &*format { + ScalarArithmeticFormatAttr::F32 => ty.deref(ctx).downcast_ref::().is_some(), + ScalarArithmeticFormatAttr::F64 => ty.deref(ctx).downcast_ref::().is_some(), + }; + if !(0..expected_operands).all(|index| type_matches(op.get_operand(index).get_type(ctx))) + || !type_matches(op.get_result(0).get_type(ctx)) + { + return verify_err!(op.loc(), "nvvm.scalar_arithmetic types do not match its format"); + } + Ok(()) + } +} + +pub(super) fn register(ctx: &mut Context) { + ScalarArithmeticFormatAttr::register(ctx); + ScalarArithmeticOperationAttr::register(ctx); + ScalarArithmeticRoundingAttr::register(ctx); + ScalarArithmeticSubnormalAttr::register(ctx); + ScalarArithmeticSaturationAttr::register(ctx); + ScalarArithmeticOp::register(ctx); +} +"##, + ); + output +} + +pub(in crate::render) fn render_dialect_scalar_math(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + r##"//! One structural operation for generated scalar math. + +use pliron::{ + attribute::Attribute, + builtin::{op_interfaces::{NOpdsInterface, NResultsInterface}, types::{FP32Type, FP64Type, IntegerType, Signedness}}, + 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_attr, pliron_op}; + +#[pliron_attr(name = "nvvm.scalar_math_format", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum ScalarMathFormatAttr { F16, F32, F64 } + +#[pliron_attr(name = "nvvm.scalar_math_operation", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum ScalarMathOperationAttr { Sin, Cos, Ex2, Lg2, Rcp, Rsqrt, Sqrt, Tanh } + +#[pliron_attr(name = "nvvm.scalar_math_precision", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum ScalarMathPrecisionAttr { Approx, Rn, Rz, Rm, Rp } + +#[pliron_attr(name = "nvvm.scalar_math_subnormal", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum ScalarMathSubnormalAttr { Preserve, Ftz } + +/// Scalar math with exact PTX modifiers. +#[pliron_op( + name = "nvvm.scalar_math", + format, + interfaces = [NOpdsInterface<1>, NResultsInterface<1>], + attributes = ( + nvvm_scalar_math_format: ScalarMathFormatAttr, + nvvm_scalar_math_operation: ScalarMathOperationAttr, + nvvm_scalar_math_precision: ScalarMathPrecisionAttr, + nvvm_scalar_math_subnormal: ScalarMathSubnormalAttr + ) +)] +pub struct ScalarMathOp; + +impl ScalarMathOp { + pub fn new(op: Ptr) -> Self { Self { op } } + + pub fn build( + ctx: &mut Context, + operand: Value, + format: ScalarMathFormatAttr, + operation: ScalarMathOperationAttr, + precision: ScalarMathPrecisionAttr, + subnormal: ScalarMathSubnormalAttr, + ) -> Ptr { + let result_ty = match &format { + ScalarMathFormatAttr::F16 => IntegerType::get(ctx, 16, Signedness::Unsigned).into(), + ScalarMathFormatAttr::F32 => FP32Type::get(ctx).into(), + ScalarMathFormatAttr::F64 => FP64Type::get(ctx).into(), + }; + let op = Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![result_ty], + vec![operand], + vec![], + 0, + ); + let this = Self { op }; + this.set_attr_nvvm_scalar_math_format(ctx, format); + this.set_attr_nvvm_scalar_math_operation(ctx, operation); + this.set_attr_nvvm_scalar_math_precision(ctx, precision); + this.set_attr_nvvm_scalar_math_subnormal(ctx, subnormal); + this.get_operation() + } +} + +impl Verify for ScalarMathOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + let Some(format) = self.get_attr_nvvm_scalar_math_format(ctx) else { + return verify_err!(op.loc(), "nvvm.scalar_math requires a format"); + }; + let Some(operation) = self.get_attr_nvvm_scalar_math_operation(ctx) else { + return verify_err!(op.loc(), "nvvm.scalar_math requires an operation"); + }; + let Some(precision) = self.get_attr_nvvm_scalar_math_precision(ctx) else { + return verify_err!(op.loc(), "nvvm.scalar_math requires a precision"); + }; + let Some(subnormal) = self.get_attr_nvvm_scalar_math_subnormal(ctx) else { + return verify_err!(op.loc(), "nvvm.scalar_math requires a subnormal mode"); + }; + let admitted = matches!( + (&*format, &*operation, &*precision, &*subnormal), +"##, + ); + for (index, record) in scalar_maths(catalog).enumerate() { + if index != 0 { + output.push_str(" |\n"); + } + writeln!( + output, + " ({}, {}, {}, {})", + scalar_math_format_attr(record), + scalar_math_operation_attr(record), + scalar_math_precision_attr(record), + scalar_math_subnormal_attr(record), + ) + .unwrap(); + } + output.push_str( + r##" ); + if !admitted { + return verify_err!(op.loc(), "nvvm.scalar_math variant is not admitted"); + } + if op.get_num_operands() != 1 || op.get_num_results() != 1 { + return verify_err!( + op.loc(), + "nvvm.scalar_math requires one operand and one result", + ); + } + let type_matches = |ty: pliron::r#type::TypeHandle| match &*format { + ScalarMathFormatAttr::F16 => ty.deref(ctx).downcast_ref::().is_some_and(|integer| integer.width() == 16), + ScalarMathFormatAttr::F32 => ty.deref(ctx).downcast_ref::().is_some(), + ScalarMathFormatAttr::F64 => ty.deref(ctx).downcast_ref::().is_some(), + }; + if !type_matches(op.get_operand(0).get_type(ctx)) + || !type_matches(op.get_result(0).get_type(ctx)) + { + return verify_err!(op.loc(), "nvvm.scalar_math types do not match its format"); + } + Ok(()) + } +} + +pub(super) fn register(ctx: &mut Context) { + ScalarMathFormatAttr::register(ctx); + ScalarMathOperationAttr::register(ctx); + ScalarMathPrecisionAttr::register(ctx); + ScalarMathSubnormalAttr::register(ctx); + ScalarMathOp::register(ctx); +} +"##, + ); + output +} + +pub(in crate::render) fn render_dialect_extended_minmax( + catalog: &CatalogFile, + hash: &str, +) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + r##"//! One structural operation for extended floating-point min/max. + +use pliron::{ + attribute::Attribute, + builtin::{ + op_interfaces::{NOpdsInterface, NResultsInterface}, + types::{FP32Type, IntegerType, Signedness}, + }, + 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_attr, pliron_op}; + +#[pliron_attr(name = "nvvm.extended_minmax_format", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum ExtendedMinMaxFormatAttr { F32, F16, Bf16, F16x2, Bf16x2 } + +#[pliron_attr(name = "nvvm.extended_minmax_operation", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum ExtendedMinMaxOperationAttr { Min, Max } + +#[pliron_attr(name = "nvvm.extended_minmax_subnormal", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum ExtendedMinMaxSubnormalAttr { Preserve, Ftz } + +#[pliron_attr(name = "nvvm.extended_minmax_nan", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum ExtendedMinMaxNanAttr { Number, Nan } + +#[pliron_attr(name = "nvvm.extended_minmax_xorsign_abs", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum ExtendedMinMaxXorSignAbsAttr { Disabled, Enabled } + +/// Extended min/max with exact PTX modifiers. +#[pliron_op( + name = "nvvm.extended_minmax", + format, + interfaces = [NOpdsInterface<2>, NResultsInterface<1>], + attributes = ( + nvvm_extended_minmax_format: ExtendedMinMaxFormatAttr, + nvvm_extended_minmax_operation: ExtendedMinMaxOperationAttr, + nvvm_extended_minmax_subnormal: ExtendedMinMaxSubnormalAttr, + nvvm_extended_minmax_nan: ExtendedMinMaxNanAttr, + nvvm_extended_minmax_xorsign_abs: ExtendedMinMaxXorSignAbsAttr + ) +)] +pub struct ExtendedMinMaxOp; + +impl ExtendedMinMaxOp { + pub fn new(op: Ptr) -> Self { Self { op } } + + #[allow(clippy::too_many_arguments)] + pub fn build( + ctx: &mut Context, + a: Value, + b: Value, + format: ExtendedMinMaxFormatAttr, + operation: ExtendedMinMaxOperationAttr, + subnormal: ExtendedMinMaxSubnormalAttr, + nan: ExtendedMinMaxNanAttr, + xorsign_abs: ExtendedMinMaxXorSignAbsAttr, + ) -> Ptr { + let result_ty = match &format { + ExtendedMinMaxFormatAttr::F32 => FP32Type::get(ctx).into(), + ExtendedMinMaxFormatAttr::F16 | ExtendedMinMaxFormatAttr::Bf16 => { + IntegerType::get(ctx, 16, Signedness::Unsigned).into() + } + ExtendedMinMaxFormatAttr::F16x2 | ExtendedMinMaxFormatAttr::Bf16x2 => { + IntegerType::get(ctx, 32, Signedness::Unsigned).into() + } + }; + let op = Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![result_ty], + vec![a, b], + vec![], + 0, + ); + let this = Self { op }; + this.set_attr_nvvm_extended_minmax_format(ctx, format); + this.set_attr_nvvm_extended_minmax_operation(ctx, operation); + this.set_attr_nvvm_extended_minmax_subnormal(ctx, subnormal); + this.set_attr_nvvm_extended_minmax_nan(ctx, nan); + this.set_attr_nvvm_extended_minmax_xorsign_abs(ctx, xorsign_abs); + this.get_operation() + } +} + +impl Verify for ExtendedMinMaxOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + let Some(format) = self.get_attr_nvvm_extended_minmax_format(ctx) else { + return verify_err!(op.loc(), "nvvm.extended_minmax requires a format"); + }; + let Some(operation) = self.get_attr_nvvm_extended_minmax_operation(ctx) else { + return verify_err!(op.loc(), "nvvm.extended_minmax requires an operation"); + }; + let Some(subnormal) = self.get_attr_nvvm_extended_minmax_subnormal(ctx) else { + return verify_err!(op.loc(), "nvvm.extended_minmax requires a subnormal mode"); + }; + let Some(nan) = self.get_attr_nvvm_extended_minmax_nan(ctx) else { + return verify_err!(op.loc(), "nvvm.extended_minmax requires a NaN mode"); + }; + let Some(xorsign_abs) = self.get_attr_nvvm_extended_minmax_xorsign_abs(ctx) else { + return verify_err!(op.loc(), "nvvm.extended_minmax requires an xorsign.abs mode"); + }; + let admitted = matches!( + (&*format, &*operation, &*subnormal, &*nan, &*xorsign_abs), +"##, + ); + for (index, record) in extended_minmax(catalog).enumerate() { + if index != 0 { + output.push_str(" |\n"); + } + writeln!( + output, + " ({}, {}, {}, {}, {})", + extended_minmax_format_attr(record), + extended_minmax_operation_attr(record), + extended_minmax_subnormal_attr(record), + extended_minmax_nan_attr(record), + extended_minmax_xorsign_abs_attr(record), + ) + .unwrap(); + } + output.push_str( + r##" ); + if !admitted { + return verify_err!(op.loc(), "nvvm.extended_minmax variant is not admitted"); + } + let type_matches = |ty: pliron::r#type::TypeHandle| match &*format { + ExtendedMinMaxFormatAttr::F32 => ty.deref(ctx).downcast_ref::().is_some(), + ExtendedMinMaxFormatAttr::F16 | ExtendedMinMaxFormatAttr::Bf16 => ty + .deref(ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 16), + ExtendedMinMaxFormatAttr::F16x2 | ExtendedMinMaxFormatAttr::Bf16x2 => ty + .deref(ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 32), + }; + if !type_matches(op.get_operand(0).get_type(ctx)) + || !type_matches(op.get_operand(1).get_type(ctx)) + || !type_matches(op.get_result(0).get_type(ctx)) + { + return verify_err!(op.loc(), "nvvm.extended_minmax types do not match its format"); + } + Ok(()) + } +} + +pub(super) fn register(ctx: &mut Context) { + ExtendedMinMaxFormatAttr::register(ctx); + ExtendedMinMaxOperationAttr::register(ctx); + ExtendedMinMaxSubnormalAttr::register(ctx); + ExtendedMinMaxNanAttr::register(ctx); + ExtendedMinMaxXorSignAbsAttr::register(ctx); + ExtendedMinMaxOp::register(ctx); +} +"##, + ); + output +} + +pub(in crate::render) fn render_dialect_prmt(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + r##"//! One structural operation for the closed generated `prmt` family. + +use pliron::{ + attribute::Attribute, + builtin::{ + op_interfaces::NResultsInterface, + types::{IntegerType, Signedness}, + }, + 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_attr, pliron_op}; + +#[pliron_attr(name = "nvvm.prmt_mode", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub enum PrmtModeAttr { + Generic, + F4e, + B4e, + Rc8, + Ecl, + Ecr, + Rc16, +} + +/// Byte permutation whose exact mode is carried by an attribute. +#[pliron_op( + name = "nvvm.prmt", + format, + interfaces = [NResultsInterface<1>], + attributes = (nvvm_prmt_mode: PrmtModeAttr) +)] +pub struct PrmtOp; + +impl PrmtOp { + pub fn new(op: Ptr) -> Self { + Self { op } + } + + pub fn build(ctx: &mut Context, operands: Vec, mode: PrmtModeAttr) -> Ptr { + let result_ty = IntegerType::get(ctx, 32, Signedness::Unsigned); + let op = Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![result_ty.into()], + operands, + vec![], + 0, + ); + let this = Self { op }; + this.set_attr_nvvm_prmt_mode(ctx, mode); + this.get_operation() + } +} + +impl Verify for PrmtOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + let Some(mode) = self.get_attr_nvvm_prmt_mode(ctx) else { + return verify_err!(op.loc(), "nvvm.prmt requires a mode attribute"); + }; + let expected_operands = match &*mode { + PrmtModeAttr::Generic | PrmtModeAttr::F4e | PrmtModeAttr::B4e => 3, + PrmtModeAttr::Rc8 | PrmtModeAttr::Ecl | PrmtModeAttr::Ecr | PrmtModeAttr::Rc16 => 2, + }; + if op.get_num_operands() != expected_operands || op.get_num_results() != 1 { + return verify_err!( + op.loc(), + "nvvm.prmt {:?} requires {} operands and one result", + mode, + expected_operands + ); + } + for index in 0..expected_operands { + let ty = op.get_operand(index).get_type(ctx); + if ty + .deref(ctx) + .downcast_ref::() + .is_none_or(|integer| integer.width() != 32) + { + return verify_err!(op.loc(), "nvvm.prmt operand {} must be a 32-bit integer", index); + } + } + let result_ty = op.get_result(0).get_type(ctx); + if result_ty + .deref(ctx) + .downcast_ref::() + .is_none_or(|integer| integer.width() != 32) + { + return verify_err!(op.loc(), "nvvm.prmt result must be a 32-bit integer"); + } + Ok(()) + } +} + +pub(super) fn register(ctx: &mut Context) { + PrmtModeAttr::register(ctx); + PrmtOp::register(ctx); +} +"##, + ); + output +} diff --git a/crates/cuda-intrinsics-gen/src/render/dialect/warp_sync.rs b/crates/cuda-intrinsics-gen/src/render/dialect/warp_sync.rs new file mode 100644 index 0000000000..0c2be7cc00 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/render/dialect/warp_sync.rs @@ -0,0 +1,1132 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + ActiveMaskAdapter, CatalogFile, CpAsyncControlOperation, CpAsyncSourceSize, + MbarrierBasicAdapter, MbarrierBasicOperation, MbarrierExtendedOperation, MbarrierStateSpace, + VoteMode, WarpBarrierAdapter, WarpMatchMode, WarpShuffleAdapter, WarpShuffleValueKind, +}; +use crate::render::common::{rust_header, source_label}; +use crate::render::families::{ + active_masks, cp_async_controls, cp_async_copies, cp_async_mbarriers, dot_product_ptx, + dot_products, elect_intrinsics, mbarrier_basics, mbarrier_extended, redux, sregs, + sync_intrinsics, vote_intrinsics, warp_barriers, warp_matches, warp_shuffles, +}; +use std::fmt::Write as _; + +pub(in crate::render) fn render_dialect_elect(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + "//! Generated warp leader-election operation.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::IntegerType,\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::{TypeHandle, Typed},\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\n", + ); + for record in elect_intrinsics(catalog) { + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!( + output, + "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<1>, NResultsInterface<2>],\n)]", + record.dialect.op_name + ) + .unwrap(); + writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); + writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); + output.push_str(" pub fn new(op: Ptr) -> Self { Self { op } }\n}\n\n"); + writeln!(output, "impl Verify for {} {{", record.dialect.op_type).unwrap(); + output.push_str( + " fn verify(&self, ctx: &Context) -> Result<(), Error> {\n let op = self.get_operation().deref(ctx);\n if op.get_num_operands() != 1 || op.get_num_results() != 2 {\n return verify_err!(op.loc(), \"nvvm.elect_sync requires one i32 operand and results [i32, i1]\");\n }\n if !is_integer_width(ctx, op.get_operand(0).get_type(ctx), 32)\n || !is_integer_width(ctx, op.get_result(0).get_type(ctx), 32)\n || !is_integer_width(ctx, op.get_result(1).get_type(ctx), 1)\n {\n return verify_err!(op.loc(), \"nvvm.elect_sync requires one i32 operand and results [i32, i1]\");\n }\n Ok(())\n }\n}\n\n", + ); + } + output.push_str( + "fn is_integer_width(ctx: &Context, ty: TypeHandle, width: u32) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == width)\n}\n\npub(super) fn register(ctx: &mut Context) {\n", + ); + for record in elect_intrinsics(catalog) { + writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); + } + output.push_str("}\n"); + output +} + +pub(in crate::render) fn render_dialect_sreg(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + "//! Structural NVVM operations for generated special-register reads.\n\nuse pliron::{\n builtin::op_interfaces::{NOpdsInterface, NResultsInterface},\n builtin::types::IntegerType,\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\n", + ); + for record in sregs(catalog) { + let width = record.scalar_width().unwrap(); + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!( + output, + "/// Catalog ID `{}`; {} returns one `i{width}` result.", + record.id, + source_label(record) + ) + .unwrap(); + writeln!( + output, + "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<0>, NResultsInterface<1>],\n)]", + record.dialect.op_name + ) + .unwrap(); + writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); + writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); + output.push_str(" pub fn new(op: Ptr) -> Self {\n"); + writeln!(output, " Self {{ op }}").unwrap(); + output.push_str(" }\n}\n"); + writeln!(output, "\nimpl Verify for {} {{", record.dialect.op_type).unwrap(); + output.push_str(" fn verify(&self, ctx: &Context) -> Result<(), Error> {\n"); + writeln!( + output, + " verify_scalar_result(ctx, self.get_operation(), {:?}, {width})", + record.dialect.op_name + ) + .unwrap(); + output.push_str(" }\n}\n\n"); + } + output.push_str( + "fn verify_scalar_result(\n ctx: &Context,\n op: Ptr,\n name: &str,\n width: u32,\n) -> Result<(), Error> {\n let op = op.deref(ctx);\n let ty = op.get_result(0).get_type(ctx);\n let ty_object = ty.deref(ctx);\n let Some(integer) = ty_object.downcast_ref::() else {\n return verify_err!(op.loc(), \"{} result must be an integer\", name);\n };\n if integer.width() != width {\n return verify_err!(\n op.loc(),\n \"{} result must be a {}-bit integer\",\n name,\n width\n );\n }\n Ok(())\n}\n\npub(super) fn register(ctx: &mut Context) {\n", + ); + for record in sregs(catalog) { + writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); + } + output.push_str("}\n"); + output +} + +pub(in crate::render) fn render_dialect_sync(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + "//! Structural operation for generated CTA synchronization.\n\nuse pliron::{\n builtin::op_interfaces::{NOpdsInterface, NResultsInterface},\n context::{Context, Ptr},\n op::Op,\n operation::Operation,\n};\nuse pliron_derive::pliron_op;\n\n", + ); + for record in sync_intrinsics(catalog) { + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!( + output, + "#[pliron_op(\n name = {:?},\n format,\n verifier = \"succ\",\n interfaces = [NOpdsInterface<0>, NResultsInterface<0>],\n)]", + record.dialect.op_name + ) + .unwrap(); + writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); + writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); + output.push_str( + " pub fn new(op: Ptr) -> Self {\n Self { op }\n }\n}\n\n", + ); + } + output.push_str("pub(super) fn register(ctx: &mut Context) {\n"); + for record in sync_intrinsics(catalog) { + writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); + } + output.push_str("}\n"); + output +} + +pub(in crate::render) fn render_dialect_vote(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + "//! Structural operations for the generated `vote.sync` family.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_integer_width(ctx: &Context, ty: pliron::r#type::TypeHandle, width: u32) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == width)\n}\n\n", + ); + for record in vote_intrinsics(catalog) { + let result_width = match record.vote.as_ref().unwrap().mode { + VoteMode::All | VoteMode::Any | VoteMode::Uni => 1, + VoteMode::Ballot => 32, + }; + let result_signedness = if result_width == 32 { + "Unsigned" + } else { + "Signless" + }; + writeln!(output, "/// {}", record.summary).unwrap(); + output.push_str( + "///\n/// Operands are `[member_mask, predicate]`. The generated verifier keeps\n/// the mask, predicate, and result types exact.\n", + ); + writeln!( + output, + "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<2>, NResultsInterface<1>],\n)]", + record.dialect.op_name + ) + .unwrap(); + writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); + writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); + writeln!( + output, + " pub fn new(op: Ptr) -> Self {{\n Self {{ op }}\n }}\n\n pub fn build(ctx: &mut Context, member_mask: Value, predicate: Value) -> Ptr {{\n let result_ty = IntegerType::get(ctx, {result_width}, Signedness::{result_signedness});\n Operation::new(\n ctx,\n Self::get_concrete_op_info(),\n vec![result_ty.into()],\n vec![member_mask, predicate],\n vec![],\n 0,\n )\n }}\n}}" + ) + .unwrap(); + writeln!(output, "\nimpl Verify for {} {{", record.dialect.op_type).unwrap(); + writeln!( + output, + " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n let op = self.get_operation().deref(ctx);\n if op.get_num_operands() != 2 || op.get_num_results() != 1 {{\n return verify_err!(op.loc(), {:?});\n }}\n if !is_integer_width(ctx, op.get_operand(0).get_type(ctx), 32)\n || !is_integer_width(ctx, op.get_operand(1).get_type(ctx), 1)\n || !is_integer_width(ctx, op.get_result(0).get_type(ctx), {result_width})\n {{\n return verify_err!(op.loc(), {:?});\n }}\n Ok(())\n }}\n}}\n", + format!( + "{} requires exactly two operands [member_mask, predicate] and one result", + record.dialect.op_name + ), + format!( + "{} requires i32 member mask, i1 predicate, and i{result_width} result", + record.dialect.op_name + ), + ) + .unwrap(); + } + output.push_str("\npub(super) fn register(ctx: &mut Context) {\n"); + for record in vote_intrinsics(catalog) { + writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); + } + output.push_str("}\n"); + output +} + +pub(in crate::render) fn render_dialect_active_mask(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + "//! Structural operation for the generated active warp mask.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\n", + ); + for record in active_masks(catalog) { + debug_assert_eq!( + record.active_mask.as_ref().unwrap().adapter, + ActiveMaskAdapter::DirectZeroOperandMask + ); + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!( + output, + "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<0>, NResultsInterface<1>],\n)]", + record.dialect.op_name + ) + .unwrap(); + writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); + writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); + output.push_str( + " pub fn new(op: Ptr) -> Self {\n Self { op }\n }\n\n pub fn build(ctx: &mut Context) -> Ptr {\n let result_ty = IntegerType::get(ctx, 32, Signedness::Unsigned);\n Operation::new(\n ctx,\n Self::get_concrete_op_info(),\n vec![result_ty.into()],\n vec![],\n vec![],\n 0,\n )\n }\n}\n", + ); + writeln!(output, "\nimpl Verify for {} {{", record.dialect.op_type).unwrap(); + writeln!( + output, + " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n let op = self.get_operation().deref(ctx);\n let valid = op.get_num_operands() == 0\n && op.get_num_results() == 1\n && op.get_result(0).get_type(ctx).deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32);\n if !valid {{\n return verify_err!(op.loc(), {:?});\n }}\n Ok(())\n }}\n}}\n", + format!("{} requires no operands and one i32 result", record.dialect.op_name) + ) + .unwrap(); + } + output.push_str("\npub(super) fn register(ctx: &mut Context) {\n"); + for record in active_masks(catalog) { + writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); + } + output.push_str("}\n"); + output +} + +pub(in crate::render) fn render_dialect_cp_async_copy(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + r#"//! Structural operations for classic `cp.async` instructions. + +use dialect_mir::{ops::MirConstantOp, types::{MirPtrType, address_space}}; +use pliron::{ + builtin::{ + op_interfaces::{NOpdsInterface, NResultsInterface}, + ops::ConstantOp, + types::{IntegerType, Signedness}, + }, + 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; + +fn verify_pointer( + ctx: &Context, + op: &Operation, + value: Value, + role: &str, + allowed_address_spaces: &[u32], +) -> Result<(), Error> { + let ty = value.get_type(ctx); + let ty = ty.deref(ctx); + let Some(pointer) = ty.downcast_ref::() else { + return verify_err!(op.loc(), "{role} must be a MIR pointer"); + }; + if !allowed_address_spaces.contains(&pointer.address_space) { + return verify_err!(op.loc(), "{role} has the wrong address space"); + } + Ok(()) +} + +fn verify_cp_async_copy( + ctx: &Context, + operation: Ptr, + name: &str, + has_source_size: bool, +) -> Result<(), Error> { + let op = operation.deref(ctx); + let expected_operands = if has_source_size { 3 } else { 2 }; + if op.get_num_operands() != expected_operands || op.get_num_results() != 0 { + return verify_err!(op.loc(), "{name} has the wrong operand or result count"); + } + verify_pointer( + ctx, + &op, + op.get_operand(0), + "shared destination", + &[address_space::GENERIC, address_space::SHARED], + )?; + verify_pointer( + ctx, + &op, + op.get_operand(1), + "global source", + &[address_space::GENERIC, address_space::GLOBAL], + )?; + if has_source_size { + let ty = op.get_operand(2).get_type(ctx); + let ty = ty.deref(ctx); + let Some(integer) = ty.downcast_ref::() else { + return verify_err!(op.loc(), "source size must be u32"); + }; + if integer.width() != 32 || integer.signedness() != Signedness::Unsigned { + return verify_err!(op.loc(), "source size must be u32"); + } + } + Ok(()) +} + +fn verify_cp_async_control( + ctx: &Context, + operation: Ptr, + name: &str, + has_immediate: bool, +) -> Result<(), Error> { + let op = operation.deref(ctx); + let expected_operands = usize::from(has_immediate); + if op.get_num_operands() != expected_operands || op.get_num_results() != 0 { + return verify_err!(op.loc(), "{name} has the wrong operand or result count"); + } + if has_immediate { + let value = op.get_operand(0); + let ty = value.get_type(ctx); + let ty = ty.deref(ctx); + let Some(integer) = ty.downcast_ref::() else { + return verify_err!(op.loc(), "maximum pending group count must be u32"); + }; + if integer.width() != 32 || integer.signedness() != Signedness::Unsigned { + return verify_err!(op.loc(), "maximum pending group count must be u32"); + } + let Some(defining_op) = value.defining_op() else { + return verify_err!(op.loc(), "maximum pending group count must be a compile-time constant"); + }; + if Operation::get_op::(defining_op, ctx).is_none() + && Operation::get_op::(defining_op, ctx).is_none() + { + return verify_err!(op.loc(), "maximum pending group count must be a compile-time constant"); + } + } + Ok(()) +} + +fn verify_cp_async_mbarrier( + ctx: &Context, + operation: Ptr, + name: &str, +) -> Result<(), Error> { + let op = operation.deref(ctx); + if op.get_num_operands() != 1 || op.get_num_results() != 0 { + return verify_err!(op.loc(), "{name} has the wrong operand or result count"); + } + let ty = op.get_operand(0).get_type(ctx); + let ty = ty.deref(ctx); + let Some(pointer) = ty.downcast_ref::() else { + return verify_err!(op.loc(), "mbarrier address must be a MIR pointer"); + }; + let pointee = pointer.pointee.deref(ctx); + let valid_pointee = pointee + .downcast_ref::() + .is_some_and(|integer| { + integer.width() == 64 && integer.signedness() == Signedness::Unsigned + }); + if !pointer.is_mutable + || !valid_pointee + || !matches!( + pointer.address_space, + address_space::GENERIC | address_space::SHARED + ) + { + return verify_err!( + op.loc(), + "mbarrier address must be a mutable generic/shared pointer to u64" + ); + } + Ok(()) +} + +"#, + ); + for record in cp_async_copies(catalog) { + let copy = record.cp_async_copy.as_ref().unwrap(); + let dynamic = copy.source_size == CpAsyncSourceSize::Runtime; + let operand_count = if dynamic { 3 } else { 2 }; + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!(output, "/// Lowers to `{}`.", record.expected_ptx).unwrap(); + writeln!( + output, + "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<{operand_count}>, NResultsInterface<0>],\n)]\npub struct {};", + record.dialect.op_name, record.dialect.op_type + ) + .unwrap(); + writeln!(output, "impl {} {{", record.dialect.op_type).unwrap(); + output.push_str(" pub fn new(op: Ptr) -> Self { Self { op } }\n\n"); + if dynamic { + output.push_str( + " pub fn build(ctx: &mut Context, shared_dst: Value, global_src: Value, source_size: Value) -> Ptr {\n Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![shared_dst, global_src, source_size], vec![], 0)\n }\n", + ); + } else { + output.push_str( + " pub fn build(ctx: &mut Context, shared_dst: Value, global_src: Value) -> Ptr {\n Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![shared_dst, global_src], vec![], 0)\n }\n", + ); + } + output.push_str("}\n\n"); + writeln!(output, "impl Verify for {} {{", record.dialect.op_type).unwrap(); + writeln!( + output, + " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n verify_cp_async_copy(ctx, self.get_operation(), {:?}, {dynamic})\n }}\n}}\n", + record.dialect.op_name + ) + .unwrap(); + } + for record in cp_async_controls(catalog) { + let control = record.cp_async_control.as_ref().unwrap(); + let has_immediate = control.operation == CpAsyncControlOperation::WaitGroup; + let operand_count = usize::from(has_immediate); + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!(output, "/// Lowers to `{}`.", record.expected_ptx).unwrap(); + writeln!( + output, + "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<{operand_count}>, NResultsInterface<0>],\n)]\npub struct {};", + record.dialect.op_name, record.dialect.op_type + ) + .unwrap(); + writeln!(output, "impl {} {{", record.dialect.op_type).unwrap(); + output.push_str(" pub fn new(op: Ptr) -> Self { Self { op } }\n\n"); + if has_immediate { + output.push_str( + " pub fn build(ctx: &mut Context, max_pending: Value) -> Ptr {\n Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![max_pending], vec![], 0)\n }\n", + ); + } else { + output.push_str( + " pub fn build(ctx: &mut Context) -> Ptr {\n Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![], vec![], 0)\n }\n", + ); + } + output.push_str("}\n\n"); + writeln!(output, "impl Verify for {} {{", record.dialect.op_type).unwrap(); + writeln!( + output, + " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n verify_cp_async_control(ctx, self.get_operation(), {:?}, {has_immediate})\n }}\n}}\n", + record.dialect.op_name + ) + .unwrap(); + } + for record in cp_async_mbarriers(catalog) { + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!(output, "/// Lowers to `{}`.", record.expected_ptx).unwrap(); + writeln!( + output, + "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<1>, NResultsInterface<0>],\n)]\npub struct {};", + record.dialect.op_name, record.dialect.op_type + ) + .unwrap(); + writeln!(output, "impl {} {{", record.dialect.op_type).unwrap(); + output.push_str( + " pub fn new(op: Ptr) -> Self { Self { op } }\n\n pub fn build(ctx: &mut Context, barrier: Value) -> Ptr {\n Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![barrier], vec![], 0)\n }\n}\n\n", + ); + writeln!(output, "impl Verify for {} {{", record.dialect.op_type).unwrap(); + writeln!( + output, + " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n verify_cp_async_mbarrier(ctx, self.get_operation(), {:?})\n }}\n}}\n", + record.dialect.op_name + ) + .unwrap(); + } + output.push_str("pub(super) fn register(ctx: &mut Context) {\n"); + for record in cp_async_copies(catalog) { + writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); + } + for record in cp_async_controls(catalog) { + writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); + } + for record in cp_async_mbarriers(catalog) { + writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); + } + output.push_str("}\n"); + output +} + +pub(in crate::render) fn render_dialect_mbarrier_basic( + catalog: &CatalogFile, + hash: &str, +) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + r#"//! Structural operations for the basic shared-memory mbarrier lifecycle. + +use dialect_mir::types::{MirPtrType, address_space}; +use pliron::{ + builtin::{ + op_interfaces::{NOpdsInterface, NResultsInterface}, + types::{IntegerType, Signedness}, + }, + 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; + +#[derive(Clone, Copy)] +enum MbarrierBasicShape { + Init, + Arrive, + TestWait, + Inval, +} + +fn is_integer( + ctx: &Context, + value: Value, + width: u32, + signedness: Signedness, +) -> bool { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_some_and(|integer| { + integer.width() == width && integer.signedness() == signedness + }) +} + +fn verify_barrier_pointer(ctx: &Context, op: &Operation, value: Value) -> Result<(), Error> { + let ty = value.get_type(ctx); + let ty = ty.deref(ctx); + let Some(pointer) = ty.downcast_ref::() else { + return verify_err!(op.loc(), "mbarrier address must be a MIR pointer"); + }; + if !matches!(pointer.address_space, address_space::GENERIC | address_space::SHARED) { + return verify_err!(op.loc(), "mbarrier address must be generic or shared"); + } + Ok(()) +} + +fn verify_mbarrier_basic( + ctx: &Context, + operation: Ptr, + name: &str, + shape: MbarrierBasicShape, +) -> Result<(), Error> { + let op = operation.deref(ctx); + let (operands, results) = match shape { + MbarrierBasicShape::Init | MbarrierBasicShape::TestWait => (2, usize::from(matches!(shape, MbarrierBasicShape::TestWait))), + MbarrierBasicShape::Arrive => (1, 1), + MbarrierBasicShape::Inval => (1, 0), + }; + if op.get_num_operands() != operands || op.get_num_results() != results { + return verify_err!(op.loc(), "{name} has the wrong operand or result count"); + } + verify_barrier_pointer(ctx, &op, op.get_operand(0))?; + match shape { + MbarrierBasicShape::Init => { + if !is_integer(ctx, op.get_operand(1), 32, Signedness::Unsigned) { + return verify_err!(op.loc(), "mbarrier expected count must be u32"); + } + } + MbarrierBasicShape::Arrive => { + if !is_integer(ctx, op.get_result(0), 64, Signedness::Unsigned) { + return verify_err!(op.loc(), "mbarrier arrival token must be u64"); + } + } + MbarrierBasicShape::TestWait => { + if !is_integer(ctx, op.get_operand(1), 64, Signedness::Unsigned) + || !is_integer(ctx, op.get_result(0), 1, Signedness::Signless) + { + return verify_err!(op.loc(), "mbarrier test-wait requires a u64 token and i1 result"); + } + } + MbarrierBasicShape::Inval => {} + } + Ok(()) +} + +"#, + ); + for record in mbarrier_basics(catalog) { + let mbarrier = record.mbarrier_basic.as_ref().unwrap(); + debug_assert_eq!(mbarrier.state_space, MbarrierStateSpace::Shared); + let (operand_count, result_count, shape, build) = match mbarrier.operation { + MbarrierBasicOperation::Init => { + debug_assert_eq!( + mbarrier.adapter, + MbarrierBasicAdapter::InitPointerCountToVoid + ); + ( + 2, + 0, + "Init", + " pub fn build(ctx: &mut Context, barrier: Value, expected_count: Value) -> Ptr {\n Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![barrier, expected_count], vec![], 0)\n }\n", + ) + } + MbarrierBasicOperation::Arrive => { + debug_assert_eq!(mbarrier.adapter, MbarrierBasicAdapter::ArrivePointerToToken); + ( + 1, + 1, + "Arrive", + " pub fn build(ctx: &mut Context, barrier: 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], vec![], 0)\n }\n", + ) + } + MbarrierBasicOperation::TestWait => { + debug_assert_eq!( + mbarrier.adapter, + MbarrierBasicAdapter::TestWaitPointerTokenToPredicate + ); + ( + 2, + 1, + "TestWait", + " pub fn build(ctx: &mut Context, barrier: Value, token: Value) -> Ptr {\n let predicate_ty = IntegerType::get(ctx, 1, Signedness::Signless);\n Operation::new(ctx, Self::get_concrete_op_info(), vec![predicate_ty.into()], vec![barrier, token], vec![], 0)\n }\n", + ) + } + MbarrierBasicOperation::Inval => { + debug_assert_eq!(mbarrier.adapter, MbarrierBasicAdapter::InvalPointerToVoid); + ( + 1, + 0, + "Inval", + " pub fn build(ctx: &mut Context, barrier: Value) -> Ptr {\n Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![barrier], vec![], 0)\n }\n", + ) + } + }; + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!(output, "/// Lowers to `{}`.", record.expected_ptx).unwrap(); + writeln!( + output, + "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<{operand_count}>, NResultsInterface<{result_count}>],\n)]\npub struct {};", + record.dialect.op_name, record.dialect.op_type + ) + .unwrap(); + writeln!(output, "impl {} {{", record.dialect.op_type).unwrap(); + output.push_str(" pub fn new(op: Ptr) -> Self { Self { op } }\n\n"); + output.push_str(build); + output.push_str("}\n\n"); + writeln!(output, "impl Verify for {} {{", record.dialect.op_type).unwrap(); + writeln!( + output, + " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n verify_mbarrier_basic(ctx, self.get_operation(), {:?}, MbarrierBasicShape::{shape})\n }}\n}}\n", + record.dialect.op_name + ) + .unwrap(); + } + output.push_str("pub(super) fn register(ctx: &mut Context) {\n"); + for record in mbarrier_basics(catalog) { + writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); + } + output.push_str("}\n"); + output +} + +pub(in crate::render) fn render_dialect_mbarrier_extended( + catalog: &CatalogFile, + hash: &str, +) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + r#"//! Structural operations for extended mbarrier and async-proxy instructions. + +use dialect_mir::types::{MirPtrType, address_space}; +use pliron::{ + builtin::{ + op_interfaces::{NOpdsInterface, NResultsInterface}, + types::{IntegerType, Signedness}, + }, + 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; + +#[derive(Clone, Copy)] +enum MbarrierExtendedShape { + PointerU32ToU64, + RawU64ToVoid, + PointerU64ToI1, + PointerU32ToI1, + ZeroToVoid, + U32ToVoid, +} + +fn is_integer(ctx: &Context, value: Value, width: u32, signedness: Signedness) -> bool { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_some_and(|integer| { + integer.width() == width && integer.signedness() == signedness + }) +} + +fn verify_barrier_pointer(ctx: &Context, op: &Operation, value: Value) -> Result<(), Error> { + let ty = value.get_type(ctx); + let ty = ty.deref(ctx); + let Some(pointer) = ty.downcast_ref::() else { + return verify_err!(op.loc(), "mbarrier address must be a MIR pointer"); + }; + if !matches!(pointer.address_space, address_space::GENERIC | address_space::SHARED) { + return verify_err!(op.loc(), "mbarrier address must be generic or shared"); + } + Ok(()) +} + +fn verify_mbarrier_extended( + ctx: &Context, + operation: Ptr, + name: &str, + shape: MbarrierExtendedShape, +) -> Result<(), Error> { + let op = operation.deref(ctx); + let (operands, results) = match shape { + MbarrierExtendedShape::PointerU32ToU64 + | MbarrierExtendedShape::PointerU64ToI1 + | MbarrierExtendedShape::PointerU32ToI1 => (2, 1), + MbarrierExtendedShape::RawU64ToVoid | MbarrierExtendedShape::U32ToVoid => (1, 0), + MbarrierExtendedShape::ZeroToVoid => (0, 0), + }; + if op.get_num_operands() != operands || op.get_num_results() != results { + return verify_err!(op.loc(), "{name} has the wrong operand or result count"); + } + match shape { + MbarrierExtendedShape::PointerU32ToU64 => { + verify_barrier_pointer(ctx, &op, op.get_operand(0))?; + if !is_integer(ctx, op.get_operand(1), 32, Signedness::Unsigned) + || !is_integer(ctx, op.get_result(0), 64, Signedness::Unsigned) + { + return verify_err!(op.loc(), "mbarrier arrival requires u32 bytes and a u64 token"); + } + } + MbarrierExtendedShape::PointerU64ToI1 => { + verify_barrier_pointer(ctx, &op, op.get_operand(0))?; + if !is_integer(ctx, op.get_operand(1), 64, Signedness::Unsigned) + || !is_integer(ctx, op.get_result(0), 1, Signedness::Signless) + { + return verify_err!(op.loc(), "mbarrier wait requires a u64 token and i1 result"); + } + } + MbarrierExtendedShape::PointerU32ToI1 => { + verify_barrier_pointer(ctx, &op, op.get_operand(0))?; + if !is_integer(ctx, op.get_operand(1), 32, Signedness::Unsigned) + || !is_integer(ctx, op.get_result(0), 1, Signedness::Signless) + { + return verify_err!(op.loc(), "mbarrier wait requires u32 parity and i1 result"); + } + } + MbarrierExtendedShape::RawU64ToVoid => { + if !is_integer(ctx, op.get_operand(0), 64, Signedness::Unsigned) { + return verify_err!(op.loc(), "remote mbarrier address must be u64"); + } + } + MbarrierExtendedShape::U32ToVoid => { + if !is_integer(ctx, op.get_operand(0), 32, Signedness::Unsigned) { + return verify_err!(op.loc(), "nanosleep duration must be u32"); + } + } + MbarrierExtendedShape::ZeroToVoid => {} + } + Ok(()) +} + +"#, + ); + for record in mbarrier_extended(catalog) { + let contract = record.mbarrier_extended.as_ref().unwrap(); + let (operand_count, result_count, shape, build) = match contract.operation { + MbarrierExtendedOperation::ArriveExpectTxCta + | MbarrierExtendedOperation::ArriveExpectTxCluster => ( + 2, + 1, + "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 => ( + 1, + 0, + "RawU64ToVoid", + " pub fn build(ctx: &mut Context, address: Value) -> Ptr {\n Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![address], vec![], 0)\n }\n", + ), + MbarrierExtendedOperation::TryWaitTokenCta => ( + 2, + 1, + "PointerU64ToI1", + " 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 => ( + 2, + 1, + "PointerU32ToI1", + " pub fn build(ctx: &mut Context, barrier: Value, parity: 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, parity], vec![], 0)\n }\n", + ), + MbarrierExtendedOperation::FenceProxyAsyncSharedCta + | MbarrierExtendedOperation::FenceMbarrierInitReleaseCluster + | MbarrierExtendedOperation::FenceProxyAsyncGenericReleaseSharedCtaCluster + | MbarrierExtendedOperation::FenceProxyAsyncGenericAcquireSharedClusterCluster => ( + 0, + 0, + "ZeroToVoid", + " pub fn build(ctx: &mut Context) -> Ptr {\n Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![], vec![], 0)\n }\n", + ), + MbarrierExtendedOperation::Nanosleep => ( + 1, + 0, + "U32ToVoid", + " pub fn build(ctx: &mut Context, ns: Value) -> Ptr {\n Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![ns], vec![], 0)\n }\n", + ), + }; + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!(output, "/// Lowers to `{}`.", record.expected_ptx).unwrap(); + writeln!( + output, + "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<{operand_count}>, NResultsInterface<{result_count}>],\n)]\npub struct {};", + record.dialect.op_name, record.dialect.op_type + ) + .unwrap(); + writeln!(output, "impl {} {{", record.dialect.op_type).unwrap(); + output.push_str(" pub fn new(op: Ptr) -> Self { Self { op } }\n\n"); + output.push_str(build); + output.push_str("}\n\n"); + writeln!(output, "impl Verify for {} {{", record.dialect.op_type).unwrap(); + writeln!( + output, + " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n verify_mbarrier_extended(ctx, self.get_operation(), {:?}, MbarrierExtendedShape::{shape})\n }}\n}}\n", + record.dialect.op_name + ) + .unwrap(); + } + output.push_str("pub(super) fn register(ctx: &mut Context) {\n"); + for record in mbarrier_extended(catalog) { + writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); + } + output.push_str("}\n"); + output +} + +pub(in crate::render) fn render_dialect_warp_match(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + "//! Structural operations for the generated `match.sync` family.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_integer_width(ctx: &Context, ty: pliron::r#type::TypeHandle, width: u32) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == width)\n}\n\n", + ); + for record in warp_matches(catalog) { + let warp_match = record.warp_match.as_ref().unwrap(); + let value_width = warp_match.value_width.bits(); + writeln!(output, "/// {}", record.summary).unwrap(); + output.push_str( + "///\n/// Operands are `[member_mask, value]`; the result is a 32-bit lane mask.\n", + ); + if warp_match.mode == WarpMatchMode::All { + output.push_str( + "/// LLVM also returns a predicate, which the established API discards.\n", + ); + } + writeln!( + output, + "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<2>, NResultsInterface<1>],\n)]", + record.dialect.op_name + ) + .unwrap(); + writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); + writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); + output.push_str( + " pub fn new(op: Ptr) -> Self {\n Self { op }\n }\n\n pub fn build(ctx: &mut Context, member_mask: Value, value: Value) -> Ptr {\n let result_ty = IntegerType::get(ctx, 32, Signedness::Unsigned);\n Operation::new(\n ctx,\n Self::get_concrete_op_info(),\n vec![result_ty.into()],\n vec![member_mask, value],\n vec![],\n 0,\n )\n }\n}\n", + ); + writeln!(output, "\nimpl Verify for {} {{", record.dialect.op_type).unwrap(); + writeln!( + output, + " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n let op = self.get_operation().deref(ctx);\n if op.get_num_operands() != 2 || op.get_num_results() != 1 {{\n return verify_err!(op.loc(), {:?});\n }}\n if !is_integer_width(ctx, op.get_operand(0).get_type(ctx), 32)\n || !is_integer_width(ctx, op.get_operand(1).get_type(ctx), {value_width})\n || !is_integer_width(ctx, op.get_result(0).get_type(ctx), 32)\n {{\n return verify_err!(op.loc(), {:?});\n }}\n Ok(())\n }}\n}}\n", + format!( + "{} requires exactly [member_mask, value] and one mask result", + record.dialect.op_name + ), + format!( + "{} requires i32 member mask, i{value_width} value, and i32 result", + record.dialect.op_name + ), + ) + .unwrap(); + } + output.push_str("\npub(super) fn register(ctx: &mut Context) {\n"); + for record in warp_matches(catalog) { + writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); + } + output.push_str("}\n"); + output +} + +pub(in crate::render) fn render_dialect_warp_barrier(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + "//! Structural operation for generated warp synchronization.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::IntegerType,\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_i32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32)\n}\n\n", + ); + for record in warp_barriers(catalog) { + debug_assert_eq!( + record.warp_barrier.as_ref().unwrap().adapter, + WarpBarrierAdapter::DirectMemberMask + ); + writeln!(output, "/// {}", record.summary).unwrap(); + output.push_str("///\n/// The operand is the 32-bit warp participation mask.\n"); + writeln!( + output, + "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<1>, NResultsInterface<0>],\n)]", + record.dialect.op_name + ) + .unwrap(); + writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); + writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); + output.push_str( + " pub fn new(op: Ptr) -> Self {\n Self { op }\n }\n\n pub fn build(ctx: &mut Context, member_mask: Value) -> Ptr {\n Operation::new(\n ctx,\n Self::get_concrete_op_info(),\n vec![],\n vec![member_mask],\n vec![],\n 0,\n )\n }\n}\n", + ); + writeln!(output, "\nimpl Verify for {} {{", record.dialect.op_type).unwrap(); + writeln!( + output, + " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n let op = self.get_operation().deref(ctx);\n if op.get_num_operands() != 1 || op.get_num_results() != 0 {{\n return verify_err!(op.loc(), {:?});\n }}\n if !is_i32(ctx, op.get_operand(0).get_type(ctx)) {{\n return verify_err!(op.loc(), {:?});\n }}\n Ok(())\n }}\n}}\n", + format!( + "{} requires exactly one member-mask operand and no results", + record.dialect.op_name + ), + format!("{} member mask must be i32", record.dialect.op_name), + ) + .unwrap(); + } + output.push_str("\npub(super) fn register(ctx: &mut Context) {\n"); + for record in warp_barriers(catalog) { + writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); + } + output.push_str("}\n"); + output +} + +pub(in crate::render) fn render_dialect_warp_shuffle(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + "//! Structural operations for the generated `shfl.sync` family.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{FP32Type, IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_i32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32)\n}\n\nfn is_i64(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 64)\n}\n\nfn is_f32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx).downcast_ref::().is_some()\n}\n\n", + ); + for record in warp_shuffles(catalog) { + let shuffle = record.warp_shuffle.as_ref().unwrap(); + debug_assert!(matches!( + (shuffle.value_kind, shuffle.adapter), + ( + WarpShuffleValueKind::I32 | WarpShuffleValueKind::F32, + WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp + ) | ( + WarpShuffleValueKind::I64, + WarpShuffleAdapter::MaskValueLaneOrDeltaSplitI64LowHighB32InsertClampReassemble + ) + )); + let (result_builder, value_check, value_label) = match shuffle.value_kind { + WarpShuffleValueKind::I32 => ( + "IntegerType::get(ctx, 32, Signedness::Unsigned).into()", + "is_i32", + "i32", + ), + WarpShuffleValueKind::F32 => ("FP32Type::get(ctx).into()", "is_f32", "f32"), + WarpShuffleValueKind::I64 => ( + "IntegerType::get(ctx, 64, Signedness::Unsigned).into()", + "is_i64", + "i64", + ), + }; + writeln!(output, "/// {}", record.summary).unwrap(); + output.push_str("///\n/// Operands are `[member_mask, value, lane_or_delta]`.\n"); + if shuffle.value_kind == WarpShuffleValueKind::I64 { + output.push_str( + "/// Lowering splits the value into two `b32` halves, shuffles both, and reassembles it.\n", + ); + } else { + output.push_str( + "/// Generated lowering inserts the fixed clamp required by the selected shuffle mode.\n", + ); + } + writeln!( + output, + "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<3>, NResultsInterface<1>],\n)]", + record.dialect.op_name + ) + .unwrap(); + writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); + writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); + writeln!( + output, + " pub fn new(op: Ptr) -> Self {{\n Self {{ op }}\n }}\n\n pub fn build(\n ctx: &mut Context,\n member_mask: Value,\n value: Value,\n lane_or_delta: Value,\n ) -> Ptr {{\n let result_ty: pliron::r#type::TypeHandle = {result_builder};\n Operation::new(\n ctx,\n Self::get_concrete_op_info(),\n vec![result_ty],\n vec![member_mask, value, lane_or_delta],\n vec![],\n 0,\n )\n }}\n}}" + ) + .unwrap(); + writeln!(output, "\nimpl Verify for {} {{", record.dialect.op_type).unwrap(); + writeln!( + output, + " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n let op = self.get_operation().deref(ctx);\n if op.get_num_operands() != 3 || op.get_num_results() != 1 {{\n return verify_err!(op.loc(), {:?});\n }}\n if !is_i32(ctx, op.get_operand(0).get_type(ctx))\n || !{value_check}(ctx, op.get_operand(1).get_type(ctx))\n || !is_i32(ctx, op.get_operand(2).get_type(ctx))\n || !{value_check}(ctx, op.get_result(0).get_type(ctx))\n {{\n return verify_err!(op.loc(), {:?});\n }}\n Ok(())\n }}\n}}\n", + format!( + "{} requires exactly [member_mask, value, lane_or_delta] and one result", + record.dialect.op_name + ), + format!( + "{} requires i32 mask/lane and {value_label} value/result", + record.dialect.op_name + ), + ) + .unwrap(); + } + output.push_str("\npub(super) fn register(ctx: &mut Context) {\n"); + for record in warp_shuffles(catalog) { + writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); + } + output.push_str("}\n"); + output +} + +pub(in crate::render) fn render_dialect_dotprod(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + "//! Structural operations for generated packed integer dot products.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_i32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32)\n}\n\n", + ); + for record in dot_products(catalog) { + let signedness = match record.rust.result.as_str() { + "i32" => "Signed", + "u32" => "Unsigned", + result => panic!("unsupported dot-product result {result}"), + }; + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!( + output, + "/// Lowers to `{}`.", + dot_product_ptx(record).replace("$0, $1, $2, $3", "%d, %a, %b, %c") + ) + .unwrap(); + writeln!( + output, + "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<3>, NResultsInterface<1>],\n)]", + record.dialect.op_name + ) + .unwrap(); + writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); + writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); + writeln!( + output, + " pub fn new(op: Ptr) -> Self {{\n Self {{ op }}\n }}\n\n pub fn build(ctx: &mut Context, a: Value, b: Value, c: Value) -> Ptr {{\n let result_ty = IntegerType::get(ctx, 32, Signedness::{signedness});\n Operation::new(\n ctx,\n Self::get_concrete_op_info(),\n vec![result_ty.into()],\n vec![a, b, c],\n vec![],\n 0,\n )\n }}\n}}" + ) + .unwrap(); + writeln!(output, "\nimpl Verify for {} {{", record.dialect.op_type).unwrap(); + writeln!( + output, + " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n let op = self.get_operation().deref(ctx);\n if op.get_num_operands() != 3 || op.get_num_results() != 1 {{\n return verify_err!(op.loc(), {:?});\n }}\n if !(0..3).all(|index| is_i32(ctx, op.get_operand(index).get_type(ctx)))\n || !is_i32(ctx, op.get_result(0).get_type(ctx))\n {{\n return verify_err!(op.loc(), {:?});\n }}\n Ok(())\n }}\n}}\n", + format!( + "{} requires exactly three operands and one result", + record.dialect.op_name + ), + format!( + "{} operands and result must be 32-bit integers", + record.dialect.op_name + ), + ) + .unwrap(); + } + output.push_str("\npub(super) fn register(ctx: &mut Context) {\n"); + for record in dot_products(catalog) { + writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); + } + output.push_str("}\n"); + output +} + +pub(in crate::render) fn render_dialect_redux(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + let has_f32 = redux(catalog).any(|record| record.dialect.results == ["f32"]); + if has_f32 { + output.push_str( + "//! Structural NVVM operations for the closed generated `redux.sync` family.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{FP32Type, IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_i32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32)\n}\n\nfn is_f32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx).downcast_ref::().is_some()\n}\n\n", + ); + } else { + output.push_str( + "//! Structural NVVM operations for the closed generated `redux.sync` family.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_i32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32)\n}\n\n", + ); + } + for record in redux(catalog) { + writeln!(output, "/// {}", record.summary).unwrap(); + output.push_str( + "///\n/// Dialect operands are `[member_mask, value]`; generated lowering adapts\n/// them to LLVM's `(value, member_mask)` signature.\n", + ); + writeln!( + output, + "#[pliron_op(\n name = {:?},\n format,\n interfaces = [NOpdsInterface<2>, NResultsInterface<1>],\n)]", + record.dialect.op_name + ) + .unwrap(); + writeln!(output, "pub struct {};", record.dialect.op_type).unwrap(); + writeln!(output, "\nimpl {} {{", record.dialect.op_type).unwrap(); + let (result_type, value_predicate, type_error) = match record.rust.result.as_str() { + "u32" => ( + "IntegerType::get(ctx, 32, Signedness::Unsigned)", + "is_i32", + format!( + "{} member mask, value, and result must be 32-bit integers", + record.dialect.op_name + ), + ), + "i32" => ( + "IntegerType::get(ctx, 32, Signedness::Signed)", + "is_i32", + format!( + "{} member mask, value, and result must be 32-bit integers", + record.dialect.op_name + ), + ), + "f32" => ( + "FP32Type::get(ctx)", + "is_f32", + format!( + "{} member mask must be a 32-bit integer and value and result must be f32", + record.dialect.op_name + ), + ), + result => panic!("unsupported redux result {result}"), + }; + writeln!( + output, + " pub fn new(op: Ptr) -> Self {{\n Self {{ op }}\n }}\n\n pub fn build(ctx: &mut Context, member_mask: Value, value: Value) -> Ptr {{\n let result_ty = {result_type};\n Operation::new(\n ctx,\n Self::get_concrete_op_info(),\n vec![result_ty.into()],\n vec![member_mask, value],\n vec![],\n 0,\n )\n }}\n}}" + ) + .unwrap(); + writeln!(output, "\nimpl Verify for {} {{", record.dialect.op_type).unwrap(); + writeln!( + output, + " fn verify(&self, ctx: &Context) -> Result<(), Error> {{\n let op = self.get_operation().deref(ctx);\n if op.get_num_operands() != 2 || op.get_num_results() != 1 {{\n return verify_err!(op.loc(), {:?});\n }}\n if !is_i32(ctx, op.get_operand(0).get_type(ctx))\n || !{value_predicate}(ctx, op.get_operand(1).get_type(ctx))\n || !{value_predicate}(ctx, op.get_result(0).get_type(ctx))\n {{\n return verify_err!(op.loc(), {:?});\n }}\n Ok(())\n }}\n}}\n", + format!( + "{} requires exactly two operands [member_mask, value] and one result", + record.dialect.op_name + ), + type_error, + ) + .unwrap(); + } + output.push_str("\npub(super) fn register(ctx: &mut Context) {\n"); + for record in redux(catalog) { + writeln!(output, " {}::register(ctx);", record.dialect.op_type).unwrap(); + } + output.push_str("}\n"); + output +} diff --git a/crates/cuda-intrinsics-gen/src/render/families/misc.rs b/crates/cuda-intrinsics-gen/src/render/families/misc.rs new file mode 100644 index 0000000000..c6abdaf74e --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/render/families/misc.rs @@ -0,0 +1,198 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, CatalogFile, CatalogIntrinsic, ClcOperation, ClusterBarrierMode, + IntrinsicBackend, SpecialRegisterObservation, SpecialRegisterOutputConstraint, + SpecialRegisterPtxType, WgmmaControlMode, +}; +use crate::render::families::{execution_controls, wgmma_controls}; +use std::fmt::Write as _; + +pub(in crate::render) fn cluster_barrier_attr(record: &CatalogIntrinsic) -> &'static str { + match record + .cluster_barrier + .as_ref() + .expect("cluster-barrier contract") + .mode + { + ClusterBarrierMode::Arrive => "ClusterBarrierModeAttr::Arrive", + ClusterBarrierMode::ArriveAligned => "ClusterBarrierModeAttr::ArriveAligned", + ClusterBarrierMode::ArriveRelaxed => "ClusterBarrierModeAttr::ArriveRelaxed", + ClusterBarrierMode::ArriveRelaxedAligned => "ClusterBarrierModeAttr::ArriveRelaxedAligned", + ClusterBarrierMode::Wait => "ClusterBarrierModeAttr::Wait", + ClusterBarrierMode::WaitAligned => "ClusterBarrierModeAttr::WaitAligned", + } +} + +pub(in crate::render) fn cluster_barrier_template(record: &CatalogIntrinsic) -> String { + format!( + "{}.{};", + record.expected_ptx.mnemonic, + record.expected_ptx.modifiers.join(".") + ) +} + +#[derive(Clone, Copy)] +pub(in crate::render) enum ClcSafetyArgNames { + RawAbi, + Compatibility, +} + +fn clc_safety_lines(operation: ClcOperation, names: ClcSafetyArgNames) -> Vec { + let (response, mbar, resp_lo, resp_hi) = match names { + ClcSafetyArgNames::RawAbi => ("`_arg0`", "`_arg1`", "`_arg0`", "`_arg1`"), + ClcSafetyArgNames::Compatibility => ("`response`", "`mbar`", "`resp_lo`", "`resp_hi`"), + }; + match operation { + ClcOperation::TryCancel => vec![ + format!( + "{response} must designate a writable, 16-byte-aligned 16-byte region in `shared::cta` memory, and {mbar} must designate a valid, aligned shared-memory mbarrier." + ), + format!( + "The calling CTA must issue `arrive_expect_tx` for {mbar} with 16 bytes before this request." + ), + "This CTA must not have observed a prior CLC try-cancel request fail.".into(), + ], + ClcOperation::TryCancelMulticast => vec![ + format!( + "{response} must designate a writable, 16-byte-aligned 16-byte region in `shared::cta` memory, and {mbar} must designate a valid, aligned shared-memory mbarrier." + ), + format!( + "Each CTA must issue `arrive_expect_tx` for its own corresponding mbarrier with 16 bytes before this request." + ), + "This CTA must not have observed a prior CLC try-cancel request fail.".into(), + "The response is written to every CTA's shared-memory window and signals every CTA's mbarrier.".into(), + "Every CTA in the cluster must still be active.".into(), + ], + ClcOperation::QueryIsCanceled => vec![format!( + "{resp_lo} and {resp_hi} must be the two halves of an opaque response observed complete from a prior CLC try-cancel request." + )], + ClcOperation::QueryGetFirstCtaidX + | ClcOperation::QueryGetFirstCtaidY + | ClcOperation::QueryGetFirstCtaidZ => vec![ + format!( + "{resp_lo} and {resp_hi} must be the two halves of an opaque response observed complete from a prior CLC try-cancel request." + ), + "The result is meaningful only when `clc_query_is_canceled` returned one for this response.".into(), + ], + } +} + +pub(in crate::render) fn render_clc_safety_lines( + output: &mut String, + operation: ClcOperation, + names: ClcSafetyArgNames, +) { + for line in clc_safety_lines(operation, names) { + writeln!(output, "/// {line}").unwrap(); + } +} + +pub(in crate::render) fn execution_control_family<'a>( + catalog: &'a CatalogFile, + family: &'a str, +) -> impl Iterator { + execution_controls(catalog).filter(move |record| record.family == family) +} + +pub(in crate::render) fn wgmma_control( + catalog: &CatalogFile, + mode: WgmmaControlMode, +) -> &CatalogIntrinsic { + wgmma_controls(catalog) + .find(|record| { + record + .wgmma_control + .as_ref() + .is_some_and(|control| control.mode == mode) + }) + .expect("complete WGMMA-control family") +} + +pub(in crate::render) fn wgmma_control_template(record: &CatalogIntrinsic) -> String { + let operand = if record + .wgmma_control + .as_ref() + .is_some_and(|control| control.mode == WgmmaControlMode::WaitGroup) + { + " $0" + } else { + "" + }; + format!( + "{}.{}{operand};", + record.expected_ptx.mnemonic, + record.expected_ptx.modifiers.join(".") + ) +} + +pub(in crate::render) fn threadfence_ptx_level(record: &CatalogIntrinsic) -> Option<&'static str> { + match record.id.as_str() { + "threadfence_block" => Some("cta"), + "threadfence" => Some("gl"), + "threadfence_system" => Some("sys"), + _ => None, + } +} + +fn special_register_ptx_type(record: &CatalogIntrinsic) -> &'static str { + match record + .special_register + .as_ref() + .expect("special-register record") + .ptx_type + { + SpecialRegisterPtxType::B32 => "b32", + SpecialRegisterPtxType::U32 => "u32", + SpecialRegisterPtxType::U64 => "u64", + } +} + +pub(in crate::render) fn special_register_output_constraint( + record: &CatalogIntrinsic, +) -> &'static str { + match record + .special_register + .as_ref() + .expect("special-register record") + .output_constraint + { + SpecialRegisterOutputConstraint::Register32 => "=r", + SpecialRegisterOutputConstraint::Register64 => "=l", + } +} + +pub(in crate::render) fn special_register_inline_template(record: &CatalogIntrinsic) -> String { + let register = match record.expected_ptx.operands.as_slice() { + [_, crate::ptx::OperandPattern::Exact { value }] => value, + _ => unreachable!("closed special-register PTX operands"), + }; + format!("mov.{} $0, {register};", special_register_ptx_type(record)) +} + +pub(in crate::render) fn special_register_asm_kind(record: &CatalogIntrinsic) -> &'static str { + match record + .special_register + .as_ref() + .expect("special-register record") + .observation + { + SpecialRegisterObservation::StablePure => "AsmKind::Pure", + SpecialRegisterObservation::VolatileObservation => "AsmKind::SideEffect", + } +} + +pub(in crate::render) fn special_register_backend_mechanism( + record: &CatalogIntrinsic, + backend: IntrinsicBackend, +) -> BackendLoweringMechanism { + record + .backend_lowerings + .iter() + .find(|route| route.backend == backend) + .expect("special-register backend route") + .mechanism +} diff --git a/crates/cuda-intrinsics-gen/src/render/families/mma.rs b/crates/cuda-intrinsics-gen/src/render/families/mma.rs new file mode 100644 index 0000000000..42aebc76c0 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/render/families/mma.rs @@ -0,0 +1,566 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + CatalogIntrinsic, LdmatrixElement, LdmatrixLayout, LdmatrixMultiplicity, LdmatrixShape, + LdmatrixStateSpace, RegisterMmaAccumulator, RegisterMmaAdapter, RegisterMmaElement, + RegisterMmaKind, RegisterMmaLayout, RegisterMmaOperation, RegisterMmaOverflow, + RegisterMmaShape, SparseMma, SparseMmaAccumulator, SparseMmaAdapter, SparseMmaElement, + SparseMmaLayout, SparseMmaMetadata, SparseMmaOverflow, SparseMmaSelector, SparseMmaShape, + StmatrixLayout, StmatrixMultiplicity, +}; +use std::fmt::Write as _; + +pub(in crate::render) fn is_blackwell_ldmatrix(record: &CatalogIntrinsic) -> bool { + record + .ldmatrix + .as_ref() + .is_some_and(|ldmatrix| ldmatrix.variant.shape != LdmatrixShape::M8n8) +} + +pub(in crate::render) const BLACKWELL_LDMATRIX_EFFECTIVE_FLOORS: &str = "sm_100a PTX 8.6, sm_100f PTX 8.8, sm_103a PTX 8.8, sm_103f PTX 8.8, sm_110a PTX 9.0, sm_110f PTX 9.0, sm_120a PTX 8.7, sm_120f PTX 8.8, sm_121a PTX 8.8, sm_121f PTX 8.8"; + +pub(in crate::render) fn ldmatrix_compat_op( + record: &CatalogIntrinsic, +) -> Option<(&'static str, &'static str)> { + match record.id.as_str() { + "ldmatrix_m8n8_x1_b16" => Some(("LdmatrixX1Op", "nvvm.ldmatrix_x1")), + "ldmatrix_m8n8_x1_trans_b16" => Some(("LdmatrixX1TransOp", "nvvm.ldmatrix_x1_trans")), + "ldmatrix_m8n8_x2_b16" => Some(("LdmatrixX2Op", "nvvm.ldmatrix_x2")), + "ldmatrix_m8n8_x2_trans_b16" => Some(("LdmatrixX2TransOp", "nvvm.ldmatrix_x2_trans")), + "ldmatrix_m8n8_x4_b16" => Some(("LdmatrixX4Op", "nvvm.ldmatrix_x4")), + "ldmatrix_m8n8_x4_trans_b16" => Some(("LdmatrixX4TransOp", "nvvm.ldmatrix_x4_trans")), + _ => None, + } +} + +pub(in crate::render) fn stmatrix_variant( + record: &CatalogIntrinsic, +) -> Option<(StmatrixMultiplicity, StmatrixLayout)> { + match record.id.as_str() { + "stmatrix_m8n8_x2_b16" => Some((StmatrixMultiplicity::X2, StmatrixLayout::Normal)), + "stmatrix_m8n8_x2_trans_b16" => { + Some((StmatrixMultiplicity::X2, StmatrixLayout::Transposed)) + } + "stmatrix_m8n8_x4_b16" => Some((StmatrixMultiplicity::X4, StmatrixLayout::Normal)), + "stmatrix_m8n8_x4_trans_b16" => { + Some((StmatrixMultiplicity::X4, StmatrixLayout::Transposed)) + } + _ => None, + } +} + +pub(in crate::render) fn stmatrix_compatibility_name(record: &CatalogIntrinsic) -> &'static str { + match record.id.as_str() { + "stmatrix_m8n8_x2_b16" => "stmatrix_m8n8_x2", + "stmatrix_m8n8_x2_trans_b16" => "stmatrix_m8n8_x2_trans", + "stmatrix_m8n8_x4_b16" => "stmatrix_m8n8_x4", + "stmatrix_m8n8_x4_trans_b16" => "stmatrix_m8n8_x4_trans", + _ => panic!("unknown stmatrix record {}", record.id), + } +} + +pub(in crate::render) fn movmatrix_template(record: &CatalogIntrinsic) -> String { + format!( + "{}.{} $0, $1;", + record.expected_ptx.mnemonic, + record.expected_ptx.modifiers.join(".") + ) +} + +pub(in crate::render) fn ldmatrix_attr_variants( + record: &CatalogIntrinsic, +) -> ( + &'static str, + &'static str, + &'static str, + &'static str, + &'static str, +) { + let variant = &record.ldmatrix.as_ref().expect("ldmatrix record").variant; + let shape = match variant.shape { + LdmatrixShape::M8n8 => "LdmatrixShapeAttr::M8n8", + LdmatrixShape::M8n16 => "LdmatrixShapeAttr::M8n16", + LdmatrixShape::M16n16 => "LdmatrixShapeAttr::M16n16", + }; + let multiplicity = match variant.multiplicity { + LdmatrixMultiplicity::X1 => "LdmatrixMultiplicityAttr::X1", + LdmatrixMultiplicity::X2 => "LdmatrixMultiplicityAttr::X2", + LdmatrixMultiplicity::X4 => "LdmatrixMultiplicityAttr::X4", + }; + let layout = match variant.layout { + LdmatrixLayout::Normal => "LdmatrixLayoutAttr::Normal", + LdmatrixLayout::Transposed => "LdmatrixLayoutAttr::Transposed", + }; + let element = match variant.element { + LdmatrixElement::B16 => "LdmatrixElementAttr::B16", + LdmatrixElement::B8 => "LdmatrixElementAttr::B8", + LdmatrixElement::B8x16B4x16P64 => "LdmatrixElementAttr::B8x16B4x16P64", + LdmatrixElement::B8x16B6x16P32 => "LdmatrixElementAttr::B8x16B6x16P32", + }; + let state_space = match variant.state_space { + LdmatrixStateSpace::Shared => "LdmatrixStateSpaceAttr::Shared", + }; + (shape, multiplicity, layout, element, state_space) +} + +pub(in crate::render) fn register_mma_effective_kind(record: &CatalogIntrinsic) -> RegisterMmaKind { + let mma = record.register_mma.as_ref().expect("register-MMA record"); + mma.kind.unwrap_or_else(|| { + if record + .llvm + .as_ref() + .is_some_and(|llvm| llvm.symbol.contains(".kind.f8f6f4.")) + { + RegisterMmaKind::F8f6f4 + } else { + RegisterMmaKind::Standard + } + }) +} + +pub(in crate::render) fn register_mma_attr_variants( + record: &CatalogIntrinsic, +) -> ( + &'static str, + &'static str, + &'static str, + &'static str, + &'static str, + &'static str, + &'static str, + &'static str, + &'static str, +) { + let mma = record.register_mma.as_ref().expect("register-MMA record"); + let shape = match mma.shape { + RegisterMmaShape::M8n8k4 => "RegisterMmaShapeAttr::M8n8k4", + RegisterMmaShape::M8n8k16 => "RegisterMmaShapeAttr::M8n8k16", + RegisterMmaShape::M8n8k32 => "RegisterMmaShapeAttr::M8n8k32", + RegisterMmaShape::M8n8k128 => "RegisterMmaShapeAttr::M8n8k128", + RegisterMmaShape::M16n8k4 => "RegisterMmaShapeAttr::M16n8k4", + RegisterMmaShape::M16n8k8 => "RegisterMmaShapeAttr::M16n8k8", + RegisterMmaShape::M16n8k16 => "RegisterMmaShapeAttr::M16n8k16", + RegisterMmaShape::M16n8k32 => "RegisterMmaShapeAttr::M16n8k32", + RegisterMmaShape::M16n8k64 => "RegisterMmaShapeAttr::M16n8k64", + RegisterMmaShape::M16n8k128 => "RegisterMmaShapeAttr::M16n8k128", + RegisterMmaShape::M16n8k256 => "RegisterMmaShapeAttr::M16n8k256", + }; + let operation = match mma.operation { + RegisterMmaOperation::Multiply => "RegisterMmaOperationAttr::Multiply", + RegisterMmaOperation::AndPopc => "RegisterMmaOperationAttr::AndPopc", + RegisterMmaOperation::XorPopc => "RegisterMmaOperationAttr::XorPopc", + }; + let kind = match register_mma_effective_kind(record) { + RegisterMmaKind::Standard => "RegisterMmaKindAttr::Standard", + RegisterMmaKind::F8f6f4 => "RegisterMmaKindAttr::F8f6f4", + RegisterMmaKind::Mxf8f6f4 => "RegisterMmaKindAttr::Mxf8f6f4", + }; + let accumulator = match mma.accumulator { + RegisterMmaAccumulator::F16 => "RegisterMmaAccumulatorAttr::F16", + RegisterMmaAccumulator::F32 => "RegisterMmaAccumulatorAttr::F32", + RegisterMmaAccumulator::F64 => "RegisterMmaAccumulatorAttr::F64", + RegisterMmaAccumulator::S32 => "RegisterMmaAccumulatorAttr::S32", + }; + let element = |element| match element { + RegisterMmaElement::Bf16 => "RegisterMmaElementAttr::Bf16", + RegisterMmaElement::E2m1 => "RegisterMmaElementAttr::E2m1", + RegisterMmaElement::E2m3 => "RegisterMmaElementAttr::E2m3", + RegisterMmaElement::E3m2 => "RegisterMmaElementAttr::E3m2", + RegisterMmaElement::E4m3 => "RegisterMmaElementAttr::E4m3", + RegisterMmaElement::E5m2 => "RegisterMmaElementAttr::E5m2", + RegisterMmaElement::F16 => "RegisterMmaElementAttr::F16", + RegisterMmaElement::Tf32 => "RegisterMmaElementAttr::Tf32", + RegisterMmaElement::F64 => "RegisterMmaElementAttr::F64", + RegisterMmaElement::B1 => "RegisterMmaElementAttr::B1", + RegisterMmaElement::S4 => "RegisterMmaElementAttr::S4", + RegisterMmaElement::U4 => "RegisterMmaElementAttr::U4", + RegisterMmaElement::S8 => "RegisterMmaElementAttr::S8", + RegisterMmaElement::U8 => "RegisterMmaElementAttr::U8", + }; + let layout = |layout| match layout { + RegisterMmaLayout::Row => "RegisterMmaLayoutAttr::Row", + RegisterMmaLayout::Col => "RegisterMmaLayoutAttr::Col", + }; + let overflow = match mma.overflow { + RegisterMmaOverflow::NotApplicable => "RegisterMmaOverflowAttr::NotApplicable", + RegisterMmaOverflow::Wrapping => "RegisterMmaOverflowAttr::Wrapping", + RegisterMmaOverflow::Satfinite => "RegisterMmaOverflowAttr::Satfinite", + }; + ( + shape, + operation, + kind, + accumulator, + element(mma.a_element), + element(mma.b_element), + layout(mma.a_layout), + layout(mma.b_layout), + overflow, + ) +} + +pub(in crate::render) fn register_mma_fragment_counts( + record: &CatalogIntrinsic, +) -> (usize, usize, usize, usize) { + match record.register_mma.as_ref().unwrap().adapter { + RegisterMmaAdapter::C2U32A2U32B1U32ToD2U32 => (2, 2, 1, 2), + RegisterMmaAdapter::C2U32A4U32B2U32ToD2U32 => (2, 4, 2, 2), + RegisterMmaAdapter::C4F32A2U32B1U32ToD4F32 => (4, 2, 1, 4), + RegisterMmaAdapter::C4F32A4U32B2U32ToD4F32 | RegisterMmaAdapter::C4I32A4U32B2U32ToD4I32 => { + (4, 4, 2, 4) + } + RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 => (4, 4, 2, 4), + RegisterMmaAdapter::C2F64A1F64B1F64ToD2F64 | RegisterMmaAdapter::C2I32A1U32B1U32ToD2I32 => { + (2, 1, 1, 2) + } + RegisterMmaAdapter::C4I32A2U32B1U32ToD4I32 => (4, 2, 1, 4), + } +} + +pub(in crate::render) fn register_mma_extra_operand_count(record: &CatalogIntrinsic) -> usize { + usize::from( + record.register_mma.as_ref().unwrap().adapter + == RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32, + ) * 6 +} + +pub(in crate::render) fn expected_ptx_head(record: &CatalogIntrinsic) -> String { + let mut head = record.expected_ptx.mnemonic.clone(); + for modifier in &record.expected_ptx.modifiers { + write!(head, ".{modifier}").unwrap(); + } + head +} + +fn register_list(first: usize, count: usize) -> String { + format!( + "{{{}}}", + (first..first + count) + .map(|index| format!("${index}")) + .collect::>() + .join(", ") + ) +} + +pub(in crate::render) fn register_mma_template(record: &CatalogIntrinsic) -> String { + let (c_count, a_count, b_count, d_count) = register_mma_fragment_counts(record); + let d = register_list(0, d_count); + let c = register_list(d_count, c_count); + let a = register_list(d_count + c_count, a_count); + let b = register_list(d_count + c_count + a_count, b_count); + if record.register_mma.as_ref().unwrap().adapter + == RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 + { + let scale_a = d_count + c_count + a_count + b_count; + let byte_id_a = scale_a + 1; + let thread_id_a = scale_a + 2; + let scale_b = scale_a + 3; + let byte_id_b = scale_a + 4; + let thread_id_b = scale_a + 5; + format!( + "{} {d}, {a}, {b}, {c}, ${scale_a}, {{${byte_id_a}, ${thread_id_a}}}, ${scale_b}, {{${byte_id_b}, ${thread_id_b}}};", + expected_ptx_head(record) + ) + } else { + format!("{} {d}, {a}, {b}, {c};", expected_ptx_head(record)) + } +} + +pub(in crate::render) fn register_mma_constraints(record: &CatalogIntrinsic) -> String { + let mma = record.register_mma.as_ref().unwrap(); + let (c_count, a_count, b_count, d_count) = register_mma_fragment_counts(record); + let (output, c, packed) = match mma.adapter { + RegisterMmaAdapter::C2U32A2U32B1U32ToD2U32 | RegisterMmaAdapter::C2U32A4U32B2U32ToD2U32 => { + ("=r", "r", "r") + } + RegisterMmaAdapter::C4F32A2U32B1U32ToD4F32 + | RegisterMmaAdapter::C4F32A4U32B2U32ToD4F32 + | RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 => ("=f", "f", "r"), + RegisterMmaAdapter::C2F64A1F64B1F64ToD2F64 => ("=d", "d", "d"), + RegisterMmaAdapter::C2I32A1U32B1U32ToD2I32 + | RegisterMmaAdapter::C4I32A4U32B2U32ToD4I32 + | RegisterMmaAdapter::C4I32A2U32B1U32ToD4I32 => ("=r", "r", "r"), + }; + let mut constraints = std::iter::repeat_n(output, d_count) + .chain(std::iter::repeat_n(c, c_count)) + .chain(std::iter::repeat_n(packed, a_count + b_count)) + .collect::>(); + if mma.adapter == RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 { + constraints.extend(["r", "h", "h", "r", "h", "h"]); + } + constraints.join(",") +} + +pub(in crate::render) fn register_mma_result_variant(record: &CatalogIntrinsic) -> &'static str { + match record.register_mma.as_ref().unwrap().adapter { + RegisterMmaAdapter::C2U32A2U32B1U32ToD2U32 | RegisterMmaAdapter::C2U32A4U32B2U32ToD2U32 => { + "GeneratedMmaResultType::I32" + } + RegisterMmaAdapter::C4F32A2U32B1U32ToD4F32 | RegisterMmaAdapter::C4F32A4U32B2U32ToD4F32 => { + "GeneratedMmaResultType::F32" + } + RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 => { + "GeneratedMmaResultType::F32" + } + RegisterMmaAdapter::C2F64A1F64B1F64ToD2F64 => "GeneratedMmaResultType::F64", + RegisterMmaAdapter::C2I32A1U32B1U32ToD2I32 + | RegisterMmaAdapter::C4I32A4U32B2U32ToD4I32 + | RegisterMmaAdapter::C4I32A2U32B1U32ToD4I32 => "GeneratedMmaResultType::I32", + } +} + +pub(in crate::render) fn sparse_mma_fragment_counts( + record: &CatalogIntrinsic, +) -> (usize, usize, usize, usize) { + match record.sparse_mma.as_ref().unwrap().adapter { + SparseMmaAdapter::C2U32A4U32B4U32MetadataU32SelectorU32ToD2U32 => (2, 4, 4, 2), + SparseMmaAdapter::C4F32A4U32B4U32MetadataU32SelectorU32ToD4F32 + | SparseMmaAdapter::C4I32A4U32B4U32MetadataU32SelectorU32ToD4I32 => (4, 4, 4, 4), + SparseMmaAdapter::C4I32A2U32B2U32MetadataU32SelectorU32ToD4I32 => (4, 2, 2, 4), + } +} + +pub(in crate::render) fn sparse_mma_selector_values(record: &CatalogIntrinsic) -> &'static [u32] { + match record.sparse_mma.as_ref().unwrap().selector { + SparseMmaSelector::ImmediateZeroOrOne => &[0, 1], + SparseMmaSelector::ImmediateZero => &[0], + } +} + +pub(in crate::render) fn sparse_mma_selector_description( + record: &CatalogIntrinsic, +) -> &'static str { + match record.sparse_mma.as_ref().unwrap().selector { + SparseMmaSelector::ImmediateZeroOrOne => "the compile-time constant `0` or `1`", + SparseMmaSelector::ImmediateZero => "the compile-time constant `0`", + } +} + +pub(in crate::render) fn sparse_mma_selector_error(record: &CatalogIntrinsic) -> &'static str { + match record.sparse_mma.as_ref().unwrap().selector { + SparseMmaSelector::ImmediateZeroOrOne => { + "sparse MMA selector must be the compile-time constant 0 or 1" + } + SparseMmaSelector::ImmediateZero => { + "sparse MMA selector must be the compile-time constant 0" + } + } +} + +pub(in crate::render) const SPARSE_MMA_STANDARD_METADATA_RULE: &str = "Every 4-bit metadata group must encode two distinct 2-bit indices; `0x0`, `0x5`, `0xa`, and `0xf` are undefined behavior."; +pub(in crate::render) const SPARSE_MMA_ORDERED_METADATA_RULE: &str = "Every 4-bit metadata group must be `0x4`, `0x8`, `0x9`, `0xc`, `0xd`, or `0xe`; any other value is undefined behavior."; + +pub(in crate::render) fn sparse_mma_metadata_rule(mma: &SparseMma) -> &'static str { + match mma.metadata { + SparseMmaMetadata::Standard => SPARSE_MMA_STANDARD_METADATA_RULE, + SparseMmaMetadata::Ordered => SPARSE_MMA_ORDERED_METADATA_RULE, + } +} + +pub(in crate::render) fn sparse_mma_import_adapter(record: &CatalogIntrinsic) -> &'static str { + match record.sparse_mma.as_ref().unwrap().adapter { + SparseMmaAdapter::C2U32A4U32B4U32MetadataU32SelectorU32ToD2U32 => { + "GeneratedMmaImportAdapter::C2U32A4U32B4U32ToD2U32" + } + SparseMmaAdapter::C4F32A4U32B4U32MetadataU32SelectorU32ToD4F32 => { + "GeneratedMmaImportAdapter::C4F32A4U32B4U32ToD4F32" + } + SparseMmaAdapter::C4I32A2U32B2U32MetadataU32SelectorU32ToD4I32 => { + "GeneratedMmaImportAdapter::C4I32A2U32B2U32ToD4I32" + } + SparseMmaAdapter::C4I32A4U32B4U32MetadataU32SelectorU32ToD4I32 => { + "GeneratedMmaImportAdapter::C4I32A4U32B4U32ToD4I32" + } + } +} + +pub(in crate::render) fn sparse_mma_attr_variants( + record: &CatalogIntrinsic, +) -> ( + &'static str, + &'static str, + &'static str, + &'static str, + &'static str, + &'static str, + &'static str, + &'static str, + &'static str, +) { + let mma = record.sparse_mma.as_ref().expect("sparse-MMA record"); + let shape = match mma.shape { + SparseMmaShape::M16n8k32 => "SparseMmaShapeAttr::M16n8k32", + SparseMmaShape::M16n8k64 => "SparseMmaShapeAttr::M16n8k64", + SparseMmaShape::M16n8k128 => "SparseMmaShapeAttr::M16n8k128", + }; + let accumulator = match mma.accumulator { + SparseMmaAccumulator::F16 => "SparseMmaAccumulatorAttr::F16", + SparseMmaAccumulator::F32 => "SparseMmaAccumulatorAttr::F32", + SparseMmaAccumulator::S32 => "SparseMmaAccumulatorAttr::S32", + }; + let element = |element| match element { + SparseMmaElement::E2m1 => "SparseMmaElementAttr::E2m1", + SparseMmaElement::E2m3 => "SparseMmaElementAttr::E2m3", + SparseMmaElement::E3m2 => "SparseMmaElementAttr::E3m2", + SparseMmaElement::E4m3 => "SparseMmaElementAttr::E4m3", + SparseMmaElement::E5m2 => "SparseMmaElementAttr::E5m2", + SparseMmaElement::S4 => "SparseMmaElementAttr::S4", + SparseMmaElement::U4 => "SparseMmaElementAttr::U4", + SparseMmaElement::S8 => "SparseMmaElementAttr::S8", + SparseMmaElement::U8 => "SparseMmaElementAttr::U8", + }; + let layout = |layout| match layout { + SparseMmaLayout::Row => "SparseMmaLayoutAttr::Row", + SparseMmaLayout::Col => "SparseMmaLayoutAttr::Col", + }; + let overflow = match mma.overflow { + SparseMmaOverflow::NotApplicable => "SparseMmaOverflowAttr::NotApplicable", + SparseMmaOverflow::Wrapping => "SparseMmaOverflowAttr::Wrapping", + SparseMmaOverflow::Satfinite => "SparseMmaOverflowAttr::Satfinite", + }; + let metadata = match mma.metadata { + SparseMmaMetadata::Standard => "SparseMmaMetadataAttr::Standard", + SparseMmaMetadata::Ordered => "SparseMmaMetadataAttr::Ordered", + }; + let selector = match mma.selector { + SparseMmaSelector::ImmediateZeroOrOne => "SparseMmaSelectorAttr::ImmediateZeroOrOne", + SparseMmaSelector::ImmediateZero => "SparseMmaSelectorAttr::ImmediateZero", + }; + ( + shape, + accumulator, + element(mma.a_element), + element(mma.b_element), + layout(mma.a_layout), + layout(mma.b_layout), + overflow, + metadata, + selector, + ) +} + +pub(in crate::render) fn sparse_mma_ptx_head(record: &CatalogIntrinsic) -> String { + let mut head = record.expected_ptx.mnemonic.clone(); + for modifier in &record.expected_ptx.modifiers { + write!(head, ".{modifier}").unwrap(); + } + head +} + +pub(in crate::render) fn sparse_mma_template(record: &CatalogIntrinsic) -> String { + let (c_count, a_count, b_count, d_count) = sparse_mma_fragment_counts(record); + let d = register_list(0, d_count); + let c = register_list(d_count, c_count); + let a = register_list(d_count + c_count, a_count); + let b = register_list(d_count + c_count + a_count, b_count); + let metadata = d_count + c_count + a_count + b_count; + let selector = metadata + 1; + format!( + "{} {d}, {a}, {b}, {c}, ${metadata}, ${selector};", + sparse_mma_ptx_head(record) + ) +} + +pub(in crate::render) fn sparse_mma_constraints(record: &CatalogIntrinsic) -> String { + let mma = record.sparse_mma.as_ref().unwrap(); + let (c_count, a_count, b_count, d_count) = sparse_mma_fragment_counts(record); + let (output, accumulator) = match mma.accumulator { + SparseMmaAccumulator::F16 => ("=r", "r"), + SparseMmaAccumulator::F32 => ("=f", "f"), + SparseMmaAccumulator::S32 => ("=r", "r"), + }; + std::iter::repeat_n(output, d_count) + .chain(std::iter::repeat_n(accumulator, c_count)) + .chain(std::iter::repeat_n("r", a_count + b_count + 1)) + .chain(std::iter::once("n")) + .collect::>() + .join(",") +} + +pub(in crate::render) fn sparse_mma_result_variant(record: &CatalogIntrinsic) -> &'static str { + match record.sparse_mma.as_ref().unwrap().accumulator { + SparseMmaAccumulator::F16 => "GeneratedMmaResultType::I32", + SparseMmaAccumulator::F32 => "GeneratedMmaResultType::F32", + SparseMmaAccumulator::S32 => "GeneratedMmaResultType::I32", + } +} + +pub(in crate::render) fn sparse_mma_carriers(record: &CatalogIntrinsic) -> (String, String) { + let mma = record.sparse_mma.as_ref().unwrap(); + let (c_count, a_count, b_count, d_count) = sparse_mma_fragment_counts(record); + let accumulator = match mma.accumulator { + SparseMmaAccumulator::F16 => "MmaCarrier::U32", + SparseMmaAccumulator::F32 => "MmaCarrier::F32", + SparseMmaAccumulator::S32 => "MmaCarrier::I32", + }; + let carrier_slice = |carriers: Vec<&str>| format!("&[{}]", carriers.join(", ")); + let operands = carrier_slice( + std::iter::repeat_n(accumulator, c_count) + .chain(std::iter::repeat_n( + "MmaCarrier::U32", + a_count + b_count + 2, + )) + .collect(), + ); + let results = carrier_slice(std::iter::repeat_n(accumulator, d_count).collect()); + (operands, results) +} + +pub(in crate::render) fn register_mma_carriers( + record: &CatalogIntrinsic, +) -> (&'static str, &'static str) { + match record.register_mma.as_ref().unwrap().adapter { + RegisterMmaAdapter::C2U32A2U32B1U32ToD2U32 => ( + "&[MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32]", + "&[MmaCarrier::U32, MmaCarrier::U32]", + ), + RegisterMmaAdapter::C2U32A4U32B2U32ToD2U32 => ( + "&[MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32]", + "&[MmaCarrier::U32, MmaCarrier::U32]", + ), + RegisterMmaAdapter::C4F32A2U32B1U32ToD4F32 => ( + "&[MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32]", + "&[MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32]", + ), + RegisterMmaAdapter::C4F32A4U32B2U32ToD4F32 => ( + "&[MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32]", + "&[MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32]", + ), + RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 => ( + "&[MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U16, MmaCarrier::U16, MmaCarrier::U32, MmaCarrier::U16, MmaCarrier::U16]", + "&[MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32]", + ), + RegisterMmaAdapter::C2F64A1F64B1F64ToD2F64 => ( + "&[MmaCarrier::F64, MmaCarrier::F64, MmaCarrier::F64, MmaCarrier::F64]", + "&[MmaCarrier::F64, MmaCarrier::F64]", + ), + RegisterMmaAdapter::C2I32A1U32B1U32ToD2I32 => ( + "&[MmaCarrier::I32, MmaCarrier::I32, MmaCarrier::U32, MmaCarrier::U32]", + "&[MmaCarrier::I32, MmaCarrier::I32]", + ), + RegisterMmaAdapter::C4I32A4U32B2U32ToD4I32 => ( + "&[MmaCarrier::I32, MmaCarrier::I32, MmaCarrier::I32, MmaCarrier::I32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32]", + "&[MmaCarrier::I32, MmaCarrier::I32, MmaCarrier::I32, MmaCarrier::I32]", + ), + RegisterMmaAdapter::C4I32A2U32B1U32ToD4I32 => ( + "&[MmaCarrier::I32, MmaCarrier::I32, MmaCarrier::I32, MmaCarrier::I32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32]", + "&[MmaCarrier::I32, MmaCarrier::I32, MmaCarrier::I32, MmaCarrier::I32]", + ), + } +} + +pub(in crate::render) fn register_mma_compat_op_type( + record: &CatalogIntrinsic, +) -> Option<&'static str> { + match record.id.as_str() { + "mma_m16n8k16_f32_bf16" => Some("MmaM16N8K16F32Bf16Op"), + "mma_m16n8k16_f32_f16" => Some("MmaM16N8K16F32F16Op"), + "mma_m16n8k8_f32_tf32" => Some("MmaM16N8K8F32Tf32Op"), + "mma_m16n8k32_s32_s8" => Some("MmaM16N8K32S32S8Op"), + "mma_m8n8k4_f64" => Some("MmaM8N8K4F64Op"), + _ => None, + } +} diff --git a/crates/cuda-intrinsics-gen/src/render/families/mod.rs b/crates/cuda-intrinsics-gen/src/render/families/mod.rs new file mode 100644 index 0000000000..33be895980 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/render/families/mod.rs @@ -0,0 +1,301 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{CatalogFile, CatalogIntrinsic}; + +mod misc; +mod mma; +mod packed; +mod scalar; +mod tcgen05; + +pub(super) use misc::*; +pub(super) use mma::*; +pub(super) use packed::*; +pub(super) use scalar::*; +pub(super) use tcgen05::*; + +pub(super) fn sregs(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "sreg") +} + +pub(super) fn ldmatrix(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "ldmatrix") +} + +pub(super) fn stmatrices(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "stmatrix") +} + +pub(super) fn register_mmas(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "register_mma") +} + +pub(super) fn sparse_mmas(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "sparse_mma") +} + +pub(super) fn packed_atomics(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "packed_atomic") +} + +pub(super) fn redux(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "redux") +} + +pub(super) fn dot_products(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "dotprod") +} + +pub(super) fn packed_alus(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "packed_alu") +} + +pub(super) fn packed_conversions(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "packed_conversion") +} + +pub(super) fn scalar_conversions(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "scalar_conversion") +} + +pub(super) fn scalar_arithmetics(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "scalar_arithmetic") +} + +pub(super) fn extended_minmax(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "extended_minmax") +} + +pub(super) fn scalar_maths(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "scalar_math") +} + +pub(super) fn prmts(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "prmt") +} + +pub(super) fn cluster_barriers(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "cluster_barrier") +} + +pub(super) fn debug_controls(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "debug_control") +} + +pub(super) fn cluster_memory(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "cluster_memory") +} + +pub(super) fn clc_intrinsics(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "clc") +} + +pub(super) fn tma_intrinsics(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "tma") +} + +pub(super) fn execution_controls(catalog: &CatalogFile) -> impl Iterator { + catalog.intrinsics.iter().filter(|record| { + matches!( + record.family.as_str(), + "counted_barrier" | "grid_dependency" | "register_control" + ) + }) +} + +pub(super) fn wgmma_controls(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "wgmma_control") +} + +pub(super) fn tcgen05_intrinsics(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "tcgen05") +} + +pub(super) fn tcgen05_mma_intrinsics( + catalog: &CatalogFile, +) -> impl Iterator { + tcgen05_intrinsics(catalog).filter(|record| { + record + .tcgen05 + .as_ref() + .is_some_and(|tcgen05| tcgen05.mma.is_some()) + }) +} + +pub(super) fn tcgen05_non_mma_intrinsics( + catalog: &CatalogFile, +) -> impl Iterator { + tcgen05_intrinsics(catalog).filter(|record| { + record + .tcgen05 + .as_ref() + .is_some_and(|tcgen05| tcgen05.mma.is_none()) + }) +} + +pub(super) fn cp_async_copies(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "cp_async_copy") +} + +pub(super) fn cp_async_controls(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "cp_async_control") +} + +pub(super) fn cp_async_mbarriers(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "cp_async_mbarrier") +} + +pub(super) fn mbarrier_basics(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "mbarrier_basic") +} + +pub(super) fn movmatrix(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "movmatrix") +} + +pub(super) fn mbarrier_extended(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "mbarrier_extended") +} + +pub(super) fn sync_intrinsics(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "sync") +} + +pub(super) fn vote_intrinsics(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "vote") +} + +pub(super) fn active_masks(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "active_mask") +} + +pub(super) fn elect_intrinsics(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "elect") +} + +pub(super) fn warp_matches(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "warp_match") +} + +pub(super) fn warp_barriers(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "warp_barrier") +} + +pub(super) fn warp_shuffles(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "warp_shuffle") +} + +pub(super) fn integer_minmaxes(catalog: &CatalogFile) -> impl Iterator { + catalog + .intrinsics + .iter() + .filter(|record| record.family == "integer_minmax") +} diff --git a/crates/cuda-intrinsics-gen/src/render/families/packed.rs b/crates/cuda-intrinsics-gen/src/render/families/packed.rs new file mode 100644 index 0000000000..6815833d5a --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/render/families/packed.rs @@ -0,0 +1,449 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, CatalogIntrinsic, DotProductAdapter, DotProductOperation, + DotProductSignedness, IntrinsicBackend, PackedAluAdapter, PackedAluFormat, PackedAluOperation, + PackedConversionAdapter, PackedConversionDestinationFormat, PackedConversionRounding, + PackedConversionSaturation, PackedConversionSourceFormat, +}; + +pub(in crate::render) fn packed_alu_format_shape( + format: PackedAluFormat, +) -> ( + &'static str, + bool, + &'static str, + &'static str, + PackedAluAdapter, +) { + match format { + PackedAluFormat::Bf16x2 => ( + "bf16x2", + false, + "u32", + "i32", + PackedAluAdapter::DirectPackedU32, + ), + PackedAluFormat::F16x2 => ( + "f16x2", + true, + "u32", + "i32", + PackedAluAdapter::DirectPackedU32, + ), + PackedAluFormat::F32x2 => ( + "f32x2", + true, + "u64", + "i64", + PackedAluAdapter::DirectPackedU64, + ), + } +} + +pub(in crate::render) fn packed_alu_width(record: &CatalogIntrinsic) -> u32 { + match record + .packed_alu + .as_ref() + .expect("packed-ALU record") + .adapter + { + PackedAluAdapter::DirectPackedU32 => 32, + PackedAluAdapter::DirectPackedU64 => 64, + } +} + +pub(in crate::render) fn packed_alu_register_constraint(record: &CatalogIntrinsic) -> &'static str { + match packed_alu_width(record) { + 32 => "r", + 64 => "l", + _ => unreachable!("closed packed-ALU carrier width"), + } +} + +pub(in crate::render) fn dot_product_ptx(record: &CatalogIntrinsic) -> &'static str { + let dot = record.dot_product.as_ref().expect("dot-product record"); + match (dot.operation, dot.signedness, dot.adapter) { + ( + DotProductOperation::Dp4a, + DotProductSignedness::Signed, + DotProductAdapter::DirectThreeOperands, + ) => "dp4a.s32.s32 $0, $1, $2, $3;", + ( + DotProductOperation::Dp4a, + DotProductSignedness::Unsigned, + DotProductAdapter::DirectThreeOperands, + ) => "dp4a.u32.u32 $0, $1, $2, $3;", + ( + DotProductOperation::Dp2a, + DotProductSignedness::Signed, + DotProductAdapter::InsertLowHalfFalse, + ) => "dp2a.lo.s32.s32 $0, $1, $2, $3;", + ( + DotProductOperation::Dp2a, + DotProductSignedness::Unsigned, + DotProductAdapter::InsertLowHalfFalse, + ) => "dp2a.lo.u32.u32 $0, $1, $2, $3;", + combination => panic!("unsupported generated dot-product recipe {combination:?}"), + } +} + +pub(in crate::render) fn packed_alu_ptx_mnemonic(record: &CatalogIntrinsic) -> &'static str { + let packed = record.packed_alu.as_ref().expect("packed-ALU record"); + match (packed.format, packed.operation, packed.adapter) { + (PackedAluFormat::Bf16x2, PackedAluOperation::Add, PackedAluAdapter::DirectPackedU32) => { + "add.rn.bf16x2" + } + (PackedAluFormat::Bf16x2, PackedAluOperation::Sub, PackedAluAdapter::DirectPackedU32) => { + "sub.rn.bf16x2" + } + (PackedAluFormat::Bf16x2, PackedAluOperation::Mul, PackedAluAdapter::DirectPackedU32) => { + "mul.rn.bf16x2" + } + (PackedAluFormat::Bf16x2, PackedAluOperation::Fma, PackedAluAdapter::DirectPackedU32) => { + "fma.rn.bf16x2" + } + ( + PackedAluFormat::Bf16x2, + PackedAluOperation::FmaRelu, + PackedAluAdapter::DirectPackedU32, + ) => "fma.rn.relu.bf16x2", + (PackedAluFormat::Bf16x2, PackedAluOperation::Min, PackedAluAdapter::DirectPackedU32) => { + "min.bf16x2" + } + (PackedAluFormat::Bf16x2, PackedAluOperation::Max, PackedAluAdapter::DirectPackedU32) => { + "max.bf16x2" + } + (PackedAluFormat::Bf16x2, PackedAluOperation::Neg, PackedAluAdapter::DirectPackedU32) => { + "neg.bf16x2" + } + (PackedAluFormat::Bf16x2, PackedAluOperation::Abs, PackedAluAdapter::DirectPackedU32) => { + "abs.bf16x2" + } + (PackedAluFormat::F16x2, PackedAluOperation::Add, PackedAluAdapter::DirectPackedU32) => { + "add.rn.f16x2" + } + (PackedAluFormat::F16x2, PackedAluOperation::Sub, PackedAluAdapter::DirectPackedU32) => { + "sub.rn.f16x2" + } + (PackedAluFormat::F16x2, PackedAluOperation::Mul, PackedAluAdapter::DirectPackedU32) => { + "mul.rn.f16x2" + } + (PackedAluFormat::F16x2, PackedAluOperation::Fma, PackedAluAdapter::DirectPackedU32) => { + "fma.rn.f16x2" + } + ( + PackedAluFormat::F16x2, + PackedAluOperation::FmaRelu, + PackedAluAdapter::DirectPackedU32, + ) => "fma.rn.relu.f16x2", + (PackedAluFormat::F16x2, PackedAluOperation::FmaFtz, PackedAluAdapter::DirectPackedU32) => { + "fma.rn.ftz.f16x2" + } + (PackedAluFormat::F16x2, PackedAluOperation::FmaSat, PackedAluAdapter::DirectPackedU32) => { + "fma.rn.sat.f16x2" + } + ( + PackedAluFormat::F16x2, + PackedAluOperation::FmaFtzSat, + PackedAluAdapter::DirectPackedU32, + ) => "fma.rn.ftz.sat.f16x2", + ( + PackedAluFormat::F16x2, + PackedAluOperation::FmaFtzRelu, + PackedAluAdapter::DirectPackedU32, + ) => "fma.rn.ftz.relu.f16x2", + (PackedAluFormat::F16x2, PackedAluOperation::Min, PackedAluAdapter::DirectPackedU32) => { + "min.f16x2" + } + (PackedAluFormat::F16x2, PackedAluOperation::Max, PackedAluAdapter::DirectPackedU32) => { + "max.f16x2" + } + (PackedAluFormat::F16x2, PackedAluOperation::Neg, PackedAluAdapter::DirectPackedU32) => { + "neg.f16x2" + } + (PackedAluFormat::F16x2, PackedAluOperation::Abs, PackedAluAdapter::DirectPackedU32) => { + "abs.f16x2" + } + (PackedAluFormat::F32x2, PackedAluOperation::Add, PackedAluAdapter::DirectPackedU64) => { + "add.rn.f32x2" + } + (PackedAluFormat::F32x2, PackedAluOperation::AddFtz, PackedAluAdapter::DirectPackedU64) => { + "add.rn.ftz.f32x2" + } + (PackedAluFormat::F32x2, PackedAluOperation::Sub, PackedAluAdapter::DirectPackedU64) => { + "sub.rn.f32x2" + } + (PackedAluFormat::F32x2, PackedAluOperation::SubFtz, PackedAluAdapter::DirectPackedU64) => { + "sub.rn.ftz.f32x2" + } + (PackedAluFormat::F32x2, PackedAluOperation::Mul, PackedAluAdapter::DirectPackedU64) => { + "mul.rn.f32x2" + } + (PackedAluFormat::F32x2, PackedAluOperation::MulFtz, PackedAluAdapter::DirectPackedU64) => { + "mul.rn.ftz.f32x2" + } + (PackedAluFormat::F32x2, PackedAluOperation::Fma, PackedAluAdapter::DirectPackedU64) => { + "fma.rn.f32x2" + } + (PackedAluFormat::F32x2, PackedAluOperation::FmaFtz, PackedAluAdapter::DirectPackedU64) => { + "fma.rn.ftz.f32x2" + } + // bf16x2 has no NVPTX selection pattern for the ftz and sat fma forms; + // `packed_alu_recipe` rejects those pairs before a record can exist. + combination => panic!("unsupported generated packed-ALU recipe {combination:?}"), + } +} + +pub(in crate::render) fn packed_conversion_destination(record: &CatalogIntrinsic) -> &'static str { + match record + .packed_conversion + .as_ref() + .expect("packed-conversion record") + .destination_format + { + PackedConversionDestinationFormat::Bf16x2 => "bf16x2", + PackedConversionDestinationFormat::E4m3x2 => "e4m3x2", + PackedConversionDestinationFormat::E5m2x2 => "e5m2x2", + PackedConversionDestinationFormat::F16x2 => "f16x2", + } +} + +pub(in crate::render) fn packed_conversion_element(record: &CatalogIntrinsic) -> &'static str { + match record + .packed_conversion + .as_ref() + .expect("packed-conversion record") + .destination_format + { + PackedConversionDestinationFormat::Bf16x2 => "bf16", + PackedConversionDestinationFormat::E4m3x2 => "e4m3", + PackedConversionDestinationFormat::E5m2x2 => "e5m2", + PackedConversionDestinationFormat::F16x2 => "f16", + } +} + +/// Whether the conversion is one of the closed, reviewed packed-conversion +/// recipes: scalar `f32` pairs narrowed to a packed format, `f16x2` narrowed to +/// packed FP8, or packed FP8 widened back to `f16x2`. +pub(in crate::render) fn packed_conversion_is_closed_recipe( + conversion: &crate::model::PackedConversion, +) -> bool { + use PackedConversionDestinationFormat as Dst; + use PackedConversionRounding as Round; + use PackedConversionSaturation as Sat; + use PackedConversionSourceFormat as Src; + + let adapter_matches = conversion.adapter + == match conversion.source_format { + Src::F32x2 => PackedConversionAdapter::ReverseHighLowOperands, + Src::E4m3x2 | Src::E5m2x2 | Src::F16x2 => PackedConversionAdapter::Identity, + }; + + adapter_matches + && matches!( + ( + conversion.source_format, + conversion.destination_format, + conversion.rounding, + conversion.saturation, + ), + (Src::F32x2, Dst::Bf16x2, Round::NearestEven, Sat::None) + | (Src::F32x2, Dst::Bf16x2, Round::NearestEven, Sat::Relu) + | (Src::F32x2, Dst::Bf16x2, Round::TowardZero, Sat::None) + | (Src::F32x2, Dst::F16x2, Round::NearestEven, Sat::None) + | (Src::F32x2, Dst::F16x2, Round::NearestEven, Sat::Relu) + | (Src::F32x2, Dst::F16x2, Round::TowardZero, Sat::None) + | (Src::F32x2, Dst::E4m3x2, Round::NearestEven, Sat::Satfinite) + | ( + Src::F32x2, + Dst::E4m3x2, + Round::NearestEven, + Sat::SatfiniteRelu + ) + | (Src::F32x2, Dst::E5m2x2, Round::NearestEven, Sat::Satfinite) + | ( + Src::F32x2, + Dst::E5m2x2, + Round::NearestEven, + Sat::SatfiniteRelu + ) + | (Src::F16x2, Dst::E4m3x2, Round::NearestEven, Sat::Satfinite) + | ( + Src::F16x2, + Dst::E4m3x2, + Round::NearestEven, + Sat::SatfiniteRelu + ) + | (Src::F16x2, Dst::E5m2x2, Round::NearestEven, Sat::Satfinite) + | ( + Src::F16x2, + Dst::E5m2x2, + Round::NearestEven, + Sat::SatfiniteRelu + ) + | (Src::E4m3x2, Dst::F16x2, Round::NearestEven, Sat::None) + | (Src::E4m3x2, Dst::F16x2, Round::NearestEven, Sat::Relu) + | (Src::E5m2x2, Dst::F16x2, Round::NearestEven, Sat::None) + | (Src::E5m2x2, Dst::F16x2, Round::NearestEven, Sat::Relu) + ) +} + +pub(in crate::render) fn packed_conversion_source( + record: &CatalogIntrinsic, +) -> PackedConversionSourceFormat { + record + .packed_conversion + .as_ref() + .expect("packed-conversion record") + .source_format +} + +/// Register width of a single packed source operand. +pub(in crate::render) fn packed_conversion_source_width(record: &CatalogIntrinsic) -> u32 { + match packed_conversion_source(record) { + PackedConversionSourceFormat::F16x2 => 32, + PackedConversionSourceFormat::E4m3x2 | PackedConversionSourceFormat::E5m2x2 => 16, + PackedConversionSourceFormat::F32x2 => { + unreachable!("f32x2 conversions do not have a single packed source") + } + } +} + +/// Rust argument types, one per source operand. +pub(in crate::render) fn packed_conversion_rust_arguments( + record: &CatalogIntrinsic, +) -> Vec<&'static str> { + match packed_conversion_source(record) { + PackedConversionSourceFormat::F32x2 => vec!["f32", "f32"], + PackedConversionSourceFormat::F16x2 => vec!["u32"], + PackedConversionSourceFormat::E4m3x2 | PackedConversionSourceFormat::E5m2x2 => vec!["u16"], + } +} + +/// Dialect operand types, one per source operand. +pub(in crate::render) fn packed_conversion_dialect_operands( + record: &CatalogIntrinsic, +) -> Vec<&'static str> { + match packed_conversion_source(record) { + PackedConversionSourceFormat::F32x2 => vec!["f32", "f32"], + PackedConversionSourceFormat::F16x2 => vec!["i32"], + PackedConversionSourceFormat::E4m3x2 | PackedConversionSourceFormat::E5m2x2 => vec!["i16"], + } +} + +pub(in crate::render) fn packed_conversion_result_width(record: &CatalogIntrinsic) -> u32 { + match record + .packed_conversion + .as_ref() + .expect("packed-conversion record") + .destination_format + { + PackedConversionDestinationFormat::Bf16x2 | PackedConversionDestinationFormat::F16x2 => 32, + PackedConversionDestinationFormat::E4m3x2 | PackedConversionDestinationFormat::E5m2x2 => 16, + } +} + +pub(in crate::render) fn packed_conversion_rust_type(record: &CatalogIntrinsic) -> &'static str { + match packed_conversion_result_width(record) { + 16 => "u16", + 32 => "u32", + _ => unreachable!("closed packed-conversion result width"), + } +} + +pub(in crate::render) fn packed_conversion_dialect_type(record: &CatalogIntrinsic) -> &'static str { + match packed_conversion_result_width(record) { + 16 => "i16", + 32 => "i32", + _ => unreachable!("closed packed-conversion result width"), + } +} + +/// Inline-asm constraint string: one result register, then one per source +/// operand. `h` is a 16-bit register, `r` a 32-bit one, and `f` an f32. +pub(in crate::render) fn packed_conversion_constraint(record: &CatalogIntrinsic) -> &'static str { + match ( + packed_conversion_result_width(record), + packed_conversion_source(record), + ) { + (16, PackedConversionSourceFormat::F32x2) => "=h,f,f", + (32, PackedConversionSourceFormat::F32x2) => "=r,f,f", + (16, PackedConversionSourceFormat::F16x2) => "=h,r", + (32, PackedConversionSourceFormat::E4m3x2 | PackedConversionSourceFormat::E5m2x2) => "=r,h", + _ => unreachable!("closed packed-conversion result width and source format"), + } +} + +/// Whether the conversion lowers through a typed NVVM intrinsic call rather +/// than inline PTX. Only the scalar-`f32` pair narrowed to FP8 does; see the +/// matching predicate in `resolve.rs` for why. +fn packed_conversion_uses_typed_nvvm(record: &CatalogIntrinsic) -> bool { + let conversion = record + .packed_conversion + .as_ref() + .expect("packed-conversion record"); + conversion.source_format == PackedConversionSourceFormat::F32x2 + && matches!( + conversion.destination_format, + PackedConversionDestinationFormat::E4m3x2 | PackedConversionDestinationFormat::E5m2x2 + ) +} + +pub(in crate::render) fn packed_conversion_lowering_name( + record: &CatalogIntrinsic, +) -> &'static str { + if packed_conversion_uses_typed_nvvm(record) { + "generated_packed_conversion_backend" + } else { + "generated_packed_conversion_inline_ptx" + } +} + +pub(in crate::render) fn packed_conversion_typed_llvm_name( + record: &CatalogIntrinsic, +) -> Option { + let route = record + .backend_lowerings + .iter() + .find(|lowering| lowering.backend == IntrinsicBackend::LlvmNvptx) + .expect("packed conversion has an LLVM-NVPTX route"); + match route.mechanism { + BackendLoweringMechanism::TypedNvvm => Some(record.llvm_identifier()), + BackendLoweringMechanism::InlinePtx => None, + } +} + +pub(in crate::render) fn packed_conversion_ptx_mnemonic(record: &CatalogIntrinsic) -> String { + let conversion = record + .packed_conversion + .as_ref() + .expect("packed-conversion record"); + debug_assert_eq!( + conversion.adapter, + match conversion.source_format { + PackedConversionSourceFormat::F32x2 => PackedConversionAdapter::ReverseHighLowOperands, + PackedConversionSourceFormat::E4m3x2 + | PackedConversionSourceFormat::E5m2x2 + | PackedConversionSourceFormat::F16x2 => PackedConversionAdapter::Identity, + } + ); + let rounding = match conversion.rounding { + PackedConversionRounding::NearestEven => "rn", + PackedConversionRounding::TowardZero => "rz", + }; + let saturation = match conversion.saturation { + PackedConversionSaturation::None => "", + PackedConversionSaturation::Relu => ".relu", + PackedConversionSaturation::Satfinite => ".satfinite", + PackedConversionSaturation::SatfiniteRelu => ".satfinite.relu", + }; + format!( + "cvt.{rounding}{saturation}.{}.{}", + packed_conversion_destination(record), + conversion.source_format.ptx_token() + ) +} diff --git a/crates/cuda-intrinsics-gen/src/render/families/scalar.rs b/crates/cuda-intrinsics-gen/src/render/families/scalar.rs new file mode 100644 index 0000000000..4f55eb7fc1 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/render/families/scalar.rs @@ -0,0 +1,319 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, CatalogIntrinsic, ExtendedMinMaxFormat, ExtendedMinMaxNan, + ExtendedMinMaxOperation, ExtendedMinMaxSubnormal, IntegerMinMaxFormat, IntegerMinMaxOperation, + IntrinsicBackend, ScalarArithmeticFormat, ScalarArithmeticOperation, ScalarArithmeticRounding, + ScalarArithmeticSaturation, ScalarArithmeticSubnormal, ScalarConversionRounding, + ScalarConversionSaturation, ScalarMathFormat, ScalarMathOperation, ScalarMathPrecision, + ScalarMathSubnormal, +}; + +pub(in crate::render) fn scalar_conversion_rounding_attr( + record: &CatalogIntrinsic, +) -> &'static str { + match record + .scalar_conversion + .as_ref() + .expect("scalar-conversion contract") + .rounding + { + ScalarConversionRounding::NearestAway => "ScalarConversionRoundingAttr::NearestAway", + ScalarConversionRounding::NearestEven => "ScalarConversionRoundingAttr::NearestEven", + ScalarConversionRounding::TowardZero => "ScalarConversionRoundingAttr::TowardZero", + } +} + +pub(in crate::render) fn scalar_conversion_saturation_attr( + record: &CatalogIntrinsic, +) -> &'static str { + match record + .scalar_conversion + .as_ref() + .expect("scalar-conversion contract") + .saturation + { + ScalarConversionSaturation::None => "ScalarConversionSaturationAttr::None", + ScalarConversionSaturation::Relu => "ScalarConversionSaturationAttr::Relu", + ScalarConversionSaturation::Satfinite => "ScalarConversionSaturationAttr::Satfinite", + ScalarConversionSaturation::ReluSatfinite => { + "ScalarConversionSaturationAttr::ReluSatfinite" + } + } +} + +pub(in crate::render) fn scalar_conversion_ptx_mnemonic(record: &CatalogIntrinsic) -> String { + format!( + "{}.{}", + record.expected_ptx.mnemonic, + record.expected_ptx.modifiers.join(".") + ) +} + +pub(in crate::render) fn extended_minmax_contract( + record: &CatalogIntrinsic, +) -> &crate::model::ExtendedMinMax { + record + .extended_minmax + .as_ref() + .expect("extended-minmax contract") +} + +pub(in crate::render) fn extended_minmax_format_attr(record: &CatalogIntrinsic) -> &'static str { + match extended_minmax_contract(record).format { + ExtendedMinMaxFormat::F32 => "ExtendedMinMaxFormatAttr::F32", + ExtendedMinMaxFormat::F16 => "ExtendedMinMaxFormatAttr::F16", + ExtendedMinMaxFormat::Bf16 => "ExtendedMinMaxFormatAttr::Bf16", + ExtendedMinMaxFormat::F16x2 => "ExtendedMinMaxFormatAttr::F16x2", + ExtendedMinMaxFormat::Bf16x2 => "ExtendedMinMaxFormatAttr::Bf16x2", + } +} + +pub(in crate::render) fn extended_minmax_operation_attr(record: &CatalogIntrinsic) -> &'static str { + match extended_minmax_contract(record).operation { + ExtendedMinMaxOperation::Min => "ExtendedMinMaxOperationAttr::Min", + ExtendedMinMaxOperation::Max => "ExtendedMinMaxOperationAttr::Max", + } +} + +pub(in crate::render) fn extended_minmax_subnormal_attr(record: &CatalogIntrinsic) -> &'static str { + match extended_minmax_contract(record).subnormal { + ExtendedMinMaxSubnormal::Preserve => "ExtendedMinMaxSubnormalAttr::Preserve", + ExtendedMinMaxSubnormal::Ftz => "ExtendedMinMaxSubnormalAttr::Ftz", + } +} + +pub(in crate::render) fn extended_minmax_nan_attr(record: &CatalogIntrinsic) -> &'static str { + match extended_minmax_contract(record).nan { + ExtendedMinMaxNan::Number => "ExtendedMinMaxNanAttr::Number", + ExtendedMinMaxNan::Nan => "ExtendedMinMaxNanAttr::Nan", + } +} + +pub(in crate::render) fn extended_minmax_xorsign_abs_attr( + record: &CatalogIntrinsic, +) -> &'static str { + if extended_minmax_contract(record).xorsign_abs { + "ExtendedMinMaxXorSignAbsAttr::Enabled" + } else { + "ExtendedMinMaxXorSignAbsAttr::Disabled" + } +} + +pub(in crate::render) fn extended_minmax_carrier(record: &CatalogIntrinsic) -> &'static str { + match extended_minmax_contract(record).format { + ExtendedMinMaxFormat::F32 => "MinMaxCarrier::F32", + ExtendedMinMaxFormat::F16 | ExtendedMinMaxFormat::Bf16 => "MinMaxCarrier::Half16", + ExtendedMinMaxFormat::F16x2 | ExtendedMinMaxFormat::Bf16x2 => "MinMaxCarrier::PackedU32", + } +} + +pub(in crate::render) fn extended_minmax_rust_type(record: &CatalogIntrinsic) -> &'static str { + match extended_minmax_contract(record).format { + ExtendedMinMaxFormat::F32 => "f32", + ExtendedMinMaxFormat::F16 | ExtendedMinMaxFormat::Bf16 => "u16", + ExtendedMinMaxFormat::F16x2 | ExtendedMinMaxFormat::Bf16x2 => "u32", + } +} + +pub(in crate::render) fn extended_minmax_ptx_mnemonic(record: &CatalogIntrinsic) -> String { + format!( + "{}.{}", + record.expected_ptx.mnemonic, + record.expected_ptx.modifiers.join(".") + ) +} + +pub(in crate::render) fn scalar_arithmetic_contract( + record: &CatalogIntrinsic, +) -> &crate::model::ScalarArithmetic { + record + .scalar_arithmetic + .as_ref() + .expect("scalar-arithmetic contract") +} + +pub(in crate::render) fn scalar_arithmetic_format_attr(record: &CatalogIntrinsic) -> &'static str { + match scalar_arithmetic_contract(record).format { + ScalarArithmeticFormat::F32 => "ScalarArithmeticFormatAttr::F32", + ScalarArithmeticFormat::F64 => "ScalarArithmeticFormatAttr::F64", + } +} + +pub(in crate::render) fn scalar_arithmetic_operation_attr( + record: &CatalogIntrinsic, +) -> &'static str { + match scalar_arithmetic_contract(record).operation { + ScalarArithmeticOperation::Mul => "ScalarArithmeticOperationAttr::Mul", + ScalarArithmeticOperation::Div => "ScalarArithmeticOperationAttr::Div", + ScalarArithmeticOperation::Fma => "ScalarArithmeticOperationAttr::Fma", + ScalarArithmeticOperation::Add => "ScalarArithmeticOperationAttr::Add", + } +} + +pub(in crate::render) fn scalar_arithmetic_rounding_attr( + record: &CatalogIntrinsic, +) -> &'static str { + match scalar_arithmetic_contract(record).rounding { + ScalarArithmeticRounding::Rn => "ScalarArithmeticRoundingAttr::Rn", + ScalarArithmeticRounding::Rz => "ScalarArithmeticRoundingAttr::Rz", + ScalarArithmeticRounding::Rm => "ScalarArithmeticRoundingAttr::Rm", + ScalarArithmeticRounding::Rp => "ScalarArithmeticRoundingAttr::Rp", + } +} + +pub(in crate::render) fn scalar_arithmetic_subnormal_attr( + record: &CatalogIntrinsic, +) -> &'static str { + match scalar_arithmetic_contract(record).subnormal { + ScalarArithmeticSubnormal::Preserve => "ScalarArithmeticSubnormalAttr::Preserve", + ScalarArithmeticSubnormal::Ftz => "ScalarArithmeticSubnormalAttr::Ftz", + } +} + +pub(in crate::render) fn scalar_arithmetic_saturation_attr( + record: &CatalogIntrinsic, +) -> &'static str { + match scalar_arithmetic_contract(record).saturation { + ScalarArithmeticSaturation::None => "ScalarArithmeticSaturationAttr::None", + ScalarArithmeticSaturation::Sat => "ScalarArithmeticSaturationAttr::Sat", + } +} + +pub(in crate::render) fn scalar_arithmetic_arity(record: &CatalogIntrinsic) -> usize { + match scalar_arithmetic_contract(record).operation { + ScalarArithmeticOperation::Mul + | ScalarArithmeticOperation::Div + | ScalarArithmeticOperation::Add => 2, + ScalarArithmeticOperation::Fma => 3, + } +} + +pub(in crate::render) fn scalar_arithmetic_rust_type(record: &CatalogIntrinsic) -> &'static str { + match scalar_arithmetic_contract(record).format { + ScalarArithmeticFormat::F32 => "f32", + ScalarArithmeticFormat::F64 => "f64", + } +} + +pub(in crate::render) fn scalar_arithmetic_llvm_type(record: &CatalogIntrinsic) -> &'static str { + match scalar_arithmetic_contract(record).format { + ScalarArithmeticFormat::F32 => "float", + ScalarArithmeticFormat::F64 => "double", + } +} + +pub(in crate::render) fn scalar_arithmetic_ptx_mnemonic(record: &CatalogIntrinsic) -> String { + format!( + "{}.{}", + record.expected_ptx.mnemonic, + record.expected_ptx.modifiers.join(".") + ) +} + +pub(in crate::render) fn scalar_arithmetic_llvm_mechanism( + record: &CatalogIntrinsic, +) -> BackendLoweringMechanism { + record + .backend_lowerings + .iter() + .find(|lowering| lowering.backend == IntrinsicBackend::LlvmNvptx) + .expect("scalar arithmetic has an LLVM-NVPTX route") + .mechanism +} + +pub(in crate::render) fn scalar_math_contract( + record: &CatalogIntrinsic, +) -> &crate::model::ScalarMath { + record.scalar_math.as_ref().expect("scalar-math contract") +} + +pub(in crate::render) fn scalar_math_format_attr(record: &CatalogIntrinsic) -> &'static str { + match scalar_math_contract(record).format { + ScalarMathFormat::F16 => "ScalarMathFormatAttr::F16", + ScalarMathFormat::F32 => "ScalarMathFormatAttr::F32", + ScalarMathFormat::F64 => "ScalarMathFormatAttr::F64", + } +} + +pub(in crate::render) fn scalar_math_operation_attr(record: &CatalogIntrinsic) -> &'static str { + match scalar_math_contract(record).operation { + ScalarMathOperation::Sin => "ScalarMathOperationAttr::Sin", + ScalarMathOperation::Cos => "ScalarMathOperationAttr::Cos", + ScalarMathOperation::Ex2 => "ScalarMathOperationAttr::Ex2", + ScalarMathOperation::Lg2 => "ScalarMathOperationAttr::Lg2", + ScalarMathOperation::Rcp => "ScalarMathOperationAttr::Rcp", + ScalarMathOperation::Rsqrt => "ScalarMathOperationAttr::Rsqrt", + ScalarMathOperation::Sqrt => "ScalarMathOperationAttr::Sqrt", + ScalarMathOperation::Tanh => "ScalarMathOperationAttr::Tanh", + } +} + +pub(in crate::render) fn scalar_math_precision_attr(record: &CatalogIntrinsic) -> &'static str { + match scalar_math_contract(record).precision { + ScalarMathPrecision::Approx => "ScalarMathPrecisionAttr::Approx", + ScalarMathPrecision::Rn => "ScalarMathPrecisionAttr::Rn", + ScalarMathPrecision::Rz => "ScalarMathPrecisionAttr::Rz", + ScalarMathPrecision::Rm => "ScalarMathPrecisionAttr::Rm", + ScalarMathPrecision::Rp => "ScalarMathPrecisionAttr::Rp", + } +} + +pub(in crate::render) fn scalar_math_subnormal_attr(record: &CatalogIntrinsic) -> &'static str { + match scalar_math_contract(record).subnormal { + ScalarMathSubnormal::Preserve => "ScalarMathSubnormalAttr::Preserve", + ScalarMathSubnormal::Ftz => "ScalarMathSubnormalAttr::Ftz", + } +} + +pub(in crate::render) fn scalar_math_llvm_type(record: &CatalogIntrinsic) -> &'static str { + match scalar_math_contract(record).format { + ScalarMathFormat::F16 => "i16", + ScalarMathFormat::F32 => "float", + ScalarMathFormat::F64 => "double", + } +} + +pub(in crate::render) fn scalar_math_ptx_mnemonic(record: &CatalogIntrinsic) -> String { + format!( + "{}.{}", + record.expected_ptx.mnemonic, + record.expected_ptx.modifiers.join(".") + ) +} + +pub(in crate::render) fn scalar_math_llvm_mechanism( + record: &CatalogIntrinsic, +) -> BackendLoweringMechanism { + record + .backend_lowerings + .iter() + .find(|lowering| lowering.backend == IntrinsicBackend::LlvmNvptx) + .expect("scalar math has an LLVM-NVPTX route") + .mechanism +} + +pub(in crate::render) fn integer_minmax_ptx_mnemonic(record: &CatalogIntrinsic) -> &'static str { + let minmax = record + .integer_minmax + .as_ref() + .expect("integer-min/max record"); + match (minmax.format, minmax.operation, minmax.relu) { + (IntegerMinMaxFormat::S32, IntegerMinMaxOperation::Min, true) => "min.relu.s32", + (IntegerMinMaxFormat::S32, IntegerMinMaxOperation::Max, true) => "max.relu.s32", + (IntegerMinMaxFormat::S16x2, IntegerMinMaxOperation::Min, false) => "min.s16x2", + (IntegerMinMaxFormat::S16x2, IntegerMinMaxOperation::Max, false) => "max.s16x2", + (IntegerMinMaxFormat::U16x2, IntegerMinMaxOperation::Min, false) => "min.u16x2", + (IntegerMinMaxFormat::U16x2, IntegerMinMaxOperation::Max, false) => "max.u16x2", + (IntegerMinMaxFormat::S16x2, IntegerMinMaxOperation::Min, true) => "min.relu.s16x2", + (IntegerMinMaxFormat::S16x2, IntegerMinMaxOperation::Max, true) => "max.relu.s16x2", + // `integer_minmax_recipe` rejects those combinations before a record + // can exist. + (IntegerMinMaxFormat::S32, _, false) | (IntegerMinMaxFormat::U16x2, _, true) => { + panic!("{} is outside the closed integer-min/max recipe", record.id) + } + } +} diff --git a/crates/cuda-intrinsics-gen/src/render/families/tcgen05.rs b/crates/cuda-intrinsics-gen/src/render/families/tcgen05.rs new file mode 100644 index 0000000000..21b8a88e81 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/render/families/tcgen05.rs @@ -0,0 +1,1398 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, CatalogHardwareAlternative, CatalogHardwareTarget, CatalogIntrinsic, + IntrinsicBackend, RuntimeValidation, Tcgen05Adapter, Tcgen05CpGroup, Tcgen05CpMember, + Tcgen05LdMultiplicity, Tcgen05LdShape, Tcgen05Mma, Tcgen05MmaAlias, Tcgen05MmaBUsage, + Tcgen05MmaForm, Tcgen05MmaKind, Tcgen05MmaSelectorLayout, Tcgen05Operation, + Tcgen05SourceContract, +}; +use crate::render::common::llvm; +use std::fmt::Write as _; + +pub(in crate::render) fn tcgen05_mma_form_name(form: Tcgen05MmaForm) -> &'static str { + match form { + Tcgen05MmaForm::Shared => "Shared", + Tcgen05MmaForm::Tensor => "Tensor", + Tcgen05MmaForm::TensorAshift => "TensorAshift", + Tcgen05MmaForm::SpShared => "SpShared", + Tcgen05MmaForm::SpTensor => "SpTensor", + Tcgen05MmaForm::SpTensorAshift => "SpTensorAshift", + Tcgen05MmaForm::WsShared => "WsShared", + Tcgen05MmaForm::WsSharedZeroColMask => "WsSharedZeroColMask", + Tcgen05MmaForm::WsSpShared => "WsSpShared", + Tcgen05MmaForm::WsSpSharedZeroColMask => "WsSpSharedZeroColMask", + Tcgen05MmaForm::WsSpTensor => "WsSpTensor", + Tcgen05MmaForm::WsSpTensorZeroColMask => "WsSpTensorZeroColMask", + Tcgen05MmaForm::WsTensor => "WsTensor", + Tcgen05MmaForm::WsTensorZeroColMask => "WsTensorZeroColMask", + } +} + +pub(in crate::render) fn tcgen05_mma_form_attr(form: Tcgen05MmaForm) -> String { + format!("Tcgen05MmaFormAttr::{}", tcgen05_mma_form_name(form)) +} + +pub(in crate::render) fn tcgen05_mma_kind_attr(kind: Tcgen05MmaKind) -> &'static str { + match kind { + Tcgen05MmaKind::F16 => "Tcgen05MmaKindAttr::F16", + Tcgen05MmaKind::Tf32 => "Tcgen05MmaKindAttr::Tf32", + Tcgen05MmaKind::F8f6f4 => "Tcgen05MmaKindAttr::F8f6f4", + Tcgen05MmaKind::I8 => "Tcgen05MmaKindAttr::I8", + } +} + +pub(in crate::render) fn tcgen05_mma_b_usage_attr(usage: Tcgen05MmaBUsage) -> &'static str { + match usage { + Tcgen05MmaBUsage::Discard => "Tcgen05MmaBUsageAttr::Discard", + Tcgen05MmaBUsage::Fill => "Tcgen05MmaBUsageAttr::Fill", + Tcgen05MmaBUsage::Use => "Tcgen05MmaBUsageAttr::Use", + Tcgen05MmaBUsage::LastUse => "Tcgen05MmaBUsageAttr::LastUse", + } +} + +pub(in crate::render) fn tcgen05_mma_runtime_parameters( + mma: &Tcgen05Mma, +) -> Vec<(&'static str, &'static str)> { + if mma.alias.is_some() && mma.form == Tcgen05MmaForm::WsTensor { + return vec![ + ("d_tmem", "u32"), + ("a_tmem", "u32"), + ("legacy_a_desc", "u64"), + ("b_desc", "u64"), + ("idesc", "u32"), + ("enable_d", "bool"), + ]; + } + let mut parameters = match mma.form { + Tcgen05MmaForm::Shared + | Tcgen05MmaForm::SpShared + | Tcgen05MmaForm::WsShared + | Tcgen05MmaForm::WsSharedZeroColMask + | Tcgen05MmaForm::WsSpShared + | Tcgen05MmaForm::WsSpSharedZeroColMask => vec![ + ("d_tmem", "u32"), + ("a_desc", "u64"), + ("b_desc", "u64"), + ("idesc", "u32"), + ("enable_d", "bool"), + ], + Tcgen05MmaForm::Tensor + | Tcgen05MmaForm::TensorAshift + | Tcgen05MmaForm::SpTensor + | Tcgen05MmaForm::SpTensorAshift + | Tcgen05MmaForm::WsSpTensor + | Tcgen05MmaForm::WsSpTensorZeroColMask + | Tcgen05MmaForm::WsTensor + | Tcgen05MmaForm::WsTensorZeroColMask => vec![ + ("d_tmem", "u32"), + ("a_tmem", "u32"), + ("b_desc", "u64"), + ("idesc", "u32"), + ("enable_d", "bool"), + ], + }; + if matches!( + mma.form, + Tcgen05MmaForm::SpShared + | Tcgen05MmaForm::SpTensor + | Tcgen05MmaForm::SpTensorAshift + | Tcgen05MmaForm::WsSpShared + | Tcgen05MmaForm::WsSpSharedZeroColMask + | Tcgen05MmaForm::WsSpTensor + | Tcgen05MmaForm::WsSpTensorZeroColMask + ) { + parameters.push(("metadata_tmem", "u32")); + } + if matches!( + mma.form, + Tcgen05MmaForm::WsSharedZeroColMask + | Tcgen05MmaForm::WsSpSharedZeroColMask + | Tcgen05MmaForm::WsSpTensorZeroColMask + | Tcgen05MmaForm::WsTensorZeroColMask + ) { + parameters.push(("zero_column_mask", "u64")); + } + parameters +} + +pub(in crate::render) fn tcgen05_mma_selector_parameters( + layout: Tcgen05MmaSelectorLayout, +) -> [(&'static str, &'static str); 3] { + match layout { + Tcgen05MmaSelectorLayout::Base { .. } => [ + ("KIND", "kind"), + ("CTA_GROUP", "cta_group"), + ("COLLECTOR_A", "collector_a"), + ], + Tcgen05MmaSelectorLayout::WarpSpecialized { .. } => [ + ("KIND", "kind"), + ("B_BUFFER", "b_buffer"), + ("B_USAGE", "b_usage"), + ], + } +} + +pub(in crate::render) fn tcgen05_mma_is_ws(form: Tcgen05MmaForm) -> bool { + matches!( + form, + Tcgen05MmaForm::WsShared + | Tcgen05MmaForm::WsSharedZeroColMask + | Tcgen05MmaForm::WsSpShared + | Tcgen05MmaForm::WsSpSharedZeroColMask + | Tcgen05MmaForm::WsSpTensor + | Tcgen05MmaForm::WsSpTensorZeroColMask + | Tcgen05MmaForm::WsTensor + | Tcgen05MmaForm::WsTensorZeroColMask + ) +} + +fn tcgen05_mma_is_sparse(form: Tcgen05MmaForm) -> bool { + matches!( + form, + Tcgen05MmaForm::SpShared + | Tcgen05MmaForm::SpTensor + | Tcgen05MmaForm::SpTensorAshift + | Tcgen05MmaForm::WsSpShared + | Tcgen05MmaForm::WsSpSharedZeroColMask + | Tcgen05MmaForm::WsSpTensor + | Tcgen05MmaForm::WsSpTensorZeroColMask + ) +} + +fn tcgen05_mma_is_tensor_a(form: Tcgen05MmaForm) -> bool { + matches!( + form, + Tcgen05MmaForm::Tensor + | Tcgen05MmaForm::TensorAshift + | Tcgen05MmaForm::SpTensor + | Tcgen05MmaForm::SpTensorAshift + | Tcgen05MmaForm::WsSpTensor + | Tcgen05MmaForm::WsSpTensorZeroColMask + | Tcgen05MmaForm::WsTensor + | Tcgen05MmaForm::WsTensorZeroColMask + ) +} + +fn tcgen05_mma_is_ashift(form: Tcgen05MmaForm) -> bool { + matches!( + form, + Tcgen05MmaForm::TensorAshift | Tcgen05MmaForm::SpTensorAshift + ) +} + +fn tcgen05_mma_has_zero_col_mask(form: Tcgen05MmaForm) -> bool { + matches!( + form, + Tcgen05MmaForm::WsSharedZeroColMask + | Tcgen05MmaForm::WsSpSharedZeroColMask + | Tcgen05MmaForm::WsSpTensorZeroColMask + | Tcgen05MmaForm::WsTensorZeroColMask + ) +} + +fn tcgen05_mma_kind_name(kind: Tcgen05MmaKind) -> &'static str { + match kind { + Tcgen05MmaKind::F16 => "f16", + Tcgen05MmaKind::Tf32 => "tf32", + Tcgen05MmaKind::F8f6f4 => "f8f6f4", + Tcgen05MmaKind::I8 => "i8", + } +} + +fn tcgen05_mma_b_usage_name(usage: Tcgen05MmaBUsage) -> &'static str { + match usage { + Tcgen05MmaBUsage::Discard => "discard", + Tcgen05MmaBUsage::LastUse => "lastuse", + Tcgen05MmaBUsage::Fill => "fill", + Tcgen05MmaBUsage::Use => "use", + } +} + +pub(in crate::render) fn tcgen05_mma_inline_asm( + form: Tcgen05MmaForm, + kind: Tcgen05MmaKind, + cta_group: u8, + collector_a: Option<&str>, + b_buffer: Option, + b_usage: Option, +) -> (String, String) { + let mut instruction = "tcgen05.mma".to_owned(); + if tcgen05_mma_is_ws(form) { + instruction.push_str(".ws"); + } + if tcgen05_mma_is_sparse(form) { + instruction.push_str(".sp"); + } + write!( + instruction, + ".cta_group::{cta_group}.kind::{}", + tcgen05_mma_kind_name(kind) + ) + .unwrap(); + if tcgen05_mma_is_ws(form) { + write!( + instruction, + ".collector::b{}::{}", + b_buffer.expect("warp-specialized B buffer"), + tcgen05_mma_b_usage_name(b_usage.expect("warp-specialized B usage")) + ) + .unwrap(); + } else { + write!( + instruction, + ".collector::a::{}", + collector_a.expect("base collector A usage") + ) + .unwrap(); + if tcgen05_mma_is_ashift(form) { + instruction.push_str(".ashift"); + } + } + + let a = if tcgen05_mma_is_tensor_a(form) { + "[$1]" + } else { + "$1" + }; + write!(instruction, " [$0], {a}, $2").unwrap(); + if tcgen05_mma_is_sparse(form) { + instruction.push_str(", [$5]"); + } + instruction.push_str(", $3, %enable_pred"); + if tcgen05_mma_has_zero_col_mask(form) { + instruction.push_str(if tcgen05_mma_is_sparse(form) { + ", $6" + } else { + ", $5" + }); + } + instruction.push(';'); + + let template = + format!("{{ .reg .pred %enable_pred; setp.ne.s32 %enable_pred, $4, 0; {instruction} }}"); + let mut constraints = vec![ + "r", + if tcgen05_mma_is_tensor_a(form) { + "r" + } else { + "l" + }, + "l", + "r", + "r", + ]; + if tcgen05_mma_is_sparse(form) { + constraints.push("r"); + } + if tcgen05_mma_has_zero_col_mask(form) { + constraints.push("l"); + } + constraints.push("~{memory}"); + (template, constraints.join(",")) +} + +pub(in crate::render) fn tcgen05_participation_doc( + operation: Tcgen05Operation, +) -> Option<&'static str> { + match operation { + Tcgen05Operation::AllocCg2 + | Tcgen05Operation::DeallocCg2 + | Tcgen05Operation::RelinquishAllocPermitCg2 => Some( + "One full warp in each peer CTA must execute this instruction; lanes within each warp must execute uniformly with the same operands.", + ), + Tcgen05Operation::Alloc + | Tcgen05Operation::Dealloc + | Tcgen05Operation::RelinquishAllocPermit + | Tcgen05Operation::Ld16x256bX8Pure + | Tcgen05Operation::Ld16x256bPure + | Tcgen05Operation::LoadWait + | Tcgen05Operation::StoreWait => { + Some("One full warp must execute this instruction uniformly with the same operands.") + } + Tcgen05Operation::Commit + | Tcgen05Operation::CommitSharedCluster + | Tcgen05Operation::CommitMulticast + | Tcgen05Operation::ShiftDown => Some("One thread in the CTA issues this instruction."), + Tcgen05Operation::CommitCg2 + | Tcgen05Operation::CommitSharedClusterCg2 + | Tcgen05Operation::CommitMulticastCg2 + | Tcgen05Operation::ShiftDownCg2 => Some( + "One thread in the CTA pair issues this instruction; the peer CTA must be active and must not have exited.", + ), + _ => None, + } +} + +pub(in crate::render) fn tcgen05_is_commit(operation: Tcgen05Operation) -> bool { + matches!( + operation, + Tcgen05Operation::Commit + | Tcgen05Operation::CommitSharedCluster + | Tcgen05Operation::CommitMulticast + | Tcgen05Operation::CommitCg2 + | Tcgen05Operation::CommitSharedClusterCg2 + | Tcgen05Operation::CommitMulticastCg2 + ) +} + +pub(in crate::render) fn tcgen05_is_multicast_commit(operation: Tcgen05Operation) -> bool { + matches!( + operation, + Tcgen05Operation::CommitMulticast | Tcgen05Operation::CommitMulticastCg2 + ) +} + +pub(in crate::render) fn tcgen05_is_shift(operation: Tcgen05Operation) -> bool { + matches!( + operation, + Tcgen05Operation::ShiftDown | Tcgen05Operation::ShiftDownCg2 + ) +} + +fn tcgen05_cp_ptx_suffix(member: Tcgen05CpMember) -> &'static str { + use Tcgen05CpMember::*; + match member { + M128x128bB4x16P64 => "128x128b.b8x16.b4x16_p64", + M128x128bB6x16P32 => "128x128b.b8x16.b6x16_p32", + M128x128b => "128x128b", + M128x256bB4x16P64 => "128x256b.b8x16.b4x16_p64", + M128x256bB6x16P32 => "128x256b.b8x16.b6x16_p32", + M32x128bWarpx4B4x16P64 => "32x128b.warpx4.b8x16.b4x16_p64", + M32x128bWarpx4B6x16P32 => "32x128b.warpx4.b8x16.b6x16_p32", + M32x128bWarpx4 => "32x128b.warpx4", + M4x256bB4x16P64 => "4x256b.b8x16.b4x16_p64", + M4x256bB6x16P32 => "4x256b.b8x16.b6x16_p32", + M4x256b => "4x256b", + M64x128bWarpx2Pair0123B4x16P64 => "64x128b.warpx2::01_23.b8x16.b4x16_p64", + M64x128bWarpx2Pair0123B6x16P32 => "64x128b.warpx2::01_23.b8x16.b6x16_p32", + M64x128bWarpx2Pair0123 => "64x128b.warpx2::01_23", + M64x128bWarpx2Pair0213B4x16P64 => "64x128b.warpx2::02_13.b8x16.b4x16_p64", + M64x128bWarpx2Pair0213B6x16P32 => "64x128b.warpx2::02_13.b8x16.b6x16_p32", + M64x128bWarpx2Pair0213 => "64x128b.warpx2::02_13", + } +} + +fn tcgen05_ld_shape_label(shape: Tcgen05LdShape) -> &'static str { + match shape { + Tcgen05LdShape::M16x32bx2 => "16x32bx2", + Tcgen05LdShape::M16x64b => "16x64b", + Tcgen05LdShape::M16x128b => "16x128b", + Tcgen05LdShape::M16x256b => "16x256b", + Tcgen05LdShape::M32x32b => "32x32b", + } +} + +fn tcgen05_ld_multiplicity_label(multiplicity: Tcgen05LdMultiplicity) -> &'static str { + match multiplicity { + Tcgen05LdMultiplicity::X1 => "x1", + Tcgen05LdMultiplicity::X2 => "x2", + Tcgen05LdMultiplicity::X4 => "x4", + Tcgen05LdMultiplicity::X8 => "x8", + Tcgen05LdMultiplicity::X16 => "x16", + Tcgen05LdMultiplicity::X32 => "x32", + Tcgen05LdMultiplicity::X64 => "x64", + Tcgen05LdMultiplicity::X128 => "x128", + } +} + +pub(in crate::render) fn tcgen05_ld_register_count(record: &CatalogIntrinsic) -> usize { + let ld = record + .tcgen05 + .as_ref() + .and_then(|tcgen05| tcgen05.ld) + .expect("generated tcgen05 load identity"); + ld.shape.register_multiplier() * ld.multiplicity.count() +} + +pub(in crate::render) fn tcgen05_st_register_count(record: &CatalogIntrinsic) -> usize { + let st = record + .tcgen05 + .as_ref() + .and_then(|tcgen05| tcgen05.st) + .expect("generated tcgen05 store identity"); + st.shape.register_multiplier() * st.multiplicity.count() +} + +fn tcgen05_mma_target_matrix_is_closed( + target: &crate::model::CatalogTargetRequirement, + backend: IntrinsicBackend, + fixed_kind: Option, +) -> bool { + let CatalogHardwareTarget::TargetMatrix { contracts } = &target.hardware else { + return false; + }; + let expected_kinds: Vec<&str> = fixed_kind.map_or_else( + || vec!["f16", "tf32", "f8f6f4", "i8"], + |kind| vec![tcgen05_mma_kind_name(kind)], + ); + if target.minimum_ptx.to_string() != "8.6" || contracts.len() != expected_kinds.len() { + return false; + } + let common = match backend { + IntrinsicBackend::LlvmNvptx => vec![ + ( + CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, + "8.6", + ), + (CatalogHardwareAlternative::FamilyTarget { sm: 100 }, "8.8"), + ( + CatalogHardwareAlternative::ExactArchitecture { sm: 101 }, + "8.6", + ), + (CatalogHardwareAlternative::FamilyTarget { sm: 101 }, "8.8"), + ( + CatalogHardwareAlternative::ExactArchitecture { sm: 103 }, + "8.8", + ), + (CatalogHardwareAlternative::FamilyTarget { sm: 103 }, "8.8"), + ( + CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, + "9.0", + ), + (CatalogHardwareAlternative::FamilyTarget { sm: 110 }, "9.0"), + ], + IntrinsicBackend::LibNvvm => vec![ + ( + CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, + "8.6", + ), + (CatalogHardwareAlternative::FamilyTarget { sm: 100 }, "8.8"), + ( + CatalogHardwareAlternative::ExactArchitecture { sm: 103 }, + "8.8", + ), + (CatalogHardwareAlternative::FamilyTarget { sm: 103 }, "8.8"), + ( + CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, + "9.0", + ), + (CatalogHardwareAlternative::FamilyTarget { sm: 110 }, "9.0"), + ], + }; + let i8 = match backend { + IntrinsicBackend::LlvmNvptx => vec![ + ( + CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, + "8.6", + ), + ( + CatalogHardwareAlternative::ExactArchitecture { sm: 101 }, + "8.6", + ), + ( + CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, + "9.0", + ), + ], + IntrinsicBackend::LibNvvm => vec![ + ( + CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, + "8.6", + ), + ( + CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, + "9.0", + ), + ], + }; + expected_kinds.into_iter().all(|kind| { + let Some(contract) = contracts.iter().find(|contract| { + contract.selectors.len() == 1 + && contract.selectors[0].name == "kind" + && contract.selectors[0].value == kind + }) else { + return false; + }; + let expected = if kind == "i8" { &i8 } else { &common }; + contract.alternatives.len() == expected.len() + && contract.alternatives.iter().zip(expected).all( + |(actual, (hardware, minimum_ptx))| { + actual.hardware == *hardware && actual.minimum_ptx.to_string() == *minimum_ptx + }, + ) + }) +} + +fn tcgen05_mma_render_contract( + record: &CatalogIntrinsic, + tcgen05: &crate::model::Tcgen05, + mma: &Tcgen05Mma, + llvm_route: &crate::model::CatalogBackendLowering, + libnvvm_route: &crate::model::CatalogBackendLowering, +) -> bool { + if tcgen05.operation != Tcgen05Operation::Mma + || tcgen05.cp.is_some() + || tcgen05.ld.is_some() + || tcgen05.st.is_some() + || tcgen05.source_contract != Tcgen05SourceContract::TablegenSelectionChangesPtx + || tcgen05.runtime_validation != RuntimeValidation::Unexecuted + || record.rust.module != "tcgen05" + || record.rust.safe + || record.rust.must_use + || record.rust.result != "()" + || record.dialect.op_type != "Tcgen05MmaOp" + || record.dialect.op_name != "nvvm.tcgen05_mma" + || !record.dialect.results.is_empty() + || record.lowering != "generated_tcgen05_mma" + || record.semantics.pure + || record.semantics.memory != "read_write" + || !record.semantics.convergent + || record.semantics.execution_scope != "thread" + || record.target.minimum_ptx.to_string() != "8.6" + || record.target.hardware != mma.llvm_target.hardware + || record.target.targets != "sm_100a|sm_101a|sm_103a|sm_110a" + || llvm_route.backend != IntrinsicBackend::LlvmNvptx + || llvm_route.mechanism != BackendLoweringMechanism::InlinePtx + || llvm_route.target != mma.llvm_target + || !tcgen05_mma_target_matrix_is_closed( + &llvm_route.target, + IntrinsicBackend::LlvmNvptx, + mma.fixed_selectors.map(|fixed| fixed.kind), + ) + || libnvvm_route.backend != IntrinsicBackend::LibNvvm + || libnvvm_route.mechanism != BackendLoweringMechanism::InlinePtx + || libnvvm_route.target != mma.libnvvm_target + || !tcgen05_mma_target_matrix_is_closed( + &libnvvm_route.target, + IntrinsicBackend::LibNvvm, + mma.fixed_selectors.map(|fixed| fixed.kind), + ) + { + return false; + } + + let mut expected_operands = vec![ + "i32".to_owned(), + if tcgen05_mma_is_tensor_a(mma.form) { + "i32".to_owned() + } else { + "i64".to_owned() + }, + "i64".to_owned(), + "i32".to_owned(), + "i1".to_owned(), + ]; + if tcgen05_mma_is_sparse(mma.form) { + expected_operands.push("i32".into()); + } + if tcgen05_mma_has_zero_col_mask(mma.form) { + expected_operands.push("i64".into()); + } + + let llvm = llvm(record); + let representative_kind = mma + .fixed_selectors + .map_or(Tcgen05MmaKind::F16, |fixed| fixed.kind); + let representative_buffer = mma.fixed_selectors.map_or(0, |fixed| fixed.b_buffer); + let representative_usage = mma + .fixed_selectors + .map_or(Tcgen05MmaBUsage::Discard, |fixed| fixed.b_usage); + let representative = if tcgen05_mma_is_ws(mma.form) { + tcgen05_mma_inline_asm( + mma.form, + representative_kind, + 1, + None, + Some(representative_buffer), + Some(representative_usage), + ) + } else { + tcgen05_mma_inline_asm( + mma.form, + representative_kind, + 1, + Some("discard"), + None, + None, + ) + }; + let instruction = representative + .0 + .find("tcgen05.mma") + .map_or("", |start| &representative.0[start..]); + let head = instruction.split_whitespace().next().unwrap_or_default(); + let mut components = head.split('.'); + let expected_mnemonic = components.next().unwrap_or_default(); + let expected_modifiers = components.map(str::to_owned).collect::>(); + let mut expected_ptx_operands = vec![crate::ptx::OperandPattern::Address]; + expected_ptx_operands.push(if tcgen05_mma_is_tensor_a(mma.form) { + crate::ptx::OperandPattern::Address + } else { + crate::ptx::OperandPattern::Register + }); + expected_ptx_operands.push(crate::ptx::OperandPattern::Register); + if tcgen05_mma_is_sparse(mma.form) { + expected_ptx_operands.push(crate::ptx::OperandPattern::Address); + } + expected_ptx_operands.extend([ + crate::ptx::OperandPattern::Register, + crate::ptx::OperandPattern::Exact { + value: "%enable_pred".into(), + }, + ]); + if tcgen05_mma_has_zero_col_mask(mma.form) { + expected_ptx_operands.push(crate::ptx::OperandPattern::Register); + } + if record.dialect.operands != expected_operands + || !llvm.results.is_empty() + || record.expected_ptx.mnemonic != expected_mnemonic + || record.expected_ptx.modifiers != expected_modifiers + || record.expected_ptx.operands != expected_ptx_operands + { + return false; + } + + match (tcgen05.adapter, mma.alias, mma.fixed_selectors) { + (Tcgen05Adapter::MmaDirectSelectors, None, None) => { + let selector_indices = match mma.selector_layout { + Tcgen05MmaSelectorLayout::Base { + kind_argument, + cta_group_argument, + collector_a_argument, + collector_a_upper_exclusive, + } => { + if collector_a_upper_exclusive + != if tcgen05_mma_is_ashift(mma.form) { + 2 + } else { + 4 + } + { + return false; + } + [kind_argument, cta_group_argument, collector_a_argument] + } + Tcgen05MmaSelectorLayout::WarpSpecialized { + kind_argument, + b_buffer_argument, + b_usage_argument, + } => [kind_argument, b_buffer_argument, b_usage_argument], + }; + let first_selector = expected_operands.len() as u8; + record.rust.arguments.len() == expected_operands.len() + 3 + && record + .rust + .arguments + .ends_with(&["u32".into(), "u32".into(), "u32".into()]) + && llvm.arguments.len() == expected_operands.len() + 3 + && selector_indices == [first_selector, first_selector + 1, first_selector + 2] + && tcgen05_mma_is_ws(mma.form) + == matches!( + mma.selector_layout, + Tcgen05MmaSelectorLayout::WarpSpecialized { .. } + ) + } + ( + Tcgen05Adapter::MmaDirectSelectors, + Some( + Tcgen05MmaAlias::E4m3 + | Tcgen05MmaAlias::E5m2 + | Tcgen05MmaAlias::E2m3 + | Tcgen05MmaAlias::E3m2 + | Tcgen05MmaAlias::E2m1, + ), + Some(fixed), + ) => { + mma.form == Tcgen05MmaForm::Shared + && fixed.kind == Tcgen05MmaKind::F8f6f4 + && fixed.b_buffer == 0 + && fixed.b_usage == Tcgen05MmaBUsage::Discard + && matches!( + mma.selector_layout, + Tcgen05MmaSelectorLayout::Base { + kind_argument: 5, + cta_group_argument: 6, + collector_a_argument: 7, + collector_a_upper_exclusive: 4, + } + ) + && record.rust.arguments == ["u32", "u64", "u64", "u32", "bool"] + && record.dialect.operands == ["i32", "i64", "i64", "i32", "i1"] + && llvm.arguments.len() == 8 + } + ( + Tcgen05Adapter::MmaWsFixedSelectorsDropLegacyADescriptor, + Some( + Tcgen05MmaAlias::E4m3 + | Tcgen05MmaAlias::E5m2 + | Tcgen05MmaAlias::E2m3 + | Tcgen05MmaAlias::E3m2 + | Tcgen05MmaAlias::E2m1, + ), + Some(fixed), + ) => { + mma.form == Tcgen05MmaForm::WsTensor + && fixed.kind == Tcgen05MmaKind::F8f6f4 + && fixed.b_buffer == 0 + && fixed.b_usage == Tcgen05MmaBUsage::Discard + && record.rust.arguments == ["u32", "u32", "u64", "u64", "u32", "bool"] + && record.dialect.operands == ["i32", "i32", "i64", "i32", "i1"] + && llvm.arguments.len() == 8 + } + _ => false, + } +} + +pub(in crate::render) fn tcgen05_render_contract(record: &CatalogIntrinsic) -> bool { + let Some(tcgen05) = &record.tcgen05 else { + return false; + }; + let [llvm_route, libnvvm_route] = record.backend_lowerings.as_slice() else { + return false; + }; + if let Some(mma) = &tcgen05.mma { + return tcgen05_mma_render_contract(record, tcgen05, mma, llvm_route, libnvvm_route); + } + let llvm_hardware = CatalogHardwareTarget::AnyOf { + alternatives: vec![ + CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 101 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 103 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, + ], + }; + let libnvvm_hardware = CatalogHardwareTarget::AnyOf { + alternatives: vec![ + CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 103 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, + ], + }; + if record.rust.module != "tcgen05" + || record.rust.must_use != (tcgen05.operation == Tcgen05Operation::Ld) + || record.lowering != "generated_tcgen05" + || !record.semantics.convergent + || record.semantics.execution_scope != tcgen05.operation.execution_scope() + || record.target.minimum_ptx.to_string() != "8.6" + || record.target.targets != "sm_100a|sm_101a|sm_103a|sm_110a" + || llvm_route.backend != IntrinsicBackend::LlvmNvptx + || llvm_route.mechanism != BackendLoweringMechanism::InlinePtx + || llvm_route.target.minimum_ptx.to_string() != "8.6" + || llvm_route.target.hardware != llvm_hardware + || libnvvm_route.backend != IntrinsicBackend::LibNvvm + || libnvvm_route.mechanism != BackendLoweringMechanism::InlinePtx + || libnvvm_route.target.minimum_ptx.to_string() != "8.6" + || libnvvm_route.target.hardware != libnvvm_hardware + || tcgen05.runtime_validation != RuntimeValidation::Unexecuted + { + return false; + } + match tcgen05.cp { + Some(cp) => { + let expected_operation = match cp.group { + Tcgen05CpGroup::Cg1 => Tcgen05Operation::CpSmemToTmem, + Tcgen05CpGroup::Cg2 => Tcgen05Operation::CpSmemToTmemCg2, + }; + let group = match cp.group { + Tcgen05CpGroup::Cg1 => "cta_group::1", + Tcgen05CpGroup::Cg2 => "cta_group::2", + }; + let modifiers = std::iter::once("cp") + .chain(std::iter::once(group)) + .chain(tcgen05_cp_ptx_suffix(cp.member).split('.')) + .collect::>(); + if tcgen05.operation != expected_operation + || tcgen05.ld.is_some() + || tcgen05.st.is_some() + || record.semantics.pure + || record.semantics.memory != "read_write" + || record.expected_ptx.mnemonic != "tcgen05" + || record.expected_ptx.modifiers != modifiers + || record.expected_ptx.operands + != [ + crate::ptx::OperandPattern::Address, + crate::ptx::OperandPattern::Register, + ] + { + return false; + } + } + None if !matches!( + tcgen05.operation, + Tcgen05Operation::CpSmemToTmem | Tcgen05Operation::CpSmemToTmemCg2 + ) => {} + None if matches!( + record.id.as_str(), + "tcgen05_cp_smem_to_tmem" | "tcgen05_cp_smem_to_tmem_cg2" + ) => {} + None => return false, + } + match tcgen05.ld { + Some(ld) => { + let count = tcgen05_ld_register_count(record); + let mut modifiers: Vec = vec![ + "ld".into(), + "sync".into(), + "aligned".into(), + tcgen05_ld_shape_label(ld.shape).into(), + tcgen05_ld_multiplicity_label(ld.multiplicity).into(), + ]; + if ld.pack16 { + modifiers.push("pack::16b".into()); + } + modifiers.push("b32".into()); + let destination = crate::ptx::OperandPattern::RegisterList { length: count }; + let mut operands = vec![destination, crate::ptx::OperandPattern::Address]; + if ld.shape == Tcgen05LdShape::M16x32bx2 { + operands.push(crate::ptx::OperandPattern::Immediate); + } + if tcgen05.operation != Tcgen05Operation::Ld + || tcgen05.cp.is_some() + || tcgen05.st.is_some() + || record.expected_ptx.mnemonic != "tcgen05" + || record.expected_ptx.modifiers != modifiers + || record.expected_ptx.operands != operands + { + return false; + } + } + None if tcgen05.operation != Tcgen05Operation::Ld => {} + None => return false, + } + match tcgen05.st { + Some(st) => { + let count = tcgen05_st_register_count(record); + let mut modifiers: Vec = vec![ + "st".into(), + "sync".into(), + "aligned".into(), + tcgen05_ld_shape_label(st.shape).into(), + tcgen05_ld_multiplicity_label(st.multiplicity).into(), + ]; + if st.unpack16 { + modifiers.push("unpack::16b".into()); + } + modifiers.push("b32".into()); + let data = crate::ptx::OperandPattern::RegisterList { length: count }; + let mut operands = vec![crate::ptx::OperandPattern::Address]; + if st.shape == Tcgen05LdShape::M16x32bx2 { + operands.push(crate::ptx::OperandPattern::Immediate); + } + operands.push(data); + if tcgen05.operation != Tcgen05Operation::St + || tcgen05.cp.is_some() + || tcgen05.ld.is_some() + || record.expected_ptx.mnemonic != "tcgen05" + || record.expected_ptx.modifiers != modifiers + || record.expected_ptx.operands != operands + { + return false; + } + } + None if tcgen05.operation != Tcgen05Operation::St => {} + None => return false, + } + let llvm = llvm(record); + match (tcgen05.operation, tcgen05.adapter) { + ( + Tcgen05Operation::Alloc | Tcgen05Operation::AllocCg2, + Tcgen05Adapter::SharedPointerColumnsToVoid, + ) => { + !record.rust.safe + && record.rust.arguments == ["*mut u32", "u32"] + && record.rust.result == "()" + && record.dialect.operands == ["ptr", "i32"] + && record.dialect.results.is_empty() + && llvm.arguments == ["shared_ptr", "i32"] + && llvm.results.is_empty() + && tcgen05.source_contract == Tcgen05SourceContract::ExactTablegenSelection + } + ( + Tcgen05Operation::Dealloc | Tcgen05Operation::DeallocCg2, + Tcgen05Adapter::TmemAddressColumnsToVoid, + ) => { + !record.rust.safe + && record.rust.arguments == ["u32", "u32"] + && record.rust.result == "()" + && record.dialect.operands == ["i32", "i32"] + && llvm.arguments == ["tmem_ptr", "i32"] + && tcgen05.source_contract == Tcgen05SourceContract::ExactTablegenSelection + } + ( + Tcgen05Operation::RelinquishAllocPermit + | Tcgen05Operation::RelinquishAllocPermitCg2 + | Tcgen05Operation::FenceBeforeThreadSync + | Tcgen05Operation::FenceAfterThreadSync + | Tcgen05Operation::LoadWait + | Tcgen05Operation::StoreWait, + Tcgen05Adapter::NoOperands, + ) => { + record.rust.safe + && record.rust.arguments.is_empty() + && record.rust.result == "()" + && record.dialect.operands.is_empty() + && llvm.arguments.is_empty() + && tcgen05.source_contract == Tcgen05SourceContract::ExactTablegenSelection + } + ( + Tcgen05Operation::Commit + | Tcgen05Operation::CommitCg2 + | Tcgen05Operation::CommitSharedCluster + | Tcgen05Operation::CommitSharedClusterCg2, + Tcgen05Adapter::BarrierPointerToVoid, + ) => { + !record.rust.safe + && record.rust.arguments == ["*mut u64"] + && record.rust.result == "()" + && record.dialect.operands == ["ptr"] + && llvm.arguments.len() == 1 + && if matches!( + tcgen05.operation, + Tcgen05Operation::Commit | Tcgen05Operation::CommitCg2 + ) { + tcgen05.source_contract == Tcgen05SourceContract::TablegenSelectionChangesPtx + && llvm.arguments == ["ptr"] + } else { + tcgen05.source_contract == Tcgen05SourceContract::ExactTablegenSelection + && llvm.arguments == ["shared_ptr"] + } + } + ( + Tcgen05Operation::MmaWsF16 | Tcgen05Operation::MmaWsBf16 | Tcgen05Operation::MmaWsTf32, + Tcgen05Adapter::MmaWsDropLegacyADescriptor, + ) => { + !record.rust.safe + && record.rust.arguments == ["u32", "u32", "u64", "u64", "u32", "bool"] + && record.rust.result == "()" + && record.dialect.operands == ["i32", "i32", "i64", "i64", "i32", "i1"] + && llvm.arguments + == [ + "tmem_ptr", "tmem_ptr", "i64", "i32", "i1", "i32", "i32", "i32", + ] + && tcgen05.source_contract == Tcgen05SourceContract::TablegenSelectionChangesPtx + && (tcgen05.operation != Tcgen05Operation::MmaWsBf16 + || (record.expected_ptx.modifiers.contains(&"kind::f16".into()) + && !record.expected_ptx.modifiers.contains(&"kind::bf16".into()))) + } + ( + Tcgen05Operation::MmaF16 | Tcgen05Operation::MmaF16Cg2, + Tcgen05Adapter::MmaInjectZeroDisableLanes, + ) => { + !record.rust.safe + && record.rust.arguments == ["u32", "u64", "u64", "u32", "bool"] + && record.rust.result == "()" + && record.dialect.operands == ["i32", "i64", "i64", "i32", "i1"] + && llvm.arguments.len() == 8 + && tcgen05.source_contract + == Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection + } + ( + Tcgen05Operation::CpSmemToTmem | Tcgen05Operation::CpSmemToTmemCg2, + Tcgen05Adapter::TmemDescriptorToVoid, + ) => { + !record.rust.safe + && record.rust.arguments == ["u32", "u64"] + && record.rust.result == "()" + && record.dialect.operands == ["i32", "i64"] + && llvm.arguments == ["tmem_ptr", "i64"] + && tcgen05.source_contract == Tcgen05SourceContract::ExactTablegenSelection + } + (Tcgen05Operation::Ld16x256bX8Pure, Tcgen05Adapter::TmemToF32x32) => { + !record.rust.safe + && record.rust.arguments == ["u32"] + && record.rust.result == "[f32; 32]" + && record.dialect.operands == ["i32"] + && record.dialect.results == vec!["f32"; 32] + && llvm.arguments == ["tmem_ptr", "i1"] + && llvm.results == ["v32i32"] + && tcgen05.source_contract + == Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection + } + (Tcgen05Operation::Ld16x256bPure, Tcgen05Adapter::TmemToF32x4) => { + !record.rust.safe + && record.rust.arguments == ["u32"] + && record.rust.result == "[f32; 4]" + && record.dialect.operands == ["i32"] + && record.dialect.results == ["f32", "f32", "f32", "f32"] + && llvm.arguments == ["tmem_ptr", "i1"] + && llvm.results == ["v4i32"] + && tcgen05.source_contract + == Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection + } + ( + Tcgen05Operation::CommitMulticast | Tcgen05Operation::CommitMulticastCg2, + Tcgen05Adapter::BarrierPointerMaskToVoid, + ) => { + !record.rust.safe + && record.rust.arguments == ["*mut u64", "u16"] + && record.rust.result == "()" + && record.dialect.operands == ["ptr", "i16"] + && llvm.arguments == ["shared_ptr", "i16"] + && tcgen05.source_contract == Tcgen05SourceContract::ExactTablegenSelection + } + ( + Tcgen05Operation::ShiftDown | Tcgen05Operation::ShiftDownCg2, + Tcgen05Adapter::TmemAddressToVoid, + ) => { + !record.rust.safe + && record.rust.arguments == ["u32"] + && record.rust.result == "()" + && record.dialect.operands == ["i32"] + && record.dialect.results.is_empty() + && llvm.arguments == ["tmem_ptr"] + && llvm.results.is_empty() + && tcgen05.source_contract == Tcgen05SourceContract::ExactTablegenSelection + } + ( + Tcgen05Operation::Ld, + Tcgen05Adapter::TmemInjectPack16ToU32Registers + | Tcgen05Adapter::TmemHalfSplitOffsetInjectPack16ToU32Registers, + ) => { + let count = tcgen05_ld_register_count(record); + let has_half_split_offset = + tcgen05.adapter == Tcgen05Adapter::TmemHalfSplitOffsetInjectPack16ToU32Registers; + let rust_result = if count == 1 { + "u32".into() + } else { + format!("[u32; {count}]") + }; + let llvm_result = if count == 1 { + "i32".into() + } else { + format!("v{count}i32") + }; + !record.rust.safe + && record.rust.arguments + == if has_half_split_offset { + vec!["u32", "i64"] + } else { + vec!["u32"] + } + && record.rust.result == rust_result + && record.dialect.operands + == if has_half_split_offset { + vec!["i32", "i64"] + } else { + vec!["i32"] + } + && record.dialect.results == vec!["i32"; count] + && llvm.arguments + == if has_half_split_offset { + vec!["tmem_ptr", "i64", "i1"] + } else { + vec!["tmem_ptr", "i1"] + } + && llvm.results == [llvm_result] + && !record.semantics.pure + && record.semantics.memory == "read" + && tcgen05.source_contract + == Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection + } + ( + Tcgen05Operation::St, + Tcgen05Adapter::TmemU32RegistersInjectUnpack16ToVoid + | Tcgen05Adapter::TmemHalfSplitOffsetU32RegistersInjectUnpack16ToVoid, + ) => { + let count = tcgen05_st_register_count(record); + let has_half_split_offset = tcgen05.adapter + == Tcgen05Adapter::TmemHalfSplitOffsetU32RegistersInjectUnpack16ToVoid; + let rust_data = if count == 1 { + "u32".into() + } else { + format!("[u32; {count}]") + }; + let llvm_data = if count == 1 { + "i32".into() + } else { + format!("v{count}i32") + }; + !record.rust.safe + && record.rust.arguments + == if has_half_split_offset { + vec!["u32", "i64", rust_data.as_str()] + } else { + vec!["u32", rust_data.as_str()] + } + && record.rust.result == "()" + && record.dialect.operands + == std::iter::once("i32") + .chain(has_half_split_offset.then_some("i64")) + .chain(std::iter::repeat_n("i32", count)) + .collect::>() + && record.dialect.results.is_empty() + && llvm.arguments + == if has_half_split_offset { + vec!["tmem_ptr", "i64", llvm_data.as_str(), "i1"] + } else { + vec!["tmem_ptr", llvm_data.as_str(), "i1"] + } + && llvm.results.is_empty() + && !record.semantics.pure + && record.semantics.memory == "write" + && tcgen05.source_contract + == Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection + } + _ => false, + } +} + +fn tcgen05_carrier_name(ty: &str) -> &'static str { + match ty { + "ptr" => "Ptr", + "i1" => "I1", + "i16" => "I16", + "i32" => "I32", + "i64" => "I64", + "f32" => "F32", + _ => panic!("unsupported tcgen05 dialect carrier {ty}"), + } +} + +pub(in crate::render) fn render_tcgen05_carrier_runs(types: &[String]) -> String { + let mut runs = Vec::<(&str, usize)>::new(); + for ty in types { + let carrier = tcgen05_carrier_name(ty); + if let Some((previous, count)) = runs.last_mut() + && *previous == carrier + { + *count += 1; + } else { + runs.push((carrier, 1)); + } + } + + let mut output = String::from("&["); + for (index, (carrier, count)) in runs.into_iter().enumerate() { + if index != 0 { + output.push_str(", "); + } + write!(output, "(Tcgen05Carrier::{carrier}, {count})").unwrap(); + } + output.push(']'); + output +} + +pub(in crate::render) fn tcgen05_inline_asm( + record: &CatalogIntrinsic, +) -> (String, String, Option) { + let operation = record.tcgen05.as_ref().expect("tcgen05 contract").operation; + match operation { + Tcgen05Operation::Alloc | Tcgen05Operation::AllocCg2 => { + let group = if operation == Tcgen05Operation::AllocCg2 { 2 } else { 1 }; + ( + format!("{{ .reg .u64 %shared64; .reg .u32 %shared32; cvta.to.shared.u64 %shared64, $0; cvt.u32.u64 %shared32, %shared64; tcgen05.alloc.cta_group::{group}.sync.aligned.shared::cta.b32 [%shared32], $1; }}"), + "l,r,~{memory}".into(), + None, + ) + } + Tcgen05Operation::Dealloc | Tcgen05Operation::DeallocCg2 => { + let group = if operation == Tcgen05Operation::DeallocCg2 { 2 } else { 1 }; + ( + format!("tcgen05.dealloc.cta_group::{group}.sync.aligned.b32 $0, $1;"), + "r,r,~{memory}".into(), + None, + ) + } + Tcgen05Operation::RelinquishAllocPermit + | Tcgen05Operation::RelinquishAllocPermitCg2 => { + let group = if operation == Tcgen05Operation::RelinquishAllocPermitCg2 { + 2 + } else { + 1 + }; + ( + format!("tcgen05.relinquish_alloc_permit.cta_group::{group}.sync.aligned;"), + "~{memory}".into(), + None, + ) + } + Tcgen05Operation::FenceBeforeThreadSync => ( + "tcgen05.fence::before_thread_sync;".into(), + "~{memory}".into(), + None, + ), + Tcgen05Operation::FenceAfterThreadSync => ( + "tcgen05.fence::after_thread_sync;".into(), + "~{memory}".into(), + None, + ), + Tcgen05Operation::Commit + | Tcgen05Operation::CommitSharedCluster + | Tcgen05Operation::CommitCg2 + | Tcgen05Operation::CommitSharedClusterCg2 => { + let group = if matches!( + operation, + Tcgen05Operation::CommitCg2 | Tcgen05Operation::CommitSharedClusterCg2 + ) { + 2 + } else { + 1 + }; + let shared = if matches!( + operation, + Tcgen05Operation::CommitSharedCluster + | Tcgen05Operation::CommitSharedClusterCg2 + ) { + ".shared::cluster" + } else { + "" + }; + ( + format!("tcgen05.commit.cta_group::{group}.mbarrier::arrive::one{shared}.b64 [$0];"), + "r,~{memory}".into(), + None, + ) + } + Tcgen05Operation::MmaWsF16 | Tcgen05Operation::MmaWsBf16 => ( + "{ .reg .pred %enable_pred; setp.ne.s32 %enable_pred, $5, 0; tcgen05.mma.ws.cta_group::1.kind::f16 [$0], [$1], $3, $4, %enable_pred; }".into(), + "r,r,l,l,r,r,~{memory}".into(), + None, + ), + Tcgen05Operation::MmaWsTf32 => ( + "{ .reg .pred %enable_pred; setp.ne.s32 %enable_pred, $5, 0; tcgen05.mma.ws.cta_group::1.kind::tf32 [$0], [$1], $3, $4, %enable_pred; }".into(), + "r,r,l,l,r,r,~{memory}".into(), + None, + ), + Tcgen05Operation::MmaF16 | Tcgen05Operation::MmaF16Cg2 => { + let group = if operation == Tcgen05Operation::MmaF16Cg2 { 2 } else { 1 }; + let zeros = std::iter::repeat_n("%z", if group == 2 { 8 } else { 4 }) + .collect::>() + .join(", "); + ( + format!("{{ .reg .pred %enable_pred; setp.ne.s32 %enable_pred, $4, 0; .reg .u32 %z; mov.u32 %z, 0; tcgen05.mma.cta_group::{group}.kind::f16 [$0], $1, $2, $3, {{{zeros}}}, %enable_pred; }}"), + "r,l,l,r,r,~{memory}".into(), + None, + ) + } + Tcgen05Operation::CpSmemToTmem | Tcgen05Operation::CpSmemToTmemCg2 => { + ( + format!( + "{}.{} [$0], $1;", + record.expected_ptx.mnemonic, + record.expected_ptx.modifiers.join(".") + ), + "r,l,~{memory}".into(), + None, + ) + } + Tcgen05Operation::Ld16x256bX8Pure | Tcgen05Operation::Ld16x256bPure => { + let count = if operation == Tcgen05Operation::Ld16x256bX8Pure { + 32 + } else { + 4 + }; + let registers = (0..count) + .map(|index| format!("${index}")) + .collect::>() + .join(","); + let constraints = std::iter::repeat_n("=f", count) + .chain(std::iter::once("r")) + .collect::>() + .join(","); + ( + format!( + "tcgen05.ld.sync.aligned.16x256b.{}.b32 {{{registers}}}, [${count}];", + if count == 32 { "x8" } else { "x1" } + ), + constraints, + Some(count), + ) + } + Tcgen05Operation::Ld => { + let count = tcgen05_ld_register_count(record); + let has_half_split_offset = record + .tcgen05 + .as_ref() + .and_then(|tcgen05| tcgen05.ld) + .is_some_and(|ld| ld.shape == Tcgen05LdShape::M16x32bx2); + let registers = (0..count) + .map(|index| format!("${index}")) + .collect::>() + .join(","); + let constraints = std::iter::repeat_n("=r", count) + .chain(if has_half_split_offset { + vec!["r", "n", "~{memory}"] + } else { + vec!["r", "~{memory}"] + }) + .collect::>() + .join(","); + ( + format!( + "{}.{} {{{registers}}}, [${count}]{};", + record.expected_ptx.mnemonic, + record.expected_ptx.modifiers.join("."), + if has_half_split_offset { + format!(", ${}", count + 1) + } else { + String::new() + } + ), + constraints, + Some(count), + ) + } + Tcgen05Operation::St => { + let count = tcgen05_st_register_count(record); + let has_half_split_offset = record + .tcgen05 + .as_ref() + .and_then(|tcgen05| tcgen05.st) + .is_some_and(|st| st.shape == Tcgen05LdShape::M16x32bx2); + let first_data = if has_half_split_offset { 2 } else { 1 }; + let registers = (first_data..first_data + count) + .map(|index| format!("${index}")) + .collect::>() + .join(","); + let constraints = std::iter::once("r") + .chain(has_half_split_offset.then_some("n")) + .chain(std::iter::repeat_n("r", count)) + .chain(std::iter::once("~{memory}")) + .collect::>() + .join(","); + ( + format!( + "{}.{} [$0], {}{{{registers}}};", + record.expected_ptx.mnemonic, + record.expected_ptx.modifiers.join("."), + if has_half_split_offset { + "$1, " + } else { + "" + } + ), + constraints, + None, + ) + } + Tcgen05Operation::LoadWait => ( + "tcgen05.wait::ld.sync.aligned;".into(), + "~{memory}".into(), + None, + ), + Tcgen05Operation::StoreWait => ( + "tcgen05.wait::st.sync.aligned;".into(), + "~{memory}".into(), + None, + ), + Tcgen05Operation::CommitMulticast | Tcgen05Operation::CommitMulticastCg2 => { + let group = if operation == Tcgen05Operation::CommitMulticastCg2 { + 2 + } else { + 1 + }; + ( + format!("tcgen05.commit.cta_group::{group}.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 [$0], $1;"), + "r,h,~{memory}".into(), + None, + ) + } + Tcgen05Operation::ShiftDown | Tcgen05Operation::ShiftDownCg2 => { + let group = if operation == Tcgen05Operation::ShiftDownCg2 { + 2 + } else { + 1 + }; + ( + format!("tcgen05.shift.cta_group::{group}.down [$0];"), + "r,~{memory}".into(), + None, + ) + } + Tcgen05Operation::Mma => unreachable!("generic MMA uses attribute-driven lowering"), + } +} diff --git a/crates/cuda-intrinsics-gen/src/render/importer.rs b/crates/cuda-intrinsics-gen/src/render/importer.rs new file mode 100644 index 0000000000..46f2d12616 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/render/importer.rs @@ -0,0 +1,2999 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + CatalogFile, CatalogIntrinsic, ClcOperation, ClusterBarrierMode, CpAsyncControlOperation, + CpAsyncMbarrierAdapter, CpAsyncSourceSize, DebugControlOperation, ExecutionControlOperation, + MbarrierBasicAdapter, MbarrierBasicOperation, MbarrierExtendedAdapter, PackedAtomicFormat, + PrmtMode, ReduxAdapter, RegisterMmaAdapter, Tcgen05LdShape, Tcgen05MmaForm, + Tcgen05MmaSelectorLayout, Tcgen05Operation, TmaAdapter, TmaOperation, VoteAdapter, + WarpBarrierAdapter, WarpShuffleAdapter, WgmmaControlMode, +}; +use crate::render::common::{intrinsic_marker, rust_header}; +use crate::render::families::{ + active_masks, clc_intrinsics, cluster_barrier_attr, cluster_barriers, cluster_memory, + cp_async_controls, cp_async_copies, cp_async_mbarriers, debug_controls, dot_products, + elect_intrinsics, execution_control_family, execution_controls, extended_minmax, + extended_minmax_format_attr, extended_minmax_nan_attr, extended_minmax_operation_attr, + extended_minmax_subnormal_attr, extended_minmax_xorsign_abs_attr, integer_minmaxes, ldmatrix, + ldmatrix_attr_variants, mbarrier_basics, mbarrier_extended, movmatrix, packed_alus, + packed_atomics, packed_conversions, prmts, redux, register_mma_attr_variants, register_mmas, + scalar_arithmetic_arity, scalar_arithmetic_format_attr, scalar_arithmetic_operation_attr, + scalar_arithmetic_rounding_attr, scalar_arithmetic_saturation_attr, + scalar_arithmetic_subnormal_attr, scalar_arithmetics, scalar_conversion_rounding_attr, + scalar_conversion_saturation_attr, scalar_conversions, scalar_math_format_attr, + scalar_math_operation_attr, scalar_math_precision_attr, scalar_math_subnormal_attr, + scalar_maths, sparse_mma_attr_variants, sparse_mma_import_adapter, sparse_mma_selector_error, + sparse_mmas, sregs, stmatrices, stmatrix_compatibility_name, stmatrix_variant, sync_intrinsics, + tcgen05_intrinsics, tcgen05_ld_register_count, tcgen05_mma_b_usage_attr, tcgen05_mma_form_attr, + tcgen05_mma_intrinsics, tcgen05_mma_kind_attr, tcgen05_non_mma_intrinsics, + tcgen05_st_register_count, tma_intrinsics, vote_intrinsics, warp_barriers, warp_matches, + warp_shuffles, wgmma_controls, +}; +use crate::render::reference::{ + render_compiler_path_patterns, render_inline_patterns, render_string_patterns, +}; +use std::fmt::Write as _; + +fn render_importer_pure_value_dispatch( + output: &mut String, + catalog: &CatalogFile, + record: &CatalogIntrinsic, +) { + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(output, &path_refs); + output.push_str(" => {\n"); + writeln!( + output, + " require_arity(name, args.len(), {}, &loc)?;", + record.rust.arguments.len() + ) + .unwrap(); + for index in 0..record.rust.arguments.len() { + let previous = if index == 0 { "prev_op" } else { "last_op" }; + writeln!( + output, + " let (arg{index}, last_op) = rvalue::translate_operand(\n ctx, body, &args[{index}], value_map, block_ptr, {previous}, loc.clone(),\n )?;" + ) + .unwrap(); + } + let arguments = (0..record.rust.arguments.len()) + .map(|index| format!("arg{index}")) + .collect::>() + .join(", "); + writeln!( + output, + " let intrinsic = {}::build(ctx, {arguments});", + record.dialect.op_type + ) + .unwrap(); + output.push_str(" intrinsic.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, intrinsic, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, intrinsic, block_ptr, last_op);\n let result = intrinsic.deref(ctx).get_result(0);\n", + ); + writeln!( + output, + " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, intrinsic, value_map, block_map, loc,\n {:?},\n )?))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n"); +} + +fn render_importer_scalar_conversion_dispatch( + output: &mut String, + catalog: &CatalogFile, + record: &CatalogIntrinsic, +) { + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(output, &path_refs); + output.push_str(" => {\n"); + output.push_str( + " require_arity(name, args.len(), 1, &loc)?;\n\ + let (value, last_op) = rvalue::translate_operand(\n\ + ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n\ + )?;\n", + ); + writeln!( + output, + " let intrinsic = ScalarConversionOp::build(ctx, value, {}, {});", + scalar_conversion_rounding_attr(record), + scalar_conversion_saturation_attr(record), + ) + .unwrap(); + output.push_str(" intrinsic.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, intrinsic, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, intrinsic, block_ptr, last_op);\n\ + let result = intrinsic.deref(ctx).get_result(0);\n", + ); + writeln!( + output, + " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, intrinsic, value_map, block_map, loc,\n {:?},\n )?))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n"); +} + +fn render_importer_scalar_arithmetic_dispatch( + output: &mut String, + catalog: &CatalogFile, + record: &CatalogIntrinsic, +) { + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(output, &path_refs); + output.push_str(" => {\n"); + let arity = scalar_arithmetic_arity(record); + writeln!( + output, + " require_arity(name, args.len(), {arity}, &loc)?;" + ) + .unwrap(); + for index in 0..arity { + let previous = if index == 0 { "prev_op" } else { "last_op" }; + writeln!( + output, + " let (arg{index}, last_op) = rvalue::translate_operand(\n ctx, body, &args[{index}], value_map, block_ptr, {previous}, loc.clone(),\n )?;" + ) + .unwrap(); + } + let arguments = (0..arity) + .map(|index| format!("arg{index}")) + .collect::>() + .join(", "); + writeln!( + output, + " let intrinsic = ScalarArithmeticOp::build(ctx, vec![{arguments}], {}, {}, {}, {}, {});", + scalar_arithmetic_format_attr(record), + scalar_arithmetic_operation_attr(record), + scalar_arithmetic_rounding_attr(record), + scalar_arithmetic_subnormal_attr(record), + scalar_arithmetic_saturation_attr(record), + ) + .unwrap(); + output.push_str(" intrinsic.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, intrinsic, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, intrinsic, block_ptr, last_op);\n\ + let result = intrinsic.deref(ctx).get_result(0);\n", + ); + writeln!( + output, + " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, intrinsic, value_map, block_map, loc,\n {:?},\n )?))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n"); +} + +fn render_importer_scalar_math_dispatch( + output: &mut String, + catalog: &CatalogFile, + record: &CatalogIntrinsic, +) { + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(output, &path_refs); + output.push_str(" => {\n"); + writeln!( + output, + " require_arity(name, args.len(), 1, &loc)?;" + ) + .unwrap(); + writeln!( + output, + " let (arg0, last_op) = rvalue::translate_operand(\n ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n )?;" + ) + .unwrap(); + writeln!( + output, + " let intrinsic = ScalarMathOp::build(ctx, arg0, {}, {}, {}, {});", + scalar_math_format_attr(record), + scalar_math_operation_attr(record), + scalar_math_precision_attr(record), + scalar_math_subnormal_attr(record), + ) + .unwrap(); + output.push_str(" intrinsic.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, intrinsic, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, intrinsic, block_ptr, last_op);\n\ + let result = intrinsic.deref(ctx).get_result(0);\n", + ); + writeln!( + output, + " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, intrinsic, value_map, block_map, loc,\n {:?},\n )?))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n"); +} + +fn render_importer_extended_minmax_dispatch( + output: &mut String, + catalog: &CatalogFile, + record: &CatalogIntrinsic, +) { + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(output, &path_refs); + output.push_str(" => {\n"); + output.push_str( + " require_arity(name, args.len(), 2, &loc)?;\n\ + let (a, last_op) = rvalue::translate_operand(\n\ + ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n\ + )?;\n\ + let (b, last_op) = rvalue::translate_operand(\n\ + ctx, body, &args[1], value_map, block_ptr, last_op, loc.clone(),\n\ + )?;\n", + ); + writeln!( + output, + " let intrinsic = ExtendedMinMaxOp::build(ctx, a, b, {}, {}, {}, {}, {});", + extended_minmax_format_attr(record), + extended_minmax_operation_attr(record), + extended_minmax_subnormal_attr(record), + extended_minmax_nan_attr(record), + extended_minmax_xorsign_abs_attr(record), + ) + .unwrap(); + output.push_str(" intrinsic.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, intrinsic, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, intrinsic, block_ptr, last_op);\n\ + let result = intrinsic.deref(ctx).get_result(0);\n", + ); + writeln!( + output, + " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, intrinsic, value_map, block_map, loc,\n {:?},\n )?))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n"); +} + +fn render_importer_elect_dispatch( + output: &mut String, + catalog: &CatalogFile, + record: &CatalogIntrinsic, +) { + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(output, &path_refs); + output.push_str(" => {\n"); + output.push_str( + " require_arity(name, args.len(), 1, &loc)?;\n\ + \n\ + let tuple_ty = crate::translator::types::translate_destination_type(\n\ + ctx, body, destination, &loc,\n\ + )?;\n\ + let (leader_ty, elected_ty) = {\n\ + let ty = tuple_ty.deref(ctx);\n\ + match ty.downcast_ref::() {\n\ + Some(tuple) if tuple.get_types().len() == 2 => {\n\ + (tuple.get_types()[0], tuple.get_types()[1])\n\ + }\n\ + _ => {\n\ + return input_err!(\n\ + loc.clone(),\n\ + TranslationErr::unsupported(\n\ + \"warp::elect_sync destination must be a (u32, bool) tuple\"\n\ + .to_owned()\n\ + )\n\ + );\n\ + }\n\ + }\n\ + };\n\ + let (mask, last_op) = rvalue::translate_operand(\n\ + ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n\ + )?;\n", + ); + writeln!( + output, + " let elect = Operation::new(\n ctx, {}::get_concrete_op_info(), vec![leader_ty, elected_ty],\n vec![mask], vec![], 0,\n );", + record.dialect.op_type + ) + .unwrap(); + output.push_str(" elect.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, elect, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, elect, block_ptr, last_op);\n\ + let leader = elect.deref(ctx).get_result(0);\n\ + let elected = elect.deref(ctx).get_result(1);\n\ + let tuple = Operation::new(\n\ + ctx, MirConstructTupleOp::get_concrete_op_info(), vec![tuple_ty],\n\ + vec![leader, elected],\n\ + vec![], 0,\n\ + );\n\ + tuple.deref_mut(ctx).set_loc(loc.clone());\n\ + tuple.insert_after(ctx, elect);\n\ + let result = tuple.deref(ctx).get_result(0);\n", + ); + writeln!( + output, + " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, tuple, value_map, block_map, loc,\n {:?},\n )?))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n"); +} + +fn render_importer_tcgen05_mma_dispatch( + output: &mut String, + catalog: &CatalogFile, + record: &CatalogIntrinsic, +) { + let mma = record + .tcgen05 + .as_ref() + .and_then(|tcgen05| tcgen05.mma.as_ref()) + .expect("tcgen05 MMA record"); + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(output, &path_refs); + output.push_str(" => {\n"); + writeln!( + output, + " require_arity(name, args.len(), {}, &loc)?;", + record.rust.arguments.len() + ) + .unwrap(); + + let runtime_indices = if mma.alias.is_some() && mma.form == Tcgen05MmaForm::WsTensor { + vec![0, 1, 3, 4, 5] + } else if mma.alias.is_some() { + (0..record.rust.arguments.len()).collect() + } else { + let selector_indices = match mma.selector_layout { + Tcgen05MmaSelectorLayout::Base { + kind_argument, + cta_group_argument, + collector_a_argument, + .. + } => [ + kind_argument as usize, + cta_group_argument as usize, + collector_a_argument as usize, + ], + Tcgen05MmaSelectorLayout::WarpSpecialized { + kind_argument, + b_buffer_argument, + b_usage_argument, + } => [ + kind_argument as usize, + b_buffer_argument as usize, + b_usage_argument as usize, + ], + }; + (0..record.rust.arguments.len()) + .filter(|index| !selector_indices.contains(index)) + .collect() + }; + writeln!( + output, + " let runtime_indices: &[usize] = &{runtime_indices:?};" + ) + .unwrap(); + output.push_str( + " let mut last_op = prev_op;\n let mut operands = Vec::with_capacity(runtime_indices.len());\n for &index in runtime_indices {\n let (value, translated) = rvalue::translate_operand(\n ctx, body, &args[index], value_map, block_ptr, last_op, loc.clone(),\n )?;\n last_op = translated;\n operands.push(value);\n }\n", + ); + + if let Some(fixed) = mma.fixed_selectors { + writeln!( + output, + " let kind = {};", + tcgen05_mma_kind_attr(fixed.kind) + ) + .unwrap(); + match mma.selector_layout { + Tcgen05MmaSelectorLayout::Base { .. } => { + output.push_str( + " let cta_group = Tcgen05MmaCtaGroupAttr::Cg1;\n let collector_a = Tcgen05MmaCollectorAAttr::Discard;\n", + ); + } + Tcgen05MmaSelectorLayout::WarpSpecialized { .. } => { + let b_buffer = match fixed.b_buffer { + 0 => "Tcgen05MmaBBufferAttr::B0", + 1 => "Tcgen05MmaBBufferAttr::B1", + 2 => "Tcgen05MmaBBufferAttr::B2", + 3 => "Tcgen05MmaBBufferAttr::B3", + _ => panic!("admitted tcgen05 MMA B buffer"), + }; + writeln!(output, " let b_buffer = {b_buffer};").unwrap(); + writeln!( + output, + " let b_usage = {};", + tcgen05_mma_b_usage_attr(fixed.b_usage) + ) + .unwrap(); + } + } + } else { + let (kind_argument, second_argument, third_argument) = match mma.selector_layout { + Tcgen05MmaSelectorLayout::Base { + kind_argument, + cta_group_argument, + collector_a_argument, + .. + } => (kind_argument, cta_group_argument, collector_a_argument), + Tcgen05MmaSelectorLayout::WarpSpecialized { + kind_argument, + b_buffer_argument, + b_usage_argument, + } => (kind_argument, b_buffer_argument, b_usage_argument), + }; + writeln!( + output, + " let kind = match generated_tcgen05_mma_selector(args.get({kind_argument}usize)) {{\n Some(0) => Tcgen05MmaKindAttr::F16,\n Some(1) => Tcgen05MmaKindAttr::Tf32,\n Some(2) => Tcgen05MmaKindAttr::F8f6f4,\n Some(3) => Tcgen05MmaKindAttr::I8,\n _ => return input_err!(loc.clone(), TranslationErr::unsupported(\"tcgen05 MMA kind must be a compile-time u32 constant in 0..=3\".to_owned())),\n }};" + ) + .unwrap(); + match mma.selector_layout { + Tcgen05MmaSelectorLayout::Base { + collector_a_upper_exclusive, + .. + } => { + writeln!( + output, + " let cta_group = match generated_tcgen05_mma_selector(args.get({second_argument}usize)) {{\n Some(1) => Tcgen05MmaCtaGroupAttr::Cg1,\n Some(2) => Tcgen05MmaCtaGroupAttr::Cg2,\n _ => {{\n return input_err!(\n loc.clone(),\n TranslationErr::unsupported(\n \"tcgen05 MMA CTA group must be the compile-time u32 constant 1 or 2\"\n .to_owned()\n )\n );\n }}\n }};" + ) + .unwrap(); + let collector_arms = if collector_a_upper_exclusive == 2 { + " Some(0) => Tcgen05MmaCollectorAAttr::Discard,\n Some(1) => Tcgen05MmaCollectorAAttr::LastUse,\n" + } else { + " Some(0) => Tcgen05MmaCollectorAAttr::Discard,\n Some(1) => Tcgen05MmaCollectorAAttr::LastUse,\n Some(2) => Tcgen05MmaCollectorAAttr::Fill,\n Some(3) => Tcgen05MmaCollectorAAttr::Use,\n" + }; + writeln!( + output, + " let collector_a = match generated_tcgen05_mma_selector(args.get({third_argument}usize)) {{\n{collector_arms} _ => return input_err!(loc.clone(), TranslationErr::unsupported(\"tcgen05 MMA collector-A selector is outside its closed range\".to_owned())),\n }};" + ) + .unwrap(); + } + Tcgen05MmaSelectorLayout::WarpSpecialized { .. } => { + writeln!( + output, + " let b_buffer = match generated_tcgen05_mma_selector(args.get({second_argument}usize)) {{\n Some(0) => Tcgen05MmaBBufferAttr::B0,\n Some(1) => Tcgen05MmaBBufferAttr::B1,\n Some(2) => Tcgen05MmaBBufferAttr::B2,\n Some(3) => Tcgen05MmaBBufferAttr::B3,\n _ => return input_err!(loc.clone(), TranslationErr::unsupported(\"tcgen05 MMA B buffer must be a compile-time u32 constant in 0..=3\".to_owned())),\n }};" + ) + .unwrap(); + writeln!( + output, + " let b_usage = match generated_tcgen05_mma_selector(args.get({third_argument}usize)) {{\n Some(0) => Tcgen05MmaBUsageAttr::Discard,\n Some(1) => Tcgen05MmaBUsageAttr::LastUse,\n Some(2) => Tcgen05MmaBUsageAttr::Fill,\n Some(3) => Tcgen05MmaBUsageAttr::Use,\n _ => return input_err!(loc.clone(), TranslationErr::unsupported(\"tcgen05 MMA B usage must be a compile-time u32 constant in 0..=3\".to_owned())),\n }};" + ) + .unwrap(); + } + } + } + + output.push_str( + " let intrinsic = Operation::new(ctx, Tcgen05MmaOp::get_concrete_op_info(), vec![], operands, vec![], 0);\n intrinsic.deref_mut(ctx).set_loc(loc.clone());\n let mma_op = Tcgen05MmaOp::new(intrinsic);\n", + ); + writeln!( + output, + " mma_op.set_attr_nvvm_tcgen05_mma_form(ctx, {});", + tcgen05_mma_form_attr(mma.form) + ) + .unwrap(); + output.push_str(" mma_op.set_attr_nvvm_tcgen05_mma_kind(ctx, kind);\n"); + match mma.selector_layout { + Tcgen05MmaSelectorLayout::Base { .. } => output.push_str( + " mma_op.set_attr_nvvm_tcgen05_mma_cta_group(ctx, cta_group);\n mma_op.set_attr_nvvm_tcgen05_mma_collector_a(ctx, collector_a);\n", + ), + Tcgen05MmaSelectorLayout::WarpSpecialized { .. } => output.push_str( + " mma_op.set_attr_nvvm_tcgen05_mma_b_buffer(ctx, b_buffer);\n mma_op.set_attr_nvvm_tcgen05_mma_b_usage(ctx, b_usage);\n", + ), + } + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, intrinsic, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, intrinsic, block_ptr, last_op);\n if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, intrinsic, block_map, loc)))\n } else {\n", + ); + writeln!( + output, + " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n }\n"); +} + +pub(super) fn render_importer(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + "//! Generated raw/compatibility path dispatch for CUDA intrinsics.\n\nuse crate::error::{TranslationErr, TranslationResult};\nuse crate::translator::{rvalue, terminator::helpers};\nuse crate::translator::values::ValueMap;\nuse dialect_nvvm::ops::{", + ); + for (index, record) in sregs(catalog).enumerate() { + if index != 0 { + output.push_str(", "); + } + output.push_str(&record.dialect.op_type); + } + if ldmatrix(catalog).next().is_some() { + if sregs(catalog).next().is_some() { + output.push_str(", "); + } + output.push_str("LdmatrixElementAttr, LdmatrixLayoutAttr, LdmatrixMultiplicityAttr, LdmatrixOp, LdmatrixShapeAttr, LdmatrixStateSpaceAttr"); + } + if let Some(record) = movmatrix(catalog).next() { + output.push_str(", "); + output.push_str(&record.dialect.op_type); + } + if register_mmas(catalog).next().is_some() { + output.push_str(", RegisterMmaAccumulatorAttr, RegisterMmaElementAttr, RegisterMmaKindAttr, RegisterMmaLayoutAttr, RegisterMmaOp, RegisterMmaOperationAttr, RegisterMmaOverflowAttr, RegisterMmaShapeAttr"); + } + if sparse_mmas(catalog).next().is_some() { + output.push_str(", SparseMmaAccumulatorAttr, SparseMmaElementAttr, SparseMmaLayoutAttr, SparseMmaMetadataAttr, SparseMmaOp, SparseMmaOverflowAttr, SparseMmaSelectorAttr, SparseMmaShapeAttr"); + } + if prmts(catalog).next().is_some() { + output.push_str(", PrmtModeAttr, PrmtOp"); + } + if scalar_conversions(catalog).next().is_some() { + output.push_str( + ", ScalarConversionOp, ScalarConversionRoundingAttr, ScalarConversionSaturationAttr", + ); + } + if scalar_arithmetics(catalog).next().is_some() { + output.push_str( + ", ScalarArithmeticFormatAttr, ScalarArithmeticOp, ScalarArithmeticOperationAttr, ScalarArithmeticRoundingAttr, ScalarArithmeticSaturationAttr, ScalarArithmeticSubnormalAttr", + ); + } + if scalar_maths(catalog).next().is_some() { + output.push_str( + ", ScalarMathFormatAttr, ScalarMathOp, ScalarMathOperationAttr, ScalarMathPrecisionAttr, ScalarMathSubnormalAttr", + ); + } + if extended_minmax(catalog).next().is_some() { + output.push_str( + ", ExtendedMinMaxFormatAttr, ExtendedMinMaxNanAttr, ExtendedMinMaxOp, ExtendedMinMaxOperationAttr, ExtendedMinMaxSubnormalAttr, ExtendedMinMaxXorSignAbsAttr", + ); + } + if cluster_barriers(catalog).next().is_some() { + output.push_str(", ClusterBarrierModeAttr, ClusterBarrierOp"); + } + if debug_controls(catalog).next().is_some() { + output.push_str(", BreakpointOp, PmEventOp, TrapOp"); + } + for record in cluster_memory(catalog) { + output.push_str(", "); + output.push_str(&record.dialect.op_type); + } + for record in clc_intrinsics(catalog) { + output.push_str(", "); + output.push_str(&record.dialect.op_type); + } + for record in execution_controls(catalog) { + output.push_str(", "); + output.push_str(&record.dialect.op_type); + } + if wgmma_controls(catalog).next().is_some() { + output.push_str( + ", WgmmaCommitGroupSyncAlignedOp, WgmmaFenceSyncAlignedOp, WgmmaWaitGroupSyncAlignedOp", + ); + } + if packed_atomics(catalog).next().is_some() { + if sregs(catalog).next().is_some() || ldmatrix(catalog).next().is_some() { + output.push_str(", "); + } + output.push_str("PackedAtomicAddOp, PackedAtomicFormatAttr"); + } + if redux(catalog).next().is_some() { + if sregs(catalog).next().is_some() + || ldmatrix(catalog).next().is_some() + || packed_atomics(catalog).next().is_some() + { + output.push_str(", "); + } + for (index, record) in redux(catalog).enumerate() { + if index != 0 { + output.push_str(", "); + } + output.push_str(&record.dialect.op_type); + } + } + if vote_intrinsics(catalog).next().is_some() { + if sregs(catalog).next().is_some() + || ldmatrix(catalog).next().is_some() + || packed_atomics(catalog).next().is_some() + || redux(catalog).next().is_some() + { + output.push_str(", "); + } + for (index, record) in vote_intrinsics(catalog).enumerate() { + if index != 0 { + output.push_str(", "); + } + output.push_str(&record.dialect.op_type); + } + } + if active_masks(catalog).next().is_some() { + output.push_str(", "); + for (index, record) in active_masks(catalog).enumerate() { + if index != 0 { + output.push_str(", "); + } + output.push_str(&record.dialect.op_type); + } + } + if warp_matches(catalog).next().is_some() { + output.push_str(", "); + for (index, record) in warp_matches(catalog).enumerate() { + if index != 0 { + output.push_str(", "); + } + output.push_str(&record.dialect.op_type); + } + } + for record in elect_intrinsics(catalog) { + output.push_str(", "); + output.push_str(&record.dialect.op_type); + } + if warp_barriers(catalog).next().is_some() { + output.push_str(", "); + for (index, record) in warp_barriers(catalog).enumerate() { + if index != 0 { + output.push_str(", "); + } + output.push_str(&record.dialect.op_type); + } + } + if warp_shuffles(catalog).next().is_some() { + output.push_str(", "); + for (index, record) in warp_shuffles(catalog).enumerate() { + if index != 0 { + output.push_str(", "); + } + output.push_str(&record.dialect.op_type); + } + } + if packed_alus(catalog).next().is_some() { + if sregs(catalog).next().is_some() + || ldmatrix(catalog).next().is_some() + || packed_atomics(catalog).next().is_some() + || redux(catalog).next().is_some() + || vote_intrinsics(catalog).next().is_some() + || active_masks(catalog).next().is_some() + || warp_matches(catalog).next().is_some() + || warp_barriers(catalog).next().is_some() + || warp_shuffles(catalog).next().is_some() + { + output.push_str(", "); + } + for (index, record) in packed_alus(catalog).enumerate() { + if index != 0 { + output.push_str(", "); + } + output.push_str(&record.dialect.op_type); + } + } + if integer_minmaxes(catalog).next().is_some() { + // packed_alu records always exist, so a separator is always needed. + output.push_str(", "); + for (index, record) in integer_minmaxes(catalog).enumerate() { + if index != 0 { + output.push_str(", "); + } + output.push_str(&record.dialect.op_type); + } + } + if packed_conversions(catalog).next().is_some() { + if sregs(catalog).next().is_some() + || ldmatrix(catalog).next().is_some() + || packed_atomics(catalog).next().is_some() + || redux(catalog).next().is_some() + || vote_intrinsics(catalog).next().is_some() + || active_masks(catalog).next().is_some() + || warp_matches(catalog).next().is_some() + || warp_barriers(catalog).next().is_some() + || warp_shuffles(catalog).next().is_some() + || packed_alus(catalog).next().is_some() + { + output.push_str(", "); + } + for (index, record) in packed_conversions(catalog).enumerate() { + if index != 0 { + output.push_str(", "); + } + output.push_str(&record.dialect.op_type); + } + } + if cp_async_copies(catalog).next().is_some() { + output.push_str(", "); + for (index, record) in cp_async_copies(catalog).enumerate() { + if index != 0 { + output.push_str(", "); + } + output.push_str(&record.dialect.op_type); + } + } + if cp_async_controls(catalog).next().is_some() { + output.push_str(", "); + for (index, record) in cp_async_controls(catalog).enumerate() { + if index != 0 { + output.push_str(", "); + } + output.push_str(&record.dialect.op_type); + } + } + if cp_async_mbarriers(catalog).next().is_some() { + output.push_str(", "); + for (index, record) in cp_async_mbarriers(catalog).enumerate() { + if index != 0 { + output.push_str(", "); + } + output.push_str(&record.dialect.op_type); + } + } + for record in mbarrier_basics(catalog) { + output.push_str(", "); + output.push_str(&record.dialect.op_type); + } + for record in mbarrier_extended(catalog) { + output.push_str(", "); + output.push_str(&record.dialect.op_type); + } + if dot_products(catalog).next().is_some() { + if sregs(catalog).next().is_some() + || ldmatrix(catalog).next().is_some() + || packed_atomics(catalog).next().is_some() + || redux(catalog).next().is_some() + || vote_intrinsics(catalog).next().is_some() + || packed_alus(catalog).next().is_some() + || packed_conversions(catalog).next().is_some() + { + output.push_str(", "); + } + for (index, record) in dot_products(catalog).enumerate() { + if index != 0 { + output.push_str(", "); + } + output.push_str(&record.dialect.op_type); + } + } + if sync_intrinsics(catalog).next().is_some() { + if sregs(catalog).next().is_some() + || ldmatrix(catalog).next().is_some() + || packed_atomics(catalog).next().is_some() + || redux(catalog).next().is_some() + || vote_intrinsics(catalog).next().is_some() + || dot_products(catalog).next().is_some() + || packed_alus(catalog).next().is_some() + || packed_conversions(catalog).next().is_some() + { + output.push_str(", "); + } + for (index, record) in sync_intrinsics(catalog).enumerate() { + if index != 0 { + output.push_str(", "); + } + output.push_str(&record.dialect.op_type); + } + } + for record in stmatrices(catalog) { + output.push_str(", "); + output.push_str(&record.dialect.op_type); + } + for record in tcgen05_non_mma_intrinsics(catalog) { + output.push_str(", "); + output.push_str(&record.dialect.op_type); + } + if tcgen05_mma_intrinsics(catalog).next().is_some() { + output.push_str(", Tcgen05MmaBBufferAttr, Tcgen05MmaBUsageAttr, Tcgen05MmaCollectorAAttr, Tcgen05MmaCtaGroupAttr, Tcgen05MmaFormAttr, Tcgen05MmaKindAttr, Tcgen05MmaOp"); + } + output.push_str( + "};\nuse pliron::basic_block::BasicBlock;\nuse pliron::context::{Context, Ptr};\nuse pliron::input_err;\nuse pliron::location::{Located, Location};\nuse pliron::op::Op;\nuse pliron::operation::Operation;\nuse rustc_public::{CrateDef, mir, ty::FnDef};\n\n", + ); + if debug_controls(catalog).next().is_some() + || execution_control_family(catalog, "register_control") + .next() + .is_some() + { + output.push_str("use dialect_mir::ops::{MirConstantOp, MirUnreachableOp};\n\n"); + } else if wgmma_controls(catalog).next().is_some() { + output.push_str("use dialect_mir::ops::MirConstantOp;\n\n"); + } + if register_mmas(catalog).next().is_some() + || sparse_mmas(catalog).next().is_some() + || tcgen05_intrinsics(catalog).next().is_some() + { + output.push_str( + "use dialect_mir::{attributes::FieldIndexAttr, ops::{MirConstructArrayOp, MirConstructStructOp, MirExtractFieldOp}, types::{MirArrayType, MirStructType}};\nuse pliron::{builtin::types::{FP32Type, FP64Type, IntegerType, Signedness}, r#type::{TypeHandle, Typed}, value::Value};\n\n", + ); + } + if tcgen05_intrinsics(catalog).next().is_some() { + output = output.replace( + "use crate::translator::{rvalue, terminator::helpers};", + "use crate::translator::{rvalue, terminator::helpers, types};", + ); + } + if elect_intrinsics(catalog).next().is_some() { + output.push_str("use dialect_mir::{ops::MirConstructTupleOp, types::MirTupleType};\n\n"); + } + writeln!( + output, + "pub const GENERATED_INTRINSIC_ABI: u32 = {};", + catalog.intrinsic_abi + ) + .unwrap(); + writeln!( + output, + "pub const GENERATED_INTRINSIC_ABI_NAMESPACE: &str = {:?};\n", + format!("__cuda_oxide_intrinsic_abi_v{}", catalog.intrinsic_abi) + ) + .unwrap(); + output.push_str( + "#[derive(Debug, Clone, PartialEq, Eq)]\npub enum RawIntrinsicIdentity {\n NotRawCrate,\n Known(String),\n UnsupportedAbi(String),\n UnknownId(String),\n}\n\n/// Classify the unrewritten raw intrinsic DefPath. rustc's ordinary name\n/// printer may prefer a public re-export path, which is source API rather than ABI.\npub fn classify_raw_intrinsic(fn_def: FnDef) -> RawIntrinsicIdentity {\n let crate_name = fn_def.krate().name.to_string();\n if !matches!(crate_name.as_str(), \"cuda_intrinsics\" | \"cuda-intrinsics\") {\n return RawIntrinsicIdentity::NotRawCrate;\n }\n\n let mut segments = Vec::new();\n let mut current = Some(fn_def.def_id());\n while let Some(def_id) = current {\n let printed = def_id.name();\n let segment = printed.as_str().rsplit(\"::\").next().unwrap_or_default();\n if segment != crate_name {\n segments.push(segment.to_owned());\n }\n current = def_id.parent();\n }\n segments.reverse();\n classify_raw_intrinsic_path(&crate_name, format!(\"{crate_name}::{}\", segments.join(\"::\")))\n}\n\nfn classify_raw_intrinsic_path(crate_name: &str, path: String) -> RawIntrinsicIdentity {\n if !matches!(crate_name, \"cuda_intrinsics\" | \"cuda-intrinsics\") {\n return RawIntrinsicIdentity::NotRawCrate;\n }\n if is_raw_generated_intrinsic_path(&path) {\n return RawIntrinsicIdentity::Known(path);\n }\n let namespace = path.split(\"::\").nth(1).unwrap_or_default();\n if namespace.starts_with(\"__cuda_oxide_intrinsic_abi_v\")\n && namespace != GENERATED_INTRINSIC_ABI_NAMESPACE\n {\n RawIntrinsicIdentity::UnsupportedAbi(path)\n } else {\n RawIntrinsicIdentity::UnknownId(path)\n }\n}\n\npub fn require_supported_raw_intrinsic(\n fn_def: FnDef,\n loc: &Location,\n) -> TranslationResult> {\n match classify_raw_intrinsic(fn_def) {\n RawIntrinsicIdentity::NotRawCrate => Ok(None),\n RawIntrinsicIdentity::Known(path) => Ok(Some(path)),\n RawIntrinsicIdentity::UnsupportedAbi(path) => input_err!(\n loc.clone(),\n TranslationErr::unsupported(format!(\n \"cuda-intrinsics ABI mismatch: `{path}` uses an unsupported intrinsic ABI; this compiler supports ABI v{GENERATED_INTRINSIC_ABI}\"\n ))\n ),\n RawIntrinsicIdentity::UnknownId(path) => input_err!(\n loc.clone(),\n TranslationErr::unsupported(format!(\n \"cuda-intrinsics ABI mismatch: `{path}` is not a known intrinsic ID in ABI v{GENERATED_INTRINSIC_ABI}\"\n ))\n ),\n }\n}\n\n", + ); + output + .push_str("pub fn is_generated_intrinsic_path(name: &str) -> bool {\n matches!(name,\n"); + render_compiler_path_patterns(&mut output, catalog, " "); + output.push_str(" )\n}\n\npub fn is_raw_generated_intrinsic_path(name: &str) -> bool {\n matches!(name,\n"); + let raw_paths: Vec<_> = catalog + .intrinsics + .iter() + .map(|record| record.rust.canonical_path.as_str()) + .collect(); + render_string_patterns(&mut output, &raw_paths, " "); + // Test-only. The dispatch arms rendered below carry each marker as a + // literal, and the op-name-keyed lookup the compiler actually calls lives + // in `cuda-oxide-codegen`; nothing outside this file's own generated tests + // reads this path-keyed table. Gate it rather than leave a 2000-line match + // sitting unreferenced in the compiler. + output.push_str(" )\n}\n\n#[cfg(test)]\npub fn generated_intrinsic_marker(name: &str) -> Option<&'static str> {\n match name {\n"); + for record in &catalog.intrinsics { + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(&mut output, &path_refs); + writeln!(output, " => Some({:?}),", intrinsic_marker(catalog, record)).unwrap(); + } + output.push_str(" _ => None,\n }\n}\n\n"); + output.push_str( + "#[allow(clippy::too_many_arguments)]\npub fn try_dispatch_generated_intrinsic(\n ctx: &mut Context,\n body: &mir::Body,\n name: &str,\n args: &[mir::Operand],\n destination: &mir::Place,\n target: &Option,\n block_ptr: Ptr,\n prev_op: Option>,\n value_map: &mut ValueMap,\n block_map: &[Ptr],\n loc: Location,\n) -> TranslationResult>> {\n match name {\n", + ); + for record in sregs(catalog) { + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(&mut output, &path_refs); + output.push_str(" => {\n"); + writeln!( + output, + " require_arity(name, args.len(), 0, &loc)?;" + ) + .unwrap(); + let helper = if record.scalar_width() == Some(64) { + "emit_generated_nvvm_intrinsic_u64" + } else { + "emit_generated_nvvm_intrinsic" + }; + writeln!(output, " Ok(Some(helpers::{helper}(").unwrap(); + writeln!( + output, + " ctx, {}::get_concrete_op_info(), {:?}, destination, target, block_ptr,", + record.dialect.op_type, + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " prev_op, value_map, block_map, loc,\n )?))\n }\n", + ); + } + for record in active_masks(catalog) { + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(&mut output, &path_refs); + output.push_str(" => {\n"); + output.push_str(" require_arity(name, args.len(), 0, &loc)?;\n"); + writeln!( + output, + " let active_mask = {}::build(ctx);", + record.dialect.op_type + ) + .unwrap(); + output.push_str(" active_mask.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, active_mask, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, active_mask, block_ptr, prev_op);\n let result = active_mask.deref(ctx).get_result(0);\n", + ); + writeln!( + output, + " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, active_mask, value_map, block_map, loc,\n {:?},\n )?))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n"); + } + for record in ldmatrix(catalog) { + let (shape, multiplicity, layout, element, state_space) = ldmatrix_attr_variants(record); + let register_count = record.ldmatrix.as_ref().unwrap().variant.register_count(); + output.push_str(" "); + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + render_inline_patterns(&mut output, &path_refs); + output.push_str(" => {\n"); + output.push_str(" require_arity(name, args.len(), 1, &loc)?;\n"); + output.push_str( + " let (address, last_op) = rvalue::translate_operand(\n ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n )?;\n", + ); + writeln!( + output, + " let load = LdmatrixOp::build(ctx, address, {shape}, {multiplicity}, {layout}, {element}, {state_space});" + ) + .unwrap(); + output.push_str(" load.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, load, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str(" helpers::insert_op(ctx, load, block_ptr, last_op);\n"); + if register_count == 1 { + output.push_str(" let value = load.deref(ctx).get_result(0);\n"); + output.push_str(" let last_op = load;\n"); + } else { + writeln!( + output, + " let (value, last_op) = helpers::bundle_generated_u32_results_as_array(ctx, load, {register_count}, loc.clone());" + ) + .unwrap(); + } + writeln!( + output, + " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, value, target, block_ptr, last_op, value_map, block_map, loc,\n {:?},\n )?))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n"); + } + for record in stmatrices(catalog) { + let (multiplicity, _) = stmatrix_variant(record).expect("stmatrix variant"); + let arity = multiplicity.register_count() + 1; + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(&mut output, &path_refs); + output.push_str(" => {\n"); + writeln!( + output, + " require_arity(name, args.len(), {arity}, &loc)?;" + ) + .unwrap(); + output.push_str( + " let mut last_op = prev_op;\n let mut operands = Vec::with_capacity(args.len());\n for arg in args {\n let (value, translated) = rvalue::translate_operand(\n ctx, body, arg, value_map, block_ptr, last_op, loc.clone(),\n )?;\n last_op = translated;\n operands.push(value);\n }\n", + ); + writeln!( + output, + " let store = Operation::new(ctx, {}::get_concrete_op_info(), vec![], operands, vec![], 0);", + record.dialect.op_type + ) + .unwrap(); + output.push_str(" store.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, store, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, store, block_ptr, last_op);\n if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, store, block_map, loc)))\n } else {\n", + ); + writeln!( + output, + " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", + format!( + "{} call without target block", + stmatrix_compatibility_name(record) + ) + ) + .unwrap(); + output.push_str(" }\n }\n"); + } + for record in register_mmas(catalog) { + let adapter = match record.register_mma.as_ref().unwrap().adapter { + RegisterMmaAdapter::C2U32A2U32B1U32ToD2U32 => { + "GeneratedMmaImportAdapter::C2U32A2U32B1U32ToD2U32" + } + RegisterMmaAdapter::C2U32A4U32B2U32ToD2U32 => { + "GeneratedMmaImportAdapter::C2U32A4U32B2U32ToD2U32" + } + RegisterMmaAdapter::C4F32A2U32B1U32ToD4F32 => { + "GeneratedMmaImportAdapter::C4F32A2U32B1U32ToD4F32" + } + RegisterMmaAdapter::C4F32A4U32B2U32ToD4F32 => { + "GeneratedMmaImportAdapter::C4F32A4U32B2U32ToD4F32" + } + RegisterMmaAdapter::C2F64A1F64B1F64ToD2F64 => { + "GeneratedMmaImportAdapter::C2F64A1F64B1F64ToD2F64" + } + RegisterMmaAdapter::C2I32A1U32B1U32ToD2I32 => { + "GeneratedMmaImportAdapter::C2I32A1U32B1U32ToD2I32" + } + RegisterMmaAdapter::C4I32A4U32B2U32ToD4I32 => { + "GeneratedMmaImportAdapter::C4I32A4U32B2U32ToD4I32" + } + RegisterMmaAdapter::C4I32A2U32B1U32ToD4I32 => { + "GeneratedMmaImportAdapter::C4I32A2U32B1U32ToD4I32" + } + RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 => { + "GeneratedMmaImportAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32" + } + }; + let ( + shape, + operation, + kind, + accumulator, + a_element, + b_element, + a_layout, + b_layout, + overflow, + ) = register_mma_attr_variants(record); + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(&mut output, &path_refs); + output.push_str(" => {\n"); + let arity = record.rust.arguments.len(); + writeln!( + output, + " require_arity(name, args.len(), {arity}, &loc)?;" + ) + .unwrap(); + writeln!( + output, + " let (operands, last_op, result_ty, result_count) = import_generated_mma_operands(ctx, body, args, block_ptr, prev_op, value_map, loc.clone(), {adapter})?;" + ) + .unwrap(); + output.push_str( + " let mma = Operation::new(ctx, RegisterMmaOp::get_concrete_op_info(), vec![result_ty; result_count], operands, vec![], 0);\n mma.deref_mut(ctx).set_loc(loc.clone());\n let mma = RegisterMmaOp::new(mma);\n", + ); + writeln!( + output, + " mma.set_attr_nvvm_register_mma_shape(ctx, {shape});\n mma.set_attr_nvvm_register_mma_operation(ctx, {operation});\n mma.set_attr_nvvm_register_mma_kind(ctx, {kind});\n mma.set_attr_nvvm_register_mma_accumulator(ctx, {accumulator});\n mma.set_attr_nvvm_register_mma_a_element(ctx, {a_element});\n mma.set_attr_nvvm_register_mma_b_element(ctx, {b_element});\n mma.set_attr_nvvm_register_mma_a_layout(ctx, {a_layout});\n mma.set_attr_nvvm_register_mma_b_layout(ctx, {b_layout});\n mma.set_attr_nvvm_register_mma_overflow(ctx, {overflow});" + ) + .unwrap(); + output.push_str(" let mma = mma.get_operation();\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, mma, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, mma, block_ptr, Some(last_op));\n let (result, last_op) = bundle_generated_mma_results(ctx, mma, result_ty, result_count, loc.clone());\n", + ); + writeln!( + output, + " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, last_op, value_map, block_map, loc,\n {:?},\n )?))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n"); + } + for record in sparse_mmas(catalog) { + let ( + shape, + accumulator, + a_element, + b_element, + a_layout, + b_layout, + overflow, + metadata, + selector, + ) = sparse_mma_attr_variants(record); + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(&mut output, &path_refs); + output.push_str(" => {\n"); + output.push_str(" require_arity(name, args.len(), 5, &loc)?;\n"); + writeln!( + output, + " if !matches!(&args[4], mir::Operand::Constant(_)) {{\n return input_err!(\n loc,\n TranslationErr::unsupported(\n {:?}.to_owned()\n )\n );\n }}", + sparse_mma_selector_error(record) + ) + .unwrap(); + writeln!( + output, + " let (mut operands, last_op, result_ty, result_count) = import_generated_mma_operands(ctx, body, args, block_ptr, prev_op, value_map, loc.clone(), {})?;", + sparse_mma_import_adapter(record) + ) + .unwrap(); + output.push_str( + " let (metadata, last_op) = rvalue::translate_operand(\n ctx, body, &args[3], value_map, block_ptr, Some(last_op), loc.clone(),\n )?;\n operands.push(metadata);\n let (selector_value, last_op) = rvalue::translate_operand(\n ctx, body, &args[4], value_map, block_ptr, last_op, loc.clone(),\n )?;\n operands.push(selector_value);\n", + ); + output.push_str( + " let mma = Operation::new(ctx, SparseMmaOp::get_concrete_op_info(), vec![result_ty; result_count], operands, vec![], 0);\n mma.deref_mut(ctx).set_loc(loc.clone());\n let mma = SparseMmaOp::new(mma);\n", + ); + writeln!( + output, + " mma.set_attr_nvvm_sparse_mma_shape(ctx, {shape});\n mma.set_attr_nvvm_sparse_mma_accumulator(ctx, {accumulator});\n mma.set_attr_nvvm_sparse_mma_a_element(ctx, {a_element});\n mma.set_attr_nvvm_sparse_mma_b_element(ctx, {b_element});\n mma.set_attr_nvvm_sparse_mma_a_layout(ctx, {a_layout});\n mma.set_attr_nvvm_sparse_mma_b_layout(ctx, {b_layout});\n mma.set_attr_nvvm_sparse_mma_overflow(ctx, {overflow});\n mma.set_attr_nvvm_sparse_mma_metadata(ctx, {metadata});\n mma.set_attr_nvvm_sparse_mma_selector(ctx, {selector});" + ) + .unwrap(); + output.push_str(" let mma = mma.get_operation();\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, mma, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, mma, block_ptr, last_op);\n let (result, last_op) = bundle_generated_mma_results(ctx, mma, result_ty, result_count, loc.clone());\n", + ); + writeln!( + output, + " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, last_op, value_map, block_map, loc,\n {:?},\n )?))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n"); + } + for record in packed_atomics(catalog) { + let format = match record.packed_atomic.as_ref().unwrap().format { + PackedAtomicFormat::F16x2 => "PackedAtomicFormatAttr::F16x2", + PackedAtomicFormat::Bf16x2 => "PackedAtomicFormatAttr::Bf16x2", + }; + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(&mut output, &path_refs); + output.push_str(" => {\n"); + output.push_str(" require_arity(name, args.len(), 2, &loc)?;\n"); + output.push_str( + " let (address, last_op) = rvalue::translate_operand(\n ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n )?;\n", + ); + output.push_str( + " let (addend, last_op) = rvalue::translate_operand(\n ctx, body, &args[1], value_map, block_ptr, last_op, loc.clone(),\n )?;\n", + ); + writeln!( + output, + " let atom = PackedAtomicAddOp::build(ctx, address, addend, {format});" + ) + .unwrap(); + output.push_str(" atom.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, atom, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str(" helpers::insert_op(ctx, atom, block_ptr, last_op);\n"); + output.push_str(" let value = atom.deref(ctx).get_result(0);\n"); + writeln!( + output, + " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, value, target, block_ptr, atom, value_map, block_map, loc,\n {:?},\n )?))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n"); + } + for record in redux(catalog) { + debug_assert_eq!( + record.redux.as_ref().unwrap().adapter, + ReduxAdapter::MaskValueToSourceMemberMask + ); + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(&mut output, &path_refs); + output.push_str(" => {\n"); + output.push_str(" require_arity(name, args.len(), 2, &loc)?;\n"); + output.push_str( + " let (member_mask, last_op) = rvalue::translate_operand(\n ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n )?;\n", + ); + output.push_str( + " let (value, last_op) = rvalue::translate_operand(\n ctx, body, &args[1], value_map, block_ptr, last_op, loc.clone(),\n )?;\n", + ); + writeln!( + output, + " let reduction = {}::build(ctx, member_mask, value);", + record.dialect.op_type + ) + .unwrap(); + output.push_str(" reduction.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, reduction, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, reduction, block_ptr, last_op);\n let result = reduction.deref(ctx).get_result(0);\n", + ); + writeln!( + output, + " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, reduction, value_map, block_map, loc,\n {:?},\n )?))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n"); + } + for record in vote_intrinsics(catalog) { + debug_assert_eq!( + record.vote.as_ref().unwrap().adapter, + VoteAdapter::DirectMaskPredicate + ); + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(&mut output, &path_refs); + output.push_str(" => {\n"); + output.push_str(" require_arity(name, args.len(), 2, &loc)?;\n"); + output.push_str( + " let (member_mask, last_op) = rvalue::translate_operand(\n ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n )?;\n", + ); + output.push_str( + " let (predicate, last_op) = rvalue::translate_operand(\n ctx, body, &args[1], value_map, block_ptr, last_op, loc.clone(),\n )?;\n", + ); + writeln!( + output, + " let vote = {}::build(ctx, member_mask, predicate);", + record.dialect.op_type + ) + .unwrap(); + output.push_str(" vote.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, vote, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, vote, block_ptr, last_op);\n let result = vote.deref(ctx).get_result(0);\n", + ); + writeln!( + output, + " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, vote, value_map, block_map, loc,\n {:?},\n )?))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n"); + } + for record in warp_matches(catalog) { + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(&mut output, &path_refs); + output.push_str(" => {\n"); + output.push_str(" require_arity(name, args.len(), 2, &loc)?;\n"); + output.push_str( + " let (member_mask, last_op) = rvalue::translate_operand(\n ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n )?;\n", + ); + output.push_str( + " let (value, last_op) = rvalue::translate_operand(\n ctx, body, &args[1], value_map, block_ptr, last_op, loc.clone(),\n )?;\n", + ); + writeln!( + output, + " let warp_match = {}::build(ctx, member_mask, value);", + record.dialect.op_type + ) + .unwrap(); + output.push_str(" warp_match.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, warp_match, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, warp_match, block_ptr, last_op);\n let result = warp_match.deref(ctx).get_result(0);\n", + ); + writeln!( + output, + " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, warp_match, value_map, block_map, loc,\n {:?},\n )?))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n"); + } + for record in elect_intrinsics(catalog) { + render_importer_elect_dispatch(&mut output, catalog, record); + } + for record in warp_barriers(catalog) { + debug_assert_eq!( + record.warp_barrier.as_ref().unwrap().adapter, + WarpBarrierAdapter::DirectMemberMask + ); + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(&mut output, &path_refs); + output.push_str(" => {\n"); + output.push_str(" require_arity(name, args.len(), 1, &loc)?;\n"); + output.push_str( + " let (member_mask, last_op) = rvalue::translate_operand(\n ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n )?;\n", + ); + writeln!( + output, + " let barrier = {}::build(ctx, member_mask);", + record.dialect.op_type + ) + .unwrap(); + output.push_str(" barrier.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, barrier, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, barrier, block_ptr, last_op);\n if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, barrier, block_map, loc)))\n } else {\n", + ); + writeln!( + output, + " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n }\n"); + } + for record in warp_shuffles(catalog) { + debug_assert!(matches!( + record.warp_shuffle.as_ref().unwrap().adapter, + WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp + | WarpShuffleAdapter::MaskValueLaneOrDeltaSplitI64LowHighB32InsertClampReassemble + )); + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(&mut output, &path_refs); + output.push_str(" => {\n"); + output.push_str(" require_arity(name, args.len(), 3, &loc)?;\n"); + output.push_str( + " let (member_mask, last_op) = rvalue::translate_operand(\n ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n )?;\n", + ); + output.push_str( + " let (value, last_op) = rvalue::translate_operand(\n ctx, body, &args[1], value_map, block_ptr, last_op, loc.clone(),\n )?;\n", + ); + output.push_str( + " let (lane_or_delta, last_op) = rvalue::translate_operand(\n ctx, body, &args[2], value_map, block_ptr, last_op, loc.clone(),\n )?;\n", + ); + writeln!( + output, + " let shuffle = {}::build(ctx, member_mask, value, lane_or_delta);", + record.dialect.op_type + ) + .unwrap(); + output.push_str(" shuffle.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, shuffle, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, shuffle, block_ptr, last_op);\n let result = shuffle.deref(ctx).get_result(0);\n", + ); + writeln!( + output, + " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, shuffle, value_map, block_map, loc,\n {:?},\n )?))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n"); + } + for record in dot_products(catalog) { + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(&mut output, &path_refs); + output.push_str(" => {\n"); + output.push_str(" require_arity(name, args.len(), 3, &loc)?;\n"); + output.push_str( + " let (a, last_op) = rvalue::translate_operand(\n ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n )?;\n", + ); + output.push_str( + " let (b, last_op) = rvalue::translate_operand(\n ctx, body, &args[1], value_map, block_ptr, last_op, loc.clone(),\n )?;\n", + ); + output.push_str( + " let (c, last_op) = rvalue::translate_operand(\n ctx, body, &args[2], value_map, block_ptr, last_op, loc.clone(),\n )?;\n", + ); + writeln!( + output, + " let dot = {}::build(ctx, a, b, c);", + record.dialect.op_type + ) + .unwrap(); + output.push_str(" dot.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, dot, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, dot, block_ptr, last_op);\n let result = dot.deref(ctx).get_result(0);\n", + ); + writeln!( + output, + " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, dot, value_map, block_map, loc,\n {:?},\n )?))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n"); + } + for record in packed_alus(catalog) { + render_importer_pure_value_dispatch(&mut output, catalog, record); + } + for record in integer_minmaxes(catalog) { + render_importer_pure_value_dispatch(&mut output, catalog, record); + } + for record in packed_conversions(catalog) { + render_importer_pure_value_dispatch(&mut output, catalog, record); + } + for record in scalar_conversions(catalog) { + render_importer_scalar_conversion_dispatch(&mut output, catalog, record); + } + for record in scalar_arithmetics(catalog) { + render_importer_scalar_arithmetic_dispatch(&mut output, catalog, record); + } + for record in scalar_maths(catalog) { + render_importer_scalar_math_dispatch(&mut output, catalog, record); + } + for record in extended_minmax(catalog) { + render_importer_extended_minmax_dispatch(&mut output, catalog, record); + } + for record in movmatrix(catalog) { + render_importer_pure_value_dispatch(&mut output, catalog, record); + } + for record in prmts(catalog) { + let mode = match record.prmt.as_ref().unwrap().mode { + PrmtMode::Generic => "PrmtModeAttr::Generic", + PrmtMode::F4e => "PrmtModeAttr::F4e", + PrmtMode::B4e => "PrmtModeAttr::B4e", + PrmtMode::Rc8 => "PrmtModeAttr::Rc8", + PrmtMode::Ecl => "PrmtModeAttr::Ecl", + PrmtMode::Ecr => "PrmtModeAttr::Ecr", + PrmtMode::Rc16 => "PrmtModeAttr::Rc16", + }; + let arity = record.rust.arguments.len(); + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(&mut output, &path_refs); + output.push_str(" => {\n"); + writeln!( + output, + " require_arity(name, args.len(), {arity}, &loc)?;" + ) + .unwrap(); + for index in 0..arity { + let previous = if index == 0 { "prev_op" } else { "last_op" }; + writeln!( + output, + " let (arg{index}, last_op) = rvalue::translate_operand(\n ctx, body, &args[{index}], value_map, block_ptr, {previous}, loc.clone(),\n )?;" + ) + .unwrap(); + } + let arguments = (0..arity) + .map(|index| format!("arg{index}")) + .collect::>() + .join(", "); + writeln!( + output, + " let prmt = PrmtOp::build(ctx, vec![{arguments}], {mode});" + ) + .unwrap(); + output.push_str(" prmt.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, prmt, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, prmt, block_ptr, last_op);\n let result = prmt.deref(ctx).get_result(0);\n", + ); + writeln!( + output, + " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, prmt, value_map, block_map, loc,\n {:?},\n )?))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n"); + } + if cluster_barriers(catalog).next().is_some() { + let arrive = cluster_barriers(catalog) + .find(|record| { + record + .cluster_barrier + .as_ref() + .is_some_and(|barrier| barrier.mode == ClusterBarrierMode::ArriveAligned) + }) + .expect("aligned cluster arrival"); + let wait = cluster_barriers(catalog) + .find(|record| { + record + .cluster_barrier + .as_ref() + .is_some_and(|barrier| barrier.mode == ClusterBarrierMode::WaitAligned) + }) + .expect("aligned cluster wait"); + output.push_str( + " \"cuda_device::cluster::cluster_sync\" => {\n require_arity(name, args.len(), 0, &loc)?;\n let arrive = ClusterBarrierOp::build(ctx, ClusterBarrierModeAttr::ArriveAligned);\n arrive.deref_mut(ctx).set_loc(loc.clone());\n", + ); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, arrive, {:?});", + intrinsic_marker(catalog, arrive) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, arrive, block_ptr, prev_op);\n let wait = ClusterBarrierOp::build(ctx, ClusterBarrierModeAttr::WaitAligned);\n wait.deref_mut(ctx).set_loc(loc.clone());\n", + ); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, wait, {:?});", + intrinsic_marker(catalog, wait) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, wait, block_ptr, Some(arrive));\n if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, wait, block_map, loc)))\n } else {\n input_err!(loc, TranslationErr::unsupported(\"cluster_sync call without target block\".to_owned()))\n }\n }\n", + ); + } + for record in cluster_barriers(catalog) { + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(&mut output, &path_refs); + output.push_str(" => {\n require_arity(name, args.len(), 0, &loc)?;\n"); + writeln!( + output, + " let barrier = ClusterBarrierOp::build(ctx, {});", + cluster_barrier_attr(record) + ) + .unwrap(); + output.push_str(" barrier.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, barrier, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, barrier, block_ptr, prev_op);\n if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, barrier, block_map, loc)))\n } else {\n", + ); + writeln!( + output, + " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n }\n"); + } + for record in cluster_memory(catalog) { + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(&mut output, &path_refs); + output.push_str(" => {\n"); + output.push_str( + " require_arity(name, args.len(), 2, &loc)?;\n\ + let (source, last_op) = rvalue::translate_operand(\n\ + ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n\ + )?;\n\ + let (rank, last_op) = rvalue::translate_operand(\n\ + ctx, body, &args[1], value_map, block_ptr, last_op, loc.clone(),\n\ + )?;\n", + ); + writeln!( + output, + " let cluster = {}::build(ctx, source, rank);", + record.dialect.op_type + ) + .unwrap(); + output.push_str(" cluster.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, cluster, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, cluster, block_ptr, last_op);\n\ + let result = cluster.deref(ctx).get_result(0);\n", + ); + writeln!( + output, + " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, cluster, value_map, block_map, loc,\n {:?},\n )?))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n"); + } + for record in debug_controls(catalog) { + let operation = record.debug_control.as_ref().unwrap().operation; + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(&mut output, &path_refs); + output.push_str(" => {\n"); + match operation { + DebugControlOperation::Trap => { + output.push_str(" require_arity(name, args.len(), 0, &loc)?;\n"); + output.push_str(" let debug = TrapOp::build(ctx);\n"); + output.push_str(" debug.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, debug, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, debug, block_ptr, prev_op);\n\ + let unreachable = Operation::new(\n\ + ctx, MirUnreachableOp::get_concrete_op_info(), vec![], vec![], vec![], 0,\n\ + );\n\ + unreachable.deref_mut(ctx).set_loc(loc);\n\ + helpers::insert_op(ctx, unreachable, block_ptr, Some(debug));\n\ + Ok(Some(unreachable))\n", + ); + } + DebugControlOperation::Breakpoint => { + output.push_str(" require_arity(name, args.len(), 0, &loc)?;\n"); + output.push_str(" let debug = BreakpointOp::build(ctx);\n"); + output.push_str(" debug.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, debug, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, debug, block_ptr, prev_op);\n\ + if let Some(target_idx) = target {\n\ + Ok(Some(helpers::emit_goto(ctx, *target_idx, debug, block_map, loc)))\n\ + } else {\n", + ); + writeln!( + output, + " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n"); + } + DebugControlOperation::Pmevent => { + output.push_str( + " require_arity(name, args.len(), 1, &loc)?;\n\ + if !matches!(&args[0], mir::Operand::Constant(_)) {\n\ + return input_err!(\n\ + loc,\n\ + TranslationErr::unsupported(\n\ + \"prof_trigger requires a compile-time constant event ID in 0..=15\".to_owned()\n\ + )\n\ + );\n\ + }\n\ + let (event_id_value, last_op) = rvalue::translate_operand(\n\ + ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n\ + )?;\n\ + let event_id = event_id_value\n\ + .defining_op()\n\ + .and_then(|defining_op| Operation::get_op::(defining_op, ctx))\n\ + .and_then(|constant| constant.get_attr_value(ctx))\n\ + .map(|value| value.value().to_u64())\n\ + .and_then(|value| u32::try_from(value).ok())\n\ + .filter(|value| *value <= 15);\n\ + let Some(event_id) = event_id else {\n\ + return input_err!(\n\ + loc,\n\ + TranslationErr::unsupported(\n\ + \"prof_trigger requires a compile-time constant event ID in 0..=15\".to_owned()\n\ + )\n\ + );\n\ + };\n\ + let debug = PmEventOp::build(ctx, event_id);\n\ + debug.deref_mut(ctx).set_loc(loc.clone());\n", + ); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, debug, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, debug, block_ptr, last_op);\n\ + if let Some(target_idx) = target {\n\ + Ok(Some(helpers::emit_goto(ctx, *target_idx, debug, block_map, loc)))\n\ + } else {\n", + ); + writeln!( + output, + " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n"); + } + } + output.push_str(" }\n"); + } + for record in clc_intrinsics(catalog) { + let query = !matches!( + record.clc.as_ref().unwrap().operation, + ClcOperation::TryCancel | ClcOperation::TryCancelMulticast + ); + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(&mut output, &path_refs); + output.push_str(" => {\n require_arity(name, args.len(), 2, &loc)?;\n"); + output.push_str( + " let (arg0, last_op) = rvalue::translate_operand(\n ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n )?;\n let (arg1, last_op) = rvalue::translate_operand(\n ctx, body, &args[1], value_map, block_ptr, last_op, loc.clone(),\n )?;\n", + ); + let result_types = if query { + "vec![pliron::builtin::types::IntegerType::get(ctx, 32, pliron::builtin::types::Signedness::Unsigned).into()]" + } else { + "vec![]" + }; + writeln!( + output, + " let intrinsic = Operation::new(ctx, {}::get_concrete_op_info(), {result_types}, vec![arg0, arg1], vec![], 0);", + record.dialect.op_type + ) + .unwrap(); + output.push_str(" intrinsic.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, intrinsic, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str(" helpers::insert_op(ctx, intrinsic, block_ptr, last_op);\n"); + if query { + output.push_str(" let result = intrinsic.deref(ctx).get_result(0);\n"); + writeln!( + output, + " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, intrinsic, value_map, block_map, loc,\n {:?},\n )?))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + } else { + output.push_str(" if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, intrinsic, block_map, loc)))\n } else {\n"); + writeln!( + output, + " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n"); + } + output.push_str(" }\n"); + } + for record in wgmma_controls(catalog) { + let control = record.wgmma_control.as_ref().unwrap(); + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(&mut output, &path_refs); + output.push_str(" => {\n"); + match control.mode { + WgmmaControlMode::Fence | WgmmaControlMode::CommitGroup => { + output.push_str(" require_arity(name, args.len(), 0, &loc)?;\n"); + writeln!( + output, + " let control = {}::build(ctx);", + record.dialect.op_type + ) + .unwrap(); + output.push_str(" control.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, control, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, control, block_ptr, prev_op);\n", + ); + } + WgmmaControlMode::WaitGroup => { + output.push_str( + " require_arity(name, args.len(), 1, &loc)?;\n\ + if !matches!(&args[0], mir::Operand::Constant(_)) {\n\ + return input_err!(\n\ + loc,\n\ + TranslationErr::unsupported(\n\ + \"wgmma_wait_group requires a compile-time constant\".to_owned()\n\ + )\n\ + );\n\ + }\n\ + let (max_pending, last_op) = rvalue::translate_operand(\n\ + ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n\ + )?;\n\ + if max_pending\n\ + .defining_op()\n\ + .and_then(|op| Operation::get_op::(op, ctx))\n\ + .is_none()\n\ + {\n\ + return input_err!(\n\ + loc,\n\ + TranslationErr::unsupported(\n\ + \"wgmma_wait_group requires a compile-time constant\".to_owned()\n\ + )\n\ + );\n\ + }\n\ + let control = WgmmaWaitGroupSyncAlignedOp::build(ctx, max_pending);\n\ + control.deref_mut(ctx).set_loc(loc.clone());\n", + ); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, control, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, control, block_ptr, last_op);\n", + ); + } + } + output.push_str( + " if let Some(target_idx) = target {\n\ + Ok(Some(helpers::emit_goto(ctx, *target_idx, control, block_map, loc)))\n\ + } else {\n", + ); + writeln!( + output, + " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n }\n"); + } + for record in cp_async_copies(catalog) { + let copy = record.cp_async_copy.as_ref().unwrap(); + let dynamic = copy.source_size == CpAsyncSourceSize::Runtime; + let arity = if dynamic { 3 } else { 2 }; + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(&mut output, &path_refs); + output.push_str(" => {\n"); + writeln!( + output, + " require_arity(name, args.len(), {arity}, &loc)?;" + ) + .unwrap(); + output.push_str( + " let (shared_dst, last_op) = rvalue::translate_operand(\n ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n )?;\n", + ); + output.push_str( + " let (global_src, last_op) = rvalue::translate_operand(\n ctx, body, &args[1], value_map, block_ptr, last_op, loc.clone(),\n )?;\n", + ); + if dynamic { + output.push_str( + " let (source_size, last_op) = rvalue::translate_operand(\n ctx, body, &args[2], value_map, block_ptr, last_op, loc.clone(),\n )?;\n", + ); + writeln!( + output, + " let copy = {}::build(ctx, shared_dst, global_src, source_size);", + record.dialect.op_type + ) + .unwrap(); + } else { + writeln!( + output, + " let copy = {}::build(ctx, shared_dst, global_src);", + record.dialect.op_type + ) + .unwrap(); + } + output.push_str(" copy.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, copy, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, copy, block_ptr, last_op);\n if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, copy, block_map, loc)))\n } else {\n", + ); + writeln!( + output, + " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n }\n"); + } + for record in cp_async_controls(catalog) { + let control = record.cp_async_control.as_ref().unwrap(); + let has_immediate = control.operation == CpAsyncControlOperation::WaitGroup; + let arity = usize::from(has_immediate); + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(&mut output, &path_refs); + output.push_str(" => {\n"); + writeln!( + output, + " require_arity(name, args.len(), {arity}, &loc)?;" + ) + .unwrap(); + if has_immediate { + output.push_str( + " if !matches!(&args[0], mir::Operand::Constant(_)) {\n return input_err!(\n loc,\n TranslationErr::unsupported(\n \"cp_async_wait_group requires a compile-time constant max_pending value\".to_owned()\n )\n );\n }\n", + ); + output.push_str( + " let (max_pending, last_op) = rvalue::translate_operand(\n ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n )?;\n", + ); + writeln!( + output, + " let control = {}::build(ctx, max_pending);", + record.dialect.op_type + ) + .unwrap(); + } else { + output.push_str(" let last_op = prev_op;\n"); + writeln!( + output, + " let control = {}::build(ctx);", + record.dialect.op_type + ) + .unwrap(); + } + output.push_str(" control.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, control, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, control, block_ptr, last_op);\n if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, control, block_map, loc)))\n } else {\n", + ); + writeln!( + output, + " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n }\n"); + } + for record in cp_async_mbarriers(catalog) { + debug_assert_eq!( + record.cp_async_mbarrier.as_ref().unwrap().adapter, + CpAsyncMbarrierAdapter::PointerToVoid + ); + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(&mut output, &path_refs); + output.push_str(" => {\n"); + output.push_str(" require_arity(name, args.len(), 1, &loc)?;\n"); + output.push_str( + " let (barrier, last_op) = rvalue::translate_operand(\n ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n )?;\n", + ); + writeln!( + output, + " let bridge = {}::build(ctx, barrier);", + record.dialect.op_type + ) + .unwrap(); + output.push_str(" bridge.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, bridge, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, bridge, block_ptr, last_op);\n if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, bridge, block_map, loc)))\n } else {\n", + ); + writeln!( + output, + " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n }\n"); + } + for record in mbarrier_basics(catalog) { + let mbarrier = record.mbarrier_basic.as_ref().unwrap(); + let (argument_names, returns_value) = match mbarrier.operation { + MbarrierBasicOperation::Init => { + debug_assert_eq!( + mbarrier.adapter, + MbarrierBasicAdapter::InitPointerCountToVoid + ); + (&["barrier", "expected_count"][..], false) + } + MbarrierBasicOperation::Arrive => { + debug_assert_eq!(mbarrier.adapter, MbarrierBasicAdapter::ArrivePointerToToken); + (&["barrier"][..], true) + } + MbarrierBasicOperation::TestWait => { + debug_assert_eq!( + mbarrier.adapter, + MbarrierBasicAdapter::TestWaitPointerTokenToPredicate + ); + (&["barrier", "token"][..], true) + } + MbarrierBasicOperation::Inval => { + debug_assert_eq!(mbarrier.adapter, MbarrierBasicAdapter::InvalPointerToVoid); + (&["barrier"][..], false) + } + }; + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(&mut output, &path_refs); + output.push_str(" => {\n"); + writeln!( + output, + " require_arity(name, args.len(), {}, &loc)?;", + argument_names.len() + ) + .unwrap(); + for (index, argument_name) in argument_names.iter().enumerate() { + let previous = if index == 0 { "prev_op" } else { "last_op" }; + writeln!( + output, + " let ({argument_name}, last_op) = rvalue::translate_operand(\n ctx, body, &args[{index}], value_map, block_ptr, {previous}, loc.clone(),\n )?;" + ) + .unwrap(); + } + writeln!( + output, + " let mbarrier = {}::build(ctx, {});", + record.dialect.op_type, + argument_names.join(", ") + ) + .unwrap(); + output.push_str(" mbarrier.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, mbarrier, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str(" helpers::insert_op(ctx, mbarrier, block_ptr, last_op);\n"); + if returns_value { + output.push_str(" let result = mbarrier.deref(ctx).get_result(0);\n"); + writeln!( + output, + " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, mbarrier, value_map, block_map, loc,\n {:?},\n )?))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + } else { + output.push_str( + " if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, mbarrier, block_map, loc)))\n } else {\n", + ); + writeln!( + output, + " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n"); + } + output.push_str(" }\n"); + } + for record in mbarrier_extended(catalog) { + let contract = record.mbarrier_extended.as_ref().unwrap(); + let (arguments, returns_value): (Vec<(usize, &str)>, bool) = match contract.adapter { + MbarrierExtendedAdapter::PointerTxCountBytesToTokenDroppingTxCount => { + (vec![(0, "barrier"), (2, "bytes")], true) + } + MbarrierExtendedAdapter::RawClusterAddressToVoid => (vec![(0, "address")], false), + MbarrierExtendedAdapter::PointerTokenToPredicate => { + (vec![(0, "barrier"), (1, "token")], true) + } + MbarrierExtendedAdapter::PointerParityToPredicate => { + (vec![(0, "barrier"), (1, "parity")], true) + } + MbarrierExtendedAdapter::ZeroOperandsToVoid => (vec![], false), + MbarrierExtendedAdapter::NanosecondsToVoid => (vec![(0, "ns")], false), + }; + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(&mut output, &path_refs); + output.push_str(" => {\n"); + writeln!( + output, + " require_arity(name, args.len(), {}, &loc)?;", + record.rust.arguments.len() + ) + .unwrap(); + if arguments.is_empty() { + output.push_str(" let last_op = prev_op;\n"); + } + for (position, (argument_index, argument_name)) in arguments.iter().enumerate() { + let previous = if position == 0 { "prev_op" } else { "last_op" }; + writeln!( + output, + " let ({argument_name}, last_op) = rvalue::translate_operand(\n ctx, body, &args[{argument_index}], value_map, block_ptr, {previous}, loc.clone(),\n )?;" + ) + .unwrap(); + } + let build_arguments = arguments + .iter() + .map(|(_, name)| *name) + .collect::>() + .join(", "); + writeln!( + output, + " let extended = {}::build(ctx{}{});", + record.dialect.op_type, + if build_arguments.is_empty() { "" } else { ", " }, + build_arguments, + ) + .unwrap(); + output.push_str(" extended.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, extended, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str(" helpers::insert_op(ctx, extended, block_ptr, last_op);\n"); + if returns_value { + output.push_str(" let result = extended.deref(ctx).get_result(0);\n"); + writeln!( + output, + " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, extended, value_map, block_map, loc,\n {:?},\n )?))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + } else { + output.push_str( + " if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, extended, block_map, loc)))\n } else {\n", + ); + writeln!( + output, + " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n"); + } + output.push_str(" }\n"); + } + for record in sync_intrinsics(catalog) { + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(&mut output, &path_refs); + output.push_str(" => {\n"); + output.push_str(" require_arity(name, args.len(), 0, &loc)?;\n"); + writeln!( + output, + " let barrier = Operation::new(ctx, {}::get_concrete_op_info(), vec![], vec![], vec![], 0);", + record.dialect.op_type + ) + .unwrap(); + output.push_str(" barrier.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, barrier, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, barrier, block_ptr, prev_op);\n if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, barrier, block_map, loc)))\n } else {\n", + ); + writeln!( + output, + " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n }\n"); + } + for record in execution_controls(catalog) { + let operation = ExecutionControlOperation::from_catalog_id(&record.id) + .expect("closed execution-control record"); + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(&mut output, &path_refs); + output.push_str(" => {\n"); + writeln!( + output, + " require_arity(name, args.len(), {}, &loc)?;", + operation.operand_count() + ) + .unwrap(); + if operation.requires_immediate_operands() { + output.push_str( + " if !matches!(&args[0], mir::Operand::Constant(_)) {\n return input_err!(\n loc,\n TranslationErr::unsupported(\n \"setmaxnreg requires a compile-time register count in 24..=256 divisible by 8\".to_owned()\n )\n );\n }\n let (register_count_value, last_op) = rvalue::translate_operand(\n ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(),\n )?;\n let register_count = register_count_value\n .defining_op()\n .and_then(|defining_op| Operation::get_op::(defining_op, ctx))\n .and_then(|constant| constant.get_attr_value(ctx))\n .map(|value| value.value().to_u64())\n .and_then(|value| u32::try_from(value).ok())\n .filter(|value| (24..=256).contains(value) && value % 8 == 0);\n let Some(register_count) = register_count else {\n return input_err!(\n loc,\n TranslationErr::unsupported(\n \"setmaxnreg requires a compile-time register count in 24..=256 divisible by 8\".to_owned()\n )\n );\n };\n", + ); + writeln!( + output, + " let control = {}::build(ctx, register_count);", + record.dialect.op_type + ) + .unwrap(); + } else if operation.operand_count() == 0 { + output.push_str(" let last_op = prev_op;\n"); + writeln!( + output, + " let control = Operation::new(ctx, {}::get_concrete_op_info(), vec![], vec![], vec![], 0);", + record.dialect.op_type + ) + .unwrap(); + } else { + output.push_str( + " let mut last_op = prev_op;\n let mut operands = Vec::with_capacity(args.len());\n for arg in args {\n let (value, translated) = rvalue::translate_operand(\n ctx, body, arg, value_map, block_ptr, last_op, loc.clone(),\n )?;\n last_op = translated;\n operands.push(value);\n }\n", + ); + writeln!( + output, + " let control = Operation::new(ctx, {}::get_concrete_op_info(), vec![], operands, vec![], 0);", + record.dialect.op_type + ) + .unwrap(); + } + output.push_str(" control.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, control, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, control, block_ptr, last_op);\n if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, control, block_map, loc)))\n } else {\n", + ); + writeln!( + output, + " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n }\n"); + } + for record in tma_intrinsics(catalog) { + let operation = record.tma.as_ref().unwrap().operation; + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(&mut output, &path_refs); + output.push_str(" => {\n"); + if matches!( + operation, + TmaOperation::WaitGroup | TmaOperation::WaitGroupRead + ) { + output.push_str( + " require_arity(name, args.len(), 1, &loc)?;\n if !matches!(args.first(), Some(mir::Operand::Constant(_))) {\n return input_err!(\n loc,\n TranslationErr::unsupported(\n \"TMA wait-group count must be a compile-time constant\".to_owned()\n )\n );\n }\n", + ); + } + let marker = intrinsic_marker(catalog, record); + match operation { + TmaOperation::G2sTile1d + | TmaOperation::G2sTile2d + | TmaOperation::G2sTile3d + | TmaOperation::G2sTile4d + | TmaOperation::G2sTile5d => { + writeln!( + output, + " Ok(Some(super::tma::emit_tma_g2s(\n ctx, body, args, target, block_ptr, prev_op, value_map, block_map, loc, {}, {marker:?},\n )?))", + operation.dimensions().unwrap() + ) + .unwrap(); + } + TmaOperation::G2sTile2dMulticast => { + writeln!( + output, + " Ok(Some(super::tma::emit_tma_g2s_multicast(\n ctx, body, args, target, block_ptr, prev_op, value_map, block_map, loc, {marker:?},\n )?))" + ) + .unwrap(); + } + TmaOperation::G2sTile2dMulticastCg2 => { + writeln!( + output, + " Ok(Some(super::tma::emit_tma_g2s_multicast_cg2(\n ctx, body, args, target, block_ptr, prev_op, value_map, block_map, loc, {marker:?},\n )?))" + ) + .unwrap(); + } + TmaOperation::S2gTile1d + | TmaOperation::S2gTile2d + | TmaOperation::S2gTile3d + | TmaOperation::S2gTile4d + | TmaOperation::S2gTile5d => { + writeln!( + output, + " Ok(Some(super::tma::emit_tma_s2g(\n ctx, body, args, target, block_ptr, prev_op, value_map, block_map, loc, {}, {marker:?},\n )?))", + operation.dimensions().unwrap() + ) + .unwrap(); + } + TmaOperation::CommitGroup => { + writeln!( + output, + " Ok(Some(super::tma::emit_tma_commit_group(\n ctx, args, target, block_ptr, prev_op, block_map, loc, {marker:?},\n )?))" + ) + .unwrap(); + } + TmaOperation::WaitGroup | TmaOperation::WaitGroupRead => { + writeln!( + output, + " Ok(Some(super::tma::emit_tma_wait_group(\n ctx, body, args, target, block_ptr, prev_op, value_map, block_map, loc, {}, {marker:?},\n )?))", + operation == TmaOperation::WaitGroupRead + ) + .unwrap(); + } + _ => { + let arity = record.dialect.operands.len(); + writeln!( + output, + " require_arity(name, args.len(), {arity}, &loc)?;" + ) + .unwrap(); + if matches!( + record.tma.as_ref().unwrap().adapter, + TmaAdapter::DescriptorOrdinalAndU32 + | TmaAdapter::DescriptorOrdinalAndU64 + | TmaAdapter::DescriptorAndImmediateU32 + ) { + output.push_str( + " if !matches!(args.get(1), Some(mir::Operand::Constant(_))) {\n return input_err!(\n loc,\n TranslationErr::unsupported(\n \"tensor-map replacement selector must be a compile-time constant\".to_owned()\n )\n );\n }\n", + ); + } + output.push_str( + " let mut last_op = prev_op;\n let mut operands = Vec::with_capacity(args.len());\n for arg in args {\n let (value, translated) = rvalue::translate_operand(\n ctx, body, arg, value_map, block_ptr, last_op, loc.clone(),\n )?;\n last_op = translated;\n operands.push(value);\n }\n", + ); + writeln!( + output, + " let intrinsic = Operation::new(ctx, dialect_nvvm::ops::{}::get_concrete_op_info(), vec![], operands, vec![], 0);", + record.dialect.op_type + ) + .unwrap(); + output.push_str(" intrinsic.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, intrinsic, {marker:?});" + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, intrinsic, block_ptr, last_op);\n if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, intrinsic, block_map, loc)))\n } else {\n", + ); + writeln!( + output, + " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n"); + } + } + output.push_str(" }\n"); + } + for record in tcgen05_intrinsics(catalog) { + let tcgen05 = record.tcgen05.as_ref().unwrap(); + if tcgen05.mma.is_some() { + render_importer_tcgen05_mma_dispatch(&mut output, catalog, record); + continue; + } + let operation = tcgen05.operation; + let has_half_split_offset = tcgen05 + .ld + .is_some_and(|ld| ld.shape == Tcgen05LdShape::M16x32bx2) + || tcgen05 + .st + .is_some_and(|st| st.shape == Tcgen05LdShape::M16x32bx2); + let mut path_refs = vec![record.rust.canonical_path.as_str()]; + path_refs.extend(record.rust.compatibility_paths.iter().map(String::as_str)); + output.push_str(" "); + render_inline_patterns(&mut output, &path_refs); + output.push_str(" => {\n"); + if operation == Tcgen05Operation::St { + let count = tcgen05_st_register_count(record); + writeln!( + output, + " require_arity(name, args.len(), {}, &loc)?;", + if has_half_split_offset { 3 } else { 2 } + ) + .unwrap(); + if has_half_split_offset { + output.push_str( + " if !matches!(args.get(1), Some(mir::Operand::Constant(_))) {\n return input_err!(\n loc,\n TranslationErr::unsupported(\n \"tcgen05 16x32bx2 half-split offset must be a compile-time constant\".to_owned()\n )\n );\n }\n", + ); + } + writeln!( + output, + " let (operands, last_op) = import_generated_tcgen05_store_operands(\n ctx, body, args, {count}, {has_half_split_offset}, block_ptr, prev_op, value_map, loc.clone(),\n )?;" + ) + .unwrap(); + writeln!( + output, + " let intrinsic = Operation::new(ctx, {}::get_concrete_op_info(), vec![], operands, vec![], 0);", + record.dialect.op_type + ) + .unwrap(); + output.push_str(" intrinsic.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, intrinsic, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, intrinsic, block_ptr, last_op);\n if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, intrinsic, block_map, loc)))\n } else {\n", + ); + writeln!( + output, + " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n }\n"); + continue; + } + let arity = record.dialect.operands.len(); + writeln!( + output, + " require_arity(name, args.len(), {arity}, &loc)?;" + ) + .unwrap(); + if has_half_split_offset { + output.push_str( + " if !matches!(args.get(1), Some(mir::Operand::Constant(_))) {\n return input_err!(\n loc,\n TranslationErr::unsupported(\n \"tcgen05 16x32bx2 half-split offset must be a compile-time constant\".to_owned()\n )\n );\n }\n", + ); + } + output.push_str( + " let mut last_op = prev_op;\n let mut operands = Vec::with_capacity(args.len());\n for arg in args {\n let (value, translated) = rvalue::translate_operand(\n ctx, body, arg, value_map, block_ptr, last_op, loc.clone(),\n )?;\n last_op = translated;\n operands.push(value);\n }\n", + ); + if has_half_split_offset { + output.push_str( + " if operands.get(1).and_then(|value| value.defining_op()).and_then(|op| Operation::get_op::(op, ctx)).is_none() {\n return input_err!(\n loc,\n TranslationErr::unsupported(\n \"tcgen05 16x32bx2 half-split offset must lower to a constant\".to_owned()\n )\n );\n }\n", + ); + } + let load = match operation { + Tcgen05Operation::Ld16x256bX8Pure => Some((32, "FP32Type::get(ctx).into()")), + Tcgen05Operation::Ld16x256bPure => Some((4, "FP32Type::get(ctx).into()")), + Tcgen05Operation::Ld => Some(( + tcgen05_ld_register_count(record), + "IntegerType::get(ctx, 32, Signedness::Unsigned).into()", + )), + _ => None, + }; + if let Some((count, result_ty)) = load { + writeln!( + output, + " let result_ty: TypeHandle = {result_ty};" + ) + .unwrap(); + writeln!( + output, + " let result_types = (0..{count}).map(|_| result_ty).collect();" + ) + .unwrap(); + writeln!( + output, + " let intrinsic = Operation::new(ctx, {}::get_concrete_op_info(), result_types, operands, vec![], 0);", + record.dialect.op_type + ) + .unwrap(); + output.push_str(" intrinsic.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, intrinsic, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output + .push_str(" helpers::insert_op(ctx, intrinsic, block_ptr, last_op);\n"); + if count == 1 { + output.push_str( + " let result = intrinsic.deref(ctx).get_result(0);\n let value = intrinsic;\n", + ); + } else { + writeln!( + output, + " let results: Vec = (0..{count}).map(|index| intrinsic.deref(ctx).get_result(index)).collect();" + ) + .unwrap(); + writeln!( + output, + " let array_ty = MirArrayType::get(ctx, result_ty, {count});" + ) + .unwrap(); + output.push_str( + " let array = Operation::new(ctx, MirConstructArrayOp::get_concrete_op_info(), vec![array_ty.into()], results, vec![], 0);\n array.deref_mut(ctx).set_loc(loc.clone());\n array.insert_after(ctx, intrinsic);\n let destination_rust_ty = match destination.ty(body.locals()) {\n Ok(ty) => ty,\n Err(error) => {\n return input_err!(\n loc,\n TranslationErr::unsupported(format!(\n \"failed to resolve destination type for intrinsic result: {error:?}\"\n ))\n );\n }\n };\n let destination_ty = types::translate_type(ctx, &destination_rust_ty)?;\n let array_result = array.deref(ctx).get_result(0);\n let (result, value) = if destination_ty == array_result.get_type(ctx) {\n (array_result, array)\n } else {\n let value = Operation::new(ctx, MirConstructStructOp::get_concrete_op_info(), vec![destination_ty], vec![array_result], vec![], 0);\n value.deref_mut(ctx).set_loc(loc.clone());\n value.insert_after(ctx, array);\n (value.deref(ctx).get_result(0), value)\n };\n helpers::set_compiler_result_bundle_marker(ctx, value);\n", + ); + } + writeln!( + output, + " Ok(Some(helpers::emit_store_result_and_goto(\n ctx, destination, result, target, block_ptr, value, value_map, block_map, loc,\n {:?},\n )?))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + } else { + writeln!( + output, + " let intrinsic = Operation::new(ctx, {}::get_concrete_op_info(), vec![], operands, vec![], 0);", + record.dialect.op_type + ) + .unwrap(); + output.push_str(" intrinsic.deref_mut(ctx).set_loc(loc.clone());\n"); + writeln!( + output, + " helpers::set_generated_intrinsic_marker(ctx, intrinsic, {:?});", + intrinsic_marker(catalog, record) + ) + .unwrap(); + output.push_str( + " helpers::insert_op(ctx, intrinsic, block_ptr, last_op);\n if let Some(target_idx) = target {\n Ok(Some(helpers::emit_goto(ctx, *target_idx, intrinsic, block_map, loc)))\n } else {\n", + ); + writeln!( + output, + " input_err!(loc, TranslationErr::unsupported({:?}.to_owned()))", + format!("{} call without target block", record.rust.name) + ) + .unwrap(); + output.push_str(" }\n"); + } + output.push_str(" }\n"); + } + output.push_str(" _ => Ok(None),\n }\n}\n\n"); + output.push_str( + "fn require_arity(\n name: &str,\n actual: usize,\n expected: usize,\n loc: &Location,\n) -> TranslationResult<()> {\n if actual != expected {\n return input_err!(\n loc.clone(),\n TranslationErr::unsupported(format!(\n \"generated intrinsic `{name}` expects {expected} arguments, got {actual}\"\n ))\n );\n }\n Ok(())\n}\n", + ); + if tcgen05_mma_intrinsics(catalog).next().is_some() { + output.push_str( + r#" + +fn generated_tcgen05_mma_selector(operand: Option<&mir::Operand>) -> Option { + use rustc_public::ty::{ConstantKind, TyConstKind}; + let mir::Operand::Constant(constant) = operand? else { return None; }; + let value: u128 = match constant.const_.kind() { + ConstantKind::Allocated(alloc) => alloc.read_uint().ok()?, + ConstantKind::Ty(value) => match value.kind() { + TyConstKind::Value(_, alloc) => alloc.read_uint().ok()?, + _ => value.eval_target_usize().ok()? as u128, + }, + ConstantKind::Unevaluated(_) | ConstantKind::Param(_) => return None, + ConstantKind::ZeroSized => return None, + }; + u32::try_from(value).ok() +} +"#, + ); + } + if tcgen05_intrinsics(catalog) + .any(|record| record.tcgen05.as_ref().unwrap().operation == Tcgen05Operation::St) + { + output.push_str( + r#" + +#[allow(clippy::too_many_arguments)] +fn import_generated_tcgen05_store_operands( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + expected_len: usize, + has_half_split_offset: bool, + block_ptr: Ptr, + prev_op: Option>, + value_map: &mut ValueMap, + loc: Location, +) -> TranslationResult<(Vec, Option>)> { + let expected_arity = if has_half_split_offset { 3 } else { 2 }; + if args.len() != expected_arity || !(1..=128).contains(&expected_len) { + return input_err!( + loc, + TranslationErr::unsupported( + "generated tcgen05 store has an invalid address, offset, or register count".to_owned() + ) + ); + } + let u32_ty: TypeHandle = IntegerType::get(ctx, 32, Signedness::Unsigned).into(); + let (address, last_op) = rvalue::translate_operand( + ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(), + )?; + let (half_split_offset, last_op) = if has_half_split_offset { + let (offset, translated) = rvalue::translate_operand( + ctx, body, &args[1], value_map, block_ptr, last_op, loc.clone(), + )?; + if offset + .defining_op() + .and_then(|op| Operation::get_op::(op, ctx)) + .is_none() + { + return input_err!( + loc, + TranslationErr::unsupported( + "tcgen05 16x32bx2 half-split offset must lower to a constant".to_owned() + ) + ); + } + (Some(offset), translated) + } else { + (None, last_op) + }; + let data_index = if has_half_split_offset { 2 } else { 1 }; + let (data, mut last_op) = rvalue::translate_operand( + ctx, body, &args[data_index], value_map, block_ptr, last_op, loc.clone(), + )?; + if expected_len == 1 { + if data.get_type(ctx) != u32_ty { + return input_err!( + loc, + TranslationErr::unsupported( + "generated tcgen05 store data must be a u32 register".to_owned() + ) + ); + } + let mut operands = vec![address]; + operands.extend(half_split_offset); + operands.push(data); + return Ok((operands, last_op)); + } + + let data_ty = data.get_type(ctx); + let direct_array = data_ty + .deref(ctx) + .downcast_ref::() + .is_some_and(|array| { + array.size() == expected_len as u64 && array.element_type() == u32_ty + }); + let array = if direct_array { + data + } else { + let wrapped_array_ty = data_ty + .deref(ctx) + .downcast_ref::() + .filter(|structure| structure.field_count() == 1) + .and_then(|structure| structure.get_field_type(0)); + let valid_wrapper = wrapped_array_ty.is_some_and(|field_ty| { + field_ty + .deref(ctx) + .downcast_ref::() + .is_some_and(|array| { + array.size() == expected_len as u64 && array.element_type() == u32_ty + }) + }); + if !valid_wrapper { + return input_err!( + loc, + TranslationErr::unsupported(format!( + "generated tcgen05 store data must contain {expected_len} u32 registers" + )) + ); + } + let array_ty = wrapped_array_ty.expect("validated tcgen05 store wrapper"); + let extract = Operation::new( + ctx, + MirExtractFieldOp::get_concrete_op_info(), + vec![array_ty], + vec![data], + vec![], + 0, + ); + extract.deref_mut(ctx).set_loc(loc.clone()); + let extract = MirExtractFieldOp::new(extract); + extract.set_attr_index(ctx, FieldIndexAttr(0)); + helpers::insert_op(ctx, extract.get_operation(), block_ptr, last_op); + last_op = Some(extract.get_operation()); + extract.get_operation().deref(ctx).get_result(0) + }; + + let mut operands = Vec::with_capacity(expected_len + 1 + usize::from(has_half_split_offset)); + operands.push(address); + operands.extend(half_split_offset); + for index in 0..expected_len { + let extract = Operation::new( + ctx, + MirExtractFieldOp::get_concrete_op_info(), + vec![u32_ty], + vec![array], + vec![], + 0, + ); + extract.deref_mut(ctx).set_loc(loc.clone()); + let extract = MirExtractFieldOp::new(extract); + extract.set_attr_index(ctx, FieldIndexAttr(index as u32)); + helpers::insert_op(ctx, extract.get_operation(), block_ptr, last_op); + last_op = Some(extract.get_operation()); + operands.push(extract.get_operation().deref(ctx).get_result(0)); + } + Ok((operands, last_op)) +} +"#, + ); + } + if register_mmas(catalog).next().is_some() || sparse_mmas(catalog).next().is_some() { + output.push_str( + r#" + +#[derive(Clone, Copy)] +enum GeneratedMmaImportAdapter { + C2U32A2U32B1U32ToD2U32, + C2U32A4U32B2U32ToD2U32, + C2U32A4U32B4U32ToD2U32, + C4F32A2U32B1U32ToD4F32, + C4F32A4U32B2U32ToD4F32, + C4F32A4U32B4U32ToD4F32, + C2F64A1F64B1F64ToD2F64, + C2I32A1U32B1U32ToD2I32, + C4I32A4U32B2U32ToD4I32, + C4I32A4U32B4U32ToD4I32, + C4I32A2U32B1U32ToD4I32, + C4I32A2U32B2U32ToD4I32, + C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32, +} + +#[allow(clippy::too_many_arguments)] +fn extract_generated_mma_array( + ctx: &mut Context, + array: Value, + expected_element_ty: TypeHandle, + expected_len: usize, + block_ptr: Ptr, + mut last_op: Option>, + loc: Location, +) -> TranslationResult<(Vec, Ptr)> { + let valid = array + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_some_and(|array_ty| { + array_ty.size() == expected_len as u64 + && array_ty.element_type() == expected_element_ty + }); + if !valid { + return input_err!( + loc, + TranslationErr::unsupported(format!( + "generated MMA fragment must be an array of {expected_len} scalar registers" + )) + ); + } + let mut registers = Vec::with_capacity(expected_len); + for index in 0..expected_len { + let extract = Operation::new( + ctx, + MirExtractFieldOp::get_concrete_op_info(), + vec![expected_element_ty], + vec![array], + vec![], + 0, + ); + extract.deref_mut(ctx).set_loc(loc.clone()); + let extract = MirExtractFieldOp::new(extract); + extract.set_attr_index(ctx, FieldIndexAttr(index as u32)); + helpers::insert_op(ctx, extract.get_operation(), block_ptr, last_op); + last_op = Some(extract.get_operation()); + registers.push(extract.get_operation().deref(ctx).get_result(0)); + } + Ok((registers, last_op.expect("non-empty generated MMA fragment"))) +} + +#[allow(clippy::too_many_arguments)] +fn import_generated_mma_operands( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + block_ptr: Ptr, + prev_op: Option>, + value_map: &mut ValueMap, + loc: Location, + adapter: GeneratedMmaImportAdapter, +) -> TranslationResult<(Vec, Ptr, TypeHandle, usize)> { + let f32_ty: TypeHandle = FP32Type::get(ctx).into(); + let f64_ty: TypeHandle = FP64Type::get(ctx).into(); + let i32_ty: TypeHandle = IntegerType::get(ctx, 32, Signedness::Signed).into(); + let u32_ty: TypeHandle = IntegerType::get(ctx, 32, Signedness::Unsigned).into(); + let (c_ty, c_count, a_ty, a_count, a_array, b_ty, b_count, b_array, result_ty, result_count) = + match adapter { + GeneratedMmaImportAdapter::C2U32A2U32B1U32ToD2U32 => + (u32_ty, 2, u32_ty, 2, true, u32_ty, 1, false, u32_ty, 2), + GeneratedMmaImportAdapter::C2U32A4U32B2U32ToD2U32 => + (u32_ty, 2, u32_ty, 4, true, u32_ty, 2, true, u32_ty, 2), + GeneratedMmaImportAdapter::C2U32A4U32B4U32ToD2U32 => + (u32_ty, 2, u32_ty, 4, true, u32_ty, 4, true, u32_ty, 2), + GeneratedMmaImportAdapter::C4F32A2U32B1U32ToD4F32 => + (f32_ty, 4, u32_ty, 2, true, u32_ty, 1, false, f32_ty, 4), + GeneratedMmaImportAdapter::C4F32A4U32B2U32ToD4F32 => + (f32_ty, 4, u32_ty, 4, true, u32_ty, 2, true, f32_ty, 4), + GeneratedMmaImportAdapter::C4F32A4U32B4U32ToD4F32 => + (f32_ty, 4, u32_ty, 4, true, u32_ty, 4, true, f32_ty, 4), + GeneratedMmaImportAdapter::C2F64A1F64B1F64ToD2F64 => + (f64_ty, 2, f64_ty, 1, false, f64_ty, 1, false, f64_ty, 2), + GeneratedMmaImportAdapter::C2I32A1U32B1U32ToD2I32 => + (i32_ty, 2, u32_ty, 1, false, u32_ty, 1, false, i32_ty, 2), + GeneratedMmaImportAdapter::C4I32A4U32B2U32ToD4I32 => + (i32_ty, 4, u32_ty, 4, true, u32_ty, 2, true, i32_ty, 4), + GeneratedMmaImportAdapter::C4I32A4U32B4U32ToD4I32 => + (i32_ty, 4, u32_ty, 4, true, u32_ty, 4, true, i32_ty, 4), + GeneratedMmaImportAdapter::C4I32A2U32B1U32ToD4I32 => + (i32_ty, 4, u32_ty, 2, true, u32_ty, 1, false, i32_ty, 4), + GeneratedMmaImportAdapter::C4I32A2U32B2U32ToD4I32 => + (i32_ty, 4, u32_ty, 2, true, u32_ty, 2, true, i32_ty, 4), + GeneratedMmaImportAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 => + (f32_ty, 4, u32_ty, 4, true, u32_ty, 2, true, f32_ty, 4), + }; + let (c_array, last_op) = rvalue::translate_operand( + ctx, body, &args[0], value_map, block_ptr, prev_op, loc.clone(), + )?; + let (mut operands, last_op) = extract_generated_mma_array( + ctx, c_array, c_ty, c_count, block_ptr, last_op, loc.clone(), + )?; + let (a_value, last_after_a) = rvalue::translate_operand( + ctx, body, &args[1], value_map, block_ptr, Some(last_op), loc.clone(), + )?; + let (a_registers, last_op) = if a_array { + extract_generated_mma_array( + ctx, a_value, a_ty, a_count, block_ptr, last_after_a, loc.clone(), + )? + } else { + (vec![a_value], last_after_a.expect("generated MMA A translation keeps predecessor")) + }; + operands.extend(a_registers); + let (b_value, last_after_b) = rvalue::translate_operand( + ctx, body, &args[2], value_map, block_ptr, Some(last_op), loc.clone(), + )?; + let (b_registers, last_op) = if b_array { + extract_generated_mma_array( + ctx, b_value, b_ty, b_count, block_ptr, last_after_b, loc.clone(), + )? + } else { + (vec![b_value], last_after_b.expect("generated MMA B translation keeps predecessor")) + }; + operands.extend(b_registers); + let mut last_op = last_op; + if matches!( + adapter, + GeneratedMmaImportAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 + ) { + for arg in &args[3..] { + let (value, translated) = rvalue::translate_operand( + ctx, + body, + arg, + value_map, + block_ptr, + Some(last_op), + loc.clone(), + )?; + last_op = translated.expect("generated block-scale MMA operand translation keeps predecessor"); + operands.push(value); + } + } + Ok((operands, last_op, result_ty, result_count)) +} + +fn bundle_generated_mma_results( + ctx: &mut Context, + mma: Ptr, + result_ty: TypeHandle, + result_count: usize, + loc: Location, +) -> (Value, Ptr) { + let results = (0..result_count) + .map(|index| mma.deref(ctx).get_result(index)) + .collect(); + let array_ty = MirArrayType::get(ctx, result_ty, result_count as u64); + let array = Operation::new( + ctx, + MirConstructArrayOp::get_concrete_op_info(), + vec![array_ty.into()], + results, + vec![], + 0, + ); + array.deref_mut(ctx).set_loc(loc); + array.insert_after(ctx, mma); + (array.deref(ctx).get_result(0), array) +} +"#, + ); + } + output.push_str("\n#[cfg(test)]\nmod tests {\n use super::*;\n\n"); + for record in &catalog.intrinsics { + writeln!( + output, + " #[test]\n fn {}_uses_only_canonical_or_compatibility_defpaths() {{", + record.id + ) + .unwrap(); + writeln!( + output, + " assert!(is_generated_intrinsic_path({:?}));", + record.rust.canonical_path + ) + .unwrap(); + writeln!( + output, + " assert!(is_raw_generated_intrinsic_path({:?}));", + record.rust.canonical_path + ) + .unwrap(); + writeln!( + output, + " assert_eq!(generated_intrinsic_marker({:?}), Some({:?}));", + record.rust.canonical_path, + intrinsic_marker(catalog, record) + ) + .unwrap(); + writeln!( + output, + " assert!(matches!(classify_raw_intrinsic_path(\"cuda_intrinsics\", {:?}.into()), RawIntrinsicIdentity::Known(_)));", + record.rust.canonical_path + ) + .unwrap(); + writeln!( + output, + " assert!(!is_generated_intrinsic_path({:?}));", + record.rust.public_path + ) + .unwrap(); + for compatibility_path in &record.rust.compatibility_paths { + writeln!( + output, + " assert!(is_generated_intrinsic_path({compatibility_path:?}));\n assert!(!is_raw_generated_intrinsic_path({compatibility_path:?}));\n assert_eq!(generated_intrinsic_marker({compatibility_path:?}), Some({:?}));", + intrinsic_marker(catalog, record) + ) + .unwrap(); + } + writeln!( + output, + " assert!(!is_generated_intrinsic_path(\"cuda_intrinsics::__cuda_oxide_intrinsic_abi_v{}::{}\"));", + catalog.intrinsic_abi + 1, + record.rust.abi_id + ) + .unwrap(); + output.push_str(" }\n"); + } + output.push_str( + "\n #[test]\n fn raw_intrinsic_identity_classification_fails_closed() {\n assert_eq!(\n classify_raw_intrinsic_path(\"serde\", \"serde::helper\".into()),\n RawIntrinsicIdentity::NotRawCrate\n );\n assert!(matches!(\n classify_raw_intrinsic_path(\"cuda_intrinsics\", \"cuda_intrinsics::__cuda_oxide_intrinsic_abi_v2::i0001\".into()),\n RawIntrinsicIdentity::UnsupportedAbi(_)\n ));\n assert!(matches!(\n classify_raw_intrinsic_path(\"cuda_intrinsics\", \"cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i9999\".into()),\n RawIntrinsicIdentity::UnknownId(_)\n ));\n assert!(matches!(\n classify_raw_intrinsic_path(\"cuda_intrinsics\", \"cuda_intrinsics::helper\".into()),\n RawIntrinsicIdentity::UnknownId(_)\n ));\n }\n", + ); + output.push_str("}\n"); + output +} diff --git a/crates/cuda-intrinsics-gen/src/render/lowering.rs b/crates/cuda-intrinsics-gen/src/render/lowering.rs new file mode 100644 index 0000000000..6aa1b0e6ed --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/render/lowering.rs @@ -0,0 +1,2152 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + ActiveMaskAdapter, BackendLoweringMechanism, CatalogFile, ClcOperation, ClusterBarrierMode, + ClusterMemoryOperation, CpAsyncCachePolicy, CpAsyncControlOperation, CpAsyncMbarrierOperation, + CpAsyncMbarrierStateSpace, CpAsyncSourceSize, DebugControlOperation, DotProductAdapter, + ExecutionControlOperation, IntrinsicBackend, MbarrierBasicAdapter, MbarrierBasicOperation, + MbarrierExtendedAdapter, PackedConversionAdapter, PackedConversionSourceFormat, PrmtAdapter, + PrmtMode, ReduxAdapter, ScalarArithmeticFormat, ScalarMathFormat, StmatrixLayout, + Tcgen05Operation, TmaOperation, TmaReductionLoadMode, TmaReductionOperation, VoteAdapter, + WarpBarrierAdapter, WarpMatchAdapter, WarpShuffleAdapter, WarpShuffleMode, + WarpShuffleValueKind, WgmmaControlMode, +}; +use crate::render::common::rust_header; +use crate::render::families::{ + active_masks, clc_intrinsics, cluster_barrier_attr, cluster_barrier_template, cluster_barriers, + cluster_memory, cp_async_controls, cp_async_copies, cp_async_mbarriers, debug_controls, + dot_product_ptx, dot_products, elect_intrinsics, execution_controls, expected_ptx_head, + extended_minmax, extended_minmax_carrier, extended_minmax_format_attr, + extended_minmax_nan_attr, extended_minmax_operation_attr, extended_minmax_ptx_mnemonic, + extended_minmax_subnormal_attr, extended_minmax_xorsign_abs_attr, integer_minmax_ptx_mnemonic, + integer_minmaxes, ldmatrix, ldmatrix_attr_variants, ldmatrix_compat_op, mbarrier_basics, + mbarrier_extended, movmatrix, movmatrix_template, packed_alu_ptx_mnemonic, packed_alu_width, + packed_alus, packed_atomics, packed_conversion_ptx_mnemonic, packed_conversion_result_width, + packed_conversion_source_width, packed_conversion_typed_llvm_name, packed_conversions, prmts, + redux, register_mma_attr_variants, register_mma_compat_op_type, register_mma_constraints, + register_mma_extra_operand_count, register_mma_fragment_counts, register_mma_result_variant, + register_mma_template, register_mmas, scalar_arithmetic_contract, + scalar_arithmetic_format_attr, scalar_arithmetic_llvm_mechanism, + scalar_arithmetic_operation_attr, scalar_arithmetic_ptx_mnemonic, + scalar_arithmetic_rounding_attr, scalar_arithmetic_saturation_attr, + scalar_arithmetic_subnormal_attr, scalar_arithmetics, scalar_conversion_ptx_mnemonic, + scalar_conversion_rounding_attr, scalar_conversion_saturation_attr, scalar_conversions, + scalar_math_contract, scalar_math_format_attr, scalar_math_llvm_mechanism, + scalar_math_operation_attr, scalar_math_precision_attr, scalar_math_ptx_mnemonic, + scalar_math_subnormal_attr, scalar_maths, sparse_mma_attr_variants, sparse_mma_constraints, + sparse_mma_fragment_counts, sparse_mma_result_variant, sparse_mma_template, sparse_mmas, + special_register_asm_kind, special_register_backend_mechanism, + special_register_inline_template, special_register_output_constraint, sregs, stmatrices, + stmatrix_variant, sync_intrinsics, tcgen05_inline_asm, tcgen05_intrinsics, + tcgen05_mma_intrinsics, tcgen05_non_mma_intrinsics, threadfence_ptx_level, tma_intrinsics, + vote_intrinsics, warp_barriers, warp_matches, warp_shuffles, wgmma_control_template, + wgmma_controls, +}; +use std::fmt::Write as _; + +pub(super) fn render_lowering(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + output.push_str( + "//! Generated conversion interfaces for admitted CUDA intrinsic families.\n\nuse crate::conversion_interface::MirToLlvmConversion;\nuse crate::convert::intrinsics::{atomic::convert_packed_atom_add, basic::convert_sreg_read_inline, common::{call_intrinsic, create_i32_const, inline_asm_convergent}, cp_async::{convert_generated_cp_async_control, convert_generated_cp_async_copy, convert_generated_cp_async_mbarrier}, dotprod::convert_generated_dot_product, ldmatrix::convert_generated_ldmatrix, mbarrier::{convert_arrive, convert_init, convert_inval, convert_test_wait}, packed::{convert_generated_packed_alu, convert_generated_packed_f32x2, convert_generated_packed_unary}, prmt::convert_generated_prmt, warp::{convert_active_mask, convert_bar_warp_sync, convert_match_all, convert_match_any, convert_redux, convert_shuffle_f32, convert_shuffle_i32, convert_shuffle_i64, convert_vote}, wmma::{convert_generated_register_mma, convert_generated_sparse_mma, GeneratedMmaResultType}};\nuse crate::{context, IntrinsicBackend};\nuse dialect_nvvm::ops::{", + ); + if debug_controls(catalog).next().is_some() { + output = output.replace( + "inline_asm_convergent}", + "inline_asm_convergent, inline_asm_sideeffect}", + ); + } + if integer_minmaxes(catalog).next().is_some() { + output = output.replace( + "ldmatrix::convert_generated_ldmatrix, ", + "integer_minmax::convert_generated_integer_minmax, ldmatrix::convert_generated_ldmatrix, ", + ); + } + if stmatrices(catalog).next().is_some() { + output = output.replace( + "common::{call_intrinsic,", + "common::{call_intrinsic, cast_to_shared_addrspace,", + ); + } + if clc_intrinsics(catalog).next().is_some() { + output = output.replace( + "cp_async::{", + "clc::{convert_generated_clc_query, convert_generated_clc_try_cancel}, cp_async::{", + ); + } + if tma_intrinsics(catalog).next().is_some() { + output = output.replace( + "prmt::convert_generated_prmt, warp::{", + "prmt::convert_generated_prmt, tma::{convert_control, convert_g2s, convert_g2s_multicast_cg2, convert_prefetch_tensormap, convert_prefetch_tile, convert_reduce_s2g, convert_s2g, convert_tensormap_fence, convert_tensormap_replace, PrefetchTileConfig, ReduceConfig}, warp::{", + ); + } + if execution_controls(catalog).next().is_some() { + output = output.replace( + "dotprod::convert_generated_dot_product, ", + "dotprod::convert_generated_dot_product, execution_control::{convert_counted_barrier, convert_grid_dependency, convert_setmaxnreg}, ", + ); + } + if scalar_conversions(catalog).next().is_some() { + output = output.replace( + "prmt::convert_generated_prmt, ", + "prmt::convert_generated_prmt, scalar_conversion::convert_generated_scalar_conversion, ", + ); + } + if scalar_arithmetics(catalog).next().is_some() { + output = output.replace( + "prmt::convert_generated_prmt, ", + "prmt::convert_generated_prmt, scalar_arithmetic::convert_generated_scalar_arithmetic, ", + ); + } + if scalar_maths(catalog).next().is_some() { + output = output.replace( + "prmt::convert_generated_prmt, ", + "prmt::convert_generated_prmt, scalar_math::convert_generated_scalar_math, ", + ); + } + if extended_minmax(catalog).next().is_some() { + output = output.replace( + "dotprod::convert_generated_dot_product, ", + "dotprod::convert_generated_dot_product, extended_minmax::{MinMaxCarrier, convert_generated_extended_minmax}, ", + ); + } + if elect_intrinsics(catalog).next().is_some() { + output = output.replace( + "warp::{convert_active_mask,", + "warp::{convert_active_mask, convert_elect_sync_inline, convert_elect_sync_typed,", + ); + } + for (index, record) in sregs(catalog).enumerate() { + if index != 0 { + output.push_str(", "); + } + output.push_str(&record.dialect.op_type); + } + if ldmatrix(catalog).next().is_some() { + if sregs(catalog).next().is_some() { + output.push_str(", "); + } + output.push_str("LdmatrixElementAttr, LdmatrixLayoutAttr, LdmatrixMultiplicityAttr, LdmatrixOp, LdmatrixShapeAttr, LdmatrixStateSpaceAttr"); + for record in ldmatrix(catalog) { + if let Some((op_type, _)) = ldmatrix_compat_op(record) { + output.push_str(", "); + output.push_str(op_type); + } + } + } + if let Some(record) = movmatrix(catalog).next() { + output.push_str(", "); + output.push_str(&record.dialect.op_type); + } + if register_mmas(catalog).next().is_some() { + for record in + register_mmas(catalog).filter(|record| register_mma_compat_op_type(record).is_some()) + { + output.push_str(", "); + output.push_str(register_mma_compat_op_type(record).unwrap()); + } + output.push_str(", RegisterMmaAccumulatorAttr, RegisterMmaElementAttr, RegisterMmaKindAttr, RegisterMmaLayoutAttr, RegisterMmaOp, RegisterMmaOperationAttr, RegisterMmaOverflowAttr, RegisterMmaShapeAttr"); + } + if sparse_mmas(catalog).next().is_some() { + output.push_str(", SparseMmaAccumulatorAttr, SparseMmaElementAttr, SparseMmaLayoutAttr, SparseMmaMetadataAttr, SparseMmaOp, SparseMmaOverflowAttr, SparseMmaSelectorAttr, SparseMmaShapeAttr"); + } + if prmts(catalog).next().is_some() { + output.push_str(", PrmtModeAttr, PrmtOp"); + } + if scalar_conversions(catalog).next().is_some() { + output.push_str( + ", ScalarConversionOp, ScalarConversionRoundingAttr, ScalarConversionSaturationAttr", + ); + } + if scalar_arithmetics(catalog).next().is_some() { + output.push_str( + ", ScalarArithmeticFormatAttr, ScalarArithmeticOp, ScalarArithmeticOperationAttr, ScalarArithmeticRoundingAttr, ScalarArithmeticSaturationAttr, ScalarArithmeticSubnormalAttr", + ); + } + if scalar_maths(catalog).next().is_some() { + output.push_str( + ", ScalarMathFormatAttr, ScalarMathOp, ScalarMathOperationAttr, ScalarMathPrecisionAttr, ScalarMathSubnormalAttr", + ); + } + if extended_minmax(catalog).next().is_some() { + output.push_str( + ", ExtendedMinMaxFormatAttr, ExtendedMinMaxNanAttr, ExtendedMinMaxOp, ExtendedMinMaxOperationAttr, ExtendedMinMaxSubnormalAttr, ExtendedMinMaxXorSignAbsAttr", + ); + } + if cluster_barriers(catalog).next().is_some() { + output.push_str(", ClusterBarrierModeAttr, ClusterBarrierOp, ClusterSyncOp"); + } + if debug_controls(catalog).next().is_some() { + output.push_str(", BreakpointOp, PmEventOp, TrapOp"); + } + for record in clc_intrinsics(catalog) { + output.push_str(", "); + output.push_str(&record.dialect.op_type); + } + if wgmma_controls(catalog).next().is_some() { + output.push_str( + ", WgmmaCommitGroupSyncAlignedOp, WgmmaFenceSyncAlignedOp, WgmmaWaitGroupSyncAlignedOp", + ); + } + for record in tcgen05_non_mma_intrinsics(catalog) { + output.push_str(", "); + output.push_str(&record.dialect.op_type); + } + if tcgen05_mma_intrinsics(catalog).next().is_some() { + output.push_str(", Tcgen05MmaBBufferAttr, Tcgen05MmaBUsageAttr, Tcgen05MmaCollectorAAttr, Tcgen05MmaCtaGroupAttr, Tcgen05MmaFormAttr, Tcgen05MmaKindAttr, Tcgen05MmaOp"); + } + if packed_atomics(catalog).next().is_some() { + if sregs(catalog).next().is_some() || ldmatrix(catalog).next().is_some() { + output.push_str(", "); + } + output.push_str("NvvmAtomAddBf16x2Op, NvvmAtomAddF16x2Op, PackedAtomicAddOp, PackedAtomicAtomicityAttr, PackedAtomicFormatAttr, PackedAtomicOrderingAttr, PackedAtomicRoundingAttr, PackedAtomicScopeAttr, PackedAtomicStateSpaceAttr, PackedAtomicSubnormalAttr"); + } + if redux(catalog).next().is_some() { + if sregs(catalog).next().is_some() + || ldmatrix(catalog).next().is_some() + || packed_atomics(catalog).next().is_some() + { + output.push_str(", "); + } + for (index, record) in redux(catalog).enumerate() { + if index != 0 { + output.push_str(", "); + } + output.push_str(&record.dialect.op_type); + } + } + if vote_intrinsics(catalog).next().is_some() { + if sregs(catalog).next().is_some() + || ldmatrix(catalog).next().is_some() + || packed_atomics(catalog).next().is_some() + || redux(catalog).next().is_some() + { + output.push_str(", "); + } + for (index, record) in vote_intrinsics(catalog).enumerate() { + if index != 0 { + output.push_str(", "); + } + output.push_str(&record.dialect.op_type); + } + } + if active_masks(catalog).next().is_some() { + output.push_str(", "); + for (index, record) in active_masks(catalog).enumerate() { + if index != 0 { + output.push_str(", "); + } + output.push_str(&record.dialect.op_type); + } + } + if warp_matches(catalog).next().is_some() { + output.push_str(", "); + for (index, record) in warp_matches(catalog).enumerate() { + if index != 0 { + output.push_str(", "); + } + output.push_str(&record.dialect.op_type); + } + } + for record in elect_intrinsics(catalog) { + output.push_str(", "); + output.push_str(&record.dialect.op_type); + } + if warp_barriers(catalog).next().is_some() { + output.push_str(", "); + for (index, record) in warp_barriers(catalog).enumerate() { + if index != 0 { + output.push_str(", "); + } + output.push_str(&record.dialect.op_type); + } + } + if dot_products(catalog).next().is_some() { + if sregs(catalog).next().is_some() + || ldmatrix(catalog).next().is_some() + || packed_atomics(catalog).next().is_some() + || redux(catalog).next().is_some() + || vote_intrinsics(catalog).next().is_some() + { + output.push_str(", "); + } + for (index, record) in dot_products(catalog).enumerate() { + if index != 0 { + output.push_str(", "); + } + output.push_str(&record.dialect.op_type); + } + } + if sync_intrinsics(catalog).next().is_some() { + if sregs(catalog).next().is_some() + || ldmatrix(catalog).next().is_some() + || packed_atomics(catalog).next().is_some() + || redux(catalog).next().is_some() + || vote_intrinsics(catalog).next().is_some() + || dot_products(catalog).next().is_some() + { + output.push_str(", "); + } + for (index, record) in sync_intrinsics(catalog).enumerate() { + if index != 0 { + output.push_str(", "); + } + output.push_str(&record.dialect.op_type); + } + } + if warp_shuffles(catalog).next().is_some() { + output.push_str(", "); + for (index, record) in warp_shuffles(catalog).enumerate() { + if index != 0 { + output.push_str(", "); + } + output.push_str(&record.dialect.op_type); + } + } + if packed_alus(catalog).next().is_some() { + if catalog + .intrinsics + .iter() + .any(|record| record.family != "packed_alu") + { + output.push_str(", "); + } + for (index, record) in packed_alus(catalog).enumerate() { + if index != 0 { + output.push_str(", "); + } + output.push_str(&record.dialect.op_type); + } + } + if integer_minmaxes(catalog).next().is_some() { + // packed_alu records always exist, so a separator is always needed. + output.push_str(", "); + for (index, record) in integer_minmaxes(catalog).enumerate() { + if index != 0 { + output.push_str(", "); + } + output.push_str(&record.dialect.op_type); + } + } + if packed_conversions(catalog).next().is_some() { + if catalog + .intrinsics + .iter() + .any(|record| record.family != "packed_conversion") + { + output.push_str(", "); + } + for (index, record) in packed_conversions(catalog).enumerate() { + if index != 0 { + output.push_str(", "); + } + output.push_str(&record.dialect.op_type); + } + } + for record in cp_async_copies(catalog) + .chain(cp_async_controls(catalog)) + .chain(cp_async_mbarriers(catalog)) + { + output.push_str(", "); + output.push_str(&record.dialect.op_type); + } + for record in mbarrier_basics(catalog) { + output.push_str(", "); + output.push_str(&record.dialect.op_type); + } + for record in cluster_memory(catalog) { + output.push_str(", "); + output.push_str(&record.dialect.op_type); + } + for record in stmatrices(catalog) { + output.push_str(", "); + output.push_str(&record.dialect.op_type); + } + for record in mbarrier_extended(catalog) { + output.push_str(", "); + output.push_str(&record.dialect.op_type); + } + for record in tma_intrinsics(catalog) { + output.push_str(", "); + output.push_str(&record.dialect.op_type); + } + for record in execution_controls(catalog) { + output.push_str(", "); + output.push_str(&record.dialect.op_type); + } + output.push_str( + "};\nuse llvm_export::{ops::AsmKind, types as llvm_types};\nuse pliron::{\n builtin::types::{IntegerType, Signedness},\n context::{Context, Ptr},\n derive::op_interface_impl,\n irbuild::{\n dialect_conversion::{DialectConversionRewriter, OperandsInfo},\n rewriter::Rewriter,\n },\n op::Op,\n operation::Operation,\n result::Result,\n};\n\n", + ); + if cluster_memory(catalog).next().is_some() { + if !output.contains("cast_to_shared_addrspace") { + output = output.replace( + "common::{call_intrinsic, ", + "common::{call_intrinsic, cast_to_shared_addrspace, ", + ); + } + output = output + .replace( + "use llvm_export::{ops::AsmKind, types as llvm_types};", + "use llvm_export::{op_interfaces::CastOpInterface, ops as llvm_ops, ops::AsmKind, types as llvm_types};", + ) + .replace( + "dialect_conversion::{DialectConversionRewriter, OperandsInfo},\n rewriter::Rewriter,", + "dialect_conversion::{DialectConversionRewriter, OperandsInfo},\n inserter::Inserter,\n rewriter::Rewriter,", + ); + } + if tcgen05_intrinsics(catalog).next().is_some() { + output = output.replace( + "use llvm_export::{ops::AsmKind, types as llvm_types};", + "use llvm_export::{ops as llvm_ops, ops::AsmKind, types as llvm_types};", + ); + output = output.replace( + "builtin::types::{IntegerType, Signedness}", + "builtin::types::{FP32Type, IntegerType, Signedness}", + ); + output = output.replace( + "dialect_conversion::{DialectConversionRewriter, OperandsInfo},\n rewriter::Rewriter,", + "dialect_conversion::{DialectConversionRewriter, OperandsInfo},\n inserter::Inserter,\n rewriter::Rewriter,", + ); + } + output.push_str("use pliron::location::Located;\n\n"); + if mbarrier_extended(catalog).next().is_some() { + if !output.contains("cast_to_shared_addrspace") { + output = output.replace( + "common::{call_intrinsic,", + "common::{call_intrinsic, cast_to_shared_addrspace,", + ); + } + output = if output.contains("inline_asm_sideeffect}") { + output.replace( + "inline_asm_sideeffect}", + "inline_asm_sideeffect, trunc_to_i1}", + ) + } else { + output.replace( + "inline_asm_convergent}", + "inline_asm_convergent, trunc_to_i1}", + ) + }; + output.push_str("use pliron::value::DefiningEntity;\n\n"); + } + output.push_str( + "fn convert_zero_operand_scalar_direct(\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n op: Ptr,\n width: u32,\n intrinsic_name: &str,\n) -> Result<()> {\n let result_ty = IntegerType::get(ctx, width, Signedness::Signless);\n let function_ty = llvm_types::FuncType::get(ctx, result_ty.into(), vec![], false);\n let call = call_intrinsic(ctx, rewriter, op, intrinsic_name, function_ty, vec![])?;\n rewriter.replace_operation(ctx, op, call);\n Ok(())\n}\n\n", + ); + if tcgen05_intrinsics(catalog).next().is_some() { + output.push_str( + r#"fn convert_generated_tcgen05_void( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + arity: usize, + template: &str, + constraints: &str, +) -> Result<()> { + let operands: Vec<_> = op.deref(ctx).operands().collect(); + if operands.len() != arity { + return pliron::input_err_noloc!( + "generated tcgen05 operation expects {arity} operands, got {}", + operands.len() + ); + } + let void_ty = llvm_types::VoidType::get(ctx); + inline_asm_convergent( + ctx, + rewriter, + op, + void_ty.into(), + operands, + template, + constraints, + ); + rewriter.erase_operation(ctx, op); + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +fn convert_generated_tcgen05_load( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + arity: usize, + count: usize, + integer_results: bool, + template: &str, + constraints: &str, +) -> Result<()> { + let operands: Vec<_> = op.deref(ctx).operands().collect(); + if operands.len() != arity || !(1..=2).contains(&arity) || !(1..=128).contains(&count) || op.deref(ctx).get_num_results() != count { + return pliron::input_err_noloc!( + "generated tcgen05 load requires 1..=2 operands and 1..=128 results" + ); + } + let scalar_ty: pliron::r#type::TypeHandle = if integer_results { + IntegerType::get(ctx, 32, Signedness::Signless).into() + } else { + FP32Type::get(ctx).into() + }; + if count == 1 { + let inline_asm = inline_asm_convergent( + ctx, + rewriter, + op, + scalar_ty, + operands, + template, + constraints, + ); + let result = inline_asm.deref(ctx).get_result(0); + rewriter.replace_operation_with_values(ctx, op, vec![result]); + return Ok(()); + } + let result_ty = llvm_types::StructType::get_unnamed( + ctx, + ( + (0..count).map(|_| scalar_ty).collect(), + llvm_types::StructLayout::Unpacked, + ), + ); + let inline_asm = inline_asm_convergent( + ctx, + rewriter, + op, + result_ty.into(), + operands, + template, + constraints, + ); + let result = inline_asm.deref(ctx).get_result(0); + let mut values = Vec::with_capacity(count); + for index in 0..count as u32 { + let extract = llvm_ops::ExtractValueOp::new(ctx, result, vec![index]) + .map_err(|error| pliron::input_error_noloc!("{}", error))?; + rewriter.insert_operation(ctx, extract.get_operation()); + values.push(extract.get_operation().deref(ctx).get_result(0)); + } + rewriter.replace_operation_with_values(ctx, op, values); + Ok(()) +} + +"#, + ); + } + if stmatrices(catalog).next().is_some() { + output.push_str( + r#"fn convert_generated_stmatrix( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + register_count: usize, + transposed: bool, + typed_intrinsic_name: &str, +) -> Result<()> { + let operands: Vec<_> = op.deref(ctx).operands().collect(); + if operands.len() != register_count + 1 || !matches!(register_count, 2 | 4) { + return pliron::input_err_noloc!( + "stmatrix requires one pointer and two or four registers" + ); + } + let void_ty = llvm_types::VoidType::get(ctx); + match context::lowering_options(ctx).intrinsic_backend { + IntrinsicBackend::LlvmNvptx => { + let shared = cast_to_shared_addrspace(ctx, rewriter, operands[0]); + let shared_ty = llvm_types::PointerType::get( + ctx, + llvm_types::address_space::SHARED, + ); + let i32_ty = IntegerType::get(ctx, 32, Signedness::Signless); + let argument_types = std::iter::once(shared_ty.into()) + .chain(std::iter::repeat_n(i32_ty.into(), register_count)) + .collect(); + let function_ty = llvm_types::FuncType::get( + ctx, + void_ty.into(), + argument_types, + false, + ); + let arguments = std::iter::once(shared) + .chain(operands.iter().skip(1).copied()) + .collect(); + call_intrinsic( + ctx, + rewriter, + op, + typed_intrinsic_name, + function_ty, + arguments, + )?; + } + IntrinsicBackend::LibNvvm => { + let registers = (1..=register_count) + .map(|index| format!("${index}")) + .collect::>() + .join(", "); + let trans = if transposed { ".trans" } else { "" }; + let template = format!( + "{{ .reg .u64 %ptr64; .reg .u32 %ptr32; cvta.to.shared.u64 %ptr64, $0; cvt.u32.u64 %ptr32, %ptr64; stmatrix.sync.aligned.m8n8.x{register_count}{trans}.shared.b16 [%ptr32], {{{registers}}}; }}" + ); + let constraints = std::iter::once("l") + .chain(std::iter::repeat_n("r", register_count)) + .chain(std::iter::once("~{memory}")) + .collect::>() + .join(","); + inline_asm_convergent( + ctx, + rewriter, + op, + void_ty.into(), + operands, + &template, + &constraints, + ); + } + } + rewriter.erase_operation(ctx, op); + Ok(()) +} + +"#, + ); + } + for record in sregs(catalog) { + writeln!( + output, + "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", + record.dialect.op_type + ) + .unwrap(); + output.push_str( + " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n", + ); + if record.special_register.is_none() { + writeln!( + output, + " convert_zero_operand_scalar_direct(ctx, rewriter, self.get_operation(), {}, {:?})", + record.scalar_width().unwrap(), + record.llvm_identifier() + ) + .unwrap(); + } else { + output.push_str(" match context::lowering_options(ctx).intrinsic_backend {\n"); + for (backend, backend_variant) in [ + (IntrinsicBackend::LlvmNvptx, "LlvmNvptx"), + (IntrinsicBackend::LibNvvm, "LibNvvm"), + ] { + writeln!( + output, + " IntrinsicBackend::{backend_variant} => {{" + ) + .unwrap(); + match special_register_backend_mechanism(record, backend) { + BackendLoweringMechanism::TypedNvvm => writeln!( + output, + " convert_zero_operand_scalar_direct(ctx, rewriter, self.get_operation(), {}, {:?})", + record.scalar_width().unwrap(), + record.llvm_identifier() + ) + .unwrap(), + BackendLoweringMechanism::InlinePtx => writeln!( + output, + " convert_sreg_read_inline(ctx, rewriter, self.get_operation(), {}, {:?}, {:?}, {})", + record.scalar_width().unwrap(), + special_register_inline_template(record), + special_register_output_constraint(record), + special_register_asm_kind(record), + ) + .unwrap(), + } + output.push_str(" }\n"); + } + output.push_str(" }\n"); + } + output.push_str(" }\n}\n\n"); + } + for record in active_masks(catalog) { + debug_assert_eq!( + record.active_mask.as_ref().unwrap().adapter, + ActiveMaskAdapter::DirectZeroOperandMask + ); + writeln!( + output, + "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", + record.dialect.op_type + ) + .unwrap(); + output.push_str( + " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n operands_info: &OperandsInfo,\n ) -> Result<()> {\n let op = self.get_operation();\n match context::lowering_options(ctx).intrinsic_backend {\n IntrinsicBackend::LlvmNvptx => {\n convert_active_mask(ctx, rewriter, op, operands_info)\n }\n IntrinsicBackend::LibNvvm => {\n let i32_ty = IntegerType::get(ctx, 32, Signedness::Signless);\n let inline_asm = inline_asm_convergent(\n ctx,\n rewriter,\n op,\n i32_ty.into(),\n vec![],\n \"activemask.b32 $0;\",\n \"=r,~{memory}\",\n );\n rewriter.replace_operation(ctx, op, inline_asm);\n Ok(())\n }\n }\n }\n}\n\n", + ); + } + if ldmatrix(catalog).next().is_some() { + output.push_str( + "#[op_interface_impl]\nimpl MirToLlvmConversion for LdmatrixOp {\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let recipe = {\n let shape = self.get_attr_nvvm_ldmatrix_shape(ctx);\n let multiplicity = self.get_attr_nvvm_ldmatrix_multiplicity(ctx);\n let layout = self.get_attr_nvvm_ldmatrix_layout(ctx);\n let element = self.get_attr_nvvm_ldmatrix_element(ctx);\n let state_space = self.get_attr_nvvm_ldmatrix_state_space(ctx);\n match (shape.as_deref(), multiplicity.as_deref(), layout.as_deref(), element.as_deref(), state_space.as_deref()) {\n", + ); + for record in ldmatrix(catalog) { + let (shape, multiplicity, layout, element, state_space) = + ldmatrix_attr_variants(record); + let variant = &record.ldmatrix.as_ref().unwrap().variant; + let register_count = variant.register_count(); + let instruction_head = expected_ptx_head(record); + let intrinsic_name = record.resolved_llvm_identifier(); + writeln!( + output, + " (Some(&{shape}), Some(&{multiplicity}), Some(&{layout}), Some(&{element}), Some(&{state_space})) => ({register_count}, {instruction_head:?}, {intrinsic_name:?})," + ) + .unwrap(); + } + output.push_str( + " _ => return pliron::input_err!(\n self.get_operation().deref(ctx).loc(),\n \"nvvm.ldmatrix variant has no generated lowering recipe\",\n ),\n }\n };\n convert_generated_ldmatrix(ctx, rewriter, self.get_operation(), recipe.0, recipe.1, recipe.2)\n }\n}\n\n", + ); + for record in ldmatrix(catalog) { + let Some((op_type, _)) = ldmatrix_compat_op(record) else { + continue; + }; + let variant = &record + .ldmatrix + .as_ref() + .expect("ldmatrix compatibility record") + .variant; + let register_count = variant.register_count(); + let instruction_head = expected_ptx_head(record); + let intrinsic_name = record.resolved_llvm_identifier(); + writeln!( + output, + "#[op_interface_impl]\nimpl MirToLlvmConversion for {op_type} {{\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {{\n convert_generated_ldmatrix(\n ctx, rewriter, self.get_operation(), {register_count}, {instruction_head:?}, {intrinsic_name:?},\n )\n }}\n}}\n" + ) + .unwrap(); + } + } + for record in stmatrices(catalog) { + let (multiplicity, layout) = stmatrix_variant(record).expect("stmatrix variant"); + let count = multiplicity.register_count(); + let transposed = layout == StmatrixLayout::Transposed; + let intrinsic_name = record.resolved_llvm_identifier(); + writeln!( + output, + "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", + record.dialect.op_type + ) + .unwrap(); + writeln!( + output, + " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {{\n convert_generated_stmatrix(\n ctx, rewriter, self.get_operation(), {count}, {transposed}, {intrinsic_name:?},\n )\n }}\n}}\n" + ) + .unwrap(); + } + if let Some(record) = movmatrix(catalog).next() { + writeln!( + output, + "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {{\n let op = self.get_operation();\n let operands: Vec<_> = op.deref(ctx).operands().collect();\n if operands.len() != 1 {{\n return pliron::input_err_noloc!(\n \"movmatrix_trans_b16 requires 1 operand, got {{}}\",\n operands.len()\n );\n }}\n let result_ty = IntegerType::get(ctx, 32, Signedness::Signless);\n let asm = inline_asm_convergent(\n ctx, rewriter, op, result_ty.into(), operands, {:?}, \"=r,r\",\n );\n rewriter.replace_operation(ctx, op, asm);\n Ok(())\n }}\n}}\n", + record.dialect.op_type, + movmatrix_template(record), + ) + .unwrap(); + } + if register_mmas(catalog).next().is_some() { + output.push_str( + "#[op_interface_impl]\nimpl MirToLlvmConversion for RegisterMmaOp {\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let operation = self.operation_or_multiply(ctx);\n let kind = self.kind_or_inferred(ctx);\n let recipe = match (\n self.get_attr_nvvm_register_mma_shape(ctx).as_deref(),\n operation,\n kind,\n self.get_attr_nvvm_register_mma_accumulator(ctx).as_deref(),\n self.get_attr_nvvm_register_mma_a_element(ctx).as_deref(),\n self.get_attr_nvvm_register_mma_b_element(ctx).as_deref(),\n self.get_attr_nvvm_register_mma_a_layout(ctx).as_deref(),\n self.get_attr_nvvm_register_mma_b_layout(ctx).as_deref(),\n self.get_attr_nvvm_register_mma_overflow(ctx).as_deref(),\n ) {\n", + ); + for record in register_mmas(catalog) { + let ( + shape, + operation, + kind, + accumulator, + a_element, + b_element, + a_layout, + b_layout, + overflow, + ) = register_mma_attr_variants(record); + let (c_count, a_count, b_count, result_count) = register_mma_fragment_counts(record); + let expected_operands = + c_count + a_count + b_count + register_mma_extra_operand_count(record); + writeln!( + output, + " (Some(&{shape}), {operation}, {kind}, Some(&{accumulator}), Some(&{a_element}), Some(&{b_element}), Some(&{a_layout}), Some(&{b_layout}), Some(&{overflow})) => ({}, {result_count}, {expected_operands}, {:?}, {:?}),", + register_mma_result_variant(record), + register_mma_template(record), + register_mma_constraints(record), + ) + .unwrap(); + } + output.push_str( + " _ => return pliron::input_err!(\n self.get_operation().deref(ctx).loc(),\n \"nvvm.register_mma variant has no generated lowering recipe\",\n ),\n };\n convert_generated_register_mma(\n ctx, rewriter, self.get_operation(), recipe.0, recipe.1, recipe.2, recipe.3, recipe.4,\n )\n }\n}\n\n", + ); + for record in + register_mmas(catalog).filter(|record| register_mma_compat_op_type(record).is_some()) + { + let op_type = register_mma_compat_op_type(record).unwrap(); + let (c_count, a_count, b_count, result_count) = register_mma_fragment_counts(record); + let expected_operands = c_count + a_count + b_count; + writeln!( + output, + "#[op_interface_impl]\nimpl MirToLlvmConversion for {op_type} {{\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {{\n convert_generated_register_mma(\n ctx, rewriter, self.get_operation(), {}, {result_count}, {expected_operands}, {:?}, {:?},\n )\n }}\n}}\n", + register_mma_result_variant(record), + register_mma_template(record), + register_mma_constraints(record), + ) + .unwrap(); + } + } + if sparse_mmas(catalog).next().is_some() { + output.push_str( + "#[op_interface_impl]\nimpl MirToLlvmConversion for SparseMmaOp {\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let recipe = match (\n self.get_attr_nvvm_sparse_mma_shape(ctx).as_deref(),\n self.get_attr_nvvm_sparse_mma_accumulator(ctx).as_deref(),\n self.get_attr_nvvm_sparse_mma_a_element(ctx).as_deref(),\n self.get_attr_nvvm_sparse_mma_b_element(ctx).as_deref(),\n self.get_attr_nvvm_sparse_mma_a_layout(ctx).as_deref(),\n self.get_attr_nvvm_sparse_mma_b_layout(ctx).as_deref(),\n self.get_attr_nvvm_sparse_mma_overflow(ctx).as_deref(),\n self.get_attr_nvvm_sparse_mma_metadata(ctx).as_deref(),\n self.get_attr_nvvm_sparse_mma_selector(ctx).as_deref(),\n ) {\n", + ); + for record in sparse_mmas(catalog) { + let ( + shape, + accumulator, + a_element, + b_element, + a_layout, + b_layout, + overflow, + metadata, + selector, + ) = sparse_mma_attr_variants(record); + let (c_count, a_count, b_count, result_count) = sparse_mma_fragment_counts(record); + let expected_operands = c_count + a_count + b_count + 2; + writeln!( + output, + " (Some(&{shape}), Some(&{accumulator}), Some(&{a_element}), Some(&{b_element}), Some(&{a_layout}), Some(&{b_layout}), Some(&{overflow}), Some(&{metadata}), Some(&{selector})) => ({}, {result_count}, {expected_operands}, {:?}, {:?}),", + sparse_mma_result_variant(record), + sparse_mma_template(record), + sparse_mma_constraints(record), + ) + .unwrap(); + } + output.push_str( + " _ => return pliron::input_err!(\n self.get_operation().deref(ctx).loc(),\n \"nvvm.sparse_mma variant has no generated lowering recipe\",\n ),\n };\n convert_generated_sparse_mma(\n ctx, rewriter, self.get_operation(), recipe.0, recipe.1, recipe.2, recipe.3, recipe.4,\n )\n }\n}\n\n", + ); + } + if prmts(catalog).next().is_some() { + output.push_str( + "#[op_interface_impl]\nimpl MirToLlvmConversion for PrmtOp {\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let recipe = match self.get_attr_nvvm_prmt_mode(ctx).as_deref() {\n", + ); + for record in prmts(catalog) { + let prmt = record.prmt.as_ref().unwrap(); + let mode = match prmt.mode { + PrmtMode::Generic => "PrmtModeAttr::Generic", + PrmtMode::F4e => "PrmtModeAttr::F4e", + PrmtMode::B4e => "PrmtModeAttr::B4e", + PrmtMode::Rc8 => "PrmtModeAttr::Rc8", + PrmtMode::Ecl => "PrmtModeAttr::Ecl", + PrmtMode::Ecr => "PrmtModeAttr::Ecr", + PrmtMode::Rc16 => "PrmtModeAttr::Rc16", + }; + let modifier = match prmt.mode { + PrmtMode::Generic => "", + PrmtMode::F4e => ".f4e", + PrmtMode::B4e => ".b4e", + PrmtMode::Rc8 => ".rc8", + PrmtMode::Ecl => ".ecl", + PrmtMode::Ecr => ".ecr", + PrmtMode::Rc16 => ".rc16", + }; + let template = match prmt.adapter { + PrmtAdapter::DirectThreeOperands => { + format!("prmt.b32{modifier} $0, $1, $2, $3;") + } + PrmtAdapter::InsertZeroSecondSource => { + format!("prmt.b32{modifier} $0, $1, 0, $2;") + } + }; + writeln!( + output, + " Some(&{mode}) => ({:?}, {:?}),", + record.llvm_identifier(), + template + ) + .unwrap(); + } + output.push_str( + " _ => return pliron::input_err!(\n self.get_operation().deref(ctx).loc(),\n \"nvvm.prmt mode has no generated lowering recipe\",\n ),\n };\n convert_generated_prmt(ctx, rewriter, self.get_operation(), recipe.0, recipe.1)\n }\n}\n\n", + ); + } + if scalar_conversions(catalog).next().is_some() { + output.push_str( + "#[op_interface_impl]\nimpl MirToLlvmConversion for ScalarConversionOp {\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let recipe = match (\n self.get_attr_nvvm_scalar_conversion_rounding(ctx).as_deref(),\n self.get_attr_nvvm_scalar_conversion_saturation(ctx).as_deref(),\n ) {\n", + ); + for record in scalar_conversions(catalog) { + writeln!( + output, + " (Some(&{}), Some(&{})) => ({:?}, {:?}),", + scalar_conversion_rounding_attr(record), + scalar_conversion_saturation_attr(record), + record.llvm_identifier(), + scalar_conversion_ptx_mnemonic(record), + ) + .unwrap(); + } + output.push_str( + " _ => return pliron::input_err!(\n self.get_operation().deref(ctx).loc(),\n \"nvvm.scalar_conversion variant has no generated lowering recipe\",\n ),\n };\n convert_generated_scalar_conversion(\n ctx, rewriter, self.get_operation(), recipe.0, recipe.1,\n )\n }\n}\n\n", + ); + } + if scalar_arithmetics(catalog).next().is_some() { + output.push_str( + "#[op_interface_impl]\nimpl MirToLlvmConversion for ScalarArithmeticOp {\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let recipe = match (\n self.get_attr_nvvm_scalar_arithmetic_format(ctx).as_deref(),\n self.get_attr_nvvm_scalar_arithmetic_operation(ctx).as_deref(),\n self.get_attr_nvvm_scalar_arithmetic_rounding(ctx).as_deref(),\n self.get_attr_nvvm_scalar_arithmetic_subnormal(ctx).as_deref(),\n self.get_attr_nvvm_scalar_arithmetic_saturation(ctx).as_deref(),\n ) {\n", + ); + for record in scalar_arithmetics(catalog) { + writeln!( + output, + " (Some(&{}), Some(&{}), Some(&{}), Some(&{}), Some(&{})) => ({:?}, {:?}, {}, {}),", + scalar_arithmetic_format_attr(record), + scalar_arithmetic_operation_attr(record), + scalar_arithmetic_rounding_attr(record), + scalar_arithmetic_subnormal_attr(record), + scalar_arithmetic_saturation_attr(record), + record.llvm_identifier(), + scalar_arithmetic_ptx_mnemonic(record), + scalar_arithmetic_contract(record).format == ScalarArithmeticFormat::F64, + scalar_arithmetic_llvm_mechanism(record) + == BackendLoweringMechanism::InlinePtx, + ) + .unwrap(); + } + output.push_str( + " _ => return pliron::input_err!(\n self.get_operation().deref(ctx).loc(),\n \"nvvm.scalar_arithmetic variant has no generated lowering recipe\",\n ),\n };\n convert_generated_scalar_arithmetic(\n ctx, rewriter, self.get_operation(), recipe.0, recipe.1, recipe.2, recipe.3,\n )\n }\n}\n\n", + ); + } + if scalar_maths(catalog).next().is_some() { + output.push_str( + "#[op_interface_impl]\nimpl MirToLlvmConversion for ScalarMathOp {\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let recipe = match (\n self.get_attr_nvvm_scalar_math_format(ctx).as_deref(),\n self.get_attr_nvvm_scalar_math_operation(ctx).as_deref(),\n self.get_attr_nvvm_scalar_math_precision(ctx).as_deref(),\n self.get_attr_nvvm_scalar_math_subnormal(ctx).as_deref(),\n ) {\n", + ); + for record in scalar_maths(catalog) { + // PTX-native records carry no LLVM symbol. The identifier is + // only consumed by the typed route, which such records never + // take (their mechanism is always inline PTX), so render an + // empty name; the lowering helper rejects it defensively. + let intrinsic_name = if record.llvm.is_some() { + record.llvm_identifier() + } else { + String::new() + }; + writeln!( + output, + " (Some(&{}), Some(&{}), Some(&{}), Some(&{})) => ({:?}, {:?}, {}, {}, {}),", + scalar_math_format_attr(record), + scalar_math_operation_attr(record), + scalar_math_precision_attr(record), + scalar_math_subnormal_attr(record), + intrinsic_name, + scalar_math_ptx_mnemonic(record), + scalar_math_contract(record).format == ScalarMathFormat::F16, + scalar_math_contract(record).format == ScalarMathFormat::F64, + scalar_math_llvm_mechanism(record) == BackendLoweringMechanism::InlinePtx, + ) + .unwrap(); + } + output.push_str( + " _ => return pliron::input_err!(\n self.get_operation().deref(ctx).loc(),\n \"nvvm.scalar_math variant has no generated lowering recipe\",\n ),\n };\n convert_generated_scalar_math(\n ctx, rewriter, self.get_operation(), recipe.0, recipe.1, recipe.2, recipe.3, recipe.4,\n )\n }\n}\n\n", + ); + } + if extended_minmax(catalog).next().is_some() { + output.push_str( + "#[op_interface_impl]\nimpl MirToLlvmConversion for ExtendedMinMaxOp {\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let recipe = match (\n self.get_attr_nvvm_extended_minmax_format(ctx).as_deref(),\n self.get_attr_nvvm_extended_minmax_operation(ctx).as_deref(),\n self.get_attr_nvvm_extended_minmax_subnormal(ctx).as_deref(),\n self.get_attr_nvvm_extended_minmax_nan(ctx).as_deref(),\n self.get_attr_nvvm_extended_minmax_xorsign_abs(ctx).as_deref(),\n ) {\n", + ); + for record in extended_minmax(catalog) { + writeln!( + output, + " (Some(&{}), Some(&{}), Some(&{}), Some(&{}), Some(&{})) => ({:?}, {}),", + extended_minmax_format_attr(record), + extended_minmax_operation_attr(record), + extended_minmax_subnormal_attr(record), + extended_minmax_nan_attr(record), + extended_minmax_xorsign_abs_attr(record), + extended_minmax_ptx_mnemonic(record), + extended_minmax_carrier(record), + ) + .unwrap(); + } + output.push_str( + " _ => return pliron::input_err!(\n self.get_operation().deref(ctx).loc(),\n \"nvvm.extended_minmax variant has no generated lowering recipe\",\n ),\n };\n convert_generated_extended_minmax(\n ctx, rewriter, self.get_operation(), recipe.0, recipe.1,\n )\n }\n}\n\n", + ); + } + if cluster_barriers(catalog).next().is_some() { + output.push_str( + "#[op_interface_impl]\nimpl MirToLlvmConversion for ClusterBarrierOp {\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let recipe = match self.get_attr_nvvm_cluster_barrier_mode(ctx).as_deref() {\n", + ); + for record in cluster_barriers(catalog) { + writeln!( + output, + " Some(&{}) => ({:?}, {:?}),", + cluster_barrier_attr(record), + record.llvm_identifier(), + cluster_barrier_template(record) + ) + .unwrap(); + } + output.push_str( + " _ => return pliron::input_err!(\n self.get_operation().deref(ctx).loc(),\n \"nvvm.cluster_barrier mode has no generated lowering recipe\",\n ),\n };\n let op = self.get_operation();\n let void_ty = llvm_types::VoidType::get(ctx);\n match context::lowering_options(ctx).intrinsic_backend {\n IntrinsicBackend::LlvmNvptx => {\n let function_ty = llvm_types::FuncType::get(ctx, void_ty.into(), vec![], false);\n call_intrinsic(ctx, rewriter, op, recipe.0, function_ty, vec![])?;\n }\n IntrinsicBackend::LibNvvm => {\n inline_asm_convergent(ctx, rewriter, op, void_ty.into(), vec![], recipe.1, \"~{memory}\");\n }\n }\n rewriter.erase_operation(ctx, op);\n Ok(())\n }\n}\n\n", + ); + + let arrive = cluster_barriers(catalog) + .find(|record| { + record + .cluster_barrier + .as_ref() + .is_some_and(|barrier| barrier.mode == ClusterBarrierMode::ArriveAligned) + }) + .expect("aligned cluster arrival"); + let wait = cluster_barriers(catalog) + .find(|record| { + record + .cluster_barrier + .as_ref() + .is_some_and(|barrier| barrier.mode == ClusterBarrierMode::WaitAligned) + }) + .expect("aligned cluster wait"); + output.push_str( + "#[op_interface_impl]\nimpl MirToLlvmConversion for ClusterSyncOp {\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let op = self.get_operation();\n let void_ty = llvm_types::VoidType::get(ctx);\n match context::lowering_options(ctx).intrinsic_backend {\n IntrinsicBackend::LlvmNvptx => {\n let function_ty = llvm_types::FuncType::get(ctx, void_ty.into(), vec![], false);\n", + ); + writeln!( + output, + " call_intrinsic(ctx, rewriter, op, {:?}, function_ty, vec![])?;", + arrive.llvm_identifier() + ) + .unwrap(); + writeln!( + output, + " call_intrinsic(ctx, rewriter, op, {:?}, function_ty, vec![])?;", + wait.llvm_identifier() + ) + .unwrap(); + output.push_str(" }\n IntrinsicBackend::LibNvvm => {\n"); + writeln!( + output, + " inline_asm_convergent(ctx, rewriter, op, void_ty.into(), vec![], {:?}, \"~{{memory}}\");", + format!( + "{} {}", + cluster_barrier_template(arrive), + cluster_barrier_template(wait) + ) + ) + .unwrap(); + output.push_str( + " }\n }\n rewriter.erase_operation(ctx, op);\n Ok(())\n }\n}\n\n", + ); + } + if wgmma_controls(catalog).next().is_some() { + output.push_str( + "fn convert_generated_wgmma_control(\n\ + ctx: &mut Context,\n\ + rewriter: &mut DialectConversionRewriter,\n\ + op: Ptr,\n\ + intrinsic_name: &str,\n\ + ptx: &str,\n\ + has_immediate: bool,\n\ + ) -> Result<()> {\n\ + let operands: Vec<_> = op.deref(ctx).operands().collect();\n\ + if operands.len() != usize::from(has_immediate) {\n\ + return pliron::input_err!(\n\ + op.deref(ctx).loc(),\n\ + \"generated WGMMA control has the wrong operand count\",\n\ + );\n\ + }\n\ + let void_ty = llvm_types::VoidType::get(ctx);\n\ + match context::lowering_options(ctx).intrinsic_backend {\n\ + IntrinsicBackend::LlvmNvptx => {\n\ + let arguments = if has_immediate {\n\ + vec![IntegerType::get(ctx, 64, Signedness::Signless).into()]\n\ + } else {\n\ + vec![]\n\ + };\n\ + let function_ty = llvm_types::FuncType::get(\n\ + ctx, void_ty.into(), arguments, false,\n\ + );\n\ + call_intrinsic(\n\ + ctx, rewriter, op, intrinsic_name, function_ty, operands,\n\ + )?;\n\ + }\n\ + IntrinsicBackend::LibNvvm => {\n\ + let constraints = if has_immediate {\n\ + \"n,~{memory}\"\n\ + } else {\n\ + \"~{memory}\"\n\ + };\n\ + inline_asm_convergent(\n\ + ctx, rewriter, op, void_ty.into(), operands, ptx, constraints,\n\ + );\n\ + }\n\ + }\n\ + rewriter.erase_operation(ctx, op);\n\ + Ok(())\n\ + }\n\n", + ); + for record in wgmma_controls(catalog) { + let has_immediate = record + .wgmma_control + .as_ref() + .is_some_and(|control| control.mode == WgmmaControlMode::WaitGroup); + writeln!( + output, + "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {{\n convert_generated_wgmma_control(\n ctx, rewriter, self.get_operation(), {:?}, {:?}, {has_immediate},\n )\n }}\n}}\n", + record.dialect.op_type, + record.llvm_identifier(), + wgmma_control_template(record), + ) + .unwrap(); + } + } + if packed_atomics(catalog).next().is_some() { + output.push_str( + "#[op_interface_impl]\nimpl MirToLlvmConversion for PackedAtomicAddOp {\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let format = self.get_attr_nvvm_packed_atomic_format(ctx);\n let state_space = self.get_attr_nvvm_packed_atomic_state_space(ctx);\n let ordering = self.get_attr_nvvm_packed_atomic_ordering(ctx);\n let scope = self.get_attr_nvvm_packed_atomic_scope(ctx);\n let rounding = self.get_attr_nvvm_packed_atomic_rounding(ctx);\n let subnormal = self.get_attr_nvvm_packed_atomic_subnormal(ctx);\n let atomicity = self.get_attr_nvvm_packed_atomic_atomicity(ctx);\n let ptx_type = match (format.as_deref(), state_space.as_deref(), ordering.as_deref(), scope.as_deref(), rounding.as_deref(), subnormal.as_deref(), atomicity.as_deref()) {\n (Some(&PackedAtomicFormatAttr::F16x2), Some(&PackedAtomicStateSpaceAttr::Global), Some(&PackedAtomicOrderingAttr::Relaxed), Some(&PackedAtomicScopeAttr::Gpu), Some(&PackedAtomicRoundingAttr::Rn), Some(&PackedAtomicSubnormalAttr::NoFtz), Some(&PackedAtomicAtomicityAttr::PerElement)) => \"f16x2\",\n (Some(&PackedAtomicFormatAttr::Bf16x2), Some(&PackedAtomicStateSpaceAttr::Global), Some(&PackedAtomicOrderingAttr::Relaxed), Some(&PackedAtomicScopeAttr::Gpu), Some(&PackedAtomicRoundingAttr::Rn), Some(&PackedAtomicSubnormalAttr::NoFtz), Some(&PackedAtomicAtomicityAttr::PerElement)) => \"bf16x2\",\n _ => return pliron::input_err!(\n self.get_operation().deref(ctx).loc(),\n \"nvvm.packed_atomic_add attributes have no generated lowering recipe\",\n ),\n };\n convert_packed_atom_add(ctx, rewriter, self.get_operation(), ptx_type)\n }\n}\n\n", + ); + output = output.replace( + " convert_packed_atom_add(ctx, rewriter, self.get_operation(), ptx_type)", + " drop((format, state_space, ordering, scope, rounding, subnormal, atomicity));\n convert_packed_atom_add(ctx, rewriter, self.get_operation(), ptx_type)", + ); + output.push_str( + "#[op_interface_impl]\nimpl MirToLlvmConversion for NvvmAtomAddF16x2Op {\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n convert_packed_atom_add(ctx, rewriter, self.get_operation(), \"f16x2\")\n }\n}\n\n#[op_interface_impl]\nimpl MirToLlvmConversion for NvvmAtomAddBf16x2Op {\n fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n convert_packed_atom_add(ctx, rewriter, self.get_operation(), \"bf16x2\")\n }\n}\n\n", + ); + } + for record in redux(catalog) { + debug_assert_eq!( + record.redux.as_ref().unwrap().adapter, + ReduxAdapter::MaskValueToSourceMemberMask + ); + writeln!( + output, + "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", + record.dialect.op_type + ) + .unwrap(); + output.push_str( + " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n operands_info: &OperandsInfo,\n ) -> Result<()> {\n", + ); + writeln!( + output, + " convert_redux(ctx, rewriter, self.get_operation(), operands_info, {:?})", + record.llvm_identifier() + ) + .unwrap(); + output.push_str(" }\n}\n\n"); + } + for record in vote_intrinsics(catalog) { + debug_assert_eq!( + record.vote.as_ref().unwrap().adapter, + VoteAdapter::DirectMaskPredicate + ); + writeln!( + output, + "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", + record.dialect.op_type + ) + .unwrap(); + output.push_str( + " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n operands_info: &OperandsInfo,\n ) -> Result<()> {\n", + ); + writeln!( + output, + " convert_vote(ctx, rewriter, self.get_operation(), operands_info, {:?})", + record.llvm_identifier() + ) + .unwrap(); + output.push_str(" }\n}\n\n"); + } + for record in warp_matches(catalog) { + let warp_match = record.warp_match.as_ref().unwrap(); + let helper = match warp_match.adapter { + WarpMatchAdapter::DirectMask => "convert_match_any", + WarpMatchAdapter::ProjectMaskDiscardPredicate => "convert_match_all", + }; + let value_width = warp_match.value_width.bits(); + writeln!( + output, + "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", + record.dialect.op_type + ) + .unwrap(); + output.push_str( + " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n operands_info: &OperandsInfo,\n ) -> Result<()> {\n", + ); + writeln!( + output, + " let value_ty = IntegerType::get(ctx, {value_width}, Signedness::Signless);" + ) + .unwrap(); + writeln!( + output, + " {helper}(ctx, rewriter, self.get_operation(), operands_info, {:?}, value_ty.into())", + record.llvm_identifier() + ) + .unwrap(); + output.push_str(" }\n}\n\n"); + } + for record in elect_intrinsics(catalog) { + writeln!( + output, + "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", + record.dialect.op_type + ) + .unwrap(); + output.push_str( + " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n operands_info: &OperandsInfo,\n ) -> Result<()> {\n match context::lowering_options(ctx).intrinsic_backend {\n", + ); + for (backend, variant) in [ + (IntrinsicBackend::LlvmNvptx, "LlvmNvptx"), + (IntrinsicBackend::LibNvvm, "LibNvvm"), + ] { + let mechanism = record + .backend_lowerings + .iter() + .find(|route| route.backend == backend) + .expect("elect backend route") + .mechanism; + writeln!(output, " IntrinsicBackend::{variant} => {{").unwrap(); + match mechanism { + BackendLoweringMechanism::TypedNvvm => writeln!( + output, + " convert_elect_sync_typed(ctx, rewriter, self.get_operation(), operands_info, {:?})", + record.llvm_identifier() + ) + .unwrap(), + BackendLoweringMechanism::InlinePtx => output.push_str( + " convert_elect_sync_inline(ctx, rewriter, self.get_operation(), operands_info)\n", + ), + } + output.push_str(" }\n"); + } + output.push_str(" }\n }\n}\n\n"); + } + for record in warp_barriers(catalog) { + debug_assert_eq!( + record.warp_barrier.as_ref().unwrap().adapter, + WarpBarrierAdapter::DirectMemberMask + ); + writeln!( + output, + "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", + record.dialect.op_type + ) + .unwrap(); + output.push_str( + " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n operands_info: &OperandsInfo,\n ) -> Result<()> {\n convert_bar_warp_sync(ctx, rewriter, self.get_operation(), operands_info)\n }\n}\n\n", + ); + } + for record in warp_shuffles(catalog) { + let shuffle = record.warp_shuffle.as_ref().unwrap(); + writeln!( + output, + "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", + record.dialect.op_type + ) + .unwrap(); + output.push_str( + " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n operands_info: &OperandsInfo,\n ) -> Result<()> {\n", + ); + match shuffle.value_kind { + WarpShuffleValueKind::I32 | WarpShuffleValueKind::F32 => { + debug_assert_eq!( + shuffle.adapter, + WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp + ); + let helper = match shuffle.value_kind { + WarpShuffleValueKind::I32 => "convert_shuffle_i32", + WarpShuffleValueKind::F32 => "convert_shuffle_f32", + WarpShuffleValueKind::I64 => unreachable!(), + }; + writeln!( + output, + " {helper}(ctx, rewriter, self.get_operation(), operands_info, {:?}, {})", + record.llvm_identifier(), + shuffle.clamp, + ) + .unwrap(); + } + WarpShuffleValueKind::I64 => { + debug_assert_eq!( + shuffle.adapter, + WarpShuffleAdapter::MaskValueLaneOrDeltaSplitI64LowHighB32InsertClampReassemble + ); + let mode = match shuffle.mode { + WarpShuffleMode::Idx => "idx", + WarpShuffleMode::Bfly => "bfly", + WarpShuffleMode::Down => "down", + WarpShuffleMode::Up => "up", + }; + writeln!( + output, + " convert_shuffle_i64(ctx, rewriter, self.get_operation(), operands_info, {mode:?}, {})", + shuffle.clamp, + ) + .unwrap(); + } + } + output.push_str(" }\n}\n\n"); + } + for record in packed_alus(catalog) { + let width = packed_alu_width(record); + writeln!( + output, + "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", + record.dialect.op_type + ) + .unwrap(); + output.push_str( + " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n", + ); + writeln!( + output, + " convert_generated_packed_alu(ctx, rewriter, self.get_operation(), {:?}, {width})", + packed_alu_ptx_mnemonic(record), + ) + .unwrap(); + output.push_str(" }\n}\n\n"); + } + for record in integer_minmaxes(catalog) { + writeln!( + output, + "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", + record.dialect.op_type + ) + .unwrap(); + output.push_str( + " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n", + ); + writeln!( + output, + " convert_generated_integer_minmax(ctx, rewriter, self.get_operation(), {:?})", + integer_minmax_ptx_mnemonic(record) + ) + .unwrap(); + output.push_str(" }\n}\n\n"); + } + for record in packed_conversions(catalog) { + let conversion = record + .packed_conversion + .as_ref() + .expect("packed-conversion record"); + writeln!( + output, + "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", + record.dialect.op_type + ) + .unwrap(); + output.push_str( + " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n", + ); + match conversion.source_format { + PackedConversionSourceFormat::F32x2 => { + debug_assert_eq!( + conversion.adapter, + PackedConversionAdapter::ReverseHighLowOperands + ); + let typed_intrinsic = packed_conversion_typed_llvm_name(record) + .map(|name| format!("Some({name:?})")) + .unwrap_or_else(|| "None".into()); + writeln!( + output, + " convert_generated_packed_f32x2(ctx, rewriter, self.get_operation(), {typed_intrinsic}, {:?}, {})", + packed_conversion_ptx_mnemonic(record), + packed_conversion_result_width(record), + ) + .unwrap(); + } + PackedConversionSourceFormat::E4m3x2 + | PackedConversionSourceFormat::E5m2x2 + | PackedConversionSourceFormat::F16x2 => { + debug_assert_eq!(conversion.adapter, PackedConversionAdapter::Identity); + writeln!( + output, + " convert_generated_packed_unary(ctx, rewriter, self.get_operation(), {:?}, {}, {})", + packed_conversion_ptx_mnemonic(record), + packed_conversion_result_width(record), + packed_conversion_source_width(record), + ) + .unwrap(); + } + } + output.push_str(" }\n}\n\n"); + } + for record in cp_async_copies(catalog) { + let copy = record.cp_async_copy.as_ref().unwrap(); + let cache_policy = match copy.cache_policy { + CpAsyncCachePolicy::Ca => "ca", + CpAsyncCachePolicy::Cg => "cg", + }; + let has_source_size = copy.source_size == CpAsyncSourceSize::Runtime; + writeln!( + output, + "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", + record.dialect.op_type + ) + .unwrap(); + output.push_str( + " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n", + ); + writeln!( + output, + " convert_generated_cp_async_copy(ctx, rewriter, self.get_operation(), {cache_policy:?}, {}, {has_source_size}, {:?})", + copy.copy_size.bytes(), + record.llvm_identifier(), + ) + .unwrap(); + output.push_str(" }\n}\n\n"); + } + for record in cp_async_controls(catalog) { + let operation = match record.cp_async_control.as_ref().unwrap().operation { + CpAsyncControlOperation::CommitGroup => "commit_group", + CpAsyncControlOperation::WaitAll => "wait_all", + CpAsyncControlOperation::WaitGroup => "wait_group", + }; + writeln!( + output, + "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", + record.dialect.op_type + ) + .unwrap(); + output.push_str( + " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n", + ); + writeln!( + output, + " convert_generated_cp_async_control(ctx, rewriter, self.get_operation(), {operation:?}, {:?})", + record.llvm_identifier(), + ) + .unwrap(); + output.push_str(" }\n}\n\n"); + } + for record in cp_async_mbarriers(catalog) { + let bridge = record.cp_async_mbarrier.as_ref().unwrap(); + let operation = match bridge.operation { + CpAsyncMbarrierOperation::Arrive => "arrive", + CpAsyncMbarrierOperation::ArriveNoInc => "arrive_no_inc", + }; + let state_space = match bridge.state_space { + CpAsyncMbarrierStateSpace::Generic => "generic", + CpAsyncMbarrierStateSpace::Shared => "shared", + }; + writeln!( + output, + "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", + record.dialect.op_type + ) + .unwrap(); + output.push_str( + " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n", + ); + writeln!( + output, + " convert_generated_cp_async_mbarrier(ctx, rewriter, self.get_operation(), {operation:?}, {state_space:?}, {:?})", + record.llvm_identifier(), + ) + .unwrap(); + output.push_str(" }\n}\n\n"); + } + for record in mbarrier_basics(catalog) { + let mbarrier = record.mbarrier_basic.as_ref().unwrap(); + let helper = match (mbarrier.operation, mbarrier.adapter) { + (MbarrierBasicOperation::Init, MbarrierBasicAdapter::InitPointerCountToVoid) => { + "convert_init" + } + (MbarrierBasicOperation::Arrive, MbarrierBasicAdapter::ArrivePointerToToken) => { + "convert_arrive" + } + ( + MbarrierBasicOperation::TestWait, + MbarrierBasicAdapter::TestWaitPointerTokenToPredicate, + ) => "convert_test_wait", + (MbarrierBasicOperation::Inval, MbarrierBasicAdapter::InvalPointerToVoid) => { + "convert_inval" + } + _ => unreachable!("resolver admitted an invalid basic mbarrier adapter"), + }; + writeln!( + output, + "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", + record.dialect.op_type + ) + .unwrap(); + output.push_str( + " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n operands_info: &OperandsInfo,\n ) -> Result<()> {\n", + ); + writeln!( + output, + " {helper}(ctx, rewriter, self.get_operation(), operands_info)" + ) + .unwrap(); + output.push_str(" }\n}\n\n"); + } + for record in cluster_memory(catalog) { + let cluster = record.cluster_memory.as_ref().unwrap(); + let (template, constraints) = + crate::resolve::cluster_memory_inline_recipe(cluster.operation); + writeln!( + output, + "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", + record.dialect.op_type + ) + .unwrap(); + output.push_str( + " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let op = self.get_operation();\n let operands: Vec<_> = op.deref(ctx).operands().collect();\n", + ); + writeln!( + output, + " if operands.len() != 2 {{\n return pliron::input_err_noloc!({:?}, operands.len());\n }}", + format!("{} requires 2 operands, got {{}}", record.rust.name) + ) + .unwrap(); + output.push_str( + " let shared_pointer = cast_to_shared_addrspace(ctx, rewriter, operands[0]);\n let rank = operands[1];\n", + ); + match cluster.operation { + ClusterMemoryOperation::MapSharedRank => { + writeln!( + output, + " let i64_ty = IntegerType::get(ctx, 64, Signedness::Signless);\n let asm = inline_asm_convergent(\n ctx, rewriter, op, i64_ty.into(), vec![shared_pointer, rank], {template:?}, {constraints:?},\n );\n let mapped = asm.deref(ctx).get_result(0);\n let cluster_shared_pointer_ty = llvm_types::PointerType::get(ctx, 7);\n let int_to_ptr = llvm_ops::IntToPtrOp::new(ctx, mapped, cluster_shared_pointer_ty.into());\n rewriter.insert_operation(ctx, int_to_ptr.get_operation());\n rewriter.replace_operation(ctx, op, int_to_ptr.get_operation());\n Ok(())" + ) + .unwrap(); + } + ClusterMemoryOperation::ReadU32 => { + writeln!( + output, + " let i32_ty = IntegerType::get(ctx, 32, Signedness::Signless);\n let asm = inline_asm_convergent(\n ctx, rewriter, op, i32_ty.into(), vec![shared_pointer, rank], {template:?}, {constraints:?},\n );\n rewriter.replace_operation(ctx, op, asm);\n Ok(())" + ) + .unwrap(); + } + } + output.push_str(" }\n}\n\n"); + } + for record in mbarrier_extended(catalog) { + let contract = record.mbarrier_extended.as_ref().unwrap(); + let (template, constraints) = + crate::resolve::mbarrier_extended_inline_recipe(contract.operation); + writeln!( + output, + "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", + record.dialect.op_type + ) + .unwrap(); + let mutability = if matches!( + contract.adapter, + MbarrierExtendedAdapter::PointerTxCountBytesToTokenDroppingTxCount + | MbarrierExtendedAdapter::PointerTokenToPredicate + | MbarrierExtendedAdapter::PointerParityToPredicate + ) { + "mut " + } else { + "" + }; + output.push_str( + " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let op = self.get_operation();\n", + ); + writeln!( + output, + " let {mutability}operands: Vec<_> = op.deref(ctx).operands().collect();" + ) + .unwrap(); + let operand_count = record.dialect.operands.len(); + let operand_count_check = if operand_count == 0 { + "!operands.is_empty()".to_owned() + } else { + format!("operands.len() != {operand_count}") + }; + writeln!( + output, + " if {operand_count_check} {{\n return pliron::input_err_noloc!({:?}, operands.len());\n }}", + format!( + "{} requires {} operands, got {{}}", + record.rust.name, + operand_count + ) + ) + .unwrap(); + if matches!( + contract.adapter, + MbarrierExtendedAdapter::PointerTxCountBytesToTokenDroppingTxCount + | MbarrierExtendedAdapter::PointerTokenToPredicate + | MbarrierExtendedAdapter::PointerParityToPredicate + ) { + output.push_str( + " operands[0] = cast_to_shared_addrspace(ctx, rewriter, operands[0]);\n", + ); + } + match contract.adapter { + MbarrierExtendedAdapter::PointerTxCountBytesToTokenDroppingTxCount => { + writeln!( + output, + " let result_ty = IntegerType::get(ctx, 64, Signedness::Signless);\n let asm = inline_asm_convergent(ctx, rewriter, op, result_ty.into(), operands, {template:?}, {constraints:?});\n rewriter.replace_operation(ctx, op, asm);\n Ok(())" + ) + .unwrap(); + } + MbarrierExtendedAdapter::PointerTokenToPredicate + | MbarrierExtendedAdapter::PointerParityToPredicate => { + writeln!( + output, + " let result_ty = IntegerType::get(ctx, 32, Signedness::Signless);\n let asm = inline_asm_convergent(ctx, rewriter, op, result_ty.into(), operands, {template:?}, {constraints:?});\n let asm_result = asm.deref(ctx).get_result(0);\n let result = trunc_to_i1(ctx, rewriter, asm_result);\n let DefiningEntity::Op(result_op) = result.defining_entity() else {{ unreachable!() }};\n rewriter.replace_operation(ctx, op, result_op);\n Ok(())" + ) + .unwrap(); + } + MbarrierExtendedAdapter::RawClusterAddressToVoid + | MbarrierExtendedAdapter::ZeroOperandsToVoid + | MbarrierExtendedAdapter::NanosecondsToVoid => { + writeln!( + output, + " let void_ty = llvm_types::VoidType::get(ctx);\n inline_asm_convergent(ctx, rewriter, op, void_ty.into(), operands, {template:?}, {constraints:?});\n rewriter.erase_operation(ctx, op);\n Ok(())" + ) + .unwrap(); + } + } + output.push_str(" }\n}\n\n"); + } + for record in dot_products(catalog) { + let insert_low_half_selector = matches!( + record.dot_product.as_ref().unwrap().adapter, + DotProductAdapter::InsertLowHalfFalse + ); + writeln!( + output, + "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", + record.dialect.op_type + ) + .unwrap(); + output.push_str( + " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n", + ); + writeln!( + output, + " convert_generated_dot_product(ctx, rewriter, self.get_operation(), {:?}, {:?}, {insert_low_half_selector})", + record.llvm_identifier(), + dot_product_ptx(record), + ) + .unwrap(); + output.push_str(" }\n}\n\n"); + } + for record in clc_intrinsics(catalog) { + writeln!( + output, + "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", + record.dialect.op_type + ) + .unwrap(); + output.push_str( + " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n operands_info: &OperandsInfo,\n ) -> Result<()> {\n", + ); + match record.clc.as_ref().unwrap().operation { + ClcOperation::TryCancel | ClcOperation::TryCancelMulticast => { + writeln!( + output, + " convert_generated_clc_try_cancel(ctx, rewriter, self.get_operation(), operands_info, {:?})", + record.llvm_identifier() + ) + .unwrap(); + } + ClcOperation::QueryIsCanceled + | ClcOperation::QueryGetFirstCtaidX + | ClcOperation::QueryGetFirstCtaidY + | ClcOperation::QueryGetFirstCtaidZ => { + writeln!( + output, + " convert_generated_clc_query(ctx, rewriter, self.get_operation(), operands_info, {:?}, {})", + record.llvm_identifier(), + record.clc.as_ref().unwrap().operation == ClcOperation::QueryIsCanceled + ) + .unwrap(); + } + } + output.push_str(" }\n}\n\n"); + } + for record in execution_controls(catalog) { + let operation = ExecutionControlOperation::from_catalog_id(&record.id) + .expect("closed execution-control record"); + writeln!( + output, + "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", + record.dialect.op_type + ) + .unwrap(); + output.push_str( + " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n operands_info: &OperandsInfo,\n ) -> Result<()> {\n", + ); + match operation { + ExecutionControlOperation::BarrierCtaSync + | ExecutionControlOperation::BarrierCtaSyncAligned + | ExecutionControlOperation::BarrierCtaArrive + | ExecutionControlOperation::BarrierCtaArriveAligned => { + let template = format!( + "{}.{} $0, $1;", + record.expected_ptx.mnemonic, + record.expected_ptx.modifiers.join(".") + ); + writeln!( + output, + " convert_counted_barrier(ctx, rewriter, self.get_operation(), operands_info, {:?}, {template:?})", + record.llvm_identifier() + ) + .unwrap(); + } + ExecutionControlOperation::GridDependencyLaunchDependents + | ExecutionControlOperation::GridDependencyWait => { + let template = format!( + "{}.{};", + record.expected_ptx.mnemonic, + record.expected_ptx.modifiers.join(".") + ); + writeln!( + output, + " convert_grid_dependency(ctx, rewriter, self.get_operation(), operands_info, {:?}, {template:?})", + record.llvm_identifier() + ) + .unwrap(); + } + ExecutionControlOperation::SetMaxNRegInc | ExecutionControlOperation::SetMaxNRegDec => { + let direction = match operation { + ExecutionControlOperation::SetMaxNRegInc => "inc", + ExecutionControlOperation::SetMaxNRegDec => "dec", + _ => unreachable!("setmaxnreg operation was matched"), + }; + output.push_str( + " let register_count = self.register_count(ctx).ok_or_else(|| pliron::input_error_noloc!(\"setmaxnreg register-count attribute is invalid\"))?;\n", + ); + writeln!( + output, + " convert_setmaxnreg(ctx, rewriter, self.get_operation(), operands_info, register_count, {:?}, {direction:?})", + record.llvm_identifier() + ) + .unwrap(); + } + } + output.push_str(" }\n}\n\n"); + } + for record in tma_intrinsics(catalog) { + let operation = record.tma.as_ref().unwrap().operation; + writeln!( + output, + "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", + record.dialect.op_type + ) + .unwrap(); + output.push_str( + " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n operands_info: &OperandsInfo,\n ) -> Result<()> {\n", + ); + match operation { + TmaOperation::G2sTile1d + | TmaOperation::G2sTile2d + | TmaOperation::G2sTile3d + | TmaOperation::G2sTile4d + | TmaOperation::G2sTile5d => { + writeln!( + output, + " convert_g2s(ctx, rewriter, self.get_operation(), operands_info, {}, false)", + operation.dimensions().unwrap() + ) + .unwrap(); + } + TmaOperation::G2sTile2dMulticast => output.push_str( + " convert_g2s(ctx, rewriter, self.get_operation(), operands_info, 2, true)\n", + ), + TmaOperation::G2sTile2dMulticastCg2 => output.push_str( + " convert_g2s_multicast_cg2(ctx, rewriter, self.get_operation(), operands_info)\n", + ), + TmaOperation::S2gTile1d + | TmaOperation::S2gTile2d + | TmaOperation::S2gTile3d + | TmaOperation::S2gTile4d + | TmaOperation::S2gTile5d => { + writeln!( + output, + " convert_s2g(ctx, rewriter, self.get_operation(), operands_info, {})", + operation.dimensions().unwrap() + ) + .unwrap(); + } + TmaOperation::Reduce => { + let reduction = record + .tma + .as_ref() + .and_then(|tma| tma.reduction.as_ref()) + .expect("TMA reduction contract"); + let reduction_name = match reduction.operation { + TmaReductionOperation::Add => "add", + TmaReductionOperation::And => "and", + TmaReductionOperation::Dec => "dec", + TmaReductionOperation::Inc => "inc", + TmaReductionOperation::Max => "max", + TmaReductionOperation::Min => "min", + TmaReductionOperation::Or => "or", + TmaReductionOperation::Xor => "xor", + }; + let load_mode = match reduction.load_mode { + TmaReductionLoadMode::Tile => "tile", + TmaReductionLoadMode::Im2col => "im2col", + }; + writeln!( + output, + " convert_reduce_s2g(ctx, rewriter, self.get_operation(), operands_info, ReduceConfig::new({}, {reduction_name:?}, {load_mode:?}, {:?}))", + reduction.dimensions, + record.resolved_llvm_identifier() + ) + .unwrap(); + } + TmaOperation::CommitGroup | TmaOperation::WaitGroup | TmaOperation::WaitGroupRead => { + let operation_name = match operation { + TmaOperation::CommitGroup => "commit_group", + TmaOperation::WaitGroup => "wait_group", + TmaOperation::WaitGroupRead => "wait_group_read", + _ => unreachable!("TMA control operation was matched"), + }; + writeln!( + output, + " convert_control(ctx, rewriter, self.get_operation(), operands_info, {operation_name:?}, {:?})", + record.resolved_llvm_identifier() + ) + .unwrap(); + } + TmaOperation::PrefetchTensorMap => { + writeln!( + output, + " convert_prefetch_tensormap(ctx, rewriter, self.get_operation(), operands_info, {:?})", + record.resolved_llvm_identifier() + ) + .unwrap(); + } + TmaOperation::PrefetchTile1d + | TmaOperation::PrefetchTile2d + | TmaOperation::PrefetchTile3d + | TmaOperation::PrefetchTile4d + | TmaOperation::PrefetchTile5d + | TmaOperation::PrefetchTileGather4TwoDimensional + | TmaOperation::PrefetchTile1dCacheHint + | TmaOperation::PrefetchTile2dCacheHint + | TmaOperation::PrefetchTile3dCacheHint + | TmaOperation::PrefetchTile4dCacheHint + | TmaOperation::PrefetchTile5dCacheHint + | TmaOperation::PrefetchTileGather4TwoDimensionalCacheHint => { + writeln!( + output, + " convert_prefetch_tile(ctx, rewriter, self.get_operation(), operands_info, PrefetchTileConfig::new({}, {}, {}, {:?}))", + operation.prefetch_coordinate_count().unwrap(), + matches!( + operation, + TmaOperation::PrefetchTileGather4TwoDimensional + | TmaOperation::PrefetchTileGather4TwoDimensionalCacheHint + ), + operation.uses_prefetch_cache_hint(), + record.resolved_llvm_identifier() + ) + .unwrap(); + } + TmaOperation::ReplaceBoxDim + | TmaOperation::ReplaceElementStride + | TmaOperation::ReplaceElementType + | TmaOperation::ReplaceFillMode + | TmaOperation::ReplaceGlobalAddress + | TmaOperation::ReplaceGlobalDim + | TmaOperation::ReplaceGlobalStride + | TmaOperation::ReplaceInterleaveLayout + | TmaOperation::ReplaceRank + | TmaOperation::ReplaceSwizzleAtomicity + | TmaOperation::ReplaceSwizzleMode => { + let (field, value_kind, ordinal, immediate) = match operation { + TmaOperation::ReplaceBoxDim => ("box_dim", "u32", true, false), + TmaOperation::ReplaceElementStride => { + ("element_stride", "u32", true, false) + } + TmaOperation::ReplaceElementType => ("elemtype", "u32", false, true), + TmaOperation::ReplaceFillMode => ("fill_mode", "u32", false, true), + TmaOperation::ReplaceGlobalAddress => { + ("global_address", "address", false, false) + } + TmaOperation::ReplaceGlobalDim => ("global_dim", "u32", true, false), + TmaOperation::ReplaceGlobalStride => { + ("global_stride", "u64", true, false) + } + TmaOperation::ReplaceInterleaveLayout => { + ("interleave_layout", "u32", false, true) + } + TmaOperation::ReplaceRank => ("rank", "u32", false, false), + TmaOperation::ReplaceSwizzleAtomicity => { + ("swizzle_atomicity", "u32", false, true) + } + TmaOperation::ReplaceSwizzleMode => { + ("swizzle_mode", "u32", false, true) + } + _ => unreachable!("TMA tensor-map replace operation was matched"), + }; + writeln!( + output, + " convert_tensormap_replace(ctx, rewriter, self.get_operation(), operands_info, {:?}, {field:?}, {value_kind:?}, {ordinal}, {immediate})", + record.resolved_llvm_identifier() + ) + .unwrap(); + } + TmaOperation::FenceProxyTensorMapAcquireCluster + | TmaOperation::FenceProxyTensorMapAcquireCta + | TmaOperation::FenceProxyTensorMapAcquireGpu + | TmaOperation::FenceProxyTensorMapAcquireSystem + | TmaOperation::FenceProxyTensorMapReleaseCluster + | TmaOperation::FenceProxyTensorMapReleaseCta + | TmaOperation::FenceProxyTensorMapReleaseGpu + | TmaOperation::FenceProxyTensorMapReleaseSystem => { + let (acquire, scope) = match operation { + TmaOperation::FenceProxyTensorMapAcquireCluster => (true, "cluster"), + TmaOperation::FenceProxyTensorMapAcquireCta => (true, "cta"), + TmaOperation::FenceProxyTensorMapAcquireGpu => (true, "gpu"), + TmaOperation::FenceProxyTensorMapAcquireSystem => (true, "sys"), + TmaOperation::FenceProxyTensorMapReleaseCluster => (false, "cluster"), + TmaOperation::FenceProxyTensorMapReleaseCta => (false, "cta"), + TmaOperation::FenceProxyTensorMapReleaseGpu => (false, "gpu"), + TmaOperation::FenceProxyTensorMapReleaseSystem => (false, "sys"), + _ => unreachable!("TMA tensor-map fence operation was matched"), + }; + writeln!( + output, + " convert_tensormap_fence(ctx, rewriter, self.get_operation(), operands_info, {:?}, {acquire}, {scope:?})", + record.resolved_llvm_identifier() + ) + .unwrap(); + } + } + output.push_str(" }\n}\n\n"); + } + for record in tcgen05_non_mma_intrinsics(catalog) { + let operation = record.tcgen05.as_ref().unwrap().operation; + let (template, constraints, result_count) = tcgen05_inline_asm(record); + writeln!( + output, + "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", + record.dialect.op_type + ) + .unwrap(); + output.push_str( + " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n", + ); + if let Some(count) = result_count { + let integer_results = operation == Tcgen05Operation::Ld; + writeln!( + output, + " convert_generated_tcgen05_load(ctx, rewriter, self.get_operation(), {}, {count}, {integer_results}, {template:?}, {constraints:?})", + record.dialect.operands.len() + ) + .unwrap(); + } else { + writeln!( + output, + " convert_generated_tcgen05_void(ctx, rewriter, self.get_operation(), {}, {template:?}, {constraints:?})", + record.dialect.operands.len() + ) + .unwrap(); + } + output.push_str(" }\n}\n\n"); + } + if tcgen05_mma_intrinsics(catalog).next().is_some() { + output.push_str( + r#"#[op_interface_impl] +impl MirToLlvmConversion for Tcgen05MmaOp { + fn convert( + &self, + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + _operands_info: &OperandsInfo, + ) -> Result<()> { + let op = self.get_operation(); + let form = self + .get_attr_nvvm_tcgen05_mma_form(ctx) + .as_deref() + .cloned(); + let kind = self + .get_attr_nvvm_tcgen05_mma_kind(ctx) + .as_deref() + .cloned(); + let cta_group = self + .get_attr_nvvm_tcgen05_mma_cta_group(ctx) + .as_deref() + .cloned(); + let collector_a = self + .get_attr_nvvm_tcgen05_mma_collector_a(ctx) + .as_deref() + .cloned(); + let b_buffer = self + .get_attr_nvvm_tcgen05_mma_b_buffer(ctx) + .as_deref() + .cloned(); + let b_usage = self + .get_attr_nvvm_tcgen05_mma_b_usage(ctx) + .as_deref() + .cloned(); + let kind = match kind { + Some(Tcgen05MmaKindAttr::F16) => "f16", + Some(Tcgen05MmaKindAttr::Tf32) => "tf32", + Some(Tcgen05MmaKindAttr::F8f6f4) => "f8f6f4", + Some(Tcgen05MmaKindAttr::I8) => "i8", + None => return pliron::input_err!( + op.deref(ctx).loc(), + "nvvm.tcgen05_mma requires a kind", + ), + }; + let (prefix, a, metadata, zero_mask, constraints, arity, base, ashift) = + match form { + Some(Tcgen05MmaFormAttr::Shared) => + ("tcgen05.mma", "$1", None, None, "r,l,l,r,r,~{memory}", 5, true, false), + Some(Tcgen05MmaFormAttr::Tensor) => + ("tcgen05.mma", "[$1]", None, None, "r,r,l,r,r,~{memory}", 5, true, false), + Some(Tcgen05MmaFormAttr::TensorAshift) => + ("tcgen05.mma", "[$1]", None, None, "r,r,l,r,r,~{memory}", 5, true, true), + Some(Tcgen05MmaFormAttr::SpShared) => + ("tcgen05.mma.sp", "$1", Some("$5"), None, "r,l,l,r,r,r,~{memory}", 6, true, false), + Some(Tcgen05MmaFormAttr::SpTensor) => + ("tcgen05.mma.sp", "[$1]", Some("$5"), None, "r,r,l,r,r,r,~{memory}", 6, true, false), + Some(Tcgen05MmaFormAttr::SpTensorAshift) => + ("tcgen05.mma.sp", "[$1]", Some("$5"), None, "r,r,l,r,r,r,~{memory}", 6, true, true), + Some(Tcgen05MmaFormAttr::WsShared) => + ("tcgen05.mma.ws", "$1", None, None, "r,l,l,r,r,~{memory}", 5, false, false), + Some(Tcgen05MmaFormAttr::WsSharedZeroColMask) => + ("tcgen05.mma.ws", "$1", None, Some("$5"), "r,l,l,r,r,l,~{memory}", 6, false, false), + Some(Tcgen05MmaFormAttr::WsSpShared) => + ("tcgen05.mma.ws.sp", "$1", Some("$5"), None, "r,l,l,r,r,r,~{memory}", 6, false, false), + Some(Tcgen05MmaFormAttr::WsSpSharedZeroColMask) => + ("tcgen05.mma.ws.sp", "$1", Some("$5"), Some("$6"), "r,l,l,r,r,r,l,~{memory}", 7, false, false), + Some(Tcgen05MmaFormAttr::WsSpTensor) => + ("tcgen05.mma.ws.sp", "[$1]", Some("$5"), None, "r,r,l,r,r,r,~{memory}", 6, false, false), + Some(Tcgen05MmaFormAttr::WsSpTensorZeroColMask) => + ("tcgen05.mma.ws.sp", "[$1]", Some("$5"), Some("$6"), "r,r,l,r,r,r,l,~{memory}", 7, false, false), + Some(Tcgen05MmaFormAttr::WsTensor) => + ("tcgen05.mma.ws", "[$1]", None, None, "r,r,l,r,r,~{memory}", 5, false, false), + Some(Tcgen05MmaFormAttr::WsTensorZeroColMask) => + ("tcgen05.mma.ws", "[$1]", None, Some("$5"), "r,r,l,r,r,l,~{memory}", 6, false, false), + None => return pliron::input_err!( + op.deref(ctx).loc(), + "nvvm.tcgen05_mma requires a form", + ), + }; + + let selector = if base { + let group = match cta_group { + Some(Tcgen05MmaCtaGroupAttr::Cg1) => 1, + Some(Tcgen05MmaCtaGroupAttr::Cg2) => 2, + None => return pliron::input_err!( + op.deref(ctx).loc(), + "base tcgen05 MMA requires a CTA-group selector", + ), + }; + let usage = match collector_a { + Some(Tcgen05MmaCollectorAAttr::Discard) => "discard", + Some(Tcgen05MmaCollectorAAttr::LastUse) => "lastuse", + Some(Tcgen05MmaCollectorAAttr::Fill) if !ashift => "fill", + Some(Tcgen05MmaCollectorAAttr::Use) if !ashift => "use", + _ => return pliron::input_err!( + op.deref(ctx).loc(), + "tcgen05 MMA collector-A selector has no generated recipe", + ), + }; + format!( + ".cta_group::{group}.kind::{kind}.collector::a::{usage}{}", + if ashift { ".ashift" } else { "" }, + ) + } else { + let buffer = match b_buffer { + Some(Tcgen05MmaBBufferAttr::B0) => 0, + Some(Tcgen05MmaBBufferAttr::B1) => 1, + Some(Tcgen05MmaBBufferAttr::B2) => 2, + Some(Tcgen05MmaBBufferAttr::B3) => 3, + None => return pliron::input_err!( + op.deref(ctx).loc(), + "warp-specialized tcgen05 MMA requires a B-buffer selector", + ), + }; + let usage = match b_usage { + Some(Tcgen05MmaBUsageAttr::Discard) => "discard", + Some(Tcgen05MmaBUsageAttr::LastUse) => "lastuse", + Some(Tcgen05MmaBUsageAttr::Fill) => "fill", + Some(Tcgen05MmaBUsageAttr::Use) => "use", + None => return pliron::input_err!( + op.deref(ctx).loc(), + "warp-specialized tcgen05 MMA requires a B-usage selector", + ), + }; + format!( + ".cta_group::1.kind::{kind}.collector::b{buffer}::{usage}" + ) + }; + + let mut operands = format!("[$0], {a}, $2"); + if let Some(metadata) = metadata { + operands.push_str(&format!(", [{metadata}]")); + } + operands.push_str(", $3, %enable_pred"); + if let Some(zero_mask) = zero_mask { + operands.push_str(&format!(", {zero_mask}")); + } + let template = format!( + "{{ .reg .pred %enable_pred; setp.ne.s32 %enable_pred, $4, 0; {prefix}{selector} {operands}; }}" + ); + convert_generated_tcgen05_void( + ctx, rewriter, op, arity, &template, constraints, + ) + } +} + +"#, + ); + } + for record in debug_controls(catalog) { + writeln!( + output, + "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", + record.dialect.op_type + ) + .unwrap(); + output.push_str( + " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let op = self.get_operation();\n let void_ty = llvm_types::VoidType::get(ctx);\n", + ); + match record.debug_control.as_ref().unwrap().operation { + DebugControlOperation::Trap => output.push_str( + " inline_asm_sideeffect(ctx, rewriter, op, void_ty.into(), vec![], \"trap;\", \"\");\n", + ), + DebugControlOperation::Breakpoint => output.push_str( + " inline_asm_sideeffect(ctx, rewriter, op, void_ty.into(), vec![], \"brkpt;\", \"\");\n", + ), + DebugControlOperation::Pmevent => output.push_str( + " let Some(event_id) = self.event_id(ctx) else {\n\ + return pliron::input_err!(\n\ + op.deref(ctx).loc(),\n\ + \"nvvm.pmevent requires a u32 event ID in 0..=15\",\n\ + );\n\ + };\n\ + let template = format!(\"pmevent {event_id};\");\n\ + inline_asm_sideeffect(ctx, rewriter, op, void_ty.into(), vec![], &template, \"\");\n", + ), + } + output.push_str(" rewriter.erase_operation(ctx, op);\n Ok(())\n }\n}\n\n"); + } + for record in sync_intrinsics(catalog) { + writeln!( + output, + "#[op_interface_impl]\nimpl MirToLlvmConversion for {} {{", + record.dialect.op_type + ) + .unwrap(); + if record.id == "sync_threads" { + output.push_str( + " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let op = self.get_operation();\n let void_ty = llvm_types::VoidType::get(ctx);\n match context::lowering_options(ctx).intrinsic_backend {\n IntrinsicBackend::LlvmNvptx => {\n let i32_ty = IntegerType::get(ctx, 32, Signedness::Signless);\n let barrier_id = create_i32_const(ctx, rewriter, 0);\n let function_ty = llvm_types::FuncType::get(ctx, void_ty.into(), vec![i32_ty.into()], false);\n", + ); + writeln!( + output, + " call_intrinsic(ctx, rewriter, op, {:?}, function_ty, vec![barrier_id])?;", + record.llvm_identifier() + ) + .unwrap(); + output.push_str( + " }\n IntrinsicBackend::LibNvvm => {\n inline_asm_convergent(ctx, rewriter, op, void_ty.into(), vec![], \"bar.sync 0;\", \"~{memory}\");\n }\n }\n rewriter.erase_operation(ctx, op);\n Ok(())\n }\n}\n\n", + ); + } else { + debug_assert!(threadfence_ptx_level(record).is_some()); + output.push_str( + " fn convert(\n &self,\n ctx: &mut Context,\n rewriter: &mut DialectConversionRewriter,\n _operands_info: &OperandsInfo,\n ) -> Result<()> {\n let op = self.get_operation();\n let void_ty = llvm_types::VoidType::get(ctx);\n let function_ty = llvm_types::FuncType::get(ctx, void_ty.into(), vec![], false);\n", + ); + writeln!( + output, + " call_intrinsic(ctx, rewriter, op, {:?}, function_ty, vec![])?;", + record.llvm_identifier() + ) + .unwrap(); + output.push_str( + " rewriter.erase_operation(ctx, op);\n Ok(())\n }\n}\n\n", + ); + } + } + output +} diff --git a/crates/cuda-intrinsics-gen/src/render/mod.rs b/crates/cuda-intrinsics-gen/src/render/mod.rs new file mode 100644 index 0000000000..c1a3aaf9a9 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/render/mod.rs @@ -0,0 +1,465 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{CatalogFile, ExtendedMinMaxFormat, IntrinsicBackend, PackedAluFormat}; +use crate::render::collector_targets::{render_collector, render_targets}; +use crate::render::common::backend_label; +use crate::render::compat::{ + render_compat_clc, render_compat_cluster_barrier, render_compat_cluster_memory, + render_compat_cluster_sreg, render_compat_counted_barrier, render_compat_cp_async_copy, + render_compat_debug_control, render_compat_dotprod, render_compat_fence, + render_compat_float_output, render_compat_grid_dependency, render_compat_integer_minmax, + render_compat_ldmatrix, render_compat_mbarrier_basic, render_compat_mbarrier_extended, + render_compat_movmatrix, render_compat_packed_alu, render_compat_packed_atomic, + render_compat_packed_conversion, render_compat_prmt, render_compat_register_control, + render_compat_register_mma, render_compat_scalar_minmax, render_compat_sparse_mma, + render_compat_special_register_module, render_compat_sreg, render_compat_stmatrix, + render_compat_tcgen05, render_compat_tma, render_compat_wgmma_control, +}; +use crate::render::dialect::{ + render_dialect_active_mask, render_dialect_clc, render_dialect_cluster_barrier, + render_dialect_cluster_memory, render_dialect_cp_async_copy, render_dialect_debug_control, + render_dialect_dotprod, render_dialect_elect, render_dialect_execution_control, + render_dialect_extended_minmax, render_dialect_integer_minmax, render_dialect_ldmatrix, + render_dialect_mbarrier_basic, render_dialect_mbarrier_extended, render_dialect_mod, + render_dialect_movmatrix, render_dialect_packed_alu, render_dialect_packed_atomic, + render_dialect_packed_conversion, render_dialect_prmt, render_dialect_redux, + render_dialect_register_mma, render_dialect_scalar_arithmetic, + render_dialect_scalar_conversion, render_dialect_scalar_math, render_dialect_sparse_mma, + render_dialect_sreg, render_dialect_stmatrix, render_dialect_sync, render_dialect_tcgen05, + render_dialect_tma, render_dialect_vote, render_dialect_warp_barrier, + render_dialect_warp_match, render_dialect_warp_shuffle, render_dialect_wgmma_control, +}; +use crate::render::families::{ + clc_intrinsics, cluster_memory, debug_controls, elect_intrinsics, execution_controls, + extended_minmax, extended_minmax_contract, integer_minmaxes, mbarrier_extended, movmatrix, + scalar_arithmetics, scalar_conversions, scalar_maths, stmatrices, sync_intrinsics, + tcgen05_intrinsics, threadfence_ptx_level, tma_intrinsics, wgmma_controls, +}; +use crate::render::importer::render_importer; +use crate::render::lowering::render_lowering; +use crate::render::probes::{render_elect_probe, render_special_register_probe}; +use crate::render::raw_abi::{render_raw_abi, render_raw_mod}; +use crate::render::reference::render_reference; +use crate::render::validate::validate_renderable; +use anyhow::Result; +use std::collections::BTreeMap; +use std::path::PathBuf; + +mod collector_targets; +mod common; +mod compat; +mod dialect; +mod families; +mod importer; +mod lowering; +mod probes; +mod raw_abi; +mod reference; +#[cfg(test)] +mod tests; +mod validate; + +pub(crate) use probes::render_probe; + +pub fn all_outputs( + catalog: &CatalogFile, + catalog_json: String, + catalog_sha256: &str, +) -> Result> { + validate_renderable(catalog)?; + let mut outputs = BTreeMap::new(); + outputs.insert("intrinsics/catalog.json".into(), catalog_json); + outputs.insert( + "crates/cuda-intrinsics/src/generated/mod.rs".into(), + render_raw_mod(catalog, catalog_sha256), + ); + outputs.insert( + format!( + "crates/cuda-intrinsics/src/generated/abi_v{}.rs", + catalog.intrinsic_abi + ) + .into(), + render_raw_abi(catalog, catalog_sha256)?, + ); + outputs.insert( + "crates/cuda-device/src/generated/register_mma.rs".into(), + render_compat_register_mma(catalog, catalog_sha256), + ); + outputs.insert( + "crates/cuda-device/src/generated/ldmatrix.rs".into(), + render_compat_ldmatrix(catalog, catalog_sha256), + ); + outputs.insert( + "crates/cuda-device/src/generated/sparse_mma.rs".into(), + render_compat_sparse_mma(catalog, catalog_sha256), + ); + outputs.insert( + "crates/cuda-device/src/generated/sreg.rs".into(), + render_compat_sreg(catalog, catalog_sha256), + ); + for module in ["debug", "grid", "shared", "warp"] { + outputs.insert( + format!("crates/cuda-device/src/generated/{module}_sreg.rs").into(), + render_compat_special_register_module(catalog, catalog_sha256, module), + ); + } + outputs.insert( + "crates/cuda-device/src/generated/cluster_sreg.rs".into(), + render_compat_cluster_sreg(catalog, catalog_sha256), + ); + if sync_intrinsics(catalog).any(|record| threadfence_ptx_level(record).is_some()) { + outputs.insert( + "crates/cuda-device/src/generated/fence.rs".into(), + render_compat_fence(catalog, catalog_sha256), + ); + } + outputs.insert( + "crates/cuda-device/src/generated/dotprod.rs".into(), + render_compat_dotprod(catalog, catalog_sha256), + ); + outputs.insert( + "crates/cuda-device/src/generated/prmt.rs".into(), + render_compat_prmt(catalog, catalog_sha256), + ); + outputs.insert( + "crates/cuda-device/src/generated/cluster_barrier.rs".into(), + render_compat_cluster_barrier(catalog, catalog_sha256), + ); + if cluster_memory(catalog).next().is_some() { + outputs.insert( + "crates/cuda-device/src/generated/cluster_memory.rs".into(), + render_compat_cluster_memory(catalog, catalog_sha256), + ); + } + if stmatrices(catalog).next().is_some() { + outputs.insert( + "crates/cuda-device/src/generated/stmatrix.rs".into(), + render_compat_stmatrix(catalog, catalog_sha256), + ); + } + if movmatrix(catalog).next().is_some() { + outputs.insert( + "crates/cuda-device/src/generated/movmatrix.rs".into(), + render_compat_movmatrix(catalog, catalog_sha256), + ); + } + if wgmma_controls(catalog).next().is_some() { + outputs.insert( + "crates/cuda-device/src/generated/wgmma_control.rs".into(), + render_compat_wgmma_control(catalog, catalog_sha256), + ); + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/wgmma_control.rs".into(), + render_dialect_wgmma_control(catalog, catalog_sha256), + ); + } + outputs.insert( + "crates/cuda-device/src/generated/atomic.rs".into(), + render_compat_packed_atomic(catalog, catalog_sha256), + ); + outputs.insert( + "crates/cuda-device/src/generated/async_copy.rs".into(), + render_compat_cp_async_copy(catalog, catalog_sha256), + ); + outputs.insert( + "crates/cuda-device/src/generated/mbarrier_basic.rs".into(), + render_compat_mbarrier_basic(catalog, catalog_sha256), + ); + if mbarrier_extended(catalog).next().is_some() { + outputs.insert( + "crates/cuda-device/src/generated/mbarrier_extended.rs".into(), + render_compat_mbarrier_extended(catalog, catalog_sha256), + ); + } + outputs.insert( + "crates/cuda-device/src/generated/bf16x2.rs".into(), + render_compat_packed_alu(catalog, catalog_sha256, PackedAluFormat::Bf16x2), + ); + for (module, format) in [ + ("f16", ExtendedMinMaxFormat::F16), + ("bf16", ExtendedMinMaxFormat::Bf16), + ] { + if extended_minmax(catalog).any(|record| extended_minmax_contract(record).format == format) + { + outputs.insert( + format!("crates/cuda-device/src/generated/{module}.rs").into(), + render_compat_scalar_minmax(catalog, catalog_sha256, module), + ); + } + } + outputs.insert( + "crates/cuda-device/src/generated/f16x2.rs".into(), + render_compat_packed_alu(catalog, catalog_sha256, PackedAluFormat::F16x2), + ); + for module in ["i16x2", "int"] { + if integer_minmaxes(catalog).any(|record| record.rust.module == module) { + outputs.insert( + format!("crates/cuda-device/src/generated/{module}.rs").into(), + render_compat_integer_minmax(catalog, catalog_sha256, module), + ); + } + } + if integer_minmaxes(catalog).next().is_some() { + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/integer_minmax.rs".into(), + render_dialect_integer_minmax(catalog, catalog_sha256), + ); + } + outputs.insert( + "crates/cuda-device/src/generated/f32x2.rs".into(), + render_compat_packed_alu(catalog, catalog_sha256, PackedAluFormat::F32x2), + ); + outputs.insert( + "crates/cuda-device/src/generated/convert.rs".into(), + render_compat_packed_conversion( + catalog, + catalog_sha256, + "cuda_device::convert::", + "convert", + ("lo", "hi"), + ), + ); + if let Some((path, source)) = render_compat_float_output(catalog, catalog_sha256) { + outputs.insert(path, source); + } + if extended_minmax(catalog).next().is_some() { + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/extended_minmax.rs".into(), + render_dialect_extended_minmax(catalog, catalog_sha256), + ); + } + if debug_controls(catalog).next().is_some() { + outputs.insert( + "crates/cuda-device/src/generated/debug_control.rs".into(), + render_compat_debug_control(catalog, catalog_sha256), + ); + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/debug_control.rs".into(), + render_dialect_debug_control(catalog, catalog_sha256), + ); + } + if clc_intrinsics(catalog).next().is_some() { + outputs.insert( + "crates/cuda-device/src/generated/clc.rs".into(), + render_compat_clc(catalog, catalog_sha256), + ); + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/clc.rs".into(), + render_dialect_clc(catalog, catalog_sha256), + ); + } + if tma_intrinsics(catalog).next().is_some() { + outputs.insert( + "crates/cuda-device/src/generated/tma.rs".into(), + render_compat_tma(catalog, catalog_sha256), + ); + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/tma.rs".into(), + render_dialect_tma(catalog, catalog_sha256), + ); + } + if execution_controls(catalog).next().is_some() { + outputs.insert( + "crates/cuda-device/src/generated/counted_barrier.rs".into(), + render_compat_counted_barrier(catalog, catalog_sha256), + ); + outputs.insert( + "crates/cuda-device/src/generated/grid_dependency.rs".into(), + render_compat_grid_dependency(catalog, catalog_sha256), + ); + outputs.insert( + "crates/cuda-device/src/generated/register_control.rs".into(), + render_compat_register_control(catalog, catalog_sha256), + ); + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/execution_control.rs".into(), + render_dialect_execution_control(catalog, catalog_sha256), + ); + } + if tcgen05_intrinsics(catalog).next().is_some() { + outputs.insert( + "crates/cuda-device/src/generated/tcgen05.rs".into(), + render_compat_tcgen05(catalog, catalog_sha256), + ); + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/tcgen05.rs".into(), + render_dialect_tcgen05(catalog, catalog_sha256), + ); + } + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/mod.rs".into(), + render_dialect_mod(catalog, catalog_sha256), + ); + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/sreg.rs".into(), + render_dialect_sreg(catalog, catalog_sha256), + ); + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/dotprod.rs".into(), + render_dialect_dotprod(catalog, catalog_sha256), + ); + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/ldmatrix.rs".into(), + render_dialect_ldmatrix(catalog, catalog_sha256), + ); + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/register_mma.rs".into(), + render_dialect_register_mma(catalog, catalog_sha256), + ); + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/sparse_mma.rs".into(), + render_dialect_sparse_mma(catalog, catalog_sha256), + ); + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/packed_atomic.rs".into(), + render_dialect_packed_atomic(catalog, catalog_sha256), + ); + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/packed_alu.rs".into(), + render_dialect_packed_alu(catalog, catalog_sha256), + ); + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/packed_conversion.rs".into(), + render_dialect_packed_conversion(catalog, catalog_sha256), + ); + if scalar_conversions(catalog).next().is_some() { + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/scalar_conversion.rs".into(), + render_dialect_scalar_conversion(catalog, catalog_sha256), + ); + } + if scalar_arithmetics(catalog).next().is_some() { + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/scalar_arithmetic.rs".into(), + render_dialect_scalar_arithmetic(catalog, catalog_sha256), + ); + } + if scalar_maths(catalog).next().is_some() { + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/scalar_math.rs".into(), + render_dialect_scalar_math(catalog, catalog_sha256), + ); + } + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/prmt.rs".into(), + render_dialect_prmt(catalog, catalog_sha256), + ); + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/cluster_barrier.rs".into(), + render_dialect_cluster_barrier(catalog, catalog_sha256), + ); + if cluster_memory(catalog).next().is_some() { + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/cluster_memory.rs".into(), + render_dialect_cluster_memory(catalog, catalog_sha256), + ); + } + if stmatrices(catalog).next().is_some() { + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/stmatrix.rs".into(), + render_dialect_stmatrix(catalog, catalog_sha256), + ); + } + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/redux.rs".into(), + render_dialect_redux(catalog, catalog_sha256), + ); + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/sync.rs".into(), + render_dialect_sync(catalog, catalog_sha256), + ); + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/vote.rs".into(), + render_dialect_vote(catalog, catalog_sha256), + ); + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/active_mask.rs".into(), + render_dialect_active_mask(catalog, catalog_sha256), + ); + if elect_intrinsics(catalog).next().is_some() { + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/elect.rs".into(), + render_dialect_elect(catalog, catalog_sha256), + ); + } + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/cp_async.rs".into(), + render_dialect_cp_async_copy(catalog, catalog_sha256), + ); + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/mbarrier_basic.rs".into(), + render_dialect_mbarrier_basic(catalog, catalog_sha256), + ); + if movmatrix(catalog).next().is_some() { + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/movmatrix.rs".into(), + render_dialect_movmatrix(catalog, catalog_sha256), + ); + } + if mbarrier_extended(catalog).next().is_some() { + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/mbarrier_extended.rs".into(), + render_dialect_mbarrier_extended(catalog, catalog_sha256), + ); + } + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/warp_match.rs".into(), + render_dialect_warp_match(catalog, catalog_sha256), + ); + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/warp_barrier.rs".into(), + render_dialect_warp_barrier(catalog, catalog_sha256), + ); + outputs.insert( + "crates/dialect-nvvm/src/ops/generated/warp_shuffle.rs".into(), + render_dialect_warp_shuffle(catalog, catalog_sha256), + ); + outputs.insert( + "crates/mir-importer/src/translator/terminator/intrinsics/generated.rs".into(), + render_importer(catalog, catalog_sha256), + ); + outputs.insert( + "crates/mir-lower/src/convert/generated_intrinsics.rs".into(), + render_lowering(catalog, catalog_sha256), + ); + outputs.insert( + "crates/rustc-codegen-cuda/src/generated_intrinsics.rs".into(), + render_collector(catalog, catalog_sha256), + ); + outputs.insert( + "crates/cuda-oxide-codegen/src/generated_intrinsic_targets.rs".into(), + render_targets(catalog, catalog_sha256), + ); + for record in &catalog.intrinsics { + if record.special_register.is_some() || record.family == "elect" { + for backend in [IntrinsicBackend::LlvmNvptx, IntrinsicBackend::LibNvvm] { + outputs.insert( + format!( + "intrinsics/probes/{}.{}.ll", + record.id, + backend_label(backend) + ) + .into(), + if record.family == "elect" { + render_elect_probe(catalog, record, catalog_sha256, backend) + } else { + render_special_register_probe(catalog, record, catalog_sha256, backend) + }, + ); + } + } else { + outputs.insert( + format!("intrinsics/probes/{}.ll", record.id).into(), + render_probe(catalog, record, catalog_sha256), + ); + } + } + outputs.insert( + "intrinsics/generated-reference.md".into(), + render_reference(catalog, catalog_sha256), + ); + Ok(outputs) +} diff --git a/crates/cuda-intrinsics-gen/src/render/probes.rs b/crates/cuda-intrinsics-gen/src/render/probes.rs new file mode 100644 index 0000000000..445a14262c --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/render/probes.rs @@ -0,0 +1,2050 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, CatalogFile, CatalogIntrinsic, ClcAdapter, ClusterMemoryOperation, + CpAsyncControlOperation, CpAsyncMbarrierStateSpace, CpAsyncSourceSize, DebugControlOperation, + DotProductAdapter, ExecutionControlOperation, ExtendedMinMaxFormat, IntrinsicBackend, + MbarrierBasicOperation, MbarrierExtendedAdapter, PackedAtomicFormat, + PackedConversionSourceFormat, ReduxAdapter, RegisterMmaAdapter, SparseMmaAccumulator, + SpecialRegisterObservation, Tcgen05LdShape, Tcgen05Mma, Tcgen05MmaBUsage, Tcgen05MmaForm, + Tcgen05MmaKind, Tcgen05Operation, TmaOperation, VoteAdapter, VoteMode, WarpBarrierAdapter, + WarpMatchMode, WarpShuffleAdapter, WarpShuffleMode, WarpShuffleValueKind, WgmmaControlMode, +}; +use crate::render::common::{backend_label, llvm, llvm_header}; +use crate::render::families::{ + extended_minmax_contract, extended_minmax_ptx_mnemonic, integer_minmax_ptx_mnemonic, + movmatrix_template, packed_alu_ptx_mnemonic, packed_alu_register_constraint, packed_alu_width, + packed_conversion_constraint, packed_conversion_dialect_type, packed_conversion_ptx_mnemonic, + packed_conversion_source, packed_conversion_source_width, packed_conversion_typed_llvm_name, + register_mma_constraints, register_mma_fragment_counts, register_mma_template, + scalar_arithmetic_arity, scalar_arithmetic_llvm_mechanism, scalar_arithmetic_llvm_type, + scalar_arithmetic_ptx_mnemonic, scalar_math_llvm_mechanism, scalar_math_llvm_type, + scalar_math_ptx_mnemonic, sparse_mma_constraints, sparse_mma_fragment_counts, + sparse_mma_selector_values, sparse_mma_template, special_register_backend_mechanism, + special_register_inline_template, special_register_output_constraint, stmatrix_variant, + tcgen05_inline_asm, tcgen05_ld_register_count, tcgen05_mma_inline_asm, tcgen05_mma_is_ws, + tcgen05_mma_runtime_parameters, tcgen05_st_register_count, threadfence_ptx_level, +}; +use std::fmt::Write as _; + +pub(super) fn render_special_register_probe( + catalog: &CatalogFile, + record: &CatalogIntrinsic, + hash: &str, + backend: IntrinsicBackend, +) -> String { + let mut output = llvm_header(catalog, hash); + output.push_str("target triple = \"nvptx64-nvidia-cuda\"\n\n"); + let width = record.scalar_width().expect("special-register width"); + let function_name = format!("probe_{}_{}", record.id, backend_label(backend)); + match special_register_backend_mechanism(record, backend) { + BackendLoweringMechanism::TypedNvvm => { + let symbol = &record + .llvm + .as_ref() + .expect("typed special-register route") + .symbol; + writeln!(output, "declare i{width} @{symbol}()\n").unwrap(); + writeln!(output, "define i{width} @{function_name}() {{").unwrap(); + writeln!(output, " %result = call i{width} @{symbol}()").unwrap(); + } + BackendLoweringMechanism::InlinePtx => { + writeln!(output, "define i{width} @{function_name}() {{").unwrap(); + let sideeffect = if record.special_register.as_ref().is_some_and(|special| { + special.observation == SpecialRegisterObservation::VolatileObservation + }) { + " sideeffect" + } else { + "" + }; + writeln!( + output, + " %result = call i{width} asm{sideeffect} {:?}, {:?}()", + special_register_inline_template(record), + special_register_output_constraint(record) + ) + .unwrap(); + } + } + writeln!(output, " ret i{width} %result\n}}").unwrap(); + output +} + +fn render_tma_probe(catalog: &CatalogFile, record: &CatalogIntrinsic, hash: &str) -> String { + let tma = record.tma.as_ref().expect("TMA contract"); + let operation = tma.operation; + let llvm = llvm(record); + let symbol = llvm.resolved_symbol.as_ref().unwrap_or(&llvm.symbol); + let mut output = llvm_header(catalog, hash); + output.push_str("target triple = \"nvptx64-nvidia-cuda\"\n\n"); + if let Some(dimensions) = tma.dimensions() { + let is_g2s = matches!( + operation, + TmaOperation::G2sTile1d + | TmaOperation::G2sTile2d + | TmaOperation::G2sTile2dMulticast + | TmaOperation::G2sTile2dMulticastCg2 + | TmaOperation::G2sTile3d + | TmaOperation::G2sTile4d + | TmaOperation::G2sTile5d + ); + let coordinates = std::iter::repeat_n("i32", dimensions) + .collect::>() + .join(", "); + let coordinate_parameters = (0..dimensions) + .map(|index| format!("i32 %coord{index}")) + .collect::>() + .join(", "); + let coordinate_arguments = (0..dimensions) + .map(|index| format!("i32 %coord{index}")) + .collect::>() + .join(", "); + if is_g2s { + let multicast = matches!( + operation, + TmaOperation::G2sTile2dMulticast | TmaOperation::G2sTile2dMulticastCg2 + ); + let declaration_coordinates = if coordinates.is_empty() { + String::new() + } else { + format!(", {coordinates}") + }; + writeln!( + output, + "declare void @{symbol}(ptr addrspace(7), ptr addrspace(3), ptr{declaration_coordinates}, i16, i64, i1, i1, i32) #0\n" + ) + .unwrap(); + let parameters = if coordinate_parameters.is_empty() { + String::new() + } else { + format!(", {coordinate_parameters}") + }; + let mask_parameter = if multicast { ", i16 %cta_mask" } else { "" }; + writeln!( + output, + "define void @probe_{}(ptr %dst_generic, ptr %barrier_generic, ptr %tensor_map{parameters}{mask_parameter}) #0 {{", + record.id + ) + .unwrap(); + output.push_str( + " %dst = addrspacecast ptr %dst_generic to ptr addrspace(7)\n %barrier = addrspacecast ptr %barrier_generic to ptr addrspace(3)\n", + ); + let arguments = if coordinate_arguments.is_empty() { + String::new() + } else { + format!(", {coordinate_arguments}") + }; + let mask = if multicast { "%cta_mask" } else { "0" }; + let group = if operation == TmaOperation::G2sTile2dMulticastCg2 { + 2 + } else { + 0 + }; + writeln!( + output, + " call void @{symbol}(ptr addrspace(7) %dst, ptr addrspace(3) %barrier, ptr %tensor_map{arguments}, i16 {mask}, i64 0, i1 {multicast}, i1 false, i32 {group}) #0" + ) + .unwrap(); + } else { + let declaration_coordinates = if coordinates.is_empty() { + String::new() + } else { + format!(", {coordinates}") + }; + writeln!( + output, + "declare void @{symbol}(ptr addrspace(3), ptr{declaration_coordinates}, i64, i1) #0\n" + ) + .unwrap(); + let parameters = if coordinate_parameters.is_empty() { + String::new() + } else { + format!(", {coordinate_parameters}") + }; + writeln!( + output, + "define void @probe_{}(ptr %src_generic, ptr %tensor_map{parameters}) #0 {{", + record.id + ) + .unwrap(); + output.push_str(" %src = addrspacecast ptr %src_generic to ptr addrspace(3)\n"); + let arguments = if coordinate_arguments.is_empty() { + String::new() + } else { + format!(", {coordinate_arguments}") + }; + writeln!( + output, + " call void @{symbol}(ptr addrspace(3) %src, ptr %tensor_map{arguments}, i64 0, i1 false) #0" + ) + .unwrap(); + } + output.push_str(" ret void\n}\n\nattributes #0 = { convergent }\n"); + } else if operation == TmaOperation::PrefetchTensorMap { + writeln!(output, "declare void @{symbol}(ptr)\n").unwrap(); + writeln!( + output, + "define void @probe_{}(ptr %tensor_map) {{", + record.id + ) + .unwrap(); + writeln!(output, " call void @{symbol}(ptr %tensor_map)").unwrap(); + output.push_str(" ret void\n}\n"); + } else if let Some(coordinate_count) = operation.prefetch_coordinate_count() { + let coordinates = std::iter::repeat_n("i32", coordinate_count) + .collect::>() + .join(", "); + let declaration_coordinates = if coordinates.is_empty() { + String::new() + } else { + format!(", {coordinates}") + }; + let coordinate_parameters = (0..coordinate_count) + .map(|index| format!("i32 %coord{index}")) + .collect::>() + .join(", "); + let parameters = if coordinate_parameters.is_empty() { + String::new() + } else { + format!(", {coordinate_parameters}") + }; + let cache_hint_parameter = if operation.uses_prefetch_cache_hint() { + ", i64 %cache_hint" + } else { + "" + }; + let coordinate_arguments = (0..coordinate_count) + .map(|index| format!("i32 %coord{index}")) + .collect::>() + .join(", "); + let arguments = if coordinate_arguments.is_empty() { + String::new() + } else { + format!(", {coordinate_arguments}") + }; + writeln!( + output, + "declare void @{symbol}(ptr{declaration_coordinates}, i64, i1) #0\n" + ) + .unwrap(); + writeln!( + output, + "define void @probe_{}(ptr %tensor_map{parameters}{cache_hint_parameter}) #0 {{", + record.id, + ) + .unwrap(); + let cache_hint = if operation.uses_prefetch_cache_hint() { + "%cache_hint" + } else { + "0" + }; + let use_cache_hint = operation.uses_prefetch_cache_hint(); + writeln!( + output, + " call void @{symbol}(ptr %tensor_map{arguments}, i64 {cache_hint}, i1 {use_cache_hint}) #0" + ) + .unwrap(); + output.push_str(" ret void\n}\n\nattributes #0 = { convergent }\n"); + } else if matches!( + operation, + TmaOperation::ReplaceBoxDim + | TmaOperation::ReplaceElementStride + | TmaOperation::ReplaceElementType + | TmaOperation::ReplaceFillMode + | TmaOperation::ReplaceGlobalAddress + | TmaOperation::ReplaceGlobalDim + | TmaOperation::ReplaceGlobalStride + | TmaOperation::ReplaceInterleaveLayout + | TmaOperation::ReplaceRank + | TmaOperation::ReplaceSwizzleAtomicity + | TmaOperation::ReplaceSwizzleMode + ) { + let (field, width, ordinal, value_ty, immediate, address) = match operation { + TmaOperation::ReplaceBoxDim => ("box_dim", "b32", true, "i32", false, false), + TmaOperation::ReplaceElementStride => { + ("element_stride", "b32", true, "i32", false, false) + } + TmaOperation::ReplaceElementType => ("elemtype", "b32", false, "i32", true, false), + TmaOperation::ReplaceFillMode => ("fill_mode", "b32", false, "i32", true, false), + TmaOperation::ReplaceGlobalAddress => { + ("global_address", "b64", false, "i64", false, true) + } + TmaOperation::ReplaceGlobalDim => ("global_dim", "b32", true, "i32", false, false), + TmaOperation::ReplaceGlobalStride => { + ("global_stride", "b64", true, "i64", false, false) + } + TmaOperation::ReplaceInterleaveLayout => { + ("interleave_layout", "b32", false, "i32", true, false) + } + TmaOperation::ReplaceRank => ("rank", "b32", false, "i32", false, false), + TmaOperation::ReplaceSwizzleAtomicity => { + ("swizzle_atomicity", "b32", false, "i32", true, false) + } + TmaOperation::ReplaceSwizzleMode => ("swizzle_mode", "b32", false, "i32", true, false), + _ => unreachable!("TMA tensor-map replace operation was matched"), + }; + let value_parameter = if immediate { + String::new() + } else if address { + ", ptr %new_address".into() + } else { + format!(", {value_ty} %new_value") + }; + writeln!( + output, + "define void @probe_{}(ptr %tensor_map{value_parameter}) {{", + record.id + ) + .unwrap(); + let mut arguments = vec!["ptr %tensor_map".to_owned()]; + let mut constraints = vec!["l"]; + if ordinal { + arguments.push("i32 0".into()); + constraints.push("n"); + } + if immediate { + arguments.push(format!("{value_ty} 0")); + constraints.push("n"); + } else if address { + arguments.push("ptr %new_address".into()); + constraints.push("l"); + } else { + arguments.push(format!("{value_ty} %new_value")); + constraints.push(if value_ty == "i64" { "l" } else { "r" }); + } + constraints.push("~{memory}"); + let template = if ordinal { + format!("tensormap.replace.tile.{field}.global.b1024.{width} [$0], $1, $2;") + } else { + format!("tensormap.replace.tile.{field}.global.b1024.{width} [$0], $1;") + }; + writeln!( + output, + " call void asm sideeffect {template:?}, {:?}({})", + constraints.join(","), + arguments.join(", ") + ) + .unwrap(); + output.push_str(" ret void\n}\n"); + } else if matches!( + operation, + TmaOperation::FenceProxyTensorMapAcquireCluster + | TmaOperation::FenceProxyTensorMapAcquireCta + | TmaOperation::FenceProxyTensorMapAcquireGpu + | TmaOperation::FenceProxyTensorMapAcquireSystem + ) { + writeln!(output, "declare void @{symbol}(ptr, i32)\n").unwrap(); + writeln!( + output, + "define void @probe_{}(ptr %tensor_map) {{", + record.id + ) + .unwrap(); + writeln!(output, " call void @{symbol}(ptr %tensor_map, i32 128)").unwrap(); + output.push_str(" ret void\n}\n"); + } else if operation == TmaOperation::CommitGroup + || matches!( + operation, + TmaOperation::FenceProxyTensorMapReleaseCluster + | TmaOperation::FenceProxyTensorMapReleaseCta + | TmaOperation::FenceProxyTensorMapReleaseGpu + | TmaOperation::FenceProxyTensorMapReleaseSystem + ) + { + writeln!(output, "declare void @{symbol}()\n").unwrap(); + writeln!(output, "define void @probe_{}() {{", record.id).unwrap(); + writeln!(output, " call void @{symbol}()").unwrap(); + output.push_str(" ret void\n}\n"); + } else { + writeln!(output, "declare void @{symbol}(i32)\n").unwrap(); + writeln!(output, "define void @probe_{}() {{", record.id).unwrap(); + writeln!(output, " call void @{symbol}(i32 0)").unwrap(); + output.push_str(" ret void\n}\n"); + } + output +} + +fn render_execution_control_probe( + catalog: &CatalogFile, + record: &CatalogIntrinsic, + hash: &str, +) -> String { + let operation = ExecutionControlOperation::from_catalog_id(&record.id) + .expect("closed execution-control operation"); + let symbol = &llvm(record).symbol; + let mut output = llvm_header(catalog, hash); + output.push_str("target triple = \"nvptx64-nvidia-cuda\"\n\n"); + match operation { + ExecutionControlOperation::BarrierCtaSync + | ExecutionControlOperation::BarrierCtaSyncAligned + | ExecutionControlOperation::BarrierCtaArrive + | ExecutionControlOperation::BarrierCtaArriveAligned => { + writeln!(output, "declare void @{symbol}(i32, i32)\n").unwrap(); + for (suffix, parameters, barrier_id, thread_count) in [ + ( + "rr", + "i32 %barrier_id, i32 %thread_count", + "%barrier_id", + "%thread_count", + ), + ("ri", "i32 %barrier_id", "%barrier_id", "32"), + ("ir", "i32 %thread_count", "1", "%thread_count"), + ("ii", "", "1", "32"), + ] { + writeln!( + output, + "define void @probe_{}_{suffix}({parameters}) #0 {{", + record.id + ) + .unwrap(); + writeln!( + output, + " call void @{symbol}(i32 {barrier_id}, i32 {thread_count})" + ) + .unwrap(); + output.push_str(" ret void\n}\n\n"); + } + output.push_str("attributes #0 = { convergent }\n"); + } + ExecutionControlOperation::GridDependencyLaunchDependents + | ExecutionControlOperation::GridDependencyWait => { + writeln!(output, "declare void @{symbol}()\n").unwrap(); + writeln!(output, "define void @probe_{}() {{", record.id).unwrap(); + writeln!(output, " call void @{symbol}()").unwrap(); + output.push_str(" ret void\n}\n"); + } + ExecutionControlOperation::SetMaxNRegInc | ExecutionControlOperation::SetMaxNRegDec => { + writeln!(output, "declare void @{symbol}(i32)\n").unwrap(); + writeln!(output, "define void @probe_{}() #0 {{", record.id).unwrap(); + writeln!(output, " call void @{symbol}(i32 64)").unwrap(); + output.push_str(" ret void\n}\n\nattributes #0 = { convergent }\n"); + } + } + output +} + +fn render_tcgen05_mma_probe( + catalog: &CatalogFile, + record: &CatalogIntrinsic, + hash: &str, + mma: &Tcgen05Mma, +) -> String { + let (kind, b_buffer, b_usage) = mma.fixed_selectors.map_or( + (Tcgen05MmaKind::F16, 0, Tcgen05MmaBUsage::Discard), + |fixed| (fixed.kind, fixed.b_buffer, fixed.b_usage), + ); + let (template, constraints) = if tcgen05_mma_is_ws(mma.form) { + tcgen05_mma_inline_asm(mma.form, kind, 1, None, Some(b_buffer), Some(b_usage)) + } else { + tcgen05_mma_inline_asm(mma.form, kind, 1, Some("discard"), None, None) + }; + let parameters = tcgen05_mma_runtime_parameters(mma); + let llvm_parameter = |(name, ty): &(&str, &str)| { + let ty = match *ty { + "u32" => "i32", + "u64" => "i64", + "bool" => "i1", + _ => unreachable!("closed tcgen05 MMA probe type"), + }; + format!("{ty} %{name}") + }; + let signature = parameters + .iter() + .map(llvm_parameter) + .collect::>() + .join(", "); + let argument_indices: Vec = + if mma.alias.is_some() && mma.form == Tcgen05MmaForm::WsTensor { + vec![0, 1, 3, 4, 5] + } else { + (0..parameters.len()).collect() + }; + let arguments = argument_indices + .iter() + .map(|&index| llvm_parameter(¶meters[index])) + .collect::>() + .join(", "); + + let mut output = llvm_header(catalog, hash); + output.push_str("target triple = \"nvptx64-nvidia-cuda\"\n\n"); + writeln!( + output, + "; LLVM source: {}; route: inline PTX; source contract: {:?}", + llvm(record).symbol, + record.tcgen05.as_ref().unwrap().source_contract + ) + .unwrap(); + writeln!( + output, + "define void @probe_{}({signature}) #0 {{", + record.id + ) + .unwrap(); + writeln!( + output, + " call void asm sideeffect {template:?}, {constraints:?}({arguments}) #0" + ) + .unwrap(); + output.push_str(" ret void\n}\n\nattributes #0 = { convergent }\n"); + output +} + +fn render_tcgen05_probe(catalog: &CatalogFile, record: &CatalogIntrinsic, hash: &str) -> String { + let tcgen05 = record.tcgen05.as_ref().expect("tcgen05 contract"); + if let Some(mma) = &tcgen05.mma { + return render_tcgen05_mma_probe(catalog, record, hash, mma); + } + let operation = tcgen05.operation; + let mut output = llvm_header(catalog, hash); + output.push_str("target triple = \"nvptx64-nvidia-cuda\"\n\n"); + writeln!( + output, + "; LLVM source: {}; route: inline PTX; source contract: {:?}", + llvm(record).symbol, + record.tcgen05.as_ref().unwrap().source_contract + ) + .unwrap(); + let (parameters, result_ty, template, constraints, arguments) = match operation { + Tcgen05Operation::Alloc | Tcgen05Operation::AllocCg2 => { + let group = if operation == Tcgen05Operation::AllocCg2 { + "2" + } else { + "1" + }; + ( + "ptr %dst, i32 %ncols".into(), + "void".into(), + format!( + "{{ .reg .u64 %shared64; .reg .u32 %shared32; cvta.to.shared.u64 %shared64, $0; cvt.u32.u64 %shared32, %shared64; tcgen05.alloc.cta_group::{group}.sync.aligned.shared::cta.b32 [%shared32], $1; }}" + ), + "l,r,~{memory}".into(), + "ptr %dst, i32 %ncols".into(), + ) + } + Tcgen05Operation::Dealloc | Tcgen05Operation::DeallocCg2 => { + let group = if operation == Tcgen05Operation::DeallocCg2 { + "2" + } else { + "1" + }; + ( + "i32 %tmem, i32 %ncols".into(), + "void".into(), + format!("tcgen05.dealloc.cta_group::{group}.sync.aligned.b32 $0, $1;"), + "r,r,~{memory}".into(), + "i32 %tmem, i32 %ncols".into(), + ) + } + Tcgen05Operation::RelinquishAllocPermit | Tcgen05Operation::RelinquishAllocPermitCg2 => { + let group = if operation == Tcgen05Operation::RelinquishAllocPermitCg2 { + "2" + } else { + "1" + }; + ( + String::new(), + "void".into(), + format!("tcgen05.relinquish_alloc_permit.cta_group::{group}.sync.aligned;"), + "~{memory}".into(), + String::new(), + ) + } + Tcgen05Operation::FenceBeforeThreadSync | Tcgen05Operation::FenceAfterThreadSync => { + let modifier = if operation == Tcgen05Operation::FenceBeforeThreadSync { + "before_thread_sync" + } else { + "after_thread_sync" + }; + ( + String::new(), + "void".into(), + format!("tcgen05.fence::{modifier};"), + "~{memory}".into(), + String::new(), + ) + } + Tcgen05Operation::Commit + | Tcgen05Operation::CommitSharedCluster + | Tcgen05Operation::CommitCg2 + | Tcgen05Operation::CommitSharedClusterCg2 => { + let group = if matches!( + operation, + Tcgen05Operation::CommitCg2 | Tcgen05Operation::CommitSharedClusterCg2 + ) { + "2" + } else { + "1" + }; + let shared = matches!( + operation, + Tcgen05Operation::CommitSharedCluster | Tcgen05Operation::CommitSharedClusterCg2 + ); + ( + "i32 %mbar".into(), + "void".into(), + format!( + "tcgen05.commit.cta_group::{group}.mbarrier::arrive::one{}.b64 [$0];", + if shared { ".shared::cluster" } else { "" } + ), + "r,~{memory}".into(), + "i32 %mbar".into(), + ) + } + Tcgen05Operation::MmaWsF16 | Tcgen05Operation::MmaWsBf16 | Tcgen05Operation::MmaWsTf32 => { + let kind = if operation == Tcgen05Operation::MmaWsTf32 { + "tf32" + } else { + "f16" + }; + ( + "i32 %d, i32 %a, i64 %legacy_a_desc, i64 %b, i32 %idesc, i1 %enable".into(), + "void".into(), + format!( + "{{ .reg .pred %enable_pred; setp.ne.s32 %enable_pred, $5, 0; tcgen05.mma.ws.cta_group::1.kind::{kind} [$0], [$1], $3, $4, %enable_pred; }}" + ), + "r,r,l,l,r,r,~{memory}".into(), + "i32 %d, i32 %a, i64 %legacy_a_desc, i64 %b, i32 %idesc, i1 %enable".into(), + ) + } + Tcgen05Operation::MmaF16 | Tcgen05Operation::MmaF16Cg2 => { + let group = if operation == Tcgen05Operation::MmaF16Cg2 { + "2" + } else { + "1" + }; + let zeros = if operation == Tcgen05Operation::MmaF16Cg2 { + "%z, %z, %z, %z, %z, %z, %z, %z" + } else { + "%z, %z, %z, %z" + }; + ( + "i32 %d, i64 %a, i64 %b, i32 %idesc, i1 %enable".into(), + "void".into(), + format!( + "{{ .reg .pred %enable_pred; setp.ne.s32 %enable_pred, $4, 0; .reg .u32 %z; mov.u32 %z, 0; tcgen05.mma.cta_group::{group}.kind::f16 [$0], $1, $2, $3, {{{zeros}}}, %enable_pred; }}" + ), + "r,l,l,r,r,~{memory}".into(), + "i32 %d, i64 %a, i64 %b, i32 %idesc, i1 %enable".into(), + ) + } + Tcgen05Operation::CpSmemToTmem | Tcgen05Operation::CpSmemToTmemCg2 => ( + "i32 %tmem, i64 %desc".into(), + "void".into(), + format!( + "{}.{} [$0], $1;", + record.expected_ptx.mnemonic, + record.expected_ptx.modifiers.join(".") + ), + "r,l,~{memory}".into(), + "i32 %tmem, i64 %desc".into(), + ), + Tcgen05Operation::Ld16x256bX8Pure | Tcgen05Operation::Ld16x256bPure => { + let count = if operation == Tcgen05Operation::Ld16x256bX8Pure { + 32 + } else { + 4 + }; + let registers = (0..count) + .map(|index| format!("${index}")) + .collect::>() + .join(","); + let constraints = std::iter::repeat_n("=f", count) + .chain(std::iter::once("r")) + .collect::>() + .join(","); + let fields = std::iter::repeat_n("float", count) + .collect::>() + .join(", "); + ( + "i32 %tmem".into(), + format!("{{ {fields} }}"), + format!( + "tcgen05.ld.sync.aligned.16x256b.{}.b32 {{{registers}}}, [${count}];", + if count == 32 { "x8" } else { "x1" } + ), + constraints, + "i32 %tmem".into(), + ) + } + Tcgen05Operation::St => { + let count = tcgen05_st_register_count(record); + let has_half_split_offset = record + .tcgen05 + .as_ref() + .and_then(|tcgen05| tcgen05.st) + .is_some_and(|st| st.shape == Tcgen05LdShape::M16x32bx2); + let parameters = std::iter::once("i32 %tmem".into()) + .chain((0..count).map(|index| format!("i32 %d{index}"))) + .collect::>() + .join(", "); + let arguments = std::iter::once("i32 %tmem".into()) + .chain(has_half_split_offset.then(|| "i64 16".into())) + .chain((0..count).map(|index| format!("i32 %d{index}"))) + .collect::>() + .join(", "); + let (template, constraints, _) = tcgen05_inline_asm(record); + (parameters, "void".into(), template, constraints, arguments) + } + Tcgen05Operation::Ld => { + let count = tcgen05_ld_register_count(record); + let has_half_split_offset = record + .tcgen05 + .as_ref() + .and_then(|tcgen05| tcgen05.ld) + .is_some_and(|ld| ld.shape == Tcgen05LdShape::M16x32bx2); + let fields = std::iter::repeat_n("i32", count) + .collect::>() + .join(", "); + let (template, constraints, _) = tcgen05_inline_asm(record); + ( + "i32 %tmem".into(), + if count == 1 { + "i32".into() + } else { + format!("{{ {fields} }}") + }, + template, + constraints, + if has_half_split_offset { + "i32 %tmem, i64 16".into() + } else { + "i32 %tmem".into() + }, + ) + } + Tcgen05Operation::LoadWait | Tcgen05Operation::StoreWait => { + let kind = if operation == Tcgen05Operation::LoadWait { + "ld" + } else { + "st" + }; + ( + String::new(), + "void".into(), + format!("tcgen05.wait::{kind}.sync.aligned;"), + "~{memory}".into(), + String::new(), + ) + } + Tcgen05Operation::CommitMulticast | Tcgen05Operation::CommitMulticastCg2 => { + let group = if operation == Tcgen05Operation::CommitMulticastCg2 { + 2 + } else { + 1 + }; + ( + "i32 %mbar, i16 %mask".into(), + "void".into(), + format!( + "tcgen05.commit.cta_group::{group}.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 [$0], $1;" + ), + "r,h,~{memory}".into(), + "i32 %mbar, i16 %mask".into(), + ) + } + Tcgen05Operation::ShiftDown | Tcgen05Operation::ShiftDownCg2 => { + let group = if operation == Tcgen05Operation::ShiftDownCg2 { + 2 + } else { + 1 + }; + ( + "i32 %tmem".into(), + "void".into(), + format!("tcgen05.shift.cta_group::{group}.down [$0];"), + "r,~{memory}".into(), + "i32 %tmem".into(), + ) + } + Tcgen05Operation::Mma => unreachable!("generic MMA probe handled above"), + }; + let (lowering_template, lowering_constraints, lowering_results) = tcgen05_inline_asm(record); + assert_eq!(template, lowering_template); + assert_eq!(constraints, lowering_constraints); + assert_eq!(result_ty != "void", lowering_results.is_some()); + writeln!( + output, + "define void @probe_{}({parameters}) #0 {{", + record.id + ) + .unwrap(); + if result_ty == "void" { + writeln!( + output, + " call void asm sideeffect {template:?}, {constraints:?}({arguments}) #0" + ) + .unwrap(); + } else { + writeln!( + output, + " %result = call {result_ty} asm sideeffect {template:?}, {constraints:?}({arguments}) #0" + ) + .unwrap(); + } + output.push_str(" ret void\n}\n\nattributes #0 = { convergent }\n"); + output +} + +pub(super) fn render_elect_probe( + catalog: &CatalogFile, + record: &CatalogIntrinsic, + hash: &str, + backend: IntrinsicBackend, +) -> String { + let mut output = llvm_header(catalog, hash); + output.push_str("target triple = \"nvptx64-nvidia-cuda\"\n\n"); + let mechanism = record + .backend_lowerings + .iter() + .find(|route| route.backend == backend) + .expect("elect backend route") + .mechanism; + match mechanism { + BackendLoweringMechanism::TypedNvvm => { + writeln!( + output, + "declare {{ i32, i1 }} @{}(i32) #0\n", + llvm(record).symbol + ) + .unwrap(); + for (suffix, parameters, mask) in + [("register", "i32 %mask", "%mask"), ("immediate", "", "-1")] + { + writeln!( + output, + "define {{ i32, i1 }} @probe_{}_{}_{suffix}({parameters}) #0 {{", + record.id, + backend_label(backend) + ) + .unwrap(); + writeln!( + output, + " %result = call {{ i32, i1 }} @{}(i32 {mask}) #0", + llvm(record).symbol + ) + .unwrap(); + output.push_str(" ret { i32, i1 } %result\n}\n"); + } + } + BackendLoweringMechanism::InlinePtx => { + for (suffix, parameters, mask) in [ + ("register", "i32 %mask", "%mask"), + ("constant_mask", "", "-1"), + ] { + writeln!( + output, + "define {{ i32, i1 }} @probe_{}_{}_{suffix}({parameters}) #0 {{", + record.id, + backend_label(backend) + ) + .unwrap(); + writeln!( + output, + " %raw = call {{ i32, i32 }} asm sideeffect \"{{ .reg .pred p; elect.sync $0|p, $2; selp.b32 $1, 1, 0, p; }}\", \"=r,=r,r\"(i32 {mask}) #0" + ) + .unwrap(); + output.push_str( + " %leader = extractvalue { i32, i32 } %raw, 0\n\ + \x20 %predicate.i32 = extractvalue { i32, i32 } %raw, 1\n\ + \x20 %predicate = trunc i32 %predicate.i32 to i1\n\ + \x20 %result.0 = insertvalue { i32, i1 } poison, i32 %leader, 0\n\ + \x20 %result.1 = insertvalue { i32, i1 } %result.0, i1 %predicate, 1\n\ + \x20 ret { i32, i1 } %result.1\n}\n", + ); + } + } + } + output.push_str("\nattributes #0 = { convergent }\n"); + output +} + +pub(crate) fn render_probe(catalog: &CatalogFile, record: &CatalogIntrinsic, hash: &str) -> String { + if record.special_register.is_some() { + return render_special_register_probe(catalog, record, hash, IntrinsicBackend::LlvmNvptx); + } + if record.tma.is_some() { + return render_tma_probe(catalog, record, hash); + } + if ExecutionControlOperation::from_catalog_id(&record.id).is_some() { + return render_execution_control_probe(catalog, record, hash); + } + if record.tcgen05.is_some() { + return render_tcgen05_probe(catalog, record, hash); + } + if record.family == "elect" { + return render_elect_probe(catalog, record, hash, IntrinsicBackend::LlvmNvptx); + } + let mut output = llvm_header(catalog, hash); + output.push_str("target triple = \"nvptx64-nvidia-cuda\"\n\n"); + if let Some(cluster) = &record.cluster_memory { + let (template, constraints) = + crate::resolve::cluster_memory_inline_recipe(cluster.operation); + match cluster.operation { + ClusterMemoryOperation::MapSharedRank => { + writeln!( + output, + "define ptr addrspace(7) @probe_{}(ptr %source_generic, i32 %rank) #0 {{", + record.id + ) + .unwrap(); + output.push_str( + " %source = addrspacecast ptr %source_generic to ptr addrspace(3)\n", + ); + writeln!( + output, + " %mapped_integer = call i64 asm sideeffect {template:?}, {constraints:?}(ptr addrspace(3) %source, i32 %rank) #0" + ) + .unwrap(); + output.push_str( + " %mapped = inttoptr i64 %mapped_integer to ptr addrspace(7)\n ret ptr addrspace(7) %mapped\n}\n\nattributes #0 = { convergent }\n", + ); + } + ClusterMemoryOperation::ReadU32 => { + writeln!( + output, + "define i32 @probe_{}(ptr %source_generic, i32 %rank) #0 {{", + record.id + ) + .unwrap(); + output.push_str( + " %source = addrspacecast ptr %source_generic to ptr addrspace(3)\n", + ); + writeln!( + output, + " %value = call i32 asm sideeffect {template:?}, {constraints:?}(ptr addrspace(3) %source, i32 %rank) #0" + ) + .unwrap(); + output.push_str(" ret i32 %value\n}\n\nattributes #0 = { convergent }\n"); + } + } + } else if let Some(clc) = &record.clc { + match clc.adapter { + ClcAdapter::GenericPointersToShared => { + writeln!( + output, + "declare void @{}(ptr addrspace(3), ptr addrspace(3))\n", + llvm(record).symbol + ) + .unwrap(); + writeln!( + output, + "define void @probe_{}(ptr %response_generic, ptr %mbarrier_generic) {{", + record.id + ) + .unwrap(); + output.push_str( + " %response = addrspacecast ptr %response_generic to ptr addrspace(3)\n\ + \x20 %mbarrier = addrspacecast ptr %mbarrier_generic to ptr addrspace(3)\n", + ); + writeln!( + output, + " call void @{}(ptr addrspace(3) %response, ptr addrspace(3) %mbarrier)", + llvm(record).symbol + ) + .unwrap(); + output.push_str(" ret void\n}\n"); + } + ClcAdapter::PairU64ToI128BoolToU32 | ClcAdapter::PairU64ToI128U32 => { + let llvm_result = if clc.adapter == ClcAdapter::PairU64ToI128BoolToU32 { + "i1" + } else { + "i32" + }; + writeln!( + output, + "declare {llvm_result} @{}(i128)\n", + llvm(record).symbol + ) + .unwrap(); + writeln!( + output, + "define i32 @probe_{}(i64 %response_low, i64 %response_high) {{", + record.id + ) + .unwrap(); + output.push_str( + " %response_low_i128 = zext i64 %response_low to i128\n\ + \x20 %response_high_i128 = zext i64 %response_high to i128\n\ + \x20 %response_high_shifted = shl i128 %response_high_i128, 64\n\ + \x20 %response = or i128 %response_low_i128, %response_high_shifted\n", + ); + writeln!( + output, + " %raw_result = call {llvm_result} @{}(i128 %response)", + llvm(record).symbol + ) + .unwrap(); + if clc.adapter == ClcAdapter::PairU64ToI128BoolToU32 { + output.push_str( + " %result = zext i1 %raw_result to i32\n\ + \x20 ret i32 %result\n}\n", + ); + } else { + output.push_str(" ret i32 %raw_result\n}\n"); + } + } + } + } else if let Some(bridge) = &record.cp_async_mbarrier { + let shared = bridge.state_space == CpAsyncMbarrierStateSpace::Shared; + let pointer_type = if shared { "ptr addrspace(3)" } else { "ptr" }; + writeln!( + output, + "declare void @{}({pointer_type})\n", + llvm(record).symbol + ) + .unwrap(); + writeln!( + output, + "define void @probe_{}(ptr %barrier_generic) {{", + record.id + ) + .unwrap(); + if shared { + output + .push_str(" %barrier = addrspacecast ptr %barrier_generic to ptr addrspace(3)\n"); + writeln!( + output, + " call void @{}(ptr addrspace(3) %barrier)", + llvm(record).symbol + ) + .unwrap(); + } else { + writeln!( + output, + " call void @{}(ptr %barrier_generic)", + llvm(record).symbol + ) + .unwrap(); + } + output.push_str(" ret void\n}\n"); + } else if let Some(mbarrier) = &record.mbarrier_basic { + match mbarrier.operation { + MbarrierBasicOperation::Init => { + writeln!( + output, + "declare void @{}(ptr addrspace(3), i32)\n", + llvm(record).symbol + ) + .unwrap(); + writeln!( + output, + "define void @probe_{}(ptr %barrier_generic, i32 %expected_count) {{", + record.id + ) + .unwrap(); + output.push_str( + " %barrier = addrspacecast ptr %barrier_generic to ptr addrspace(3)\n", + ); + writeln!( + output, + " call void @{}(ptr addrspace(3) %barrier, i32 %expected_count)", + llvm(record).symbol + ) + .unwrap(); + output.push_str(" ret void\n}\n"); + } + MbarrierBasicOperation::Arrive => { + writeln!( + output, + "declare i64 @{}(ptr addrspace(3))\n", + llvm(record).symbol + ) + .unwrap(); + writeln!( + output, + "define i64 @probe_{}(ptr %barrier_generic) {{", + record.id + ) + .unwrap(); + output.push_str( + " %barrier = addrspacecast ptr %barrier_generic to ptr addrspace(3)\n", + ); + writeln!( + output, + " %token = call i64 @{}(ptr addrspace(3) %barrier)", + llvm(record).symbol + ) + .unwrap(); + output.push_str(" ret i64 %token\n}\n"); + } + MbarrierBasicOperation::TestWait => { + writeln!( + output, + "define i1 @probe_{}(ptr %barrier_generic, i64 %token) #0 {{", + record.id + ) + .unwrap(); + output.push_str( + " %barrier = addrspacecast ptr %barrier_generic to ptr addrspace(3)\n\ + \x20 %result_i32 = call i32 asm sideeffect \"{ .reg .pred %p0; mbarrier.test_wait.shared.b64 %p0, [$1], $2; selp.b32 $0, 1, 0, %p0; }\", \"=r,l,l,~{memory}\"(ptr addrspace(3) %barrier, i64 %token) #0\n\ + \x20 %result = trunc i32 %result_i32 to i1\n\ + \x20 ret i1 %result\n\ + }\n\n\ + attributes #0 = { convergent }\n", + ); + } + MbarrierBasicOperation::Inval => { + writeln!( + output, + "declare void @{}(ptr addrspace(3))\n", + llvm(record).symbol + ) + .unwrap(); + writeln!( + output, + "define void @probe_{}(ptr %barrier_generic) {{", + record.id + ) + .unwrap(); + output.push_str( + " %barrier = addrspacecast ptr %barrier_generic to ptr addrspace(3)\n", + ); + writeln!( + output, + " call void @{}(ptr addrspace(3) %barrier)", + llvm(record).symbol + ) + .unwrap(); + output.push_str(" ret void\n}\n"); + } + } + } else if record.movmatrix.is_some() { + writeln!(output, "define i32 @probe_{}(i32 %value) #0 {{", record.id).unwrap(); + writeln!( + output, + " %result = call i32 asm {:?}, \"=r,r\"(i32 %value) #0", + movmatrix_template(record), + ) + .unwrap(); + output.push_str(" ret i32 %result\n}\n\nattributes #0 = { convergent }\n"); + } else if let Some(mbarrier) = &record.mbarrier_extended { + let (template, constraints) = + crate::resolve::mbarrier_extended_inline_recipe(mbarrier.operation); + match mbarrier.adapter { + MbarrierExtendedAdapter::PointerTxCountBytesToTokenDroppingTxCount => { + writeln!( + output, + "define i64 @probe_{}(ptr %barrier_generic, i32 %bytes) #0 {{", + record.id + ) + .unwrap(); + output.push_str( + " %barrier = addrspacecast ptr %barrier_generic to ptr addrspace(3)\n", + ); + writeln!( + output, + " %token = call i64 asm sideeffect {template:?}, {constraints:?}(ptr addrspace(3) %barrier, i32 %bytes) #0" + ) + .unwrap(); + output.push_str(" ret i64 %token\n}\n\nattributes #0 = { convergent }\n"); + } + MbarrierExtendedAdapter::PointerTokenToPredicate => { + writeln!( + output, + "define i1 @probe_{}(ptr %barrier_generic, i64 %token) #0 {{", + record.id + ) + .unwrap(); + output.push_str( + " %barrier = addrspacecast ptr %barrier_generic to ptr addrspace(3)\n", + ); + writeln!( + output, + " %result_i32 = call i32 asm sideeffect {template:?}, {constraints:?}(ptr addrspace(3) %barrier, i64 %token) #0" + ) + .unwrap(); + output.push_str( + " %result = trunc i32 %result_i32 to i1\n ret i1 %result\n}\n\nattributes #0 = { convergent }\n", + ); + } + MbarrierExtendedAdapter::PointerParityToPredicate => { + writeln!( + output, + "define i1 @probe_{}(ptr %barrier_generic, i32 %parity) #0 {{", + record.id + ) + .unwrap(); + output.push_str( + " %barrier = addrspacecast ptr %barrier_generic to ptr addrspace(3)\n", + ); + writeln!( + output, + " %result_i32 = call i32 asm sideeffect {template:?}, {constraints:?}(ptr addrspace(3) %barrier, i32 %parity) #0" + ) + .unwrap(); + output.push_str( + " %result = trunc i32 %result_i32 to i1\n ret i1 %result\n}\n\nattributes #0 = { convergent }\n", + ); + } + MbarrierExtendedAdapter::RawClusterAddressToVoid => { + writeln!( + output, + "define void @probe_{}(i64 %address) #0 {{", + record.id + ) + .unwrap(); + writeln!( + output, + " call void asm sideeffect {template:?}, {constraints:?}(i64 %address) #0" + ) + .unwrap(); + output.push_str(" ret void\n}\n\nattributes #0 = { convergent }\n"); + } + MbarrierExtendedAdapter::ZeroOperandsToVoid => { + writeln!(output, "define void @probe_{}() #0 {{", record.id).unwrap(); + writeln!( + output, + " call void asm sideeffect {template:?}, {constraints:?}() #0" + ) + .unwrap(); + output.push_str(" ret void\n}\n\nattributes #0 = { convergent }\n"); + } + MbarrierExtendedAdapter::NanosecondsToVoid => { + writeln!(output, "define void @probe_{}(i32 %ns) #0 {{", record.id).unwrap(); + writeln!( + output, + " call void asm sideeffect {template:?}, {constraints:?}(i32 %ns) #0" + ) + .unwrap(); + output.push_str(" ret void\n}\n\nattributes #0 = { convergent }\n"); + } + } + } else if record.integer_minmax.is_some() { + let parameters = "i32 %arg0, i32 %arg1"; + writeln!(output, "define i32 @probe_{}({parameters}) {{", record.id).unwrap(); + writeln!( + output, + " %result = call i32 asm \"{} $0, $1, $2;\", \"=r,r,r\"(i32 %arg0, i32 %arg1)", + integer_minmax_ptx_mnemonic(record) + ) + .unwrap(); + output.push_str(" ret i32 %result\n}\n"); + } else if record.packed_alu.is_some() { + let arity = record.rust.arguments.len(); + let width = packed_alu_width(record); + let llvm_type = format!("i{width}"); + let register_constraint = packed_alu_register_constraint(record); + let parameters = (0..arity) + .map(|index| format!("{llvm_type} %arg{index}")) + .collect::>() + .join(", "); + let arguments = (0..arity) + .map(|index| format!("{llvm_type} %arg{index}")) + .collect::>() + .join(", "); + let operands = (0..=arity) + .map(|index| format!("${index}")) + .collect::>() + .join(", "); + let constraints = std::iter::once(format!("={register_constraint}")) + .chain(std::iter::repeat_n(register_constraint.to_owned(), arity)) + .collect::>() + .join(","); + writeln!( + output, + "define {llvm_type} @probe_{}({parameters}) {{", + record.id + ) + .unwrap(); + writeln!( + output, + " %result = call {llvm_type} asm \"{} {operands};\", \"{constraints}\"({arguments})", + packed_alu_ptx_mnemonic(record) + ) + .unwrap(); + writeln!(output, " ret {llvm_type} %result\n}}").unwrap(); + } else if record.packed_conversion.is_some() { + let result_ty = packed_conversion_dialect_type(record); + if packed_conversion_typed_llvm_name(record).is_some() { + let llvm = llvm(record); + let symbol = llvm.resolved_symbol.as_deref().unwrap_or(&llvm.symbol); + writeln!(output, "declare {result_ty} @{symbol}(float, float)\n").unwrap(); + writeln!( + output, + "define {result_ty} @probe_{}(float %low, float %high) {{", + record.id, + ) + .unwrap(); + writeln!( + output, + " %result = call {result_ty} @{symbol}(float %high, float %low)" + ) + .unwrap(); + } else if packed_conversion_source(record) == PackedConversionSourceFormat::F32x2 { + writeln!( + output, + "define {result_ty} @probe_{}(float %low, float %high) {{", + record.id, + ) + .unwrap(); + writeln!( + output, + " %result = call {result_ty} asm \"{} $0, $2, $1;\", \"{}\"(float %low, float %high)", + packed_conversion_ptx_mnemonic(record), + packed_conversion_constraint(record), + ) + .unwrap(); + } else { + // A packed source arrives in one integer register, so the probe + // takes a single operand and needs no reordering. + let source_ty = format!("i{}", packed_conversion_source_width(record)); + writeln!( + output, + "define {result_ty} @probe_{}({source_ty} %packed) {{", + record.id, + ) + .unwrap(); + writeln!( + output, + " %result = call {result_ty} asm \"{} $0, $1;\", \"{}\"({source_ty} %packed)", + packed_conversion_ptx_mnemonic(record), + packed_conversion_constraint(record), + ) + .unwrap(); + } + writeln!(output, " ret {result_ty} %result\n}}").unwrap(); + } else if let Some(packed) = &record.packed_atomic { + let format = match packed.format { + PackedAtomicFormat::F16x2 => "f16x2", + PackedAtomicFormat::Bf16x2 => "bf16x2", + }; + writeln!( + output, + "define i32 @probe_{}(ptr %address, i32 %addend) {{", + record.id + ) + .unwrap(); + writeln!( + output, + " %old = call i32 asm sideeffect \"atom.global.add.noftz.{format} $0, [$1], $2;\", \"=r,l,r,~{{memory}}\"(ptr %address, i32 %addend)" + ) + .unwrap(); + output.push_str(" ret i32 %old\n}\n"); + } else if let Some(copy) = &record.cp_async_copy { + let dynamic_source_size = copy.source_size == CpAsyncSourceSize::Runtime; + let declaration_arguments = if dynamic_source_size { + "ptr addrspace(3), ptr addrspace(1), i32" + } else { + "ptr addrspace(3), ptr addrspace(1)" + }; + writeln!( + output, + "declare void @{}({declaration_arguments})", + llvm(record).symbol + ) + .unwrap(); + output.push('\n'); + if dynamic_source_size { + writeln!( + output, + "define void @probe_{}_register(ptr %shared_generic, ptr %global_generic, i32 %source_size) {{", + record.id + ) + .unwrap(); + output.push_str( + " %shared = addrspacecast ptr %shared_generic to ptr addrspace(3)\n %global = addrspacecast ptr %global_generic to ptr addrspace(1)\n", + ); + writeln!( + output, + " call void @{}(ptr addrspace(3) %shared, ptr addrspace(1) %global, i32 %source_size)", + llvm(record).symbol + ) + .unwrap(); + output.push_str(" ret void\n}\n"); + writeln!( + output, + "define void @probe_{}_immediate(ptr %shared_generic, ptr %global_generic) {{", + record.id + ) + .unwrap(); + output.push_str( + " %shared = addrspacecast ptr %shared_generic to ptr addrspace(3)\n %global = addrspacecast ptr %global_generic to ptr addrspace(1)\n", + ); + writeln!( + output, + " call void @{}(ptr addrspace(3) %shared, ptr addrspace(1) %global, i32 3)", + llvm(record).symbol + ) + .unwrap(); + output.push_str(" ret void\n}\n"); + } else { + writeln!( + output, + "define void @probe_{}(ptr %shared_generic, ptr %global_generic) {{", + record.id + ) + .unwrap(); + output.push_str( + " %shared = addrspacecast ptr %shared_generic to ptr addrspace(3)\n %global = addrspacecast ptr %global_generic to ptr addrspace(1)\n", + ); + writeln!( + output, + " call void @{}(ptr addrspace(3) %shared, ptr addrspace(1) %global)", + llvm(record).symbol + ) + .unwrap(); + output.push_str(" ret void\n}\n"); + } + } else if let Some(control) = &record.cp_async_control { + let has_immediate = control.operation == CpAsyncControlOperation::WaitGroup; + let declaration_arguments = if has_immediate { "i32" } else { "" }; + writeln!( + output, + "declare void @{}({declaration_arguments})", + llvm(record).symbol + ) + .unwrap(); + output.push('\n'); + writeln!(output, "define void @probe_{}() {{", record.id).unwrap(); + if has_immediate { + writeln!(output, " call void @{}(i32 3)", llvm(record).symbol).unwrap(); + } else { + writeln!(output, " call void @{}()", llvm(record).symbol).unwrap(); + } + output.push_str(" ret void\n}\n"); + } else if let Some(control) = &record.wgmma_control { + if control.mode == WgmmaControlMode::WaitGroup { + writeln!( + output, + "declare void @{}(i64 immarg) #0\n", + llvm(record).symbol + ) + .unwrap(); + } else { + writeln!(output, "declare void @{}() #0\n", llvm(record).symbol).unwrap(); + } + writeln!(output, "define void @probe_{}() #0 {{", record.id).unwrap(); + if control.mode == WgmmaControlMode::WaitGroup { + writeln!(output, " call void @{}(i64 0) #0", llvm(record).symbol,).unwrap(); + } else { + writeln!(output, " call void @{}() #0", llvm(record).symbol,).unwrap(); + } + output.push_str(" ret void\n}\n\nattributes #0 = { convergent }\n"); + } else if record.cluster_barrier.is_some() { + writeln!(output, "declare void @{}() #0", llvm(record).symbol).unwrap(); + output.push('\n'); + writeln!(output, "define void @probe_{}() #0 {{", record.id).unwrap(); + writeln!(output, " call void @{}() #0", llvm(record).symbol).unwrap(); + output.push_str(" ret void\n}\n\nattributes #0 = { convergent }\n"); + } else if record.family == "sync" { + let arguments = if record.id == "sync_threads" { + "i32" + } else { + debug_assert!(threadfence_ptx_level(record).is_some()); + "" + }; + writeln!(output, "declare void @{}({arguments})", llvm(record).symbol).unwrap(); + output.push('\n'); + writeln!(output, "define void @probe_{}() {{", record.id).unwrap(); + let call_arguments = if record.id == "sync_threads" { + "i32 0" + } else { + "" + }; + writeln!( + output, + " call void @{}({call_arguments})", + llvm(record).symbol + ) + .unwrap(); + output.push_str(" ret void\n}\n"); + } else if let Some(warp_barrier) = &record.warp_barrier { + debug_assert_eq!(warp_barrier.adapter, WarpBarrierAdapter::DirectMemberMask); + writeln!(output, "declare void @{}(i32)", llvm(record).symbol).unwrap(); + output.push('\n'); + writeln!( + output, + "define void @probe_{}(i32 %member_mask) {{", + record.id + ) + .unwrap(); + writeln!( + output, + " call void @{}(i32 %member_mask)", + llvm(record).symbol + ) + .unwrap(); + output.push_str(" ret void\n}\n"); + writeln!(output, "define void @probe_{}_immediate() {{", record.id).unwrap(); + writeln!(output, " call void @{}(i32 -1)", llvm(record).symbol).unwrap(); + output.push_str(" ret void\n}\n"); + } else if let Some(dot) = &record.dot_product { + match dot.adapter { + DotProductAdapter::DirectThreeOperands => { + writeln!( + output, + "declare i32 @{}(i32, i32, i32)", + llvm(record).symbol + ) + .unwrap(); + output.push('\n'); + writeln!( + output, + "define i32 @probe_{}(i32 %a, i32 %b, i32 %c) {{", + record.id + ) + .unwrap(); + writeln!( + output, + " %result = call i32 @{}(i32 %a, i32 %b, i32 %c)", + llvm(record).symbol + ) + .unwrap(); + } + DotProductAdapter::InsertLowHalfFalse => { + writeln!( + output, + "declare i32 @{}(i32, i32, i1, i32)", + llvm(record).symbol + ) + .unwrap(); + output.push('\n'); + writeln!( + output, + "define i32 @probe_{}(i32 %a, i32 %b, i32 %c) {{", + record.id + ) + .unwrap(); + writeln!( + output, + " %result = call i32 @{}(i32 %a, i32 %b, i1 false, i32 %c)", + llvm(record).symbol + ) + .unwrap(); + } + } + output.push_str(" ret i32 %result\n}\n"); + } else if let Some(vote) = &record.vote { + debug_assert_eq!(vote.adapter, VoteAdapter::DirectMaskPredicate); + let result_ty = match vote.mode { + VoteMode::All | VoteMode::Any | VoteMode::Uni => "i1", + VoteMode::Ballot => "i32", + }; + writeln!( + output, + "declare {result_ty} @{}(i32, i1)", + llvm(record).symbol + ) + .unwrap(); + output.push('\n'); + writeln!( + output, + "define {result_ty} @probe_{}(i32 %member_mask, i1 %predicate) {{", + record.id + ) + .unwrap(); + writeln!( + output, + " %result = call {result_ty} @{}(i32 %member_mask, i1 %predicate)", + llvm(record).symbol + ) + .unwrap(); + writeln!(output, " ret {result_ty} %result\n}}").unwrap(); + writeln!( + output, + "define {result_ty} @probe_{}_immediate(i1 %predicate) {{", + record.id + ) + .unwrap(); + writeln!( + output, + " %result = call {result_ty} @{}(i32 -1, i1 %predicate)", + llvm(record).symbol + ) + .unwrap(); + writeln!(output, " ret {result_ty} %result\n}}").unwrap(); + } else if let Some(warp_match) = &record.warp_match { + let value_ty = format!("i{}", warp_match.value_width.bits()); + let result_ty = match warp_match.mode { + WarpMatchMode::Any => "i32".to_owned(), + WarpMatchMode::All => "{ i32, i1 }".to_owned(), + }; + writeln!( + output, + "declare {result_ty} @{}(i32, {value_ty})", + llvm(record).symbol + ) + .unwrap(); + output.push('\n'); + let forms = [ + ( + "rr", + "i32 %member_mask, ", + "i32 %member_mask", + format!("{value_ty} %value"), + ), + ("ri", "", "i32 -1", format!("{value_ty} %value")), + ( + "ir", + "i32 %member_mask", + "i32 %member_mask", + format!("{value_ty} 7"), + ), + ("ii", "", "i32 -1", format!("{value_ty} 7")), + ]; + for (suffix, first_parameter, mask, value) in forms { + let parameters = match suffix { + "rr" => format!("{first_parameter}{value_ty} %value"), + "ri" => format!("{value_ty} %value"), + "ir" => first_parameter.to_owned(), + "ii" => String::new(), + _ => unreachable!(), + }; + writeln!( + output, + "define {result_ty} @probe_{}_{suffix}({parameters}) {{", + record.id + ) + .unwrap(); + writeln!( + output, + " %result = call {result_ty} @{}({mask}, {value})", + llvm(record).symbol + ) + .unwrap(); + writeln!(output, " ret {result_ty} %result\n}}").unwrap(); + } + } else if let Some(warp_shuffle) = &record.warp_shuffle { + if warp_shuffle.value_kind == WarpShuffleValueKind::I64 { + debug_assert_eq!( + warp_shuffle.adapter, + WarpShuffleAdapter::MaskValueLaneOrDeltaSplitI64LowHighB32InsertClampReassemble + ); + let mode = match warp_shuffle.mode { + WarpShuffleMode::Idx => "idx", + WarpShuffleMode::Bfly => "bfly", + WarpShuffleMode::Down => "down", + WarpShuffleMode::Up => "up", + }; + let asm = format!( + "{{ .reg .b32 lo; .reg .b32 hi; mov.b64 {{lo, hi}}, $1; shfl.sync.{mode}.b32 lo, lo, $2, {}, $3; shfl.sync.{mode}.b32 hi, hi, $2, {}, $3; mov.b64 $0, {{lo, hi}}; }}", + warp_shuffle.clamp, warp_shuffle.clamp + ); + writeln!( + output, + "define i64 @probe_{}(i32 %member_mask, i64 %value, i32 %lane) #0 {{", + record.id + ) + .unwrap(); + writeln!( + output, + " %result = call i64 asm sideeffect {asm:?}, \"=l,l,r,r\"(i64 %value, i32 %lane, i32 %member_mask) #0" + ) + .unwrap(); + output.push_str(" ret i64 %result\n}\n\nattributes #0 = { convergent }\n"); + } else { + debug_assert_eq!( + warp_shuffle.adapter, + WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp + ); + let value_ty = match warp_shuffle.value_kind { + WarpShuffleValueKind::I32 => "i32", + WarpShuffleValueKind::F32 => "float", + WarpShuffleValueKind::I64 => unreachable!(), + }; + writeln!( + output, + "declare {value_ty} @{}(i32, {value_ty}, i32, i32)", + llvm(record).symbol + ) + .unwrap(); + output.push('\n'); + let forms = [ + ( + "rr", + format!("i32 %member_mask, {value_ty} %value, i32 %lane"), + "i32 %member_mask", + "i32 %lane", + ), + ( + "ri", + format!("{value_ty} %value, i32 %lane"), + "i32 -1", + "i32 %lane", + ), + ( + "ir", + format!("i32 %member_mask, {value_ty} %value"), + "i32 %member_mask", + "i32 1", + ), + ("ii", format!("{value_ty} %value"), "i32 -1", "i32 1"), + ]; + for (suffix, parameters, member_mask, lane) in forms { + writeln!( + output, + "define {value_ty} @probe_{}_{suffix}({parameters}) {{", + record.id + ) + .unwrap(); + writeln!( + output, + " %result = call {value_ty} @{}({member_mask}, {value_ty} %value, {lane}, i32 {})", + llvm(record).symbol, + warp_shuffle.clamp, + ) + .unwrap(); + writeln!(output, " ret {value_ty} %result\n}}").unwrap(); + } + } + } else if let Some(redux) = &record.redux { + debug_assert_eq!(redux.adapter, ReduxAdapter::MaskValueToSourceMemberMask); + let value_type = match llvm(record).results[0].as_str() { + "i32" => "i32", + "f32" => "float", + other => panic!("unsupported redux value type {other}"), + }; + writeln!( + output, + "declare {value_type} @{}({value_type}, i32)", + llvm(record).symbol + ) + .unwrap(); + output.push('\n'); + writeln!( + output, + "define {value_type} @probe_{}(i32 %member_mask, {value_type} %value) {{", + record.id + ) + .unwrap(); + writeln!( + output, + " %result = call {value_type} @{}({value_type} %value, i32 %member_mask)", + llvm(record).symbol + ) + .unwrap(); + writeln!(output, " ret {value_type} %result\n}}").unwrap(); + } else if let Some(sparse_mma) = &record.sparse_mma { + let accumulator = match sparse_mma.accumulator { + SparseMmaAccumulator::F16 => "i32", + SparseMmaAccumulator::F32 => "float", + SparseMmaAccumulator::S32 => "i32", + }; + let (c_count, a_count, b_count, d_count) = sparse_mma_fragment_counts(record); + let result = format!( + "{{ {} }}", + std::iter::repeat_n(accumulator, d_count) + .collect::>() + .join(", ") + ); + let parameters = (0..c_count) + .map(|index| format!("{accumulator} %c{index}")) + .chain((0..a_count).map(|index| format!("i32 %a{index}"))) + .chain((0..b_count).map(|index| format!("i32 %b{index}"))) + .chain(std::iter::once("i32 %metadata".to_owned())) + .collect::>() + .join(", "); + let arguments = (0..c_count) + .map(|index| format!("{accumulator} %c{index}")) + .chain((0..a_count).map(|index| format!("i32 %a{index}"))) + .chain((0..b_count).map(|index| format!("i32 %b{index}"))) + .chain(std::iter::once("i32 %metadata".to_owned())) + .collect::>() + .join(", "); + for selector in sparse_mma_selector_values(record) { + writeln!( + output, + "define {result} @probe_{}_selector_{selector}({parameters}) #0 {{", + record.id + ) + .unwrap(); + writeln!( + output, + " %result = call {result} asm sideeffect {:?}, {:?}({arguments}, i32 {selector}) #0", + sparse_mma_template(record), + sparse_mma_constraints(record), + ) + .unwrap(); + writeln!(output, " ret {result} %result\n}}").unwrap(); + } + output.push_str("\nattributes #0 = { convergent }\n"); + } else if let Some(mma) = &record.register_mma { + let (c_count, a_count, b_count, d_count) = register_mma_fragment_counts(record); + let (c_type, packed_type, result_type) = match mma.adapter { + RegisterMmaAdapter::C2U32A2U32B1U32ToD2U32 + | RegisterMmaAdapter::C2U32A4U32B2U32ToD2U32 => ("i32", "i32", "i32"), + RegisterMmaAdapter::C4F32A2U32B1U32ToD4F32 + | RegisterMmaAdapter::C4F32A4U32B2U32ToD4F32 + | RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 => { + ("float", "i32", "float") + } + RegisterMmaAdapter::C2F64A1F64B1F64ToD2F64 => ("double", "double", "double"), + RegisterMmaAdapter::C2I32A1U32B1U32ToD2I32 + | RegisterMmaAdapter::C4I32A4U32B2U32ToD4I32 + | RegisterMmaAdapter::C4I32A2U32B1U32ToD4I32 => ("i32", "i32", "i32"), + }; + let mut parameter_values = (0..c_count) + .map(|index| format!("{c_type} %c{index}")) + .chain((0..a_count).map(|index| format!("{packed_type} %a{index}"))) + .chain((0..b_count).map(|index| format!("{packed_type} %b{index}"))) + .collect::>(); + let mut argument_values = (0..c_count) + .map(|index| format!("{c_type} %c{index}")) + .chain((0..a_count).map(|index| format!("{packed_type} %a{index}"))) + .chain((0..b_count).map(|index| format!("{packed_type} %b{index}"))) + .collect::>(); + if mma.adapter == RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 { + let block_scale_parameters = [ + "i32 %scale_a", + "i16 %byte_id_a", + "i16 %thread_id_a", + "i32 %scale_b", + "i16 %byte_id_b", + "i16 %thread_id_b", + ]; + parameter_values.extend(block_scale_parameters.map(str::to_owned)); + argument_values.extend(block_scale_parameters.map(str::to_owned)); + } + let parameters = parameter_values.join(", "); + let arguments = argument_values.join(", "); + let result = format!( + "{{ {} }}", + std::iter::repeat_n(result_type, d_count) + .collect::>() + .join(", ") + ); + writeln!( + output, + "define {result} @probe_{}({parameters}) #0 {{", + record.id + ) + .unwrap(); + writeln!( + output, + " %result = call {result} asm sideeffect {:?}, {:?}({arguments}) #0", + register_mma_template(record), + register_mma_constraints(record), + ) + .unwrap(); + output.push_str(" ret "); + output.push_str(&result); + output.push_str(" %result\n}\n\nattributes #0 = { convergent }\n"); + } else if let Some(debug) = &record.debug_control { + let template = match debug.operation { + DebugControlOperation::Trap => "trap;", + DebugControlOperation::Breakpoint => "brkpt;", + DebugControlOperation::Pmevent => "pmevent 15;", + }; + writeln!(output, "define void @probe_{}() {{", record.id).unwrap(); + writeln!(output, " call void asm sideeffect {:?}, \"\"()", template).unwrap(); + output.push_str(" ret void\n}\n"); + } else if record.family == "stmatrix" { + let (multiplicity, _) = stmatrix_variant(record).expect("stmatrix variant"); + let count = multiplicity.register_count(); + let symbol = llvm(record) + .resolved_symbol + .as_ref() + .expect("stmatrix resolved symbol"); + let declaration = std::iter::once("ptr addrspace(3)".to_owned()) + .chain(std::iter::repeat_n("i32".to_owned(), count)) + .collect::>() + .join(", "); + let parameters = std::iter::once("ptr %generic".to_owned()) + .chain((0..count).map(|index| format!("i32 %r{index}"))) + .collect::>() + .join(", "); + let arguments = std::iter::once("ptr addrspace(3) %shared".to_owned()) + .chain((0..count).map(|index| format!("i32 %r{index}"))) + .collect::>() + .join(", "); + writeln!(output, "declare void @{symbol}({declaration})\n").unwrap(); + writeln!(output, "define void @probe_{}({parameters}) {{", record.id).unwrap(); + output.push_str(" %shared = addrspacecast ptr %generic to ptr addrspace(3)\n"); + writeln!(output, " call void @{symbol}({arguments})").unwrap(); + output.push_str(" ret void\n}\n"); + } else if record.prmt.is_some() { + let arity = record.rust.arguments.len(); + let declaration = std::iter::repeat_n("i32", arity) + .collect::>() + .join(", "); + let parameters = (0..arity) + .map(|index| format!("i32 %arg{index}")) + .collect::>() + .join(", "); + let arguments = (0..arity) + .map(|index| format!("i32 %arg{index}")) + .collect::>() + .join(", "); + writeln!( + output, + "declare i32 @{}({declaration})\n", + llvm(record).symbol + ) + .unwrap(); + writeln!(output, "define i32 @probe_{}({parameters}) {{", record.id).unwrap(); + writeln!( + output, + " %result = call i32 @{}({arguments})", + llvm(record).symbol + ) + .unwrap(); + output.push_str(" ret i32 %result\n}\n"); + } else if let Some(ldmatrix) = &record.ldmatrix { + let register_count = ldmatrix.variant.register_count(); + let result_ty = if register_count == 1 { + "i32".to_owned() + } else { + format!( + "{{ {} }}", + std::iter::repeat_n("i32", register_count) + .collect::>() + .join(", ") + ) + }; + let symbol = record + .llvm + .as_ref() + .expect("ldmatrix LLVM facts") + .resolved_symbol + .as_ref() + .expect("ldmatrix resolved symbol"); + writeln!(output, "declare {result_ty} @{symbol}(ptr addrspace(3))").unwrap(); + output.push('\n'); + writeln!( + output, + "define {result_ty} @probe_{}(ptr %generic) {{", + record.id + ) + .unwrap(); + output.push_str(" %shared = addrspacecast ptr %generic to ptr addrspace(3)\n"); + writeln!( + output, + " %result = call {result_ty} @{symbol}(ptr addrspace(3) %shared)" + ) + .unwrap(); + writeln!(output, " ret {result_ty} %result\n}}\n").unwrap(); + } else if record.scalar_arithmetic.is_some() { + let ty = scalar_arithmetic_llvm_type(record); + let arity = scalar_arithmetic_arity(record); + let parameters = (0..arity) + .map(|index| format!("{ty} %arg{index}")) + .collect::>() + .join(", "); + let arguments = (0..arity) + .map(|index| format!("{ty} %arg{index}")) + .collect::>() + .join(", "); + if scalar_arithmetic_llvm_mechanism(record) == BackendLoweringMechanism::TypedNvvm { + let declaration = std::iter::repeat_n(ty, arity) + .collect::>() + .join(", "); + writeln!( + output, + "declare {ty} @{}({declaration})\n", + llvm(record).symbol + ) + .unwrap(); + } + writeln!(output, "define {ty} @probe_{}({parameters}) {{", record.id).unwrap(); + match scalar_arithmetic_llvm_mechanism(record) { + BackendLoweringMechanism::TypedNvvm => { + writeln!( + output, + " %result = call {ty} @{}({arguments})", + llvm(record).symbol + ) + .unwrap(); + } + BackendLoweringMechanism::InlinePtx => { + let register = if ty == "double" { "d" } else { "f" }; + let constraints = std::iter::once(format!("={register}")) + .chain(std::iter::repeat_n(register.to_owned(), arity)) + .collect::>() + .join(","); + let operands = (0..=arity) + .map(|index| format!("${index}")) + .collect::>() + .join(", "); + writeln!( + output, + " %result = call {ty} asm {:?}, {:?}({arguments})", + format!("{} {operands};", scalar_arithmetic_ptx_mnemonic(record)), + constraints, + ) + .unwrap(); + } + } + writeln!(output, " ret {ty} %result\n}}").unwrap(); + } else if record.scalar_math.is_some() { + let ty = scalar_math_llvm_type(record); + let parameters = format!("{ty} %arg0"); + let arguments = format!("{ty} %arg0"); + if scalar_math_llvm_mechanism(record) == BackendLoweringMechanism::TypedNvvm { + writeln!(output, "declare {ty} @{}({ty})\n", llvm(record).symbol).unwrap(); + } + writeln!(output, "define {ty} @probe_{}({parameters}) {{", record.id).unwrap(); + match scalar_math_llvm_mechanism(record) { + BackendLoweringMechanism::TypedNvvm => { + writeln!( + output, + " %result = call {ty} @{}({arguments})", + llvm(record).symbol + ) + .unwrap(); + } + BackendLoweringMechanism::InlinePtx => { + let register = match ty { + "double" => "d", + "i16" => "h", + _ => "f", + }; + writeln!( + output, + " %result = call {ty} asm {:?}, {:?}({arguments})", + format!("{} $0, $1;", scalar_math_ptx_mnemonic(record)), + format!("={register},{register}"), + ) + .unwrap(); + } + } + writeln!(output, " ret {ty} %result\n}}").unwrap(); + } else if record.extended_minmax.is_some() { + let (ty, register) = match extended_minmax_contract(record).format { + ExtendedMinMaxFormat::F32 => ("float", "f"), + ExtendedMinMaxFormat::F16 | ExtendedMinMaxFormat::Bf16 => ("i16", "h"), + ExtendedMinMaxFormat::F16x2 | ExtendedMinMaxFormat::Bf16x2 => ("i32", "r"), + }; + writeln!( + output, + "define {ty} @probe_{}({ty} %a, {ty} %b) {{", + record.id + ) + .unwrap(); + writeln!( + output, + " %result = call {ty} asm {:?}, {:?}({ty} %a, {ty} %b)", + format!("{} $0, $1, $2;", extended_minmax_ptx_mnemonic(record)), + format!("={register},{register},{register}"), + ) + .unwrap(); + writeln!(output, " ret {ty} %result\n}}").unwrap(); + } else if record.scalar_conversion.is_some() { + writeln!(output, "declare i32 @{}(float)", llvm(record).symbol).unwrap(); + output.push('\n'); + writeln!(output, "define i32 @probe_{}(float %value) {{", record.id).unwrap(); + writeln!( + output, + " %result = call i32 @{}(float %value)", + llvm(record).symbol + ) + .unwrap(); + output.push_str(" ret i32 %result\n}\n"); + } else if let Some(width) = record.scalar_width() { + writeln!(output, "declare i{width} @{}()", llvm(record).symbol).unwrap(); + output.push('\n'); + writeln!(output, "define i{width} @probe_{}() {{", record.id).unwrap(); + writeln!( + output, + " %result = call i{width} @{}()", + llvm(record).symbol + ) + .unwrap(); + writeln!(output, " ret i{width} %result\n}}\n").unwrap(); + } else { + unreachable!("generated intrinsic has no probe renderer: {}", record.id); + } + output +} diff --git a/crates/cuda-intrinsics-gen/src/render/raw_abi.rs b/crates/cuda-intrinsics-gen/src/render/raw_abi.rs new file mode 100644 index 0000000000..460399d395 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/render/raw_abi.rs @@ -0,0 +1,580 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + CatalogFile, ClusterBarrierOrdering, ClusterMemoryOperation, CpAsyncControlOperation, + CpAsyncMbarrierOperation, CpAsyncSourceSize, ExecutionControlOperation, LdmatrixParticipation, + LdmatrixShape, MbarrierBasicOperation, MbarrierExtendedAdapter, RegisterMmaAdapter, + RegisterMmaOverflow, SparseMmaOverflow, Tcgen05LdShape, Tcgen05Operation, TmaAdapter, + WarpShuffleValueKind, WgmmaControlMode, +}; +use crate::render::common::{hardware_target_label, rust_header, source_label}; +use crate::render::families::{ + ClcSafetyArgNames, is_blackwell_ldmatrix, render_clc_safety_lines, sparse_mma_fragment_counts, + sparse_mma_metadata_rule, sparse_mma_selector_description, stmatrix_variant, tcgen05_is_commit, + tcgen05_is_multicast_commit, tcgen05_is_shift, tcgen05_participation_doc, +}; +use crate::render::reference::modules; +use anyhow::Result; +use std::fmt::Write as _; + +pub(super) fn render_raw_mod(catalog: &CatalogFile, hash: &str) -> String { + let mut output = rust_header(catalog, hash); + let abi_module = format!("__cuda_oxide_intrinsic_abi_v{}", catalog.intrinsic_abi); + writeln!( + output, + "#[doc(hidden)]\n#[path = \"abi_v{}.rs\"]\npub mod {abi_module};\n", + catalog.intrinsic_abi + ) + .unwrap(); + for module in modules(catalog) { + writeln!( + output, + "/// Generated `{module}` intrinsic source API.\npub mod {module} {{" + ) + .unwrap(); + for record in catalog + .intrinsics + .iter() + .filter(|record| record.rust.module == module) + { + writeln!( + output, + " pub use crate::{abi_module}::{} as {};", + record.rust.abi_id, record.rust.name + ) + .unwrap(); + } + output.push_str("}\n\n"); + } + output.push_str("#[cfg(test)]\n#[allow(clippy::type_complexity)]\nmod tests {\n"); + for record in &catalog.intrinsics { + let arguments = record.rust.arguments.join(", "); + writeln!( + output, + " #[test]\n fn public_{}_reexports_abi_{}() {{", + record.rust.name, record.rust.abi_id + ) + .unwrap(); + writeln!( + output, + " let public: {}fn({}) -> {} = super::{}::{};", + if record.rust.safe { "" } else { "unsafe " }, + arguments, + record.rust.result, + record.rust.module, + record.rust.name + ) + .unwrap(); + writeln!( + output, + " let canonical: {}fn({}) -> {} = super::{abi_module}::{};", + if record.rust.safe { "" } else { "unsafe " }, + arguments, + record.rust.result, + record.rust.abi_id + ) + .unwrap(); + output.push_str(" assert_eq!(public as usize, canonical as usize);\n }\n"); + } + output.push_str("}\n"); + output +} + +pub(super) fn render_raw_abi(catalog: &CatalogFile, hash: &str) -> Result { + let mut output = rust_header(catalog, hash); + output.push_str("//! Raw ABI functions recognized by cuda-oxide.\n\n"); + for record in &catalog.intrinsics { + writeln!(output, "/// {}", record.summary).unwrap(); + writeln!(output, "///").unwrap(); + writeln!( + output, + "/// Catalog ID: `{}`. Source: {}; expects PTX `{}`.", + record.id, + source_label(record), + record.expected_ptx + ) + .unwrap(); + if is_blackwell_ldmatrix(record) { + writeln!( + output, + "/// Available on `{}` targets. Instruction floor PTX {}; the selected target may require a newer PTX version.", + hardware_target_label(&record.target.hardware), + record.target.minimum_ptx + ) + .unwrap(); + } else { + writeln!( + output, + "/// Available on `{}` targets from PTX {}.", + hardware_target_label(&record.target.hardware), + record.target.minimum_ptx + ) + .unwrap(); + } + if let Some(operation) = record.tcgen05.as_ref().map(|tcgen05| tcgen05.operation) + && let Some(participation) = tcgen05_participation_doc(operation) + { + writeln!(output, "/// {participation}").unwrap(); + } + if record.tcgen05.is_some() { + output.push_str( + "/// All tcgen05 operations in the kernel must use the same CTA-group mode.\n", + ); + } + if let Some(reason) = &record.rust.safe_allowlist_reason { + writeln!(output, "/// Safe because {reason}").unwrap(); + } + if !record.rust.safe { + output.push_str("///\n/// # Safety\n"); + if let Some(ldmatrix) = &record.ldmatrix { + let variant = &ldmatrix.variant; + let contributing_lanes = variant.multiplicity.register_count() * 8; + let readable_bytes = if variant.shape == LdmatrixShape::M16n16 { + 32 + } else { + 16 + }; + let participation = match ldmatrix.safety.participation { + LdmatrixParticipation::AllWarpLanesSameInstruction => { + "All 32 warp lanes must execute the same instruction." + } + LdmatrixParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes => { + "All 32 warp lanes must execute the same instruction and qualifiers; no lane may have exited." + } + }; + writeln!( + output, + "/// {participation} PTX maps {contributing_lanes} lane-provided row addresses for this {:?} variant; addresses may alias, but each used address must be 16-byte aligned and have {readable_bytes} readable shared-memory bytes.", + variant.multiplicity + ) + .unwrap(); + if matches!( + ldmatrix.safety.address_contract, + crate::model::LdmatrixAddressContract::WarpLaneAddressesMappedByMultiplicitySixteenByteAlignedSixteenBytesReadableWithSm75Replication + ) && contributing_lanes < 32 + { + output.push_str("/// For portable sm_75 behavior, otherwise-unused lanes must also carry valid addresses replicated from the contributing lanes.\n"); + } + output.push_str( + "/// This weak memory operation does not replace a required barrier or fence.\n", + ); + } else if let Some((multiplicity, _)) = stmatrix_variant(record) { + let address_lanes = multiplicity.register_count() * 8; + writeln!( + output, + "/// All 32 warp lanes must execute the same instruction. The first {address_lanes} lanes must provide valid, 16-byte-aligned shared-memory row addresses." + ) + .unwrap(); + output.push_str( + "/// Register operands must contain the packed b16 fragments required by the selected layout.\n/// This weak memory operation does not replace a required barrier or fence.\n", + ); + } else if let Some(mma) = &record.sparse_mma { + let (c_count, a_count, b_count, _) = sparse_mma_fragment_counts(record); + output.push_str( + "/// All 32 warp lanes must execute the same sparse MMA instruction with the same qualifiers, and no lane may have exited.\n", + ); + writeln!( + output, + "/// `_arg0`, `_arg1`, and `_arg2` hold this lane's {c_count}-register C, {a_count}-register A, and {b_count}-register B fragments; `_arg3` holds its sparse metadata." + ) + .unwrap(); + writeln!( + output, + "/// `_arg4` must be {}.", + sparse_mma_selector_description(record) + ) + .unwrap(); + output.push_str("/// The operation is register-only and is not a memory fence.\n"); + writeln!(output, "/// {}", sparse_mma_metadata_rule(mma)).unwrap(); + writeln!( + output, + "/// See the [PTX sparse MMA fragment layouts]({}).", + record.target.ptx_isa_url + ) + .unwrap(); + match mma.overflow { + SparseMmaOverflow::NotApplicable => {} + SparseMmaOverflow::Wrapping => { + output.push_str( + "/// Signed accumulator overflow wraps because this form omits `.satfinite`.\n", + ); + } + SparseMmaOverflow::Satfinite => output.push_str( + "/// Signed accumulator overflow clamps to the finite `i32` range.\n", + ), + } + } else if let Some(mma) = &record.register_mma { + output.push_str( + "/// All 32 warp lanes must execute the same `mma.sync.aligned` instruction with the same qualifiers, and no lane may have exited.\n\ + /// `_arg0`, `_arg1`, and `_arg2` must contain this lane's C, A, and B fragments in the documented PTX layout.\n\ + /// The operation is register-only and is not a memory fence.\n", + ); + if mma.adapter == RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 + { + output.push_str( + "/// `_arg3` and `_arg6` contain this lane's packed A and B scale data. `_arg4`/`_arg5` and `_arg7`/`_arg8` are the corresponding byte/thread selectors.\n\ + /// For `scale_vec::1X`, A and B byte selectors must be in `0..=3`, the A thread selector in `0..=1`, and the B thread selector in `0..=3`; other values make the PTX operation undefined.\n", + ); + } + writeln!( + output, + "/// See the [PTX MMA fragment layouts]({}).", + record.target.ptx_isa_url + ) + .unwrap(); + if mma.overflow == RegisterMmaOverflow::Wrapping { + output.push_str("/// Signed accumulator overflow wraps because this form omits `.satfinite`.\n"); + } else if mma.overflow == RegisterMmaOverflow::Satfinite { + output.push_str( + "/// Signed accumulator overflow clamps to the finite `i32` range.\n", + ); + } + } else if record.movmatrix.is_some() { + output.push_str( + "/// All 32 warp lanes must execute the same instruction, and no lane may have exited.\n\ + /// `_arg0` must contain this lane's two packed b16 input elements.\n\ + /// The operation is register-only and is not a memory fence.\n", + ); + } else if record.redux.is_some() { + output.push_str( + "/// The executing lane must be named in `mask`. Every non-exited lane named in `mask` must execute the same `redux.sync` operation with the same qualifiers and mask.\n\ + /// The instruction waits for those lanes; violating this participation contract makes the PTX operation undefined.\n", + ); + } else if record.warp_barrier.is_some() { + output.push_str( + "/// The executing lane must be named in `mask`. Every non-exited lane named in `mask` must execute the same `bar.warp.sync` operation with the same mask.\n\ + /// On `sm_6x` and earlier, all lanes named in `mask` must execute the barrier in convergence, and no lane outside `mask` may be active when it executes.\n\ + /// The barrier orders memory accesses among participating lanes; violating the participation contract makes the PTX operation undefined.\n", + ); + } else if record.warp_shuffle.is_some() { + output.push_str( + "/// The executing lane must be named in `mask`. Every non-exited lane named in `mask` must execute the same `shfl.sync` operation with the same qualifiers and mask.\n\ + /// On `sm_6x` and earlier, all lanes named in `mask` must execute in convergence, and no lane outside `mask` may be active.\n\ + /// If the computed source lane is in range, it must be active and named in `mask`; otherwise the result is undefined. If PTX marks the computed source out of range, the calling lane's input is copied.\n", + ); + if record + .warp_shuffle + .as_ref() + .is_some_and(|shuffle| shuffle.value_kind == WarpShuffleValueKind::I64) + { + output.push_str( + "/// The 64-bit value is moved by two `b32` shuffles in one convergent block.\n", + ); + } + } else if record.warp_match.is_some() { + output.push_str( + "/// The executing lane must be named in `mask`. Every non-exited lane named in `mask` must execute the same `match.sync` operation with the same qualifiers and mask.\n\ + /// Violating this participation contract makes the PTX operation undefined.\n", + ); + } else if record.family == "elect" { + output.push_str( + "/// The executing lane must be named in `mask`. Every non-exited lane named in `mask` must execute the same `elect.sync` operation with the same mask.\n\ + /// Violating this participation contract makes the PTX operation undefined.\n", + ); + } else if record.vote.is_some() { + output.push_str( + "/// The executing lane must be named in `mask`. Every non-exited lane named in `mask` must execute the same `vote.sync` operation with the same qualifiers and mask.\n\ + /// On `sm_6x` and earlier, all lanes named in `mask` must execute in convergence, and no lane outside `mask` may be active.\n\ + /// Violating this participation contract makes the PTX operation undefined.\n", + ); + } else if record.family == "sync" { + output.push_str( + "/// Every active thread in the CTA must reach the same barrier. Calling it from divergent control flow can deadlock the CTA.\n", + ); + } else if record.family == "cluster_barrier" { + let barrier = record + .cluster_barrier + .as_ref() + .expect("cluster-barrier contract"); + output.push_str( + "/// Each non-exited cluster thread must arrive exactly once before the barrier completes, then execute the matching wait.\n", + ); + if barrier.aligned { + output.push_str( + "/// Every non-exited thread in the warp must execute this aligned instruction with identical control flow.\n", + ); + } + if barrier.ordering == ClusterBarrierOrdering::Relaxed { + output.push_str( + "/// This relaxed arrival does not publish earlier memory accesses; add the required cluster-scope fence before it.\n", + ); + } + } else if let Some(cluster) = &record.cluster_memory { + match cluster.operation { + ClusterMemoryOperation::MapSharedRank => output.push_str( + "/// `_arg0` must point into CTA shared memory, and `_arg1` must name a rank in the same live cluster.\n\ + /// The result is a cluster-shared pointer in address space 7. Dereferencing it performs a remote DSMEM access; ordinary loads and stores compile to `ld.shared::cluster` and `st.shared::cluster`. The target CTA must remain live and synchronization must make the access race-free.\n", + ), + ClusterMemoryOperation::ReadU32 => output.push_str( + "/// `_arg0` must point to an aligned readable `u32` in CTA shared memory, and `_arg1` must name a rank in the same live cluster.\n\ + /// The target CTA must publish the value with the required cluster synchronization before this weak load.\n", + ), + } + } else if let Some(control) = &record.wgmma_control { + output.push_str( + "/// All four warps in the warpgroup must execute this instruction in convergence with the same control flow.\n", + ); + match control.mode { + WgmmaControlMode::Fence => output.push_str( + "/// Issue it after register writes and before the WGMMA operations that consume those registers.\n", + ), + WgmmaControlMode::CommitGroup => output.push_str( + "/// Issue it only after the warpgroup has submitted the WGMMA operations to commit.\n", + ), + WgmmaControlMode::WaitGroup => output.push_str( + "/// `_arg0` must be a compile-time constant shared by the whole warpgroup.\n", + ), + } + } else if let Some(tma) = &record.tma { + match tma.adapter { + TmaAdapter::G2sPointersCoordinatesBarrierInjectDefaults + | TmaAdapter::G2sPointersCoordinatesBarrierMaskInjectDefaults => output + .push_str( + "/// The destination and barrier must be valid shared-memory objects, and the tensor map must be a live descriptor for this dimensionality.\n\ + /// Keep every object alive until the asynchronous copy completes. Only the designated issuing thread may start this transfer.\n", + ), + TmaAdapter::S2gPointersCoordinatesInjectDefaults => output.push_str( + "/// The source must name a live shared-memory tile, and the tensor map must be a live descriptor for this dimensionality.\n\ + /// Keep both objects alive until the committed bulk-copy group completes.\n", + ), + TmaAdapter::ReductionPointersCoordinatesInjectDefaults => output.push_str( + "/// The source must name a live shared-memory tile, and the tensor map must describe a compatible global tensor destination for this dimensionality.\n\ + /// Keep both objects alive until the committed asynchronous reduction completes.\n", + ), + TmaAdapter::DescriptorPointer + | TmaAdapter::DescriptorCoordinatesInjectDefaults + | TmaAdapter::DescriptorCoordinatesCacheHintInjectFlag => output.push_str( + "/// The tensor-map pointer must name a live, correctly encoded descriptor for this operation.\n", + ), + TmaAdapter::DescriptorAndAddressPointers + | TmaAdapter::DescriptorOrdinalAndU32 + | TmaAdapter::DescriptorOrdinalAndU64 + | TmaAdapter::DescriptorAndImmediateU32 + | TmaAdapter::DescriptorAndRuntimeU32 => output.push_str( + "/// The tensor-map pointer must name a writable, 128-byte descriptor in global memory, and every replacement value must satisfy the PTX field contract.\n", + ), + TmaAdapter::DescriptorPointerInjectBytes => output.push_str( + "/// The tensor-map pointer must name a live 128-byte descriptor covered by the matching generic-proxy release fence.\n", + ), + TmaAdapter::NoOperands | TmaAdapter::CompileTimeConstantMaxPending => {} + } + } else if let Some(operation) = ExecutionControlOperation::from_catalog_id(&record.id) { + match operation { + ExecutionControlOperation::BarrierCtaSync + | ExecutionControlOperation::BarrierCtaSyncAligned + | ExecutionControlOperation::BarrierCtaArrive + | ExecutionControlOperation::BarrierCtaArriveAligned => output.push_str( + "/// `_arg0` must identify a CTA barrier and `_arg1` must be the compatible expected thread count used by every participant.\n", + ), + ExecutionControlOperation::GridDependencyLaunchDependents + | ExecutionControlOperation::GridDependencyWait => output.push_str( + "/// The kernel launch must participate in a valid programmatic dependent-launch protocol.\n", + ), + ExecutionControlOperation::SetMaxNRegInc + | ExecutionControlOperation::SetMaxNRegDec => output.push_str( + "/// `_arg0` must be a compile-time multiple of eight in `24..=256`, and every thread in the warpgroup must execute the same operation and count.\n", + ), + } + } else if let Some(bridge) = &record.cp_async_mbarrier { + output.push_str( + "/// `_arg0` must point to a live, initialized, eight-byte-aligned mbarrier object in shared memory.\n\ + /// The issuing thread must have prior `cp.async` operations, and the object must remain valid until they complete.\n", + ); + match bridge.operation { + CpAsyncMbarrierOperation::Arrive => output.push_str( + "/// This instruction increments the pending count before scheduling the asynchronous arrival; that increment must not exceed the barrier's pending-count limit.\n", + ), + CpAsyncMbarrierOperation::ArriveNoInc => output.push_str( + "/// The barrier's initial pending count must already include the asynchronous arrival because this form does not increment it.\n", + ), + } + } else if let Some(mbarrier) = &record.mbarrier_basic { + output.push_str( + "/// `_arg0` must point to a live, eight-byte-aligned mbarrier object in shared memory.\n", + ); + match mbarrier.operation { + MbarrierBasicOperation::Init => output.push_str( + "/// Exactly one thread may initialize the object. `_arg1` must be in `1..=0xFFFFF` and include every arrival in the phase.\n\ + /// The object must be uninitialized or invalidated, with no concurrent barrier operation.\n", + ), + MbarrierBasicOperation::Arrive => output.push_str( + "/// The object must be initialized, and this arrival must be included in the current phase's expected count.\n\ + /// Use the returned token only with the same object and phase.\n", + ), + MbarrierBasicOperation::TestWait => output.push_str( + "/// The object must be initialized. `_arg1` must be a token returned for the same object and phase.\n", + ), + MbarrierBasicOperation::Inval => output.push_str( + "/// The object must be initialized and no thread or asynchronous operation may still use it.\n\ + /// Exactly one thread may invalidate the object.\n", + ), + } + } else if let Some(mbarrier) = &record.mbarrier_extended { + output.push_str( + "/// This is a convergent, side-effecting operation with a compiler memory clobber.\n", + ); + match mbarrier.adapter { + MbarrierExtendedAdapter::PointerTxCountBytesToTokenDroppingTxCount => { + output.push_str( + "/// `_arg0` must point to a live, initialized, eight-byte-aligned mbarrier in shared memory. `_arg2` must match the asynchronous transaction bytes.\n\ + /// `_arg1` is retained for compatibility and is not a PTX operand.\n", + ); + } + MbarrierExtendedAdapter::RawClusterAddressToVoid => output.push_str( + "/// `_arg0` must be a live remote shared-memory barrier address obtained through cluster address mapping.\n", + ), + MbarrierExtendedAdapter::PointerTokenToPredicate => output.push_str( + "/// `_arg0` must point to a live initialized barrier. `_arg1` must be a token for that barrier and phase.\n", + ), + MbarrierExtendedAdapter::PointerParityToPredicate => output.push_str( + "/// `_arg0` must point to a live initialized barrier. `_arg1` must be the expected phase parity.\n", + ), + MbarrierExtendedAdapter::ZeroOperandsToVoid => output.push_str( + "/// The call site must satisfy the scope and proxy-ordering contract named by the operation.\n", + ), + MbarrierExtendedAdapter::NanosecondsToVoid => output.push_str( + "/// `_arg0` is an approximate suspension duration; hardware may resume the thread earlier.\n", + ), + } + } else if let Some(copy) = &record.cp_async_copy { + let bytes = copy.copy_size.bytes(); + writeln!( + output, + "/// `_arg0` must point to {bytes} writable bytes in shared memory and be aligned to {bytes} bytes." + ) + .unwrap(); + if copy.source_size == CpAsyncSourceSize::Runtime { + writeln!( + output, + "/// `_arg2` must be at most {bytes}; `_arg1` must point to that many readable bytes in global memory and be aligned to {bytes} bytes." + ) + .unwrap(); + } else { + writeln!( + output, + "/// `_arg1` must point to {bytes} readable bytes in global memory and be aligned to {bytes} bytes." + ) + .unwrap(); + } + output.push_str( + "/// Both ranges must remain valid, the source must remain unchanged, and the destination must not be accessed until this copy completes.\n\ + /// The issuing thread must use a matching `cp.async` completion operation. Synchronize threads after completion before another thread accesses the destination.\n\ + /// User-authored completion assembly must include a compiler memory clobber.\n", + ); + } else if let Some(control) = &record.cp_async_control { + if control.operation == CpAsyncControlOperation::WaitGroup { + output.push_str( + "/// `_arg0` must be a compile-time constant. Access only destinations whose copy groups this wait completes.\n", + ); + } else if control.operation == CpAsyncControlOperation::WaitAll { + output.push_str( + "/// This waits only for copies issued by the executing thread. Synchronize threads before another thread accesses a completed destination.\n", + ); + } else { + output.push_str( + "/// This commits only copies issued by the executing thread and does not wait for completion.\n", + ); + } + } else if let Some(tcgen05) = &record.tcgen05 { + if tcgen05_is_commit(tcgen05.operation) { + if tcgen05_is_multicast_commit(tcgen05.operation) { + output.push_str( + "/// `_arg0` must point to a live initialized cluster-shared mbarrier. `_arg1` must select valid CTA ranks in its cluster.\n", + ); + } else { + output.push_str( + "/// `_arg0` must point to a live initialized mbarrier valid for this CTA-group mode.\n", + ); + } + output.push_str( + "/// The same thread that issued the tracked asynchronous tcgen05 operations must issue this commit.\n", + ); + } else if tcgen05_is_shift(tcgen05.operation) { + output.push_str( + "/// `_arg0` must name a live tensor-memory allocation, and its lane component must be a multiple of 32.\n\ + /// Completion must be tracked by a matching commit from that same thread and observed through the selected mbarrier before relying on shifted data.\n", + ); + } else if tcgen05.operation == Tcgen05Operation::Ld { + output.push_str( + "/// `_arg0` must name a live tensor-memory allocation covering the selected tile.\n\ + /// All active warp lanes must execute convergently with the same address.\n\ + /// Complete the matching tensor-memory load wait before consuming the returned registers.\n", + ); + if tcgen05 + .ld + .is_some_and(|ld| ld.shape == Tcgen05LdShape::M16x32bx2) + { + output.push_str( + "/// `_arg1` must be a compile-time constant; inline PTX encodes its low 32 bits.\n", + ); + } + } else if tcgen05.operation == Tcgen05Operation::St { + output.push_str( + "/// `_arg0` must name a live tensor-memory allocation covering the selected tile.\n\ + /// All active warp lanes must execute convergently with the same address.\n\ + /// Complete the matching tensor-memory store wait before relying on completion or reusing the affected storage.\n", + ); + if tcgen05 + .st + .is_some_and(|st| st.shape == Tcgen05LdShape::M16x32bx2) + { + output.push_str( + "/// `_arg1` must be a compile-time constant; inline PTX encodes its low 32 bits.\n", + ); + } + } else { + output.push_str( + "/// The caller must satisfy the tcgen05 address, lifetime, and participation rules.\n", + ); + } + } else if record.packed_atomic.is_some() { + output.push_str( + "/// `addr` must designate four writable bytes in global memory and be naturally aligned to four bytes.\n\ + /// Do not overlap this operation with a whole-word atomic or any non-atomic access to either 16-bit lane.\n\ + /// Racing atomics must use scopes that include each other; this relaxed GPU-scope operation is not atomic with host/system access.\n", + ); + } else if let Some(clc) = &record.clc { + render_clc_safety_lines(&mut output, clc.operation, ClcSafetyArgNames::RawAbi); + } else { + anyhow::bail!( + "{} ({}) is unsafe but has no dedicated family safety renderer", + record.id, + record.family + ); + } + } + if record.rust.arguments.len() > 7 { + output.push_str("#[allow(clippy::too_many_arguments)]\n"); + } + if record.rust.must_use { + output.push_str("#[must_use]\n"); + } + output.push_str("#[inline(never)]\n"); + let safety = if record.rust.safe { "" } else { "unsafe " }; + let arguments = record + .rust + .arguments + .iter() + .enumerate() + .map(|(index, ty)| format!("_arg{index}: {ty}")) + .collect::>() + .join(", "); + writeln!( + output, + "pub {safety}fn {}({arguments}) -> {} {{", + record.rust.abi_id, record.rust.result, + ) + .unwrap(); + writeln!( + output, + " unreachable!(\"generated CUDA intrinsic `{}` executed outside device compilation\")", + record.rust.canonical_path + ) + .unwrap(); + output.push_str("}\n\n"); + } + Ok(output) +} diff --git a/crates/cuda-intrinsics-gen/src/render/reference.rs b/crates/cuda-intrinsics-gen/src/render/reference.rs new file mode 100644 index 0000000000..98f59a0cfa --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/render/reference.rs @@ -0,0 +1,733 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, CatalogFile, ClusterBarrierOrdering, ClusterMemorySourceContract, + CpAsyncMbarrierOperation, CpAsyncMbarrierStateSpace, EvidenceArtifactKind, IntrinsicBackend, + MbarrierBasicOperation, MbarrierExtendedSourceContract, PackedAtomicFormat, + PackedConversionRounding, PackedConversionSaturation, RegisterMmaAdapter, RegisterMmaOperation, + RegisterMmaOverflow, RuntimeValidation, SparseMmaOverflow, WarpMatchAdapter, + WarpShuffleValueKind, WgmmaControlMode, +}; +use crate::render::common::{ + backend_label, evidence_stage_label, hardware_target_label, llvm, lowering_mechanism_label, + markdown_header, source_label, +}; +use crate::render::families::{ + BLACKWELL_LDMATRIX_EFFECTIVE_FLOORS, active_masks, cluster_barriers, cluster_memory, + cp_async_mbarriers, expected_ptx_head, is_blackwell_ldmatrix, mbarrier_basics, + mbarrier_extended, movmatrix, packed_alu_format_shape, packed_alu_ptx_mnemonic, packed_alus, + packed_atomics, packed_conversion_destination, packed_conversion_ptx_mnemonic, + packed_conversion_typed_llvm_name, packed_conversions, redux, register_mmas, + scalar_arithmetic_llvm_mechanism, scalar_arithmetic_ptx_mnemonic, scalar_arithmetics, + scalar_conversion_ptx_mnemonic, scalar_conversions, sparse_mma_metadata_rule, + sparse_mma_ptx_head, sparse_mma_selector_description, sparse_mmas, sync_intrinsics, + threadfence_ptx_level, vote_intrinsics, warp_barriers, warp_matches, warp_shuffles, + wgmma_controls, +}; +use std::collections::BTreeSet; +use std::fmt::Write as _; + +/// Escape a value for a GFM table cell. +/// +/// A `|` splits the row before inline parsing runs, so it splits the cell even +/// inside a code span. GFM's own rule is that a literal pipe must be written +/// `\|` wherever it appears. Three catalog entries carry a literal pipe today: +/// the `elect.sync` and `match.all.sync` PTX patterns write their two +/// destinations as ``. Their rows rendered with a seventh +/// cell, which dropped the backend-evidence column. +fn escape_table_cell(value: impl std::fmt::Display) -> String { + value.to_string().replace('|', "\\|") +} + +pub(super) fn render_reference(catalog: &CatalogFile, hash: &str) -> String { + let mut output = markdown_header(catalog, hash); + output.push_str( + "# Generated Intrinsic Reference\n\nThis table is generated from the resolved catalog. Evidence stages below distinguish backend code generation, terminal validation, and GPU runtime execution; no runtime claim is made unless an executed stage is recorded.\n\n| Rust function | CUDA operation | Source | Effects | Availability | Backend evidence |\n|:--|:--|:--|:--|:--|:--|\n", + ); + for record in &catalog.intrinsics { + let safety = if record.rust.safe { "safe" } else { "unsafe" }; + let availability = if is_blackwell_ldmatrix(record) { + format!( + "Instruction floor PTX {} on {}; effective LLVM target floors: {BLACKWELL_LDMATRIX_EFFECTIVE_FLOORS}", + record.target.minimum_ptx, + hardware_target_label(&record.target.hardware) + ) + } else { + format!( + "PTX {} on {}", + record.target.minimum_ptx, + hardware_target_label(&record.target.hardware) + ) + }; + writeln!( + output, + "| `{}` | `{}` | {} | {safety}; scope {}; {}; memory {}; convergent {} | {availability} ([PTX ISA {}, {}]({})) | {} on `{}`; expects `{}` (`{}`, SHA-256 `{}`) |", + escape_table_cell(&record.rust.public_path), + escape_table_cell(&record.dialect.op_name), + escape_table_cell(source_label(record)), + escape_table_cell(&record.semantics.execution_scope), + if record.semantics.pure { "pure" } else { "impure" }, + escape_table_cell(&record.semantics.memory), + record.semantics.convergent, + escape_table_cell(&record.target.ptx_isa_version), + escape_table_cell(&record.target.ptx_isa_section), + escape_table_cell(&record.target.ptx_isa_url), + escape_table_cell(&record.backend.status), + escape_table_cell(&record.backend.profile), + escape_table_cell(&record.expected_ptx), + escape_table_cell(&record.backend.version), + escape_table_cell(&record.backend.sha256), + ) + .unwrap(); + } + output.push_str("\n## Compiler identity and compatibility paths\n\n"); + for record in &catalog.intrinsics { + writeln!(output, "- `{}` (`{}`)", record.id, record.rust.abi_id).unwrap(); + writeln!( + output, + " - canonical compiler path: `{}`", + record.rust.canonical_path + ) + .unwrap(); + writeln!( + output, + " - public source path: `{}`", + record.rust.public_path + ) + .unwrap(); + for path in &record.rust.compatibility_paths { + writeln!(output, " - compatibility compiler path: `{path}`").unwrap(); + } + } + output.push_str("\n## Register-MMA contracts\n\n"); + for record in register_mmas(catalog) { + let mma = record.register_mma.as_ref().unwrap(); + let operation = match mma.operation { + RegisterMmaOperation::Multiply => "multiply-accumulate", + RegisterMmaOperation::AndPopc => "AND, population count, and accumulate", + RegisterMmaOperation::XorPopc => "XOR, population count, and accumulate", + }; + let overflow = match mma.overflow { + RegisterMmaOverflow::NotApplicable => "not applicable", + RegisterMmaOverflow::Wrapping => "wrapping", + RegisterMmaOverflow::Satfinite => "finite saturation", + }; + let runtime = match mma.runtime_validation { + RuntimeValidation::Unexecuted => "not executed on a GPU", + RuntimeValidation::Executed => "executed on a GPU", + }; + if mma.adapter == RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 { + writeln!( + output, + "- `{}` takes C, A, B, scale-A data/selectors, and scale-B data/selectors in PTX operand order, performs {operation}, and lowers to one convergent, register-only `{}` instruction. For `scale_vec::1X`, byte selectors must be in `0..=3`, the A thread selector in `0..=1`, and the B thread selector in `0..=3`. Every non-exited warp lane must execute the same instruction and qualifiers. Integer overflow is {overflow}; runtime validation is {runtime}.", + record.id, + expected_ptx_head(record), + ) + .unwrap(); + } else { + writeln!( + output, + "- `{}` takes fragments in C, A, B order, performs {operation}, and lowers to one convergent, register-only `{}` instruction. Every non-exited warp lane must execute the same instruction and qualifiers. Integer overflow is {overflow}; runtime validation is {runtime}.", + record.id, + expected_ptx_head(record), + ) + .unwrap(); + } + } + output.push_str("\n## Sparse-MMA contracts\n\n"); + for record in sparse_mmas(catalog) { + let mma = record.sparse_mma.as_ref().unwrap(); + let overflow = match mma.overflow { + SparseMmaOverflow::NotApplicable => "Overflow mode is not applicable.", + SparseMmaOverflow::Wrapping => "Integer overflow wraps.", + SparseMmaOverflow::Satfinite => "Integer overflow uses finite saturation.", + }; + let runtime = match mma.runtime_validation { + RuntimeValidation::Unexecuted => "not executed on a GPU", + RuntimeValidation::Executed => "executed on a GPU", + }; + let metadata = sparse_mma_metadata_rule(mma); + writeln!( + output, + "- `{}` takes fragments in C, A, B, metadata, selector order and lowers to one convergent, register-only `{}` instruction. Its LLVM source record uses A, B, C, metadata, selector order. The selector must be {}. Every non-exited warp lane must execute the same instruction and qualifiers. {metadata} {overflow} Runtime validation is {runtime}.", + record.id, + sparse_mma_ptx_head(record), + sparse_mma_selector_description(record), + ) + .unwrap(); + } + output.push_str("\n## Packed-atomic contracts\n\n"); + for record in packed_atomics(catalog) { + let packed = record.packed_atomic.as_ref().unwrap(); + let format = match packed.format { + PackedAtomicFormat::F16x2 => "f16x2", + PackedAtomicFormat::Bf16x2 => "bf16x2", + }; + writeln!( + output, + "- `{}` (`{format}`): the native PTX instruction starts at `sm_{}`; cuda-oxide admits it from {}. Omitted `.sem` / `.scope` mean relaxed GPU scope. Each 16-bit element rounds to nearest-even and `.noftz` preserves subnormal inputs and results. The elements are atomic independently, so the returned `u32` contains old per-element values that may not form one coherent pair. The pointer must address four writable, four-byte-aligned global bytes; do not mix whole-word or non-atomic overlapping access, and every racing atomic must use a mutually inclusive scope.", + record.id, + packed.native_minimum_sm, + hardware_target_label(&record.target.hardware), + ) + .unwrap(); + } + output.push_str("\n## Redux contracts\n\n"); + for record in redux(catalog) { + writeln!( + output, + "- `{}`: raw and dialect operands are `[member_mask, value]`, adapted to LLVM `(value, membermask)`. The executing lane must be named in the mask, and every non-exited named lane must execute the same instruction with the same qualifiers and mask.", + record.id, + ) + .unwrap(); + } + output.push_str("\n## Packed-ALU contracts\n\n"); + for record in packed_alus(catalog) { + let packed = record.packed_alu.as_ref().unwrap(); + let (format, _, carrier, _, _) = packed_alu_format_shape(packed.format); + let backend_floors = record + .backend_lowerings + .iter() + .map(|lowering| { + let backend = match lowering.backend { + IntrinsicBackend::LlvmNvptx => "LLVM-NVPTX", + IntrinsicBackend::LibNvvm => "libNVVM", + }; + format!( + "{backend} PTX {} on {}", + lowering.target.minimum_ptx, + hardware_target_label(&lowering.target.hardware), + ) + }) + .collect::>() + .join("; "); + writeln!( + output, + "- `{}` carries one packed `{format}` value in a `{carrier}` and lowers to one pure `{}` instruction. The native instruction starts at PTX {} / `sm_{}`; cuda-oxide admits it from {}. Backend profile floors: {backend_floors}.", + record.id, + packed_alu_ptx_mnemonic(record), + record.target.minimum_ptx, + packed.native_minimum_sm, + hardware_target_label(&record.target.hardware), + ) + .unwrap(); + } + output.push_str("\n## Packed-conversion contract\n\n"); + for record in packed_conversions(catalog) { + let conversion = record.packed_conversion.as_ref().unwrap(); + let rounding = match conversion.rounding { + PackedConversionRounding::NearestEven => "nearest-even", + PackedConversionRounding::TowardZero => "toward-zero", + }; + let saturation = match conversion.saturation { + PackedConversionSaturation::None => "without saturation", + PackedConversionSaturation::Relu => "with ReLU", + PackedConversionSaturation::Satfinite => "with finite saturation", + PackedConversionSaturation::SatfiniteRelu => "with finite saturation and ReLU", + }; + if packed_conversion_typed_llvm_name(record).is_some() { + writeln!( + output, + "- `{}` converts two `f32` inputs to packed `{}` using {rounding} rounding {saturation}. LLVM-NVPTX uses typed `{}` with `[high, low]` inputs; libNVVM uses pure `{}` inline PTX. The public first input becomes the low lane and the second becomes the high lane.", + record.id, + packed_conversion_destination(record), + llvm(record).symbol, + packed_conversion_ptx_mnemonic(record), + ) + .unwrap(); + } else { + writeln!( + output, + "- `{}` converts two `f32` inputs to packed `{}` using {rounding} rounding {saturation}. It lowers to pure `{}` inline PTX. The first input becomes the low lane and the second becomes the high lane, so PTX prints the inputs in reverse order.", + record.id, + packed_conversion_destination(record), + packed_conversion_ptx_mnemonic(record), + ) + .unwrap(); + } + } + output.push_str("\n## Scalar TF32-conversion contract\n\n"); + for record in scalar_conversions(catalog) { + writeln!( + output, + "- `{}` converts one `f32` value to raw TF32 bits with `{}`. LLVM-NVPTX uses the typed intrinsic; libNVVM uses pure inline PTX so the native target floor is preserved.", + record.id, + scalar_conversion_ptx_mnemonic(record), + ) + .unwrap(); + } + output.push_str("\n## Explicit-rounding scalar-arithmetic contracts\n\n"); + for record in scalar_arithmetics(catalog) { + let llvm_route = match scalar_arithmetic_llvm_mechanism(record) { + BackendLoweringMechanism::TypedNvvm => "the typed intrinsic", + BackendLoweringMechanism::InlinePtx => "pure inline PTX", + }; + writeln!( + output, + "- `{}` performs `{}`. LLVM-NVPTX uses {llvm_route}; libNVVM uses pure inline PTX.", + record.id, + scalar_arithmetic_ptx_mnemonic(record), + ) + .unwrap(); + } + output.push_str("\n## Warp vote contracts\n\n"); + for record in vote_intrinsics(catalog) { + writeln!( + output, + "- `{}` keeps raw and dialect operands in `[member_mask, predicate]` order. The executing lane must be named in the mask, and every non-exited named lane must execute the same `vote.sync` instruction with the same qualifiers and mask. On `sm_6x` and earlier, all named lanes must execute in convergence and no unnamed lane may be active. Both immediate and register member masks are admitted.", + record.id, + ) + .unwrap(); + } + output.push_str("\n## Active-mask contract\n\n"); + for record in active_masks(catalog) { + writeln!( + output, + "- `{}` observes the lanes active at the instruction. LLVM uses the typed intrinsic; libNVVM uses reviewed convergent, side-effecting inline PTX because that backend does not select the intrinsic.", + record.id, + ) + .unwrap(); + } + output.push_str("\n## Warp-match contracts\n\n"); + for record in warp_matches(catalog) { + let adapter = match record.warp_match.as_ref().unwrap().adapter { + WarpMatchAdapter::DirectMask => "returns LLVM's mask directly", + WarpMatchAdapter::ProjectMaskDiscardPredicate => { + "projects field 0 from LLVM's `{i32, i1}` result" + } + }; + writeln!( + output, + "- `{}` keeps operands in `[member_mask, value]` order and {adapter}. The executing lane must be named in the mask, and every non-exited named lane must execute the same `match.sync` operation with the same qualifiers and mask. All register/immediate value and mask forms are admitted.", + record.id, + ) + .unwrap(); + } + output.push_str("\n## Warp-barrier contract\n\n"); + for record in warp_barriers(catalog) { + writeln!( + output, + "- `{}` passes the 32-bit member mask directly to the typed LLVM intrinsic on both backends. The executing lane must be named in the mask, and every non-exited named lane must execute the same `bar.warp.sync` operation with the same mask. On `sm_6x` and earlier, all named lanes must execute the barrier in convergence, and no unnamed lane may be active when it executes. The barrier orders memory accesses among participating lanes. Both immediate and register masks are admitted.", + record.id, + ) + .unwrap(); + } + output.push_str("\n## Warp-shuffle contracts\n\n"); + for record in warp_shuffles(catalog) { + let shuffle = record.warp_shuffle.as_ref().unwrap(); + let lowering = match shuffle.value_kind { + WarpShuffleValueKind::I32 | WarpShuffleValueKind::F32 => { + "Register and immediate lane/mask forms are admitted." + } + WarpShuffleValueKind::I64 => { + "One convergent, side-effecting inline-PTX block splits `i64`, performs low then high `b32` shuffles with the same register lane/mask, and reassembles it." + } + }; + writeln!( + output, + "- `{}` keeps raw and dialect operands in `[member_mask, value, lane_or_delta]` order and inserts clamp `{}` during lowering. The executing lane must be named in the mask, and every non-exited named lane must execute the same `shfl.sync` operation with the same qualifiers and mask. On `sm_6x` and earlier, all named lanes must execute in convergence and no unnamed lane may be active. A computed in-range source must be active and named; if PTX marks it out of range, the calling lane's input is copied. {lowering}", + record.id, + shuffle.clamp, + ) + .unwrap(); + } + output.push_str("\n## Cluster-barrier contracts\n\n"); + for record in cluster_barriers(catalog) { + let barrier = record.cluster_barrier.as_ref().unwrap(); + let ordering = match barrier.ordering { + ClusterBarrierOrdering::Release => "release", + ClusterBarrierOrdering::Relaxed => "relaxed", + ClusterBarrierOrdering::Acquire => "acquire", + }; + let alignment = if barrier.aligned { + "Every non-exited warp thread must execute the same aligned instruction in identical control flow." + } else { + "The instruction is not aligned." + }; + writeln!( + output, + "- `{}` has {ordering} ordering. Each non-exited cluster thread must arrive exactly once before completion, then execute the matching wait. {alignment}", + record.id, + ) + .unwrap(); + } + if cluster_barriers(catalog).next().is_some() { + output.push_str( + "- `cuda_device::cluster::cluster_sync` is the generated compatibility operation: aligned arrive followed by aligned wait.\n", + ); + } + if cluster_memory(catalog).next().is_some() { + output.push_str("\n## Cluster-memory contracts\n\n"); + for record in cluster_memory(catalog) { + match record.cluster_memory.as_ref().unwrap().source_contract { + ClusterMemorySourceContract::LlvmMapaSharedClusterAs7IdentityInlinePtx => { + writeln!( + output, + "- `{}` preserves LLVM 22's address-space-7 result. Both backends use exact convergent `mapa.shared::cluster.u64` inline PTX, and ordinary Rust dereference of the mapped pointer lowers through cluster shared memory.", + record.id + ) + .unwrap(); + } + ClusterMemorySourceContract::PtxNativeMapaThenWeakClusterLoad => { + writeln!( + output, + "- `{}` has no one-to-one LLVM intrinsic. Both backends use one convergent `mapa.shared::cluster.u64` plus `ld.shared::cluster.u32` block with a compiler memory clobber.", + record.id + ) + .unwrap(); + } + } + } + } + if movmatrix(catalog).next().is_some() { + output.push_str("\n## movmatrix contracts\n\n"); + for record in movmatrix(catalog) { + writeln!( + output, + "- `{}` transposes one packed b16 fragment in registers. All warp lanes must execute the same instruction, and no lane may have exited.", + record.id, + ) + .unwrap(); + } + } + if mbarrier_extended(catalog).next().is_some() { + output.push_str("\n## Extended mbarrier contracts\n\n"); + for record in mbarrier_extended(catalog) { + let contract = record.mbarrier_extended.as_ref().unwrap(); + let source = match contract.source_contract { + MbarrierExtendedSourceContract::LlvmImported => "LLVM 22 TableGen", + MbarrierExtendedSourceContract::PtxNativeRawClusterAddress => { + "PTX-native raw-address contract" + } + }; + writeln!( + output, + "- `{}` keeps its established CUDA-device signature and exact convergent inline-PTX memory clobber. Its source identity is {source}.", + record.id, + ) + .unwrap(); + } + } + if wgmma_controls(catalog).next().is_some() { + output.push_str("\n## WGMMA-control contracts\n\n"); + for record in wgmma_controls(catalog) { + let detail = match record.wgmma_control.as_ref().unwrap().mode { + WgmmaControlMode::Fence => { + "It orders register writes before later WGMMA operations." + } + WgmmaControlMode::CommitGroup => { + "It commits the warpgroup's prior uncommitted WGMMA operations." + } + WgmmaControlMode::WaitGroup => { + "Its `u64` argument must be a compile-time constant." + } + }; + writeln!( + output, + "- `{}` must be executed by all four warps in the warpgroup with the same control flow. {detail}", + record.id, + ) + .unwrap(); + } + } + if sync_intrinsics(catalog).any(|record| threadfence_ptx_level(record).is_some()) { + output.push_str("\n## Synchronization contracts\n\n"); + } else { + output.push_str("\n## CTA synchronization contracts\n\n"); + } + for record in sync_intrinsics(catalog) { + if record.id == "sync_threads" { + writeln!( + output, + "- `{}` inserts the fixed barrier ID `0`. Every active CTA thread must reach the same barrier; divergent use can deadlock the CTA.", + record.id, + ) + .unwrap(); + } else { + let level = threadfence_ptx_level(record) + .expect("resolver admitted an unknown generated synchronization record"); + writeln!( + output, + "- `{}` emits `membar.{level}` through the reviewed typed NVVM route on both backends.", + record.id, + ) + .unwrap(); + } + } + output.push_str("\n## cp.async mbarrier contracts\n\n"); + for record in cp_async_mbarriers(catalog) { + let bridge = record.cp_async_mbarrier.as_ref().unwrap(); + let address = match bridge.state_space { + CpAsyncMbarrierStateSpace::Generic => "generic", + CpAsyncMbarrierStateSpace::Shared => "explicit shared", + }; + let counting = match bridge.operation { + CpAsyncMbarrierOperation::Arrive => { + "increments the pending count before the later asynchronous decrement" + } + CpAsyncMbarrierOperation::ArriveNoInc => { + "does not increment the pending count, so initialization must already include the later asynchronous decrement" + } + }; + writeln!( + output, + "- `{}` uses {address} addressing and {counting}. It associates only the executing thread's prior `cp.async` operations with a live, eight-byte-aligned shared-memory barrier. LLVM-NVPTX uses the typed intrinsic; libNVVM uses reviewed convergent, side-effecting inline PTX with a memory clobber.", + record.id, + ) + .unwrap(); + } + output.push_str("\n## Basic mbarrier contracts\n\n"); + for record in mbarrier_basics(catalog) { + let mbarrier = record.mbarrier_basic.as_ref().unwrap(); + let contract = match mbarrier.operation { + MbarrierBasicOperation::Init => { + "initializes one eight-byte-aligned shared-memory object exactly once with an expected count in `1..=0xFFFFF`" + } + MbarrierBasicOperation::Arrive => { + "records one expected arrival and returns a token for the same barrier phase" + } + MbarrierBasicOperation::TestWait => { + "tests one token from the same barrier phase with convergent, side-effecting inline PTX" + } + MbarrierBasicOperation::Inval => { + "invalidates an initialized barrier only after all users are finished" + } + }; + let mechanism = if mbarrier.operation == MbarrierBasicOperation::TestWait { + "inline PTX on both backends" + } else { + "the typed NVVM intrinsic with LLVM-NVPTX and inline PTX with libNVVM" + }; + writeln!( + output, + "- `{}` {contract}. It lowers through {mechanism}; generic MIR pointers are normalized to shared address space during lowering.", + record.id, + ) + .unwrap(); + } + output.push_str("\n## Imported LLVM properties and result facts\n\n"); + for record in &catalog.intrinsics { + let Some(llvm) = &record.llvm else { + writeln!( + output, + "- `{}`: PTX-native source; no LLVM record, symbol, properties, or selection facts.", + record.id + ) + .unwrap(); + continue; + }; + let properties = record + .llvm + .as_ref() + .unwrap() + .properties + .iter() + .map(|property| format!("`{property}`")) + .collect::>() + .join(", "); + let range = record + .llvm + .as_ref() + .unwrap() + .result_facts + .range + .as_ref() + .map(|range| format!("[{}, {})", range.lower, range.upper_exclusive)) + .unwrap_or_else(|| "none".to_owned()); + let selection_records = record + .selections + .iter() + .map(|selection| format!("`{}`", selection.source_record)) + .collect::>() + .into_iter() + .collect::>() + .join(", "); + let selection_predicates = record + .selections + .iter() + .flat_map(|selection| selection.predicates.iter()) + .map(|predicate| format!("`{predicate}`")) + .collect::>() + .into_iter() + .collect::>() + .join(", "); + writeln!( + output, + "- `{}`: properties {}; result `NoUndef` {}; half-open result range `{}`; selection records {}; selection predicates {}", + record.id, + properties, + llvm.result_facts.no_undef, + range, + selection_records, + if selection_predicates.is_empty() { + "none" + } else { + &selection_predicates + } + ) + .unwrap(); + } + output.push_str("\n## Backend-specific lowering evidence\n\n"); + for record in &catalog.intrinsics { + if record.backend_lowerings.is_empty() { + continue; + } + let runtime = record + .ldmatrix + .as_ref() + .map(|record| format!("{:?}", record.safety.runtime_validation).to_lowercase()) + .or_else(|| { + record + .packed_atomic + .as_ref() + .map(|record| format!("{:?}", record.runtime_validation).to_lowercase()) + }) + .or_else(|| { + record + .mbarrier_basic + .as_ref() + .map(|record| format!("{:?}", record.runtime_validation).to_lowercase()) + }) + .or_else(|| { + record + .movmatrix + .as_ref() + .map(|record| format!("{:?}", record.runtime_validation).to_lowercase()) + }) + .or_else(|| { + record + .mbarrier_extended + .as_ref() + .map(|record| format!("{:?}", record.runtime_validation).to_lowercase()) + }) + .or_else(|| { + record + .cp_async_mbarrier + .as_ref() + .map(|record| format!("{:?}", record.runtime_validation).to_lowercase()) + }) + .or_else(|| { + record + .register_mma + .as_ref() + .map(|record| format!("{:?}", record.runtime_validation).to_lowercase()) + }) + .or_else(|| { + record + .sparse_mma + .as_ref() + .map(|record| format!("{:?}", record.runtime_validation).to_lowercase()) + }) + .or_else(|| { + record + .debug_control + .as_ref() + .map(|record| format!("{:?}", record.runtime_validation).to_lowercase()) + }) + .or_else(|| { + record + .cluster_memory + .as_ref() + .map(|record| format!("{:?}", record.runtime_validation).to_lowercase()) + }) + .or_else(|| { + record + .tcgen05 + .as_ref() + .map(|record| format!("{:?}", record.runtime_validation).to_lowercase()) + }) + .or_else(|| (record.family == "stmatrix").then(|| "unexecuted".to_owned())) + .unwrap_or_else(|| "not recorded".to_owned()); + writeln!(output, "- `{}`: runtime `{runtime}`", record.id).unwrap(); + for lowering in &record.backend_lowerings { + writeln!( + output, + " - `{}` uses `{}` from profile `{}` at PTX {} / {}: status `{}` (`{}`, SHA-256 `{}`)", + backend_label(lowering.backend), + lowering_mechanism_label(lowering.mechanism), + lowering.evidence_profile, + lowering.target.minimum_ptx, + hardware_target_label(&lowering.target.hardware), + lowering.status, + lowering.version, + lowering.sha256, + ) + .unwrap(); + for stage in &lowering.stages { + let tool = match ( + stage.tool_path.as_deref(), + stage.tool_version.as_deref(), + stage.tool_sha256.as_deref(), + ) { + (Some(path), Some(version), Some(sha256)) => { + format!(" Tool `{path}` reports `{version}` (SHA-256 `{sha256}`).") + } + _ => String::new(), + }; + let artifact = match stage.artifact_kind { + Some(EvidenceArtifactKind::Cubin) => " Artifact `cubin`.", + None => "", + }; + writeln!( + output, + " - {} on `{}`: `{}` — {}{}{}", + evidence_stage_label(stage.stage), + stage.targets.join(", "), + stage.outcome, + stage.detail, + tool, + artifact, + ) + .unwrap(); + } + } + } + output +} + +pub(super) fn render_compiler_path_patterns( + output: &mut String, + catalog: &CatalogFile, + indent: &str, +) { + let paths: Vec<_> = catalog + .intrinsics + .iter() + .flat_map(|record| { + std::iter::once(record.rust.canonical_path.as_str()) + .chain(record.rust.compatibility_paths.iter().map(String::as_str)) + }) + .collect(); + render_string_patterns(output, &paths, indent); +} + +pub(super) fn render_string_patterns(output: &mut String, values: &[&str], indent: &str) { + for (index, value) in values.iter().enumerate() { + if index == 0 { + writeln!(output, "{indent}{value:?}").unwrap(); + } else { + writeln!(output, "{indent}| {value:?}").unwrap(); + } + } +} + +pub(super) fn render_inline_patterns(output: &mut String, values: &[&str]) { + for (index, value) in values.iter().enumerate() { + if index != 0 { + output.push_str(" | "); + } + write!(output, "{value:?}").unwrap(); + } +} + +#[allow(dead_code)] +pub(super) fn modules(catalog: &CatalogFile) -> BTreeSet<&str> { + catalog + .intrinsics + .iter() + .map(|record| record.rust.module.as_str()) + .collect() +} diff --git a/crates/cuda-intrinsics-gen/src/render/tests/mma.rs b/crates/cuda-intrinsics-gen/src/render/tests/mma.rs new file mode 100644 index 0000000000..5b1d4c4570 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/render/tests/mma.rs @@ -0,0 +1,1160 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use super::*; + +use crate::model::{ + RegisterMmaAdapter, RegisterMmaCompatibilitySource, RegisterMmaKind, RuntimeValidation, + SparseMmaAccumulator, SparseMmaAdapter, SparseMmaShape, +}; +use crate::render::collector_targets::generated_intrinsic_variant; +use crate::render::common::intrinsic_marker; +use crate::render::families::{ + BLACKWELL_LDMATRIX_EFFECTIVE_FLOORS, SPARSE_MMA_ORDERED_METADATA_RULE, + SPARSE_MMA_STANDARD_METADATA_RULE, ldmatrix, register_mma_attr_variants, + register_mma_constraints, register_mma_template, register_mmas, sparse_mma_carriers, + sparse_mma_constraints, sparse_mma_fragment_counts, sparse_mma_selector_values, + sparse_mma_template, sparse_mmas, +}; +use crate::util::read_json; +use std::path::Path; + +#[test] +fn stmatrix_rendering_preserves_all_four_public_and_backend_contracts() { + let catalog = catalog_with_stmatrix(); + validate_renderable(&catalog).unwrap(); + assert_eq!(stmatrices(&catalog).count(), 4); + + let compatibility = render_compat_stmatrix(&catalog, "test-hash"); + for signature in [ + "pub unsafe fn stmatrix_m8n8_x2(smem_ptr: *mut u8, r0: u32, r1: u32)", + "pub unsafe fn stmatrix_m8n8_x2_trans(smem_ptr: *mut u8, r0: u32, r1: u32)", + "pub unsafe fn stmatrix_m8n8_x4(smem_ptr: *mut u8, r0: u32, r1: u32, r2: u32, r3: u32)", + "pub unsafe fn stmatrix_m8n8_x4_trans(smem_ptr: *mut u8, r0: u32, r1: u32, r2: u32, r3: u32)", + ] { + assert!(compatibility.contains(signature)); + } + assert!(compatibility.contains("All warp lanes must participate")); + + let dialect_mod = render_dialect_mod(&catalog, "test-hash"); + assert!(dialect_mod.contains("mod stmatrix;")); + assert!(dialect_mod.contains("stmatrix::register(ctx)")); + let dialect = render_dialect_stmatrix(&catalog, "test-hash"); + for (op_type, op_name, operands) in [ + ("StmatrixM8n8X2Op", "nvvm.stmatrix_m8n8_x2", 3), + ("StmatrixM8n8X2TransOp", "nvvm.stmatrix_m8n8_x2_trans", 3), + ("StmatrixM8n8X4Op", "nvvm.stmatrix_m8n8_x4", 5), + ("StmatrixM8n8X4TransOp", "nvvm.stmatrix_m8n8_x4_trans", 5), + ] { + assert!(dialect.contains(&format!("pub struct {op_type};"))); + assert!(dialect.contains(&format!("name = \"{op_name}\""))); + assert!(dialect.contains(&format!("NOpdsInterface<{operands}>"))); + assert!(dialect.contains(&format!("{op_type}::register(ctx)"))); + } + + let importer = render_importer(&catalog, "test-hash"); + for (path, op_type, marker) in [ + ( + "cuda_device::tcgen05::stmatrix_m8n8_x2", + "StmatrixM8n8X2Op", + "v1:i0301", + ), + ( + "cuda_device::tcgen05::stmatrix_m8n8_x2_trans", + "StmatrixM8n8X2TransOp", + "v1:i0302", + ), + ( + "cuda_device::tcgen05::stmatrix_m8n8_x4", + "StmatrixM8n8X4Op", + "v1:i0303", + ), + ( + "cuda_device::tcgen05::stmatrix_m8n8_x4_trans", + "StmatrixM8n8X4TransOp", + "v1:i0304", + ), + ] { + assert!(importer.contains(path)); + assert!(importer.contains(&format!("{op_type}::get_concrete_op_info()"))); + assert!(importer.contains(&format!( + "set_generated_intrinsic_marker(ctx, store, \"{marker}\")" + ))); + } + + let lowering = render_lowering(&catalog, "test-hash"); + assert_eq!( + lowering + .matches("impl MirToLlvmConversion for StmatrixM8n8") + .count(), + 4 + ); + assert!(lowering.contains("IntrinsicBackend::LlvmNvptx")); + assert!(lowering.contains("IntrinsicBackend::LibNvvm")); + assert!(lowering.contains("cast_to_shared_addrspace")); + assert!(lowering.contains("llvm_nvvm_stmatrix_sync_aligned_m8n8_x2_b16_p3")); + assert!(lowering.contains("llvm_nvvm_stmatrix_sync_aligned_m8n8_x4_trans_b16_p3")); + assert!(lowering.contains("stmatrix.sync.aligned.m8n8.x{register_count}{trans}.shared.b16")); + assert!(lowering.contains("std::iter::once(\"~{memory}\")")); + + let x2 = stmatrices(&catalog) + .find(|record| record.id == "stmatrix_m8n8_x2_b16") + .unwrap(); + let probe = render_probe(&catalog, x2, "test-hash"); + assert!(probe.contains( + "declare void @llvm.nvvm.stmatrix.sync.aligned.m8n8.x2.b16.p3(ptr addrspace(3), i32, i32)" + )); + assert!(probe.contains("%shared = addrspacecast ptr %generic to ptr addrspace(3)")); + assert!(probe.contains( + "call void @llvm.nvvm.stmatrix.sync.aligned.m8n8.x2.b16.p3(ptr addrspace(3) %shared, i32 %r0, i32 %r1)" + )); + + let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); + assert!(outputs.contains_key(&PathBuf::from( + "crates/cuda-device/src/generated/stmatrix.rs" + ))); + assert!(outputs.contains_key(&PathBuf::from( + "crates/dialect-nvvm/src/ops/generated/stmatrix.rs" + ))); +} + +#[test] +fn ldmatrix_family_lowering_uses_one_attribute_dispatch_impl() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + + let dialect = render_dialect_ldmatrix(&catalog, "test-hash"); + let rendered = render_lowering(&catalog, "test-hash"); + assert_eq!( + rendered + .matches("impl MirToLlvmConversion for LdmatrixOp") + .count(), + 1 + ); + assert!(rendered.contains("LdmatrixMultiplicityAttr::X4")); + assert!(rendered.contains("LdmatrixMultiplicityAttr::X2")); + assert!(rendered.contains("LdmatrixMultiplicityAttr::X1")); + assert!(rendered.contains("LdmatrixLayoutAttr::Transposed")); + assert!(rendered.contains("llvm_nvvm_ldmatrix_sync_aligned_m8n8_x2_trans_b16_p3")); + for (op_type, op_name, count, instruction_head, intrinsic) in [ + ( + "LdmatrixX1Op", + "nvvm.ldmatrix_x1", + 1, + "ldmatrix.sync.aligned.m8n8.x1.shared.b16", + "llvm_nvvm_ldmatrix_sync_aligned_m8n8_x1_b16_p3", + ), + ( + "LdmatrixX1TransOp", + "nvvm.ldmatrix_x1_trans", + 1, + "ldmatrix.sync.aligned.m8n8.x1.trans.shared.b16", + "llvm_nvvm_ldmatrix_sync_aligned_m8n8_x1_trans_b16_p3", + ), + ( + "LdmatrixX2Op", + "nvvm.ldmatrix_x2", + 2, + "ldmatrix.sync.aligned.m8n8.x2.shared.b16", + "llvm_nvvm_ldmatrix_sync_aligned_m8n8_x2_b16_p3", + ), + ( + "LdmatrixX2TransOp", + "nvvm.ldmatrix_x2_trans", + 2, + "ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16", + "llvm_nvvm_ldmatrix_sync_aligned_m8n8_x2_trans_b16_p3", + ), + ( + "LdmatrixX4Op", + "nvvm.ldmatrix_x4", + 4, + "ldmatrix.sync.aligned.m8n8.x4.shared.b16", + "llvm_nvvm_ldmatrix_sync_aligned_m8n8_x4_b16_p3", + ), + ( + "LdmatrixX4TransOp", + "nvvm.ldmatrix_x4_trans", + 4, + "ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16", + "llvm_nvvm_ldmatrix_sync_aligned_m8n8_x4_trans_b16_p3", + ), + ] { + assert!(dialect.contains(&format!("pub struct {op_type};"))); + assert!(dialect.contains(&format!("name = {op_name:?}"))); + assert!(dialect.contains(&format!("NResultsInterface<{count}>"))); + assert!(dialect.contains(&format!("{op_type}::register(ctx);"))); + assert_eq!( + rendered + .matches(&format!("impl MirToLlvmConversion for {op_type}")) + .count(), + 1 + ); + assert!(rendered.contains(&format!( + "self.get_operation(), {count}, {instruction_head:?}, {intrinsic:?}" + ))); + } + + let importer = render_importer(&catalog, "test-hash"); + assert!(!importer.contains("LdmatrixX1Op")); + assert!(!importer.contains("LdmatrixX4TransOp")); + let targets = render_targets(&catalog, "test-hash"); + for op_name in [ + "nvvm.ldmatrix_x1", + "nvvm.ldmatrix_x1_trans", + "nvvm.ldmatrix_x2", + "nvvm.ldmatrix_x2_trans", + "nvvm.ldmatrix_x4", + "nvvm.ldmatrix_x4_trans", + ] { + assert!(targets.contains(&format!("{op_name:?} => Some("))); + } + assert_eq!(ldmatrix(&catalog).count(), 18); + assert!(dialect.contains("LdmatrixShapeAttr::M16n16")); + assert!(dialect.contains("LdmatrixShapeAttr::M8n16")); + assert!(dialect.contains("LdmatrixElementAttr::B8x16B4x16P64")); + assert!( + rendered.contains( + "llvm__nvvm_dldmatrix_dsync_daligned_dm16n16_dx1_dtrans_db8x16_db4x16_up64_dp3" + ) + ); + assert!(rendered.contains("ldmatrix.sync.aligned.m16n16.x2.trans.shared.b8x16.b6x16_p32")); + + let raw = render_raw_abi(&catalog, "test-hash").unwrap(); + assert!(raw.contains( + "Instruction floor PTX 8.6; the selected target may require a newer PTX version." + )); + assert!(raw.contains("no lane may have exited")); + assert!(raw.contains("have 32 readable shared-memory bytes")); + assert!(raw.contains("have 16 readable shared-memory bytes")); + + let reference = render_reference(&catalog, "test-hash"); + assert!(reference.contains(BLACKWELL_LDMATRIX_EFFECTIVE_FLOORS)); +} + +#[test] +fn ldmatrix_x1_probe_keeps_its_pointer_operand() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + let record = catalog + .intrinsics + .iter() + .find(|record| record.id == "ldmatrix_m8n8_x1_b16") + .unwrap(); + + let rendered = render_probe(&catalog, record, "test-hash"); + assert!( + rendered.contains( + "declare i32 @llvm.nvvm.ldmatrix.sync.aligned.m8n8.x1.b16.p3(ptr addrspace(3))" + ) + ); + assert!(rendered.contains("define i32 @probe_ldmatrix_m8n8_x1_b16(ptr %generic)")); + assert!(!rendered.contains("@llvm.nvvm.ldmatrix.sync.aligned.m8n8.x1.b16.p3()")); +} + +#[test] +fn movmatrix_rendering_owns_dialect_import_and_lowering() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let mut catalog = crate::resolve::resolve(&repo_root).unwrap(); + let mut record = catalog + .intrinsics + .iter() + .find(|record| record.id == "packed_atomic_add_f16x2") + .unwrap() + .clone(); + record.id = "movmatrix_trans_b16".into(); + record.operation_key = "movmatrix.m8n8.trans.b16".into(); + record.family = "movmatrix".into(); + record.rust.abi_id = "i0305".into(); + record.rust.module = "matrix".into(); + record.rust.name = "movmatrix_trans_b16".into(); + record.rust.arguments = vec!["u32".into()]; + record.rust.result = "u32".into(); + record.rust.safe = false; + record.rust.must_use = true; + record.rust.canonical_path = "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0305".into(); + record.rust.public_path = "cuda_intrinsics::matrix::movmatrix_trans_b16".into(); + record.rust.compatibility_paths = vec!["cuda_device::wmma::movmatrix_trans_b16".into()]; + record.dialect.op_type = "MovmatrixTransB16Op".into(); + record.dialect.op_name = "nvvm.movmatrix_trans_b16".into(); + record.dialect.operands = vec!["i32".into()]; + record.dialect.results = vec!["i32".into()]; + record.semantics.pure = false; + record.semantics.memory = "inaccessible_read_write".into(); + record.semantics.convergent = true; + record.semantics.execution_scope = "warp".into(); + record.packed_atomic = None; + record.movmatrix = Some(crate::model::Movmatrix { + participation: + crate::model::MovmatrixParticipation::AllWarpLanesSameInstructionNoExitedLanes, + adapter: crate::model::MovmatrixAdapter::PackedB16x2U32ToPackedB16x2U32, + runtime_validation: RuntimeValidation::Unexecuted, + }); + record.lowering = "generated_movmatrix_inline_ptx".into(); + record.expected_ptx = crate::ptx::InstructionPattern::new( + "movmatrix", + &["sync", "aligned", "m8n8", "trans", "b16"], + vec![ + crate::ptx::OperandPattern::Register, + crate::ptx::OperandPattern::Register, + ], + ); + record.summary = "Transposes one packed b16 matrix fragment across a warp.".into(); + catalog.intrinsics.push(record); + + validate_renderable(&catalog).unwrap(); + let dialect_mod = render_dialect_mod(&catalog, "test-hash"); + assert!(dialect_mod.contains("mod movmatrix;")); + assert!(dialect_mod.contains("movmatrix::register(ctx);")); + + let dialect = render_dialect_movmatrix(&catalog, "test-hash"); + assert!(dialect.contains("pub struct MovmatrixTransB16Op;")); + assert!(dialect.contains("name = \"nvvm.movmatrix_trans_b16\"")); + assert!(dialect.contains("MovmatrixTransB16Op::register(ctx);")); + + let compatibility = render_compat_movmatrix(&catalog, "test-hash"); + assert!(compatibility.contains("pub unsafe fn movmatrix_trans_b16(value: u32) -> u32")); + assert!(compatibility.contains("All 32 warp lanes must execute the same call")); + + let importer = render_importer(&catalog, "test-hash"); + assert!(importer.contains("cuda_device::wmma::movmatrix_trans_b16")); + assert!(importer.contains("MovmatrixTransB16Op::build(ctx, arg0)")); + + let lowering = render_lowering(&catalog, "test-hash"); + assert!(lowering.contains("impl MirToLlvmConversion for MovmatrixTransB16Op")); + assert!(lowering.contains("movmatrix.sync.aligned.m8n8.trans.b16 $0, $1;")); + assert!(lowering.contains("\"=r,r\"")); + assert!(lowering.contains("inline_asm_convergent")); + + let probe_record = movmatrix(&catalog).next().unwrap(); + let probe = render_probe(&catalog, probe_record, "test-hash"); + assert!(probe.contains("call i32 asm \"movmatrix.sync.aligned.m8n8.trans.b16 $0, $1;\"")); + assert!(probe.contains("attributes #0 = { convergent }")); + + let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); + assert!(outputs.contains_key(&PathBuf::from( + "crates/cuda-device/src/generated/movmatrix.rs" + ))); +} + +#[test] +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(), 1016); + let records: Vec<_> = register_mmas(&catalog).collect(); + assert_eq!(records.len(), 154); + let generated_records = records + .iter() + .copied() + .filter(|record| { + record.register_mma.as_ref().unwrap().compatibility_source + == RegisterMmaCompatibilitySource::GeneratedStub + }) + .collect::>(); + let existing_records = records + .iter() + .copied() + .filter(|record| { + record.register_mma.as_ref().unwrap().compatibility_source + == RegisterMmaCompatibilitySource::ExistingStub + }) + .collect::>(); + assert_eq!(generated_records.len(), 149); + assert_eq!(existing_records.len(), 5); + + let raw = render_raw_abi(&catalog, "test-hash").unwrap(); + for record in &records { + assert!(raw.contains(&format!("pub unsafe fn {}(", record.rust.abi_id))); + } + assert!(raw.contains("no lane may have exited")); + assert!(raw.contains("Signed accumulator overflow wraps")); + assert!(raw.contains("Signed accumulator overflow clamps")); + + let compatibility = render_compat_register_mma(&catalog, "test-hash"); + assert_eq!(compatibility.matches("pub unsafe fn ").count(), 149); + for record in generated_records { + let argument_names: &[&str] = if record.register_mma.as_ref().unwrap().adapter + == RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 + { + &[ + "c", + "a", + "b", + "scale_a", + "byte_id_a", + "thread_id_a", + "scale_b", + "byte_id_b", + "thread_id_b", + ] + } else { + &["c", "a", "b"] + }; + let arguments = argument_names + .iter() + .zip(&record.rust.arguments) + .map(|(name, ty)| format!("{name}: {ty}")) + .collect::>() + .join(", "); + assert!(compatibility.contains(&format!( + "pub unsafe fn {}({arguments}) -> {} {{", + record.rust.name, record.rust.result + ))); + } + for record in existing_records { + assert!(!compatibility.contains(&format!("pub unsafe fn {}(", record.rust.name))); + } + + let dialect = render_dialect_register_mma(&catalog, "test-hash"); + assert_eq!(dialect.matches("pub struct RegisterMmaOp").count(), 1); + assert!(dialect.contains("let recipe: (&[MmaCarrier], &[MmaCarrier])")); + assert!(dialect.contains("RegisterMmaShapeAttr::M8n8k4")); + assert!(dialect.contains("RegisterMmaShapeAttr::M8n8k16")); + assert!(dialect.contains("RegisterMmaShapeAttr::M8n8k32")); + assert!(dialect.contains("RegisterMmaShapeAttr::M8n8k128")); + assert!(dialect.contains("RegisterMmaShapeAttr::M16n8k4")); + assert!(dialect.contains("RegisterMmaShapeAttr::M16n8k32")); + assert!(dialect.contains("RegisterMmaShapeAttr::M16n8k64")); + assert!(dialect.contains("RegisterMmaShapeAttr::M16n8k128")); + assert!(dialect.contains("RegisterMmaShapeAttr::M16n8k256")); + assert!(dialect.contains("RegisterMmaOperationAttr::Multiply")); + assert!(dialect.contains("RegisterMmaOperationAttr::AndPopc")); + assert!(dialect.contains("RegisterMmaOperationAttr::XorPopc")); + assert!(dialect.contains("pub enum RegisterMmaKindAttr { Standard, F8f6f4, Mxf8f6f4 }")); + assert!(dialect.contains("kind_or_inferred")); + assert!(dialect.contains("RegisterMmaAccumulatorAttr::F16")); + assert!(dialect.contains("operation_or_multiply")); + assert!(dialect.contains("MmaCarrier::I32 | MmaCarrier::U16 | MmaCarrier::U32")); + assert!(dialect.contains("if matches!(carrier, MmaCarrier::U16)")); + assert!(dialect.contains("RegisterMmaElementAttr::B1")); + assert!(dialect.contains("RegisterMmaElementAttr::E2m1")); + assert!(dialect.contains("RegisterMmaElementAttr::E5m2")); + assert!(dialect.contains("RegisterMmaElementAttr::S4")); + assert!(dialect.contains("RegisterMmaElementAttr::U4")); + assert!(dialect.contains("RegisterMmaElementAttr::U8")); + assert!(dialect.contains("RegisterMmaOverflowAttr::Wrapping")); + assert!(dialect.contains("RegisterMmaOverflowAttr::Satfinite")); + assert!(dialect.contains( + "&[MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32]" + )); + for (op_type, op_name) in [ + ("MmaM16N8K16F32Bf16Op", "nvvm.mma_m16n8k16_f32_bf16"), + ("MmaM16N8K16F32F16Op", "nvvm.mma_m16n8k16_f32_f16"), + ("MmaM16N8K8F32Tf32Op", "nvvm.mma_m16n8k8_f32_tf32"), + ("MmaM16N8K32S32S8Op", "nvvm.mma_m16n8k32_s32_s8"), + ("MmaM8N8K4F64Op", "nvvm.mma_m8n8k4_f64"), + ] { + assert!(dialect.contains(&format!("pub struct {op_type};"))); + assert!(dialect.contains(&format!("name = {op_name:?}"))); + assert!(dialect.contains(&format!("{op_type}::register(ctx);"))); + } + + let importer = render_importer(&catalog, "test-hash"); + assert!(importer.contains("enum GeneratedMmaImportAdapter")); + assert!(importer.contains("C2U32A4U32B2U32ToD2U32")); + assert!(importer.contains("C2U32A2U32B1U32ToD2U32")); + assert!(importer.contains("C4F32A2U32B1U32ToD4F32")); + assert!(importer.contains("C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32")); + assert!(importer.contains("ctx, b_value, b_ty, b_count, block_ptr, last_after_b, loc.clone()")); + assert!(importer.contains("C4I32A2U32B1U32ToD4I32")); + assert!(importer.contains("C2I32A1U32B1U32ToD2I32")); + assert!(importer.contains("(i32_ty, 2, u32_ty, 1, false, u32_ty, 1, false, i32_ty, 2)")); + assert!(importer.contains("(i32_ty, 4, u32_ty, 2, true, u32_ty, 1, false, i32_ty, 4)")); + assert!(importer.contains("(u32_ty, 2, u32_ty, 4, true, u32_ty, 2, true, u32_ty, 2)")); + assert!(importer.contains("import_generated_mma_operands")); + assert!(importer.contains("bundle_generated_mma_results")); + assert!(importer.contains("set_attr_nvvm_register_mma_operation")); + assert!(importer.contains("set_attr_nvvm_register_mma_kind")); + for record in &records { + assert!(importer.contains(&record.rust.canonical_path)); + assert!(importer.contains(&record.rust.compatibility_paths[0])); + assert!(importer.contains(&format!( + "set_generated_intrinsic_marker(ctx, mma, {:?})", + intrinsic_marker(&catalog, record) + ))); + } + + let lowering = render_lowering(&catalog, "test-hash"); + assert_eq!( + lowering + .matches("impl MirToLlvmConversion for RegisterMmaOp") + .count(), + 1 + ); + assert!(lowering.contains("convert_generated_register_mma")); + assert!(lowering.contains("operation_or_multiply")); + assert!(lowering.contains("kind_or_inferred")); + assert!(lowering.contains("mma.sync.aligned.m16n8k16.row.col.f16.e4m3.e4m3.f16")); + assert!(lowering.contains("mma.sync.aligned.m16n8k4.row.col.f32.tf32.tf32.f32")); + assert!(lowering.contains("mma.sync.aligned.m16n8k8.row.col.f32.bf16.bf16.f32")); + assert!(lowering.contains("mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16")); + assert!(lowering.contains("mma.sync.aligned.m16n8k32.row.col.f32.e5m2.e5m2.f32")); + assert!(lowering.contains("=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r")); + assert!(lowering.contains("=d,=d,d,d,d,d")); + assert!(lowering.contains("=r,=r,=r,=r,r,r,r,r,r,r,r")); + assert!(lowering.contains("=r,=r,r,r,r,r")); + assert!(lowering.contains("=r,=r,r,r,r,r,r,r,r,r")); + assert!(lowering.contains("mma.sync.aligned.m8n8k16.row.col.s32.s8.u8.s32")); + assert!(lowering.contains("mma.sync.aligned.m8n8k16.row.col.satfinite.s32.u8.s8.s32")); + assert!(lowering.contains("mma.sync.aligned.m8n8k32.row.col.s32.s4.u4.s32")); + assert!(lowering.contains("mma.sync.aligned.m8n8k32.row.col.satfinite.s32.u4.s4.s32")); + assert!(lowering.contains("mma.sync.aligned.m16n8k16.row.col.s32.s8.u8.s32")); + assert!(lowering.contains("mma.sync.aligned.m16n8k32.row.col.satfinite.s32.u8.s8.s32")); + assert!(lowering.contains("mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32")); + assert!(lowering.contains("mma.sync.aligned.m16n8k32.row.col.s32.s4.u4.s32")); + assert!(lowering.contains("mma.sync.aligned.m16n8k32.row.col.satfinite.s32.u4.s4.s32")); + assert!(lowering.contains("mma.sync.aligned.m16n8k64.row.col.s32.s4.u4.s32")); + assert!(lowering.contains("mma.sync.aligned.m16n8k64.row.col.satfinite.s32.u4.s4.s32")); + assert!(lowering.contains("mma.sync.aligned.m8n8k128.row.col.s32.b1.b1.s32.xor.popc")); + assert!(lowering.contains("mma.sync.aligned.m8n8k128.row.col.s32.b1.b1.s32.and.popc")); + assert!(lowering.contains("mma.sync.aligned.m16n8k128.row.col.s32.b1.b1.s32.xor.popc")); + assert!(lowering.contains("mma.sync.aligned.m16n8k128.row.col.s32.b1.b1.s32.and.popc")); + assert!(lowering.contains("mma.sync.aligned.m16n8k256.row.col.s32.b1.b1.s32.xor.popc")); + assert!(lowering.contains("mma.sync.aligned.m16n8k256.row.col.s32.b1.b1.s32.and.popc")); + for op_type in [ + "MmaM16N8K16F32Bf16Op", + "MmaM16N8K16F32F16Op", + "MmaM16N8K8F32Tf32Op", + "MmaM16N8K32S32S8Op", + "MmaM8N8K4F64Op", + ] { + assert_eq!( + lowering + .matches(&format!("impl MirToLlvmConversion for {op_type}")) + .count(), + 1 + ); + } + assert!(lowering.contains( + r#"(GeneratedMmaResultType::I32, 4, 7, "mma.sync.aligned.m16n8k32.row.col.s32.s4.u4.s32 {$0, $1, $2, $3}, {$8, $9}, {$10}, {$4, $5, $6, $7};", "=r,=r,=r,=r,r,r,r,r,r,r,r")"# + )); + assert!(lowering.contains( + r#"(GeneratedMmaResultType::I32, 4, 10, "mma.sync.aligned.m16n8k64.row.col.satfinite.s32.u4.s4.s32 {$0, $1, $2, $3}, {$8, $9, $10, $11}, {$12, $13}, {$4, $5, $6, $7};", "=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,r")"# + )); + assert!(lowering.contains( + r#"(GeneratedMmaResultType::I32, 2, 4, "mma.sync.aligned.m8n8k128.row.col.s32.b1.b1.s32.xor.popc {$0, $1}, {$4}, {$5}, {$2, $3};", "=r,=r,r,r,r,r")"# + )); + assert!(lowering.contains( + r#"(GeneratedMmaResultType::I32, 4, 7, "mma.sync.aligned.m16n8k128.row.col.s32.b1.b1.s32.and.popc {$0, $1, $2, $3}, {$8, $9}, {$10}, {$4, $5, $6, $7};", "=r,=r,=r,=r,r,r,r,r,r,r,r")"# + )); + assert!(lowering.contains( + r#"(GeneratedMmaResultType::I32, 4, 10, "mma.sync.aligned.m16n8k256.row.col.s32.b1.b1.s32.xor.popc {$0, $1, $2, $3}, {$8, $9, $10, $11}, {$12, $13}, {$4, $5, $6, $7};", "=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,r")"# + )); + assert!(lowering.contains( + r#"(GeneratedMmaResultType::F32, 4, 10, "mma.sync.aligned.m16n8k32.row.col.kind::f8f6f4.f32.e2m1.e2m1.f32 {$0, $1, $2, $3}, {$8, $9, $10, $11}, {$12, $13}, {$4, $5, $6, $7};", "=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r")"# + )); + assert!(lowering.contains( + r#"(GeneratedMmaResultType::F32, 4, 10, "mma.sync.aligned.m16n8k32.row.col.kind::f8f6f4.f32.e5m2.e5m2.f32 {$0, $1, $2, $3}, {$8, $9, $10, $11}, {$12, $13}, {$4, $5, $6, $7};", "=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r")"# + )); + assert!(lowering.contains( + r#"(GeneratedMmaResultType::I32, 2, 8, "mma.sync.aligned.m16n8k32.row.col.kind::f8f6f4.f16.e2m1.e2m1.f16 {$0, $1}, {$4, $5, $6, $7}, {$8, $9}, {$2, $3};", "=r,=r,r,r,r,r,r,r,r,r")"# + )); + assert!(lowering.contains( + r#"(GeneratedMmaResultType::I32, 2, 8, "mma.sync.aligned.m16n8k32.row.col.kind::f8f6f4.f16.e5m2.e5m2.f16 {$0, $1}, {$4, $5, $6, $7}, {$8, $9}, {$2, $3};", "=r,=r,r,r,r,r,r,r,r,r")"# + )); + assert!(lowering.contains( + r#"(GeneratedMmaResultType::F32, 4, 16, "mma.sync.aligned.m16n8k32.row.col.kind::mxf8f6f4.block_scale.f32.e2m1.e2m1.f32.ue8m0 {$0, $1, $2, $3}, {$8, $9, $10, $11}, {$12, $13}, {$4, $5, $6, $7}, $14, {$15, $16}, $17, {$18, $19};", "=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r,r,h,h,r,h,h")"# + )); + + let targets = render_targets(&catalog, "test-hash"); + assert!(targets.contains("GeneratedIntrinsicVariant::RegisterMma")); + assert!(targets.contains("GeneratedRegisterMmaShape::M8n8k4")); + assert!(targets.contains("GeneratedRegisterMmaShape::M16n8k4")); + assert!(targets.contains( + "pub enum GeneratedRegisterMmaShape { M8n8k4, M8n8k16, M8n8k32, M8n8k128, M16n8k4, M16n8k8, M16n8k16, M16n8k32, M16n8k64, M16n8k128, M16n8k256 }" + )); + assert!(targets.contains( + "GeneratedRegisterMmaShape::M8n8k16 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M8n8k16)" + )); + assert!(targets.contains( + "GeneratedRegisterMmaShape::M8n8k32 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M8n8k32)" + )); + assert!(targets.contains( + "GeneratedRegisterMmaShape::M8n8k128 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M8n8k128)" + )); + assert!(targets.contains("GeneratedRegisterMmaShape::M16n8k32")); + assert!(targets.contains( + "GeneratedRegisterMmaShape::M16n8k64 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M16n8k64)" + )); + assert!(targets.contains( + "GeneratedRegisterMmaShape::M16n8k128 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M16n8k128)" + )); + assert!(targets.contains( + "GeneratedRegisterMmaShape::M16n8k256 => op.get_attr_nvvm_register_mma_shape(ctx).as_deref() == Some(&RegisterMmaShapeAttr::M16n8k256)" + )); + assert!(targets.contains("GeneratedRegisterMmaOperation::Multiply")); + assert!(targets.contains("GeneratedRegisterMmaOperation::AndPopc")); + assert!(targets.contains("GeneratedRegisterMmaOperation::XorPopc")); + assert!(targets.contains("pub enum GeneratedRegisterMmaKind { Standard, F8f6f4, Mxf8f6f4 }")); + assert!(targets.contains("kind: GeneratedRegisterMmaKind::Standard")); + assert!(targets.contains("kind: GeneratedRegisterMmaKind::F8f6f4")); + assert!(targets.contains("kind: GeneratedRegisterMmaKind::Mxf8f6f4")); + assert!(targets.contains("kind_or_inferred")); + assert!(targets.contains("RegisterMmaKindAttr::Standard")); + assert!(targets.contains("RegisterMmaKindAttr::F8f6f4")); + assert!(targets.contains("RegisterMmaKindAttr::Mxf8f6f4")); + assert!(targets.contains("GeneratedRegisterMmaAccumulator::F16")); + assert!(targets.contains("operation: GeneratedRegisterMmaOperation::AndPopc")); + assert!(targets.contains("operation: GeneratedRegisterMmaOperation::XorPopc")); + assert!(targets.contains("operation: mma_operation")); + assert!(targets.contains("operation_or_multiply")); + assert!(targets.contains("GeneratedRegisterMmaElement::B1")); + assert!(targets.contains("GeneratedRegisterMmaElement::E2m1")); + assert!(targets.contains("GeneratedRegisterMmaElement::E5m2")); + assert!(targets.contains("GeneratedRegisterMmaElement::S4")); + assert!(targets.contains("GeneratedRegisterMmaElement::U4")); + assert!(targets.contains("GeneratedRegisterMmaElement::U8")); + assert!(targets.contains( + "a_element: GeneratedRegisterMmaElement::S4, b_element: GeneratedRegisterMmaElement::U4" + )); + assert!(targets.contains( + "a_element: GeneratedRegisterMmaElement::U4, b_element: GeneratedRegisterMmaElement::S4" + )); + assert!(targets.contains( + "a_element: GeneratedRegisterMmaElement::S8, b_element: GeneratedRegisterMmaElement::U8" + )); + assert!(targets.contains( + "a_element: GeneratedRegisterMmaElement::U8, b_element: GeneratedRegisterMmaElement::S8" + )); + assert!(targets.contains("overflow: GeneratedRegisterMmaOverflow::Satfinite")); + assert!(targets.contains("get_attr_nvvm_register_mma_overflow")); + assert!(targets.contains("minimum_ptx: GeneratedPtxVersion::from_encoded(87)")); + assert!(targets.contains("minimum_ptx: GeneratedPtxVersion::from_encoded(84)")); + assert!(targets.contains("GeneratedHardwareAlternative::ExactArchitecture(120)")); + assert!(targets.contains("GeneratedHardwareAlternative::FamilyTarget(121)")); + + let first_dense = records + .iter() + .find(|record| record.id == "mma_m16n8k32_f32_e2m1_e2m1") + .unwrap(); + let last_dense = records + .iter() + .find(|record| record.id == "mma_m16n8k32_f32_e5m2_e5m2") + .unwrap(); + assert_eq!(first_dense.rust.abi_id, "i0454"); + assert_eq!(last_dense.rust.abi_id, "i0478"); + assert_eq!( + register_mma_constraints(first_dense), + "=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r" + ); + let first_dense_f16 = records + .iter() + .find(|record| record.id == "mma_m16n8k32_f16_e2m1_e2m1") + .unwrap(); + let last_dense_f16 = records + .iter() + .find(|record| record.id == "mma_m16n8k32_f16_e5m2_e5m2") + .unwrap(); + assert_eq!(first_dense_f16.rust.abi_id, "i0479"); + assert_eq!(last_dense_f16.rust.abi_id, "i0503"); + assert_eq!( + register_mma_constraints(first_dense_f16), + "=r,=r,r,r,r,r,r,r,r,r" + ); + assert!(render_probe(&catalog, first_dense_f16, "test-hash").contains("define { i32, i32 }")); + + let first_mxf8f6f4 = records + .iter() + .find(|record| record.id == "mma_m16n8k32_mxf8f6f4_f32_e2m1_e2m1") + .unwrap(); + let last_mxf8f6f4 = records + .iter() + .find(|record| record.id == "mma_m16n8k32_mxf8f6f4_f32_e5m2_e5m2") + .unwrap(); + assert_eq!(first_mxf8f6f4.rust.abi_id, "i0858"); + assert_eq!(last_mxf8f6f4.rust.abi_id, "i0882"); + assert_eq!( + register_mma_constraints(first_mxf8f6f4), + "=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r,r,h,h,r,h,h" + ); + assert!( + render_probe(&catalog, first_mxf8f6f4, "test-hash").contains("kind::mxf8f6f4.block_scale") + ); + let first_mxf8f6f4_compatibility = format!("pub unsafe fn {}(", first_mxf8f6f4.rust.name); + let first_mxf8f6f4_compatibility = compatibility + .find(&first_mxf8f6f4_compatibility) + .expect("generated MXF8F6F4 compatibility wrapper"); + assert!( + compatibility[..first_mxf8f6f4_compatibility] + .ends_with("#[allow(clippy::too_many_arguments)]\n#[must_use]\n#[inline(never)]\n") + ); + + for record in &records { + let probe = render_probe(&catalog, record, "test-hash"); + assert!(probe.contains("asm sideeffect")); + assert!(probe.contains("attributes #0 = { convergent }")); + assert!(probe.contains(®ister_mma_template(record))); + assert!(probe.contains(®ister_mma_constraints(record))); + } + + let reference = render_reference(&catalog, "test-hash"); + assert!(reference.contains("## Register-MMA contracts")); + assert!(reference.contains("performs XOR, population count, and accumulate")); + assert!(reference.contains("performs AND, population count, and accumulate")); + assert!(reference.contains("runtime validation is not executed on a GPU")); + for record in &records { + assert!(reference.contains(&format!("- `{}`: runtime `unexecuted`", record.id))); + } + + let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); + assert!(outputs.contains_key(&PathBuf::from( + "crates/cuda-device/src/generated/register_mma.rs" + ))); + assert!(outputs.contains_key(&PathBuf::from( + "crates/dialect-nvvm/src/ops/generated/register_mma.rs" + ))); +} + +#[test] +fn register_mma_kind_is_explicit_and_legacy_ir_stays_compatible() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let mut catalog: CatalogFile = read_json(&repo_root.join("intrinsics/catalog.json")).unwrap(); + let old = catalog + .intrinsics + .iter() + .find(|record| record.id == "mma_m16n8k32_f16_e4m3_e4m3") + .unwrap() + .clone(); + assert_eq!(old.register_mma.as_ref().unwrap().kind, None); + assert_eq!( + register_mma_attr_variants(&old).2, + "RegisterMmaKindAttr::F8f6f4" + ); + assert!(generated_intrinsic_variant(&old).contains("kind: GeneratedRegisterMmaKind::F8f6f4")); + + let mut standard = old.clone(); + standard.id = "mma_m16n8k32_fp8_f16_e4m3_e4m3".into(); + standard.operation_key = "matrix.mma.m16n8k32.row.col.standard_fp8.f16.e4m3.e4m3.f16".into(); + standard.rust.abi_id = "i9999".into(); + standard.rust.name = standard.id.clone(); + standard.rust.canonical_path = format!("cuda_intrinsics::matrix::{}", standard.id); + standard.rust.public_path = standard.rust.canonical_path.clone(); + standard.rust.compatibility_paths = vec![format!("cuda_device::wmma::{}", standard.id)]; + standard.llvm.as_mut().unwrap().symbol = + "llvm.nvvm.mma.m16n8k32.row.col.f16.e4m3.e4m3.f16".into(); + standard.register_mma.as_mut().unwrap().kind = Some(RegisterMmaKind::Standard); + standard + .expected_ptx + .modifiers + .retain(|modifier| modifier != "kind::f8f6f4"); + assert_eq!( + register_mma_attr_variants(&standard).2, + "RegisterMmaKindAttr::Standard" + ); + assert!( + generated_intrinsic_variant(&standard).contains("kind: GeneratedRegisterMmaKind::Standard") + ); + catalog.intrinsics.push(standard); + + let dialect = render_dialect_register_mma(&catalog, "test-hash"); + assert!(dialect.contains("pub fn kind_or_inferred")); + assert!( + dialect.contains("let old_f8f6f4 = self.get_attr_nvvm_register_mma_shape(ctx).as_deref()") + ); + assert!(dialect.contains("== Some(&RegisterMmaShapeAttr::M16n8k32)")); + assert!( + dialect + .contains("&& low_format(self.get_attr_nvvm_register_mma_a_element(ctx).as_deref())") + ); + assert!( + dialect + .contains("&& low_format(self.get_attr_nvvm_register_mma_b_element(ctx).as_deref())") + ); + assert!(dialect.contains("RegisterMmaKindAttr::F8f6f4")); + assert!(dialect.contains("RegisterMmaKindAttr::Standard")); + + let importer = render_importer(&catalog, "test-hash"); + assert!(importer.contains("set_attr_nvvm_register_mma_kind(ctx, RegisterMmaKindAttr::F8f6f4)")); + assert!( + importer.contains("set_attr_nvvm_register_mma_kind(ctx, RegisterMmaKindAttr::Standard)") + ); + + let lowering = render_lowering(&catalog, "test-hash"); + assert!(lowering.contains("let kind = self.kind_or_inferred(ctx)")); + assert!(lowering.contains("kind::f8f6f4.f16.e4m3.e4m3.f16")); + assert!(lowering.contains("row.col.f16.e4m3.e4m3.f16")); + + let targets = render_targets(&catalog, "test-hash"); + assert!(targets.contains("kind: GeneratedRegisterMmaKind::F8f6f4")); + assert!(targets.contains("kind: GeneratedRegisterMmaKind::Standard")); + assert!(targets.contains( + "GeneratedRegisterMmaKind::Standard => op.kind_or_inferred(ctx) == RegisterMmaKindAttr::Standard" + )); + assert!(targets.contains( + "GeneratedRegisterMmaKind::F8f6f4 => op.kind_or_inferred(ctx) == RegisterMmaKindAttr::F8f6f4" + )); + assert!(targets.contains("kind_matches && accumulator_matches")); +} + +#[test] +fn sparse_mma_rendering_enforces_selector_and_keeps_family_distinct() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + validate_renderable(&catalog).unwrap(); + let records = sparse_mmas(&catalog).collect::>(); + let k32 = records + .iter() + .filter(|record| record.sparse_mma.as_ref().unwrap().shape == SparseMmaShape::M16n8k32) + .count(); + let k64 = records + .iter() + .filter(|record| record.sparse_mma.as_ref().unwrap().shape == SparseMmaShape::M16n8k64) + .count(); + let k128 = records + .iter() + .filter(|record| record.sparse_mma.as_ref().unwrap().shape == SparseMmaShape::M16n8k128) + .count(); + assert_eq!((records.len(), k32, k64, k128), (114, 16, 82, 16)); + let standard_k64 = records + .iter() + .copied() + .find(|record| record.id == "mma_sp_m16n8k64_s32_s8") + .unwrap(); + assert_eq!( + standard_k64.sparse_mma.as_ref().unwrap().adapter, + SparseMmaAdapter::C4I32A4U32B4U32MetadataU32SelectorU32ToD4I32 + ); + assert_eq!( + standard_k64.sparse_mma.as_ref().unwrap().llvm_adapter, + crate::model::SparseMmaLlvmAdapter::A4I32B4I32C4I32MetadataI32SelectorI32ToD4I32 + ); + assert_eq!( + sparse_mma_template(standard_k64), + "mma.sp.sync.aligned.m16n8k64.row.col.s32.s8.s8.s32 {$0, $1, $2, $3}, {$8, $9, $10, $11}, {$12, $13, $14, $15}, {$4, $5, $6, $7}, $16, $17;" + ); + assert_eq!( + sparse_mma_constraints(standard_k64), + "=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,r,r,r,r,n" + ); + let standard_int4_k64 = records + .iter() + .copied() + .find(|record| record.id == "mma_sp_m16n8k64_s32_s4") + .unwrap(); + assert_eq!( + sparse_mma_template(standard_int4_k64), + "mma.sp.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 {$0, $1, $2, $3}, {$8, $9}, {$10, $11}, {$4, $5, $6, $7}, $12, $13;" + ); + assert_eq!( + sparse_mma_constraints(standard_int4_k64), + "=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,n" + ); + let ordered_int4_k64 = records + .iter() + .copied() + .find(|record| record.id == "mma_sp_ordered_metadata_m16n8k64_s32_s4") + .unwrap(); + assert_eq!( + ordered_int4_k64.sparse_mma.as_ref().unwrap().adapter, + SparseMmaAdapter::C4I32A2U32B2U32MetadataU32SelectorU32ToD4I32 + ); + assert_eq!( + ordered_int4_k64.sparse_mma.as_ref().unwrap().llvm_adapter, + crate::model::SparseMmaLlvmAdapter::A2I32B2I32C4I32MetadataI32SelectorI32ToD4I32 + ); + assert_eq!( + sparse_mma_template(ordered_int4_k64), + "mma.sp::ordered_metadata.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 {$0, $1, $2, $3}, {$8, $9}, {$10, $11}, {$4, $5, $6, $7}, $12, $13;" + ); + assert_eq!( + sparse_mma_constraints(ordered_int4_k64), + "=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,n" + ); + let ordered_int4_k128 = records + .iter() + .copied() + .find(|record| record.id == "mma_sp_ordered_metadata_m16n8k128_s32_s4") + .unwrap(); + assert_eq!( + ordered_int4_k128.sparse_mma.as_ref().unwrap().adapter, + SparseMmaAdapter::C4I32A4U32B4U32MetadataU32SelectorU32ToD4I32 + ); + assert_eq!( + ordered_int4_k128.sparse_mma.as_ref().unwrap().llvm_adapter, + crate::model::SparseMmaLlvmAdapter::A4I32B4I32C4I32MetadataI32SelectorI32ToD4I32 + ); + assert_eq!( + sparse_mma_template(ordered_int4_k128), + "mma.sp::ordered_metadata.sync.aligned.m16n8k128.row.col.s32.s4.s4.s32 {$0, $1, $2, $3}, {$8, $9, $10, $11}, {$12, $13, $14, $15}, {$4, $5, $6, $7}, $16, $17;" + ); + assert_eq!( + sparse_mma_constraints(ordered_int4_k128), + "=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,r,r,r,r,n" + ); + let standard_int4_k128 = records + .iter() + .copied() + .find(|record| record.id == "mma_sp_m16n8k128_s32_s4") + .unwrap(); + assert_eq!( + sparse_mma_template(standard_int4_k128), + "mma.sp.sync.aligned.m16n8k128.row.col.s32.s4.s4.s32 {$0, $1, $2, $3}, {$8, $9, $10, $11}, {$12, $13, $14, $15}, {$4, $5, $6, $7}, $16, $17;" + ); + assert_eq!( + sparse_mma_constraints(standard_int4_k128), + "=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,r,r,r,r,n" + ); + let ordered_f8f6f4 = records + .iter() + .copied() + .find(|record| { + record.id == "mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m1_e2m1_f32" + }) + .unwrap(); + assert_eq!( + ordered_f8f6f4.sparse_mma.as_ref().unwrap().adapter, + SparseMmaAdapter::C4F32A4U32B4U32MetadataU32SelectorU32ToD4F32 + ); + assert_eq!( + sparse_mma_template(ordered_f8f6f4), + "mma.sp::ordered_metadata.sync.aligned.m16n8k64.row.col.kind::f8f6f4.f32.e2m1.e2m1.f32 {$0, $1, $2, $3}, {$8, $9, $10, $11}, {$12, $13, $14, $15}, {$4, $5, $6, $7}, $16, $17;" + ); + assert_eq!( + sparse_mma_constraints(ordered_f8f6f4), + "=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r,r,r,r,n" + ); + assert_eq!( + sparse_mma_carriers(ordered_f8f6f4), + ( + "&[MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32]".into(), + "&[MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32, MmaCarrier::F32]".into(), + ) + ); + let ordered_f8f6f4_f16 = records + .iter() + .copied() + .find(|record| { + record.id == "mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e2m1_f16" + }) + .unwrap(); + assert_eq!( + ordered_f8f6f4_f16.sparse_mma.as_ref().unwrap().adapter, + SparseMmaAdapter::C2U32A4U32B4U32MetadataU32SelectorU32ToD2U32 + ); + assert_eq!(sparse_mma_fragment_counts(ordered_f8f6f4_f16), (2, 4, 4, 2)); + assert_eq!( + sparse_mma_template(ordered_f8f6f4_f16), + "mma.sp::ordered_metadata.sync.aligned.m16n8k64.row.col.kind::f8f6f4.f16.e2m1.e2m1.f16 {$0, $1}, {$4, $5, $6, $7}, {$8, $9, $10, $11}, {$2, $3}, $12, $13;" + ); + assert_eq!( + sparse_mma_constraints(ordered_f8f6f4_f16), + "=r,=r,r,r,r,r,r,r,r,r,r,r,r,n" + ); + assert_eq!( + sparse_mma_carriers(ordered_f8f6f4_f16), + ( + "&[MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32, MmaCarrier::U32]".into(), + "&[MmaCarrier::U32, MmaCarrier::U32]".into(), + ) + ); + assert_eq!(sparse_mma_selector_values(ordered_f8f6f4_f16), &[0]); + + let raw = render_raw_abi(&catalog, "test-hash").unwrap(); + assert!(raw.contains("must be the compile-time constant `0` or `1`")); + assert!(raw.contains("must be the compile-time constant `0`")); + for record in &records { + let (c_count, a_count, b_count, d_count) = sparse_mma_fragment_counts(record); + let scalar = match record.sparse_mma.as_ref().unwrap().accumulator { + SparseMmaAccumulator::F16 => "u32", + SparseMmaAccumulator::F32 => "f32", + SparseMmaAccumulator::S32 => "i32", + }; + assert!(raw.contains(&format!( + "pub unsafe fn {}(_arg0: [{scalar}; {c_count}], _arg1: [u32; {a_count}], _arg2: [u32; {b_count}], _arg3: u32, _arg4: u32) -> [{scalar}; {d_count}]", + record.rust.abi_id, + ))); + } + + let compatibility = render_compat_sparse_mma(&catalog, "test-hash"); + assert_eq!(compatibility.matches("pub unsafe fn ").count(), 114); + assert!( + compatibility + .contains("c: [i32; 4], a: [u32; 2], b: [u32; 2], metadata: u32, selector: u32") + ); + assert!( + compatibility + .contains("c: [i32; 4], a: [u32; 4], b: [u32; 4], metadata: u32, selector: u32") + ); + assert!( + compatibility + .contains("c: [f32; 4], a: [u32; 4], b: [u32; 4], metadata: u32, selector: u32") + ); + assert!( + compatibility + .contains("c: [u32; 2], a: [u32; 4], b: [u32; 4], metadata: u32, selector: u32") + ); + + let dialect = render_dialect_sparse_mma(&catalog, "test-hash"); + assert_eq!(dialect.matches("pub struct SparseMmaOp").count(), 1); + assert!(dialect.contains("MmaCarrier::F32")); + assert!(dialect.contains("MmaCarrier::I32")); + assert!(dialect.contains("MmaCarrier::U32")); + assert!(dialect.contains("operands.len() != expected_operands")); + assert!(dialect.contains("SparseMmaShapeAttr { M16n8k32, M16n8k64, M16n8k128 }")); + assert!(dialect.contains("SparseMmaAccumulatorAttr { F16, F32, S32 }")); + assert!(dialect.contains("SparseMmaSelectorAttr { ImmediateZeroOrOne, ImmediateZero }")); + assert!(dialect.contains("SparseMmaElementAttr::E2m1")); + assert!(dialect.contains("SparseMmaElementAttr::E2m3")); + assert!(dialect.contains("SparseMmaElementAttr::E3m2")); + assert!(dialect.contains("SparseMmaElementAttr::E4m3")); + assert!(dialect.contains("SparseMmaElementAttr::E5m2")); + assert!(dialect.contains("SparseMmaElementAttr::S4")); + assert!(dialect.contains("SparseMmaElementAttr::U4")); + assert!(dialect.contains("SparseMmaElementAttr::S8")); + assert!(dialect.contains("SparseMmaElementAttr::U8")); + assert!(dialect.contains("SparseMmaOverflowAttr::Wrapping")); + assert!(dialect.contains("SparseMmaOverflowAttr::Satfinite")); + assert!(dialect.contains("SparseMmaOverflowAttr::NotApplicable")); + assert!(dialect.contains("SparseMmaMetadataAttr { Standard, Ordered }")); + + let importer = render_importer(&catalog, "test-hash"); + assert!(importer.contains("sparse MMA selector must be the compile-time constant 0 or 1")); + assert!(importer.contains("sparse MMA selector must be the compile-time constant 0")); + assert!(importer.contains("GeneratedMmaImportAdapter::C4I32A2U32B2U32ToD4I32")); + assert!(importer.contains("GeneratedMmaImportAdapter::C4I32A4U32B4U32ToD4I32")); + assert!(importer.contains("GeneratedMmaImportAdapter::C4F32A4U32B4U32ToD4F32")); + assert!(importer.contains("GeneratedMmaImportAdapter::C2U32A4U32B4U32ToD2U32")); + assert!(importer.contains("(u32_ty, 2, u32_ty, 4, true, u32_ty, 4, true, u32_ty, 2)")); + assert!(importer.contains("let (c_array, last_op) = rvalue::translate_operand(")); + assert!(importer.contains("ctx, body, &args[0], value_map")); + assert!(importer.contains("let (a_value, last_after_a) = rvalue::translate_operand(")); + assert!(importer.contains("ctx, body, &args[1], value_map")); + assert!(importer.contains("let (b_value, last_after_b) = rvalue::translate_operand(")); + assert!(importer.contains("ctx, body, &args[2], value_map")); + assert!(importer.contains("operands.push(metadata)")); + assert!(importer.contains("operands.push(selector_value)")); + assert!(importer.contains("SparseMmaOp::get_concrete_op_info")); + + let lowering = render_lowering(&catalog, "test-hash"); + assert_eq!( + lowering + .matches("impl MirToLlvmConversion for SparseMmaOp") + .count(), + 1 + ); + assert!(lowering.contains("convert_generated_sparse_mma")); + for record in &records { + assert!(lowering.contains(&sparse_mma_template(record))); + assert!(lowering.contains(&sparse_mma_constraints(record))); + } + + let targets = render_targets(&catalog, "test-hash"); + assert!(targets.contains("GeneratedIntrinsicVariant::SparseMma")); + assert!(targets.contains("Operation::get_op::")); + assert!(targets.contains("get_attr_nvvm_sparse_mma_selector")); + assert!(targets.contains("GeneratedSparseMmaSelector::ImmediateZeroOrOne")); + assert!(targets.contains("GeneratedSparseMmaSelector::ImmediateZero")); + assert!(targets.contains("GeneratedSparseMmaShape::M16n8k64")); + assert!(targets.contains("SparseMmaShapeAttr::M16n8k64")); + assert!(targets.contains("GeneratedSparseMmaShape::M16n8k128")); + assert!(targets.contains("SparseMmaShapeAttr::M16n8k128")); + assert!(targets.contains("GeneratedSparseMmaElement::S4")); + assert!(targets.contains("GeneratedSparseMmaElement::U4")); + assert!(targets.contains("GeneratedSparseMmaElement::E2m1")); + assert!(targets.contains("GeneratedSparseMmaElement::E5m2")); + assert!(targets.contains("SparseMmaElementAttr::S4")); + assert!(targets.contains("SparseMmaElementAttr::U4")); + assert!(targets.contains("GeneratedSparseMmaMetadata::Standard")); + assert!(targets.contains("SparseMmaMetadataAttr::Standard")); + assert!(targets.contains("GeneratedSparseMmaMetadata::Ordered")); + assert!(targets.contains("SparseMmaMetadataAttr::Ordered")); + assert!(targets.contains("GeneratedSparseMmaAccumulator::F16")); + assert!(targets.contains("GeneratedSparseMmaAccumulator::F32")); + assert!(targets.contains("GeneratedSparseMmaOverflow::NotApplicable")); + assert!(targets.contains("GeneratedHardwareAlternative::ExactArchitecture(120)")); + + assert_eq!(raw.matches(SPARSE_MMA_STANDARD_METADATA_RULE).count(), 32); + assert_eq!(raw.matches(SPARSE_MMA_ORDERED_METADATA_RULE).count(), 82); + assert_eq!( + compatibility + .matches(SPARSE_MMA_STANDARD_METADATA_RULE) + .count(), + 32 + ); + assert_eq!( + compatibility + .matches(SPARSE_MMA_ORDERED_METADATA_RULE) + .count(), + 82 + ); + assert!( + lowering.contains("mma.sp::ordered_metadata.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32") + ); + assert!( + lowering.contains("mma.sp::ordered_metadata.sync.aligned.m16n8k64.row.col.s32.s8.s8.s32") + ); + assert!( + lowering.contains("mma.sp::ordered_metadata.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32") + ); + assert!( + lowering.contains("mma.sp::ordered_metadata.sync.aligned.m16n8k128.row.col.s32.s4.s4.s32") + ); + assert!(lowering.contains("mma.sp.sync.aligned.m16n8k64.row.col.s32.s8.s8.s32")); + assert!(lowering.contains("mma.sp.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32")); + assert!(lowering.contains("mma.sp.sync.aligned.m16n8k128.row.col.s32.s4.s4.s32")); + assert!(lowering.contains( + "mma.sp::ordered_metadata.sync.aligned.m16n8k64.row.col.kind::f8f6f4.f32.e2m1.e2m1.f32" + )); + assert!(lowering.contains( + r#"(GeneratedMmaResultType::F32, 4, 14, "mma.sp::ordered_metadata.sync.aligned.m16n8k64.row.col.kind::f8f6f4.f32.e2m1.e2m1.f32 {$0, $1, $2, $3}, {$8, $9, $10, $11}, {$12, $13, $14, $15}, {$4, $5, $6, $7}, $16, $17;", "=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r,r,r,r,n")"# + )); + assert!(lowering.contains( + r#"(GeneratedMmaResultType::I32, 2, 12, "mma.sp::ordered_metadata.sync.aligned.m16n8k64.row.col.kind::f8f6f4.f16.e2m1.e2m1.f16 {$0, $1}, {$4, $5, $6, $7}, {$8, $9, $10, $11}, {$2, $3}, $12, $13;", "=r,=r,r,r,r,r,r,r,r,r,r,r,r,n")"# + )); + let f16_probe = render_probe(&catalog, ordered_f8f6f4_f16, "test-hash"); + assert!(f16_probe.contains( + "define { i32, i32 } @probe_mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e2m1_f16_selector_0" + )); + assert!(!f16_probe.contains("_selector_1")); + + for record in &records { + let probe = render_probe(&catalog, record, "test-hash"); + assert!(probe.contains(&format!("probe_{}_selector_0", record.id))); + let selectors = sparse_mma_selector_values(record); + assert_eq!( + probe.contains(&format!("probe_{}_selector_1", record.id)), + selectors.contains(&1) + ); + assert_eq!(probe.matches("asm sideeffect").count(), selectors.len()); + assert!(probe.contains(&sparse_mma_template(record))); + assert!(probe.contains(&sparse_mma_constraints(record))); + assert!(probe.contains("attributes #0 = { convergent }")); + } + + let reference = render_reference(&catalog, "test-hash"); + assert!(reference.contains("## Sparse-MMA contracts")); + assert!(reference.contains("C, A, B, metadata, selector order")); + assert!(reference.contains("LLVM source record uses A, B, C, metadata, selector order")); + assert_eq!( + reference.matches(SPARSE_MMA_STANDARD_METADATA_RULE).count(), + 32 + ); + assert_eq!( + reference.matches(SPARSE_MMA_ORDERED_METADATA_RULE).count(), + 82 + ); + let sparse_reference = reference + .split("## Sparse-MMA contracts") + .nth(1) + .unwrap() + .split("## Packed-atomic contracts") + .next() + .unwrap(); + assert!(sparse_reference.contains("Overflow mode is not applicable.")); + assert!(!sparse_reference.contains("Integer overflow is not applicable")); + for record in &records { + assert!(reference.contains(&format!("- `{}`: runtime `unexecuted`", record.id))); + } + + let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); + assert!(outputs.contains_key(&PathBuf::from( + "crates/cuda-device/src/generated/sparse_mma.rs" + ))); + assert!(outputs.contains_key(&PathBuf::from( + "crates/dialect-nvvm/src/ops/generated/sparse_mma.rs" + ))); +} diff --git a/crates/cuda-intrinsics-gen/src/render/tests/mod.rs b/crates/cuda-intrinsics-gen/src/render/tests/mod.rs new file mode 100644 index 0000000000..729db95fa8 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/render/tests/mod.rs @@ -0,0 +1,241 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use super::*; +use crate::model::ImportedSelectionConstraints; +use crate::ptx::{InstructionPattern, OperandPattern}; + +use crate::model::{ + BackendLoweringMechanism, CatalogHardwareAlternative, CatalogHardwareTarget, CatalogSelection, + WgmmaControlAdapter, WgmmaControlMode, WgmmaControlParticipation, +}; +use std::path::Path; + +mod mma; +mod scalar_misc; +mod tensor; +mod warp_memory; + +fn catalog_with_debug_controls() -> CatalogFile { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + assert_eq!(debug_controls(&catalog).count(), 3); + catalog +} + +fn catalog_with_stmatrix() -> CatalogFile { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + assert_eq!(stmatrices(&catalog).count(), 4); + catalog +} + +fn catalog_with_clc() -> CatalogFile { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::test_catalog_with_clc(&repo_root).unwrap(); + assert_eq!(clc_intrinsics(&catalog).count(), 6); + catalog +} + +fn catalog_with_tma() -> CatalogFile { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::test_catalog_with_tma(&repo_root).unwrap(); + assert_eq!(tma_intrinsics(&catalog).count(), 111); + catalog +} + +fn catalog_with_wgmma_controls() -> CatalogFile { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let mut catalog = crate::resolve::resolve(&repo_root).unwrap(); + if wgmma_controls(&catalog).count() == 3 { + return catalog; + } + let template = catalog + .intrinsics + .iter() + .find(|record| record.cluster_barrier.is_some()) + .expect("cluster-barrier template") + .clone(); + + let recipes = [ + ( + WgmmaControlMode::Fence, + "i0317", + "wgmma_fence", + "WgmmaFenceSyncAlignedOp", + "nvvm.wgmma_fence_sync_aligned", + "llvm.nvvm.wgmma.fence.sync.aligned", + "WGMMA_FENCE_SYNC_ALIGNED", + "fence.sync.aligned", + WgmmaControlAdapter::NoArguments, + "cuda_device::wgmma::wgmma_fence", + ), + ( + WgmmaControlMode::CommitGroup, + "i0318", + "wgmma_commit_group", + "WgmmaCommitGroupSyncAlignedOp", + "nvvm.wgmma_commit_group_sync_aligned", + "llvm.nvvm.wgmma.commit_group.sync.aligned", + "WGMMA_COMMIT_GROUP_SYNC_ALIGNED", + "commit_group.sync.aligned", + WgmmaControlAdapter::NoArguments, + "cuda_device::wgmma::wgmma_commit_group", + ), + ( + WgmmaControlMode::WaitGroup, + "i0319", + "wgmma_wait_group", + "WgmmaWaitGroupSyncAlignedOp", + "nvvm.wgmma_wait_group_sync_aligned", + "llvm.nvvm.wgmma.wait_group.sync.aligned", + "WGMMA_WAIT_GROUP_SYNC_ALIGNED", + "wait_group.sync.aligned", + WgmmaControlAdapter::ConstGenericU32ToI64Immediate, + "cuda_device::wgmma::__wgmma_wait_group", + ), + ]; + + for ( + mode, + abi_id, + id, + op_type, + op_name, + llvm_symbol, + selection_record, + suffix, + adapter, + compatibility_path, + ) in recipes + { + let wait = mode == WgmmaControlMode::WaitGroup; + let mut record = template.clone(); + record.id = id.into(); + record.operation_key = format!("wgmma.control.{suffix}"); + record.family = "wgmma_control".into(); + record.source = crate::model::IntrinsicSource::LlvmImported { + source_record: format!("int_nvvm_wgmma_{}", suffix.replace('.', "_")), + }; + record.selections = vec![CatalogSelection { + source_record: selection_record.into(), + asm: format!("wgmma.{suffix};"), + predicates: vec!["Subtarget->getPTXVersion() >= 80".into(), "hasSM90a".into()], + constraints: ImportedSelectionConstraints::default(), + }]; + record.rust.abi_id = abi_id.into(); + record.rust.module = "wgmma".into(); + record.rust.name = id.into(); + record.rust.arguments = if wait { vec!["u64".into()] } else { vec![] }; + record.rust.result = "()".into(); + record.rust.safe = false; + record.rust.must_use = false; + record.rust.safe_allowlist_reason = None; + record.rust.canonical_path = + format!("cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::{abi_id}"); + record.rust.public_path = format!("cuda_intrinsics::wgmma::{id}"); + record.rust.compatibility_paths = vec![compatibility_path.into()]; + record.dialect.op_type = op_type.into(); + record.dialect.op_name = op_name.into(); + record.dialect.operands = if wait { vec!["i64".into()] } else { vec![] }; + record.dialect.results.clear(); + let llvm = record.llvm.as_mut().unwrap(); + llvm.symbol = llvm_symbol.into(); + llvm.resolved_symbol = None; + llvm.arguments = if wait { vec!["i64".into()] } else { vec![] }; + llvm.results.clear(); + llvm.properties = if wait { + vec!["ImmArg".into(), "IntrConvergent".into()] + } else { + vec!["IntrConvergent".into()] + }; + record.semantics.pure = false; + record.semantics.memory = "read_write".into(); + record.semantics.convergent = true; + record.semantics.execution_scope = "warpgroup".into(); + record.target.minimum_ptx = "8.0".parse().unwrap(); + record.target.hardware = CatalogHardwareTarget::AnyOf { + alternatives: vec![CatalogHardwareAlternative::ExactArchitecture { sm: 90 }], + }; + record.target.ptx_result = "()".into(); + record.target.targets = "sm_90a".into(); + record.target.ptx_isa_version = "9.3".into(); + record.target.ptx_isa_section = "WGMMA control".into(); + record.target.ptx_isa_url = "https://docs.nvidia.com/cuda/parallel-thread-execution/#asynchronous-warpgroup-level-matrix-instructions".into(); + for lowering in &mut record.backend_lowerings { + lowering.target.minimum_ptx = "8.0".parse().unwrap(); + lowering.target.hardware = record.target.hardware.clone(); + lowering.mechanism = match lowering.backend { + IntrinsicBackend::LlvmNvptx => BackendLoweringMechanism::TypedNvvm, + IntrinsicBackend::LibNvvm => BackendLoweringMechanism::InlinePtx, + }; + } + record.packed_atomic = None; + record.redux = None; + record.vote = None; + record.active_mask = None; + record.warp_match = None; + record.warp_barrier = None; + record.warp_shuffle = None; + record.dot_product = None; + record.packed_alu = None; + record.packed_conversion = None; + record.cp_async_copy = None; + record.cp_async_control = None; + record.cp_async_mbarrier = None; + record.mbarrier_basic = None; + record.register_mma = None; + record.sparse_mma = None; + record.prmt = None; + record.cluster_barrier = None; + record.wgmma_control = Some(crate::model::WgmmaControl { + mode, + adapter, + participation: WgmmaControlParticipation::WarpgroupAllThreadsSameInstruction, + }); + record.special_register = None; + record.ldmatrix = None; + record.lowering = "generated_wgmma_control".into(); + record.expected_ptx = InstructionPattern { + mnemonic: "wgmma".into(), + modifiers: suffix.split('.').map(str::to_owned).collect(), + operands: if wait { + vec![OperandPattern::Immediate] + } else { + vec![] + }, + }; + record.summary = format!("Generated {id} test record."); + catalog.intrinsics.push(record); + } + catalog +} + +fn raw_abi_item<'a>(raw: &'a str, id: &str) -> &'a str { + let marker = format!("Catalog ID: `{id}`"); + let marker_offset = raw.find(&marker).unwrap(); + let start = raw[..marker_offset] + .rfind("\n\n") + .map_or(0, |offset| offset + 2); + let end = raw[marker_offset..] + .find("\n}\n") + .map_or(raw.len(), |offset| marker_offset + offset + 3); + &raw[start..end] +} + +fn raw_abi_safety_block<'a>(raw: &'a str, id: &str) -> &'a str { + let item = raw_abi_item(raw, id); + let marker = "/// # Safety\n"; + let start = item.find(marker).unwrap() + marker.len(); + let end = item[start..].find("#[").unwrap() + start; + &item[start..end] +} + +fn catalog_with_tcgen05() -> CatalogFile { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::test_catalog_with_tcgen05(&repo_root).unwrap(); + assert_eq!(tcgen05_intrinsics(&catalog).count(), 233); + catalog +} diff --git a/crates/cuda-intrinsics-gen/src/render/tests/scalar_misc.rs b/crates/cuda-intrinsics-gen/src/render/tests/scalar_misc.rs new file mode 100644 index 0000000000..4bd49ea060 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/render/tests/scalar_misc.rs @@ -0,0 +1,1094 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use super::*; +use crate::model::ImportedSelectionConstraints; + +use crate::model::{ + BackendLoweringMechanism, CatalogSelection, ClcAdapter, DebugControlAdapter, + ExtendedMinMaxAdapter, ImportedAddressSpace, PackedConversionSourceFormat, RuntimeValidation, + ScalarArithmeticOperation, +}; +use crate::render::collector_targets::generated_selection_alternatives; +use crate::render::common::llvm; +use crate::render::compat::render_compat_float; +use crate::render::families::{ + dot_products, packed_alu_register_constraint, packed_alus, packed_conversion_constraint, + packed_conversion_ptx_mnemonic, packed_conversion_source, packed_conversion_source_width, + packed_conversion_typed_llvm_name, packed_conversions, prmts, scalar_conversion_ptx_mnemonic, + scalar_conversion_rounding_attr, scalar_conversion_saturation_attr, +}; +use std::collections::BTreeSet; +use std::path::Path; + +#[test] +fn unsafe_safety_docs_never_reference_missing_addr() { + let catalog = catalog_with_clc(); + let raw = render_raw_abi(&catalog, "test-hash").unwrap(); + let offenders = catalog + .intrinsics + .iter() + .filter(|record| !record.rust.safe) + .filter(|record| raw_abi_item(&raw, &record.id).contains("/// `addr`")) + .filter(|record| { + !record + .rust + .arguments + .iter() + .any(|argument| argument.starts_with('*')) + }) + .map(|record| record.id.as_str()) + .collect::>(); + assert!( + offenders.is_empty(), + "unsafe safety docs reference `addr` without a pointer argument: {}", + offenders.join(", ") + ); +} + +#[test] +fn unsafe_safety_blocks_are_exclusive_to_one_family() { + let catalog = catalog_with_clc(); + let raw = render_raw_abi(&catalog, "test-hash").unwrap(); + let mut families_by_block = BTreeMap::<&str, BTreeSet<&str>>::new(); + for record in catalog.intrinsics.iter().filter(|record| !record.rust.safe) { + families_by_block + .entry(raw_abi_safety_block(&raw, &record.id)) + .or_default() + .insert(&record.family); + } + let collisions = families_by_block + .values() + .filter(|families| families.len() != 1) + .collect::>(); + assert!( + collisions.is_empty(), + "unsafe safety blocks shared across families: {collisions:?}; render a legitimate shared contract through a shared arm" + ); +} + +#[test] +fn selection_alternatives_keep_predicates_and_constraints_grouped() { + let selections = vec![ + CatalogSelection { + source_record: "SELECT_A".into(), + asm: "op.a $dst;".into(), + predicates: vec!["HasA".into(), "HasCommon".into()], + constraints: ImportedSelectionConstraints { + address_space: Some(ImportedAddressSpace::Generic), + immediate_bindings: vec![], + }, + }, + CatalogSelection { + source_record: "SELECT_B".into(), + asm: "op.b $dst;".into(), + predicates: vec!["HasB".into()], + constraints: ImportedSelectionConstraints { + address_space: Some(ImportedAddressSpace::Shared), + immediate_bindings: vec![], + }, + }, + ]; + + let rendered = generated_selection_alternatives(&selections); + assert_eq!(rendered.matches("GeneratedSelectionAlternative").count(), 2); + assert!(rendered.contains( + "source_record: \"SELECT_A\", asm: \"op.a $dst;\", predicates: &[\"HasA\", \"HasCommon\"], constraints: GeneratedSelectionConstraints { address_space: Some(GeneratedSelectionAddressSpace::Generic), immediate_bindings: &[] }" + )); + assert!(rendered.contains( + "source_record: \"SELECT_B\", asm: \"op.b $dst;\", predicates: &[\"HasB\"], constraints: GeneratedSelectionConstraints { address_space: Some(GeneratedSelectionAddressSpace::Shared), immediate_bindings: &[] }" + )); +} + +#[test] +fn packed_alu_and_conversion_render_exact_pure_inline_ptx_adapters() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + validate_renderable(&catalog).unwrap(); + assert_eq!(packed_alus(&catalog).count(), 30); + assert_eq!(packed_conversions(&catalog).count(), 18); + + let dialect = render_dialect_packed_alu(&catalog, "test-hash"); + for op in [ + "FmaBf16x2Op", + "FmaReluBf16x2Op", + "AddBf16x2Op", + "SubBf16x2Op", + "MulBf16x2Op", + "MinBf16x2Op", + "MaxBf16x2Op", + "NegBf16x2Op", + "AbsBf16x2Op", + "FmaF16x2Op", + "FmaReluF16x2Op", + "AddF16x2Op", + "SubF16x2Op", + "MulF16x2Op", + "MinF16x2Op", + "MaxF16x2Op", + "NegF16x2Op", + "AbsF16x2Op", + "AddF32x2Op", + "AddFtzF32x2Op", + "SubF32x2Op", + "SubFtzF32x2Op", + "MulF32x2Op", + "MulFtzF32x2Op", + "FmaF32x2Op", + "FmaFtzF32x2Op", + ] { + assert!(dialect.contains(&format!("pub struct {op}"))); + assert!(dialect.contains(&format!("{op}::register(ctx)"))); + } + let conversion_dialect = render_dialect_packed_conversion(&catalog, "test-hash"); + for op in [ + "CvtF32x2Bf16x2Op", + "CvtF16x2F32Op", + "CvtRzF16x2F32Op", + "CvtRnReluF16x2F32Op", + "CvtRnReluBf16x2F32Op", + "CvtRzBf16x2F32Op", + "CvtRnSatfiniteE4m3x2F32Op", + "CvtRnSatfiniteReluE4m3x2F32Op", + "CvtRnSatfiniteE5m2x2F32Op", + "CvtRnSatfiniteReluE5m2x2F32Op", + ] { + assert!(conversion_dialect.contains(&format!("pub struct {op}"))); + assert!(conversion_dialect.contains(&format!("{op}::register(ctx)"))); + } + assert!(conversion_dialect.contains("low f16 lane")); + assert!(conversion_dialect.contains("low bf16 lane")); + assert!(conversion_dialect.contains("low e4m3 lane")); + assert!(conversion_dialect.contains("low e5m2 lane")); + assert!(conversion_dialect.contains("vec![low, high]")); + + let importer = render_importer(&catalog, "test-hash"); + assert!(importer.contains("cuda_device::bf16x2::fma_bf16x2")); + assert!(importer.contains("cuda_device::f16x2::fma_f16x2")); + assert!(importer.contains("cuda_device::convert::cvt_bf16x2_f32")); + assert!(importer.contains("cuda_device::convert::cvt_f16x2_f32")); + assert!(importer.contains("cuda_device::convert::cvt_rz_bf16x2_f32")); + assert!(importer.contains("cuda_device::convert::cvt_rn_satfinite_e4m3x2_f32")); + assert!(importer.contains("cuda_device::convert::cvt_rn_satfinite_relu_e5m2x2_f32")); + assert!(importer.contains("FmaBf16x2Op::build(ctx, arg0, arg1, arg2)")); + assert!(importer.contains("CvtF32x2Bf16x2Op::build(ctx, arg0, arg1)")); + + let lowering = render_lowering(&catalog, "test-hash"); + for mnemonic in [ + "fma.rn.bf16x2", + "fma.rn.relu.bf16x2", + "add.rn.bf16x2", + "sub.rn.bf16x2", + "mul.rn.bf16x2", + "min.bf16x2", + "max.bf16x2", + "neg.bf16x2", + "abs.bf16x2", + "fma.rn.f16x2", + "fma.rn.relu.f16x2", + "add.rn.f16x2", + "sub.rn.f16x2", + "mul.rn.f16x2", + "min.f16x2", + "max.f16x2", + "neg.f16x2", + "abs.f16x2", + ] { + assert!(lowering.contains(&format!( + "convert_generated_packed_alu(ctx, rewriter, self.get_operation(), \"{mnemonic}\", 32)" + ))); + } + for mnemonic in [ + "add.rn.f32x2", + "add.rn.ftz.f32x2", + "sub.rn.f32x2", + "sub.rn.ftz.f32x2", + "mul.rn.f32x2", + "mul.rn.ftz.f32x2", + "fma.rn.f32x2", + "fma.rn.ftz.f32x2", + ] { + assert!(lowering.contains(&format!( + "convert_generated_packed_alu(ctx, rewriter, self.get_operation(), \"{mnemonic}\", 64)" + ))); + } + for mnemonic in [ + "cvt.rn.bf16x2.f32", + "cvt.rn.f16x2.f32", + "cvt.rz.f16x2.f32", + "cvt.rn.relu.f16x2.f32", + "cvt.rn.relu.bf16x2.f32", + "cvt.rz.bf16x2.f32", + ] { + assert!(lowering.contains(&format!( + "convert_generated_packed_f32x2(ctx, rewriter, self.get_operation(), None, \"{mnemonic}\", 32)" + ))); + } + for (intrinsic, mnemonic) in [ + ("llvm_nvvm_ff_to_e4m3x2_rn", "cvt.rn.satfinite.e4m3x2.f32"), + ( + "llvm_nvvm_ff_to_e4m3x2_rn_relu", + "cvt.rn.satfinite.relu.e4m3x2.f32", + ), + ("llvm_nvvm_ff_to_e5m2x2_rn", "cvt.rn.satfinite.e5m2x2.f32"), + ( + "llvm_nvvm_ff_to_e5m2x2_rn_relu", + "cvt.rn.satfinite.relu.e5m2x2.f32", + ), + ] { + assert!(lowering.contains(&format!( + "convert_generated_packed_f32x2(ctx, rewriter, self.get_operation(), Some(\"{intrinsic}\"), \"{mnemonic}\", 16)" + ))); + } + + for record in packed_alus(&catalog) { + let probe = render_probe(&catalog, record, "test-hash"); + let register = packed_alu_register_constraint(record); + let constraints = std::iter::once(format!("={register}")) + .chain(std::iter::repeat_n( + register.to_owned(), + record.rust.arguments.len(), + )) + .collect::>() + .join(","); + assert!(probe.contains(&format!("\", \"{constraints}\""))); + assert!(!probe.contains("asm sideeffect")); + assert!(!probe.contains("~{memory}")); + } + let mut typed_probe_count = 0; + for conversion in packed_conversions(&catalog) { + let probe = render_probe(&catalog, conversion, "test-hash"); + if packed_conversion_typed_llvm_name(conversion).is_some() { + typed_probe_count += 1; + let llvm = llvm(conversion); + let symbol = llvm.resolved_symbol.as_deref().unwrap_or(&llvm.symbol); + assert!(probe.contains(&format!("declare i16 @{symbol}(float, float)"))); + assert!(probe.contains(&format!("call i16 @{symbol}(float %high, float %low)"))); + assert!(!probe.contains(" asm ")); + } else if packed_conversion_source(conversion) == PackedConversionSourceFormat::F32x2 { + assert!(probe.contains(&format!( + "asm \"{} $0, $2, $1;\", \"{}\"(float %low, float %high)", + packed_conversion_ptx_mnemonic(conversion), + packed_conversion_constraint(conversion), + ))); + assert!(!probe.contains("declare ")); + } else { + // One packed source operand, so no high/low reordering. + let source_ty = format!("i{}", packed_conversion_source_width(conversion)); + assert!(probe.contains(&format!( + "asm \"{} $0, $1;\", \"{}\"({source_ty} %packed)", + packed_conversion_ptx_mnemonic(conversion), + packed_conversion_constraint(conversion), + ))); + assert!(!probe.contains("declare ")); + } + assert!(!probe.contains("asm sideeffect")); + } + assert_eq!(typed_probe_count, 4); + + let raw = render_raw_abi(&catalog, "test-hash").unwrap(); + assert!(raw.contains("pub fn i0062(_arg0: u32, _arg1: u32, _arg2: u32) -> u32")); + assert!(raw.contains("pub fn i0071(_arg0: f32, _arg1: f32) -> u32")); + assert!(raw.contains("pub fn i0072(_arg0: u32, _arg1: u32, _arg2: u32) -> u32")); + assert!(raw.contains("pub fn i0085(_arg0: f32, _arg1: f32) -> u32")); + assert!(raw.contains("pub fn i0259(_arg0: f32, _arg1: f32) -> u16")); + assert!(raw.contains("pub fn i0262(_arg0: f32, _arg1: f32) -> u16")); + assert!(raw.contains("pub fn i0995(_arg0: u64, _arg1: u64) -> u64")); + assert!(raw.contains("pub fn i1002(_arg0: u64, _arg1: u64, _arg2: u64) -> u64")); + assert!(!raw.contains("#[must_use]\n#[inline(never)]\npub fn i0062")); + let f16_raw = raw.find("pub fn i0072").unwrap(); + assert!(raw[..f16_raw].ends_with("#[must_use]\n#[inline(never)]\n")); + + let compatibility = render_compat_packed_alu(&catalog, "test-hash", PackedAluFormat::Bf16x2); + assert!(compatibility.contains("pub fn fma_bf16x2(arg0: u32, arg1: u32, arg2: u32)")); + assert!(!compatibility.contains("fma_f16x2")); + assert!(compatibility.contains("let _ = arg0;")); + assert!(!compatibility.contains("let _ = (arg0);")); + let compatibility = render_compat_packed_alu(&catalog, "test-hash", PackedAluFormat::F16x2); + let f16_compat = compatibility.find("pub fn fma_f16x2").unwrap(); + assert!(compatibility[..f16_compat].ends_with("#[must_use]\n#[inline(never)]\n")); + assert!(!compatibility.contains("fma_bf16x2")); + let compatibility = render_compat_packed_alu(&catalog, "test-hash", PackedAluFormat::F32x2); + assert!(compatibility.contains("pub fn add_f32x2(arg0: u64, arg1: u64) -> u64")); + assert!(compatibility.contains("pub fn fma_ftz_f32x2(arg0: u64, arg1: u64, arg2: u64) -> u64")); + assert!(!compatibility.contains("fma_f16x2")); + let reference = render_reference(&catalog, "test-hash"); + assert!(reference.contains("`fma_f16x2` carries one packed `f16x2` value in a `u32`")); + assert!(reference.contains( + "native instruction starts at PTX 4.2 / `sm_53`; cuda-oxide admits it from sm_70+" + )); + assert!(reference.contains("LLVM-NVPTX PTX 6.0 on sm_70+")); + assert!(reference.contains("libNVVM PTX 4.2 on sm_75+")); + assert!(reference.contains( + "`cvt_rn_relu_f16x2_f32` converts two `f32` inputs to packed `f16x2` using nearest-even rounding with ReLU" + )); + assert!(reference.contains("pure `cvt.rz.bf16x2.f32` inline PTX")); + assert!(reference.contains( + "LLVM-NVPTX uses typed `llvm.nvvm.ff.to.e4m3x2.rn` with `[high, low]` inputs; libNVVM uses pure `cvt.rn.satfinite.e4m3x2.f32` inline PTX" + )); + let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); + assert!(outputs.contains_key(&PathBuf::from("crates/cuda-device/src/generated/bf16x2.rs"))); + assert!(outputs.contains_key(&PathBuf::from("crates/cuda-device/src/generated/f16x2.rs"))); + assert!(outputs.contains_key(&PathBuf::from("crates/cuda-device/src/generated/f32x2.rs"))); + assert!(outputs.contains_key(&PathBuf::from( + "crates/cuda-device/src/generated/convert.rs" + ))); + let compatibility = render_compat_packed_conversion( + &catalog, + "test-hash", + "cuda_device::convert::", + "convert", + ("lo", "hi"), + ); + assert!(compatibility.contains("pub fn cvt_bf16x2_f32(lo: f32, hi: f32) -> u32")); + assert!(compatibility.contains("pub fn cvt_f16x2_f32(lo: f32, hi: f32) -> u32")); + assert!(compatibility.contains("pub fn cvt_rz_bf16x2_f32(lo: f32, hi: f32) -> u32")); + assert!(compatibility.contains("pub fn cvt_rn_satfinite_e4m3x2_f32(lo: f32, hi: f32) -> u16")); + assert!( + compatibility.contains("pub fn cvt_rn_satfinite_relu_e5m2x2_f32(lo: f32, hi: f32) -> u16") + ); + assert!(!compatibility.contains("pub fn cvt_f32x2_bf16x2(")); + assert!(!outputs.contains_key(&PathBuf::from( + "crates/cuda-device/src/generated/tcgen05_conversion.rs" + ))); +} + +#[test] +fn scalar_tf32_conversions_render_one_exact_attribute_carrier() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + validate_renderable(&catalog).unwrap(); + let records = scalar_conversions(&catalog).collect::>(); + assert_eq!(records.len(), 10); + + for record in &records { + assert_eq!(record.rust.arguments, ["f32"]); + assert_eq!(record.rust.result, "u32"); + assert!(record.backend_lowerings.iter().any(|lowering| { + lowering.backend == IntrinsicBackend::LlvmNvptx + && lowering.mechanism == BackendLoweringMechanism::TypedNvvm + })); + assert!(record.backend_lowerings.iter().any(|lowering| { + lowering.backend == IntrinsicBackend::LibNvvm + && lowering.mechanism == BackendLoweringMechanism::InlinePtx + })); + } + + let dialect = render_dialect_scalar_conversion(&catalog, "test-hash"); + assert_eq!(dialect.matches("pub struct ScalarConversionOp;").count(), 1); + assert!(dialect.contains("pub enum ScalarConversionRoundingAttr")); + assert!(dialect.contains("pub enum ScalarConversionSaturationAttr")); + assert_eq!( + dialect + .matches("ScalarConversionOp::register(ctx);") + .count(), + 1 + ); + for record in &records { + assert!(dialect.contains(scalar_conversion_rounding_attr(record))); + assert!(dialect.contains(scalar_conversion_saturation_attr(record))); + } + + let importer = render_importer(&catalog, "test-hash"); + assert_eq!( + importer + .matches("let intrinsic = ScalarConversionOp::build") + .count(), + 10 + ); + let lowering = render_lowering(&catalog, "test-hash"); + assert_eq!( + lowering + .matches("impl MirToLlvmConversion for ScalarConversionOp") + .count(), + 1 + ); + for record in &records { + assert!(lowering.contains(&record.llvm_identifier())); + assert!(lowering.contains(&scalar_conversion_ptx_mnemonic(record))); + let probe = render_probe(&catalog, record, "test-hash"); + assert!(probe.contains(&format!("declare i32 @{}(float)", llvm(record).symbol))); + } + + let targets = render_targets(&catalog, "test-hash"); + let production_targets = targets.split("\n#[cfg(test)]").next().unwrap(); + assert_eq!( + production_targets + .matches("GeneratedIntrinsicVariant::ScalarConversion {") + .count(), + 11 + ); + assert!(targets.contains("Operation::get_op::")); + assert!(targets.contains("GeneratedScalarConversionRounding::NearestAway")); + assert!(targets.contains("GeneratedScalarConversionSaturation::ReluSatfinite")); + + let compatibility = render_compat_packed_conversion( + &catalog, + "test-hash", + "cuda_device::convert::", + "convert", + ("lo", "hi"), + ); + assert!(compatibility.contains("pub fn cvt_rna_tf32_f32(value: f32) -> u32")); + assert!(compatibility.contains("pub fn cvt_rz_relu_satfinite_tf32_f32(value: f32) -> u32")); + let raw = render_raw_abi(&catalog, "test-hash").unwrap(); + assert!(raw.contains("pub fn i0368(_arg0: f32) -> u32")); + assert!(raw.contains("pub fn i0377(_arg0: f32) -> u32")); +} + +#[test] +fn packed_atomic_family_uses_one_attribute_dispatch_impl() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + let rendered = render_lowering(&catalog, "test-hash"); + assert_eq!( + rendered + .matches("impl MirToLlvmConversion for PackedAtomicAddOp") + .count(), + 1 + ); + assert_eq!( + rendered.matches("PackedAtomicFormatAttr::F16x2)").count(), + 1 + ); + assert_eq!( + rendered.matches("PackedAtomicFormatAttr::Bf16x2)").count(), + 1 + ); + for (op_type, ptx_type) in [ + ("NvvmAtomAddF16x2Op", "f16x2"), + ("NvvmAtomAddBf16x2Op", "bf16x2"), + ] { + assert_eq!( + rendered + .matches(&format!("impl MirToLlvmConversion for {op_type}")) + .count(), + 1 + ); + assert!(rendered.contains(&format!( + "convert_packed_atom_add(ctx, rewriter, self.get_operation(), \"{ptx_type}\")" + ))); + } + + let dialect = render_dialect_packed_atomic(&catalog, "test-hash"); + for (op_type, op_name) in [ + ("NvvmAtomAddF16x2Op", "nvvm.atom_add_f16x2"), + ("NvvmAtomAddBf16x2Op", "nvvm.atom_add_bf16x2"), + ] { + assert!(dialect.contains(&format!("pub struct {op_type};"))); + assert!(dialect.contains(&format!("name = \"{op_name}\""))); + assert!(dialect.contains(&format!("{op_type}::register(ctx);"))); + } +} + +#[test] +fn packed_atomic_raw_abi_is_unsafe_and_must_use() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + let rendered = render_raw_abi(&catalog, "test-hash").unwrap(); + for abi_id in ["i0014", "i0015"] { + let signature = format!("pub unsafe fn {abi_id}(_arg0: *mut u32, _arg1: u32) -> u32"); + let index = rendered.find(&signature).unwrap(); + assert!(rendered[..index].ends_with("#[must_use]\n#[inline(never)]\n")); + } +} + +#[test] +fn packed_atomic_compatibility_preserves_paths_signature_and_safety() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + let rendered = render_compat_packed_atomic(&catalog, "test-hash"); + + for name in ["atom_add_f16x2", "atom_add_bf16x2"] { + let signature = format!("pub unsafe fn {name}(addr: *mut u32, val: u32) -> u32"); + let index = rendered.find(&signature).unwrap(); + assert!(rendered[..index].ends_with("#[must_use]\n#[inline(never)]\n")); + assert!(rendered.contains(&format!( + "unreachable!(\"{name} called outside CUDA kernel context\")" + ))); + } + assert!(rendered.contains("relaxed GPU-scope operation")); + assert!(rendered.contains("low lane first")); + assert!(rendered.contains("may not form one old 32-bit snapshot")); + assert!(rendered.contains("Requires PTX 6.2 and `sm_70+`")); + assert!(rendered.contains("Requires PTX 7.8 and `sm_90+`")); + assert!(rendered.contains("four writable, four-byte-aligned bytes in global memory")); + assert!(rendered.contains("whole-word atomic or non-atomic lane access")); + assert!(rendered.contains("Racing atomics must use mutually inclusive scopes")); + + let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); + assert_eq!( + outputs.get(&PathBuf::from("crates/cuda-device/src/generated/atomic.rs")), + Some(&rendered) + ); +} + +#[test] +fn dot_product_rendering_preserves_stable_paths_and_low_selector() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + validate_renderable(&catalog).unwrap(); + assert_eq!(dot_products(&catalog).count(), 4); + + let compatibility = render_compat_dotprod(&catalog, "test-hash"); + for name in ["dp4a_s32", "dp4a_u32", "dp2a_s32", "dp2a_u32"] { + assert!(compatibility.contains(&format!("pub fn {name}("))); + } + + let dialect = render_dialect_dotprod(&catalog, "test-hash"); + assert!(dialect.contains("pub struct Dp4aS32Op")); + assert!(dialect.contains("pub struct Dp2aU32Op")); + assert!(dialect.contains("NOpdsInterface<3>")); + + let importer = render_importer(&catalog, "test-hash"); + assert!(importer.contains("cuda_device::dotprod::dp2a_s32")); + assert!(importer.contains("let dot = Dp2aS32Op::build(ctx, a, b, c)")); + assert!(importer.contains("set_generated_intrinsic_marker(ctx, dot, \"v1:i0032\")")); + + let lowering = render_lowering(&catalog, "test-hash"); + assert!(lowering.contains("impl MirToLlvmConversion for Dp4aS32Op")); + assert!(lowering.contains("\"llvm_nvvm_idp2a_s_s\"")); + assert!(lowering.contains("\"dp2a.lo.s32.s32 $0, $1, $2, $3;\"")); + assert!(lowering.contains("\"dp2a.lo.s32.s32 $0, $1, $2, $3;\", true)")); + + let low = dot_products(&catalog) + .find(|record| record.id == "dp2a_s32") + .unwrap(); + let probe = render_probe(&catalog, low, "test-hash"); + assert!(probe.contains("call i32 @llvm.nvvm.idp2a.s.s(i32 %a, i32 %b, i1 false, i32 %c)")); + assert!(!probe.contains("i1 true")); + + let target = render_targets(&catalog, "test-hash"); + assert!(target.contains("GeneratedImmediateBinding { argument_index: 2, value: 0 }")); + assert!(!target.contains("GeneratedImmediateBinding { argument_index: 2, value: -1 }")); +} + +#[test] +fn prmt_rendering_keeps_modes_and_zero_source_exact() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + validate_renderable(&catalog).unwrap(); + assert_eq!(prmts(&catalog).count(), 7); + + let dialect = render_dialect_prmt(&catalog, "test-hash"); + assert_eq!(dialect.matches("pub struct PrmtOp").count(), 1); + for mode in ["Generic", "F4e", "B4e", "Rc8", "Ecl", "Ecr", "Rc16"] { + assert!(dialect.contains(mode)); + } + + let lowering = render_lowering(&catalog, "test-hash"); + for mode in ["rc8", "ecl", "ecr", "rc16"] { + assert!(lowering.contains(&format!("prmt.b32.{mode} $0, $1, 0, $2;"))); + } + assert!(lowering.contains("prmt.b32 $0, $1, $2, $3;")); + assert!(lowering.contains("prmt.b32.f4e $0, $1, $2, $3;")); + assert!(lowering.contains("prmt.b32.b4e $0, $1, $2, $3;")); + assert!(lowering.contains("convert_generated_prmt")); + + let targets = render_targets(&catalog, "test-hash"); + assert!(targets.contains("GeneratedIntrinsicVariant::Prmt")); + assert!(targets.contains("Operation::get_op::")); + assert!(targets.contains("GeneratedPrmtMode::Rc16")); + + let compatibility = render_compat_prmt(&catalog, "test-hash"); + assert_eq!(compatibility.matches("pub fn prmt").count(), 7); +} + +#[test] +fn selected_special_register_probes_use_the_llvm_route() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + + let clock = catalog + .intrinsics + .iter() + .find(|record| record.id == "clock") + .unwrap(); + let typed = render_probe(&catalog, clock, "test-hash"); + assert!(typed.contains("declare i32 @llvm.nvvm.read.ptx.sreg.clock()")); + assert!(typed.contains("call i32 @llvm.nvvm.read.ptx.sreg.clock()")); + + let gridid = catalog + .intrinsics + .iter() + .find(|record| record.id == "gridid") + .unwrap(); + let inline = render_probe(&catalog, gridid, "test-hash"); + assert!(inline.contains("define i64 @probe_gridid_llvm_nvptx()")); + assert!(inline.contains("call i64 asm \"mov.u64 $0, %gridid;\", \"=l\"()")); +} + +#[test] +fn clc_rendering_preserves_api_and_uses_typed_llvm_routes() { + let catalog = catalog_with_clc(); + validate_renderable(&catalog).unwrap(); + + let compatibility = render_compat_clc(&catalog, "test-hash"); + assert!( + compatibility + .contains("pub unsafe fn clc_try_cancel(response: *mut u8, mbar: *mut Barrier)") + ); + assert!( + compatibility.contains( + "pub unsafe fn clc_try_cancel_multicast(response: *mut u8, mbar: *mut Barrier)" + ) + ); + assert!( + compatibility + .contains("pub unsafe fn clc_query_is_canceled(resp_lo: u64, resp_hi: u64) -> u32") + ); + + let dialect = render_dialect_clc(&catalog, "test-hash"); + for op in [ + "ClcTryCancelOp", + "ClcTryCancelMulticastOp", + "ClcQueryIsCanceledOp", + "ClcQueryGetFirstCtaidXOp", + "ClcQueryGetFirstCtaidYOp", + "ClcQueryGetFirstCtaidZOp", + ] { + assert_eq!(dialect.matches(&format!("pub struct {op}")).count(), 1); + assert!(dialect.contains(&format!("{op}::register(ctx)"))); + } + assert_eq!(dialect.matches("NResultsInterface<0>").count(), 2); + assert_eq!(dialect.matches("NResultsInterface<1>").count(), 4); + assert_eq!(dialect.matches("impl Verify for ClcQuery").count(), 4); + assert!(dialect.contains("is_integer_width(ctx, op.get_operand(0).get_type(ctx), 64)")); + assert!(dialect.contains("is_integer_width(ctx, op.get_result(0).get_type(ctx), 32)")); + + let importer = render_importer(&catalog, "test-hash"); + assert!(importer.contains("cuda_device::clc::clc_try_cancel")); + assert!(importer.contains("ClcTryCancelOp::get_concrete_op_info()")); + assert!(importer.contains("ClcQueryIsCanceledOp::get_concrete_op_info()")); + assert!(importer.contains("helpers::emit_store_result_and_goto")); + + let lowering = render_lowering(&catalog, "test-hash"); + assert!( + lowering.contains("clc::{convert_generated_clc_query, convert_generated_clc_try_cancel}") + ); + assert!(lowering.contains( + "convert_generated_clc_try_cancel(ctx, rewriter, self.get_operation(), operands_info, \"llvm__nvvm_dclusterlaunchcontrol_dtry_ucancel_dasync_dshared\")" + )); + assert!(lowering.contains( + "convert_generated_clc_query(ctx, rewriter, self.get_operation(), operands_info, \"llvm__nvvm_dclusterlaunchcontrol_dquery_ucancel_dis_ucanceled\", true)" + )); + + let request = clc_intrinsics(&catalog) + .find(|record| record.id == "clc_try_cancel") + .unwrap(); + let request_probe = render_probe(&catalog, request, "test-hash"); + assert_eq!(request_probe.matches("addrspacecast ptr").count(), 2); + assert!(request_probe.contains( + "call void @llvm.nvvm.clusterlaunchcontrol.try_cancel.async.shared(ptr addrspace(3) %response, ptr addrspace(3) %mbarrier)" + )); + + let query = clc_intrinsics(&catalog) + .find(|record| record.id == "clc_query_is_canceled") + .unwrap(); + let query_probe = render_probe(&catalog, query, "test-hash"); + assert!(query_probe.contains("%response_high_shifted = shl i128 %response_high_i128, 64")); + assert!(query_probe.contains("%response = or i128 %response_low_i128, %response_high_shifted")); + assert!(query_probe.contains("%result = zext i1 %raw_result to i32")); + + let targets = render_targets(&catalog, "test-hash"); + assert!(targets.contains( + "GeneratedHardwareTarget::AnyOf(&[GeneratedHardwareAlternative::ExactArchitecture(100), GeneratedHardwareAlternative::ExactArchitecture(101), GeneratedHardwareAlternative::ExactArchitecture(103), GeneratedHardwareAlternative::ExactArchitecture(110), GeneratedHardwareAlternative::ExactArchitecture(120), GeneratedHardwareAlternative::ExactArchitecture(121)])" + )); + + let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); + assert!(outputs.contains_key(&PathBuf::from("crates/cuda-device/src/generated/clc.rs"))); + assert!(outputs.contains_key(&PathBuf::from( + "crates/dialect-nvvm/src/ops/generated/clc.rs" + ))); + + let mut wrong_adapter = catalog; + wrong_adapter + .intrinsics + .iter_mut() + .find(|record| record.id == "clc_query_is_canceled") + .unwrap() + .clc + .as_mut() + .unwrap() + .adapter = ClcAdapter::PairU64ToI128U32; + assert!(validate_renderable(&wrong_adapter).is_err()); +} + +#[test] +fn debug_control_rendering_preserves_api_immediates_and_side_effects() { + let catalog = catalog_with_debug_controls(); + validate_renderable(&catalog).unwrap(); + assert_eq!(debug_controls(&catalog).count(), 3); + + let compatibility = render_compat_debug_control(&catalog, "test-hash"); + assert!(compatibility.contains("pub fn trap() -> !")); + assert!(compatibility.contains("pub fn breakpoint()")); + assert!(compatibility.contains("pub fn prof_trigger()")); + assert!(compatibility.contains("pub(crate) fn __prof_trigger(_event_id: u32)")); + assert!(compatibility.contains("__prof_trigger(N);")); + + let dialect = render_dialect_debug_control(&catalog, "test-hash"); + for op in ["TrapOp", "BreakpointOp", "PmEventOp"] { + assert_eq!(dialect.matches(&format!("pub struct {op}")).count(), 1); + assert!(dialect.contains(&format!("{op}::register(ctx)"))); + } + assert!(dialect.contains("event_id: u32")); + assert!(dialect.contains("pub fn new_with_event_id")); + assert!(dialect.contains("pub fn get_event_id")); + assert!(dialect.contains("filter(|value| *value <= 15)")); + assert!(dialect.contains("requires a u32 event ID in 0..=15")); + + let importer = render_importer(&catalog, "test-hash"); + assert!(importer.contains("cuda_device::debug::__prof_trigger")); + assert!(importer.contains("mir::Operand::Constant")); + assert!(importer.contains("u32::try_from(value)")); + assert!(importer.contains("filter(|value| *value <= 15)")); + assert!(importer.contains("PmEventOp::build(ctx, event_id)")); + assert!(importer.contains("MirUnreachableOp::get_concrete_op_info()")); + assert!(importer.contains("prof_trigger requires a compile-time constant event ID in 0..=15")); + + let lowering = render_lowering(&catalog, "test-hash"); + for op in ["TrapOp", "BreakpointOp", "PmEventOp"] { + assert!(lowering.contains(&format!("impl MirToLlvmConversion for {op}"))); + } + assert!(lowering.contains( + "inline_asm_sideeffect(ctx, rewriter, op, void_ty.into(), vec![], \"trap;\", \"\")" + )); + assert!(lowering.contains( + "inline_asm_sideeffect(ctx, rewriter, op, void_ty.into(), vec![], \"brkpt;\", \"\")" + )); + assert!(lowering.contains("let template = format!(\"pmevent {event_id};\")")); + + for record in debug_controls(&catalog) { + let probe = render_probe(&catalog, record, "test-hash"); + assert!(probe.contains("call void asm sideeffect")); + assert!(!probe.contains("attributes #0 = { convergent }")); + } + + let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); + assert!(outputs.contains_key(&PathBuf::from( + "crates/cuda-device/src/generated/debug_control.rs" + ))); + assert!(outputs.contains_key(&PathBuf::from( + "crates/dialect-nvvm/src/ops/generated/debug_control.rs" + ))); + + let mut wrong_adapter = catalog.clone(); + wrong_adapter + .intrinsics + .iter_mut() + .find(|record| record.id == "pmevent") + .unwrap() + .debug_control + .as_mut() + .unwrap() + .adapter = DebugControlAdapter::Direct; + assert!(validate_renderable(&wrong_adapter).is_err()); + + let mut wrong_runtime = catalog; + wrong_runtime + .intrinsics + .iter_mut() + .find(|record| record.id == "trap") + .unwrap() + .debug_control + .as_mut() + .unwrap() + .runtime_validation = RuntimeValidation::Executed; + assert!(validate_renderable(&wrong_runtime).is_err()); +} + +#[test] +fn scalar_arithmetic_rendering_uses_one_closed_carrier() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + validate_renderable(&catalog).unwrap(); + let records = scalar_arithmetics(&catalog).collect::>(); + assert_eq!(records.len(), 64); + assert!(records.iter().all(|record| { + record.semantics.pure + == (record.scalar_arithmetic.as_ref().unwrap().operation + != ScalarArithmeticOperation::Div) + })); + + let compatibility = render_compat_float(&catalog, "test-hash"); + assert!(compatibility.contains("pub fn mul_rn_f64(arg0: f64, arg1: f64) -> f64")); + assert!( + compatibility.contains("pub fn fma_rp_ftz_sat_f32(arg0: f32, arg1: f32, arg2: f32) -> f32") + ); + assert!(compatibility.contains("pub fn add_rp_ftz_sat_f32(arg0: f32, arg1: f32) -> f32")); + assert_eq!(compatibility.matches("#[must_use]").count(), 113); + + let dialect = render_dialect_scalar_arithmetic(&catalog, "test-hash"); + assert_eq!(dialect.matches("pub struct ScalarArithmeticOp").count(), 1); + for attr in [ + "ScalarArithmeticFormatAttr", + "ScalarArithmeticOperationAttr", + "ScalarArithmeticRoundingAttr", + "ScalarArithmeticSubnormalAttr", + "ScalarArithmeticSaturationAttr", + ] { + assert!(dialect.contains(attr)); + assert!(dialect.contains(&format!("{attr}::register(ctx)"))); + } + assert!(dialect.contains("variant is not admitted")); + assert!(dialect.contains("ScalarArithmeticOperationAttr::Fma => 3")); + + let importer = render_importer(&catalog, "test-hash"); + assert_eq!( + importer.matches("ScalarArithmeticOp::build(ctx").count(), + 64 + ); + assert!(importer.contains("cuda_device::float::fma_rp_ftz_sat_f32")); + assert!(importer.contains("cuda_device::float::add_rp_ftz_sat_f32")); + + let lowering = render_lowering(&catalog, "test-hash"); + assert_eq!( + lowering + .matches("impl MirToLlvmConversion for ScalarArithmeticOp") + .count(), + 1 + ); + assert!(lowering.contains("llvm_nvvm_mul_rn_d")); + assert!(lowering.contains("fma.rp.ftz.sat.f32")); + assert!(lowering.contains("llvm_nvvm_add_rn_d")); + assert!(lowering.contains("add.rp.sat.ftz.f32")); + assert!(lowering.contains("recipe.0, recipe.1, recipe.2, recipe.3")); + + let targets = render_targets(&catalog, "test-hash"); + assert!(targets.contains("GeneratedIntrinsicVariant::ScalarArithmetic")); + assert!(targets.contains("GeneratedScalarArithmeticFormat::F64")); + assert!(targets.contains("GeneratedScalarArithmeticSaturation::Sat")); + assert!(targets.contains("Operation::get_op::")); + + let f64_mul = records + .iter() + .copied() + .find(|record| record.id == "mul_rn_f64") + .unwrap(); + let f64_probe = render_probe(&catalog, f64_mul, "test-hash"); + assert!(f64_probe.contains("declare double @llvm.nvvm.mul.rn.d(double, double)")); + let f32_sat = records + .iter() + .copied() + .find(|record| record.id == "fma_rp_ftz_sat_f32") + .unwrap(); + let f32_probe = render_probe(&catalog, f32_sat, "test-hash"); + assert!( + f32_probe.contains("call float asm \"fma.rp.ftz.sat.f32 $0, $1, $2, $3;\", \"=f,f,f,f\"") + ); + assert!(!f32_probe.contains("@llvm.nvvm.fma.rp.ftz.sat.f")); + + let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); + assert!(outputs.contains_key(&PathBuf::from("crates/cuda-device/src/generated/float.rs"))); + assert!(outputs.contains_key(&PathBuf::from( + "crates/dialect-nvvm/src/ops/generated/scalar_arithmetic.rs" + ))); +} + +#[test] +fn extended_minmax_rendering_is_closed_pure_and_exact() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + validate_renderable(&catalog).unwrap(); + let records = extended_minmax(&catalog).collect::>(); + assert_eq!(records.len(), 52); + assert_eq!( + records + .iter() + .filter(|record| record.target.minimum_ptx.encoded() == 70) + .count(), + 20 + ); + assert_eq!( + records + .iter() + .filter(|record| record.target.minimum_ptx.encoded() == 72) + .count(), + 32 + ); + for id in ["min_f16x2", "max_f16x2", "min_bf16x2", "max_bf16x2"] { + assert_eq!( + catalog + .intrinsics + .iter() + .find(|record| record.id == id) + .unwrap() + .family, + "packed_alu" + ); + assert!(records.iter().all(|record| record.id != id)); + } + + let f16 = render_compat_packed_alu(&catalog, "test-hash", PackedAluFormat::F16x2); + assert!(f16.contains("pub fn min_ftz_nan_xorsign_abs_f16x2(a: u32, b: u32) -> u32")); + let bf16 = render_compat_packed_alu(&catalog, "test-hash", PackedAluFormat::Bf16x2); + assert!(bf16.contains("pub fn max_nan_xorsign_abs_bf16x2(a: u32, b: u32) -> u32")); + let float = render_compat_float(&catalog, "test-hash"); + assert!(float.contains("pub fn min_ftz_nan_xorsign_abs_f32(a: f32, b: f32) -> f32")); + + let dialect = render_dialect_extended_minmax(&catalog, "test-hash"); + assert_eq!(dialect.matches("pub struct ExtendedMinMaxOp").count(), 1); + assert_eq!( + dialect + .matches(" (ExtendedMinMaxFormatAttr::") + .count(), + 52 + ); + for attr in [ + "ExtendedMinMaxFormatAttr", + "ExtendedMinMaxOperationAttr", + "ExtendedMinMaxSubnormalAttr", + "ExtendedMinMaxNanAttr", + "ExtendedMinMaxXorSignAbsAttr", + ] { + assert!(dialect.contains(attr)); + assert!(dialect.contains(&format!("{attr}::register(ctx)"))); + } + assert!(dialect.contains("variant is not admitted")); + + let importer = render_importer(&catalog, "test-hash"); + assert_eq!(importer.matches("ExtendedMinMaxOp::build(ctx").count(), 52); + assert!(importer.contains("cuda_device::float::min_xorsign_abs_f32")); + assert!(importer.contains("cuda_device::bf16x2::max_nan_bf16x2")); + + let lowering = render_lowering(&catalog, "test-hash"); + assert_eq!( + lowering + .matches("impl MirToLlvmConversion for ExtendedMinMaxOp") + .count(), + 1 + ); + assert!(lowering.contains("convert_generated_extended_minmax")); + assert!(lowering.contains("min.ftz.NaN.xorsign.abs.f32")); + assert!(lowering.contains("max.NaN.xorsign.abs.bf16x2")); + + let targets = render_targets(&catalog, "test-hash"); + assert!(targets.contains("GeneratedIntrinsicVariant::ExtendedMinMax")); + assert!(targets.contains("GeneratedPtxVersion::from_encoded(72)")); + assert!(targets.contains("Operation::get_op::")); + assert!(targets.contains("ExtendedMinMaxXorSignAbsAttr::Enabled")); + + let packed_record = records + .iter() + .copied() + .find(|record| record.id == "min_ftz_nan_f16x2") + .unwrap(); + let packed_probe = render_probe(&catalog, packed_record, "test-hash"); + assert!(packed_probe.contains("call i32 asm \"min.ftz.NaN.f16x2 $0, $1, $2;\", \"=r,r,r\"")); + assert!(!packed_probe.contains("sideeffect")); + assert!(!packed_probe.contains("convergent")); + let f32_record = records + .iter() + .copied() + .find(|record| record.id == "max_ftz_nan_xorsign_abs_f32") + .unwrap(); + let f32_probe = render_probe(&catalog, f32_record, "test-hash"); + assert!( + f32_probe + .contains("call float asm \"max.ftz.NaN.xorsign.abs.f32 $0, $1, $2;\", \"=f,f,f\"") + ); + // Scalar 16-bit forms ride in `h` registers, not the packed `r` pair. + let half_record = records + .iter() + .copied() + .find(|record| record.id == "min_ftz_nan_f16") + .unwrap(); + let half_probe = render_probe(&catalog, half_record, "test-hash"); + assert!(half_probe.contains("call i16 asm \"min.ftz.NaN.f16 $0, $1, $2;\", \"=h,h,h\"")); + let bf16_record = records + .iter() + .copied() + .find(|record| record.id == "max_nan_xorsign_abs_bf16") + .unwrap(); + let bf16_probe = render_probe(&catalog, bf16_record, "test-hash"); + assert!( + bf16_probe.contains("call i16 asm \"max.NaN.xorsign.abs.bf16 $0, $1, $2;\", \"=h,h,h\"") + ); + assert!(lowering.contains("MinMaxCarrier::Half16")); + + let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); + assert!(outputs.contains_key(&PathBuf::from( + "crates/dialect-nvvm/src/ops/generated/extended_minmax.rs" + ))); + for module in ["float", "f16", "bf16", "f16x2", "bf16x2"] { + assert!(outputs.contains_key(&PathBuf::from(format!( + "crates/cuda-device/src/generated/{module}.rs" + )))); + } + + let mut wrong_adapter = catalog; + wrong_adapter + .intrinsics + .iter_mut() + .find(|record| record.id == "min_xorsign_abs_f32") + .unwrap() + .extended_minmax + .as_mut() + .unwrap() + .adapter = ExtendedMinMaxAdapter::DirectPackedU32; + assert!(validate_renderable(&wrong_adapter).is_err()); +} + +#[test] +fn extended_minmax_alone_renders_float_wrappers() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let mut catalog = crate::resolve::resolve(&repo_root).unwrap(); + catalog + .intrinsics + .retain(|record| record.family == "extended_minmax"); + + let (path, float) = render_compat_float_output(&catalog, "test-hash") + .expect("extended min/max alone must render float wrappers"); + assert_eq!( + path, + PathBuf::from("crates/cuda-device/src/generated/float.rs") + ); + assert!(float.contains("pub fn min_xorsign_abs_f32(a: f32, b: f32) -> f32")); + assert!(!float.contains("pub fn add_rn_f32")); +} + +/// Every reference row must carry exactly the six cells of its header. +/// +/// A `|` in cell data splits the row before inline parsing, so it splits +/// the cell even inside a code span. Three catalog entries write a PTX +/// destination pair as ``, and their rows rendered with +/// a seventh cell, which drops the backend-evidence column. Counting +/// unescaped pipes per row is the direct check. +#[test] +fn reference_rows_have_one_cell_per_header_column() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + let reference = render_reference(&catalog, "test-hash"); + + let mut rows = 0usize; + let mut carried_a_pipe = 0usize; + for line in reference + .lines() + .filter(|line| line.starts_with("| `cuda_intrinsics::")) + { + rows += 1; + if line.contains("\\|") { + carried_a_pipe += 1; + } + // Delimiters only: drop every escaped pipe first, then a six-column + // row leaves seven. + let delimiters = line.replace("\\|", "").matches('|').count(); + assert_eq!( + delimiters, 7, + "row has {delimiters} cell delimiters, expected 7: {line}" + ); + } + + assert!( + rows > 900, + "expected the full reference, rendered {rows} rows" + ); + assert!( + carried_a_pipe >= 3, + "expected the elect.sync and match.all.sync patterns to need escaping, \ + found {carried_a_pipe} rows with an escaped pipe" + ); +} diff --git a/crates/cuda-intrinsics-gen/src/render/tests/tensor.rs b/crates/cuda-intrinsics-gen/src/render/tests/tensor.rs new file mode 100644 index 0000000000..12ca94f437 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/render/tests/tensor.rs @@ -0,0 +1,1133 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use super::*; +use crate::model::{CatalogTargetAlternative, TargetSelectorBinding}; + +use crate::model::{ + CatalogHardwareAlternative, CatalogHardwareTarget, CatalogTargetContract, IntrinsicSource, + Tcgen05Adapter, Tcgen05MmaBUsage, Tcgen05MmaForm, Tcgen05MmaKind, TmaAdapter, +}; +use crate::render::common::{generated_hardware_target, hardware_target_label, llvm}; +use crate::render::families::{cluster_barriers, tcgen05_mma_inline_asm}; +use std::path::Path; + +#[test] +fn paired_target_matrix_renders_selectors_and_ptx_floors() { + let target = CatalogHardwareTarget::TargetMatrix { + contracts: vec![CatalogTargetContract { + selectors: vec![TargetSelectorBinding { + name: "kind".into(), + value: "i8".into(), + }], + alternatives: vec![ + CatalogTargetAlternative { + minimum_ptx: "8.8".parse().unwrap(), + hardware: CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, + }, + CatalogTargetAlternative { + minimum_ptx: "9.0".parse().unwrap(), + hardware: CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, + }, + ], + }], + }; + + assert_eq!( + hardware_target_label(&target), + "sm_100a at PTX 8.8 or sm_110a at PTX 9.0 for kind=i8" + ); + assert_eq!( + generated_hardware_target(&target), + "GeneratedHardwareTarget::TargetMatrix { contracts: &[GeneratedTargetContract { selectors: &[GeneratedTargetSelectorBinding { name: \"kind\", value: \"i8\" }], alternatives: &[GeneratedTargetAlternative { minimum_ptx: GeneratedPtxVersion::from_encoded(88), hardware: GeneratedHardwareAlternative::ExactArchitecture(100) }, GeneratedTargetAlternative { minimum_ptx: GeneratedPtxVersion::from_encoded(90), hardware: GeneratedHardwareAlternative::ExactArchitecture(110) }] }] }" + ); +} + +#[test] +fn tcgen05_mma_inline_asm_closes_selectors_and_operand_order() { + for (kind, spelling) in [ + (Tcgen05MmaKind::F16, "kind::f16"), + (Tcgen05MmaKind::Tf32, "kind::tf32"), + (Tcgen05MmaKind::F8f6f4, "kind::f8f6f4"), + (Tcgen05MmaKind::I8, "kind::i8"), + ] { + let (template, _) = tcgen05_mma_inline_asm( + Tcgen05MmaForm::WsTensor, + kind, + 1, + None, + Some(0), + Some(Tcgen05MmaBUsage::Discard), + ); + assert!(template.contains(spelling)); + } + for (usage, spelling) in [ + (Tcgen05MmaBUsage::Discard, "collector::b0::discard"), + (Tcgen05MmaBUsage::LastUse, "collector::b0::lastuse"), + (Tcgen05MmaBUsage::Fill, "collector::b0::fill"), + (Tcgen05MmaBUsage::Use, "collector::b0::use"), + ] { + let (template, _) = tcgen05_mma_inline_asm( + Tcgen05MmaForm::WsTensor, + Tcgen05MmaKind::F16, + 1, + None, + Some(0), + Some(usage), + ); + assert!(template.contains(spelling)); + } + + for (form, operands, constraints) in [ + ( + Tcgen05MmaForm::WsTensor, + "[$0], [$1], $2, $3, %enable_pred;", + "r,r,l,r,r,~{memory}", + ), + ( + Tcgen05MmaForm::WsSpTensor, + "[$0], [$1], $2, [$5], $3, %enable_pred;", + "r,r,l,r,r,r,~{memory}", + ), + ( + Tcgen05MmaForm::WsTensorZeroColMask, + "[$0], [$1], $2, $3, %enable_pred, $5;", + "r,r,l,r,r,l,~{memory}", + ), + ( + Tcgen05MmaForm::WsSpTensorZeroColMask, + "[$0], [$1], $2, [$5], $3, %enable_pred, $6;", + "r,r,l,r,r,r,l,~{memory}", + ), + ] { + let (template, actual_constraints) = tcgen05_mma_inline_asm( + form, + Tcgen05MmaKind::F16, + 1, + None, + Some(3), + Some(Tcgen05MmaBUsage::LastUse), + ); + assert!(template.contains(operands)); + assert_eq!(actual_constraints, constraints); + } +} + +#[test] +fn wgmma_controls_render_closed_compatibility_and_backend_routes() { + let catalog = catalog_with_wgmma_controls(); + validate_renderable(&catalog).unwrap(); + assert_eq!(wgmma_controls(&catalog).count(), 3); + + let compatibility = render_compat_wgmma_control(&catalog, "test-hash"); + assert!(compatibility.contains("pub fn wgmma_fence()")); + assert!(compatibility.contains("pub fn wgmma_commit_group()")); + assert!(compatibility.contains("pub fn wgmma_wait_group()")); + assert!(compatibility.contains("__wgmma_wait_group(N as u64);")); + assert!(compatibility.contains("pub(crate) fn __wgmma_wait_group(_max_pending: u64)")); + + let raw = render_raw_abi(&catalog, "test-hash").unwrap(); + assert!(raw.contains("pub unsafe fn i0317()")); + assert!(raw.contains("pub unsafe fn i0318()")); + assert!(raw.contains("pub unsafe fn i0319(_arg0: u64)")); + let raw_mod = render_raw_mod(&catalog, "test-hash"); + assert!( + raw_mod.contains("pub use crate::__cuda_oxide_intrinsic_abi_v1::i0317 as wgmma_fence;") + ); + assert!( + raw_mod + .contains("pub use crate::__cuda_oxide_intrinsic_abi_v1::i0319 as wgmma_wait_group;") + ); + + let dialect = render_dialect_wgmma_control(&catalog, "test-hash"); + for op in [ + "WgmmaFenceSyncAlignedOp", + "WgmmaCommitGroupSyncAlignedOp", + "WgmmaWaitGroupSyncAlignedOp", + ] { + assert!(dialect.contains(&format!("pub struct {op};"))); + assert!(dialect.contains(&format!("{op}::register(ctx);"))); + } + assert!(dialect.contains("operand must be i64")); + + let importer = render_importer(&catalog, "test-hash"); + assert!(importer.contains("cuda_device::wgmma::wgmma_fence")); + assert!(importer.contains("cuda_device::wgmma::wgmma_commit_group")); + assert!(importer.contains("cuda_device::wgmma::__wgmma_wait_group")); + assert!(importer.contains("wgmma_wait_group requires a compile-time constant")); + assert!(importer.contains("WgmmaWaitGroupSyncAlignedOp::build(ctx, max_pending)")); + + let lowering = render_lowering(&catalog, "test-hash"); + assert!(lowering.contains("llvm_nvvm_wgmma_fence_sync_aligned")); + assert!(lowering.contains("llvm__nvvm_dwgmma_dcommit_ugroup_dsync_daligned")); + assert!(lowering.contains("llvm__nvvm_dwgmma_dwait_ugroup_dsync_daligned")); + assert!(lowering.contains("IntrinsicBackend::LlvmNvptx")); + assert!(lowering.contains("IntrinsicBackend::LibNvvm")); + assert!(lowering.contains("wgmma.fence.sync.aligned;")); + assert!(lowering.contains("wgmma.commit_group.sync.aligned;")); + assert!(lowering.contains("wgmma.wait_group.sync.aligned $0;")); + assert!(lowering.contains("\"n,~{memory}\"")); + + let targets = render_targets(&catalog, "test-hash"); + assert!( + targets.contains("pub enum GeneratedWgmmaControlMode { Fence, CommitGroup, WaitGroup }") + ); + assert!(targets.contains("GeneratedIntrinsicVariant::WgmmaControl")); + assert!(targets.contains("WgmmaFenceSyncAlignedOp>(operation, ctx).is_some()")); + + for record in wgmma_controls(&catalog) { + let probe = render_probe(&catalog, record, "test-hash"); + assert!(probe.contains(&format!("@{}", llvm(record).symbol))); + assert!(!probe.contains("asm sideeffect")); + assert!(probe.contains("attributes #0 = { convergent }")); + } + + let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); + assert!(outputs.contains_key(&PathBuf::from( + "crates/cuda-device/src/generated/wgmma_control.rs" + ))); + assert!(outputs.contains_key(&PathBuf::from( + "crates/dialect-nvvm/src/ops/generated/wgmma_control.rs" + ))); + let reference = render_reference(&catalog, "test-hash"); + assert!(reference.contains("## WGMMA-control contracts")); +} + +#[test] +fn cluster_composites_use_generated_private_sreg_leaves() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + validate_renderable(&catalog).unwrap(); + + let leaves = render_compat_cluster_sreg(&catalog, "test-hash"); + assert_eq!(leaves.matches("pub(crate) fn __cluster_").count(), 6); + for name in [ + "__cluster_idxX", + "__cluster_idxY", + "__cluster_idxZ", + "__cluster_grid_dimX", + "__cluster_grid_dimY", + "__cluster_grid_dimZ", + ] { + assert!(leaves.contains(&format!("pub(crate) fn {name}() -> u32"))); + } + assert!(!leaves.contains("pub fn cluster_ctaidX")); + assert!(!leaves.contains("pub fn cluster_nctaidX")); + + let importer = render_importer(&catalog, "test-hash"); + assert!(importer.contains("cuda_device::cluster::__cluster_idxX")); + assert!(importer.contains("cuda_device::cluster::__cluster_grid_dimZ")); + + let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); + assert!(outputs.contains_key(&PathBuf::from( + "crates/cuda-device/src/generated/cluster_sreg.rs" + ))); +} + +#[test] +fn cluster_barrier_rendering_owns_compatibility_sync_and_both_backends() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + assert_eq!(cluster_barriers(&catalog).count(), 6); + + let dialect = render_dialect_cluster_barrier(&catalog, "test-hash"); + assert_eq!(dialect.matches("pub struct ClusterBarrierOp").count(), 1); + assert_eq!(dialect.matches("pub struct ClusterSyncOp").count(), 1); + assert!(dialect.contains("name = \"nvvm.cluster_sync\"")); + assert!(dialect.contains("ClusterSyncOp::register(ctx);")); + + let importer = render_importer(&catalog, "test-hash"); + assert!(importer.contains("\"cuda_device::cluster::cluster_sync\" =>")); + assert!(importer.contains("ClusterBarrierModeAttr::ArriveAligned")); + assert!(importer.contains("ClusterBarrierModeAttr::WaitAligned")); + assert!(!importer.contains("ClusterSyncOp")); + + let lowering = render_lowering(&catalog, "test-hash"); + assert_eq!( + lowering + .matches("impl MirToLlvmConversion for ClusterSyncOp") + .count(), + 1 + ); + assert!(lowering.contains("llvm_nvvm_barrier_cluster_arrive_aligned")); + assert!(lowering.contains("llvm_nvvm_barrier_cluster_wait_aligned")); + assert!(lowering.contains("barrier.cluster.arrive.aligned; barrier.cluster.wait.aligned;")); + + let reference = render_reference(&catalog, "test-hash"); + assert!(reference.contains("## Cluster-barrier contracts")); + assert!( + reference.contains( + "`cuda_device::cluster::cluster_sync` is the generated compatibility operation" + ) + ); + + let targets = render_targets(&catalog, "test-hash"); + assert!(targets.contains( + "GeneratedIntrinsicVariant::ClusterBarrier { mode: GeneratedClusterBarrierMode::ArriveRelaxedAligned }" + )); + assert!(targets.contains("Operation::get_op::")); + assert!(targets.contains("ClusterBarrierModeAttr::WaitAligned")); +} + +#[test] +fn cluster_memory_rendering_preserves_cluster_shared_addrspace_and_composite_load() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + validate_renderable(&catalog).unwrap(); + let records = cluster_memory(&catalog).collect::>(); + assert_eq!(records.len(), 2); + + let compatibility = render_compat_cluster_memory(&catalog, "test-hash"); + for signature in [ + "pub unsafe fn map_shared_rank(local_ptr: *const T, target_rank: u32) -> *const T", + "pub unsafe fn map_shared_rank_mut(local_ptr: *mut T, target_rank: u32) -> *mut T", + "pub unsafe fn dsmem_read_u32(local_ptr: *const u32, target_rank: u32) -> u32", + ] { + assert!(compatibility.contains(signature), "missing {signature}"); + } + assert!(compatibility.contains("cluster-shared pointer in address space 7")); + + let dialect_mod = render_dialect_mod(&catalog, "test-hash"); + assert!(dialect_mod.contains("mod cluster_memory;")); + assert!(dialect_mod.contains("pub use cluster_memory::*;")); + assert!(dialect_mod.contains("cluster_memory::register(ctx);")); + + let dialect = render_dialect_cluster_memory(&catalog, "test-hash"); + assert!(dialect.contains("pub struct MapaSharedClusterOp")); + assert!(dialect.contains("pub struct DsmemReadU32Op")); + assert!(dialect.contains("source_ptr.pointer_kind()")); + assert!(dialect.contains("address_space::CLUSTER_SHARED")); + assert!(dialect.contains("MirPointerKind::RawConst | MirPointerKind::RawMut")); + assert!(dialect.contains("result_ptr.pointer_kind() != source_ptr.pointer_kind()")); + assert!(dialect.contains( + "nvvm.mapa_shared_cluster requires a raw source pointer and must preserve its pointee, mutability, and raw kind while returning addrspace(7)" + )); + assert!(dialect.contains("nvvm.dsmem_read_u32 result must be u32")); + assert!(dialect.contains("MapaSharedClusterOp::register(ctx);")); + assert!(dialect.contains("DsmemReadU32Op::register(ctx);")); + + let importer = render_importer(&catalog, "test-hash"); + for path in [ + "cuda_device::cluster::map_shared_rank", + "cuda_device::cluster::map_shared_rank_mut", + "cuda_device::cluster::dsmem_read_u32", + ] { + assert!(importer.contains(path)); + } + assert!(importer.contains("MapaSharedClusterOp::build(ctx, source, rank)")); + assert!(importer.contains("DsmemReadU32Op::build(ctx, source, rank)")); + + let lowering = render_lowering(&catalog, "test-hash"); + assert!(lowering.contains("cast_to_shared_addrspace")); + assert!(lowering.contains("llvm_ops::IntToPtrOp::new")); + assert!(lowering.contains("llvm_types::PointerType::get(ctx, 7)")); + assert!(lowering.contains("mapa.shared::cluster.u64 $0, $1, $2;")); + assert!(lowering.contains( + "{ .reg .u64 %mapped; mapa.shared::cluster.u64 %mapped, $1, $2; ld.shared::cluster.u32 $0, [%mapped]; }" + )); + assert!(lowering.contains("\"=l,l,r\"")); + assert!(lowering.contains("\"=r,l,r,~{memory}\"")); + + let map = records + .iter() + .find(|record| record.id == "map_shared_rank") + .unwrap(); + assert_eq!(map.llvm.as_ref().unwrap().results, ["shared_cluster_ptr"]); + let map_probe = render_probe(&catalog, map, "test-hash"); + assert!(map_probe.contains("define ptr addrspace(7)")); + assert!(map_probe.contains("inttoptr i64")); + assert!(map_probe.contains("asm sideeffect")); + assert!(!map_probe.contains("~{memory}")); + assert!(map_probe.contains("attributes #0 = { convergent }")); + + let read = records + .iter() + .find(|record| record.id == "dsmem_read_u32") + .unwrap(); + assert!(read.llvm.is_none()); + let read_probe = render_probe(&catalog, read, "test-hash"); + assert!(read_probe.contains("ld.shared::cluster.u32")); + assert!(read_probe.contains("~{memory}")); + assert!(read_probe.contains("attributes #0 = { convergent }")); + + let raw = render_raw_abi(&catalog, "test-hash").unwrap(); + assert!(raw.contains("pub unsafe fn i0320(_arg0: *const u8, _arg1: u32) -> *const u8")); + assert!(raw.contains("pub unsafe fn i0321(_arg0: *const u32, _arg1: u32) -> u32")); + assert!(raw.contains( + "ordinary loads and stores compile to `ld.shared::cluster` and `st.shared::cluster`" + )); + + let reference = render_reference(&catalog, "test-hash"); + assert!(reference.contains("## Cluster-memory contracts")); + assert!(reference.contains("preserves LLVM 22's address-space-7 result")); + assert!(reference.contains("has no one-to-one LLVM intrinsic")); + + let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); + assert!(outputs.contains_key(&PathBuf::from( + "crates/cuda-device/src/generated/cluster_memory.rs" + ))); + assert!(outputs.contains_key(&PathBuf::from( + "crates/dialect-nvvm/src/ops/generated/cluster_memory.rs" + ))); + + let mut wrong_as = catalog.clone(); + wrong_as + .intrinsics + .iter_mut() + .find(|record| record.id == "map_shared_rank") + .unwrap() + .llvm + .as_mut() + .unwrap() + .results = vec!["shared_ptr".into()]; + assert!(validate_renderable(&wrong_as).is_err()); + + let mut wrong_source = catalog; + wrong_source + .intrinsics + .iter_mut() + .find(|record| record.id == "dsmem_read_u32") + .unwrap() + .source = IntrinsicSource::LlvmImported { + source_record: "invented".into(), + }; + assert!(validate_renderable(&wrong_source).is_err()); +} + +#[test] +fn tma_rendering_preserves_api_and_injects_backend_defaults() { + let catalog = catalog_with_tma(); + validate_renderable(&catalog).unwrap(); + + let compatibility = render_compat_tma(&catalog, "test-hash"); + assert!(compatibility.contains( + "pub unsafe fn cp_async_bulk_tensor_1d_g2s(dst: *mut u8, tensor_map: *const TmaDescriptor, coord0: i32, barrier: *mut Barrier)" + )); + assert!(compatibility.contains( + "pub unsafe fn cp_async_bulk_tensor_2d_g2s_multicast(dst: *mut u8, tensor_map: *const TmaDescriptor, coord0: i32, coord1: i32, barrier: *mut Barrier, cta_mask: u16)" + )); + assert!(compatibility.contains( + "pub unsafe fn cp_async_bulk_tensor_5d_s2g(src: *const u8, tensor_map: *const TmaDescriptor, coord0: i32, coord1: i32, coord2: i32, coord3: i32, coord4: i32)" + )); + assert!(compatibility.contains("pub fn cp_async_bulk_commit_group()")); + assert!(compatibility.contains("pub fn cp_async_bulk_wait_group(n: u32)")); + assert!(compatibility.contains( + "pub unsafe fn cp_async_bulk_tensor_reduce_add_tile_2d(src: *const u8, tensor_map: *const TmaDescriptor, coord0: i32, coord1: i32)" + )); + assert!(compatibility.contains( + "pub unsafe fn cp_async_bulk_tensor_reduce_xor_im2col_5d(src: *const u8, tensor_map: *const TmaDescriptor, coord0: i32, coord1: i32, coord2: i32, coord3: i32, coord4: i32)" + )); + for dimensions in 1..=5 { + assert!(compatibility.contains(&format!( + "pub unsafe fn cp_async_bulk_prefetch_tensor_{dimensions}d_l2(" + ))); + assert!(compatibility.contains(&format!( + "pub unsafe fn cp_async_bulk_prefetch_tensor_{dimensions}d_l2_cache_hint(" + ))); + } + assert!(compatibility.contains("pub unsafe fn cp_async_bulk_prefetch_tensor_gather4_2d_l2(")); + assert!( + compatibility + .contains("pub unsafe fn cp_async_bulk_prefetch_tensor_gather4_2d_l2_cache_hint(") + ); + assert!(compatibility.contains("pub unsafe fn tensormap_replace_swizzle_atomicity(")); + assert!(compatibility.contains("pub unsafe fn fence_proxy_tensormap_generic_acquire_system(")); + assert!(compatibility.contains("pub fn fence_proxy_tensormap_generic_release_system()")); + + let dialect = render_dialect_tma(&catalog, "test-hash"); + assert_eq!(dialect.matches("pub struct ").count(), 111); + assert_eq!(dialect.matches("NResultsInterface<0>").count(), 111); + assert!(dialect.contains("NOpdsInterface<10>")); + assert!(dialect.contains("CpAsyncBulkWaitGroupReadOp::register(ctx)")); + assert!(dialect.contains("CpAsyncBulkPrefetchTensorGather4TwoDimensionalL2Op::register(ctx)")); + assert!( + dialect + .contains("CpAsyncBulkPrefetchTensorGather4TwoDimensionalL2CacheHintOp::register(ctx)") + ); + assert!(dialect.contains("ReplaceTensorMapSwizzleAtomicityOp::register(ctx)")); + assert!(dialect.contains("FenceProxyTensorMapGenericReleaseSystemOp::register(ctx)")); + + let dialect_mod = render_dialect_mod(&catalog, "test-hash"); + assert!(dialect_mod.contains("mod tma;")); + assert!(dialect_mod.contains("pub use tma::*;")); + assert!(dialect_mod.contains("tma::register(ctx);")); + + let importer = render_importer(&catalog, "test-hash"); + assert!(importer.contains("cuda_device::tma::cp_async_bulk_tensor_1d_g2s")); + assert!(importer.contains("super::tma::emit_tma_g2s(")); + assert!(importer.contains("super::tma::emit_tma_g2s_multicast_cg2(")); + assert!(importer.contains("super::tma::emit_tma_s2g(")); + assert!(importer.contains("TMA wait-group count must be a compile-time constant")); + assert!(importer.contains("require_arity(name, args.len(), 1, &loc)?;")); + assert!(importer.contains("args.first(), Some(mir::Operand::Constant(_))")); + assert!(importer.contains("\"v1:i0328\"")); + assert!(importer.contains("cp_async_bulk_tensor_reduce_add_tile_2d")); + assert!( + importer.contains( + "dialect_nvvm::ops::CpAsyncBulkTensorReduceAddTile2dOp::get_concrete_op_info()" + ) + ); + + let lowering = render_lowering(&catalog, "test-hash"); + assert!(lowering.contains("CpAsyncBulkTensorG2sTile1dOp")); + assert!(lowering.contains( + "tma::{convert_control, convert_g2s, convert_g2s_multicast_cg2, convert_prefetch_tensormap, convert_prefetch_tile, convert_reduce_s2g, convert_s2g, convert_tensormap_fence, convert_tensormap_replace, PrefetchTileConfig, ReduceConfig}" + )); + assert!( + lowering + .contains("convert_g2s(ctx, rewriter, self.get_operation(), operands_info, 5, false)") + ); + assert!( + lowering.contains("convert_s2g(ctx, rewriter, self.get_operation(), operands_info, 5)") + ); + assert!(lowering.contains( + "convert_control(ctx, rewriter, self.get_operation(), operands_info, \"commit_group\", \"llvm_nvvm_cp_async_bulk_commit_group\")" + )); + assert!(lowering.contains( + "convert_reduce_s2g(ctx, rewriter, self.get_operation(), operands_info, ReduceConfig::new(2, \"add\", \"tile\", \"llvm_nvvm_cp_async_bulk_tensor_reduce_add_tile_2d\"))" + )); + assert!(lowering.contains( + "convert_prefetch_tile(ctx, rewriter, self.get_operation(), operands_info, PrefetchTileConfig::new(5, false, false, \"llvm_nvvm_cp_async_bulk_tensor_prefetch_tile_5d\"))" + )); + assert!(lowering.contains( + "convert_tensormap_replace(ctx, rewriter, self.get_operation(), operands_info, \"llvm_nvvm_tensormap_replace_global_address\", \"global_address\", \"address\", false, false)" + )); + let release_system = tma_intrinsics(&catalog) + .find(|record| record.id == "fence_proxy_tensormap_generic_release_system") + .unwrap(); + assert!(lowering.contains(&format!( + "convert_tensormap_fence(ctx, rewriter, self.get_operation(), operands_info, {:?}, false, \"sys\")", + release_system.resolved_llvm_identifier() + ))); + + let g2s = tma_intrinsics(&catalog) + .find(|record| record.id == "cp_async_bulk_tensor_2d_g2s_multicast_cg2") + .unwrap(); + let g2s_probe = render_probe(&catalog, g2s, "test-hash"); + assert!(g2s_probe.contains("i16 %cta_mask, i64 0, i1 true, i1 false, i32 2")); + assert!(g2s_probe.contains("addrspacecast ptr %dst_generic to ptr addrspace(7)")); + + let s2g = tma_intrinsics(&catalog) + .find(|record| record.id == "cp_async_bulk_tensor_2d_s2g") + .unwrap(); + let s2g_probe = render_probe(&catalog, s2g, "test-hash"); + assert!(s2g_probe.contains("ptr %tensor_map, i32 %coord0, i32 %coord1, i64 0, i1 false")); + + let reduce = tma_intrinsics(&catalog) + .find(|record| record.id == "cp_async_bulk_tensor_reduce_add_tile_2d") + .unwrap(); + let reduce_probe = render_probe(&catalog, reduce, "test-hash"); + assert!(reduce_probe.contains( + "declare void @llvm.nvvm.cp.async.bulk.tensor.reduce.add.tile.2d(ptr addrspace(3), ptr, i32, i32, i64, i1)" + )); + assert!(reduce_probe.contains( + "call void @llvm.nvvm.cp.async.bulk.tensor.reduce.add.tile.2d(ptr addrspace(3) %src, ptr %tensor_map, i32 %coord0, i32 %coord1, i64 0, i1 false)" + )); + + for stem in [ + "cp_async_bulk_prefetch_tensor_1d_l2", + "cp_async_bulk_prefetch_tensor_2d_l2", + "cp_async_bulk_prefetch_tensor_3d_l2", + "cp_async_bulk_prefetch_tensor_4d_l2", + "cp_async_bulk_prefetch_tensor_5d_l2", + "cp_async_bulk_prefetch_tensor_gather4_2d_l2", + ] { + let prefetch_plain = tma_intrinsics(&catalog) + .find(|record| record.id == stem) + .unwrap(); + let prefetch_plain_probe = render_probe(&catalog, prefetch_plain, "test-hash"); + assert!(prefetch_plain_probe.contains("i64 0, i1 false")); + assert!(!prefetch_plain_probe.contains("%cache_hint")); + + let cache_hint_id = format!("{stem}_cache_hint"); + let prefetch_cache_hint = tma_intrinsics(&catalog) + .find(|record| record.id == cache_hint_id) + .unwrap(); + let prefetch_cache_hint_probe = render_probe(&catalog, prefetch_cache_hint, "test-hash"); + assert!(prefetch_cache_hint_probe.contains("i64 %cache_hint")); + assert!(prefetch_cache_hint_probe.contains("i64 %cache_hint, i1 true")); + } + + let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); + assert!(outputs.contains_key(&PathBuf::from("crates/cuda-device/src/generated/tma.rs"))); + assert!(outputs.contains_key(&PathBuf::from( + "crates/dialect-nvvm/src/ops/generated/tma.rs" + ))); + + let mut wrong_adapter = catalog; + wrong_adapter + .intrinsics + .iter_mut() + .find(|record| record.id == "cp_async_bulk_tensor_1d_g2s") + .unwrap() + .tma + .as_mut() + .unwrap() + .adapter = TmaAdapter::NoOperands; + assert!(validate_renderable(&wrong_adapter).is_err()); +} + +#[test] +fn tcgen05_rendering_preserves_api_and_inline_ptx_routes() { + let catalog = catalog_with_tcgen05(); + validate_renderable(&catalog).unwrap(); + + let compatibility = render_compat_tcgen05(&catalog, "test-hash"); + assert!( + compatibility + .contains("pub unsafe fn tcgen05_alloc(dst_smem: *mut u32, n_cols: u32) -> ()") + ); + assert!(compatibility.contains( + "pub unsafe fn tcgen05_mma_ws_f16(d_tmem: u32, a_tmem: u32, a_desc: u64, b_desc: u64, idesc: u32, enable_d: bool) -> ()" + )); + for alias in ["e4m3", "e5m2", "e2m3", "e3m2", "e2m1"] { + assert!(compatibility.contains(&format!( + "pub unsafe fn tcgen05_mma_{alias}(d_tmem: u32, a_desc: u64, b_desc: u64, idesc: u32, enable_d: bool)" + ))); + } + assert!( + compatibility + .contains("pub unsafe fn tcgen05_ld_16x256b_x8_pure(tmem_addr: u32) -> TmemF32x32") + ); + assert!(compatibility.contains("pub fn tcgen05_load_wait() -> ()")); + assert!(compatibility.contains("pub fn tcgen05_relinquish_alloc_permit() -> ()")); + assert!(compatibility.contains("pub fn tcgen05_relinquish_alloc_permit_cg2() -> ()")); + assert!( + compatibility.contains( + "One full warp must execute this instruction uniformly with the same operands." + ) + ); + assert!(compatibility.contains( + "One full warp in each peer CTA must execute this instruction; lanes within each warp must execute uniformly with the same operands." + )); + assert!(compatibility.contains( + "pub unsafe fn tcgen05_commit_multicast_cg2(mbar: *mut u64, cta_mask: u16) -> ()" + )); + assert!(compatibility.contains( + "pub unsafe fn tcgen05_cp_128x128b_b4x16_p64(tmem_addr: u32, smem_desc: u64) -> ()" + )); + assert!(compatibility.contains( + "pub unsafe fn tcgen05_cp_64x128b_warpx2_02_13_cg2(tmem_addr: u32, smem_desc: u64) -> ()" + )); + assert!( + compatibility.contains("pub unsafe fn tcgen05_ld_16x64b_x1_raw(tmem_addr: u32) -> u32") + ); + assert!( + compatibility + .contains("pub unsafe fn tcgen05_ld_16x64b_x2_raw(tmem_addr: u32) -> CuSimd") + ); + assert!(compatibility.contains( + "pub unsafe fn tcgen05_ld_32x32b_x128_pack16(tmem_addr: u32) -> CuSimd" + )); + assert!( + compatibility + .contains("pub unsafe fn tcgen05_st_16x64b_x1_raw(tmem_addr: u32, data: u32) -> ()") + ); + assert!(compatibility.contains( + "pub unsafe fn tcgen05_st_16x64b_x2_raw(tmem_addr: u32, data: CuSimd) -> ()" + )); + assert!(compatibility.contains( + "pub unsafe fn tcgen05_st_32x32b_x128_unpack16(tmem_addr: u32, data: CuSimd) -> ()" + )); + assert!(compatibility.contains( + "pub unsafe fn tcgen05_ld_16x32bx2_x1_raw(tmem_addr: u32) -> u32" + )); + assert!(compatibility.contains( + "pub(crate) unsafe fn __tcgen05_ld_16x32bx2_x1_raw(_tmem_addr: u32, _half_split_offset: i64) -> u32" + )); + assert!( + compatibility.contains( + "unsafe { __tcgen05_ld_16x32bx2_x1_raw(tmem_addr, HALF_SPLIT_OFFSET as i64) }" + ) + ); + assert!(compatibility.contains( + "pub unsafe fn tcgen05_st_16x32bx2_x128_unpack16(tmem_addr: u32, data: CuSimd)" + )); + assert!(compatibility.contains( + "pub(crate) unsafe fn __tcgen05_st_16x32bx2_x128_unpack16(_tmem_addr: u32, _half_split_offset: i64, _data: CuSimd)" + )); + assert!( + compatibility.contains( + "pub unsafe fn tcgen05_commit_multicast(mbar: *mut u64, cta_mask: u16) -> ()" + ) + ); + assert!(compatibility.contains("pub unsafe fn tcgen05_shift_down(tmem_addr: u32) -> ()")); + assert!(compatibility.contains("pub unsafe fn tcgen05_shift_down_cg2(tmem_addr: u32) -> ()")); + assert_eq!( + compatibility + .matches("/// One thread in the CTA issues this instruction.\n") + .count(), + 4 + ); + assert_eq!( + compatibility + .matches( + "/// One thread in the CTA pair issues this instruction; the peer CTA must be active and must not have exited.\n", + ) + .count(), + 4 + ); + assert_eq!( + compatibility + .matches( + "/// The same thread that issued the tracked asynchronous tcgen05 operations must issue this commit.\n", + ) + .count(), + 6 + ); + assert_eq!( + compatibility + .matches( + "/// Completion must be tracked by a matching commit from that same thread and observed through the selected mbarrier before relying on shifted data.\n", + ) + .count(), + 2 + ); + assert_eq!( + compatibility + .matches( + "/// All tcgen05 operations in the kernel must use the same CTA-group mode.\n", + ) + .count(), + 233 + ); + assert!(compatibility.contains("/// `KIND` is 0=f16, 1=tf32, 2=f8f6f4, or 3=i8.")); + assert!(compatibility.contains( + "/// `CTA_GROUP` is 1 or 2. `COLLECTOR_A` is 0=discard, 1=lastuse, 2=fill, or 3=use." + )); + assert!(compatibility.contains( + "/// `B_BUFFER` is 0 through 3. `B_USAGE` is 0=discard, 1=lastuse, 2=fill, or 3=use." + )); + assert_eq!( + compatibility + .matches("/// `legacy_a_desc` is kept for compatibility; tensor A uses `a_tmem`.\n",) + .count(), + 5 + ); + assert_eq!( + compatibility + .matches("/// This uses kind f8f6f4, CTA group 1, and collector a::discard.\n",) + .count(), + 5 + ); + assert!(!compatibility.contains("tcgen05_mma_ws_f16_with_collector")); + + let raw = render_raw_abi(&catalog, "test-hash").unwrap(); + assert!(raw.contains("pub unsafe fn i0578(_arg0: u32, _arg1: u64) -> ()")); + assert!(raw.contains("pub unsafe fn i0611(_arg0: u32, _arg1: u64) -> ()")); + assert!(raw.contains("pub unsafe fn i0612(_arg0: u32) -> u32")); + assert!(raw.contains("pub unsafe fn i0614(_arg0: u32) -> [u32; 2]")); + assert!(raw.contains("pub unsafe fn i0669(_arg0: u32) -> [u32; 128]")); + assert!(raw.contains("pub unsafe fn i0670(_arg0: u32, _arg1: u32) -> ()")); + assert!(raw.contains("pub unsafe fn i0672(_arg0: u32, _arg1: [u32; 2]) -> ()")); + assert!(raw.contains("pub unsafe fn i0727(_arg0: u32, _arg1: [u32; 128]) -> ()")); + assert!(raw.contains("pub unsafe fn i0728(_arg0: u32, _arg1: i64) -> u32")); + assert!(raw.contains("pub unsafe fn i0743(_arg0: u32, _arg1: i64) -> [u32; 128]")); + assert!(raw.contains("pub unsafe fn i0759(_arg0: u32, _arg1: i64, _arg2: [u32; 128]) -> ()")); + assert!(raw.contains("pub unsafe fn i0760(_arg0: *mut u64, _arg1: u16) -> ()")); + assert!(raw.contains("pub unsafe fn i0761(_arg0: u32) -> ()")); + assert!(raw.contains("pub unsafe fn i0762(_arg0: u32) -> ()")); + let mma = raw.find("pub unsafe fn i0763").unwrap(); + assert!(raw[..mma].ends_with("#[allow(clippy::too_many_arguments)]\n#[inline(never)]\n")); + assert!(raw.contains("lane component must be a multiple of 32")); + assert!(raw.contains("`_arg0` must name a live tensor-memory allocation")); + assert_eq!( + raw.matches("/// One thread in the CTA issues this instruction.\n") + .count(), + 4 + ); + assert_eq!( + raw.matches( + "/// One thread in the CTA pair issues this instruction; the peer CTA must be active and must not have exited.\n", + ) + .count(), + 4 + ); + assert_eq!( + raw.matches( + "/// The same thread that issued the tracked asynchronous tcgen05 operations must issue this commit.\n", + ) + .count(), + 6 + ); + assert_eq!( + raw.matches( + "/// Completion must be tracked by a matching commit from that same thread and observed through the selected mbarrier before relying on shifted data.\n", + ) + .count(), + 2 + ); + assert_eq!( + raw.matches( + "/// All tcgen05 operations in the kernel must use the same CTA-group mode.\n", + ) + .count(), + 233 + ); + assert!(raw.contains("pub fn i0345() -> ()")); + assert!(raw.contains("pub fn i0357() -> ()")); + assert!(raw.contains("pub fn i0358() -> ()")); + assert!(raw.contains("pub fn i0361() -> ()")); + assert!(raw.contains("Complete the matching tensor-memory load wait")); + assert!(raw.contains("Complete the matching tensor-memory store wait")); + let raw_mod = render_raw_mod(&catalog, "test-hash"); + assert!(raw_mod.contains( + "pub use crate::__cuda_oxide_intrinsic_abi_v1::i0578 as tcgen05_cp_128x128b_b4x16_p64;" + )); + assert!(raw_mod.contains( + "pub use crate::__cuda_oxide_intrinsic_abi_v1::i0611 as tcgen05_cp_64x128b_warpx2_02_13_cg2;" + )); + assert!(raw_mod.contains( + "pub use crate::__cuda_oxide_intrinsic_abi_v1::i0727 as tcgen05_st_32x32b_x128_unpack16;" + )); + assert!(raw_mod.contains( + "pub use crate::__cuda_oxide_intrinsic_abi_v1::i0728 as tcgen05_ld_16x32bx2_x1_raw;" + )); + assert!(raw_mod.contains( + "pub use crate::__cuda_oxide_intrinsic_abi_v1::i0759 as tcgen05_st_16x32bx2_x128_unpack16;" + )); + assert!(raw_mod.contains( + "pub use crate::__cuda_oxide_intrinsic_abi_v1::i0760 as tcgen05_commit_multicast;" + )); + assert!( + raw_mod + .contains("pub use crate::__cuda_oxide_intrinsic_abi_v1::i0761 as tcgen05_shift_down;") + ); + assert!(raw_mod.contains( + "pub use crate::__cuda_oxide_intrinsic_abi_v1::i0762 as tcgen05_shift_down_cg2;" + )); + + let dialect = render_dialect_tcgen05(&catalog, "test-hash"); + assert_eq!(dialect.matches("pub struct Tcgen05").count(), 210); + assert_eq!(dialect.matches("impl Verify for Tcgen05").count(), 210); + assert_eq!(dialect.matches("::register(ctx)").count(), 216); + assert_eq!(dialect.matches(" Some(1),").count(), 32); + assert_eq!(dialect.matches("verifier = \"succ\"").count(), 6); + assert!(dialect.contains("Operation::get_op::")); + assert!(dialect.contains("Operation::get_op::")); + assert!(dialect.contains("&[(Tcgen05Carrier::I32, 1), (Tcgen05Carrier::I64, 1)],")); + assert!(dialect.contains( + "&[(Tcgen05Carrier::I32, 1), (Tcgen05Carrier::I64, 1), (Tcgen05Carrier::I32, 128)]," + )); + assert!(dialect.contains("pub struct Tcgen05Ld16x64bX1RawOp")); + assert!(dialect.contains("pub struct Tcgen05Ld32x32bX128Pack16Op")); + assert!(dialect.contains("pub struct Tcgen05St16x64bX1RawOp")); + assert!(dialect.contains("pub struct Tcgen05St32x32bX128Unpack16Op")); + assert!(dialect.contains("pub struct Tcgen05Ld16x32bx2X1RawOp")); + assert!(dialect.contains("pub struct Tcgen05St16x32bx2X128Unpack16Op")); + assert!(dialect.contains("pub struct Tcgen05CommitMulticastOp")); + assert!(dialect.contains("pub struct Tcgen05ShiftDownOp")); + assert!(dialect.contains("pub struct Tcgen05ShiftDownCg2Op")); + assert!(dialect.contains("NResultsInterface<128>")); + assert!(dialect.contains("NResultsInterface<32>")); + assert!(dialect.contains("NResultsInterface<4>")); + assert!(dialect.contains("NOpdsInterface<6>")); + + let dialect_mod = render_dialect_mod(&catalog, "test-hash"); + assert!(dialect_mod.contains("mod tcgen05;")); + assert!(dialect_mod.contains("pub use tcgen05::*;")); + assert!(dialect_mod.contains("tcgen05::register(ctx);")); + + let importer = render_importer(&catalog, "test-hash"); + assert!(importer.contains("cuda_device::tcgen05::tcgen05_alloc")); + assert!(importer.contains("Tcgen05AllocOp::get_concrete_op_info()")); + assert!(importer.contains("Tcgen05Ld16x256bX8PureOp::get_concrete_op_info()")); + assert!(importer.contains("MirConstructStructOp::get_concrete_op_info()")); + assert!(importer.contains("Tcgen05Ld16x64bX1RawOp::get_concrete_op_info()")); + assert!(importer.contains("Tcgen05Ld32x32bX128Pack16Op::get_concrete_op_info()")); + assert!(importer.contains("Tcgen05St16x64bX1RawOp::get_concrete_op_info()")); + assert!(importer.contains("Tcgen05St32x32bX128Unpack16Op::get_concrete_op_info()")); + assert!(importer.contains("import_generated_tcgen05_store_operands(")); + assert!(importer.contains("downcast_ref::()")); + assert!( + importer.contains("generated tcgen05 store data must contain {expected_len} u32 registers") + ); + assert!(importer.contains("destination_ty == array_result.get_type(ctx)")); + assert!(!importer.contains("super::tcgen05::emit_")); + assert!(importer.contains("\"v1:i0343\"")); + assert!(importer.contains("\"v1:i0366\"")); + assert!(importer.contains("cuda_device::tcgen05::tcgen05_cp_128x128b_b4x16_p64")); + assert!(importer.contains("Tcgen05Cp128x128bB4x16P64Op::get_concrete_op_info()")); + assert!(importer.contains("\"v1:i0578\"")); + assert!(importer.contains("\"v1:i0611\"")); + assert!(importer.contains("\"v1:i0670\"")); + assert!(importer.contains("\"v1:i0727\"")); + assert!(importer.contains("cuda_device::tcgen05::__tcgen05_ld_16x32bx2_x1_raw")); + assert!(importer.contains("cuda_device::tcgen05::__tcgen05_st_16x32bx2_x128_unpack16")); + assert!( + importer.contains("tcgen05 16x32bx2 half-split offset must be a compile-time constant") + ); + assert!(importer.contains("half-split offset must lower to a constant")); + assert!(importer.contains("args, 128, true, block_ptr")); + assert!(importer.contains("\"v1:i0728\"")); + assert!(importer.contains("\"v1:i0759\"")); + assert!(importer.contains("\"v1:i0760\"")); + assert!(importer.contains("\"v1:i0761\"")); + assert!(importer.contains("\"v1:i0762\"")); + + let lowering = render_lowering(&catalog, "test-hash"); + assert!(!lowering.contains("use crate::convert::intrinsics::tcgen05;")); + assert!(lowering.contains("impl MirToLlvmConversion for Tcgen05AllocOp")); + assert!( + lowering + .contains("tcgen05.mma.ws.cta_group::1.kind::f16 [$0], [$1], $3, $4, %enable_pred;") + ); + assert!(!lowering.contains(".kind::bf16")); + assert!(lowering.contains("convert_generated_tcgen05_load(")); + assert!(lowering.contains("tcgen05.ld.sync.aligned.16x64b.x1.b32 {$0}, [$1];")); + assert!(lowering.contains("tcgen05.ld.sync.aligned.16x64b.x1.pack::16b.b32 {$0}, [$1];")); + assert!(lowering.contains("\"=r,r,~{memory}\"")); + assert!(lowering.contains( + "convert_generated_tcgen05_load(ctx, rewriter, self.get_operation(), 1, 1, true" + )); + assert!(lowering.contains("convert_generated_tcgen05_void(")); + assert!(lowering.contains("tcgen05.st.sync.aligned.16x64b.x1.b32 [$0], {$1};")); + assert!(lowering.contains("tcgen05.st.sync.aligned.32x32b.x128.unpack::16b.b32 [$0], {$1,$2")); + assert!(lowering.contains("\"r,r,~{memory}\"")); + assert!(lowering.contains("tcgen05.ld.sync.aligned.16x32bx2.x1.b32 {$0}, [$1], $2;")); + assert!(lowering.contains("tcgen05.st.sync.aligned.16x32bx2.x1.b32 [$0], $1, {$2};")); + assert!(lowering.contains("\"=r,r,n,~{memory}\"")); + assert!(lowering.contains("\"r,n,r,~{memory}\"")); + assert!(lowering.contains( + "convert_generated_tcgen05_load(ctx, rewriter, self.get_operation(), 2, 1, true" + )); + assert!(lowering.contains("builtin::types::{FP32Type, IntegerType, Signedness}")); + assert!(lowering.contains("ops as llvm_ops")); + assert!(lowering.contains("inserter::Inserter")); + assert!(lowering.contains("tcgen05.cp.cta_group::1.128x128b.b8x16.b4x16_p64 [$0], $1;")); + assert!(lowering.contains("tcgen05.cp.cta_group::1.32x128b.warpx4 [$0], $1;")); + assert!( + lowering + .contains("tcgen05.cp.cta_group::2.64x128b.warpx2::01_23.b8x16.b6x16_p32 [$0], $1;") + ); + assert!(lowering.contains("tcgen05.cp.cta_group::1.64x128b.warpx2::02_13 [$0], $1;")); + assert!(lowering.contains("\"r,l,~{memory}\"")); + assert!(lowering.contains( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 [$0], $1;" + )); + assert!(lowering.contains("tcgen05.shift.cta_group::1.down [$0];")); + assert!(lowering.contains("tcgen05.shift.cta_group::2.down [$0];")); + assert!(lowering.contains("\"r,h,~{memory}\"")); + + let bf16 = tcgen05_intrinsics(&catalog) + .find(|record| record.id == "tcgen05_mma_ws_bf16") + .unwrap(); + let bf16_probe = render_probe(&catalog, bf16, "test-hash"); + assert!(bf16_probe.contains("route: inline PTX")); + assert!(bf16_probe.contains(".kind::f16 [$0]")); + assert!(!bf16_probe.contains(".kind::bf16")); + + let base_commit = tcgen05_intrinsics(&catalog) + .find(|record| record.id == "tcgen05_commit") + .unwrap(); + let base_commit_probe = render_probe(&catalog, base_commit, "test-hash"); + assert!(base_commit_probe.contains("mbarrier::arrive::one.b64")); + assert!(!base_commit_probe.contains("one.shared::cluster.b64")); + + let shared_commit = tcgen05_intrinsics(&catalog) + .find(|record| record.id == "tcgen05_commit_shared_cluster") + .unwrap(); + let shared_commit_probe = render_probe(&catalog, shared_commit, "test-hash"); + assert!(shared_commit_probe.contains("one.shared::cluster.b64")); + + for (id, spelling) in [ + ( + "tcgen05_commit_multicast", + "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 [$0], $1;", + ), + ( + "tcgen05_shift_down", + "tcgen05.shift.cta_group::1.down [$0];", + ), + ( + "tcgen05_shift_down_cg2", + "tcgen05.shift.cta_group::2.down [$0];", + ), + ] { + let record = tcgen05_intrinsics(&catalog) + .find(|record| record.id == id) + .unwrap(); + let probe = render_probe(&catalog, record, "test-hash"); + assert!(probe.contains(spelling)); + assert!(probe.contains("asm sideeffect")); + assert!(probe.contains("attributes #0 = { convergent }")); + } + + for (id, spelling) in [ + ( + "tcgen05_cp_128x128b_b4x16_p64", + "tcgen05.cp.cta_group::1.128x128b.b8x16.b4x16_p64", + ), + ( + "tcgen05_cp_32x128b_warpx4", + "tcgen05.cp.cta_group::1.32x128b.warpx4", + ), + ( + "tcgen05_cp_64x128b_warpx2_01_23_b6x16_p32_cg2", + "tcgen05.cp.cta_group::2.64x128b.warpx2::01_23.b8x16.b6x16_p32", + ), + ( + "tcgen05_cp_64x128b_warpx2_02_13", + "tcgen05.cp.cta_group::1.64x128b.warpx2::02_13", + ), + ] { + let record = tcgen05_intrinsics(&catalog) + .find(|record| record.id == id) + .unwrap(); + let probe = render_probe(&catalog, record, "test-hash"); + assert!(probe.contains(spelling)); + assert!(probe.contains("asm sideeffect")); + assert!(probe.contains("\"r,l,~{memory}\"")); + assert!(probe.contains("attributes #0 = { convergent }")); + } + + let scalar_load = tcgen05_intrinsics(&catalog) + .find(|record| record.id == "tcgen05_ld_16x64b_x1_raw") + .unwrap(); + let scalar_load_probe = render_probe(&catalog, scalar_load, "test-hash"); + assert!(scalar_load_probe.contains("%result = call i32 asm sideeffect")); + assert!(scalar_load_probe.contains("tcgen05.ld.sync.aligned.16x64b.x1.b32")); + assert!(scalar_load_probe.contains("\"=r,r,~{memory}\"")); + + let packed_load = tcgen05_intrinsics(&catalog) + .find(|record| record.id == "tcgen05_ld_32x32b_x128_pack16") + .unwrap(); + let packed_load_probe = render_probe(&catalog, packed_load, "test-hash"); + assert!(packed_load_probe.contains("tcgen05.ld.sync.aligned.32x32b.x128.pack::16b.b32")); + assert!(packed_load_probe.contains("{ i32, i32")); + + let scalar_store = tcgen05_intrinsics(&catalog) + .find(|record| record.id == "tcgen05_st_16x64b_x1_raw") + .unwrap(); + let scalar_store_probe = render_probe(&catalog, scalar_store, "test-hash"); + assert!(scalar_store_probe.contains("call void asm sideeffect")); + assert!(scalar_store_probe.contains("tcgen05.st.sync.aligned.16x64b.x1.b32")); + assert!(scalar_store_probe.contains("\"r,r,~{memory}\"")); + + let unpacked_store = tcgen05_intrinsics(&catalog) + .find(|record| record.id == "tcgen05_st_32x32b_x128_unpack16") + .unwrap(); + let unpacked_store_probe = render_probe(&catalog, unpacked_store, "test-hash"); + assert!(unpacked_store_probe.contains("tcgen05.st.sync.aligned.32x32b.x128.unpack::16b.b32")); + assert!(unpacked_store_probe.contains("i32 %d127")); + + let offset_load = tcgen05_intrinsics(&catalog) + .find(|record| record.id == "tcgen05_ld_16x32bx2_x1_pack16") + .unwrap(); + let offset_load_probe = render_probe(&catalog, offset_load, "test-hash"); + assert!(offset_load_probe.contains("i64 16")); + assert!(offset_load_probe.contains("[$1], $2;")); + assert!(offset_load_probe.contains("\"=r,r,n,~{memory}\"")); + + let offset_store = tcgen05_intrinsics(&catalog) + .find(|record| record.id == "tcgen05_st_16x32bx2_x1_unpack16") + .unwrap(); + let offset_store_probe = render_probe(&catalog, offset_store, "test-hash"); + assert!(offset_store_probe.contains("i64 16")); + assert!(offset_store_probe.contains("[$0], $1, {$2};")); + assert!(offset_store_probe.contains("\"r,n,r,~{memory}\"")); + + let target = render_targets(&catalog, "test-hash"); + assert!(target.contains("\"v1:i0343\"")); + assert!(target.contains("\"v1:i0578\"")); + assert!(target.contains("\"v1:i0611\"")); + assert!(target.contains("\"v1:i0727\"")); + assert!(target.contains("\"v1:i0728\"")); + assert!(target.contains("\"v1:i0759\"")); + assert!(target.contains( + "Tcgen05MmaBUsageAttr, Tcgen05MmaCollectorAAttr, Tcgen05MmaCtaGroupAttr, Tcgen05MmaFormAttr" + )); + assert!(target.contains( + "GeneratedHardwareTarget::AnyOf(&[GeneratedHardwareAlternative::ExactArchitecture(100), GeneratedHardwareAlternative::ExactArchitecture(101), GeneratedHardwareAlternative::ExactArchitecture(103), GeneratedHardwareAlternative::ExactArchitecture(110)])" + )); + assert!(target.contains( + "GeneratedHardwareTarget::AnyOf(&[GeneratedHardwareAlternative::ExactArchitecture(100), GeneratedHardwareAlternative::ExactArchitecture(103), GeneratedHardwareAlternative::ExactArchitecture(110)])" + )); + + let reference = render_reference(&catalog, "test-hash"); + assert!(reference.contains("- `tcgen05_st_16x64b_x1_raw`: runtime `unexecuted`")); + assert!(!reference.contains("- `tcgen05_st_16x64b_x1_raw`: runtime `not recorded`")); + + let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); + assert!(outputs.contains_key(&PathBuf::from( + "crates/cuda-device/src/generated/tcgen05.rs" + ))); + assert!(outputs.contains_key(&PathBuf::from( + "crates/dialect-nvvm/src/ops/generated/tcgen05.rs" + ))); + + let mut wrong_adapter = catalog.clone(); + wrong_adapter + .intrinsics + .iter_mut() + .find(|record| record.id == "tcgen05_alloc") + .unwrap() + .tcgen05 + .as_mut() + .unwrap() + .adapter = Tcgen05Adapter::NoOperands; + assert!(validate_renderable(&wrong_adapter).is_err()); + + let mut wrong_libnvvm_target = catalog.clone(); + wrong_libnvvm_target + .intrinsics + .iter_mut() + .find(|record| record.id == "tcgen05_cp_128x128b_b4x16_p64") + .unwrap() + .backend_lowerings[1] + .target + .hardware = CatalogHardwareTarget::AnyOf { + alternatives: vec![ + CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 101 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 103 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, + ], + }; + assert!(validate_renderable(&wrong_libnvvm_target).is_err()); + + let mut wrong_load_selector = catalog.clone(); + wrong_load_selector + .intrinsics + .iter_mut() + .find(|record| record.id == "tcgen05_ld_16x64b_x1_pack16") + .unwrap() + .expected_ptx + .modifiers + .remove(5); + assert!(validate_renderable(&wrong_load_selector).is_err()); + + let mut wrong_store_selector = catalog.clone(); + wrong_store_selector + .intrinsics + .iter_mut() + .find(|record| record.id == "tcgen05_st_16x64b_x1_unpack16") + .unwrap() + .expected_ptx + .modifiers + .remove(5); + assert!(validate_renderable(&wrong_store_selector).is_err()); + + let mut wrong_offset_scope = catalog.clone(); + wrong_offset_scope + .intrinsics + .iter_mut() + .find(|record| record.id == "tcgen05_ld_16x32bx2_x1_raw") + .unwrap() + .semantics + .execution_scope = "thread".into(); + assert!(validate_renderable(&wrong_offset_scope).is_err()); + + let mut wrong_copy_spelling = catalog; + wrong_copy_spelling + .intrinsics + .iter_mut() + .find(|record| record.id == "tcgen05_cp_128x128b_b4x16_p64") + .unwrap() + .expected_ptx + .modifiers + .remove(3); + assert!(validate_renderable(&wrong_copy_spelling).is_err()); +} diff --git a/crates/cuda-intrinsics-gen/src/render/tests/warp_memory.rs b/crates/cuda-intrinsics-gen/src/render/tests/warp_memory.rs new file mode 100644 index 0000000000..964abc82db --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/render/tests/warp_memory.rs @@ -0,0 +1,964 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use super::*; + +use crate::model::{ + BackendLoweringMechanism, IntrinsicSource, MbarrierBasicOperation, ReduxAdapter, + WarpBarrierAdapter, WarpShuffleMode, WarpShuffleValueKind, +}; +use crate::render::common::{intrinsic_marker, llvm}; +use crate::render::families::{ + active_masks, cp_async_controls, cp_async_copies, cp_async_mbarriers, mbarrier_basics, redux, + vote_intrinsics, warp_barriers, warp_matches, warp_shuffles, +}; +use std::collections::BTreeSet; +use std::path::Path; + +#[test] +fn cp_async_rendering_preserves_compatibility_dispatch_and_backend_routes() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + validate_renderable(&catalog).unwrap(); + assert_eq!(cp_async_copies(&catalog).count(), 8); + assert_eq!(cp_async_controls(&catalog).count(), 3); + assert_eq!(cp_async_mbarriers(&catalog).count(), 4); + + let compatibility = render_compat_cp_async_copy(&catalog, "test-hash"); + for signature in [ + "pub unsafe fn cp_async_ca_4(_shared_dst: *mut u32, _global_src: *const u32)", + "pub unsafe fn cp_async_ca_8(_shared_dst: *mut u32, _global_src: *const u32)", + "pub unsafe fn cp_async_ca_16(_shared_dst: *mut u32, _global_src: *const u32)", + "pub unsafe fn cp_async_ca_zfill_4(_shared_dst: *mut u32, _global_src: *const u8, _src_size: u32)", + "pub unsafe fn cp_async_ca_zfill_8(_shared_dst: *mut u32, _global_src: *const u8, _src_size: u32)", + "pub unsafe fn cp_async_ca_zfill_16(_shared_dst: *mut u32, _global_src: *const u8, _src_size: u32)", + "pub unsafe fn cp_async_cg_16(_shared_dst: *mut u32, _global_src: *const u32)", + "pub unsafe fn cp_async_cg_zfill_16(_shared_dst: *mut u32, _global_src: *const u8, _src_size: u32)", + "pub unsafe fn cp_async_commit_group()", + "pub unsafe fn cp_async_wait_all()", + "pub unsafe fn cp_async_wait_group(_max_pending: u32)", + "pub unsafe fn cp_async_mbarrier_arrive(_barrier: *mut crate::barrier::Barrier)", + "pub unsafe fn cp_async_mbarrier_arrive_shared(_barrier: *mut crate::barrier::Barrier)", + "pub unsafe fn cp_async_mbarrier_arrive_noinc(_barrier: *mut crate::barrier::Barrier)", + "pub unsafe fn cp_async_mbarrier_arrive_noinc_shared(_barrier: *mut crate::barrier::Barrier)", + ] { + assert!(compatibility.contains(signature)); + } + + let dialect = render_dialect_cp_async_copy(&catalog, "test-hash"); + let importer = render_importer(&catalog, "test-hash"); + let lowering = render_lowering(&catalog, "test-hash"); + let targets = render_targets(&catalog, "test-hash"); + for record in cp_async_copies(&catalog) + .chain(cp_async_controls(&catalog)) + .chain(cp_async_mbarriers(&catalog)) + { + assert!(dialect.contains(&format!("pub struct {}", record.dialect.op_type))); + assert!(dialect.contains(&format!("{}::register(ctx)", record.dialect.op_type))); + assert!(importer.contains(&record.rust.canonical_path)); + for path in &record.rust.compatibility_paths { + assert!(importer.contains(path)); + } + assert!(importer.contains(&intrinsic_marker(&catalog, record))); + assert!(lowering.contains(&format!( + "impl MirToLlvmConversion for {}", + record.dialect.op_type + ))); + assert!(lowering.contains(&record.llvm_identifier())); + assert!( + record + .backend_lowerings + .iter() + .any(|entry| entry.backend == IntrinsicBackend::LlvmNvptx) + ); + assert!( + record + .backend_lowerings + .iter() + .any(|entry| entry.backend == IntrinsicBackend::LibNvvm) + ); + assert!(targets.contains(&format!("id: {:?}", record.id))); + } + assert_eq!(dialect.matches("::register(ctx);").count(), 15); + assert!(lowering.contains("convert_generated_cp_async_copy")); + assert!(lowering.contains("convert_generated_cp_async_control")); + assert!(lowering.contains("convert_generated_cp_async_mbarrier")); +} + +#[test] +fn cp_async_mbarrier_rendering_preserves_counting_and_state_space_routes() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + let records: Vec<_> = cp_async_mbarriers(&catalog).collect(); + assert_eq!(records.len(), 4); + + let raw = render_raw_abi(&catalog, "test-hash").unwrap(); + assert!(raw.contains("pub unsafe fn i0101(_arg0: *mut u64)")); + assert!(raw.contains("live, initialized, eight-byte-aligned mbarrier")); + assert!(raw.contains("initial pending count must already include")); + + let compatibility = render_compat_cp_async_copy(&catalog, "test-hash"); + assert_eq!( + compatibility + .matches("That increment must not exceed the barrier's pending-count limit.") + .count(), + 2 + ); + + let dialect = render_dialect_cp_async_copy(&catalog, "test-hash"); + assert!(dialect.contains("mutable generic/shared pointer to u64")); + let importer = render_importer(&catalog, "test-hash"); + let lowering = render_lowering(&catalog, "test-hash"); + for record in &records { + assert!(importer.contains(&format!( + "let bridge = {}::build(ctx, barrier)", + record.dialect.op_type + ))); + assert!(lowering.contains(&record.llvm_identifier())); + let routes: BTreeSet<_> = record + .backend_lowerings + .iter() + .map(|route| (route.backend, route.mechanism)) + .collect(); + assert_eq!( + routes, + BTreeSet::from([ + ( + IntrinsicBackend::LlvmNvptx, + BackendLoweringMechanism::TypedNvvm, + ), + ( + IntrinsicBackend::LibNvvm, + BackendLoweringMechanism::InlinePtx, + ), + ]) + ); + } + assert!(lowering.contains("\"arrive\", \"generic\"")); + assert!(lowering.contains("\"arrive\", \"shared\"")); + assert!(lowering.contains("\"arrive_no_inc\", \"generic\"")); + assert!(lowering.contains("\"arrive_no_inc\", \"shared\"")); + + let generic = records + .iter() + .find(|record| record.id == "cp_async_mbarrier_arrive") + .unwrap(); + let generic_probe = render_probe(&catalog, generic, "test-hash"); + assert!(generic_probe.contains("declare void @llvm.nvvm.cp.async.mbarrier.arrive(ptr)")); + assert!(!generic_probe.contains("addrspacecast")); + let shared = records + .iter() + .find(|record| record.id == "cp_async_mbarrier_arrive_shared") + .unwrap(); + let shared_probe = render_probe(&catalog, shared, "test-hash"); + assert!(shared_probe.contains("ptr addrspace(3)")); + assert!(shared_probe.contains("addrspacecast ptr %barrier_generic")); + + let reference = render_reference(&catalog, "test-hash"); + assert!(reference.contains("## cp.async mbarrier contracts")); + assert!(reference.contains("cp_async_mbarrier_arrive_noinc`: runtime `unexecuted`")); + + let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); + assert!(outputs.contains_key(&PathBuf::from( + "crates/cuda-device/src/generated/async_copy.rs" + ))); + assert!(outputs.contains_key(&PathBuf::from( + "crates/dialect-nvvm/src/ops/generated/cp_async.rs" + ))); + assert!(!outputs.contains_key(&PathBuf::from( + "crates/dialect-nvvm/src/ops/generated/cp_async_mbarrier.rs" + ))); +} + +#[test] +fn redux_rendering_preserves_mask_first_api_and_source_first_llvm_order() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + assert_eq!(catalog.schema, crate::resolve::CATALOG_SCHEMA); + validate_renderable(&catalog).unwrap(); + + assert_eq!(redux(&catalog).count(), 16); + let record = redux(&catalog).next().unwrap(); + assert_eq!( + record.redux.as_ref().unwrap().adapter, + ReduxAdapter::MaskValueToSourceMemberMask + ); + + let dialect = render_dialect_redux(&catalog, "test-hash"); + assert!(dialect.contains("name = \"nvvm.redux_sync_add\"")); + assert!(dialect.contains("name = \"nvvm.redux_sync_min\"")); + assert!(dialect.contains("vec![member_mask, value]")); + assert!(dialect.contains("ReduxSyncAddOp::register(ctx)")); + assert!(dialect.contains("Signedness::Signed")); + + assert!(dialect.contains("name = \"nvvm.redux_sync_fmin\"")); + assert!(dialect.contains("types::{FP32Type, IntegerType, Signedness}")); + assert!(dialect.contains("let result_ty = FP32Type::get(ctx);")); + assert!(dialect.contains("is_f32(ctx, op.get_operand(1).get_type(ctx))")); + assert!(dialect.contains("is_f32(ctx, op.get_result(0).get_type(ctx))")); + + let importer = render_importer(&catalog, "test-hash"); + assert!(importer.contains("cuda_device::warp::redux_sync_add")); + assert!(importer.contains("let (member_mask, last_op)")); + assert!(importer.contains("let reduction = ReduxSyncAddOp::build(ctx, member_mask, value)")); + assert!(importer.contains("set_generated_intrinsic_marker(ctx, reduction, \"v1:i0017\")")); + + let lowering = render_lowering(&catalog, "test-hash"); + assert!(lowering.contains("impl MirToLlvmConversion for ReduxSyncAddOp")); + assert!(lowering.contains("convert_redux(ctx, rewriter, self.get_operation(), operands_info")); + assert!(lowering.contains("\"llvm_nvvm_redux_sync_add\"")); + + let probe = render_probe(&catalog, record, "test-hash"); + assert!(probe.contains("define i32 @probe_redux_sync_add(i32 %member_mask, i32 %value)")); + assert!(probe.contains("call i32 @llvm.nvvm.redux.sync.add(i32 %value, i32 %member_mask)")); + + let f32_record = redux(&catalog) + .find(|record| record.id == "redux_sync_min_f32") + .unwrap(); + let f32_probe = render_probe(&catalog, f32_record, "test-hash"); + assert!( + f32_probe + .contains("define float @probe_redux_sync_min_f32(i32 %member_mask, float %value)") + ); + assert!( + f32_probe.contains("call float @llvm.nvvm.redux.sync.fmin(float %value, i32 %member_mask)") + ); + + let raw = render_raw_abi(&catalog, "test-hash").unwrap(); + let signature = "pub unsafe fn i0017(_arg0: u32, _arg1: u32) -> u32"; + let index = raw.find(signature).unwrap(); + assert!(raw[..index].ends_with("#[must_use]\n#[inline(never)]\n")); + assert!(raw.contains("pub unsafe fn i0019(_arg0: u32, _arg1: i32) -> i32")); + assert!(raw.contains("pub unsafe fn i0024(_arg0: u32, _arg1: u32) -> u32")); + assert!(raw.contains("pub unsafe fn i1003(_arg0: u32, _arg1: f32) -> f32")); + assert!(raw.contains("The executing lane must be named in `mask`")); +} + +#[test] +fn vote_rendering_keeps_types_selection_pairs_and_raw_only_uni() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + validate_renderable(&catalog).unwrap(); + assert_eq!(vote_intrinsics(&catalog).count(), 4); + + let dialect = render_dialect_vote(&catalog, "test-hash"); + for op in [ + "VoteSyncAllOp", + "VoteSyncAnyOp", + "VoteSyncBallotOp", + "VoteSyncUniOp", + ] { + assert!(dialect.contains(&format!("pub struct {op}"))); + assert!(dialect.contains(&format!("{op}::register(ctx)"))); + } + assert!(dialect.contains("requires i32 member mask, i1 predicate, and i1 result")); + assert!(dialect.contains("requires i32 member mask, i1 predicate, and i32 result")); + + let importer = render_importer(&catalog, "test-hash"); + assert!(importer.contains("cuda_device::warp::all_sync")); + assert!(!importer.contains("cuda_device::warp::uni_sync")); + assert!(importer.contains("let vote = VoteSyncUniOp::build(ctx, member_mask, predicate)")); + assert!(importer.contains("set_generated_intrinsic_marker(ctx, vote, \"v1:i0043\")")); + + let lowering = render_lowering(&catalog, "test-hash"); + assert!(lowering.contains("impl MirToLlvmConversion for VoteSyncUniOp")); + assert!(lowering.contains("\"llvm_nvvm_vote_uni_sync\"")); + assert!(lowering.contains("convert_vote(ctx, rewriter, self.get_operation(), operands_info")); + + let record = vote_intrinsics(&catalog) + .find(|record| record.id == "ballot_sync") + .unwrap(); + assert_eq!(record.selections.len(), 2); + let probe = render_probe(&catalog, record, "test-hash"); + assert!(probe.contains("define i32 @probe_ballot_sync(i32 %member_mask, i1 %predicate)")); + assert!(probe.contains("define i32 @probe_ballot_sync_immediate(i1 %predicate)")); + assert!(probe.contains("i32 -1, i1 %predicate")); + + let raw = render_raw_abi(&catalog, "test-hash").unwrap(); + for abi_id in ["i0040", "i0041", "i0042", "i0043"] { + assert!(raw.contains(&format!("pub unsafe fn {abi_id}"))); + } + assert!(raw.contains("Every non-exited lane named in `mask`")); +} + +#[test] +fn active_mask_and_warp_match_rendering_preserves_backend_and_abi_contracts() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + validate_renderable(&catalog).unwrap(); + assert_eq!(active_masks(&catalog).count(), 1); + assert_eq!(warp_matches(&catalog).count(), 4); + + let dialect = render_dialect_active_mask(&catalog, "test-hash"); + assert!(dialect.contains("pub struct ActiveMaskOp")); + assert!(dialect.contains("NOpdsInterface<0>, NResultsInterface<1>")); + + let match_dialect = render_dialect_warp_match(&catalog, "test-hash"); + for op in [ + "MatchAnySyncI32Op", + "MatchAnySyncI64Op", + "MatchAllSyncI32Op", + "MatchAllSyncI64Op", + ] { + assert!(match_dialect.contains(&format!("pub struct {op}"))); + assert!(match_dialect.contains(&format!("{op}::register(ctx)"))); + } + + let lowering = render_lowering(&catalog, "test-hash"); + assert!(lowering.contains("IntrinsicBackend::LlvmNvptx => {")); + assert!(lowering.contains("convert_active_mask(ctx, rewriter, op, operands_info)")); + assert!(lowering.contains("IntrinsicBackend::LibNvvm =>")); + assert!(lowering.contains("\"activemask.b32 $0;\"")); + assert!(lowering.contains("\"=r,~{memory}\"")); + assert!(lowering.contains("convert_match_any(")); + assert!(lowering.contains("convert_match_all(")); + assert!(lowering.contains("\"llvm_nvvm_match_all_sync_i64p\"")); + + let match_all = warp_matches(&catalog) + .find(|record| record.id == "match_all_sync") + .unwrap(); + let probe = render_probe(&catalog, match_all, "test-hash"); + for suffix in ["rr", "ri", "ir", "ii"] { + assert!(probe.contains(&format!("@probe_match_all_sync_{suffix}"))); + } + assert!(probe.contains("declare { i32, i1 } @llvm.nvvm.match.all.sync.i32p")); + + let raw = render_raw_abi(&catalog, "test-hash").unwrap(); + assert!(raw.contains("pub fn i0044() -> u32")); + assert!(!raw.contains("pub unsafe fn i0044() -> u32")); + assert!(raw.contains("pub unsafe fn i0045(_arg0: u32, _arg1: u32) -> u32")); + assert!(raw.contains("pub unsafe fn i0046(_arg0: u32, _arg1: u64) -> u32")); + assert!(raw.contains("pub unsafe fn i0047(_arg0: u32, _arg1: u32) -> u32")); + assert!(raw.contains("pub unsafe fn i0048(_arg0: u32, _arg1: u64) -> u32")); +} + +#[test] +fn warp_barrier_rendering_preserves_mask_and_void_contracts() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + validate_renderable(&catalog).unwrap(); + assert_eq!(warp_barriers(&catalog).count(), 1); + + let record = warp_barriers(&catalog).next().unwrap(); + assert_eq!(record.id, "sync_mask"); + assert_eq!( + record.warp_barrier.as_ref().unwrap().adapter, + WarpBarrierAdapter::DirectMemberMask + ); + + let dialect_mod = render_dialect_mod(&catalog, "test-hash"); + assert!(dialect_mod.contains("mod warp_barrier;")); + assert!(dialect_mod.contains("warp_barrier::register(ctx)")); + + let dialect = render_dialect_warp_barrier(&catalog, "test-hash"); + assert!(dialect.contains("pub struct BarWarpSyncOp")); + assert!(dialect.contains("NOpdsInterface<1>, NResultsInterface<0>")); + assert!(dialect.contains("vec![member_mask]")); + assert!(dialect.contains("op.get_num_operands() != 1 || op.get_num_results() != 0")); + assert!(dialect.contains("if !is_i32(ctx, op.get_operand(0).get_type(ctx))")); + assert!(dialect.contains("requires exactly one member-mask operand and no results")); + assert!(dialect.contains("member mask must be i32")); + assert!(dialect.contains("BarWarpSyncOp::register(ctx)")); + + let importer = render_importer(&catalog, "test-hash"); + assert!(importer.contains("cuda_device::warp::sync_mask")); + assert!(importer.contains("let barrier = BarWarpSyncOp::build(ctx, member_mask)")); + assert!(importer.contains("set_generated_intrinsic_marker(ctx, barrier, \"v1:i0049\")")); + assert!(importer.contains("helpers::emit_goto(ctx, *target_idx, barrier")); + + let lowering = render_lowering(&catalog, "test-hash"); + assert!(lowering.contains("impl MirToLlvmConversion for BarWarpSyncOp")); + assert!( + lowering + .contains("convert_bar_warp_sync(ctx, rewriter, self.get_operation(), operands_info)") + ); + + let probe = render_probe(&catalog, record, "test-hash"); + assert!(probe.contains("declare void @llvm.nvvm.bar.warp.sync(i32)")); + assert!(probe.contains("define void @probe_sync_mask(i32 %member_mask)")); + assert!(probe.contains("call void @llvm.nvvm.bar.warp.sync(i32 %member_mask)")); + assert!(probe.contains("define void @probe_sync_mask_immediate()")); + assert!(probe.contains("call void @llvm.nvvm.bar.warp.sync(i32 -1)")); + + let raw = render_raw_abi(&catalog, "test-hash").unwrap(); + assert!(raw.contains("pub unsafe fn i0049(_arg0: u32) -> ()")); + assert!(raw.contains("On `sm_6x` and earlier")); + assert!(raw.contains("no lane outside `mask` may be active")); + assert!(raw.contains("The barrier orders memory accesses among participating lanes")); + + let reference = render_reference(&catalog, "test-hash"); + assert!(reference.contains("## Warp-barrier contract")); + assert!(reference.contains("no unnamed lane may be active")); + assert!(reference.contains("Both immediate and register masks are admitted")); +} + +#[test] +fn warp_shuffle_rendering_owns_all_i32_f32_and_i64_modes() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + validate_renderable(&catalog).unwrap(); + assert_eq!(warp_shuffles(&catalog).count(), 12); + + let dialect_mod = render_dialect_mod(&catalog, "test-hash"); + assert!(dialect_mod.contains("mod warp_shuffle;")); + assert!(dialect_mod.contains("warp_shuffle::register(ctx)")); + + let dialect = render_dialect_warp_shuffle(&catalog, "test-hash"); + for op in [ + "ShflSyncIdxI32Op", + "ShflSyncBflyI32Op", + "ShflSyncDownI32Op", + "ShflSyncUpI32Op", + "ShflSyncIdxF32Op", + "ShflSyncBflyF32Op", + "ShflSyncDownF32Op", + "ShflSyncUpF32Op", + "ShflSyncIdxI64Op", + "ShflSyncBflyI64Op", + "ShflSyncDownI64Op", + "ShflSyncUpI64Op", + ] { + assert!(dialect.contains(&format!("pub struct {op}"))); + assert!(dialect.contains(&format!("{op}::register(ctx)"))); + } + assert!(dialect.contains("vec![member_mask, value, lane_or_delta]")); + assert!(dialect.contains("requires i32 mask/lane and i32 value/result")); + assert!(dialect.contains("requires i32 mask/lane and f32 value/result")); + assert!(dialect.contains("requires i32 mask/lane and i64 value/result")); + assert!(dialect.contains("fn is_i64")); + assert!(dialect.contains("integer.width() == 64")); + assert!(dialect.contains("IntegerType::get(ctx, 64, Signedness::Unsigned)")); + + let importer = render_importer(&catalog, "test-hash"); + assert!(importer.contains("cuda_device::warp::shuffle_sync")); + assert!(importer.contains("cuda_device::warp::shuffle_up_f32_sync")); + assert!( + importer.contains( + "let shuffle = ShflSyncIdxI32Op::build(ctx, member_mask, value, lane_or_delta)" + ) + ); + assert!(importer.contains("set_generated_intrinsic_marker(ctx, shuffle, \"v1:i0050\")")); + for (name, op, marker) in [ + ("shuffle_u64_sync", "ShflSyncIdxI64Op", "v1:i0058"), + ("shuffle_xor_u64_sync", "ShflSyncBflyI64Op", "v1:i0059"), + ("shuffle_down_u64_sync", "ShflSyncDownI64Op", "v1:i0060"), + ("shuffle_up_u64_sync", "ShflSyncUpI64Op", "v1:i0061"), + ] { + assert!(importer.contains(&format!( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::{}", + marker.strip_prefix("v1:").unwrap() + ))); + assert!(importer.contains(&format!("cuda_device::warp::{name}"))); + assert!(importer.contains(&format!( + "let shuffle = {op}::build(ctx, member_mask, value, lane_or_delta)" + ))); + assert!(importer.contains(&format!( + "set_generated_intrinsic_marker(ctx, shuffle, {marker:?})" + ))); + } + + let lowering = render_lowering(&catalog, "test-hash"); + assert!(lowering.contains("impl MirToLlvmConversion for ShflSyncIdxI32Op")); + assert!(lowering.contains("convert_shuffle_i32(ctx, rewriter")); + assert!(lowering.contains("\"llvm_nvvm_shfl_sync_idx_i32\", 31)")); + assert!(lowering.contains("impl MirToLlvmConversion for ShflSyncUpF32Op")); + assert!(lowering.contains("convert_shuffle_f32(ctx, rewriter")); + assert!(lowering.contains("\"llvm_nvvm_shfl_sync_up_f32\", 0)")); + + for (op, mode, clamp) in [ + ("ShflSyncIdxI64Op", "idx", 31), + ("ShflSyncBflyI64Op", "bfly", 31), + ("ShflSyncDownI64Op", "down", 31), + ("ShflSyncUpI64Op", "up", 0), + ] { + assert!(lowering.contains(&format!("impl MirToLlvmConversion for {op}"))); + assert!(lowering.contains(&format!( + "convert_shuffle_i64(ctx, rewriter, self.get_operation(), operands_info, {mode:?}, {clamp})" + ))); + } + assert!(!lowering.contains("llvm_nvvm_shfl_sync_idx_i64")); + assert!(!lowering.contains("llvm_nvvm_shfl_sync_bfly_i64")); + assert!(!lowering.contains("llvm_nvvm_shfl_sync_down_i64")); + assert!(!lowering.contains("llvm_nvvm_shfl_sync_up_i64")); + + for record in warp_shuffles(&catalog) { + let probe = render_probe(&catalog, record, "test-hash"); + let shuffle = record.warp_shuffle.as_ref().unwrap(); + if shuffle.value_kind == WarpShuffleValueKind::I64 { + let mode = match shuffle.mode { + WarpShuffleMode::Idx => "idx", + WarpShuffleMode::Bfly => "bfly", + WarpShuffleMode::Down => "down", + WarpShuffleMode::Up => "up", + }; + let asm = format!( + "{{ .reg .b32 lo; .reg .b32 hi; mov.b64 {{lo, hi}}, $1; shfl.sync.{mode}.b32 lo, lo, $2, {}, $3; shfl.sync.{mode}.b32 hi, hi, $2, {}, $3; mov.b64 $0, {{lo, hi}}; }}", + shuffle.clamp, shuffle.clamp + ); + assert!(probe.contains(&format!( + "define i64 @probe_{}(i32 %member_mask, i64 %value, i32 %lane) #0", + record.id + ))); + assert!(probe.contains(&format!( + "call i64 asm sideeffect {asm:?}, \"=l,l,r,r\"(i64 %value, i32 %lane, i32 %member_mask) #0" + ))); + assert_eq!(probe.matches("asm sideeffect").count(), 1); + assert_eq!(probe.matches("attributes #0 = { convergent }").count(), 1); + assert!(!probe.contains("declare i64 @llvm.nvvm.shfl")); + for suffix in ["rr", "ri", "ir", "ii"] { + assert!(!probe.contains(&format!("@probe_{}_{suffix}", record.id))); + } + } else { + for suffix in ["rr", "ri", "ir", "ii"] { + assert!(probe.contains(&format!("@probe_{}_{suffix}", record.id))); + } + assert!(probe.contains(&format!(", i32 {})", shuffle.clamp))); + } + } + + let raw = render_raw_abi(&catalog, "test-hash").unwrap(); + assert!(raw.contains("pub unsafe fn i0050(_arg0: u32, _arg1: u32, _arg2: u32) -> u32")); + assert!(raw.contains("pub unsafe fn i0057(_arg0: u32, _arg1: f32, _arg2: u32) -> f32")); + for abi_id in ["i0058", "i0059", "i0060", "i0061"] { + assert!(raw.contains(&format!( + "pub unsafe fn {abi_id}(_arg0: u32, _arg1: u64, _arg2: u32) -> u64" + ))); + } + assert!(raw.contains("If the computed source lane is in range")); + assert!(raw.contains( + "If PTX marks the computed source out of range, the calling lane's input is copied" + )); + assert!(raw.contains("two `b32` shuffles in one convergent block")); + + let reference = render_reference(&catalog, "test-hash"); + assert!(reference.contains("## Warp-shuffle contracts")); + assert!(reference.contains("inserts clamp `31` during lowering")); + assert!(reference.contains("inserts clamp `0` during lowering")); + assert!(reference.contains("One convergent, side-effecting inline-PTX block splits `i64`")); + assert!(reference.contains("PTX-native source; no LLVM record")); + + let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); + assert!(outputs.contains_key(&PathBuf::from( + "crates/dialect-nvvm/src/ops/generated/warp_shuffle.rs" + ))); +} + +#[test] +fn sync_rendering_keeps_barrier_and_threadfence_contracts_explicit() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + validate_renderable(&catalog).unwrap(); + let record = sync_intrinsics(&catalog).next().unwrap(); + assert_eq!(sync_intrinsics(&catalog).count(), 4); + + let raw = render_raw_abi(&catalog, "test-hash").unwrap(); + assert!(raw.contains("pub unsafe fn i0034() -> ()")); + assert!(raw.contains("Every active thread in the CTA must reach the same barrier")); + + let dialect = render_dialect_sync(&catalog, "test-hash"); + assert!(dialect.contains("pub struct Barrier0Op")); + assert!(dialect.contains("NOpdsInterface<0>, NResultsInterface<0>")); + assert!(dialect.contains("Barrier0Op::register(ctx)")); + + let importer = render_importer(&catalog, "test-hash"); + assert!(importer.contains("cuda_device::thread::sync_threads")); + assert!(importer.contains("cuda_device::sync_threads")); + assert!(importer.contains("Barrier0Op::get_concrete_op_info()")); + assert!(importer.contains("set_generated_intrinsic_marker(ctx, barrier, \"v1:i0034\")")); + assert!(importer.contains("helpers::emit_goto(ctx, *target_idx, barrier")); + + let lowering = render_lowering(&catalog, "test-hash"); + assert!(lowering.contains("impl MirToLlvmConversion for Barrier0Op")); + assert!(lowering.contains("create_i32_const(ctx, rewriter, 0)")); + assert!(lowering.contains("\"llvm_nvvm_barrier_cta_sync_aligned_all\"")); + assert!(lowering.contains("IntrinsicBackend::LlvmNvptx")); + assert!(lowering.contains("IntrinsicBackend::LibNvvm")); + assert!(lowering.contains("\"bar.sync 0;\", \"~{memory}\"")); + + let probe = render_probe(&catalog, record, "test-hash"); + assert!(probe.contains("declare void @llvm.nvvm.barrier.cta.sync.aligned.all(i32)")); + assert!(probe.contains("call void @llvm.nvvm.barrier.cta.sync.aligned.all(i32 0)")); + + let targets = render_targets(&catalog, "test-hash"); + assert!(targets.contains("id: \"sync_threads\", abi_id: \"i0034\"")); + assert!(targets.contains("source_record: \"BARRIER_CTA_SYNC_ALIGNED_ALL_i\"")); + assert!(targets.contains( + "backend: GeneratedIntrinsicBackend::LlvmNvptx, requirement: GeneratedTargetRequirement { minimum_ptx: GeneratedPtxVersion::from_encoded(32), hardware: GeneratedHardwareTarget::AnyOf(&[GeneratedHardwareAlternative::MinimumSm(20)]) }" + )); + assert!(targets.contains( + "backend: GeneratedIntrinsicBackend::LibNvvm, requirement: GeneratedTargetRequirement { minimum_ptx: GeneratedPtxVersion::from_encoded(10), hardware: GeneratedHardwareTarget::AnyOf(&[GeneratedHardwareAlternative::MinimumSm(75)]) }" + )); + assert!( + !record + .selections + .iter() + .any(|selection| selection.source_record.ends_with("_r")) + ); + + let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); + assert!(outputs.contains_key(&PathBuf::from( + "crates/dialect-nvvm/src/ops/generated/sync.rs" + ))); + assert!(outputs.contains_key(&PathBuf::from("crates/cuda-device/src/generated/fence.rs"))); + + let compatibility = render_compat_fence(&catalog, "test-hash"); + for (id, abi_id, op_type, op_name, llvm_identifier, ptx) in [ + ( + "threadfence_block", + "i0298", + "ThreadfenceBlockOp", + "nvvm.threadfence_block", + "llvm_nvvm_membar_cta", + "membar.cta;", + ), + ( + "threadfence", + "i0299", + "ThreadfenceOp", + "nvvm.threadfence", + "llvm_nvvm_membar_gl", + "membar.gl;", + ), + ( + "threadfence_system", + "i0300", + "ThreadfenceSystemOp", + "nvvm.threadfence_system", + "llvm_nvvm_membar_sys", + "membar.sys;", + ), + ] { + let record = sync_intrinsics(&catalog) + .find(|record| record.id == id) + .unwrap(); + assert_eq!(record.rust.abi_id, abi_id); + assert!(compatibility.contains(&format!("pub fn {id}()"))); + assert!(dialect.contains(&format!("pub struct {op_type}"))); + assert!(dialect.contains(&format!("name = \"{op_name}\""))); + assert!(dialect.contains(&format!("{op_type}::register(ctx)"))); + assert!(importer.contains(&format!("cuda_device::fence::{id}"))); + assert!(importer.contains(&format!("cuda_device::{id}"))); + assert!(importer.contains(&format!( + "set_generated_intrinsic_marker(ctx, barrier, \"v1:{abi_id}\")" + ))); + assert!(lowering.contains(&format!("impl MirToLlvmConversion for {op_type}"))); + assert!(lowering.contains(&format!("\"{llvm_identifier}\""))); + + let probe = render_probe(&catalog, record, "test-hash"); + assert!(probe.contains(&format!("declare void @{}()", llvm(record).symbol))); + assert!(probe.contains(&format!("call void @{}()", llvm(record).symbol))); + assert!(targets.contains(&format!("id: \"{id}\", abi_id: \"{abi_id}\""))); + assert!(targets.contains(&format!("asm: \"{ptx}\""))); + } + for ptx in ["membar.cta;", "membar.gl;", "membar.sys;"] { + assert!(!lowering.contains(&format!("\"{ptx}\""))); + } +} + +#[test] +fn basic_mbarrier_rendering_preserves_existing_paths_shapes_and_routes() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + validate_renderable(&catalog).unwrap(); + assert_eq!(mbarrier_basics(&catalog).count(), 4); + + let compatibility = render_compat_mbarrier_basic(&catalog, "test-hash"); + for signature in [ + "pub unsafe fn mbarrier_init(bar: *mut Barrier, expected_count: u32)", + "pub unsafe fn mbarrier_arrive(bar: *const Barrier) -> u64", + "pub unsafe fn mbarrier_test_wait(bar: *const Barrier, token: u64) -> bool", + "pub unsafe fn mbarrier_inval(bar: *mut Barrier)", + ] { + assert!(compatibility.contains(signature)); + } + assert!(compatibility.contains("eight-byte-aligned `Barrier` in shared memory")); + assert!(compatibility.contains("expected_count` must be in `1..=0xFFFFF`")); + let arrive = compatibility.find("pub unsafe fn mbarrier_arrive").unwrap(); + assert!(compatibility[..arrive].ends_with("#[inline(never)]\n")); + assert!(!compatibility.contains("#[must_use]")); + + let dialect_mod = render_dialect_mod(&catalog, "test-hash"); + assert!(dialect_mod.contains("mod mbarrier_basic;")); + assert!(dialect_mod.contains("mbarrier_basic::register(ctx)")); + let dialect = render_dialect_mbarrier_basic(&catalog, "test-hash"); + assert!(dialect.contains("address_space::GENERIC | address_space::SHARED")); + assert!(dialect.contains("mbarrier expected count must be u32")); + assert!(dialect.contains("mbarrier arrival token must be u64")); + assert!(dialect.contains("mbarrier test-wait requires a u64 token and i1 result")); + + let importer = render_importer(&catalog, "test-hash"); + let lowering = render_lowering(&catalog, "test-hash"); + let targets = render_targets(&catalog, "test-hash"); + let raw = render_raw_abi(&catalog, "test-hash").unwrap(); + for record in mbarrier_basics(&catalog) { + assert!(dialect.contains(&format!("pub struct {}", record.dialect.op_type))); + assert!(dialect.contains(&format!("{}::register(ctx)", record.dialect.op_type))); + assert!(importer.contains(&record.rust.canonical_path)); + assert!(importer.contains(&format!( + "set_generated_intrinsic_marker(ctx, mbarrier, {:?})", + intrinsic_marker(&catalog, record) + ))); + assert!(lowering.contains(&format!( + "impl MirToLlvmConversion for {}", + record.dialect.op_type + ))); + assert!(targets.contains(&format!("id: {:?}", record.id))); + assert!(raw.contains(&format!("pub unsafe fn {}", record.rust.abi_id))); + + let probe = render_probe(&catalog, record, "test-hash"); + assert!( + probe.contains("%barrier = addrspacecast ptr %barrier_generic to ptr addrspace(3)") + ); + let mbarrier = record.mbarrier_basic.as_ref().unwrap(); + let llvm_route = record + .backend_lowerings + .iter() + .find(|lowering| lowering.backend == IntrinsicBackend::LlvmNvptx) + .unwrap(); + let libnvvm_route = record + .backend_lowerings + .iter() + .find(|lowering| lowering.backend == IntrinsicBackend::LibNvvm) + .unwrap(); + let expected_llvm_mechanism = match mbarrier.operation { + MbarrierBasicOperation::TestWait => BackendLoweringMechanism::InlinePtx, + _ => BackendLoweringMechanism::TypedNvvm, + }; + assert_eq!(llvm_route.mechanism, expected_llvm_mechanism); + assert_eq!(libnvvm_route.mechanism, BackendLoweringMechanism::InlinePtx); + match mbarrier.operation { + MbarrierBasicOperation::Init => { + assert!( + dialect.contains("MbarrierInitSharedOp::build(ctx, barrier, expected_count)") + || importer + .contains("MbarrierInitSharedOp::build(ctx, barrier, expected_count)") + ); + assert!( + lowering.contains( + "convert_init(ctx, rewriter, self.get_operation(), operands_info)" + ) + ); + assert!(probe.contains(&format!( + "declare void @{}(ptr addrspace(3), i32)", + llvm(record).symbol + ))); + } + MbarrierBasicOperation::Arrive => { + assert!(importer.contains("MbarrierArriveSharedOp::build(ctx, barrier)")); + assert!(lowering.contains( + "convert_arrive(ctx, rewriter, self.get_operation(), operands_info)" + )); + assert!(probe.contains(&format!( + "declare i64 @{}(ptr addrspace(3))", + llvm(record).symbol + ))); + } + MbarrierBasicOperation::TestWait => { + assert!(importer.contains("MbarrierTestWaitSharedOp::build(ctx, barrier, token)")); + assert!(lowering.contains( + "convert_test_wait(ctx, rewriter, self.get_operation(), operands_info)" + )); + assert!(probe.contains("mbarrier.test_wait.shared.b64")); + assert!(probe.contains("asm sideeffect")); + assert!(probe.contains("attributes #0 = { convergent }")); + assert!(!probe.contains(&format!("declare i1 @{}", llvm(record).symbol))); + } + MbarrierBasicOperation::Inval => { + assert!(importer.contains("MbarrierInvalSharedOp::build(ctx, barrier)")); + assert!( + lowering.contains( + "convert_inval(ctx, rewriter, self.get_operation(), operands_info)" + ) + ); + assert!(probe.contains(&format!( + "declare void @{}(ptr addrspace(3))", + llvm(record).symbol + ))); + } + } + } + + let reference = render_reference(&catalog, "test-hash"); + assert!(reference.contains("## Basic mbarrier contracts")); + assert!(reference.contains("inline PTX on both backends")); + assert!(reference.contains("typed NVVM intrinsic with LLVM-NVPTX")); + assert!(reference.contains("inline PTX with libNVVM")); + + let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); + assert!(outputs.contains_key(&PathBuf::from( + "crates/cuda-device/src/generated/mbarrier_basic.rs" + ))); + assert!(outputs.contains_key(&PathBuf::from( + "crates/dialect-nvvm/src/ops/generated/mbarrier_basic.rs" + ))); +} + +#[test] +fn extended_mbarrier_rendering_preserves_all_manual_contracts() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + validate_renderable(&catalog).unwrap(); + let records = mbarrier_extended(&catalog).collect::>(); + assert_eq!(records.len(), 11); + + 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_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", + "pub unsafe fn fence_proxy_async_shared_cta()", + "pub unsafe fn fence_mbarrier_init_release_cluster()", + "pub unsafe fn fence_proxy_async_generic_release_shared_cta_cluster()", + "pub unsafe fn fence_proxy_async_generic_acquire_shared_cluster_cluster()", + "pub unsafe fn nanosleep(ns: u32)", + ] { + assert!(compatibility.contains(signature), "missing {signature}"); + } + assert!(!compatibility.contains("#[must_use]")); + + let dialect_mod = render_dialect_mod(&catalog, "test-hash"); + assert!(dialect_mod.contains("mod mbarrier_extended;")); + assert!(dialect_mod.contains("pub use mbarrier_extended::*;")); + assert!(dialect_mod.contains("mbarrier_extended::register(ctx);")); + + let dialect = render_dialect_mbarrier_extended(&catalog, "test-hash"); + let importer = render_importer(&catalog, "test-hash"); + let lowering = render_lowering(&catalog, "test-hash"); + let raw = render_raw_abi(&catalog, "test-hash").unwrap(); + for record in &records { + assert!(dialect.contains(&format!("pub struct {}", record.dialect.op_type))); + assert!(dialect.contains(&format!("{}::register(ctx);", record.dialect.op_type))); + assert!(importer.contains(&record.rust.canonical_path)); + assert!(importer.contains(&record.rust.compatibility_paths[0])); + assert!(lowering.contains(&format!( + "impl MirToLlvmConversion for {}", + record.dialect.op_type + ))); + assert!(raw.contains(&format!("pub unsafe fn {}", record.rust.abi_id))); + for route in &record.backend_lowerings { + assert_eq!(route.mechanism, BackendLoweringMechanism::InlinePtx); + } + let (template, constraints) = crate::resolve::mbarrier_extended_inline_recipe( + record.mbarrier_extended.as_ref().unwrap().operation, + ); + assert!(lowering.contains(&format!("{template:?}, {constraints:?}"))); + let probe = render_probe(&catalog, record, "test-hash"); + assert!(probe.contains("asm sideeffect")); + assert!(probe.contains(template)); + assert!(probe.contains("~{memory}")); + assert!(probe.contains("attributes #0 = { convergent }")); + } + + assert!(dialect.contains("address_space::GENERIC | address_space::SHARED")); + assert!(dialect.contains("mbarrier arrival requires u32 bytes and a u64 token")); + assert!(dialect.contains("mbarrier wait requires a u64 token and i1 result")); + assert!(dialect.contains("mbarrier wait requires u32 parity and i1 result")); + assert!(dialect.contains("remote mbarrier address must be u64")); + assert!(dialect.contains("nanosleep duration must be u32")); + + assert!(importer.contains("MbarrierArriveExpectTxSharedOp::build(ctx, barrier, bytes)")); + assert!(importer.contains("MbarrierArriveExpectTxClusterOp::build(ctx, barrier, bytes)")); + assert!( + !importer.contains("MbarrierArriveExpectTxSharedOp::build(ctx, barrier, tx_count, bytes)") + ); + assert!(importer.contains("MbarrierArriveClusterOp::build(ctx, address)")); + 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("FenceProxyAsyncSharedCtaOp::build(ctx)")); + assert!(importer.contains("FenceMbarrierInitReleaseClusterOp::build(ctx)")); + assert!(importer.contains("NanosleepOp::build(ctx, ns)")); + + assert!(lowering.contains("cast_to_shared_addrspace")); + assert!(lowering.contains("trunc_to_i1")); + assert!(lowering.contains("DefiningEntity::Op(result_op)")); + assert!(lowering.contains("rewriter.erase_operation(ctx, op)")); + assert!(lowering.contains("mbarrier.arrive.release.cluster.shared::cluster.b64 _, [$0];")); + let remote = records + .iter() + .find(|record| record.id == "mbarrier_arrive_cluster") + .unwrap(); + assert!(remote.llvm.is_none()); + assert!(matches!(remote.source, IntrinsicSource::PtxNative { .. })); + assert!(!render_probe(&catalog, remote, "test-hash").contains("declare")); + + let reference = render_reference(&catalog, "test-hash"); + assert!(reference.contains("## Extended mbarrier contracts")); + assert!(reference.contains("LLVM 22 TableGen")); + assert!(reference.contains("PTX-native raw-address contract")); + + let outputs = all_outputs(&catalog, "{}\n".into(), "test-hash").unwrap(); + assert!(outputs.contains_key(&PathBuf::from( + "crates/cuda-device/src/generated/mbarrier_extended.rs" + ))); + assert!(outputs.contains_key(&PathBuf::from( + "crates/dialect-nvvm/src/ops/generated/mbarrier_extended.rs" + ))); +} + +#[test] +fn lowering_imports_cast_to_shared_addrspace_exactly_once_across_families() { + // The stmatrix, cluster_memory, and mbarrier_extended families all + // insert the `cast_to_shared_addrspace` helper into the same + // `common::{call_intrinsic, ...}` use group. An unguarded second + // insert duplicates the name inside one use group, which is a hard + // rustc error (E0252) in the raw generator output. The rustfmt pass + // run over generated files silently dedupes the duplicate, so assert + // on the raw (pre-rustfmt) output to catch a regressed guard loudly. + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + assert!(stmatrices(&catalog).next().is_some()); + assert!(cluster_memory(&catalog).next().is_some()); + assert!(mbarrier_extended(&catalog).next().is_some()); + + let import_count = |lowering: &str| { + // Everything before the first rendered helper function is the + // module header plus the use block; helper call sites in the + // function bodies below must not count. + let bodies_start = lowering + .find("fn convert_zero_operand_scalar_direct") + .expect("lowering output must contain the first helper function"); + lowering[..bodies_start] + .matches("cast_to_shared_addrspace") + .count() + }; + + assert_eq!( + import_count(&render_lowering(&catalog, "test-hash")), + 1, + "cast_to_shared_addrspace must be imported exactly once in the \ + raw lowering output when multiple families need it" + ); + + // The guard must not suppress the import either: with stmatrix + // absent, cluster_memory still needs the helper and must insert it. + let mut without_stmatrix = catalog; + without_stmatrix + .intrinsics + .retain(|record| record.family != "stmatrix"); + assert!(stmatrices(&without_stmatrix).next().is_none()); + assert_eq!( + import_count(&render_lowering(&without_stmatrix, "test-hash")), + 1, + "cast_to_shared_addrspace must still be imported when only \ + cluster_memory and mbarrier_extended need it" + ); +} + +#[test] +fn elect_dialect_verifies_its_scalar_shape() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + let dialect = render_dialect_elect(&catalog, "test-hash"); + + assert!(dialect.contains("impl Verify for ElectSyncOp")); + assert!(dialect.contains("is_integer_width(ctx, op.get_operand(0).get_type(ctx), 32)")); + assert!(dialect.contains("is_integer_width(ctx, op.get_result(0).get_type(ctx), 32)")); + assert!(dialect.contains("is_integer_width(ctx, op.get_result(1).get_type(ctx), 1)")); + assert!(!dialect.contains("verifier = \"succ\"")); +} diff --git a/crates/cuda-intrinsics-gen/src/render/validate.rs b/crates/cuda-intrinsics-gen/src/render/validate.rs new file mode 100644 index 0000000000..4b61703361 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/render/validate.rs @@ -0,0 +1,1369 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, CatalogFile, CatalogHardwareAlternative, CatalogHardwareTarget, + ClcAdapter, ClcOperation, ClusterBarrierMode, ClusterBarrierOrdering, ClusterMemoryAdapter, + ClusterMemoryOperation, ClusterMemorySourceContract, CpAsyncMbarrierAdapter, + CpAsyncMbarrierStateSpace, DebugControlAdapter, DebugControlOperation, + ExecutionControlOperation, ExtendedMinMaxAdapter, ExtendedMinMaxFormat, IntegerMinMaxFormat, + IntrinsicBackend, IntrinsicSource, LdmatrixElement, MbarrierBasicAdapter, + MbarrierBasicOperation, MbarrierExtendedAdapter, MbarrierExtendedOperation, + MbarrierExtendedSourceContract, MbarrierStateSpace, RuntimeValidation, ScalarArithmeticFormat, + ScalarArithmeticOperation, ScalarMathFormat, SparseMmaAccumulator, SpecialRegisterObservation, + TmaAdapter, TmaOperation, WarpBarrierAdapter, WarpShuffleAdapter, WarpShuffleOperandEncoding, + WarpShuffleValueKind, WgmmaControlAdapter, WgmmaControlMode, WgmmaControlParticipation, +}; +use crate::render::common::llvm; +use crate::render::families::{ + packed_alu_format_shape, packed_conversion_dialect_operands, packed_conversion_dialect_type, + packed_conversion_is_closed_recipe, packed_conversion_lowering_name, + packed_conversion_rust_arguments, packed_conversion_rust_type, + register_mma_extra_operand_count, sparse_mma_fragment_counts, stmatrix_variant, + tcgen05_render_contract, threadfence_ptx_level, +}; +use anyhow::{Result, ensure}; +use std::collections::BTreeSet; + +pub(super) fn validate_renderable(catalog: &CatalogFile) -> Result<()> { + ensure!( + !catalog.intrinsics.is_empty(), + "catalog contains no intrinsics" + ); + for record in &catalog.intrinsics { + match record.family.as_str() { + "sreg" => { + if let Some(special) = &record.special_register { + ensure!( + matches!(record.rust.module.as_str(), "debug" | "grid" | "thread" | "warp" | "shared") + && record.rust.arguments.is_empty() + && record.dialect.operands.is_empty() + && record.dialect.results + == [format!("i{}", special.result_width.bits())] + && record.scalar_width() == Some(special.result_width.bits()) + && record.lowering == "generated_special_register" + && record.backend_lowerings.len() == 2 + && record.semantics.pure + == (special.observation + == SpecialRegisterObservation::StablePure), + "{} is outside the closed generated special-register recipe", + record.id + ); + } else { + ensure!( + record.rust.module == "sreg" + && record.rust.arguments.is_empty() + && llvm(record).arguments.is_empty() + && record.lowering == "direct_nvvm" + && record.scalar_width().is_some(), + "{} is outside the zero-operand scalar direct-NVVM sreg recipe", + record.id + ); + } + } + "ldmatrix" => ensure!( + record.rust.module == "matrix" + && record.lowering == "generated_ldmatrix" + && record.ldmatrix.as_ref().is_some_and(|ldmatrix| { + record.rust.arguments + == [if ldmatrix.variant.element == LdmatrixElement::B16 { + "*const u32" + } else { + "*const u8" + }] + }), + "{} is outside the generated ldmatrix recipe", + record.id + ), + "stmatrix" => ensure!( + stmatrix_variant(record).is_some() + && record.rust.module == "matrix" + && !record.rust.safe + && !record.rust.must_use + && record.rust.result == "()" + && record.dialect.results.is_empty() + && record.llvm.as_ref().is_some_and(|llvm| llvm.results.is_empty()) + && record.semantics.memory == "write" + && record.semantics.convergent + && record.lowering == "generated_stmatrix", + "{} is outside the closed generated stmatrix recipe", + record.id + ), + "register_mma" => ensure!( + record.rust.module == "matrix" + && record.rust.arguments.len() == 3 + register_mma_extra_operand_count(record) + && !record.rust.safe + && record.rust.must_use + && record.semantics.memory == "none" + && record.semantics.convergent + && record.dialect.op_type == "RegisterMmaOp" + && record.dialect.op_name == "nvvm.register_mma" + && record.lowering == "generated_register_mma" + && record.register_mma.is_some(), + "{} is outside the closed generated register-MMA recipe", + record.id + ), + "movmatrix" => ensure!( + record.rust.module == "matrix" + && record.rust.arguments == ["u32"] + && record.rust.result == "u32" + && !record.rust.safe + && record.rust.must_use + && record.dialect.op_type == "MovmatrixTransB16Op" + && record.dialect.op_name == "nvvm.movmatrix_trans_b16" + && record.dialect.operands == ["i32"] + && record.dialect.results == ["i32"] + && !record.semantics.pure + && record.semantics.memory == "inaccessible_read_write" + && record.semantics.convergent + && record.lowering == "generated_movmatrix_inline_ptx" + && record.movmatrix.is_some(), + "{} is outside the generated movmatrix recipe", + record.id + ), + "sparse_mma" => { + ensure!( + record.sparse_mma.is_some(), + "{} is outside the closed generated sparse-MMA recipe", + record.id + ); + let (c_count, a_count, b_count, d_count) = + sparse_mma_fragment_counts(record); + let accumulator = match record.sparse_mma.as_ref().unwrap().accumulator { + SparseMmaAccumulator::F16 => "u32", + SparseMmaAccumulator::F32 => "f32", + SparseMmaAccumulator::S32 => "i32", + }; + let arguments = [ + format!("[{accumulator}; {c_count}]"), + format!("[u32; {a_count}]"), + format!("[u32; {b_count}]"), + "u32".to_owned(), + "u32".to_owned(), + ]; + ensure!( + record.rust.module == "matrix" + && record.rust.arguments == arguments + && record.rust.result == format!("[{accumulator}; {d_count}]") + && !record.rust.safe + && record.rust.must_use + && record.semantics.memory == "none" + && record.semantics.convergent + && record.dialect.op_type == "SparseMmaOp" + && record.dialect.op_name == "nvvm.sparse_mma" + && record.lowering == "generated_sparse_mma" + && record.sparse_mma.is_some(), + "{} is outside the closed generated sparse-MMA recipe", + record.id + ) + } + "packed_atomic" => ensure!( + record.rust.module == "atomic" + && record.rust.arguments == ["*mut u32", "u32"] + && record.rust.result == "u32" + && record.rust.must_use + && record.llvm.is_none() + && record.lowering == "generated_packed_atomic_inline_ptx" + && record.packed_atomic.is_some(), + "{} is outside the closed generated packed-atomic recipe", + record.id + ), + "redux" => ensure!( + record.rust.module == "warp" + && matches!(record.rust.arguments.as_slice(), [mask, value] + if mask == "u32" && value == &record.rust.result) + && matches!(record.rust.result.as_str(), "u32" | "i32" | "f32") + && !record.rust.safe + && matches!(record.dialect.results.as_slice(), [result] + if matches!(result.as_str(), "i32" | "f32")) + && matches!(record.dialect.results.as_slice(), [result] + if record.llvm.as_ref().is_some_and(|llvm| { + llvm.arguments == [result.as_str(), "i32"] + && llvm.results == [result.as_str()] + }) && record.dialect.operands == ["i32", result.as_str()]) + && record.lowering == "generated_redux" + && record.redux.is_some(), + "{} is outside the closed generated redux recipe", + record.id + ), + "dotprod" => ensure!( + record.rust.module == "dotprod" + && matches!(record.rust.arguments.as_slice(), [a, b, c] + if a == "u32" && b == "u32" && matches!(c.as_str(), "u32" | "i32")) + && record.rust.result == *record.rust.arguments.last().unwrap() + && record.rust.safe + && !record.rust.must_use + && record.llvm.as_ref().is_some_and(|llvm| { + llvm.results == ["i32"] + && (matches!(llvm.arguments.as_slice(), [a, b, c] + if a == "i32" && b == "i32" && c == "i32") + || matches!(llvm.arguments.as_slice(), [a, b, selector, c] + if a == "i32" && b == "i32" && selector == "i1" && c == "i32")) + }) + && record.dialect.operands == ["i32", "i32", "i32"] + && record.dialect.results == ["i32"] + && record.lowering == "generated_dotprod" + && record.dot_product.is_some(), + "{} is outside the closed generated dot-product recipe", + record.id + ), + "packed_alu" => ensure!( + record.packed_alu.as_ref().is_some_and(|packed| { + let (module, must_use, rust_type, dialect_type, adapter) = + packed_alu_format_shape(packed.format); + record.rust.module == module + && record.rust.must_use == must_use + && packed.adapter == adapter + && record + .rust + .arguments + .iter() + .all(|argument| argument == rust_type) + && record.rust.result == rust_type + && record + .dialect + .operands + .iter() + .all(|operand| operand == dialect_type) + && record.dialect.results == [dialect_type] + }) + && (1..=3).contains(&record.rust.arguments.len()) + && record.rust.safe + && record.dialect.operands.len() == record.rust.arguments.len() + && record.lowering == "generated_packed_alu_inline_ptx", + "{} is outside the closed generated packed-ALU recipe", + record.id + ), + "packed_conversion" => ensure!( + record.rust.module == "convert" + && record.rust.arguments == packed_conversion_rust_arguments(record) + && record.rust.result == packed_conversion_rust_type(record) + && record.rust.safe + && !record.rust.must_use + && record.dialect.operands == packed_conversion_dialect_operands(record) + && record.dialect.results == [packed_conversion_dialect_type(record)] + && record.lowering == packed_conversion_lowering_name(record) + && record + .packed_conversion + .as_ref() + .is_some_and(packed_conversion_is_closed_recipe), + "{} is outside the closed generated packed-conversion recipe", + record.id + ), + "scalar_conversion" => ensure!( + record.rust.module == "convert" + && record.rust.arguments == ["f32"] + && record.rust.result == "u32" + && record.rust.safe + && record.rust.must_use + && record.dialect.op_type == "ScalarConversionOp" + && record.dialect.op_name == "nvvm.scalar_conversion" + && record.dialect.operands == ["f32"] + && record.dialect.results == ["i32"] + && record.lowering == "generated_scalar_conversion" + && record.scalar_conversion.is_some(), + "{} is outside the closed generated scalar-conversion recipe", + record.id + ), + "scalar_arithmetic" => ensure!( + record.scalar_arithmetic.as_ref().is_some_and(|arithmetic| { + let ty = match arithmetic.format { + ScalarArithmeticFormat::F32 => "f32", + ScalarArithmeticFormat::F64 => "f64", + }; + let arity = match arithmetic.operation { + ScalarArithmeticOperation::Mul + | ScalarArithmeticOperation::Div + | ScalarArithmeticOperation::Add => 2, + ScalarArithmeticOperation::Fma => 3, + }; + record.rust.arguments == vec![ty; arity] + && record.rust.result == ty + && record.dialect.operands == vec![ty; arity] + && record.dialect.results == [ty] + && record.semantics.pure + == (arithmetic.operation != ScalarArithmeticOperation::Div) + }) + && record.rust.module == "float" + && record.rust.safe + && record.rust.must_use + && record.dialect.op_type == "ScalarArithmeticOp" + && record.dialect.op_name == "nvvm.scalar_arithmetic" + && record.semantics.memory == "none" + && !record.semantics.convergent + && record.lowering == "generated_scalar_arithmetic", + "{} is outside the closed generated scalar-arithmetic recipe", + record.id + ), + "scalar_math" => ensure!( + record.scalar_math.as_ref().is_some_and(|math| { + let (rust_ty, dialect_ty) = match math.format { + ScalarMathFormat::F16 => ("u16", "i16"), + ScalarMathFormat::F32 => ("f32", "f32"), + ScalarMathFormat::F64 => ("f64", "f64"), + }; + record.rust.arguments == [rust_ty] + && record.rust.result == rust_ty + && record.dialect.operands == [dialect_ty] + && record.dialect.results == [dialect_ty] + && record.semantics.pure + }) + && record.rust.module == "float" + && record.rust.safe + && record.rust.must_use + && record.dialect.op_type == "ScalarMathOp" + && record.dialect.op_name == "nvvm.scalar_math" + && record.semantics.memory == "none" + && !record.semantics.convergent + && record.lowering == "generated_scalar_math", + "{} is outside the closed generated scalar-math recipe", + record.id + ), + "extended_minmax" => ensure!( + record.extended_minmax.as_ref().is_some_and(|minmax| { + let (module, rust_type, dialect_type, adapter) = match minmax.format { + ExtendedMinMaxFormat::F32 => ( + "float", + "f32", + "f32", + ExtendedMinMaxAdapter::DirectF32, + ), + ExtendedMinMaxFormat::F16 => ( + "f16", + "u16", + "i16", + ExtendedMinMaxAdapter::DirectHalfU16, + ), + ExtendedMinMaxFormat::Bf16 => ( + "bf16", + "u16", + "i16", + ExtendedMinMaxAdapter::DirectHalfU16, + ), + ExtendedMinMaxFormat::F16x2 => ( + "f16x2", + "u32", + "i32", + ExtendedMinMaxAdapter::DirectPackedU32, + ), + ExtendedMinMaxFormat::Bf16x2 => ( + "bf16x2", + "u32", + "i32", + ExtendedMinMaxAdapter::DirectPackedU32, + ), + }; + record.rust.module == module + && record.rust.arguments == [rust_type, rust_type] + && record.rust.result == rust_type + && record.dialect.operands == [dialect_type, dialect_type] + && record.dialect.results == [dialect_type] + && minmax.adapter == adapter + }) + && record.rust.safe + && record.rust.must_use + && record.dialect.op_type == "ExtendedMinMaxOp" + && record.dialect.op_name == "nvvm.extended_minmax" + && record.semantics.pure + && record.semantics.memory == "none" + && !record.semantics.convergent + && record.lowering == "generated_extended_minmax", + "{} is outside the closed generated extended-minmax recipe", + record.id + ), + "prmt" => ensure!( + record.rust.module == "prmt" + && matches!(record.rust.arguments.len(), 2 | 3) + && record.rust.arguments.iter().all(|argument| argument == "u32") + && record.rust.result == "u32" + && record.rust.safe + && record.rust.must_use + && record.dialect.op_type == "PrmtOp" + && record.dialect.op_name == "nvvm.prmt" + && record.dialect.operands.len() == record.rust.arguments.len() + && record.dialect.operands.iter().all(|operand| operand == "i32") + && record.dialect.results == ["i32"] + && record.lowering == "generated_prmt" + && record.prmt.is_some(), + "{} is outside the closed generated prmt recipe", + record.id + ), + "cluster_barrier" => ensure!( + matches!( + record.id.as_str(), + "barrier_cluster_arrive" + | "barrier_cluster_arrive_aligned" + | "barrier_cluster_arrive_relaxed" + | "barrier_cluster_arrive_relaxed_aligned" + | "barrier_cluster_wait" + | "barrier_cluster_wait_aligned" + ) + && record.rust.module == "cluster" + && record.rust.arguments.is_empty() + && record.rust.result == "()" + && !record.rust.safe + && !record.rust.must_use + && record.dialect.op_type == "ClusterBarrierOp" + && record.dialect.op_name == "nvvm.cluster_barrier" + && record.dialect.operands.is_empty() + && record.dialect.results.is_empty() + && record.llvm.as_ref().is_some_and(|llvm| { + llvm.arguments.is_empty() && llvm.results.is_empty() + }) + && record.cluster_barrier.as_ref().is_some_and(|barrier| { + matches!( + (barrier.mode, barrier.ordering, barrier.aligned), + ( + ClusterBarrierMode::Arrive, + ClusterBarrierOrdering::Release, + false + ) | ( + ClusterBarrierMode::ArriveAligned, + ClusterBarrierOrdering::Release, + true + ) | ( + ClusterBarrierMode::ArriveRelaxed, + ClusterBarrierOrdering::Relaxed, + false + ) | ( + ClusterBarrierMode::ArriveRelaxedAligned, + ClusterBarrierOrdering::Relaxed, + true + ) | ( + ClusterBarrierMode::Wait, + ClusterBarrierOrdering::Acquire, + false + ) | ( + ClusterBarrierMode::WaitAligned, + ClusterBarrierOrdering::Acquire, + true + ) + ) + }) + && record.lowering == "generated_cluster_barrier", + "{} is outside the closed generated cluster-barrier recipe", + record.id + ), + "cluster_memory" => ensure!( + record.rust.module == "cluster" + && !record.rust.safe + && record.rust.must_use + && record.semantics.convergent + && record.semantics.execution_scope == "cluster" + && record.target.minimum_ptx.to_string() == "7.8" + && matches!( + &record.target.hardware, + CatalogHardwareTarget::AnyOf { alternatives } + if alternatives.as_slice() + == [CatalogHardwareAlternative::MinimumSm { sm: 90 }] + ) + && record.lowering == "generated_cluster_memory_inline_ptx" + && record.backend_lowerings.len() == 2 + && record.backend_lowerings.iter().all(|lowering| { + lowering.mechanism == BackendLoweringMechanism::InlinePtx + && lowering.target.minimum_ptx.to_string() == "7.8" + && matches!( + &lowering.target.hardware, + CatalogHardwareTarget::AnyOf { alternatives } + if alternatives.as_slice() + == [CatalogHardwareAlternative::MinimumSm { sm: 90 }] + ) + }) + && record.cluster_memory.as_ref().is_some_and(|cluster| { + cluster.runtime_validation == RuntimeValidation::Unexecuted + && match (cluster.operation, cluster.adapter, cluster.source_contract) { + ( + ClusterMemoryOperation::MapSharedRank, + ClusterMemoryAdapter::GenericConstAndMutPointerRankToSamePointer, + ClusterMemorySourceContract::LlvmMapaSharedClusterAs7IdentityInlinePtx, + ) => { + record.id == "map_shared_rank" + && record.rust.arguments == ["*const u8", "u32"] + && record.rust.result == "*const u8" + && record.rust.compatibility_paths + == [ + "cuda_device::cluster::map_shared_rank", + "cuda_device::cluster::map_shared_rank_mut", + ] + && record.dialect.op_type == "MapaSharedClusterOp" + && record.dialect.op_name == "nvvm.mapa_shared_cluster" + && record.dialect.operands == ["ptr", "i32"] + && record.dialect.results == ["ptr"] + && record.semantics.memory == "none" + && matches!( + record.source, + IntrinsicSource::LlvmImported { ref source_record } + if source_record == "int_nvvm_mapa_shared_cluster" + ) + && record.llvm.as_ref().is_some_and(|llvm| { + llvm.symbol == "llvm.nvvm.mapa.shared.cluster" + && llvm.arguments == ["shared_ptr", "i32"] + && llvm.results == ["shared_cluster_ptr"] + && llvm.properties + == [ + "IntrNoMem", + "IntrSpeculatable", + "NoCapture", + ] + }) + && record + .selections + .iter() + .map(|selection| selection.source_record.as_str()) + .collect::>() + == BTreeSet::from([ + "mapa_shared_cluster_64", + "mapa_shared_cluster_64i", + ]) + } + ( + ClusterMemoryOperation::ReadU32, + ClusterMemoryAdapter::ConstU32PointerRankToU32, + ClusterMemorySourceContract::PtxNativeMapaThenWeakClusterLoad, + ) => { + record.id == "dsmem_read_u32" + && record.rust.arguments == ["*const u32", "u32"] + && record.rust.result == "u32" + && record.rust.compatibility_paths + == ["cuda_device::cluster::dsmem_read_u32"] + && record.dialect.op_type == "DsmemReadU32Op" + && record.dialect.op_name == "nvvm.dsmem_read_u32" + && record.dialect.operands == ["ptr", "i32"] + && record.dialect.results == ["i32"] + && record.semantics.memory == "read" + && record.llvm.is_none() + && record.selections.is_empty() + && matches!( + record.source, + IntrinsicSource::PtxNative { ref instruction } + if instruction + == "mapa.shared::cluster.u64 + ld.shared::cluster.u32" + ) + } + _ => false, + } + }), + "{} is outside the closed generated cluster-memory recipe", + record.id + ), + "debug_control" => ensure!( + record.rust.module == "debug" + && record.rust.safe + && !record.rust.must_use + && record.dialect.operands.is_empty() + && record.dialect.results.is_empty() + && record.llvm.is_none() + && record.lowering == "generated_debug_control" + && record.debug_control.as_ref().is_some_and(|debug| { + debug.runtime_validation == RuntimeValidation::Unexecuted + && match debug.operation { + DebugControlOperation::Trap => { + debug.adapter == DebugControlAdapter::Direct + && record.rust.arguments.is_empty() + && record.rust.result == "!" + && record.dialect.op_type == "TrapOp" + && record.dialect.op_name == "nvvm.trap" + } + DebugControlOperation::Breakpoint => { + debug.adapter == DebugControlAdapter::Direct + && record.rust.arguments.is_empty() + && record.rust.result == "()" + && record.dialect.op_type == "BreakpointOp" + && record.dialect.op_name == "nvvm.brkpt" + } + DebugControlOperation::Pmevent => { + debug.adapter + == DebugControlAdapter::ConstGenericToImmediateU32 + && record.rust.arguments == ["u32"] + && record.rust.result == "()" + && record.dialect.op_type == "PmEventOp" + && record.dialect.op_name == "nvvm.pmevent" + } + } + }), + "{} is outside the closed generated debug-control recipe", + record.id + ), + "clc" => ensure!( + record.rust.module == "clc" + && !record.rust.safe + && !record.rust.must_use + && record.lowering == "generated_clc" + && record.llvm.is_some() + && record.clc.as_ref().is_some_and(|clc| { + clc.runtime_validation == RuntimeValidation::Unexecuted + && match (clc.operation, clc.adapter) { + ( + ClcOperation::TryCancel + | ClcOperation::TryCancelMulticast, + ClcAdapter::GenericPointersToShared, + ) => { + record.rust.arguments == ["*mut u8", "*mut u64"] + && record.rust.result == "()" + && record.dialect.operands == ["ptr", "ptr"] + && record.dialect.results.is_empty() + && record.llvm.as_ref().is_some_and(|llvm| { + llvm.arguments == ["shared_ptr", "shared_ptr"] + && llvm.results.is_empty() + }) + } + ( + ClcOperation::QueryIsCanceled, + ClcAdapter::PairU64ToI128BoolToU32, + ) => { + record.rust.arguments == ["u64", "u64"] + && record.rust.result == "u32" + && record.dialect.operands == ["i64", "i64"] + && record.dialect.results == ["i32"] + && record.llvm.as_ref().is_some_and(|llvm| { + llvm.arguments == ["i128"] + && llvm.results == ["i1"] + }) + } + ( + ClcOperation::QueryGetFirstCtaidX + | ClcOperation::QueryGetFirstCtaidY + | ClcOperation::QueryGetFirstCtaidZ, + ClcAdapter::PairU64ToI128U32, + ) => { + record.rust.arguments == ["u64", "u64"] + && record.rust.result == "u32" + && record.dialect.operands == ["i64", "i64"] + && record.dialect.results == ["i32"] + && record.llvm.as_ref().is_some_and(|llvm| { + llvm.arguments == ["i128"] + && llvm.results == ["i32"] + }) + } + _ => false, + } + }), + "{} is outside the closed generated CLC recipe", + record.id + ), + "wgmma_control" => ensure!( + record.rust.module == "wgmma" + && record.rust.result == "()" + && !record.rust.safe + && !record.rust.must_use + && record.dialect.results.is_empty() + && record.llvm.as_ref().is_some_and(|llvm| llvm.results.is_empty()) + && !record.semantics.pure + && record.semantics.memory == "read_write" + && record.semantics.convergent + && record.semantics.execution_scope == "warpgroup" + && record.target.minimum_ptx.encoded() == 80 + && record.target.hardware + == CatalogHardwareTarget::AnyOf { + alternatives: vec![CatalogHardwareAlternative::ExactArchitecture { + sm: 90, + }], + } + && record.backend_lowerings.len() == 2 + && record.backend_lowerings.iter().any(|lowering| { + lowering.backend == IntrinsicBackend::LlvmNvptx + && lowering.mechanism == BackendLoweringMechanism::TypedNvvm + }) + && record.backend_lowerings.iter().any(|lowering| { + lowering.backend == IntrinsicBackend::LibNvvm + && lowering.mechanism == BackendLoweringMechanism::InlinePtx + }) + && record.lowering == "generated_wgmma_control" + && record.wgmma_control.as_ref().is_some_and(|control| { + control.participation + == WgmmaControlParticipation::WarpgroupAllThreadsSameInstruction + && match control.mode { + WgmmaControlMode::Fence => { + control.adapter == WgmmaControlAdapter::NoArguments + && record.id == "wgmma_fence" + && record.rust.arguments.is_empty() + && record.dialect.op_type == "WgmmaFenceSyncAlignedOp" + && record.dialect.op_name + == "nvvm.wgmma_fence_sync_aligned" + && record.dialect.operands.is_empty() + && record.llvm.as_ref().is_some_and(|llvm| { + llvm.arguments.is_empty() + }) + } + WgmmaControlMode::CommitGroup => { + control.adapter == WgmmaControlAdapter::NoArguments + && record.id == "wgmma_commit_group" + && record.rust.arguments.is_empty() + && record.dialect.op_type + == "WgmmaCommitGroupSyncAlignedOp" + && record.dialect.op_name + == "nvvm.wgmma_commit_group_sync_aligned" + && record.dialect.operands.is_empty() + && record.llvm.as_ref().is_some_and(|llvm| { + llvm.arguments.is_empty() + }) + } + WgmmaControlMode::WaitGroup => { + control.adapter + == WgmmaControlAdapter::ConstGenericU32ToI64Immediate + && record.id == "wgmma_wait_group" + && record.rust.arguments == ["u64"] + && record.rust.compatibility_paths + == ["cuda_device::wgmma::__wgmma_wait_group"] + && record.dialect.op_type + == "WgmmaWaitGroupSyncAlignedOp" + && record.dialect.op_name + == "nvvm.wgmma_wait_group_sync_aligned" + && record.dialect.operands == ["i64"] + && record.llvm.as_ref().is_some_and(|llvm| { + llvm.arguments == ["i64"] + }) + } + } + }), + "{} is outside the closed generated WGMMA-control recipe", + record.id + ), + "tma" => ensure!( + record.rust.module == "tma" + && record.rust.result == "()" + && !record.rust.must_use + && record.dialect.results.is_empty() + && record.llvm.as_ref().is_some_and(|llvm| llvm.results.is_empty()) + && record.lowering == "generated_tma" + && record.tma.as_ref().is_some_and(|tma| { + tma.runtime_validation == RuntimeValidation::Unexecuted + && match (tma.operation, tma.adapter) { + ( + TmaOperation::G2sTile1d + | TmaOperation::G2sTile2d + | TmaOperation::G2sTile3d + | TmaOperation::G2sTile4d + | TmaOperation::G2sTile5d, + TmaAdapter::G2sPointersCoordinatesBarrierInjectDefaults, + ) => { + !record.rust.safe + && record.semantics.convergent + && record.semantics.memory == "read_write" + } + ( + TmaOperation::G2sTile2dMulticast + | TmaOperation::G2sTile2dMulticastCg2, + TmaAdapter::G2sPointersCoordinatesBarrierMaskInjectDefaults, + ) => { + !record.rust.safe + && record.semantics.convergent + && record.semantics.memory == "read_write" + } + ( + TmaOperation::S2gTile1d + | TmaOperation::S2gTile2d + | TmaOperation::S2gTile3d + | TmaOperation::S2gTile4d + | TmaOperation::S2gTile5d, + TmaAdapter::S2gPointersCoordinatesInjectDefaults, + ) => { + !record.rust.safe + && record.semantics.convergent + && record.semantics.memory == "read_write" + } + ( + TmaOperation::Reduce, + TmaAdapter::ReductionPointersCoordinatesInjectDefaults, + ) => { + !record.rust.safe + && record.semantics.convergent + && record.semantics.memory == "read_write" + && tma.reduction.is_some() + } + (TmaOperation::CommitGroup, TmaAdapter::NoOperands) => { + record.rust.safe + && !record.semantics.convergent + && record.semantics.memory == "read_write" + } + ( + TmaOperation::WaitGroup | TmaOperation::WaitGroupRead, + TmaAdapter::CompileTimeConstantMaxPending, + ) => { + record.rust.safe + && !record.semantics.convergent + && record.semantics.memory == "read_write" + } + ( + TmaOperation::PrefetchTensorMap, + TmaAdapter::DescriptorPointer, + ) => { + !record.rust.safe + && !record.semantics.convergent + && record.semantics.memory == "read" + } + ( + TmaOperation::PrefetchTile1d + | TmaOperation::PrefetchTile2d + | TmaOperation::PrefetchTile3d + | TmaOperation::PrefetchTile4d + | TmaOperation::PrefetchTile5d + | TmaOperation::PrefetchTileGather4TwoDimensional, + TmaAdapter::DescriptorCoordinatesInjectDefaults, + ) => { + !record.rust.safe + && record.semantics.convergent + && record.semantics.memory == "read" + } + ( + TmaOperation::PrefetchTile1dCacheHint + | TmaOperation::PrefetchTile2dCacheHint + | TmaOperation::PrefetchTile3dCacheHint + | TmaOperation::PrefetchTile4dCacheHint + | TmaOperation::PrefetchTile5dCacheHint + | TmaOperation::PrefetchTileGather4TwoDimensionalCacheHint, + TmaAdapter::DescriptorCoordinatesCacheHintInjectFlag, + ) => { + !record.rust.safe + && record.semantics.convergent + && record.semantics.memory == "read" + } + ( + TmaOperation::ReplaceGlobalAddress, + TmaAdapter::DescriptorAndAddressPointers, + ) + | ( + TmaOperation::ReplaceBoxDim + | TmaOperation::ReplaceElementStride + | TmaOperation::ReplaceGlobalDim, + TmaAdapter::DescriptorOrdinalAndU32, + ) + | ( + TmaOperation::ReplaceGlobalStride, + TmaAdapter::DescriptorOrdinalAndU64, + ) + | ( + TmaOperation::ReplaceElementType + | TmaOperation::ReplaceFillMode + | TmaOperation::ReplaceInterleaveLayout + | TmaOperation::ReplaceSwizzleAtomicity + | TmaOperation::ReplaceSwizzleMode, + TmaAdapter::DescriptorAndImmediateU32, + ) + | ( + TmaOperation::ReplaceRank, + TmaAdapter::DescriptorAndRuntimeU32, + ) => { + !record.rust.safe + && !record.semantics.convergent + && record.semantics.memory == "write" + } + ( + TmaOperation::FenceProxyTensorMapAcquireCluster + | TmaOperation::FenceProxyTensorMapAcquireCta + | TmaOperation::FenceProxyTensorMapAcquireGpu + | TmaOperation::FenceProxyTensorMapAcquireSystem, + TmaAdapter::DescriptorPointerInjectBytes, + ) => { + !record.rust.safe + && !record.semantics.convergent + && record.semantics.memory == "read_write" + } + ( + TmaOperation::FenceProxyTensorMapReleaseCluster + | TmaOperation::FenceProxyTensorMapReleaseCta + | TmaOperation::FenceProxyTensorMapReleaseGpu + | TmaOperation::FenceProxyTensorMapReleaseSystem, + TmaAdapter::NoOperands, + ) => { + record.rust.safe + && !record.semantics.convergent + && record.semantics.memory == "read_write" + } + _ => false, + } + }), + "{} is outside the closed generated TMA recipe", + record.id + ), + "counted_barrier" | "grid_dependency" | "register_control" => ensure!( + ExecutionControlOperation::from_catalog_id(&record.id).is_some_and(|operation| { + record.family == operation.family() + && record.rust.result == "()" + && !record.rust.safe + && !record.rust.must_use + && record.dialect.results.is_empty() + && record + .llvm + .as_ref() + .is_some_and(|llvm| llvm.results.is_empty()) + && record.lowering == "generated_execution_control" + && match operation { + ExecutionControlOperation::BarrierCtaSync + | ExecutionControlOperation::BarrierCtaSyncAligned + | ExecutionControlOperation::BarrierCtaArrive + | ExecutionControlOperation::BarrierCtaArriveAligned => { + record.rust.module == "barrier" + && record.rust.arguments == ["u32", "u32"] + && record.dialect.operands == ["i32", "i32"] + && record.llvm.as_ref().is_some_and(|llvm| { + llvm.arguments == ["i32", "i32"] + }) + } + ExecutionControlOperation::GridDependencyLaunchDependents + | ExecutionControlOperation::GridDependencyWait => { + record.rust.module == "grid" + && record.rust.arguments.is_empty() + && record.dialect.operands.is_empty() + && record + .llvm + .as_ref() + .is_some_and(|llvm| llvm.arguments.is_empty()) + } + ExecutionControlOperation::SetMaxNRegInc + | ExecutionControlOperation::SetMaxNRegDec => { + record.rust.module == "thread" + && record.rust.arguments == ["u32"] + && record.dialect.operands.is_empty() + && record + .llvm + .as_ref() + .is_some_and(|llvm| llvm.arguments == ["i32"]) + } + } + }), + "{} is outside the closed generated execution-control recipe", + record.id + ), + "tcgen05" => ensure!( + tcgen05_render_contract(record), + "{} is outside the closed generated tcgen05 recipe", + record.id + ), + "cp_async_copy" => ensure!( + record.rust.module == "async_copy" + && record.rust.result == "()" + && !record.rust.safe + && !record.rust.must_use + && record.llvm.as_ref().is_some_and(|llvm| { + matches!(llvm.arguments.as_slice(), [dst, src] + if dst == "shared_ptr" && src == "global_ptr") + || matches!(llvm.arguments.as_slice(), [dst, src, size] + if dst == "shared_ptr" && src == "global_ptr" && size == "i32") + }) + && record.llvm.as_ref().is_some_and(|llvm| llvm.results.is_empty()) + && record.dialect.results.is_empty() + && record.lowering == "generated_cp_async_copy" + && record.cp_async_copy.is_some(), + "{} is outside the closed generated cp.async copy recipe", + record.id + ), + "cp_async_control" => ensure!( + record.rust.module == "async_copy" + && record.rust.result == "()" + && !record.rust.safe + && !record.rust.must_use + && record.llvm.as_ref().is_some_and(|llvm| { + llvm.results.is_empty() + && (llvm.arguments.is_empty() || llvm.arguments == ["i32"]) + }) + && record.dialect.results.is_empty() + && record.lowering == "generated_cp_async_control" + && record.cp_async_control.is_some(), + "{} is outside the closed generated cp.async control recipe", + record.id + ), + "cp_async_mbarrier" => ensure!( + record.rust.module == "async_copy" + && record.rust.arguments == ["*mut u64"] + && record.rust.result == "()" + && !record.rust.safe + && !record.rust.must_use + && record.dialect.operands == ["ptr"] + && record.dialect.results.is_empty() + && record.lowering == "generated_cp_async_mbarrier" + && record.cp_async_mbarrier.as_ref().is_some_and(|bridge| { + bridge.adapter == CpAsyncMbarrierAdapter::PointerToVoid + && record.llvm.as_ref().is_some_and(|llvm| { + llvm.results.is_empty() + && match bridge.state_space { + CpAsyncMbarrierStateSpace::Generic => { + llvm.arguments == ["ptr"] + } + CpAsyncMbarrierStateSpace::Shared => { + llvm.arguments == ["shared_ptr"] + } + } + }) + }), + "{} is outside the closed generated cp.async mbarrier recipe", + record.id + ), + "mbarrier_basic" => ensure!( + record.rust.module == "barrier" + && !record.rust.safe + && record.lowering == "generated_mbarrier_basic" + && record.mbarrier_basic.as_ref().is_some_and(|mbarrier| { + mbarrier.state_space == MbarrierStateSpace::Shared + && match (mbarrier.operation, mbarrier.adapter) { + ( + MbarrierBasicOperation::Init, + MbarrierBasicAdapter::InitPointerCountToVoid, + ) => { + record.rust.arguments == ["*mut u64", "u32"] + && record.rust.result == "()" + && !record.rust.must_use + && record.llvm.as_ref().is_some_and(|llvm| { + llvm.arguments == ["shared_ptr", "i32"] + && llvm.results.is_empty() + }) + && record.dialect.operands == ["ptr", "i32"] + && record.dialect.results.is_empty() + } + ( + MbarrierBasicOperation::Arrive, + MbarrierBasicAdapter::ArrivePointerToToken, + ) => { + record.rust.arguments == ["*const u64"] + && record.rust.result == "u64" + && record.rust.must_use + && record.llvm.as_ref().is_some_and(|llvm| { + llvm.arguments == ["shared_ptr"] + && llvm.results == ["i64"] + }) + && record.dialect.operands == ["ptr"] + && record.dialect.results == ["i64"] + } + ( + MbarrierBasicOperation::TestWait, + MbarrierBasicAdapter::TestWaitPointerTokenToPredicate, + ) => { + record.rust.arguments == ["*const u64", "u64"] + && record.rust.result == "bool" + && record.rust.must_use + && record.llvm.as_ref().is_some_and(|llvm| { + llvm.arguments == ["shared_ptr", "i64"] + && llvm.results == ["i1"] + }) + && record.dialect.operands == ["ptr", "i64"] + && record.dialect.results == ["i1"] + } + ( + MbarrierBasicOperation::Inval, + MbarrierBasicAdapter::InvalPointerToVoid, + ) => { + record.rust.arguments == ["*mut u64"] + && record.rust.result == "()" + && !record.rust.must_use + && record.llvm.as_ref().is_some_and(|llvm| { + llvm.arguments == ["shared_ptr"] + && llvm.results.is_empty() + }) + && record.dialect.operands == ["ptr"] + && record.dialect.results.is_empty() + } + _ => false, + } + }), + "{} is outside the closed generated basic mbarrier recipe", + record.id + ), + "mbarrier_extended" => ensure!( + record.rust.module == "barrier" + && !record.rust.safe + && record.semantics.memory == "read_write" + && record.semantics.convergent + && record.lowering == "generated_mbarrier_extended_inline_ptx" + && record.backend_lowerings.iter().all(|lowering| { + lowering.mechanism == BackendLoweringMechanism::InlinePtx + }) + && record.mbarrier_extended.as_ref().is_some_and(|mbarrier| { + let source_matches = match mbarrier.source_contract { + MbarrierExtendedSourceContract::LlvmImported => record.llvm.is_some(), + MbarrierExtendedSourceContract::PtxNativeRawClusterAddress => { + record.llvm.is_none() + && matches!(record.source, IntrinsicSource::PtxNative { .. }) + } + }; + source_matches + && match (mbarrier.operation, mbarrier.adapter) { + ( + MbarrierExtendedOperation::ArriveExpectTxCta, + MbarrierExtendedAdapter::PointerTxCountBytesToTokenDroppingTxCount, + ) => record.id == "mbarrier_arrive_expect_tx", + ( + MbarrierExtendedOperation::ArriveExpectTxCluster, + MbarrierExtendedAdapter::PointerTxCountBytesToTokenDroppingTxCount, + ) => record.id == "mbarrier_arrive_expect_tx_cluster", + ( + MbarrierExtendedOperation::ArriveRemoteCluster, + MbarrierExtendedAdapter::RawClusterAddressToVoid, + ) => record.id == "mbarrier_arrive_cluster", + ( + MbarrierExtendedOperation::TryWaitTokenCta, + MbarrierExtendedAdapter::PointerTokenToPredicate, + ) => record.id == "mbarrier_try_wait", + ( + MbarrierExtendedOperation::TryWaitParityCta, + MbarrierExtendedAdapter::PointerParityToPredicate, + ) => record.id == "mbarrier_try_wait_parity", + ( + MbarrierExtendedOperation::TryWaitParityCluster, + MbarrierExtendedAdapter::PointerParityToPredicate, + ) => record.id == "mbarrier_try_wait_parity_cluster", + ( + MbarrierExtendedOperation::FenceProxyAsyncSharedCta, + MbarrierExtendedAdapter::ZeroOperandsToVoid, + ) => record.id == "fence_proxy_async_shared_cta", + ( + MbarrierExtendedOperation::FenceMbarrierInitReleaseCluster, + MbarrierExtendedAdapter::ZeroOperandsToVoid, + ) => record.id == "fence_mbarrier_init_release_cluster", + ( + MbarrierExtendedOperation::FenceProxyAsyncGenericReleaseSharedCtaCluster, + MbarrierExtendedAdapter::ZeroOperandsToVoid, + ) => record.id + == "fence_proxy_async_generic_release_shared_cta_cluster", + ( + MbarrierExtendedOperation::FenceProxyAsyncGenericAcquireSharedClusterCluster, + MbarrierExtendedAdapter::ZeroOperandsToVoid, + ) => record.id + == "fence_proxy_async_generic_acquire_shared_cluster_cluster", + ( + MbarrierExtendedOperation::Nanosleep, + MbarrierExtendedAdapter::NanosecondsToVoid, + ) => record.id == "nanosleep", + _ => false, + } + }), + "{} is outside the closed generated extended-mbarrier recipe", + record.id + ), + "sync" => { + if record.id == "sync_threads" { + ensure!( + record.rust.module == "thread" + && record.rust.name == "sync_threads" + && record.rust.arguments.is_empty() + && record.rust.result == "()" + && !record.rust.safe + && !record.rust.must_use + && record.llvm.as_ref().is_some_and(|llvm| { + llvm.symbol == "llvm.nvvm.barrier.cta.sync.aligned.all" + && llvm.arguments == ["i32"] + && llvm.results.is_empty() + }) + && record.dialect.op_type == "Barrier0Op" + && record.dialect.op_name == "nvvm.barrier0" + && record.dialect.operands.is_empty() + && record.dialect.results.is_empty() + && record.lowering == "generated_sync_threads", + "{} is outside the fixed-zero generated sync_threads recipe", + record.id + ); + } else { + ensure!( + threadfence_ptx_level(record).is_some() + && record.rust.module == "fence" + && record.rust.name == record.id + && record.rust.arguments.is_empty() + && record.rust.result == "()" + && record.rust.safe + && !record.rust.must_use + && record.llvm.as_ref().is_some_and(|llvm| { + llvm.arguments.is_empty() && llvm.results.is_empty() + }) + && record.dialect.operands.is_empty() + && record.dialect.results.is_empty() + && record.lowering == "direct_nvvm" + && record.semantics.memory == "read_write" + && !record.semantics.pure + && !record.semantics.convergent, + "{} is outside the closed generated thread-fence recipe", + record.id + ); + } + } + "vote" => ensure!( + record.rust.module == "warp" + && record.rust.arguments == ["u32", "bool"] + && matches!(record.rust.result.as_str(), "bool" | "u32") + && !record.rust.safe + && record.rust.must_use + && record.llvm.as_ref().is_some_and(|llvm| { + llvm.arguments == ["i32", "i1"] + && matches!(llvm.results.as_slice(), [result] + if result == "i1" || result == "i32") + }) + && record.dialect.operands == ["i32", "i1"] + && matches!(record.dialect.results.as_slice(), [result] + if result == "i1" || result == "i32") + && record.lowering == "generated_vote" + && record.vote.is_some(), + "{} is outside the closed generated vote.sync recipe", + record.id + ), + "active_mask" => ensure!( + record.id == "active_mask" + && record.rust.module == "warp" + && record.rust.arguments.is_empty() + && record.rust.result == "u32" + && record.rust.safe + && record.rust.must_use + && record.llvm.as_ref().is_some_and(|llvm| { + llvm.arguments.is_empty() && llvm.results == ["i32"] + }) + && record.dialect.operands.is_empty() + && record.dialect.results == ["i32"] + && record.lowering == "generated_active_mask" + && record.active_mask.is_some(), + "{} is outside the closed generated active-mask recipe", + record.id + ), + "warp_match" => ensure!( + record.rust.module == "warp" + && matches!(record.rust.arguments.as_slice(), [mask, value] + if mask == "u32" && matches!(value.as_str(), "u32" | "u64")) + && record.rust.result == "u32" + && !record.rust.safe + && record.rust.must_use + && record.llvm.as_ref().is_some_and(|llvm| { + matches!(llvm.arguments.as_slice(), [mask, value] + if mask == "i32" && matches!(value.as_str(), "i32" | "i64")) + && (matches!(llvm.results.as_slice(), [mask] if mask == "i32") + || matches!(llvm.results.as_slice(), [mask, predicate] + if mask == "i32" && predicate == "i1")) + }) + && matches!(record.dialect.operands.as_slice(), [mask, value] + if mask == "i32" && matches!(value.as_str(), "i32" | "i64")) + && record.dialect.results == ["i32"] + && record.lowering == "generated_warp_match" + && record.warp_match.is_some(), + "{} is outside the closed generated match.sync recipe", + record.id + ), + "elect" => ensure!( + record.id == "elect_sync" + && record.rust.module == "warp" + && record.rust.name == "elect_sync" + && record.rust.arguments == ["u32"] + && record.rust.result == "(u32, bool)" + && !record.rust.safe + && record.rust.must_use + && record.llvm.as_ref().is_some_and(|llvm| { + llvm.symbol == "llvm.nvvm.elect.sync" + && llvm.arguments == ["i32"] + && llvm.results == ["i32", "i1"] + }) + && record.dialect.op_type == "ElectSyncOp" + && record.dialect.op_name == "nvvm.elect_sync" + && record.dialect.operands == ["i32"] + && record.dialect.results == ["i32", "i1"] + && record.lowering == "generated_elect", + "{} is outside the closed generated elect.sync recipe", + record.id + ), + "warp_barrier" => ensure!( + record.id == "sync_mask" + && record.rust.module == "warp" + && record.rust.name == "sync_mask" + && record.rust.arguments == ["u32"] + && record.rust.result == "()" + && !record.rust.safe + && !record.rust.must_use + && record.llvm.as_ref().is_some_and(|llvm| { + llvm.symbol == "llvm.nvvm.bar.warp.sync" + && llvm.arguments == ["i32"] + && llvm.results.is_empty() + }) + && record.dialect.op_type == "BarWarpSyncOp" + && record.dialect.op_name == "nvvm.bar_warp_sync" + && record.dialect.operands == ["i32"] + && record.dialect.results.is_empty() + && record.lowering == "generated_warp_barrier" + && record.warp_barrier.as_ref().is_some_and(|barrier| { + barrier.adapter == WarpBarrierAdapter::DirectMemberMask + }), + "{} is outside the closed generated bar.warp.sync recipe", + record.id + ), + "warp_shuffle" => ensure!( + record.rust.module == "warp" + && matches!(record.rust.arguments.as_slice(), [mask, value, lane] + if mask == "u32" + && lane == "u32" + && matches!(value.as_str(), "u32" | "f32" | "u64") + && value == &record.rust.result) + && !record.rust.safe + && record.rust.must_use + && matches!(record.dialect.operands.as_slice(), [mask, value, lane] + if mask == "i32" + && lane == "i32" + && matches!(value.as_str(), "i32" | "f32" | "i64")) + && matches!(record.dialect.results.as_slice(), [result] + if result == &record.dialect.operands[1]) + && record.warp_shuffle.as_ref().is_some_and(|shuffle| { + match shuffle.value_kind { + WarpShuffleValueKind::I32 | WarpShuffleValueKind::F32 => { + let (value_ty, rust_ty) = match shuffle.value_kind { + WarpShuffleValueKind::I32 => ("i32", "u32"), + WarpShuffleValueKind::F32 => ("f32", "f32"), + WarpShuffleValueKind::I64 => unreachable!(), + }; + shuffle.adapter + == WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp + && shuffle.lane_encoding + == WarpShuffleOperandEncoding::RegisterOrImmediate + && shuffle.mask_encoding + == WarpShuffleOperandEncoding::RegisterOrImmediate + && record.lowering == "generated_warp_shuffle" + && record.rust.result == rust_ty + && record.dialect.operands[1] == value_ty + && record.llvm.as_ref().is_some_and(|llvm| { + matches!(llvm.arguments.as_slice(), [mask, value, lane, clamp] + if mask == "i32" + && value == value_ty + && lane == "i32" + && clamp == "i32") + && llvm.results == [value_ty] + }) + } + WarpShuffleValueKind::I64 => { + shuffle.adapter + == WarpShuffleAdapter::MaskValueLaneOrDeltaSplitI64LowHighB32InsertClampReassemble + && shuffle.lane_encoding + == WarpShuffleOperandEncoding::RegisterOnly + && shuffle.mask_encoding + == WarpShuffleOperandEncoding::RegisterOnly + && record.rust.result == "u64" + && record.dialect.operands[1] == "i64" + && record.lowering + == "generated_warp_shuffle_i64_inline_ptx" + && record.llvm.is_none() + && matches!(record.source, IntrinsicSource::PtxNative { .. }) + } + } + }), + "{} is outside the closed generated shfl.sync recipe", + record.id + ), + "integer_minmax" => ensure!( + record.integer_minmax.as_ref().is_some_and(|minmax| { + let module = match minmax.format { + IntegerMinMaxFormat::S32 => "int", + IntegerMinMaxFormat::S16x2 | IntegerMinMaxFormat::U16x2 => "i16x2", + }; + record.rust.module == module + }) + && record.rust.arguments.len() == 2 + && record + .rust + .arguments + .iter() + .all(|argument| argument == &record.rust.result) + && matches!(record.rust.result.as_str(), "i32" | "u32") + && record.rust.safe + && record.rust.must_use + && record.dialect.operands == ["i32", "i32"] + && record.dialect.results == ["i32"] + && record.lowering == "generated_integer_minmax_inline_ptx", + "{} is outside the closed generated integer-min/max recipe", + record.id + ), + family => ensure!(false, "{} has unrenderable family {family}", record.id), + }; + } + Ok(()) +} From 37b510f53e859b58f791c17db0d142f0ad6055d9 Mon Sep 17 00:00:00 2001 From: nihalpasham Date: Fri, 28 Aug 2026 22:22:20 +0530 Subject: [PATCH 10/18] refactor(intrinsics-gen): split resolve.rs into resolve/ modules --- crates/cuda-intrinsics-gen/src/resolve.rs | 42000 ---------------- .../src/resolve/abi_ledger.rs | 231 + .../cuda-intrinsics-gen/src/resolve/driver.rs | 765 + .../src/resolve/evidence.rs | 1504 + .../src/resolve/families/clc.rs | 472 + .../src/resolve/families/cluster.rs | 791 + .../src/resolve/families/cp_async.rs | 1121 + .../src/resolve/families/debug_wgmma.rs | 708 + .../src/resolve/families/execution_control.rs | 633 + .../src/resolve/families/matrix_move.rs | 776 + .../src/resolve/families/mbarrier_extended.rs | 835 + .../src/resolve/families/minmax.rs | 891 + .../src/resolve/families/mma/mod.rs | 14 + .../src/resolve/families/mma/register.rs | 1043 + .../families/mma/register_admissions.rs | 1454 + .../src/resolve/families/mma/sparse.rs | 907 + .../src/resolve/families/mma/validate.rs | 736 + .../src/resolve/families/mod.rs | 48 + .../src/resolve/families/packed_alu.rs | 1158 + .../src/resolve/families/packed_conversion.rs | 928 + .../src/resolve/families/prmt.rs | 410 + .../src/resolve/families/redux_dotprod.rs | 627 + .../src/resolve/families/scalar_arithmetic.rs | 661 + .../src/resolve/families/scalar_conversion.rs | 490 + .../src/resolve/families/scalar_math.rs | 692 + .../src/resolve/families/special_register.rs | 1459 + .../src/resolve/families/tcgen05/admission.rs | 850 + .../src/resolve/families/tcgen05/ldst_cp.rs | 1081 + .../src/resolve/families/tcgen05/mma.rs | 937 + .../src/resolve/families/tcgen05/mod.rs | 14 + .../src/resolve/families/tcgen05/recipes.rs | 914 + .../src/resolve/families/threadfence.rs | 363 + .../src/resolve/families/tma.rs | 1930 + .../src/resolve/families/warp.rs | 1702 + .../cuda-intrinsics-gen/src/resolve/guards.rs | 602 + .../src/resolve/materialize.rs | 490 + crates/cuda-intrinsics-gen/src/resolve/mod.rs | 32 + .../src/resolve/overlay.rs | 858 + .../cuda-intrinsics-gen/src/resolve/policy.rs | 612 + .../src/resolve/targets.rs | 852 + .../src/resolve/tests/evidence.rs | 1129 + .../src/resolve/tests/fixtures.rs | 1904 + .../src/resolve/tests/misc_families.rs | 1123 + .../src/resolve/tests/mma.rs | 2064 + .../src/resolve/tests/mod.rs | 16 + .../src/resolve/tests/overlay_schema.rs | 736 + .../src/resolve/tests/scalar_packed.rs | 1041 + .../src/resolve/tests/sreg_threadfence.rs | 486 + .../src/resolve/tests/targets_abi.rs | 813 + .../src/resolve/tests/tcgen05.rs | 1366 + .../src/resolve/tests/tma_clc_cluster.rs | 680 + .../src/resolve/tests/warp.rs | 1008 + 52 files changed, 42957 insertions(+), 42000 deletions(-) delete mode 100644 crates/cuda-intrinsics-gen/src/resolve.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/abi_ledger.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/driver.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/evidence.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/clc.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/cluster.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/cp_async.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/debug_wgmma.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/execution_control.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/matrix_move.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/mbarrier_extended.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/minmax.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/mma/mod.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/mma/register.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/mma/register_admissions.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/mma/sparse.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/mma/validate.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/mod.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/packed_alu.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/packed_conversion.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/prmt.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/redux_dotprod.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/scalar_arithmetic.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/scalar_conversion.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/scalar_math.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/special_register.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/tcgen05/admission.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/tcgen05/ldst_cp.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/tcgen05/mma.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/tcgen05/mod.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/tcgen05/recipes.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/threadfence.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/tma.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/families/warp.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/guards.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/materialize.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/mod.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/overlay.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/policy.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/targets.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/tests/evidence.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/tests/fixtures.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/tests/misc_families.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/tests/mma.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/tests/mod.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/tests/overlay_schema.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/tests/scalar_packed.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/tests/sreg_threadfence.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/tests/targets_abi.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/tests/tcgen05.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/tests/tma_clc_cluster.rs create mode 100644 crates/cuda-intrinsics-gen/src/resolve/tests/warp.rs diff --git a/crates/cuda-intrinsics-gen/src/resolve.rs b/crates/cuda-intrinsics-gen/src/resolve.rs deleted file mode 100644 index 9ffeb4a5a2..0000000000 --- a/crates/cuda-intrinsics-gen/src/resolve.rs +++ /dev/null @@ -1,42000 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -use crate::extract::{IMPORTED_SCHEMA, read_upstream_lock}; -use crate::model::{ - AbiLedgerEntry, AbiLedgerFile, AbiRawRustSignature, ActiveMaskAdapter, ActiveMaskObservation, - BackendLoweringMechanism, CatalogBackend, CatalogBackendLowering, CatalogDialect, CatalogFile, - CatalogHalfOpenRange, CatalogHardwareAlternative, CatalogHardwareTarget, CatalogInputs, - CatalogIntrinsic, CatalogLdmatrix, CatalogLlvm, CatalogLlvmResultFacts, CatalogRust, - CatalogSelection, CatalogSemantics, CatalogSource, CatalogTarget, CatalogTargetAlternative, - CatalogTargetContract, CatalogTargetRequirement, Clc, ClcAdapter, ClcAdmission, ClcOperation, - ClusterBarrier, ClusterBarrierAdmission, ClusterBarrierMode, ClusterBarrierOrdering, - ClusterMemory, ClusterMemoryAdapter, ClusterMemoryAdmission, ClusterMemoryOperation, - ClusterMemorySourceContract, ClusterSregAdmission, CpAsyncAdapter, CpAsyncCachePolicy, - CpAsyncControlAdapter, CpAsyncControlOperation, CpAsyncCopySize, CpAsyncMbarrierAdapter, - CpAsyncMbarrierOperation, CpAsyncMbarrierStateSpace, CpAsyncSourceSize, DebugControl, - DebugControlAdapter, DebugControlAdmission, DebugControlOperation, DotProductAdapter, - DotProductOperation, DotProductSignedness, EvidenceArtifactKind, EvidenceFile, EvidenceFileV6, - EvidenceMatrix, EvidenceMatrixTemplate, EvidenceRecord, EvidenceRecordDefaults, EvidenceStage, - EvidenceStageKind, ExecutionControlOperation, ExtendedMinMax, ExtendedMinMaxAdapter, - ExtendedMinMaxAdmission, ExtendedMinMaxFormat, ExtendedMinMaxNan, ExtendedMinMaxOperation, - ExtendedMinMaxSubnormal, ImportedAddressSpace, ImportedFile, ImportedIntrinsic, - IntegerMinMaxFormat, IntegerMinMaxOperation, IntrinsicBackend, IntrinsicSource, - LdmatrixAdapter, LdmatrixAddressContract, LdmatrixElement, LdmatrixLayout, LdmatrixMemoryOrder, - LdmatrixMultiplicity, LdmatrixParticipation, LdmatrixShape, LdmatrixStateSpace, MaskEncoding, - MatchOperandEncoding, MbarrierBasicAdapter, MbarrierBasicOperation, MbarrierExtended, - MbarrierExtendedAdapter, MbarrierExtendedAdmission, MbarrierExtendedOperation, - MbarrierExtendedSourceContract, MbarrierStateSpace, MovmatrixAdapter, MovmatrixParticipation, - OverlayBackendLowering, OverlayFile, OverlayIntrinsic, OverlayShardFile, PackedAluAdapter, - PackedAluFormat, PackedAluOperation, PackedAtomicAccessContract, PackedAtomicAdapter, - PackedAtomicAtomicity, PackedAtomicCodegenContract, PackedAtomicFormat, PackedAtomicOperation, - PackedAtomicOrdering, PackedAtomicPointerContract, PackedAtomicReturnContract, - PackedAtomicRounding, PackedAtomicScope, PackedAtomicScopeContract, PackedAtomicStateSpace, - PackedAtomicSubnormal, PackedConversionAdapter, PackedConversionDestinationFormat, - PackedConversionFp8Admission, PackedConversionFp8Direction, PackedConversionFp8F16x2Admission, - PackedConversionFp8Format, PackedConversionRounding, PackedConversionSaturation, - PackedConversionSourceFormat, PreSm70MemberMaskRule, Prmt, PrmtAdapter, PrmtAdmission, - PrmtMode, PtxVersion, ReduxAdapter, ReduxOperation, ReduxParticipation, RegisterMma, - RegisterMmaAccumulator, RegisterMmaAdapter, RegisterMmaAmpereFloatAdmission, - RegisterMmaAmpereFloatVariant, RegisterMmaBinaryAdmission, RegisterMmaCompatibilitySource, - RegisterMmaElement, RegisterMmaF8F6F4Admission, RegisterMmaFp8Admission, - RegisterMmaIntegerAdmission, RegisterMmaKind, RegisterMmaLayout, RegisterMmaOperation, - RegisterMmaOverflow, RegisterMmaParticipation, RegisterMmaShape, RuntimeValidation, - ScalarArithmetic, ScalarArithmeticAdmission, ScalarArithmeticFormat, ScalarArithmeticOperation, - ScalarArithmeticRounding, ScalarArithmeticSaturation, ScalarArithmeticSubnormal, - ScalarConversion, ScalarConversionAdapter, ScalarConversionAdmission, - ScalarConversionDestinationFormat, ScalarConversionResultRepresentation, - ScalarConversionRounding, ScalarConversionSaturation, ScalarConversionSourceFormat, ScalarMath, - ScalarMathAdmission, ScalarMathFormat, ScalarMathOperation, ScalarMathPrecision, - ScalarMathSubnormal, SparseMma, SparseMmaAccumulator, SparseMmaAdapter, - SparseMmaCompatibilitySource, SparseMmaElement, SparseMmaF8F6F4Admission, - SparseMmaF8F6F4F16Admission, SparseMmaIntegerAdmission, SparseMmaLayout, SparseMmaLlvmAdapter, - SparseMmaMetadata, SparseMmaOverflow, SparseMmaParticipation, SparseMmaSelector, - SparseMmaShape, SpecialRegister, SpecialRegisterAdmission, SpecialRegisterKind, - SpecialRegisterLlvmExclusion, SpecialRegisterLlvmExclusionReason, SpecialRegisterObservation, - SpecialRegisterOutputConstraint, SpecialRegisterPtxType, SpecialRegisterWidth, - StmatrixAdmission, StmatrixLayout, StmatrixMultiplicity, TargetContract, TargetSelectorBinding, - Tcgen05, Tcgen05Adapter, Tcgen05Admission, Tcgen05Cp, Tcgen05CpAdmissionVariant, - Tcgen05CpGroup, Tcgen05CpMember, Tcgen05Ld, Tcgen05LdAdmissionVariant, Tcgen05LdMultiplicity, - Tcgen05LdShape, Tcgen05Mma, Tcgen05MmaAdmissionVariant, Tcgen05MmaAlias, Tcgen05MmaBUsage, - Tcgen05MmaFixedSelectors, Tcgen05MmaForm, Tcgen05MmaKind, Tcgen05MmaSelectorLayout, - Tcgen05Operation, Tcgen05SourceContract, Tcgen05St, Tcgen05StAdmissionVariant, - ThreadfenceAdmission, ThreadfenceScope, Tma, TmaAdapter, TmaAdmission, TmaOperation, - TmaReduction, TmaReductionAdmissionVariant, TmaReductionLoadMode, TmaReductionOperation, - VoteAdapter, VoteMode, VoteParticipation, WarpBarrierAdapter, WarpBarrierMaskEncoding, - WarpBarrierMemoryOrdering, WarpBarrierParticipation, WarpMatchAdapter, WarpMatchMode, - WarpMatchParticipation, WarpMatchValueWidth, WarpShuffleAdapter, WarpShuffleMode, - WarpShuffleOperandEncoding, WarpShuffleParticipation, WarpShuffleSourceLane, - WarpShuffleValueKind, WgmmaControl, WgmmaControlAdapter, WgmmaControlAdmission, - WgmmaControlMode, WgmmaControlParticipation, -}; -use crate::ptx::{InstructionPattern, OperandPattern}; -use crate::util::{read_json, sha256_bytes, sha256_file}; -use anyhow::{Context, Result, bail, ensure}; -use cuda_target_spec::{CudaArch, recorded_ptx_floor}; -use ptx_parse::ParseError; -use std::collections::{BTreeMap, BTreeSet}; -use std::fs; -use std::path::{Component, Path, PathBuf}; - -const OVERLAY_SCHEMA: u32 = 44; -const MINIMUM_OVERLAY_SHARD_SCHEMA: u32 = 26; -const OVERLAY_SHARD_SCHEMA: u32 = 62; -const REGISTER_MMA_F8F6F4_SHARD_SCHEMA: u32 = 46; -const REGISTER_MMA_F8F6F4_F16_SHARD_SCHEMA: u32 = 47; -const REGISTER_MMA_MXF8F6F4_SHARD_SCHEMA: u32 = 60; -const REGISTER_MMA_FP8_SHARD_SCHEMA: u32 = 48; -const REGISTER_MMA_AMPERE_FLOAT_SHARD_SCHEMA: u32 = 49; -const SPARSE_MMA_F8F6F4_SHARD_SCHEMA: u32 = 27; -const SPARSE_MMA_F8F6F4_F16_SHARD_SCHEMA: u32 = 50; -const PRMT_SHARD_SCHEMA: u32 = 28; -const PACKED_CONVERSION_FP8_SHARD_SCHEMA: u32 = 29; -const PACKED_CONVERSION_FP8_F16X2_SHARD_SCHEMA: u32 = 59; -const CLUSTER_SREG_SHARD_SCHEMA: u32 = 30; -const CLUSTER_BARRIER_SHARD_SCHEMA: u32 = 31; -const SPECIAL_REGISTER_SHARD_SCHEMA: u32 = 32; -const DEBUG_CONTROL_SHARD_SCHEMA: u32 = 33; -const THREADFENCE_SHARD_SCHEMA: u32 = 34; -const STMATRIX_SHARD_SCHEMA: u32 = 35; -const CLUSTER_MEMORY_SHARD_SCHEMA: u32 = 39; -const CLC_SHARD_SCHEMA: u32 = 40; -const TMA_SHARD_SCHEMA: u32 = 61; -const TMA_REDUCTION_SHARD_SCHEMA: u32 = 62; -const MBARRIER_EXTENDED_SHARD_SCHEMA: u32 = 40; -const WGMMA_CONTROL_SHARD_SCHEMA: u32 = 38; -const TCGEN05_SHARD_SCHEMA: u32 = 42; -const SCALAR_CONVERSION_SHARD_SCHEMA: u32 = 43; -const SCALAR_ARITHMETIC_SHARD_SCHEMA: u32 = 45; -const EXTENDED_MINMAX_SHARD_SCHEMA: u32 = 60; -const TCGEN05_CP_SHARD_SCHEMA: u32 = 52; -const TCGEN05_LD_SHARD_SCHEMA: u32 = 53; -const TCGEN05_ST_SHARD_SCHEMA: u32 = 54; -const TCGEN05_OFFSET_LDST_SHARD_SCHEMA: u32 = 55; -const TCGEN05_CONTROL_SHARD_SCHEMA: u32 = 56; -const TCGEN05_MMA_SHARD_SCHEMA: u32 = 57; -const SCALAR_MATH_SHARD_SCHEMA: u32 = 58; -pub(crate) const CATALOG_SCHEMA: u32 = 46; -const BLACKWELL_LDMATRIX_LLVM_TARGETS: &str = - "sm_100a|sm_100f|sm_103a|sm_103f|sm_110a|sm_110f|sm_120a|sm_120f|sm_121a|sm_121f"; -const BLACKWELL_LDMATRIX_LIBNVVM_TARGETS: &str = BLACKWELL_LDMATRIX_LLVM_TARGETS; -const TMA_BLACKWELL_TARGETS: &str = "sm_100a|sm_101a|sm_103a|sm_110a"; -const TENSOR_MAP_REPLACE_TARGETS: &str = - "sm_100a|sm_100f|sm_103a|sm_103f|sm_110a|sm_110f|sm_120a|sm_120f|sm_121a|sm_121f|sm_90a"; -const TCGEN05_LLVM_TARGETS: &str = "sm_100a|sm_101a|sm_103a|sm_110a"; -const TCGEN05_LIBNVVM_TARGETS: &str = "sm_100a|sm_103a|sm_110a"; -const TCGEN05_MMA_FORMS: [Tcgen05MmaForm; 14] = [ - Tcgen05MmaForm::Shared, - Tcgen05MmaForm::Tensor, - Tcgen05MmaForm::TensorAshift, - Tcgen05MmaForm::SpShared, - Tcgen05MmaForm::SpTensor, - Tcgen05MmaForm::SpTensorAshift, - Tcgen05MmaForm::WsShared, - Tcgen05MmaForm::WsSharedZeroColMask, - Tcgen05MmaForm::WsSpShared, - Tcgen05MmaForm::WsSpSharedZeroColMask, - Tcgen05MmaForm::WsSpTensor, - Tcgen05MmaForm::WsSpTensorZeroColMask, - Tcgen05MmaForm::WsTensor, - Tcgen05MmaForm::WsTensorZeroColMask, -]; -const TCGEN05_MMA_ALIASES: [Tcgen05MmaAlias; 5] = [ - Tcgen05MmaAlias::E4m3, - Tcgen05MmaAlias::E5m2, - Tcgen05MmaAlias::E2m3, - Tcgen05MmaAlias::E3m2, - Tcgen05MmaAlias::E2m1, -]; -const TCGEN05_MMA_KINDS: [Tcgen05MmaKind; 4] = [ - Tcgen05MmaKind::F16, - Tcgen05MmaKind::F8f6f4, - Tcgen05MmaKind::I8, - Tcgen05MmaKind::Tf32, -]; -const TCGEN05_MMA_DIALECT_OP_TYPE: &str = "Tcgen05MmaOp"; -const TCGEN05_MMA_DIALECT_OP_NAME: &str = "nvvm.tcgen05_mma"; -const REGISTER_MMA_F8F6F4_TARGETS: &str = "sm_120a|sm_120f|sm_121a|sm_121f"; -const SPARSE_MMA_F8F6F4_TARGETS: &str = "sm_120a|sm_120f|sm_121a|sm_121f"; - -struct ResolutionBase { - overlay: OverlayFile, - imported: ImportedFile, - source: CatalogSource, - imported_sha256: String, - overlay_sha256: String, - abi_ledger_sha256: String, -} - -#[derive(Debug)] -pub(crate) struct CandidateResolution { - pub catalog: CatalogFile, - pub mechanism: BackendLoweringMechanism, - pub requirement: CatalogTargetRequirement, -} - -fn primary_evidence_profile<'a>( - policy: &'a OverlayIntrinsic, - default_profile: &'a str, -) -> Result<&'a str> { - let needs_family_profile = policy.tma.is_some() - || ExecutionControlOperation::from_catalog_id(&policy.id).is_some() - || policy - .tcgen05 - .as_ref() - .and_then(|tcgen05| tcgen05.mma.as_ref()) - .is_some(); - if !needs_family_profile { - return Ok(default_profile); - } - policy - .backend_lowerings - .iter() - .find(|lowering| lowering.backend == IntrinsicBackend::LlvmNvptx) - .map(|lowering| lowering.evidence_profile.as_str()) - .with_context(|| format!("{} has no LLVM family evidence route", policy.id)) -} - -pub fn resolve(repo_root: &Path) -> Result { - let base = load_resolution_base(repo_root)?; - let ResolutionBase { - overlay, - imported, - source, - imported_sha256, - overlay_sha256, - abi_ledger_sha256, - } = base; - let imported_by_record = index_imported_intrinsics(&imported)?; - let llvm_revision = source.llvm_revision.clone(); - let (evidence_files, evidence_hashes) = read_evidence(repo_root)?; - let evidence_by_profile_id = index_evidence(&evidence_files, &llvm_revision)?; - - let mut intrinsics = Vec::with_capacity(overlay.intrinsics.len()); - for policy in &overlay.intrinsics { - let source = resolve_policy_source(policy)?; - let declaration = resolve_imported_declaration(policy, &source, &imported_by_record)?; - validate_special_register_llvm_exclusion(policy, &imported_by_record)?; - validate_policy(policy, &source, declaration, overlay.intrinsic_abi)?; - let primary_profile = primary_evidence_profile(policy, &overlay.backend_profile)?; - let evidence = evidence_by_profile_id - .get(&(primary_profile, policy.id.as_str())) - .with_context(|| { - format!( - "intrinsic {} has no evidence record in selected profile {}", - policy.id, primary_profile - ) - })?; - validate_evidence(policy, evidence, None)?; - let backend_lowerings = resolve_backend_lowerings(policy, &evidence_by_profile_id)?; - intrinsics.push(resolve_record( - policy, - source, - declaration, - evidence.record, - primary_profile, - evidence.backend_version, - evidence.backend_sha256, - backend_lowerings, - overlay.intrinsic_abi, - )?); - } - for (_, evidence_id) in evidence_by_profile_id.keys() { - ensure!( - overlay - .intrinsics - .iter() - .any(|record| record.id == *evidence_id), - "evidence exists for unknown catalog ID {evidence_id}" - ); - } - - Ok(CatalogFile { - schema: CATALOG_SCHEMA, - catalog_version: overlay.catalog_version, - intrinsic_abi: overlay.intrinsic_abi, - generator_version: env!("CARGO_PKG_VERSION").to_owned(), - source, - inputs: CatalogInputs { - imported_sha256, - overlay_sha256, - abi_ledger_sha256, - evidence_sha256: evidence_hashes, - }, - intrinsics, - }) -} - -#[cfg(test)] -pub(crate) fn test_catalog_with_clc(repo_root: &Path) -> Result { - let mut catalog = resolve(repo_root)?; - let active_count = catalog - .intrinsics - .iter() - .filter(|record| record.family == "clc") - .count(); - if active_count != 0 { - ensure!( - active_count == 6, - "active CLC family must contain six records" - ); - return Ok(catalog); - } - let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json"))?; - let imported_by_record = index_imported_intrinsics(&imported)?; - let operations = [ - ClcOperation::TryCancel, - ClcOperation::TryCancelMulticast, - ClcOperation::QueryIsCanceled, - ClcOperation::QueryGetFirstCtaidX, - ClcOperation::QueryGetFirstCtaidY, - ClcOperation::QueryGetFirstCtaidZ, - ]; - let admission = ClcAdmission { - llvm_evidence_profile: "llvm-clc-test".into(), - libnvvm_evidence_profile: "libnvvm-clc-test".into(), - runtime_validation: RuntimeValidation::Unexecuted, - variants: operations - .into_iter() - .map(|operation| crate::model::ClcAdmissionVariant { - abi_id: clc_recipe(operation).abi_id.into(), - operation, - }) - .collect(), - }; - for policy in expand_clc_admission(&admission)? { - let source = resolve_policy_source(&policy)?; - let declaration = resolve_imported_declaration(&policy, &source, &imported_by_record)?; - validate_policy(&policy, &source, declaration, catalog.intrinsic_abi)?; - let backend_lowerings = policy - .backend_lowerings - .iter() - .map(|lowering| { - Ok(CatalogBackendLowering { - backend: lowering.backend, - mechanism: lowering.mechanism, - evidence_profile: lowering.evidence_profile.clone(), - target: backend_target_requirement(&policy, lowering)?, - version: "test".into(), - sha256: "0".repeat(64), - artifact_path: None, - build_id_prefix: None, - status: "validated".into(), - stages: vec![], - }) - }) - .collect::>>()?; - catalog.intrinsics.push(materialize_record( - &policy, - source, - declaration, - CatalogBackend { - profile: "clc-test".into(), - version: "test".into(), - sha256: "0".repeat(64), - status: "validated".into(), - target_triple: "nvptx64-nvidia-cuda".into(), - gpu_target: "sm_100".into(), - ptx_feature: "+ptx86".into(), - }, - backend_lowerings, - catalog.intrinsic_abi, - )?); - } - catalog - .intrinsics - .sort_by(|left, right| left.id.cmp(&right.id)); - Ok(catalog) -} - -#[cfg(test)] -pub(crate) fn test_catalog_with_tma(repo_root: &Path) -> Result { - let mut catalog = resolve(repo_root)?; - let active_count = catalog - .intrinsics - .iter() - .filter(|record| record.tma.is_some()) - .count(); - let expected_count = TMA_OPERATIONS.len() + tma_reduction_matrix().len(); - if active_count == expected_count { - return Ok(catalog); - } - ensure!( - active_count == 0 || active_count == TMA_OPERATIONS.len(), - "active TMA catalog has {active_count} records; expected 0, {} base records, or {expected_count} total records", - TMA_OPERATIONS.len() - ); - let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json"))?; - let imported_by_record = index_imported_intrinsics(&imported)?; - let admission = TmaAdmission { - llvm_evidence_profile: "llvm-tma-test".into(), - libnvvm_evidence_profile: "libnvvm-tma-test".into(), - reduce_llvm_evidence_profile: Some("llvm-tma-reduce-test".into()), - reduce_libnvvm_evidence_profile: Some("libnvvm-tma-reduce-test".into()), - runtime_validation: RuntimeValidation::Unexecuted, - variants: TMA_OPERATIONS - .into_iter() - .map(|operation| crate::model::TmaAdmissionVariant { - abi_id: tma_recipe(operation).abi_id.into(), - operation, - }) - .collect(), - reduce_variants: tma_reduction_admission_variants(), - }; - for policy in expand_tma_admission(&admission)? { - if catalog - .intrinsics - .iter() - .any(|record| record.id.as_str() == policy.id.as_str()) - { - continue; - } - let source = resolve_policy_source(&policy)?; - let declaration = resolve_imported_declaration(&policy, &source, &imported_by_record)?; - validate_policy(&policy, &source, declaration, catalog.intrinsic_abi)?; - let backend_lowerings = policy - .backend_lowerings - .iter() - .map(|lowering| { - Ok(CatalogBackendLowering { - backend: lowering.backend, - mechanism: lowering.mechanism, - evidence_profile: lowering.evidence_profile.clone(), - target: backend_target_requirement(&policy, lowering)?, - version: "test".into(), - sha256: "0".repeat(64), - artifact_path: None, - build_id_prefix: None, - status: "validated".into(), - stages: vec![], - }) - }) - .collect::>>()?; - catalog.intrinsics.push(materialize_record( - &policy, - source, - declaration, - CatalogBackend { - profile: "tma-test".into(), - version: "test".into(), - sha256: "0".repeat(64), - status: "validated".into(), - target_triple: "nvptx64-nvidia-cuda".into(), - gpu_target: "sm_90".into(), - ptx_feature: "+ptx80".into(), - }, - backend_lowerings, - catalog.intrinsic_abi, - )?); - } - catalog - .intrinsics - .sort_by(|left, right| left.id.cmp(&right.id)); - Ok(catalog) -} - -#[cfg(test)] -pub(crate) fn test_catalog_with_tcgen05(repo_root: &Path) -> Result { - let mut catalog = resolve(repo_root)?; - let active_count = catalog - .intrinsics - .iter() - .filter(|record| record.family == "tcgen05") - .count(); - match active_count { - 233 => return Ok(catalog), - 0 => {} - count => bail!("active tcgen05 catalog has {count} of 233 records"), - } - let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json"))?; - let imported_by_record = index_imported_intrinsics(&imported)?; - let operations = [ - Tcgen05Operation::Alloc, - Tcgen05Operation::Dealloc, - Tcgen05Operation::RelinquishAllocPermit, - Tcgen05Operation::FenceBeforeThreadSync, - Tcgen05Operation::FenceAfterThreadSync, - Tcgen05Operation::Commit, - Tcgen05Operation::CommitSharedCluster, - Tcgen05Operation::MmaWsF16, - Tcgen05Operation::MmaF16, - Tcgen05Operation::MmaWsBf16, - Tcgen05Operation::MmaWsTf32, - Tcgen05Operation::CpSmemToTmem, - Tcgen05Operation::Ld16x256bX8Pure, - Tcgen05Operation::Ld16x256bPure, - Tcgen05Operation::LoadWait, - Tcgen05Operation::StoreWait, - Tcgen05Operation::AllocCg2, - Tcgen05Operation::DeallocCg2, - Tcgen05Operation::RelinquishAllocPermitCg2, - Tcgen05Operation::MmaF16Cg2, - Tcgen05Operation::CommitCg2, - Tcgen05Operation::CommitSharedClusterCg2, - Tcgen05Operation::CommitMulticastCg2, - Tcgen05Operation::CpSmemToTmemCg2, - Tcgen05Operation::CommitMulticast, - Tcgen05Operation::ShiftDown, - Tcgen05Operation::ShiftDownCg2, - ]; - let admission = Tcgen05Admission { - llvm_evidence_profile: "llvm-tcgen05-test".into(), - libnvvm_evidence_profile: "libnvvm-tcgen05-test".into(), - cp_llvm_evidence_profile: None, - cp_libnvvm_evidence_profile: None, - ld_llvm_evidence_profile: None, - ld_libnvvm_evidence_profile: None, - st_llvm_evidence_profile: None, - st_libnvvm_evidence_profile: None, - offset_llvm_evidence_profile: Some("llvm-tcgen05-offset-test".into()), - offset_libnvvm_evidence_profile: Some("libnvvm-tcgen05-offset-test".into()), - control_llvm_evidence_profile: Some("llvm-tcgen05-control-test".into()), - control_libnvvm_evidence_profile: Some("libnvvm-tcgen05-control-test".into()), - mma_llvm_evidence_profile: None, - mma_libnvvm_evidence_profile: None, - mma_llvm_target_contracts: vec![], - mma_libnvvm_target_contracts: vec![], - runtime_validation: RuntimeValidation::Unexecuted, - variants: operations - .into_iter() - .map(|operation| crate::model::Tcgen05AdmissionVariant { - abi_id: tcgen05_recipe(operation).abi_id.into(), - operation, - }) - .collect(), - cp_variants: vec![], - ld_variants: vec![], - st_variants: vec![], - ld_offset_variants: TCGEN05_OFFSET_LDST_VARIANTS - .into_iter() - .flat_map(|(shape, multiplicity)| { - [false, true] - .into_iter() - .map(move |pack16| (shape, multiplicity, pack16)) - }) - .enumerate() - .map( - |(index, (shape, multiplicity, pack16))| Tcgen05LdAdmissionVariant { - abi_id: format!("i{:04}", 728 + index), - shape, - multiplicity, - pack16, - }, - ) - .collect(), - st_offset_variants: TCGEN05_OFFSET_LDST_VARIANTS - .into_iter() - .flat_map(|(shape, multiplicity)| { - [false, true] - .into_iter() - .map(move |unpack16| (shape, multiplicity, unpack16)) - }) - .enumerate() - .map( - |(index, (shape, multiplicity, unpack16))| Tcgen05StAdmissionVariant { - abi_id: format!("i{:04}", 744 + index), - shape, - multiplicity, - unpack16, - }, - ) - .collect(), - mma_variants: vec![], - }; - for policy in expand_tcgen05_admission(&admission)? { - let source = resolve_policy_source(&policy)?; - let declaration = resolve_imported_declaration(&policy, &source, &imported_by_record)?; - validate_policy(&policy, &source, declaration, catalog.intrinsic_abi)?; - let backend_lowerings = policy - .backend_lowerings - .iter() - .map(|lowering| { - Ok(CatalogBackendLowering { - backend: lowering.backend, - mechanism: lowering.mechanism, - evidence_profile: lowering.evidence_profile.clone(), - target: backend_target_requirement(&policy, lowering)?, - version: "test".into(), - sha256: "0".repeat(64), - artifact_path: None, - build_id_prefix: None, - status: "validated".into(), - stages: vec![], - }) - }) - .collect::>>()?; - catalog.intrinsics.push(materialize_record( - &policy, - source, - declaration, - CatalogBackend { - profile: "tcgen05-test".into(), - version: "test".into(), - sha256: "0".repeat(64), - status: "validated".into(), - target_triple: "nvptx64-nvidia-cuda".into(), - gpu_target: "sm_100a".into(), - ptx_feature: "+ptx86".into(), - }, - backend_lowerings, - catalog.intrinsic_abi, - )?); - } - catalog - .intrinsics - .sort_by(|left, right| left.id.cmp(&right.id)); - Ok(catalog) -} - -fn load_resolution_base(repo_root: &Path) -> Result { - let lock = read_upstream_lock(repo_root)?; - let imported_path = repo_root.join("intrinsics/imported.json"); - let overlay_path = repo_root.join("intrinsics/overlay.toml"); - let imported: ImportedFile = read_json(&imported_path)?; - let (mut overlay, overlay_sha256) = read_overlay(repo_root, &overlay_path)?; - let ledger_path = repo_root.join(format!("intrinsics/abi-v{}.toml", overlay.intrinsic_abi)); - let ledger_text = fs::read_to_string(&ledger_path) - .with_context(|| format!("read {}", ledger_path.display()))?; - let ledger: AbiLedgerFile = - toml::from_str(&ledger_text).with_context(|| format!("parse {}", ledger_path.display()))?; - - ensure!( - imported.schema == IMPORTED_SCHEMA, - "unsupported imported.json schema {}", - imported.schema - ); - ensure!( - overlay.schema == OVERLAY_SCHEMA, - "unsupported overlay.toml schema {}", - overlay.schema - ); - ensure!( - overlay.intrinsic_abi > 0, - "intrinsic_abi must be a positive integer" - ); - ensure!( - imported.source.llvm_revision == lock.llvm.revision, - "imported facts use LLVM {}, but upstream.lock pins {}", - imported.source.llvm_revision, - lock.llvm.revision - ); - ensure!( - imported.source.llvm_tblgen_source_revision == lock.llvm.revision, - "imported facts were not produced by llvm-tblgen built from the pinned source" - ); - ensure!( - imported.source.llvm_tblgen_version == lock.llvm_tblgen.version_line, - "imported facts use llvm-tblgen {:?}, but upstream.lock pins {:?}", - imported.source.llvm_tblgen_version, - lock.llvm_tblgen.version_line - ); - ensure!( - imported.source.intrinsics_json_sha256 == lock.dumps.intrinsics_sha256, - "imported intrinsic dump hash does not match upstream.lock" - ); - ensure!( - imported.source.nvptx_json_sha256 == lock.dumps.nvptx_sha256, - "imported NVPTX dump hash does not match upstream.lock" - ); - let imported_sha256 = sha256_file(&imported_path)?; - ensure!( - imported_sha256 == lock.dumps.normalized_imported_sha256, - "normalized imported.json hash mismatch: upstream.lock records {}, found {}; regenerate from the pinned dumps, and refresh the lock explicitly only for an intentional normalizer change", - lock.dumps.normalized_imported_sha256, - imported_sha256 - ); - - bind_generated_abi_ids(&mut overlay, &ledger)?; - overlay - .intrinsics - .sort_by(|left, right| left.id.cmp(&right.id)); - validate_execution_control_family_completeness(&overlay.intrinsics)?; - validate_unique_overlay(&overlay.intrinsics, overlay.intrinsic_abi)?; - validate_abi_ledger(&overlay, &ledger)?; - Ok(ResolutionBase { - overlay, - imported, - source: CatalogSource { - llvm_repository: lock.llvm.repository, - llvm_revision: lock.llvm.revision, - llvm_tblgen_version: lock.llvm_tblgen.version_line, - llvm_tblgen_source_revision: lock - .llvm_tblgen - .built_from_llvm_revision - .context("pinned llvm-tblgen has no source revision")?, - }, - imported_sha256, - overlay_sha256, - abi_ledger_sha256: sha256_file(&ledger_path)?, - }) -} - -fn index_imported_intrinsics( - imported: &ImportedFile, -) -> Result> { - let imported_by_record: BTreeMap<_, _> = imported - .intrinsics - .iter() - .map(|intrinsic| (intrinsic.source_record.as_str(), intrinsic)) - .collect(); - ensure!( - imported_by_record.len() == imported.intrinsics.len(), - "imported.json contains duplicate source records" - ); - Ok(imported_by_record) -} - -fn resolve_imported_declaration<'a>( - policy: &OverlayIntrinsic, - source: &IntrinsicSource, - imported_by_record: &'a BTreeMap<&str, &'a ImportedIntrinsic>, -) -> Result> { - match source { - IntrinsicSource::LlvmImported { source_record } => Ok(Some( - *imported_by_record - .get(source_record.as_str()) - .with_context(|| { - format!( - "overlay intrinsic {} references missing imported record {}", - policy.id, source_record - ) - })?, - )), - IntrinsicSource::PtxNative { .. } => Ok(None), - } -} - -pub(crate) fn resolve_candidate( - repo_root: &Path, - intrinsic_id: &str, - backend_version: &str, - backend_sha256: &str, - gpu_target: &str, - ptx_feature: &str, -) -> Result { - let base = load_resolution_base(repo_root)?; - let imported_by_record = index_imported_intrinsics(&base.imported)?; - let policy = base - .overlay - .intrinsics - .iter() - .find(|policy| policy.id == intrinsic_id || policy.abi_id == intrinsic_id) - .with_context(|| format!("unknown overlay intrinsic {intrinsic_id}"))?; - let source = resolve_policy_source(policy)?; - let declaration = resolve_imported_declaration(policy, &source, &imported_by_record)?; - validate_special_register_llvm_exclusion(policy, &imported_by_record)?; - validate_policy(policy, &source, declaration, base.overlay.intrinsic_abi)?; - ensure!( - !backend_version.trim().is_empty() - && backend_sha256.len() == 64 - && backend_sha256.bytes().all(|byte| byte.is_ascii_hexdigit()), - "candidate backend identity is incomplete" - ); - let (mechanism, target) = candidate_llvm_route(policy)?; - validate_candidate_target(policy, &target, gpu_target, ptx_feature)?; - - let record = materialize_record( - policy, - source, - declaration, - CatalogBackend { - profile: "candidate-comparison".into(), - version: backend_version.into(), - sha256: backend_sha256.into(), - status: "candidate".into(), - target_triple: "nvptx64-nvidia-cuda".into(), - gpu_target: gpu_target.into(), - ptx_feature: ptx_feature.into(), - }, - Vec::new(), - base.overlay.intrinsic_abi, - )?; - Ok(CandidateResolution { - catalog: CatalogFile { - schema: CATALOG_SCHEMA, - catalog_version: base.overlay.catalog_version, - intrinsic_abi: base.overlay.intrinsic_abi, - generator_version: env!("CARGO_PKG_VERSION").to_owned(), - source: base.source, - inputs: CatalogInputs { - imported_sha256: base.imported_sha256, - overlay_sha256: base.overlay_sha256, - abi_ledger_sha256: base.abi_ledger_sha256, - evidence_sha256: Vec::new(), - }, - intrinsics: vec![record], - }, - mechanism, - requirement: target, - }) -} - -fn candidate_llvm_route( - policy: &OverlayIntrinsic, -) -> Result<(BackendLoweringMechanism, CatalogTargetRequirement)> { - let routes = policy - .backend_lowerings - .iter() - .filter(|lowering| lowering.backend == IntrinsicBackend::LlvmNvptx) - .collect::>(); - ensure!( - routes.len() <= 1, - "{} has more than one LLVM-NVPTX route", - policy.id - ); - if let Some(route) = routes.first() { - return Ok((route.mechanism, backend_target_requirement(policy, route)?)); - } - ensure!( - matches!( - resolve_policy_source(policy)?, - IntrinsicSource::LlvmImported { .. } - ), - "{} has no LLVM-NVPTX candidate route", - policy.id - ); - Ok(( - BackendLoweringMechanism::TypedNvvm, - CatalogTargetRequirement { - minimum_ptx: parse_ptx_version(&policy.minimum_ptx, &policy.id)?, - hardware: parse_hardware_target(policy)?, - }, - )) -} - -fn validate_candidate_target( - policy: &OverlayIntrinsic, - requirement: &CatalogTargetRequirement, - gpu_target: &str, - ptx_feature: &str, -) -> Result<()> { - ensure!( - gpu_target.starts_with("sm_"), - "candidate GPU target {gpu_target:?} must use sm_NN, sm_NNa, or sm_NNf" - ); - let hardware = parse_stage_hardware(gpu_target).with_context(|| { - format!("candidate GPU target {gpu_target:?} must use sm_NN, sm_NNa, or sm_NNf") - })?; - ensure!( - describe_stage_hardware(hardware) == gpu_target, - "candidate GPU target {gpu_target:?} is not canonical" - ); - let ptx = parse_candidate_ptx_feature(ptx_feature)?; - let effective_minimum_ptx = if is_f8f6f4_mma_target_matrix_policy(policy) { - f8f6f4_llvm_ptx_floor(hardware)? - } else { - target_requirement_ptx_floor(requirement, hardware, true).with_context(|| { - format!( - "candidate GPU target {gpu_target} does not satisfy {} hardware requirement {:?}", - policy.id, requirement.hardware - ) - })? - }; - ensure!( - ptx.encoded() >= effective_minimum_ptx, - "candidate target {gpu_target} / {ptx_feature} is below {} PTX floor {}.{}", - policy.id, - effective_minimum_ptx / 10, - effective_minimum_ptx % 10 - ); - let hardware_matches = target_requirement_ptx_floor(requirement, hardware, true).is_some(); - ensure!( - hardware_matches, - "candidate GPU target {gpu_target} does not satisfy {} hardware requirement {:?}", - policy.id, - requirement.hardware - ); - Ok(()) -} - -fn target_requirement_ptx_floor( - requirement: &CatalogTargetRequirement, - hardware: CatalogHardwareAlternative, - allow_forward_minimum: bool, -) -> Option { - match &requirement.hardware { - CatalogHardwareTarget::All => Some(requirement.minimum_ptx.encoded()), - CatalogHardwareTarget::AnyOf { alternatives } => alternatives - .iter() - .any(|expected| { - selected_stage_hardware_matches(hardware, *expected, allow_forward_minimum) - }) - .then(|| requirement.minimum_ptx.encoded()), - CatalogHardwareTarget::TargetMatrix { contracts } => contracts - .iter() - .flat_map(|contract| contract.alternatives.iter()) - .filter(|alternative| { - selected_stage_hardware_matches( - hardware, - alternative.hardware, - allow_forward_minimum, - ) - }) - .map(|alternative| alternative.minimum_ptx.encoded()) - .min(), - } -} - -fn parse_candidate_ptx_feature(value: &str) -> Result { - let digits = value - .strip_prefix("+ptx") - .with_context(|| format!("candidate PTX feature {value:?} must use +ptxNN"))?; - ensure!( - matches!(digits.len(), 2 | 3) && digits.bytes().all(|byte| byte.is_ascii_digit()), - "candidate PTX feature {value:?} must use +ptxNN" - ); - let split = digits.len() - 1; - let version = format!("{}.{}", &digits[..split], &digits[split..]); - version - .parse() - .map_err(|reason: String| anyhow::anyhow!("candidate PTX feature {value:?}: {reason}")) -} - -fn read_overlay(repo_root: &Path, manifest_path: &Path) -> Result<(OverlayFile, String)> { - let manifest_bytes = - fs::read(manifest_path).with_context(|| format!("read {}", manifest_path.display()))?; - let mut overlay: OverlayFile = toml::from_slice(&manifest_bytes) - .with_context(|| format!("parse {}", manifest_path.display()))?; - - ensure!( - overlay.intrinsics.is_empty(), - "overlay.toml must list family shards instead of inline intrinsic records" - ); - ensure!( - !overlay.shards.is_empty(), - "overlay.toml must list at least one family shard" - ); - - let mut previous = None; - let mut seen = BTreeSet::new(); - let mut hash_input = Vec::new(); - append_overlay_hash_input(&mut hash_input, "intrinsics/overlay.toml", &manifest_bytes); - - for shard_name in &overlay.shards { - validate_overlay_shard_path(shard_name)?; - ensure!( - seen.insert(shard_name.as_str()), - "overlay.toml lists duplicate shard {shard_name}" - ); - if let Some(previous) = previous { - ensure!( - previous < shard_name.as_str(), - "overlay.toml shards must be sorted" - ); - } - previous = Some(shard_name.as_str()); - - let relative = Path::new("intrinsics").join(shard_name); - let path = repo_root.join(&relative); - let bytes = fs::read(&path).with_context(|| format!("read {}", path.display()))?; - let mut shard: OverlayShardFile = - toml::from_slice(&bytes).with_context(|| format!("parse {}", path.display()))?; - validate_overlay_shard_schema(&shard, &path)?; - let int4_mma_admission = shard.register_mma_int4.take(); - let int8_mma_admission = shard.register_mma_int8.take(); - let binary_mma_admission = shard.register_mma_b1.take(); - let f8f6f4_f32_mma_admission = shard.register_mma_f8f6f4_f32.take(); - let f8f6f4_f16_mma_admission = shard.register_mma_f8f6f4_f16.take(); - let mxf8f6f4_f32_mma_admission = shard.register_mma_mxf8f6f4_f32.take(); - let fp8_mma_admission = shard.register_mma_fp8.take(); - let ampere_float_mma_admission = shard.register_mma_ampere_float.take(); - let sparse_mma_admission = shard.sparse_mma_integer.take(); - let sparse_mma_f8f6f4_admission = shard.sparse_mma_f8f6f4_f32.take(); - let sparse_mma_f8f6f4_f16_admission = shard.sparse_mma_f8f6f4_f16.take(); - let prmt_admission = shard.prmt.take(); - let packed_conversion_fp8_admission = shard.packed_conversion_fp8.take(); - let packed_conversion_fp8_f16x2_admission = shard.packed_conversion_fp8_f16x2.take(); - let scalar_conversion_admission = shard.scalar_conversion.take(); - let scalar_arithmetic_admission = shard.scalar_arithmetic.take(); - let scalar_math_admission = shard.scalar_math.take(); - let extended_minmax_admission = shard.extended_minmax.take(); - let cluster_sreg_admission = shard.cluster_sreg.take(); - let cluster_barrier_admission = shard.cluster_barrier.take(); - let mbarrier_extended_admission = shard.mbarrier_extended.take(); - let special_register_admission = shard.special_registers.take(); - let debug_control_admission = shard.debug_control.take(); - let threadfence_admission = shard.threadfence.take(); - let cluster_memory_admission = shard.cluster_memory.take(); - let stmatrix_admission = shard.stmatrix.take(); - let clc_admission = shard.clc.take(); - let wgmma_control_admission = shard.wgmma_controls.take(); - let tma_admission = shard.tma.take(); - let tcgen05_admission = shard.tcgen05.take(); - let compact_mma_count = usize::from(int4_mma_admission.is_some()) - + usize::from(int8_mma_admission.is_some()) - + usize::from(binary_mma_admission.is_some()) - + usize::from(f8f6f4_f32_mma_admission.is_some()) - + usize::from(f8f6f4_f16_mma_admission.is_some()) - + usize::from(mxf8f6f4_f32_mma_admission.is_some()) - + usize::from(fp8_mma_admission.is_some()) - + usize::from(ampere_float_mma_admission.is_some()) - + usize::from(sparse_mma_admission.is_some()) - + usize::from(sparse_mma_f8f6f4_admission.is_some()) - + usize::from(sparse_mma_f8f6f4_f16_admission.is_some()); - ensure!( - compact_mma_count <= 1, - "overlay shard {} contains more than one compact MMA admission", - path.display() - ); - let integer_mma_admission = int4_mma_admission - .map(|admission| (RegisterMmaIntegerKind::Int4, admission)) - .or_else(|| { - int8_mma_admission.map(|admission| (RegisterMmaIntegerKind::Int8, admission)) - }); - if let Some((kind, admission)) = integer_mma_admission { - ensure!( - shard.family == "register_mma" && shard.intrinsics.is_empty(), - "compact integer MMA admission must be the only content of a register_mma shard" - ); - shard.intrinsics = expand_register_mma_integer_admission(kind, &admission)?; - } - if let Some(admission) = binary_mma_admission { - ensure!( - shard.family == "register_mma" && shard.intrinsics.is_empty(), - "compact binary MMA admission must be the only content of a register_mma shard" - ); - shard.intrinsics = expand_register_mma_binary_admission(&admission)?; - } - if let Some(admission) = f8f6f4_f32_mma_admission { - ensure!( - shard.family == "register_mma" && shard.intrinsics.is_empty(), - "compact dense f8f6f4 MMA admission must be the only content of a register_mma shard" - ); - shard.intrinsics = - expand_register_mma_f8f6f4_admission(&admission, RegisterMmaAccumulator::F32)?; - } - if let Some(admission) = f8f6f4_f16_mma_admission { - ensure!( - shard.family == "register_mma" && shard.intrinsics.is_empty(), - "compact dense f8f6f4 MMA admission must be the only content of a register_mma shard" - ); - shard.intrinsics = - expand_register_mma_f8f6f4_admission(&admission, RegisterMmaAccumulator::F16)?; - } - if let Some(admission) = mxf8f6f4_f32_mma_admission { - ensure!( - shard.family == "register_mma" && shard.intrinsics.is_empty(), - "compact dense mxf8f6f4 MMA admission must be the only content of a register_mma shard" - ); - shard.intrinsics = expand_register_mma_mxf8f6f4_admission(&admission)?; - } - if let Some(admission) = fp8_mma_admission { - ensure!( - shard.family == "register_mma" && shard.intrinsics.is_empty(), - "compact standard FP8 MMA admission must be the only content of a register_mma shard" - ); - shard.intrinsics = expand_register_mma_fp8_admission(&admission)?; - } - if let Some(admission) = ampere_float_mma_admission { - ensure!( - shard.family == "register_mma" && shard.intrinsics.is_empty(), - "compact Ampere floating-point MMA admission must be the only content of a register_mma shard" - ); - shard.intrinsics = expand_register_mma_ampere_float_admission(&admission)?; - } - if let Some(admission) = sparse_mma_admission { - ensure!( - shard.family == "sparse_mma" && shard.intrinsics.is_empty(), - "compact sparse MMA admission must be the only content of a sparse_mma shard" - ); - shard.intrinsics = expand_sparse_mma_integer_admission(&admission)?; - } - if let Some(admission) = sparse_mma_f8f6f4_admission { - ensure!( - shard.family == "sparse_mma" && shard.intrinsics.is_empty(), - "compact sparse f8f6f4 MMA admission must be the only content of a sparse_mma shard" - ); - shard.intrinsics = expand_sparse_mma_f8f6f4_admission(&admission)?; - } - if let Some(admission) = sparse_mma_f8f6f4_f16_admission { - ensure!( - shard.family == "sparse_mma" && shard.intrinsics.is_empty(), - "compact sparse f8f6f4 F16 MMA admission must be the only content of a sparse_mma shard" - ); - shard.intrinsics = expand_sparse_mma_f8f6f4_f16_admission(&admission)?; - } - if let Some(admission) = prmt_admission { - ensure!( - shard.family == "prmt" && shard.intrinsics.is_empty(), - "compact prmt admission must be the only content of a prmt shard" - ); - shard.intrinsics = expand_prmt_admission(&admission)?; - } - if let Some(admission) = packed_conversion_fp8_admission { - ensure!( - shard.family == "packed_conversion" && shard.intrinsics.is_empty(), - "compact FP8 conversion admission must be the only content of a packed_conversion shard" - ); - shard.intrinsics = expand_packed_conversion_fp8_admission(&admission)?; - } - if let Some(admission) = packed_conversion_fp8_f16x2_admission { - ensure!( - shard.family == "packed_conversion" && shard.intrinsics.is_empty(), - "compact FP8 f16x2 conversion admission must be the only content of a packed_conversion shard" - ); - shard.intrinsics = expand_packed_conversion_fp8_f16x2_admission(&admission)?; - } - if let Some(admission) = scalar_conversion_admission { - ensure!( - shard.family == "scalar_conversion" && shard.intrinsics.is_empty(), - "compact scalar-conversion admission must be the only content of its shard" - ); - shard.intrinsics = expand_scalar_conversion_admission(&admission)?; - } - if let Some(admission) = scalar_arithmetic_admission { - ensure!( - shard.family == "scalar_arithmetic" && shard.intrinsics.is_empty(), - "compact scalar-arithmetic admission must be the only content of its shard" - ); - shard.intrinsics = expand_scalar_arithmetic_admission(&admission)?; - } - if let Some(admission) = scalar_math_admission { - ensure!( - shard.family == "scalar_math" && shard.intrinsics.is_empty(), - "compact scalar-math admission must be the only content of its shard" - ); - shard.intrinsics = expand_scalar_math_admission(&admission)?; - } - if let Some(admission) = extended_minmax_admission { - ensure!( - shard.family == "extended_minmax" && shard.intrinsics.is_empty(), - "compact extended-minmax admission must be the only content of its shard" - ); - shard.intrinsics = expand_extended_minmax_admission(&admission)?; - } - if let Some(admission) = cluster_sreg_admission { - ensure!( - shard.family == "sreg" && shard.intrinsics.is_empty() && compact_mma_count == 0, - "compact cluster-sreg admission must be the only content of an sreg shard" - ); - shard.intrinsics = expand_cluster_sreg_admission(&admission)?; - } - if let Some(admission) = cluster_barrier_admission { - ensure!( - shard.family == "cluster_barrier" && shard.intrinsics.is_empty(), - "compact cluster-barrier admission must be the only content of its shard" - ); - shard.intrinsics = expand_cluster_barrier_admission(&admission)?; - } - if let Some(admission) = mbarrier_extended_admission { - ensure!( - shard.family == "mbarrier_extended" && shard.intrinsics.is_empty(), - "compact extended-mbarrier admission must be the only content of its shard" - ); - shard.intrinsics = expand_mbarrier_extended_admission(&admission)?; - } - if let Some(admission) = special_register_admission { - ensure!( - shard.family == "sreg" && shard.intrinsics.is_empty(), - "compact special-register admission must be the only content of an sreg shard" - ); - shard.intrinsics = expand_special_register_admission(&admission)?; - } - if let Some(admission) = debug_control_admission { - ensure!( - shard.family == "debug_control" && shard.intrinsics.is_empty(), - "compact debug-control admission must be the only content of a debug_control shard" - ); - shard.intrinsics = expand_debug_control_admission(&admission)?; - } - if let Some(admission) = threadfence_admission { - ensure!( - shard.family == "sync" && shard.intrinsics.is_empty(), - "compact threadfence admission must be the only content of a sync shard" - ); - shard.intrinsics = expand_threadfence_admission(&admission)?; - } - if let Some(admission) = cluster_memory_admission { - ensure!( - shard.family == "cluster_memory" && shard.intrinsics.is_empty(), - "compact cluster-memory admission must be the only content of its shard" - ); - shard.intrinsics = expand_cluster_memory_admission(&admission)?; - } - if let Some(admission) = stmatrix_admission { - ensure!( - shard.family == "stmatrix" && shard.intrinsics.is_empty(), - "compact stmatrix admission must be the only content of its shard" - ); - shard.intrinsics = expand_stmatrix_admission(&admission)?; - } - if let Some(admission) = clc_admission { - ensure!( - shard.family == "clc" && shard.intrinsics.is_empty(), - "compact CLC admission must be the only content of a clc shard" - ); - shard.intrinsics = expand_clc_admission(&admission)?; - } - if let Some(admission) = wgmma_control_admission { - ensure!( - shard.family == "wgmma_control" && shard.intrinsics.is_empty(), - "compact WGMMA-control admission must be the only content of its shard" - ); - shard.intrinsics = expand_wgmma_control_admission(&admission)?; - } - if let Some(admission) = tma_admission { - ensure!( - shard.family == "tma" && shard.intrinsics.is_empty(), - "compact TMA admission must be the only content of a tma shard" - ); - shard.intrinsics = expand_tma_admission(&admission)?; - } - if let Some(admission) = tcgen05_admission { - ensure!( - shard.family == "tcgen05" && shard.intrinsics.is_empty(), - "compact tcgen05 admission must be the only content of a tcgen05 shard" - ); - shard.intrinsics = expand_tcgen05_admission(&admission)?; - } - ensure!( - !shard.intrinsics.is_empty(), - "overlay shard {} contains no intrinsic records", - path.display() - ); - for record in &shard.intrinsics { - ensure!( - record.family == shard.family, - "overlay shard {} declares family {}, but intrinsic {} uses family {}", - path.display(), - shard.family, - record.id, - record.family - ); - } - - append_overlay_hash_input( - &mut hash_input, - relative - .to_str() - .context("overlay shard path is not valid UTF-8")?, - &bytes, - ); - overlay.intrinsics.extend(shard.intrinsics); - } - - Ok((overlay, sha256_bytes(&hash_input))) -} - -fn validate_overlay_shard_schema(shard: &OverlayShardFile, path: &Path) -> Result<()> { - validate_overlay_shard_schema_with_max(shard, path, OVERLAY_SHARD_SCHEMA) -} - -fn validate_overlay_shard_schema_with_max( - shard: &OverlayShardFile, - path: &Path, - maximum_schema: u32, -) -> Result<()> { - ensure!( - (MINIMUM_OVERLAY_SHARD_SCHEMA..=maximum_schema).contains(&shard.schema), - "unsupported overlay shard schema {} in {}", - shard.schema, - path.display() - ); - ensure!( - shard.sparse_mma_f8f6f4_f32.is_none() || shard.schema >= SPARSE_MMA_F8F6F4_SHARD_SCHEMA, - "compact sparse f8f6f4 MMA admission requires overlay shard schema {}", - SPARSE_MMA_F8F6F4_SHARD_SCHEMA - ); - ensure!( - shard.sparse_mma_f8f6f4_f16.is_none() || shard.schema >= SPARSE_MMA_F8F6F4_F16_SHARD_SCHEMA, - "compact sparse f8f6f4 F16 MMA admission requires overlay shard schema {}", - SPARSE_MMA_F8F6F4_F16_SHARD_SCHEMA - ); - ensure!( - shard.register_mma_f8f6f4_f32.is_none() || shard.schema >= REGISTER_MMA_F8F6F4_SHARD_SCHEMA, - "compact dense f8f6f4 MMA admission requires overlay shard schema {}", - REGISTER_MMA_F8F6F4_SHARD_SCHEMA - ); - ensure!( - shard.register_mma_f8f6f4_f16.is_none() - || shard.schema >= REGISTER_MMA_F8F6F4_F16_SHARD_SCHEMA, - "compact dense f8f6f4 F16 MMA admission requires overlay shard schema {}", - REGISTER_MMA_F8F6F4_F16_SHARD_SCHEMA - ); - ensure!( - shard.register_mma_mxf8f6f4_f32.is_none() - || shard.schema >= REGISTER_MMA_MXF8F6F4_SHARD_SCHEMA, - "compact dense mxf8f6f4 MMA admission requires overlay shard schema {}", - REGISTER_MMA_MXF8F6F4_SHARD_SCHEMA - ); - ensure!( - shard.register_mma_fp8.is_none() || shard.schema >= REGISTER_MMA_FP8_SHARD_SCHEMA, - "compact standard FP8 MMA admission requires overlay shard schema {}", - REGISTER_MMA_FP8_SHARD_SCHEMA - ); - ensure!( - shard.register_mma_ampere_float.is_none() - || shard.schema >= REGISTER_MMA_AMPERE_FLOAT_SHARD_SCHEMA, - "compact Ampere floating-point MMA admission requires overlay shard schema {}", - REGISTER_MMA_AMPERE_FLOAT_SHARD_SCHEMA - ); - ensure!( - shard.prmt.is_none() || shard.schema >= PRMT_SHARD_SCHEMA, - "compact prmt admission requires overlay shard schema {}", - PRMT_SHARD_SCHEMA - ); - ensure!( - shard.packed_conversion_fp8.is_none() || shard.schema >= PACKED_CONVERSION_FP8_SHARD_SCHEMA, - "compact FP8 conversion admission requires overlay shard schema {}", - PACKED_CONVERSION_FP8_SHARD_SCHEMA - ); - ensure!( - shard.packed_conversion_fp8_f16x2.is_none() - || shard.schema >= PACKED_CONVERSION_FP8_F16X2_SHARD_SCHEMA, - "compact FP8 f16x2 conversion admission requires overlay shard schema {}", - PACKED_CONVERSION_FP8_F16X2_SHARD_SCHEMA - ); - ensure!( - shard.scalar_conversion.is_none() || shard.schema >= SCALAR_CONVERSION_SHARD_SCHEMA, - "compact scalar-conversion admission requires overlay shard schema {}", - SCALAR_CONVERSION_SHARD_SCHEMA - ); - ensure!( - shard.scalar_arithmetic.is_none() || shard.schema >= SCALAR_ARITHMETIC_SHARD_SCHEMA, - "compact scalar-arithmetic admission requires overlay shard schema {}", - SCALAR_ARITHMETIC_SHARD_SCHEMA - ); - ensure!( - shard.scalar_math.is_none() || shard.schema >= SCALAR_MATH_SHARD_SCHEMA, - "compact scalar-math admission requires overlay shard schema {}", - SCALAR_MATH_SHARD_SCHEMA - ); - ensure!( - shard.extended_minmax.is_none() || shard.schema >= EXTENDED_MINMAX_SHARD_SCHEMA, - "compact extended-minmax admission requires overlay shard schema {}", - EXTENDED_MINMAX_SHARD_SCHEMA - ); - ensure!( - shard.cluster_sreg.is_none() || shard.schema >= CLUSTER_SREG_SHARD_SCHEMA, - "compact cluster-sreg admission requires overlay shard schema {}", - CLUSTER_SREG_SHARD_SCHEMA - ); - ensure!( - shard.cluster_barrier.is_none() || shard.schema >= CLUSTER_BARRIER_SHARD_SCHEMA, - "compact cluster-barrier admission requires overlay shard schema {}", - CLUSTER_BARRIER_SHARD_SCHEMA - ); - ensure!( - shard.special_registers.is_none() || shard.schema >= SPECIAL_REGISTER_SHARD_SCHEMA, - "compact special-register admission requires overlay shard schema {}", - SPECIAL_REGISTER_SHARD_SCHEMA - ); - ensure!( - shard.debug_control.is_none() || shard.schema >= DEBUG_CONTROL_SHARD_SCHEMA, - "compact debug-control admission requires overlay shard schema {}", - DEBUG_CONTROL_SHARD_SCHEMA - ); - ensure!( - shard.threadfence.is_none() || shard.schema >= THREADFENCE_SHARD_SCHEMA, - "compact threadfence admission requires overlay shard schema {}", - THREADFENCE_SHARD_SCHEMA - ); - ensure!( - shard.cluster_memory.is_none() || shard.schema >= CLUSTER_MEMORY_SHARD_SCHEMA, - "compact cluster-memory admission requires overlay shard schema {}", - CLUSTER_MEMORY_SHARD_SCHEMA - ); - ensure!( - shard.stmatrix.is_none() || shard.schema >= STMATRIX_SHARD_SCHEMA, - "compact stmatrix admission requires overlay shard schema {}", - STMATRIX_SHARD_SCHEMA - ); - ensure!( - shard.clc.is_none() || shard.schema >= CLC_SHARD_SCHEMA, - "compact CLC admission requires overlay shard schema {}", - CLC_SHARD_SCHEMA - ); - ensure!( - shard.tma.is_none() || shard.schema >= TMA_SHARD_SCHEMA, - "compact TMA admission requires overlay shard schema {}", - TMA_SHARD_SCHEMA - ); - ensure!( - shard - .tma - .as_ref() - .is_none_or(|admission| admission.reduce_variants.is_empty()) - || shard.schema >= TMA_REDUCTION_SHARD_SCHEMA, - "compact TMA reduction admission requires overlay shard schema {}", - TMA_REDUCTION_SHARD_SCHEMA - ); - ensure!( - shard.mbarrier_extended.is_none() || shard.schema >= MBARRIER_EXTENDED_SHARD_SCHEMA, - "compact extended-mbarrier admission requires overlay shard schema {}", - MBARRIER_EXTENDED_SHARD_SCHEMA - ); - ensure!( - shard.wgmma_controls.is_none() || shard.schema >= WGMMA_CONTROL_SHARD_SCHEMA, - "compact WGMMA-control admission requires overlay shard schema {}", - WGMMA_CONTROL_SHARD_SCHEMA - ); - ensure!( - shard.tcgen05.is_none() || shard.schema >= TCGEN05_SHARD_SCHEMA, - "compact tcgen05 admission requires overlay shard schema {}", - TCGEN05_SHARD_SCHEMA - ); - ensure!( - shard - .tcgen05 - .as_ref() - .is_none_or(|admission| admission.cp_variants.is_empty()) - || shard.schema >= TCGEN05_CP_SHARD_SCHEMA, - "compact tcgen05 copy admission requires overlay shard schema {}", - TCGEN05_CP_SHARD_SCHEMA - ); - ensure!( - shard - .tcgen05 - .as_ref() - .is_none_or(|admission| admission.ld_variants.is_empty()) - || shard.schema >= TCGEN05_LD_SHARD_SCHEMA, - "compact tcgen05 load admission requires overlay shard schema {}", - TCGEN05_LD_SHARD_SCHEMA - ); - ensure!( - shard - .tcgen05 - .as_ref() - .is_none_or(|admission| admission.st_variants.is_empty()) - || shard.schema >= TCGEN05_ST_SHARD_SCHEMA, - "compact tcgen05 store admission requires overlay shard schema {}", - TCGEN05_ST_SHARD_SCHEMA - ); - ensure!( - shard.tcgen05.as_ref().is_none_or(|admission| { - admission.ld_offset_variants.is_empty() && admission.st_offset_variants.is_empty() - }) || shard.schema >= TCGEN05_OFFSET_LDST_SHARD_SCHEMA, - "compact tcgen05 offset load/store admission requires overlay shard schema {}", - TCGEN05_OFFSET_LDST_SHARD_SCHEMA - ); - ensure!( - shard.tcgen05.as_ref().is_none_or(|admission| { - admission.control_llvm_evidence_profile.is_none() - && admission.control_libnvvm_evidence_profile.is_none() - && !admission.variants.iter().any(|variant| { - matches!( - variant.operation, - Tcgen05Operation::CommitMulticast - | Tcgen05Operation::ShiftDown - | Tcgen05Operation::ShiftDownCg2 - ) - }) - }) || shard.schema >= TCGEN05_CONTROL_SHARD_SCHEMA, - "compact tcgen05 control admission requires overlay shard schema {}", - TCGEN05_CONTROL_SHARD_SCHEMA - ); - if let Some(admission) = &shard.tcgen05 - && shard.schema >= TCGEN05_CONTROL_SHARD_SCHEMA - { - ensure!( - admission - .control_llvm_evidence_profile - .as_deref() - .is_some_and(|profile| !profile.trim().is_empty()) - && admission - .control_libnvvm_evidence_profile - .as_deref() - .is_some_and(|profile| !profile.trim().is_empty()) - && [ - Tcgen05Operation::CommitMulticast, - Tcgen05Operation::ShiftDown, - Tcgen05Operation::ShiftDownCg2, - ] - .into_iter() - .all(|operation| admission - .variants - .iter() - .any(|variant| variant.operation == operation)), - "compact tcgen05 schema {} requires all three control variants and both backend evidence profiles", - TCGEN05_CONTROL_SHARD_SCHEMA - ); - } - ensure!( - shard.tcgen05.as_ref().is_none_or(|admission| { - admission.mma_variants.is_empty() - && admission.mma_llvm_evidence_profile.is_none() - && admission.mma_libnvvm_evidence_profile.is_none() - && admission.mma_llvm_target_contracts.is_empty() - && admission.mma_libnvvm_target_contracts.is_empty() - }) || shard.schema >= TCGEN05_MMA_SHARD_SCHEMA, - "compact tcgen05 MMA admission requires overlay shard schema {}", - TCGEN05_MMA_SHARD_SCHEMA - ); - Ok(()) -} - -fn validate_overlay_shard_path(path: &str) -> Result<()> { - let path = Path::new(path); - ensure!( - path.extension().and_then(|extension| extension.to_str()) == Some("toml"), - "overlay shard path must name a TOML file: {}", - path.display() - ); - let components: Vec<_> = path.components().collect(); - ensure!( - components.len() >= 2 && components[0] == Component::Normal("overlay".as_ref()), - "overlay shard path must stay under intrinsics/overlay: {}", - path.display() - ); - ensure!( - components - .iter() - .all(|component| matches!(component, Component::Normal(_))), - "overlay shard path contains a non-normal component: {}", - path.display() - ); - Ok(()) -} - -fn append_overlay_hash_input(output: &mut Vec, path: &str, contents: &[u8]) { - output.extend_from_slice(&(path.len() as u64).to_le_bytes()); - output.extend_from_slice(path.as_bytes()); - output.extend_from_slice(&(contents.len() as u64).to_le_bytes()); - output.extend_from_slice(contents); -} - -fn shares_tma_2d_g2s_symbol(record: &OverlayIntrinsic, symbol: &str) -> bool { - symbol == "llvm.nvvm.cp.async.bulk.tensor.g2s.tile.2d" - && record.tma.as_ref().is_some_and(|tma| { - matches!( - tma.operation, - TmaOperation::G2sTile2d - | TmaOperation::G2sTile2dMulticast - | TmaOperation::G2sTile2dMulticastCg2 - ) - }) -} - -fn shares_tma_prefetch_tile_symbol(record: &OverlayIntrinsic, symbol: &str) -> bool { - let Some(tma) = record.tma.as_ref() else { - return false; - }; - let expected_symbol = match tma.operation { - TmaOperation::PrefetchTile1d | TmaOperation::PrefetchTile1dCacheHint => { - "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.1d" - } - TmaOperation::PrefetchTile2d | TmaOperation::PrefetchTile2dCacheHint => { - "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.2d" - } - TmaOperation::PrefetchTile3d | TmaOperation::PrefetchTile3dCacheHint => { - "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.3d" - } - TmaOperation::PrefetchTile4d | TmaOperation::PrefetchTile4dCacheHint => { - "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.4d" - } - TmaOperation::PrefetchTile5d | TmaOperation::PrefetchTile5dCacheHint => { - "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.5d" - } - TmaOperation::PrefetchTileGather4TwoDimensional - | TmaOperation::PrefetchTileGather4TwoDimensionalCacheHint => { - "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.gather4.2d" - } - _ => return false, - }; - symbol == expected_symbol -} - -fn shares_tcgen05_mma_symbol(record: &OverlayIntrinsic, symbol: &str) -> bool { - record.tcgen05.as_ref().is_some_and(|tcgen05| { - (symbol == "llvm.nvvm.tcgen05.mma.ws.tensor" - && matches!( - tcgen05.operation, - Tcgen05Operation::MmaWsF16 - | Tcgen05Operation::MmaWsBf16 - | Tcgen05Operation::MmaWsTf32 - )) - || tcgen05 - .mma - .as_ref() - .is_some_and(|mma| symbol == tcgen05_mma_llvm_symbol(mma.form)) - }) -} - -fn shares_tcgen05_ld_symbol(record: &OverlayIntrinsic, symbol: &str) -> bool { - let Some(tcgen05) = record.tcgen05.as_ref() else { - return false; - }; - if let Some(ld) = tcgen05.ld { - return tcgen05.operation == Tcgen05Operation::Ld - && record.source_record.as_deref() == Some(tcgen05_ld_source_record(ld).as_str()) - && symbol == tcgen05_ld_llvm_symbol(ld); - } - matches!( - (tcgen05.operation, record.source_record.as_deref(), symbol), - ( - Tcgen05Operation::Ld16x256bPure, - Some("int_nvvm_tcgen05_ld_16x256b_x1"), - "llvm.nvvm.tcgen05.ld.16x256b.x1" - ) | ( - Tcgen05Operation::Ld16x256bX8Pure, - Some("int_nvvm_tcgen05_ld_16x256b_x8"), - "llvm.nvvm.tcgen05.ld.16x256b.x8" - ) - ) -} - -fn shares_tcgen05_st_symbol(record: &OverlayIntrinsic, symbol: &str) -> bool { - record.tcgen05.as_ref().is_some_and(|tcgen05| { - tcgen05.st.is_some_and(|st| { - tcgen05.operation == Tcgen05Operation::St - && record.source_record.as_deref() == Some(tcgen05_st_source_record(st).as_str()) - && symbol == tcgen05_st_llvm_symbol(st) - }) - }) -} - -fn validate_unique_overlay(records: &[OverlayIntrinsic], intrinsic_abi: u32) -> Result<()> { - let mut ids = BTreeSet::new(); - let mut abi_ids = BTreeSet::new(); - let mut operation_keys = BTreeSet::new(); - let mut paths = BTreeSet::new(); - let mut op_variants = BTreeSet::new(); - let mut op_type_names = BTreeMap::new(); - let mut symbol_bases = BTreeMap::new(); - let mut symbols = BTreeSet::new(); - let mut rust_items = BTreeSet::new(); - for record in records { - insert_unique(&mut ids, &record.id, "catalog ID")?; - validate_abi_id(&record.abi_id)?; - insert_unique(&mut abi_ids, &record.abi_id, "intrinsic ABI ID")?; - validate_operation_key(&record.operation_key)?; - insert_unique( - &mut operation_keys, - &record.operation_key, - "intrinsic operation key", - )?; - if let Some(previous_name) = op_type_names.insert( - record.dialect_op_type.as_str(), - record.dialect_op_name.as_str(), - ) { - ensure!( - previous_name == record.dialect_op_name, - "dialect op type {} maps to both {} and {}", - record.dialect_op_type, - previous_name, - record.dialect_op_name - ); - } - let op_variant = format!( - "{}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}", - record.dialect_op_name, - record.ldmatrix_variant, - record.packed_atomic, - record.redux, - record.vote, - record.active_mask, - record.warp_match, - record.warp_barrier, - record.warp_shuffle, - record.dot_product, - record.cp_async_copy, - record.cp_async_control, - record.cp_async_mbarrier, - record.mbarrier_basic, - record.movmatrix, - record.mbarrier_extended, - record.register_mma, - record.sparse_mma, - record.prmt, - record.cluster_barrier, - record.cluster_memory, - record.scalar_conversion, - record.scalar_arithmetic, - record.scalar_math, - record.extended_minmax, - record.tcgen05, - ); - insert_unique(&mut op_variants, &op_variant, "dialect op variant")?; - if let Some(symbol) = &record.llvm_symbol { - let is_resolved = record.resolved_llvm_symbol.is_some(); - let shares_reviewed_symbol = shares_tma_2d_g2s_symbol(record, symbol) - || shares_tma_prefetch_tile_symbol(record, symbol) - || shares_tcgen05_mma_symbol(record, symbol) - || shares_tcgen05_ld_symbol(record, symbol) - || shares_tcgen05_st_symbol(record, symbol); - if let Some((previous_was_resolved, previous_shared_symbol)) = - symbol_bases.insert(symbol, (is_resolved, shares_reviewed_symbol)) - { - ensure!( - (previous_was_resolved && is_resolved) - || (previous_shared_symbol && shares_reviewed_symbol), - "duplicate LLVM symbol {symbol} is reused without a resolved symbol" - ); - } - if !shares_reviewed_symbol { - insert_unique( - &mut symbols, - record.resolved_llvm_symbol.as_ref().unwrap_or(symbol), - "resolved LLVM symbol", - )?; - } - } - insert_unique( - &mut rust_items, - &format!("{}::{}", record.rust_module, record.rust_name), - "raw Rust item", - )?; - insert_unique( - &mut paths, - &canonical_rust_path(intrinsic_abi, &record.abi_id), - "canonical Rust path", - )?; - insert_unique(&mut paths, &record.public_rust_path, "public Rust path")?; - for path in &record.compatibility_rust_paths { - insert_unique(&mut paths, path, "compatibility Rust path")?; - } - } - Ok(()) -} - -struct AbiLedgerIndex<'a> { - by_catalog_id: BTreeMap<&'a str, &'a AbiLedgerEntry>, -} - -impl<'a> AbiLedgerIndex<'a> { - fn new(ledger: &'a AbiLedgerFile) -> Result { - let mut by_catalog_id = BTreeMap::new(); - for entry in &ledger.entries { - ensure!( - by_catalog_id - .insert(entry.catalog_id.as_str(), entry) - .is_none(), - "duplicate ABI ledger catalog ID: {}", - entry.catalog_id - ); - } - Ok(Self { by_catalog_id }) - } - - fn active(&self, catalog_id: &str) -> Result<&'a AbiLedgerEntry> { - let entry = self - .by_catalog_id - .get(catalog_id) - .copied() - .with_context(|| format!("generated intrinsic {catalog_id} has no ABI ledger entry"))?; - ensure!( - entry.status == "active", - "generated intrinsic {catalog_id} maps to non-active ABI ledger entry {}", - entry.abi_id - ); - Ok(entry) - } -} - -fn bind_generated_abi_ids(overlay: &mut OverlayFile, ledger: &AbiLedgerFile) -> Result<()> { - let index = AbiLedgerIndex::new(ledger)?; - for record in overlay - .intrinsics - .iter_mut() - .filter(|record| record.abi_id.is_empty()) - { - let entry = index.active(&record.id)?; - ensure!( - entry.operation_key == record.operation_key, - "generated intrinsic {} operation key mismatch: ledger {:?}, derived {:?}", - record.id, - entry.operation_key, - record.operation_key - ); - let derived_signature = raw_rust_signature(record); - ensure!( - entry.raw_rust_signature == derived_signature, - "generated intrinsic {} raw Rust signature mismatch: ledger {:?}, derived {:?}", - record.id, - entry.raw_rust_signature, - derived_signature - ); - record.abi_id.clone_from(&entry.abi_id); - } - Ok(()) -} - -fn validate_abi_ledger(overlay: &OverlayFile, ledger: &AbiLedgerFile) -> Result<()> { - ensure!( - ledger.schema == 1, - "unsupported ABI ledger schema {}", - ledger.schema - ); - ensure!( - ledger.intrinsic_abi == overlay.intrinsic_abi, - "ABI ledger v{} does not match overlay ABI v{}", - ledger.intrinsic_abi, - overlay.intrinsic_abi - ); - ensure!(!ledger.entries.is_empty(), "ABI ledger contains no entries"); - - let overlay_by_abi_id: BTreeMap<_, _> = overlay - .intrinsics - .iter() - .map(|record| (record.abi_id.as_str(), record)) - .collect(); - let mut abi_ids = BTreeSet::new(); - let mut catalog_ids = BTreeSet::new(); - let mut operation_keys = BTreeSet::new(); - let mut previous_abi_id: Option<&str> = None; - for entry in &ledger.entries { - validate_abi_id(&entry.abi_id)?; - if let Some(previous) = previous_abi_id { - ensure!( - previous < entry.abi_id.as_str(), - "ABI ledger IDs must be unique and append-only in ascending order: {} follows {}", - entry.abi_id, - previous - ); - } - previous_abi_id = Some(&entry.abi_id); - insert_unique(&mut abi_ids, &entry.abi_id, "ABI ledger ID")?; - insert_unique(&mut catalog_ids, &entry.catalog_id, "ABI ledger catalog ID")?; - validate_operation_key(&entry.operation_key)?; - insert_unique( - &mut operation_keys, - &entry.operation_key, - "ABI ledger operation key", - )?; - ensure!( - !entry.catalog_id.is_empty() - && !entry.raw_rust_signature.result.is_empty() - && entry - .raw_rust_signature - .arguments - .iter() - .all(|argument| !argument.is_empty()), - "ABI ledger entry {} has incomplete identity data", - entry.abi_id - ); - - let overlay_record = overlay_by_abi_id.get(entry.abi_id.as_str()).copied(); - match entry.status.as_str() { - "active" => { - let record = overlay_record.with_context(|| { - format!( - "active ABI ledger entry {} ({}) has no overlay record", - entry.abi_id, entry.catalog_id - ) - })?; - validate_active_ledger_entry(entry, record)?; - } - "tombstone" => ensure!( - overlay_record.is_none(), - "tombstoned ABI ID {} cannot reappear in the overlay", - entry.abi_id - ), - status => bail!( - "ABI ledger entry {} has unsupported status {status:?}; expected active or tombstone", - entry.abi_id - ), - } - } - - for record in &overlay.intrinsics { - ensure!( - abi_ids.contains(&record.abi_id), - "overlay intrinsic {} uses ABI ID {} with no ledger entry", - record.id, - record.abi_id - ); - } - Ok(()) -} - -fn validate_active_ledger_entry(entry: &AbiLedgerEntry, record: &OverlayIntrinsic) -> Result<()> { - let comparisons = [ - ("catalog ID", entry.catalog_id.as_str(), record.id.as_str()), - ( - "operation key", - entry.operation_key.as_str(), - record.operation_key.as_str(), - ), - ]; - for (field, ledger_value, overlay_value) in comparisons { - ensure!( - ledger_value == overlay_value, - "ABI ledger {} {field} mismatch: ledger {ledger_value:?}, overlay {overlay_value:?}", - entry.abi_id - ); - } - let expected_signature = raw_rust_signature(record); - ensure!( - entry.raw_rust_signature == expected_signature, - "ABI ledger {} raw Rust signature mismatch: ledger {:?}, overlay {:?}", - entry.abi_id, - entry.raw_rust_signature, - expected_signature - ); - Ok(()) -} - -fn raw_rust_signature(record: &OverlayIntrinsic) -> AbiRawRustSignature { - AbiRawRustSignature { - safe: record.safe, - arguments: record.rust_arguments.clone(), - result: record.rust_result.clone(), - } -} - -fn validate_abi_id(abi_id: &str) -> Result<()> { - ensure!( - abi_id.len() == 5 - && abi_id.starts_with('i') - && abi_id[1..].bytes().all(|byte| byte.is_ascii_digit()), - "intrinsic ABI ID `{abi_id}` must use the stable `iNNNN` form" - ); - Ok(()) -} - -pub(crate) fn validate_operation_key(operation_key: &str) -> Result<()> { - ensure!( - !operation_key.is_empty() - && operation_key.split('.').all(|segment| { - !segment.is_empty() - && segment.bytes().all(|byte| { - byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'_' - }) - }), - "intrinsic operation key `{operation_key}` must contain dot-separated lowercase semantic components" - ); - Ok(()) -} - -fn canonical_rust_path(intrinsic_abi: u32, abi_id: &str) -> String { - format!("cuda_intrinsics::__cuda_oxide_intrinsic_abi_v{intrinsic_abi}::{abi_id}") -} - -fn insert_unique(set: &mut BTreeSet, value: &str, kind: &str) -> Result<()> { - ensure!(set.insert(value.to_owned()), "duplicate {kind}: {value}"); - Ok(()) -} - -fn resolve_policy_source(policy: &OverlayIntrinsic) -> Result { - match (&policy.source, &policy.source_record) { - (None, Some(source_record)) => Ok(IntrinsicSource::LlvmImported { - source_record: source_record.clone(), - }), - (Some(source @ IntrinsicSource::PtxNative { .. }), None) => Ok(source.clone()), - (Some(IntrinsicSource::LlvmImported { source_record }), None) => { - ensure!( - !source_record.trim().is_empty(), - "{} has an empty imported LLVM source record", - policy.id - ); - Ok(IntrinsicSource::LlvmImported { - source_record: source_record.clone(), - }) - } - (Some(_), Some(_)) => bail!( - "{} mixes tagged source provenance with the legacy source_record field", - policy.id - ), - (None, None) => bail!("{} has no intrinsic source provenance", policy.id), - } -} - -fn validate_policy( - policy: &OverlayIntrinsic, - source: &IntrinsicSource, - declaration: Option<&ImportedIntrinsic>, - intrinsic_abi: u32, -) -> Result<()> { - validate_abi_id(&policy.abi_id)?; - parse_ptx_version(&policy.minimum_ptx, &policy.id)?; - parse_hardware_target(policy)?; - policy.expected_ptx.validate().map_err(|reason| { - anyhow::anyhow!( - "{} has an invalid expected PTX pattern: {reason}", - policy.id - ) - })?; - let public_path = format!( - "cuda_intrinsics::{}::{}", - policy.rust_module, policy.rust_name - ); - ensure!( - policy.public_rust_path == public_path, - "{} public Rust path must be {}", - policy.id, - public_path - ); - let canonical_path = canonical_rust_path(intrinsic_abi, &policy.abi_id); - ensure!( - canonical_path != policy.public_rust_path - && !policy - .compatibility_rust_paths - .iter() - .any(|path| path == &canonical_path || path == &policy.public_rust_path), - "{} must keep canonical, public, and compatibility Rust paths distinct", - policy.id - ); - match (source, declaration) { - (IntrinsicSource::LlvmImported { .. }, Some(declaration)) => { - ensure!( - policy.llvm_symbol.as_deref() == Some(declaration.llvm_name.as_str()), - "{} LLVM symbol mismatch: imported {}, overlay {:?}", - policy.id, - declaration.llvm_name, - policy.llvm_symbol - ); - ensure!( - declaration.arguments == policy.llvm_arguments, - "{} LLVM argument signature mismatch: imported {:?}, overlay {:?}", - policy.id, - declaration.arguments, - policy.llvm_arguments - ); - ensure!( - declaration.results == policy.llvm_results, - "{} LLVM result signature mismatch: imported {:?}, overlay {:?}", - policy.id, - declaration.results, - policy.llvm_results - ); - } - (IntrinsicSource::PtxNative { instruction }, None) => ensure!( - !instruction.trim().is_empty() - && policy.llvm_symbol.is_none() - && policy.resolved_llvm_symbol.is_none() - && policy.llvm_arguments.is_empty() - && policy.llvm_results.is_empty(), - "{} PTX-native source must not invent LLVM source facts", - policy.id - ), - _ => bail!( - "{} source kind and imported declaration disagree", - policy.id - ), - } - match policy.family.as_str() { - "sreg" => validate_sreg_policy(policy, source, declaration)?, - "ldmatrix" => validate_ldmatrix_policy( - policy, - declaration.context("ldmatrix requires imported LLVM declaration")?, - )?, - "stmatrix" => validate_stmatrix_policy( - policy, - declaration.context("stmatrix requires imported LLVM declaration")?, - )?, - "packed_atomic" => validate_packed_atomic_policy(policy, source)?, - "redux" => validate_redux_policy( - policy, - declaration.context("redux requires imported LLVM declaration")?, - )?, - "dotprod" => validate_dot_product_policy( - policy, - declaration.context("dotprod requires imported LLVM declaration")?, - )?, - "sync" => validate_sync_policy( - policy, - declaration.context("sync requires imported LLVM declaration")?, - )?, - "vote" => validate_vote_policy( - policy, - declaration.context("vote requires imported LLVM declaration")?, - )?, - "active_mask" => validate_active_mask_policy( - policy, - declaration.context("active_mask requires imported LLVM declaration")?, - )?, - "warp_match" => validate_warp_match_policy( - policy, - declaration.context("warp_match requires imported LLVM declaration")?, - )?, - "elect" => validate_elect_policy( - policy, - declaration.context("elect requires imported LLVM declaration")?, - )?, - "warp_barrier" => validate_warp_barrier_policy( - policy, - declaration.context("warp_barrier requires imported LLVM declaration")?, - )?, - "warp_shuffle" => validate_warp_shuffle_policy(policy, declaration)?, - "packed_alu" => validate_packed_alu_policy(policy, source, declaration)?, - "integer_minmax" => validate_integer_minmax_policy(policy, source, declaration)?, - "packed_conversion" => validate_packed_conversion_policy(policy, source, declaration)?, - "scalar_conversion" => validate_scalar_conversion_policy( - policy, - declaration.context("scalar_conversion requires imported LLVM declaration")?, - )?, - "scalar_arithmetic" => validate_scalar_arithmetic_policy( - policy, - declaration.context("scalar_arithmetic requires imported LLVM declaration")?, - )?, - "scalar_math" => validate_scalar_math_policy(policy, declaration)?, - "extended_minmax" => validate_extended_minmax_policy( - policy, - declaration.context("extended_minmax requires imported LLVM declaration")?, - )?, - "cp_async_copy" => validate_cp_async_copy_policy( - policy, - declaration.context("cp_async_copy requires imported LLVM declaration")?, - )?, - "cp_async_control" => validate_cp_async_control_policy( - policy, - declaration.context("cp_async_control requires imported LLVM declaration")?, - )?, - "cp_async_mbarrier" => validate_cp_async_mbarrier_policy( - policy, - declaration.context("cp_async_mbarrier requires imported LLVM declaration")?, - )?, - "mbarrier_basic" => validate_mbarrier_basic_policy( - policy, - declaration.context("mbarrier_basic requires imported LLVM declaration")?, - )?, - "movmatrix" => validate_movmatrix_policy(policy, source)?, - "mbarrier_extended" => validate_mbarrier_extended_policy(policy, source, declaration)?, - "register_mma" => validate_register_mma_policy( - policy, - declaration.context("register_mma requires imported LLVM declaration")?, - )?, - "sparse_mma" => validate_sparse_mma_policy( - policy, - declaration.context("sparse_mma requires imported LLVM declaration")?, - )?, - "prmt" => validate_prmt_policy( - policy, - declaration.context("prmt requires imported LLVM declaration")?, - )?, - "cluster_barrier" => validate_cluster_barrier_policy( - policy, - declaration.context("cluster_barrier requires imported LLVM declaration")?, - )?, - "debug_control" => validate_debug_control_policy(policy, source)?, - "cluster_memory" => validate_cluster_memory_policy(policy, source, declaration)?, - "clc" => validate_clc_policy( - policy, - declaration.context("clc requires imported LLVM declaration")?, - )?, - "wgmma_control" => validate_wgmma_control_policy( - policy, - declaration.context("wgmma_control requires imported LLVM declaration")?, - )?, - "tma" => validate_tma_policy( - policy, - declaration.context("tma requires imported LLVM declaration")?, - )?, - "counted_barrier" | "grid_dependency" | "register_control" => { - validate_execution_control_policy( - policy, - declaration.context("execution-control requires imported LLVM declaration")?, - )? - } - "tcgen05" => validate_tcgen05_policy( - policy, - declaration.context("tcgen05 requires imported LLVM declaration")?, - )?, - family => bail!("{} uses unsupported generated family {family:?}", policy.id), - } - ensure!( - (policy.family == "movmatrix") == policy.movmatrix.is_some(), - "{} mixes the movmatrix contract with another generated family", - policy.id - ); - ensure!( - (policy.family == "mbarrier_extended") == policy.mbarrier_extended.is_some(), - "{} mixes the extended-mbarrier contract with another generated family", - policy.id - ); - ensure!( - (policy.family == "register_mma") == policy.register_mma.is_some(), - "{} mixes the register-MMA contract with another generated family", - policy.id - ); - ensure!( - (policy.family == "sparse_mma") == policy.sparse_mma.is_some(), - "{} mixes the sparse-MMA contract with another generated family", - policy.id - ); - ensure!( - (policy.family == "prmt") == policy.prmt.is_some(), - "{} mixes the prmt contract with another generated family", - policy.id - ); - ensure!( - (policy.family == "cluster_barrier") == policy.cluster_barrier.is_some(), - "{} mixes the cluster-barrier contract with another generated family", - policy.id - ); - ensure!( - (policy.family == "wgmma_control") == policy.wgmma_control.is_some(), - "{} mixes the WGMMA-control contract with another generated family", - policy.id - ); - ensure!( - (policy.family == "tma") == policy.tma.is_some(), - "{} mixes the TMA contract with another generated family", - policy.id - ); - ensure!( - policy.special_register.is_none() || policy.family == "sreg", - "{} mixes the special-register contract with another generated family", - policy.id - ); - ensure!( - (policy.family == "debug_control") == policy.debug_control.is_some(), - "{} mixes the debug-control contract with another generated family", - policy.id - ); - ensure!( - (policy.family == "cluster_memory") == policy.cluster_memory.is_some(), - "{} mixes the cluster-memory contract with another generated family", - policy.id - ); - ensure!( - (policy.family == "clc") == policy.clc.is_some(), - "{} mixes the CLC contract with another generated family", - policy.id - ); - ensure!( - (policy.family == "tma") == policy.tma.is_some(), - "{} mixes the TMA contract with another generated family", - policy.id - ); - ensure!( - (policy.family == "tcgen05") == policy.tcgen05.is_some(), - "{} mixes the tcgen05 contract with another generated family", - policy.id - ); - ensure!( - (policy.family == "scalar_conversion") == policy.scalar_conversion.is_some(), - "{} mixes the scalar-conversion contract with another generated family", - policy.id - ); - ensure!( - (policy.family == "scalar_arithmetic") == policy.scalar_arithmetic.is_some(), - "{} mixes the scalar-arithmetic contract with another generated family", - policy.id - ); - ensure!( - (policy.family == "scalar_math") == policy.scalar_math.is_some(), - "{} mixes the scalar-math contract with another generated family", - policy.id - ); - ensure!( - (policy.family == "extended_minmax") == policy.extended_minmax.is_some(), - "{} mixes the extended-minmax contract with another generated family", - policy.id - ); - ensure!( - !policy.execution_scope.trim().is_empty(), - "{} has no execution scope", - policy.id - ); - ensure!( - !policy.ptx_isa_version.trim().is_empty() - && !policy.ptx_isa_section.trim().is_empty() - && policy.ptx_isa_url.starts_with("https://docs.nvidia.com/"), - "{} has incomplete or non-authoritative PTX ISA provenance", - policy.id - ); - match (policy.safe, policy.safe_allowlist_reason.as_deref()) { - (true, Some(reason)) if !reason.trim().is_empty() => {} - (true, _) => bail!( - "{} is safe but has no nonempty safe_allowlist_reason", - policy.id - ), - (false, Some(reason)) if !reason.trim().is_empty() => bail!( - "{} is unsafe but has a safe_allowlist_reason; safe exceptions apply only to safe items", - policy.id - ), - (false, _) => {} - } - if let Some(declaration) = declaration { - if policy.pure { - ensure!( - declaration - .classes - .iter() - .any(|class| class == "NVVMPureIntrinsic") - || (matches!( - policy.family.as_str(), - "packed_alu" | "scalar_arithmetic" | "extended_minmax" - ) && declaration - .properties - .iter() - .any(|property| property == "IntrNoMem") - && declaration - .properties - .iter() - .any(|property| property == "IntrSpeculatable")) - || (policy.family == "scalar_math" - && declaration - .properties - .iter() - .any(|property| property == "IntrNoMem")), - "{} is marked pure, but its imported declaration is not an NVVMPureIntrinsic", - policy.id - ); - } - if policy.memory == "none" { - ensure!( - declaration - .properties - .iter() - .any(|property| property == "IntrNoMem"), - "{} is marked no-memory, but its imported declaration lacks IntrNoMem", - policy.id - ); - } - let imported_convergent = declaration - .properties - .iter() - .any(|property| property == "IntrConvergent"); - let convergence_supplied_by_ptx = - (matches!(policy.family.as_str(), "register_mma" | "sparse_mma") - && (policy.register_mma.is_some() || policy.sparse_mma.is_some()) - || (policy.family == "cluster_memory" - && policy.cluster_memory.is_some() - && policy.backend_lowerings.iter().all(|lowering| { - lowering.mechanism == BackendLoweringMechanism::InlinePtx - })) - || (policy.family == "mbarrier_extended" - && policy.mbarrier_extended.is_some() - && policy.backend_lowerings.iter().all(|lowering| { - lowering.mechanism == BackendLoweringMechanism::InlinePtx - })) - || (policy.family == "tcgen05" - && policy.tcgen05.is_some() - && policy.backend_lowerings.iter().all(|lowering| { - lowering.mechanism == BackendLoweringMechanism::InlinePtx - }))) - && policy.convergent - && !imported_convergent; - ensure!( - imported_convergent == policy.convergent || convergence_supplied_by_ptx, - "{} convergence mismatch: imported {}, overlay {}", - policy.id, - imported_convergent, - policy.convergent - ); - let selectionless_closed_family = (policy.family == "packed_conversion" - && policy.packed_conversion.is_some()) - || (policy.family == "prmt" && policy.prmt.is_some()) - || (policy.family == "clc" - && policy.clc.as_ref().is_some_and(|clc| { - matches!( - clc.operation, - ClcOperation::QueryIsCanceled - | ClcOperation::QueryGetFirstCtaidX - | ClcOperation::QueryGetFirstCtaidY - | ClcOperation::QueryGetFirstCtaidZ - ) - })) - || (policy.family == "sreg" - && policy.special_register.as_ref().is_some_and(|special| { - matches!( - special.register, - SpecialRegisterKind::Envreg1 | SpecialRegisterKind::Envreg2 - ) - })) - || policy.family == "stmatrix" - || (policy.family == "tma" - && policy.tma.as_ref().is_some_and(|tma| { - matches!( - tma.operation, - TmaOperation::Reduce - | TmaOperation::PrefetchTensorMap - | TmaOperation::ReplaceBoxDim - | TmaOperation::ReplaceElementStride - | TmaOperation::ReplaceElementType - | TmaOperation::ReplaceFillMode - | TmaOperation::ReplaceGlobalAddress - | TmaOperation::ReplaceGlobalDim - | TmaOperation::ReplaceGlobalStride - | TmaOperation::ReplaceInterleaveLayout - | TmaOperation::ReplaceRank - | TmaOperation::ReplaceSwizzleAtomicity - | TmaOperation::ReplaceSwizzleMode - ) - })) - || (policy.family == "tcgen05" - && policy.tcgen05.as_ref().is_some_and(|tcgen05| { - tcgen05.source_contract - == Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection - })) - || (policy.family == "scalar_math" - && policy.scalar_math.as_ref().is_some_and(|sm| { - // Tanh is absent here only because it is PTX-native (no - // imported declaration), so this check never sees it. - matches!( - sm.operation, - ScalarMathOperation::Sin - | ScalarMathOperation::Cos - | ScalarMathOperation::Ex2 - | ScalarMathOperation::Lg2 - | ScalarMathOperation::Rsqrt - ) - })); - ensure!( - !declaration.selections.is_empty() || selectionless_closed_family, - "{} has a declaration but no NVPTX TableGen selection record", - policy.id - ); - let mut matching_selections = Vec::new(); - for selection in &declaration.selections { - if selection_matches_policy(policy, selection)? { - matching_selections.push(selection); - } - } - let expected_selection_count = match policy.family.as_str() { - "vote" | "warp_barrier" | "elect" => 2, - "warp_match" => 4, - "warp_shuffle" => 8, - "counted_barrier" => 4, - "packed_conversion" | "prmt" | "stmatrix" => 0, - "tma" - if policy.tma.as_ref().is_some_and(|tma| { - matches!( - tma.operation, - TmaOperation::Reduce - | TmaOperation::PrefetchTensorMap - | TmaOperation::ReplaceBoxDim - | TmaOperation::ReplaceElementStride - | TmaOperation::ReplaceElementType - | TmaOperation::ReplaceFillMode - | TmaOperation::ReplaceGlobalAddress - | TmaOperation::ReplaceGlobalDim - | TmaOperation::ReplaceGlobalStride - | TmaOperation::ReplaceInterleaveLayout - | TmaOperation::ReplaceRank - | TmaOperation::ReplaceSwizzleAtomicity - | TmaOperation::ReplaceSwizzleMode - ) - }) => - { - 0 - } - "tcgen05" - if policy - .tcgen05 - .as_ref() - .and_then(|tcgen05| tcgen05.mma.as_ref()) - .is_some() => - { - let mma = policy.tcgen05.as_ref().unwrap().mma.as_ref().unwrap(); - if mma.alias.is_some() { - 1 - } else if tcgen05_mma_is_ws(mma.form) { - 64 - } else if tcgen05_mma_is_ashift(mma.form) { - 16 - } else { - 32 - } - } - "tcgen05" - if policy.tcgen05.as_ref().is_some_and(|tcgen05| { - tcgen05.source_contract != Tcgen05SourceContract::ExactTablegenSelection - }) => - { - 0 - } - "clc" - if policy.clc.as_ref().is_some_and(|clc| { - matches!( - clc.operation, - ClcOperation::QueryIsCanceled - | ClcOperation::QueryGetFirstCtaidX - | ClcOperation::QueryGetFirstCtaidY - | ClcOperation::QueryGetFirstCtaidZ - ) - }) => - { - 0 - } - "sreg" - if policy.special_register.as_ref().is_some_and(|special| { - matches!( - special.register, - SpecialRegisterKind::Envreg1 | SpecialRegisterKind::Envreg2 - ) - }) => - { - 0 - } - "mbarrier_extended" - if policy.mbarrier_extended.as_ref().is_some_and(|mbarrier| { - matches!( - mbarrier.operation, - MbarrierExtendedOperation::ArriveExpectTxCta - | MbarrierExtendedOperation::ArriveExpectTxCluster - | MbarrierExtendedOperation::TryWaitParityCta - | MbarrierExtendedOperation::TryWaitParityCluster - ) - }) => - { - 0 - } - "mbarrier_extended" if policy.id == "nanosleep" => 2, - "cp_async_copy" - if policy - .cp_async_copy - .as_ref() - .is_some_and(|copy| copy.source_size == CpAsyncSourceSize::Runtime) => - { - 2 - } - "cluster_memory" if policy.id == "map_shared_rank" => 2, - "scalar_math" - if policy.scalar_math.as_ref().is_some_and(|sm| { - matches!( - sm.operation, - ScalarMathOperation::Sin - | ScalarMathOperation::Cos - | ScalarMathOperation::Ex2 - | ScalarMathOperation::Lg2 - | ScalarMathOperation::Rsqrt - ) - }) => - { - 0 - } - _ => 1, - }; - ensure!( - matching_selections.len() == expected_selection_count, - "{} expected PTX {:?} does not agree with its closed imported selection set", - policy.id, - policy.expected_ptx - ); - for selection in matching_selections { - validate_selected_target_predicates(policy, selection)?; - } - } - Ok(()) -} - -fn selection_matches_policy( - policy: &OverlayIntrinsic, - selection: &crate::model::ImportedSelection, -) -> Result { - if policy.family == "wgmma_control" { - let Some(control) = &policy.wgmma_control else { - return Ok(false); - }; - let recipe = wgmma_control_recipe(control.mode); - let selection_shape_matches = if control.mode == WgmmaControlMode::WaitGroup { - selection.asm == "wgmma.wait_group.sync.aligned \t$n;" - } else { - policy.expected_ptx.matches(&selection.asm)? - }; - return Ok(selection.source_record == recipe.selection_record - && selection_shape_matches - && selection.predicates == ["Subtarget->getPTXVersion() >= 80", "hasSM90a"] - && selection.constraints.is_empty()); - } - - if policy.family == "tma" { - return Ok(selection_matches_tma_policy(policy, selection)); - } - if matches!( - policy.family.as_str(), - "counted_barrier" | "grid_dependency" | "register_control" - ) { - let Some(operation) = ExecutionControlOperation::from_catalog_id(&policy.id) else { - return Ok(false); - }; - let recipe = execution_control_recipe(operation); - return Ok(recipe - .selection_records - .contains(&selection.source_record.as_str()) - && selection.asm == recipe.selection_asm - && selection.predicates == recipe.selection_predicates - && selection.constraints.is_empty()); - } - if policy.family == "tcgen05" { - let Some(tcgen05) = &policy.tcgen05 else { - return Ok(false); - }; - if let Some(mma) = &tcgen05.mma { - let expected = if mma.alias.is_some() { - BTreeSet::from([if tcgen05_mma_is_ws(mma.form) { - tcgen05_mma_selection_asm( - mma.form, - Tcgen05MmaKind::F8f6f4, - 1, - None, - Some(0), - Some(Tcgen05MmaBUsage::Discard), - ) - } else { - tcgen05_mma_selection_asm( - mma.form, - Tcgen05MmaKind::F8f6f4, - 1, - Some("discard"), - None, - None, - ) - }]) - } else { - tcgen05_mma_valid_selection_asms(mma.form) - }; - return Ok(expected.contains(&selection.asm) - && selection.predicates - == [if selection.asm.contains(".kind::i8.") { - "Subtarget->hasTcgen05MMAI8Kind()" - } else { - "Subtarget->hasTcgen05InstSupport()" - }] - && selection.constraints.is_empty()); - } - if tcgen05.source_contract != Tcgen05SourceContract::ExactTablegenSelection { - return Ok(false); - } - if let Some(cp) = tcgen05.cp { - let recipe = tcgen05_cp_member_recipe(cp.member); - let group = match cp.group { - Tcgen05CpGroup::Cg1 => 1, - Tcgen05CpGroup::Cg2 => 2, - }; - return Ok(selection.source_record - == format!("TCGEN05_CP_{}_cg{group}", recipe.selection_stem) - && selection.asm - == format!( - "tcgen05.cp.cta_group::{group}.{} \t[$tmem_addr], $sdesc;", - recipe.ptx_suffix - ) - && selection.predicates == ["Subtarget->hasTcgen05InstSupport()"] - && selection.constraints.is_empty()); - } - let recipe = tcgen05_recipe(tcgen05.operation); - return Ok( - recipe.selection_record == Some(selection.source_record.as_str()) - && recipe.selection_asm == Some(selection.asm.as_str()) - && selection.constraints.is_empty(), - ); - } - if policy.family == "sparse_mma" { - let Some(last) = policy.expected_ptx.operands.last() else { - return Ok(false); - }; - if *last != OperandPattern::Immediate { - return Ok(false); - } - let mut selection_shape = policy.expected_ptx.clone(); - *selection_shape.operands.last_mut().unwrap() = OperandPattern::RegisterOrImmediate; - return Ok(selection_shape.matches(&selection.asm)? && selection.constraints.is_empty()); - } - if policy.family == "sync" { - if policy.id == "sync_threads" { - return Ok(selection.source_record == "BARRIER_CTA_SYNC_ALIGNED_ALL_i" - && selection.asm == "bar.sync \t$i;" - && selection.predicates.is_empty() - && selection.constraints.is_empty()); - } - let Some(scope) = threadfence_scope_for_id(&policy.id) else { - return Ok(false); - }; - let recipe = threadfence_recipe(scope); - return Ok(selection.source_record == recipe.selection_record - && selection.asm == format!("membar.{};", recipe.ptx_level) - && selection.predicates.is_empty() - && selection.constraints.is_empty()); - } - - if policy.family == "vote" { - let Some(vote) = &policy.vote else { - return Ok(false); - }; - let recipe = vote_recipe(vote.mode); - return Ok([recipe.immediate_selection, recipe.register_selection] - .contains(&selection.source_record.as_str()) - && policy.expected_ptx.matches(&selection.asm)? - && selection.constraints.address_space.is_none() - && selection.constraints.immediate_bindings.is_empty()); - } - - if policy.family == "warp_match" { - let Some(warp_match) = &policy.warp_match else { - return Ok(false); - }; - let recipe = warp_match_recipe(warp_match.mode, warp_match.value_width); - return Ok(recipe - .selections - .contains(&selection.source_record.as_str()) - && policy.expected_ptx.matches(&selection.asm)? - && selection.constraints.is_empty()); - } - - if policy.family == "elect" { - return Ok(["INT_ELECT_SYNC_I", "INT_ELECT_SYNC_R"] - .contains(&selection.source_record.as_str()) - && selection.asm == "elect.sync \t$dest|$pred, $mask;" - && selection.constraints.is_empty()); - } - - if policy.family == "warp_barrier" { - return Ok(policy.warp_barrier.is_some() - && ["INT_BAR_WARP_SYNC_I", "INT_BAR_WARP_SYNC_R"] - .contains(&selection.source_record.as_str()) - && policy.expected_ptx.matches(&selection.asm)? - && selection.constraints.is_empty()); - } - - if policy.family == "warp_shuffle" { - let Some(shuffle) = &policy.warp_shuffle else { - return Ok(false); - }; - let recipe = warp_shuffle_recipe(shuffle.mode, shuffle.value_kind); - return Ok(selection.asm - == format!( - "shfl.sync.{}.b32 \t$dst, $src, $offset, $mask, $threadmask;", - recipe.ptx_mode - ) - && selection.constraints.is_empty()); - } - - if policy.family == "cp_async_copy" { - let Some(copy) = &policy.cp_async_copy else { - return Ok(false); - }; - let Some(recipe) = cp_async_copy_recipe(copy) else { - return Ok(false); - }; - return Ok(recipe - .selections - .contains(&selection.source_record.as_str()) - && policy.expected_ptx.matches(&selection.asm)? - && selection.constraints.is_empty()); - } - - if policy.family == "cp_async_control" { - let Some(control) = &policy.cp_async_control else { - return Ok(false); - }; - let recipe = cp_async_control_recipe(control.operation); - let instruction_matches = if control.operation == CpAsyncControlOperation::WaitGroup { - selection.asm == "cp.async.wait_group \t$n;" - } else { - policy.expected_ptx.matches(&selection.asm)? - }; - return Ok(selection.source_record == recipe.selection - && instruction_matches - && selection.constraints.is_empty()); - } - - if policy.family == "cp_async_mbarrier" { - let Some(bridge) = &policy.cp_async_mbarrier else { - return Ok(false); - }; - let recipe = cp_async_mbarrier_recipe(bridge.operation, bridge.state_space); - return Ok(selection.source_record == recipe.selection - && selection.asm == recipe.selection_asm - && selection.constraints.is_empty()); - } - - if policy.family == "mbarrier_basic" { - let Some(mbarrier) = &policy.mbarrier_basic else { - return Ok(false); - }; - let recipe = mbarrier_basic_recipe(mbarrier.operation); - return Ok(selection.source_record == recipe.selection - && policy.expected_ptx.matches(&selection.asm)? - && selection.constraints.is_empty()); - } - - if !policy.expected_ptx.matches(&selection.asm)? - || policy - .selected_address_space - .is_some_and(|address_space| selection.constraints.address_space != Some(address_space)) - { - return Ok(false); - } - - let Some(dot_product) = &policy.dot_product else { - return Ok(true); - }; - if selection.constraints.address_space.is_some() { - return Ok(false); - } - Ok(match dot_product.adapter { - DotProductAdapter::DirectThreeOperands => { - selection.constraints.immediate_bindings.is_empty() - } - DotProductAdapter::InsertLowHalfFalse => { - selection.constraints.immediate_bindings.len() == 1 - && selection.constraints.immediate_bindings[0].argument_index == 2 - && selection.constraints.immediate_bindings[0].value == 0 - } - }) -} - -fn selection_matches_tma_policy( - policy: &OverlayIntrinsic, - selection: &crate::model::ImportedSelection, -) -> bool { - let Some(tma) = &policy.tma else { - return false; - }; - let operation = tma.operation; - if matches!( - operation, - TmaOperation::Reduce - | TmaOperation::PrefetchTensorMap - | TmaOperation::ReplaceBoxDim - | TmaOperation::ReplaceElementStride - | TmaOperation::ReplaceElementType - | TmaOperation::ReplaceFillMode - | TmaOperation::ReplaceGlobalAddress - | TmaOperation::ReplaceGlobalDim - | TmaOperation::ReplaceGlobalStride - | TmaOperation::ReplaceInterleaveLayout - | TmaOperation::ReplaceRank - | TmaOperation::ReplaceSwizzleAtomicity - | TmaOperation::ReplaceSwizzleMode - ) { - return false; - } - let (source_record, asm) = match operation { - TmaOperation::G2sTile1d => ( - "TMA_G2S_TILE_CG0_1D", - "cp.async.bulk.tensor.1d.shared::cluster.global.tile.mbarrier::complete_tx::bytes$cg [$dst], [$tmap, {{$d0}}], [$mbar];", - ), - TmaOperation::G2sTile2d => ( - "TMA_G2S_TILE_CG0_2D", - "cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes$cg [$dst], [$tmap, {{$d0, $d1}}], [$mbar];", - ), - TmaOperation::G2sTile2dMulticast => ( - "TMA_G2S_TILE_CG0_2D_MC", - "cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.multicast::cluster$cg [$dst], [$tmap, {{$d0, $d1}}], [$mbar], $mc;", - ), - TmaOperation::G2sTile2dMulticastCg2 => ( - "TMA_G2S_TILE_2D_MC", - "cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.multicast::cluster$cg [$dst], [$tmap, {{$d0, $d1}}], [$mbar], $mc;", - ), - TmaOperation::G2sTile3d => ( - "TMA_G2S_TILE_CG0_3D", - "cp.async.bulk.tensor.3d.shared::cluster.global.tile.mbarrier::complete_tx::bytes$cg [$dst], [$tmap, {{$d0, $d1, $d2}}], [$mbar];", - ), - TmaOperation::G2sTile4d => ( - "TMA_G2S_TILE_CG0_4D", - "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes$cg [$dst], [$tmap, {{$d0, $d1, $d2, $d3}}], [$mbar];", - ), - TmaOperation::G2sTile5d => ( - "TMA_G2S_TILE_CG0_5D", - "cp.async.bulk.tensor.5d.shared::cluster.global.tile.mbarrier::complete_tx::bytes$cg [$dst], [$tmap, {{$d0, $d1, $d2, $d3, $d4}}], [$mbar];", - ), - TmaOperation::S2gTile1d => ( - "TMA_TENSOR_S2G_TILE_1D", - "cp.async.bulk.tensor.1d.global.shared::cta.tile.bulk_group [$tmap, {{$d0}}], [$src];", - ), - TmaOperation::S2gTile2d => ( - "TMA_TENSOR_S2G_TILE_2D", - "cp.async.bulk.tensor.2d.global.shared::cta.tile.bulk_group [$tmap, {{$d0, $d1}}], [$src];", - ), - TmaOperation::S2gTile3d => ( - "TMA_TENSOR_S2G_TILE_3D", - "cp.async.bulk.tensor.3d.global.shared::cta.tile.bulk_group [$tmap, {{$d0, $d1, $d2}}], [$src];", - ), - TmaOperation::S2gTile4d => ( - "TMA_TENSOR_S2G_TILE_4D", - "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group [$tmap, {{$d0, $d1, $d2, $d3}}], [$src];", - ), - TmaOperation::S2gTile5d => ( - "TMA_TENSOR_S2G_TILE_5D", - "cp.async.bulk.tensor.5d.global.shared::cta.tile.bulk_group [$tmap, {{$d0, $d1, $d2, $d3, $d4}}], [$src];", - ), - TmaOperation::CommitGroup => ("CP_ASYNC_BULK_COMMIT_GROUP", "cp.async.bulk.commit_group;"), - TmaOperation::WaitGroup => ("CP_ASYNC_BULK_WAIT_GROUP", "cp.async.bulk.wait_group \t$n;"), - TmaOperation::WaitGroupRead => ( - "CP_ASYNC_BULK_WAIT_GROUP_READ", - "cp.async.bulk.wait_group.read \t$n;", - ), - TmaOperation::PrefetchTile1d => ( - "TMA_TENSOR_PF_TILE_1D", - "cp.async.bulk.prefetch.tensor.1d.L2.global.tile [$tmap, {{$d0}}];", - ), - TmaOperation::PrefetchTile2d => ( - "TMA_TENSOR_PF_TILE_2D", - "cp.async.bulk.prefetch.tensor.2d.L2.global.tile [$tmap, {{$d0, $d1}}];", - ), - TmaOperation::PrefetchTile3d => ( - "TMA_TENSOR_PF_TILE_3D", - "cp.async.bulk.prefetch.tensor.3d.L2.global.tile [$tmap, {{$d0, $d1, $d2}}];", - ), - TmaOperation::PrefetchTile4d => ( - "TMA_TENSOR_PF_TILE_4D", - "cp.async.bulk.prefetch.tensor.4d.L2.global.tile [$tmap, {{$d0, $d1, $d2, $d3}}];", - ), - TmaOperation::PrefetchTile5d => ( - "TMA_TENSOR_PF_TILE_5D", - "cp.async.bulk.prefetch.tensor.5d.L2.global.tile [$tmap, {{$d0, $d1, $d2, $d3, $d4}}];", - ), - TmaOperation::PrefetchTileGather4TwoDimensional => ( - "TMA_TENSOR_PF_TILE_GATHER4_2D", - "cp.async.bulk.prefetch.tensor.2d.L2.global.tile::gather4 [$tmap, {{$d0, $d1, $d2, $d3, $d4}}];", - ), - TmaOperation::PrefetchTile1dCacheHint => ( - "TMA_TENSOR_PF_TILE_1D_CH", - "cp.async.bulk.prefetch.tensor.1d.L2.global.tile.L2::cache_hint [$tmap, {{$d0}}], $ch;", - ), - TmaOperation::PrefetchTile2dCacheHint => ( - "TMA_TENSOR_PF_TILE_2D_CH", - "cp.async.bulk.prefetch.tensor.2d.L2.global.tile.L2::cache_hint [$tmap, {{$d0, $d1}}], $ch;", - ), - TmaOperation::PrefetchTile3dCacheHint => ( - "TMA_TENSOR_PF_TILE_3D_CH", - "cp.async.bulk.prefetch.tensor.3d.L2.global.tile.L2::cache_hint [$tmap, {{$d0, $d1, $d2}}], $ch;", - ), - TmaOperation::PrefetchTile4dCacheHint => ( - "TMA_TENSOR_PF_TILE_4D_CH", - "cp.async.bulk.prefetch.tensor.4d.L2.global.tile.L2::cache_hint [$tmap, {{$d0, $d1, $d2, $d3}}], $ch;", - ), - TmaOperation::PrefetchTile5dCacheHint => ( - "TMA_TENSOR_PF_TILE_5D_CH", - "cp.async.bulk.prefetch.tensor.5d.L2.global.tile.L2::cache_hint [$tmap, {{$d0, $d1, $d2, $d3, $d4}}], $ch;", - ), - TmaOperation::PrefetchTileGather4TwoDimensionalCacheHint => ( - "TMA_TENSOR_PF_TILE_GATHER4_2D_CH", - "cp.async.bulk.prefetch.tensor.2d.L2.global.tile::gather4.L2::cache_hint [$tmap, {{$d0, $d1, $d2, $d3, $d4}}], $ch;", - ), - TmaOperation::FenceProxyTensorMapAcquireCluster => ( - "INT_FENCE_PROXY_TENSORMAP_GENERIC_ACQUIRE_CLUSTER", - "fence.proxy.tensormap::generic.acquire.cluster [$addr], 128;", - ), - TmaOperation::FenceProxyTensorMapAcquireCta => ( - "INT_FENCE_PROXY_TENSORMAP_GENERIC_ACQUIRE_CTA", - "fence.proxy.tensormap::generic.acquire.cta [$addr], 128;", - ), - TmaOperation::FenceProxyTensorMapAcquireGpu => ( - "INT_FENCE_PROXY_TENSORMAP_GENERIC_ACQUIRE_GPU", - "fence.proxy.tensormap::generic.acquire.gpu [$addr], 128;", - ), - TmaOperation::FenceProxyTensorMapAcquireSystem => ( - "INT_FENCE_PROXY_TENSORMAP_GENERIC_ACQUIRE_SYS", - "fence.proxy.tensormap::generic.acquire.sys [$addr], 128;", - ), - TmaOperation::FenceProxyTensorMapReleaseCluster => ( - "INT_FENCE_PROXY_TENSORMAP_GENERIC_RELEASE_CLUSTER", - "fence.proxy.tensormap::generic.release.cluster;", - ), - TmaOperation::FenceProxyTensorMapReleaseCta => ( - "INT_FENCE_PROXY_TENSORMAP_GENERIC_RELEASE_CTA", - "fence.proxy.tensormap::generic.release.cta;", - ), - TmaOperation::FenceProxyTensorMapReleaseGpu => ( - "INT_FENCE_PROXY_TENSORMAP_GENERIC_RELEASE_GPU", - "fence.proxy.tensormap::generic.release.gpu;", - ), - TmaOperation::FenceProxyTensorMapReleaseSystem => ( - "INT_FENCE_PROXY_TENSORMAP_GENERIC_RELEASE_SYS", - "fence.proxy.tensormap::generic.release.sys;", - ), - TmaOperation::Reduce - | TmaOperation::PrefetchTensorMap - | TmaOperation::ReplaceBoxDim - | TmaOperation::ReplaceElementStride - | TmaOperation::ReplaceElementType - | TmaOperation::ReplaceFillMode - | TmaOperation::ReplaceGlobalAddress - | TmaOperation::ReplaceGlobalDim - | TmaOperation::ReplaceGlobalStride - | TmaOperation::ReplaceInterleaveLayout - | TmaOperation::ReplaceRank - | TmaOperation::ReplaceSwizzleAtomicity - | TmaOperation::ReplaceSwizzleMode => unreachable!("selectionless TMA operation"), - }; - - let mut immediate_bindings = Vec::new(); - if matches!( - operation, - TmaOperation::G2sTile1d - | TmaOperation::G2sTile2d - | TmaOperation::G2sTile2dMulticast - | TmaOperation::G2sTile2dMulticastCg2 - | TmaOperation::G2sTile3d - | TmaOperation::G2sTile4d - | TmaOperation::G2sTile5d - ) { - let dimensions = operation.dimensions().unwrap(); - immediate_bindings.push(crate::model::ImportedImmediateBinding { - argument_index: dimensions + 5, - value: if matches!( - operation, - TmaOperation::G2sTile2dMulticast | TmaOperation::G2sTile2dMulticastCg2 - ) { - -1 - } else { - 0 - }, - }); - immediate_bindings.push(crate::model::ImportedImmediateBinding { - argument_index: dimensions + 6, - value: 0, - }); - } else if matches!( - operation, - TmaOperation::S2gTile1d - | TmaOperation::S2gTile2d - | TmaOperation::S2gTile3d - | TmaOperation::S2gTile4d - | TmaOperation::S2gTile5d - ) { - immediate_bindings.push(crate::model::ImportedImmediateBinding { - argument_index: operation.dimensions().unwrap() + 3, - value: 0, - }); - } else if let Some(coordinate_count) = operation.prefetch_coordinate_count() { - immediate_bindings.push(crate::model::ImportedImmediateBinding { - argument_index: coordinate_count + 2, - value: if operation.uses_prefetch_cache_hint() { - -1 - } else { - 0 - }, - }); - } - - selection.source_record == source_record - && selection.asm == asm - && selection.constraints - == crate::model::ImportedSelectionConstraints { - address_space: None, - immediate_bindings, - } -} - -#[derive(Clone, Copy)] -struct ThreadfenceRecipe { - scope: ThreadfenceScope, - abi_id: &'static str, - id: &'static str, - operation_key: &'static str, - source_record: &'static str, - selection_record: &'static str, - llvm_symbol: &'static str, - dialect_op_type: &'static str, - dialect_op_name: &'static str, - ptx_level: &'static str, - execution_scope: &'static str, - minimum_ptx: &'static str, - minimum_sm: Option<&'static str>, - summary: &'static str, -} - -fn threadfence_recipe(scope: ThreadfenceScope) -> ThreadfenceRecipe { - match scope { - ThreadfenceScope::Cta => ThreadfenceRecipe { - scope, - abi_id: "i0298", - id: "threadfence_block", - operation_key: "memory.fence.cta.sc", - source_record: "int_nvvm_membar_cta", - selection_record: "INT_MEMBAR_CTA", - llvm_symbol: "llvm.nvvm.membar.cta", - dialect_op_type: "ThreadfenceBlockOp", - dialect_op_name: "nvvm.threadfence_block", - ptx_level: "cta", - execution_scope: "cta", - minimum_ptx: "1.4", - minimum_sm: None, - summary: "Orders this thread's memory operations for observers in its CTA.", - }, - ThreadfenceScope::Device => ThreadfenceRecipe { - scope, - abi_id: "i0299", - id: "threadfence", - operation_key: "memory.fence.device.sc", - source_record: "int_nvvm_membar_gl", - selection_record: "INT_MEMBAR_GL", - llvm_symbol: "llvm.nvvm.membar.gl", - dialect_op_type: "ThreadfenceOp", - dialect_op_name: "nvvm.threadfence", - ptx_level: "gl", - execution_scope: "device", - minimum_ptx: "1.4", - minimum_sm: None, - summary: "Orders this thread's memory operations for observers on its GPU.", - }, - ThreadfenceScope::System => ThreadfenceRecipe { - scope, - abi_id: "i0300", - id: "threadfence_system", - operation_key: "memory.fence.system.sc", - source_record: "int_nvvm_membar_sys", - selection_record: "INT_MEMBAR_SYS", - llvm_symbol: "llvm.nvvm.membar.sys", - dialect_op_type: "ThreadfenceSystemOp", - dialect_op_name: "nvvm.threadfence_system", - ptx_level: "sys", - execution_scope: "system", - minimum_ptx: "2.0", - minimum_sm: Some("sm_20"), - summary: "Orders this thread's memory operations for system-wide observers.", - }, - } -} - -fn threadfence_scope_for_id(id: &str) -> Option { - match id { - "threadfence_block" => Some(ThreadfenceScope::Cta), - "threadfence" => Some(ThreadfenceScope::Device), - "threadfence_system" => Some(ThreadfenceScope::System), - _ => None, - } -} - -fn expand_threadfence_admission(admission: &ThreadfenceAdmission) -> Result> { - ensure!( - admission.runtime_validation == RuntimeValidation::Unexecuted, - "threadfence runtime validation may be marked executed only with GPU evidence" - ); - ensure!( - !admission.llvm_evidence_profile.trim().is_empty() - && !admission.libnvvm_evidence_profile.trim().is_empty(), - "compact threadfence admission requires both backend evidence profiles" - ); - let expected_scopes = [ - ThreadfenceScope::Cta, - ThreadfenceScope::Device, - ThreadfenceScope::System, - ]; - let actual_scopes = admission - .variants - .iter() - .map(|variant| variant.scope) - .collect::>(); - ensure!( - actual_scopes == expected_scopes, - "compact threadfence admission must contain each reviewed scope exactly once in canonical order" - ); - - admission - .variants - .iter() - .map(|variant| { - let recipe = threadfence_recipe(variant.scope); - ensure!( - variant.abi_id == recipe.abi_id, - "{} must keep reserved ABI ID {}", - recipe.id, - recipe.abi_id - ); - Ok(OverlayIntrinsic { - id: recipe.id.into(), - abi_id: variant.abi_id.clone(), - operation_key: recipe.operation_key.into(), - family: "sync".into(), - source: None, - source_record: Some(recipe.source_record.into()), - rust_module: "fence".into(), - rust_name: recipe.id.into(), - rust_arguments: vec![], - rust_result: "()".into(), - safe: true, - must_use: false, - safe_allowlist_reason: Some( - "a fence only orders the calling thread's memory operations and has no caller preconditions" - .into(), - ), - public_rust_path: format!("cuda_intrinsics::fence::{}", recipe.id), - compatibility_rust_paths: vec![ - format!("cuda_device::fence::{}", recipe.id), - format!("cuda_device::{}", recipe.id), - ], - dialect_op_type: recipe.dialect_op_type.into(), - dialect_op_name: recipe.dialect_op_name.into(), - dialect_operands: vec![], - dialect_results: vec![], - llvm_symbol: Some(recipe.llvm_symbol.into()), - resolved_llvm_symbol: None, - llvm_arguments: vec![], - llvm_results: vec![], - pure: false, - memory: "read_write".into(), - convergent: false, - execution_scope: recipe.execution_scope.into(), - minimum_ptx: recipe.minimum_ptx.into(), - minimum_sm: recipe.minimum_sm.map(Into::into), - ptx_result: "()".into(), - targets: "all".into(), - ptx_isa_version: "9.3".into(), - ptx_isa_section: - "9.7.14.4 Parallel Synchronization and Communication Instructions: membar / fence" - .into(), - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-membar-fence".into(), - lowering: "direct_nvvm".into(), - backend_lowerings: vec![ - OverlayBackendLowering { - backend: IntrinsicBackend::LlvmNvptx, - mechanism: BackendLoweringMechanism::TypedNvvm, - evidence_profile: admission.llvm_evidence_profile.clone(), - targets: None, - minimum_ptx: Some("3.2".into()), - minimum_sm: Some("sm_20".into()), - }, - OverlayBackendLowering { - backend: IntrinsicBackend::LibNvvm, - mechanism: BackendLoweringMechanism::TypedNvvm, - evidence_profile: admission.libnvvm_evidence_profile.clone(), - targets: None, - minimum_ptx: Some("7.0".into()), - minimum_sm: Some("sm_80".into()), - }, - ], - packed_atomic: None, - redux: None, - vote: None, - active_mask: None, - warp_match: None, - warp_barrier: None, - warp_shuffle: None, - dot_product: None, - packed_alu: None, - integer_minmax: None, - packed_conversion: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cp_async_copy: None, - cp_async_control: None, - cp_async_mbarrier: None, - mbarrier_basic: None, - movmatrix: None, - mbarrier_extended: None, - register_mma: None, - sparse_mma: None, - prmt: None, - cluster_barrier: None, - wgmma_control: None, - special_register: None, - debug_control: None, - cluster_memory: None, - clc: None, - tma: None, - tcgen05: None, - ldmatrix_variant: None, - ldmatrix_safety: None, - ldmatrix_adapter: None, - selected_address_space: None, - expected_ptx: InstructionPattern { - mnemonic: "membar".into(), - modifiers: vec![recipe.ptx_level.into()], - operands: vec![], - }, - summary: recipe.summary.into(), - }) - }) - .collect() -} - -fn validate_threadfence_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, - scope: ThreadfenceScope, -) -> Result<()> { - let recipe = threadfence_recipe(scope); - ensure!( - recipe.scope == scope - && policy.id == recipe.id - && policy.abi_id == recipe.abi_id - && policy.operation_key == recipe.operation_key - && policy.source.is_none() - && policy.source_record.as_deref() == Some(recipe.source_record) - && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) - && policy.resolved_llvm_symbol.is_none(), - "{} threadfence identity does not match its closed recipe", - policy.id - ); - ensure!( - policy.rust_module == "fence" - && policy.rust_name == recipe.id - && policy.rust_arguments.is_empty() - && policy.rust_result == "()" - && policy.safe - && !policy.must_use - && policy.public_rust_path == format!("cuda_intrinsics::fence::{}", recipe.id) - && policy.compatibility_rust_paths - == [ - format!("cuda_device::fence::{}", recipe.id), - format!("cuda_device::{}", recipe.id), - ], - "{} threadfence Rust API does not match its closed recipe", - policy.id - ); - ensure!( - policy.dialect_op_type == recipe.dialect_op_type - && policy.dialect_op_name == recipe.dialect_op_name - && policy.dialect_operands.is_empty() - && policy.dialect_results.is_empty() - && policy.llvm_arguments.is_empty() - && policy.llvm_results.is_empty() - && policy.lowering == "direct_nvvm", - "{} threadfence carrier or lowering does not match its closed recipe", - policy.id - ); - ensure!( - declaration.properties == ["IntrNoCallback"] - && !policy.pure - && policy.memory == "read_write" - && !policy.convergent - && policy.execution_scope == recipe.execution_scope, - "{} threadfence effects disagree with the imported declaration", - policy.id - ); - ensure!( - policy.minimum_ptx == recipe.minimum_ptx - && policy.minimum_sm.as_deref() == recipe.minimum_sm - && policy.ptx_result == "()" - && policy.targets == "all" - && policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section - == "9.7.14.4 Parallel Synchronization and Communication Instructions: membar / fence" - && policy.ptx_isa_url - == "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-membar-fence", - "{} threadfence target floor or PTX provenance changed", - policy.id - ); - ensure!( - policy.expected_ptx - == (InstructionPattern { - mnemonic: "membar".into(), - modifiers: vec![recipe.ptx_level.into()], - operands: vec![], - }), - "{} expected PTX does not match its closed threadfence scope", - policy.id - ); - ensure!( - policy.backend_lowerings.len() == 2 - && policy.backend_lowerings.iter().any(|route| { - route.backend == IntrinsicBackend::LlvmNvptx - && route.mechanism == BackendLoweringMechanism::TypedNvvm - && route.minimum_ptx.as_deref() == Some("3.2") - && route.minimum_sm.as_deref() == Some("sm_20") - }) - && policy.backend_lowerings.iter().any(|route| { - route.backend == IntrinsicBackend::LibNvvm - && route.mechanism == BackendLoweringMechanism::TypedNvvm - && route.minimum_ptx.as_deref() == Some("7.0") - && route.minimum_sm.as_deref() == Some("sm_80") - }), - "{} must keep both reviewed typed threadfence routes", - policy.id - ); - ensure!( - policy.ldmatrix_variant.is_none() - && policy.ldmatrix_safety.is_none() - && policy.ldmatrix_adapter.is_none() - && policy.packed_atomic.is_none() - && policy.redux.is_none() - && policy.vote.is_none() - && policy.active_mask.is_none() - && policy.warp_match.is_none() - && policy.warp_barrier.is_none() - && policy.warp_shuffle.is_none() - && policy.dot_product.is_none() - && policy.packed_alu.is_none() - && policy.packed_conversion.is_none() - && policy.cp_async_copy.is_none() - && policy.cp_async_control.is_none() - && policy.cp_async_mbarrier.is_none() - && policy.mbarrier_basic.is_none() - && policy.register_mma.is_none() - && policy.sparse_mma.is_none() - && policy.prmt.is_none() - && policy.cluster_barrier.is_none() - && policy.special_register.is_none() - && policy.debug_control.is_none() - && policy.selected_address_space.is_none(), - "{} mixes another generated-family contract with threadfence", - policy.id - ); - Ok(()) -} - -#[derive(Clone, Copy)] -struct ClcRecipe { - operation: ClcOperation, - abi_id: &'static str, - id: &'static str, - operation_key: &'static str, - source_record: &'static str, - llvm_symbol: &'static str, - rust_arguments: &'static [&'static str], - llvm_arguments: &'static [&'static str], - llvm_results: &'static [&'static str], - dialect_op_type: &'static str, - dialect_op_name: &'static str, - dialect_operands: &'static [&'static str], - dialect_results: &'static [&'static str], - adapter: ClcAdapter, - modifiers: &'static [&'static str], - operands: &'static [OperandPattern], - targets: &'static str, - minimum_sm: Option<&'static str>, - pure: bool, - memory: &'static str, - execution_scope: &'static str, - summary: &'static str, -} - -fn clc_recipe(operation: ClcOperation) -> ClcRecipe { - const TRY_OPERANDS: &[OperandPattern] = &[OperandPattern::Address, OperandPattern::Address]; - const QUERY_OPERANDS: &[OperandPattern] = &[OperandPattern::Register, OperandPattern::Register]; - let (abi_id, id, operation_key, source_record, llvm_symbol, op_type, op_name) = match operation - { - ClcOperation::TryCancel => ( - "i0322", - "clc_try_cancel", - "cluster.launch_control.try_cancel", - "int_nvvm_clusterlaunchcontrol_try_cancel_async_shared", - "llvm.nvvm.clusterlaunchcontrol.try_cancel.async.shared", - "ClcTryCancelOp", - "nvvm.clc_try_cancel", - ), - ClcOperation::TryCancelMulticast => ( - "i0323", - "clc_try_cancel_multicast", - "cluster.launch_control.try_cancel.multicast", - "int_nvvm_clusterlaunchcontrol_try_cancel_async_multicast_shared", - "llvm.nvvm.clusterlaunchcontrol.try_cancel.async.multicast.shared", - "ClcTryCancelMulticastOp", - "nvvm.clc_try_cancel_multicast", - ), - ClcOperation::QueryIsCanceled => ( - "i0324", - "clc_query_is_canceled", - "cluster.launch_control.query.is_canceled", - "int_nvvm_clusterlaunchcontrol_query_cancel_is_canceled", - "llvm.nvvm.clusterlaunchcontrol.query_cancel.is_canceled", - "ClcQueryIsCanceledOp", - "nvvm.clc_query_is_canceled", - ), - ClcOperation::QueryGetFirstCtaidX => ( - "i0325", - "clc_query_get_first_ctaid_x", - "cluster.launch_control.query.first_ctaid.x", - "int_nvvm_clusterlaunchcontrol_query_cancel_get_first_ctaid_x", - "llvm.nvvm.clusterlaunchcontrol.query_cancel.get_first_ctaid.x", - "ClcQueryGetFirstCtaidXOp", - "nvvm.clc_query_get_first_ctaid_x", - ), - ClcOperation::QueryGetFirstCtaidY => ( - "i0326", - "clc_query_get_first_ctaid_y", - "cluster.launch_control.query.first_ctaid.y", - "int_nvvm_clusterlaunchcontrol_query_cancel_get_first_ctaid_y", - "llvm.nvvm.clusterlaunchcontrol.query_cancel.get_first_ctaid.y", - "ClcQueryGetFirstCtaidYOp", - "nvvm.clc_query_get_first_ctaid_y", - ), - ClcOperation::QueryGetFirstCtaidZ => ( - "i0327", - "clc_query_get_first_ctaid_z", - "cluster.launch_control.query.first_ctaid.z", - "int_nvvm_clusterlaunchcontrol_query_cancel_get_first_ctaid_z", - "llvm.nvvm.clusterlaunchcontrol.query_cancel.get_first_ctaid.z", - "ClcQueryGetFirstCtaidZOp", - "nvvm.clc_query_get_first_ctaid_z", - ), - }; - match operation { - ClcOperation::TryCancel | ClcOperation::TryCancelMulticast => ClcRecipe { - operation, - abi_id, - id, - operation_key, - source_record, - llvm_symbol, - rust_arguments: &["*mut u8", "*mut u64"], - llvm_arguments: &["shared_ptr", "shared_ptr"], - llvm_results: &[], - dialect_op_type: op_type, - dialect_op_name: op_name, - dialect_operands: &["ptr", "ptr"], - dialect_results: &[], - adapter: ClcAdapter::GenericPointersToShared, - modifiers: if operation == ClcOperation::TryCancel { - &[ - "try_cancel", - "async", - "shared::cta", - "mbarrier::complete_tx::bytes", - "b128", - ] - } else { - &[ - "try_cancel", - "async", - "shared::cta", - "mbarrier::complete_tx::bytes", - "multicast::cluster::all", - "b128", - ] - }, - operands: TRY_OPERANDS, - targets: if operation == ClcOperation::TryCancel { - "all" - } else { - // LLVM 22 exposes 101a; CUDA 13.3 exposes the other toolkit names. - "sm_100a|sm_101a|sm_103a|sm_110a|sm_120a|sm_121a" - }, - minimum_sm: if operation == ClcOperation::TryCancel { - Some("sm_100") - } else { - None - }, - pure: false, - memory: "read_write", - execution_scope: "cta", - summary: if operation == ClcOperation::TryCancel { - "Requests one pending CTA and writes its response to shared memory." - } else { - "Requests one pending CTA and multicasts its response across the cluster." - }, - }, - ClcOperation::QueryIsCanceled - | ClcOperation::QueryGetFirstCtaidX - | ClcOperation::QueryGetFirstCtaidY - | ClcOperation::QueryGetFirstCtaidZ => { - let (adapter, modifiers, summary) = match operation { - ClcOperation::QueryIsCanceled => ( - ClcAdapter::PairU64ToI128BoolToU32, - &["query_cancel", "is_canceled", "pred", "b128"] as &[_], - "Returns whether the Cluster Launch Control request was canceled.", - ), - ClcOperation::QueryGetFirstCtaidX => ( - ClcAdapter::PairU64ToI128U32, - &["query_cancel", "get_first_ctaid::x", "b32", "b128"] as &[_], - "Returns the X coordinate from a successful cancellation response.", - ), - ClcOperation::QueryGetFirstCtaidY => ( - ClcAdapter::PairU64ToI128U32, - &["query_cancel", "get_first_ctaid::y", "b32", "b128"] as &[_], - "Returns the Y coordinate from a successful cancellation response.", - ), - ClcOperation::QueryGetFirstCtaidZ => ( - ClcAdapter::PairU64ToI128U32, - &["query_cancel", "get_first_ctaid::z", "b32", "b128"] as &[_], - "Returns the Z coordinate from a successful cancellation response.", - ), - _ => unreachable!(), - }; - ClcRecipe { - operation, - abi_id, - id, - operation_key, - source_record, - llvm_symbol, - rust_arguments: &["u64", "u64"], - llvm_arguments: &["i128"], - llvm_results: if operation == ClcOperation::QueryIsCanceled { - &["i1"] - } else { - &["i32"] - }, - dialect_op_type: op_type, - dialect_op_name: op_name, - dialect_operands: &["i64", "i64"], - dialect_results: &["i32"], - adapter, - modifiers, - operands: QUERY_OPERANDS, - targets: "all", - minimum_sm: Some("sm_100"), - pure: true, - memory: "none", - execution_scope: "thread", - summary, - } - } - } -} - -fn expand_clc_admission(admission: &ClcAdmission) -> Result> { - ensure!( - admission.runtime_validation == RuntimeValidation::Unexecuted, - "CLC runtime validation may be marked executed only with GPU evidence" - ); - ensure!( - !admission.llvm_evidence_profile.trim().is_empty() - && !admission.libnvvm_evidence_profile.trim().is_empty(), - "compact CLC admission requires both backend evidence profiles" - ); - let expected = [ - ClcOperation::TryCancel, - ClcOperation::TryCancelMulticast, - ClcOperation::QueryIsCanceled, - ClcOperation::QueryGetFirstCtaidX, - ClcOperation::QueryGetFirstCtaidY, - ClcOperation::QueryGetFirstCtaidZ, - ]; - ensure!( - admission - .variants - .iter() - .map(|variant| variant.operation) - .eq(expected), - "compact CLC admission must list all six operations in canonical order" - ); - - admission - .variants - .iter() - .map(|variant| { - let recipe = clc_recipe(variant.operation); - ensure!( - variant.abi_id == recipe.abi_id, - "{} must keep reserved ABI ID {}", - recipe.id, - recipe.abi_id - ); - let query = recipe.pure; - Ok(OverlayIntrinsic { - id: recipe.id.into(), - abi_id: variant.abi_id.clone(), - operation_key: recipe.operation_key.into(), - family: "clc".into(), - source: None, - source_record: Some(recipe.source_record.into()), - rust_module: "clc".into(), - rust_name: recipe.id.into(), - rust_arguments: recipe.rust_arguments.iter().map(|value| (*value).into()).collect(), - rust_result: if query { "u32".into() } else { "()".into() }, - safe: false, - must_use: false, - safe_allowlist_reason: None, - public_rust_path: format!("cuda_intrinsics::clc::{}", recipe.id), - compatibility_rust_paths: vec![format!("cuda_device::clc::{}", recipe.id)], - dialect_op_type: recipe.dialect_op_type.into(), - dialect_op_name: recipe.dialect_op_name.into(), - dialect_operands: recipe.dialect_operands.iter().map(|value| (*value).into()).collect(), - dialect_results: recipe.dialect_results.iter().map(|value| (*value).into()).collect(), - llvm_symbol: Some(recipe.llvm_symbol.into()), - resolved_llvm_symbol: None, - llvm_arguments: recipe.llvm_arguments.iter().map(|value| (*value).into()).collect(), - llvm_results: recipe.llvm_results.iter().map(|value| (*value).into()).collect(), - pure: recipe.pure, - memory: recipe.memory.into(), - convergent: false, - execution_scope: recipe.execution_scope.into(), - minimum_ptx: "8.6".into(), - minimum_sm: recipe.minimum_sm.map(Into::into), - ptx_result: if query { "u32".into() } else { "()".into() }, - targets: recipe.targets.into(), - ptx_isa_version: "9.3".into(), - ptx_isa_section: "9.7.14.18-19 Cluster Launch Control".into(), - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-clusterlaunchcontrol-try-cancel".into(), - lowering: "generated_clc".into(), - backend_lowerings: vec![ - OverlayBackendLowering { - backend: IntrinsicBackend::LlvmNvptx, - mechanism: BackendLoweringMechanism::TypedNvvm, - evidence_profile: admission.llvm_evidence_profile.clone(), - targets: None, - minimum_ptx: Some("8.6".into()), - minimum_sm: recipe.minimum_sm.map(Into::into), - }, - OverlayBackendLowering { - backend: IntrinsicBackend::LibNvvm, - mechanism: BackendLoweringMechanism::TypedNvvm, - evidence_profile: admission.libnvvm_evidence_profile.clone(), - targets: None, - minimum_ptx: Some("8.6".into()), - minimum_sm: recipe.minimum_sm.map(Into::into), - }, - ], - packed_atomic: None, - redux: None, - vote: None, - active_mask: None, - warp_match: None, - warp_barrier: None, - warp_shuffle: None, - dot_product: None, - packed_alu: None, - integer_minmax: None, - packed_conversion: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cp_async_copy: None, - cp_async_control: None, - cp_async_mbarrier: None, - mbarrier_basic: None, - movmatrix: None, - mbarrier_extended: None, - register_mma: None, - sparse_mma: None, - prmt: None, - cluster_barrier: None, - wgmma_control: None, - special_register: None, - debug_control: None, - cluster_memory: None, - clc: Some(Clc { - operation: recipe.operation, - adapter: recipe.adapter, - runtime_validation: admission.runtime_validation, - }), - tma: None, - tcgen05: None, - ldmatrix_variant: None, - ldmatrix_safety: None, - ldmatrix_adapter: None, - selected_address_space: None, - expected_ptx: InstructionPattern { - mnemonic: "clusterlaunchcontrol".into(), - modifiers: recipe.modifiers.iter().map(|value| (*value).into()).collect(), - operands: if query { - vec![ - OperandPattern::Register, - OperandPattern::Exact { value: "%clc_handle".into() }, - ] - } else { - recipe.operands.to_vec() - }, - }, - summary: recipe.summary.into(), - }) - }) - .collect() -} - -fn validate_clc_policy(policy: &OverlayIntrinsic, declaration: &ImportedIntrinsic) -> Result<()> { - let clc = policy - .clc - .as_ref() - .with_context(|| format!("{} has no closed CLC contract", policy.id))?; - let recipe = clc_recipe(clc.operation); - let query = recipe.pure; - ensure!( - policy.id == recipe.id - && policy.abi_id == recipe.abi_id - && policy.operation_key == recipe.operation_key - && policy.source.is_none() - && policy.source_record.as_deref() == Some(recipe.source_record) - && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) - && policy.resolved_llvm_symbol.is_none() - && declaration.source_record == recipe.source_record - && declaration.llvm_name == recipe.llvm_symbol, - "{} CLC identity changed", - policy.id - ); - ensure!( - policy.rust_module == "clc" - && policy.rust_name == recipe.id - && policy.rust_arguments == recipe.rust_arguments - && policy.rust_result == if query { "u32" } else { "()" } - && !policy.safe - && !policy.must_use - && policy.safe_allowlist_reason.is_none() - && policy.public_rust_path == format!("cuda_intrinsics::clc::{}", recipe.id) - && policy.compatibility_rust_paths == [format!("cuda_device::clc::{}", recipe.id)], - "{} CLC Rust API changed", - policy.id - ); - ensure!( - policy.dialect_op_type == recipe.dialect_op_type - && policy.dialect_op_name == recipe.dialect_op_name - && policy.dialect_operands == recipe.dialect_operands - && policy.dialect_results == recipe.dialect_results - && policy.llvm_arguments == recipe.llvm_arguments - && policy.llvm_results == recipe.llvm_results - && policy.lowering == "generated_clc", - "{} CLC carrier or LLVM adapter changed", - policy.id - ); - ensure!( - declaration.arguments == recipe.llvm_arguments - && declaration.results == recipe.llvm_results - && declaration.properties - == if query { - vec!["IntrNoMem", "IntrSpeculatable"] - } else { - vec!["IntrArgMemOnly", "IntrHasSideEffects"] - }, - "{} imported CLC declaration changed", - policy.id - ); - ensure!( - policy.pure == recipe.pure - && policy.memory == recipe.memory - && !policy.convergent - && policy.execution_scope == recipe.execution_scope - && clc.adapter == recipe.adapter - && clc.runtime_validation == RuntimeValidation::Unexecuted, - "{} CLC semantics changed", - policy.id - ); - ensure!( - policy.minimum_ptx == "8.6" - && policy.minimum_sm.as_deref() == recipe.minimum_sm - && policy.targets == recipe.targets - && policy.ptx_isa_version == "9.3" - && policy.ptx_result == if query { "u32" } else { "()" }, - "{} CLC target contract changed", - policy.id - ); - ensure!( - policy.expected_ptx.mnemonic == "clusterlaunchcontrol" - && policy.expected_ptx.modifiers == recipe.modifiers - && policy.expected_ptx.operands - == if query { - vec![ - OperandPattern::Register, - OperandPattern::Exact { - value: "%clc_handle".into(), - }, - ] - } else { - recipe.operands.to_vec() - } - && policy.backend_lowerings.len() == 2 - && policy.backend_lowerings.iter().all(|route| { - route.mechanism == BackendLoweringMechanism::TypedNvvm - && route.minimum_ptx.as_deref() == Some("8.6") - && route.minimum_sm.as_deref() == recipe.minimum_sm - && !route.evidence_profile.trim().is_empty() - }), - "{} CLC PTX shape or backend route changed", - policy.id - ); - ensure_no_other_family_contract(policy, "CLC")?; - Ok(()) -} - -const EXECUTION_CONTROL_OPERATIONS: [ExecutionControlOperation; 8] = [ - ExecutionControlOperation::BarrierCtaSync, - ExecutionControlOperation::BarrierCtaSyncAligned, - ExecutionControlOperation::BarrierCtaArrive, - ExecutionControlOperation::BarrierCtaArriveAligned, - ExecutionControlOperation::GridDependencyLaunchDependents, - ExecutionControlOperation::GridDependencyWait, - ExecutionControlOperation::SetMaxNRegInc, - ExecutionControlOperation::SetMaxNRegDec, -]; - -fn validate_execution_control_family_completeness(intrinsics: &[OverlayIntrinsic]) -> Result<()> { - let has_execution_control_family = intrinsics.iter().any(|policy| { - matches!( - policy.family.as_str(), - "counted_barrier" | "grid_dependency" | "register_control" - ) - }); - if !has_execution_control_family { - return Ok(()); - } - let actual = intrinsics - .iter() - .filter_map(|policy| ExecutionControlOperation::from_catalog_id(&policy.id)) - .collect::>(); - let expected = EXECUTION_CONTROL_OPERATIONS - .into_iter() - .collect::>(); - ensure!( - actual == expected, - "execution-control overlay must admit all four counted barriers, both grid-dependency controls, and both setmaxnreg operations" - ); - for family in ["counted_barrier", "grid_dependency", "register_control"] { - ensure!( - intrinsics - .iter() - .filter(|policy| policy.family == family) - .all(|policy| { - ExecutionControlOperation::from_catalog_id(&policy.id) - .is_some_and(|operation| operation.family() == family) - }), - "{family} overlay contains an operation outside its closed instruction family" - ); - } - Ok(()) -} - -#[derive(Clone)] -struct ExecutionControlRecipe { - abi_id: &'static str, - id: &'static str, - operation_key: &'static str, - source_record: &'static str, - llvm_symbol: &'static str, - rust_module: &'static str, - rust_arguments: &'static [&'static str], - compatibility_path: &'static str, - dialect_op_type: &'static str, - dialect_op_name: &'static str, - dialect_operands: &'static [&'static str], - llvm_arguments: &'static [&'static str], - imported_classes: &'static [&'static str], - imported_properties: &'static [&'static str], - memory: &'static str, - convergent: bool, - execution_scope: &'static str, - minimum_ptx: &'static str, - minimum_sm: Option<&'static str>, - targets: &'static str, - ptx_isa_section: &'static str, - ptx_isa_url: &'static str, - mnemonic: &'static str, - modifiers: &'static [&'static str], - operands: &'static [OperandPattern], - selection_records: &'static [&'static str], - selection_asm: &'static str, - selection_predicates: &'static [&'static str], - summary: &'static str, -} - -fn execution_control_recipe(operation: ExecutionControlOperation) -> ExecutionControlRecipe { - use ExecutionControlOperation::*; - - const EMPTY: &[&str] = &[]; - const I32: &[&str] = &["i32"]; - const I32_I32: &[&str] = &["i32", "i32"]; - const U32: &[&str] = &["u32"]; - const U32_U32: &[&str] = &["u32", "u32"]; - const BASE_CLASSES: &[&str] = &["SDPatternOperator", "Intrinsic"]; - const DEFAULT_CLASSES: &[&str] = &["SDPatternOperator", "Intrinsic", "DefaultAttrsIntrinsic"]; - const BARRIER_PROPERTIES: &[&str] = &["IntrConvergent", "IntrNoCallback"]; - const GRID_PROPERTIES: &[&str] = &["IntrHasSideEffects", "IntrNoMem"]; - const SETMAX_PROPERTIES: &[&str] = &[ - "ImmArg", - "IntrConvergent", - "IntrHasSideEffects", - "IntrNoMem", - "Range", - ]; - const TWO_REGISTER_OR_IMMEDIATE: &[OperandPattern] = &[ - OperandPattern::RegisterOrImmediate, - OperandPattern::RegisterOrImmediate, - ]; - const ONE_IMMEDIATE: &[OperandPattern] = &[OperandPattern::Immediate]; - const NO_OPERANDS: &[OperandPattern] = &[]; - const BARRIER_SYNC_SELECTIONS: &[&str] = &[ - "BARRIER_CTA_SYNC_ii", - "BARRIER_CTA_SYNC_ir", - "BARRIER_CTA_SYNC_ri", - "BARRIER_CTA_SYNC_rr", - ]; - const BAR_SYNC_SELECTIONS: &[&str] = &[ - "BARRIER_CTA_SYNC_ALIGNED_ii", - "BARRIER_CTA_SYNC_ALIGNED_ir", - "BARRIER_CTA_SYNC_ALIGNED_ri", - "BARRIER_CTA_SYNC_ALIGNED_rr", - ]; - const BARRIER_ARRIVE_SELECTIONS: &[&str] = &[ - "BARRIER_CTA_ARRIVE_ii", - "BARRIER_CTA_ARRIVE_ir", - "BARRIER_CTA_ARRIVE_ri", - "BARRIER_CTA_ARRIVE_rr", - ]; - const BAR_ARRIVE_SELECTIONS: &[&str] = &[ - "BARRIER_CTA_ARRIVE_ALIGNED_ii", - "BARRIER_CTA_ARRIVE_ALIGNED_ir", - "BARRIER_CTA_ARRIVE_ALIGNED_ri", - "BARRIER_CTA_ARRIVE_ALIGNED_rr", - ]; - const PTX_60: &[&str] = &["Subtarget->getPTXVersion() >= 60"]; - const GRID_PREDICATES: &[&str] = &[ - "Subtarget->getSmVersion() >= 90", - "Subtarget->getPTXVersion() >= 78", - ]; - const SETMAX_PREDICATES: &[&str] = &["Subtarget->hasSetMaxNRegSupport()"]; - const SYNC_MODIFIERS: &[&str] = &["sync"]; - const ARRIVE_MODIFIERS: &[&str] = &["arrive"]; - const LAUNCH_DEPENDENTS_MODIFIERS: &[&str] = &["launch_dependents"]; - const WAIT_MODIFIERS: &[&str] = &["wait"]; - const SETMAX_INC_MODIFIERS: &[&str] = &["inc", "sync", "aligned", "u32"]; - const SETMAX_DEC_MODIFIERS: &[&str] = &["dec", "sync", "aligned", "u32"]; - const GRID_LAUNCH_SELECTIONS: &[&str] = &["GRIDDEPCONTROL_LAUNCH_DEPENDENTS"]; - const GRID_WAIT_SELECTIONS: &[&str] = &["GRIDDEPCONTROL_WAIT"]; - const SETMAX_INC_SELECTIONS: &[&str] = &["anonymous_22212"]; - const SETMAX_DEC_SELECTIONS: &[&str] = &["anonymous_22214"]; - const BAR_SECTION: &str = - "9.7.14.1 Parallel Synchronization and Communication Instructions: bar, barrier"; - const BAR_URL: &str = "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-bar-barrier"; - const GRID_SECTION: &str = - "9.7.15.1 Parallel Synchronization and Communication Instructions: griddepcontrol"; - const GRID_URL: &str = "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-griddepcontrol"; - const SETMAX_SECTION: &str = "9.7.16.5 Miscellaneous Instructions: setmaxnreg"; - const SETMAX_URL: &str = "https://docs.nvidia.com/cuda/parallel-thread-execution/#miscellaneous-instructions-setmaxnreg"; - - let ( - abi_id, - id, - operation_key, - source_record, - llvm_symbol, - rust_module, - rust_arguments, - compatibility_path, - dialect_op_type, - dialect_op_name, - dialect_operands, - llvm_arguments, - imported_classes, - imported_properties, - memory, - convergent, - execution_scope, - minimum_ptx, - minimum_sm, - targets, - ptx_isa_section, - ptx_isa_url, - mnemonic, - modifiers, - operands, - selection_records, - selection_asm, - selection_predicates, - summary, - ) = match operation { - BarrierCtaSync => ( - "i0883", - "barrier_cta_sync", - "synchronization.cta.barrier.count", - "int_nvvm_barrier_cta_sync_count", - "llvm.nvvm.barrier.cta.sync.count", - "barrier", - U32_U32, - "cuda_device::barrier::barrier_cta_sync", - "BarrierCtaSyncCountOp", - "nvvm.barrier_cta_sync_count", - I32_I32, - I32_I32, - BASE_CLASSES, - BARRIER_PROPERTIES, - "read_write", - true, - "cta", - "6.0", - Some("sm_30"), - "all", - BAR_SECTION, - BAR_URL, - "barrier", - SYNC_MODIFIERS, - TWO_REGISTER_OR_IMMEDIATE, - BARRIER_SYNC_SELECTIONS, - "barrier.sync \t$i, $j;", - PTX_60, - "Synchronizes the requested number of CTA threads at a numbered barrier.", - ), - BarrierCtaSyncAligned => ( - "i0884", - "barrier_cta_sync_aligned", - "synchronization.cta.bar.aligned.count", - "int_nvvm_barrier_cta_sync_aligned_count", - "llvm.nvvm.barrier.cta.sync.aligned.count", - "barrier", - U32_U32, - "cuda_device::barrier::barrier_cta_sync_aligned", - "BarrierCtaSyncAlignedCountOp", - "nvvm.barrier_cta_sync_aligned_count", - I32_I32, - I32_I32, - BASE_CLASSES, - BARRIER_PROPERTIES, - "read_write", - true, - "cta", - "1.0", - None, - "all", - BAR_SECTION, - BAR_URL, - "bar", - SYNC_MODIFIERS, - TWO_REGISTER_OR_IMMEDIATE, - BAR_SYNC_SELECTIONS, - "bar.sync \t$i, $j;", - EMPTY, - "Synchronizes the requested number of CTA threads at an aligned numbered barrier.", - ), - BarrierCtaArrive => ( - "i0885", - "barrier_cta_arrive", - "synchronization.cta.barrier.arrive.count", - "int_nvvm_barrier_cta_arrive_count", - "llvm.nvvm.barrier.cta.arrive.count", - "barrier", - U32_U32, - "cuda_device::barrier::barrier_cta_arrive", - "BarrierCtaArriveCountOp", - "nvvm.barrier_cta_arrive_count", - I32_I32, - I32_I32, - BASE_CLASSES, - BARRIER_PROPERTIES, - "read_write", - true, - "cta", - "6.0", - Some("sm_30"), - "all", - BAR_SECTION, - BAR_URL, - "barrier", - ARRIVE_MODIFIERS, - TWO_REGISTER_OR_IMMEDIATE, - BARRIER_ARRIVE_SELECTIONS, - "barrier.arrive \t$i, $j;", - PTX_60, - "Signals arrival of the requested number of CTA threads at a numbered barrier without waiting.", - ), - BarrierCtaArriveAligned => ( - "i0886", - "barrier_cta_arrive_aligned", - "synchronization.cta.bar.arrive.aligned.count", - "int_nvvm_barrier_cta_arrive_aligned_count", - "llvm.nvvm.barrier.cta.arrive.aligned.count", - "barrier", - U32_U32, - "cuda_device::barrier::barrier_cta_arrive_aligned", - "BarrierCtaArriveAlignedCountOp", - "nvvm.barrier_cta_arrive_aligned_count", - I32_I32, - I32_I32, - BASE_CLASSES, - BARRIER_PROPERTIES, - "read_write", - true, - "cta", - "1.0", - None, - "all", - BAR_SECTION, - BAR_URL, - "bar", - ARRIVE_MODIFIERS, - TWO_REGISTER_OR_IMMEDIATE, - BAR_ARRIVE_SELECTIONS, - "bar.arrive \t$i, $j;", - EMPTY, - "Signals arrival of the requested number of CTA threads at an aligned numbered barrier without waiting.", - ), - GridDependencyLaunchDependents => ( - "i0913", - "grid_dependency_launch_dependents", - "synchronization.grid.dependency.launch_dependents", - "int_nvvm_griddepcontrol_launch_dependents", - "llvm.nvvm.griddepcontrol.launch.dependents", - "grid", - EMPTY, - "cuda_device::grid::dependency::trigger_dependents", - "GridDependencyLaunchDependentsOp", - "nvvm.grid_dependency_launch_dependents", - EMPTY, - EMPTY, - BASE_CLASSES, - GRID_PROPERTIES, - "none", - false, - "grid", - "7.8", - Some("sm_90"), - "all", - GRID_SECTION, - GRID_URL, - "griddepcontrol", - LAUNCH_DEPENDENTS_MODIFIERS, - NO_OPERANDS, - GRID_LAUNCH_SELECTIONS, - "griddepcontrol.launch_dependents;", - GRID_PREDICATES, - "Makes dependent grids eligible to launch before the current grid completes.", - ), - GridDependencyWait => ( - "i0914", - "grid_dependency_wait", - "synchronization.grid.dependency.wait", - "int_nvvm_griddepcontrol_wait", - "llvm.nvvm.griddepcontrol.wait", - "grid", - EMPTY, - "cuda_device::grid::dependency::wait", - "GridDependencyWaitOp", - "nvvm.grid_dependency_wait", - EMPTY, - EMPTY, - BASE_CLASSES, - GRID_PROPERTIES, - "none", - false, - "grid", - "7.8", - Some("sm_90"), - "all", - GRID_SECTION, - GRID_URL, - "griddepcontrol", - WAIT_MODIFIERS, - NO_OPERANDS, - GRID_WAIT_SELECTIONS, - "griddepcontrol.wait;", - GRID_PREDICATES, - "Waits until all prerequisite grids have completed.", - ), - SetMaxNRegInc => ( - "i0915", - "setmaxnreg_inc", - "register_control.warpgroup.setmaxnreg.increment", - "int_nvvm_setmaxnreg_inc_sync_aligned_u32", - "llvm.nvvm.setmaxnreg.inc.sync.aligned.u32", - "thread", - U32, - "cuda_device::thread::__setmaxnreg_inc", - "SetMaxNRegIncOp", - "nvvm.setmaxnreg_inc", - EMPTY, - I32, - DEFAULT_CLASSES, - SETMAX_PROPERTIES, - "none", - true, - "warpgroup", - "8.0", - None, - TENSOR_MAP_REPLACE_TARGETS, - SETMAX_SECTION, - SETMAX_URL, - "setmaxnreg", - SETMAX_INC_MODIFIERS, - ONE_IMMEDIATE, - SETMAX_INC_SELECTIONS, - "setmaxnreg.inc.sync.aligned.u32 \t$reg_count;", - SETMAX_PREDICATES, - "Increases the executing warpgroup's maximum register allocation to the immediate count.", - ), - SetMaxNRegDec => ( - "i0916", - "setmaxnreg_dec", - "register_control.warpgroup.setmaxnreg.decrement", - "int_nvvm_setmaxnreg_dec_sync_aligned_u32", - "llvm.nvvm.setmaxnreg.dec.sync.aligned.u32", - "thread", - U32, - "cuda_device::thread::__setmaxnreg_dec", - "SetMaxNRegDecOp", - "nvvm.setmaxnreg_dec", - EMPTY, - I32, - DEFAULT_CLASSES, - SETMAX_PROPERTIES, - "none", - true, - "warpgroup", - "8.0", - None, - TENSOR_MAP_REPLACE_TARGETS, - SETMAX_SECTION, - SETMAX_URL, - "setmaxnreg", - SETMAX_DEC_MODIFIERS, - ONE_IMMEDIATE, - SETMAX_DEC_SELECTIONS, - "setmaxnreg.dec.sync.aligned.u32 \t$reg_count;", - SETMAX_PREDICATES, - "Decreases the executing warpgroup's maximum register allocation to the immediate count.", - ), - }; - - ExecutionControlRecipe { - abi_id, - id, - operation_key, - source_record, - llvm_symbol, - rust_module, - rust_arguments, - compatibility_path, - dialect_op_type, - dialect_op_name, - dialect_operands, - llvm_arguments, - imported_classes, - imported_properties, - memory, - convergent, - execution_scope, - minimum_ptx, - minimum_sm, - targets, - ptx_isa_section, - ptx_isa_url, - mnemonic, - modifiers, - operands, - selection_records, - selection_asm, - selection_predicates, - summary, - } -} - -fn execution_control_backend_floor( - operation: ExecutionControlOperation, - backend: IntrinsicBackend, -) -> (&'static str, Option<&'static str>) { - use ExecutionControlOperation::*; - - match (operation, backend) { - (BarrierCtaSync | BarrierCtaArrive, IntrinsicBackend::LlvmNvptx) => ("6.0", Some("sm_30")), - (BarrierCtaSync | BarrierCtaArrive, IntrinsicBackend::LibNvvm) => ("6.0", Some("sm_75")), - (BarrierCtaSyncAligned | BarrierCtaArriveAligned, IntrinsicBackend::LlvmNvptx) => { - ("3.2", Some("sm_20")) - } - (BarrierCtaSyncAligned | BarrierCtaArriveAligned, IntrinsicBackend::LibNvvm) => { - ("1.0", Some("sm_75")) - } - _ => { - let recipe = execution_control_recipe(operation); - (recipe.minimum_ptx, recipe.minimum_sm) - } - } -} - -fn validate_execution_control_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, -) -> Result<()> { - let operation = ExecutionControlOperation::from_catalog_id(&policy.id) - .with_context(|| format!("{} is not a closed execution-control operation", policy.id))?; - let recipe = execution_control_recipe(operation); - ensure!( - policy.family == operation.family() - && policy.abi_id == recipe.abi_id - && policy.operation_key == recipe.operation_key - && policy.summary == recipe.summary - && policy.source.is_none() - && policy.source_record.as_deref() == Some(recipe.source_record) - && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) - && policy.resolved_llvm_symbol.is_none() - && declaration.source_record == recipe.source_record - && declaration.llvm_name == recipe.llvm_symbol, - "{} execution-control identity changed", - policy.id - ); - ensure!( - policy.rust_module == recipe.rust_module - && policy.rust_name == recipe.id - && policy.rust_arguments == recipe.rust_arguments - && policy.rust_result == "()" - && !policy.safe - && !policy.must_use - && policy.safe_allowlist_reason.is_none() - && policy.public_rust_path - == format!("cuda_intrinsics::{}::{}", recipe.rust_module, recipe.id) - && policy.compatibility_rust_paths == [recipe.compatibility_path], - "{} execution-control Rust API changed", - policy.id - ); - ensure!( - policy.dialect_op_type == recipe.dialect_op_type - && policy.dialect_op_name == recipe.dialect_op_name - && policy.dialect_operands == recipe.dialect_operands - && policy.dialect_results.is_empty() - && policy.llvm_arguments == recipe.llvm_arguments - && policy.llvm_results.is_empty() - && declaration.arguments == recipe.llvm_arguments - && declaration.results.is_empty() - && policy.lowering == "generated_execution_control", - "{} execution-control carrier or LLVM adapter changed", - policy.id - ); - ensure!( - declaration.classes == recipe.imported_classes - && declaration.properties == recipe.imported_properties, - "{} imported execution-control declaration changed", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == recipe.memory - && policy.convergent == recipe.convergent - && policy.execution_scope == recipe.execution_scope, - "{} execution-control semantics changed", - policy.id - ); - ensure!( - policy.minimum_ptx == recipe.minimum_ptx - && policy.minimum_sm.as_deref() == recipe.minimum_sm - && policy.targets == recipe.targets - && policy.ptx_result == "()" - && policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section == recipe.ptx_isa_section - && policy.ptx_isa_url == recipe.ptx_isa_url - && policy.expected_ptx.mnemonic == recipe.mnemonic - && policy.expected_ptx.modifiers == recipe.modifiers - && policy.expected_ptx.operands == recipe.operands, - "{} execution-control target or PTX contract changed", - policy.id - ); - let valid_route = |backend, mechanism| { - let (minimum_ptx, minimum_sm) = execution_control_backend_floor(operation, backend); - policy.backend_lowerings.iter().any(|route| { - route.backend == backend - && route.mechanism == mechanism - && route.minimum_ptx.as_deref() == Some(minimum_ptx) - && route.minimum_sm.as_deref() == minimum_sm - && !route.evidence_profile.trim().is_empty() - }) - }; - ensure!( - policy.backend_lowerings.len() == 2 - && valid_route( - IntrinsicBackend::LlvmNvptx, - BackendLoweringMechanism::TypedNvvm, - ) - && valid_route( - IntrinsicBackend::LibNvvm, - BackendLoweringMechanism::InlinePtx, - ), - "{} execution-control backend route changed", - policy.id - ); - let expected_selections = recipe - .selection_records - .iter() - .copied() - .collect::>(); - let actual_selections = declaration - .selections - .iter() - .map(|selection| selection.source_record.as_str()) - .collect::>(); - ensure!( - declaration.selections.len() == recipe.selection_records.len() - && actual_selections == expected_selections - && declaration.selections.iter().all(|selection| { - selection.asm == recipe.selection_asm - && selection.predicates == recipe.selection_predicates - && selection.constraints.is_empty() - }), - "{} imported execution-control selection family changed", - policy.id - ); - ensure_no_other_family_contract(policy, "execution-control")?; - Ok(()) -} - -const TMA_OPERATIONS: [TmaOperation; 47] = [ - TmaOperation::G2sTile1d, - TmaOperation::G2sTile2d, - TmaOperation::G2sTile2dMulticast, - TmaOperation::G2sTile2dMulticastCg2, - TmaOperation::G2sTile3d, - TmaOperation::G2sTile4d, - TmaOperation::G2sTile5d, - TmaOperation::S2gTile1d, - TmaOperation::S2gTile2d, - TmaOperation::S2gTile3d, - TmaOperation::S2gTile4d, - TmaOperation::S2gTile5d, - TmaOperation::CommitGroup, - TmaOperation::WaitGroup, - TmaOperation::WaitGroupRead, - TmaOperation::PrefetchTensorMap, - TmaOperation::PrefetchTile1d, - TmaOperation::PrefetchTile2d, - TmaOperation::PrefetchTile3d, - TmaOperation::PrefetchTile4d, - TmaOperation::PrefetchTile5d, - TmaOperation::PrefetchTileGather4TwoDimensional, - TmaOperation::ReplaceBoxDim, - TmaOperation::ReplaceElementStride, - TmaOperation::ReplaceElementType, - TmaOperation::ReplaceFillMode, - TmaOperation::ReplaceGlobalAddress, - TmaOperation::ReplaceGlobalDim, - TmaOperation::ReplaceGlobalStride, - TmaOperation::ReplaceInterleaveLayout, - TmaOperation::ReplaceRank, - TmaOperation::ReplaceSwizzleAtomicity, - TmaOperation::ReplaceSwizzleMode, - TmaOperation::FenceProxyTensorMapAcquireCluster, - TmaOperation::FenceProxyTensorMapAcquireCta, - TmaOperation::FenceProxyTensorMapAcquireGpu, - TmaOperation::FenceProxyTensorMapAcquireSystem, - TmaOperation::FenceProxyTensorMapReleaseCluster, - TmaOperation::FenceProxyTensorMapReleaseCta, - TmaOperation::FenceProxyTensorMapReleaseGpu, - TmaOperation::FenceProxyTensorMapReleaseSystem, - TmaOperation::PrefetchTile1dCacheHint, - TmaOperation::PrefetchTile2dCacheHint, - TmaOperation::PrefetchTile3dCacheHint, - TmaOperation::PrefetchTile4dCacheHint, - TmaOperation::PrefetchTile5dCacheHint, - TmaOperation::PrefetchTileGather4TwoDimensionalCacheHint, -]; - -const TMA_REDUCTION_OPERATIONS: [TmaReductionOperation; 8] = [ - TmaReductionOperation::Add, - TmaReductionOperation::And, - TmaReductionOperation::Dec, - TmaReductionOperation::Inc, - TmaReductionOperation::Max, - TmaReductionOperation::Min, - TmaReductionOperation::Or, - TmaReductionOperation::Xor, -]; - -fn tma_reduction_matrix() -> Vec { - let mut reductions = Vec::with_capacity(64); - for operation in TMA_REDUCTION_OPERATIONS { - for dimensions in 1..=5 { - reductions.push(TmaReduction { - operation, - load_mode: TmaReductionLoadMode::Tile, - dimensions, - }); - } - for dimensions in 3..=5 { - reductions.push(TmaReduction { - operation, - load_mode: TmaReductionLoadMode::Im2col, - dimensions, - }); - } - } - reductions -} - -#[cfg(test)] -fn tma_reduction_admission_variants() -> Vec { - tma_reduction_matrix() - .into_iter() - .enumerate() - .map(|(index, reduction)| TmaReductionAdmissionVariant { - // Preserve the repository's current ledger assignments in this test fixture. - abi_id: format!("i{:04}", 923 + index), - operation: reduction.operation, - load_mode: reduction.load_mode, - dimensions: reduction.dimensions, - }) - .collect() -} - -fn tma_reduction_operation_name(operation: TmaReductionOperation) -> (&'static str, &'static str) { - match operation { - TmaReductionOperation::Add => ("add", "Add"), - TmaReductionOperation::And => ("and", "And"), - TmaReductionOperation::Dec => ("dec", "Dec"), - TmaReductionOperation::Inc => ("inc", "Inc"), - TmaReductionOperation::Max => ("max", "Max"), - TmaReductionOperation::Min => ("min", "Min"), - TmaReductionOperation::Or => ("or", "Or"), - TmaReductionOperation::Xor => ("xor", "Xor"), - } -} - -fn tma_reduction_load_mode_name( - load_mode: TmaReductionLoadMode, -) -> (&'static str, &'static str, &'static str) { - match load_mode { - TmaReductionLoadMode::Tile => ("tile", "Tile", "tile"), - TmaReductionLoadMode::Im2col => ("im2col", "Im2col", "im2col_no_offs"), - } -} - -struct TmaReductionRecipe { - id: String, - operation_key: String, - source_record: String, - llvm_symbol: String, - rust_arguments: Vec<&'static str>, - dialect_op_type: String, - dialect_op_name: String, - dialect_operands: Vec<&'static str>, - llvm_arguments: Vec<&'static str>, - modifiers: Vec, - operands: Vec, - summary: String, -} - -fn tma_reduction_recipe(reduction: TmaReduction) -> Result { - ensure!( - (1..=5).contains(&reduction.dimensions), - "TMA reduction dimensionality must be in 1..=5" - ); - ensure!( - reduction.load_mode != TmaReductionLoadMode::Im2col || reduction.dimensions >= 3, - "TMA reduction im2col mode requires at least three dimensions" - ); - - let (operation, operation_camel) = tma_reduction_operation_name(reduction.operation); - let (load_mode, load_mode_camel, ptx_load_mode) = - tma_reduction_load_mode_name(reduction.load_mode); - let dimensions = reduction.dimensions as usize; - let id = format!( - "cp_async_bulk_tensor_reduce_{operation}_{load_mode}_{}d", - reduction.dimensions - ); - let mut rust_arguments = vec!["*const u8", "*const u8"]; - rust_arguments.extend(std::iter::repeat_n("i32", dimensions)); - let mut dialect_operands = vec!["ptr", "ptr"]; - dialect_operands.extend(std::iter::repeat_n("i32", dimensions)); - let mut llvm_arguments = vec!["shared_ptr", "ptr"]; - llvm_arguments.extend(std::iter::repeat_n("i32", dimensions)); - llvm_arguments.extend(["i64", "i1"]); - - Ok(TmaReductionRecipe { - operation_key: format!( - "memory.reduce.async.bulk.tensor.{operation}.{load_mode}.{}d", - reduction.dimensions - ), - source_record: format!( - "int_nvvm_cp_async_bulk_tensor_reduce_{operation}_{load_mode}_{}d", - reduction.dimensions - ), - llvm_symbol: format!( - "llvm.nvvm.cp.async.bulk.tensor.reduce.{operation}.{load_mode}.{}d", - reduction.dimensions - ), - dialect_op_type: format!( - "CpAsyncBulkTensorReduce{operation_camel}{load_mode_camel}{}dOp", - reduction.dimensions - ), - dialect_op_name: format!( - "nvvm.cp_async_bulk_tensor_reduce_{operation}_{load_mode}_{}d", - reduction.dimensions - ), - rust_arguments, - dialect_operands, - llvm_arguments, - modifiers: vec![ - "reduce".into(), - "async".into(), - "bulk".into(), - "tensor".into(), - format!("{}d", reduction.dimensions), - "global".into(), - "shared::cta".into(), - operation.into(), - ptx_load_mode.into(), - "bulk_group".into(), - ], - operands: vec![OperandPattern::Address, OperandPattern::Address], - summary: format!( - "Starts a TMA tensor {operation} reduction from shared to global memory in {load_mode} mode." - ), - id, - }) -} - -fn expand_tma_reduction_variant( - admission: &TmaAdmission, - variant: &TmaReductionAdmissionVariant, -) -> Result { - validate_abi_id(&variant.abi_id)?; - let reduction = TmaReduction { - operation: variant.operation, - load_mode: variant.load_mode, - dimensions: variant.dimensions, - }; - let recipe = tma_reduction_recipe(reduction)?; - - Ok(OverlayIntrinsic { - id: recipe.id.clone(), - abi_id: variant.abi_id.clone(), - operation_key: recipe.operation_key.clone(), - family: "tma".into(), - source: None, - source_record: Some(recipe.source_record.clone()), - rust_module: "tma".into(), - rust_name: recipe.id.clone(), - rust_arguments: recipe - .rust_arguments - .iter() - .map(|value| (*value).into()) - .collect(), - rust_result: "()".into(), - safe: false, - must_use: false, - safe_allowlist_reason: None, - public_rust_path: format!("cuda_intrinsics::tma::{}", recipe.id), - compatibility_rust_paths: vec![format!("cuda_device::tma::{}", recipe.id)], - dialect_op_type: recipe.dialect_op_type.clone(), - dialect_op_name: recipe.dialect_op_name.clone(), - dialect_operands: recipe - .dialect_operands - .iter() - .map(|value| (*value).into()) - .collect(), - dialect_results: vec![], - llvm_symbol: Some(recipe.llvm_symbol.clone()), - resolved_llvm_symbol: None, - llvm_arguments: recipe - .llvm_arguments - .iter() - .map(|value| (*value).into()) - .collect(), - llvm_results: vec![], - pure: false, - memory: "read_write".into(), - convergent: true, - execution_scope: "thread".into(), - minimum_ptx: "8.0".into(), - minimum_sm: Some("sm_90".into()), - ptx_result: "()".into(), - targets: "all".into(), - ptx_isa_version: "9.3".into(), - ptx_isa_section: - "9.7.9.26.5.3 Data Movement and Conversion Instructions: cp.reduce.async.bulk.tensor" - .into(), - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-reduce-async-bulk-tensor".into(), - lowering: "generated_tma".into(), - backend_lowerings: vec![ - OverlayBackendLowering { - backend: IntrinsicBackend::LlvmNvptx, - mechanism: BackendLoweringMechanism::TypedNvvm, - evidence_profile: admission - .reduce_llvm_evidence_profile - .as_ref() - .expect("validated TMA reduction LLVM evidence profile") - .clone(), - targets: None, - minimum_ptx: Some("8.0".into()), - minimum_sm: Some("sm_90".into()), - }, - OverlayBackendLowering { - backend: IntrinsicBackend::LibNvvm, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: admission - .reduce_libnvvm_evidence_profile - .as_ref() - .expect("validated TMA reduction libNVVM evidence profile") - .clone(), - targets: None, - minimum_ptx: Some("8.0".into()), - minimum_sm: Some("sm_90".into()), - }, - ], - packed_atomic: None, - redux: None, - vote: None, - active_mask: None, - warp_match: None, - warp_barrier: None, - warp_shuffle: None, - dot_product: None, - packed_alu: None, - integer_minmax: None, - packed_conversion: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cp_async_copy: None, - cp_async_control: None, - cp_async_mbarrier: None, - mbarrier_basic: None, - movmatrix: None, - mbarrier_extended: None, - register_mma: None, - sparse_mma: None, - prmt: None, - cluster_barrier: None, - wgmma_control: None, - special_register: None, - debug_control: None, - cluster_memory: None, - clc: None, - tma: Some(Tma { - operation: TmaOperation::Reduce, - reduction: Some(reduction), - adapter: TmaAdapter::ReductionPointersCoordinatesInjectDefaults, - runtime_validation: admission.runtime_validation, - }), - tcgen05: None, - ldmatrix_variant: None, - ldmatrix_safety: None, - ldmatrix_adapter: None, - selected_address_space: None, - expected_ptx: InstructionPattern { - mnemonic: "cp".into(), - modifiers: recipe.modifiers, - operands: recipe.operands, - }, - summary: recipe.summary, - }) -} - -struct TmaRecipe { - operation: TmaOperation, - abi_id: &'static str, - id: &'static str, - operation_key: &'static str, - source_record: &'static str, - llvm_symbol: &'static str, - resolved_llvm_symbol: Option, - selected_address_space: Option, - llvm_mechanism: BackendLoweringMechanism, - rust_arguments: Vec<&'static str>, - dialect_op_type: &'static str, - dialect_op_name: &'static str, - dialect_operands: Vec<&'static str>, - llvm_arguments: Vec<&'static str>, - adapter: TmaAdapter, - safe: bool, - safe_reason: Option<&'static str>, - convergent: bool, - minimum_ptx: &'static str, - minimum_sm: Option<&'static str>, - targets: &'static str, - memory: &'static str, - ptx_isa_section: &'static str, - ptx_isa_url: &'static str, - mnemonic: &'static str, - modifiers: Vec, - operands: Vec, - summary: &'static str, -} - -fn tma_recipe(operation: TmaOperation) -> TmaRecipe { - let (abi_id, id, operation_key, source_record, llvm_symbol, op_type, op_name) = match operation - { - TmaOperation::Reduce => unreachable!("TMA reductions use tma_reduction_recipe"), - TmaOperation::G2sTile1d => ( - "i0328", - "cp_async_bulk_tensor_1d_g2s", - "memory.copy.async.bulk.tensor.g2s.tile.1d", - "int_nvvm_cp_async_bulk_tensor_g2s_tile_1d", - "llvm.nvvm.cp.async.bulk.tensor.g2s.tile.1d", - "CpAsyncBulkTensorG2sTile1dOp", - "nvvm.cp_async_bulk_tensor_g2s_tile_1d", - ), - TmaOperation::G2sTile2d => ( - "i0329", - "cp_async_bulk_tensor_2d_g2s", - "memory.copy.async.bulk.tensor.g2s.tile.2d", - "int_nvvm_cp_async_bulk_tensor_g2s_tile_2d", - "llvm.nvvm.cp.async.bulk.tensor.g2s.tile.2d", - "CpAsyncBulkTensorG2sTile2dOp", - "nvvm.cp_async_bulk_tensor_g2s_tile_2d", - ), - TmaOperation::G2sTile2dMulticast => ( - "i0330", - "cp_async_bulk_tensor_2d_g2s_multicast", - "memory.copy.async.bulk.tensor.g2s.tile.2d.multicast", - "int_nvvm_cp_async_bulk_tensor_g2s_tile_2d", - "llvm.nvvm.cp.async.bulk.tensor.g2s.tile.2d", - "CpAsyncBulkTensorG2sTile2dMulticastOp", - "nvvm.cp_async_bulk_tensor_g2s_tile_2d_multicast", - ), - TmaOperation::G2sTile2dMulticastCg2 => ( - "i0331", - "cp_async_bulk_tensor_2d_g2s_multicast_cg2", - "memory.copy.async.bulk.tensor.g2s.tile.2d.multicast.cta_group_2", - "int_nvvm_cp_async_bulk_tensor_g2s_tile_2d", - "llvm.nvvm.cp.async.bulk.tensor.g2s.tile.2d", - "CpAsyncBulkTensorG2sTile2dMulticastCg2Op", - "nvvm.cp_async_bulk_tensor_g2s_tile_2d_multicast_cg2", - ), - TmaOperation::G2sTile3d => ( - "i0332", - "cp_async_bulk_tensor_3d_g2s", - "memory.copy.async.bulk.tensor.g2s.tile.3d", - "int_nvvm_cp_async_bulk_tensor_g2s_tile_3d", - "llvm.nvvm.cp.async.bulk.tensor.g2s.tile.3d", - "CpAsyncBulkTensorG2sTile3dOp", - "nvvm.cp_async_bulk_tensor_g2s_tile_3d", - ), - TmaOperation::G2sTile4d => ( - "i0333", - "cp_async_bulk_tensor_4d_g2s", - "memory.copy.async.bulk.tensor.g2s.tile.4d", - "int_nvvm_cp_async_bulk_tensor_g2s_tile_4d", - "llvm.nvvm.cp.async.bulk.tensor.g2s.tile.4d", - "CpAsyncBulkTensorG2sTile4dOp", - "nvvm.cp_async_bulk_tensor_g2s_tile_4d", - ), - TmaOperation::G2sTile5d => ( - "i0334", - "cp_async_bulk_tensor_5d_g2s", - "memory.copy.async.bulk.tensor.g2s.tile.5d", - "int_nvvm_cp_async_bulk_tensor_g2s_tile_5d", - "llvm.nvvm.cp.async.bulk.tensor.g2s.tile.5d", - "CpAsyncBulkTensorG2sTile5dOp", - "nvvm.cp_async_bulk_tensor_g2s_tile_5d", - ), - TmaOperation::S2gTile1d => ( - "i0335", - "cp_async_bulk_tensor_1d_s2g", - "memory.copy.async.bulk.tensor.s2g.tile.1d", - "int_nvvm_cp_async_bulk_tensor_s2g_tile_1d", - "llvm.nvvm.cp.async.bulk.tensor.s2g.tile.1d", - "CpAsyncBulkTensorS2gTile1dOp", - "nvvm.cp_async_bulk_tensor_s2g_tile_1d", - ), - TmaOperation::S2gTile2d => ( - "i0336", - "cp_async_bulk_tensor_2d_s2g", - "memory.copy.async.bulk.tensor.s2g.tile.2d", - "int_nvvm_cp_async_bulk_tensor_s2g_tile_2d", - "llvm.nvvm.cp.async.bulk.tensor.s2g.tile.2d", - "CpAsyncBulkTensorS2gTile2dOp", - "nvvm.cp_async_bulk_tensor_s2g_tile_2d", - ), - TmaOperation::S2gTile3d => ( - "i0337", - "cp_async_bulk_tensor_3d_s2g", - "memory.copy.async.bulk.tensor.s2g.tile.3d", - "int_nvvm_cp_async_bulk_tensor_s2g_tile_3d", - "llvm.nvvm.cp.async.bulk.tensor.s2g.tile.3d", - "CpAsyncBulkTensorS2gTile3dOp", - "nvvm.cp_async_bulk_tensor_s2g_tile_3d", - ), - TmaOperation::S2gTile4d => ( - "i0338", - "cp_async_bulk_tensor_4d_s2g", - "memory.copy.async.bulk.tensor.s2g.tile.4d", - "int_nvvm_cp_async_bulk_tensor_s2g_tile_4d", - "llvm.nvvm.cp.async.bulk.tensor.s2g.tile.4d", - "CpAsyncBulkTensorS2gTile4dOp", - "nvvm.cp_async_bulk_tensor_s2g_tile_4d", - ), - TmaOperation::S2gTile5d => ( - "i0339", - "cp_async_bulk_tensor_5d_s2g", - "memory.copy.async.bulk.tensor.s2g.tile.5d", - "int_nvvm_cp_async_bulk_tensor_s2g_tile_5d", - "llvm.nvvm.cp.async.bulk.tensor.s2g.tile.5d", - "CpAsyncBulkTensorS2gTile5dOp", - "nvvm.cp_async_bulk_tensor_s2g_tile_5d", - ), - TmaOperation::CommitGroup => ( - "i0340", - "cp_async_bulk_commit_group", - "memory.copy.async.bulk.group.commit", - "int_nvvm_cp_async_bulk_commit_group", - "llvm.nvvm.cp.async.bulk.commit.group", - "CpAsyncBulkCommitGroupOp", - "nvvm.cp_async_bulk_commit_group", - ), - TmaOperation::WaitGroup => ( - "i0341", - "cp_async_bulk_wait_group", - "memory.copy.async.bulk.group.wait_max_pending", - "int_nvvm_cp_async_bulk_wait_group", - "llvm.nvvm.cp.async.bulk.wait.group", - "CpAsyncBulkWaitGroupOp", - "nvvm.cp_async_bulk_wait_group", - ), - TmaOperation::WaitGroupRead => ( - "i0342", - "cp_async_bulk_wait_group_read", - "memory.copy.async.bulk.group.wait_read_max_pending", - "int_nvvm_cp_async_bulk_wait_group_read", - "llvm.nvvm.cp.async.bulk.wait.group.read", - "CpAsyncBulkWaitGroupReadOp", - "nvvm.cp_async_bulk_wait_group_read", - ), - TmaOperation::PrefetchTensorMap => ( - "i0887", - "prefetch_tma_descriptor", - "memory.prefetch.tensor_map", - "int_nvvm_prefetch_tensormap", - "llvm.nvvm.prefetch.tensormap", - "PrefetchTensorMapOp", - "nvvm.prefetch_tensormap", - ), - TmaOperation::PrefetchTile1d => ( - "i0888", - "cp_async_bulk_prefetch_tensor_1d_l2", - "memory.prefetch.async.bulk.tensor.global.tile.1d.l2", - "int_nvvm_cp_async_bulk_tensor_prefetch_tile_1d", - "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.1d", - "CpAsyncBulkPrefetchTensor1dL2Op", - "nvvm.cp_async_bulk_prefetch_tensor_1d_l2", - ), - TmaOperation::PrefetchTile2d => ( - "i0889", - "cp_async_bulk_prefetch_tensor_2d_l2", - "memory.prefetch.async.bulk.tensor.global.tile.2d.l2", - "int_nvvm_cp_async_bulk_tensor_prefetch_tile_2d", - "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.2d", - "CpAsyncBulkPrefetchTensor2dL2Op", - "nvvm.cp_async_bulk_prefetch_tensor_2d_l2", - ), - TmaOperation::PrefetchTile3d => ( - "i0890", - "cp_async_bulk_prefetch_tensor_3d_l2", - "memory.prefetch.async.bulk.tensor.global.tile.3d.l2", - "int_nvvm_cp_async_bulk_tensor_prefetch_tile_3d", - "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.3d", - "CpAsyncBulkPrefetchTensor3dL2Op", - "nvvm.cp_async_bulk_prefetch_tensor_3d_l2", - ), - TmaOperation::PrefetchTile4d => ( - "i0891", - "cp_async_bulk_prefetch_tensor_4d_l2", - "memory.prefetch.async.bulk.tensor.global.tile.4d.l2", - "int_nvvm_cp_async_bulk_tensor_prefetch_tile_4d", - "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.4d", - "CpAsyncBulkPrefetchTensor4dL2Op", - "nvvm.cp_async_bulk_prefetch_tensor_4d_l2", - ), - TmaOperation::PrefetchTile5d => ( - "i0892", - "cp_async_bulk_prefetch_tensor_5d_l2", - "memory.prefetch.async.bulk.tensor.global.tile.5d.l2", - "int_nvvm_cp_async_bulk_tensor_prefetch_tile_5d", - "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.5d", - "CpAsyncBulkPrefetchTensor5dL2Op", - "nvvm.cp_async_bulk_prefetch_tensor_5d_l2", - ), - TmaOperation::PrefetchTileGather4TwoDimensional => ( - "i0893", - "cp_async_bulk_prefetch_tensor_gather4_2d_l2", - "memory.prefetch.async.bulk.tensor.global.tile.gather4.2d.l2", - "int_nvvm_cp_async_bulk_tensor_prefetch_tile_gather4_2d", - "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.gather4.2d", - "CpAsyncBulkPrefetchTensorGather4TwoDimensionalL2Op", - "nvvm.cp_async_bulk_prefetch_tensor_gather4_2d_l2", - ), - TmaOperation::PrefetchTile1dCacheHint => ( - "i0917", - "cp_async_bulk_prefetch_tensor_1d_l2_cache_hint", - "memory.prefetch.async.bulk.tensor.global.tile.1d.l2.cache_hint", - "int_nvvm_cp_async_bulk_tensor_prefetch_tile_1d", - "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.1d", - "CpAsyncBulkPrefetchTensor1dL2CacheHintOp", - "nvvm.cp_async_bulk_prefetch_tensor_1d_l2_cache_hint", - ), - TmaOperation::PrefetchTile2dCacheHint => ( - "i0918", - "cp_async_bulk_prefetch_tensor_2d_l2_cache_hint", - "memory.prefetch.async.bulk.tensor.global.tile.2d.l2.cache_hint", - "int_nvvm_cp_async_bulk_tensor_prefetch_tile_2d", - "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.2d", - "CpAsyncBulkPrefetchTensor2dL2CacheHintOp", - "nvvm.cp_async_bulk_prefetch_tensor_2d_l2_cache_hint", - ), - TmaOperation::PrefetchTile3dCacheHint => ( - "i0919", - "cp_async_bulk_prefetch_tensor_3d_l2_cache_hint", - "memory.prefetch.async.bulk.tensor.global.tile.3d.l2.cache_hint", - "int_nvvm_cp_async_bulk_tensor_prefetch_tile_3d", - "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.3d", - "CpAsyncBulkPrefetchTensor3dL2CacheHintOp", - "nvvm.cp_async_bulk_prefetch_tensor_3d_l2_cache_hint", - ), - TmaOperation::PrefetchTile4dCacheHint => ( - "i0920", - "cp_async_bulk_prefetch_tensor_4d_l2_cache_hint", - "memory.prefetch.async.bulk.tensor.global.tile.4d.l2.cache_hint", - "int_nvvm_cp_async_bulk_tensor_prefetch_tile_4d", - "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.4d", - "CpAsyncBulkPrefetchTensor4dL2CacheHintOp", - "nvvm.cp_async_bulk_prefetch_tensor_4d_l2_cache_hint", - ), - TmaOperation::PrefetchTile5dCacheHint => ( - "i0921", - "cp_async_bulk_prefetch_tensor_5d_l2_cache_hint", - "memory.prefetch.async.bulk.tensor.global.tile.5d.l2.cache_hint", - "int_nvvm_cp_async_bulk_tensor_prefetch_tile_5d", - "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.5d", - "CpAsyncBulkPrefetchTensor5dL2CacheHintOp", - "nvvm.cp_async_bulk_prefetch_tensor_5d_l2_cache_hint", - ), - TmaOperation::PrefetchTileGather4TwoDimensionalCacheHint => ( - "i0922", - "cp_async_bulk_prefetch_tensor_gather4_2d_l2_cache_hint", - "memory.prefetch.async.bulk.tensor.global.tile.gather4.2d.l2.cache_hint", - "int_nvvm_cp_async_bulk_tensor_prefetch_tile_gather4_2d", - "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.gather4.2d", - "CpAsyncBulkPrefetchTensorGather4TwoDimensionalL2CacheHintOp", - "nvvm.cp_async_bulk_prefetch_tensor_gather4_2d_l2_cache_hint", - ), - TmaOperation::ReplaceBoxDim => ( - "i0894", - "tensormap_replace_box_dim", - "memory.tensor_map.replace.box_dim", - "int_nvvm_tensormap_replace_box_dim", - "llvm.nvvm.tensormap.replace.box.dim", - "ReplaceTensorMapBoxDimOp", - "nvvm.tensormap_replace_box_dim", - ), - TmaOperation::ReplaceElementStride => ( - "i0895", - "tensormap_replace_element_stride", - "memory.tensor_map.replace.element_stride", - "int_nvvm_tensormap_replace_element_stride", - "llvm.nvvm.tensormap.replace.element.stride", - "ReplaceTensorMapElementStrideOp", - "nvvm.tensormap_replace_element_stride", - ), - TmaOperation::ReplaceElementType => ( - "i0896", - "tensormap_replace_element_type", - "memory.tensor_map.replace.element_type", - "int_nvvm_tensormap_replace_elemtype", - "llvm.nvvm.tensormap.replace.elemtype", - "ReplaceTensorMapElementTypeOp", - "nvvm.tensormap_replace_element_type", - ), - TmaOperation::ReplaceFillMode => ( - "i0897", - "tensormap_replace_fill_mode", - "memory.tensor_map.replace.fill_mode", - "int_nvvm_tensormap_replace_fill_mode", - "llvm.nvvm.tensormap.replace.fill.mode", - "ReplaceTensorMapFillModeOp", - "nvvm.tensormap_replace_fill_mode", - ), - TmaOperation::ReplaceGlobalAddress => ( - "i0898", - "tensormap_replace_global_address", - "memory.tensor_map.replace.global_address", - "int_nvvm_tensormap_replace_global_address", - "llvm.nvvm.tensormap.replace.global.address", - "ReplaceTensorMapGlobalAddressOp", - "nvvm.tensormap_replace_global_address", - ), - TmaOperation::ReplaceGlobalDim => ( - "i0899", - "tensormap_replace_global_dim", - "memory.tensor_map.replace.global_dim", - "int_nvvm_tensormap_replace_global_dim", - "llvm.nvvm.tensormap.replace.global.dim", - "ReplaceTensorMapGlobalDimOp", - "nvvm.tensormap_replace_global_dim", - ), - TmaOperation::ReplaceGlobalStride => ( - "i0900", - "tensormap_replace_global_stride", - "memory.tensor_map.replace.global_stride", - "int_nvvm_tensormap_replace_global_stride", - "llvm.nvvm.tensormap.replace.global.stride", - "ReplaceTensorMapGlobalStrideOp", - "nvvm.tensormap_replace_global_stride", - ), - TmaOperation::ReplaceInterleaveLayout => ( - "i0901", - "tensormap_replace_interleave_layout", - "memory.tensor_map.replace.interleave_layout", - "int_nvvm_tensormap_replace_interleave_layout", - "llvm.nvvm.tensormap.replace.interleave.layout", - "ReplaceTensorMapInterleaveLayoutOp", - "nvvm.tensormap_replace_interleave_layout", - ), - TmaOperation::ReplaceRank => ( - "i0902", - "tensormap_replace_rank", - "memory.tensor_map.replace.rank", - "int_nvvm_tensormap_replace_rank", - "llvm.nvvm.tensormap.replace.rank", - "ReplaceTensorMapRankOp", - "nvvm.tensormap_replace_rank", - ), - TmaOperation::ReplaceSwizzleAtomicity => ( - "i0903", - "tensormap_replace_swizzle_atomicity", - "memory.tensor_map.replace.swizzle_atomicity", - "int_nvvm_tensormap_replace_swizzle_atomicity", - "llvm.nvvm.tensormap.replace.swizzle.atomicity", - "ReplaceTensorMapSwizzleAtomicityOp", - "nvvm.tensormap_replace_swizzle_atomicity", - ), - TmaOperation::ReplaceSwizzleMode => ( - "i0904", - "tensormap_replace_swizzle_mode", - "memory.tensor_map.replace.swizzle_mode", - "int_nvvm_tensormap_replace_swizzle_mode", - "llvm.nvvm.tensormap.replace.swizzle.mode", - "ReplaceTensorMapSwizzleModeOp", - "nvvm.tensormap_replace_swizzle_mode", - ), - TmaOperation::FenceProxyTensorMapAcquireCluster => ( - "i0905", - "fence_proxy_tensormap_generic_acquire_cluster", - "memory.fence.proxy.tensor_map.generic.acquire.cluster", - "int_nvvm_fence_proxy_tensormap_generic_acquire_cluster", - "llvm.nvvm.fence.proxy.tensormap_generic.acquire.cluster", - "FenceProxyTensorMapGenericAcquireClusterOp", - "nvvm.fence_proxy_tensormap_generic_acquire_cluster", - ), - TmaOperation::FenceProxyTensorMapAcquireCta => ( - "i0906", - "fence_proxy_tensormap_generic_acquire_cta", - "memory.fence.proxy.tensor_map.generic.acquire.cta", - "int_nvvm_fence_proxy_tensormap_generic_acquire_cta", - "llvm.nvvm.fence.proxy.tensormap_generic.acquire.cta", - "FenceProxyTensorMapGenericAcquireCtaOp", - "nvvm.fence_proxy_tensormap_generic_acquire_cta", - ), - TmaOperation::FenceProxyTensorMapAcquireGpu => ( - "i0907", - "fence_proxy_tensormap_generic_acquire_gpu", - "memory.fence.proxy.tensor_map.generic.acquire.gpu", - "int_nvvm_fence_proxy_tensormap_generic_acquire_gpu", - "llvm.nvvm.fence.proxy.tensormap_generic.acquire.gpu", - "FenceProxyTensorMapGenericAcquireGpuOp", - "nvvm.fence_proxy_tensormap_generic_acquire_gpu", - ), - TmaOperation::FenceProxyTensorMapAcquireSystem => ( - "i0908", - "fence_proxy_tensormap_generic_acquire_system", - "memory.fence.proxy.tensor_map.generic.acquire.system", - "int_nvvm_fence_proxy_tensormap_generic_acquire_sys", - "llvm.nvvm.fence.proxy.tensormap_generic.acquire.sys", - "FenceProxyTensorMapGenericAcquireSystemOp", - "nvvm.fence_proxy_tensormap_generic_acquire_system", - ), - TmaOperation::FenceProxyTensorMapReleaseCluster => ( - "i0909", - "fence_proxy_tensormap_generic_release_cluster", - "memory.fence.proxy.tensor_map.generic.release.cluster", - "int_nvvm_fence_proxy_tensormap_generic_release_cluster", - "llvm.nvvm.fence.proxy.tensormap_generic.release.cluster", - "FenceProxyTensorMapGenericReleaseClusterOp", - "nvvm.fence_proxy_tensormap_generic_release_cluster", - ), - TmaOperation::FenceProxyTensorMapReleaseCta => ( - "i0910", - "fence_proxy_tensormap_generic_release_cta", - "memory.fence.proxy.tensor_map.generic.release.cta", - "int_nvvm_fence_proxy_tensormap_generic_release_cta", - "llvm.nvvm.fence.proxy.tensormap_generic.release.cta", - "FenceProxyTensorMapGenericReleaseCtaOp", - "nvvm.fence_proxy_tensormap_generic_release_cta", - ), - TmaOperation::FenceProxyTensorMapReleaseGpu => ( - "i0911", - "fence_proxy_tensormap_generic_release_gpu", - "memory.fence.proxy.tensor_map.generic.release.gpu", - "int_nvvm_fence_proxy_tensormap_generic_release_gpu", - "llvm.nvvm.fence.proxy.tensormap_generic.release.gpu", - "FenceProxyTensorMapGenericReleaseGpuOp", - "nvvm.fence_proxy_tensormap_generic_release_gpu", - ), - TmaOperation::FenceProxyTensorMapReleaseSystem => ( - "i0912", - "fence_proxy_tensormap_generic_release_system", - "memory.fence.proxy.tensor_map.generic.release.system", - "int_nvvm_fence_proxy_tensormap_generic_release_sys", - "llvm.nvvm.fence.proxy.tensormap_generic.release.sys", - "FenceProxyTensorMapGenericReleaseSystemOp", - "nvvm.fence_proxy_tensormap_generic_release_system", - ), - }; - - let dimensions = operation.dimensions(); - let is_g2s = matches!( - operation, - TmaOperation::G2sTile1d - | TmaOperation::G2sTile2d - | TmaOperation::G2sTile2dMulticast - | TmaOperation::G2sTile2dMulticastCg2 - | TmaOperation::G2sTile3d - | TmaOperation::G2sTile4d - | TmaOperation::G2sTile5d - ); - let is_s2g = matches!( - operation, - TmaOperation::S2gTile1d - | TmaOperation::S2gTile2d - | TmaOperation::S2gTile3d - | TmaOperation::S2gTile4d - | TmaOperation::S2gTile5d - ); - let multicast = matches!( - operation, - TmaOperation::G2sTile2dMulticast | TmaOperation::G2sTile2dMulticastCg2 - ); - let cg2 = operation == TmaOperation::G2sTile2dMulticastCg2; - let prefetch_coordinates = operation.prefetch_coordinate_count(); - let is_release_fence = matches!( - operation, - TmaOperation::FenceProxyTensorMapReleaseCluster - | TmaOperation::FenceProxyTensorMapReleaseCta - | TmaOperation::FenceProxyTensorMapReleaseGpu - | TmaOperation::FenceProxyTensorMapReleaseSystem - ); - let is_acquire_fence = matches!( - operation, - TmaOperation::FenceProxyTensorMapAcquireCluster - | TmaOperation::FenceProxyTensorMapAcquireCta - | TmaOperation::FenceProxyTensorMapAcquireGpu - | TmaOperation::FenceProxyTensorMapAcquireSystem - ); - - let mut rust_arguments = Vec::new(); - let mut dialect_operands = Vec::new(); - let mut llvm_arguments = Vec::new(); - let (adapter, safe, safe_reason, convergent, summary) = if is_g2s { - rust_arguments.extend(["*mut u8", "*const u8"]); - rust_arguments.extend(std::iter::repeat_n("i32", dimensions.unwrap())); - rust_arguments.push("*mut u64"); - if multicast { - rust_arguments.push("u16"); - } - dialect_operands.extend(["ptr", "ptr", "ptr"]); - dialect_operands.extend(std::iter::repeat_n("i32", dimensions.unwrap())); - dialect_operands.extend(["i16", "i64"]); - llvm_arguments.extend(["shared_cluster_ptr", "shared_ptr", "ptr"]); - llvm_arguments.extend(std::iter::repeat_n("i32", dimensions.unwrap())); - llvm_arguments.extend(["i16", "i64", "i1", "i1", "i32"]); - ( - if multicast { - TmaAdapter::G2sPointersCoordinatesBarrierMaskInjectDefaults - } else { - TmaAdapter::G2sPointersCoordinatesBarrierInjectDefaults - }, - false, - None, - true, - if multicast { - "Starts a multicast TMA tile copy from global to cluster shared memory." - } else { - "Starts a TMA tile copy from global to cluster shared memory." - }, - ) - } else if is_s2g { - rust_arguments.extend(["*const u8", "*const u8"]); - rust_arguments.extend(std::iter::repeat_n("i32", dimensions.unwrap())); - dialect_operands.extend(["ptr", "ptr"]); - dialect_operands.extend(std::iter::repeat_n("i32", dimensions.unwrap())); - llvm_arguments.extend(["shared_ptr", "ptr"]); - llvm_arguments.extend(std::iter::repeat_n("i32", dimensions.unwrap())); - llvm_arguments.extend(["i64", "i1"]); - ( - TmaAdapter::S2gPointersCoordinatesInjectDefaults, - false, - None, - true, - "Starts a TMA tile copy from shared to global memory.", - ) - } else if operation == TmaOperation::CommitGroup || is_release_fence { - ( - TmaAdapter::NoOperands, - true, - Some(if operation == TmaOperation::CommitGroup { - "committing this thread's pending bulk-copy group has no Rust memory-safety precondition." - } else { - "publishing prior tensor-map writes has no Rust memory-safety precondition." - }), - false, - if operation == TmaOperation::CommitGroup { - "Commits this thread's pending asynchronous bulk copies as one group." - } else { - "Publishes prior generic-proxy tensor-map writes to the selected tensor-map proxy scope." - }, - ) - } else if operation == TmaOperation::PrefetchTensorMap { - rust_arguments.push("*const u8"); - dialect_operands.push("ptr"); - llvm_arguments.push("anyptr"); - ( - TmaAdapter::DescriptorPointer, - false, - None, - false, - "Prefetches a live tensor-map descriptor into the tensor-map cache.", - ) - } else if let Some(coordinate_count) = prefetch_coordinates { - rust_arguments.push("*const u8"); - rust_arguments.extend(std::iter::repeat_n("i32", coordinate_count)); - dialect_operands.push("ptr"); - dialect_operands.extend(std::iter::repeat_n("i32", coordinate_count)); - if operation.uses_prefetch_cache_hint() { - rust_arguments.push("u64"); - dialect_operands.push("i64"); - } - llvm_arguments.push("ptr"); - llvm_arguments.extend(std::iter::repeat_n("i32", coordinate_count)); - llvm_arguments.extend(["i64", "i1"]); - ( - if operation.uses_prefetch_cache_hint() { - TmaAdapter::DescriptorCoordinatesCacheHintInjectFlag - } else { - TmaAdapter::DescriptorCoordinatesInjectDefaults - }, - false, - None, - true, - if operation.uses_prefetch_cache_hint() { - "Prefetches one tensor tile through a tensor map into L2 using an explicit cache hint." - } else { - "Prefetches one tensor tile through a tensor map into L2." - }, - ) - } else if operation == TmaOperation::ReplaceGlobalAddress { - rust_arguments.extend(["*mut u8", "*const u8"]); - dialect_operands.extend(["ptr", "ptr"]); - llvm_arguments.extend(["anyptr", "i64"]); - ( - TmaAdapter::DescriptorAndAddressPointers, - false, - None, - false, - "Replaces the global base address stored in a writable tensor-map descriptor.", - ) - } else if matches!( - operation, - TmaOperation::ReplaceBoxDim - | TmaOperation::ReplaceElementStride - | TmaOperation::ReplaceGlobalDim - ) { - rust_arguments.extend(["*mut u8", "u32", "u32"]); - dialect_operands.extend(["ptr", "i32", "i32"]); - llvm_arguments.extend(["anyptr", "i32", "i32"]); - ( - TmaAdapter::DescriptorOrdinalAndU32, - false, - None, - false, - "Replaces one indexed 32-bit field in a writable tensor-map descriptor.", - ) - } else if operation == TmaOperation::ReplaceGlobalStride { - rust_arguments.extend(["*mut u8", "u32", "u64"]); - dialect_operands.extend(["ptr", "i32", "i64"]); - llvm_arguments.extend(["anyptr", "i32", "i64"]); - ( - TmaAdapter::DescriptorOrdinalAndU64, - false, - None, - false, - "Replaces one indexed 64-bit stride in a writable tensor-map descriptor.", - ) - } else if matches!( - operation, - TmaOperation::ReplaceElementType - | TmaOperation::ReplaceFillMode - | TmaOperation::ReplaceInterleaveLayout - | TmaOperation::ReplaceSwizzleAtomicity - | TmaOperation::ReplaceSwizzleMode - ) { - rust_arguments.extend(["*mut u8", "u32"]); - dialect_operands.extend(["ptr", "i32"]); - llvm_arguments.extend(["anyptr", "i32"]); - ( - TmaAdapter::DescriptorAndImmediateU32, - false, - None, - false, - "Replaces one immediate-valued field in a writable tensor-map descriptor.", - ) - } else if operation == TmaOperation::ReplaceRank { - rust_arguments.extend(["*mut u8", "u32"]); - dialect_operands.extend(["ptr", "i32"]); - llvm_arguments.extend(["anyptr", "i32"]); - ( - TmaAdapter::DescriptorAndRuntimeU32, - false, - None, - false, - "Replaces the rank stored in a writable tensor-map descriptor.", - ) - } else if is_acquire_fence { - rust_arguments.push("*const u8"); - dialect_operands.push("ptr"); - llvm_arguments.extend(["ptr", "i32"]); - ( - TmaAdapter::DescriptorPointerInjectBytes, - false, - None, - false, - "Acquires a published tensor-map descriptor into the selected tensor-map proxy scope.", - ) - } else { - rust_arguments.push("u32"); - dialect_operands.push("i32"); - llvm_arguments.push("i32"); - ( - TmaAdapter::CompileTimeConstantMaxPending, - true, - Some( - "waiting on this thread's bulk-copy groups has no Rust memory-safety precondition.", - ), - false, - if operation == TmaOperation::WaitGroupRead { - "Waits for bulk-copy groups and completes their reads." - } else { - "Waits until at most the requested bulk-copy groups remain pending." - }, - ) - }; - - let is_replace = matches!( - operation, - TmaOperation::ReplaceBoxDim - | TmaOperation::ReplaceElementStride - | TmaOperation::ReplaceElementType - | TmaOperation::ReplaceFillMode - | TmaOperation::ReplaceGlobalAddress - | TmaOperation::ReplaceGlobalDim - | TmaOperation::ReplaceGlobalStride - | TmaOperation::ReplaceInterleaveLayout - | TmaOperation::ReplaceRank - | TmaOperation::ReplaceSwizzleAtomicity - | TmaOperation::ReplaceSwizzleMode - ); - let is_fence = is_acquire_fence || is_release_fence; - let descriptor_control = operation == TmaOperation::PrefetchTensorMap - || prefetch_coordinates.is_some() - || is_replace - || is_fence; - let blackwell_tma = cg2 - || matches!( - operation, - TmaOperation::PrefetchTileGather4TwoDimensional - | TmaOperation::PrefetchTileGather4TwoDimensionalCacheHint - ); - - let (mnemonic, modifiers, operands) = if is_g2s { - let mut modifiers = vec![ - "async".into(), - "bulk".into(), - "tensor".into(), - format!("{}d", dimensions.unwrap()), - "shared::cluster".into(), - "global".into(), - "tile".into(), - "mbarrier::complete_tx::bytes".into(), - ]; - if multicast { - modifiers.push("multicast::cluster".into()); - } - if cg2 { - modifiers.push("cta_group::2".into()); - } - let mut operands = vec![ - OperandPattern::Address, - OperandPattern::Address, - OperandPattern::Address, - ]; - if multicast { - operands.push(OperandPattern::Register); - } - ("cp", modifiers, operands) - } else if is_s2g { - ( - "cp", - vec![ - "async".into(), - "bulk".into(), - "tensor".into(), - format!("{}d", dimensions.unwrap()), - "global".into(), - "shared::cta".into(), - "tile".into(), - "bulk_group".into(), - ], - vec![OperandPattern::Address, OperandPattern::Address], - ) - } else if operation == TmaOperation::CommitGroup { - ( - "cp", - vec!["async".into(), "bulk".into(), "commit_group".into()], - vec![], - ) - } else if matches!( - operation, - TmaOperation::WaitGroup | TmaOperation::WaitGroupRead - ) { - let mut modifiers = vec!["async".into(), "bulk".into(), "wait_group".into()]; - if operation == TmaOperation::WaitGroupRead { - modifiers.push("read".into()); - } - ("cp", modifiers, vec![OperandPattern::Immediate]) - } else if operation == TmaOperation::PrefetchTensorMap { - ( - "prefetch", - vec!["tensormap".into()], - vec![OperandPattern::Address], - ) - } else if prefetch_coordinates.is_some() { - let (dimensionality, tile) = if matches!( - operation, - TmaOperation::PrefetchTileGather4TwoDimensional - | TmaOperation::PrefetchTileGather4TwoDimensionalCacheHint - ) { - ("2d", "tile::gather4") - } else { - ( - match prefetch_coordinates.unwrap() { - 1 => "1d", - 2 => "2d", - 3 => "3d", - 4 => "4d", - 5 => "5d", - _ => unreachable!("closed TMA prefetch dimensionality"), - }, - "tile", - ) - }; - let mut modifiers = vec![ - "async".into(), - "bulk".into(), - "prefetch".into(), - "tensor".into(), - dimensionality.into(), - "L2".into(), - "global".into(), - tile.into(), - ]; - let mut operands = vec![OperandPattern::Address]; - if operation.uses_prefetch_cache_hint() { - modifiers.push("L2::cache_hint".into()); - operands.push(OperandPattern::Register); - } - ("cp", modifiers, operands) - } else if is_replace { - let (field, width, operands) = match operation { - TmaOperation::ReplaceGlobalAddress => ( - "global_address", - "b64", - vec![OperandPattern::Address, OperandPattern::Register], - ), - TmaOperation::ReplaceRank => ( - "rank", - "b32", - vec![OperandPattern::Address, OperandPattern::Register], - ), - TmaOperation::ReplaceBoxDim => ( - "box_dim", - "b32", - vec![ - OperandPattern::Address, - OperandPattern::Immediate, - OperandPattern::Register, - ], - ), - TmaOperation::ReplaceElementStride => ( - "element_stride", - "b32", - vec![ - OperandPattern::Address, - OperandPattern::Immediate, - OperandPattern::Register, - ], - ), - TmaOperation::ReplaceGlobalDim => ( - "global_dim", - "b32", - vec![ - OperandPattern::Address, - OperandPattern::Immediate, - OperandPattern::Register, - ], - ), - TmaOperation::ReplaceGlobalStride => ( - "global_stride", - "b64", - vec![ - OperandPattern::Address, - OperandPattern::Immediate, - OperandPattern::Register, - ], - ), - TmaOperation::ReplaceElementType => ( - "elemtype", - "b32", - vec![OperandPattern::Address, OperandPattern::Immediate], - ), - TmaOperation::ReplaceFillMode => ( - "fill_mode", - "b32", - vec![OperandPattern::Address, OperandPattern::Immediate], - ), - TmaOperation::ReplaceInterleaveLayout => ( - "interleave_layout", - "b32", - vec![OperandPattern::Address, OperandPattern::Immediate], - ), - TmaOperation::ReplaceSwizzleAtomicity => ( - "swizzle_atomicity", - "b32", - vec![OperandPattern::Address, OperandPattern::Immediate], - ), - TmaOperation::ReplaceSwizzleMode => ( - "swizzle_mode", - "b32", - vec![OperandPattern::Address, OperandPattern::Immediate], - ), - _ => unreachable!("TMA tensor-map replace operation was matched"), - }; - ( - "tensormap", - vec![ - "replace".into(), - "tile".into(), - field.into(), - "global".into(), - "b1024".into(), - width.into(), - ], - operands, - ) - } else if is_fence { - let (semantics, scope) = match operation { - TmaOperation::FenceProxyTensorMapAcquireCluster => ("acquire", "cluster"), - TmaOperation::FenceProxyTensorMapAcquireCta => ("acquire", "cta"), - TmaOperation::FenceProxyTensorMapAcquireGpu => ("acquire", "gpu"), - TmaOperation::FenceProxyTensorMapAcquireSystem => ("acquire", "sys"), - TmaOperation::FenceProxyTensorMapReleaseCluster => ("release", "cluster"), - TmaOperation::FenceProxyTensorMapReleaseCta => ("release", "cta"), - TmaOperation::FenceProxyTensorMapReleaseGpu => ("release", "gpu"), - TmaOperation::FenceProxyTensorMapReleaseSystem => ("release", "sys"), - _ => unreachable!("TMA tensor-map fence operation was matched"), - }; - ( - "fence", - vec![ - "proxy".into(), - "tensormap::generic".into(), - semantics.into(), - scope.into(), - ], - if is_acquire_fence { - vec![ - OperandPattern::Address, - OperandPattern::Exact { - value: "128".into(), - }, - ] - } else { - vec![] - }, - ) - } else { - unreachable!("TMA operation category was matched") - }; - - let minimum_ptx = if blackwell_tma || operation == TmaOperation::ReplaceSwizzleAtomicity { - "8.6" - } else if is_replace || is_fence { - "8.3" - } else { - "8.0" - }; - let (minimum_sm, targets) = if blackwell_tma { - (None, TMA_BLACKWELL_TARGETS) - } else if operation == TmaOperation::ReplaceSwizzleAtomicity { - (None, BLACKWELL_LDMATRIX_LLVM_TARGETS) - } else if is_replace { - (None, TENSOR_MAP_REPLACE_TARGETS) - } else { - (Some("sm_90"), "all") - }; - let memory = if operation == TmaOperation::PrefetchTensorMap || prefetch_coordinates.is_some() { - "read" - } else if is_replace { - "write" - } else { - "read_write" - }; - let polymorphic_descriptor = operation == TmaOperation::PrefetchTensorMap; - - TmaRecipe { - operation, - abi_id, - id, - operation_key, - source_record, - llvm_symbol, - resolved_llvm_symbol: polymorphic_descriptor.then(|| format!("{llvm_symbol}.p0")), - selected_address_space: polymorphic_descriptor.then_some(ImportedAddressSpace::Generic), - llvm_mechanism: if is_replace { - BackendLoweringMechanism::InlinePtx - } else { - BackendLoweringMechanism::TypedNvvm - }, - rust_arguments, - dialect_op_type: op_type, - dialect_op_name: op_name, - dialect_operands, - llvm_arguments, - adapter, - safe, - safe_reason, - convergent, - minimum_ptx, - minimum_sm, - targets, - memory, - ptx_isa_section: if descriptor_control { - "Tensor-map descriptor and asynchronous bulk tensor operations" - } else { - "9.7.9.26.5 Asynchronous bulk tensor copy" - }, - ptx_isa_url: if descriptor_control { - "https://docs.nvidia.com/cuda/parallel-thread-execution/#tensor-map" - } else { - "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async-bulk-tensor" - }, - mnemonic, - modifiers, - operands, - summary, - } -} - -fn expand_tma_admission(admission: &TmaAdmission) -> Result> { - ensure!( - admission.runtime_validation == RuntimeValidation::Unexecuted, - "TMA runtime validation may be marked executed only with GPU evidence" - ); - ensure!( - !admission.llvm_evidence_profile.trim().is_empty() - && !admission.libnvvm_evidence_profile.trim().is_empty(), - "compact TMA admission requires both backend evidence profiles" - ); - ensure!( - admission - .variants - .iter() - .map(|variant| variant.operation) - .eq(TMA_OPERATIONS), - "compact TMA admission must list all {} operations in canonical order", - TMA_OPERATIONS.len() - ); - - let mut records = admission - .variants - .iter() - .map(|variant| { - let recipe = tma_recipe(variant.operation); - ensure!( - variant.abi_id == recipe.abi_id, - "{} must keep reserved ABI ID {}", - recipe.id, - recipe.abi_id - ); - Ok(OverlayIntrinsic { - id: recipe.id.into(), - abi_id: variant.abi_id.clone(), - operation_key: recipe.operation_key.into(), - family: "tma".into(), - source: None, - source_record: Some(recipe.source_record.into()), - rust_module: "tma".into(), - rust_name: recipe.id.into(), - rust_arguments: recipe - .rust_arguments - .iter() - .map(|value| (*value).into()) - .collect(), - rust_result: "()".into(), - safe: recipe.safe, - must_use: false, - safe_allowlist_reason: recipe.safe_reason.map(Into::into), - public_rust_path: format!("cuda_intrinsics::tma::{}", recipe.id), - compatibility_rust_paths: vec![format!("cuda_device::tma::{}", recipe.id)], - dialect_op_type: recipe.dialect_op_type.into(), - dialect_op_name: recipe.dialect_op_name.into(), - dialect_operands: recipe - .dialect_operands - .iter() - .map(|value| (*value).into()) - .collect(), - dialect_results: vec![], - llvm_symbol: Some(recipe.llvm_symbol.into()), - resolved_llvm_symbol: recipe.resolved_llvm_symbol.clone(), - llvm_arguments: recipe - .llvm_arguments - .iter() - .map(|value| (*value).into()) - .collect(), - llvm_results: vec![], - pure: false, - memory: recipe.memory.into(), - convergent: recipe.convergent, - execution_scope: "thread".into(), - minimum_ptx: recipe.minimum_ptx.into(), - minimum_sm: recipe.minimum_sm.map(Into::into), - ptx_result: "()".into(), - targets: recipe.targets.into(), - ptx_isa_version: "9.3".into(), - ptx_isa_section: recipe.ptx_isa_section.into(), - ptx_isa_url: recipe.ptx_isa_url.into(), - lowering: "generated_tma".into(), - backend_lowerings: vec![ - OverlayBackendLowering { - backend: IntrinsicBackend::LlvmNvptx, - mechanism: recipe.llvm_mechanism, - evidence_profile: admission.llvm_evidence_profile.clone(), - targets: None, - minimum_ptx: Some(recipe.minimum_ptx.into()), - minimum_sm: recipe.minimum_sm.map(Into::into), - }, - OverlayBackendLowering { - backend: IntrinsicBackend::LibNvvm, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: admission.libnvvm_evidence_profile.clone(), - targets: None, - minimum_ptx: Some(recipe.minimum_ptx.into()), - minimum_sm: recipe.minimum_sm.map(Into::into), - }, - ], - packed_atomic: None, - redux: None, - vote: None, - active_mask: None, - warp_match: None, - warp_barrier: None, - warp_shuffle: None, - dot_product: None, - packed_alu: None, - integer_minmax: None, - packed_conversion: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cp_async_copy: None, - cp_async_control: None, - cp_async_mbarrier: None, - mbarrier_basic: None, - movmatrix: None, - mbarrier_extended: None, - register_mma: None, - sparse_mma: None, - prmt: None, - cluster_barrier: None, - wgmma_control: None, - special_register: None, - debug_control: None, - cluster_memory: None, - clc: None, - tma: Some(Tma { - operation: recipe.operation, - reduction: None, - adapter: recipe.adapter, - runtime_validation: admission.runtime_validation, - }), - tcgen05: None, - ldmatrix_variant: None, - ldmatrix_safety: None, - ldmatrix_adapter: None, - selected_address_space: recipe.selected_address_space, - expected_ptx: InstructionPattern { - mnemonic: recipe.mnemonic.into(), - modifiers: recipe.modifiers, - operands: recipe.operands, - }, - summary: recipe.summary.into(), - }) - }) - .collect::>>()?; - - if !admission.reduce_variants.is_empty() { - ensure!( - admission.reduce_llvm_evidence_profile.is_some(), - "compact TMA reduction admission requires reduce_llvm_evidence_profile" - ); - ensure!( - admission.reduce_libnvvm_evidence_profile.is_some(), - "compact TMA reduction admission requires reduce_libnvvm_evidence_profile" - ); - - let expected_reductions = tma_reduction_matrix(); - ensure!( - admission.reduce_variants.len() == expected_reductions.len() - && admission - .reduce_variants - .iter() - .zip(expected_reductions.iter()) - .all(|(variant, expected)| { - variant.operation == expected.operation - && variant.load_mode == expected.load_mode - && variant.dimensions == expected.dimensions - }), - "compact TMA reduction admission must list all 64 operations in canonical order" - ); - for variant in &admission.reduce_variants { - records.push(expand_tma_reduction_variant(admission, variant)?); - } - } - Ok(records) -} - -fn validate_tma_policy(policy: &OverlayIntrinsic, declaration: &ImportedIntrinsic) -> Result<()> { - let tma = policy - .tma - .as_ref() - .with_context(|| format!("{} has no closed TMA contract", policy.id))?; - if tma.operation == TmaOperation::Reduce { - return validate_tma_reduction_policy(policy, declaration, tma); - } - ensure!( - tma.reduction.is_none(), - "{} non-reduction TMA operation carries a reduction contract", - policy.id - ); - let recipe = tma_recipe(tma.operation); - ensure!( - policy.id == recipe.id - && policy.abi_id == recipe.abi_id - && policy.operation_key == recipe.operation_key - && policy.source.is_none() - && policy.source_record.as_deref() == Some(recipe.source_record) - && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) - && policy.resolved_llvm_symbol == recipe.resolved_llvm_symbol - && declaration.source_record == recipe.source_record - && declaration.llvm_name == recipe.llvm_symbol, - "{} TMA identity changed", - policy.id - ); - ensure!( - policy.rust_module == "tma" - && policy.rust_name == recipe.id - && policy.rust_arguments == recipe.rust_arguments - && policy.rust_result == "()" - && policy.safe == recipe.safe - && !policy.must_use - && policy.safe_allowlist_reason.as_deref() == recipe.safe_reason - && policy.public_rust_path == format!("cuda_intrinsics::tma::{}", recipe.id) - && policy.compatibility_rust_paths == [format!("cuda_device::tma::{}", recipe.id)], - "{} TMA Rust API changed", - policy.id - ); - ensure!( - policy.dialect_op_type == recipe.dialect_op_type - && policy.dialect_op_name == recipe.dialect_op_name - && policy.dialect_operands == recipe.dialect_operands - && policy.dialect_results.is_empty() - && policy.llvm_arguments == recipe.llvm_arguments - && policy.llvm_results.is_empty() - && declaration.arguments == recipe.llvm_arguments - && declaration.results.is_empty() - && policy.selected_address_space == recipe.selected_address_space - && policy.lowering == "generated_tma", - "{} TMA carrier or LLVM adapter changed", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == recipe.memory - && policy.convergent == recipe.convergent - && policy.execution_scope == "thread" - && tma.adapter == recipe.adapter - && tma.runtime_validation == RuntimeValidation::Unexecuted, - "{} TMA semantics changed", - policy.id - ); - ensure!( - policy.minimum_ptx == recipe.minimum_ptx - && policy.minimum_sm.as_deref() == recipe.minimum_sm - && policy.targets == recipe.targets - && policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section == recipe.ptx_isa_section - && policy.ptx_isa_url == recipe.ptx_isa_url - && policy.ptx_result == "()" - && policy.expected_ptx.mnemonic == recipe.mnemonic - && policy.expected_ptx.modifiers == recipe.modifiers - && policy.expected_ptx.operands == recipe.operands, - "{} TMA target or PTX contract changed", - policy.id - ); - let valid_route = |backend, mechanism| { - policy.backend_lowerings.iter().any(|route| { - route.backend == backend - && route.mechanism == mechanism - && route.minimum_ptx.as_deref() == Some(recipe.minimum_ptx) - && route.minimum_sm.as_deref() == recipe.minimum_sm - && !route.evidence_profile.trim().is_empty() - }) - }; - ensure!( - policy.backend_lowerings.len() == 2 - && valid_route(IntrinsicBackend::LlvmNvptx, recipe.llvm_mechanism,) - && valid_route( - IntrinsicBackend::LibNvvm, - BackendLoweringMechanism::InlinePtx, - ), - "{} TMA backend route changed", - policy.id - ); - let expected_properties = tma_imported_properties(tma.operation); - ensure!( - declaration.properties == expected_properties, - "{} imported TMA declaration changed", - policy.id - ); - ensure_no_other_family_contract(policy, "TMA")?; - Ok(()) -} - -fn validate_tma_reduction_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, - tma: &Tma, -) -> Result<()> { - ensure!( - tma.operation == TmaOperation::Reduce, - "{} is not a TMA reduction", - policy.id - ); - let reduction = tma - .reduction - .with_context(|| format!("{} has no TMA reduction contract", policy.id))?; - let recipe = tma_reduction_recipe(reduction)?; - ensure!( - policy.id == recipe.id - && policy.operation_key == recipe.operation_key - && policy.source.is_none() - && policy.source_record.as_deref() == Some(recipe.source_record.as_str()) - && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol.as_str()) - && policy.resolved_llvm_symbol.is_none() - && declaration.source_record == recipe.source_record - && declaration.llvm_name == recipe.llvm_symbol, - "{} TMA reduction identity changed", - policy.id - ); - ensure!( - policy.rust_module == "tma" - && policy.rust_name == recipe.id - && policy.rust_arguments == recipe.rust_arguments - && policy.rust_result == "()" - && !policy.safe - && !policy.must_use - && policy.safe_allowlist_reason.is_none() - && policy.public_rust_path == format!("cuda_intrinsics::tma::{}", recipe.id) - && policy.compatibility_rust_paths == [format!("cuda_device::tma::{}", recipe.id)], - "{} TMA reduction Rust API changed", - policy.id - ); - ensure!( - policy.dialect_op_type == recipe.dialect_op_type - && policy.dialect_op_name == recipe.dialect_op_name - && policy.dialect_operands == recipe.dialect_operands - && policy.dialect_results.is_empty() - && policy.llvm_arguments == recipe.llvm_arguments - && policy.llvm_results.is_empty() - && declaration.arguments == recipe.llvm_arguments - && declaration.results.is_empty() - && declaration.selections.is_empty() - && policy.lowering == "generated_tma", - "{} TMA reduction carrier or LLVM adapter changed", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == "read_write" - && policy.convergent - && policy.execution_scope == "thread" - && tma.adapter == TmaAdapter::ReductionPointersCoordinatesInjectDefaults - && tma.runtime_validation == RuntimeValidation::Unexecuted, - "{} TMA reduction semantics changed", - policy.id - ); - ensure!( - policy.minimum_ptx == "8.0" - && policy.minimum_sm.as_deref() == Some("sm_90") - && policy.targets == "all" - && policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section - == "9.7.9.26.5.3 Data Movement and Conversion Instructions: cp.reduce.async.bulk.tensor" - && policy.ptx_isa_url - == "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-reduce-async-bulk-tensor" - && policy.ptx_result == "()" - && policy.expected_ptx.mnemonic == "cp" - && policy.expected_ptx.modifiers == recipe.modifiers - && policy.expected_ptx.operands == recipe.operands, - "{} TMA reduction target or PTX contract changed", - policy.id - ); - let valid_route = |backend, mechanism| { - policy.backend_lowerings.iter().any(|route| { - route.backend == backend - && route.mechanism == mechanism - && route.minimum_ptx.as_deref() == Some("8.0") - && route.minimum_sm.as_deref() == Some("sm_90") - && !route.evidence_profile.trim().is_empty() - }) - }; - ensure!( - policy.backend_lowerings.len() == 2 - && valid_route( - IntrinsicBackend::LlvmNvptx, - BackendLoweringMechanism::TypedNvvm, - ) - && valid_route( - IntrinsicBackend::LibNvvm, - BackendLoweringMechanism::InlinePtx, - ), - "{} TMA reduction backend route changed", - policy.id - ); - ensure!( - declaration.properties == tma_reduction_imported_properties(reduction.dimensions as usize), - "{} imported TMA reduction declaration changed", - policy.id - ); - ensure_no_other_family_contract(policy, "TMA")?; - Ok(()) -} - -struct Tcgen05Recipe { - operation: Tcgen05Operation, - abi_id: &'static str, - id: &'static str, - operation_key: &'static str, - source_record: &'static str, - llvm_symbol: &'static str, - rust_arguments: &'static [&'static str], - rust_result: &'static str, - dialect_op_type: &'static str, - dialect_op_name: &'static str, - dialect_operands: &'static [&'static str], - dialect_results: &'static [&'static str], - llvm_arguments: &'static [&'static str], - llvm_results: &'static [&'static str], - imported_classes: &'static [&'static str], - imported_properties: &'static [&'static str], - adapter: Tcgen05Adapter, - source_contract: Tcgen05SourceContract, - safe: bool, - safe_reason: Option<&'static str>, - memory: &'static str, - modifiers: Vec, - operands: Vec, - selection_record: Option<&'static str>, - selection_asm: Option<&'static str>, - summary: &'static str, -} - -fn tcgen05_recipe(operation: Tcgen05Operation) -> Tcgen05Recipe { - const EMPTY: &[&str] = &[]; - const F32_X4: &[&str] = &["f32", "f32", "f32", "f32"]; - const F32_X32: &[&str] = &["f32"; 32]; - const BASE_CLASSES: &[&str] = &["SDPatternOperator", "Intrinsic"]; - const MMA_CLASSES: &[&str] = &[ - "SDPatternOperator", - "Intrinsic", - "DefaultAttrsIntrinsic", - "DefaultAttrsIntrinsicFlags", - ]; - const LOAD_CLASSES: &[&str] = &["SDPatternOperator", "Intrinsic", "NVVM_TCGEN05_LD"]; - const CONVERGENT_ARG_MEMORY: &[&str] = &["IntrArgMemOnly", "IntrConvergent", "NoCapture"]; - const CONVERGENT_INACCESSIBLE_ARG_MEMORY: &[&str] = &[ - "IntrConvergent", - "IntrInaccessibleMemOrArgMemOnly", - "NoCapture", - ]; - const CONVERGENT_INACCESSIBLE_MEMORY: &[&str] = &["IntrConvergent", "IntrInaccessibleMemOnly"]; - const ALLOC_PROPERTIES: &[&str] = &[ - "IntrConvergent", - "IntrInaccessibleMemOrArgMemOnly", - "NoCapture", - "WriteOnly", - ]; - const FENCE_PROPERTIES: &[&str] = &["IntrHasSideEffects", "IntrNoMem"]; - const MMA_WS_PROPERTIES: &[&str] = &[ - "ImmArg", - "ImmArg", - "ImmArg", - "IntrArgMemOnly", - "Range", - "Range", - "Range", - "ReadOnly", - "WriteOnly", - ]; - const MMA_CG1_PROPERTIES: &[&str] = &[ - "ImmArg", - "ImmArg", - "IntrArgMemOnly", - "Range", - "Range", - "WriteOnly", - ]; - const LOAD_PROPERTIES: &[&str] = &[ - "ImmArg", - "IntrArgMemOnly", - "IntrConvergent", - "NoCapture", - ]; - - let ( - abi_id, - id, - operation_key, - source_record, - llvm_symbol, - dialect_op_type, - dialect_op_name, - adapter, - source_contract, - ) = match operation { - Tcgen05Operation::Alloc => ( - "i0343", - "tcgen05_alloc", - "tcgen05.alloc.cg1", - "int_nvvm_tcgen05_alloc_shared_cg1", - "llvm.nvvm.tcgen05.alloc.shared.cg1", - "Tcgen05AllocOp", - "nvvm.tcgen05_alloc", - Tcgen05Adapter::SharedPointerColumnsToVoid, - Tcgen05SourceContract::ExactTablegenSelection, - ), - Tcgen05Operation::Dealloc => ( - "i0344", - "tcgen05_dealloc", - "tcgen05.dealloc.cg1", - "int_nvvm_tcgen05_dealloc_cg1", - "llvm.nvvm.tcgen05.dealloc.cg1", - "Tcgen05DeallocOp", - "nvvm.tcgen05_dealloc", - Tcgen05Adapter::TmemAddressColumnsToVoid, - Tcgen05SourceContract::ExactTablegenSelection, - ), - Tcgen05Operation::RelinquishAllocPermit => ( - "i0345", - "tcgen05_relinquish_alloc_permit", - "tcgen05.relinquish_alloc_permit.cg1", - "int_nvvm_tcgen05_relinq_alloc_permit_cg1", - "llvm.nvvm.tcgen05.relinq.alloc.permit.cg1", - "Tcgen05RelinquishAllocPermitOp", - "nvvm.tcgen05_relinquish_alloc_permit", - Tcgen05Adapter::NoOperands, - Tcgen05SourceContract::ExactTablegenSelection, - ), - Tcgen05Operation::FenceBeforeThreadSync => ( - "i0346", - "tcgen05_fence_before_thread_sync", - "tcgen05.fence.before_thread_sync", - "int_nvvm_tcgen05_fence_before_thread_sync", - "llvm.nvvm.tcgen05.fence.before.thread.sync", - "Tcgen05FenceBeforeThreadSyncOp", - "nvvm.tcgen05_fence_before_thread_sync", - Tcgen05Adapter::NoOperands, - Tcgen05SourceContract::ExactTablegenSelection, - ), - Tcgen05Operation::FenceAfterThreadSync => ( - "i0347", - "tcgen05_fence_after_thread_sync", - "tcgen05.fence.after.thread.sync", - "int_nvvm_tcgen05_fence_after_thread_sync", - "llvm.nvvm.tcgen05.fence.after.thread.sync", - "Tcgen05FenceAfterThreadSyncOp", - "nvvm.tcgen05_fence_after_thread_sync", - Tcgen05Adapter::NoOperands, - Tcgen05SourceContract::ExactTablegenSelection, - ), - Tcgen05Operation::Commit => ( - "i0348", - "tcgen05_commit", - "tcgen05.commit.cg1", - "int_nvvm_tcgen05_commit_cg1", - "llvm.nvvm.tcgen05.commit.cg1", - "Tcgen05CommitOp", - "nvvm.tcgen05_commit", - Tcgen05Adapter::BarrierPointerToVoid, - Tcgen05SourceContract::TablegenSelectionChangesPtx, - ), - Tcgen05Operation::CommitSharedCluster => ( - "i0349", - "tcgen05_commit_shared_cluster", - "tcgen05.commit.shared_cluster.cg1", - "int_nvvm_tcgen05_commit_shared_cg1", - "llvm.nvvm.tcgen05.commit.shared.cg1", - "Tcgen05CommitSharedClusterOp", - "nvvm.tcgen05_commit_shared_cluster", - Tcgen05Adapter::BarrierPointerToVoid, - Tcgen05SourceContract::ExactTablegenSelection, - ), - Tcgen05Operation::MmaWsF16 => ( - "i0350", - "tcgen05_mma_ws_f16", - "tcgen05.mma.ws.f16.cg1", - "int_nvvm_tcgen05_mma_ws_tensor", - "llvm.nvvm.tcgen05.mma.ws.tensor", - "Tcgen05MmaWsF16Op", - "nvvm.tcgen05_mma_ws_f16", - Tcgen05Adapter::MmaWsDropLegacyADescriptor, - Tcgen05SourceContract::TablegenSelectionChangesPtx, - ), - Tcgen05Operation::MmaF16 => ( - "i0351", - "tcgen05_mma_f16", - "tcgen05.mma.f16.cg1", - "int_nvvm_tcgen05_mma_shared_disable_output_lane_cg1", - "llvm.nvvm.tcgen05.mma.shared.disable_output_lane.cg1", - "Tcgen05MmaF16Op", - "nvvm.tcgen05_mma_f16", - Tcgen05Adapter::MmaInjectZeroDisableLanes, - Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection, - ), - Tcgen05Operation::MmaWsBf16 => ( - "i0352", - "tcgen05_mma_ws_bf16", - "tcgen05.mma.ws.bf16.cg1", - "int_nvvm_tcgen05_mma_ws_tensor", - "llvm.nvvm.tcgen05.mma.ws.tensor", - "Tcgen05MmaWsBf16Op", - "nvvm.tcgen05_mma_ws_bf16", - Tcgen05Adapter::MmaWsDropLegacyADescriptor, - Tcgen05SourceContract::TablegenSelectionChangesPtx, - ), - Tcgen05Operation::MmaWsTf32 => ( - "i0353", - "tcgen05_mma_ws_tf32", - "tcgen05.mma.ws.tf32.cg1", - "int_nvvm_tcgen05_mma_ws_tensor", - "llvm.nvvm.tcgen05.mma.ws.tensor", - "Tcgen05MmaWsTf32Op", - "nvvm.tcgen05_mma_ws_tf32", - Tcgen05Adapter::MmaWsDropLegacyADescriptor, - Tcgen05SourceContract::TablegenSelectionChangesPtx, - ), - Tcgen05Operation::CpSmemToTmem => ( - "i0354", - "tcgen05_cp_smem_to_tmem", - "tcgen05.cp.128x256b.cg1", - "int_nvvm_tcgen05_cp_128x256b_cg1", - "llvm.nvvm.tcgen05.cp.128x256b.cg1", - "Tcgen05CpSmemToTmemOp", - "nvvm.tcgen05_cp_smem_to_tmem", - Tcgen05Adapter::TmemDescriptorToVoid, - Tcgen05SourceContract::ExactTablegenSelection, - ), - Tcgen05Operation::Ld16x256bX8Pure => ( - "i0355", - "tcgen05_ld_16x256b_x8_pure", - "tcgen05.ld.16x256b.x8", - "int_nvvm_tcgen05_ld_16x256b_x8", - "llvm.nvvm.tcgen05.ld.16x256b.x8", - "Tcgen05Ld16x256bX8PureOp", - "nvvm.tcgen05_ld_16x256b_x8_pure", - Tcgen05Adapter::TmemToF32x32, - Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection, - ), - Tcgen05Operation::Ld16x256bPure => ( - "i0356", - "tcgen05_ld_16x256b_pure", - "tcgen05.ld.16x256b.x1", - "int_nvvm_tcgen05_ld_16x256b_x1", - "llvm.nvvm.tcgen05.ld.16x256b.x1", - "Tcgen05Ld16x256bPureOp", - "nvvm.tcgen05_ld_16x256b_pure", - Tcgen05Adapter::TmemToF32x4, - Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection, - ), - Tcgen05Operation::LoadWait => ( - "i0357", - "tcgen05_load_wait", - "tcgen05.wait.ld", - "int_nvvm_tcgen05_wait_ld", - "llvm.nvvm.tcgen05.wait.ld", - "Tcgen05LoadWaitOp", - "nvvm.tcgen05_load_wait", - Tcgen05Adapter::NoOperands, - Tcgen05SourceContract::ExactTablegenSelection, - ), - Tcgen05Operation::StoreWait => ( - "i0358", - "tcgen05_store_wait", - "tcgen05.wait.st", - "int_nvvm_tcgen05_wait_st", - "llvm.nvvm.tcgen05.wait.st", - "Tcgen05StoreWaitOp", - "nvvm.tcgen05_store_wait", - Tcgen05Adapter::NoOperands, - Tcgen05SourceContract::ExactTablegenSelection, - ), - Tcgen05Operation::AllocCg2 => ( - "i0359", - "tcgen05_alloc_cg2", - "tcgen05.alloc.cg2", - "int_nvvm_tcgen05_alloc_shared_cg2", - "llvm.nvvm.tcgen05.alloc.shared.cg2", - "Tcgen05AllocCg2Op", - "nvvm.tcgen05_alloc_cg2", - Tcgen05Adapter::SharedPointerColumnsToVoid, - Tcgen05SourceContract::ExactTablegenSelection, - ), - Tcgen05Operation::DeallocCg2 => ( - "i0360", - "tcgen05_dealloc_cg2", - "tcgen05.dealloc.cg2", - "int_nvvm_tcgen05_dealloc_cg2", - "llvm.nvvm.tcgen05.dealloc.cg2", - "Tcgen05DeallocCg2Op", - "nvvm.tcgen05_dealloc_cg2", - Tcgen05Adapter::TmemAddressColumnsToVoid, - Tcgen05SourceContract::ExactTablegenSelection, - ), - Tcgen05Operation::RelinquishAllocPermitCg2 => ( - "i0361", - "tcgen05_relinquish_alloc_permit_cg2", - "tcgen05.relinquish_alloc_permit.cg2", - "int_nvvm_tcgen05_relinq_alloc_permit_cg2", - "llvm.nvvm.tcgen05.relinq.alloc.permit.cg2", - "Tcgen05RelinquishAllocPermitCg2Op", - "nvvm.tcgen05_relinquish_alloc_permit_cg2", - Tcgen05Adapter::NoOperands, - Tcgen05SourceContract::ExactTablegenSelection, - ), - Tcgen05Operation::MmaF16Cg2 => ( - "i0362", - "tcgen05_mma_f16_cg2", - "tcgen05.mma.f16.cg2", - "int_nvvm_tcgen05_mma_shared_disable_output_lane_cg2", - "llvm.nvvm.tcgen05.mma.shared.disable_output_lane.cg2", - "Tcgen05MmaF16Cg2Op", - "nvvm.tcgen05_mma_f16_cg2", - Tcgen05Adapter::MmaInjectZeroDisableLanes, - Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection, - ), - Tcgen05Operation::CommitCg2 => ( - "i0363", - "tcgen05_commit_cg2", - "tcgen05.commit.cg2", - "int_nvvm_tcgen05_commit_cg2", - "llvm.nvvm.tcgen05.commit.cg2", - "Tcgen05CommitCg2Op", - "nvvm.tcgen05_commit_cg2", - Tcgen05Adapter::BarrierPointerToVoid, - Tcgen05SourceContract::TablegenSelectionChangesPtx, - ), - Tcgen05Operation::CommitSharedClusterCg2 => ( - "i0364", - "tcgen05_commit_shared_cluster_cg2", - "tcgen05.commit.shared_cluster.cg2", - "int_nvvm_tcgen05_commit_shared_cg2", - "llvm.nvvm.tcgen05.commit.shared.cg2", - "Tcgen05CommitSharedClusterCg2Op", - "nvvm.tcgen05_commit_shared_cluster_cg2", - Tcgen05Adapter::BarrierPointerToVoid, - Tcgen05SourceContract::ExactTablegenSelection, - ), - Tcgen05Operation::CommitMulticastCg2 => ( - "i0365", - "tcgen05_commit_multicast_cg2", - "tcgen05.commit.multicast.cg2", - "int_nvvm_tcgen05_commit_mc_shared_cg2", - "llvm.nvvm.tcgen05.commit.mc.shared.cg2", - "Tcgen05CommitMulticastCg2Op", - "nvvm.tcgen05_commit_multicast_cg2", - Tcgen05Adapter::BarrierPointerMaskToVoid, - Tcgen05SourceContract::ExactTablegenSelection, - ), - Tcgen05Operation::CpSmemToTmemCg2 => ( - "i0366", - "tcgen05_cp_smem_to_tmem_cg2", - "tcgen05.cp.128x256b.cg2", - "int_nvvm_tcgen05_cp_128x256b_cg2", - "llvm.nvvm.tcgen05.cp.128x256b.cg2", - "Tcgen05CpSmemToTmemCg2Op", - "nvvm.tcgen05_cp_smem_to_tmem_cg2", - Tcgen05Adapter::TmemDescriptorToVoid, - Tcgen05SourceContract::ExactTablegenSelection, - ), - Tcgen05Operation::CommitMulticast => ( - "i0760", - "tcgen05_commit_multicast", - "tcgen05.commit.multicast.cg1", - "int_nvvm_tcgen05_commit_mc_shared_cg1", - "llvm.nvvm.tcgen05.commit.mc.shared.cg1", - "Tcgen05CommitMulticastOp", - "nvvm.tcgen05_commit_multicast", - Tcgen05Adapter::BarrierPointerMaskToVoid, - Tcgen05SourceContract::ExactTablegenSelection, - ), - Tcgen05Operation::ShiftDown => ( - "i0761", - "tcgen05_shift_down", - "tcgen05.shift.down.cg1", - "int_nvvm_tcgen05_shift_down_cg1", - "llvm.nvvm.tcgen05.shift.down.cg1", - "Tcgen05ShiftDownOp", - "nvvm.tcgen05_shift_down", - Tcgen05Adapter::TmemAddressToVoid, - Tcgen05SourceContract::ExactTablegenSelection, - ), - Tcgen05Operation::ShiftDownCg2 => ( - "i0762", - "tcgen05_shift_down_cg2", - "tcgen05.shift.down.cg2", - "int_nvvm_tcgen05_shift_down_cg2", - "llvm.nvvm.tcgen05.shift.down.cg2", - "Tcgen05ShiftDownCg2Op", - "nvvm.tcgen05_shift_down_cg2", - Tcgen05Adapter::TmemAddressToVoid, - Tcgen05SourceContract::ExactTablegenSelection, - ), - Tcgen05Operation::Ld | Tcgen05Operation::St | Tcgen05Operation::Mma => { - unreachable!("tcgen05 load/store variants use their compact recipes") - } - }; - - let cg2 = matches!( - operation, - Tcgen05Operation::AllocCg2 - | Tcgen05Operation::DeallocCg2 - | Tcgen05Operation::RelinquishAllocPermitCg2 - | Tcgen05Operation::MmaF16Cg2 - | Tcgen05Operation::CommitCg2 - | Tcgen05Operation::CommitSharedClusterCg2 - | Tcgen05Operation::CommitMulticastCg2 - | Tcgen05Operation::CpSmemToTmemCg2 - | Tcgen05Operation::ShiftDownCg2 - ); - let group = if cg2 { "cta_group::2" } else { "cta_group::1" }; - let ( - rust_arguments, - rust_result, - dialect_operands, - dialect_results, - llvm_arguments, - llvm_results, - imported_classes, - imported_properties, - safe, - safe_reason, - memory, - modifiers, - operands, - selection_record, - selection_asm, - summary, - ) = match operation { - Tcgen05Operation::Alloc | Tcgen05Operation::AllocCg2 => ( - &["*mut u32", "u32"] as &[_], - "()", - &["ptr", "i32"] as &[_], - EMPTY, - &["shared_ptr", "i32"] as &[_], - EMPTY, - BASE_CLASSES, - ALLOC_PROPERTIES, - false, - None, - "write", - vec![ - "alloc".into(), - group.into(), - "sync".into(), - "aligned".into(), - "shared::cta".into(), - "b32".into(), - ], - vec![OperandPattern::Address, OperandPattern::Register], - Some(if cg2 { - "TCGEN05_ALLOC_S64_CG2" - } else { - "TCGEN05_ALLOC_S64_CG1" - }), - Some(if cg2 { - "tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 \t[$dst], $ncols;" - } else { - "tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 \t[$dst], $ncols;" - }), - "Allocates tensor-memory columns and writes their address to shared memory.", - ), - Tcgen05Operation::Dealloc | Tcgen05Operation::DeallocCg2 => ( - &["u32", "u32"] as &[_], - "()", - &["i32", "i32"] as &[_], - EMPTY, - &["tmem_ptr", "i32"] as &[_], - EMPTY, - BASE_CLASSES, - CONVERGENT_ARG_MEMORY, - false, - None, - "read_write", - vec![ - "dealloc".into(), - group.into(), - "sync".into(), - "aligned".into(), - "b32".into(), - ], - vec![OperandPattern::Register, OperandPattern::Register], - Some(if cg2 { - "TCGEN05_DEALLOC_CG2" - } else { - "TCGEN05_DEALLOC_CG1" - }), - Some(if cg2 { - "tcgen05.dealloc.cta_group::2.sync.aligned.b32 \t$tmem_addr, $ncols;" - } else { - "tcgen05.dealloc.cta_group::1.sync.aligned.b32 \t$tmem_addr, $ncols;" - }), - "Releases tensor-memory columns allocated to the CTA group.", - ), - Tcgen05Operation::RelinquishAllocPermit | Tcgen05Operation::RelinquishAllocPermitCg2 => ( - EMPTY, - "()", - EMPTY, - EMPTY, - EMPTY, - EMPTY, - BASE_CLASSES, - CONVERGENT_INACCESSIBLE_MEMORY, - true, - Some( - "The operation has no Rust memory operand and only releases this CTA group's allocation permit.", - ), - "read_write", - vec![ - "relinquish_alloc_permit".into(), - group.into(), - "sync".into(), - "aligned".into(), - ], - vec![], - Some(if cg2 { - "TCGEN05_RELINQ_CG2" - } else { - "TCGEN05_RELINQ_CG1" - }), - Some(if cg2 { - "tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;" - } else { - "tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;" - }), - "Releases the CTA group's tensor-memory allocation permit.", - ), - Tcgen05Operation::FenceBeforeThreadSync | Tcgen05Operation::FenceAfterThreadSync => { - let before = operation == Tcgen05Operation::FenceBeforeThreadSync; - ( - EMPTY, - "()", - EMPTY, - EMPTY, - EMPTY, - EMPTY, - BASE_CLASSES, - FENCE_PROPERTIES, - true, - Some("The operation only orders the calling thread's tcgen05 accesses."), - "none", - vec![if before { - "fence::before_thread_sync".into() - } else { - "fence::after_thread_sync".into() - }], - vec![], - Some(if before { - "tcgen05_fence_before_thread_sync" - } else { - "tcgen05_fence_after_thread_sync" - }), - Some(if before { - "tcgen05.fence::before_thread_sync;" - } else { - "tcgen05.fence::after_thread_sync;" - }), - if before { - "Orders prior tcgen05 accesses before thread synchronization." - } else { - "Orders later tcgen05 accesses after thread synchronization." - }, - ) - } - Tcgen05Operation::Commit - | Tcgen05Operation::CommitCg2 - | Tcgen05Operation::CommitSharedCluster - | Tcgen05Operation::CommitSharedClusterCg2 => { - let shared = matches!( - operation, - Tcgen05Operation::CommitSharedCluster | Tcgen05Operation::CommitSharedClusterCg2 - ); - let llvm_args = if shared { - &["shared_ptr"] as &[_] - } else { - &["ptr"] as &[_] - }; - let record = match (cg2, shared) { - (false, false) => "TCGEN05_COMMIT_CG1", - (true, false) => "TCGEN05_COMMIT_CG2", - (false, true) => "TCGEN05_COMMIT_S64_CG1", - (true, true) => "TCGEN05_COMMIT_S64_CG2", - }; - let asm = match cg2 { - false => { - "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 \t[$mbar];" - } - true => { - "tcgen05.commit.cta_group::2.mbarrier::arrive::one.shared::cluster.b64 \t[$mbar];" - } - }; - let mut modifiers = vec![ - "commit".into(), - group.into(), - "mbarrier::arrive::one".into(), - ]; - if shared { - modifiers.push("shared::cluster".into()); - } - modifiers.push("b64".into()); - ( - &["*mut u64"] as &[_], - "()", - &["ptr"] as &[_], - EMPTY, - llvm_args, - EMPTY, - BASE_CLASSES, - CONVERGENT_INACCESSIBLE_ARG_MEMORY, - false, - None, - "read_write", - modifiers, - vec![OperandPattern::Address], - Some(record), - Some(asm), - if shared { - "Commits tcgen05 completion to a cluster-shared mbarrier." - } else { - "Commits tcgen05 completion to an mbarrier." - }, - ) - } - Tcgen05Operation::MmaWsF16 | Tcgen05Operation::MmaWsBf16 | Tcgen05Operation::MmaWsTf32 => { - let kind = if operation == Tcgen05Operation::MmaWsTf32 { - "kind::tf32" - } else { - "kind::f16" - }; - ( - &["u32", "u32", "u64", "u64", "u32", "bool"] as &[_], - "()", - &["i32", "i32", "i64", "i64", "i32", "i1"] as &[_], - EMPTY, - &[ - "tmem_ptr", "tmem_ptr", "i64", "i32", "i1", "i32", "i32", "i32", - ] as &[_], - EMPTY, - MMA_CLASSES, - MMA_WS_PROPERTIES, - false, - None, - "read_write", - vec![ - "mma".into(), - "ws".into(), - "cta_group::1".into(), - kind.into(), - ], - vec![ - OperandPattern::Address, - OperandPattern::Address, - OperandPattern::Register, - OperandPattern::Register, - OperandPattern::Exact { - value: "%enable_pred".into(), - }, - ], - None, - None, - match operation { - Tcgen05Operation::MmaWsF16 => "Issues warp-specialized f16 tensor-memory MMA.", - Tcgen05Operation::MmaWsBf16 => { - "Issues warp-specialized bf16 tensor-memory MMA using the f16 instruction class." - } - _ => "Issues warp-specialized tf32 tensor-memory MMA.", - }, - ) - } - Tcgen05Operation::MmaF16 | Tcgen05Operation::MmaF16Cg2 => ( - &["u32", "u64", "u64", "u32", "bool"] as &[_], - "()", - &["i32", "i64", "i64", "i32", "i1"] as &[_], - EMPTY, - if cg2 { - &["tmem_ptr", "i64", "i64", "i32", "i1", "v8i32", "i32", "i32"] as &[_] - } else { - &["tmem_ptr", "i64", "i64", "i32", "i1", "v4i32", "i32", "i32"] as &[_] - }, - EMPTY, - MMA_CLASSES, - MMA_CG1_PROPERTIES, - false, - None, - "write", - vec!["mma".into(), group.into(), "kind::f16".into()], - vec![ - OperandPattern::Address, - OperandPattern::Register, - OperandPattern::Register, - OperandPattern::Register, - OperandPattern::Exact { - value: if cg2 { - "{%z, %z, %z, %z, %z, %z, %z, %z}" - } else { - "{%z, %z, %z, %z}" - } - .into(), - }, - OperandPattern::Exact { - value: "%enable_pred".into(), - }, - ], - None, - None, - "Issues f16 tensor-memory MMA with zeroed disable-output-lane controls.", - ), - Tcgen05Operation::CpSmemToTmem | Tcgen05Operation::CpSmemToTmemCg2 => ( - &["u32", "u64"] as &[_], - "()", - &["i32", "i64"] as &[_], - EMPTY, - &["tmem_ptr", "i64"] as &[_], - EMPTY, - BASE_CLASSES, - CONVERGENT_INACCESSIBLE_ARG_MEMORY, - false, - None, - "read_write", - vec!["cp".into(), group.into(), "128x256b".into()], - vec![OperandPattern::Address, OperandPattern::Register], - Some(if cg2 { - "TCGEN05_CP_128x256b_cg2" - } else { - "TCGEN05_CP_128x256b_cg1" - }), - Some(if cg2 { - "tcgen05.cp.cta_group::2.128x256b \t[$tmem_addr], $sdesc;" - } else { - "tcgen05.cp.cta_group::1.128x256b \t[$tmem_addr], $sdesc;" - }), - "Copies one 128x256-bit tile from shared memory to tensor memory.", - ), - Tcgen05Operation::Ld16x256bX8Pure | Tcgen05Operation::Ld16x256bPure => { - let x8 = operation == Tcgen05Operation::Ld16x256bX8Pure; - ( - &["u32"] as &[_], - if x8 { "[f32; 32]" } else { "[f32; 4]" }, - &["i32"] as &[_], - if x8 { F32_X32 } else { F32_X4 }, - &["tmem_ptr", "i1"] as &[_], - if x8 { - &["v32i32"] as &[_] - } else { - &["v4i32"] as &[_] - }, - LOAD_CLASSES, - LOAD_PROPERTIES, - false, - None, - "read", - vec![ - "ld".into(), - "sync".into(), - "aligned".into(), - "16x256b".into(), - if x8 { "x8".into() } else { "x1".into() }, - "b32".into(), - ], - vec![ - OperandPattern::RegisterList { - length: if x8 { 32 } else { 4 }, - }, - OperandPattern::Address, - ], - None, - None, - if x8 { - "Loads 32 f32 register values from tensor memory." - } else { - "Loads four f32 register values from tensor memory." - }, - ) - } - Tcgen05Operation::LoadWait | Tcgen05Operation::StoreWait => { - let load = operation == Tcgen05Operation::LoadWait; - ( - EMPTY, - "()", - EMPTY, - EMPTY, - EMPTY, - EMPTY, - BASE_CLASSES, - CONVERGENT_INACCESSIBLE_MEMORY, - true, - Some("The operation only waits for the calling thread's prior tcgen05 access."), - "read_write", - vec![ - if load { - "wait::ld".into() - } else { - "wait::st".into() - }, - "sync".into(), - "aligned".into(), - ], - vec![], - Some(if load { - "tcgen05_wait_ld" - } else { - "tcgen05_wait_st" - }), - Some(if load { - "tcgen05.wait::ld.sync.aligned;" - } else { - "tcgen05.wait::st.sync.aligned;" - }), - if load { - "Waits for prior tensor-memory loads to complete." - } else { - "Waits for prior tensor-memory stores to complete." - }, - ) - } - Tcgen05Operation::CommitMulticast | Tcgen05Operation::CommitMulticastCg2 => ( - &["*mut u64", "u16"] as &[_], - "()", - &["ptr", "i16"] as &[_], - EMPTY, - &["shared_ptr", "i16"] as &[_], - EMPTY, - BASE_CLASSES, - CONVERGENT_INACCESSIBLE_ARG_MEMORY, - false, - None, - "read_write", - vec![ - "commit".into(), - group.into(), - "mbarrier::arrive::one".into(), - "shared::cluster".into(), - "multicast::cluster".into(), - "b64".into(), - ], - vec![OperandPattern::Address, OperandPattern::Register], - Some(if cg2 { - "TCGEN05_COMMIT_S64_CG2_MC" - } else { - "TCGEN05_COMMIT_S64_CG1_MC" - }), - Some(if cg2 { - "tcgen05.commit.cta_group::2.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 \t[$mbar], $mc;" - } else { - "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 \t[$mbar], $mc;" - }), - "Commits tcgen05 completion to the selected cluster mbarriers.", - ), - Tcgen05Operation::ShiftDown | Tcgen05Operation::ShiftDownCg2 => ( - &["u32"] as &[_], - "()", - &["i32"] as &[_], - EMPTY, - &["tmem_ptr"] as &[_], - EMPTY, - BASE_CLASSES, - CONVERGENT_ARG_MEMORY, - false, - None, - "read_write", - vec!["shift".into(), group.into(), "down".into()], - vec![OperandPattern::Address], - Some(if cg2 { - "TCGEN05_SHIFT_CG2" - } else { - "TCGEN05_SHIFT_CG1" - }), - Some(if cg2 { - "tcgen05.shift.cta_group::2.down \t[$tmem_addr];" - } else { - "tcgen05.shift.cta_group::1.down \t[$tmem_addr];" - }), - "Shifts tensor-memory rows down by one row.", - ), - Tcgen05Operation::Ld | Tcgen05Operation::St | Tcgen05Operation::Mma => { - unreachable!("tcgen05 load/store variants use their compact recipes") - } - }; - - Tcgen05Recipe { - operation, - abi_id, - id, - operation_key, - source_record, - llvm_symbol, - rust_arguments, - rust_result, - dialect_op_type, - dialect_op_name, - dialect_operands, - dialect_results, - llvm_arguments, - llvm_results, - imported_classes, - imported_properties, - adapter, - source_contract, - safe, - safe_reason, - memory, - modifiers, - operands, - selection_record, - selection_asm, - summary, - } -} - -fn tcgen05_mma_form_name(form: Tcgen05MmaForm) -> &'static str { - match form { - Tcgen05MmaForm::Shared => "shared", - Tcgen05MmaForm::Tensor => "tensor", - Tcgen05MmaForm::TensorAshift => "tensor_ashift", - Tcgen05MmaForm::SpShared => "sp_shared", - Tcgen05MmaForm::SpTensor => "sp_tensor", - Tcgen05MmaForm::SpTensorAshift => "sp_tensor_ashift", - Tcgen05MmaForm::WsShared => "ws_shared", - Tcgen05MmaForm::WsSharedZeroColMask => "ws_shared_zero_col_mask", - Tcgen05MmaForm::WsSpShared => "ws_sp_shared", - Tcgen05MmaForm::WsSpSharedZeroColMask => "ws_sp_shared_zero_col_mask", - Tcgen05MmaForm::WsSpTensor => "ws_sp_tensor", - Tcgen05MmaForm::WsSpTensorZeroColMask => "ws_sp_tensor_zero_col_mask", - Tcgen05MmaForm::WsTensor => "ws_tensor", - Tcgen05MmaForm::WsTensorZeroColMask => "ws_tensor_zero_col_mask", - } -} - -fn tcgen05_mma_alias_name(alias: Tcgen05MmaAlias) -> &'static str { - match alias { - Tcgen05MmaAlias::E4m3 => "e4m3", - Tcgen05MmaAlias::E5m2 => "e5m2", - Tcgen05MmaAlias::E2m3 => "e2m3", - Tcgen05MmaAlias::E3m2 => "e3m2", - Tcgen05MmaAlias::E2m1 => "e2m1", - } -} - -fn tcgen05_mma_kind_name(kind: Tcgen05MmaKind) -> &'static str { - match kind { - Tcgen05MmaKind::F16 => "f16", - Tcgen05MmaKind::Tf32 => "tf32", - Tcgen05MmaKind::F8f6f4 => "f8f6f4", - Tcgen05MmaKind::I8 => "i8", - } -} - -fn tcgen05_mma_b_usage_name(usage: Tcgen05MmaBUsage) -> &'static str { - match usage { - Tcgen05MmaBUsage::Discard => "discard", - Tcgen05MmaBUsage::Fill => "fill", - Tcgen05MmaBUsage::Use => "use", - Tcgen05MmaBUsage::LastUse => "lastuse", - } -} - -fn tcgen05_mma_is_ws(form: Tcgen05MmaForm) -> bool { - matches!( - form, - Tcgen05MmaForm::WsShared - | Tcgen05MmaForm::WsSharedZeroColMask - | Tcgen05MmaForm::WsSpShared - | Tcgen05MmaForm::WsSpSharedZeroColMask - | Tcgen05MmaForm::WsSpTensor - | Tcgen05MmaForm::WsSpTensorZeroColMask - | Tcgen05MmaForm::WsTensor - | Tcgen05MmaForm::WsTensorZeroColMask - ) -} - -fn tcgen05_mma_is_sparse(form: Tcgen05MmaForm) -> bool { - matches!( - form, - Tcgen05MmaForm::SpShared - | Tcgen05MmaForm::SpTensor - | Tcgen05MmaForm::SpTensorAshift - | Tcgen05MmaForm::WsSpShared - | Tcgen05MmaForm::WsSpSharedZeroColMask - | Tcgen05MmaForm::WsSpTensor - | Tcgen05MmaForm::WsSpTensorZeroColMask - ) -} - -fn tcgen05_mma_is_tensor_a(form: Tcgen05MmaForm) -> bool { - matches!( - form, - Tcgen05MmaForm::Tensor - | Tcgen05MmaForm::TensorAshift - | Tcgen05MmaForm::SpTensor - | Tcgen05MmaForm::SpTensorAshift - | Tcgen05MmaForm::WsSpTensor - | Tcgen05MmaForm::WsSpTensorZeroColMask - | Tcgen05MmaForm::WsTensor - | Tcgen05MmaForm::WsTensorZeroColMask - ) -} - -fn tcgen05_mma_is_ashift(form: Tcgen05MmaForm) -> bool { - matches!( - form, - Tcgen05MmaForm::TensorAshift | Tcgen05MmaForm::SpTensorAshift - ) -} - -fn tcgen05_mma_has_zero_col_mask(form: Tcgen05MmaForm) -> bool { - matches!( - form, - Tcgen05MmaForm::WsSharedZeroColMask - | Tcgen05MmaForm::WsSpSharedZeroColMask - | Tcgen05MmaForm::WsSpTensorZeroColMask - | Tcgen05MmaForm::WsTensorZeroColMask - ) -} - -fn tcgen05_mma_source_record(form: Tcgen05MmaForm) -> String { - format!("int_nvvm_tcgen05_mma_{}", tcgen05_mma_form_name(form)) -} - -fn tcgen05_mma_llvm_symbol(form: Tcgen05MmaForm) -> String { - let suffix = match form { - Tcgen05MmaForm::Shared => "shared", - Tcgen05MmaForm::Tensor => "tensor", - Tcgen05MmaForm::TensorAshift => "tensor.ashift", - Tcgen05MmaForm::SpShared => "sp.shared", - Tcgen05MmaForm::SpTensor => "sp.tensor", - Tcgen05MmaForm::SpTensorAshift => "sp.tensor.ashift", - Tcgen05MmaForm::WsShared => "ws.shared", - Tcgen05MmaForm::WsSharedZeroColMask => "ws.shared.zero_col_mask", - Tcgen05MmaForm::WsSpShared => "ws.sp.shared", - Tcgen05MmaForm::WsSpSharedZeroColMask => "ws.sp.shared.zero_col_mask", - Tcgen05MmaForm::WsSpTensor => "ws.sp.tensor", - Tcgen05MmaForm::WsSpTensorZeroColMask => "ws.sp.tensor.zero_col_mask", - Tcgen05MmaForm::WsTensor => "ws.tensor", - Tcgen05MmaForm::WsTensorZeroColMask => "ws.tensor.zero_col_mask", - }; - format!("llvm.nvvm.tcgen05.mma.{suffix}") -} - -fn tcgen05_mma_llvm_arguments(form: Tcgen05MmaForm) -> Vec { - let mut arguments = vec![ - "tmem_ptr".into(), - if tcgen05_mma_is_tensor_a(form) { - "tmem_ptr".into() - } else { - "i64".into() - }, - "i64".into(), - "i32".into(), - "i1".into(), - ]; - if tcgen05_mma_is_sparse(form) { - arguments.push("tmem_ptr".into()); - } - if tcgen05_mma_has_zero_col_mask(form) { - arguments.push("i64".into()); - } - arguments.extend(["i32".into(), "i32".into(), "i32".into()]); - arguments -} - -fn tcgen05_mma_selector_layout(form: Tcgen05MmaForm) -> Tcgen05MmaSelectorLayout { - let first = - 5 + u8::from(tcgen05_mma_is_sparse(form)) + u8::from(tcgen05_mma_has_zero_col_mask(form)); - if tcgen05_mma_is_ws(form) { - Tcgen05MmaSelectorLayout::WarpSpecialized { - kind_argument: first, - b_buffer_argument: first + 1, - b_usage_argument: first + 2, - } - } else { - Tcgen05MmaSelectorLayout::Base { - kind_argument: first, - cta_group_argument: first + 1, - collector_a_argument: first + 2, - collector_a_upper_exclusive: if tcgen05_mma_is_ashift(form) { 2 } else { 4 }, - } - } -} - -fn tcgen05_mma_imported_properties(form: Tcgen05MmaForm) -> Vec { - let layout = tcgen05_mma_selector_layout(form); - let (kind, second, third, second_lower, second_upper, third_upper) = match layout { - Tcgen05MmaSelectorLayout::Base { - kind_argument, - cta_group_argument, - collector_a_argument, - collector_a_upper_exclusive, - } => ( - kind_argument, - cta_group_argument, - collector_a_argument, - 1, - 3, - collector_a_upper_exclusive, - ), - Tcgen05MmaSelectorLayout::WarpSpecialized { - kind_argument, - b_buffer_argument, - b_usage_argument, - } => (kind_argument, b_buffer_argument, b_usage_argument, 0, 4, 4), - }; - let mut properties = vec![ - format!("ImmArg"), - format!("ImmArg"), - format!("ImmArg"), - "IntrArgMemOnly".into(), - format!("Range"), - format!("Range"), - format!("Range"), - "WriteOnly".into(), - ]; - if tcgen05_mma_is_tensor_a(form) { - properties.push("ReadOnly".into()); - } - if form == Tcgen05MmaForm::Tensor { - properties.extend([ - "ArgInfo".into(), - "ArgInfo".into(), - "ArgInfo".into(), - ]); - } - properties.sort(); - properties -} - -fn tcgen05_mma_selection_asm( - form: Tcgen05MmaForm, - kind: Tcgen05MmaKind, - cta_group: u8, - collector_a: Option<&str>, - b_buffer: Option, - b_usage: Option, -) -> String { - let mut head = "tcgen05.mma".to_owned(); - if tcgen05_mma_is_ws(form) { - head.push_str(".ws"); - } - if tcgen05_mma_is_sparse(form) { - head.push_str(".sp"); - } - head.push_str(&format!( - ".cta_group::{cta_group}.kind::{}", - tcgen05_mma_kind_name(kind) - )); - if tcgen05_mma_is_ws(form) { - head.push_str(&format!( - ".collector::b{}::{}", - b_buffer.expect("warp-specialized B buffer"), - tcgen05_mma_b_usage_name(b_usage.expect("warp-specialized B usage")) - )); - } else { - head.push_str(&format!( - ".collector::a::{}", - collector_a.expect("base collector A usage") - )); - if tcgen05_mma_is_ashift(form) { - head.push_str(".ashift"); - } - } - - let a = if tcgen05_mma_is_tensor_a(form) { - "[$a]" - } else { - "$a" - }; - let mut operands = format!("[$dtmem], {a}, $b"); - if tcgen05_mma_is_sparse(form) { - operands.push_str(", [$spmetadata]"); - } - operands.push_str(", $idesc, $enable_inp_d"); - if tcgen05_mma_has_zero_col_mask(form) { - operands.push_str(", $zero_col_mask"); - } - format!("{head} {operands};") -} - -fn tcgen05_mma_all_selection_asms(form: Tcgen05MmaForm) -> BTreeSet { - if tcgen05_mma_is_ws(form) { - return [ - Tcgen05MmaBUsage::Discard, - Tcgen05MmaBUsage::LastUse, - Tcgen05MmaBUsage::Fill, - Tcgen05MmaBUsage::Use, - ] - .into_iter() - .flat_map(|usage| { - TCGEN05_MMA_KINDS.into_iter().flat_map(move |kind| { - (0..4).map(move |buffer| { - tcgen05_mma_selection_asm(form, kind, 1, None, Some(buffer), Some(usage)) - }) - }) - }) - .collect(); - } - TCGEN05_MMA_KINDS - .into_iter() - .flat_map(|kind| { - (1..=2).flat_map(move |group| { - ["discard", "lastuse", "fill", "use"].map(move |usage| { - tcgen05_mma_selection_asm(form, kind, group, Some(usage), None, None) - }) - }) - }) - .collect() -} - -fn tcgen05_mma_valid_selection_asms(form: Tcgen05MmaForm) -> BTreeSet { - if !tcgen05_mma_is_ashift(form) { - return tcgen05_mma_all_selection_asms(form); - } - TCGEN05_MMA_KINDS - .into_iter() - .flat_map(|kind| { - (1..=2).flat_map(move |group| { - ["discard", "lastuse"].map(move |usage| { - tcgen05_mma_selection_asm(form, kind, group, Some(usage), None, None) - }) - }) - }) - .collect() -} - -fn tcgen05_mma_target_contract( - kind: Tcgen05MmaKind, - alternatives: &[(&str, &str)], -) -> TargetContract { - TargetContract { - selectors: vec![TargetSelectorBinding { - name: "kind".into(), - value: tcgen05_mma_kind_name(kind).into(), - }], - alternatives: alternatives - .iter() - .map( - |(target, minimum_ptx)| crate::model::TargetContractAlternative { - target: (*target).into(), - minimum_ptx: (*minimum_ptx).into(), - }, - ) - .collect(), - } -} - -fn expected_tcgen05_mma_target_contracts(backend: IntrinsicBackend) -> Vec { - const LLVM_COMMON: &[(&str, &str)] = &[ - ("sm_100a", "8.6"), - ("sm_100f", "8.8"), - ("sm_101a", "8.6"), - ("sm_101f", "8.8"), - ("sm_103a", "8.8"), - ("sm_103f", "8.8"), - ("sm_110a", "9.0"), - ("sm_110f", "9.0"), - ]; - const LLVM_I8: &[(&str, &str)] = &[("sm_100a", "8.6"), ("sm_101a", "8.6"), ("sm_110a", "9.0")]; - const LIBNVVM_COMMON: &[(&str, &str)] = &[ - ("sm_100a", "8.6"), - ("sm_100f", "8.8"), - ("sm_103a", "8.8"), - ("sm_103f", "8.8"), - ("sm_110a", "9.0"), - ("sm_110f", "9.0"), - ]; - const LIBNVVM_I8: &[(&str, &str)] = &[("sm_100a", "8.6"), ("sm_110a", "9.0")]; - TCGEN05_MMA_KINDS - .into_iter() - .map(|kind| { - let alternatives = match (backend, kind) { - (IntrinsicBackend::LlvmNvptx, Tcgen05MmaKind::I8) => LLVM_I8, - (IntrinsicBackend::LlvmNvptx, _) => LLVM_COMMON, - (IntrinsicBackend::LibNvvm, Tcgen05MmaKind::I8) => LIBNVVM_I8, - (IntrinsicBackend::LibNvvm, _) => LIBNVVM_COMMON, - }; - tcgen05_mma_target_contract(kind, alternatives) - }) - .collect() -} - -fn tcgen05_mma_expected_ptx( - form: Tcgen05MmaForm, - alias: Option, -) -> InstructionPattern { - let kind = if alias.is_some() { - Tcgen05MmaKind::F8f6f4 - } else { - Tcgen05MmaKind::F16 - }; - let asm = if tcgen05_mma_is_ws(form) { - tcgen05_mma_selection_asm( - form, - kind, - 1, - None, - Some(0), - Some(Tcgen05MmaBUsage::Discard), - ) - } else { - tcgen05_mma_selection_asm(form, kind, 1, Some("discard"), None, None) - }; - let head = asm.split_whitespace().next().unwrap(); - let mut components = head.split('.'); - let mnemonic = components.next().unwrap().into(); - let modifiers = components.map(str::to_owned).collect(); - let mut operands = vec![OperandPattern::Address]; - operands.push(if tcgen05_mma_is_tensor_a(form) { - OperandPattern::Address - } else { - OperandPattern::Register - }); - operands.push(OperandPattern::Register); - if tcgen05_mma_is_sparse(form) { - operands.push(OperandPattern::Address); - } - operands.extend([ - OperandPattern::Register, - OperandPattern::Exact { - value: "%enable_pred".into(), - }, - ]); - if tcgen05_mma_has_zero_col_mask(form) { - operands.push(OperandPattern::Register); - } - InstructionPattern { - mnemonic, - modifiers, - operands, - } -} - -fn tcgen05_mma_rust_arguments(form: Tcgen05MmaForm, alias: Option) -> Vec { - if alias.is_some() && form == Tcgen05MmaForm::WsTensor { - return ["u32", "u32", "u64", "u64", "u32", "bool"] - .into_iter() - .map(str::to_owned) - .collect(); - } - let mut arguments = tcgen05_mma_llvm_arguments(form); - if alias.is_some() { - arguments.truncate(arguments.len() - 3); - } - arguments - .iter() - .map(|argument| match argument.as_str() { - "tmem_ptr" | "i32" => "u32", - "i64" => "u64", - "i1" => "bool", - _ => unreachable!("closed tcgen05 MMA LLVM type"), - }) - .map(str::to_owned) - .collect() -} - -fn tcgen05_mma_dialect_operands( - form: Tcgen05MmaForm, - _alias: Option, -) -> Vec { - let mut arguments = tcgen05_mma_llvm_arguments(form); - arguments.truncate(arguments.len() - 3); - arguments - .iter() - .map(|argument| match argument.as_str() { - "tmem_ptr" | "i32" => "i32", - "i64" => "i64", - "i1" => "i1", - _ => unreachable!("closed tcgen05 MMA LLVM type"), - }) - .map(str::to_owned) - .collect() -} - -fn tcgen05_mma_public_id(form: Tcgen05MmaForm, alias: Option) -> String { - alias.map_or_else( - || format!("tcgen05_mma_{}", tcgen05_mma_form_name(form)), - |alias| { - let alias = tcgen05_mma_alias_name(alias); - if form == Tcgen05MmaForm::WsTensor { - format!("tcgen05_mma_ws_{alias}") - } else { - format!("tcgen05_mma_{alias}") - } - }, - ) -} - -fn tcgen05_mma_operation_key(form: Tcgen05MmaForm, alias: Option) -> String { - let base = tcgen05_mma_form_name(form).replace('_', "."); - alias.map_or_else( - || format!("tcgen05.mma.{base}"), - |alias| format!("tcgen05.mma.{base}.{}", tcgen05_mma_alias_name(alias)), - ) -} - -fn tcgen05_mma_adapter(form: Tcgen05MmaForm, alias: Option) -> Tcgen05Adapter { - match alias { - Some(_) if form == Tcgen05MmaForm::WsTensor => { - Tcgen05Adapter::MmaWsFixedSelectorsDropLegacyADescriptor - } - _ => Tcgen05Adapter::MmaDirectSelectors, - } -} - -fn materialize_tcgen05_mma_variant( - admission: &Tcgen05Admission, - variant: &Tcgen05MmaAdmissionVariant, - llvm_target: &CatalogTargetRequirement, - libnvvm_target: &CatalogTargetRequirement, -) -> OverlayIntrinsic { - let form = variant.form; - let alias = variant.alias; - let id = tcgen05_mma_public_id(form, alias); - let operation_key = tcgen05_mma_operation_key(form, alias); - let llvm_arguments = tcgen05_mma_llvm_arguments(form); - let rust_arguments = tcgen05_mma_rust_arguments(form, alias); - let dialect_operands = tcgen05_mma_dialect_operands(form, alias); - let fixed_selectors = alias.map(|_| Tcgen05MmaFixedSelectors { - kind: Tcgen05MmaKind::F8f6f4, - b_buffer: 0, - b_usage: Tcgen05MmaBUsage::Discard, - }); - - OverlayIntrinsic { - id: id.clone(), - abi_id: variant.abi_id.clone(), - operation_key, - family: "tcgen05".into(), - source: None, - source_record: Some(tcgen05_mma_source_record(form)), - rust_module: "tcgen05".into(), - rust_name: id.clone(), - rust_arguments, - rust_result: "()".into(), - safe: false, - must_use: false, - safe_allowlist_reason: None, - public_rust_path: format!("cuda_intrinsics::tcgen05::{id}"), - compatibility_rust_paths: vec![if alias.is_some() { - format!("cuda_device::tcgen05::{id}") - } else { - format!("cuda_device::tcgen05::__{id}") - }], - dialect_op_type: TCGEN05_MMA_DIALECT_OP_TYPE.into(), - dialect_op_name: TCGEN05_MMA_DIALECT_OP_NAME.into(), - dialect_operands, - dialect_results: vec![], - llvm_symbol: Some(tcgen05_mma_llvm_symbol(form)), - resolved_llvm_symbol: None, - llvm_arguments, - llvm_results: vec![], - pure: false, - memory: "read_write".into(), - // TableGen omits IntrConvergent, but collective MMA must not move across control flow. - convergent: true, - execution_scope: "thread".into(), - minimum_ptx: "8.6".into(), - minimum_sm: None, - ptx_result: "()".into(), - targets: TCGEN05_LLVM_TARGETS.into(), - ptx_isa_version: "8.6".into(), - ptx_isa_section: "Tensor Memory tcgen05 instructions".into(), - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#tensor-memory".into(), - lowering: "generated_tcgen05_mma".into(), - backend_lowerings: vec![ - OverlayBackendLowering { - backend: IntrinsicBackend::LlvmNvptx, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: admission - .mma_llvm_evidence_profile - .as_ref() - .expect("validated tcgen05 MMA LLVM evidence profile") - .clone(), - targets: None, - minimum_ptx: None, - minimum_sm: None, - }, - OverlayBackendLowering { - backend: IntrinsicBackend::LibNvvm, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: admission - .mma_libnvvm_evidence_profile - .as_ref() - .expect("validated tcgen05 MMA libNVVM evidence profile") - .clone(), - targets: None, - minimum_ptx: None, - minimum_sm: None, - }, - ], - packed_atomic: None, - redux: None, - vote: None, - active_mask: None, - warp_match: None, - warp_barrier: None, - warp_shuffle: None, - dot_product: None, - packed_alu: None, - integer_minmax: None, - packed_conversion: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cp_async_copy: None, - cp_async_control: None, - cp_async_mbarrier: None, - mbarrier_basic: None, - movmatrix: None, - mbarrier_extended: None, - register_mma: None, - sparse_mma: None, - prmt: None, - cluster_barrier: None, - wgmma_control: None, - special_register: None, - debug_control: None, - cluster_memory: None, - clc: None, - tma: None, - tcgen05: Some(Tcgen05 { - operation: Tcgen05Operation::Mma, - cp: None, - ld: None, - st: None, - mma: Some(Tcgen05Mma { - form, - selector_layout: tcgen05_mma_selector_layout(form), - fixed_selectors, - alias, - llvm_target: llvm_target.clone(), - libnvvm_target: libnvvm_target.clone(), - }), - adapter: tcgen05_mma_adapter(form, alias), - source_contract: Tcgen05SourceContract::TablegenSelectionChangesPtx, - runtime_validation: admission.runtime_validation, - }), - ldmatrix_variant: None, - ldmatrix_safety: None, - ldmatrix_adapter: None, - selected_address_space: None, - expected_ptx: tcgen05_mma_expected_ptx(form, alias), - summary: match (form, alias) { - (Tcgen05MmaForm::WsTensor, Some(_)) => { - "Issues one f8f6f4 warp-specialized tensor-memory MMA.".into() - } - (Tcgen05MmaForm::Shared, Some(_)) => { - "Issues one f8f6f4 standard tensor-memory MMA.".into() - } - _ => "Issues one selector-controlled tensor-memory MMA.".into(), - }, - } -} - -struct Tcgen05CpMemberRecipe { - llvm_suffix: &'static str, - ptx_suffix: &'static str, - op_suffix: &'static str, - selection_stem: &'static str, -} - -const TCGEN05_CP_MEMBERS: [Tcgen05CpMember; 17] = [ - Tcgen05CpMember::M128x128bB4x16P64, - Tcgen05CpMember::M128x128bB6x16P32, - Tcgen05CpMember::M128x128b, - Tcgen05CpMember::M128x256bB4x16P64, - Tcgen05CpMember::M128x256bB6x16P32, - Tcgen05CpMember::M32x128bWarpx4B4x16P64, - Tcgen05CpMember::M32x128bWarpx4B6x16P32, - Tcgen05CpMember::M32x128bWarpx4, - Tcgen05CpMember::M4x256bB4x16P64, - Tcgen05CpMember::M4x256bB6x16P32, - Tcgen05CpMember::M4x256b, - Tcgen05CpMember::M64x128bWarpx2Pair0123B4x16P64, - Tcgen05CpMember::M64x128bWarpx2Pair0123B6x16P32, - Tcgen05CpMember::M64x128bWarpx2Pair0123, - Tcgen05CpMember::M64x128bWarpx2Pair0213B4x16P64, - Tcgen05CpMember::M64x128bWarpx2Pair0213B6x16P32, - Tcgen05CpMember::M64x128bWarpx2Pair0213, -]; - -fn tcgen05_cp_member_recipe(member: Tcgen05CpMember) -> Tcgen05CpMemberRecipe { - use Tcgen05CpMember::*; - let (llvm_suffix, ptx_suffix, op_suffix, selection_stem) = match member { - M128x128bB4x16P64 => ( - "128x128b.b4x16_p64", - "128x128b.b8x16.b4x16_p64", - "128x128bB4x16P64", - "128x128bb4x16_p64", - ), - M128x128bB6x16P32 => ( - "128x128b.b6x16_p32", - "128x128b.b8x16.b6x16_p32", - "128x128bB6x16P32", - "128x128bb6x16_p32", - ), - M128x128b => ("128x128b", "128x128b", "128x128b", "128x128b"), - M128x256bB4x16P64 => ( - "128x256b.b4x16_p64", - "128x256b.b8x16.b4x16_p64", - "128x256bB4x16P64", - "128x256bb4x16_p64", - ), - M128x256bB6x16P32 => ( - "128x256b.b6x16_p32", - "128x256b.b8x16.b6x16_p32", - "128x256bB6x16P32", - "128x256bb6x16_p32", - ), - M32x128bWarpx4B4x16P64 => ( - "32x128b_warpx4.b4x16_p64", - "32x128b.warpx4.b8x16.b4x16_p64", - "32x128bWarpx4B4x16P64", - "32x128b4x16_p64", - ), - M32x128bWarpx4B6x16P32 => ( - "32x128b_warpx4.b6x16_p32", - "32x128b.warpx4.b8x16.b6x16_p32", - "32x128bWarpx4B6x16P32", - "32x128b6x16_p32", - ), - M32x128bWarpx4 => ( - "32x128b_warpx4", - "32x128b.warpx4", - "32x128bWarpx4", - "32x128", - ), - M4x256bB4x16P64 => ( - "4x256b.b4x16_p64", - "4x256b.b8x16.b4x16_p64", - "4x256bB4x16P64", - "4x256bb4x16_p64", - ), - M4x256bB6x16P32 => ( - "4x256b.b6x16_p32", - "4x256b.b8x16.b6x16_p32", - "4x256bB6x16P32", - "4x256bb6x16_p32", - ), - M4x256b => ("4x256b", "4x256b", "4x256b", "4x256b"), - M64x128bWarpx2Pair0123B4x16P64 => ( - "64x128b_warpx2_01_23.b4x16_p64", - "64x128b.warpx2::01_23.b8x16.b4x16_p64", - "64x128bWarpx2Pair0123B4x16P64", - "64x128_2b4x16_p64", - ), - M64x128bWarpx2Pair0123B6x16P32 => ( - "64x128b_warpx2_01_23.b6x16_p32", - "64x128b.warpx2::01_23.b8x16.b6x16_p32", - "64x128bWarpx2Pair0123B6x16P32", - "64x128_2b6x16_p32", - ), - M64x128bWarpx2Pair0123 => ( - "64x128b_warpx2_01_23", - "64x128b.warpx2::01_23", - "64x128bWarpx2Pair0123", - "64x128_2", - ), - M64x128bWarpx2Pair0213B4x16P64 => ( - "64x128b_warpx2_02_13.b4x16_p64", - "64x128b.warpx2::02_13.b8x16.b4x16_p64", - "64x128bWarpx2Pair0213B4x16P64", - "64x128_1b4x16_p64", - ), - M64x128bWarpx2Pair0213B6x16P32 => ( - "64x128b_warpx2_02_13.b6x16_p32", - "64x128b.warpx2::02_13.b8x16.b6x16_p32", - "64x128bWarpx2Pair0213B6x16P32", - "64x128_1b6x16_p32", - ), - M64x128bWarpx2Pair0213 => ( - "64x128b_warpx2_02_13", - "64x128b.warpx2::02_13", - "64x128bWarpx2Pair0213", - "64x128_1", - ), - }; - Tcgen05CpMemberRecipe { - llvm_suffix, - ptx_suffix, - op_suffix, - selection_stem, - } -} - -fn materialize_tcgen05_cp_variant( - base: &OverlayIntrinsic, - admission: &Tcgen05Admission, - variant: &Tcgen05CpAdmissionVariant, -) -> OverlayIntrinsic { - let recipe = tcgen05_cp_member_recipe(variant.member); - let group = match variant.group { - Tcgen05CpGroup::Cg1 => 1, - Tcgen05CpGroup::Cg2 => 2, - }; - let group_suffix = if group == 1 { "" } else { "_cg2" }; - let id_suffix = recipe.llvm_suffix.replace('.', "_"); - let id = format!("tcgen05_cp_{id_suffix}{group_suffix}"); - let mut record = base.clone(); - record.id = id.clone(); - record.abi_id = variant.abi_id.clone(); - record.operation_key = format!("tcgen05.cp.{}.cg{group}", recipe.llvm_suffix); - record.source_record = Some(format!("int_nvvm_tcgen05_cp_{}_cg{group}", id_suffix)); - record.rust_name = id.clone(); - record.public_rust_path = format!("cuda_intrinsics::tcgen05::{id}"); - record.compatibility_rust_paths = vec![format!("cuda_device::tcgen05::{id}")]; - record.dialect_op_type = format!( - "Tcgen05Cp{}{}Op", - recipe.op_suffix, - if group == 1 { "" } else { "Cg2" } - ); - record.dialect_op_name = format!("nvvm.{id}"); - record.llvm_symbol = Some(format!( - "llvm.nvvm.tcgen05.cp.{}.cg{group}", - recipe.llvm_suffix - )); - record.backend_lowerings[0].evidence_profile = admission - .cp_llvm_evidence_profile - .as_ref() - .expect("validated tcgen05 copy LLVM evidence profile") - .clone(); - record.backend_lowerings[1].evidence_profile = admission - .cp_libnvvm_evidence_profile - .as_ref() - .expect("validated tcgen05 copy libNVVM evidence profile") - .clone(); - record.tcgen05 = Some(Tcgen05 { - operation: if group == 1 { - Tcgen05Operation::CpSmemToTmem - } else { - Tcgen05Operation::CpSmemToTmemCg2 - }, - cp: Some(Tcgen05Cp { - member: variant.member, - group: variant.group, - }), - ld: None, - st: None, - mma: None, - adapter: Tcgen05Adapter::TmemDescriptorToVoid, - source_contract: Tcgen05SourceContract::ExactTablegenSelection, - runtime_validation: admission.runtime_validation, - }); - record.expected_ptx.modifiers = std::iter::once("cp".into()) - .chain(std::iter::once(format!("cta_group::{group}"))) - .chain(recipe.ptx_suffix.split('.').map(Into::into)) - .collect(); - record.summary = format!( - "Copies one {} tile from shared memory to tensor memory.", - recipe.ptx_suffix - ); - record -} - -const TCGEN05_LD_VARIANTS: [(Tcgen05LdShape, Tcgen05LdMultiplicity); 29] = { - use Tcgen05LdMultiplicity::*; - use Tcgen05LdShape::*; - [ - (M16x64b, X1), - (M16x64b, X2), - (M16x64b, X4), - (M16x64b, X8), - (M16x64b, X16), - (M16x64b, X32), - (M16x64b, X64), - (M16x64b, X128), - (M16x128b, X1), - (M16x128b, X2), - (M16x128b, X4), - (M16x128b, X8), - (M16x128b, X16), - (M16x128b, X32), - (M16x128b, X64), - (M16x256b, X1), - (M16x256b, X2), - (M16x256b, X4), - (M16x256b, X8), - (M16x256b, X16), - (M16x256b, X32), - (M32x32b, X1), - (M32x32b, X2), - (M32x32b, X4), - (M32x32b, X8), - (M32x32b, X16), - (M32x32b, X32), - (M32x32b, X64), - (M32x32b, X128), - ] -}; - -const TCGEN05_OFFSET_LDST_VARIANTS: [(Tcgen05LdShape, Tcgen05LdMultiplicity); 8] = { - use Tcgen05LdMultiplicity::*; - use Tcgen05LdShape::*; - [ - (M16x32bx2, X1), - (M16x32bx2, X2), - (M16x32bx2, X4), - (M16x32bx2, X8), - (M16x32bx2, X16), - (M16x32bx2, X32), - (M16x32bx2, X64), - (M16x32bx2, X128), - ] -}; - -fn tcgen05_ld_shape_name(shape: Tcgen05LdShape) -> &'static str { - match shape { - Tcgen05LdShape::M16x32bx2 => "16x32bx2", - Tcgen05LdShape::M16x64b => "16x64b", - Tcgen05LdShape::M16x128b => "16x128b", - Tcgen05LdShape::M16x256b => "16x256b", - Tcgen05LdShape::M32x32b => "32x32b", - } -} - -fn tcgen05_ld_multiplicity_name(multiplicity: Tcgen05LdMultiplicity) -> &'static str { - match multiplicity { - Tcgen05LdMultiplicity::X1 => "x1", - Tcgen05LdMultiplicity::X2 => "x2", - Tcgen05LdMultiplicity::X4 => "x4", - Tcgen05LdMultiplicity::X8 => "x8", - Tcgen05LdMultiplicity::X16 => "x16", - Tcgen05LdMultiplicity::X32 => "x32", - Tcgen05LdMultiplicity::X64 => "x64", - Tcgen05LdMultiplicity::X128 => "x128", - } -} - -fn tcgen05_ld_register_count(ld: Tcgen05Ld) -> usize { - ld.shape.register_multiplier() * ld.multiplicity.count() -} - -fn tcgen05_ld_id(ld: Tcgen05Ld) -> String { - format!( - "tcgen05_ld_{}_{}_{}", - tcgen05_ld_shape_name(ld.shape), - tcgen05_ld_multiplicity_name(ld.multiplicity), - if ld.pack16 { "pack16" } else { "raw" } - ) -} - -fn tcgen05_ld_source_record(ld: Tcgen05Ld) -> String { - format!( - "int_nvvm_tcgen05_ld_{}_{}", - tcgen05_ld_shape_name(ld.shape), - tcgen05_ld_multiplicity_name(ld.multiplicity) - ) -} - -fn tcgen05_ld_llvm_symbol(ld: Tcgen05Ld) -> String { - format!( - "llvm.nvvm.tcgen05.ld.{}.{}", - tcgen05_ld_shape_name(ld.shape), - tcgen05_ld_multiplicity_name(ld.multiplicity) - ) -} - -fn tcgen05_ld_rust_result(ld: Tcgen05Ld) -> String { - match tcgen05_ld_register_count(ld) { - 1 => "u32".into(), - count => format!("[u32; {count}]"), - } -} - -fn tcgen05_ld_llvm_result(ld: Tcgen05Ld) -> String { - match tcgen05_ld_register_count(ld) { - 1 => "i32".into(), - count => format!("v{count}i32"), - } -} - -fn tcgen05_ld_op_type(ld: Tcgen05Ld) -> String { - let shape = tcgen05_ld_shape_name(ld.shape); - let multiplicity = tcgen05_ld_multiplicity_name(ld.multiplicity) - .strip_prefix('x') - .unwrap(); - format!( - "Tcgen05Ld{shape}X{multiplicity}{}Op", - if ld.pack16 { "Pack16" } else { "Raw" } - ) -} - -fn tcgen05_ld_modifiers(ld: Tcgen05Ld) -> Vec { - let mut modifiers = vec![ - "ld".into(), - "sync".into(), - "aligned".into(), - tcgen05_ld_shape_name(ld.shape).into(), - tcgen05_ld_multiplicity_name(ld.multiplicity).into(), - ]; - if ld.pack16 { - modifiers.push("pack::16b".into()); - } - modifiers.push("b32".into()); - modifiers -} - -fn tcgen05_ld_operands(ld: Tcgen05Ld) -> Vec { - let result = OperandPattern::RegisterList { - length: tcgen05_ld_register_count(ld), - }; - let mut operands = vec![result, OperandPattern::Address]; - if ld.shape == Tcgen05LdShape::M16x32bx2 { - operands.push(OperandPattern::Immediate); - } - operands -} - -fn materialize_tcgen05_ld_variant( - base: &OverlayIntrinsic, - admission: &Tcgen05Admission, - variant: &Tcgen05LdAdmissionVariant, -) -> OverlayIntrinsic { - let ld = Tcgen05Ld { - shape: variant.shape, - multiplicity: variant.multiplicity, - pack16: variant.pack16, - }; - let id = tcgen05_ld_id(ld); - let register_count = tcgen05_ld_register_count(ld); - let rust_result = tcgen05_ld_rust_result(ld); - let llvm_result = tcgen05_ld_llvm_result(ld); - let mode = if ld.pack16 { "pack16" } else { "raw" }; - let mut record = base.clone(); - record.id = id.clone(); - record.abi_id = variant.abi_id.clone(); - record.operation_key = format!( - "tcgen05.ld.{}.{}.{}", - tcgen05_ld_shape_name(ld.shape), - tcgen05_ld_multiplicity_name(ld.multiplicity), - mode - ); - record.source_record = Some(tcgen05_ld_source_record(ld)); - record.rust_name = id.clone(); - let has_half_split_offset = ld.shape == Tcgen05LdShape::M16x32bx2; - record.rust_arguments = if has_half_split_offset { - vec!["u32".into(), "i64".into()] - } else { - vec!["u32".into()] - }; - record.rust_result = rust_result.clone(); - record.must_use = true; - record.public_rust_path = format!("cuda_intrinsics::tcgen05::{id}"); - record.compatibility_rust_paths = vec![format!( - "cuda_device::tcgen05::{}", - if has_half_split_offset { - format!("__{id}") - } else { - id.clone() - } - )]; - record.dialect_op_type = tcgen05_ld_op_type(ld); - record.dialect_op_name = format!("nvvm.{id}"); - record.dialect_operands = if has_half_split_offset { - vec!["i32".into(), "i64".into()] - } else { - vec!["i32".into()] - }; - record.dialect_results = vec!["i32".into(); register_count]; - record.llvm_symbol = Some(tcgen05_ld_llvm_symbol(ld)); - record.resolved_llvm_symbol = None; - record.llvm_arguments = if has_half_split_offset { - vec!["tmem_ptr".into(), "i64".into(), "i1".into()] - } else { - vec!["tmem_ptr".into(), "i1".into()] - }; - record.llvm_results = vec![llvm_result]; - record.ptx_result = rust_result; - record.execution_scope = Tcgen05Operation::Ld.execution_scope().into(); - record.backend_lowerings[0].evidence_profile = if has_half_split_offset { - &admission.offset_llvm_evidence_profile - } else { - &admission.ld_llvm_evidence_profile - } - .as_ref() - .expect("validated tcgen05 load LLVM evidence profile") - .clone(); - record.backend_lowerings[1].evidence_profile = if has_half_split_offset { - &admission.offset_libnvvm_evidence_profile - } else { - &admission.ld_libnvvm_evidence_profile - } - .as_ref() - .expect("validated tcgen05 load libNVVM evidence profile") - .clone(); - record.tcgen05 = Some(Tcgen05 { - operation: Tcgen05Operation::Ld, - cp: None, - ld: Some(ld), - st: None, - mma: None, - adapter: if has_half_split_offset { - Tcgen05Adapter::TmemHalfSplitOffsetInjectPack16ToU32Registers - } else { - Tcgen05Adapter::TmemInjectPack16ToU32Registers - }, - source_contract: Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection, - runtime_validation: admission.runtime_validation, - }); - record.expected_ptx.modifiers = tcgen05_ld_modifiers(ld); - record.expected_ptx.operands = tcgen05_ld_operands(ld); - record.summary = format!( - "Loads {register_count} {} 32-bit register value{} from tensor memory.", - if ld.pack16 { "packed" } else { "raw" }, - if register_count == 1 { "" } else { "s" } - ); - record -} - -const TCGEN05_ST_VARIANTS: [(Tcgen05LdShape, Tcgen05LdMultiplicity); 29] = TCGEN05_LD_VARIANTS; - -fn tcgen05_st_register_count(st: Tcgen05St) -> usize { - st.shape.register_multiplier() * st.multiplicity.count() -} - -fn tcgen05_st_id(st: Tcgen05St) -> String { - format!( - "tcgen05_st_{}_{}_{}", - tcgen05_ld_shape_name(st.shape), - tcgen05_ld_multiplicity_name(st.multiplicity), - if st.unpack16 { "unpack16" } else { "raw" } - ) -} - -fn tcgen05_st_source_record(st: Tcgen05St) -> String { - format!( - "int_nvvm_tcgen05_st_{}_{}", - tcgen05_ld_shape_name(st.shape), - tcgen05_ld_multiplicity_name(st.multiplicity) - ) -} - -fn tcgen05_st_llvm_symbol(st: Tcgen05St) -> String { - format!( - "llvm.nvvm.tcgen05.st.{}.{}", - tcgen05_ld_shape_name(st.shape), - tcgen05_ld_multiplicity_name(st.multiplicity) - ) -} - -fn tcgen05_st_rust_data(st: Tcgen05St) -> String { - match tcgen05_st_register_count(st) { - 1 => "u32".into(), - count => format!("[u32; {count}]"), - } -} - -fn tcgen05_st_llvm_data(st: Tcgen05St) -> String { - match tcgen05_st_register_count(st) { - 1 => "i32".into(), - count => format!("v{count}i32"), - } -} - -fn tcgen05_st_op_type(st: Tcgen05St) -> String { - let shape = tcgen05_ld_shape_name(st.shape); - let multiplicity = tcgen05_ld_multiplicity_name(st.multiplicity) - .strip_prefix('x') - .unwrap(); - format!( - "Tcgen05St{shape}X{multiplicity}{}Op", - if st.unpack16 { "Unpack16" } else { "Raw" } - ) -} - -fn tcgen05_st_modifiers(st: Tcgen05St) -> Vec { - let mut modifiers = vec![ - "st".into(), - "sync".into(), - "aligned".into(), - tcgen05_ld_shape_name(st.shape).into(), - tcgen05_ld_multiplicity_name(st.multiplicity).into(), - ]; - if st.unpack16 { - modifiers.push("unpack::16b".into()); - } - modifiers.push("b32".into()); - modifiers -} - -fn tcgen05_st_operands(st: Tcgen05St) -> Vec { - let data = OperandPattern::RegisterList { - length: tcgen05_st_register_count(st), - }; - let mut operands = vec![OperandPattern::Address]; - if st.shape == Tcgen05LdShape::M16x32bx2 { - operands.push(OperandPattern::Immediate); - } - operands.push(data); - operands -} - -fn materialize_tcgen05_st_variant( - base: &OverlayIntrinsic, - admission: &Tcgen05Admission, - variant: &Tcgen05StAdmissionVariant, -) -> OverlayIntrinsic { - let st = Tcgen05St { - shape: variant.shape, - multiplicity: variant.multiplicity, - unpack16: variant.unpack16, - }; - let id = tcgen05_st_id(st); - let register_count = tcgen05_st_register_count(st); - let rust_data = tcgen05_st_rust_data(st); - let llvm_data = tcgen05_st_llvm_data(st); - let mode = if st.unpack16 { "unpack16" } else { "raw" }; - let mut record = base.clone(); - record.id = id.clone(); - record.abi_id = variant.abi_id.clone(); - record.operation_key = format!( - "tcgen05.st.{}.{}.{}", - tcgen05_ld_shape_name(st.shape), - tcgen05_ld_multiplicity_name(st.multiplicity), - mode - ); - record.source_record = Some(tcgen05_st_source_record(st)); - record.rust_name = id.clone(); - let has_half_split_offset = st.shape == Tcgen05LdShape::M16x32bx2; - record.rust_arguments = if has_half_split_offset { - vec!["u32".into(), "i64".into(), rust_data] - } else { - vec!["u32".into(), rust_data] - }; - record.rust_result = "()".into(); - record.must_use = false; - record.public_rust_path = format!("cuda_intrinsics::tcgen05::{id}"); - record.compatibility_rust_paths = vec![format!( - "cuda_device::tcgen05::{}", - if has_half_split_offset { - format!("__{id}") - } else { - id.clone() - } - )]; - record.dialect_op_type = tcgen05_st_op_type(st); - record.dialect_op_name = format!("nvvm.{id}"); - record.dialect_operands = std::iter::once("i32".into()) - .chain(has_half_split_offset.then(|| "i64".into())) - .chain(std::iter::repeat_n("i32".into(), register_count)) - .collect(); - record.dialect_results.clear(); - record.llvm_symbol = Some(tcgen05_st_llvm_symbol(st)); - record.resolved_llvm_symbol = None; - record.llvm_arguments = if has_half_split_offset { - vec!["tmem_ptr".into(), "i64".into(), llvm_data, "i1".into()] - } else { - vec!["tmem_ptr".into(), llvm_data, "i1".into()] - }; - record.llvm_results.clear(); - record.memory = "write".into(); - record.ptx_result = "()".into(); - record.execution_scope = Tcgen05Operation::St.execution_scope().into(); - record.backend_lowerings[0].evidence_profile = if has_half_split_offset { - &admission.offset_llvm_evidence_profile - } else { - &admission.st_llvm_evidence_profile - } - .as_ref() - .expect("validated tcgen05 store LLVM evidence profile") - .clone(); - record.backend_lowerings[1].evidence_profile = if has_half_split_offset { - &admission.offset_libnvvm_evidence_profile - } else { - &admission.st_libnvvm_evidence_profile - } - .as_ref() - .expect("validated tcgen05 store libNVVM evidence profile") - .clone(); - record.tcgen05 = Some(Tcgen05 { - operation: Tcgen05Operation::St, - cp: None, - ld: None, - st: Some(st), - mma: None, - adapter: if has_half_split_offset { - Tcgen05Adapter::TmemHalfSplitOffsetU32RegistersInjectUnpack16ToVoid - } else { - Tcgen05Adapter::TmemU32RegistersInjectUnpack16ToVoid - }, - source_contract: Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection, - runtime_validation: admission.runtime_validation, - }); - record.expected_ptx.modifiers = tcgen05_st_modifiers(st); - record.expected_ptx.operands = tcgen05_st_operands(st); - record.summary = format!( - "Stores {register_count} {} 32-bit register value{} to tensor memory.", - if st.unpack16 { "unpacked" } else { "raw" }, - if register_count == 1 { "" } else { "s" } - ); - record -} - -fn expand_tcgen05_admission(admission: &Tcgen05Admission) -> Result> { - ensure!( - admission.runtime_validation == RuntimeValidation::Unexecuted, - "tcgen05 runtime validation may be marked executed only with GPU evidence" - ); - ensure!( - !admission.llvm_evidence_profile.trim().is_empty() - && !admission.libnvvm_evidence_profile.trim().is_empty(), - "compact tcgen05 admission requires both backend evidence profiles" - ); - let expected = [ - Tcgen05Operation::Alloc, - Tcgen05Operation::Dealloc, - Tcgen05Operation::RelinquishAllocPermit, - Tcgen05Operation::FenceBeforeThreadSync, - Tcgen05Operation::FenceAfterThreadSync, - Tcgen05Operation::Commit, - Tcgen05Operation::CommitSharedCluster, - Tcgen05Operation::MmaWsF16, - Tcgen05Operation::MmaF16, - Tcgen05Operation::MmaWsBf16, - Tcgen05Operation::MmaWsTf32, - Tcgen05Operation::CpSmemToTmem, - Tcgen05Operation::Ld16x256bX8Pure, - Tcgen05Operation::Ld16x256bPure, - Tcgen05Operation::LoadWait, - Tcgen05Operation::StoreWait, - Tcgen05Operation::AllocCg2, - Tcgen05Operation::DeallocCg2, - Tcgen05Operation::RelinquishAllocPermitCg2, - Tcgen05Operation::MmaF16Cg2, - Tcgen05Operation::CommitCg2, - Tcgen05Operation::CommitSharedClusterCg2, - Tcgen05Operation::CommitMulticastCg2, - Tcgen05Operation::CpSmemToTmemCg2, - Tcgen05Operation::CommitMulticast, - Tcgen05Operation::ShiftDown, - Tcgen05Operation::ShiftDownCg2, - ]; - let has_control_variants = admission.variants.iter().any(|variant| { - matches!( - variant.operation, - Tcgen05Operation::CommitMulticast - | Tcgen05Operation::ShiftDown - | Tcgen05Operation::ShiftDownCg2 - ) - }); - if has_control_variants { - ensure!( - admission - .control_llvm_evidence_profile - .as_deref() - .is_some_and(|profile| !profile.trim().is_empty()) - && admission - .control_libnvvm_evidence_profile - .as_deref() - .is_some_and(|profile| !profile.trim().is_empty()), - "compact tcgen05 control admission requires both backend evidence profiles" - ); - } - let expected = if has_control_variants { - &expected[..] - } else { - &expected[..24] - }; - ensure!( - admission - .variants - .iter() - .map(|variant| variant.operation) - .eq(expected.iter().copied()), - "compact tcgen05 admission must list all 24 base operations or all 27 current operations in canonical order" - ); - - let mut records = admission - .variants - .iter() - .map(|variant| { - let recipe = tcgen05_recipe(variant.operation); - let control = matches!( - variant.operation, - Tcgen05Operation::CommitMulticast - | Tcgen05Operation::ShiftDown - | Tcgen05Operation::ShiftDownCg2 - ); - ensure!( - variant.abi_id == recipe.abi_id, - "{} must keep reserved ABI ID {}", - recipe.id, - recipe.abi_id - ); - Ok(OverlayIntrinsic { - id: recipe.id.into(), - abi_id: variant.abi_id.clone(), - operation_key: recipe.operation_key.into(), - family: "tcgen05".into(), - source: None, - source_record: Some(recipe.source_record.into()), - rust_module: "tcgen05".into(), - rust_name: recipe.id.into(), - rust_arguments: recipe - .rust_arguments - .iter() - .map(|value| (*value).into()) - .collect(), - rust_result: recipe.rust_result.into(), - safe: recipe.safe, - must_use: false, - safe_allowlist_reason: recipe.safe_reason.map(Into::into), - public_rust_path: format!("cuda_intrinsics::tcgen05::{}", recipe.id), - compatibility_rust_paths: vec![format!("cuda_device::tcgen05::{}", recipe.id)], - dialect_op_type: recipe.dialect_op_type.into(), - dialect_op_name: recipe.dialect_op_name.into(), - dialect_operands: recipe - .dialect_operands - .iter() - .map(|value| (*value).into()) - .collect(), - dialect_results: recipe - .dialect_results - .iter() - .map(|value| (*value).into()) - .collect(), - llvm_symbol: Some(recipe.llvm_symbol.into()), - resolved_llvm_symbol: None, - llvm_arguments: recipe - .llvm_arguments - .iter() - .map(|value| (*value).into()) - .collect(), - llvm_results: recipe - .llvm_results - .iter() - .map(|value| (*value).into()) - .collect(), - pure: false, - memory: recipe.memory.into(), - convergent: true, - execution_scope: recipe.operation.execution_scope().into(), - minimum_ptx: "8.6".into(), - minimum_sm: None, - ptx_result: recipe.rust_result.into(), - targets: TCGEN05_LLVM_TARGETS.into(), - ptx_isa_version: "8.6".into(), - ptx_isa_section: "Tensor Memory tcgen05 instructions".into(), - ptx_isa_url: - "https://docs.nvidia.com/cuda/parallel-thread-execution/#tensor-memory".into(), - lowering: "generated_tcgen05".into(), - backend_lowerings: vec![ - OverlayBackendLowering { - backend: IntrinsicBackend::LlvmNvptx, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: if control { - admission - .control_llvm_evidence_profile - .as_ref() - .expect("validated tcgen05 control LLVM evidence profile") - .clone() - } else { - admission.llvm_evidence_profile.clone() - }, - targets: None, - minimum_ptx: Some("8.6".into()), - minimum_sm: None, - }, - OverlayBackendLowering { - backend: IntrinsicBackend::LibNvvm, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: if control { - admission - .control_libnvvm_evidence_profile - .as_ref() - .expect("validated tcgen05 control libNVVM evidence profile") - .clone() - } else { - admission.libnvvm_evidence_profile.clone() - }, - targets: Some(TCGEN05_LIBNVVM_TARGETS.into()), - minimum_ptx: Some("8.6".into()), - minimum_sm: None, - }, - ], - packed_atomic: None, - redux: None, - vote: None, - active_mask: None, - warp_match: None, - warp_barrier: None, - warp_shuffle: None, - dot_product: None, - packed_alu: None, - integer_minmax: None, - packed_conversion: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cp_async_copy: None, - cp_async_control: None, - cp_async_mbarrier: None, - mbarrier_basic: None, - movmatrix: None, - mbarrier_extended: None, - register_mma: None, - sparse_mma: None, - prmt: None, - cluster_barrier: None, - wgmma_control: None, - special_register: None, - debug_control: None, - cluster_memory: None, - clc: None, - tma: None, - tcgen05: Some(Tcgen05 { - operation: recipe.operation, - cp: None, - ld: None, - st: None, - mma: None, - adapter: recipe.adapter, - source_contract: recipe.source_contract, - runtime_validation: admission.runtime_validation, - }), - ldmatrix_variant: None, - ldmatrix_safety: None, - ldmatrix_adapter: None, - selected_address_space: None, - expected_ptx: InstructionPattern { - mnemonic: "tcgen05".into(), - modifiers: recipe.modifiers, - operands: recipe.operands, - }, - summary: recipe.summary.into(), - }) - }) - .collect::>>()?; - - if admission.cp_variants.is_empty() { - ensure!( - admission.cp_llvm_evidence_profile.is_none() - && admission.cp_libnvvm_evidence_profile.is_none(), - "tcgen05 copy evidence profiles require admitted copy variants" - ); - } else { - ensure!( - [ - Tcgen05MmaBUsage::Discard, - Tcgen05MmaBUsage::LastUse, - Tcgen05MmaBUsage::Fill, - Tcgen05MmaBUsage::Use, - ] - .map(Tcgen05MmaBUsage::selector_value) - == [0, 1, 2, 3], - "tcgen05 MMA B-usage selector mapping changed" - ); - ensure!( - admission - .cp_llvm_evidence_profile - .as_deref() - .is_some_and(|profile| !profile.trim().is_empty()) - && admission - .cp_libnvvm_evidence_profile - .as_deref() - .is_some_and(|profile| !profile.trim().is_empty()), - "compact tcgen05 copy admission requires both backend evidence profiles" - ); - - use Tcgen05CpGroup::{Cg1, Cg2}; - let expected_cp = TCGEN05_CP_MEMBERS - .into_iter() - .flat_map(|member| [Cg1, Cg2].into_iter().map(move |group| (member, group))) - .collect::>(); - ensure!( - admission - .cp_variants - .iter() - .map(|variant| (variant.member, variant.group)) - .eq(expected_cp), - "compact tcgen05 copy admission must list all 34 variants in canonical order" - ); - for variant in &admission.cp_variants { - validate_abi_id(&variant.abi_id)?; - let base_id = if variant.group == Cg1 { - "tcgen05_cp_smem_to_tmem" - } else { - "tcgen05_cp_smem_to_tmem_cg2" - }; - let base = records - .iter() - .find(|record| record.id == base_id) - .unwrap() - .clone(); - records.push(materialize_tcgen05_cp_variant(&base, admission, variant)); - } - } - - if admission.ld_variants.is_empty() { - ensure!( - admission.ld_llvm_evidence_profile.is_none() - && admission.ld_libnvvm_evidence_profile.is_none(), - "tcgen05 load evidence profiles require admitted load variants" - ); - } else { - ensure!( - admission - .ld_llvm_evidence_profile - .as_deref() - .is_some_and(|profile| !profile.trim().is_empty()) - && admission - .ld_libnvvm_evidence_profile - .as_deref() - .is_some_and(|profile| !profile.trim().is_empty()), - "compact tcgen05 load admission requires both backend evidence profiles" - ); - let expected_ld = TCGEN05_LD_VARIANTS - .into_iter() - .flat_map(|(shape, multiplicity)| { - [false, true] - .into_iter() - .map(move |pack16| (shape, multiplicity, pack16)) - }) - .collect::>(); - ensure!( - admission - .ld_variants - .iter() - .map(|variant| (variant.shape, variant.multiplicity, variant.pack16)) - .eq(expected_ld), - "compact tcgen05 load admission must list all 58 variants in canonical order" - ); - let base = records - .iter() - .find(|record| record.id == "tcgen05_ld_16x256b_pure") - .expect("closed tcgen05 load base") - .clone(); - let first_load = records.len(); - for variant in &admission.ld_variants { - validate_abi_id(&variant.abi_id)?; - records.push(materialize_tcgen05_ld_variant(&base, admission, variant)); - } - for pair in records[first_load..].chunks_exact(2) { - let raw = pair[0].tcgen05.as_ref().and_then(|tcgen05| tcgen05.ld); - let packed = pair[1].tcgen05.as_ref().and_then(|tcgen05| tcgen05.ld); - ensure!( - raw.is_some_and(|ld| !ld.pack16) - && packed.is_some_and(|ld| ld.pack16) - && raw.map(|mut ld| { - ld.pack16 = true; - ld - }) == packed - && pair[0].source_record == pair[1].source_record - && pair[0].llvm_symbol == pair[1].llvm_symbol, - "tcgen05 load source sharing must pair raw and pack16 leaves" - ); - } - } - - if admission.st_variants.is_empty() { - ensure!( - admission.st_llvm_evidence_profile.is_none() - && admission.st_libnvvm_evidence_profile.is_none(), - "tcgen05 store evidence profiles require admitted store variants" - ); - } else { - ensure!( - admission - .st_llvm_evidence_profile - .as_deref() - .is_some_and(|profile| !profile.trim().is_empty()) - && admission - .st_libnvvm_evidence_profile - .as_deref() - .is_some_and(|profile| !profile.trim().is_empty()), - "compact tcgen05 store admission requires both backend evidence profiles" - ); - let expected_st = TCGEN05_ST_VARIANTS - .into_iter() - .flat_map(|(shape, multiplicity)| { - [false, true] - .into_iter() - .map(move |unpack16| (shape, multiplicity, unpack16)) - }) - .collect::>(); - ensure!( - admission - .st_variants - .iter() - .map(|variant| (variant.shape, variant.multiplicity, variant.unpack16)) - .eq(expected_st), - "compact tcgen05 store admission must list all 58 variants in canonical order" - ); - let base = records - .iter() - .find(|record| record.id == "tcgen05_ld_16x256b_pure") - .expect("closed tcgen05 store base") - .clone(); - let first_store = records.len(); - for variant in &admission.st_variants { - validate_abi_id(&variant.abi_id)?; - records.push(materialize_tcgen05_st_variant(&base, admission, variant)); - } - for pair in records[first_store..].chunks_exact(2) { - let raw = pair[0].tcgen05.as_ref().and_then(|tcgen05| tcgen05.st); - let unpacked = pair[1].tcgen05.as_ref().and_then(|tcgen05| tcgen05.st); - ensure!( - raw.is_some_and(|st| !st.unpack16) - && unpacked.is_some_and(|st| st.unpack16) - && raw.map(|mut st| { - st.unpack16 = true; - st - }) == unpacked - && pair[0].source_record == pair[1].source_record - && pair[0].llvm_symbol == pair[1].llvm_symbol, - "tcgen05 store source sharing must pair raw and unpack16 leaves" - ); - } - } - - if admission.ld_offset_variants.is_empty() && admission.st_offset_variants.is_empty() { - ensure!( - admission.offset_llvm_evidence_profile.is_none() - && admission.offset_libnvvm_evidence_profile.is_none(), - "tcgen05 offset evidence profiles require admitted offset load/store variants" - ); - } else { - ensure!( - admission - .offset_llvm_evidence_profile - .as_deref() - .is_some_and(|profile| !profile.trim().is_empty()) - && admission - .offset_libnvvm_evidence_profile - .as_deref() - .is_some_and(|profile| !profile.trim().is_empty()), - "compact tcgen05 offset admission requires both backend evidence profiles" - ); - let expected_ld = TCGEN05_OFFSET_LDST_VARIANTS - .into_iter() - .flat_map(|(shape, multiplicity)| { - [false, true] - .into_iter() - .map(move |pack16| (shape, multiplicity, pack16)) - }) - .collect::>(); - ensure!( - admission - .ld_offset_variants - .iter() - .map(|variant| (variant.shape, variant.multiplicity, variant.pack16)) - .eq(expected_ld), - "compact tcgen05 offset load admission must list all 16 variants in canonical order" - ); - let expected_st = TCGEN05_OFFSET_LDST_VARIANTS - .into_iter() - .flat_map(|(shape, multiplicity)| { - [false, true] - .into_iter() - .map(move |unpack16| (shape, multiplicity, unpack16)) - }) - .collect::>(); - ensure!( - admission - .st_offset_variants - .iter() - .map(|variant| (variant.shape, variant.multiplicity, variant.unpack16)) - .eq(expected_st), - "compact tcgen05 offset store admission must list all 16 variants in canonical order" - ); - - let base = records - .iter() - .find(|record| record.id == "tcgen05_ld_16x256b_pure") - .expect("closed tcgen05 offset load/store base") - .clone(); - let first_load = records.len(); - for variant in &admission.ld_offset_variants { - validate_abi_id(&variant.abi_id)?; - records.push(materialize_tcgen05_ld_variant(&base, admission, variant)); - } - for pair in records[first_load..].chunks_exact(2) { - let raw = pair[0].tcgen05.as_ref().and_then(|tcgen05| tcgen05.ld); - let packed = pair[1].tcgen05.as_ref().and_then(|tcgen05| tcgen05.ld); - ensure!( - raw.is_some_and(|ld| !ld.pack16) - && packed.is_some_and(|ld| ld.pack16) - && raw.map(|mut ld| { - ld.pack16 = true; - ld - }) == packed - && pair[0].source_record == pair[1].source_record - && pair[0].llvm_symbol == pair[1].llvm_symbol, - "tcgen05 offset load source sharing must pair raw and pack16 leaves" - ); - } - - let first_store = records.len(); - for variant in &admission.st_offset_variants { - validate_abi_id(&variant.abi_id)?; - records.push(materialize_tcgen05_st_variant(&base, admission, variant)); - } - for pair in records[first_store..].chunks_exact(2) { - let raw = pair[0].tcgen05.as_ref().and_then(|tcgen05| tcgen05.st); - let unpacked = pair[1].tcgen05.as_ref().and_then(|tcgen05| tcgen05.st); - ensure!( - raw.is_some_and(|st| !st.unpack16) - && unpacked.is_some_and(|st| st.unpack16) - && raw.map(|mut st| { - st.unpack16 = true; - st - }) == unpacked - && pair[0].source_record == pair[1].source_record - && pair[0].llvm_symbol == pair[1].llvm_symbol, - "tcgen05 offset store source sharing must pair raw and unpack16 leaves" - ); - } - } - - if admission.mma_variants.is_empty() { - ensure!( - admission.mma_llvm_evidence_profile.is_none() - && admission.mma_libnvvm_evidence_profile.is_none() - && admission.mma_llvm_target_contracts.is_empty() - && admission.mma_libnvvm_target_contracts.is_empty(), - "tcgen05 MMA profiles and target contracts require admitted variants" - ); - } else { - ensure!( - admission - .mma_llvm_evidence_profile - .as_deref() - .is_some_and(|profile| !profile.trim().is_empty()) - && admission - .mma_libnvvm_evidence_profile - .as_deref() - .is_some_and(|profile| !profile.trim().is_empty()), - "compact tcgen05 MMA admission requires both backend evidence profiles" - ); - let expected_variants = TCGEN05_MMA_FORMS - .into_iter() - .map(|form| (form, None)) - .chain( - TCGEN05_MMA_ALIASES - .into_iter() - .map(|alias| (Tcgen05MmaForm::WsTensor, Some(alias))), - ) - .chain( - TCGEN05_MMA_ALIASES - .into_iter() - .map(|alias| (Tcgen05MmaForm::Shared, Some(alias))), - ) - .collect::>(); - ensure!( - admission - .mma_variants - .iter() - .map(|variant| (variant.form, variant.alias)) - .eq(expected_variants), - "compact tcgen05 MMA admission must list all 24 APIs in canonical order" - ); - for variant in &admission.mma_variants { - validate_abi_id(&variant.abi_id)?; - } - let expected_llvm = expected_tcgen05_mma_target_contracts(IntrinsicBackend::LlvmNvptx); - let expected_libnvvm = expected_tcgen05_mma_target_contracts(IntrinsicBackend::LibNvvm); - ensure!( - admission.mma_llvm_target_contracts == expected_llvm - && admission.mma_libnvvm_target_contracts == expected_libnvvm, - "tcgen05 MMA target contracts changed from the reviewed backend matrices" - ); - let llvm_target = - resolve_target_contracts("tcgen05 MMA LLVM", &admission.mma_llvm_target_contracts)?; - let libnvvm_target = resolve_target_contracts( - "tcgen05 MMA libNVVM", - &admission.mma_libnvvm_target_contracts, - )?; - let fixed_selector = [TargetSelectorBinding { - name: "kind".into(), - value: "f8f6f4".into(), - }]; - let fixed_llvm = resolve_target_contract( - "tcgen05 MMA LLVM alias", - &fixed_selector, - &admission.mma_llvm_target_contracts, - )?; - let fixed_libnvvm = resolve_target_contract( - "tcgen05 MMA libNVVM alias", - &fixed_selector, - &admission.mma_libnvvm_target_contracts, - )?; - for variant in &admission.mma_variants { - let (llvm, libnvvm) = if variant.alias.is_some() { - (&fixed_llvm, &fixed_libnvvm) - } else { - (&llvm_target, &libnvvm_target) - }; - records.push(materialize_tcgen05_mma_variant( - admission, variant, llvm, libnvvm, - )); - } - } - Ok(records) -} - -fn validate_tcgen05_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, -) -> Result<()> { - let tcgen05 = policy - .tcgen05 - .as_ref() - .with_context(|| format!("{} has no closed tcgen05 contract", policy.id))?; - if let Some(mma) = &tcgen05.mma { - return validate_tcgen05_mma_policy(policy, declaration, tcgen05, mma); - } - if let Some(ld) = tcgen05.ld { - return validate_tcgen05_ld_policy(policy, declaration, tcgen05, ld); - } - if let Some(st) = tcgen05.st { - return validate_tcgen05_st_policy(policy, declaration, tcgen05, st); - } - if let Some(cp) = tcgen05.cp { - return validate_tcgen05_cp_policy(policy, declaration, tcgen05, cp); - } - ensure!( - !matches!( - tcgen05.operation, - Tcgen05Operation::Ld | Tcgen05Operation::St | Tcgen05Operation::Mma - ), - "{} has a tcgen05 load/store/MMA operation without its closed identity", - policy.id - ); - let recipe = tcgen05_recipe(tcgen05.operation); - ensure!( - policy.id == recipe.id - && policy.abi_id == recipe.abi_id - && policy.operation_key == recipe.operation_key - && policy.source.is_none() - && policy.source_record.as_deref() == Some(recipe.source_record) - && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) - && policy.resolved_llvm_symbol.is_none() - && declaration.source_record == recipe.source_record - && declaration.llvm_name == recipe.llvm_symbol, - "{} tcgen05 identity changed", - policy.id - ); - ensure!( - policy.rust_module == "tcgen05" - && policy.rust_name == recipe.id - && policy.rust_arguments == recipe.rust_arguments - && policy.rust_result == recipe.rust_result - && policy.safe == recipe.safe - && !policy.must_use - && policy.safe_allowlist_reason.as_deref() == recipe.safe_reason - && policy.public_rust_path == format!("cuda_intrinsics::tcgen05::{}", recipe.id) - && policy.compatibility_rust_paths == [format!("cuda_device::tcgen05::{}", recipe.id)], - "{} tcgen05 Rust API changed", - policy.id - ); - ensure!( - policy.dialect_op_type == recipe.dialect_op_type - && policy.dialect_op_name == recipe.dialect_op_name - && policy.dialect_operands == recipe.dialect_operands - && policy.dialect_results == recipe.dialect_results - && policy.llvm_arguments == recipe.llvm_arguments - && policy.llvm_results == recipe.llvm_results - && declaration.arguments == recipe.llvm_arguments - && declaration.results == recipe.llvm_results - && declaration.classes == recipe.imported_classes - && declaration.properties == recipe.imported_properties - && policy.lowering == "generated_tcgen05", - "{} tcgen05 carrier or imported declaration changed", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == recipe.memory - && policy.convergent - && policy.execution_scope == recipe.operation.execution_scope() - && tcgen05.ld.is_none() - && tcgen05.st.is_none() - && tcgen05.adapter == recipe.adapter - && tcgen05.source_contract == recipe.source_contract - && tcgen05.runtime_validation == RuntimeValidation::Unexecuted, - "{} tcgen05 semantics changed", - policy.id - ); - ensure!( - policy.minimum_ptx == "8.6" - && policy.minimum_sm.is_none() - && policy.targets == TCGEN05_LLVM_TARGETS - && policy.ptx_isa_version == "8.6" - && policy.ptx_result == recipe.rust_result - && policy.expected_ptx.mnemonic == "tcgen05" - && policy.expected_ptx.modifiers == recipe.modifiers - && policy.expected_ptx.operands == recipe.operands, - "{} tcgen05 target or PTX contract changed", - policy.id - ); - validate_tcgen05_backend_routes(policy, "tcgen05")?; - validate_tcgen05_source_contract(&recipe, declaration)?; - ensure_no_other_family_contract(policy, "tcgen05")?; - Ok(()) -} - -fn validate_tcgen05_mma_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, - tcgen05: &Tcgen05, - mma: &Tcgen05Mma, -) -> Result<()> { - let form = mma.form; - let alias = mma.alias; - ensure!( - match alias { - Some(alias) => { - matches!(form, Tcgen05MmaForm::Shared | Tcgen05MmaForm::WsTensor) - && TCGEN05_MMA_ALIASES.contains(&alias) - } - None => TCGEN05_MMA_FORMS.contains(&form), - }, - "{} has an unsupported tcgen05 MMA identity", - policy.id - ); - let id = tcgen05_mma_public_id(form, alias); - let operation_key = tcgen05_mma_operation_key(form, alias); - let llvm_arguments = tcgen05_mma_llvm_arguments(form); - let expected_fixed = alias.map(|_| Tcgen05MmaFixedSelectors { - kind: Tcgen05MmaKind::F8f6f4, - b_buffer: 0, - b_usage: Tcgen05MmaBUsage::Discard, - }); - ensure!( - policy.id == id - && policy.operation_key == operation_key - && policy.source.is_none() - && policy.source_record.as_deref() == Some(tcgen05_mma_source_record(form).as_str()) - && policy.llvm_symbol.as_deref() == Some(tcgen05_mma_llvm_symbol(form).as_str()) - && policy.resolved_llvm_symbol.is_none() - && declaration.source_record == tcgen05_mma_source_record(form) - && declaration.llvm_name == tcgen05_mma_llvm_symbol(form), - "{} tcgen05 MMA identity changed", - policy.id - ); - ensure!( - policy.rust_module == "tcgen05" - && policy.rust_name == id - && policy.rust_arguments == tcgen05_mma_rust_arguments(form, alias) - && policy.rust_result == "()" - && !policy.safe - && !policy.must_use - && policy.safe_allowlist_reason.is_none() - && policy.public_rust_path == format!("cuda_intrinsics::tcgen05::{id}") - && policy.compatibility_rust_paths - == [if alias.is_some() { - format!("cuda_device::tcgen05::{id}") - } else { - format!("cuda_device::tcgen05::__{id}") - }], - "{} tcgen05 MMA Rust API changed", - policy.id - ); - ensure!( - policy.dialect_op_type == TCGEN05_MMA_DIALECT_OP_TYPE - && policy.dialect_op_name == TCGEN05_MMA_DIALECT_OP_NAME - && policy.dialect_operands == tcgen05_mma_dialect_operands(form, alias) - && policy.dialect_results.is_empty() - && policy.llvm_arguments == llvm_arguments - && policy.llvm_results.is_empty() - && declaration.arguments == llvm_arguments - && declaration.results.is_empty() - && declaration.classes - == [ - "SDPatternOperator", - "Intrinsic", - "DefaultAttrsIntrinsic", - "DefaultAttrsIntrinsicFlags", - ] - && declaration.properties == tcgen05_mma_imported_properties(form) - && policy.lowering == "generated_tcgen05_mma", - "{} tcgen05 MMA carrier or imported declaration changed", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == "read_write" - && policy.convergent - && policy.execution_scope == "thread" - && tcgen05.operation == Tcgen05Operation::Mma - && tcgen05.cp.is_none() - && tcgen05.ld.is_none() - && tcgen05.st.is_none() - && mma.selector_layout == tcgen05_mma_selector_layout(form) - && mma.fixed_selectors == expected_fixed - && tcgen05.adapter == tcgen05_mma_adapter(form, alias) - && tcgen05.source_contract == Tcgen05SourceContract::TablegenSelectionChangesPtx - && tcgen05.runtime_validation == RuntimeValidation::Unexecuted, - "{} tcgen05 MMA semantics or selector contract changed", - policy.id - ); - - let expected_llvm_contracts = - expected_tcgen05_mma_target_contracts(IntrinsicBackend::LlvmNvptx); - let expected_libnvvm_contracts = - expected_tcgen05_mma_target_contracts(IntrinsicBackend::LibNvvm); - let selected = [TargetSelectorBinding { - name: "kind".into(), - value: "f8f6f4".into(), - }]; - let expected_llvm = if alias.is_some() { - resolve_target_contract( - "tcgen05 MMA LLVM alias", - &selected, - &expected_llvm_contracts, - )? - } else { - resolve_target_contracts("tcgen05 MMA LLVM", &expected_llvm_contracts)? - }; - let expected_libnvvm = if alias.is_some() { - resolve_target_contract( - "tcgen05 MMA libNVVM alias", - &selected, - &expected_libnvvm_contracts, - )? - } else { - resolve_target_contracts("tcgen05 MMA libNVVM", &expected_libnvvm_contracts)? - }; - ensure!( - mma.llvm_target == expected_llvm - && mma.libnvvm_target == expected_libnvvm - && policy.minimum_ptx == "8.6" - && policy.minimum_sm.is_none() - && policy.targets == TCGEN05_LLVM_TARGETS - && policy.ptx_result == "()" - && policy.ptx_isa_version == "8.6" - && policy.expected_ptx == tcgen05_mma_expected_ptx(form, alias), - "{} tcgen05 MMA target or PTX contract changed", - policy.id - ); - let backend_pairs = policy - .backend_lowerings - .iter() - .map(|route| (route.backend, route.mechanism)) - .collect::>(); - ensure!( - policy.backend_lowerings.len() == 2 - && backend_pairs - == BTreeSet::from([ - ( - IntrinsicBackend::LlvmNvptx, - BackendLoweringMechanism::InlinePtx, - ), - ( - IntrinsicBackend::LibNvvm, - BackendLoweringMechanism::InlinePtx, - ), - ]) - && policy.backend_lowerings.iter().all(|route| { - !route.evidence_profile.trim().is_empty() - && route.targets.is_none() - && route.minimum_ptx.is_none() - && route.minimum_sm.is_none() - }), - "{} tcgen05 MMA backend routes changed", - policy.id - ); - - let expected_all = tcgen05_mma_all_selection_asms(form); - let expected_valid = if alias.is_some() { - BTreeSet::from([if tcgen05_mma_is_ws(form) { - tcgen05_mma_selection_asm( - form, - Tcgen05MmaKind::F8f6f4, - 1, - None, - Some(0), - Some(Tcgen05MmaBUsage::Discard), - ) - } else { - tcgen05_mma_selection_asm(form, Tcgen05MmaKind::F8f6f4, 1, Some("discard"), None, None) - }]) - } else { - tcgen05_mma_valid_selection_asms(form) - }; - let mut source_records = BTreeSet::new(); - let actual_all = declaration - .selections - .iter() - .map(|selection| { - ensure!( - !selection.source_record.is_empty() - && source_records.insert(selection.source_record.as_str()) - && selection.constraints.is_empty(), - "{} tcgen05 MMA selection provenance changed", - policy.id - ); - let predicate = if selection.asm.contains(".kind::i8.") { - "Subtarget->hasTcgen05MMAI8Kind()" - } else { - "Subtarget->hasTcgen05InstSupport()" - }; - ensure!( - selection.predicates == [predicate], - "{} tcgen05 MMA selection predicate changed", - policy.id - ); - Ok(selection.asm.clone()) - }) - .collect::>>()?; - ensure!( - actual_all == expected_all - && declaration.selections.len() == expected_all.len() - && expected_valid.is_subset(&expected_all) - && expected_valid.len() - == if alias.is_some() { - 1 - } else if tcgen05_mma_is_ws(form) { - 64 - } else if tcgen05_mma_is_ashift(form) { - 16 - } else { - 32 - }, - "{} tcgen05 MMA selection matrix changed", - policy.id - ); - ensure_no_other_family_contract(policy, "tcgen05 MMA")?; - Ok(()) -} - -fn validate_tcgen05_ld_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, - tcgen05: &Tcgen05, - ld: Tcgen05Ld, -) -> Result<()> { - let has_half_split_offset = ld.shape == Tcgen05LdShape::M16x32bx2; - let variants: &[(Tcgen05LdShape, Tcgen05LdMultiplicity)] = if has_half_split_offset { - &TCGEN05_OFFSET_LDST_VARIANTS - } else { - &TCGEN05_LD_VARIANTS - }; - ensure!( - variants.contains(&(ld.shape, ld.multiplicity)), - "{} has an unsupported tcgen05 load identity", - policy.id - ); - let id = tcgen05_ld_id(ld); - let source_record = tcgen05_ld_source_record(ld); - let llvm_symbol = tcgen05_ld_llvm_symbol(ld); - let rust_result = tcgen05_ld_rust_result(ld); - let llvm_result = tcgen05_ld_llvm_result(ld); - let register_count = tcgen05_ld_register_count(ld); - let mode = if ld.pack16 { "pack16" } else { "raw" }; - let expected_rust_arguments = if has_half_split_offset { - vec!["u32", "i64"] - } else { - vec!["u32"] - }; - let compatibility_name = if has_half_split_offset { - format!("__{id}") - } else { - id.clone() - }; - let expected_dialect_operands = if has_half_split_offset { - vec!["i32", "i64"] - } else { - vec!["i32"] - }; - let expected_llvm_arguments = if has_half_split_offset { - vec!["tmem_ptr", "i64", "i1"] - } else { - vec!["tmem_ptr", "i1"] - }; - let expected_properties = if has_half_split_offset { - vec![ - "ImmArg", - "ImmArg", - "IntrArgMemOnly", - "IntrConvergent", - "NoCapture", - ] - } else { - vec![ - "ImmArg", - "IntrArgMemOnly", - "IntrConvergent", - "NoCapture", - ] - }; - ensure!( - policy.id == id - && policy.operation_key - == format!( - "tcgen05.ld.{}.{}.{}", - tcgen05_ld_shape_name(ld.shape), - tcgen05_ld_multiplicity_name(ld.multiplicity), - mode - ) - && policy.source.is_none() - && policy.source_record.as_deref() == Some(source_record.as_str()) - && policy.llvm_symbol.as_deref() == Some(llvm_symbol.as_str()) - && policy.resolved_llvm_symbol.is_none() - && declaration.source_record == source_record - && declaration.llvm_name == llvm_symbol, - "{} tcgen05 load identity changed", - policy.id - ); - ensure!( - policy.rust_module == "tcgen05" - && policy.rust_name == id - && policy.rust_arguments == expected_rust_arguments - && policy.rust_result == rust_result - && !policy.safe - && policy.must_use - && policy.safe_allowlist_reason.is_none() - && policy.public_rust_path == format!("cuda_intrinsics::tcgen05::{id}") - && policy.compatibility_rust_paths - == [format!("cuda_device::tcgen05::{compatibility_name}")], - "{} tcgen05 load Rust API changed", - policy.id - ); - ensure!( - policy.dialect_op_type == tcgen05_ld_op_type(ld) - && policy.dialect_op_name == format!("nvvm.{id}") - && policy.dialect_operands == expected_dialect_operands - && policy.dialect_results == vec!["i32"; register_count] - && policy.llvm_arguments == expected_llvm_arguments - && policy.llvm_results == [llvm_result.as_str()] - && declaration.arguments == expected_llvm_arguments - && declaration.results == [llvm_result.as_str()] - && declaration.classes == ["SDPatternOperator", "Intrinsic", "NVVM_TCGEN05_LD"] - && declaration.properties == expected_properties - && declaration.selections.is_empty() - && policy.lowering == "generated_tcgen05", - "{} tcgen05 load carrier or imported declaration changed", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == "read" - && policy.convergent - && policy.execution_scope == Tcgen05Operation::Ld.execution_scope() - && tcgen05.operation == Tcgen05Operation::Ld - && tcgen05.cp.is_none() - && tcgen05.ld == Some(ld) - && tcgen05.st.is_none() - && tcgen05.adapter - == if has_half_split_offset { - Tcgen05Adapter::TmemHalfSplitOffsetInjectPack16ToU32Registers - } else { - Tcgen05Adapter::TmemInjectPack16ToU32Registers - } - && tcgen05.source_contract == Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection - && tcgen05.runtime_validation == RuntimeValidation::Unexecuted, - "{} tcgen05 load semantics changed", - policy.id - ); - ensure!( - policy.minimum_ptx == "8.6" - && policy.minimum_sm.is_none() - && policy.targets == TCGEN05_LLVM_TARGETS - && policy.ptx_isa_version == "8.6" - && policy.ptx_result == rust_result - && policy.expected_ptx.mnemonic == "tcgen05" - && policy.expected_ptx.modifiers == tcgen05_ld_modifiers(ld) - && policy.expected_ptx.operands == tcgen05_ld_operands(ld), - "{} tcgen05 load target or PTX contract changed", - policy.id - ); - validate_tcgen05_backend_routes(policy, "tcgen05 load")?; - ensure_no_other_family_contract(policy, "tcgen05 load")?; - Ok(()) -} - -fn validate_tcgen05_st_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, - tcgen05: &Tcgen05, - st: Tcgen05St, -) -> Result<()> { - let has_half_split_offset = st.shape == Tcgen05LdShape::M16x32bx2; - let variants: &[(Tcgen05LdShape, Tcgen05LdMultiplicity)] = if has_half_split_offset { - &TCGEN05_OFFSET_LDST_VARIANTS - } else { - &TCGEN05_ST_VARIANTS - }; - ensure!( - variants.contains(&(st.shape, st.multiplicity)), - "{} has an unsupported tcgen05 store identity", - policy.id - ); - let id = tcgen05_st_id(st); - let source_record = tcgen05_st_source_record(st); - let llvm_symbol = tcgen05_st_llvm_symbol(st); - let rust_data = tcgen05_st_rust_data(st); - let llvm_data = tcgen05_st_llvm_data(st); - let register_count = tcgen05_st_register_count(st); - let mode = if st.unpack16 { "unpack16" } else { "raw" }; - let expected_rust_arguments = if has_half_split_offset { - vec!["u32", "i64", rust_data.as_str()] - } else { - vec!["u32", rust_data.as_str()] - }; - let compatibility_name = if has_half_split_offset { - format!("__{id}") - } else { - id.clone() - }; - let expected_dialect_operands = std::iter::once("i32".to_owned()) - .chain(has_half_split_offset.then(|| "i64".to_owned())) - .chain(std::iter::repeat_n("i32".to_owned(), register_count)) - .collect::>(); - let expected_llvm_arguments = if has_half_split_offset { - vec!["tmem_ptr", "i64", llvm_data.as_str(), "i1"] - } else { - vec!["tmem_ptr", llvm_data.as_str(), "i1"] - }; - let expected_properties = if has_half_split_offset { - vec![ - "ImmArg", - "ImmArg", - "IntrArgMemOnly", - "IntrConvergent", - "NoCapture", - ] - } else { - vec![ - "ImmArg", - "IntrArgMemOnly", - "IntrConvergent", - "NoCapture", - ] - }; - ensure!( - policy.id == id - && policy.operation_key - == format!( - "tcgen05.st.{}.{}.{}", - tcgen05_ld_shape_name(st.shape), - tcgen05_ld_multiplicity_name(st.multiplicity), - mode - ) - && policy.source.is_none() - && policy.source_record.as_deref() == Some(source_record.as_str()) - && policy.llvm_symbol.as_deref() == Some(llvm_symbol.as_str()) - && policy.resolved_llvm_symbol.is_none() - && declaration.source_record == source_record - && declaration.llvm_name == llvm_symbol, - "{} tcgen05 store identity changed", - policy.id - ); - ensure!( - policy.rust_module == "tcgen05" - && policy.rust_name == id - && policy.rust_arguments == expected_rust_arguments - && policy.rust_result == "()" - && !policy.safe - && !policy.must_use - && policy.safe_allowlist_reason.is_none() - && policy.public_rust_path == format!("cuda_intrinsics::tcgen05::{id}") - && policy.compatibility_rust_paths - == [format!("cuda_device::tcgen05::{compatibility_name}")], - "{} tcgen05 store Rust API changed", - policy.id - ); - ensure!( - policy.dialect_op_type == tcgen05_st_op_type(st) - && policy.dialect_op_name == format!("nvvm.{id}") - && policy.dialect_operands == expected_dialect_operands - && policy.dialect_results.is_empty() - && policy.llvm_arguments == expected_llvm_arguments - && policy.llvm_results.is_empty() - && declaration.arguments == expected_llvm_arguments - && declaration.results.is_empty() - && declaration.classes == ["SDPatternOperator", "Intrinsic", "NVVM_TCGEN05_ST"] - && declaration.properties == expected_properties - && declaration.selections.is_empty() - && policy.lowering == "generated_tcgen05", - "{} tcgen05 store carrier or imported declaration changed", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == "write" - && policy.convergent - && policy.execution_scope == Tcgen05Operation::St.execution_scope() - && tcgen05.operation == Tcgen05Operation::St - && tcgen05.cp.is_none() - && tcgen05.ld.is_none() - && tcgen05.st == Some(st) - && tcgen05.adapter - == if has_half_split_offset { - Tcgen05Adapter::TmemHalfSplitOffsetU32RegistersInjectUnpack16ToVoid - } else { - Tcgen05Adapter::TmemU32RegistersInjectUnpack16ToVoid - } - && tcgen05.source_contract == Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection - && tcgen05.runtime_validation == RuntimeValidation::Unexecuted, - "{} tcgen05 store semantics changed", - policy.id - ); - ensure!( - policy.minimum_ptx == "8.6" - && policy.minimum_sm.is_none() - && policy.targets == TCGEN05_LLVM_TARGETS - && policy.ptx_isa_version == "8.6" - && policy.ptx_result == "()" - && policy.expected_ptx.mnemonic == "tcgen05" - && policy.expected_ptx.modifiers == tcgen05_st_modifiers(st) - && policy.expected_ptx.operands == tcgen05_st_operands(st), - "{} tcgen05 store target or PTX contract changed", - policy.id - ); - validate_tcgen05_backend_routes(policy, "tcgen05 store")?; - ensure_no_other_family_contract(policy, "tcgen05 store")?; - Ok(()) -} - -fn validate_tcgen05_cp_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, - tcgen05: &Tcgen05, - cp: Tcgen05Cp, -) -> Result<()> { - ensure!( - TCGEN05_CP_MEMBERS.contains(&cp.member), - "{} has an unsupported tcgen05 copy identity", - policy.id - ); - let recipe = tcgen05_cp_member_recipe(cp.member); - let group = match cp.group { - Tcgen05CpGroup::Cg1 => 1, - Tcgen05CpGroup::Cg2 => 2, - }; - let group_suffix = if group == 1 { "" } else { "_cg2" }; - let id_suffix = recipe.llvm_suffix.replace('.', "_"); - let id = format!("tcgen05_cp_{id_suffix}{group_suffix}"); - let operation = if group == 1 { - Tcgen05Operation::CpSmemToTmem - } else { - Tcgen05Operation::CpSmemToTmemCg2 - }; - let op_type = format!( - "Tcgen05Cp{}{}Op", - recipe.op_suffix, - if group == 1 { "" } else { "Cg2" } - ); - let source_record = format!("int_nvvm_tcgen05_cp_{}_cg{group}", id_suffix); - let llvm_symbol = format!("llvm.nvvm.tcgen05.cp.{}.cg{group}", recipe.llvm_suffix); - let modifiers = std::iter::once("cp".into()) - .chain(std::iter::once(format!("cta_group::{group}"))) - .chain(recipe.ptx_suffix.split('.').map(Into::into)) - .collect::>(); - ensure!( - policy.id == id - && policy.operation_key == format!("tcgen05.cp.{}.cg{group}", recipe.llvm_suffix) - && policy.source.is_none() - && policy.source_record.as_deref() == Some(source_record.as_str()) - && policy.llvm_symbol.as_deref() == Some(llvm_symbol.as_str()) - && policy.resolved_llvm_symbol.is_none() - && declaration.source_record == source_record - && declaration.llvm_name == llvm_symbol, - "{} tcgen05 copy identity changed", - policy.id - ); - ensure!( - policy.rust_module == "tcgen05" - && policy.rust_name == id - && policy.rust_arguments == ["u32", "u64"] - && policy.rust_result == "()" - && !policy.safe - && !policy.must_use - && policy.safe_allowlist_reason.is_none() - && policy.public_rust_path == format!("cuda_intrinsics::tcgen05::{id}") - && policy.compatibility_rust_paths == [format!("cuda_device::tcgen05::{id}")], - "{} tcgen05 copy Rust API changed", - policy.id - ); - ensure!( - policy.dialect_op_type == op_type - && policy.dialect_op_name == format!("nvvm.{id}") - && policy.dialect_operands == ["i32", "i64"] - && policy.dialect_results.is_empty() - && policy.llvm_arguments == ["tmem_ptr", "i64"] - && policy.llvm_results.is_empty() - && declaration.arguments == ["tmem_ptr", "i64"] - && declaration.results.is_empty() - && declaration.classes == ["SDPatternOperator", "Intrinsic"] - && declaration.properties - == [ - "IntrConvergent", - "IntrInaccessibleMemOrArgMemOnly", - "NoCapture", - ] - && policy.lowering == "generated_tcgen05", - "{} tcgen05 copy carrier or declaration changed", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == "read_write" - && policy.convergent - && policy.execution_scope == "thread" - && tcgen05.operation == operation - && tcgen05.ld.is_none() - && tcgen05.st.is_none() - && tcgen05.adapter == Tcgen05Adapter::TmemDescriptorToVoid - && tcgen05.source_contract == Tcgen05SourceContract::ExactTablegenSelection - && tcgen05.runtime_validation == RuntimeValidation::Unexecuted, - "{} tcgen05 copy semantics changed", - policy.id - ); - ensure!( - policy.minimum_ptx == "8.6" - && policy.minimum_sm.is_none() - && policy.targets == TCGEN05_LLVM_TARGETS - && policy.ptx_isa_version == "8.6" - && policy.ptx_result == "()" - && policy.expected_ptx.mnemonic == "tcgen05" - && policy.expected_ptx.modifiers == modifiers - && policy.expected_ptx.operands == [OperandPattern::Address, OperandPattern::Register], - "{} tcgen05 copy target or PTX contract changed", - policy.id - ); - validate_tcgen05_backend_routes(policy, "tcgen05 copy")?; - ensure!( - declaration.selections.len() == 1, - "{} must keep one exact tcgen05 copy selection", - policy.id - ); - let selection = &declaration.selections[0]; - ensure!( - selection.source_record == format!("TCGEN05_CP_{}_cg{group}", recipe.selection_stem) - && selection.asm - == format!( - "tcgen05.cp.cta_group::{group}.{} \t[$tmem_addr], $sdesc;", - recipe.ptx_suffix - ) - && selection.predicates == ["Subtarget->hasTcgen05InstSupport()"] - && selection.constraints.is_empty(), - "{} exact tcgen05 copy selection changed", - policy.id - ); - ensure_no_other_family_contract(policy, "tcgen05 copy")?; - Ok(()) -} - -fn validate_tcgen05_backend_routes(policy: &OverlayIntrinsic, family: &str) -> Result<()> { - ensure!( - policy.backend_lowerings.len() == 2 - && policy.backend_lowerings[0].backend == IntrinsicBackend::LlvmNvptx - && policy.backend_lowerings[1].backend == IntrinsicBackend::LibNvvm - && policy.backend_lowerings[0].targets.is_none() - && policy.backend_lowerings[1].targets.as_deref() == Some(TCGEN05_LIBNVVM_TARGETS) - && policy.backend_lowerings.iter().all(|route| { - route.mechanism == BackendLoweringMechanism::InlinePtx - && route.minimum_ptx.as_deref() == Some("8.6") - && route.minimum_sm.is_none() - && !route.evidence_profile.trim().is_empty() - }), - "{} {family} backend route changed", - policy.id - ); - Ok(()) -} - -fn validate_tcgen05_source_contract( - recipe: &Tcgen05Recipe, - declaration: &ImportedIntrinsic, -) -> Result<()> { - match recipe.source_contract { - Tcgen05SourceContract::ExactTablegenSelection => { - ensure!( - declaration.selections.len() == 1, - "{} must keep one exact tcgen05 selection", - recipe.id - ); - let selection = &declaration.selections[0]; - let expected_predicate = if matches!( - recipe.operation, - Tcgen05Operation::ShiftDown | Tcgen05Operation::ShiftDownCg2 - ) { - "Subtarget->hasTcgen05ShiftSupport()" - } else { - "Subtarget->hasTcgen05InstSupport()" - }; - ensure!( - recipe.selection_record == Some(selection.source_record.as_str()) - && recipe.selection_asm == Some(selection.asm.as_str()) - && selection.predicates == [expected_predicate] - && selection.constraints.is_empty(), - "{} exact tcgen05 selection changed", - recipe.id - ); - } - Tcgen05SourceContract::TablegenSelectionChangesPtx - if matches!( - recipe.operation, - Tcgen05Operation::Commit | Tcgen05Operation::CommitCg2 - ) => - { - ensure!( - declaration.selections.len() == 1, - "{} must keep one canonical commit selection", - recipe.id - ); - let selection = &declaration.selections[0]; - ensure!( - recipe.selection_record == Some(selection.source_record.as_str()) - && recipe.selection_asm == Some(selection.asm.as_str()) - && selection.predicates == ["Subtarget->hasTcgen05InstSupport()"] - && selection.constraints.is_empty() - && !recipe.operands.is_empty(), - "{} canonical commit selection changed", - recipe.id - ); - } - Tcgen05SourceContract::TablegenSelectionChangesPtx => { - ensure!( - declaration.selections.len() == 64, - "{} must keep the 64 collector selections", - recipe.id - ); - let actual = declaration - .selections - .iter() - .map(|selection| { - ensure!( - selection.constraints.is_empty(), - "{} collector selection gained constraints", - recipe.id - ); - let predicate = if selection.asm.contains(".kind::i8.") { - "Subtarget->hasTcgen05MMAI8Kind()" - } else { - "Subtarget->hasTcgen05InstSupport()" - }; - ensure!( - selection.predicates == [predicate], - "{} collector predicate changed", - recipe.id - ); - Ok(selection.asm.clone()) - }) - .collect::>>()?; - let expected = ["f16", "tf32", "f8f6f4", "i8"] - .into_iter() - .flat_map(|kind| { - ["b0", "b1", "b2", "b3"].into_iter().flat_map(move |collector| { - ["discard", "fill", "use", "lastuse"].into_iter().map(move |action| { - format!("tcgen05.mma.ws.cta_group::1.kind::{kind}.collector::{collector}::{action} [$dtmem], [$a], $b, $idesc, $enable_inp_d;") - }) - }) - }) - .collect::>(); - ensure!( - actual == expected, - "{} collector selection matrix changed", - recipe.id - ); - } - Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection => ensure!( - declaration.selections.is_empty(), - "{} unexpectedly gained a TableGen selection; review the backend route", - recipe.id - ), - } - Ok(()) -} - -fn tma_imported_properties(operation: TmaOperation) -> Vec { - assert_ne!(operation, TmaOperation::Reduce); - let dimensions = operation.dimensions(); - if matches!( - operation, - TmaOperation::G2sTile1d - | TmaOperation::G2sTile2d - | TmaOperation::G2sTile2dMulticast - | TmaOperation::G2sTile2dMulticastCg2 - | TmaOperation::G2sTile3d - | TmaOperation::G2sTile4d - | TmaOperation::G2sTile5d - ) { - let dimensions = dimensions.unwrap(); - let mut properties = vec![ - format!("ImmArg", dimensions + 5), - format!("ImmArg", dimensions + 6), - format!("ImmArg", dimensions + 7), - "IntrConvergent".into(), - format!("Range", dimensions + 7), - "ReadOnly".into(), - "WriteOnly".into(), - ]; - properties.sort(); - return properties; - } - if matches!( - operation, - TmaOperation::S2gTile1d - | TmaOperation::S2gTile2d - | TmaOperation::S2gTile3d - | TmaOperation::S2gTile4d - | TmaOperation::S2gTile5d - ) { - return vec![ - format!("ImmArg", dimensions.unwrap() + 3), - "IntrConvergent".into(), - "ReadOnly".into(), - "ReadOnly".into(), - ]; - } - if operation == TmaOperation::CommitGroup { - return vec![]; - } - if matches!( - operation, - TmaOperation::WaitGroup | TmaOperation::WaitGroupRead - ) { - return vec!["ImmArg".into()]; - } - if operation == TmaOperation::PrefetchTensorMap { - return vec![ - "IntrArgMemOnly".into(), - "NoCapture".into(), - "ReadOnly".into(), - ]; - } - if let Some(coordinate_count) = operation.prefetch_coordinate_count() { - return vec![ - format!("ImmArg", coordinate_count + 2), - "IntrConvergent".into(), - "ReadOnly".into(), - ]; - } - if matches!( - operation, - TmaOperation::FenceProxyTensorMapAcquireCluster - | TmaOperation::FenceProxyTensorMapAcquireCta - | TmaOperation::FenceProxyTensorMapAcquireGpu - | TmaOperation::FenceProxyTensorMapAcquireSystem - ) { - return vec![ - "ImmArg".into(), - "IntrArgMemOnly".into(), - "IntrNoCallback".into(), - "Range".into(), - ]; - } - if matches!( - operation, - TmaOperation::FenceProxyTensorMapReleaseCluster - | TmaOperation::FenceProxyTensorMapReleaseCta - | TmaOperation::FenceProxyTensorMapReleaseGpu - | TmaOperation::FenceProxyTensorMapReleaseSystem - ) { - return vec!["IntrNoCallback".into()]; - } - - let mut properties = vec![ - "IntrArgMemOnly".into(), - "IntrWriteMem".into(), - "NoCapture".into(), - ]; - match operation { - TmaOperation::ReplaceBoxDim - | TmaOperation::ReplaceElementStride - | TmaOperation::ReplaceGlobalDim - | TmaOperation::ReplaceGlobalStride => { - properties.push("ImmArg".into()); - } - TmaOperation::ReplaceElementType => { - properties.extend([ - "ArgInfo".into(), - "ImmArg".into(), - "Range".into(), - ]); - } - TmaOperation::ReplaceFillMode => { - properties.extend([ - "ArgInfo".into(), - "ImmArg".into(), - "Range".into(), - ]); - } - TmaOperation::ReplaceInterleaveLayout => { - properties.extend([ - "ArgInfo".into(), - "ImmArg".into(), - "Range".into(), - ]); - } - TmaOperation::ReplaceSwizzleAtomicity => { - properties.extend([ - "ArgInfo".into(), - "ImmArg".into(), - "Range".into(), - ]); - } - TmaOperation::ReplaceSwizzleMode => { - properties.extend([ - "ArgInfo".into(), - "ImmArg".into(), - "Range".into(), - ]); - } - TmaOperation::ReplaceGlobalAddress | TmaOperation::ReplaceRank => {} - _ => unreachable!("closed TMA operation property contract"), - } - properties.sort(); - properties -} - -fn tma_reduction_imported_properties(dimensions: usize) -> Vec { - vec![ - format!("ImmArg", dimensions + 3), - "IntrConvergent".into(), - "ReadOnly".into(), - "ReadOnly".into(), - ] -} - -fn validate_sync_policy(policy: &OverlayIntrinsic, declaration: &ImportedIntrinsic) -> Result<()> { - if let Some(scope) = threadfence_scope_for_id(&policy.id) { - return validate_threadfence_policy(policy, declaration, scope); - } - ensure!( - policy.id == "sync_threads" - && policy.abi_id == "i0034" - && policy.operation_key == "synchronization.cta.barrier.aligned.all" - && policy.source.is_none() - && policy.source_record.as_deref() == Some("int_nvvm_barrier_cta_sync_aligned_all") - && policy.llvm_symbol.as_deref() == Some("llvm.nvvm.barrier.cta.sync.aligned.all") - && policy.resolved_llvm_symbol.is_none(), - "{} sync identity does not match the closed sync_threads recipe", - policy.id - ); - ensure!( - policy.rust_module == "thread" - && policy.rust_name == "sync_threads" - && policy.rust_arguments.is_empty() - && policy.rust_result == "()" - && !policy.safe - && !policy.must_use - && policy.safe_allowlist_reason.is_none() - && policy.public_rust_path == "cuda_intrinsics::thread::sync_threads" - && policy.compatibility_rust_paths - == [ - "cuda_device::thread::sync_threads", - "cuda_device::sync_threads", - ], - "{} must preserve the unsafe sync_threads raw API and both cuda-device compatibility paths", - policy.id - ); - ensure!( - policy.dialect_op_type == "Barrier0Op" - && policy.dialect_op_name == "nvvm.barrier0" - && policy.dialect_operands.is_empty() - && policy.dialect_results.is_empty() - && policy.llvm_arguments == ["i32"] - && policy.llvm_results.is_empty() - && policy.lowering == "generated_sync_threads", - "{} is outside the fixed-zero sync_threads lowering recipe", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == "read_write" - && policy.convergent - && policy.execution_scope == "cta" - && policy.minimum_ptx == "1.0" - && policy.minimum_sm.is_none() - && policy.ptx_result == "()" - && policy.targets == "all", - "{} sync effects or native target floor disagree with the closed recipe", - policy.id - ); - ensure!( - policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section - == "9.7.14.1 Parallel Synchronization and Communication Instructions: bar, barrier" - && policy.ptx_isa_url - == "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-bar-barrier", - "{} sync PTX provenance disagrees with the reviewed recipe", - policy.id - ); - ensure!( - declaration.properties == ["IntrConvergent", "IntrNoCallback"], - "{} sync properties disagree with the imported LLVM declaration", - policy.id - ); - ensure!( - policy.ldmatrix_variant.is_none() - && policy.ldmatrix_safety.is_none() - && policy.ldmatrix_adapter.is_none() - && policy.packed_atomic.is_none() - && policy.redux.is_none() - && policy.vote.is_none() - && policy.active_mask.is_none() - && policy.warp_match.is_none() - && policy.warp_barrier.is_none() - && policy.warp_shuffle.is_none() - && policy.dot_product.is_none() - && policy.packed_alu.is_none() - && policy.packed_conversion.is_none() - && policy.cp_async_copy.is_none() - && policy.cp_async_control.is_none() - && policy.mbarrier_basic.is_none() - && policy.selected_address_space.is_none(), - "{} mixes another generated-family contract with sync", - policy.id - ); - ensure!( - policy.expected_ptx.mnemonic == "bar" - && policy.expected_ptx.modifiers == ["sync"] - && policy.expected_ptx.operands == [OperandPattern::Exact { value: "0".into() }], - "{} expected PTX does not match literal bar.sync 0", - policy.id - ); - - let backend_pairs: BTreeSet<_> = policy - .backend_lowerings - .iter() - .map(|lowering| (lowering.backend, lowering.mechanism)) - .collect(); - ensure!( - policy.backend_lowerings.len() == 2 - && backend_pairs - == BTreeSet::from([ - ( - IntrinsicBackend::LlvmNvptx, - BackendLoweringMechanism::TypedNvvm, - ), - ( - IntrinsicBackend::LibNvvm, - BackendLoweringMechanism::InlinePtx, - ), - ]), - "{} must define exactly the reviewed LLVM typed and libNVVM inline-PTX routes", - policy.id - ); - for lowering in &policy.backend_lowerings { - let floor_matches = match lowering.backend { - IntrinsicBackend::LlvmNvptx => { - lowering.mechanism == BackendLoweringMechanism::TypedNvvm - && lowering.minimum_ptx.as_deref() == Some("3.2") - && lowering.minimum_sm.as_deref() == Some("sm_20") - } - IntrinsicBackend::LibNvvm => { - lowering.mechanism == BackendLoweringMechanism::InlinePtx - && lowering.minimum_ptx.is_none() - && lowering.minimum_sm.as_deref() == Some("sm_75") - } - }; - ensure!( - floor_matches && !lowering.evidence_profile.trim().is_empty(), - "{} backend {:?} does not carry its reviewed sync profile floor", - policy.id, - lowering.backend - ); - } - Ok(()) -} - -fn validate_vote_policy(policy: &OverlayIntrinsic, declaration: &ImportedIntrinsic) -> Result<()> { - let vote = policy - .vote - .as_ref() - .with_context(|| format!("{} has no closed vote contract", policy.id))?; - let recipe = vote_recipe(vote.mode); - ensure!( - vote.participation - == VoteParticipation::ExecutingLaneNamedAllNamedLanesSameInstructionAndMask - && vote.legacy_pre_sm70 - == PreSm70MemberMaskRule::AllNamedLanesConvergedAndOnlyNamedLanesActive - && vote.adapter == VoteAdapter::DirectMaskPredicate - && vote.mask_encoding == MaskEncoding::RegisterOrImmediate, - "{} requests an unsupported vote participation, pre-sm70 rule, adapter, or mask encoding", - policy.id - ); - ensure!( - policy.id == recipe.id - && policy.abi_id == recipe.abi_id - && policy.operation_key == recipe.operation_key - && policy.source.is_none() - && policy.source_record.as_deref() == Some(recipe.source_record) - && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) - && policy.resolved_llvm_symbol.is_none(), - "{} vote identity does not match its closed mode recipe", - policy.id - ); - let expected_compatibility_paths: Vec = if recipe.has_compatibility_path { - vec![format!("cuda_device::warp::{}", recipe.rust_name)] - } else { - vec![] - }; - ensure!( - policy.rust_module == "warp" - && policy.rust_name == recipe.rust_name - && policy.rust_arguments == ["u32", "bool"] - && policy.rust_result == recipe.rust_result - && !policy.safe - && policy.must_use - && policy.safe_allowlist_reason.is_none() - && policy.public_rust_path == format!("cuda_intrinsics::warp::{}", recipe.rust_name) - && policy.compatibility_rust_paths == expected_compatibility_paths, - "{} must preserve its unsafe must-use vote raw API and reviewed compatibility path", - policy.id - ); - ensure!( - policy.dialect_op_type == recipe.dialect_op_type - && policy.dialect_op_name == recipe.dialect_op_name - && policy.dialect_operands == ["i32", "i1"] - && policy.dialect_results == [recipe.llvm_result] - && policy.llvm_arguments == ["i32", "i1"] - && policy.llvm_results == [recipe.llvm_result] - && policy.lowering == "generated_vote", - "{} is outside the closed two-operand vote lowering recipe", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == "inaccessible_read_write" - && policy.convergent - && policy.execution_scope == "warp" - && policy.minimum_ptx == "6.0" - && policy.minimum_sm.as_deref() == Some("sm_30") - && policy.ptx_result == recipe.rust_result - && policy.targets == "all", - "{} vote effects, carrier, or target floor disagree with its mode recipe", - policy.id - ); - ensure!( - policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section == "9.7.14.10 Warp Vote Instructions: vote.sync" - && policy.ptx_isa_url - == "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-vote-sync", - "{} vote PTX provenance disagrees with the reviewed recipe", - policy.id - ); - ensure!( - declaration.properties - == [ - "IntrConvergent", - "IntrInaccessibleMemOnly", - "IntrNoCallback", - ], - "{} vote memory and convergence effects disagree with the imported declaration", - policy.id - ); - ensure!( - policy.backend_lowerings.is_empty() - && policy.packed_atomic.is_none() - && policy.ldmatrix_variant.is_none() - && policy.ldmatrix_safety.is_none() - && policy.ldmatrix_adapter.is_none() - && policy.redux.is_none() - && policy.active_mask.is_none() - && policy.warp_match.is_none() - && policy.warp_barrier.is_none() - && policy.warp_shuffle.is_none() - && policy.dot_product.is_none() - && policy.packed_alu.is_none() - && policy.packed_conversion.is_none() - && policy.cp_async_copy.is_none() - && policy.cp_async_control.is_none() - && policy.mbarrier_basic.is_none() - && policy.selected_address_space.is_none(), - "{} mixes another generated-family contract with vote", - policy.id - ); - ensure!( - policy.expected_ptx.mnemonic == "vote" - && policy.expected_ptx.modifiers == ["sync", recipe.ptx_mode, recipe.ptx_type] - && policy.expected_ptx.operands - == [ - OperandPattern::Register, - OperandPattern::Register, - OperandPattern::RegisterOrImmediate, - ], - "{} expected PTX does not match its closed vote mode recipe", - policy.id - ); - - let expected_selection_records = - BTreeSet::from([recipe.immediate_selection, recipe.register_selection]); - let actual_selection_records: BTreeSet<_> = declaration - .selections - .iter() - .map(|selection| selection.source_record.as_str()) - .collect(); - ensure!( - declaration.selections.len() == 2 && actual_selection_records == expected_selection_records, - "{} vote declaration must contain exactly its immediate/register selection pair", - policy.id - ); - let expected_asm = format!( - "vote.sync.{}.{} \t$dest, $pred, $mask;", - recipe.ptx_mode, recipe.ptx_type - ); - for selection in &declaration.selections { - ensure!( - selection.asm == expected_asm - && selection.predicates - == [ - "Subtarget->getPTXVersion() >= 60", - "Subtarget->getSmVersion() >= 30", - ] - && selection.constraints.is_empty(), - "{} vote immediate/register selections disagree on PTX shape, target predicates, or constraints", - policy.id - ); - } - Ok(()) -} - -struct VoteRecipe { - id: &'static str, - abi_id: &'static str, - operation_key: &'static str, - source_record: &'static str, - llvm_symbol: &'static str, - rust_name: &'static str, - rust_result: &'static str, - dialect_op_type: &'static str, - dialect_op_name: &'static str, - llvm_result: &'static str, - ptx_mode: &'static str, - ptx_type: &'static str, - immediate_selection: &'static str, - register_selection: &'static str, - has_compatibility_path: bool, -} - -fn vote_recipe(mode: VoteMode) -> VoteRecipe { - match mode { - VoteMode::All => VoteRecipe { - id: "all_sync", - abi_id: "i0040", - operation_key: "warp.vote.sync.all.pred", - source_record: "int_nvvm_vote_all_sync", - llvm_symbol: "llvm.nvvm.vote.all.sync", - rust_name: "all_sync", - rust_result: "bool", - dialect_op_type: "VoteSyncAllOp", - dialect_op_name: "nvvm.vote_sync_all", - llvm_result: "i1", - ptx_mode: "all", - ptx_type: "pred", - immediate_selection: "VOTE_SYNC_ALLi", - register_selection: "VOTE_SYNC_ALLr", - has_compatibility_path: true, - }, - VoteMode::Any => VoteRecipe { - id: "any_sync", - abi_id: "i0041", - operation_key: "warp.vote.sync.any.pred", - source_record: "int_nvvm_vote_any_sync", - llvm_symbol: "llvm.nvvm.vote.any.sync", - rust_name: "any_sync", - rust_result: "bool", - dialect_op_type: "VoteSyncAnyOp", - dialect_op_name: "nvvm.vote_sync_any", - llvm_result: "i1", - ptx_mode: "any", - ptx_type: "pred", - immediate_selection: "VOTE_SYNC_ANYi", - register_selection: "VOTE_SYNC_ANYr", - has_compatibility_path: true, - }, - VoteMode::Ballot => VoteRecipe { - id: "ballot_sync", - abi_id: "i0042", - operation_key: "warp.vote.sync.ballot.b32", - source_record: "int_nvvm_vote_ballot_sync", - llvm_symbol: "llvm.nvvm.vote.ballot.sync", - rust_name: "ballot_sync", - rust_result: "u32", - dialect_op_type: "VoteSyncBallotOp", - dialect_op_name: "nvvm.vote_sync_ballot", - llvm_result: "i32", - ptx_mode: "ballot", - ptx_type: "b32", - immediate_selection: "VOTE_SYNC_BALLOTi", - register_selection: "VOTE_SYNC_BALLOTr", - has_compatibility_path: true, - }, - VoteMode::Uni => VoteRecipe { - id: "uni_sync", - abi_id: "i0043", - operation_key: "warp.vote.sync.uni.pred", - source_record: "int_nvvm_vote_uni_sync", - llvm_symbol: "llvm.nvvm.vote.uni.sync", - rust_name: "uni_sync", - rust_result: "bool", - dialect_op_type: "VoteSyncUniOp", - dialect_op_name: "nvvm.vote_sync_uni", - llvm_result: "i1", - ptx_mode: "uni", - ptx_type: "pred", - immediate_selection: "VOTE_SYNC_UNIi", - register_selection: "VOTE_SYNC_UNIr", - has_compatibility_path: false, - }, - } -} - -fn validate_active_mask_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, -) -> Result<()> { - let active_mask = policy - .active_mask - .as_ref() - .with_context(|| format!("{} has no closed active-mask contract", policy.id))?; - ensure!( - active_mask.observation == ActiveMaskObservation::ExecutingLanesAtInstruction - && active_mask.adapter == ActiveMaskAdapter::DirectZeroOperandMask, - "{} requests an unsupported active-mask observation or adapter", - policy.id - ); - ensure!( - policy.id == "active_mask" - && policy.abi_id == "i0044" - && policy.operation_key == "warp.active_mask" - && policy.source.is_none() - && policy.source_record.as_deref() == Some("int_nvvm_activemask") - && policy.llvm_symbol.as_deref() == Some("llvm.nvvm.activemask") - && policy.resolved_llvm_symbol.is_none(), - "{} active-mask identity does not match the closed recipe", - policy.id - ); - ensure!( - policy.rust_module == "warp" - && policy.rust_name == "active_mask" - && policy.rust_arguments.is_empty() - && policy.rust_result == "u32" - && policy.safe - && policy.must_use - && policy - .safe_allowlist_reason - .as_deref() - .is_some_and(|reason| !reason.is_empty()) - && policy.public_rust_path == "cuda_intrinsics::warp::active_mask" - && policy.compatibility_rust_paths == ["cuda_device::warp::active_mask"], - "{} must preserve its safe must-use raw and compatibility APIs", - policy.id - ); - ensure!( - policy.dialect_op_type == "ActiveMaskOp" - && policy.dialect_op_name == "nvvm.activemask" - && policy.dialect_operands.is_empty() - && policy.dialect_results == ["i32"] - && policy.llvm_arguments.is_empty() - && policy.llvm_results == ["i32"] - && policy.lowering == "generated_active_mask", - "{} is outside the closed zero-operand active-mask lowering recipe", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == "inaccessible_read_write" - && policy.convergent - && policy.execution_scope == "warp" - && policy.minimum_ptx == "6.2" - && policy.minimum_sm.as_deref() == Some("sm_30") - && policy.ptx_result == "u32" - && policy.targets == "all", - "{} active-mask effects or target floor disagree with the closed recipe", - policy.id - ); - ensure!( - policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section - == "9.7.14.12 Parallel Synchronization and Communication Instructions: activemask" - && policy.ptx_isa_url - == "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-activemask", - "{} active-mask PTX provenance disagrees with the reviewed recipe", - policy.id - ); - ensure!( - declaration.properties - == [ - "IntrConvergent", - "IntrHasSideEffects", - "IntrInaccessibleMemOnly", - "IntrNoCallback", - ] - && declaration.selections.len() == 1 - && declaration.selections[0].source_record == "ACTIVEMASK" - && declaration.selections[0].asm == "activemask.b32 \t$dest;" - && declaration.selections[0].predicates - == [ - "Subtarget->getSmVersion() >= 30", - "Subtarget->getPTXVersion() >= 62", - ] - && declaration.selections[0].constraints.is_empty(), - "{} active-mask declaration or selection facts changed", - policy.id - ); - ensure!( - policy.expected_ptx.mnemonic == "activemask" - && policy.expected_ptx.modifiers == ["b32"] - && policy.expected_ptx.operands == [OperandPattern::Register], - "{} expected PTX does not match activemask.b32", - policy.id - ); - ensure!( - policy.packed_atomic.is_none() - && policy.redux.is_none() - && policy.vote.is_none() - && policy.warp_match.is_none() - && policy.warp_barrier.is_none() - && policy.warp_shuffle.is_none() - && policy.dot_product.is_none() - && policy.packed_alu.is_none() - && policy.packed_conversion.is_none() - && policy.cp_async_copy.is_none() - && policy.cp_async_control.is_none() - && policy.mbarrier_basic.is_none() - && policy.ldmatrix_variant.is_none() - && policy.ldmatrix_safety.is_none() - && policy.ldmatrix_adapter.is_none() - && policy.selected_address_space.is_none(), - "{} mixes another generated-family contract with active_mask", - policy.id - ); - - let backend_pairs: BTreeSet<_> = policy - .backend_lowerings - .iter() - .map(|lowering| (lowering.backend, lowering.mechanism)) - .collect(); - ensure!( - policy.backend_lowerings.len() == 2 - && backend_pairs - == BTreeSet::from([ - ( - IntrinsicBackend::LlvmNvptx, - BackendLoweringMechanism::TypedNvvm, - ), - ( - IntrinsicBackend::LibNvvm, - BackendLoweringMechanism::InlinePtx, - ), - ]), - "{} must keep the LLVM typed and libNVVM inline-PTX routes explicit", - policy.id - ); - for lowering in &policy.backend_lowerings { - let floor_matches = match lowering.backend { - IntrinsicBackend::LlvmNvptx => { - lowering.mechanism == BackendLoweringMechanism::TypedNvvm - && lowering.minimum_ptx.as_deref() == Some("6.2") - && lowering.minimum_sm.as_deref() == Some("sm_30") - } - IntrinsicBackend::LibNvvm => { - lowering.mechanism == BackendLoweringMechanism::InlinePtx - && lowering.minimum_ptx.is_none() - && lowering.minimum_sm.as_deref() == Some("sm_75") - } - }; - ensure!( - floor_matches && !lowering.evidence_profile.trim().is_empty(), - "{} backend {:?} does not carry its reviewed active-mask floor", - policy.id, - lowering.backend - ); - } - Ok(()) -} - -struct WarpMatchRecipe { - id: &'static str, - abi_id: &'static str, - operation_key: &'static str, - source_record: &'static str, - llvm_symbol: &'static str, - rust_name: &'static str, - rust_value: &'static str, - llvm_value: &'static str, - dialect_op_type: &'static str, - dialect_op_name: &'static str, - ptx_mode: &'static str, - ptx_type: &'static str, - selections: [&'static str; 4], - adapter: WarpMatchAdapter, -} - -fn warp_match_recipe(mode: WarpMatchMode, width: WarpMatchValueWidth) -> WarpMatchRecipe { - match (mode, width) { - (WarpMatchMode::Any, WarpMatchValueWidth::B32) => WarpMatchRecipe { - id: "match_any_sync", - abi_id: "i0045", - operation_key: "warp.match.sync.any.b32", - source_record: "int_nvvm_match_any_sync_i32", - llvm_symbol: "llvm.nvvm.match.any.sync.i32", - rust_name: "match_any_sync", - rust_value: "u32", - llvm_value: "i32", - dialect_op_type: "MatchAnySyncI32Op", - dialect_op_name: "nvvm.match_any_sync_i32", - ptx_mode: "any", - ptx_type: "b32", - selections: [ - "MATCH_ANY_SYNC_32ii", - "MATCH_ANY_SYNC_32ir", - "MATCH_ANY_SYNC_32ri", - "MATCH_ANY_SYNC_32rr", - ], - adapter: WarpMatchAdapter::DirectMask, - }, - (WarpMatchMode::Any, WarpMatchValueWidth::B64) => WarpMatchRecipe { - id: "match_any_i64_sync", - abi_id: "i0046", - operation_key: "warp.match.sync.any.b64", - source_record: "int_nvvm_match_any_sync_i64", - llvm_symbol: "llvm.nvvm.match.any.sync.i64", - rust_name: "match_any_i64_sync", - rust_value: "u64", - llvm_value: "i64", - dialect_op_type: "MatchAnySyncI64Op", - dialect_op_name: "nvvm.match_any_sync_i64", - ptx_mode: "any", - ptx_type: "b64", - selections: [ - "MATCH_ANY_SYNC_64ii", - "MATCH_ANY_SYNC_64ir", - "MATCH_ANY_SYNC_64ri", - "MATCH_ANY_SYNC_64rr", - ], - adapter: WarpMatchAdapter::DirectMask, - }, - (WarpMatchMode::All, WarpMatchValueWidth::B32) => WarpMatchRecipe { - id: "match_all_sync", - abi_id: "i0047", - operation_key: "warp.match.sync.all.b32", - source_record: "int_nvvm_match_all_sync_i32p", - llvm_symbol: "llvm.nvvm.match.all.sync.i32p", - rust_name: "match_all_sync", - rust_value: "u32", - llvm_value: "i32", - dialect_op_type: "MatchAllSyncI32Op", - dialect_op_name: "nvvm.match_all_sync_i32", - ptx_mode: "all", - ptx_type: "b32", - selections: [ - "MATCH_ALLP_SYNC_32ii", - "MATCH_ALLP_SYNC_32ir", - "MATCH_ALLP_SYNC_32ri", - "MATCH_ALLP_SYNC_32rr", - ], - adapter: WarpMatchAdapter::ProjectMaskDiscardPredicate, - }, - (WarpMatchMode::All, WarpMatchValueWidth::B64) => WarpMatchRecipe { - id: "match_all_i64_sync", - abi_id: "i0048", - operation_key: "warp.match.sync.all.b64", - source_record: "int_nvvm_match_all_sync_i64p", - llvm_symbol: "llvm.nvvm.match.all.sync.i64p", - rust_name: "match_all_i64_sync", - rust_value: "u64", - llvm_value: "i64", - dialect_op_type: "MatchAllSyncI64Op", - dialect_op_name: "nvvm.match_all_sync_i64", - ptx_mode: "all", - ptx_type: "b64", - selections: [ - "MATCH_ALLP_SYNC_64ii", - "MATCH_ALLP_SYNC_64ir", - "MATCH_ALLP_SYNC_64ri", - "MATCH_ALLP_SYNC_64rr", - ], - adapter: WarpMatchAdapter::ProjectMaskDiscardPredicate, - }, - } -} - -fn validate_warp_match_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, -) -> Result<()> { - let warp_match = policy - .warp_match - .as_ref() - .with_context(|| format!("{} has no closed warp-match contract", policy.id))?; - let recipe = warp_match_recipe(warp_match.mode, warp_match.value_width); - ensure!( - warp_match.participation - == WarpMatchParticipation::ExecutingLaneNamedAllNamedLanesSameInstructionAndMask - && warp_match.adapter == recipe.adapter - && warp_match.value_encoding == MatchOperandEncoding::RegisterOrImmediate - && warp_match.mask_encoding == MatchOperandEncoding::RegisterOrImmediate, - "{} requests an unsupported warp-match participation, adapter, or encoding", - policy.id - ); - ensure!( - policy.id == recipe.id - && policy.abi_id == recipe.abi_id - && policy.operation_key == recipe.operation_key - && policy.source.is_none() - && policy.source_record.as_deref() == Some(recipe.source_record) - && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) - && policy.resolved_llvm_symbol.is_none(), - "{} warp-match identity does not match its closed recipe", - policy.id - ); - ensure!( - policy.rust_module == "warp" - && policy.rust_name == recipe.rust_name - && policy.rust_arguments == ["u32", recipe.rust_value] - && policy.rust_result == "u32" - && !policy.safe - && policy.must_use - && policy.safe_allowlist_reason.is_none() - && policy.public_rust_path == format!("cuda_intrinsics::warp::{}", recipe.rust_name) - && policy.compatibility_rust_paths - == [format!("cuda_device::warp::{}", recipe.rust_name)], - "{} must preserve its unsafe raw and stable compatibility paths", - policy.id - ); - let expected_llvm_results: &[&str] = match warp_match.mode { - WarpMatchMode::Any => &["i32"], - WarpMatchMode::All => &["i32", "i1"], - }; - ensure!( - policy.dialect_op_type == recipe.dialect_op_type - && policy.dialect_op_name == recipe.dialect_op_name - && policy.dialect_operands == ["i32", recipe.llvm_value] - && policy.dialect_results == ["i32"] - && policy.llvm_arguments == ["i32", recipe.llvm_value] - && policy.llvm_results == expected_llvm_results - && policy.lowering == "generated_warp_match", - "{} is outside the closed two-operand warp-match lowering recipe", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == "inaccessible_read_write" - && policy.convergent - && policy.execution_scope == "warp" - && policy.minimum_ptx == "6.0" - && policy.minimum_sm.as_deref() == Some("sm_70") - && policy.ptx_result == "u32" - && policy.targets == "all", - "{} warp-match effects, carrier, or target floor disagree with its recipe", - policy.id - ); - ensure!( - policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section - == "9.7.14.11 Parallel Synchronization and Communication Instructions: match.sync" - && policy.ptx_isa_url - == "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-match-sync", - "{} warp-match PTX provenance disagrees with the reviewed recipe", - policy.id - ); - ensure!( - declaration.properties - == [ - "IntrConvergent", - "IntrInaccessibleMemOnly", - "IntrNoCallback", - ], - "{} warp-match effects disagree with the imported declaration", - policy.id - ); - ensure!( - policy.backend_lowerings.is_empty() - && policy.packed_atomic.is_none() - && policy.redux.is_none() - && policy.vote.is_none() - && policy.active_mask.is_none() - && policy.dot_product.is_none() - && policy.packed_alu.is_none() - && policy.packed_conversion.is_none() - && policy.cp_async_copy.is_none() - && policy.cp_async_control.is_none() - && policy.mbarrier_basic.is_none() - && policy.warp_barrier.is_none() - && policy.warp_shuffle.is_none() - && policy.ldmatrix_variant.is_none() - && policy.ldmatrix_safety.is_none() - && policy.ldmatrix_adapter.is_none() - && policy.selected_address_space.is_none(), - "{} mixes another generated-family contract with warp_match", - policy.id - ); - let destination = match warp_match.mode { - WarpMatchMode::Any => OperandPattern::Register, - WarpMatchMode::All => OperandPattern::RegisterPredicatePair, - }; - ensure!( - policy.expected_ptx.mnemonic == "match" - && policy.expected_ptx.modifiers == [recipe.ptx_mode, "sync", recipe.ptx_type] - && policy.expected_ptx.operands - == [ - destination, - OperandPattern::RegisterOrImmediate, - OperandPattern::RegisterOrImmediate, - ], - "{} expected PTX does not match its closed match.sync recipe", - policy.id - ); - let actual_selection_records: BTreeSet<_> = declaration - .selections - .iter() - .map(|selection| selection.source_record.as_str()) - .collect(); - ensure!( - declaration.selections.len() == 4 - && actual_selection_records == BTreeSet::from(recipe.selections), - "{} warp-match declaration must contain exactly ii/ir/ri/rr selections", - policy.id - ); - let destination = if warp_match.mode == WarpMatchMode::All { - "$dest|$pred" - } else { - "$dest" - }; - let expected_asm = format!( - "match.{}.sync.{} \t{}, $value, $mask;", - recipe.ptx_mode, recipe.ptx_type, destination - ); - for selection in &declaration.selections { - ensure!( - selection.asm == expected_asm - && selection.predicates - == [ - "Subtarget->getPTXVersion() >= 60", - "Subtarget->getSmVersion() >= 70", - ] - && selection.constraints.is_empty(), - "{} warp-match selections disagree on PTX shape, predicates, or constraints", - policy.id - ); - } - Ok(()) -} - -fn validate_elect_policy(policy: &OverlayIntrinsic, declaration: &ImportedIntrinsic) -> Result<()> { - ensure!( - policy.id == "elect_sync" - && policy.abi_id == "i0367" - && policy.operation_key == "warp.elect.sync" - && policy.source.is_none() - && policy.source_record.as_deref() == Some("int_nvvm_elect_sync") - && policy.llvm_symbol.as_deref() == Some("llvm.nvvm.elect.sync") - && policy.resolved_llvm_symbol.is_none(), - "{} elect identity does not match the closed recipe", - policy.id - ); - ensure!( - policy.rust_module == "warp" - && policy.rust_name == "elect_sync" - && policy.rust_arguments == ["u32"] - && policy.rust_result == "(u32, bool)" - && !policy.safe - && policy.must_use - && policy.safe_allowlist_reason.is_none() - && policy.public_rust_path == "cuda_intrinsics::warp::elect_sync" - && policy.compatibility_rust_paths == ["cuda_device::warp::elect_sync"], - "{} must keep its unsafe raw API and stable compatibility path", - policy.id - ); - ensure!( - policy.dialect_op_type == "ElectSyncOp" - && policy.dialect_op_name == "nvvm.elect_sync" - && policy.dialect_operands == ["i32"] - && policy.dialect_results == ["i32", "i1"] - && policy.llvm_arguments == ["i32"] - && policy.llvm_results == ["i32", "i1"] - && policy.lowering == "generated_elect", - "{} is outside the closed elect lowering recipe", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == "inaccessible_read_write" - && policy.convergent - && policy.execution_scope == "warp" - && policy.minimum_ptx == "8.0" - && policy.minimum_sm.as_deref() == Some("sm_90") - && policy.ptx_result == "(u32, bool)" - && policy.targets == "all", - "{} elect effects or target floor disagree with the closed recipe", - policy.id - ); - ensure!( - policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section - == "9.7.14.15 Parallel Synchronization and Communication Instructions: elect.sync" - && policy.ptx_isa_url - == "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-elect-sync", - "{} elect PTX provenance disagrees with the reviewed recipe", - policy.id - ); - ensure!( - declaration.classes == ["SDPatternOperator", "Intrinsic", "DefaultAttrsIntrinsic"] - && declaration.properties == ["IntrConvergent", "IntrInaccessibleMemOnly"] - && declaration.arguments == ["i32"] - && declaration.results == ["i32", "i1"], - "{} elect declaration facts changed", - policy.id - ); - ensure!( - declaration.selections.len() == 2 - && declaration.selections.iter().all(|selection| { - ["INT_ELECT_SYNC_I", "INT_ELECT_SYNC_R"].contains(&selection.source_record.as_str()) - && selection.asm == "elect.sync \t$dest|$pred, $mask;" - && selection.predicates - == [ - "Subtarget->getPTXVersion() >= 80", - "Subtarget->getSmVersion() >= 90", - ] - && selection.constraints.is_empty() - }), - "{} elect selections changed", - policy.id - ); - ensure!( - policy.expected_ptx.mnemonic == "elect" - && policy.expected_ptx.modifiers == ["sync"] - && policy.expected_ptx.operands - == [ - OperandPattern::RegisterPredicatePair, - OperandPattern::RegisterOrImmediate, - ], - "{} expected PTX does not match elect.sync", - policy.id - ); - ensure!( - policy.backend_lowerings.len() == 2 - && policy.backend_lowerings.iter().any(|route| { - route.backend == IntrinsicBackend::LlvmNvptx - && route.mechanism == BackendLoweringMechanism::TypedNvvm - && route.minimum_ptx.as_deref() == Some("8.0") - && route.minimum_sm.as_deref() == Some("sm_90") - && !route.evidence_profile.trim().is_empty() - }) - && policy.backend_lowerings.iter().any(|route| { - route.backend == IntrinsicBackend::LibNvvm - && route.mechanism == BackendLoweringMechanism::InlinePtx - && route.minimum_ptx.as_deref() == Some("8.0") - && route.minimum_sm.as_deref() == Some("sm_90") - && !route.evidence_profile.trim().is_empty() - }), - "{} must keep the LLVM typed and libNVVM inline-PTX routes explicit", - policy.id - ); - ensure!( - policy.packed_atomic.is_none() - && policy.redux.is_none() - && policy.vote.is_none() - && policy.active_mask.is_none() - && policy.warp_match.is_none() - && policy.warp_barrier.is_none() - && policy.warp_shuffle.is_none() - && policy.dot_product.is_none() - && policy.packed_alu.is_none() - && policy.packed_conversion.is_none() - && policy.cp_async_copy.is_none() - && policy.cp_async_control.is_none() - && policy.cp_async_mbarrier.is_none() - && policy.mbarrier_basic.is_none() - && policy.movmatrix.is_none() - && policy.mbarrier_extended.is_none() - && policy.register_mma.is_none() - && policy.sparse_mma.is_none() - && policy.prmt.is_none() - && policy.cluster_barrier.is_none() - && policy.wgmma_control.is_none() - && policy.special_register.is_none() - && policy.debug_control.is_none() - && policy.cluster_memory.is_none() - && policy.clc.is_none() - && policy.ldmatrix_variant.is_none() - && policy.ldmatrix_safety.is_none() - && policy.ldmatrix_adapter.is_none() - && policy.selected_address_space.is_none(), - "{} mixes another generated-family contract with elect", - policy.id - ); - Ok(()) -} - -fn validate_warp_barrier_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, -) -> Result<()> { - let barrier = policy - .warp_barrier - .as_ref() - .with_context(|| format!("{} has no closed warp-barrier contract", policy.id))?; - ensure!( - barrier.participation - == WarpBarrierParticipation::ExecutingLaneNamedAllNamedLanesSameInstructionAndMask - && barrier.legacy_pre_sm70 - == PreSm70MemberMaskRule::AllNamedLanesConvergedAndOnlyNamedLanesActive - && barrier.adapter == WarpBarrierAdapter::DirectMemberMask - && barrier.mask_encoding == WarpBarrierMaskEncoding::RegisterOrImmediate - && barrier.memory_ordering == WarpBarrierMemoryOrdering::ParticipatingLanes, - "{} requests an unsupported warp-barrier participation, legacy rule, adapter, mask encoding, or memory ordering", - policy.id - ); - ensure!( - policy.id == "sync_mask" - && policy.abi_id == "i0049" - && policy.operation_key == "warp.barrier.sync.masked" - && policy.source.is_none() - && policy.source_record.as_deref() == Some("int_nvvm_bar_warp_sync") - && policy.llvm_symbol.as_deref() == Some("llvm.nvvm.bar.warp.sync") - && policy.resolved_llvm_symbol.is_none(), - "{} warp-barrier identity does not match the closed sync_mask recipe", - policy.id - ); - ensure!( - policy.rust_module == "warp" - && policy.rust_name == "sync_mask" - && policy.rust_arguments == ["u32"] - && policy.rust_result == "()" - && !policy.safe - && !policy.must_use - && policy.safe_allowlist_reason.is_none() - && policy.public_rust_path == "cuda_intrinsics::warp::sync_mask" - && policy.compatibility_rust_paths == ["cuda_device::warp::sync_mask"], - "{} must keep its unsafe raw API and safe cuda-device compatibility path distinct", - policy.id - ); - ensure!( - policy.dialect_op_type == "BarWarpSyncOp" - && policy.dialect_op_name == "nvvm.bar_warp_sync" - && policy.dialect_operands == ["i32"] - && policy.dialect_results.is_empty() - && policy.llvm_arguments == ["i32"] - && policy.llvm_results.is_empty() - && policy.lowering == "generated_warp_barrier", - "{} is outside the closed one-mask warp-barrier lowering recipe", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == "read_write" - && policy.convergent - && policy.execution_scope == "warp" - && policy.minimum_ptx == "6.0" - && policy.minimum_sm.as_deref() == Some("sm_30") - && policy.ptx_result == "()" - && policy.targets == "all", - "{} warp-barrier effects or target floor disagree with the closed recipe", - policy.id - ); - ensure!( - policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section - == "9.7.14.2 Parallel Synchronization and Communication Instructions: bar.warp.sync" - && policy.ptx_isa_url - == "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-bar-warp-sync", - "{} warp-barrier PTX provenance disagrees with the reviewed recipe", - policy.id - ); - ensure!( - declaration.properties == ["IntrConvergent", "IntrNoCallback"], - "{} warp-barrier effects disagree with the imported declaration", - policy.id - ); - ensure!( - policy.packed_atomic.is_none() - && policy.redux.is_none() - && policy.vote.is_none() - && policy.active_mask.is_none() - && policy.warp_match.is_none() - && policy.warp_shuffle.is_none() - && policy.dot_product.is_none() - && policy.packed_alu.is_none() - && policy.packed_conversion.is_none() - && policy.cp_async_copy.is_none() - && policy.cp_async_control.is_none() - && policy.mbarrier_basic.is_none() - && policy.ldmatrix_variant.is_none() - && policy.ldmatrix_safety.is_none() - && policy.ldmatrix_adapter.is_none() - && policy.selected_address_space.is_none(), - "{} mixes another generated-family contract with warp_barrier", - policy.id - ); - ensure!( - policy.expected_ptx.mnemonic == "bar" - && policy.expected_ptx.modifiers == ["warp", "sync"] - && policy.expected_ptx.operands == [OperandPattern::RegisterOrImmediate], - "{} expected PTX does not match bar.warp.sync mask", - policy.id - ); - - let backend_pairs: BTreeSet<_> = policy - .backend_lowerings - .iter() - .map(|lowering| (lowering.backend, lowering.mechanism)) - .collect(); - ensure!( - policy.backend_lowerings.len() == 2 - && backend_pairs - == BTreeSet::from([ - ( - IntrinsicBackend::LlvmNvptx, - BackendLoweringMechanism::TypedNvvm, - ), - ( - IntrinsicBackend::LibNvvm, - BackendLoweringMechanism::TypedNvvm, - ), - ]), - "{} must define exactly the reviewed typed LLVM and libNVVM routes", - policy.id - ); - for lowering in &policy.backend_lowerings { - let floor_matches = match lowering.backend { - IntrinsicBackend::LlvmNvptx => { - lowering.mechanism == BackendLoweringMechanism::TypedNvvm - && lowering.minimum_ptx.as_deref() == Some("6.0") - && lowering.minimum_sm.as_deref() == Some("sm_30") - } - IntrinsicBackend::LibNvvm => { - lowering.mechanism == BackendLoweringMechanism::TypedNvvm - && lowering.minimum_ptx.as_deref() == Some("6.0") - && lowering.minimum_sm.as_deref() == Some("sm_75") - } - }; - ensure!( - floor_matches && !lowering.evidence_profile.trim().is_empty(), - "{} backend {:?} does not carry its reviewed warp-barrier profile floor", - policy.id, - lowering.backend - ); - } - - let expected_selection_records = BTreeSet::from(["INT_BAR_WARP_SYNC_I", "INT_BAR_WARP_SYNC_R"]); - let actual_selection_records: BTreeSet<_> = declaration - .selections - .iter() - .map(|selection| selection.source_record.as_str()) - .collect(); - ensure!( - declaration.selections.len() == 2 && actual_selection_records == expected_selection_records, - "{} warp-barrier declaration must contain exactly its immediate/register selection pair", - policy.id - ); - for selection in &declaration.selections { - ensure!( - selection.asm == "bar.warp.sync \t$i;" - && selection.predicates - == [ - "Subtarget->getPTXVersion() >= 60", - "Subtarget->getSmVersion() >= 30", - ] - && selection.constraints.is_empty(), - "{} warp-barrier selections disagree on PTX shape, target predicates, or constraints", - policy.id - ); - } - Ok(()) -} - -fn validate_warp_shuffle_policy( - policy: &OverlayIntrinsic, - declaration: Option<&ImportedIntrinsic>, -) -> Result<()> { - let shuffle = policy - .warp_shuffle - .as_ref() - .with_context(|| format!("{} has no closed warp-shuffle contract", policy.id))?; - let recipe = warp_shuffle_recipe(shuffle.mode, shuffle.value_kind); - ensure!( - shuffle.participation - == WarpShuffleParticipation::ExecutingLaneNamedAllNamedLanesSameInstructionAndMask - && shuffle.legacy_pre_sm70 - == PreSm70MemberMaskRule::AllNamedLanesConvergedAndOnlyNamedLanesActive - && shuffle.source_lane - == WarpShuffleSourceLane::InRangeSourceActiveAndNamedOutOfRangeCopiesSelf - && shuffle.adapter == recipe.adapter - && shuffle.clamp == recipe.clamp - && shuffle.lane_encoding == recipe.operand_encoding - && shuffle.mask_encoding == recipe.operand_encoding, - "{} requests an unsupported warp-shuffle semantic or operand contract", - policy.id - ); - - let source_matches = match recipe.source { - WarpShuffleRecipeSource::LlvmImported { - source_record, - llvm_symbol, - } => { - policy.source.is_none() - && policy.source_record.as_deref() == Some(source_record) - && policy.llvm_symbol.as_deref() == Some(llvm_symbol) - && policy.resolved_llvm_symbol.is_none() - } - WarpShuffleRecipeSource::PtxNative { instruction } => { - policy.source - == Some(IntrinsicSource::PtxNative { - instruction: instruction.into(), - }) - && policy.source_record.is_none() - && policy.llvm_symbol.is_none() - && policy.resolved_llvm_symbol.is_none() - && policy.llvm_arguments.is_empty() - && policy.llvm_results.is_empty() - } - }; - ensure!( - policy.id == recipe.id - && policy.abi_id == recipe.abi_id - && policy.operation_key == recipe.operation_key - && source_matches, - "{} warp-shuffle identity does not match its closed mode and value recipe", - policy.id - ); - ensure!( - policy.rust_module == "warp" - && policy.rust_name == recipe.rust_name - && policy.rust_arguments == ["u32", recipe.rust_value, "u32"] - && policy.rust_result == recipe.rust_value - && !policy.safe - && policy.must_use - && policy.safe_allowlist_reason.is_none() - && policy.public_rust_path == format!("cuda_intrinsics::warp::{}", recipe.rust_name) - && policy.compatibility_rust_paths - == [format!("cuda_device::warp::{}", recipe.rust_name)], - "{} must preserve its unsafe must-use warp-shuffle raw API and compatibility path", - policy.id - ); - ensure!( - policy.dialect_op_type == recipe.dialect_op_type - && policy.dialect_op_name == recipe.dialect_op_name - && policy.dialect_operands == ["i32", recipe.dialect_value, "i32"] - && policy.dialect_results == [recipe.dialect_value] - && policy.lowering == recipe.lowering - && match recipe.source { - WarpShuffleRecipeSource::LlvmImported { .. } => { - policy.llvm_arguments == ["i32", recipe.dialect_value, "i32", "i32"] - && policy.llvm_results == [recipe.dialect_value] - } - WarpShuffleRecipeSource::PtxNative { .. } => { - policy.llvm_arguments.is_empty() && policy.llvm_results.is_empty() - } - }, - "{} is outside the closed warp-shuffle lowering recipe", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == "inaccessible_read_write" - && policy.convergent - && policy.execution_scope == "warp" - && policy.minimum_ptx == "6.0" - && policy.minimum_sm.as_deref() == Some("sm_30") - && policy.ptx_result == recipe.rust_value - && policy.targets == "all", - "{} warp-shuffle effects, carrier, or target floor disagree with its recipe", - policy.id - ); - ensure!( - policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section - == "9.7.9.6 Data Movement and Conversion Instructions: shfl.sync" - && policy.ptx_isa_url - == "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-shfl-sync", - "{} warp-shuffle PTX provenance disagrees with the reviewed recipe", - policy.id - ); - if let Some(declaration) = declaration { - ensure!( - matches!(recipe.source, WarpShuffleRecipeSource::LlvmImported { .. }) - && declaration.classes - == [ - "ClangBuiltin", - "NVVMBuiltin", - "SDPatternOperator", - "Intrinsic" - ] - && declaration.properties - == [ - "IntrConvergent", - "IntrInaccessibleMemOnly", - "IntrNoCallback", - ], - "{} warp-shuffle class or effects disagree with the imported declaration", - policy.id - ); - } else { - ensure!( - matches!(recipe.source, WarpShuffleRecipeSource::PtxNative { .. }), - "{} imported warp shuffle is missing its LLVM declaration", - policy.id - ); - } - ensure!( - policy.packed_atomic.is_none() - && policy.redux.is_none() - && policy.vote.is_none() - && policy.active_mask.is_none() - && policy.warp_match.is_none() - && policy.warp_barrier.is_none() - && policy.dot_product.is_none() - && policy.packed_alu.is_none() - && policy.packed_conversion.is_none() - && policy.cp_async_copy.is_none() - && policy.cp_async_control.is_none() - && policy.mbarrier_basic.is_none() - && policy.ldmatrix_variant.is_none() - && policy.ldmatrix_safety.is_none() - && policy.ldmatrix_adapter.is_none() - && policy.selected_address_space.is_none(), - "{} mixes another generated-family contract with warp_shuffle", - policy.id - ); - let expected_operands = match recipe.adapter { - WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp => vec![ - OperandPattern::Register, - OperandPattern::Register, - OperandPattern::RegisterOrImmediate, - OperandPattern::Exact { - value: recipe.clamp.to_string(), - }, - OperandPattern::RegisterOrImmediate, - ], - WarpShuffleAdapter::MaskValueLaneOrDeltaSplitI64LowHighB32InsertClampReassemble => vec![ - OperandPattern::Exact { value: "lo".into() }, - OperandPattern::Exact { value: "lo".into() }, - OperandPattern::Register, - OperandPattern::Exact { - value: recipe.clamp.to_string(), - }, - OperandPattern::Register, - ], - }; - ensure!( - policy.expected_ptx.mnemonic == "shfl" - && policy.expected_ptx.modifiers == ["sync", recipe.ptx_mode, "b32"] - && policy.expected_ptx.operands == expected_operands, - "{} expected PTX does not match its closed shfl.sync recipe", - policy.id - ); - - let backend_pairs: BTreeSet<_> = policy - .backend_lowerings - .iter() - .map(|lowering| (lowering.backend, lowering.mechanism)) - .collect(); - ensure!( - policy.backend_lowerings.len() == 2 - && backend_pairs - == BTreeSet::from([ - (IntrinsicBackend::LlvmNvptx, recipe.backend_mechanism), - (IntrinsicBackend::LibNvvm, recipe.backend_mechanism), - ]), - "{} must define exactly the reviewed LLVM and libNVVM routes", - policy.id - ); - for lowering in &policy.backend_lowerings { - let floor_matches = match lowering.backend { - IntrinsicBackend::LlvmNvptx => { - lowering.mechanism == recipe.backend_mechanism - && lowering.minimum_ptx.as_deref() == Some("6.0") - && lowering.minimum_sm.as_deref() == Some("sm_30") - } - IntrinsicBackend::LibNvvm => { - lowering.mechanism == recipe.backend_mechanism - && lowering.minimum_ptx.as_deref() == Some("6.0") - && lowering.minimum_sm.as_deref() == Some("sm_75") - } - }; - ensure!( - floor_matches && !lowering.evidence_profile.trim().is_empty(), - "{} backend {:?} does not carry its reviewed warp-shuffle profile floor", - policy.id, - lowering.backend - ); - } - - if let Some(declaration) = declaration { - let selection_records: BTreeSet<_> = declaration - .selections - .iter() - .map(|selection| selection.source_record.as_str()) - .collect(); - ensure!( - declaration.selections.len() == 8 - && selection_records.len() == 8 - && selection_records - .iter() - .all(|source_record| !source_record.trim().is_empty()), - "{} warp-shuffle declaration must contain exactly eight distinct operand-encoding selections", - policy.id - ); - let expected_asm = format!( - "shfl.sync.{}.b32 \t$dst, $src, $offset, $mask, $threadmask;", - recipe.ptx_mode - ); - for selection in &declaration.selections { - ensure!( - selection.asm == expected_asm - && selection.predicates - == [ - "Subtarget->getPTXVersion() >= 60", - "Subtarget->getSmVersion() >= 30", - ] - && selection.constraints.is_empty(), - "{} warp-shuffle selections disagree on PTX shape, target predicates, or constraints", - policy.id - ); - } - } - Ok(()) -} - -#[derive(Clone, Copy)] -enum WarpShuffleRecipeSource { - LlvmImported { - source_record: &'static str, - llvm_symbol: &'static str, - }, - PtxNative { - instruction: &'static str, - }, -} - -struct WarpShuffleRecipe { - id: &'static str, - abi_id: &'static str, - operation_key: &'static str, - source: WarpShuffleRecipeSource, - rust_name: &'static str, - rust_value: &'static str, - dialect_value: &'static str, - dialect_op_type: &'static str, - dialect_op_name: &'static str, - ptx_mode: &'static str, - clamp: u32, - adapter: WarpShuffleAdapter, - operand_encoding: WarpShuffleOperandEncoding, - lowering: &'static str, - backend_mechanism: BackendLoweringMechanism, -} - -fn warp_shuffle_recipe( - mode: WarpShuffleMode, - value_kind: WarpShuffleValueKind, -) -> WarpShuffleRecipe { - match (mode, value_kind) { - (WarpShuffleMode::Idx, WarpShuffleValueKind::I32) => WarpShuffleRecipe { - id: "shuffle_sync", - abi_id: "i0050", - operation_key: "warp.shuffle.sync.idx.i32", - source: WarpShuffleRecipeSource::LlvmImported { - source_record: "int_nvvm_shfl_sync_idx_i32", - llvm_symbol: "llvm.nvvm.shfl.sync.idx.i32", - }, - rust_name: "shuffle_sync", - rust_value: "u32", - dialect_value: "i32", - dialect_op_type: "ShflSyncIdxI32Op", - dialect_op_name: "nvvm.shfl_sync_idx_i32", - ptx_mode: "idx", - clamp: 31, - adapter: WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp, - operand_encoding: WarpShuffleOperandEncoding::RegisterOrImmediate, - lowering: "generated_warp_shuffle", - backend_mechanism: BackendLoweringMechanism::TypedNvvm, - }, - (WarpShuffleMode::Bfly, WarpShuffleValueKind::I32) => WarpShuffleRecipe { - id: "shuffle_xor_sync", - abi_id: "i0051", - operation_key: "warp.shuffle.sync.bfly.i32", - source: WarpShuffleRecipeSource::LlvmImported { - source_record: "int_nvvm_shfl_sync_bfly_i32", - llvm_symbol: "llvm.nvvm.shfl.sync.bfly.i32", - }, - rust_name: "shuffle_xor_sync", - rust_value: "u32", - dialect_value: "i32", - dialect_op_type: "ShflSyncBflyI32Op", - dialect_op_name: "nvvm.shfl_sync_bfly_i32", - ptx_mode: "bfly", - clamp: 31, - adapter: WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp, - operand_encoding: WarpShuffleOperandEncoding::RegisterOrImmediate, - lowering: "generated_warp_shuffle", - backend_mechanism: BackendLoweringMechanism::TypedNvvm, - }, - (WarpShuffleMode::Down, WarpShuffleValueKind::I32) => WarpShuffleRecipe { - id: "shuffle_down_sync", - abi_id: "i0052", - operation_key: "warp.shuffle.sync.down.i32", - source: WarpShuffleRecipeSource::LlvmImported { - source_record: "int_nvvm_shfl_sync_down_i32", - llvm_symbol: "llvm.nvvm.shfl.sync.down.i32", - }, - rust_name: "shuffle_down_sync", - rust_value: "u32", - dialect_value: "i32", - dialect_op_type: "ShflSyncDownI32Op", - dialect_op_name: "nvvm.shfl_sync_down_i32", - ptx_mode: "down", - clamp: 31, - adapter: WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp, - operand_encoding: WarpShuffleOperandEncoding::RegisterOrImmediate, - lowering: "generated_warp_shuffle", - backend_mechanism: BackendLoweringMechanism::TypedNvvm, - }, - (WarpShuffleMode::Up, WarpShuffleValueKind::I32) => WarpShuffleRecipe { - id: "shuffle_up_sync", - abi_id: "i0053", - operation_key: "warp.shuffle.sync.up.i32", - source: WarpShuffleRecipeSource::LlvmImported { - source_record: "int_nvvm_shfl_sync_up_i32", - llvm_symbol: "llvm.nvvm.shfl.sync.up.i32", - }, - rust_name: "shuffle_up_sync", - rust_value: "u32", - dialect_value: "i32", - dialect_op_type: "ShflSyncUpI32Op", - dialect_op_name: "nvvm.shfl_sync_up_i32", - ptx_mode: "up", - clamp: 0, - adapter: WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp, - operand_encoding: WarpShuffleOperandEncoding::RegisterOrImmediate, - lowering: "generated_warp_shuffle", - backend_mechanism: BackendLoweringMechanism::TypedNvvm, - }, - (WarpShuffleMode::Idx, WarpShuffleValueKind::F32) => WarpShuffleRecipe { - id: "shuffle_f32_sync", - abi_id: "i0054", - operation_key: "warp.shuffle.sync.idx.f32", - source: WarpShuffleRecipeSource::LlvmImported { - source_record: "int_nvvm_shfl_sync_idx_f32", - llvm_symbol: "llvm.nvvm.shfl.sync.idx.f32", - }, - rust_name: "shuffle_f32_sync", - rust_value: "f32", - dialect_value: "f32", - dialect_op_type: "ShflSyncIdxF32Op", - dialect_op_name: "nvvm.shfl_sync_idx_f32", - ptx_mode: "idx", - clamp: 31, - adapter: WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp, - operand_encoding: WarpShuffleOperandEncoding::RegisterOrImmediate, - lowering: "generated_warp_shuffle", - backend_mechanism: BackendLoweringMechanism::TypedNvvm, - }, - (WarpShuffleMode::Bfly, WarpShuffleValueKind::F32) => WarpShuffleRecipe { - id: "shuffle_xor_f32_sync", - abi_id: "i0055", - operation_key: "warp.shuffle.sync.bfly.f32", - source: WarpShuffleRecipeSource::LlvmImported { - source_record: "int_nvvm_shfl_sync_bfly_f32", - llvm_symbol: "llvm.nvvm.shfl.sync.bfly.f32", - }, - rust_name: "shuffle_xor_f32_sync", - rust_value: "f32", - dialect_value: "f32", - dialect_op_type: "ShflSyncBflyF32Op", - dialect_op_name: "nvvm.shfl_sync_bfly_f32", - ptx_mode: "bfly", - clamp: 31, - adapter: WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp, - operand_encoding: WarpShuffleOperandEncoding::RegisterOrImmediate, - lowering: "generated_warp_shuffle", - backend_mechanism: BackendLoweringMechanism::TypedNvvm, - }, - (WarpShuffleMode::Down, WarpShuffleValueKind::F32) => WarpShuffleRecipe { - id: "shuffle_down_f32_sync", - abi_id: "i0056", - operation_key: "warp.shuffle.sync.down.f32", - source: WarpShuffleRecipeSource::LlvmImported { - source_record: "int_nvvm_shfl_sync_down_f32", - llvm_symbol: "llvm.nvvm.shfl.sync.down.f32", - }, - rust_name: "shuffle_down_f32_sync", - rust_value: "f32", - dialect_value: "f32", - dialect_op_type: "ShflSyncDownF32Op", - dialect_op_name: "nvvm.shfl_sync_down_f32", - ptx_mode: "down", - clamp: 31, - adapter: WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp, - operand_encoding: WarpShuffleOperandEncoding::RegisterOrImmediate, - lowering: "generated_warp_shuffle", - backend_mechanism: BackendLoweringMechanism::TypedNvvm, - }, - (WarpShuffleMode::Up, WarpShuffleValueKind::F32) => WarpShuffleRecipe { - id: "shuffle_up_f32_sync", - abi_id: "i0057", - operation_key: "warp.shuffle.sync.up.f32", - source: WarpShuffleRecipeSource::LlvmImported { - source_record: "int_nvvm_shfl_sync_up_f32", - llvm_symbol: "llvm.nvvm.shfl.sync.up.f32", - }, - rust_name: "shuffle_up_f32_sync", - rust_value: "f32", - dialect_value: "f32", - dialect_op_type: "ShflSyncUpF32Op", - dialect_op_name: "nvvm.shfl_sync_up_f32", - ptx_mode: "up", - clamp: 0, - adapter: WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp, - operand_encoding: WarpShuffleOperandEncoding::RegisterOrImmediate, - lowering: "generated_warp_shuffle", - backend_mechanism: BackendLoweringMechanism::TypedNvvm, - }, - (WarpShuffleMode::Idx, WarpShuffleValueKind::I64) => WarpShuffleRecipe { - id: "shuffle_u64_sync", - abi_id: "i0058", - operation_key: "warp.shuffle.sync.idx.i64", - source: WarpShuffleRecipeSource::PtxNative { - instruction: "shfl.sync.idx.b32", - }, - rust_name: "shuffle_u64_sync", - rust_value: "u64", - dialect_value: "i64", - dialect_op_type: "ShflSyncIdxI64Op", - dialect_op_name: "nvvm.shfl_sync_idx_i64", - ptx_mode: "idx", - clamp: 31, - adapter: - WarpShuffleAdapter::MaskValueLaneOrDeltaSplitI64LowHighB32InsertClampReassemble, - operand_encoding: WarpShuffleOperandEncoding::RegisterOnly, - lowering: "generated_warp_shuffle_i64_inline_ptx", - backend_mechanism: BackendLoweringMechanism::InlinePtx, - }, - (WarpShuffleMode::Bfly, WarpShuffleValueKind::I64) => WarpShuffleRecipe { - id: "shuffle_xor_u64_sync", - abi_id: "i0059", - operation_key: "warp.shuffle.sync.bfly.i64", - source: WarpShuffleRecipeSource::PtxNative { - instruction: "shfl.sync.bfly.b32", - }, - rust_name: "shuffle_xor_u64_sync", - rust_value: "u64", - dialect_value: "i64", - dialect_op_type: "ShflSyncBflyI64Op", - dialect_op_name: "nvvm.shfl_sync_bfly_i64", - ptx_mode: "bfly", - clamp: 31, - adapter: - WarpShuffleAdapter::MaskValueLaneOrDeltaSplitI64LowHighB32InsertClampReassemble, - operand_encoding: WarpShuffleOperandEncoding::RegisterOnly, - lowering: "generated_warp_shuffle_i64_inline_ptx", - backend_mechanism: BackendLoweringMechanism::InlinePtx, - }, - (WarpShuffleMode::Down, WarpShuffleValueKind::I64) => WarpShuffleRecipe { - id: "shuffle_down_u64_sync", - abi_id: "i0060", - operation_key: "warp.shuffle.sync.down.i64", - source: WarpShuffleRecipeSource::PtxNative { - instruction: "shfl.sync.down.b32", - }, - rust_name: "shuffle_down_u64_sync", - rust_value: "u64", - dialect_value: "i64", - dialect_op_type: "ShflSyncDownI64Op", - dialect_op_name: "nvvm.shfl_sync_down_i64", - ptx_mode: "down", - clamp: 31, - adapter: - WarpShuffleAdapter::MaskValueLaneOrDeltaSplitI64LowHighB32InsertClampReassemble, - operand_encoding: WarpShuffleOperandEncoding::RegisterOnly, - lowering: "generated_warp_shuffle_i64_inline_ptx", - backend_mechanism: BackendLoweringMechanism::InlinePtx, - }, - (WarpShuffleMode::Up, WarpShuffleValueKind::I64) => WarpShuffleRecipe { - id: "shuffle_up_u64_sync", - abi_id: "i0061", - operation_key: "warp.shuffle.sync.up.i64", - source: WarpShuffleRecipeSource::PtxNative { - instruction: "shfl.sync.up.b32", - }, - rust_name: "shuffle_up_u64_sync", - rust_value: "u64", - dialect_value: "i64", - dialect_op_type: "ShflSyncUpI64Op", - dialect_op_name: "nvvm.shfl_sync_up_i64", - ptx_mode: "up", - clamp: 0, - adapter: - WarpShuffleAdapter::MaskValueLaneOrDeltaSplitI64LowHighB32InsertClampReassemble, - operand_encoding: WarpShuffleOperandEncoding::RegisterOnly, - lowering: "generated_warp_shuffle_i64_inline_ptx", - backend_mechanism: BackendLoweringMechanism::InlinePtx, - }, - } -} - -fn validate_selected_target_predicates( - policy: &OverlayIntrinsic, - selection: &crate::model::ImportedSelection, -) -> Result<()> { - if policy.family == "wgmma_control" { - ensure!( - selection.predicates == ["Subtarget->getPTXVersion() >= 80", "hasSM90a"] - && parse_ptx_version(&policy.minimum_ptx, &policy.id)?.encoded() == 80 - && parse_hardware_target(policy)? - == CatalogHardwareTarget::AnyOf { - alternatives: vec![CatalogHardwareAlternative::ExactArchitecture { - sm: 90, - }], - }, - "{} WGMMA-control selection must retain the PTX 8.0 and sm_90a gates", - policy.id - ); - return Ok(()); - } - if policy.family == "register_control" { - let operation = ExecutionControlOperation::from_catalog_id(&policy.id) - .with_context(|| format!("{} has no closed register-control operation", policy.id))?; - ensure!( - matches!( - operation, - ExecutionControlOperation::SetMaxNRegInc | ExecutionControlOperation::SetMaxNRegDec - ) && selection.predicates == ["Subtarget->hasSetMaxNRegSupport()"] - && parse_ptx_version(&policy.minimum_ptx, &policy.id)?.encoded() == 80 - && policy.minimum_sm.is_none() - && policy.targets == TENSOR_MAP_REPLACE_TARGETS, - "{} setmaxnreg selection must retain the helper predicate and reviewed PTX 8.0 target matrix", - policy.id - ); - return Ok(()); - } - - let mma_family = matches!(policy.family.as_str(), "register_mma" | "sparse_mma"); - let tcgen05_mma = policy - .tcgen05 - .as_ref() - .and_then(|tcgen05| tcgen05.mma.as_ref()); - let mut imported_ptx: Option = None; - let mut imported_sm: Option = None; - let mut has_dot_instructions = false; - let mut has_clc_multicast_support = false; - let mut has_tma_blackwell_support = false; - let mut has_tcgen05_support = false; - let mut has_tcgen05_shift_support = false; - let mut has_tcgen05_mma_i8_support = false; - let mut has_ldstmatrix_blackwell_support = false; - let mut has_mma_block_scale_support = false; - let mut has_redux_sync_f32_support = false; - for predicate in &selection.predicates { - if let Some(value) = predicate.strip_prefix("Subtarget->getPTXVersion() >= ") { - let value = value.parse::().with_context(|| { - format!("{} has malformed PTX predicate {predicate:?}", policy.id) - })?; - ensure!( - imported_ptx.is_none() || mma_family, - "{} has duplicate PTX predicates", - policy.id - ); - imported_ptx = Some(imported_ptx.unwrap_or_default().max(value)); - } else if let Some(value) = predicate.strip_prefix("Subtarget->getSmVersion() >= ") { - let value = value.parse::().with_context(|| { - format!("{} has malformed SM predicate {predicate:?}", policy.id) - })?; - ensure!( - imported_sm.is_none() || mma_family, - "{} has duplicate SM predicates", - policy.id - ); - imported_sm = Some(imported_sm.unwrap_or_default().max(value)); - } else if predicate == "hasDotInstructions" { - ensure!( - policy.family == "dotprod", - "{} selected instruction uses dot-product target gating outside the dotprod family", - policy.id - ); - ensure!( - !has_dot_instructions && imported_ptx.is_none() && imported_sm.is_none(), - "{} has duplicate or conflicting dot-product target predicates", - policy.id - ); - has_dot_instructions = true; - imported_ptx = Some(50); - imported_sm = Some(61); - } else if predicate == "Subtarget->hasClusterLaunchControlTryCancelMulticastSupport()" { - ensure!( - policy.family == "clc" - && policy - .clc - .as_ref() - .is_some_and(|clc| { clc.operation == ClcOperation::TryCancelMulticast }), - "{} uses the CLC multicast target predicate outside that operation", - policy.id - ); - ensure!( - !has_clc_multicast_support - && imported_ptx.is_none() - && imported_sm.is_none() - && !has_dot_instructions, - "{} has duplicate or conflicting CLC multicast target predicates", - policy.id - ); - has_clc_multicast_support = true; - } else if predicate == "Subtarget->hasTMABlackwellSupport()" { - ensure!( - policy.family == "tma" - && policy.tma.as_ref().is_some_and(|tma| { - matches!( - tma.operation, - TmaOperation::G2sTile2dMulticastCg2 - | TmaOperation::PrefetchTileGather4TwoDimensional - | TmaOperation::PrefetchTileGather4TwoDimensionalCacheHint - ) - }), - "{} uses the Blackwell TMA target predicate outside the cta_group::2 operation", - policy.id - ); - ensure!( - !has_tma_blackwell_support - && imported_ptx.is_none() - && imported_sm.is_none() - && !has_dot_instructions - && !has_clc_multicast_support, - "{} has duplicate or conflicting Blackwell TMA target predicates", - policy.id - ); - has_tma_blackwell_support = true; - } else if predicate == "Subtarget->hasTcgen05InstSupport()" { - ensure!( - policy.family == "tcgen05" && policy.tcgen05.is_some(), - "{} uses the tcgen05 target predicate outside that family", - policy.id - ); - ensure!( - !has_tcgen05_support - && !has_tcgen05_shift_support - && !has_tcgen05_mma_i8_support - && imported_ptx.is_none() - && imported_sm.is_none() - && !has_dot_instructions - && !has_clc_multicast_support - && !has_tma_blackwell_support, - "{} has duplicate or conflicting tcgen05 target predicates", - policy.id - ); - has_tcgen05_support = true; - } else if predicate == "Subtarget->hasTcgen05MMAI8Kind()" { - ensure!( - tcgen05_mma.is_some() && selection.asm.contains(".kind::i8."), - "{} uses the tcgen05 I8 predicate outside an I8 MMA selection", - policy.id - ); - ensure!( - !has_tcgen05_mma_i8_support - && !has_tcgen05_support - && !has_tcgen05_shift_support - && imported_ptx.is_none() - && imported_sm.is_none() - && !has_dot_instructions - && !has_clc_multicast_support - && !has_tma_blackwell_support, - "{} has duplicate or conflicting tcgen05 I8 predicates", - policy.id - ); - has_tcgen05_mma_i8_support = true; - } else if predicate == "Subtarget->hasTcgen05ShiftSupport()" { - ensure!( - policy.family == "tcgen05" - && policy.tcgen05.as_ref().is_some_and(|tcgen05| { - matches!( - tcgen05.operation, - Tcgen05Operation::ShiftDown | Tcgen05Operation::ShiftDownCg2 - ) - }), - "{} uses the tcgen05 shift target predicate outside a shift operation", - policy.id - ); - ensure!( - !has_tcgen05_shift_support - && !has_tcgen05_support - && !has_tcgen05_mma_i8_support - && imported_ptx.is_none() - && imported_sm.is_none() - && !has_dot_instructions - && !has_clc_multicast_support - && !has_tma_blackwell_support, - "{} has duplicate or conflicting tcgen05 shift target predicates", - policy.id - ); - has_tcgen05_shift_support = true; - } else if predicate == "Subtarget->hasLdStmatrixBlackwellSupport()" { - ensure!( - policy.family == "ldmatrix" - && policy.ldmatrix_variant.as_ref().is_some_and(|variant| { - variant.shape != LdmatrixShape::M8n8 - && variant.element != LdmatrixElement::B16 - }), - "{} uses the Blackwell ld/stmatrix predicate outside a Blackwell ldmatrix variant", - policy.id - ); - ensure!( - !has_ldstmatrix_blackwell_support - && imported_ptx.is_none() - && imported_sm.is_none() - && !has_dot_instructions - && !has_clc_multicast_support - && !has_tma_blackwell_support - && !has_tcgen05_support - && !has_tcgen05_shift_support, - "{} has duplicate or conflicting Blackwell ldmatrix target predicates", - policy.id - ); - has_ldstmatrix_blackwell_support = true; - } else if predicate == "Subtarget->hasMMABlockScale()" { - ensure!( - matches!(policy.family.as_str(), "register_mma" | "sparse_mma") - && (policy.register_mma.is_some() || policy.sparse_mma.is_some()), - "{} uses the MMA block-scale target predicate outside a closed MMA family", - policy.id - ); - ensure!( - !has_mma_block_scale_support - && imported_ptx.is_none() - && imported_sm.is_none() - && !has_dot_instructions - && !has_clc_multicast_support - && !has_tma_blackwell_support - && !has_tcgen05_support - && !has_tcgen05_shift_support - && !has_ldstmatrix_blackwell_support, - "{} has duplicate or conflicting MMA block-scale target predicates", - policy.id - ); - has_mma_block_scale_support = true; - } else if predicate == "Subtarget->hasReduxSyncF32()" { - ensure!( - policy.family == "redux" - && policy - .redux - .as_ref() - .is_some_and(|redux| redux_recipe(redux.operation).value_type == "f32"), - "{} uses the f32 redux target predicate outside an f32 redux operation", - policy.id - ); - ensure!( - !has_redux_sync_f32_support && imported_ptx.is_none() && imported_sm.is_none(), - "{} has duplicate or conflicting f32 redux target predicates", - policy.id - ); - has_redux_sync_f32_support = true; - } else { - bail!( - "{} selected instruction has unsupported target predicate {predicate:?}; target gates must fail closed", - policy.id - ); - } - } - let overlay_ptx = parse_ptx_version(&policy.minimum_ptx, &policy.id)?.encoded(); - if let Some(mma) = tcgen05_mma { - let i8_selection = selection.asm.contains(".kind::i8."); - ensure!( - selection.predicates.len() == 1 - && ((i8_selection && has_tcgen05_mma_i8_support && !has_tcgen05_support) - || (!i8_selection && has_tcgen05_support && !has_tcgen05_mma_i8_support)) - && !has_tcgen05_shift_support, - "{} tcgen05 MMA selection does not retain its exact kind predicate", - policy.id - ); - let contracts = expected_tcgen05_mma_target_contracts(IntrinsicBackend::LlvmNvptx); - let expected_target = if let Some(fixed) = mma.fixed_selectors { - resolve_target_contract( - "tcgen05 MMA selected predicate", - &[TargetSelectorBinding { - name: "kind".into(), - value: tcgen05_mma_kind_name(fixed.kind).into(), - }], - &contracts, - )? - } else { - resolve_target_contracts("tcgen05 MMA selected predicate", &contracts)? - }; - ensure!( - mma.llvm_target == expected_target, - "{} tcgen05 MMA predicate does not map to its closed LLVM target matrix", - policy.id - ); - return Ok(()); - } - if has_mma_block_scale_support { - let (minimum_ptx_matches, target_matches) = match policy.family.as_str() { - "register_mma" => ( - overlay_ptx == 87, - policy.targets == "sm_120a|sm_120f|sm_121a|sm_121f" && policy.minimum_sm.is_none(), - ), - "sparse_mma" => ( - overlay_ptx == 87, - policy.minimum_sm.is_none() - && policy.sparse_mma.as_ref().is_some_and(|mma| { - policy.targets - == match mma.accumulator { - SparseMmaAccumulator::F16 | SparseMmaAccumulator::F32 => { - SPARSE_MMA_F8F6F4_TARGETS - } - SparseMmaAccumulator::S32 => return false, - } - }), - ), - _ => (false, false), - }; - ensure!( - minimum_ptx_matches && target_matches, - "{} MMA block-scale predicate requires its reviewed Blackwell target matrix", - policy.id - ); - return Ok(()); - } - if has_redux_sync_f32_support { - ensure!( - selection.predicates.len() == 1 - && overlay_ptx == 86 - && policy.targets == REDUX_F32_TARGETS - && policy.minimum_sm.is_none(), - "{} f32 redux selection must carry only the hasReduxSyncF32 predicate and its reviewed Blackwell target matrix", - policy.id - ); - return Ok(()); - } - // MMA uses reviewed inline PTX. Its imported predicates gate LLVM's typed - // selection, while the closed recipe and terminal evidence set the native - // PTX floor. - if mma_family { - return Ok(()); - } - if let Some(imported_ptx) = imported_ptx { - ensure!( - overlay_ptx == imported_ptx, - "{} minimum PTX {} disagrees with selected instruction predicate PTX {}", - policy.id, - policy.minimum_ptx, - format_args!("{}.{}", imported_ptx / 10, imported_ptx % 10) - ); - } - if let Some(imported_sm) = imported_sm { - if let Some(packed) = &policy.packed_alu { - ensure!( - packed.native_minimum_sm == imported_sm, - "{} native minimum SM {} disagrees with selected instruction predicate sm_{}", - policy.id, - packed.native_minimum_sm, - imported_sm - ); - } else { - let overlay_target = parse_hardware_target(policy)?; - let target_matches = overlay_target - == CatalogHardwareTarget::AnyOf { - alternatives: vec![CatalogHardwareAlternative::MinimumSm { sm: imported_sm }], - }; - ensure!( - target_matches, - "{} minimum SM {:?} disagrees with selected instruction predicate sm_{}", - policy.id, - policy.minimum_sm, - imported_sm - ); - } - } - if policy.family == "ldmatrix" { - if policy - .ldmatrix_variant - .as_ref() - .is_some_and(|variant| variant.shape == LdmatrixShape::M8n8) - { - ensure!( - imported_ptx.is_some() && imported_sm.is_some() && selection.predicates.len() == 2, - "{} classic ldmatrix selection must carry exactly its PTX and SM predicates", - policy.id - ); - } else { - ensure!( - has_ldstmatrix_blackwell_support - && selection.predicates.len() == 1 - && policy.targets == BLACKWELL_LDMATRIX_LLVM_TARGETS, - "{} Blackwell ldmatrix selection must retain its helper predicate and reviewed exact targets", - policy.id - ); - } - } else if policy.family == "dotprod" { - ensure!( - has_dot_instructions && selection.predicates.len() == 1, - "{} dotprod selection must carry only the hasDotInstructions predicate", - policy.id - ); - } else if policy.family == "clc" { - match policy.clc.as_ref().map(|clc| clc.operation) { - Some(ClcOperation::TryCancel) => ensure!( - imported_ptx.is_some() && imported_sm.is_some() && selection.predicates.len() == 2, - "{} selection must carry exactly its PTX and SM predicates", - policy.id - ), - Some(ClcOperation::TryCancelMulticast) => ensure!( - has_clc_multicast_support - && selection.predicates.len() == 1 - && parse_hardware_target(policy)? - == CatalogHardwareTarget::AnyOf { - alternatives: vec![ - CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 101 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 103 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 120 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 121 }, - ], - }, - "{} multicast target predicate must map to the reviewed exact architectures", - policy.id - ), - _ => bail!( - "{} query operation unexpectedly has an instruction selection", - policy.id - ), - } - } else if policy.family == "tma" { - if policy.tma.as_ref().is_some_and(|tma| { - matches!( - tma.operation, - TmaOperation::G2sTile2dMulticastCg2 - | TmaOperation::PrefetchTileGather4TwoDimensional - | TmaOperation::PrefetchTileGather4TwoDimensionalCacheHint - ) - }) { - ensure!( - has_tma_blackwell_support - && selection.predicates.len() == 1 - && parse_hardware_target(policy)? - == CatalogHardwareTarget::AnyOf { - alternatives: vec![ - CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 101 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 103 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, - ], - }, - "{} Blackwell TMA predicate must map to the reviewed exact architectures", - policy.id - ); - } else { - ensure!( - imported_ptx.is_some() && imported_sm.is_some() && selection.predicates.len() == 2, - "{} TMA selection must carry exactly its PTX and SM predicates", - policy.id - ); - } - } else if policy.family == "tcgen05" { - let shift = policy.tcgen05.as_ref().is_some_and(|tcgen05| { - matches!( - tcgen05.operation, - Tcgen05Operation::ShiftDown | Tcgen05Operation::ShiftDownCg2 - ) - }); - ensure!( - ((shift && has_tcgen05_shift_support && !has_tcgen05_support) - || (!shift && has_tcgen05_support && !has_tcgen05_shift_support)) - && selection.predicates.len() == 1 - && parse_hardware_target(policy)? - == CatalogHardwareTarget::AnyOf { - alternatives: vec![ - CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 101 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 103 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, - ], - }, - "{} tcgen05 selection must use the reviewed accelerated architectures", - policy.id - ); - } else if matches!( - policy.family.as_str(), - "vote" - | "active_mask" - | "warp_match" - | "elect" - | "warp_barrier" - | "warp_shuffle" - | "cp_async_copy" - | "cp_async_control" - | "mbarrier_basic" - | "cluster_barrier" - | "scalar_conversion" - | "grid_dependency" - ) { - ensure!( - imported_ptx.is_some() && imported_sm.is_some() && selection.predicates.len() == 2, - "{} selection must carry exactly its PTX and SM predicates", - policy.id - ); - } - Ok(()) -} - -const REVIEWED_SPECIAL_REGISTERS: [SpecialRegisterKind; 12] = [ - SpecialRegisterKind::Clock, - SpecialRegisterKind::Clock64, - SpecialRegisterKind::Globaltimer, - SpecialRegisterKind::Envreg1, - SpecialRegisterKind::Envreg2, - SpecialRegisterKind::Smid, - SpecialRegisterKind::Nsmid, - SpecialRegisterKind::Gridid, - SpecialRegisterKind::Warpid, - SpecialRegisterKind::Nwarpid, - SpecialRegisterKind::DynamicSmemSize, - SpecialRegisterKind::TotalSmemSize, -]; - -#[derive(Clone, Copy)] -struct SpecialRegisterRecipe { - kind: SpecialRegisterKind, - id: &'static str, - operation_key: &'static str, - source_record: Option<&'static str>, - llvm_symbol: Option<&'static str>, - rust_module: &'static str, - compatibility_paths: &'static [&'static str], - dialect_op_type: &'static str, - dialect_op_name: &'static str, - register_spelling: &'static str, - observation: SpecialRegisterObservation, - result_width: SpecialRegisterWidth, - ptx_type: SpecialRegisterPtxType, - output_constraint: SpecialRegisterOutputConstraint, - llvm_mechanism: BackendLoweringMechanism, - libnvvm_mechanism: BackendLoweringMechanism, - minimum_ptx: &'static str, - minimum_sm: Option<&'static str>, - execution_scope: &'static str, - ptx_isa_section: &'static str, - ptx_isa_url: &'static str, - selection_record: Option<&'static str>, - selection_asm: Option<&'static str>, - summary: &'static str, -} - -fn special_register_recipe(kind: SpecialRegisterKind) -> SpecialRegisterRecipe { - use BackendLoweringMechanism::{InlinePtx, TypedNvvm}; - use SpecialRegisterKind::*; - use SpecialRegisterObservation::{StablePure, VolatileObservation}; - use SpecialRegisterOutputConstraint::{Register32, Register64}; - use SpecialRegisterPtxType::{B32, U32, U64}; - use SpecialRegisterWidth::{B32 as Width32, B64 as Width64}; - - match kind { - Clock => SpecialRegisterRecipe { - kind, - id: "clock", - operation_key: "debug.clock", - source_record: Some("int_nvvm_read_ptx_sreg_clock"), - llvm_symbol: Some("llvm.nvvm.read.ptx.sreg.clock"), - rust_module: "debug", - compatibility_paths: &["cuda_device::debug::clock"], - dialect_op_type: "ReadPtxSregClockOp", - dialect_op_name: "nvvm.read_ptx_sreg_clock", - register_spelling: "%clock", - observation: VolatileObservation, - result_width: Width32, - ptx_type: U32, - output_constraint: Register32, - llvm_mechanism: TypedNvvm, - libnvvm_mechanism: TypedNvvm, - minimum_ptx: "1.0", - minimum_sm: None, - execution_scope: "sm", - ptx_isa_section: "10.23 Special Registers: %clock, %clock_hi", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-clock-clock-hi", - selection_record: Some("SREG_CLOCK"), - selection_asm: Some("mov.u32 \t$d, %clock;"), - summary: "Samples the current SM's 32-bit clock counter.", - }, - Clock64 => SpecialRegisterRecipe { - kind, - id: "clock64", - operation_key: "debug.clock64", - source_record: Some("int_nvvm_read_ptx_sreg_clock64"), - llvm_symbol: Some("llvm.nvvm.read.ptx.sreg.clock64"), - rust_module: "debug", - compatibility_paths: &["cuda_device::debug::clock64"], - dialect_op_type: "ReadPtxSregClock64Op", - dialect_op_name: "nvvm.read_ptx_sreg_clock64", - register_spelling: "%clock64", - observation: VolatileObservation, - result_width: Width64, - ptx_type: U64, - output_constraint: Register64, - llvm_mechanism: TypedNvvm, - libnvvm_mechanism: TypedNvvm, - minimum_ptx: "2.0", - minimum_sm: Some("sm_20"), - execution_scope: "sm", - ptx_isa_section: "10.24 Special Registers: %clock64", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-clock64", - selection_record: Some("SREG_CLOCK64"), - selection_asm: Some("mov.u64 \t$d, %clock64;"), - summary: "Samples the current SM's 64-bit clock counter.", - }, - Globaltimer => SpecialRegisterRecipe { - kind, - id: "globaltimer", - operation_key: "debug.global_timer", - source_record: Some("int_nvvm_read_ptx_sreg_globaltimer"), - llvm_symbol: Some("llvm.nvvm.read.ptx.sreg.globaltimer"), - rust_module: "debug", - compatibility_paths: &["cuda_device::debug::globaltimer"], - dialect_op_type: "ReadPtxSregGlobaltimerOp", - dialect_op_name: "nvvm.read_ptx_sreg_globaltimer", - register_spelling: "%globaltimer", - observation: VolatileObservation, - result_width: Width64, - ptx_type: U64, - output_constraint: Register64, - llvm_mechanism: TypedNvvm, - libnvvm_mechanism: InlinePtx, - minimum_ptx: "3.1", - minimum_sm: Some("sm_30"), - execution_scope: "device", - ptx_isa_section: "10.28 Special Registers: %globaltimer, %globaltimer_lo, %globaltimer_hi", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-globaltimer", - selection_record: Some("SREG_GLOBALTIMER"), - selection_asm: Some("mov.u64 \t$d, %globaltimer;"), - summary: "Samples the device-wide 64-bit global timer.", - }, - Envreg1 => SpecialRegisterRecipe { - kind, - id: "envreg1", - operation_key: "grid.environment_register.1", - source_record: Some("int_nvvm_read_ptx_sreg_envreg1"), - llvm_symbol: Some("llvm.nvvm.read.ptx.sreg.envreg1"), - rust_module: "grid", - compatibility_paths: &["cuda_device::grid::envreg1"], - dialect_op_type: "ReadPtxSregEnvReg1Op", - dialect_op_name: "nvvm.read_ptx_sreg_envreg1", - register_spelling: "%envreg1", - observation: StablePure, - result_width: Width32, - ptx_type: B32, - output_constraint: Register32, - llvm_mechanism: TypedNvvm, - libnvvm_mechanism: TypedNvvm, - minimum_ptx: "2.1", - minimum_sm: None, - execution_scope: "grid", - ptx_isa_section: "10.27 Special Registers: %envreg<32>", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-envreg", - selection_record: None, - selection_asm: None, - summary: "Reads PTX environment register 1.", - }, - Envreg2 => SpecialRegisterRecipe { - kind, - id: "envreg2", - operation_key: "grid.environment_register.2", - source_record: Some("int_nvvm_read_ptx_sreg_envreg2"), - llvm_symbol: Some("llvm.nvvm.read.ptx.sreg.envreg2"), - rust_module: "grid", - compatibility_paths: &["cuda_device::grid::envreg2"], - dialect_op_type: "ReadPtxSregEnvReg2Op", - dialect_op_name: "nvvm.read_ptx_sreg_envreg2", - register_spelling: "%envreg2", - observation: StablePure, - result_width: Width32, - ptx_type: B32, - output_constraint: Register32, - llvm_mechanism: TypedNvvm, - libnvvm_mechanism: TypedNvvm, - minimum_ptx: "2.1", - minimum_sm: None, - execution_scope: "grid", - ptx_isa_section: "10.27 Special Registers: %envreg<32>", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-envreg", - selection_record: None, - selection_asm: None, - summary: "Reads PTX environment register 2.", - }, - Smid => SpecialRegisterRecipe { - kind, - id: "smid", - operation_key: "execution.sm_identifier", - source_record: Some("int_nvvm_read_ptx_sreg_smid"), - llvm_symbol: Some("llvm.nvvm.read.ptx.sreg.smid"), - rust_module: "thread", - compatibility_paths: &["cuda_device::thread::smid", "cuda_device::smid"], - dialect_op_type: "ReadPtxSregSmIdOp", - dialect_op_name: "nvvm.read_ptx_sreg_smid", - register_spelling: "%smid", - observation: VolatileObservation, - result_width: Width32, - ptx_type: U32, - output_constraint: Register32, - llvm_mechanism: InlinePtx, - libnvvm_mechanism: InlinePtx, - minimum_ptx: "1.3", - minimum_sm: None, - execution_scope: "thread", - ptx_isa_section: "10.8 Special Registers: %smid", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-smid", - selection_record: Some("SREG_SMID"), - selection_asm: Some("mov.u32 \t$d, %smid;"), - summary: "Samples the SM currently executing this thread.", - }, - Nsmid => SpecialRegisterRecipe { - kind, - id: "nsmid", - operation_key: "execution.sm_identifier_bound", - source_record: Some("int_nvvm_read_ptx_sreg_nsmid"), - llvm_symbol: Some("llvm.nvvm.read.ptx.sreg.nsmid"), - rust_module: "thread", - compatibility_paths: &["cuda_device::thread::nsmid", "cuda_device::nsmid"], - dialect_op_type: "ReadPtxSregNsmIdOp", - dialect_op_name: "nvvm.read_ptx_sreg_nsmid", - register_spelling: "%nsmid", - observation: StablePure, - result_width: Width32, - ptx_type: U32, - output_constraint: Register32, - llvm_mechanism: TypedNvvm, - libnvvm_mechanism: TypedNvvm, - minimum_ptx: "2.0", - minimum_sm: Some("sm_20"), - execution_scope: "device", - ptx_isa_section: "10.9 Special Registers: %nsmid", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-nsmid", - selection_record: Some("SREG_NSMID"), - selection_asm: Some("mov.u32 \t$d, %nsmid;"), - summary: "Returns the upper bound for SM identifiers.", - }, - Gridid => SpecialRegisterRecipe { - kind, - id: "gridid", - operation_key: "launch.grid_identifier", - source_record: None, - llvm_symbol: None, - rust_module: "thread", - compatibility_paths: &["cuda_device::thread::gridid", "cuda_device::gridid"], - dialect_op_type: "ReadPtxSregGridIdOp", - dialect_op_name: "nvvm.read_ptx_sreg_gridid", - register_spelling: "%gridid", - observation: StablePure, - result_width: Width64, - ptx_type: U64, - output_constraint: Register64, - llvm_mechanism: InlinePtx, - libnvvm_mechanism: InlinePtx, - minimum_ptx: "3.0", - minimum_sm: Some("sm_30"), - execution_scope: "grid", - ptx_isa_section: "10.10 Special Registers: %gridid", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-gridid", - selection_record: None, - selection_asm: None, - summary: "Returns the full 64-bit temporal grid identifier.", - }, - Warpid => SpecialRegisterRecipe { - kind, - id: "warpid", - operation_key: "warp.hardware_identifier", - source_record: Some("int_nvvm_read_ptx_sreg_warpid"), - llvm_symbol: Some("llvm.nvvm.read.ptx.sreg.warpid"), - rust_module: "warp", - compatibility_paths: &["cuda_device::warp::warpid"], - dialect_op_type: "ReadPtxSregWarpIdOp", - dialect_op_name: "nvvm.read_ptx_sreg_warpid", - register_spelling: "%warpid", - observation: VolatileObservation, - result_width: Width32, - ptx_type: U32, - output_constraint: Register32, - llvm_mechanism: InlinePtx, - libnvvm_mechanism: InlinePtx, - minimum_ptx: "1.3", - minimum_sm: None, - execution_scope: "cta", - ptx_isa_section: "10.4 Special Registers: %warpid", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-warpid", - selection_record: Some("SREG_WARPID"), - selection_asm: Some("mov.u32 \t$d, %warpid;"), - summary: "Samples the hardware warp currently executing this thread.", - }, - Nwarpid => SpecialRegisterRecipe { - kind, - id: "nwarpid", - operation_key: "warp.hardware_identifier_bound", - source_record: Some("int_nvvm_read_ptx_sreg_nwarpid"), - llvm_symbol: Some("llvm.nvvm.read.ptx.sreg.nwarpid"), - rust_module: "warp", - compatibility_paths: &["cuda_device::warp::nwarpid"], - dialect_op_type: "ReadPtxSregNwarpIdOp", - dialect_op_name: "nvvm.read_ptx_sreg_nwarpid", - register_spelling: "%nwarpid", - observation: StablePure, - result_width: Width32, - ptx_type: U32, - output_constraint: Register32, - llvm_mechanism: TypedNvvm, - libnvvm_mechanism: TypedNvvm, - minimum_ptx: "2.0", - minimum_sm: Some("sm_20"), - execution_scope: "cta", - ptx_isa_section: "10.5 Special Registers: %nwarpid", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-nwarpid", - selection_record: Some("SREG_NWARPID"), - selection_asm: Some("mov.u32 \t$d, %nwarpid;"), - summary: "Returns the upper bound for hardware warp identifiers.", - }, - DynamicSmemSize => SpecialRegisterRecipe { - kind, - id: "dynamic_smem_size", - operation_key: "shared.dynamic_size", - source_record: Some("int_nvvm_read_ptx_sreg_dynamic_smem_size"), - llvm_symbol: Some("llvm.nvvm.read.ptx.sreg.dynamic_smem_size"), - rust_module: "shared", - compatibility_paths: &["cuda_device::shared::dynamic_smem_size"], - dialect_op_type: "ReadPtxSregDynamicSmemSizeOp", - dialect_op_name: "nvvm.read_ptx_sreg_dynamic_smem_size", - register_spelling: "%dynamic_smem_size", - observation: StablePure, - result_width: Width32, - ptx_type: U32, - output_constraint: Register32, - llvm_mechanism: InlinePtx, - libnvvm_mechanism: InlinePtx, - minimum_ptx: "4.1", - minimum_sm: Some("sm_20"), - execution_scope: "cta", - ptx_isa_section: "10.32 Special Registers: %dynamic_smem_size", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-dynamic-smem-size", - selection_record: Some("INT_PTX_SREG_DYNAMIC_SMEM_SIZE"), - selection_asm: Some("mov.u32 \t$d, %dynamic_smem_size;"), - summary: "Returns the launch-time dynamic shared-memory size in bytes.", - }, - TotalSmemSize => SpecialRegisterRecipe { - kind, - id: "total_smem_size", - operation_key: "shared.total_size", - source_record: Some("int_nvvm_read_ptx_sreg_total_smem_size"), - llvm_symbol: Some("llvm.nvvm.read.ptx.sreg.total_smem_size"), - rust_module: "shared", - compatibility_paths: &["cuda_device::shared::total_smem_size"], - dialect_op_type: "ReadPtxSregTotalSmemSizeOp", - dialect_op_name: "nvvm.read_ptx_sreg_total_smem_size", - register_spelling: "%total_smem_size", - observation: StablePure, - result_width: Width32, - ptx_type: U32, - output_constraint: Register32, - llvm_mechanism: InlinePtx, - libnvvm_mechanism: InlinePtx, - minimum_ptx: "4.1", - minimum_sm: Some("sm_20"), - execution_scope: "cta", - ptx_isa_section: "10.30 Special Registers: %total_smem_size", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-total-smem-size", - selection_record: Some("INT_PTX_SREG_TOTAL_SMEM_SIZE"), - selection_asm: Some("mov.u32 \t$d, %total_smem_size;"), - summary: "Returns the total user shared-memory allocation in bytes.", - }, - } -} - -fn special_register_ptx_type(ptx_type: SpecialRegisterPtxType) -> &'static str { - match ptx_type { - SpecialRegisterPtxType::B32 => "b32", - SpecialRegisterPtxType::U32 => "u32", - SpecialRegisterPtxType::U64 => "u64", - } -} - -fn special_register_backend_floor( - recipe: SpecialRegisterRecipe, - backend: IntrinsicBackend, -) -> (Option<&'static str>, Option<&'static str>) { - match backend { - IntrinsicBackend::LlvmNvptx => { - let minimum_ptx = if matches!(recipe.minimum_ptx, "4.1") { - "4.1" - } else { - "3.2" - }; - let minimum_sm = if recipe.minimum_sm == Some("sm_30") { - "sm_30" - } else { - "sm_20" - }; - (Some(minimum_ptx), Some(minimum_sm)) - } - IntrinsicBackend::LibNvvm => (None, Some("sm_75")), - } -} - -fn special_register_contract(recipe: SpecialRegisterRecipe) -> SpecialRegister { - let llvm_exclusion = - (recipe.kind == SpecialRegisterKind::Gridid).then(|| SpecialRegisterLlvmExclusion { - source_record: "int_nvvm_read_ptx_sreg_gridid".into(), - llvm_symbol: "llvm.nvvm.read.ptx.sreg.gridid".into(), - imported_result_width: SpecialRegisterWidth::B32, - reason: SpecialRegisterLlvmExclusionReason::ResultWidthMismatch, - }); - SpecialRegister { - register: recipe.kind, - observation: recipe.observation, - result_width: recipe.result_width, - ptx_type: recipe.ptx_type, - output_constraint: recipe.output_constraint, - llvm_exclusion, - } -} - -fn expand_special_register_admission( - admission: &SpecialRegisterAdmission, -) -> Result> { - ensure!( - admission.runtime_validation == RuntimeValidation::Unexecuted, - "special-register runtime validation may be marked executed only with GPU evidence" - ); - ensure!( - !admission.llvm_evidence_profile.trim().is_empty() - && !admission.libnvvm_evidence_profile.trim().is_empty(), - "compact special-register admission requires both backend evidence profiles" - ); - ensure!( - admission.registers == REVIEWED_SPECIAL_REGISTERS - && admission.product_count == REVIEWED_SPECIAL_REGISTERS.len(), - "compact special-register admission must list the canonical 12 registers exactly once and in order" - ); - - admission - .registers - .iter() - .copied() - .map(|kind| { - let recipe = special_register_recipe(kind); - let width = recipe.result_width.bits(); - let rust_result = format!("u{width}"); - let dialect_result = format!("i{width}"); - let source = match recipe.source_record { - Some(_) => None, - None => Some(IntrinsicSource::PtxNative { - instruction: format!( - "mov.{} {}", - special_register_ptx_type(recipe.ptx_type), - recipe.register_spelling - ), - }), - }; - Ok(OverlayIntrinsic { - id: recipe.id.into(), - abi_id: String::new(), - operation_key: recipe.operation_key.into(), - family: "sreg".into(), - source, - source_record: recipe.source_record.map(str::to_owned), - rust_module: recipe.rust_module.into(), - rust_name: recipe.id.into(), - rust_arguments: vec![], - rust_result: rust_result.clone(), - safe: true, - must_use: false, - safe_allowlist_reason: Some( - "reading this special register has no caller obligations.".into(), - ), - public_rust_path: format!("cuda_intrinsics::{}::{}", recipe.rust_module, recipe.id), - compatibility_rust_paths: recipe - .compatibility_paths - .iter() - .map(|path| (*path).into()) - .collect(), - dialect_op_type: recipe.dialect_op_type.into(), - dialect_op_name: recipe.dialect_op_name.into(), - dialect_operands: vec![], - dialect_results: vec![dialect_result.clone()], - llvm_symbol: recipe.llvm_symbol.map(str::to_owned), - resolved_llvm_symbol: None, - llvm_arguments: vec![], - llvm_results: recipe - .source_record - .map(|_| dialect_result) - .into_iter() - .collect(), - pure: recipe.observation == SpecialRegisterObservation::StablePure, - memory: if matches!( - recipe.kind, - SpecialRegisterKind::Clock - | SpecialRegisterKind::Clock64 - | SpecialRegisterKind::Globaltimer - ) { - "inaccessible_read_write".into() - } else { - "none".into() - }, - convergent: false, - execution_scope: recipe.execution_scope.into(), - minimum_ptx: recipe.minimum_ptx.into(), - minimum_sm: recipe.minimum_sm.map(str::to_owned), - ptx_result: rust_result, - targets: "all".into(), - ptx_isa_version: "9.3".into(), - ptx_isa_section: recipe.ptx_isa_section.into(), - ptx_isa_url: recipe.ptx_isa_url.into(), - lowering: "generated_special_register".into(), - backend_lowerings: [ - ( - IntrinsicBackend::LlvmNvptx, - recipe.llvm_mechanism, - &admission.llvm_evidence_profile, - ), - ( - IntrinsicBackend::LibNvvm, - recipe.libnvvm_mechanism, - &admission.libnvvm_evidence_profile, - ), - ] - .into_iter() - .map( - |(backend, mechanism, evidence_profile)| OverlayBackendLowering { - minimum_ptx: special_register_backend_floor(recipe, backend) - .0 - .map(str::to_owned), - minimum_sm: special_register_backend_floor(recipe, backend) - .1 - .map(str::to_owned), - backend, - mechanism, - evidence_profile: evidence_profile.clone(), - targets: None, - }, - ) - .collect(), - packed_atomic: None, - redux: None, - vote: None, - active_mask: None, - warp_match: None, - warp_barrier: None, - warp_shuffle: None, - dot_product: None, - packed_alu: None, - integer_minmax: None, - packed_conversion: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cp_async_copy: None, - cp_async_control: None, - cp_async_mbarrier: None, - mbarrier_basic: None, - movmatrix: None, - mbarrier_extended: None, - register_mma: None, - sparse_mma: None, - prmt: None, - cluster_barrier: None, - wgmma_control: None, - special_register: Some(special_register_contract(recipe)), - debug_control: None, - cluster_memory: None, - clc: None, - tma: None, - tcgen05: None, - ldmatrix_variant: None, - ldmatrix_safety: None, - ldmatrix_adapter: None, - selected_address_space: None, - expected_ptx: InstructionPattern { - mnemonic: "mov".into(), - modifiers: vec![special_register_ptx_type(recipe.ptx_type).into()], - operands: vec![ - OperandPattern::Register, - OperandPattern::Exact { - value: recipe.register_spelling.into(), - }, - ], - }, - summary: recipe.summary.into(), - }) - }) - .collect() -} - -fn validate_special_register_policy( - policy: &OverlayIntrinsic, - source: &IntrinsicSource, - declaration: Option<&ImportedIntrinsic>, -) -> Result<()> { - let special = policy - .special_register - .as_ref() - .with_context(|| format!("{} has no closed special-register contract", policy.id))?; - let recipe = special_register_recipe(special.register); - ensure!( - special == &special_register_contract(recipe), - "{} special-register width, PTX type, constraint, observation, or LLVM exclusion changed", - policy.id - ); - let expected_source = match recipe.source_record { - Some(source_record) => IntrinsicSource::LlvmImported { - source_record: source_record.into(), - }, - None => IntrinsicSource::PtxNative { - instruction: format!( - "mov.{} {}", - special_register_ptx_type(recipe.ptx_type), - recipe.register_spelling - ), - }, - }; - ensure!( - policy.id == recipe.id - && policy.operation_key == recipe.operation_key - && source == &expected_source - && policy.source_record.as_deref() == recipe.source_record - && policy.llvm_symbol.as_deref() == recipe.llvm_symbol - && policy.resolved_llvm_symbol.is_none() - && policy.llvm_arguments.is_empty(), - "{} special-register identity or source changed", - policy.id - ); - let width = recipe.result_width.bits(); - let rust_result = format!("u{width}"); - let dialect_result = format!("i{width}"); - let expected_llvm_results = recipe - .source_record - .map(|_| dialect_result.clone()) - .into_iter() - .collect::>(); - ensure!( - policy.rust_module == recipe.rust_module - && policy.rust_name == recipe.id - && policy.rust_arguments.is_empty() - && policy.rust_result == rust_result - && policy.safe - && !policy.must_use - && policy.public_rust_path - == format!("cuda_intrinsics::{}::{}", recipe.rust_module, recipe.id) - && policy.compatibility_rust_paths - == recipe - .compatibility_paths - .iter() - .map(|path| (*path).to_owned()) - .collect::>(), - "{} special-register Rust API changed", - policy.id - ); - ensure!( - policy.dialect_op_type == recipe.dialect_op_type - && policy.dialect_op_name == recipe.dialect_op_name - && policy.dialect_operands.is_empty() - && policy.dialect_results == [dialect_result.as_str()] - && policy.llvm_results == expected_llvm_results - && policy.lowering == "generated_special_register", - "{} special-register carrier, result width, or lowering changed", - policy.id - ); - let expected_pure = recipe.observation == SpecialRegisterObservation::StablePure; - let expected_memory = if matches!( - recipe.kind, - SpecialRegisterKind::Clock - | SpecialRegisterKind::Clock64 - | SpecialRegisterKind::Globaltimer - ) { - "inaccessible_read_write" - } else { - "none" - }; - ensure!( - policy.pure == expected_pure - && policy.memory == expected_memory - && !policy.convergent - && policy.execution_scope == recipe.execution_scope, - "{} special-register observation or effects changed", - policy.id - ); - ensure!( - policy.minimum_ptx == recipe.minimum_ptx - && policy.minimum_sm.as_deref() == recipe.minimum_sm - && policy.ptx_result == rust_result - && policy.targets == "all" - && policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section == recipe.ptx_isa_section - && policy.ptx_isa_url == recipe.ptx_isa_url, - "{} special-register target floor or PTX provenance changed", - policy.id - ); - ensure!( - policy.expected_ptx - == InstructionPattern { - mnemonic: "mov".into(), - modifiers: vec![special_register_ptx_type(recipe.ptx_type).into()], - operands: vec![ - OperandPattern::Register, - OperandPattern::Exact { - value: recipe.register_spelling.into(), - }, - ], - }, - "{} special-register PTX shape changed", - policy.id - ); - let expected_routes = [ - ( - IntrinsicBackend::LlvmNvptx, - recipe.llvm_mechanism, - special_register_backend_floor(recipe, IntrinsicBackend::LlvmNvptx), - ), - ( - IntrinsicBackend::LibNvvm, - recipe.libnvvm_mechanism, - special_register_backend_floor(recipe, IntrinsicBackend::LibNvvm), - ), - ]; - ensure!( - policy.backend_lowerings.len() == expected_routes.len() - && policy.backend_lowerings.iter().zip(expected_routes).all( - |(actual, (backend, mechanism, (minimum_ptx, minimum_sm)))| { - actual.backend == backend - && actual.mechanism == mechanism - && !actual.evidence_profile.trim().is_empty() - && actual.minimum_ptx.as_deref() == minimum_ptx - && actual.minimum_sm.as_deref() == minimum_sm - } - ), - "{} special-register backend routes changed", - policy.id - ); - ensure!( - policy.ldmatrix_variant.is_none() - && policy.ldmatrix_safety.is_none() - && policy.ldmatrix_adapter.is_none() - && policy.packed_atomic.is_none() - && policy.redux.is_none() - && policy.vote.is_none() - && policy.active_mask.is_none() - && policy.warp_match.is_none() - && policy.warp_barrier.is_none() - && policy.warp_shuffle.is_none() - && policy.dot_product.is_none() - && policy.packed_alu.is_none() - && policy.packed_conversion.is_none() - && policy.cp_async_copy.is_none() - && policy.cp_async_control.is_none() - && policy.cp_async_mbarrier.is_none() - && policy.mbarrier_basic.is_none() - && policy.register_mma.is_none() - && policy.sparse_mma.is_none() - && policy.prmt.is_none() - && policy.selected_address_space.is_none(), - "{} mixes another generated-family contract with a special register", - policy.id - ); - - match (recipe.source_record, declaration) { - (None, None) => {} - (Some(_), Some(declaration)) => { - let timer = matches!( - recipe.kind, - SpecialRegisterKind::Clock - | SpecialRegisterKind::Clock64 - | SpecialRegisterKind::Globaltimer - ); - let expected_properties: &[&str] = if timer { - &[ - "IntrInaccessibleMemOnly", - "IntrNoCallback", - "IntrNoFree", - "IntrWillReturn", - "NoUndef", - ] - } else { - &["IntrNoMem", "IntrSpeculatable", "NoUndef"] - }; - ensure!( - declaration.arguments.is_empty() - && declaration.results == [dialect_result.as_str()] - && declaration.properties - == expected_properties - .iter() - .map(|property| (*property).to_owned()) - .collect::>() - && if timer { - declaration - .classes - .iter() - .any(|class| class == "PTXReadNCSRegIntrinsic") - } else { - declaration - .classes - .iter() - .any(|class| class == "NVVMPureIntrinsic") - }, - "{} imported special-register signature, class, or properties changed", - policy.id - ); - match (recipe.selection_record, recipe.selection_asm) { - (None, None) => ensure!( - declaration.selections.is_empty(), - "{} selectionless environment-register contract changed", - policy.id - ), - (Some(selection_record), Some(selection_asm)) => ensure!( - declaration.selections.len() == 1 - && declaration.selections[0].source_record == selection_record - && declaration.selections[0].asm == selection_asm - && declaration.selections[0].predicates.is_empty() - && declaration.selections[0].constraints.is_empty(), - "{} imported special-register selection changed", - policy.id - ), - _ => unreachable!("closed special-register selection recipe"), - } - } - _ => bail!( - "{} special-register source and imported declaration disagree", - policy.id - ), - } - Ok(()) -} - -fn validate_special_register_llvm_exclusion( - policy: &OverlayIntrinsic, - imported_by_record: &BTreeMap<&str, &ImportedIntrinsic>, -) -> Result<()> { - let Some(exclusion) = policy - .special_register - .as_ref() - .and_then(|special| special.llvm_exclusion.as_ref()) - else { - return Ok(()); - }; - let declaration = imported_by_record - .get(exclusion.source_record.as_str()) - .with_context(|| { - format!( - "{} excludes missing imported LLVM record {}", - policy.id, exclusion.source_record - ) - })?; - ensure!( - policy.id == "gridid" - && exclusion.reason == SpecialRegisterLlvmExclusionReason::ResultWidthMismatch - && exclusion.imported_result_width == SpecialRegisterWidth::B32 - && declaration.llvm_name == exclusion.llvm_symbol - && declaration.arguments.is_empty() - && declaration.results == ["i32"] - && declaration.properties == ["IntrNoMem", "IntrSpeculatable", "NoUndef"] - && declaration.selections.len() == 1 - && declaration.selections[0].source_record == "SREG_GRIDID" - && declaration.selections[0].asm == "mov.u32 \t$d, %gridid;" - && declaration.selections[0].predicates.is_empty() - && declaration.selections[0].constraints.is_empty() - && policy.rust_result == "u64" - && policy.dialect_results == ["i64"], - "{} LLVM exclusion no longer proves the reviewed i32-to-u64 gridid width mismatch", - policy.id - ); - Ok(()) -} - -fn validate_sreg_policy( - policy: &OverlayIntrinsic, - source: &IntrinsicSource, - declaration: Option<&ImportedIntrinsic>, -) -> Result<()> { - if policy.special_register.is_some() { - return validate_special_register_policy(policy, source, declaration); - } - let declaration = declaration.context("sreg requires imported LLVM declaration")?; - ensure!( - policy.rust_arguments.is_empty() && policy.llvm_arguments.is_empty(), - "{} is not a zero-operand intrinsic; the sreg recipe cannot lower it", - policy.id - ); - ensure!( - matches!(policy.rust_result.as_str(), "u32" | "u64"), - "{} has unsupported raw scalar result {}", - policy.id, - policy.rust_result - ); - let expected_llvm_result = match policy.rust_result.as_str() { - "u32" => "i32", - "u64" => "i64", - _ => unreachable!(), - }; - ensure!( - policy.llvm_results == [expected_llvm_result] - && policy.ptx_result == policy.rust_result - && policy.lowering == "direct_nvvm", - "{} has a signature or lowering outside the scalar direct-NVVM sreg recipe", - policy.id - ); - ensure!( - policy.resolved_llvm_symbol.is_none() - && policy.backend_lowerings.is_empty() - && policy.special_register.is_none(), - "{} uses a backend contract outside the direct-NVVM sreg recipe", - policy.id - ); - ensure_no_other_family_contract(policy, "sreg")?; - if policy.id.starts_with("lanemask_") { - validate_lanemask_policy(policy, declaration)?; - } - if is_cluster_sreg_source(&declaration.source_record) { - validate_cluster_sreg_policy(policy, declaration)?; - } - Ok(()) -} - -fn is_cluster_sreg_source(source_record: &str) -> bool { - source_record == "int_nvvm_read_ptx_sreg_cluster_ctarank" - || source_record == "int_nvvm_read_ptx_sreg_cluster_nctarank" - || source_record.starts_with("int_nvvm_read_ptx_sreg_cluster_ctaid_") - || source_record.starts_with("int_nvvm_read_ptx_sreg_cluster_nctaid_") - || source_record.starts_with("int_nvvm_read_ptx_sreg_clusterid_") - || source_record.starts_with("int_nvvm_read_ptx_sreg_nclusterid_") -} - -#[derive(Clone)] -struct ClusterSregRecipe { - id: String, - operation_key: String, - source_suffix: String, - llvm_suffix: String, - selection_record: String, - ptx_register: String, - compatibility_path: Option, - op_type: String, - scope: &'static str, - section: &'static str, - anchor: &'static str, - range: Option<&'static str>, - safe_reason: String, - summary: String, -} - -#[derive(Clone, Copy)] -struct ClusterSregXyzFamilyRecipe { - id_prefix: &'static str, - operation_key_prefix: &'static str, - source_prefix: &'static str, - llvm_prefix: &'static str, - selection_prefix: &'static str, - ptx_prefix: &'static str, - compatibility_prefix: Option<&'static str>, - op_type_prefix: &'static str, - scope: &'static str, - section: &'static str, - anchor: &'static str, - x_range: &'static str, - yz_range: &'static str, - safe_reason: &'static str, - summary: &'static str, -} - -const CLUSTER_SREG_AXES: [&str; 3] = ["x", "y", "z"]; - -const CLUSTER_SREG_XYZ_FAMILIES: [ClusterSregXyzFamilyRecipe; 4] = [ - ClusterSregXyzFamilyRecipe { - id_prefix: "cluster_block_idx", - operation_key_prefix: "launch.cluster.block_index", - source_prefix: "cluster_ctaid_", - llvm_prefix: "cluster.ctaid.", - selection_prefix: "INT_PTX_SREG_CLUSTER_CTAID_", - ptx_prefix: "%cluster_ctaid.", - compatibility_prefix: Some("cuda_device::cluster::cluster_ctaid"), - op_type_prefix: "ReadPtxSregClusterCtaid", - scope: "cta", - section: "10.14 Special Registers: %cluster_ctaid", - anchor: "cluster-ctaid", - x_range: "Range", - yz_range: "Range", - safe_reason: "reading the read-only block index within its cluster has no caller obligations", - summary: "Returns the block's {axis} index within its thread block cluster.", - }, - ClusterSregXyzFamilyRecipe { - id_prefix: "cluster_dim", - operation_key_prefix: "launch.cluster.dimension", - source_prefix: "cluster_nctaid_", - llvm_prefix: "cluster.nctaid.", - selection_prefix: "INT_PTX_SREG_CLUSTER_NCTAID_", - ptx_prefix: "%cluster_nctaid.", - compatibility_prefix: Some("cuda_device::cluster::cluster_nctaid"), - op_type_prefix: "ReadPtxSregClusterNctaid", - scope: "cluster", - section: "10.15 Special Registers: %cluster_nctaid", - anchor: "cluster-nctaid", - x_range: "Range", - yz_range: "Range", - safe_reason: "reading the read-only cluster dimension has no caller obligations", - summary: "Returns the number of blocks in the cluster's {axis} dimension.", - }, - ClusterSregXyzFamilyRecipe { - id_prefix: "cluster_idx", - operation_key_prefix: "launch.cluster.index", - source_prefix: "clusterid_", - llvm_prefix: "clusterid.", - selection_prefix: "INT_PTX_SREG_CLUSTERID_", - ptx_prefix: "%clusterid.", - compatibility_prefix: Some("cuda_device::cluster::__cluster_idx"), - op_type_prefix: "ReadPtxSregClusterId", - scope: "cluster", - section: "10.12 Special Registers: %clusterid", - anchor: "clusterid", - x_range: "Range", - yz_range: "Range", - safe_reason: "reading the read-only cluster index has no caller obligations", - summary: "Returns the cluster's {axis} index within the grid.", - }, - ClusterSregXyzFamilyRecipe { - id_prefix: "cluster_grid_dim", - operation_key_prefix: "launch.cluster.grid_dimension", - source_prefix: "nclusterid_", - llvm_prefix: "nclusterid.", - selection_prefix: "INT_PTX_SREG_NCLUSTERID_", - ptx_prefix: "%nclusterid.", - compatibility_prefix: Some("cuda_device::cluster::__cluster_grid_dim"), - op_type_prefix: "ReadPtxSregNclusterId", - scope: "grid", - section: "10.13 Special Registers: %nclusterid", - anchor: "nclusterid", - x_range: "Range", - yz_range: "Range", - safe_reason: "reading the read-only cluster-grid dimension has no caller obligations", - summary: "Returns the number of clusters in the grid's {axis} dimension.", - }, -]; - -fn cluster_sreg_recipes() -> Vec { - let mut recipes = Vec::with_capacity(14); - for family in CLUSTER_SREG_XYZ_FAMILIES { - for axis in CLUSTER_SREG_AXES { - let axis_upper = axis.to_ascii_uppercase(); - recipes.push(ClusterSregRecipe { - id: format!("{}_{axis}", family.id_prefix), - operation_key: format!("{}.{axis}", family.operation_key_prefix), - source_suffix: format!("{}{axis}", family.source_prefix), - llvm_suffix: format!("{}{axis}", family.llvm_prefix), - selection_record: format!("{}{axis}", family.selection_prefix), - ptx_register: format!("{}{axis}", family.ptx_prefix), - compatibility_path: family - .compatibility_prefix - .map(|prefix| format!("{prefix}{axis_upper}")), - op_type: format!("{}{axis_upper}Op", family.op_type_prefix), - scope: family.scope, - section: family.section, - anchor: family.anchor, - range: Some(if axis == "x" { - family.x_range - } else { - family.yz_range - }), - safe_reason: family.safe_reason.into(), - summary: family.summary.replace("{axis}", &axis_upper), - }); - } - } - recipes.extend([ - ClusterSregRecipe { - id: "cluster_block_rank".into(), - operation_key: "launch.cluster.block_rank".into(), - source_suffix: "cluster_ctarank".into(), - llvm_suffix: "cluster.ctarank".into(), - selection_record: "INT_PTX_SREG_CLUSTER_CTARANK".into(), - ptx_register: "%cluster_ctarank".into(), - compatibility_path: None, - op_type: "ReadPtxSregClusterCtarankOp".into(), - scope: "cta", - section: "10.16 Special Registers: %cluster_ctarank", - anchor: "cluster-ctarank", - range: None, - safe_reason: - "reading the read-only block rank within its cluster has no caller obligations" - .into(), - summary: "Returns the block's linear rank within its thread block cluster.".into(), - }, - ClusterSregRecipe { - id: "cluster_block_count".into(), - operation_key: "launch.cluster.block_count".into(), - source_suffix: "cluster_nctarank".into(), - llvm_suffix: "cluster.nctarank".into(), - selection_record: "INT_PTX_SREG_CLUSTER_NCTARANK".into(), - ptx_register: "%cluster_nctarank".into(), - compatibility_path: None, - op_type: "ReadPtxSregClusterNctarankOp".into(), - scope: "cluster", - section: "10.17 Special Registers: %cluster_nctarank", - anchor: "cluster-nctarank", - range: None, - safe_reason: "reading the read-only block count has no caller obligations".into(), - summary: "Returns the total number of blocks in the thread block cluster.".into(), - }, - ]); - recipes -} - -fn expand_cluster_sreg_admission( - admission: &ClusterSregAdmission, -) -> Result> { - ensure!( - admission.axes == CLUSTER_SREG_AXES, - "cluster-sreg axes must be exactly x, y, z" - ); - ensure!( - admission.xyz_product_count == 12 && admission.record_count == 14, - "cluster-sreg admission must expand to 12 xyz and 14 total records" - ); - let recipes = cluster_sreg_recipes(); - ensure!( - recipes.len() == admission.record_count, - "cluster-sreg recipe count disagrees with its admission" - ); - Ok(recipes.into_iter().map(cluster_sreg_policy).collect()) -} - -fn cluster_sreg_policy(recipe: ClusterSregRecipe) -> OverlayIntrinsic { - let compatibility_rust_paths = recipe.compatibility_path.iter().cloned().collect(); - OverlayIntrinsic { - id: recipe.id.clone(), - abi_id: String::new(), - operation_key: recipe.operation_key, - family: "sreg".into(), - source: None, - source_record: Some(format!("int_nvvm_read_ptx_sreg_{}", recipe.source_suffix)), - rust_module: "sreg".into(), - rust_name: recipe.id.clone(), - rust_arguments: vec![], - rust_result: "u32".into(), - safe: true, - must_use: false, - safe_allowlist_reason: Some(recipe.safe_reason), - public_rust_path: format!("cuda_intrinsics::sreg::{}", recipe.id), - compatibility_rust_paths, - dialect_op_type: recipe.op_type, - dialect_op_name: format!("nvvm.read_ptx_sreg_{}", recipe.source_suffix), - dialect_operands: vec![], - dialect_results: vec![], - llvm_symbol: Some(format!("llvm.nvvm.read.ptx.sreg.{}", recipe.llvm_suffix)), - resolved_llvm_symbol: None, - llvm_arguments: vec![], - llvm_results: vec!["i32".into()], - pure: true, - memory: "none".into(), - convergent: false, - execution_scope: recipe.scope.into(), - minimum_ptx: "7.8".into(), - minimum_sm: Some("sm_90".into()), - ptx_result: "u32".into(), - targets: "all".into(), - ptx_isa_version: "9.3".into(), - ptx_isa_section: recipe.section.into(), - ptx_isa_url: format!( - "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-{}", - recipe.anchor - ), - lowering: "direct_nvvm".into(), - backend_lowerings: vec![], - packed_atomic: None, - redux: None, - vote: None, - active_mask: None, - warp_match: None, - warp_barrier: None, - warp_shuffle: None, - dot_product: None, - packed_alu: None, - integer_minmax: None, - packed_conversion: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cp_async_copy: None, - cp_async_control: None, - cp_async_mbarrier: None, - mbarrier_basic: None, - movmatrix: None, - mbarrier_extended: None, - register_mma: None, - sparse_mma: None, - prmt: None, - cluster_barrier: None, - wgmma_control: None, - special_register: None, - debug_control: None, - cluster_memory: None, - clc: None, - tma: None, - tcgen05: None, - ldmatrix_variant: None, - ldmatrix_safety: None, - ldmatrix_adapter: None, - selected_address_space: None, - expected_ptx: InstructionPattern { - mnemonic: "mov".into(), - modifiers: vec!["u32".into()], - operands: vec![ - OperandPattern::Register, - OperandPattern::Exact { - value: recipe.ptx_register, - }, - ], - }, - summary: recipe.summary, - } -} - -fn validate_cluster_sreg_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, -) -> Result<()> { - ensure!( - !declaration.source_record.ends_with("_w"), - "{} selects an unused always-zero fourth cluster-register component", - policy.id - ); - let recipe = cluster_sreg_recipes() - .into_iter() - .find(|recipe| recipe.id == policy.id) - .with_context(|| format!("{} is not a reviewed cluster special register", policy.id))?; - - let source_record = format!("int_nvvm_read_ptx_sreg_{}", recipe.source_suffix); - let llvm_symbol = format!("llvm.nvvm.read.ptx.sreg.{}", recipe.llvm_suffix); - let compatibility_paths = recipe - .compatibility_path - .iter() - .cloned() - .collect::>(); - ensure!( - policy.operation_key == recipe.operation_key - && policy.source.is_none() - && policy.source_record.as_deref() == Some(source_record.as_str()) - && policy.llvm_symbol.as_deref() == Some(llvm_symbol.as_str()) - && policy.resolved_llvm_symbol.is_none(), - "{} cluster-register identity does not match its closed recipe", - policy.id - ); - ensure!( - policy.rust_module == "sreg" - && policy.rust_name == policy.id - && policy.rust_arguments.is_empty() - && policy.rust_result == "u32" - && policy.safe - && !policy.must_use - && policy - .safe_allowlist_reason - .as_deref() - .is_some_and(|reason| !reason.is_empty()) - && policy.public_rust_path == format!("cuda_intrinsics::sreg::{}", policy.id) - && policy.compatibility_rust_paths == compatibility_paths, - "{} must preserve its reviewed raw and compatibility APIs", - policy.id - ); - ensure!( - policy.dialect_op_type == recipe.op_type - && policy.dialect_op_name == format!("nvvm.read_ptx_sreg_{}", recipe.source_suffix) - && policy.dialect_operands.is_empty() - && policy.dialect_results.is_empty() - && policy.llvm_arguments.is_empty() - && policy.llvm_results == ["i32"] - && policy.lowering == "direct_nvvm", - "{} is outside the closed cluster-register lowering recipe", - policy.id - ); - ensure!( - policy.pure - && policy.memory == "none" - && !policy.convergent - && policy.execution_scope == recipe.scope - && policy.minimum_ptx == "7.8" - && policy.minimum_sm.as_deref() == Some("sm_90") - && policy.ptx_result == "u32" - && policy.targets == "all", - "{} cluster-register effects or target floor disagree with PTX", - policy.id - ); - ensure!( - policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section == recipe.section - && policy.ptx_isa_url - == format!( - "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-{}", - recipe.anchor - ), - "{} cluster-register PTX provenance disagrees with the reviewed recipe", - policy.id - ); - let mut properties = vec!["IntrNoMem", "IntrSpeculatable", "NoUndef"]; - properties.extend(recipe.range); - ensure!( - declaration.arguments.is_empty() - && declaration.results == ["i32"] - && declaration.classes - == [ - "SDPatternOperator", - "Intrinsic", - "DefaultAttrsIntrinsic", - "NVVMPureIntrinsic", - "PTXReadSRegIntrinsicNB_r32", - ] - && declaration.properties == properties, - "{} declaration shape or properties disagree with LLVM TableGen", - policy.id - ); - let [selection] = declaration.selections.as_slice() else { - bail!("{} must have exactly one LLVM selection", policy.id); - }; - ensure!( - selection.source_record == recipe.selection_record - && selection.asm == format!("mov.u32 \t$d, {};", recipe.ptx_register) - && selection.predicates - == [ - "Subtarget->getSmVersion() >= 90", - "Subtarget->getPTXVersion() >= 78", - ] - && selection.constraints.address_space.is_none() - && selection.constraints.immediate_bindings.is_empty(), - "{} selector disagrees with LLVM TableGen", - policy.id - ); - ensure!( - policy.expected_ptx.mnemonic == "mov" - && policy.expected_ptx.modifiers == ["u32"] - && policy.expected_ptx.operands - == [ - OperandPattern::Register, - OperandPattern::Exact { - value: recipe.ptx_register, - }, - ], - "{} expected PTX does not match its cluster register", - policy.id - ); - Ok(()) -} - -fn validate_lanemask_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, -) -> Result<()> { - let (suffix, abi_id, section, op_type) = match policy.id.as_str() { - "lanemask_lt" => ("lt", "i0035", "10.13", "ReadPtxSregLanemaskLtOp"), - "lanemask_le" => ("le", "i0036", "10.12", "ReadPtxSregLanemaskLeOp"), - "lanemask_eq" => ("eq", "i0037", "10.11", "ReadPtxSregLanemaskEqOp"), - "lanemask_ge" => ("ge", "i0038", "10.14", "ReadPtxSregLanemaskGeOp"), - "lanemask_gt" => ("gt", "i0039", "10.15", "ReadPtxSregLanemaskGtOp"), - _ => bail!("{} is not a reviewed lane-mask special register", policy.id), - }; - ensure!( - policy.abi_id == abi_id - && policy.operation_key == format!("warp.lane_mask.{suffix}") - && policy.source.is_none() - && policy.source_record.as_deref() - == Some(format!("int_nvvm_read_ptx_sreg_lanemask_{suffix}").as_str()) - && policy.llvm_symbol.as_deref() - == Some(format!("llvm.nvvm.read.ptx.sreg.lanemask.{suffix}").as_str()) - && policy.resolved_llvm_symbol.is_none(), - "{} lane-mask identity does not match its closed recipe", - policy.id - ); - ensure!( - policy.rust_module == "sreg" - && policy.rust_name == policy.id - && policy.rust_arguments.is_empty() - && policy.rust_result == "u32" - && policy.safe - && policy.must_use - && policy - .safe_allowlist_reason - .as_deref() - .is_some_and(|reason| !reason.is_empty()) - && policy.public_rust_path == format!("cuda_intrinsics::sreg::{}", policy.id) - && policy.compatibility_rust_paths == [format!("cuda_device::warp::{}", policy.id)], - "{} must preserve its safe must-use raw and compatibility APIs", - policy.id - ); - ensure!( - policy.dialect_op_type == op_type - && policy.dialect_op_name == format!("nvvm.read_ptx_sreg_lanemask_{suffix}") - && policy.dialect_operands.is_empty() - && policy.dialect_results.is_empty() - && policy.llvm_arguments.is_empty() - && policy.llvm_results == ["i32"] - && policy.lowering == "direct_nvvm", - "{} is outside the closed lane-mask lowering recipe", - policy.id - ); - ensure!( - policy.pure - && policy.memory == "none" - && !policy.convergent - && policy.execution_scope == "thread" - && policy.minimum_ptx == "2.0" - && policy.minimum_sm.as_deref() == Some("sm_20") - && policy.ptx_result == "u32" - && policy.targets == "all", - "{} lane-mask effects or target floor disagree with PTX", - policy.id - ); - ensure!( - policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section == format!("{section} Special Registers: %lanemask_{suffix}") - && policy.ptx_isa_url - == format!( - "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-lanemask-{suffix}" - ), - "{} lane-mask PTX provenance disagrees with the reviewed recipe", - policy.id - ); - ensure!( - declaration.properties == ["IntrNoMem", "IntrSpeculatable", "NoUndef"], - "{} lane-mask properties disagree with the imported LLVM declaration", - policy.id - ); - ensure!( - policy.expected_ptx.mnemonic == "mov" - && policy.expected_ptx.modifiers == ["u32"] - && policy.expected_ptx.operands - == [ - OperandPattern::Register, - OperandPattern::Exact { - value: format!("%lanemask_{suffix}"), - }, - ], - "{} expected PTX does not match its lane-mask register", - policy.id - ); - Ok(()) -} - -fn validate_redux_policy(policy: &OverlayIntrinsic, declaration: &ImportedIntrinsic) -> Result<()> { - let redux = policy - .redux - .as_ref() - .with_context(|| format!("{} has no closed redux contract", policy.id))?; - let recipe = redux_recipe(redux.operation); - ensure!( - redux.participation - == ReduxParticipation::ExecutingLaneNamedAllNamedLanesSameInstructionAndMask - && redux.adapter == ReduxAdapter::MaskValueToSourceMemberMask, - "{} requests an unsupported redux participation contract or operand adapter", - policy.id - ); - ensure!( - policy.id == recipe.id - && policy.operation_key == recipe.operation_key - && policy.source_record.as_deref() == Some(recipe.source_record) - && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) - && policy.resolved_llvm_symbol.is_none(), - "{} redux identity does not match its closed operation recipe", - policy.id - ); - ensure!( - policy.rust_module == "warp" - && policy.rust_name == recipe.rust_name - && policy.rust_arguments == ["u32", recipe.rust_value] - && policy.rust_result == recipe.rust_value - && !policy.safe - && policy.must_use - && policy.safe_allowlist_reason.is_none() - && policy.public_rust_path == format!("cuda_intrinsics::warp::{}", recipe.rust_name) - && policy.compatibility_rust_paths - == [format!("cuda_device::warp::{}", recipe.rust_name)], - "{} must preserve the unsafe must-use redux raw API and legacy cuda-device compatibility DefPath", - policy.id - ); - ensure!( - policy.dialect_op_type == recipe.dialect_op_type - && policy.dialect_op_name == recipe.dialect_op_name - && policy.dialect_operands == ["i32", recipe.value_type] - && policy.dialect_results == [recipe.value_type] - && policy.llvm_arguments == [recipe.value_type, "i32"] - && policy.llvm_results == [recipe.value_type] - && policy.lowering == "generated_redux", - "{} is outside the generated two-operand redux recipe", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == "inaccessible_read_write" - && policy.convergent - && policy.execution_scope == "warp" - && policy.minimum_ptx == recipe.minimum_ptx - && policy.minimum_sm.as_deref() == recipe.minimum_sm - && policy.ptx_result == recipe.rust_value - && policy.targets == recipe.targets, - "{} redux effects, carrier, or target floor disagree with its operation recipe", - policy.id - ); - ensure!( - declaration - .properties - .iter() - .any(|property| property == "IntrConvergent") - && declaration - .properties - .iter() - .any(|property| property == "IntrInaccessibleMemOnly") - && declaration - .properties - .iter() - .any(|property| property == "IntrNoCallback") - && !declaration.properties.iter().any(|property| matches!( - property.as_str(), - "IntrNoMem" | "IntrReadMem" | "IntrWriteMem" - )), - "{} redux memory and convergence effects disagree with the imported declaration", - policy.id - ); - ensure!( - policy.backend_lowerings.is_empty() - && policy.packed_atomic.is_none() - && policy.ldmatrix_variant.is_none() - && policy.ldmatrix_safety.is_none() - && policy.ldmatrix_adapter.is_none() - && policy.dot_product.is_none() - && policy.packed_alu.is_none() - && policy.packed_conversion.is_none() - && policy.cp_async_copy.is_none() - && policy.cp_async_control.is_none() - && policy.mbarrier_basic.is_none() - && policy.vote.is_none() - && policy.active_mask.is_none() - && policy.warp_match.is_none() - && policy.warp_barrier.is_none() - && policy.warp_shuffle.is_none() - && policy.selected_address_space.is_none(), - "{} mixes another generated-family contract with redux", - policy.id - ); - ensure!( - policy.expected_ptx.mnemonic == "redux" - && policy.expected_ptx.modifiers == recipe.ptx_modifiers - && policy.expected_ptx.operands - == [ - OperandPattern::Register, - OperandPattern::Register, - OperandPattern::Register, - ], - "{} expected PTX does not match its closed redux operation recipe", - policy.id - ); - Ok(()) -} - -// Pinned LLVM's hasReduxSyncF32(): accel {100} at PTX 8.6, family {100} at -// PTX 8.8. The catalog carries the accel floor with the a/f target union. -const REDUX_F32_TARGETS: &str = "sm_100a|sm_100f|sm_103a|sm_103f"; - -struct ReduxRecipe { - id: &'static str, - operation_key: &'static str, - source_record: &'static str, - llvm_symbol: &'static str, - rust_name: &'static str, - rust_value: &'static str, - value_type: &'static str, - dialect_op_type: &'static str, - dialect_op_name: &'static str, - ptx_modifiers: &'static [&'static str], - minimum_ptx: &'static str, - minimum_sm: Option<&'static str>, - targets: &'static str, -} - -fn redux_recipe(operation: ReduxOperation) -> ReduxRecipe { - match operation { - ReduxOperation::Add => ReduxRecipe { - id: "redux_sync_add", - operation_key: "warp.redux.sync.add.wrap32", - source_record: "int_nvvm_redux_sync_add", - llvm_symbol: "llvm.nvvm.redux.sync.add", - rust_name: "redux_sync_add", - rust_value: "u32", - value_type: "i32", - dialect_op_type: "ReduxSyncAddOp", - dialect_op_name: "nvvm.redux_sync_add", - ptx_modifiers: &["sync", "add", "s32"], - minimum_ptx: "7.0", - minimum_sm: Some("sm_80"), - targets: "all", - }, - ReduxOperation::Umin => ReduxRecipe { - id: "redux_sync_min_u32", - operation_key: "warp.redux.sync.min.u32", - source_record: "int_nvvm_redux_sync_umin", - llvm_symbol: "llvm.nvvm.redux.sync.umin", - rust_name: "redux_sync_min_u32", - rust_value: "u32", - value_type: "i32", - dialect_op_type: "ReduxSyncUminOp", - dialect_op_name: "nvvm.redux_sync_umin", - ptx_modifiers: &["sync", "min", "u32"], - minimum_ptx: "7.0", - minimum_sm: Some("sm_80"), - targets: "all", - }, - ReduxOperation::Min => ReduxRecipe { - id: "redux_sync_min_i32", - operation_key: "warp.redux.sync.min.s32", - source_record: "int_nvvm_redux_sync_min", - llvm_symbol: "llvm.nvvm.redux.sync.min", - rust_name: "redux_sync_min_i32", - rust_value: "i32", - value_type: "i32", - dialect_op_type: "ReduxSyncMinOp", - dialect_op_name: "nvvm.redux_sync_min", - ptx_modifiers: &["sync", "min", "s32"], - minimum_ptx: "7.0", - minimum_sm: Some("sm_80"), - targets: "all", - }, - ReduxOperation::Umax => ReduxRecipe { - id: "redux_sync_max_u32", - operation_key: "warp.redux.sync.max.u32", - source_record: "int_nvvm_redux_sync_umax", - llvm_symbol: "llvm.nvvm.redux.sync.umax", - rust_name: "redux_sync_max_u32", - rust_value: "u32", - value_type: "i32", - dialect_op_type: "ReduxSyncUmaxOp", - dialect_op_name: "nvvm.redux_sync_umax", - ptx_modifiers: &["sync", "max", "u32"], - minimum_ptx: "7.0", - minimum_sm: Some("sm_80"), - targets: "all", - }, - ReduxOperation::Max => ReduxRecipe { - id: "redux_sync_max_i32", - operation_key: "warp.redux.sync.max.s32", - source_record: "int_nvvm_redux_sync_max", - llvm_symbol: "llvm.nvvm.redux.sync.max", - rust_name: "redux_sync_max_i32", - rust_value: "i32", - value_type: "i32", - dialect_op_type: "ReduxSyncMaxOp", - dialect_op_name: "nvvm.redux_sync_max", - ptx_modifiers: &["sync", "max", "s32"], - minimum_ptx: "7.0", - minimum_sm: Some("sm_80"), - targets: "all", - }, - ReduxOperation::And => ReduxRecipe { - id: "redux_sync_and", - operation_key: "warp.redux.sync.and.b32", - source_record: "int_nvvm_redux_sync_and", - llvm_symbol: "llvm.nvvm.redux.sync.and", - rust_name: "redux_sync_and", - rust_value: "u32", - value_type: "i32", - dialect_op_type: "ReduxSyncAndOp", - dialect_op_name: "nvvm.redux_sync_and", - ptx_modifiers: &["sync", "and", "b32"], - minimum_ptx: "7.0", - minimum_sm: Some("sm_80"), - targets: "all", - }, - ReduxOperation::Or => ReduxRecipe { - id: "redux_sync_or", - operation_key: "warp.redux.sync.or.b32", - source_record: "int_nvvm_redux_sync_or", - llvm_symbol: "llvm.nvvm.redux.sync.or", - rust_name: "redux_sync_or", - rust_value: "u32", - value_type: "i32", - dialect_op_type: "ReduxSyncOrOp", - dialect_op_name: "nvvm.redux_sync_or", - ptx_modifiers: &["sync", "or", "b32"], - minimum_ptx: "7.0", - minimum_sm: Some("sm_80"), - targets: "all", - }, - ReduxOperation::Xor => ReduxRecipe { - id: "redux_sync_xor", - operation_key: "warp.redux.sync.xor.b32", - source_record: "int_nvvm_redux_sync_xor", - llvm_symbol: "llvm.nvvm.redux.sync.xor", - rust_name: "redux_sync_xor", - rust_value: "u32", - value_type: "i32", - dialect_op_type: "ReduxSyncXorOp", - dialect_op_name: "nvvm.redux_sync_xor", - ptx_modifiers: &["sync", "xor", "b32"], - minimum_ptx: "7.0", - minimum_sm: Some("sm_80"), - targets: "all", - }, - ReduxOperation::Fmin => ReduxRecipe { - id: "redux_sync_min_f32", - operation_key: "warp.redux.sync.min.f32", - source_record: "int_nvvm_redux_sync_fmin", - llvm_symbol: "llvm.nvvm.redux.sync.fmin", - rust_name: "redux_sync_min_f32", - rust_value: "f32", - value_type: "f32", - dialect_op_type: "ReduxSyncFminOp", - dialect_op_name: "nvvm.redux_sync_fmin", - ptx_modifiers: &["sync", "min", "f32"], - minimum_ptx: "8.6", - minimum_sm: None, - targets: REDUX_F32_TARGETS, - }, - ReduxOperation::FminNan => ReduxRecipe { - id: "redux_sync_min_nan_f32", - operation_key: "warp.redux.sync.min.nan.f32", - source_record: "int_nvvm_redux_sync_fmin_NaN", - llvm_symbol: "llvm.nvvm.redux.sync.fmin.NaN", - rust_name: "redux_sync_min_nan_f32", - rust_value: "f32", - value_type: "f32", - dialect_op_type: "ReduxSyncFminNanOp", - dialect_op_name: "nvvm.redux_sync_fmin_nan", - ptx_modifiers: &["sync", "min", "NaN", "f32"], - minimum_ptx: "8.6", - minimum_sm: None, - targets: REDUX_F32_TARGETS, - }, - ReduxOperation::FminAbs => ReduxRecipe { - id: "redux_sync_min_abs_f32", - operation_key: "warp.redux.sync.min.abs.f32", - source_record: "int_nvvm_redux_sync_fmin_abs", - llvm_symbol: "llvm.nvvm.redux.sync.fmin.abs", - rust_name: "redux_sync_min_abs_f32", - rust_value: "f32", - value_type: "f32", - dialect_op_type: "ReduxSyncFminAbsOp", - dialect_op_name: "nvvm.redux_sync_fmin_abs", - ptx_modifiers: &["sync", "min", "abs", "f32"], - minimum_ptx: "8.6", - minimum_sm: None, - targets: REDUX_F32_TARGETS, - }, - ReduxOperation::FminAbsNan => ReduxRecipe { - id: "redux_sync_min_abs_nan_f32", - operation_key: "warp.redux.sync.min.abs.nan.f32", - source_record: "int_nvvm_redux_sync_fmin_abs_NaN", - llvm_symbol: "llvm.nvvm.redux.sync.fmin.abs.NaN", - rust_name: "redux_sync_min_abs_nan_f32", - rust_value: "f32", - value_type: "f32", - dialect_op_type: "ReduxSyncFminAbsNanOp", - dialect_op_name: "nvvm.redux_sync_fmin_abs_nan", - ptx_modifiers: &["sync", "min", "abs", "NaN", "f32"], - minimum_ptx: "8.6", - minimum_sm: None, - targets: REDUX_F32_TARGETS, - }, - ReduxOperation::Fmax => ReduxRecipe { - id: "redux_sync_max_f32", - operation_key: "warp.redux.sync.max.f32", - source_record: "int_nvvm_redux_sync_fmax", - llvm_symbol: "llvm.nvvm.redux.sync.fmax", - rust_name: "redux_sync_max_f32", - rust_value: "f32", - value_type: "f32", - dialect_op_type: "ReduxSyncFmaxOp", - dialect_op_name: "nvvm.redux_sync_fmax", - ptx_modifiers: &["sync", "max", "f32"], - minimum_ptx: "8.6", - minimum_sm: None, - targets: REDUX_F32_TARGETS, - }, - ReduxOperation::FmaxNan => ReduxRecipe { - id: "redux_sync_max_nan_f32", - operation_key: "warp.redux.sync.max.nan.f32", - source_record: "int_nvvm_redux_sync_fmax_NaN", - llvm_symbol: "llvm.nvvm.redux.sync.fmax.NaN", - rust_name: "redux_sync_max_nan_f32", - rust_value: "f32", - value_type: "f32", - dialect_op_type: "ReduxSyncFmaxNanOp", - dialect_op_name: "nvvm.redux_sync_fmax_nan", - ptx_modifiers: &["sync", "max", "NaN", "f32"], - minimum_ptx: "8.6", - minimum_sm: None, - targets: REDUX_F32_TARGETS, - }, - ReduxOperation::FmaxAbs => ReduxRecipe { - id: "redux_sync_max_abs_f32", - operation_key: "warp.redux.sync.max.abs.f32", - source_record: "int_nvvm_redux_sync_fmax_abs", - llvm_symbol: "llvm.nvvm.redux.sync.fmax.abs", - rust_name: "redux_sync_max_abs_f32", - rust_value: "f32", - value_type: "f32", - dialect_op_type: "ReduxSyncFmaxAbsOp", - dialect_op_name: "nvvm.redux_sync_fmax_abs", - ptx_modifiers: &["sync", "max", "abs", "f32"], - minimum_ptx: "8.6", - minimum_sm: None, - targets: REDUX_F32_TARGETS, - }, - ReduxOperation::FmaxAbsNan => ReduxRecipe { - id: "redux_sync_max_abs_nan_f32", - operation_key: "warp.redux.sync.max.abs.nan.f32", - source_record: "int_nvvm_redux_sync_fmax_abs_NaN", - llvm_symbol: "llvm.nvvm.redux.sync.fmax.abs.NaN", - rust_name: "redux_sync_max_abs_nan_f32", - rust_value: "f32", - value_type: "f32", - dialect_op_type: "ReduxSyncFmaxAbsNanOp", - dialect_op_name: "nvvm.redux_sync_fmax_abs_nan", - ptx_modifiers: &["sync", "max", "abs", "NaN", "f32"], - minimum_ptx: "8.6", - minimum_sm: None, - targets: REDUX_F32_TARGETS, - }, - } -} - -fn validate_dot_product_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, -) -> Result<()> { - let dot_product = policy - .dot_product - .as_ref() - .with_context(|| format!("{} has no closed dot-product contract", policy.id))?; - let recipe = dot_product_recipe(dot_product.operation, dot_product.signedness); - ensure!( - dot_product.adapter == recipe.adapter, - "{} dot-product source adapter does not match its operation", - policy.id - ); - ensure!( - policy.id == recipe.id - && policy.operation_key == recipe.operation_key - && policy.source.is_none() - && policy.source_record.as_deref() == Some(recipe.source_record) - && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) - && policy.resolved_llvm_symbol.is_none(), - "{} dot-product identity does not match its closed operation recipe", - policy.id - ); - ensure!( - policy.rust_module == "dotprod" - && policy.rust_name == recipe.rust_name - && policy.rust_arguments == ["u32", "u32", recipe.rust_value] - && policy.rust_result == recipe.rust_value - && policy.safe - && !policy.must_use - && policy.public_rust_path == format!("cuda_intrinsics::dotprod::{}", recipe.rust_name) - && policy.compatibility_rust_paths - == [format!("cuda_device::dotprod::{}", recipe.rust_name)], - "{} must preserve the safe, non-must-use dotprod raw API and legacy cuda-device DefPath", - policy.id - ); - ensure!( - policy.dialect_op_type == recipe.dialect_op_type - && policy.dialect_op_name == recipe.dialect_op_name - && policy.dialect_operands == ["i32", "i32", "i32"] - && policy.dialect_results == ["i32"] - && policy.llvm_arguments == recipe.llvm_arguments - && policy.llvm_results == ["i32"] - && policy.lowering == "generated_dotprod", - "{} is outside the closed three-operand dot-product lowering recipe", - policy.id - ); - ensure!( - policy.pure - && policy.memory == "none" - && !policy.convergent - && policy.execution_scope == "thread" - && policy.minimum_ptx == "5.0" - && policy.minimum_sm.as_deref() == Some("sm_61") - && policy.ptx_result == recipe.rust_value - && policy.targets == "all", - "{} dot-product effects, carrier, or target floor disagree with its operation recipe", - policy.id - ); - ensure!( - declaration - .classes - .iter() - .any(|class| class == "NVVMPureIntrinsic") - && declaration.properties == recipe.llvm_properties, - "{} dot-product effects or immediate contract disagree with the imported declaration", - policy.id - ); - ensure!( - policy.ldmatrix_variant.is_none() - && policy.ldmatrix_safety.is_none() - && policy.ldmatrix_adapter.is_none() - && policy.packed_atomic.is_none() - && policy.redux.is_none() - && policy.vote.is_none() - && policy.active_mask.is_none() - && policy.warp_match.is_none() - && policy.warp_barrier.is_none() - && policy.warp_shuffle.is_none() - && policy.packed_alu.is_none() - && policy.packed_conversion.is_none() - && policy.cp_async_copy.is_none() - && policy.cp_async_control.is_none() - && policy.mbarrier_basic.is_none() - && policy.selected_address_space.is_none(), - "{} mixes another generated-family contract with dotprod", - policy.id - ); - ensure!( - policy.expected_ptx.mnemonic == recipe.ptx_mnemonic - && policy.expected_ptx.modifiers == recipe.ptx_modifiers - && policy.expected_ptx.operands - == [ - OperandPattern::Register, - OperandPattern::Register, - OperandPattern::Register, - OperandPattern::Register, - ], - "{} expected PTX does not match its closed dot-product recipe", - policy.id - ); - - let backend_pairs: BTreeSet<_> = policy - .backend_lowerings - .iter() - .map(|lowering| (lowering.backend, lowering.mechanism)) - .collect(); - ensure!( - policy.backend_lowerings.len() == 2 - && backend_pairs - == BTreeSet::from([ - ( - IntrinsicBackend::LlvmNvptx, - BackendLoweringMechanism::TypedNvvm, - ), - ( - IntrinsicBackend::LibNvvm, - BackendLoweringMechanism::InlinePtx, - ), - ]), - "{} must define exactly the reviewed LLVM typed and libNVVM inline-PTX routes", - policy.id - ); - for lowering in &policy.backend_lowerings { - let floor_matches = match lowering.backend { - IntrinsicBackend::LlvmNvptx => { - lowering.mechanism == BackendLoweringMechanism::TypedNvvm - && lowering.minimum_ptx.is_none() - && lowering.minimum_sm.is_none() - } - IntrinsicBackend::LibNvvm => { - lowering.mechanism == BackendLoweringMechanism::InlinePtx - && lowering.minimum_ptx.is_none() - && lowering.minimum_sm.as_deref() == Some("sm_75") - } - }; - ensure!( - floor_matches && !lowering.evidence_profile.trim().is_empty(), - "{} backend {:?} does not carry its reviewed dot-product profile floor", - policy.id, - lowering.backend - ); - } - Ok(()) -} - -struct DotProductRecipe { - id: &'static str, - operation_key: &'static str, - source_record: &'static str, - llvm_symbol: &'static str, - rust_name: &'static str, - rust_value: &'static str, - dialect_op_type: &'static str, - dialect_op_name: &'static str, - llvm_arguments: &'static [&'static str], - llvm_properties: &'static [&'static str], - adapter: DotProductAdapter, - ptx_mnemonic: &'static str, - ptx_modifiers: &'static [&'static str], -} - -fn dot_product_recipe( - operation: DotProductOperation, - signedness: DotProductSignedness, -) -> DotProductRecipe { - match (operation, signedness) { - (DotProductOperation::Dp4a, DotProductSignedness::Signed) => DotProductRecipe { - id: "dp4a_s32", - operation_key: "integer.dot_product.dp4a.s32", - source_record: "int_nvvm_idp4a_s_s", - llvm_symbol: "llvm.nvvm.idp4a.s.s", - rust_name: "dp4a_s32", - rust_value: "i32", - dialect_op_type: "Dp4aS32Op", - dialect_op_name: "nvvm.dp4a_s32", - llvm_arguments: &["i32", "i32", "i32"], - llvm_properties: &["IntrNoMem", "IntrSpeculatable"], - adapter: DotProductAdapter::DirectThreeOperands, - ptx_mnemonic: "dp4a", - ptx_modifiers: &["s32", "s32"], - }, - (DotProductOperation::Dp4a, DotProductSignedness::Unsigned) => DotProductRecipe { - id: "dp4a_u32", - operation_key: "integer.dot_product.dp4a.u32", - source_record: "int_nvvm_idp4a_u_u", - llvm_symbol: "llvm.nvvm.idp4a.u.u", - rust_name: "dp4a_u32", - rust_value: "u32", - dialect_op_type: "Dp4aU32Op", - dialect_op_name: "nvvm.dp4a_u32", - llvm_arguments: &["i32", "i32", "i32"], - llvm_properties: &["IntrNoMem", "IntrSpeculatable"], - adapter: DotProductAdapter::DirectThreeOperands, - ptx_mnemonic: "dp4a", - ptx_modifiers: &["u32", "u32"], - }, - (DotProductOperation::Dp2a, DotProductSignedness::Signed) => DotProductRecipe { - id: "dp2a_s32", - operation_key: "integer.dot_product.dp2a.lo.s32", - source_record: "int_nvvm_idp2a_s_s", - llvm_symbol: "llvm.nvvm.idp2a.s.s", - rust_name: "dp2a_s32", - rust_value: "i32", - dialect_op_type: "Dp2aS32Op", - dialect_op_name: "nvvm.dp2a_s32", - llvm_arguments: &["i32", "i32", "i1", "i32"], - llvm_properties: &["ImmArg", "IntrNoMem", "IntrSpeculatable"], - adapter: DotProductAdapter::InsertLowHalfFalse, - ptx_mnemonic: "dp2a", - ptx_modifiers: &["lo", "s32", "s32"], - }, - (DotProductOperation::Dp2a, DotProductSignedness::Unsigned) => DotProductRecipe { - id: "dp2a_u32", - operation_key: "integer.dot_product.dp2a.lo.u32", - source_record: "int_nvvm_idp2a_u_u", - llvm_symbol: "llvm.nvvm.idp2a.u.u", - rust_name: "dp2a_u32", - rust_value: "u32", - dialect_op_type: "Dp2aU32Op", - dialect_op_name: "nvvm.dp2a_u32", - llvm_arguments: &["i32", "i32", "i1", "i32"], - llvm_properties: &["ImmArg", "IntrNoMem", "IntrSpeculatable"], - adapter: DotProductAdapter::InsertLowHalfFalse, - ptx_mnemonic: "dp2a", - ptx_modifiers: &["lo", "u32", "u32"], - }, - } -} - -#[derive(Clone, Copy)] -struct StmatrixRecipe { - multiplicity: StmatrixMultiplicity, - layout: StmatrixLayout, - abi_id: &'static str, - id: &'static str, - operation_key: &'static str, - source_record: &'static str, - llvm_symbol: &'static str, - compatibility_name: &'static str, - dialect_op_type: &'static str, - dialect_op_name: &'static str, - summary: &'static str, -} - -fn stmatrix_recipe(multiplicity: StmatrixMultiplicity, layout: StmatrixLayout) -> StmatrixRecipe { - match (multiplicity, layout) { - (StmatrixMultiplicity::X2, StmatrixLayout::Normal) => StmatrixRecipe { - multiplicity, - layout, - abi_id: "i0301", - id: "stmatrix_m8n8_x2_b16", - operation_key: "matrix.stmatrix.m8n8.x2.normal.b16.shared", - source_record: "int_nvvm_stmatrix_sync_aligned_m8n8_x2_b16", - llvm_symbol: "llvm.nvvm.stmatrix.sync.aligned.m8n8.x2.b16", - compatibility_name: "stmatrix_m8n8_x2", - dialect_op_type: "StmatrixM8n8X2Op", - dialect_op_name: "nvvm.stmatrix_m8n8_x2", - summary: "Stores two 8×8 b16 matrix fragments cooperatively to shared memory.", - }, - (StmatrixMultiplicity::X2, StmatrixLayout::Transposed) => StmatrixRecipe { - multiplicity, - layout, - abi_id: "i0302", - id: "stmatrix_m8n8_x2_trans_b16", - operation_key: "matrix.stmatrix.m8n8.x2.transposed.b16.shared", - source_record: "int_nvvm_stmatrix_sync_aligned_m8n8_x2_trans_b16", - llvm_symbol: "llvm.nvvm.stmatrix.sync.aligned.m8n8.x2.trans.b16", - compatibility_name: "stmatrix_m8n8_x2_trans", - dialect_op_type: "StmatrixM8n8X2TransOp", - dialect_op_name: "nvvm.stmatrix_m8n8_x2_trans", - summary: "Stores two transposed 8×8 b16 matrix fragments cooperatively to shared memory.", - }, - (StmatrixMultiplicity::X4, StmatrixLayout::Normal) => StmatrixRecipe { - multiplicity, - layout, - abi_id: "i0303", - id: "stmatrix_m8n8_x4_b16", - operation_key: "matrix.stmatrix.m8n8.x4.normal.b16.shared", - source_record: "int_nvvm_stmatrix_sync_aligned_m8n8_x4_b16", - llvm_symbol: "llvm.nvvm.stmatrix.sync.aligned.m8n8.x4.b16", - compatibility_name: "stmatrix_m8n8_x4", - dialect_op_type: "StmatrixM8n8X4Op", - dialect_op_name: "nvvm.stmatrix_m8n8_x4", - summary: "Stores four 8×8 b16 matrix fragments cooperatively to shared memory.", - }, - (StmatrixMultiplicity::X4, StmatrixLayout::Transposed) => StmatrixRecipe { - multiplicity, - layout, - abi_id: "i0304", - id: "stmatrix_m8n8_x4_trans_b16", - operation_key: "matrix.stmatrix.m8n8.x4.transposed.b16.shared", - source_record: "int_nvvm_stmatrix_sync_aligned_m8n8_x4_trans_b16", - llvm_symbol: "llvm.nvvm.stmatrix.sync.aligned.m8n8.x4.trans.b16", - compatibility_name: "stmatrix_m8n8_x4_trans", - dialect_op_type: "StmatrixM8n8X4TransOp", - dialect_op_name: "nvvm.stmatrix_m8n8_x4_trans", - summary: "Stores four transposed 8×8 b16 matrix fragments cooperatively to shared memory.", - }, - } -} - -fn stmatrix_variant_for_id(id: &str) -> Option<(StmatrixMultiplicity, StmatrixLayout)> { - match id { - "stmatrix_m8n8_x2_b16" => Some((StmatrixMultiplicity::X2, StmatrixLayout::Normal)), - "stmatrix_m8n8_x2_trans_b16" => { - Some((StmatrixMultiplicity::X2, StmatrixLayout::Transposed)) - } - "stmatrix_m8n8_x4_b16" => Some((StmatrixMultiplicity::X4, StmatrixLayout::Normal)), - "stmatrix_m8n8_x4_trans_b16" => { - Some((StmatrixMultiplicity::X4, StmatrixLayout::Transposed)) - } - _ => None, - } -} - -fn expand_stmatrix_admission(admission: &StmatrixAdmission) -> Result> { - ensure!( - admission.runtime_validation == RuntimeValidation::Unexecuted, - "stmatrix runtime validation may be marked executed only with GPU evidence" - ); - ensure!( - !admission.llvm_evidence_profile.trim().is_empty() - && !admission.libnvvm_evidence_profile.trim().is_empty(), - "compact stmatrix admission requires both backend evidence profiles" - ); - let expected_variants = [ - (StmatrixMultiplicity::X2, StmatrixLayout::Normal), - (StmatrixMultiplicity::X2, StmatrixLayout::Transposed), - (StmatrixMultiplicity::X4, StmatrixLayout::Normal), - (StmatrixMultiplicity::X4, StmatrixLayout::Transposed), - ]; - let actual_variants = admission - .variants - .iter() - .map(|variant| (variant.multiplicity, variant.layout)) - .collect::>(); - ensure!( - actual_variants == expected_variants, - "compact stmatrix admission must contain the four reviewed variants in canonical order" - ); - - admission - .variants - .iter() - .map(|variant| { - let recipe = stmatrix_recipe(variant.multiplicity, variant.layout); - ensure!( - variant.abi_id == recipe.abi_id, - "{} must keep reserved ABI ID {}", - recipe.id, - recipe.abi_id - ); - let count = recipe.multiplicity.register_count(); - let mut rust_arguments = vec!["*mut u8".to_owned()]; - rust_arguments.extend(std::iter::repeat_n("u32".to_owned(), count)); - let mut dialect_operands = vec!["ptr".to_owned()]; - dialect_operands.extend(std::iter::repeat_n("i32".to_owned(), count)); - let mut llvm_arguments = vec!["anyptr".to_owned()]; - llvm_arguments.extend(std::iter::repeat_n("i32".to_owned(), count)); - let multiplicity = match recipe.multiplicity { - StmatrixMultiplicity::X2 => "x2", - StmatrixMultiplicity::X4 => "x4", - }; - let mut modifiers = vec![ - "sync".into(), - "aligned".into(), - "m8n8".into(), - multiplicity.into(), - ]; - if recipe.layout == StmatrixLayout::Transposed { - modifiers.push("trans".into()); - } - modifiers.extend(["shared".into(), "b16".into()]); - Ok(OverlayIntrinsic { - id: recipe.id.into(), - abi_id: variant.abi_id.clone(), - operation_key: recipe.operation_key.into(), - family: "stmatrix".into(), - source: None, - source_record: Some(recipe.source_record.into()), - rust_module: "matrix".into(), - rust_name: recipe.id.into(), - rust_arguments, - rust_result: "()".into(), - safe: false, - must_use: false, - safe_allowlist_reason: None, - public_rust_path: format!("cuda_intrinsics::matrix::{}", recipe.id), - compatibility_rust_paths: vec![format!( - "cuda_device::tcgen05::{}", - recipe.compatibility_name - )], - dialect_op_type: recipe.dialect_op_type.into(), - dialect_op_name: recipe.dialect_op_name.into(), - dialect_operands, - dialect_results: vec![], - llvm_symbol: Some(recipe.llvm_symbol.into()), - resolved_llvm_symbol: Some(format!("{}.p3", recipe.llvm_symbol)), - llvm_arguments, - llvm_results: vec![], - pure: false, - memory: "write".into(), - convergent: true, - execution_scope: "warp".into(), - minimum_ptx: "7.8".into(), - minimum_sm: Some("sm_90".into()), - ptx_result: "()".into(), - targets: "all".into(), - ptx_isa_version: "9.3".into(), - ptx_isa_section: "9.7.14.5.16 Warp-level matrix store instruction: stmatrix" - .into(), - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-stmatrix".into(), - lowering: "generated_stmatrix".into(), - backend_lowerings: vec![ - OverlayBackendLowering { - backend: IntrinsicBackend::LlvmNvptx, - mechanism: BackendLoweringMechanism::TypedNvvm, - evidence_profile: admission.llvm_evidence_profile.clone(), - targets: None, - minimum_ptx: Some("7.8".into()), - minimum_sm: Some("sm_90".into()), - }, - OverlayBackendLowering { - backend: IntrinsicBackend::LibNvvm, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: admission.libnvvm_evidence_profile.clone(), - targets: None, - minimum_ptx: Some("7.8".into()), - minimum_sm: Some("sm_90".into()), - }, - ], - packed_atomic: None, - redux: None, - vote: None, - active_mask: None, - warp_match: None, - warp_barrier: None, - warp_shuffle: None, - dot_product: None, - packed_alu: None, - integer_minmax: None, - packed_conversion: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cp_async_copy: None, - cp_async_control: None, - cp_async_mbarrier: None, - mbarrier_basic: None, - movmatrix: None, - mbarrier_extended: None, - register_mma: None, - sparse_mma: None, - prmt: None, - cluster_barrier: None, - wgmma_control: None, - special_register: None, - debug_control: None, - cluster_memory: None, - clc: None, - tma: None, - tcgen05: None, - ldmatrix_variant: None, - ldmatrix_safety: None, - ldmatrix_adapter: None, - selected_address_space: Some(ImportedAddressSpace::Shared), - expected_ptx: InstructionPattern { - mnemonic: "stmatrix".into(), - modifiers, - operands: vec![ - OperandPattern::Address, - OperandPattern::RegisterList { length: count }, - ], - }, - summary: recipe.summary.into(), - }) - }) - .collect() -} - -fn validate_stmatrix_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, -) -> Result<()> { - let (multiplicity, layout) = stmatrix_variant_for_id(&policy.id) - .with_context(|| format!("{} has no closed stmatrix recipe", policy.id))?; - let recipe = stmatrix_recipe(multiplicity, layout); - let count = multiplicity.register_count(); - let mut rust_arguments = vec!["*mut u8".to_owned()]; - rust_arguments.extend(std::iter::repeat_n("u32".to_owned(), count)); - let mut dialect_operands = vec!["ptr".to_owned()]; - dialect_operands.extend(std::iter::repeat_n("i32".to_owned(), count)); - let mut llvm_arguments = vec!["anyptr".to_owned()]; - llvm_arguments.extend(std::iter::repeat_n("i32".to_owned(), count)); - - ensure!( - policy.abi_id == recipe.abi_id - && policy.operation_key == recipe.operation_key - && policy.source.is_none() - && policy.source_record.as_deref() == Some(recipe.source_record) - && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) - && policy.resolved_llvm_symbol.as_deref() - == Some(format!("{}.p3", recipe.llvm_symbol).as_str()), - "{} stmatrix identity does not match its closed recipe", - policy.id - ); - ensure!( - policy.rust_module == "matrix" - && policy.rust_name == recipe.id - && policy.rust_arguments == rust_arguments - && policy.rust_result == "()" - && !policy.safe - && !policy.must_use - && policy.safe_allowlist_reason.is_none() - && policy.public_rust_path == format!("cuda_intrinsics::matrix::{}", recipe.id) - && policy.compatibility_rust_paths - == [format!( - "cuda_device::tcgen05::{}", - recipe.compatibility_name - )], - "{} stmatrix Rust API does not match its closed recipe", - policy.id - ); - ensure!( - policy.dialect_op_type == recipe.dialect_op_type - && policy.dialect_op_name == recipe.dialect_op_name - && policy.dialect_operands == dialect_operands - && policy.dialect_results.is_empty() - && policy.llvm_arguments == llvm_arguments - && policy.llvm_results.is_empty() - && policy.lowering == "generated_stmatrix" - && policy.selected_address_space == Some(ImportedAddressSpace::Shared), - "{} stmatrix carriers or lowering do not match its closed recipe", - policy.id - ); - ensure!( - declaration.properties - == [ - "IntrArgMemOnly", - "IntrConvergent", - "IntrNoCallback", - "IntrWriteMem", - "NoCapture", - "WriteOnly", - ] - && !policy.pure - && policy.memory == "write" - && policy.convergent - && policy.execution_scope == "warp", - "{} stmatrix effects disagree with the imported declaration", - policy.id - ); - ensure!( - policy.minimum_ptx == "7.8" - && policy.minimum_sm.as_deref() == Some("sm_90") - && policy.ptx_result == "()" - && policy.targets == "all" - && policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section - == "9.7.14.5.16 Warp-level matrix store instruction: stmatrix" - && policy.ptx_isa_url - == "https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-stmatrix", - "{} stmatrix target floor or PTX provenance changed", - policy.id - ); - let multiplicity_name = match multiplicity { - StmatrixMultiplicity::X2 => "x2", - StmatrixMultiplicity::X4 => "x4", - }; - let mut modifiers = vec![ - "sync".into(), - "aligned".into(), - "m8n8".into(), - multiplicity_name.into(), - ]; - if layout == StmatrixLayout::Transposed { - modifiers.push("trans".into()); - } - modifiers.extend(["shared".into(), "b16".into()]); - ensure!( - policy.expected_ptx - == (InstructionPattern { - mnemonic: "stmatrix".into(), - modifiers, - operands: vec![ - OperandPattern::Address, - OperandPattern::RegisterList { length: count }, - ], - }), - "{} expected PTX does not match its closed stmatrix shape", - policy.id - ); - ensure!( - policy.backend_lowerings.len() == 2 - && policy.backend_lowerings.iter().all(|route| { - !route.evidence_profile.trim().is_empty() - && route.minimum_ptx.as_deref() == Some("7.8") - && route.minimum_sm.as_deref() == Some("sm_90") - }) - && policy.backend_lowerings.iter().any(|route| { - route.backend == IntrinsicBackend::LlvmNvptx - && route.mechanism == BackendLoweringMechanism::TypedNvvm - }) - && policy.backend_lowerings.iter().any(|route| { - route.backend == IntrinsicBackend::LibNvvm - && route.mechanism == BackendLoweringMechanism::InlinePtx - }), - "{} must keep its reviewed typed-NVVM and inline-PTX routes", - policy.id - ); - ensure!( - policy.ldmatrix_variant.is_none() - && policy.ldmatrix_safety.is_none() - && policy.ldmatrix_adapter.is_none() - && policy.packed_atomic.is_none() - && policy.redux.is_none() - && policy.vote.is_none() - && policy.active_mask.is_none() - && policy.warp_match.is_none() - && policy.warp_barrier.is_none() - && policy.warp_shuffle.is_none() - && policy.dot_product.is_none() - && policy.packed_alu.is_none() - && policy.packed_conversion.is_none() - && policy.cp_async_copy.is_none() - && policy.cp_async_control.is_none() - && policy.cp_async_mbarrier.is_none() - && policy.mbarrier_basic.is_none() - && policy.register_mma.is_none() - && policy.sparse_mma.is_none() - && policy.prmt.is_none() - && policy.cluster_barrier.is_none() - && policy.special_register.is_none() - && policy.debug_control.is_none(), - "{} mixes another generated-family contract with stmatrix", - policy.id - ); - Ok(()) -} - -fn validate_ldmatrix_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, -) -> Result<()> { - let variant = policy - .ldmatrix_variant - .as_ref() - .with_context(|| format!("{} has no closed ldmatrix variant", policy.id))?; - let safety = policy - .ldmatrix_safety - .as_ref() - .with_context(|| format!("{} has no ldmatrix safety contract", policy.id))?; - let classic = variant.shape == LdmatrixShape::M8n8 && variant.element == LdmatrixElement::B16; - let blackwell = matches!( - (variant.shape, variant.layout, variant.element), - ( - LdmatrixShape::M8n16, - LdmatrixLayout::Normal, - LdmatrixElement::B8x16B4x16P64 | LdmatrixElement::B8x16B6x16P32 - ) | ( - LdmatrixShape::M16n16, - LdmatrixLayout::Transposed, - LdmatrixElement::B8 | LdmatrixElement::B8x16B4x16P64 | LdmatrixElement::B8x16B6x16P32 - ) - ); - ensure!( - variant.shape != LdmatrixShape::M16n16 || variant.multiplicity != LdmatrixMultiplicity::X4, - "{} requests unsupported m16n16.x4 ldmatrix", - policy.id - ); - ensure!( - (classic || blackwell) && variant.state_space == LdmatrixStateSpace::Shared, - "{} requests an unsupported ldmatrix shape, element, or state space", - policy.id - ); - let expected_participation = if classic { - LdmatrixParticipation::AllWarpLanesSameInstruction - } else { - LdmatrixParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes - }; - let expected_address_contract = match variant.shape { - LdmatrixShape::M8n8 => { - LdmatrixAddressContract::WarpLaneAddressesMappedByMultiplicitySixteenByteAlignedSixteenBytesReadableWithSm75Replication - } - LdmatrixShape::M8n16 => { - LdmatrixAddressContract::WarpLaneAddressesMappedByMultiplicitySixteenByteAlignedSixteenBytesReadable - } - LdmatrixShape::M16n16 => { - LdmatrixAddressContract::WarpLaneAddressesMappedByMultiplicitySixteenByteAlignedThirtyTwoBytesReadable - } - }; - ensure!( - safety.participation == expected_participation - && safety.address_contract == expected_address_contract - && safety.memory_order == LdmatrixMemoryOrder::Weak - && safety.runtime_validation == RuntimeValidation::Unexecuted, - "{} has an unsupported ldmatrix safety contract", - policy.id - ); - let count = variant.register_count(); - let count_name = match variant.multiplicity { - LdmatrixMultiplicity::X1 => "x1", - LdmatrixMultiplicity::X2 => "x2", - LdmatrixMultiplicity::X4 => "x4", - }; - let trans_record = match variant.layout { - LdmatrixLayout::Normal => "", - LdmatrixLayout::Transposed => "_trans", - }; - let trans_symbol = match variant.layout { - LdmatrixLayout::Normal => "", - LdmatrixLayout::Transposed => ".trans", - }; - let layout_name = match variant.layout { - LdmatrixLayout::Normal => "normal", - LdmatrixLayout::Transposed => "transposed", - }; - let shape_name = match variant.shape { - LdmatrixShape::M8n8 => "m8n8", - LdmatrixShape::M8n16 => "m8n16", - LdmatrixShape::M16n16 => "m16n16", - }; - let (element_record, element_symbol) = match variant.element { - LdmatrixElement::B16 => ("b16", "b16"), - LdmatrixElement::B8 => ("b8", "b8"), - LdmatrixElement::B8x16B4x16P64 => ("b8x16_b4x16_p64", "b8x16.b4x16_p64"), - LdmatrixElement::B8x16B6x16P32 => ("b8x16_b6x16_p32", "b8x16.b6x16_p32"), - }; - let expected_source = format!( - "int_nvvm_ldmatrix_sync_aligned_{shape_name}_{count_name}{trans_record}_{element_record}" - ); - let expected_symbol = format!( - "llvm.nvvm.ldmatrix.sync.aligned.{shape_name}.{count_name}{trans_symbol}.{element_symbol}" - ); - let expected_name = - format!("ldmatrix_{shape_name}_{count_name}{trans_record}_{element_record}"); - let expected_result = if count == 1 { - "u32".to_owned() - } else { - format!("[u32; {count}]") - }; - let expected_adapter = if count == 1 { - LdmatrixAdapter::SingleResultDirect - } else { - LdmatrixAdapter::MultipleResultsToArray - }; - ensure!( - policy.source_record.as_deref() == Some(expected_source.as_str()) - && policy.llvm_symbol.as_deref() == Some(expected_symbol.as_str()), - "{} ldmatrix variant does not match its imported source record or base LLVM symbol", - policy.id - ); - ensure!( - policy.resolved_llvm_symbol.as_deref() == Some(format!("{expected_symbol}.p3").as_str()), - "{} must keep the imported base symbol distinct from the resolved `.p3` overload", - policy.id - ); - ensure!( - policy.rust_arguments == [if classic { "*const u32" } else { "*const u8" }] - && policy.rust_result == expected_result - && policy.llvm_arguments == ["anyptr"] - && policy.llvm_results == vec!["i32"; count] - && policy.ptx_result == policy.rust_result, - "{} ldmatrix Rust, imported LLVM, and PTX carrier signatures disagree", - policy.id - ); - ensure!( - policy.id == expected_name - && policy.operation_key - == format!( - "matrix.ldmatrix.{shape_name}.{count_name}.{layout_name}.{element_record}.shared" - ) - && policy.rust_module == "matrix" - && policy.rust_name == expected_name - && !policy.safe - && policy.must_use != classic - && policy.safe_allowlist_reason.is_none() - && { - let expected_compatibility_paths = if classic { - vec![format!( - "cuda_device::wmma::ldmatrix_{count_name}{trans_record}" - )] - } else if matches!( - policy.id.as_str(), - "ldmatrix_m8n16_x2_b8x16_b4x16_p64" | "ldmatrix_m8n16_x4_b8x16_b4x16_p64" - ) { - vec![format!("cuda_device::wmma::{expected_name}")] - } else { - vec![] - }; - policy.compatibility_rust_paths == expected_compatibility_paths - } - && policy.lowering == "generated_ldmatrix" - && policy.ldmatrix_adapter == Some(expected_adapter) - && policy.selected_address_space == Some(ImportedAddressSpace::Shared), - "{} must preserve the closed raw/compatibility ldmatrix API, result adapter, and shared selection", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == "read" - && policy.convergent - && policy.execution_scope == "warp" - && declaration - .properties - .iter() - .any(|property| property == "IntrArgMemOnly") - && declaration - .properties - .iter() - .any(|property| property == "IntrReadMem"), - "{} ldmatrix effects disagree with the imported declaration", - policy.id - ); - let backend_pairs: BTreeSet<_> = policy - .backend_lowerings - .iter() - .map(|lowering| (lowering.backend, lowering.mechanism)) - .collect(); - ensure!( - backend_pairs - == BTreeSet::from([ - ( - IntrinsicBackend::LlvmNvptx, - BackendLoweringMechanism::TypedNvvm - ), - ( - IntrinsicBackend::LibNvvm, - BackendLoweringMechanism::InlinePtx - ), - ]), - "{} must define exactly the reviewed LLVM typed and libNVVM inline-PTX lowerings", - policy.id - ); - ensure!( - policy - .backend_lowerings - .iter() - .all(|lowering| !lowering.evidence_profile.trim().is_empty()), - "{} backend lowering omits its evidence profile", - policy.id - ); - if blackwell { - ensure!( - policy.minimum_ptx == "8.6" - && policy.minimum_sm.is_none() - && policy.targets == BLACKWELL_LDMATRIX_LLVM_TARGETS, - "{} must retain the reviewed PTX 8.6 floor and exact LLVM Blackwell target set", - policy.id - ); - let llvm = policy - .backend_lowerings - .iter() - .find(|lowering| lowering.backend == IntrinsicBackend::LlvmNvptx) - .expect("validated LLVM route"); - let libnvvm = policy - .backend_lowerings - .iter() - .find(|lowering| lowering.backend == IntrinsicBackend::LibNvvm) - .expect("validated libNVVM route"); - ensure!( - llvm.targets.is_none() - && llvm.minimum_sm.is_none() - && backend_target_requirement(policy, llvm)? - == CatalogTargetRequirement { - minimum_ptx: parse_ptx_version("8.6", &policy.id)?, - hardware: parse_hardware_target_fields( - &policy.id, - BLACKWELL_LDMATRIX_LLVM_TARGETS, - None, - )?, - } - && libnvvm.targets.as_deref() == Some(BLACKWELL_LDMATRIX_LIBNVVM_TARGETS) - && libnvvm.minimum_sm.is_none() - && backend_target_requirement(policy, libnvvm)? - == CatalogTargetRequirement { - minimum_ptx: parse_ptx_version("8.6", &policy.id)?, - hardware: parse_hardware_target_fields( - &policy.id, - BLACKWELL_LDMATRIX_LIBNVVM_TARGETS, - None, - )?, - }, - "{} backend routes do not preserve the reviewed LLVM/libNVVM Blackwell target split", - policy.id - ); - } - ensure!( - policy.packed_atomic.is_none() - && policy.redux.is_none() - && policy.vote.is_none() - && policy.active_mask.is_none() - && policy.warp_match.is_none() - && policy.warp_barrier.is_none() - && policy.warp_shuffle.is_none() - && policy.dot_product.is_none() - && policy.packed_alu.is_none() - && policy.packed_conversion.is_none() - && policy.cp_async_copy.is_none() - && policy.cp_async_control.is_none() - && policy.mbarrier_basic.is_none(), - "{} mixes another generated-family contract with ldmatrix", - policy.id - ); - Ok(()) -} - -fn validate_packed_atomic_policy( - policy: &OverlayIntrinsic, - source: &IntrinsicSource, -) -> Result<()> { - let packed = policy - .packed_atomic - .as_ref() - .with_context(|| format!("{} has no closed packed-atomic contract", policy.id))?; - ensure!( - packed.operation == PackedAtomicOperation::Add - && packed.state_space == PackedAtomicStateSpace::Global - && packed.ordering == PackedAtomicOrdering::Relaxed - && packed.scope == PackedAtomicScope::Gpu - && packed.rounding == PackedAtomicRounding::NearestEven - && packed.subnormal == PackedAtomicSubnormal::Preserve - && packed.atomicity == PackedAtomicAtomicity::PerElement - && packed.pointer_contract == PackedAtomicPointerContract::MutableGlobalU32Aligned4 - && packed.access_contract - == PackedAtomicAccessContract::NoMixedWholeWordOrNonAtomicAccess - && packed.scope_contract == PackedAtomicScopeContract::RacingAtomicsMutuallyInclusive - && packed.codegen_contract == PackedAtomicCodegenContract::ExactNativeInstruction - && packed.return_contract - == PackedAtomicReturnContract::OldValuesPerElementMayBeNoncoherent - && packed.adapter == PackedAtomicAdapter::OldPackedU32, - "{} requests an unsupported packed-atomic semantic or safety contract", - policy.id - ); - let (format, native_sm, minimum_ptx, minimum_sm, public_name) = match packed.format { - PackedAtomicFormat::F16x2 => ("f16x2", 60, "6.2", "sm_70", "atom_add_f16x2"), - PackedAtomicFormat::Bf16x2 => ("bf16x2", 90, "7.8", "sm_90", "atom_add_bf16x2"), - }; - ensure!( - packed.native_minimum_sm == native_sm, - "{} PTX-native hardware floor does not match the selected packed format", - policy.id - ); - ensure!( - source - == &IntrinsicSource::PtxNative { - instruction: format!("atom.global.add.noftz.{format}"), - }, - "{} PTX-native source does not match its packed format", - policy.id - ); - ensure!( - policy.rust_module == "atomic" - && policy.rust_name == public_name - && policy.rust_arguments == ["*mut u32", "u32"] - && policy.rust_result == "u32" - && !policy.safe - && policy.must_use - && policy.safe_allowlist_reason.is_none() - && policy.public_rust_path == format!("cuda_intrinsics::atomic::{public_name}") - && policy.compatibility_rust_paths == [format!("cuda_device::atomic::{public_name}")], - "{} must preserve the unsafe must-use packed atomic raw/compatibility API", - policy.id - ); - ensure!( - policy.dialect_op_type == "PackedAtomicAddOp" - && policy.dialect_op_name == "nvvm.packed_atomic_add" - && policy.dialect_operands == ["ptr", "i32"] - && policy.dialect_results == ["i32"] - && policy.lowering == "generated_packed_atomic_inline_ptx", - "{} is outside the one closed packed-atomic dialect recipe", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == "read_write" - && !policy.convergent - && policy.execution_scope == "thread" - && policy.minimum_ptx == minimum_ptx - && policy.minimum_sm.as_deref() == Some(minimum_sm) - && policy.targets == "all" - && policy.ptx_result == "u32" - && policy.ldmatrix_variant.is_none() - && policy.ldmatrix_safety.is_none() - && policy.ldmatrix_adapter.is_none() - && policy.redux.is_none() - && policy.vote.is_none() - && policy.active_mask.is_none() - && policy.warp_match.is_none() - && policy.warp_barrier.is_none() - && policy.warp_shuffle.is_none() - && policy.dot_product.is_none() - && policy.packed_alu.is_none() - && policy.packed_conversion.is_none() - && policy.cp_async_copy.is_none() - && policy.cp_async_control.is_none() - && policy.mbarrier_basic.is_none() - && policy.selected_address_space.is_none(), - "{} packed-atomic effects, carrier, or native target floor disagree", - policy.id - ); - ensure!( - policy.expected_ptx.mnemonic == "atom" - && policy.expected_ptx.modifiers == ["global", "add", "noftz", format] - && policy.expected_ptx.operands - == [ - OperandPattern::Register, - OperandPattern::Address, - OperandPattern::Register, - ], - "{} expected PTX must match the exact packed global add spelling", - policy.id - ); - let backend_pairs: BTreeSet<_> = policy - .backend_lowerings - .iter() - .map(|lowering| (lowering.backend, lowering.mechanism)) - .collect(); - ensure!( - backend_pairs - == BTreeSet::from([ - ( - IntrinsicBackend::LlvmNvptx, - BackendLoweringMechanism::InlinePtx, - ), - ( - IntrinsicBackend::LibNvvm, - BackendLoweringMechanism::InlinePtx, - ), - ]), - "{} must define exactly the reviewed LLVM-NVPTX and libNVVM inline-PTX routes", - policy.id - ); - for lowering in &policy.backend_lowerings { - let expected_sm = match (packed.format, lowering.backend) { - (PackedAtomicFormat::F16x2, IntrinsicBackend::LlvmNvptx) => "sm_70", - (PackedAtomicFormat::F16x2, IntrinsicBackend::LibNvvm) => "sm_75", - (PackedAtomicFormat::Bf16x2, _) => "sm_90", - }; - ensure!( - lowering.minimum_ptx.as_deref() == Some(minimum_ptx) - && lowering.minimum_sm.as_deref() == Some(expected_sm) - && !lowering.evidence_profile.trim().is_empty(), - "{} backend {:?} does not carry its exact reviewed profile floor", - policy.id, - lowering.backend - ); - } - Ok(()) -} - -enum PackedAluRecipeSource { - Imported { - record: &'static str, - symbol: &'static str, - resolved_symbol: Option<&'static str>, - arguments: &'static [&'static str], - results: &'static [&'static str], - properties: &'static [&'static str], - selection: &'static str, - selection_asm: &'static str, - }, - PtxNative, -} - -struct PackedAluRecipe { - id: &'static str, - abi_id: &'static str, - operation_key: &'static str, - rust_name: &'static str, - dialect_op_type: &'static str, - dialect_op_name: &'static str, - arity: usize, - must_use: bool, - ptx_mnemonic: &'static str, - modifiers: &'static [&'static str], - native_minimum_sm: u16, - minimum_ptx: &'static str, - minimum_sm: &'static str, - ptx_isa_section: &'static str, - ptx_isa_url: &'static str, - source: PackedAluRecipeSource, -} - -/// Returns the closed recipe for a packed-ALU (format, operation) pair. -/// -/// `None` means the pair is outside the family: the operation exists for some -/// other format but this one has no reviewed lowering for it. -fn packed_alu_recipe( - format: PackedAluFormat, - operation: PackedAluOperation, -) -> Option { - match format { - PackedAluFormat::Bf16x2 => packed_bf16x2_alu_recipe(operation), - PackedAluFormat::F16x2 => packed_f16x2_alu_recipe(operation), - PackedAluFormat::F32x2 => packed_f32x2_alu_recipe(operation), - } -} - -fn packed_f32x2_alu_recipe(operation: PackedAluOperation) -> Option { - let ( - id, - abi_id, - operation_key, - rust_name, - dialect_op_type, - dialect_op_name, - arity, - head, - modifiers, - section, - url, - ) = match operation { - PackedAluOperation::Add => ( - "add_f32x2", - "i0995", - "packed.alu.f32x2.add", - "add_f32x2", - "AddF32x2Op", - "nvvm.add_f32x2", - 2, - "add.rn.f32x2", - &["rn", "f32x2"][..], - "9.7.3.1 Floating Point Instructions: add", - "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-add", - ), - PackedAluOperation::AddFtz => ( - "add_ftz_f32x2", - "i0996", - "packed.alu.f32x2.add.ftz", - "add_ftz_f32x2", - "AddFtzF32x2Op", - "nvvm.add_ftz_f32x2", - 2, - "add.rn.ftz.f32x2", - &["rn", "ftz", "f32x2"][..], - "9.7.3.1 Floating Point Instructions: add", - "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-add", - ), - PackedAluOperation::Sub => ( - "sub_f32x2", - "i0997", - "packed.alu.f32x2.sub", - "sub_f32x2", - "SubF32x2Op", - "nvvm.sub_f32x2", - 2, - "sub.rn.f32x2", - &["rn", "f32x2"][..], - "9.7.3.2 Floating Point Instructions: sub", - "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-sub", - ), - PackedAluOperation::SubFtz => ( - "sub_ftz_f32x2", - "i0998", - "packed.alu.f32x2.sub.ftz", - "sub_ftz_f32x2", - "SubFtzF32x2Op", - "nvvm.sub_ftz_f32x2", - 2, - "sub.rn.ftz.f32x2", - &["rn", "ftz", "f32x2"][..], - "9.7.3.2 Floating Point Instructions: sub", - "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-sub", - ), - PackedAluOperation::Mul => ( - "mul_f32x2", - "i0999", - "packed.alu.f32x2.mul", - "mul_f32x2", - "MulF32x2Op", - "nvvm.mul_f32x2", - 2, - "mul.rn.f32x2", - &["rn", "f32x2"][..], - "9.7.3.3 Floating Point Instructions: mul", - "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-mul", - ), - PackedAluOperation::MulFtz => ( - "mul_ftz_f32x2", - "i1000", - "packed.alu.f32x2.mul.ftz", - "mul_ftz_f32x2", - "MulFtzF32x2Op", - "nvvm.mul_ftz_f32x2", - 2, - "mul.rn.ftz.f32x2", - &["rn", "ftz", "f32x2"][..], - "9.7.3.3 Floating Point Instructions: mul", - "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-mul", - ), - PackedAluOperation::Fma => ( - "fma_f32x2", - "i1001", - "packed.alu.f32x2.fma", - "fma_f32x2", - "FmaF32x2Op", - "nvvm.fma_f32x2", - 3, - "fma.rn.f32x2", - &["rn", "f32x2"][..], - "9.7.3.4 Floating Point Instructions: fma", - "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-fma", - ), - PackedAluOperation::FmaFtz => ( - "fma_ftz_f32x2", - "i1002", - "packed.alu.f32x2.fma.ftz", - "fma_ftz_f32x2", - "FmaFtzF32x2Op", - "nvvm.fma_ftz_f32x2", - 3, - "fma.rn.ftz.f32x2", - &["rn", "ftz", "f32x2"][..], - "9.7.3.4 Floating Point Instructions: fma", - "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-fma", - ), - _ => return None, - }; - Some(PackedAluRecipe { - id, - abi_id, - operation_key, - rust_name, - dialect_op_type, - dialect_op_name, - arity, - must_use: true, - ptx_mnemonic: head, - modifiers, - native_minimum_sm: 100, - minimum_ptx: "8.6", - minimum_sm: "sm_100", - ptx_isa_section: section, - ptx_isa_url: url, - source: PackedAluRecipeSource::PtxNative, - }) -} - -fn packed_bf16x2_alu_recipe(operation: PackedAluOperation) -> Option { - const PURE: &[&str] = &["IntrNoMem", "IntrSpeculatable"]; - const COMMUTATIVE_PURE: &[&str] = &["Commutative", "IntrNoMem", "IntrSpeculatable"]; - Some(match operation { - PackedAluOperation::Fma => PackedAluRecipe { - id: "fma_bf16x2", - abi_id: "i0062", - operation_key: "packed.alu.bf16x2.fma", - rust_name: "fma_bf16x2", - dialect_op_type: "FmaBf16x2Op", - dialect_op_name: "nvvm.fma_bf16x2", - arity: 3, - must_use: false, - ptx_mnemonic: "fma.rn.bf16x2", - modifiers: &["rn", "bf16x2"], - native_minimum_sm: 80, - minimum_ptx: "7.0", - minimum_sm: "sm_80", - ptx_isa_section: "9.7.4.4 Half Precision Floating Point Instructions: fma", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-fma", - source: PackedAluRecipeSource::Imported { - record: "int_nvvm_fma_rn_bf16x2", - symbol: "llvm.nvvm.fma.rn.bf16x2", - resolved_symbol: None, - arguments: &["v2bf16", "v2bf16", "v2bf16"], - results: &["v2bf16"], - properties: PURE, - selection: "INT_NVVM_FMA_rn_bf16x2", - selection_asm: "fma.rn.bf16x2 \t$dst, $src0, $src1, $src2;", - }, - }, - PackedAluOperation::FmaRelu => PackedAluRecipe { - id: "fma_relu_bf16x2", - abi_id: "i0063", - operation_key: "packed.alu.bf16x2.fma.relu", - rust_name: "fma_relu_bf16x2", - dialect_op_type: "FmaReluBf16x2Op", - dialect_op_name: "nvvm.fma_relu_bf16x2", - arity: 3, - must_use: false, - ptx_mnemonic: "fma.rn.relu.bf16x2", - modifiers: &["rn", "relu", "bf16x2"], - native_minimum_sm: 80, - minimum_ptx: "7.0", - minimum_sm: "sm_80", - ptx_isa_section: "9.7.4.4 Half Precision Floating Point Instructions: fma", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-fma", - source: PackedAluRecipeSource::Imported { - record: "int_nvvm_fma_rn_relu_bf16x2", - symbol: "llvm.nvvm.fma.rn.relu.bf16x2", - resolved_symbol: None, - arguments: &["v2bf16", "v2bf16", "v2bf16"], - results: &["v2bf16"], - properties: PURE, - selection: "INT_NVVM_FMA_rn_relu_bf16x2", - selection_asm: "fma.rn.relu.bf16x2 \t$dst, $src0, $src1, $src2;", - }, - }, - PackedAluOperation::Add => PackedAluRecipe { - id: "add_bf16x2", - abi_id: "i0064", - operation_key: "packed.alu.bf16x2.add", - rust_name: "add_bf16x2", - dialect_op_type: "AddBf16x2Op", - dialect_op_name: "nvvm.add_bf16x2", - arity: 2, - must_use: false, - ptx_mnemonic: "add.rn.bf16x2", - modifiers: &["rn", "bf16x2"], - native_minimum_sm: 90, - minimum_ptx: "7.8", - minimum_sm: "sm_90", - ptx_isa_section: "9.7.4.1 Half Precision Floating Point Instructions: add", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-add", - source: PackedAluRecipeSource::PtxNative, - }, - PackedAluOperation::Sub => PackedAluRecipe { - id: "sub_bf16x2", - abi_id: "i0065", - operation_key: "packed.alu.bf16x2.sub", - rust_name: "sub_bf16x2", - dialect_op_type: "SubBf16x2Op", - dialect_op_name: "nvvm.sub_bf16x2", - arity: 2, - must_use: false, - ptx_mnemonic: "sub.rn.bf16x2", - modifiers: &["rn", "bf16x2"], - native_minimum_sm: 90, - minimum_ptx: "7.8", - minimum_sm: "sm_90", - ptx_isa_section: "9.7.4.2 Half Precision Floating Point Instructions: sub", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-sub", - source: PackedAluRecipeSource::PtxNative, - }, - PackedAluOperation::Mul => PackedAluRecipe { - id: "mul_bf16x2", - abi_id: "i0066", - operation_key: "packed.alu.bf16x2.mul", - rust_name: "mul_bf16x2", - dialect_op_type: "MulBf16x2Op", - dialect_op_name: "nvvm.mul_bf16x2", - arity: 2, - must_use: false, - ptx_mnemonic: "mul.rn.bf16x2", - modifiers: &["rn", "bf16x2"], - native_minimum_sm: 90, - minimum_ptx: "7.8", - minimum_sm: "sm_90", - ptx_isa_section: "9.7.4.3 Half Precision Floating Point Instructions: mul", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-mul", - source: PackedAluRecipeSource::PtxNative, - }, - PackedAluOperation::Min => PackedAluRecipe { - id: "min_bf16x2", - abi_id: "i0067", - operation_key: "packed.alu.bf16x2.min", - rust_name: "min_bf16x2", - dialect_op_type: "MinBf16x2Op", - dialect_op_name: "nvvm.min_bf16x2", - arity: 2, - must_use: false, - ptx_mnemonic: "min.bf16x2", - modifiers: &["bf16x2"], - native_minimum_sm: 80, - minimum_ptx: "7.0", - minimum_sm: "sm_80", - ptx_isa_section: "9.7.4.7 Half Precision Floating Point Instructions: min", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-min", - source: PackedAluRecipeSource::Imported { - record: "int_nvvm_fmin_bf16x2", - symbol: "llvm.nvvm.fmin.bf16x2", - resolved_symbol: None, - arguments: &["v2bf16", "v2bf16"], - results: &["v2bf16"], - properties: COMMUTATIVE_PURE, - selection: "INT_NVVM_FMIN_bf16x2", - selection_asm: "min.bf16x2 \t$dst, $src0, $src1;", - }, - }, - PackedAluOperation::Max => PackedAluRecipe { - id: "max_bf16x2", - abi_id: "i0068", - operation_key: "packed.alu.bf16x2.max", - rust_name: "max_bf16x2", - dialect_op_type: "MaxBf16x2Op", - dialect_op_name: "nvvm.max_bf16x2", - arity: 2, - must_use: false, - ptx_mnemonic: "max.bf16x2", - modifiers: &["bf16x2"], - native_minimum_sm: 80, - minimum_ptx: "7.0", - minimum_sm: "sm_80", - ptx_isa_section: "9.7.4.8 Half Precision Floating Point Instructions: max", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-max", - source: PackedAluRecipeSource::Imported { - record: "int_nvvm_fmax_bf16x2", - symbol: "llvm.nvvm.fmax.bf16x2", - resolved_symbol: None, - arguments: &["v2bf16", "v2bf16"], - results: &["v2bf16"], - properties: COMMUTATIVE_PURE, - selection: "INT_NVVM_FMAN_bf16x2", - selection_asm: "max.bf16x2 \t$dst, $src0, $src1;", - }, - }, - PackedAluOperation::Neg => PackedAluRecipe { - id: "neg_bf16x2", - abi_id: "i0069", - operation_key: "packed.alu.bf16x2.neg", - rust_name: "neg_bf16x2", - dialect_op_type: "NegBf16x2Op", - dialect_op_name: "nvvm.neg_bf16x2", - arity: 1, - must_use: false, - ptx_mnemonic: "neg.bf16x2", - modifiers: &["bf16x2"], - native_minimum_sm: 80, - minimum_ptx: "7.0", - minimum_sm: "sm_80", - ptx_isa_section: "9.7.4.5 Half Precision Floating Point Instructions: neg", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-neg", - source: PackedAluRecipeSource::Imported { - record: "int_nvvm_neg_bf16x2", - symbol: "llvm.nvvm.neg.bf16x2", - resolved_symbol: None, - arguments: &["v2bf16"], - results: &["v2bf16"], - properties: PURE, - selection: "INT_NVVM_NEG_BF16X2", - selection_asm: "neg.bf16x2 \t$dst, $src0;", - }, - }, - PackedAluOperation::Abs => PackedAluRecipe { - id: "abs_bf16x2", - abi_id: "i0070", - operation_key: "packed.alu.bf16x2.abs", - rust_name: "abs_bf16x2", - dialect_op_type: "AbsBf16x2Op", - dialect_op_name: "nvvm.abs_bf16x2", - arity: 1, - must_use: false, - ptx_mnemonic: "abs.bf16x2", - modifiers: &["bf16x2"], - native_minimum_sm: 80, - minimum_ptx: "7.0", - minimum_sm: "sm_80", - ptx_isa_section: "9.7.4.6 Half Precision Floating Point Instructions: abs", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-abs", - source: PackedAluRecipeSource::Imported { - record: "int_nvvm_fabs", - symbol: "llvm.nvvm.fabs", - resolved_symbol: Some("llvm.nvvm.fabs.v2bf16"), - arguments: &["anonymous_14"], - results: &["anyfloat"], - properties: PURE, - selection: "ABS_BF16X2", - selection_asm: "abs.bf16x2 \t$dst, $src0;", - }, - }, - // The f32x2-only operation keys are outside this admission. The ftz - // and sat fma forms also do not exist for bf16x2 in the PTX ISA at - // all: ptxas rejects hand-written `fma.rn.ftz.bf16x2` with - // "Illegal modifier '.ftz' for instruction 'fma'", and accepts - // `fma.rn.relu.bf16x2` beside it. LLVM declares the intrinsics anyway, - // so instruction selection is what fails first ("Cannot select: - // intrinsic %llvm.nvvm.fma.rn.ftz.bf16x2"), but the assembler is the - // authority here. Reject them rather than admit a recipe that has no - // instruction to lower to. - PackedAluOperation::AddFtz - | PackedAluOperation::SubFtz - | PackedAluOperation::MulFtz - | PackedAluOperation::FmaFtz - | PackedAluOperation::FmaSat - | PackedAluOperation::FmaFtzSat - | PackedAluOperation::FmaFtzRelu => return None, - }) -} - -fn packed_f16x2_alu_recipe(operation: PackedAluOperation) -> Option { - const PURE: &[&str] = &["IntrNoMem", "IntrSpeculatable"]; - const COMMUTATIVE_PURE: &[&str] = &["Commutative", "IntrNoMem", "IntrSpeculatable"]; - Some(match operation { - PackedAluOperation::Fma => PackedAluRecipe { - id: "fma_f16x2", - abi_id: "i0072", - operation_key: "packed.alu.f16x2.fma", - rust_name: "fma_f16x2", - dialect_op_type: "FmaF16x2Op", - dialect_op_name: "nvvm.fma_f16x2", - arity: 3, - must_use: true, - ptx_mnemonic: "fma.rn.f16x2", - modifiers: &["rn", "f16x2"], - native_minimum_sm: 53, - minimum_ptx: "4.2", - minimum_sm: "sm_70", - ptx_isa_section: "9.7.4.4 Half Precision Floating Point Instructions: fma", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-fma", - source: PackedAluRecipeSource::Imported { - record: "int_nvvm_fma_rn_f16x2", - symbol: "llvm.nvvm.fma.rn.f16x2", - resolved_symbol: None, - arguments: &["v2f16", "v2f16", "v2f16"], - results: &["v2f16"], - properties: PURE, - selection: "INT_NVVM_FMA_rn_f16x2", - selection_asm: "fma.rn.f16x2 \t$dst, $src0, $src1, $src2;", - }, - }, - PackedAluOperation::FmaFtz => PackedAluRecipe { - id: "fma_ftz_f16x2", - abi_id: "i0854", - operation_key: "packed.alu.f16x2.fma.ftz", - rust_name: "fma_ftz_f16x2", - dialect_op_type: "FmaFtzF16x2Op", - dialect_op_name: "nvvm.fma_ftz_f16x2", - arity: 3, - must_use: true, - ptx_mnemonic: "fma.rn.ftz.f16x2", - modifiers: &["rn", "ftz", "f16x2"], - native_minimum_sm: 53, - minimum_ptx: "4.2", - minimum_sm: "sm_70", - ptx_isa_section: "9.7.4.4 Half Precision Floating Point Instructions: fma", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-fma", - source: PackedAluRecipeSource::Imported { - record: "int_nvvm_fma_rn_ftz_f16x2", - symbol: "llvm.nvvm.fma.rn.ftz.f16x2", - resolved_symbol: None, - arguments: &["v2f16", "v2f16", "v2f16"], - results: &["v2f16"], - properties: PURE, - selection: "INT_NVVM_FMA_rn_ftz_f16x2", - selection_asm: "fma.rn.ftz.f16x2 \t$dst, $src0, $src1, $src2;", - }, - }, - PackedAluOperation::FmaSat => PackedAluRecipe { - id: "fma_sat_f16x2", - abi_id: "i0855", - operation_key: "packed.alu.f16x2.fma.sat", - rust_name: "fma_sat_f16x2", - dialect_op_type: "FmaSatF16x2Op", - dialect_op_name: "nvvm.fma_sat_f16x2", - arity: 3, - must_use: true, - ptx_mnemonic: "fma.rn.sat.f16x2", - modifiers: &["rn", "sat", "f16x2"], - native_minimum_sm: 53, - minimum_ptx: "4.2", - minimum_sm: "sm_70", - ptx_isa_section: "9.7.4.4 Half Precision Floating Point Instructions: fma", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-fma", - source: PackedAluRecipeSource::Imported { - record: "int_nvvm_fma_rn_sat_f16x2", - symbol: "llvm.nvvm.fma.rn.sat.f16x2", - resolved_symbol: None, - arguments: &["v2f16", "v2f16", "v2f16"], - results: &["v2f16"], - properties: PURE, - selection: "INT_NVVM_FMA_rn_sat_f16x2", - selection_asm: "fma.rn.sat.f16x2 \t$dst, $src0, $src1, $src2;", - }, - }, - PackedAluOperation::FmaFtzSat => PackedAluRecipe { - id: "fma_ftz_sat_f16x2", - abi_id: "i0856", - operation_key: "packed.alu.f16x2.fma.ftz.sat", - rust_name: "fma_ftz_sat_f16x2", - dialect_op_type: "FmaFtzSatF16x2Op", - dialect_op_name: "nvvm.fma_ftz_sat_f16x2", - arity: 3, - must_use: true, - ptx_mnemonic: "fma.rn.ftz.sat.f16x2", - modifiers: &["rn", "ftz", "sat", "f16x2"], - native_minimum_sm: 53, - minimum_ptx: "4.2", - minimum_sm: "sm_70", - ptx_isa_section: "9.7.4.4 Half Precision Floating Point Instructions: fma", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-fma", - source: PackedAluRecipeSource::Imported { - record: "int_nvvm_fma_rn_ftz_sat_f16x2", - symbol: "llvm.nvvm.fma.rn.ftz.sat.f16x2", - resolved_symbol: None, - arguments: &["v2f16", "v2f16", "v2f16"], - results: &["v2f16"], - properties: PURE, - selection: "INT_NVVM_FMA_rn_ftz_sat_f16x2", - selection_asm: "fma.rn.ftz.sat.f16x2 \t$dst, $src0, $src1, $src2;", - }, - }, - PackedAluOperation::FmaRelu => PackedAluRecipe { - id: "fma_relu_f16x2", - abi_id: "i0073", - operation_key: "packed.alu.f16x2.fma.relu", - rust_name: "fma_relu_f16x2", - dialect_op_type: "FmaReluF16x2Op", - dialect_op_name: "nvvm.fma_relu_f16x2", - arity: 3, - must_use: true, - ptx_mnemonic: "fma.rn.relu.f16x2", - modifiers: &["rn", "relu", "f16x2"], - native_minimum_sm: 80, - minimum_ptx: "7.0", - minimum_sm: "sm_80", - ptx_isa_section: "9.7.4.4 Half Precision Floating Point Instructions: fma", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-fma", - source: PackedAluRecipeSource::Imported { - record: "int_nvvm_fma_rn_relu_f16x2", - symbol: "llvm.nvvm.fma.rn.relu.f16x2", - resolved_symbol: None, - arguments: &["v2f16", "v2f16", "v2f16"], - results: &["v2f16"], - properties: PURE, - selection: "INT_NVVM_FMA_rn_relu_f16x2", - selection_asm: "fma.rn.relu.f16x2 \t$dst, $src0, $src1, $src2;", - }, - }, - PackedAluOperation::FmaFtzRelu => PackedAluRecipe { - id: "fma_ftz_relu_f16x2", - abi_id: "i0857", - operation_key: "packed.alu.f16x2.fma.ftz.relu", - rust_name: "fma_ftz_relu_f16x2", - dialect_op_type: "FmaFtzReluF16x2Op", - dialect_op_name: "nvvm.fma_ftz_relu_f16x2", - arity: 3, - must_use: true, - ptx_mnemonic: "fma.rn.ftz.relu.f16x2", - modifiers: &["rn", "ftz", "relu", "f16x2"], - native_minimum_sm: 80, - minimum_ptx: "7.0", - minimum_sm: "sm_80", - ptx_isa_section: "9.7.4.4 Half Precision Floating Point Instructions: fma", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-fma", - source: PackedAluRecipeSource::Imported { - record: "int_nvvm_fma_rn_ftz_relu_f16x2", - symbol: "llvm.nvvm.fma.rn.ftz.relu.f16x2", - resolved_symbol: None, - arguments: &["v2f16", "v2f16", "v2f16"], - results: &["v2f16"], - properties: PURE, - selection: "INT_NVVM_FMA_rn_ftz_relu_f16x2", - selection_asm: "fma.rn.ftz.relu.f16x2 \t$dst, $src0, $src1, $src2;", - }, - }, - PackedAluOperation::Add => PackedAluRecipe { - id: "add_f16x2", - abi_id: "i0074", - operation_key: "packed.alu.f16x2.add", - rust_name: "add_f16x2", - dialect_op_type: "AddF16x2Op", - dialect_op_name: "nvvm.add_f16x2", - arity: 2, - must_use: true, - ptx_mnemonic: "add.rn.f16x2", - modifiers: &["rn", "f16x2"], - native_minimum_sm: 53, - minimum_ptx: "4.2", - minimum_sm: "sm_70", - ptx_isa_section: "9.7.4.1 Half Precision Floating Point Instructions: add", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-add", - source: PackedAluRecipeSource::PtxNative, - }, - PackedAluOperation::Sub => PackedAluRecipe { - id: "sub_f16x2", - abi_id: "i0075", - operation_key: "packed.alu.f16x2.sub", - rust_name: "sub_f16x2", - dialect_op_type: "SubF16x2Op", - dialect_op_name: "nvvm.sub_f16x2", - arity: 2, - must_use: true, - ptx_mnemonic: "sub.rn.f16x2", - modifiers: &["rn", "f16x2"], - native_minimum_sm: 53, - minimum_ptx: "4.2", - minimum_sm: "sm_70", - ptx_isa_section: "9.7.4.2 Half Precision Floating Point Instructions: sub", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-sub", - source: PackedAluRecipeSource::PtxNative, - }, - PackedAluOperation::Mul => PackedAluRecipe { - id: "mul_f16x2", - abi_id: "i0076", - operation_key: "packed.alu.f16x2.mul", - rust_name: "mul_f16x2", - dialect_op_type: "MulF16x2Op", - dialect_op_name: "nvvm.mul_f16x2", - arity: 2, - must_use: true, - ptx_mnemonic: "mul.rn.f16x2", - modifiers: &["rn", "f16x2"], - native_minimum_sm: 53, - minimum_ptx: "4.2", - minimum_sm: "sm_70", - ptx_isa_section: "9.7.4.3 Half Precision Floating Point Instructions: mul", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-mul", - source: PackedAluRecipeSource::PtxNative, - }, - PackedAluOperation::Min => PackedAluRecipe { - id: "min_f16x2", - abi_id: "i0077", - operation_key: "packed.alu.f16x2.min", - rust_name: "min_f16x2", - dialect_op_type: "MinF16x2Op", - dialect_op_name: "nvvm.min_f16x2", - arity: 2, - must_use: true, - ptx_mnemonic: "min.f16x2", - modifiers: &["f16x2"], - native_minimum_sm: 80, - minimum_ptx: "7.0", - minimum_sm: "sm_80", - ptx_isa_section: "9.7.4.7 Half Precision Floating Point Instructions: min", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-min", - source: PackedAluRecipeSource::Imported { - record: "int_nvvm_fmin_f16x2", - symbol: "llvm.nvvm.fmin.f16x2", - resolved_symbol: None, - arguments: &["v2f16", "v2f16"], - results: &["v2f16"], - properties: COMMUTATIVE_PURE, - selection: "INT_NVVM_FMIN_f16x2", - selection_asm: "min.f16x2 \t$dst, $src0, $src1;", - }, - }, - PackedAluOperation::Max => PackedAluRecipe { - id: "max_f16x2", - abi_id: "i0078", - operation_key: "packed.alu.f16x2.max", - rust_name: "max_f16x2", - dialect_op_type: "MaxF16x2Op", - dialect_op_name: "nvvm.max_f16x2", - arity: 2, - must_use: true, - ptx_mnemonic: "max.f16x2", - modifiers: &["f16x2"], - native_minimum_sm: 80, - minimum_ptx: "7.0", - minimum_sm: "sm_80", - ptx_isa_section: "9.7.4.8 Half Precision Floating Point Instructions: max", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-max", - source: PackedAluRecipeSource::Imported { - record: "int_nvvm_fmax_f16x2", - symbol: "llvm.nvvm.fmax.f16x2", - resolved_symbol: None, - arguments: &["v2f16", "v2f16"], - results: &["v2f16"], - properties: COMMUTATIVE_PURE, - selection: "INT_NVVM_FMAN_f16x2", - selection_asm: "max.f16x2 \t$dst, $src0, $src1;", - }, - }, - PackedAluOperation::Neg => PackedAluRecipe { - id: "neg_f16x2", - abi_id: "i0079", - operation_key: "packed.alu.f16x2.neg", - rust_name: "neg_f16x2", - dialect_op_type: "NegF16x2Op", - dialect_op_name: "nvvm.neg_f16x2", - arity: 1, - must_use: true, - ptx_mnemonic: "neg.f16x2", - modifiers: &["f16x2"], - native_minimum_sm: 53, - minimum_ptx: "6.0", - minimum_sm: "sm_70", - ptx_isa_section: "9.7.4.5 Half Precision Floating Point Instructions: neg", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-neg", - source: PackedAluRecipeSource::PtxNative, - }, - PackedAluOperation::Abs => PackedAluRecipe { - id: "abs_f16x2", - abi_id: "i0080", - operation_key: "packed.alu.f16x2.abs", - rust_name: "abs_f16x2", - dialect_op_type: "AbsF16x2Op", - dialect_op_name: "nvvm.abs_f16x2", - arity: 1, - must_use: true, - ptx_mnemonic: "abs.f16x2", - modifiers: &["f16x2"], - native_minimum_sm: 53, - minimum_ptx: "6.5", - minimum_sm: "sm_70", - ptx_isa_section: "9.7.4.6 Half Precision Floating Point Instructions: abs", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-abs", - source: PackedAluRecipeSource::Imported { - record: "int_nvvm_fabs", - symbol: "llvm.nvvm.fabs", - resolved_symbol: Some("llvm.nvvm.fabs.v2f16"), - arguments: &["anonymous_14"], - results: &["anyfloat"], - properties: PURE, - selection: "ABS_F16X2", - selection_asm: "abs.f16x2 \t$dst, $src0;", - }, - }, - // These operations were added for the f32x2 admission. Keeping them - // outside the f16x2 recipe avoids silently expanding this PR's scope, - // even where a similarly spelled PTX form may exist. - PackedAluOperation::AddFtz | PackedAluOperation::SubFtz | PackedAluOperation::MulFtz => { - return None; - } - }) -} - -/// Per-backend PTX and SM floors, which can sit above the recipe's own floor. -/// -/// Takes the caller's already-resolved recipe rather than looking it up again, -/// so this cannot be reached with a pair the family rejects. -fn packed_alu_backend_floor( - recipe: &PackedAluRecipe, - format: PackedAluFormat, - operation: PackedAluOperation, - backend: IntrinsicBackend, -) -> (&'static str, &'static str) { - match (format, operation, backend) { - ( - PackedAluFormat::F16x2, - PackedAluOperation::Fma - | PackedAluOperation::Add - | PackedAluOperation::Sub - | PackedAluOperation::Mul, - IntrinsicBackend::LlvmNvptx, - ) => ("6.0", "sm_70"), - ( - PackedAluFormat::F16x2, - PackedAluOperation::Fma - | PackedAluOperation::Add - | PackedAluOperation::Sub - | PackedAluOperation::Mul, - IntrinsicBackend::LibNvvm, - ) => ("4.2", "sm_75"), - // The ftz and sat fma forms share the plain fma floors: PTX 4.2 natively, - // but the pinned LLVM backend needs 6.0 and CUDA 13.3 ptxas no longer - // targets sm_70, so libNVVM is checked at sm_75. - ( - PackedAluFormat::F16x2, - PackedAluOperation::FmaFtz | PackedAluOperation::FmaSat | PackedAluOperation::FmaFtzSat, - IntrinsicBackend::LlvmNvptx, - ) => ("6.0", "sm_70"), - ( - PackedAluFormat::F16x2, - PackedAluOperation::FmaFtz | PackedAluOperation::FmaSat | PackedAluOperation::FmaFtzSat, - IntrinsicBackend::LibNvvm, - ) => ("4.2", "sm_75"), - (PackedAluFormat::F16x2, PackedAluOperation::Neg, IntrinsicBackend::LlvmNvptx) => { - ("6.0", "sm_70") - } - (PackedAluFormat::F16x2, PackedAluOperation::Neg, IntrinsicBackend::LibNvvm) => { - ("6.0", "sm_75") - } - (PackedAluFormat::F16x2, PackedAluOperation::Abs, IntrinsicBackend::LlvmNvptx) => { - ("6.5", "sm_70") - } - (PackedAluFormat::F16x2, PackedAluOperation::Abs, IntrinsicBackend::LibNvvm) => { - ("6.5", "sm_75") - } - _ => (recipe.minimum_ptx, recipe.minimum_sm), - } -} - -fn validate_packed_alu_policy( - policy: &OverlayIntrinsic, - source: &IntrinsicSource, - declaration: Option<&ImportedIntrinsic>, -) -> Result<()> { - let packed = policy - .packed_alu - .as_ref() - .with_context(|| format!("{} has no closed packed-ALU contract", policy.id))?; - let recipe = packed_alu_recipe(packed.format, packed.operation) - .with_context(|| format!("{} is outside the closed packed-ALU recipe", policy.id))?; - let (rust_module, rust_type, dialect_type, expected_adapter) = match packed.format { - PackedAluFormat::Bf16x2 => ("bf16x2", "u32", "i32", PackedAluAdapter::DirectPackedU32), - PackedAluFormat::F16x2 => ("f16x2", "u32", "i32", PackedAluAdapter::DirectPackedU32), - PackedAluFormat::F32x2 => ("f32x2", "u64", "i64", PackedAluAdapter::DirectPackedU64), - }; - ensure!( - packed.adapter == expected_adapter, - "{} requests an unsupported packed-ALU adapter", - policy.id - ); - let rust_arguments = vec![rust_type; recipe.arity]; - let dialect_operands = vec![dialect_type; recipe.arity]; - ensure!( - policy.id == recipe.id - && policy.abi_id == recipe.abi_id - && policy.operation_key == recipe.operation_key, - "{} packed-ALU identity does not match its closed operation recipe", - policy.id - ); - ensure!( - policy.rust_module == rust_module - && policy.rust_name == recipe.rust_name - && policy.rust_arguments == rust_arguments - && policy.rust_result == rust_type - && policy.safe - && policy.must_use == recipe.must_use - && policy - .safe_allowlist_reason - .as_deref() - .is_some_and(|reason| !reason.trim().is_empty()) - && policy.public_rust_path - == format!("cuda_intrinsics::{rust_module}::{}", recipe.rust_name) - && policy.compatibility_rust_paths - == [format!("cuda_device::{rust_module}::{}", recipe.rust_name)], - "{} must preserve its reviewed safe packed-ALU API", - policy.id - ); - ensure!( - policy.dialect_op_type == recipe.dialect_op_type - && policy.dialect_op_name == recipe.dialect_op_name - && policy.dialect_operands == dialect_operands - && policy.dialect_results == [dialect_type] - && policy.lowering == "generated_packed_alu_inline_ptx", - "{} is outside the closed packed-ALU dialect and lowering recipe", - policy.id - ); - ensure!( - policy.pure - && policy.memory == "none" - && !policy.convergent - && policy.execution_scope == "thread" - && policy.minimum_ptx == recipe.minimum_ptx - && policy.minimum_sm.as_deref() == Some(recipe.minimum_sm) - && policy.ptx_result == rust_type - && policy.targets == "all" - && packed.native_minimum_sm == recipe.native_minimum_sm, - "{} packed-ALU effects, carrier, or target floor disagree", - policy.id - ); - ensure!( - policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section == recipe.ptx_isa_section - && policy.ptx_isa_url == recipe.ptx_isa_url, - "{} packed-ALU PTX provenance does not match its reviewed instruction section", - policy.id - ); - let expected_operands = vec![OperandPattern::Register; recipe.arity + 1]; - ensure!( - policy.expected_ptx.mnemonic - == recipe.ptx_mnemonic.split('.').next().expect("PTX mnemonic") - && policy.expected_ptx.modifiers == recipe.modifiers - && policy.expected_ptx.operands == expected_operands, - "{} expected PTX does not match its exact packed-ALU instruction", - policy.id - ); - - match &recipe.source { - PackedAluRecipeSource::PtxNative => { - ensure!( - source - == &IntrinsicSource::PtxNative { - instruction: recipe.ptx_mnemonic.to_owned(), - } - && declaration.is_none(), - "{} packed-ALU source does not match its PTX-native recipe", - policy.id - ); - } - PackedAluRecipeSource::Imported { - record, - symbol, - resolved_symbol, - arguments, - results, - properties, - selection, - selection_asm, - } => { - let declaration = declaration.context("imported packed ALU has no declaration")?; - ensure!( - source - == &IntrinsicSource::LlvmImported { - source_record: (*record).to_owned(), - } - && policy.llvm_symbol.as_deref() == Some(*symbol) - && policy.resolved_llvm_symbol.as_deref() == *resolved_symbol - && policy.llvm_arguments == *arguments - && policy.llvm_results == *results, - "{} packed-ALU LLVM source or signature changed", - policy.id - ); - let matching_selections: Vec<_> = declaration - .selections - .iter() - .filter(|candidate| candidate.source_record == *selection) - .collect(); - let expected_selection_count = if *record == "int_nvvm_fabs" { 6 } else { 1 }; - ensure!( - declaration.properties == *properties - && declaration.selections.len() == expected_selection_count - && matching_selections.len() == 1 - && matching_selections[0].asm == *selection_asm - && matching_selections[0].predicates - == [ - format!("Subtarget->getSmVersion() >= {}", recipe.native_minimum_sm), - format!( - "Subtarget->getPTXVersion() >= {}", - recipe.minimum_ptx.replace('.', "") - ), - ] - && matching_selections[0].constraints.is_empty(), - "{} packed-ALU imported properties or selection changed", - policy.id - ); - } - } - let llvm_floor = packed_alu_backend_floor( - &recipe, - packed.format, - packed.operation, - IntrinsicBackend::LlvmNvptx, - ); - let libnvvm_floor = packed_alu_backend_floor( - &recipe, - packed.format, - packed.operation, - IntrinsicBackend::LibNvvm, - ); - ensure_exact_inline_ptx_backends( - policy, - [ - ( - IntrinsicBackend::LlvmNvptx, - llvm_floor.0, - Some(llvm_floor.1), - ), - ( - IntrinsicBackend::LibNvvm, - libnvvm_floor.0, - Some(libnvvm_floor.1), - ), - ], - "packed-ALU", - )?; - ensure_no_other_family_contract(policy, "packed-ALU")?; - Ok(()) -} - -struct PackedConversionRecipe { - id: &'static str, - abi_id: &'static str, - operation_key: &'static str, - rust_name: &'static str, - compatibility_path: &'static str, - dialect_op_type: &'static str, - dialect_op_name: &'static str, - source_record: &'static str, - llvm_symbol: &'static str, - llvm_result: &'static str, - summary: &'static str, -} - -struct CpAsyncCopyRecipe { - id: &'static str, - abi_id: &'static str, - operation_key: &'static str, - rust_name: &'static str, - dialect_op_type: &'static str, - dialect_op_name: &'static str, - source_record: &'static str, - llvm_symbol: &'static str, - selections: &'static [&'static str], - summary: &'static str, -} - -fn cp_async_copy_recipe(copy: &crate::model::CpAsyncCopy) -> Option { - match (copy.cache_policy, copy.copy_size, copy.source_size) { - (CpAsyncCachePolicy::Ca, CpAsyncCopySize::B4, CpAsyncSourceSize::Full) => { - Some(CpAsyncCopyRecipe { - id: "cp_async_ca_4", - abi_id: "i0086", - operation_key: "memory.copy.async.global_to_shared.ca.4.full", - rust_name: "cp_async_ca_4", - dialect_op_type: "CpAsyncCa4Op", - dialect_op_name: "nvvm.cp_async_ca_4", - source_record: "int_nvvm_cp_async_ca_shared_global_4", - llvm_symbol: "llvm.nvvm.cp.async.ca.shared.global.4", - selections: &["CP_ASYNC_CA_SHARED_GLOBAL_4"], - summary: "Starts a four-byte cache-all asynchronous copy from global to shared memory.", - }) - } - (CpAsyncCachePolicy::Ca, CpAsyncCopySize::B4, CpAsyncSourceSize::Runtime) => { - Some(CpAsyncCopyRecipe { - id: "cp_async_ca_zfill_4", - abi_id: "i0087", - operation_key: "memory.copy.async.global_to_shared.ca.4.runtime_source_size", - rust_name: "cp_async_ca_zfill_4", - dialect_op_type: "CpAsyncCaZfill4Op", - dialect_op_name: "nvvm.cp_async_ca_zfill_4", - source_record: "int_nvvm_cp_async_ca_shared_global_4_s", - llvm_symbol: "llvm.nvvm.cp.async.ca.shared.global.4.s", - selections: &[ - "CP_ASYNC_CA_SHARED_GLOBAL_4_s", - "CP_ASYNC_CA_SHARED_GLOBAL_4_si", - ], - summary: "Starts a four-byte cache-all asynchronous copy and zero-fills bytes beyond the runtime source size.", - }) - } - (CpAsyncCachePolicy::Ca, CpAsyncCopySize::B8, CpAsyncSourceSize::Full) => { - Some(CpAsyncCopyRecipe { - id: "cp_async_ca_8", - abi_id: "i0088", - operation_key: "memory.copy.async.global_to_shared.ca.8.full", - rust_name: "cp_async_ca_8", - dialect_op_type: "CpAsyncCa8Op", - dialect_op_name: "nvvm.cp_async_ca_8", - source_record: "int_nvvm_cp_async_ca_shared_global_8", - llvm_symbol: "llvm.nvvm.cp.async.ca.shared.global.8", - selections: &["CP_ASYNC_CA_SHARED_GLOBAL_8"], - summary: "Starts an eight-byte cache-all asynchronous copy from global to shared memory.", - }) - } - (CpAsyncCachePolicy::Ca, CpAsyncCopySize::B8, CpAsyncSourceSize::Runtime) => { - Some(CpAsyncCopyRecipe { - id: "cp_async_ca_zfill_8", - abi_id: "i0089", - operation_key: "memory.copy.async.global_to_shared.ca.8.runtime_source_size", - rust_name: "cp_async_ca_zfill_8", - dialect_op_type: "CpAsyncCaZfill8Op", - dialect_op_name: "nvvm.cp_async_ca_zfill_8", - source_record: "int_nvvm_cp_async_ca_shared_global_8_s", - llvm_symbol: "llvm.nvvm.cp.async.ca.shared.global.8.s", - selections: &[ - "CP_ASYNC_CA_SHARED_GLOBAL_8_s", - "CP_ASYNC_CA_SHARED_GLOBAL_8_si", - ], - summary: "Starts an eight-byte cache-all asynchronous copy and zero-fills bytes beyond the runtime source size.", - }) - } - (CpAsyncCachePolicy::Ca, CpAsyncCopySize::B16, CpAsyncSourceSize::Full) => { - Some(CpAsyncCopyRecipe { - id: "cp_async_ca_16", - abi_id: "i0090", - operation_key: "memory.copy.async.global_to_shared.ca.16.full", - rust_name: "cp_async_ca_16", - dialect_op_type: "CpAsyncCa16Op", - dialect_op_name: "nvvm.cp_async_ca_16", - source_record: "int_nvvm_cp_async_ca_shared_global_16", - llvm_symbol: "llvm.nvvm.cp.async.ca.shared.global.16", - selections: &["CP_ASYNC_CA_SHARED_GLOBAL_16"], - summary: "Starts a sixteen-byte cache-all asynchronous copy from global to shared memory.", - }) - } - (CpAsyncCachePolicy::Ca, CpAsyncCopySize::B16, CpAsyncSourceSize::Runtime) => { - Some(CpAsyncCopyRecipe { - id: "cp_async_ca_zfill_16", - abi_id: "i0091", - operation_key: "memory.copy.async.global_to_shared.ca.16.runtime_source_size", - rust_name: "cp_async_ca_zfill_16", - dialect_op_type: "CpAsyncCaZfill16Op", - dialect_op_name: "nvvm.cp_async_ca_zfill_16", - source_record: "int_nvvm_cp_async_ca_shared_global_16_s", - llvm_symbol: "llvm.nvvm.cp.async.ca.shared.global.16.s", - selections: &[ - "CP_ASYNC_CA_SHARED_GLOBAL_16_s", - "CP_ASYNC_CA_SHARED_GLOBAL_16_si", - ], - summary: "Starts a sixteen-byte cache-all asynchronous copy and zero-fills bytes beyond the runtime source size.", - }) - } - (CpAsyncCachePolicy::Cg, CpAsyncCopySize::B16, CpAsyncSourceSize::Full) => { - Some(CpAsyncCopyRecipe { - id: "cp_async_cg_16", - abi_id: "i0092", - operation_key: "memory.copy.async.global_to_shared.cg.16.full", - rust_name: "cp_async_cg_16", - dialect_op_type: "CpAsyncCg16Op", - dialect_op_name: "nvvm.cp_async_cg_16", - source_record: "int_nvvm_cp_async_cg_shared_global_16", - llvm_symbol: "llvm.nvvm.cp.async.cg.shared.global.16", - selections: &["CP_ASYNC_CG_SHARED_GLOBAL_16"], - summary: "Starts a sixteen-byte cache-global asynchronous copy from global to shared memory.", - }) - } - (CpAsyncCachePolicy::Cg, CpAsyncCopySize::B16, CpAsyncSourceSize::Runtime) => { - Some(CpAsyncCopyRecipe { - id: "cp_async_cg_zfill_16", - abi_id: "i0093", - operation_key: "memory.copy.async.global_to_shared.cg.16.runtime_source_size", - rust_name: "cp_async_cg_zfill_16", - dialect_op_type: "CpAsyncCgZfill16Op", - dialect_op_name: "nvvm.cp_async_cg_zfill_16", - source_record: "int_nvvm_cp_async_cg_shared_global_16_s", - llvm_symbol: "llvm.nvvm.cp.async.cg.shared.global.16.s", - selections: &[ - "CP_ASYNC_CG_SHARED_GLOBAL_16_s", - "CP_ASYNC_CG_SHARED_GLOBAL_16_si", - ], - summary: "Starts a sixteen-byte cache-global asynchronous copy and zero-fills bytes beyond the runtime source size.", - }) - } - _ => None, - } -} - -fn validate_cp_async_copy_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, -) -> Result<()> { - let copy = policy - .cp_async_copy - .as_ref() - .with_context(|| format!("{} has no closed cp.async copy contract", policy.id))?; - let recipe = cp_async_copy_recipe(copy).with_context(|| { - format!( - "{} requests an unsupported classic cp.async copy form", - policy.id - ) - })?; - let dynamic_source_size = copy.source_size == CpAsyncSourceSize::Runtime; - let expected_adapter = if dynamic_source_size { - CpAsyncAdapter::DirectPointersAndSourceSize - } else { - CpAsyncAdapter::DirectPointers - }; - ensure!( - copy.adapter == expected_adapter, - "{} cp.async source-size form and adapter disagree", - policy.id - ); - ensure!( - copy.runtime_validation == RuntimeValidation::Unexecuted, - "{} cannot claim unrecorded cp.async runtime validation", - policy.id - ); - ensure!( - policy.id == recipe.id - && policy.abi_id == recipe.abi_id - && policy.operation_key == recipe.operation_key - && policy.source.is_none() - && policy.source_record.as_deref() == Some(recipe.source_record) - && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) - && policy.resolved_llvm_symbol.is_none(), - "{} cp.async identity does not match its closed recipe", - policy.id - ); - let rust_arguments = if dynamic_source_size { - vec!["*mut u32", "*const u8", "u32"] - } else { - vec!["*mut u32", "*const u32"] - }; - ensure!( - policy.rust_module == "async_copy" - && policy.rust_name == recipe.rust_name - && policy.rust_arguments == rust_arguments - && policy.rust_result == "()" - && !policy.safe - && !policy.must_use - && policy.safe_allowlist_reason.is_none() - && policy.public_rust_path - == format!("cuda_intrinsics::async_copy::{}", recipe.rust_name) - && policy.compatibility_rust_paths - == [format!("cuda_device::async_copy::{}", recipe.rust_name)], - "{} must preserve its unsafe cp.async raw and compatibility API", - policy.id - ); - let llvm_arguments = if dynamic_source_size { - vec!["shared_ptr", "global_ptr", "i32"] - } else { - vec!["shared_ptr", "global_ptr"] - }; - let dialect_operands = if dynamic_source_size { - vec!["ptr", "ptr", "i32"] - } else { - vec!["ptr", "ptr"] - }; - ensure!( - policy.dialect_op_type == recipe.dialect_op_type - && policy.dialect_op_name == recipe.dialect_op_name - && policy.dialect_operands == dialect_operands - && policy.dialect_results.is_empty() - && policy.llvm_arguments == llvm_arguments - && policy.llvm_results.is_empty() - && policy.lowering == "generated_cp_async_copy", - "{} is outside the closed cp.async signature and lowering recipe", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == "read_write" - && !policy.convergent - && policy.execution_scope == "thread" - && policy.minimum_ptx == "7.0" - && policy.minimum_sm.as_deref() == Some("sm_80") - && policy.ptx_result == "()" - && policy.targets == "all", - "{} cp.async effects or target floor disagree with the closed recipe", - policy.id - ); - ensure!( - policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section - == "9.7.9.26.3.1 Data Movement and Conversion Instructions: cp.async" - && policy.ptx_isa_url - == "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async", - "{} cp.async PTX provenance disagrees with the reviewed recipe", - policy.id - ); - ensure!( - policy.summary == recipe.summary, - "{} cp.async summary does not match its closed recipe", - policy.id - ); - ensure!( - declaration.properties - == [ - "IntrArgMemOnly", - "IntrNoCallback", - "NoAlias", - "NoAlias", - "ReadOnly", - "WriteOnly", - ], - "{} cp.async effects disagree with the imported LLVM declaration", - policy.id - ); - let cache = match copy.cache_policy { - CpAsyncCachePolicy::Ca => "ca", - CpAsyncCachePolicy::Cg => "cg", - }; - let mut operands = vec![ - OperandPattern::Address, - OperandPattern::Address, - OperandPattern::Exact { - value: copy.copy_size.bytes().to_string(), - }, - ]; - if dynamic_source_size { - operands.push(OperandPattern::RegisterOrImmediate); - } - ensure!( - policy.expected_ptx.mnemonic == "cp" - && policy.expected_ptx.modifiers == ["async", cache, "shared", "global"] - && policy.expected_ptx.operands == operands, - "{} expected PTX does not match its cp.async recipe", - policy.id - ); - let actual_selections: BTreeSet<_> = declaration - .selections - .iter() - .map(|selection| selection.source_record.as_str()) - .collect(); - ensure!( - actual_selections == recipe.selections.iter().copied().collect(), - "{} imported cp.async selection set changed", - policy.id - ); - let backend_pairs: BTreeSet<_> = policy - .backend_lowerings - .iter() - .map(|lowering| (lowering.backend, lowering.mechanism)) - .collect(); - ensure!( - policy.backend_lowerings.len() == 2 - && backend_pairs - == BTreeSet::from([ - ( - IntrinsicBackend::LlvmNvptx, - BackendLoweringMechanism::TypedNvvm, - ), - ( - IntrinsicBackend::LibNvvm, - BackendLoweringMechanism::InlinePtx, - ), - ]) - && policy.backend_lowerings.iter().all(|lowering| { - lowering.minimum_ptx.is_none() - && lowering.minimum_sm.is_none() - && !lowering.evidence_profile.trim().is_empty() - }), - "{} must define the reviewed typed-LLVM and inline-PTX cp.async routes", - policy.id - ); - ensure_no_other_family_contract(policy, "cp_async_copy")?; - ensure!( - policy.ldmatrix_variant.is_none() - && policy.ldmatrix_safety.is_none() - && policy.ldmatrix_adapter.is_none() - && policy.selected_address_space.is_none(), - "{} mixes ldmatrix state with cp_async_copy", - policy.id - ); - Ok(()) -} - -struct CpAsyncControlRecipe { - id: &'static str, - abi_id: &'static str, - operation_key: &'static str, - rust_name: &'static str, - dialect_op_type: &'static str, - dialect_op_name: &'static str, - source_record: &'static str, - llvm_symbol: &'static str, - selection: &'static str, - ptx_modifier: &'static str, - summary: &'static str, -} - -fn cp_async_control_recipe(operation: CpAsyncControlOperation) -> CpAsyncControlRecipe { - match operation { - CpAsyncControlOperation::CommitGroup => CpAsyncControlRecipe { - id: "cp_async_commit_group", - abi_id: "i0094", - operation_key: "memory.copy.async.group.commit", - rust_name: "cp_async_commit_group", - dialect_op_type: "CpAsyncCommitGroupOp", - dialect_op_name: "nvvm.cp_async_commit_group", - source_record: "int_nvvm_cp_async_commit_group", - llvm_symbol: "llvm.nvvm.cp.async.commit.group", - selection: "CP_ASYNC_COMMIT_GROUP", - ptx_modifier: "commit_group", - summary: "Commits this thread's uncommitted asynchronous copies as one group.", - }, - CpAsyncControlOperation::WaitAll => CpAsyncControlRecipe { - id: "cp_async_wait_all", - abi_id: "i0095", - operation_key: "memory.copy.async.group.wait_all", - rust_name: "cp_async_wait_all", - dialect_op_type: "CpAsyncWaitAllOp", - dialect_op_name: "nvvm.cp_async_wait_all", - source_record: "int_nvvm_cp_async_wait_all", - llvm_symbol: "llvm.nvvm.cp.async.wait.all", - selection: "CP_ASYNC_WAIT_ALL", - ptx_modifier: "wait_all", - summary: "Waits for all asynchronous copy groups issued by this thread.", - }, - CpAsyncControlOperation::WaitGroup => CpAsyncControlRecipe { - id: "cp_async_wait_group", - abi_id: "i0096", - operation_key: "memory.copy.async.group.wait_max_pending", - rust_name: "cp_async_wait_group", - dialect_op_type: "CpAsyncWaitGroupOp", - dialect_op_name: "nvvm.cp_async_wait_group", - source_record: "int_nvvm_cp_async_wait_group", - llvm_symbol: "llvm.nvvm.cp.async.wait.group", - selection: "CP_ASYNC_WAIT_GROUP", - ptx_modifier: "wait_group", - summary: "Waits until at most the compile-time constant number of recent asynchronous copy groups remain pending.", - }, - } -} - -fn validate_cp_async_control_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, -) -> Result<()> { - let control = policy - .cp_async_control - .as_ref() - .with_context(|| format!("{} has no closed cp.async control contract", policy.id))?; - let recipe = cp_async_control_recipe(control.operation); - let has_operand = control.operation == CpAsyncControlOperation::WaitGroup; - let expected_adapter = if has_operand { - CpAsyncControlAdapter::CompileTimeConstantMaxPending - } else { - CpAsyncControlAdapter::NoOperands - }; - ensure!( - control.adapter == expected_adapter, - "{} cp.async control and adapter disagree", - policy.id - ); - ensure!( - control.runtime_validation == RuntimeValidation::Unexecuted, - "{} cannot claim unrecorded cp.async control runtime validation", - policy.id - ); - ensure!( - policy.id == recipe.id - && policy.abi_id == recipe.abi_id - && policy.operation_key == recipe.operation_key - && policy.source.is_none() - && policy.source_record.as_deref() == Some(recipe.source_record) - && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) - && policy.resolved_llvm_symbol.is_none(), - "{} cp.async control identity does not match its closed recipe", - policy.id - ); - let rust_arguments = if has_operand { vec!["u32"] } else { vec![] }; - let dialect_operands = if has_operand { vec!["i32"] } else { vec![] }; - ensure!( - policy.rust_module == "async_copy" - && policy.rust_name == recipe.rust_name - && policy.rust_arguments == rust_arguments - && policy.rust_result == "()" - && !policy.safe - && !policy.must_use - && policy.safe_allowlist_reason.is_none() - && policy.public_rust_path - == format!("cuda_intrinsics::async_copy::{}", recipe.rust_name) - && policy.compatibility_rust_paths - == [format!("cuda_device::async_copy::{}", recipe.rust_name)] - && policy.dialect_op_type == recipe.dialect_op_type - && policy.dialect_op_name == recipe.dialect_op_name - && policy.dialect_operands == dialect_operands - && policy.dialect_results.is_empty() - && policy.llvm_arguments == dialect_operands - && policy.llvm_results.is_empty() - && policy.lowering == "generated_cp_async_control", - "{} is outside the closed cp.async control API and lowering recipe", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == "read_write" - && !policy.convergent - && policy.execution_scope == "thread" - && policy.minimum_ptx == "7.0" - && policy.minimum_sm.as_deref() == Some("sm_80") - && policy.ptx_result == "()" - && policy.targets == "all", - "{} cp.async control effects or target floor disagree with the closed recipe", - policy.id - ); - let (ptx_isa_section, ptx_isa_url) = match control.operation { - CpAsyncControlOperation::CommitGroup => ( - "9.7.9.26.3.2 Data Movement and Conversion Instructions: cp.async.commit_group", - "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async-commit-group", - ), - CpAsyncControlOperation::WaitAll | CpAsyncControlOperation::WaitGroup => ( - "9.7.9.26.3.3 Data Movement and Conversion Instructions: cp.async.wait_group / cp.async.wait_all", - "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async-wait-group-cp-async-wait-all", - ), - }; - ensure!( - policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section == ptx_isa_section - && policy.ptx_isa_url == ptx_isa_url, - "{} cp.async control PTX provenance disagrees with the reviewed recipe", - policy.id - ); - ensure!( - policy.summary == recipe.summary, - "{} cp.async control summary does not match its closed recipe", - policy.id - ); - let expected_properties: Vec = if has_operand { - vec!["ImmArg".into()] - } else { - vec![] - }; - ensure!( - declaration.properties == expected_properties, - "{} cp.async control properties disagree with the imported declaration", - policy.id - ); - let operands = if has_operand { - vec![OperandPattern::Immediate] - } else { - vec![] - }; - ensure!( - policy.expected_ptx.mnemonic == "cp" - && policy.expected_ptx.modifiers == ["async", recipe.ptx_modifier] - && policy.expected_ptx.operands == operands - && declaration.selections.len() == 1 - && declaration.selections[0].source_record == recipe.selection, - "{} expected PTX or imported selection disagrees with its cp.async control recipe", - policy.id - ); - let backend_pairs: BTreeSet<_> = policy - .backend_lowerings - .iter() - .map(|lowering| (lowering.backend, lowering.mechanism)) - .collect(); - ensure!( - policy.backend_lowerings.len() == 2 - && backend_pairs - == BTreeSet::from([ - ( - IntrinsicBackend::LlvmNvptx, - BackendLoweringMechanism::TypedNvvm, - ), - ( - IntrinsicBackend::LibNvvm, - BackendLoweringMechanism::InlinePtx, - ), - ]) - && policy.backend_lowerings.iter().all(|lowering| { - lowering.minimum_ptx.is_none() - && lowering.minimum_sm.is_none() - && !lowering.evidence_profile.trim().is_empty() - }), - "{} must define the reviewed typed-LLVM and inline-PTX cp.async control routes", - policy.id - ); - ensure_no_other_family_contract(policy, "cp_async_control")?; - ensure!( - policy.ldmatrix_variant.is_none() - && policy.ldmatrix_safety.is_none() - && policy.ldmatrix_adapter.is_none() - && policy.selected_address_space.is_none(), - "{} mixes ldmatrix state with cp_async_control", - policy.id - ); - Ok(()) -} - -struct CpAsyncMbarrierRecipe { - id: &'static str, - abi_id: &'static str, - operation_key: &'static str, - dialect_op_type: &'static str, - dialect_op_name: &'static str, - source_record: &'static str, - llvm_symbol: &'static str, - llvm_argument: &'static str, - selection: &'static str, - selection_asm: &'static str, - modifiers: &'static [&'static str], - summary: &'static str, -} - -fn cp_async_mbarrier_recipe( - operation: CpAsyncMbarrierOperation, - state_space: CpAsyncMbarrierStateSpace, -) -> CpAsyncMbarrierRecipe { - match (operation, state_space) { - (CpAsyncMbarrierOperation::Arrive, CpAsyncMbarrierStateSpace::Generic) => { - CpAsyncMbarrierRecipe { - id: "cp_async_mbarrier_arrive", - abi_id: "i0101", - operation_key: "memory.copy.async.mbarrier.arrive.generic.increment", - dialect_op_type: "CpAsyncMbarrierArriveOp", - dialect_op_name: "nvvm.cp_async_mbarrier_arrive", - source_record: "int_nvvm_cp_async_mbarrier_arrive", - llvm_symbol: "llvm.nvvm.cp.async.mbarrier.arrive", - llvm_argument: "ptr", - selection: "CP_ASYNC_MBARRIER_ARRIVE", - selection_asm: "cp.async.mbarrier.arrive.b64 \t[$addr];", - modifiers: &["async", "mbarrier", "arrive", "b64"], - summary: "Associates this thread's prior asynchronous copies with a shared-memory barrier using balanced pending-count accounting.", - } - } - (CpAsyncMbarrierOperation::ArriveNoInc, CpAsyncMbarrierStateSpace::Generic) => { - CpAsyncMbarrierRecipe { - id: "cp_async_mbarrier_arrive_noinc", - abi_id: "i0103", - operation_key: "memory.copy.async.mbarrier.arrive.generic.no_increment", - dialect_op_type: "CpAsyncMbarrierArriveNoIncOp", - dialect_op_name: "nvvm.cp_async_mbarrier_arrive_noinc", - source_record: "int_nvvm_cp_async_mbarrier_arrive_noinc", - llvm_symbol: "llvm.nvvm.cp.async.mbarrier.arrive.noinc", - llvm_argument: "ptr", - selection: "CP_ASYNC_MBARRIER_ARRIVE_NOINC", - selection_asm: "cp.async.mbarrier.arrive.noinc.b64 \t[$addr];", - modifiers: &["async", "mbarrier", "arrive", "noinc", "b64"], - summary: "Associates this thread's prior asynchronous copies with a shared-memory barrier without incrementing its pending count.", - } - } - (CpAsyncMbarrierOperation::ArriveNoInc, CpAsyncMbarrierStateSpace::Shared) => { - CpAsyncMbarrierRecipe { - id: "cp_async_mbarrier_arrive_noinc_shared", - abi_id: "i0104", - operation_key: "memory.copy.async.mbarrier.arrive.shared.no_increment", - dialect_op_type: "CpAsyncMbarrierArriveNoIncSharedOp", - dialect_op_name: "nvvm.cp_async_mbarrier_arrive_noinc_shared", - source_record: "int_nvvm_cp_async_mbarrier_arrive_noinc_shared", - llvm_symbol: "llvm.nvvm.cp.async.mbarrier.arrive.noinc.shared", - llvm_argument: "shared_ptr", - selection: "CP_ASYNC_MBARRIER_ARRIVE_NOINC_SHARED", - selection_asm: "cp.async.mbarrier.arrive.noinc.shared.b64 \t[$addr];", - modifiers: &["async", "mbarrier", "arrive", "noinc", "shared", "b64"], - summary: "Associates this thread's prior asynchronous copies with a shared-address barrier without incrementing its pending count.", - } - } - (CpAsyncMbarrierOperation::Arrive, CpAsyncMbarrierStateSpace::Shared) => { - CpAsyncMbarrierRecipe { - id: "cp_async_mbarrier_arrive_shared", - abi_id: "i0102", - operation_key: "memory.copy.async.mbarrier.arrive.shared.increment", - dialect_op_type: "CpAsyncMbarrierArriveSharedOp", - dialect_op_name: "nvvm.cp_async_mbarrier_arrive_shared", - source_record: "int_nvvm_cp_async_mbarrier_arrive_shared", - llvm_symbol: "llvm.nvvm.cp.async.mbarrier.arrive.shared", - llvm_argument: "shared_ptr", - selection: "CP_ASYNC_MBARRIER_ARRIVE_SHARED", - selection_asm: "cp.async.mbarrier.arrive.shared.b64 \t[$addr];", - modifiers: &["async", "mbarrier", "arrive", "shared", "b64"], - summary: "Associates this thread's prior asynchronous copies with a shared-address barrier using balanced pending-count accounting.", - } - } - } -} - -fn validate_cp_async_mbarrier_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, -) -> Result<()> { - let bridge = policy - .cp_async_mbarrier - .as_ref() - .with_context(|| format!("{} has no closed cp.async mbarrier contract", policy.id))?; - let recipe = cp_async_mbarrier_recipe(bridge.operation, bridge.state_space); - ensure!( - bridge.adapter == CpAsyncMbarrierAdapter::PointerToVoid, - "{} cp.async mbarrier adapter changed", - policy.id - ); - ensure!( - bridge.runtime_validation == RuntimeValidation::Unexecuted, - "{} cannot claim unrecorded cp.async mbarrier runtime validation", - policy.id - ); - ensure!( - policy.id == recipe.id - && policy.abi_id == recipe.abi_id - && policy.operation_key == recipe.operation_key - && policy.source.is_none() - && policy.source_record.as_deref() == Some(recipe.source_record) - && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) - && policy.resolved_llvm_symbol.is_none(), - "{} cp.async mbarrier identity does not match its closed recipe", - policy.id - ); - ensure!( - policy.rust_module == "async_copy" - && policy.rust_name == recipe.id - && policy.rust_arguments == ["*mut u64"] - && policy.rust_result == "()" - && !policy.safe - && !policy.must_use - && policy.safe_allowlist_reason.is_none() - && policy.public_rust_path == format!("cuda_intrinsics::async_copy::{}", recipe.id) - && policy.compatibility_rust_paths - == [format!("cuda_device::async_copy::{}", recipe.id)], - "{} is outside the closed cp.async mbarrier Rust API", - policy.id - ); - ensure!( - policy.dialect_op_type == recipe.dialect_op_type - && policy.dialect_op_name == recipe.dialect_op_name - && policy.dialect_operands == ["ptr"] - && policy.dialect_results.is_empty() - && policy.llvm_arguments == [recipe.llvm_argument] - && policy.llvm_results.is_empty() - && policy.lowering == "generated_cp_async_mbarrier", - "{} is outside the closed cp.async mbarrier signature and lowering recipe", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == "read_write" - && policy.convergent - && policy.execution_scope == "thread" - && policy.minimum_ptx == "7.0" - && policy.minimum_sm.as_deref() == Some("sm_80") - && policy.ptx_result == "()" - && policy.targets == "all", - "{} cp.async mbarrier effects or target floor disagree with the closed recipe", - policy.id - ); - ensure!( - policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section - == "9.7.14.16.18 Parallel Synchronization and Communication Instructions: cp.async.mbarrier.arrive" - && policy.ptx_isa_url - == "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-cp-async-mbarrier-arrive", - "{} cp.async mbarrier PTX provenance disagrees with the reviewed recipe", - policy.id - ); - ensure!( - policy.summary == recipe.summary, - "{} cp.async mbarrier summary does not match its closed recipe", - policy.id - ); - ensure!( - declaration.properties == ["IntrConvergent", "IntrNoCallback"], - "{} cp.async mbarrier properties disagree with the imported declaration", - policy.id - ); - ensure!( - policy.expected_ptx.mnemonic == "cp" - && policy - .expected_ptx - .modifiers - .iter() - .map(String::as_str) - .eq(recipe.modifiers.iter().copied()) - && policy.expected_ptx.operands == [OperandPattern::Address], - "{} expected PTX does not match its cp.async mbarrier recipe", - policy.id - ); - ensure!( - declaration.selections.len() == 1 - && declaration.selections[0].source_record == recipe.selection - && declaration.selections[0].asm == recipe.selection_asm - && declaration.selections[0].predicates - == [ - "Subtarget->getSmVersion() >= 80", - "Subtarget->getPTXVersion() >= 70", - ] - && declaration.selections[0].constraints.is_empty(), - "{} imported cp.async mbarrier selection changed", - policy.id - ); - let backend_pairs: BTreeSet<_> = policy - .backend_lowerings - .iter() - .map(|lowering| (lowering.backend, lowering.mechanism)) - .collect(); - ensure!( - policy.backend_lowerings.len() == 2 - && backend_pairs - == BTreeSet::from([ - ( - IntrinsicBackend::LlvmNvptx, - BackendLoweringMechanism::TypedNvvm, - ), - ( - IntrinsicBackend::LibNvvm, - BackendLoweringMechanism::InlinePtx, - ), - ]) - && policy.backend_lowerings.iter().all(|lowering| { - lowering.minimum_ptx.is_none() - && lowering.minimum_sm.is_none() - && !lowering.evidence_profile.trim().is_empty() - }), - "{} must define the reviewed typed-LLVM and inline-PTX cp.async mbarrier routes", - policy.id - ); - ensure_no_other_family_contract(policy, "cp_async_mbarrier")?; - Ok(()) -} - -struct MbarrierBasicRecipe { - id: &'static str, - abi_id: &'static str, - operation_key: &'static str, - rust_arguments: &'static [&'static str], - rust_result: &'static str, - must_use: bool, - adapter: MbarrierBasicAdapter, - dialect_op_type: &'static str, - dialect_op_name: &'static str, - dialect_operands: &'static [&'static str], - dialect_results: &'static [&'static str], - source_record: &'static str, - llvm_symbol: &'static str, - llvm_arguments: &'static [&'static str], - llvm_results: &'static [&'static str], - memory: &'static str, - ptx_result: &'static str, - selection: &'static str, - selection_asm: &'static str, - ptx_modifier: &'static str, - ptx_isa_section: &'static str, - ptx_isa_url: &'static str, - llvm_nvptx_mechanism: BackendLoweringMechanism, - lib_nvvm_mechanism: BackendLoweringMechanism, - properties: &'static [&'static str], - summary: &'static str, -} - -fn mbarrier_basic_recipe(operation: MbarrierBasicOperation) -> MbarrierBasicRecipe { - match operation { - MbarrierBasicOperation::Init => MbarrierBasicRecipe { - id: "mbarrier_init", - abi_id: "i0097", - operation_key: "barrier.mbarrier.init.shared.cta", - rust_arguments: &["*mut u64", "u32"], - rust_result: "()", - must_use: false, - adapter: MbarrierBasicAdapter::InitPointerCountToVoid, - dialect_op_type: "MbarrierInitSharedOp", - dialect_op_name: "nvvm.mbarrier_init_shared", - dialect_operands: &["ptr", "i32"], - dialect_results: &[], - source_record: "int_nvvm_mbarrier_init_shared", - llvm_symbol: "llvm.nvvm.mbarrier.init.shared", - llvm_arguments: &["shared_ptr", "i32"], - llvm_results: &[], - memory: "read_write", - ptx_result: "()", - selection: "MBARRIER_INIT_SHARED", - selection_asm: "mbarrier.init.shared.b64 \t[$addr], $count;", - ptx_modifier: "init", - ptx_isa_section: "9.7.14.16.12 Parallel Synchronization and Communication Instructions: mbarrier.init", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-mbarrier-init", - llvm_nvptx_mechanism: BackendLoweringMechanism::TypedNvvm, - lib_nvvm_mechanism: BackendLoweringMechanism::InlinePtx, - properties: &["IntrConvergent", "IntrNoCallback"], - summary: "Initializes a CTA shared-memory barrier with the expected arrival count.", - }, - MbarrierBasicOperation::Arrive => MbarrierBasicRecipe { - id: "mbarrier_arrive", - abi_id: "i0098", - operation_key: "barrier.mbarrier.arrive.shared.cta", - rust_arguments: &["*const u64"], - rust_result: "u64", - must_use: true, - adapter: MbarrierBasicAdapter::ArrivePointerToToken, - dialect_op_type: "MbarrierArriveSharedOp", - dialect_op_name: "nvvm.mbarrier_arrive_shared", - dialect_operands: &["ptr"], - dialect_results: &["i64"], - source_record: "int_nvvm_mbarrier_arrive_shared", - llvm_symbol: "llvm.nvvm.mbarrier.arrive.shared", - llvm_arguments: &["shared_ptr"], - llvm_results: &["i64"], - memory: "read_write", - ptx_result: "u64", - selection: "MBARRIER_ARRIVE_SHARED", - selection_asm: "mbarrier.arrive.shared.b64 \t$state, [$addr];", - ptx_modifier: "arrive", - ptx_isa_section: "9.7.14.16.16 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", - llvm_nvptx_mechanism: BackendLoweringMechanism::TypedNvvm, - lib_nvvm_mechanism: BackendLoweringMechanism::InlinePtx, - properties: &["IntrConvergent", "IntrNoCallback"], - summary: "Arrives at a CTA shared-memory barrier and returns its phase token.", - }, - MbarrierBasicOperation::TestWait => MbarrierBasicRecipe { - id: "mbarrier_test_wait", - abi_id: "i0099", - operation_key: "barrier.mbarrier.test_wait.shared.cta", - rust_arguments: &["*const u64", "u64"], - rust_result: "bool", - must_use: true, - adapter: MbarrierBasicAdapter::TestWaitPointerTokenToPredicate, - dialect_op_type: "MbarrierTestWaitSharedOp", - dialect_op_name: "nvvm.mbarrier_test_wait_shared", - dialect_operands: &["ptr", "i64"], - dialect_results: &["i1"], - source_record: "int_nvvm_mbarrier_test_wait_shared", - llvm_symbol: "llvm.nvvm.mbarrier.test.wait.shared", - llvm_arguments: &["shared_ptr", "i64"], - llvm_results: &["i1"], - memory: "read_write", - ptx_result: "bool", - selection: "MBARRIER_TEST_WAIT_SHARED", - selection_asm: "mbarrier.test_wait.shared.b64 \t$res, [$addr], $state;", - ptx_modifier: "test_wait", - ptx_isa_section: "9.7.14.16.19 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", - llvm_nvptx_mechanism: BackendLoweringMechanism::InlinePtx, - lib_nvvm_mechanism: BackendLoweringMechanism::InlinePtx, - properties: &["IntrConvergent", "IntrNoCallback"], - summary: "Tests whether the CTA shared-memory barrier phase for a token is complete.", - }, - MbarrierBasicOperation::Inval => MbarrierBasicRecipe { - id: "mbarrier_inval", - abi_id: "i0100", - operation_key: "barrier.mbarrier.inval.shared.cta", - rust_arguments: &["*mut u64"], - rust_result: "()", - must_use: false, - adapter: MbarrierBasicAdapter::InvalPointerToVoid, - dialect_op_type: "MbarrierInvalSharedOp", - dialect_op_name: "nvvm.mbarrier_inval_shared", - dialect_operands: &["ptr"], - dialect_results: &[], - source_record: "int_nvvm_mbarrier_inval_shared", - llvm_symbol: "llvm.nvvm.mbarrier.inval.shared", - llvm_arguments: &["shared_ptr"], - llvm_results: &[], - memory: "write", - ptx_result: "()", - selection: "MBARRIER_INVAL_SHARED", - selection_asm: "mbarrier.inval.shared.b64 \t[$addr];", - ptx_modifier: "inval", - ptx_isa_section: "9.7.14.16.13 Parallel Synchronization and Communication Instructions: mbarrier.inval", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-mbarrier-inval", - llvm_nvptx_mechanism: BackendLoweringMechanism::TypedNvvm, - lib_nvvm_mechanism: BackendLoweringMechanism::InlinePtx, - properties: &[ - "IntrArgMemOnly", - "IntrConvergent", - "IntrNoCallback", - "IntrWriteMem", - "NoCapture", - "WriteOnly", - ], - summary: "Invalidates a CTA shared-memory barrier after its users have finished.", - }, - } -} - -fn mbarrier_expected_operands(operation: MbarrierBasicOperation) -> Vec { - match operation { - MbarrierBasicOperation::Init => { - vec![OperandPattern::Address, OperandPattern::RegisterOrImmediate] - } - MbarrierBasicOperation::Arrive => { - vec![OperandPattern::Register, OperandPattern::Address] - } - MbarrierBasicOperation::TestWait => vec![ - OperandPattern::Register, - OperandPattern::Address, - OperandPattern::Register, - ], - MbarrierBasicOperation::Inval => vec![OperandPattern::Address], - } -} - -fn validate_mbarrier_basic_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, -) -> Result<()> { - let mbarrier = policy - .mbarrier_basic - .as_ref() - .with_context(|| format!("{} has no closed basic mbarrier contract", policy.id))?; - let recipe = mbarrier_basic_recipe(mbarrier.operation); - ensure!( - mbarrier.state_space == MbarrierStateSpace::Shared && mbarrier.adapter == recipe.adapter, - "{} mbarrier operation, state space, and adapter disagree", - policy.id - ); - ensure!( - mbarrier.runtime_validation == RuntimeValidation::Unexecuted, - "{} cannot claim unrecorded mbarrier runtime validation", - policy.id - ); - ensure!( - policy.id == recipe.id - && policy.abi_id == recipe.abi_id - && policy.operation_key == recipe.operation_key - && policy.source.is_none() - && policy.source_record.as_deref() == Some(recipe.source_record) - && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) - && policy.resolved_llvm_symbol.is_none(), - "{} mbarrier identity does not match its closed recipe", - policy.id - ); - ensure!( - policy.rust_module == "barrier" - && policy.rust_name == recipe.id - && policy - .rust_arguments - .iter() - .map(String::as_str) - .eq(recipe.rust_arguments.iter().copied()) - && policy.rust_result == recipe.rust_result - && !policy.safe - && policy.must_use == recipe.must_use - && policy.safe_allowlist_reason.is_none() - && policy.public_rust_path == format!("cuda_intrinsics::barrier::{}", recipe.id) - && policy.compatibility_rust_paths == [format!("cuda_device::barrier::{}", recipe.id)], - "{} must preserve its unsafe mbarrier raw and compatibility API", - policy.id - ); - ensure!( - policy.dialect_op_type == recipe.dialect_op_type - && policy.dialect_op_name == recipe.dialect_op_name - && policy - .dialect_operands - .iter() - .map(String::as_str) - .eq(recipe.dialect_operands.iter().copied()) - && policy - .dialect_results - .iter() - .map(String::as_str) - .eq(recipe.dialect_results.iter().copied()) - && policy - .llvm_arguments - .iter() - .map(String::as_str) - .eq(recipe.llvm_arguments.iter().copied()) - && policy - .llvm_results - .iter() - .map(String::as_str) - .eq(recipe.llvm_results.iter().copied()) - && policy.lowering == "generated_mbarrier_basic", - "{} is outside the closed mbarrier signature and lowering recipe", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == recipe.memory - && policy.convergent - && policy.execution_scope == "cta" - && policy.minimum_ptx == "7.0" - && policy.minimum_sm.as_deref() == Some("sm_80") - && policy.ptx_result == recipe.ptx_result - && policy.targets == "all", - "{} mbarrier effects or target floor disagree with the closed recipe", - policy.id - ); - ensure!( - policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section == recipe.ptx_isa_section - && policy.ptx_isa_url == recipe.ptx_isa_url, - "{} mbarrier PTX provenance disagrees with the reviewed recipe", - policy.id - ); - ensure!( - policy.summary == recipe.summary, - "{} mbarrier summary does not match its closed recipe", - policy.id - ); - ensure!( - declaration - .properties - .iter() - .map(String::as_str) - .eq(recipe.properties.iter().copied()), - "{} mbarrier properties disagree with the imported LLVM declaration", - policy.id - ); - ensure!( - policy.expected_ptx.mnemonic == "mbarrier" - && policy.expected_ptx.modifiers == [recipe.ptx_modifier, "shared", "b64"] - && policy.expected_ptx.operands == mbarrier_expected_operands(mbarrier.operation), - "{} expected PTX does not match its mbarrier recipe", - policy.id - ); - ensure!( - declaration.selections.len() == 1 - && declaration.selections[0].source_record == recipe.selection - && declaration.selections[0].asm == recipe.selection_asm - && declaration.selections[0].predicates - == [ - "Subtarget->getSmVersion() >= 80", - "Subtarget->getPTXVersion() >= 70", - ] - && declaration.selections[0].constraints.is_empty(), - "{} imported mbarrier selection changed", - policy.id - ); - let backend_pairs: BTreeSet<_> = policy - .backend_lowerings - .iter() - .map(|lowering| (lowering.backend, lowering.mechanism)) - .collect(); - ensure!( - policy.backend_lowerings.len() == 2 - && backend_pairs - == BTreeSet::from([ - (IntrinsicBackend::LlvmNvptx, recipe.llvm_nvptx_mechanism), - (IntrinsicBackend::LibNvvm, recipe.lib_nvvm_mechanism), - ]) - && policy.backend_lowerings.iter().all(|lowering| { - lowering.minimum_ptx.is_none() - && lowering.minimum_sm.is_none() - && !lowering.evidence_profile.trim().is_empty() - }), - "{} must define exactly the reviewed mbarrier backend routes", - policy.id - ); - ensure_no_other_family_contract(policy, "mbarrier_basic")?; - Ok(()) -} - -fn packed_conversion_recipe( - conversion: &crate::model::PackedConversion, -) -> Option { - match conversion.source_format { - PackedConversionSourceFormat::F32x2 => packed_conversion_recipe_f32x2(conversion), - PackedConversionSourceFormat::E4m3x2 - | PackedConversionSourceFormat::E5m2x2 - | PackedConversionSourceFormat::F16x2 => packed_conversion_recipe_fp8_f16x2(conversion), - } -} - -/// Recipes for the packed FP8 conversions whose other side is `f16x2`. -/// -/// Keyed on source as well as destination: unpacking to `f16x2` shares its -/// destination, rounding, and saturation with the scalar-`f32` recipe for -/// `cvt.rn.f16x2.f32`, so the source format is what separates them. -fn packed_conversion_recipe_fp8_f16x2( - conversion: &crate::model::PackedConversion, -) -> Option { - match ( - conversion.source_format, - conversion.destination_format, - conversion.rounding, - conversion.saturation, - ) { - ( - PackedConversionSourceFormat::F16x2, - PackedConversionDestinationFormat::E4m3x2, - PackedConversionRounding::NearestEven, - PackedConversionSaturation::Satfinite, - ) => Some(PackedConversionRecipe { - id: "cvt_rn_satfinite_e4m3x2_f16x2", - abi_id: "i0822", - operation_key: "packed.convert.f16x2.e4m3x2.nearest_even.satfinite", - rust_name: "cvt_rn_satfinite_e4m3x2_f16x2", - compatibility_path: "cuda_device::convert::cvt_rn_satfinite_e4m3x2_f16x2", - dialect_op_type: "CvtRnSatfiniteE4m3x2F16x2Op", - dialect_op_name: "nvvm.cvt_rn_satfinite_e4m3x2_f16x2", - source_record: "int_nvvm_f16x2_to_e4m3x2_rn", - llvm_symbol: "llvm.nvvm.f16x2.to.e4m3x2.rn", - llvm_result: "i16", - summary: "Converts packed f16x2 to packed e4m3x2 with nearest-even finite saturation, preserving half order.", - }), - ( - PackedConversionSourceFormat::F16x2, - PackedConversionDestinationFormat::E4m3x2, - PackedConversionRounding::NearestEven, - PackedConversionSaturation::SatfiniteRelu, - ) => Some(PackedConversionRecipe { - id: "cvt_rn_satfinite_relu_e4m3x2_f16x2", - abi_id: "i0823", - operation_key: "packed.convert.f16x2.e4m3x2.nearest_even.satfinite.relu", - rust_name: "cvt_rn_satfinite_relu_e4m3x2_f16x2", - compatibility_path: "cuda_device::convert::cvt_rn_satfinite_relu_e4m3x2_f16x2", - dialect_op_type: "CvtRnSatfiniteReluE4m3x2F16x2Op", - dialect_op_name: "nvvm.cvt_rn_satfinite_relu_e4m3x2_f16x2", - source_record: "int_nvvm_f16x2_to_e4m3x2_rn_relu", - llvm_symbol: "llvm.nvvm.f16x2.to.e4m3x2.rn.relu", - llvm_result: "i16", - summary: "Converts packed f16x2 to packed e4m3x2 with nearest-even finite saturation and ReLU, preserving half order.", - }), - ( - PackedConversionSourceFormat::F16x2, - PackedConversionDestinationFormat::E5m2x2, - PackedConversionRounding::NearestEven, - PackedConversionSaturation::Satfinite, - ) => Some(PackedConversionRecipe { - id: "cvt_rn_satfinite_e5m2x2_f16x2", - abi_id: "i0824", - operation_key: "packed.convert.f16x2.e5m2x2.nearest_even.satfinite", - rust_name: "cvt_rn_satfinite_e5m2x2_f16x2", - compatibility_path: "cuda_device::convert::cvt_rn_satfinite_e5m2x2_f16x2", - dialect_op_type: "CvtRnSatfiniteE5m2x2F16x2Op", - dialect_op_name: "nvvm.cvt_rn_satfinite_e5m2x2_f16x2", - source_record: "int_nvvm_f16x2_to_e5m2x2_rn", - llvm_symbol: "llvm.nvvm.f16x2.to.e5m2x2.rn", - llvm_result: "i16", - summary: "Converts packed f16x2 to packed e5m2x2 with nearest-even finite saturation, preserving half order.", - }), - ( - PackedConversionSourceFormat::F16x2, - PackedConversionDestinationFormat::E5m2x2, - PackedConversionRounding::NearestEven, - PackedConversionSaturation::SatfiniteRelu, - ) => Some(PackedConversionRecipe { - id: "cvt_rn_satfinite_relu_e5m2x2_f16x2", - abi_id: "i0825", - operation_key: "packed.convert.f16x2.e5m2x2.nearest_even.satfinite.relu", - rust_name: "cvt_rn_satfinite_relu_e5m2x2_f16x2", - compatibility_path: "cuda_device::convert::cvt_rn_satfinite_relu_e5m2x2_f16x2", - dialect_op_type: "CvtRnSatfiniteReluE5m2x2F16x2Op", - dialect_op_name: "nvvm.cvt_rn_satfinite_relu_e5m2x2_f16x2", - source_record: "int_nvvm_f16x2_to_e5m2x2_rn_relu", - llvm_symbol: "llvm.nvvm.f16x2.to.e5m2x2.rn.relu", - llvm_result: "i16", - summary: "Converts packed f16x2 to packed e5m2x2 with nearest-even finite saturation and ReLU, preserving half order.", - }), - ( - PackedConversionSourceFormat::E4m3x2, - PackedConversionDestinationFormat::F16x2, - PackedConversionRounding::NearestEven, - PackedConversionSaturation::None, - ) => Some(PackedConversionRecipe { - id: "cvt_rn_f16x2_e4m3x2", - abi_id: "i0826", - operation_key: "packed.convert.e4m3x2.f16x2.nearest_even", - rust_name: "cvt_rn_f16x2_e4m3x2", - compatibility_path: "cuda_device::convert::cvt_rn_f16x2_e4m3x2", - dialect_op_type: "CvtRnF16x2E4m3x2Op", - dialect_op_name: "nvvm.cvt_rn_f16x2_e4m3x2", - source_record: "int_nvvm_e4m3x2_to_f16x2_rn", - llvm_symbol: "llvm.nvvm.e4m3x2.to.f16x2.rn", - llvm_result: "v2f16", - summary: "Converts packed e4m3x2 to packed f16x2, preserving byte order.", - }), - ( - PackedConversionSourceFormat::E4m3x2, - PackedConversionDestinationFormat::F16x2, - PackedConversionRounding::NearestEven, - PackedConversionSaturation::Relu, - ) => Some(PackedConversionRecipe { - id: "cvt_rn_relu_f16x2_e4m3x2", - abi_id: "i0827", - operation_key: "packed.convert.e4m3x2.f16x2.nearest_even.relu", - rust_name: "cvt_rn_relu_f16x2_e4m3x2", - compatibility_path: "cuda_device::convert::cvt_rn_relu_f16x2_e4m3x2", - dialect_op_type: "CvtRnReluF16x2E4m3x2Op", - dialect_op_name: "nvvm.cvt_rn_relu_f16x2_e4m3x2", - source_record: "int_nvvm_e4m3x2_to_f16x2_rn_relu", - llvm_symbol: "llvm.nvvm.e4m3x2.to.f16x2.rn.relu", - llvm_result: "v2f16", - summary: "Converts packed e4m3x2 to packed f16x2 with ReLU, preserving byte order.", - }), - ( - PackedConversionSourceFormat::E5m2x2, - PackedConversionDestinationFormat::F16x2, - PackedConversionRounding::NearestEven, - PackedConversionSaturation::None, - ) => Some(PackedConversionRecipe { - id: "cvt_rn_f16x2_e5m2x2", - abi_id: "i0828", - operation_key: "packed.convert.e5m2x2.f16x2.nearest_even", - rust_name: "cvt_rn_f16x2_e5m2x2", - compatibility_path: "cuda_device::convert::cvt_rn_f16x2_e5m2x2", - dialect_op_type: "CvtRnF16x2E5m2x2Op", - dialect_op_name: "nvvm.cvt_rn_f16x2_e5m2x2", - source_record: "int_nvvm_e5m2x2_to_f16x2_rn", - llvm_symbol: "llvm.nvvm.e5m2x2.to.f16x2.rn", - llvm_result: "v2f16", - summary: "Converts packed e5m2x2 to packed f16x2, preserving byte order.", - }), - ( - PackedConversionSourceFormat::E5m2x2, - PackedConversionDestinationFormat::F16x2, - PackedConversionRounding::NearestEven, - PackedConversionSaturation::Relu, - ) => Some(PackedConversionRecipe { - id: "cvt_rn_relu_f16x2_e5m2x2", - abi_id: "i0829", - operation_key: "packed.convert.e5m2x2.f16x2.nearest_even.relu", - rust_name: "cvt_rn_relu_f16x2_e5m2x2", - compatibility_path: "cuda_device::convert::cvt_rn_relu_f16x2_e5m2x2", - dialect_op_type: "CvtRnReluF16x2E5m2x2Op", - dialect_op_name: "nvvm.cvt_rn_relu_f16x2_e5m2x2", - source_record: "int_nvvm_e5m2x2_to_f16x2_rn_relu", - llvm_symbol: "llvm.nvvm.e5m2x2.to.f16x2.rn.relu", - llvm_result: "v2f16", - summary: "Converts packed e5m2x2 to packed f16x2 with ReLU, preserving byte order.", - }), - _ => None, - } -} - -fn packed_conversion_recipe_f32x2( - conversion: &crate::model::PackedConversion, -) -> Option { - match ( - conversion.destination_format, - conversion.rounding, - conversion.saturation, - ) { - ( - PackedConversionDestinationFormat::Bf16x2, - PackedConversionRounding::NearestEven, - PackedConversionSaturation::None, - ) => Some(PackedConversionRecipe { - id: "cvt_f32x2_bf16x2", - abi_id: "i0071", - operation_key: "packed.convert.f32x2.bf16x2.nearest_even", - rust_name: "cvt_f32x2_bf16x2", - compatibility_path: "cuda_device::convert::cvt_bf16x2_f32", - dialect_op_type: "CvtF32x2Bf16x2Op", - dialect_op_name: "nvvm.cvt_f32x2_bf16x2", - source_record: "int_nvvm_ff2bf16x2_rn", - llvm_symbol: "llvm.nvvm.ff2bf16x2.rn", - llvm_result: "v2bf16", - summary: "Converts two f32 values to packed bf16x2 with the first argument in the low half.", - }), - ( - PackedConversionDestinationFormat::F16x2, - PackedConversionRounding::NearestEven, - PackedConversionSaturation::None, - ) => Some(PackedConversionRecipe { - id: "cvt_f16x2_f32", - abi_id: "i0081", - operation_key: "packed.convert.f32x2.f16x2.nearest_even", - rust_name: "cvt_f16x2_f32", - compatibility_path: "cuda_device::convert::cvt_f16x2_f32", - dialect_op_type: "CvtF16x2F32Op", - dialect_op_name: "nvvm.cvt_f16x2_f32", - source_record: "int_nvvm_ff2f16x2_rn", - llvm_symbol: "llvm.nvvm.ff2f16x2.rn", - llvm_result: "v2f16", - summary: "Converts two f32 values to packed f16x2 with nearest-even rounding and the first argument in the low half.", - }), - ( - PackedConversionDestinationFormat::F16x2, - PackedConversionRounding::TowardZero, - PackedConversionSaturation::None, - ) => Some(PackedConversionRecipe { - id: "cvt_rz_f16x2_f32", - abi_id: "i0082", - operation_key: "packed.convert.f32x2.f16x2.toward_zero", - rust_name: "cvt_rz_f16x2_f32", - compatibility_path: "cuda_device::convert::cvt_rz_f16x2_f32", - dialect_op_type: "CvtRzF16x2F32Op", - dialect_op_name: "nvvm.cvt_rz_f16x2_f32", - source_record: "int_nvvm_ff2f16x2_rz", - llvm_symbol: "llvm.nvvm.ff2f16x2.rz", - llvm_result: "v2f16", - summary: "Converts two f32 values to packed f16x2 with toward-zero rounding and the first argument in the low half.", - }), - ( - PackedConversionDestinationFormat::F16x2, - PackedConversionRounding::NearestEven, - PackedConversionSaturation::Relu, - ) => Some(PackedConversionRecipe { - id: "cvt_rn_relu_f16x2_f32", - abi_id: "i0083", - operation_key: "packed.convert.f32x2.f16x2.nearest_even.relu", - rust_name: "cvt_rn_relu_f16x2_f32", - compatibility_path: "cuda_device::convert::cvt_rn_relu_f16x2_f32", - dialect_op_type: "CvtRnReluF16x2F32Op", - dialect_op_name: "nvvm.cvt_rn_relu_f16x2_f32", - source_record: "int_nvvm_ff2f16x2_rn_relu", - llvm_symbol: "llvm.nvvm.ff2f16x2.rn.relu", - llvm_result: "v2f16", - summary: "Converts two f32 values to packed f16x2 with nearest-even rounding, ReLU, and the first argument in the low half.", - }), - ( - PackedConversionDestinationFormat::Bf16x2, - PackedConversionRounding::NearestEven, - PackedConversionSaturation::Relu, - ) => Some(PackedConversionRecipe { - id: "cvt_rn_relu_bf16x2_f32", - abi_id: "i0084", - operation_key: "packed.convert.f32x2.bf16x2.nearest_even.relu", - rust_name: "cvt_rn_relu_bf16x2_f32", - compatibility_path: "cuda_device::convert::cvt_rn_relu_bf16x2_f32", - dialect_op_type: "CvtRnReluBf16x2F32Op", - dialect_op_name: "nvvm.cvt_rn_relu_bf16x2_f32", - source_record: "int_nvvm_ff2bf16x2_rn_relu", - llvm_symbol: "llvm.nvvm.ff2bf16x2.rn.relu", - llvm_result: "v2bf16", - summary: "Converts two f32 values to packed bf16x2 with nearest-even rounding, ReLU, and the first argument in the low half.", - }), - ( - PackedConversionDestinationFormat::Bf16x2, - PackedConversionRounding::TowardZero, - PackedConversionSaturation::None, - ) => Some(PackedConversionRecipe { - id: "cvt_rz_bf16x2_f32", - abi_id: "i0085", - operation_key: "packed.convert.f32x2.bf16x2.toward_zero", - rust_name: "cvt_rz_bf16x2_f32", - compatibility_path: "cuda_device::convert::cvt_rz_bf16x2_f32", - dialect_op_type: "CvtRzBf16x2F32Op", - dialect_op_name: "nvvm.cvt_rz_bf16x2_f32", - source_record: "int_nvvm_ff2bf16x2_rz", - llvm_symbol: "llvm.nvvm.ff2bf16x2.rz", - llvm_result: "v2bf16", - summary: "Converts two f32 values to packed bf16x2 with toward-zero rounding and the first argument in the low half.", - }), - ( - PackedConversionDestinationFormat::E4m3x2, - PackedConversionRounding::NearestEven, - PackedConversionSaturation::Satfinite, - ) => Some(PackedConversionRecipe { - id: "cvt_rn_satfinite_e4m3x2_f32", - abi_id: "i0259", - operation_key: "packed.convert.f32x2.e4m3x2.nearest_even.satfinite", - rust_name: "cvt_rn_satfinite_e4m3x2_f32", - compatibility_path: "cuda_device::convert::cvt_rn_satfinite_e4m3x2_f32", - dialect_op_type: "CvtRnSatfiniteE4m3x2F32Op", - dialect_op_name: "nvvm.cvt_rn_satfinite_e4m3x2_f32", - source_record: "int_nvvm_ff_to_e4m3x2_rn", - llvm_symbol: "llvm.nvvm.ff.to.e4m3x2.rn", - llvm_result: "i16", - summary: "Converts two f32 values to packed e4m3x2 with nearest-even finite saturation and the first argument in the low byte.", - }), - ( - PackedConversionDestinationFormat::E4m3x2, - PackedConversionRounding::NearestEven, - PackedConversionSaturation::SatfiniteRelu, - ) => Some(PackedConversionRecipe { - id: "cvt_rn_satfinite_relu_e4m3x2_f32", - abi_id: "i0260", - operation_key: "packed.convert.f32x2.e4m3x2.nearest_even.satfinite.relu", - rust_name: "cvt_rn_satfinite_relu_e4m3x2_f32", - compatibility_path: "cuda_device::convert::cvt_rn_satfinite_relu_e4m3x2_f32", - dialect_op_type: "CvtRnSatfiniteReluE4m3x2F32Op", - dialect_op_name: "nvvm.cvt_rn_satfinite_relu_e4m3x2_f32", - source_record: "int_nvvm_ff_to_e4m3x2_rn_relu", - llvm_symbol: "llvm.nvvm.ff.to.e4m3x2.rn.relu", - llvm_result: "i16", - summary: "Converts two f32 values to packed e4m3x2 with nearest-even finite saturation, ReLU, and the first argument in the low byte.", - }), - ( - PackedConversionDestinationFormat::E5m2x2, - PackedConversionRounding::NearestEven, - PackedConversionSaturation::Satfinite, - ) => Some(PackedConversionRecipe { - id: "cvt_rn_satfinite_e5m2x2_f32", - abi_id: "i0261", - operation_key: "packed.convert.f32x2.e5m2x2.nearest_even.satfinite", - rust_name: "cvt_rn_satfinite_e5m2x2_f32", - compatibility_path: "cuda_device::convert::cvt_rn_satfinite_e5m2x2_f32", - dialect_op_type: "CvtRnSatfiniteE5m2x2F32Op", - dialect_op_name: "nvvm.cvt_rn_satfinite_e5m2x2_f32", - source_record: "int_nvvm_ff_to_e5m2x2_rn", - llvm_symbol: "llvm.nvvm.ff.to.e5m2x2.rn", - llvm_result: "i16", - summary: "Converts two f32 values to packed e5m2x2 with nearest-even finite saturation and the first argument in the low byte.", - }), - ( - PackedConversionDestinationFormat::E5m2x2, - PackedConversionRounding::NearestEven, - PackedConversionSaturation::SatfiniteRelu, - ) => Some(PackedConversionRecipe { - id: "cvt_rn_satfinite_relu_e5m2x2_f32", - abi_id: "i0262", - operation_key: "packed.convert.f32x2.e5m2x2.nearest_even.satfinite.relu", - rust_name: "cvt_rn_satfinite_relu_e5m2x2_f32", - compatibility_path: "cuda_device::convert::cvt_rn_satfinite_relu_e5m2x2_f32", - dialect_op_type: "CvtRnSatfiniteReluE5m2x2F32Op", - dialect_op_name: "nvvm.cvt_rn_satfinite_relu_e5m2x2_f32", - source_record: "int_nvvm_ff_to_e5m2x2_rn_relu", - llvm_symbol: "llvm.nvvm.ff.to.e5m2x2.rn.relu", - llvm_result: "i16", - summary: "Converts two f32 values to packed e5m2x2 with nearest-even finite saturation, ReLU, and the first argument in the low byte.", - }), - _ => None, - } -} - -fn packed_conversion_ptx_modifiers( - conversion: &crate::model::PackedConversion, -) -> Vec<&'static str> { - let rounding = match conversion.rounding { - PackedConversionRounding::NearestEven => "rn", - PackedConversionRounding::TowardZero => "rz", - }; - let format = match conversion.destination_format { - PackedConversionDestinationFormat::Bf16x2 => "bf16x2", - PackedConversionDestinationFormat::E4m3x2 => "e4m3x2", - PackedConversionDestinationFormat::E5m2x2 => "e5m2x2", - PackedConversionDestinationFormat::F16x2 => "f16x2", - }; - let mut modifiers = vec![rounding]; - match conversion.saturation { - PackedConversionSaturation::None => {} - PackedConversionSaturation::Relu => modifiers.push("relu"), - PackedConversionSaturation::Satfinite => modifiers.push("satfinite"), - PackedConversionSaturation::SatfiniteRelu => modifiers.extend(["satfinite", "relu"]), - } - modifiers.extend([format, conversion.source_format.ptx_token()]); - modifiers -} - -/// Whether either side of the conversion is a packed FP8 format. -fn packed_conversion_uses_fp8(conversion: &crate::model::PackedConversion) -> bool { - matches!( - conversion.source_format, - PackedConversionSourceFormat::E4m3x2 | PackedConversionSourceFormat::E5m2x2 - ) || matches!( - conversion.destination_format, - PackedConversionDestinationFormat::E4m3x2 | PackedConversionDestinationFormat::E5m2x2 - ) -} - -/// Source operand types, in Rust, dialect, and LLVM spellings. -/// -/// A packed source arrives in one register; `f32x2` names two scalar operands. -fn packed_conversion_source_types( - conversion: &crate::model::PackedConversion, -) -> (Vec, Vec, Vec) { - match conversion.source_format { - PackedConversionSourceFormat::F32x2 => ( - vec!["f32".into(), "f32".into()], - vec!["f32".into(), "f32".into()], - vec!["f32".into(), "f32".into()], - ), - PackedConversionSourceFormat::F16x2 => { - (vec!["u32".into()], vec!["i32".into()], vec!["v2f16".into()]) - } - PackedConversionSourceFormat::E4m3x2 | PackedConversionSourceFormat::E5m2x2 => { - (vec!["u16".into()], vec!["i16".into()], vec!["i16".into()]) - } - } -} - -fn packed_conversion_result_width(conversion: &crate::model::PackedConversion) -> u32 { - match conversion.destination_format { - PackedConversionDestinationFormat::Bf16x2 | PackedConversionDestinationFormat::F16x2 => 32, - PackedConversionDestinationFormat::E4m3x2 | PackedConversionDestinationFormat::E5m2x2 => 16, - } -} - -fn packed_conversion_floor( - conversion: &crate::model::PackedConversion, -) -> (&'static str, &'static str) { - // FP8 on either side carries the Ada floor, including when FP8 is the - // source and the destination is the older `f16x2`. - if packed_conversion_uses_fp8(conversion) { - return ("8.1", "sm_89"); - } - match conversion.destination_format { - PackedConversionDestinationFormat::Bf16x2 | PackedConversionDestinationFormat::F16x2 => { - ("7.0", "sm_80") - } - PackedConversionDestinationFormat::E4m3x2 | PackedConversionDestinationFormat::E5m2x2 => { - ("8.1", "sm_89") - } - } -} - -fn packed_conversion_backend_mechanism( - conversion: &crate::model::PackedConversion, - backend: IntrinsicBackend, -) -> BackendLoweringMechanism { - match (packed_conversion_uses_typed_nvvm(conversion), backend) { - (true, IntrinsicBackend::LlvmNvptx) => BackendLoweringMechanism::TypedNvvm, - _ => BackendLoweringMechanism::InlinePtx, - } -} - -/// Whether the conversion lowers through a typed NVVM intrinsic call. -/// -/// Only the scalar-`f32` pair does. The `f16x2` sources and destinations carry -/// their halves in one integer register on the MIR side, while the typed -/// intrinsics are declared over `<2 x half>`; routing them through inline PTX -/// avoids a packed-half MIR type for no change in emitted code. The recorded -/// evidence assembles both routes to a byte-identical cubin, so the typed route -/// can be switched on later without changing the PTX contract. -fn packed_conversion_uses_typed_nvvm(conversion: &crate::model::PackedConversion) -> bool { - conversion.source_format == PackedConversionSourceFormat::F32x2 - && matches!( - conversion.destination_format, - PackedConversionDestinationFormat::E4m3x2 | PackedConversionDestinationFormat::E5m2x2 - ) -} - -fn packed_conversion_lowering(conversion: &crate::model::PackedConversion) -> &'static str { - if packed_conversion_uses_typed_nvvm(conversion) { - "generated_packed_conversion_backend" - } else { - "generated_packed_conversion_inline_ptx" - } -} - -fn expand_packed_conversion_fp8_admission( - admission: &PackedConversionFp8Admission, -) -> Result> { - ensure!( - admission.runtime_validation == RuntimeValidation::Unexecuted, - "FP8 conversion runtime validation may be marked executed only with GPU evidence" - ); - ensure!( - admission.destination_formats - == [ - PackedConversionDestinationFormat::E4m3x2, - PackedConversionDestinationFormat::E5m2x2, - ], - "compact FP8 conversion admission must list the canonical two formats" - ); - ensure!( - admission.saturations - == [ - PackedConversionSaturation::Satfinite, - PackedConversionSaturation::SatfiniteRelu, - ], - "compact FP8 conversion admission must list base and ReLU finite saturation" - ); - ensure!( - admission.product_count - == admission - .destination_formats - .len() - .checked_mul(admission.saturations.len()) - .context("compact FP8 conversion product count overflow")? - && admission.product_count == 4, - "compact FP8 conversion product_count must be exactly 4" - ); - - let mut records = Vec::with_capacity(admission.product_count); - for &destination_format in &admission.destination_formats { - for &saturation in &admission.saturations { - let conversion = crate::model::PackedConversion { - source_format: PackedConversionSourceFormat::F32x2, - destination_format, - rounding: PackedConversionRounding::NearestEven, - saturation, - adapter: PackedConversionAdapter::ReverseHighLowOperands, - }; - records.push(packed_conversion_overlay_record( - conversion, - &admission.llvm_evidence_profile, - &admission.libnvvm_evidence_profile, - )?); - } - } - ensure!(records.len() == admission.product_count); - Ok(records) -} - -fn expand_packed_conversion_fp8_f16x2_admission( - admission: &PackedConversionFp8F16x2Admission, -) -> Result> { - ensure!( - admission.runtime_validation == RuntimeValidation::Unexecuted, - "FP8 f16x2 conversion runtime validation may be marked executed only with GPU evidence" - ); - ensure!( - admission.fp8_formats - == [ - PackedConversionFp8Format::E4m3x2, - PackedConversionFp8Format::E5m2x2, - ], - "compact FP8 f16x2 conversion admission must list the canonical two formats" - ); - ensure!( - admission.directions - == [ - PackedConversionFp8Direction::Pack, - PackedConversionFp8Direction::Unpack, - ], - "compact FP8 f16x2 conversion admission must list both conversion directions" - ); - ensure!( - admission.relu_variants, - "compact FP8 f16x2 conversion admission must admit the ReLU variants" - ); - ensure!( - admission.product_count - == admission - .fp8_formats - .len() - .checked_mul(admission.directions.len()) - .and_then(|count| count.checked_mul(2)) - .context("compact FP8 f16x2 conversion product count overflow")? - && admission.product_count == 8, - "compact FP8 f16x2 conversion product_count must be exactly 8" - ); - - let mut records = Vec::with_capacity(admission.product_count); - for &fp8_format in &admission.fp8_formats { - for &direction in &admission.directions { - for relu in [false, true] { - let conversion = match direction { - // Narrowing to FP8 always saturates to finite. - PackedConversionFp8Direction::Pack => crate::model::PackedConversion { - source_format: PackedConversionSourceFormat::F16x2, - destination_format: match fp8_format { - PackedConversionFp8Format::E4m3x2 => { - PackedConversionDestinationFormat::E4m3x2 - } - PackedConversionFp8Format::E5m2x2 => { - PackedConversionDestinationFormat::E5m2x2 - } - }, - rounding: PackedConversionRounding::NearestEven, - saturation: if relu { - PackedConversionSaturation::SatfiniteRelu - } else { - PackedConversionSaturation::Satfinite - }, - adapter: PackedConversionAdapter::Identity, - }, - // Widening back to f16x2 is exact, so it carries no - // saturation modifier of its own. - PackedConversionFp8Direction::Unpack => crate::model::PackedConversion { - source_format: match fp8_format { - PackedConversionFp8Format::E4m3x2 => { - PackedConversionSourceFormat::E4m3x2 - } - PackedConversionFp8Format::E5m2x2 => { - PackedConversionSourceFormat::E5m2x2 - } - }, - destination_format: PackedConversionDestinationFormat::F16x2, - rounding: PackedConversionRounding::NearestEven, - saturation: if relu { - PackedConversionSaturation::Relu - } else { - PackedConversionSaturation::None - }, - adapter: PackedConversionAdapter::Identity, - }, - }; - records.push(packed_conversion_overlay_record( - conversion, - &admission.llvm_evidence_profile, - &admission.libnvvm_evidence_profile, - )?); - } - } - } - ensure!(records.len() == admission.product_count); - Ok(records) -} - -fn packed_conversion_overlay_record( - conversion: crate::model::PackedConversion, - llvm_evidence_profile: &str, - libnvvm_evidence_profile: &str, -) -> Result { - let recipe = packed_conversion_recipe(&conversion) - .context("compact FP8 conversion is outside the closed recipe set")?; - let result_width = packed_conversion_result_width(&conversion); - let rust_result = format!("u{result_width}"); - let dialect_result = format!("i{result_width}"); - let (minimum_ptx, minimum_sm) = packed_conversion_floor(&conversion); - let (rust_arguments, dialect_operands, llvm_arguments) = - packed_conversion_source_types(&conversion); - // `cvt` writes one destination and reads every source operand. - let ptx_operands = conversion - .source_format - .operand_count() - .checked_add(1) - .context("packed-conversion operand count overflow")?; - Ok(OverlayIntrinsic { - id: recipe.id.into(), - abi_id: String::new(), - operation_key: recipe.operation_key.into(), - family: "packed_conversion".into(), - source: None, - source_record: Some(recipe.source_record.into()), - rust_module: "convert".into(), - rust_name: recipe.rust_name.into(), - rust_arguments, - rust_result: rust_result.clone(), - safe: true, - must_use: false, - safe_allowlist_reason: Some("This conversion has no caller obligations.".into()), - public_rust_path: format!("cuda_intrinsics::convert::{}", recipe.rust_name), - compatibility_rust_paths: vec![recipe.compatibility_path.into()], - dialect_op_type: recipe.dialect_op_type.into(), - dialect_op_name: recipe.dialect_op_name.into(), - dialect_operands, - dialect_results: vec![dialect_result], - llvm_symbol: Some(recipe.llvm_symbol.into()), - resolved_llvm_symbol: None, - llvm_arguments, - llvm_results: vec![recipe.llvm_result.into()], - pure: true, - memory: "none".into(), - convergent: false, - execution_scope: "thread".into(), - minimum_ptx: minimum_ptx.into(), - minimum_sm: Some(minimum_sm.into()), - ptx_result: rust_result, - targets: "all".into(), - ptx_isa_version: "9.3".into(), - ptx_isa_section: "9.7.9.22 Data Movement and Conversion Instructions: cvt".into(), - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cvt".into(), - lowering: packed_conversion_lowering(&conversion).into(), - backend_lowerings: [ - (IntrinsicBackend::LlvmNvptx, llvm_evidence_profile), - (IntrinsicBackend::LibNvvm, libnvvm_evidence_profile), - ] - .into_iter() - .map(|(backend, evidence_profile)| OverlayBackendLowering { - backend, - mechanism: packed_conversion_backend_mechanism(&conversion, backend), - evidence_profile: evidence_profile.into(), - targets: None, - minimum_ptx: Some(minimum_ptx.into()), - minimum_sm: Some(minimum_sm.into()), - }) - .collect(), - packed_atomic: None, - redux: None, - vote: None, - active_mask: None, - warp_match: None, - warp_barrier: None, - warp_shuffle: None, - dot_product: None, - packed_alu: None, - integer_minmax: None, - packed_conversion: Some(conversion.clone()), - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cp_async_copy: None, - cp_async_control: None, - cp_async_mbarrier: None, - mbarrier_basic: None, - movmatrix: None, - mbarrier_extended: None, - register_mma: None, - sparse_mma: None, - prmt: None, - cluster_barrier: None, - wgmma_control: None, - special_register: None, - debug_control: None, - cluster_memory: None, - clc: None, - tma: None, - tcgen05: None, - ldmatrix_variant: None, - ldmatrix_safety: None, - ldmatrix_adapter: None, - selected_address_space: None, - expected_ptx: InstructionPattern { - mnemonic: "cvt".into(), - modifiers: packed_conversion_ptx_modifiers(&conversion) - .into_iter() - .map(str::to_owned) - .collect(), - operands: vec![OperandPattern::Register; ptx_operands], - }, - summary: recipe.summary.into(), - }) -} - -fn validate_packed_conversion_policy( - policy: &OverlayIntrinsic, - source: &IntrinsicSource, - declaration: Option<&ImportedIntrinsic>, -) -> Result<()> { - let conversion = policy - .packed_conversion - .as_ref() - .with_context(|| format!("{} has no closed packed-conversion contract", policy.id))?; - // The adapter is a function of the source arity: two scalar operands are - // reversed so the first Rust argument lands in the low half, while a single - // packed operand is forwarded unchanged. - let expected_adapter = match conversion.source_format { - PackedConversionSourceFormat::F32x2 => PackedConversionAdapter::ReverseHighLowOperands, - PackedConversionSourceFormat::E4m3x2 - | PackedConversionSourceFormat::E5m2x2 - | PackedConversionSourceFormat::F16x2 => PackedConversionAdapter::Identity, - }; - ensure!( - conversion.adapter == expected_adapter, - "{} requests an unsupported packed-conversion source or adapter", - policy.id - ); - let recipe = packed_conversion_recipe(conversion).with_context(|| { - format!( - "{} requests an unsupported packed-conversion source, destination, rounding, or saturation combination", - policy.id - ) - })?; - let result_width = packed_conversion_result_width(conversion); - let rust_result = format!("u{result_width}"); - let dialect_result = format!("i{result_width}"); - let (minimum_ptx, minimum_sm) = packed_conversion_floor(conversion); - let (expected_rust_arguments, expected_dialect_operands, expected_llvm_arguments) = - packed_conversion_source_types(conversion); - ensure!( - policy.id == recipe.id - && policy.abi_id == recipe.abi_id - && policy.operation_key == recipe.operation_key - && source - == &IntrinsicSource::LlvmImported { - source_record: recipe.source_record.into(), - } - && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) - && policy.resolved_llvm_symbol.is_none() - && policy.llvm_arguments == expected_llvm_arguments - && policy.llvm_results == [recipe.llvm_result], - "{} packed-conversion identity or LLVM source changed", - policy.id - ); - let declaration = declaration.context("packed conversion has no imported declaration")?; - ensure!( - declaration.properties == ["IntrNoMem", "IntrSpeculatable"] - && declaration.selections.is_empty(), - "{} selectionless packed-conversion declaration changed", - policy.id - ); - ensure!( - policy.rust_module == "convert" - && policy.rust_name == recipe.rust_name - && policy.rust_arguments == expected_rust_arguments - && policy.rust_result == rust_result - && policy.safe - && !policy.must_use - && policy - .safe_allowlist_reason - .as_deref() - .is_some_and(|reason| !reason.trim().is_empty()) - && policy.public_rust_path == format!("cuda_intrinsics::convert::{}", recipe.rust_name) - && policy.compatibility_rust_paths == [recipe.compatibility_path], - "{} must preserve its safe non-must-use conversion API", - policy.id - ); - ensure!( - policy.dialect_op_type == recipe.dialect_op_type - && policy.dialect_op_name == recipe.dialect_op_name - && policy.dialect_operands == expected_dialect_operands - && policy.dialect_results == [dialect_result.as_str()] - && policy.lowering == packed_conversion_lowering(conversion), - "{} is outside the closed packed-conversion dialect and lowering recipe", - policy.id - ); - ensure!( - policy.pure - && policy.memory == "none" - && !policy.convergent - && policy.execution_scope == "thread" - && policy.minimum_ptx == minimum_ptx - && policy.minimum_sm.as_deref() == Some(minimum_sm) - && policy.ptx_result == rust_result - && policy.targets == "all" - && policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section == "9.7.9.22 Data Movement and Conversion Instructions: cvt" - && policy.ptx_isa_url - == "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cvt", - "{} packed-conversion effects, carrier, target floor, or PTX provenance disagree", - policy.id - ); - let expected_ptx_operands = vec![ - OperandPattern::Register; - conversion - .source_format - .operand_count() - .checked_add(1) - .context("packed-conversion operand count overflow")? - ]; - ensure!( - policy.expected_ptx.mnemonic == "cvt" - && policy.expected_ptx.modifiers == packed_conversion_ptx_modifiers(conversion) - && policy.expected_ptx.operands == expected_ptx_operands, - "{} expected PTX does not match its closed conversion recipe", - policy.id - ); - ensure!( - policy.summary == recipe.summary, - "{} packed-conversion summary does not match its closed recipe", - policy.id - ); - let backend_pairs = policy - .backend_lowerings - .iter() - .map(|lowering| (lowering.backend, lowering.mechanism)) - .collect::>(); - let expected_pairs = [IntrinsicBackend::LlvmNvptx, IntrinsicBackend::LibNvvm] - .map(|backend| { - ( - backend, - packed_conversion_backend_mechanism(conversion, backend), - ) - }) - .into_iter() - .collect::>(); - ensure!( - policy.backend_lowerings.len() == 2 && backend_pairs == expected_pairs, - "{} must define exactly the reviewed packed-conversion backend routes", - policy.id - ); - for lowering in &policy.backend_lowerings { - ensure!( - lowering.minimum_ptx.as_deref() == Some(minimum_ptx) - && lowering.minimum_sm.as_deref() == Some(minimum_sm) - && !lowering.evidence_profile.trim().is_empty(), - "{} backend {:?} does not carry its exact packed-conversion floor", - policy.id, - lowering.backend - ); - } - ensure_no_other_family_contract(policy, "packed conversion")?; - Ok(()) -} - -#[derive(Clone, Copy)] -struct ScalarConversionRecipe { - id: &'static str, - abi_id: &'static str, - operation_key: &'static str, - source_record: &'static str, - llvm_symbol: &'static str, - rust_name: &'static str, - selection_record: &'static str, - selection_asm: &'static str, - minimum_ptx: &'static str, - minimum_sm: &'static str, - ptx_modifiers: &'static [&'static str], -} - -fn scalar_conversion_recipe( - rounding: ScalarConversionRounding, - saturation: ScalarConversionSaturation, -) -> Option { - use ScalarConversionRounding::{NearestAway, NearestEven, TowardZero}; - use ScalarConversionSaturation::{None, Relu, ReluSatfinite, Satfinite}; - - Some(match (rounding, saturation) { - (NearestAway, None) => ScalarConversionRecipe { - id: "cvt_rna_tf32_f32", - abi_id: "i0368", - operation_key: "convert.f32.tf32.rna", - source_record: "int_nvvm_f2tf32_rna", - llvm_symbol: "llvm.nvvm.f2tf32.rna", - rust_name: "cvt_rna_tf32_f32", - selection_record: "CVT_to_tf32_rna", - selection_asm: "cvt.rna.tf32.f32 \t$dst, $src;", - minimum_ptx: "7.0", - minimum_sm: "sm_80", - ptx_modifiers: &["rna", "tf32", "f32"], - }, - (NearestAway, Satfinite) => ScalarConversionRecipe { - id: "cvt_rna_satfinite_tf32_f32", - abi_id: "i0369", - operation_key: "convert.f32.tf32.rna.satfinite", - source_record: "int_nvvm_f2tf32_rna_satfinite", - llvm_symbol: "llvm.nvvm.f2tf32.rna.satfinite", - rust_name: "cvt_rna_satfinite_tf32_f32", - selection_record: "CVT_to_tf32_rna_satf", - selection_asm: "cvt.rna.satfinite.tf32.f32 \t$dst, $src;", - minimum_ptx: "8.1", - minimum_sm: "sm_80", - ptx_modifiers: &["rna", "satfinite", "tf32", "f32"], - }, - (NearestEven, None) => ScalarConversionRecipe { - id: "cvt_rn_tf32_f32", - abi_id: "i0370", - operation_key: "convert.f32.tf32.rn", - source_record: "int_nvvm_f2tf32_rn", - llvm_symbol: "llvm.nvvm.f2tf32.rn", - rust_name: "cvt_rn_tf32_f32", - selection_record: "CVT_to_tf32_rn", - selection_asm: "cvt.rn.tf32.f32 \t$dst, $src;", - minimum_ptx: "7.8", - minimum_sm: "sm_90", - ptx_modifiers: &["rn", "tf32", "f32"], - }, - (NearestEven, Relu) => ScalarConversionRecipe { - id: "cvt_rn_relu_tf32_f32", - abi_id: "i0371", - operation_key: "convert.f32.tf32.rn.relu", - source_record: "int_nvvm_f2tf32_rn_relu", - llvm_symbol: "llvm.nvvm.f2tf32.rn.relu", - rust_name: "cvt_rn_relu_tf32_f32", - selection_record: "CVT_to_tf32_rn_relu", - selection_asm: "cvt.rn.relu.tf32.f32 \t$dst, $src;", - minimum_ptx: "7.8", - minimum_sm: "sm_90", - ptx_modifiers: &["rn", "relu", "tf32", "f32"], - }, - (NearestEven, Satfinite) => ScalarConversionRecipe { - id: "cvt_rn_satfinite_tf32_f32", - abi_id: "i0372", - operation_key: "convert.f32.tf32.rn.satfinite", - source_record: "int_nvvm_f2tf32_rn_satfinite", - llvm_symbol: "llvm.nvvm.f2tf32.rn.satfinite", - rust_name: "cvt_rn_satfinite_tf32_f32", - selection_record: "CVT_to_tf32_rn_satf", - selection_asm: "cvt.rn.satfinite.tf32.f32 \t$dst, $src;", - minimum_ptx: "8.6", - minimum_sm: "sm_100", - ptx_modifiers: &["rn", "satfinite", "tf32", "f32"], - }, - (NearestEven, ReluSatfinite) => ScalarConversionRecipe { - id: "cvt_rn_relu_satfinite_tf32_f32", - abi_id: "i0373", - operation_key: "convert.f32.tf32.rn.relu.satfinite", - source_record: "int_nvvm_f2tf32_rn_relu_satfinite", - llvm_symbol: "llvm.nvvm.f2tf32.rn.relu.satfinite", - rust_name: "cvt_rn_relu_satfinite_tf32_f32", - selection_record: "CVT_to_tf32_rn_relu_satf", - selection_asm: "cvt.rn.relu.satfinite.tf32.f32 \t$dst, $src;", - minimum_ptx: "8.6", - minimum_sm: "sm_100", - ptx_modifiers: &["rn", "relu", "satfinite", "tf32", "f32"], - }, - (TowardZero, None) => ScalarConversionRecipe { - id: "cvt_rz_tf32_f32", - abi_id: "i0374", - operation_key: "convert.f32.tf32.rz", - source_record: "int_nvvm_f2tf32_rz", - llvm_symbol: "llvm.nvvm.f2tf32.rz", - rust_name: "cvt_rz_tf32_f32", - selection_record: "CVT_to_tf32_rz", - selection_asm: "cvt.rz.tf32.f32 \t$dst, $src;", - minimum_ptx: "7.8", - minimum_sm: "sm_90", - ptx_modifiers: &["rz", "tf32", "f32"], - }, - (TowardZero, Relu) => ScalarConversionRecipe { - id: "cvt_rz_relu_tf32_f32", - abi_id: "i0375", - operation_key: "convert.f32.tf32.rz.relu", - source_record: "int_nvvm_f2tf32_rz_relu", - llvm_symbol: "llvm.nvvm.f2tf32.rz.relu", - rust_name: "cvt_rz_relu_tf32_f32", - selection_record: "CVT_to_tf32_rz_relu", - selection_asm: "cvt.rz.relu.tf32.f32 \t$dst, $src;", - minimum_ptx: "7.8", - minimum_sm: "sm_90", - ptx_modifiers: &["rz", "relu", "tf32", "f32"], - }, - (TowardZero, Satfinite) => ScalarConversionRecipe { - id: "cvt_rz_satfinite_tf32_f32", - abi_id: "i0376", - operation_key: "convert.f32.tf32.rz.satfinite", - source_record: "int_nvvm_f2tf32_rz_satfinite", - llvm_symbol: "llvm.nvvm.f2tf32.rz.satfinite", - rust_name: "cvt_rz_satfinite_tf32_f32", - selection_record: "CVT_to_tf32_rz_satf", - selection_asm: "cvt.rz.satfinite.tf32.f32 \t$dst, $src;", - minimum_ptx: "8.6", - minimum_sm: "sm_100", - ptx_modifiers: &["rz", "satfinite", "tf32", "f32"], - }, - (TowardZero, ReluSatfinite) => ScalarConversionRecipe { - id: "cvt_rz_relu_satfinite_tf32_f32", - abi_id: "i0377", - operation_key: "convert.f32.tf32.rz.relu.satfinite", - source_record: "int_nvvm_f2tf32_rz_relu_satfinite", - llvm_symbol: "llvm.nvvm.f2tf32.rz.relu.satfinite", - rust_name: "cvt_rz_relu_satfinite_tf32_f32", - selection_record: "CVT_to_tf32_rz_relu_satf", - selection_asm: "cvt.rz.relu.satfinite.tf32.f32 \t$dst, $src;", - minimum_ptx: "8.6", - minimum_sm: "sm_100", - ptx_modifiers: &["rz", "relu", "satfinite", "tf32", "f32"], - }, - _ => return Option::None, - }) -} - -fn expand_scalar_conversion_admission( - admission: &ScalarConversionAdmission, -) -> Result> { - ensure!( - admission.runtime_validation == RuntimeValidation::Unexecuted, - "scalar-conversion runtime may be marked executed only with GPU evidence" - ); - let expected = [ - ( - ScalarConversionRounding::NearestAway, - ScalarConversionSaturation::None, - ), - ( - ScalarConversionRounding::NearestAway, - ScalarConversionSaturation::Satfinite, - ), - ( - ScalarConversionRounding::NearestEven, - ScalarConversionSaturation::None, - ), - ( - ScalarConversionRounding::NearestEven, - ScalarConversionSaturation::Relu, - ), - ( - ScalarConversionRounding::NearestEven, - ScalarConversionSaturation::Satfinite, - ), - ( - ScalarConversionRounding::NearestEven, - ScalarConversionSaturation::ReluSatfinite, - ), - ( - ScalarConversionRounding::TowardZero, - ScalarConversionSaturation::None, - ), - ( - ScalarConversionRounding::TowardZero, - ScalarConversionSaturation::Relu, - ), - ( - ScalarConversionRounding::TowardZero, - ScalarConversionSaturation::Satfinite, - ), - ( - ScalarConversionRounding::TowardZero, - ScalarConversionSaturation::ReluSatfinite, - ), - ]; - let actual = admission - .variants - .iter() - .map(|variant| (variant.rounding, variant.saturation)) - .collect::>(); - ensure!( - actual == expected, - "compact scalar-conversion admission must list the canonical ten variants" - ); - - admission - .variants - .iter() - .map(|variant| { - let recipe = scalar_conversion_recipe(variant.rounding, variant.saturation) - .context("scalar conversion is outside the closed recipe set")?; - ensure!( - variant.abi_id == recipe.abi_id, - "{} must reserve ABI ID {}", - recipe.id, - recipe.abi_id - ); - scalar_conversion_overlay_record( - recipe, - admission, - variant.rounding, - variant.saturation, - ) - }) - .collect() -} - -fn scalar_conversion_overlay_record( - recipe: ScalarConversionRecipe, - admission: &ScalarConversionAdmission, - rounding: ScalarConversionRounding, - saturation: ScalarConversionSaturation, -) -> Result { - Ok(OverlayIntrinsic { - id: recipe.id.into(), - abi_id: recipe.abi_id.into(), - operation_key: recipe.operation_key.into(), - family: "scalar_conversion".into(), - source: None, - source_record: Some(recipe.source_record.into()), - rust_module: "convert".into(), - rust_name: recipe.rust_name.into(), - rust_arguments: vec!["f32".into()], - rust_result: "u32".into(), - safe: true, - must_use: true, - safe_allowlist_reason: Some("This conversion has no caller obligations.".into()), - public_rust_path: format!("cuda_intrinsics::convert::{}", recipe.rust_name), - compatibility_rust_paths: vec![format!( - "cuda_device::convert::{}", - recipe.rust_name - )], - dialect_op_type: "ScalarConversionOp".into(), - dialect_op_name: "nvvm.scalar_conversion".into(), - dialect_operands: vec!["f32".into()], - dialect_results: vec!["i32".into()], - llvm_symbol: Some(recipe.llvm_symbol.into()), - resolved_llvm_symbol: None, - llvm_arguments: vec!["f32".into()], - llvm_results: vec!["i32".into()], - pure: true, - memory: "none".into(), - convergent: false, - execution_scope: "thread".into(), - minimum_ptx: recipe.minimum_ptx.into(), - minimum_sm: Some(recipe.minimum_sm.into()), - ptx_result: "u32".into(), - targets: "all".into(), - ptx_isa_version: "9.3".into(), - ptx_isa_section: "9.7.9.22 Data Movement and Conversion Instructions: cvt".into(), - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cvt".into(), - lowering: "generated_scalar_conversion".into(), - backend_lowerings: [ - (IntrinsicBackend::LlvmNvptx, &admission.llvm_evidence_profile), - (IntrinsicBackend::LibNvvm, &admission.libnvvm_evidence_profile), - ] - .into_iter() - .map(|(backend, evidence_profile)| OverlayBackendLowering { - backend, - mechanism: match backend { - IntrinsicBackend::LlvmNvptx => BackendLoweringMechanism::TypedNvvm, - IntrinsicBackend::LibNvvm => BackendLoweringMechanism::InlinePtx, - }, - evidence_profile: evidence_profile.clone(), - targets: None, - minimum_ptx: Some(recipe.minimum_ptx.into()), - minimum_sm: Some(recipe.minimum_sm.into()), - }) - .collect(), - packed_atomic: None, - redux: None, - vote: None, - active_mask: None, - warp_match: None, - warp_barrier: None, - warp_shuffle: None, - dot_product: None, - packed_alu: None, - integer_minmax: None, - packed_conversion: None, - scalar_conversion: Some(ScalarConversion { - source_format: ScalarConversionSourceFormat::F32, - destination_format: ScalarConversionDestinationFormat::Tf32, - rounding, - saturation, - result_representation: ScalarConversionResultRepresentation::RawU32Bits, - adapter: ScalarConversionAdapter::DirectF32ToRawU32Bits, - runtime_validation: admission.runtime_validation, - }), - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cp_async_copy: None, - cp_async_control: None, - cp_async_mbarrier: None, - mbarrier_basic: None, - movmatrix: None, - mbarrier_extended: None, - register_mma: None, - sparse_mma: None, - prmt: None, - cluster_barrier: None, - wgmma_control: None, - special_register: None, - debug_control: None, - cluster_memory: None, - clc: None, - tma: None, - tcgen05: None, - ldmatrix_variant: None, - ldmatrix_safety: None, - ldmatrix_adapter: None, - selected_address_space: None, - expected_ptx: InstructionPattern { - mnemonic: "cvt".into(), - modifiers: recipe.ptx_modifiers.iter().map(|value| (*value).into()).collect(), - operands: vec![OperandPattern::Register, OperandPattern::Register], - }, - summary: format!( - "Converts one f32 value with {} and returns raw TF32 bits.", - recipe.ptx_modifiers[0] - ), - }) -} - -fn validate_scalar_conversion_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, -) -> Result<()> { - let conversion = policy - .scalar_conversion - .as_ref() - .with_context(|| format!("{} has no scalar-conversion contract", policy.id))?; - let recipe = scalar_conversion_recipe(conversion.rounding, conversion.saturation) - .with_context(|| { - format!( - "{} is outside the closed scalar-conversion recipe", - policy.id - ) - })?; - ensure!( - conversion.source_format == ScalarConversionSourceFormat::F32 - && conversion.destination_format == ScalarConversionDestinationFormat::Tf32 - && conversion.result_representation == ScalarConversionResultRepresentation::RawU32Bits - && conversion.adapter == ScalarConversionAdapter::DirectF32ToRawU32Bits - && conversion.runtime_validation == RuntimeValidation::Unexecuted, - "{} changed its scalar-conversion representation or adapter", - policy.id - ); - ensure!( - policy.id == recipe.id - && policy.abi_id == recipe.abi_id - && policy.operation_key == recipe.operation_key - && policy.source_record.as_deref() == Some(recipe.source_record) - && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) - && policy.resolved_llvm_symbol.is_none() - && policy.llvm_arguments == ["f32"] - && policy.llvm_results == ["i32"], - "{} scalar-conversion identity or LLVM source changed", - policy.id - ); - ensure!( - declaration.properties == ["IntrNoMem", "IntrSpeculatable"] - && declaration.selections.len() == 1 - && declaration.selections[0].source_record == recipe.selection_record - && declaration.selections[0].asm == recipe.selection_asm - && declaration.selections[0].constraints.is_empty(), - "{} imported scalar-conversion declaration or selection changed", - policy.id - ); - ensure!( - policy.rust_module == "convert" - && policy.rust_name == recipe.rust_name - && policy.rust_arguments == ["f32"] - && policy.rust_result == "u32" - && policy.safe - && policy.must_use - && policy.compatibility_rust_paths - == [format!("cuda_device::convert::{}", recipe.rust_name)] - && policy.dialect_op_type == "ScalarConversionOp" - && policy.dialect_op_name == "nvvm.scalar_conversion" - && policy.dialect_operands == ["f32"] - && policy.dialect_results == ["i32"] - && policy.lowering == "generated_scalar_conversion", - "{} changed its scalar-conversion API, carrier, or lowering", - policy.id - ); - ensure!( - policy.pure - && policy.memory == "none" - && !policy.convergent - && policy.execution_scope == "thread" - && policy.minimum_ptx == recipe.minimum_ptx - && policy.minimum_sm.as_deref() == Some(recipe.minimum_sm) - && policy.ptx_result == "u32" - && policy.targets == "all", - "{} scalar-conversion effects or target floor changed", - policy.id - ); - ensure!( - policy.expected_ptx - == InstructionPattern { - mnemonic: "cvt".into(), - modifiers: recipe - .ptx_modifiers - .iter() - .map(|value| (*value).into()) - .collect(), - operands: vec![OperandPattern::Register, OperandPattern::Register], - }, - "{} expected PTX changed", - policy.id - ); - let expected_backends = [ - ( - IntrinsicBackend::LlvmNvptx, - BackendLoweringMechanism::TypedNvvm, - ), - ( - IntrinsicBackend::LibNvvm, - BackendLoweringMechanism::InlinePtx, - ), - ]; - ensure!( - policy.backend_lowerings.len() == 2 - && expected_backends.into_iter().all(|(backend, mechanism)| { - policy.backend_lowerings.iter().any(|lowering| { - lowering.backend == backend - && lowering.mechanism == mechanism - && lowering.minimum_ptx.as_deref() == Some(recipe.minimum_ptx) - && lowering.minimum_sm.as_deref() == Some(recipe.minimum_sm) - && !lowering.evidence_profile.trim().is_empty() - }) - }), - "{} must define the typed LLVM and inline-PTX libNVVM routes", - policy.id - ); - validate_selected_target_predicates(policy, &declaration.selections[0])?; - ensure_no_other_family_contract(policy, "scalar conversion")?; - Ok(()) -} - -#[derive(Clone)] -struct ScalarArithmeticRecipe { - id: String, - operation_key: String, - source_record: String, - llvm_symbol: String, - selection_record: String, - selection_asm: String, - rust_type: &'static str, - argument_count: usize, - properties: Vec<&'static str>, - ptx_modifiers: Vec, - ptx_isa_section: &'static str, - ptx_isa_url: &'static str, -} - -type ScalarArithmeticVariant = ( - ScalarArithmeticFormat, - ScalarArithmeticOperation, - ScalarArithmeticRounding, - ScalarArithmeticSubnormal, - ScalarArithmeticSaturation, -); - -fn canonical_scalar_arithmetic_variants() -> Vec { - use ScalarArithmeticFormat::{F32, F64}; - use ScalarArithmeticOperation::{Add, Div, Fma, Mul}; - use ScalarArithmeticRounding::{Rm, Rn, Rp, Rz}; - use ScalarArithmeticSaturation::{None, Sat}; - use ScalarArithmeticSubnormal::{Ftz, Preserve}; - - let roundings = [Rn, Rz, Rm, Rp]; - let mut variants = Vec::with_capacity(64); - for operation in [Mul, Div, Fma] { - for rounding in roundings { - variants.push((F64, operation, rounding, Preserve, None)); - } - } - for operation in [Mul, Div] { - for rounding in roundings { - variants.push((F32, operation, rounding, Preserve, None)); - variants.push((F32, operation, rounding, Ftz, None)); - } - } - for rounding in roundings { - variants.push((F32, Fma, rounding, Preserve, None)); - variants.push((F32, Fma, rounding, Ftz, None)); - variants.push((F32, Fma, rounding, Preserve, Sat)); - variants.push((F32, Fma, rounding, Ftz, Sat)); - } - for rounding in roundings { - variants.push((F64, Add, rounding, Preserve, None)); - } - for rounding in roundings { - variants.push((F32, Add, rounding, Preserve, None)); - variants.push((F32, Add, rounding, Ftz, None)); - variants.push((F32, Add, rounding, Preserve, Sat)); - variants.push((F32, Add, rounding, Ftz, Sat)); - } - variants -} - -fn scalar_arithmetic_format_name(format: ScalarArithmeticFormat) -> &'static str { - match format { - ScalarArithmeticFormat::F32 => "f32", - ScalarArithmeticFormat::F64 => "f64", - } -} - -fn scalar_arithmetic_operation_name(operation: ScalarArithmeticOperation) -> &'static str { - match operation { - ScalarArithmeticOperation::Mul => "mul", - ScalarArithmeticOperation::Div => "div", - ScalarArithmeticOperation::Fma => "fma", - ScalarArithmeticOperation::Add => "add", - } -} - -fn scalar_arithmetic_rounding_name(rounding: ScalarArithmeticRounding) -> &'static str { - match rounding { - ScalarArithmeticRounding::Rn => "rn", - ScalarArithmeticRounding::Rz => "rz", - ScalarArithmeticRounding::Rm => "rm", - ScalarArithmeticRounding::Rp => "rp", - } -} - -fn scalar_arithmetic_recipe(variant: ScalarArithmeticVariant) -> Option { - if !canonical_scalar_arithmetic_variants().contains(&variant) { - return None; - } - let (format, operation, rounding, subnormal, saturation) = variant; - let operation_name = scalar_arithmetic_operation_name(operation); - let rounding_name = scalar_arithmetic_rounding_name(rounding); - let format_name = scalar_arithmetic_format_name(format); - let source_format = match format { - ScalarArithmeticFormat::F32 => "f", - ScalarArithmeticFormat::F64 => "d", - }; - - let mut modifier_names = vec![rounding_name]; - if subnormal == ScalarArithmeticSubnormal::Ftz { - modifier_names.push("ftz"); - } - if saturation == ScalarArithmeticSaturation::Sat { - modifier_names.push("sat"); - } - - let modifier_id = modifier_names.join("_"); - let modifier_symbol = modifier_names.join("."); - let id = format!("{operation_name}_{modifier_id}_{format_name}"); - let source_record = format!("int_nvvm_{operation_name}_{modifier_id}_{source_format}"); - let llvm_symbol = format!("llvm.nvvm.{operation_name}.{modifier_symbol}.{source_format}"); - let mut ptx_modifier_names = modifier_names.clone(); - if operation == ScalarArithmeticOperation::Add - && subnormal == ScalarArithmeticSubnormal::Ftz - && saturation == ScalarArithmeticSaturation::Sat - { - ptx_modifier_names.swap(1, 2); - } - let ptx_modifiers = ptx_modifier_names - .iter() - .copied() - .chain(std::iter::once(format_name)) - .map(str::to_owned) - .collect::>(); - let selection_record = match operation { - ScalarArithmeticOperation::Fma => { - format!("INT_NVVM_FMA_{modifier_id}_{format_name}") - } - ScalarArithmeticOperation::Mul | ScalarArithmeticOperation::Div => format!( - "INT_NVVM_{}_{}_{}", - operation_name.to_ascii_uppercase(), - modifier_id.to_ascii_uppercase(), - source_format.to_ascii_uppercase() - ), - ScalarArithmeticOperation::Add => { - let selection_modifier = ptx_modifier_names.join("_").to_ascii_uppercase(); - format!( - "INT_NVVM_ADD_{selection_modifier}_{}", - source_format.to_ascii_uppercase() - ) - } - }; - let argument_count = if operation == ScalarArithmeticOperation::Fma { - 3 - } else { - 2 - }; - let operands = (0..argument_count) - .map(|argument| format!("$src{argument}")) - .collect::>() - .join(", "); - let ptx_modifier_symbol = ptx_modifier_names.join("."); - let selection_asm = - format!("{operation_name}.{ptx_modifier_symbol}.{format_name} \t$dst, {operands};"); - let properties = match (operation, format) { - (ScalarArithmeticOperation::Mul | ScalarArithmeticOperation::Add, _) => { - vec!["Commutative", "IntrNoMem", "IntrSpeculatable"] - } - (ScalarArithmeticOperation::Div, _) => vec!["IntrNoMem"], - (ScalarArithmeticOperation::Fma, ScalarArithmeticFormat::F32) => { - vec!["IntrNoMem", "IntrSpeculatable"] - } - (ScalarArithmeticOperation::Fma, ScalarArithmeticFormat::F64) => { - vec!["IntrNoCreateUndefOrPoison", "IntrNoMem", "IntrSpeculatable"] - } - }; - let (ptx_isa_section, ptx_isa_url) = match operation { - ScalarArithmeticOperation::Mul => ( - "9.7.3.5 Floating Point Instructions: mul", - "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-mul", - ), - ScalarArithmeticOperation::Fma => ( - "9.7.3.6 Floating Point Instructions: fma", - "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-fma", - ), - ScalarArithmeticOperation::Div => ( - "9.7.3.8 Floating Point Instructions: div", - "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-div", - ), - ScalarArithmeticOperation::Add => ( - "9.7.3.3 Floating Point Instructions: add", - "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-add", - ), - }; - - Some(ScalarArithmeticRecipe { - id, - operation_key: format!( - "scalar.arithmetic.{operation_name}.{modifier_symbol}.{format_name}" - ), - source_record, - llvm_symbol, - selection_record, - selection_asm, - rust_type: format_name, - argument_count, - properties, - ptx_modifiers, - ptx_isa_section, - ptx_isa_url, - }) -} - -fn expand_scalar_arithmetic_admission( - admission: &ScalarArithmeticAdmission, -) -> Result> { - ensure!( - admission.runtime_validation == RuntimeValidation::Unexecuted, - "scalar-arithmetic runtime may be marked executed only with GPU evidence" - ); - let expected = canonical_scalar_arithmetic_variants(); - let actual = admission - .variants - .iter() - .map(|variant| { - ( - variant.format, - variant.operation, - variant.rounding, - variant.subnormal, - variant.saturation, - ) - }) - .collect::>(); - ensure!( - actual == expected, - "compact scalar-arithmetic admission must list the canonical 64 variants" - ); - - admission - .variants - .iter() - .map(|variant| { - validate_abi_id(&variant.abi_id)?; - let identity = ( - variant.format, - variant.operation, - variant.rounding, - variant.subnormal, - variant.saturation, - ); - let recipe = scalar_arithmetic_recipe(identity) - .context("scalar arithmetic is outside the closed recipe set")?; - scalar_arithmetic_overlay_record(recipe, admission, identity, &variant.abi_id) - }) - .collect() -} - -fn scalar_arithmetic_overlay_record( - recipe: ScalarArithmeticRecipe, - admission: &ScalarArithmeticAdmission, - variant: ScalarArithmeticVariant, - abi_id: &str, -) -> Result { - let (format, operation, rounding, subnormal, saturation) = variant; - let rust_arguments = vec![recipe.rust_type.into(); recipe.argument_count]; - let ptx_operands = vec![OperandPattern::Register; recipe.argument_count + 1]; - let summary = format!( - "Computes scalar {} with explicit {} rounding.", - scalar_arithmetic_operation_name(operation), - scalar_arithmetic_rounding_name(rounding) - ); - Ok(OverlayIntrinsic { - id: recipe.id.clone(), - abi_id: abi_id.into(), - operation_key: recipe.operation_key, - family: "scalar_arithmetic".into(), - source: None, - source_record: Some(recipe.source_record), - rust_module: "float".into(), - rust_name: recipe.id.clone(), - rust_arguments: rust_arguments.clone(), - rust_result: recipe.rust_type.into(), - safe: true, - must_use: true, - safe_allowlist_reason: Some("Scalar arithmetic has no caller obligations.".into()), - public_rust_path: format!("cuda_intrinsics::float::{}", recipe.id), - compatibility_rust_paths: vec![format!("cuda_device::float::{}", recipe.id)], - dialect_op_type: "ScalarArithmeticOp".into(), - dialect_op_name: "nvvm.scalar_arithmetic".into(), - dialect_operands: rust_arguments.clone(), - dialect_results: vec![recipe.rust_type.into()], - llvm_symbol: Some(recipe.llvm_symbol), - resolved_llvm_symbol: None, - llvm_arguments: rust_arguments, - llvm_results: vec![recipe.rust_type.into()], - pure: operation != ScalarArithmeticOperation::Div, - memory: "none".into(), - convergent: false, - execution_scope: "thread".into(), - minimum_ptx: "7.0".into(), - minimum_sm: Some("sm_80".into()), - ptx_result: recipe.rust_type.into(), - targets: "all".into(), - ptx_isa_version: "9.3".into(), - ptx_isa_section: recipe.ptx_isa_section.into(), - ptx_isa_url: recipe.ptx_isa_url.into(), - lowering: "generated_scalar_arithmetic".into(), - backend_lowerings: [ - ( - IntrinsicBackend::LlvmNvptx, - &admission.llvm_evidence_profile, - ), - ( - IntrinsicBackend::LibNvvm, - &admission.libnvvm_evidence_profile, - ), - ] - .into_iter() - .map(|(backend, evidence_profile)| OverlayBackendLowering { - backend, - mechanism: match backend { - IntrinsicBackend::LlvmNvptx if saturation == ScalarArithmeticSaturation::Sat => { - // LLVM 21 has no typed intrinsic for these LLVM 22 forms. - BackendLoweringMechanism::InlinePtx - } - IntrinsicBackend::LlvmNvptx => BackendLoweringMechanism::TypedNvvm, - IntrinsicBackend::LibNvvm => BackendLoweringMechanism::InlinePtx, - }, - evidence_profile: evidence_profile.clone(), - targets: None, - minimum_ptx: Some("7.0".into()), - minimum_sm: Some("sm_80".into()), - }) - .collect(), - packed_atomic: None, - redux: None, - vote: None, - active_mask: None, - warp_match: None, - warp_barrier: None, - warp_shuffle: None, - dot_product: None, - packed_alu: None, - integer_minmax: None, - packed_conversion: None, - scalar_conversion: None, - scalar_arithmetic: Some(ScalarArithmetic { - format, - operation, - rounding, - subnormal, - saturation, - runtime_validation: admission.runtime_validation, - }), - scalar_math: None, - extended_minmax: None, - cp_async_copy: None, - cp_async_control: None, - cp_async_mbarrier: None, - mbarrier_basic: None, - movmatrix: None, - mbarrier_extended: None, - register_mma: None, - sparse_mma: None, - prmt: None, - cluster_barrier: None, - wgmma_control: None, - special_register: None, - debug_control: None, - cluster_memory: None, - clc: None, - tma: None, - tcgen05: None, - ldmatrix_variant: None, - ldmatrix_safety: None, - ldmatrix_adapter: None, - selected_address_space: None, - expected_ptx: InstructionPattern { - mnemonic: scalar_arithmetic_operation_name(operation).into(), - modifiers: recipe.ptx_modifiers, - operands: ptx_operands, - }, - summary, - }) -} - -fn validate_scalar_arithmetic_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, -) -> Result<()> { - let arithmetic = policy - .scalar_arithmetic - .as_ref() - .with_context(|| format!("{} has no scalar-arithmetic contract", policy.id))?; - let variant = ( - arithmetic.format, - arithmetic.operation, - arithmetic.rounding, - arithmetic.subnormal, - arithmetic.saturation, - ); - let recipe = scalar_arithmetic_recipe(variant).with_context(|| { - format!( - "{} is outside the closed scalar-arithmetic recipe", - policy.id - ) - })?; - ensure!( - arithmetic.runtime_validation == RuntimeValidation::Unexecuted, - "{} scalar-arithmetic runtime may be executed only with GPU evidence", - policy.id - ); - let signature = vec![recipe.rust_type.to_owned(); recipe.argument_count]; - ensure!( - policy.id == recipe.id - && policy.operation_key == recipe.operation_key - && policy.source_record.as_deref() == Some(recipe.source_record.as_str()) - && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol.as_str()) - && policy.resolved_llvm_symbol.is_none() - && policy.llvm_arguments == signature - && policy.llvm_results == [recipe.rust_type], - "{} scalar-arithmetic identity or LLVM source changed", - policy.id - ); - let expected_properties = recipe - .properties - .iter() - .map(|property| (*property).to_owned()) - .collect::>(); - ensure!( - declaration.properties == expected_properties, - "{} imported scalar-arithmetic properties changed", - policy.id - ); - let direct = declaration - .selections - .first() - .with_context(|| format!("{} has no imported scalar selection", policy.id))?; - ensure!( - direct.source_record == recipe.selection_record - && direct.asm == recipe.selection_asm - && direct.predicates.is_empty() - && direct.constraints.is_empty(), - "{} direct scalar-arithmetic selection changed", - policy.id - ); - if arithmetic.operation == ScalarArithmeticOperation::Add { - let rounding = scalar_arithmetic_rounding_name(arithmetic.rounding); - let source_modifiers = [ - Some(rounding), - (arithmetic.subnormal == ScalarArithmeticSubnormal::Ftz).then_some("ftz"), - (arithmetic.saturation == ScalarArithmeticSaturation::Sat).then_some("sat"), - ] - .into_iter() - .flatten() - .collect::>() - .join("_"); - let mut expected = Vec::new(); - if arithmetic.format == ScalarArithmeticFormat::F32 - && arithmetic.subnormal == ScalarArithmeticSubnormal::Preserve - { - let saturation = if arithmetic.saturation == ScalarArithmeticSaturation::Sat { - "_sat" - } else { - "" - }; - for instruction in ["ADD", "SUB"] { - for input_format in ["bf16", "f16"] { - expected.push(( - format!( - "INT_NVVM_MIXED_{instruction}_{rounding}{saturation}_f32_{input_format}" - ), - format!( - "{}.{rounding}{}.f32.{input_format} \t$dst, $a, $b;", - instruction.to_ascii_lowercase(), - if arithmetic.saturation == ScalarArithmeticSaturation::Sat { - ".sat" - } else { - "" - } - ), - vec![ - "Subtarget->getSmVersion() >= 100".to_owned(), - "Subtarget->getPTXVersion() >= 86".to_owned(), - ], - )); - } - } - } - let source_format = match arithmetic.format { - ScalarArithmeticFormat::F32 => "F", - ScalarArithmeticFormat::F64 => "D", - }; - expected.push(( - format!("INT_NVVM_SUB_{source_modifiers}_{source_format}"), - format!("sub.{} \t$dst, $a, $b;", recipe.ptx_modifiers.join(".")), - Vec::new(), - )); - ensure!( - declaration.selections.len() == expected.len() + 1, - "{} must retain one direct add and the reviewed non-add alternatives", - policy.id - ); - for (selection, (source_record, asm, predicates)) in - declaration.selections[1..].iter().zip(expected) - { - ensure!( - selection.source_record == source_record - && selection.asm == asm - && selection.predicates == predicates - && selection.constraints.is_empty(), - "{} non-add scalar selection changed", - policy.id - ); - } - } else if arithmetic.format == ScalarArithmeticFormat::F32 - && arithmetic.operation == ScalarArithmeticOperation::Fma - && arithmetic.subnormal == ScalarArithmeticSubnormal::Preserve - { - ensure!( - declaration.selections.len() == 3, - "{} must retain one direct and two mixed-input selections", - policy.id - ); - let rounding = scalar_arithmetic_rounding_name(arithmetic.rounding); - let saturation = if arithmetic.saturation == ScalarArithmeticSaturation::Sat { - "_sat" - } else { - "" - }; - for (selection, input_format) in declaration.selections[1..].iter().zip(["bf16", "f16"]) { - ensure!( - selection.source_record - == format!("INT_NVVM_MIXED_FMA_{rounding}{saturation}_f32_{input_format}") - && selection.asm - == format!( - "fma.{rounding}{}.f32.{input_format} \t$dst, $a, $b, $c;", - if arithmetic.saturation == ScalarArithmeticSaturation::Sat { - ".sat" - } else { - "" - } - ) - && selection.predicates - == [ - "Subtarget->getSmVersion() >= 100", - "Subtarget->getPTXVersion() >= 86", - ] - && selection.constraints.is_empty(), - "{} mixed-input scalar-arithmetic selection changed", - policy.id - ); - } - } else { - ensure!( - declaration.selections.len() == 1, - "{} gained an unreviewed scalar-arithmetic selection", - policy.id - ); - } - let mut selected = Vec::new(); - for selection in &declaration.selections { - if selection_matches_policy(policy, selection)? { - selected.push(selection); - } - } - ensure!( - selected.len() == 1 && selected[0].source_record == recipe.selection_record, - "{} must select only its direct scalar arithmetic instruction", - policy.id - ); - ensure!( - policy.rust_module == "float" - && policy.rust_name == recipe.id - && policy.rust_arguments == signature - && policy.rust_result == recipe.rust_type - && policy.safe - && policy.must_use - && policy.compatibility_rust_paths == [format!("cuda_device::float::{}", recipe.id)] - && policy.dialect_op_type == "ScalarArithmeticOp" - && policy.dialect_op_name == "nvvm.scalar_arithmetic" - && policy.dialect_operands == signature - && policy.dialect_results == [recipe.rust_type] - && policy.lowering == "generated_scalar_arithmetic", - "{} changed its scalar-arithmetic API, carrier, or lowering", - policy.id - ); - ensure!( - policy.pure == (arithmetic.operation != ScalarArithmeticOperation::Div) - && policy.memory == "none" - && !policy.convergent - && policy.execution_scope == "thread" - && policy.minimum_ptx == "7.0" - && policy.minimum_sm.as_deref() == Some("sm_80") - && policy.ptx_result == recipe.rust_type - && policy.targets == "all" - && policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section == recipe.ptx_isa_section - && policy.ptx_isa_url == recipe.ptx_isa_url, - "{} scalar-arithmetic effects, provenance, or target floor changed", - policy.id - ); - ensure!( - policy.expected_ptx - == InstructionPattern { - mnemonic: scalar_arithmetic_operation_name(arithmetic.operation).into(), - modifiers: recipe.ptx_modifiers.clone(), - operands: vec![OperandPattern::Register; recipe.argument_count + 1], - }, - "{} expected scalar-arithmetic PTX changed", - policy.id - ); - let llvm_mechanism = if arithmetic.saturation == ScalarArithmeticSaturation::Sat { - BackendLoweringMechanism::InlinePtx - } else { - BackendLoweringMechanism::TypedNvvm - }; - let expected_backends = [ - (IntrinsicBackend::LlvmNvptx, llvm_mechanism), - ( - IntrinsicBackend::LibNvvm, - BackendLoweringMechanism::InlinePtx, - ), - ]; - ensure!( - policy.backend_lowerings.len() == 2 - && expected_backends.into_iter().all(|(backend, mechanism)| { - policy.backend_lowerings.iter().any(|lowering| { - lowering.backend == backend - && lowering.mechanism == mechanism - && lowering.minimum_ptx.as_deref() == Some("7.0") - && lowering.minimum_sm.as_deref() == Some("sm_80") - && !lowering.evidence_profile.trim().is_empty() - }) - }), - "{} has the wrong reviewed scalar-arithmetic backend routes", - policy.id - ); - validate_selected_target_predicates(policy, direct)?; - ensure_no_other_family_contract(policy, "scalar arithmetic")?; - Ok(()) -} - -type ScalarMathVariant = ( - ScalarMathFormat, - ScalarMathOperation, - ScalarMathPrecision, - ScalarMathSubnormal, -); - -#[derive(Clone, PartialEq, Eq)] -enum ScalarMathRecipeSource { - /// Bound to a monomorphic tblgen record in the pinned import. - Imported { - source_record: String, - llvm_symbol: String, - }, - /// Bound to an overloaded (polymorphic) tblgen record: the import - /// carries `anonymous_14`/`anyfloat` signature tokens and the concrete - /// f32 instantiation is recorded as the resolved symbol. Only ex2 uses - /// this (LLVM 22 models it as `int_nvvm_ex2_approx{,_ftz}` without a - /// per-format record). - ImportedOverloaded { - source_record: String, - llvm_symbol: String, - resolved_llvm_symbol: String, - }, - /// No record exists in the pinned tblgen import at all. Only tanh uses - /// this: llc selects `llvm.nvvm.tanh.approx.f32` via NVVMIntrinsic-class - /// matching, but the import exports no record for it, so the op is - /// admitted directly against the PTX instruction. - PtxNative { instruction: String }, -} - -struct ScalarMathRecipe { - id: String, - operation_key: String, - source: ScalarMathRecipeSource, - rust_type: &'static str, - dialect_type: &'static str, - minimum_ptx: &'static str, - minimum_sm: &'static str, - properties: Vec<&'static str>, - ptx_modifiers: Vec, - ptx_isa_section: &'static str, - ptx_isa_url: &'static str, - /// The generator's tblgen import found no DAG selection pattern for the - /// intrinsic, so the lowering routes it through inline PTX. Note this is - /// an import limitation, not an llc one: llc still selects these - /// intrinsics through NVVMIntrinsic-class pattern matching, which the - /// evidence import cannot see. Promoting them to typed calls once the - /// import understands those patterns would reopen them to LLVM - /// optimization. - force_inline_ptx: bool, -} - -fn scalar_math_format_name(format: ScalarMathFormat) -> &'static str { - match format { - ScalarMathFormat::F16 => "f16", - ScalarMathFormat::F32 => "f32", - ScalarMathFormat::F64 => "f64", - } -} - -fn scalar_math_operation_name(operation: ScalarMathOperation) -> &'static str { - match operation { - ScalarMathOperation::Sin => "sin", - ScalarMathOperation::Cos => "cos", - ScalarMathOperation::Ex2 => "ex2", - ScalarMathOperation::Lg2 => "lg2", - ScalarMathOperation::Rcp => "rcp", - ScalarMathOperation::Rsqrt => "rsqrt", - ScalarMathOperation::Sqrt => "sqrt", - ScalarMathOperation::Tanh => "tanh", - } -} - -fn scalar_math_precision_name(precision: ScalarMathPrecision) -> &'static str { - match precision { - ScalarMathPrecision::Approx => "approx", - ScalarMathPrecision::Rn => "rn", - ScalarMathPrecision::Rz => "rz", - ScalarMathPrecision::Rm => "rm", - ScalarMathPrecision::Rp => "rp", - } -} - -fn canonical_scalar_math_variants() -> Vec { - use ScalarMathFormat::{F16, F32, F64}; - use ScalarMathOperation::{Cos, Ex2, Lg2, Rcp, Rsqrt, Sin, Sqrt, Tanh}; - use ScalarMathPrecision::{Approx, Rm, Rn, Rp, Rz}; - use ScalarMathSubnormal::{Ftz, Preserve}; - - vec![ - // sin: approx f32 only - (F32, Sin, Approx, Preserve), - (F32, Sin, Approx, Ftz), - // cos: approx f32 only - (F32, Cos, Approx, Preserve), - (F32, Cos, Approx, Ftz), - // lg2: approx f32 only (f64 approx produces invalid PTX) - (F32, Lg2, Approx, Preserve), - (F32, Lg2, Approx, Ftz), - // rcp: approx only with ftz, rounded for both formats - (F32, Rcp, Approx, Ftz), - (F32, Rcp, Rn, Preserve), - (F32, Rcp, Rn, Ftz), - (F32, Rcp, Rz, Preserve), - (F32, Rcp, Rz, Ftz), - (F32, Rcp, Rm, Preserve), - (F32, Rcp, Rm, Ftz), - (F32, Rcp, Rp, Preserve), - (F32, Rcp, Rp, Ftz), - (F64, Rcp, Approx, Ftz), - (F64, Rcp, Rn, Preserve), - (F64, Rcp, Rz, Preserve), - (F64, Rcp, Rm, Preserve), - (F64, Rcp, Rp, Preserve), - // rsqrt: approx only (PTX has no rounded rsqrt). The f64+ftz variant - // is valid PTX (`rsqrt.approx.ftz.f64` assembles for sm_80+, and LLVM - // selects `llvm.nvvm.rsqrt.approx.ftz.d` directly); it is deferred - // only because it has not been probed under the pinned evidence - // profile yet, not because the instruction is invalid. - (F32, Rsqrt, Approx, Preserve), - (F32, Rsqrt, Approx, Ftz), - (F64, Rsqrt, Approx, Preserve), - // sqrt: approx f32 only, rounded for both formats - (F32, Sqrt, Approx, Preserve), - (F32, Sqrt, Approx, Ftz), - (F32, Sqrt, Rn, Preserve), - (F32, Sqrt, Rn, Ftz), - (F32, Sqrt, Rz, Preserve), - (F32, Sqrt, Rz, Ftz), - (F32, Sqrt, Rm, Preserve), - (F32, Sqrt, Rm, Ftz), - (F32, Sqrt, Rp, Preserve), - (F32, Sqrt, Rp, Ftz), - (F64, Sqrt, Rn, Preserve), - (F64, Sqrt, Rz, Preserve), - (F64, Sqrt, Rm, Preserve), - (F64, Sqrt, Rp, Preserve), - // ex2: approx f32 only (PTX has no ex2.approx.f64). Appended after - // the original 37 so existing ABI ids stay stable. - (F32, Ex2, Approx, Preserve), - (F32, Ex2, Approx, Ftz), - // tanh: approx f32 only; the instruction has no ftz form (PTX ISA - // Table 29) and no rounded variants. Hardware floor is sm_75; the - // family contract gates it at the attested sm_80 evidence floor. - (F32, Tanh, Approx, Preserve), - // ex2.approx.f16 starts at PTX 7.0 / sm_75. LLVM 22 rejects the - // .ftz.f16 spelling; .ftz.bf16 is a separate PTX 7.8 / sm_90 variant. - (F16, Ex2, Approx, Preserve), - ] -} - -fn scalar_math_recipe(variant: ScalarMathVariant) -> Option { - if !canonical_scalar_math_variants().contains(&variant) { - return None; - } - let (format, operation, precision, subnormal) = variant; - let operation_name = scalar_math_operation_name(operation); - let precision_name = scalar_math_precision_name(precision); - let format_name = scalar_math_format_name(format); - let source_format = match format { - ScalarMathFormat::F16 => "f16", - ScalarMathFormat::F32 => "f", - ScalarMathFormat::F64 => "d", - }; - - let mut modifier_names = vec![precision_name]; - if subnormal == ScalarMathSubnormal::Ftz { - modifier_names.push("ftz"); - } - - let modifier_id = modifier_names.join("_"); - let modifier_symbol = modifier_names.join("."); - let id = format!("{operation_name}_{modifier_id}_{format_name}"); - let source = match operation { - // LLVM 22 models ex2 as overloaded records without a per-format - // suffix (`int_nvvm_ex2_approx{,_ftz}` over anyfloat); bind those - // for declaration-level checks and record the concrete f32 - // instantiation as the resolved symbol. - ScalarMathOperation::Ex2 => ScalarMathRecipeSource::ImportedOverloaded { - source_record: format!("int_nvvm_{operation_name}_{modifier_id}"), - llvm_symbol: format!("llvm.nvvm.{operation_name}.{modifier_symbol}"), - resolved_llvm_symbol: format!( - "llvm.nvvm.{operation_name}.{modifier_symbol}.{format_name}" - ), - }, - // LLVM 22.1.2's tblgen export has no record for tanh at all, so the - // op is admitted directly against the PTX instruction. - ScalarMathOperation::Tanh => ScalarMathRecipeSource::PtxNative { - instruction: format!("{operation_name}.{modifier_symbol}.{format_name}"), - }, - _ => ScalarMathRecipeSource::Imported { - source_record: format!("int_nvvm_{operation_name}_{modifier_id}_{source_format}"), - llvm_symbol: format!("llvm.nvvm.{operation_name}.{modifier_symbol}.{source_format}"), - }, - }; - let ptx_modifiers = modifier_names - .iter() - .copied() - .chain(std::iter::once(format_name)) - .map(str::to_owned) - .collect::>(); - - // Operations whose llvm.nvvm.* records carry no *imported* DAG selection - // pattern in the pinned evidence (sel=0). llc itself still selects them - // (via NVVMIntrinsic-class patterns the tblgen import cannot see), but - // the evidence-driven contract only admits typed calls backed by an - // imported selection, so these route through inline PTX for now. - let force_inline_ptx = matches!( - operation, - ScalarMathOperation::Sin - | ScalarMathOperation::Cos - | ScalarMathOperation::Ex2 - | ScalarMathOperation::Lg2 - | ScalarMathOperation::Rsqrt - | ScalarMathOperation::Tanh - ); - - let (ptx_isa_section, ptx_isa_url) = match (operation, format) { - (ScalarMathOperation::Ex2, ScalarMathFormat::F16) => ( - "9.7.4.10 Half Precision Floating Point Instructions: ex2", - "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-ex2", - ), - (ScalarMathOperation::Sin, _) => ( - "9.7.3.10 Floating Point Instructions: sin", - "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-sin", - ), - (ScalarMathOperation::Cos, _) => ( - "9.7.3.11 Floating Point Instructions: cos", - "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-cos", - ), - (ScalarMathOperation::Lg2, _) => ( - "9.7.3.12 Floating Point Instructions: lg2", - "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-lg2", - ), - (ScalarMathOperation::Rcp, _) => ( - "9.7.3.7 Floating Point Instructions: rcp", - "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-rcp", - ), - (ScalarMathOperation::Rsqrt, _) => ( - "9.7.3.14 Floating Point Instructions: rsqrt", - "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-rsqrt", - ), - (ScalarMathOperation::Sqrt, _) => ( - "9.7.3.9 Floating Point Instructions: sqrt", - "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-sqrt", - ), - (ScalarMathOperation::Ex2, _) => ( - "9.7.3.13 Floating Point Instructions: ex2", - "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-ex2", - ), - (ScalarMathOperation::Tanh, _) => ( - "9.7.3.15 Floating Point Instructions: tanh", - "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-tanh", - ), - }; - - let (rust_type, dialect_type, minimum_ptx, minimum_sm) = match format { - ScalarMathFormat::F16 => ("u16", "i16", "7.0", "sm_75"), - ScalarMathFormat::F32 => ("f32", "f32", "7.0", "sm_80"), - ScalarMathFormat::F64 => ("f64", "f64", "7.0", "sm_80"), - }; - Some(ScalarMathRecipe { - id, - operation_key: format!("scalar.math.{operation_name}.{modifier_symbol}.{format_name}"), - source, - rust_type, - dialect_type, - minimum_ptx, - minimum_sm, - properties: vec!["IntrNoMem"], - ptx_modifiers, - ptx_isa_section, - ptx_isa_url, - force_inline_ptx, - }) -} - -fn expand_scalar_math_admission(admission: &ScalarMathAdmission) -> Result> { - ensure!( - admission.runtime_validation == RuntimeValidation::Unexecuted, - "scalar-math runtime may be marked executed only with GPU evidence" - ); - let expected = canonical_scalar_math_variants(); - let actual = admission - .variants - .iter() - .map(|variant| { - ( - variant.format, - variant.operation, - variant.precision, - variant.subnormal, - ) - }) - .collect::>(); - ensure!( - actual == expected, - "compact scalar-math admission must list the canonical 41 variants" - ); - - admission - .variants - .iter() - .map(|variant| { - validate_abi_id(&variant.abi_id)?; - let identity = ( - variant.format, - variant.operation, - variant.precision, - variant.subnormal, - ); - let recipe = scalar_math_recipe(identity) - .context("scalar math is outside the closed recipe set")?; - let libnvvm_evidence_profile = variant - .libnvvm_evidence_profile - .as_ref() - .unwrap_or(&admission.libnvvm_evidence_profile); - ensure!( - !libnvvm_evidence_profile.trim().is_empty(), - "scalar-math libNVVM evidence profile must not be empty" - ); - scalar_math_overlay_record( - recipe, - admission, - identity, - &variant.abi_id, - libnvvm_evidence_profile, - ) - }) - .collect() -} - -fn scalar_math_overlay_record( - recipe: ScalarMathRecipe, - admission: &ScalarMathAdmission, - variant: ScalarMathVariant, - abi_id: &str, - libnvvm_evidence_profile: &str, -) -> Result { - let (format, operation, precision, subnormal) = variant; - let ptx_operands = vec![OperandPattern::Register; 2]; // 1 result + 1 operand - let summary = format!( - "Computes unary {} with {} precision.", - scalar_math_operation_name(operation), - scalar_math_precision_name(precision), - ); - let (source, source_record, llvm_symbol, resolved_llvm_symbol, llvm_arguments, llvm_results) = - match &recipe.source { - ScalarMathRecipeSource::Imported { - source_record, - llvm_symbol, - } => ( - None, - Some(source_record.clone()), - Some(llvm_symbol.clone()), - None, - vec![recipe.rust_type.to_owned()], - vec![recipe.rust_type.to_owned()], - ), - // The polymorphic signature tokens mirror the imported - // overloaded record verbatim (anyfloat over anonymous_14), the - // same shape packed_alu uses for llvm.nvvm.fabs. - ScalarMathRecipeSource::ImportedOverloaded { - source_record, - llvm_symbol, - resolved_llvm_symbol, - } => ( - None, - Some(source_record.clone()), - Some(llvm_symbol.clone()), - Some(resolved_llvm_symbol.clone()), - vec!["anonymous_14".to_owned()], - vec!["anyfloat".to_owned()], - ), - ScalarMathRecipeSource::PtxNative { instruction } => ( - Some(IntrinsicSource::PtxNative { - instruction: instruction.clone(), - }), - None, - None, - None, - Vec::new(), - Vec::new(), - ), - }; - Ok(OverlayIntrinsic { - id: recipe.id.clone(), - abi_id: abi_id.into(), - operation_key: recipe.operation_key, - family: "scalar_math".into(), - source, - source_record, - rust_module: "float".into(), - rust_name: recipe.id.clone(), - rust_arguments: vec![recipe.rust_type.into()], - rust_result: recipe.rust_type.into(), - safe: true, - must_use: true, - safe_allowlist_reason: Some("Scalar math has no caller obligations.".into()), - public_rust_path: format!("cuda_intrinsics::float::{}", recipe.id), - compatibility_rust_paths: vec![format!("cuda_device::float::{}", recipe.id)], - dialect_op_type: "ScalarMathOp".into(), - dialect_op_name: "nvvm.scalar_math".into(), - dialect_operands: vec![recipe.dialect_type.into()], - dialect_results: vec![recipe.dialect_type.into()], - llvm_symbol, - resolved_llvm_symbol, - llvm_arguments, - llvm_results, - pure: true, - memory: "none".into(), - convergent: false, - execution_scope: "thread".into(), - minimum_ptx: recipe.minimum_ptx.into(), - minimum_sm: Some(recipe.minimum_sm.into()), - ptx_result: recipe.rust_type.into(), - targets: "all".into(), - ptx_isa_version: "9.3".into(), - ptx_isa_section: recipe.ptx_isa_section.into(), - ptx_isa_url: recipe.ptx_isa_url.into(), - lowering: "generated_scalar_math".into(), - backend_lowerings: [ - ( - IntrinsicBackend::LlvmNvptx, - admission.llvm_evidence_profile.as_str(), - ), - (IntrinsicBackend::LibNvvm, libnvvm_evidence_profile), - ] - .into_iter() - .map( - |(backend, evidence_profile): (IntrinsicBackend, &str)| OverlayBackendLowering { - backend, - mechanism: match backend { - IntrinsicBackend::LlvmNvptx if recipe.force_inline_ptx => { - BackendLoweringMechanism::InlinePtx - } - IntrinsicBackend::LlvmNvptx => BackendLoweringMechanism::TypedNvvm, - IntrinsicBackend::LibNvvm => BackendLoweringMechanism::InlinePtx, - }, - evidence_profile: evidence_profile.to_owned(), - targets: None, - minimum_ptx: Some(recipe.minimum_ptx.into()), - minimum_sm: Some(recipe.minimum_sm.into()), - }, - ) - .collect(), - packed_atomic: None, - redux: None, - vote: None, - active_mask: None, - warp_match: None, - warp_barrier: None, - warp_shuffle: None, - dot_product: None, - packed_alu: None, - integer_minmax: None, - packed_conversion: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: Some(ScalarMath { - format, - operation, - precision, - subnormal, - runtime_validation: admission.runtime_validation, - }), - extended_minmax: None, - cp_async_copy: None, - cp_async_control: None, - cp_async_mbarrier: None, - mbarrier_basic: None, - movmatrix: None, - mbarrier_extended: None, - register_mma: None, - sparse_mma: None, - prmt: None, - cluster_barrier: None, - wgmma_control: None, - special_register: None, - debug_control: None, - cluster_memory: None, - clc: None, - tma: None, - tcgen05: None, - ldmatrix_variant: None, - ldmatrix_safety: None, - ldmatrix_adapter: None, - selected_address_space: None, - expected_ptx: InstructionPattern { - mnemonic: scalar_math_operation_name(operation).into(), - modifiers: recipe.ptx_modifiers, - operands: ptx_operands, - }, - summary, - }) -} - -fn validate_scalar_math_policy( - policy: &OverlayIntrinsic, - declaration: Option<&ImportedIntrinsic>, -) -> Result<()> { - let math = policy - .scalar_math - .as_ref() - .with_context(|| format!("{} has no scalar-math contract", policy.id))?; - let variant = (math.format, math.operation, math.precision, math.subnormal); - let recipe = scalar_math_recipe(variant) - .with_context(|| format!("{} is outside the closed scalar-math recipe", policy.id))?; - ensure!( - math.runtime_validation == RuntimeValidation::Unexecuted, - "{} scalar-math runtime may be executed only with GPU evidence", - policy.id - ); - let signature = vec![recipe.rust_type.to_owned()]; - let source_matches = match &recipe.source { - ScalarMathRecipeSource::Imported { - source_record, - llvm_symbol, - } => { - policy.source.is_none() - && policy.source_record.as_deref() == Some(source_record.as_str()) - && policy.llvm_symbol.as_deref() == Some(llvm_symbol.as_str()) - && policy.resolved_llvm_symbol.is_none() - && policy.llvm_arguments == signature - && policy.llvm_results == [recipe.rust_type] - } - ScalarMathRecipeSource::ImportedOverloaded { - source_record, - llvm_symbol, - resolved_llvm_symbol, - } => { - policy.source.is_none() - && policy.source_record.as_deref() == Some(source_record.as_str()) - && policy.llvm_symbol.as_deref() == Some(llvm_symbol.as_str()) - && policy.resolved_llvm_symbol.as_deref() == Some(resolved_llvm_symbol.as_str()) - && policy.llvm_arguments == ["anonymous_14"] - && policy.llvm_results == ["anyfloat"] - } - ScalarMathRecipeSource::PtxNative { instruction } => { - policy.source - == Some(IntrinsicSource::PtxNative { - instruction: instruction.clone(), - }) - && policy.source_record.is_none() - && policy.llvm_symbol.is_none() - && policy.resolved_llvm_symbol.is_none() - && policy.llvm_arguments.is_empty() - && policy.llvm_results.is_empty() - } - }; - ensure!( - policy.id == recipe.id && policy.operation_key == recipe.operation_key && source_matches, - "{} scalar-math identity or LLVM source changed", - policy.id - ); - ensure!( - declaration.is_none() == matches!(recipe.source, ScalarMathRecipeSource::PtxNative { .. }), - "{} scalar-math source kind and imported declaration disagree", - policy.id - ); - if let Some(declaration) = declaration { - let expected_properties = recipe - .properties - .iter() - .map(|property| (*property).to_owned()) - .collect::>(); - ensure!( - declaration.properties == expected_properties, - "{} imported scalar-math properties changed", - policy.id - ); - ensure!( - declaration.selections.len() <= 1, - "{} gained an unreviewed scalar-math selection", - policy.id - ); - if let Some(direct) = declaration.selections.first() { - ensure!( - direct.predicates.is_empty() && direct.constraints.is_empty(), - "{} direct scalar-math selection changed", - policy.id - ); - } - } - ensure!( - policy.rust_module == "float" - && policy.rust_name == recipe.id - && policy.rust_arguments == signature - && policy.rust_result == recipe.rust_type - && policy.safe - && policy.must_use - && policy.compatibility_rust_paths == [format!("cuda_device::float::{}", recipe.id)] - && policy.dialect_op_type == "ScalarMathOp" - && policy.dialect_op_name == "nvvm.scalar_math" - && policy.dialect_operands == [recipe.dialect_type] - && policy.dialect_results == [recipe.dialect_type] - && policy.lowering == "generated_scalar_math", - "{} changed its scalar-math API, carrier, or lowering", - policy.id - ); - ensure!( - policy.pure - && policy.memory == "none" - && !policy.convergent - && policy.execution_scope == "thread" - && policy.minimum_ptx == recipe.minimum_ptx - && policy.minimum_sm.as_deref() == Some(recipe.minimum_sm) - && policy.ptx_result == recipe.rust_type - && policy.targets == "all" - && policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section == recipe.ptx_isa_section - && policy.ptx_isa_url == recipe.ptx_isa_url, - "{} scalar-math effects, provenance, or target floor changed", - policy.id - ); - ensure!( - policy.expected_ptx - == InstructionPattern { - mnemonic: scalar_math_operation_name(math.operation).into(), - modifiers: recipe.ptx_modifiers.clone(), - operands: vec![OperandPattern::Register; 2], - }, - "{} expected scalar-math PTX changed", - policy.id - ); - let llvm_mechanism = if recipe.force_inline_ptx { - BackendLoweringMechanism::InlinePtx - } else { - BackendLoweringMechanism::TypedNvvm - }; - let expected_backends = [ - ( - IntrinsicBackend::LlvmNvptx, - llvm_mechanism, - recipe.minimum_sm, - ), - ( - IntrinsicBackend::LibNvvm, - BackendLoweringMechanism::InlinePtx, - recipe.minimum_sm, - ), - ]; - ensure!( - policy.backend_lowerings.len() == 2 - && expected_backends - .into_iter() - .all(|(backend, mechanism, minimum_sm)| { - policy.backend_lowerings.iter().any(|lowering| { - lowering.backend == backend - && lowering.mechanism == mechanism - && lowering.minimum_ptx.as_deref() == Some(recipe.minimum_ptx) - && lowering.minimum_sm.as_deref() == Some(minimum_sm) - && !lowering.evidence_profile.trim().is_empty() - }) - }), - "{} has the wrong reviewed scalar-math backend routes (expected {} / {})", - policy.id, - recipe.minimum_ptx, - recipe.minimum_sm - ); - if let Some(direct) = declaration.and_then(|declaration| declaration.selections.first()) { - validate_selected_target_predicates(policy, direct)?; - } - ensure_no_other_family_contract(policy, "scalar math")?; - Ok(()) -} - -type ExtendedMinMaxVariant = ( - ExtendedMinMaxFormat, - ExtendedMinMaxOperation, - ExtendedMinMaxSubnormal, - ExtendedMinMaxNan, - bool, -); - -struct ExtendedMinMaxRecipe { - id: String, - operation_key: String, - source_record: String, - llvm_symbol: String, - llvm_type: &'static str, - rust_module: &'static str, - rust_type: &'static str, - dialect_type: &'static str, - adapter: ExtendedMinMaxAdapter, - selection_record: String, - selection_asm: String, - predicates: Vec, - classes: Vec<&'static str>, - modifiers: Vec, - minimum_ptx: &'static str, - minimum_sm: &'static str, - ptx_isa_section: &'static str, - ptx_isa_url: &'static str, -} - -fn canonical_extended_minmax_variants() -> Vec { - use ExtendedMinMaxFormat::{Bf16, Bf16x2, F16, F16x2, F32}; - use ExtendedMinMaxNan::{Nan, Number}; - use ExtendedMinMaxOperation::{Max, Min}; - use ExtendedMinMaxSubnormal::{Ftz, Preserve}; - - let one_operation = |operation| { - [ - (F16x2, operation, Ftz, Number, false), - (F16x2, operation, Ftz, Nan, false), - (F32, operation, Ftz, Nan, true), - (F16x2, operation, Ftz, Nan, true), - (F32, operation, Ftz, Number, true), - (F16x2, operation, Ftz, Number, true), - (Bf16x2, operation, Preserve, Nan, false), - (F16x2, operation, Preserve, Nan, false), - (Bf16x2, operation, Preserve, Nan, true), - (F32, operation, Preserve, Nan, true), - (F16x2, operation, Preserve, Nan, true), - (Bf16x2, operation, Preserve, Number, true), - (F32, operation, Preserve, Number, true), - (F16x2, operation, Preserve, Number, true), - ] - }; - // The scalar 16-bit forms. LLVM declares every `f16` modifier combination - // and the four `bf16` combinations that do not request `ftz`; the four - // `bf16` `ftz` declarations exist but have no NVPTX selection pattern and - // fail instruction selection, so they are deliberately absent here. - // - // ABI identity is intentionally absent from this canonical list. Existing - // and future IDs are carried by admission entries and enforced by the - // global append-only ABI ledger. - let scalar_halves = |operation| { - [ - (F16, operation, Preserve, Number, false), - (F16, operation, Ftz, Number, false), - (F16, operation, Preserve, Nan, false), - (F16, operation, Ftz, Nan, false), - (F16, operation, Preserve, Number, true), - (F16, operation, Ftz, Number, true), - (F16, operation, Preserve, Nan, true), - (F16, operation, Ftz, Nan, true), - (Bf16, operation, Preserve, Number, false), - (Bf16, operation, Preserve, Nan, false), - (Bf16, operation, Preserve, Number, true), - (Bf16, operation, Preserve, Nan, true), - ] - }; - one_operation(Min) - .into_iter() - .chain(one_operation(Max)) - .chain(scalar_halves(Min)) - .chain(scalar_halves(Max)) - .collect() -} - -/// Joins an operation name, its modifiers, and its format into one identifier. -/// -/// `min.f16` and `min.bf16` carry no modifier at all, so the separators have to -/// come from joining the parts that are present rather than from a format -/// string with fixed separators. -fn extended_minmax_joined_name( - leading: &str, - trailing: &str, - modifiers: &[&str], - separator: &str, -) -> String { - let mut parts = vec![leading]; - parts.extend(modifiers.iter().copied()); - parts.push(trailing); - parts.join(separator) -} - -fn extended_minmax_format_name(format: ExtendedMinMaxFormat) -> &'static str { - match format { - ExtendedMinMaxFormat::F32 => "f32", - ExtendedMinMaxFormat::F16 => "f16", - ExtendedMinMaxFormat::Bf16 => "bf16", - ExtendedMinMaxFormat::F16x2 => "f16x2", - ExtendedMinMaxFormat::Bf16x2 => "bf16x2", - } -} - -fn extended_minmax_operation_name(operation: ExtendedMinMaxOperation) -> &'static str { - match operation { - ExtendedMinMaxOperation::Min => "min", - ExtendedMinMaxOperation::Max => "max", - } -} - -fn extended_minmax_recipe(variant: ExtendedMinMaxVariant) -> Option { - if !canonical_extended_minmax_variants().contains(&variant) { - return None; - } - let (format, operation, subnormal, nan, xorsign_abs) = variant; - let operation_name = extended_minmax_operation_name(operation); - let format_name = extended_minmax_format_name(format); - let source_type = if format == ExtendedMinMaxFormat::F32 { - "f" - } else { - format_name - }; - let mut source_modifiers = Vec::new(); - let mut ptx_modifiers = Vec::new(); - if subnormal == ExtendedMinMaxSubnormal::Ftz { - source_modifiers.push("ftz"); - ptx_modifiers.push("ftz".to_owned()); - } - if nan == ExtendedMinMaxNan::Nan { - source_modifiers.push("nan"); - ptx_modifiers.push("NaN".to_owned()); - } - if xorsign_abs { - source_modifiers.extend(["xorsign", "abs"]); - ptx_modifiers.extend(["xorsign".to_owned(), "abs".to_owned()]); - } - let id = extended_minmax_joined_name(operation_name, format_name, &source_modifiers, "_"); - let intrinsic_name = format!("f{operation_name}"); - let source_record = format!( - "int_nvvm_{}", - extended_minmax_joined_name(&intrinsic_name, source_type, &source_modifiers, "_") - ); - let llvm_symbol = format!( - "llvm.nvvm.{}", - extended_minmax_joined_name(&intrinsic_name, source_type, &source_modifiers, ".") - ); - let selection_record = if format == ExtendedMinMaxFormat::F32 { - format!( - "INT_NVVM_{}", - source_record - .trim_start_matches("int_nvvm_") - .to_ascii_uppercase() - ) - } else { - let prefix = match operation { - ExtendedMinMaxOperation::Min => "FMIN", - ExtendedMinMaxOperation::Max => "FMAN", - }; - let selection_modifiers = source_modifiers - .iter() - .map(|modifier| if *modifier == "nan" { "NaN" } else { modifier }) - .collect::>(); - format!( - "INT_NVVM_{}", - extended_minmax_joined_name(prefix, format_name, &selection_modifiers, "_") - ) - }; - let ptx_format = if format == ExtendedMinMaxFormat::F32 { - "f32" - } else { - format_name - }; - ptx_modifiers.push(ptx_format.to_owned()); - let selection_asm = format!( - "{operation_name}.{} \t$dst, $src0, $src1;", - ptx_modifiers.join(".") - ); - let (minimum_ptx, minimum_sm, predicates) = if xorsign_abs { - ( - "7.2", - "sm_86", - vec![ - "Subtarget->getPTXVersion() >= 72".to_owned(), - "Subtarget->getSmVersion() >= 86".to_owned(), - ], - ) - } else { - ( - "7.0", - "sm_80", - vec![ - "Subtarget->getSmVersion() >= 80".to_owned(), - "Subtarget->getPTXVersion() >= 70".to_owned(), - ], - ) - }; - let (llvm_type, rust_module, rust_type, dialect_type, adapter, classes) = match format { - ExtendedMinMaxFormat::F32 => ( - "f32", - "float", - "f32", - "f32", - ExtendedMinMaxAdapter::DirectF32, - vec![ - "ClangBuiltin", - "NVVMBuiltin", - "SDPatternOperator", - "Intrinsic", - "DefaultAttrsIntrinsic", - ], - ), - ExtendedMinMaxFormat::F16 => ( - "f16", - "f16", - "u16", - "i16", - ExtendedMinMaxAdapter::DirectHalfU16, - vec!["SDPatternOperator", "Intrinsic", "DefaultAttrsIntrinsic"], - ), - ExtendedMinMaxFormat::Bf16 => ( - "bf16", - "bf16", - "u16", - "i16", - ExtendedMinMaxAdapter::DirectHalfU16, - vec![ - "ClangBuiltin", - "NVVMBuiltin", - "SDPatternOperator", - "Intrinsic", - "DefaultAttrsIntrinsic", - ], - ), - ExtendedMinMaxFormat::F16x2 => ( - "v2f16", - "f16x2", - "u32", - "i32", - ExtendedMinMaxAdapter::DirectPackedU32, - vec!["SDPatternOperator", "Intrinsic", "DefaultAttrsIntrinsic"], - ), - ExtendedMinMaxFormat::Bf16x2 => ( - "v2bf16", - "bf16x2", - "u32", - "i32", - ExtendedMinMaxAdapter::DirectPackedU32, - vec![ - "ClangBuiltin", - "NVVMBuiltin", - "SDPatternOperator", - "Intrinsic", - "DefaultAttrsIntrinsic", - ], - ), - }; - let (ptx_isa_section, ptx_isa_url) = match (format, operation) { - (ExtendedMinMaxFormat::F32, ExtendedMinMaxOperation::Min) => ( - "9.7.3.9 Floating Point Instructions: min", - "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-min", - ), - (ExtendedMinMaxFormat::F32, ExtendedMinMaxOperation::Max) => ( - "9.7.3.10 Floating Point Instructions: max", - "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-max", - ), - (_, ExtendedMinMaxOperation::Min) => ( - "9.7.4.7 Half Precision Floating Point Instructions: min", - "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-min", - ), - (_, ExtendedMinMaxOperation::Max) => ( - "9.7.4.8 Half Precision Floating Point Instructions: max", - "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-max", - ), - }; - Some(ExtendedMinMaxRecipe { - id, - operation_key: format!( - "floating.minmax.{format_name}.{operation_name}.{}.{}.{}", - match subnormal { - ExtendedMinMaxSubnormal::Preserve => "preserve", - ExtendedMinMaxSubnormal::Ftz => "ftz", - }, - match nan { - ExtendedMinMaxNan::Number => "number", - ExtendedMinMaxNan::Nan => "nan", - }, - if xorsign_abs { "xorsign_abs" } else { "direct" } - ), - source_record, - llvm_symbol, - llvm_type, - rust_module, - rust_type, - dialect_type, - adapter, - selection_record, - selection_asm, - predicates, - classes, - modifiers: ptx_modifiers, - minimum_ptx, - minimum_sm, - ptx_isa_section, - ptx_isa_url, - }) -} - -fn expand_extended_minmax_admission( - admission: &ExtendedMinMaxAdmission, -) -> Result> { - ensure!( - admission.runtime_validation == RuntimeValidation::Unexecuted, - "extended-minmax runtime may be marked executed only with GPU evidence" - ); - let expected = canonical_extended_minmax_variants(); - let actual = admission - .variants - .iter() - .map(|variant| { - ( - variant.format, - variant.operation, - variant.subnormal, - variant.nan, - variant.xorsign_abs, - ) - }) - .collect::>(); - ensure!( - actual == expected, - "compact extended-minmax admission must list the exact canonical 28 variants" - ); - admission - .variants - .iter() - .map(|variant| { - validate_abi_id(&variant.abi_id)?; - let identity = ( - variant.format, - variant.operation, - variant.subnormal, - variant.nan, - variant.xorsign_abs, - ); - let recipe = extended_minmax_recipe(identity) - .context("extended min/max is outside the closed recipe set")?; - extended_minmax_overlay_record(recipe, admission, identity, &variant.abi_id) - }) - .collect() -} - -fn extended_minmax_overlay_record( - recipe: ExtendedMinMaxRecipe, - admission: &ExtendedMinMaxAdmission, - variant: ExtendedMinMaxVariant, - abi_id: &str, -) -> Result { - let (format, operation, subnormal, nan, xorsign_abs) = variant; - let rust_arguments = vec![recipe.rust_type.to_owned(); 2]; - let dialect_operands = vec![recipe.dialect_type.to_owned(); 2]; - Ok(OverlayIntrinsic { - id: recipe.id.clone(), - abi_id: abi_id.into(), - operation_key: recipe.operation_key, - family: "extended_minmax".into(), - source: None, - source_record: Some(recipe.source_record), - rust_module: recipe.rust_module.into(), - rust_name: recipe.id.clone(), - rust_arguments: rust_arguments.clone(), - rust_result: recipe.rust_type.into(), - safe: true, - must_use: true, - safe_allowlist_reason: Some("Floating-point min/max has no caller obligations.".into()), - public_rust_path: format!("cuda_intrinsics::{}::{}", recipe.rust_module, recipe.id), - compatibility_rust_paths: vec![format!( - "cuda_device::{}::{}", - recipe.rust_module, recipe.id - )], - dialect_op_type: "ExtendedMinMaxOp".into(), - dialect_op_name: "nvvm.extended_minmax".into(), - dialect_operands: dialect_operands.clone(), - dialect_results: vec![recipe.dialect_type.into()], - llvm_symbol: Some(recipe.llvm_symbol), - resolved_llvm_symbol: None, - llvm_arguments: vec![recipe.llvm_type.into(); 2], - llvm_results: vec![recipe.llvm_type.into()], - pure: true, - memory: "none".into(), - convergent: false, - execution_scope: "thread".into(), - minimum_ptx: recipe.minimum_ptx.into(), - minimum_sm: Some(recipe.minimum_sm.into()), - ptx_result: recipe.rust_type.into(), - targets: "all".into(), - ptx_isa_version: "9.3".into(), - ptx_isa_section: recipe.ptx_isa_section.into(), - ptx_isa_url: recipe.ptx_isa_url.into(), - lowering: "generated_extended_minmax".into(), - backend_lowerings: [ - ( - IntrinsicBackend::LlvmNvptx, - &admission.llvm_evidence_profile, - ), - ( - IntrinsicBackend::LibNvvm, - &admission.libnvvm_evidence_profile, - ), - ] - .into_iter() - .map(|(backend, evidence_profile)| OverlayBackendLowering { - backend, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: evidence_profile.clone(), - targets: None, - minimum_ptx: Some(recipe.minimum_ptx.into()), - minimum_sm: Some(recipe.minimum_sm.into()), - }) - .collect(), - packed_atomic: None, - redux: None, - vote: None, - active_mask: None, - warp_match: None, - warp_barrier: None, - warp_shuffle: None, - dot_product: None, - packed_alu: None, - integer_minmax: None, - packed_conversion: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: Some(ExtendedMinMax { - format, - operation, - subnormal, - nan, - xorsign_abs, - adapter: recipe.adapter, - runtime_validation: admission.runtime_validation, - }), - cp_async_copy: None, - cp_async_control: None, - cp_async_mbarrier: None, - mbarrier_basic: None, - movmatrix: None, - mbarrier_extended: None, - register_mma: None, - sparse_mma: None, - prmt: None, - cluster_barrier: None, - wgmma_control: None, - special_register: None, - debug_control: None, - cluster_memory: None, - clc: None, - tma: None, - tcgen05: None, - ldmatrix_variant: None, - ldmatrix_safety: None, - ldmatrix_adapter: None, - selected_address_space: None, - expected_ptx: InstructionPattern { - mnemonic: extended_minmax_operation_name(operation).into(), - modifiers: recipe.modifiers, - operands: vec![OperandPattern::Register; 3], - }, - summary: format!( - "Computes extended {} for {} values.", - extended_minmax_operation_name(operation), - extended_minmax_format_name(format) - ), - }) -} - -fn validate_extended_minmax_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, -) -> Result<()> { - let minmax = policy - .extended_minmax - .as_ref() - .with_context(|| format!("{} has no extended-minmax contract", policy.id))?; - let identity = ( - minmax.format, - minmax.operation, - minmax.subnormal, - minmax.nan, - minmax.xorsign_abs, - ); - let recipe = extended_minmax_recipe(identity) - .with_context(|| format!("{} is outside the closed extended-minmax recipe", policy.id))?; - ensure!( - minmax.runtime_validation == RuntimeValidation::Unexecuted, - "{} extended-minmax runtime may be executed only with GPU evidence", - policy.id - ); - let expected_adapter = match minmax.format { - ExtendedMinMaxFormat::F32 => ExtendedMinMaxAdapter::DirectF32, - ExtendedMinMaxFormat::F16 | ExtendedMinMaxFormat::Bf16 => { - ExtendedMinMaxAdapter::DirectHalfU16 - } - ExtendedMinMaxFormat::F16x2 | ExtendedMinMaxFormat::Bf16x2 => { - ExtendedMinMaxAdapter::DirectPackedU32 - } - }; - ensure!( - minmax.adapter == expected_adapter, - "{} extended-minmax adapter does not match its format", - policy.id - ); - ensure!( - policy.id == recipe.id - && policy.operation_key == recipe.operation_key - && policy.source_record.as_deref() == Some(recipe.source_record.as_str()) - && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol.as_str()) - && policy.resolved_llvm_symbol.is_none() - && policy.llvm_arguments == vec![recipe.llvm_type; 2] - && policy.llvm_results == [recipe.llvm_type], - "{} extended-minmax identity or LLVM source changed", - policy.id - ); - ensure!( - declaration.classes == recipe.classes - && declaration.properties == ["Commutative", "IntrNoMem", "IntrSpeculatable"] - && declaration.selections.len() == 1, - "{} imported extended-minmax classes, properties, or selection count changed", - policy.id - ); - let selection = &declaration.selections[0]; - ensure!( - selection.source_record == recipe.selection_record - && selection.asm == recipe.selection_asm - && selection.predicates == recipe.predicates - && selection.constraints.is_empty(), - "{} imported extended-minmax instruction selection changed", - policy.id - ); - let rust_arguments = vec![recipe.rust_type; 2]; - let dialect_operands = vec![recipe.dialect_type; 2]; - ensure!( - policy.rust_module == recipe.rust_module - && policy.rust_name == recipe.id - && policy.rust_arguments == rust_arguments - && policy.rust_result == recipe.rust_type - && policy.safe - && policy.must_use - && policy.public_rust_path - == format!("cuda_intrinsics::{}::{}", recipe.rust_module, recipe.id) - && policy.compatibility_rust_paths - == [format!( - "cuda_device::{}::{}", - recipe.rust_module, recipe.id - )] - && policy.dialect_op_type == "ExtendedMinMaxOp" - && policy.dialect_op_name == "nvvm.extended_minmax" - && policy.dialect_operands == dialect_operands - && policy.dialect_results == [recipe.dialect_type] - && policy.lowering == "generated_extended_minmax", - "{} changed its extended-minmax API, carrier, or lowering", - policy.id - ); - ensure!( - policy.pure - && policy.memory == "none" - && !policy.convergent - && policy.execution_scope == "thread" - && policy.minimum_ptx == recipe.minimum_ptx - && policy.minimum_sm.as_deref() == Some(recipe.minimum_sm) - && policy.ptx_result == recipe.rust_type - && policy.targets == "all" - && policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section == recipe.ptx_isa_section - && policy.ptx_isa_url == recipe.ptx_isa_url, - "{} extended-minmax effects, provenance, or target floor changed", - policy.id - ); - ensure!( - policy.expected_ptx - == InstructionPattern { - mnemonic: extended_minmax_operation_name(minmax.operation).into(), - modifiers: recipe.modifiers, - operands: vec![OperandPattern::Register; 3], - }, - "{} expected extended-minmax PTX changed", - policy.id - ); - ensure_exact_inline_ptx_backends( - policy, - [ - ( - IntrinsicBackend::LlvmNvptx, - recipe.minimum_ptx, - Some(recipe.minimum_sm), - ), - ( - IntrinsicBackend::LibNvvm, - recipe.minimum_ptx, - Some(recipe.minimum_sm), - ), - ], - "extended-minmax", - )?; - validate_selected_target_predicates(policy, selection)?; - ensure_no_other_family_contract(policy, "extended min/max")?; - Ok(()) -} - -struct RegisterMmaRecipe { - id: &'static str, - abi_id: &'static str, - operation_key: &'static str, - source_record: &'static str, - llvm_symbol: &'static str, - rust_arguments: &'static [&'static str], - rust_result: &'static str, - dialect_op_type: &'static str, - dialect_op_name: &'static str, - dialect_operands: &'static [&'static str], - dialect_results: &'static [&'static str], - llvm_arguments: &'static [&'static str], - llvm_results: &'static [&'static str], - adapter: RegisterMmaAdapter, - compatibility_source: RegisterMmaCompatibilitySource, - minimum_ptx: &'static str, - minimum_sm: &'static str, - ptx_modifiers: Vec<&'static str>, - ptx_register_counts: [usize; 4], -} - -fn integer_register_mma_recipe(mma: &RegisterMma, common: bool) -> Option { - use RegisterMmaAdapter::{ - C2I32A1U32B1U32ToD2I32, C4I32A2U32B1U32ToD4I32, C4I32A4U32B2U32ToD4I32, - }; - use RegisterMmaCompatibilitySource::{ExistingStub, GeneratedStub}; - use RegisterMmaElement::{S4, S8, U4, U8}; - use RegisterMmaOverflow::{Satfinite, Wrapping}; - use RegisterMmaShape::{M8n8k16, M8n8k32, M16n8k16, M16n8k32, M16n8k64}; - - if !common - || mma.operation != RegisterMmaOperation::Multiply - || mma.accumulator != RegisterMmaAccumulator::S32 - { - return None; - } - - let (id, abi_id, operation_key, source_record, llvm_symbol, compatibility_source) = - match (mma.shape, mma.a_element, mma.b_element, mma.overflow) { - (M16n8k32, S8, S8, Wrapping) => ( - "mma_m16n8k32_s32_s8", - "i0108", - "matrix.mma.m16n8k32.row.col.s32.s8.s8.s32.wrapping", - "int_nvvm_mma_m16n8k32_row_col_s8", - "llvm.nvvm.mma.m16n8k32.row.col.s8", - ExistingStub, - ), - (M16n8k16, S8, S8, Wrapping) => ( - "mma_m16n8k16_s32_s8", - "i0110", - "matrix.mma.m16n8k16.row.col.s32.s8.s8.s32.wrapping", - "int_nvvm_mma_m16n8k16_row_col_s8", - "llvm.nvvm.mma.m16n8k16.row.col.s8", - GeneratedStub, - ), - (M16n8k16, S8, U8, Wrapping) => ( - "mma_m16n8k16_s32_s8_u8", - "i0111", - "matrix.mma.m16n8k16.row.col.s32.s8.u8.s32.wrapping", - "int_nvvm_mma_m16n8k16_row_col_s8_u8", - "llvm.nvvm.mma.m16n8k16.row.col.s8.u8", - GeneratedStub, - ), - (M16n8k16, U8, U8, Wrapping) => ( - "mma_m16n8k16_s32_u8", - "i0112", - "matrix.mma.m16n8k16.row.col.s32.u8.u8.s32.wrapping", - "int_nvvm_mma_m16n8k16_row_col_u8", - "llvm.nvvm.mma.m16n8k16.row.col.u8", - GeneratedStub, - ), - (M16n8k16, U8, S8, Wrapping) => ( - "mma_m16n8k16_s32_u8_s8", - "i0113", - "matrix.mma.m16n8k16.row.col.s32.u8.s8.s32.wrapping", - "int_nvvm_mma_m16n8k16_row_col_u8_s8", - "llvm.nvvm.mma.m16n8k16.row.col.u8.s8", - GeneratedStub, - ), - (M16n8k32, S8, U8, Wrapping) => ( - "mma_m16n8k32_s32_s8_u8", - "i0114", - "matrix.mma.m16n8k32.row.col.s32.s8.u8.s32.wrapping", - "int_nvvm_mma_m16n8k32_row_col_s8_u8", - "llvm.nvvm.mma.m16n8k32.row.col.s8.u8", - GeneratedStub, - ), - (M16n8k32, U8, U8, Wrapping) => ( - "mma_m16n8k32_s32_u8", - "i0115", - "matrix.mma.m16n8k32.row.col.s32.u8.u8.s32.wrapping", - "int_nvvm_mma_m16n8k32_row_col_u8", - "llvm.nvvm.mma.m16n8k32.row.col.u8", - GeneratedStub, - ), - (M16n8k32, U8, S8, Wrapping) => ( - "mma_m16n8k32_s32_u8_s8", - "i0116", - "matrix.mma.m16n8k32.row.col.s32.u8.s8.s32.wrapping", - "int_nvvm_mma_m16n8k32_row_col_u8_s8", - "llvm.nvvm.mma.m16n8k32.row.col.u8.s8", - GeneratedStub, - ), - (M16n8k16, S8, S8, Satfinite) => ( - "mma_m16n8k16_s32_s8_satfinite", - "i0117", - "matrix.mma.m16n8k16.row.col.s32.s8.s8.s32.satfinite", - "int_nvvm_mma_m16n8k16_row_col_satfinite_s8", - "llvm.nvvm.mma.m16n8k16.row.col.satfinite.s8", - GeneratedStub, - ), - (M16n8k16, S8, U8, Satfinite) => ( - "mma_m16n8k16_s32_s8_u8_satfinite", - "i0118", - "matrix.mma.m16n8k16.row.col.s32.s8.u8.s32.satfinite", - "int_nvvm_mma_m16n8k16_row_col_satfinite_s8_u8", - "llvm.nvvm.mma.m16n8k16.row.col.satfinite.s8.u8", - GeneratedStub, - ), - (M16n8k16, U8, U8, Satfinite) => ( - "mma_m16n8k16_s32_u8_satfinite", - "i0119", - "matrix.mma.m16n8k16.row.col.s32.u8.u8.s32.satfinite", - "int_nvvm_mma_m16n8k16_row_col_satfinite_u8", - "llvm.nvvm.mma.m16n8k16.row.col.satfinite.u8", - GeneratedStub, - ), - (M16n8k16, U8, S8, Satfinite) => ( - "mma_m16n8k16_s32_u8_s8_satfinite", - "i0120", - "matrix.mma.m16n8k16.row.col.s32.u8.s8.s32.satfinite", - "int_nvvm_mma_m16n8k16_row_col_satfinite_u8_s8", - "llvm.nvvm.mma.m16n8k16.row.col.satfinite.u8.s8", - GeneratedStub, - ), - (M16n8k32, S8, S8, Satfinite) => ( - "mma_m16n8k32_s32_s8_satfinite", - "i0121", - "matrix.mma.m16n8k32.row.col.s32.s8.s8.s32.satfinite", - "int_nvvm_mma_m16n8k32_row_col_satfinite_s8", - "llvm.nvvm.mma.m16n8k32.row.col.satfinite.s8", - GeneratedStub, - ), - (M16n8k32, S8, U8, Satfinite) => ( - "mma_m16n8k32_s32_s8_u8_satfinite", - "i0122", - "matrix.mma.m16n8k32.row.col.s32.s8.u8.s32.satfinite", - "int_nvvm_mma_m16n8k32_row_col_satfinite_s8_u8", - "llvm.nvvm.mma.m16n8k32.row.col.satfinite.s8.u8", - GeneratedStub, - ), - (M16n8k32, U8, U8, Satfinite) => ( - "mma_m16n8k32_s32_u8_satfinite", - "i0123", - "matrix.mma.m16n8k32.row.col.s32.u8.u8.s32.satfinite", - "int_nvvm_mma_m16n8k32_row_col_satfinite_u8", - "llvm.nvvm.mma.m16n8k32.row.col.satfinite.u8", - GeneratedStub, - ), - (M16n8k32, U8, S8, Satfinite) => ( - "mma_m16n8k32_s32_u8_s8_satfinite", - "i0124", - "matrix.mma.m16n8k32.row.col.s32.u8.s8.s32.satfinite", - "int_nvvm_mma_m16n8k32_row_col_satfinite_u8_s8", - "llvm.nvvm.mma.m16n8k32.row.col.satfinite.u8.s8", - GeneratedStub, - ), - (M8n8k16, S8, S8, Wrapping) => ( - "mma_m8n8k16_s32_s8", - "i0125", - "matrix.mma.m8n8k16.row.col.s32.s8.s8.s32.wrapping", - "int_nvvm_mma_m8n8k16_row_col_s8", - "llvm.nvvm.mma.m8n8k16.row.col.s8", - GeneratedStub, - ), - (M8n8k16, S8, U8, Wrapping) => ( - "mma_m8n8k16_s32_s8_u8", - "i0126", - "matrix.mma.m8n8k16.row.col.s32.s8.u8.s32.wrapping", - "int_nvvm_mma_m8n8k16_row_col_s8_u8", - "llvm.nvvm.mma.m8n8k16.row.col.s8.u8", - GeneratedStub, - ), - (M8n8k16, U8, U8, Wrapping) => ( - "mma_m8n8k16_s32_u8", - "i0127", - "matrix.mma.m8n8k16.row.col.s32.u8.u8.s32.wrapping", - "int_nvvm_mma_m8n8k16_row_col_u8", - "llvm.nvvm.mma.m8n8k16.row.col.u8", - GeneratedStub, - ), - (M8n8k16, U8, S8, Wrapping) => ( - "mma_m8n8k16_s32_u8_s8", - "i0128", - "matrix.mma.m8n8k16.row.col.s32.u8.s8.s32.wrapping", - "int_nvvm_mma_m8n8k16_row_col_u8_s8", - "llvm.nvvm.mma.m8n8k16.row.col.u8.s8", - GeneratedStub, - ), - (M8n8k16, S8, S8, Satfinite) => ( - "mma_m8n8k16_s32_s8_satfinite", - "i0129", - "matrix.mma.m8n8k16.row.col.s32.s8.s8.s32.satfinite", - "int_nvvm_mma_m8n8k16_row_col_satfinite_s8", - "llvm.nvvm.mma.m8n8k16.row.col.satfinite.s8", - GeneratedStub, - ), - (M8n8k16, S8, U8, Satfinite) => ( - "mma_m8n8k16_s32_s8_u8_satfinite", - "i0130", - "matrix.mma.m8n8k16.row.col.s32.s8.u8.s32.satfinite", - "int_nvvm_mma_m8n8k16_row_col_satfinite_s8_u8", - "llvm.nvvm.mma.m8n8k16.row.col.satfinite.s8.u8", - GeneratedStub, - ), - (M8n8k16, U8, U8, Satfinite) => ( - "mma_m8n8k16_s32_u8_satfinite", - "i0131", - "matrix.mma.m8n8k16.row.col.s32.u8.u8.s32.satfinite", - "int_nvvm_mma_m8n8k16_row_col_satfinite_u8", - "llvm.nvvm.mma.m8n8k16.row.col.satfinite.u8", - GeneratedStub, - ), - (M8n8k16, U8, S8, Satfinite) => ( - "mma_m8n8k16_s32_u8_s8_satfinite", - "i0132", - "matrix.mma.m8n8k16.row.col.s32.u8.s8.s32.satfinite", - "int_nvvm_mma_m8n8k16_row_col_satfinite_u8_s8", - "llvm.nvvm.mma.m8n8k16.row.col.satfinite.u8.s8", - GeneratedStub, - ), - (M8n8k32, S4, S4, Wrapping) => ( - "mma_m8n8k32_s32_s4", - "i0133", - "matrix.mma.m8n8k32.row.col.s32.s4.s4.s32.wrapping", - "int_nvvm_mma_m8n8k32_row_col_s4", - "llvm.nvvm.mma.m8n8k32.row.col.s4", - GeneratedStub, - ), - (M8n8k32, S4, U4, Wrapping) => ( - "mma_m8n8k32_s32_s4_u4", - "i0134", - "matrix.mma.m8n8k32.row.col.s32.s4.u4.s32.wrapping", - "int_nvvm_mma_m8n8k32_row_col_s4_u4", - "llvm.nvvm.mma.m8n8k32.row.col.s4.u4", - GeneratedStub, - ), - (M8n8k32, U4, U4, Wrapping) => ( - "mma_m8n8k32_s32_u4", - "i0135", - "matrix.mma.m8n8k32.row.col.s32.u4.u4.s32.wrapping", - "int_nvvm_mma_m8n8k32_row_col_u4", - "llvm.nvvm.mma.m8n8k32.row.col.u4", - GeneratedStub, - ), - (M8n8k32, U4, S4, Wrapping) => ( - "mma_m8n8k32_s32_u4_s4", - "i0136", - "matrix.mma.m8n8k32.row.col.s32.u4.s4.s32.wrapping", - "int_nvvm_mma_m8n8k32_row_col_u4_s4", - "llvm.nvvm.mma.m8n8k32.row.col.u4.s4", - GeneratedStub, - ), - (M8n8k32, S4, S4, Satfinite) => ( - "mma_m8n8k32_s32_s4_satfinite", - "i0137", - "matrix.mma.m8n8k32.row.col.s32.s4.s4.s32.satfinite", - "int_nvvm_mma_m8n8k32_row_col_satfinite_s4", - "llvm.nvvm.mma.m8n8k32.row.col.satfinite.s4", - GeneratedStub, - ), - (M8n8k32, S4, U4, Satfinite) => ( - "mma_m8n8k32_s32_s4_u4_satfinite", - "i0138", - "matrix.mma.m8n8k32.row.col.s32.s4.u4.s32.satfinite", - "int_nvvm_mma_m8n8k32_row_col_satfinite_s4_u4", - "llvm.nvvm.mma.m8n8k32.row.col.satfinite.s4.u4", - GeneratedStub, - ), - (M8n8k32, U4, U4, Satfinite) => ( - "mma_m8n8k32_s32_u4_satfinite", - "i0139", - "matrix.mma.m8n8k32.row.col.s32.u4.u4.s32.satfinite", - "int_nvvm_mma_m8n8k32_row_col_satfinite_u4", - "llvm.nvvm.mma.m8n8k32.row.col.satfinite.u4", - GeneratedStub, - ), - (M8n8k32, U4, S4, Satfinite) => ( - "mma_m8n8k32_s32_u4_s4_satfinite", - "i0140", - "matrix.mma.m8n8k32.row.col.s32.u4.s4.s32.satfinite", - "int_nvvm_mma_m8n8k32_row_col_satfinite_u4_s4", - "llvm.nvvm.mma.m8n8k32.row.col.satfinite.u4.s4", - GeneratedStub, - ), - (M16n8k32, S4, S4, Wrapping) => ( - "mma_m16n8k32_s32_s4", - "i0141", - "matrix.mma.m16n8k32.row.col.s32.s4.s4.s32.wrapping", - "int_nvvm_mma_m16n8k32_row_col_s4", - "llvm.nvvm.mma.m16n8k32.row.col.s4", - GeneratedStub, - ), - (M16n8k32, S4, U4, Wrapping) => ( - "mma_m16n8k32_s32_s4_u4", - "i0142", - "matrix.mma.m16n8k32.row.col.s32.s4.u4.s32.wrapping", - "int_nvvm_mma_m16n8k32_row_col_s4_u4", - "llvm.nvvm.mma.m16n8k32.row.col.s4.u4", - GeneratedStub, - ), - (M16n8k32, U4, U4, Wrapping) => ( - "mma_m16n8k32_s32_u4", - "i0143", - "matrix.mma.m16n8k32.row.col.s32.u4.u4.s32.wrapping", - "int_nvvm_mma_m16n8k32_row_col_u4", - "llvm.nvvm.mma.m16n8k32.row.col.u4", - GeneratedStub, - ), - (M16n8k32, U4, S4, Wrapping) => ( - "mma_m16n8k32_s32_u4_s4", - "i0144", - "matrix.mma.m16n8k32.row.col.s32.u4.s4.s32.wrapping", - "int_nvvm_mma_m16n8k32_row_col_u4_s4", - "llvm.nvvm.mma.m16n8k32.row.col.u4.s4", - GeneratedStub, - ), - (M16n8k64, S4, S4, Wrapping) => ( - "mma_m16n8k64_s32_s4", - "i0145", - "matrix.mma.m16n8k64.row.col.s32.s4.s4.s32.wrapping", - "int_nvvm_mma_m16n8k64_row_col_s4", - "llvm.nvvm.mma.m16n8k64.row.col.s4", - GeneratedStub, - ), - (M16n8k64, S4, U4, Wrapping) => ( - "mma_m16n8k64_s32_s4_u4", - "i0146", - "matrix.mma.m16n8k64.row.col.s32.s4.u4.s32.wrapping", - "int_nvvm_mma_m16n8k64_row_col_s4_u4", - "llvm.nvvm.mma.m16n8k64.row.col.s4.u4", - GeneratedStub, - ), - (M16n8k64, U4, U4, Wrapping) => ( - "mma_m16n8k64_s32_u4", - "i0147", - "matrix.mma.m16n8k64.row.col.s32.u4.u4.s32.wrapping", - "int_nvvm_mma_m16n8k64_row_col_u4", - "llvm.nvvm.mma.m16n8k64.row.col.u4", - GeneratedStub, - ), - (M16n8k64, U4, S4, Wrapping) => ( - "mma_m16n8k64_s32_u4_s4", - "i0148", - "matrix.mma.m16n8k64.row.col.s32.u4.s4.s32.wrapping", - "int_nvvm_mma_m16n8k64_row_col_u4_s4", - "llvm.nvvm.mma.m16n8k64.row.col.u4.s4", - GeneratedStub, - ), - (M16n8k32, S4, S4, Satfinite) => ( - "mma_m16n8k32_s32_s4_satfinite", - "i0149", - "matrix.mma.m16n8k32.row.col.s32.s4.s4.s32.satfinite", - "int_nvvm_mma_m16n8k32_row_col_satfinite_s4", - "llvm.nvvm.mma.m16n8k32.row.col.satfinite.s4", - GeneratedStub, - ), - (M16n8k32, S4, U4, Satfinite) => ( - "mma_m16n8k32_s32_s4_u4_satfinite", - "i0150", - "matrix.mma.m16n8k32.row.col.s32.s4.u4.s32.satfinite", - "int_nvvm_mma_m16n8k32_row_col_satfinite_s4_u4", - "llvm.nvvm.mma.m16n8k32.row.col.satfinite.s4.u4", - GeneratedStub, - ), - (M16n8k32, U4, U4, Satfinite) => ( - "mma_m16n8k32_s32_u4_satfinite", - "i0151", - "matrix.mma.m16n8k32.row.col.s32.u4.u4.s32.satfinite", - "int_nvvm_mma_m16n8k32_row_col_satfinite_u4", - "llvm.nvvm.mma.m16n8k32.row.col.satfinite.u4", - GeneratedStub, - ), - (M16n8k32, U4, S4, Satfinite) => ( - "mma_m16n8k32_s32_u4_s4_satfinite", - "i0152", - "matrix.mma.m16n8k32.row.col.s32.u4.s4.s32.satfinite", - "int_nvvm_mma_m16n8k32_row_col_satfinite_u4_s4", - "llvm.nvvm.mma.m16n8k32.row.col.satfinite.u4.s4", - GeneratedStub, - ), - (M16n8k64, S4, S4, Satfinite) => ( - "mma_m16n8k64_s32_s4_satfinite", - "i0153", - "matrix.mma.m16n8k64.row.col.s32.s4.s4.s32.satfinite", - "int_nvvm_mma_m16n8k64_row_col_satfinite_s4", - "llvm.nvvm.mma.m16n8k64.row.col.satfinite.s4", - GeneratedStub, - ), - (M16n8k64, S4, U4, Satfinite) => ( - "mma_m16n8k64_s32_s4_u4_satfinite", - "i0154", - "matrix.mma.m16n8k64.row.col.s32.s4.u4.s32.satfinite", - "int_nvvm_mma_m16n8k64_row_col_satfinite_s4_u4", - "llvm.nvvm.mma.m16n8k64.row.col.satfinite.s4.u4", - GeneratedStub, - ), - (M16n8k64, U4, U4, Satfinite) => ( - "mma_m16n8k64_s32_u4_satfinite", - "i0155", - "matrix.mma.m16n8k64.row.col.s32.u4.u4.s32.satfinite", - "int_nvvm_mma_m16n8k64_row_col_satfinite_u4", - "llvm.nvvm.mma.m16n8k64.row.col.satfinite.u4", - GeneratedStub, - ), - (M16n8k64, U4, S4, Satfinite) => ( - "mma_m16n8k64_s32_u4_s4_satfinite", - "i0156", - "matrix.mma.m16n8k64.row.col.s32.u4.s4.s32.satfinite", - "int_nvvm_mma_m16n8k64_row_col_satfinite_u4_s4", - "llvm.nvvm.mma.m16n8k64.row.col.satfinite.u4.s4", - GeneratedStub, - ), - _ => return None, - }; - - let int4 = matches!(mma.a_element, S4 | U4); - let (rust_arguments, dialect_operands, llvm_arguments, adapter, shape, register_counts) = - match (mma.shape, int4) { - (M8n8k16, false) | (M8n8k32, true) => ( - &["[i32; 2]", "u32", "u32"] as &'static [&'static str], - &["i32", "i32", "i32", "i32"] as &'static [&'static str], - &["i32", "i32", "i32", "i32"] as &'static [&'static str], - C2I32A1U32B1U32ToD2I32, - match mma.shape { - M8n8k16 => "m8n8k16", - M8n8k32 => "m8n8k32", - _ => unreachable!(), - }, - [2, 1, 1, 2], - ), - (M16n8k16, false) | (M16n8k32, true) => ( - &["[i32; 4]", "[u32; 2]", "u32"] as &'static [&'static str], - &["i32", "i32", "i32", "i32", "i32", "i32", "i32"] as &'static [&'static str], - &["i32", "i32", "i32", "i32", "i32", "i32", "i32"] as &'static [&'static str], - C4I32A2U32B1U32ToD4I32, - match mma.shape { - M16n8k16 => "m16n8k16", - M16n8k32 => "m16n8k32", - _ => unreachable!(), - }, - [4, 2, 1, 4], - ), - (M16n8k32, false) | (M16n8k64, true) => ( - &["[i32; 4]", "[u32; 4]", "[u32; 2]"] as &'static [&'static str], - &[ - "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", - ] as &'static [&'static str], - &[ - "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", - ] as &'static [&'static str], - C4I32A4U32B2U32ToD4I32, - match mma.shape { - M16n8k32 => "m16n8k32", - M16n8k64 => "m16n8k64", - _ => unreachable!(), - }, - [4, 4, 2, 4], - ), - _ => return None, - }; - if mma.adapter != adapter { - return None; - } - - let (rust_result, dialect_results, llvm_results, minimum_ptx, minimum_sm) = match mma.shape { - M8n8k16 | M8n8k32 => ( - "[i32; 2]", - &["i32", "i32"] as &'static [&'static str], - &["i32", "i32"] as &'static [&'static str], - "6.5", - "sm_75", - ), - M16n8k16 | M16n8k32 | M16n8k64 => ( - "[i32; 4]", - &["i32", "i32", "i32", "i32"] as &'static [&'static str], - &["i32", "i32", "i32", "i32"] as &'static [&'static str], - "7.0", - "sm_80", - ), - _ => return None, - }; - - let element = |element| match element { - S4 => Some("s4"), - U4 => Some("u4"), - S8 => Some("s8"), - U8 => Some("u8"), - _ => None, - }; - let mut ptx_modifiers = vec!["sync", "aligned", shape, "row", "col"]; - if mma.overflow == Satfinite { - ptx_modifiers.push("satfinite"); - } - ptx_modifiers.extend([ - "s32", - element(mma.a_element)?, - element(mma.b_element)?, - "s32", - ]); - - Some(RegisterMmaRecipe { - id, - abi_id, - operation_key, - source_record, - llvm_symbol, - rust_arguments, - rust_result, - dialect_op_type: "RegisterMmaOp", - dialect_op_name: "nvvm.register_mma", - dialect_operands, - dialect_results, - llvm_arguments, - llvm_results, - adapter, - compatibility_source, - minimum_ptx, - minimum_sm, - ptx_modifiers, - ptx_register_counts: register_counts, - }) -} - -fn binary_register_mma_recipe(mma: &RegisterMma, common: bool) -> Option { - use RegisterMmaAdapter::{ - C2I32A1U32B1U32ToD2I32, C4I32A2U32B1U32ToD4I32, C4I32A4U32B2U32ToD4I32, - }; - use RegisterMmaOperation::{AndPopc, XorPopc}; - use RegisterMmaShape::{M8n8k128, M16n8k128, M16n8k256}; - - if !common - || mma.accumulator != RegisterMmaAccumulator::S32 - || mma.a_element != RegisterMmaElement::B1 - || mma.b_element != RegisterMmaElement::B1 - || mma.overflow != RegisterMmaOverflow::Wrapping - || mma.compatibility_source != RegisterMmaCompatibilitySource::GeneratedStub - { - return None; - } - - let (id, abi_id, operation_key, source_record, llvm_symbol, operation_name) = - match (mma.shape, mma.operation) { - (M8n8k128, XorPopc) => ( - "mma_m8n8k128_s32_b1_xor_popc", - "i0157", - "matrix.mma.m8n8k128.row.col.s32.b1.b1.s32.xor.popc", - "int_nvvm_mma_xor_popc_m8n8k128_row_col_b1", - "llvm.nvvm.mma.xor.popc.m8n8k128.row.col.b1", - "xor", - ), - (M16n8k128, XorPopc) => ( - "mma_m16n8k128_s32_b1_xor_popc", - "i0158", - "matrix.mma.m16n8k128.row.col.s32.b1.b1.s32.xor.popc", - "int_nvvm_mma_xor_popc_m16n8k128_row_col_b1", - "llvm.nvvm.mma.xor.popc.m16n8k128.row.col.b1", - "xor", - ), - (M16n8k256, XorPopc) => ( - "mma_m16n8k256_s32_b1_xor_popc", - "i0159", - "matrix.mma.m16n8k256.row.col.s32.b1.b1.s32.xor.popc", - "int_nvvm_mma_xor_popc_m16n8k256_row_col_b1", - "llvm.nvvm.mma.xor.popc.m16n8k256.row.col.b1", - "xor", - ), - (M8n8k128, AndPopc) => ( - "mma_m8n8k128_s32_b1_and_popc", - "i0160", - "matrix.mma.m8n8k128.row.col.s32.b1.b1.s32.and.popc", - "int_nvvm_mma_and_popc_m8n8k128_row_col_b1", - "llvm.nvvm.mma.and.popc.m8n8k128.row.col.b1", - "and", - ), - (M16n8k128, AndPopc) => ( - "mma_m16n8k128_s32_b1_and_popc", - "i0161", - "matrix.mma.m16n8k128.row.col.s32.b1.b1.s32.and.popc", - "int_nvvm_mma_and_popc_m16n8k128_row_col_b1", - "llvm.nvvm.mma.and.popc.m16n8k128.row.col.b1", - "and", - ), - (M16n8k256, AndPopc) => ( - "mma_m16n8k256_s32_b1_and_popc", - "i0162", - "matrix.mma.m16n8k256.row.col.s32.b1.b1.s32.and.popc", - "int_nvvm_mma_and_popc_m16n8k256_row_col_b1", - "llvm.nvvm.mma.and.popc.m16n8k256.row.col.b1", - "and", - ), - _ => return None, - }; - - let ( - rust_arguments, - dialect_operands, - llvm_arguments, - rust_result, - result_types, - adapter, - counts, - ) = match mma.shape { - M8n8k128 => ( - &["[i32; 2]", "u32", "u32"] as &'static [&'static str], - &["i32", "i32", "i32", "i32"] as &'static [&'static str], - &["i32", "i32", "i32", "i32"] as &'static [&'static str], - "[i32; 2]", - &["i32", "i32"] as &'static [&'static str], - C2I32A1U32B1U32ToD2I32, - [2, 1, 1, 2], - ), - M16n8k128 => ( - &["[i32; 4]", "[u32; 2]", "u32"] as &'static [&'static str], - &["i32", "i32", "i32", "i32", "i32", "i32", "i32"] as &'static [&'static str], - &["i32", "i32", "i32", "i32", "i32", "i32", "i32"] as &'static [&'static str], - "[i32; 4]", - &["i32", "i32", "i32", "i32"] as &'static [&'static str], - C4I32A2U32B1U32ToD4I32, - [4, 2, 1, 4], - ), - M16n8k256 => ( - &["[i32; 4]", "[u32; 4]", "[u32; 2]"] as &'static [&'static str], - &[ - "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", - ] as &'static [&'static str], - &[ - "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", - ] as &'static [&'static str], - "[i32; 4]", - &["i32", "i32", "i32", "i32"] as &'static [&'static str], - C4I32A4U32B2U32ToD4I32, - [4, 4, 2, 4], - ), - _ => return None, - }; - if mma.adapter != adapter { - return None; - } - - let (minimum_ptx, minimum_sm) = match (mma.operation, mma.shape) { - (XorPopc, M8n8k128) => ("7.0", "sm_75"), - (XorPopc, M16n8k128 | M16n8k256) => ("7.0", "sm_80"), - (AndPopc, M8n8k128 | M16n8k128 | M16n8k256) => ("7.1", "sm_80"), - _ => return None, - }; - let shape = match mma.shape { - M8n8k128 => "m8n8k128", - M16n8k128 => "m16n8k128", - M16n8k256 => "m16n8k256", - _ => return None, - }; - - Some(RegisterMmaRecipe { - id, - abi_id, - operation_key, - source_record, - llvm_symbol, - rust_arguments, - rust_result, - dialect_op_type: "RegisterMmaOp", - dialect_op_name: "nvvm.register_mma", - dialect_operands, - dialect_results: result_types, - llvm_arguments, - llvm_results: result_types, - adapter, - compatibility_source: RegisterMmaCompatibilitySource::GeneratedStub, - minimum_ptx, - minimum_sm, - ptx_modifiers: vec![ - "sync", - "aligned", - shape, - "row", - "col", - "s32", - "b1", - "b1", - "s32", - operation_name, - "popc", - ], - ptx_register_counts: counts, - }) -} - -fn register_mma_recipe(mma: &RegisterMma) -> Option { - use RegisterMmaAccumulator::{F16 as F16Accumulator, F32, F64}; - use RegisterMmaAdapter::{ - C2F64A1F64B1F64ToD2F64, C2U32A2U32B1U32ToD2U32, C2U32A4U32B2U32ToD2U32, - C4F32A2U32B1U32ToD4F32, C4F32A4U32B2U32ToD4F32, - }; - use RegisterMmaElement::{Bf16, F16 as F16Element, F64 as F64Element, Tf32}; - use RegisterMmaShape::{M8n8k4, M16n8k4, M16n8k8, M16n8k16}; - - let common = mma.a_layout == RegisterMmaLayout::Row - && mma.b_layout == RegisterMmaLayout::Col - && mma.participation - == RegisterMmaParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes; - if let Some(recipe) = integer_register_mma_recipe(mma, common) { - return Some(recipe); - } - if let Some(recipe) = binary_register_mma_recipe(mma, common) { - return Some(recipe); - } - if mma.operation != RegisterMmaOperation::Multiply { - return None; - } - match ( - mma.shape, - mma.accumulator, - mma.a_element, - mma.b_element, - mma.overflow, - mma.adapter, - common, - ) { - ( - M16n8k4, - F32, - Tf32, - Tf32, - RegisterMmaOverflow::NotApplicable, - C4F32A2U32B1U32ToD4F32, - true, - ) => Some(RegisterMmaRecipe { - id: "mma_m16n8k4_f32_tf32", - abi_id: "i0520", - operation_key: "matrix.mma.m16n8k4.row.col.f32.tf32.tf32.f32", - source_record: "int_nvvm_mma_m16n8k4_row_col_tf32", - llvm_symbol: "llvm.nvvm.mma.m16n8k4.row.col.tf32", - rust_arguments: &["[f32; 4]", "[u32; 2]", "u32"], - rust_result: "[f32; 4]", - dialect_op_type: "RegisterMmaOp", - dialect_op_name: "nvvm.register_mma", - dialect_operands: &["f32", "f32", "f32", "f32", "i32", "i32", "i32"], - dialect_results: &["f32", "f32", "f32", "f32"], - llvm_arguments: &["i32", "i32", "i32", "f32", "f32", "f32", "f32"], - llvm_results: &["f32", "f32", "f32", "f32"], - adapter: C4F32A2U32B1U32ToD4F32, - compatibility_source: RegisterMmaCompatibilitySource::GeneratedStub, - minimum_ptx: "7.0", - minimum_sm: "sm_80", - ptx_modifiers: vec![ - "sync", "aligned", "m16n8k4", "row", "col", "f32", "tf32", "tf32", "f32", - ], - ptx_register_counts: [4, 2, 1, 4], - }), - ( - M16n8k8, - F16Accumulator, - F16Element, - F16Element, - RegisterMmaOverflow::NotApplicable, - C2U32A2U32B1U32ToD2U32, - true, - ) => Some(RegisterMmaRecipe { - id: "mma_m16n8k8_f16_f16", - abi_id: "i0521", - operation_key: "matrix.mma.m16n8k8.row.col.f16.f16.f16.f16", - source_record: "int_nvvm_mma_m16n8k8_row_col_f16_f16", - llvm_symbol: "llvm.nvvm.mma.m16n8k8.row.col.f16.f16", - rust_arguments: &["[u32; 2]", "[u32; 2]", "u32"], - rust_result: "[u32; 2]", - dialect_op_type: "RegisterMmaOp", - dialect_op_name: "nvvm.register_mma", - dialect_operands: &["i32", "i32", "i32", "i32", "i32"], - dialect_results: &["i32", "i32"], - llvm_arguments: &["v2f16", "v2f16", "v2f16", "v2f16", "v2f16"], - llvm_results: &["v2f16", "v2f16"], - adapter: C2U32A2U32B1U32ToD2U32, - compatibility_source: RegisterMmaCompatibilitySource::GeneratedStub, - minimum_ptx: "7.0", - minimum_sm: "sm_80", - ptx_modifiers: vec![ - "sync", "aligned", "m16n8k8", "row", "col", "f16", "f16", "f16", "f16", - ], - ptx_register_counts: [2, 2, 1, 2], - }), - ( - M16n8k8, - F32, - Bf16, - Bf16, - RegisterMmaOverflow::NotApplicable, - C4F32A2U32B1U32ToD4F32, - true, - ) => Some(RegisterMmaRecipe { - id: "mma_m16n8k8_f32_bf16", - abi_id: "i0522", - operation_key: "matrix.mma.m16n8k8.row.col.f32.bf16.bf16.f32", - source_record: "int_nvvm_mma_m16n8k8_row_col_bf16", - llvm_symbol: "llvm.nvvm.mma.m16n8k8.row.col.bf16", - rust_arguments: &["[f32; 4]", "[u32; 2]", "u32"], - rust_result: "[f32; 4]", - dialect_op_type: "RegisterMmaOp", - dialect_op_name: "nvvm.register_mma", - dialect_operands: &["f32", "f32", "f32", "f32", "i32", "i32", "i32"], - dialect_results: &["f32", "f32", "f32", "f32"], - llvm_arguments: &["i32", "i32", "i32", "f32", "f32", "f32", "f32"], - llvm_results: &["f32", "f32", "f32", "f32"], - adapter: C4F32A2U32B1U32ToD4F32, - compatibility_source: RegisterMmaCompatibilitySource::GeneratedStub, - minimum_ptx: "7.0", - minimum_sm: "sm_80", - ptx_modifiers: vec![ - "sync", "aligned", "m16n8k8", "row", "col", "f32", "bf16", "bf16", "f32", - ], - ptx_register_counts: [4, 2, 1, 4], - }), - ( - M16n8k8, - F32, - F16Element, - F16Element, - RegisterMmaOverflow::NotApplicable, - C4F32A2U32B1U32ToD4F32, - true, - ) => Some(RegisterMmaRecipe { - id: "mma_m16n8k8_f32_f16", - abi_id: "i0523", - operation_key: "matrix.mma.m16n8k8.row.col.f32.f16.f16.f32", - source_record: "int_nvvm_mma_m16n8k8_row_col_f32_f32", - llvm_symbol: "llvm.nvvm.mma.m16n8k8.row.col.f32.f32", - rust_arguments: &["[f32; 4]", "[u32; 2]", "u32"], - rust_result: "[f32; 4]", - dialect_op_type: "RegisterMmaOp", - dialect_op_name: "nvvm.register_mma", - dialect_operands: &["f32", "f32", "f32", "f32", "i32", "i32", "i32"], - dialect_results: &["f32", "f32", "f32", "f32"], - llvm_arguments: &["v2f16", "v2f16", "v2f16", "f32", "f32", "f32", "f32"], - llvm_results: &["f32", "f32", "f32", "f32"], - adapter: C4F32A2U32B1U32ToD4F32, - compatibility_source: RegisterMmaCompatibilitySource::GeneratedStub, - minimum_ptx: "7.0", - minimum_sm: "sm_80", - ptx_modifiers: vec![ - "sync", "aligned", "m16n8k8", "row", "col", "f32", "f16", "f16", "f32", - ], - ptx_register_counts: [4, 2, 1, 4], - }), - ( - M16n8k16, - F16Accumulator, - F16Element, - F16Element, - RegisterMmaOverflow::NotApplicable, - C2U32A4U32B2U32ToD2U32, - true, - ) => Some(RegisterMmaRecipe { - id: "mma_m16n8k16_f16_f16", - abi_id: "i0524", - operation_key: "matrix.mma.m16n8k16.row.col.f16.f16.f16.f16", - source_record: "int_nvvm_mma_m16n8k16_row_col_f16_f16", - llvm_symbol: "llvm.nvvm.mma.m16n8k16.row.col.f16.f16", - rust_arguments: &["[u32; 2]", "[u32; 4]", "[u32; 2]"], - rust_result: "[u32; 2]", - dialect_op_type: "RegisterMmaOp", - dialect_op_name: "nvvm.register_mma", - dialect_operands: &["i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32"], - dialect_results: &["i32", "i32"], - llvm_arguments: &[ - "v2f16", "v2f16", "v2f16", "v2f16", "v2f16", "v2f16", "v2f16", "v2f16", - ], - llvm_results: &["v2f16", "v2f16"], - adapter: C2U32A4U32B2U32ToD2U32, - compatibility_source: RegisterMmaCompatibilitySource::GeneratedStub, - minimum_ptx: "7.0", - minimum_sm: "sm_80", - ptx_modifiers: vec![ - "sync", "aligned", "m16n8k16", "row", "col", "f16", "f16", "f16", "f16", - ], - ptx_register_counts: [2, 4, 2, 2], - }), - ( - M16n8k16, - F32, - Bf16, - Bf16, - RegisterMmaOverflow::NotApplicable, - C4F32A4U32B2U32ToD4F32, - true, - ) => Some(RegisterMmaRecipe { - id: "mma_m16n8k16_f32_bf16", - abi_id: "i0105", - operation_key: "matrix.mma.m16n8k16.row.col.f32.bf16.bf16.f32", - source_record: "int_nvvm_mma_m16n8k16_row_col_bf16", - llvm_symbol: "llvm.nvvm.mma.m16n8k16.row.col.bf16", - rust_arguments: &["[f32; 4]", "[u32; 4]", "[u32; 2]"], - rust_result: "[f32; 4]", - dialect_op_type: "RegisterMmaOp", - dialect_op_name: "nvvm.register_mma", - dialect_operands: &[ - "f32", "f32", "f32", "f32", "i32", "i32", "i32", "i32", "i32", "i32", - ], - dialect_results: &["f32", "f32", "f32", "f32"], - llvm_arguments: &[ - "i32", "i32", "i32", "i32", "i32", "i32", "f32", "f32", "f32", "f32", - ], - llvm_results: &["f32", "f32", "f32", "f32"], - adapter: C4F32A4U32B2U32ToD4F32, - compatibility_source: RegisterMmaCompatibilitySource::ExistingStub, - minimum_ptx: "7.0", - minimum_sm: "sm_80", - ptx_modifiers: vec![ - "sync", "aligned", "m16n8k16", "row", "col", "f32", "bf16", "bf16", "f32", - ], - ptx_register_counts: [4, 4, 2, 4], - }), - ( - M16n8k16, - F32, - F16Element, - F16Element, - RegisterMmaOverflow::NotApplicable, - C4F32A4U32B2U32ToD4F32, - true, - ) => Some(RegisterMmaRecipe { - id: "mma_m16n8k16_f32_f16", - abi_id: "i0106", - operation_key: "matrix.mma.m16n8k16.row.col.f32.f16.f16.f32", - source_record: "int_nvvm_mma_m16n8k16_row_col_f32_f32", - llvm_symbol: "llvm.nvvm.mma.m16n8k16.row.col.f32.f32", - rust_arguments: &["[f32; 4]", "[u32; 4]", "[u32; 2]"], - rust_result: "[f32; 4]", - dialect_op_type: "RegisterMmaOp", - dialect_op_name: "nvvm.register_mma", - dialect_operands: &[ - "f32", "f32", "f32", "f32", "i32", "i32", "i32", "i32", "i32", "i32", - ], - dialect_results: &["f32", "f32", "f32", "f32"], - llvm_arguments: &[ - "v2f16", "v2f16", "v2f16", "v2f16", "v2f16", "v2f16", "f32", "f32", "f32", "f32", - ], - llvm_results: &["f32", "f32", "f32", "f32"], - adapter: C4F32A4U32B2U32ToD4F32, - compatibility_source: RegisterMmaCompatibilitySource::ExistingStub, - minimum_ptx: "7.0", - minimum_sm: "sm_80", - ptx_modifiers: vec![ - "sync", "aligned", "m16n8k16", "row", "col", "f32", "f16", "f16", "f32", - ], - ptx_register_counts: [4, 4, 2, 4], - }), - ( - M16n8k8, - F32, - Tf32, - Tf32, - RegisterMmaOverflow::NotApplicable, - C4F32A4U32B2U32ToD4F32, - true, - ) => Some(RegisterMmaRecipe { - id: "mma_m16n8k8_f32_tf32", - abi_id: "i0107", - operation_key: "matrix.mma.m16n8k8.row.col.f32.tf32.tf32.f32", - source_record: "int_nvvm_mma_m16n8k8_row_col_tf32", - llvm_symbol: "llvm.nvvm.mma.m16n8k8.row.col.tf32", - rust_arguments: &["[f32; 4]", "[u32; 4]", "[u32; 2]"], - rust_result: "[f32; 4]", - dialect_op_type: "RegisterMmaOp", - dialect_op_name: "nvvm.register_mma", - dialect_operands: &[ - "f32", "f32", "f32", "f32", "i32", "i32", "i32", "i32", "i32", "i32", - ], - dialect_results: &["f32", "f32", "f32", "f32"], - llvm_arguments: &[ - "i32", "i32", "i32", "i32", "i32", "i32", "f32", "f32", "f32", "f32", - ], - llvm_results: &["f32", "f32", "f32", "f32"], - adapter: C4F32A4U32B2U32ToD4F32, - compatibility_source: RegisterMmaCompatibilitySource::ExistingStub, - minimum_ptx: "7.0", - minimum_sm: "sm_80", - ptx_modifiers: vec![ - "sync", "aligned", "m16n8k8", "row", "col", "f32", "tf32", "tf32", "f32", - ], - ptx_register_counts: [4, 4, 2, 4], - }), - ( - M8n8k4, - F64, - F64Element, - F64Element, - RegisterMmaOverflow::NotApplicable, - C2F64A1F64B1F64ToD2F64, - true, - ) => Some(RegisterMmaRecipe { - id: "mma_m8n8k4_f64", - abi_id: "i0109", - operation_key: "matrix.mma.m8n8k4.row.col.f64.f64.f64.f64", - source_record: "int_nvvm_mma_m8n8k4_row_col_f64", - llvm_symbol: "llvm.nvvm.mma.m8n8k4.row.col.f64", - rust_arguments: &["[f64; 2]", "f64", "f64"], - rust_result: "[f64; 2]", - dialect_op_type: "RegisterMmaOp", - dialect_op_name: "nvvm.register_mma", - dialect_operands: &["f64", "f64", "f64", "f64"], - dialect_results: &["f64", "f64"], - llvm_arguments: &["f64", "f64", "f64", "f64"], - llvm_results: &["f64", "f64"], - adapter: C2F64A1F64B1F64ToD2F64, - compatibility_source: RegisterMmaCompatibilitySource::ExistingStub, - minimum_ptx: "7.0", - minimum_sm: "sm_80", - ptx_modifiers: vec![ - "sync", "aligned", "m8n8k4", "row", "col", "f64", "f64", "f64", "f64", - ], - ptx_register_counts: [2, 1, 1, 2], - }), - _ => None, - } -} - -#[derive(Debug, Clone, Copy)] -enum RegisterMmaIntegerKind { - Int4, - Int8, -} - -impl RegisterMmaIntegerKind { - fn label(self) -> &'static str { - match self { - Self::Int4 => "INT4", - Self::Int8 => "INT8", - } - } - - fn supports(self, shape: RegisterMmaShape, element: RegisterMmaElement) -> bool { - match self { - Self::Int4 => { - matches!( - shape, - RegisterMmaShape::M8n8k32 - | RegisterMmaShape::M16n8k32 - | RegisterMmaShape::M16n8k64 - ) && matches!(element, RegisterMmaElement::S4 | RegisterMmaElement::U4) - } - Self::Int8 => { - matches!( - shape, - RegisterMmaShape::M8n8k16 - | RegisterMmaShape::M16n8k16 - | RegisterMmaShape::M16n8k32 - ) && matches!(element, RegisterMmaElement::S8 | RegisterMmaElement::U8) - } - } - } -} - -fn expand_register_mma_integer_admission( - kind: RegisterMmaIntegerKind, - admission: &RegisterMmaIntegerAdmission, -) -> Result> { - use RegisterMmaAdapter::{ - C2I32A1U32B1U32ToD2I32, C4I32A2U32B1U32ToD4I32, C4I32A4U32B2U32ToD4I32, - }; - use RegisterMmaCompatibilitySource::GeneratedStub; - use RegisterMmaLayout::{Col, Row}; - use RegisterMmaParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes; - use RegisterMmaShape::{M8n8k16, M8n8k32, M16n8k16, M16n8k32, M16n8k64}; - - ensure!( - !admission.variants.is_empty(), - "compact {} MMA admission has no variants", - kind.label() - ); - ensure!( - admission.runtime_validation == RuntimeValidation::Unexecuted, - "{} MMA runtime validation may be marked executed only with GPU evidence", - kind.label() - ); - - let mut seen = BTreeSet::new(); - let mut records = Vec::with_capacity(admission.variants.len()); - for variant in &admission.variants { - ensure!( - seen.insert(( - variant.shape, - variant.a_element, - variant.b_element, - variant.overflow, - )), - "compact {} MMA admission contains a duplicate variant", - kind.label() - ); - ensure!( - kind.supports(variant.shape, variant.a_element) - && kind.supports(variant.shape, variant.b_element), - "compact {} MMA admission contains an unsupported shape or element", - kind.label() - ); - let adapter = - match (kind, variant.shape) { - (RegisterMmaIntegerKind::Int8, M8n8k16) - | (RegisterMmaIntegerKind::Int4, M8n8k32) => C2I32A1U32B1U32ToD2I32, - (RegisterMmaIntegerKind::Int8, M16n8k16) - | (RegisterMmaIntegerKind::Int4, M16n8k32) => C4I32A2U32B1U32ToD4I32, - (RegisterMmaIntegerKind::Int8, M16n8k32) - | (RegisterMmaIntegerKind::Int4, M16n8k64) => C4I32A4U32B2U32ToD4I32, - _ => bail!( - "compact {} MMA admission contains an unsupported shape", - kind.label() - ), - }; - let mma = RegisterMma { - shape: variant.shape, - operation: RegisterMmaOperation::Multiply, - kind: None, - accumulator: RegisterMmaAccumulator::S32, - a_element: variant.a_element, - b_element: variant.b_element, - a_layout: Row, - b_layout: Col, - overflow: variant.overflow, - participation: AllWarpLanesSameInstructionAndQualifiersNoExitedLanes, - adapter, - compatibility_source: GeneratedStub, - runtime_validation: admission.runtime_validation, - }; - let recipe = register_mma_recipe(&mma).with_context(|| { - format!( - "compact {} MMA admission requests a variant outside the closed recipe set", - kind.label() - ) - })?; - ensure!( - recipe.compatibility_source == GeneratedStub, - "compact {} MMA admission may only add generated compatibility stubs", - kind.label() - ); - - let element = |element| match (kind, element) { - (RegisterMmaIntegerKind::Int4, RegisterMmaElement::S4) - | (RegisterMmaIntegerKind::Int8, RegisterMmaElement::S8) => Ok("signed"), - (RegisterMmaIntegerKind::Int4, RegisterMmaElement::U4) - | (RegisterMmaIntegerKind::Int8, RegisterMmaElement::U8) => Ok("unsigned"), - _ => bail!( - "compact {} MMA admission contains an unsupported element", - kind.label() - ), - }; - let overflow = match variant.overflow { - RegisterMmaOverflow::Wrapping => "wrapping", - RegisterMmaOverflow::Satfinite => "saturating", - RegisterMmaOverflow::NotApplicable => { - bail!( - "compact {} MMA admission requires an integer overflow mode", - kind.label() - ) - } - }; - let summary = format!( - "Multiplies warp-distributed {} A and {} B {} fragments and adds a {overflow} s32 accumulator.", - element(variant.a_element)?, - element(variant.b_element)?, - kind.label(), - ); - - records.push(OverlayIntrinsic { - id: recipe.id.into(), - abi_id: recipe.abi_id.into(), - operation_key: recipe.operation_key.into(), - family: "register_mma".into(), - source: None, - source_record: Some(recipe.source_record.into()), - rust_module: "matrix".into(), - rust_name: recipe.id.into(), - rust_arguments: recipe - .rust_arguments - .iter() - .map(|value| (*value).into()) - .collect(), - rust_result: recipe.rust_result.into(), - safe: false, - must_use: true, - safe_allowlist_reason: None, - public_rust_path: format!("cuda_intrinsics::matrix::{}", recipe.id), - compatibility_rust_paths: vec![format!("cuda_device::wmma::{}", recipe.id)], - dialect_op_type: recipe.dialect_op_type.into(), - dialect_op_name: recipe.dialect_op_name.into(), - dialect_operands: recipe - .dialect_operands - .iter() - .map(|value| (*value).into()) - .collect(), - dialect_results: recipe - .dialect_results - .iter() - .map(|value| (*value).into()) - .collect(), - llvm_symbol: Some(recipe.llvm_symbol.into()), - resolved_llvm_symbol: None, - llvm_arguments: recipe - .llvm_arguments - .iter() - .map(|value| (*value).into()) - .collect(), - llvm_results: recipe - .llvm_results - .iter() - .map(|value| (*value).into()) - .collect(), - pure: false, - memory: "none".into(), - convergent: true, - execution_scope: "warp".into(), - minimum_ptx: recipe.minimum_ptx.into(), - minimum_sm: Some(recipe.minimum_sm.into()), - ptx_result: recipe.rust_result.into(), - targets: "all".into(), - ptx_isa_version: "9.3".into(), - ptx_isa_section: "9.7.15.5.14 Multiply-and-Accumulate Instruction: mma".into(), - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma".into(), - lowering: "generated_register_mma".into(), - backend_lowerings: vec![ - OverlayBackendLowering { - backend: IntrinsicBackend::LlvmNvptx, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: admission.llvm_evidence_profile.clone(), - targets: None, - minimum_ptx: Some(recipe.minimum_ptx.into()), - minimum_sm: Some(recipe.minimum_sm.into()), - }, - OverlayBackendLowering { - backend: IntrinsicBackend::LibNvvm, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: admission.libnvvm_evidence_profile.clone(), - targets: None, - minimum_ptx: Some(recipe.minimum_ptx.into()), - minimum_sm: Some(recipe.minimum_sm.into()), - }, - ], - packed_atomic: None, - redux: None, - vote: None, - active_mask: None, - warp_match: None, - warp_barrier: None, - warp_shuffle: None, - dot_product: None, - packed_alu: None, - integer_minmax: None, - packed_conversion: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cp_async_copy: None, - cp_async_control: None, - cp_async_mbarrier: None, - mbarrier_basic: None, - movmatrix: None, - mbarrier_extended: None, - register_mma: Some(mma), - sparse_mma: None, - prmt: None, - cluster_barrier: None, - wgmma_control: None, - special_register: None, - debug_control: None, - cluster_memory: None, - clc: None, - tma: None, - tcgen05: None, - ldmatrix_variant: None, - ldmatrix_safety: None, - ldmatrix_adapter: None, - selected_address_space: None, - expected_ptx: InstructionPattern { - mnemonic: "mma".into(), - modifiers: recipe - .ptx_modifiers - .iter() - .map(|value| (*value).into()) - .collect(), - operands: recipe - .ptx_register_counts - .map(|length| OperandPattern::RegisterList { length }) - .into(), - }, - summary, - }); - } - Ok(records) -} - -const REGISTER_MMA_F8F6F4_ELEMENTS: [RegisterMmaElement; 5] = [ - RegisterMmaElement::E2m1, - RegisterMmaElement::E2m3, - RegisterMmaElement::E3m2, - RegisterMmaElement::E4m3, - RegisterMmaElement::E5m2, -]; - -fn register_mma_f8f6f4_element_name(element: RegisterMmaElement) -> Option<&'static str> { - match element { - RegisterMmaElement::E2m1 => Some("e2m1"), - RegisterMmaElement::E2m3 => Some("e2m3"), - RegisterMmaElement::E3m2 => Some("e3m2"), - RegisterMmaElement::E4m3 => Some("e4m3"), - RegisterMmaElement::E5m2 => Some("e5m2"), - _ => None, - } -} - -fn is_dense_f8f6f4_register_mma_policy(policy: &OverlayIntrinsic) -> bool { - policy.family == "register_mma" - && policy.targets == REGISTER_MMA_F8F6F4_TARGETS - && policy.register_mma.as_ref().is_some_and(|mma| { - mma.kind != Some(RegisterMmaKind::Standard) - && mma.shape == RegisterMmaShape::M16n8k32 - && mma.operation == RegisterMmaOperation::Multiply - && matches!( - mma.accumulator, - RegisterMmaAccumulator::F16 | RegisterMmaAccumulator::F32 - ) - && register_mma_f8f6f4_element_name(mma.a_element).is_some() - && register_mma_f8f6f4_element_name(mma.b_element).is_some() - && mma.a_layout == RegisterMmaLayout::Row - && mma.b_layout == RegisterMmaLayout::Col - && mma.overflow == RegisterMmaOverflow::NotApplicable - }) -} - -fn is_sparse_f8f6f4_f16_policy(policy: &OverlayIntrinsic) -> bool { - policy.family == "sparse_mma" - && policy.targets == SPARSE_MMA_F8F6F4_TARGETS - && policy.sparse_mma.as_ref().is_some_and(|mma| { - mma.shape == SparseMmaShape::M16n8k64 - && mma.accumulator == SparseMmaAccumulator::F16 - && SPARSE_MMA_F8F6F4_ELEMENTS.contains(&mma.a_element) - && SPARSE_MMA_F8F6F4_ELEMENTS.contains(&mma.b_element) - && mma.a_layout == SparseMmaLayout::Row - && mma.b_layout == SparseMmaLayout::Col - && mma.overflow == SparseMmaOverflow::NotApplicable - && mma.metadata == SparseMmaMetadata::Ordered - && mma.selector == SparseMmaSelector::ImmediateZero - }) -} - -fn is_f8f6f4_mma_target_matrix_policy(policy: &OverlayIntrinsic) -> bool { - is_dense_f8f6f4_register_mma_policy(policy) || is_sparse_f8f6f4_f16_policy(policy) -} - -#[derive(Clone, Copy)] -struct RegisterMmaF8F6F4Contract { - accumulator: RegisterMmaAccumulator, - scalar_name: &'static str, - rust_arguments: [&'static str; 3], - rust_result: &'static str, - dialect_operands: &'static [&'static str], - dialect_results: &'static [&'static str], - llvm_arguments: &'static [&'static str], - llvm_results: &'static [&'static str], - adapter: RegisterMmaAdapter, - ptx_register_counts: [usize; 4], - summary_accumulator: &'static str, -} - -fn register_mma_f8f6f4_contract( - accumulator: RegisterMmaAccumulator, -) -> Result { - match accumulator { - RegisterMmaAccumulator::F16 => Ok(RegisterMmaF8F6F4Contract { - accumulator, - scalar_name: "f16", - rust_arguments: ["[u32; 2]", "[u32; 4]", "[u32; 2]"], - rust_result: "[u32; 2]", - dialect_operands: &["i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32"], - dialect_results: &["i32", "i32"], - llvm_arguments: &["i32", "i32", "i32", "i32", "i32", "i32", "v2f16", "v2f16"], - llvm_results: &["v2f16", "v2f16"], - adapter: RegisterMmaAdapter::C2U32A4U32B2U32ToD2U32, - ptx_register_counts: [2, 4, 2, 2], - summary_accumulator: "F16", - }), - RegisterMmaAccumulator::F32 => Ok(RegisterMmaF8F6F4Contract { - accumulator, - scalar_name: "f32", - rust_arguments: ["[f32; 4]", "[u32; 4]", "[u32; 2]"], - rust_result: "[f32; 4]", - dialect_operands: &[ - "f32", "f32", "f32", "f32", "i32", "i32", "i32", "i32", "i32", "i32", - ], - dialect_results: &["f32", "f32", "f32", "f32"], - llvm_arguments: &[ - "i32", "i32", "i32", "i32", "i32", "i32", "f32", "f32", "f32", "f32", - ], - llvm_results: &["f32", "f32", "f32", "f32"], - adapter: RegisterMmaAdapter::C4F32A4U32B2U32ToD4F32, - ptx_register_counts: [4, 4, 2, 4], - summary_accumulator: "F32", - }), - _ => bail!("unsupported dense f8f6f4 MMA accumulator {accumulator:?}"), - } -} - -fn expand_register_mma_f8f6f4_admission( - admission: &RegisterMmaF8F6F4Admission, - accumulator: RegisterMmaAccumulator, -) -> Result> { - let contract = register_mma_f8f6f4_contract(accumulator)?; - ensure!( - admission.runtime_validation == RuntimeValidation::Unexecuted, - "dense f8f6f4 MMA runtime validation may be marked executed only with GPU evidence" - ); - ensure!( - !admission.llvm_evidence_profile.trim().is_empty() - && !admission.libnvvm_evidence_profile.trim().is_empty(), - "dense f8f6f4 MMA admission requires both backend evidence profiles" - ); - ensure!( - admission.a_elements == REGISTER_MMA_F8F6F4_ELEMENTS - && admission.b_elements == REGISTER_MMA_F8F6F4_ELEMENTS - && admission.product_count == 25, - "dense f8f6f4 MMA admission must contain the canonical 5 by 5 element matrix" - ); - let expected_targets = ["sm_120a", "sm_120f", "sm_121a", "sm_121f"]; - ensure!( - admission - .targets - .iter() - .map(String::as_str) - .eq(expected_targets), - "dense f8f6f4 MMA admission must retain the reviewed Blackwell target set" - ); - - let mut records = Vec::with_capacity(admission.product_count); - for &a_element in &admission.a_elements { - for &b_element in &admission.b_elements { - let a = register_mma_f8f6f4_element_name(a_element) - .expect("validated dense f8f6f4 A element"); - let b = register_mma_f8f6f4_element_name(b_element) - .expect("validated dense f8f6f4 B element"); - let scalar = contract.scalar_name; - let id = format!("mma_m16n8k32_{scalar}_{a}_{b}"); - let source_record = - format!("int_nvvm_mma_m16n8k32_row_col_kind_f8f6f4_{scalar}_{a}_{b}_{scalar}"); - let llvm_symbol = - format!("llvm.nvvm.mma.m16n8k32.row.col.kind.f8f6f4.{scalar}.{a}.{b}.{scalar}"); - let mma = RegisterMma { - shape: RegisterMmaShape::M16n8k32, - operation: RegisterMmaOperation::Multiply, - kind: None, - accumulator: contract.accumulator, - a_element, - b_element, - a_layout: RegisterMmaLayout::Row, - b_layout: RegisterMmaLayout::Col, - overflow: RegisterMmaOverflow::NotApplicable, - participation: - RegisterMmaParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes, - adapter: contract.adapter, - compatibility_source: RegisterMmaCompatibilitySource::GeneratedStub, - runtime_validation: admission.runtime_validation, - }; - - records.push(OverlayIntrinsic { - id: id.clone(), - abi_id: String::new(), - operation_key: format!( - "matrix.mma.m16n8k32.row.col.kind_f8f6f4.{scalar}.{a}.{b}.{scalar}" - ), - family: "register_mma".into(), - source: None, - source_record: Some(source_record), - rust_module: "matrix".into(), - rust_name: id.clone(), - rust_arguments: contract.rust_arguments.map(Into::into).into(), - rust_result: contract.rust_result.into(), - safe: false, - must_use: true, - safe_allowlist_reason: None, - public_rust_path: format!("cuda_intrinsics::matrix::{id}"), - compatibility_rust_paths: vec![format!("cuda_device::wmma::{id}")], - dialect_op_type: "RegisterMmaOp".into(), - dialect_op_name: "nvvm.register_mma".into(), - dialect_operands: contract - .dialect_operands - .iter() - .map(|value| (*value).into()) - .collect(), - dialect_results: contract - .dialect_results - .iter() - .map(|value| (*value).into()) - .collect(), - llvm_symbol: Some(llvm_symbol), - resolved_llvm_symbol: None, - llvm_arguments: contract - .llvm_arguments - .iter() - .map(|value| (*value).into()) - .collect(), - llvm_results: contract - .llvm_results - .iter() - .map(|value| (*value).into()) - .collect(), - pure: false, - memory: "none".into(), - convergent: true, - execution_scope: "warp".into(), - minimum_ptx: "8.7".into(), - minimum_sm: None, - ptx_result: contract.rust_result.into(), - targets: REGISTER_MMA_F8F6F4_TARGETS.into(), - ptx_isa_version: "9.3".into(), - ptx_isa_section: "9.7.15.5.14 Multiply-and-Accumulate Instruction: mma".into(), - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma".into(), - lowering: "generated_register_mma".into(), - backend_lowerings: vec![ - OverlayBackendLowering { - backend: IntrinsicBackend::LlvmNvptx, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: admission.llvm_evidence_profile.clone(), - targets: None, - minimum_ptx: None, - minimum_sm: None, - }, - OverlayBackendLowering { - backend: IntrinsicBackend::LibNvvm, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: admission.libnvvm_evidence_profile.clone(), - targets: None, - minimum_ptx: None, - minimum_sm: None, - }, - ], - packed_atomic: None, - redux: None, - vote: None, - active_mask: None, - warp_match: None, - warp_barrier: None, - warp_shuffle: None, - dot_product: None, - packed_alu: None, - integer_minmax: None, - packed_conversion: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cp_async_copy: None, - cp_async_control: None, - cp_async_mbarrier: None, - mbarrier_basic: None, - movmatrix: None, - mbarrier_extended: None, - register_mma: Some(mma), - sparse_mma: None, - prmt: None, - cluster_barrier: None, - wgmma_control: None, - special_register: None, - debug_control: None, - cluster_memory: None, - clc: None, - tma: None, - tcgen05: None, - ldmatrix_variant: None, - ldmatrix_safety: None, - ldmatrix_adapter: None, - selected_address_space: None, - expected_ptx: InstructionPattern { - mnemonic: "mma".into(), - modifiers: [ - "sync", - "aligned", - "m16n8k32", - "row", - "col", - "kind::f8f6f4", - scalar, - a, - b, - scalar, - ] - .into_iter() - .map(Into::into) - .collect(), - operands: contract - .ptx_register_counts - .map(|length| OperandPattern::RegisterList { length }) - .into(), - }, - summary: format!( - "Multiplies warp-distributed {a} A and {b} B fragments and adds an {} accumulator.", - contract.summary_accumulator - ), - }); - } - } - Ok(records) -} - -fn expand_register_mma_mxf8f6f4_admission( - admission: &RegisterMmaF8F6F4Admission, -) -> Result> { - ensure!( - admission.runtime_validation == RuntimeValidation::Unexecuted, - "dense mxf8f6f4 MMA runtime validation may be marked executed only with GPU evidence" - ); - ensure!( - !admission.llvm_evidence_profile.trim().is_empty() - && !admission.libnvvm_evidence_profile.trim().is_empty(), - "dense mxf8f6f4 MMA admission requires both backend evidence profiles" - ); - ensure!( - admission.a_elements == REGISTER_MMA_F8F6F4_ELEMENTS - && admission.b_elements == REGISTER_MMA_F8F6F4_ELEMENTS - && admission.product_count == 25, - "dense mxf8f6f4 MMA admission must contain the canonical 5 by 5 element matrix" - ); - let expected_targets = ["sm_120a", "sm_120f", "sm_121a", "sm_121f"]; - ensure!( - admission - .targets - .iter() - .map(String::as_str) - .eq(expected_targets), - "dense mxf8f6f4 MMA admission must retain the reviewed Blackwell target set" - ); - - let rust_arguments = [ - "[f32; 4]", "[u32; 4]", "[u32; 2]", "u32", "u16", "u16", "u32", "u16", "u16", - ]; - let dialect_operands = [ - "f32", "f32", "f32", "f32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i16", "i16", - "i32", "i16", "i16", - ]; - let llvm_arguments = [ - "i32", "i32", "i32", "i32", "i32", "i32", "f32", "f32", "f32", "f32", "i32", "i16", "i16", - "i32", "i16", "i16", - ]; - let dialect_results = ["f32", "f32", "f32", "f32"]; - - let mut records = Vec::with_capacity(admission.product_count); - for &a_element in &admission.a_elements { - for &b_element in &admission.b_elements { - let a = register_mma_f8f6f4_element_name(a_element) - .expect("validated dense mxf8f6f4 A element"); - let b = register_mma_f8f6f4_element_name(b_element) - .expect("validated dense mxf8f6f4 B element"); - let id = format!("mma_m16n8k32_mxf8f6f4_f32_{a}_{b}"); - let source_record = - format!("int_nvvm_mma_block_scale_m16n8k32_row_col_mxf8f6f4_f32_{a}_{b}_f32_ue8m0"); - let llvm_symbol = format!( - "llvm.nvvm.mma.block.scale.m16n8k32.row.col.mxf8f6f4.f32.{a}.{b}.f32.ue8m0" - ); - let mma = RegisterMma { - shape: RegisterMmaShape::M16n8k32, - operation: RegisterMmaOperation::Multiply, - kind: Some(RegisterMmaKind::Mxf8f6f4), - accumulator: RegisterMmaAccumulator::F32, - a_element, - b_element, - a_layout: RegisterMmaLayout::Row, - b_layout: RegisterMmaLayout::Col, - overflow: RegisterMmaOverflow::NotApplicable, - participation: - RegisterMmaParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes, - adapter: RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32, - compatibility_source: RegisterMmaCompatibilitySource::GeneratedStub, - runtime_validation: admission.runtime_validation, - }; - - records.push(OverlayIntrinsic { - id: id.clone(), - abi_id: String::new(), - operation_key: format!( - "matrix.mma.m16n8k32.row.col.kind_mxf8f6f4.scale_vec_1x.f32.{a}.{b}.f32.ue8m0" - ), - family: "register_mma".into(), - source: None, - source_record: Some(source_record), - rust_module: "matrix".into(), - rust_name: id.clone(), - rust_arguments: rust_arguments.map(Into::into).into(), - rust_result: "[f32; 4]".into(), - safe: false, - must_use: true, - safe_allowlist_reason: None, - public_rust_path: format!("cuda_intrinsics::matrix::{id}"), - compatibility_rust_paths: vec![format!("cuda_device::wmma::{id}")], - dialect_op_type: "RegisterMmaOp".into(), - dialect_op_name: "nvvm.register_mma".into(), - dialect_operands: dialect_operands.map(Into::into).into(), - dialect_results: dialect_results.map(Into::into).into(), - llvm_symbol: Some(llvm_symbol), - resolved_llvm_symbol: None, - llvm_arguments: llvm_arguments.map(Into::into).into(), - llvm_results: dialect_results.map(Into::into).into(), - pure: false, - memory: "none".into(), - convergent: true, - execution_scope: "warp".into(), - minimum_ptx: "8.7".into(), - minimum_sm: None, - ptx_result: "[f32; 4]".into(), - targets: REGISTER_MMA_F8F6F4_TARGETS.into(), - ptx_isa_version: "9.3".into(), - ptx_isa_section: - "9.7.15.5.14 Multiply-and-Accumulate Instruction: mma".into(), - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma".into(), - lowering: "generated_register_mma".into(), - backend_lowerings: vec![ - OverlayBackendLowering { - backend: IntrinsicBackend::LlvmNvptx, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: admission.llvm_evidence_profile.clone(), - targets: None, - minimum_ptx: None, - minimum_sm: None, - }, - OverlayBackendLowering { - backend: IntrinsicBackend::LibNvvm, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: admission.libnvvm_evidence_profile.clone(), - targets: None, - minimum_ptx: None, - minimum_sm: None, - }, - ], - packed_atomic: None, - redux: None, - vote: None, - active_mask: None, - warp_match: None, - warp_barrier: None, - warp_shuffle: None, - dot_product: None, - packed_alu: None, - integer_minmax: None, - packed_conversion: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cp_async_copy: None, - cp_async_control: None, - cp_async_mbarrier: None, - mbarrier_basic: None, - movmatrix: None, - mbarrier_extended: None, - register_mma: Some(mma), - sparse_mma: None, - prmt: None, - cluster_barrier: None, - wgmma_control: None, - special_register: None, - debug_control: None, - cluster_memory: None, - clc: None, - tma: None, - tcgen05: None, - ldmatrix_variant: None, - ldmatrix_safety: None, - ldmatrix_adapter: None, - selected_address_space: None, - expected_ptx: InstructionPattern { - mnemonic: "mma".into(), - modifiers: [ - "sync", - "aligned", - "m16n8k32", - "row", - "col", - "kind::mxf8f6f4", - "block_scale", - "f32", - a, - b, - "f32", - "ue8m0", - ] - .into_iter() - .map(Into::into) - .collect(), - operands: vec![ - OperandPattern::RegisterList { length: 4 }, - OperandPattern::RegisterList { length: 4 }, - OperandPattern::RegisterList { length: 2 }, - OperandPattern::RegisterList { length: 4 }, - OperandPattern::Register, - OperandPattern::RegisterList { length: 2 }, - OperandPattern::Register, - OperandPattern::RegisterList { length: 2 }, - ], - }, - summary: format!( - "Multiplies block-scaled warp-distributed {a} A and {b} B fragments and adds an F32 accumulator." - ), - }); - } - } - Ok(records) -} - -const REGISTER_MMA_FP8_SHAPES: [RegisterMmaShape; 2] = - [RegisterMmaShape::M16n8k16, RegisterMmaShape::M16n8k32]; -const REGISTER_MMA_FP8_ACCUMULATORS: [RegisterMmaAccumulator; 2] = - [RegisterMmaAccumulator::F16, RegisterMmaAccumulator::F32]; -const REGISTER_MMA_FP8_ELEMENTS: [RegisterMmaElement; 2] = - [RegisterMmaElement::E4m3, RegisterMmaElement::E5m2]; - -fn register_mma_fp8_element_name(element: RegisterMmaElement) -> Option<&'static str> { - match element { - RegisterMmaElement::E4m3 => Some("e4m3"), - RegisterMmaElement::E5m2 => Some("e5m2"), - _ => None, - } -} - -fn register_mma_fp8_shape_contract( - shape: RegisterMmaShape, -) -> Result<(&'static str, usize, usize)> { - match shape { - RegisterMmaShape::M16n8k16 => Ok(("m16n8k16", 2, 1)), - RegisterMmaShape::M16n8k32 => Ok(("m16n8k32", 4, 2)), - _ => bail!("unsupported standard FP8 register-MMA shape {shape:?}"), - } -} - -fn register_mma_fp8_minimum_ptx( - shape: RegisterMmaShape, - accumulator: RegisterMmaAccumulator, -) -> &'static str { - // ptxas requires PTX 8.7 for the k32 F16 form despite LLVM's 8.4 predicate. - match (shape, accumulator) { - (RegisterMmaShape::M16n8k32, RegisterMmaAccumulator::F32) => "8.4", - _ => "8.7", - } -} - -fn expand_register_mma_fp8_admission( - admission: &RegisterMmaFp8Admission, -) -> Result> { - ensure!( - admission.runtime_validation == RuntimeValidation::Unexecuted, - "standard FP8 MMA runtime validation may be marked executed only with GPU evidence" - ); - ensure!( - !admission.llvm_evidence_profile.trim().is_empty() - && !admission.libnvvm_evidence_profile.trim().is_empty(), - "standard FP8 MMA admission requires both backend evidence profiles" - ); - ensure!( - admission.shapes == REGISTER_MMA_FP8_SHAPES - && admission.accumulators == REGISTER_MMA_FP8_ACCUMULATORS - && admission.a_elements == REGISTER_MMA_FP8_ELEMENTS - && admission.b_elements == REGISTER_MMA_FP8_ELEMENTS - && admission.product_count == 16, - "standard FP8 MMA admission must retain the canonical 2 by 2 by 2 by 2 matrix" - ); - - let mut records = Vec::with_capacity(admission.product_count); - for &shape in &admission.shapes { - let (shape_name, a_count, b_count) = register_mma_fp8_shape_contract(shape)?; - for &accumulator in &admission.accumulators { - let minimum_ptx = register_mma_fp8_minimum_ptx(shape, accumulator); - let (scalar, rust_arguments, rust_result, dialect_operands, dialect_results, adapter) = - match accumulator { - RegisterMmaAccumulator::F16 => ( - "f16", - vec![ - "[u32; 2]".into(), - format!("[u32; {a_count}]"), - if b_count == 1 { - "u32".into() - } else { - "[u32; 2]".into() - }, - ], - "[u32; 2]".to_owned(), - [ - vec!["i32".into(); 2], - vec!["i32".into(); a_count], - vec!["i32".into(); b_count], - ] - .concat(), - vec!["i32".into(); 2], - if shape == RegisterMmaShape::M16n8k16 { - RegisterMmaAdapter::C2U32A2U32B1U32ToD2U32 - } else { - RegisterMmaAdapter::C2U32A4U32B2U32ToD2U32 - }, - ), - RegisterMmaAccumulator::F32 => ( - "f32", - vec![ - "[f32; 4]".into(), - format!("[u32; {a_count}]"), - if b_count == 1 { - "u32".into() - } else { - "[u32; 2]".into() - }, - ], - "[f32; 4]".to_owned(), - [ - vec!["f32".into(); 4], - vec!["i32".into(); a_count], - vec!["i32".into(); b_count], - ] - .concat(), - vec!["f32".into(); 4], - if shape == RegisterMmaShape::M16n8k16 { - RegisterMmaAdapter::C4F32A2U32B1U32ToD4F32 - } else { - RegisterMmaAdapter::C4F32A4U32B2U32ToD4F32 - }, - ), - _ => unreachable!("validated standard FP8 accumulator"), - }; - let llvm_arguments = [ - vec!["i32".into(); a_count + b_count], - if accumulator == RegisterMmaAccumulator::F16 { - vec!["v2f16".into(); 2] - } else { - vec!["f32".into(); 4] - }, - ] - .concat(); - let llvm_results = if accumulator == RegisterMmaAccumulator::F16 { - vec!["v2f16".into(); 2] - } else { - vec!["f32".into(); 4] - }; - let result_count = dialect_results.len(); - - for &a_element in &admission.a_elements { - for &b_element in &admission.b_elements { - let a = register_mma_fp8_element_name(a_element) - .expect("validated standard FP8 A element"); - let b = register_mma_fp8_element_name(b_element) - .expect("validated standard FP8 B element"); - let id = format!("mma_{shape_name}_fp8_{scalar}_{a}_{b}"); - let source_record = - format!("int_nvvm_mma_{shape_name}_row_col_{scalar}_{a}_{b}_{scalar}"); - let llvm_symbol = - format!("llvm.nvvm.mma.{shape_name}.row.col.{scalar}.{a}.{b}.{scalar}"); - let mma = RegisterMma { - shape, - operation: RegisterMmaOperation::Multiply, - kind: Some(RegisterMmaKind::Standard), - accumulator, - a_element, - b_element, - a_layout: RegisterMmaLayout::Row, - b_layout: RegisterMmaLayout::Col, - overflow: RegisterMmaOverflow::NotApplicable, - participation: RegisterMmaParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes, - adapter, - compatibility_source: RegisterMmaCompatibilitySource::GeneratedStub, - runtime_validation: admission.runtime_validation, - }; - records.push(OverlayIntrinsic { - id: id.clone(), - abi_id: String::new(), - operation_key: format!( - "matrix.mma.{shape_name}.row.col.standard_fp8.{scalar}.{a}.{b}.{scalar}" - ), - family: "register_mma".into(), - source: None, - source_record: Some(source_record), - rust_module: "matrix".into(), - rust_name: id.clone(), - rust_arguments: rust_arguments.clone(), - rust_result: rust_result.clone(), - safe: false, - must_use: true, - safe_allowlist_reason: None, - public_rust_path: format!("cuda_intrinsics::matrix::{id}"), - compatibility_rust_paths: vec![format!("cuda_device::wmma::{id}")], - dialect_op_type: "RegisterMmaOp".into(), - dialect_op_name: "nvvm.register_mma".into(), - dialect_operands: dialect_operands.clone(), - dialect_results: dialect_results.clone(), - llvm_symbol: Some(llvm_symbol), - resolved_llvm_symbol: None, - llvm_arguments: llvm_arguments.clone(), - llvm_results: llvm_results.clone(), - pure: false, - memory: "none".into(), - convergent: true, - execution_scope: "warp".into(), - minimum_ptx: minimum_ptx.into(), - minimum_sm: Some("sm_89".into()), - ptx_result: rust_result.clone(), - targets: "all".into(), - ptx_isa_version: "9.3".into(), - ptx_isa_section: "9.7.15.5.14 Multiply-and-Accumulate Instruction: mma".into(), - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma".into(), - lowering: "generated_register_mma".into(), - backend_lowerings: [ - (IntrinsicBackend::LlvmNvptx, &admission.llvm_evidence_profile), - (IntrinsicBackend::LibNvvm, &admission.libnvvm_evidence_profile), - ] - .into_iter() - .map(|(backend, evidence_profile)| OverlayBackendLowering { - backend, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: evidence_profile.clone(), - targets: None, - minimum_ptx: Some(minimum_ptx.into()), - minimum_sm: Some("sm_89".into()), - }) - .collect(), - packed_atomic: None, - redux: None, - vote: None, - active_mask: None, - warp_match: None, - warp_barrier: None, - warp_shuffle: None, - dot_product: None, - packed_alu: None, - integer_minmax: None, - packed_conversion: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cp_async_copy: None, - cp_async_control: None, - cp_async_mbarrier: None, - mbarrier_basic: None, - movmatrix: None, - mbarrier_extended: None, - register_mma: Some(mma), - sparse_mma: None, - prmt: None, - cluster_barrier: None, - wgmma_control: None, - special_register: None, - debug_control: None, - cluster_memory: None, - clc: None, - tma: None, - tcgen05: None, - ldmatrix_variant: None, - ldmatrix_safety: None, - ldmatrix_adapter: None, - selected_address_space: None, - expected_ptx: InstructionPattern { - mnemonic: "mma".into(), - modifiers: ["sync", "aligned", shape_name, "row", "col", scalar, a, b, scalar] - .into_iter() - .map(Into::into) - .collect(), - operands: [result_count, a_count, b_count, result_count] - .map(|length| OperandPattern::RegisterList { length }) - .into(), - }, - summary: format!( - "Multiplies warp-distributed {a} A and {b} B fragments and adds an {scalar} accumulator." - ), - }); - } - } - } - } - ensure!(records.len() == admission.product_count); - Ok(records) -} - -const REGISTER_MMA_AMPERE_FLOAT_VARIANTS: [RegisterMmaAmpereFloatVariant; 5] = [ - RegisterMmaAmpereFloatVariant { - shape: RegisterMmaShape::M16n8k4, - accumulator: RegisterMmaAccumulator::F32, - element: RegisterMmaElement::Tf32, - }, - RegisterMmaAmpereFloatVariant { - shape: RegisterMmaShape::M16n8k8, - accumulator: RegisterMmaAccumulator::F16, - element: RegisterMmaElement::F16, - }, - RegisterMmaAmpereFloatVariant { - shape: RegisterMmaShape::M16n8k8, - accumulator: RegisterMmaAccumulator::F32, - element: RegisterMmaElement::Bf16, - }, - RegisterMmaAmpereFloatVariant { - shape: RegisterMmaShape::M16n8k8, - accumulator: RegisterMmaAccumulator::F32, - element: RegisterMmaElement::F16, - }, - RegisterMmaAmpereFloatVariant { - shape: RegisterMmaShape::M16n8k16, - accumulator: RegisterMmaAccumulator::F16, - element: RegisterMmaElement::F16, - }, -]; - -fn expand_register_mma_ampere_float_admission( - admission: &RegisterMmaAmpereFloatAdmission, -) -> Result> { - ensure!( - admission.runtime_validation == RuntimeValidation::Unexecuted, - "Ampere floating-point MMA runtime validation requires GPU evidence" - ); - ensure!( - !admission.llvm_evidence_profile.trim().is_empty() - && !admission.libnvvm_evidence_profile.trim().is_empty(), - "Ampere floating-point MMA admission requires both backend evidence profiles" - ); - ensure!( - admission.product_count == REGISTER_MMA_AMPERE_FLOAT_VARIANTS.len() - && admission.variants == REGISTER_MMA_AMPERE_FLOAT_VARIANTS, - "Ampere floating-point MMA admission must retain the five reviewed variants in canonical order" - ); - - admission - .variants - .iter() - .map(|variant| { - let adapter = match (variant.shape, variant.accumulator) { - (RegisterMmaShape::M16n8k4, RegisterMmaAccumulator::F32) - | (RegisterMmaShape::M16n8k8, RegisterMmaAccumulator::F32) => { - RegisterMmaAdapter::C4F32A2U32B1U32ToD4F32 - } - (RegisterMmaShape::M16n8k8, RegisterMmaAccumulator::F16) => { - RegisterMmaAdapter::C2U32A2U32B1U32ToD2U32 - } - (RegisterMmaShape::M16n8k16, RegisterMmaAccumulator::F16) => { - RegisterMmaAdapter::C2U32A4U32B2U32ToD2U32 - } - _ => bail!("Ampere floating-point MMA admission contains an unsupported carrier"), - }; - let mma = RegisterMma { - shape: variant.shape, - operation: RegisterMmaOperation::Multiply, - kind: None, - accumulator: variant.accumulator, - a_element: variant.element, - b_element: variant.element, - a_layout: RegisterMmaLayout::Row, - b_layout: RegisterMmaLayout::Col, - overflow: RegisterMmaOverflow::NotApplicable, - participation: RegisterMmaParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes, - adapter, - compatibility_source: RegisterMmaCompatibilitySource::GeneratedStub, - runtime_validation: admission.runtime_validation, - }; - let recipe = register_mma_recipe(&mma).with_context(|| { - "Ampere floating-point MMA admission requests a variant outside the closed recipe set" - })?; - Ok(OverlayIntrinsic { - id: recipe.id.into(), - abi_id: String::new(), - operation_key: recipe.operation_key.into(), - family: "register_mma".into(), - source: None, - source_record: Some(recipe.source_record.into()), - rust_module: "matrix".into(), - rust_name: recipe.id.into(), - rust_arguments: recipe.rust_arguments.iter().map(|value| (*value).into()).collect(), - rust_result: recipe.rust_result.into(), - safe: false, - must_use: true, - safe_allowlist_reason: None, - public_rust_path: format!("cuda_intrinsics::matrix::{}", recipe.id), - compatibility_rust_paths: vec![format!("cuda_device::wmma::{}", recipe.id)], - dialect_op_type: recipe.dialect_op_type.into(), - dialect_op_name: recipe.dialect_op_name.into(), - dialect_operands: recipe.dialect_operands.iter().map(|value| (*value).into()).collect(), - dialect_results: recipe.dialect_results.iter().map(|value| (*value).into()).collect(), - llvm_symbol: Some(recipe.llvm_symbol.into()), - resolved_llvm_symbol: None, - llvm_arguments: recipe.llvm_arguments.iter().map(|value| (*value).into()).collect(), - llvm_results: recipe.llvm_results.iter().map(|value| (*value).into()).collect(), - pure: false, - memory: "none".into(), - convergent: true, - execution_scope: "warp".into(), - minimum_ptx: recipe.minimum_ptx.into(), - minimum_sm: Some(recipe.minimum_sm.into()), - ptx_result: recipe.rust_result.into(), - targets: "all".into(), - ptx_isa_version: "9.3".into(), - ptx_isa_section: "9.7.15.5.14 Multiply-and-Accumulate Instruction: mma".into(), - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma".into(), - lowering: "generated_register_mma".into(), - backend_lowerings: [ - (IntrinsicBackend::LlvmNvptx, &admission.llvm_evidence_profile), - (IntrinsicBackend::LibNvvm, &admission.libnvvm_evidence_profile), - ] - .into_iter() - .map(|(backend, evidence_profile)| OverlayBackendLowering { - backend, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: evidence_profile.clone(), - targets: None, - minimum_ptx: Some(recipe.minimum_ptx.into()), - minimum_sm: Some(recipe.minimum_sm.into()), - }) - .collect(), - packed_atomic: None, - redux: None, - vote: None, - active_mask: None, - warp_match: None, - warp_barrier: None, - warp_shuffle: None, - dot_product: None, - packed_alu: None, - integer_minmax: None, - packed_conversion: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cp_async_copy: None, - cp_async_control: None, - cp_async_mbarrier: None, - mbarrier_basic: None, - movmatrix: None, - mbarrier_extended: None, - register_mma: Some(mma), - sparse_mma: None, - prmt: None, - cluster_barrier: None, - wgmma_control: None, - special_register: None, - debug_control: None, - cluster_memory: None, - clc: None, - tma: None, - tcgen05: None, - ldmatrix_variant: None, - ldmatrix_safety: None, - ldmatrix_adapter: None, - selected_address_space: None, - expected_ptx: InstructionPattern { - mnemonic: "mma".into(), - modifiers: recipe.ptx_modifiers.iter().map(|value| (*value).into()).collect(), - operands: recipe - .ptx_register_counts - .map(|length| OperandPattern::RegisterList { length }) - .into(), - }, - summary: format!("Executes the {} Ampere floating-point MMA form.", recipe.id), - }) - }) - .collect() -} - -fn expand_register_mma_binary_admission( - admission: &RegisterMmaBinaryAdmission, -) -> Result> { - use RegisterMmaAdapter::{ - C2I32A1U32B1U32ToD2I32, C4I32A2U32B1U32ToD4I32, C4I32A4U32B2U32ToD4I32, - }; - use RegisterMmaLayout::{Col, Row}; - use RegisterMmaOperation::{AndPopc, XorPopc}; - use RegisterMmaParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes; - use RegisterMmaShape::{M8n8k128, M16n8k128, M16n8k256}; - - ensure!( - !admission.variants.is_empty(), - "compact binary MMA admission has no variants" - ); - ensure!( - admission.runtime_validation == RuntimeValidation::Unexecuted, - "binary MMA runtime validation may be marked executed only with GPU evidence" - ); - - let mut seen = BTreeSet::new(); - let mut records = Vec::with_capacity(admission.variants.len()); - for variant in &admission.variants { - ensure!( - seen.insert((variant.shape, variant.operation)), - "compact binary MMA admission contains a duplicate variant" - ); - ensure!( - matches!(variant.operation, AndPopc | XorPopc), - "compact binary MMA admission contains a non-binary operation" - ); - let adapter = match variant.shape { - M8n8k128 => C2I32A1U32B1U32ToD2I32, - M16n8k128 => C4I32A2U32B1U32ToD4I32, - M16n8k256 => C4I32A4U32B2U32ToD4I32, - _ => bail!("compact binary MMA admission contains an unsupported shape"), - }; - let mma = RegisterMma { - shape: variant.shape, - operation: variant.operation, - kind: None, - accumulator: RegisterMmaAccumulator::S32, - a_element: RegisterMmaElement::B1, - b_element: RegisterMmaElement::B1, - a_layout: Row, - b_layout: Col, - overflow: RegisterMmaOverflow::Wrapping, - participation: AllWarpLanesSameInstructionAndQualifiersNoExitedLanes, - adapter, - compatibility_source: RegisterMmaCompatibilitySource::GeneratedStub, - runtime_validation: admission.runtime_validation, - }; - let recipe = register_mma_recipe(&mma).with_context( - || "compact binary MMA admission requests a variant outside the closed recipe set", - )?; - let operation = match variant.operation { - AndPopc => "AND and population count", - XorPopc => "XOR and population count", - RegisterMmaOperation::Multiply => unreachable!(), - }; - let summary = format!( - "Computes warp-distributed binary matrix products with {operation}, then adds a wrapping s32 accumulator." - ); - - records.push(OverlayIntrinsic { - id: recipe.id.into(), - abi_id: recipe.abi_id.into(), - operation_key: recipe.operation_key.into(), - family: "register_mma".into(), - source: None, - source_record: Some(recipe.source_record.into()), - rust_module: "matrix".into(), - rust_name: recipe.id.into(), - rust_arguments: recipe - .rust_arguments - .iter() - .map(|value| (*value).into()) - .collect(), - rust_result: recipe.rust_result.into(), - safe: false, - must_use: true, - safe_allowlist_reason: None, - public_rust_path: format!("cuda_intrinsics::matrix::{}", recipe.id), - compatibility_rust_paths: vec![format!("cuda_device::wmma::{}", recipe.id)], - dialect_op_type: recipe.dialect_op_type.into(), - dialect_op_name: recipe.dialect_op_name.into(), - dialect_operands: recipe - .dialect_operands - .iter() - .map(|value| (*value).into()) - .collect(), - dialect_results: recipe - .dialect_results - .iter() - .map(|value| (*value).into()) - .collect(), - llvm_symbol: Some(recipe.llvm_symbol.into()), - resolved_llvm_symbol: None, - llvm_arguments: recipe - .llvm_arguments - .iter() - .map(|value| (*value).into()) - .collect(), - llvm_results: recipe - .llvm_results - .iter() - .map(|value| (*value).into()) - .collect(), - pure: false, - memory: "none".into(), - convergent: true, - execution_scope: "warp".into(), - minimum_ptx: recipe.minimum_ptx.into(), - minimum_sm: Some(recipe.minimum_sm.into()), - ptx_result: recipe.rust_result.into(), - targets: "all".into(), - ptx_isa_version: "9.3".into(), - ptx_isa_section: "9.7.15.5.14 Multiply-and-Accumulate Instruction: mma".into(), - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma".into(), - lowering: "generated_register_mma".into(), - backend_lowerings: vec![ - OverlayBackendLowering { - backend: IntrinsicBackend::LlvmNvptx, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: admission.llvm_evidence_profile.clone(), - targets: None, - minimum_ptx: Some(recipe.minimum_ptx.into()), - minimum_sm: Some(recipe.minimum_sm.into()), - }, - OverlayBackendLowering { - backend: IntrinsicBackend::LibNvvm, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: admission.libnvvm_evidence_profile.clone(), - targets: None, - minimum_ptx: Some(recipe.minimum_ptx.into()), - minimum_sm: Some(recipe.minimum_sm.into()), - }, - ], - packed_atomic: None, - redux: None, - vote: None, - active_mask: None, - warp_match: None, - warp_barrier: None, - warp_shuffle: None, - dot_product: None, - packed_alu: None, - integer_minmax: None, - packed_conversion: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cp_async_copy: None, - cp_async_control: None, - cp_async_mbarrier: None, - mbarrier_basic: None, - movmatrix: None, - mbarrier_extended: None, - register_mma: Some(mma), - sparse_mma: None, - prmt: None, - cluster_barrier: None, - wgmma_control: None, - special_register: None, - debug_control: None, - cluster_memory: None, - clc: None, - tma: None, - tcgen05: None, - ldmatrix_variant: None, - ldmatrix_safety: None, - ldmatrix_adapter: None, - selected_address_space: None, - expected_ptx: InstructionPattern { - mnemonic: "mma".into(), - modifiers: recipe - .ptx_modifiers - .iter() - .map(|value| (*value).into()) - .collect(), - operands: recipe - .ptx_register_counts - .map(|length| OperandPattern::RegisterList { length }) - .into(), - }, - summary, - }); - } - Ok(records) -} - -#[derive(Clone, Copy)] -struct SparseMmaCarrierRecipe { - adapter: SparseMmaAdapter, - llvm_adapter: SparseMmaLlvmAdapter, - accumulator: SparseMmaAccumulator, - selector: SparseMmaSelector, - a_registers: usize, - b_registers: usize, -} - -impl SparseMmaCarrierRecipe { - fn accumulator_registers(self) -> usize { - match self.accumulator { - SparseMmaAccumulator::F16 => 2, - SparseMmaAccumulator::F32 | SparseMmaAccumulator::S32 => 4, - } - } - - fn operand_count(self) -> usize { - self.accumulator_registers() + self.a_registers + self.b_registers + 2 - } - - fn selector_index(self) -> usize { - self.operand_count() - 1 - } - - fn selector_upper_exclusive(self) -> u8 { - match self.selector { - SparseMmaSelector::ImmediateZeroOrOne => 2, - SparseMmaSelector::ImmediateZero => 1, - } - } - - fn rust_arguments(self) -> Vec { - let accumulator = match self.accumulator { - SparseMmaAccumulator::F16 => "[u32; 2]", - SparseMmaAccumulator::F32 => "[f32; 4]", - SparseMmaAccumulator::S32 => "[i32; 4]", - }; - vec![ - accumulator.into(), - format!("[u32; {}]", self.a_registers), - format!("[u32; {}]", self.b_registers), - "u32".into(), - "u32".into(), - ] - } - - fn dialect_operands(self) -> Vec { - let accumulator = match self.accumulator { - SparseMmaAccumulator::F16 => "i32", - SparseMmaAccumulator::F32 => "f32", - SparseMmaAccumulator::S32 => "i32", - }; - let accumulator_registers = self.accumulator_registers(); - std::iter::repeat_n(accumulator.to_owned(), accumulator_registers) - .chain(std::iter::repeat_n( - "u32".to_owned(), - self.operand_count() - accumulator_registers, - )) - .collect() - } - - fn llvm_arguments(self) -> Vec { - let accumulator = match self.accumulator { - SparseMmaAccumulator::F16 => "v2f16", - SparseMmaAccumulator::F32 => "f32", - SparseMmaAccumulator::S32 => "i32", - }; - std::iter::repeat_n("i32".to_owned(), self.a_registers + self.b_registers) - .chain(std::iter::repeat_n( - accumulator.to_owned(), - self.accumulator_registers(), - )) - .chain(std::iter::repeat_n("i32".to_owned(), 2)) - .collect() - } - - fn rust_result(self) -> String { - match self.accumulator { - SparseMmaAccumulator::F16 => "[u32; 2]", - SparseMmaAccumulator::F32 => "[f32; 4]", - SparseMmaAccumulator::S32 => "[i32; 4]", - } - .into() - } - - fn dialect_results(self) -> Vec { - match self.accumulator { - SparseMmaAccumulator::F16 => vec!["i32".into(); 2], - SparseMmaAccumulator::F32 => vec!["f32".into(); 4], - SparseMmaAccumulator::S32 => vec!["i32".into(); 4], - } - } - - fn llvm_results(self) -> Vec { - match self.accumulator { - SparseMmaAccumulator::F16 => vec!["v2f16".into(); 2], - SparseMmaAccumulator::F32 => vec!["f32".into(); 4], - SparseMmaAccumulator::S32 => vec!["i32".into(); 4], - } - } - - fn expected_ptx_operands(self) -> Vec { - vec![ - OperandPattern::RegisterList { - length: self.accumulator_registers(), - }, - OperandPattern::RegisterList { - length: self.a_registers, - }, - OperandPattern::RegisterList { - length: self.b_registers, - }, - OperandPattern::RegisterList { - length: self.accumulator_registers(), - }, - OperandPattern::Register, - OperandPattern::Immediate, - ] - } - - fn imported_properties(self) -> Vec { - let selector = self.selector_index(); - vec![ - format!("ImmArg"), - "IntrNoCallback".into(), - "IntrNoMem".into(), - format!("Range", self.selector_upper_exclusive()), - ] - } -} - -fn sparse_mma_carrier_recipe( - shape: SparseMmaShape, - a_element: SparseMmaElement, - b_element: SparseMmaElement, -) -> Option { - use SparseMmaElement::{E2m1, E2m3, E3m2, E4m3, E5m2, S4, S8, U4, U8}; - use SparseMmaShape::{M16n8k32, M16n8k64, M16n8k128}; - - match (shape, a_element, b_element) { - (M16n8k32, S8 | U8, S8 | U8) | (M16n8k64, S4 | U4, S4 | U4) => { - Some(SparseMmaCarrierRecipe { - adapter: SparseMmaAdapter::C4I32A2U32B2U32MetadataU32SelectorU32ToD4I32, - llvm_adapter: SparseMmaLlvmAdapter::A2I32B2I32C4I32MetadataI32SelectorI32ToD4I32, - accumulator: SparseMmaAccumulator::S32, - selector: SparseMmaSelector::ImmediateZeroOrOne, - a_registers: 2, - b_registers: 2, - }) - } - (M16n8k64, S8 | U8, S8 | U8) => Some(SparseMmaCarrierRecipe { - adapter: SparseMmaAdapter::C4I32A4U32B4U32MetadataU32SelectorU32ToD4I32, - llvm_adapter: SparseMmaLlvmAdapter::A4I32B4I32C4I32MetadataI32SelectorI32ToD4I32, - accumulator: SparseMmaAccumulator::S32, - selector: SparseMmaSelector::ImmediateZero, - a_registers: 4, - b_registers: 4, - }), - (M16n8k128, S4 | U4, S4 | U4) => Some(SparseMmaCarrierRecipe { - adapter: SparseMmaAdapter::C4I32A4U32B4U32MetadataU32SelectorU32ToD4I32, - llvm_adapter: SparseMmaLlvmAdapter::A4I32B4I32C4I32MetadataI32SelectorI32ToD4I32, - accumulator: SparseMmaAccumulator::S32, - selector: SparseMmaSelector::ImmediateZero, - a_registers: 4, - b_registers: 4, - }), - (M16n8k64, E2m1 | E2m3 | E3m2 | E4m3 | E5m2, E2m1 | E2m3 | E3m2 | E4m3 | E5m2) => { - Some(SparseMmaCarrierRecipe { - adapter: SparseMmaAdapter::C4F32A4U32B4U32MetadataU32SelectorU32ToD4F32, - llvm_adapter: SparseMmaLlvmAdapter::A4I32B4I32C4F32MetadataI32SelectorI32ToD4F32, - accumulator: SparseMmaAccumulator::F32, - selector: SparseMmaSelector::ImmediateZero, - a_registers: 4, - b_registers: 4, - }) - } - _ => None, - } -} - -fn sparse_mma_f8f6f4_f16_carrier_recipe( - shape: SparseMmaShape, - a_element: SparseMmaElement, - b_element: SparseMmaElement, -) -> Option { - use SparseMmaElement::{E2m1, E2m3, E3m2, E4m3, E5m2}; - - if shape != SparseMmaShape::M16n8k64 - || !matches!(a_element, E2m1 | E2m3 | E3m2 | E4m3 | E5m2) - || !matches!(b_element, E2m1 | E2m3 | E3m2 | E4m3 | E5m2) - { - return None; - } - Some(SparseMmaCarrierRecipe { - adapter: SparseMmaAdapter::C2U32A4U32B4U32MetadataU32SelectorU32ToD2U32, - llvm_adapter: SparseMmaLlvmAdapter::A4I32B4I32C2V2F16MetadataI32SelectorI32ToD2V2F16, - accumulator: SparseMmaAccumulator::F16, - selector: SparseMmaSelector::ImmediateZero, - a_registers: 4, - b_registers: 4, - }) -} - -struct SparseMmaIdentity { - id: String, - operation_key: String, - source_record: String, - llvm_symbol: String, - ptx_modifiers: Vec<&'static str>, -} - -struct SparseMmaRecipe { - identity: SparseMmaIdentity, - carrier: SparseMmaCarrierRecipe, -} - -fn sparse_mma_shape_name(shape: SparseMmaShape) -> &'static str { - match shape { - SparseMmaShape::M16n8k32 => "m16n8k32", - SparseMmaShape::M16n8k64 => "m16n8k64", - SparseMmaShape::M16n8k128 => "m16n8k128", - } -} - -fn sparse_mma_element_name(element: SparseMmaElement) -> &'static str { - match element { - SparseMmaElement::E2m1 => "e2m1", - SparseMmaElement::E2m3 => "e2m3", - SparseMmaElement::E3m2 => "e3m2", - SparseMmaElement::E4m3 => "e4m3", - SparseMmaElement::E5m2 => "e5m2", - SparseMmaElement::S4 => "s4", - SparseMmaElement::U4 => "u4", - SparseMmaElement::S8 => "s8", - SparseMmaElement::U8 => "u8", - } -} - -fn sparse_mma_identity(mma: &SparseMma) -> SparseMmaIdentity { - let shape = sparse_mma_shape_name(mma.shape); - let a_element = sparse_mma_element_name(mma.a_element); - let b_element = sparse_mma_element_name(mma.b_element); - if matches!( - mma.accumulator, - SparseMmaAccumulator::F16 | SparseMmaAccumulator::F32 - ) { - let scalar = match mma.accumulator { - SparseMmaAccumulator::F16 => "f16", - SparseMmaAccumulator::F32 => "f32", - SparseMmaAccumulator::S32 => unreachable!(), - }; - return SparseMmaIdentity { - id: format!( - "mma_sp_ordered_metadata_{shape}_kind_f8f6f4_{scalar}_{a_element}_{b_element}_{scalar}" - ), - operation_key: format!( - "matrix.mma.sp.{shape}.row.col.kind_f8f6f4.{scalar}.{a_element}.{b_element}.{scalar}.not_applicable.ordered_metadata" - ), - source_record: format!( - "int_nvvm_mma_sp_ordered_metadata_{shape}_row_col_kind_f8f6f4_{scalar}_{a_element}_{b_element}_{scalar}" - ), - llvm_symbol: format!( - "llvm.nvvm.mma.sp.ordered.metadata.{shape}.row.col.kind.f8f6f4.{scalar}.{a_element}.{b_element}.{scalar}" - ), - ptx_modifiers: vec![ - "sp::ordered_metadata", - "sync", - "aligned", - shape, - "row", - "col", - "kind::f8f6f4", - scalar, - a_element, - b_element, - scalar, - ], - }; - } - let compact_elements = if mma.a_element == mma.b_element { - a_element.to_owned() - } else { - format!("{a_element}_{b_element}") - }; - let dotted_elements = if mma.a_element == mma.b_element { - a_element.to_owned() - } else { - format!("{a_element}.{b_element}") - }; - let metadata_id_prefix = match mma.metadata { - SparseMmaMetadata::Standard => "", - SparseMmaMetadata::Ordered => "ordered_metadata_", - }; - let metadata_source_prefix = metadata_id_prefix; - let metadata_symbol_prefix = match mma.metadata { - SparseMmaMetadata::Standard => "", - SparseMmaMetadata::Ordered => "ordered.metadata.", - }; - let metadata_key = match mma.metadata { - SparseMmaMetadata::Standard => "standard_metadata", - SparseMmaMetadata::Ordered => "ordered_metadata", - }; - let overflow_id_suffix = match mma.overflow { - SparseMmaOverflow::NotApplicable => unreachable!("integer sparse MMA has overflow"), - SparseMmaOverflow::Wrapping => "", - SparseMmaOverflow::Satfinite => "_satfinite", - }; - let overflow_source_prefix = match mma.overflow { - SparseMmaOverflow::NotApplicable => unreachable!("integer sparse MMA has overflow"), - SparseMmaOverflow::Wrapping => "", - SparseMmaOverflow::Satfinite => "satfinite_", - }; - let overflow_symbol_prefix = match mma.overflow { - SparseMmaOverflow::NotApplicable => unreachable!("integer sparse MMA has overflow"), - SparseMmaOverflow::Wrapping => "", - SparseMmaOverflow::Satfinite => "satfinite.", - }; - let overflow_key = match mma.overflow { - SparseMmaOverflow::NotApplicable => unreachable!("integer sparse MMA has overflow"), - SparseMmaOverflow::Wrapping => "wrapping", - SparseMmaOverflow::Satfinite => "satfinite", - }; - - let mut ptx_modifiers = vec![ - match mma.metadata { - SparseMmaMetadata::Standard => "sp", - SparseMmaMetadata::Ordered => "sp::ordered_metadata", - }, - "sync", - "aligned", - shape, - "row", - "col", - ]; - if mma.overflow == SparseMmaOverflow::Satfinite { - ptx_modifiers.push("satfinite"); - } - ptx_modifiers.extend(["s32", a_element, b_element, "s32"]); - - SparseMmaIdentity { - id: format!( - "mma_sp_{metadata_id_prefix}{shape}_s32_{compact_elements}{overflow_id_suffix}" - ), - operation_key: format!( - "matrix.mma.sp.{shape}.row.col.s32.{a_element}.{b_element}.s32.{overflow_key}.{metadata_key}" - ), - source_record: format!( - "int_nvvm_mma_sp_{metadata_source_prefix}{shape}_row_col_{overflow_source_prefix}{compact_elements}" - ), - llvm_symbol: format!( - "llvm.nvvm.mma.sp.{metadata_symbol_prefix}{shape}.row.col.{overflow_symbol_prefix}{dotted_elements}" - ), - ptx_modifiers, - } -} - -fn sparse_mma_recipe(mma: &SparseMma) -> Option { - let carrier = if mma.accumulator == SparseMmaAccumulator::F16 { - sparse_mma_f8f6f4_f16_carrier_recipe(mma.shape, mma.a_element, mma.b_element)? - } else { - sparse_mma_carrier_recipe(mma.shape, mma.a_element, mma.b_element)? - }; - - let scalar_contract = match carrier.accumulator { - SparseMmaAccumulator::F16 | SparseMmaAccumulator::F32 => { - mma.accumulator == carrier.accumulator - && mma.overflow == SparseMmaOverflow::NotApplicable - && mma.metadata == SparseMmaMetadata::Ordered - } - SparseMmaAccumulator::S32 => { - mma.accumulator == SparseMmaAccumulator::S32 - && matches!( - mma.overflow, - SparseMmaOverflow::Wrapping | SparseMmaOverflow::Satfinite - ) - } - }; - if !scalar_contract - || mma.a_layout != SparseMmaLayout::Row - || mma.b_layout != SparseMmaLayout::Col - || mma.selector != carrier.selector - || mma.participation - != SparseMmaParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes - || mma.adapter != carrier.adapter - || mma.llvm_adapter != carrier.llvm_adapter - || mma.compatibility_source != SparseMmaCompatibilitySource::GeneratedStub - { - return None; - } - - Some(SparseMmaRecipe { - identity: sparse_mma_identity(mma), - carrier, - }) -} - -fn sparse_mma_minimum_ptx(mma: &SparseMma) -> &'static str { - if matches!( - mma.accumulator, - SparseMmaAccumulator::F16 | SparseMmaAccumulator::F32 - ) { - return "8.7"; - } - match mma.metadata { - SparseMmaMetadata::Standard => "7.1", - SparseMmaMetadata::Ordered => "8.5", - } -} - -fn sparse_mma_hardware(mma: &SparseMma) -> (&'static str, Option<&'static str>) { - match mma.accumulator { - SparseMmaAccumulator::F16 | SparseMmaAccumulator::F32 => (SPARSE_MMA_F8F6F4_TARGETS, None), - SparseMmaAccumulator::S32 => ("all", Some("sm_80")), - } -} - -fn sparse_mma_ptx_section(_: SparseMmaMetadata) -> &'static str { - "9.7.15.6.3 Multiply-and-Accumulate Instruction: mma.sp" -} - -fn expand_sparse_mma_integer_admission( - admission: &SparseMmaIntegerAdmission, -) -> Result> { - ensure!( - !admission.variants.is_empty(), - "compact sparse integer MMA admission has no variants" - ); - ensure!( - admission.runtime_validation == RuntimeValidation::Unexecuted, - "sparse integer MMA runtime validation may be marked executed only with GPU evidence" - ); - - let mut seen = BTreeSet::new(); - let mut records = Vec::with_capacity(admission.variants.len()); - for variant in &admission.variants { - ensure!( - seen.insert(( - variant.shape, - variant.a_element, - variant.b_element, - variant.overflow, - )), - "compact sparse integer MMA admission contains a duplicate variant" - ); - let carrier = sparse_mma_carrier_recipe( - variant.shape, - variant.a_element, - variant.b_element, - ) - .with_context( - || "compact sparse integer MMA admission uses unsupported or mixed-width elements", - )?; - let mma = SparseMma { - shape: variant.shape, - accumulator: SparseMmaAccumulator::S32, - a_element: variant.a_element, - b_element: variant.b_element, - a_layout: SparseMmaLayout::Row, - b_layout: SparseMmaLayout::Col, - overflow: variant.overflow, - metadata: admission.metadata, - selector: carrier.selector, - participation: - SparseMmaParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes, - adapter: carrier.adapter, - llvm_adapter: carrier.llvm_adapter, - compatibility_source: SparseMmaCompatibilitySource::GeneratedStub, - runtime_validation: admission.runtime_validation, - }; - let recipe = sparse_mma_recipe(&mma).with_context( - || "compact sparse integer MMA admission requests a variant outside the closed recipe set", - )?; - let signedness = |element| match element { - SparseMmaElement::S4 => "signed", - SparseMmaElement::U4 => "unsigned", - SparseMmaElement::S8 => "signed", - SparseMmaElement::U8 => "unsigned", - _ => unreachable!("integer admission rejects floating formats"), - }; - let width = match (variant.a_element, variant.b_element) { - ( - SparseMmaElement::S4 | SparseMmaElement::U4, - SparseMmaElement::S4 | SparseMmaElement::U4, - ) => "INT4", - ( - SparseMmaElement::S8 | SparseMmaElement::U8, - SparseMmaElement::S8 | SparseMmaElement::U8, - ) => "INT8", - _ => unreachable!("carrier selection rejects mixed element widths"), - }; - let overflow = match variant.overflow { - SparseMmaOverflow::NotApplicable => { - unreachable!("integer admission rejects inapplicable overflow") - } - SparseMmaOverflow::Wrapping => "wrapping", - SparseMmaOverflow::Satfinite => "saturating", - }; - let metadata = match admission.metadata { - SparseMmaMetadata::Standard => "", - SparseMmaMetadata::Ordered => " with ordered sparsity metadata", - }; - let summary = format!( - "Multiplies warp-distributed sparse {} A and {} B {width} fragments{metadata} and adds a {overflow} s32 accumulator.", - signedness(variant.a_element), - signedness(variant.b_element), - ); - records.push(sparse_mma_overlay_record( - String::new(), - mma, - recipe, - &admission.llvm_evidence_profile, - &admission.libnvvm_evidence_profile, - summary, - )); - } - Ok(records) -} - -fn expand_sparse_mma_f8f6f4_admission( - admission: &SparseMmaF8F6F4Admission, -) -> Result> { - ensure!( - admission.runtime_validation == RuntimeValidation::Unexecuted, - "sparse f8f6f4 MMA runtime validation may be marked executed only with GPU evidence" - ); - - let formats = vec![ - SparseMmaElement::E2m1, - SparseMmaElement::E2m3, - SparseMmaElement::E3m2, - SparseMmaElement::E4m3, - SparseMmaElement::E5m2, - ]; - ensure!( - admission.a_elements == formats, - "compact sparse f8f6f4 MMA admission must list the canonical five A formats" - ); - ensure!( - admission.b_elements == formats, - "compact sparse f8f6f4 MMA admission must list the canonical five B formats" - ); - ensure!( - admission.product_count - == admission - .a_elements - .len() - .checked_mul(admission.b_elements.len()) - .context("compact sparse f8f6f4 MMA admission product count overflow")? - && admission.product_count == 25, - "compact sparse f8f6f4 MMA admission product_count must be exactly 25" - ); - let mut records = Vec::with_capacity(admission.product_count); - for &a_element in &admission.a_elements { - for &b_element in &admission.b_elements { - let carrier = sparse_mma_carrier_recipe(SparseMmaShape::M16n8k64, a_element, b_element) - .with_context( - || "compact sparse f8f6f4 MMA admission uses an unsupported format", - )?; - ensure!( - carrier.accumulator == SparseMmaAccumulator::F32, - "compact sparse f8f6f4 MMA admission contains an integer format" - ); - let mma = SparseMma { - shape: SparseMmaShape::M16n8k64, - accumulator: SparseMmaAccumulator::F32, - a_element, - b_element, - a_layout: SparseMmaLayout::Row, - b_layout: SparseMmaLayout::Col, - overflow: SparseMmaOverflow::NotApplicable, - metadata: SparseMmaMetadata::Ordered, - selector: carrier.selector, - participation: - SparseMmaParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes, - adapter: carrier.adapter, - llvm_adapter: carrier.llvm_adapter, - compatibility_source: SparseMmaCompatibilitySource::GeneratedStub, - runtime_validation: admission.runtime_validation, - }; - let recipe = sparse_mma_recipe(&mma).with_context(|| { - "compact sparse f8f6f4 MMA admission requests a variant outside the closed recipe set" - })?; - let summary = format!( - "Multiplies warp-distributed sparse {} A and {} B fragments and adds an f32 accumulator.", - sparse_mma_element_name(a_element), - sparse_mma_element_name(b_element), - ); - records.push(sparse_mma_overlay_record( - String::new(), - mma, - recipe, - &admission.llvm_evidence_profile, - &admission.libnvvm_evidence_profile, - summary, - )); - } - } - ensure!(records.len() == admission.product_count); - Ok(records) -} - -const SPARSE_MMA_F8F6F4_ELEMENTS: [SparseMmaElement; 5] = [ - SparseMmaElement::E2m1, - SparseMmaElement::E2m3, - SparseMmaElement::E3m2, - SparseMmaElement::E4m3, - SparseMmaElement::E5m2, -]; - -fn expand_sparse_mma_f8f6f4_f16_admission( - admission: &SparseMmaF8F6F4F16Admission, -) -> Result> { - ensure!( - admission.runtime_validation == RuntimeValidation::Unexecuted, - "sparse f8f6f4 F16 MMA runtime validation may be marked executed only with GPU evidence" - ); - ensure!( - !admission.llvm_evidence_profile.trim().is_empty() - && !admission.libnvvm_evidence_profile.trim().is_empty(), - "sparse f8f6f4 F16 MMA admission requires both backend evidence profiles" - ); - ensure!( - admission.a_elements == SPARSE_MMA_F8F6F4_ELEMENTS - && admission.b_elements == SPARSE_MMA_F8F6F4_ELEMENTS - && admission.product_count == 25, - "sparse f8f6f4 F16 MMA admission must contain the canonical 5 by 5 element matrix" - ); - - let mut records = Vec::with_capacity(admission.product_count); - for &a_element in &admission.a_elements { - for &b_element in &admission.b_elements { - let carrier = sparse_mma_f8f6f4_f16_carrier_recipe( - SparseMmaShape::M16n8k64, - a_element, - b_element, - ) - .context("compact sparse f8f6f4 F16 MMA admission uses an unsupported format")?; - let mma = SparseMma { - shape: SparseMmaShape::M16n8k64, - accumulator: SparseMmaAccumulator::F16, - a_element, - b_element, - a_layout: SparseMmaLayout::Row, - b_layout: SparseMmaLayout::Col, - overflow: SparseMmaOverflow::NotApplicable, - metadata: SparseMmaMetadata::Ordered, - selector: SparseMmaSelector::ImmediateZero, - participation: - SparseMmaParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes, - adapter: carrier.adapter, - llvm_adapter: carrier.llvm_adapter, - compatibility_source: SparseMmaCompatibilitySource::GeneratedStub, - runtime_validation: admission.runtime_validation, - }; - let recipe = sparse_mma_recipe(&mma).context( - "compact sparse f8f6f4 F16 MMA admission requests a variant outside the closed recipe set", - )?; - let summary = format!( - "Multiplies warp-distributed sparse {} A and {} B fragments and adds a packed F16 accumulator.", - sparse_mma_element_name(a_element), - sparse_mma_element_name(b_element), - ); - records.push(sparse_mma_overlay_record( - String::new(), - mma, - recipe, - &admission.llvm_evidence_profile, - &admission.libnvvm_evidence_profile, - summary, - )); - } - } - ensure!(records.len() == admission.product_count); - Ok(records) -} - -fn sparse_mma_overlay_record( - abi_id: String, - mma: SparseMma, - recipe: SparseMmaRecipe, - llvm_evidence_profile: &str, - libnvvm_evidence_profile: &str, - summary: String, -) -> OverlayIntrinsic { - let identity = &recipe.identity; - let minimum_ptx = sparse_mma_minimum_ptx(&mma); - let (targets, minimum_sm) = sparse_mma_hardware(&mma); - OverlayIntrinsic { - id: identity.id.clone(), - abi_id, - operation_key: identity.operation_key.clone(), - family: "sparse_mma".into(), - source: None, - source_record: Some(identity.source_record.clone()), - rust_module: "matrix".into(), - rust_name: identity.id.clone(), - rust_arguments: recipe.carrier.rust_arguments(), - rust_result: recipe.carrier.rust_result(), - safe: false, - must_use: true, - safe_allowlist_reason: None, - public_rust_path: format!("cuda_intrinsics::matrix::{}", identity.id), - compatibility_rust_paths: vec![format!("cuda_device::wmma::{}", identity.id)], - dialect_op_type: "SparseMmaOp".into(), - dialect_op_name: "nvvm.sparse_mma".into(), - dialect_operands: recipe.carrier.dialect_operands(), - dialect_results: recipe.carrier.dialect_results(), - llvm_symbol: Some(identity.llvm_symbol.clone()), - resolved_llvm_symbol: None, - llvm_arguments: recipe.carrier.llvm_arguments(), - llvm_results: recipe.carrier.llvm_results(), - pure: false, - memory: "none".into(), - convergent: true, - execution_scope: "warp".into(), - minimum_ptx: minimum_ptx.into(), - minimum_sm: minimum_sm.map(str::to_owned), - ptx_result: recipe.carrier.rust_result(), - targets: targets.into(), - ptx_isa_version: "9.3".into(), - ptx_isa_section: sparse_mma_ptx_section(mma.metadata).into(), - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma-sp".into(), - lowering: "generated_sparse_mma".into(), - backend_lowerings: [ - (IntrinsicBackend::LlvmNvptx, llvm_evidence_profile), - (IntrinsicBackend::LibNvvm, libnvvm_evidence_profile), - ] - .into_iter() - .map(|(backend, evidence_profile)| OverlayBackendLowering { - backend, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: evidence_profile.into(), - targets: None, - minimum_ptx: Some(minimum_ptx.into()), - minimum_sm: minimum_sm.map(str::to_owned), - }) - .collect(), - packed_atomic: None, - redux: None, - vote: None, - active_mask: None, - warp_match: None, - warp_barrier: None, - warp_shuffle: None, - dot_product: None, - packed_alu: None, - integer_minmax: None, - packed_conversion: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cp_async_copy: None, - cp_async_control: None, - cp_async_mbarrier: None, - mbarrier_basic: None, - movmatrix: None, - mbarrier_extended: None, - register_mma: None, - sparse_mma: Some(mma), - prmt: None, - cluster_barrier: None, - wgmma_control: None, - special_register: None, - debug_control: None, - cluster_memory: None, - clc: None, - tma: None, - tcgen05: None, - ldmatrix_variant: None, - ldmatrix_safety: None, - ldmatrix_adapter: None, - selected_address_space: None, - expected_ptx: InstructionPattern { - mnemonic: "mma".into(), - modifiers: identity - .ptx_modifiers - .iter() - .map(|value| (*value).into()) - .collect(), - operands: recipe.carrier.expected_ptx_operands(), - }, - summary, - } -} - -fn validate_sparse_mma_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, -) -> Result<()> { - let mma = policy - .sparse_mma - .as_ref() - .with_context(|| format!("{} has no closed sparse-MMA contract", policy.id))?; - let recipe = sparse_mma_recipe(mma) - .with_context(|| format!("{} requests an unsupported sparse-MMA variant", policy.id))?; - let identity = &recipe.identity; - let minimum_ptx = sparse_mma_minimum_ptx(mma); - let (targets, minimum_sm) = sparse_mma_hardware(mma); - ensure!( - policy.id == identity.id - && policy.operation_key == identity.operation_key - && policy.source.is_none() - && policy.source_record.as_deref() == Some(identity.source_record.as_str()) - && policy.llvm_symbol.as_deref() == Some(identity.llvm_symbol.as_str()) - && policy.resolved_llvm_symbol.is_none(), - "{} sparse-MMA identity does not match its closed recipe", - policy.id - ); - ensure!( - policy.rust_module == "matrix" - && policy.rust_name == identity.id - && policy.public_rust_path == format!("cuda_intrinsics::matrix::{}", identity.id) - && policy.rust_arguments == recipe.carrier.rust_arguments() - && policy.rust_result == recipe.carrier.rust_result() - && !policy.safe - && policy.must_use - && policy.safe_allowlist_reason.is_none() - && policy.compatibility_rust_paths == [format!("cuda_device::wmma::{}", identity.id)], - "{} must preserve its unsafe must-use Rust sparse-MMA API", - policy.id - ); - ensure!( - policy.dialect_op_type == "SparseMmaOp" - && policy.dialect_op_name == "nvvm.sparse_mma" - && policy.dialect_operands == recipe.carrier.dialect_operands() - && policy.dialect_results == recipe.carrier.dialect_results() - && policy.llvm_arguments == recipe.carrier.llvm_arguments() - && policy.llvm_results == recipe.carrier.llvm_results() - && policy.ptx_result == recipe.carrier.rust_result() - && policy.lowering == "generated_sparse_mma", - "{} sparse-MMA carrier or lowering adapter disagrees with its recipe", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == "none" - && policy.convergent - && policy.execution_scope == "warp" - && policy.minimum_ptx == minimum_ptx - && policy.minimum_sm.as_deref() == minimum_sm - && policy.targets == targets, - "{} sparse-MMA effects or target floor disagree with PTX", - policy.id - ); - ensure!( - policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section == sparse_mma_ptx_section(mma.metadata) - && policy.ptx_isa_url - == "https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma-sp", - "{} sparse-MMA PTX provenance disagrees with the reviewed recipe", - policy.id - ); - ensure!( - declaration.classes == ["SDPatternOperator", "Intrinsic", "NVVM_MMA_SP"] - && declaration.properties == recipe.carrier.imported_properties() - && declaration.selections.len() == 1 - && (!matches!( - mma.accumulator, - SparseMmaAccumulator::F16 | SparseMmaAccumulator::F32 - ) || declaration.selections[0].predicates == ["Subtarget->hasMMABlockScale()"]) - && selection_matches_policy(policy, &declaration.selections[0])?, - "{} imported sparse MMA declaration changed its class, immediate range, properties, or exact selection contract", - policy.id - ); - ensure!( - policy.expected_ptx.mnemonic == "mma" - && policy.expected_ptx.modifiers == identity.ptx_modifiers - && policy.expected_ptx.operands == recipe.carrier.expected_ptx_operands(), - "{} expected PTX does not match its exact sparse-MMA spelling", - policy.id - ); - ensure_exact_inline_ptx_backends( - policy, - [ - (IntrinsicBackend::LlvmNvptx, minimum_ptx, minimum_sm), - (IntrinsicBackend::LibNvvm, minimum_ptx, minimum_sm), - ], - "sparse MMA", - )?; - ensure_no_other_family_contract(policy, "sparse MMA")?; - Ok(()) -} - -fn validate_register_mma_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, -) -> Result<()> { - let mma = policy - .register_mma - .as_ref() - .with_context(|| format!("{} has no closed register-MMA contract", policy.id))?; - match mma.kind { - Some(RegisterMmaKind::Standard) => { - return validate_register_mma_fp8_policy(policy, declaration, mma); - } - Some(RegisterMmaKind::F8f6f4) => { - return validate_register_mma_f8f6f4_policy(policy, declaration, mma); - } - Some(RegisterMmaKind::Mxf8f6f4) => { - return validate_register_mma_mxf8f6f4_policy(policy, declaration, mma); - } - None if register_mma_f8f6f4_element_name(mma.a_element).is_some() - || register_mma_f8f6f4_element_name(mma.b_element).is_some() => - { - return validate_register_mma_f8f6f4_policy(policy, declaration, mma); - } - None => {} - } - let recipe = register_mma_recipe(mma) - .with_context(|| format!("{} requests an unsupported register-MMA variant", policy.id))?; - let abi_matches = matches!( - recipe.id, - "mma_m16n8k4_f32_tf32" - | "mma_m16n8k8_f16_f16" - | "mma_m16n8k8_f32_bf16" - | "mma_m16n8k8_f32_f16" - | "mma_m16n8k16_f16_f16" - ) || policy.abi_id == recipe.abi_id; - ensure!( - policy.id == recipe.id - && abi_matches - && policy.operation_key == recipe.operation_key - && policy.source.is_none() - && policy.source_record.as_deref() == Some(recipe.source_record) - && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) - && policy.resolved_llvm_symbol.is_none(), - "{} register-MMA identity does not match its closed recipe", - policy.id - ); - ensure!( - policy.rust_module == "matrix" - && policy.rust_name == recipe.id - && policy.rust_arguments == recipe.rust_arguments - && policy.rust_result == recipe.rust_result - && !policy.safe - && policy.must_use - && policy.safe_allowlist_reason.is_none() - && policy.compatibility_rust_paths == [format!("cuda_device::wmma::{}", recipe.id)], - "{} must preserve its unsafe must-use Rust MMA API", - policy.id - ); - ensure!( - policy.dialect_op_type == recipe.dialect_op_type - && policy.dialect_op_name == recipe.dialect_op_name - && policy.dialect_operands == recipe.dialect_operands - && policy.dialect_results == recipe.dialect_results - && policy.llvm_arguments == recipe.llvm_arguments - && policy.llvm_results == recipe.llvm_results - && policy.ptx_result == recipe.rust_result - && mma.adapter == recipe.adapter - && mma.compatibility_source == recipe.compatibility_source - && policy.lowering == "generated_register_mma", - "{} register-MMA carrier or lowering adapter disagrees with its recipe", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == "none" - && policy.convergent - && policy.execution_scope == "warp" - && policy.minimum_ptx == recipe.minimum_ptx - && policy.minimum_sm.as_deref() == Some(recipe.minimum_sm) - && policy.targets == "all", - "{} register-MMA effects or target floor disagree with PTX", - policy.id - ); - ensure!( - policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section == "9.7.15.5.14 Multiply-and-Accumulate Instruction: mma" - && policy.ptx_isa_url - == "https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma", - "{} register-MMA PTX provenance disagrees with the reviewed recipe", - policy.id - ); - ensure!( - declaration.classes.iter().any(|class| class == "NVVM_MMA") - && declaration.properties == ["IntrNoCallback", "IntrNoMem"] - && declaration.selections.len() == 1 - && selection_matches_policy(policy, &declaration.selections[0])?, - "{} imported MMA declaration changed its class, properties, or exact selection contract", - policy.id - ); - // These are imported facts. The BF16 selection is underconstrained. - let ampere_float_predicates = match recipe.id { - "mma_m16n8k4_f32_tf32" | "mma_m16n8k16_f16_f16" => Some([ - "Subtarget->getSmVersion() >= 80", - "Subtarget->getPTXVersion() >= 70", - ]), - "mma_m16n8k8_f16_f16" | "mma_m16n8k8_f32_bf16" | "mma_m16n8k8_f32_f16" => Some([ - "Subtarget->getPTXVersion() >= 65", - "Subtarget->getSmVersion() >= 75", - ]), - _ => None, - }; - if let Some(predicates) = ampere_float_predicates { - let selection = &declaration.selections[0]; - ensure!( - declaration.classes == ["SDPatternOperator", "Intrinsic", "NVVM_MMA"] - && selection.predicates == predicates - && selection.constraints.is_empty() - && mma.kind.is_none() - && mma.runtime_validation == RuntimeValidation::Unexecuted, - "{} imported Ampere floating-point MMA contract changed", - policy.id - ); - } - ensure!( - policy.expected_ptx.mnemonic == "mma" - && policy.expected_ptx.modifiers == recipe.ptx_modifiers - && policy.expected_ptx.operands - == recipe - .ptx_register_counts - .map(|length| OperandPattern::RegisterList { length }), - "{} expected PTX does not match its exact register-MMA spelling", - policy.id - ); - ensure_exact_inline_ptx_backends( - policy, - [ - ( - IntrinsicBackend::LlvmNvptx, - recipe.minimum_ptx, - Some(recipe.minimum_sm), - ), - ( - IntrinsicBackend::LibNvvm, - recipe.minimum_ptx, - Some(recipe.minimum_sm), - ), - ], - "register MMA", - )?; - ensure_no_other_family_contract(policy, "register MMA")?; - Ok(()) -} - -fn validate_register_mma_fp8_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, - mma: &RegisterMma, -) -> Result<()> { - ensure!( - REGISTER_MMA_FP8_SHAPES.contains(&mma.shape), - "{} has an unsupported standard FP8 shape", - policy.id - ); - ensure!( - REGISTER_MMA_FP8_ACCUMULATORS.contains(&mma.accumulator), - "{} has an unsupported standard FP8 accumulator", - policy.id - ); - ensure!( - REGISTER_MMA_FP8_ELEMENTS.contains(&mma.a_element), - "{} has an unsupported standard FP8 A element", - policy.id - ); - ensure!( - REGISTER_MMA_FP8_ELEMENTS.contains(&mma.b_element), - "{} has an unsupported standard FP8 B element", - policy.id - ); - let (shape_name, a_count, b_count) = register_mma_fp8_shape_contract(mma.shape)?; - let minimum_ptx = register_mma_fp8_minimum_ptx(mma.shape, mma.accumulator); - let a = register_mma_fp8_element_name(mma.a_element).expect("validated A element"); - let b = register_mma_fp8_element_name(mma.b_element).expect("validated B element"); - let ( - scalar, - rust_arguments, - rust_result, - dialect_operands, - dialect_results, - llvm_arguments, - llvm_results, - adapter, - ) = match mma.accumulator { - RegisterMmaAccumulator::F16 => ( - "f16", - vec![ - "[u32; 2]".into(), - format!("[u32; {a_count}]"), - if b_count == 1 { - "u32".into() - } else { - "[u32; 2]".into() - }, - ], - "[u32; 2]".to_owned(), - [ - vec!["i32".to_owned(); 2], - vec!["i32".to_owned(); a_count], - vec!["i32".to_owned(); b_count], - ] - .concat(), - vec!["i32".to_owned(); 2], - [ - vec!["i32".to_owned(); a_count + b_count], - vec!["v2f16".to_owned(); 2], - ] - .concat(), - vec!["v2f16".to_owned(); 2], - if mma.shape == RegisterMmaShape::M16n8k16 { - RegisterMmaAdapter::C2U32A2U32B1U32ToD2U32 - } else { - RegisterMmaAdapter::C2U32A4U32B2U32ToD2U32 - }, - ), - RegisterMmaAccumulator::F32 => ( - "f32", - vec![ - "[f32; 4]".into(), - format!("[u32; {a_count}]"), - if b_count == 1 { - "u32".into() - } else { - "[u32; 2]".into() - }, - ], - "[f32; 4]".to_owned(), - [ - vec!["f32".to_owned(); 4], - vec!["i32".to_owned(); a_count], - vec!["i32".to_owned(); b_count], - ] - .concat(), - vec!["f32".to_owned(); 4], - [ - vec!["i32".to_owned(); a_count + b_count], - vec!["f32".to_owned(); 4], - ] - .concat(), - vec!["f32".to_owned(); 4], - if mma.shape == RegisterMmaShape::M16n8k16 { - RegisterMmaAdapter::C4F32A2U32B1U32ToD4F32 - } else { - RegisterMmaAdapter::C4F32A4U32B2U32ToD4F32 - }, - ), - _ => unreachable!("validated accumulator"), - }; - let expected_id = format!("mma_{shape_name}_fp8_{scalar}_{a}_{b}"); - let expected_source = format!("int_nvvm_mma_{shape_name}_row_col_{scalar}_{a}_{b}_{scalar}"); - let expected_symbol = format!("llvm.nvvm.mma.{shape_name}.row.col.{scalar}.{a}.{b}.{scalar}"); - - ensure!( - policy.id == expected_id - && policy.operation_key - == format!( - "matrix.mma.{shape_name}.row.col.standard_fp8.{scalar}.{a}.{b}.{scalar}" - ) - && policy.source.is_none() - && policy.source_record.as_deref() == Some(expected_source.as_str()) - && policy.llvm_symbol.as_deref() == Some(expected_symbol.as_str()) - && policy.resolved_llvm_symbol.is_none(), - "{} standard FP8 MMA identity changed", - policy.id - ); - ensure!( - policy.rust_module == "matrix" - && policy.rust_name == expected_id - && policy.rust_arguments == rust_arguments - && policy.rust_result == rust_result - && !policy.safe - && policy.must_use - && policy.safe_allowlist_reason.is_none() - && policy.public_rust_path == format!("cuda_intrinsics::matrix::{expected_id}") - && policy.compatibility_rust_paths == [format!("cuda_device::wmma::{expected_id}")], - "{} must preserve its unsafe must-use standard FP8 API", - policy.id - ); - ensure!( - mma.kind == Some(RegisterMmaKind::Standard) - && mma.operation == RegisterMmaOperation::Multiply - && mma.a_layout == RegisterMmaLayout::Row - && mma.b_layout == RegisterMmaLayout::Col - && mma.overflow == RegisterMmaOverflow::NotApplicable - && mma.participation - == RegisterMmaParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes - && mma.adapter == adapter - && mma.compatibility_source == RegisterMmaCompatibilitySource::GeneratedStub - && mma.runtime_validation == RuntimeValidation::Unexecuted - && policy.dialect_op_type == "RegisterMmaOp" - && policy.dialect_op_name == "nvvm.register_mma" - && policy.dialect_operands == dialect_operands - && policy.dialect_results == dialect_results - && policy.llvm_arguments == llvm_arguments - && policy.llvm_results == llvm_results - && policy.ptx_result == rust_result - && policy.lowering == "generated_register_mma", - "{} standard FP8 MMA carrier or lowering changed", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == "none" - && policy.convergent - && policy.execution_scope == "warp" - && policy.minimum_ptx == minimum_ptx - && policy.minimum_sm.as_deref() == Some("sm_89") - && policy.targets == "all", - "{} standard FP8 effects or target floor changed", - policy.id - ); - ensure!( - policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section == "9.7.15.5.14 Multiply-and-Accumulate Instruction: mma" - && policy.ptx_isa_url - == "https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma", - "{} standard FP8 PTX provenance changed", - policy.id - ); - let [selection] = declaration.selections.as_slice() else { - bail!( - "{} must retain exactly one standard FP8 selection", - policy.id - ); - }; - let expected_predicates = if mma.shape == RegisterMmaShape::M16n8k16 { - [ - "Subtarget->getSmVersion() >= 89", - "Subtarget->getPTXVersion() >= 87", - ] - } else { - [ - "Subtarget->getSmVersion() >= 89", - "Subtarget->getPTXVersion() >= 84", - ] - }; - ensure!( - declaration.classes == ["SDPatternOperator", "Intrinsic", "NVVM_MMA"] - && declaration.properties == ["IntrNoCallback", "IntrNoMem"] - && selection_matches_policy(policy, selection)? - && selection.predicates == expected_predicates - && selection.constraints.is_empty(), - "{} imported standard FP8 declaration or selection changed", - policy.id - ); - ensure!( - policy.expected_ptx - == InstructionPattern { - mnemonic: "mma".into(), - modifiers: [ - "sync", "aligned", shape_name, "row", "col", scalar, a, b, scalar - ] - .into_iter() - .map(Into::into) - .collect(), - operands: [ - dialect_results.len(), - a_count, - b_count, - dialect_results.len() - ] - .map(|length| OperandPattern::RegisterList { length }) - .into(), - }, - "{} expected standard FP8 PTX changed", - policy.id - ); - ensure_exact_inline_ptx_backends( - policy, - [ - (IntrinsicBackend::LlvmNvptx, minimum_ptx, Some("sm_89")), - (IntrinsicBackend::LibNvvm, minimum_ptx, Some("sm_89")), - ], - "standard FP8 register MMA", - )?; - ensure_no_other_family_contract(policy, "standard FP8 register MMA")?; - Ok(()) -} - -fn validate_register_mma_f8f6f4_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, - mma: &RegisterMma, -) -> Result<()> { - let contract = register_mma_f8f6f4_contract(mma.accumulator)?; - let a = register_mma_f8f6f4_element_name(mma.a_element) - .with_context(|| format!("{} has a non-f8f6f4 A format", policy.id))?; - let b = register_mma_f8f6f4_element_name(mma.b_element) - .with_context(|| format!("{} has a non-f8f6f4 B format", policy.id))?; - let scalar = contract.scalar_name; - let expected_id = format!("mma_m16n8k32_{scalar}_{a}_{b}"); - let expected_source = - format!("int_nvvm_mma_m16n8k32_row_col_kind_f8f6f4_{scalar}_{a}_{b}_{scalar}"); - let expected_symbol = - format!("llvm.nvvm.mma.m16n8k32.row.col.kind.f8f6f4.{scalar}.{a}.{b}.{scalar}"); - ensure!( - policy.id == expected_id - && policy.operation_key - == format!("matrix.mma.m16n8k32.row.col.kind_f8f6f4.{scalar}.{a}.{b}.{scalar}") - && policy.source.is_none() - && policy.source_record.as_deref() == Some(expected_source.as_str()) - && policy.llvm_symbol.as_deref() == Some(expected_symbol.as_str()) - && policy.resolved_llvm_symbol.is_none(), - "{} dense f8f6f4 MMA identity changed", - policy.id - ); - ensure!( - policy.rust_module == "matrix" - && policy.rust_name == expected_id - && policy.rust_arguments == contract.rust_arguments - && policy.rust_result == contract.rust_result - && !policy.safe - && policy.must_use - && policy.safe_allowlist_reason.is_none() - && policy.public_rust_path == format!("cuda_intrinsics::matrix::{expected_id}") - && policy.compatibility_rust_paths == [format!("cuda_device::wmma::{expected_id}")], - "{} must preserve the unsafe must-use dense f8f6f4 API", - policy.id - ); - ensure!( - policy.dialect_op_type == "RegisterMmaOp" - && policy.dialect_op_name == "nvvm.register_mma" - && policy.dialect_operands == contract.dialect_operands - && policy.dialect_results == contract.dialect_results - && policy.llvm_arguments == contract.llvm_arguments - && policy.llvm_results == contract.llvm_results - && policy.ptx_result == contract.rust_result - && mma.shape == RegisterMmaShape::M16n8k32 - && matches!(mma.kind, None | Some(RegisterMmaKind::F8f6f4)) - && mma.operation == RegisterMmaOperation::Multiply - && mma.accumulator == contract.accumulator - && mma.a_layout == RegisterMmaLayout::Row - && mma.b_layout == RegisterMmaLayout::Col - && mma.overflow == RegisterMmaOverflow::NotApplicable - && mma.participation - == RegisterMmaParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes - && mma.adapter == contract.adapter - && mma.compatibility_source == RegisterMmaCompatibilitySource::GeneratedStub - && mma.runtime_validation == RuntimeValidation::Unexecuted - && policy.lowering == "generated_register_mma", - "{} dense f8f6f4 carrier or lowering changed", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == "none" - && policy.convergent - && policy.execution_scope == "warp" - && policy.minimum_ptx == "8.7" - && policy.minimum_sm.is_none() - && policy.targets == REGISTER_MMA_F8F6F4_TARGETS, - "{} dense f8f6f4 effects or exact target set changed", - policy.id - ); - ensure!( - policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section == "9.7.15.5.14 Multiply-and-Accumulate Instruction: mma" - && policy.ptx_isa_url - == "https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma", - "{} dense f8f6f4 PTX provenance changed", - policy.id - ); - let [selection] = declaration.selections.as_slice() else { - bail!( - "{} must retain exactly one imported dense f8f6f4 instruction selection", - policy.id - ); - }; - ensure!( - declaration.classes == ["SDPatternOperator", "Intrinsic", "NVVM_MMA"] - && declaration.properties == ["IntrNoCallback", "IntrNoMem"] - && selection_matches_policy(policy, selection)? - && selection.predicates == ["Subtarget->hasMMABlockScale()"] - && selection.constraints.is_empty(), - "{} imported dense f8f6f4 declaration or selection changed", - policy.id - ); - ensure!( - policy.expected_ptx - == InstructionPattern { - mnemonic: "mma".into(), - modifiers: [ - "sync", - "aligned", - "m16n8k32", - "row", - "col", - "kind::f8f6f4", - scalar, - a, - b, - scalar, - ] - .into_iter() - .map(Into::into) - .collect(), - operands: contract - .ptx_register_counts - .map(|length| OperandPattern::RegisterList { length }) - .into(), - }, - "{} expected dense f8f6f4 PTX changed", - policy.id - ); - let backend_pairs: BTreeSet<_> = policy - .backend_lowerings - .iter() - .map(|lowering| (lowering.backend, lowering.mechanism)) - .collect(); - ensure!( - policy.backend_lowerings.len() == 2 - && backend_pairs - == BTreeSet::from([ - ( - IntrinsicBackend::LlvmNvptx, - BackendLoweringMechanism::InlinePtx, - ), - ( - IntrinsicBackend::LibNvvm, - BackendLoweringMechanism::InlinePtx, - ), - ]) - && policy.backend_lowerings.iter().all(|lowering| { - lowering.targets.is_none() - && lowering.minimum_ptx.is_none() - && lowering.minimum_sm.is_none() - && !lowering.evidence_profile.trim().is_empty() - }), - "{} must inherit the exact reviewed target set on both inline-PTX routes", - policy.id - ); - ensure_no_other_family_contract(policy, "dense f8f6f4 register MMA")?; - Ok(()) -} - -fn validate_register_mma_mxf8f6f4_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, - mma: &RegisterMma, -) -> Result<()> { - let a = register_mma_f8f6f4_element_name(mma.a_element) - .with_context(|| format!("{} has a non-mxf8f6f4 A format", policy.id))?; - let b = register_mma_f8f6f4_element_name(mma.b_element) - .with_context(|| format!("{} has a non-mxf8f6f4 B format", policy.id))?; - let expected_id = format!("mma_m16n8k32_mxf8f6f4_f32_{a}_{b}"); - let expected_source = - format!("int_nvvm_mma_block_scale_m16n8k32_row_col_mxf8f6f4_f32_{a}_{b}_f32_ue8m0"); - let expected_symbol = - format!("llvm.nvvm.mma.block.scale.m16n8k32.row.col.mxf8f6f4.f32.{a}.{b}.f32.ue8m0"); - let rust_arguments = [ - "[f32; 4]", "[u32; 4]", "[u32; 2]", "u32", "u16", "u16", "u32", "u16", "u16", - ]; - let dialect_operands = [ - "f32", "f32", "f32", "f32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i16", "i16", - "i32", "i16", "i16", - ]; - let llvm_arguments = [ - "i32", "i32", "i32", "i32", "i32", "i32", "f32", "f32", "f32", "f32", "i32", "i16", "i16", - "i32", "i16", "i16", - ]; - let results = ["f32", "f32", "f32", "f32"]; - ensure!( - policy.id == expected_id - && policy.operation_key - == format!( - "matrix.mma.m16n8k32.row.col.kind_mxf8f6f4.scale_vec_1x.f32.{a}.{b}.f32.ue8m0" - ) - && policy.source.is_none() - && policy.source_record.as_deref() == Some(expected_source.as_str()) - && policy.llvm_symbol.as_deref() == Some(expected_symbol.as_str()) - && policy.resolved_llvm_symbol.is_none(), - "{} dense mxf8f6f4 MMA identity changed", - policy.id - ); - ensure!( - policy.rust_module == "matrix" - && policy.rust_name == expected_id - && policy.rust_arguments == rust_arguments - && policy.rust_result == "[f32; 4]" - && !policy.safe - && policy.must_use - && policy.safe_allowlist_reason.is_none() - && policy.public_rust_path == format!("cuda_intrinsics::matrix::{expected_id}") - && policy.compatibility_rust_paths == [format!("cuda_device::wmma::{expected_id}")], - "{} must preserve the unsafe must-use dense mxf8f6f4 API", - policy.id - ); - ensure!( - policy.dialect_op_type == "RegisterMmaOp" - && policy.dialect_op_name == "nvvm.register_mma" - && policy.dialect_operands == dialect_operands - && policy.dialect_results == results - && policy.llvm_arguments == llvm_arguments - && policy.llvm_results == results - && policy.ptx_result == "[f32; 4]" - && mma.shape == RegisterMmaShape::M16n8k32 - && mma.kind == Some(RegisterMmaKind::Mxf8f6f4) - && mma.operation == RegisterMmaOperation::Multiply - && mma.accumulator == RegisterMmaAccumulator::F32 - && mma.a_layout == RegisterMmaLayout::Row - && mma.b_layout == RegisterMmaLayout::Col - && mma.overflow == RegisterMmaOverflow::NotApplicable - && mma.participation - == RegisterMmaParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes - && mma.adapter == RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 - && mma.compatibility_source == RegisterMmaCompatibilitySource::GeneratedStub - && mma.runtime_validation == RuntimeValidation::Unexecuted - && policy.lowering == "generated_register_mma", - "{} dense mxf8f6f4 carrier or lowering changed", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == "none" - && policy.convergent - && policy.execution_scope == "warp" - && policy.minimum_ptx == "8.7" - && policy.minimum_sm.is_none() - && policy.targets == REGISTER_MMA_F8F6F4_TARGETS, - "{} dense mxf8f6f4 effects or exact target set changed", - policy.id - ); - ensure!( - policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section == "9.7.15.5.14 Multiply-and-Accumulate Instruction: mma" - && policy.ptx_isa_url - == "https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma", - "{} dense mxf8f6f4 PTX provenance changed", - policy.id - ); - let [selection] = declaration.selections.as_slice() else { - bail!( - "{} must retain exactly one imported dense mxf8f6f4 instruction selection", - policy.id - ); - }; - ensure!( - declaration.classes == ["SDPatternOperator", "Intrinsic", "NVVM_MMA_BLOCK_SCALE"] - && declaration.properties == ["IntrNoCallback", "IntrNoMem"] - && selection_matches_policy(policy, selection)? - && selection.predicates == ["Subtarget->hasMMABlockScale()"] - && selection.constraints.is_empty(), - "{} imported dense mxf8f6f4 declaration or selection changed", - policy.id - ); - ensure!( - policy.expected_ptx - == InstructionPattern { - mnemonic: "mma".into(), - modifiers: [ - "sync", - "aligned", - "m16n8k32", - "row", - "col", - "kind::mxf8f6f4", - "block_scale", - "f32", - a, - b, - "f32", - "ue8m0", - ] - .into_iter() - .map(Into::into) - .collect(), - operands: vec![ - OperandPattern::RegisterList { length: 4 }, - OperandPattern::RegisterList { length: 4 }, - OperandPattern::RegisterList { length: 2 }, - OperandPattern::RegisterList { length: 4 }, - OperandPattern::Register, - OperandPattern::RegisterList { length: 2 }, - OperandPattern::Register, - OperandPattern::RegisterList { length: 2 }, - ], - }, - "{} expected dense mxf8f6f4 PTX changed", - policy.id - ); - let backend_pairs: BTreeSet<_> = policy - .backend_lowerings - .iter() - .map(|lowering| (lowering.backend, lowering.mechanism)) - .collect(); - ensure!( - policy.backend_lowerings.len() == 2 - && backend_pairs - == BTreeSet::from([ - ( - IntrinsicBackend::LlvmNvptx, - BackendLoweringMechanism::InlinePtx, - ), - ( - IntrinsicBackend::LibNvvm, - BackendLoweringMechanism::InlinePtx, - ), - ]) - && policy.backend_lowerings.iter().all(|lowering| { - lowering.targets.is_none() - && lowering.minimum_ptx.is_none() - && lowering.minimum_sm.is_none() - && !lowering.evidence_profile.trim().is_empty() - }), - "{} must inherit the exact reviewed target set on both inline-PTX routes", - policy.id - ); - ensure_no_other_family_contract(policy, "dense mxf8f6f4 register MMA")?; - Ok(()) -} - -#[derive(Clone, Copy)] -struct PrmtRecipe { - abi_id: &'static str, - id: &'static str, - operation_key: &'static str, - source_record: &'static str, - llvm_symbol: &'static str, - modifier: Option<&'static str>, - adapter: PrmtAdapter, - summary: &'static str, -} - -fn prmt_recipe(mode: PrmtMode) -> PrmtRecipe { - use PrmtAdapter::{DirectThreeOperands, InsertZeroSecondSource}; - match mode { - PrmtMode::Generic => PrmtRecipe { - abi_id: "i0252", - id: "prmt", - operation_key: "integer.prmt.b32", - source_record: "int_nvvm_prmt", - llvm_symbol: "llvm.nvvm.prmt", - modifier: None, - adapter: DirectThreeOperands, - summary: "Permutes bytes selected from two 32-bit inputs.", - }, - PrmtMode::F4e => PrmtRecipe { - abi_id: "i0253", - id: "prmt_f4e", - operation_key: "integer.prmt.b32.f4e", - source_record: "int_nvvm_prmt_f4e", - llvm_symbol: "llvm.nvvm.prmt.f4e", - modifier: Some("f4e"), - adapter: DirectThreeOperands, - summary: "Permutes bytes with the forward four-byte extract mode.", - }, - PrmtMode::B4e => PrmtRecipe { - abi_id: "i0254", - id: "prmt_b4e", - operation_key: "integer.prmt.b32.b4e", - source_record: "int_nvvm_prmt_b4e", - llvm_symbol: "llvm.nvvm.prmt.b4e", - modifier: Some("b4e"), - adapter: DirectThreeOperands, - summary: "Permutes bytes with the backward four-byte extract mode.", - }, - PrmtMode::Rc8 => PrmtRecipe { - abi_id: "i0255", - id: "prmt_rc8", - operation_key: "integer.prmt.b32.rc8", - source_record: "int_nvvm_prmt_rc8", - llvm_symbol: "llvm.nvvm.prmt.rc8", - modifier: Some("rc8"), - adapter: InsertZeroSecondSource, - summary: "Replicates a selected byte across the 32-bit result.", - }, - PrmtMode::Ecl => PrmtRecipe { - abi_id: "i0256", - id: "prmt_ecl", - operation_key: "integer.prmt.b32.ecl", - source_record: "int_nvvm_prmt_ecl", - llvm_symbol: "llvm.nvvm.prmt.ecl", - modifier: Some("ecl"), - adapter: InsertZeroSecondSource, - summary: "Clamps a byte extract toward the least-significant byte.", - }, - PrmtMode::Ecr => PrmtRecipe { - abi_id: "i0257", - id: "prmt_ecr", - operation_key: "integer.prmt.b32.ecr", - source_record: "int_nvvm_prmt_ecr", - llvm_symbol: "llvm.nvvm.prmt.ecr", - modifier: Some("ecr"), - adapter: InsertZeroSecondSource, - summary: "Clamps a byte extract toward the most-significant byte.", - }, - PrmtMode::Rc16 => PrmtRecipe { - abi_id: "i0258", - id: "prmt_rc16", - operation_key: "integer.prmt.b32.rc16", - source_record: "int_nvvm_prmt_rc16", - llvm_symbol: "llvm.nvvm.prmt.rc16", - modifier: Some("rc16"), - adapter: InsertZeroSecondSource, - summary: "Replicates a selected 16-bit half across the result.", - }, - } -} - -fn expand_prmt_admission(admission: &PrmtAdmission) -> Result> { - ensure!( - admission.runtime_validation == RuntimeValidation::Unexecuted, - "prmt runtime validation may be marked executed only with GPU evidence" - ); - ensure!( - !admission.llvm_evidence_profile.trim().is_empty() - && !admission.libnvvm_evidence_profile.trim().is_empty(), - "compact prmt admission requires both backend evidence profiles" - ); - let expected_modes = BTreeSet::from([ - PrmtMode::Generic, - PrmtMode::F4e, - PrmtMode::B4e, - PrmtMode::Rc8, - PrmtMode::Ecl, - PrmtMode::Ecr, - PrmtMode::Rc16, - ]); - let actual_modes: BTreeSet<_> = admission - .variants - .iter() - .map(|variant| variant.mode) - .collect(); - ensure!( - admission.variants.len() == expected_modes.len() && actual_modes == expected_modes, - "compact prmt admission must contain each of the seven reviewed modes exactly once" - ); - - admission - .variants - .iter() - .map(|variant| { - let recipe = prmt_recipe(variant.mode); - ensure!( - variant.abi_id == recipe.abi_id, - "{} must keep reserved ABI ID {}", - recipe.id, - recipe.abi_id - ); - let three_operands = recipe.adapter == PrmtAdapter::DirectThreeOperands; - let rust_arguments = vec!["u32".into(); if three_operands { 3 } else { 2 }]; - let llvm_arguments = vec!["i32".into(); if three_operands { 3 } else { 2 }]; - let mut modifiers = vec!["b32".into()]; - if let Some(modifier) = recipe.modifier { - modifiers.push(modifier.into()); - } - let operands = if three_operands { - vec![ - OperandPattern::Register, - OperandPattern::Register, - OperandPattern::Register, - OperandPattern::Register, - ] - } else { - vec![ - OperandPattern::Register, - OperandPattern::Register, - OperandPattern::Exact { value: "0".into() }, - OperandPattern::Register, - ] - }; - Ok(OverlayIntrinsic { - id: recipe.id.into(), - abi_id: variant.abi_id.clone(), - operation_key: recipe.operation_key.into(), - family: "prmt".into(), - source: None, - source_record: Some(recipe.source_record.into()), - rust_module: "prmt".into(), - rust_name: recipe.id.into(), - rust_arguments, - rust_result: "u32".into(), - safe: true, - must_use: true, - safe_allowlist_reason: Some( - "it only permutes register bytes and has no caller preconditions.".into(), - ), - public_rust_path: format!("cuda_intrinsics::prmt::{}", recipe.id), - compatibility_rust_paths: vec![format!("cuda_device::prmt::{}", recipe.id)], - dialect_op_type: "PrmtOp".into(), - dialect_op_name: "nvvm.prmt".into(), - dialect_operands: llvm_arguments.clone(), - dialect_results: vec!["i32".into()], - llvm_symbol: Some(recipe.llvm_symbol.into()), - resolved_llvm_symbol: None, - llvm_arguments, - llvm_results: vec!["i32".into()], - pure: true, - memory: "none".into(), - convergent: false, - execution_scope: "thread".into(), - minimum_ptx: "2.0".into(), - minimum_sm: Some("sm_20".into()), - ptx_result: "u32".into(), - targets: "all".into(), - ptx_isa_version: "9.3".into(), - ptx_isa_section: "9.7.9.7 Data Movement and Conversion Instructions: prmt".into(), - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-prmt".into(), - lowering: "generated_prmt".into(), - backend_lowerings: vec![ - OverlayBackendLowering { - backend: IntrinsicBackend::LlvmNvptx, - mechanism: BackendLoweringMechanism::TypedNvvm, - evidence_profile: admission.llvm_evidence_profile.clone(), - targets: None, - minimum_ptx: Some("3.2".into()), - minimum_sm: Some("sm_20".into()), - }, - OverlayBackendLowering { - backend: IntrinsicBackend::LibNvvm, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: admission.libnvvm_evidence_profile.clone(), - targets: None, - minimum_ptx: None, - minimum_sm: Some("sm_75".into()), - }, - ], - packed_atomic: None, - redux: None, - vote: None, - active_mask: None, - warp_match: None, - warp_barrier: None, - warp_shuffle: None, - dot_product: None, - packed_alu: None, - integer_minmax: None, - packed_conversion: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cp_async_copy: None, - cp_async_control: None, - cp_async_mbarrier: None, - mbarrier_basic: None, - movmatrix: None, - mbarrier_extended: None, - register_mma: None, - sparse_mma: None, - prmt: Some(Prmt { - mode: variant.mode, - adapter: recipe.adapter, - }), - cluster_barrier: None, - wgmma_control: None, - special_register: None, - debug_control: None, - cluster_memory: None, - clc: None, - tma: None, - tcgen05: None, - ldmatrix_variant: None, - ldmatrix_safety: None, - ldmatrix_adapter: None, - selected_address_space: None, - expected_ptx: InstructionPattern { - mnemonic: "prmt".into(), - modifiers, - operands, - }, - summary: recipe.summary.into(), - }) - }) - .collect() -} - -fn validate_prmt_policy(policy: &OverlayIntrinsic, declaration: &ImportedIntrinsic) -> Result<()> { - let prmt = policy - .prmt - .as_ref() - .with_context(|| format!("{} has no closed prmt contract", policy.id))?; - let recipe = prmt_recipe(prmt.mode); - let three_operands = recipe.adapter == PrmtAdapter::DirectThreeOperands; - ensure!( - prmt.adapter == recipe.adapter - && policy.id == recipe.id - && policy.abi_id == recipe.abi_id - && policy.operation_key == recipe.operation_key - && policy.source.is_none() - && policy.source_record.as_deref() == Some(recipe.source_record) - && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) - && policy.resolved_llvm_symbol.is_none(), - "{} prmt identity does not match its closed recipe", - policy.id - ); - ensure!( - policy.rust_module == "prmt" - && policy.rust_name == recipe.id - && policy.rust_arguments == vec!["u32"; if three_operands { 3 } else { 2 }] - && policy.rust_result == "u32" - && policy.safe - && policy.must_use - && policy.public_rust_path == format!("cuda_intrinsics::prmt::{}", recipe.id) - && policy.compatibility_rust_paths == [format!("cuda_device::prmt::{}", recipe.id)], - "{} prmt Rust API does not match its closed recipe", - policy.id - ); - ensure!( - policy.dialect_op_type == "PrmtOp" - && policy.dialect_op_name == "nvvm.prmt" - && policy.dialect_operands == vec!["i32"; if three_operands { 3 } else { 2 }] - && policy.dialect_results == ["i32"] - && policy.llvm_arguments == policy.dialect_operands - && policy.llvm_results == ["i32"] - && policy.lowering == "generated_prmt", - "{} prmt carrier or lowering does not match its closed recipe", - policy.id - ); - ensure!( - declaration - .classes - .iter() - .any(|class| class == "NVVMPureIntrinsic") - && declaration.properties == ["IntrNoMem", "IntrSpeculatable"] - && policy.pure - && policy.memory == "none" - && !policy.convergent - && policy.execution_scope == "thread", - "{} prmt effects disagree with the imported declaration", - policy.id - ); - ensure!( - policy.minimum_ptx == "2.0" - && policy.minimum_sm.as_deref() == Some("sm_20") - && policy.targets == "all" - && policy.ptx_result == "u32" - && policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section == "9.7.9.7 Data Movement and Conversion Instructions: prmt" - && policy.ptx_isa_url - == "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-prmt", - "{} prmt target floor or PTX provenance changed", - policy.id - ); - let mut modifiers = vec!["b32"]; - if let Some(modifier) = recipe.modifier { - modifiers.push(modifier); - } - let expected_operands = if three_operands { - vec![ - OperandPattern::Register, - OperandPattern::Register, - OperandPattern::Register, - OperandPattern::Register, - ] - } else { - vec![ - OperandPattern::Register, - OperandPattern::Register, - OperandPattern::Exact { value: "0".into() }, - OperandPattern::Register, - ] - }; - ensure!( - policy.expected_ptx.mnemonic == "prmt" - && policy.expected_ptx.modifiers == modifiers - && policy.expected_ptx.operands == expected_operands, - "{} expected PTX does not match its closed prmt mode", - policy.id - ); - let backend_pairs: BTreeSet<_> = policy - .backend_lowerings - .iter() - .map(|lowering| (lowering.backend, lowering.mechanism)) - .collect(); - ensure!( - policy.backend_lowerings.len() == 2 - && backend_pairs - == BTreeSet::from([ - ( - IntrinsicBackend::LlvmNvptx, - BackendLoweringMechanism::TypedNvvm, - ), - ( - IntrinsicBackend::LibNvvm, - BackendLoweringMechanism::InlinePtx, - ), - ]), - "{} must define exactly the reviewed prmt backend routes", - policy.id - ); - for lowering in &policy.backend_lowerings { - let floor_matches = match lowering.backend { - IntrinsicBackend::LlvmNvptx => { - lowering.minimum_ptx.as_deref() == Some("3.2") - && lowering.minimum_sm.as_deref() == Some("sm_20") - } - IntrinsicBackend::LibNvvm => { - lowering.minimum_ptx.is_none() && lowering.minimum_sm.as_deref() == Some("sm_75") - } - }; - ensure!( - floor_matches && !lowering.evidence_profile.trim().is_empty(), - "{} backend {:?} does not carry its reviewed prmt floor", - policy.id, - lowering.backend - ); - } - ensure_no_other_family_contract(policy, "prmt")?; - Ok(()) -} - -#[derive(Clone, Copy)] -struct ClusterBarrierRecipe { - mode: ClusterBarrierMode, - abi_id: &'static str, - id: &'static str, - suffix: &'static str, - minimum_ptx: &'static str, - ordering: ClusterBarrierOrdering, - aligned: bool, - summary: &'static str, -} - -fn cluster_barrier_recipe(mode: ClusterBarrierMode) -> ClusterBarrierRecipe { - match mode { - ClusterBarrierMode::Arrive => ClusterBarrierRecipe { - mode, - abi_id: "i0277", - id: "barrier_cluster_arrive", - suffix: "arrive", - minimum_ptx: "7.8", - ordering: ClusterBarrierOrdering::Release, - aligned: false, - summary: "Arrives at the cluster barrier with release ordering.", - }, - ClusterBarrierMode::ArriveAligned => ClusterBarrierRecipe { - mode, - abi_id: "i0278", - id: "barrier_cluster_arrive_aligned", - suffix: "arrive.aligned", - minimum_ptx: "7.8", - ordering: ClusterBarrierOrdering::Release, - aligned: true, - summary: "Arrives at the cluster barrier in aligned mode with release ordering.", - }, - ClusterBarrierMode::ArriveRelaxed => ClusterBarrierRecipe { - mode, - abi_id: "i0279", - id: "barrier_cluster_arrive_relaxed", - suffix: "arrive.relaxed", - minimum_ptx: "8.0", - ordering: ClusterBarrierOrdering::Relaxed, - aligned: false, - summary: "Arrives at the cluster barrier without a release guarantee.", - }, - ClusterBarrierMode::ArriveRelaxedAligned => ClusterBarrierRecipe { - mode, - abi_id: "i0280", - id: "barrier_cluster_arrive_relaxed_aligned", - suffix: "arrive.relaxed.aligned", - minimum_ptx: "8.0", - ordering: ClusterBarrierOrdering::Relaxed, - aligned: true, - summary: "Arrives at the cluster barrier in aligned mode without a release guarantee.", - }, - ClusterBarrierMode::Wait => ClusterBarrierRecipe { - mode, - abi_id: "i0281", - id: "barrier_cluster_wait", - suffix: "wait", - minimum_ptx: "7.8", - ordering: ClusterBarrierOrdering::Acquire, - aligned: false, - summary: "Waits at the cluster barrier with acquire ordering.", - }, - ClusterBarrierMode::WaitAligned => ClusterBarrierRecipe { - mode, - abi_id: "i0282", - id: "barrier_cluster_wait_aligned", - suffix: "wait.aligned", - minimum_ptx: "7.8", - ordering: ClusterBarrierOrdering::Acquire, - aligned: true, - summary: "Waits at the cluster barrier in aligned mode with acquire ordering.", - }, - } -} - -fn expand_cluster_barrier_admission( - admission: &ClusterBarrierAdmission, -) -> Result> { - ensure!( - admission.runtime_validation == RuntimeValidation::Unexecuted, - "cluster-barrier runtime validation may be marked executed only with GPU evidence" - ); - ensure!( - !admission.llvm_evidence_profile.trim().is_empty() - && !admission.libnvvm_evidence_profile.trim().is_empty(), - "compact cluster-barrier admission requires both backend evidence profiles" - ); - let expected_modes = BTreeSet::from([ - ClusterBarrierMode::Arrive, - ClusterBarrierMode::ArriveAligned, - ClusterBarrierMode::ArriveRelaxed, - ClusterBarrierMode::ArriveRelaxedAligned, - ClusterBarrierMode::Wait, - ClusterBarrierMode::WaitAligned, - ]); - let actual_modes: BTreeSet<_> = admission - .variants - .iter() - .map(|variant| variant.mode) - .collect(); - ensure!( - admission.variants.len() == expected_modes.len() && actual_modes == expected_modes, - "compact cluster-barrier admission must contain each reviewed mode exactly once" - ); - - admission - .variants - .iter() - .map(|variant| { - let recipe = cluster_barrier_recipe(variant.mode); - ensure!( - variant.abi_id == recipe.abi_id, - "{} must keep reserved ABI ID {}", - recipe.id, - recipe.abi_id - ); - let source_record = format!("int_nvvm_barrier_cluster_{}", recipe.suffix.replace('.', "_")); - let llvm_symbol = format!("llvm.nvvm.barrier.cluster.{}", recipe.suffix); - let modifiers: Vec = recipe.suffix.split('.').map(str::to_owned).collect(); - Ok(OverlayIntrinsic { - id: recipe.id.into(), - abi_id: variant.abi_id.clone(), - operation_key: format!("cluster.barrier.{}", recipe.suffix), - family: "cluster_barrier".into(), - source: None, - source_record: Some(source_record), - rust_module: "cluster".into(), - rust_name: recipe.id.into(), - rust_arguments: vec![], - rust_result: "()".into(), - safe: false, - must_use: false, - safe_allowlist_reason: None, - public_rust_path: format!("cuda_intrinsics::cluster::{}", recipe.id), - compatibility_rust_paths: vec![format!("cuda_device::cluster::{}", recipe.id)], - dialect_op_type: "ClusterBarrierOp".into(), - dialect_op_name: "nvvm.cluster_barrier".into(), - dialect_operands: vec![], - dialect_results: vec![], - llvm_symbol: Some(llvm_symbol), - resolved_llvm_symbol: None, - llvm_arguments: vec![], - llvm_results: vec![], - pure: false, - memory: "read_write".into(), - convergent: true, - execution_scope: "cluster".into(), - minimum_ptx: recipe.minimum_ptx.into(), - minimum_sm: Some("sm_90".into()), - ptx_result: "()".into(), - targets: "all".into(), - ptx_isa_version: "9.3".into(), - ptx_isa_section: - "Parallel Synchronization and Communication Instructions: barrier.cluster" - .into(), - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-barrier-cluster".into(), - lowering: "generated_cluster_barrier".into(), - backend_lowerings: vec![ - OverlayBackendLowering { - backend: IntrinsicBackend::LlvmNvptx, - mechanism: BackendLoweringMechanism::TypedNvvm, - evidence_profile: admission.llvm_evidence_profile.clone(), - targets: None, - minimum_ptx: Some(recipe.minimum_ptx.into()), - minimum_sm: Some("sm_90".into()), - }, - OverlayBackendLowering { - backend: IntrinsicBackend::LibNvvm, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: admission.libnvvm_evidence_profile.clone(), - targets: None, - minimum_ptx: Some(recipe.minimum_ptx.into()), - minimum_sm: Some("sm_90".into()), - }, - ], - packed_atomic: None, - redux: None, - vote: None, - active_mask: None, - warp_match: None, - warp_barrier: None, - warp_shuffle: None, - dot_product: None, - packed_alu: None, - integer_minmax: None, - packed_conversion: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cp_async_copy: None, - cp_async_control: None, - cp_async_mbarrier: None, - mbarrier_basic: None, - movmatrix: None, - mbarrier_extended: None, - register_mma: None, - sparse_mma: None, - prmt: None, - cluster_barrier: Some(ClusterBarrier { - mode: recipe.mode, - ordering: recipe.ordering, - aligned: recipe.aligned, - }), - wgmma_control: None, - special_register: None, - debug_control: None, - cluster_memory: None, - clc: None, - tma: None, - tcgen05: None, - ldmatrix_variant: None, - ldmatrix_safety: None, - ldmatrix_adapter: None, - selected_address_space: None, - expected_ptx: InstructionPattern { - mnemonic: "barrier".into(), - modifiers: std::iter::once("cluster".into()).chain(modifiers).collect(), - operands: vec![], - }, - summary: recipe.summary.into(), - }) - }) - .collect() -} - -fn validate_cluster_barrier_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, -) -> Result<()> { - let barrier = policy - .cluster_barrier - .as_ref() - .with_context(|| format!("{} has no closed cluster-barrier contract", policy.id))?; - let recipe = cluster_barrier_recipe(barrier.mode); - let source_record = format!( - "int_nvvm_barrier_cluster_{}", - recipe.suffix.replace('.', "_") - ); - let llvm_symbol = format!("llvm.nvvm.barrier.cluster.{}", recipe.suffix); - ensure!( - barrier.ordering == recipe.ordering - && barrier.aligned == recipe.aligned - && policy.id == recipe.id - && policy.abi_id == recipe.abi_id - && policy.operation_key == format!("cluster.barrier.{}", recipe.suffix) - && policy.source.is_none() - && policy.source_record.as_deref() == Some(source_record.as_str()) - && policy.llvm_symbol.as_deref() == Some(llvm_symbol.as_str()) - && policy.resolved_llvm_symbol.is_none(), - "{} identity or semantics do not match its closed cluster-barrier recipe", - policy.id - ); - ensure!( - policy.rust_module == "cluster" - && policy.rust_name == recipe.id - && policy.rust_arguments.is_empty() - && policy.rust_result == "()" - && !policy.safe - && !policy.must_use - && policy.public_rust_path == format!("cuda_intrinsics::cluster::{}", recipe.id) - && policy.compatibility_rust_paths == [format!("cuda_device::cluster::{}", recipe.id)], - "{} Rust API does not match its closed cluster-barrier recipe", - policy.id - ); - ensure!( - policy.dialect_op_type == "ClusterBarrierOp" - && policy.dialect_op_name == "nvvm.cluster_barrier" - && policy.dialect_operands.is_empty() - && policy.dialect_results.is_empty() - && policy.llvm_arguments.is_empty() - && policy.llvm_results.is_empty() - && policy.lowering == "generated_cluster_barrier", - "{} carrier or lowering does not match its closed cluster-barrier recipe", - policy.id - ); - ensure!( - declaration.classes == ["SDPatternOperator", "Intrinsic"] - && declaration.properties == ["IntrConvergent", "IntrNoCallback"] - && !policy.pure - && policy.memory == "read_write" - && policy.convergent - && policy.execution_scope == "cluster", - "{} effects disagree with the imported cluster-barrier declaration", - policy.id - ); - ensure!( - policy.minimum_ptx == recipe.minimum_ptx - && policy.minimum_sm.as_deref() == Some("sm_90") - && policy.targets == "all" - && policy.ptx_result == "()" - && policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section - == "Parallel Synchronization and Communication Instructions: barrier.cluster" - && policy.ptx_isa_url - == "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-barrier-cluster", - "{} target floor or PTX provenance changed", - policy.id - ); - let expected_modifiers: Vec<_> = std::iter::once("cluster") - .chain(recipe.suffix.split('.')) - .collect(); - ensure!( - policy.expected_ptx.mnemonic == "barrier" - && policy.expected_ptx.modifiers == expected_modifiers - && policy.expected_ptx.operands.is_empty(), - "{} expected PTX does not match its exact cluster-barrier spelling", - policy.id - ); - ensure!( - (recipe.ordering == ClusterBarrierOrdering::Relaxed) == recipe.suffix.contains(".relaxed") - && recipe.aligned == recipe.suffix.ends_with(".aligned") - && matches!( - (recipe.mode, recipe.ordering), - ( - ClusterBarrierMode::Arrive | ClusterBarrierMode::ArriveAligned, - ClusterBarrierOrdering::Release - ) | ( - ClusterBarrierMode::ArriveRelaxed | ClusterBarrierMode::ArriveRelaxedAligned, - ClusterBarrierOrdering::Relaxed - ) | ( - ClusterBarrierMode::Wait | ClusterBarrierMode::WaitAligned, - ClusterBarrierOrdering::Acquire - ) - ), - "{} cluster-barrier semantic recipe is inconsistent", - policy.id - ); - let backend_pairs: BTreeSet<_> = policy - .backend_lowerings - .iter() - .map(|lowering| (lowering.backend, lowering.mechanism)) - .collect(); - ensure!( - policy.backend_lowerings.len() == 2 - && backend_pairs - == BTreeSet::from([ - ( - IntrinsicBackend::LlvmNvptx, - BackendLoweringMechanism::TypedNvvm, - ), - ( - IntrinsicBackend::LibNvvm, - BackendLoweringMechanism::InlinePtx, - ), - ]) - && policy.backend_lowerings.iter().all(|lowering| { - lowering.minimum_ptx.as_deref() == Some(recipe.minimum_ptx) - && lowering.minimum_sm.as_deref() == Some("sm_90") - && !lowering.evidence_profile.trim().is_empty() - }), - "{} must define exactly the reviewed cluster-barrier backend routes", - policy.id - ); - ensure_no_other_family_contract(policy, "cluster barrier")?; - Ok(()) -} - -#[derive(Clone, Copy)] -struct DebugControlRecipe { - id: &'static str, - operation_key: &'static str, - rust_name: &'static str, - rust_arguments: &'static [&'static str], - rust_result: &'static str, - compatibility_path: &'static str, - op_type: &'static str, - op_name: &'static str, - instruction: &'static str, - minimum_ptx: &'static str, - minimum_sm: Option<&'static str>, - section: &'static str, - anchor: &'static str, - adapter: DebugControlAdapter, - summary: &'static str, -} - -fn debug_control_recipe(operation: DebugControlOperation) -> DebugControlRecipe { - match operation { - DebugControlOperation::Trap => DebugControlRecipe { - id: "trap", - operation_key: "debug.control.trap", - rust_name: "trap", - rust_arguments: &[], - rust_result: "!", - compatibility_path: "cuda_device::debug::trap", - op_type: "TrapOp", - op_name: "nvvm.trap", - instruction: "trap", - minimum_ptx: "1.0", - minimum_sm: None, - section: "9.7.20.4 Miscellaneous Instructions: trap", - anchor: "miscellaneous-instructions-trap", - adapter: DebugControlAdapter::Direct, - summary: "Aborts device execution and reports an interrupt to the host.", - }, - DebugControlOperation::Breakpoint => DebugControlRecipe { - id: "breakpoint", - operation_key: "debug.control.breakpoint", - rust_name: "breakpoint", - rust_arguments: &[], - rust_result: "()", - compatibility_path: "cuda_device::debug::breakpoint", - op_type: "BreakpointOp", - op_name: "nvvm.brkpt", - instruction: "brkpt", - minimum_ptx: "1.0", - minimum_sm: Some("sm_11"), - section: "9.7.20.1 Miscellaneous Instructions: brkpt", - anchor: "miscellaneous-instructions-brkpt", - adapter: DebugControlAdapter::Direct, - summary: "Suspends device execution for a debugger breakpoint.", - }, - DebugControlOperation::Pmevent => DebugControlRecipe { - id: "pmevent", - operation_key: "debug.profiler.event", - rust_name: "pmevent", - rust_arguments: &["u32"], - rust_result: "()", - compatibility_path: "cuda_device::debug::__prof_trigger", - op_type: "PmEventOp", - op_name: "nvvm.pmevent", - instruction: "pmevent", - minimum_ptx: "1.4", - minimum_sm: None, - section: "9.7.20.3 Miscellaneous Instructions: pmevent", - anchor: "miscellaneous-instructions-pmevent", - adapter: DebugControlAdapter::ConstGenericToImmediateU32, - summary: "Triggers one compile-time-selected performance monitor event.", - }, - } -} - -#[derive(Clone, Copy)] -struct WgmmaControlRecipe { - mode: WgmmaControlMode, - abi_id: &'static str, - id: &'static str, - operation_key: &'static str, - source_record: &'static str, - selection_record: &'static str, - llvm_symbol: &'static str, - compatibility_path: &'static str, - dialect_op_type: &'static str, - dialect_op_name: &'static str, - suffix: &'static str, - adapter: WgmmaControlAdapter, - summary: &'static str, -} - -fn wgmma_control_recipe(mode: WgmmaControlMode) -> WgmmaControlRecipe { - match mode { - WgmmaControlMode::Fence => WgmmaControlRecipe { - mode, - abi_id: "i0317", - id: "wgmma_fence", - operation_key: "wgmma.control.fence.sync.aligned", - source_record: "int_nvvm_wgmma_fence_sync_aligned", - selection_record: "WGMMA_FENCE_SYNC_ALIGNED", - llvm_symbol: "llvm.nvvm.wgmma.fence.sync.aligned", - compatibility_path: "cuda_device::wgmma::wgmma_fence", - dialect_op_type: "WgmmaFenceSyncAlignedOp", - dialect_op_name: "nvvm.wgmma_fence_sync_aligned", - suffix: "fence.sync.aligned", - adapter: WgmmaControlAdapter::NoArguments, - summary: "Orders register accesses before later WGMMA operations.", - }, - WgmmaControlMode::CommitGroup => WgmmaControlRecipe { - mode, - abi_id: "i0318", - id: "wgmma_commit_group", - operation_key: "wgmma.control.commit_group.sync.aligned", - source_record: "int_nvvm_wgmma_commit_group_sync_aligned", - selection_record: "WGMMA_COMMIT_GROUP_SYNC_ALIGNED", - llvm_symbol: "llvm.nvvm.wgmma.commit_group.sync.aligned", - compatibility_path: "cuda_device::wgmma::wgmma_commit_group", - dialect_op_type: "WgmmaCommitGroupSyncAlignedOp", - dialect_op_name: "nvvm.wgmma_commit_group_sync_aligned", - suffix: "commit_group.sync.aligned", - adapter: WgmmaControlAdapter::NoArguments, - summary: "Commits prior uncommitted WGMMA operations as one group.", - }, - WgmmaControlMode::WaitGroup => WgmmaControlRecipe { - mode, - abi_id: "i0319", - id: "wgmma_wait_group", - operation_key: "wgmma.control.wait_group.sync.aligned", - source_record: "int_nvvm_wgmma_wait_group_sync_aligned", - selection_record: "WGMMA_WAIT_GROUP_SYNC_ALIGNED", - llvm_symbol: "llvm.nvvm.wgmma.wait_group.sync.aligned", - compatibility_path: "cuda_device::wgmma::__wgmma_wait_group", - dialect_op_type: "WgmmaWaitGroupSyncAlignedOp", - dialect_op_name: "nvvm.wgmma_wait_group_sync_aligned", - suffix: "wait_group.sync.aligned", - adapter: WgmmaControlAdapter::ConstGenericU32ToI64Immediate, - summary: "Waits until at most the requested number of WGMMA groups remain pending.", - }, - } -} - -fn expand_debug_control_admission( - admission: &DebugControlAdmission, -) -> Result> { - ensure!( - admission.runtime_validation == RuntimeValidation::Unexecuted, - "debug-control runtime validation may be marked executed only with GPU evidence" - ); - ensure!( - !admission.llvm_evidence_profile.trim().is_empty() - && !admission.libnvvm_evidence_profile.trim().is_empty(), - "compact debug-control admission requires both backend evidence profiles" - ); - ensure!( - admission.operations - == [ - DebugControlOperation::Trap, - DebugControlOperation::Breakpoint, - DebugControlOperation::Pmevent, - ], - "compact debug-control admission must list trap, breakpoint, and pmevent exactly once in canonical order" - ); - ensure!( - admission.abi_ids.len() == admission.operations.len(), - "pending debug-control admission needs exactly three ABI IDs before aggregation" - ); - let unique_abi_ids = admission.abi_ids.iter().collect::>(); - ensure!( - unique_abi_ids.len() == admission.abi_ids.len(), - "debug-control ABI IDs must be unique" - ); - - admission - .operations - .iter() - .zip(&admission.abi_ids) - .map(|(&operation, abi_id)| { - validate_abi_id(abi_id)?; - let recipe = debug_control_recipe(operation); - let immediate = operation == DebugControlOperation::Pmevent; - Ok(OverlayIntrinsic { - id: recipe.id.into(), - abi_id: abi_id.clone(), - operation_key: recipe.operation_key.into(), - family: "debug_control".into(), - source: Some(IntrinsicSource::PtxNative { - instruction: recipe.instruction.into(), - }), - source_record: None, - rust_module: "debug".into(), - rust_name: recipe.rust_name.into(), - rust_arguments: recipe - .rust_arguments - .iter() - .map(|value| (*value).into()) - .collect(), - rust_result: recipe.rust_result.into(), - safe: true, - must_use: false, - safe_allowlist_reason: Some( - match operation { - DebugControlOperation::Trap => { - "aborting the kernel has no memory-safety preconditions" - } - DebugControlOperation::Breakpoint => { - "requesting a debugger breakpoint has no memory-safety preconditions" - } - DebugControlOperation::Pmevent => { - "the importer accepts only the documented immediate event range" - } - } - .into(), - ), - public_rust_path: format!("cuda_intrinsics::debug::{}", recipe.rust_name), - compatibility_rust_paths: vec![recipe.compatibility_path.into()], - dialect_op_type: recipe.op_type.into(), - dialect_op_name: recipe.op_name.into(), - dialect_operands: vec![], - dialect_results: vec![], - llvm_symbol: None, - resolved_llvm_symbol: None, - llvm_arguments: vec![], - llvm_results: vec![], - pure: false, - memory: "none".into(), - convergent: false, - execution_scope: "thread".into(), - minimum_ptx: recipe.minimum_ptx.into(), - minimum_sm: recipe.minimum_sm.map(Into::into), - ptx_result: "()".into(), - targets: "all".into(), - ptx_isa_version: "9.3".into(), - ptx_isa_section: recipe.section.into(), - ptx_isa_url: format!( - "https://docs.nvidia.com/cuda/parallel-thread-execution/#{}", - recipe.anchor - ), - lowering: "generated_debug_control".into(), - backend_lowerings: vec![ - OverlayBackendLowering { - backend: IntrinsicBackend::LlvmNvptx, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: admission.llvm_evidence_profile.clone(), - targets: None, - minimum_ptx: Some("3.2".into()), - minimum_sm: Some("sm_20".into()), - }, - OverlayBackendLowering { - backend: IntrinsicBackend::LibNvvm, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: admission.libnvvm_evidence_profile.clone(), - targets: None, - // PTX floor is the inline-PTX mechanism floor, same as - // the LlvmNvptx entry above; the hardware floor stays at - // the probed sm_75 (backend-codegen evidence must sit - // exactly at the hardware floor, and CUDA 13 cannot - // probe older targets). Writing the probe PTX version - // (9.3) here made every panic path unbuildable on the - // NVVM-IR route: the floor exceeded the newest PTX - // version cuda-oxide can request. - minimum_ptx: Some("3.2".into()), - minimum_sm: Some("sm_75".into()), - }, - ], - packed_atomic: None, - redux: None, - vote: None, - active_mask: None, - warp_match: None, - warp_barrier: None, - warp_shuffle: None, - dot_product: None, - packed_alu: None, - integer_minmax: None, - packed_conversion: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cp_async_copy: None, - cp_async_control: None, - cp_async_mbarrier: None, - mbarrier_basic: None, - movmatrix: None, - mbarrier_extended: None, - register_mma: None, - sparse_mma: None, - prmt: None, - cluster_barrier: None, - wgmma_control: None, - special_register: None, - debug_control: Some(DebugControl { - operation, - adapter: recipe.adapter, - runtime_validation: admission.runtime_validation, - }), - cluster_memory: None, - clc: None, - tma: None, - tcgen05: None, - ldmatrix_variant: None, - ldmatrix_safety: None, - ldmatrix_adapter: None, - selected_address_space: None, - expected_ptx: InstructionPattern { - mnemonic: recipe.instruction.into(), - modifiers: vec![], - operands: if immediate { - vec![OperandPattern::Immediate] - } else { - vec![] - }, - }, - summary: recipe.summary.into(), - }) - }) - .collect() -} - -fn validate_debug_control_policy( - policy: &OverlayIntrinsic, - source: &IntrinsicSource, -) -> Result<()> { - let debug = policy - .debug_control - .as_ref() - .with_context(|| format!("{} has no closed debug-control contract", policy.id))?; - let recipe = debug_control_recipe(debug.operation); - let immediate = debug.operation == DebugControlOperation::Pmevent; - ensure!( - debug.adapter == recipe.adapter - && debug.runtime_validation == RuntimeValidation::Unexecuted - && policy.id == recipe.id - && policy.operation_key == recipe.operation_key - && source - == &IntrinsicSource::PtxNative { - instruction: recipe.instruction.into(), - } - && policy.source_record.is_none() - && policy.llvm_symbol.is_none() - && policy.resolved_llvm_symbol.is_none() - && policy.llvm_arguments.is_empty() - && policy.llvm_results.is_empty(), - "{} identity must remain PTX-native and match its closed debug-control recipe", - policy.id - ); - ensure!( - policy.rust_module == "debug" - && policy.rust_name == recipe.rust_name - && policy.rust_arguments == recipe.rust_arguments - && policy.rust_result == recipe.rust_result - && policy.safe - && !policy.must_use - && policy.public_rust_path == format!("cuda_intrinsics::debug::{}", recipe.rust_name) - && policy.compatibility_rust_paths == [recipe.compatibility_path], - "{} Rust API or compatibility adapter changed", - policy.id - ); - ensure!( - policy.dialect_op_type == recipe.op_type - && policy.dialect_op_name == recipe.op_name - && policy.dialect_operands.is_empty() - && policy.dialect_results.is_empty() - && policy.lowering == "generated_debug_control", - "{} dialect carrier or lowering changed", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == "none" - && !policy.convergent - && policy.execution_scope == "thread", - "{} debug-control effects changed", - policy.id - ); - ensure!( - policy.minimum_ptx == recipe.minimum_ptx - && policy.minimum_sm.as_deref() == recipe.minimum_sm - && policy.targets == "all" - && policy.ptx_result == "()" - && policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section == recipe.section - && policy.ptx_isa_url - == format!( - "https://docs.nvidia.com/cuda/parallel-thread-execution/#{}", - recipe.anchor - ), - "{} native target floor or PTX provenance changed", - policy.id - ); - ensure!( - policy.expected_ptx.mnemonic == recipe.instruction - && policy.expected_ptx.modifiers.is_empty() - && policy.expected_ptx.operands - == if immediate { - vec![OperandPattern::Immediate] - } else { - vec![] - }, - "{} expected PTX changed", - policy.id - ); - ensure_exact_inline_ptx_backends( - policy, - [ - (IntrinsicBackend::LlvmNvptx, "3.2", Some("sm_20")), - (IntrinsicBackend::LibNvvm, "3.2", Some("sm_75")), - ], - "debug-control", - )?; - ensure_no_other_family_contract(policy, "debug-control")?; - Ok(()) -} - -#[derive(Clone)] -struct ClusterMemoryRecipe { - operation: ClusterMemoryOperation, - abi_id: &'static str, - id: &'static str, - operation_key: &'static str, - source_record: Option<&'static str>, - llvm_symbol: Option<&'static str>, - ptx_native_instruction: Option<&'static str>, - rust_arguments: &'static [&'static str], - rust_result: &'static str, - compatibility_paths: &'static [&'static str], - dialect_op_type: &'static str, - dialect_op_name: &'static str, - dialect_operands: &'static [&'static str], - dialect_results: &'static [&'static str], - llvm_arguments: &'static [&'static str], - llvm_results: &'static [&'static str], - adapter: ClusterMemoryAdapter, - source_contract: ClusterMemorySourceContract, - memory: &'static str, - expected_ptx: InstructionPattern, - inline_ptx: &'static str, - inline_constraints: &'static str, - ptx_isa_section: &'static str, - ptx_isa_anchor: &'static str, - summary: &'static str, -} - -fn cluster_memory_recipe(operation: ClusterMemoryOperation) -> ClusterMemoryRecipe { - match operation { - ClusterMemoryOperation::MapSharedRank => ClusterMemoryRecipe { - operation, - abi_id: "i0320", - id: "map_shared_rank", - operation_key: "cluster.shared_address.map_rank", - source_record: Some("int_nvvm_mapa_shared_cluster"), - llvm_symbol: Some("llvm.nvvm.mapa.shared.cluster"), - ptx_native_instruction: None, - rust_arguments: &["*const u8", "u32"], - rust_result: "*const u8", - compatibility_paths: &[ - "cuda_device::cluster::map_shared_rank", - "cuda_device::cluster::map_shared_rank_mut", - ], - dialect_op_type: "MapaSharedClusterOp", - dialect_op_name: "nvvm.mapa_shared_cluster", - dialect_operands: &["ptr", "i32"], - dialect_results: &["ptr"], - llvm_arguments: &["shared_ptr", "i32"], - llvm_results: &["shared_cluster_ptr"], - adapter: ClusterMemoryAdapter::GenericConstAndMutPointerRankToSamePointer, - source_contract: ClusterMemorySourceContract::LlvmMapaSharedClusterAs7IdentityInlinePtx, - memory: "none", - expected_ptx: InstructionPattern { - mnemonic: "mapa".into(), - modifiers: vec!["shared::cluster".into(), "u64".into()], - operands: vec![ - OperandPattern::Register, - OperandPattern::Register, - OperandPattern::Register, - ], - }, - inline_ptx: "mapa.shared::cluster.u64 $0, $1, $2;", - inline_constraints: "=l,l,r", - ptx_isa_section: "9.7.9.24 Data Movement and Conversion Instructions: mapa", - ptx_isa_anchor: "data-movement-and-conversion-instructions-mapa", - summary: "Maps a CTA-shared address to the same offset in another cluster rank.", - }, - ClusterMemoryOperation::ReadU32 => ClusterMemoryRecipe { - operation, - abi_id: "i0321", - id: "dsmem_read_u32", - operation_key: "cluster.shared_memory.map_rank_then_read_u32", - source_record: None, - llvm_symbol: None, - ptx_native_instruction: Some("mapa.shared::cluster.u64 + ld.shared::cluster.u32"), - rust_arguments: &["*const u32", "u32"], - rust_result: "u32", - compatibility_paths: &["cuda_device::cluster::dsmem_read_u32"], - dialect_op_type: "DsmemReadU32Op", - dialect_op_name: "nvvm.dsmem_read_u32", - dialect_operands: &["ptr", "i32"], - dialect_results: &["i32"], - llvm_arguments: &[], - llvm_results: &[], - adapter: ClusterMemoryAdapter::ConstU32PointerRankToU32, - source_contract: ClusterMemorySourceContract::PtxNativeMapaThenWeakClusterLoad, - memory: "read", - expected_ptx: InstructionPattern { - mnemonic: "ld".into(), - modifiers: vec!["shared::cluster".into(), "u32".into()], - operands: vec![OperandPattern::Register, OperandPattern::Address], - }, - inline_ptx: "{ .reg .u64 %mapped; mapa.shared::cluster.u64 %mapped, $1, $2; ld.shared::cluster.u32 $0, [%mapped]; }", - inline_constraints: "=r,l,r,~{memory}", - ptx_isa_section: "9.7.9.8 Data Movement and Conversion Instructions: ld", - ptx_isa_anchor: "data-movement-and-conversion-instructions-ld", - summary: "Maps a CTA-shared address to another cluster rank and reads one weak u32 value.", - }, - } -} - -pub(crate) fn cluster_memory_inline_recipe( - operation: ClusterMemoryOperation, -) -> (&'static str, &'static str) { - let recipe = cluster_memory_recipe(operation); - (recipe.inline_ptx, recipe.inline_constraints) -} - -fn expand_cluster_memory_admission( - admission: &ClusterMemoryAdmission, -) -> Result> { - ensure!( - admission.runtime_validation == RuntimeValidation::Unexecuted, - "cluster-memory runtime validation may be marked executed only with GPU evidence" - ); - ensure!( - !admission.llvm_evidence_profile.trim().is_empty() - && !admission.libnvvm_evidence_profile.trim().is_empty(), - "compact cluster-memory admission requires both backend evidence profiles" - ); - let expected = [ - ClusterMemoryOperation::MapSharedRank, - ClusterMemoryOperation::ReadU32, - ]; - ensure!( - admission.variants.len() == expected.len() - && admission - .variants - .iter() - .map(|variant| variant.operation) - .eq(expected), - "compact cluster-memory admission must list map_shared_rank and read_u32 once in canonical order" - ); - - admission - .variants - .iter() - .map(|variant| { - let recipe = cluster_memory_recipe(variant.operation); - ensure!( - variant.abi_id == recipe.abi_id, - "{} must keep reserved ABI ID {}", - recipe.id, - recipe.abi_id - ); - let source = - recipe - .ptx_native_instruction - .map(|instruction| IntrinsicSource::PtxNative { - instruction: instruction.into(), - }); - Ok(OverlayIntrinsic { - id: recipe.id.into(), - abi_id: variant.abi_id.clone(), - operation_key: recipe.operation_key.into(), - family: "cluster_memory".into(), - source, - source_record: recipe.source_record.map(Into::into), - rust_module: "cluster".into(), - rust_name: recipe.id.into(), - rust_arguments: recipe - .rust_arguments - .iter() - .map(|value| (*value).into()) - .collect(), - rust_result: recipe.rust_result.into(), - safe: false, - must_use: true, - safe_allowlist_reason: None, - public_rust_path: format!("cuda_intrinsics::cluster::{}", recipe.id), - compatibility_rust_paths: recipe - .compatibility_paths - .iter() - .map(|path| (*path).into()) - .collect(), - dialect_op_type: recipe.dialect_op_type.into(), - dialect_op_name: recipe.dialect_op_name.into(), - dialect_operands: recipe - .dialect_operands - .iter() - .map(|value| (*value).into()) - .collect(), - dialect_results: recipe - .dialect_results - .iter() - .map(|value| (*value).into()) - .collect(), - llvm_symbol: recipe.llvm_symbol.map(Into::into), - resolved_llvm_symbol: None, - llvm_arguments: recipe - .llvm_arguments - .iter() - .map(|value| (*value).into()) - .collect(), - llvm_results: recipe - .llvm_results - .iter() - .map(|value| (*value).into()) - .collect(), - pure: false, - memory: recipe.memory.into(), - convergent: true, - execution_scope: "cluster".into(), - minimum_ptx: "7.8".into(), - minimum_sm: Some("sm_90".into()), - ptx_result: recipe.rust_result.into(), - targets: "all".into(), - ptx_isa_version: "9.3".into(), - ptx_isa_section: recipe.ptx_isa_section.into(), - ptx_isa_url: format!( - "https://docs.nvidia.com/cuda/parallel-thread-execution/#{}", - recipe.ptx_isa_anchor - ), - lowering: "generated_cluster_memory_inline_ptx".into(), - backend_lowerings: [IntrinsicBackend::LlvmNvptx, IntrinsicBackend::LibNvvm] - .into_iter() - .map(|backend| OverlayBackendLowering { - backend, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: match backend { - IntrinsicBackend::LlvmNvptx => admission.llvm_evidence_profile.clone(), - IntrinsicBackend::LibNvvm => admission.libnvvm_evidence_profile.clone(), - }, - targets: None, - minimum_ptx: Some("7.8".into()), - minimum_sm: Some("sm_90".into()), - }) - .collect(), - packed_atomic: None, - redux: None, - vote: None, - active_mask: None, - warp_match: None, - warp_barrier: None, - warp_shuffle: None, - dot_product: None, - packed_alu: None, - integer_minmax: None, - packed_conversion: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cp_async_copy: None, - cp_async_control: None, - cp_async_mbarrier: None, - mbarrier_basic: None, - movmatrix: None, - mbarrier_extended: None, - register_mma: None, - sparse_mma: None, - prmt: None, - cluster_barrier: None, - wgmma_control: None, - special_register: None, - debug_control: None, - cluster_memory: Some(ClusterMemory { - operation: recipe.operation, - adapter: recipe.adapter, - source_contract: recipe.source_contract, - runtime_validation: admission.runtime_validation, - }), - clc: None, - tma: None, - tcgen05: None, - ldmatrix_variant: None, - ldmatrix_safety: None, - ldmatrix_adapter: None, - selected_address_space: None, - expected_ptx: recipe.expected_ptx, - summary: recipe.summary.into(), - }) - }) - .collect() -} - -fn validate_cluster_memory_policy( - policy: &OverlayIntrinsic, - source: &IntrinsicSource, - declaration: Option<&ImportedIntrinsic>, -) -> Result<()> { - let cluster = policy - .cluster_memory - .as_ref() - .with_context(|| format!("{} has no closed cluster-memory contract", policy.id))?; - let recipe = cluster_memory_recipe(cluster.operation); - ensure!( - cluster.adapter == recipe.adapter - && cluster.source_contract == recipe.source_contract - && cluster.runtime_validation == RuntimeValidation::Unexecuted - && policy.id == recipe.id - && policy.operation_key == recipe.operation_key, - "{} does not match its closed cluster-memory identity", - policy.id - ); - match cluster.source_contract { - ClusterMemorySourceContract::LlvmMapaSharedClusterAs7IdentityInlinePtx => { - let declaration = declaration.context("map_shared_rank requires its LLVM identity")?; - ensure!( - matches!(source, IntrinsicSource::LlvmImported { source_record } - if source_record == "int_nvvm_mapa_shared_cluster") - && policy.source.is_none() - && policy.source_record.as_deref() == recipe.source_record - && policy.llvm_symbol.as_deref() == recipe.llvm_symbol - && declaration.arguments == ["shared_ptr", "i32"] - && declaration.results == ["shared_cluster_ptr"] - && declaration.properties - == ["IntrNoMem", "IntrSpeculatable", "NoCapture"], - "{} must retain the AS7-returning LLVM mapa record as identity only", - policy.id - ); - let mut selections = Vec::new(); - for selection in &declaration.selections { - if selection_matches_policy(policy, selection)? { - selections.push(selection); - } - } - ensure!( - selections.len() == 2 - && selections - .iter() - .map(|selection| selection.source_record.as_str()) - .collect::>() - == BTreeSet::from(["mapa_shared_cluster_64", "mapa_shared_cluster_64i",]) - && selections.iter().all(|selection| { - selection.asm == "mapa.shared::cluster.u64 \t$d, $a, $b;" - && selection.predicates - == [ - "Subtarget->getSmVersion() >= 90", - "Subtarget->getPTXVersion() >= 78", - ] - }), - "{} must retain both exact 64-bit mapa selections", - policy.id - ); - } - ClusterMemorySourceContract::PtxNativeMapaThenWeakClusterLoad => ensure!( - matches!(source, IntrinsicSource::PtxNative { instruction } - if Some(instruction.as_str()) == recipe.ptx_native_instruction) - && policy.source_record.is_none() - && policy.llvm_symbol.is_none() - && declaration.is_none() - && policy.llvm_arguments.is_empty() - && policy.llvm_results.is_empty(), - "{} must remain a PTX-native mapa plus weak cluster-load composite", - policy.id - ), - } - ensure!( - policy.rust_module == "cluster" - && policy.rust_name == recipe.id - && policy.rust_arguments == recipe.rust_arguments - && policy.rust_result == recipe.rust_result - && !policy.safe - && policy.must_use - && policy.public_rust_path == format!("cuda_intrinsics::cluster::{}", recipe.id) - && policy.compatibility_rust_paths == recipe.compatibility_paths, - "{} Rust API or compatibility paths changed", - policy.id - ); - ensure!( - policy.dialect_op_type == recipe.dialect_op_type - && policy.dialect_op_name == recipe.dialect_op_name - && policy.dialect_operands == recipe.dialect_operands - && policy.dialect_results == recipe.dialect_results - && policy.llvm_arguments == recipe.llvm_arguments - && policy.llvm_results == recipe.llvm_results - && policy.lowering == "generated_cluster_memory_inline_ptx", - "{} dialect carrier or AS7 source boundary changed", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == recipe.memory - && policy.convergent - && policy.execution_scope == "cluster", - "{} effects changed", - policy.id - ); - ensure!( - policy.minimum_ptx == "7.8" - && policy.minimum_sm.as_deref() == Some("sm_90") - && policy.targets == "all" - && policy.ptx_result == recipe.rust_result - && policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section == recipe.ptx_isa_section - && policy.ptx_isa_url - == format!( - "https://docs.nvidia.com/cuda/parallel-thread-execution/#{}", - recipe.ptx_isa_anchor - ), - "{} PTX provenance or target floor changed", - policy.id - ); - ensure!( - policy.expected_ptx == recipe.expected_ptx, - "{} expected PTX changed", - policy.id - ); - ensure_exact_inline_ptx_backends( - policy, - [ - (IntrinsicBackend::LlvmNvptx, "7.8", Some("sm_90")), - (IntrinsicBackend::LibNvvm, "7.8", Some("sm_90")), - ], - "cluster-memory", - )?; - ensure_no_other_family_contract(policy, "cluster-memory")?; - Ok(()) -} - -#[derive(Clone)] -struct MbarrierExtendedRecipe { - operation: MbarrierExtendedOperation, - abi_id: &'static str, - id: &'static str, - operation_key: &'static str, - source_record: Option<&'static str>, - llvm_symbol: Option<&'static str>, - ptx_native_instruction: Option<&'static str>, - rust_arguments: &'static [&'static str], - rust_result: &'static str, - must_use: bool, - dialect_op_type: &'static str, - dialect_op_name: &'static str, - dialect_operands: &'static [&'static str], - dialect_results: &'static [&'static str], - llvm_arguments: &'static [&'static str], - llvm_results: &'static [&'static str], - llvm_properties: &'static [&'static str], - adapter: MbarrierExtendedAdapter, - source_contract: MbarrierExtendedSourceContract, - execution_scope: &'static str, - minimum_ptx: &'static str, - minimum_sm: &'static str, - ptx_result: &'static str, - expected_ptx: InstructionPattern, - inline_ptx: &'static str, - inline_constraints: &'static str, - ptx_isa_section: &'static str, - ptx_isa_url: &'static str, - summary: &'static str, -} - -fn mbarrier_extended_recipe(operation: MbarrierExtendedOperation) -> MbarrierExtendedRecipe { - let instruction = |modifiers: &[&str], operands| InstructionPattern { - mnemonic: if modifiers.first() == Some(&"nanosleep") { - "nanosleep".into() - } else if modifiers.first() == Some(&"fence") { - "fence".into() - } else { - "mbarrier".into() - }, - modifiers: modifiers[1..].iter().map(|value| (*value).into()).collect(), - operands, - }; - match operation { - MbarrierExtendedOperation::ArriveExpectTxCta => MbarrierExtendedRecipe { - operation, - abi_id: "i0306", - id: "mbarrier_arrive_expect_tx", - operation_key: "barrier.mbarrier.arrive.expect_tx.shared.cta.release.cta", - source_record: Some("int_nvvm_mbarrier_arrive_expect_tx_scope_cta_space_cta"), - llvm_symbol: Some("llvm.nvvm.mbarrier.arrive.expect.tx.scope.cta.space.cta"), - ptx_native_instruction: None, - rust_arguments: &["*const u64", "u32", "u32"], - rust_result: "u64", - must_use: true, - dialect_op_type: "MbarrierArriveExpectTxSharedOp", - dialect_op_name: "nvvm.mbarrier_arrive_expect_tx_shared", - dialect_operands: &["ptr", "i32"], - dialect_results: &["i64"], - llvm_arguments: &["shared_ptr", "i32"], - llvm_results: &["i64"], - llvm_properties: &["IntrConvergent", "IntrNoCallback"], - adapter: MbarrierExtendedAdapter::PointerTxCountBytesToTokenDroppingTxCount, - source_contract: MbarrierExtendedSourceContract::LlvmImported, - execution_scope: "cta", - minimum_ptx: "8.0", - minimum_sm: "sm_90", - ptx_result: "u64", - expected_ptx: instruction( - &[ - "mbarrier", - "arrive", - "expect_tx", - "release", - "cta", - "shared::cta", - "b64", - ], - vec![ - OperandPattern::Register, - OperandPattern::Address, - OperandPattern::Register, - ], - ), - inline_ptx: "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 $0, [$1], $2;", - inline_constraints: "=l,l,r,~{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 CTA-shared barrier and adds expected transaction bytes.", - }, - MbarrierExtendedOperation::ArriveExpectTxCluster => MbarrierExtendedRecipe { - operation, - abi_id: "i0307", - id: "mbarrier_arrive_expect_tx_cluster", - operation_key: "barrier.mbarrier.arrive.expect_tx.shared.cta.relaxed.cluster", - source_record: Some( - "int_nvvm_mbarrier_arrive_expect_tx_relaxed_scope_cluster_space_cta", - ), - llvm_symbol: Some( - "llvm.nvvm.mbarrier.arrive.expect.tx.relaxed.scope.cluster.space.cta", - ), - ptx_native_instruction: None, - rust_arguments: &["*const u64", "u32", "u32"], - rust_result: "u64", - must_use: true, - dialect_op_type: "MbarrierArriveExpectTxClusterOp", - dialect_op_name: "nvvm.mbarrier_arrive_expect_tx_cluster", - dialect_operands: &["ptr", "i32"], - dialect_results: &["i64"], - llvm_arguments: &["shared_ptr", "i32"], - llvm_results: &["i64"], - llvm_properties: &["IntrArgMemOnly", "IntrConvergent", "IntrNoCallback"], - adapter: MbarrierExtendedAdapter::PointerTxCountBytesToTokenDroppingTxCount, - source_contract: MbarrierExtendedSourceContract::LlvmImported, - execution_scope: "cluster", - minimum_ptx: "8.6", - minimum_sm: "sm_90", - ptx_result: "u64", - expected_ptx: instruction( - &[ - "mbarrier", - "arrive", - "expect_tx", - "relaxed", - "cluster", - "shared::cta", - "b64", - ], - vec![ - OperandPattern::Register, - OperandPattern::Address, - OperandPattern::Register, - ], - ), - inline_ptx: "mbarrier.arrive.expect_tx.relaxed.cluster.shared::cta.b64 $0, [$1], $2;", - inline_constraints: "=l,l,r,~{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 CTA-shared barrier with cluster-scope transaction tracking.", - }, - MbarrierExtendedOperation::ArriveRemoteCluster => MbarrierExtendedRecipe { - operation, - abi_id: "i0308", - id: "mbarrier_arrive_cluster", - operation_key: "barrier.mbarrier.arrive.shared.cluster.release.cluster.raw_address", - source_record: None, - llvm_symbol: None, - ptx_native_instruction: Some("mbarrier.arrive.release.cluster.shared::cluster.b64"), - rust_arguments: &["u64"], - rust_result: "()", - must_use: false, - dialect_op_type: "MbarrierArriveClusterOp", - dialect_op_name: "nvvm.mbarrier_arrive_cluster", - dialect_operands: &["i64"], - dialect_results: &[], - llvm_arguments: &[], - llvm_results: &[], - llvm_properties: &[], - adapter: MbarrierExtendedAdapter::RawClusterAddressToVoid, - source_contract: MbarrierExtendedSourceContract::PtxNativeRawClusterAddress, - execution_scope: "cluster", - minimum_ptx: "8.0", - minimum_sm: "sm_90", - ptx_result: "()", - expected_ptx: instruction( - &[ - "mbarrier", - "arrive", - "release", - "cluster", - "shared::cluster", - "b64", - ], - vec![ - OperandPattern::Exact { value: "_".into() }, - OperandPattern::Address, - ], - ), - inline_ptx: "mbarrier.arrive.release.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 through its raw address.", - }, - MbarrierExtendedOperation::TryWaitTokenCta => MbarrierExtendedRecipe { - operation, - abi_id: "i0309", - id: "mbarrier_try_wait", - operation_key: "barrier.mbarrier.try_wait.shared.cta.token", - source_record: Some("int_nvvm_mbarrier_try_wait_scope_cta_space_cta"), - llvm_symbol: Some("llvm.nvvm.mbarrier.try.wait.scope.cta.space.cta"), - ptx_native_instruction: None, - rust_arguments: &["*const u64", "u64"], - rust_result: "bool", - must_use: true, - dialect_op_type: "MbarrierTryWaitSharedOp", - dialect_op_name: "nvvm.mbarrier_try_wait_shared", - dialect_operands: &["ptr", "i64"], - dialect_results: &["i1"], - llvm_arguments: &["shared_ptr", "i64"], - llvm_results: &["i1"], - llvm_properties: &["IntrConvergent", "IntrNoCallback", "NoCapture"], - adapter: MbarrierExtendedAdapter::PointerTokenToPredicate, - source_contract: MbarrierExtendedSourceContract::LlvmImported, - execution_scope: "cta", - minimum_ptx: "7.8", - minimum_sm: "sm_90", - ptx_result: "bool", - expected_ptx: instruction( - &["mbarrier", "try_wait", "shared", "b64"], - vec![ - OperandPattern::Register, - OperandPattern::Address, - OperandPattern::Register, - ], - ), - inline_ptx: "{ .reg .pred %p0; mbarrier.try_wait.shared.b64 %p0, [$1], $2; selp.b32 $0, 1, 0, %p0; }", - inline_constraints: "=r,l,l,~{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 a CTA-shared barrier token with a scheduling hint.", - }, - MbarrierExtendedOperation::TryWaitParityCta => MbarrierExtendedRecipe { - operation, - abi_id: "i0310", - id: "mbarrier_try_wait_parity", - operation_key: "barrier.mbarrier.try_wait.parity.shared.cta", - source_record: Some("int_nvvm_mbarrier_try_wait_parity_scope_cta_space_cta"), - llvm_symbol: Some("llvm.nvvm.mbarrier.try.wait.parity.scope.cta.space.cta"), - ptx_native_instruction: None, - rust_arguments: &["*const u64", "u32"], - rust_result: "bool", - must_use: true, - dialect_op_type: "MbarrierTryWaitParitySharedOp", - dialect_op_name: "nvvm.mbarrier_try_wait_parity_shared", - dialect_operands: &["ptr", "i32"], - dialect_results: &["i1"], - llvm_arguments: &["shared_ptr", "i32"], - llvm_results: &["i1"], - llvm_properties: &["IntrConvergent", "IntrNoCallback", "NoCapture"], - adapter: MbarrierExtendedAdapter::PointerParityToPredicate, - source_contract: MbarrierExtendedSourceContract::LlvmImported, - execution_scope: "cta", - minimum_ptx: "7.8", - minimum_sm: "sm_90", - ptx_result: "bool", - expected_ptx: instruction( - &["mbarrier", "try_wait", "parity", "shared::cta", "b64"], - vec![ - OperandPattern::Register, - OperandPattern::Address, - OperandPattern::Register, - ], - ), - inline_ptx: "{ .reg .pred %p0; mbarrier.try_wait.parity.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 a CTA-shared barrier phase by parity.", - }, - MbarrierExtendedOperation::TryWaitParityCluster => MbarrierExtendedRecipe { - operation, - abi_id: "i0311", - id: "mbarrier_try_wait_parity_cluster", - operation_key: "barrier.mbarrier.try_wait.parity.shared.cta.acquire.cluster", - source_record: Some("int_nvvm_mbarrier_try_wait_parity_scope_cluster_space_cta"), - llvm_symbol: Some("llvm.nvvm.mbarrier.try.wait.parity.scope.cluster.space.cta"), - ptx_native_instruction: None, - rust_arguments: &["*const u64", "u32"], - rust_result: "bool", - must_use: true, - dialect_op_type: "MbarrierTryWaitParityClusterOp", - dialect_op_name: "nvvm.mbarrier_try_wait_parity_cluster", - dialect_operands: &["ptr", "i32"], - dialect_results: &["i1"], - llvm_arguments: &["shared_ptr", "i32"], - llvm_results: &["i1"], - llvm_properties: &["IntrConvergent", "IntrNoCallback", "NoCapture"], - adapter: MbarrierExtendedAdapter::PointerParityToPredicate, - source_contract: MbarrierExtendedSourceContract::LlvmImported, - execution_scope: "cluster", - minimum_ptx: "8.0", - minimum_sm: "sm_90", - ptx_result: "bool", - expected_ptx: instruction( - &[ - "mbarrier", - "try_wait", - "parity", - "acquire", - "cluster", - "shared::cta", - "b64", - ], - vec![ - OperandPattern::Register, - OperandPattern::Address, - OperandPattern::Register, - ], - ), - inline_ptx: "{ .reg .pred %p0; mbarrier.try_wait.parity.acquire.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 with cluster-scope acquire ordering.", - }, - MbarrierExtendedOperation::FenceProxyAsyncSharedCta => MbarrierExtendedRecipe { - operation, - abi_id: "i0312", - id: "fence_proxy_async_shared_cta", - operation_key: "fence.proxy.async.shared.cta", - source_record: Some("int_nvvm_fence_proxy_async_shared_cta"), - llvm_symbol: Some("llvm.nvvm.fence.proxy.async.shared_cta"), - ptx_native_instruction: None, - rust_arguments: &[], - rust_result: "()", - must_use: false, - dialect_op_type: "FenceProxyAsyncSharedCtaOp", - dialect_op_name: "nvvm.fence_proxy_async_shared_cta", - dialect_operands: &[], - dialect_results: &[], - llvm_arguments: &[], - llvm_results: &[], - llvm_properties: &["IntrNoCallback"], - adapter: MbarrierExtendedAdapter::ZeroOperandsToVoid, - source_contract: MbarrierExtendedSourceContract::LlvmImported, - execution_scope: "cta", - minimum_ptx: "8.0", - minimum_sm: "sm_90", - ptx_result: "()", - expected_ptx: instruction(&["fence", "proxy", "async", "shared::cta"], vec![]), - inline_ptx: "fence.proxy.async.shared::cta;", - inline_constraints: "~{memory}", - ptx_isa_section: "Parallel Synchronization and Communication Instructions: fence.proxy", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-fence-proxy", - summary: "Makes CTA-shared generic-proxy writes visible to the async proxy.", - }, - MbarrierExtendedOperation::FenceMbarrierInitReleaseCluster => MbarrierExtendedRecipe { - operation, - abi_id: "i0313", - id: "fence_mbarrier_init_release_cluster", - operation_key: "fence.mbarrier_init.release.cluster", - source_record: Some("int_nvvm_fence_mbarrier_init_release_cluster"), - llvm_symbol: Some("llvm.nvvm.fence.mbarrier_init.release.cluster"), - ptx_native_instruction: None, - rust_arguments: &[], - rust_result: "()", - must_use: false, - dialect_op_type: "FenceMbarrierInitReleaseClusterOp", - dialect_op_name: "nvvm.fence_mbarrier_init_release_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.0", - minimum_sm: "sm_90", - ptx_result: "()", - expected_ptx: instruction(&["fence", "mbarrier_init", "release", "cluster"], vec![]), - inline_ptx: "fence.mbarrier_init.release.cluster;", - inline_constraints: "~{memory}", - ptx_isa_section: "Parallel Synchronization and Communication Instructions: fence.mbarrier_init", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-fence-mbarrier-init", - summary: "Releases mbarrier initialization at cluster scope.", - }, - MbarrierExtendedOperation::FenceProxyAsyncGenericReleaseSharedCtaCluster => { - MbarrierExtendedRecipe { - operation, - abi_id: "i0314", - id: "fence_proxy_async_generic_release_shared_cta_cluster", - operation_key: "fence.proxy.async_generic.release.sync_restrict.shared_cta.cluster", - source_record: Some( - "int_nvvm_fence_proxy_async_generic_release_sync_restrict_space_cta_scope_cluster", - ), - llvm_symbol: Some( - "llvm.nvvm.fence.proxy.async_generic.release.sync_restrict.space.cta.scope.cluster", - ), - ptx_native_instruction: None, - rust_arguments: &[], - rust_result: "()", - must_use: false, - dialect_op_type: "FenceProxyAsyncGenericReleaseSharedCtaClusterOp", - dialect_op_name: "nvvm.fence_proxy_async_generic_release_shared_cta_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", - "proxy", - "async::generic", - "release", - "sync_restrict::shared::cta", - "cluster", - ], - vec![], - ), - inline_ptx: "fence.proxy.async::generic.release.sync_restrict::shared::cta.cluster;", - inline_constraints: "~{memory}", - ptx_isa_section: "Parallel Synchronization and Communication Instructions: fence.proxy", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-fence-proxy", - summary: "Releases CTA-shared generic-proxy writes to the async proxy at cluster scope.", - } - } - MbarrierExtendedOperation::FenceProxyAsyncGenericAcquireSharedClusterCluster => { - MbarrierExtendedRecipe { - operation, - abi_id: "i0315", - id: "fence_proxy_async_generic_acquire_shared_cluster_cluster", - operation_key: "fence.proxy.async_generic.acquire.sync_restrict.shared_cluster.cluster", - source_record: Some( - "int_nvvm_fence_proxy_async_generic_acquire_sync_restrict_space_cluster_scope_cluster", - ), - llvm_symbol: Some( - "llvm.nvvm.fence.proxy.async_generic.acquire.sync_restrict.space.cluster.scope.cluster", - ), - ptx_native_instruction: None, - rust_arguments: &[], - rust_result: "()", - must_use: false, - dialect_op_type: "FenceProxyAsyncGenericAcquireSharedClusterClusterOp", - dialect_op_name: "nvvm.fence_proxy_async_generic_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", - "proxy", - "async::generic", - "acquire", - "sync_restrict::shared::cluster", - "cluster", - ], - vec![], - ), - inline_ptx: "fence.proxy.async::generic.acquire.sync_restrict::shared::cluster.cluster;", - inline_constraints: "~{memory}", - ptx_isa_section: "Parallel Synchronization and Communication Instructions: fence.proxy", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-fence-proxy", - summary: "Acquires cluster-shared async-proxy writes through the generic proxy.", - } - } - MbarrierExtendedOperation::Nanosleep => MbarrierExtendedRecipe { - operation, - abi_id: "i0316", - id: "nanosleep", - operation_key: "thread.nanosleep.u32", - source_record: Some("int_nvvm_nanosleep"), - llvm_symbol: Some("llvm.nvvm.nanosleep"), - ptx_native_instruction: None, - rust_arguments: &["u32"], - rust_result: "()", - must_use: false, - dialect_op_type: "NanosleepOp", - dialect_op_name: "nvvm.nanosleep", - dialect_operands: &["i32"], - dialect_results: &[], - llvm_arguments: &["i32"], - llvm_results: &[], - llvm_properties: &["IntrConvergent", "IntrHasSideEffects", "IntrNoMem"], - adapter: MbarrierExtendedAdapter::NanosecondsToVoid, - source_contract: MbarrierExtendedSourceContract::LlvmImported, - execution_scope: "thread", - minimum_ptx: "6.3", - minimum_sm: "sm_70", - ptx_result: "()", - expected_ptx: instruction( - &["nanosleep", "u32"], - vec![OperandPattern::RegisterOrImmediate], - ), - inline_ptx: "nanosleep.u32 $0;", - inline_constraints: "r,~{memory}", - ptx_isa_section: "Parallel Synchronization and Communication Instructions: nanosleep", - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-nanosleep", - summary: "Suspends the executing thread for approximately the requested nanoseconds.", - }, - } -} - -pub(crate) fn mbarrier_extended_inline_recipe( - operation: MbarrierExtendedOperation, -) -> (&'static str, &'static str) { - let recipe = mbarrier_extended_recipe(operation); - (recipe.inline_ptx, recipe.inline_constraints) -} - -fn mbarrier_extended_backend_floor( - operation: MbarrierExtendedOperation, - backend: IntrinsicBackend, -) -> (&'static str, &'static str) { - let recipe = mbarrier_extended_recipe(operation); - match (operation, backend) { - (MbarrierExtendedOperation::Nanosleep, IntrinsicBackend::LibNvvm) => ("6.3", "sm_75"), - _ => (recipe.minimum_ptx, recipe.minimum_sm), - } -} - -fn expand_mbarrier_extended_admission( - admission: &MbarrierExtendedAdmission, -) -> Result> { - ensure!( - admission.runtime_validation == RuntimeValidation::Unexecuted, - "extended-mbarrier runtime validation may be marked executed only with GPU evidence" - ); - ensure!( - !admission.llvm_evidence_profile.trim().is_empty() - && !admission.libnvvm_evidence_profile.trim().is_empty(), - "compact extended-mbarrier admission requires both backend evidence profiles" - ); - let expected_operations = BTreeSet::from([ - MbarrierExtendedOperation::ArriveExpectTxCta, - MbarrierExtendedOperation::ArriveExpectTxCluster, - MbarrierExtendedOperation::ArriveRemoteCluster, - MbarrierExtendedOperation::TryWaitTokenCta, - MbarrierExtendedOperation::TryWaitParityCta, - MbarrierExtendedOperation::TryWaitParityCluster, - MbarrierExtendedOperation::FenceProxyAsyncSharedCta, - MbarrierExtendedOperation::FenceMbarrierInitReleaseCluster, - MbarrierExtendedOperation::FenceProxyAsyncGenericReleaseSharedCtaCluster, - MbarrierExtendedOperation::FenceProxyAsyncGenericAcquireSharedClusterCluster, - MbarrierExtendedOperation::Nanosleep, - ]); - let actual_operations = admission - .variants - .iter() - .map(|variant| variant.operation) - .collect::>(); - ensure!( - admission.variants.len() == expected_operations.len() - && actual_operations == expected_operations, - "compact extended-mbarrier admission must contain each reviewed operation exactly once" - ); - - admission - .variants - .iter() - .map(|variant| { - let recipe = mbarrier_extended_recipe(variant.operation); - ensure!( - variant.abi_id == recipe.abi_id, - "{} must keep reserved ABI ID {}", - recipe.id, - recipe.abi_id - ); - let source = - recipe - .ptx_native_instruction - .map(|instruction| IntrinsicSource::PtxNative { - instruction: instruction.into(), - }); - Ok(OverlayIntrinsic { - id: recipe.id.into(), - abi_id: variant.abi_id.clone(), - operation_key: recipe.operation_key.into(), - family: "mbarrier_extended".into(), - source, - source_record: recipe.source_record.map(str::to_owned), - rust_module: "barrier".into(), - rust_name: recipe.id.into(), - rust_arguments: recipe - .rust_arguments - .iter() - .map(|value| (*value).into()) - .collect(), - rust_result: recipe.rust_result.into(), - safe: false, - must_use: recipe.must_use, - safe_allowlist_reason: None, - public_rust_path: format!("cuda_intrinsics::barrier::{}", recipe.id), - compatibility_rust_paths: vec![format!("cuda_device::barrier::{}", recipe.id)], - dialect_op_type: recipe.dialect_op_type.into(), - dialect_op_name: recipe.dialect_op_name.into(), - dialect_operands: recipe - .dialect_operands - .iter() - .map(|value| (*value).into()) - .collect(), - dialect_results: recipe - .dialect_results - .iter() - .map(|value| (*value).into()) - .collect(), - llvm_symbol: recipe.llvm_symbol.map(str::to_owned), - resolved_llvm_symbol: None, - llvm_arguments: recipe - .llvm_arguments - .iter() - .map(|value| (*value).into()) - .collect(), - llvm_results: recipe - .llvm_results - .iter() - .map(|value| (*value).into()) - .collect(), - pure: false, - memory: "read_write".into(), - convergent: true, - execution_scope: recipe.execution_scope.into(), - minimum_ptx: recipe.minimum_ptx.into(), - minimum_sm: Some(recipe.minimum_sm.into()), - ptx_result: recipe.ptx_result.into(), - targets: "all".into(), - ptx_isa_version: "9.3".into(), - ptx_isa_section: recipe.ptx_isa_section.into(), - ptx_isa_url: recipe.ptx_isa_url.into(), - lowering: "generated_mbarrier_extended_inline_ptx".into(), - backend_lowerings: [IntrinsicBackend::LlvmNvptx, IntrinsicBackend::LibNvvm] - .into_iter() - .map(|backend| { - let (minimum_ptx, minimum_sm) = - mbarrier_extended_backend_floor(recipe.operation, backend); - OverlayBackendLowering { - backend, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: match backend { - IntrinsicBackend::LlvmNvptx => { - admission.llvm_evidence_profile.clone() - } - IntrinsicBackend::LibNvvm => { - admission.libnvvm_evidence_profile.clone() - } - }, - targets: None, - minimum_ptx: Some(minimum_ptx.into()), - minimum_sm: Some(minimum_sm.into()), - } - }) - .collect(), - packed_atomic: None, - redux: None, - vote: None, - active_mask: None, - warp_match: None, - warp_barrier: None, - warp_shuffle: None, - dot_product: None, - packed_alu: None, - integer_minmax: None, - packed_conversion: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cp_async_copy: None, - cp_async_control: None, - cp_async_mbarrier: None, - mbarrier_basic: None, - movmatrix: None, - mbarrier_extended: Some(MbarrierExtended { - operation: recipe.operation, - adapter: recipe.adapter, - source_contract: recipe.source_contract, - runtime_validation: admission.runtime_validation, - }), - register_mma: None, - sparse_mma: None, - prmt: None, - cluster_barrier: None, - wgmma_control: None, - special_register: None, - debug_control: None, - cluster_memory: None, - clc: None, - tma: None, - tcgen05: None, - ldmatrix_variant: None, - ldmatrix_safety: None, - ldmatrix_adapter: None, - selected_address_space: None, - expected_ptx: recipe.expected_ptx, - summary: recipe.summary.into(), - }) - }) - .collect() -} - -fn validate_mbarrier_extended_policy( - policy: &OverlayIntrinsic, - source: &IntrinsicSource, - declaration: Option<&ImportedIntrinsic>, -) -> Result<()> { - let contract = policy - .mbarrier_extended - .as_ref() - .with_context(|| format!("{} has no closed extended-mbarrier contract", policy.id))?; - let recipe = mbarrier_extended_recipe(contract.operation); - ensure!( - contract.adapter == recipe.adapter - && contract.source_contract == recipe.source_contract - && policy.id == recipe.id - && policy.abi_id == recipe.abi_id - && policy.operation_key == recipe.operation_key, - "{} identity or adapter does not match its closed extended-mbarrier recipe", - policy.id - ); - match recipe.source_contract { - MbarrierExtendedSourceContract::LlvmImported => ensure!( - matches!( - source, - IntrinsicSource::LlvmImported { source_record } - if Some(source_record.as_str()) == recipe.source_record - ) && policy.source.is_none() - && policy.source_record.as_deref() == recipe.source_record - && policy.llvm_symbol.as_deref() == recipe.llvm_symbol - && declaration.is_some_and(|record| { - record.source_record == recipe.source_record.unwrap() - && record.properties == recipe.llvm_properties - }), - "{} LLVM source contract changed", - policy.id - ), - MbarrierExtendedSourceContract::PtxNativeRawClusterAddress => ensure!( - matches!( - source, - IntrinsicSource::PtxNative { instruction } - if Some(instruction.as_str()) == recipe.ptx_native_instruction - ) && policy.source_record.is_none() - && policy.llvm_symbol.is_none() - && declaration.is_none(), - "{} must remain the PTX-native raw-cluster-address carrier", - policy.id - ), - } - ensure!( - policy.rust_module == "barrier" - && policy.rust_name == recipe.id - && policy.rust_arguments == recipe.rust_arguments - && policy.rust_result == recipe.rust_result - && !policy.safe - && policy.must_use == recipe.must_use - && policy.public_rust_path == format!("cuda_intrinsics::barrier::{}", recipe.id) - && policy.compatibility_rust_paths == [format!("cuda_device::barrier::{}", recipe.id)], - "{} Rust API does not match its closed extended-mbarrier recipe", - policy.id - ); - ensure!( - policy.dialect_op_type == recipe.dialect_op_type - && policy.dialect_op_name == recipe.dialect_op_name - && policy.dialect_operands == recipe.dialect_operands - && policy.dialect_results == recipe.dialect_results - && policy.llvm_arguments == recipe.llvm_arguments - && policy.llvm_results == recipe.llvm_results - && policy.lowering == "generated_mbarrier_extended_inline_ptx", - "{} carrier or lowering does not match its closed extended-mbarrier recipe", - policy.id - ); - ensure!( - !policy.pure - && policy.memory == "read_write" - && policy.convergent - && policy.execution_scope == recipe.execution_scope, - "{} convergence, memory clobber, or execution scope changed", - policy.id - ); - ensure!( - policy.minimum_ptx == recipe.minimum_ptx - && policy.minimum_sm.as_deref() == Some(recipe.minimum_sm) - && policy.ptx_result == recipe.ptx_result - && policy.targets == "all" - && policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section == recipe.ptx_isa_section - && policy.ptx_isa_url == recipe.ptx_isa_url, - "{} target floor or PTX provenance changed", - policy.id - ); - ensure!( - policy.expected_ptx == recipe.expected_ptx, - "{} expected PTX does not match its closed extended-mbarrier recipe", - policy.id - ); - ensure_exact_inline_ptx_backends( - policy, - [ - ( - IntrinsicBackend::LlvmNvptx, - mbarrier_extended_backend_floor(recipe.operation, IntrinsicBackend::LlvmNvptx).0, - Some( - mbarrier_extended_backend_floor(recipe.operation, IntrinsicBackend::LlvmNvptx) - .1, - ), - ), - ( - IntrinsicBackend::LibNvvm, - mbarrier_extended_backend_floor(recipe.operation, IntrinsicBackend::LibNvvm).0, - Some( - mbarrier_extended_backend_floor(recipe.operation, IntrinsicBackend::LibNvvm).1, - ), - ), - ], - "extended-mbarrier", - )?; - ensure_no_other_family_contract(policy, "extended mbarrier")?; - Ok(()) -} - -fn expand_wgmma_control_admission( - admission: &WgmmaControlAdmission, -) -> Result> { - ensure!( - admission.runtime_validation == RuntimeValidation::Unexecuted, - "WGMMA-control runtime validation may be marked executed only with GPU evidence" - ); - ensure!( - !admission.llvm_evidence_profile.trim().is_empty() - && !admission.libnvvm_evidence_profile.trim().is_empty(), - "compact WGMMA-control admission requires both backend evidence profiles" - ); - let expected_modes = [ - WgmmaControlMode::Fence, - WgmmaControlMode::CommitGroup, - WgmmaControlMode::WaitGroup, - ]; - let actual_modes = admission - .variants - .iter() - .map(|variant| variant.mode) - .collect::>(); - ensure!( - actual_modes == expected_modes, - "compact WGMMA-control admission must contain each reviewed mode once in canonical order" - ); - - admission - .variants - .iter() - .map(|variant| { - let recipe = wgmma_control_recipe(variant.mode); - ensure!( - variant.abi_id == recipe.abi_id, - "{} must keep reserved ABI ID {}", - recipe.id, - recipe.abi_id - ); - let wait = recipe.mode == WgmmaControlMode::WaitGroup; - Ok(OverlayIntrinsic { - id: recipe.id.into(), - abi_id: variant.abi_id.clone(), - operation_key: recipe.operation_key.into(), - family: "wgmma_control".into(), - source: None, - source_record: Some(recipe.source_record.into()), - rust_module: "wgmma".into(), - rust_name: recipe.id.into(), - rust_arguments: if wait { vec!["u64".into()] } else { vec![] }, - rust_result: "()".into(), - safe: false, - must_use: false, - safe_allowlist_reason: None, - public_rust_path: format!("cuda_intrinsics::wgmma::{}", recipe.id), - compatibility_rust_paths: vec![recipe.compatibility_path.into()], - dialect_op_type: recipe.dialect_op_type.into(), - dialect_op_name: recipe.dialect_op_name.into(), - dialect_operands: if wait { vec!["i64".into()] } else { vec![] }, - dialect_results: vec![], - llvm_symbol: Some(recipe.llvm_symbol.into()), - resolved_llvm_symbol: None, - llvm_arguments: if wait { vec!["i64".into()] } else { vec![] }, - llvm_results: vec![], - pure: false, - memory: "read_write".into(), - convergent: true, - execution_scope: "warpgroup".into(), - minimum_ptx: "8.0".into(), - minimum_sm: None, - ptx_result: "()".into(), - targets: "sm_90a".into(), - ptx_isa_version: "9.3".into(), - ptx_isa_section: - "Asynchronous Warpgroup Level Matrix Instructions: WGMMA control".into(), - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#asynchronous-warpgroup-level-matrix-instructions".into(), - lowering: "generated_wgmma_control".into(), - backend_lowerings: vec![ - OverlayBackendLowering { - backend: IntrinsicBackend::LlvmNvptx, - mechanism: BackendLoweringMechanism::TypedNvvm, - evidence_profile: admission.llvm_evidence_profile.clone(), - targets: None, - minimum_ptx: Some("8.0".into()), - minimum_sm: None, - }, - OverlayBackendLowering { - backend: IntrinsicBackend::LibNvvm, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: admission.libnvvm_evidence_profile.clone(), - targets: None, - minimum_ptx: Some("8.0".into()), - minimum_sm: None, - }, - ], - packed_atomic: None, - redux: None, - vote: None, - active_mask: None, - warp_match: None, - warp_barrier: None, - warp_shuffle: None, - dot_product: None, - packed_alu: None, - integer_minmax: None, - packed_conversion: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cp_async_copy: None, - cp_async_control: None, - cp_async_mbarrier: None, - mbarrier_basic: None, - movmatrix: None, - mbarrier_extended: None, - register_mma: None, - sparse_mma: None, - prmt: None, - cluster_barrier: None, - wgmma_control: Some(WgmmaControl { - mode: recipe.mode, - adapter: recipe.adapter, - participation: WgmmaControlParticipation::WarpgroupAllThreadsSameInstruction, - }), - special_register: None, - debug_control: None, - cluster_memory: None, - clc: None, - tma: None, - tcgen05: None, - ldmatrix_variant: None, - ldmatrix_safety: None, - ldmatrix_adapter: None, - selected_address_space: None, - expected_ptx: InstructionPattern { - mnemonic: "wgmma".into(), - modifiers: recipe.suffix.split('.').map(str::to_owned).collect(), - operands: if wait { - vec![OperandPattern::Immediate] - } else { - vec![] - }, - }, - summary: recipe.summary.into(), - }) - }) - .collect() -} - -fn validate_wgmma_control_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, -) -> Result<()> { - let control = policy - .wgmma_control - .as_ref() - .with_context(|| format!("{} has no closed WGMMA-control contract", policy.id))?; - let recipe = wgmma_control_recipe(control.mode); - let wait = recipe.mode == WgmmaControlMode::WaitGroup; - ensure!( - control.adapter == recipe.adapter - && control.participation - == WgmmaControlParticipation::WarpgroupAllThreadsSameInstruction - && policy.id == recipe.id - && policy.abi_id == recipe.abi_id - && policy.operation_key == recipe.operation_key - && policy.source.is_none() - && policy.source_record.as_deref() == Some(recipe.source_record) - && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) - && policy.resolved_llvm_symbol.is_none(), - "{} identity or semantics do not match its closed WGMMA-control recipe", - policy.id - ); - ensure!( - policy.rust_module == "wgmma" - && policy.rust_name == recipe.id - && policy.rust_arguments - == if wait { - vec!["u64"] - } else { - Vec::<&str>::new() - } - && policy.rust_result == "()" - && !policy.safe - && !policy.must_use - && policy.public_rust_path == format!("cuda_intrinsics::wgmma::{}", recipe.id) - && policy.compatibility_rust_paths == [recipe.compatibility_path], - "{} Rust API does not match its closed WGMMA-control recipe", - policy.id - ); - let expected_arguments = if wait { - vec!["i64"] - } else { - Vec::<&str>::new() - }; - ensure!( - policy.dialect_op_type == recipe.dialect_op_type - && policy.dialect_op_name == recipe.dialect_op_name - && policy.dialect_operands == expected_arguments - && policy.dialect_results.is_empty() - && policy.llvm_arguments == expected_arguments - && policy.llvm_results.is_empty() - && policy.lowering == "generated_wgmma_control", - "{} carrier or lowering does not match its closed WGMMA-control recipe", - policy.id - ); - let expected_properties = if wait { - vec!["ImmArg", "IntrConvergent"] - } else { - vec!["IntrConvergent"] - }; - ensure!( - declaration.classes == ["SDPatternOperator", "Intrinsic"] - && declaration.properties == expected_properties - && !policy.pure - && policy.memory == "read_write" - && policy.convergent - && policy.execution_scope == "warpgroup", - "{} effects disagree with the imported WGMMA-control declaration", - policy.id - ); - ensure!( - policy.minimum_ptx == "8.0" - && policy.minimum_sm.is_none() - && policy.targets == "sm_90a" - && policy.ptx_result == "()" - && policy.ptx_isa_version == "9.3" - && policy.ptx_isa_url - == "https://docs.nvidia.com/cuda/parallel-thread-execution/#asynchronous-warpgroup-level-matrix-instructions", - "{} target floor or PTX provenance changed", - policy.id - ); - ensure!( - policy.expected_ptx.mnemonic == "wgmma" - && policy.expected_ptx.modifiers == recipe.suffix.split('.').collect::>() - && policy.expected_ptx.operands - == if wait { - vec![OperandPattern::Immediate] - } else { - vec![] - }, - "{} expected PTX does not match its exact WGMMA-control spelling", - policy.id - ); - let backend_pairs: BTreeSet<_> = policy - .backend_lowerings - .iter() - .map(|lowering| (lowering.backend, lowering.mechanism)) - .collect(); - ensure!( - policy.backend_lowerings.len() == 2 - && backend_pairs - == BTreeSet::from([ - ( - IntrinsicBackend::LlvmNvptx, - BackendLoweringMechanism::TypedNvvm, - ), - ( - IntrinsicBackend::LibNvvm, - BackendLoweringMechanism::InlinePtx, - ), - ]) - && policy.backend_lowerings.iter().all(|lowering| { - lowering.minimum_ptx.as_deref() == Some("8.0") - && lowering.minimum_sm.is_none() - && !lowering.evidence_profile.trim().is_empty() - }), - "{} must define exactly the reviewed WGMMA-control backend routes", - policy.id - ); - ensure_no_other_family_contract(policy, "WGMMA control")?; - Ok(()) -} - -fn ensure_exact_inline_ptx_backends( - policy: &OverlayIntrinsic, - requirements: [(IntrinsicBackend, &str, Option<&str>); 2], - family: &str, -) -> Result<()> { - let backend_pairs: BTreeSet<_> = policy - .backend_lowerings - .iter() - .map(|lowering| (lowering.backend, lowering.mechanism)) - .collect(); - ensure!( - policy.backend_lowerings.len() == 2 - && backend_pairs - == BTreeSet::from([ - ( - IntrinsicBackend::LlvmNvptx, - BackendLoweringMechanism::InlinePtx, - ), - ( - IntrinsicBackend::LibNvvm, - BackendLoweringMechanism::InlinePtx, - ), - ]), - "{} must define exactly two reviewed {family} inline-PTX routes", - policy.id - ); - let requirements: BTreeMap<_, _> = requirements - .into_iter() - .map(|(backend, ptx, minimum_sm)| (backend, (ptx, minimum_sm))) - .collect(); - for lowering in &policy.backend_lowerings { - let (minimum_ptx, minimum_sm) = requirements[&lowering.backend]; - ensure!( - lowering.minimum_ptx.as_deref() == Some(minimum_ptx) - && lowering.minimum_sm.as_deref() == minimum_sm - && !lowering.evidence_profile.trim().is_empty(), - "{} backend {:?} does not carry its exact {family} floor", - policy.id, - lowering.backend - ); - } - Ok(()) -} - -fn validate_movmatrix_policy(policy: &OverlayIntrinsic, source: &IntrinsicSource) -> Result<()> { - let contract = policy - .movmatrix - .as_ref() - .context("movmatrix requires its closed contract")?; - ensure!( - policy.id == "movmatrix_trans_b16" - && policy.operation_key == "movmatrix.m8n8.trans.b16" - && matches!( - source, - IntrinsicSource::PtxNative { instruction } - if instruction == "movmatrix.sync.aligned.m8n8.trans.b16" - ) - && policy.rust_module == "matrix" - && policy.rust_name == "movmatrix_trans_b16" - && policy.rust_arguments == ["u32"] - && policy.rust_result == "u32" - && !policy.safe - && policy.must_use - && policy.public_rust_path == "cuda_intrinsics::matrix::movmatrix_trans_b16" - && policy.compatibility_rust_paths == ["cuda_device::wmma::movmatrix_trans_b16"] - && policy.dialect_op_type == "MovmatrixTransB16Op" - && policy.dialect_op_name == "nvvm.movmatrix_trans_b16" - && policy.dialect_operands == ["i32"] - && policy.dialect_results == ["i32"] - && !policy.pure - && policy.memory == "inaccessible_read_write" - && policy.convergent - && policy.execution_scope == "warp" - && policy.minimum_ptx == "7.8" - && policy.minimum_sm.as_deref() == Some("sm_75") - && policy.ptx_result == "u32" - && policy.targets == "all" - && policy.lowering == "generated_movmatrix_inline_ptx" - && policy.expected_ptx - == InstructionPattern { - mnemonic: "movmatrix".into(), - modifiers: ["sync", "aligned", "m8n8", "trans", "b16"] - .into_iter() - .map(str::to_owned) - .collect(), - operands: vec![OperandPattern::Register, OperandPattern::Register], - }, - "{} is outside the closed movmatrix recipe", - policy.id - ); - ensure!( - contract.participation == MovmatrixParticipation::AllWarpLanesSameInstructionNoExitedLanes - && contract.adapter == MovmatrixAdapter::PackedB16x2U32ToPackedB16x2U32, - "{} has an unreviewed movmatrix safety or adapter contract", - policy.id - ); - ensure_exact_inline_ptx_backends( - policy, - [ - (IntrinsicBackend::LlvmNvptx, "7.8", Some("sm_75")), - (IntrinsicBackend::LibNvvm, "7.8", Some("sm_75")), - ], - "movmatrix", - )?; - ensure_no_other_family_contract(policy, "movmatrix")?; - Ok(()) -} - -struct IntegerMinMaxRecipe { - id: &'static str, - abi_id: &'static str, - operation_key: &'static str, - rust_module: &'static str, - rust_name: &'static str, - /// Rust and PTX carrier type: `i32` for the scalar forms, `u32` for the - /// packed pairs. - scalar: &'static str, - dialect_op_type: &'static str, - dialect_op_name: &'static str, - ptx_mnemonic: &'static str, - modifiers: &'static [&'static str], -} - -/// Returns the closed recipe for an extended integer min/max variant. -/// -/// `None` means the combination is outside the family: `.relu` exists only -/// for the signed formats, and the plain scalar forms are ordinary codegen -/// (`min.s32`/`max.s32` from generic Rust min/max), not intrinsics. -fn integer_minmax_recipe( - format: IntegerMinMaxFormat, - operation: IntegerMinMaxOperation, - relu: bool, -) -> Option { - use IntegerMinMaxFormat as Format; - use IntegerMinMaxOperation as Operation; - Some(match (format, operation, relu) { - (Format::S32, Operation::Min, true) => IntegerMinMaxRecipe { - id: "min_relu_s32", - abi_id: "i0987", - operation_key: "integer.minmax.s32.min.relu", - rust_module: "int", - rust_name: "min_relu_s32", - scalar: "i32", - dialect_op_type: "MinReluS32Op", - dialect_op_name: "nvvm.min_relu_s32", - ptx_mnemonic: "min.relu.s32", - modifiers: &["relu", "s32"], - }, - (Format::S32, Operation::Max, true) => IntegerMinMaxRecipe { - id: "max_relu_s32", - abi_id: "i0988", - operation_key: "integer.minmax.s32.max.relu", - rust_module: "int", - rust_name: "max_relu_s32", - scalar: "i32", - dialect_op_type: "MaxReluS32Op", - dialect_op_name: "nvvm.max_relu_s32", - ptx_mnemonic: "max.relu.s32", - modifiers: &["relu", "s32"], - }, - (Format::S16x2, Operation::Min, false) => IntegerMinMaxRecipe { - id: "min_s16x2", - abi_id: "i0989", - operation_key: "integer.minmax.s16x2.min", - rust_module: "i16x2", - rust_name: "min_s16x2", - scalar: "u32", - dialect_op_type: "MinS16x2Op", - dialect_op_name: "nvvm.min_s16x2", - ptx_mnemonic: "min.s16x2", - modifiers: &["s16x2"], - }, - (Format::S16x2, Operation::Max, false) => IntegerMinMaxRecipe { - id: "max_s16x2", - abi_id: "i0990", - operation_key: "integer.minmax.s16x2.max", - rust_module: "i16x2", - rust_name: "max_s16x2", - scalar: "u32", - dialect_op_type: "MaxS16x2Op", - dialect_op_name: "nvvm.max_s16x2", - ptx_mnemonic: "max.s16x2", - modifiers: &["s16x2"], - }, - (Format::U16x2, Operation::Min, false) => IntegerMinMaxRecipe { - id: "min_u16x2", - abi_id: "i0991", - operation_key: "integer.minmax.u16x2.min", - rust_module: "i16x2", - rust_name: "min_u16x2", - scalar: "u32", - dialect_op_type: "MinU16x2Op", - dialect_op_name: "nvvm.min_u16x2", - ptx_mnemonic: "min.u16x2", - modifiers: &["u16x2"], - }, - (Format::U16x2, Operation::Max, false) => IntegerMinMaxRecipe { - id: "max_u16x2", - abi_id: "i0992", - operation_key: "integer.minmax.u16x2.max", - rust_module: "i16x2", - rust_name: "max_u16x2", - scalar: "u32", - dialect_op_type: "MaxU16x2Op", - dialect_op_name: "nvvm.max_u16x2", - ptx_mnemonic: "max.u16x2", - modifiers: &["u16x2"], - }, - (Format::S16x2, Operation::Min, true) => IntegerMinMaxRecipe { - id: "min_relu_s16x2", - abi_id: "i0993", - operation_key: "integer.minmax.s16x2.min.relu", - rust_module: "i16x2", - rust_name: "min_relu_s16x2", - scalar: "u32", - dialect_op_type: "MinReluS16x2Op", - dialect_op_name: "nvvm.min_relu_s16x2", - ptx_mnemonic: "min.relu.s16x2", - modifiers: &["relu", "s16x2"], - }, - (Format::S16x2, Operation::Max, true) => IntegerMinMaxRecipe { - id: "max_relu_s16x2", - abi_id: "i0994", - operation_key: "integer.minmax.s16x2.max.relu", - rust_module: "i16x2", - rust_name: "max_relu_s16x2", - scalar: "u32", - dialect_op_type: "MaxReluS16x2Op", - dialect_op_name: "nvvm.max_relu_s16x2", - ptx_mnemonic: "max.relu.s16x2", - modifiers: &["relu", "s16x2"], - }, - (Format::S32, _, false) | (Format::U16x2, _, true) => return None, - }) -} - -fn integer_minmax_isa_reference(operation: IntegerMinMaxOperation) -> (&'static str, &'static str) { - match operation { - IntegerMinMaxOperation::Min => ( - "9.7.1.13 Integer Arithmetic Instructions: min", - "https://docs.nvidia.com/cuda/parallel-thread-execution/#integer-arithmetic-instructions-min", - ), - IntegerMinMaxOperation::Max => ( - "9.7.1.14 Integer Arithmetic Instructions: max", - "https://docs.nvidia.com/cuda/parallel-thread-execution/#integer-arithmetic-instructions-max", - ), - } -} - -fn validate_integer_minmax_policy( - policy: &OverlayIntrinsic, - source: &IntrinsicSource, - declaration: Option<&ImportedIntrinsic>, -) -> Result<()> { - let minmax = policy - .integer_minmax - .as_ref() - .with_context(|| format!("{} has no closed integer-min/max contract", policy.id))?; - let recipe = integer_minmax_recipe(minmax.format, minmax.operation, minmax.relu) - .with_context(|| format!("{} is outside the closed integer-min/max recipe", policy.id))?; - ensure!( - policy.id == recipe.id - && policy.abi_id == recipe.abi_id - && policy.operation_key == recipe.operation_key, - "{} integer-min/max identity does not match its closed operation recipe", - policy.id - ); - ensure!( - policy.rust_module == recipe.rust_module - && policy.rust_name == recipe.rust_name - && policy.rust_arguments == [recipe.scalar, recipe.scalar] - && policy.rust_result == recipe.scalar - && policy.safe - && policy.must_use - && policy - .safe_allowlist_reason - .as_deref() - .is_some_and(|reason| !reason.trim().is_empty()) - && policy.public_rust_path - == format!( - "cuda_intrinsics::{}::{}", - recipe.rust_module, recipe.rust_name - ) - && policy.compatibility_rust_paths - == [format!( - "cuda_device::{}::{}", - recipe.rust_module, recipe.rust_name - )], - "{} must preserve its reviewed safe integer-min/max API", - policy.id - ); - ensure!( - policy.dialect_op_type == recipe.dialect_op_type - && policy.dialect_op_name == recipe.dialect_op_name - && policy.dialect_operands == ["i32", "i32"] - && policy.dialect_results == ["i32"] - && policy.lowering == "generated_integer_minmax_inline_ptx", - "{} is outside the closed integer-min/max dialect and lowering recipe", - policy.id - ); - ensure!( - policy.pure - && policy.memory == "none" - && !policy.convergent - && policy.execution_scope == "thread" - && policy.minimum_ptx == "8.0" - && policy.minimum_sm.as_deref() == Some("sm_90") - && policy.ptx_result == recipe.scalar - && policy.targets == "all" - && minmax.native_minimum_sm == 90, - "{} integer-min/max effects, carrier, or target floor disagree", - policy.id - ); - let (isa_section, isa_url) = integer_minmax_isa_reference(minmax.operation); - ensure!( - policy.ptx_isa_version == "9.3" - && policy.ptx_isa_section == isa_section - && policy.ptx_isa_url == isa_url, - "{} integer-min/max PTX provenance does not match its reviewed instruction section", - policy.id - ); - ensure!( - policy.expected_ptx.mnemonic - == recipe.ptx_mnemonic.split('.').next().expect("PTX mnemonic") - && policy.expected_ptx.modifiers == recipe.modifiers - && policy.expected_ptx.operands == vec![OperandPattern::Register; 3], - "{} expected PTX does not match its exact integer-min/max instruction", - policy.id - ); - ensure!( - source - == &IntrinsicSource::PtxNative { - instruction: recipe.ptx_mnemonic.to_owned(), - } - && declaration.is_none(), - "{} integer-min/max source does not match its PTX-native recipe", - policy.id - ); - ensure!( - policy.backend_lowerings.len() == 1 - && policy.backend_lowerings[0].backend == IntrinsicBackend::LlvmNvptx - && policy.backend_lowerings[0].mechanism == BackendLoweringMechanism::InlinePtx - && policy.backend_lowerings[0].targets.is_none() - && policy.backend_lowerings[0].minimum_ptx.is_none() - && policy.backend_lowerings[0].minimum_sm.is_none() - && !policy.backend_lowerings[0] - .evidence_profile - .trim() - .is_empty(), - "{} integer-min/max backend route changed", - policy.id - ); - ensure_no_other_family_contract(policy, "integer_minmax")?; - Ok(()) -} - -fn ensure_no_other_family_contract(policy: &OverlayIntrinsic, family: &str) -> Result<()> { - ensure!( - policy.packed_atomic.is_none() - && policy.redux.is_none() - && policy.vote.is_none() - && policy.active_mask.is_none() - && policy.warp_match.is_none() - && policy.warp_barrier.is_none() - && policy.warp_shuffle.is_none() - && policy.dot_product.is_none() - && policy.ldmatrix_variant.is_none() - && policy.ldmatrix_safety.is_none() - && policy.ldmatrix_adapter.is_none() - && (policy.family == "tma" || policy.selected_address_space.is_none()) - && (policy.family == "packed_alu") == policy.packed_alu.is_some() - && (policy.family == "integer_minmax") == policy.integer_minmax.is_some() - && (policy.family == "packed_conversion") == policy.packed_conversion.is_some() - && (policy.family == "scalar_conversion") == policy.scalar_conversion.is_some() - && (policy.family == "scalar_arithmetic") == policy.scalar_arithmetic.is_some() - && (policy.family == "scalar_math") == policy.scalar_math.is_some() - && (policy.family == "extended_minmax") == policy.extended_minmax.is_some() - && (policy.family == "cp_async_copy") == policy.cp_async_copy.is_some() - && (policy.family == "cp_async_control") == policy.cp_async_control.is_some() - && (policy.family == "cp_async_mbarrier") == policy.cp_async_mbarrier.is_some() - && (policy.family == "mbarrier_basic") == policy.mbarrier_basic.is_some() - && (policy.family == "movmatrix") == policy.movmatrix.is_some() - && (policy.family == "mbarrier_extended") == policy.mbarrier_extended.is_some() - && (policy.family == "register_mma") == policy.register_mma.is_some() - && (policy.family == "sparse_mma") == policy.sparse_mma.is_some() - && (policy.family == "prmt") == policy.prmt.is_some() - && (policy.family == "cluster_barrier") == policy.cluster_barrier.is_some() - && (policy.family == "wgmma_control") == policy.wgmma_control.is_some() - && (policy.family == "debug_control") == policy.debug_control.is_some() - && (policy.family == "cluster_memory") == policy.cluster_memory.is_some() - && (policy.family == "clc") == policy.clc.is_some() - && (policy.family == "tma") == policy.tma.is_some() - && (policy.family == "tcgen05") == policy.tcgen05.is_some(), - "{} mixes another generated-family contract with {family}", - policy.id - ); - Ok(()) -} - -fn parse_ptx_version(value: &str, intrinsic_id: &str) -> Result { - value - .parse() - .map_err(|reason: String| anyhow::anyhow!("{intrinsic_id} minimum_ptx {value:?}: {reason}")) -} - -fn parse_hardware_target(policy: &OverlayIntrinsic) -> Result { - parse_hardware_target_fields(&policy.id, &policy.targets, policy.minimum_sm.as_deref()) -} - -/// Resolve every selector contract without merging target sets. -pub(crate) fn resolve_target_contracts( - intrinsic_id: &str, - contracts: &[TargetContract], -) -> Result { - ensure!( - !contracts.is_empty(), - "{intrinsic_id} has no reviewed target contracts" - ); - - let expected_selector_names = contracts[0] - .selectors - .iter() - .map(|selector| selector.name.as_str()) - .collect::>(); - let mut resolved_contracts = Vec::with_capacity(contracts.len()); - for contract in contracts { - validate_selector_bindings( - intrinsic_id, - "target-contract selectors", - &contract.selectors, - )?; - ensure!( - contract - .selectors - .iter() - .map(|selector| selector.name.as_str()) - .eq(expected_selector_names.iter().copied()), - "{intrinsic_id} target contracts must use one closed selector schema" - ); - ensure!( - !contract.alternatives.is_empty(), - "{intrinsic_id} target contract {:?} has no alternatives", - contract.selectors - ); - - let alternatives = contract - .alternatives - .iter() - .map(|alternative| { - Ok(CatalogTargetAlternative { - minimum_ptx: parse_ptx_version(&alternative.minimum_ptx, intrinsic_id)?, - hardware: parse_target_contract_hardware(intrinsic_id, &alternative.target)?, - }) - }) - .collect::>>()?; - ensure!( - alternatives - .windows(2) - .all(|pair| target_hardware_sort_key(pair[0].hardware) - < target_hardware_sort_key(pair[1].hardware)), - "{intrinsic_id} target contract {:?} alternatives must have unique, sorted hardware targets", - contract.selectors - ); - let minimum_count = alternatives - .iter() - .filter(|alternative| { - matches!( - alternative.hardware, - CatalogHardwareAlternative::MinimumSm { .. } - ) - }) - .count(); - ensure!( - minimum_count == 0 || alternatives.len() == 1, - "{intrinsic_id} target contract {:?} cannot mix a minimum-SM range with other alternatives", - contract.selectors - ); - resolved_contracts.push(CatalogTargetContract { - selectors: contract.selectors.clone(), - alternatives, - }); - } - ensure!( - resolved_contracts - .windows(2) - .all(|pair| pair[0].selectors < pair[1].selectors), - "{intrinsic_id} target contracts must have unique, sorted selector bindings" - ); - - let minimum_ptx = resolved_contracts - .iter() - .flat_map(|contract| contract.alternatives.iter()) - .map(|alternative| alternative.minimum_ptx) - .min() - .unwrap(); - Ok(CatalogTargetRequirement { - minimum_ptx, - hardware: CatalogHardwareTarget::TargetMatrix { - contracts: resolved_contracts, - }, - }) -} - -/// Resolve one fixed selector tuple from a closed target matrix. -pub(crate) fn resolve_target_contract( - intrinsic_id: &str, - selected: &[TargetSelectorBinding], - contracts: &[TargetContract], -) -> Result { - validate_selector_bindings(intrinsic_id, "selected target selectors", selected)?; - let resolved = resolve_target_contracts(intrinsic_id, contracts)?; - let CatalogHardwareTarget::TargetMatrix { contracts } = resolved.hardware else { - unreachable!("target-contract resolution always returns a matrix") - }; - - let matching = contracts - .into_iter() - .filter(|contract| contract.selectors == selected) - .collect::>(); - ensure!( - matching.len() == 1, - "{intrinsic_id} selected target selectors {:?} must match exactly one reviewed contract", - selected - ); - let contract = matching.into_iter().next().unwrap(); - let minimum_ptx = contract - .alternatives - .iter() - .map(|alternative| alternative.minimum_ptx) - .min() - .unwrap(); - Ok(CatalogTargetRequirement { - minimum_ptx, - hardware: CatalogHardwareTarget::TargetMatrix { - contracts: vec![contract], - }, - }) -} - -fn validate_selector_bindings( - intrinsic_id: &str, - label: &str, - selectors: &[TargetSelectorBinding], -) -> Result<()> { - for selector in selectors { - ensure!( - is_target_selector_token(&selector.name) && is_target_selector_token(&selector.value), - "{intrinsic_id} {label} must use lowercase snake-case names and values" - ); - } - ensure!( - selectors.windows(2).all(|pair| pair[0].name < pair[1].name), - "{intrinsic_id} {label} must have unique, sorted names" - ); - Ok(()) -} - -fn is_target_selector_token(value: &str) -> bool { - value - .bytes() - .next() - .is_some_and(|byte| byte.is_ascii_lowercase()) - && value - .bytes() - .all(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'_') - && !value.ends_with('_') - && !value.contains("__") -} - -fn parse_target_contract_hardware( - intrinsic_id: &str, - target: &str, -) -> Result { - if let Some(minimum) = target.strip_suffix('+') { - let sm = parse_sm_spelling(intrinsic_id, "target contract", minimum, None)?; - return Ok(CatalogHardwareAlternative::MinimumSm { sm }); - } - parse_exact_hardware_alternative(intrinsic_id, target) -} - -fn target_hardware_sort_key(hardware: CatalogHardwareAlternative) -> (u16, u8) { - match hardware { - CatalogHardwareAlternative::MinimumSm { sm } => (sm, 0), - CatalogHardwareAlternative::ExactArchitecture { sm } => (sm, 1), - CatalogHardwareAlternative::FamilyTarget { sm } => (sm, 2), - } -} - -fn parse_hardware_target_fields( - intrinsic_id: &str, - targets: &str, - minimum_sm: Option<&str>, -) -> Result { - if targets == "all" { - let Some(minimum_sm) = minimum_sm else { - return Ok(CatalogHardwareTarget::All); - }; - let sm = parse_sm_spelling(intrinsic_id, "minimum_sm", minimum_sm, None)?; - return Ok(CatalogHardwareTarget::AnyOf { - alternatives: vec![CatalogHardwareAlternative::MinimumSm { sm }], - }); - } - - if targets.contains('|') { - ensure!( - minimum_sm.is_none(), - "{} target alternatives {:?} cannot be combined with minimum_sm", - intrinsic_id, - targets - ); - let spellings = targets.split('|').collect::>(); - ensure!( - spellings.len() >= 2, - "{} target alternatives must contain at least two targets", - intrinsic_id - ); - ensure!( - spellings.windows(2).all(|pair| pair[0] < pair[1]), - "{} target alternatives must be unique and sorted", - intrinsic_id - ); - return Ok(CatalogHardwareTarget::AnyOf { - alternatives: spellings - .into_iter() - .map(|spelling| parse_exact_hardware_alternative(intrinsic_id, spelling)) - .collect::>>()?, - }); - } - - ensure!( - minimum_sm.is_none(), - "{} exact targets {:?} cannot be combined with minimum_sm", - intrinsic_id, - targets - ); - Ok(CatalogHardwareTarget::AnyOf { - alternatives: vec![parse_exact_hardware_alternative(intrinsic_id, targets)?], - }) -} - -fn parse_exact_hardware_alternative( - intrinsic_id: &str, - target: &str, -) -> Result { - let suffix = target - .chars() - .last() - .filter(|suffix| matches!(suffix, 'a' | 'f')); - let Some(suffix) = suffix else { - bail!( - "{} targets {:?} must be `all`, exact `sm_Na`, or family `sm_Nf`", - intrinsic_id, - target - ); - }; - let sm = parse_sm_spelling(intrinsic_id, "targets", target, Some(suffix))?; - Ok(match suffix { - 'a' => CatalogHardwareAlternative::ExactArchitecture { sm }, - 'f' => CatalogHardwareAlternative::FamilyTarget { sm }, - _ => unreachable!(), - }) -} - -fn parse_sm_spelling( - intrinsic_id: &str, - field: &str, - value: &str, - suffix: Option, -) -> Result { - let body = value.strip_prefix("sm_").with_context(|| { - format!("{intrinsic_id} {field} {value:?} must use canonical sm_NN spelling") - })?; - let digits = match suffix { - Some(suffix) => body.strip_suffix(suffix).with_context(|| { - format!("{intrinsic_id} {field} {value:?} has the wrong target suffix") - })?, - None => body, - }; - ensure!( - matches!(digits.len(), 2 | 3) && digits.bytes().all(|byte| byte.is_ascii_digit()), - "{} {} {:?} must use canonical sm_NN{} spelling", - intrinsic_id, - field, - value, - suffix.map_or("", |suffix| if suffix == 'a' { "a" } else { "f" }) - ); - let sm: u16 = digits - .parse() - .with_context(|| format!("{intrinsic_id} {field} target is too large"))?; - let canonical = match suffix { - Some(suffix) => format!("sm_{sm}{suffix}"), - None => format!("sm_{sm}"), - }; - ensure!( - sm > 0 && canonical == value, - "{} {} {:?} is not canonical", - intrinsic_id, - field, - value - ); - Ok(sm) -} - -fn backend_target_requirement( - policy: &OverlayIntrinsic, - lowering: &crate::model::OverlayBackendLowering, -) -> Result { - if let Some(mma) = policy - .tcgen05 - .as_ref() - .and_then(|tcgen05| tcgen05.mma.as_ref()) - { - return Ok(match lowering.backend { - IntrinsicBackend::LlvmNvptx => mma.llvm_target.clone(), - IntrinsicBackend::LibNvvm => mma.libnvvm_target.clone(), - }); - } - let minimum_ptx = lowering - .minimum_ptx - .as_deref() - .unwrap_or(&policy.minimum_ptx); - let targets = lowering.targets.as_deref().unwrap_or(&policy.targets); - let minimum_sm = lowering.minimum_sm.as_deref().or_else(|| { - if lowering.targets.is_none() { - policy.minimum_sm.as_deref() - } else { - None - } - }); - Ok(CatalogTargetRequirement { - minimum_ptx: parse_ptx_version(minimum_ptx, &policy.id)?, - hardware: parse_hardware_target_fields(&policy.id, targets, minimum_sm)?, - }) -} - -fn read_evidence_file(path: &Path) -> Result { - let bytes = fs::read(path).with_context(|| format!("read {}", path.display()))?; - parse_evidence_bytes(&bytes, &path.display().to_string()) -} - -fn parse_evidence_bytes(bytes: &[u8], source: &str) -> Result { - #[derive(serde::Deserialize)] - struct Schema { - schema: u32, - } - - let schema: Schema = - serde_json::from_slice(bytes).with_context(|| format!("parse JSON {source}"))?; - match schema.schema { - 2..=5 => serde_json::from_slice(bytes) - .with_context(|| format!("parse legacy evidence JSON {source}")), - 6 => { - let file: EvidenceFileV6 = serde_json::from_slice(bytes) - .with_context(|| format!("parse matrix evidence JSON {source}"))?; - expand_evidence_file_v6(file) - .with_context(|| format!("expand matrix evidence {source}")) - } - _ => bail!("unsupported evidence schema in {source}"), - } -} - -fn expand_evidence_file_v6(file: EvidenceFileV6) -> Result { - ensure!(file.schema == 6, "matrix evidence must use schema 6"); - ensure!( - !file.records.is_empty() || !file.matrices.is_empty(), - "schema-6 evidence contains no records or matrices" - ); - reject_default_placeholders(&file.defaults, "evidence defaults", false)?; - - let mut fixture_by_id = BTreeMap::new(); - let mut fixture_coverage = BTreeMap::new(); - for fixture in &file.fixtures { - ensure!( - is_safe_matrix_token(&fixture.id), - "evidence fixture ID {:?} is not a safe token", - fixture.id - ); - ensure!( - fixture.coverage_count > 0 && !fixture.stages.is_empty(), - "evidence fixture {} has no coverage or stages", - fixture.id - ); - reject_stage_placeholders(&fixture.stages, &format!("fixture {}", fixture.id))?; - ensure!( - fixture_by_id.insert(fixture.id.as_str(), fixture).is_none(), - "duplicate evidence fixture ID {}", - fixture.id - ); - fixture_coverage.insert(fixture.id.as_str(), 0usize); - } - - let mut records = file.records; - let mut record_ids = BTreeSet::new(); - for record in &records { - ensure!( - record_ids.insert(record.id.clone()), - "duplicate expanded evidence ID {}", - record.id - ); - validate_stage_pairs(&record.stages, &record.id)?; - } - - for matrix in &file.matrices { - let (expanded, referenced_fixtures) = - expand_evidence_matrix(matrix, &file.defaults, &fixture_by_id)?; - for fixture_id in referenced_fixtures { - *fixture_coverage - .get_mut(fixture_id.as_str()) - .expect("validated fixture reference") += expanded.len(); - } - for record in expanded { - ensure!( - record_ids.insert(record.id.clone()), - "duplicate expanded evidence ID {}", - record.id - ); - records.push(record); - } - } - - for fixture in &file.fixtures { - let actual = fixture_coverage[fixture.id.as_str()]; - ensure!( - actual > 0, - "evidence fixture {} is not referenced by any matrix", - fixture.id - ); - ensure!( - actual == fixture.coverage_count, - "evidence fixture {} covers {actual} expanded records, expected {}", - fixture.id, - fixture.coverage_count - ); - } - - Ok(EvidenceFile { - schema: file.schema, - backend_profile: file.backend_profile, - backend_kind: file.backend_kind, - llvm_revision: file.llvm_revision, - backend_version: file.backend_version, - backend_sha256: file.backend_sha256, - artifact_path: file.artifact_path, - build_id_prefix: file.build_id_prefix, - nvvm_ir_version: file.nvvm_ir_version, - debug_ir_version: file.debug_ir_version, - records, - }) -} - -fn expand_evidence_matrix( - matrix: &EvidenceMatrix, - defaults: &EvidenceRecordDefaults, - fixtures: &BTreeMap<&str, &crate::model::EvidenceFixture>, -) -> Result<(Vec, Vec)> { - ensure!(!matrix.axes.is_empty(), "evidence matrix has no axes"); - let mut previous_axis: Option<&str> = None; - let mut product_count = 1usize; - let mut bindings = vec![BTreeMap::::new()]; - for axis in &matrix.axes { - ensure!( - is_safe_matrix_token(&axis.name), - "evidence matrix axis {:?} is not a safe token", - axis.name - ); - if let Some(previous) = previous_axis { - ensure!( - previous < axis.name.as_str(), - "evidence matrix axes must be unique and sorted: {} follows {}", - axis.name, - previous - ); - } - previous_axis = Some(&axis.name); - ensure!( - !axis.values.is_empty(), - "evidence matrix axis {} has no values", - axis.name - ); - let mut values = BTreeSet::new(); - for value in &axis.values { - ensure!( - is_safe_matrix_token(value), - "evidence matrix axis {} has unsafe value {:?}", - axis.name, - value - ); - ensure!( - values.insert(value.as_str()), - "evidence matrix axis {} has duplicate value {:?}", - axis.name, - value - ); - } - product_count = product_count - .checked_mul(axis.values.len()) - .context("evidence matrix product count overflow")?; - let mut next = Vec::with_capacity(product_count); - for binding in bindings { - for value in &axis.values { - let mut expanded = binding.clone(); - expanded.insert(axis.name.clone(), value.clone()); - next.push(expanded); - } - } - bindings = next; - } - ensure!( - product_count == matrix.product_count, - "evidence matrix expands to {product_count} records, expected {}", - matrix.product_count - ); - ensure!( - !matrix.fixtures.is_empty(), - "evidence matrix references no shared fixture" - ); - - let mut fixture_ids = BTreeSet::new(); - let mut previous_fixture: Option<&str> = None; - let mut fixture_stages = Vec::new(); - for fixture_id in &matrix.fixtures { - if let Some(previous) = previous_fixture { - ensure!( - previous < fixture_id.as_str(), - "evidence matrix fixtures must be unique and sorted: {fixture_id} follows {previous}" - ); - } - previous_fixture = Some(fixture_id); - let fixture = fixtures - .get(fixture_id.as_str()) - .with_context(|| format!("evidence matrix references unknown fixture {fixture_id}"))?; - fixture_ids.insert(fixture_id.clone()); - fixture_stages.extend(fixture.stages.iter().cloned()); - } - - reject_default_placeholders(&matrix.template.facts, "matrix template facts", true)?; - validate_matrix_identity(&matrix.template)?; - let axis_names = matrix - .axes - .iter() - .map(|axis| axis.name.as_str()) - .collect::>(); - let mut used_axes = BTreeSet::new(); - let mut records = Vec::with_capacity(product_count); - for binding in &bindings { - let record = materialize_evidence_record( - &matrix.template, - defaults, - binding, - &mut used_axes, - &fixture_stages, - )?; - validate_stage_pairs(&record.stages, &record.id)?; - records.push(record); - } - for axis in axis_names { - ensure!( - used_axes.contains(axis), - "evidence matrix declares unused axis {axis}" - ); - } - Ok((records, fixture_ids.into_iter().collect())) -} - -fn validate_matrix_identity(template: &EvidenceMatrixTemplate) -> Result<()> { - ensure!( - !template.id.is_empty(), - "evidence matrix template has an empty ID" - ); - match (&template.source, &template.source_record) { - (Some(_), None) | (None, Some(_)) => {} - (Some(_), Some(_)) => bail!("evidence matrix template mixes source forms"), - (None, None) => bail!("evidence matrix template has no source"), - } - reject_disallowed_placeholder(&template.expected_ptx.mnemonic, "PTX mnemonic")?; - Ok(()) -} - -fn materialize_evidence_record( - template: &EvidenceMatrixTemplate, - defaults: &EvidenceRecordDefaults, - bindings: &BTreeMap, - used_axes: &mut BTreeSet, - fixture_stages: &[EvidenceStage], -) -> Result { - let id = expand_axis_placeholders(&template.id, bindings, used_axes, "evidence ID")?; - let source = template - .source - .as_ref() - .map(|source| expand_evidence_source(source, bindings, used_axes)) - .transpose()?; - let source_record = template - .source_record - .as_deref() - .map(|value| expand_axis_placeholders(value, bindings, used_axes, "source record")) - .transpose()?; - let llvm_symbol = template - .llvm_symbol - .as_deref() - .map(|value| expand_axis_placeholders(value, bindings, used_axes, "LLVM symbol")) - .transpose()?; - validate_expanded_matrix_identity( - &id, - source.as_ref(), - source_record.as_deref(), - llvm_symbol.as_deref(), - )?; - let resolved_llvm_symbol = select_fact( - &template.facts.resolved_llvm_symbol, - &defaults.resolved_llvm_symbol, - ) - .map(|value| expand_axis_placeholders(&value, bindings, used_axes, "resolved LLVM symbol")) - .transpose()?; - let mut expected_ptx = template.expected_ptx.clone(); - for modifier in &mut expected_ptx.modifiers { - *modifier = expand_axis_placeholders(modifier, bindings, used_axes, "PTX modifier")?; - } - for operand in &mut expected_ptx.operands { - if let OperandPattern::Exact { value } = operand { - *value = expand_axis_placeholders(value, bindings, used_axes, "exact PTX operand")?; - } - } - - let mut stages = defaults.stages.clone(); - stages.extend(template.facts.stages.iter().cloned()); - stages.extend(fixture_stages.iter().cloned()); - let target_triple = required_fact( - select_fact(&template.facts.target_triple, &defaults.target_triple), - &id, - "target triple", - )?; - let gpu_target = required_fact( - select_fact(&template.facts.gpu_target, &defaults.gpu_target), - &id, - "GPU target", - )?; - let ptx_feature = required_fact( - select_fact(&template.facts.ptx_feature, &defaults.ptx_feature), - &id, - "PTX feature", - )?; - let status = required_fact( - select_fact(&template.facts.status, &defaults.status), - &id, - "status", - )?; - let source_is_ptx_native = matches!(source, Some(IntrinsicSource::PtxNative { .. })); - Ok(EvidenceRecord { - id, - source, - source_record, - llvm_symbol, - resolved_llvm_symbol, - llvm_arguments: select_fact(&template.facts.llvm_arguments, &defaults.llvm_arguments) - .unwrap_or_default(), - llvm_results: select_fact(&template.facts.llvm_results, &defaults.llvm_results) - .unwrap_or_default(), - concrete_llvm_arguments: select_fact( - &template.facts.concrete_llvm_arguments, - &defaults.concrete_llvm_arguments, - ) - .unwrap_or_default(), - concrete_llvm_results: select_fact( - &template.facts.concrete_llvm_results, - &defaults.concrete_llvm_results, - ) - .unwrap_or_default(), - target_triple, - gpu_target, - ptx_feature, - status, - stages, - // Declaration canonicalization is a statement about an imported - // LLVM declaration; a PTX-native record has none, so the file-level - // default must not invent the fact for it. - declaration_attributes_canonicalized: if source_is_ptx_native { - template.facts.declaration_attributes_canonicalized - } else { - template - .facts - .declaration_attributes_canonicalized - .or(defaults.declaration_attributes_canonicalized) - }, - runtime_validation: template - .facts - .runtime_validation - .or(defaults.runtime_validation), - expected_ptx, - }) -} - -fn validate_expanded_matrix_identity( - id: &str, - source: Option<&IntrinsicSource>, - source_record: Option<&str>, - llvm_symbol: Option<&str>, -) -> Result<()> { - ensure!(!id.is_empty(), "expanded evidence has an empty ID"); - let imported_source = match (source, source_record) { - (Some(IntrinsicSource::LlvmImported { source_record }), None) => { - ensure!( - !source_record.is_empty(), - "expanded evidence {id} has an empty source record" - ); - true - } - (Some(IntrinsicSource::PtxNative { instruction }), None) => { - ensure!( - !instruction.is_empty(), - "expanded evidence {id} has an empty PTX source" - ); - false - } - (None, Some(source_record)) => { - ensure!( - !source_record.is_empty(), - "expanded evidence {id} has an empty source record" - ); - true - } - _ => unreachable!("matrix source shape was validated before expansion"), - }; - if imported_source { - ensure!( - llvm_symbol.is_some_and(|symbol| !symbol.is_empty()), - "expanded imported evidence {id} has no LLVM symbol" - ); - } else { - ensure!( - llvm_symbol.is_none(), - "expanded PTX-native evidence {id} invents an LLVM symbol" - ); - } - Ok(()) -} - -fn select_fact(specific: &Option, default: &Option) -> Option { - specific.clone().or_else(|| default.clone()) -} - -fn required_fact(value: Option, id: &str, field: &str) -> Result { - let value = value.with_context(|| format!("expanded evidence {id} has no {field}"))?; - ensure!( - !value.trim().is_empty(), - "expanded evidence {id} has an empty {field}" - ); - Ok(value) -} - -fn expand_evidence_source( - source: &IntrinsicSource, - bindings: &BTreeMap, - used_axes: &mut BTreeSet, -) -> Result { - Ok(match source { - IntrinsicSource::LlvmImported { source_record } => IntrinsicSource::LlvmImported { - source_record: expand_axis_placeholders( - source_record, - bindings, - used_axes, - "tagged source record", - )?, - }, - IntrinsicSource::PtxNative { instruction } => IntrinsicSource::PtxNative { - instruction: expand_axis_placeholders( - instruction, - bindings, - used_axes, - "PTX-native source", - )?, - }, - }) -} - -fn expand_axis_placeholders( - value: &str, - bindings: &BTreeMap, - used_axes: &mut BTreeSet, - field: &str, -) -> Result { - let mut output = String::with_capacity(value.len()); - let mut rest = value; - while let Some(position) = rest.find('$') { - output.push_str(&rest[..position]); - let placeholder = &rest[position..]; - ensure!( - placeholder.starts_with("${"), - "{field} contains malformed matrix placeholder {value:?}" - ); - let close = placeholder - .find('}') - .with_context(|| format!("{field} contains an unterminated matrix placeholder"))?; - let axis = &placeholder[2..close]; - ensure!( - is_safe_matrix_token(axis), - "{field} contains malformed matrix axis {axis:?}" - ); - let replacement = bindings - .get(axis) - .with_context(|| format!("{field} references unknown matrix axis {axis}"))?; - output.push_str(replacement); - used_axes.insert(axis.to_owned()); - rest = &placeholder[close + 1..]; - } - output.push_str(rest); - Ok(output) -} - -fn is_safe_matrix_token(value: &str) -> bool { - let mut bytes = value.bytes(); - bytes.next().is_some_and(|byte| byte.is_ascii_lowercase()) - && bytes.all(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'_') -} - -fn reject_default_placeholders( - defaults: &EvidenceRecordDefaults, - context: &str, - allow_resolved_symbol: bool, -) -> Result<()> { - if !allow_resolved_symbol && let Some(value) = &defaults.resolved_llvm_symbol { - reject_disallowed_placeholder(value, context)?; - } - for value in defaults - .target_triple - .iter() - .chain(defaults.gpu_target.iter()) - .chain(defaults.ptx_feature.iter()) - .chain(defaults.status.iter()) - .chain(defaults.llvm_arguments.iter().flatten()) - .chain(defaults.llvm_results.iter().flatten()) - .chain(defaults.concrete_llvm_arguments.iter().flatten()) - .chain(defaults.concrete_llvm_results.iter().flatten()) - { - reject_disallowed_placeholder(value, context)?; - } - reject_stage_placeholders(&defaults.stages, context) -} - -fn reject_stage_placeholders(stages: &[EvidenceStage], context: &str) -> Result<()> { - for stage in stages { - for value in stage - .targets - .iter() - .chain(std::iter::once(&stage.representation)) - .chain(std::iter::once(&stage.detail)) - .chain(stage.tool_path.iter()) - .chain(stage.tool_version.iter()) - .chain(stage.tool_sha256.iter()) - { - reject_disallowed_placeholder(value, context)?; - } - } - Ok(()) -} - -fn reject_disallowed_placeholder(value: &str, field: &str) -> Result<()> { - ensure!( - !value.contains("${"), - "{field} cannot contain matrix placeholders" - ); - Ok(()) -} - -fn validate_stage_pairs(stages: &[EvidenceStage], id: &str) -> Result<()> { - let mut identities = Vec::new(); - for stage in stages { - let identity = (stage.stage, stage.mechanism, stage.targets.clone()); - ensure!( - !identities.contains(&identity), - "expanded evidence {id} has conflicting duplicate {:?}/{:?} stage targets {:?}", - stage.stage, - stage.mechanism, - stage.targets - ); - identities.push(identity); - } - Ok(()) -} - -fn read_evidence(repo_root: &Path) -> Result<(Vec, Vec)> { - let directory = repo_root.join("intrinsics/evidence"); - let mut paths: Vec = fs::read_dir(&directory) - .with_context(|| format!("read {}", directory.display()))? - .filter_map(|entry| entry.ok().map(|entry| entry.path())) - .filter(|path| path.extension().and_then(|ext| ext.to_str()) == Some("json")) - .collect(); - paths.sort(); - ensure!( - !paths.is_empty(), - "no evidence JSON files in {}", - directory.display() - ); - let mut files = Vec::with_capacity(paths.len()); - let mut hashes = Vec::with_capacity(paths.len()); - for path in paths { - let file = read_evidence_file(&path)?; - let name = path.file_name().unwrap().to_string_lossy(); - hashes.push(format!("{name}:{}", sha256_file(&path)?)); - files.push(file); - } - Ok((files, hashes)) -} - -#[derive(Debug, Clone, Copy)] -struct IndexedEvidence<'a> { - file: &'a EvidenceFile, - record: &'a EvidenceRecord, - backend_version: &'a str, - backend_sha256: &'a str, -} - -fn index_evidence<'a>( - files: &'a [EvidenceFile], - llvm_revision: &str, -) -> Result>> { - let mut result = BTreeMap::new(); - for file in files { - ensure!( - !file.backend_profile.trim().is_empty() && !file.llvm_revision.trim().is_empty(), - "evidence file has no concrete backend profile or LLVM revision" - ); - ensure!( - !file.backend_version.trim().is_empty() && !file.backend_sha256.trim().is_empty(), - "evidence does not identify the backend binary" - ); - if file.backend_kind != Some(IntrinsicBackend::LibNvvm) { - ensure!( - file.llvm_revision == llvm_revision, - "selected evidence LLVM revision {} does not match pinned {}", - file.llvm_revision, - llvm_revision - ); - } - for record in &file.records { - ensure!( - result - .insert( - (file.backend_profile.as_str(), record.id.as_str()), - IndexedEvidence { - file, - record, - backend_version: &file.backend_version, - backend_sha256: &file.backend_sha256, - }, - ) - .is_none(), - "duplicate evidence for catalog ID {}", - record.id - ); - } - } - Ok(result) -} - -fn validate_evidence( - policy: &OverlayIntrinsic, - evidence: &IndexedEvidence<'_>, - lowering: Option<&crate::model::OverlayBackendLowering>, -) -> Result<()> { - let record = evidence.record; - record.expected_ptx.validate().map_err(|reason| { - anyhow::anyhow!( - "{} evidence has an invalid expected PTX pattern: {reason}", - policy.id - ) - })?; - let policy_source = resolve_policy_source(policy)?; - let evidence_source = match (&record.source, &record.source_record) { - (None, Some(source_record)) => IntrinsicSource::LlvmImported { - source_record: source_record.clone(), - }, - (Some(source), None) => source.clone(), - (Some(_), Some(_)) => bail!( - "{} evidence mixes tagged source with legacy source_record", - policy.id - ), - (None, None) => bail!("{} evidence has no source provenance", policy.id), - }; - ensure!( - evidence_source == policy_source, - "{} evidence source provenance mismatch", - policy.id - ); - ensure!( - record.llvm_symbol == policy.llvm_symbol - && record.llvm_arguments == policy.llvm_arguments - && record.llvm_results == policy.llvm_results, - "{} evidence signature mismatch", - policy.id - ); - if matches!(policy_source, IntrinsicSource::PtxNative { .. }) { - ensure!( - record.llvm_symbol.is_none() - && record.resolved_llvm_symbol.is_none() - && record.llvm_arguments.is_empty() - && record.llvm_results.is_empty() - && record.concrete_llvm_arguments.is_empty() - && record.concrete_llvm_results.is_empty() - && record.declaration_attributes_canonicalized.is_none(), - "{} PTX-native evidence must not invent LLVM declaration facts", - policy.id - ); - } - ensure!( - record.expected_ptx == policy.expected_ptx, - "{} evidence PTX expectation mismatch", - policy.id - ); - ensure!( - matches!(record.status.as_str(), "lowered" | "validated" | "executed"), - "{} evidence status {} is too weak to generate a lowering", - policy.id, - record.status - ); - ensure!( - !record.target_triple.is_empty() - && !record.gpu_target.is_empty() - && !record.ptx_feature.is_empty(), - "{} evidence omits its full target profile", - policy.id - ); - if let Some(lowering) = lowering { - ensure!( - evidence.file.backend_kind == Some(lowering.backend), - "{} evidence profile {} has the wrong backend kind", - policy.id, - evidence.file.backend_profile - ); - match record.status.as_str() { - "executed" => ensure!( - record.runtime_validation == Some(RuntimeValidation::Executed), - "{} executed evidence must record runtime_validation = executed", - policy.id - ), - _ => ensure!( - record.runtime_validation == Some(RuntimeValidation::Unexecuted), - "{} non-executed backend evidence must record runtime_validation = unexecuted", - policy.id - ), - } - ensure!( - !record.stages.is_empty(), - "{} backend evidence omits compilation stages", - policy.id - ); - ensure!( - record.stages.iter().any(|stage| { - stage.stage == EvidenceStageKind::BackendCodegen - && stage.mechanism == Some(lowering.mechanism) - && stage.outcome == "succeeded" - }), - "{} evidence has no successful backend-codegen stage for {:?}", - policy.id, - lowering.mechanism - ); - validate_selected_stage_targets(policy, record, lowering)?; - if lowering.mechanism == BackendLoweringMechanism::TypedNvvm { - validate_typed_llvm_evidence(policy, record)?; - } - validate_packed_conversion_backend_evidence(policy, record, lowering)?; - validate_scalar_conversion_backend_evidence(policy, record, lowering)?; - validate_scalar_arithmetic_backend_evidence(policy, record, lowering)?; - validate_inline_ptx_fallback_evidence(policy, record, lowering)?; - if lowering.backend == IntrinsicBackend::LlvmNvptx - && matches!(record.status.as_str(), "validated" | "executed") - { - ensure!( - has_valid_ptx_assembly_stage(record, lowering.mechanism), - "{} validated LLVM-NVPTX evidence requires a successful PTX-assembly stage with exact tool identity", - policy.id - ); - } else if lowering.backend == IntrinsicBackend::LibNvvm - && matches!(record.status.as_str(), "validated" | "executed") - { - ensure!( - has_valid_cubin_device_link_stage(record, lowering.mechanism), - "{} validated libNVVM evidence requires a successful cubin-producing device-link stage with exact tool identity", - policy.id - ); - } - if record.status == "executed" { - ensure!( - record.stages.iter().any(|stage| { - stage.stage == EvidenceStageKind::Runtime - && stage.mechanism == Some(lowering.mechanism) - && stage.outcome == "succeeded" - }), - "{} executed evidence requires a successful runtime stage for the selected mechanism", - policy.id - ); - } - } - Ok(()) -} - -fn validate_inline_ptx_fallback_evidence( - policy: &OverlayIntrinsic, - record: &EvidenceRecord, - lowering: &crate::model::OverlayBackendLowering, -) -> Result<()> { - if !matches!( - policy.family.as_str(), - "cluster_barrier" | "wgmma_control" | "elect" - ) || lowering.backend != IntrinsicBackend::LibNvvm - { - return Ok(()); - } - for stage in [ - EvidenceStageKind::BackendCodegen, - EvidenceStageKind::DeviceLink, - ] { - ensure!( - record.stages.iter().any(|candidate| { - candidate.stage == stage - && candidate.mechanism == Some(BackendLoweringMechanism::TypedNvvm) - && candidate.outcome == "failed" - }), - "{} libNVVM inline-PTX evidence must record the failed typed-NVVM {:?} comparison", - policy.id, - stage - ); - } - ensure!( - !record.stages.iter().any(|candidate| { - candidate.stage == EvidenceStageKind::DeviceLink - && candidate.mechanism == Some(BackendLoweringMechanism::TypedNvvm) - && candidate.outcome == "succeeded" - }), - "{} libNVVM evidence cannot select inline PTX after a successful typed-NVVM terminal", - policy.id - ); - Ok(()) -} - -fn validate_packed_conversion_backend_evidence( - policy: &OverlayIntrinsic, - record: &EvidenceRecord, - lowering: &crate::model::OverlayBackendLowering, -) -> Result<()> { - if policy.family != "packed_conversion" { - return Ok(()); - } - match lowering.backend { - IntrinsicBackend::LlvmNvptx => { - validate_typed_llvm_evidence(policy, record)?; - for stage in [ - EvidenceStageKind::DeclarationCanonicalization, - EvidenceStageKind::BackendCodegen, - ] { - ensure!( - successful_stage(record, BackendLoweringMechanism::TypedNvvm, stage).is_some(), - "{} LLVM packed-conversion evidence requires a successful auxiliary typed-NVVM {:?} stage", - policy.id, - stage - ); - } - ensure!( - has_valid_ptx_assembly_stage(record, BackendLoweringMechanism::TypedNvvm), - "{} LLVM packed-conversion evidence requires a successful auxiliary typed-NVVM PTX-assembly stage with exact tool identity", - policy.id - ); - ensure!( - matches!(record.status.as_str(), "validated" | "executed"), - "{} LLVM packed-conversion evidence requires validated evidence status for its auxiliary typed-NVVM terminal stage", - policy.id - ); - let typed_lowering = crate::model::OverlayBackendLowering { - mechanism: BackendLoweringMechanism::TypedNvvm, - ..lowering.clone() - }; - validate_selected_stage_targets(policy, record, &typed_lowering)?; - Ok(()) - } - IntrinsicBackend::LibNvvm => { - ensure!( - record.resolved_llvm_symbol.is_none() - && record.concrete_llvm_arguments.is_empty() - && record.concrete_llvm_results.is_empty() - && record.declaration_attributes_canonicalized.is_none() - && !record.stages.iter().any(|stage| { - stage.mechanism == Some(BackendLoweringMechanism::TypedNvvm) - }), - "{} libNVVM inline-PTX evidence must not claim typed LLVM support", - policy.id - ); - Ok(()) - } - } -} - -fn validate_scalar_conversion_backend_evidence( - policy: &OverlayIntrinsic, - record: &EvidenceRecord, - lowering: &crate::model::OverlayBackendLowering, -) -> Result<()> { - if policy.family != "scalar_conversion" { - return Ok(()); - } - match lowering.backend { - IntrinsicBackend::LlvmNvptx => ensure!( - successful_stage( - record, - BackendLoweringMechanism::TypedNvvm, - EvidenceStageKind::DeclarationCanonicalization, - ) - .is_some(), - "{} LLVM scalar-conversion evidence must canonicalize the typed declaration", - policy.id - ), - IntrinsicBackend::LibNvvm => ensure!( - record.resolved_llvm_symbol.is_none() - && record.concrete_llvm_arguments.is_empty() - && record.concrete_llvm_results.is_empty() - && record.declaration_attributes_canonicalized.is_none(), - "{} libNVVM scalar-conversion evidence must describe the selected inline-PTX route", - policy.id - ), - }; - Ok(()) -} - -fn validate_scalar_arithmetic_backend_evidence( - policy: &OverlayIntrinsic, - record: &EvidenceRecord, - lowering: &crate::model::OverlayBackendLowering, -) -> Result<()> { - if policy.family != "scalar_arithmetic" { - return Ok(()); - } - match (lowering.backend, lowering.mechanism) { - (IntrinsicBackend::LlvmNvptx, BackendLoweringMechanism::TypedNvvm) => ensure!( - successful_stage( - record, - BackendLoweringMechanism::TypedNvvm, - EvidenceStageKind::DeclarationCanonicalization, - ) - .is_some(), - "{} typed LLVM scalar-arithmetic evidence must canonicalize the declaration", - policy.id - ), - (IntrinsicBackend::LlvmNvptx, BackendLoweringMechanism::InlinePtx) => { - validate_typed_llvm_evidence(policy, record)?; - ensure!( - successful_stage( - record, - BackendLoweringMechanism::TypedNvvm, - EvidenceStageKind::DeclarationCanonicalization, - ) - .is_some(), - "{} inline LLVM scalar-arithmetic evidence must canonicalize its imported declaration", - policy.id - ); - } - (IntrinsicBackend::LibNvvm, _) => ensure!( - record.resolved_llvm_symbol.is_none() - && record.concrete_llvm_arguments.is_empty() - && record.concrete_llvm_results.is_empty() - && record.declaration_attributes_canonicalized.is_none(), - "{} libNVVM scalar-arithmetic evidence must describe the selected inline-PTX route", - policy.id - ), - }; - Ok(()) -} - -fn validate_typed_llvm_evidence(policy: &OverlayIntrinsic, record: &EvidenceRecord) -> Result<()> { - let concrete_arguments = policy - .llvm_arguments - .iter() - .map(|argument| { - match argument.as_str() { - "shared_cluster_ptr" => return Ok("ptr addrspace(7)".into()), - "shared_ptr" => return Ok("ptr addrspace(3)".into()), - "global_ptr" => return Ok("ptr addrspace(1)".into()), - "ptr" => return Ok("ptr".into()), - "anyptr" => {} - _ => return Ok(argument.clone()), - } - match policy.selected_address_space.with_context(|| { - format!( - "{} has a polymorphic LLVM pointer without a selected address space", - policy.id - ) - })? { - ImportedAddressSpace::Generic => Ok("ptr".into()), - ImportedAddressSpace::Shared => Ok("ptr addrspace(3)".into()), - } - }) - .collect::>>()?; - ensure!( - record.resolved_llvm_symbol == policy.resolved_llvm_symbol - && record.concrete_llvm_arguments == concrete_arguments - && record.concrete_llvm_results == policy.llvm_results - && record.declaration_attributes_canonicalized == Some(true), - "{} typed LLVM evidence does not prove its resolved signature and canonical declaration attributes", - policy.id - ); - Ok(()) -} - -fn validate_selected_stage_targets( - policy: &OverlayIntrinsic, - record: &EvidenceRecord, - lowering: &crate::model::OverlayBackendLowering, -) -> Result<()> { - let requirement = backend_target_requirement(policy, lowering)?; - let blackwell_ldmatrix = policy.family == "ldmatrix" - && policy - .ldmatrix_variant - .as_ref() - .is_some_and(|variant| variant.shape != LdmatrixShape::M8n8); - let f8f6f4_mma = is_f8f6f4_mma_target_matrix_policy(policy); - let paired_alternatives = match &requirement.hardware { - CatalogHardwareTarget::TargetMatrix { contracts } => Some( - contracts - .iter() - .flat_map(|contract| contract.alternatives.iter().cloned()) - .collect::>(), - ), - _ => None, - }; - let target_matrix_evidence = blackwell_ldmatrix || f8f6f4_mma || paired_alternatives.is_some(); - if !target_matrix_evidence { - let mut pairs = Vec::new(); - for stage in record.stages.iter().filter(|stage| { - stage.mechanism == Some(lowering.mechanism) && stage.outcome == "succeeded" - }) { - let pair = (stage.stage, stage.mechanism); - ensure!( - !pairs.contains(&pair), - "{} has multiple target-specific {:?}/{:?} stages outside reviewed target-matrix evidence", - policy.id, - stage.stage, - stage.mechanism - ); - pairs.push(pair); - } - } - for stage in &record.stages { - ensure!( - !stage.targets.is_empty(), - "{} evidence stage {:?} has no targets", - policy.id, - stage.stage - ); - for target in &stage.targets { - ensure!( - is_normalized_stage_target(target), - "{} evidence stage has unsupported target spelling {target:?}", - policy.id - ); - } - } - - let terminal_kind = match lowering.backend { - IntrinsicBackend::LlvmNvptx => EvidenceStageKind::PtxAssembly, - IntrinsicBackend::LibNvvm => EvidenceStageKind::DeviceLink, - }; - let expected_ptx = requirement.minimum_ptx.encoded(); - let expected_hardware = match &requirement.hardware { - CatalogHardwareTarget::AnyOf { alternatives } if !alternatives.is_empty() => { - alternatives.clone() - } - CatalogHardwareTarget::TargetMatrix { contracts } if !contracts.is_empty() => { - let mut alternatives = contracts - .iter() - .flat_map(|contract| contract.alternatives.iter()) - .map(|alternative| alternative.hardware) - .collect::>(); - alternatives.sort_by_key(|hardware| target_hardware_sort_key(*hardware)); - alternatives.dedup(); - alternatives - } - _ => bail!( - "{} selected backend stages require a hardware target", - policy.id - ), - }; - if target_matrix_evidence { - return validate_target_matrix_stage_targets( - policy, - record, - lowering, - terminal_kind, - &expected_hardware, - expected_ptx, - paired_alternatives.as_deref(), - ); - } - let backend = successful_stage( - record, - lowering.mechanism, - EvidenceStageKind::BackendCodegen, - ) - .with_context(|| { - format!( - "{} has no successful selected backend-codegen stage", - policy.id - ) - })?; - let mut required_stages = vec![backend]; - if matches!(record.status.as_str(), "validated" | "executed") { - required_stages.push( - successful_stage(record, lowering.mechanism, terminal_kind).with_context(|| { - format!("{} has no successful selected terminal stage", policy.id) - })?, - ); - } - for stage in required_stages { - let (hardware, ptx) = selected_stage_floor(stage)?; - let allow_forward_minimum = stage.stage != EvidenceStageKind::BackendCodegen; - let hardware_matches = expected_hardware.iter().any(|expected| { - selected_stage_hardware_matches(hardware, *expected, allow_forward_minimum) - }); - let ptx_matches = match lowering.backend { - IntrinsicBackend::LlvmNvptx => ptx == expected_ptx, - IntrinsicBackend::LibNvvm => ptx >= expected_ptx, - }; - ensure!( - hardware_matches && ptx_matches, - "{} evidence stage {:?} targets {} / PTX {}.{} instead of a compatible target at catalog floor {} / PTX {}.{}", - policy.id, - stage.stage, - describe_stage_hardware(hardware), - ptx / 10, - ptx % 10, - expected_hardware - .iter() - .map(|hardware| describe_stage_hardware(*hardware)) - .collect::>() - .join(" or "), - expected_ptx / 10, - expected_ptx % 10 - ); - } - if record.status == "executed" { - let runtime = successful_stage(record, lowering.mechanism, EvidenceStageKind::Runtime) - .with_context(|| { - format!( - "{} executed evidence has no successful runtime stage", - policy.id - ) - })?; - let (hardware, ptx) = selected_stage_floor(runtime)?; - let ptx_matches = match lowering.backend { - IntrinsicBackend::LlvmNvptx => ptx == expected_ptx, - IntrinsicBackend::LibNvvm => ptx >= expected_ptx, - }; - ensure!( - expected_hardware - .iter() - .any(|expected| selected_stage_hardware_matches(hardware, *expected, true)) - && ptx_matches, - "{} runtime stage target does not satisfy its catalog floor", - policy.id - ); - } - Ok(()) -} - -fn validate_target_matrix_stage_targets( - policy: &OverlayIntrinsic, - record: &EvidenceRecord, - lowering: &OverlayBackendLowering, - terminal_kind: EvidenceStageKind, - expected_hardware: &[CatalogHardwareAlternative], - minimum_ptx: u16, - paired_alternatives: Option<&[CatalogTargetAlternative]>, -) -> Result<()> { - let mut required_kinds = vec![EvidenceStageKind::BackendCodegen]; - if matches!(record.status.as_str(), "validated" | "executed") { - if lowering.backend == IntrinsicBackend::LibNvvm { - required_kinds.push(EvidenceStageKind::PtxAssembly); - } - required_kinds.push(terminal_kind); - } - - for kind in required_kinds { - let stages = record - .stages - .iter() - .filter(|stage| { - stage.stage == kind - && stage.mechanism == Some(lowering.mechanism) - && stage.outcome == "succeeded" - }) - .collect::>(); - ensure!( - stages.len() == expected_hardware.len(), - "{} target-matrix {:?} evidence must contain one structured stage for each of its {} reviewed targets", - policy.id, - kind, - expected_hardware.len() - ); - - let mut covered = Vec::new(); - for stage in stages { - let (hardware, ptx) = selected_stage_floor(stage)?; - ensure!( - expected_hardware.contains(&hardware) && !covered.contains(&hardware), - "{} target-matrix {:?} evidence has an unexpected or duplicate target {}", - policy.id, - kind, - describe_stage_hardware(hardware) - ); - covered.push(hardware); - let exact_floor = - target_matrix_ptx_floor(policy, hardware, paired_alternatives, false)?; - let ptx_matches = match lowering.backend { - IntrinsicBackend::LlvmNvptx => ptx == exact_floor, - IntrinsicBackend::LibNvvm => { - ptx >= paired_alternatives.map_or(minimum_ptx, |_| exact_floor) - } - }; - ensure!( - ptx_matches, - "{} target-matrix {:?} evidence records the wrong PTX floor for {}", - policy.id, - kind, - describe_stage_hardware(hardware) - ); - if matches!( - kind, - EvidenceStageKind::PtxAssembly | EvidenceStageKind::DeviceLink - ) { - ensure!( - stage - .tool_path - .as_deref() - .is_some_and(|value| !value.is_empty()) - && stage - .tool_version - .as_deref() - .is_some_and(|value| !value.is_empty()) - && stage - .tool_sha256 - .as_deref() - .is_some_and(|value| !value.is_empty()), - "{} target-matrix {:?} evidence for {} does not identify its tool", - policy.id, - kind, - describe_stage_hardware(hardware) - ); - } - if kind == EvidenceStageKind::DeviceLink { - ensure!( - stage.artifact_kind == Some(EvidenceArtifactKind::Cubin), - "{} target-matrix device-link evidence for {} is not a cubin", - policy.id, - describe_stage_hardware(hardware) - ); - } - } - } - - if record.status == "executed" { - let runtime = record - .stages - .iter() - .filter(|stage| { - stage.stage == EvidenceStageKind::Runtime - && stage.mechanism == Some(lowering.mechanism) - && stage.outcome == "succeeded" - }) - .collect::>(); - ensure!( - runtime.len() == 1, - "{} executed target-matrix evidence requires one successful runtime stage", - policy.id - ); - let (hardware, ptx) = selected_stage_floor(runtime[0])?; - ensure!( - expected_hardware - .iter() - .any(|expected| { selected_stage_hardware_matches(hardware, *expected, true) }), - "{} runtime target {} is outside its target matrix", - policy.id, - describe_stage_hardware(hardware) - ); - let floor = target_matrix_ptx_floor(policy, hardware, paired_alternatives, true)?; - let ptx_matches = match lowering.backend { - IntrinsicBackend::LlvmNvptx => ptx == floor, - IntrinsicBackend::LibNvvm => ptx >= floor, - }; - ensure!( - ptx_matches, - "{} runtime target {} records PTX {}.{} instead of its paired floor {}.{}", - policy.id, - describe_stage_hardware(hardware), - ptx / 10, - ptx % 10, - floor / 10, - floor % 10 - ); - } - Ok(()) -} - -fn target_matrix_ptx_floor( - policy: &OverlayIntrinsic, - hardware: CatalogHardwareAlternative, - paired_alternatives: Option<&[CatalogTargetAlternative]>, - allow_forward_minimum: bool, -) -> Result { - if let Some(alternatives) = paired_alternatives { - return alternatives - .iter() - .filter(|alternative| { - selected_stage_hardware_matches( - hardware, - alternative.hardware, - allow_forward_minimum, - ) - }) - .map(|alternative| alternative.minimum_ptx.encoded()) - .min() - .with_context(|| { - format!( - "{} target matrix has no PTX floor for {}", - policy.id, - describe_stage_hardware(hardware) - ) - }); - } - if is_f8f6f4_mma_target_matrix_policy(policy) { - f8f6f4_llvm_ptx_floor(hardware) - } else { - blackwell_ldmatrix_llvm_ptx_floor(hardware) - } -} - -fn f8f6f4_llvm_ptx_floor(hardware: CatalogHardwareAlternative) -> Result { - match hardware { - CatalogHardwareAlternative::ExactArchitecture { sm: 120 } - | CatalogHardwareAlternative::FamilyTarget { sm: 120 } - | CatalogHardwareAlternative::ExactArchitecture { sm: 121 } - | CatalogHardwareAlternative::FamilyTarget { sm: 121 } => { - recorded_stage_ptx_floor(hardware, 87) - } - _ => bail!( - "{} is not a reviewed f8f6f4 MMA target", - describe_stage_hardware(hardware) - ), - } -} - -fn blackwell_ldmatrix_llvm_ptx_floor(hardware: CatalogHardwareAlternative) -> Result { - match hardware { - CatalogHardwareAlternative::ExactArchitecture { - sm: 100 | 103 | 110 | 120 | 121, - } - | CatalogHardwareAlternative::FamilyTarget { - sm: 100 | 103 | 120 | 121, - } - | CatalogHardwareAlternative::FamilyTarget { sm: 110 } => { - recorded_stage_ptx_floor(hardware, 86) - } - _ => bail!( - "{} is not a reviewed Blackwell ldmatrix target", - describe_stage_hardware(hardware) - ), - } -} - -fn recorded_stage_ptx_floor( - hardware: CatalogHardwareAlternative, - instruction_floor: u16, -) -> Result { - let (sm, suffix) = match hardware { - CatalogHardwareAlternative::MinimumSm { sm } => (sm, None), - CatalogHardwareAlternative::ExactArchitecture { sm } => (sm, Some('a')), - CatalogHardwareAlternative::FamilyTarget { sm } => (sm, Some('f')), - }; - let arch = CudaArch::new(u32::from(sm), suffix)?; - Ok(instruction_floor.max(recorded_ptx_floor(&arch)?)) -} - -fn successful_stage( - record: &EvidenceRecord, - mechanism: BackendLoweringMechanism, - kind: EvidenceStageKind, -) -> Option<&crate::model::EvidenceStage> { - record.stages.iter().find(|stage| { - stage.stage == kind && stage.mechanism == Some(mechanism) && stage.outcome == "succeeded" - }) -} - -fn is_normalized_stage_target(target: &str) -> bool { - if let Some(value) = target.strip_prefix("ptx") { - return value.len() == 2 && value.bytes().all(|byte| byte.is_ascii_digit()); - } - parse_stage_hardware(target).is_some() -} - -fn parse_stage_hardware(target: &str) -> Option { - let arch = target.parse::().ok()?; - let canonical = if target.starts_with("sm_") { - arch.sm() - } else if target.starts_with("compute_") { - arch.compute() - } else { - return None; - }; - if canonical != target || !matches!(arch.capability().to_string().len(), 2 | 3) { - return None; - } - let sm = u16::try_from(arch.capability()).ok()?; - Some(match arch.suffix() { - None => CatalogHardwareAlternative::MinimumSm { sm }, - Some('a') => CatalogHardwareAlternative::ExactArchitecture { sm }, - Some('f') => CatalogHardwareAlternative::FamilyTarget { sm }, - _ => unreachable!(), - }) -} - -fn selected_stage_hardware_matches( - actual: CatalogHardwareAlternative, - expected: CatalogHardwareAlternative, - allow_forward_minimum: bool, -) -> bool { - match expected { - CatalogHardwareAlternative::MinimumSm { sm: expected } => { - if allow_forward_minimum { - match actual { - CatalogHardwareAlternative::MinimumSm { sm } - | CatalogHardwareAlternative::ExactArchitecture { sm } - | CatalogHardwareAlternative::FamilyTarget { sm } => sm >= expected, - } - } else { - actual == CatalogHardwareAlternative::MinimumSm { sm: expected } - } - } - CatalogHardwareAlternative::ExactArchitecture { .. } - | CatalogHardwareAlternative::FamilyTarget { .. } => actual == expected, - } -} - -fn describe_stage_hardware(hardware: CatalogHardwareAlternative) -> String { - match hardware { - CatalogHardwareAlternative::MinimumSm { sm } => format!("sm_{sm}"), - CatalogHardwareAlternative::ExactArchitecture { sm } => format!("sm_{sm}a"), - CatalogHardwareAlternative::FamilyTarget { sm } => format!("sm_{sm}f"), - } -} - -fn selected_stage_floor( - stage: &crate::model::EvidenceStage, -) -> Result<(CatalogHardwareAlternative, u16)> { - let mut hardware = None; - let mut ptx = None; - for target in &stage.targets { - if let Some(value) = target.strip_prefix("ptx") { - let value = value.parse::()?; - ensure!( - ptx.replace(value).is_none(), - "stage has duplicate PTX targets" - ); - } else { - let value = parse_stage_hardware(target) - .with_context(|| format!("stage has unsupported target spelling {target:?}"))?; - ensure!( - hardware.replace(value).is_none(), - "stage has duplicate architecture targets" - ); - } - } - Ok(( - hardware.context("selected stage has no architecture target")?, - ptx.context("selected stage has no PTX target")?, - )) -} - -fn has_valid_ptx_assembly_stage( - record: &EvidenceRecord, - mechanism: BackendLoweringMechanism, -) -> bool { - has_valid_tool_stage(record, mechanism, EvidenceStageKind::PtxAssembly) -} - -fn has_valid_cubin_device_link_stage( - record: &EvidenceRecord, - mechanism: BackendLoweringMechanism, -) -> bool { - has_valid_tool_stage(record, mechanism, EvidenceStageKind::DeviceLink) - && successful_stage(record, mechanism, EvidenceStageKind::DeviceLink) - .is_some_and(|stage| stage.artifact_kind == Some(EvidenceArtifactKind::Cubin)) -} - -fn has_valid_tool_stage( - record: &EvidenceRecord, - mechanism: BackendLoweringMechanism, - stage_kind: EvidenceStageKind, -) -> bool { - record.stages.iter().any(|stage| { - stage.stage == stage_kind - && stage.mechanism == Some(mechanism) - && stage.outcome == "succeeded" - && stage - .tool_path - .as_deref() - .is_some_and(|value| !value.is_empty()) - && stage - .tool_version - .as_deref() - .is_some_and(|value| !value.is_empty()) - && stage.tool_sha256.as_deref().is_some_and(|value| { - value.len() == 64 && value.bytes().all(|byte| byte.is_ascii_hexdigit()) - }) - }) -} - -fn resolve_backend_lowerings( - policy: &OverlayIntrinsic, - evidence_by_profile_id: &BTreeMap<(&str, &str), IndexedEvidence<'_>>, -) -> Result> { - let mut resolved = Vec::with_capacity(policy.backend_lowerings.len()); - let mut runtime_states = Vec::with_capacity(policy.backend_lowerings.len()); - for lowering in &policy.backend_lowerings { - let evidence = evidence_by_profile_id - .get(&(lowering.evidence_profile.as_str(), policy.id.as_str())) - .with_context(|| { - format!( - "{} has no evidence in backend profile {}", - policy.id, lowering.evidence_profile - ) - })?; - validate_evidence(policy, evidence, Some(lowering))?; - runtime_states.push(evidence.record.runtime_validation); - resolved.push(CatalogBackendLowering { - backend: lowering.backend, - mechanism: lowering.mechanism, - evidence_profile: lowering.evidence_profile.clone(), - target: backend_target_requirement(policy, lowering)?, - version: evidence.backend_version.to_owned(), - sha256: evidence.backend_sha256.to_owned(), - artifact_path: evidence.file.artifact_path.clone(), - build_id_prefix: evidence.file.build_id_prefix.clone(), - status: evidence.record.status.clone(), - stages: evidence.record.stages.clone(), - }); - } - if let Some(safety) = &policy.ldmatrix_safety { - match safety.runtime_validation { - RuntimeValidation::Unexecuted => ensure!( - runtime_states - .iter() - .all(|state| *state == Some(RuntimeValidation::Unexecuted)), - "{} overlay says runtime is unexecuted but backend evidence disagrees", - policy.id - ), - RuntimeValidation::Executed => ensure!( - runtime_states.contains(&Some(RuntimeValidation::Executed)), - "{} overlay says runtime is executed but no backend evidence has an executed runtime stage", - policy.id - ), - } - } - if policy.family == "stmatrix" { - ensure!( - runtime_states - .iter() - .all(|state| *state == Some(RuntimeValidation::Unexecuted)), - "{} stmatrix source contract is unexecuted but backend evidence disagrees", - policy.id - ); - } - if let Some(packed) = &policy.packed_atomic { - match packed.runtime_validation { - RuntimeValidation::Unexecuted => ensure!( - runtime_states - .iter() - .all(|state| *state == Some(RuntimeValidation::Unexecuted)), - "{} packed-atomic runtime is unexecuted but backend evidence disagrees", - policy.id - ), - RuntimeValidation::Executed => ensure!( - runtime_states.contains(&Some(RuntimeValidation::Executed)), - "{} packed-atomic runtime is executed but no backend evidence records execution", - policy.id - ), - } - } - if let Some(mbarrier) = &policy.mbarrier_basic { - match mbarrier.runtime_validation { - RuntimeValidation::Unexecuted => ensure!( - runtime_states - .iter() - .all(|state| *state == Some(RuntimeValidation::Unexecuted)), - "{} mbarrier runtime is unexecuted but backend evidence disagrees", - policy.id - ), - RuntimeValidation::Executed => ensure!( - runtime_states.contains(&Some(RuntimeValidation::Executed)), - "{} mbarrier runtime is executed but no backend evidence records execution", - policy.id - ), - } - } - if let Some(movmatrix) = &policy.movmatrix { - match movmatrix.runtime_validation { - RuntimeValidation::Unexecuted => ensure!( - runtime_states - .iter() - .all(|state| *state == Some(RuntimeValidation::Unexecuted)), - "{} movmatrix runtime is unexecuted but backend evidence disagrees", - policy.id - ), - RuntimeValidation::Executed => ensure!( - runtime_states.contains(&Some(RuntimeValidation::Executed)), - "{} movmatrix runtime is executed but no backend evidence records execution", - policy.id - ), - } - } - if let Some(mbarrier) = &policy.mbarrier_extended { - match mbarrier.runtime_validation { - RuntimeValidation::Unexecuted => ensure!( - runtime_states - .iter() - .all(|state| *state == Some(RuntimeValidation::Unexecuted)), - "{} extended-mbarrier runtime is unexecuted but backend evidence disagrees", - policy.id - ), - RuntimeValidation::Executed => ensure!( - runtime_states.contains(&Some(RuntimeValidation::Executed)), - "{} extended-mbarrier runtime is executed but no backend evidence records execution", - policy.id - ), - } - } - if let Some(bridge) = &policy.cp_async_mbarrier { - match bridge.runtime_validation { - RuntimeValidation::Unexecuted => ensure!( - runtime_states - .iter() - .all(|state| *state == Some(RuntimeValidation::Unexecuted)), - "{} cp.async mbarrier runtime is unexecuted but backend evidence disagrees", - policy.id - ), - RuntimeValidation::Executed => ensure!( - runtime_states.contains(&Some(RuntimeValidation::Executed)), - "{} cp.async mbarrier runtime is executed but no backend evidence records execution", - policy.id - ), - } - } - if let Some(mma) = &policy.register_mma { - match mma.runtime_validation { - RuntimeValidation::Unexecuted => ensure!( - runtime_states - .iter() - .all(|state| *state == Some(RuntimeValidation::Unexecuted)), - "{} register-MMA runtime is unexecuted but backend evidence disagrees", - policy.id - ), - RuntimeValidation::Executed => ensure!( - runtime_states.contains(&Some(RuntimeValidation::Executed)), - "{} register-MMA runtime is executed but no backend evidence records execution", - policy.id - ), - } - } - if let Some(mma) = &policy.sparse_mma { - match mma.runtime_validation { - RuntimeValidation::Unexecuted => ensure!( - runtime_states - .iter() - .all(|state| *state == Some(RuntimeValidation::Unexecuted)), - "{} sparse-MMA runtime is unexecuted but backend evidence disagrees", - policy.id - ), - RuntimeValidation::Executed => ensure!( - runtime_states.contains(&Some(RuntimeValidation::Executed)), - "{} sparse-MMA runtime is executed but no backend evidence records execution", - policy.id - ), - } - } - if let Some(debug) = &policy.debug_control { - match debug.runtime_validation { - RuntimeValidation::Unexecuted => ensure!( - runtime_states - .iter() - .all(|state| *state == Some(RuntimeValidation::Unexecuted)), - "{} debug-control runtime is unexecuted but backend evidence disagrees", - policy.id - ), - RuntimeValidation::Executed => ensure!( - runtime_states.contains(&Some(RuntimeValidation::Executed)), - "{} debug-control runtime is executed but no backend evidence records execution", - policy.id - ), - } - } - if let Some(cluster_memory) = &policy.cluster_memory { - match cluster_memory.runtime_validation { - RuntimeValidation::Unexecuted => ensure!( - runtime_states - .iter() - .all(|state| *state == Some(RuntimeValidation::Unexecuted)), - "{} cluster-memory runtime is unexecuted but backend evidence disagrees", - policy.id - ), - RuntimeValidation::Executed => ensure!( - runtime_states.contains(&Some(RuntimeValidation::Executed)), - "{} cluster-memory runtime is executed but no backend evidence records execution", - policy.id - ), - } - } - if let Some(clc) = &policy.clc { - match clc.runtime_validation { - RuntimeValidation::Unexecuted => ensure!( - runtime_states - .iter() - .all(|state| *state == Some(RuntimeValidation::Unexecuted)), - "{} CLC runtime is unexecuted but backend evidence disagrees", - policy.id - ), - RuntimeValidation::Executed => ensure!( - runtime_states.contains(&Some(RuntimeValidation::Executed)), - "{} CLC runtime is executed but no backend evidence records execution", - policy.id - ), - } - } - if let Some(conversion) = &policy.scalar_conversion { - match conversion.runtime_validation { - RuntimeValidation::Unexecuted => ensure!( - runtime_states - .iter() - .all(|state| *state == Some(RuntimeValidation::Unexecuted)), - "{} scalar-conversion runtime is unexecuted but backend evidence disagrees", - policy.id - ), - RuntimeValidation::Executed => ensure!( - runtime_states.contains(&Some(RuntimeValidation::Executed)), - "{} scalar-conversion runtime is executed but no backend evidence records execution", - policy.id - ), - } - } - if let Some(arithmetic) = &policy.scalar_arithmetic { - match arithmetic.runtime_validation { - RuntimeValidation::Unexecuted => ensure!( - runtime_states - .iter() - .all(|state| *state == Some(RuntimeValidation::Unexecuted)), - "{} scalar-arithmetic runtime is unexecuted but backend evidence disagrees", - policy.id - ), - RuntimeValidation::Executed => ensure!( - runtime_states.contains(&Some(RuntimeValidation::Executed)), - "{} scalar-arithmetic runtime is executed but no backend evidence records execution", - policy.id - ), - } - } - resolved.sort_by_key(|lowering| lowering.backend); - Ok(resolved) -} - -#[allow(clippy::too_many_arguments)] -fn resolve_record( - policy: &OverlayIntrinsic, - source: IntrinsicSource, - declaration: Option<&ImportedIntrinsic>, - evidence: &EvidenceRecord, - backend_profile: &str, - backend_version: &str, - backend_sha256: &str, - backend_lowerings: Vec, - intrinsic_abi: u32, -) -> Result { - materialize_record( - policy, - source, - declaration, - CatalogBackend { - profile: backend_profile.to_owned(), - version: backend_version.to_owned(), - sha256: backend_sha256.to_owned(), - status: evidence.status.clone(), - target_triple: evidence.target_triple.clone(), - gpu_target: evidence.gpu_target.clone(), - ptx_feature: evidence.ptx_feature.clone(), - }, - backend_lowerings, - intrinsic_abi, - ) -} - -fn materialize_record( - policy: &OverlayIntrinsic, - source: IntrinsicSource, - declaration: Option<&ImportedIntrinsic>, - backend: CatalogBackend, - backend_lowerings: Vec, - intrinsic_abi: u32, -) -> Result { - let native_target = if let Some(mma) = policy - .tcgen05 - .as_ref() - .and_then(|tcgen05| tcgen05.mma.as_ref()) - { - mma.llvm_target.clone() - } else { - CatalogTargetRequirement { - minimum_ptx: parse_ptx_version(&policy.minimum_ptx, &policy.id)?, - hardware: parse_hardware_target(policy)?, - } - }; - let llvm = if let Some(declaration) = declaration { - Some(CatalogLlvm { - symbol: policy - .llvm_symbol - .clone() - .expect("validated imported LLVM symbol"), - resolved_symbol: policy.resolved_llvm_symbol.clone(), - arguments: policy.llvm_arguments.clone(), - results: policy.llvm_results.clone(), - properties: declaration.properties.clone(), - result_facts: imported_result_facts(&declaration.properties)?, - }) - } else { - None - }; - let preserves_empty_dialect_signature = (policy.family == "sync" - && policy.id == "sync_threads") - || matches!( - ExecutionControlOperation::from_catalog_id(&policy.id), - Some( - ExecutionControlOperation::SetMaxNRegInc | ExecutionControlOperation::SetMaxNRegDec - ) - ); - let dialect_operands = - if policy.dialect_operands.is_empty() && !preserves_empty_dialect_signature { - policy.llvm_arguments.clone() - } else { - policy.dialect_operands.clone() - }; - let dialect_results = if policy.dialect_results.is_empty() && !preserves_empty_dialect_signature - { - policy.llvm_results.clone() - } else { - policy.dialect_results.clone() - }; - let mut selections = Vec::new(); - for selection in declaration - .into_iter() - .flat_map(|declaration| declaration.selections.iter()) - { - if selection_matches_policy(policy, selection)? { - selections.push(CatalogSelection { - source_record: selection.source_record.clone(), - asm: selection.asm.clone(), - predicates: selection.predicates.clone(), - constraints: selection.constraints.clone(), - }); - } - } - Ok(CatalogIntrinsic { - id: policy.id.clone(), - operation_key: policy.operation_key.clone(), - family: policy.family.clone(), - source, - selections, - rust: CatalogRust { - abi_id: policy.abi_id.clone(), - module: policy.rust_module.clone(), - name: policy.rust_name.clone(), - arguments: policy.rust_arguments.clone(), - result: policy.rust_result.clone(), - safe: policy.safe, - must_use: policy.must_use, - safe_allowlist_reason: policy.safe_allowlist_reason.clone(), - canonical_path: canonical_rust_path(intrinsic_abi, &policy.abi_id), - public_path: policy.public_rust_path.clone(), - compatibility_paths: policy.compatibility_rust_paths.clone(), - }, - dialect: CatalogDialect { - op_type: policy.dialect_op_type.clone(), - op_name: policy.dialect_op_name.clone(), - operands: dialect_operands, - results: dialect_results, - }, - llvm, - semantics: CatalogSemantics { - pure: policy.pure, - memory: policy.memory.clone(), - convergent: policy.convergent, - execution_scope: policy.execution_scope.clone(), - }, - target: CatalogTarget { - minimum_ptx: native_target.minimum_ptx, - hardware: native_target.hardware, - ptx_result: policy.ptx_result.clone(), - targets: policy.targets.clone(), - ptx_isa_version: policy.ptx_isa_version.clone(), - ptx_isa_section: policy.ptx_isa_section.clone(), - ptx_isa_url: policy.ptx_isa_url.clone(), - }, - backend, - backend_lowerings, - packed_atomic: policy.packed_atomic.clone(), - redux: policy.redux.clone(), - vote: policy.vote.clone(), - active_mask: policy.active_mask.clone(), - warp_match: policy.warp_match.clone(), - warp_barrier: policy.warp_barrier.clone(), - warp_shuffle: policy.warp_shuffle.clone(), - dot_product: policy.dot_product.clone(), - packed_alu: policy.packed_alu.clone(), - integer_minmax: policy.integer_minmax.clone(), - packed_conversion: policy.packed_conversion.clone(), - scalar_conversion: policy.scalar_conversion.clone(), - scalar_arithmetic: policy.scalar_arithmetic.clone(), - scalar_math: policy.scalar_math.clone(), - extended_minmax: policy.extended_minmax.clone(), - cp_async_copy: policy.cp_async_copy.clone(), - cp_async_control: policy.cp_async_control.clone(), - cp_async_mbarrier: policy.cp_async_mbarrier.clone(), - mbarrier_basic: policy.mbarrier_basic.clone(), - movmatrix: policy.movmatrix.clone(), - mbarrier_extended: policy.mbarrier_extended.clone(), - register_mma: policy.register_mma.clone(), - sparse_mma: policy.sparse_mma.clone(), - prmt: policy.prmt.clone(), - cluster_barrier: policy.cluster_barrier.clone(), - wgmma_control: policy.wgmma_control.clone(), - special_register: policy.special_register.clone(), - debug_control: policy.debug_control.clone(), - cluster_memory: policy.cluster_memory.clone(), - clc: policy.clc.clone(), - tma: policy.tma.clone(), - tcgen05: policy.tcgen05.clone(), - ldmatrix: policy - .ldmatrix_variant - .clone() - .map(|variant| CatalogLdmatrix { - variant, - safety: policy - .ldmatrix_safety - .clone() - .expect("validated ldmatrix safety"), - adapter: policy.ldmatrix_adapter.expect("validated ldmatrix adapter"), - selected_address_space: policy - .selected_address_space - .expect("validated ldmatrix address space"), - }), - lowering: policy.lowering.clone(), - expected_ptx: policy.expected_ptx.clone(), - summary: policy.summary.clone(), - }) -} - -fn imported_result_facts(properties: &[String]) -> Result { - let no_undef = properties.iter().any(|property| property == "NoUndef"); - let mut range = None; - for property in properties { - let Some(bounds) = property - .strip_prefix("Range')) - else { - continue; - }; - let (lower, upper_exclusive) = bounds - .split_once(',') - .with_context(|| format!("malformed return range property {property:?}"))?; - ensure!( - !lower.is_empty() && !upper_exclusive.is_empty(), - "malformed return range property {property:?}" - ); - ensure!(range.is_none(), "duplicate return range properties"); - range = Some(CatalogHalfOpenRange { - lower: lower.to_owned(), - upper_exclusive: upper_exclusive.to_owned(), - }); - } - Ok(CatalogLlvmResultFacts { no_undef, range }) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::model::{ImportedSelection, ScalarConversionAdmissionVariant}; - - #[test] - fn f32_redux_recipes_match_pinned_llvm_records() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - let declarations: BTreeMap<_, _> = imported - .intrinsics - .iter() - .map(|record| (record.source_record.as_str(), record)) - .collect(); - let cases = [ - ( - ReduxOperation::Fmin, - "redux.sync.min.f32 \t$dst, $src, $mask;", - &["sync", "min", "f32"][..], - ), - ( - ReduxOperation::FminNan, - "redux.sync.min.NaN.f32 \t$dst, $src, $mask;", - &["sync", "min", "NaN", "f32"], - ), - ( - ReduxOperation::FminAbs, - "redux.sync.min.abs.f32 \t$dst, $src, $mask;", - &["sync", "min", "abs", "f32"], - ), - ( - ReduxOperation::FminAbsNan, - "redux.sync.min.abs.NaN.f32 \t$dst, $src, $mask;", - &["sync", "min", "abs", "NaN", "f32"], - ), - ( - ReduxOperation::Fmax, - "redux.sync.max.f32 \t$dst, $src, $mask;", - &["sync", "max", "f32"], - ), - ( - ReduxOperation::FmaxNan, - "redux.sync.max.NaN.f32 \t$dst, $src, $mask;", - &["sync", "max", "NaN", "f32"], - ), - ( - ReduxOperation::FmaxAbs, - "redux.sync.max.abs.f32 \t$dst, $src, $mask;", - &["sync", "max", "abs", "f32"], - ), - ( - ReduxOperation::FmaxAbsNan, - "redux.sync.max.abs.NaN.f32 \t$dst, $src, $mask;", - &["sync", "max", "abs", "NaN", "f32"], - ), - ]; - - for (operation, expected_asm, expected_modifiers) in cases { - let recipe = redux_recipe(operation); - let declaration = declarations[recipe.source_record]; - assert_eq!(declaration.llvm_name, recipe.llvm_symbol); - assert_eq!(declaration.arguments, ["f32", "i32"]); - assert_eq!(declaration.results, ["f32"]); - assert_eq!(declaration.selections.len(), 1); - assert_eq!(declaration.selections[0].asm, expected_asm); - assert_eq!( - declaration.selections[0].predicates, - ["Subtarget->hasReduxSyncF32()"] - ); - assert_eq!(recipe.ptx_modifiers, expected_modifiers); - assert_eq!(recipe.minimum_ptx, "8.6"); - assert_eq!(recipe.minimum_sm, None); - assert_eq!(recipe.targets, REDUX_F32_TARGETS); - } - } - - fn sreg_pattern(special_register: &str) -> InstructionPattern { - InstructionPattern::new( - "mov", - &["u32"], - vec![ - OperandPattern::Register, - OperandPattern::Exact { - value: special_register.into(), - }, - ], - ) - } - - fn policy() -> OverlayIntrinsic { - OverlayIntrinsic { - id: "thread_idx_x".into(), - abi_id: "i0001".into(), - operation_key: "launch.thread_index.x".into(), - family: "sreg".into(), - source: None, - source_record: Some("int_nvvm_read_ptx_sreg_tid_x".into()), - rust_module: "sreg".into(), - rust_name: "thread_idx_x".into(), - rust_arguments: vec![], - rust_result: "u32".into(), - safe: true, - must_use: false, - safe_allowlist_reason: Some("no caller obligations".into()), - public_rust_path: "cuda_intrinsics::sreg::thread_idx_x".into(), - compatibility_rust_paths: vec!["cuda_device::thread::threadIdx_x".into()], - dialect_op_type: "ReadPtxSregTidXOp".into(), - dialect_op_name: "nvvm.read_ptx_sreg_tid_x".into(), - dialect_operands: vec![], - dialect_results: vec![], - llvm_symbol: Some("llvm.nvvm.read.ptx.sreg.tid.x".into()), - resolved_llvm_symbol: None, - llvm_arguments: vec![], - llvm_results: vec!["i32".into()], - pure: true, - memory: "none".into(), - convergent: false, - execution_scope: "thread".into(), - minimum_ptx: "2.0".into(), - minimum_sm: None, - ptx_result: "u32".into(), - targets: "all".into(), - ptx_isa_version: "9.3".into(), - ptx_isa_section: "10.1 Special Registers: %tid".into(), - ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/".into(), - lowering: "direct_nvvm".into(), - backend_lowerings: vec![], - packed_atomic: None, - redux: None, - vote: None, - active_mask: None, - warp_match: None, - warp_barrier: None, - warp_shuffle: None, - dot_product: None, - packed_alu: None, - integer_minmax: None, - packed_conversion: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cp_async_copy: None, - cp_async_control: None, - cp_async_mbarrier: None, - mbarrier_basic: None, - movmatrix: None, - mbarrier_extended: None, - register_mma: None, - sparse_mma: None, - prmt: None, - cluster_barrier: None, - wgmma_control: None, - special_register: None, - debug_control: None, - cluster_memory: None, - clc: None, - tma: None, - tcgen05: None, - ldmatrix_variant: None, - ldmatrix_safety: None, - ldmatrix_adapter: None, - selected_address_space: None, - expected_ptx: sreg_pattern("%tid.x"), - summary: "thread index".into(), - } - } - - fn distinct_policy() -> OverlayIntrinsic { - let mut record = policy(); - record.id = "thread_idx_y".into(); - record.abi_id = "i0002".into(); - record.operation_key = "launch.thread_index.y".into(); - record.source_record = Some("int_nvvm_read_ptx_sreg_tid_y".into()); - record.rust_name = "thread_idx_y".into(); - record.public_rust_path = "cuda_intrinsics::sreg::thread_idx_y".into(); - record.compatibility_rust_paths = vec!["cuda_device::thread::threadIdx_y".into()]; - record.dialect_op_type = "ReadPtxSregTidYOp".into(); - record.dialect_op_name = "nvvm.read_ptx_sreg_tid_y".into(); - record.llvm_symbol = Some("llvm.nvvm.read.ptx.sreg.tid.y".into()); - record.expected_ptx = sreg_pattern("%tid.y"); - record - } - - fn movmatrix_policy() -> OverlayIntrinsic { - let mut record = policy(); - record.id = "movmatrix_trans_b16".into(); - record.abi_id = "i0305".into(); - record.operation_key = "movmatrix.m8n8.trans.b16".into(); - record.family = "movmatrix".into(); - record.source = Some(IntrinsicSource::PtxNative { - instruction: "movmatrix.sync.aligned.m8n8.trans.b16".into(), - }); - record.source_record = None; - record.rust_module = "matrix".into(); - record.rust_name = "movmatrix_trans_b16".into(); - record.rust_arguments = vec!["u32".into()]; - record.rust_result = "u32".into(); - record.safe = false; - record.must_use = true; - record.safe_allowlist_reason = None; - record.public_rust_path = "cuda_intrinsics::matrix::movmatrix_trans_b16".into(); - record.compatibility_rust_paths = vec!["cuda_device::wmma::movmatrix_trans_b16".into()]; - record.dialect_op_type = "MovmatrixTransB16Op".into(); - record.dialect_op_name = "nvvm.movmatrix_trans_b16".into(); - record.dialect_operands = vec!["i32".into()]; - record.dialect_results = vec!["i32".into()]; - record.llvm_symbol = None; - record.resolved_llvm_symbol = None; - record.llvm_arguments.clear(); - record.llvm_results.clear(); - record.pure = false; - record.memory = "inaccessible_read_write".into(); - record.convergent = true; - record.execution_scope = "warp".into(); - record.minimum_ptx = "7.8".into(); - record.minimum_sm = Some("sm_75".into()); - record.ptx_result = "u32".into(); - record.targets = "all".into(); - record.ptx_isa_section = - "9.7.15.5.17 Warp-level matrix transpose instruction: movmatrix".into(); - record.lowering = "generated_movmatrix_inline_ptx".into(); - record.backend_lowerings = [IntrinsicBackend::LlvmNvptx, IntrinsicBackend::LibNvvm] - .into_iter() - .map(|backend| OverlayBackendLowering { - backend, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: match backend { - IntrinsicBackend::LlvmNvptx => "llvm-test", - IntrinsicBackend::LibNvvm => "libnvvm-test", - } - .into(), - targets: None, - minimum_ptx: Some("7.8".into()), - minimum_sm: Some("sm_75".into()), - }) - .collect(); - record.movmatrix = Some(crate::model::Movmatrix { - participation: MovmatrixParticipation::AllWarpLanesSameInstructionNoExitedLanes, - adapter: MovmatrixAdapter::PackedB16x2U32ToPackedB16x2U32, - runtime_validation: RuntimeValidation::Unexecuted, - }); - record.expected_ptx = InstructionPattern::new( - "movmatrix", - &["sync", "aligned", "m8n8", "trans", "b16"], - vec![OperandPattern::Register, OperandPattern::Register], - ); - record.summary = "Transposes one packed b16 matrix fragment across a warp.".into(); - record - } - - fn declaration() -> ImportedIntrinsic { - ImportedIntrinsic { - source_record: "int_nvvm_read_ptx_sreg_tid_x".into(), - llvm_name: "llvm.nvvm.read.ptx.sreg.tid.x".into(), - arguments: vec![], - results: vec!["i32".into()], - classes: vec!["NVVMPureIntrinsic".into()], - properties: vec![ - "IntrNoMem".into(), - "IntrSpeculatable".into(), - "NoUndef".into(), - "Range".into(), - ], - selections: vec![ImportedSelection { - source_record: "INT_PTX_SREG_TID_x".into(), - asm: "mov.u32 $d, %tid.x;".into(), - predicates: vec![], - constraints: Default::default(), - }], - } - } - - fn evidence() -> EvidenceRecord { - EvidenceRecord { - id: "thread_idx_x".into(), - source: None, - source_record: Some("int_nvvm_read_ptx_sreg_tid_x".into()), - llvm_symbol: Some("llvm.nvvm.read.ptx.sreg.tid.x".into()), - resolved_llvm_symbol: None, - llvm_arguments: vec![], - llvm_results: vec!["i32".into()], - concrete_llvm_arguments: vec![], - concrete_llvm_results: vec![], - target_triple: "nvptx64-nvidia-cuda".into(), - gpu_target: "sm_70".into(), - ptx_feature: "+ptx60".into(), - status: "lowered".into(), - stages: vec![], - declaration_attributes_canonicalized: None, - runtime_validation: None, - expected_ptx: sreg_pattern("%tid.x"), - } - } - - fn validate_test_evidence(policy: &OverlayIntrinsic, record: EvidenceRecord) -> Result<()> { - let file = EvidenceFile { - schema: 3, - backend_profile: "test".into(), - backend_kind: None, - llvm_revision: "test".into(), - backend_version: "LLVM version test".into(), - backend_sha256: "0123456789abcdef".into(), - artifact_path: None, - build_id_prefix: None, - nvvm_ir_version: None, - debug_ir_version: None, - records: vec![record], - }; - let indexed = IndexedEvidence { - file: &file, - record: &file.records[0], - backend_version: &file.backend_version, - backend_sha256: &file.backend_sha256, - }; - validate_evidence(policy, &indexed, None) - } - - fn shared_matrix_stage() -> EvidenceStage { - EvidenceStage { - targets: vec!["sm_80".into(), "ptx71".into()], - representation: "shared fixture".into(), - stage: EvidenceStageKind::BackendCodegen, - mechanism: Some(BackendLoweringMechanism::InlinePtx), - outcome: "succeeded".into(), - detail: "$dst remains fixture text".into(), - artifact_kind: None, - tool_path: None, - tool_version: None, - tool_sha256: None, - } - } - - fn synthetic_matrix_json() -> serde_json::Value { - serde_json::json!({ - "schema": 6, - "backend_profile": "matrix-test", - "backend_kind": "llvm_nvptx", - "llvm_revision": "test", - "backend_version": "LLVM matrix test", - "backend_sha256": "0123456789abcdef", - "defaults": { - "llvm_arguments": ["i32"], - "llvm_results": ["i32"], - "target_triple": "nvptx64-nvidia-cuda", - "gpu_target": "sm_80", - "ptx_feature": "+ptx71", - "status": "lowered" - }, - "fixtures": [{ - "id": "shared", - "coverage_count": 2, - "stages": [{ - "targets": ["sm_80", "ptx71"], - "representation": "shared fixture", - "stage": "backend_codegen", - "mechanism": "inline_ptx", - "outcome": "succeeded", - "detail": "$dst remains fixture text" - }] - }], - "matrices": [{ - "axes": [{ - "name": "element", - "values": ["s8", "u8"] - }], - "product_count": 2, - "fixtures": ["shared"], - "template": { - "id": "synthetic_${element}", - "source_record": "int_synthetic_${element}", - "llvm_symbol": "llvm.synthetic.${element}", - "expected_ptx": { - "mnemonic": "mma", - "modifiers": ["sync", "${element}"], - "operands": [{"kind": "register"}] - } - } - }], - "records": [{ - "id": "synthetic_explicit", - "source_record": "int_synthetic_explicit", - "llvm_symbol": "llvm.synthetic.explicit", - "llvm_arguments": ["i32"], - "llvm_results": ["i32"], - "target_triple": "nvptx64-nvidia-cuda", - "gpu_target": "sm_80", - "ptx_feature": "+ptx71", - "status": "lowered", - "expected_ptx": { - "mnemonic": "mma", - "modifiers": ["sync", "explicit"], - "operands": [{"kind": "register"}] - } - }] - }) - } - - fn policy_matrix_json() -> serde_json::Value { - serde_json::json!({ - "schema": 6, - "backend_profile": "matrix-test", - "llvm_revision": "test", - "backend_version": "LLVM matrix test", - "backend_sha256": "0123456789abcdef", - "defaults": { - "llvm_arguments": [], - "llvm_results": ["i32"], - "target_triple": "nvptx64-nvidia-cuda", - "gpu_target": "sm_70", - "ptx_feature": "+ptx60", - "status": "lowered" - }, - "fixtures": [{ - "id": "policy_fixture", - "coverage_count": 1, - "stages": [{ - "targets": ["sm_70", "ptx60"], - "representation": "policy fixture", - "stage": "backend_codegen", - "mechanism": "typed_nvvm", - "outcome": "succeeded", - "detail": "shared policy fixture" - }] - }], - "matrices": [{ - "axes": [{ - "name": "axis", - "values": ["x"] - }], - "product_count": 1, - "fixtures": ["policy_fixture"], - "template": { - "id": "thread_idx_${axis}", - "source_record": "int_nvvm_read_ptx_sreg_tid_${axis}", - "llvm_symbol": "llvm.nvvm.read.ptx.sreg.tid.${axis}", - "expected_ptx": { - "mnemonic": "mov", - "modifiers": ["u32"], - "operands": [ - {"kind": "register"}, - {"kind": "exact", "value": "%tid.${axis}"} - ] - } - } - }] - }) - } - - fn parse_synthetic_evidence(value: &serde_json::Value) -> Result { - parse_evidence_bytes(&serde_json::to_vec(value).unwrap(), "synthetic evidence") - } - - fn assert_synthetic_evidence_error(value: &serde_json::Value, expected: &str) { - let error = parse_synthetic_evidence(value).unwrap_err(); - let message = format!("{error:#}"); - assert!( - message.contains(expected), - "expected {expected:?} in {message:?}" - ); - } - - fn overlay_file(records: Vec) -> OverlayFile { - OverlayFile { - schema: OVERLAY_SCHEMA, - catalog_version: "test".into(), - intrinsic_abi: 1, - backend_profile: "test".into(), - shards: vec![], - intrinsics: records, - } - } - - fn bind_pinned_abi_ids(repo_root: &Path, overlay: &mut OverlayFile) { - let ledger_path = repo_root.join(format!("intrinsics/abi-v{}.toml", overlay.intrinsic_abi)); - let ledger: AbiLedgerFile = - toml::from_str(&std::fs::read_to_string(ledger_path).unwrap()).unwrap(); - bind_generated_abi_ids(overlay, &ledger).unwrap(); - } - - fn validate_imported_policy( - policy: &OverlayIntrinsic, - declaration: &ImportedIntrinsic, - ) -> Result<()> { - let source = resolve_policy_source(policy)?; - validate_policy(policy, &source, Some(declaration), 1) - } - - fn pinned_active_mask_and_warp_match_records() - -> BTreeMap { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let (overlay, _) = - read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")).unwrap(); - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - let declarations: BTreeMap<_, _> = imported - .intrinsics - .into_iter() - .map(|record| (record.source_record.clone(), record)) - .collect(); - - overlay - .intrinsics - .into_iter() - .filter(|record| matches!(record.family.as_str(), "active_mask" | "warp_match")) - .map(|policy| { - let declaration = declarations[policy.source_record.as_deref().unwrap()].clone(); - (policy.id.clone(), (policy, declaration)) - }) - .collect() - } - - fn pinned_mbarrier_basic_records() -> BTreeMap { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let (overlay, _) = - read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")).unwrap(); - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - let declarations: BTreeMap<_, _> = imported - .intrinsics - .into_iter() - .map(|record| (record.source_record.clone(), record)) - .collect(); - - overlay - .intrinsics - .into_iter() - .filter(|record| record.family == "mbarrier_basic") - .map(|policy| { - let declaration = declarations[policy.source_record.as_deref().unwrap()].clone(); - (policy.id.clone(), (policy, declaration)) - }) - .collect() - } - - fn pinned_cp_async_mbarrier_records() -> BTreeMap - { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let (overlay, _) = - read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")).unwrap(); - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - let declarations: BTreeMap<_, _> = imported - .intrinsics - .into_iter() - .map(|record| (record.source_record.clone(), record)) - .collect(); - - overlay - .intrinsics - .into_iter() - .filter(|record| record.family == "cp_async_mbarrier") - .map(|policy| { - let declaration = declarations[policy.source_record.as_deref().unwrap()].clone(); - (policy.id.clone(), (policy, declaration)) - }) - .collect() - } - - fn packed_policy(id: &str) -> OverlayIntrinsic { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")) - .unwrap() - .0 - .intrinsics - .into_iter() - .find(|record| record.id == id) - .unwrap() - } - - fn packed_alu_policy( - format: PackedAluFormat, - operation: PackedAluOperation, - ) -> OverlayIntrinsic { - let recipe = packed_alu_recipe(format, operation).expect("test recipe pair"); - let (rust_module, rust_type, dialect_type, adapter) = match format { - PackedAluFormat::Bf16x2 => ("bf16x2", "u32", "i32", PackedAluAdapter::DirectPackedU32), - PackedAluFormat::F16x2 => ("f16x2", "u32", "i32", PackedAluAdapter::DirectPackedU32), - PackedAluFormat::F32x2 => ("f32x2", "u64", "i64", PackedAluAdapter::DirectPackedU64), - }; - let mut record = policy(); - record.id = recipe.id.into(); - record.abi_id = recipe.abi_id.into(); - record.operation_key = recipe.operation_key.into(); - record.family = "packed_alu".into(); - match &recipe.source { - PackedAluRecipeSource::Imported { - record: source_record, - symbol, - resolved_symbol, - arguments, - results, - .. - } => { - record.source = None; - record.source_record = Some((*source_record).into()); - record.llvm_symbol = Some((*symbol).into()); - record.resolved_llvm_symbol = resolved_symbol.map(str::to_owned); - record.llvm_arguments = arguments.iter().map(|value| (*value).into()).collect(); - record.llvm_results = results.iter().map(|value| (*value).into()).collect(); - } - PackedAluRecipeSource::PtxNative => { - record.source = Some(IntrinsicSource::PtxNative { - instruction: recipe.ptx_mnemonic.into(), - }); - record.source_record = None; - record.llvm_symbol = None; - record.resolved_llvm_symbol = None; - record.llvm_arguments.clear(); - record.llvm_results.clear(); - } - } - record.rust_module = rust_module.into(); - record.rust_name = recipe.rust_name.into(); - record.rust_arguments = vec![rust_type.into(); recipe.arity]; - record.rust_result = rust_type.into(); - record.safe = true; - record.must_use = recipe.must_use; - record.safe_allowlist_reason = Some("the operation has no caller obligations".into()); - record.public_rust_path = format!("cuda_intrinsics::{rust_module}::{}", recipe.rust_name); - record.compatibility_rust_paths = - vec![format!("cuda_device::{rust_module}::{}", recipe.rust_name)]; - record.dialect_op_type = recipe.dialect_op_type.into(); - record.dialect_op_name = recipe.dialect_op_name.into(); - record.dialect_operands = vec![dialect_type.into(); recipe.arity]; - record.dialect_results = vec![dialect_type.into()]; - record.pure = true; - record.memory = "none".into(); - record.convergent = false; - record.execution_scope = "thread".into(); - record.minimum_ptx = recipe.minimum_ptx.into(); - record.minimum_sm = Some(recipe.minimum_sm.into()); - record.ptx_result = rust_type.into(); - record.ptx_isa_section = recipe.ptx_isa_section.into(); - record.ptx_isa_url = recipe.ptx_isa_url.into(); - record.lowering = "generated_packed_alu_inline_ptx".into(); - record.backend_lowerings = [IntrinsicBackend::LlvmNvptx, IntrinsicBackend::LibNvvm] - .into_iter() - .map(|backend| { - let (minimum_ptx, minimum_sm) = - packed_alu_backend_floor(&recipe, format, operation, backend); - crate::model::OverlayBackendLowering { - backend, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: format!("{backend:?}-test"), - targets: None, - minimum_ptx: Some(minimum_ptx.into()), - minimum_sm: Some(minimum_sm.into()), - } - }) - .collect(); - record.packed_alu = Some(crate::model::PackedAlu { - format, - native_minimum_sm: recipe.native_minimum_sm, - operation, - adapter, - }); - record.expected_ptx = InstructionPattern::new( - recipe.ptx_mnemonic.split('.').next().unwrap(), - recipe.modifiers, - vec![OperandPattern::Register; recipe.arity + 1], - ); - record.summary = format!("packed {rust_module} arithmetic"); - record - } - - fn packed_alu_declaration( - format: PackedAluFormat, - operation: PackedAluOperation, - ) -> Option { - let recipe = packed_alu_recipe(format, operation).expect("test recipe pair"); - let PackedAluRecipeSource::Imported { - record, - symbol, - arguments, - results, - properties, - selection, - selection_asm, - .. - } = recipe.source - else { - return None; - }; - let classes = if matches!(operation, PackedAluOperation::Min | PackedAluOperation::Max) { - vec!["Intrinsic".into()] - } else { - vec!["Intrinsic".into(), "NVVMPureIntrinsic".into()] - }; - let mut selections = vec![ImportedSelection { - source_record: selection.into(), - asm: selection_asm.into(), - predicates: vec![ - format!("Subtarget->getSmVersion() >= {}", recipe.native_minimum_sm), - format!( - "Subtarget->getPTXVersion() >= {}", - recipe.minimum_ptx.replace('.', "") - ), - ], - constraints: Default::default(), - }]; - if operation == PackedAluOperation::Abs { - selections.extend((0..5).map(|index| ImportedSelection { - source_record: format!("OTHER_ABS_{index}"), - asm: "abs.f32 $dst, $src0;".into(), - predicates: vec![], - constraints: Default::default(), - })); - } - Some(ImportedIntrinsic { - source_record: record.into(), - llvm_name: symbol.into(), - arguments: arguments.iter().map(|value| (*value).into()).collect(), - results: results.iter().map(|value| (*value).into()).collect(), - classes, - properties: properties.iter().map(|value| (*value).into()).collect(), - selections, - }) - } - - fn packed_conversion_policy( - destination_format: PackedConversionDestinationFormat, - rounding: PackedConversionRounding, - saturation: PackedConversionSaturation, - ) -> OverlayIntrinsic { - let conversion = crate::model::PackedConversion { - source_format: PackedConversionSourceFormat::F32x2, - destination_format, - rounding, - saturation, - adapter: PackedConversionAdapter::ReverseHighLowOperands, - }; - let recipe = packed_conversion_recipe(&conversion).expect("test packed-conversion recipe"); - let mut record = policy(); - record.id = recipe.id.into(); - record.abi_id = recipe.abi_id.into(); - record.operation_key = recipe.operation_key.into(); - record.family = "packed_conversion".into(); - record.source_record = Some(recipe.source_record.into()); - record.rust_module = "convert".into(); - record.rust_name = recipe.rust_name.into(); - record.rust_arguments = vec!["f32".into(), "f32".into()]; - let result_width = packed_conversion_result_width(&conversion); - record.rust_result = format!("u{result_width}"); - record.safe = true; - record.must_use = false; - record.safe_allowlist_reason = Some("the operation has no caller obligations".into()); - record.public_rust_path = format!("cuda_intrinsics::convert::{}", recipe.rust_name); - record.compatibility_rust_paths = vec![recipe.compatibility_path.into()]; - record.dialect_op_type = recipe.dialect_op_type.into(); - record.dialect_op_name = recipe.dialect_op_name.into(); - record.dialect_operands = vec!["f32".into(), "f32".into()]; - record.dialect_results = vec![format!("i{result_width}")]; - record.llvm_symbol = Some(recipe.llvm_symbol.into()); - record.llvm_arguments = vec!["f32".into(), "f32".into()]; - record.llvm_results = vec![recipe.llvm_result.into()]; - record.pure = true; - record.memory = "none".into(); - record.convergent = false; - record.execution_scope = "thread".into(); - let (minimum_ptx, minimum_sm) = packed_conversion_floor(&conversion); - record.minimum_ptx = minimum_ptx.into(); - record.minimum_sm = Some(minimum_sm.into()); - record.ptx_result = format!("u{result_width}"); - record.ptx_isa_section = "9.7.9.22 Data Movement and Conversion Instructions: cvt".into(); - record.ptx_isa_url = "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cvt".into(); - record.lowering = packed_conversion_lowering(&conversion).into(); - record.backend_lowerings = [IntrinsicBackend::LlvmNvptx, IntrinsicBackend::LibNvvm] - .into_iter() - .map(|backend| OverlayBackendLowering { - backend, - mechanism: packed_conversion_backend_mechanism(&conversion, backend), - evidence_profile: "test".into(), - targets: None, - minimum_ptx: Some(minimum_ptx.into()), - minimum_sm: Some(minimum_sm.into()), - }) - .collect(); - let modifiers = packed_conversion_ptx_modifiers(&conversion); - record.packed_conversion = Some(conversion); - record.expected_ptx = - InstructionPattern::new("cvt", &modifiers, vec![OperandPattern::Register; 3]); - record.summary = recipe.summary.into(); - record - } - - fn packed_conversion_declaration(policy: &OverlayIntrinsic) -> ImportedIntrinsic { - ImportedIntrinsic { - source_record: policy.source_record.clone().unwrap(), - llvm_name: policy.llvm_symbol.clone().unwrap(), - arguments: policy.llvm_arguments.clone(), - results: policy.llvm_results.clone(), - classes: vec!["Intrinsic".into(), "NVVMPureIntrinsic".into()], - properties: vec!["IntrNoMem".into(), "IntrSpeculatable".into()], - selections: vec![], - } - } - - fn packed_conversion_evidence(policy: &OverlayIntrinsic) -> EvidenceRecord { - let mut record = evidence(); - record.id = policy.id.clone(); - record.source_record = policy.source_record.clone(); - record.llvm_symbol = policy.llvm_symbol.clone(); - record.resolved_llvm_symbol = policy.resolved_llvm_symbol.clone(); - record.llvm_arguments = policy.llvm_arguments.clone(); - record.llvm_results = policy.llvm_results.clone(); - record.concrete_llvm_arguments = policy.llvm_arguments.clone(); - record.concrete_llvm_results = policy.llvm_results.clone(); - record.declaration_attributes_canonicalized = Some(true); - record.expected_ptx = policy.expected_ptx.clone(); - record - } - - fn redux_policy() -> OverlayIntrinsic { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")) - .unwrap() - .0 - .intrinsics - .into_iter() - .find(|record| record.id == "redux_sync_add") - .unwrap() - } - - fn redux_declaration() -> ImportedIntrinsic { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let text = std::fs::read_to_string(repo_root.join("intrinsics/imported.json")).unwrap(); - serde_json::from_str::(&text) - .unwrap() - .intrinsics - .into_iter() - .find(|record| record.source_record == "int_nvvm_redux_sync_add") - .unwrap() - } - - fn sync_policy() -> OverlayIntrinsic { - packed_policy("sync_threads") - } - - fn sync_declaration() -> ImportedIntrinsic { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let text = std::fs::read_to_string(repo_root.join("intrinsics/imported.json")).unwrap(); - serde_json::from_str::(&text) - .unwrap() - .intrinsics - .into_iter() - .find(|record| record.source_record == "int_nvvm_barrier_cta_sync_aligned_all") - .unwrap() - } - - fn warp_barrier_policy() -> OverlayIntrinsic { - packed_policy("sync_mask") - } - - fn warp_barrier_declaration() -> ImportedIntrinsic { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let text = std::fs::read_to_string(repo_root.join("intrinsics/imported.json")).unwrap(); - serde_json::from_str::(&text) - .unwrap() - .intrinsics - .into_iter() - .find(|record| record.source_record == "int_nvvm_bar_warp_sync") - .unwrap() - } - - fn vote_policy(mode: VoteMode) -> OverlayIntrinsic { - let recipe = vote_recipe(mode); - let mut record = policy(); - record.id = recipe.id.into(); - record.abi_id = recipe.abi_id.into(); - record.operation_key = recipe.operation_key.into(); - record.family = "vote".into(); - record.source_record = Some(recipe.source_record.into()); - record.rust_module = "warp".into(); - record.rust_name = recipe.rust_name.into(); - record.rust_arguments = vec!["u32".into(), "bool".into()]; - record.rust_result = recipe.rust_result.into(); - record.safe = false; - record.must_use = true; - record.safe_allowlist_reason = None; - record.public_rust_path = format!("cuda_intrinsics::warp::{}", recipe.rust_name); - record.compatibility_rust_paths = if recipe.has_compatibility_path { - vec![format!("cuda_device::warp::{}", recipe.rust_name)] - } else { - vec![] - }; - record.dialect_op_type = recipe.dialect_op_type.into(); - record.dialect_op_name = recipe.dialect_op_name.into(); - record.dialect_operands = vec!["i32".into(), "i1".into()]; - record.dialect_results = vec![recipe.llvm_result.into()]; - record.llvm_symbol = Some(recipe.llvm_symbol.into()); - record.llvm_arguments = vec!["i32".into(), "i1".into()]; - record.llvm_results = vec![recipe.llvm_result.into()]; - record.pure = false; - record.memory = "inaccessible_read_write".into(); - record.convergent = true; - record.execution_scope = "warp".into(); - record.minimum_ptx = "6.0".into(); - record.minimum_sm = Some("sm_30".into()); - record.ptx_result = recipe.rust_result.into(); - record.ptx_isa_section = "9.7.14.10 Warp Vote Instructions: vote.sync".into(); - record.ptx_isa_url = "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-vote-sync".into(); - record.lowering = "generated_vote".into(); - record.vote = Some(crate::model::Vote { - mode, - participation: VoteParticipation::ExecutingLaneNamedAllNamedLanesSameInstructionAndMask, - legacy_pre_sm70: PreSm70MemberMaskRule::AllNamedLanesConvergedAndOnlyNamedLanesActive, - adapter: VoteAdapter::DirectMaskPredicate, - mask_encoding: MaskEncoding::RegisterOrImmediate, - }); - record.expected_ptx = InstructionPattern::new( - "vote", - &["sync", recipe.ptx_mode, recipe.ptx_type], - vec![ - OperandPattern::Register, - OperandPattern::Register, - OperandPattern::RegisterOrImmediate, - ], - ); - record.summary = "warp vote".into(); - record - } - - fn vote_declaration(mode: VoteMode) -> ImportedIntrinsic { - let recipe = vote_recipe(mode); - let selection = |source_record: &str| ImportedSelection { - source_record: source_record.into(), - asm: format!( - "vote.sync.{}.{} \t$dest, $pred, $mask;", - recipe.ptx_mode, recipe.ptx_type - ), - predicates: vec![ - "Subtarget->getPTXVersion() >= 60".into(), - "Subtarget->getSmVersion() >= 30".into(), - ], - constraints: Default::default(), - }; - ImportedIntrinsic { - source_record: recipe.source_record.into(), - llvm_name: recipe.llvm_symbol.into(), - arguments: vec!["i32".into(), "i1".into()], - results: vec![recipe.llvm_result.into()], - classes: vec![ - "ClangBuiltin".into(), - "NVVMBuiltin".into(), - "SDPatternOperator".into(), - "Intrinsic".into(), - ], - properties: vec![ - "IntrConvergent".into(), - "IntrInaccessibleMemOnly".into(), - "IntrNoCallback".into(), - ], - selections: vec![ - selection(recipe.immediate_selection), - selection(recipe.register_selection), - ], - } - } - - fn warp_shuffle_policy( - mode: WarpShuffleMode, - value_kind: WarpShuffleValueKind, - ) -> OverlayIntrinsic { - let recipe = warp_shuffle_recipe(mode, value_kind); - let mut record = policy(); - record.id = recipe.id.into(); - record.abi_id = recipe.abi_id.into(); - record.operation_key = recipe.operation_key.into(); - record.family = "warp_shuffle".into(); - match recipe.source { - WarpShuffleRecipeSource::LlvmImported { - source_record, - llvm_symbol, - } => { - record.source_record = Some(source_record.into()); - record.llvm_symbol = Some(llvm_symbol.into()); - record.llvm_arguments = vec![ - "i32".into(), - recipe.dialect_value.into(), - "i32".into(), - "i32".into(), - ]; - record.llvm_results = vec![recipe.dialect_value.into()]; - } - WarpShuffleRecipeSource::PtxNative { instruction } => { - record.source = Some(IntrinsicSource::PtxNative { - instruction: instruction.into(), - }); - record.source_record = None; - record.llvm_symbol = None; - record.llvm_arguments.clear(); - record.llvm_results.clear(); - } - } - record.rust_module = "warp".into(); - record.rust_name = recipe.rust_name.into(); - record.rust_arguments = vec!["u32".into(), recipe.rust_value.into(), "u32".into()]; - record.rust_result = recipe.rust_value.into(); - record.safe = false; - record.must_use = true; - record.safe_allowlist_reason = None; - record.public_rust_path = format!("cuda_intrinsics::warp::{}", recipe.rust_name); - record.compatibility_rust_paths = vec![format!("cuda_device::warp::{}", recipe.rust_name)]; - record.dialect_op_type = recipe.dialect_op_type.into(); - record.dialect_op_name = recipe.dialect_op_name.into(); - record.dialect_operands = vec!["i32".into(), recipe.dialect_value.into(), "i32".into()]; - record.dialect_results = vec![recipe.dialect_value.into()]; - record.pure = false; - record.memory = "inaccessible_read_write".into(); - record.convergent = true; - record.execution_scope = "warp".into(); - record.minimum_ptx = "6.0".into(); - record.minimum_sm = Some("sm_30".into()); - record.ptx_result = recipe.rust_value.into(); - record.targets = "all".into(); - record.ptx_isa_section = - "9.7.9.6 Data Movement and Conversion Instructions: shfl.sync".into(); - record.ptx_isa_url = "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-shfl-sync".into(); - record.lowering = recipe.lowering.into(); - record.backend_lowerings = vec![ - crate::model::OverlayBackendLowering { - backend: IntrinsicBackend::LlvmNvptx, - mechanism: recipe.backend_mechanism, - evidence_profile: "llvm-test".into(), - targets: None, - minimum_ptx: Some("6.0".into()), - minimum_sm: Some("sm_30".into()), - }, - crate::model::OverlayBackendLowering { - backend: IntrinsicBackend::LibNvvm, - mechanism: recipe.backend_mechanism, - evidence_profile: "libnvvm-test".into(), - targets: None, - minimum_ptx: Some("6.0".into()), - minimum_sm: Some("sm_75".into()), - }, - ]; - record.warp_shuffle = Some(crate::model::WarpShuffle { - mode, - value_kind, - participation: - WarpShuffleParticipation::ExecutingLaneNamedAllNamedLanesSameInstructionAndMask, - legacy_pre_sm70: PreSm70MemberMaskRule::AllNamedLanesConvergedAndOnlyNamedLanesActive, - source_lane: WarpShuffleSourceLane::InRangeSourceActiveAndNamedOutOfRangeCopiesSelf, - adapter: recipe.adapter, - clamp: recipe.clamp, - lane_encoding: recipe.operand_encoding, - mask_encoding: recipe.operand_encoding, - }); - let operands = match recipe.adapter { - WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp => vec![ - OperandPattern::Register, - OperandPattern::Register, - OperandPattern::RegisterOrImmediate, - OperandPattern::Exact { - value: recipe.clamp.to_string(), - }, - OperandPattern::RegisterOrImmediate, - ], - WarpShuffleAdapter::MaskValueLaneOrDeltaSplitI64LowHighB32InsertClampReassemble => { - vec![ - OperandPattern::Exact { value: "lo".into() }, - OperandPattern::Exact { value: "lo".into() }, - OperandPattern::Register, - OperandPattern::Exact { - value: recipe.clamp.to_string(), - }, - OperandPattern::Register, - ] - } - }; - record.expected_ptx = - InstructionPattern::new("shfl", &["sync", recipe.ptx_mode, "b32"], operands); - record.summary = "synchronized warp shuffle".into(); - record - } - - fn warp_shuffle_declaration( - mode: WarpShuffleMode, - value_kind: WarpShuffleValueKind, - ) -> ImportedIntrinsic { - let recipe = warp_shuffle_recipe(mode, value_kind); - let WarpShuffleRecipeSource::LlvmImported { - source_record, - llvm_symbol, - } = recipe.source - else { - panic!("PTX-native i64 shuffles have no imported declaration"); - }; - let selections = (0..8) - .map(|index| ImportedSelection { - source_record: format!("anonymous_test_{index}"), - asm: format!( - "shfl.sync.{}.b32 \t$dst, $src, $offset, $mask, $threadmask;", - recipe.ptx_mode - ), - predicates: vec![ - "Subtarget->getPTXVersion() >= 60".into(), - "Subtarget->getSmVersion() >= 30".into(), - ], - constraints: Default::default(), - }) - .collect(); - ImportedIntrinsic { - source_record: source_record.into(), - llvm_name: llvm_symbol.into(), - arguments: vec![ - "i32".into(), - recipe.dialect_value.into(), - "i32".into(), - "i32".into(), - ], - results: vec![recipe.dialect_value.into()], - classes: vec![ - "ClangBuiltin".into(), - "NVVMBuiltin".into(), - "SDPatternOperator".into(), - "Intrinsic".into(), - ], - properties: vec![ - "IntrConvergent".into(), - "IntrInaccessibleMemOnly".into(), - "IntrNoCallback".into(), - ], - selections, - } - } - - fn sync_evidence(policy: &OverlayIntrinsic) -> EvidenceRecord { - let mut record = evidence(); - record.id = policy.id.clone(); - record.source_record = policy.source_record.clone(); - record.llvm_symbol = policy.llvm_symbol.clone(); - record.llvm_arguments = policy.llvm_arguments.clone(); - record.llvm_results = policy.llvm_results.clone(); - record.expected_ptx = policy.expected_ptx.clone(); - record - } - - fn dot_product_policy( - operation: DotProductOperation, - signedness: DotProductSignedness, - ) -> OverlayIntrinsic { - let recipe = dot_product_recipe(operation, signedness); - let mut record = policy(); - record.id = recipe.id.into(); - record.abi_id = match (operation, signedness) { - (DotProductOperation::Dp4a, DotProductSignedness::Signed) => "i0030", - (DotProductOperation::Dp4a, DotProductSignedness::Unsigned) => "i0031", - (DotProductOperation::Dp2a, DotProductSignedness::Signed) => "i0032", - (DotProductOperation::Dp2a, DotProductSignedness::Unsigned) => "i0033", - } - .into(); - record.operation_key = recipe.operation_key.into(); - record.family = "dotprod".into(); - record.source = None; - record.source_record = Some(recipe.source_record.into()); - record.rust_module = "dotprod".into(); - record.rust_name = recipe.rust_name.into(); - record.rust_arguments = vec!["u32".into(), "u32".into(), recipe.rust_value.into()]; - record.rust_result = recipe.rust_value.into(); - record.safe = true; - record.must_use = false; - record.safe_allowlist_reason = Some( - "per-thread integer arithmetic has no memory, pointer, or participation obligations" - .into(), - ); - record.public_rust_path = format!("cuda_intrinsics::dotprod::{}", recipe.rust_name); - record.compatibility_rust_paths = - vec![format!("cuda_device::dotprod::{}", recipe.rust_name)]; - record.dialect_op_type = recipe.dialect_op_type.into(); - record.dialect_op_name = recipe.dialect_op_name.into(); - record.dialect_operands = vec!["i32".into(), "i32".into(), "i32".into()]; - record.dialect_results = vec!["i32".into()]; - record.llvm_symbol = Some(recipe.llvm_symbol.into()); - record.resolved_llvm_symbol = None; - record.llvm_arguments = recipe - .llvm_arguments - .iter() - .map(|argument| (*argument).into()) - .collect(); - record.llvm_results = vec!["i32".into()]; - record.pure = true; - record.memory = "none".into(); - record.convergent = false; - record.execution_scope = "thread".into(); - record.minimum_ptx = "5.0".into(); - record.minimum_sm = Some("sm_61".into()); - record.ptx_result = recipe.rust_value.into(); - record.targets = "all".into(); - record.lowering = "generated_dotprod".into(); - record.backend_lowerings = vec![ - crate::model::OverlayBackendLowering { - backend: IntrinsicBackend::LlvmNvptx, - mechanism: BackendLoweringMechanism::TypedNvvm, - evidence_profile: "llvm-test".into(), - targets: None, - minimum_ptx: None, - minimum_sm: None, - }, - crate::model::OverlayBackendLowering { - backend: IntrinsicBackend::LibNvvm, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: "libnvvm-test".into(), - targets: None, - minimum_ptx: None, - minimum_sm: Some("sm_75".into()), - }, - ]; - record.dot_product = Some(crate::model::DotProduct { - operation, - signedness, - adapter: recipe.adapter, - }); - record.expected_ptx = InstructionPattern::new( - recipe.ptx_mnemonic, - recipe.ptx_modifiers, - vec![OperandPattern::Register; 4], - ); - record.summary = "packed integer dot product".into(); - record - } - - fn dot_product_declaration( - operation: DotProductOperation, - signedness: DotProductSignedness, - ) -> ImportedIntrinsic { - let recipe = dot_product_recipe(operation, signedness); - let selection = |source_record: &str, half: Option<(&str, i64)>| ImportedSelection { - source_record: source_record.into(), - asm: format!( - "{}.{} $dst, $a, $b, $c;", - recipe.ptx_mnemonic, - match half { - Some((name, _)) => { - let types = &recipe.ptx_modifiers[1..]; - format!("{name}.{}", types.join(".")) - } - None => recipe.ptx_modifiers.join("."), - } - ), - predicates: vec!["hasDotInstructions".into()], - constraints: crate::model::ImportedSelectionConstraints { - address_space: None, - immediate_bindings: half - .map(|(_, value)| { - vec![crate::model::ImportedImmediateBinding { - argument_index: 2, - value, - }] - }) - .unwrap_or_default(), - }, - }; - let selections = match operation { - DotProductOperation::Dp4a => vec![selection("DOT4", None)], - DotProductOperation::Dp2a => vec![ - selection("DOT2_hi", Some(("hi", -1))), - selection("DOT2_lo", Some(("lo", 0))), - ], - }; - ImportedIntrinsic { - source_record: recipe.source_record.into(), - llvm_name: recipe.llvm_symbol.into(), - arguments: recipe - .llvm_arguments - .iter() - .map(|argument| (*argument).into()) - .collect(), - results: vec!["i32".into()], - classes: vec!["NVVMPureIntrinsic".into()], - properties: recipe - .llvm_properties - .iter() - .map(|property| (*property).into()) - .collect(), - selections, - } - } - - fn dot_product_evidence(policy: &OverlayIntrinsic) -> EvidenceRecord { - let mut record = evidence(); - record.id = policy.id.clone(); - record.source_record = policy.source_record.clone(); - record.llvm_symbol = policy.llvm_symbol.clone(); - record.llvm_arguments = policy.llvm_arguments.clone(); - record.llvm_results = policy.llvm_results.clone(); - record.concrete_llvm_arguments = policy.llvm_arguments.clone(); - record.concrete_llvm_results = policy.llvm_results.clone(); - record.declaration_attributes_canonicalized = Some(true); - record.gpu_target = "sm_61".into(); - record.ptx_feature = "+ptx50".into(); - record.expected_ptx = policy.expected_ptx.clone(); - record - } - - fn validate_ptx_native_policy(policy: &OverlayIntrinsic) -> Result<()> { - let source = resolve_policy_source(policy)?; - validate_policy(policy, &source, None, 1) - } - - fn ledger_entry(record: &OverlayIntrinsic) -> AbiLedgerEntry { - AbiLedgerEntry { - abi_id: record.abi_id.clone(), - status: "active".into(), - catalog_id: record.id.clone(), - operation_key: record.operation_key.clone(), - raw_rust_signature: raw_rust_signature(record), - } - } - - fn ledger(entries: Vec) -> AbiLedgerFile { - AbiLedgerFile { - schema: 1, - intrinsic_abi: 1, - entries, - } - } - - #[test] - fn duplicate_values_are_rejected() { - let mut values = BTreeSet::new(); - insert_unique(&mut values, "thread_idx_x", "catalog ID").unwrap(); - let error = insert_unique(&mut values, "thread_idx_x", "catalog ID").unwrap_err(); - assert!(error.to_string().contains("duplicate catalog ID")); - } - - #[test] - fn legacy_evidence_schema_is_unchanged_and_rejects_matrix_fields() { - let legacy = EvidenceFile { - schema: 5, - backend_profile: "legacy".into(), - backend_kind: None, - llvm_revision: "test".into(), - backend_version: "LLVM legacy test".into(), - backend_sha256: "0123456789abcdef".into(), - artifact_path: None, - build_id_prefix: None, - nvvm_ir_version: None, - debug_ir_version: None, - records: vec![evidence()], - }; - let bytes = serde_json::to_vec(&legacy).unwrap(); - assert_eq!(parse_evidence_bytes(&bytes, "legacy").unwrap(), legacy); - - let mut with_matrix_field = serde_json::to_value(&legacy).unwrap(); - with_matrix_field["matrices"] = serde_json::json!([]); - let error = parse_synthetic_evidence(&with_matrix_field).unwrap_err(); - assert!(error.to_string().contains("legacy evidence")); - assert!(format!("{error:#}").contains("unknown field")); - } - - #[test] - fn compact_evidence_matrix_equals_explicit_records() { - let expanded = parse_synthetic_evidence(&synthetic_matrix_json()).unwrap(); - let mut expected = vec![EvidenceRecord { - id: "synthetic_explicit".into(), - source: None, - source_record: Some("int_synthetic_explicit".into()), - llvm_symbol: Some("llvm.synthetic.explicit".into()), - resolved_llvm_symbol: None, - llvm_arguments: vec!["i32".into()], - llvm_results: vec!["i32".into()], - concrete_llvm_arguments: vec![], - concrete_llvm_results: vec![], - target_triple: "nvptx64-nvidia-cuda".into(), - gpu_target: "sm_80".into(), - ptx_feature: "+ptx71".into(), - status: "lowered".into(), - stages: vec![], - declaration_attributes_canonicalized: None, - runtime_validation: None, - expected_ptx: InstructionPattern { - mnemonic: "mma".into(), - modifiers: vec!["sync".into(), "explicit".into()], - operands: vec![OperandPattern::Register], - }, - }]; - expected.extend(["s8", "u8"].into_iter().map(|element| EvidenceRecord { - id: format!("synthetic_{element}"), - source: None, - source_record: Some(format!("int_synthetic_{element}")), - llvm_symbol: Some(format!("llvm.synthetic.{element}")), - resolved_llvm_symbol: None, - llvm_arguments: vec!["i32".into()], - llvm_results: vec!["i32".into()], - concrete_llvm_arguments: vec![], - concrete_llvm_results: vec![], - target_triple: "nvptx64-nvidia-cuda".into(), - gpu_target: "sm_80".into(), - ptx_feature: "+ptx71".into(), - status: "lowered".into(), - stages: vec![shared_matrix_stage()], - declaration_attributes_canonicalized: None, - runtime_validation: None, - expected_ptx: InstructionPattern { - mnemonic: "mma".into(), - modifiers: vec!["sync".into(), element.into()], - operands: vec![OperandPattern::Register], - }, - })); - assert_eq!(expanded.schema, 6); - assert_eq!(expanded.records, expected); - } - - #[test] - fn matrix_identity_mutations_reach_existing_evidence_validation() { - let mut expanded = parse_synthetic_evidence(&policy_matrix_json()).unwrap(); - let record = expanded.records.pop().unwrap(); - validate_test_evidence(&policy(), record.clone()).unwrap(); - - let mut wrong_source = record.clone(); - wrong_source.source_record = Some("int_nvvm_read_ptx_sreg_tid_y".into()); - assert!( - validate_test_evidence(&policy(), wrong_source) - .unwrap_err() - .to_string() - .contains("source provenance mismatch") - ); - - let mut wrong_symbol = record.clone(); - wrong_symbol.llvm_symbol = Some("llvm.nvvm.read.ptx.sreg.tid.y".into()); - assert!( - validate_test_evidence(&policy(), wrong_symbol) - .unwrap_err() - .to_string() - .contains("signature mismatch") - ); - - let mut wrong_signature = record.clone(); - wrong_signature.llvm_arguments.push("i32".into()); - assert!( - validate_test_evidence(&policy(), wrong_signature) - .unwrap_err() - .to_string() - .contains("signature mismatch") - ); - - let mut wrong_ptx = record; - wrong_ptx.expected_ptx.modifiers.push("changed".into()); - assert!( - validate_test_evidence(&policy(), wrong_ptx) - .unwrap_err() - .to_string() - .contains("PTX expectation mismatch") - ); - } - - #[test] - fn evidence_matrix_rejects_bad_counts_fixtures_placeholders_and_collisions() { - let base = synthetic_matrix_json(); - - let mut bad_product = base.clone(); - bad_product["matrices"][0]["product_count"] = 3.into(); - assert_synthetic_evidence_error(&bad_product, "expands to 2 records"); - - let mut unknown_fixture = base.clone(); - unknown_fixture["matrices"][0]["fixtures"][0] = "missing".into(); - assert_synthetic_evidence_error(&unknown_fixture, "unknown fixture"); - - let mut uncovered_fixture = base.clone(); - let extra = uncovered_fixture["fixtures"][0].clone(); - uncovered_fixture["fixtures"] - .as_array_mut() - .unwrap() - .push(extra); - uncovered_fixture["fixtures"][1]["id"] = "unused".into(); - assert_synthetic_evidence_error(&uncovered_fixture, "not referenced"); - - let mut wrong_coverage = base.clone(); - wrong_coverage["fixtures"][0]["coverage_count"] = 1.into(); - assert_synthetic_evidence_error(&wrong_coverage, "covers 2 expanded records"); - - let mut malformed = base.clone(); - malformed["matrices"][0]["template"]["id"] = "synthetic_$element".into(); - assert_synthetic_evidence_error(&malformed, "malformed matrix placeholder"); - - let mut unknown_axis = base.clone(); - unknown_axis["matrices"][0]["template"]["id"] = "synthetic_${other}".into(); - assert_synthetic_evidence_error(&unknown_axis, "unknown matrix axis"); - - let mut collision = base.clone(); - collision["matrices"][0]["template"]["id"] = "synthetic".into(); - assert_synthetic_evidence_error(&collision, "duplicate expanded evidence ID"); - } - - #[test] - fn exact_operand_matrix_placeholders_fail_closed() { - let base = policy_matrix_json(); - - let mut unknown = base.clone(); - unknown["matrices"][0]["template"]["expected_ptx"]["operands"][1]["value"] = - "%tid.${other}".into(); - assert_synthetic_evidence_error(&unknown, "unknown matrix axis other"); - - let mut unterminated = base.clone(); - unterminated["matrices"][0]["template"]["expected_ptx"]["operands"][1]["value"] = - "%tid.${axis".into(); - assert_synthetic_evidence_error(&unterminated, "unterminated matrix placeholder"); - - let mut disallowed = base; - disallowed["matrices"][0]["template"]["expected_ptx"]["mnemonic"] = "mov.${axis}".into(); - assert_synthetic_evidence_error(&disallowed, "PTX mnemonic cannot contain"); - } - - #[test] - fn evidence_matrix_rejects_bad_axes_fixture_ids_and_stage_conflicts() { - let base = synthetic_matrix_json(); - - let mut no_fixture = base.clone(); - no_fixture["matrices"][0]["fixtures"] = serde_json::json!([]); - assert_synthetic_evidence_error(&no_fixture, "references no shared fixture"); - - let mut empty_axes = base.clone(); - empty_axes["matrices"][0]["axes"] = serde_json::json!([]); - assert_synthetic_evidence_error(&empty_axes, "has no axes"); - - let mut duplicate_axis = base.clone(); - let axis = duplicate_axis["matrices"][0]["axes"][0].clone(); - duplicate_axis["matrices"][0]["axes"] - .as_array_mut() - .unwrap() - .push(axis); - duplicate_axis["matrices"][0]["product_count"] = 4.into(); - assert_synthetic_evidence_error(&duplicate_axis, "axes must be unique and sorted"); - - let mut empty_values = base.clone(); - empty_values["matrices"][0]["axes"][0]["values"] = serde_json::json!([]); - assert_synthetic_evidence_error(&empty_values, "has no values"); - - let mut empty_axis_name = base.clone(); - empty_axis_name["matrices"][0]["axes"][0]["name"] = "".into(); - assert_synthetic_evidence_error(&empty_axis_name, "is not a safe token"); - - let mut empty_value = base.clone(); - empty_value["matrices"][0]["axes"][0]["values"][0] = "".into(); - assert_synthetic_evidence_error(&empty_value, "unsafe value"); - - let mut duplicate_value = base.clone(); - duplicate_value["matrices"][0]["axes"][0]["values"][1] = "s8".into(); - assert_synthetic_evidence_error(&duplicate_value, "duplicate value"); - - let mut unsafe_value = base.clone(); - unsafe_value["matrices"][0]["axes"][0]["values"][0] = "../s8".into(); - assert_synthetic_evidence_error(&unsafe_value, "unsafe value"); - - let mut unused_axis = base.clone(); - unused_axis["matrices"][0]["axes"] - .as_array_mut() - .unwrap() - .push(serde_json::json!({"name": "other", "values": ["x"]})); - assert_synthetic_evidence_error(&unused_axis, "unused axis"); - - let mut duplicate_fixture = base.clone(); - let fixture = duplicate_fixture["fixtures"][0].clone(); - duplicate_fixture["fixtures"] - .as_array_mut() - .unwrap() - .push(fixture); - assert_synthetic_evidence_error(&duplicate_fixture, "duplicate evidence fixture ID"); - - let mut fixture_placeholder = base.clone(); - fixture_placeholder["fixtures"][0]["stages"][0]["detail"] = "covers ${element}".into(); - assert_synthetic_evidence_error(&fixture_placeholder, "cannot contain matrix placeholders"); - - let mut missing_symbol = base.clone(); - missing_symbol["matrices"][0]["template"] - .as_object_mut() - .unwrap() - .remove("llvm_symbol"); - assert_synthetic_evidence_error(&missing_symbol, "missing field `llvm_symbol`"); - - let mut conflicting_stage = base; - conflicting_stage["matrices"][0]["template"]["facts"]["stages"] = - conflicting_stage["fixtures"][0]["stages"].clone(); - assert_synthetic_evidence_error(&conflicting_stage, "conflicting duplicate"); - } - - #[test] - fn overloaded_symbols_require_distinct_resolved_identities() { - let bf16 = packed_alu_policy(PackedAluFormat::Bf16x2, PackedAluOperation::Abs); - let f16 = packed_alu_policy(PackedAluFormat::F16x2, PackedAluOperation::Abs); - validate_unique_overlay(&[bf16.clone(), f16.clone()], 1).unwrap(); - - let mut unresolved = f16.clone(); - unresolved.resolved_llvm_symbol = None; - let error = validate_unique_overlay(&[bf16.clone(), unresolved], 1).unwrap_err(); - assert!(error.to_string().contains("without a resolved symbol")); - - let mut duplicate = f16; - duplicate.resolved_llvm_symbol = bf16.resolved_llvm_symbol.clone(); - let error = validate_unique_overlay(&[bf16, duplicate], 1).unwrap_err(); - assert!(error.to_string().contains("duplicate resolved LLVM symbol")); - } - - #[test] - fn overlay_manifest_loads_sorted_family_shards() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let (overlay, hash) = - read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")).unwrap(); - assert_eq!(overlay.schema, OVERLAY_SCHEMA); - assert_eq!(overlay.shards.len(), 64); - assert_eq!(overlay.intrinsics.len(), 1016); - assert_eq!( - overlay - .intrinsics - .iter() - .filter(|record| record.family == "scalar_arithmetic") - .count(), - 64 - ); - assert_eq!( - overlay - .intrinsics - .iter() - .filter(|record| record.family == "scalar_conversion") - .count(), - 10 - ); - assert_eq!( - overlay - .intrinsics - .iter() - .filter(|record| record.family == "prmt") - .count(), - 7 - ); - assert_eq!( - overlay - .intrinsics - .iter() - .filter(|record| record.family == "debug_control") - .count(), - 3 - ); - assert_eq!( - overlay - .intrinsics - .iter() - .filter(|record| record.family == "stmatrix") - .count(), - 4 - ); - assert_eq!( - overlay - .intrinsics - .iter() - .filter(|record| record.family == "packed_alu") - .count(), - 30 - ); - assert_eq!( - overlay - .intrinsics - .iter() - .filter(|record| record.family == "packed_conversion") - .count(), - 18 - ); - assert_eq!( - overlay - .intrinsics - .iter() - .filter(|record| record.family == "active_mask") - .count(), - 1 - ); - assert_eq!( - overlay - .intrinsics - .iter() - .filter(|record| record.family == "dotprod") - .count(), - 4 - ); - assert_eq!( - overlay - .intrinsics - .iter() - .filter(|record| record.family == "ldmatrix") - .count(), - 18 - ); - assert_eq!( - overlay - .intrinsics - .iter() - .filter(|record| record.family == "register_mma") - .count(), - 154 - ); - assert_eq!( - overlay - .intrinsics - .iter() - .filter(|record| record.family == "sparse_mma") - .count(), - 114 - ); - assert_eq!( - overlay - .intrinsics - .iter() - .filter(|record| record.family == "cp_async_copy") - .count(), - 8 - ); - assert_eq!( - overlay - .intrinsics - .iter() - .filter(|record| record.family == "cp_async_control") - .count(), - 3 - ); - assert_eq!( - overlay - .intrinsics - .iter() - .filter(|record| record.family == "cp_async_mbarrier") - .count(), - 4 - ); - assert_eq!( - overlay - .intrinsics - .iter() - .filter(|record| record.family == "mbarrier_basic") - .count(), - 4 - ); - assert_eq!( - overlay - .intrinsics - .iter() - .filter(|record| record.family == "cluster_memory") - .count(), - 2 - ); - assert_eq!( - overlay - .intrinsics - .iter() - .filter(|record| record.family == "clc") - .count(), - 6 - ); - assert_eq!( - overlay - .intrinsics - .iter() - .filter(|record| record.family == "sync") - .count(), - 4 - ); - assert_eq!( - overlay - .intrinsics - .iter() - .filter(|record| record.family == "vote") - .count(), - 4 - ); - assert_eq!( - overlay - .intrinsics - .iter() - .filter(|record| record.family == "warp_barrier") - .count(), - 1 - ); - assert_eq!( - overlay - .intrinsics - .iter() - .filter(|record| record.family == "warp_match") - .count(), - 4 - ); - assert_eq!( - overlay - .intrinsics - .iter() - .filter(|record| record.family == "warp_shuffle") - .count(), - 12 - ); - assert_eq!(hash.len(), 64); - - for invalid in [ - "../outside.toml", - "/absolute.toml", - "other/family.toml", - "overlay/../outside.toml", - "overlay/not-toml.json", - ] { - assert!(validate_overlay_shard_path(invalid).is_err(), "{invalid}"); - } - } - - fn test_f8f6f4_admission() -> SparseMmaF8F6F4Admission { - let formats = vec![ - SparseMmaElement::E2m1, - SparseMmaElement::E2m3, - SparseMmaElement::E3m2, - SparseMmaElement::E4m3, - SparseMmaElement::E5m2, - ]; - SparseMmaF8F6F4Admission { - llvm_evidence_profile: "llvm-test".into(), - libnvvm_evidence_profile: "libnvvm-test".into(), - runtime_validation: RuntimeValidation::Unexecuted, - a_elements: formats.clone(), - b_elements: formats, - product_count: 25, - } - } - - fn test_sparse_mma_f8f6f4_f16_admission() -> SparseMmaF8F6F4F16Admission { - SparseMmaF8F6F4F16Admission { - llvm_evidence_profile: "llvm-test".into(), - libnvvm_evidence_profile: "libnvvm-test".into(), - runtime_validation: RuntimeValidation::Unexecuted, - _legacy_first_abi_id: Some("i0525".into()), - a_elements: SPARSE_MMA_F8F6F4_ELEMENTS.into(), - b_elements: SPARSE_MMA_F8F6F4_ELEMENTS.into(), - product_count: 25, - } - } - - fn test_register_mma_f8f6f4_admission( - _accumulator: RegisterMmaAccumulator, - ) -> RegisterMmaF8F6F4Admission { - let formats = REGISTER_MMA_F8F6F4_ELEMENTS.to_vec(); - RegisterMmaF8F6F4Admission { - llvm_evidence_profile: "llvm-test".into(), - libnvvm_evidence_profile: "libnvvm-test".into(), - runtime_validation: RuntimeValidation::Unexecuted, - _legacy_first_abi_id: None, - a_elements: formats.clone(), - b_elements: formats, - product_count: 25, - targets: ["sm_120a", "sm_120f", "sm_121a", "sm_121f"] - .map(Into::into) - .into(), - } - } - - fn test_register_mma_fp8_admission() -> RegisterMmaFp8Admission { - RegisterMmaFp8Admission { - llvm_evidence_profile: "llvm-fp8-test".into(), - libnvvm_evidence_profile: "libnvvm-fp8-test".into(), - runtime_validation: RuntimeValidation::Unexecuted, - _legacy_first_abi_id: Some("i0504".into()), - shapes: REGISTER_MMA_FP8_SHAPES.into(), - accumulators: REGISTER_MMA_FP8_ACCUMULATORS.into(), - a_elements: REGISTER_MMA_FP8_ELEMENTS.into(), - b_elements: REGISTER_MMA_FP8_ELEMENTS.into(), - product_count: 16, - } - } - - fn test_register_mma_ampere_float_admission() -> RegisterMmaAmpereFloatAdmission { - RegisterMmaAmpereFloatAdmission { - llvm_evidence_profile: "llvm-ampere-float-test".into(), - libnvvm_evidence_profile: "libnvvm-ampere-float-test".into(), - runtime_validation: RuntimeValidation::Unexecuted, - _legacy_first_abi_id: Some("i0520".into()), - product_count: 5, - variants: REGISTER_MMA_AMPERE_FLOAT_VARIANTS.into(), - } - } - - fn test_prmt_admission() -> PrmtAdmission { - let variants = [ - PrmtMode::Generic, - PrmtMode::F4e, - PrmtMode::B4e, - PrmtMode::Rc8, - PrmtMode::Ecl, - PrmtMode::Ecr, - PrmtMode::Rc16, - ] - .map(|mode| crate::model::PrmtAdmissionVariant { - abi_id: prmt_recipe(mode).abi_id.into(), - mode, - }) - .into(); - PrmtAdmission { - llvm_evidence_profile: "llvm-test".into(), - libnvvm_evidence_profile: "libnvvm-test".into(), - runtime_validation: RuntimeValidation::Unexecuted, - variants, - } - } - - fn test_fp8_conversion_admission() -> PackedConversionFp8Admission { - PackedConversionFp8Admission { - llvm_evidence_profile: "llvm-test".into(), - libnvvm_evidence_profile: "libnvvm-test".into(), - runtime_validation: RuntimeValidation::Unexecuted, - destination_formats: vec![ - PackedConversionDestinationFormat::E4m3x2, - PackedConversionDestinationFormat::E5m2x2, - ], - saturations: vec![ - PackedConversionSaturation::Satfinite, - PackedConversionSaturation::SatfiniteRelu, - ], - product_count: 4, - } - } - - fn test_fp8_f16x2_conversion_admission() -> PackedConversionFp8F16x2Admission { - PackedConversionFp8F16x2Admission { - llvm_evidence_profile: "llvm-test".into(), - libnvvm_evidence_profile: "libnvvm-test".into(), - runtime_validation: RuntimeValidation::Unexecuted, - fp8_formats: vec![ - PackedConversionFp8Format::E4m3x2, - PackedConversionFp8Format::E5m2x2, - ], - directions: vec![ - PackedConversionFp8Direction::Pack, - PackedConversionFp8Direction::Unpack, - ], - relu_variants: true, - product_count: 8, - } - } - - fn test_cluster_barrier_admission() -> ClusterBarrierAdmission { - let variants = [ - ClusterBarrierMode::Arrive, - ClusterBarrierMode::ArriveAligned, - ClusterBarrierMode::ArriveRelaxed, - ClusterBarrierMode::ArriveRelaxedAligned, - ClusterBarrierMode::Wait, - ClusterBarrierMode::WaitAligned, - ] - .map(|mode| crate::model::ClusterBarrierAdmissionVariant { - abi_id: cluster_barrier_recipe(mode).abi_id.into(), - mode, - }) - .into(); - ClusterBarrierAdmission { - llvm_evidence_profile: "llvm-test".into(), - libnvvm_evidence_profile: "libnvvm-test".into(), - runtime_validation: RuntimeValidation::Unexecuted, - variants, - } - } - - fn test_wgmma_control_admission() -> WgmmaControlAdmission { - let variants = [ - WgmmaControlMode::Fence, - WgmmaControlMode::CommitGroup, - WgmmaControlMode::WaitGroup, - ] - .map(|mode| crate::model::WgmmaControlAdmissionVariant { - abi_id: wgmma_control_recipe(mode).abi_id.into(), - mode, - }) - .into(); - WgmmaControlAdmission { - llvm_evidence_profile: "llvm-test".into(), - libnvvm_evidence_profile: "libnvvm-test".into(), - runtime_validation: RuntimeValidation::Unexecuted, - variants, - } - } - - fn test_special_register_admission() -> SpecialRegisterAdmission { - SpecialRegisterAdmission { - llvm_evidence_profile: "rust-llvm-22.1.2-1cb4e383".into(), - libnvvm_evidence_profile: "cuda-13.3-libnvvm-13.3.33".into(), - runtime_validation: RuntimeValidation::Unexecuted, - registers: REVIEWED_SPECIAL_REGISTERS.into(), - product_count: REVIEWED_SPECIAL_REGISTERS.len(), - } - } - - fn test_debug_control_admission() -> DebugControlAdmission { - DebugControlAdmission { - llvm_evidence_profile: "llvm-debug-test".into(), - libnvvm_evidence_profile: "libnvvm-debug-test".into(), - runtime_validation: RuntimeValidation::Unexecuted, - operations: vec![ - DebugControlOperation::Trap, - DebugControlOperation::Breakpoint, - DebugControlOperation::Pmevent, - ], - abi_ids: vec!["i9001".into(), "i9002".into(), "i9003".into()], - } - } - - fn test_clc_admission() -> ClcAdmission { - let operations = [ - ClcOperation::TryCancel, - ClcOperation::TryCancelMulticast, - ClcOperation::QueryIsCanceled, - ClcOperation::QueryGetFirstCtaidX, - ClcOperation::QueryGetFirstCtaidY, - ClcOperation::QueryGetFirstCtaidZ, - ]; - ClcAdmission { - llvm_evidence_profile: "llvm-clc-test".into(), - libnvvm_evidence_profile: "libnvvm-clc-test".into(), - runtime_validation: RuntimeValidation::Unexecuted, - variants: operations - .into_iter() - .map(|operation| crate::model::ClcAdmissionVariant { - abi_id: clc_recipe(operation).abi_id.into(), - operation, - }) - .collect(), - } - } - - fn test_tma_admission() -> TmaAdmission { - TmaAdmission { - llvm_evidence_profile: "llvm-tma-test".into(), - libnvvm_evidence_profile: "libnvvm-tma-test".into(), - reduce_llvm_evidence_profile: Some("llvm-tma-reduce-test".into()), - reduce_libnvvm_evidence_profile: Some("libnvvm-tma-reduce-test".into()), - runtime_validation: RuntimeValidation::Unexecuted, - variants: TMA_OPERATIONS - .into_iter() - .map(|operation| crate::model::TmaAdmissionVariant { - abi_id: tma_recipe(operation).abi_id.into(), - operation, - }) - .collect(), - reduce_variants: tma_reduction_admission_variants(), - } - } - - fn test_tcgen05_admission() -> Tcgen05Admission { - let operations = [ - Tcgen05Operation::Alloc, - Tcgen05Operation::Dealloc, - Tcgen05Operation::RelinquishAllocPermit, - Tcgen05Operation::FenceBeforeThreadSync, - Tcgen05Operation::FenceAfterThreadSync, - Tcgen05Operation::Commit, - Tcgen05Operation::CommitSharedCluster, - Tcgen05Operation::MmaWsF16, - Tcgen05Operation::MmaF16, - Tcgen05Operation::MmaWsBf16, - Tcgen05Operation::MmaWsTf32, - Tcgen05Operation::CpSmemToTmem, - Tcgen05Operation::Ld16x256bX8Pure, - Tcgen05Operation::Ld16x256bPure, - Tcgen05Operation::LoadWait, - Tcgen05Operation::StoreWait, - Tcgen05Operation::AllocCg2, - Tcgen05Operation::DeallocCg2, - Tcgen05Operation::RelinquishAllocPermitCg2, - Tcgen05Operation::MmaF16Cg2, - Tcgen05Operation::CommitCg2, - Tcgen05Operation::CommitSharedClusterCg2, - Tcgen05Operation::CommitMulticastCg2, - Tcgen05Operation::CpSmemToTmemCg2, - Tcgen05Operation::CommitMulticast, - Tcgen05Operation::ShiftDown, - Tcgen05Operation::ShiftDownCg2, - ]; - Tcgen05Admission { - llvm_evidence_profile: "llvm-tcgen05-test".into(), - libnvvm_evidence_profile: "libnvvm-tcgen05-test".into(), - cp_llvm_evidence_profile: None, - cp_libnvvm_evidence_profile: None, - ld_llvm_evidence_profile: None, - ld_libnvvm_evidence_profile: None, - st_llvm_evidence_profile: None, - st_libnvvm_evidence_profile: None, - offset_llvm_evidence_profile: None, - offset_libnvvm_evidence_profile: None, - control_llvm_evidence_profile: Some("llvm-tcgen05-control-test".into()), - control_libnvvm_evidence_profile: Some("libnvvm-tcgen05-control-test".into()), - mma_llvm_evidence_profile: None, - mma_libnvvm_evidence_profile: None, - mma_llvm_target_contracts: vec![], - mma_libnvvm_target_contracts: vec![], - runtime_validation: RuntimeValidation::Unexecuted, - variants: operations - .into_iter() - .map(|operation| crate::model::Tcgen05AdmissionVariant { - abi_id: tcgen05_recipe(operation).abi_id.into(), - operation, - }) - .collect(), - cp_variants: vec![], - ld_variants: vec![], - st_variants: vec![], - ld_offset_variants: vec![], - st_offset_variants: vec![], - mma_variants: vec![], - } - } - - fn without_tcgen05_control(mut admission: Tcgen05Admission) -> Tcgen05Admission { - admission.variants.truncate(24); - admission.control_llvm_evidence_profile = None; - admission.control_libnvvm_evidence_profile = None; - admission - } - - fn test_tcgen05_cp_admission() -> Tcgen05Admission { - let mut admission = test_tcgen05_admission(); - admission.cp_llvm_evidence_profile = Some("llvm-tcgen05-cp-test".into()); - admission.cp_libnvvm_evidence_profile = Some("libnvvm-tcgen05-cp-test".into()); - admission.cp_variants = TCGEN05_CP_MEMBERS - .into_iter() - .flat_map(|member| { - [Tcgen05CpGroup::Cg1, Tcgen05CpGroup::Cg2] - .into_iter() - .map(move |group| (member, group)) - }) - .enumerate() - .map(|(index, (member, group))| Tcgen05CpAdmissionVariant { - abi_id: format!("i{:04}", 578 + index), - member, - group, - }) - .collect(); - admission - } - - fn test_tcgen05_ld_admission() -> Tcgen05Admission { - let mut admission = test_tcgen05_cp_admission(); - admission.ld_llvm_evidence_profile = Some("llvm-tcgen05-ld-test".into()); - admission.ld_libnvvm_evidence_profile = Some("libnvvm-tcgen05-ld-test".into()); - admission.ld_variants = TCGEN05_LD_VARIANTS - .into_iter() - .flat_map(|(shape, multiplicity)| { - [false, true] - .into_iter() - .map(move |pack16| (shape, multiplicity, pack16)) - }) - .enumerate() - .map( - |(index, (shape, multiplicity, pack16))| Tcgen05LdAdmissionVariant { - abi_id: format!("i{:04}", 612 + index), - shape, - multiplicity, - pack16, - }, - ) - .collect(); - admission - } - - fn test_tcgen05_st_admission() -> Tcgen05Admission { - let mut admission = test_tcgen05_ld_admission(); - admission.st_llvm_evidence_profile = Some("llvm-tcgen05-st-test".into()); - admission.st_libnvvm_evidence_profile = Some("libnvvm-tcgen05-st-test".into()); - admission.st_variants = TCGEN05_ST_VARIANTS - .into_iter() - .flat_map(|(shape, multiplicity)| { - [false, true] - .into_iter() - .map(move |unpack16| (shape, multiplicity, unpack16)) - }) - .enumerate() - .map( - |(index, (shape, multiplicity, unpack16))| Tcgen05StAdmissionVariant { - abi_id: format!("i{:04}", 670 + index), - shape, - multiplicity, - unpack16, - }, - ) - .collect(); - admission - } - - fn test_tcgen05_offset_admission() -> Tcgen05Admission { - let mut admission = test_tcgen05_st_admission(); - admission.offset_llvm_evidence_profile = Some("llvm-tcgen05-offset-test".into()); - admission.offset_libnvvm_evidence_profile = Some("libnvvm-tcgen05-offset-test".into()); - admission.ld_offset_variants = TCGEN05_OFFSET_LDST_VARIANTS - .into_iter() - .flat_map(|(shape, multiplicity)| { - [false, true] - .into_iter() - .map(move |pack16| (shape, multiplicity, pack16)) - }) - .enumerate() - .map( - |(index, (shape, multiplicity, pack16))| Tcgen05LdAdmissionVariant { - abi_id: format!("i{:04}", 728 + index), - shape, - multiplicity, - pack16, - }, - ) - .collect(); - admission.st_offset_variants = TCGEN05_OFFSET_LDST_VARIANTS - .into_iter() - .flat_map(|(shape, multiplicity)| { - [false, true] - .into_iter() - .map(move |unpack16| (shape, multiplicity, unpack16)) - }) - .enumerate() - .map( - |(index, (shape, multiplicity, unpack16))| Tcgen05StAdmissionVariant { - abi_id: format!("i{:04}", 744 + index), - shape, - multiplicity, - unpack16, - }, - ) - .collect(); - admission - } - - fn test_tcgen05_mma_admission() -> Tcgen05Admission { - let mut admission = test_tcgen05_admission(); - admission.mma_llvm_evidence_profile = Some("llvm-tcgen05-mma-test".into()); - admission.mma_libnvvm_evidence_profile = Some("libnvvm-tcgen05-mma-test".into()); - admission.mma_llvm_target_contracts = - expected_tcgen05_mma_target_contracts(IntrinsicBackend::LlvmNvptx); - admission.mma_libnvvm_target_contracts = - expected_tcgen05_mma_target_contracts(IntrinsicBackend::LibNvvm); - admission.mma_variants = TCGEN05_MMA_FORMS - .into_iter() - .map(|form| (form, None)) - .chain( - TCGEN05_MMA_ALIASES - .into_iter() - .map(|alias| (Tcgen05MmaForm::WsTensor, Some(alias))), - ) - .enumerate() - .map(|(index, (form, alias))| Tcgen05MmaAdmissionVariant { - abi_id: format!("i{:04}", 763 + index), - form, - alias, - }) - .chain( - TCGEN05_MMA_ALIASES - .into_iter() - .enumerate() - .map(|(index, alias)| Tcgen05MmaAdmissionVariant { - abi_id: format!("i{:04}", 1011 + index), - form: Tcgen05MmaForm::Shared, - alias: Some(alias), - }), - ) - .collect(); - admission - } - - fn assert_tcgen05_backend_target_split(record: &OverlayIntrinsic) { - let llvm = &record.backend_lowerings[0]; - let libnvvm = &record.backend_lowerings[1]; - assert_eq!(record.targets, TCGEN05_LLVM_TARGETS); - assert_eq!(llvm.backend, IntrinsicBackend::LlvmNvptx); - assert_eq!(llvm.targets, None); - assert_eq!(libnvvm.backend, IntrinsicBackend::LibNvvm); - assert_eq!(libnvvm.targets.as_deref(), Some(TCGEN05_LIBNVVM_TARGETS)); - assert_eq!( - backend_target_requirement(record, llvm).unwrap().hardware, - CatalogHardwareTarget::AnyOf { - alternatives: vec![ - CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 101 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 103 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, - ], - } - ); - assert_eq!( - backend_target_requirement(record, libnvvm) - .unwrap() - .hardware, - CatalogHardwareTarget::AnyOf { - alternatives: vec![ - CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 103 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, - ], - } - ); - } - - fn test_threadfence_admission() -> ThreadfenceAdmission { - ThreadfenceAdmission { - llvm_evidence_profile: "llvm-test".into(), - libnvvm_evidence_profile: "libnvvm-test".into(), - runtime_validation: RuntimeValidation::Unexecuted, - variants: vec![ - crate::model::ThreadfenceAdmissionVariant { - abi_id: "i0298".into(), - scope: ThreadfenceScope::Cta, - }, - crate::model::ThreadfenceAdmissionVariant { - abi_id: "i0299".into(), - scope: ThreadfenceScope::Device, - }, - crate::model::ThreadfenceAdmissionVariant { - abi_id: "i0300".into(), - scope: ThreadfenceScope::System, - }, - ], - } - } - - fn test_cluster_memory_admission() -> ClusterMemoryAdmission { - ClusterMemoryAdmission { - llvm_evidence_profile: "llvm-cluster-memory-test".into(), - libnvvm_evidence_profile: "libnvvm-cluster-memory-test".into(), - runtime_validation: RuntimeValidation::Unexecuted, - variants: vec![ - crate::model::ClusterMemoryAdmissionVariant { - abi_id: "i0320".into(), - operation: ClusterMemoryOperation::MapSharedRank, - }, - crate::model::ClusterMemoryAdmissionVariant { - abi_id: "i0321".into(), - operation: ClusterMemoryOperation::ReadU32, - }, - ], - } - } - - fn test_stmatrix_admission() -> StmatrixAdmission { - let variants = [ - (StmatrixMultiplicity::X2, StmatrixLayout::Normal, "i0301"), - ( - StmatrixMultiplicity::X2, - StmatrixLayout::Transposed, - "i0302", - ), - (StmatrixMultiplicity::X4, StmatrixLayout::Normal, "i0303"), - ( - StmatrixMultiplicity::X4, - StmatrixLayout::Transposed, - "i0304", - ), - ] - .map( - |(multiplicity, layout, abi_id)| crate::model::StmatrixAdmissionVariant { - abi_id: abi_id.into(), - multiplicity, - layout, - }, - ) - .into(); - StmatrixAdmission { - llvm_evidence_profile: "llvm-test".into(), - libnvvm_evidence_profile: "libnvvm-test".into(), - runtime_validation: RuntimeValidation::Unexecuted, - variants, - } - } - - fn test_mbarrier_extended_admission() -> MbarrierExtendedAdmission { - let variants = [ - MbarrierExtendedOperation::ArriveExpectTxCta, - MbarrierExtendedOperation::ArriveExpectTxCluster, - MbarrierExtendedOperation::ArriveRemoteCluster, - MbarrierExtendedOperation::TryWaitTokenCta, - MbarrierExtendedOperation::TryWaitParityCta, - MbarrierExtendedOperation::TryWaitParityCluster, - MbarrierExtendedOperation::FenceProxyAsyncSharedCta, - MbarrierExtendedOperation::FenceMbarrierInitReleaseCluster, - MbarrierExtendedOperation::FenceProxyAsyncGenericReleaseSharedCtaCluster, - MbarrierExtendedOperation::FenceProxyAsyncGenericAcquireSharedClusterCluster, - MbarrierExtendedOperation::Nanosleep, - ] - .map(|operation| crate::model::MbarrierExtendedAdmissionVariant { - abi_id: mbarrier_extended_recipe(operation).abi_id.into(), - operation, - }) - .into(); - MbarrierExtendedAdmission { - llvm_evidence_profile: "llvm-test".into(), - libnvvm_evidence_profile: "libnvvm-test".into(), - runtime_validation: RuntimeValidation::Unexecuted, - variants, - } - } - - #[test] - fn overlay_shard_schema_range_is_composable_and_new_fields_fail_closed() { - let shard = |schema, sparse_mma_f8f6f4_f32, prmt| OverlayShardFile { - schema, - family: "sparse_mma".into(), - intrinsics: vec![], - register_mma_int4: None, - register_mma_int8: None, - register_mma_b1: None, - register_mma_f8f6f4_f32: None, - register_mma_f8f6f4_f16: None, - register_mma_mxf8f6f4_f32: None, - register_mma_fp8: None, - register_mma_ampere_float: None, - sparse_mma_integer: None, - sparse_mma_f8f6f4_f32, - sparse_mma_f8f6f4_f16: None, - prmt, - packed_conversion_fp8: None, - packed_conversion_fp8_f16x2: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cluster_sreg: None, - cluster_barrier: None, - mbarrier_extended: None, - special_registers: None, - debug_control: None, - threadfence: None, - cluster_memory: None, - stmatrix: None, - clc: None, - wgmma_controls: None, - tma: None, - tcgen05: None, - }; - let path = Path::new("intrinsics/overlay/test.toml"); - validate_overlay_shard_schema(&shard(26, None, None), path).unwrap(); - validate_overlay_shard_schema(&shard(27, None, None), path).unwrap(); - validate_overlay_shard_schema(&shard(28, None, None), path).unwrap(); - validate_overlay_shard_schema(&shard(29, None, None), path).unwrap(); - validate_overlay_shard_schema(&shard(30, None, None), path).unwrap(); - validate_overlay_shard_schema(&shard(31, None, None), path).unwrap(); - validate_overlay_shard_schema(&shard(32, None, None), path).unwrap(); - validate_overlay_shard_schema(&shard(33, None, None), path).unwrap(); - validate_overlay_shard_schema(&shard(34, None, None), path).unwrap(); - validate_overlay_shard_schema(&shard(35, None, None), path).unwrap(); - validate_overlay_shard_schema(&shard(27, Some(test_f8f6f4_admission()), None), path) - .unwrap(); - validate_overlay_shard_schema_with_max( - &shard(27, Some(test_f8f6f4_admission()), None), - path, - 30, - ) - .unwrap(); - validate_overlay_shard_schema(&shard(28, None, Some(test_prmt_admission())), path).unwrap(); - validate_overlay_shard_schema_with_max( - &shard(28, None, Some(test_prmt_admission())), - path, - 30, - ) - .unwrap(); - - assert!(validate_overlay_shard_schema(&shard(25, None, None), path).is_err()); - for schema in 35..=OVERLAY_SHARD_SCHEMA { - validate_overlay_shard_schema(&shard(schema, None, None), path).unwrap(); - } - assert!( - validate_overlay_shard_schema(&shard(OVERLAY_SHARD_SCHEMA + 1, None, None), path) - .is_err() - ); - let error = - validate_overlay_shard_schema(&shard(26, Some(test_f8f6f4_admission()), None), path) - .unwrap_err(); - assert!( - error - .to_string() - .contains("requires overlay shard schema 27") - ); - let error = - validate_overlay_shard_schema(&shard(27, None, Some(test_prmt_admission())), path) - .unwrap_err(); - assert!( - error - .to_string() - .contains("requires overlay shard schema 28") - ); - let mut f16_mma = shard(REGISTER_MMA_F8F6F4_F16_SHARD_SCHEMA, None, None); - f16_mma.family = "register_mma".into(); - f16_mma.register_mma_f8f6f4_f16 = Some(test_register_mma_f8f6f4_admission( - RegisterMmaAccumulator::F16, - )); - validate_overlay_shard_schema(&f16_mma, path).unwrap(); - f16_mma.schema -= 1; - assert!( - validate_overlay_shard_schema(&f16_mma, path) - .unwrap_err() - .to_string() - .contains("requires overlay shard schema 47") - ); - - let mut sparse_f16_mma = shard(SPARSE_MMA_F8F6F4_F16_SHARD_SCHEMA, None, None); - sparse_f16_mma.sparse_mma_f8f6f4_f16 = Some(test_sparse_mma_f8f6f4_f16_admission()); - validate_overlay_shard_schema(&sparse_f16_mma, path).unwrap(); - sparse_f16_mma.schema -= 1; - assert!( - validate_overlay_shard_schema(&sparse_f16_mma, path) - .unwrap_err() - .to_string() - .contains("requires overlay shard schema 50") - ); - - let mut standard_fp8_mma = shard(REGISTER_MMA_FP8_SHARD_SCHEMA, None, None); - standard_fp8_mma.family = "register_mma".into(); - standard_fp8_mma.register_mma_fp8 = Some(test_register_mma_fp8_admission()); - validate_overlay_shard_schema(&standard_fp8_mma, path).unwrap(); - standard_fp8_mma.schema -= 1; - assert!( - validate_overlay_shard_schema(&standard_fp8_mma, path) - .unwrap_err() - .to_string() - .contains("requires overlay shard schema 48") - ); - - let fp8_shard = |schema| OverlayShardFile { - schema, - family: "packed_conversion".into(), - intrinsics: vec![], - register_mma_int4: None, - register_mma_int8: None, - register_mma_b1: None, - register_mma_f8f6f4_f32: None, - register_mma_f8f6f4_f16: None, - register_mma_mxf8f6f4_f32: None, - register_mma_fp8: None, - register_mma_ampere_float: None, - sparse_mma_integer: None, - sparse_mma_f8f6f4_f32: None, - sparse_mma_f8f6f4_f16: None, - prmt: None, - packed_conversion_fp8: Some(test_fp8_conversion_admission()), - packed_conversion_fp8_f16x2: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cluster_sreg: None, - cluster_barrier: None, - mbarrier_extended: None, - special_registers: None, - debug_control: None, - threadfence: None, - cluster_memory: None, - stmatrix: None, - clc: None, - wgmma_controls: None, - tma: None, - tcgen05: None, - }; - validate_overlay_shard_schema(&fp8_shard(29), path).unwrap(); - validate_overlay_shard_schema_with_max(&fp8_shard(29), path, 30).unwrap(); - let error = validate_overlay_shard_schema(&fp8_shard(28), path).unwrap_err(); - assert!( - error - .to_string() - .contains("requires overlay shard schema 29") - ); - - let cluster_shard = OverlayShardFile { - schema: 31, - family: "cluster_barrier".into(), - intrinsics: vec![], - register_mma_int4: None, - register_mma_int8: None, - register_mma_b1: None, - register_mma_f8f6f4_f32: None, - register_mma_f8f6f4_f16: None, - register_mma_mxf8f6f4_f32: None, - register_mma_fp8: None, - register_mma_ampere_float: None, - sparse_mma_integer: None, - sparse_mma_f8f6f4_f32: None, - sparse_mma_f8f6f4_f16: None, - prmt: None, - packed_conversion_fp8: None, - packed_conversion_fp8_f16x2: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cluster_sreg: None, - cluster_barrier: Some(test_cluster_barrier_admission()), - mbarrier_extended: None, - special_registers: None, - debug_control: None, - threadfence: None, - cluster_memory: None, - stmatrix: None, - clc: None, - wgmma_controls: None, - tma: None, - tcgen05: None, - }; - validate_overlay_shard_schema_with_max(&cluster_shard, path, 31).unwrap(); - let mut old_cluster_shard = cluster_shard; - old_cluster_shard.schema = 30; - assert!( - validate_overlay_shard_schema_with_max(&old_cluster_shard, path, 31) - .unwrap_err() - .to_string() - .contains("requires overlay shard schema 31") - ); - - let extended_shard = OverlayShardFile { - schema: MBARRIER_EXTENDED_SHARD_SCHEMA, - family: "mbarrier_extended".into(), - intrinsics: vec![], - register_mma_int4: None, - register_mma_int8: None, - register_mma_b1: None, - register_mma_f8f6f4_f32: None, - register_mma_f8f6f4_f16: None, - register_mma_mxf8f6f4_f32: None, - register_mma_fp8: None, - register_mma_ampere_float: None, - sparse_mma_integer: None, - sparse_mma_f8f6f4_f32: None, - sparse_mma_f8f6f4_f16: None, - prmt: None, - packed_conversion_fp8: None, - packed_conversion_fp8_f16x2: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cluster_sreg: None, - cluster_barrier: None, - mbarrier_extended: Some(test_mbarrier_extended_admission()), - special_registers: None, - debug_control: None, - threadfence: None, - cluster_memory: None, - stmatrix: None, - clc: None, - wgmma_controls: None, - tma: None, - tcgen05: None, - }; - validate_overlay_shard_schema_with_max( - &extended_shard, - path, - MBARRIER_EXTENDED_SHARD_SCHEMA, - ) - .unwrap(); - let mut old_extended_shard = extended_shard; - old_extended_shard.schema = MBARRIER_EXTENDED_SHARD_SCHEMA - 1; - assert!( - validate_overlay_shard_schema_with_max( - &old_extended_shard, - path, - MBARRIER_EXTENDED_SHARD_SCHEMA, - ) - .unwrap_err() - .to_string() - .contains("requires overlay shard schema 40") - ); - } - - #[test] - fn stmatrix_admission_is_closed_and_uses_schema_35() { - let shard = |schema| { - toml::from_str::(&format!( - r#" -schema = {schema} -family = "stmatrix" - -[stmatrix] -llvm_evidence_profile = "llvm-test" -libnvvm_evidence_profile = "libnvvm-test" -runtime_validation = "unexecuted" - -[[stmatrix.variant]] -abi_id = "i0301" -multiplicity = "x2" -layout = "normal" - -[[stmatrix.variant]] -abi_id = "i0302" -multiplicity = "x2" -layout = "transposed" - -[[stmatrix.variant]] -abi_id = "i0303" -multiplicity = "x4" -layout = "normal" - -[[stmatrix.variant]] -abi_id = "i0304" -multiplicity = "x4" -layout = "transposed" -"# - )) - .unwrap() - }; - let path = Path::new("intrinsics/overlay/stmatrix.toml"); - - let old = shard(STMATRIX_SHARD_SCHEMA - 1); - assert!( - validate_overlay_shard_schema(&old, path) - .unwrap_err() - .to_string() - .contains("requires overlay shard schema 35") - ); - - let current = shard(STMATRIX_SHARD_SCHEMA); - validate_overlay_shard_schema(¤t, path).unwrap(); - let admission = current.stmatrix.unwrap(); - let records = expand_stmatrix_admission(&admission).unwrap(); - assert_eq!( - records - .iter() - .map(|record| (record.abi_id.as_str(), record.id.as_str())) - .collect::>(), - [ - ("i0301", "stmatrix_m8n8_x2_b16"), - ("i0302", "stmatrix_m8n8_x2_trans_b16"), - ("i0303", "stmatrix_m8n8_x4_b16"), - ("i0304", "stmatrix_m8n8_x4_trans_b16"), - ] - ); - - let mut reordered = admission.clone(); - reordered.variants.swap(0, 1); - assert!(expand_stmatrix_admission(&reordered).is_err()); - - let mut wrong_id = admission.clone(); - wrong_id.variants[0].abi_id = "i0302".into(); - assert!(expand_stmatrix_admission(&wrong_id).is_err()); - - let mut executed = admission; - executed.runtime_validation = RuntimeValidation::Executed; - assert!(expand_stmatrix_admission(&executed).is_err()); - } - - #[test] - fn pinned_stmatrix_records_match_llvm_and_reject_contract_drift() { - let records = expand_stmatrix_admission(&test_stmatrix_admission()).unwrap(); - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - let declarations = imported - .intrinsics - .iter() - .map(|record| (record.source_record.as_str(), record)) - .collect::>(); - - for record in &records { - let declaration = declarations[record.source_record.as_deref().unwrap()]; - assert!(declaration.selections.is_empty()); - validate_imported_policy(record, declaration).unwrap(); - } - - let declaration = declarations["int_nvvm_stmatrix_sync_aligned_m8n8_x2_b16"]; - let mut changed = records[0].clone(); - changed.memory = "read_write".into(); - assert!(validate_imported_policy(&changed, declaration).is_err()); - - let mut changed = records[0].clone(); - changed.convergent = false; - assert!(validate_imported_policy(&changed, declaration).is_err()); - - let mut changed = records[0].clone(); - changed.backend_lowerings[0].mechanism = BackendLoweringMechanism::InlinePtx; - assert!(validate_imported_policy(&changed, declaration).is_err()); - - let mut changed = records[0].clone(); - changed.minimum_sm = Some("sm_80".into()); - assert!(validate_imported_policy(&changed, declaration).is_err()); - } - - #[test] - fn cluster_sreg_admission_uses_its_fixed_introduction_schema() { - let shard = |schema| { - toml::from_str::(&format!( - r#" -schema = {schema} -family = "sreg" - -[cluster_sreg] -axes = ["x", "y", "z"] -xyz_product_count = 12 -record_count = 14 -"# - )) - .unwrap() - }; - let path = Path::new("intrinsics/overlay/sreg_cluster.toml"); - - let old = shard(CLUSTER_SREG_SHARD_SCHEMA - 1); - assert!( - validate_overlay_shard_schema(&old, path) - .unwrap_err() - .to_string() - .contains("requires overlay shard schema 30") - ); - - let current = shard(CLUSTER_SREG_SHARD_SCHEMA); - validate_overlay_shard_schema(¤t, path).unwrap(); - let admission = current.cluster_sreg.unwrap(); - assert_eq!(expand_cluster_sreg_admission(&admission).unwrap().len(), 14); - - let mut wrong_axes = admission.clone(); - wrong_axes.axes.swap(0, 1); - assert!(expand_cluster_sreg_admission(&wrong_axes).is_err()); - - let mut wrong_count = admission; - wrong_count.record_count = 13; - assert!(expand_cluster_sreg_admission(&wrong_count).is_err()); - } - - #[test] - fn special_register_admission_is_closed_and_schema_gated() { - let admission = test_special_register_admission(); - let records = expand_special_register_admission(&admission).unwrap(); - assert_eq!(records.len(), 12); - assert_eq!( - records - .iter() - .map(|record| record.id.as_str()) - .collect::>(), - [ - "clock", - "clock64", - "globaltimer", - "envreg1", - "envreg2", - "smid", - "nsmid", - "gridid", - "warpid", - "nwarpid", - "dynamic_smem_size", - "total_smem_size", - ] - ); - - let mut reordered = admission.clone(); - reordered.registers.swap(0, 1); - assert!(expand_special_register_admission(&reordered).is_err()); - - let mut wrong_count = admission.clone(); - wrong_count.product_count -= 1; - assert!(expand_special_register_admission(&wrong_count).is_err()); - - let mut executed = admission.clone(); - executed.runtime_validation = RuntimeValidation::Executed; - assert!(expand_special_register_admission(&executed).is_err()); - - let shard = |schema| OverlayShardFile { - schema, - family: "sreg".into(), - intrinsics: vec![], - register_mma_int4: None, - register_mma_int8: None, - register_mma_b1: None, - register_mma_f8f6f4_f32: None, - register_mma_f8f6f4_f16: None, - register_mma_mxf8f6f4_f32: None, - register_mma_fp8: None, - register_mma_ampere_float: None, - sparse_mma_integer: None, - sparse_mma_f8f6f4_f32: None, - sparse_mma_f8f6f4_f16: None, - prmt: None, - packed_conversion_fp8: None, - packed_conversion_fp8_f16x2: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cluster_sreg: None, - cluster_barrier: None, - mbarrier_extended: None, - special_registers: Some(admission.clone()), - debug_control: None, - threadfence: None, - cluster_memory: None, - stmatrix: None, - clc: None, - wgmma_controls: None, - tma: None, - tcgen05: None, - }; - let path = Path::new("intrinsics/overlay/sreg_special.toml"); - validate_overlay_shard_schema(&shard(SPECIAL_REGISTER_SHARD_SCHEMA), path).unwrap(); - assert!( - validate_overlay_shard_schema(&shard(SPECIAL_REGISTER_SHARD_SCHEMA - 1), path) - .unwrap_err() - .to_string() - .contains("requires overlay shard schema 32") - ); - } - - #[test] - fn pinned_special_registers_preserve_apis_widths_and_backend_routes() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let base = load_resolution_base(&repo_root).unwrap(); - let declarations = index_imported_intrinsics(&base.imported).unwrap(); - let records = base - .overlay - .intrinsics - .iter() - .filter(|record| record.special_register.is_some()) - .collect::>(); - assert_eq!(records.len(), 12); - assert_eq!( - records - .iter() - .map(|record| record.abi_id.clone()) - .collect::>(), - (283..=294) - .map(|id| format!("i{id:04}")) - .collect::>() - ); - - for record in &records { - let source = resolve_policy_source(record).unwrap(); - let declaration = resolve_imported_declaration(record, &source, &declarations).unwrap(); - validate_special_register_policy(record, &source, declaration).unwrap(); - validate_special_register_llvm_exclusion(record, &declarations).unwrap(); - } - - for (id, section, anchor) in [ - ("clock", "10.23", "special-registers-clock-clock-hi"), - ("clock64", "10.24", "special-registers-clock64"), - ("globaltimer", "10.28", "special-registers-globaltimer"), - ("envreg1", "10.27", "special-registers-envreg"), - ("envreg2", "10.27", "special-registers-envreg"), - ("smid", "10.8", "special-registers-smid"), - ("nsmid", "10.9", "special-registers-nsmid"), - ("gridid", "10.10", "special-registers-gridid"), - ("warpid", "10.4", "special-registers-warpid"), - ("nwarpid", "10.5", "special-registers-nwarpid"), - ( - "dynamic_smem_size", - "10.32", - "special-registers-dynamic-smem-size", - ), - ( - "total_smem_size", - "10.30", - "special-registers-total-smem-size", - ), - ] { - let record = records.iter().find(|record| record.id == id).unwrap(); - assert!(record.ptx_isa_section.starts_with(section)); - assert!(record.ptx_isa_url.ends_with(anchor)); - } - - let gridid = records.iter().find(|record| record.id == "gridid").unwrap(); - assert_eq!(gridid.rust_result, "u64"); - assert_eq!(gridid.dialect_results, ["i64"]); - assert!(matches!( - resolve_policy_source(gridid).unwrap(), - IntrinsicSource::PtxNative { .. } - )); - assert!( - gridid - .special_register - .as_ref() - .unwrap() - .llvm_exclusion - .is_some() - ); - - let clock = records.iter().find(|record| record.id == "clock").unwrap(); - let source = resolve_policy_source(clock).unwrap(); - let declaration = resolve_imported_declaration(clock, &source, &declarations) - .unwrap() - .unwrap(); - - let mut wrong_effects = (*clock).clone(); - wrong_effects.memory = "none".into(); - assert!( - validate_special_register_policy(&wrong_effects, &source, Some(declaration)).is_err() - ); - - let mut wrong_contract = (*clock).clone(); - wrong_contract - .special_register - .as_mut() - .unwrap() - .output_constraint = SpecialRegisterOutputConstraint::Register64; - assert!( - validate_special_register_policy(&wrong_contract, &source, Some(declaration)).is_err() - ); - - let mut wrong_route = (*clock).clone(); - wrong_route.backend_lowerings[0].mechanism = BackendLoweringMechanism::InlinePtx; - assert!( - validate_special_register_policy(&wrong_route, &source, Some(declaration)).is_err() - ); - } - - #[test] - fn special_register_evidence_validates_both_backend_routes() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let policies = - expand_special_register_admission(&test_special_register_admission()).unwrap(); - let mut evidence_files = vec![ - read_evidence_file( - &repo_root - .join("intrinsics/evidence/rust-llvm-22.1.2-1cb4e383-special-registers.json"), - ) - .unwrap(), - read_evidence_file( - &repo_root - .join("intrinsics/evidence/cuda-13.3-libnvvm-13.3.33-special-registers.json"), - ) - .unwrap(), - ]; - let llvm_revision = "1cb4e3833c1919c2e6fb579a23ac0e2b22587b7e"; - let indexed = index_evidence(&evidence_files, llvm_revision).unwrap(); - for policy in &policies { - for lowering in &policy.backend_lowerings { - let evidence = indexed - .get(&(lowering.evidence_profile.as_str(), policy.id.as_str())) - .unwrap(); - validate_evidence(policy, evidence, Some(lowering)).unwrap(); - } - } - - let libnvvm = evidence_files - .iter_mut() - .find(|file| file.backend_kind == Some(IntrinsicBackend::LibNvvm)) - .unwrap(); - libnvvm - .records - .iter_mut() - .find(|record| record.id == "gridid") - .unwrap() - .stages - .retain(|stage| stage.stage != EvidenceStageKind::DeviceLink); - let indexed = index_evidence(&evidence_files, llvm_revision).unwrap(); - let policy = policies - .iter() - .find(|policy| policy.id == "gridid") - .unwrap(); - let lowering = policy - .backend_lowerings - .iter() - .find(|lowering| lowering.backend == IntrinsicBackend::LibNvvm) - .unwrap(); - let evidence = indexed - .get(&(lowering.evidence_profile.as_str(), policy.id.as_str())) - .unwrap(); - assert!(validate_evidence(policy, evidence, Some(lowering)).is_err()); - } - - #[test] - fn threadfence_admission_is_closed_and_uses_schema_34() { - let shard = |schema| { - toml::from_str::(&format!( - r#" -schema = {schema} -family = "sync" - -[threadfence] -llvm_evidence_profile = "llvm-test" -libnvvm_evidence_profile = "libnvvm-test" -runtime_validation = "unexecuted" - -[[threadfence.variant]] -abi_id = "i0298" -scope = "cta" - -[[threadfence.variant]] -abi_id = "i0299" -scope = "device" - -[[threadfence.variant]] -abi_id = "i0300" -scope = "system" -"# - )) - .unwrap() - }; - let path = Path::new("intrinsics/overlay/threadfence.toml"); - - let old = shard(THREADFENCE_SHARD_SCHEMA - 1); - assert!( - validate_overlay_shard_schema(&old, path) - .unwrap_err() - .to_string() - .contains("requires overlay shard schema 34") - ); - - let current = shard(THREADFENCE_SHARD_SCHEMA); - validate_overlay_shard_schema(¤t, path).unwrap(); - let admission = current.threadfence.unwrap(); - let records = expand_threadfence_admission(&admission).unwrap(); - assert_eq!(records.len(), 3); - assert_eq!( - records - .iter() - .map(|record| (record.abi_id.as_str(), record.id.as_str())) - .collect::>(), - [ - ("i0298", "threadfence_block"), - ("i0299", "threadfence"), - ("i0300", "threadfence_system"), - ] - ); - - let mut reordered = admission.clone(); - reordered.variants.swap(0, 1); - assert!(expand_threadfence_admission(&reordered).is_err()); - - let mut wrong_id = admission.clone(); - wrong_id.variants[0].abi_id = "i0300".into(); - assert!(expand_threadfence_admission(&wrong_id).is_err()); - - let mut executed = admission.clone(); - executed.runtime_validation = RuntimeValidation::Executed; - assert!(expand_threadfence_admission(&executed).is_err()); - } - - #[test] - fn pinned_threadfences_match_llvm_and_reject_contract_drift() { - let records = expand_threadfence_admission(&test_threadfence_admission()).unwrap(); - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - let declarations = imported - .intrinsics - .iter() - .map(|record| (record.source_record.as_str(), record)) - .collect::>(); - - for record in &records { - let declaration = declarations[record.source_record.as_deref().unwrap()]; - validate_imported_policy(record, declaration).unwrap(); - assert_eq!( - declaration - .selections - .iter() - .filter(|selection| selection_matches_policy(record, selection).unwrap()) - .count(), - 1 - ); - } - - let declaration = declarations["int_nvvm_membar_cta"]; - let mut changed = records[0].clone(); - changed.memory = "none".into(); - assert!(validate_imported_policy(&changed, declaration).is_err()); - - let mut changed = records[0].clone(); - changed.convergent = true; - assert!(validate_imported_policy(&changed, declaration).is_err()); - - let mut changed = records[0].clone(); - changed.backend_lowerings[0].mechanism = BackendLoweringMechanism::InlinePtx; - assert!(validate_imported_policy(&changed, declaration).is_err()); - - let mut changed = records[0].clone(); - changed.minimum_ptx = "2.0".into(); - assert!(validate_imported_policy(&changed, declaration).is_err()); - } - - #[test] - fn pinned_cluster_sregs_preserve_helpers_and_reject_unused_w_components() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let (mut overlay, _) = - read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")).unwrap(); - bind_pinned_abi_ids(&repo_root, &mut overlay); - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - let declarations = imported - .intrinsics - .iter() - .map(|record| (record.source_record.as_str(), record)) - .collect::>(); - let records = overlay - .intrinsics - .iter() - .filter(|record| { - record - .source_record - .as_deref() - .is_some_and(is_cluster_sreg_source) - }) - .collect::>(); - - assert_eq!(records.len(), 14); - let actual_abi_ids = records - .iter() - .map(|record| record.abi_id.clone()) - .collect::>(); - let expected_abi_ids = (263..=276) - .map(|id| format!("i{id:04}")) - .collect::>(); - assert_eq!(actual_abi_ids, expected_abi_ids); - for record in &records { - let declaration = declarations[record.source_record.as_deref().unwrap()]; - validate_imported_policy(record, declaration).unwrap(); - assert!(!declaration.source_record.ends_with("_w")); - } - - let compatibility_paths = records - .iter() - .flat_map(|record| record.compatibility_rust_paths.iter().map(String::as_str)) - .collect::>(); - assert_eq!( - compatibility_paths, - [ - "cuda_device::cluster::cluster_ctaidX", - "cuda_device::cluster::cluster_ctaidY", - "cuda_device::cluster::cluster_ctaidZ", - "cuda_device::cluster::cluster_nctaidX", - "cuda_device::cluster::cluster_nctaidY", - "cuda_device::cluster::cluster_nctaidZ", - "cuda_device::cluster::__cluster_grid_dimX", - "cuda_device::cluster::__cluster_grid_dimY", - "cuda_device::cluster::__cluster_grid_dimZ", - "cuda_device::cluster::__cluster_idxX", - "cuda_device::cluster::__cluster_idxY", - "cuda_device::cluster::__cluster_idxZ", - ] - .into_iter() - .collect() - ); - - let source = resolve_policy_source(records[0]).unwrap(); - let error = validate_sreg_policy( - records[0], - &source, - Some(declarations["int_nvvm_read_ptx_sreg_cluster_ctaid_w"]), - ) - .unwrap_err(); - assert!(error.to_string().contains("unused always-zero")); - - let mut mixed = (*records[0]).clone(); - mixed.sparse_mma = overlay - .intrinsics - .iter() - .find_map(|record| record.sparse_mma.clone()); - let source = resolve_policy_source(&mixed).unwrap(); - let error = validate_sreg_policy( - &mixed, - &source, - Some(declarations[mixed.source_record.as_deref().unwrap()]), - ) - .unwrap_err(); - assert!(error.to_string().contains("mixes another generated-family")); - } - - #[test] - fn compact_f8f6f4_axes_require_the_exact_canonical_matrix() { - let records = expand_sparse_mma_f8f6f4_admission(&test_f8f6f4_admission()).unwrap(); - assert_eq!(records.len(), 25); - assert_eq!( - records[0].id, - "mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m1_e2m1_f32" - ); - assert_eq!(records[0].rust_arguments[0], "[f32; 4]"); - assert_eq!(records[0].dialect_results, ["f32"; 4]); - assert_eq!(records[0].llvm_results, ["f32"; 4]); - - let mut missing = test_f8f6f4_admission(); - missing.a_elements.pop(); - assert!(expand_sparse_mma_f8f6f4_admission(&missing).is_err()); - - let mut duplicate = test_f8f6f4_admission(); - duplicate.a_elements[4] = SparseMmaElement::E4m3; - assert!(expand_sparse_mma_f8f6f4_admission(&duplicate).is_err()); - - let mut extra = test_f8f6f4_admission(); - extra.b_elements.push(SparseMmaElement::S4); - assert!(expand_sparse_mma_f8f6f4_admission(&extra).is_err()); - - let mut unsorted = test_f8f6f4_admission(); - unsorted.b_elements.swap(0, 1); - assert!(expand_sparse_mma_f8f6f4_admission(&unsorted).is_err()); - - let mut wrong_count = test_f8f6f4_admission(); - wrong_count.product_count = 24; - assert!(expand_sparse_mma_f8f6f4_admission(&wrong_count).is_err()); - } - - #[test] - fn compact_sparse_f8f6f4_f16_admission_is_closed_and_ordered() { - let admission = test_sparse_mma_f8f6f4_f16_admission(); - let records = expand_sparse_mma_f8f6f4_f16_admission(&admission).unwrap(); - assert_eq!(records.len(), 25); - assert!(records.iter().all(|record| record.abi_id.is_empty())); - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let mut bound = overlay_file(records.clone()); - bind_pinned_abi_ids(&repo_root, &mut bound); - assert_eq!( - bound - .intrinsics - .iter() - .map(|record| record.abi_id.clone()) - .collect::>(), - (525..=549) - .map(|id| format!("i{id:04}")) - .collect::>() - ); - assert_eq!( - records[0].id, - "mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e2m1_f16" - ); - assert_eq!( - records[24].id, - "mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e5m2_e5m2_f16" - ); - let first = &records[0]; - assert_eq!( - first.source_record.as_deref(), - Some("int_nvvm_mma_sp_ordered_metadata_m16n8k64_row_col_kind_f8f6f4_f16_e2m1_e2m1_f16") - ); - assert_eq!( - first.llvm_symbol.as_deref(), - Some( - "llvm.nvvm.mma.sp.ordered.metadata.m16n8k64.row.col.kind.f8f6f4.f16.e2m1.e2m1.f16" - ) - ); - assert_eq!( - first.rust_arguments, - ["[u32; 2]", "[u32; 4]", "[u32; 4]", "u32", "u32"] - ); - assert_eq!(first.rust_result, "[u32; 2]"); - assert_eq!(first.dialect_results, ["i32"; 2]); - assert_eq!(first.llvm_results, ["v2f16"; 2]); - assert_eq!( - first.llvm_arguments, - [ - "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", "v2f16", "v2f16", "i32", - "i32" - ] - ); - assert_eq!(first.minimum_ptx, "8.7"); - assert_eq!(first.minimum_sm, None); - assert_eq!(first.targets, SPARSE_MMA_F8F6F4_TARGETS); - assert!(first.convergent && !first.pure); - assert!(first.backend_lowerings.iter().all(|route| { - route.mechanism == BackendLoweringMechanism::InlinePtx - && route.minimum_ptx.as_deref() == Some("8.7") - && route.minimum_sm.is_none() - })); - let mma = first.sparse_mma.as_ref().unwrap(); - assert_eq!(mma.accumulator, SparseMmaAccumulator::F16); - assert_eq!(mma.selector, SparseMmaSelector::ImmediateZero); - assert_eq!( - mma.adapter, - SparseMmaAdapter::C2U32A4U32B4U32MetadataU32SelectorU32ToD2U32 - ); - assert_eq!( - mma.llvm_adapter, - SparseMmaLlvmAdapter::A4I32B4I32C2V2F16MetadataI32SelectorI32ToD2V2F16 - ); - assert_eq!( - first.compatibility_rust_paths, - ["cuda_device::wmma::mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e2m1_f16"] - ); - assert_eq!( - first.expected_ptx.operands, - [ - OperandPattern::RegisterList { length: 2 }, - OperandPattern::RegisterList { length: 4 }, - OperandPattern::RegisterList { length: 4 }, - OperandPattern::RegisterList { length: 2 }, - OperandPattern::Register, - OperandPattern::Immediate, - ] - ); - - let mut missing = admission.clone(); - missing.a_elements.pop(); - assert!(expand_sparse_mma_f8f6f4_f16_admission(&missing).is_err()); - - let mut reordered = admission.clone(); - reordered.b_elements.swap(0, 1); - assert!(expand_sparse_mma_f8f6f4_f16_admission(&reordered).is_err()); - - let mut wrong_count = admission.clone(); - wrong_count.product_count = 24; - assert!(expand_sparse_mma_f8f6f4_f16_admission(&wrong_count).is_err()); - - let mut legacy_first_id = admission.clone(); - legacy_first_id._legacy_first_abi_id = Some("i9999".into()); - assert!( - expand_sparse_mma_f8f6f4_f16_admission(&legacy_first_id) - .unwrap() - .iter() - .all(|record| record.abi_id.is_empty()) - ); - - let mut missing_evidence = admission.clone(); - missing_evidence.llvm_evidence_profile.clear(); - assert!(expand_sparse_mma_f8f6f4_f16_admission(&missing_evidence).is_err()); - - let mut executed = admission; - executed.runtime_validation = RuntimeValidation::Executed; - assert!(expand_sparse_mma_f8f6f4_f16_admission(&executed).is_err()); - } - - #[test] - fn sparse_f8f6f4_f16_candidate_floor_uses_the_resolved_policy() { - let policy = - expand_sparse_mma_f8f6f4_f16_admission(&test_sparse_mma_f8f6f4_f16_admission()) - .unwrap() - .remove(0); - let (_, requirement) = candidate_llvm_route(&policy).unwrap(); - - validate_candidate_target(&policy, &requirement, "sm_120a", "+ptx87").unwrap(); - validate_candidate_target(&policy, &requirement, "sm_120f", "+ptx88").unwrap(); - validate_candidate_target(&policy, &requirement, "sm_121a", "+ptx88").unwrap(); - validate_candidate_target(&policy, &requirement, "sm_121f", "+ptx88").unwrap(); - for target in ["sm_120f", "sm_121a", "sm_121f"] { - assert!( - validate_candidate_target(&policy, &requirement, target, "+ptx87").is_err(), - "{target} must require PTX 8.8" - ); - } - } - - #[test] - fn compact_dense_f8f6f4_admission_is_closed() { - for (accumulator, first, first_id, last_id, rust_arguments, rust_result, adapter) in [ - ( - RegisterMmaAccumulator::F32, - 454, - "mma_m16n8k32_f32_e2m1_e2m1", - "mma_m16n8k32_f32_e5m2_e5m2", - ["[f32; 4]", "[u32; 4]", "[u32; 2]"], - "[f32; 4]", - RegisterMmaAdapter::C4F32A4U32B2U32ToD4F32, - ), - ( - RegisterMmaAccumulator::F16, - 479, - "mma_m16n8k32_f16_e2m1_e2m1", - "mma_m16n8k32_f16_e5m2_e5m2", - ["[u32; 2]", "[u32; 4]", "[u32; 2]"], - "[u32; 2]", - RegisterMmaAdapter::C2U32A4U32B2U32ToD2U32, - ), - ] { - let admission = test_register_mma_f8f6f4_admission(accumulator); - let records = expand_register_mma_f8f6f4_admission(&admission, accumulator).unwrap(); - assert_eq!(records.len(), 25); - assert!(records.iter().all(|record| record.abi_id.is_empty())); - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let mut bound = overlay_file(records.clone()); - bind_pinned_abi_ids(&repo_root, &mut bound); - assert_eq!( - bound - .intrinsics - .iter() - .map(|record| record.abi_id.clone()) - .collect::>(), - (first..first + 25) - .map(|id| format!("i{id:04}")) - .collect::>() - ); - assert_eq!(records[0].id, first_id); - assert_eq!(records[24].id, last_id); - assert!(records.iter().all(|record| { - record.rust_arguments == rust_arguments - && record.rust_result == rust_result - && record.register_mma.as_ref().unwrap().adapter == adapter - && record.targets == REGISTER_MMA_F8F6F4_TARGETS - && record.minimum_ptx == "8.7" - && record.minimum_sm.is_none() - && record.backend_lowerings.len() == 2 - && record - .backend_lowerings - .iter() - .all(|route| route.mechanism == BackendLoweringMechanism::InlinePtx) - })); - - let mut missing = admission.clone(); - missing.a_elements.pop(); - assert!(expand_register_mma_f8f6f4_admission(&missing, accumulator).is_err()); - - let mut reordered = admission.clone(); - reordered.b_elements.swap(0, 1); - assert!(expand_register_mma_f8f6f4_admission(&reordered, accumulator).is_err()); - - let mut wrong_count = admission.clone(); - wrong_count.product_count = 24; - assert!(expand_register_mma_f8f6f4_admission(&wrong_count, accumulator).is_err()); - - let mut legacy_first_id = admission.clone(); - legacy_first_id._legacy_first_abi_id = Some("i9999".into()); - assert!( - expand_register_mma_f8f6f4_admission(&legacy_first_id, accumulator) - .unwrap() - .iter() - .all(|record| record.abi_id.is_empty()) - ); - - let mut reordered_targets = admission.clone(); - reordered_targets.targets.swap(0, 1); - assert!(expand_register_mma_f8f6f4_admission(&reordered_targets, accumulator).is_err()); - - let mut missing_evidence = admission.clone(); - missing_evidence.llvm_evidence_profile.clear(); - assert!(expand_register_mma_f8f6f4_admission(&missing_evidence, accumulator).is_err()); - - let mut executed = admission; - executed.runtime_validation = RuntimeValidation::Executed; - assert!(expand_register_mma_f8f6f4_admission(&executed, accumulator).is_err()); - } - - let admission = test_register_mma_f8f6f4_admission(RegisterMmaAccumulator::F32); - assert!( - expand_register_mma_f8f6f4_admission(&admission, RegisterMmaAccumulator::F64).is_err() - ); - } - - #[test] - fn compact_standard_fp8_admission_is_closed_and_ordered() { - let admission = test_register_mma_fp8_admission(); - let records = expand_register_mma_fp8_admission(&admission).unwrap(); - assert_eq!(records.len(), 16); - assert!(records.iter().all(|record| record.abi_id.is_empty())); - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let mut bound = overlay_file(records.clone()); - bind_pinned_abi_ids(&repo_root, &mut bound); - assert_eq!( - bound - .intrinsics - .iter() - .map(|record| record.abi_id.clone()) - .collect::>(), - (504..=519) - .map(|id| format!("i{id:04}")) - .collect::>() - ); - - let mut index = 0; - for &shape in ®ISTER_MMA_FP8_SHAPES { - let (shape_name, a_count, b_count) = register_mma_fp8_shape_contract(shape).unwrap(); - for &accumulator in ®ISTER_MMA_FP8_ACCUMULATORS { - let (scalar, arguments, result, adapter, result_count) = match accumulator { - RegisterMmaAccumulator::F16 => ( - "f16", - vec![ - "[u32; 2]".to_owned(), - format!("[u32; {a_count}]"), - if b_count == 1 { - "u32".to_owned() - } else { - "[u32; 2]".to_owned() - }, - ], - "[u32; 2]", - if shape == RegisterMmaShape::M16n8k16 { - RegisterMmaAdapter::C2U32A2U32B1U32ToD2U32 - } else { - RegisterMmaAdapter::C2U32A4U32B2U32ToD2U32 - }, - 2, - ), - RegisterMmaAccumulator::F32 => ( - "f32", - vec![ - "[f32; 4]".to_owned(), - format!("[u32; {a_count}]"), - if b_count == 1 { - "u32".to_owned() - } else { - "[u32; 2]".to_owned() - }, - ], - "[f32; 4]", - if shape == RegisterMmaShape::M16n8k16 { - RegisterMmaAdapter::C4F32A2U32B1U32ToD4F32 - } else { - RegisterMmaAdapter::C4F32A4U32B2U32ToD4F32 - }, - 4, - ), - _ => unreachable!(), - }; - for &a_element in ®ISTER_MMA_FP8_ELEMENTS { - for &b_element in ®ISTER_MMA_FP8_ELEMENTS { - let record = &records[index]; - let a = register_mma_fp8_element_name(a_element).unwrap(); - let b = register_mma_fp8_element_name(b_element).unwrap(); - assert_eq!(record.id, format!("mma_{shape_name}_fp8_{scalar}_{a}_{b}")); - assert_eq!( - record.operation_key, - format!( - "matrix.mma.{shape_name}.row.col.standard_fp8.{scalar}.{a}.{b}.{scalar}" - ) - ); - assert_eq!( - record.source_record.as_deref(), - Some( - format!( - "int_nvvm_mma_{shape_name}_row_col_{scalar}_{a}_{b}_{scalar}" - ) - .as_str() - ) - ); - assert_eq!(record.rust_arguments, arguments); - assert_eq!(record.rust_result, result); - assert_eq!( - record.minimum_ptx, - register_mma_fp8_minimum_ptx(shape, accumulator) - ); - assert_eq!(record.minimum_sm.as_deref(), Some("sm_89")); - assert_eq!( - record.expected_ptx.operands, - [result_count, a_count, b_count, result_count] - .map(|length| OperandPattern::RegisterList { length }) - ); - let mma = record.register_mma.as_ref().unwrap(); - assert_eq!(mma.kind, Some(RegisterMmaKind::Standard)); - assert_eq!(mma.adapter, adapter); - index += 1; - } - } - } - } - assert_eq!(index, records.len()); - - let mut reordered = admission.clone(); - reordered.shapes.swap(0, 1); - assert!(expand_register_mma_fp8_admission(&reordered).is_err()); - let mut reordered = admission.clone(); - reordered.accumulators.swap(0, 1); - assert!(expand_register_mma_fp8_admission(&reordered).is_err()); - let mut reordered = admission.clone(); - reordered.a_elements.swap(0, 1); - assert!(expand_register_mma_fp8_admission(&reordered).is_err()); - let mut reordered = admission.clone(); - reordered.b_elements.swap(0, 1); - assert!(expand_register_mma_fp8_admission(&reordered).is_err()); - let mut wrong = admission.clone(); - wrong.product_count = 15; - assert!(expand_register_mma_fp8_admission(&wrong).is_err()); - let mut legacy_first_id = admission.clone(); - legacy_first_id._legacy_first_abi_id = Some("i9999".into()); - assert!( - expand_register_mma_fp8_admission(&legacy_first_id) - .unwrap() - .iter() - .all(|record| record.abi_id.is_empty()) - ); - let mut wrong = admission; - wrong.runtime_validation = RuntimeValidation::Executed; - assert!(expand_register_mma_fp8_admission(&wrong).is_err()); - } - - #[test] - fn compact_ampere_float_mma_admission_is_closed_and_ordered() { - let admission = test_register_mma_ampere_float_admission(); - let records = expand_register_mma_ampere_float_admission(&admission).unwrap(); - assert_eq!(records.len(), 5); - assert!(records.iter().all(|record| record.abi_id.is_empty())); - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let mut bound = overlay_file(records.clone()); - bind_pinned_abi_ids(&repo_root, &mut bound); - assert_eq!( - bound - .intrinsics - .iter() - .map(|record| (record.abi_id.as_str(), record.id.as_str())) - .collect::>(), - [ - ("i0520", "mma_m16n8k4_f32_tf32"), - ("i0521", "mma_m16n8k8_f16_f16"), - ("i0522", "mma_m16n8k8_f32_bf16"), - ("i0523", "mma_m16n8k8_f32_f16"), - ("i0524", "mma_m16n8k16_f16_f16"), - ] - ); - for record in &records { - assert_eq!(record.minimum_ptx, "7.0"); - assert_eq!(record.minimum_sm.as_deref(), Some("sm_80")); - assert_eq!(record.backend_lowerings.len(), 2); - assert!(record.backend_lowerings.iter().all(|lowering| { - lowering.mechanism == BackendLoweringMechanism::InlinePtx - && lowering.minimum_ptx.as_deref() == Some("7.0") - && lowering.minimum_sm.as_deref() == Some("sm_80") - })); - let mma = record.register_mma.as_ref().unwrap(); - assert_eq!(mma.kind, None); - assert_eq!( - mma.compatibility_source, - RegisterMmaCompatibilitySource::GeneratedStub - ); - } - - let mut reordered = admission.clone(); - reordered.variants.swap(0, 1); - assert!(expand_register_mma_ampere_float_admission(&reordered).is_err()); - let mut legacy_first_id = admission.clone(); - legacy_first_id._legacy_first_abi_id = Some("i9999".into()); - assert!( - expand_register_mma_ampere_float_admission(&legacy_first_id) - .unwrap() - .iter() - .all(|record| record.abi_id.is_empty()) - ); - let mut wrong = admission.clone(); - wrong.product_count = 4; - assert!(expand_register_mma_ampere_float_admission(&wrong).is_err()); - let mut wrong = admission.clone(); - wrong.llvm_evidence_profile.clear(); - assert!(expand_register_mma_ampere_float_admission(&wrong).is_err()); - let mut wrong = admission; - wrong.runtime_validation = RuntimeValidation::Executed; - assert!(expand_register_mma_ampere_float_admission(&wrong).is_err()); - - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let path = repo_root.join("intrinsics/overlay/register_mma_ampere_float.toml"); - let bytes = fs::read(&path).unwrap(); - let mut shard: OverlayShardFile = toml::from_slice(&bytes).unwrap(); - validate_overlay_shard_schema(&shard, &path).unwrap(); - shard.schema = REGISTER_MMA_AMPERE_FLOAT_SHARD_SCHEMA - 1; - assert!( - validate_overlay_shard_schema(&shard, &path) - .unwrap_err() - .to_string() - .contains("requires overlay shard schema 49") - ); - } - - #[test] - fn ampere_float_mma_policies_match_llvm_and_fail_closed() { - let records = - expand_register_mma_ampere_float_admission(&test_register_mma_ampere_float_admission()) - .unwrap(); - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let mut overlay = overlay_file(records); - bind_pinned_abi_ids(&repo_root, &mut overlay); - let records = overlay.intrinsics; - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - let declarations = imported - .intrinsics - .iter() - .map(|record| (record.source_record.as_str(), record)) - .collect::>(); - - for record in &records { - let declaration = declarations[record.source_record.as_deref().unwrap()]; - validate_imported_policy(record, declaration).unwrap(); - let expected_predicates = match record.id.as_str() { - "mma_m16n8k4_f32_tf32" | "mma_m16n8k16_f16_f16" => [ - "Subtarget->getSmVersion() >= 80", - "Subtarget->getPTXVersion() >= 70", - ], - _ => [ - "Subtarget->getPTXVersion() >= 65", - "Subtarget->getSmVersion() >= 75", - ], - }; - assert_eq!(declaration.selections[0].predicates, expected_predicates); - assert_eq!(record.minimum_ptx, "7.0"); - assert_eq!(record.minimum_sm.as_deref(), Some("sm_80")); - } - - let valid = &records[2]; - assert_eq!(valid.id, "mma_m16n8k8_f32_bf16"); - assert_eq!( - declarations[valid.source_record.as_deref().unwrap()].selections[0].predicates, - [ - "Subtarget->getPTXVersion() >= 65", - "Subtarget->getSmVersion() >= 75", - ] - ); - assert_eq!(valid.minimum_ptx, "7.0"); - assert_eq!(valid.minimum_sm.as_deref(), Some("sm_80")); - assert!(valid.backend_lowerings.iter().all(|lowering| { - lowering.minimum_ptx.as_deref() == Some("7.0") - && lowering.minimum_sm.as_deref() == Some("sm_80") - })); - let declaration = declarations[valid.source_record.as_deref().unwrap()]; - let mut wrong = valid.clone(); - wrong.minimum_ptx = "6.5".into(); - assert!(validate_imported_policy(&wrong, declaration).is_err()); - let mut wrong = valid.clone(); - wrong.register_mma.as_mut().unwrap().adapter = RegisterMmaAdapter::C4F32A4U32B2U32ToD4F32; - assert!(validate_imported_policy(&wrong, declaration).is_err()); - let mut wrong = valid.clone(); - wrong.register_mma.as_mut().unwrap().kind = Some(RegisterMmaKind::Standard); - assert!(validate_imported_policy(&wrong, declaration).is_err()); - let mut wrong_declaration = declaration.clone(); - wrong_declaration.selections[0].predicates[0] = "Subtarget->getPTXVersion() >= 70".into(); - assert!(validate_imported_policy(valid, &wrong_declaration).is_err()); - } - - #[test] - fn standard_fp8_policies_match_llvm_and_fail_closed() { - let records = - expand_register_mma_fp8_admission(&test_register_mma_fp8_admission()).unwrap(); - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let mut overlay = overlay_file(records); - bind_pinned_abi_ids(&repo_root, &mut overlay); - let records = overlay.intrinsics; - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - let declarations = imported - .intrinsics - .iter() - .map(|record| (record.source_record.as_str(), record)) - .collect::>(); - - for record in &records { - let declaration = declarations[record.source_record.as_deref().unwrap()]; - validate_imported_policy(record, declaration).unwrap(); - let mma = record.register_mma.as_ref().unwrap(); - let source_floor = declaration.selections[0] - .predicates - .iter() - .find(|predicate| predicate.contains("getPTXVersion")) - .unwrap(); - if mma.shape == RegisterMmaShape::M16n8k32 { - assert!(source_floor.ends_with(">= 84")); - if mma.accumulator == RegisterMmaAccumulator::F16 { - assert_eq!(record.minimum_ptx, "8.7"); - } else { - assert_eq!(record.minimum_ptx, "8.4"); - } - } else { - assert!(source_floor.ends_with(">= 87")); - assert_eq!(record.minimum_ptx, "8.7"); - } - } - - let valid = &records[0]; - let declaration = declarations[valid.source_record.as_deref().unwrap()]; - let mut wrong = valid.clone(); - wrong.register_mma.as_mut().unwrap().kind = None; - assert!(validate_imported_policy(&wrong, declaration).is_err()); - let mut wrong = valid.clone(); - wrong.register_mma.as_mut().unwrap().adapter = RegisterMmaAdapter::C2U32A4U32B2U32ToD2U32; - assert!(validate_imported_policy(&wrong, declaration).is_err()); - let mut wrong_declaration = declaration.clone(); - wrong_declaration.selections[0].predicates[1] = "Subtarget->getPTXVersion() >= 84".into(); - assert!(validate_imported_policy(valid, &wrong_declaration).is_err()); - let mut wrong = valid.clone(); - wrong.minimum_ptx = "8.4".into(); - assert!(validate_imported_policy(&wrong, declaration).is_err()); - - let mut old = expand_register_mma_f8f6f4_admission( - &test_register_mma_f8f6f4_admission(RegisterMmaAccumulator::F32), - RegisterMmaAccumulator::F32, - ) - .unwrap() - .remove(18); - let mut standard = records[12].clone(); - assert_ne!(old.id, standard.id); - standard.id = old.id.clone(); - assert!(validate_unique_overlay(&[old.clone(), standard], 1).is_err()); - old.register_mma.as_mut().unwrap().kind = Some(RegisterMmaKind::F8f6f4); - assert_eq!( - old.register_mma.as_ref().unwrap().kind, - Some(RegisterMmaKind::F8f6f4) - ); - } - - #[test] - fn standard_fp8_resolves_exact_routes_and_target_floors() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let catalog = resolve(&repo_root).unwrap(); - let records = catalog - .intrinsics - .iter() - .filter(|record| { - record.rust.abi_id.as_str() >= "i0504" && record.rust.abi_id.as_str() <= "i0519" - }) - .collect::>(); - let expected = [ - ( - "mma_m16n8k16_fp8_f16_e4m3_e4m3", - "i0504", - RegisterMmaShape::M16n8k16, - RegisterMmaAccumulator::F16, - 87, - ), - ( - "mma_m16n8k16_fp8_f16_e4m3_e5m2", - "i0505", - RegisterMmaShape::M16n8k16, - RegisterMmaAccumulator::F16, - 87, - ), - ( - "mma_m16n8k16_fp8_f16_e5m2_e4m3", - "i0506", - RegisterMmaShape::M16n8k16, - RegisterMmaAccumulator::F16, - 87, - ), - ( - "mma_m16n8k16_fp8_f16_e5m2_e5m2", - "i0507", - RegisterMmaShape::M16n8k16, - RegisterMmaAccumulator::F16, - 87, - ), - ( - "mma_m16n8k16_fp8_f32_e4m3_e4m3", - "i0508", - RegisterMmaShape::M16n8k16, - RegisterMmaAccumulator::F32, - 87, - ), - ( - "mma_m16n8k16_fp8_f32_e4m3_e5m2", - "i0509", - RegisterMmaShape::M16n8k16, - RegisterMmaAccumulator::F32, - 87, - ), - ( - "mma_m16n8k16_fp8_f32_e5m2_e4m3", - "i0510", - RegisterMmaShape::M16n8k16, - RegisterMmaAccumulator::F32, - 87, - ), - ( - "mma_m16n8k16_fp8_f32_e5m2_e5m2", - "i0511", - RegisterMmaShape::M16n8k16, - RegisterMmaAccumulator::F32, - 87, - ), - ( - "mma_m16n8k32_fp8_f16_e4m3_e4m3", - "i0512", - RegisterMmaShape::M16n8k32, - RegisterMmaAccumulator::F16, - 87, - ), - ( - "mma_m16n8k32_fp8_f16_e4m3_e5m2", - "i0513", - RegisterMmaShape::M16n8k32, - RegisterMmaAccumulator::F16, - 87, - ), - ( - "mma_m16n8k32_fp8_f16_e5m2_e4m3", - "i0514", - RegisterMmaShape::M16n8k32, - RegisterMmaAccumulator::F16, - 87, - ), - ( - "mma_m16n8k32_fp8_f16_e5m2_e5m2", - "i0515", - RegisterMmaShape::M16n8k32, - RegisterMmaAccumulator::F16, - 87, - ), - ( - "mma_m16n8k32_fp8_f32_e4m3_e4m3", - "i0516", - RegisterMmaShape::M16n8k32, - RegisterMmaAccumulator::F32, - 84, - ), - ( - "mma_m16n8k32_fp8_f32_e4m3_e5m2", - "i0517", - RegisterMmaShape::M16n8k32, - RegisterMmaAccumulator::F32, - 84, - ), - ( - "mma_m16n8k32_fp8_f32_e5m2_e4m3", - "i0518", - RegisterMmaShape::M16n8k32, - RegisterMmaAccumulator::F32, - 84, - ), - ( - "mma_m16n8k32_fp8_f32_e5m2_e5m2", - "i0519", - RegisterMmaShape::M16n8k32, - RegisterMmaAccumulator::F32, - 84, - ), - ]; - assert_eq!(records.len(), expected.len()); - let sm_89 = CatalogHardwareTarget::AnyOf { - alternatives: vec![CatalogHardwareAlternative::MinimumSm { sm: 89 }], - }; - let expected_routes = [ - ( - IntrinsicBackend::LlvmNvptx, - BackendLoweringMechanism::InlinePtx, - "rust-llvm-22.1.2-1cb4e383", - ), - ( - IntrinsicBackend::LibNvvm, - BackendLoweringMechanism::InlinePtx, - "cuda-13.3-libnvvm-13.3.33", - ), - ]; - let mut floor_groups = BTreeMap::new(); - - for (record, &(id, abi_id, shape, accumulator, minimum_ptx)) in - records.iter().zip(&expected) - { - assert_eq!(record.id, id); - assert_eq!(record.rust.abi_id, abi_id); - let mma = record.register_mma.as_ref().unwrap(); - assert_eq!((mma.shape, mma.accumulator), (shape, accumulator)); - assert_eq!(record.target.minimum_ptx.encoded(), minimum_ptx); - assert_eq!(record.target.hardware, sm_89); - assert_eq!( - record - .backend_lowerings - .iter() - .map(|route| ( - route.backend, - route.mechanism, - route.evidence_profile.as_str() - )) - .collect::>(), - expected_routes - ); - for route in &record.backend_lowerings { - assert_eq!(route.target.minimum_ptx.encoded(), minimum_ptx); - assert_eq!(route.target.hardware, sm_89); - } - *floor_groups - .entry((shape, accumulator, minimum_ptx)) - .or_insert(0) += 1; - - if shape == RegisterMmaShape::M16n8k32 && accumulator == RegisterMmaAccumulator::F16 { - assert_eq!( - record - .selections - .iter() - .flat_map(|selection| &selection.predicates) - .filter(|predicate| predicate.contains("getPTXVersion")) - .map(String::as_str) - .collect::>(), - ["Subtarget->getPTXVersion() >= 84"] - ); - let llvm_route = &record.backend_lowerings[0]; - assert!(llvm_route.stages.iter().any(|stage| { - stage.stage == EvidenceStageKind::PtxAssembly - && stage.mechanism == Some(BackendLoweringMechanism::InlinePtx) - && stage.outcome == "succeeded" - && stage.targets == ["sm_89", "ptx87"] - && stage.artifact_kind == Some(EvidenceArtifactKind::Cubin) - })); - assert!(llvm_route.stages.iter().any(|stage| { - stage.stage == EvidenceStageKind::PtxAssembly - && stage.mechanism == Some(BackendLoweringMechanism::InlinePtx) - && stage.outcome == "failed" - && stage.targets == ["sm_89", "ptx86"] - })); - } - } - assert_eq!( - floor_groups, - BTreeMap::from([ - ( - (RegisterMmaShape::M16n8k16, RegisterMmaAccumulator::F16, 87), - 4 - ), - ( - (RegisterMmaShape::M16n8k16, RegisterMmaAccumulator::F32, 87), - 4 - ), - ( - (RegisterMmaShape::M16n8k32, RegisterMmaAccumulator::F16, 87), - 4 - ), - ( - (RegisterMmaShape::M16n8k32, RegisterMmaAccumulator::F32, 84), - 4 - ), - ]) - ); - } - - #[test] - fn existing_catalog_intrinsics_keep_their_json_shape_without_kind() { - fn abi_number(abi_id: &str) -> u16 { - validate_abi_id(abi_id).unwrap(); - abi_id[1..].parse().unwrap() - } - - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let original: serde_json::Value = - read_json(&repo_root.join("intrinsics/catalog.json")).unwrap(); - let catalog: CatalogFile = read_json(&repo_root.join("intrinsics/catalog.json")).unwrap(); - let serialized = serde_json::to_value(&catalog).unwrap(); - let original_by_abi = original["intrinsics"] - .as_array() - .unwrap() - .iter() - .map(|record| { - let abi_id = record["rust"]["abi_id"].as_str().unwrap(); - (abi_number(abi_id), record) - }) - .collect::>(); - let serialized_by_abi = serialized["intrinsics"] - .as_array() - .unwrap() - .iter() - .map(|record| { - let abi_id = record["rust"]["abi_id"].as_str().unwrap(); - (abi_number(abi_id), record) - }) - .collect::>(); - for abi in 1..=503 { - assert_eq!(original_by_abi.get(&abi), serialized_by_abi.get(&abi)); - } - - let existing = catalog - .intrinsics - .iter() - .filter(|record| abi_number(&record.rust.abi_id) <= 503) - .collect::>(); - assert_eq!(existing.len(), 503); - assert!(existing.iter().all(|record| { - record - .register_mma - .as_ref() - .is_none_or(|mma| mma.kind.is_none()) - })); - - let standard_fp8 = catalog - .intrinsics - .iter() - .filter(|record| (504..=519).contains(&abi_number(&record.rust.abi_id))) - .collect::>(); - assert!(standard_fp8.is_empty() || standard_fp8.len() == 16); - if standard_fp8.len() == 16 { - assert!(standard_fp8.iter().all(|record| { - record - .register_mma - .as_ref() - .is_some_and(|mma| mma.kind == Some(RegisterMmaKind::Standard)) - })); - for abi in 504..=519 { - assert_eq!(original_by_abi[&abi]["register_mma"]["kind"], "standard"); - } - } - } - - #[test] - fn dense_f8f6f4_candidate_floor_uses_the_resolved_policy() { - for accumulator in [RegisterMmaAccumulator::F16, RegisterMmaAccumulator::F32] { - let policy = expand_register_mma_f8f6f4_admission( - &test_register_mma_f8f6f4_admission(accumulator), - accumulator, - ) - .unwrap() - .remove(0); - let (_, requirement) = candidate_llvm_route(&policy).unwrap(); - - validate_candidate_target(&policy, &requirement, "sm_120a", "+ptx87").unwrap(); - validate_candidate_target(&policy, &requirement, "sm_120f", "+ptx88").unwrap(); - validate_candidate_target(&policy, &requirement, "sm_121a", "+ptx88").unwrap(); - validate_candidate_target(&policy, &requirement, "sm_121f", "+ptx88").unwrap(); - for target in ["sm_120f", "sm_121a", "sm_121f"] { - assert!( - validate_candidate_target(&policy, &requirement, target, "+ptx87").is_err(), - "{target} must require PTX 8.8" - ); - } - } - } - - #[test] - fn compact_prmt_admission_requires_every_mode_and_reserved_abi_id() { - assert_eq!( - expand_prmt_admission(&test_prmt_admission()).unwrap().len(), - 7 - ); - - let mut missing = test_prmt_admission(); - missing.variants.pop(); - assert!(expand_prmt_admission(&missing).is_err()); - - let mut duplicate = test_prmt_admission(); - duplicate.variants[6].mode = PrmtMode::Rc8; - assert!(expand_prmt_admission(&duplicate).is_err()); - - let mut wrong_abi = test_prmt_admission(); - wrong_abi.variants[0].abi_id = "i9999".into(); - assert!(expand_prmt_admission(&wrong_abi).is_err()); - } - - #[test] - fn compact_fp8_conversion_axes_require_the_exact_closed_product() { - let records = - expand_packed_conversion_fp8_admission(&test_fp8_conversion_admission()).unwrap(); - assert_eq!(records.len(), 4); - assert_eq!(records[0].id, "cvt_rn_satfinite_e4m3x2_f32"); - assert_eq!(records[1].id, "cvt_rn_satfinite_relu_e4m3x2_f32"); - assert_eq!(records[2].id, "cvt_rn_satfinite_e5m2x2_f32"); - assert_eq!(records[3].id, "cvt_rn_satfinite_relu_e5m2x2_f32"); - assert!(records.iter().all(|record| { - record.rust_result == "u16" - && record.dialect_results == ["i16"] - && record.llvm_results == ["i16"] - && record.minimum_ptx == "8.1" - && record.minimum_sm.as_deref() == Some("sm_89") - && record.pure - && !record.convergent - })); - - let mut missing_format = test_fp8_conversion_admission(); - missing_format.destination_formats.pop(); - assert!(expand_packed_conversion_fp8_admission(&missing_format).is_err()); - - let mut reversed_formats = test_fp8_conversion_admission(); - reversed_formats.destination_formats.reverse(); - assert!(expand_packed_conversion_fp8_admission(&reversed_formats).is_err()); - - let mut missing_saturation = test_fp8_conversion_admission(); - missing_saturation.saturations.pop(); - assert!(expand_packed_conversion_fp8_admission(&missing_saturation).is_err()); - - let mut reversed_saturations = test_fp8_conversion_admission(); - reversed_saturations.saturations.reverse(); - assert!(expand_packed_conversion_fp8_admission(&reversed_saturations).is_err()); - - let mut wrong_count = test_fp8_conversion_admission(); - wrong_count.product_count = 3; - assert!(expand_packed_conversion_fp8_admission(&wrong_count).is_err()); - - let mut executed = test_fp8_conversion_admission(); - executed.runtime_validation = RuntimeValidation::Executed; - assert!(expand_packed_conversion_fp8_admission(&executed).is_err()); - } - - #[test] - fn compact_fp8_f16x2_conversion_axes_require_the_exact_closed_product() { - let records = - expand_packed_conversion_fp8_f16x2_admission(&test_fp8_f16x2_conversion_admission()) - .unwrap(); - assert_eq!(records.len(), 8); - assert_eq!(records[0].id, "cvt_rn_satfinite_e4m3x2_f16x2"); - assert_eq!(records[1].id, "cvt_rn_satfinite_relu_e4m3x2_f16x2"); - assert_eq!(records[2].id, "cvt_rn_f16x2_e4m3x2"); - assert_eq!(records[3].id, "cvt_rn_relu_f16x2_e4m3x2"); - assert_eq!(records[4].id, "cvt_rn_satfinite_e5m2x2_f16x2"); - assert_eq!(records[5].id, "cvt_rn_satfinite_relu_e5m2x2_f16x2"); - assert_eq!(records[6].id, "cvt_rn_f16x2_e5m2x2"); - assert_eq!(records[7].id, "cvt_rn_relu_f16x2_e5m2x2"); - - // Every form takes one packed operand and carries the Ada floor, - // including the unpacks whose destination is the older f16x2. - assert!(records.iter().all(|record| { - record.rust_arguments.len() == 1 - && record.dialect_operands.len() == 1 - && record.llvm_arguments.len() == 1 - && record.expected_ptx.operands.len() == 2 - && record.minimum_ptx == "8.1" - && record.minimum_sm.as_deref() == Some("sm_89") - && record.pure - && !record.convergent - })); - - // Packing narrows f16x2 to a 16-bit FP8 pair; unpacking widens back. - assert!( - records - .iter() - .filter(|record| record.id.ends_with("_f16x2")) - .all(|record| { - record.rust_arguments == ["u32"] - && record.rust_result == "u16" - && record.llvm_arguments == ["v2f16"] - && record.llvm_results == ["i16"] - }) - ); - assert!( - records - .iter() - .filter(|record| !record.id.ends_with("_f16x2")) - .all(|record| { - record.rust_arguments == ["u16"] - && record.rust_result == "u32" - && record.llvm_arguments == ["i16"] - && record.llvm_results == ["v2f16"] - }) - ); - - let mut missing_format = test_fp8_f16x2_conversion_admission(); - missing_format.fp8_formats.pop(); - assert!(expand_packed_conversion_fp8_f16x2_admission(&missing_format).is_err()); - - let mut reversed_formats = test_fp8_f16x2_conversion_admission(); - reversed_formats.fp8_formats.reverse(); - assert!(expand_packed_conversion_fp8_f16x2_admission(&reversed_formats).is_err()); - - let mut missing_direction = test_fp8_f16x2_conversion_admission(); - missing_direction.directions.pop(); - assert!(expand_packed_conversion_fp8_f16x2_admission(&missing_direction).is_err()); - - let mut reversed_directions = test_fp8_f16x2_conversion_admission(); - reversed_directions.directions.reverse(); - assert!(expand_packed_conversion_fp8_f16x2_admission(&reversed_directions).is_err()); - - let mut without_relu = test_fp8_f16x2_conversion_admission(); - without_relu.relu_variants = false; - assert!(expand_packed_conversion_fp8_f16x2_admission(&without_relu).is_err()); - - let mut wrong_count = test_fp8_f16x2_conversion_admission(); - wrong_count.product_count = 4; - assert!(expand_packed_conversion_fp8_f16x2_admission(&wrong_count).is_err()); - - let mut executed = test_fp8_f16x2_conversion_admission(); - executed.runtime_validation = RuntimeValidation::Executed; - assert!(expand_packed_conversion_fp8_f16x2_admission(&executed).is_err()); - } - - #[test] - fn compact_debug_control_admission_is_closed() { - let records = expand_debug_control_admission(&test_debug_control_admission()).unwrap(); - assert_eq!(records.len(), 3); - assert_eq!(records[0].id, "trap"); - assert_eq!(records[1].id, "breakpoint"); - assert_eq!(records[2].id, "pmevent"); - for record in &records { - validate_ptx_native_policy(record).unwrap(); - assert_eq!(record.backend_lowerings.len(), 2); - assert!( - record - .debug_control - .as_ref() - .is_some_and(|debug| debug.runtime_validation == RuntimeValidation::Unexecuted) - ); - } - assert_eq!(records[0].minimum_ptx, "1.0"); - assert_eq!(records[0].minimum_sm, None); - assert_eq!(records[1].minimum_ptx, "1.0"); - assert_eq!(records[1].minimum_sm.as_deref(), Some("sm_11")); - assert_eq!(records[2].minimum_ptx, "1.4"); - assert_eq!( - records[2].expected_ptx.operands, - [OperandPattern::Immediate] - ); - - let mut pending = test_debug_control_admission(); - pending.abi_ids.clear(); - assert!(expand_debug_control_admission(&pending).is_err()); - - let mut missing = test_debug_control_admission(); - missing.operations.pop(); - assert!(expand_debug_control_admission(&missing).is_err()); - - let mut duplicate_operation = test_debug_control_admission(); - duplicate_operation.operations[2] = DebugControlOperation::Breakpoint; - assert!(expand_debug_control_admission(&duplicate_operation).is_err()); - - let mut duplicate_id = test_debug_control_admission(); - duplicate_id.abi_ids[2] = duplicate_id.abi_ids[1].clone(); - assert!(expand_debug_control_admission(&duplicate_id).is_err()); - - let mut malformed_id = test_debug_control_admission(); - malformed_id.abi_ids[0] = "debug1".into(); - assert!(expand_debug_control_admission(&malformed_id).is_err()); - - let mut executed = test_debug_control_admission(); - executed.runtime_validation = RuntimeValidation::Executed; - assert!(expand_debug_control_admission(&executed).is_err()); - - let mut wrong_source = records[0].clone(); - wrong_source.source = Some(IntrinsicSource::LlvmImported { - source_record: "invented".into(), - }); - assert!(validate_ptx_native_policy(&wrong_source).is_err()); - - let mut wrong_adapter = records[2].clone(); - wrong_adapter.debug_control.as_mut().unwrap().adapter = DebugControlAdapter::Direct; - assert!(validate_ptx_native_policy(&wrong_adapter).is_err()); - - let mut wrong_immediate = records[2].clone(); - wrong_immediate.expected_ptx.operands = vec![OperandPattern::Register]; - assert!(validate_ptx_native_policy(&wrong_immediate).is_err()); - - let mut wrong_floor = records[1].clone(); - wrong_floor.backend_lowerings[0].minimum_sm = Some("sm_75".into()); - assert!(validate_ptx_native_policy(&wrong_floor).is_err()); - } - - #[test] - fn compact_clc_admission_matches_llvm_and_fails_closed() { - let records = expand_clc_admission(&test_clc_admission()).unwrap(); - assert_eq!(records.len(), 6); - assert_eq!( - records - .iter() - .map(|record| (record.abi_id.as_str(), record.id.as_str())) - .collect::>(), - [ - ("i0322", "clc_try_cancel"), - ("i0323", "clc_try_cancel_multicast"), - ("i0324", "clc_query_is_canceled"), - ("i0325", "clc_query_get_first_ctaid_x"), - ("i0326", "clc_query_get_first_ctaid_y"), - ("i0327", "clc_query_get_first_ctaid_z"), - ] - ); - - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - let declarations = imported - .intrinsics - .iter() - .map(|record| (record.source_record.as_str(), record)) - .collect::>(); - for record in &records { - let declaration = declarations[record.source_record.as_deref().unwrap()]; - validate_imported_policy(record, declaration).unwrap(); - } - - assert_eq!( - parse_hardware_target(&records[1]).unwrap(), - CatalogHardwareTarget::AnyOf { - alternatives: vec![ - CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 101 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 103 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 120 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 121 }, - ], - } - ); - - let mut missing = test_clc_admission(); - missing.variants.pop(); - assert!(expand_clc_admission(&missing).is_err()); - - let mut reordered = test_clc_admission(); - reordered.variants.swap(0, 1); - assert!(expand_clc_admission(&reordered).is_err()); - - let mut wrong_abi = test_clc_admission(); - wrong_abi.variants[0].abi_id = "i9999".into(); - assert!(expand_clc_admission(&wrong_abi).is_err()); - - let mut executed = test_clc_admission(); - executed.runtime_validation = RuntimeValidation::Executed; - assert!(expand_clc_admission(&executed).is_err()); - - let declaration = declarations[records[2].source_record.as_deref().unwrap()]; - let mut wrong_adapter = records[2].clone(); - wrong_adapter.clc.as_mut().unwrap().adapter = ClcAdapter::PairU64ToI128U32; - assert!(validate_imported_policy(&wrong_adapter, declaration).is_err()); - - let mut unsorted_targets = records[1].clone(); - unsorted_targets.targets = "sm_120a|sm_100a".into(); - assert!(parse_hardware_target(&unsorted_targets).is_err()); - - let mut duplicate_targets = records[1].clone(); - duplicate_targets.targets = "sm_100a|sm_100a".into(); - assert!(parse_hardware_target(&duplicate_targets).is_err()); - } - - #[test] - fn clc_compact_schema_is_reserved_for_aggregation() { - let shard = |schema| OverlayShardFile { - schema, - family: "clc".into(), - intrinsics: vec![], - register_mma_int4: None, - register_mma_int8: None, - register_mma_b1: None, - register_mma_f8f6f4_f32: None, - register_mma_f8f6f4_f16: None, - register_mma_mxf8f6f4_f32: None, - register_mma_fp8: None, - register_mma_ampere_float: None, - sparse_mma_integer: None, - sparse_mma_f8f6f4_f32: None, - sparse_mma_f8f6f4_f16: None, - prmt: None, - packed_conversion_fp8: None, - packed_conversion_fp8_f16x2: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cluster_sreg: None, - cluster_barrier: None, - mbarrier_extended: None, - special_registers: None, - debug_control: None, - threadfence: None, - cluster_memory: None, - stmatrix: None, - clc: Some(test_clc_admission()), - wgmma_controls: None, - tma: None, - tcgen05: None, - }; - let path = Path::new("intrinsics/overlay/clc.toml"); - validate_overlay_shard_schema_with_max(&shard(CLC_SHARD_SCHEMA), path, CLC_SHARD_SCHEMA) - .unwrap(); - assert!( - validate_overlay_shard_schema_with_max( - &shard(CLC_SHARD_SCHEMA - 1), - path, - CLC_SHARD_SCHEMA, - ) - .unwrap_err() - .to_string() - .contains("requires overlay shard schema 40") - ); - } - - #[test] - fn compact_tma_admission_matches_llvm_and_fails_closed() { - let records = expand_tma_admission(&test_tma_admission()).unwrap(); - assert_eq!( - records.len(), - TMA_OPERATIONS.len() + tma_reduction_matrix().len() - ); - assert!(records.iter().take(TMA_OPERATIONS.len()).all(|record| { - let operation = record.tma.as_ref().unwrap().operation; - let recipe = tma_recipe(operation); - record.backend_lowerings.iter().any(|route| { - route.backend == IntrinsicBackend::LlvmNvptx - && route.mechanism == recipe.llvm_mechanism - }) && record.backend_lowerings.iter().any(|route| { - route.backend == IntrinsicBackend::LibNvvm - && route.mechanism == BackendLoweringMechanism::InlinePtx - }) - })); - assert_eq!( - records - .iter() - .filter(|record| { - record - .tma - .as_ref() - .unwrap() - .operation - .prefetch_coordinate_count() - .is_some() - }) - .count(), - 12, - "the tile-prefetch subfamily must contain every plain/cache-hint pair" - ); - assert_eq!( - records - .iter() - .take(TMA_OPERATIONS.len()) - .map(|record| (record.abi_id.as_str(), record.id.as_str())) - .collect::>(), - [ - ("i0328", "cp_async_bulk_tensor_1d_g2s"), - ("i0329", "cp_async_bulk_tensor_2d_g2s"), - ("i0330", "cp_async_bulk_tensor_2d_g2s_multicast"), - ("i0331", "cp_async_bulk_tensor_2d_g2s_multicast_cg2"), - ("i0332", "cp_async_bulk_tensor_3d_g2s"), - ("i0333", "cp_async_bulk_tensor_4d_g2s"), - ("i0334", "cp_async_bulk_tensor_5d_g2s"), - ("i0335", "cp_async_bulk_tensor_1d_s2g"), - ("i0336", "cp_async_bulk_tensor_2d_s2g"), - ("i0337", "cp_async_bulk_tensor_3d_s2g"), - ("i0338", "cp_async_bulk_tensor_4d_s2g"), - ("i0339", "cp_async_bulk_tensor_5d_s2g"), - ("i0340", "cp_async_bulk_commit_group"), - ("i0341", "cp_async_bulk_wait_group"), - ("i0342", "cp_async_bulk_wait_group_read"), - ("i0887", "prefetch_tma_descriptor"), - ("i0888", "cp_async_bulk_prefetch_tensor_1d_l2"), - ("i0889", "cp_async_bulk_prefetch_tensor_2d_l2"), - ("i0890", "cp_async_bulk_prefetch_tensor_3d_l2"), - ("i0891", "cp_async_bulk_prefetch_tensor_4d_l2"), - ("i0892", "cp_async_bulk_prefetch_tensor_5d_l2"), - ("i0893", "cp_async_bulk_prefetch_tensor_gather4_2d_l2"), - ("i0894", "tensormap_replace_box_dim"), - ("i0895", "tensormap_replace_element_stride"), - ("i0896", "tensormap_replace_element_type"), - ("i0897", "tensormap_replace_fill_mode"), - ("i0898", "tensormap_replace_global_address"), - ("i0899", "tensormap_replace_global_dim"), - ("i0900", "tensormap_replace_global_stride"), - ("i0901", "tensormap_replace_interleave_layout"), - ("i0902", "tensormap_replace_rank"), - ("i0903", "tensormap_replace_swizzle_atomicity"), - ("i0904", "tensormap_replace_swizzle_mode"), - ("i0905", "fence_proxy_tensormap_generic_acquire_cluster"), - ("i0906", "fence_proxy_tensormap_generic_acquire_cta"), - ("i0907", "fence_proxy_tensormap_generic_acquire_gpu"), - ("i0908", "fence_proxy_tensormap_generic_acquire_system"), - ("i0909", "fence_proxy_tensormap_generic_release_cluster"), - ("i0910", "fence_proxy_tensormap_generic_release_cta"), - ("i0911", "fence_proxy_tensormap_generic_release_gpu"), - ("i0912", "fence_proxy_tensormap_generic_release_system"), - ("i0917", "cp_async_bulk_prefetch_tensor_1d_l2_cache_hint"), - ("i0918", "cp_async_bulk_prefetch_tensor_2d_l2_cache_hint"), - ("i0919", "cp_async_bulk_prefetch_tensor_3d_l2_cache_hint"), - ("i0920", "cp_async_bulk_prefetch_tensor_4d_l2_cache_hint"), - ("i0921", "cp_async_bulk_prefetch_tensor_5d_l2_cache_hint"), - ( - "i0922", - "cp_async_bulk_prefetch_tensor_gather4_2d_l2_cache_hint" - ), - ] - ); - - let reductions = &records[TMA_OPERATIONS.len()..]; - assert_eq!(reductions.len(), 64); - assert_eq!( - (reductions[0].abi_id.as_str(), reductions[0].id.as_str()), - ("i0923", "cp_async_bulk_tensor_reduce_add_tile_1d") - ); - assert_eq!( - (reductions[7].abi_id.as_str(), reductions[7].id.as_str()), - ("i0930", "cp_async_bulk_tensor_reduce_add_im2col_5d") - ); - assert_eq!( - (reductions[63].abi_id.as_str(), reductions[63].id.as_str()), - ("i0986", "cp_async_bulk_tensor_reduce_xor_im2col_5d") - ); - assert!(reductions.iter().all(|record| { - record.tma.as_ref().is_some_and(|tma| { - tma.operation == TmaOperation::Reduce - && tma.reduction.is_some() - && tma.adapter == TmaAdapter::ReductionPointersCoordinatesInjectDefaults - }) - })); - - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - let declarations = imported - .intrinsics - .iter() - .map(|record| (record.source_record.as_str(), record)) - .collect::>(); - for record in &records { - let declaration = declarations[record.source_record.as_deref().unwrap()]; - validate_imported_policy(record, declaration).unwrap(); - } - - assert_eq!( - parse_hardware_target(&records[3]).unwrap(), - CatalogHardwareTarget::AnyOf { - alternatives: vec![ - CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 101 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 103 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, - ], - } - ); - - let mut missing = test_tma_admission(); - missing.variants.pop(); - assert!(expand_tma_admission(&missing).is_err()); - - let mut reordered = test_tma_admission(); - reordered.variants.swap(0, 1); - assert!(expand_tma_admission(&reordered).is_err()); - - let mut wrong_abi = test_tma_admission(); - wrong_abi.variants[0].abi_id = "i9999".into(); - assert!(expand_tma_admission(&wrong_abi).is_err()); - - let mut missing_reduction = test_tma_admission(); - missing_reduction.reduce_variants.pop(); - assert!(expand_tma_admission(&missing_reduction).is_err()); - - let mut reordered_reduction = test_tma_admission(); - reordered_reduction.reduce_variants.swap(0, 1); - assert!(expand_tma_admission(&reordered_reduction).is_err()); - - let mut non_contiguous_reduction_abi = test_tma_admission(); - non_contiguous_reduction_abi.reduce_variants[0].abi_id = "i9999".into(); - let records = expand_tma_admission(&non_contiguous_reduction_abi).unwrap(); - assert_eq!(records[TMA_OPERATIONS.len()].abi_id, "i9999"); - - let mut executed = test_tma_admission(); - executed.runtime_validation = RuntimeValidation::Executed; - assert!(expand_tma_admission(&executed).is_err()); - - let declaration = declarations[records[0].source_record.as_deref().unwrap()]; - let mut wrong_adapter = records[0].clone(); - wrong_adapter.tma.as_mut().unwrap().adapter = TmaAdapter::NoOperands; - assert!(validate_imported_policy(&wrong_adapter, declaration).is_err()); - } - - #[test] - fn tma_compact_schema_is_reserved_for_aggregation() { - let shard = |schema: u32, include_reductions: bool| { - let mut admission = test_tma_admission(); - if !include_reductions { - admission.reduce_variants.clear(); - } - OverlayShardFile { - schema, - family: "tma".into(), - intrinsics: vec![], - register_mma_int4: None, - register_mma_int8: None, - register_mma_b1: None, - register_mma_f8f6f4_f32: None, - register_mma_f8f6f4_f16: None, - register_mma_mxf8f6f4_f32: None, - register_mma_fp8: None, - register_mma_ampere_float: None, - sparse_mma_integer: None, - sparse_mma_f8f6f4_f32: None, - sparse_mma_f8f6f4_f16: None, - prmt: None, - packed_conversion_fp8: None, - packed_conversion_fp8_f16x2: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cluster_sreg: None, - cluster_barrier: None, - mbarrier_extended: None, - special_registers: None, - debug_control: None, - threadfence: None, - cluster_memory: None, - stmatrix: None, - clc: None, - wgmma_controls: None, - tma: Some(admission), - tcgen05: None, - } - }; - let path = Path::new("intrinsics/overlay/tma.toml"); - - validate_overlay_shard_schema_with_max( - &shard(TMA_SHARD_SCHEMA, false), - path, - TMA_REDUCTION_SHARD_SCHEMA, - ) - .unwrap(); - assert!( - validate_overlay_shard_schema_with_max( - &shard(TMA_SHARD_SCHEMA - 1, false), - path, - TMA_REDUCTION_SHARD_SCHEMA, - ) - .unwrap_err() - .to_string() - .contains("compact TMA admission requires overlay shard schema 61") - ); - - validate_overlay_shard_schema_with_max( - &shard(TMA_REDUCTION_SHARD_SCHEMA, true), - path, - TMA_REDUCTION_SHARD_SCHEMA, - ) - .unwrap(); - assert!( - validate_overlay_shard_schema_with_max( - &shard(TMA_REDUCTION_SHARD_SCHEMA - 1, true), - path, - TMA_REDUCTION_SHARD_SCHEMA, - ) - .unwrap_err() - .to_string() - .contains("compact TMA reduction admission requires overlay shard schema 62") - ); - } - - #[test] - fn compact_tcgen05_admission_matches_llvm_and_fails_closed() { - let records = expand_tcgen05_admission(&test_tcgen05_admission()).unwrap(); - assert_eq!(records.len(), 27); - assert_eq!( - records - .iter() - .map(|record| (record.abi_id.as_str(), record.id.as_str())) - .collect::>(), - [ - ("i0343", "tcgen05_alloc"), - ("i0344", "tcgen05_dealloc"), - ("i0345", "tcgen05_relinquish_alloc_permit"), - ("i0346", "tcgen05_fence_before_thread_sync"), - ("i0347", "tcgen05_fence_after_thread_sync"), - ("i0348", "tcgen05_commit"), - ("i0349", "tcgen05_commit_shared_cluster"), - ("i0350", "tcgen05_mma_ws_f16"), - ("i0351", "tcgen05_mma_f16"), - ("i0352", "tcgen05_mma_ws_bf16"), - ("i0353", "tcgen05_mma_ws_tf32"), - ("i0354", "tcgen05_cp_smem_to_tmem"), - ("i0355", "tcgen05_ld_16x256b_x8_pure"), - ("i0356", "tcgen05_ld_16x256b_pure"), - ("i0357", "tcgen05_load_wait"), - ("i0358", "tcgen05_store_wait"), - ("i0359", "tcgen05_alloc_cg2"), - ("i0360", "tcgen05_dealloc_cg2"), - ("i0361", "tcgen05_relinquish_alloc_permit_cg2"), - ("i0362", "tcgen05_mma_f16_cg2"), - ("i0363", "tcgen05_commit_cg2"), - ("i0364", "tcgen05_commit_shared_cluster_cg2"), - ("i0365", "tcgen05_commit_multicast_cg2"), - ("i0366", "tcgen05_cp_smem_to_tmem_cg2"), - ("i0760", "tcgen05_commit_multicast"), - ("i0761", "tcgen05_shift_down"), - ("i0762", "tcgen05_shift_down_cg2"), - ] - ); - - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - let declarations = imported - .intrinsics - .iter() - .map(|record| (record.source_record.as_str(), record)) - .collect::>(); - for record in &records { - let declaration = declarations[record.source_record.as_deref().unwrap()]; - validate_imported_policy(record, declaration).unwrap(); - assert_tcgen05_backend_target_split(record); - assert_eq!( - record.execution_scope, - record.tcgen05.as_ref().unwrap().operation.execution_scope() - ); - } - - let bf16 = records - .iter() - .find(|record| record.id == "tcgen05_mma_ws_bf16") - .unwrap(); - assert!(bf16.expected_ptx.modifiers.contains(&"kind::f16".into())); - assert!(!bf16.expected_ptx.modifiers.contains(&"kind::bf16".into())); - assert_eq!( - bf16.tcgen05.as_ref().unwrap().source_contract, - Tcgen05SourceContract::TablegenSelectionChangesPtx - ); - - let mma_f16 = records - .iter() - .find(|record| record.id == "tcgen05_mma_f16") - .unwrap(); - assert!(mma_f16.expected_ptx.matches( - "tcgen05.mma.cta_group::1.kind::f16 [%r1], %rd1, %rd2, %r2, {%z, %z, %z, %z}, %enable_pred;" - ).unwrap()); - assert!(!mma_f16.expected_ptx.matches( - "tcgen05.mma.cta_group::1.kind::f16 [%r1], %rd1, %rd2, %r2, {%clock64, %z, %z, %z}, %enable_pred;" - ).unwrap()); - assert!(!mma_f16.expected_ptx.matches( - "tcgen05.mma.cta_group::1.kind::f16 [%r1], %rd1, %rd2, %r2, {%z, %z, %z, %z}, %other_pred;" - ).unwrap()); - - let mma_f16_cg2 = records - .iter() - .find(|record| record.id == "tcgen05_mma_f16_cg2") - .unwrap(); - assert!(mma_f16_cg2.expected_ptx.matches( - "tcgen05.mma.cta_group::2.kind::f16 [%r1], %rd1, %rd2, %r2, {%z, %z, %z, %z, %z, %z, %z, %z}, %enable_pred;" - ).unwrap()); - assert!( - !records - .iter() - .any(|record| { record.rust_name == "tcgen05_mma_ws_f16_with_collector" }) - ); - assert_eq!( - parse_hardware_target(&records[0]).unwrap(), - CatalogHardwareTarget::AnyOf { - alternatives: vec![ - CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 101 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 103 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, - ], - } - ); - - let legacy_records = - expand_tcgen05_admission(&without_tcgen05_control(test_tcgen05_admission())).unwrap(); - assert_eq!(legacy_records.len(), 24); - assert!(legacy_records.iter().all(|record| { - record - .backend_lowerings - .iter() - .map(|lowering| lowering.evidence_profile.as_str()) - .eq(["llvm-tcgen05-test", "libnvvm-tcgen05-test"]) - })); - - let multicast = records - .iter() - .find(|record| record.id == "tcgen05_commit_multicast") - .unwrap(); - assert_eq!(multicast.abi_id, "i0760"); - assert_eq!(multicast.rust_arguments, ["*mut u64", "u16"]); - assert_eq!(multicast.llvm_arguments, ["shared_ptr", "i16"]); - assert_eq!(multicast.execution_scope, "thread"); - assert_eq!( - multicast.backend_lowerings[0].evidence_profile, - "llvm-tcgen05-control-test" - ); - assert_eq!( - multicast.backend_lowerings[1].evidence_profile, - "libnvvm-tcgen05-control-test" - ); - - for (id, group) in [ - ("tcgen05_shift_down", "cta_group::1"), - ("tcgen05_shift_down_cg2", "cta_group::2"), - ] { - let shift = records.iter().find(|record| record.id == id).unwrap(); - assert_eq!(shift.rust_arguments, ["u32"]); - assert_eq!(shift.llvm_arguments, ["tmem_ptr"]); - assert_eq!(shift.execution_scope, "thread"); - assert_eq!(shift.memory, "read_write"); - assert_eq!(shift.expected_ptx.modifiers, ["shift", group, "down"]); - assert_eq!(shift.expected_ptx.operands, [OperandPattern::Address]); - } - - let mut missing = test_tcgen05_admission(); - missing.variants.pop(); - assert!(expand_tcgen05_admission(&missing).is_err()); - - let mut reordered = test_tcgen05_admission(); - reordered.variants.swap(0, 1); - assert!(expand_tcgen05_admission(&reordered).is_err()); - - let mut wrong_abi = test_tcgen05_admission(); - wrong_abi.variants[0].abi_id = "i9999".into(); - assert!(expand_tcgen05_admission(&wrong_abi).is_err()); - - let mut executed = test_tcgen05_admission(); - executed.runtime_validation = RuntimeValidation::Executed; - assert!(expand_tcgen05_admission(&executed).is_err()); - - let mut missing_control_evidence = test_tcgen05_admission(); - missing_control_evidence.control_llvm_evidence_profile = None; - assert!(expand_tcgen05_admission(&missing_control_evidence).is_err()); - - let mut partial_control = without_tcgen05_control(test_tcgen05_admission()); - partial_control - .variants - .push(test_tcgen05_admission().variants[24].clone()); - assert!(expand_tcgen05_admission(&partial_control).is_err()); - - let shift = records - .iter() - .find(|record| record.id == "tcgen05_shift_down") - .unwrap(); - let mut wrong_shift_declaration = - declarations[shift.source_record.as_deref().unwrap()].clone(); - wrong_shift_declaration.selections[0].predicates = - vec!["Subtarget->hasTcgen05InstSupport()".into()]; - assert!(validate_imported_policy(shift, &wrong_shift_declaration).is_err()); - - let declaration = declarations[records[0].source_record.as_deref().unwrap()]; - let mut wrong_adapter = records[0].clone(); - wrong_adapter.tcgen05.as_mut().unwrap().adapter = Tcgen05Adapter::NoOperands; - assert!(validate_imported_policy(&wrong_adapter, declaration).is_err()); - - let mut changed_declaration = declaration.clone(); - changed_declaration.properties.pop(); - assert!(validate_imported_policy(&records[0], &changed_declaration).is_err()); - - let mut broadened_libnvvm = records[0].clone(); - broadened_libnvvm.backend_lowerings[1].targets = None; - assert!(validate_imported_policy(&broadened_libnvvm, declaration).is_err()); - - let mut wrong_alloc_scope = records[0].clone(); - wrong_alloc_scope.execution_scope = "thread".into(); - assert!(validate_imported_policy(&wrong_alloc_scope, declaration).is_err()); - - let pure_load = &records[12]; - let pure_declaration = declarations[pure_load.source_record.as_deref().unwrap()]; - let mut wrong_pure_load_scope = pure_load.clone(); - wrong_pure_load_scope.execution_scope = "thread".into(); - assert!(validate_imported_policy(&wrong_pure_load_scope, pure_declaration).is_err()); - } - - #[test] - fn compact_tcgen05_mma_admission_closes_sources_selectors_and_targets() { - let records = expand_tcgen05_admission(&test_tcgen05_mma_admission()).unwrap(); - assert_eq!(records.len(), 51); - let mma = &records[27..]; - assert_eq!(mma.len(), 24); - assert_eq!( - mma.iter() - .map(|record| (record.abi_id.as_str(), record.id.as_str())) - .collect::>(), - [ - ("i0763", "tcgen05_mma_shared"), - ("i0764", "tcgen05_mma_tensor"), - ("i0765", "tcgen05_mma_tensor_ashift"), - ("i0766", "tcgen05_mma_sp_shared"), - ("i0767", "tcgen05_mma_sp_tensor"), - ("i0768", "tcgen05_mma_sp_tensor_ashift"), - ("i0769", "tcgen05_mma_ws_shared"), - ("i0770", "tcgen05_mma_ws_shared_zero_col_mask"), - ("i0771", "tcgen05_mma_ws_sp_shared"), - ("i0772", "tcgen05_mma_ws_sp_shared_zero_col_mask"), - ("i0773", "tcgen05_mma_ws_sp_tensor"), - ("i0774", "tcgen05_mma_ws_sp_tensor_zero_col_mask"), - ("i0775", "tcgen05_mma_ws_tensor"), - ("i0776", "tcgen05_mma_ws_tensor_zero_col_mask"), - ("i0777", "tcgen05_mma_ws_e4m3"), - ("i0778", "tcgen05_mma_ws_e5m2"), - ("i0779", "tcgen05_mma_ws_e2m3"), - ("i0780", "tcgen05_mma_ws_e3m2"), - ("i0781", "tcgen05_mma_ws_e2m1"), - ("i1011", "tcgen05_mma_e4m3"), - ("i1012", "tcgen05_mma_e5m2"), - ("i1013", "tcgen05_mma_e2m3"), - ("i1014", "tcgen05_mma_e3m2"), - ("i1015", "tcgen05_mma_e2m1"), - ] - ); - assert!(mma.iter().all(|record| { - record.dialect_op_type == TCGEN05_MMA_DIALECT_OP_TYPE - && record.dialect_op_name == TCGEN05_MMA_DIALECT_OP_NAME - })); - assert_eq!( - [ - Tcgen05MmaBUsage::Discard, - Tcgen05MmaBUsage::LastUse, - Tcgen05MmaBUsage::Fill, - Tcgen05MmaBUsage::Use, - ] - .map(Tcgen05MmaBUsage::selector_value), - [0, 1, 2, 3] - ); - - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - let declarations = imported - .intrinsics - .iter() - .map(|record| (record.source_record.as_str(), record)) - .collect::>(); - for record in mma { - let declaration = declarations[record.source_record.as_deref().unwrap()]; - validate_imported_policy(record, declaration).unwrap(); - assert_eq!(record.memory, "read_write"); - assert!(record.convergent); - assert_eq!(record.execution_scope, "thread"); - assert!(!record.safe); - assert_eq!( - record - .expected_ptx - .operands - .iter() - .filter(|operand| { - **operand - == (OperandPattern::Exact { - value: "%enable_pred".into(), - }) - }) - .count(), - 1 - ); - } - - let dense = &mma[0]; - assert!(dense.expected_ptx.matches( - "tcgen05.mma.cta_group::1.kind::f16.collector::a::discard [%r1], %rd1, %rd2, %r2, %enable_pred;" - ).unwrap()); - assert!(!dense.expected_ptx.matches( - "tcgen05.mma.cta_group::1.kind::f16.collector::a::discard [%r1], %rd1, %rd2, %r2, %clock64;" - ).unwrap()); - assert!(mma[4].expected_ptx.matches( - "tcgen05.mma.sp.cta_group::1.kind::f16.collector::a::discard [%r1], [%r2], %rd1, [%r3], %r4, %enable_pred;" - ).unwrap()); - assert!(mma[7].expected_ptx.matches( - "tcgen05.mma.ws.cta_group::1.kind::f16.collector::b0::discard [%r1], %rd1, %rd2, %r2, %enable_pred, %rd3;" - ).unwrap()); - assert!(mma[14].expected_ptx.matches( - "tcgen05.mma.ws.cta_group::1.kind::f8f6f4.collector::b0::discard [%r1], [%r2], %rd1, %r3, %enable_pred;" - ).unwrap()); - assert!(mma[19].expected_ptx.matches( - "tcgen05.mma.cta_group::1.kind::f8f6f4.collector::a::discard [%r1], %rd1, %rd2, %r2, %enable_pred;" - ).unwrap()); - - let canonical = &mma[..14]; - assert_eq!( - canonical - .iter() - .map(|record| record.source_record.as_deref().unwrap()) - .collect::>() - .len(), - 14 - ); - let selected_count = |record: &OverlayIntrinsic| { - declarations[record.source_record.as_deref().unwrap()] - .selections - .iter() - .filter(|selection| selection_matches_policy(record, selection).unwrap()) - .count() - }; - let newly_matched = canonical - .iter() - .filter(|record| record.id != "tcgen05_mma_ws_tensor") - .map(selected_count) - .sum::(); - assert_eq!(newly_matched, 608); - assert_eq!(selected_count(&canonical[12]), 64); - assert_eq!(canonical.iter().map(selected_count).sum::(), 672); - assert_eq!(mma[14..].iter().map(selected_count).sum::(), 10); - assert_eq!( - [2, 5] - .into_iter() - .map(|index| { - declarations[canonical[index].source_record.as_deref().unwrap()] - .selections - .len() - - selected_count(&canonical[index]) - }) - .sum::(), - 32 - ); - - for record in canonical { - assert_eq!( - record.compatibility_rust_paths, - [format!("cuda_device::tcgen05::__{}", record.id)] - ); - assert_eq!( - &record.rust_arguments[record.rust_arguments.len() - 3..], - ["u32", "u32", "u32"] - ); - let contract = record.tcgen05.as_ref().unwrap().mma.as_ref().unwrap(); - let CatalogHardwareTarget::TargetMatrix { contracts } = &contract.llvm_target.hardware - else { - panic!("generic LLVM tcgen05 MMA target must be a matrix") - }; - assert_eq!(contracts.len(), 4); - assert_eq!( - contracts - .iter() - .map(|contract| contract.alternatives.len()) - .collect::>(), - [8, 8, 3, 8] - ); - let CatalogHardwareTarget::TargetMatrix { contracts } = - &contract.libnvvm_target.hardware - else { - panic!("generic libNVVM tcgen05 MMA target must be a matrix") - }; - assert_eq!(contracts.len(), 4); - assert_eq!( - contracts - .iter() - .map(|contract| contract.alternatives.len()) - .collect::>(), - [6, 6, 2, 6] - ); - } - for record in &mma[14..19] { - assert_eq!( - record.compatibility_rust_paths, - [format!("cuda_device::tcgen05::{}", record.id)] - ); - assert_eq!( - record.rust_arguments, - ["u32", "u32", "u64", "u64", "u32", "bool"] - ); - assert_eq!(record.dialect_operands, ["i32", "i32", "i64", "i32", "i1"]); - let contract = record.tcgen05.as_ref().unwrap().mma.as_ref().unwrap(); - assert_eq!( - contract.fixed_selectors, - Some(Tcgen05MmaFixedSelectors { - kind: Tcgen05MmaKind::F8f6f4, - b_buffer: 0, - b_usage: Tcgen05MmaBUsage::Discard, - }) - ); - let CatalogHardwareTarget::TargetMatrix { contracts } = &contract.llvm_target.hardware - else { - panic!("fixed LLVM tcgen05 MMA target must be a matrix") - }; - assert_eq!(contracts.len(), 1); - assert_eq!(contracts[0].selectors[0].value, "f8f6f4"); - } - for record in &mma[19..] { - assert_eq!( - record.compatibility_rust_paths, - [format!("cuda_device::tcgen05::{}", record.id)] - ); - assert_eq!(record.rust_arguments, ["u32", "u64", "u64", "u32", "bool"]); - assert_eq!(record.dialect_operands, ["i32", "i64", "i64", "i32", "i1"]); - let contract = record.tcgen05.as_ref().unwrap().mma.as_ref().unwrap(); - assert_eq!(contract.form, Tcgen05MmaForm::Shared); - assert_eq!( - contract.fixed_selectors, - Some(Tcgen05MmaFixedSelectors { - kind: Tcgen05MmaKind::F8f6f4, - b_buffer: 0, - b_usage: Tcgen05MmaBUsage::Discard, - }) - ); - assert_eq!( - record.tcgen05.as_ref().unwrap().adapter, - Tcgen05Adapter::MmaDirectSelectors - ); - let CatalogHardwareTarget::TargetMatrix { contracts } = &contract.llvm_target.hardware - else { - panic!("fixed LLVM tcgen05 MMA target must be a matrix") - }; - assert_eq!(contracts.len(), 1); - assert_eq!(contracts[0].selectors[0].value, "f8f6f4"); - } - - for legacy in [ - "tcgen05_mma_ws_f16", - "tcgen05_mma_ws_bf16", - "tcgen05_mma_ws_tf32", - ] { - let old = &records[..27] - .iter() - .find(|record| record.id == legacy) - .unwrap(); - assert!(old.tcgen05.as_ref().unwrap().mma.is_none()); - assert!(old.abi_id.starts_with("i035")); - assert_eq!( - old.expected_ptx.operands.last(), - Some(&OperandPattern::Exact { - value: "%enable_pred".into(), - }) - ); - } - let legacy_f16 = records - .iter() - .find(|record| record.id == "tcgen05_mma_ws_f16") - .unwrap(); - assert!( - legacy_f16 - .expected_ptx - .matches( - "tcgen05.mma.ws.cta_group::1.kind::f16 [%r1], [%r2], %rd1, %r3, %enable_pred;" - ) - .unwrap() - ); - assert!( - !legacy_f16 - .expected_ptx - .matches( - "tcgen05.mma.ws.cta_group::1.kind::f16 [%r1], [%r2], %rd1, %r3, %other_pred;" - ) - .unwrap() - ); - } - - #[test] - fn compact_tcgen05_mma_admission_fails_closed_on_drift() { - let mut missing = test_tcgen05_mma_admission(); - missing.mma_variants.pop(); - assert!(expand_tcgen05_admission(&missing).is_err()); - - let mut reordered = test_tcgen05_mma_admission(); - reordered.mma_variants.swap(0, 1); - assert!(expand_tcgen05_admission(&reordered).is_err()); - - let mut non_contiguous_abi = test_tcgen05_mma_admission(); - non_contiguous_abi.mma_variants[0].abi_id = "i9999".into(); - let records = expand_tcgen05_admission(&non_contiguous_abi).unwrap(); - assert!(records.iter().any(|record| record.abi_id == "i9999")); - - let mut missing_target = test_tcgen05_mma_admission(); - missing_target.mma_llvm_target_contracts[0] - .alternatives - .pop(); - assert!(expand_tcgen05_admission(&missing_target).is_err()); - - let mut broadened_i8 = test_tcgen05_mma_admission(); - broadened_i8.mma_libnvvm_target_contracts[2] - .alternatives - .insert( - 1, - crate::model::TargetContractAlternative { - target: "sm_103a".into(), - minimum_ptx: "8.8".into(), - }, - ); - assert!(expand_tcgen05_admission(&broadened_i8).is_err()); - - let records = expand_tcgen05_admission(&test_tcgen05_mma_admission()).unwrap(); - let record = records - .iter() - .find(|record| record.id == "tcgen05_mma_tensor_ashift") - .unwrap(); - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - let declaration = imported - .intrinsics - .iter() - .find(|declaration| { - Some(declaration.source_record.as_str()) == record.source_record.as_deref() - }) - .unwrap(); - - let mut wrong_layout = record.clone(); - let mma = wrong_layout.tcgen05.as_mut().unwrap().mma.as_mut().unwrap(); - mma.selector_layout = Tcgen05MmaSelectorLayout::Base { - kind_argument: 6, - cta_group_argument: 5, - collector_a_argument: 7, - collector_a_upper_exclusive: 2, - }; - assert!(validate_imported_policy(&wrong_layout, declaration).is_err()); - - let mut wrong_properties = declaration.clone(); - wrong_properties.properties.pop(); - assert!(validate_imported_policy(record, &wrong_properties).is_err()); - - let mut wrong_range = declaration.clone(); - let range = wrong_range - .properties - .iter_mut() - .find(|property| property.as_str() == "Range") - .unwrap(); - *range = "Range".into(); - assert!(validate_imported_policy(record, &wrong_range).is_err()); - - let mut wrong_predicate = declaration.clone(); - wrong_predicate.selections[0].predicates = vec!["hasSM100a".into()]; - assert!(validate_imported_policy(record, &wrong_predicate).is_err()); - - let mut wrong_memory = record.clone(); - wrong_memory.memory = "write".into(); - assert!(validate_imported_policy(&wrong_memory, declaration).is_err()); - - let mut wrong_convergence = record.clone(); - wrong_convergence.convergent = false; - assert!(validate_imported_policy(&wrong_convergence, declaration).is_err()); - - let mut split_dialect_op = record.clone(); - split_dialect_op.dialect_op_type = "Tcgen05MmaTensorAshiftOp".into(); - split_dialect_op.dialect_op_name = "nvvm.tcgen05_mma_tensor_ashift".into(); - assert!(validate_imported_policy(&split_dialect_op, declaration).is_err()); - } - - #[test] - fn compact_tcgen05_copy_admission_matches_all_llvm_records_and_fails_closed() { - let records = expand_tcgen05_admission(&test_tcgen05_cp_admission()).unwrap(); - assert_eq!(records.len(), 61); - let copies = &records[27..]; - assert_eq!(copies.len(), 34); - assert_eq!( - (copies[0].abi_id.as_str(), copies[0].id.as_str()), - ("i0578", "tcgen05_cp_128x128b_b4x16_p64") - ); - assert_eq!( - (copies[33].abi_id.as_str(), copies[33].id.as_str()), - ("i0611", "tcgen05_cp_64x128b_warpx2_02_13_cg2") - ); - - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - let declarations = imported - .intrinsics - .iter() - .map(|record| (record.source_record.as_str(), record)) - .collect::>(); - for record in copies { - let declaration = declarations[record.source_record.as_deref().unwrap()]; - validate_imported_policy(record, declaration).unwrap(); - assert_tcgen05_backend_target_split(record); - assert_eq!(record.rust_arguments, ["u32", "u64"]); - assert_eq!(record.rust_result, "()"); - assert!(!record.safe); - assert!(record.convergent); - assert_eq!(record.memory, "read_write"); - assert_eq!(record.targets, TCGEN05_LLVM_TARGETS); - } - - let packed = copies - .iter() - .find(|record| record.id == "tcgen05_cp_128x128b_b4x16_p64") - .unwrap(); - assert_eq!( - packed.expected_ptx.modifiers, - ["cp", "cta_group::1", "128x128b", "b8x16", "b4x16_p64"] - ); - let warpx4 = copies - .iter() - .find(|record| record.id == "tcgen05_cp_32x128b_warpx4") - .unwrap(); - assert_eq!( - warpx4.expected_ptx.modifiers, - ["cp", "cta_group::1", "32x128b", "warpx4"] - ); - let pair_01_23 = copies - .iter() - .find(|record| record.id == "tcgen05_cp_64x128b_warpx2_01_23_b6x16_p32_cg2") - .unwrap(); - assert_eq!( - pair_01_23.expected_ptx.modifiers, - [ - "cp", - "cta_group::2", - "64x128b", - "warpx2::01_23", - "b8x16", - "b6x16_p32" - ] - ); - let pair_02_13 = copies - .iter() - .find(|record| record.id == "tcgen05_cp_64x128b_warpx2_02_13") - .unwrap(); - assert_eq!( - pair_02_13.expected_ptx.modifiers, - ["cp", "cta_group::1", "64x128b", "warpx2::02_13"] - ); - - let mut missing = test_tcgen05_cp_admission(); - missing.cp_variants.pop(); - assert!(expand_tcgen05_admission(&missing).is_err()); - - let mut reordered = test_tcgen05_cp_admission(); - reordered.cp_variants.swap(0, 1); - assert!(expand_tcgen05_admission(&reordered).is_err()); - - let mut non_contiguous_abi = test_tcgen05_cp_admission(); - non_contiguous_abi.cp_variants[0].abi_id = "i9999".into(); - let records = expand_tcgen05_admission(&non_contiguous_abi).unwrap(); - assert!(records.iter().any(|record| record.abi_id == "i9999")); - - let mut missing_evidence = test_tcgen05_cp_admission(); - missing_evidence.cp_llvm_evidence_profile = None; - assert!(expand_tcgen05_admission(&missing_evidence).is_err()); - - let declaration = declarations[copies[0].source_record.as_deref().unwrap()]; - let mut wrong_spelling = copies[0].clone(); - wrong_spelling.expected_ptx.modifiers.remove(3); - assert!(validate_imported_policy(&wrong_spelling, declaration).is_err()); - - let mut broadened_libnvvm = copies[0].clone(); - broadened_libnvvm.backend_lowerings[1].targets = None; - assert!(validate_imported_policy(&broadened_libnvvm, declaration).is_err()); - } - - #[test] - fn compact_tcgen05_load_admission_matches_all_llvm_records_and_fails_closed() { - let records = expand_tcgen05_admission(&test_tcgen05_ld_admission()).unwrap(); - assert_eq!(records.len(), 119); - let loads = &records[61..]; - assert_eq!(loads.len(), 58); - assert_eq!( - (loads[0].abi_id.as_str(), loads[0].id.as_str()), - ("i0612", "tcgen05_ld_16x64b_x1_raw") - ); - assert_eq!( - (loads[1].abi_id.as_str(), loads[1].id.as_str()), - ("i0613", "tcgen05_ld_16x64b_x1_pack16") - ); - assert_eq!( - (loads[57].abi_id.as_str(), loads[57].id.as_str()), - ("i0669", "tcgen05_ld_32x32b_x128_pack16") - ); - - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - let declarations = imported - .intrinsics - .iter() - .map(|record| (record.source_record.as_str(), record)) - .collect::>(); - let mut source_counts = BTreeMap::new(); - for record in loads { - let source_record = record.source_record.as_deref().unwrap(); - *source_counts.entry(source_record).or_insert(0) += 1; - let declaration = declarations[source_record]; - validate_imported_policy(record, declaration).unwrap(); - assert_tcgen05_backend_target_split(record); - assert_eq!(record.rust_arguments, ["u32"]); - assert!(record.must_use); - assert!(!record.safe); - assert!(!record.pure); - assert_eq!(record.memory, "read"); - assert!(record.convergent); - assert_eq!(record.execution_scope, "warp"); - assert!( - !record - .rust_arguments - .iter() - .any(|argument| argument == "bool") - ); - assert_eq!(record.dialect_operands, ["i32"]); - assert_eq!(record.llvm_arguments, ["tmem_ptr", "i1"]); - assert!(declaration.selections.is_empty()); - } - assert_eq!(source_counts.len(), 29); - assert!(source_counts.values().all(|count| *count == 2)); - assert_eq!( - records - .iter() - .filter(|record| { - record.source_record.as_deref() == Some("int_nvvm_tcgen05_ld_16x256b_x1") - }) - .count(), - 3 - ); - assert_eq!( - records - .iter() - .filter(|record| { - record.source_record.as_deref() == Some("int_nvvm_tcgen05_ld_16x256b_x8") - }) - .count(), - 3 - ); - validate_unique_overlay(&records, 1).unwrap(); - - let scalar_raw = &loads[0]; - assert_eq!(scalar_raw.rust_result, "u32"); - assert_eq!(scalar_raw.dialect_results, ["i32"]); - assert_eq!(scalar_raw.llvm_results, ["i32"]); - assert_eq!( - scalar_raw.expected_ptx.modifiers, - ["ld", "sync", "aligned", "16x64b", "x1", "b32"] - ); - assert_eq!( - scalar_raw.expected_ptx.operands, - [ - OperandPattern::RegisterList { length: 1 }, - OperandPattern::Address, - ] - ); - assert!( - scalar_raw - .expected_ptx - .matches("tcgen05.ld.sync.aligned.16x64b.x1.b32 {%r1}, [%r2];") - .unwrap() - ); - assert!( - !scalar_raw - .expected_ptx - .matches("tcgen05.ld.sync.aligned.16x64b.x1.b32 %r1, [%r2];") - .unwrap() - ); - let scalar_pack = &loads[1]; - assert_eq!( - scalar_pack.expected_ptx.modifiers, - ["ld", "sync", "aligned", "16x64b", "x1", "pack::16b", "b32",] - ); - let largest = &loads[57]; - assert_eq!(largest.rust_result, "[u32; 128]"); - assert_eq!(largest.dialect_results.len(), 128); - assert_eq!(largest.llvm_results, ["v128i32"]); - assert_eq!( - largest.expected_ptx.operands, - [ - OperandPattern::RegisterList { length: 128 }, - OperandPattern::Address, - ] - ); - - let mut missing = test_tcgen05_ld_admission(); - missing.ld_variants.pop(); - assert!(expand_tcgen05_admission(&missing).is_err()); - - let mut reordered = test_tcgen05_ld_admission(); - reordered.ld_variants.swap(0, 1); - assert!(expand_tcgen05_admission(&reordered).is_err()); - - let mut non_contiguous_abi = test_tcgen05_ld_admission(); - non_contiguous_abi.ld_variants[0].abi_id = "i9999".into(); - let records = expand_tcgen05_admission(&non_contiguous_abi).unwrap(); - assert!(records.iter().any(|record| record.abi_id == "i9999")); - - let mut missing_evidence = test_tcgen05_ld_admission(); - missing_evidence.ld_llvm_evidence_profile = None; - assert!(expand_tcgen05_admission(&missing_evidence).is_err()); - - let declaration = declarations[loads[0].source_record.as_deref().unwrap()]; - let mut wrong_selector = loads[0].clone(); - wrong_selector - .tcgen05 - .as_mut() - .unwrap() - .ld - .as_mut() - .unwrap() - .pack16 = true; - assert!(validate_imported_policy(&wrong_selector, declaration).is_err()); - - let mut changed_declaration = declaration.clone(); - changed_declaration.properties.pop(); - assert!(validate_imported_policy(&loads[0], &changed_declaration).is_err()); - - let mut wrong_scope = loads[0].clone(); - wrong_scope.execution_scope = "thread".into(); - assert!(validate_imported_policy(&wrong_scope, declaration).is_err()); - - let mut unreviewed_sharing = records.clone(); - unreviewed_sharing[62].tcgen05.as_mut().unwrap().operation = Tcgen05Operation::Alloc; - assert!(validate_unique_overlay(&unreviewed_sharing, 1).is_err()); - } - - #[test] - fn compact_tcgen05_store_admission_matches_all_llvm_records_and_fails_closed() { - let records = expand_tcgen05_admission(&test_tcgen05_st_admission()).unwrap(); - assert_eq!(records.len(), 177); - let stores = &records[119..]; - assert_eq!(stores.len(), 58); - assert_eq!( - (stores[0].abi_id.as_str(), stores[0].id.as_str()), - ("i0670", "tcgen05_st_16x64b_x1_raw") - ); - assert_eq!( - (stores[1].abi_id.as_str(), stores[1].id.as_str()), - ("i0671", "tcgen05_st_16x64b_x1_unpack16") - ); - assert_eq!( - (stores[57].abi_id.as_str(), stores[57].id.as_str()), - ("i0727", "tcgen05_st_32x32b_x128_unpack16") - ); - - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - let declarations = imported - .intrinsics - .iter() - .map(|record| (record.source_record.as_str(), record)) - .collect::>(); - let mut source_counts = BTreeMap::new(); - for record in stores { - let source_record = record.source_record.as_deref().unwrap(); - *source_counts.entry(source_record).or_insert(0) += 1; - let declaration = declarations[source_record]; - validate_imported_policy(record, declaration).unwrap(); - assert_tcgen05_backend_target_split(record); - assert_eq!(record.rust_arguments.len(), 2); - assert_eq!(record.rust_result, "()"); - assert!(!record.must_use); - assert!(!record.safe); - assert!(!record.pure); - assert_eq!(record.memory, "write"); - assert!(record.convergent); - assert_eq!(record.execution_scope, "warp"); - assert_eq!(record.dialect_results, Vec::::new()); - assert_eq!(record.llvm_results, Vec::::new()); - assert_eq!(record.llvm_arguments, declaration.arguments); - assert_eq!( - declaration.properties, - [ - "ImmArg", - "IntrArgMemOnly", - "IntrConvergent", - "NoCapture", - ] - ); - assert!(declaration.selections.is_empty()); - } - assert_eq!(source_counts.len(), 29); - assert!(source_counts.values().all(|count| *count == 2)); - validate_unique_overlay(&records, 1).unwrap(); - - let scalar_raw = &stores[0]; - assert_eq!(scalar_raw.rust_arguments, ["u32", "u32"]); - assert_eq!(scalar_raw.dialect_operands, ["i32", "i32"]); - assert_eq!(scalar_raw.llvm_arguments, ["tmem_ptr", "i32", "i1"]); - assert_eq!( - scalar_raw.expected_ptx.modifiers, - ["st", "sync", "aligned", "16x64b", "x1", "b32"] - ); - assert_eq!( - scalar_raw.expected_ptx.operands, - [ - OperandPattern::Address, - OperandPattern::RegisterList { length: 1 }, - ] - ); - assert!( - scalar_raw - .expected_ptx - .matches("tcgen05.st.sync.aligned.16x64b.x1.b32 [%r1], {%r2};") - .unwrap() - ); - assert!( - !scalar_raw - .expected_ptx - .matches("tcgen05.st.sync.aligned.16x64b.x1.b32 [%r1], %r2;") - .unwrap() - ); - assert_eq!( - stores[1].expected_ptx.modifiers, - [ - "st", - "sync", - "aligned", - "16x64b", - "x1", - "unpack::16b", - "b32", - ] - ); - let largest = &stores[57]; - assert_eq!(largest.rust_arguments, ["u32", "[u32; 128]"]); - assert_eq!(largest.dialect_operands.len(), 129); - assert_eq!(largest.llvm_arguments, ["tmem_ptr", "v128i32", "i1"]); - assert_eq!( - largest.expected_ptx.operands, - [ - OperandPattern::Address, - OperandPattern::RegisterList { length: 128 }, - ] - ); - - let mut missing = test_tcgen05_st_admission(); - missing.st_variants.pop(); - assert!(expand_tcgen05_admission(&missing).is_err()); - - let mut reordered = test_tcgen05_st_admission(); - reordered.st_variants.swap(0, 1); - assert!(expand_tcgen05_admission(&reordered).is_err()); - - let mut non_contiguous_abi = test_tcgen05_st_admission(); - non_contiguous_abi.st_variants[0].abi_id = "i9999".into(); - let records = expand_tcgen05_admission(&non_contiguous_abi).unwrap(); - assert!(records.iter().any(|record| record.abi_id == "i9999")); - - let mut missing_evidence = test_tcgen05_st_admission(); - missing_evidence.st_llvm_evidence_profile = None; - assert!(expand_tcgen05_admission(&missing_evidence).is_err()); - - let declaration = declarations[stores[0].source_record.as_deref().unwrap()]; - let mut wrong_selector = stores[0].clone(); - wrong_selector - .tcgen05 - .as_mut() - .unwrap() - .st - .as_mut() - .unwrap() - .unpack16 = true; - assert!(validate_imported_policy(&wrong_selector, declaration).is_err()); - - let mut changed_declaration = declaration.clone(); - changed_declaration.properties.pop(); - assert!(validate_imported_policy(&stores[0], &changed_declaration).is_err()); - - let mut wrong_scope = stores[0].clone(); - wrong_scope.execution_scope = "thread".into(); - assert!(validate_imported_policy(&wrong_scope, declaration).is_err()); - - let mut unreviewed_sharing = records.clone(); - unreviewed_sharing[120].tcgen05.as_mut().unwrap().operation = Tcgen05Operation::Alloc; - assert!(validate_unique_overlay(&unreviewed_sharing, 1).is_err()); - } - - #[test] - fn compact_tcgen05_offset_admission_is_exact_and_fails_closed() { - let records = expand_tcgen05_admission(&test_tcgen05_offset_admission()).unwrap(); - assert_eq!(records.len(), 209); - let loads = &records[177..193]; - let stores = &records[193..209]; - assert_eq!( - (loads[0].abi_id.as_str(), loads[0].id.as_str()), - ("i0728", "tcgen05_ld_16x32bx2_x1_raw") - ); - assert_eq!( - (loads[15].abi_id.as_str(), loads[15].id.as_str()), - ("i0743", "tcgen05_ld_16x32bx2_x128_pack16") - ); - assert_eq!( - (stores[0].abi_id.as_str(), stores[0].id.as_str()), - ("i0744", "tcgen05_st_16x32bx2_x1_raw") - ); - assert_eq!( - (stores[15].abi_id.as_str(), stores[15].id.as_str()), - ("i0759", "tcgen05_st_16x32bx2_x128_unpack16") - ); - - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - let declarations = imported - .intrinsics - .iter() - .map(|record| (record.source_record.as_str(), record)) - .collect::>(); - let mut load_sources = BTreeMap::new(); - for record in loads { - let source = record.source_record.as_deref().unwrap(); - *load_sources.entry(source).or_insert(0) += 1; - let declaration = declarations[source]; - validate_imported_policy(record, declaration).unwrap(); - assert_tcgen05_backend_target_split(record); - assert_eq!(record.rust_arguments, ["u32", "i64"]); - assert_eq!(record.dialect_operands, ["i32", "i64"]); - assert_eq!(record.llvm_arguments, ["tmem_ptr", "i64", "i1"]); - assert_eq!(record.execution_scope, "warp"); - assert!(record.must_use); - assert_eq!( - record.compatibility_rust_paths, - [format!("cuda_device::tcgen05::__{}", record.id)] - ); - assert_eq!( - record.expected_ptx.operands.last(), - Some(&OperandPattern::Immediate) - ); - assert_eq!( - declaration.properties, - [ - "ImmArg", - "ImmArg", - "IntrArgMemOnly", - "IntrConvergent", - "NoCapture", - ] - ); - } - assert_eq!(load_sources.len(), 8); - assert!(load_sources.values().all(|count| *count == 2)); - - let mut store_sources = BTreeMap::new(); - for record in stores { - let source = record.source_record.as_deref().unwrap(); - *store_sources.entry(source).or_insert(0) += 1; - let declaration = declarations[source]; - validate_imported_policy(record, declaration).unwrap(); - assert_tcgen05_backend_target_split(record); - assert_eq!(record.rust_arguments[0..2], ["u32", "i64"]); - assert_eq!(record.dialect_operands[0..2], ["i32", "i64"]); - assert_eq!(record.llvm_arguments[0..2], ["tmem_ptr", "i64"]); - assert_eq!(record.execution_scope, "warp"); - assert!(!record.must_use); - assert_eq!( - record.compatibility_rust_paths, - [format!("cuda_device::tcgen05::__{}", record.id)] - ); - assert_eq!( - record.expected_ptx.operands.get(1), - Some(&OperandPattern::Immediate) - ); - assert_eq!( - declaration.properties, - [ - "ImmArg", - "ImmArg", - "IntrArgMemOnly", - "IntrConvergent", - "NoCapture", - ] - ); - } - assert_eq!(store_sources.len(), 8); - assert!(store_sources.values().all(|count| *count == 2)); - validate_unique_overlay(&records, 1).unwrap(); - - let scalar_pack = &loads[1]; - assert_eq!( - scalar_pack.expected_ptx.operands, - [ - OperandPattern::RegisterList { length: 1 }, - OperandPattern::Address, - OperandPattern::Immediate, - ] - ); - assert!( - scalar_pack - .expected_ptx - .matches("tcgen05.ld.sync.aligned.16x32bx2.x1.pack::16b.b32 {%r1}, [%r2], 16;") - .unwrap() - ); - - let scalar_unpack = &stores[1]; - assert_eq!( - scalar_unpack.expected_ptx.operands, - [ - OperandPattern::Address, - OperandPattern::Immediate, - OperandPattern::RegisterList { length: 1 }, - ] - ); - assert!( - scalar_unpack - .expected_ptx - .matches("tcgen05.st.sync.aligned.16x32bx2.x1.unpack::16b.b32 [%r1], 16, {%r2};") - .unwrap() - ); - - let mut missing = test_tcgen05_offset_admission(); - missing.ld_offset_variants.pop(); - assert!(expand_tcgen05_admission(&missing).is_err()); - - let mut reordered = test_tcgen05_offset_admission(); - reordered.st_offset_variants.swap(0, 1); - assert!(expand_tcgen05_admission(&reordered).is_err()); - - let mut non_contiguous_abi = test_tcgen05_offset_admission(); - non_contiguous_abi.ld_offset_variants[0].abi_id = "i9999".into(); - let records = expand_tcgen05_admission(&non_contiguous_abi).unwrap(); - assert!(records.iter().any(|record| record.abi_id == "i9999")); - - let mut missing_evidence = test_tcgen05_offset_admission(); - missing_evidence.offset_llvm_evidence_profile = None; - assert!(expand_tcgen05_admission(&missing_evidence).is_err()); - - let load_declaration = declarations[loads[0].source_record.as_deref().unwrap()]; - let mut wrong_offset_type = loads[0].clone(); - wrong_offset_type.rust_arguments[1] = "i32".into(); - assert!(validate_imported_policy(&wrong_offset_type, load_declaration).is_err()); - - let mut wrong_scope = loads[0].clone(); - wrong_scope.execution_scope = "thread".into(); - assert!(validate_imported_policy(&wrong_scope, load_declaration).is_err()); - - let mut missing_immediate = loads[0].clone(); - missing_immediate.expected_ptx.operands.pop(); - assert!(validate_imported_policy(&missing_immediate, load_declaration).is_err()); - - let store_declaration = declarations[stores[0].source_record.as_deref().unwrap()]; - let mut wrong_adapter = stores[0].clone(); - wrong_adapter.tcgen05.as_mut().unwrap().adapter = - Tcgen05Adapter::TmemU32RegistersInjectUnpack16ToVoid; - assert!(validate_imported_policy(&wrong_adapter, store_declaration).is_err()); - - let mut changed_declaration = store_declaration.clone(); - changed_declaration.properties.remove(0); - assert!(validate_imported_policy(&stores[0], &changed_declaration).is_err()); - } - - #[test] - fn tcgen05_compact_schema_is_reserved_for_aggregation() { - let shard = |schema, admission| OverlayShardFile { - schema, - family: "tcgen05".into(), - intrinsics: vec![], - register_mma_int4: None, - register_mma_int8: None, - register_mma_b1: None, - register_mma_f8f6f4_f32: None, - register_mma_f8f6f4_f16: None, - register_mma_mxf8f6f4_f32: None, - register_mma_fp8: None, - register_mma_ampere_float: None, - sparse_mma_integer: None, - sparse_mma_f8f6f4_f32: None, - sparse_mma_f8f6f4_f16: None, - prmt: None, - packed_conversion_fp8: None, - packed_conversion_fp8_f16x2: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cluster_sreg: None, - cluster_barrier: None, - mbarrier_extended: None, - special_registers: None, - debug_control: None, - threadfence: None, - cluster_memory: None, - stmatrix: None, - clc: None, - wgmma_controls: None, - tma: None, - tcgen05: Some(admission), - }; - let path = Path::new("intrinsics/overlay/tcgen05.toml"); - validate_overlay_shard_schema_with_max( - &shard( - TCGEN05_SHARD_SCHEMA, - without_tcgen05_control(test_tcgen05_admission()), - ), - path, - TCGEN05_CP_SHARD_SCHEMA, - ) - .unwrap(); - assert!( - validate_overlay_shard_schema_with_max( - &shard( - TCGEN05_SHARD_SCHEMA - 1, - without_tcgen05_control(test_tcgen05_admission()), - ), - path, - TCGEN05_CP_SHARD_SCHEMA, - ) - .unwrap_err() - .to_string() - .contains("requires overlay shard schema 42") - ); - validate_overlay_shard_schema_with_max( - &shard( - TCGEN05_CP_SHARD_SCHEMA, - without_tcgen05_control(test_tcgen05_cp_admission()), - ), - path, - TCGEN05_CP_SHARD_SCHEMA, - ) - .unwrap(); - assert!( - validate_overlay_shard_schema_with_max( - &shard( - TCGEN05_CP_SHARD_SCHEMA - 1, - without_tcgen05_control(test_tcgen05_cp_admission()), - ), - path, - TCGEN05_CP_SHARD_SCHEMA, - ) - .unwrap_err() - .to_string() - .contains("requires overlay shard schema 52") - ); - validate_overlay_shard_schema_with_max( - &shard( - TCGEN05_LD_SHARD_SCHEMA, - without_tcgen05_control(test_tcgen05_ld_admission()), - ), - path, - TCGEN05_LD_SHARD_SCHEMA, - ) - .unwrap(); - assert!( - validate_overlay_shard_schema_with_max( - &shard( - TCGEN05_LD_SHARD_SCHEMA - 1, - without_tcgen05_control(test_tcgen05_ld_admission()), - ), - path, - TCGEN05_LD_SHARD_SCHEMA, - ) - .unwrap_err() - .to_string() - .contains("requires overlay shard schema 53") - ); - validate_overlay_shard_schema_with_max( - &shard( - TCGEN05_ST_SHARD_SCHEMA, - without_tcgen05_control(test_tcgen05_st_admission()), - ), - path, - TCGEN05_ST_SHARD_SCHEMA, - ) - .unwrap(); - assert!( - validate_overlay_shard_schema_with_max( - &shard( - TCGEN05_ST_SHARD_SCHEMA - 1, - without_tcgen05_control(test_tcgen05_st_admission()), - ), - path, - TCGEN05_ST_SHARD_SCHEMA, - ) - .unwrap_err() - .to_string() - .contains("requires overlay shard schema 54") - ); - validate_overlay_shard_schema_with_max( - &shard( - TCGEN05_OFFSET_LDST_SHARD_SCHEMA, - without_tcgen05_control(test_tcgen05_offset_admission()), - ), - path, - TCGEN05_OFFSET_LDST_SHARD_SCHEMA, - ) - .unwrap(); - assert!( - validate_overlay_shard_schema_with_max( - &shard( - TCGEN05_OFFSET_LDST_SHARD_SCHEMA - 1, - without_tcgen05_control(test_tcgen05_offset_admission()), - ), - path, - TCGEN05_OFFSET_LDST_SHARD_SCHEMA, - ) - .unwrap_err() - .to_string() - .contains("requires overlay shard schema 55") - ); - validate_overlay_shard_schema_with_max( - &shard( - TCGEN05_CONTROL_SHARD_SCHEMA, - test_tcgen05_offset_admission(), - ), - path, - TCGEN05_CONTROL_SHARD_SCHEMA, - ) - .unwrap(); - assert!( - validate_overlay_shard_schema_with_max( - &shard( - TCGEN05_CONTROL_SHARD_SCHEMA, - without_tcgen05_control(test_tcgen05_offset_admission()), - ), - path, - TCGEN05_CONTROL_SHARD_SCHEMA, - ) - .unwrap_err() - .to_string() - .contains("requires all three control variants and both backend evidence profiles") - ); - assert!( - validate_overlay_shard_schema_with_max( - &shard( - TCGEN05_CONTROL_SHARD_SCHEMA - 1, - test_tcgen05_offset_admission(), - ), - path, - TCGEN05_CONTROL_SHARD_SCHEMA, - ) - .unwrap_err() - .to_string() - .contains("requires overlay shard schema 56") - ); - validate_overlay_shard_schema_with_max( - &shard(TCGEN05_MMA_SHARD_SCHEMA, test_tcgen05_mma_admission()), - path, - TCGEN05_MMA_SHARD_SCHEMA, - ) - .unwrap(); - assert!( - validate_overlay_shard_schema_with_max( - &shard(TCGEN05_MMA_SHARD_SCHEMA - 1, test_tcgen05_mma_admission(),), - path, - TCGEN05_MMA_SHARD_SCHEMA, - ) - .unwrap_err() - .to_string() - .contains("requires overlay shard schema 57") - ); - } - - #[test] - fn debug_control_compact_schema_is_reserved_for_aggregation() { - let shard = |schema| OverlayShardFile { - schema, - family: "debug_control".into(), - intrinsics: vec![], - register_mma_int4: None, - register_mma_int8: None, - register_mma_b1: None, - register_mma_f8f6f4_f32: None, - register_mma_f8f6f4_f16: None, - register_mma_mxf8f6f4_f32: None, - register_mma_fp8: None, - register_mma_ampere_float: None, - sparse_mma_integer: None, - sparse_mma_f8f6f4_f32: None, - sparse_mma_f8f6f4_f16: None, - prmt: None, - packed_conversion_fp8: None, - packed_conversion_fp8_f16x2: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cluster_sreg: None, - cluster_barrier: None, - mbarrier_extended: None, - special_registers: None, - debug_control: Some(test_debug_control_admission()), - threadfence: None, - cluster_memory: None, - stmatrix: None, - clc: None, - wgmma_controls: None, - tma: None, - tcgen05: None, - }; - let path = Path::new("intrinsics/overlay/debug_control.toml"); - validate_overlay_shard_schema_with_max(&shard(33), path, 33).unwrap(); - assert!(validate_overlay_shard_schema_with_max(&shard(33), path, 32).is_err()); - let error = validate_overlay_shard_schema_with_max(&shard(32), path, 33).unwrap_err(); - assert!( - error - .to_string() - .contains("requires overlay shard schema 33") - ); - } - - #[test] - fn cluster_memory_compact_schema_is_reserved_and_fail_closed() { - let shard = |schema| OverlayShardFile { - schema, - family: "cluster_memory".into(), - intrinsics: vec![], - register_mma_int4: None, - register_mma_int8: None, - register_mma_b1: None, - register_mma_f8f6f4_f32: None, - register_mma_f8f6f4_f16: None, - register_mma_mxf8f6f4_f32: None, - register_mma_fp8: None, - register_mma_ampere_float: None, - sparse_mma_integer: None, - sparse_mma_f8f6f4_f32: None, - sparse_mma_f8f6f4_f16: None, - prmt: None, - packed_conversion_fp8: None, - packed_conversion_fp8_f16x2: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cluster_sreg: None, - cluster_barrier: None, - mbarrier_extended: None, - special_registers: None, - debug_control: None, - threadfence: None, - cluster_memory: Some(test_cluster_memory_admission()), - stmatrix: None, - clc: None, - wgmma_controls: None, - tma: None, - tcgen05: None, - }; - let path = Path::new("intrinsics/overlay/cluster_memory.toml"); - validate_overlay_shard_schema_with_max( - &shard(CLUSTER_MEMORY_SHARD_SCHEMA), - path, - CLUSTER_MEMORY_SHARD_SCHEMA, - ) - .unwrap(); - assert!( - validate_overlay_shard_schema_with_max( - &shard(CLUSTER_MEMORY_SHARD_SCHEMA), - path, - CLUSTER_MEMORY_SHARD_SCHEMA - 1, - ) - .is_err() - ); - let error = validate_overlay_shard_schema_with_max( - &shard(CLUSTER_MEMORY_SHARD_SCHEMA - 1), - path, - CLUSTER_MEMORY_SHARD_SCHEMA, - ) - .unwrap_err(); - assert!( - error - .to_string() - .contains("requires overlay shard schema 39") - ); - } - - #[test] - fn cluster_memory_admission_preserves_mapa_identity_and_ptx_native_read() { - let admission = test_cluster_memory_admission(); - let records = expand_cluster_memory_admission(&admission).unwrap(); - assert_eq!(records.len(), 2); - assert_eq!(records[0].id, "map_shared_rank"); - assert_eq!(records[1].id, "dsmem_read_u32"); - assert_eq!(records[0].abi_id, "i0320"); - assert_eq!(records[1].abi_id, "i0321"); - assert_eq!( - cluster_memory_inline_recipe(ClusterMemoryOperation::MapSharedRank), - ("mapa.shared::cluster.u64 $0, $1, $2;", "=l,l,r") - ); - assert_eq!( - cluster_memory_inline_recipe(ClusterMemoryOperation::ReadU32), - ( - "{ .reg .u64 %mapped; mapa.shared::cluster.u64 %mapped, $1, $2; ld.shared::cluster.u32 $0, [%mapped]; }", - "=r,l,r,~{memory}" - ) - ); - - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - let mapa = &records[0]; - let declaration = imported - .intrinsics - .iter() - .find(|declaration| declaration.source_record == "int_nvvm_mapa_shared_cluster") - .unwrap(); - assert_eq!(declaration.arguments, ["shared_ptr", "i32"]); - assert_eq!(declaration.results, ["shared_cluster_ptr"]); - assert_eq!( - declaration.properties, - ["IntrNoMem", "IntrSpeculatable", "NoCapture"] - ); - validate_imported_policy(mapa, declaration).unwrap(); - assert_eq!( - declaration - .selections - .iter() - .filter(|selection| selection.asm.starts_with("mapa.shared::cluster.u64")) - .map(|selection| selection.source_record.as_str()) - .collect::>(), - BTreeSet::from(["mapa_shared_cluster_64", "mapa_shared_cluster_64i"]) - ); - - let read = &records[1]; - validate_ptx_native_policy(read).unwrap(); - assert!(read.source_record.is_none()); - assert!(read.llvm_symbol.is_none()); - assert!(matches!( - resolve_policy_source(read).unwrap(), - IntrinsicSource::PtxNative { .. } - )); - assert_eq!(read.memory, "read"); - - let mut wrong_adapter = mapa.clone(); - wrong_adapter.cluster_memory.as_mut().unwrap().adapter = - ClusterMemoryAdapter::ConstU32PointerRankToU32; - assert!(validate_imported_policy(&wrong_adapter, declaration).is_err()); - - let mut typed_as3 = mapa.clone(); - typed_as3.llvm_results = vec!["shared_ptr".into()]; - assert!(validate_imported_policy(&typed_as3, declaration).is_err()); - - let mut wrong_route = mapa.clone(); - wrong_route.backend_lowerings[0].mechanism = BackendLoweringMechanism::TypedNvvm; - assert!(validate_imported_policy(&wrong_route, declaration).is_err()); - - let mut wrong_floor = read.clone(); - wrong_floor.minimum_sm = Some("sm_80".into()); - assert!(validate_ptx_native_policy(&wrong_floor).is_err()); - - let mut missing = admission.clone(); - missing.variants.pop(); - assert!(expand_cluster_memory_admission(&missing).is_err()); - - let mut duplicate = admission.clone(); - duplicate.variants[1].operation = ClusterMemoryOperation::MapSharedRank; - assert!(expand_cluster_memory_admission(&duplicate).is_err()); - - let mut wrong_abi = admission.clone(); - wrong_abi.variants[0].abi_id = "i9999".into(); - assert!(expand_cluster_memory_admission(&wrong_abi).is_err()); - - let mut executed = admission; - executed.runtime_validation = RuntimeValidation::Executed; - assert!(expand_cluster_memory_admission(&executed).is_err()); - } - - #[test] - fn active_debug_control_sources_parse_and_prove_both_backend_routes() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let shard_path = repo_root.join("intrinsics/overlay/debug_control.toml"); - let shard: OverlayShardFile = - toml::from_str(&fs::read_to_string(&shard_path).unwrap()).unwrap(); - validate_overlay_shard_schema_with_max(&shard, &shard_path, 33).unwrap(); - let admission = shard.debug_control.unwrap(); - assert_eq!(admission.abi_ids, ["i0295", "i0296", "i0297"]); - let records = expand_debug_control_admission(&admission).unwrap(); - - let evidence = vec![ - read_evidence_file( - &repo_root.join("intrinsics/evidence/rust-llvm-22.1.2-1cb4e383-debug-control.json"), - ) - .unwrap(), - read_evidence_file( - &repo_root.join("intrinsics/evidence/cuda-13.3-libnvvm-13.3.33-debug-control.json"), - ) - .unwrap(), - ]; - let indexed = - index_evidence(&evidence, "1cb4e3833c1919c2e6fb579a23ac0e2b22587b7e").unwrap(); - for record in &records { - let routes = resolve_backend_lowerings(record, &indexed).unwrap(); - assert_eq!(routes.len(), 2); - assert!(routes.iter().all(|route| { - route.mechanism == BackendLoweringMechanism::InlinePtx - && route.status == "validated" - })); - } - } - - #[test] - fn compact_cluster_barrier_admission_and_semantics_fail_closed() { - let records = expand_cluster_barrier_admission(&test_cluster_barrier_admission()).unwrap(); - assert_eq!(records.len(), 6); - - let imported: ImportedFile = read_json( - &Path::new(env!("CARGO_MANIFEST_DIR")) - .join("../..") - .join("intrinsics/imported.json"), - ) - .unwrap(); - for record in &records { - let declaration = imported - .intrinsics - .iter() - .find(|declaration| { - Some(declaration.source_record.as_str()) == record.source_record.as_deref() - }) - .unwrap(); - validate_imported_policy(record, declaration).unwrap(); - } - - let declaration_for = |record: &OverlayIntrinsic| { - imported - .intrinsics - .iter() - .find(|declaration| { - Some(declaration.source_record.as_str()) == record.source_record.as_deref() - }) - .unwrap() - }; - let base = records - .iter() - .find(|record| { - record - .cluster_barrier - .as_ref() - .is_some_and(|barrier| barrier.mode == ClusterBarrierMode::ArriveAligned) - }) - .unwrap(); - - let mut wrong_mode = base.clone(); - wrong_mode.cluster_barrier.as_mut().unwrap().mode = ClusterBarrierMode::WaitAligned; - assert!(validate_imported_policy(&wrong_mode, declaration_for(base)).is_err()); - - let mut wrong_order = base.clone(); - wrong_order.cluster_barrier.as_mut().unwrap().ordering = ClusterBarrierOrdering::Relaxed; - assert!(validate_imported_policy(&wrong_order, declaration_for(base)).is_err()); - - let mut wrong_alignment = base.clone(); - wrong_alignment.cluster_barrier.as_mut().unwrap().aligned = false; - assert!(validate_imported_policy(&wrong_alignment, declaration_for(base)).is_err()); - - let mut missing = test_cluster_barrier_admission(); - missing.variants.pop(); - assert!(expand_cluster_barrier_admission(&missing).is_err()); - - let mut duplicate = test_cluster_barrier_admission(); - duplicate.variants[5].mode = ClusterBarrierMode::Arrive; - assert!(expand_cluster_barrier_admission(&duplicate).is_err()); - - let mut wrong_abi = test_cluster_barrier_admission(); - wrong_abi.variants[0].abi_id = "i9999".into(); - assert!(expand_cluster_barrier_admission(&wrong_abi).is_err()); - } - - #[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 - .iter() - .map(|record| record.abi_id.as_str()) - .collect::>(), - (306..=316) - .map(|id| format!("i{id:04}")) - .collect::>() - ); - - let imported: ImportedFile = read_json( - &Path::new(env!("CARGO_MANIFEST_DIR")) - .join("../..") - .join("intrinsics/imported.json"), - ) - .unwrap(); - for record in &records { - let contract = record.mbarrier_extended.as_ref().unwrap(); - let (template, constraints) = mbarrier_extended_inline_recipe(contract.operation); - assert!(template.ends_with(';') || template.ends_with("; }")); - assert!(constraints.contains("~{memory}")); - assert!(record.convergent && record.memory == "read_write"); - assert!( - record - .backend_lowerings - .iter() - .all(|lowering| { lowering.mechanism == BackendLoweringMechanism::InlinePtx }) - ); - match contract.source_contract { - MbarrierExtendedSourceContract::LlvmImported => { - let declaration = imported - .intrinsics - .iter() - .find(|declaration| { - Some(declaration.source_record.as_str()) - == record.source_record.as_deref() - }) - .unwrap(); - validate_imported_policy(record, declaration).unwrap(); - } - MbarrierExtendedSourceContract::PtxNativeRawClusterAddress => { - validate_ptx_native_policy(record).unwrap(); - assert_eq!(record.rust_arguments, ["u64"]); - assert_eq!(record.dialect_operands, ["i64"]); - } - } - } - - let base = records - .iter() - .find(|record| record.id == "mbarrier_arrive_expect_tx") - .unwrap(); - let declaration = imported - .intrinsics - .iter() - .find(|declaration| { - Some(declaration.source_record.as_str()) == base.source_record.as_deref() - }) - .unwrap(); - let mut wrong_adapter = base.clone(); - wrong_adapter.mbarrier_extended.as_mut().unwrap().adapter = - MbarrierExtendedAdapter::PointerTokenToPredicate; - assert!(validate_imported_policy(&wrong_adapter, declaration).is_err()); - let mut wrong_floor = base.clone(); - wrong_floor.minimum_ptx = "8.6".into(); - assert!(validate_imported_policy(&wrong_floor, declaration).is_err()); - let mut lost_clobber = base.clone(); - lost_clobber.memory = "none".into(); - assert!(validate_imported_policy(&lost_clobber, declaration).is_err()); - - let remote = records - .iter() - .find(|record| record.id == "mbarrier_arrive_cluster") - .unwrap(); - let incompatible = imported - .intrinsics - .iter() - .find(|declaration| { - declaration.source_record == "int_nvvm_mbarrier_arrive_scope_cluster_space_cluster" - }) - .unwrap(); - assert_eq!(incompatible.arguments, ["shared_cluster_ptr", "i32"]); - assert!( - validate_mbarrier_extended_policy( - remote, - &IntrinsicSource::LlvmImported { - source_record: incompatible.source_record.clone(), - }, - Some(incompatible), - ) - .is_err() - ); - - let mut missing = test_mbarrier_extended_admission(); - missing.variants.pop(); - assert!(expand_mbarrier_extended_admission(&missing).is_err()); - let mut duplicate = test_mbarrier_extended_admission(); - duplicate.variants[10].operation = MbarrierExtendedOperation::ArriveExpectTxCta; - assert!(expand_mbarrier_extended_admission(&duplicate).is_err()); - let mut wrong_abi = test_mbarrier_extended_admission(); - wrong_abi.variants[0].abi_id = "i9999".into(); - assert!(expand_mbarrier_extended_admission(&wrong_abi).is_err()); - } - - #[test] - fn compact_wgmma_control_admission_and_semantics_fail_closed() { - let records = expand_wgmma_control_admission(&test_wgmma_control_admission()).unwrap(); - assert_eq!(records.len(), 3); - - let imported: ImportedFile = read_json( - &Path::new(env!("CARGO_MANIFEST_DIR")) - .join("../..") - .join("intrinsics/imported.json"), - ) - .unwrap(); - let declaration_for = |record: &OverlayIntrinsic| { - imported - .intrinsics - .iter() - .find(|declaration| { - Some(declaration.source_record.as_str()) == record.source_record.as_deref() - }) - .unwrap() - }; - for record in &records { - validate_imported_policy(record, declaration_for(record)).unwrap(); - assert_eq!(record.targets, "sm_90a"); - assert_eq!(record.minimum_ptx, "8.0"); - assert_eq!(record.backend_lowerings.len(), 2); - } - - let wait = records - .iter() - .find(|record| { - record - .wgmma_control - .as_ref() - .is_some_and(|control| control.mode == WgmmaControlMode::WaitGroup) - }) - .unwrap(); - assert_eq!(wait.rust_arguments, ["u64"]); - assert_eq!(wait.dialect_operands, ["i64"]); - assert_eq!( - wait.compatibility_rust_paths, - ["cuda_device::wgmma::__wgmma_wait_group"] - ); - - let mut wrong_adapter = wait.clone(); - wrong_adapter.wgmma_control.as_mut().unwrap().adapter = WgmmaControlAdapter::NoArguments; - assert!(validate_imported_policy(&wrong_adapter, declaration_for(wait)).is_err()); - - let mut wrong_participation = wait.clone(); - wrong_participation.execution_scope = "warp".into(); - assert!(validate_imported_policy(&wrong_participation, declaration_for(wait)).is_err()); - - let mut wrong_route = wait.clone(); - wrong_route.backend_lowerings[1].mechanism = BackendLoweringMechanism::TypedNvvm; - assert!(validate_imported_policy(&wrong_route, declaration_for(wait)).is_err()); - - let mut wrong_target = wait.clone(); - wrong_target.targets = "all".into(); - wrong_target.minimum_sm = Some("sm_90".into()); - assert!(validate_imported_policy(&wrong_target, declaration_for(wait)).is_err()); - - let mut missing = test_wgmma_control_admission(); - missing.variants.pop(); - assert!(expand_wgmma_control_admission(&missing).is_err()); - - let mut reversed = test_wgmma_control_admission(); - reversed.variants.reverse(); - assert!(expand_wgmma_control_admission(&reversed).is_err()); - - let mut duplicate = test_wgmma_control_admission(); - duplicate.variants[2].mode = WgmmaControlMode::Fence; - assert!(expand_wgmma_control_admission(&duplicate).is_err()); - - let mut wrong_abi = test_wgmma_control_admission(); - wrong_abi.variants[0].abi_id = "i9999".into(); - assert!(expand_wgmma_control_admission(&wrong_abi).is_err()); - - let mut executed = test_wgmma_control_admission(); - executed.runtime_validation = RuntimeValidation::Executed; - assert!(expand_wgmma_control_admission(&executed).is_err()); - } - - #[test] - fn wgmma_control_compact_schema_is_reserved_for_aggregation() { - let shard = OverlayShardFile { - schema: WGMMA_CONTROL_SHARD_SCHEMA, - family: "wgmma_control".into(), - intrinsics: vec![], - register_mma_int4: None, - register_mma_int8: None, - register_mma_b1: None, - register_mma_f8f6f4_f32: None, - register_mma_f8f6f4_f16: None, - register_mma_mxf8f6f4_f32: None, - register_mma_fp8: None, - register_mma_ampere_float: None, - sparse_mma_integer: None, - sparse_mma_f8f6f4_f32: None, - sparse_mma_f8f6f4_f16: None, - prmt: None, - packed_conversion_fp8: None, - packed_conversion_fp8_f16x2: None, - scalar_conversion: None, - scalar_arithmetic: None, - scalar_math: None, - extended_minmax: None, - cluster_sreg: None, - cluster_barrier: None, - mbarrier_extended: None, - special_registers: None, - debug_control: None, - threadfence: None, - cluster_memory: None, - stmatrix: None, - clc: None, - wgmma_controls: Some(test_wgmma_control_admission()), - tma: None, - tcgen05: None, - }; - let path = Path::new("intrinsics/overlay/wgmma_control.toml"); - validate_overlay_shard_schema_with_max(&shard, path, WGMMA_CONTROL_SHARD_SCHEMA).unwrap(); - let mut old = shard; - old.schema -= 1; - assert!( - validate_overlay_shard_schema_with_max(&old, path, WGMMA_CONTROL_SHARD_SCHEMA) - .unwrap_err() - .to_string() - .contains("requires overlay shard schema 38") - ); - } - - #[test] - fn cluster_barrier_evidence_validates_both_backend_routes() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let mut admission = test_cluster_barrier_admission(); - admission.llvm_evidence_profile = "rust-llvm-22.1.2-1cb4e383".into(); - admission.libnvvm_evidence_profile = "cuda-13.3-libnvvm-13.3.33-cluster-barrier".into(); - let policies = expand_cluster_barrier_admission(&admission).unwrap(); - let evidence_files = vec![ - read_evidence_file( - &repo_root - .join("intrinsics/evidence/rust-llvm-22.1.2-1cb4e383-cluster-barrier.json"), - ) - .unwrap(), - read_evidence_file( - &repo_root - .join("intrinsics/evidence/cuda-13.3-libnvvm-13.3.33-cluster-barrier.json"), - ) - .unwrap(), - ]; - let indexed = - index_evidence(&evidence_files, "1cb4e3833c1919c2e6fb579a23ac0e2b22587b7e").unwrap(); - - for policy in &policies { - for lowering in &policy.backend_lowerings { - let evidence = indexed - .get(&(lowering.evidence_profile.as_str(), policy.id.as_str())) - .unwrap(); - validate_evidence(policy, evidence, Some(lowering)).unwrap(); - } - } - - let mut missing_typed_failure = evidence_files.clone(); - let libnvvm = missing_typed_failure - .iter_mut() - .find(|file| file.backend_kind == Some(IntrinsicBackend::LibNvvm)) - .unwrap(); - for record in &mut libnvvm.records { - record.stages.retain(|stage| { - stage.mechanism != Some(BackendLoweringMechanism::TypedNvvm) - || stage.stage != EvidenceStageKind::DeviceLink - }); - } - let indexed = index_evidence( - &missing_typed_failure, - "1cb4e3833c1919c2e6fb579a23ac0e2b22587b7e", - ) - .unwrap(); - let policy = &policies[0]; - let lowering = policy - .backend_lowerings - .iter() - .find(|lowering| lowering.backend == IntrinsicBackend::LibNvvm) - .unwrap(); - let evidence = indexed - .get(&(lowering.evidence_profile.as_str(), policy.id.as_str())) - .unwrap(); - assert!(validate_evidence(policy, evidence, Some(lowering)).is_err()); - } - - #[test] - fn pinned_register_mma_records_match_the_closed_recipes_and_fail_closed() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let (mut overlay, _) = - read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")).unwrap(); - bind_pinned_abi_ids(&repo_root, &mut overlay); - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - let declarations: BTreeMap<_, _> = imported - .intrinsics - .iter() - .map(|record| (record.source_record.as_str(), record)) - .collect(); - let records: Vec<_> = overlay - .intrinsics - .iter() - .filter(|record| record.family == "register_mma") - .collect(); - assert_eq!(records.len(), 154); - - let dense_f8f6f4_records = records - .iter() - .copied() - .filter(|record| { - record.register_mma.as_ref().is_some_and(|mma| { - matches!(mma.kind, None | Some(RegisterMmaKind::F8f6f4)) - && register_mma_f8f6f4_element_name(mma.a_element).is_some() - && register_mma_f8f6f4_element_name(mma.b_element).is_some() - }) - }) - .collect::>(); - assert_eq!(dense_f8f6f4_records.len(), 50); - for (accumulator, first_abi, last_abi, adapter) in [ - ( - RegisterMmaAccumulator::F32, - "i0454", - "i0478", - RegisterMmaAdapter::C4F32A4U32B2U32ToD4F32, - ), - ( - RegisterMmaAccumulator::F16, - "i0479", - "i0503", - RegisterMmaAdapter::C2U32A4U32B2U32ToD2U32, - ), - ] { - let family = dense_f8f6f4_records - .iter() - .copied() - .filter(|record| record.register_mma.as_ref().unwrap().accumulator == accumulator) - .collect::>(); - assert_eq!(family.len(), 25); - assert_eq!(family[0].abi_id, first_abi); - assert_eq!(family[24].abi_id, last_abi); - assert!( - family - .iter() - .all(|record| record.register_mma.as_ref().unwrap().adapter == adapter) - ); - } - assert!(dense_f8f6f4_records.iter().all(|record| { - record.minimum_ptx == "8.7" - && record.minimum_sm.is_none() - && record.targets == REGISTER_MMA_F8F6F4_TARGETS - })); - - let mxf8f6f4_records = records - .iter() - .copied() - .filter(|record| { - record - .register_mma - .as_ref() - .is_some_and(|mma| mma.kind == Some(RegisterMmaKind::Mxf8f6f4)) - }) - .collect::>(); - assert_eq!(mxf8f6f4_records.len(), 25); - assert_eq!(mxf8f6f4_records[0].abi_id, "i0858"); - assert_eq!(mxf8f6f4_records[24].abi_id, "i0882"); - assert!(mxf8f6f4_records.iter().all(|record| { - let mma = record.register_mma.as_ref().unwrap(); - mma.shape == RegisterMmaShape::M16n8k32 - && mma.accumulator == RegisterMmaAccumulator::F32 - && mma.adapter == RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 - && record.rust_arguments - == [ - "[f32; 4]", "[u32; 4]", "[u32; 2]", "u32", "u16", "u16", "u32", "u16", - "u16", - ] - && record.rust_result == "[f32; 4]" - && record.minimum_ptx == "8.7" - && record.minimum_sm.is_none() - && record.targets == REGISTER_MMA_F8F6F4_TARGETS - && record - .expected_ptx - .modifiers - .iter() - .any(|modifier| modifier == "kind::mxf8f6f4") - })); - - let integer_records: Vec<_> = records - .iter() - .copied() - .filter(|record| { - record.register_mma.as_ref().is_some_and(|mma| { - mma.operation == RegisterMmaOperation::Multiply - && mma.accumulator == RegisterMmaAccumulator::S32 - }) - }) - .collect(); - assert_eq!(integer_records.len(), 48); - let binary_records = records - .iter() - .copied() - .filter(|record| { - record - .register_mma - .as_ref() - .is_some_and(|mma| mma.operation != RegisterMmaOperation::Multiply) - }) - .collect::>(); - assert_eq!(binary_records.len(), 6); - let int8_records = integer_records - .iter() - .copied() - .filter(|record| { - let mma = record.register_mma.as_ref().unwrap(); - matches!( - mma.a_element, - RegisterMmaElement::S8 | RegisterMmaElement::U8 - ) && matches!( - mma.b_element, - RegisterMmaElement::S8 | RegisterMmaElement::U8 - ) - }) - .collect::>(); - assert_eq!(int8_records.len(), 24); - let int4_records = integer_records - .iter() - .copied() - .filter(|record| { - let mma = record.register_mma.as_ref().unwrap(); - matches!( - mma.a_element, - RegisterMmaElement::S4 | RegisterMmaElement::U4 - ) && matches!( - mma.b_element, - RegisterMmaElement::S4 | RegisterMmaElement::U4 - ) - }) - .collect::>(); - assert_eq!(int4_records.len(), 24); - let actual_variants = integer_records - .iter() - .map(|record| { - let mma = record.register_mma.as_ref().unwrap(); - (mma.shape, mma.a_element, mma.b_element, mma.overflow) - }) - .collect::>(); - let expected_int8_variants = [ - RegisterMmaShape::M8n8k16, - RegisterMmaShape::M16n8k16, - RegisterMmaShape::M16n8k32, - ] - .into_iter() - .flat_map(|shape| { - [RegisterMmaElement::S8, RegisterMmaElement::U8] - .into_iter() - .flat_map(move |a_element| { - [RegisterMmaElement::S8, RegisterMmaElement::U8] - .into_iter() - .flat_map(move |b_element| { - [ - RegisterMmaOverflow::Wrapping, - RegisterMmaOverflow::Satfinite, - ] - .into_iter() - .map(move |overflow| (shape, a_element, b_element, overflow)) - }) - }) - }) - .collect::>(); - let expected_int4_variants = [ - RegisterMmaShape::M8n8k32, - RegisterMmaShape::M16n8k32, - RegisterMmaShape::M16n8k64, - ] - .into_iter() - .flat_map(|shape| { - [RegisterMmaElement::S4, RegisterMmaElement::U4] - .into_iter() - .flat_map(move |a_element| { - [RegisterMmaElement::S4, RegisterMmaElement::U4] - .into_iter() - .flat_map(move |b_element| { - [ - RegisterMmaOverflow::Wrapping, - RegisterMmaOverflow::Satfinite, - ] - .into_iter() - .map(move |overflow| (shape, a_element, b_element, overflow)) - }) - }) - }) - .collect::>(); - let expected_variants = expected_int8_variants - .union(&expected_int4_variants) - .copied() - .collect::>(); - assert_eq!(actual_variants, expected_variants); - assert_eq!( - integer_records - .iter() - .filter(|record| { - record.register_mma.as_ref().unwrap().compatibility_source - == RegisterMmaCompatibilitySource::GeneratedStub - }) - .count(), - 47 - ); - - let actual_binary_variants = binary_records - .iter() - .map(|record| { - let mma = record.register_mma.as_ref().unwrap(); - (mma.shape, mma.operation) - }) - .collect::>(); - let expected_binary_variants = [ - RegisterMmaShape::M8n8k128, - RegisterMmaShape::M16n8k128, - RegisterMmaShape::M16n8k256, - ] - .into_iter() - .flat_map(|shape| { - [RegisterMmaOperation::XorPopc, RegisterMmaOperation::AndPopc] - .into_iter() - .map(move |operation| (shape, operation)) - }) - .collect::>(); - assert_eq!(actual_binary_variants, expected_binary_variants); - assert!(binary_records.iter().all(|record| { - let mma = record.register_mma.as_ref().unwrap(); - mma.accumulator == RegisterMmaAccumulator::S32 - && mma.a_element == RegisterMmaElement::B1 - && mma.b_element == RegisterMmaElement::B1 - && mma.overflow == RegisterMmaOverflow::Wrapping - && mma.compatibility_source == RegisterMmaCompatibilitySource::GeneratedStub - && record.expected_ptx.modifiers.ends_with(&[ - match mma.operation { - RegisterMmaOperation::XorPopc => "xor".into(), - RegisterMmaOperation::AndPopc => "and".into(), - RegisterMmaOperation::Multiply => unreachable!(), - }, - "popc".into(), - ]) - })); - - for record in &binary_records { - let mma = record.register_mma.as_ref().unwrap(); - let (arguments, result, adapter) = match mma.shape { - RegisterMmaShape::M8n8k128 => ( - &["[i32; 2]", "u32", "u32"] as &[_], - "[i32; 2]", - RegisterMmaAdapter::C2I32A1U32B1U32ToD2I32, - ), - RegisterMmaShape::M16n8k128 => ( - &["[i32; 4]", "[u32; 2]", "u32"] as &[_], - "[i32; 4]", - RegisterMmaAdapter::C4I32A2U32B1U32ToD4I32, - ), - RegisterMmaShape::M16n8k256 => ( - &["[i32; 4]", "[u32; 4]", "[u32; 2]"] as &[_], - "[i32; 4]", - RegisterMmaAdapter::C4I32A4U32B2U32ToD4I32, - ), - _ => unreachable!(), - }; - assert_eq!(record.rust_arguments, arguments); - assert_eq!(record.rust_result, result); - assert_eq!(mma.adapter, adapter); - let expected_floor = match (mma.shape, mma.operation) { - (RegisterMmaShape::M8n8k128, RegisterMmaOperation::XorPopc) => ("7.0", "sm_75"), - (_, RegisterMmaOperation::XorPopc) => ("7.0", "sm_80"), - (_, RegisterMmaOperation::AndPopc) => ("7.1", "sm_80"), - _ => unreachable!(), - }; - assert_eq!(record.minimum_ptx, expected_floor.0); - assert_eq!(record.minimum_sm.as_deref(), Some(expected_floor.1)); - } - - for record in integer_records.iter().filter(|record| { - matches!( - record.register_mma.as_ref().unwrap().shape, - RegisterMmaShape::M8n8k16 | RegisterMmaShape::M8n8k32 - ) - }) { - assert_eq!(record.rust_arguments, ["[i32; 2]", "u32", "u32"]); - assert_eq!(record.rust_result, "[i32; 2]"); - assert_eq!(record.minimum_ptx, "6.5"); - assert_eq!(record.minimum_sm.as_deref(), Some("sm_75")); - assert_eq!( - record.register_mma.as_ref().unwrap().adapter, - RegisterMmaAdapter::C2I32A1U32B1U32ToD2I32 - ); - } - - for record in int4_records.iter().filter(|record| { - record.register_mma.as_ref().unwrap().shape == RegisterMmaShape::M16n8k32 - }) { - assert_eq!(record.rust_arguments, ["[i32; 4]", "[u32; 2]", "u32"]); - assert_eq!(record.rust_result, "[i32; 4]"); - assert_eq!(record.minimum_ptx, "7.0"); - assert_eq!(record.minimum_sm.as_deref(), Some("sm_80")); - assert_eq!( - record.register_mma.as_ref().unwrap().adapter, - RegisterMmaAdapter::C4I32A2U32B1U32ToD4I32 - ); - } - - for record in int4_records.iter().filter(|record| { - record.register_mma.as_ref().unwrap().shape == RegisterMmaShape::M16n8k64 - }) { - assert_eq!(record.rust_arguments, ["[i32; 4]", "[u32; 4]", "[u32; 2]"]); - assert_eq!(record.rust_result, "[i32; 4]"); - assert_eq!(record.minimum_ptx, "7.0"); - assert_eq!(record.minimum_sm.as_deref(), Some("sm_80")); - assert_eq!( - record.register_mma.as_ref().unwrap().adapter, - RegisterMmaAdapter::C4I32A4U32B2U32ToD4I32 - ); - } - - let actual_int4_abi_ids = int4_records - .iter() - .map(|record| record.abi_id.as_str()) - .collect::>(); - let expected_int4_abi_ids = (133..=156) - .map(|id| format!("i{id:04}")) - .collect::>(); - assert_eq!( - actual_int4_abi_ids, - expected_int4_abi_ids - .iter() - .map(String::as_str) - .collect::>() - ); - - let int8_k32 = int8_records - .iter() - .find(|record| { - record.register_mma.as_ref().unwrap().shape == RegisterMmaShape::M16n8k32 - }) - .unwrap(); - let int4_k32 = int4_records - .iter() - .find(|record| { - record.register_mma.as_ref().unwrap().shape == RegisterMmaShape::M16n8k32 - }) - .unwrap(); - assert_eq!( - int8_k32.register_mma.as_ref().unwrap().adapter, - RegisterMmaAdapter::C4I32A4U32B2U32ToD4I32 - ); - assert_eq!( - int4_k32.register_mma.as_ref().unwrap().adapter, - RegisterMmaAdapter::C4I32A2U32B1U32ToD4I32 - ); - - for policy in &records { - let declaration = declarations[policy.source_record.as_deref().unwrap()]; - assert_eq!(declaration.selections.len(), 1); - assert!( - selection_matches_policy(policy, &declaration.selections[0]).unwrap(), - "{}", - policy.id - ); - validate_imported_policy(policy, declaration).unwrap(); - } - - let valid = records[0]; - let declaration = declarations[valid.source_record.as_deref().unwrap()]; - - let mut non_convergent = valid.clone(); - non_convergent.convergent = false; - assert!( - validate_imported_policy(&non_convergent, declaration) - .unwrap_err() - .to_string() - .contains("effects") - ); - - let mut typed_route = valid.clone(); - typed_route.backend_lowerings[0].mechanism = BackendLoweringMechanism::TypedNvvm; - assert!(validate_imported_policy(&typed_route, declaration).is_err()); - - let mut selectionless = declaration.clone(); - selectionless.selections.clear(); - assert!(validate_imported_policy(valid, &selectionless).is_err()); - - let mut crossed_variant = valid.clone(); - crossed_variant.register_mma.as_mut().unwrap().a_element = RegisterMmaElement::F16; - assert!(validate_imported_policy(&crossed_variant, declaration).is_err()); - - let generated = int8_records - .iter() - .copied() - .find(|record| record.id == "mma_m16n8k16_s32_s8_u8_satfinite") - .unwrap(); - let generated_declaration = declarations[generated.source_record.as_deref().unwrap()]; - - let mut wrong_stub_owner = generated.clone(); - wrong_stub_owner - .register_mma - .as_mut() - .unwrap() - .compatibility_source = RegisterMmaCompatibilitySource::ExistingStub; - assert!(validate_imported_policy(&wrong_stub_owner, generated_declaration).is_err()); - - let mut wrong_b_element = generated.clone(); - wrong_b_element.register_mma.as_mut().unwrap().b_element = RegisterMmaElement::S8; - assert!(validate_imported_policy(&wrong_b_element, generated_declaration).is_err()); - - let mut wrong_overflow = generated.clone(); - wrong_overflow.register_mma.as_mut().unwrap().overflow = RegisterMmaOverflow::Wrapping; - assert!(validate_imported_policy(&wrong_overflow, generated_declaration).is_err()); - - let mut wrong_shape = generated.clone(); - wrong_shape.register_mma.as_mut().unwrap().shape = RegisterMmaShape::M16n8k32; - assert!(validate_imported_policy(&wrong_shape, generated_declaration).is_err()); - - let mut wrong_adapter = generated.clone(); - wrong_adapter.register_mma.as_mut().unwrap().adapter = - RegisterMmaAdapter::C4I32A4U32B2U32ToD4I32; - assert!(validate_imported_policy(&wrong_adapter, generated_declaration).is_err()); - - let binary = binary_records - .iter() - .copied() - .find(|record| record.id == "mma_m8n8k128_s32_b1_xor_popc") - .unwrap(); - let binary_declaration = declarations[binary.source_record.as_deref().unwrap()]; - - let mut wrong_binary_operation = binary.clone(); - wrong_binary_operation - .register_mma - .as_mut() - .unwrap() - .operation = RegisterMmaOperation::AndPopc; - assert!(validate_imported_policy(&wrong_binary_operation, binary_declaration).is_err()); - - let mut wrong_binary_floor = binary.clone(); - wrong_binary_floor.minimum_sm = Some("sm_80".into()); - assert!(validate_imported_policy(&wrong_binary_floor, binary_declaration).is_err()); - - let mut wrong_binary_element = binary.clone(); - wrong_binary_element - .register_mma - .as_mut() - .unwrap() - .a_element = RegisterMmaElement::U4; - assert!(validate_imported_policy(&wrong_binary_element, binary_declaration).is_err()); - } - - #[test] - fn pinned_sparse_mma_records_close_shape_specific_selectors_and_ranges() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let (mut overlay, _) = - read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")).unwrap(); - bind_pinned_abi_ids(&repo_root, &mut overlay); - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - let declarations: BTreeMap<_, _> = imported - .intrinsics - .iter() - .map(|record| (record.source_record.as_str(), record)) - .collect(); - let records = overlay - .intrinsics - .iter() - .filter(|record| record.family == "sparse_mma") - .collect::>(); - assert_eq!(records.len(), 114); - assert_eq!( - records - .iter() - .map(|record| record.abi_id.clone()) - .collect::>(), - (163..=251) - .chain(525..=549) - .map(|id| format!("i{id:04}")) - .collect::>() - ); - - let mut derived_ids = BTreeSet::new(); - let mut derived_operation_keys = BTreeSet::new(); - let mut derived_source_records = BTreeSet::new(); - let mut derived_llvm_symbols = BTreeSet::new(); - for record in &records { - let identity = &sparse_mma_recipe(record.sparse_mma.as_ref().unwrap()) - .unwrap() - .identity; - assert_eq!(record.id, identity.id); - assert_eq!(record.operation_key, identity.operation_key); - assert_eq!( - record.source_record.as_deref(), - Some(identity.source_record.as_str()) - ); - assert_eq!( - record.llvm_symbol.as_deref(), - Some(identity.llvm_symbol.as_str()) - ); - assert_eq!(record.expected_ptx.modifiers, identity.ptx_modifiers); - assert!(derived_ids.insert(identity.id.clone())); - assert!(derived_operation_keys.insert(identity.operation_key.clone())); - assert!(derived_source_records.insert(identity.source_record.clone())); - assert!(derived_llvm_symbols.insert(identity.llvm_symbol.clone())); - } - assert_eq!(derived_ids.len(), 114); - assert_eq!(derived_operation_keys.len(), 114); - assert_eq!(derived_source_records.len(), 114); - assert_eq!(derived_llvm_symbols.len(), 114); - - let integer_records = records - .iter() - .copied() - .filter(|record| { - record.sparse_mma.as_ref().unwrap().accumulator == SparseMmaAccumulator::S32 - }) - .collect::>(); - assert_eq!(integer_records.len(), 64); - let f32_records = records - .iter() - .copied() - .filter(|record| { - record.sparse_mma.as_ref().unwrap().accumulator == SparseMmaAccumulator::F32 - }) - .collect::>(); - assert_eq!(f32_records.len(), 25); - let f16_records = records - .iter() - .copied() - .filter(|record| { - record.sparse_mma.as_ref().unwrap().accumulator == SparseMmaAccumulator::F16 - }) - .collect::>(); - assert_eq!(f16_records.len(), 25); - - let variants = integer_records - .iter() - .map(|record| { - let mma = record.sparse_mma.as_ref().unwrap(); - let carrier = - sparse_mma_carrier_recipe(mma.shape, mma.a_element, mma.b_element).unwrap(); - assert_eq!(mma.accumulator, SparseMmaAccumulator::S32); - assert_eq!(mma.selector, carrier.selector); - assert_eq!(mma.adapter, carrier.adapter); - assert_eq!(mma.llvm_adapter, carrier.llvm_adapter); - assert_eq!(record.rust_arguments, carrier.rust_arguments()); - assert_eq!(record.dialect_operands, carrier.dialect_operands()); - assert_eq!(record.llvm_arguments, carrier.llvm_arguments()); - assert_eq!( - record.expected_ptx.operands, - carrier.expected_ptx_operands() - ); - assert_eq!(record.minimum_ptx, sparse_mma_minimum_ptx(mma)); - assert_eq!(record.minimum_sm.as_deref(), Some("sm_80")); - assert_eq!( - record.expected_ptx.operands.last(), - Some(&OperandPattern::Immediate) - ); - assert_eq!( - record.expected_ptx.modifiers.first().map(String::as_str), - Some(match mma.metadata { - SparseMmaMetadata::Standard => "sp", - SparseMmaMetadata::Ordered => "sp::ordered_metadata", - }) - ); - ( - mma.shape, - mma.a_element, - mma.b_element, - mma.overflow, - mma.metadata, - ) - }) - .collect::>(); - let mut expected_variants = BTreeSet::new(); - for shape in [SparseMmaShape::M16n8k32, SparseMmaShape::M16n8k64] { - let metadata = match shape { - SparseMmaShape::M16n8k32 => [ - Some(SparseMmaMetadata::Standard), - Some(SparseMmaMetadata::Ordered), - ], - SparseMmaShape::M16n8k64 => [ - Some(SparseMmaMetadata::Standard), - Some(SparseMmaMetadata::Ordered), - ], - SparseMmaShape::M16n8k128 => [None, None], - }; - for a_element in [SparseMmaElement::S8, SparseMmaElement::U8] { - for b_element in [SparseMmaElement::S8, SparseMmaElement::U8] { - for overflow in [SparseMmaOverflow::Wrapping, SparseMmaOverflow::Satfinite] { - for metadata in metadata.into_iter().flatten() { - expected_variants - .insert((shape, a_element, b_element, overflow, metadata)); - } - } - } - } - } - for shape in [SparseMmaShape::M16n8k64, SparseMmaShape::M16n8k128] { - for a_element in [SparseMmaElement::S4, SparseMmaElement::U4] { - for b_element in [SparseMmaElement::S4, SparseMmaElement::U4] { - for overflow in [SparseMmaOverflow::Wrapping, SparseMmaOverflow::Satfinite] { - for metadata in [SparseMmaMetadata::Standard, SparseMmaMetadata::Ordered] { - expected_variants - .insert((shape, a_element, b_element, overflow, metadata)); - } - } - } - } - } - assert_eq!(variants, expected_variants); - - let f8f6f4_formats = [ - SparseMmaElement::E2m1, - SparseMmaElement::E2m3, - SparseMmaElement::E3m2, - SparseMmaElement::E4m3, - SparseMmaElement::E5m2, - ]; - assert_eq!( - f32_records - .iter() - .map(|record| { - let mma = record.sparse_mma.as_ref().unwrap(); - (mma.a_element, mma.b_element) - }) - .collect::>(), - f8f6f4_formats - .into_iter() - .flat_map(|a| f8f6f4_formats.into_iter().map(move |b| (a, b))) - .collect() - ); - assert_eq!( - f32_records - .iter() - .map(|record| record.abi_id.as_str()) - .collect::>(), - (227..=251) - .map(|id| format!("i{id:04}")) - .collect::>() - .iter() - .map(String::as_str) - .collect() - ); - assert_eq!( - f16_records - .iter() - .map(|record| { - let mma = record.sparse_mma.as_ref().unwrap(); - (mma.a_element, mma.b_element) - }) - .collect::>(), - f8f6f4_formats - .into_iter() - .flat_map(|a| f8f6f4_formats.into_iter().map(move |b| (a, b))) - .collect() - ); - assert_eq!( - f16_records - .iter() - .map(|record| record.abi_id.as_str()) - .collect::>(), - (525..=549) - .map(|id| format!("i{id:04}")) - .collect::>() - .iter() - .map(String::as_str) - .collect() - ); - for record in &f32_records { - let mma = record.sparse_mma.as_ref().unwrap(); - assert_eq!(mma.shape, SparseMmaShape::M16n8k64); - assert_eq!(mma.accumulator, SparseMmaAccumulator::F32); - assert_eq!(mma.overflow, SparseMmaOverflow::NotApplicable); - assert_eq!(mma.metadata, SparseMmaMetadata::Ordered); - assert_eq!(mma.selector, SparseMmaSelector::ImmediateZero); - assert_eq!( - mma.adapter, - SparseMmaAdapter::C4F32A4U32B4U32MetadataU32SelectorU32ToD4F32 - ); - assert_eq!( - mma.llvm_adapter, - SparseMmaLlvmAdapter::A4I32B4I32C4F32MetadataI32SelectorI32ToD4F32 - ); - assert_eq!( - record.rust_arguments, - ["[f32; 4]", "[u32; 4]", "[u32; 4]", "u32", "u32"] - ); - assert_eq!(record.rust_result, "[f32; 4]"); - assert_eq!( - record.dialect_operands, - [ - "f32", "f32", "f32", "f32", "u32", "u32", "u32", "u32", "u32", "u32", "u32", - "u32", "u32", "u32" - ] - ); - assert_eq!(record.dialect_results, ["f32", "f32", "f32", "f32"]); - assert_eq!( - record.llvm_arguments, - [ - "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", "f32", "f32", "f32", - "f32", "i32", "i32" - ] - ); - assert_eq!(record.llvm_results, ["f32", "f32", "f32", "f32"]); - assert_eq!(record.minimum_ptx, "8.7"); - assert_eq!(record.minimum_sm, None); - // Same contract as the F16 accumulator: both float forms are gated - // on `hasMMABlockScale()`, so neither is narrower than the other. - assert_eq!(record.targets, SPARSE_MMA_F8F6F4_TARGETS); - assert_eq!(record.backend_lowerings.len(), 2); - assert!(record.backend_lowerings.iter().all(|lowering| { - lowering.mechanism == BackendLoweringMechanism::InlinePtx - && lowering.minimum_ptx.as_deref() == Some("8.7") - && lowering.minimum_sm.is_none() - })); - assert_eq!( - record.expected_ptx.operands, - [ - OperandPattern::RegisterList { length: 4 }, - OperandPattern::RegisterList { length: 4 }, - OperandPattern::RegisterList { length: 4 }, - OperandPattern::RegisterList { length: 4 }, - OperandPattern::Register, - OperandPattern::Immediate, - ] - ); - } - - for record in &f16_records { - let mma = record.sparse_mma.as_ref().unwrap(); - assert_eq!(mma.shape, SparseMmaShape::M16n8k64); - assert_eq!(mma.accumulator, SparseMmaAccumulator::F16); - assert_eq!(mma.overflow, SparseMmaOverflow::NotApplicable); - assert_eq!(mma.metadata, SparseMmaMetadata::Ordered); - assert_eq!(mma.selector, SparseMmaSelector::ImmediateZero); - assert_eq!( - mma.adapter, - SparseMmaAdapter::C2U32A4U32B4U32MetadataU32SelectorU32ToD2U32 - ); - assert_eq!( - mma.llvm_adapter, - SparseMmaLlvmAdapter::A4I32B4I32C2V2F16MetadataI32SelectorI32ToD2V2F16 - ); - assert_eq!( - record.rust_arguments, - ["[u32; 2]", "[u32; 4]", "[u32; 4]", "u32", "u32"] - ); - assert_eq!(record.rust_result, "[u32; 2]"); - assert_eq!( - record.dialect_operands, - [ - "i32", "i32", "u32", "u32", "u32", "u32", "u32", "u32", "u32", "u32", "u32", - "u32" - ] - ); - assert_eq!(record.dialect_results, ["i32", "i32"]); - assert_eq!( - record.llvm_arguments, - [ - "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", "v2f16", "v2f16", - "i32", "i32" - ] - ); - assert_eq!(record.llvm_results, ["v2f16", "v2f16"]); - assert_eq!(record.minimum_ptx, "8.7"); - assert_eq!(record.minimum_sm, None); - assert_eq!(record.targets, SPARSE_MMA_F8F6F4_TARGETS); - assert!(record.convergent && !record.pure); - assert!(record.backend_lowerings.iter().all(|lowering| { - lowering.mechanism == BackendLoweringMechanism::InlinePtx - && lowering.minimum_ptx.as_deref() == Some("8.7") - && lowering.minimum_sm.is_none() - })); - assert_eq!( - record.expected_ptx.operands, - [ - OperandPattern::RegisterList { length: 2 }, - OperandPattern::RegisterList { length: 4 }, - OperandPattern::RegisterList { length: 4 }, - OperandPattern::RegisterList { length: 2 }, - OperandPattern::Register, - OperandPattern::Immediate, - ] - ); - } - - for policy in &records { - let declaration = declarations[policy.source_record.as_deref().unwrap()]; - assert_eq!(declaration.selections.len(), 1); - assert!( - selection_matches_policy(policy, &declaration.selections[0]).unwrap(), - "{}", - policy.id - ); - validate_imported_policy(policy, declaration).unwrap(); - } - - let mut selectionless = declarations[records[0].source_record.as_deref().unwrap()].clone(); - selectionless.selections.clear(); - assert!(validate_imported_policy(records[0], &selectionless).is_err()); - - for (id, range_prefix, wrong_range) in [ - ("mma_sp_m16n8k32_s32_s8", "Range"), - ("mma_sp_m16n8k64_s32_s8", "Range"), - ( - "mma_sp_ordered_metadata_m16n8k64_s32_s4", - "Range", - ), - ("mma_sp_m16n8k64_s32_s4", "Range"), - ( - "mma_sp_ordered_metadata_m16n8k128_s32_s4", - "Range", - ), - ("mma_sp_m16n8k128_s32_s4", "Range"), - ( - "mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m1_e2m1_f32", - "Range", - ), - ( - "mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e2m1_f16", - "Range", - ), - ] { - let valid = records - .iter() - .copied() - .find(|record| record.id == id) - .unwrap(); - let declaration = declarations[valid.source_record.as_deref().unwrap()]; - - let mut runtime_selector = valid.clone(); - *runtime_selector.expected_ptx.operands.last_mut().unwrap() = OperandPattern::Register; - assert!( - validate_imported_policy(&runtime_selector, declaration) - .unwrap_err() - .to_string() - .contains("exact selection") - ); - - let mut wrong_declaration = declaration.clone(); - *wrong_declaration - .properties - .iter_mut() - .find(|property| property.starts_with(range_prefix)) - .unwrap() = wrong_range.into(); - assert!( - validate_imported_policy(valid, &wrong_declaration) - .unwrap_err() - .to_string() - .contains("immediate range") - ); - } - - let k64 = records - .iter() - .copied() - .find(|record| record.id == "mma_sp_m16n8k64_s32_s8") - .unwrap(); - let k64_declaration = declarations[k64.source_record.as_deref().unwrap()]; - assert_eq!(k64.minimum_ptx, "7.1"); - let ordered_k64 = records - .iter() - .copied() - .find(|record| record.id == "mma_sp_ordered_metadata_m16n8k64_s32_s8") - .unwrap(); - assert_eq!(ordered_k64.minimum_ptx, "8.5"); - - let f8f6f4 = f32_records[0]; - let f8f6f4_declaration = declarations[f8f6f4.source_record.as_deref().unwrap()]; - let mut widened_family = f8f6f4.clone(); - widened_family.targets = "sm_120f".into(); - assert!(validate_imported_policy(&widened_family, f8f6f4_declaration).is_err()); - let mut widened_architecture = f8f6f4.clone(); - widened_architecture.targets = "sm_121a".into(); - assert!(validate_imported_policy(&widened_architecture, f8f6f4_declaration).is_err()); - - let f8f6f4_f16 = f16_records[0]; - for block_scale in [f8f6f4, f8f6f4_f16] { - let mut missing_predicate = - declarations[block_scale.source_record.as_deref().unwrap()].clone(); - missing_predicate.selections[0].predicates.clear(); - assert!( - validate_imported_policy(block_scale, &missing_predicate) - .unwrap_err() - .to_string() - .contains("exact selection") - ); - } - - let ordered_k64_int4 = records - .iter() - .copied() - .find(|record| record.id == "mma_sp_ordered_metadata_m16n8k64_s32_s4") - .unwrap(); - let ordered_k64_int4_declaration = - declarations[ordered_k64_int4.source_record.as_deref().unwrap()]; - assert_eq!(ordered_k64_int4.minimum_ptx, "8.5"); - assert_eq!(ordered_k64_int4.rust_arguments[1], "[u32; 2]"); - assert_eq!(ordered_k64_int4.rust_arguments[2], "[u32; 2]"); - assert_eq!(ordered_k64_int4.llvm_arguments.len(), 10); - assert_eq!( - ordered_k64_int4.sparse_mma.as_ref().unwrap().selector, - SparseMmaSelector::ImmediateZeroOrOne - ); - assert_eq!( - ordered_k64_int4.sparse_mma.as_ref().unwrap().adapter, - SparseMmaAdapter::C4I32A2U32B2U32MetadataU32SelectorU32ToD4I32 - ); - assert_eq!( - ordered_k64_int4.sparse_mma.as_ref().unwrap().llvm_adapter, - SparseMmaLlvmAdapter::A2I32B2I32C4I32MetadataI32SelectorI32ToD4I32 - ); - let standard_k64_int4 = records - .iter() - .copied() - .find(|record| record.id == "mma_sp_m16n8k64_s32_s4") - .unwrap(); - assert_eq!(standard_k64_int4.minimum_ptx, "7.1"); - assert_eq!( - standard_k64_int4.rust_arguments, - ordered_k64_int4.rust_arguments - ); - assert_eq!( - standard_k64_int4.llvm_arguments, - ordered_k64_int4.llvm_arguments - ); - - let ordered_k128_int4 = records - .iter() - .copied() - .find(|record| record.id == "mma_sp_ordered_metadata_m16n8k128_s32_s4") - .unwrap(); - let ordered_k128_int4_declaration = - declarations[ordered_k128_int4.source_record.as_deref().unwrap()]; - assert_eq!(ordered_k128_int4.minimum_ptx, "8.5"); - assert_eq!(ordered_k128_int4.rust_arguments[1], "[u32; 4]"); - assert_eq!(ordered_k128_int4.rust_arguments[2], "[u32; 4]"); - assert_eq!(ordered_k128_int4.llvm_arguments.len(), 14); - assert_eq!( - ordered_k128_int4.sparse_mma.as_ref().unwrap().selector, - SparseMmaSelector::ImmediateZero - ); - assert_eq!( - ordered_k128_int4.sparse_mma.as_ref().unwrap().adapter, - SparseMmaAdapter::C4I32A4U32B4U32MetadataU32SelectorU32ToD4I32 - ); - assert_eq!( - ordered_k128_int4.sparse_mma.as_ref().unwrap().llvm_adapter, - SparseMmaLlvmAdapter::A4I32B4I32C4I32MetadataI32SelectorI32ToD4I32 - ); - let standard_k128_int4 = records - .iter() - .copied() - .find(|record| record.id == "mma_sp_m16n8k128_s32_s4") - .unwrap(); - assert_eq!(standard_k128_int4.minimum_ptx, "7.1"); - assert_eq!( - standard_k128_int4.rust_arguments, - ordered_k128_int4.rust_arguments - ); - assert_eq!( - standard_k128_int4.llvm_arguments, - ordered_k128_int4.llvm_arguments - ); - - let mut wrong_k128_selector = ordered_k128_int4.clone(); - wrong_k128_selector.sparse_mma.as_mut().unwrap().selector = - SparseMmaSelector::ImmediateZeroOrOne; - assert!( - validate_imported_policy(&wrong_k128_selector, ordered_k128_int4_declaration).is_err() - ); - - let mut mixed_k128_width = ordered_k128_int4.clone(); - mixed_k128_width.sparse_mma.as_mut().unwrap().b_element = SparseMmaElement::U8; - assert!( - validate_imported_policy(&mixed_k128_width, ordered_k128_int4_declaration) - .unwrap_err() - .to_string() - .contains("unsupported sparse-MMA variant") - ); - - let mut mixed_width = ordered_k64_int4.clone(); - mixed_width.sparse_mma.as_mut().unwrap().b_element = SparseMmaElement::U8; - assert!( - validate_imported_policy(&mixed_width, ordered_k64_int4_declaration) - .unwrap_err() - .to_string() - .contains("unsupported sparse-MMA variant") - ); - - let mut wrong_k64_selector = k64.clone(); - wrong_k64_selector.sparse_mma.as_mut().unwrap().selector = - SparseMmaSelector::ImmediateZeroOrOne; - assert!(validate_imported_policy(&wrong_k64_selector, k64_declaration).is_err()); - - let mut wrong_k64_adapter = k64.clone(); - wrong_k64_adapter.sparse_mma.as_mut().unwrap().adapter = - SparseMmaAdapter::C4I32A2U32B2U32MetadataU32SelectorU32ToD4I32; - assert!(validate_imported_policy(&wrong_k64_adapter, k64_declaration).is_err()); - - let mut wrong_k64_llvm_adapter = k64.clone(); - wrong_k64_llvm_adapter - .sparse_mma - .as_mut() - .unwrap() - .llvm_adapter = SparseMmaLlvmAdapter::A2I32B2I32C4I32MetadataI32SelectorI32ToD4I32; - assert!(validate_imported_policy(&wrong_k64_llvm_adapter, k64_declaration).is_err()); - - let mut wrong_k64_shape = k64.clone(); - wrong_k64_shape.sparse_mma.as_mut().unwrap().shape = SparseMmaShape::M16n8k32; - assert!(validate_imported_policy(&wrong_k64_shape, k64_declaration).is_err()); - - let mut wrong_k64_carriers = k64.clone(); - wrong_k64_carriers.dialect_operands.pop(); - assert!(validate_imported_policy(&wrong_k64_carriers, k64_declaration).is_err()); - - let mut wrong_k64_lowering = k64.clone(); - wrong_k64_lowering.lowering = "generated_register_mma".into(); - assert!(validate_imported_policy(&wrong_k64_lowering, k64_declaration).is_err()); - - let mut mismatched_metadata_identity = k64.clone(); - mismatched_metadata_identity - .sparse_mma - .as_mut() - .unwrap() - .metadata = SparseMmaMetadata::Ordered; - assert!(validate_imported_policy(&mismatched_metadata_identity, k64_declaration).is_err()); - } - - #[test] - fn cp_async_copy_recipe_admits_only_classic_llvm_forms() { - let cases = [ - ( - CpAsyncCachePolicy::Ca, - CpAsyncCopySize::B4, - CpAsyncSourceSize::Full, - Some("cp_async_ca_4"), - ), - ( - CpAsyncCachePolicy::Ca, - CpAsyncCopySize::B4, - CpAsyncSourceSize::Runtime, - Some("cp_async_ca_zfill_4"), - ), - ( - CpAsyncCachePolicy::Ca, - CpAsyncCopySize::B8, - CpAsyncSourceSize::Full, - Some("cp_async_ca_8"), - ), - ( - CpAsyncCachePolicy::Ca, - CpAsyncCopySize::B8, - CpAsyncSourceSize::Runtime, - Some("cp_async_ca_zfill_8"), - ), - ( - CpAsyncCachePolicy::Ca, - CpAsyncCopySize::B16, - CpAsyncSourceSize::Full, - Some("cp_async_ca_16"), - ), - ( - CpAsyncCachePolicy::Ca, - CpAsyncCopySize::B16, - CpAsyncSourceSize::Runtime, - Some("cp_async_ca_zfill_16"), - ), - ( - CpAsyncCachePolicy::Cg, - CpAsyncCopySize::B4, - CpAsyncSourceSize::Full, - None, - ), - ( - CpAsyncCachePolicy::Cg, - CpAsyncCopySize::B4, - CpAsyncSourceSize::Runtime, - None, - ), - ( - CpAsyncCachePolicy::Cg, - CpAsyncCopySize::B8, - CpAsyncSourceSize::Full, - None, - ), - ( - CpAsyncCachePolicy::Cg, - CpAsyncCopySize::B8, - CpAsyncSourceSize::Runtime, - None, - ), - ( - CpAsyncCachePolicy::Cg, - CpAsyncCopySize::B16, - CpAsyncSourceSize::Full, - Some("cp_async_cg_16"), - ), - ( - CpAsyncCachePolicy::Cg, - CpAsyncCopySize::B16, - CpAsyncSourceSize::Runtime, - Some("cp_async_cg_zfill_16"), - ), - ]; - - for (cache_policy, copy_size, source_size, expected) in cases { - let copy = crate::model::CpAsyncCopy { - cache_policy, - copy_size, - source_size, - adapter: if source_size == CpAsyncSourceSize::Runtime { - CpAsyncAdapter::DirectPointersAndSourceSize - } else { - CpAsyncAdapter::DirectPointers - }, - runtime_validation: RuntimeValidation::Unexecuted, - }; - assert_eq!( - cp_async_copy_recipe(©).map(|recipe| recipe.id), - expected - ); - } - } - - #[test] - fn pinned_cp_async_records_match_the_closed_recipes() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let (overlay, _) = - read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")).unwrap(); - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - let declarations: BTreeMap<_, _> = imported - .intrinsics - .iter() - .map(|record| (record.source_record.as_str(), record)) - .collect(); - let policies: Vec<_> = overlay - .intrinsics - .iter() - .filter(|record| matches!(record.family.as_str(), "cp_async_copy" | "cp_async_control")) - .collect(); - - assert_eq!(policies.len(), 11); - for policy in policies { - let declaration = declarations[policy.source_record.as_deref().unwrap()]; - validate_imported_policy(policy, declaration).unwrap(); - } - } - - #[test] - fn pinned_cp_async_mbarrier_records_match_the_closed_recipes() { - let records = pinned_cp_async_mbarrier_records(); - assert_eq!(records.len(), 4); - - for (policy, declaration) in records.values() { - validate_imported_policy(policy, declaration).unwrap(); - } - } - - #[test] - fn cp_async_mbarrier_recipes_fail_closed() { - let records = pinned_cp_async_mbarrier_records(); - let (arrive, declaration) = &records["cp_async_mbarrier_arrive"]; - let reject = - |policy: &OverlayIntrinsic, declaration: &ImportedIntrinsic, expected: &str| { - let error = validate_imported_policy(policy, declaration).unwrap_err(); - let message = error.to_string(); - assert!(message.contains(expected), "unexpected error: {message}"); - }; - - let mut wrong_symbol = arrive.clone(); - wrong_symbol.llvm_symbol = Some("llvm.nvvm.cp.async.mbarrier.changed".into()); - reject(&wrong_symbol, declaration, "LLVM symbol mismatch"); - - let mut wrong_signature = arrive.clone(); - wrong_signature.rust_arguments = vec!["*const u64".into()]; - reject( - &wrong_signature, - declaration, - "closed cp.async mbarrier Rust API", - ); - - let mut wrong_operation = arrive.clone(); - wrong_operation - .cp_async_mbarrier - .as_mut() - .unwrap() - .operation = CpAsyncMbarrierOperation::ArriveNoInc; - reject(&wrong_operation, declaration, "identity does not match"); - - let mut wrong_state_space = arrive.clone(); - wrong_state_space - .cp_async_mbarrier - .as_mut() - .unwrap() - .state_space = CpAsyncMbarrierStateSpace::Shared; - reject(&wrong_state_space, declaration, "identity does not match"); - - let mut wrong_adapter = arrive.clone(); - wrong_adapter.cp_async_mbarrier.as_mut().unwrap().adapter = - CpAsyncMbarrierAdapter::PointerToVoid; - wrong_adapter.rust_result = "u64".into(); - reject( - &wrong_adapter, - declaration, - "closed cp.async mbarrier Rust API", - ); - - let mut executed_without_evidence = arrive.clone(); - executed_without_evidence - .cp_async_mbarrier - .as_mut() - .unwrap() - .runtime_validation = RuntimeValidation::Executed; - reject( - &executed_without_evidence, - declaration, - "unrecorded cp.async mbarrier runtime validation", - ); - - let mut wrong_properties = declaration.clone(); - wrong_properties.properties.pop(); - reject(arrive, &wrong_properties, "cp.async mbarrier properties"); - - let mut wrong_selection = declaration.clone(); - wrong_selection.selections[0].source_record = "CP_ASYNC_MBARRIER_CHANGED".into(); - reject( - arrive, - &wrong_selection, - "imported cp.async mbarrier selection changed", - ); - - let mut wrong_floor = arrive.clone(); - wrong_floor.minimum_sm = Some("sm_90".into()); - reject(&wrong_floor, declaration, "effects or target floor"); - - let mut wrong_llvm_route = arrive.clone(); - wrong_llvm_route - .backend_lowerings - .iter_mut() - .find(|lowering| lowering.backend == IntrinsicBackend::LlvmNvptx) - .unwrap() - .mechanism = BackendLoweringMechanism::InlinePtx; - reject(&wrong_llvm_route, declaration, "reviewed typed-LLVM"); - - let mut wrong_lib_route = arrive.clone(); - wrong_lib_route - .backend_lowerings - .iter_mut() - .find(|lowering| lowering.backend == IntrinsicBackend::LibNvvm) - .unwrap() - .mechanism = BackendLoweringMechanism::TypedNvvm; - reject(&wrong_lib_route, declaration, "reviewed typed-LLVM"); - - let mut mixed_family = arrive.clone(); - mixed_family.cp_async_control = Some(crate::model::CpAsyncControl { - operation: CpAsyncControlOperation::CommitGroup, - adapter: CpAsyncControlAdapter::NoOperands, - runtime_validation: RuntimeValidation::Unexecuted, - }); - reject( - &mixed_family, - declaration, - "mixes another generated-family contract", - ); - } - - #[test] - fn pinned_mbarrier_basic_records_match_the_closed_recipes() { - let records = pinned_mbarrier_basic_records(); - assert_eq!(records.len(), 4); - - for (policy, declaration) in records.values() { - validate_imported_policy(policy, declaration).unwrap(); - } - } - - #[test] - fn mbarrier_basic_recipes_fail_closed() { - let records = pinned_mbarrier_basic_records(); - let (init, init_declaration) = &records["mbarrier_init"]; - let reject = - |policy: &OverlayIntrinsic, declaration: &ImportedIntrinsic, expected: &str| { - let error = validate_imported_policy(policy, declaration).unwrap_err(); - let message = error.to_string(); - assert!(message.contains(expected), "unexpected error: {message}"); - }; - - let mut wrong_symbol = init.clone(); - wrong_symbol.llvm_symbol = Some("llvm.nvvm.mbarrier.init.changed".into()); - reject(&wrong_symbol, init_declaration, "LLVM symbol mismatch"); - - let mut wrong_signature = init.clone(); - wrong_signature.rust_arguments = vec!["*mut u32".into(), "u32".into()]; - reject( - &wrong_signature, - init_declaration, - "unsafe mbarrier raw and compatibility API", - ); - - let mut wrong_operation = init.clone(); - wrong_operation.mbarrier_basic.as_mut().unwrap().operation = MbarrierBasicOperation::Arrive; - reject( - &wrong_operation, - init_declaration, - "operation, state space, and adapter disagree", - ); - - let mut wrong_adapter = init.clone(); - wrong_adapter.mbarrier_basic.as_mut().unwrap().adapter = - MbarrierBasicAdapter::InvalPointerToVoid; - reject( - &wrong_adapter, - init_declaration, - "operation, state space, and adapter disagree", - ); - - let mut executed_without_evidence = init.clone(); - executed_without_evidence - .mbarrier_basic - .as_mut() - .unwrap() - .runtime_validation = RuntimeValidation::Executed; - reject( - &executed_without_evidence, - init_declaration, - "unrecorded mbarrier runtime validation", - ); - - let mut wrong_properties = init_declaration.clone(); - wrong_properties.properties.pop(); - reject(init, &wrong_properties, "mbarrier properties"); - - let mut wrong_selection = init_declaration.clone(); - wrong_selection.selections[0].source_record = "MBARRIER_INIT_CHANGED".into(); - reject( - init, - &wrong_selection, - "imported mbarrier selection changed", - ); - - let mut wrong_ptx_floor = init.clone(); - wrong_ptx_floor.minimum_ptx = "7.1".into(); - reject( - &wrong_ptx_floor, - init_declaration, - "effects or target floor", - ); - - let mut wrong_sm_floor = init.clone(); - wrong_sm_floor.minimum_sm = Some("sm_90".into()); - reject(&wrong_sm_floor, init_declaration, "effects or target floor"); - - let mut wrong_llvm_route = init.clone(); - wrong_llvm_route - .backend_lowerings - .iter_mut() - .find(|lowering| lowering.backend == IntrinsicBackend::LlvmNvptx) - .unwrap() - .mechanism = BackendLoweringMechanism::InlinePtx; - reject( - &wrong_llvm_route, - init_declaration, - "reviewed mbarrier backend routes", - ); - - let mut wrong_lib_nvvm_route = init.clone(); - wrong_lib_nvvm_route - .backend_lowerings - .iter_mut() - .find(|lowering| lowering.backend == IntrinsicBackend::LibNvvm) - .unwrap() - .mechanism = BackendLoweringMechanism::TypedNvvm; - reject( - &wrong_lib_nvvm_route, - init_declaration, - "reviewed mbarrier backend routes", - ); - - let mut route_with_unreviewed_floor = init.clone(); - route_with_unreviewed_floor.backend_lowerings[0].minimum_sm = Some("sm_90".into()); - reject( - &route_with_unreviewed_floor, - init_declaration, - "reviewed mbarrier backend routes", - ); - - let mut mixed_family = init.clone(); - mixed_family.cp_async_control = Some(crate::model::CpAsyncControl { - operation: CpAsyncControlOperation::CommitGroup, - adapter: CpAsyncControlAdapter::NoOperands, - runtime_validation: RuntimeValidation::Unexecuted, - }); - reject( - &mixed_family, - init_declaration, - "mixes another generated-family contract", - ); - } - - #[test] - fn pinned_active_mask_and_warp_match_recipes_resolve() { - let records = pinned_active_mask_and_warp_match_records(); - assert_eq!(records.len(), 5); - - for (policy, declaration) in records.values() { - validate_imported_policy(policy, declaration).unwrap(); - } - } - - #[test] - fn active_mask_and_warp_match_recipes_fail_closed() { - let records = pinned_active_mask_and_warp_match_records(); - let reject = - |policy: &OverlayIntrinsic, declaration: &ImportedIntrinsic, expected: &str| { - let error = validate_imported_policy(policy, declaration).unwrap_err(); - let message = error.to_string(); - assert!(message.contains(expected), "unexpected error: {message}"); - }; - - let (active_mask, active_mask_declaration) = &records["active_mask"]; - let mut wrong_identity = active_mask.clone(); - wrong_identity.operation_key = "warp.active_mask.changed".into(); - reject( - &wrong_identity, - active_mask_declaration, - "active-mask identity", - ); - - let mut wrong_effects = active_mask.clone(); - wrong_effects.memory = "none".into(); - reject( - &wrong_effects, - active_mask_declaration, - "active-mask effects or target floor", - ); - - let (match_any, match_any_declaration) = &records["match_any_sync"]; - let mut wrong_adapter = match_any.clone(); - wrong_adapter.warp_match.as_mut().unwrap().adapter = - WarpMatchAdapter::ProjectMaskDiscardPredicate; - reject( - &wrong_adapter, - match_any_declaration, - "warp-match participation, adapter, or encoding", - ); - - let (match_all, match_all_declaration) = &records["match_all_sync"]; - let mut wrong_projection = match_all.clone(); - wrong_projection.dialect_results.push("i1".into()); - reject( - &wrong_projection, - match_all_declaration, - "two-operand warp-match lowering recipe", - ); - - let (match_any_i64, match_any_i64_declaration) = &records["match_any_i64_sync"]; - let mut incomplete_selections = match_any_i64_declaration.clone(); - incomplete_selections.selections.pop(); - reject( - match_any_i64, - &incomplete_selections, - "exactly ii/ir/ri/rr selections", - ); - - let (match_all_i64, match_all_i64_declaration) = &records["match_all_i64_sync"]; - let mut wrong_predicates = match_all_i64_declaration.clone(); - wrong_predicates.selections[0].predicates[0] = "Subtarget->getPTXVersion() >= 61".into(); - reject( - match_all_i64, - &wrong_predicates, - "PTX shape, predicates, or constraints", - ); - } - - #[test] - fn ptx_native_source_provenance_fails_closed() { - let mut mixed = packed_policy("packed_atomic_add_f16x2"); - mixed.source_record = Some("invented_llvm_record".into()); - assert!( - resolve_policy_source(&mixed) - .unwrap_err() - .to_string() - .contains("mixes tagged source provenance") - ); - - let mut fake_llvm = packed_policy("packed_atomic_add_f16x2"); - fake_llvm.llvm_symbol = Some("llvm.fake.packed.atomic".into()); - fake_llvm.llvm_arguments = vec!["ptr".into(), "i32".into()]; - fake_llvm.llvm_results = vec!["i32".into()]; - assert!( - validate_ptx_native_policy(&fake_llvm) - .unwrap_err() - .to_string() - .contains("must not invent LLVM source facts") - ); - - let mut wrong_instruction = packed_policy("packed_atomic_add_f16x2"); - wrong_instruction.source = Some(IntrinsicSource::PtxNative { - instruction: "atom.global.add.noftz.bf16x2".into(), - }); - assert!( - validate_ptx_native_policy(&wrong_instruction) - .unwrap_err() - .to_string() - .contains("does not match its packed format") - ); - - let mut wrong_kind = packed_policy("packed_atomic_add_f16x2"); - wrong_kind.source = Some(IntrinsicSource::LlvmImported { - source_record: "invented_llvm_record".into(), - }); - assert!( - validate_ptx_native_policy(&wrong_kind) - .unwrap_err() - .to_string() - .contains("source kind and imported declaration disagree") - ); - } - - #[test] - fn movmatrix_recipe_is_exact_and_fails_closed() { - let valid = movmatrix_policy(); - validate_ptx_native_policy(&valid).unwrap(); - - let reject = |policy: &OverlayIntrinsic, expected: &str| { - let message = validate_ptx_native_policy(policy).unwrap_err().to_string(); - assert!(message.contains(expected), "unexpected error: {message}"); - }; - - let mut wrong_shape = valid.clone(); - wrong_shape.expected_ptx.modifiers[2] = "m16n8".into(); - reject(&wrong_shape, "closed movmatrix recipe"); - - let mut wrong_participation = valid.clone(); - wrong_participation.convergent = false; - reject(&wrong_participation, "closed movmatrix recipe"); - - // A warp collective is not a function of its own operand, so the - // pure/no-memory pair every other collective avoids must stay rejected - // here too. - let mut wrong_purity = valid.clone(); - wrong_purity.pure = true; - reject(&wrong_purity, "closed movmatrix recipe"); - - let mut wrong_memory = valid.clone(); - wrong_memory.memory = "none".into(); - reject(&wrong_memory, "closed movmatrix recipe"); - - let mut wrong_floor = valid.clone(); - wrong_floor.backend_lowerings[0].minimum_ptx = Some("8.0".into()); - reject(&wrong_floor, "exact movmatrix floor"); - - let mut mixed = valid; - mixed.warp_barrier = Some(crate::model::WarpBarrier { - participation: - WarpBarrierParticipation::ExecutingLaneNamedAllNamedLanesSameInstructionAndMask, - legacy_pre_sm70: PreSm70MemberMaskRule::AllNamedLanesConvergedAndOnlyNamedLanesActive, - adapter: WarpBarrierAdapter::DirectMemberMask, - mask_encoding: WarpBarrierMaskEncoding::RegisterOrImmediate, - memory_ordering: WarpBarrierMemoryOrdering::ParticipatingLanes, - }); - reject(&mixed, "mixes another generated-family contract"); - } - - #[test] - fn packed_atomic_closed_semantics_reject_every_unreviewed_mutation() { - let valid = packed_policy("packed_atomic_add_f16x2"); - validate_ptx_native_policy(&valid).unwrap(); - - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let overlay = - std::fs::read_to_string(repo_root.join("intrinsics/overlay/packed_atomic.toml")) - .unwrap(); - for (field, accepted, rejected) in [ - ( - "state space", - "state_space = \"global\"", - "state_space = \"shared\"", - ), - ( - "ordering", - "ordering = \"relaxed\"", - "ordering = \"acquire\"", - ), - ("scope", "scope = \"gpu\"", "scope = \"system\""), - ( - "rounding", - "rounding = \"nearest_even\"", - "rounding = \"toward_zero\"", - ), - ( - "subnormal", - "subnormal = \"preserve\"", - "subnormal = \"flush\"", - ), - ( - "atomicity", - "atomicity = \"per_element\"", - "atomicity = \"coherent_pair\"", - ), - ( - "pointer safety", - "pointer_contract = \"mutable_global_u32_aligned4\"", - "pointer_contract = \"unaligned\"", - ), - ( - "mixed access safety", - "access_contract = \"no_mixed_whole_word_or_non_atomic_access\"", - "access_contract = \"mixed_access_allowed\"", - ), - ( - "scope safety", - "scope_contract = \"racing_atomics_mutually_inclusive\"", - "scope_contract = \"scope_unchecked\"", - ), - ( - "codegen", - "codegen_contract = \"exact_native_instruction\"", - "codegen_contract = \"semantic_equivalence\"", - ), - ] { - let mutated = overlay.replacen(accepted, rejected, 1); - let error = toml::from_str::(&mutated).unwrap_err(); - assert!( - error - .to_string() - .contains(rejected.split(" = ").next().unwrap()), - "{field} mutation did not fail closed: {error}" - ); - } - - let mut safe = valid; - safe.safe = true; - safe.safe_allowlist_reason = Some("incorrectly claims no caller obligations".into()); - assert!( - validate_ptx_native_policy(&safe) - .unwrap_err() - .to_string() - .contains("unsafe must-use packed atomic") - ); - } - - #[test] - fn redux_contract_validates_effects_participation_and_operand_adapter() { - let valid = redux_policy(); - let imported = redux_declaration(); - validate_imported_policy(&valid, &imported).unwrap(); - - assert_eq!( - valid.redux.as_ref().unwrap().adapter, - ReduxAdapter::MaskValueToSourceMemberMask - ); - - let mut missing_contract = valid.clone(); - missing_contract.redux = None; - assert!( - validate_imported_policy(&missing_contract, &imported) - .unwrap_err() - .to_string() - .contains("closed redux contract") - ); - - let mut wrong_effect = valid.clone(); - wrong_effect.memory = "none".into(); - assert!( - validate_imported_policy(&wrong_effect, &imported) - .unwrap_err() - .to_string() - .contains("redux effects") - ); - - let mut missing_imported_effect = imported; - missing_imported_effect - .properties - .retain(|property| property != "IntrInaccessibleMemOnly"); - assert!( - validate_imported_policy(&valid, &missing_imported_effect) - .unwrap_err() - .to_string() - .contains("memory and convergence effects") - ); - } - - #[test] - fn vote_modes_keep_exact_abi_identity_and_both_selection_encodings() { - for mode in [ - VoteMode::All, - VoteMode::Any, - VoteMode::Ballot, - VoteMode::Uni, - ] { - let policy = vote_policy(mode); - let declaration = vote_declaration(mode); - validate_imported_policy(&policy, &declaration).unwrap(); - assert_eq!( - policy.vote.as_ref().unwrap().legacy_pre_sm70, - PreSm70MemberMaskRule::AllNamedLanesConvergedAndOnlyNamedLanesActive - ); - - let selected: Vec<_> = declaration - .selections - .iter() - .filter(|selection| selection_matches_policy(&policy, selection).unwrap()) - .collect(); - assert_eq!(selected.len(), 2); - assert!(selected.iter().any(|selection| { - selection.source_record == vote_recipe(mode).immediate_selection - })); - assert!(selected.iter().any(|selection| { - selection.source_record == vote_recipe(mode).register_selection - })); - - let mut record = evidence(); - record.id = policy.id.clone(); - record.source_record = policy.source_record.clone(); - record.llvm_symbol = policy.llvm_symbol.clone(); - record.llvm_arguments = policy.llvm_arguments.clone(); - record.llvm_results = policy.llvm_results.clone(); - record.expected_ptx = policy.expected_ptx.clone(); - let resolved = resolve_record( - &policy, - resolve_policy_source(&policy).unwrap(), - Some(&declaration), - &record, - "test", - "LLVM version test", - "0123456789abcdef", - vec![], - 1, - ) - .unwrap(); - assert_eq!(resolved.selections.len(), 2); - assert_eq!(resolved.vote, policy.vote); - } - } - - #[test] - fn vote_contract_rejects_unreviewed_identity_effect_and_selection_changes() { - let valid = vote_policy(VoteMode::All); - let declaration = vote_declaration(VoteMode::All); - - let mut wrong_abi = valid.clone(); - wrong_abi.abi_id = "i0041".into(); - assert!( - validate_imported_policy(&wrong_abi, &declaration) - .unwrap_err() - .to_string() - .contains("vote identity") - ); - - let mut safe = valid.clone(); - safe.safe = true; - safe.safe_allowlist_reason = Some("incorrectly hides participation obligations".into()); - assert!( - validate_imported_policy(&safe, &declaration) - .unwrap_err() - .to_string() - .contains("unsafe must-use vote") - ); - - let mut wrong_memory = valid.clone(); - wrong_memory.memory = "none".into(); - assert!( - validate_imported_policy(&wrong_memory, &declaration) - .unwrap_err() - .to_string() - .contains("vote effects") - ); - - let mut register_only_mask = valid.clone(); - register_only_mask.expected_ptx.operands[2] = OperandPattern::Register; - assert!( - validate_imported_policy(®ister_only_mask, &declaration) - .unwrap_err() - .to_string() - .contains("expected PTX") - ); - - let mut one_selection = declaration.clone(); - one_selection.selections.pop(); - assert!( - validate_imported_policy(&valid, &one_selection) - .unwrap_err() - .to_string() - .contains("immediate/register selection pair") - ); - - let mut different_predicates = declaration; - different_predicates.selections[1].predicates[0] = - "Subtarget->getPTXVersion() >= 61".into(); - assert!( - validate_imported_policy(&valid, &different_predicates) - .unwrap_err() - .to_string() - .contains("disagree on PTX shape") - ); - } - - #[test] - fn uni_vote_is_raw_only_while_existing_votes_keep_compatibility_paths() { - for mode in [VoteMode::All, VoteMode::Any, VoteMode::Ballot] { - assert_eq!(vote_policy(mode).compatibility_rust_paths.len(), 1); - } - let uni = vote_policy(VoteMode::Uni); - assert!(uni.compatibility_rust_paths.is_empty()); - - let mut invented_compatibility_path = uni.clone(); - invented_compatibility_path.compatibility_rust_paths = - vec!["cuda_device::warp::uni_sync".into()]; - assert!( - validate_imported_policy( - &invented_compatibility_path, - &vote_declaration(VoteMode::Uni), - ) - .unwrap_err() - .to_string() - .contains("reviewed compatibility path") - ); - } - - #[test] - fn warp_shuffle_variants_keep_exact_identity_clamp_and_eight_selections() { - for (mode, value_kind, clamp) in [ - (WarpShuffleMode::Idx, WarpShuffleValueKind::I32, 31), - (WarpShuffleMode::Bfly, WarpShuffleValueKind::I32, 31), - (WarpShuffleMode::Down, WarpShuffleValueKind::I32, 31), - (WarpShuffleMode::Up, WarpShuffleValueKind::I32, 0), - (WarpShuffleMode::Idx, WarpShuffleValueKind::F32, 31), - (WarpShuffleMode::Bfly, WarpShuffleValueKind::F32, 31), - (WarpShuffleMode::Down, WarpShuffleValueKind::F32, 31), - (WarpShuffleMode::Up, WarpShuffleValueKind::F32, 0), - ] { - let policy = warp_shuffle_policy(mode, value_kind); - let declaration = warp_shuffle_declaration(mode, value_kind); - validate_imported_policy(&policy, &declaration).unwrap(); - - assert_eq!(policy.warp_shuffle.as_ref().unwrap().clamp, clamp); - assert_eq!( - declaration - .selections - .iter() - .filter(|selection| selection_matches_policy(&policy, selection).unwrap()) - .count(), - 8 - ); - - let mut record = evidence(); - record.id = policy.id.clone(); - record.source_record = policy.source_record.clone(); - record.llvm_symbol = policy.llvm_symbol.clone(); - record.llvm_arguments = policy.llvm_arguments.clone(); - record.llvm_results = policy.llvm_results.clone(); - record.expected_ptx = policy.expected_ptx.clone(); - let resolved = resolve_record( - &policy, - resolve_policy_source(&policy).unwrap(), - Some(&declaration), - &record, - "test", - "LLVM version test", - "0123456789abcdef", - vec![], - 1, - ) - .unwrap(); - assert_eq!(resolved.selections.len(), 8); - assert_eq!(resolved.warp_shuffle, policy.warp_shuffle); - } - } - - #[test] - fn pinned_warp_shuffle_records_match_the_closed_recipes() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let (overlay, _) = - read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")).unwrap(); - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - let declarations: BTreeMap<_, _> = imported - .intrinsics - .iter() - .map(|record| (record.source_record.as_str(), record)) - .collect(); - let all_policies: Vec<_> = overlay - .intrinsics - .iter() - .filter(|record| record.family == "warp_shuffle") - .collect(); - assert_eq!(all_policies.len(), 12); - let native_policies: Vec<_> = all_policies - .iter() - .copied() - .filter(|record| record.source.is_some()) - .collect(); - assert_eq!(native_policies.len(), 4); - for policy in native_policies { - validate_ptx_native_policy(policy).unwrap(); - } - let policies: Vec<_> = overlay - .intrinsics - .iter() - .filter(|record| record.family == "warp_shuffle" && record.source_record.is_some()) - .collect(); - assert_eq!(policies.len(), 8); - for policy in policies { - let declaration = declarations[policy.source_record.as_deref().unwrap()]; - validate_imported_policy(policy, declaration).unwrap(); - } - } - - #[test] - fn pinned_llvm_has_no_direct_i64_or_f64_shuffle_record() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - let direct_64: Vec<_> = imported - .intrinsics - .iter() - .filter(|record| { - record.llvm_name.starts_with("llvm.nvvm.shfl") - && (record.llvm_name.contains(".i64") || record.llvm_name.contains(".f64")) - }) - .map(|record| record.llvm_name.as_str()) - .collect(); - assert!( - direct_64.is_empty(), - "unexpected LLVM records: {direct_64:?}" - ); - } - - #[test] - fn i64_warp_shuffle_recipes_are_exact_ptx_native_pairs() { - let cases = [ - ( - WarpShuffleMode::Idx, - "shuffle_u64_sync", - "i0058", - "warp.shuffle.sync.idx.i64", - "shfl.sync.idx.b32", - "idx", - 31, - "ShflSyncIdxI64Op", - "nvvm.shfl_sync_idx_i64", - ), - ( - WarpShuffleMode::Bfly, - "shuffle_xor_u64_sync", - "i0059", - "warp.shuffle.sync.bfly.i64", - "shfl.sync.bfly.b32", - "bfly", - 31, - "ShflSyncBflyI64Op", - "nvvm.shfl_sync_bfly_i64", - ), - ( - WarpShuffleMode::Down, - "shuffle_down_u64_sync", - "i0060", - "warp.shuffle.sync.down.i64", - "shfl.sync.down.b32", - "down", - 31, - "ShflSyncDownI64Op", - "nvvm.shfl_sync_down_i64", - ), - ( - WarpShuffleMode::Up, - "shuffle_up_u64_sync", - "i0061", - "warp.shuffle.sync.up.i64", - "shfl.sync.up.b32", - "up", - 0, - "ShflSyncUpI64Op", - "nvvm.shfl_sync_up_i64", - ), - ]; - - for (mode, id, abi_id, operation_key, instruction, ptx_mode, clamp, op_type, op_name) in - cases - { - let policy = warp_shuffle_policy(mode, WarpShuffleValueKind::I64); - validate_ptx_native_policy(&policy).unwrap(); - - assert_eq!(policy.id, id); - assert_eq!(policy.abi_id, abi_id); - assert_eq!(policy.operation_key, operation_key); - assert_eq!( - policy.source, - Some(IntrinsicSource::PtxNative { - instruction: instruction.into(), - }) - ); - assert!(policy.source_record.is_none()); - assert!(policy.llvm_symbol.is_none()); - assert!(policy.resolved_llvm_symbol.is_none()); - assert!(policy.llvm_arguments.is_empty()); - assert!(policy.llvm_results.is_empty()); - assert_eq!(policy.rust_arguments, ["u32", "u64", "u32"]); - assert_eq!(policy.rust_result, "u64"); - assert!(!policy.safe); - assert!(policy.must_use); - assert_eq!(policy.dialect_op_type, op_type); - assert_eq!(policy.dialect_op_name, op_name); - assert_eq!(policy.dialect_operands, ["i32", "i64", "i32"]); - assert_eq!(policy.dialect_results, ["i64"]); - assert_eq!(policy.lowering, "generated_warp_shuffle_i64_inline_ptx"); - - let shuffle = policy.warp_shuffle.as_ref().unwrap(); - assert_eq!(shuffle.clamp, clamp); - assert_eq!( - shuffle.adapter, - WarpShuffleAdapter::MaskValueLaneOrDeltaSplitI64LowHighB32InsertClampReassemble - ); - assert_eq!( - shuffle.lane_encoding, - WarpShuffleOperandEncoding::RegisterOnly - ); - assert_eq!( - shuffle.mask_encoding, - WarpShuffleOperandEncoding::RegisterOnly - ); - assert_eq!( - policy.expected_ptx, - InstructionPattern::new( - "shfl", - &["sync", ptx_mode, "b32"], - vec![ - OperandPattern::Exact { value: "lo".into() }, - OperandPattern::Exact { value: "lo".into() }, - OperandPattern::Register, - OperandPattern::Exact { - value: clamp.to_string(), - }, - OperandPattern::Register, - ], - ) - ); - - let routes: BTreeMap<_, _> = policy - .backend_lowerings - .iter() - .map(|route| (route.backend, route)) - .collect(); - assert_eq!(routes.len(), 2); - for backend in [IntrinsicBackend::LlvmNvptx, IntrinsicBackend::LibNvvm] { - assert_eq!( - routes[&backend].mechanism, - BackendLoweringMechanism::InlinePtx - ); - assert_eq!(routes[&backend].minimum_ptx.as_deref(), Some("6.0")); - } - assert_eq!( - routes[&IntrinsicBackend::LlvmNvptx].minimum_sm.as_deref(), - Some("sm_30") - ); - assert_eq!( - routes[&IntrinsicBackend::LibNvvm].minimum_sm.as_deref(), - Some("sm_75") - ); - - let mut record = evidence(); - record.id = policy.id.clone(); - record.source = policy.source.clone(); - record.source_record = None; - record.llvm_symbol = None; - record.llvm_arguments.clear(); - record.llvm_results.clear(); - record.expected_ptx = policy.expected_ptx.clone(); - let resolved = resolve_record( - &policy, - resolve_policy_source(&policy).unwrap(), - None, - &record, - "test", - "LLVM version test", - "0123456789abcdef", - vec![], - 1, - ) - .unwrap(); - assert!(resolved.llvm.is_none()); - assert!(resolved.selections.is_empty()); - assert_eq!(resolved.warp_shuffle, policy.warp_shuffle); - } - } - - #[test] - fn i64_warp_shuffle_contract_rejects_unreviewed_changes() { - let valid = warp_shuffle_policy(WarpShuffleMode::Idx, WarpShuffleValueKind::I64); - validate_ptx_native_policy(&valid).unwrap(); - - let reject = |policy: &OverlayIntrinsic, expected: &str| { - let error = validate_ptx_native_policy(policy).unwrap_err(); - let message = error.to_string(); - assert!(message.contains(expected), "unexpected error: {message}"); - }; - - let mut fabricated_llvm = valid.clone(); - fabricated_llvm.source = None; - fabricated_llvm.source_record = Some("int_nvvm_shfl_sync_idx_i64".into()); - fabricated_llvm.llvm_symbol = Some("llvm.nvvm.shfl.sync.idx.i64".into()); - fabricated_llvm.llvm_arguments = - vec!["i32".into(), "i64".into(), "i32".into(), "i32".into()]; - fabricated_llvm.llvm_results = vec!["i64".into()]; - reject( - &fabricated_llvm, - "source kind and imported declaration disagree", - ); - - let mut wrong_source = valid.clone(); - wrong_source.source = Some(IntrinsicSource::PtxNative { - instruction: "shfl.sync.down.b32".into(), - }); - reject(&wrong_source, "warp-shuffle identity"); - - let mut wrong_adapter = valid.clone(); - wrong_adapter.warp_shuffle.as_mut().unwrap().adapter = - WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp; - reject(&wrong_adapter, "semantic or operand contract"); - - let mut wrong_mode = valid.clone(); - wrong_mode.warp_shuffle.as_mut().unwrap().mode = WarpShuffleMode::Up; - wrong_mode.warp_shuffle.as_mut().unwrap().clamp = 0; - wrong_mode.expected_ptx.modifiers[1] = "up".into(); - wrong_mode.expected_ptx.operands[3] = OperandPattern::Exact { value: "0".into() }; - reject(&wrong_mode, "warp-shuffle identity"); - - let mut wrong_clamp = valid.clone(); - wrong_clamp.warp_shuffle.as_mut().unwrap().clamp = 0; - reject(&wrong_clamp, "semantic or operand contract"); - - let mut broad_encoding = valid.clone(); - broad_encoding.warp_shuffle.as_mut().unwrap().lane_encoding = - WarpShuffleOperandEncoding::RegisterOrImmediate; - reject(&broad_encoding, "semantic or operand contract"); - - let mut typed_backend = valid.clone(); - typed_backend.backend_lowerings[0].mechanism = BackendLoweringMechanism::TypedNvvm; - reject(&typed_backend, "reviewed LLVM and libNVVM routes"); - - let mut wrong_native_floor = valid.clone(); - wrong_native_floor.minimum_sm = Some("sm_70".into()); - reject(&wrong_native_floor, "target floor"); - - let mut wrong_profile_floor = valid.clone(); - wrong_profile_floor - .backend_lowerings - .iter_mut() - .find(|route| route.backend == IntrinsicBackend::LibNvvm) - .unwrap() - .minimum_sm = Some("sm_80".into()); - reject(&wrong_profile_floor, "profile floor"); - - let mut safe = valid.clone(); - safe.safe = true; - safe.safe_allowlist_reason = Some("incorrectly hides participation obligations".into()); - reject(&safe, "unsafe must-use warp-shuffle"); - - let mut wrong_ptx = valid; - wrong_ptx.expected_ptx.operands[0] = OperandPattern::Register; - reject(&wrong_ptx, "closed shfl.sync recipe"); - } - - #[test] - fn warp_shuffle_contract_rejects_unreviewed_policy_changes() { - let valid = warp_shuffle_policy(WarpShuffleMode::Idx, WarpShuffleValueKind::I32); - let declaration = warp_shuffle_declaration(WarpShuffleMode::Idx, WarpShuffleValueKind::I32); - - let reject_policy = |policy: &OverlayIntrinsic, expected: &str| { - let error = match validate_imported_policy(policy, &declaration) { - Ok(()) => panic!("{expected} mutation was accepted"), - Err(error) => error, - }; - let message = error.to_string(); - assert!(message.contains(expected), "unexpected error: {message}"); - }; - - let mut wrong_identity = valid.clone(); - wrong_identity.operation_key = "warp.shuffle.sync.idx.changed".into(); - reject_policy(&wrong_identity, "warp-shuffle identity"); - - let mut safe = valid.clone(); - safe.safe = true; - safe.safe_allowlist_reason = Some("incorrectly hides participation obligations".into()); - reject_policy(&safe, "unsafe must-use warp-shuffle"); - - let mut wrong_signature = valid.clone(); - wrong_signature.dialect_operands.pop(); - reject_policy(&wrong_signature, "closed warp-shuffle lowering recipe"); - - let mut wrong_clamp = valid.clone(); - wrong_clamp.warp_shuffle.as_mut().unwrap().clamp = 0; - reject_policy(&wrong_clamp, "semantic or operand contract"); - - let mut missing_contract = valid.clone(); - missing_contract.warp_shuffle = None; - reject_policy(&missing_contract, "closed warp-shuffle contract"); - - let mut mixed_contract = valid.clone(); - mixed_contract.vote = vote_policy(VoteMode::All).vote; - reject_policy(&mixed_contract, "mixes another generated-family contract"); - - let mut wrong_backend_floor = valid.clone(); - wrong_backend_floor - .backend_lowerings - .iter_mut() - .find(|lowering| lowering.backend == IntrinsicBackend::LibNvvm) - .unwrap() - .minimum_sm = Some("sm_80".into()); - reject_policy(&wrong_backend_floor, "profile floor"); - } - - #[test] - fn warp_shuffle_contract_rejects_selection_drift() { - let valid = warp_shuffle_policy(WarpShuffleMode::Down, WarpShuffleValueKind::F32); - let declaration = - warp_shuffle_declaration(WarpShuffleMode::Down, WarpShuffleValueKind::F32); - let reject = |declaration: &ImportedIntrinsic, expected: &str| { - let error = validate_imported_policy(&valid, declaration).unwrap_err(); - let message = error.to_string(); - assert!(message.contains(expected), "unexpected error: {message}"); - }; - - let mut missing_selection = declaration.clone(); - missing_selection.selections.pop(); - reject( - &missing_selection, - "eight distinct operand-encoding selections", - ); - - let mut duplicate_selection = declaration.clone(); - duplicate_selection.selections[7].source_record = - duplicate_selection.selections[0].source_record.clone(); - reject( - &duplicate_selection, - "eight distinct operand-encoding selections", - ); - - let mut empty_selection_name = declaration.clone(); - empty_selection_name.selections[7].source_record.clear(); - reject( - &empty_selection_name, - "eight distinct operand-encoding selections", - ); - - let mut wrong_asm = declaration.clone(); - wrong_asm.selections[0].asm = - "shfl.sync.up.b32 \t$dst, $src, $offset, $mask, $threadmask;".into(); - reject(&wrong_asm, "selections disagree on PTX shape"); - - let mut wrong_predicate = declaration.clone(); - wrong_predicate.selections[0].predicates[0] = "Subtarget->getPTXVersion() >= 61".into(); - reject(&wrong_predicate, "selections disagree on PTX shape"); - - let mut constrained = declaration; - constrained.selections[0] - .constraints - .immediate_bindings - .push(crate::model::ImportedImmediateBinding { - argument_index: 2, - value: 1, - }); - reject(&constrained, "selections disagree on PTX shape"); - - let mut wrong_classes = - warp_shuffle_declaration(WarpShuffleMode::Down, WarpShuffleValueKind::F32); - wrong_classes.classes.pop(); - reject(&wrong_classes, "class or effects"); - } - - #[test] - fn sync_threads_selects_only_the_fixed_immediate_barrier_recipe() { - let policy = sync_policy(); - let declaration = sync_declaration(); - validate_imported_policy(&policy, &declaration).unwrap(); - - let selected: Vec<_> = declaration - .selections - .iter() - .filter(|selection| selection_matches_policy(&policy, selection).unwrap()) - .collect(); - assert_eq!(selected.len(), 1); - assert_eq!(selected[0].source_record, "BARRIER_CTA_SYNC_ALIGNED_ALL_i"); - assert_eq!(selected[0].asm, "bar.sync \t$i;"); - assert!(policy.expected_ptx.matches("bar.sync 0;").unwrap()); - assert!(!policy.expected_ptx.matches(&selected[0].asm).unwrap()); - assert_eq!(policy.minimum_ptx, "1.0"); - assert!(policy.minimum_sm.is_none()); - let llvm_route = policy - .backend_lowerings - .iter() - .find(|lowering| lowering.backend == IntrinsicBackend::LlvmNvptx) - .unwrap(); - assert_eq!(llvm_route.minimum_ptx.as_deref(), Some("3.2")); - assert_eq!(llvm_route.minimum_sm.as_deref(), Some("sm_20")); - - let resolved = resolve_record( - &policy, - resolve_policy_source(&policy).unwrap(), - Some(&declaration), - &sync_evidence(&policy), - "test", - "LLVM version test", - "0123456789abcdef", - vec![], - 1, - ) - .unwrap(); - assert!(resolved.dialect.operands.is_empty()); - assert!(resolved.dialect.results.is_empty()); - assert_eq!(resolved.selections.len(), 1); - assert_eq!( - resolved.selections[0].source_record, - "BARRIER_CTA_SYNC_ALIGNED_ALL_i" - ); - } - - #[test] - fn sync_threads_recipe_rejects_unreviewed_selection_effect_and_floor_changes() { - let valid = sync_policy(); - let declaration = sync_declaration(); - - let mut register_only = declaration.clone(); - register_only - .selections - .retain(|selection| selection.source_record.ends_with("_r")); - assert!( - validate_imported_policy(&valid, ®ister_only) - .unwrap_err() - .to_string() - .contains("does not agree") - ); - - let mut wrong_properties = declaration.clone(); - wrong_properties.properties.pop(); - assert!( - validate_imported_policy(&valid, &wrong_properties) - .unwrap_err() - .to_string() - .contains("sync properties") - ); - - let mut wrong_source = valid.clone(); - wrong_source.source_record = Some("int_nvvm_barrier0".into()); - assert!( - validate_imported_policy(&wrong_source, &declaration) - .unwrap_err() - .to_string() - .contains("sync identity") - ); - - let mut wrong_signature = valid.clone(); - wrong_signature.llvm_arguments.clear(); - assert!( - validate_imported_policy(&wrong_signature, &declaration) - .unwrap_err() - .to_string() - .contains("LLVM argument signature mismatch") - ); - - let mut wrong_path = valid.clone(); - wrong_path.compatibility_rust_paths.swap(0, 1); - assert!( - validate_imported_policy(&wrong_path, &declaration) - .unwrap_err() - .to_string() - .contains("both cuda-device compatibility paths") - ); - - let mut safe = valid.clone(); - safe.safe = true; - safe.safe_allowlist_reason = Some("incorrectly hides the participation contract".into()); - assert!( - validate_imported_policy(&safe, &declaration) - .unwrap_err() - .to_string() - .contains("unsafe sync_threads raw API") - ); - - let mut wrong_effect = valid.clone(); - wrong_effect.memory = "none".into(); - assert!( - validate_imported_policy(&wrong_effect, &declaration) - .unwrap_err() - .to_string() - .contains("sync effects") - ); - - let mut native_floor = valid.clone(); - native_floor.minimum_sm = Some("sm_75".into()); - assert!( - validate_imported_policy(&native_floor, &declaration) - .unwrap_err() - .to_string() - .contains("native target floor") - ); - - let mut missing_profile_floor = valid; - missing_profile_floor - .backend_lowerings - .iter_mut() - .find(|lowering| lowering.backend == IntrinsicBackend::LibNvvm) - .unwrap() - .minimum_sm = None; - assert!( - validate_imported_policy(&missing_profile_floor, &declaration) - .unwrap_err() - .to_string() - .contains("profile floor") - ); - - let mut wrong_llvm_floor = sync_policy(); - wrong_llvm_floor - .backend_lowerings - .iter_mut() - .find(|lowering| lowering.backend == IntrinsicBackend::LlvmNvptx) - .unwrap() - .minimum_ptx = None; - assert!( - validate_imported_policy(&wrong_llvm_floor, &declaration) - .unwrap_err() - .to_string() - .contains("profile floor") - ); - } - - #[test] - fn sync_mask_matches_the_closed_warp_barrier_recipe() { - let policy = warp_barrier_policy(); - let declaration = warp_barrier_declaration(); - validate_imported_policy(&policy, &declaration).unwrap(); - - let selected: Vec<_> = declaration - .selections - .iter() - .filter(|selection| selection_matches_policy(&policy, selection).unwrap()) - .collect(); - assert_eq!(selected.len(), 2); - assert_eq!( - selected - .iter() - .map(|selection| selection.source_record.as_str()) - .collect::>(), - BTreeSet::from(["INT_BAR_WARP_SYNC_I", "INT_BAR_WARP_SYNC_R"]) - ); - - let mut record = evidence(); - record.id = policy.id.clone(); - record.source_record = policy.source_record.clone(); - record.llvm_symbol = policy.llvm_symbol.clone(); - record.llvm_arguments = policy.llvm_arguments.clone(); - record.llvm_results = policy.llvm_results.clone(); - record.expected_ptx = policy.expected_ptx.clone(); - let resolved = resolve_record( - &policy, - resolve_policy_source(&policy).unwrap(), - Some(&declaration), - &record, - "test", - "LLVM version test", - "0123456789abcdef", - vec![], - 1, - ) - .unwrap(); - assert_eq!(resolved.selections.len(), 2); - assert_eq!(resolved.warp_barrier, policy.warp_barrier); - } - - #[test] - fn sync_mask_recipe_rejects_unreviewed_contract_and_selection_changes() { - let valid = warp_barrier_policy(); - let declaration = warp_barrier_declaration(); - - let mut wrong_identity = valid.clone(); - wrong_identity.id = "bar_warp_sync".into(); - assert!( - validate_imported_policy(&wrong_identity, &declaration) - .unwrap_err() - .to_string() - .contains("warp-barrier identity") - ); - - let mut missing_contract = valid.clone(); - missing_contract.warp_barrier = None; - assert!( - validate_imported_policy(&missing_contract, &declaration) - .unwrap_err() - .to_string() - .contains("closed warp-barrier contract") - ); - - let mut safe_raw_api = valid.clone(); - safe_raw_api.safe = true; - safe_raw_api.safe_allowlist_reason = Some("incorrectly hides participation rules".into()); - assert!( - validate_imported_policy(&safe_raw_api, &declaration) - .unwrap_err() - .to_string() - .contains("unsafe raw API") - ); - - let mut wrong_memory = valid.clone(); - wrong_memory.memory = "none".into(); - assert!( - validate_imported_policy(&wrong_memory, &declaration) - .unwrap_err() - .to_string() - .contains("effects or target floor") - ); - - let mut register_only = valid.clone(); - register_only.expected_ptx.operands[0] = OperandPattern::Register; - assert!( - validate_imported_policy(®ister_only, &declaration) - .unwrap_err() - .to_string() - .contains("expected PTX") - ); - - let mut one_selection = declaration.clone(); - one_selection.selections.pop(); - assert!( - validate_imported_policy(&valid, &one_selection) - .unwrap_err() - .to_string() - .contains("immediate/register selection pair") - ); - - let mut wrong_predicate = declaration.clone(); - wrong_predicate.selections[1].predicates[0] = "Subtarget->getPTXVersion() >= 61".into(); - assert!( - validate_imported_policy(&valid, &wrong_predicate) - .unwrap_err() - .to_string() - .contains("selections disagree") - ); - - let mut missing_libnvvm_floor = valid; - missing_libnvvm_floor - .backend_lowerings - .iter_mut() - .find(|lowering| lowering.backend == IntrinsicBackend::LibNvvm) - .unwrap() - .minimum_sm = None; - assert!( - validate_imported_policy(&missing_libnvvm_floor, &declaration) - .unwrap_err() - .to_string() - .contains("profile floor") - ); - } - - #[test] - fn every_redux_variant_matches_its_closed_recipe() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let (overlay, _) = - read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")).unwrap(); - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - assert_eq!(imported.schema, IMPORTED_SCHEMA); - let declarations: BTreeMap<_, _> = imported - .intrinsics - .iter() - .map(|record| (record.source_record.as_str(), record)) - .collect(); - let redux: Vec<_> = overlay - .intrinsics - .iter() - .filter(|record| record.family == "redux") - .collect(); - assert_eq!(redux.len(), 16); - - for policy in redux { - let declaration = declarations - .get(policy.source_record.as_deref().unwrap()) - .unwrap(); - validate_imported_policy(policy, declaration).unwrap(); - } - - let mut mismatched = packed_policy("redux_sync_min_u32"); - mismatched.redux.as_mut().unwrap().operation = ReduxOperation::Umax; - let declaration = declarations["int_nvvm_redux_sync_umin"]; - assert!( - validate_imported_policy(&mismatched, declaration) - .unwrap_err() - .to_string() - .contains("closed operation recipe") - ); - } - - #[test] - fn every_dot_product_variant_matches_its_closed_recipe() { - let variants = [ - ( - DotProductOperation::Dp4a, - DotProductSignedness::Signed, - "dp4a_s32", - "int_nvvm_idp4a_s_s", - "integer.dot_product.dp4a.s32", - ), - ( - DotProductOperation::Dp4a, - DotProductSignedness::Unsigned, - "dp4a_u32", - "int_nvvm_idp4a_u_u", - "integer.dot_product.dp4a.u32", - ), - ( - DotProductOperation::Dp2a, - DotProductSignedness::Signed, - "dp2a_s32", - "int_nvvm_idp2a_s_s", - "integer.dot_product.dp2a.lo.s32", - ), - ( - DotProductOperation::Dp2a, - DotProductSignedness::Unsigned, - "dp2a_u32", - "int_nvvm_idp2a_u_u", - "integer.dot_product.dp2a.lo.u32", - ), - ]; - - for (operation, signedness, id, source_record, operation_key) in variants { - let policy = dot_product_policy(operation, signedness); - let declaration = dot_product_declaration(operation, signedness); - assert_eq!(policy.id, id); - assert_eq!(policy.source_record.as_deref(), Some(source_record)); - assert_eq!(policy.operation_key, operation_key); - validate_imported_policy(&policy, &declaration).unwrap(); - } - } - - #[test] - fn pinned_dot_product_records_match_the_reviewed_overlay() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let (overlay, _) = - read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")).unwrap(); - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - let declarations: BTreeMap<_, _> = imported - .intrinsics - .iter() - .map(|record| (record.source_record.as_str(), record)) - .collect(); - let dot_products: Vec<_> = overlay - .intrinsics - .iter() - .filter(|record| record.family == "dotprod") - .collect(); - assert_eq!(dot_products.len(), 4); - - for policy in dot_products { - let declaration = declarations[policy.source_record.as_deref().unwrap()]; - validate_imported_policy(policy, declaration).unwrap(); - let selected: Vec<_> = declaration - .selections - .iter() - .filter(|selection| selection_matches_policy(policy, selection).unwrap()) - .collect(); - assert_eq!(selected.len(), 1); - if policy.id.starts_with("dp2a") { - assert_eq!(selected[0].constraints.immediate_bindings[0].value, 0); - } - } - } - - #[test] - fn dp2a_selects_only_the_reviewed_low_half_binding() { - let policy = dot_product_policy(DotProductOperation::Dp2a, DotProductSignedness::Signed); - let declaration = - dot_product_declaration(DotProductOperation::Dp2a, DotProductSignedness::Signed); - let resolved = resolve_record( - &policy, - resolve_policy_source(&policy).unwrap(), - Some(&declaration), - &dot_product_evidence(&policy), - "test", - "LLVM version test", - "0123456789abcdef", - vec![], - 1, - ) - .unwrap(); - - assert_eq!(resolved.selections.len(), 1); - assert_eq!(resolved.selections[0].source_record, "DOT2_lo"); - assert_eq!( - resolved.selections[0].constraints.immediate_bindings, - [crate::model::ImportedImmediateBinding { - argument_index: 2, - value: 0, - }] - ); - assert_eq!( - resolved.dot_product.as_ref().unwrap().adapter, - DotProductAdapter::InsertLowHalfFalse - ); - - let mut wrong_binding = declaration; - wrong_binding.selections[1].constraints.immediate_bindings[0].value = -1; - let error = validate_imported_policy(&policy, &wrong_binding).unwrap_err(); - assert!(error.to_string().contains("does not agree")); - } - - #[test] - fn dot_product_recipe_rejects_unreviewed_api_and_adapter_changes() { - let valid = dot_product_policy(DotProductOperation::Dp2a, DotProductSignedness::Unsigned); - let declaration = - dot_product_declaration(DotProductOperation::Dp2a, DotProductSignedness::Unsigned); - - let mut wrong_adapter = valid.clone(); - wrong_adapter.dot_product.as_mut().unwrap().adapter = - DotProductAdapter::DirectThreeOperands; - assert!( - validate_imported_policy(&wrong_adapter, &declaration) - .unwrap_err() - .to_string() - .contains("source adapter") - ); - - let mut must_use = valid.clone(); - must_use.must_use = true; - assert!( - validate_imported_policy(&must_use, &declaration) - .unwrap_err() - .to_string() - .contains("non-must-use") - ); - - let mut wrong_llvm_signature = valid; - wrong_llvm_signature.llvm_arguments = vec!["i32".into(); 3]; - assert!( - validate_imported_policy(&wrong_llvm_signature, &declaration) - .unwrap_err() - .to_string() - .contains("LLVM argument signature mismatch") - ); - } - - #[test] - fn dot_product_target_predicate_is_closed_to_ptx50_and_sm61() { - let policy = dot_product_policy(DotProductOperation::Dp4a, DotProductSignedness::Signed); - let selection = - &dot_product_declaration(DotProductOperation::Dp4a, DotProductSignedness::Signed) - .selections[0]; - validate_selected_target_predicates(&policy, selection).unwrap(); - - let mut wrong_ptx = policy.clone(); - wrong_ptx.minimum_ptx = "5.1".into(); - assert!( - validate_selected_target_predicates(&wrong_ptx, selection) - .unwrap_err() - .to_string() - .contains("minimum PTX") - ); - - let mut wrong_sm = policy; - wrong_sm.minimum_sm = Some("sm_60".into()); - assert!( - validate_selected_target_predicates(&wrong_sm, selection) - .unwrap_err() - .to_string() - .contains("minimum SM") - ); - } - - #[test] - fn typed_evidence_accepts_direct_scalar_intrinsic_signatures() { - let policy = dot_product_policy(DotProductOperation::Dp2a, DotProductSignedness::Signed); - let mut record = dot_product_evidence(&policy); - validate_typed_llvm_evidence(&policy, &record).unwrap(); - - record.concrete_llvm_arguments.remove(2); - let error = validate_typed_llvm_evidence(&policy, &record).unwrap_err(); - assert!(error.to_string().contains("resolved signature")); - } - - #[test] - fn typed_evidence_maps_cluster_shared_pointers_to_address_space_seven() { - let mut policy = - dot_product_policy(DotProductOperation::Dp2a, DotProductSignedness::Signed); - let mut record = dot_product_evidence(&policy); - policy.llvm_arguments = vec![ - "shared_cluster_ptr".into(), - "shared_ptr".into(), - "ptr".into(), - ]; - record.concrete_llvm_arguments = vec![ - "ptr addrspace(7)".into(), - "ptr addrspace(3)".into(), - "ptr".into(), - ]; - - validate_typed_llvm_evidence(&policy, &record).unwrap(); - } - - #[test] - fn packed_conversion_evidence_separates_llvm_declaration_facts_from_libnvvm() { - for (destination, result) in [ - (PackedConversionDestinationFormat::Bf16x2, "v2bf16"), - (PackedConversionDestinationFormat::F16x2, "v2f16"), - ] { - let policy = packed_conversion_policy( - destination, - PackedConversionRounding::NearestEven, - PackedConversionSaturation::None, - ); - let llvm = policy - .backend_lowerings - .iter() - .find(|lowering| lowering.backend == IntrinsicBackend::LlvmNvptx) - .unwrap(); - let mut record = packed_conversion_evidence(&policy); - record.status = "validated".into(); - record.stages = [ - EvidenceStageKind::DeclarationCanonicalization, - EvidenceStageKind::BackendCodegen, - EvidenceStageKind::PtxAssembly, - ] - .into_iter() - .map(|stage| { - evidence_stage( - stage, - BackendLoweringMechanism::TypedNvvm, - &["sm_80", "ptx70"], - ) - }) - .collect(); - let assembly = record - .stages - .iter_mut() - .find(|stage| stage.stage == EvidenceStageKind::PtxAssembly) - .unwrap(); - assembly.tool_path = Some("/usr/local/cuda/bin/ptxas".into()); - assembly.tool_version = Some("CUDA 13.3 V13.3.33".into()); - assembly.tool_sha256 = - Some("7fdd01a4cf50e30746da98989c9272a907f491e6fd7fecfda14642e4375f88fb".into()); - assert_eq!(record.concrete_llvm_results, [result]); - validate_packed_conversion_backend_evidence(&policy, &record, llvm).unwrap(); - - let mut lowered = record.clone(); - lowered.status = "lowered".into(); - let error = - validate_packed_conversion_backend_evidence(&policy, &lowered, llvm).unwrap_err(); - assert!( - error.to_string().contains("validated evidence status"), - "{error:#}" - ); - - for required in [ - EvidenceStageKind::DeclarationCanonicalization, - EvidenceStageKind::BackendCodegen, - EvidenceStageKind::PtxAssembly, - ] { - let mut missing = record.clone(); - missing.stages.retain(|stage| stage.stage != required); - let error = validate_packed_conversion_backend_evidence(&policy, &missing, llvm) - .unwrap_err(); - assert!( - error - .to_string() - .contains("successful auxiliary typed-NVVM"), - "{error:#}" - ); - - let mut failed = record.clone(); - failed - .stages - .iter_mut() - .find(|stage| stage.stage == required) - .unwrap() - .outcome = "failed".into(); - let error = validate_packed_conversion_backend_evidence(&policy, &failed, llvm) - .unwrap_err(); - assert!( - error - .to_string() - .contains("successful auxiliary typed-NVVM"), - "{error:#}" - ); - - let mut wrong_mechanism = record.clone(); - wrong_mechanism - .stages - .iter_mut() - .find(|stage| stage.stage == required) - .unwrap() - .mechanism = Some(BackendLoweringMechanism::InlinePtx); - let error = - validate_packed_conversion_backend_evidence(&policy, &wrong_mechanism, llvm) - .unwrap_err(); - assert!( - error - .to_string() - .contains("successful auxiliary typed-NVVM"), - "{error:#}" - ); - } - - let mut missing_tool_identity = record.clone(); - missing_tool_identity - .stages - .iter_mut() - .find(|stage| stage.stage == EvidenceStageKind::PtxAssembly) - .unwrap() - .tool_sha256 = None; - let error = - validate_packed_conversion_backend_evidence(&policy, &missing_tool_identity, llvm) - .unwrap_err(); - assert!( - error.to_string().contains("exact tool identity"), - "{error:#}" - ); - - for stage_kind in [ - EvidenceStageKind::BackendCodegen, - EvidenceStageKind::PtxAssembly, - ] { - let mut wrong_floor = record.clone(); - wrong_floor - .stages - .iter_mut() - .find(|stage| stage.stage == stage_kind) - .unwrap() - .targets = vec!["sm_75".into(), "ptx70".into()]; - let error = - validate_packed_conversion_backend_evidence(&policy, &wrong_floor, llvm) - .unwrap_err(); - assert!( - error.to_string().contains("catalog floor sm_80"), - "{error:#}" - ); - } - - record.declaration_attributes_canonicalized = None; - let error = - validate_packed_conversion_backend_evidence(&policy, &record, llvm).unwrap_err(); - assert!( - error - .to_string() - .contains("canonical declaration attributes") - ); - } - - let policy = packed_conversion_policy( - PackedConversionDestinationFormat::Bf16x2, - PackedConversionRounding::NearestEven, - PackedConversionSaturation::None, - ); - let libnvvm = policy - .backend_lowerings - .iter() - .find(|lowering| lowering.backend == IntrinsicBackend::LibNvvm) - .unwrap(); - let mut record = packed_conversion_evidence(&policy); - record.concrete_llvm_arguments.clear(); - record.concrete_llvm_results.clear(); - record.declaration_attributes_canonicalized = None; - validate_packed_conversion_backend_evidence(&policy, &record, libnvvm).unwrap(); - - record.concrete_llvm_arguments = policy.llvm_arguments.clone(); - let error = - validate_packed_conversion_backend_evidence(&policy, &record, libnvvm).unwrap_err(); - assert!( - error - .to_string() - .contains("must not claim typed LLVM support") - ); - - record.concrete_llvm_arguments.clear(); - record.stages.push(evidence_stage( - EvidenceStageKind::BackendCodegen, - BackendLoweringMechanism::TypedNvvm, - &["sm_80", "ptx70"], - )); - let error = - validate_packed_conversion_backend_evidence(&policy, &record, libnvvm).unwrap_err(); - assert!( - error - .to_string() - .contains("must not claim typed LLVM support") - ); - } - - #[test] - fn duplicate_identity_surfaces_are_rejected_independently() { - let first = policy(); - - let mut second = distinct_policy(); - second.abi_id = first.abi_id.clone(); - assert!( - validate_unique_overlay(&[first.clone(), second], 1) - .unwrap_err() - .to_string() - .contains("duplicate intrinsic ABI ID") - ); - - let mut second = distinct_policy(); - second.operation_key = first.operation_key.clone(); - assert!( - validate_unique_overlay(&[first.clone(), second], 1) - .unwrap_err() - .to_string() - .contains("duplicate intrinsic operation key") - ); - - let mut second = distinct_policy(); - second.public_rust_path = first.public_rust_path.clone(); - assert!( - validate_unique_overlay(&[first.clone(), second], 1) - .unwrap_err() - .to_string() - .contains("duplicate public Rust path") - ); - - let mut second = distinct_policy(); - second.dialect_op_name = first.dialect_op_name.clone(); - assert!( - validate_unique_overlay(&[first.clone(), second], 1) - .unwrap_err() - .to_string() - .contains("duplicate dialect op variant") - ); - - let mut second = distinct_policy(); - second.llvm_symbol = first.llvm_symbol.clone(); - assert!( - validate_unique_overlay(&[first, second], 1) - .unwrap_err() - .to_string() - .contains("duplicate LLVM symbol") - ); - } - - #[test] - fn signature_and_evidence_mismatches_are_rejected() { - let mut imported = declaration(); - imported.results = vec!["i64".into()]; - assert!( - validate_imported_policy(&policy(), &imported) - .unwrap_err() - .to_string() - .contains("LLVM result signature mismatch") - ); - - let mut backend_evidence = evidence(); - backend_evidence.llvm_results = vec!["i64".into()]; - assert!( - validate_test_evidence(&policy(), backend_evidence) - .unwrap_err() - .to_string() - .contains("evidence signature mismatch") - ); - - let mut backend_evidence = evidence(); - backend_evidence.expected_ptx = sreg_pattern("%tid.y"); - assert!( - validate_test_evidence(&policy(), backend_evidence) - .unwrap_err() - .to_string() - .contains("evidence PTX expectation mismatch") - ); - } - - #[test] - fn validated_llvm_evidence_requires_exact_ptxas_identity() { - let mut record = evidence(); - record.status = "validated".into(); - record.stages.push(crate::model::EvidenceStage { - targets: vec!["sm_75".into()], - representation: "probe PTX".into(), - stage: EvidenceStageKind::PtxAssembly, - mechanism: Some(BackendLoweringMechanism::TypedNvvm), - outcome: "succeeded".into(), - detail: "accepted".into(), - artifact_kind: None, - tool_path: Some("/usr/local/cuda/bin/ptxas".into()), - tool_version: Some("CUDA 13.3 V13.3.33".into()), - tool_sha256: Some( - "7fdd01a4cf50e30746da98989c9272a907f491e6fd7fecfda14642e4375f88fb".into(), - ), - }); - assert!(has_valid_ptx_assembly_stage( - &record, - BackendLoweringMechanism::TypedNvvm - )); - - let stage = record.stages.last_mut().unwrap(); - stage.tool_path = None; - assert!(!has_valid_ptx_assembly_stage( - &record, - BackendLoweringMechanism::TypedNvvm - )); - record.stages.clear(); - assert!(!has_valid_ptx_assembly_stage( - &record, - BackendLoweringMechanism::TypedNvvm - )); - } - - #[test] - fn validated_libnvvm_evidence_requires_a_real_cubin_terminal() { - let mut record = evidence(); - record.stages.push(crate::model::EvidenceStage { - targets: vec!["sm_90".into(), "ptx78".into()], - representation: "linked output".into(), - stage: EvidenceStageKind::DeviceLink, - mechanism: Some(BackendLoweringMechanism::InlinePtx), - outcome: "succeeded".into(), - detail: "test".into(), - artifact_kind: None, - tool_path: Some("/usr/local/cuda-13.3/lib64/libnvJitLink.so.13.3.33".into()), - tool_version: Some("V13.3.33".into()), - tool_sha256: Some( - "3ba1e744347cd68617b862eccfd98b125482e882b7a6319f42abc9a768513db8".into(), - ), - }); - assert!(!has_valid_cubin_device_link_stage( - &record, - BackendLoweringMechanism::InlinePtx - )); - record.stages[0].artifact_kind = Some(EvidenceArtifactKind::Cubin); - assert!(has_valid_cubin_device_link_stage( - &record, - BackendLoweringMechanism::InlinePtx - )); - } - - fn evidence_stage( - stage: EvidenceStageKind, - mechanism: BackendLoweringMechanism, - targets: &[&str], - ) -> crate::model::EvidenceStage { - crate::model::EvidenceStage { - targets: targets.iter().map(|target| (*target).into()).collect(), - representation: "test".into(), - stage, - mechanism: Some(mechanism), - outcome: "succeeded".into(), - detail: "test".into(), - artifact_kind: None, - tool_path: None, - tool_version: None, - tool_sha256: None, - } - } - - #[test] - fn backend_stage_targets_and_executed_status_are_monotonic() { - let mut target_policy = policy(); - target_policy.minimum_ptx = "6.5".into(); - target_policy.minimum_sm = Some("sm_75".into()); - let lowering = crate::model::OverlayBackendLowering { - backend: IntrinsicBackend::LlvmNvptx, - mechanism: BackendLoweringMechanism::TypedNvvm, - evidence_profile: "test".into(), - targets: None, - minimum_ptx: None, - minimum_sm: None, - }; - let mut record = evidence(); - record.status = "validated".into(); - record.runtime_validation = Some(RuntimeValidation::Unexecuted); - record.stages = vec![ - evidence_stage( - EvidenceStageKind::BackendCodegen, - BackendLoweringMechanism::TypedNvvm, - &["sm_75", "ptx65"], - ), - evidence_stage( - EvidenceStageKind::PtxAssembly, - BackendLoweringMechanism::TypedNvvm, - &["sm_75", "ptx65"], - ), - ]; - validate_selected_stage_targets(&target_policy, &record, &lowering).unwrap(); - - record.stages[0].targets = vec!["sm_75a".into(), "ptx65".into()]; - assert!(validate_selected_stage_targets(&target_policy, &record, &lowering).is_err()); - record.stages[0].targets = vec!["sm_75".into(), "ptx65".into()]; - - record.stages[1].targets = vec!["sm_80".into(), "ptx65".into()]; - validate_selected_stage_targets(&target_policy, &record, &lowering).unwrap(); - - record.stages[1].targets = vec!["sm_90a".into(), "ptx65".into()]; - validate_selected_stage_targets(&target_policy, &record, &lowering).unwrap(); - - record.stages[1].targets = vec!["sm_74".into(), "ptx65".into()]; - assert!( - validate_selected_stage_targets(&target_policy, &record, &lowering) - .unwrap_err() - .to_string() - .contains("catalog floor sm_75") - ); - - record.stages[1].targets = vec!["sm_75".into(), "ptx65".into()]; - record.status = "executed".into(); - record.runtime_validation = Some(RuntimeValidation::Executed); - assert!( - validate_selected_stage_targets(&target_policy, &record, &lowering) - .unwrap_err() - .to_string() - .contains("runtime stage") - ); - } - - #[test] - fn exact_and_family_evidence_targets_match_at_every_stage() { - let lowering = crate::model::OverlayBackendLowering { - backend: IntrinsicBackend::LlvmNvptx, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: "test".into(), - targets: None, - minimum_ptx: None, - minimum_sm: None, - }; - - for (target, wrong_targets) in [ - ("sm_120a", ["sm_120", "sm_120f", "sm_121a"]), - ("sm_120f", ["sm_120", "sm_120a", "sm_121f"]), - ] { - let mut target_policy = policy(); - target_policy.minimum_ptx = "8.7".into(); - target_policy.targets = target.into(); - let mut record = evidence(); - record.status = "validated".into(); - record.runtime_validation = Some(RuntimeValidation::Unexecuted); - record.stages = vec![ - evidence_stage( - EvidenceStageKind::BackendCodegen, - BackendLoweringMechanism::InlinePtx, - &[target, "ptx87"], - ), - evidence_stage( - EvidenceStageKind::PtxAssembly, - BackendLoweringMechanism::InlinePtx, - &[target, "ptx87"], - ), - ]; - validate_selected_stage_targets(&target_policy, &record, &lowering).unwrap(); - - for wrong in wrong_targets { - record.stages[0].targets = vec![wrong.into(), "ptx87".into()]; - let error = validate_selected_stage_targets(&target_policy, &record, &lowering) - .unwrap_err() - .to_string(); - assert!(error.contains(target), "{error}"); - } - record.stages[0].targets = vec![target.into(), "ptx87".into()]; - for wrong in wrong_targets { - record.stages[1].targets = vec![wrong.into(), "ptx87".into()]; - let error = validate_selected_stage_targets(&target_policy, &record, &lowering) - .unwrap_err() - .to_string(); - assert!(error.contains(target), "{error}"); - } - } - } - - #[test] - fn suffixed_evidence_target_spellings_are_normalized() { - for target in ["sm_120a", "compute_120a", "sm_120f", "compute_120f"] { - assert!(is_normalized_stage_target(target), "{target}"); - } - for target in ["sm_120", "compute_120", "ptx87"] { - assert!(is_normalized_stage_target(target), "{target}"); - } - for target in ["sm_0120a", "sm_120af", "sm_120x", "compute_120A"] { - assert!(!is_normalized_stage_target(target), "{target}"); - } - } - - #[test] - fn libnvvm_stage_may_report_newer_ptx_than_the_native_instruction_floor() { - let mut target_policy = policy(); - target_policy.minimum_ptx = "1.0".into(); - target_policy.minimum_sm = None; - let lowering = crate::model::OverlayBackendLowering { - backend: IntrinsicBackend::LibNvvm, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: "test".into(), - targets: None, - minimum_ptx: None, - minimum_sm: Some("sm_75".into()), - }; - let mut record = evidence(); - record.stages = vec![evidence_stage( - EvidenceStageKind::BackendCodegen, - BackendLoweringMechanism::InlinePtx, - &["sm_75", "ptx93"], - )]; - validate_selected_stage_targets(&target_policy, &record, &lowering).unwrap(); - - record.stages[0].targets = vec!["sm_75".into(), "ptx09".into()]; - assert!( - validate_selected_stage_targets(&target_policy, &record, &lowering) - .unwrap_err() - .to_string() - .contains("catalog floor sm_75 / PTX 1.0") - ); - - let llvm_lowering = crate::model::OverlayBackendLowering { - backend: IntrinsicBackend::LlvmNvptx, - mechanism: BackendLoweringMechanism::TypedNvvm, - evidence_profile: "test".into(), - targets: None, - minimum_ptx: Some("3.2".into()), - minimum_sm: Some("sm_20".into()), - }; - record.stages = vec![evidence_stage( - EvidenceStageKind::BackendCodegen, - BackendLoweringMechanism::TypedNvvm, - &["sm_20", "ptx93"], - )]; - assert!( - validate_selected_stage_targets(&target_policy, &record, &llvm_lowering) - .unwrap_err() - .to_string() - .contains("catalog floor sm_20 / PTX 3.2") - ); - } - - #[test] - fn imported_selection_must_match_the_full_ptx_shape() { - let mut imported = declaration(); - imported.selections[0].asm = "mov.u32 $d, %tid.xy;".into(); - let error = validate_imported_policy(&policy(), &imported).unwrap_err(); - assert!(error.to_string().contains("does not agree")); - - imported.selections[0].asm = "mov.u32.relaxed $d, %tid.x;".into(); - let error = validate_imported_policy(&policy(), &imported).unwrap_err(); - assert!(error.to_string().contains("does not agree")); - - imported.selections[0].asm = "mov.u32 $d, %tid.x;".into(); - validate_imported_policy(&policy(), &imported).unwrap(); - } - - #[test] - fn selected_target_predicates_fail_closed() { - let selection = ImportedSelection { - source_record: "selection".into(), - asm: "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {{$r0, $r1, $r2, $r3}}, [$src];".into(), - predicates: vec![ - "Subtarget->getPTXVersion() >= 65".into(), - "Subtarget->getSmVersion() >= 75".into(), - ], - constraints: Default::default(), - }; - - let mut too_low_ptx = policy(); - too_low_ptx.minimum_ptx = "6.4".into(); - too_low_ptx.minimum_sm = Some("sm_75".into()); - assert!( - validate_selected_target_predicates(&too_low_ptx, &selection) - .unwrap_err() - .to_string() - .contains("minimum PTX") - ); - - let mut too_low_sm = policy(); - too_low_sm.minimum_ptx = "6.5".into(); - too_low_sm.minimum_sm = Some("sm_74".into()); - assert!( - validate_selected_target_predicates(&too_low_sm, &selection) - .unwrap_err() - .to_string() - .contains("minimum SM") - ); - - let mut unknown = selection; - unknown - .predicates - .push("Subtarget->hasMysteryFeature()".into()); - assert!( - validate_selected_target_predicates(&too_low_sm, &unknown) - .unwrap_err() - .to_string() - .contains("fail closed") - ); - } - - #[test] - fn blackwell_ldmatrix_predicate_requires_the_reviewed_exact_target_set() { - let mut record = policy(); - record.family = "ldmatrix".into(); - record.minimum_ptx = "8.6".into(); - record.minimum_sm = None; - record.targets = BLACKWELL_LDMATRIX_LLVM_TARGETS.into(); - record.ldmatrix_variant = Some(crate::model::LdmatrixVariant { - shape: LdmatrixShape::M16n16, - multiplicity: LdmatrixMultiplicity::X1, - layout: LdmatrixLayout::Transposed, - element: LdmatrixElement::B8, - state_space: LdmatrixStateSpace::Shared, - }); - let selection = ImportedSelection { - source_record: "selection".into(), - asm: "ldmatrix.sync.aligned.m16n16.x1.trans.shared.b8 {{$r0, $r1}}, [$src];".into(), - predicates: vec!["Subtarget->hasLdStmatrixBlackwellSupport()".into()], - constraints: Default::default(), - }; - - validate_selected_target_predicates(&record, &selection).unwrap(); - - record.targets = "sm_100a|sm_101a".into(); - assert!(validate_selected_target_predicates(&record, &selection).is_err()); - - record.targets = BLACKWELL_LDMATRIX_LLVM_TARGETS.into(); - let mut conflicting = selection; - conflicting - .predicates - .push("Subtarget->getPTXVersion() >= 86".into()); - assert!(validate_selected_target_predicates(&record, &conflicting).is_err()); - } - - #[test] - fn f32_redux_predicate_requires_the_reviewed_exact_target_matrix() { - let mut record = policy(); - record.family = "redux".into(); - record.minimum_ptx = "8.6".into(); - record.minimum_sm = None; - record.targets = REDUX_F32_TARGETS.into(); - record.redux = Some(crate::model::Redux { - operation: ReduxOperation::Fmin, - participation: - ReduxParticipation::ExecutingLaneNamedAllNamedLanesSameInstructionAndMask, - adapter: ReduxAdapter::MaskValueToSourceMemberMask, - }); - let selection = ImportedSelection { - source_record: "selection".into(), - asm: "redux.sync.min.f32 \t$dst, $src, $mask;".into(), - predicates: vec!["Subtarget->hasReduxSyncF32()".into()], - constraints: Default::default(), - }; - - validate_selected_target_predicates(&record, &selection).unwrap(); - - record.targets = - "sm_100a|sm_100f|sm_103a|sm_103f|sm_110a|sm_110f|sm_120a|sm_120f|sm_121a|sm_121f" - .into(); - assert!(validate_selected_target_predicates(&record, &selection).is_err()); - - record.targets = REDUX_F32_TARGETS.into(); - record.minimum_ptx = "8.8".into(); - assert!(validate_selected_target_predicates(&record, &selection).is_err()); - - record.minimum_ptx = "8.6".into(); - let mut conflicting = selection; - conflicting - .predicates - .push("Subtarget->getPTXVersion() >= 88".into()); - assert!(validate_selected_target_predicates(&record, &conflicting).is_err()); - } - - #[test] - fn safe_record_requires_an_allowlist_reason() { - let mut record = policy(); - record.safe_allowlist_reason = None; - assert!( - validate_imported_policy(&record, &declaration()) - .unwrap_err() - .to_string() - .contains("safe_allowlist_reason") - ); - } - - #[test] - fn intrinsic_abi_identity_is_stable_and_explicit() { - let policy = policy(); - let declaration = declaration(); - let resolved = resolve_record( - &policy, - resolve_policy_source(&policy).unwrap(), - Some(&declaration), - &evidence(), - "test", - "LLVM version test", - "0123456789abcdef", - vec![], - 1, - ) - .unwrap(); - - assert_eq!(resolved.rust.abi_id, "i0001"); - assert_eq!( - resolved.rust.canonical_path, - "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0001" - ); - assert_eq!( - resolved.rust.public_path, - "cuda_intrinsics::sreg::thread_idx_x" - ); - assert_eq!( - resolved.rust.compatibility_paths, - ["cuda_device::thread::threadIdx_x"] - ); - assert_eq!( - resolved.llvm.as_ref().unwrap().properties, - [ - "IntrNoMem", - "IntrSpeculatable", - "NoUndef", - "Range" - ] - ); - assert!(resolved.llvm.as_ref().unwrap().result_facts.no_undef); - assert_eq!( - resolved.llvm.as_ref().unwrap().result_facts.range, - Some(CatalogHalfOpenRange { - lower: "0".into(), - upper_exclusive: "1024".into(), - }) - ); - assert_eq!(resolved.backend.version, "LLVM version test"); - assert_eq!(resolved.backend.sha256, "0123456789abcdef"); - } - - #[test] - fn malformed_intrinsic_abi_ids_are_rejected() { - for abi_id in ["thread_idx_x", "i1", "x0001", "i00a1"] { - let mut record = policy(); - record.abi_id = abi_id.into(); - let error = validate_unique_overlay(&[record], 1).unwrap_err(); - assert!(error.to_string().contains("stable `iNNNN` form")); - } - } - - #[test] - fn ptx_versions_are_parsed_once_and_serialize_compatibly() { - for (text, encoded) in [("2.0", 20), ("6.5", 65), ("10.0", 100)] { - let version = parse_ptx_version(text, "test").unwrap(); - assert_eq!(version.encoded(), encoded); - assert_eq!( - serde_json::to_string(&version).unwrap(), - format!("\"{text}\"") - ); - assert_eq!( - serde_json::from_str::(&format!("\"{text}\"")).unwrap(), - version - ); - } - for malformed in ["6", "6.05", " 6.5", "06.5", "6.5 "] { - assert!(parse_ptx_version(malformed, "test").is_err(), "{malformed}"); - } - } - - #[test] - fn hardware_targets_are_parsed_without_losing_suffix_semantics() { - let all = policy(); - assert_eq!( - parse_hardware_target(&all).unwrap(), - CatalogHardwareTarget::All - ); - - let mut minimum = policy(); - minimum.minimum_sm = Some("sm_75".into()); - assert_eq!( - parse_hardware_target(&minimum).unwrap(), - CatalogHardwareTarget::AnyOf { - alternatives: vec![CatalogHardwareAlternative::MinimumSm { sm: 75 }], - } - ); - - let mut exact = policy(); - exact.targets = "sm_120a".into(); - assert_eq!( - parse_hardware_target(&exact).unwrap(), - CatalogHardwareTarget::AnyOf { - alternatives: vec![CatalogHardwareAlternative::ExactArchitecture { sm: 120 }], - } - ); - - let mut family = policy(); - family.targets = "sm_120f".into(); - assert_eq!( - parse_hardware_target(&family).unwrap(), - CatalogHardwareTarget::AnyOf { - alternatives: vec![CatalogHardwareAlternative::FamilyTarget { sm: 120 }], - } - ); - } - - #[test] - fn stage_hardware_shared_parser_preserves_canonical_evidence_language() { - for (target, expected) in [ - ("sm_75", CatalogHardwareAlternative::MinimumSm { sm: 75 }), - ( - "sm_120a", - CatalogHardwareAlternative::ExactArchitecture { sm: 120 }, - ), - ( - "compute_120f", - CatalogHardwareAlternative::FamilyTarget { sm: 120 }, - ), - ] { - assert_eq!(parse_stage_hardware(target), Some(expected), "{target}"); - } - for target in ["sm_090", "sm_1000"] { - assert_eq!(parse_stage_hardware(target), None, "{target}"); - } - } - - #[test] - fn reviewed_target_floors_equal_shared_backend_derivation() { - for (target, floor) in [ - ("sm_120a", 87), - ("sm_120f", 88), - ("sm_121a", 88), - ("sm_121f", 88), - ] { - let hardware = parse_stage_hardware(target).unwrap(); - assert_eq!(f8f6f4_llvm_ptx_floor(hardware).unwrap(), floor, "{target}"); - } - for (target, floor) in [ - ("sm_100a", 86), - ("sm_100f", 88), - ("sm_103a", 88), - ("sm_103f", 88), - ("sm_110a", 90), - ("sm_110f", 90), - ("sm_120a", 87), - ("sm_120f", 88), - ("sm_121a", 88), - ("sm_121f", 88), - ] { - let hardware = parse_stage_hardware(target).unwrap(); - assert_eq!( - blackwell_ldmatrix_llvm_ptx_floor(hardware).unwrap(), - floor, - "{target}" - ); - } - } - - fn selector(name: &str, value: &str) -> TargetSelectorBinding { - TargetSelectorBinding { - name: name.into(), - value: value.into(), - } - } - - fn target_contract( - selectors: Vec, - alternatives: &[(&str, &str)], - ) -> TargetContract { - TargetContract { - selectors, - alternatives: alternatives - .iter() - .map( - |(target, minimum_ptx)| crate::model::TargetContractAlternative { - target: (*target).into(), - minimum_ptx: (*minimum_ptx).into(), - }, - ) - .collect(), - } - } - - #[test] - fn target_contracts_keep_selector_specific_ptx_hardware_pairs() { - let contracts = [ - target_contract( - vec![selector("kind", "f16")], - &[ - ("sm_100a", "8.6"), - ("sm_101a", "8.6"), - ("sm_103a", "8.8"), - ("sm_110a", "9.0"), - ], - ), - target_contract( - vec![selector("kind", "i8")], - &[("sm_100a", "8.6"), ("sm_101a", "8.6"), ("sm_110a", "9.0")], - ), - ]; - let full_requirement = resolve_target_contracts("tcgen05_mma", &contracts).unwrap(); - assert_eq!(full_requirement.minimum_ptx.encoded(), 86); - let CatalogHardwareTarget::TargetMatrix { - contracts: full_contracts, - } = &full_requirement.hardware - else { - panic!("full target contracts must remain a matrix") - }; - assert_eq!(full_contracts.len(), 2); - assert_eq!(full_contracts[0].selectors, [selector("kind", "f16")]); - assert_eq!(full_contracts[1].selectors, [selector("kind", "i8")]); - - let requirement = - resolve_target_contract("tcgen05_mma", &[selector("kind", "i8")], &contracts).unwrap(); - - assert_eq!(requirement.minimum_ptx.encoded(), 86); - assert_eq!( - requirement.hardware, - CatalogHardwareTarget::TargetMatrix { - contracts: vec![CatalogTargetContract { - selectors: vec![selector("kind", "i8")], - alternatives: vec![ - CatalogTargetAlternative { - minimum_ptx: "8.6".parse().unwrap(), - hardware: CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, - }, - CatalogTargetAlternative { - minimum_ptx: "8.6".parse().unwrap(), - hardware: CatalogHardwareAlternative::ExactArchitecture { sm: 101 }, - }, - CatalogTargetAlternative { - minimum_ptx: "9.0".parse().unwrap(), - hardware: CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, - }, - ], - }], - } - ); - - let mut policy = policy(); - policy.id = "tcgen05_mma".into(); - validate_candidate_target(&policy, &requirement, "sm_100a", "+ptx86").unwrap(); - assert!(validate_candidate_target(&policy, &requirement, "sm_103a", "+ptx88").is_err()); - assert!(validate_candidate_target(&policy, &requirement, "sm_110a", "+ptx88").is_err()); - validate_candidate_target(&policy, &requirement, "sm_110a", "+ptx90").unwrap(); - - let libnvvm = [target_contract( - vec![selector("kind", "i8")], - &[("sm_100a", "8.6"), ("sm_110a", "9.0")], - )]; - let libnvvm_requirement = - resolve_target_contract("tcgen05_mma", &[selector("kind", "i8")], &libnvvm).unwrap(); - assert_ne!(requirement, libnvvm_requirement); - validate_candidate_target(&policy, &requirement, "sm_101a", "+ptx86").unwrap(); - assert!( - validate_candidate_target(&policy, &libnvvm_requirement, "sm_101a", "+ptx86").is_err() - ); - } - - #[test] - fn target_contract_selection_and_shape_fail_closed() { - resolve_target_contract( - "tcgen05_mma", - &[selector("kind", "f16")], - &[target_contract( - vec![selector("kind", "f16")], - &[("sm_100a", "8.6"), ("sm_100f", "8.8"), ("sm_103a", "8.8")], - )], - ) - .unwrap(); - - let valid = target_contract( - vec![selector("kind", "f16"), selector("scale_d", "false")], - &[("sm_100a", "8.6"), ("sm_103a", "8.8")], - ); - assert!( - resolve_target_contract( - "tcgen05_mma", - &[selector("kind", "i8"), selector("scale_d", "false")], - std::slice::from_ref(&valid), - ) - .unwrap_err() - .to_string() - .contains("exactly one reviewed contract") - ); - - for invalid in [ - target_contract( - vec![selector("scale_d", "false"), selector("kind", "f16")], - &[("sm_100a", "8.6")], - ), - target_contract(vec![selector("kind", "F16")], &[("sm_100a", "8.6")]), - target_contract(vec![selector("kind_", "f16")], &[("sm_100a", "8.6")]), - target_contract( - vec![selector("kind", "f16")], - &[("sm_103a", "8.8"), ("sm_100a", "8.6")], - ), - target_contract( - vec![selector("kind", "f16")], - &[("sm_100+", "8.6"), ("sm_103a", "8.8")], - ), - target_contract(vec![selector("kind", "f16")], &[("sm_100a", "8.60")]), - ] { - assert!( - resolve_target_contract( - "tcgen05_mma", - &invalid.selectors, - std::slice::from_ref(&invalid) - ) - .is_err(), - "{invalid:?}" - ); - } - } - - #[test] - fn legacy_target_requirement_json_is_unchanged() { - let requirement = CatalogTargetRequirement { - minimum_ptx: "8.6".parse().unwrap(), - hardware: CatalogHardwareTarget::AnyOf { - alternatives: vec![CatalogHardwareAlternative::ExactArchitecture { sm: 100 }], - }, - }; - assert_eq!( - serde_json::to_value(requirement).unwrap(), - serde_json::json!({ - "minimum_ptx": "8.6", - "hardware": { - "kind": "any_of", - "alternatives": [{ "kind": "exact_architecture", "sm": 100 }] - } - }) - ); - } - - #[test] - fn malformed_or_conflicting_hardware_targets_are_rejected() { - for malformed in [ - "sm_120", - "sm_120af", - "sm_120A", - "sm_0120a", - "sm_0a", - "sm_120+", - "compute_120a", - "all ", - ] { - let mut record = policy(); - record.targets = malformed.into(); - assert!(parse_hardware_target(&record).is_err(), "{malformed}"); - } - - let mut suffixed_minimum = policy(); - suffixed_minimum.minimum_sm = Some("sm_90a".into()); - assert!(parse_hardware_target(&suffixed_minimum).is_err()); - - for target in ["sm_120a", "sm_120f"] { - let mut conflicting = policy(); - conflicting.targets = target.into(); - conflicting.minimum_sm = Some("sm_120".into()); - let error = parse_hardware_target(&conflicting).unwrap_err().to_string(); - assert!(error.contains("cannot be combined"), "{error}"); - } - } - - #[test] - fn exact_inline_ptx_routes_can_inherit_exact_or_family_targets() { - for target in ["sm_120a", "sm_120f"] { - let mut record = policy(); - record.minimum_ptx = "8.7".into(); - record.targets = target.into(); - record.backend_lowerings = [IntrinsicBackend::LlvmNvptx, IntrinsicBackend::LibNvvm] - .into_iter() - .map(|backend| crate::model::OverlayBackendLowering { - backend, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: "test".into(), - targets: None, - minimum_ptx: Some("8.7".into()), - minimum_sm: None, - }) - .collect(); - - ensure_exact_inline_ptx_backends( - &record, - [ - (IntrinsicBackend::LlvmNvptx, "8.7", None), - (IntrinsicBackend::LibNvvm, "8.7", None), - ], - "test", - ) - .unwrap(); - for lowering in &record.backend_lowerings { - assert_eq!( - backend_target_requirement(&record, lowering) - .unwrap() - .hardware, - parse_hardware_target(&record).unwrap() - ); - } - - record.backend_lowerings[0].minimum_sm = Some("sm_120".into()); - assert!( - ensure_exact_inline_ptx_backends( - &record, - [ - (IntrinsicBackend::LlvmNvptx, "8.7", None), - (IntrinsicBackend::LibNvvm, "8.7", None), - ], - "test", - ) - .is_err() - ); - } - } - - #[test] - fn backend_route_target_override_is_exact_and_does_not_inherit_the_record_sm_floor() { - let mut record = policy(); - record.minimum_ptx = "8.6".into(); - record.minimum_sm = Some("sm_75".into()); - record.targets = "all".into(); - let lowering = crate::model::OverlayBackendLowering { - backend: IntrinsicBackend::LibNvvm, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: "test".into(), - targets: Some("sm_100a|sm_120a".into()), - minimum_ptx: None, - minimum_sm: None, - }; - - assert_eq!( - backend_target_requirement(&record, &lowering).unwrap(), - CatalogTargetRequirement { - minimum_ptx: parse_ptx_version("8.6", &record.id).unwrap(), - hardware: CatalogHardwareTarget::AnyOf { - alternatives: vec![ - CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 120 }, - ], - }, - } - ); - } - - #[test] - fn blackwell_ldmatrix_evidence_covers_every_target_with_its_effective_ptx_floor() { - let mut policy = policy(); - policy.family = "ldmatrix".into(); - policy.minimum_ptx = "8.6".into(); - policy.minimum_sm = None; - policy.targets = BLACKWELL_LDMATRIX_LLVM_TARGETS.into(); - policy.ldmatrix_variant = Some(crate::model::LdmatrixVariant { - shape: LdmatrixShape::M16n16, - multiplicity: LdmatrixMultiplicity::X1, - layout: LdmatrixLayout::Transposed, - element: LdmatrixElement::B8, - state_space: LdmatrixStateSpace::Shared, - }); - let lowering = OverlayBackendLowering { - backend: IntrinsicBackend::LlvmNvptx, - mechanism: BackendLoweringMechanism::TypedNvvm, - evidence_profile: "test".into(), - targets: None, - minimum_ptx: None, - minimum_sm: None, - }; - let mut record = evidence(); - record.status = "validated".into(); - record.stages.clear(); - for (target, ptx) in [ - ("sm_100a", "ptx86"), - ("sm_100f", "ptx88"), - ("sm_103a", "ptx88"), - ("sm_103f", "ptx88"), - ("sm_110a", "ptx90"), - ("sm_110f", "ptx90"), - ("sm_120a", "ptx87"), - ("sm_120f", "ptx88"), - ("sm_121a", "ptx88"), - ("sm_121f", "ptx88"), - ] { - record.stages.push(evidence_stage( - EvidenceStageKind::BackendCodegen, - BackendLoweringMechanism::TypedNvvm, - &[target, ptx], - )); - let mut assembly = evidence_stage( - EvidenceStageKind::PtxAssembly, - BackendLoweringMechanism::TypedNvvm, - &[target, ptx], - ); - assembly.tool_path = Some("/tool/ptxas".into()); - assembly.tool_version = Some("test".into()); - assembly.tool_sha256 = Some("0".repeat(64)); - record.stages.push(assembly); - } - validate_selected_stage_targets(&policy, &record, &lowering).unwrap(); - - record.stages.retain(|stage| { - !(stage.stage == EvidenceStageKind::PtxAssembly - && stage.targets.iter().any(|target| target == "sm_121f")) - }); - assert!( - validate_selected_stage_targets(&policy, &record, &lowering) - .unwrap_err() - .to_string() - .contains("one structured stage for each") - ); - } - - #[test] - fn paired_target_evidence_checks_backend_and_runtime_floors() { - let policy = policy(); - let hardware = vec![ - CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, - CatalogHardwareAlternative::ExactArchitecture { sm: 103 }, - ]; - let paired = vec![ - CatalogTargetAlternative { - minimum_ptx: "8.6".parse().unwrap(), - hardware: hardware[0], - }, - CatalogTargetAlternative { - minimum_ptx: "8.8".parse().unwrap(), - hardware: hardware[1], - }, - ]; - let llvm = OverlayBackendLowering { - backend: IntrinsicBackend::LlvmNvptx, - mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: "test".into(), - targets: None, - minimum_ptx: None, - minimum_sm: None, - }; - let mut llvm_record = evidence(); - llvm_record.status = "validated".into(); - llvm_record.stages.clear(); - for (target, ptx) in [("sm_100a", "ptx86"), ("sm_103a", "ptx88")] { - llvm_record.stages.push(evidence_stage( - EvidenceStageKind::BackendCodegen, - llvm.mechanism, - &[target, ptx], - )); - let mut assembly = evidence_stage( - EvidenceStageKind::PtxAssembly, - llvm.mechanism, - &[target, ptx], - ); - assembly.tool_path = Some("/tool/ptxas".into()); - assembly.tool_version = Some("test".into()); - assembly.tool_sha256 = Some("0".repeat(64)); - llvm_record.stages.push(assembly); - } - validate_target_matrix_stage_targets( - &policy, - &llvm_record, - &llvm, - EvidenceStageKind::PtxAssembly, - &hardware, - 86, - Some(&paired), - ) - .unwrap(); - - let mut wrong_floor = llvm_record.clone(); - wrong_floor - .stages - .iter_mut() - .find(|stage| { - stage.stage == EvidenceStageKind::BackendCodegen - && stage.targets.iter().any(|target| target == "sm_103a") - }) - .unwrap() - .targets = vec!["sm_103a".into(), "ptx86".into()]; - assert!( - validate_target_matrix_stage_targets( - &policy, - &wrong_floor, - &llvm, - EvidenceStageKind::PtxAssembly, - &hardware, - 86, - Some(&paired), - ) - .unwrap_err() - .to_string() - .contains("wrong PTX floor") - ); - - let libnvvm = OverlayBackendLowering { - backend: IntrinsicBackend::LibNvvm, - ..llvm.clone() - }; - let mut libnvvm_record = evidence(); - libnvvm_record.status = "validated".into(); - libnvvm_record.stages.clear(); - for (target, ptx) in [("sm_100a", "ptx88"), ("sm_103a", "ptx90")] { - libnvvm_record.stages.push(evidence_stage( - EvidenceStageKind::BackendCodegen, - libnvvm.mechanism, - &[target, ptx], - )); - let mut assembly = evidence_stage( - EvidenceStageKind::PtxAssembly, - libnvvm.mechanism, - &[target, ptx], - ); - assembly.tool_path = Some("/tool/ptxas".into()); - assembly.tool_version = Some("test".into()); - assembly.tool_sha256 = Some("0".repeat(64)); - libnvvm_record.stages.push(assembly); - let mut link = evidence_stage( - EvidenceStageKind::DeviceLink, - libnvvm.mechanism, - &[target, ptx], - ); - link.artifact_kind = Some(EvidenceArtifactKind::Cubin); - link.tool_path = Some("/tool/nvlink".into()); - link.tool_version = Some("test".into()); - link.tool_sha256 = Some("0".repeat(64)); - libnvvm_record.stages.push(link); - } - validate_target_matrix_stage_targets( - &policy, - &libnvvm_record, - &libnvvm, - EvidenceStageKind::DeviceLink, - &hardware, - 86, - Some(&paired), - ) - .unwrap(); - - let mut executed = llvm_record; - executed.status = "executed".into(); - executed.stages.push(evidence_stage( - EvidenceStageKind::Runtime, - llvm.mechanism, - &["sm_103a", "ptx88"], - )); - validate_target_matrix_stage_targets( - &policy, - &executed, - &llvm, - EvidenceStageKind::PtxAssembly, - &hardware, - 86, - Some(&paired), - ) - .unwrap(); - - let mut wrong_runtime_ptx = executed.clone(); - wrong_runtime_ptx.stages.last_mut().unwrap().targets = - vec!["sm_103a".into(), "ptx87".into()]; - assert!( - validate_target_matrix_stage_targets( - &policy, - &wrong_runtime_ptx, - &llvm, - EvidenceStageKind::PtxAssembly, - &hardware, - 86, - Some(&paired), - ) - .unwrap_err() - .to_string() - .contains("paired floor") - ); - - let mut wrong_runtime_hardware = executed; - wrong_runtime_hardware.stages.last_mut().unwrap().targets = - vec!["sm_110a".into(), "ptx90".into()]; - assert!( - validate_target_matrix_stage_targets( - &policy, - &wrong_runtime_hardware, - &llvm, - EvidenceStageKind::PtxAssembly, - &hardware, - 86, - Some(&paired), - ) - .unwrap_err() - .to_string() - .contains("outside its target matrix") - ); - } - - #[test] - fn sparse_f8f6f4_f16_evidence_covers_every_target_and_floor() { - let policy = - expand_sparse_mma_f8f6f4_f16_admission(&test_sparse_mma_f8f6f4_f16_admission()) - .unwrap() - .remove(0); - let lowering = policy - .backend_lowerings - .iter() - .find(|lowering| lowering.backend == IntrinsicBackend::LlvmNvptx) - .unwrap(); - let mut record = evidence(); - record.status = "validated".into(); - record.stages.clear(); - for (target, ptx) in [ - ("sm_120a", "ptx87"), - ("sm_120f", "ptx88"), - ("sm_121a", "ptx88"), - ("sm_121f", "ptx88"), - ] { - record.stages.push(evidence_stage( - EvidenceStageKind::BackendCodegen, - BackendLoweringMechanism::InlinePtx, - &[target, ptx], - )); - let mut assembly = evidence_stage( - EvidenceStageKind::PtxAssembly, - BackendLoweringMechanism::InlinePtx, - &[target, ptx], - ); - assembly.artifact_kind = Some(EvidenceArtifactKind::Cubin); - assembly.tool_path = Some("/tool/ptxas".into()); - assembly.tool_version = Some("test".into()); - assembly.tool_sha256 = Some("0".repeat(64)); - record.stages.push(assembly); - } - validate_selected_stage_targets(&policy, &record, lowering).unwrap(); - - let valid = record.clone(); - record.stages.retain(|stage| { - !(stage.stage == EvidenceStageKind::PtxAssembly - && stage.targets.iter().any(|target| target == "sm_121f")) - }); - assert!( - validate_selected_stage_targets(&policy, &record, lowering) - .unwrap_err() - .to_string() - .contains("one structured stage for each") - ); - - record = valid; - record - .stages - .iter_mut() - .find(|stage| { - stage.stage == EvidenceStageKind::BackendCodegen - && stage.targets.iter().any(|target| target == "sm_120f") - }) - .unwrap() - .targets = vec!["sm_120f".into(), "ptx87".into()]; - assert!( - validate_selected_stage_targets(&policy, &record, lowering) - .unwrap_err() - .to_string() - .contains("wrong PTX floor") - ); - } - - #[test] - fn non_blackwell_evidence_rejects_multiple_target_specific_stage_pairs() { - let policy = policy(); - let lowering = OverlayBackendLowering { - backend: IntrinsicBackend::LlvmNvptx, - mechanism: BackendLoweringMechanism::TypedNvvm, - evidence_profile: "test".into(), - targets: None, - minimum_ptx: None, - minimum_sm: None, - }; - let mut record = evidence(); - record.stages = vec![ - evidence_stage( - EvidenceStageKind::BackendCodegen, - BackendLoweringMechanism::TypedNvvm, - &["sm_20", "ptx20"], - ), - evidence_stage( - EvidenceStageKind::BackendCodegen, - BackendLoweringMechanism::TypedNvvm, - &["sm_21", "ptx20"], - ), - ]; - - assert!( - validate_selected_stage_targets(&policy, &record, &lowering) - .unwrap_err() - .to_string() - .contains("outside reviewed target-matrix evidence") - ); - } - - #[test] - fn abi_ledger_requires_exact_active_identity() { - let record = policy(); - let frozen_entry = ledger_entry(&record); - validate_abi_ledger( - &overlay_file(vec![record.clone()]), - &ledger(vec![frozen_entry.clone()]), - ) - .unwrap(); - - let mut reassigned = record.clone(); - reassigned.id = "different_catalog_id".into(); - let error = validate_abi_ledger( - &overlay_file(vec![reassigned]), - &ledger(vec![frozen_entry.clone()]), - ) - .unwrap_err(); - assert!(error.to_string().contains("catalog ID mismatch")); - - let mut reassigned = record.clone(); - reassigned.operation_key = "launch.block_index.x".into(); - let error = validate_abi_ledger( - &overlay_file(vec![reassigned]), - &ledger(vec![frozen_entry.clone()]), - ) - .unwrap_err(); - assert!(error.to_string().contains("operation key mismatch")); - - for mutate in [ - |record: &mut OverlayIntrinsic| record.safe = false, - |record: &mut OverlayIntrinsic| record.rust_arguments.push("u32".into()), - |record: &mut OverlayIntrinsic| record.rust_result = "u64".into(), - ] { - let mut changed_signature = record.clone(); - mutate(&mut changed_signature); - let error = validate_abi_ledger( - &overlay_file(vec![changed_signature]), - &ledger(vec![frozen_entry.clone()]), - ) - .unwrap_err(); - assert!(error.to_string().contains("raw Rust signature mismatch")); - } - } - - #[test] - fn generated_abi_binding_uses_catalog_identity_not_axis_position() { - let mut record = policy(); - let mut frozen = ledger_entry(&record); - frozen.abi_id = "i9001".into(); - record.abi_id.clear(); - let mut overlay = overlay_file(vec![record]); - - bind_generated_abi_ids(&mut overlay, &ledger(vec![frozen])).unwrap(); - - assert_eq!(overlay.intrinsics[0].abi_id, "i9001"); - } - - #[test] - fn generated_abi_binding_rejects_missing_tombstoned_or_ambiguous_identity() { - let record = policy(); - let mut unbound = record.clone(); - unbound.abi_id.clear(); - - let error = bind_generated_abi_ids( - &mut overlay_file(vec![unbound.clone()]), - &ledger(vec![ledger_entry(&distinct_policy())]), - ) - .unwrap_err(); - assert!(error.to_string().contains("has no ABI ledger entry")); - - let mut tombstone = ledger_entry(&record); - tombstone.status = "tombstone".into(); - let error = bind_generated_abi_ids( - &mut overlay_file(vec![unbound.clone()]), - &ledger(vec![tombstone]), - ) - .unwrap_err(); - assert!(error.to_string().contains("non-active ABI ledger entry")); - - let first = ledger_entry(&record); - let mut duplicate = first.clone(); - duplicate.abi_id = "i9002".into(); - let error = bind_generated_abi_ids( - &mut overlay_file(vec![unbound]), - &ledger(vec![first, duplicate]), - ) - .unwrap_err(); - assert!( - error - .to_string() - .contains("duplicate ABI ledger catalog ID") - ); - } - - #[test] - fn generated_abi_binding_checks_derived_operation_and_raw_signature() { - let record = policy(); - let mut unbound = record.clone(); - unbound.abi_id.clear(); - - let mut wrong_operation = ledger_entry(&record); - wrong_operation.operation_key = "launch.block_index.x".into(); - let error = bind_generated_abi_ids( - &mut overlay_file(vec![unbound.clone()]), - &ledger(vec![wrong_operation]), - ) - .unwrap_err(); - assert!(error.to_string().contains("operation key mismatch")); - - let mut wrong_signature = ledger_entry(&record); - wrong_signature - .raw_rust_signature - .arguments - .push("u32".into()); - let error = bind_generated_abi_ids( - &mut overlay_file(vec![unbound]), - &ledger(vec![wrong_signature]), - ) - .unwrap_err(); - assert!(error.to_string().contains("raw Rust signature mismatch")); - } - - #[test] - fn abi_ledger_does_not_freeze_public_or_backend_implementation_details() { - let record = policy(); - let frozen_entry = ledger_entry(&record); - let mut evolved = record.clone(); - evolved.rust_module = "coordinates".into(); - evolved.rust_name = "thread_x".into(); - evolved.public_rust_path = "cuda_intrinsics::coordinates::thread_x".into(); - evolved.llvm_symbol = Some("llvm.nvvm.backend.v2.tid.x".into()); - evolved.llvm_arguments = vec!["i8".into()]; - evolved.llvm_results = vec!["i64".into()]; - evolved.dialect_op_type = "ReadThreadIndexXOpV2".into(); - evolved.dialect_op_name = "nvvm.read_thread_index_x_v2".into(); - evolved.lowering = "backend_v2_adapter".into(); - - validate_abi_ledger(&overlay_file(vec![evolved]), &ledger(vec![frozen_entry])).unwrap(); - } - - #[test] - fn tombstoned_or_unlisted_abi_ids_cannot_reappear() { - let record = policy(); - let mut tombstone = ledger_entry(&record); - tombstone.status = "tombstone".into(); - let error = validate_abi_ledger( - &overlay_file(vec![record.clone()]), - &ledger(vec![tombstone]), - ) - .unwrap_err(); - assert!(error.to_string().contains("cannot reappear")); - - let error = validate_abi_ledger(&overlay_file(vec![record]), &ledger(vec![])).unwrap_err(); - assert!(error.to_string().contains("contains no entries")); - } - - #[test] - fn every_active_ledger_entry_requires_an_overlay_record() { - let record = policy(); - let error = - validate_abi_ledger(&overlay_file(vec![]), &ledger(vec![ledger_entry(&record)])) - .unwrap_err(); - assert!(error.to_string().contains("has no overlay record")); - } - - #[test] - fn return_range_properties_are_half_open_and_unique() { - let facts = - imported_result_facts(&["NoUndef".into(), "Range".into()]).unwrap(); - assert!(facts.no_undef); - let range = facts.range.unwrap(); - assert_eq!(range.lower, "1"); - assert_eq!(range.upper_exclusive, "1025"); - - let duplicate = - imported_result_facts(&["Range".into(), "Range".into()]) - .unwrap_err(); - assert!(duplicate.to_string().contains("duplicate return range")); - } - - #[test] - fn packed_alu_recipes_accept_only_the_reviewed_source_shape_and_floor() { - let operations = [ - PackedAluOperation::Add, - PackedAluOperation::Sub, - PackedAluOperation::Mul, - PackedAluOperation::Fma, - PackedAluOperation::FmaRelu, - PackedAluOperation::Min, - PackedAluOperation::Max, - PackedAluOperation::Neg, - PackedAluOperation::Abs, - ]; - for format in [PackedAluFormat::Bf16x2, PackedAluFormat::F16x2] { - for operation in operations { - let policy = packed_alu_policy(format, operation); - match packed_alu_declaration(format, operation) { - Some(declaration) => validate_imported_policy(&policy, &declaration).unwrap(), - None => validate_ptx_native_policy(&policy).unwrap(), - } - } - } - - let declaration = - packed_alu_declaration(PackedAluFormat::Bf16x2, PackedAluOperation::Fma).unwrap(); - let reject_imported = |policy: &OverlayIntrinsic, message: &str| { - let error = validate_imported_policy(policy, &declaration).unwrap_err(); - assert!(error.to_string().contains(message), "{error:#}"); - }; - - let valid = packed_alu_policy(PackedAluFormat::Bf16x2, PackedAluOperation::Fma); - let mut wrong_source = valid.clone(); - wrong_source.source_record = Some("int_nvvm_fma_rn_bf16".into()); - reject_imported(&wrong_source, "source"); - - let mut wrong_format = valid.clone(); - wrong_format.packed_alu.as_mut().unwrap().format = PackedAluFormat::F16x2; - reject_imported(&wrong_format, "identity"); - - let mut wrong_operation = valid.clone(); - wrong_operation.packed_alu.as_mut().unwrap().operation = PackedAluOperation::Max; - reject_imported(&wrong_operation, "identity"); - - let mut wrong_floor = valid.clone(); - wrong_floor.minimum_sm = Some("sm_90".into()); - reject_imported(&wrong_floor, "target floor"); - - let mut wrong_effects = valid.clone(); - wrong_effects.memory = "read".into(); - reject_imported(&wrong_effects, "effects"); - - let mut wrong_section = valid.clone(); - wrong_section.ptx_isa_section = "9.7.4 Floating Point Instructions".into(); - reject_imported(&wrong_section, "PTX provenance"); - - let mut wrong_url = valid.clone(); - wrong_url.ptx_isa_url = - "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions" - .into(); - reject_imported(&wrong_url, "PTX provenance"); - - let mut wrong_adapter = valid.clone(); - wrong_adapter.lowering = "direct_nvvm".into(); - reject_imported(&wrong_adapter, "lowering recipe"); - - let mut wrong_backend = valid; - wrong_backend.backend_lowerings[0].mechanism = BackendLoweringMechanism::TypedNvvm; - reject_imported(&wrong_backend, "inline-PTX routes"); - - let mut wrong_native = packed_alu_policy(PackedAluFormat::Bf16x2, PackedAluOperation::Add); - wrong_native.source = Some(IntrinsicSource::PtxNative { - instruction: "add.bf16x2".into(), - }); - let error = validate_ptx_native_policy(&wrong_native).unwrap_err(); - assert!(error.to_string().contains("PTX-native recipe")); - - let mut invented_llvm = packed_alu_policy(PackedAluFormat::F16x2, PackedAluOperation::Add); - invented_llvm.llvm_symbol = Some("llvm.nvvm.add.rn.f16x2".into()); - let error = validate_ptx_native_policy(&invented_llvm).unwrap_err(); - assert!( - error - .to_string() - .contains("must not invent LLVM source facts") - ); - - let mut unreviewed_modifier = - packed_alu_policy(PackedAluFormat::F16x2, PackedAluOperation::Add); - unreviewed_modifier.expected_ptx.modifiers = - vec!["rn".into(), "ftz".into(), "f16x2".into()]; - let error = validate_ptx_native_policy(&unreviewed_modifier).unwrap_err(); - assert!(error.to_string().contains("exact packed-ALU instruction")); - - let mut wrong_arity = packed_alu_policy(PackedAluFormat::F16x2, PackedAluOperation::Add); - wrong_arity.expected_ptx.operands.pop(); - let error = validate_ptx_native_policy(&wrong_arity).unwrap_err(); - assert!(error.to_string().contains("exact packed-ALU instruction")); - - let f16_declaration = - packed_alu_declaration(PackedAluFormat::F16x2, PackedAluOperation::Fma).unwrap(); - let reject_f16 = |policy: &OverlayIntrinsic, message: &str| { - let error = validate_imported_policy(policy, &f16_declaration).unwrap_err(); - assert!(error.to_string().contains(message), "{error:#}"); - }; - let f16 = packed_alu_policy(PackedAluFormat::F16x2, PackedAluOperation::Fma); - - let mut wrong_signature = f16.clone(); - wrong_signature.llvm_arguments = vec!["v2bf16".into(); 3]; - reject_f16(&wrong_signature, "LLVM argument signature mismatch"); - - let mut missing_must_use = f16.clone(); - missing_must_use.must_use = false; - reject_f16(&missing_must_use, "reviewed safe packed-ALU API"); - - let mut wrong_native_floor = f16.clone(); - wrong_native_floor - .packed_alu - .as_mut() - .unwrap() - .native_minimum_sm = 70; - reject_f16(&wrong_native_floor, "target floor"); - - let mut wrong_backend_floor = f16; - wrong_backend_floor.backend_lowerings[0].minimum_ptx = Some("4.2".into()); - reject_f16(&wrong_backend_floor, "exact packed-ALU floor"); - - let abs_declaration = - packed_alu_declaration(PackedAluFormat::F16x2, PackedAluOperation::Abs).unwrap(); - let mut wrong_abs = packed_alu_policy(PackedAluFormat::F16x2, PackedAluOperation::Abs); - wrong_abs.resolved_llvm_symbol = Some("llvm.nvvm.fabs.v2bf16".into()); - let error = validate_imported_policy(&wrong_abs, &abs_declaration).unwrap_err(); - assert!(error.to_string().contains("LLVM source or signature")); - } - - #[test] - fn pinned_packed_alu_records_match_the_closed_recipes() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let (overlay, _) = - read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")).unwrap(); - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - let declarations: BTreeMap<_, _> = imported - .intrinsics - .iter() - .map(|record| (record.source_record.as_str(), record)) - .collect(); - let packed: Vec<_> = overlay - .intrinsics - .iter() - .filter(|record| record.family == "packed_alu") - .collect(); - assert_eq!(packed.len(), 30); - for policy in packed { - let source = resolve_policy_source(policy).unwrap(); - let declaration = policy - .source_record - .as_deref() - .map(|record| declarations[record]); - validate_policy(policy, &source, declaration, 1).unwrap(); - } - } - - #[test] - fn pinned_packed_conversion_records_match_the_closed_recipes() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let overlay = load_resolution_base(&repo_root).unwrap().overlay; - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - let declarations: BTreeMap<_, _> = imported - .intrinsics - .iter() - .map(|record| (record.source_record.as_str(), record)) - .collect(); - let packed: Vec<_> = overlay - .intrinsics - .iter() - .filter(|record| record.family == "packed_conversion") - .collect(); - assert_eq!(packed.len(), 18); - for policy in packed { - let source = resolve_policy_source(policy).unwrap(); - let declaration = policy - .source_record - .as_deref() - .map(|record| declarations[record]); - validate_policy(policy, &source, declaration, 1).unwrap(); - } - } - - #[test] - fn selectionless_packed_conversion_is_admitted_only_by_its_closed_recipe() { - let cases = [ - ( - PackedConversionDestinationFormat::Bf16x2, - PackedConversionRounding::NearestEven, - PackedConversionSaturation::None, - ), - ( - PackedConversionDestinationFormat::F16x2, - PackedConversionRounding::NearestEven, - PackedConversionSaturation::None, - ), - ( - PackedConversionDestinationFormat::F16x2, - PackedConversionRounding::TowardZero, - PackedConversionSaturation::None, - ), - ( - PackedConversionDestinationFormat::F16x2, - PackedConversionRounding::NearestEven, - PackedConversionSaturation::Relu, - ), - ( - PackedConversionDestinationFormat::Bf16x2, - PackedConversionRounding::NearestEven, - PackedConversionSaturation::Relu, - ), - ( - PackedConversionDestinationFormat::Bf16x2, - PackedConversionRounding::TowardZero, - PackedConversionSaturation::None, - ), - ( - PackedConversionDestinationFormat::E4m3x2, - PackedConversionRounding::NearestEven, - PackedConversionSaturation::Satfinite, - ), - ( - PackedConversionDestinationFormat::E4m3x2, - PackedConversionRounding::NearestEven, - PackedConversionSaturation::SatfiniteRelu, - ), - ( - PackedConversionDestinationFormat::E5m2x2, - PackedConversionRounding::NearestEven, - PackedConversionSaturation::Satfinite, - ), - ( - PackedConversionDestinationFormat::E5m2x2, - PackedConversionRounding::NearestEven, - PackedConversionSaturation::SatfiniteRelu, - ), - ]; - for (destination, rounding, saturation) in cases { - let policy = packed_conversion_policy(destination, rounding, saturation); - let declaration = packed_conversion_declaration(&policy); - validate_imported_policy(&policy, &declaration).unwrap(); - } - - let valid = packed_conversion_policy( - PackedConversionDestinationFormat::Bf16x2, - PackedConversionRounding::NearestEven, - PackedConversionSaturation::None, - ); - let declaration = packed_conversion_declaration(&valid); - - let reject = |policy: &OverlayIntrinsic, declaration: &ImportedIntrinsic, message: &str| { - let error = validate_imported_policy(policy, declaration).unwrap_err(); - assert!(error.to_string().contains(message), "{error:#}"); - }; - - let mut wrong_source = valid.clone(); - wrong_source.source_record = Some("int_nvvm_ff2bf16x2_rz".into()); - reject(&wrong_source, &declaration, "identity or LLVM source"); - - let mut wrong_floor = valid.clone(); - wrong_floor.minimum_ptx = "7.8".into(); - reject(&wrong_floor, &declaration, "target floor"); - - let mut wrong_effect = valid.clone(); - wrong_effect.pure = false; - reject(&wrong_effect, &declaration, "effects"); - - let mut wrong_section = valid.clone(); - wrong_section.ptx_isa_section = "9.7.9 Data Movement and Conversion Instructions".into(); - reject(&wrong_section, &declaration, "PTX provenance"); - - let mut wrong_url = valid.clone(); - wrong_url.ptx_isa_url = - "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions" - .into(); - reject(&wrong_url, &declaration, "PTX provenance"); - - let mut wrong_shape = valid.clone(); - wrong_shape.expected_ptx.modifiers.swap(1, 2); - reject(&wrong_shape, &declaration, "expected PTX does not match"); - - let mut wrong_identity = valid.clone(); - wrong_identity.id = "cvt_f16x2_f32".into(); - reject(&wrong_identity, &declaration, "identity or LLVM source"); - - let mut unsupported = valid.clone(); - let conversion = unsupported.packed_conversion.as_mut().unwrap(); - conversion.rounding = PackedConversionRounding::TowardZero; - conversion.saturation = PackedConversionSaturation::Relu; - reject( - &unsupported, - &declaration, - "unsupported packed-conversion source, destination", - ); - - let mut wrong_compatibility = valid.clone(); - wrong_compatibility.compatibility_rust_paths = - vec!["cuda_device::tcgen05::cvt_f32x2_bf16x2".into()]; - reject(&wrong_compatibility, &declaration, "conversion API"); - - let mut wrong_result = valid.clone(); - wrong_result.llvm_results = vec!["v2f16".into()]; - reject(&wrong_result, &declaration, "result signature mismatch"); - - let mut selected = declaration.clone(); - selected.selections.push(ImportedSelection { - source_record: "UNREVIEWED".into(), - asm: "cvt.rn.bf16x2.f32 $d, $a, $b;".into(), - predicates: vec![], - constraints: Default::default(), - }); - reject(&valid, &selected, "selectionless"); - } - - #[test] - fn candidate_targets_are_canonical_and_satisfy_every_floor() { - let policy = policy(); - let requirement = CatalogTargetRequirement { - minimum_ptx: "7.0".parse().unwrap(), - hardware: CatalogHardwareTarget::AnyOf { - alternatives: vec![CatalogHardwareAlternative::MinimumSm { sm: 80 }], - }, - }; - validate_candidate_target(&policy, &requirement, "sm_80", "+ptx70").unwrap(); - validate_candidate_target(&policy, &requirement, "sm_90a", "+ptx86").unwrap(); - assert!( - validate_candidate_target(&policy, &requirement, "sm_75", "+ptx70") - .unwrap_err() - .to_string() - .contains("hardware requirement") - ); - assert!( - validate_candidate_target(&policy, &requirement, "sm_80", "+ptx69") - .unwrap_err() - .to_string() - .contains("PTX floor") - ); - for malformed in ["compute_80", "sm_080", "sm_80x"] { - assert!( - validate_candidate_target(&policy, &requirement, malformed, "+ptx70").is_err(), - "{malformed}" - ); - } - for malformed in ["ptx70", "+ptx7", "+ptx070"] { - assert!( - validate_candidate_target(&policy, &requirement, "sm_80", malformed).is_err(), - "{malformed}" - ); - } - - let exact = CatalogTargetRequirement { - minimum_ptx: "8.7".parse().unwrap(), - hardware: CatalogHardwareTarget::AnyOf { - alternatives: vec![CatalogHardwareAlternative::ExactArchitecture { sm: 120 }], - }, - }; - validate_candidate_target(&policy, &exact, "sm_120a", "+ptx87").unwrap(); - assert!(validate_candidate_target(&policy, &exact, "sm_120a", "+ptx86").is_err()); - assert!(validate_candidate_target(&policy, &exact, "sm_120", "+ptx87").is_err()); - assert!(validate_candidate_target(&policy, &exact, "sm_120f", "+ptx87").is_err()); - - let family = CatalogTargetRequirement { - minimum_ptx: "8.7".parse().unwrap(), - hardware: CatalogHardwareTarget::AnyOf { - alternatives: vec![CatalogHardwareAlternative::FamilyTarget { sm: 120 }], - }, - }; - validate_candidate_target(&policy, &family, "sm_120f", "+ptx87").unwrap(); - assert!(validate_candidate_target(&policy, &family, "sm_120a", "+ptx87").is_err()); - } - - #[test] - fn scalar_conversion_admission_is_closed_and_backend_specific() { - let variants = [ - ( - "i0368", - ScalarConversionRounding::NearestAway, - ScalarConversionSaturation::None, - ), - ( - "i0369", - ScalarConversionRounding::NearestAway, - ScalarConversionSaturation::Satfinite, - ), - ( - "i0370", - ScalarConversionRounding::NearestEven, - ScalarConversionSaturation::None, - ), - ( - "i0371", - ScalarConversionRounding::NearestEven, - ScalarConversionSaturation::Relu, - ), - ( - "i0372", - ScalarConversionRounding::NearestEven, - ScalarConversionSaturation::Satfinite, - ), - ( - "i0373", - ScalarConversionRounding::NearestEven, - ScalarConversionSaturation::ReluSatfinite, - ), - ( - "i0374", - ScalarConversionRounding::TowardZero, - ScalarConversionSaturation::None, - ), - ( - "i0375", - ScalarConversionRounding::TowardZero, - ScalarConversionSaturation::Relu, - ), - ( - "i0376", - ScalarConversionRounding::TowardZero, - ScalarConversionSaturation::Satfinite, - ), - ( - "i0377", - ScalarConversionRounding::TowardZero, - ScalarConversionSaturation::ReluSatfinite, - ), - ]; - let admission = ScalarConversionAdmission { - llvm_evidence_profile: "llvm-scalar".into(), - libnvvm_evidence_profile: "libnvvm-scalar".into(), - runtime_validation: RuntimeValidation::Unexecuted, - variants: variants - .iter() - .map( - |(abi_id, rounding, saturation)| ScalarConversionAdmissionVariant { - abi_id: (*abi_id).into(), - rounding: *rounding, - saturation: *saturation, - }, - ) - .collect(), - }; - - let records = expand_scalar_conversion_admission(&admission).unwrap(); - assert_eq!(records.len(), 10); - for (record, (abi_id, rounding, saturation)) in records.iter().zip(variants) { - assert_eq!(record.abi_id, abi_id); - assert_eq!(record.rust_arguments, ["f32"]); - assert_eq!(record.rust_result, "u32"); - assert_eq!( - record.compatibility_rust_paths, - [format!("cuda_device::convert::{}", record.rust_name)] - ); - let conversion = record.scalar_conversion.as_ref().unwrap(); - assert_eq!(conversion.rounding, rounding); - assert_eq!(conversion.saturation, saturation); - assert!(record.backend_lowerings.iter().any(|lowering| { - lowering.backend == IntrinsicBackend::LlvmNvptx - && lowering.mechanism == BackendLoweringMechanism::TypedNvvm - })); - assert!(record.backend_lowerings.iter().any(|lowering| { - lowering.backend == IntrinsicBackend::LibNvvm - && lowering.mechanism == BackendLoweringMechanism::InlinePtx - })); - } - - let mut reordered = admission.clone(); - reordered.variants.swap(0, 1); - assert!( - expand_scalar_conversion_admission(&reordered) - .unwrap_err() - .to_string() - .contains("canonical ten variants") - ); - } - - #[test] - fn scalar_arithmetic_admission_is_closed_and_selects_only_direct_ptx() { - let variants = canonical_scalar_arithmetic_variants(); - let admission = ScalarArithmeticAdmission { - llvm_evidence_profile: "llvm-scalar-arithmetic".into(), - libnvvm_evidence_profile: "libnvvm-scalar-arithmetic".into(), - runtime_validation: RuntimeValidation::Unexecuted, - variants: variants - .iter() - .copied() - .enumerate() - .map( - |(index, variant)| crate::model::ScalarArithmeticAdmissionVariant { - abi_id: format!("i{:04}", 390 + index), - format: variant.0, - operation: variant.1, - rounding: variant.2, - subnormal: variant.3, - saturation: variant.4, - }, - ) - .collect(), - }; - let records = expand_scalar_arithmetic_admission(&admission).unwrap(); - assert_eq!(records.len(), 64); - assert_eq!(records.first().unwrap().id, "mul_rn_f64"); - assert_eq!(records.first().unwrap().abi_id, "i0390"); - assert_eq!(records.last().unwrap().id, "add_rp_ftz_sat_f32"); - assert_eq!(records.last().unwrap().abi_id, "i0453"); - validate_unique_overlay(&records, 1).unwrap(); - let llvm_inline_count = records - .iter() - .filter(|record| { - record.backend_lowerings.iter().any(|lowering| { - lowering.backend == IntrinsicBackend::LlvmNvptx - && lowering.mechanism == BackendLoweringMechanism::InlinePtx - }) - }) - .count(); - assert_eq!(llvm_inline_count, 16); - assert!(records.iter().all(|record| { - let expected = - if record.scalar_arithmetic.as_ref().is_some_and(|arithmetic| { - arithmetic.saturation == ScalarArithmeticSaturation::Sat - }) { - BackendLoweringMechanism::InlinePtx - } else { - BackendLoweringMechanism::TypedNvvm - }; - record.backend_lowerings.iter().any(|lowering| { - lowering.backend == IntrinsicBackend::LlvmNvptx && lowering.mechanism == expected - }) - })); - - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - let declarations = imported - .intrinsics - .into_iter() - .map(|record| (record.source_record.clone(), record)) - .collect::>(); - let mut mixed_count = 0; - let mut add_two_selection_count = 0; - let mut add_six_selection_count = 0; - for record in &records { - let declaration = &declarations[record.source_record.as_ref().unwrap()]; - validate_imported_policy(record, declaration).unwrap(); - let selected = declaration - .selections - .iter() - .filter(|selection| selection_matches_policy(record, selection).unwrap()) - .collect::>(); - assert_eq!(selected.len(), 1, "{}", record.id); - if declaration.selections.len() == 3 { - mixed_count += 1; - } - if record - .scalar_arithmetic - .as_ref() - .is_some_and(|contract| contract.operation == ScalarArithmeticOperation::Add) - { - match declaration.selections.len() { - 2 => add_two_selection_count += 1, - 6 => add_six_selection_count += 1, - count => panic!("unexpected add selection count {count}"), - } - } - } - assert_eq!(mixed_count, 8); - assert_eq!(add_two_selection_count, 12); - assert_eq!(add_six_selection_count, 8); - let add_ftz_sat = records - .iter() - .find(|record| record.id == "add_rn_ftz_sat_f32") - .unwrap(); - assert_eq!( - add_ftz_sat.expected_ptx.modifiers, - ["rn", "sat", "ftz", "f32"] - ); - - let mut non_contiguous_abi = admission.clone(); - non_contiguous_abi.variants[0].abi_id = "i9999".into(); - let non_contiguous_records = - expand_scalar_arithmetic_admission(&non_contiguous_abi).unwrap(); - assert_eq!(non_contiguous_records[0].abi_id, "i9999"); - - let mut reordered = admission.clone(); - reordered.variants.swap(0, 1); - assert!( - expand_scalar_arithmetic_admission(&reordered) - .unwrap_err() - .to_string() - .contains("canonical 64 variants") - ); - assert!( - scalar_arithmetic_recipe(( - ScalarArithmeticFormat::F64, - ScalarArithmeticOperation::Fma, - ScalarArithmeticRounding::Rn, - ScalarArithmeticSubnormal::Ftz, - ScalarArithmeticSaturation::None, - )) - .is_none() - ); - } - - #[test] - fn scalar_math_admission_keeps_semantic_order_but_accepts_non_contiguous_abi_ids() { - let variants = canonical_scalar_math_variants(); - let admission = ScalarMathAdmission { - llvm_evidence_profile: "llvm-scalar-math".into(), - libnvvm_evidence_profile: "libnvvm-scalar-math".into(), - runtime_validation: RuntimeValidation::Unexecuted, - variants: variants - .iter() - .copied() - .enumerate() - .map( - |(index, variant)| crate::model::ScalarMathAdmissionVariant { - abi_id: format!("i{:04}", 782 + index), - libnvvm_evidence_profile: (index == 40).then(|| "libnvvm-ex2-f16".into()), - format: variant.0, - operation: variant.1, - precision: variant.2, - subnormal: variant.3, - }, - ) - .collect(), - }; - - let records = expand_scalar_math_admission(&admission).unwrap(); - assert_eq!(records.len(), 41); - assert_eq!(records.first().unwrap().abi_id, "i0782"); - assert_eq!(records.last().unwrap().abi_id, "i0822"); - assert!( - records - .iter() - .all(|record| record.backend_lowerings.len() == 2) - ); - assert_eq!( - records.last().unwrap().backend_lowerings[1].evidence_profile, - "libnvvm-ex2-f16" - ); - assert!(records[..40].iter().all(|record| { - record.backend_lowerings[1].evidence_profile == "libnvvm-scalar-math" - })); - - let mut non_contiguous_abi = admission.clone(); - non_contiguous_abi.variants[40].abi_id = "i9999".into(); - let records = expand_scalar_math_admission(&non_contiguous_abi).unwrap(); - assert_eq!(records.last().unwrap().abi_id, "i9999"); - - let mut reordered = admission; - reordered.variants.swap(0, 1); - assert!( - expand_scalar_math_admission(&reordered) - .unwrap_err() - .to_string() - .contains("canonical 41 variants") - ); - } - - #[test] - fn checked_in_scalar_math_overlay_pins_the_current_libnvvm_override_set() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let (overlay, _) = - read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")).unwrap(); - let scalar_math = overlay - .intrinsics - .iter() - .filter(|record| record.family == "scalar_math") - .collect::>(); - assert_eq!(scalar_math.len(), 41); - assert!( - scalar_math - .iter() - .all(|record| record.backend_lowerings.len() == 2) - ); - - let family_profile = "cuda-13.3-libnvvm-13.3.33-scalar-math"; - let overrides = scalar_math - .iter() - .filter_map(|record| { - let libnvvm = record - .backend_lowerings - .iter() - .find(|route| route.backend == IntrinsicBackend::LibNvvm) - .unwrap(); - (libnvvm.evidence_profile != family_profile).then_some(record.id.as_str()) - }) - .collect::>(); - assert_eq!(overrides, ["ex2_approx_f16"]); - } - - #[test] - fn extended_minmax_admission_is_exact_and_fail_closed() { - let variants = canonical_extended_minmax_variants(); - let admission = ExtendedMinMaxAdmission { - llvm_evidence_profile: "llvm-extended-minmax".into(), - libnvvm_evidence_profile: "libnvvm-extended-minmax".into(), - runtime_validation: RuntimeValidation::Unexecuted, - variants: variants - .iter() - .copied() - .enumerate() - .map( - |(index, variant)| crate::model::ExtendedMinMaxAdmissionVariant { - abi_id: if index < 28 { - format!("i{:04}", 550 + index) - } else { - format!("i{:04}", 830 + index - 28) - }, - format: variant.0, - operation: variant.1, - subnormal: variant.2, - nan: variant.3, - xorsign_abs: variant.4, - }, - ) - .collect(), - }; - let records = expand_extended_minmax_admission(&admission).unwrap(); - assert_eq!(records.len(), 52); - assert_eq!(records.first().unwrap().id, "min_ftz_f16x2"); - assert_eq!(records.first().unwrap().abi_id, "i0550"); - // The first reserved block ends here; the scalar 16-bit forms admitted - // afterwards continue at a second base rather than at i0578. - assert_eq!(records[27].id, "max_xorsign_abs_f16x2"); - assert_eq!(records[27].abi_id, "i0577"); - assert_eq!(records[28].id, "min_f16"); - assert_eq!(records[28].abi_id, "i0830"); - assert_eq!(records.last().unwrap().id, "max_nan_xorsign_abs_bf16"); - assert_eq!(records.last().unwrap().abi_id, "i0853"); - validate_unique_overlay(&records, 1).unwrap(); - assert_eq!( - records - .iter() - .filter(|record| record.minimum_sm.as_deref() == Some("sm_80")) - .count(), - 20 - ); - assert_eq!( - records - .iter() - .filter(|record| record.minimum_sm.as_deref() == Some("sm_86")) - .count(), - 32 - ); - assert!(records.iter().all(|record| { - record.backend_lowerings.len() == 2 - && record - .backend_lowerings - .iter() - .all(|lowering| lowering.mechanism == BackendLoweringMechanism::InlinePtx) - })); - - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let imported: ImportedFile = - read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); - let declarations = imported - .intrinsics - .into_iter() - .map(|record| (record.source_record.clone(), record)) - .collect::>(); - for record in &records { - let declaration = &declarations[record.source_record.as_ref().unwrap()]; - assert_eq!(declaration.selections.len(), 1, "{}", record.id); - validate_imported_policy(record, declaration).unwrap(); - assert_eq!( - declaration - .selections - .iter() - .filter(|selection| selection_matches_policy(record, selection).unwrap()) - .count(), - 1, - "{}", - record.id - ); - } - let packed_max = records - .iter() - .find(|record| record.id == "max_nan_bf16x2") - .unwrap(); - assert_eq!( - declarations[packed_max.source_record.as_ref().unwrap()].selections[0].source_record, - "INT_NVVM_FMAN_NaN_bf16x2" - ); - assert_eq!( - declarations["int_nvvm_fmax_ftz_nan_xorsign_abs_f"].selections[0].source_record, - "INT_NVVM_FMAX_FTZ_NAN_XORSIGN_ABS_F" - ); - assert!( - declarations["int_nvvm_fmax_ftz_nan_xorsign_abs_f"].selections[0] - .asm - .contains(".NaN.") - ); - for excluded in [ - "int_nvvm_fmin_ftz_bf16x2", - "int_nvvm_fmax_ftz_nan_bf16x2", - "int_nvvm_fmin_f", - "int_nvvm_fmax_nan_f", - ] { - assert!(declarations[excluded].selections.is_empty(), "{excluded}"); - assert!( - records - .iter() - .all(|record| record.source_record.as_deref() != Some(excluded)), - "{excluded} was admitted without an instruction selection" - ); - } - - let catalog: CatalogFile = read_json(&repo_root.join("intrinsics/catalog.json")).unwrap(); - for id in ["min_f16x2", "max_f16x2", "min_bf16x2", "max_bf16x2"] { - assert_eq!( - catalog - .intrinsics - .iter() - .find(|record| record.id == id) - .unwrap() - .family, - "packed_alu" - ); - assert!(records.iter().all(|record| record.id != id)); - } - - let mut non_contiguous_abi = admission.clone(); - non_contiguous_abi.variants[0].abi_id = "i9999".into(); - let non_contiguous_records = expand_extended_minmax_admission(&non_contiguous_abi).unwrap(); - assert_eq!(non_contiguous_records[0].abi_id, "i9999"); - - let mut reordered = admission.clone(); - reordered.variants.swap(0, 1); - assert!( - expand_extended_minmax_admission(&reordered) - .unwrap_err() - .to_string() - .contains("exact canonical 28 variants") - ); - assert!( - extended_minmax_recipe(( - ExtendedMinMaxFormat::Bf16x2, - ExtendedMinMaxOperation::Min, - ExtendedMinMaxSubnormal::Ftz, - ExtendedMinMaxNan::Nan, - false, - )) - .is_none() - ); - let mut changed = declarations["int_nvvm_fmin_ftz_f16x2"].clone(); - changed.selections[0].predicates.swap(0, 1); - assert!(validate_imported_policy(&records[0], &changed).is_err()); - } - - struct CandidateTestRepo(PathBuf); - - impl Drop for CandidateTestRepo { - fn drop(&mut self) { - let _ = fs::remove_dir_all(&self.0); - } - } - - fn repo_without_evidence() -> CandidateTestRepo { - use std::sync::atomic::{AtomicU64, Ordering}; - - static NEXT: AtomicU64 = AtomicU64::new(0); - let source = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let root = std::env::temp_dir().join(format!( - "cuda-intrinsics-candidate-{}-{}", - std::process::id(), - NEXT.fetch_add(1, Ordering::Relaxed) - )); - let input = root.join("intrinsics"); - fs::create_dir_all(input.join("overlay")).unwrap(); - for name in [ - "upstream.lock", - "imported.json", - "overlay.toml", - "abi-v1.toml", - ] { - fs::copy(source.join("intrinsics").join(name), input.join(name)).unwrap(); - } - for entry in fs::read_dir(source.join("intrinsics/overlay")).unwrap() { - let entry = entry.unwrap(); - if entry.path().extension().and_then(|value| value.to_str()) == Some("toml") { - fs::copy(entry.path(), input.join("overlay").join(entry.file_name())).unwrap(); - } - } - CandidateTestRepo(root) - } - - #[test] - fn candidate_resolution_is_the_only_path_that_can_omit_evidence() { - let repo = repo_without_evidence(); - let candidate = resolve_candidate( - &repo.0, - "thread_idx_x", - "LLVM version candidate", - &"a".repeat(64), - "sm_80", - "+ptx70", - ) - .unwrap(); - assert_eq!(candidate.catalog.intrinsics.len(), 1); - assert_eq!(candidate.catalog.intrinsics[0].id, "thread_idx_x"); - assert_eq!(candidate.catalog.intrinsics[0].backend.status, "candidate"); - - let scalar = resolve_candidate( - &repo.0, - "i0390", - "LLVM version candidate", - &"a".repeat(64), - "sm_80", - "+ptx70", - ) - .unwrap(); - assert_eq!(scalar.catalog.intrinsics[0].id, "mul_rn_f64"); - assert!(candidate.catalog.inputs.evidence_sha256.is_empty()); - - for (dense_id, abi_id) in [ - ("mma_m16n8k32_f32_e2m1_e2m1", "i0454"), - ("mma_m16n8k32_f16_e2m1_e2m1", "i0479"), - ] { - let dense_by_name = resolve_candidate( - &repo.0, - dense_id, - "LLVM version candidate", - &"a".repeat(64), - "sm_120f", - "+ptx88", - ) - .unwrap(); - let dense_by_abi = resolve_candidate( - &repo.0, - abi_id, - "LLVM version candidate", - &"a".repeat(64), - "sm_120f", - "+ptx88", - ) - .unwrap(); - assert_eq!( - dense_by_name.catalog.intrinsics, - dense_by_abi.catalog.intrinsics - ); - for lookup in [dense_id, abi_id] { - let error = resolve_candidate( - &repo.0, - lookup, - "LLVM version candidate", - &"a".repeat(64), - "sm_120f", - "+ptx87", - ) - .unwrap_err(); - assert!(error.to_string().contains(dense_id), "{error:#}"); - } - } - - let error = resolve(&repo.0).unwrap_err(); - assert!( - error.to_string().contains("intrinsics/evidence"), - "{error:#}" - ); - let error = resolve_candidate( - &repo.0, - "not_an_intrinsic", - "LLVM version candidate", - &"a".repeat(64), - "sm_80", - "+ptx70", - ) - .unwrap_err(); - assert!(error.to_string().contains("unknown overlay intrinsic")); - } - - #[test] - fn candidate_resolution_cannot_change_normal_catalog_bytes() { - let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); - let before = crate::util::pretty_json(&resolve(&repo_root).unwrap()).unwrap(); - resolve_candidate( - &repo_root, - "thread_idx_x", - "LLVM version candidate", - &"a".repeat(64), - "sm_80", - "+ptx70", - ) - .unwrap(); - let after = crate::util::pretty_json(&resolve(&repo_root).unwrap()).unwrap(); - assert_eq!(before.as_bytes(), after.as_bytes()); - } -} diff --git a/crates/cuda-intrinsics-gen/src/resolve/abi_ledger.rs b/crates/cuda-intrinsics-gen/src/resolve/abi_ledger.rs new file mode 100644 index 0000000000..0e2167de11 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/abi_ledger.rs @@ -0,0 +1,231 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + AbiLedgerEntry, AbiLedgerFile, AbiRawRustSignature, OverlayFile, OverlayIntrinsic, +}; +use anyhow::{Context, Result, bail, ensure}; +use std::collections::{BTreeMap, BTreeSet}; + +use super::guards::*; + +pub(super) struct AbiLedgerIndex<'a> { + by_catalog_id: BTreeMap<&'a str, &'a AbiLedgerEntry>, +} + +impl<'a> AbiLedgerIndex<'a> { + fn new(ledger: &'a AbiLedgerFile) -> Result { + let mut by_catalog_id = BTreeMap::new(); + for entry in &ledger.entries { + ensure!( + by_catalog_id + .insert(entry.catalog_id.as_str(), entry) + .is_none(), + "duplicate ABI ledger catalog ID: {}", + entry.catalog_id + ); + } + Ok(Self { by_catalog_id }) + } + + fn active(&self, catalog_id: &str) -> Result<&'a AbiLedgerEntry> { + let entry = self + .by_catalog_id + .get(catalog_id) + .copied() + .with_context(|| format!("generated intrinsic {catalog_id} has no ABI ledger entry"))?; + ensure!( + entry.status == "active", + "generated intrinsic {catalog_id} maps to non-active ABI ledger entry {}", + entry.abi_id + ); + Ok(entry) + } +} + +pub(super) fn bind_generated_abi_ids( + overlay: &mut OverlayFile, + ledger: &AbiLedgerFile, +) -> Result<()> { + let index = AbiLedgerIndex::new(ledger)?; + for record in overlay + .intrinsics + .iter_mut() + .filter(|record| record.abi_id.is_empty()) + { + let entry = index.active(&record.id)?; + ensure!( + entry.operation_key == record.operation_key, + "generated intrinsic {} operation key mismatch: ledger {:?}, derived {:?}", + record.id, + entry.operation_key, + record.operation_key + ); + let derived_signature = raw_rust_signature(record); + ensure!( + entry.raw_rust_signature == derived_signature, + "generated intrinsic {} raw Rust signature mismatch: ledger {:?}, derived {:?}", + record.id, + entry.raw_rust_signature, + derived_signature + ); + record.abi_id.clone_from(&entry.abi_id); + } + Ok(()) +} + +pub(super) fn validate_abi_ledger(overlay: &OverlayFile, ledger: &AbiLedgerFile) -> Result<()> { + ensure!( + ledger.schema == 1, + "unsupported ABI ledger schema {}", + ledger.schema + ); + ensure!( + ledger.intrinsic_abi == overlay.intrinsic_abi, + "ABI ledger v{} does not match overlay ABI v{}", + ledger.intrinsic_abi, + overlay.intrinsic_abi + ); + ensure!(!ledger.entries.is_empty(), "ABI ledger contains no entries"); + + let overlay_by_abi_id: BTreeMap<_, _> = overlay + .intrinsics + .iter() + .map(|record| (record.abi_id.as_str(), record)) + .collect(); + let mut abi_ids = BTreeSet::new(); + let mut catalog_ids = BTreeSet::new(); + let mut operation_keys = BTreeSet::new(); + let mut previous_abi_id: Option<&str> = None; + for entry in &ledger.entries { + validate_abi_id(&entry.abi_id)?; + if let Some(previous) = previous_abi_id { + ensure!( + previous < entry.abi_id.as_str(), + "ABI ledger IDs must be unique and append-only in ascending order: {} follows {}", + entry.abi_id, + previous + ); + } + previous_abi_id = Some(&entry.abi_id); + insert_unique(&mut abi_ids, &entry.abi_id, "ABI ledger ID")?; + insert_unique(&mut catalog_ids, &entry.catalog_id, "ABI ledger catalog ID")?; + validate_operation_key(&entry.operation_key)?; + insert_unique( + &mut operation_keys, + &entry.operation_key, + "ABI ledger operation key", + )?; + ensure!( + !entry.catalog_id.is_empty() + && !entry.raw_rust_signature.result.is_empty() + && entry + .raw_rust_signature + .arguments + .iter() + .all(|argument| !argument.is_empty()), + "ABI ledger entry {} has incomplete identity data", + entry.abi_id + ); + + let overlay_record = overlay_by_abi_id.get(entry.abi_id.as_str()).copied(); + match entry.status.as_str() { + "active" => { + let record = overlay_record.with_context(|| { + format!( + "active ABI ledger entry {} ({}) has no overlay record", + entry.abi_id, entry.catalog_id + ) + })?; + validate_active_ledger_entry(entry, record)?; + } + "tombstone" => ensure!( + overlay_record.is_none(), + "tombstoned ABI ID {} cannot reappear in the overlay", + entry.abi_id + ), + status => bail!( + "ABI ledger entry {} has unsupported status {status:?}; expected active or tombstone", + entry.abi_id + ), + } + } + + for record in &overlay.intrinsics { + ensure!( + abi_ids.contains(&record.abi_id), + "overlay intrinsic {} uses ABI ID {} with no ledger entry", + record.id, + record.abi_id + ); + } + Ok(()) +} + +pub(super) fn validate_active_ledger_entry( + entry: &AbiLedgerEntry, + record: &OverlayIntrinsic, +) -> Result<()> { + let comparisons = [ + ("catalog ID", entry.catalog_id.as_str(), record.id.as_str()), + ( + "operation key", + entry.operation_key.as_str(), + record.operation_key.as_str(), + ), + ]; + for (field, ledger_value, overlay_value) in comparisons { + ensure!( + ledger_value == overlay_value, + "ABI ledger {} {field} mismatch: ledger {ledger_value:?}, overlay {overlay_value:?}", + entry.abi_id + ); + } + let expected_signature = raw_rust_signature(record); + ensure!( + entry.raw_rust_signature == expected_signature, + "ABI ledger {} raw Rust signature mismatch: ledger {:?}, overlay {:?}", + entry.abi_id, + entry.raw_rust_signature, + expected_signature + ); + Ok(()) +} + +pub(super) fn raw_rust_signature(record: &OverlayIntrinsic) -> AbiRawRustSignature { + AbiRawRustSignature { + safe: record.safe, + arguments: record.rust_arguments.clone(), + result: record.rust_result.clone(), + } +} + +pub(super) fn validate_abi_id(abi_id: &str) -> Result<()> { + ensure!( + abi_id.len() == 5 + && abi_id.starts_with('i') + && abi_id[1..].bytes().all(|byte| byte.is_ascii_digit()), + "intrinsic ABI ID `{abi_id}` must use the stable `iNNNN` form" + ); + Ok(()) +} + +pub(crate) fn validate_operation_key(operation_key: &str) -> Result<()> { + ensure!( + !operation_key.is_empty() + && operation_key.split('.').all(|segment| { + !segment.is_empty() + && segment.bytes().all(|byte| { + byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'_' + }) + }), + "intrinsic operation key `{operation_key}` must contain dot-separated lowercase semantic components" + ); + Ok(()) +} + +pub(super) fn canonical_rust_path(intrinsic_abi: u32, abi_id: &str) -> String { + format!("cuda_intrinsics::__cuda_oxide_intrinsic_abi_v{intrinsic_abi}::{abi_id}") +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/driver.rs b/crates/cuda-intrinsics-gen/src/resolve/driver.rs new file mode 100644 index 0000000000..41a224c375 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/driver.rs @@ -0,0 +1,765 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::extract::{IMPORTED_SCHEMA, read_upstream_lock}; +use crate::model::{ + AbiLedgerFile, BackendLoweringMechanism, CatalogBackend, CatalogFile, + CatalogHardwareAlternative, CatalogHardwareTarget, CatalogInputs, CatalogSource, + CatalogTargetRequirement, ExecutionControlOperation, ImportedFile, ImportedIntrinsic, + IntrinsicBackend, IntrinsicSource, OverlayFile, OverlayIntrinsic, PtxVersion, +}; +#[cfg(test)] +use crate::model::{ + CatalogBackendLowering, ClcAdmission, ClcOperation, RuntimeValidation, Tcgen05Admission, + Tcgen05LdAdmissionVariant, Tcgen05Operation, Tcgen05StAdmissionVariant, TmaAdmission, +}; +use crate::util::{read_json, sha256_file}; +#[cfg(test)] +use anyhow::bail; +use anyhow::{Context, Result, ensure}; +use std::collections::BTreeMap; +use std::fs; +use std::path::Path; + +use super::abi_ledger::*; +use super::evidence::*; +use super::families::*; +use super::materialize::*; +use super::overlay::*; +use super::policy::*; +use super::targets::*; + +pub(super) struct ResolutionBase { + pub(super) overlay: OverlayFile, + pub(super) imported: ImportedFile, + pub(super) source: CatalogSource, + pub(super) imported_sha256: String, + pub(super) overlay_sha256: String, + pub(super) abi_ledger_sha256: String, +} + +#[derive(Debug)] +pub(crate) struct CandidateResolution { + pub catalog: CatalogFile, + pub mechanism: BackendLoweringMechanism, + pub requirement: CatalogTargetRequirement, +} + +pub(super) fn primary_evidence_profile<'a>( + policy: &'a OverlayIntrinsic, + default_profile: &'a str, +) -> Result<&'a str> { + let needs_family_profile = policy.tma.is_some() + || ExecutionControlOperation::from_catalog_id(&policy.id).is_some() + || policy + .tcgen05 + .as_ref() + .and_then(|tcgen05| tcgen05.mma.as_ref()) + .is_some(); + if !needs_family_profile { + return Ok(default_profile); + } + policy + .backend_lowerings + .iter() + .find(|lowering| lowering.backend == IntrinsicBackend::LlvmNvptx) + .map(|lowering| lowering.evidence_profile.as_str()) + .with_context(|| format!("{} has no LLVM family evidence route", policy.id)) +} + +pub fn resolve(repo_root: &Path) -> Result { + let base = load_resolution_base(repo_root)?; + let ResolutionBase { + overlay, + imported, + source, + imported_sha256, + overlay_sha256, + abi_ledger_sha256, + } = base; + let imported_by_record = index_imported_intrinsics(&imported)?; + let llvm_revision = source.llvm_revision.clone(); + let (evidence_files, evidence_hashes) = read_evidence(repo_root)?; + let evidence_by_profile_id = index_evidence(&evidence_files, &llvm_revision)?; + + let mut intrinsics = Vec::with_capacity(overlay.intrinsics.len()); + for policy in &overlay.intrinsics { + let source = resolve_policy_source(policy)?; + let declaration = resolve_imported_declaration(policy, &source, &imported_by_record)?; + validate_special_register_llvm_exclusion(policy, &imported_by_record)?; + validate_policy(policy, &source, declaration, overlay.intrinsic_abi)?; + let primary_profile = primary_evidence_profile(policy, &overlay.backend_profile)?; + let evidence = evidence_by_profile_id + .get(&(primary_profile, policy.id.as_str())) + .with_context(|| { + format!( + "intrinsic {} has no evidence record in selected profile {}", + policy.id, primary_profile + ) + })?; + validate_evidence(policy, evidence, None)?; + let backend_lowerings = resolve_backend_lowerings(policy, &evidence_by_profile_id)?; + intrinsics.push(resolve_record( + policy, + source, + declaration, + evidence.record, + primary_profile, + evidence.backend_version, + evidence.backend_sha256, + backend_lowerings, + overlay.intrinsic_abi, + )?); + } + for (_, evidence_id) in evidence_by_profile_id.keys() { + ensure!( + overlay + .intrinsics + .iter() + .any(|record| record.id == *evidence_id), + "evidence exists for unknown catalog ID {evidence_id}" + ); + } + + Ok(CatalogFile { + schema: CATALOG_SCHEMA, + catalog_version: overlay.catalog_version, + intrinsic_abi: overlay.intrinsic_abi, + generator_version: env!("CARGO_PKG_VERSION").to_owned(), + source, + inputs: CatalogInputs { + imported_sha256, + overlay_sha256, + abi_ledger_sha256, + evidence_sha256: evidence_hashes, + }, + intrinsics, + }) +} + +#[cfg(test)] +pub(crate) fn test_catalog_with_clc(repo_root: &Path) -> Result { + let mut catalog = resolve(repo_root)?; + let active_count = catalog + .intrinsics + .iter() + .filter(|record| record.family == "clc") + .count(); + if active_count != 0 { + ensure!( + active_count == 6, + "active CLC family must contain six records" + ); + return Ok(catalog); + } + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json"))?; + let imported_by_record = index_imported_intrinsics(&imported)?; + let operations = [ + ClcOperation::TryCancel, + ClcOperation::TryCancelMulticast, + ClcOperation::QueryIsCanceled, + ClcOperation::QueryGetFirstCtaidX, + ClcOperation::QueryGetFirstCtaidY, + ClcOperation::QueryGetFirstCtaidZ, + ]; + let admission = ClcAdmission { + llvm_evidence_profile: "llvm-clc-test".into(), + libnvvm_evidence_profile: "libnvvm-clc-test".into(), + runtime_validation: RuntimeValidation::Unexecuted, + variants: operations + .into_iter() + .map(|operation| crate::model::ClcAdmissionVariant { + abi_id: clc_recipe(operation).abi_id.into(), + operation, + }) + .collect(), + }; + for policy in expand_clc_admission(&admission)? { + let source = resolve_policy_source(&policy)?; + let declaration = resolve_imported_declaration(&policy, &source, &imported_by_record)?; + validate_policy(&policy, &source, declaration, catalog.intrinsic_abi)?; + let backend_lowerings = policy + .backend_lowerings + .iter() + .map(|lowering| { + Ok(CatalogBackendLowering { + backend: lowering.backend, + mechanism: lowering.mechanism, + evidence_profile: lowering.evidence_profile.clone(), + target: backend_target_requirement(&policy, lowering)?, + version: "test".into(), + sha256: "0".repeat(64), + artifact_path: None, + build_id_prefix: None, + status: "validated".into(), + stages: vec![], + }) + }) + .collect::>>()?; + catalog.intrinsics.push(materialize_record( + &policy, + source, + declaration, + CatalogBackend { + profile: "clc-test".into(), + version: "test".into(), + sha256: "0".repeat(64), + status: "validated".into(), + target_triple: "nvptx64-nvidia-cuda".into(), + gpu_target: "sm_100".into(), + ptx_feature: "+ptx86".into(), + }, + backend_lowerings, + catalog.intrinsic_abi, + )?); + } + catalog + .intrinsics + .sort_by(|left, right| left.id.cmp(&right.id)); + Ok(catalog) +} + +#[cfg(test)] +pub(crate) fn test_catalog_with_tma(repo_root: &Path) -> Result { + let mut catalog = resolve(repo_root)?; + let active_count = catalog + .intrinsics + .iter() + .filter(|record| record.tma.is_some()) + .count(); + let expected_count = TMA_OPERATIONS.len() + tma_reduction_matrix().len(); + if active_count == expected_count { + return Ok(catalog); + } + ensure!( + active_count == 0 || active_count == TMA_OPERATIONS.len(), + "active TMA catalog has {active_count} records; expected 0, {} base records, or {expected_count} total records", + TMA_OPERATIONS.len() + ); + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json"))?; + let imported_by_record = index_imported_intrinsics(&imported)?; + let admission = TmaAdmission { + llvm_evidence_profile: "llvm-tma-test".into(), + libnvvm_evidence_profile: "libnvvm-tma-test".into(), + reduce_llvm_evidence_profile: Some("llvm-tma-reduce-test".into()), + reduce_libnvvm_evidence_profile: Some("libnvvm-tma-reduce-test".into()), + runtime_validation: RuntimeValidation::Unexecuted, + variants: TMA_OPERATIONS + .into_iter() + .map(|operation| crate::model::TmaAdmissionVariant { + abi_id: tma_recipe(operation).abi_id.into(), + operation, + }) + .collect(), + reduce_variants: tma_reduction_admission_variants(), + }; + for policy in expand_tma_admission(&admission)? { + if catalog + .intrinsics + .iter() + .any(|record| record.id.as_str() == policy.id.as_str()) + { + continue; + } + let source = resolve_policy_source(&policy)?; + let declaration = resolve_imported_declaration(&policy, &source, &imported_by_record)?; + validate_policy(&policy, &source, declaration, catalog.intrinsic_abi)?; + let backend_lowerings = policy + .backend_lowerings + .iter() + .map(|lowering| { + Ok(CatalogBackendLowering { + backend: lowering.backend, + mechanism: lowering.mechanism, + evidence_profile: lowering.evidence_profile.clone(), + target: backend_target_requirement(&policy, lowering)?, + version: "test".into(), + sha256: "0".repeat(64), + artifact_path: None, + build_id_prefix: None, + status: "validated".into(), + stages: vec![], + }) + }) + .collect::>>()?; + catalog.intrinsics.push(materialize_record( + &policy, + source, + declaration, + CatalogBackend { + profile: "tma-test".into(), + version: "test".into(), + sha256: "0".repeat(64), + status: "validated".into(), + target_triple: "nvptx64-nvidia-cuda".into(), + gpu_target: "sm_90".into(), + ptx_feature: "+ptx80".into(), + }, + backend_lowerings, + catalog.intrinsic_abi, + )?); + } + catalog + .intrinsics + .sort_by(|left, right| left.id.cmp(&right.id)); + Ok(catalog) +} + +#[cfg(test)] +pub(crate) fn test_catalog_with_tcgen05(repo_root: &Path) -> Result { + let mut catalog = resolve(repo_root)?; + let active_count = catalog + .intrinsics + .iter() + .filter(|record| record.family == "tcgen05") + .count(); + match active_count { + 233 => return Ok(catalog), + 0 => {} + count => bail!("active tcgen05 catalog has {count} of 233 records"), + } + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json"))?; + let imported_by_record = index_imported_intrinsics(&imported)?; + let operations = [ + Tcgen05Operation::Alloc, + Tcgen05Operation::Dealloc, + Tcgen05Operation::RelinquishAllocPermit, + Tcgen05Operation::FenceBeforeThreadSync, + Tcgen05Operation::FenceAfterThreadSync, + Tcgen05Operation::Commit, + Tcgen05Operation::CommitSharedCluster, + Tcgen05Operation::MmaWsF16, + Tcgen05Operation::MmaF16, + Tcgen05Operation::MmaWsBf16, + Tcgen05Operation::MmaWsTf32, + Tcgen05Operation::CpSmemToTmem, + Tcgen05Operation::Ld16x256bX8Pure, + Tcgen05Operation::Ld16x256bPure, + Tcgen05Operation::LoadWait, + Tcgen05Operation::StoreWait, + Tcgen05Operation::AllocCg2, + Tcgen05Operation::DeallocCg2, + Tcgen05Operation::RelinquishAllocPermitCg2, + Tcgen05Operation::MmaF16Cg2, + Tcgen05Operation::CommitCg2, + Tcgen05Operation::CommitSharedClusterCg2, + Tcgen05Operation::CommitMulticastCg2, + Tcgen05Operation::CpSmemToTmemCg2, + Tcgen05Operation::CommitMulticast, + Tcgen05Operation::ShiftDown, + Tcgen05Operation::ShiftDownCg2, + ]; + let admission = Tcgen05Admission { + llvm_evidence_profile: "llvm-tcgen05-test".into(), + libnvvm_evidence_profile: "libnvvm-tcgen05-test".into(), + cp_llvm_evidence_profile: None, + cp_libnvvm_evidence_profile: None, + ld_llvm_evidence_profile: None, + ld_libnvvm_evidence_profile: None, + st_llvm_evidence_profile: None, + st_libnvvm_evidence_profile: None, + offset_llvm_evidence_profile: Some("llvm-tcgen05-offset-test".into()), + offset_libnvvm_evidence_profile: Some("libnvvm-tcgen05-offset-test".into()), + control_llvm_evidence_profile: Some("llvm-tcgen05-control-test".into()), + control_libnvvm_evidence_profile: Some("libnvvm-tcgen05-control-test".into()), + mma_llvm_evidence_profile: None, + mma_libnvvm_evidence_profile: None, + mma_llvm_target_contracts: vec![], + mma_libnvvm_target_contracts: vec![], + runtime_validation: RuntimeValidation::Unexecuted, + variants: operations + .into_iter() + .map(|operation| crate::model::Tcgen05AdmissionVariant { + abi_id: tcgen05_recipe(operation).abi_id.into(), + operation, + }) + .collect(), + cp_variants: vec![], + ld_variants: vec![], + st_variants: vec![], + ld_offset_variants: TCGEN05_OFFSET_LDST_VARIANTS + .into_iter() + .flat_map(|(shape, multiplicity)| { + [false, true] + .into_iter() + .map(move |pack16| (shape, multiplicity, pack16)) + }) + .enumerate() + .map( + |(index, (shape, multiplicity, pack16))| Tcgen05LdAdmissionVariant { + abi_id: format!("i{:04}", 728 + index), + shape, + multiplicity, + pack16, + }, + ) + .collect(), + st_offset_variants: TCGEN05_OFFSET_LDST_VARIANTS + .into_iter() + .flat_map(|(shape, multiplicity)| { + [false, true] + .into_iter() + .map(move |unpack16| (shape, multiplicity, unpack16)) + }) + .enumerate() + .map( + |(index, (shape, multiplicity, unpack16))| Tcgen05StAdmissionVariant { + abi_id: format!("i{:04}", 744 + index), + shape, + multiplicity, + unpack16, + }, + ) + .collect(), + mma_variants: vec![], + }; + for policy in expand_tcgen05_admission(&admission)? { + let source = resolve_policy_source(&policy)?; + let declaration = resolve_imported_declaration(&policy, &source, &imported_by_record)?; + validate_policy(&policy, &source, declaration, catalog.intrinsic_abi)?; + let backend_lowerings = policy + .backend_lowerings + .iter() + .map(|lowering| { + Ok(CatalogBackendLowering { + backend: lowering.backend, + mechanism: lowering.mechanism, + evidence_profile: lowering.evidence_profile.clone(), + target: backend_target_requirement(&policy, lowering)?, + version: "test".into(), + sha256: "0".repeat(64), + artifact_path: None, + build_id_prefix: None, + status: "validated".into(), + stages: vec![], + }) + }) + .collect::>>()?; + catalog.intrinsics.push(materialize_record( + &policy, + source, + declaration, + CatalogBackend { + profile: "tcgen05-test".into(), + version: "test".into(), + sha256: "0".repeat(64), + status: "validated".into(), + target_triple: "nvptx64-nvidia-cuda".into(), + gpu_target: "sm_100a".into(), + ptx_feature: "+ptx86".into(), + }, + backend_lowerings, + catalog.intrinsic_abi, + )?); + } + catalog + .intrinsics + .sort_by(|left, right| left.id.cmp(&right.id)); + Ok(catalog) +} + +pub(super) fn load_resolution_base(repo_root: &Path) -> Result { + let lock = read_upstream_lock(repo_root)?; + let imported_path = repo_root.join("intrinsics/imported.json"); + let overlay_path = repo_root.join("intrinsics/overlay.toml"); + let imported: ImportedFile = read_json(&imported_path)?; + let (mut overlay, overlay_sha256) = read_overlay(repo_root, &overlay_path)?; + let ledger_path = repo_root.join(format!("intrinsics/abi-v{}.toml", overlay.intrinsic_abi)); + let ledger_text = fs::read_to_string(&ledger_path) + .with_context(|| format!("read {}", ledger_path.display()))?; + let ledger: AbiLedgerFile = + toml::from_str(&ledger_text).with_context(|| format!("parse {}", ledger_path.display()))?; + + ensure!( + imported.schema == IMPORTED_SCHEMA, + "unsupported imported.json schema {}", + imported.schema + ); + ensure!( + overlay.schema == OVERLAY_SCHEMA, + "unsupported overlay.toml schema {}", + overlay.schema + ); + ensure!( + overlay.intrinsic_abi > 0, + "intrinsic_abi must be a positive integer" + ); + ensure!( + imported.source.llvm_revision == lock.llvm.revision, + "imported facts use LLVM {}, but upstream.lock pins {}", + imported.source.llvm_revision, + lock.llvm.revision + ); + ensure!( + imported.source.llvm_tblgen_source_revision == lock.llvm.revision, + "imported facts were not produced by llvm-tblgen built from the pinned source" + ); + ensure!( + imported.source.llvm_tblgen_version == lock.llvm_tblgen.version_line, + "imported facts use llvm-tblgen {:?}, but upstream.lock pins {:?}", + imported.source.llvm_tblgen_version, + lock.llvm_tblgen.version_line + ); + ensure!( + imported.source.intrinsics_json_sha256 == lock.dumps.intrinsics_sha256, + "imported intrinsic dump hash does not match upstream.lock" + ); + ensure!( + imported.source.nvptx_json_sha256 == lock.dumps.nvptx_sha256, + "imported NVPTX dump hash does not match upstream.lock" + ); + let imported_sha256 = sha256_file(&imported_path)?; + ensure!( + imported_sha256 == lock.dumps.normalized_imported_sha256, + "normalized imported.json hash mismatch: upstream.lock records {}, found {}; regenerate from the pinned dumps, and refresh the lock explicitly only for an intentional normalizer change", + lock.dumps.normalized_imported_sha256, + imported_sha256 + ); + + bind_generated_abi_ids(&mut overlay, &ledger)?; + overlay + .intrinsics + .sort_by(|left, right| left.id.cmp(&right.id)); + validate_execution_control_family_completeness(&overlay.intrinsics)?; + validate_unique_overlay(&overlay.intrinsics, overlay.intrinsic_abi)?; + validate_abi_ledger(&overlay, &ledger)?; + Ok(ResolutionBase { + overlay, + imported, + source: CatalogSource { + llvm_repository: lock.llvm.repository, + llvm_revision: lock.llvm.revision, + llvm_tblgen_version: lock.llvm_tblgen.version_line, + llvm_tblgen_source_revision: lock + .llvm_tblgen + .built_from_llvm_revision + .context("pinned llvm-tblgen has no source revision")?, + }, + imported_sha256, + overlay_sha256, + abi_ledger_sha256: sha256_file(&ledger_path)?, + }) +} + +pub(super) fn index_imported_intrinsics( + imported: &ImportedFile, +) -> Result> { + let imported_by_record: BTreeMap<_, _> = imported + .intrinsics + .iter() + .map(|intrinsic| (intrinsic.source_record.as_str(), intrinsic)) + .collect(); + ensure!( + imported_by_record.len() == imported.intrinsics.len(), + "imported.json contains duplicate source records" + ); + Ok(imported_by_record) +} + +pub(super) fn resolve_imported_declaration<'a>( + policy: &OverlayIntrinsic, + source: &IntrinsicSource, + imported_by_record: &'a BTreeMap<&str, &'a ImportedIntrinsic>, +) -> Result> { + match source { + IntrinsicSource::LlvmImported { source_record } => Ok(Some( + *imported_by_record + .get(source_record.as_str()) + .with_context(|| { + format!( + "overlay intrinsic {} references missing imported record {}", + policy.id, source_record + ) + })?, + )), + IntrinsicSource::PtxNative { .. } => Ok(None), + } +} + +pub(crate) fn resolve_candidate( + repo_root: &Path, + intrinsic_id: &str, + backend_version: &str, + backend_sha256: &str, + gpu_target: &str, + ptx_feature: &str, +) -> Result { + let base = load_resolution_base(repo_root)?; + let imported_by_record = index_imported_intrinsics(&base.imported)?; + let policy = base + .overlay + .intrinsics + .iter() + .find(|policy| policy.id == intrinsic_id || policy.abi_id == intrinsic_id) + .with_context(|| format!("unknown overlay intrinsic {intrinsic_id}"))?; + let source = resolve_policy_source(policy)?; + let declaration = resolve_imported_declaration(policy, &source, &imported_by_record)?; + validate_special_register_llvm_exclusion(policy, &imported_by_record)?; + validate_policy(policy, &source, declaration, base.overlay.intrinsic_abi)?; + ensure!( + !backend_version.trim().is_empty() + && backend_sha256.len() == 64 + && backend_sha256.bytes().all(|byte| byte.is_ascii_hexdigit()), + "candidate backend identity is incomplete" + ); + let (mechanism, target) = candidate_llvm_route(policy)?; + validate_candidate_target(policy, &target, gpu_target, ptx_feature)?; + + let record = materialize_record( + policy, + source, + declaration, + CatalogBackend { + profile: "candidate-comparison".into(), + version: backend_version.into(), + sha256: backend_sha256.into(), + status: "candidate".into(), + target_triple: "nvptx64-nvidia-cuda".into(), + gpu_target: gpu_target.into(), + ptx_feature: ptx_feature.into(), + }, + Vec::new(), + base.overlay.intrinsic_abi, + )?; + Ok(CandidateResolution { + catalog: CatalogFile { + schema: CATALOG_SCHEMA, + catalog_version: base.overlay.catalog_version, + intrinsic_abi: base.overlay.intrinsic_abi, + generator_version: env!("CARGO_PKG_VERSION").to_owned(), + source: base.source, + inputs: CatalogInputs { + imported_sha256: base.imported_sha256, + overlay_sha256: base.overlay_sha256, + abi_ledger_sha256: base.abi_ledger_sha256, + evidence_sha256: Vec::new(), + }, + intrinsics: vec![record], + }, + mechanism, + requirement: target, + }) +} + +pub(super) fn candidate_llvm_route( + policy: &OverlayIntrinsic, +) -> Result<(BackendLoweringMechanism, CatalogTargetRequirement)> { + let routes = policy + .backend_lowerings + .iter() + .filter(|lowering| lowering.backend == IntrinsicBackend::LlvmNvptx) + .collect::>(); + ensure!( + routes.len() <= 1, + "{} has more than one LLVM-NVPTX route", + policy.id + ); + if let Some(route) = routes.first() { + return Ok((route.mechanism, backend_target_requirement(policy, route)?)); + } + ensure!( + matches!( + resolve_policy_source(policy)?, + IntrinsicSource::LlvmImported { .. } + ), + "{} has no LLVM-NVPTX candidate route", + policy.id + ); + Ok(( + BackendLoweringMechanism::TypedNvvm, + CatalogTargetRequirement { + minimum_ptx: parse_ptx_version(&policy.minimum_ptx, &policy.id)?, + hardware: parse_hardware_target(policy)?, + }, + )) +} + +pub(super) fn validate_candidate_target( + policy: &OverlayIntrinsic, + requirement: &CatalogTargetRequirement, + gpu_target: &str, + ptx_feature: &str, +) -> Result<()> { + ensure!( + gpu_target.starts_with("sm_"), + "candidate GPU target {gpu_target:?} must use sm_NN, sm_NNa, or sm_NNf" + ); + let hardware = parse_stage_hardware(gpu_target).with_context(|| { + format!("candidate GPU target {gpu_target:?} must use sm_NN, sm_NNa, or sm_NNf") + })?; + ensure!( + describe_stage_hardware(hardware) == gpu_target, + "candidate GPU target {gpu_target:?} is not canonical" + ); + let ptx = parse_candidate_ptx_feature(ptx_feature)?; + let effective_minimum_ptx = if is_f8f6f4_mma_target_matrix_policy(policy) { + f8f6f4_llvm_ptx_floor(hardware)? + } else { + target_requirement_ptx_floor(requirement, hardware, true).with_context(|| { + format!( + "candidate GPU target {gpu_target} does not satisfy {} hardware requirement {:?}", + policy.id, requirement.hardware + ) + })? + }; + ensure!( + ptx.encoded() >= effective_minimum_ptx, + "candidate target {gpu_target} / {ptx_feature} is below {} PTX floor {}.{}", + policy.id, + effective_minimum_ptx / 10, + effective_minimum_ptx % 10 + ); + let hardware_matches = target_requirement_ptx_floor(requirement, hardware, true).is_some(); + ensure!( + hardware_matches, + "candidate GPU target {gpu_target} does not satisfy {} hardware requirement {:?}", + policy.id, + requirement.hardware + ); + Ok(()) +} + +pub(super) fn target_requirement_ptx_floor( + requirement: &CatalogTargetRequirement, + hardware: CatalogHardwareAlternative, + allow_forward_minimum: bool, +) -> Option { + match &requirement.hardware { + CatalogHardwareTarget::All => Some(requirement.minimum_ptx.encoded()), + CatalogHardwareTarget::AnyOf { alternatives } => alternatives + .iter() + .any(|expected| { + selected_stage_hardware_matches(hardware, *expected, allow_forward_minimum) + }) + .then(|| requirement.minimum_ptx.encoded()), + CatalogHardwareTarget::TargetMatrix { contracts } => contracts + .iter() + .flat_map(|contract| contract.alternatives.iter()) + .filter(|alternative| { + selected_stage_hardware_matches( + hardware, + alternative.hardware, + allow_forward_minimum, + ) + }) + .map(|alternative| alternative.minimum_ptx.encoded()) + .min(), + } +} + +pub(super) fn parse_candidate_ptx_feature(value: &str) -> Result { + let digits = value + .strip_prefix("+ptx") + .with_context(|| format!("candidate PTX feature {value:?} must use +ptxNN"))?; + ensure!( + matches!(digits.len(), 2 | 3) && digits.bytes().all(|byte| byte.is_ascii_digit()), + "candidate PTX feature {value:?} must use +ptxNN" + ); + let split = digits.len() - 1; + let version = format!("{}.{}", &digits[..split], &digits[split..]); + version + .parse() + .map_err(|reason: String| anyhow::anyhow!("candidate PTX feature {value:?}: {reason}")) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/evidence.rs b/crates/cuda-intrinsics-gen/src/resolve/evidence.rs new file mode 100644 index 0000000000..30387a6c9a --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/evidence.rs @@ -0,0 +1,1504 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, CatalogHardwareAlternative, CatalogHardwareTarget, + CatalogTargetAlternative, EvidenceArtifactKind, EvidenceFile, EvidenceFileV6, EvidenceMatrix, + EvidenceMatrixTemplate, EvidenceRecord, EvidenceRecordDefaults, EvidenceStage, + EvidenceStageKind, ImportedAddressSpace, IntrinsicBackend, IntrinsicSource, LdmatrixShape, + OverlayBackendLowering, OverlayIntrinsic, RuntimeValidation, +}; +use crate::ptx::OperandPattern; +use crate::util::sha256_file; +use anyhow::{Context, Result, bail, ensure}; +use cuda_target_spec::{CudaArch, recorded_ptx_floor}; +use std::collections::{BTreeMap, BTreeSet}; +use std::fs; +use std::path::{Path, PathBuf}; + +use super::families::*; +use super::policy::*; +use super::targets::*; + +pub(super) fn read_evidence_file(path: &Path) -> Result { + let bytes = fs::read(path).with_context(|| format!("read {}", path.display()))?; + parse_evidence_bytes(&bytes, &path.display().to_string()) +} + +pub(super) fn parse_evidence_bytes(bytes: &[u8], source: &str) -> Result { + #[derive(serde::Deserialize)] + struct Schema { + schema: u32, + } + + let schema: Schema = + serde_json::from_slice(bytes).with_context(|| format!("parse JSON {source}"))?; + match schema.schema { + 2..=5 => serde_json::from_slice(bytes) + .with_context(|| format!("parse legacy evidence JSON {source}")), + 6 => { + let file: EvidenceFileV6 = serde_json::from_slice(bytes) + .with_context(|| format!("parse matrix evidence JSON {source}"))?; + expand_evidence_file_v6(file) + .with_context(|| format!("expand matrix evidence {source}")) + } + _ => bail!("unsupported evidence schema in {source}"), + } +} + +pub(super) fn expand_evidence_file_v6(file: EvidenceFileV6) -> Result { + ensure!(file.schema == 6, "matrix evidence must use schema 6"); + ensure!( + !file.records.is_empty() || !file.matrices.is_empty(), + "schema-6 evidence contains no records or matrices" + ); + reject_default_placeholders(&file.defaults, "evidence defaults", false)?; + + let mut fixture_by_id = BTreeMap::new(); + let mut fixture_coverage = BTreeMap::new(); + for fixture in &file.fixtures { + ensure!( + is_safe_matrix_token(&fixture.id), + "evidence fixture ID {:?} is not a safe token", + fixture.id + ); + ensure!( + fixture.coverage_count > 0 && !fixture.stages.is_empty(), + "evidence fixture {} has no coverage or stages", + fixture.id + ); + reject_stage_placeholders(&fixture.stages, &format!("fixture {}", fixture.id))?; + ensure!( + fixture_by_id.insert(fixture.id.as_str(), fixture).is_none(), + "duplicate evidence fixture ID {}", + fixture.id + ); + fixture_coverage.insert(fixture.id.as_str(), 0usize); + } + + let mut records = file.records; + let mut record_ids = BTreeSet::new(); + for record in &records { + ensure!( + record_ids.insert(record.id.clone()), + "duplicate expanded evidence ID {}", + record.id + ); + validate_stage_pairs(&record.stages, &record.id)?; + } + + for matrix in &file.matrices { + let (expanded, referenced_fixtures) = + expand_evidence_matrix(matrix, &file.defaults, &fixture_by_id)?; + for fixture_id in referenced_fixtures { + *fixture_coverage + .get_mut(fixture_id.as_str()) + .expect("validated fixture reference") += expanded.len(); + } + for record in expanded { + ensure!( + record_ids.insert(record.id.clone()), + "duplicate expanded evidence ID {}", + record.id + ); + records.push(record); + } + } + + for fixture in &file.fixtures { + let actual = fixture_coverage[fixture.id.as_str()]; + ensure!( + actual > 0, + "evidence fixture {} is not referenced by any matrix", + fixture.id + ); + ensure!( + actual == fixture.coverage_count, + "evidence fixture {} covers {actual} expanded records, expected {}", + fixture.id, + fixture.coverage_count + ); + } + + Ok(EvidenceFile { + schema: file.schema, + backend_profile: file.backend_profile, + backend_kind: file.backend_kind, + llvm_revision: file.llvm_revision, + backend_version: file.backend_version, + backend_sha256: file.backend_sha256, + artifact_path: file.artifact_path, + build_id_prefix: file.build_id_prefix, + nvvm_ir_version: file.nvvm_ir_version, + debug_ir_version: file.debug_ir_version, + records, + }) +} + +pub(super) fn expand_evidence_matrix( + matrix: &EvidenceMatrix, + defaults: &EvidenceRecordDefaults, + fixtures: &BTreeMap<&str, &crate::model::EvidenceFixture>, +) -> Result<(Vec, Vec)> { + ensure!(!matrix.axes.is_empty(), "evidence matrix has no axes"); + let mut previous_axis: Option<&str> = None; + let mut product_count = 1usize; + let mut bindings = vec![BTreeMap::::new()]; + for axis in &matrix.axes { + ensure!( + is_safe_matrix_token(&axis.name), + "evidence matrix axis {:?} is not a safe token", + axis.name + ); + if let Some(previous) = previous_axis { + ensure!( + previous < axis.name.as_str(), + "evidence matrix axes must be unique and sorted: {} follows {}", + axis.name, + previous + ); + } + previous_axis = Some(&axis.name); + ensure!( + !axis.values.is_empty(), + "evidence matrix axis {} has no values", + axis.name + ); + let mut values = BTreeSet::new(); + for value in &axis.values { + ensure!( + is_safe_matrix_token(value), + "evidence matrix axis {} has unsafe value {:?}", + axis.name, + value + ); + ensure!( + values.insert(value.as_str()), + "evidence matrix axis {} has duplicate value {:?}", + axis.name, + value + ); + } + product_count = product_count + .checked_mul(axis.values.len()) + .context("evidence matrix product count overflow")?; + let mut next = Vec::with_capacity(product_count); + for binding in bindings { + for value in &axis.values { + let mut expanded = binding.clone(); + expanded.insert(axis.name.clone(), value.clone()); + next.push(expanded); + } + } + bindings = next; + } + ensure!( + product_count == matrix.product_count, + "evidence matrix expands to {product_count} records, expected {}", + matrix.product_count + ); + ensure!( + !matrix.fixtures.is_empty(), + "evidence matrix references no shared fixture" + ); + + let mut fixture_ids = BTreeSet::new(); + let mut previous_fixture: Option<&str> = None; + let mut fixture_stages = Vec::new(); + for fixture_id in &matrix.fixtures { + if let Some(previous) = previous_fixture { + ensure!( + previous < fixture_id.as_str(), + "evidence matrix fixtures must be unique and sorted: {fixture_id} follows {previous}" + ); + } + previous_fixture = Some(fixture_id); + let fixture = fixtures + .get(fixture_id.as_str()) + .with_context(|| format!("evidence matrix references unknown fixture {fixture_id}"))?; + fixture_ids.insert(fixture_id.clone()); + fixture_stages.extend(fixture.stages.iter().cloned()); + } + + reject_default_placeholders(&matrix.template.facts, "matrix template facts", true)?; + validate_matrix_identity(&matrix.template)?; + let axis_names = matrix + .axes + .iter() + .map(|axis| axis.name.as_str()) + .collect::>(); + let mut used_axes = BTreeSet::new(); + let mut records = Vec::with_capacity(product_count); + for binding in &bindings { + let record = materialize_evidence_record( + &matrix.template, + defaults, + binding, + &mut used_axes, + &fixture_stages, + )?; + validate_stage_pairs(&record.stages, &record.id)?; + records.push(record); + } + for axis in axis_names { + ensure!( + used_axes.contains(axis), + "evidence matrix declares unused axis {axis}" + ); + } + Ok((records, fixture_ids.into_iter().collect())) +} + +pub(super) fn validate_matrix_identity(template: &EvidenceMatrixTemplate) -> Result<()> { + ensure!( + !template.id.is_empty(), + "evidence matrix template has an empty ID" + ); + match (&template.source, &template.source_record) { + (Some(_), None) | (None, Some(_)) => {} + (Some(_), Some(_)) => bail!("evidence matrix template mixes source forms"), + (None, None) => bail!("evidence matrix template has no source"), + } + reject_disallowed_placeholder(&template.expected_ptx.mnemonic, "PTX mnemonic")?; + Ok(()) +} + +pub(super) fn materialize_evidence_record( + template: &EvidenceMatrixTemplate, + defaults: &EvidenceRecordDefaults, + bindings: &BTreeMap, + used_axes: &mut BTreeSet, + fixture_stages: &[EvidenceStage], +) -> Result { + let id = expand_axis_placeholders(&template.id, bindings, used_axes, "evidence ID")?; + let source = template + .source + .as_ref() + .map(|source| expand_evidence_source(source, bindings, used_axes)) + .transpose()?; + let source_record = template + .source_record + .as_deref() + .map(|value| expand_axis_placeholders(value, bindings, used_axes, "source record")) + .transpose()?; + let llvm_symbol = template + .llvm_symbol + .as_deref() + .map(|value| expand_axis_placeholders(value, bindings, used_axes, "LLVM symbol")) + .transpose()?; + validate_expanded_matrix_identity( + &id, + source.as_ref(), + source_record.as_deref(), + llvm_symbol.as_deref(), + )?; + let resolved_llvm_symbol = select_fact( + &template.facts.resolved_llvm_symbol, + &defaults.resolved_llvm_symbol, + ) + .map(|value| expand_axis_placeholders(&value, bindings, used_axes, "resolved LLVM symbol")) + .transpose()?; + let mut expected_ptx = template.expected_ptx.clone(); + for modifier in &mut expected_ptx.modifiers { + *modifier = expand_axis_placeholders(modifier, bindings, used_axes, "PTX modifier")?; + } + for operand in &mut expected_ptx.operands { + if let OperandPattern::Exact { value } = operand { + *value = expand_axis_placeholders(value, bindings, used_axes, "exact PTX operand")?; + } + } + + let mut stages = defaults.stages.clone(); + stages.extend(template.facts.stages.iter().cloned()); + stages.extend(fixture_stages.iter().cloned()); + let target_triple = required_fact( + select_fact(&template.facts.target_triple, &defaults.target_triple), + &id, + "target triple", + )?; + let gpu_target = required_fact( + select_fact(&template.facts.gpu_target, &defaults.gpu_target), + &id, + "GPU target", + )?; + let ptx_feature = required_fact( + select_fact(&template.facts.ptx_feature, &defaults.ptx_feature), + &id, + "PTX feature", + )?; + let status = required_fact( + select_fact(&template.facts.status, &defaults.status), + &id, + "status", + )?; + let source_is_ptx_native = matches!(source, Some(IntrinsicSource::PtxNative { .. })); + Ok(EvidenceRecord { + id, + source, + source_record, + llvm_symbol, + resolved_llvm_symbol, + llvm_arguments: select_fact(&template.facts.llvm_arguments, &defaults.llvm_arguments) + .unwrap_or_default(), + llvm_results: select_fact(&template.facts.llvm_results, &defaults.llvm_results) + .unwrap_or_default(), + concrete_llvm_arguments: select_fact( + &template.facts.concrete_llvm_arguments, + &defaults.concrete_llvm_arguments, + ) + .unwrap_or_default(), + concrete_llvm_results: select_fact( + &template.facts.concrete_llvm_results, + &defaults.concrete_llvm_results, + ) + .unwrap_or_default(), + target_triple, + gpu_target, + ptx_feature, + status, + stages, + // Declaration canonicalization is a statement about an imported + // LLVM declaration; a PTX-native record has none, so the file-level + // default must not invent the fact for it. + declaration_attributes_canonicalized: if source_is_ptx_native { + template.facts.declaration_attributes_canonicalized + } else { + template + .facts + .declaration_attributes_canonicalized + .or(defaults.declaration_attributes_canonicalized) + }, + runtime_validation: template + .facts + .runtime_validation + .or(defaults.runtime_validation), + expected_ptx, + }) +} + +pub(super) fn validate_expanded_matrix_identity( + id: &str, + source: Option<&IntrinsicSource>, + source_record: Option<&str>, + llvm_symbol: Option<&str>, +) -> Result<()> { + ensure!(!id.is_empty(), "expanded evidence has an empty ID"); + let imported_source = match (source, source_record) { + (Some(IntrinsicSource::LlvmImported { source_record }), None) => { + ensure!( + !source_record.is_empty(), + "expanded evidence {id} has an empty source record" + ); + true + } + (Some(IntrinsicSource::PtxNative { instruction }), None) => { + ensure!( + !instruction.is_empty(), + "expanded evidence {id} has an empty PTX source" + ); + false + } + (None, Some(source_record)) => { + ensure!( + !source_record.is_empty(), + "expanded evidence {id} has an empty source record" + ); + true + } + _ => unreachable!("matrix source shape was validated before expansion"), + }; + if imported_source { + ensure!( + llvm_symbol.is_some_and(|symbol| !symbol.is_empty()), + "expanded imported evidence {id} has no LLVM symbol" + ); + } else { + ensure!( + llvm_symbol.is_none(), + "expanded PTX-native evidence {id} invents an LLVM symbol" + ); + } + Ok(()) +} + +pub(super) fn select_fact(specific: &Option, default: &Option) -> Option { + specific.clone().or_else(|| default.clone()) +} + +pub(super) fn required_fact(value: Option, id: &str, field: &str) -> Result { + let value = value.with_context(|| format!("expanded evidence {id} has no {field}"))?; + ensure!( + !value.trim().is_empty(), + "expanded evidence {id} has an empty {field}" + ); + Ok(value) +} + +pub(super) fn expand_evidence_source( + source: &IntrinsicSource, + bindings: &BTreeMap, + used_axes: &mut BTreeSet, +) -> Result { + Ok(match source { + IntrinsicSource::LlvmImported { source_record } => IntrinsicSource::LlvmImported { + source_record: expand_axis_placeholders( + source_record, + bindings, + used_axes, + "tagged source record", + )?, + }, + IntrinsicSource::PtxNative { instruction } => IntrinsicSource::PtxNative { + instruction: expand_axis_placeholders( + instruction, + bindings, + used_axes, + "PTX-native source", + )?, + }, + }) +} + +pub(super) fn expand_axis_placeholders( + value: &str, + bindings: &BTreeMap, + used_axes: &mut BTreeSet, + field: &str, +) -> Result { + let mut output = String::with_capacity(value.len()); + let mut rest = value; + while let Some(position) = rest.find('$') { + output.push_str(&rest[..position]); + let placeholder = &rest[position..]; + ensure!( + placeholder.starts_with("${"), + "{field} contains malformed matrix placeholder {value:?}" + ); + let close = placeholder + .find('}') + .with_context(|| format!("{field} contains an unterminated matrix placeholder"))?; + let axis = &placeholder[2..close]; + ensure!( + is_safe_matrix_token(axis), + "{field} contains malformed matrix axis {axis:?}" + ); + let replacement = bindings + .get(axis) + .with_context(|| format!("{field} references unknown matrix axis {axis}"))?; + output.push_str(replacement); + used_axes.insert(axis.to_owned()); + rest = &placeholder[close + 1..]; + } + output.push_str(rest); + Ok(output) +} + +pub(super) fn is_safe_matrix_token(value: &str) -> bool { + let mut bytes = value.bytes(); + bytes.next().is_some_and(|byte| byte.is_ascii_lowercase()) + && bytes.all(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'_') +} + +pub(super) fn reject_default_placeholders( + defaults: &EvidenceRecordDefaults, + context: &str, + allow_resolved_symbol: bool, +) -> Result<()> { + if !allow_resolved_symbol && let Some(value) = &defaults.resolved_llvm_symbol { + reject_disallowed_placeholder(value, context)?; + } + for value in defaults + .target_triple + .iter() + .chain(defaults.gpu_target.iter()) + .chain(defaults.ptx_feature.iter()) + .chain(defaults.status.iter()) + .chain(defaults.llvm_arguments.iter().flatten()) + .chain(defaults.llvm_results.iter().flatten()) + .chain(defaults.concrete_llvm_arguments.iter().flatten()) + .chain(defaults.concrete_llvm_results.iter().flatten()) + { + reject_disallowed_placeholder(value, context)?; + } + reject_stage_placeholders(&defaults.stages, context) +} + +pub(super) fn reject_stage_placeholders(stages: &[EvidenceStage], context: &str) -> Result<()> { + for stage in stages { + for value in stage + .targets + .iter() + .chain(std::iter::once(&stage.representation)) + .chain(std::iter::once(&stage.detail)) + .chain(stage.tool_path.iter()) + .chain(stage.tool_version.iter()) + .chain(stage.tool_sha256.iter()) + { + reject_disallowed_placeholder(value, context)?; + } + } + Ok(()) +} + +pub(super) fn reject_disallowed_placeholder(value: &str, field: &str) -> Result<()> { + ensure!( + !value.contains("${"), + "{field} cannot contain matrix placeholders" + ); + Ok(()) +} + +pub(super) fn validate_stage_pairs(stages: &[EvidenceStage], id: &str) -> Result<()> { + let mut identities = Vec::new(); + for stage in stages { + let identity = (stage.stage, stage.mechanism, stage.targets.clone()); + ensure!( + !identities.contains(&identity), + "expanded evidence {id} has conflicting duplicate {:?}/{:?} stage targets {:?}", + stage.stage, + stage.mechanism, + stage.targets + ); + identities.push(identity); + } + Ok(()) +} + +pub(super) fn read_evidence(repo_root: &Path) -> Result<(Vec, Vec)> { + let directory = repo_root.join("intrinsics/evidence"); + let mut paths: Vec = fs::read_dir(&directory) + .with_context(|| format!("read {}", directory.display()))? + .filter_map(|entry| entry.ok().map(|entry| entry.path())) + .filter(|path| path.extension().and_then(|ext| ext.to_str()) == Some("json")) + .collect(); + paths.sort(); + ensure!( + !paths.is_empty(), + "no evidence JSON files in {}", + directory.display() + ); + let mut files = Vec::with_capacity(paths.len()); + let mut hashes = Vec::with_capacity(paths.len()); + for path in paths { + let file = read_evidence_file(&path)?; + let name = path.file_name().unwrap().to_string_lossy(); + hashes.push(format!("{name}:{}", sha256_file(&path)?)); + files.push(file); + } + Ok((files, hashes)) +} + +#[derive(Debug, Clone, Copy)] +pub(super) struct IndexedEvidence<'a> { + pub(super) file: &'a EvidenceFile, + pub(super) record: &'a EvidenceRecord, + pub(super) backend_version: &'a str, + pub(super) backend_sha256: &'a str, +} + +pub(super) fn index_evidence<'a>( + files: &'a [EvidenceFile], + llvm_revision: &str, +) -> Result>> { + let mut result = BTreeMap::new(); + for file in files { + ensure!( + !file.backend_profile.trim().is_empty() && !file.llvm_revision.trim().is_empty(), + "evidence file has no concrete backend profile or LLVM revision" + ); + ensure!( + !file.backend_version.trim().is_empty() && !file.backend_sha256.trim().is_empty(), + "evidence does not identify the backend binary" + ); + if file.backend_kind != Some(IntrinsicBackend::LibNvvm) { + ensure!( + file.llvm_revision == llvm_revision, + "selected evidence LLVM revision {} does not match pinned {}", + file.llvm_revision, + llvm_revision + ); + } + for record in &file.records { + ensure!( + result + .insert( + (file.backend_profile.as_str(), record.id.as_str()), + IndexedEvidence { + file, + record, + backend_version: &file.backend_version, + backend_sha256: &file.backend_sha256, + }, + ) + .is_none(), + "duplicate evidence for catalog ID {}", + record.id + ); + } + } + Ok(result) +} + +pub(super) fn validate_evidence( + policy: &OverlayIntrinsic, + evidence: &IndexedEvidence<'_>, + lowering: Option<&crate::model::OverlayBackendLowering>, +) -> Result<()> { + let record = evidence.record; + record.expected_ptx.validate().map_err(|reason| { + anyhow::anyhow!( + "{} evidence has an invalid expected PTX pattern: {reason}", + policy.id + ) + })?; + let policy_source = resolve_policy_source(policy)?; + let evidence_source = match (&record.source, &record.source_record) { + (None, Some(source_record)) => IntrinsicSource::LlvmImported { + source_record: source_record.clone(), + }, + (Some(source), None) => source.clone(), + (Some(_), Some(_)) => bail!( + "{} evidence mixes tagged source with legacy source_record", + policy.id + ), + (None, None) => bail!("{} evidence has no source provenance", policy.id), + }; + ensure!( + evidence_source == policy_source, + "{} evidence source provenance mismatch", + policy.id + ); + ensure!( + record.llvm_symbol == policy.llvm_symbol + && record.llvm_arguments == policy.llvm_arguments + && record.llvm_results == policy.llvm_results, + "{} evidence signature mismatch", + policy.id + ); + if matches!(policy_source, IntrinsicSource::PtxNative { .. }) { + ensure!( + record.llvm_symbol.is_none() + && record.resolved_llvm_symbol.is_none() + && record.llvm_arguments.is_empty() + && record.llvm_results.is_empty() + && record.concrete_llvm_arguments.is_empty() + && record.concrete_llvm_results.is_empty() + && record.declaration_attributes_canonicalized.is_none(), + "{} PTX-native evidence must not invent LLVM declaration facts", + policy.id + ); + } + ensure!( + record.expected_ptx == policy.expected_ptx, + "{} evidence PTX expectation mismatch", + policy.id + ); + ensure!( + matches!(record.status.as_str(), "lowered" | "validated" | "executed"), + "{} evidence status {} is too weak to generate a lowering", + policy.id, + record.status + ); + ensure!( + !record.target_triple.is_empty() + && !record.gpu_target.is_empty() + && !record.ptx_feature.is_empty(), + "{} evidence omits its full target profile", + policy.id + ); + if let Some(lowering) = lowering { + ensure!( + evidence.file.backend_kind == Some(lowering.backend), + "{} evidence profile {} has the wrong backend kind", + policy.id, + evidence.file.backend_profile + ); + match record.status.as_str() { + "executed" => ensure!( + record.runtime_validation == Some(RuntimeValidation::Executed), + "{} executed evidence must record runtime_validation = executed", + policy.id + ), + _ => ensure!( + record.runtime_validation == Some(RuntimeValidation::Unexecuted), + "{} non-executed backend evidence must record runtime_validation = unexecuted", + policy.id + ), + } + ensure!( + !record.stages.is_empty(), + "{} backend evidence omits compilation stages", + policy.id + ); + ensure!( + record.stages.iter().any(|stage| { + stage.stage == EvidenceStageKind::BackendCodegen + && stage.mechanism == Some(lowering.mechanism) + && stage.outcome == "succeeded" + }), + "{} evidence has no successful backend-codegen stage for {:?}", + policy.id, + lowering.mechanism + ); + validate_selected_stage_targets(policy, record, lowering)?; + if lowering.mechanism == BackendLoweringMechanism::TypedNvvm { + validate_typed_llvm_evidence(policy, record)?; + } + validate_packed_conversion_backend_evidence(policy, record, lowering)?; + validate_scalar_conversion_backend_evidence(policy, record, lowering)?; + validate_scalar_arithmetic_backend_evidence(policy, record, lowering)?; + validate_inline_ptx_fallback_evidence(policy, record, lowering)?; + if lowering.backend == IntrinsicBackend::LlvmNvptx + && matches!(record.status.as_str(), "validated" | "executed") + { + ensure!( + has_valid_ptx_assembly_stage(record, lowering.mechanism), + "{} validated LLVM-NVPTX evidence requires a successful PTX-assembly stage with exact tool identity", + policy.id + ); + } else if lowering.backend == IntrinsicBackend::LibNvvm + && matches!(record.status.as_str(), "validated" | "executed") + { + ensure!( + has_valid_cubin_device_link_stage(record, lowering.mechanism), + "{} validated libNVVM evidence requires a successful cubin-producing device-link stage with exact tool identity", + policy.id + ); + } + if record.status == "executed" { + ensure!( + record.stages.iter().any(|stage| { + stage.stage == EvidenceStageKind::Runtime + && stage.mechanism == Some(lowering.mechanism) + && stage.outcome == "succeeded" + }), + "{} executed evidence requires a successful runtime stage for the selected mechanism", + policy.id + ); + } + } + Ok(()) +} + +pub(super) fn validate_inline_ptx_fallback_evidence( + policy: &OverlayIntrinsic, + record: &EvidenceRecord, + lowering: &crate::model::OverlayBackendLowering, +) -> Result<()> { + if !matches!( + policy.family.as_str(), + "cluster_barrier" | "wgmma_control" | "elect" + ) || lowering.backend != IntrinsicBackend::LibNvvm + { + return Ok(()); + } + for stage in [ + EvidenceStageKind::BackendCodegen, + EvidenceStageKind::DeviceLink, + ] { + ensure!( + record.stages.iter().any(|candidate| { + candidate.stage == stage + && candidate.mechanism == Some(BackendLoweringMechanism::TypedNvvm) + && candidate.outcome == "failed" + }), + "{} libNVVM inline-PTX evidence must record the failed typed-NVVM {:?} comparison", + policy.id, + stage + ); + } + ensure!( + !record.stages.iter().any(|candidate| { + candidate.stage == EvidenceStageKind::DeviceLink + && candidate.mechanism == Some(BackendLoweringMechanism::TypedNvvm) + && candidate.outcome == "succeeded" + }), + "{} libNVVM evidence cannot select inline PTX after a successful typed-NVVM terminal", + policy.id + ); + Ok(()) +} + +pub(super) fn validate_packed_conversion_backend_evidence( + policy: &OverlayIntrinsic, + record: &EvidenceRecord, + lowering: &crate::model::OverlayBackendLowering, +) -> Result<()> { + if policy.family != "packed_conversion" { + return Ok(()); + } + match lowering.backend { + IntrinsicBackend::LlvmNvptx => { + validate_typed_llvm_evidence(policy, record)?; + for stage in [ + EvidenceStageKind::DeclarationCanonicalization, + EvidenceStageKind::BackendCodegen, + ] { + ensure!( + successful_stage(record, BackendLoweringMechanism::TypedNvvm, stage).is_some(), + "{} LLVM packed-conversion evidence requires a successful auxiliary typed-NVVM {:?} stage", + policy.id, + stage + ); + } + ensure!( + has_valid_ptx_assembly_stage(record, BackendLoweringMechanism::TypedNvvm), + "{} LLVM packed-conversion evidence requires a successful auxiliary typed-NVVM PTX-assembly stage with exact tool identity", + policy.id + ); + ensure!( + matches!(record.status.as_str(), "validated" | "executed"), + "{} LLVM packed-conversion evidence requires validated evidence status for its auxiliary typed-NVVM terminal stage", + policy.id + ); + let typed_lowering = crate::model::OverlayBackendLowering { + mechanism: BackendLoweringMechanism::TypedNvvm, + ..lowering.clone() + }; + validate_selected_stage_targets(policy, record, &typed_lowering)?; + Ok(()) + } + IntrinsicBackend::LibNvvm => { + ensure!( + record.resolved_llvm_symbol.is_none() + && record.concrete_llvm_arguments.is_empty() + && record.concrete_llvm_results.is_empty() + && record.declaration_attributes_canonicalized.is_none() + && !record.stages.iter().any(|stage| { + stage.mechanism == Some(BackendLoweringMechanism::TypedNvvm) + }), + "{} libNVVM inline-PTX evidence must not claim typed LLVM support", + policy.id + ); + Ok(()) + } + } +} + +pub(super) fn validate_scalar_conversion_backend_evidence( + policy: &OverlayIntrinsic, + record: &EvidenceRecord, + lowering: &crate::model::OverlayBackendLowering, +) -> Result<()> { + if policy.family != "scalar_conversion" { + return Ok(()); + } + match lowering.backend { + IntrinsicBackend::LlvmNvptx => ensure!( + successful_stage( + record, + BackendLoweringMechanism::TypedNvvm, + EvidenceStageKind::DeclarationCanonicalization, + ) + .is_some(), + "{} LLVM scalar-conversion evidence must canonicalize the typed declaration", + policy.id + ), + IntrinsicBackend::LibNvvm => ensure!( + record.resolved_llvm_symbol.is_none() + && record.concrete_llvm_arguments.is_empty() + && record.concrete_llvm_results.is_empty() + && record.declaration_attributes_canonicalized.is_none(), + "{} libNVVM scalar-conversion evidence must describe the selected inline-PTX route", + policy.id + ), + }; + Ok(()) +} + +pub(super) fn validate_scalar_arithmetic_backend_evidence( + policy: &OverlayIntrinsic, + record: &EvidenceRecord, + lowering: &crate::model::OverlayBackendLowering, +) -> Result<()> { + if policy.family != "scalar_arithmetic" { + return Ok(()); + } + match (lowering.backend, lowering.mechanism) { + (IntrinsicBackend::LlvmNvptx, BackendLoweringMechanism::TypedNvvm) => ensure!( + successful_stage( + record, + BackendLoweringMechanism::TypedNvvm, + EvidenceStageKind::DeclarationCanonicalization, + ) + .is_some(), + "{} typed LLVM scalar-arithmetic evidence must canonicalize the declaration", + policy.id + ), + (IntrinsicBackend::LlvmNvptx, BackendLoweringMechanism::InlinePtx) => { + validate_typed_llvm_evidence(policy, record)?; + ensure!( + successful_stage( + record, + BackendLoweringMechanism::TypedNvvm, + EvidenceStageKind::DeclarationCanonicalization, + ) + .is_some(), + "{} inline LLVM scalar-arithmetic evidence must canonicalize its imported declaration", + policy.id + ); + } + (IntrinsicBackend::LibNvvm, _) => ensure!( + record.resolved_llvm_symbol.is_none() + && record.concrete_llvm_arguments.is_empty() + && record.concrete_llvm_results.is_empty() + && record.declaration_attributes_canonicalized.is_none(), + "{} libNVVM scalar-arithmetic evidence must describe the selected inline-PTX route", + policy.id + ), + }; + Ok(()) +} + +pub(super) fn validate_typed_llvm_evidence( + policy: &OverlayIntrinsic, + record: &EvidenceRecord, +) -> Result<()> { + let concrete_arguments = policy + .llvm_arguments + .iter() + .map(|argument| { + match argument.as_str() { + "shared_cluster_ptr" => return Ok("ptr addrspace(7)".into()), + "shared_ptr" => return Ok("ptr addrspace(3)".into()), + "global_ptr" => return Ok("ptr addrspace(1)".into()), + "ptr" => return Ok("ptr".into()), + "anyptr" => {} + _ => return Ok(argument.clone()), + } + match policy.selected_address_space.with_context(|| { + format!( + "{} has a polymorphic LLVM pointer without a selected address space", + policy.id + ) + })? { + ImportedAddressSpace::Generic => Ok("ptr".into()), + ImportedAddressSpace::Shared => Ok("ptr addrspace(3)".into()), + } + }) + .collect::>>()?; + ensure!( + record.resolved_llvm_symbol == policy.resolved_llvm_symbol + && record.concrete_llvm_arguments == concrete_arguments + && record.concrete_llvm_results == policy.llvm_results + && record.declaration_attributes_canonicalized == Some(true), + "{} typed LLVM evidence does not prove its resolved signature and canonical declaration attributes", + policy.id + ); + Ok(()) +} + +pub(super) fn validate_selected_stage_targets( + policy: &OverlayIntrinsic, + record: &EvidenceRecord, + lowering: &crate::model::OverlayBackendLowering, +) -> Result<()> { + let requirement = backend_target_requirement(policy, lowering)?; + let blackwell_ldmatrix = policy.family == "ldmatrix" + && policy + .ldmatrix_variant + .as_ref() + .is_some_and(|variant| variant.shape != LdmatrixShape::M8n8); + let f8f6f4_mma = is_f8f6f4_mma_target_matrix_policy(policy); + let paired_alternatives = match &requirement.hardware { + CatalogHardwareTarget::TargetMatrix { contracts } => Some( + contracts + .iter() + .flat_map(|contract| contract.alternatives.iter().cloned()) + .collect::>(), + ), + _ => None, + }; + let target_matrix_evidence = blackwell_ldmatrix || f8f6f4_mma || paired_alternatives.is_some(); + if !target_matrix_evidence { + let mut pairs = Vec::new(); + for stage in record.stages.iter().filter(|stage| { + stage.mechanism == Some(lowering.mechanism) && stage.outcome == "succeeded" + }) { + let pair = (stage.stage, stage.mechanism); + ensure!( + !pairs.contains(&pair), + "{} has multiple target-specific {:?}/{:?} stages outside reviewed target-matrix evidence", + policy.id, + stage.stage, + stage.mechanism + ); + pairs.push(pair); + } + } + for stage in &record.stages { + ensure!( + !stage.targets.is_empty(), + "{} evidence stage {:?} has no targets", + policy.id, + stage.stage + ); + for target in &stage.targets { + ensure!( + is_normalized_stage_target(target), + "{} evidence stage has unsupported target spelling {target:?}", + policy.id + ); + } + } + + let terminal_kind = match lowering.backend { + IntrinsicBackend::LlvmNvptx => EvidenceStageKind::PtxAssembly, + IntrinsicBackend::LibNvvm => EvidenceStageKind::DeviceLink, + }; + let expected_ptx = requirement.minimum_ptx.encoded(); + let expected_hardware = match &requirement.hardware { + CatalogHardwareTarget::AnyOf { alternatives } if !alternatives.is_empty() => { + alternatives.clone() + } + CatalogHardwareTarget::TargetMatrix { contracts } if !contracts.is_empty() => { + let mut alternatives = contracts + .iter() + .flat_map(|contract| contract.alternatives.iter()) + .map(|alternative| alternative.hardware) + .collect::>(); + alternatives.sort_by_key(|hardware| target_hardware_sort_key(*hardware)); + alternatives.dedup(); + alternatives + } + _ => bail!( + "{} selected backend stages require a hardware target", + policy.id + ), + }; + if target_matrix_evidence { + return validate_target_matrix_stage_targets( + policy, + record, + lowering, + terminal_kind, + &expected_hardware, + expected_ptx, + paired_alternatives.as_deref(), + ); + } + let backend = successful_stage( + record, + lowering.mechanism, + EvidenceStageKind::BackendCodegen, + ) + .with_context(|| { + format!( + "{} has no successful selected backend-codegen stage", + policy.id + ) + })?; + let mut required_stages = vec![backend]; + if matches!(record.status.as_str(), "validated" | "executed") { + required_stages.push( + successful_stage(record, lowering.mechanism, terminal_kind).with_context(|| { + format!("{} has no successful selected terminal stage", policy.id) + })?, + ); + } + for stage in required_stages { + let (hardware, ptx) = selected_stage_floor(stage)?; + let allow_forward_minimum = stage.stage != EvidenceStageKind::BackendCodegen; + let hardware_matches = expected_hardware.iter().any(|expected| { + selected_stage_hardware_matches(hardware, *expected, allow_forward_minimum) + }); + let ptx_matches = match lowering.backend { + IntrinsicBackend::LlvmNvptx => ptx == expected_ptx, + IntrinsicBackend::LibNvvm => ptx >= expected_ptx, + }; + ensure!( + hardware_matches && ptx_matches, + "{} evidence stage {:?} targets {} / PTX {}.{} instead of a compatible target at catalog floor {} / PTX {}.{}", + policy.id, + stage.stage, + describe_stage_hardware(hardware), + ptx / 10, + ptx % 10, + expected_hardware + .iter() + .map(|hardware| describe_stage_hardware(*hardware)) + .collect::>() + .join(" or "), + expected_ptx / 10, + expected_ptx % 10 + ); + } + if record.status == "executed" { + let runtime = successful_stage(record, lowering.mechanism, EvidenceStageKind::Runtime) + .with_context(|| { + format!( + "{} executed evidence has no successful runtime stage", + policy.id + ) + })?; + let (hardware, ptx) = selected_stage_floor(runtime)?; + let ptx_matches = match lowering.backend { + IntrinsicBackend::LlvmNvptx => ptx == expected_ptx, + IntrinsicBackend::LibNvvm => ptx >= expected_ptx, + }; + ensure!( + expected_hardware + .iter() + .any(|expected| selected_stage_hardware_matches(hardware, *expected, true)) + && ptx_matches, + "{} runtime stage target does not satisfy its catalog floor", + policy.id + ); + } + Ok(()) +} + +pub(super) fn validate_target_matrix_stage_targets( + policy: &OverlayIntrinsic, + record: &EvidenceRecord, + lowering: &OverlayBackendLowering, + terminal_kind: EvidenceStageKind, + expected_hardware: &[CatalogHardwareAlternative], + minimum_ptx: u16, + paired_alternatives: Option<&[CatalogTargetAlternative]>, +) -> Result<()> { + let mut required_kinds = vec![EvidenceStageKind::BackendCodegen]; + if matches!(record.status.as_str(), "validated" | "executed") { + if lowering.backend == IntrinsicBackend::LibNvvm { + required_kinds.push(EvidenceStageKind::PtxAssembly); + } + required_kinds.push(terminal_kind); + } + + for kind in required_kinds { + let stages = record + .stages + .iter() + .filter(|stage| { + stage.stage == kind + && stage.mechanism == Some(lowering.mechanism) + && stage.outcome == "succeeded" + }) + .collect::>(); + ensure!( + stages.len() == expected_hardware.len(), + "{} target-matrix {:?} evidence must contain one structured stage for each of its {} reviewed targets", + policy.id, + kind, + expected_hardware.len() + ); + + let mut covered = Vec::new(); + for stage in stages { + let (hardware, ptx) = selected_stage_floor(stage)?; + ensure!( + expected_hardware.contains(&hardware) && !covered.contains(&hardware), + "{} target-matrix {:?} evidence has an unexpected or duplicate target {}", + policy.id, + kind, + describe_stage_hardware(hardware) + ); + covered.push(hardware); + let exact_floor = + target_matrix_ptx_floor(policy, hardware, paired_alternatives, false)?; + let ptx_matches = match lowering.backend { + IntrinsicBackend::LlvmNvptx => ptx == exact_floor, + IntrinsicBackend::LibNvvm => { + ptx >= paired_alternatives.map_or(minimum_ptx, |_| exact_floor) + } + }; + ensure!( + ptx_matches, + "{} target-matrix {:?} evidence records the wrong PTX floor for {}", + policy.id, + kind, + describe_stage_hardware(hardware) + ); + if matches!( + kind, + EvidenceStageKind::PtxAssembly | EvidenceStageKind::DeviceLink + ) { + ensure!( + stage + .tool_path + .as_deref() + .is_some_and(|value| !value.is_empty()) + && stage + .tool_version + .as_deref() + .is_some_and(|value| !value.is_empty()) + && stage + .tool_sha256 + .as_deref() + .is_some_and(|value| !value.is_empty()), + "{} target-matrix {:?} evidence for {} does not identify its tool", + policy.id, + kind, + describe_stage_hardware(hardware) + ); + } + if kind == EvidenceStageKind::DeviceLink { + ensure!( + stage.artifact_kind == Some(EvidenceArtifactKind::Cubin), + "{} target-matrix device-link evidence for {} is not a cubin", + policy.id, + describe_stage_hardware(hardware) + ); + } + } + } + + if record.status == "executed" { + let runtime = record + .stages + .iter() + .filter(|stage| { + stage.stage == EvidenceStageKind::Runtime + && stage.mechanism == Some(lowering.mechanism) + && stage.outcome == "succeeded" + }) + .collect::>(); + ensure!( + runtime.len() == 1, + "{} executed target-matrix evidence requires one successful runtime stage", + policy.id + ); + let (hardware, ptx) = selected_stage_floor(runtime[0])?; + ensure!( + expected_hardware + .iter() + .any(|expected| { selected_stage_hardware_matches(hardware, *expected, true) }), + "{} runtime target {} is outside its target matrix", + policy.id, + describe_stage_hardware(hardware) + ); + let floor = target_matrix_ptx_floor(policy, hardware, paired_alternatives, true)?; + let ptx_matches = match lowering.backend { + IntrinsicBackend::LlvmNvptx => ptx == floor, + IntrinsicBackend::LibNvvm => ptx >= floor, + }; + ensure!( + ptx_matches, + "{} runtime target {} records PTX {}.{} instead of its paired floor {}.{}", + policy.id, + describe_stage_hardware(hardware), + ptx / 10, + ptx % 10, + floor / 10, + floor % 10 + ); + } + Ok(()) +} + +pub(super) fn target_matrix_ptx_floor( + policy: &OverlayIntrinsic, + hardware: CatalogHardwareAlternative, + paired_alternatives: Option<&[CatalogTargetAlternative]>, + allow_forward_minimum: bool, +) -> Result { + if let Some(alternatives) = paired_alternatives { + return alternatives + .iter() + .filter(|alternative| { + selected_stage_hardware_matches( + hardware, + alternative.hardware, + allow_forward_minimum, + ) + }) + .map(|alternative| alternative.minimum_ptx.encoded()) + .min() + .with_context(|| { + format!( + "{} target matrix has no PTX floor for {}", + policy.id, + describe_stage_hardware(hardware) + ) + }); + } + if is_f8f6f4_mma_target_matrix_policy(policy) { + f8f6f4_llvm_ptx_floor(hardware) + } else { + blackwell_ldmatrix_llvm_ptx_floor(hardware) + } +} + +pub(super) fn f8f6f4_llvm_ptx_floor(hardware: CatalogHardwareAlternative) -> Result { + match hardware { + CatalogHardwareAlternative::ExactArchitecture { sm: 120 } + | CatalogHardwareAlternative::FamilyTarget { sm: 120 } + | CatalogHardwareAlternative::ExactArchitecture { sm: 121 } + | CatalogHardwareAlternative::FamilyTarget { sm: 121 } => { + recorded_stage_ptx_floor(hardware, 87) + } + _ => bail!( + "{} is not a reviewed f8f6f4 MMA target", + describe_stage_hardware(hardware) + ), + } +} + +pub(super) fn blackwell_ldmatrix_llvm_ptx_floor( + hardware: CatalogHardwareAlternative, +) -> Result { + match hardware { + CatalogHardwareAlternative::ExactArchitecture { + sm: 100 | 103 | 110 | 120 | 121, + } + | CatalogHardwareAlternative::FamilyTarget { + sm: 100 | 103 | 120 | 121, + } + | CatalogHardwareAlternative::FamilyTarget { sm: 110 } => { + recorded_stage_ptx_floor(hardware, 86) + } + _ => bail!( + "{} is not a reviewed Blackwell ldmatrix target", + describe_stage_hardware(hardware) + ), + } +} + +pub(super) fn recorded_stage_ptx_floor( + hardware: CatalogHardwareAlternative, + instruction_floor: u16, +) -> Result { + let (sm, suffix) = match hardware { + CatalogHardwareAlternative::MinimumSm { sm } => (sm, None), + CatalogHardwareAlternative::ExactArchitecture { sm } => (sm, Some('a')), + CatalogHardwareAlternative::FamilyTarget { sm } => (sm, Some('f')), + }; + let arch = CudaArch::new(u32::from(sm), suffix)?; + Ok(instruction_floor.max(recorded_ptx_floor(&arch)?)) +} + +pub(super) fn successful_stage( + record: &EvidenceRecord, + mechanism: BackendLoweringMechanism, + kind: EvidenceStageKind, +) -> Option<&crate::model::EvidenceStage> { + record.stages.iter().find(|stage| { + stage.stage == kind && stage.mechanism == Some(mechanism) && stage.outcome == "succeeded" + }) +} + +pub(super) fn is_normalized_stage_target(target: &str) -> bool { + if let Some(value) = target.strip_prefix("ptx") { + return value.len() == 2 && value.bytes().all(|byte| byte.is_ascii_digit()); + } + parse_stage_hardware(target).is_some() +} + +pub(super) fn parse_stage_hardware(target: &str) -> Option { + let arch = target.parse::().ok()?; + let canonical = if target.starts_with("sm_") { + arch.sm() + } else if target.starts_with("compute_") { + arch.compute() + } else { + return None; + }; + if canonical != target || !matches!(arch.capability().to_string().len(), 2 | 3) { + return None; + } + let sm = u16::try_from(arch.capability()).ok()?; + Some(match arch.suffix() { + None => CatalogHardwareAlternative::MinimumSm { sm }, + Some('a') => CatalogHardwareAlternative::ExactArchitecture { sm }, + Some('f') => CatalogHardwareAlternative::FamilyTarget { sm }, + _ => unreachable!(), + }) +} + +pub(super) fn selected_stage_hardware_matches( + actual: CatalogHardwareAlternative, + expected: CatalogHardwareAlternative, + allow_forward_minimum: bool, +) -> bool { + match expected { + CatalogHardwareAlternative::MinimumSm { sm: expected } => { + if allow_forward_minimum { + match actual { + CatalogHardwareAlternative::MinimumSm { sm } + | CatalogHardwareAlternative::ExactArchitecture { sm } + | CatalogHardwareAlternative::FamilyTarget { sm } => sm >= expected, + } + } else { + actual == CatalogHardwareAlternative::MinimumSm { sm: expected } + } + } + CatalogHardwareAlternative::ExactArchitecture { .. } + | CatalogHardwareAlternative::FamilyTarget { .. } => actual == expected, + } +} + +pub(super) fn describe_stage_hardware(hardware: CatalogHardwareAlternative) -> String { + match hardware { + CatalogHardwareAlternative::MinimumSm { sm } => format!("sm_{sm}"), + CatalogHardwareAlternative::ExactArchitecture { sm } => format!("sm_{sm}a"), + CatalogHardwareAlternative::FamilyTarget { sm } => format!("sm_{sm}f"), + } +} + +pub(super) fn selected_stage_floor( + stage: &crate::model::EvidenceStage, +) -> Result<(CatalogHardwareAlternative, u16)> { + let mut hardware = None; + let mut ptx = None; + for target in &stage.targets { + if let Some(value) = target.strip_prefix("ptx") { + let value = value.parse::()?; + ensure!( + ptx.replace(value).is_none(), + "stage has duplicate PTX targets" + ); + } else { + let value = parse_stage_hardware(target) + .with_context(|| format!("stage has unsupported target spelling {target:?}"))?; + ensure!( + hardware.replace(value).is_none(), + "stage has duplicate architecture targets" + ); + } + } + Ok(( + hardware.context("selected stage has no architecture target")?, + ptx.context("selected stage has no PTX target")?, + )) +} + +pub(super) fn has_valid_ptx_assembly_stage( + record: &EvidenceRecord, + mechanism: BackendLoweringMechanism, +) -> bool { + has_valid_tool_stage(record, mechanism, EvidenceStageKind::PtxAssembly) +} + +pub(super) fn has_valid_cubin_device_link_stage( + record: &EvidenceRecord, + mechanism: BackendLoweringMechanism, +) -> bool { + has_valid_tool_stage(record, mechanism, EvidenceStageKind::DeviceLink) + && successful_stage(record, mechanism, EvidenceStageKind::DeviceLink) + .is_some_and(|stage| stage.artifact_kind == Some(EvidenceArtifactKind::Cubin)) +} + +pub(super) fn has_valid_tool_stage( + record: &EvidenceRecord, + mechanism: BackendLoweringMechanism, + stage_kind: EvidenceStageKind, +) -> bool { + record.stages.iter().any(|stage| { + stage.stage == stage_kind + && stage.mechanism == Some(mechanism) + && stage.outcome == "succeeded" + && stage + .tool_path + .as_deref() + .is_some_and(|value| !value.is_empty()) + && stage + .tool_version + .as_deref() + .is_some_and(|value| !value.is_empty()) + && stage.tool_sha256.as_deref().is_some_and(|value| { + value.len() == 64 && value.bytes().all(|byte| byte.is_ascii_hexdigit()) + }) + }) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/clc.rs b/crates/cuda-intrinsics-gen/src/resolve/families/clc.rs new file mode 100644 index 0000000000..146af8f710 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/clc.rs @@ -0,0 +1,472 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, Clc, ClcAdapter, ClcAdmission, ClcOperation, ImportedIntrinsic, + IntrinsicBackend, OverlayBackendLowering, OverlayIntrinsic, RuntimeValidation, +}; +use crate::ptx::{InstructionPattern, OperandPattern}; +use anyhow::{Context, Result, ensure}; + +use crate::resolve::guards::*; + +#[derive(Clone, Copy)] +pub(in crate::resolve) struct ClcRecipe { + pub(in crate::resolve) operation: ClcOperation, + pub(in crate::resolve) abi_id: &'static str, + pub(in crate::resolve) id: &'static str, + pub(in crate::resolve) operation_key: &'static str, + pub(in crate::resolve) source_record: &'static str, + pub(in crate::resolve) llvm_symbol: &'static str, + pub(in crate::resolve) rust_arguments: &'static [&'static str], + pub(in crate::resolve) llvm_arguments: &'static [&'static str], + pub(in crate::resolve) llvm_results: &'static [&'static str], + pub(in crate::resolve) dialect_op_type: &'static str, + pub(in crate::resolve) dialect_op_name: &'static str, + pub(in crate::resolve) dialect_operands: &'static [&'static str], + pub(in crate::resolve) dialect_results: &'static [&'static str], + pub(in crate::resolve) adapter: ClcAdapter, + pub(in crate::resolve) modifiers: &'static [&'static str], + pub(in crate::resolve) operands: &'static [OperandPattern], + pub(in crate::resolve) targets: &'static str, + pub(in crate::resolve) minimum_sm: Option<&'static str>, + pub(in crate::resolve) pure: bool, + pub(in crate::resolve) memory: &'static str, + pub(in crate::resolve) execution_scope: &'static str, + pub(in crate::resolve) summary: &'static str, +} + +pub(in crate::resolve) fn clc_recipe(operation: ClcOperation) -> ClcRecipe { + const TRY_OPERANDS: &[OperandPattern] = &[OperandPattern::Address, OperandPattern::Address]; + const QUERY_OPERANDS: &[OperandPattern] = &[OperandPattern::Register, OperandPattern::Register]; + let (abi_id, id, operation_key, source_record, llvm_symbol, op_type, op_name) = match operation + { + ClcOperation::TryCancel => ( + "i0322", + "clc_try_cancel", + "cluster.launch_control.try_cancel", + "int_nvvm_clusterlaunchcontrol_try_cancel_async_shared", + "llvm.nvvm.clusterlaunchcontrol.try_cancel.async.shared", + "ClcTryCancelOp", + "nvvm.clc_try_cancel", + ), + ClcOperation::TryCancelMulticast => ( + "i0323", + "clc_try_cancel_multicast", + "cluster.launch_control.try_cancel.multicast", + "int_nvvm_clusterlaunchcontrol_try_cancel_async_multicast_shared", + "llvm.nvvm.clusterlaunchcontrol.try_cancel.async.multicast.shared", + "ClcTryCancelMulticastOp", + "nvvm.clc_try_cancel_multicast", + ), + ClcOperation::QueryIsCanceled => ( + "i0324", + "clc_query_is_canceled", + "cluster.launch_control.query.is_canceled", + "int_nvvm_clusterlaunchcontrol_query_cancel_is_canceled", + "llvm.nvvm.clusterlaunchcontrol.query_cancel.is_canceled", + "ClcQueryIsCanceledOp", + "nvvm.clc_query_is_canceled", + ), + ClcOperation::QueryGetFirstCtaidX => ( + "i0325", + "clc_query_get_first_ctaid_x", + "cluster.launch_control.query.first_ctaid.x", + "int_nvvm_clusterlaunchcontrol_query_cancel_get_first_ctaid_x", + "llvm.nvvm.clusterlaunchcontrol.query_cancel.get_first_ctaid.x", + "ClcQueryGetFirstCtaidXOp", + "nvvm.clc_query_get_first_ctaid_x", + ), + ClcOperation::QueryGetFirstCtaidY => ( + "i0326", + "clc_query_get_first_ctaid_y", + "cluster.launch_control.query.first_ctaid.y", + "int_nvvm_clusterlaunchcontrol_query_cancel_get_first_ctaid_y", + "llvm.nvvm.clusterlaunchcontrol.query_cancel.get_first_ctaid.y", + "ClcQueryGetFirstCtaidYOp", + "nvvm.clc_query_get_first_ctaid_y", + ), + ClcOperation::QueryGetFirstCtaidZ => ( + "i0327", + "clc_query_get_first_ctaid_z", + "cluster.launch_control.query.first_ctaid.z", + "int_nvvm_clusterlaunchcontrol_query_cancel_get_first_ctaid_z", + "llvm.nvvm.clusterlaunchcontrol.query_cancel.get_first_ctaid.z", + "ClcQueryGetFirstCtaidZOp", + "nvvm.clc_query_get_first_ctaid_z", + ), + }; + match operation { + ClcOperation::TryCancel | ClcOperation::TryCancelMulticast => ClcRecipe { + operation, + abi_id, + id, + operation_key, + source_record, + llvm_symbol, + rust_arguments: &["*mut u8", "*mut u64"], + llvm_arguments: &["shared_ptr", "shared_ptr"], + llvm_results: &[], + dialect_op_type: op_type, + dialect_op_name: op_name, + dialect_operands: &["ptr", "ptr"], + dialect_results: &[], + adapter: ClcAdapter::GenericPointersToShared, + modifiers: if operation == ClcOperation::TryCancel { + &[ + "try_cancel", + "async", + "shared::cta", + "mbarrier::complete_tx::bytes", + "b128", + ] + } else { + &[ + "try_cancel", + "async", + "shared::cta", + "mbarrier::complete_tx::bytes", + "multicast::cluster::all", + "b128", + ] + }, + operands: TRY_OPERANDS, + targets: if operation == ClcOperation::TryCancel { + "all" + } else { + // LLVM 22 exposes 101a; CUDA 13.3 exposes the other toolkit names. + "sm_100a|sm_101a|sm_103a|sm_110a|sm_120a|sm_121a" + }, + minimum_sm: if operation == ClcOperation::TryCancel { + Some("sm_100") + } else { + None + }, + pure: false, + memory: "read_write", + execution_scope: "cta", + summary: if operation == ClcOperation::TryCancel { + "Requests one pending CTA and writes its response to shared memory." + } else { + "Requests one pending CTA and multicasts its response across the cluster." + }, + }, + ClcOperation::QueryIsCanceled + | ClcOperation::QueryGetFirstCtaidX + | ClcOperation::QueryGetFirstCtaidY + | ClcOperation::QueryGetFirstCtaidZ => { + let (adapter, modifiers, summary) = match operation { + ClcOperation::QueryIsCanceled => ( + ClcAdapter::PairU64ToI128BoolToU32, + &["query_cancel", "is_canceled", "pred", "b128"] as &[_], + "Returns whether the Cluster Launch Control request was canceled.", + ), + ClcOperation::QueryGetFirstCtaidX => ( + ClcAdapter::PairU64ToI128U32, + &["query_cancel", "get_first_ctaid::x", "b32", "b128"] as &[_], + "Returns the X coordinate from a successful cancellation response.", + ), + ClcOperation::QueryGetFirstCtaidY => ( + ClcAdapter::PairU64ToI128U32, + &["query_cancel", "get_first_ctaid::y", "b32", "b128"] as &[_], + "Returns the Y coordinate from a successful cancellation response.", + ), + ClcOperation::QueryGetFirstCtaidZ => ( + ClcAdapter::PairU64ToI128U32, + &["query_cancel", "get_first_ctaid::z", "b32", "b128"] as &[_], + "Returns the Z coordinate from a successful cancellation response.", + ), + _ => unreachable!(), + }; + ClcRecipe { + operation, + abi_id, + id, + operation_key, + source_record, + llvm_symbol, + rust_arguments: &["u64", "u64"], + llvm_arguments: &["i128"], + llvm_results: if operation == ClcOperation::QueryIsCanceled { + &["i1"] + } else { + &["i32"] + }, + dialect_op_type: op_type, + dialect_op_name: op_name, + dialect_operands: &["i64", "i64"], + dialect_results: &["i32"], + adapter, + modifiers, + operands: QUERY_OPERANDS, + targets: "all", + minimum_sm: Some("sm_100"), + pure: true, + memory: "none", + execution_scope: "thread", + summary, + } + } + } +} + +pub(in crate::resolve) fn expand_clc_admission( + admission: &ClcAdmission, +) -> Result> { + ensure!( + admission.runtime_validation == RuntimeValidation::Unexecuted, + "CLC runtime validation may be marked executed only with GPU evidence" + ); + ensure!( + !admission.llvm_evidence_profile.trim().is_empty() + && !admission.libnvvm_evidence_profile.trim().is_empty(), + "compact CLC admission requires both backend evidence profiles" + ); + let expected = [ + ClcOperation::TryCancel, + ClcOperation::TryCancelMulticast, + ClcOperation::QueryIsCanceled, + ClcOperation::QueryGetFirstCtaidX, + ClcOperation::QueryGetFirstCtaidY, + ClcOperation::QueryGetFirstCtaidZ, + ]; + ensure!( + admission + .variants + .iter() + .map(|variant| variant.operation) + .eq(expected), + "compact CLC admission must list all six operations in canonical order" + ); + + admission + .variants + .iter() + .map(|variant| { + let recipe = clc_recipe(variant.operation); + ensure!( + variant.abi_id == recipe.abi_id, + "{} must keep reserved ABI ID {}", + recipe.id, + recipe.abi_id + ); + let query = recipe.pure; + Ok(OverlayIntrinsic { + id: recipe.id.into(), + abi_id: variant.abi_id.clone(), + operation_key: recipe.operation_key.into(), + family: "clc".into(), + source: None, + source_record: Some(recipe.source_record.into()), + rust_module: "clc".into(), + rust_name: recipe.id.into(), + rust_arguments: recipe.rust_arguments.iter().map(|value| (*value).into()).collect(), + rust_result: if query { "u32".into() } else { "()".into() }, + safe: false, + must_use: false, + safe_allowlist_reason: None, + public_rust_path: format!("cuda_intrinsics::clc::{}", recipe.id), + compatibility_rust_paths: vec![format!("cuda_device::clc::{}", recipe.id)], + dialect_op_type: recipe.dialect_op_type.into(), + dialect_op_name: recipe.dialect_op_name.into(), + dialect_operands: recipe.dialect_operands.iter().map(|value| (*value).into()).collect(), + dialect_results: recipe.dialect_results.iter().map(|value| (*value).into()).collect(), + llvm_symbol: Some(recipe.llvm_symbol.into()), + resolved_llvm_symbol: None, + llvm_arguments: recipe.llvm_arguments.iter().map(|value| (*value).into()).collect(), + llvm_results: recipe.llvm_results.iter().map(|value| (*value).into()).collect(), + pure: recipe.pure, + memory: recipe.memory.into(), + convergent: false, + execution_scope: recipe.execution_scope.into(), + minimum_ptx: "8.6".into(), + minimum_sm: recipe.minimum_sm.map(Into::into), + ptx_result: if query { "u32".into() } else { "()".into() }, + targets: recipe.targets.into(), + ptx_isa_version: "9.3".into(), + ptx_isa_section: "9.7.14.18-19 Cluster Launch Control".into(), + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-clusterlaunchcontrol-try-cancel".into(), + lowering: "generated_clc".into(), + backend_lowerings: vec![ + OverlayBackendLowering { + backend: IntrinsicBackend::LlvmNvptx, + mechanism: BackendLoweringMechanism::TypedNvvm, + evidence_profile: admission.llvm_evidence_profile.clone(), + targets: None, + minimum_ptx: Some("8.6".into()), + minimum_sm: recipe.minimum_sm.map(Into::into), + }, + OverlayBackendLowering { + backend: IntrinsicBackend::LibNvvm, + mechanism: BackendLoweringMechanism::TypedNvvm, + evidence_profile: admission.libnvvm_evidence_profile.clone(), + targets: None, + minimum_ptx: Some("8.6".into()), + minimum_sm: recipe.minimum_sm.map(Into::into), + }, + ], + packed_atomic: None, + redux: None, + vote: None, + active_mask: None, + warp_match: None, + warp_barrier: None, + warp_shuffle: None, + dot_product: None, + packed_alu: None, + integer_minmax: None, + packed_conversion: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cp_async_copy: None, + cp_async_control: None, + cp_async_mbarrier: None, + mbarrier_basic: None, + movmatrix: None, + mbarrier_extended: None, + register_mma: None, + sparse_mma: None, + prmt: None, + cluster_barrier: None, + wgmma_control: None, + special_register: None, + debug_control: None, + cluster_memory: None, + clc: Some(Clc { + operation: recipe.operation, + adapter: recipe.adapter, + runtime_validation: admission.runtime_validation, + }), + tma: None, + tcgen05: None, + ldmatrix_variant: None, + ldmatrix_safety: None, + ldmatrix_adapter: None, + selected_address_space: None, + expected_ptx: InstructionPattern { + mnemonic: "clusterlaunchcontrol".into(), + modifiers: recipe.modifiers.iter().map(|value| (*value).into()).collect(), + operands: if query { + vec![ + OperandPattern::Register, + OperandPattern::Exact { value: "%clc_handle".into() }, + ] + } else { + recipe.operands.to_vec() + }, + }, + summary: recipe.summary.into(), + }) + }) + .collect() +} + +pub(in crate::resolve) fn validate_clc_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, +) -> Result<()> { + let clc = policy + .clc + .as_ref() + .with_context(|| format!("{} has no closed CLC contract", policy.id))?; + let recipe = clc_recipe(clc.operation); + let query = recipe.pure; + ensure!( + policy.id == recipe.id + && policy.abi_id == recipe.abi_id + && policy.operation_key == recipe.operation_key + && policy.source.is_none() + && policy.source_record.as_deref() == Some(recipe.source_record) + && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) + && policy.resolved_llvm_symbol.is_none() + && declaration.source_record == recipe.source_record + && declaration.llvm_name == recipe.llvm_symbol, + "{} CLC identity changed", + policy.id + ); + ensure!( + policy.rust_module == "clc" + && policy.rust_name == recipe.id + && policy.rust_arguments == recipe.rust_arguments + && policy.rust_result == if query { "u32" } else { "()" } + && !policy.safe + && !policy.must_use + && policy.safe_allowlist_reason.is_none() + && policy.public_rust_path == format!("cuda_intrinsics::clc::{}", recipe.id) + && policy.compatibility_rust_paths == [format!("cuda_device::clc::{}", recipe.id)], + "{} CLC Rust API changed", + policy.id + ); + ensure!( + policy.dialect_op_type == recipe.dialect_op_type + && policy.dialect_op_name == recipe.dialect_op_name + && policy.dialect_operands == recipe.dialect_operands + && policy.dialect_results == recipe.dialect_results + && policy.llvm_arguments == recipe.llvm_arguments + && policy.llvm_results == recipe.llvm_results + && policy.lowering == "generated_clc", + "{} CLC carrier or LLVM adapter changed", + policy.id + ); + ensure!( + declaration.arguments == recipe.llvm_arguments + && declaration.results == recipe.llvm_results + && declaration.properties + == if query { + vec!["IntrNoMem", "IntrSpeculatable"] + } else { + vec!["IntrArgMemOnly", "IntrHasSideEffects"] + }, + "{} imported CLC declaration changed", + policy.id + ); + ensure!( + policy.pure == recipe.pure + && policy.memory == recipe.memory + && !policy.convergent + && policy.execution_scope == recipe.execution_scope + && clc.adapter == recipe.adapter + && clc.runtime_validation == RuntimeValidation::Unexecuted, + "{} CLC semantics changed", + policy.id + ); + ensure!( + policy.minimum_ptx == "8.6" + && policy.minimum_sm.as_deref() == recipe.minimum_sm + && policy.targets == recipe.targets + && policy.ptx_isa_version == "9.3" + && policy.ptx_result == if query { "u32" } else { "()" }, + "{} CLC target contract changed", + policy.id + ); + ensure!( + policy.expected_ptx.mnemonic == "clusterlaunchcontrol" + && policy.expected_ptx.modifiers == recipe.modifiers + && policy.expected_ptx.operands + == if query { + vec![ + OperandPattern::Register, + OperandPattern::Exact { + value: "%clc_handle".into(), + }, + ] + } else { + recipe.operands.to_vec() + } + && policy.backend_lowerings.len() == 2 + && policy.backend_lowerings.iter().all(|route| { + route.mechanism == BackendLoweringMechanism::TypedNvvm + && route.minimum_ptx.as_deref() == Some("8.6") + && route.minimum_sm.as_deref() == recipe.minimum_sm + && !route.evidence_profile.trim().is_empty() + }), + "{} CLC PTX shape or backend route changed", + policy.id + ); + ensure_no_other_family_contract(policy, "CLC")?; + Ok(()) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/cluster.rs b/crates/cuda-intrinsics-gen/src/resolve/families/cluster.rs new file mode 100644 index 0000000000..9a07382a9d --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/cluster.rs @@ -0,0 +1,791 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, ClusterBarrier, ClusterBarrierAdmission, ClusterBarrierMode, + ClusterBarrierOrdering, ClusterMemory, ClusterMemoryAdapter, ClusterMemoryAdmission, + ClusterMemoryOperation, ClusterMemorySourceContract, ImportedIntrinsic, IntrinsicBackend, + IntrinsicSource, OverlayBackendLowering, OverlayIntrinsic, RuntimeValidation, +}; +use crate::ptx::{InstructionPattern, OperandPattern}; +use anyhow::{Context, Result, ensure}; +use std::collections::BTreeSet; + +use crate::resolve::guards::*; + +#[derive(Clone, Copy)] +pub(in crate::resolve) struct ClusterBarrierRecipe { + pub(in crate::resolve) mode: ClusterBarrierMode, + pub(in crate::resolve) abi_id: &'static str, + pub(in crate::resolve) id: &'static str, + pub(in crate::resolve) suffix: &'static str, + pub(in crate::resolve) minimum_ptx: &'static str, + pub(in crate::resolve) ordering: ClusterBarrierOrdering, + pub(in crate::resolve) aligned: bool, + pub(in crate::resolve) summary: &'static str, +} + +pub(in crate::resolve) fn cluster_barrier_recipe(mode: ClusterBarrierMode) -> ClusterBarrierRecipe { + match mode { + ClusterBarrierMode::Arrive => ClusterBarrierRecipe { + mode, + abi_id: "i0277", + id: "barrier_cluster_arrive", + suffix: "arrive", + minimum_ptx: "7.8", + ordering: ClusterBarrierOrdering::Release, + aligned: false, + summary: "Arrives at the cluster barrier with release ordering.", + }, + ClusterBarrierMode::ArriveAligned => ClusterBarrierRecipe { + mode, + abi_id: "i0278", + id: "barrier_cluster_arrive_aligned", + suffix: "arrive.aligned", + minimum_ptx: "7.8", + ordering: ClusterBarrierOrdering::Release, + aligned: true, + summary: "Arrives at the cluster barrier in aligned mode with release ordering.", + }, + ClusterBarrierMode::ArriveRelaxed => ClusterBarrierRecipe { + mode, + abi_id: "i0279", + id: "barrier_cluster_arrive_relaxed", + suffix: "arrive.relaxed", + minimum_ptx: "8.0", + ordering: ClusterBarrierOrdering::Relaxed, + aligned: false, + summary: "Arrives at the cluster barrier without a release guarantee.", + }, + ClusterBarrierMode::ArriveRelaxedAligned => ClusterBarrierRecipe { + mode, + abi_id: "i0280", + id: "barrier_cluster_arrive_relaxed_aligned", + suffix: "arrive.relaxed.aligned", + minimum_ptx: "8.0", + ordering: ClusterBarrierOrdering::Relaxed, + aligned: true, + summary: "Arrives at the cluster barrier in aligned mode without a release guarantee.", + }, + ClusterBarrierMode::Wait => ClusterBarrierRecipe { + mode, + abi_id: "i0281", + id: "barrier_cluster_wait", + suffix: "wait", + minimum_ptx: "7.8", + ordering: ClusterBarrierOrdering::Acquire, + aligned: false, + summary: "Waits at the cluster barrier with acquire ordering.", + }, + ClusterBarrierMode::WaitAligned => ClusterBarrierRecipe { + mode, + abi_id: "i0282", + id: "barrier_cluster_wait_aligned", + suffix: "wait.aligned", + minimum_ptx: "7.8", + ordering: ClusterBarrierOrdering::Acquire, + aligned: true, + summary: "Waits at the cluster barrier in aligned mode with acquire ordering.", + }, + } +} + +pub(in crate::resolve) fn expand_cluster_barrier_admission( + admission: &ClusterBarrierAdmission, +) -> Result> { + ensure!( + admission.runtime_validation == RuntimeValidation::Unexecuted, + "cluster-barrier runtime validation may be marked executed only with GPU evidence" + ); + ensure!( + !admission.llvm_evidence_profile.trim().is_empty() + && !admission.libnvvm_evidence_profile.trim().is_empty(), + "compact cluster-barrier admission requires both backend evidence profiles" + ); + let expected_modes = BTreeSet::from([ + ClusterBarrierMode::Arrive, + ClusterBarrierMode::ArriveAligned, + ClusterBarrierMode::ArriveRelaxed, + ClusterBarrierMode::ArriveRelaxedAligned, + ClusterBarrierMode::Wait, + ClusterBarrierMode::WaitAligned, + ]); + let actual_modes: BTreeSet<_> = admission + .variants + .iter() + .map(|variant| variant.mode) + .collect(); + ensure!( + admission.variants.len() == expected_modes.len() && actual_modes == expected_modes, + "compact cluster-barrier admission must contain each reviewed mode exactly once" + ); + + admission + .variants + .iter() + .map(|variant| { + let recipe = cluster_barrier_recipe(variant.mode); + ensure!( + variant.abi_id == recipe.abi_id, + "{} must keep reserved ABI ID {}", + recipe.id, + recipe.abi_id + ); + let source_record = format!("int_nvvm_barrier_cluster_{}", recipe.suffix.replace('.', "_")); + let llvm_symbol = format!("llvm.nvvm.barrier.cluster.{}", recipe.suffix); + let modifiers: Vec = recipe.suffix.split('.').map(str::to_owned).collect(); + Ok(OverlayIntrinsic { + id: recipe.id.into(), + abi_id: variant.abi_id.clone(), + operation_key: format!("cluster.barrier.{}", recipe.suffix), + family: "cluster_barrier".into(), + source: None, + source_record: Some(source_record), + rust_module: "cluster".into(), + rust_name: recipe.id.into(), + rust_arguments: vec![], + rust_result: "()".into(), + safe: false, + must_use: false, + safe_allowlist_reason: None, + public_rust_path: format!("cuda_intrinsics::cluster::{}", recipe.id), + compatibility_rust_paths: vec![format!("cuda_device::cluster::{}", recipe.id)], + dialect_op_type: "ClusterBarrierOp".into(), + dialect_op_name: "nvvm.cluster_barrier".into(), + dialect_operands: vec![], + dialect_results: vec![], + llvm_symbol: Some(llvm_symbol), + resolved_llvm_symbol: None, + llvm_arguments: vec![], + llvm_results: vec![], + pure: false, + memory: "read_write".into(), + convergent: true, + execution_scope: "cluster".into(), + minimum_ptx: recipe.minimum_ptx.into(), + minimum_sm: Some("sm_90".into()), + ptx_result: "()".into(), + targets: "all".into(), + ptx_isa_version: "9.3".into(), + ptx_isa_section: + "Parallel Synchronization and Communication Instructions: barrier.cluster" + .into(), + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-barrier-cluster".into(), + lowering: "generated_cluster_barrier".into(), + backend_lowerings: vec![ + OverlayBackendLowering { + backend: IntrinsicBackend::LlvmNvptx, + mechanism: BackendLoweringMechanism::TypedNvvm, + evidence_profile: admission.llvm_evidence_profile.clone(), + targets: None, + minimum_ptx: Some(recipe.minimum_ptx.into()), + minimum_sm: Some("sm_90".into()), + }, + OverlayBackendLowering { + backend: IntrinsicBackend::LibNvvm, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: admission.libnvvm_evidence_profile.clone(), + targets: None, + minimum_ptx: Some(recipe.minimum_ptx.into()), + minimum_sm: Some("sm_90".into()), + }, + ], + packed_atomic: None, + redux: None, + vote: None, + active_mask: None, + warp_match: None, + warp_barrier: None, + warp_shuffle: None, + dot_product: None, + packed_alu: None, + integer_minmax: None, + packed_conversion: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cp_async_copy: None, + cp_async_control: None, + cp_async_mbarrier: None, + mbarrier_basic: None, + movmatrix: None, + mbarrier_extended: None, + register_mma: None, + sparse_mma: None, + prmt: None, + cluster_barrier: Some(ClusterBarrier { + mode: recipe.mode, + ordering: recipe.ordering, + aligned: recipe.aligned, + }), + wgmma_control: None, + special_register: None, + debug_control: None, + cluster_memory: None, + clc: None, + tma: None, + tcgen05: None, + ldmatrix_variant: None, + ldmatrix_safety: None, + ldmatrix_adapter: None, + selected_address_space: None, + expected_ptx: InstructionPattern { + mnemonic: "barrier".into(), + modifiers: std::iter::once("cluster".into()).chain(modifiers).collect(), + operands: vec![], + }, + summary: recipe.summary.into(), + }) + }) + .collect() +} + +pub(in crate::resolve) fn validate_cluster_barrier_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, +) -> Result<()> { + let barrier = policy + .cluster_barrier + .as_ref() + .with_context(|| format!("{} has no closed cluster-barrier contract", policy.id))?; + let recipe = cluster_barrier_recipe(barrier.mode); + let source_record = format!( + "int_nvvm_barrier_cluster_{}", + recipe.suffix.replace('.', "_") + ); + let llvm_symbol = format!("llvm.nvvm.barrier.cluster.{}", recipe.suffix); + ensure!( + barrier.ordering == recipe.ordering + && barrier.aligned == recipe.aligned + && policy.id == recipe.id + && policy.abi_id == recipe.abi_id + && policy.operation_key == format!("cluster.barrier.{}", recipe.suffix) + && policy.source.is_none() + && policy.source_record.as_deref() == Some(source_record.as_str()) + && policy.llvm_symbol.as_deref() == Some(llvm_symbol.as_str()) + && policy.resolved_llvm_symbol.is_none(), + "{} identity or semantics do not match its closed cluster-barrier recipe", + policy.id + ); + ensure!( + policy.rust_module == "cluster" + && policy.rust_name == recipe.id + && policy.rust_arguments.is_empty() + && policy.rust_result == "()" + && !policy.safe + && !policy.must_use + && policy.public_rust_path == format!("cuda_intrinsics::cluster::{}", recipe.id) + && policy.compatibility_rust_paths == [format!("cuda_device::cluster::{}", recipe.id)], + "{} Rust API does not match its closed cluster-barrier recipe", + policy.id + ); + ensure!( + policy.dialect_op_type == "ClusterBarrierOp" + && policy.dialect_op_name == "nvvm.cluster_barrier" + && policy.dialect_operands.is_empty() + && policy.dialect_results.is_empty() + && policy.llvm_arguments.is_empty() + && policy.llvm_results.is_empty() + && policy.lowering == "generated_cluster_barrier", + "{} carrier or lowering does not match its closed cluster-barrier recipe", + policy.id + ); + ensure!( + declaration.classes == ["SDPatternOperator", "Intrinsic"] + && declaration.properties == ["IntrConvergent", "IntrNoCallback"] + && !policy.pure + && policy.memory == "read_write" + && policy.convergent + && policy.execution_scope == "cluster", + "{} effects disagree with the imported cluster-barrier declaration", + policy.id + ); + ensure!( + policy.minimum_ptx == recipe.minimum_ptx + && policy.minimum_sm.as_deref() == Some("sm_90") + && policy.targets == "all" + && policy.ptx_result == "()" + && policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section + == "Parallel Synchronization and Communication Instructions: barrier.cluster" + && policy.ptx_isa_url + == "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-barrier-cluster", + "{} target floor or PTX provenance changed", + policy.id + ); + let expected_modifiers: Vec<_> = std::iter::once("cluster") + .chain(recipe.suffix.split('.')) + .collect(); + ensure!( + policy.expected_ptx.mnemonic == "barrier" + && policy.expected_ptx.modifiers == expected_modifiers + && policy.expected_ptx.operands.is_empty(), + "{} expected PTX does not match its exact cluster-barrier spelling", + policy.id + ); + ensure!( + (recipe.ordering == ClusterBarrierOrdering::Relaxed) == recipe.suffix.contains(".relaxed") + && recipe.aligned == recipe.suffix.ends_with(".aligned") + && matches!( + (recipe.mode, recipe.ordering), + ( + ClusterBarrierMode::Arrive | ClusterBarrierMode::ArriveAligned, + ClusterBarrierOrdering::Release + ) | ( + ClusterBarrierMode::ArriveRelaxed | ClusterBarrierMode::ArriveRelaxedAligned, + ClusterBarrierOrdering::Relaxed + ) | ( + ClusterBarrierMode::Wait | ClusterBarrierMode::WaitAligned, + ClusterBarrierOrdering::Acquire + ) + ), + "{} cluster-barrier semantic recipe is inconsistent", + policy.id + ); + let backend_pairs: BTreeSet<_> = policy + .backend_lowerings + .iter() + .map(|lowering| (lowering.backend, lowering.mechanism)) + .collect(); + ensure!( + policy.backend_lowerings.len() == 2 + && backend_pairs + == BTreeSet::from([ + ( + IntrinsicBackend::LlvmNvptx, + BackendLoweringMechanism::TypedNvvm, + ), + ( + IntrinsicBackend::LibNvvm, + BackendLoweringMechanism::InlinePtx, + ), + ]) + && policy.backend_lowerings.iter().all(|lowering| { + lowering.minimum_ptx.as_deref() == Some(recipe.minimum_ptx) + && lowering.minimum_sm.as_deref() == Some("sm_90") + && !lowering.evidence_profile.trim().is_empty() + }), + "{} must define exactly the reviewed cluster-barrier backend routes", + policy.id + ); + ensure_no_other_family_contract(policy, "cluster barrier")?; + Ok(()) +} + +#[derive(Clone)] +pub(in crate::resolve) struct ClusterMemoryRecipe { + operation: ClusterMemoryOperation, + abi_id: &'static str, + id: &'static str, + operation_key: &'static str, + source_record: Option<&'static str>, + llvm_symbol: Option<&'static str>, + ptx_native_instruction: Option<&'static str>, + rust_arguments: &'static [&'static str], + rust_result: &'static str, + compatibility_paths: &'static [&'static str], + dialect_op_type: &'static str, + dialect_op_name: &'static str, + dialect_operands: &'static [&'static str], + dialect_results: &'static [&'static str], + llvm_arguments: &'static [&'static str], + llvm_results: &'static [&'static str], + adapter: ClusterMemoryAdapter, + source_contract: ClusterMemorySourceContract, + memory: &'static str, + expected_ptx: InstructionPattern, + inline_ptx: &'static str, + inline_constraints: &'static str, + ptx_isa_section: &'static str, + ptx_isa_anchor: &'static str, + summary: &'static str, +} + +pub(in crate::resolve) fn cluster_memory_recipe( + operation: ClusterMemoryOperation, +) -> ClusterMemoryRecipe { + match operation { + ClusterMemoryOperation::MapSharedRank => ClusterMemoryRecipe { + operation, + abi_id: "i0320", + id: "map_shared_rank", + operation_key: "cluster.shared_address.map_rank", + source_record: Some("int_nvvm_mapa_shared_cluster"), + llvm_symbol: Some("llvm.nvvm.mapa.shared.cluster"), + ptx_native_instruction: None, + rust_arguments: &["*const u8", "u32"], + rust_result: "*const u8", + compatibility_paths: &[ + "cuda_device::cluster::map_shared_rank", + "cuda_device::cluster::map_shared_rank_mut", + ], + dialect_op_type: "MapaSharedClusterOp", + dialect_op_name: "nvvm.mapa_shared_cluster", + dialect_operands: &["ptr", "i32"], + dialect_results: &["ptr"], + llvm_arguments: &["shared_ptr", "i32"], + llvm_results: &["shared_cluster_ptr"], + adapter: ClusterMemoryAdapter::GenericConstAndMutPointerRankToSamePointer, + source_contract: ClusterMemorySourceContract::LlvmMapaSharedClusterAs7IdentityInlinePtx, + memory: "none", + expected_ptx: InstructionPattern { + mnemonic: "mapa".into(), + modifiers: vec!["shared::cluster".into(), "u64".into()], + operands: vec![ + OperandPattern::Register, + OperandPattern::Register, + OperandPattern::Register, + ], + }, + inline_ptx: "mapa.shared::cluster.u64 $0, $1, $2;", + inline_constraints: "=l,l,r", + ptx_isa_section: "9.7.9.24 Data Movement and Conversion Instructions: mapa", + ptx_isa_anchor: "data-movement-and-conversion-instructions-mapa", + summary: "Maps a CTA-shared address to the same offset in another cluster rank.", + }, + ClusterMemoryOperation::ReadU32 => ClusterMemoryRecipe { + operation, + abi_id: "i0321", + id: "dsmem_read_u32", + operation_key: "cluster.shared_memory.map_rank_then_read_u32", + source_record: None, + llvm_symbol: None, + ptx_native_instruction: Some("mapa.shared::cluster.u64 + ld.shared::cluster.u32"), + rust_arguments: &["*const u32", "u32"], + rust_result: "u32", + compatibility_paths: &["cuda_device::cluster::dsmem_read_u32"], + dialect_op_type: "DsmemReadU32Op", + dialect_op_name: "nvvm.dsmem_read_u32", + dialect_operands: &["ptr", "i32"], + dialect_results: &["i32"], + llvm_arguments: &[], + llvm_results: &[], + adapter: ClusterMemoryAdapter::ConstU32PointerRankToU32, + source_contract: ClusterMemorySourceContract::PtxNativeMapaThenWeakClusterLoad, + memory: "read", + expected_ptx: InstructionPattern { + mnemonic: "ld".into(), + modifiers: vec!["shared::cluster".into(), "u32".into()], + operands: vec![OperandPattern::Register, OperandPattern::Address], + }, + inline_ptx: "{ .reg .u64 %mapped; mapa.shared::cluster.u64 %mapped, $1, $2; ld.shared::cluster.u32 $0, [%mapped]; }", + inline_constraints: "=r,l,r,~{memory}", + ptx_isa_section: "9.7.9.8 Data Movement and Conversion Instructions: ld", + ptx_isa_anchor: "data-movement-and-conversion-instructions-ld", + summary: "Maps a CTA-shared address to another cluster rank and reads one weak u32 value.", + }, + } +} + +pub(crate) fn cluster_memory_inline_recipe( + operation: ClusterMemoryOperation, +) -> (&'static str, &'static str) { + let recipe = cluster_memory_recipe(operation); + (recipe.inline_ptx, recipe.inline_constraints) +} + +pub(in crate::resolve) fn expand_cluster_memory_admission( + admission: &ClusterMemoryAdmission, +) -> Result> { + ensure!( + admission.runtime_validation == RuntimeValidation::Unexecuted, + "cluster-memory runtime validation may be marked executed only with GPU evidence" + ); + ensure!( + !admission.llvm_evidence_profile.trim().is_empty() + && !admission.libnvvm_evidence_profile.trim().is_empty(), + "compact cluster-memory admission requires both backend evidence profiles" + ); + let expected = [ + ClusterMemoryOperation::MapSharedRank, + ClusterMemoryOperation::ReadU32, + ]; + ensure!( + admission.variants.len() == expected.len() + && admission + .variants + .iter() + .map(|variant| variant.operation) + .eq(expected), + "compact cluster-memory admission must list map_shared_rank and read_u32 once in canonical order" + ); + + admission + .variants + .iter() + .map(|variant| { + let recipe = cluster_memory_recipe(variant.operation); + ensure!( + variant.abi_id == recipe.abi_id, + "{} must keep reserved ABI ID {}", + recipe.id, + recipe.abi_id + ); + let source = + recipe + .ptx_native_instruction + .map(|instruction| IntrinsicSource::PtxNative { + instruction: instruction.into(), + }); + Ok(OverlayIntrinsic { + id: recipe.id.into(), + abi_id: variant.abi_id.clone(), + operation_key: recipe.operation_key.into(), + family: "cluster_memory".into(), + source, + source_record: recipe.source_record.map(Into::into), + rust_module: "cluster".into(), + rust_name: recipe.id.into(), + rust_arguments: recipe + .rust_arguments + .iter() + .map(|value| (*value).into()) + .collect(), + rust_result: recipe.rust_result.into(), + safe: false, + must_use: true, + safe_allowlist_reason: None, + public_rust_path: format!("cuda_intrinsics::cluster::{}", recipe.id), + compatibility_rust_paths: recipe + .compatibility_paths + .iter() + .map(|path| (*path).into()) + .collect(), + dialect_op_type: recipe.dialect_op_type.into(), + dialect_op_name: recipe.dialect_op_name.into(), + dialect_operands: recipe + .dialect_operands + .iter() + .map(|value| (*value).into()) + .collect(), + dialect_results: recipe + .dialect_results + .iter() + .map(|value| (*value).into()) + .collect(), + llvm_symbol: recipe.llvm_symbol.map(Into::into), + resolved_llvm_symbol: None, + llvm_arguments: recipe + .llvm_arguments + .iter() + .map(|value| (*value).into()) + .collect(), + llvm_results: recipe + .llvm_results + .iter() + .map(|value| (*value).into()) + .collect(), + pure: false, + memory: recipe.memory.into(), + convergent: true, + execution_scope: "cluster".into(), + minimum_ptx: "7.8".into(), + minimum_sm: Some("sm_90".into()), + ptx_result: recipe.rust_result.into(), + targets: "all".into(), + ptx_isa_version: "9.3".into(), + ptx_isa_section: recipe.ptx_isa_section.into(), + ptx_isa_url: format!( + "https://docs.nvidia.com/cuda/parallel-thread-execution/#{}", + recipe.ptx_isa_anchor + ), + lowering: "generated_cluster_memory_inline_ptx".into(), + backend_lowerings: [IntrinsicBackend::LlvmNvptx, IntrinsicBackend::LibNvvm] + .into_iter() + .map(|backend| OverlayBackendLowering { + backend, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: match backend { + IntrinsicBackend::LlvmNvptx => admission.llvm_evidence_profile.clone(), + IntrinsicBackend::LibNvvm => admission.libnvvm_evidence_profile.clone(), + }, + targets: None, + minimum_ptx: Some("7.8".into()), + minimum_sm: Some("sm_90".into()), + }) + .collect(), + packed_atomic: None, + redux: None, + vote: None, + active_mask: None, + warp_match: None, + warp_barrier: None, + warp_shuffle: None, + dot_product: None, + packed_alu: None, + integer_minmax: None, + packed_conversion: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cp_async_copy: None, + cp_async_control: None, + cp_async_mbarrier: None, + mbarrier_basic: None, + movmatrix: None, + mbarrier_extended: None, + register_mma: None, + sparse_mma: None, + prmt: None, + cluster_barrier: None, + wgmma_control: None, + special_register: None, + debug_control: None, + cluster_memory: Some(ClusterMemory { + operation: recipe.operation, + adapter: recipe.adapter, + source_contract: recipe.source_contract, + runtime_validation: admission.runtime_validation, + }), + clc: None, + tma: None, + tcgen05: None, + ldmatrix_variant: None, + ldmatrix_safety: None, + ldmatrix_adapter: None, + selected_address_space: None, + expected_ptx: recipe.expected_ptx, + summary: recipe.summary.into(), + }) + }) + .collect() +} + +pub(in crate::resolve) fn validate_cluster_memory_policy( + policy: &OverlayIntrinsic, + source: &IntrinsicSource, + declaration: Option<&ImportedIntrinsic>, +) -> Result<()> { + let cluster = policy + .cluster_memory + .as_ref() + .with_context(|| format!("{} has no closed cluster-memory contract", policy.id))?; + let recipe = cluster_memory_recipe(cluster.operation); + ensure!( + cluster.adapter == recipe.adapter + && cluster.source_contract == recipe.source_contract + && cluster.runtime_validation == RuntimeValidation::Unexecuted + && policy.id == recipe.id + && policy.operation_key == recipe.operation_key, + "{} does not match its closed cluster-memory identity", + policy.id + ); + match cluster.source_contract { + ClusterMemorySourceContract::LlvmMapaSharedClusterAs7IdentityInlinePtx => { + let declaration = declaration.context("map_shared_rank requires its LLVM identity")?; + ensure!( + matches!(source, IntrinsicSource::LlvmImported { source_record } + if source_record == "int_nvvm_mapa_shared_cluster") + && policy.source.is_none() + && policy.source_record.as_deref() == recipe.source_record + && policy.llvm_symbol.as_deref() == recipe.llvm_symbol + && declaration.arguments == ["shared_ptr", "i32"] + && declaration.results == ["shared_cluster_ptr"] + && declaration.properties + == ["IntrNoMem", "IntrSpeculatable", "NoCapture"], + "{} must retain the AS7-returning LLVM mapa record as identity only", + policy.id + ); + let mut selections = Vec::new(); + for selection in &declaration.selections { + if selection_matches_policy(policy, selection)? { + selections.push(selection); + } + } + ensure!( + selections.len() == 2 + && selections + .iter() + .map(|selection| selection.source_record.as_str()) + .collect::>() + == BTreeSet::from(["mapa_shared_cluster_64", "mapa_shared_cluster_64i",]) + && selections.iter().all(|selection| { + selection.asm == "mapa.shared::cluster.u64 \t$d, $a, $b;" + && selection.predicates + == [ + "Subtarget->getSmVersion() >= 90", + "Subtarget->getPTXVersion() >= 78", + ] + }), + "{} must retain both exact 64-bit mapa selections", + policy.id + ); + } + ClusterMemorySourceContract::PtxNativeMapaThenWeakClusterLoad => ensure!( + matches!(source, IntrinsicSource::PtxNative { instruction } + if Some(instruction.as_str()) == recipe.ptx_native_instruction) + && policy.source_record.is_none() + && policy.llvm_symbol.is_none() + && declaration.is_none() + && policy.llvm_arguments.is_empty() + && policy.llvm_results.is_empty(), + "{} must remain a PTX-native mapa plus weak cluster-load composite", + policy.id + ), + } + ensure!( + policy.rust_module == "cluster" + && policy.rust_name == recipe.id + && policy.rust_arguments == recipe.rust_arguments + && policy.rust_result == recipe.rust_result + && !policy.safe + && policy.must_use + && policy.public_rust_path == format!("cuda_intrinsics::cluster::{}", recipe.id) + && policy.compatibility_rust_paths == recipe.compatibility_paths, + "{} Rust API or compatibility paths changed", + policy.id + ); + ensure!( + policy.dialect_op_type == recipe.dialect_op_type + && policy.dialect_op_name == recipe.dialect_op_name + && policy.dialect_operands == recipe.dialect_operands + && policy.dialect_results == recipe.dialect_results + && policy.llvm_arguments == recipe.llvm_arguments + && policy.llvm_results == recipe.llvm_results + && policy.lowering == "generated_cluster_memory_inline_ptx", + "{} dialect carrier or AS7 source boundary changed", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == recipe.memory + && policy.convergent + && policy.execution_scope == "cluster", + "{} effects changed", + policy.id + ); + ensure!( + policy.minimum_ptx == "7.8" + && policy.minimum_sm.as_deref() == Some("sm_90") + && policy.targets == "all" + && policy.ptx_result == recipe.rust_result + && policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section == recipe.ptx_isa_section + && policy.ptx_isa_url + == format!( + "https://docs.nvidia.com/cuda/parallel-thread-execution/#{}", + recipe.ptx_isa_anchor + ), + "{} PTX provenance or target floor changed", + policy.id + ); + ensure!( + policy.expected_ptx == recipe.expected_ptx, + "{} expected PTX changed", + policy.id + ); + ensure_exact_inline_ptx_backends( + policy, + [ + (IntrinsicBackend::LlvmNvptx, "7.8", Some("sm_90")), + (IntrinsicBackend::LibNvvm, "7.8", Some("sm_90")), + ], + "cluster-memory", + )?; + ensure_no_other_family_contract(policy, "cluster-memory")?; + Ok(()) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/cp_async.rs b/crates/cuda-intrinsics-gen/src/resolve/families/cp_async.rs new file mode 100644 index 0000000000..faf83aedbd --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/cp_async.rs @@ -0,0 +1,1121 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, CpAsyncAdapter, CpAsyncCachePolicy, CpAsyncControlAdapter, + CpAsyncControlOperation, CpAsyncCopySize, CpAsyncMbarrierAdapter, CpAsyncMbarrierOperation, + CpAsyncMbarrierStateSpace, CpAsyncSourceSize, ImportedIntrinsic, IntrinsicBackend, + MbarrierBasicAdapter, MbarrierBasicOperation, MbarrierStateSpace, OverlayIntrinsic, + RuntimeValidation, +}; +use crate::ptx::OperandPattern; +use anyhow::{Context, Result, ensure}; +use std::collections::BTreeSet; + +use crate::resolve::guards::*; + +pub(in crate::resolve) struct CpAsyncCopyRecipe { + pub(in crate::resolve) id: &'static str, + pub(in crate::resolve) abi_id: &'static str, + pub(in crate::resolve) operation_key: &'static str, + pub(in crate::resolve) rust_name: &'static str, + pub(in crate::resolve) dialect_op_type: &'static str, + pub(in crate::resolve) dialect_op_name: &'static str, + pub(in crate::resolve) source_record: &'static str, + pub(in crate::resolve) llvm_symbol: &'static str, + pub(in crate::resolve) selections: &'static [&'static str], + pub(in crate::resolve) summary: &'static str, +} + +pub(in crate::resolve) fn cp_async_copy_recipe( + copy: &crate::model::CpAsyncCopy, +) -> Option { + match (copy.cache_policy, copy.copy_size, copy.source_size) { + (CpAsyncCachePolicy::Ca, CpAsyncCopySize::B4, CpAsyncSourceSize::Full) => { + Some(CpAsyncCopyRecipe { + id: "cp_async_ca_4", + abi_id: "i0086", + operation_key: "memory.copy.async.global_to_shared.ca.4.full", + rust_name: "cp_async_ca_4", + dialect_op_type: "CpAsyncCa4Op", + dialect_op_name: "nvvm.cp_async_ca_4", + source_record: "int_nvvm_cp_async_ca_shared_global_4", + llvm_symbol: "llvm.nvvm.cp.async.ca.shared.global.4", + selections: &["CP_ASYNC_CA_SHARED_GLOBAL_4"], + summary: "Starts a four-byte cache-all asynchronous copy from global to shared memory.", + }) + } + (CpAsyncCachePolicy::Ca, CpAsyncCopySize::B4, CpAsyncSourceSize::Runtime) => { + Some(CpAsyncCopyRecipe { + id: "cp_async_ca_zfill_4", + abi_id: "i0087", + operation_key: "memory.copy.async.global_to_shared.ca.4.runtime_source_size", + rust_name: "cp_async_ca_zfill_4", + dialect_op_type: "CpAsyncCaZfill4Op", + dialect_op_name: "nvvm.cp_async_ca_zfill_4", + source_record: "int_nvvm_cp_async_ca_shared_global_4_s", + llvm_symbol: "llvm.nvvm.cp.async.ca.shared.global.4.s", + selections: &[ + "CP_ASYNC_CA_SHARED_GLOBAL_4_s", + "CP_ASYNC_CA_SHARED_GLOBAL_4_si", + ], + summary: "Starts a four-byte cache-all asynchronous copy and zero-fills bytes beyond the runtime source size.", + }) + } + (CpAsyncCachePolicy::Ca, CpAsyncCopySize::B8, CpAsyncSourceSize::Full) => { + Some(CpAsyncCopyRecipe { + id: "cp_async_ca_8", + abi_id: "i0088", + operation_key: "memory.copy.async.global_to_shared.ca.8.full", + rust_name: "cp_async_ca_8", + dialect_op_type: "CpAsyncCa8Op", + dialect_op_name: "nvvm.cp_async_ca_8", + source_record: "int_nvvm_cp_async_ca_shared_global_8", + llvm_symbol: "llvm.nvvm.cp.async.ca.shared.global.8", + selections: &["CP_ASYNC_CA_SHARED_GLOBAL_8"], + summary: "Starts an eight-byte cache-all asynchronous copy from global to shared memory.", + }) + } + (CpAsyncCachePolicy::Ca, CpAsyncCopySize::B8, CpAsyncSourceSize::Runtime) => { + Some(CpAsyncCopyRecipe { + id: "cp_async_ca_zfill_8", + abi_id: "i0089", + operation_key: "memory.copy.async.global_to_shared.ca.8.runtime_source_size", + rust_name: "cp_async_ca_zfill_8", + dialect_op_type: "CpAsyncCaZfill8Op", + dialect_op_name: "nvvm.cp_async_ca_zfill_8", + source_record: "int_nvvm_cp_async_ca_shared_global_8_s", + llvm_symbol: "llvm.nvvm.cp.async.ca.shared.global.8.s", + selections: &[ + "CP_ASYNC_CA_SHARED_GLOBAL_8_s", + "CP_ASYNC_CA_SHARED_GLOBAL_8_si", + ], + summary: "Starts an eight-byte cache-all asynchronous copy and zero-fills bytes beyond the runtime source size.", + }) + } + (CpAsyncCachePolicy::Ca, CpAsyncCopySize::B16, CpAsyncSourceSize::Full) => { + Some(CpAsyncCopyRecipe { + id: "cp_async_ca_16", + abi_id: "i0090", + operation_key: "memory.copy.async.global_to_shared.ca.16.full", + rust_name: "cp_async_ca_16", + dialect_op_type: "CpAsyncCa16Op", + dialect_op_name: "nvvm.cp_async_ca_16", + source_record: "int_nvvm_cp_async_ca_shared_global_16", + llvm_symbol: "llvm.nvvm.cp.async.ca.shared.global.16", + selections: &["CP_ASYNC_CA_SHARED_GLOBAL_16"], + summary: "Starts a sixteen-byte cache-all asynchronous copy from global to shared memory.", + }) + } + (CpAsyncCachePolicy::Ca, CpAsyncCopySize::B16, CpAsyncSourceSize::Runtime) => { + Some(CpAsyncCopyRecipe { + id: "cp_async_ca_zfill_16", + abi_id: "i0091", + operation_key: "memory.copy.async.global_to_shared.ca.16.runtime_source_size", + rust_name: "cp_async_ca_zfill_16", + dialect_op_type: "CpAsyncCaZfill16Op", + dialect_op_name: "nvvm.cp_async_ca_zfill_16", + source_record: "int_nvvm_cp_async_ca_shared_global_16_s", + llvm_symbol: "llvm.nvvm.cp.async.ca.shared.global.16.s", + selections: &[ + "CP_ASYNC_CA_SHARED_GLOBAL_16_s", + "CP_ASYNC_CA_SHARED_GLOBAL_16_si", + ], + summary: "Starts a sixteen-byte cache-all asynchronous copy and zero-fills bytes beyond the runtime source size.", + }) + } + (CpAsyncCachePolicy::Cg, CpAsyncCopySize::B16, CpAsyncSourceSize::Full) => { + Some(CpAsyncCopyRecipe { + id: "cp_async_cg_16", + abi_id: "i0092", + operation_key: "memory.copy.async.global_to_shared.cg.16.full", + rust_name: "cp_async_cg_16", + dialect_op_type: "CpAsyncCg16Op", + dialect_op_name: "nvvm.cp_async_cg_16", + source_record: "int_nvvm_cp_async_cg_shared_global_16", + llvm_symbol: "llvm.nvvm.cp.async.cg.shared.global.16", + selections: &["CP_ASYNC_CG_SHARED_GLOBAL_16"], + summary: "Starts a sixteen-byte cache-global asynchronous copy from global to shared memory.", + }) + } + (CpAsyncCachePolicy::Cg, CpAsyncCopySize::B16, CpAsyncSourceSize::Runtime) => { + Some(CpAsyncCopyRecipe { + id: "cp_async_cg_zfill_16", + abi_id: "i0093", + operation_key: "memory.copy.async.global_to_shared.cg.16.runtime_source_size", + rust_name: "cp_async_cg_zfill_16", + dialect_op_type: "CpAsyncCgZfill16Op", + dialect_op_name: "nvvm.cp_async_cg_zfill_16", + source_record: "int_nvvm_cp_async_cg_shared_global_16_s", + llvm_symbol: "llvm.nvvm.cp.async.cg.shared.global.16.s", + selections: &[ + "CP_ASYNC_CG_SHARED_GLOBAL_16_s", + "CP_ASYNC_CG_SHARED_GLOBAL_16_si", + ], + summary: "Starts a sixteen-byte cache-global asynchronous copy and zero-fills bytes beyond the runtime source size.", + }) + } + _ => None, + } +} + +pub(in crate::resolve) fn validate_cp_async_copy_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, +) -> Result<()> { + let copy = policy + .cp_async_copy + .as_ref() + .with_context(|| format!("{} has no closed cp.async copy contract", policy.id))?; + let recipe = cp_async_copy_recipe(copy).with_context(|| { + format!( + "{} requests an unsupported classic cp.async copy form", + policy.id + ) + })?; + let dynamic_source_size = copy.source_size == CpAsyncSourceSize::Runtime; + let expected_adapter = if dynamic_source_size { + CpAsyncAdapter::DirectPointersAndSourceSize + } else { + CpAsyncAdapter::DirectPointers + }; + ensure!( + copy.adapter == expected_adapter, + "{} cp.async source-size form and adapter disagree", + policy.id + ); + ensure!( + copy.runtime_validation == RuntimeValidation::Unexecuted, + "{} cannot claim unrecorded cp.async runtime validation", + policy.id + ); + ensure!( + policy.id == recipe.id + && policy.abi_id == recipe.abi_id + && policy.operation_key == recipe.operation_key + && policy.source.is_none() + && policy.source_record.as_deref() == Some(recipe.source_record) + && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) + && policy.resolved_llvm_symbol.is_none(), + "{} cp.async identity does not match its closed recipe", + policy.id + ); + let rust_arguments = if dynamic_source_size { + vec!["*mut u32", "*const u8", "u32"] + } else { + vec!["*mut u32", "*const u32"] + }; + ensure!( + policy.rust_module == "async_copy" + && policy.rust_name == recipe.rust_name + && policy.rust_arguments == rust_arguments + && policy.rust_result == "()" + && !policy.safe + && !policy.must_use + && policy.safe_allowlist_reason.is_none() + && policy.public_rust_path + == format!("cuda_intrinsics::async_copy::{}", recipe.rust_name) + && policy.compatibility_rust_paths + == [format!("cuda_device::async_copy::{}", recipe.rust_name)], + "{} must preserve its unsafe cp.async raw and compatibility API", + policy.id + ); + let llvm_arguments = if dynamic_source_size { + vec!["shared_ptr", "global_ptr", "i32"] + } else { + vec!["shared_ptr", "global_ptr"] + }; + let dialect_operands = if dynamic_source_size { + vec!["ptr", "ptr", "i32"] + } else { + vec!["ptr", "ptr"] + }; + ensure!( + policy.dialect_op_type == recipe.dialect_op_type + && policy.dialect_op_name == recipe.dialect_op_name + && policy.dialect_operands == dialect_operands + && policy.dialect_results.is_empty() + && policy.llvm_arguments == llvm_arguments + && policy.llvm_results.is_empty() + && policy.lowering == "generated_cp_async_copy", + "{} is outside the closed cp.async signature and lowering recipe", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == "read_write" + && !policy.convergent + && policy.execution_scope == "thread" + && policy.minimum_ptx == "7.0" + && policy.minimum_sm.as_deref() == Some("sm_80") + && policy.ptx_result == "()" + && policy.targets == "all", + "{} cp.async effects or target floor disagree with the closed recipe", + policy.id + ); + ensure!( + policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section + == "9.7.9.26.3.1 Data Movement and Conversion Instructions: cp.async" + && policy.ptx_isa_url + == "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async", + "{} cp.async PTX provenance disagrees with the reviewed recipe", + policy.id + ); + ensure!( + policy.summary == recipe.summary, + "{} cp.async summary does not match its closed recipe", + policy.id + ); + ensure!( + declaration.properties + == [ + "IntrArgMemOnly", + "IntrNoCallback", + "NoAlias", + "NoAlias", + "ReadOnly", + "WriteOnly", + ], + "{} cp.async effects disagree with the imported LLVM declaration", + policy.id + ); + let cache = match copy.cache_policy { + CpAsyncCachePolicy::Ca => "ca", + CpAsyncCachePolicy::Cg => "cg", + }; + let mut operands = vec![ + OperandPattern::Address, + OperandPattern::Address, + OperandPattern::Exact { + value: copy.copy_size.bytes().to_string(), + }, + ]; + if dynamic_source_size { + operands.push(OperandPattern::RegisterOrImmediate); + } + ensure!( + policy.expected_ptx.mnemonic == "cp" + && policy.expected_ptx.modifiers == ["async", cache, "shared", "global"] + && policy.expected_ptx.operands == operands, + "{} expected PTX does not match its cp.async recipe", + policy.id + ); + let actual_selections: BTreeSet<_> = declaration + .selections + .iter() + .map(|selection| selection.source_record.as_str()) + .collect(); + ensure!( + actual_selections == recipe.selections.iter().copied().collect(), + "{} imported cp.async selection set changed", + policy.id + ); + let backend_pairs: BTreeSet<_> = policy + .backend_lowerings + .iter() + .map(|lowering| (lowering.backend, lowering.mechanism)) + .collect(); + ensure!( + policy.backend_lowerings.len() == 2 + && backend_pairs + == BTreeSet::from([ + ( + IntrinsicBackend::LlvmNvptx, + BackendLoweringMechanism::TypedNvvm, + ), + ( + IntrinsicBackend::LibNvvm, + BackendLoweringMechanism::InlinePtx, + ), + ]) + && policy.backend_lowerings.iter().all(|lowering| { + lowering.minimum_ptx.is_none() + && lowering.minimum_sm.is_none() + && !lowering.evidence_profile.trim().is_empty() + }), + "{} must define the reviewed typed-LLVM and inline-PTX cp.async routes", + policy.id + ); + ensure_no_other_family_contract(policy, "cp_async_copy")?; + ensure!( + policy.ldmatrix_variant.is_none() + && policy.ldmatrix_safety.is_none() + && policy.ldmatrix_adapter.is_none() + && policy.selected_address_space.is_none(), + "{} mixes ldmatrix state with cp_async_copy", + policy.id + ); + Ok(()) +} + +pub(in crate::resolve) struct CpAsyncControlRecipe { + pub(in crate::resolve) id: &'static str, + pub(in crate::resolve) abi_id: &'static str, + pub(in crate::resolve) operation_key: &'static str, + pub(in crate::resolve) rust_name: &'static str, + pub(in crate::resolve) dialect_op_type: &'static str, + pub(in crate::resolve) dialect_op_name: &'static str, + pub(in crate::resolve) source_record: &'static str, + pub(in crate::resolve) llvm_symbol: &'static str, + pub(in crate::resolve) selection: &'static str, + pub(in crate::resolve) ptx_modifier: &'static str, + pub(in crate::resolve) summary: &'static str, +} + +pub(in crate::resolve) fn cp_async_control_recipe( + operation: CpAsyncControlOperation, +) -> CpAsyncControlRecipe { + match operation { + CpAsyncControlOperation::CommitGroup => CpAsyncControlRecipe { + id: "cp_async_commit_group", + abi_id: "i0094", + operation_key: "memory.copy.async.group.commit", + rust_name: "cp_async_commit_group", + dialect_op_type: "CpAsyncCommitGroupOp", + dialect_op_name: "nvvm.cp_async_commit_group", + source_record: "int_nvvm_cp_async_commit_group", + llvm_symbol: "llvm.nvvm.cp.async.commit.group", + selection: "CP_ASYNC_COMMIT_GROUP", + ptx_modifier: "commit_group", + summary: "Commits this thread's uncommitted asynchronous copies as one group.", + }, + CpAsyncControlOperation::WaitAll => CpAsyncControlRecipe { + id: "cp_async_wait_all", + abi_id: "i0095", + operation_key: "memory.copy.async.group.wait_all", + rust_name: "cp_async_wait_all", + dialect_op_type: "CpAsyncWaitAllOp", + dialect_op_name: "nvvm.cp_async_wait_all", + source_record: "int_nvvm_cp_async_wait_all", + llvm_symbol: "llvm.nvvm.cp.async.wait.all", + selection: "CP_ASYNC_WAIT_ALL", + ptx_modifier: "wait_all", + summary: "Waits for all asynchronous copy groups issued by this thread.", + }, + CpAsyncControlOperation::WaitGroup => CpAsyncControlRecipe { + id: "cp_async_wait_group", + abi_id: "i0096", + operation_key: "memory.copy.async.group.wait_max_pending", + rust_name: "cp_async_wait_group", + dialect_op_type: "CpAsyncWaitGroupOp", + dialect_op_name: "nvvm.cp_async_wait_group", + source_record: "int_nvvm_cp_async_wait_group", + llvm_symbol: "llvm.nvvm.cp.async.wait.group", + selection: "CP_ASYNC_WAIT_GROUP", + ptx_modifier: "wait_group", + summary: "Waits until at most the compile-time constant number of recent asynchronous copy groups remain pending.", + }, + } +} + +pub(in crate::resolve) fn validate_cp_async_control_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, +) -> Result<()> { + let control = policy + .cp_async_control + .as_ref() + .with_context(|| format!("{} has no closed cp.async control contract", policy.id))?; + let recipe = cp_async_control_recipe(control.operation); + let has_operand = control.operation == CpAsyncControlOperation::WaitGroup; + let expected_adapter = if has_operand { + CpAsyncControlAdapter::CompileTimeConstantMaxPending + } else { + CpAsyncControlAdapter::NoOperands + }; + ensure!( + control.adapter == expected_adapter, + "{} cp.async control and adapter disagree", + policy.id + ); + ensure!( + control.runtime_validation == RuntimeValidation::Unexecuted, + "{} cannot claim unrecorded cp.async control runtime validation", + policy.id + ); + ensure!( + policy.id == recipe.id + && policy.abi_id == recipe.abi_id + && policy.operation_key == recipe.operation_key + && policy.source.is_none() + && policy.source_record.as_deref() == Some(recipe.source_record) + && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) + && policy.resolved_llvm_symbol.is_none(), + "{} cp.async control identity does not match its closed recipe", + policy.id + ); + let rust_arguments = if has_operand { vec!["u32"] } else { vec![] }; + let dialect_operands = if has_operand { vec!["i32"] } else { vec![] }; + ensure!( + policy.rust_module == "async_copy" + && policy.rust_name == recipe.rust_name + && policy.rust_arguments == rust_arguments + && policy.rust_result == "()" + && !policy.safe + && !policy.must_use + && policy.safe_allowlist_reason.is_none() + && policy.public_rust_path + == format!("cuda_intrinsics::async_copy::{}", recipe.rust_name) + && policy.compatibility_rust_paths + == [format!("cuda_device::async_copy::{}", recipe.rust_name)] + && policy.dialect_op_type == recipe.dialect_op_type + && policy.dialect_op_name == recipe.dialect_op_name + && policy.dialect_operands == dialect_operands + && policy.dialect_results.is_empty() + && policy.llvm_arguments == dialect_operands + && policy.llvm_results.is_empty() + && policy.lowering == "generated_cp_async_control", + "{} is outside the closed cp.async control API and lowering recipe", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == "read_write" + && !policy.convergent + && policy.execution_scope == "thread" + && policy.minimum_ptx == "7.0" + && policy.minimum_sm.as_deref() == Some("sm_80") + && policy.ptx_result == "()" + && policy.targets == "all", + "{} cp.async control effects or target floor disagree with the closed recipe", + policy.id + ); + let (ptx_isa_section, ptx_isa_url) = match control.operation { + CpAsyncControlOperation::CommitGroup => ( + "9.7.9.26.3.2 Data Movement and Conversion Instructions: cp.async.commit_group", + "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async-commit-group", + ), + CpAsyncControlOperation::WaitAll | CpAsyncControlOperation::WaitGroup => ( + "9.7.9.26.3.3 Data Movement and Conversion Instructions: cp.async.wait_group / cp.async.wait_all", + "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async-wait-group-cp-async-wait-all", + ), + }; + ensure!( + policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section == ptx_isa_section + && policy.ptx_isa_url == ptx_isa_url, + "{} cp.async control PTX provenance disagrees with the reviewed recipe", + policy.id + ); + ensure!( + policy.summary == recipe.summary, + "{} cp.async control summary does not match its closed recipe", + policy.id + ); + let expected_properties: Vec = if has_operand { + vec!["ImmArg".into()] + } else { + vec![] + }; + ensure!( + declaration.properties == expected_properties, + "{} cp.async control properties disagree with the imported declaration", + policy.id + ); + let operands = if has_operand { + vec![OperandPattern::Immediate] + } else { + vec![] + }; + ensure!( + policy.expected_ptx.mnemonic == "cp" + && policy.expected_ptx.modifiers == ["async", recipe.ptx_modifier] + && policy.expected_ptx.operands == operands + && declaration.selections.len() == 1 + && declaration.selections[0].source_record == recipe.selection, + "{} expected PTX or imported selection disagrees with its cp.async control recipe", + policy.id + ); + let backend_pairs: BTreeSet<_> = policy + .backend_lowerings + .iter() + .map(|lowering| (lowering.backend, lowering.mechanism)) + .collect(); + ensure!( + policy.backend_lowerings.len() == 2 + && backend_pairs + == BTreeSet::from([ + ( + IntrinsicBackend::LlvmNvptx, + BackendLoweringMechanism::TypedNvvm, + ), + ( + IntrinsicBackend::LibNvvm, + BackendLoweringMechanism::InlinePtx, + ), + ]) + && policy.backend_lowerings.iter().all(|lowering| { + lowering.minimum_ptx.is_none() + && lowering.minimum_sm.is_none() + && !lowering.evidence_profile.trim().is_empty() + }), + "{} must define the reviewed typed-LLVM and inline-PTX cp.async control routes", + policy.id + ); + ensure_no_other_family_contract(policy, "cp_async_control")?; + ensure!( + policy.ldmatrix_variant.is_none() + && policy.ldmatrix_safety.is_none() + && policy.ldmatrix_adapter.is_none() + && policy.selected_address_space.is_none(), + "{} mixes ldmatrix state with cp_async_control", + policy.id + ); + Ok(()) +} + +pub(in crate::resolve) struct CpAsyncMbarrierRecipe { + pub(in crate::resolve) id: &'static str, + pub(in crate::resolve) abi_id: &'static str, + pub(in crate::resolve) operation_key: &'static str, + pub(in crate::resolve) dialect_op_type: &'static str, + pub(in crate::resolve) dialect_op_name: &'static str, + pub(in crate::resolve) source_record: &'static str, + pub(in crate::resolve) llvm_symbol: &'static str, + pub(in crate::resolve) llvm_argument: &'static str, + pub(in crate::resolve) selection: &'static str, + pub(in crate::resolve) selection_asm: &'static str, + pub(in crate::resolve) modifiers: &'static [&'static str], + pub(in crate::resolve) summary: &'static str, +} + +pub(in crate::resolve) fn cp_async_mbarrier_recipe( + operation: CpAsyncMbarrierOperation, + state_space: CpAsyncMbarrierStateSpace, +) -> CpAsyncMbarrierRecipe { + match (operation, state_space) { + (CpAsyncMbarrierOperation::Arrive, CpAsyncMbarrierStateSpace::Generic) => { + CpAsyncMbarrierRecipe { + id: "cp_async_mbarrier_arrive", + abi_id: "i0101", + operation_key: "memory.copy.async.mbarrier.arrive.generic.increment", + dialect_op_type: "CpAsyncMbarrierArriveOp", + dialect_op_name: "nvvm.cp_async_mbarrier_arrive", + source_record: "int_nvvm_cp_async_mbarrier_arrive", + llvm_symbol: "llvm.nvvm.cp.async.mbarrier.arrive", + llvm_argument: "ptr", + selection: "CP_ASYNC_MBARRIER_ARRIVE", + selection_asm: "cp.async.mbarrier.arrive.b64 \t[$addr];", + modifiers: &["async", "mbarrier", "arrive", "b64"], + summary: "Associates this thread's prior asynchronous copies with a shared-memory barrier using balanced pending-count accounting.", + } + } + (CpAsyncMbarrierOperation::ArriveNoInc, CpAsyncMbarrierStateSpace::Generic) => { + CpAsyncMbarrierRecipe { + id: "cp_async_mbarrier_arrive_noinc", + abi_id: "i0103", + operation_key: "memory.copy.async.mbarrier.arrive.generic.no_increment", + dialect_op_type: "CpAsyncMbarrierArriveNoIncOp", + dialect_op_name: "nvvm.cp_async_mbarrier_arrive_noinc", + source_record: "int_nvvm_cp_async_mbarrier_arrive_noinc", + llvm_symbol: "llvm.nvvm.cp.async.mbarrier.arrive.noinc", + llvm_argument: "ptr", + selection: "CP_ASYNC_MBARRIER_ARRIVE_NOINC", + selection_asm: "cp.async.mbarrier.arrive.noinc.b64 \t[$addr];", + modifiers: &["async", "mbarrier", "arrive", "noinc", "b64"], + summary: "Associates this thread's prior asynchronous copies with a shared-memory barrier without incrementing its pending count.", + } + } + (CpAsyncMbarrierOperation::ArriveNoInc, CpAsyncMbarrierStateSpace::Shared) => { + CpAsyncMbarrierRecipe { + id: "cp_async_mbarrier_arrive_noinc_shared", + abi_id: "i0104", + operation_key: "memory.copy.async.mbarrier.arrive.shared.no_increment", + dialect_op_type: "CpAsyncMbarrierArriveNoIncSharedOp", + dialect_op_name: "nvvm.cp_async_mbarrier_arrive_noinc_shared", + source_record: "int_nvvm_cp_async_mbarrier_arrive_noinc_shared", + llvm_symbol: "llvm.nvvm.cp.async.mbarrier.arrive.noinc.shared", + llvm_argument: "shared_ptr", + selection: "CP_ASYNC_MBARRIER_ARRIVE_NOINC_SHARED", + selection_asm: "cp.async.mbarrier.arrive.noinc.shared.b64 \t[$addr];", + modifiers: &["async", "mbarrier", "arrive", "noinc", "shared", "b64"], + summary: "Associates this thread's prior asynchronous copies with a shared-address barrier without incrementing its pending count.", + } + } + (CpAsyncMbarrierOperation::Arrive, CpAsyncMbarrierStateSpace::Shared) => { + CpAsyncMbarrierRecipe { + id: "cp_async_mbarrier_arrive_shared", + abi_id: "i0102", + operation_key: "memory.copy.async.mbarrier.arrive.shared.increment", + dialect_op_type: "CpAsyncMbarrierArriveSharedOp", + dialect_op_name: "nvvm.cp_async_mbarrier_arrive_shared", + source_record: "int_nvvm_cp_async_mbarrier_arrive_shared", + llvm_symbol: "llvm.nvvm.cp.async.mbarrier.arrive.shared", + llvm_argument: "shared_ptr", + selection: "CP_ASYNC_MBARRIER_ARRIVE_SHARED", + selection_asm: "cp.async.mbarrier.arrive.shared.b64 \t[$addr];", + modifiers: &["async", "mbarrier", "arrive", "shared", "b64"], + summary: "Associates this thread's prior asynchronous copies with a shared-address barrier using balanced pending-count accounting.", + } + } + } +} + +pub(in crate::resolve) fn validate_cp_async_mbarrier_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, +) -> Result<()> { + let bridge = policy + .cp_async_mbarrier + .as_ref() + .with_context(|| format!("{} has no closed cp.async mbarrier contract", policy.id))?; + let recipe = cp_async_mbarrier_recipe(bridge.operation, bridge.state_space); + ensure!( + bridge.adapter == CpAsyncMbarrierAdapter::PointerToVoid, + "{} cp.async mbarrier adapter changed", + policy.id + ); + ensure!( + bridge.runtime_validation == RuntimeValidation::Unexecuted, + "{} cannot claim unrecorded cp.async mbarrier runtime validation", + policy.id + ); + ensure!( + policy.id == recipe.id + && policy.abi_id == recipe.abi_id + && policy.operation_key == recipe.operation_key + && policy.source.is_none() + && policy.source_record.as_deref() == Some(recipe.source_record) + && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) + && policy.resolved_llvm_symbol.is_none(), + "{} cp.async mbarrier identity does not match its closed recipe", + policy.id + ); + ensure!( + policy.rust_module == "async_copy" + && policy.rust_name == recipe.id + && policy.rust_arguments == ["*mut u64"] + && policy.rust_result == "()" + && !policy.safe + && !policy.must_use + && policy.safe_allowlist_reason.is_none() + && policy.public_rust_path == format!("cuda_intrinsics::async_copy::{}", recipe.id) + && policy.compatibility_rust_paths + == [format!("cuda_device::async_copy::{}", recipe.id)], + "{} is outside the closed cp.async mbarrier Rust API", + policy.id + ); + ensure!( + policy.dialect_op_type == recipe.dialect_op_type + && policy.dialect_op_name == recipe.dialect_op_name + && policy.dialect_operands == ["ptr"] + && policy.dialect_results.is_empty() + && policy.llvm_arguments == [recipe.llvm_argument] + && policy.llvm_results.is_empty() + && policy.lowering == "generated_cp_async_mbarrier", + "{} is outside the closed cp.async mbarrier signature and lowering recipe", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == "read_write" + && policy.convergent + && policy.execution_scope == "thread" + && policy.minimum_ptx == "7.0" + && policy.minimum_sm.as_deref() == Some("sm_80") + && policy.ptx_result == "()" + && policy.targets == "all", + "{} cp.async mbarrier effects or target floor disagree with the closed recipe", + policy.id + ); + ensure!( + policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section + == "9.7.14.16.18 Parallel Synchronization and Communication Instructions: cp.async.mbarrier.arrive" + && policy.ptx_isa_url + == "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-cp-async-mbarrier-arrive", + "{} cp.async mbarrier PTX provenance disagrees with the reviewed recipe", + policy.id + ); + ensure!( + policy.summary == recipe.summary, + "{} cp.async mbarrier summary does not match its closed recipe", + policy.id + ); + ensure!( + declaration.properties == ["IntrConvergent", "IntrNoCallback"], + "{} cp.async mbarrier properties disagree with the imported declaration", + policy.id + ); + ensure!( + policy.expected_ptx.mnemonic == "cp" + && policy + .expected_ptx + .modifiers + .iter() + .map(String::as_str) + .eq(recipe.modifiers.iter().copied()) + && policy.expected_ptx.operands == [OperandPattern::Address], + "{} expected PTX does not match its cp.async mbarrier recipe", + policy.id + ); + ensure!( + declaration.selections.len() == 1 + && declaration.selections[0].source_record == recipe.selection + && declaration.selections[0].asm == recipe.selection_asm + && declaration.selections[0].predicates + == [ + "Subtarget->getSmVersion() >= 80", + "Subtarget->getPTXVersion() >= 70", + ] + && declaration.selections[0].constraints.is_empty(), + "{} imported cp.async mbarrier selection changed", + policy.id + ); + let backend_pairs: BTreeSet<_> = policy + .backend_lowerings + .iter() + .map(|lowering| (lowering.backend, lowering.mechanism)) + .collect(); + ensure!( + policy.backend_lowerings.len() == 2 + && backend_pairs + == BTreeSet::from([ + ( + IntrinsicBackend::LlvmNvptx, + BackendLoweringMechanism::TypedNvvm, + ), + ( + IntrinsicBackend::LibNvvm, + BackendLoweringMechanism::InlinePtx, + ), + ]) + && policy.backend_lowerings.iter().all(|lowering| { + lowering.minimum_ptx.is_none() + && lowering.minimum_sm.is_none() + && !lowering.evidence_profile.trim().is_empty() + }), + "{} must define the reviewed typed-LLVM and inline-PTX cp.async mbarrier routes", + policy.id + ); + ensure_no_other_family_contract(policy, "cp_async_mbarrier")?; + Ok(()) +} + +pub(in crate::resolve) struct MbarrierBasicRecipe { + pub(in crate::resolve) id: &'static str, + pub(in crate::resolve) abi_id: &'static str, + pub(in crate::resolve) operation_key: &'static str, + pub(in crate::resolve) rust_arguments: &'static [&'static str], + pub(in crate::resolve) rust_result: &'static str, + pub(in crate::resolve) must_use: bool, + pub(in crate::resolve) adapter: MbarrierBasicAdapter, + pub(in crate::resolve) dialect_op_type: &'static str, + pub(in crate::resolve) dialect_op_name: &'static str, + pub(in crate::resolve) dialect_operands: &'static [&'static str], + pub(in crate::resolve) dialect_results: &'static [&'static str], + pub(in crate::resolve) source_record: &'static str, + pub(in crate::resolve) llvm_symbol: &'static str, + pub(in crate::resolve) llvm_arguments: &'static [&'static str], + pub(in crate::resolve) llvm_results: &'static [&'static str], + pub(in crate::resolve) memory: &'static str, + pub(in crate::resolve) ptx_result: &'static str, + pub(in crate::resolve) selection: &'static str, + pub(in crate::resolve) selection_asm: &'static str, + pub(in crate::resolve) ptx_modifier: &'static str, + pub(in crate::resolve) ptx_isa_section: &'static str, + pub(in crate::resolve) ptx_isa_url: &'static str, + pub(in crate::resolve) llvm_nvptx_mechanism: BackendLoweringMechanism, + pub(in crate::resolve) lib_nvvm_mechanism: BackendLoweringMechanism, + pub(in crate::resolve) properties: &'static [&'static str], + pub(in crate::resolve) summary: &'static str, +} + +pub(in crate::resolve) fn mbarrier_basic_recipe( + operation: MbarrierBasicOperation, +) -> MbarrierBasicRecipe { + match operation { + MbarrierBasicOperation::Init => MbarrierBasicRecipe { + id: "mbarrier_init", + abi_id: "i0097", + operation_key: "barrier.mbarrier.init.shared.cta", + rust_arguments: &["*mut u64", "u32"], + rust_result: "()", + must_use: false, + adapter: MbarrierBasicAdapter::InitPointerCountToVoid, + dialect_op_type: "MbarrierInitSharedOp", + dialect_op_name: "nvvm.mbarrier_init_shared", + dialect_operands: &["ptr", "i32"], + dialect_results: &[], + source_record: "int_nvvm_mbarrier_init_shared", + llvm_symbol: "llvm.nvvm.mbarrier.init.shared", + llvm_arguments: &["shared_ptr", "i32"], + llvm_results: &[], + memory: "read_write", + ptx_result: "()", + selection: "MBARRIER_INIT_SHARED", + selection_asm: "mbarrier.init.shared.b64 \t[$addr], $count;", + ptx_modifier: "init", + ptx_isa_section: "9.7.14.16.12 Parallel Synchronization and Communication Instructions: mbarrier.init", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-mbarrier-init", + llvm_nvptx_mechanism: BackendLoweringMechanism::TypedNvvm, + lib_nvvm_mechanism: BackendLoweringMechanism::InlinePtx, + properties: &["IntrConvergent", "IntrNoCallback"], + summary: "Initializes a CTA shared-memory barrier with the expected arrival count.", + }, + MbarrierBasicOperation::Arrive => MbarrierBasicRecipe { + id: "mbarrier_arrive", + abi_id: "i0098", + operation_key: "barrier.mbarrier.arrive.shared.cta", + rust_arguments: &["*const u64"], + rust_result: "u64", + must_use: true, + adapter: MbarrierBasicAdapter::ArrivePointerToToken, + dialect_op_type: "MbarrierArriveSharedOp", + dialect_op_name: "nvvm.mbarrier_arrive_shared", + dialect_operands: &["ptr"], + dialect_results: &["i64"], + source_record: "int_nvvm_mbarrier_arrive_shared", + llvm_symbol: "llvm.nvvm.mbarrier.arrive.shared", + llvm_arguments: &["shared_ptr"], + llvm_results: &["i64"], + memory: "read_write", + ptx_result: "u64", + selection: "MBARRIER_ARRIVE_SHARED", + selection_asm: "mbarrier.arrive.shared.b64 \t$state, [$addr];", + ptx_modifier: "arrive", + ptx_isa_section: "9.7.14.16.16 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", + llvm_nvptx_mechanism: BackendLoweringMechanism::TypedNvvm, + lib_nvvm_mechanism: BackendLoweringMechanism::InlinePtx, + properties: &["IntrConvergent", "IntrNoCallback"], + summary: "Arrives at a CTA shared-memory barrier and returns its phase token.", + }, + MbarrierBasicOperation::TestWait => MbarrierBasicRecipe { + id: "mbarrier_test_wait", + abi_id: "i0099", + operation_key: "barrier.mbarrier.test_wait.shared.cta", + rust_arguments: &["*const u64", "u64"], + rust_result: "bool", + must_use: true, + adapter: MbarrierBasicAdapter::TestWaitPointerTokenToPredicate, + dialect_op_type: "MbarrierTestWaitSharedOp", + dialect_op_name: "nvvm.mbarrier_test_wait_shared", + dialect_operands: &["ptr", "i64"], + dialect_results: &["i1"], + source_record: "int_nvvm_mbarrier_test_wait_shared", + llvm_symbol: "llvm.nvvm.mbarrier.test.wait.shared", + llvm_arguments: &["shared_ptr", "i64"], + llvm_results: &["i1"], + memory: "read_write", + ptx_result: "bool", + selection: "MBARRIER_TEST_WAIT_SHARED", + selection_asm: "mbarrier.test_wait.shared.b64 \t$res, [$addr], $state;", + ptx_modifier: "test_wait", + ptx_isa_section: "9.7.14.16.19 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", + llvm_nvptx_mechanism: BackendLoweringMechanism::InlinePtx, + lib_nvvm_mechanism: BackendLoweringMechanism::InlinePtx, + properties: &["IntrConvergent", "IntrNoCallback"], + summary: "Tests whether the CTA shared-memory barrier phase for a token is complete.", + }, + MbarrierBasicOperation::Inval => MbarrierBasicRecipe { + id: "mbarrier_inval", + abi_id: "i0100", + operation_key: "barrier.mbarrier.inval.shared.cta", + rust_arguments: &["*mut u64"], + rust_result: "()", + must_use: false, + adapter: MbarrierBasicAdapter::InvalPointerToVoid, + dialect_op_type: "MbarrierInvalSharedOp", + dialect_op_name: "nvvm.mbarrier_inval_shared", + dialect_operands: &["ptr"], + dialect_results: &[], + source_record: "int_nvvm_mbarrier_inval_shared", + llvm_symbol: "llvm.nvvm.mbarrier.inval.shared", + llvm_arguments: &["shared_ptr"], + llvm_results: &[], + memory: "write", + ptx_result: "()", + selection: "MBARRIER_INVAL_SHARED", + selection_asm: "mbarrier.inval.shared.b64 \t[$addr];", + ptx_modifier: "inval", + ptx_isa_section: "9.7.14.16.13 Parallel Synchronization and Communication Instructions: mbarrier.inval", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-mbarrier-inval", + llvm_nvptx_mechanism: BackendLoweringMechanism::TypedNvvm, + lib_nvvm_mechanism: BackendLoweringMechanism::InlinePtx, + properties: &[ + "IntrArgMemOnly", + "IntrConvergent", + "IntrNoCallback", + "IntrWriteMem", + "NoCapture", + "WriteOnly", + ], + summary: "Invalidates a CTA shared-memory barrier after its users have finished.", + }, + } +} + +pub(in crate::resolve) fn mbarrier_expected_operands( + operation: MbarrierBasicOperation, +) -> Vec { + match operation { + MbarrierBasicOperation::Init => { + vec![OperandPattern::Address, OperandPattern::RegisterOrImmediate] + } + MbarrierBasicOperation::Arrive => { + vec![OperandPattern::Register, OperandPattern::Address] + } + MbarrierBasicOperation::TestWait => vec![ + OperandPattern::Register, + OperandPattern::Address, + OperandPattern::Register, + ], + MbarrierBasicOperation::Inval => vec![OperandPattern::Address], + } +} + +pub(in crate::resolve) fn validate_mbarrier_basic_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, +) -> Result<()> { + let mbarrier = policy + .mbarrier_basic + .as_ref() + .with_context(|| format!("{} has no closed basic mbarrier contract", policy.id))?; + let recipe = mbarrier_basic_recipe(mbarrier.operation); + ensure!( + mbarrier.state_space == MbarrierStateSpace::Shared && mbarrier.adapter == recipe.adapter, + "{} mbarrier operation, state space, and adapter disagree", + policy.id + ); + ensure!( + mbarrier.runtime_validation == RuntimeValidation::Unexecuted, + "{} cannot claim unrecorded mbarrier runtime validation", + policy.id + ); + ensure!( + policy.id == recipe.id + && policy.abi_id == recipe.abi_id + && policy.operation_key == recipe.operation_key + && policy.source.is_none() + && policy.source_record.as_deref() == Some(recipe.source_record) + && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) + && policy.resolved_llvm_symbol.is_none(), + "{} mbarrier identity does not match its closed recipe", + policy.id + ); + ensure!( + policy.rust_module == "barrier" + && policy.rust_name == recipe.id + && policy + .rust_arguments + .iter() + .map(String::as_str) + .eq(recipe.rust_arguments.iter().copied()) + && policy.rust_result == recipe.rust_result + && !policy.safe + && policy.must_use == recipe.must_use + && policy.safe_allowlist_reason.is_none() + && policy.public_rust_path == format!("cuda_intrinsics::barrier::{}", recipe.id) + && policy.compatibility_rust_paths == [format!("cuda_device::barrier::{}", recipe.id)], + "{} must preserve its unsafe mbarrier raw and compatibility API", + policy.id + ); + ensure!( + policy.dialect_op_type == recipe.dialect_op_type + && policy.dialect_op_name == recipe.dialect_op_name + && policy + .dialect_operands + .iter() + .map(String::as_str) + .eq(recipe.dialect_operands.iter().copied()) + && policy + .dialect_results + .iter() + .map(String::as_str) + .eq(recipe.dialect_results.iter().copied()) + && policy + .llvm_arguments + .iter() + .map(String::as_str) + .eq(recipe.llvm_arguments.iter().copied()) + && policy + .llvm_results + .iter() + .map(String::as_str) + .eq(recipe.llvm_results.iter().copied()) + && policy.lowering == "generated_mbarrier_basic", + "{} is outside the closed mbarrier signature and lowering recipe", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == recipe.memory + && policy.convergent + && policy.execution_scope == "cta" + && policy.minimum_ptx == "7.0" + && policy.minimum_sm.as_deref() == Some("sm_80") + && policy.ptx_result == recipe.ptx_result + && policy.targets == "all", + "{} mbarrier effects or target floor disagree with the closed recipe", + policy.id + ); + ensure!( + policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section == recipe.ptx_isa_section + && policy.ptx_isa_url == recipe.ptx_isa_url, + "{} mbarrier PTX provenance disagrees with the reviewed recipe", + policy.id + ); + ensure!( + policy.summary == recipe.summary, + "{} mbarrier summary does not match its closed recipe", + policy.id + ); + ensure!( + declaration + .properties + .iter() + .map(String::as_str) + .eq(recipe.properties.iter().copied()), + "{} mbarrier properties disagree with the imported LLVM declaration", + policy.id + ); + ensure!( + policy.expected_ptx.mnemonic == "mbarrier" + && policy.expected_ptx.modifiers == [recipe.ptx_modifier, "shared", "b64"] + && policy.expected_ptx.operands == mbarrier_expected_operands(mbarrier.operation), + "{} expected PTX does not match its mbarrier recipe", + policy.id + ); + ensure!( + declaration.selections.len() == 1 + && declaration.selections[0].source_record == recipe.selection + && declaration.selections[0].asm == recipe.selection_asm + && declaration.selections[0].predicates + == [ + "Subtarget->getSmVersion() >= 80", + "Subtarget->getPTXVersion() >= 70", + ] + && declaration.selections[0].constraints.is_empty(), + "{} imported mbarrier selection changed", + policy.id + ); + let backend_pairs: BTreeSet<_> = policy + .backend_lowerings + .iter() + .map(|lowering| (lowering.backend, lowering.mechanism)) + .collect(); + ensure!( + policy.backend_lowerings.len() == 2 + && backend_pairs + == BTreeSet::from([ + (IntrinsicBackend::LlvmNvptx, recipe.llvm_nvptx_mechanism), + (IntrinsicBackend::LibNvvm, recipe.lib_nvvm_mechanism), + ]) + && policy.backend_lowerings.iter().all(|lowering| { + lowering.minimum_ptx.is_none() + && lowering.minimum_sm.is_none() + && !lowering.evidence_profile.trim().is_empty() + }), + "{} must define exactly the reviewed mbarrier backend routes", + policy.id + ); + ensure_no_other_family_contract(policy, "mbarrier_basic")?; + Ok(()) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/debug_wgmma.rs b/crates/cuda-intrinsics-gen/src/resolve/families/debug_wgmma.rs new file mode 100644 index 0000000000..2e61a3d18c --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/debug_wgmma.rs @@ -0,0 +1,708 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, DebugControl, DebugControlAdapter, DebugControlAdmission, + DebugControlOperation, ImportedIntrinsic, IntrinsicBackend, IntrinsicSource, + OverlayBackendLowering, OverlayIntrinsic, RuntimeValidation, WgmmaControl, WgmmaControlAdapter, + WgmmaControlAdmission, WgmmaControlMode, WgmmaControlParticipation, +}; +use crate::ptx::{InstructionPattern, OperandPattern}; +use anyhow::{Context, Result, ensure}; +use std::collections::BTreeSet; + +use crate::resolve::abi_ledger::*; +use crate::resolve::guards::*; + +#[derive(Clone, Copy)] +pub(in crate::resolve) struct DebugControlRecipe { + id: &'static str, + operation_key: &'static str, + rust_name: &'static str, + rust_arguments: &'static [&'static str], + rust_result: &'static str, + compatibility_path: &'static str, + op_type: &'static str, + op_name: &'static str, + instruction: &'static str, + minimum_ptx: &'static str, + minimum_sm: Option<&'static str>, + section: &'static str, + anchor: &'static str, + adapter: DebugControlAdapter, + summary: &'static str, +} + +pub(in crate::resolve) fn debug_control_recipe( + operation: DebugControlOperation, +) -> DebugControlRecipe { + match operation { + DebugControlOperation::Trap => DebugControlRecipe { + id: "trap", + operation_key: "debug.control.trap", + rust_name: "trap", + rust_arguments: &[], + rust_result: "!", + compatibility_path: "cuda_device::debug::trap", + op_type: "TrapOp", + op_name: "nvvm.trap", + instruction: "trap", + minimum_ptx: "1.0", + minimum_sm: None, + section: "9.7.20.4 Miscellaneous Instructions: trap", + anchor: "miscellaneous-instructions-trap", + adapter: DebugControlAdapter::Direct, + summary: "Aborts device execution and reports an interrupt to the host.", + }, + DebugControlOperation::Breakpoint => DebugControlRecipe { + id: "breakpoint", + operation_key: "debug.control.breakpoint", + rust_name: "breakpoint", + rust_arguments: &[], + rust_result: "()", + compatibility_path: "cuda_device::debug::breakpoint", + op_type: "BreakpointOp", + op_name: "nvvm.brkpt", + instruction: "brkpt", + minimum_ptx: "1.0", + minimum_sm: Some("sm_11"), + section: "9.7.20.1 Miscellaneous Instructions: brkpt", + anchor: "miscellaneous-instructions-brkpt", + adapter: DebugControlAdapter::Direct, + summary: "Suspends device execution for a debugger breakpoint.", + }, + DebugControlOperation::Pmevent => DebugControlRecipe { + id: "pmevent", + operation_key: "debug.profiler.event", + rust_name: "pmevent", + rust_arguments: &["u32"], + rust_result: "()", + compatibility_path: "cuda_device::debug::__prof_trigger", + op_type: "PmEventOp", + op_name: "nvvm.pmevent", + instruction: "pmevent", + minimum_ptx: "1.4", + minimum_sm: None, + section: "9.7.20.3 Miscellaneous Instructions: pmevent", + anchor: "miscellaneous-instructions-pmevent", + adapter: DebugControlAdapter::ConstGenericToImmediateU32, + summary: "Triggers one compile-time-selected performance monitor event.", + }, + } +} + +#[derive(Clone, Copy)] +pub(in crate::resolve) struct WgmmaControlRecipe { + pub(in crate::resolve) mode: WgmmaControlMode, + pub(in crate::resolve) abi_id: &'static str, + pub(in crate::resolve) id: &'static str, + pub(in crate::resolve) operation_key: &'static str, + pub(in crate::resolve) source_record: &'static str, + pub(in crate::resolve) selection_record: &'static str, + pub(in crate::resolve) llvm_symbol: &'static str, + pub(in crate::resolve) compatibility_path: &'static str, + pub(in crate::resolve) dialect_op_type: &'static str, + pub(in crate::resolve) dialect_op_name: &'static str, + pub(in crate::resolve) suffix: &'static str, + pub(in crate::resolve) adapter: WgmmaControlAdapter, + pub(in crate::resolve) summary: &'static str, +} + +pub(in crate::resolve) fn wgmma_control_recipe(mode: WgmmaControlMode) -> WgmmaControlRecipe { + match mode { + WgmmaControlMode::Fence => WgmmaControlRecipe { + mode, + abi_id: "i0317", + id: "wgmma_fence", + operation_key: "wgmma.control.fence.sync.aligned", + source_record: "int_nvvm_wgmma_fence_sync_aligned", + selection_record: "WGMMA_FENCE_SYNC_ALIGNED", + llvm_symbol: "llvm.nvvm.wgmma.fence.sync.aligned", + compatibility_path: "cuda_device::wgmma::wgmma_fence", + dialect_op_type: "WgmmaFenceSyncAlignedOp", + dialect_op_name: "nvvm.wgmma_fence_sync_aligned", + suffix: "fence.sync.aligned", + adapter: WgmmaControlAdapter::NoArguments, + summary: "Orders register accesses before later WGMMA operations.", + }, + WgmmaControlMode::CommitGroup => WgmmaControlRecipe { + mode, + abi_id: "i0318", + id: "wgmma_commit_group", + operation_key: "wgmma.control.commit_group.sync.aligned", + source_record: "int_nvvm_wgmma_commit_group_sync_aligned", + selection_record: "WGMMA_COMMIT_GROUP_SYNC_ALIGNED", + llvm_symbol: "llvm.nvvm.wgmma.commit_group.sync.aligned", + compatibility_path: "cuda_device::wgmma::wgmma_commit_group", + dialect_op_type: "WgmmaCommitGroupSyncAlignedOp", + dialect_op_name: "nvvm.wgmma_commit_group_sync_aligned", + suffix: "commit_group.sync.aligned", + adapter: WgmmaControlAdapter::NoArguments, + summary: "Commits prior uncommitted WGMMA operations as one group.", + }, + WgmmaControlMode::WaitGroup => WgmmaControlRecipe { + mode, + abi_id: "i0319", + id: "wgmma_wait_group", + operation_key: "wgmma.control.wait_group.sync.aligned", + source_record: "int_nvvm_wgmma_wait_group_sync_aligned", + selection_record: "WGMMA_WAIT_GROUP_SYNC_ALIGNED", + llvm_symbol: "llvm.nvvm.wgmma.wait_group.sync.aligned", + compatibility_path: "cuda_device::wgmma::__wgmma_wait_group", + dialect_op_type: "WgmmaWaitGroupSyncAlignedOp", + dialect_op_name: "nvvm.wgmma_wait_group_sync_aligned", + suffix: "wait_group.sync.aligned", + adapter: WgmmaControlAdapter::ConstGenericU32ToI64Immediate, + summary: "Waits until at most the requested number of WGMMA groups remain pending.", + }, + } +} + +pub(in crate::resolve) fn expand_debug_control_admission( + admission: &DebugControlAdmission, +) -> Result> { + ensure!( + admission.runtime_validation == RuntimeValidation::Unexecuted, + "debug-control runtime validation may be marked executed only with GPU evidence" + ); + ensure!( + !admission.llvm_evidence_profile.trim().is_empty() + && !admission.libnvvm_evidence_profile.trim().is_empty(), + "compact debug-control admission requires both backend evidence profiles" + ); + ensure!( + admission.operations + == [ + DebugControlOperation::Trap, + DebugControlOperation::Breakpoint, + DebugControlOperation::Pmevent, + ], + "compact debug-control admission must list trap, breakpoint, and pmevent exactly once in canonical order" + ); + ensure!( + admission.abi_ids.len() == admission.operations.len(), + "pending debug-control admission needs exactly three ABI IDs before aggregation" + ); + let unique_abi_ids = admission.abi_ids.iter().collect::>(); + ensure!( + unique_abi_ids.len() == admission.abi_ids.len(), + "debug-control ABI IDs must be unique" + ); + + admission + .operations + .iter() + .zip(&admission.abi_ids) + .map(|(&operation, abi_id)| { + validate_abi_id(abi_id)?; + let recipe = debug_control_recipe(operation); + let immediate = operation == DebugControlOperation::Pmevent; + Ok(OverlayIntrinsic { + id: recipe.id.into(), + abi_id: abi_id.clone(), + operation_key: recipe.operation_key.into(), + family: "debug_control".into(), + source: Some(IntrinsicSource::PtxNative { + instruction: recipe.instruction.into(), + }), + source_record: None, + rust_module: "debug".into(), + rust_name: recipe.rust_name.into(), + rust_arguments: recipe + .rust_arguments + .iter() + .map(|value| (*value).into()) + .collect(), + rust_result: recipe.rust_result.into(), + safe: true, + must_use: false, + safe_allowlist_reason: Some( + match operation { + DebugControlOperation::Trap => { + "aborting the kernel has no memory-safety preconditions" + } + DebugControlOperation::Breakpoint => { + "requesting a debugger breakpoint has no memory-safety preconditions" + } + DebugControlOperation::Pmevent => { + "the importer accepts only the documented immediate event range" + } + } + .into(), + ), + public_rust_path: format!("cuda_intrinsics::debug::{}", recipe.rust_name), + compatibility_rust_paths: vec![recipe.compatibility_path.into()], + dialect_op_type: recipe.op_type.into(), + dialect_op_name: recipe.op_name.into(), + dialect_operands: vec![], + dialect_results: vec![], + llvm_symbol: None, + resolved_llvm_symbol: None, + llvm_arguments: vec![], + llvm_results: vec![], + pure: false, + memory: "none".into(), + convergent: false, + execution_scope: "thread".into(), + minimum_ptx: recipe.minimum_ptx.into(), + minimum_sm: recipe.minimum_sm.map(Into::into), + ptx_result: "()".into(), + targets: "all".into(), + ptx_isa_version: "9.3".into(), + ptx_isa_section: recipe.section.into(), + ptx_isa_url: format!( + "https://docs.nvidia.com/cuda/parallel-thread-execution/#{}", + recipe.anchor + ), + lowering: "generated_debug_control".into(), + backend_lowerings: vec![ + OverlayBackendLowering { + backend: IntrinsicBackend::LlvmNvptx, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: admission.llvm_evidence_profile.clone(), + targets: None, + minimum_ptx: Some("3.2".into()), + minimum_sm: Some("sm_20".into()), + }, + OverlayBackendLowering { + backend: IntrinsicBackend::LibNvvm, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: admission.libnvvm_evidence_profile.clone(), + targets: None, + // PTX floor is the inline-PTX mechanism floor, same as + // the LlvmNvptx entry above; the hardware floor stays at + // the probed sm_75 (backend-codegen evidence must sit + // exactly at the hardware floor, and CUDA 13 cannot + // probe older targets). Writing the probe PTX version + // (9.3) here made every panic path unbuildable on the + // NVVM-IR route: the floor exceeded the newest PTX + // version cuda-oxide can request. + minimum_ptx: Some("3.2".into()), + minimum_sm: Some("sm_75".into()), + }, + ], + packed_atomic: None, + redux: None, + vote: None, + active_mask: None, + warp_match: None, + warp_barrier: None, + warp_shuffle: None, + dot_product: None, + packed_alu: None, + integer_minmax: None, + packed_conversion: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cp_async_copy: None, + cp_async_control: None, + cp_async_mbarrier: None, + mbarrier_basic: None, + movmatrix: None, + mbarrier_extended: None, + register_mma: None, + sparse_mma: None, + prmt: None, + cluster_barrier: None, + wgmma_control: None, + special_register: None, + debug_control: Some(DebugControl { + operation, + adapter: recipe.adapter, + runtime_validation: admission.runtime_validation, + }), + cluster_memory: None, + clc: None, + tma: None, + tcgen05: None, + ldmatrix_variant: None, + ldmatrix_safety: None, + ldmatrix_adapter: None, + selected_address_space: None, + expected_ptx: InstructionPattern { + mnemonic: recipe.instruction.into(), + modifiers: vec![], + operands: if immediate { + vec![OperandPattern::Immediate] + } else { + vec![] + }, + }, + summary: recipe.summary.into(), + }) + }) + .collect() +} + +pub(in crate::resolve) fn validate_debug_control_policy( + policy: &OverlayIntrinsic, + source: &IntrinsicSource, +) -> Result<()> { + let debug = policy + .debug_control + .as_ref() + .with_context(|| format!("{} has no closed debug-control contract", policy.id))?; + let recipe = debug_control_recipe(debug.operation); + let immediate = debug.operation == DebugControlOperation::Pmevent; + ensure!( + debug.adapter == recipe.adapter + && debug.runtime_validation == RuntimeValidation::Unexecuted + && policy.id == recipe.id + && policy.operation_key == recipe.operation_key + && source + == &IntrinsicSource::PtxNative { + instruction: recipe.instruction.into(), + } + && policy.source_record.is_none() + && policy.llvm_symbol.is_none() + && policy.resolved_llvm_symbol.is_none() + && policy.llvm_arguments.is_empty() + && policy.llvm_results.is_empty(), + "{} identity must remain PTX-native and match its closed debug-control recipe", + policy.id + ); + ensure!( + policy.rust_module == "debug" + && policy.rust_name == recipe.rust_name + && policy.rust_arguments == recipe.rust_arguments + && policy.rust_result == recipe.rust_result + && policy.safe + && !policy.must_use + && policy.public_rust_path == format!("cuda_intrinsics::debug::{}", recipe.rust_name) + && policy.compatibility_rust_paths == [recipe.compatibility_path], + "{} Rust API or compatibility adapter changed", + policy.id + ); + ensure!( + policy.dialect_op_type == recipe.op_type + && policy.dialect_op_name == recipe.op_name + && policy.dialect_operands.is_empty() + && policy.dialect_results.is_empty() + && policy.lowering == "generated_debug_control", + "{} dialect carrier or lowering changed", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == "none" + && !policy.convergent + && policy.execution_scope == "thread", + "{} debug-control effects changed", + policy.id + ); + ensure!( + policy.minimum_ptx == recipe.minimum_ptx + && policy.minimum_sm.as_deref() == recipe.minimum_sm + && policy.targets == "all" + && policy.ptx_result == "()" + && policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section == recipe.section + && policy.ptx_isa_url + == format!( + "https://docs.nvidia.com/cuda/parallel-thread-execution/#{}", + recipe.anchor + ), + "{} native target floor or PTX provenance changed", + policy.id + ); + ensure!( + policy.expected_ptx.mnemonic == recipe.instruction + && policy.expected_ptx.modifiers.is_empty() + && policy.expected_ptx.operands + == if immediate { + vec![OperandPattern::Immediate] + } else { + vec![] + }, + "{} expected PTX changed", + policy.id + ); + ensure_exact_inline_ptx_backends( + policy, + [ + (IntrinsicBackend::LlvmNvptx, "3.2", Some("sm_20")), + (IntrinsicBackend::LibNvvm, "3.2", Some("sm_75")), + ], + "debug-control", + )?; + ensure_no_other_family_contract(policy, "debug-control")?; + Ok(()) +} + +pub(in crate::resolve) fn expand_wgmma_control_admission( + admission: &WgmmaControlAdmission, +) -> Result> { + ensure!( + admission.runtime_validation == RuntimeValidation::Unexecuted, + "WGMMA-control runtime validation may be marked executed only with GPU evidence" + ); + ensure!( + !admission.llvm_evidence_profile.trim().is_empty() + && !admission.libnvvm_evidence_profile.trim().is_empty(), + "compact WGMMA-control admission requires both backend evidence profiles" + ); + let expected_modes = [ + WgmmaControlMode::Fence, + WgmmaControlMode::CommitGroup, + WgmmaControlMode::WaitGroup, + ]; + let actual_modes = admission + .variants + .iter() + .map(|variant| variant.mode) + .collect::>(); + ensure!( + actual_modes == expected_modes, + "compact WGMMA-control admission must contain each reviewed mode once in canonical order" + ); + + admission + .variants + .iter() + .map(|variant| { + let recipe = wgmma_control_recipe(variant.mode); + ensure!( + variant.abi_id == recipe.abi_id, + "{} must keep reserved ABI ID {}", + recipe.id, + recipe.abi_id + ); + let wait = recipe.mode == WgmmaControlMode::WaitGroup; + Ok(OverlayIntrinsic { + id: recipe.id.into(), + abi_id: variant.abi_id.clone(), + operation_key: recipe.operation_key.into(), + family: "wgmma_control".into(), + source: None, + source_record: Some(recipe.source_record.into()), + rust_module: "wgmma".into(), + rust_name: recipe.id.into(), + rust_arguments: if wait { vec!["u64".into()] } else { vec![] }, + rust_result: "()".into(), + safe: false, + must_use: false, + safe_allowlist_reason: None, + public_rust_path: format!("cuda_intrinsics::wgmma::{}", recipe.id), + compatibility_rust_paths: vec![recipe.compatibility_path.into()], + dialect_op_type: recipe.dialect_op_type.into(), + dialect_op_name: recipe.dialect_op_name.into(), + dialect_operands: if wait { vec!["i64".into()] } else { vec![] }, + dialect_results: vec![], + llvm_symbol: Some(recipe.llvm_symbol.into()), + resolved_llvm_symbol: None, + llvm_arguments: if wait { vec!["i64".into()] } else { vec![] }, + llvm_results: vec![], + pure: false, + memory: "read_write".into(), + convergent: true, + execution_scope: "warpgroup".into(), + minimum_ptx: "8.0".into(), + minimum_sm: None, + ptx_result: "()".into(), + targets: "sm_90a".into(), + ptx_isa_version: "9.3".into(), + ptx_isa_section: + "Asynchronous Warpgroup Level Matrix Instructions: WGMMA control".into(), + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#asynchronous-warpgroup-level-matrix-instructions".into(), + lowering: "generated_wgmma_control".into(), + backend_lowerings: vec![ + OverlayBackendLowering { + backend: IntrinsicBackend::LlvmNvptx, + mechanism: BackendLoweringMechanism::TypedNvvm, + evidence_profile: admission.llvm_evidence_profile.clone(), + targets: None, + minimum_ptx: Some("8.0".into()), + minimum_sm: None, + }, + OverlayBackendLowering { + backend: IntrinsicBackend::LibNvvm, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: admission.libnvvm_evidence_profile.clone(), + targets: None, + minimum_ptx: Some("8.0".into()), + minimum_sm: None, + }, + ], + packed_atomic: None, + redux: None, + vote: None, + active_mask: None, + warp_match: None, + warp_barrier: None, + warp_shuffle: None, + dot_product: None, + packed_alu: None, + integer_minmax: None, + packed_conversion: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cp_async_copy: None, + cp_async_control: None, + cp_async_mbarrier: None, + mbarrier_basic: None, + movmatrix: None, + mbarrier_extended: None, + register_mma: None, + sparse_mma: None, + prmt: None, + cluster_barrier: None, + wgmma_control: Some(WgmmaControl { + mode: recipe.mode, + adapter: recipe.adapter, + participation: WgmmaControlParticipation::WarpgroupAllThreadsSameInstruction, + }), + special_register: None, + debug_control: None, + cluster_memory: None, + clc: None, + tma: None, + tcgen05: None, + ldmatrix_variant: None, + ldmatrix_safety: None, + ldmatrix_adapter: None, + selected_address_space: None, + expected_ptx: InstructionPattern { + mnemonic: "wgmma".into(), + modifiers: recipe.suffix.split('.').map(str::to_owned).collect(), + operands: if wait { + vec![OperandPattern::Immediate] + } else { + vec![] + }, + }, + summary: recipe.summary.into(), + }) + }) + .collect() +} + +pub(in crate::resolve) fn validate_wgmma_control_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, +) -> Result<()> { + let control = policy + .wgmma_control + .as_ref() + .with_context(|| format!("{} has no closed WGMMA-control contract", policy.id))?; + let recipe = wgmma_control_recipe(control.mode); + let wait = recipe.mode == WgmmaControlMode::WaitGroup; + ensure!( + control.adapter == recipe.adapter + && control.participation + == WgmmaControlParticipation::WarpgroupAllThreadsSameInstruction + && policy.id == recipe.id + && policy.abi_id == recipe.abi_id + && policy.operation_key == recipe.operation_key + && policy.source.is_none() + && policy.source_record.as_deref() == Some(recipe.source_record) + && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) + && policy.resolved_llvm_symbol.is_none(), + "{} identity or semantics do not match its closed WGMMA-control recipe", + policy.id + ); + ensure!( + policy.rust_module == "wgmma" + && policy.rust_name == recipe.id + && policy.rust_arguments + == if wait { + vec!["u64"] + } else { + Vec::<&str>::new() + } + && policy.rust_result == "()" + && !policy.safe + && !policy.must_use + && policy.public_rust_path == format!("cuda_intrinsics::wgmma::{}", recipe.id) + && policy.compatibility_rust_paths == [recipe.compatibility_path], + "{} Rust API does not match its closed WGMMA-control recipe", + policy.id + ); + let expected_arguments = if wait { + vec!["i64"] + } else { + Vec::<&str>::new() + }; + ensure!( + policy.dialect_op_type == recipe.dialect_op_type + && policy.dialect_op_name == recipe.dialect_op_name + && policy.dialect_operands == expected_arguments + && policy.dialect_results.is_empty() + && policy.llvm_arguments == expected_arguments + && policy.llvm_results.is_empty() + && policy.lowering == "generated_wgmma_control", + "{} carrier or lowering does not match its closed WGMMA-control recipe", + policy.id + ); + let expected_properties = if wait { + vec!["ImmArg", "IntrConvergent"] + } else { + vec!["IntrConvergent"] + }; + ensure!( + declaration.classes == ["SDPatternOperator", "Intrinsic"] + && declaration.properties == expected_properties + && !policy.pure + && policy.memory == "read_write" + && policy.convergent + && policy.execution_scope == "warpgroup", + "{} effects disagree with the imported WGMMA-control declaration", + policy.id + ); + ensure!( + policy.minimum_ptx == "8.0" + && policy.minimum_sm.is_none() + && policy.targets == "sm_90a" + && policy.ptx_result == "()" + && policy.ptx_isa_version == "9.3" + && policy.ptx_isa_url + == "https://docs.nvidia.com/cuda/parallel-thread-execution/#asynchronous-warpgroup-level-matrix-instructions", + "{} target floor or PTX provenance changed", + policy.id + ); + ensure!( + policy.expected_ptx.mnemonic == "wgmma" + && policy.expected_ptx.modifiers == recipe.suffix.split('.').collect::>() + && policy.expected_ptx.operands + == if wait { + vec![OperandPattern::Immediate] + } else { + vec![] + }, + "{} expected PTX does not match its exact WGMMA-control spelling", + policy.id + ); + let backend_pairs: BTreeSet<_> = policy + .backend_lowerings + .iter() + .map(|lowering| (lowering.backend, lowering.mechanism)) + .collect(); + ensure!( + policy.backend_lowerings.len() == 2 + && backend_pairs + == BTreeSet::from([ + ( + IntrinsicBackend::LlvmNvptx, + BackendLoweringMechanism::TypedNvvm, + ), + ( + IntrinsicBackend::LibNvvm, + BackendLoweringMechanism::InlinePtx, + ), + ]) + && policy.backend_lowerings.iter().all(|lowering| { + lowering.minimum_ptx.as_deref() == Some("8.0") + && lowering.minimum_sm.is_none() + && !lowering.evidence_profile.trim().is_empty() + }), + "{} must define exactly the reviewed WGMMA-control backend routes", + policy.id + ); + ensure_no_other_family_contract(policy, "WGMMA control")?; + Ok(()) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/execution_control.rs b/crates/cuda-intrinsics-gen/src/resolve/families/execution_control.rs new file mode 100644 index 0000000000..70c07709cc --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/execution_control.rs @@ -0,0 +1,633 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, ExecutionControlOperation, ImportedIntrinsic, IntrinsicBackend, + OverlayIntrinsic, +}; +use crate::ptx::OperandPattern; +use anyhow::{Context, Result, ensure}; +use std::collections::BTreeSet; + +use super::*; +use crate::resolve::guards::*; + +pub(in crate::resolve) const EXECUTION_CONTROL_OPERATIONS: [ExecutionControlOperation; 8] = [ + ExecutionControlOperation::BarrierCtaSync, + ExecutionControlOperation::BarrierCtaSyncAligned, + ExecutionControlOperation::BarrierCtaArrive, + ExecutionControlOperation::BarrierCtaArriveAligned, + ExecutionControlOperation::GridDependencyLaunchDependents, + ExecutionControlOperation::GridDependencyWait, + ExecutionControlOperation::SetMaxNRegInc, + ExecutionControlOperation::SetMaxNRegDec, +]; + +pub(in crate::resolve) fn validate_execution_control_family_completeness( + intrinsics: &[OverlayIntrinsic], +) -> Result<()> { + let has_execution_control_family = intrinsics.iter().any(|policy| { + matches!( + policy.family.as_str(), + "counted_barrier" | "grid_dependency" | "register_control" + ) + }); + if !has_execution_control_family { + return Ok(()); + } + let actual = intrinsics + .iter() + .filter_map(|policy| ExecutionControlOperation::from_catalog_id(&policy.id)) + .collect::>(); + let expected = EXECUTION_CONTROL_OPERATIONS + .into_iter() + .collect::>(); + ensure!( + actual == expected, + "execution-control overlay must admit all four counted barriers, both grid-dependency controls, and both setmaxnreg operations" + ); + for family in ["counted_barrier", "grid_dependency", "register_control"] { + ensure!( + intrinsics + .iter() + .filter(|policy| policy.family == family) + .all(|policy| { + ExecutionControlOperation::from_catalog_id(&policy.id) + .is_some_and(|operation| operation.family() == family) + }), + "{family} overlay contains an operation outside its closed instruction family" + ); + } + Ok(()) +} + +#[derive(Clone)] +pub(in crate::resolve) struct ExecutionControlRecipe { + pub(in crate::resolve) abi_id: &'static str, + pub(in crate::resolve) id: &'static str, + pub(in crate::resolve) operation_key: &'static str, + pub(in crate::resolve) source_record: &'static str, + pub(in crate::resolve) llvm_symbol: &'static str, + pub(in crate::resolve) rust_module: &'static str, + pub(in crate::resolve) rust_arguments: &'static [&'static str], + pub(in crate::resolve) compatibility_path: &'static str, + pub(in crate::resolve) dialect_op_type: &'static str, + pub(in crate::resolve) dialect_op_name: &'static str, + pub(in crate::resolve) dialect_operands: &'static [&'static str], + pub(in crate::resolve) llvm_arguments: &'static [&'static str], + pub(in crate::resolve) imported_classes: &'static [&'static str], + pub(in crate::resolve) imported_properties: &'static [&'static str], + pub(in crate::resolve) memory: &'static str, + pub(in crate::resolve) convergent: bool, + pub(in crate::resolve) execution_scope: &'static str, + pub(in crate::resolve) minimum_ptx: &'static str, + pub(in crate::resolve) minimum_sm: Option<&'static str>, + pub(in crate::resolve) targets: &'static str, + pub(in crate::resolve) ptx_isa_section: &'static str, + pub(in crate::resolve) ptx_isa_url: &'static str, + pub(in crate::resolve) mnemonic: &'static str, + pub(in crate::resolve) modifiers: &'static [&'static str], + pub(in crate::resolve) operands: &'static [OperandPattern], + pub(in crate::resolve) selection_records: &'static [&'static str], + pub(in crate::resolve) selection_asm: &'static str, + pub(in crate::resolve) selection_predicates: &'static [&'static str], + pub(in crate::resolve) summary: &'static str, +} + +pub(in crate::resolve) fn execution_control_recipe( + operation: ExecutionControlOperation, +) -> ExecutionControlRecipe { + use ExecutionControlOperation::*; + + const EMPTY: &[&str] = &[]; + const I32: &[&str] = &["i32"]; + const I32_I32: &[&str] = &["i32", "i32"]; + const U32: &[&str] = &["u32"]; + const U32_U32: &[&str] = &["u32", "u32"]; + const BASE_CLASSES: &[&str] = &["SDPatternOperator", "Intrinsic"]; + const DEFAULT_CLASSES: &[&str] = &["SDPatternOperator", "Intrinsic", "DefaultAttrsIntrinsic"]; + const BARRIER_PROPERTIES: &[&str] = &["IntrConvergent", "IntrNoCallback"]; + const GRID_PROPERTIES: &[&str] = &["IntrHasSideEffects", "IntrNoMem"]; + const SETMAX_PROPERTIES: &[&str] = &[ + "ImmArg", + "IntrConvergent", + "IntrHasSideEffects", + "IntrNoMem", + "Range", + ]; + const TWO_REGISTER_OR_IMMEDIATE: &[OperandPattern] = &[ + OperandPattern::RegisterOrImmediate, + OperandPattern::RegisterOrImmediate, + ]; + const ONE_IMMEDIATE: &[OperandPattern] = &[OperandPattern::Immediate]; + const NO_OPERANDS: &[OperandPattern] = &[]; + const BARRIER_SYNC_SELECTIONS: &[&str] = &[ + "BARRIER_CTA_SYNC_ii", + "BARRIER_CTA_SYNC_ir", + "BARRIER_CTA_SYNC_ri", + "BARRIER_CTA_SYNC_rr", + ]; + const BAR_SYNC_SELECTIONS: &[&str] = &[ + "BARRIER_CTA_SYNC_ALIGNED_ii", + "BARRIER_CTA_SYNC_ALIGNED_ir", + "BARRIER_CTA_SYNC_ALIGNED_ri", + "BARRIER_CTA_SYNC_ALIGNED_rr", + ]; + const BARRIER_ARRIVE_SELECTIONS: &[&str] = &[ + "BARRIER_CTA_ARRIVE_ii", + "BARRIER_CTA_ARRIVE_ir", + "BARRIER_CTA_ARRIVE_ri", + "BARRIER_CTA_ARRIVE_rr", + ]; + const BAR_ARRIVE_SELECTIONS: &[&str] = &[ + "BARRIER_CTA_ARRIVE_ALIGNED_ii", + "BARRIER_CTA_ARRIVE_ALIGNED_ir", + "BARRIER_CTA_ARRIVE_ALIGNED_ri", + "BARRIER_CTA_ARRIVE_ALIGNED_rr", + ]; + const PTX_60: &[&str] = &["Subtarget->getPTXVersion() >= 60"]; + const GRID_PREDICATES: &[&str] = &[ + "Subtarget->getSmVersion() >= 90", + "Subtarget->getPTXVersion() >= 78", + ]; + const SETMAX_PREDICATES: &[&str] = &["Subtarget->hasSetMaxNRegSupport()"]; + const SYNC_MODIFIERS: &[&str] = &["sync"]; + const ARRIVE_MODIFIERS: &[&str] = &["arrive"]; + const LAUNCH_DEPENDENTS_MODIFIERS: &[&str] = &["launch_dependents"]; + const WAIT_MODIFIERS: &[&str] = &["wait"]; + const SETMAX_INC_MODIFIERS: &[&str] = &["inc", "sync", "aligned", "u32"]; + const SETMAX_DEC_MODIFIERS: &[&str] = &["dec", "sync", "aligned", "u32"]; + const GRID_LAUNCH_SELECTIONS: &[&str] = &["GRIDDEPCONTROL_LAUNCH_DEPENDENTS"]; + const GRID_WAIT_SELECTIONS: &[&str] = &["GRIDDEPCONTROL_WAIT"]; + const SETMAX_INC_SELECTIONS: &[&str] = &["anonymous_22212"]; + const SETMAX_DEC_SELECTIONS: &[&str] = &["anonymous_22214"]; + const BAR_SECTION: &str = + "9.7.14.1 Parallel Synchronization and Communication Instructions: bar, barrier"; + const BAR_URL: &str = "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-bar-barrier"; + const GRID_SECTION: &str = + "9.7.15.1 Parallel Synchronization and Communication Instructions: griddepcontrol"; + const GRID_URL: &str = "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-griddepcontrol"; + const SETMAX_SECTION: &str = "9.7.16.5 Miscellaneous Instructions: setmaxnreg"; + const SETMAX_URL: &str = "https://docs.nvidia.com/cuda/parallel-thread-execution/#miscellaneous-instructions-setmaxnreg"; + + let ( + abi_id, + id, + operation_key, + source_record, + llvm_symbol, + rust_module, + rust_arguments, + compatibility_path, + dialect_op_type, + dialect_op_name, + dialect_operands, + llvm_arguments, + imported_classes, + imported_properties, + memory, + convergent, + execution_scope, + minimum_ptx, + minimum_sm, + targets, + ptx_isa_section, + ptx_isa_url, + mnemonic, + modifiers, + operands, + selection_records, + selection_asm, + selection_predicates, + summary, + ) = match operation { + BarrierCtaSync => ( + "i0883", + "barrier_cta_sync", + "synchronization.cta.barrier.count", + "int_nvvm_barrier_cta_sync_count", + "llvm.nvvm.barrier.cta.sync.count", + "barrier", + U32_U32, + "cuda_device::barrier::barrier_cta_sync", + "BarrierCtaSyncCountOp", + "nvvm.barrier_cta_sync_count", + I32_I32, + I32_I32, + BASE_CLASSES, + BARRIER_PROPERTIES, + "read_write", + true, + "cta", + "6.0", + Some("sm_30"), + "all", + BAR_SECTION, + BAR_URL, + "barrier", + SYNC_MODIFIERS, + TWO_REGISTER_OR_IMMEDIATE, + BARRIER_SYNC_SELECTIONS, + "barrier.sync \t$i, $j;", + PTX_60, + "Synchronizes the requested number of CTA threads at a numbered barrier.", + ), + BarrierCtaSyncAligned => ( + "i0884", + "barrier_cta_sync_aligned", + "synchronization.cta.bar.aligned.count", + "int_nvvm_barrier_cta_sync_aligned_count", + "llvm.nvvm.barrier.cta.sync.aligned.count", + "barrier", + U32_U32, + "cuda_device::barrier::barrier_cta_sync_aligned", + "BarrierCtaSyncAlignedCountOp", + "nvvm.barrier_cta_sync_aligned_count", + I32_I32, + I32_I32, + BASE_CLASSES, + BARRIER_PROPERTIES, + "read_write", + true, + "cta", + "1.0", + None, + "all", + BAR_SECTION, + BAR_URL, + "bar", + SYNC_MODIFIERS, + TWO_REGISTER_OR_IMMEDIATE, + BAR_SYNC_SELECTIONS, + "bar.sync \t$i, $j;", + EMPTY, + "Synchronizes the requested number of CTA threads at an aligned numbered barrier.", + ), + BarrierCtaArrive => ( + "i0885", + "barrier_cta_arrive", + "synchronization.cta.barrier.arrive.count", + "int_nvvm_barrier_cta_arrive_count", + "llvm.nvvm.barrier.cta.arrive.count", + "barrier", + U32_U32, + "cuda_device::barrier::barrier_cta_arrive", + "BarrierCtaArriveCountOp", + "nvvm.barrier_cta_arrive_count", + I32_I32, + I32_I32, + BASE_CLASSES, + BARRIER_PROPERTIES, + "read_write", + true, + "cta", + "6.0", + Some("sm_30"), + "all", + BAR_SECTION, + BAR_URL, + "barrier", + ARRIVE_MODIFIERS, + TWO_REGISTER_OR_IMMEDIATE, + BARRIER_ARRIVE_SELECTIONS, + "barrier.arrive \t$i, $j;", + PTX_60, + "Signals arrival of the requested number of CTA threads at a numbered barrier without waiting.", + ), + BarrierCtaArriveAligned => ( + "i0886", + "barrier_cta_arrive_aligned", + "synchronization.cta.bar.arrive.aligned.count", + "int_nvvm_barrier_cta_arrive_aligned_count", + "llvm.nvvm.barrier.cta.arrive.aligned.count", + "barrier", + U32_U32, + "cuda_device::barrier::barrier_cta_arrive_aligned", + "BarrierCtaArriveAlignedCountOp", + "nvvm.barrier_cta_arrive_aligned_count", + I32_I32, + I32_I32, + BASE_CLASSES, + BARRIER_PROPERTIES, + "read_write", + true, + "cta", + "1.0", + None, + "all", + BAR_SECTION, + BAR_URL, + "bar", + ARRIVE_MODIFIERS, + TWO_REGISTER_OR_IMMEDIATE, + BAR_ARRIVE_SELECTIONS, + "bar.arrive \t$i, $j;", + EMPTY, + "Signals arrival of the requested number of CTA threads at an aligned numbered barrier without waiting.", + ), + GridDependencyLaunchDependents => ( + "i0913", + "grid_dependency_launch_dependents", + "synchronization.grid.dependency.launch_dependents", + "int_nvvm_griddepcontrol_launch_dependents", + "llvm.nvvm.griddepcontrol.launch.dependents", + "grid", + EMPTY, + "cuda_device::grid::dependency::trigger_dependents", + "GridDependencyLaunchDependentsOp", + "nvvm.grid_dependency_launch_dependents", + EMPTY, + EMPTY, + BASE_CLASSES, + GRID_PROPERTIES, + "none", + false, + "grid", + "7.8", + Some("sm_90"), + "all", + GRID_SECTION, + GRID_URL, + "griddepcontrol", + LAUNCH_DEPENDENTS_MODIFIERS, + NO_OPERANDS, + GRID_LAUNCH_SELECTIONS, + "griddepcontrol.launch_dependents;", + GRID_PREDICATES, + "Makes dependent grids eligible to launch before the current grid completes.", + ), + GridDependencyWait => ( + "i0914", + "grid_dependency_wait", + "synchronization.grid.dependency.wait", + "int_nvvm_griddepcontrol_wait", + "llvm.nvvm.griddepcontrol.wait", + "grid", + EMPTY, + "cuda_device::grid::dependency::wait", + "GridDependencyWaitOp", + "nvvm.grid_dependency_wait", + EMPTY, + EMPTY, + BASE_CLASSES, + GRID_PROPERTIES, + "none", + false, + "grid", + "7.8", + Some("sm_90"), + "all", + GRID_SECTION, + GRID_URL, + "griddepcontrol", + WAIT_MODIFIERS, + NO_OPERANDS, + GRID_WAIT_SELECTIONS, + "griddepcontrol.wait;", + GRID_PREDICATES, + "Waits until all prerequisite grids have completed.", + ), + SetMaxNRegInc => ( + "i0915", + "setmaxnreg_inc", + "register_control.warpgroup.setmaxnreg.increment", + "int_nvvm_setmaxnreg_inc_sync_aligned_u32", + "llvm.nvvm.setmaxnreg.inc.sync.aligned.u32", + "thread", + U32, + "cuda_device::thread::__setmaxnreg_inc", + "SetMaxNRegIncOp", + "nvvm.setmaxnreg_inc", + EMPTY, + I32, + DEFAULT_CLASSES, + SETMAX_PROPERTIES, + "none", + true, + "warpgroup", + "8.0", + None, + TENSOR_MAP_REPLACE_TARGETS, + SETMAX_SECTION, + SETMAX_URL, + "setmaxnreg", + SETMAX_INC_MODIFIERS, + ONE_IMMEDIATE, + SETMAX_INC_SELECTIONS, + "setmaxnreg.inc.sync.aligned.u32 \t$reg_count;", + SETMAX_PREDICATES, + "Increases the executing warpgroup's maximum register allocation to the immediate count.", + ), + SetMaxNRegDec => ( + "i0916", + "setmaxnreg_dec", + "register_control.warpgroup.setmaxnreg.decrement", + "int_nvvm_setmaxnreg_dec_sync_aligned_u32", + "llvm.nvvm.setmaxnreg.dec.sync.aligned.u32", + "thread", + U32, + "cuda_device::thread::__setmaxnreg_dec", + "SetMaxNRegDecOp", + "nvvm.setmaxnreg_dec", + EMPTY, + I32, + DEFAULT_CLASSES, + SETMAX_PROPERTIES, + "none", + true, + "warpgroup", + "8.0", + None, + TENSOR_MAP_REPLACE_TARGETS, + SETMAX_SECTION, + SETMAX_URL, + "setmaxnreg", + SETMAX_DEC_MODIFIERS, + ONE_IMMEDIATE, + SETMAX_DEC_SELECTIONS, + "setmaxnreg.dec.sync.aligned.u32 \t$reg_count;", + SETMAX_PREDICATES, + "Decreases the executing warpgroup's maximum register allocation to the immediate count.", + ), + }; + + ExecutionControlRecipe { + abi_id, + id, + operation_key, + source_record, + llvm_symbol, + rust_module, + rust_arguments, + compatibility_path, + dialect_op_type, + dialect_op_name, + dialect_operands, + llvm_arguments, + imported_classes, + imported_properties, + memory, + convergent, + execution_scope, + minimum_ptx, + minimum_sm, + targets, + ptx_isa_section, + ptx_isa_url, + mnemonic, + modifiers, + operands, + selection_records, + selection_asm, + selection_predicates, + summary, + } +} + +pub(in crate::resolve) fn execution_control_backend_floor( + operation: ExecutionControlOperation, + backend: IntrinsicBackend, +) -> (&'static str, Option<&'static str>) { + use ExecutionControlOperation::*; + + match (operation, backend) { + (BarrierCtaSync | BarrierCtaArrive, IntrinsicBackend::LlvmNvptx) => ("6.0", Some("sm_30")), + (BarrierCtaSync | BarrierCtaArrive, IntrinsicBackend::LibNvvm) => ("6.0", Some("sm_75")), + (BarrierCtaSyncAligned | BarrierCtaArriveAligned, IntrinsicBackend::LlvmNvptx) => { + ("3.2", Some("sm_20")) + } + (BarrierCtaSyncAligned | BarrierCtaArriveAligned, IntrinsicBackend::LibNvvm) => { + ("1.0", Some("sm_75")) + } + _ => { + let recipe = execution_control_recipe(operation); + (recipe.minimum_ptx, recipe.minimum_sm) + } + } +} + +pub(in crate::resolve) fn validate_execution_control_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, +) -> Result<()> { + let operation = ExecutionControlOperation::from_catalog_id(&policy.id) + .with_context(|| format!("{} is not a closed execution-control operation", policy.id))?; + let recipe = execution_control_recipe(operation); + ensure!( + policy.family == operation.family() + && policy.abi_id == recipe.abi_id + && policy.operation_key == recipe.operation_key + && policy.summary == recipe.summary + && policy.source.is_none() + && policy.source_record.as_deref() == Some(recipe.source_record) + && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) + && policy.resolved_llvm_symbol.is_none() + && declaration.source_record == recipe.source_record + && declaration.llvm_name == recipe.llvm_symbol, + "{} execution-control identity changed", + policy.id + ); + ensure!( + policy.rust_module == recipe.rust_module + && policy.rust_name == recipe.id + && policy.rust_arguments == recipe.rust_arguments + && policy.rust_result == "()" + && !policy.safe + && !policy.must_use + && policy.safe_allowlist_reason.is_none() + && policy.public_rust_path + == format!("cuda_intrinsics::{}::{}", recipe.rust_module, recipe.id) + && policy.compatibility_rust_paths == [recipe.compatibility_path], + "{} execution-control Rust API changed", + policy.id + ); + ensure!( + policy.dialect_op_type == recipe.dialect_op_type + && policy.dialect_op_name == recipe.dialect_op_name + && policy.dialect_operands == recipe.dialect_operands + && policy.dialect_results.is_empty() + && policy.llvm_arguments == recipe.llvm_arguments + && policy.llvm_results.is_empty() + && declaration.arguments == recipe.llvm_arguments + && declaration.results.is_empty() + && policy.lowering == "generated_execution_control", + "{} execution-control carrier or LLVM adapter changed", + policy.id + ); + ensure!( + declaration.classes == recipe.imported_classes + && declaration.properties == recipe.imported_properties, + "{} imported execution-control declaration changed", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == recipe.memory + && policy.convergent == recipe.convergent + && policy.execution_scope == recipe.execution_scope, + "{} execution-control semantics changed", + policy.id + ); + ensure!( + policy.minimum_ptx == recipe.minimum_ptx + && policy.minimum_sm.as_deref() == recipe.minimum_sm + && policy.targets == recipe.targets + && policy.ptx_result == "()" + && policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section == recipe.ptx_isa_section + && policy.ptx_isa_url == recipe.ptx_isa_url + && policy.expected_ptx.mnemonic == recipe.mnemonic + && policy.expected_ptx.modifiers == recipe.modifiers + && policy.expected_ptx.operands == recipe.operands, + "{} execution-control target or PTX contract changed", + policy.id + ); + let valid_route = |backend, mechanism| { + let (minimum_ptx, minimum_sm) = execution_control_backend_floor(operation, backend); + policy.backend_lowerings.iter().any(|route| { + route.backend == backend + && route.mechanism == mechanism + && route.minimum_ptx.as_deref() == Some(minimum_ptx) + && route.minimum_sm.as_deref() == minimum_sm + && !route.evidence_profile.trim().is_empty() + }) + }; + ensure!( + policy.backend_lowerings.len() == 2 + && valid_route( + IntrinsicBackend::LlvmNvptx, + BackendLoweringMechanism::TypedNvvm, + ) + && valid_route( + IntrinsicBackend::LibNvvm, + BackendLoweringMechanism::InlinePtx, + ), + "{} execution-control backend route changed", + policy.id + ); + let expected_selections = recipe + .selection_records + .iter() + .copied() + .collect::>(); + let actual_selections = declaration + .selections + .iter() + .map(|selection| selection.source_record.as_str()) + .collect::>(); + ensure!( + declaration.selections.len() == recipe.selection_records.len() + && actual_selections == expected_selections + && declaration.selections.iter().all(|selection| { + selection.asm == recipe.selection_asm + && selection.predicates == recipe.selection_predicates + && selection.constraints.is_empty() + }), + "{} imported execution-control selection family changed", + policy.id + ); + ensure_no_other_family_contract(policy, "execution-control")?; + Ok(()) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/matrix_move.rs b/crates/cuda-intrinsics-gen/src/resolve/families/matrix_move.rs new file mode 100644 index 0000000000..4eef5aec67 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/matrix_move.rs @@ -0,0 +1,776 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, CatalogTargetRequirement, ImportedAddressSpace, ImportedIntrinsic, + IntrinsicBackend, IntrinsicSource, LdmatrixAdapter, LdmatrixAddressContract, LdmatrixElement, + LdmatrixLayout, LdmatrixMemoryOrder, LdmatrixMultiplicity, LdmatrixParticipation, + LdmatrixShape, LdmatrixStateSpace, MovmatrixAdapter, MovmatrixParticipation, + OverlayBackendLowering, OverlayIntrinsic, RuntimeValidation, StmatrixAdmission, StmatrixLayout, + StmatrixMultiplicity, +}; +use crate::ptx::{InstructionPattern, OperandPattern}; +use anyhow::{Context, Result, ensure}; +use std::collections::BTreeSet; + +use crate::resolve::guards::*; +use crate::resolve::targets::*; + +pub(in crate::resolve) const BLACKWELL_LDMATRIX_LLVM_TARGETS: &str = + "sm_100a|sm_100f|sm_103a|sm_103f|sm_110a|sm_110f|sm_120a|sm_120f|sm_121a|sm_121f"; +pub(in crate::resolve) const BLACKWELL_LDMATRIX_LIBNVVM_TARGETS: &str = + BLACKWELL_LDMATRIX_LLVM_TARGETS; +#[derive(Clone, Copy)] +pub(in crate::resolve) struct StmatrixRecipe { + multiplicity: StmatrixMultiplicity, + layout: StmatrixLayout, + abi_id: &'static str, + id: &'static str, + operation_key: &'static str, + source_record: &'static str, + llvm_symbol: &'static str, + compatibility_name: &'static str, + dialect_op_type: &'static str, + dialect_op_name: &'static str, + summary: &'static str, +} + +pub(in crate::resolve) fn stmatrix_recipe( + multiplicity: StmatrixMultiplicity, + layout: StmatrixLayout, +) -> StmatrixRecipe { + match (multiplicity, layout) { + (StmatrixMultiplicity::X2, StmatrixLayout::Normal) => StmatrixRecipe { + multiplicity, + layout, + abi_id: "i0301", + id: "stmatrix_m8n8_x2_b16", + operation_key: "matrix.stmatrix.m8n8.x2.normal.b16.shared", + source_record: "int_nvvm_stmatrix_sync_aligned_m8n8_x2_b16", + llvm_symbol: "llvm.nvvm.stmatrix.sync.aligned.m8n8.x2.b16", + compatibility_name: "stmatrix_m8n8_x2", + dialect_op_type: "StmatrixM8n8X2Op", + dialect_op_name: "nvvm.stmatrix_m8n8_x2", + summary: "Stores two 8×8 b16 matrix fragments cooperatively to shared memory.", + }, + (StmatrixMultiplicity::X2, StmatrixLayout::Transposed) => StmatrixRecipe { + multiplicity, + layout, + abi_id: "i0302", + id: "stmatrix_m8n8_x2_trans_b16", + operation_key: "matrix.stmatrix.m8n8.x2.transposed.b16.shared", + source_record: "int_nvvm_stmatrix_sync_aligned_m8n8_x2_trans_b16", + llvm_symbol: "llvm.nvvm.stmatrix.sync.aligned.m8n8.x2.trans.b16", + compatibility_name: "stmatrix_m8n8_x2_trans", + dialect_op_type: "StmatrixM8n8X2TransOp", + dialect_op_name: "nvvm.stmatrix_m8n8_x2_trans", + summary: "Stores two transposed 8×8 b16 matrix fragments cooperatively to shared memory.", + }, + (StmatrixMultiplicity::X4, StmatrixLayout::Normal) => StmatrixRecipe { + multiplicity, + layout, + abi_id: "i0303", + id: "stmatrix_m8n8_x4_b16", + operation_key: "matrix.stmatrix.m8n8.x4.normal.b16.shared", + source_record: "int_nvvm_stmatrix_sync_aligned_m8n8_x4_b16", + llvm_symbol: "llvm.nvvm.stmatrix.sync.aligned.m8n8.x4.b16", + compatibility_name: "stmatrix_m8n8_x4", + dialect_op_type: "StmatrixM8n8X4Op", + dialect_op_name: "nvvm.stmatrix_m8n8_x4", + summary: "Stores four 8×8 b16 matrix fragments cooperatively to shared memory.", + }, + (StmatrixMultiplicity::X4, StmatrixLayout::Transposed) => StmatrixRecipe { + multiplicity, + layout, + abi_id: "i0304", + id: "stmatrix_m8n8_x4_trans_b16", + operation_key: "matrix.stmatrix.m8n8.x4.transposed.b16.shared", + source_record: "int_nvvm_stmatrix_sync_aligned_m8n8_x4_trans_b16", + llvm_symbol: "llvm.nvvm.stmatrix.sync.aligned.m8n8.x4.trans.b16", + compatibility_name: "stmatrix_m8n8_x4_trans", + dialect_op_type: "StmatrixM8n8X4TransOp", + dialect_op_name: "nvvm.stmatrix_m8n8_x4_trans", + summary: "Stores four transposed 8×8 b16 matrix fragments cooperatively to shared memory.", + }, + } +} + +pub(in crate::resolve) fn stmatrix_variant_for_id( + id: &str, +) -> Option<(StmatrixMultiplicity, StmatrixLayout)> { + match id { + "stmatrix_m8n8_x2_b16" => Some((StmatrixMultiplicity::X2, StmatrixLayout::Normal)), + "stmatrix_m8n8_x2_trans_b16" => { + Some((StmatrixMultiplicity::X2, StmatrixLayout::Transposed)) + } + "stmatrix_m8n8_x4_b16" => Some((StmatrixMultiplicity::X4, StmatrixLayout::Normal)), + "stmatrix_m8n8_x4_trans_b16" => { + Some((StmatrixMultiplicity::X4, StmatrixLayout::Transposed)) + } + _ => None, + } +} + +pub(in crate::resolve) fn expand_stmatrix_admission( + admission: &StmatrixAdmission, +) -> Result> { + ensure!( + admission.runtime_validation == RuntimeValidation::Unexecuted, + "stmatrix runtime validation may be marked executed only with GPU evidence" + ); + ensure!( + !admission.llvm_evidence_profile.trim().is_empty() + && !admission.libnvvm_evidence_profile.trim().is_empty(), + "compact stmatrix admission requires both backend evidence profiles" + ); + let expected_variants = [ + (StmatrixMultiplicity::X2, StmatrixLayout::Normal), + (StmatrixMultiplicity::X2, StmatrixLayout::Transposed), + (StmatrixMultiplicity::X4, StmatrixLayout::Normal), + (StmatrixMultiplicity::X4, StmatrixLayout::Transposed), + ]; + let actual_variants = admission + .variants + .iter() + .map(|variant| (variant.multiplicity, variant.layout)) + .collect::>(); + ensure!( + actual_variants == expected_variants, + "compact stmatrix admission must contain the four reviewed variants in canonical order" + ); + + admission + .variants + .iter() + .map(|variant| { + let recipe = stmatrix_recipe(variant.multiplicity, variant.layout); + ensure!( + variant.abi_id == recipe.abi_id, + "{} must keep reserved ABI ID {}", + recipe.id, + recipe.abi_id + ); + let count = recipe.multiplicity.register_count(); + let mut rust_arguments = vec!["*mut u8".to_owned()]; + rust_arguments.extend(std::iter::repeat_n("u32".to_owned(), count)); + let mut dialect_operands = vec!["ptr".to_owned()]; + dialect_operands.extend(std::iter::repeat_n("i32".to_owned(), count)); + let mut llvm_arguments = vec!["anyptr".to_owned()]; + llvm_arguments.extend(std::iter::repeat_n("i32".to_owned(), count)); + let multiplicity = match recipe.multiplicity { + StmatrixMultiplicity::X2 => "x2", + StmatrixMultiplicity::X4 => "x4", + }; + let mut modifiers = vec![ + "sync".into(), + "aligned".into(), + "m8n8".into(), + multiplicity.into(), + ]; + if recipe.layout == StmatrixLayout::Transposed { + modifiers.push("trans".into()); + } + modifiers.extend(["shared".into(), "b16".into()]); + Ok(OverlayIntrinsic { + id: recipe.id.into(), + abi_id: variant.abi_id.clone(), + operation_key: recipe.operation_key.into(), + family: "stmatrix".into(), + source: None, + source_record: Some(recipe.source_record.into()), + rust_module: "matrix".into(), + rust_name: recipe.id.into(), + rust_arguments, + rust_result: "()".into(), + safe: false, + must_use: false, + safe_allowlist_reason: None, + public_rust_path: format!("cuda_intrinsics::matrix::{}", recipe.id), + compatibility_rust_paths: vec![format!( + "cuda_device::tcgen05::{}", + recipe.compatibility_name + )], + dialect_op_type: recipe.dialect_op_type.into(), + dialect_op_name: recipe.dialect_op_name.into(), + dialect_operands, + dialect_results: vec![], + llvm_symbol: Some(recipe.llvm_symbol.into()), + resolved_llvm_symbol: Some(format!("{}.p3", recipe.llvm_symbol)), + llvm_arguments, + llvm_results: vec![], + pure: false, + memory: "write".into(), + convergent: true, + execution_scope: "warp".into(), + minimum_ptx: "7.8".into(), + minimum_sm: Some("sm_90".into()), + ptx_result: "()".into(), + targets: "all".into(), + ptx_isa_version: "9.3".into(), + ptx_isa_section: "9.7.14.5.16 Warp-level matrix store instruction: stmatrix" + .into(), + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-stmatrix".into(), + lowering: "generated_stmatrix".into(), + backend_lowerings: vec![ + OverlayBackendLowering { + backend: IntrinsicBackend::LlvmNvptx, + mechanism: BackendLoweringMechanism::TypedNvvm, + evidence_profile: admission.llvm_evidence_profile.clone(), + targets: None, + minimum_ptx: Some("7.8".into()), + minimum_sm: Some("sm_90".into()), + }, + OverlayBackendLowering { + backend: IntrinsicBackend::LibNvvm, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: admission.libnvvm_evidence_profile.clone(), + targets: None, + minimum_ptx: Some("7.8".into()), + minimum_sm: Some("sm_90".into()), + }, + ], + packed_atomic: None, + redux: None, + vote: None, + active_mask: None, + warp_match: None, + warp_barrier: None, + warp_shuffle: None, + dot_product: None, + packed_alu: None, + integer_minmax: None, + packed_conversion: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cp_async_copy: None, + cp_async_control: None, + cp_async_mbarrier: None, + mbarrier_basic: None, + movmatrix: None, + mbarrier_extended: None, + register_mma: None, + sparse_mma: None, + prmt: None, + cluster_barrier: None, + wgmma_control: None, + special_register: None, + debug_control: None, + cluster_memory: None, + clc: None, + tma: None, + tcgen05: None, + ldmatrix_variant: None, + ldmatrix_safety: None, + ldmatrix_adapter: None, + selected_address_space: Some(ImportedAddressSpace::Shared), + expected_ptx: InstructionPattern { + mnemonic: "stmatrix".into(), + modifiers, + operands: vec![ + OperandPattern::Address, + OperandPattern::RegisterList { length: count }, + ], + }, + summary: recipe.summary.into(), + }) + }) + .collect() +} + +pub(in crate::resolve) fn validate_stmatrix_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, +) -> Result<()> { + let (multiplicity, layout) = stmatrix_variant_for_id(&policy.id) + .with_context(|| format!("{} has no closed stmatrix recipe", policy.id))?; + let recipe = stmatrix_recipe(multiplicity, layout); + let count = multiplicity.register_count(); + let mut rust_arguments = vec!["*mut u8".to_owned()]; + rust_arguments.extend(std::iter::repeat_n("u32".to_owned(), count)); + let mut dialect_operands = vec!["ptr".to_owned()]; + dialect_operands.extend(std::iter::repeat_n("i32".to_owned(), count)); + let mut llvm_arguments = vec!["anyptr".to_owned()]; + llvm_arguments.extend(std::iter::repeat_n("i32".to_owned(), count)); + + ensure!( + policy.abi_id == recipe.abi_id + && policy.operation_key == recipe.operation_key + && policy.source.is_none() + && policy.source_record.as_deref() == Some(recipe.source_record) + && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) + && policy.resolved_llvm_symbol.as_deref() + == Some(format!("{}.p3", recipe.llvm_symbol).as_str()), + "{} stmatrix identity does not match its closed recipe", + policy.id + ); + ensure!( + policy.rust_module == "matrix" + && policy.rust_name == recipe.id + && policy.rust_arguments == rust_arguments + && policy.rust_result == "()" + && !policy.safe + && !policy.must_use + && policy.safe_allowlist_reason.is_none() + && policy.public_rust_path == format!("cuda_intrinsics::matrix::{}", recipe.id) + && policy.compatibility_rust_paths + == [format!( + "cuda_device::tcgen05::{}", + recipe.compatibility_name + )], + "{} stmatrix Rust API does not match its closed recipe", + policy.id + ); + ensure!( + policy.dialect_op_type == recipe.dialect_op_type + && policy.dialect_op_name == recipe.dialect_op_name + && policy.dialect_operands == dialect_operands + && policy.dialect_results.is_empty() + && policy.llvm_arguments == llvm_arguments + && policy.llvm_results.is_empty() + && policy.lowering == "generated_stmatrix" + && policy.selected_address_space == Some(ImportedAddressSpace::Shared), + "{} stmatrix carriers or lowering do not match its closed recipe", + policy.id + ); + ensure!( + declaration.properties + == [ + "IntrArgMemOnly", + "IntrConvergent", + "IntrNoCallback", + "IntrWriteMem", + "NoCapture", + "WriteOnly", + ] + && !policy.pure + && policy.memory == "write" + && policy.convergent + && policy.execution_scope == "warp", + "{} stmatrix effects disagree with the imported declaration", + policy.id + ); + ensure!( + policy.minimum_ptx == "7.8" + && policy.minimum_sm.as_deref() == Some("sm_90") + && policy.ptx_result == "()" + && policy.targets == "all" + && policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section + == "9.7.14.5.16 Warp-level matrix store instruction: stmatrix" + && policy.ptx_isa_url + == "https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-stmatrix", + "{} stmatrix target floor or PTX provenance changed", + policy.id + ); + let multiplicity_name = match multiplicity { + StmatrixMultiplicity::X2 => "x2", + StmatrixMultiplicity::X4 => "x4", + }; + let mut modifiers = vec![ + "sync".into(), + "aligned".into(), + "m8n8".into(), + multiplicity_name.into(), + ]; + if layout == StmatrixLayout::Transposed { + modifiers.push("trans".into()); + } + modifiers.extend(["shared".into(), "b16".into()]); + ensure!( + policy.expected_ptx + == (InstructionPattern { + mnemonic: "stmatrix".into(), + modifiers, + operands: vec![ + OperandPattern::Address, + OperandPattern::RegisterList { length: count }, + ], + }), + "{} expected PTX does not match its closed stmatrix shape", + policy.id + ); + ensure!( + policy.backend_lowerings.len() == 2 + && policy.backend_lowerings.iter().all(|route| { + !route.evidence_profile.trim().is_empty() + && route.minimum_ptx.as_deref() == Some("7.8") + && route.minimum_sm.as_deref() == Some("sm_90") + }) + && policy.backend_lowerings.iter().any(|route| { + route.backend == IntrinsicBackend::LlvmNvptx + && route.mechanism == BackendLoweringMechanism::TypedNvvm + }) + && policy.backend_lowerings.iter().any(|route| { + route.backend == IntrinsicBackend::LibNvvm + && route.mechanism == BackendLoweringMechanism::InlinePtx + }), + "{} must keep its reviewed typed-NVVM and inline-PTX routes", + policy.id + ); + ensure!( + policy.ldmatrix_variant.is_none() + && policy.ldmatrix_safety.is_none() + && policy.ldmatrix_adapter.is_none() + && policy.packed_atomic.is_none() + && policy.redux.is_none() + && policy.vote.is_none() + && policy.active_mask.is_none() + && policy.warp_match.is_none() + && policy.warp_barrier.is_none() + && policy.warp_shuffle.is_none() + && policy.dot_product.is_none() + && policy.packed_alu.is_none() + && policy.packed_conversion.is_none() + && policy.cp_async_copy.is_none() + && policy.cp_async_control.is_none() + && policy.cp_async_mbarrier.is_none() + && policy.mbarrier_basic.is_none() + && policy.register_mma.is_none() + && policy.sparse_mma.is_none() + && policy.prmt.is_none() + && policy.cluster_barrier.is_none() + && policy.special_register.is_none() + && policy.debug_control.is_none(), + "{} mixes another generated-family contract with stmatrix", + policy.id + ); + Ok(()) +} + +pub(in crate::resolve) fn validate_ldmatrix_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, +) -> Result<()> { + let variant = policy + .ldmatrix_variant + .as_ref() + .with_context(|| format!("{} has no closed ldmatrix variant", policy.id))?; + let safety = policy + .ldmatrix_safety + .as_ref() + .with_context(|| format!("{} has no ldmatrix safety contract", policy.id))?; + let classic = variant.shape == LdmatrixShape::M8n8 && variant.element == LdmatrixElement::B16; + let blackwell = matches!( + (variant.shape, variant.layout, variant.element), + ( + LdmatrixShape::M8n16, + LdmatrixLayout::Normal, + LdmatrixElement::B8x16B4x16P64 | LdmatrixElement::B8x16B6x16P32 + ) | ( + LdmatrixShape::M16n16, + LdmatrixLayout::Transposed, + LdmatrixElement::B8 | LdmatrixElement::B8x16B4x16P64 | LdmatrixElement::B8x16B6x16P32 + ) + ); + ensure!( + variant.shape != LdmatrixShape::M16n16 || variant.multiplicity != LdmatrixMultiplicity::X4, + "{} requests unsupported m16n16.x4 ldmatrix", + policy.id + ); + ensure!( + (classic || blackwell) && variant.state_space == LdmatrixStateSpace::Shared, + "{} requests an unsupported ldmatrix shape, element, or state space", + policy.id + ); + let expected_participation = if classic { + LdmatrixParticipation::AllWarpLanesSameInstruction + } else { + LdmatrixParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes + }; + let expected_address_contract = match variant.shape { + LdmatrixShape::M8n8 => { + LdmatrixAddressContract::WarpLaneAddressesMappedByMultiplicitySixteenByteAlignedSixteenBytesReadableWithSm75Replication + } + LdmatrixShape::M8n16 => { + LdmatrixAddressContract::WarpLaneAddressesMappedByMultiplicitySixteenByteAlignedSixteenBytesReadable + } + LdmatrixShape::M16n16 => { + LdmatrixAddressContract::WarpLaneAddressesMappedByMultiplicitySixteenByteAlignedThirtyTwoBytesReadable + } + }; + ensure!( + safety.participation == expected_participation + && safety.address_contract == expected_address_contract + && safety.memory_order == LdmatrixMemoryOrder::Weak + && safety.runtime_validation == RuntimeValidation::Unexecuted, + "{} has an unsupported ldmatrix safety contract", + policy.id + ); + let count = variant.register_count(); + let count_name = match variant.multiplicity { + LdmatrixMultiplicity::X1 => "x1", + LdmatrixMultiplicity::X2 => "x2", + LdmatrixMultiplicity::X4 => "x4", + }; + let trans_record = match variant.layout { + LdmatrixLayout::Normal => "", + LdmatrixLayout::Transposed => "_trans", + }; + let trans_symbol = match variant.layout { + LdmatrixLayout::Normal => "", + LdmatrixLayout::Transposed => ".trans", + }; + let layout_name = match variant.layout { + LdmatrixLayout::Normal => "normal", + LdmatrixLayout::Transposed => "transposed", + }; + let shape_name = match variant.shape { + LdmatrixShape::M8n8 => "m8n8", + LdmatrixShape::M8n16 => "m8n16", + LdmatrixShape::M16n16 => "m16n16", + }; + let (element_record, element_symbol) = match variant.element { + LdmatrixElement::B16 => ("b16", "b16"), + LdmatrixElement::B8 => ("b8", "b8"), + LdmatrixElement::B8x16B4x16P64 => ("b8x16_b4x16_p64", "b8x16.b4x16_p64"), + LdmatrixElement::B8x16B6x16P32 => ("b8x16_b6x16_p32", "b8x16.b6x16_p32"), + }; + let expected_source = format!( + "int_nvvm_ldmatrix_sync_aligned_{shape_name}_{count_name}{trans_record}_{element_record}" + ); + let expected_symbol = format!( + "llvm.nvvm.ldmatrix.sync.aligned.{shape_name}.{count_name}{trans_symbol}.{element_symbol}" + ); + let expected_name = + format!("ldmatrix_{shape_name}_{count_name}{trans_record}_{element_record}"); + let expected_result = if count == 1 { + "u32".to_owned() + } else { + format!("[u32; {count}]") + }; + let expected_adapter = if count == 1 { + LdmatrixAdapter::SingleResultDirect + } else { + LdmatrixAdapter::MultipleResultsToArray + }; + ensure!( + policy.source_record.as_deref() == Some(expected_source.as_str()) + && policy.llvm_symbol.as_deref() == Some(expected_symbol.as_str()), + "{} ldmatrix variant does not match its imported source record or base LLVM symbol", + policy.id + ); + ensure!( + policy.resolved_llvm_symbol.as_deref() == Some(format!("{expected_symbol}.p3").as_str()), + "{} must keep the imported base symbol distinct from the resolved `.p3` overload", + policy.id + ); + ensure!( + policy.rust_arguments == [if classic { "*const u32" } else { "*const u8" }] + && policy.rust_result == expected_result + && policy.llvm_arguments == ["anyptr"] + && policy.llvm_results == vec!["i32"; count] + && policy.ptx_result == policy.rust_result, + "{} ldmatrix Rust, imported LLVM, and PTX carrier signatures disagree", + policy.id + ); + ensure!( + policy.id == expected_name + && policy.operation_key + == format!( + "matrix.ldmatrix.{shape_name}.{count_name}.{layout_name}.{element_record}.shared" + ) + && policy.rust_module == "matrix" + && policy.rust_name == expected_name + && !policy.safe + && policy.must_use != classic + && policy.safe_allowlist_reason.is_none() + && { + let expected_compatibility_paths = if classic { + vec![format!( + "cuda_device::wmma::ldmatrix_{count_name}{trans_record}" + )] + } else if matches!( + policy.id.as_str(), + "ldmatrix_m8n16_x2_b8x16_b4x16_p64" | "ldmatrix_m8n16_x4_b8x16_b4x16_p64" + ) { + vec![format!("cuda_device::wmma::{expected_name}")] + } else { + vec![] + }; + policy.compatibility_rust_paths == expected_compatibility_paths + } + && policy.lowering == "generated_ldmatrix" + && policy.ldmatrix_adapter == Some(expected_adapter) + && policy.selected_address_space == Some(ImportedAddressSpace::Shared), + "{} must preserve the closed raw/compatibility ldmatrix API, result adapter, and shared selection", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == "read" + && policy.convergent + && policy.execution_scope == "warp" + && declaration + .properties + .iter() + .any(|property| property == "IntrArgMemOnly") + && declaration + .properties + .iter() + .any(|property| property == "IntrReadMem"), + "{} ldmatrix effects disagree with the imported declaration", + policy.id + ); + let backend_pairs: BTreeSet<_> = policy + .backend_lowerings + .iter() + .map(|lowering| (lowering.backend, lowering.mechanism)) + .collect(); + ensure!( + backend_pairs + == BTreeSet::from([ + ( + IntrinsicBackend::LlvmNvptx, + BackendLoweringMechanism::TypedNvvm + ), + ( + IntrinsicBackend::LibNvvm, + BackendLoweringMechanism::InlinePtx + ), + ]), + "{} must define exactly the reviewed LLVM typed and libNVVM inline-PTX lowerings", + policy.id + ); + ensure!( + policy + .backend_lowerings + .iter() + .all(|lowering| !lowering.evidence_profile.trim().is_empty()), + "{} backend lowering omits its evidence profile", + policy.id + ); + if blackwell { + ensure!( + policy.minimum_ptx == "8.6" + && policy.minimum_sm.is_none() + && policy.targets == BLACKWELL_LDMATRIX_LLVM_TARGETS, + "{} must retain the reviewed PTX 8.6 floor and exact LLVM Blackwell target set", + policy.id + ); + let llvm = policy + .backend_lowerings + .iter() + .find(|lowering| lowering.backend == IntrinsicBackend::LlvmNvptx) + .expect("validated LLVM route"); + let libnvvm = policy + .backend_lowerings + .iter() + .find(|lowering| lowering.backend == IntrinsicBackend::LibNvvm) + .expect("validated libNVVM route"); + ensure!( + llvm.targets.is_none() + && llvm.minimum_sm.is_none() + && backend_target_requirement(policy, llvm)? + == CatalogTargetRequirement { + minimum_ptx: parse_ptx_version("8.6", &policy.id)?, + hardware: parse_hardware_target_fields( + &policy.id, + BLACKWELL_LDMATRIX_LLVM_TARGETS, + None, + )?, + } + && libnvvm.targets.as_deref() == Some(BLACKWELL_LDMATRIX_LIBNVVM_TARGETS) + && libnvvm.minimum_sm.is_none() + && backend_target_requirement(policy, libnvvm)? + == CatalogTargetRequirement { + minimum_ptx: parse_ptx_version("8.6", &policy.id)?, + hardware: parse_hardware_target_fields( + &policy.id, + BLACKWELL_LDMATRIX_LIBNVVM_TARGETS, + None, + )?, + }, + "{} backend routes do not preserve the reviewed LLVM/libNVVM Blackwell target split", + policy.id + ); + } + ensure!( + policy.packed_atomic.is_none() + && policy.redux.is_none() + && policy.vote.is_none() + && policy.active_mask.is_none() + && policy.warp_match.is_none() + && policy.warp_barrier.is_none() + && policy.warp_shuffle.is_none() + && policy.dot_product.is_none() + && policy.packed_alu.is_none() + && policy.packed_conversion.is_none() + && policy.cp_async_copy.is_none() + && policy.cp_async_control.is_none() + && policy.mbarrier_basic.is_none(), + "{} mixes another generated-family contract with ldmatrix", + policy.id + ); + Ok(()) +} + +pub(in crate::resolve) fn validate_movmatrix_policy( + policy: &OverlayIntrinsic, + source: &IntrinsicSource, +) -> Result<()> { + let contract = policy + .movmatrix + .as_ref() + .context("movmatrix requires its closed contract")?; + ensure!( + policy.id == "movmatrix_trans_b16" + && policy.operation_key == "movmatrix.m8n8.trans.b16" + && matches!( + source, + IntrinsicSource::PtxNative { instruction } + if instruction == "movmatrix.sync.aligned.m8n8.trans.b16" + ) + && policy.rust_module == "matrix" + && policy.rust_name == "movmatrix_trans_b16" + && policy.rust_arguments == ["u32"] + && policy.rust_result == "u32" + && !policy.safe + && policy.must_use + && policy.public_rust_path == "cuda_intrinsics::matrix::movmatrix_trans_b16" + && policy.compatibility_rust_paths == ["cuda_device::wmma::movmatrix_trans_b16"] + && policy.dialect_op_type == "MovmatrixTransB16Op" + && policy.dialect_op_name == "nvvm.movmatrix_trans_b16" + && policy.dialect_operands == ["i32"] + && policy.dialect_results == ["i32"] + && !policy.pure + && policy.memory == "inaccessible_read_write" + && policy.convergent + && policy.execution_scope == "warp" + && policy.minimum_ptx == "7.8" + && policy.minimum_sm.as_deref() == Some("sm_75") + && policy.ptx_result == "u32" + && policy.targets == "all" + && policy.lowering == "generated_movmatrix_inline_ptx" + && policy.expected_ptx + == InstructionPattern { + mnemonic: "movmatrix".into(), + modifiers: ["sync", "aligned", "m8n8", "trans", "b16"] + .into_iter() + .map(str::to_owned) + .collect(), + operands: vec![OperandPattern::Register, OperandPattern::Register], + }, + "{} is outside the closed movmatrix recipe", + policy.id + ); + ensure!( + contract.participation == MovmatrixParticipation::AllWarpLanesSameInstructionNoExitedLanes + && contract.adapter == MovmatrixAdapter::PackedB16x2U32ToPackedB16x2U32, + "{} has an unreviewed movmatrix safety or adapter contract", + policy.id + ); + ensure_exact_inline_ptx_backends( + policy, + [ + (IntrinsicBackend::LlvmNvptx, "7.8", Some("sm_75")), + (IntrinsicBackend::LibNvvm, "7.8", Some("sm_75")), + ], + "movmatrix", + )?; + ensure_no_other_family_contract(policy, "movmatrix")?; + Ok(()) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/mbarrier_extended.rs b/crates/cuda-intrinsics-gen/src/resolve/families/mbarrier_extended.rs new file mode 100644 index 0000000000..50800ed4e4 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/mbarrier_extended.rs @@ -0,0 +1,835 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, ImportedIntrinsic, IntrinsicBackend, IntrinsicSource, + MbarrierExtended, MbarrierExtendedAdapter, MbarrierExtendedAdmission, + MbarrierExtendedOperation, MbarrierExtendedSourceContract, OverlayBackendLowering, + OverlayIntrinsic, RuntimeValidation, +}; +use crate::ptx::{InstructionPattern, OperandPattern}; +use anyhow::{Context, Result, ensure}; +use std::collections::BTreeSet; + +use crate::resolve::guards::*; + +#[derive(Clone)] +pub(in crate::resolve) struct MbarrierExtendedRecipe { + pub(in crate::resolve) operation: MbarrierExtendedOperation, + pub(in crate::resolve) abi_id: &'static str, + pub(in crate::resolve) id: &'static str, + pub(in crate::resolve) operation_key: &'static str, + pub(in crate::resolve) source_record: Option<&'static str>, + pub(in crate::resolve) llvm_symbol: Option<&'static str>, + pub(in crate::resolve) ptx_native_instruction: Option<&'static str>, + pub(in crate::resolve) rust_arguments: &'static [&'static str], + pub(in crate::resolve) rust_result: &'static str, + pub(in crate::resolve) must_use: bool, + pub(in crate::resolve) dialect_op_type: &'static str, + pub(in crate::resolve) dialect_op_name: &'static str, + pub(in crate::resolve) dialect_operands: &'static [&'static str], + pub(in crate::resolve) dialect_results: &'static [&'static str], + pub(in crate::resolve) llvm_arguments: &'static [&'static str], + pub(in crate::resolve) llvm_results: &'static [&'static str], + pub(in crate::resolve) llvm_properties: &'static [&'static str], + pub(in crate::resolve) adapter: MbarrierExtendedAdapter, + pub(in crate::resolve) source_contract: MbarrierExtendedSourceContract, + pub(in crate::resolve) execution_scope: &'static str, + pub(in crate::resolve) minimum_ptx: &'static str, + pub(in crate::resolve) minimum_sm: &'static str, + pub(in crate::resolve) ptx_result: &'static str, + pub(in crate::resolve) expected_ptx: InstructionPattern, + pub(in crate::resolve) inline_ptx: &'static str, + pub(in crate::resolve) inline_constraints: &'static str, + pub(in crate::resolve) ptx_isa_section: &'static str, + pub(in crate::resolve) ptx_isa_url: &'static str, + pub(in crate::resolve) summary: &'static str, +} + +pub(in crate::resolve) fn mbarrier_extended_recipe( + operation: MbarrierExtendedOperation, +) -> MbarrierExtendedRecipe { + let instruction = |modifiers: &[&str], operands| InstructionPattern { + mnemonic: if modifiers.first() == Some(&"nanosleep") { + "nanosleep".into() + } else if modifiers.first() == Some(&"fence") { + "fence".into() + } else { + "mbarrier".into() + }, + modifiers: modifiers[1..].iter().map(|value| (*value).into()).collect(), + operands, + }; + match operation { + MbarrierExtendedOperation::ArriveExpectTxCta => MbarrierExtendedRecipe { + operation, + abi_id: "i0306", + id: "mbarrier_arrive_expect_tx", + operation_key: "barrier.mbarrier.arrive.expect_tx.shared.cta.release.cta", + source_record: Some("int_nvvm_mbarrier_arrive_expect_tx_scope_cta_space_cta"), + llvm_symbol: Some("llvm.nvvm.mbarrier.arrive.expect.tx.scope.cta.space.cta"), + ptx_native_instruction: None, + rust_arguments: &["*const u64", "u32", "u32"], + rust_result: "u64", + must_use: true, + dialect_op_type: "MbarrierArriveExpectTxSharedOp", + dialect_op_name: "nvvm.mbarrier_arrive_expect_tx_shared", + dialect_operands: &["ptr", "i32"], + dialect_results: &["i64"], + llvm_arguments: &["shared_ptr", "i32"], + llvm_results: &["i64"], + llvm_properties: &["IntrConvergent", "IntrNoCallback"], + adapter: MbarrierExtendedAdapter::PointerTxCountBytesToTokenDroppingTxCount, + source_contract: MbarrierExtendedSourceContract::LlvmImported, + execution_scope: "cta", + minimum_ptx: "8.0", + minimum_sm: "sm_90", + ptx_result: "u64", + expected_ptx: instruction( + &[ + "mbarrier", + "arrive", + "expect_tx", + "release", + "cta", + "shared::cta", + "b64", + ], + vec![ + OperandPattern::Register, + OperandPattern::Address, + OperandPattern::Register, + ], + ), + inline_ptx: "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 $0, [$1], $2;", + inline_constraints: "=l,l,r,~{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 CTA-shared barrier and adds expected transaction bytes.", + }, + MbarrierExtendedOperation::ArriveExpectTxCluster => MbarrierExtendedRecipe { + operation, + abi_id: "i0307", + id: "mbarrier_arrive_expect_tx_cluster", + operation_key: "barrier.mbarrier.arrive.expect_tx.shared.cta.relaxed.cluster", + source_record: Some( + "int_nvvm_mbarrier_arrive_expect_tx_relaxed_scope_cluster_space_cta", + ), + llvm_symbol: Some( + "llvm.nvvm.mbarrier.arrive.expect.tx.relaxed.scope.cluster.space.cta", + ), + ptx_native_instruction: None, + rust_arguments: &["*const u64", "u32", "u32"], + rust_result: "u64", + must_use: true, + dialect_op_type: "MbarrierArriveExpectTxClusterOp", + dialect_op_name: "nvvm.mbarrier_arrive_expect_tx_cluster", + dialect_operands: &["ptr", "i32"], + dialect_results: &["i64"], + llvm_arguments: &["shared_ptr", "i32"], + llvm_results: &["i64"], + llvm_properties: &["IntrArgMemOnly", "IntrConvergent", "IntrNoCallback"], + adapter: MbarrierExtendedAdapter::PointerTxCountBytesToTokenDroppingTxCount, + source_contract: MbarrierExtendedSourceContract::LlvmImported, + execution_scope: "cluster", + minimum_ptx: "8.6", + minimum_sm: "sm_90", + ptx_result: "u64", + expected_ptx: instruction( + &[ + "mbarrier", + "arrive", + "expect_tx", + "relaxed", + "cluster", + "shared::cta", + "b64", + ], + vec![ + OperandPattern::Register, + OperandPattern::Address, + OperandPattern::Register, + ], + ), + inline_ptx: "mbarrier.arrive.expect_tx.relaxed.cluster.shared::cta.b64 $0, [$1], $2;", + inline_constraints: "=l,l,r,~{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 CTA-shared barrier with cluster-scope transaction tracking.", + }, + MbarrierExtendedOperation::ArriveRemoteCluster => MbarrierExtendedRecipe { + operation, + abi_id: "i0308", + id: "mbarrier_arrive_cluster", + operation_key: "barrier.mbarrier.arrive.shared.cluster.release.cluster.raw_address", + source_record: None, + llvm_symbol: None, + ptx_native_instruction: Some("mbarrier.arrive.release.cluster.shared::cluster.b64"), + rust_arguments: &["u64"], + rust_result: "()", + must_use: false, + dialect_op_type: "MbarrierArriveClusterOp", + dialect_op_name: "nvvm.mbarrier_arrive_cluster", + dialect_operands: &["i64"], + dialect_results: &[], + llvm_arguments: &[], + llvm_results: &[], + llvm_properties: &[], + adapter: MbarrierExtendedAdapter::RawClusterAddressToVoid, + source_contract: MbarrierExtendedSourceContract::PtxNativeRawClusterAddress, + execution_scope: "cluster", + minimum_ptx: "8.0", + minimum_sm: "sm_90", + ptx_result: "()", + expected_ptx: instruction( + &[ + "mbarrier", + "arrive", + "release", + "cluster", + "shared::cluster", + "b64", + ], + vec![ + OperandPattern::Exact { value: "_".into() }, + OperandPattern::Address, + ], + ), + inline_ptx: "mbarrier.arrive.release.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 through its raw address.", + }, + MbarrierExtendedOperation::TryWaitTokenCta => MbarrierExtendedRecipe { + operation, + abi_id: "i0309", + id: "mbarrier_try_wait", + operation_key: "barrier.mbarrier.try_wait.shared.cta.token", + source_record: Some("int_nvvm_mbarrier_try_wait_scope_cta_space_cta"), + llvm_symbol: Some("llvm.nvvm.mbarrier.try.wait.scope.cta.space.cta"), + ptx_native_instruction: None, + rust_arguments: &["*const u64", "u64"], + rust_result: "bool", + must_use: true, + dialect_op_type: "MbarrierTryWaitSharedOp", + dialect_op_name: "nvvm.mbarrier_try_wait_shared", + dialect_operands: &["ptr", "i64"], + dialect_results: &["i1"], + llvm_arguments: &["shared_ptr", "i64"], + llvm_results: &["i1"], + llvm_properties: &["IntrConvergent", "IntrNoCallback", "NoCapture"], + adapter: MbarrierExtendedAdapter::PointerTokenToPredicate, + source_contract: MbarrierExtendedSourceContract::LlvmImported, + execution_scope: "cta", + minimum_ptx: "7.8", + minimum_sm: "sm_90", + ptx_result: "bool", + expected_ptx: instruction( + &["mbarrier", "try_wait", "shared", "b64"], + vec![ + OperandPattern::Register, + OperandPattern::Address, + OperandPattern::Register, + ], + ), + inline_ptx: "{ .reg .pred %p0; mbarrier.try_wait.shared.b64 %p0, [$1], $2; selp.b32 $0, 1, 0, %p0; }", + inline_constraints: "=r,l,l,~{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 a CTA-shared barrier token with a scheduling hint.", + }, + MbarrierExtendedOperation::TryWaitParityCta => MbarrierExtendedRecipe { + operation, + abi_id: "i0310", + id: "mbarrier_try_wait_parity", + operation_key: "barrier.mbarrier.try_wait.parity.shared.cta", + source_record: Some("int_nvvm_mbarrier_try_wait_parity_scope_cta_space_cta"), + llvm_symbol: Some("llvm.nvvm.mbarrier.try.wait.parity.scope.cta.space.cta"), + ptx_native_instruction: None, + rust_arguments: &["*const u64", "u32"], + rust_result: "bool", + must_use: true, + dialect_op_type: "MbarrierTryWaitParitySharedOp", + dialect_op_name: "nvvm.mbarrier_try_wait_parity_shared", + dialect_operands: &["ptr", "i32"], + dialect_results: &["i1"], + llvm_arguments: &["shared_ptr", "i32"], + llvm_results: &["i1"], + llvm_properties: &["IntrConvergent", "IntrNoCallback", "NoCapture"], + adapter: MbarrierExtendedAdapter::PointerParityToPredicate, + source_contract: MbarrierExtendedSourceContract::LlvmImported, + execution_scope: "cta", + minimum_ptx: "7.8", + minimum_sm: "sm_90", + ptx_result: "bool", + expected_ptx: instruction( + &["mbarrier", "try_wait", "parity", "shared::cta", "b64"], + vec![ + OperandPattern::Register, + OperandPattern::Address, + OperandPattern::Register, + ], + ), + inline_ptx: "{ .reg .pred %p0; mbarrier.try_wait.parity.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 a CTA-shared barrier phase by parity.", + }, + MbarrierExtendedOperation::TryWaitParityCluster => MbarrierExtendedRecipe { + operation, + abi_id: "i0311", + id: "mbarrier_try_wait_parity_cluster", + operation_key: "barrier.mbarrier.try_wait.parity.shared.cta.acquire.cluster", + source_record: Some("int_nvvm_mbarrier_try_wait_parity_scope_cluster_space_cta"), + llvm_symbol: Some("llvm.nvvm.mbarrier.try.wait.parity.scope.cluster.space.cta"), + ptx_native_instruction: None, + rust_arguments: &["*const u64", "u32"], + rust_result: "bool", + must_use: true, + dialect_op_type: "MbarrierTryWaitParityClusterOp", + dialect_op_name: "nvvm.mbarrier_try_wait_parity_cluster", + dialect_operands: &["ptr", "i32"], + dialect_results: &["i1"], + llvm_arguments: &["shared_ptr", "i32"], + llvm_results: &["i1"], + llvm_properties: &["IntrConvergent", "IntrNoCallback", "NoCapture"], + adapter: MbarrierExtendedAdapter::PointerParityToPredicate, + source_contract: MbarrierExtendedSourceContract::LlvmImported, + execution_scope: "cluster", + minimum_ptx: "8.0", + minimum_sm: "sm_90", + ptx_result: "bool", + expected_ptx: instruction( + &[ + "mbarrier", + "try_wait", + "parity", + "acquire", + "cluster", + "shared::cta", + "b64", + ], + vec![ + OperandPattern::Register, + OperandPattern::Address, + OperandPattern::Register, + ], + ), + inline_ptx: "{ .reg .pred %p0; mbarrier.try_wait.parity.acquire.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 with cluster-scope acquire ordering.", + }, + MbarrierExtendedOperation::FenceProxyAsyncSharedCta => MbarrierExtendedRecipe { + operation, + abi_id: "i0312", + id: "fence_proxy_async_shared_cta", + operation_key: "fence.proxy.async.shared.cta", + source_record: Some("int_nvvm_fence_proxy_async_shared_cta"), + llvm_symbol: Some("llvm.nvvm.fence.proxy.async.shared_cta"), + ptx_native_instruction: None, + rust_arguments: &[], + rust_result: "()", + must_use: false, + dialect_op_type: "FenceProxyAsyncSharedCtaOp", + dialect_op_name: "nvvm.fence_proxy_async_shared_cta", + dialect_operands: &[], + dialect_results: &[], + llvm_arguments: &[], + llvm_results: &[], + llvm_properties: &["IntrNoCallback"], + adapter: MbarrierExtendedAdapter::ZeroOperandsToVoid, + source_contract: MbarrierExtendedSourceContract::LlvmImported, + execution_scope: "cta", + minimum_ptx: "8.0", + minimum_sm: "sm_90", + ptx_result: "()", + expected_ptx: instruction(&["fence", "proxy", "async", "shared::cta"], vec![]), + inline_ptx: "fence.proxy.async.shared::cta;", + inline_constraints: "~{memory}", + ptx_isa_section: "Parallel Synchronization and Communication Instructions: fence.proxy", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-fence-proxy", + summary: "Makes CTA-shared generic-proxy writes visible to the async proxy.", + }, + MbarrierExtendedOperation::FenceMbarrierInitReleaseCluster => MbarrierExtendedRecipe { + operation, + abi_id: "i0313", + id: "fence_mbarrier_init_release_cluster", + operation_key: "fence.mbarrier_init.release.cluster", + source_record: Some("int_nvvm_fence_mbarrier_init_release_cluster"), + llvm_symbol: Some("llvm.nvvm.fence.mbarrier_init.release.cluster"), + ptx_native_instruction: None, + rust_arguments: &[], + rust_result: "()", + must_use: false, + dialect_op_type: "FenceMbarrierInitReleaseClusterOp", + dialect_op_name: "nvvm.fence_mbarrier_init_release_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.0", + minimum_sm: "sm_90", + ptx_result: "()", + expected_ptx: instruction(&["fence", "mbarrier_init", "release", "cluster"], vec![]), + inline_ptx: "fence.mbarrier_init.release.cluster;", + inline_constraints: "~{memory}", + ptx_isa_section: "Parallel Synchronization and Communication Instructions: fence.mbarrier_init", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-fence-mbarrier-init", + summary: "Releases mbarrier initialization at cluster scope.", + }, + MbarrierExtendedOperation::FenceProxyAsyncGenericReleaseSharedCtaCluster => { + MbarrierExtendedRecipe { + operation, + abi_id: "i0314", + id: "fence_proxy_async_generic_release_shared_cta_cluster", + operation_key: "fence.proxy.async_generic.release.sync_restrict.shared_cta.cluster", + source_record: Some( + "int_nvvm_fence_proxy_async_generic_release_sync_restrict_space_cta_scope_cluster", + ), + llvm_symbol: Some( + "llvm.nvvm.fence.proxy.async_generic.release.sync_restrict.space.cta.scope.cluster", + ), + ptx_native_instruction: None, + rust_arguments: &[], + rust_result: "()", + must_use: false, + dialect_op_type: "FenceProxyAsyncGenericReleaseSharedCtaClusterOp", + dialect_op_name: "nvvm.fence_proxy_async_generic_release_shared_cta_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", + "proxy", + "async::generic", + "release", + "sync_restrict::shared::cta", + "cluster", + ], + vec![], + ), + inline_ptx: "fence.proxy.async::generic.release.sync_restrict::shared::cta.cluster;", + inline_constraints: "~{memory}", + ptx_isa_section: "Parallel Synchronization and Communication Instructions: fence.proxy", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-fence-proxy", + summary: "Releases CTA-shared generic-proxy writes to the async proxy at cluster scope.", + } + } + MbarrierExtendedOperation::FenceProxyAsyncGenericAcquireSharedClusterCluster => { + MbarrierExtendedRecipe { + operation, + abi_id: "i0315", + id: "fence_proxy_async_generic_acquire_shared_cluster_cluster", + operation_key: "fence.proxy.async_generic.acquire.sync_restrict.shared_cluster.cluster", + source_record: Some( + "int_nvvm_fence_proxy_async_generic_acquire_sync_restrict_space_cluster_scope_cluster", + ), + llvm_symbol: Some( + "llvm.nvvm.fence.proxy.async_generic.acquire.sync_restrict.space.cluster.scope.cluster", + ), + ptx_native_instruction: None, + rust_arguments: &[], + rust_result: "()", + must_use: false, + dialect_op_type: "FenceProxyAsyncGenericAcquireSharedClusterClusterOp", + dialect_op_name: "nvvm.fence_proxy_async_generic_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", + "proxy", + "async::generic", + "acquire", + "sync_restrict::shared::cluster", + "cluster", + ], + vec![], + ), + inline_ptx: "fence.proxy.async::generic.acquire.sync_restrict::shared::cluster.cluster;", + inline_constraints: "~{memory}", + ptx_isa_section: "Parallel Synchronization and Communication Instructions: fence.proxy", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-fence-proxy", + summary: "Acquires cluster-shared async-proxy writes through the generic proxy.", + } + } + MbarrierExtendedOperation::Nanosleep => MbarrierExtendedRecipe { + operation, + abi_id: "i0316", + id: "nanosleep", + operation_key: "thread.nanosleep.u32", + source_record: Some("int_nvvm_nanosleep"), + llvm_symbol: Some("llvm.nvvm.nanosleep"), + ptx_native_instruction: None, + rust_arguments: &["u32"], + rust_result: "()", + must_use: false, + dialect_op_type: "NanosleepOp", + dialect_op_name: "nvvm.nanosleep", + dialect_operands: &["i32"], + dialect_results: &[], + llvm_arguments: &["i32"], + llvm_results: &[], + llvm_properties: &["IntrConvergent", "IntrHasSideEffects", "IntrNoMem"], + adapter: MbarrierExtendedAdapter::NanosecondsToVoid, + source_contract: MbarrierExtendedSourceContract::LlvmImported, + execution_scope: "thread", + minimum_ptx: "6.3", + minimum_sm: "sm_70", + ptx_result: "()", + expected_ptx: instruction( + &["nanosleep", "u32"], + vec![OperandPattern::RegisterOrImmediate], + ), + inline_ptx: "nanosleep.u32 $0;", + inline_constraints: "r,~{memory}", + ptx_isa_section: "Parallel Synchronization and Communication Instructions: nanosleep", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-nanosleep", + summary: "Suspends the executing thread for approximately the requested nanoseconds.", + }, + } +} + +pub(crate) fn mbarrier_extended_inline_recipe( + operation: MbarrierExtendedOperation, +) -> (&'static str, &'static str) { + let recipe = mbarrier_extended_recipe(operation); + (recipe.inline_ptx, recipe.inline_constraints) +} + +pub(in crate::resolve) fn mbarrier_extended_backend_floor( + operation: MbarrierExtendedOperation, + backend: IntrinsicBackend, +) -> (&'static str, &'static str) { + let recipe = mbarrier_extended_recipe(operation); + match (operation, backend) { + (MbarrierExtendedOperation::Nanosleep, IntrinsicBackend::LibNvvm) => ("6.3", "sm_75"), + _ => (recipe.minimum_ptx, recipe.minimum_sm), + } +} + +pub(in crate::resolve) fn expand_mbarrier_extended_admission( + admission: &MbarrierExtendedAdmission, +) -> Result> { + ensure!( + admission.runtime_validation == RuntimeValidation::Unexecuted, + "extended-mbarrier runtime validation may be marked executed only with GPU evidence" + ); + ensure!( + !admission.llvm_evidence_profile.trim().is_empty() + && !admission.libnvvm_evidence_profile.trim().is_empty(), + "compact extended-mbarrier admission requires both backend evidence profiles" + ); + let expected_operations = BTreeSet::from([ + MbarrierExtendedOperation::ArriveExpectTxCta, + MbarrierExtendedOperation::ArriveExpectTxCluster, + MbarrierExtendedOperation::ArriveRemoteCluster, + MbarrierExtendedOperation::TryWaitTokenCta, + MbarrierExtendedOperation::TryWaitParityCta, + MbarrierExtendedOperation::TryWaitParityCluster, + MbarrierExtendedOperation::FenceProxyAsyncSharedCta, + MbarrierExtendedOperation::FenceMbarrierInitReleaseCluster, + MbarrierExtendedOperation::FenceProxyAsyncGenericReleaseSharedCtaCluster, + MbarrierExtendedOperation::FenceProxyAsyncGenericAcquireSharedClusterCluster, + MbarrierExtendedOperation::Nanosleep, + ]); + let actual_operations = admission + .variants + .iter() + .map(|variant| variant.operation) + .collect::>(); + ensure!( + admission.variants.len() == expected_operations.len() + && actual_operations == expected_operations, + "compact extended-mbarrier admission must contain each reviewed operation exactly once" + ); + + admission + .variants + .iter() + .map(|variant| { + let recipe = mbarrier_extended_recipe(variant.operation); + ensure!( + variant.abi_id == recipe.abi_id, + "{} must keep reserved ABI ID {}", + recipe.id, + recipe.abi_id + ); + let source = + recipe + .ptx_native_instruction + .map(|instruction| IntrinsicSource::PtxNative { + instruction: instruction.into(), + }); + Ok(OverlayIntrinsic { + id: recipe.id.into(), + abi_id: variant.abi_id.clone(), + operation_key: recipe.operation_key.into(), + family: "mbarrier_extended".into(), + source, + source_record: recipe.source_record.map(str::to_owned), + rust_module: "barrier".into(), + rust_name: recipe.id.into(), + rust_arguments: recipe + .rust_arguments + .iter() + .map(|value| (*value).into()) + .collect(), + rust_result: recipe.rust_result.into(), + safe: false, + must_use: recipe.must_use, + safe_allowlist_reason: None, + public_rust_path: format!("cuda_intrinsics::barrier::{}", recipe.id), + compatibility_rust_paths: vec![format!("cuda_device::barrier::{}", recipe.id)], + dialect_op_type: recipe.dialect_op_type.into(), + dialect_op_name: recipe.dialect_op_name.into(), + dialect_operands: recipe + .dialect_operands + .iter() + .map(|value| (*value).into()) + .collect(), + dialect_results: recipe + .dialect_results + .iter() + .map(|value| (*value).into()) + .collect(), + llvm_symbol: recipe.llvm_symbol.map(str::to_owned), + resolved_llvm_symbol: None, + llvm_arguments: recipe + .llvm_arguments + .iter() + .map(|value| (*value).into()) + .collect(), + llvm_results: recipe + .llvm_results + .iter() + .map(|value| (*value).into()) + .collect(), + pure: false, + memory: "read_write".into(), + convergent: true, + execution_scope: recipe.execution_scope.into(), + minimum_ptx: recipe.minimum_ptx.into(), + minimum_sm: Some(recipe.minimum_sm.into()), + ptx_result: recipe.ptx_result.into(), + targets: "all".into(), + ptx_isa_version: "9.3".into(), + ptx_isa_section: recipe.ptx_isa_section.into(), + ptx_isa_url: recipe.ptx_isa_url.into(), + lowering: "generated_mbarrier_extended_inline_ptx".into(), + backend_lowerings: [IntrinsicBackend::LlvmNvptx, IntrinsicBackend::LibNvvm] + .into_iter() + .map(|backend| { + let (minimum_ptx, minimum_sm) = + mbarrier_extended_backend_floor(recipe.operation, backend); + OverlayBackendLowering { + backend, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: match backend { + IntrinsicBackend::LlvmNvptx => { + admission.llvm_evidence_profile.clone() + } + IntrinsicBackend::LibNvvm => { + admission.libnvvm_evidence_profile.clone() + } + }, + targets: None, + minimum_ptx: Some(minimum_ptx.into()), + minimum_sm: Some(minimum_sm.into()), + } + }) + .collect(), + packed_atomic: None, + redux: None, + vote: None, + active_mask: None, + warp_match: None, + warp_barrier: None, + warp_shuffle: None, + dot_product: None, + packed_alu: None, + integer_minmax: None, + packed_conversion: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cp_async_copy: None, + cp_async_control: None, + cp_async_mbarrier: None, + mbarrier_basic: None, + movmatrix: None, + mbarrier_extended: Some(MbarrierExtended { + operation: recipe.operation, + adapter: recipe.adapter, + source_contract: recipe.source_contract, + runtime_validation: admission.runtime_validation, + }), + register_mma: None, + sparse_mma: None, + prmt: None, + cluster_barrier: None, + wgmma_control: None, + special_register: None, + debug_control: None, + cluster_memory: None, + clc: None, + tma: None, + tcgen05: None, + ldmatrix_variant: None, + ldmatrix_safety: None, + ldmatrix_adapter: None, + selected_address_space: None, + expected_ptx: recipe.expected_ptx, + summary: recipe.summary.into(), + }) + }) + .collect() +} + +pub(in crate::resolve) fn validate_mbarrier_extended_policy( + policy: &OverlayIntrinsic, + source: &IntrinsicSource, + declaration: Option<&ImportedIntrinsic>, +) -> Result<()> { + let contract = policy + .mbarrier_extended + .as_ref() + .with_context(|| format!("{} has no closed extended-mbarrier contract", policy.id))?; + let recipe = mbarrier_extended_recipe(contract.operation); + ensure!( + contract.adapter == recipe.adapter + && contract.source_contract == recipe.source_contract + && policy.id == recipe.id + && policy.abi_id == recipe.abi_id + && policy.operation_key == recipe.operation_key, + "{} identity or adapter does not match its closed extended-mbarrier recipe", + policy.id + ); + match recipe.source_contract { + MbarrierExtendedSourceContract::LlvmImported => ensure!( + matches!( + source, + IntrinsicSource::LlvmImported { source_record } + if Some(source_record.as_str()) == recipe.source_record + ) && policy.source.is_none() + && policy.source_record.as_deref() == recipe.source_record + && policy.llvm_symbol.as_deref() == recipe.llvm_symbol + && declaration.is_some_and(|record| { + record.source_record == recipe.source_record.unwrap() + && record.properties == recipe.llvm_properties + }), + "{} LLVM source contract changed", + policy.id + ), + MbarrierExtendedSourceContract::PtxNativeRawClusterAddress => ensure!( + matches!( + source, + IntrinsicSource::PtxNative { instruction } + if Some(instruction.as_str()) == recipe.ptx_native_instruction + ) && policy.source_record.is_none() + && policy.llvm_symbol.is_none() + && declaration.is_none(), + "{} must remain the PTX-native raw-cluster-address carrier", + policy.id + ), + } + ensure!( + policy.rust_module == "barrier" + && policy.rust_name == recipe.id + && policy.rust_arguments == recipe.rust_arguments + && policy.rust_result == recipe.rust_result + && !policy.safe + && policy.must_use == recipe.must_use + && policy.public_rust_path == format!("cuda_intrinsics::barrier::{}", recipe.id) + && policy.compatibility_rust_paths == [format!("cuda_device::barrier::{}", recipe.id)], + "{} Rust API does not match its closed extended-mbarrier recipe", + policy.id + ); + ensure!( + policy.dialect_op_type == recipe.dialect_op_type + && policy.dialect_op_name == recipe.dialect_op_name + && policy.dialect_operands == recipe.dialect_operands + && policy.dialect_results == recipe.dialect_results + && policy.llvm_arguments == recipe.llvm_arguments + && policy.llvm_results == recipe.llvm_results + && policy.lowering == "generated_mbarrier_extended_inline_ptx", + "{} carrier or lowering does not match its closed extended-mbarrier recipe", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == "read_write" + && policy.convergent + && policy.execution_scope == recipe.execution_scope, + "{} convergence, memory clobber, or execution scope changed", + policy.id + ); + ensure!( + policy.minimum_ptx == recipe.minimum_ptx + && policy.minimum_sm.as_deref() == Some(recipe.minimum_sm) + && policy.ptx_result == recipe.ptx_result + && policy.targets == "all" + && policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section == recipe.ptx_isa_section + && policy.ptx_isa_url == recipe.ptx_isa_url, + "{} target floor or PTX provenance changed", + policy.id + ); + ensure!( + policy.expected_ptx == recipe.expected_ptx, + "{} expected PTX does not match its closed extended-mbarrier recipe", + policy.id + ); + ensure_exact_inline_ptx_backends( + policy, + [ + ( + IntrinsicBackend::LlvmNvptx, + mbarrier_extended_backend_floor(recipe.operation, IntrinsicBackend::LlvmNvptx).0, + Some( + mbarrier_extended_backend_floor(recipe.operation, IntrinsicBackend::LlvmNvptx) + .1, + ), + ), + ( + IntrinsicBackend::LibNvvm, + mbarrier_extended_backend_floor(recipe.operation, IntrinsicBackend::LibNvvm).0, + Some( + mbarrier_extended_backend_floor(recipe.operation, IntrinsicBackend::LibNvvm).1, + ), + ), + ], + "extended-mbarrier", + )?; + ensure_no_other_family_contract(policy, "extended mbarrier")?; + Ok(()) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/minmax.rs b/crates/cuda-intrinsics-gen/src/resolve/families/minmax.rs new file mode 100644 index 0000000000..7a4478bd2a --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/minmax.rs @@ -0,0 +1,891 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, ExtendedMinMax, ExtendedMinMaxAdapter, ExtendedMinMaxAdmission, + ExtendedMinMaxFormat, ExtendedMinMaxNan, ExtendedMinMaxOperation, ExtendedMinMaxSubnormal, + ImportedIntrinsic, IntegerMinMaxFormat, IntegerMinMaxOperation, IntrinsicBackend, + IntrinsicSource, OverlayBackendLowering, OverlayIntrinsic, RuntimeValidation, +}; +use crate::ptx::{InstructionPattern, OperandPattern}; +use anyhow::{Context, Result, ensure}; + +use crate::resolve::abi_ledger::*; +use crate::resolve::guards::*; +use crate::resolve::targets::*; + +pub(in crate::resolve) type ExtendedMinMaxVariant = ( + ExtendedMinMaxFormat, + ExtendedMinMaxOperation, + ExtendedMinMaxSubnormal, + ExtendedMinMaxNan, + bool, +); + +pub(in crate::resolve) struct ExtendedMinMaxRecipe { + id: String, + operation_key: String, + source_record: String, + llvm_symbol: String, + llvm_type: &'static str, + rust_module: &'static str, + rust_type: &'static str, + dialect_type: &'static str, + adapter: ExtendedMinMaxAdapter, + selection_record: String, + selection_asm: String, + predicates: Vec, + classes: Vec<&'static str>, + modifiers: Vec, + minimum_ptx: &'static str, + minimum_sm: &'static str, + ptx_isa_section: &'static str, + ptx_isa_url: &'static str, +} + +pub(in crate::resolve) fn canonical_extended_minmax_variants() -> Vec { + use ExtendedMinMaxFormat::{Bf16, Bf16x2, F16, F16x2, F32}; + use ExtendedMinMaxNan::{Nan, Number}; + use ExtendedMinMaxOperation::{Max, Min}; + use ExtendedMinMaxSubnormal::{Ftz, Preserve}; + + let one_operation = |operation| { + [ + (F16x2, operation, Ftz, Number, false), + (F16x2, operation, Ftz, Nan, false), + (F32, operation, Ftz, Nan, true), + (F16x2, operation, Ftz, Nan, true), + (F32, operation, Ftz, Number, true), + (F16x2, operation, Ftz, Number, true), + (Bf16x2, operation, Preserve, Nan, false), + (F16x2, operation, Preserve, Nan, false), + (Bf16x2, operation, Preserve, Nan, true), + (F32, operation, Preserve, Nan, true), + (F16x2, operation, Preserve, Nan, true), + (Bf16x2, operation, Preserve, Number, true), + (F32, operation, Preserve, Number, true), + (F16x2, operation, Preserve, Number, true), + ] + }; + // The scalar 16-bit forms. LLVM declares every `f16` modifier combination + // and the four `bf16` combinations that do not request `ftz`; the four + // `bf16` `ftz` declarations exist but have no NVPTX selection pattern and + // fail instruction selection, so they are deliberately absent here. + // + // ABI identity is intentionally absent from this canonical list. Existing + // and future IDs are carried by admission entries and enforced by the + // global append-only ABI ledger. + let scalar_halves = |operation| { + [ + (F16, operation, Preserve, Number, false), + (F16, operation, Ftz, Number, false), + (F16, operation, Preserve, Nan, false), + (F16, operation, Ftz, Nan, false), + (F16, operation, Preserve, Number, true), + (F16, operation, Ftz, Number, true), + (F16, operation, Preserve, Nan, true), + (F16, operation, Ftz, Nan, true), + (Bf16, operation, Preserve, Number, false), + (Bf16, operation, Preserve, Nan, false), + (Bf16, operation, Preserve, Number, true), + (Bf16, operation, Preserve, Nan, true), + ] + }; + one_operation(Min) + .into_iter() + .chain(one_operation(Max)) + .chain(scalar_halves(Min)) + .chain(scalar_halves(Max)) + .collect() +} + +/// Joins an operation name, its modifiers, and its format into one identifier. +/// +/// `min.f16` and `min.bf16` carry no modifier at all, so the separators have to +/// come from joining the parts that are present rather than from a format +/// string with fixed separators. +pub(in crate::resolve) fn extended_minmax_joined_name( + leading: &str, + trailing: &str, + modifiers: &[&str], + separator: &str, +) -> String { + let mut parts = vec![leading]; + parts.extend(modifiers.iter().copied()); + parts.push(trailing); + parts.join(separator) +} + +pub(in crate::resolve) fn extended_minmax_format_name( + format: ExtendedMinMaxFormat, +) -> &'static str { + match format { + ExtendedMinMaxFormat::F32 => "f32", + ExtendedMinMaxFormat::F16 => "f16", + ExtendedMinMaxFormat::Bf16 => "bf16", + ExtendedMinMaxFormat::F16x2 => "f16x2", + ExtendedMinMaxFormat::Bf16x2 => "bf16x2", + } +} + +pub(in crate::resolve) fn extended_minmax_operation_name( + operation: ExtendedMinMaxOperation, +) -> &'static str { + match operation { + ExtendedMinMaxOperation::Min => "min", + ExtendedMinMaxOperation::Max => "max", + } +} + +pub(in crate::resolve) fn extended_minmax_recipe( + variant: ExtendedMinMaxVariant, +) -> Option { + if !canonical_extended_minmax_variants().contains(&variant) { + return None; + } + let (format, operation, subnormal, nan, xorsign_abs) = variant; + let operation_name = extended_minmax_operation_name(operation); + let format_name = extended_minmax_format_name(format); + let source_type = if format == ExtendedMinMaxFormat::F32 { + "f" + } else { + format_name + }; + let mut source_modifiers = Vec::new(); + let mut ptx_modifiers = Vec::new(); + if subnormal == ExtendedMinMaxSubnormal::Ftz { + source_modifiers.push("ftz"); + ptx_modifiers.push("ftz".to_owned()); + } + if nan == ExtendedMinMaxNan::Nan { + source_modifiers.push("nan"); + ptx_modifiers.push("NaN".to_owned()); + } + if xorsign_abs { + source_modifiers.extend(["xorsign", "abs"]); + ptx_modifiers.extend(["xorsign".to_owned(), "abs".to_owned()]); + } + let id = extended_minmax_joined_name(operation_name, format_name, &source_modifiers, "_"); + let intrinsic_name = format!("f{operation_name}"); + let source_record = format!( + "int_nvvm_{}", + extended_minmax_joined_name(&intrinsic_name, source_type, &source_modifiers, "_") + ); + let llvm_symbol = format!( + "llvm.nvvm.{}", + extended_minmax_joined_name(&intrinsic_name, source_type, &source_modifiers, ".") + ); + let selection_record = if format == ExtendedMinMaxFormat::F32 { + format!( + "INT_NVVM_{}", + source_record + .trim_start_matches("int_nvvm_") + .to_ascii_uppercase() + ) + } else { + let prefix = match operation { + ExtendedMinMaxOperation::Min => "FMIN", + ExtendedMinMaxOperation::Max => "FMAN", + }; + let selection_modifiers = source_modifiers + .iter() + .map(|modifier| if *modifier == "nan" { "NaN" } else { modifier }) + .collect::>(); + format!( + "INT_NVVM_{}", + extended_minmax_joined_name(prefix, format_name, &selection_modifiers, "_") + ) + }; + let ptx_format = if format == ExtendedMinMaxFormat::F32 { + "f32" + } else { + format_name + }; + ptx_modifiers.push(ptx_format.to_owned()); + let selection_asm = format!( + "{operation_name}.{} \t$dst, $src0, $src1;", + ptx_modifiers.join(".") + ); + let (minimum_ptx, minimum_sm, predicates) = if xorsign_abs { + ( + "7.2", + "sm_86", + vec![ + "Subtarget->getPTXVersion() >= 72".to_owned(), + "Subtarget->getSmVersion() >= 86".to_owned(), + ], + ) + } else { + ( + "7.0", + "sm_80", + vec![ + "Subtarget->getSmVersion() >= 80".to_owned(), + "Subtarget->getPTXVersion() >= 70".to_owned(), + ], + ) + }; + let (llvm_type, rust_module, rust_type, dialect_type, adapter, classes) = match format { + ExtendedMinMaxFormat::F32 => ( + "f32", + "float", + "f32", + "f32", + ExtendedMinMaxAdapter::DirectF32, + vec![ + "ClangBuiltin", + "NVVMBuiltin", + "SDPatternOperator", + "Intrinsic", + "DefaultAttrsIntrinsic", + ], + ), + ExtendedMinMaxFormat::F16 => ( + "f16", + "f16", + "u16", + "i16", + ExtendedMinMaxAdapter::DirectHalfU16, + vec!["SDPatternOperator", "Intrinsic", "DefaultAttrsIntrinsic"], + ), + ExtendedMinMaxFormat::Bf16 => ( + "bf16", + "bf16", + "u16", + "i16", + ExtendedMinMaxAdapter::DirectHalfU16, + vec![ + "ClangBuiltin", + "NVVMBuiltin", + "SDPatternOperator", + "Intrinsic", + "DefaultAttrsIntrinsic", + ], + ), + ExtendedMinMaxFormat::F16x2 => ( + "v2f16", + "f16x2", + "u32", + "i32", + ExtendedMinMaxAdapter::DirectPackedU32, + vec!["SDPatternOperator", "Intrinsic", "DefaultAttrsIntrinsic"], + ), + ExtendedMinMaxFormat::Bf16x2 => ( + "v2bf16", + "bf16x2", + "u32", + "i32", + ExtendedMinMaxAdapter::DirectPackedU32, + vec![ + "ClangBuiltin", + "NVVMBuiltin", + "SDPatternOperator", + "Intrinsic", + "DefaultAttrsIntrinsic", + ], + ), + }; + let (ptx_isa_section, ptx_isa_url) = match (format, operation) { + (ExtendedMinMaxFormat::F32, ExtendedMinMaxOperation::Min) => ( + "9.7.3.9 Floating Point Instructions: min", + "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-min", + ), + (ExtendedMinMaxFormat::F32, ExtendedMinMaxOperation::Max) => ( + "9.7.3.10 Floating Point Instructions: max", + "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-max", + ), + (_, ExtendedMinMaxOperation::Min) => ( + "9.7.4.7 Half Precision Floating Point Instructions: min", + "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-min", + ), + (_, ExtendedMinMaxOperation::Max) => ( + "9.7.4.8 Half Precision Floating Point Instructions: max", + "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-max", + ), + }; + Some(ExtendedMinMaxRecipe { + id, + operation_key: format!( + "floating.minmax.{format_name}.{operation_name}.{}.{}.{}", + match subnormal { + ExtendedMinMaxSubnormal::Preserve => "preserve", + ExtendedMinMaxSubnormal::Ftz => "ftz", + }, + match nan { + ExtendedMinMaxNan::Number => "number", + ExtendedMinMaxNan::Nan => "nan", + }, + if xorsign_abs { "xorsign_abs" } else { "direct" } + ), + source_record, + llvm_symbol, + llvm_type, + rust_module, + rust_type, + dialect_type, + adapter, + selection_record, + selection_asm, + predicates, + classes, + modifiers: ptx_modifiers, + minimum_ptx, + minimum_sm, + ptx_isa_section, + ptx_isa_url, + }) +} + +pub(in crate::resolve) fn expand_extended_minmax_admission( + admission: &ExtendedMinMaxAdmission, +) -> Result> { + ensure!( + admission.runtime_validation == RuntimeValidation::Unexecuted, + "extended-minmax runtime may be marked executed only with GPU evidence" + ); + let expected = canonical_extended_minmax_variants(); + let actual = admission + .variants + .iter() + .map(|variant| { + ( + variant.format, + variant.operation, + variant.subnormal, + variant.nan, + variant.xorsign_abs, + ) + }) + .collect::>(); + ensure!( + actual == expected, + "compact extended-minmax admission must list the exact canonical 28 variants" + ); + admission + .variants + .iter() + .map(|variant| { + validate_abi_id(&variant.abi_id)?; + let identity = ( + variant.format, + variant.operation, + variant.subnormal, + variant.nan, + variant.xorsign_abs, + ); + let recipe = extended_minmax_recipe(identity) + .context("extended min/max is outside the closed recipe set")?; + extended_minmax_overlay_record(recipe, admission, identity, &variant.abi_id) + }) + .collect() +} + +pub(in crate::resolve) fn extended_minmax_overlay_record( + recipe: ExtendedMinMaxRecipe, + admission: &ExtendedMinMaxAdmission, + variant: ExtendedMinMaxVariant, + abi_id: &str, +) -> Result { + let (format, operation, subnormal, nan, xorsign_abs) = variant; + let rust_arguments = vec![recipe.rust_type.to_owned(); 2]; + let dialect_operands = vec![recipe.dialect_type.to_owned(); 2]; + Ok(OverlayIntrinsic { + id: recipe.id.clone(), + abi_id: abi_id.into(), + operation_key: recipe.operation_key, + family: "extended_minmax".into(), + source: None, + source_record: Some(recipe.source_record), + rust_module: recipe.rust_module.into(), + rust_name: recipe.id.clone(), + rust_arguments: rust_arguments.clone(), + rust_result: recipe.rust_type.into(), + safe: true, + must_use: true, + safe_allowlist_reason: Some("Floating-point min/max has no caller obligations.".into()), + public_rust_path: format!("cuda_intrinsics::{}::{}", recipe.rust_module, recipe.id), + compatibility_rust_paths: vec![format!( + "cuda_device::{}::{}", + recipe.rust_module, recipe.id + )], + dialect_op_type: "ExtendedMinMaxOp".into(), + dialect_op_name: "nvvm.extended_minmax".into(), + dialect_operands: dialect_operands.clone(), + dialect_results: vec![recipe.dialect_type.into()], + llvm_symbol: Some(recipe.llvm_symbol), + resolved_llvm_symbol: None, + llvm_arguments: vec![recipe.llvm_type.into(); 2], + llvm_results: vec![recipe.llvm_type.into()], + pure: true, + memory: "none".into(), + convergent: false, + execution_scope: "thread".into(), + minimum_ptx: recipe.minimum_ptx.into(), + minimum_sm: Some(recipe.minimum_sm.into()), + ptx_result: recipe.rust_type.into(), + targets: "all".into(), + ptx_isa_version: "9.3".into(), + ptx_isa_section: recipe.ptx_isa_section.into(), + ptx_isa_url: recipe.ptx_isa_url.into(), + lowering: "generated_extended_minmax".into(), + backend_lowerings: [ + ( + IntrinsicBackend::LlvmNvptx, + &admission.llvm_evidence_profile, + ), + ( + IntrinsicBackend::LibNvvm, + &admission.libnvvm_evidence_profile, + ), + ] + .into_iter() + .map(|(backend, evidence_profile)| OverlayBackendLowering { + backend, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: evidence_profile.clone(), + targets: None, + minimum_ptx: Some(recipe.minimum_ptx.into()), + minimum_sm: Some(recipe.minimum_sm.into()), + }) + .collect(), + packed_atomic: None, + redux: None, + vote: None, + active_mask: None, + warp_match: None, + warp_barrier: None, + warp_shuffle: None, + dot_product: None, + packed_alu: None, + integer_minmax: None, + packed_conversion: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: Some(ExtendedMinMax { + format, + operation, + subnormal, + nan, + xorsign_abs, + adapter: recipe.adapter, + runtime_validation: admission.runtime_validation, + }), + cp_async_copy: None, + cp_async_control: None, + cp_async_mbarrier: None, + mbarrier_basic: None, + movmatrix: None, + mbarrier_extended: None, + register_mma: None, + sparse_mma: None, + prmt: None, + cluster_barrier: None, + wgmma_control: None, + special_register: None, + debug_control: None, + cluster_memory: None, + clc: None, + tma: None, + tcgen05: None, + ldmatrix_variant: None, + ldmatrix_safety: None, + ldmatrix_adapter: None, + selected_address_space: None, + expected_ptx: InstructionPattern { + mnemonic: extended_minmax_operation_name(operation).into(), + modifiers: recipe.modifiers, + operands: vec![OperandPattern::Register; 3], + }, + summary: format!( + "Computes extended {} for {} values.", + extended_minmax_operation_name(operation), + extended_minmax_format_name(format) + ), + }) +} + +pub(in crate::resolve) fn validate_extended_minmax_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, +) -> Result<()> { + let minmax = policy + .extended_minmax + .as_ref() + .with_context(|| format!("{} has no extended-minmax contract", policy.id))?; + let identity = ( + minmax.format, + minmax.operation, + minmax.subnormal, + minmax.nan, + minmax.xorsign_abs, + ); + let recipe = extended_minmax_recipe(identity) + .with_context(|| format!("{} is outside the closed extended-minmax recipe", policy.id))?; + ensure!( + minmax.runtime_validation == RuntimeValidation::Unexecuted, + "{} extended-minmax runtime may be executed only with GPU evidence", + policy.id + ); + let expected_adapter = match minmax.format { + ExtendedMinMaxFormat::F32 => ExtendedMinMaxAdapter::DirectF32, + ExtendedMinMaxFormat::F16 | ExtendedMinMaxFormat::Bf16 => { + ExtendedMinMaxAdapter::DirectHalfU16 + } + ExtendedMinMaxFormat::F16x2 | ExtendedMinMaxFormat::Bf16x2 => { + ExtendedMinMaxAdapter::DirectPackedU32 + } + }; + ensure!( + minmax.adapter == expected_adapter, + "{} extended-minmax adapter does not match its format", + policy.id + ); + ensure!( + policy.id == recipe.id + && policy.operation_key == recipe.operation_key + && policy.source_record.as_deref() == Some(recipe.source_record.as_str()) + && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol.as_str()) + && policy.resolved_llvm_symbol.is_none() + && policy.llvm_arguments == vec![recipe.llvm_type; 2] + && policy.llvm_results == [recipe.llvm_type], + "{} extended-minmax identity or LLVM source changed", + policy.id + ); + ensure!( + declaration.classes == recipe.classes + && declaration.properties == ["Commutative", "IntrNoMem", "IntrSpeculatable"] + && declaration.selections.len() == 1, + "{} imported extended-minmax classes, properties, or selection count changed", + policy.id + ); + let selection = &declaration.selections[0]; + ensure!( + selection.source_record == recipe.selection_record + && selection.asm == recipe.selection_asm + && selection.predicates == recipe.predicates + && selection.constraints.is_empty(), + "{} imported extended-minmax instruction selection changed", + policy.id + ); + let rust_arguments = vec![recipe.rust_type; 2]; + let dialect_operands = vec![recipe.dialect_type; 2]; + ensure!( + policy.rust_module == recipe.rust_module + && policy.rust_name == recipe.id + && policy.rust_arguments == rust_arguments + && policy.rust_result == recipe.rust_type + && policy.safe + && policy.must_use + && policy.public_rust_path + == format!("cuda_intrinsics::{}::{}", recipe.rust_module, recipe.id) + && policy.compatibility_rust_paths + == [format!( + "cuda_device::{}::{}", + recipe.rust_module, recipe.id + )] + && policy.dialect_op_type == "ExtendedMinMaxOp" + && policy.dialect_op_name == "nvvm.extended_minmax" + && policy.dialect_operands == dialect_operands + && policy.dialect_results == [recipe.dialect_type] + && policy.lowering == "generated_extended_minmax", + "{} changed its extended-minmax API, carrier, or lowering", + policy.id + ); + ensure!( + policy.pure + && policy.memory == "none" + && !policy.convergent + && policy.execution_scope == "thread" + && policy.minimum_ptx == recipe.minimum_ptx + && policy.minimum_sm.as_deref() == Some(recipe.minimum_sm) + && policy.ptx_result == recipe.rust_type + && policy.targets == "all" + && policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section == recipe.ptx_isa_section + && policy.ptx_isa_url == recipe.ptx_isa_url, + "{} extended-minmax effects, provenance, or target floor changed", + policy.id + ); + ensure!( + policy.expected_ptx + == InstructionPattern { + mnemonic: extended_minmax_operation_name(minmax.operation).into(), + modifiers: recipe.modifiers, + operands: vec![OperandPattern::Register; 3], + }, + "{} expected extended-minmax PTX changed", + policy.id + ); + ensure_exact_inline_ptx_backends( + policy, + [ + ( + IntrinsicBackend::LlvmNvptx, + recipe.minimum_ptx, + Some(recipe.minimum_sm), + ), + ( + IntrinsicBackend::LibNvvm, + recipe.minimum_ptx, + Some(recipe.minimum_sm), + ), + ], + "extended-minmax", + )?; + validate_selected_target_predicates(policy, selection)?; + ensure_no_other_family_contract(policy, "extended min/max")?; + Ok(()) +} + +pub(in crate::resolve) struct IntegerMinMaxRecipe { + id: &'static str, + abi_id: &'static str, + operation_key: &'static str, + rust_module: &'static str, + rust_name: &'static str, + /// Rust and PTX carrier type: `i32` for the scalar forms, `u32` for the + /// packed pairs. + scalar: &'static str, + dialect_op_type: &'static str, + dialect_op_name: &'static str, + ptx_mnemonic: &'static str, + modifiers: &'static [&'static str], +} + +/// Returns the closed recipe for an extended integer min/max variant. +/// +/// `None` means the combination is outside the family: `.relu` exists only +/// for the signed formats, and the plain scalar forms are ordinary codegen +/// (`min.s32`/`max.s32` from generic Rust min/max), not intrinsics. +pub(in crate::resolve) fn integer_minmax_recipe( + format: IntegerMinMaxFormat, + operation: IntegerMinMaxOperation, + relu: bool, +) -> Option { + use IntegerMinMaxFormat as Format; + use IntegerMinMaxOperation as Operation; + Some(match (format, operation, relu) { + (Format::S32, Operation::Min, true) => IntegerMinMaxRecipe { + id: "min_relu_s32", + abi_id: "i0987", + operation_key: "integer.minmax.s32.min.relu", + rust_module: "int", + rust_name: "min_relu_s32", + scalar: "i32", + dialect_op_type: "MinReluS32Op", + dialect_op_name: "nvvm.min_relu_s32", + ptx_mnemonic: "min.relu.s32", + modifiers: &["relu", "s32"], + }, + (Format::S32, Operation::Max, true) => IntegerMinMaxRecipe { + id: "max_relu_s32", + abi_id: "i0988", + operation_key: "integer.minmax.s32.max.relu", + rust_module: "int", + rust_name: "max_relu_s32", + scalar: "i32", + dialect_op_type: "MaxReluS32Op", + dialect_op_name: "nvvm.max_relu_s32", + ptx_mnemonic: "max.relu.s32", + modifiers: &["relu", "s32"], + }, + (Format::S16x2, Operation::Min, false) => IntegerMinMaxRecipe { + id: "min_s16x2", + abi_id: "i0989", + operation_key: "integer.minmax.s16x2.min", + rust_module: "i16x2", + rust_name: "min_s16x2", + scalar: "u32", + dialect_op_type: "MinS16x2Op", + dialect_op_name: "nvvm.min_s16x2", + ptx_mnemonic: "min.s16x2", + modifiers: &["s16x2"], + }, + (Format::S16x2, Operation::Max, false) => IntegerMinMaxRecipe { + id: "max_s16x2", + abi_id: "i0990", + operation_key: "integer.minmax.s16x2.max", + rust_module: "i16x2", + rust_name: "max_s16x2", + scalar: "u32", + dialect_op_type: "MaxS16x2Op", + dialect_op_name: "nvvm.max_s16x2", + ptx_mnemonic: "max.s16x2", + modifiers: &["s16x2"], + }, + (Format::U16x2, Operation::Min, false) => IntegerMinMaxRecipe { + id: "min_u16x2", + abi_id: "i0991", + operation_key: "integer.minmax.u16x2.min", + rust_module: "i16x2", + rust_name: "min_u16x2", + scalar: "u32", + dialect_op_type: "MinU16x2Op", + dialect_op_name: "nvvm.min_u16x2", + ptx_mnemonic: "min.u16x2", + modifiers: &["u16x2"], + }, + (Format::U16x2, Operation::Max, false) => IntegerMinMaxRecipe { + id: "max_u16x2", + abi_id: "i0992", + operation_key: "integer.minmax.u16x2.max", + rust_module: "i16x2", + rust_name: "max_u16x2", + scalar: "u32", + dialect_op_type: "MaxU16x2Op", + dialect_op_name: "nvvm.max_u16x2", + ptx_mnemonic: "max.u16x2", + modifiers: &["u16x2"], + }, + (Format::S16x2, Operation::Min, true) => IntegerMinMaxRecipe { + id: "min_relu_s16x2", + abi_id: "i0993", + operation_key: "integer.minmax.s16x2.min.relu", + rust_module: "i16x2", + rust_name: "min_relu_s16x2", + scalar: "u32", + dialect_op_type: "MinReluS16x2Op", + dialect_op_name: "nvvm.min_relu_s16x2", + ptx_mnemonic: "min.relu.s16x2", + modifiers: &["relu", "s16x2"], + }, + (Format::S16x2, Operation::Max, true) => IntegerMinMaxRecipe { + id: "max_relu_s16x2", + abi_id: "i0994", + operation_key: "integer.minmax.s16x2.max.relu", + rust_module: "i16x2", + rust_name: "max_relu_s16x2", + scalar: "u32", + dialect_op_type: "MaxReluS16x2Op", + dialect_op_name: "nvvm.max_relu_s16x2", + ptx_mnemonic: "max.relu.s16x2", + modifiers: &["relu", "s16x2"], + }, + (Format::S32, _, false) | (Format::U16x2, _, true) => return None, + }) +} + +pub(in crate::resolve) fn integer_minmax_isa_reference( + operation: IntegerMinMaxOperation, +) -> (&'static str, &'static str) { + match operation { + IntegerMinMaxOperation::Min => ( + "9.7.1.13 Integer Arithmetic Instructions: min", + "https://docs.nvidia.com/cuda/parallel-thread-execution/#integer-arithmetic-instructions-min", + ), + IntegerMinMaxOperation::Max => ( + "9.7.1.14 Integer Arithmetic Instructions: max", + "https://docs.nvidia.com/cuda/parallel-thread-execution/#integer-arithmetic-instructions-max", + ), + } +} + +pub(in crate::resolve) fn validate_integer_minmax_policy( + policy: &OverlayIntrinsic, + source: &IntrinsicSource, + declaration: Option<&ImportedIntrinsic>, +) -> Result<()> { + let minmax = policy + .integer_minmax + .as_ref() + .with_context(|| format!("{} has no closed integer-min/max contract", policy.id))?; + let recipe = integer_minmax_recipe(minmax.format, minmax.operation, minmax.relu) + .with_context(|| format!("{} is outside the closed integer-min/max recipe", policy.id))?; + ensure!( + policy.id == recipe.id + && policy.abi_id == recipe.abi_id + && policy.operation_key == recipe.operation_key, + "{} integer-min/max identity does not match its closed operation recipe", + policy.id + ); + ensure!( + policy.rust_module == recipe.rust_module + && policy.rust_name == recipe.rust_name + && policy.rust_arguments == [recipe.scalar, recipe.scalar] + && policy.rust_result == recipe.scalar + && policy.safe + && policy.must_use + && policy + .safe_allowlist_reason + .as_deref() + .is_some_and(|reason| !reason.trim().is_empty()) + && policy.public_rust_path + == format!( + "cuda_intrinsics::{}::{}", + recipe.rust_module, recipe.rust_name + ) + && policy.compatibility_rust_paths + == [format!( + "cuda_device::{}::{}", + recipe.rust_module, recipe.rust_name + )], + "{} must preserve its reviewed safe integer-min/max API", + policy.id + ); + ensure!( + policy.dialect_op_type == recipe.dialect_op_type + && policy.dialect_op_name == recipe.dialect_op_name + && policy.dialect_operands == ["i32", "i32"] + && policy.dialect_results == ["i32"] + && policy.lowering == "generated_integer_minmax_inline_ptx", + "{} is outside the closed integer-min/max dialect and lowering recipe", + policy.id + ); + ensure!( + policy.pure + && policy.memory == "none" + && !policy.convergent + && policy.execution_scope == "thread" + && policy.minimum_ptx == "8.0" + && policy.minimum_sm.as_deref() == Some("sm_90") + && policy.ptx_result == recipe.scalar + && policy.targets == "all" + && minmax.native_minimum_sm == 90, + "{} integer-min/max effects, carrier, or target floor disagree", + policy.id + ); + let (isa_section, isa_url) = integer_minmax_isa_reference(minmax.operation); + ensure!( + policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section == isa_section + && policy.ptx_isa_url == isa_url, + "{} integer-min/max PTX provenance does not match its reviewed instruction section", + policy.id + ); + ensure!( + policy.expected_ptx.mnemonic + == recipe.ptx_mnemonic.split('.').next().expect("PTX mnemonic") + && policy.expected_ptx.modifiers == recipe.modifiers + && policy.expected_ptx.operands == vec![OperandPattern::Register; 3], + "{} expected PTX does not match its exact integer-min/max instruction", + policy.id + ); + ensure!( + source + == &IntrinsicSource::PtxNative { + instruction: recipe.ptx_mnemonic.to_owned(), + } + && declaration.is_none(), + "{} integer-min/max source does not match its PTX-native recipe", + policy.id + ); + ensure!( + policy.backend_lowerings.len() == 1 + && policy.backend_lowerings[0].backend == IntrinsicBackend::LlvmNvptx + && policy.backend_lowerings[0].mechanism == BackendLoweringMechanism::InlinePtx + && policy.backend_lowerings[0].targets.is_none() + && policy.backend_lowerings[0].minimum_ptx.is_none() + && policy.backend_lowerings[0].minimum_sm.is_none() + && !policy.backend_lowerings[0] + .evidence_profile + .trim() + .is_empty(), + "{} integer-min/max backend route changed", + policy.id + ); + ensure_no_other_family_contract(policy, "integer_minmax")?; + Ok(()) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/mma/mod.rs b/crates/cuda-intrinsics-gen/src/resolve/families/mma/mod.rs new file mode 100644 index 0000000000..c4583a5153 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/mma/mod.rs @@ -0,0 +1,14 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +pub(in crate::resolve) mod register; +pub(in crate::resolve) mod register_admissions; +pub(in crate::resolve) mod sparse; +pub(in crate::resolve) mod validate; + +pub(in crate::resolve) use register::*; +pub(in crate::resolve) use register_admissions::*; +pub(in crate::resolve) use sparse::*; +pub(in crate::resolve) use validate::*; diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/mma/register.rs b/crates/cuda-intrinsics-gen/src/resolve/families/mma/register.rs new file mode 100644 index 0000000000..55b69af7ee --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/mma/register.rs @@ -0,0 +1,1043 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + RegisterMma, RegisterMmaAccumulator, RegisterMmaAdapter, RegisterMmaCompatibilitySource, + RegisterMmaElement, RegisterMmaLayout, RegisterMmaOperation, RegisterMmaOverflow, + RegisterMmaParticipation, RegisterMmaShape, +}; + +pub(in crate::resolve) struct RegisterMmaRecipe { + pub(in crate::resolve) id: &'static str, + pub(in crate::resolve) abi_id: &'static str, + pub(in crate::resolve) operation_key: &'static str, + pub(in crate::resolve) source_record: &'static str, + pub(in crate::resolve) llvm_symbol: &'static str, + pub(in crate::resolve) rust_arguments: &'static [&'static str], + pub(in crate::resolve) rust_result: &'static str, + pub(in crate::resolve) dialect_op_type: &'static str, + pub(in crate::resolve) dialect_op_name: &'static str, + pub(in crate::resolve) dialect_operands: &'static [&'static str], + pub(in crate::resolve) dialect_results: &'static [&'static str], + pub(in crate::resolve) llvm_arguments: &'static [&'static str], + pub(in crate::resolve) llvm_results: &'static [&'static str], + pub(in crate::resolve) adapter: RegisterMmaAdapter, + pub(in crate::resolve) compatibility_source: RegisterMmaCompatibilitySource, + pub(in crate::resolve) minimum_ptx: &'static str, + pub(in crate::resolve) minimum_sm: &'static str, + pub(in crate::resolve) ptx_modifiers: Vec<&'static str>, + pub(in crate::resolve) ptx_register_counts: [usize; 4], +} + +pub(in crate::resolve) fn integer_register_mma_recipe( + mma: &RegisterMma, + common: bool, +) -> Option { + use RegisterMmaAdapter::{ + C2I32A1U32B1U32ToD2I32, C4I32A2U32B1U32ToD4I32, C4I32A4U32B2U32ToD4I32, + }; + use RegisterMmaCompatibilitySource::{ExistingStub, GeneratedStub}; + use RegisterMmaElement::{S4, S8, U4, U8}; + use RegisterMmaOverflow::{Satfinite, Wrapping}; + use RegisterMmaShape::{M8n8k16, M8n8k32, M16n8k16, M16n8k32, M16n8k64}; + + if !common + || mma.operation != RegisterMmaOperation::Multiply + || mma.accumulator != RegisterMmaAccumulator::S32 + { + return None; + } + + let (id, abi_id, operation_key, source_record, llvm_symbol, compatibility_source) = + match (mma.shape, mma.a_element, mma.b_element, mma.overflow) { + (M16n8k32, S8, S8, Wrapping) => ( + "mma_m16n8k32_s32_s8", + "i0108", + "matrix.mma.m16n8k32.row.col.s32.s8.s8.s32.wrapping", + "int_nvvm_mma_m16n8k32_row_col_s8", + "llvm.nvvm.mma.m16n8k32.row.col.s8", + ExistingStub, + ), + (M16n8k16, S8, S8, Wrapping) => ( + "mma_m16n8k16_s32_s8", + "i0110", + "matrix.mma.m16n8k16.row.col.s32.s8.s8.s32.wrapping", + "int_nvvm_mma_m16n8k16_row_col_s8", + "llvm.nvvm.mma.m16n8k16.row.col.s8", + GeneratedStub, + ), + (M16n8k16, S8, U8, Wrapping) => ( + "mma_m16n8k16_s32_s8_u8", + "i0111", + "matrix.mma.m16n8k16.row.col.s32.s8.u8.s32.wrapping", + "int_nvvm_mma_m16n8k16_row_col_s8_u8", + "llvm.nvvm.mma.m16n8k16.row.col.s8.u8", + GeneratedStub, + ), + (M16n8k16, U8, U8, Wrapping) => ( + "mma_m16n8k16_s32_u8", + "i0112", + "matrix.mma.m16n8k16.row.col.s32.u8.u8.s32.wrapping", + "int_nvvm_mma_m16n8k16_row_col_u8", + "llvm.nvvm.mma.m16n8k16.row.col.u8", + GeneratedStub, + ), + (M16n8k16, U8, S8, Wrapping) => ( + "mma_m16n8k16_s32_u8_s8", + "i0113", + "matrix.mma.m16n8k16.row.col.s32.u8.s8.s32.wrapping", + "int_nvvm_mma_m16n8k16_row_col_u8_s8", + "llvm.nvvm.mma.m16n8k16.row.col.u8.s8", + GeneratedStub, + ), + (M16n8k32, S8, U8, Wrapping) => ( + "mma_m16n8k32_s32_s8_u8", + "i0114", + "matrix.mma.m16n8k32.row.col.s32.s8.u8.s32.wrapping", + "int_nvvm_mma_m16n8k32_row_col_s8_u8", + "llvm.nvvm.mma.m16n8k32.row.col.s8.u8", + GeneratedStub, + ), + (M16n8k32, U8, U8, Wrapping) => ( + "mma_m16n8k32_s32_u8", + "i0115", + "matrix.mma.m16n8k32.row.col.s32.u8.u8.s32.wrapping", + "int_nvvm_mma_m16n8k32_row_col_u8", + "llvm.nvvm.mma.m16n8k32.row.col.u8", + GeneratedStub, + ), + (M16n8k32, U8, S8, Wrapping) => ( + "mma_m16n8k32_s32_u8_s8", + "i0116", + "matrix.mma.m16n8k32.row.col.s32.u8.s8.s32.wrapping", + "int_nvvm_mma_m16n8k32_row_col_u8_s8", + "llvm.nvvm.mma.m16n8k32.row.col.u8.s8", + GeneratedStub, + ), + (M16n8k16, S8, S8, Satfinite) => ( + "mma_m16n8k16_s32_s8_satfinite", + "i0117", + "matrix.mma.m16n8k16.row.col.s32.s8.s8.s32.satfinite", + "int_nvvm_mma_m16n8k16_row_col_satfinite_s8", + "llvm.nvvm.mma.m16n8k16.row.col.satfinite.s8", + GeneratedStub, + ), + (M16n8k16, S8, U8, Satfinite) => ( + "mma_m16n8k16_s32_s8_u8_satfinite", + "i0118", + "matrix.mma.m16n8k16.row.col.s32.s8.u8.s32.satfinite", + "int_nvvm_mma_m16n8k16_row_col_satfinite_s8_u8", + "llvm.nvvm.mma.m16n8k16.row.col.satfinite.s8.u8", + GeneratedStub, + ), + (M16n8k16, U8, U8, Satfinite) => ( + "mma_m16n8k16_s32_u8_satfinite", + "i0119", + "matrix.mma.m16n8k16.row.col.s32.u8.u8.s32.satfinite", + "int_nvvm_mma_m16n8k16_row_col_satfinite_u8", + "llvm.nvvm.mma.m16n8k16.row.col.satfinite.u8", + GeneratedStub, + ), + (M16n8k16, U8, S8, Satfinite) => ( + "mma_m16n8k16_s32_u8_s8_satfinite", + "i0120", + "matrix.mma.m16n8k16.row.col.s32.u8.s8.s32.satfinite", + "int_nvvm_mma_m16n8k16_row_col_satfinite_u8_s8", + "llvm.nvvm.mma.m16n8k16.row.col.satfinite.u8.s8", + GeneratedStub, + ), + (M16n8k32, S8, S8, Satfinite) => ( + "mma_m16n8k32_s32_s8_satfinite", + "i0121", + "matrix.mma.m16n8k32.row.col.s32.s8.s8.s32.satfinite", + "int_nvvm_mma_m16n8k32_row_col_satfinite_s8", + "llvm.nvvm.mma.m16n8k32.row.col.satfinite.s8", + GeneratedStub, + ), + (M16n8k32, S8, U8, Satfinite) => ( + "mma_m16n8k32_s32_s8_u8_satfinite", + "i0122", + "matrix.mma.m16n8k32.row.col.s32.s8.u8.s32.satfinite", + "int_nvvm_mma_m16n8k32_row_col_satfinite_s8_u8", + "llvm.nvvm.mma.m16n8k32.row.col.satfinite.s8.u8", + GeneratedStub, + ), + (M16n8k32, U8, U8, Satfinite) => ( + "mma_m16n8k32_s32_u8_satfinite", + "i0123", + "matrix.mma.m16n8k32.row.col.s32.u8.u8.s32.satfinite", + "int_nvvm_mma_m16n8k32_row_col_satfinite_u8", + "llvm.nvvm.mma.m16n8k32.row.col.satfinite.u8", + GeneratedStub, + ), + (M16n8k32, U8, S8, Satfinite) => ( + "mma_m16n8k32_s32_u8_s8_satfinite", + "i0124", + "matrix.mma.m16n8k32.row.col.s32.u8.s8.s32.satfinite", + "int_nvvm_mma_m16n8k32_row_col_satfinite_u8_s8", + "llvm.nvvm.mma.m16n8k32.row.col.satfinite.u8.s8", + GeneratedStub, + ), + (M8n8k16, S8, S8, Wrapping) => ( + "mma_m8n8k16_s32_s8", + "i0125", + "matrix.mma.m8n8k16.row.col.s32.s8.s8.s32.wrapping", + "int_nvvm_mma_m8n8k16_row_col_s8", + "llvm.nvvm.mma.m8n8k16.row.col.s8", + GeneratedStub, + ), + (M8n8k16, S8, U8, Wrapping) => ( + "mma_m8n8k16_s32_s8_u8", + "i0126", + "matrix.mma.m8n8k16.row.col.s32.s8.u8.s32.wrapping", + "int_nvvm_mma_m8n8k16_row_col_s8_u8", + "llvm.nvvm.mma.m8n8k16.row.col.s8.u8", + GeneratedStub, + ), + (M8n8k16, U8, U8, Wrapping) => ( + "mma_m8n8k16_s32_u8", + "i0127", + "matrix.mma.m8n8k16.row.col.s32.u8.u8.s32.wrapping", + "int_nvvm_mma_m8n8k16_row_col_u8", + "llvm.nvvm.mma.m8n8k16.row.col.u8", + GeneratedStub, + ), + (M8n8k16, U8, S8, Wrapping) => ( + "mma_m8n8k16_s32_u8_s8", + "i0128", + "matrix.mma.m8n8k16.row.col.s32.u8.s8.s32.wrapping", + "int_nvvm_mma_m8n8k16_row_col_u8_s8", + "llvm.nvvm.mma.m8n8k16.row.col.u8.s8", + GeneratedStub, + ), + (M8n8k16, S8, S8, Satfinite) => ( + "mma_m8n8k16_s32_s8_satfinite", + "i0129", + "matrix.mma.m8n8k16.row.col.s32.s8.s8.s32.satfinite", + "int_nvvm_mma_m8n8k16_row_col_satfinite_s8", + "llvm.nvvm.mma.m8n8k16.row.col.satfinite.s8", + GeneratedStub, + ), + (M8n8k16, S8, U8, Satfinite) => ( + "mma_m8n8k16_s32_s8_u8_satfinite", + "i0130", + "matrix.mma.m8n8k16.row.col.s32.s8.u8.s32.satfinite", + "int_nvvm_mma_m8n8k16_row_col_satfinite_s8_u8", + "llvm.nvvm.mma.m8n8k16.row.col.satfinite.s8.u8", + GeneratedStub, + ), + (M8n8k16, U8, U8, Satfinite) => ( + "mma_m8n8k16_s32_u8_satfinite", + "i0131", + "matrix.mma.m8n8k16.row.col.s32.u8.u8.s32.satfinite", + "int_nvvm_mma_m8n8k16_row_col_satfinite_u8", + "llvm.nvvm.mma.m8n8k16.row.col.satfinite.u8", + GeneratedStub, + ), + (M8n8k16, U8, S8, Satfinite) => ( + "mma_m8n8k16_s32_u8_s8_satfinite", + "i0132", + "matrix.mma.m8n8k16.row.col.s32.u8.s8.s32.satfinite", + "int_nvvm_mma_m8n8k16_row_col_satfinite_u8_s8", + "llvm.nvvm.mma.m8n8k16.row.col.satfinite.u8.s8", + GeneratedStub, + ), + (M8n8k32, S4, S4, Wrapping) => ( + "mma_m8n8k32_s32_s4", + "i0133", + "matrix.mma.m8n8k32.row.col.s32.s4.s4.s32.wrapping", + "int_nvvm_mma_m8n8k32_row_col_s4", + "llvm.nvvm.mma.m8n8k32.row.col.s4", + GeneratedStub, + ), + (M8n8k32, S4, U4, Wrapping) => ( + "mma_m8n8k32_s32_s4_u4", + "i0134", + "matrix.mma.m8n8k32.row.col.s32.s4.u4.s32.wrapping", + "int_nvvm_mma_m8n8k32_row_col_s4_u4", + "llvm.nvvm.mma.m8n8k32.row.col.s4.u4", + GeneratedStub, + ), + (M8n8k32, U4, U4, Wrapping) => ( + "mma_m8n8k32_s32_u4", + "i0135", + "matrix.mma.m8n8k32.row.col.s32.u4.u4.s32.wrapping", + "int_nvvm_mma_m8n8k32_row_col_u4", + "llvm.nvvm.mma.m8n8k32.row.col.u4", + GeneratedStub, + ), + (M8n8k32, U4, S4, Wrapping) => ( + "mma_m8n8k32_s32_u4_s4", + "i0136", + "matrix.mma.m8n8k32.row.col.s32.u4.s4.s32.wrapping", + "int_nvvm_mma_m8n8k32_row_col_u4_s4", + "llvm.nvvm.mma.m8n8k32.row.col.u4.s4", + GeneratedStub, + ), + (M8n8k32, S4, S4, Satfinite) => ( + "mma_m8n8k32_s32_s4_satfinite", + "i0137", + "matrix.mma.m8n8k32.row.col.s32.s4.s4.s32.satfinite", + "int_nvvm_mma_m8n8k32_row_col_satfinite_s4", + "llvm.nvvm.mma.m8n8k32.row.col.satfinite.s4", + GeneratedStub, + ), + (M8n8k32, S4, U4, Satfinite) => ( + "mma_m8n8k32_s32_s4_u4_satfinite", + "i0138", + "matrix.mma.m8n8k32.row.col.s32.s4.u4.s32.satfinite", + "int_nvvm_mma_m8n8k32_row_col_satfinite_s4_u4", + "llvm.nvvm.mma.m8n8k32.row.col.satfinite.s4.u4", + GeneratedStub, + ), + (M8n8k32, U4, U4, Satfinite) => ( + "mma_m8n8k32_s32_u4_satfinite", + "i0139", + "matrix.mma.m8n8k32.row.col.s32.u4.u4.s32.satfinite", + "int_nvvm_mma_m8n8k32_row_col_satfinite_u4", + "llvm.nvvm.mma.m8n8k32.row.col.satfinite.u4", + GeneratedStub, + ), + (M8n8k32, U4, S4, Satfinite) => ( + "mma_m8n8k32_s32_u4_s4_satfinite", + "i0140", + "matrix.mma.m8n8k32.row.col.s32.u4.s4.s32.satfinite", + "int_nvvm_mma_m8n8k32_row_col_satfinite_u4_s4", + "llvm.nvvm.mma.m8n8k32.row.col.satfinite.u4.s4", + GeneratedStub, + ), + (M16n8k32, S4, S4, Wrapping) => ( + "mma_m16n8k32_s32_s4", + "i0141", + "matrix.mma.m16n8k32.row.col.s32.s4.s4.s32.wrapping", + "int_nvvm_mma_m16n8k32_row_col_s4", + "llvm.nvvm.mma.m16n8k32.row.col.s4", + GeneratedStub, + ), + (M16n8k32, S4, U4, Wrapping) => ( + "mma_m16n8k32_s32_s4_u4", + "i0142", + "matrix.mma.m16n8k32.row.col.s32.s4.u4.s32.wrapping", + "int_nvvm_mma_m16n8k32_row_col_s4_u4", + "llvm.nvvm.mma.m16n8k32.row.col.s4.u4", + GeneratedStub, + ), + (M16n8k32, U4, U4, Wrapping) => ( + "mma_m16n8k32_s32_u4", + "i0143", + "matrix.mma.m16n8k32.row.col.s32.u4.u4.s32.wrapping", + "int_nvvm_mma_m16n8k32_row_col_u4", + "llvm.nvvm.mma.m16n8k32.row.col.u4", + GeneratedStub, + ), + (M16n8k32, U4, S4, Wrapping) => ( + "mma_m16n8k32_s32_u4_s4", + "i0144", + "matrix.mma.m16n8k32.row.col.s32.u4.s4.s32.wrapping", + "int_nvvm_mma_m16n8k32_row_col_u4_s4", + "llvm.nvvm.mma.m16n8k32.row.col.u4.s4", + GeneratedStub, + ), + (M16n8k64, S4, S4, Wrapping) => ( + "mma_m16n8k64_s32_s4", + "i0145", + "matrix.mma.m16n8k64.row.col.s32.s4.s4.s32.wrapping", + "int_nvvm_mma_m16n8k64_row_col_s4", + "llvm.nvvm.mma.m16n8k64.row.col.s4", + GeneratedStub, + ), + (M16n8k64, S4, U4, Wrapping) => ( + "mma_m16n8k64_s32_s4_u4", + "i0146", + "matrix.mma.m16n8k64.row.col.s32.s4.u4.s32.wrapping", + "int_nvvm_mma_m16n8k64_row_col_s4_u4", + "llvm.nvvm.mma.m16n8k64.row.col.s4.u4", + GeneratedStub, + ), + (M16n8k64, U4, U4, Wrapping) => ( + "mma_m16n8k64_s32_u4", + "i0147", + "matrix.mma.m16n8k64.row.col.s32.u4.u4.s32.wrapping", + "int_nvvm_mma_m16n8k64_row_col_u4", + "llvm.nvvm.mma.m16n8k64.row.col.u4", + GeneratedStub, + ), + (M16n8k64, U4, S4, Wrapping) => ( + "mma_m16n8k64_s32_u4_s4", + "i0148", + "matrix.mma.m16n8k64.row.col.s32.u4.s4.s32.wrapping", + "int_nvvm_mma_m16n8k64_row_col_u4_s4", + "llvm.nvvm.mma.m16n8k64.row.col.u4.s4", + GeneratedStub, + ), + (M16n8k32, S4, S4, Satfinite) => ( + "mma_m16n8k32_s32_s4_satfinite", + "i0149", + "matrix.mma.m16n8k32.row.col.s32.s4.s4.s32.satfinite", + "int_nvvm_mma_m16n8k32_row_col_satfinite_s4", + "llvm.nvvm.mma.m16n8k32.row.col.satfinite.s4", + GeneratedStub, + ), + (M16n8k32, S4, U4, Satfinite) => ( + "mma_m16n8k32_s32_s4_u4_satfinite", + "i0150", + "matrix.mma.m16n8k32.row.col.s32.s4.u4.s32.satfinite", + "int_nvvm_mma_m16n8k32_row_col_satfinite_s4_u4", + "llvm.nvvm.mma.m16n8k32.row.col.satfinite.s4.u4", + GeneratedStub, + ), + (M16n8k32, U4, U4, Satfinite) => ( + "mma_m16n8k32_s32_u4_satfinite", + "i0151", + "matrix.mma.m16n8k32.row.col.s32.u4.u4.s32.satfinite", + "int_nvvm_mma_m16n8k32_row_col_satfinite_u4", + "llvm.nvvm.mma.m16n8k32.row.col.satfinite.u4", + GeneratedStub, + ), + (M16n8k32, U4, S4, Satfinite) => ( + "mma_m16n8k32_s32_u4_s4_satfinite", + "i0152", + "matrix.mma.m16n8k32.row.col.s32.u4.s4.s32.satfinite", + "int_nvvm_mma_m16n8k32_row_col_satfinite_u4_s4", + "llvm.nvvm.mma.m16n8k32.row.col.satfinite.u4.s4", + GeneratedStub, + ), + (M16n8k64, S4, S4, Satfinite) => ( + "mma_m16n8k64_s32_s4_satfinite", + "i0153", + "matrix.mma.m16n8k64.row.col.s32.s4.s4.s32.satfinite", + "int_nvvm_mma_m16n8k64_row_col_satfinite_s4", + "llvm.nvvm.mma.m16n8k64.row.col.satfinite.s4", + GeneratedStub, + ), + (M16n8k64, S4, U4, Satfinite) => ( + "mma_m16n8k64_s32_s4_u4_satfinite", + "i0154", + "matrix.mma.m16n8k64.row.col.s32.s4.u4.s32.satfinite", + "int_nvvm_mma_m16n8k64_row_col_satfinite_s4_u4", + "llvm.nvvm.mma.m16n8k64.row.col.satfinite.s4.u4", + GeneratedStub, + ), + (M16n8k64, U4, U4, Satfinite) => ( + "mma_m16n8k64_s32_u4_satfinite", + "i0155", + "matrix.mma.m16n8k64.row.col.s32.u4.u4.s32.satfinite", + "int_nvvm_mma_m16n8k64_row_col_satfinite_u4", + "llvm.nvvm.mma.m16n8k64.row.col.satfinite.u4", + GeneratedStub, + ), + (M16n8k64, U4, S4, Satfinite) => ( + "mma_m16n8k64_s32_u4_s4_satfinite", + "i0156", + "matrix.mma.m16n8k64.row.col.s32.u4.s4.s32.satfinite", + "int_nvvm_mma_m16n8k64_row_col_satfinite_u4_s4", + "llvm.nvvm.mma.m16n8k64.row.col.satfinite.u4.s4", + GeneratedStub, + ), + _ => return None, + }; + + let int4 = matches!(mma.a_element, S4 | U4); + let (rust_arguments, dialect_operands, llvm_arguments, adapter, shape, register_counts) = + match (mma.shape, int4) { + (M8n8k16, false) | (M8n8k32, true) => ( + &["[i32; 2]", "u32", "u32"] as &'static [&'static str], + &["i32", "i32", "i32", "i32"] as &'static [&'static str], + &["i32", "i32", "i32", "i32"] as &'static [&'static str], + C2I32A1U32B1U32ToD2I32, + match mma.shape { + M8n8k16 => "m8n8k16", + M8n8k32 => "m8n8k32", + _ => unreachable!(), + }, + [2, 1, 1, 2], + ), + (M16n8k16, false) | (M16n8k32, true) => ( + &["[i32; 4]", "[u32; 2]", "u32"] as &'static [&'static str], + &["i32", "i32", "i32", "i32", "i32", "i32", "i32"] as &'static [&'static str], + &["i32", "i32", "i32", "i32", "i32", "i32", "i32"] as &'static [&'static str], + C4I32A2U32B1U32ToD4I32, + match mma.shape { + M16n8k16 => "m16n8k16", + M16n8k32 => "m16n8k32", + _ => unreachable!(), + }, + [4, 2, 1, 4], + ), + (M16n8k32, false) | (M16n8k64, true) => ( + &["[i32; 4]", "[u32; 4]", "[u32; 2]"] as &'static [&'static str], + &[ + "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", + ] as &'static [&'static str], + &[ + "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", + ] as &'static [&'static str], + C4I32A4U32B2U32ToD4I32, + match mma.shape { + M16n8k32 => "m16n8k32", + M16n8k64 => "m16n8k64", + _ => unreachable!(), + }, + [4, 4, 2, 4], + ), + _ => return None, + }; + if mma.adapter != adapter { + return None; + } + + let (rust_result, dialect_results, llvm_results, minimum_ptx, minimum_sm) = match mma.shape { + M8n8k16 | M8n8k32 => ( + "[i32; 2]", + &["i32", "i32"] as &'static [&'static str], + &["i32", "i32"] as &'static [&'static str], + "6.5", + "sm_75", + ), + M16n8k16 | M16n8k32 | M16n8k64 => ( + "[i32; 4]", + &["i32", "i32", "i32", "i32"] as &'static [&'static str], + &["i32", "i32", "i32", "i32"] as &'static [&'static str], + "7.0", + "sm_80", + ), + _ => return None, + }; + + let element = |element| match element { + S4 => Some("s4"), + U4 => Some("u4"), + S8 => Some("s8"), + U8 => Some("u8"), + _ => None, + }; + let mut ptx_modifiers = vec!["sync", "aligned", shape, "row", "col"]; + if mma.overflow == Satfinite { + ptx_modifiers.push("satfinite"); + } + ptx_modifiers.extend([ + "s32", + element(mma.a_element)?, + element(mma.b_element)?, + "s32", + ]); + + Some(RegisterMmaRecipe { + id, + abi_id, + operation_key, + source_record, + llvm_symbol, + rust_arguments, + rust_result, + dialect_op_type: "RegisterMmaOp", + dialect_op_name: "nvvm.register_mma", + dialect_operands, + dialect_results, + llvm_arguments, + llvm_results, + adapter, + compatibility_source, + minimum_ptx, + minimum_sm, + ptx_modifiers, + ptx_register_counts: register_counts, + }) +} + +pub(in crate::resolve) fn binary_register_mma_recipe( + mma: &RegisterMma, + common: bool, +) -> Option { + use RegisterMmaAdapter::{ + C2I32A1U32B1U32ToD2I32, C4I32A2U32B1U32ToD4I32, C4I32A4U32B2U32ToD4I32, + }; + use RegisterMmaOperation::{AndPopc, XorPopc}; + use RegisterMmaShape::{M8n8k128, M16n8k128, M16n8k256}; + + if !common + || mma.accumulator != RegisterMmaAccumulator::S32 + || mma.a_element != RegisterMmaElement::B1 + || mma.b_element != RegisterMmaElement::B1 + || mma.overflow != RegisterMmaOverflow::Wrapping + || mma.compatibility_source != RegisterMmaCompatibilitySource::GeneratedStub + { + return None; + } + + let (id, abi_id, operation_key, source_record, llvm_symbol, operation_name) = + match (mma.shape, mma.operation) { + (M8n8k128, XorPopc) => ( + "mma_m8n8k128_s32_b1_xor_popc", + "i0157", + "matrix.mma.m8n8k128.row.col.s32.b1.b1.s32.xor.popc", + "int_nvvm_mma_xor_popc_m8n8k128_row_col_b1", + "llvm.nvvm.mma.xor.popc.m8n8k128.row.col.b1", + "xor", + ), + (M16n8k128, XorPopc) => ( + "mma_m16n8k128_s32_b1_xor_popc", + "i0158", + "matrix.mma.m16n8k128.row.col.s32.b1.b1.s32.xor.popc", + "int_nvvm_mma_xor_popc_m16n8k128_row_col_b1", + "llvm.nvvm.mma.xor.popc.m16n8k128.row.col.b1", + "xor", + ), + (M16n8k256, XorPopc) => ( + "mma_m16n8k256_s32_b1_xor_popc", + "i0159", + "matrix.mma.m16n8k256.row.col.s32.b1.b1.s32.xor.popc", + "int_nvvm_mma_xor_popc_m16n8k256_row_col_b1", + "llvm.nvvm.mma.xor.popc.m16n8k256.row.col.b1", + "xor", + ), + (M8n8k128, AndPopc) => ( + "mma_m8n8k128_s32_b1_and_popc", + "i0160", + "matrix.mma.m8n8k128.row.col.s32.b1.b1.s32.and.popc", + "int_nvvm_mma_and_popc_m8n8k128_row_col_b1", + "llvm.nvvm.mma.and.popc.m8n8k128.row.col.b1", + "and", + ), + (M16n8k128, AndPopc) => ( + "mma_m16n8k128_s32_b1_and_popc", + "i0161", + "matrix.mma.m16n8k128.row.col.s32.b1.b1.s32.and.popc", + "int_nvvm_mma_and_popc_m16n8k128_row_col_b1", + "llvm.nvvm.mma.and.popc.m16n8k128.row.col.b1", + "and", + ), + (M16n8k256, AndPopc) => ( + "mma_m16n8k256_s32_b1_and_popc", + "i0162", + "matrix.mma.m16n8k256.row.col.s32.b1.b1.s32.and.popc", + "int_nvvm_mma_and_popc_m16n8k256_row_col_b1", + "llvm.nvvm.mma.and.popc.m16n8k256.row.col.b1", + "and", + ), + _ => return None, + }; + + let ( + rust_arguments, + dialect_operands, + llvm_arguments, + rust_result, + result_types, + adapter, + counts, + ) = match mma.shape { + M8n8k128 => ( + &["[i32; 2]", "u32", "u32"] as &'static [&'static str], + &["i32", "i32", "i32", "i32"] as &'static [&'static str], + &["i32", "i32", "i32", "i32"] as &'static [&'static str], + "[i32; 2]", + &["i32", "i32"] as &'static [&'static str], + C2I32A1U32B1U32ToD2I32, + [2, 1, 1, 2], + ), + M16n8k128 => ( + &["[i32; 4]", "[u32; 2]", "u32"] as &'static [&'static str], + &["i32", "i32", "i32", "i32", "i32", "i32", "i32"] as &'static [&'static str], + &["i32", "i32", "i32", "i32", "i32", "i32", "i32"] as &'static [&'static str], + "[i32; 4]", + &["i32", "i32", "i32", "i32"] as &'static [&'static str], + C4I32A2U32B1U32ToD4I32, + [4, 2, 1, 4], + ), + M16n8k256 => ( + &["[i32; 4]", "[u32; 4]", "[u32; 2]"] as &'static [&'static str], + &[ + "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", + ] as &'static [&'static str], + &[ + "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", + ] as &'static [&'static str], + "[i32; 4]", + &["i32", "i32", "i32", "i32"] as &'static [&'static str], + C4I32A4U32B2U32ToD4I32, + [4, 4, 2, 4], + ), + _ => return None, + }; + if mma.adapter != adapter { + return None; + } + + let (minimum_ptx, minimum_sm) = match (mma.operation, mma.shape) { + (XorPopc, M8n8k128) => ("7.0", "sm_75"), + (XorPopc, M16n8k128 | M16n8k256) => ("7.0", "sm_80"), + (AndPopc, M8n8k128 | M16n8k128 | M16n8k256) => ("7.1", "sm_80"), + _ => return None, + }; + let shape = match mma.shape { + M8n8k128 => "m8n8k128", + M16n8k128 => "m16n8k128", + M16n8k256 => "m16n8k256", + _ => return None, + }; + + Some(RegisterMmaRecipe { + id, + abi_id, + operation_key, + source_record, + llvm_symbol, + rust_arguments, + rust_result, + dialect_op_type: "RegisterMmaOp", + dialect_op_name: "nvvm.register_mma", + dialect_operands, + dialect_results: result_types, + llvm_arguments, + llvm_results: result_types, + adapter, + compatibility_source: RegisterMmaCompatibilitySource::GeneratedStub, + minimum_ptx, + minimum_sm, + ptx_modifiers: vec![ + "sync", + "aligned", + shape, + "row", + "col", + "s32", + "b1", + "b1", + "s32", + operation_name, + "popc", + ], + ptx_register_counts: counts, + }) +} + +pub(in crate::resolve) fn register_mma_recipe(mma: &RegisterMma) -> Option { + use RegisterMmaAccumulator::{F16 as F16Accumulator, F32, F64}; + use RegisterMmaAdapter::{ + C2F64A1F64B1F64ToD2F64, C2U32A2U32B1U32ToD2U32, C2U32A4U32B2U32ToD2U32, + C4F32A2U32B1U32ToD4F32, C4F32A4U32B2U32ToD4F32, + }; + use RegisterMmaElement::{Bf16, F16 as F16Element, F64 as F64Element, Tf32}; + use RegisterMmaShape::{M8n8k4, M16n8k4, M16n8k8, M16n8k16}; + + let common = mma.a_layout == RegisterMmaLayout::Row + && mma.b_layout == RegisterMmaLayout::Col + && mma.participation + == RegisterMmaParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes; + if let Some(recipe) = integer_register_mma_recipe(mma, common) { + return Some(recipe); + } + if let Some(recipe) = binary_register_mma_recipe(mma, common) { + return Some(recipe); + } + if mma.operation != RegisterMmaOperation::Multiply { + return None; + } + match ( + mma.shape, + mma.accumulator, + mma.a_element, + mma.b_element, + mma.overflow, + mma.adapter, + common, + ) { + ( + M16n8k4, + F32, + Tf32, + Tf32, + RegisterMmaOverflow::NotApplicable, + C4F32A2U32B1U32ToD4F32, + true, + ) => Some(RegisterMmaRecipe { + id: "mma_m16n8k4_f32_tf32", + abi_id: "i0520", + operation_key: "matrix.mma.m16n8k4.row.col.f32.tf32.tf32.f32", + source_record: "int_nvvm_mma_m16n8k4_row_col_tf32", + llvm_symbol: "llvm.nvvm.mma.m16n8k4.row.col.tf32", + rust_arguments: &["[f32; 4]", "[u32; 2]", "u32"], + rust_result: "[f32; 4]", + dialect_op_type: "RegisterMmaOp", + dialect_op_name: "nvvm.register_mma", + dialect_operands: &["f32", "f32", "f32", "f32", "i32", "i32", "i32"], + dialect_results: &["f32", "f32", "f32", "f32"], + llvm_arguments: &["i32", "i32", "i32", "f32", "f32", "f32", "f32"], + llvm_results: &["f32", "f32", "f32", "f32"], + adapter: C4F32A2U32B1U32ToD4F32, + compatibility_source: RegisterMmaCompatibilitySource::GeneratedStub, + minimum_ptx: "7.0", + minimum_sm: "sm_80", + ptx_modifiers: vec![ + "sync", "aligned", "m16n8k4", "row", "col", "f32", "tf32", "tf32", "f32", + ], + ptx_register_counts: [4, 2, 1, 4], + }), + ( + M16n8k8, + F16Accumulator, + F16Element, + F16Element, + RegisterMmaOverflow::NotApplicable, + C2U32A2U32B1U32ToD2U32, + true, + ) => Some(RegisterMmaRecipe { + id: "mma_m16n8k8_f16_f16", + abi_id: "i0521", + operation_key: "matrix.mma.m16n8k8.row.col.f16.f16.f16.f16", + source_record: "int_nvvm_mma_m16n8k8_row_col_f16_f16", + llvm_symbol: "llvm.nvvm.mma.m16n8k8.row.col.f16.f16", + rust_arguments: &["[u32; 2]", "[u32; 2]", "u32"], + rust_result: "[u32; 2]", + dialect_op_type: "RegisterMmaOp", + dialect_op_name: "nvvm.register_mma", + dialect_operands: &["i32", "i32", "i32", "i32", "i32"], + dialect_results: &["i32", "i32"], + llvm_arguments: &["v2f16", "v2f16", "v2f16", "v2f16", "v2f16"], + llvm_results: &["v2f16", "v2f16"], + adapter: C2U32A2U32B1U32ToD2U32, + compatibility_source: RegisterMmaCompatibilitySource::GeneratedStub, + minimum_ptx: "7.0", + minimum_sm: "sm_80", + ptx_modifiers: vec![ + "sync", "aligned", "m16n8k8", "row", "col", "f16", "f16", "f16", "f16", + ], + ptx_register_counts: [2, 2, 1, 2], + }), + ( + M16n8k8, + F32, + Bf16, + Bf16, + RegisterMmaOverflow::NotApplicable, + C4F32A2U32B1U32ToD4F32, + true, + ) => Some(RegisterMmaRecipe { + id: "mma_m16n8k8_f32_bf16", + abi_id: "i0522", + operation_key: "matrix.mma.m16n8k8.row.col.f32.bf16.bf16.f32", + source_record: "int_nvvm_mma_m16n8k8_row_col_bf16", + llvm_symbol: "llvm.nvvm.mma.m16n8k8.row.col.bf16", + rust_arguments: &["[f32; 4]", "[u32; 2]", "u32"], + rust_result: "[f32; 4]", + dialect_op_type: "RegisterMmaOp", + dialect_op_name: "nvvm.register_mma", + dialect_operands: &["f32", "f32", "f32", "f32", "i32", "i32", "i32"], + dialect_results: &["f32", "f32", "f32", "f32"], + llvm_arguments: &["i32", "i32", "i32", "f32", "f32", "f32", "f32"], + llvm_results: &["f32", "f32", "f32", "f32"], + adapter: C4F32A2U32B1U32ToD4F32, + compatibility_source: RegisterMmaCompatibilitySource::GeneratedStub, + minimum_ptx: "7.0", + minimum_sm: "sm_80", + ptx_modifiers: vec![ + "sync", "aligned", "m16n8k8", "row", "col", "f32", "bf16", "bf16", "f32", + ], + ptx_register_counts: [4, 2, 1, 4], + }), + ( + M16n8k8, + F32, + F16Element, + F16Element, + RegisterMmaOverflow::NotApplicable, + C4F32A2U32B1U32ToD4F32, + true, + ) => Some(RegisterMmaRecipe { + id: "mma_m16n8k8_f32_f16", + abi_id: "i0523", + operation_key: "matrix.mma.m16n8k8.row.col.f32.f16.f16.f32", + source_record: "int_nvvm_mma_m16n8k8_row_col_f32_f32", + llvm_symbol: "llvm.nvvm.mma.m16n8k8.row.col.f32.f32", + rust_arguments: &["[f32; 4]", "[u32; 2]", "u32"], + rust_result: "[f32; 4]", + dialect_op_type: "RegisterMmaOp", + dialect_op_name: "nvvm.register_mma", + dialect_operands: &["f32", "f32", "f32", "f32", "i32", "i32", "i32"], + dialect_results: &["f32", "f32", "f32", "f32"], + llvm_arguments: &["v2f16", "v2f16", "v2f16", "f32", "f32", "f32", "f32"], + llvm_results: &["f32", "f32", "f32", "f32"], + adapter: C4F32A2U32B1U32ToD4F32, + compatibility_source: RegisterMmaCompatibilitySource::GeneratedStub, + minimum_ptx: "7.0", + minimum_sm: "sm_80", + ptx_modifiers: vec![ + "sync", "aligned", "m16n8k8", "row", "col", "f32", "f16", "f16", "f32", + ], + ptx_register_counts: [4, 2, 1, 4], + }), + ( + M16n8k16, + F16Accumulator, + F16Element, + F16Element, + RegisterMmaOverflow::NotApplicable, + C2U32A4U32B2U32ToD2U32, + true, + ) => Some(RegisterMmaRecipe { + id: "mma_m16n8k16_f16_f16", + abi_id: "i0524", + operation_key: "matrix.mma.m16n8k16.row.col.f16.f16.f16.f16", + source_record: "int_nvvm_mma_m16n8k16_row_col_f16_f16", + llvm_symbol: "llvm.nvvm.mma.m16n8k16.row.col.f16.f16", + rust_arguments: &["[u32; 2]", "[u32; 4]", "[u32; 2]"], + rust_result: "[u32; 2]", + dialect_op_type: "RegisterMmaOp", + dialect_op_name: "nvvm.register_mma", + dialect_operands: &["i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32"], + dialect_results: &["i32", "i32"], + llvm_arguments: &[ + "v2f16", "v2f16", "v2f16", "v2f16", "v2f16", "v2f16", "v2f16", "v2f16", + ], + llvm_results: &["v2f16", "v2f16"], + adapter: C2U32A4U32B2U32ToD2U32, + compatibility_source: RegisterMmaCompatibilitySource::GeneratedStub, + minimum_ptx: "7.0", + minimum_sm: "sm_80", + ptx_modifiers: vec![ + "sync", "aligned", "m16n8k16", "row", "col", "f16", "f16", "f16", "f16", + ], + ptx_register_counts: [2, 4, 2, 2], + }), + ( + M16n8k16, + F32, + Bf16, + Bf16, + RegisterMmaOverflow::NotApplicable, + C4F32A4U32B2U32ToD4F32, + true, + ) => Some(RegisterMmaRecipe { + id: "mma_m16n8k16_f32_bf16", + abi_id: "i0105", + operation_key: "matrix.mma.m16n8k16.row.col.f32.bf16.bf16.f32", + source_record: "int_nvvm_mma_m16n8k16_row_col_bf16", + llvm_symbol: "llvm.nvvm.mma.m16n8k16.row.col.bf16", + rust_arguments: &["[f32; 4]", "[u32; 4]", "[u32; 2]"], + rust_result: "[f32; 4]", + dialect_op_type: "RegisterMmaOp", + dialect_op_name: "nvvm.register_mma", + dialect_operands: &[ + "f32", "f32", "f32", "f32", "i32", "i32", "i32", "i32", "i32", "i32", + ], + dialect_results: &["f32", "f32", "f32", "f32"], + llvm_arguments: &[ + "i32", "i32", "i32", "i32", "i32", "i32", "f32", "f32", "f32", "f32", + ], + llvm_results: &["f32", "f32", "f32", "f32"], + adapter: C4F32A4U32B2U32ToD4F32, + compatibility_source: RegisterMmaCompatibilitySource::ExistingStub, + minimum_ptx: "7.0", + minimum_sm: "sm_80", + ptx_modifiers: vec![ + "sync", "aligned", "m16n8k16", "row", "col", "f32", "bf16", "bf16", "f32", + ], + ptx_register_counts: [4, 4, 2, 4], + }), + ( + M16n8k16, + F32, + F16Element, + F16Element, + RegisterMmaOverflow::NotApplicable, + C4F32A4U32B2U32ToD4F32, + true, + ) => Some(RegisterMmaRecipe { + id: "mma_m16n8k16_f32_f16", + abi_id: "i0106", + operation_key: "matrix.mma.m16n8k16.row.col.f32.f16.f16.f32", + source_record: "int_nvvm_mma_m16n8k16_row_col_f32_f32", + llvm_symbol: "llvm.nvvm.mma.m16n8k16.row.col.f32.f32", + rust_arguments: &["[f32; 4]", "[u32; 4]", "[u32; 2]"], + rust_result: "[f32; 4]", + dialect_op_type: "RegisterMmaOp", + dialect_op_name: "nvvm.register_mma", + dialect_operands: &[ + "f32", "f32", "f32", "f32", "i32", "i32", "i32", "i32", "i32", "i32", + ], + dialect_results: &["f32", "f32", "f32", "f32"], + llvm_arguments: &[ + "v2f16", "v2f16", "v2f16", "v2f16", "v2f16", "v2f16", "f32", "f32", "f32", "f32", + ], + llvm_results: &["f32", "f32", "f32", "f32"], + adapter: C4F32A4U32B2U32ToD4F32, + compatibility_source: RegisterMmaCompatibilitySource::ExistingStub, + minimum_ptx: "7.0", + minimum_sm: "sm_80", + ptx_modifiers: vec![ + "sync", "aligned", "m16n8k16", "row", "col", "f32", "f16", "f16", "f32", + ], + ptx_register_counts: [4, 4, 2, 4], + }), + ( + M16n8k8, + F32, + Tf32, + Tf32, + RegisterMmaOverflow::NotApplicable, + C4F32A4U32B2U32ToD4F32, + true, + ) => Some(RegisterMmaRecipe { + id: "mma_m16n8k8_f32_tf32", + abi_id: "i0107", + operation_key: "matrix.mma.m16n8k8.row.col.f32.tf32.tf32.f32", + source_record: "int_nvvm_mma_m16n8k8_row_col_tf32", + llvm_symbol: "llvm.nvvm.mma.m16n8k8.row.col.tf32", + rust_arguments: &["[f32; 4]", "[u32; 4]", "[u32; 2]"], + rust_result: "[f32; 4]", + dialect_op_type: "RegisterMmaOp", + dialect_op_name: "nvvm.register_mma", + dialect_operands: &[ + "f32", "f32", "f32", "f32", "i32", "i32", "i32", "i32", "i32", "i32", + ], + dialect_results: &["f32", "f32", "f32", "f32"], + llvm_arguments: &[ + "i32", "i32", "i32", "i32", "i32", "i32", "f32", "f32", "f32", "f32", + ], + llvm_results: &["f32", "f32", "f32", "f32"], + adapter: C4F32A4U32B2U32ToD4F32, + compatibility_source: RegisterMmaCompatibilitySource::ExistingStub, + minimum_ptx: "7.0", + minimum_sm: "sm_80", + ptx_modifiers: vec![ + "sync", "aligned", "m16n8k8", "row", "col", "f32", "tf32", "tf32", "f32", + ], + ptx_register_counts: [4, 4, 2, 4], + }), + ( + M8n8k4, + F64, + F64Element, + F64Element, + RegisterMmaOverflow::NotApplicable, + C2F64A1F64B1F64ToD2F64, + true, + ) => Some(RegisterMmaRecipe { + id: "mma_m8n8k4_f64", + abi_id: "i0109", + operation_key: "matrix.mma.m8n8k4.row.col.f64.f64.f64.f64", + source_record: "int_nvvm_mma_m8n8k4_row_col_f64", + llvm_symbol: "llvm.nvvm.mma.m8n8k4.row.col.f64", + rust_arguments: &["[f64; 2]", "f64", "f64"], + rust_result: "[f64; 2]", + dialect_op_type: "RegisterMmaOp", + dialect_op_name: "nvvm.register_mma", + dialect_operands: &["f64", "f64", "f64", "f64"], + dialect_results: &["f64", "f64"], + llvm_arguments: &["f64", "f64", "f64", "f64"], + llvm_results: &["f64", "f64"], + adapter: C2F64A1F64B1F64ToD2F64, + compatibility_source: RegisterMmaCompatibilitySource::ExistingStub, + minimum_ptx: "7.0", + minimum_sm: "sm_80", + ptx_modifiers: vec![ + "sync", "aligned", "m8n8k4", "row", "col", "f64", "f64", "f64", "f64", + ], + ptx_register_counts: [2, 1, 1, 2], + }), + _ => None, + } +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/mma/register_admissions.rs b/crates/cuda-intrinsics-gen/src/resolve/families/mma/register_admissions.rs new file mode 100644 index 0000000000..c0a2d4b3d8 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/mma/register_admissions.rs @@ -0,0 +1,1454 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, IntrinsicBackend, OverlayBackendLowering, OverlayIntrinsic, + RegisterMma, RegisterMmaAccumulator, RegisterMmaAdapter, RegisterMmaAmpereFloatAdmission, + RegisterMmaAmpereFloatVariant, RegisterMmaBinaryAdmission, RegisterMmaCompatibilitySource, + RegisterMmaElement, RegisterMmaF8F6F4Admission, RegisterMmaFp8Admission, + RegisterMmaIntegerAdmission, RegisterMmaKind, RegisterMmaLayout, RegisterMmaOperation, + RegisterMmaOverflow, RegisterMmaParticipation, RegisterMmaShape, RuntimeValidation, + SparseMmaAccumulator, SparseMmaLayout, SparseMmaMetadata, SparseMmaOverflow, SparseMmaSelector, + SparseMmaShape, +}; +use crate::ptx::{InstructionPattern, OperandPattern}; +use anyhow::{Context, Result, bail, ensure}; +use std::collections::BTreeSet; + +use super::*; + +pub(in crate::resolve) const REGISTER_MMA_F8F6F4_TARGETS: &str = "sm_120a|sm_120f|sm_121a|sm_121f"; +#[derive(Debug, Clone, Copy)] +pub(in crate::resolve) enum RegisterMmaIntegerKind { + Int4, + Int8, +} + +impl RegisterMmaIntegerKind { + fn label(self) -> &'static str { + match self { + Self::Int4 => "INT4", + Self::Int8 => "INT8", + } + } + + fn supports(self, shape: RegisterMmaShape, element: RegisterMmaElement) -> bool { + match self { + Self::Int4 => { + matches!( + shape, + RegisterMmaShape::M8n8k32 + | RegisterMmaShape::M16n8k32 + | RegisterMmaShape::M16n8k64 + ) && matches!(element, RegisterMmaElement::S4 | RegisterMmaElement::U4) + } + Self::Int8 => { + matches!( + shape, + RegisterMmaShape::M8n8k16 + | RegisterMmaShape::M16n8k16 + | RegisterMmaShape::M16n8k32 + ) && matches!(element, RegisterMmaElement::S8 | RegisterMmaElement::U8) + } + } + } +} + +pub(in crate::resolve) fn expand_register_mma_integer_admission( + kind: RegisterMmaIntegerKind, + admission: &RegisterMmaIntegerAdmission, +) -> Result> { + use RegisterMmaAdapter::{ + C2I32A1U32B1U32ToD2I32, C4I32A2U32B1U32ToD4I32, C4I32A4U32B2U32ToD4I32, + }; + use RegisterMmaCompatibilitySource::GeneratedStub; + use RegisterMmaLayout::{Col, Row}; + use RegisterMmaParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes; + use RegisterMmaShape::{M8n8k16, M8n8k32, M16n8k16, M16n8k32, M16n8k64}; + + ensure!( + !admission.variants.is_empty(), + "compact {} MMA admission has no variants", + kind.label() + ); + ensure!( + admission.runtime_validation == RuntimeValidation::Unexecuted, + "{} MMA runtime validation may be marked executed only with GPU evidence", + kind.label() + ); + + let mut seen = BTreeSet::new(); + let mut records = Vec::with_capacity(admission.variants.len()); + for variant in &admission.variants { + ensure!( + seen.insert(( + variant.shape, + variant.a_element, + variant.b_element, + variant.overflow, + )), + "compact {} MMA admission contains a duplicate variant", + kind.label() + ); + ensure!( + kind.supports(variant.shape, variant.a_element) + && kind.supports(variant.shape, variant.b_element), + "compact {} MMA admission contains an unsupported shape or element", + kind.label() + ); + let adapter = + match (kind, variant.shape) { + (RegisterMmaIntegerKind::Int8, M8n8k16) + | (RegisterMmaIntegerKind::Int4, M8n8k32) => C2I32A1U32B1U32ToD2I32, + (RegisterMmaIntegerKind::Int8, M16n8k16) + | (RegisterMmaIntegerKind::Int4, M16n8k32) => C4I32A2U32B1U32ToD4I32, + (RegisterMmaIntegerKind::Int8, M16n8k32) + | (RegisterMmaIntegerKind::Int4, M16n8k64) => C4I32A4U32B2U32ToD4I32, + _ => bail!( + "compact {} MMA admission contains an unsupported shape", + kind.label() + ), + }; + let mma = RegisterMma { + shape: variant.shape, + operation: RegisterMmaOperation::Multiply, + kind: None, + accumulator: RegisterMmaAccumulator::S32, + a_element: variant.a_element, + b_element: variant.b_element, + a_layout: Row, + b_layout: Col, + overflow: variant.overflow, + participation: AllWarpLanesSameInstructionAndQualifiersNoExitedLanes, + adapter, + compatibility_source: GeneratedStub, + runtime_validation: admission.runtime_validation, + }; + let recipe = register_mma_recipe(&mma).with_context(|| { + format!( + "compact {} MMA admission requests a variant outside the closed recipe set", + kind.label() + ) + })?; + ensure!( + recipe.compatibility_source == GeneratedStub, + "compact {} MMA admission may only add generated compatibility stubs", + kind.label() + ); + + let element = |element| match (kind, element) { + (RegisterMmaIntegerKind::Int4, RegisterMmaElement::S4) + | (RegisterMmaIntegerKind::Int8, RegisterMmaElement::S8) => Ok("signed"), + (RegisterMmaIntegerKind::Int4, RegisterMmaElement::U4) + | (RegisterMmaIntegerKind::Int8, RegisterMmaElement::U8) => Ok("unsigned"), + _ => bail!( + "compact {} MMA admission contains an unsupported element", + kind.label() + ), + }; + let overflow = match variant.overflow { + RegisterMmaOverflow::Wrapping => "wrapping", + RegisterMmaOverflow::Satfinite => "saturating", + RegisterMmaOverflow::NotApplicable => { + bail!( + "compact {} MMA admission requires an integer overflow mode", + kind.label() + ) + } + }; + let summary = format!( + "Multiplies warp-distributed {} A and {} B {} fragments and adds a {overflow} s32 accumulator.", + element(variant.a_element)?, + element(variant.b_element)?, + kind.label(), + ); + + records.push(OverlayIntrinsic { + id: recipe.id.into(), + abi_id: recipe.abi_id.into(), + operation_key: recipe.operation_key.into(), + family: "register_mma".into(), + source: None, + source_record: Some(recipe.source_record.into()), + rust_module: "matrix".into(), + rust_name: recipe.id.into(), + rust_arguments: recipe + .rust_arguments + .iter() + .map(|value| (*value).into()) + .collect(), + rust_result: recipe.rust_result.into(), + safe: false, + must_use: true, + safe_allowlist_reason: None, + public_rust_path: format!("cuda_intrinsics::matrix::{}", recipe.id), + compatibility_rust_paths: vec![format!("cuda_device::wmma::{}", recipe.id)], + dialect_op_type: recipe.dialect_op_type.into(), + dialect_op_name: recipe.dialect_op_name.into(), + dialect_operands: recipe + .dialect_operands + .iter() + .map(|value| (*value).into()) + .collect(), + dialect_results: recipe + .dialect_results + .iter() + .map(|value| (*value).into()) + .collect(), + llvm_symbol: Some(recipe.llvm_symbol.into()), + resolved_llvm_symbol: None, + llvm_arguments: recipe + .llvm_arguments + .iter() + .map(|value| (*value).into()) + .collect(), + llvm_results: recipe + .llvm_results + .iter() + .map(|value| (*value).into()) + .collect(), + pure: false, + memory: "none".into(), + convergent: true, + execution_scope: "warp".into(), + minimum_ptx: recipe.minimum_ptx.into(), + minimum_sm: Some(recipe.minimum_sm.into()), + ptx_result: recipe.rust_result.into(), + targets: "all".into(), + ptx_isa_version: "9.3".into(), + ptx_isa_section: "9.7.15.5.14 Multiply-and-Accumulate Instruction: mma".into(), + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma".into(), + lowering: "generated_register_mma".into(), + backend_lowerings: vec![ + OverlayBackendLowering { + backend: IntrinsicBackend::LlvmNvptx, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: admission.llvm_evidence_profile.clone(), + targets: None, + minimum_ptx: Some(recipe.minimum_ptx.into()), + minimum_sm: Some(recipe.minimum_sm.into()), + }, + OverlayBackendLowering { + backend: IntrinsicBackend::LibNvvm, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: admission.libnvvm_evidence_profile.clone(), + targets: None, + minimum_ptx: Some(recipe.minimum_ptx.into()), + minimum_sm: Some(recipe.minimum_sm.into()), + }, + ], + packed_atomic: None, + redux: None, + vote: None, + active_mask: None, + warp_match: None, + warp_barrier: None, + warp_shuffle: None, + dot_product: None, + packed_alu: None, + integer_minmax: None, + packed_conversion: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cp_async_copy: None, + cp_async_control: None, + cp_async_mbarrier: None, + mbarrier_basic: None, + movmatrix: None, + mbarrier_extended: None, + register_mma: Some(mma), + sparse_mma: None, + prmt: None, + cluster_barrier: None, + wgmma_control: None, + special_register: None, + debug_control: None, + cluster_memory: None, + clc: None, + tma: None, + tcgen05: None, + ldmatrix_variant: None, + ldmatrix_safety: None, + ldmatrix_adapter: None, + selected_address_space: None, + expected_ptx: InstructionPattern { + mnemonic: "mma".into(), + modifiers: recipe + .ptx_modifiers + .iter() + .map(|value| (*value).into()) + .collect(), + operands: recipe + .ptx_register_counts + .map(|length| OperandPattern::RegisterList { length }) + .into(), + }, + summary, + }); + } + Ok(records) +} + +pub(in crate::resolve) const REGISTER_MMA_F8F6F4_ELEMENTS: [RegisterMmaElement; 5] = [ + RegisterMmaElement::E2m1, + RegisterMmaElement::E2m3, + RegisterMmaElement::E3m2, + RegisterMmaElement::E4m3, + RegisterMmaElement::E5m2, +]; + +pub(in crate::resolve) fn register_mma_f8f6f4_element_name( + element: RegisterMmaElement, +) -> Option<&'static str> { + match element { + RegisterMmaElement::E2m1 => Some("e2m1"), + RegisterMmaElement::E2m3 => Some("e2m3"), + RegisterMmaElement::E3m2 => Some("e3m2"), + RegisterMmaElement::E4m3 => Some("e4m3"), + RegisterMmaElement::E5m2 => Some("e5m2"), + _ => None, + } +} + +pub(in crate::resolve) fn is_dense_f8f6f4_register_mma_policy(policy: &OverlayIntrinsic) -> bool { + policy.family == "register_mma" + && policy.targets == REGISTER_MMA_F8F6F4_TARGETS + && policy.register_mma.as_ref().is_some_and(|mma| { + mma.kind != Some(RegisterMmaKind::Standard) + && mma.shape == RegisterMmaShape::M16n8k32 + && mma.operation == RegisterMmaOperation::Multiply + && matches!( + mma.accumulator, + RegisterMmaAccumulator::F16 | RegisterMmaAccumulator::F32 + ) + && register_mma_f8f6f4_element_name(mma.a_element).is_some() + && register_mma_f8f6f4_element_name(mma.b_element).is_some() + && mma.a_layout == RegisterMmaLayout::Row + && mma.b_layout == RegisterMmaLayout::Col + && mma.overflow == RegisterMmaOverflow::NotApplicable + }) +} + +pub(in crate::resolve) fn is_sparse_f8f6f4_f16_policy(policy: &OverlayIntrinsic) -> bool { + policy.family == "sparse_mma" + && policy.targets == SPARSE_MMA_F8F6F4_TARGETS + && policy.sparse_mma.as_ref().is_some_and(|mma| { + mma.shape == SparseMmaShape::M16n8k64 + && mma.accumulator == SparseMmaAccumulator::F16 + && SPARSE_MMA_F8F6F4_ELEMENTS.contains(&mma.a_element) + && SPARSE_MMA_F8F6F4_ELEMENTS.contains(&mma.b_element) + && mma.a_layout == SparseMmaLayout::Row + && mma.b_layout == SparseMmaLayout::Col + && mma.overflow == SparseMmaOverflow::NotApplicable + && mma.metadata == SparseMmaMetadata::Ordered + && mma.selector == SparseMmaSelector::ImmediateZero + }) +} + +pub(in crate::resolve) fn is_f8f6f4_mma_target_matrix_policy(policy: &OverlayIntrinsic) -> bool { + is_dense_f8f6f4_register_mma_policy(policy) || is_sparse_f8f6f4_f16_policy(policy) +} + +#[derive(Clone, Copy)] +pub(in crate::resolve) struct RegisterMmaF8F6F4Contract { + pub(in crate::resolve) accumulator: RegisterMmaAccumulator, + pub(in crate::resolve) scalar_name: &'static str, + pub(in crate::resolve) rust_arguments: [&'static str; 3], + pub(in crate::resolve) rust_result: &'static str, + pub(in crate::resolve) dialect_operands: &'static [&'static str], + pub(in crate::resolve) dialect_results: &'static [&'static str], + pub(in crate::resolve) llvm_arguments: &'static [&'static str], + pub(in crate::resolve) llvm_results: &'static [&'static str], + pub(in crate::resolve) adapter: RegisterMmaAdapter, + pub(in crate::resolve) ptx_register_counts: [usize; 4], + pub(in crate::resolve) summary_accumulator: &'static str, +} + +pub(in crate::resolve) fn register_mma_f8f6f4_contract( + accumulator: RegisterMmaAccumulator, +) -> Result { + match accumulator { + RegisterMmaAccumulator::F16 => Ok(RegisterMmaF8F6F4Contract { + accumulator, + scalar_name: "f16", + rust_arguments: ["[u32; 2]", "[u32; 4]", "[u32; 2]"], + rust_result: "[u32; 2]", + dialect_operands: &["i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32"], + dialect_results: &["i32", "i32"], + llvm_arguments: &["i32", "i32", "i32", "i32", "i32", "i32", "v2f16", "v2f16"], + llvm_results: &["v2f16", "v2f16"], + adapter: RegisterMmaAdapter::C2U32A4U32B2U32ToD2U32, + ptx_register_counts: [2, 4, 2, 2], + summary_accumulator: "F16", + }), + RegisterMmaAccumulator::F32 => Ok(RegisterMmaF8F6F4Contract { + accumulator, + scalar_name: "f32", + rust_arguments: ["[f32; 4]", "[u32; 4]", "[u32; 2]"], + rust_result: "[f32; 4]", + dialect_operands: &[ + "f32", "f32", "f32", "f32", "i32", "i32", "i32", "i32", "i32", "i32", + ], + dialect_results: &["f32", "f32", "f32", "f32"], + llvm_arguments: &[ + "i32", "i32", "i32", "i32", "i32", "i32", "f32", "f32", "f32", "f32", + ], + llvm_results: &["f32", "f32", "f32", "f32"], + adapter: RegisterMmaAdapter::C4F32A4U32B2U32ToD4F32, + ptx_register_counts: [4, 4, 2, 4], + summary_accumulator: "F32", + }), + _ => bail!("unsupported dense f8f6f4 MMA accumulator {accumulator:?}"), + } +} + +pub(in crate::resolve) fn expand_register_mma_f8f6f4_admission( + admission: &RegisterMmaF8F6F4Admission, + accumulator: RegisterMmaAccumulator, +) -> Result> { + let contract = register_mma_f8f6f4_contract(accumulator)?; + ensure!( + admission.runtime_validation == RuntimeValidation::Unexecuted, + "dense f8f6f4 MMA runtime validation may be marked executed only with GPU evidence" + ); + ensure!( + !admission.llvm_evidence_profile.trim().is_empty() + && !admission.libnvvm_evidence_profile.trim().is_empty(), + "dense f8f6f4 MMA admission requires both backend evidence profiles" + ); + ensure!( + admission.a_elements == REGISTER_MMA_F8F6F4_ELEMENTS + && admission.b_elements == REGISTER_MMA_F8F6F4_ELEMENTS + && admission.product_count == 25, + "dense f8f6f4 MMA admission must contain the canonical 5 by 5 element matrix" + ); + let expected_targets = ["sm_120a", "sm_120f", "sm_121a", "sm_121f"]; + ensure!( + admission + .targets + .iter() + .map(String::as_str) + .eq(expected_targets), + "dense f8f6f4 MMA admission must retain the reviewed Blackwell target set" + ); + + let mut records = Vec::with_capacity(admission.product_count); + for &a_element in &admission.a_elements { + for &b_element in &admission.b_elements { + let a = register_mma_f8f6f4_element_name(a_element) + .expect("validated dense f8f6f4 A element"); + let b = register_mma_f8f6f4_element_name(b_element) + .expect("validated dense f8f6f4 B element"); + let scalar = contract.scalar_name; + let id = format!("mma_m16n8k32_{scalar}_{a}_{b}"); + let source_record = + format!("int_nvvm_mma_m16n8k32_row_col_kind_f8f6f4_{scalar}_{a}_{b}_{scalar}"); + let llvm_symbol = + format!("llvm.nvvm.mma.m16n8k32.row.col.kind.f8f6f4.{scalar}.{a}.{b}.{scalar}"); + let mma = RegisterMma { + shape: RegisterMmaShape::M16n8k32, + operation: RegisterMmaOperation::Multiply, + kind: None, + accumulator: contract.accumulator, + a_element, + b_element, + a_layout: RegisterMmaLayout::Row, + b_layout: RegisterMmaLayout::Col, + overflow: RegisterMmaOverflow::NotApplicable, + participation: + RegisterMmaParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes, + adapter: contract.adapter, + compatibility_source: RegisterMmaCompatibilitySource::GeneratedStub, + runtime_validation: admission.runtime_validation, + }; + + records.push(OverlayIntrinsic { + id: id.clone(), + abi_id: String::new(), + operation_key: format!( + "matrix.mma.m16n8k32.row.col.kind_f8f6f4.{scalar}.{a}.{b}.{scalar}" + ), + family: "register_mma".into(), + source: None, + source_record: Some(source_record), + rust_module: "matrix".into(), + rust_name: id.clone(), + rust_arguments: contract.rust_arguments.map(Into::into).into(), + rust_result: contract.rust_result.into(), + safe: false, + must_use: true, + safe_allowlist_reason: None, + public_rust_path: format!("cuda_intrinsics::matrix::{id}"), + compatibility_rust_paths: vec![format!("cuda_device::wmma::{id}")], + dialect_op_type: "RegisterMmaOp".into(), + dialect_op_name: "nvvm.register_mma".into(), + dialect_operands: contract + .dialect_operands + .iter() + .map(|value| (*value).into()) + .collect(), + dialect_results: contract + .dialect_results + .iter() + .map(|value| (*value).into()) + .collect(), + llvm_symbol: Some(llvm_symbol), + resolved_llvm_symbol: None, + llvm_arguments: contract + .llvm_arguments + .iter() + .map(|value| (*value).into()) + .collect(), + llvm_results: contract + .llvm_results + .iter() + .map(|value| (*value).into()) + .collect(), + pure: false, + memory: "none".into(), + convergent: true, + execution_scope: "warp".into(), + minimum_ptx: "8.7".into(), + minimum_sm: None, + ptx_result: contract.rust_result.into(), + targets: REGISTER_MMA_F8F6F4_TARGETS.into(), + ptx_isa_version: "9.3".into(), + ptx_isa_section: "9.7.15.5.14 Multiply-and-Accumulate Instruction: mma".into(), + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma".into(), + lowering: "generated_register_mma".into(), + backend_lowerings: vec![ + OverlayBackendLowering { + backend: IntrinsicBackend::LlvmNvptx, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: admission.llvm_evidence_profile.clone(), + targets: None, + minimum_ptx: None, + minimum_sm: None, + }, + OverlayBackendLowering { + backend: IntrinsicBackend::LibNvvm, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: admission.libnvvm_evidence_profile.clone(), + targets: None, + minimum_ptx: None, + minimum_sm: None, + }, + ], + packed_atomic: None, + redux: None, + vote: None, + active_mask: None, + warp_match: None, + warp_barrier: None, + warp_shuffle: None, + dot_product: None, + packed_alu: None, + integer_minmax: None, + packed_conversion: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cp_async_copy: None, + cp_async_control: None, + cp_async_mbarrier: None, + mbarrier_basic: None, + movmatrix: None, + mbarrier_extended: None, + register_mma: Some(mma), + sparse_mma: None, + prmt: None, + cluster_barrier: None, + wgmma_control: None, + special_register: None, + debug_control: None, + cluster_memory: None, + clc: None, + tma: None, + tcgen05: None, + ldmatrix_variant: None, + ldmatrix_safety: None, + ldmatrix_adapter: None, + selected_address_space: None, + expected_ptx: InstructionPattern { + mnemonic: "mma".into(), + modifiers: [ + "sync", + "aligned", + "m16n8k32", + "row", + "col", + "kind::f8f6f4", + scalar, + a, + b, + scalar, + ] + .into_iter() + .map(Into::into) + .collect(), + operands: contract + .ptx_register_counts + .map(|length| OperandPattern::RegisterList { length }) + .into(), + }, + summary: format!( + "Multiplies warp-distributed {a} A and {b} B fragments and adds an {} accumulator.", + contract.summary_accumulator + ), + }); + } + } + Ok(records) +} + +pub(in crate::resolve) fn expand_register_mma_mxf8f6f4_admission( + admission: &RegisterMmaF8F6F4Admission, +) -> Result> { + ensure!( + admission.runtime_validation == RuntimeValidation::Unexecuted, + "dense mxf8f6f4 MMA runtime validation may be marked executed only with GPU evidence" + ); + ensure!( + !admission.llvm_evidence_profile.trim().is_empty() + && !admission.libnvvm_evidence_profile.trim().is_empty(), + "dense mxf8f6f4 MMA admission requires both backend evidence profiles" + ); + ensure!( + admission.a_elements == REGISTER_MMA_F8F6F4_ELEMENTS + && admission.b_elements == REGISTER_MMA_F8F6F4_ELEMENTS + && admission.product_count == 25, + "dense mxf8f6f4 MMA admission must contain the canonical 5 by 5 element matrix" + ); + let expected_targets = ["sm_120a", "sm_120f", "sm_121a", "sm_121f"]; + ensure!( + admission + .targets + .iter() + .map(String::as_str) + .eq(expected_targets), + "dense mxf8f6f4 MMA admission must retain the reviewed Blackwell target set" + ); + + let rust_arguments = [ + "[f32; 4]", "[u32; 4]", "[u32; 2]", "u32", "u16", "u16", "u32", "u16", "u16", + ]; + let dialect_operands = [ + "f32", "f32", "f32", "f32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i16", "i16", + "i32", "i16", "i16", + ]; + let llvm_arguments = [ + "i32", "i32", "i32", "i32", "i32", "i32", "f32", "f32", "f32", "f32", "i32", "i16", "i16", + "i32", "i16", "i16", + ]; + let dialect_results = ["f32", "f32", "f32", "f32"]; + + let mut records = Vec::with_capacity(admission.product_count); + for &a_element in &admission.a_elements { + for &b_element in &admission.b_elements { + let a = register_mma_f8f6f4_element_name(a_element) + .expect("validated dense mxf8f6f4 A element"); + let b = register_mma_f8f6f4_element_name(b_element) + .expect("validated dense mxf8f6f4 B element"); + let id = format!("mma_m16n8k32_mxf8f6f4_f32_{a}_{b}"); + let source_record = + format!("int_nvvm_mma_block_scale_m16n8k32_row_col_mxf8f6f4_f32_{a}_{b}_f32_ue8m0"); + let llvm_symbol = format!( + "llvm.nvvm.mma.block.scale.m16n8k32.row.col.mxf8f6f4.f32.{a}.{b}.f32.ue8m0" + ); + let mma = RegisterMma { + shape: RegisterMmaShape::M16n8k32, + operation: RegisterMmaOperation::Multiply, + kind: Some(RegisterMmaKind::Mxf8f6f4), + accumulator: RegisterMmaAccumulator::F32, + a_element, + b_element, + a_layout: RegisterMmaLayout::Row, + b_layout: RegisterMmaLayout::Col, + overflow: RegisterMmaOverflow::NotApplicable, + participation: + RegisterMmaParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes, + adapter: RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32, + compatibility_source: RegisterMmaCompatibilitySource::GeneratedStub, + runtime_validation: admission.runtime_validation, + }; + + records.push(OverlayIntrinsic { + id: id.clone(), + abi_id: String::new(), + operation_key: format!( + "matrix.mma.m16n8k32.row.col.kind_mxf8f6f4.scale_vec_1x.f32.{a}.{b}.f32.ue8m0" + ), + family: "register_mma".into(), + source: None, + source_record: Some(source_record), + rust_module: "matrix".into(), + rust_name: id.clone(), + rust_arguments: rust_arguments.map(Into::into).into(), + rust_result: "[f32; 4]".into(), + safe: false, + must_use: true, + safe_allowlist_reason: None, + public_rust_path: format!("cuda_intrinsics::matrix::{id}"), + compatibility_rust_paths: vec![format!("cuda_device::wmma::{id}")], + dialect_op_type: "RegisterMmaOp".into(), + dialect_op_name: "nvvm.register_mma".into(), + dialect_operands: dialect_operands.map(Into::into).into(), + dialect_results: dialect_results.map(Into::into).into(), + llvm_symbol: Some(llvm_symbol), + resolved_llvm_symbol: None, + llvm_arguments: llvm_arguments.map(Into::into).into(), + llvm_results: dialect_results.map(Into::into).into(), + pure: false, + memory: "none".into(), + convergent: true, + execution_scope: "warp".into(), + minimum_ptx: "8.7".into(), + minimum_sm: None, + ptx_result: "[f32; 4]".into(), + targets: REGISTER_MMA_F8F6F4_TARGETS.into(), + ptx_isa_version: "9.3".into(), + ptx_isa_section: + "9.7.15.5.14 Multiply-and-Accumulate Instruction: mma".into(), + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma".into(), + lowering: "generated_register_mma".into(), + backend_lowerings: vec![ + OverlayBackendLowering { + backend: IntrinsicBackend::LlvmNvptx, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: admission.llvm_evidence_profile.clone(), + targets: None, + minimum_ptx: None, + minimum_sm: None, + }, + OverlayBackendLowering { + backend: IntrinsicBackend::LibNvvm, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: admission.libnvvm_evidence_profile.clone(), + targets: None, + minimum_ptx: None, + minimum_sm: None, + }, + ], + packed_atomic: None, + redux: None, + vote: None, + active_mask: None, + warp_match: None, + warp_barrier: None, + warp_shuffle: None, + dot_product: None, + packed_alu: None, + integer_minmax: None, + packed_conversion: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cp_async_copy: None, + cp_async_control: None, + cp_async_mbarrier: None, + mbarrier_basic: None, + movmatrix: None, + mbarrier_extended: None, + register_mma: Some(mma), + sparse_mma: None, + prmt: None, + cluster_barrier: None, + wgmma_control: None, + special_register: None, + debug_control: None, + cluster_memory: None, + clc: None, + tma: None, + tcgen05: None, + ldmatrix_variant: None, + ldmatrix_safety: None, + ldmatrix_adapter: None, + selected_address_space: None, + expected_ptx: InstructionPattern { + mnemonic: "mma".into(), + modifiers: [ + "sync", + "aligned", + "m16n8k32", + "row", + "col", + "kind::mxf8f6f4", + "block_scale", + "f32", + a, + b, + "f32", + "ue8m0", + ] + .into_iter() + .map(Into::into) + .collect(), + operands: vec![ + OperandPattern::RegisterList { length: 4 }, + OperandPattern::RegisterList { length: 4 }, + OperandPattern::RegisterList { length: 2 }, + OperandPattern::RegisterList { length: 4 }, + OperandPattern::Register, + OperandPattern::RegisterList { length: 2 }, + OperandPattern::Register, + OperandPattern::RegisterList { length: 2 }, + ], + }, + summary: format!( + "Multiplies block-scaled warp-distributed {a} A and {b} B fragments and adds an F32 accumulator." + ), + }); + } + } + Ok(records) +} + +pub(in crate::resolve) const REGISTER_MMA_FP8_SHAPES: [RegisterMmaShape; 2] = + [RegisterMmaShape::M16n8k16, RegisterMmaShape::M16n8k32]; +pub(in crate::resolve) const REGISTER_MMA_FP8_ACCUMULATORS: [RegisterMmaAccumulator; 2] = + [RegisterMmaAccumulator::F16, RegisterMmaAccumulator::F32]; +pub(in crate::resolve) const REGISTER_MMA_FP8_ELEMENTS: [RegisterMmaElement; 2] = + [RegisterMmaElement::E4m3, RegisterMmaElement::E5m2]; + +pub(in crate::resolve) fn register_mma_fp8_element_name( + element: RegisterMmaElement, +) -> Option<&'static str> { + match element { + RegisterMmaElement::E4m3 => Some("e4m3"), + RegisterMmaElement::E5m2 => Some("e5m2"), + _ => None, + } +} + +pub(in crate::resolve) fn register_mma_fp8_shape_contract( + shape: RegisterMmaShape, +) -> Result<(&'static str, usize, usize)> { + match shape { + RegisterMmaShape::M16n8k16 => Ok(("m16n8k16", 2, 1)), + RegisterMmaShape::M16n8k32 => Ok(("m16n8k32", 4, 2)), + _ => bail!("unsupported standard FP8 register-MMA shape {shape:?}"), + } +} + +pub(in crate::resolve) fn register_mma_fp8_minimum_ptx( + shape: RegisterMmaShape, + accumulator: RegisterMmaAccumulator, +) -> &'static str { + // ptxas requires PTX 8.7 for the k32 F16 form despite LLVM's 8.4 predicate. + match (shape, accumulator) { + (RegisterMmaShape::M16n8k32, RegisterMmaAccumulator::F32) => "8.4", + _ => "8.7", + } +} + +pub(in crate::resolve) fn expand_register_mma_fp8_admission( + admission: &RegisterMmaFp8Admission, +) -> Result> { + ensure!( + admission.runtime_validation == RuntimeValidation::Unexecuted, + "standard FP8 MMA runtime validation may be marked executed only with GPU evidence" + ); + ensure!( + !admission.llvm_evidence_profile.trim().is_empty() + && !admission.libnvvm_evidence_profile.trim().is_empty(), + "standard FP8 MMA admission requires both backend evidence profiles" + ); + ensure!( + admission.shapes == REGISTER_MMA_FP8_SHAPES + && admission.accumulators == REGISTER_MMA_FP8_ACCUMULATORS + && admission.a_elements == REGISTER_MMA_FP8_ELEMENTS + && admission.b_elements == REGISTER_MMA_FP8_ELEMENTS + && admission.product_count == 16, + "standard FP8 MMA admission must retain the canonical 2 by 2 by 2 by 2 matrix" + ); + + let mut records = Vec::with_capacity(admission.product_count); + for &shape in &admission.shapes { + let (shape_name, a_count, b_count) = register_mma_fp8_shape_contract(shape)?; + for &accumulator in &admission.accumulators { + let minimum_ptx = register_mma_fp8_minimum_ptx(shape, accumulator); + let (scalar, rust_arguments, rust_result, dialect_operands, dialect_results, adapter) = + match accumulator { + RegisterMmaAccumulator::F16 => ( + "f16", + vec![ + "[u32; 2]".into(), + format!("[u32; {a_count}]"), + if b_count == 1 { + "u32".into() + } else { + "[u32; 2]".into() + }, + ], + "[u32; 2]".to_owned(), + [ + vec!["i32".into(); 2], + vec!["i32".into(); a_count], + vec!["i32".into(); b_count], + ] + .concat(), + vec!["i32".into(); 2], + if shape == RegisterMmaShape::M16n8k16 { + RegisterMmaAdapter::C2U32A2U32B1U32ToD2U32 + } else { + RegisterMmaAdapter::C2U32A4U32B2U32ToD2U32 + }, + ), + RegisterMmaAccumulator::F32 => ( + "f32", + vec![ + "[f32; 4]".into(), + format!("[u32; {a_count}]"), + if b_count == 1 { + "u32".into() + } else { + "[u32; 2]".into() + }, + ], + "[f32; 4]".to_owned(), + [ + vec!["f32".into(); 4], + vec!["i32".into(); a_count], + vec!["i32".into(); b_count], + ] + .concat(), + vec!["f32".into(); 4], + if shape == RegisterMmaShape::M16n8k16 { + RegisterMmaAdapter::C4F32A2U32B1U32ToD4F32 + } else { + RegisterMmaAdapter::C4F32A4U32B2U32ToD4F32 + }, + ), + _ => unreachable!("validated standard FP8 accumulator"), + }; + let llvm_arguments = [ + vec!["i32".into(); a_count + b_count], + if accumulator == RegisterMmaAccumulator::F16 { + vec!["v2f16".into(); 2] + } else { + vec!["f32".into(); 4] + }, + ] + .concat(); + let llvm_results = if accumulator == RegisterMmaAccumulator::F16 { + vec!["v2f16".into(); 2] + } else { + vec!["f32".into(); 4] + }; + let result_count = dialect_results.len(); + + for &a_element in &admission.a_elements { + for &b_element in &admission.b_elements { + let a = register_mma_fp8_element_name(a_element) + .expect("validated standard FP8 A element"); + let b = register_mma_fp8_element_name(b_element) + .expect("validated standard FP8 B element"); + let id = format!("mma_{shape_name}_fp8_{scalar}_{a}_{b}"); + let source_record = + format!("int_nvvm_mma_{shape_name}_row_col_{scalar}_{a}_{b}_{scalar}"); + let llvm_symbol = + format!("llvm.nvvm.mma.{shape_name}.row.col.{scalar}.{a}.{b}.{scalar}"); + let mma = RegisterMma { + shape, + operation: RegisterMmaOperation::Multiply, + kind: Some(RegisterMmaKind::Standard), + accumulator, + a_element, + b_element, + a_layout: RegisterMmaLayout::Row, + b_layout: RegisterMmaLayout::Col, + overflow: RegisterMmaOverflow::NotApplicable, + participation: RegisterMmaParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes, + adapter, + compatibility_source: RegisterMmaCompatibilitySource::GeneratedStub, + runtime_validation: admission.runtime_validation, + }; + records.push(OverlayIntrinsic { + id: id.clone(), + abi_id: String::new(), + operation_key: format!( + "matrix.mma.{shape_name}.row.col.standard_fp8.{scalar}.{a}.{b}.{scalar}" + ), + family: "register_mma".into(), + source: None, + source_record: Some(source_record), + rust_module: "matrix".into(), + rust_name: id.clone(), + rust_arguments: rust_arguments.clone(), + rust_result: rust_result.clone(), + safe: false, + must_use: true, + safe_allowlist_reason: None, + public_rust_path: format!("cuda_intrinsics::matrix::{id}"), + compatibility_rust_paths: vec![format!("cuda_device::wmma::{id}")], + dialect_op_type: "RegisterMmaOp".into(), + dialect_op_name: "nvvm.register_mma".into(), + dialect_operands: dialect_operands.clone(), + dialect_results: dialect_results.clone(), + llvm_symbol: Some(llvm_symbol), + resolved_llvm_symbol: None, + llvm_arguments: llvm_arguments.clone(), + llvm_results: llvm_results.clone(), + pure: false, + memory: "none".into(), + convergent: true, + execution_scope: "warp".into(), + minimum_ptx: minimum_ptx.into(), + minimum_sm: Some("sm_89".into()), + ptx_result: rust_result.clone(), + targets: "all".into(), + ptx_isa_version: "9.3".into(), + ptx_isa_section: "9.7.15.5.14 Multiply-and-Accumulate Instruction: mma".into(), + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma".into(), + lowering: "generated_register_mma".into(), + backend_lowerings: [ + (IntrinsicBackend::LlvmNvptx, &admission.llvm_evidence_profile), + (IntrinsicBackend::LibNvvm, &admission.libnvvm_evidence_profile), + ] + .into_iter() + .map(|(backend, evidence_profile)| OverlayBackendLowering { + backend, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: evidence_profile.clone(), + targets: None, + minimum_ptx: Some(minimum_ptx.into()), + minimum_sm: Some("sm_89".into()), + }) + .collect(), + packed_atomic: None, + redux: None, + vote: None, + active_mask: None, + warp_match: None, + warp_barrier: None, + warp_shuffle: None, + dot_product: None, + packed_alu: None, + integer_minmax: None, + packed_conversion: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cp_async_copy: None, + cp_async_control: None, + cp_async_mbarrier: None, + mbarrier_basic: None, + movmatrix: None, + mbarrier_extended: None, + register_mma: Some(mma), + sparse_mma: None, + prmt: None, + cluster_barrier: None, + wgmma_control: None, + special_register: None, + debug_control: None, + cluster_memory: None, + clc: None, + tma: None, + tcgen05: None, + ldmatrix_variant: None, + ldmatrix_safety: None, + ldmatrix_adapter: None, + selected_address_space: None, + expected_ptx: InstructionPattern { + mnemonic: "mma".into(), + modifiers: ["sync", "aligned", shape_name, "row", "col", scalar, a, b, scalar] + .into_iter() + .map(Into::into) + .collect(), + operands: [result_count, a_count, b_count, result_count] + .map(|length| OperandPattern::RegisterList { length }) + .into(), + }, + summary: format!( + "Multiplies warp-distributed {a} A and {b} B fragments and adds an {scalar} accumulator." + ), + }); + } + } + } + } + ensure!(records.len() == admission.product_count); + Ok(records) +} + +pub(in crate::resolve) const REGISTER_MMA_AMPERE_FLOAT_VARIANTS: [RegisterMmaAmpereFloatVariant; + 5] = [ + RegisterMmaAmpereFloatVariant { + shape: RegisterMmaShape::M16n8k4, + accumulator: RegisterMmaAccumulator::F32, + element: RegisterMmaElement::Tf32, + }, + RegisterMmaAmpereFloatVariant { + shape: RegisterMmaShape::M16n8k8, + accumulator: RegisterMmaAccumulator::F16, + element: RegisterMmaElement::F16, + }, + RegisterMmaAmpereFloatVariant { + shape: RegisterMmaShape::M16n8k8, + accumulator: RegisterMmaAccumulator::F32, + element: RegisterMmaElement::Bf16, + }, + RegisterMmaAmpereFloatVariant { + shape: RegisterMmaShape::M16n8k8, + accumulator: RegisterMmaAccumulator::F32, + element: RegisterMmaElement::F16, + }, + RegisterMmaAmpereFloatVariant { + shape: RegisterMmaShape::M16n8k16, + accumulator: RegisterMmaAccumulator::F16, + element: RegisterMmaElement::F16, + }, +]; + +pub(in crate::resolve) fn expand_register_mma_ampere_float_admission( + admission: &RegisterMmaAmpereFloatAdmission, +) -> Result> { + ensure!( + admission.runtime_validation == RuntimeValidation::Unexecuted, + "Ampere floating-point MMA runtime validation requires GPU evidence" + ); + ensure!( + !admission.llvm_evidence_profile.trim().is_empty() + && !admission.libnvvm_evidence_profile.trim().is_empty(), + "Ampere floating-point MMA admission requires both backend evidence profiles" + ); + ensure!( + admission.product_count == REGISTER_MMA_AMPERE_FLOAT_VARIANTS.len() + && admission.variants == REGISTER_MMA_AMPERE_FLOAT_VARIANTS, + "Ampere floating-point MMA admission must retain the five reviewed variants in canonical order" + ); + + admission + .variants + .iter() + .map(|variant| { + let adapter = match (variant.shape, variant.accumulator) { + (RegisterMmaShape::M16n8k4, RegisterMmaAccumulator::F32) + | (RegisterMmaShape::M16n8k8, RegisterMmaAccumulator::F32) => { + RegisterMmaAdapter::C4F32A2U32B1U32ToD4F32 + } + (RegisterMmaShape::M16n8k8, RegisterMmaAccumulator::F16) => { + RegisterMmaAdapter::C2U32A2U32B1U32ToD2U32 + } + (RegisterMmaShape::M16n8k16, RegisterMmaAccumulator::F16) => { + RegisterMmaAdapter::C2U32A4U32B2U32ToD2U32 + } + _ => bail!("Ampere floating-point MMA admission contains an unsupported carrier"), + }; + let mma = RegisterMma { + shape: variant.shape, + operation: RegisterMmaOperation::Multiply, + kind: None, + accumulator: variant.accumulator, + a_element: variant.element, + b_element: variant.element, + a_layout: RegisterMmaLayout::Row, + b_layout: RegisterMmaLayout::Col, + overflow: RegisterMmaOverflow::NotApplicable, + participation: RegisterMmaParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes, + adapter, + compatibility_source: RegisterMmaCompatibilitySource::GeneratedStub, + runtime_validation: admission.runtime_validation, + }; + let recipe = register_mma_recipe(&mma).with_context(|| { + "Ampere floating-point MMA admission requests a variant outside the closed recipe set" + })?; + Ok(OverlayIntrinsic { + id: recipe.id.into(), + abi_id: String::new(), + operation_key: recipe.operation_key.into(), + family: "register_mma".into(), + source: None, + source_record: Some(recipe.source_record.into()), + rust_module: "matrix".into(), + rust_name: recipe.id.into(), + rust_arguments: recipe.rust_arguments.iter().map(|value| (*value).into()).collect(), + rust_result: recipe.rust_result.into(), + safe: false, + must_use: true, + safe_allowlist_reason: None, + public_rust_path: format!("cuda_intrinsics::matrix::{}", recipe.id), + compatibility_rust_paths: vec![format!("cuda_device::wmma::{}", recipe.id)], + dialect_op_type: recipe.dialect_op_type.into(), + dialect_op_name: recipe.dialect_op_name.into(), + dialect_operands: recipe.dialect_operands.iter().map(|value| (*value).into()).collect(), + dialect_results: recipe.dialect_results.iter().map(|value| (*value).into()).collect(), + llvm_symbol: Some(recipe.llvm_symbol.into()), + resolved_llvm_symbol: None, + llvm_arguments: recipe.llvm_arguments.iter().map(|value| (*value).into()).collect(), + llvm_results: recipe.llvm_results.iter().map(|value| (*value).into()).collect(), + pure: false, + memory: "none".into(), + convergent: true, + execution_scope: "warp".into(), + minimum_ptx: recipe.minimum_ptx.into(), + minimum_sm: Some(recipe.minimum_sm.into()), + ptx_result: recipe.rust_result.into(), + targets: "all".into(), + ptx_isa_version: "9.3".into(), + ptx_isa_section: "9.7.15.5.14 Multiply-and-Accumulate Instruction: mma".into(), + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma".into(), + lowering: "generated_register_mma".into(), + backend_lowerings: [ + (IntrinsicBackend::LlvmNvptx, &admission.llvm_evidence_profile), + (IntrinsicBackend::LibNvvm, &admission.libnvvm_evidence_profile), + ] + .into_iter() + .map(|(backend, evidence_profile)| OverlayBackendLowering { + backend, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: evidence_profile.clone(), + targets: None, + minimum_ptx: Some(recipe.minimum_ptx.into()), + minimum_sm: Some(recipe.minimum_sm.into()), + }) + .collect(), + packed_atomic: None, + redux: None, + vote: None, + active_mask: None, + warp_match: None, + warp_barrier: None, + warp_shuffle: None, + dot_product: None, + packed_alu: None, + integer_minmax: None, + packed_conversion: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cp_async_copy: None, + cp_async_control: None, + cp_async_mbarrier: None, + mbarrier_basic: None, + movmatrix: None, + mbarrier_extended: None, + register_mma: Some(mma), + sparse_mma: None, + prmt: None, + cluster_barrier: None, + wgmma_control: None, + special_register: None, + debug_control: None, + cluster_memory: None, + clc: None, + tma: None, + tcgen05: None, + ldmatrix_variant: None, + ldmatrix_safety: None, + ldmatrix_adapter: None, + selected_address_space: None, + expected_ptx: InstructionPattern { + mnemonic: "mma".into(), + modifiers: recipe.ptx_modifiers.iter().map(|value| (*value).into()).collect(), + operands: recipe + .ptx_register_counts + .map(|length| OperandPattern::RegisterList { length }) + .into(), + }, + summary: format!("Executes the {} Ampere floating-point MMA form.", recipe.id), + }) + }) + .collect() +} + +pub(in crate::resolve) fn expand_register_mma_binary_admission( + admission: &RegisterMmaBinaryAdmission, +) -> Result> { + use RegisterMmaAdapter::{ + C2I32A1U32B1U32ToD2I32, C4I32A2U32B1U32ToD4I32, C4I32A4U32B2U32ToD4I32, + }; + use RegisterMmaLayout::{Col, Row}; + use RegisterMmaOperation::{AndPopc, XorPopc}; + use RegisterMmaParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes; + use RegisterMmaShape::{M8n8k128, M16n8k128, M16n8k256}; + + ensure!( + !admission.variants.is_empty(), + "compact binary MMA admission has no variants" + ); + ensure!( + admission.runtime_validation == RuntimeValidation::Unexecuted, + "binary MMA runtime validation may be marked executed only with GPU evidence" + ); + + let mut seen = BTreeSet::new(); + let mut records = Vec::with_capacity(admission.variants.len()); + for variant in &admission.variants { + ensure!( + seen.insert((variant.shape, variant.operation)), + "compact binary MMA admission contains a duplicate variant" + ); + ensure!( + matches!(variant.operation, AndPopc | XorPopc), + "compact binary MMA admission contains a non-binary operation" + ); + let adapter = match variant.shape { + M8n8k128 => C2I32A1U32B1U32ToD2I32, + M16n8k128 => C4I32A2U32B1U32ToD4I32, + M16n8k256 => C4I32A4U32B2U32ToD4I32, + _ => bail!("compact binary MMA admission contains an unsupported shape"), + }; + let mma = RegisterMma { + shape: variant.shape, + operation: variant.operation, + kind: None, + accumulator: RegisterMmaAccumulator::S32, + a_element: RegisterMmaElement::B1, + b_element: RegisterMmaElement::B1, + a_layout: Row, + b_layout: Col, + overflow: RegisterMmaOverflow::Wrapping, + participation: AllWarpLanesSameInstructionAndQualifiersNoExitedLanes, + adapter, + compatibility_source: RegisterMmaCompatibilitySource::GeneratedStub, + runtime_validation: admission.runtime_validation, + }; + let recipe = register_mma_recipe(&mma).with_context( + || "compact binary MMA admission requests a variant outside the closed recipe set", + )?; + let operation = match variant.operation { + AndPopc => "AND and population count", + XorPopc => "XOR and population count", + RegisterMmaOperation::Multiply => unreachable!(), + }; + let summary = format!( + "Computes warp-distributed binary matrix products with {operation}, then adds a wrapping s32 accumulator." + ); + + records.push(OverlayIntrinsic { + id: recipe.id.into(), + abi_id: recipe.abi_id.into(), + operation_key: recipe.operation_key.into(), + family: "register_mma".into(), + source: None, + source_record: Some(recipe.source_record.into()), + rust_module: "matrix".into(), + rust_name: recipe.id.into(), + rust_arguments: recipe + .rust_arguments + .iter() + .map(|value| (*value).into()) + .collect(), + rust_result: recipe.rust_result.into(), + safe: false, + must_use: true, + safe_allowlist_reason: None, + public_rust_path: format!("cuda_intrinsics::matrix::{}", recipe.id), + compatibility_rust_paths: vec![format!("cuda_device::wmma::{}", recipe.id)], + dialect_op_type: recipe.dialect_op_type.into(), + dialect_op_name: recipe.dialect_op_name.into(), + dialect_operands: recipe + .dialect_operands + .iter() + .map(|value| (*value).into()) + .collect(), + dialect_results: recipe + .dialect_results + .iter() + .map(|value| (*value).into()) + .collect(), + llvm_symbol: Some(recipe.llvm_symbol.into()), + resolved_llvm_symbol: None, + llvm_arguments: recipe + .llvm_arguments + .iter() + .map(|value| (*value).into()) + .collect(), + llvm_results: recipe + .llvm_results + .iter() + .map(|value| (*value).into()) + .collect(), + pure: false, + memory: "none".into(), + convergent: true, + execution_scope: "warp".into(), + minimum_ptx: recipe.minimum_ptx.into(), + minimum_sm: Some(recipe.minimum_sm.into()), + ptx_result: recipe.rust_result.into(), + targets: "all".into(), + ptx_isa_version: "9.3".into(), + ptx_isa_section: "9.7.15.5.14 Multiply-and-Accumulate Instruction: mma".into(), + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma".into(), + lowering: "generated_register_mma".into(), + backend_lowerings: vec![ + OverlayBackendLowering { + backend: IntrinsicBackend::LlvmNvptx, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: admission.llvm_evidence_profile.clone(), + targets: None, + minimum_ptx: Some(recipe.minimum_ptx.into()), + minimum_sm: Some(recipe.minimum_sm.into()), + }, + OverlayBackendLowering { + backend: IntrinsicBackend::LibNvvm, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: admission.libnvvm_evidence_profile.clone(), + targets: None, + minimum_ptx: Some(recipe.minimum_ptx.into()), + minimum_sm: Some(recipe.minimum_sm.into()), + }, + ], + packed_atomic: None, + redux: None, + vote: None, + active_mask: None, + warp_match: None, + warp_barrier: None, + warp_shuffle: None, + dot_product: None, + packed_alu: None, + integer_minmax: None, + packed_conversion: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cp_async_copy: None, + cp_async_control: None, + cp_async_mbarrier: None, + mbarrier_basic: None, + movmatrix: None, + mbarrier_extended: None, + register_mma: Some(mma), + sparse_mma: None, + prmt: None, + cluster_barrier: None, + wgmma_control: None, + special_register: None, + debug_control: None, + cluster_memory: None, + clc: None, + tma: None, + tcgen05: None, + ldmatrix_variant: None, + ldmatrix_safety: None, + ldmatrix_adapter: None, + selected_address_space: None, + expected_ptx: InstructionPattern { + mnemonic: "mma".into(), + modifiers: recipe + .ptx_modifiers + .iter() + .map(|value| (*value).into()) + .collect(), + operands: recipe + .ptx_register_counts + .map(|length| OperandPattern::RegisterList { length }) + .into(), + }, + summary, + }); + } + Ok(records) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/mma/sparse.rs b/crates/cuda-intrinsics-gen/src/resolve/families/mma/sparse.rs new file mode 100644 index 0000000000..4749642eb8 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/mma/sparse.rs @@ -0,0 +1,907 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, ImportedIntrinsic, IntrinsicBackend, OverlayBackendLowering, + OverlayIntrinsic, RuntimeValidation, SparseMma, SparseMmaAccumulator, SparseMmaAdapter, + SparseMmaCompatibilitySource, SparseMmaElement, SparseMmaF8F6F4Admission, + SparseMmaF8F6F4F16Admission, SparseMmaIntegerAdmission, SparseMmaLayout, SparseMmaLlvmAdapter, + SparseMmaMetadata, SparseMmaOverflow, SparseMmaParticipation, SparseMmaSelector, + SparseMmaShape, +}; +use crate::ptx::{InstructionPattern, OperandPattern}; +use anyhow::{Context, Result, ensure}; +use std::collections::BTreeSet; + +use crate::resolve::guards::*; + +pub(in crate::resolve) const SPARSE_MMA_F8F6F4_TARGETS: &str = "sm_120a|sm_120f|sm_121a|sm_121f"; +#[derive(Clone, Copy)] +pub(in crate::resolve) struct SparseMmaCarrierRecipe { + pub(in crate::resolve) adapter: SparseMmaAdapter, + pub(in crate::resolve) llvm_adapter: SparseMmaLlvmAdapter, + pub(in crate::resolve) accumulator: SparseMmaAccumulator, + pub(in crate::resolve) selector: SparseMmaSelector, + pub(in crate::resolve) a_registers: usize, + pub(in crate::resolve) b_registers: usize, +} + +impl SparseMmaCarrierRecipe { + fn accumulator_registers(self) -> usize { + match self.accumulator { + SparseMmaAccumulator::F16 => 2, + SparseMmaAccumulator::F32 | SparseMmaAccumulator::S32 => 4, + } + } + + fn operand_count(self) -> usize { + self.accumulator_registers() + self.a_registers + self.b_registers + 2 + } + + fn selector_index(self) -> usize { + self.operand_count() - 1 + } + + fn selector_upper_exclusive(self) -> u8 { + match self.selector { + SparseMmaSelector::ImmediateZeroOrOne => 2, + SparseMmaSelector::ImmediateZero => 1, + } + } + + pub(in crate::resolve) fn rust_arguments(self) -> Vec { + let accumulator = match self.accumulator { + SparseMmaAccumulator::F16 => "[u32; 2]", + SparseMmaAccumulator::F32 => "[f32; 4]", + SparseMmaAccumulator::S32 => "[i32; 4]", + }; + vec![ + accumulator.into(), + format!("[u32; {}]", self.a_registers), + format!("[u32; {}]", self.b_registers), + "u32".into(), + "u32".into(), + ] + } + + pub(in crate::resolve) fn dialect_operands(self) -> Vec { + let accumulator = match self.accumulator { + SparseMmaAccumulator::F16 => "i32", + SparseMmaAccumulator::F32 => "f32", + SparseMmaAccumulator::S32 => "i32", + }; + let accumulator_registers = self.accumulator_registers(); + std::iter::repeat_n(accumulator.to_owned(), accumulator_registers) + .chain(std::iter::repeat_n( + "u32".to_owned(), + self.operand_count() - accumulator_registers, + )) + .collect() + } + + pub(in crate::resolve) fn llvm_arguments(self) -> Vec { + let accumulator = match self.accumulator { + SparseMmaAccumulator::F16 => "v2f16", + SparseMmaAccumulator::F32 => "f32", + SparseMmaAccumulator::S32 => "i32", + }; + std::iter::repeat_n("i32".to_owned(), self.a_registers + self.b_registers) + .chain(std::iter::repeat_n( + accumulator.to_owned(), + self.accumulator_registers(), + )) + .chain(std::iter::repeat_n("i32".to_owned(), 2)) + .collect() + } + + fn rust_result(self) -> String { + match self.accumulator { + SparseMmaAccumulator::F16 => "[u32; 2]", + SparseMmaAccumulator::F32 => "[f32; 4]", + SparseMmaAccumulator::S32 => "[i32; 4]", + } + .into() + } + + fn dialect_results(self) -> Vec { + match self.accumulator { + SparseMmaAccumulator::F16 => vec!["i32".into(); 2], + SparseMmaAccumulator::F32 => vec!["f32".into(); 4], + SparseMmaAccumulator::S32 => vec!["i32".into(); 4], + } + } + + fn llvm_results(self) -> Vec { + match self.accumulator { + SparseMmaAccumulator::F16 => vec!["v2f16".into(); 2], + SparseMmaAccumulator::F32 => vec!["f32".into(); 4], + SparseMmaAccumulator::S32 => vec!["i32".into(); 4], + } + } + + pub(in crate::resolve) fn expected_ptx_operands(self) -> Vec { + vec![ + OperandPattern::RegisterList { + length: self.accumulator_registers(), + }, + OperandPattern::RegisterList { + length: self.a_registers, + }, + OperandPattern::RegisterList { + length: self.b_registers, + }, + OperandPattern::RegisterList { + length: self.accumulator_registers(), + }, + OperandPattern::Register, + OperandPattern::Immediate, + ] + } + + fn imported_properties(self) -> Vec { + let selector = self.selector_index(); + vec![ + format!("ImmArg"), + "IntrNoCallback".into(), + "IntrNoMem".into(), + format!("Range", self.selector_upper_exclusive()), + ] + } +} + +pub(in crate::resolve) fn sparse_mma_carrier_recipe( + shape: SparseMmaShape, + a_element: SparseMmaElement, + b_element: SparseMmaElement, +) -> Option { + use SparseMmaElement::{E2m1, E2m3, E3m2, E4m3, E5m2, S4, S8, U4, U8}; + use SparseMmaShape::{M16n8k32, M16n8k64, M16n8k128}; + + match (shape, a_element, b_element) { + (M16n8k32, S8 | U8, S8 | U8) | (M16n8k64, S4 | U4, S4 | U4) => { + Some(SparseMmaCarrierRecipe { + adapter: SparseMmaAdapter::C4I32A2U32B2U32MetadataU32SelectorU32ToD4I32, + llvm_adapter: SparseMmaLlvmAdapter::A2I32B2I32C4I32MetadataI32SelectorI32ToD4I32, + accumulator: SparseMmaAccumulator::S32, + selector: SparseMmaSelector::ImmediateZeroOrOne, + a_registers: 2, + b_registers: 2, + }) + } + (M16n8k64, S8 | U8, S8 | U8) => Some(SparseMmaCarrierRecipe { + adapter: SparseMmaAdapter::C4I32A4U32B4U32MetadataU32SelectorU32ToD4I32, + llvm_adapter: SparseMmaLlvmAdapter::A4I32B4I32C4I32MetadataI32SelectorI32ToD4I32, + accumulator: SparseMmaAccumulator::S32, + selector: SparseMmaSelector::ImmediateZero, + a_registers: 4, + b_registers: 4, + }), + (M16n8k128, S4 | U4, S4 | U4) => Some(SparseMmaCarrierRecipe { + adapter: SparseMmaAdapter::C4I32A4U32B4U32MetadataU32SelectorU32ToD4I32, + llvm_adapter: SparseMmaLlvmAdapter::A4I32B4I32C4I32MetadataI32SelectorI32ToD4I32, + accumulator: SparseMmaAccumulator::S32, + selector: SparseMmaSelector::ImmediateZero, + a_registers: 4, + b_registers: 4, + }), + (M16n8k64, E2m1 | E2m3 | E3m2 | E4m3 | E5m2, E2m1 | E2m3 | E3m2 | E4m3 | E5m2) => { + Some(SparseMmaCarrierRecipe { + adapter: SparseMmaAdapter::C4F32A4U32B4U32MetadataU32SelectorU32ToD4F32, + llvm_adapter: SparseMmaLlvmAdapter::A4I32B4I32C4F32MetadataI32SelectorI32ToD4F32, + accumulator: SparseMmaAccumulator::F32, + selector: SparseMmaSelector::ImmediateZero, + a_registers: 4, + b_registers: 4, + }) + } + _ => None, + } +} + +pub(in crate::resolve) fn sparse_mma_f8f6f4_f16_carrier_recipe( + shape: SparseMmaShape, + a_element: SparseMmaElement, + b_element: SparseMmaElement, +) -> Option { + use SparseMmaElement::{E2m1, E2m3, E3m2, E4m3, E5m2}; + + if shape != SparseMmaShape::M16n8k64 + || !matches!(a_element, E2m1 | E2m3 | E3m2 | E4m3 | E5m2) + || !matches!(b_element, E2m1 | E2m3 | E3m2 | E4m3 | E5m2) + { + return None; + } + Some(SparseMmaCarrierRecipe { + adapter: SparseMmaAdapter::C2U32A4U32B4U32MetadataU32SelectorU32ToD2U32, + llvm_adapter: SparseMmaLlvmAdapter::A4I32B4I32C2V2F16MetadataI32SelectorI32ToD2V2F16, + accumulator: SparseMmaAccumulator::F16, + selector: SparseMmaSelector::ImmediateZero, + a_registers: 4, + b_registers: 4, + }) +} + +pub(in crate::resolve) struct SparseMmaIdentity { + pub(in crate::resolve) id: String, + pub(in crate::resolve) operation_key: String, + pub(in crate::resolve) source_record: String, + pub(in crate::resolve) llvm_symbol: String, + pub(in crate::resolve) ptx_modifiers: Vec<&'static str>, +} + +pub(in crate::resolve) struct SparseMmaRecipe { + pub(in crate::resolve) identity: SparseMmaIdentity, + pub(in crate::resolve) carrier: SparseMmaCarrierRecipe, +} + +pub(in crate::resolve) fn sparse_mma_shape_name(shape: SparseMmaShape) -> &'static str { + match shape { + SparseMmaShape::M16n8k32 => "m16n8k32", + SparseMmaShape::M16n8k64 => "m16n8k64", + SparseMmaShape::M16n8k128 => "m16n8k128", + } +} + +pub(in crate::resolve) fn sparse_mma_element_name(element: SparseMmaElement) -> &'static str { + match element { + SparseMmaElement::E2m1 => "e2m1", + SparseMmaElement::E2m3 => "e2m3", + SparseMmaElement::E3m2 => "e3m2", + SparseMmaElement::E4m3 => "e4m3", + SparseMmaElement::E5m2 => "e5m2", + SparseMmaElement::S4 => "s4", + SparseMmaElement::U4 => "u4", + SparseMmaElement::S8 => "s8", + SparseMmaElement::U8 => "u8", + } +} + +pub(in crate::resolve) fn sparse_mma_identity(mma: &SparseMma) -> SparseMmaIdentity { + let shape = sparse_mma_shape_name(mma.shape); + let a_element = sparse_mma_element_name(mma.a_element); + let b_element = sparse_mma_element_name(mma.b_element); + if matches!( + mma.accumulator, + SparseMmaAccumulator::F16 | SparseMmaAccumulator::F32 + ) { + let scalar = match mma.accumulator { + SparseMmaAccumulator::F16 => "f16", + SparseMmaAccumulator::F32 => "f32", + SparseMmaAccumulator::S32 => unreachable!(), + }; + return SparseMmaIdentity { + id: format!( + "mma_sp_ordered_metadata_{shape}_kind_f8f6f4_{scalar}_{a_element}_{b_element}_{scalar}" + ), + operation_key: format!( + "matrix.mma.sp.{shape}.row.col.kind_f8f6f4.{scalar}.{a_element}.{b_element}.{scalar}.not_applicable.ordered_metadata" + ), + source_record: format!( + "int_nvvm_mma_sp_ordered_metadata_{shape}_row_col_kind_f8f6f4_{scalar}_{a_element}_{b_element}_{scalar}" + ), + llvm_symbol: format!( + "llvm.nvvm.mma.sp.ordered.metadata.{shape}.row.col.kind.f8f6f4.{scalar}.{a_element}.{b_element}.{scalar}" + ), + ptx_modifiers: vec![ + "sp::ordered_metadata", + "sync", + "aligned", + shape, + "row", + "col", + "kind::f8f6f4", + scalar, + a_element, + b_element, + scalar, + ], + }; + } + let compact_elements = if mma.a_element == mma.b_element { + a_element.to_owned() + } else { + format!("{a_element}_{b_element}") + }; + let dotted_elements = if mma.a_element == mma.b_element { + a_element.to_owned() + } else { + format!("{a_element}.{b_element}") + }; + let metadata_id_prefix = match mma.metadata { + SparseMmaMetadata::Standard => "", + SparseMmaMetadata::Ordered => "ordered_metadata_", + }; + let metadata_source_prefix = metadata_id_prefix; + let metadata_symbol_prefix = match mma.metadata { + SparseMmaMetadata::Standard => "", + SparseMmaMetadata::Ordered => "ordered.metadata.", + }; + let metadata_key = match mma.metadata { + SparseMmaMetadata::Standard => "standard_metadata", + SparseMmaMetadata::Ordered => "ordered_metadata", + }; + let overflow_id_suffix = match mma.overflow { + SparseMmaOverflow::NotApplicable => unreachable!("integer sparse MMA has overflow"), + SparseMmaOverflow::Wrapping => "", + SparseMmaOverflow::Satfinite => "_satfinite", + }; + let overflow_source_prefix = match mma.overflow { + SparseMmaOverflow::NotApplicable => unreachable!("integer sparse MMA has overflow"), + SparseMmaOverflow::Wrapping => "", + SparseMmaOverflow::Satfinite => "satfinite_", + }; + let overflow_symbol_prefix = match mma.overflow { + SparseMmaOverflow::NotApplicable => unreachable!("integer sparse MMA has overflow"), + SparseMmaOverflow::Wrapping => "", + SparseMmaOverflow::Satfinite => "satfinite.", + }; + let overflow_key = match mma.overflow { + SparseMmaOverflow::NotApplicable => unreachable!("integer sparse MMA has overflow"), + SparseMmaOverflow::Wrapping => "wrapping", + SparseMmaOverflow::Satfinite => "satfinite", + }; + + let mut ptx_modifiers = vec![ + match mma.metadata { + SparseMmaMetadata::Standard => "sp", + SparseMmaMetadata::Ordered => "sp::ordered_metadata", + }, + "sync", + "aligned", + shape, + "row", + "col", + ]; + if mma.overflow == SparseMmaOverflow::Satfinite { + ptx_modifiers.push("satfinite"); + } + ptx_modifiers.extend(["s32", a_element, b_element, "s32"]); + + SparseMmaIdentity { + id: format!( + "mma_sp_{metadata_id_prefix}{shape}_s32_{compact_elements}{overflow_id_suffix}" + ), + operation_key: format!( + "matrix.mma.sp.{shape}.row.col.s32.{a_element}.{b_element}.s32.{overflow_key}.{metadata_key}" + ), + source_record: format!( + "int_nvvm_mma_sp_{metadata_source_prefix}{shape}_row_col_{overflow_source_prefix}{compact_elements}" + ), + llvm_symbol: format!( + "llvm.nvvm.mma.sp.{metadata_symbol_prefix}{shape}.row.col.{overflow_symbol_prefix}{dotted_elements}" + ), + ptx_modifiers, + } +} + +pub(in crate::resolve) fn sparse_mma_recipe(mma: &SparseMma) -> Option { + let carrier = if mma.accumulator == SparseMmaAccumulator::F16 { + sparse_mma_f8f6f4_f16_carrier_recipe(mma.shape, mma.a_element, mma.b_element)? + } else { + sparse_mma_carrier_recipe(mma.shape, mma.a_element, mma.b_element)? + }; + + let scalar_contract = match carrier.accumulator { + SparseMmaAccumulator::F16 | SparseMmaAccumulator::F32 => { + mma.accumulator == carrier.accumulator + && mma.overflow == SparseMmaOverflow::NotApplicable + && mma.metadata == SparseMmaMetadata::Ordered + } + SparseMmaAccumulator::S32 => { + mma.accumulator == SparseMmaAccumulator::S32 + && matches!( + mma.overflow, + SparseMmaOverflow::Wrapping | SparseMmaOverflow::Satfinite + ) + } + }; + if !scalar_contract + || mma.a_layout != SparseMmaLayout::Row + || mma.b_layout != SparseMmaLayout::Col + || mma.selector != carrier.selector + || mma.participation + != SparseMmaParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes + || mma.adapter != carrier.adapter + || mma.llvm_adapter != carrier.llvm_adapter + || mma.compatibility_source != SparseMmaCompatibilitySource::GeneratedStub + { + return None; + } + + Some(SparseMmaRecipe { + identity: sparse_mma_identity(mma), + carrier, + }) +} + +pub(in crate::resolve) fn sparse_mma_minimum_ptx(mma: &SparseMma) -> &'static str { + if matches!( + mma.accumulator, + SparseMmaAccumulator::F16 | SparseMmaAccumulator::F32 + ) { + return "8.7"; + } + match mma.metadata { + SparseMmaMetadata::Standard => "7.1", + SparseMmaMetadata::Ordered => "8.5", + } +} + +pub(in crate::resolve) fn sparse_mma_hardware( + mma: &SparseMma, +) -> (&'static str, Option<&'static str>) { + match mma.accumulator { + SparseMmaAccumulator::F16 | SparseMmaAccumulator::F32 => (SPARSE_MMA_F8F6F4_TARGETS, None), + SparseMmaAccumulator::S32 => ("all", Some("sm_80")), + } +} + +pub(in crate::resolve) fn sparse_mma_ptx_section(_: SparseMmaMetadata) -> &'static str { + "9.7.15.6.3 Multiply-and-Accumulate Instruction: mma.sp" +} + +pub(in crate::resolve) fn expand_sparse_mma_integer_admission( + admission: &SparseMmaIntegerAdmission, +) -> Result> { + ensure!( + !admission.variants.is_empty(), + "compact sparse integer MMA admission has no variants" + ); + ensure!( + admission.runtime_validation == RuntimeValidation::Unexecuted, + "sparse integer MMA runtime validation may be marked executed only with GPU evidence" + ); + + let mut seen = BTreeSet::new(); + let mut records = Vec::with_capacity(admission.variants.len()); + for variant in &admission.variants { + ensure!( + seen.insert(( + variant.shape, + variant.a_element, + variant.b_element, + variant.overflow, + )), + "compact sparse integer MMA admission contains a duplicate variant" + ); + let carrier = sparse_mma_carrier_recipe( + variant.shape, + variant.a_element, + variant.b_element, + ) + .with_context( + || "compact sparse integer MMA admission uses unsupported or mixed-width elements", + )?; + let mma = SparseMma { + shape: variant.shape, + accumulator: SparseMmaAccumulator::S32, + a_element: variant.a_element, + b_element: variant.b_element, + a_layout: SparseMmaLayout::Row, + b_layout: SparseMmaLayout::Col, + overflow: variant.overflow, + metadata: admission.metadata, + selector: carrier.selector, + participation: + SparseMmaParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes, + adapter: carrier.adapter, + llvm_adapter: carrier.llvm_adapter, + compatibility_source: SparseMmaCompatibilitySource::GeneratedStub, + runtime_validation: admission.runtime_validation, + }; + let recipe = sparse_mma_recipe(&mma).with_context( + || "compact sparse integer MMA admission requests a variant outside the closed recipe set", + )?; + let signedness = |element| match element { + SparseMmaElement::S4 => "signed", + SparseMmaElement::U4 => "unsigned", + SparseMmaElement::S8 => "signed", + SparseMmaElement::U8 => "unsigned", + _ => unreachable!("integer admission rejects floating formats"), + }; + let width = match (variant.a_element, variant.b_element) { + ( + SparseMmaElement::S4 | SparseMmaElement::U4, + SparseMmaElement::S4 | SparseMmaElement::U4, + ) => "INT4", + ( + SparseMmaElement::S8 | SparseMmaElement::U8, + SparseMmaElement::S8 | SparseMmaElement::U8, + ) => "INT8", + _ => unreachable!("carrier selection rejects mixed element widths"), + }; + let overflow = match variant.overflow { + SparseMmaOverflow::NotApplicable => { + unreachable!("integer admission rejects inapplicable overflow") + } + SparseMmaOverflow::Wrapping => "wrapping", + SparseMmaOverflow::Satfinite => "saturating", + }; + let metadata = match admission.metadata { + SparseMmaMetadata::Standard => "", + SparseMmaMetadata::Ordered => " with ordered sparsity metadata", + }; + let summary = format!( + "Multiplies warp-distributed sparse {} A and {} B {width} fragments{metadata} and adds a {overflow} s32 accumulator.", + signedness(variant.a_element), + signedness(variant.b_element), + ); + records.push(sparse_mma_overlay_record( + String::new(), + mma, + recipe, + &admission.llvm_evidence_profile, + &admission.libnvvm_evidence_profile, + summary, + )); + } + Ok(records) +} + +pub(in crate::resolve) fn expand_sparse_mma_f8f6f4_admission( + admission: &SparseMmaF8F6F4Admission, +) -> Result> { + ensure!( + admission.runtime_validation == RuntimeValidation::Unexecuted, + "sparse f8f6f4 MMA runtime validation may be marked executed only with GPU evidence" + ); + + let formats = vec![ + SparseMmaElement::E2m1, + SparseMmaElement::E2m3, + SparseMmaElement::E3m2, + SparseMmaElement::E4m3, + SparseMmaElement::E5m2, + ]; + ensure!( + admission.a_elements == formats, + "compact sparse f8f6f4 MMA admission must list the canonical five A formats" + ); + ensure!( + admission.b_elements == formats, + "compact sparse f8f6f4 MMA admission must list the canonical five B formats" + ); + ensure!( + admission.product_count + == admission + .a_elements + .len() + .checked_mul(admission.b_elements.len()) + .context("compact sparse f8f6f4 MMA admission product count overflow")? + && admission.product_count == 25, + "compact sparse f8f6f4 MMA admission product_count must be exactly 25" + ); + let mut records = Vec::with_capacity(admission.product_count); + for &a_element in &admission.a_elements { + for &b_element in &admission.b_elements { + let carrier = sparse_mma_carrier_recipe(SparseMmaShape::M16n8k64, a_element, b_element) + .with_context( + || "compact sparse f8f6f4 MMA admission uses an unsupported format", + )?; + ensure!( + carrier.accumulator == SparseMmaAccumulator::F32, + "compact sparse f8f6f4 MMA admission contains an integer format" + ); + let mma = SparseMma { + shape: SparseMmaShape::M16n8k64, + accumulator: SparseMmaAccumulator::F32, + a_element, + b_element, + a_layout: SparseMmaLayout::Row, + b_layout: SparseMmaLayout::Col, + overflow: SparseMmaOverflow::NotApplicable, + metadata: SparseMmaMetadata::Ordered, + selector: carrier.selector, + participation: + SparseMmaParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes, + adapter: carrier.adapter, + llvm_adapter: carrier.llvm_adapter, + compatibility_source: SparseMmaCompatibilitySource::GeneratedStub, + runtime_validation: admission.runtime_validation, + }; + let recipe = sparse_mma_recipe(&mma).with_context(|| { + "compact sparse f8f6f4 MMA admission requests a variant outside the closed recipe set" + })?; + let summary = format!( + "Multiplies warp-distributed sparse {} A and {} B fragments and adds an f32 accumulator.", + sparse_mma_element_name(a_element), + sparse_mma_element_name(b_element), + ); + records.push(sparse_mma_overlay_record( + String::new(), + mma, + recipe, + &admission.llvm_evidence_profile, + &admission.libnvvm_evidence_profile, + summary, + )); + } + } + ensure!(records.len() == admission.product_count); + Ok(records) +} + +pub(in crate::resolve) const SPARSE_MMA_F8F6F4_ELEMENTS: [SparseMmaElement; 5] = [ + SparseMmaElement::E2m1, + SparseMmaElement::E2m3, + SparseMmaElement::E3m2, + SparseMmaElement::E4m3, + SparseMmaElement::E5m2, +]; + +pub(in crate::resolve) fn expand_sparse_mma_f8f6f4_f16_admission( + admission: &SparseMmaF8F6F4F16Admission, +) -> Result> { + ensure!( + admission.runtime_validation == RuntimeValidation::Unexecuted, + "sparse f8f6f4 F16 MMA runtime validation may be marked executed only with GPU evidence" + ); + ensure!( + !admission.llvm_evidence_profile.trim().is_empty() + && !admission.libnvvm_evidence_profile.trim().is_empty(), + "sparse f8f6f4 F16 MMA admission requires both backend evidence profiles" + ); + ensure!( + admission.a_elements == SPARSE_MMA_F8F6F4_ELEMENTS + && admission.b_elements == SPARSE_MMA_F8F6F4_ELEMENTS + && admission.product_count == 25, + "sparse f8f6f4 F16 MMA admission must contain the canonical 5 by 5 element matrix" + ); + + let mut records = Vec::with_capacity(admission.product_count); + for &a_element in &admission.a_elements { + for &b_element in &admission.b_elements { + let carrier = sparse_mma_f8f6f4_f16_carrier_recipe( + SparseMmaShape::M16n8k64, + a_element, + b_element, + ) + .context("compact sparse f8f6f4 F16 MMA admission uses an unsupported format")?; + let mma = SparseMma { + shape: SparseMmaShape::M16n8k64, + accumulator: SparseMmaAccumulator::F16, + a_element, + b_element, + a_layout: SparseMmaLayout::Row, + b_layout: SparseMmaLayout::Col, + overflow: SparseMmaOverflow::NotApplicable, + metadata: SparseMmaMetadata::Ordered, + selector: SparseMmaSelector::ImmediateZero, + participation: + SparseMmaParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes, + adapter: carrier.adapter, + llvm_adapter: carrier.llvm_adapter, + compatibility_source: SparseMmaCompatibilitySource::GeneratedStub, + runtime_validation: admission.runtime_validation, + }; + let recipe = sparse_mma_recipe(&mma).context( + "compact sparse f8f6f4 F16 MMA admission requests a variant outside the closed recipe set", + )?; + let summary = format!( + "Multiplies warp-distributed sparse {} A and {} B fragments and adds a packed F16 accumulator.", + sparse_mma_element_name(a_element), + sparse_mma_element_name(b_element), + ); + records.push(sparse_mma_overlay_record( + String::new(), + mma, + recipe, + &admission.llvm_evidence_profile, + &admission.libnvvm_evidence_profile, + summary, + )); + } + } + ensure!(records.len() == admission.product_count); + Ok(records) +} + +pub(in crate::resolve) fn sparse_mma_overlay_record( + abi_id: String, + mma: SparseMma, + recipe: SparseMmaRecipe, + llvm_evidence_profile: &str, + libnvvm_evidence_profile: &str, + summary: String, +) -> OverlayIntrinsic { + let identity = &recipe.identity; + let minimum_ptx = sparse_mma_minimum_ptx(&mma); + let (targets, minimum_sm) = sparse_mma_hardware(&mma); + OverlayIntrinsic { + id: identity.id.clone(), + abi_id, + operation_key: identity.operation_key.clone(), + family: "sparse_mma".into(), + source: None, + source_record: Some(identity.source_record.clone()), + rust_module: "matrix".into(), + rust_name: identity.id.clone(), + rust_arguments: recipe.carrier.rust_arguments(), + rust_result: recipe.carrier.rust_result(), + safe: false, + must_use: true, + safe_allowlist_reason: None, + public_rust_path: format!("cuda_intrinsics::matrix::{}", identity.id), + compatibility_rust_paths: vec![format!("cuda_device::wmma::{}", identity.id)], + dialect_op_type: "SparseMmaOp".into(), + dialect_op_name: "nvvm.sparse_mma".into(), + dialect_operands: recipe.carrier.dialect_operands(), + dialect_results: recipe.carrier.dialect_results(), + llvm_symbol: Some(identity.llvm_symbol.clone()), + resolved_llvm_symbol: None, + llvm_arguments: recipe.carrier.llvm_arguments(), + llvm_results: recipe.carrier.llvm_results(), + pure: false, + memory: "none".into(), + convergent: true, + execution_scope: "warp".into(), + minimum_ptx: minimum_ptx.into(), + minimum_sm: minimum_sm.map(str::to_owned), + ptx_result: recipe.carrier.rust_result(), + targets: targets.into(), + ptx_isa_version: "9.3".into(), + ptx_isa_section: sparse_mma_ptx_section(mma.metadata).into(), + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma-sp".into(), + lowering: "generated_sparse_mma".into(), + backend_lowerings: [ + (IntrinsicBackend::LlvmNvptx, llvm_evidence_profile), + (IntrinsicBackend::LibNvvm, libnvvm_evidence_profile), + ] + .into_iter() + .map(|(backend, evidence_profile)| OverlayBackendLowering { + backend, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: evidence_profile.into(), + targets: None, + minimum_ptx: Some(minimum_ptx.into()), + minimum_sm: minimum_sm.map(str::to_owned), + }) + .collect(), + packed_atomic: None, + redux: None, + vote: None, + active_mask: None, + warp_match: None, + warp_barrier: None, + warp_shuffle: None, + dot_product: None, + packed_alu: None, + integer_minmax: None, + packed_conversion: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cp_async_copy: None, + cp_async_control: None, + cp_async_mbarrier: None, + mbarrier_basic: None, + movmatrix: None, + mbarrier_extended: None, + register_mma: None, + sparse_mma: Some(mma), + prmt: None, + cluster_barrier: None, + wgmma_control: None, + special_register: None, + debug_control: None, + cluster_memory: None, + clc: None, + tma: None, + tcgen05: None, + ldmatrix_variant: None, + ldmatrix_safety: None, + ldmatrix_adapter: None, + selected_address_space: None, + expected_ptx: InstructionPattern { + mnemonic: "mma".into(), + modifiers: identity + .ptx_modifiers + .iter() + .map(|value| (*value).into()) + .collect(), + operands: recipe.carrier.expected_ptx_operands(), + }, + summary, + } +} + +pub(in crate::resolve) fn validate_sparse_mma_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, +) -> Result<()> { + let mma = policy + .sparse_mma + .as_ref() + .with_context(|| format!("{} has no closed sparse-MMA contract", policy.id))?; + let recipe = sparse_mma_recipe(mma) + .with_context(|| format!("{} requests an unsupported sparse-MMA variant", policy.id))?; + let identity = &recipe.identity; + let minimum_ptx = sparse_mma_minimum_ptx(mma); + let (targets, minimum_sm) = sparse_mma_hardware(mma); + ensure!( + policy.id == identity.id + && policy.operation_key == identity.operation_key + && policy.source.is_none() + && policy.source_record.as_deref() == Some(identity.source_record.as_str()) + && policy.llvm_symbol.as_deref() == Some(identity.llvm_symbol.as_str()) + && policy.resolved_llvm_symbol.is_none(), + "{} sparse-MMA identity does not match its closed recipe", + policy.id + ); + ensure!( + policy.rust_module == "matrix" + && policy.rust_name == identity.id + && policy.public_rust_path == format!("cuda_intrinsics::matrix::{}", identity.id) + && policy.rust_arguments == recipe.carrier.rust_arguments() + && policy.rust_result == recipe.carrier.rust_result() + && !policy.safe + && policy.must_use + && policy.safe_allowlist_reason.is_none() + && policy.compatibility_rust_paths == [format!("cuda_device::wmma::{}", identity.id)], + "{} must preserve its unsafe must-use Rust sparse-MMA API", + policy.id + ); + ensure!( + policy.dialect_op_type == "SparseMmaOp" + && policy.dialect_op_name == "nvvm.sparse_mma" + && policy.dialect_operands == recipe.carrier.dialect_operands() + && policy.dialect_results == recipe.carrier.dialect_results() + && policy.llvm_arguments == recipe.carrier.llvm_arguments() + && policy.llvm_results == recipe.carrier.llvm_results() + && policy.ptx_result == recipe.carrier.rust_result() + && policy.lowering == "generated_sparse_mma", + "{} sparse-MMA carrier or lowering adapter disagrees with its recipe", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == "none" + && policy.convergent + && policy.execution_scope == "warp" + && policy.minimum_ptx == minimum_ptx + && policy.minimum_sm.as_deref() == minimum_sm + && policy.targets == targets, + "{} sparse-MMA effects or target floor disagree with PTX", + policy.id + ); + ensure!( + policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section == sparse_mma_ptx_section(mma.metadata) + && policy.ptx_isa_url + == "https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma-sp", + "{} sparse-MMA PTX provenance disagrees with the reviewed recipe", + policy.id + ); + ensure!( + declaration.classes == ["SDPatternOperator", "Intrinsic", "NVVM_MMA_SP"] + && declaration.properties == recipe.carrier.imported_properties() + && declaration.selections.len() == 1 + && (!matches!( + mma.accumulator, + SparseMmaAccumulator::F16 | SparseMmaAccumulator::F32 + ) || declaration.selections[0].predicates == ["Subtarget->hasMMABlockScale()"]) + && selection_matches_policy(policy, &declaration.selections[0])?, + "{} imported sparse MMA declaration changed its class, immediate range, properties, or exact selection contract", + policy.id + ); + ensure!( + policy.expected_ptx.mnemonic == "mma" + && policy.expected_ptx.modifiers == identity.ptx_modifiers + && policy.expected_ptx.operands == recipe.carrier.expected_ptx_operands(), + "{} expected PTX does not match its exact sparse-MMA spelling", + policy.id + ); + ensure_exact_inline_ptx_backends( + policy, + [ + (IntrinsicBackend::LlvmNvptx, minimum_ptx, minimum_sm), + (IntrinsicBackend::LibNvvm, minimum_ptx, minimum_sm), + ], + "sparse MMA", + )?; + ensure_no_other_family_contract(policy, "sparse MMA")?; + Ok(()) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/mma/validate.rs b/crates/cuda-intrinsics-gen/src/resolve/families/mma/validate.rs new file mode 100644 index 0000000000..b76e4112bd --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/mma/validate.rs @@ -0,0 +1,736 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, ImportedIntrinsic, IntrinsicBackend, OverlayIntrinsic, RegisterMma, + RegisterMmaAccumulator, RegisterMmaAdapter, RegisterMmaCompatibilitySource, RegisterMmaKind, + RegisterMmaLayout, RegisterMmaOperation, RegisterMmaOverflow, RegisterMmaParticipation, + RegisterMmaShape, RuntimeValidation, +}; +use crate::ptx::{InstructionPattern, OperandPattern}; +use anyhow::{Context, Result, bail, ensure}; +use std::collections::BTreeSet; + +use super::*; +use crate::resolve::guards::*; + +pub(in crate::resolve) fn validate_register_mma_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, +) -> Result<()> { + let mma = policy + .register_mma + .as_ref() + .with_context(|| format!("{} has no closed register-MMA contract", policy.id))?; + match mma.kind { + Some(RegisterMmaKind::Standard) => { + return validate_register_mma_fp8_policy(policy, declaration, mma); + } + Some(RegisterMmaKind::F8f6f4) => { + return validate_register_mma_f8f6f4_policy(policy, declaration, mma); + } + Some(RegisterMmaKind::Mxf8f6f4) => { + return validate_register_mma_mxf8f6f4_policy(policy, declaration, mma); + } + None if register_mma_f8f6f4_element_name(mma.a_element).is_some() + || register_mma_f8f6f4_element_name(mma.b_element).is_some() => + { + return validate_register_mma_f8f6f4_policy(policy, declaration, mma); + } + None => {} + } + let recipe = register_mma_recipe(mma) + .with_context(|| format!("{} requests an unsupported register-MMA variant", policy.id))?; + let abi_matches = matches!( + recipe.id, + "mma_m16n8k4_f32_tf32" + | "mma_m16n8k8_f16_f16" + | "mma_m16n8k8_f32_bf16" + | "mma_m16n8k8_f32_f16" + | "mma_m16n8k16_f16_f16" + ) || policy.abi_id == recipe.abi_id; + ensure!( + policy.id == recipe.id + && abi_matches + && policy.operation_key == recipe.operation_key + && policy.source.is_none() + && policy.source_record.as_deref() == Some(recipe.source_record) + && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) + && policy.resolved_llvm_symbol.is_none(), + "{} register-MMA identity does not match its closed recipe", + policy.id + ); + ensure!( + policy.rust_module == "matrix" + && policy.rust_name == recipe.id + && policy.rust_arguments == recipe.rust_arguments + && policy.rust_result == recipe.rust_result + && !policy.safe + && policy.must_use + && policy.safe_allowlist_reason.is_none() + && policy.compatibility_rust_paths == [format!("cuda_device::wmma::{}", recipe.id)], + "{} must preserve its unsafe must-use Rust MMA API", + policy.id + ); + ensure!( + policy.dialect_op_type == recipe.dialect_op_type + && policy.dialect_op_name == recipe.dialect_op_name + && policy.dialect_operands == recipe.dialect_operands + && policy.dialect_results == recipe.dialect_results + && policy.llvm_arguments == recipe.llvm_arguments + && policy.llvm_results == recipe.llvm_results + && policy.ptx_result == recipe.rust_result + && mma.adapter == recipe.adapter + && mma.compatibility_source == recipe.compatibility_source + && policy.lowering == "generated_register_mma", + "{} register-MMA carrier or lowering adapter disagrees with its recipe", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == "none" + && policy.convergent + && policy.execution_scope == "warp" + && policy.minimum_ptx == recipe.minimum_ptx + && policy.minimum_sm.as_deref() == Some(recipe.minimum_sm) + && policy.targets == "all", + "{} register-MMA effects or target floor disagree with PTX", + policy.id + ); + ensure!( + policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section == "9.7.15.5.14 Multiply-and-Accumulate Instruction: mma" + && policy.ptx_isa_url + == "https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma", + "{} register-MMA PTX provenance disagrees with the reviewed recipe", + policy.id + ); + ensure!( + declaration.classes.iter().any(|class| class == "NVVM_MMA") + && declaration.properties == ["IntrNoCallback", "IntrNoMem"] + && declaration.selections.len() == 1 + && selection_matches_policy(policy, &declaration.selections[0])?, + "{} imported MMA declaration changed its class, properties, or exact selection contract", + policy.id + ); + // These are imported facts. The BF16 selection is underconstrained. + let ampere_float_predicates = match recipe.id { + "mma_m16n8k4_f32_tf32" | "mma_m16n8k16_f16_f16" => Some([ + "Subtarget->getSmVersion() >= 80", + "Subtarget->getPTXVersion() >= 70", + ]), + "mma_m16n8k8_f16_f16" | "mma_m16n8k8_f32_bf16" | "mma_m16n8k8_f32_f16" => Some([ + "Subtarget->getPTXVersion() >= 65", + "Subtarget->getSmVersion() >= 75", + ]), + _ => None, + }; + if let Some(predicates) = ampere_float_predicates { + let selection = &declaration.selections[0]; + ensure!( + declaration.classes == ["SDPatternOperator", "Intrinsic", "NVVM_MMA"] + && selection.predicates == predicates + && selection.constraints.is_empty() + && mma.kind.is_none() + && mma.runtime_validation == RuntimeValidation::Unexecuted, + "{} imported Ampere floating-point MMA contract changed", + policy.id + ); + } + ensure!( + policy.expected_ptx.mnemonic == "mma" + && policy.expected_ptx.modifiers == recipe.ptx_modifiers + && policy.expected_ptx.operands + == recipe + .ptx_register_counts + .map(|length| OperandPattern::RegisterList { length }), + "{} expected PTX does not match its exact register-MMA spelling", + policy.id + ); + ensure_exact_inline_ptx_backends( + policy, + [ + ( + IntrinsicBackend::LlvmNvptx, + recipe.minimum_ptx, + Some(recipe.minimum_sm), + ), + ( + IntrinsicBackend::LibNvvm, + recipe.minimum_ptx, + Some(recipe.minimum_sm), + ), + ], + "register MMA", + )?; + ensure_no_other_family_contract(policy, "register MMA")?; + Ok(()) +} + +pub(in crate::resolve) fn validate_register_mma_fp8_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, + mma: &RegisterMma, +) -> Result<()> { + ensure!( + REGISTER_MMA_FP8_SHAPES.contains(&mma.shape), + "{} has an unsupported standard FP8 shape", + policy.id + ); + ensure!( + REGISTER_MMA_FP8_ACCUMULATORS.contains(&mma.accumulator), + "{} has an unsupported standard FP8 accumulator", + policy.id + ); + ensure!( + REGISTER_MMA_FP8_ELEMENTS.contains(&mma.a_element), + "{} has an unsupported standard FP8 A element", + policy.id + ); + ensure!( + REGISTER_MMA_FP8_ELEMENTS.contains(&mma.b_element), + "{} has an unsupported standard FP8 B element", + policy.id + ); + let (shape_name, a_count, b_count) = register_mma_fp8_shape_contract(mma.shape)?; + let minimum_ptx = register_mma_fp8_minimum_ptx(mma.shape, mma.accumulator); + let a = register_mma_fp8_element_name(mma.a_element).expect("validated A element"); + let b = register_mma_fp8_element_name(mma.b_element).expect("validated B element"); + let ( + scalar, + rust_arguments, + rust_result, + dialect_operands, + dialect_results, + llvm_arguments, + llvm_results, + adapter, + ) = match mma.accumulator { + RegisterMmaAccumulator::F16 => ( + "f16", + vec![ + "[u32; 2]".into(), + format!("[u32; {a_count}]"), + if b_count == 1 { + "u32".into() + } else { + "[u32; 2]".into() + }, + ], + "[u32; 2]".to_owned(), + [ + vec!["i32".to_owned(); 2], + vec!["i32".to_owned(); a_count], + vec!["i32".to_owned(); b_count], + ] + .concat(), + vec!["i32".to_owned(); 2], + [ + vec!["i32".to_owned(); a_count + b_count], + vec!["v2f16".to_owned(); 2], + ] + .concat(), + vec!["v2f16".to_owned(); 2], + if mma.shape == RegisterMmaShape::M16n8k16 { + RegisterMmaAdapter::C2U32A2U32B1U32ToD2U32 + } else { + RegisterMmaAdapter::C2U32A4U32B2U32ToD2U32 + }, + ), + RegisterMmaAccumulator::F32 => ( + "f32", + vec![ + "[f32; 4]".into(), + format!("[u32; {a_count}]"), + if b_count == 1 { + "u32".into() + } else { + "[u32; 2]".into() + }, + ], + "[f32; 4]".to_owned(), + [ + vec!["f32".to_owned(); 4], + vec!["i32".to_owned(); a_count], + vec!["i32".to_owned(); b_count], + ] + .concat(), + vec!["f32".to_owned(); 4], + [ + vec!["i32".to_owned(); a_count + b_count], + vec!["f32".to_owned(); 4], + ] + .concat(), + vec!["f32".to_owned(); 4], + if mma.shape == RegisterMmaShape::M16n8k16 { + RegisterMmaAdapter::C4F32A2U32B1U32ToD4F32 + } else { + RegisterMmaAdapter::C4F32A4U32B2U32ToD4F32 + }, + ), + _ => unreachable!("validated accumulator"), + }; + let expected_id = format!("mma_{shape_name}_fp8_{scalar}_{a}_{b}"); + let expected_source = format!("int_nvvm_mma_{shape_name}_row_col_{scalar}_{a}_{b}_{scalar}"); + let expected_symbol = format!("llvm.nvvm.mma.{shape_name}.row.col.{scalar}.{a}.{b}.{scalar}"); + + ensure!( + policy.id == expected_id + && policy.operation_key + == format!( + "matrix.mma.{shape_name}.row.col.standard_fp8.{scalar}.{a}.{b}.{scalar}" + ) + && policy.source.is_none() + && policy.source_record.as_deref() == Some(expected_source.as_str()) + && policy.llvm_symbol.as_deref() == Some(expected_symbol.as_str()) + && policy.resolved_llvm_symbol.is_none(), + "{} standard FP8 MMA identity changed", + policy.id + ); + ensure!( + policy.rust_module == "matrix" + && policy.rust_name == expected_id + && policy.rust_arguments == rust_arguments + && policy.rust_result == rust_result + && !policy.safe + && policy.must_use + && policy.safe_allowlist_reason.is_none() + && policy.public_rust_path == format!("cuda_intrinsics::matrix::{expected_id}") + && policy.compatibility_rust_paths == [format!("cuda_device::wmma::{expected_id}")], + "{} must preserve its unsafe must-use standard FP8 API", + policy.id + ); + ensure!( + mma.kind == Some(RegisterMmaKind::Standard) + && mma.operation == RegisterMmaOperation::Multiply + && mma.a_layout == RegisterMmaLayout::Row + && mma.b_layout == RegisterMmaLayout::Col + && mma.overflow == RegisterMmaOverflow::NotApplicable + && mma.participation + == RegisterMmaParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes + && mma.adapter == adapter + && mma.compatibility_source == RegisterMmaCompatibilitySource::GeneratedStub + && mma.runtime_validation == RuntimeValidation::Unexecuted + && policy.dialect_op_type == "RegisterMmaOp" + && policy.dialect_op_name == "nvvm.register_mma" + && policy.dialect_operands == dialect_operands + && policy.dialect_results == dialect_results + && policy.llvm_arguments == llvm_arguments + && policy.llvm_results == llvm_results + && policy.ptx_result == rust_result + && policy.lowering == "generated_register_mma", + "{} standard FP8 MMA carrier or lowering changed", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == "none" + && policy.convergent + && policy.execution_scope == "warp" + && policy.minimum_ptx == minimum_ptx + && policy.minimum_sm.as_deref() == Some("sm_89") + && policy.targets == "all", + "{} standard FP8 effects or target floor changed", + policy.id + ); + ensure!( + policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section == "9.7.15.5.14 Multiply-and-Accumulate Instruction: mma" + && policy.ptx_isa_url + == "https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma", + "{} standard FP8 PTX provenance changed", + policy.id + ); + let [selection] = declaration.selections.as_slice() else { + bail!( + "{} must retain exactly one standard FP8 selection", + policy.id + ); + }; + let expected_predicates = if mma.shape == RegisterMmaShape::M16n8k16 { + [ + "Subtarget->getSmVersion() >= 89", + "Subtarget->getPTXVersion() >= 87", + ] + } else { + [ + "Subtarget->getSmVersion() >= 89", + "Subtarget->getPTXVersion() >= 84", + ] + }; + ensure!( + declaration.classes == ["SDPatternOperator", "Intrinsic", "NVVM_MMA"] + && declaration.properties == ["IntrNoCallback", "IntrNoMem"] + && selection_matches_policy(policy, selection)? + && selection.predicates == expected_predicates + && selection.constraints.is_empty(), + "{} imported standard FP8 declaration or selection changed", + policy.id + ); + ensure!( + policy.expected_ptx + == InstructionPattern { + mnemonic: "mma".into(), + modifiers: [ + "sync", "aligned", shape_name, "row", "col", scalar, a, b, scalar + ] + .into_iter() + .map(Into::into) + .collect(), + operands: [ + dialect_results.len(), + a_count, + b_count, + dialect_results.len() + ] + .map(|length| OperandPattern::RegisterList { length }) + .into(), + }, + "{} expected standard FP8 PTX changed", + policy.id + ); + ensure_exact_inline_ptx_backends( + policy, + [ + (IntrinsicBackend::LlvmNvptx, minimum_ptx, Some("sm_89")), + (IntrinsicBackend::LibNvvm, minimum_ptx, Some("sm_89")), + ], + "standard FP8 register MMA", + )?; + ensure_no_other_family_contract(policy, "standard FP8 register MMA")?; + Ok(()) +} + +pub(in crate::resolve) fn validate_register_mma_f8f6f4_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, + mma: &RegisterMma, +) -> Result<()> { + let contract = register_mma_f8f6f4_contract(mma.accumulator)?; + let a = register_mma_f8f6f4_element_name(mma.a_element) + .with_context(|| format!("{} has a non-f8f6f4 A format", policy.id))?; + let b = register_mma_f8f6f4_element_name(mma.b_element) + .with_context(|| format!("{} has a non-f8f6f4 B format", policy.id))?; + let scalar = contract.scalar_name; + let expected_id = format!("mma_m16n8k32_{scalar}_{a}_{b}"); + let expected_source = + format!("int_nvvm_mma_m16n8k32_row_col_kind_f8f6f4_{scalar}_{a}_{b}_{scalar}"); + let expected_symbol = + format!("llvm.nvvm.mma.m16n8k32.row.col.kind.f8f6f4.{scalar}.{a}.{b}.{scalar}"); + ensure!( + policy.id == expected_id + && policy.operation_key + == format!("matrix.mma.m16n8k32.row.col.kind_f8f6f4.{scalar}.{a}.{b}.{scalar}") + && policy.source.is_none() + && policy.source_record.as_deref() == Some(expected_source.as_str()) + && policy.llvm_symbol.as_deref() == Some(expected_symbol.as_str()) + && policy.resolved_llvm_symbol.is_none(), + "{} dense f8f6f4 MMA identity changed", + policy.id + ); + ensure!( + policy.rust_module == "matrix" + && policy.rust_name == expected_id + && policy.rust_arguments == contract.rust_arguments + && policy.rust_result == contract.rust_result + && !policy.safe + && policy.must_use + && policy.safe_allowlist_reason.is_none() + && policy.public_rust_path == format!("cuda_intrinsics::matrix::{expected_id}") + && policy.compatibility_rust_paths == [format!("cuda_device::wmma::{expected_id}")], + "{} must preserve the unsafe must-use dense f8f6f4 API", + policy.id + ); + ensure!( + policy.dialect_op_type == "RegisterMmaOp" + && policy.dialect_op_name == "nvvm.register_mma" + && policy.dialect_operands == contract.dialect_operands + && policy.dialect_results == contract.dialect_results + && policy.llvm_arguments == contract.llvm_arguments + && policy.llvm_results == contract.llvm_results + && policy.ptx_result == contract.rust_result + && mma.shape == RegisterMmaShape::M16n8k32 + && matches!(mma.kind, None | Some(RegisterMmaKind::F8f6f4)) + && mma.operation == RegisterMmaOperation::Multiply + && mma.accumulator == contract.accumulator + && mma.a_layout == RegisterMmaLayout::Row + && mma.b_layout == RegisterMmaLayout::Col + && mma.overflow == RegisterMmaOverflow::NotApplicable + && mma.participation + == RegisterMmaParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes + && mma.adapter == contract.adapter + && mma.compatibility_source == RegisterMmaCompatibilitySource::GeneratedStub + && mma.runtime_validation == RuntimeValidation::Unexecuted + && policy.lowering == "generated_register_mma", + "{} dense f8f6f4 carrier or lowering changed", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == "none" + && policy.convergent + && policy.execution_scope == "warp" + && policy.minimum_ptx == "8.7" + && policy.minimum_sm.is_none() + && policy.targets == REGISTER_MMA_F8F6F4_TARGETS, + "{} dense f8f6f4 effects or exact target set changed", + policy.id + ); + ensure!( + policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section == "9.7.15.5.14 Multiply-and-Accumulate Instruction: mma" + && policy.ptx_isa_url + == "https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma", + "{} dense f8f6f4 PTX provenance changed", + policy.id + ); + let [selection] = declaration.selections.as_slice() else { + bail!( + "{} must retain exactly one imported dense f8f6f4 instruction selection", + policy.id + ); + }; + ensure!( + declaration.classes == ["SDPatternOperator", "Intrinsic", "NVVM_MMA"] + && declaration.properties == ["IntrNoCallback", "IntrNoMem"] + && selection_matches_policy(policy, selection)? + && selection.predicates == ["Subtarget->hasMMABlockScale()"] + && selection.constraints.is_empty(), + "{} imported dense f8f6f4 declaration or selection changed", + policy.id + ); + ensure!( + policy.expected_ptx + == InstructionPattern { + mnemonic: "mma".into(), + modifiers: [ + "sync", + "aligned", + "m16n8k32", + "row", + "col", + "kind::f8f6f4", + scalar, + a, + b, + scalar, + ] + .into_iter() + .map(Into::into) + .collect(), + operands: contract + .ptx_register_counts + .map(|length| OperandPattern::RegisterList { length }) + .into(), + }, + "{} expected dense f8f6f4 PTX changed", + policy.id + ); + let backend_pairs: BTreeSet<_> = policy + .backend_lowerings + .iter() + .map(|lowering| (lowering.backend, lowering.mechanism)) + .collect(); + ensure!( + policy.backend_lowerings.len() == 2 + && backend_pairs + == BTreeSet::from([ + ( + IntrinsicBackend::LlvmNvptx, + BackendLoweringMechanism::InlinePtx, + ), + ( + IntrinsicBackend::LibNvvm, + BackendLoweringMechanism::InlinePtx, + ), + ]) + && policy.backend_lowerings.iter().all(|lowering| { + lowering.targets.is_none() + && lowering.minimum_ptx.is_none() + && lowering.minimum_sm.is_none() + && !lowering.evidence_profile.trim().is_empty() + }), + "{} must inherit the exact reviewed target set on both inline-PTX routes", + policy.id + ); + ensure_no_other_family_contract(policy, "dense f8f6f4 register MMA")?; + Ok(()) +} + +pub(in crate::resolve) fn validate_register_mma_mxf8f6f4_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, + mma: &RegisterMma, +) -> Result<()> { + let a = register_mma_f8f6f4_element_name(mma.a_element) + .with_context(|| format!("{} has a non-mxf8f6f4 A format", policy.id))?; + let b = register_mma_f8f6f4_element_name(mma.b_element) + .with_context(|| format!("{} has a non-mxf8f6f4 B format", policy.id))?; + let expected_id = format!("mma_m16n8k32_mxf8f6f4_f32_{a}_{b}"); + let expected_source = + format!("int_nvvm_mma_block_scale_m16n8k32_row_col_mxf8f6f4_f32_{a}_{b}_f32_ue8m0"); + let expected_symbol = + format!("llvm.nvvm.mma.block.scale.m16n8k32.row.col.mxf8f6f4.f32.{a}.{b}.f32.ue8m0"); + let rust_arguments = [ + "[f32; 4]", "[u32; 4]", "[u32; 2]", "u32", "u16", "u16", "u32", "u16", "u16", + ]; + let dialect_operands = [ + "f32", "f32", "f32", "f32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i16", "i16", + "i32", "i16", "i16", + ]; + let llvm_arguments = [ + "i32", "i32", "i32", "i32", "i32", "i32", "f32", "f32", "f32", "f32", "i32", "i16", "i16", + "i32", "i16", "i16", + ]; + let results = ["f32", "f32", "f32", "f32"]; + ensure!( + policy.id == expected_id + && policy.operation_key + == format!( + "matrix.mma.m16n8k32.row.col.kind_mxf8f6f4.scale_vec_1x.f32.{a}.{b}.f32.ue8m0" + ) + && policy.source.is_none() + && policy.source_record.as_deref() == Some(expected_source.as_str()) + && policy.llvm_symbol.as_deref() == Some(expected_symbol.as_str()) + && policy.resolved_llvm_symbol.is_none(), + "{} dense mxf8f6f4 MMA identity changed", + policy.id + ); + ensure!( + policy.rust_module == "matrix" + && policy.rust_name == expected_id + && policy.rust_arguments == rust_arguments + && policy.rust_result == "[f32; 4]" + && !policy.safe + && policy.must_use + && policy.safe_allowlist_reason.is_none() + && policy.public_rust_path == format!("cuda_intrinsics::matrix::{expected_id}") + && policy.compatibility_rust_paths == [format!("cuda_device::wmma::{expected_id}")], + "{} must preserve the unsafe must-use dense mxf8f6f4 API", + policy.id + ); + ensure!( + policy.dialect_op_type == "RegisterMmaOp" + && policy.dialect_op_name == "nvvm.register_mma" + && policy.dialect_operands == dialect_operands + && policy.dialect_results == results + && policy.llvm_arguments == llvm_arguments + && policy.llvm_results == results + && policy.ptx_result == "[f32; 4]" + && mma.shape == RegisterMmaShape::M16n8k32 + && mma.kind == Some(RegisterMmaKind::Mxf8f6f4) + && mma.operation == RegisterMmaOperation::Multiply + && mma.accumulator == RegisterMmaAccumulator::F32 + && mma.a_layout == RegisterMmaLayout::Row + && mma.b_layout == RegisterMmaLayout::Col + && mma.overflow == RegisterMmaOverflow::NotApplicable + && mma.participation + == RegisterMmaParticipation::AllWarpLanesSameInstructionAndQualifiersNoExitedLanes + && mma.adapter == RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 + && mma.compatibility_source == RegisterMmaCompatibilitySource::GeneratedStub + && mma.runtime_validation == RuntimeValidation::Unexecuted + && policy.lowering == "generated_register_mma", + "{} dense mxf8f6f4 carrier or lowering changed", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == "none" + && policy.convergent + && policy.execution_scope == "warp" + && policy.minimum_ptx == "8.7" + && policy.minimum_sm.is_none() + && policy.targets == REGISTER_MMA_F8F6F4_TARGETS, + "{} dense mxf8f6f4 effects or exact target set changed", + policy.id + ); + ensure!( + policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section == "9.7.15.5.14 Multiply-and-Accumulate Instruction: mma" + && policy.ptx_isa_url + == "https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma", + "{} dense mxf8f6f4 PTX provenance changed", + policy.id + ); + let [selection] = declaration.selections.as_slice() else { + bail!( + "{} must retain exactly one imported dense mxf8f6f4 instruction selection", + policy.id + ); + }; + ensure!( + declaration.classes == ["SDPatternOperator", "Intrinsic", "NVVM_MMA_BLOCK_SCALE"] + && declaration.properties == ["IntrNoCallback", "IntrNoMem"] + && selection_matches_policy(policy, selection)? + && selection.predicates == ["Subtarget->hasMMABlockScale()"] + && selection.constraints.is_empty(), + "{} imported dense mxf8f6f4 declaration or selection changed", + policy.id + ); + ensure!( + policy.expected_ptx + == InstructionPattern { + mnemonic: "mma".into(), + modifiers: [ + "sync", + "aligned", + "m16n8k32", + "row", + "col", + "kind::mxf8f6f4", + "block_scale", + "f32", + a, + b, + "f32", + "ue8m0", + ] + .into_iter() + .map(Into::into) + .collect(), + operands: vec![ + OperandPattern::RegisterList { length: 4 }, + OperandPattern::RegisterList { length: 4 }, + OperandPattern::RegisterList { length: 2 }, + OperandPattern::RegisterList { length: 4 }, + OperandPattern::Register, + OperandPattern::RegisterList { length: 2 }, + OperandPattern::Register, + OperandPattern::RegisterList { length: 2 }, + ], + }, + "{} expected dense mxf8f6f4 PTX changed", + policy.id + ); + let backend_pairs: BTreeSet<_> = policy + .backend_lowerings + .iter() + .map(|lowering| (lowering.backend, lowering.mechanism)) + .collect(); + ensure!( + policy.backend_lowerings.len() == 2 + && backend_pairs + == BTreeSet::from([ + ( + IntrinsicBackend::LlvmNvptx, + BackendLoweringMechanism::InlinePtx, + ), + ( + IntrinsicBackend::LibNvvm, + BackendLoweringMechanism::InlinePtx, + ), + ]) + && policy.backend_lowerings.iter().all(|lowering| { + lowering.targets.is_none() + && lowering.minimum_ptx.is_none() + && lowering.minimum_sm.is_none() + && !lowering.evidence_profile.trim().is_empty() + }), + "{} must inherit the exact reviewed target set on both inline-PTX routes", + policy.id + ); + ensure_no_other_family_contract(policy, "dense mxf8f6f4 register MMA")?; + Ok(()) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/mod.rs b/crates/cuda-intrinsics-gen/src/resolve/families/mod.rs new file mode 100644 index 0000000000..797c0b060d --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/mod.rs @@ -0,0 +1,48 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +pub(super) mod clc; +pub(super) mod cluster; +pub(super) mod cp_async; +pub(super) mod debug_wgmma; +pub(super) mod execution_control; +pub(super) mod matrix_move; +pub(super) mod mbarrier_extended; +pub(super) mod minmax; +pub(super) mod mma; +pub(super) mod packed_alu; +pub(super) mod packed_conversion; +pub(super) mod prmt; +pub(super) mod redux_dotprod; +pub(super) mod scalar_arithmetic; +pub(super) mod scalar_conversion; +pub(super) mod scalar_math; +pub(super) mod special_register; +pub(super) mod tcgen05; +pub(super) mod threadfence; +pub(super) mod tma; +pub(super) mod warp; + +pub(super) use clc::*; +pub(super) use cluster::*; +pub(super) use cp_async::*; +pub(super) use debug_wgmma::*; +pub(super) use execution_control::*; +pub(super) use matrix_move::*; +pub(super) use mbarrier_extended::*; +pub(super) use minmax::*; +pub(super) use mma::*; +pub(super) use packed_alu::*; +pub(super) use packed_conversion::*; +pub(super) use prmt::*; +pub(super) use redux_dotprod::*; +pub(super) use scalar_arithmetic::*; +pub(super) use scalar_conversion::*; +pub(super) use scalar_math::*; +pub(super) use special_register::*; +pub(super) use tcgen05::*; +pub(super) use threadfence::*; +pub(super) use tma::*; +pub(super) use warp::*; diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/packed_alu.rs b/crates/cuda-intrinsics-gen/src/resolve/families/packed_alu.rs new file mode 100644 index 0000000000..91e49c5d5f --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/packed_alu.rs @@ -0,0 +1,1158 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, ImportedIntrinsic, IntrinsicBackend, IntrinsicSource, + OverlayIntrinsic, PackedAluAdapter, PackedAluFormat, PackedAluOperation, + PackedAtomicAccessContract, PackedAtomicAdapter, PackedAtomicAtomicity, + PackedAtomicCodegenContract, PackedAtomicFormat, PackedAtomicOperation, PackedAtomicOrdering, + PackedAtomicPointerContract, PackedAtomicReturnContract, PackedAtomicRounding, + PackedAtomicScope, PackedAtomicScopeContract, PackedAtomicStateSpace, PackedAtomicSubnormal, +}; +use crate::ptx::OperandPattern; +use anyhow::{Context, Result, ensure}; +use std::collections::BTreeSet; + +use crate::resolve::guards::*; + +pub(in crate::resolve) fn validate_packed_atomic_policy( + policy: &OverlayIntrinsic, + source: &IntrinsicSource, +) -> Result<()> { + let packed = policy + .packed_atomic + .as_ref() + .with_context(|| format!("{} has no closed packed-atomic contract", policy.id))?; + ensure!( + packed.operation == PackedAtomicOperation::Add + && packed.state_space == PackedAtomicStateSpace::Global + && packed.ordering == PackedAtomicOrdering::Relaxed + && packed.scope == PackedAtomicScope::Gpu + && packed.rounding == PackedAtomicRounding::NearestEven + && packed.subnormal == PackedAtomicSubnormal::Preserve + && packed.atomicity == PackedAtomicAtomicity::PerElement + && packed.pointer_contract == PackedAtomicPointerContract::MutableGlobalU32Aligned4 + && packed.access_contract + == PackedAtomicAccessContract::NoMixedWholeWordOrNonAtomicAccess + && packed.scope_contract == PackedAtomicScopeContract::RacingAtomicsMutuallyInclusive + && packed.codegen_contract == PackedAtomicCodegenContract::ExactNativeInstruction + && packed.return_contract + == PackedAtomicReturnContract::OldValuesPerElementMayBeNoncoherent + && packed.adapter == PackedAtomicAdapter::OldPackedU32, + "{} requests an unsupported packed-atomic semantic or safety contract", + policy.id + ); + let (format, native_sm, minimum_ptx, minimum_sm, public_name) = match packed.format { + PackedAtomicFormat::F16x2 => ("f16x2", 60, "6.2", "sm_70", "atom_add_f16x2"), + PackedAtomicFormat::Bf16x2 => ("bf16x2", 90, "7.8", "sm_90", "atom_add_bf16x2"), + }; + ensure!( + packed.native_minimum_sm == native_sm, + "{} PTX-native hardware floor does not match the selected packed format", + policy.id + ); + ensure!( + source + == &IntrinsicSource::PtxNative { + instruction: format!("atom.global.add.noftz.{format}"), + }, + "{} PTX-native source does not match its packed format", + policy.id + ); + ensure!( + policy.rust_module == "atomic" + && policy.rust_name == public_name + && policy.rust_arguments == ["*mut u32", "u32"] + && policy.rust_result == "u32" + && !policy.safe + && policy.must_use + && policy.safe_allowlist_reason.is_none() + && policy.public_rust_path == format!("cuda_intrinsics::atomic::{public_name}") + && policy.compatibility_rust_paths == [format!("cuda_device::atomic::{public_name}")], + "{} must preserve the unsafe must-use packed atomic raw/compatibility API", + policy.id + ); + ensure!( + policy.dialect_op_type == "PackedAtomicAddOp" + && policy.dialect_op_name == "nvvm.packed_atomic_add" + && policy.dialect_operands == ["ptr", "i32"] + && policy.dialect_results == ["i32"] + && policy.lowering == "generated_packed_atomic_inline_ptx", + "{} is outside the one closed packed-atomic dialect recipe", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == "read_write" + && !policy.convergent + && policy.execution_scope == "thread" + && policy.minimum_ptx == minimum_ptx + && policy.minimum_sm.as_deref() == Some(minimum_sm) + && policy.targets == "all" + && policy.ptx_result == "u32" + && policy.ldmatrix_variant.is_none() + && policy.ldmatrix_safety.is_none() + && policy.ldmatrix_adapter.is_none() + && policy.redux.is_none() + && policy.vote.is_none() + && policy.active_mask.is_none() + && policy.warp_match.is_none() + && policy.warp_barrier.is_none() + && policy.warp_shuffle.is_none() + && policy.dot_product.is_none() + && policy.packed_alu.is_none() + && policy.packed_conversion.is_none() + && policy.cp_async_copy.is_none() + && policy.cp_async_control.is_none() + && policy.mbarrier_basic.is_none() + && policy.selected_address_space.is_none(), + "{} packed-atomic effects, carrier, or native target floor disagree", + policy.id + ); + ensure!( + policy.expected_ptx.mnemonic == "atom" + && policy.expected_ptx.modifiers == ["global", "add", "noftz", format] + && policy.expected_ptx.operands + == [ + OperandPattern::Register, + OperandPattern::Address, + OperandPattern::Register, + ], + "{} expected PTX must match the exact packed global add spelling", + policy.id + ); + let backend_pairs: BTreeSet<_> = policy + .backend_lowerings + .iter() + .map(|lowering| (lowering.backend, lowering.mechanism)) + .collect(); + ensure!( + backend_pairs + == BTreeSet::from([ + ( + IntrinsicBackend::LlvmNvptx, + BackendLoweringMechanism::InlinePtx, + ), + ( + IntrinsicBackend::LibNvvm, + BackendLoweringMechanism::InlinePtx, + ), + ]), + "{} must define exactly the reviewed LLVM-NVPTX and libNVVM inline-PTX routes", + policy.id + ); + for lowering in &policy.backend_lowerings { + let expected_sm = match (packed.format, lowering.backend) { + (PackedAtomicFormat::F16x2, IntrinsicBackend::LlvmNvptx) => "sm_70", + (PackedAtomicFormat::F16x2, IntrinsicBackend::LibNvvm) => "sm_75", + (PackedAtomicFormat::Bf16x2, _) => "sm_90", + }; + ensure!( + lowering.minimum_ptx.as_deref() == Some(minimum_ptx) + && lowering.minimum_sm.as_deref() == Some(expected_sm) + && !lowering.evidence_profile.trim().is_empty(), + "{} backend {:?} does not carry its exact reviewed profile floor", + policy.id, + lowering.backend + ); + } + Ok(()) +} + +pub(in crate::resolve) enum PackedAluRecipeSource { + Imported { + record: &'static str, + symbol: &'static str, + resolved_symbol: Option<&'static str>, + arguments: &'static [&'static str], + results: &'static [&'static str], + properties: &'static [&'static str], + selection: &'static str, + selection_asm: &'static str, + }, + PtxNative, +} + +pub(in crate::resolve) struct PackedAluRecipe { + pub(in crate::resolve) id: &'static str, + pub(in crate::resolve) abi_id: &'static str, + pub(in crate::resolve) operation_key: &'static str, + pub(in crate::resolve) rust_name: &'static str, + pub(in crate::resolve) dialect_op_type: &'static str, + pub(in crate::resolve) dialect_op_name: &'static str, + pub(in crate::resolve) arity: usize, + pub(in crate::resolve) must_use: bool, + pub(in crate::resolve) ptx_mnemonic: &'static str, + pub(in crate::resolve) modifiers: &'static [&'static str], + pub(in crate::resolve) native_minimum_sm: u16, + pub(in crate::resolve) minimum_ptx: &'static str, + pub(in crate::resolve) minimum_sm: &'static str, + pub(in crate::resolve) ptx_isa_section: &'static str, + pub(in crate::resolve) ptx_isa_url: &'static str, + pub(in crate::resolve) source: PackedAluRecipeSource, +} + +/// Returns the closed recipe for a packed-ALU (format, operation) pair. +/// +/// `None` means the pair is outside the family: the operation exists for some +/// other format but this one has no reviewed lowering for it. +pub(in crate::resolve) fn packed_alu_recipe( + format: PackedAluFormat, + operation: PackedAluOperation, +) -> Option { + match format { + PackedAluFormat::Bf16x2 => packed_bf16x2_alu_recipe(operation), + PackedAluFormat::F16x2 => packed_f16x2_alu_recipe(operation), + PackedAluFormat::F32x2 => packed_f32x2_alu_recipe(operation), + } +} + +pub(in crate::resolve) fn packed_f32x2_alu_recipe( + operation: PackedAluOperation, +) -> Option { + let ( + id, + abi_id, + operation_key, + rust_name, + dialect_op_type, + dialect_op_name, + arity, + head, + modifiers, + section, + url, + ) = match operation { + PackedAluOperation::Add => ( + "add_f32x2", + "i0995", + "packed.alu.f32x2.add", + "add_f32x2", + "AddF32x2Op", + "nvvm.add_f32x2", + 2, + "add.rn.f32x2", + &["rn", "f32x2"][..], + "9.7.3.1 Floating Point Instructions: add", + "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-add", + ), + PackedAluOperation::AddFtz => ( + "add_ftz_f32x2", + "i0996", + "packed.alu.f32x2.add.ftz", + "add_ftz_f32x2", + "AddFtzF32x2Op", + "nvvm.add_ftz_f32x2", + 2, + "add.rn.ftz.f32x2", + &["rn", "ftz", "f32x2"][..], + "9.7.3.1 Floating Point Instructions: add", + "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-add", + ), + PackedAluOperation::Sub => ( + "sub_f32x2", + "i0997", + "packed.alu.f32x2.sub", + "sub_f32x2", + "SubF32x2Op", + "nvvm.sub_f32x2", + 2, + "sub.rn.f32x2", + &["rn", "f32x2"][..], + "9.7.3.2 Floating Point Instructions: sub", + "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-sub", + ), + PackedAluOperation::SubFtz => ( + "sub_ftz_f32x2", + "i0998", + "packed.alu.f32x2.sub.ftz", + "sub_ftz_f32x2", + "SubFtzF32x2Op", + "nvvm.sub_ftz_f32x2", + 2, + "sub.rn.ftz.f32x2", + &["rn", "ftz", "f32x2"][..], + "9.7.3.2 Floating Point Instructions: sub", + "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-sub", + ), + PackedAluOperation::Mul => ( + "mul_f32x2", + "i0999", + "packed.alu.f32x2.mul", + "mul_f32x2", + "MulF32x2Op", + "nvvm.mul_f32x2", + 2, + "mul.rn.f32x2", + &["rn", "f32x2"][..], + "9.7.3.3 Floating Point Instructions: mul", + "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-mul", + ), + PackedAluOperation::MulFtz => ( + "mul_ftz_f32x2", + "i1000", + "packed.alu.f32x2.mul.ftz", + "mul_ftz_f32x2", + "MulFtzF32x2Op", + "nvvm.mul_ftz_f32x2", + 2, + "mul.rn.ftz.f32x2", + &["rn", "ftz", "f32x2"][..], + "9.7.3.3 Floating Point Instructions: mul", + "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-mul", + ), + PackedAluOperation::Fma => ( + "fma_f32x2", + "i1001", + "packed.alu.f32x2.fma", + "fma_f32x2", + "FmaF32x2Op", + "nvvm.fma_f32x2", + 3, + "fma.rn.f32x2", + &["rn", "f32x2"][..], + "9.7.3.4 Floating Point Instructions: fma", + "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-fma", + ), + PackedAluOperation::FmaFtz => ( + "fma_ftz_f32x2", + "i1002", + "packed.alu.f32x2.fma.ftz", + "fma_ftz_f32x2", + "FmaFtzF32x2Op", + "nvvm.fma_ftz_f32x2", + 3, + "fma.rn.ftz.f32x2", + &["rn", "ftz", "f32x2"][..], + "9.7.3.4 Floating Point Instructions: fma", + "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-fma", + ), + _ => return None, + }; + Some(PackedAluRecipe { + id, + abi_id, + operation_key, + rust_name, + dialect_op_type, + dialect_op_name, + arity, + must_use: true, + ptx_mnemonic: head, + modifiers, + native_minimum_sm: 100, + minimum_ptx: "8.6", + minimum_sm: "sm_100", + ptx_isa_section: section, + ptx_isa_url: url, + source: PackedAluRecipeSource::PtxNative, + }) +} + +pub(in crate::resolve) fn packed_bf16x2_alu_recipe( + operation: PackedAluOperation, +) -> Option { + const PURE: &[&str] = &["IntrNoMem", "IntrSpeculatable"]; + const COMMUTATIVE_PURE: &[&str] = &["Commutative", "IntrNoMem", "IntrSpeculatable"]; + Some(match operation { + PackedAluOperation::Fma => PackedAluRecipe { + id: "fma_bf16x2", + abi_id: "i0062", + operation_key: "packed.alu.bf16x2.fma", + rust_name: "fma_bf16x2", + dialect_op_type: "FmaBf16x2Op", + dialect_op_name: "nvvm.fma_bf16x2", + arity: 3, + must_use: false, + ptx_mnemonic: "fma.rn.bf16x2", + modifiers: &["rn", "bf16x2"], + native_minimum_sm: 80, + minimum_ptx: "7.0", + minimum_sm: "sm_80", + ptx_isa_section: "9.7.4.4 Half Precision Floating Point Instructions: fma", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-fma", + source: PackedAluRecipeSource::Imported { + record: "int_nvvm_fma_rn_bf16x2", + symbol: "llvm.nvvm.fma.rn.bf16x2", + resolved_symbol: None, + arguments: &["v2bf16", "v2bf16", "v2bf16"], + results: &["v2bf16"], + properties: PURE, + selection: "INT_NVVM_FMA_rn_bf16x2", + selection_asm: "fma.rn.bf16x2 \t$dst, $src0, $src1, $src2;", + }, + }, + PackedAluOperation::FmaRelu => PackedAluRecipe { + id: "fma_relu_bf16x2", + abi_id: "i0063", + operation_key: "packed.alu.bf16x2.fma.relu", + rust_name: "fma_relu_bf16x2", + dialect_op_type: "FmaReluBf16x2Op", + dialect_op_name: "nvvm.fma_relu_bf16x2", + arity: 3, + must_use: false, + ptx_mnemonic: "fma.rn.relu.bf16x2", + modifiers: &["rn", "relu", "bf16x2"], + native_minimum_sm: 80, + minimum_ptx: "7.0", + minimum_sm: "sm_80", + ptx_isa_section: "9.7.4.4 Half Precision Floating Point Instructions: fma", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-fma", + source: PackedAluRecipeSource::Imported { + record: "int_nvvm_fma_rn_relu_bf16x2", + symbol: "llvm.nvvm.fma.rn.relu.bf16x2", + resolved_symbol: None, + arguments: &["v2bf16", "v2bf16", "v2bf16"], + results: &["v2bf16"], + properties: PURE, + selection: "INT_NVVM_FMA_rn_relu_bf16x2", + selection_asm: "fma.rn.relu.bf16x2 \t$dst, $src0, $src1, $src2;", + }, + }, + PackedAluOperation::Add => PackedAluRecipe { + id: "add_bf16x2", + abi_id: "i0064", + operation_key: "packed.alu.bf16x2.add", + rust_name: "add_bf16x2", + dialect_op_type: "AddBf16x2Op", + dialect_op_name: "nvvm.add_bf16x2", + arity: 2, + must_use: false, + ptx_mnemonic: "add.rn.bf16x2", + modifiers: &["rn", "bf16x2"], + native_minimum_sm: 90, + minimum_ptx: "7.8", + minimum_sm: "sm_90", + ptx_isa_section: "9.7.4.1 Half Precision Floating Point Instructions: add", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-add", + source: PackedAluRecipeSource::PtxNative, + }, + PackedAluOperation::Sub => PackedAluRecipe { + id: "sub_bf16x2", + abi_id: "i0065", + operation_key: "packed.alu.bf16x2.sub", + rust_name: "sub_bf16x2", + dialect_op_type: "SubBf16x2Op", + dialect_op_name: "nvvm.sub_bf16x2", + arity: 2, + must_use: false, + ptx_mnemonic: "sub.rn.bf16x2", + modifiers: &["rn", "bf16x2"], + native_minimum_sm: 90, + minimum_ptx: "7.8", + minimum_sm: "sm_90", + ptx_isa_section: "9.7.4.2 Half Precision Floating Point Instructions: sub", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-sub", + source: PackedAluRecipeSource::PtxNative, + }, + PackedAluOperation::Mul => PackedAluRecipe { + id: "mul_bf16x2", + abi_id: "i0066", + operation_key: "packed.alu.bf16x2.mul", + rust_name: "mul_bf16x2", + dialect_op_type: "MulBf16x2Op", + dialect_op_name: "nvvm.mul_bf16x2", + arity: 2, + must_use: false, + ptx_mnemonic: "mul.rn.bf16x2", + modifiers: &["rn", "bf16x2"], + native_minimum_sm: 90, + minimum_ptx: "7.8", + minimum_sm: "sm_90", + ptx_isa_section: "9.7.4.3 Half Precision Floating Point Instructions: mul", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-mul", + source: PackedAluRecipeSource::PtxNative, + }, + PackedAluOperation::Min => PackedAluRecipe { + id: "min_bf16x2", + abi_id: "i0067", + operation_key: "packed.alu.bf16x2.min", + rust_name: "min_bf16x2", + dialect_op_type: "MinBf16x2Op", + dialect_op_name: "nvvm.min_bf16x2", + arity: 2, + must_use: false, + ptx_mnemonic: "min.bf16x2", + modifiers: &["bf16x2"], + native_minimum_sm: 80, + minimum_ptx: "7.0", + minimum_sm: "sm_80", + ptx_isa_section: "9.7.4.7 Half Precision Floating Point Instructions: min", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-min", + source: PackedAluRecipeSource::Imported { + record: "int_nvvm_fmin_bf16x2", + symbol: "llvm.nvvm.fmin.bf16x2", + resolved_symbol: None, + arguments: &["v2bf16", "v2bf16"], + results: &["v2bf16"], + properties: COMMUTATIVE_PURE, + selection: "INT_NVVM_FMIN_bf16x2", + selection_asm: "min.bf16x2 \t$dst, $src0, $src1;", + }, + }, + PackedAluOperation::Max => PackedAluRecipe { + id: "max_bf16x2", + abi_id: "i0068", + operation_key: "packed.alu.bf16x2.max", + rust_name: "max_bf16x2", + dialect_op_type: "MaxBf16x2Op", + dialect_op_name: "nvvm.max_bf16x2", + arity: 2, + must_use: false, + ptx_mnemonic: "max.bf16x2", + modifiers: &["bf16x2"], + native_minimum_sm: 80, + minimum_ptx: "7.0", + minimum_sm: "sm_80", + ptx_isa_section: "9.7.4.8 Half Precision Floating Point Instructions: max", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-max", + source: PackedAluRecipeSource::Imported { + record: "int_nvvm_fmax_bf16x2", + symbol: "llvm.nvvm.fmax.bf16x2", + resolved_symbol: None, + arguments: &["v2bf16", "v2bf16"], + results: &["v2bf16"], + properties: COMMUTATIVE_PURE, + selection: "INT_NVVM_FMAN_bf16x2", + selection_asm: "max.bf16x2 \t$dst, $src0, $src1;", + }, + }, + PackedAluOperation::Neg => PackedAluRecipe { + id: "neg_bf16x2", + abi_id: "i0069", + operation_key: "packed.alu.bf16x2.neg", + rust_name: "neg_bf16x2", + dialect_op_type: "NegBf16x2Op", + dialect_op_name: "nvvm.neg_bf16x2", + arity: 1, + must_use: false, + ptx_mnemonic: "neg.bf16x2", + modifiers: &["bf16x2"], + native_minimum_sm: 80, + minimum_ptx: "7.0", + minimum_sm: "sm_80", + ptx_isa_section: "9.7.4.5 Half Precision Floating Point Instructions: neg", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-neg", + source: PackedAluRecipeSource::Imported { + record: "int_nvvm_neg_bf16x2", + symbol: "llvm.nvvm.neg.bf16x2", + resolved_symbol: None, + arguments: &["v2bf16"], + results: &["v2bf16"], + properties: PURE, + selection: "INT_NVVM_NEG_BF16X2", + selection_asm: "neg.bf16x2 \t$dst, $src0;", + }, + }, + PackedAluOperation::Abs => PackedAluRecipe { + id: "abs_bf16x2", + abi_id: "i0070", + operation_key: "packed.alu.bf16x2.abs", + rust_name: "abs_bf16x2", + dialect_op_type: "AbsBf16x2Op", + dialect_op_name: "nvvm.abs_bf16x2", + arity: 1, + must_use: false, + ptx_mnemonic: "abs.bf16x2", + modifiers: &["bf16x2"], + native_minimum_sm: 80, + minimum_ptx: "7.0", + minimum_sm: "sm_80", + ptx_isa_section: "9.7.4.6 Half Precision Floating Point Instructions: abs", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-abs", + source: PackedAluRecipeSource::Imported { + record: "int_nvvm_fabs", + symbol: "llvm.nvvm.fabs", + resolved_symbol: Some("llvm.nvvm.fabs.v2bf16"), + arguments: &["anonymous_14"], + results: &["anyfloat"], + properties: PURE, + selection: "ABS_BF16X2", + selection_asm: "abs.bf16x2 \t$dst, $src0;", + }, + }, + // The f32x2-only operation keys are outside this admission. The ftz + // and sat fma forms also do not exist for bf16x2 in the PTX ISA at + // all: ptxas rejects hand-written `fma.rn.ftz.bf16x2` with + // "Illegal modifier '.ftz' for instruction 'fma'", and accepts + // `fma.rn.relu.bf16x2` beside it. LLVM declares the intrinsics anyway, + // so instruction selection is what fails first ("Cannot select: + // intrinsic %llvm.nvvm.fma.rn.ftz.bf16x2"), but the assembler is the + // authority here. Reject them rather than admit a recipe that has no + // instruction to lower to. + PackedAluOperation::AddFtz + | PackedAluOperation::SubFtz + | PackedAluOperation::MulFtz + | PackedAluOperation::FmaFtz + | PackedAluOperation::FmaSat + | PackedAluOperation::FmaFtzSat + | PackedAluOperation::FmaFtzRelu => return None, + }) +} + +pub(in crate::resolve) fn packed_f16x2_alu_recipe( + operation: PackedAluOperation, +) -> Option { + const PURE: &[&str] = &["IntrNoMem", "IntrSpeculatable"]; + const COMMUTATIVE_PURE: &[&str] = &["Commutative", "IntrNoMem", "IntrSpeculatable"]; + Some(match operation { + PackedAluOperation::Fma => PackedAluRecipe { + id: "fma_f16x2", + abi_id: "i0072", + operation_key: "packed.alu.f16x2.fma", + rust_name: "fma_f16x2", + dialect_op_type: "FmaF16x2Op", + dialect_op_name: "nvvm.fma_f16x2", + arity: 3, + must_use: true, + ptx_mnemonic: "fma.rn.f16x2", + modifiers: &["rn", "f16x2"], + native_minimum_sm: 53, + minimum_ptx: "4.2", + minimum_sm: "sm_70", + ptx_isa_section: "9.7.4.4 Half Precision Floating Point Instructions: fma", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-fma", + source: PackedAluRecipeSource::Imported { + record: "int_nvvm_fma_rn_f16x2", + symbol: "llvm.nvvm.fma.rn.f16x2", + resolved_symbol: None, + arguments: &["v2f16", "v2f16", "v2f16"], + results: &["v2f16"], + properties: PURE, + selection: "INT_NVVM_FMA_rn_f16x2", + selection_asm: "fma.rn.f16x2 \t$dst, $src0, $src1, $src2;", + }, + }, + PackedAluOperation::FmaFtz => PackedAluRecipe { + id: "fma_ftz_f16x2", + abi_id: "i0854", + operation_key: "packed.alu.f16x2.fma.ftz", + rust_name: "fma_ftz_f16x2", + dialect_op_type: "FmaFtzF16x2Op", + dialect_op_name: "nvvm.fma_ftz_f16x2", + arity: 3, + must_use: true, + ptx_mnemonic: "fma.rn.ftz.f16x2", + modifiers: &["rn", "ftz", "f16x2"], + native_minimum_sm: 53, + minimum_ptx: "4.2", + minimum_sm: "sm_70", + ptx_isa_section: "9.7.4.4 Half Precision Floating Point Instructions: fma", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-fma", + source: PackedAluRecipeSource::Imported { + record: "int_nvvm_fma_rn_ftz_f16x2", + symbol: "llvm.nvvm.fma.rn.ftz.f16x2", + resolved_symbol: None, + arguments: &["v2f16", "v2f16", "v2f16"], + results: &["v2f16"], + properties: PURE, + selection: "INT_NVVM_FMA_rn_ftz_f16x2", + selection_asm: "fma.rn.ftz.f16x2 \t$dst, $src0, $src1, $src2;", + }, + }, + PackedAluOperation::FmaSat => PackedAluRecipe { + id: "fma_sat_f16x2", + abi_id: "i0855", + operation_key: "packed.alu.f16x2.fma.sat", + rust_name: "fma_sat_f16x2", + dialect_op_type: "FmaSatF16x2Op", + dialect_op_name: "nvvm.fma_sat_f16x2", + arity: 3, + must_use: true, + ptx_mnemonic: "fma.rn.sat.f16x2", + modifiers: &["rn", "sat", "f16x2"], + native_minimum_sm: 53, + minimum_ptx: "4.2", + minimum_sm: "sm_70", + ptx_isa_section: "9.7.4.4 Half Precision Floating Point Instructions: fma", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-fma", + source: PackedAluRecipeSource::Imported { + record: "int_nvvm_fma_rn_sat_f16x2", + symbol: "llvm.nvvm.fma.rn.sat.f16x2", + resolved_symbol: None, + arguments: &["v2f16", "v2f16", "v2f16"], + results: &["v2f16"], + properties: PURE, + selection: "INT_NVVM_FMA_rn_sat_f16x2", + selection_asm: "fma.rn.sat.f16x2 \t$dst, $src0, $src1, $src2;", + }, + }, + PackedAluOperation::FmaFtzSat => PackedAluRecipe { + id: "fma_ftz_sat_f16x2", + abi_id: "i0856", + operation_key: "packed.alu.f16x2.fma.ftz.sat", + rust_name: "fma_ftz_sat_f16x2", + dialect_op_type: "FmaFtzSatF16x2Op", + dialect_op_name: "nvvm.fma_ftz_sat_f16x2", + arity: 3, + must_use: true, + ptx_mnemonic: "fma.rn.ftz.sat.f16x2", + modifiers: &["rn", "ftz", "sat", "f16x2"], + native_minimum_sm: 53, + minimum_ptx: "4.2", + minimum_sm: "sm_70", + ptx_isa_section: "9.7.4.4 Half Precision Floating Point Instructions: fma", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-fma", + source: PackedAluRecipeSource::Imported { + record: "int_nvvm_fma_rn_ftz_sat_f16x2", + symbol: "llvm.nvvm.fma.rn.ftz.sat.f16x2", + resolved_symbol: None, + arguments: &["v2f16", "v2f16", "v2f16"], + results: &["v2f16"], + properties: PURE, + selection: "INT_NVVM_FMA_rn_ftz_sat_f16x2", + selection_asm: "fma.rn.ftz.sat.f16x2 \t$dst, $src0, $src1, $src2;", + }, + }, + PackedAluOperation::FmaRelu => PackedAluRecipe { + id: "fma_relu_f16x2", + abi_id: "i0073", + operation_key: "packed.alu.f16x2.fma.relu", + rust_name: "fma_relu_f16x2", + dialect_op_type: "FmaReluF16x2Op", + dialect_op_name: "nvvm.fma_relu_f16x2", + arity: 3, + must_use: true, + ptx_mnemonic: "fma.rn.relu.f16x2", + modifiers: &["rn", "relu", "f16x2"], + native_minimum_sm: 80, + minimum_ptx: "7.0", + minimum_sm: "sm_80", + ptx_isa_section: "9.7.4.4 Half Precision Floating Point Instructions: fma", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-fma", + source: PackedAluRecipeSource::Imported { + record: "int_nvvm_fma_rn_relu_f16x2", + symbol: "llvm.nvvm.fma.rn.relu.f16x2", + resolved_symbol: None, + arguments: &["v2f16", "v2f16", "v2f16"], + results: &["v2f16"], + properties: PURE, + selection: "INT_NVVM_FMA_rn_relu_f16x2", + selection_asm: "fma.rn.relu.f16x2 \t$dst, $src0, $src1, $src2;", + }, + }, + PackedAluOperation::FmaFtzRelu => PackedAluRecipe { + id: "fma_ftz_relu_f16x2", + abi_id: "i0857", + operation_key: "packed.alu.f16x2.fma.ftz.relu", + rust_name: "fma_ftz_relu_f16x2", + dialect_op_type: "FmaFtzReluF16x2Op", + dialect_op_name: "nvvm.fma_ftz_relu_f16x2", + arity: 3, + must_use: true, + ptx_mnemonic: "fma.rn.ftz.relu.f16x2", + modifiers: &["rn", "ftz", "relu", "f16x2"], + native_minimum_sm: 80, + minimum_ptx: "7.0", + minimum_sm: "sm_80", + ptx_isa_section: "9.7.4.4 Half Precision Floating Point Instructions: fma", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-fma", + source: PackedAluRecipeSource::Imported { + record: "int_nvvm_fma_rn_ftz_relu_f16x2", + symbol: "llvm.nvvm.fma.rn.ftz.relu.f16x2", + resolved_symbol: None, + arguments: &["v2f16", "v2f16", "v2f16"], + results: &["v2f16"], + properties: PURE, + selection: "INT_NVVM_FMA_rn_ftz_relu_f16x2", + selection_asm: "fma.rn.ftz.relu.f16x2 \t$dst, $src0, $src1, $src2;", + }, + }, + PackedAluOperation::Add => PackedAluRecipe { + id: "add_f16x2", + abi_id: "i0074", + operation_key: "packed.alu.f16x2.add", + rust_name: "add_f16x2", + dialect_op_type: "AddF16x2Op", + dialect_op_name: "nvvm.add_f16x2", + arity: 2, + must_use: true, + ptx_mnemonic: "add.rn.f16x2", + modifiers: &["rn", "f16x2"], + native_minimum_sm: 53, + minimum_ptx: "4.2", + minimum_sm: "sm_70", + ptx_isa_section: "9.7.4.1 Half Precision Floating Point Instructions: add", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-add", + source: PackedAluRecipeSource::PtxNative, + }, + PackedAluOperation::Sub => PackedAluRecipe { + id: "sub_f16x2", + abi_id: "i0075", + operation_key: "packed.alu.f16x2.sub", + rust_name: "sub_f16x2", + dialect_op_type: "SubF16x2Op", + dialect_op_name: "nvvm.sub_f16x2", + arity: 2, + must_use: true, + ptx_mnemonic: "sub.rn.f16x2", + modifiers: &["rn", "f16x2"], + native_minimum_sm: 53, + minimum_ptx: "4.2", + minimum_sm: "sm_70", + ptx_isa_section: "9.7.4.2 Half Precision Floating Point Instructions: sub", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-sub", + source: PackedAluRecipeSource::PtxNative, + }, + PackedAluOperation::Mul => PackedAluRecipe { + id: "mul_f16x2", + abi_id: "i0076", + operation_key: "packed.alu.f16x2.mul", + rust_name: "mul_f16x2", + dialect_op_type: "MulF16x2Op", + dialect_op_name: "nvvm.mul_f16x2", + arity: 2, + must_use: true, + ptx_mnemonic: "mul.rn.f16x2", + modifiers: &["rn", "f16x2"], + native_minimum_sm: 53, + minimum_ptx: "4.2", + minimum_sm: "sm_70", + ptx_isa_section: "9.7.4.3 Half Precision Floating Point Instructions: mul", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-mul", + source: PackedAluRecipeSource::PtxNative, + }, + PackedAluOperation::Min => PackedAluRecipe { + id: "min_f16x2", + abi_id: "i0077", + operation_key: "packed.alu.f16x2.min", + rust_name: "min_f16x2", + dialect_op_type: "MinF16x2Op", + dialect_op_name: "nvvm.min_f16x2", + arity: 2, + must_use: true, + ptx_mnemonic: "min.f16x2", + modifiers: &["f16x2"], + native_minimum_sm: 80, + minimum_ptx: "7.0", + minimum_sm: "sm_80", + ptx_isa_section: "9.7.4.7 Half Precision Floating Point Instructions: min", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-min", + source: PackedAluRecipeSource::Imported { + record: "int_nvvm_fmin_f16x2", + symbol: "llvm.nvvm.fmin.f16x2", + resolved_symbol: None, + arguments: &["v2f16", "v2f16"], + results: &["v2f16"], + properties: COMMUTATIVE_PURE, + selection: "INT_NVVM_FMIN_f16x2", + selection_asm: "min.f16x2 \t$dst, $src0, $src1;", + }, + }, + PackedAluOperation::Max => PackedAluRecipe { + id: "max_f16x2", + abi_id: "i0078", + operation_key: "packed.alu.f16x2.max", + rust_name: "max_f16x2", + dialect_op_type: "MaxF16x2Op", + dialect_op_name: "nvvm.max_f16x2", + arity: 2, + must_use: true, + ptx_mnemonic: "max.f16x2", + modifiers: &["f16x2"], + native_minimum_sm: 80, + minimum_ptx: "7.0", + minimum_sm: "sm_80", + ptx_isa_section: "9.7.4.8 Half Precision Floating Point Instructions: max", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-max", + source: PackedAluRecipeSource::Imported { + record: "int_nvvm_fmax_f16x2", + symbol: "llvm.nvvm.fmax.f16x2", + resolved_symbol: None, + arguments: &["v2f16", "v2f16"], + results: &["v2f16"], + properties: COMMUTATIVE_PURE, + selection: "INT_NVVM_FMAN_f16x2", + selection_asm: "max.f16x2 \t$dst, $src0, $src1;", + }, + }, + PackedAluOperation::Neg => PackedAluRecipe { + id: "neg_f16x2", + abi_id: "i0079", + operation_key: "packed.alu.f16x2.neg", + rust_name: "neg_f16x2", + dialect_op_type: "NegF16x2Op", + dialect_op_name: "nvvm.neg_f16x2", + arity: 1, + must_use: true, + ptx_mnemonic: "neg.f16x2", + modifiers: &["f16x2"], + native_minimum_sm: 53, + minimum_ptx: "6.0", + minimum_sm: "sm_70", + ptx_isa_section: "9.7.4.5 Half Precision Floating Point Instructions: neg", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-neg", + source: PackedAluRecipeSource::PtxNative, + }, + PackedAluOperation::Abs => PackedAluRecipe { + id: "abs_f16x2", + abi_id: "i0080", + operation_key: "packed.alu.f16x2.abs", + rust_name: "abs_f16x2", + dialect_op_type: "AbsF16x2Op", + dialect_op_name: "nvvm.abs_f16x2", + arity: 1, + must_use: true, + ptx_mnemonic: "abs.f16x2", + modifiers: &["f16x2"], + native_minimum_sm: 53, + minimum_ptx: "6.5", + minimum_sm: "sm_70", + ptx_isa_section: "9.7.4.6 Half Precision Floating Point Instructions: abs", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-abs", + source: PackedAluRecipeSource::Imported { + record: "int_nvvm_fabs", + symbol: "llvm.nvvm.fabs", + resolved_symbol: Some("llvm.nvvm.fabs.v2f16"), + arguments: &["anonymous_14"], + results: &["anyfloat"], + properties: PURE, + selection: "ABS_F16X2", + selection_asm: "abs.f16x2 \t$dst, $src0;", + }, + }, + // These operations were added for the f32x2 admission. Keeping them + // outside the f16x2 recipe avoids silently expanding this PR's scope, + // even where a similarly spelled PTX form may exist. + PackedAluOperation::AddFtz | PackedAluOperation::SubFtz | PackedAluOperation::MulFtz => { + return None; + } + }) +} + +/// Per-backend PTX and SM floors, which can sit above the recipe's own floor. +/// +/// Takes the caller's already-resolved recipe rather than looking it up again, +/// so this cannot be reached with a pair the family rejects. +pub(in crate::resolve) fn packed_alu_backend_floor( + recipe: &PackedAluRecipe, + format: PackedAluFormat, + operation: PackedAluOperation, + backend: IntrinsicBackend, +) -> (&'static str, &'static str) { + match (format, operation, backend) { + ( + PackedAluFormat::F16x2, + PackedAluOperation::Fma + | PackedAluOperation::Add + | PackedAluOperation::Sub + | PackedAluOperation::Mul, + IntrinsicBackend::LlvmNvptx, + ) => ("6.0", "sm_70"), + ( + PackedAluFormat::F16x2, + PackedAluOperation::Fma + | PackedAluOperation::Add + | PackedAluOperation::Sub + | PackedAluOperation::Mul, + IntrinsicBackend::LibNvvm, + ) => ("4.2", "sm_75"), + // The ftz and sat fma forms share the plain fma floors: PTX 4.2 natively, + // but the pinned LLVM backend needs 6.0 and CUDA 13.3 ptxas no longer + // targets sm_70, so libNVVM is checked at sm_75. + ( + PackedAluFormat::F16x2, + PackedAluOperation::FmaFtz | PackedAluOperation::FmaSat | PackedAluOperation::FmaFtzSat, + IntrinsicBackend::LlvmNvptx, + ) => ("6.0", "sm_70"), + ( + PackedAluFormat::F16x2, + PackedAluOperation::FmaFtz | PackedAluOperation::FmaSat | PackedAluOperation::FmaFtzSat, + IntrinsicBackend::LibNvvm, + ) => ("4.2", "sm_75"), + (PackedAluFormat::F16x2, PackedAluOperation::Neg, IntrinsicBackend::LlvmNvptx) => { + ("6.0", "sm_70") + } + (PackedAluFormat::F16x2, PackedAluOperation::Neg, IntrinsicBackend::LibNvvm) => { + ("6.0", "sm_75") + } + (PackedAluFormat::F16x2, PackedAluOperation::Abs, IntrinsicBackend::LlvmNvptx) => { + ("6.5", "sm_70") + } + (PackedAluFormat::F16x2, PackedAluOperation::Abs, IntrinsicBackend::LibNvvm) => { + ("6.5", "sm_75") + } + _ => (recipe.minimum_ptx, recipe.minimum_sm), + } +} + +pub(in crate::resolve) fn validate_packed_alu_policy( + policy: &OverlayIntrinsic, + source: &IntrinsicSource, + declaration: Option<&ImportedIntrinsic>, +) -> Result<()> { + let packed = policy + .packed_alu + .as_ref() + .with_context(|| format!("{} has no closed packed-ALU contract", policy.id))?; + let recipe = packed_alu_recipe(packed.format, packed.operation) + .with_context(|| format!("{} is outside the closed packed-ALU recipe", policy.id))?; + let (rust_module, rust_type, dialect_type, expected_adapter) = match packed.format { + PackedAluFormat::Bf16x2 => ("bf16x2", "u32", "i32", PackedAluAdapter::DirectPackedU32), + PackedAluFormat::F16x2 => ("f16x2", "u32", "i32", PackedAluAdapter::DirectPackedU32), + PackedAluFormat::F32x2 => ("f32x2", "u64", "i64", PackedAluAdapter::DirectPackedU64), + }; + ensure!( + packed.adapter == expected_adapter, + "{} requests an unsupported packed-ALU adapter", + policy.id + ); + let rust_arguments = vec![rust_type; recipe.arity]; + let dialect_operands = vec![dialect_type; recipe.arity]; + ensure!( + policy.id == recipe.id + && policy.abi_id == recipe.abi_id + && policy.operation_key == recipe.operation_key, + "{} packed-ALU identity does not match its closed operation recipe", + policy.id + ); + ensure!( + policy.rust_module == rust_module + && policy.rust_name == recipe.rust_name + && policy.rust_arguments == rust_arguments + && policy.rust_result == rust_type + && policy.safe + && policy.must_use == recipe.must_use + && policy + .safe_allowlist_reason + .as_deref() + .is_some_and(|reason| !reason.trim().is_empty()) + && policy.public_rust_path + == format!("cuda_intrinsics::{rust_module}::{}", recipe.rust_name) + && policy.compatibility_rust_paths + == [format!("cuda_device::{rust_module}::{}", recipe.rust_name)], + "{} must preserve its reviewed safe packed-ALU API", + policy.id + ); + ensure!( + policy.dialect_op_type == recipe.dialect_op_type + && policy.dialect_op_name == recipe.dialect_op_name + && policy.dialect_operands == dialect_operands + && policy.dialect_results == [dialect_type] + && policy.lowering == "generated_packed_alu_inline_ptx", + "{} is outside the closed packed-ALU dialect and lowering recipe", + policy.id + ); + ensure!( + policy.pure + && policy.memory == "none" + && !policy.convergent + && policy.execution_scope == "thread" + && policy.minimum_ptx == recipe.minimum_ptx + && policy.minimum_sm.as_deref() == Some(recipe.minimum_sm) + && policy.ptx_result == rust_type + && policy.targets == "all" + && packed.native_minimum_sm == recipe.native_minimum_sm, + "{} packed-ALU effects, carrier, or target floor disagree", + policy.id + ); + ensure!( + policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section == recipe.ptx_isa_section + && policy.ptx_isa_url == recipe.ptx_isa_url, + "{} packed-ALU PTX provenance does not match its reviewed instruction section", + policy.id + ); + let expected_operands = vec![OperandPattern::Register; recipe.arity + 1]; + ensure!( + policy.expected_ptx.mnemonic + == recipe.ptx_mnemonic.split('.').next().expect("PTX mnemonic") + && policy.expected_ptx.modifiers == recipe.modifiers + && policy.expected_ptx.operands == expected_operands, + "{} expected PTX does not match its exact packed-ALU instruction", + policy.id + ); + + match &recipe.source { + PackedAluRecipeSource::PtxNative => { + ensure!( + source + == &IntrinsicSource::PtxNative { + instruction: recipe.ptx_mnemonic.to_owned(), + } + && declaration.is_none(), + "{} packed-ALU source does not match its PTX-native recipe", + policy.id + ); + } + PackedAluRecipeSource::Imported { + record, + symbol, + resolved_symbol, + arguments, + results, + properties, + selection, + selection_asm, + } => { + let declaration = declaration.context("imported packed ALU has no declaration")?; + ensure!( + source + == &IntrinsicSource::LlvmImported { + source_record: (*record).to_owned(), + } + && policy.llvm_symbol.as_deref() == Some(*symbol) + && policy.resolved_llvm_symbol.as_deref() == *resolved_symbol + && policy.llvm_arguments == *arguments + && policy.llvm_results == *results, + "{} packed-ALU LLVM source or signature changed", + policy.id + ); + let matching_selections: Vec<_> = declaration + .selections + .iter() + .filter(|candidate| candidate.source_record == *selection) + .collect(); + let expected_selection_count = if *record == "int_nvvm_fabs" { 6 } else { 1 }; + ensure!( + declaration.properties == *properties + && declaration.selections.len() == expected_selection_count + && matching_selections.len() == 1 + && matching_selections[0].asm == *selection_asm + && matching_selections[0].predicates + == [ + format!("Subtarget->getSmVersion() >= {}", recipe.native_minimum_sm), + format!( + "Subtarget->getPTXVersion() >= {}", + recipe.minimum_ptx.replace('.', "") + ), + ] + && matching_selections[0].constraints.is_empty(), + "{} packed-ALU imported properties or selection changed", + policy.id + ); + } + } + let llvm_floor = packed_alu_backend_floor( + &recipe, + packed.format, + packed.operation, + IntrinsicBackend::LlvmNvptx, + ); + let libnvvm_floor = packed_alu_backend_floor( + &recipe, + packed.format, + packed.operation, + IntrinsicBackend::LibNvvm, + ); + ensure_exact_inline_ptx_backends( + policy, + [ + ( + IntrinsicBackend::LlvmNvptx, + llvm_floor.0, + Some(llvm_floor.1), + ), + ( + IntrinsicBackend::LibNvvm, + libnvvm_floor.0, + Some(libnvvm_floor.1), + ), + ], + "packed-ALU", + )?; + ensure_no_other_family_contract(policy, "packed-ALU")?; + Ok(()) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/packed_conversion.rs b/crates/cuda-intrinsics-gen/src/resolve/families/packed_conversion.rs new file mode 100644 index 0000000000..3f42abc9c1 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/packed_conversion.rs @@ -0,0 +1,928 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, ImportedIntrinsic, IntrinsicBackend, IntrinsicSource, + OverlayBackendLowering, OverlayIntrinsic, PackedConversionAdapter, + PackedConversionDestinationFormat, PackedConversionFp8Admission, PackedConversionFp8Direction, + PackedConversionFp8F16x2Admission, PackedConversionFp8Format, PackedConversionRounding, + PackedConversionSaturation, PackedConversionSourceFormat, RuntimeValidation, +}; +use crate::ptx::{InstructionPattern, OperandPattern}; +use anyhow::{Context, Result, ensure}; +use std::collections::BTreeSet; + +use crate::resolve::guards::*; + +pub(in crate::resolve) struct PackedConversionRecipe { + pub(in crate::resolve) id: &'static str, + pub(in crate::resolve) abi_id: &'static str, + pub(in crate::resolve) operation_key: &'static str, + pub(in crate::resolve) rust_name: &'static str, + pub(in crate::resolve) compatibility_path: &'static str, + pub(in crate::resolve) dialect_op_type: &'static str, + pub(in crate::resolve) dialect_op_name: &'static str, + pub(in crate::resolve) source_record: &'static str, + pub(in crate::resolve) llvm_symbol: &'static str, + pub(in crate::resolve) llvm_result: &'static str, + pub(in crate::resolve) summary: &'static str, +} + +pub(in crate::resolve) fn packed_conversion_recipe( + conversion: &crate::model::PackedConversion, +) -> Option { + match conversion.source_format { + PackedConversionSourceFormat::F32x2 => packed_conversion_recipe_f32x2(conversion), + PackedConversionSourceFormat::E4m3x2 + | PackedConversionSourceFormat::E5m2x2 + | PackedConversionSourceFormat::F16x2 => packed_conversion_recipe_fp8_f16x2(conversion), + } +} + +/// Recipes for the packed FP8 conversions whose other side is `f16x2`. +/// +/// Keyed on source as well as destination: unpacking to `f16x2` shares its +/// destination, rounding, and saturation with the scalar-`f32` recipe for +/// `cvt.rn.f16x2.f32`, so the source format is what separates them. +pub(in crate::resolve) fn packed_conversion_recipe_fp8_f16x2( + conversion: &crate::model::PackedConversion, +) -> Option { + match ( + conversion.source_format, + conversion.destination_format, + conversion.rounding, + conversion.saturation, + ) { + ( + PackedConversionSourceFormat::F16x2, + PackedConversionDestinationFormat::E4m3x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::Satfinite, + ) => Some(PackedConversionRecipe { + id: "cvt_rn_satfinite_e4m3x2_f16x2", + abi_id: "i0822", + operation_key: "packed.convert.f16x2.e4m3x2.nearest_even.satfinite", + rust_name: "cvt_rn_satfinite_e4m3x2_f16x2", + compatibility_path: "cuda_device::convert::cvt_rn_satfinite_e4m3x2_f16x2", + dialect_op_type: "CvtRnSatfiniteE4m3x2F16x2Op", + dialect_op_name: "nvvm.cvt_rn_satfinite_e4m3x2_f16x2", + source_record: "int_nvvm_f16x2_to_e4m3x2_rn", + llvm_symbol: "llvm.nvvm.f16x2.to.e4m3x2.rn", + llvm_result: "i16", + summary: "Converts packed f16x2 to packed e4m3x2 with nearest-even finite saturation, preserving half order.", + }), + ( + PackedConversionSourceFormat::F16x2, + PackedConversionDestinationFormat::E4m3x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::SatfiniteRelu, + ) => Some(PackedConversionRecipe { + id: "cvt_rn_satfinite_relu_e4m3x2_f16x2", + abi_id: "i0823", + operation_key: "packed.convert.f16x2.e4m3x2.nearest_even.satfinite.relu", + rust_name: "cvt_rn_satfinite_relu_e4m3x2_f16x2", + compatibility_path: "cuda_device::convert::cvt_rn_satfinite_relu_e4m3x2_f16x2", + dialect_op_type: "CvtRnSatfiniteReluE4m3x2F16x2Op", + dialect_op_name: "nvvm.cvt_rn_satfinite_relu_e4m3x2_f16x2", + source_record: "int_nvvm_f16x2_to_e4m3x2_rn_relu", + llvm_symbol: "llvm.nvvm.f16x2.to.e4m3x2.rn.relu", + llvm_result: "i16", + summary: "Converts packed f16x2 to packed e4m3x2 with nearest-even finite saturation and ReLU, preserving half order.", + }), + ( + PackedConversionSourceFormat::F16x2, + PackedConversionDestinationFormat::E5m2x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::Satfinite, + ) => Some(PackedConversionRecipe { + id: "cvt_rn_satfinite_e5m2x2_f16x2", + abi_id: "i0824", + operation_key: "packed.convert.f16x2.e5m2x2.nearest_even.satfinite", + rust_name: "cvt_rn_satfinite_e5m2x2_f16x2", + compatibility_path: "cuda_device::convert::cvt_rn_satfinite_e5m2x2_f16x2", + dialect_op_type: "CvtRnSatfiniteE5m2x2F16x2Op", + dialect_op_name: "nvvm.cvt_rn_satfinite_e5m2x2_f16x2", + source_record: "int_nvvm_f16x2_to_e5m2x2_rn", + llvm_symbol: "llvm.nvvm.f16x2.to.e5m2x2.rn", + llvm_result: "i16", + summary: "Converts packed f16x2 to packed e5m2x2 with nearest-even finite saturation, preserving half order.", + }), + ( + PackedConversionSourceFormat::F16x2, + PackedConversionDestinationFormat::E5m2x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::SatfiniteRelu, + ) => Some(PackedConversionRecipe { + id: "cvt_rn_satfinite_relu_e5m2x2_f16x2", + abi_id: "i0825", + operation_key: "packed.convert.f16x2.e5m2x2.nearest_even.satfinite.relu", + rust_name: "cvt_rn_satfinite_relu_e5m2x2_f16x2", + compatibility_path: "cuda_device::convert::cvt_rn_satfinite_relu_e5m2x2_f16x2", + dialect_op_type: "CvtRnSatfiniteReluE5m2x2F16x2Op", + dialect_op_name: "nvvm.cvt_rn_satfinite_relu_e5m2x2_f16x2", + source_record: "int_nvvm_f16x2_to_e5m2x2_rn_relu", + llvm_symbol: "llvm.nvvm.f16x2.to.e5m2x2.rn.relu", + llvm_result: "i16", + summary: "Converts packed f16x2 to packed e5m2x2 with nearest-even finite saturation and ReLU, preserving half order.", + }), + ( + PackedConversionSourceFormat::E4m3x2, + PackedConversionDestinationFormat::F16x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::None, + ) => Some(PackedConversionRecipe { + id: "cvt_rn_f16x2_e4m3x2", + abi_id: "i0826", + operation_key: "packed.convert.e4m3x2.f16x2.nearest_even", + rust_name: "cvt_rn_f16x2_e4m3x2", + compatibility_path: "cuda_device::convert::cvt_rn_f16x2_e4m3x2", + dialect_op_type: "CvtRnF16x2E4m3x2Op", + dialect_op_name: "nvvm.cvt_rn_f16x2_e4m3x2", + source_record: "int_nvvm_e4m3x2_to_f16x2_rn", + llvm_symbol: "llvm.nvvm.e4m3x2.to.f16x2.rn", + llvm_result: "v2f16", + summary: "Converts packed e4m3x2 to packed f16x2, preserving byte order.", + }), + ( + PackedConversionSourceFormat::E4m3x2, + PackedConversionDestinationFormat::F16x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::Relu, + ) => Some(PackedConversionRecipe { + id: "cvt_rn_relu_f16x2_e4m3x2", + abi_id: "i0827", + operation_key: "packed.convert.e4m3x2.f16x2.nearest_even.relu", + rust_name: "cvt_rn_relu_f16x2_e4m3x2", + compatibility_path: "cuda_device::convert::cvt_rn_relu_f16x2_e4m3x2", + dialect_op_type: "CvtRnReluF16x2E4m3x2Op", + dialect_op_name: "nvvm.cvt_rn_relu_f16x2_e4m3x2", + source_record: "int_nvvm_e4m3x2_to_f16x2_rn_relu", + llvm_symbol: "llvm.nvvm.e4m3x2.to.f16x2.rn.relu", + llvm_result: "v2f16", + summary: "Converts packed e4m3x2 to packed f16x2 with ReLU, preserving byte order.", + }), + ( + PackedConversionSourceFormat::E5m2x2, + PackedConversionDestinationFormat::F16x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::None, + ) => Some(PackedConversionRecipe { + id: "cvt_rn_f16x2_e5m2x2", + abi_id: "i0828", + operation_key: "packed.convert.e5m2x2.f16x2.nearest_even", + rust_name: "cvt_rn_f16x2_e5m2x2", + compatibility_path: "cuda_device::convert::cvt_rn_f16x2_e5m2x2", + dialect_op_type: "CvtRnF16x2E5m2x2Op", + dialect_op_name: "nvvm.cvt_rn_f16x2_e5m2x2", + source_record: "int_nvvm_e5m2x2_to_f16x2_rn", + llvm_symbol: "llvm.nvvm.e5m2x2.to.f16x2.rn", + llvm_result: "v2f16", + summary: "Converts packed e5m2x2 to packed f16x2, preserving byte order.", + }), + ( + PackedConversionSourceFormat::E5m2x2, + PackedConversionDestinationFormat::F16x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::Relu, + ) => Some(PackedConversionRecipe { + id: "cvt_rn_relu_f16x2_e5m2x2", + abi_id: "i0829", + operation_key: "packed.convert.e5m2x2.f16x2.nearest_even.relu", + rust_name: "cvt_rn_relu_f16x2_e5m2x2", + compatibility_path: "cuda_device::convert::cvt_rn_relu_f16x2_e5m2x2", + dialect_op_type: "CvtRnReluF16x2E5m2x2Op", + dialect_op_name: "nvvm.cvt_rn_relu_f16x2_e5m2x2", + source_record: "int_nvvm_e5m2x2_to_f16x2_rn_relu", + llvm_symbol: "llvm.nvvm.e5m2x2.to.f16x2.rn.relu", + llvm_result: "v2f16", + summary: "Converts packed e5m2x2 to packed f16x2 with ReLU, preserving byte order.", + }), + _ => None, + } +} + +pub(in crate::resolve) fn packed_conversion_recipe_f32x2( + conversion: &crate::model::PackedConversion, +) -> Option { + match ( + conversion.destination_format, + conversion.rounding, + conversion.saturation, + ) { + ( + PackedConversionDestinationFormat::Bf16x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::None, + ) => Some(PackedConversionRecipe { + id: "cvt_f32x2_bf16x2", + abi_id: "i0071", + operation_key: "packed.convert.f32x2.bf16x2.nearest_even", + rust_name: "cvt_f32x2_bf16x2", + compatibility_path: "cuda_device::convert::cvt_bf16x2_f32", + dialect_op_type: "CvtF32x2Bf16x2Op", + dialect_op_name: "nvvm.cvt_f32x2_bf16x2", + source_record: "int_nvvm_ff2bf16x2_rn", + llvm_symbol: "llvm.nvvm.ff2bf16x2.rn", + llvm_result: "v2bf16", + summary: "Converts two f32 values to packed bf16x2 with the first argument in the low half.", + }), + ( + PackedConversionDestinationFormat::F16x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::None, + ) => Some(PackedConversionRecipe { + id: "cvt_f16x2_f32", + abi_id: "i0081", + operation_key: "packed.convert.f32x2.f16x2.nearest_even", + rust_name: "cvt_f16x2_f32", + compatibility_path: "cuda_device::convert::cvt_f16x2_f32", + dialect_op_type: "CvtF16x2F32Op", + dialect_op_name: "nvvm.cvt_f16x2_f32", + source_record: "int_nvvm_ff2f16x2_rn", + llvm_symbol: "llvm.nvvm.ff2f16x2.rn", + llvm_result: "v2f16", + summary: "Converts two f32 values to packed f16x2 with nearest-even rounding and the first argument in the low half.", + }), + ( + PackedConversionDestinationFormat::F16x2, + PackedConversionRounding::TowardZero, + PackedConversionSaturation::None, + ) => Some(PackedConversionRecipe { + id: "cvt_rz_f16x2_f32", + abi_id: "i0082", + operation_key: "packed.convert.f32x2.f16x2.toward_zero", + rust_name: "cvt_rz_f16x2_f32", + compatibility_path: "cuda_device::convert::cvt_rz_f16x2_f32", + dialect_op_type: "CvtRzF16x2F32Op", + dialect_op_name: "nvvm.cvt_rz_f16x2_f32", + source_record: "int_nvvm_ff2f16x2_rz", + llvm_symbol: "llvm.nvvm.ff2f16x2.rz", + llvm_result: "v2f16", + summary: "Converts two f32 values to packed f16x2 with toward-zero rounding and the first argument in the low half.", + }), + ( + PackedConversionDestinationFormat::F16x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::Relu, + ) => Some(PackedConversionRecipe { + id: "cvt_rn_relu_f16x2_f32", + abi_id: "i0083", + operation_key: "packed.convert.f32x2.f16x2.nearest_even.relu", + rust_name: "cvt_rn_relu_f16x2_f32", + compatibility_path: "cuda_device::convert::cvt_rn_relu_f16x2_f32", + dialect_op_type: "CvtRnReluF16x2F32Op", + dialect_op_name: "nvvm.cvt_rn_relu_f16x2_f32", + source_record: "int_nvvm_ff2f16x2_rn_relu", + llvm_symbol: "llvm.nvvm.ff2f16x2.rn.relu", + llvm_result: "v2f16", + summary: "Converts two f32 values to packed f16x2 with nearest-even rounding, ReLU, and the first argument in the low half.", + }), + ( + PackedConversionDestinationFormat::Bf16x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::Relu, + ) => Some(PackedConversionRecipe { + id: "cvt_rn_relu_bf16x2_f32", + abi_id: "i0084", + operation_key: "packed.convert.f32x2.bf16x2.nearest_even.relu", + rust_name: "cvt_rn_relu_bf16x2_f32", + compatibility_path: "cuda_device::convert::cvt_rn_relu_bf16x2_f32", + dialect_op_type: "CvtRnReluBf16x2F32Op", + dialect_op_name: "nvvm.cvt_rn_relu_bf16x2_f32", + source_record: "int_nvvm_ff2bf16x2_rn_relu", + llvm_symbol: "llvm.nvvm.ff2bf16x2.rn.relu", + llvm_result: "v2bf16", + summary: "Converts two f32 values to packed bf16x2 with nearest-even rounding, ReLU, and the first argument in the low half.", + }), + ( + PackedConversionDestinationFormat::Bf16x2, + PackedConversionRounding::TowardZero, + PackedConversionSaturation::None, + ) => Some(PackedConversionRecipe { + id: "cvt_rz_bf16x2_f32", + abi_id: "i0085", + operation_key: "packed.convert.f32x2.bf16x2.toward_zero", + rust_name: "cvt_rz_bf16x2_f32", + compatibility_path: "cuda_device::convert::cvt_rz_bf16x2_f32", + dialect_op_type: "CvtRzBf16x2F32Op", + dialect_op_name: "nvvm.cvt_rz_bf16x2_f32", + source_record: "int_nvvm_ff2bf16x2_rz", + llvm_symbol: "llvm.nvvm.ff2bf16x2.rz", + llvm_result: "v2bf16", + summary: "Converts two f32 values to packed bf16x2 with toward-zero rounding and the first argument in the low half.", + }), + ( + PackedConversionDestinationFormat::E4m3x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::Satfinite, + ) => Some(PackedConversionRecipe { + id: "cvt_rn_satfinite_e4m3x2_f32", + abi_id: "i0259", + operation_key: "packed.convert.f32x2.e4m3x2.nearest_even.satfinite", + rust_name: "cvt_rn_satfinite_e4m3x2_f32", + compatibility_path: "cuda_device::convert::cvt_rn_satfinite_e4m3x2_f32", + dialect_op_type: "CvtRnSatfiniteE4m3x2F32Op", + dialect_op_name: "nvvm.cvt_rn_satfinite_e4m3x2_f32", + source_record: "int_nvvm_ff_to_e4m3x2_rn", + llvm_symbol: "llvm.nvvm.ff.to.e4m3x2.rn", + llvm_result: "i16", + summary: "Converts two f32 values to packed e4m3x2 with nearest-even finite saturation and the first argument in the low byte.", + }), + ( + PackedConversionDestinationFormat::E4m3x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::SatfiniteRelu, + ) => Some(PackedConversionRecipe { + id: "cvt_rn_satfinite_relu_e4m3x2_f32", + abi_id: "i0260", + operation_key: "packed.convert.f32x2.e4m3x2.nearest_even.satfinite.relu", + rust_name: "cvt_rn_satfinite_relu_e4m3x2_f32", + compatibility_path: "cuda_device::convert::cvt_rn_satfinite_relu_e4m3x2_f32", + dialect_op_type: "CvtRnSatfiniteReluE4m3x2F32Op", + dialect_op_name: "nvvm.cvt_rn_satfinite_relu_e4m3x2_f32", + source_record: "int_nvvm_ff_to_e4m3x2_rn_relu", + llvm_symbol: "llvm.nvvm.ff.to.e4m3x2.rn.relu", + llvm_result: "i16", + summary: "Converts two f32 values to packed e4m3x2 with nearest-even finite saturation, ReLU, and the first argument in the low byte.", + }), + ( + PackedConversionDestinationFormat::E5m2x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::Satfinite, + ) => Some(PackedConversionRecipe { + id: "cvt_rn_satfinite_e5m2x2_f32", + abi_id: "i0261", + operation_key: "packed.convert.f32x2.e5m2x2.nearest_even.satfinite", + rust_name: "cvt_rn_satfinite_e5m2x2_f32", + compatibility_path: "cuda_device::convert::cvt_rn_satfinite_e5m2x2_f32", + dialect_op_type: "CvtRnSatfiniteE5m2x2F32Op", + dialect_op_name: "nvvm.cvt_rn_satfinite_e5m2x2_f32", + source_record: "int_nvvm_ff_to_e5m2x2_rn", + llvm_symbol: "llvm.nvvm.ff.to.e5m2x2.rn", + llvm_result: "i16", + summary: "Converts two f32 values to packed e5m2x2 with nearest-even finite saturation and the first argument in the low byte.", + }), + ( + PackedConversionDestinationFormat::E5m2x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::SatfiniteRelu, + ) => Some(PackedConversionRecipe { + id: "cvt_rn_satfinite_relu_e5m2x2_f32", + abi_id: "i0262", + operation_key: "packed.convert.f32x2.e5m2x2.nearest_even.satfinite.relu", + rust_name: "cvt_rn_satfinite_relu_e5m2x2_f32", + compatibility_path: "cuda_device::convert::cvt_rn_satfinite_relu_e5m2x2_f32", + dialect_op_type: "CvtRnSatfiniteReluE5m2x2F32Op", + dialect_op_name: "nvvm.cvt_rn_satfinite_relu_e5m2x2_f32", + source_record: "int_nvvm_ff_to_e5m2x2_rn_relu", + llvm_symbol: "llvm.nvvm.ff.to.e5m2x2.rn.relu", + llvm_result: "i16", + summary: "Converts two f32 values to packed e5m2x2 with nearest-even finite saturation, ReLU, and the first argument in the low byte.", + }), + _ => None, + } +} + +pub(in crate::resolve) fn packed_conversion_ptx_modifiers( + conversion: &crate::model::PackedConversion, +) -> Vec<&'static str> { + let rounding = match conversion.rounding { + PackedConversionRounding::NearestEven => "rn", + PackedConversionRounding::TowardZero => "rz", + }; + let format = match conversion.destination_format { + PackedConversionDestinationFormat::Bf16x2 => "bf16x2", + PackedConversionDestinationFormat::E4m3x2 => "e4m3x2", + PackedConversionDestinationFormat::E5m2x2 => "e5m2x2", + PackedConversionDestinationFormat::F16x2 => "f16x2", + }; + let mut modifiers = vec![rounding]; + match conversion.saturation { + PackedConversionSaturation::None => {} + PackedConversionSaturation::Relu => modifiers.push("relu"), + PackedConversionSaturation::Satfinite => modifiers.push("satfinite"), + PackedConversionSaturation::SatfiniteRelu => modifiers.extend(["satfinite", "relu"]), + } + modifiers.extend([format, conversion.source_format.ptx_token()]); + modifiers +} + +/// Whether either side of the conversion is a packed FP8 format. +pub(in crate::resolve) fn packed_conversion_uses_fp8( + conversion: &crate::model::PackedConversion, +) -> bool { + matches!( + conversion.source_format, + PackedConversionSourceFormat::E4m3x2 | PackedConversionSourceFormat::E5m2x2 + ) || matches!( + conversion.destination_format, + PackedConversionDestinationFormat::E4m3x2 | PackedConversionDestinationFormat::E5m2x2 + ) +} + +/// Source operand types, in Rust, dialect, and LLVM spellings. +/// +/// A packed source arrives in one register; `f32x2` names two scalar operands. +pub(in crate::resolve) fn packed_conversion_source_types( + conversion: &crate::model::PackedConversion, +) -> (Vec, Vec, Vec) { + match conversion.source_format { + PackedConversionSourceFormat::F32x2 => ( + vec!["f32".into(), "f32".into()], + vec!["f32".into(), "f32".into()], + vec!["f32".into(), "f32".into()], + ), + PackedConversionSourceFormat::F16x2 => { + (vec!["u32".into()], vec!["i32".into()], vec!["v2f16".into()]) + } + PackedConversionSourceFormat::E4m3x2 | PackedConversionSourceFormat::E5m2x2 => { + (vec!["u16".into()], vec!["i16".into()], vec!["i16".into()]) + } + } +} + +pub(in crate::resolve) fn packed_conversion_result_width( + conversion: &crate::model::PackedConversion, +) -> u32 { + match conversion.destination_format { + PackedConversionDestinationFormat::Bf16x2 | PackedConversionDestinationFormat::F16x2 => 32, + PackedConversionDestinationFormat::E4m3x2 | PackedConversionDestinationFormat::E5m2x2 => 16, + } +} + +pub(in crate::resolve) fn packed_conversion_floor( + conversion: &crate::model::PackedConversion, +) -> (&'static str, &'static str) { + // FP8 on either side carries the Ada floor, including when FP8 is the + // source and the destination is the older `f16x2`. + if packed_conversion_uses_fp8(conversion) { + return ("8.1", "sm_89"); + } + match conversion.destination_format { + PackedConversionDestinationFormat::Bf16x2 | PackedConversionDestinationFormat::F16x2 => { + ("7.0", "sm_80") + } + PackedConversionDestinationFormat::E4m3x2 | PackedConversionDestinationFormat::E5m2x2 => { + ("8.1", "sm_89") + } + } +} + +pub(in crate::resolve) fn packed_conversion_backend_mechanism( + conversion: &crate::model::PackedConversion, + backend: IntrinsicBackend, +) -> BackendLoweringMechanism { + match (packed_conversion_uses_typed_nvvm(conversion), backend) { + (true, IntrinsicBackend::LlvmNvptx) => BackendLoweringMechanism::TypedNvvm, + _ => BackendLoweringMechanism::InlinePtx, + } +} + +/// Whether the conversion lowers through a typed NVVM intrinsic call. +/// +/// Only the scalar-`f32` pair does. The `f16x2` sources and destinations carry +/// their halves in one integer register on the MIR side, while the typed +/// intrinsics are declared over `<2 x half>`; routing them through inline PTX +/// avoids a packed-half MIR type for no change in emitted code. The recorded +/// evidence assembles both routes to a byte-identical cubin, so the typed route +/// can be switched on later without changing the PTX contract. +pub(in crate::resolve) fn packed_conversion_uses_typed_nvvm( + conversion: &crate::model::PackedConversion, +) -> bool { + conversion.source_format == PackedConversionSourceFormat::F32x2 + && matches!( + conversion.destination_format, + PackedConversionDestinationFormat::E4m3x2 | PackedConversionDestinationFormat::E5m2x2 + ) +} + +pub(in crate::resolve) fn packed_conversion_lowering( + conversion: &crate::model::PackedConversion, +) -> &'static str { + if packed_conversion_uses_typed_nvvm(conversion) { + "generated_packed_conversion_backend" + } else { + "generated_packed_conversion_inline_ptx" + } +} + +pub(in crate::resolve) fn expand_packed_conversion_fp8_admission( + admission: &PackedConversionFp8Admission, +) -> Result> { + ensure!( + admission.runtime_validation == RuntimeValidation::Unexecuted, + "FP8 conversion runtime validation may be marked executed only with GPU evidence" + ); + ensure!( + admission.destination_formats + == [ + PackedConversionDestinationFormat::E4m3x2, + PackedConversionDestinationFormat::E5m2x2, + ], + "compact FP8 conversion admission must list the canonical two formats" + ); + ensure!( + admission.saturations + == [ + PackedConversionSaturation::Satfinite, + PackedConversionSaturation::SatfiniteRelu, + ], + "compact FP8 conversion admission must list base and ReLU finite saturation" + ); + ensure!( + admission.product_count + == admission + .destination_formats + .len() + .checked_mul(admission.saturations.len()) + .context("compact FP8 conversion product count overflow")? + && admission.product_count == 4, + "compact FP8 conversion product_count must be exactly 4" + ); + + let mut records = Vec::with_capacity(admission.product_count); + for &destination_format in &admission.destination_formats { + for &saturation in &admission.saturations { + let conversion = crate::model::PackedConversion { + source_format: PackedConversionSourceFormat::F32x2, + destination_format, + rounding: PackedConversionRounding::NearestEven, + saturation, + adapter: PackedConversionAdapter::ReverseHighLowOperands, + }; + records.push(packed_conversion_overlay_record( + conversion, + &admission.llvm_evidence_profile, + &admission.libnvvm_evidence_profile, + )?); + } + } + ensure!(records.len() == admission.product_count); + Ok(records) +} + +pub(in crate::resolve) fn expand_packed_conversion_fp8_f16x2_admission( + admission: &PackedConversionFp8F16x2Admission, +) -> Result> { + ensure!( + admission.runtime_validation == RuntimeValidation::Unexecuted, + "FP8 f16x2 conversion runtime validation may be marked executed only with GPU evidence" + ); + ensure!( + admission.fp8_formats + == [ + PackedConversionFp8Format::E4m3x2, + PackedConversionFp8Format::E5m2x2, + ], + "compact FP8 f16x2 conversion admission must list the canonical two formats" + ); + ensure!( + admission.directions + == [ + PackedConversionFp8Direction::Pack, + PackedConversionFp8Direction::Unpack, + ], + "compact FP8 f16x2 conversion admission must list both conversion directions" + ); + ensure!( + admission.relu_variants, + "compact FP8 f16x2 conversion admission must admit the ReLU variants" + ); + ensure!( + admission.product_count + == admission + .fp8_formats + .len() + .checked_mul(admission.directions.len()) + .and_then(|count| count.checked_mul(2)) + .context("compact FP8 f16x2 conversion product count overflow")? + && admission.product_count == 8, + "compact FP8 f16x2 conversion product_count must be exactly 8" + ); + + let mut records = Vec::with_capacity(admission.product_count); + for &fp8_format in &admission.fp8_formats { + for &direction in &admission.directions { + for relu in [false, true] { + let conversion = match direction { + // Narrowing to FP8 always saturates to finite. + PackedConversionFp8Direction::Pack => crate::model::PackedConversion { + source_format: PackedConversionSourceFormat::F16x2, + destination_format: match fp8_format { + PackedConversionFp8Format::E4m3x2 => { + PackedConversionDestinationFormat::E4m3x2 + } + PackedConversionFp8Format::E5m2x2 => { + PackedConversionDestinationFormat::E5m2x2 + } + }, + rounding: PackedConversionRounding::NearestEven, + saturation: if relu { + PackedConversionSaturation::SatfiniteRelu + } else { + PackedConversionSaturation::Satfinite + }, + adapter: PackedConversionAdapter::Identity, + }, + // Widening back to f16x2 is exact, so it carries no + // saturation modifier of its own. + PackedConversionFp8Direction::Unpack => crate::model::PackedConversion { + source_format: match fp8_format { + PackedConversionFp8Format::E4m3x2 => { + PackedConversionSourceFormat::E4m3x2 + } + PackedConversionFp8Format::E5m2x2 => { + PackedConversionSourceFormat::E5m2x2 + } + }, + destination_format: PackedConversionDestinationFormat::F16x2, + rounding: PackedConversionRounding::NearestEven, + saturation: if relu { + PackedConversionSaturation::Relu + } else { + PackedConversionSaturation::None + }, + adapter: PackedConversionAdapter::Identity, + }, + }; + records.push(packed_conversion_overlay_record( + conversion, + &admission.llvm_evidence_profile, + &admission.libnvvm_evidence_profile, + )?); + } + } + } + ensure!(records.len() == admission.product_count); + Ok(records) +} + +pub(in crate::resolve) fn packed_conversion_overlay_record( + conversion: crate::model::PackedConversion, + llvm_evidence_profile: &str, + libnvvm_evidence_profile: &str, +) -> Result { + let recipe = packed_conversion_recipe(&conversion) + .context("compact FP8 conversion is outside the closed recipe set")?; + let result_width = packed_conversion_result_width(&conversion); + let rust_result = format!("u{result_width}"); + let dialect_result = format!("i{result_width}"); + let (minimum_ptx, minimum_sm) = packed_conversion_floor(&conversion); + let (rust_arguments, dialect_operands, llvm_arguments) = + packed_conversion_source_types(&conversion); + // `cvt` writes one destination and reads every source operand. + let ptx_operands = conversion + .source_format + .operand_count() + .checked_add(1) + .context("packed-conversion operand count overflow")?; + Ok(OverlayIntrinsic { + id: recipe.id.into(), + abi_id: String::new(), + operation_key: recipe.operation_key.into(), + family: "packed_conversion".into(), + source: None, + source_record: Some(recipe.source_record.into()), + rust_module: "convert".into(), + rust_name: recipe.rust_name.into(), + rust_arguments, + rust_result: rust_result.clone(), + safe: true, + must_use: false, + safe_allowlist_reason: Some("This conversion has no caller obligations.".into()), + public_rust_path: format!("cuda_intrinsics::convert::{}", recipe.rust_name), + compatibility_rust_paths: vec![recipe.compatibility_path.into()], + dialect_op_type: recipe.dialect_op_type.into(), + dialect_op_name: recipe.dialect_op_name.into(), + dialect_operands, + dialect_results: vec![dialect_result], + llvm_symbol: Some(recipe.llvm_symbol.into()), + resolved_llvm_symbol: None, + llvm_arguments, + llvm_results: vec![recipe.llvm_result.into()], + pure: true, + memory: "none".into(), + convergent: false, + execution_scope: "thread".into(), + minimum_ptx: minimum_ptx.into(), + minimum_sm: Some(minimum_sm.into()), + ptx_result: rust_result, + targets: "all".into(), + ptx_isa_version: "9.3".into(), + ptx_isa_section: "9.7.9.22 Data Movement and Conversion Instructions: cvt".into(), + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cvt".into(), + lowering: packed_conversion_lowering(&conversion).into(), + backend_lowerings: [ + (IntrinsicBackend::LlvmNvptx, llvm_evidence_profile), + (IntrinsicBackend::LibNvvm, libnvvm_evidence_profile), + ] + .into_iter() + .map(|(backend, evidence_profile)| OverlayBackendLowering { + backend, + mechanism: packed_conversion_backend_mechanism(&conversion, backend), + evidence_profile: evidence_profile.into(), + targets: None, + minimum_ptx: Some(minimum_ptx.into()), + minimum_sm: Some(minimum_sm.into()), + }) + .collect(), + packed_atomic: None, + redux: None, + vote: None, + active_mask: None, + warp_match: None, + warp_barrier: None, + warp_shuffle: None, + dot_product: None, + packed_alu: None, + integer_minmax: None, + packed_conversion: Some(conversion.clone()), + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cp_async_copy: None, + cp_async_control: None, + cp_async_mbarrier: None, + mbarrier_basic: None, + movmatrix: None, + mbarrier_extended: None, + register_mma: None, + sparse_mma: None, + prmt: None, + cluster_barrier: None, + wgmma_control: None, + special_register: None, + debug_control: None, + cluster_memory: None, + clc: None, + tma: None, + tcgen05: None, + ldmatrix_variant: None, + ldmatrix_safety: None, + ldmatrix_adapter: None, + selected_address_space: None, + expected_ptx: InstructionPattern { + mnemonic: "cvt".into(), + modifiers: packed_conversion_ptx_modifiers(&conversion) + .into_iter() + .map(str::to_owned) + .collect(), + operands: vec![OperandPattern::Register; ptx_operands], + }, + summary: recipe.summary.into(), + }) +} + +pub(in crate::resolve) fn validate_packed_conversion_policy( + policy: &OverlayIntrinsic, + source: &IntrinsicSource, + declaration: Option<&ImportedIntrinsic>, +) -> Result<()> { + let conversion = policy + .packed_conversion + .as_ref() + .with_context(|| format!("{} has no closed packed-conversion contract", policy.id))?; + // The adapter is a function of the source arity: two scalar operands are + // reversed so the first Rust argument lands in the low half, while a single + // packed operand is forwarded unchanged. + let expected_adapter = match conversion.source_format { + PackedConversionSourceFormat::F32x2 => PackedConversionAdapter::ReverseHighLowOperands, + PackedConversionSourceFormat::E4m3x2 + | PackedConversionSourceFormat::E5m2x2 + | PackedConversionSourceFormat::F16x2 => PackedConversionAdapter::Identity, + }; + ensure!( + conversion.adapter == expected_adapter, + "{} requests an unsupported packed-conversion source or adapter", + policy.id + ); + let recipe = packed_conversion_recipe(conversion).with_context(|| { + format!( + "{} requests an unsupported packed-conversion source, destination, rounding, or saturation combination", + policy.id + ) + })?; + let result_width = packed_conversion_result_width(conversion); + let rust_result = format!("u{result_width}"); + let dialect_result = format!("i{result_width}"); + let (minimum_ptx, minimum_sm) = packed_conversion_floor(conversion); + let (expected_rust_arguments, expected_dialect_operands, expected_llvm_arguments) = + packed_conversion_source_types(conversion); + ensure!( + policy.id == recipe.id + && policy.abi_id == recipe.abi_id + && policy.operation_key == recipe.operation_key + && source + == &IntrinsicSource::LlvmImported { + source_record: recipe.source_record.into(), + } + && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) + && policy.resolved_llvm_symbol.is_none() + && policy.llvm_arguments == expected_llvm_arguments + && policy.llvm_results == [recipe.llvm_result], + "{} packed-conversion identity or LLVM source changed", + policy.id + ); + let declaration = declaration.context("packed conversion has no imported declaration")?; + ensure!( + declaration.properties == ["IntrNoMem", "IntrSpeculatable"] + && declaration.selections.is_empty(), + "{} selectionless packed-conversion declaration changed", + policy.id + ); + ensure!( + policy.rust_module == "convert" + && policy.rust_name == recipe.rust_name + && policy.rust_arguments == expected_rust_arguments + && policy.rust_result == rust_result + && policy.safe + && !policy.must_use + && policy + .safe_allowlist_reason + .as_deref() + .is_some_and(|reason| !reason.trim().is_empty()) + && policy.public_rust_path == format!("cuda_intrinsics::convert::{}", recipe.rust_name) + && policy.compatibility_rust_paths == [recipe.compatibility_path], + "{} must preserve its safe non-must-use conversion API", + policy.id + ); + ensure!( + policy.dialect_op_type == recipe.dialect_op_type + && policy.dialect_op_name == recipe.dialect_op_name + && policy.dialect_operands == expected_dialect_operands + && policy.dialect_results == [dialect_result.as_str()] + && policy.lowering == packed_conversion_lowering(conversion), + "{} is outside the closed packed-conversion dialect and lowering recipe", + policy.id + ); + ensure!( + policy.pure + && policy.memory == "none" + && !policy.convergent + && policy.execution_scope == "thread" + && policy.minimum_ptx == minimum_ptx + && policy.minimum_sm.as_deref() == Some(minimum_sm) + && policy.ptx_result == rust_result + && policy.targets == "all" + && policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section == "9.7.9.22 Data Movement and Conversion Instructions: cvt" + && policy.ptx_isa_url + == "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cvt", + "{} packed-conversion effects, carrier, target floor, or PTX provenance disagree", + policy.id + ); + let expected_ptx_operands = vec![ + OperandPattern::Register; + conversion + .source_format + .operand_count() + .checked_add(1) + .context("packed-conversion operand count overflow")? + ]; + ensure!( + policy.expected_ptx.mnemonic == "cvt" + && policy.expected_ptx.modifiers == packed_conversion_ptx_modifiers(conversion) + && policy.expected_ptx.operands == expected_ptx_operands, + "{} expected PTX does not match its closed conversion recipe", + policy.id + ); + ensure!( + policy.summary == recipe.summary, + "{} packed-conversion summary does not match its closed recipe", + policy.id + ); + let backend_pairs = policy + .backend_lowerings + .iter() + .map(|lowering| (lowering.backend, lowering.mechanism)) + .collect::>(); + let expected_pairs = [IntrinsicBackend::LlvmNvptx, IntrinsicBackend::LibNvvm] + .map(|backend| { + ( + backend, + packed_conversion_backend_mechanism(conversion, backend), + ) + }) + .into_iter() + .collect::>(); + ensure!( + policy.backend_lowerings.len() == 2 && backend_pairs == expected_pairs, + "{} must define exactly the reviewed packed-conversion backend routes", + policy.id + ); + for lowering in &policy.backend_lowerings { + ensure!( + lowering.minimum_ptx.as_deref() == Some(minimum_ptx) + && lowering.minimum_sm.as_deref() == Some(minimum_sm) + && !lowering.evidence_profile.trim().is_empty(), + "{} backend {:?} does not carry its exact packed-conversion floor", + policy.id, + lowering.backend + ); + } + ensure_no_other_family_contract(policy, "packed conversion")?; + Ok(()) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/prmt.rs b/crates/cuda-intrinsics-gen/src/resolve/families/prmt.rs new file mode 100644 index 0000000000..2c25bda706 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/prmt.rs @@ -0,0 +1,410 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, ImportedIntrinsic, IntrinsicBackend, OverlayBackendLowering, + OverlayIntrinsic, Prmt, PrmtAdapter, PrmtAdmission, PrmtMode, RuntimeValidation, +}; +use crate::ptx::{InstructionPattern, OperandPattern}; +use anyhow::{Context, Result, ensure}; +use std::collections::BTreeSet; + +use crate::resolve::guards::*; + +#[derive(Clone, Copy)] +pub(in crate::resolve) struct PrmtRecipe { + pub(in crate::resolve) abi_id: &'static str, + pub(in crate::resolve) id: &'static str, + pub(in crate::resolve) operation_key: &'static str, + pub(in crate::resolve) source_record: &'static str, + pub(in crate::resolve) llvm_symbol: &'static str, + pub(in crate::resolve) modifier: Option<&'static str>, + pub(in crate::resolve) adapter: PrmtAdapter, + pub(in crate::resolve) summary: &'static str, +} + +pub(in crate::resolve) fn prmt_recipe(mode: PrmtMode) -> PrmtRecipe { + use PrmtAdapter::{DirectThreeOperands, InsertZeroSecondSource}; + match mode { + PrmtMode::Generic => PrmtRecipe { + abi_id: "i0252", + id: "prmt", + operation_key: "integer.prmt.b32", + source_record: "int_nvvm_prmt", + llvm_symbol: "llvm.nvvm.prmt", + modifier: None, + adapter: DirectThreeOperands, + summary: "Permutes bytes selected from two 32-bit inputs.", + }, + PrmtMode::F4e => PrmtRecipe { + abi_id: "i0253", + id: "prmt_f4e", + operation_key: "integer.prmt.b32.f4e", + source_record: "int_nvvm_prmt_f4e", + llvm_symbol: "llvm.nvvm.prmt.f4e", + modifier: Some("f4e"), + adapter: DirectThreeOperands, + summary: "Permutes bytes with the forward four-byte extract mode.", + }, + PrmtMode::B4e => PrmtRecipe { + abi_id: "i0254", + id: "prmt_b4e", + operation_key: "integer.prmt.b32.b4e", + source_record: "int_nvvm_prmt_b4e", + llvm_symbol: "llvm.nvvm.prmt.b4e", + modifier: Some("b4e"), + adapter: DirectThreeOperands, + summary: "Permutes bytes with the backward four-byte extract mode.", + }, + PrmtMode::Rc8 => PrmtRecipe { + abi_id: "i0255", + id: "prmt_rc8", + operation_key: "integer.prmt.b32.rc8", + source_record: "int_nvvm_prmt_rc8", + llvm_symbol: "llvm.nvvm.prmt.rc8", + modifier: Some("rc8"), + adapter: InsertZeroSecondSource, + summary: "Replicates a selected byte across the 32-bit result.", + }, + PrmtMode::Ecl => PrmtRecipe { + abi_id: "i0256", + id: "prmt_ecl", + operation_key: "integer.prmt.b32.ecl", + source_record: "int_nvvm_prmt_ecl", + llvm_symbol: "llvm.nvvm.prmt.ecl", + modifier: Some("ecl"), + adapter: InsertZeroSecondSource, + summary: "Clamps a byte extract toward the least-significant byte.", + }, + PrmtMode::Ecr => PrmtRecipe { + abi_id: "i0257", + id: "prmt_ecr", + operation_key: "integer.prmt.b32.ecr", + source_record: "int_nvvm_prmt_ecr", + llvm_symbol: "llvm.nvvm.prmt.ecr", + modifier: Some("ecr"), + adapter: InsertZeroSecondSource, + summary: "Clamps a byte extract toward the most-significant byte.", + }, + PrmtMode::Rc16 => PrmtRecipe { + abi_id: "i0258", + id: "prmt_rc16", + operation_key: "integer.prmt.b32.rc16", + source_record: "int_nvvm_prmt_rc16", + llvm_symbol: "llvm.nvvm.prmt.rc16", + modifier: Some("rc16"), + adapter: InsertZeroSecondSource, + summary: "Replicates a selected 16-bit half across the result.", + }, + } +} + +pub(in crate::resolve) fn expand_prmt_admission( + admission: &PrmtAdmission, +) -> Result> { + ensure!( + admission.runtime_validation == RuntimeValidation::Unexecuted, + "prmt runtime validation may be marked executed only with GPU evidence" + ); + ensure!( + !admission.llvm_evidence_profile.trim().is_empty() + && !admission.libnvvm_evidence_profile.trim().is_empty(), + "compact prmt admission requires both backend evidence profiles" + ); + let expected_modes = BTreeSet::from([ + PrmtMode::Generic, + PrmtMode::F4e, + PrmtMode::B4e, + PrmtMode::Rc8, + PrmtMode::Ecl, + PrmtMode::Ecr, + PrmtMode::Rc16, + ]); + let actual_modes: BTreeSet<_> = admission + .variants + .iter() + .map(|variant| variant.mode) + .collect(); + ensure!( + admission.variants.len() == expected_modes.len() && actual_modes == expected_modes, + "compact prmt admission must contain each of the seven reviewed modes exactly once" + ); + + admission + .variants + .iter() + .map(|variant| { + let recipe = prmt_recipe(variant.mode); + ensure!( + variant.abi_id == recipe.abi_id, + "{} must keep reserved ABI ID {}", + recipe.id, + recipe.abi_id + ); + let three_operands = recipe.adapter == PrmtAdapter::DirectThreeOperands; + let rust_arguments = vec!["u32".into(); if three_operands { 3 } else { 2 }]; + let llvm_arguments = vec!["i32".into(); if three_operands { 3 } else { 2 }]; + let mut modifiers = vec!["b32".into()]; + if let Some(modifier) = recipe.modifier { + modifiers.push(modifier.into()); + } + let operands = if three_operands { + vec![ + OperandPattern::Register, + OperandPattern::Register, + OperandPattern::Register, + OperandPattern::Register, + ] + } else { + vec![ + OperandPattern::Register, + OperandPattern::Register, + OperandPattern::Exact { value: "0".into() }, + OperandPattern::Register, + ] + }; + Ok(OverlayIntrinsic { + id: recipe.id.into(), + abi_id: variant.abi_id.clone(), + operation_key: recipe.operation_key.into(), + family: "prmt".into(), + source: None, + source_record: Some(recipe.source_record.into()), + rust_module: "prmt".into(), + rust_name: recipe.id.into(), + rust_arguments, + rust_result: "u32".into(), + safe: true, + must_use: true, + safe_allowlist_reason: Some( + "it only permutes register bytes and has no caller preconditions.".into(), + ), + public_rust_path: format!("cuda_intrinsics::prmt::{}", recipe.id), + compatibility_rust_paths: vec![format!("cuda_device::prmt::{}", recipe.id)], + dialect_op_type: "PrmtOp".into(), + dialect_op_name: "nvvm.prmt".into(), + dialect_operands: llvm_arguments.clone(), + dialect_results: vec!["i32".into()], + llvm_symbol: Some(recipe.llvm_symbol.into()), + resolved_llvm_symbol: None, + llvm_arguments, + llvm_results: vec!["i32".into()], + pure: true, + memory: "none".into(), + convergent: false, + execution_scope: "thread".into(), + minimum_ptx: "2.0".into(), + minimum_sm: Some("sm_20".into()), + ptx_result: "u32".into(), + targets: "all".into(), + ptx_isa_version: "9.3".into(), + ptx_isa_section: "9.7.9.7 Data Movement and Conversion Instructions: prmt".into(), + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-prmt".into(), + lowering: "generated_prmt".into(), + backend_lowerings: vec![ + OverlayBackendLowering { + backend: IntrinsicBackend::LlvmNvptx, + mechanism: BackendLoweringMechanism::TypedNvvm, + evidence_profile: admission.llvm_evidence_profile.clone(), + targets: None, + minimum_ptx: Some("3.2".into()), + minimum_sm: Some("sm_20".into()), + }, + OverlayBackendLowering { + backend: IntrinsicBackend::LibNvvm, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: admission.libnvvm_evidence_profile.clone(), + targets: None, + minimum_ptx: None, + minimum_sm: Some("sm_75".into()), + }, + ], + packed_atomic: None, + redux: None, + vote: None, + active_mask: None, + warp_match: None, + warp_barrier: None, + warp_shuffle: None, + dot_product: None, + packed_alu: None, + integer_minmax: None, + packed_conversion: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cp_async_copy: None, + cp_async_control: None, + cp_async_mbarrier: None, + mbarrier_basic: None, + movmatrix: None, + mbarrier_extended: None, + register_mma: None, + sparse_mma: None, + prmt: Some(Prmt { + mode: variant.mode, + adapter: recipe.adapter, + }), + cluster_barrier: None, + wgmma_control: None, + special_register: None, + debug_control: None, + cluster_memory: None, + clc: None, + tma: None, + tcgen05: None, + ldmatrix_variant: None, + ldmatrix_safety: None, + ldmatrix_adapter: None, + selected_address_space: None, + expected_ptx: InstructionPattern { + mnemonic: "prmt".into(), + modifiers, + operands, + }, + summary: recipe.summary.into(), + }) + }) + .collect() +} + +pub(in crate::resolve) fn validate_prmt_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, +) -> Result<()> { + let prmt = policy + .prmt + .as_ref() + .with_context(|| format!("{} has no closed prmt contract", policy.id))?; + let recipe = prmt_recipe(prmt.mode); + let three_operands = recipe.adapter == PrmtAdapter::DirectThreeOperands; + ensure!( + prmt.adapter == recipe.adapter + && policy.id == recipe.id + && policy.abi_id == recipe.abi_id + && policy.operation_key == recipe.operation_key + && policy.source.is_none() + && policy.source_record.as_deref() == Some(recipe.source_record) + && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) + && policy.resolved_llvm_symbol.is_none(), + "{} prmt identity does not match its closed recipe", + policy.id + ); + ensure!( + policy.rust_module == "prmt" + && policy.rust_name == recipe.id + && policy.rust_arguments == vec!["u32"; if three_operands { 3 } else { 2 }] + && policy.rust_result == "u32" + && policy.safe + && policy.must_use + && policy.public_rust_path == format!("cuda_intrinsics::prmt::{}", recipe.id) + && policy.compatibility_rust_paths == [format!("cuda_device::prmt::{}", recipe.id)], + "{} prmt Rust API does not match its closed recipe", + policy.id + ); + ensure!( + policy.dialect_op_type == "PrmtOp" + && policy.dialect_op_name == "nvvm.prmt" + && policy.dialect_operands == vec!["i32"; if three_operands { 3 } else { 2 }] + && policy.dialect_results == ["i32"] + && policy.llvm_arguments == policy.dialect_operands + && policy.llvm_results == ["i32"] + && policy.lowering == "generated_prmt", + "{} prmt carrier or lowering does not match its closed recipe", + policy.id + ); + ensure!( + declaration + .classes + .iter() + .any(|class| class == "NVVMPureIntrinsic") + && declaration.properties == ["IntrNoMem", "IntrSpeculatable"] + && policy.pure + && policy.memory == "none" + && !policy.convergent + && policy.execution_scope == "thread", + "{} prmt effects disagree with the imported declaration", + policy.id + ); + ensure!( + policy.minimum_ptx == "2.0" + && policy.minimum_sm.as_deref() == Some("sm_20") + && policy.targets == "all" + && policy.ptx_result == "u32" + && policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section == "9.7.9.7 Data Movement and Conversion Instructions: prmt" + && policy.ptx_isa_url + == "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-prmt", + "{} prmt target floor or PTX provenance changed", + policy.id + ); + let mut modifiers = vec!["b32"]; + if let Some(modifier) = recipe.modifier { + modifiers.push(modifier); + } + let expected_operands = if three_operands { + vec![ + OperandPattern::Register, + OperandPattern::Register, + OperandPattern::Register, + OperandPattern::Register, + ] + } else { + vec![ + OperandPattern::Register, + OperandPattern::Register, + OperandPattern::Exact { value: "0".into() }, + OperandPattern::Register, + ] + }; + ensure!( + policy.expected_ptx.mnemonic == "prmt" + && policy.expected_ptx.modifiers == modifiers + && policy.expected_ptx.operands == expected_operands, + "{} expected PTX does not match its closed prmt mode", + policy.id + ); + let backend_pairs: BTreeSet<_> = policy + .backend_lowerings + .iter() + .map(|lowering| (lowering.backend, lowering.mechanism)) + .collect(); + ensure!( + policy.backend_lowerings.len() == 2 + && backend_pairs + == BTreeSet::from([ + ( + IntrinsicBackend::LlvmNvptx, + BackendLoweringMechanism::TypedNvvm, + ), + ( + IntrinsicBackend::LibNvvm, + BackendLoweringMechanism::InlinePtx, + ), + ]), + "{} must define exactly the reviewed prmt backend routes", + policy.id + ); + for lowering in &policy.backend_lowerings { + let floor_matches = match lowering.backend { + IntrinsicBackend::LlvmNvptx => { + lowering.minimum_ptx.as_deref() == Some("3.2") + && lowering.minimum_sm.as_deref() == Some("sm_20") + } + IntrinsicBackend::LibNvvm => { + lowering.minimum_ptx.is_none() && lowering.minimum_sm.as_deref() == Some("sm_75") + } + }; + ensure!( + floor_matches && !lowering.evidence_profile.trim().is_empty(), + "{} backend {:?} does not carry its reviewed prmt floor", + policy.id, + lowering.backend + ); + } + ensure_no_other_family_contract(policy, "prmt")?; + Ok(()) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/redux_dotprod.rs b/crates/cuda-intrinsics-gen/src/resolve/families/redux_dotprod.rs new file mode 100644 index 0000000000..2794b77211 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/redux_dotprod.rs @@ -0,0 +1,627 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, DotProductAdapter, DotProductOperation, DotProductSignedness, + ImportedIntrinsic, IntrinsicBackend, OverlayIntrinsic, ReduxAdapter, ReduxOperation, + ReduxParticipation, +}; +use crate::ptx::OperandPattern; +use anyhow::{Context, Result, ensure}; +use std::collections::BTreeSet; + +pub(in crate::resolve) fn validate_redux_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, +) -> Result<()> { + let redux = policy + .redux + .as_ref() + .with_context(|| format!("{} has no closed redux contract", policy.id))?; + let recipe = redux_recipe(redux.operation); + ensure!( + redux.participation + == ReduxParticipation::ExecutingLaneNamedAllNamedLanesSameInstructionAndMask + && redux.adapter == ReduxAdapter::MaskValueToSourceMemberMask, + "{} requests an unsupported redux participation contract or operand adapter", + policy.id + ); + ensure!( + policy.id == recipe.id + && policy.operation_key == recipe.operation_key + && policy.source_record.as_deref() == Some(recipe.source_record) + && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) + && policy.resolved_llvm_symbol.is_none(), + "{} redux identity does not match its closed operation recipe", + policy.id + ); + ensure!( + policy.rust_module == "warp" + && policy.rust_name == recipe.rust_name + && policy.rust_arguments == ["u32", recipe.rust_value] + && policy.rust_result == recipe.rust_value + && !policy.safe + && policy.must_use + && policy.safe_allowlist_reason.is_none() + && policy.public_rust_path == format!("cuda_intrinsics::warp::{}", recipe.rust_name) + && policy.compatibility_rust_paths + == [format!("cuda_device::warp::{}", recipe.rust_name)], + "{} must preserve the unsafe must-use redux raw API and legacy cuda-device compatibility DefPath", + policy.id + ); + ensure!( + policy.dialect_op_type == recipe.dialect_op_type + && policy.dialect_op_name == recipe.dialect_op_name + && policy.dialect_operands == ["i32", recipe.value_type] + && policy.dialect_results == [recipe.value_type] + && policy.llvm_arguments == [recipe.value_type, "i32"] + && policy.llvm_results == [recipe.value_type] + && policy.lowering == "generated_redux", + "{} is outside the generated two-operand redux recipe", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == "inaccessible_read_write" + && policy.convergent + && policy.execution_scope == "warp" + && policy.minimum_ptx == recipe.minimum_ptx + && policy.minimum_sm.as_deref() == recipe.minimum_sm + && policy.ptx_result == recipe.rust_value + && policy.targets == recipe.targets, + "{} redux effects, carrier, or target floor disagree with its operation recipe", + policy.id + ); + ensure!( + declaration + .properties + .iter() + .any(|property| property == "IntrConvergent") + && declaration + .properties + .iter() + .any(|property| property == "IntrInaccessibleMemOnly") + && declaration + .properties + .iter() + .any(|property| property == "IntrNoCallback") + && !declaration.properties.iter().any(|property| matches!( + property.as_str(), + "IntrNoMem" | "IntrReadMem" | "IntrWriteMem" + )), + "{} redux memory and convergence effects disagree with the imported declaration", + policy.id + ); + ensure!( + policy.backend_lowerings.is_empty() + && policy.packed_atomic.is_none() + && policy.ldmatrix_variant.is_none() + && policy.ldmatrix_safety.is_none() + && policy.ldmatrix_adapter.is_none() + && policy.dot_product.is_none() + && policy.packed_alu.is_none() + && policy.packed_conversion.is_none() + && policy.cp_async_copy.is_none() + && policy.cp_async_control.is_none() + && policy.mbarrier_basic.is_none() + && policy.vote.is_none() + && policy.active_mask.is_none() + && policy.warp_match.is_none() + && policy.warp_barrier.is_none() + && policy.warp_shuffle.is_none() + && policy.selected_address_space.is_none(), + "{} mixes another generated-family contract with redux", + policy.id + ); + ensure!( + policy.expected_ptx.mnemonic == "redux" + && policy.expected_ptx.modifiers == recipe.ptx_modifiers + && policy.expected_ptx.operands + == [ + OperandPattern::Register, + OperandPattern::Register, + OperandPattern::Register, + ], + "{} expected PTX does not match its closed redux operation recipe", + policy.id + ); + Ok(()) +} + +// Pinned LLVM's hasReduxSyncF32(): accel {100} at PTX 8.6, family {100} at +// PTX 8.8. The catalog carries the accel floor with the a/f target union. +pub(in crate::resolve) const REDUX_F32_TARGETS: &str = "sm_100a|sm_100f|sm_103a|sm_103f"; + +pub(in crate::resolve) struct ReduxRecipe { + pub(in crate::resolve) id: &'static str, + pub(in crate::resolve) operation_key: &'static str, + pub(in crate::resolve) source_record: &'static str, + pub(in crate::resolve) llvm_symbol: &'static str, + pub(in crate::resolve) rust_name: &'static str, + pub(in crate::resolve) rust_value: &'static str, + pub(in crate::resolve) value_type: &'static str, + pub(in crate::resolve) dialect_op_type: &'static str, + pub(in crate::resolve) dialect_op_name: &'static str, + pub(in crate::resolve) ptx_modifiers: &'static [&'static str], + pub(in crate::resolve) minimum_ptx: &'static str, + pub(in crate::resolve) minimum_sm: Option<&'static str>, + pub(in crate::resolve) targets: &'static str, +} + +pub(in crate::resolve) fn redux_recipe(operation: ReduxOperation) -> ReduxRecipe { + match operation { + ReduxOperation::Add => ReduxRecipe { + id: "redux_sync_add", + operation_key: "warp.redux.sync.add.wrap32", + source_record: "int_nvvm_redux_sync_add", + llvm_symbol: "llvm.nvvm.redux.sync.add", + rust_name: "redux_sync_add", + rust_value: "u32", + value_type: "i32", + dialect_op_type: "ReduxSyncAddOp", + dialect_op_name: "nvvm.redux_sync_add", + ptx_modifiers: &["sync", "add", "s32"], + minimum_ptx: "7.0", + minimum_sm: Some("sm_80"), + targets: "all", + }, + ReduxOperation::Umin => ReduxRecipe { + id: "redux_sync_min_u32", + operation_key: "warp.redux.sync.min.u32", + source_record: "int_nvvm_redux_sync_umin", + llvm_symbol: "llvm.nvvm.redux.sync.umin", + rust_name: "redux_sync_min_u32", + rust_value: "u32", + value_type: "i32", + dialect_op_type: "ReduxSyncUminOp", + dialect_op_name: "nvvm.redux_sync_umin", + ptx_modifiers: &["sync", "min", "u32"], + minimum_ptx: "7.0", + minimum_sm: Some("sm_80"), + targets: "all", + }, + ReduxOperation::Min => ReduxRecipe { + id: "redux_sync_min_i32", + operation_key: "warp.redux.sync.min.s32", + source_record: "int_nvvm_redux_sync_min", + llvm_symbol: "llvm.nvvm.redux.sync.min", + rust_name: "redux_sync_min_i32", + rust_value: "i32", + value_type: "i32", + dialect_op_type: "ReduxSyncMinOp", + dialect_op_name: "nvvm.redux_sync_min", + ptx_modifiers: &["sync", "min", "s32"], + minimum_ptx: "7.0", + minimum_sm: Some("sm_80"), + targets: "all", + }, + ReduxOperation::Umax => ReduxRecipe { + id: "redux_sync_max_u32", + operation_key: "warp.redux.sync.max.u32", + source_record: "int_nvvm_redux_sync_umax", + llvm_symbol: "llvm.nvvm.redux.sync.umax", + rust_name: "redux_sync_max_u32", + rust_value: "u32", + value_type: "i32", + dialect_op_type: "ReduxSyncUmaxOp", + dialect_op_name: "nvvm.redux_sync_umax", + ptx_modifiers: &["sync", "max", "u32"], + minimum_ptx: "7.0", + minimum_sm: Some("sm_80"), + targets: "all", + }, + ReduxOperation::Max => ReduxRecipe { + id: "redux_sync_max_i32", + operation_key: "warp.redux.sync.max.s32", + source_record: "int_nvvm_redux_sync_max", + llvm_symbol: "llvm.nvvm.redux.sync.max", + rust_name: "redux_sync_max_i32", + rust_value: "i32", + value_type: "i32", + dialect_op_type: "ReduxSyncMaxOp", + dialect_op_name: "nvvm.redux_sync_max", + ptx_modifiers: &["sync", "max", "s32"], + minimum_ptx: "7.0", + minimum_sm: Some("sm_80"), + targets: "all", + }, + ReduxOperation::And => ReduxRecipe { + id: "redux_sync_and", + operation_key: "warp.redux.sync.and.b32", + source_record: "int_nvvm_redux_sync_and", + llvm_symbol: "llvm.nvvm.redux.sync.and", + rust_name: "redux_sync_and", + rust_value: "u32", + value_type: "i32", + dialect_op_type: "ReduxSyncAndOp", + dialect_op_name: "nvvm.redux_sync_and", + ptx_modifiers: &["sync", "and", "b32"], + minimum_ptx: "7.0", + minimum_sm: Some("sm_80"), + targets: "all", + }, + ReduxOperation::Or => ReduxRecipe { + id: "redux_sync_or", + operation_key: "warp.redux.sync.or.b32", + source_record: "int_nvvm_redux_sync_or", + llvm_symbol: "llvm.nvvm.redux.sync.or", + rust_name: "redux_sync_or", + rust_value: "u32", + value_type: "i32", + dialect_op_type: "ReduxSyncOrOp", + dialect_op_name: "nvvm.redux_sync_or", + ptx_modifiers: &["sync", "or", "b32"], + minimum_ptx: "7.0", + minimum_sm: Some("sm_80"), + targets: "all", + }, + ReduxOperation::Xor => ReduxRecipe { + id: "redux_sync_xor", + operation_key: "warp.redux.sync.xor.b32", + source_record: "int_nvvm_redux_sync_xor", + llvm_symbol: "llvm.nvvm.redux.sync.xor", + rust_name: "redux_sync_xor", + rust_value: "u32", + value_type: "i32", + dialect_op_type: "ReduxSyncXorOp", + dialect_op_name: "nvvm.redux_sync_xor", + ptx_modifiers: &["sync", "xor", "b32"], + minimum_ptx: "7.0", + minimum_sm: Some("sm_80"), + targets: "all", + }, + ReduxOperation::Fmin => ReduxRecipe { + id: "redux_sync_min_f32", + operation_key: "warp.redux.sync.min.f32", + source_record: "int_nvvm_redux_sync_fmin", + llvm_symbol: "llvm.nvvm.redux.sync.fmin", + rust_name: "redux_sync_min_f32", + rust_value: "f32", + value_type: "f32", + dialect_op_type: "ReduxSyncFminOp", + dialect_op_name: "nvvm.redux_sync_fmin", + ptx_modifiers: &["sync", "min", "f32"], + minimum_ptx: "8.6", + minimum_sm: None, + targets: REDUX_F32_TARGETS, + }, + ReduxOperation::FminNan => ReduxRecipe { + id: "redux_sync_min_nan_f32", + operation_key: "warp.redux.sync.min.nan.f32", + source_record: "int_nvvm_redux_sync_fmin_NaN", + llvm_symbol: "llvm.nvvm.redux.sync.fmin.NaN", + rust_name: "redux_sync_min_nan_f32", + rust_value: "f32", + value_type: "f32", + dialect_op_type: "ReduxSyncFminNanOp", + dialect_op_name: "nvvm.redux_sync_fmin_nan", + ptx_modifiers: &["sync", "min", "NaN", "f32"], + minimum_ptx: "8.6", + minimum_sm: None, + targets: REDUX_F32_TARGETS, + }, + ReduxOperation::FminAbs => ReduxRecipe { + id: "redux_sync_min_abs_f32", + operation_key: "warp.redux.sync.min.abs.f32", + source_record: "int_nvvm_redux_sync_fmin_abs", + llvm_symbol: "llvm.nvvm.redux.sync.fmin.abs", + rust_name: "redux_sync_min_abs_f32", + rust_value: "f32", + value_type: "f32", + dialect_op_type: "ReduxSyncFminAbsOp", + dialect_op_name: "nvvm.redux_sync_fmin_abs", + ptx_modifiers: &["sync", "min", "abs", "f32"], + minimum_ptx: "8.6", + minimum_sm: None, + targets: REDUX_F32_TARGETS, + }, + ReduxOperation::FminAbsNan => ReduxRecipe { + id: "redux_sync_min_abs_nan_f32", + operation_key: "warp.redux.sync.min.abs.nan.f32", + source_record: "int_nvvm_redux_sync_fmin_abs_NaN", + llvm_symbol: "llvm.nvvm.redux.sync.fmin.abs.NaN", + rust_name: "redux_sync_min_abs_nan_f32", + rust_value: "f32", + value_type: "f32", + dialect_op_type: "ReduxSyncFminAbsNanOp", + dialect_op_name: "nvvm.redux_sync_fmin_abs_nan", + ptx_modifiers: &["sync", "min", "abs", "NaN", "f32"], + minimum_ptx: "8.6", + minimum_sm: None, + targets: REDUX_F32_TARGETS, + }, + ReduxOperation::Fmax => ReduxRecipe { + id: "redux_sync_max_f32", + operation_key: "warp.redux.sync.max.f32", + source_record: "int_nvvm_redux_sync_fmax", + llvm_symbol: "llvm.nvvm.redux.sync.fmax", + rust_name: "redux_sync_max_f32", + rust_value: "f32", + value_type: "f32", + dialect_op_type: "ReduxSyncFmaxOp", + dialect_op_name: "nvvm.redux_sync_fmax", + ptx_modifiers: &["sync", "max", "f32"], + minimum_ptx: "8.6", + minimum_sm: None, + targets: REDUX_F32_TARGETS, + }, + ReduxOperation::FmaxNan => ReduxRecipe { + id: "redux_sync_max_nan_f32", + operation_key: "warp.redux.sync.max.nan.f32", + source_record: "int_nvvm_redux_sync_fmax_NaN", + llvm_symbol: "llvm.nvvm.redux.sync.fmax.NaN", + rust_name: "redux_sync_max_nan_f32", + rust_value: "f32", + value_type: "f32", + dialect_op_type: "ReduxSyncFmaxNanOp", + dialect_op_name: "nvvm.redux_sync_fmax_nan", + ptx_modifiers: &["sync", "max", "NaN", "f32"], + minimum_ptx: "8.6", + minimum_sm: None, + targets: REDUX_F32_TARGETS, + }, + ReduxOperation::FmaxAbs => ReduxRecipe { + id: "redux_sync_max_abs_f32", + operation_key: "warp.redux.sync.max.abs.f32", + source_record: "int_nvvm_redux_sync_fmax_abs", + llvm_symbol: "llvm.nvvm.redux.sync.fmax.abs", + rust_name: "redux_sync_max_abs_f32", + rust_value: "f32", + value_type: "f32", + dialect_op_type: "ReduxSyncFmaxAbsOp", + dialect_op_name: "nvvm.redux_sync_fmax_abs", + ptx_modifiers: &["sync", "max", "abs", "f32"], + minimum_ptx: "8.6", + minimum_sm: None, + targets: REDUX_F32_TARGETS, + }, + ReduxOperation::FmaxAbsNan => ReduxRecipe { + id: "redux_sync_max_abs_nan_f32", + operation_key: "warp.redux.sync.max.abs.nan.f32", + source_record: "int_nvvm_redux_sync_fmax_abs_NaN", + llvm_symbol: "llvm.nvvm.redux.sync.fmax.abs.NaN", + rust_name: "redux_sync_max_abs_nan_f32", + rust_value: "f32", + value_type: "f32", + dialect_op_type: "ReduxSyncFmaxAbsNanOp", + dialect_op_name: "nvvm.redux_sync_fmax_abs_nan", + ptx_modifiers: &["sync", "max", "abs", "NaN", "f32"], + minimum_ptx: "8.6", + minimum_sm: None, + targets: REDUX_F32_TARGETS, + }, + } +} + +pub(in crate::resolve) fn validate_dot_product_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, +) -> Result<()> { + let dot_product = policy + .dot_product + .as_ref() + .with_context(|| format!("{} has no closed dot-product contract", policy.id))?; + let recipe = dot_product_recipe(dot_product.operation, dot_product.signedness); + ensure!( + dot_product.adapter == recipe.adapter, + "{} dot-product source adapter does not match its operation", + policy.id + ); + ensure!( + policy.id == recipe.id + && policy.operation_key == recipe.operation_key + && policy.source.is_none() + && policy.source_record.as_deref() == Some(recipe.source_record) + && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) + && policy.resolved_llvm_symbol.is_none(), + "{} dot-product identity does not match its closed operation recipe", + policy.id + ); + ensure!( + policy.rust_module == "dotprod" + && policy.rust_name == recipe.rust_name + && policy.rust_arguments == ["u32", "u32", recipe.rust_value] + && policy.rust_result == recipe.rust_value + && policy.safe + && !policy.must_use + && policy.public_rust_path == format!("cuda_intrinsics::dotprod::{}", recipe.rust_name) + && policy.compatibility_rust_paths + == [format!("cuda_device::dotprod::{}", recipe.rust_name)], + "{} must preserve the safe, non-must-use dotprod raw API and legacy cuda-device DefPath", + policy.id + ); + ensure!( + policy.dialect_op_type == recipe.dialect_op_type + && policy.dialect_op_name == recipe.dialect_op_name + && policy.dialect_operands == ["i32", "i32", "i32"] + && policy.dialect_results == ["i32"] + && policy.llvm_arguments == recipe.llvm_arguments + && policy.llvm_results == ["i32"] + && policy.lowering == "generated_dotprod", + "{} is outside the closed three-operand dot-product lowering recipe", + policy.id + ); + ensure!( + policy.pure + && policy.memory == "none" + && !policy.convergent + && policy.execution_scope == "thread" + && policy.minimum_ptx == "5.0" + && policy.minimum_sm.as_deref() == Some("sm_61") + && policy.ptx_result == recipe.rust_value + && policy.targets == "all", + "{} dot-product effects, carrier, or target floor disagree with its operation recipe", + policy.id + ); + ensure!( + declaration + .classes + .iter() + .any(|class| class == "NVVMPureIntrinsic") + && declaration.properties == recipe.llvm_properties, + "{} dot-product effects or immediate contract disagree with the imported declaration", + policy.id + ); + ensure!( + policy.ldmatrix_variant.is_none() + && policy.ldmatrix_safety.is_none() + && policy.ldmatrix_adapter.is_none() + && policy.packed_atomic.is_none() + && policy.redux.is_none() + && policy.vote.is_none() + && policy.active_mask.is_none() + && policy.warp_match.is_none() + && policy.warp_barrier.is_none() + && policy.warp_shuffle.is_none() + && policy.packed_alu.is_none() + && policy.packed_conversion.is_none() + && policy.cp_async_copy.is_none() + && policy.cp_async_control.is_none() + && policy.mbarrier_basic.is_none() + && policy.selected_address_space.is_none(), + "{} mixes another generated-family contract with dotprod", + policy.id + ); + ensure!( + policy.expected_ptx.mnemonic == recipe.ptx_mnemonic + && policy.expected_ptx.modifiers == recipe.ptx_modifiers + && policy.expected_ptx.operands + == [ + OperandPattern::Register, + OperandPattern::Register, + OperandPattern::Register, + OperandPattern::Register, + ], + "{} expected PTX does not match its closed dot-product recipe", + policy.id + ); + + let backend_pairs: BTreeSet<_> = policy + .backend_lowerings + .iter() + .map(|lowering| (lowering.backend, lowering.mechanism)) + .collect(); + ensure!( + policy.backend_lowerings.len() == 2 + && backend_pairs + == BTreeSet::from([ + ( + IntrinsicBackend::LlvmNvptx, + BackendLoweringMechanism::TypedNvvm, + ), + ( + IntrinsicBackend::LibNvvm, + BackendLoweringMechanism::InlinePtx, + ), + ]), + "{} must define exactly the reviewed LLVM typed and libNVVM inline-PTX routes", + policy.id + ); + for lowering in &policy.backend_lowerings { + let floor_matches = match lowering.backend { + IntrinsicBackend::LlvmNvptx => { + lowering.mechanism == BackendLoweringMechanism::TypedNvvm + && lowering.minimum_ptx.is_none() + && lowering.minimum_sm.is_none() + } + IntrinsicBackend::LibNvvm => { + lowering.mechanism == BackendLoweringMechanism::InlinePtx + && lowering.minimum_ptx.is_none() + && lowering.minimum_sm.as_deref() == Some("sm_75") + } + }; + ensure!( + floor_matches && !lowering.evidence_profile.trim().is_empty(), + "{} backend {:?} does not carry its reviewed dot-product profile floor", + policy.id, + lowering.backend + ); + } + Ok(()) +} + +pub(in crate::resolve) struct DotProductRecipe { + pub(in crate::resolve) id: &'static str, + pub(in crate::resolve) operation_key: &'static str, + pub(in crate::resolve) source_record: &'static str, + pub(in crate::resolve) llvm_symbol: &'static str, + pub(in crate::resolve) rust_name: &'static str, + pub(in crate::resolve) rust_value: &'static str, + pub(in crate::resolve) dialect_op_type: &'static str, + pub(in crate::resolve) dialect_op_name: &'static str, + pub(in crate::resolve) llvm_arguments: &'static [&'static str], + pub(in crate::resolve) llvm_properties: &'static [&'static str], + pub(in crate::resolve) adapter: DotProductAdapter, + pub(in crate::resolve) ptx_mnemonic: &'static str, + pub(in crate::resolve) ptx_modifiers: &'static [&'static str], +} + +pub(in crate::resolve) fn dot_product_recipe( + operation: DotProductOperation, + signedness: DotProductSignedness, +) -> DotProductRecipe { + match (operation, signedness) { + (DotProductOperation::Dp4a, DotProductSignedness::Signed) => DotProductRecipe { + id: "dp4a_s32", + operation_key: "integer.dot_product.dp4a.s32", + source_record: "int_nvvm_idp4a_s_s", + llvm_symbol: "llvm.nvvm.idp4a.s.s", + rust_name: "dp4a_s32", + rust_value: "i32", + dialect_op_type: "Dp4aS32Op", + dialect_op_name: "nvvm.dp4a_s32", + llvm_arguments: &["i32", "i32", "i32"], + llvm_properties: &["IntrNoMem", "IntrSpeculatable"], + adapter: DotProductAdapter::DirectThreeOperands, + ptx_mnemonic: "dp4a", + ptx_modifiers: &["s32", "s32"], + }, + (DotProductOperation::Dp4a, DotProductSignedness::Unsigned) => DotProductRecipe { + id: "dp4a_u32", + operation_key: "integer.dot_product.dp4a.u32", + source_record: "int_nvvm_idp4a_u_u", + llvm_symbol: "llvm.nvvm.idp4a.u.u", + rust_name: "dp4a_u32", + rust_value: "u32", + dialect_op_type: "Dp4aU32Op", + dialect_op_name: "nvvm.dp4a_u32", + llvm_arguments: &["i32", "i32", "i32"], + llvm_properties: &["IntrNoMem", "IntrSpeculatable"], + adapter: DotProductAdapter::DirectThreeOperands, + ptx_mnemonic: "dp4a", + ptx_modifiers: &["u32", "u32"], + }, + (DotProductOperation::Dp2a, DotProductSignedness::Signed) => DotProductRecipe { + id: "dp2a_s32", + operation_key: "integer.dot_product.dp2a.lo.s32", + source_record: "int_nvvm_idp2a_s_s", + llvm_symbol: "llvm.nvvm.idp2a.s.s", + rust_name: "dp2a_s32", + rust_value: "i32", + dialect_op_type: "Dp2aS32Op", + dialect_op_name: "nvvm.dp2a_s32", + llvm_arguments: &["i32", "i32", "i1", "i32"], + llvm_properties: &["ImmArg", "IntrNoMem", "IntrSpeculatable"], + adapter: DotProductAdapter::InsertLowHalfFalse, + ptx_mnemonic: "dp2a", + ptx_modifiers: &["lo", "s32", "s32"], + }, + (DotProductOperation::Dp2a, DotProductSignedness::Unsigned) => DotProductRecipe { + id: "dp2a_u32", + operation_key: "integer.dot_product.dp2a.lo.u32", + source_record: "int_nvvm_idp2a_u_u", + llvm_symbol: "llvm.nvvm.idp2a.u.u", + rust_name: "dp2a_u32", + rust_value: "u32", + dialect_op_type: "Dp2aU32Op", + dialect_op_name: "nvvm.dp2a_u32", + llvm_arguments: &["i32", "i32", "i1", "i32"], + llvm_properties: &["ImmArg", "IntrNoMem", "IntrSpeculatable"], + adapter: DotProductAdapter::InsertLowHalfFalse, + ptx_mnemonic: "dp2a", + ptx_modifiers: &["lo", "u32", "u32"], + }, + } +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/scalar_arithmetic.rs b/crates/cuda-intrinsics-gen/src/resolve/families/scalar_arithmetic.rs new file mode 100644 index 0000000000..17f39ba0bb --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/scalar_arithmetic.rs @@ -0,0 +1,661 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, ImportedIntrinsic, IntrinsicBackend, OverlayBackendLowering, + OverlayIntrinsic, RuntimeValidation, ScalarArithmetic, ScalarArithmeticAdmission, + ScalarArithmeticFormat, ScalarArithmeticOperation, ScalarArithmeticRounding, + ScalarArithmeticSaturation, ScalarArithmeticSubnormal, +}; +use crate::ptx::{InstructionPattern, OperandPattern}; +use anyhow::{Context, Result, ensure}; + +use crate::resolve::abi_ledger::*; +use crate::resolve::guards::*; +use crate::resolve::targets::*; + +#[derive(Clone)] +pub(in crate::resolve) struct ScalarArithmeticRecipe { + id: String, + operation_key: String, + source_record: String, + llvm_symbol: String, + selection_record: String, + selection_asm: String, + rust_type: &'static str, + argument_count: usize, + properties: Vec<&'static str>, + ptx_modifiers: Vec, + ptx_isa_section: &'static str, + ptx_isa_url: &'static str, +} + +pub(in crate::resolve) type ScalarArithmeticVariant = ( + ScalarArithmeticFormat, + ScalarArithmeticOperation, + ScalarArithmeticRounding, + ScalarArithmeticSubnormal, + ScalarArithmeticSaturation, +); + +pub(in crate::resolve) fn canonical_scalar_arithmetic_variants() -> Vec { + use ScalarArithmeticFormat::{F32, F64}; + use ScalarArithmeticOperation::{Add, Div, Fma, Mul}; + use ScalarArithmeticRounding::{Rm, Rn, Rp, Rz}; + use ScalarArithmeticSaturation::{None, Sat}; + use ScalarArithmeticSubnormal::{Ftz, Preserve}; + + let roundings = [Rn, Rz, Rm, Rp]; + let mut variants = Vec::with_capacity(64); + for operation in [Mul, Div, Fma] { + for rounding in roundings { + variants.push((F64, operation, rounding, Preserve, None)); + } + } + for operation in [Mul, Div] { + for rounding in roundings { + variants.push((F32, operation, rounding, Preserve, None)); + variants.push((F32, operation, rounding, Ftz, None)); + } + } + for rounding in roundings { + variants.push((F32, Fma, rounding, Preserve, None)); + variants.push((F32, Fma, rounding, Ftz, None)); + variants.push((F32, Fma, rounding, Preserve, Sat)); + variants.push((F32, Fma, rounding, Ftz, Sat)); + } + for rounding in roundings { + variants.push((F64, Add, rounding, Preserve, None)); + } + for rounding in roundings { + variants.push((F32, Add, rounding, Preserve, None)); + variants.push((F32, Add, rounding, Ftz, None)); + variants.push((F32, Add, rounding, Preserve, Sat)); + variants.push((F32, Add, rounding, Ftz, Sat)); + } + variants +} + +pub(in crate::resolve) fn scalar_arithmetic_format_name( + format: ScalarArithmeticFormat, +) -> &'static str { + match format { + ScalarArithmeticFormat::F32 => "f32", + ScalarArithmeticFormat::F64 => "f64", + } +} + +pub(in crate::resolve) fn scalar_arithmetic_operation_name( + operation: ScalarArithmeticOperation, +) -> &'static str { + match operation { + ScalarArithmeticOperation::Mul => "mul", + ScalarArithmeticOperation::Div => "div", + ScalarArithmeticOperation::Fma => "fma", + ScalarArithmeticOperation::Add => "add", + } +} + +pub(in crate::resolve) fn scalar_arithmetic_rounding_name( + rounding: ScalarArithmeticRounding, +) -> &'static str { + match rounding { + ScalarArithmeticRounding::Rn => "rn", + ScalarArithmeticRounding::Rz => "rz", + ScalarArithmeticRounding::Rm => "rm", + ScalarArithmeticRounding::Rp => "rp", + } +} + +pub(in crate::resolve) fn scalar_arithmetic_recipe( + variant: ScalarArithmeticVariant, +) -> Option { + if !canonical_scalar_arithmetic_variants().contains(&variant) { + return None; + } + let (format, operation, rounding, subnormal, saturation) = variant; + let operation_name = scalar_arithmetic_operation_name(operation); + let rounding_name = scalar_arithmetic_rounding_name(rounding); + let format_name = scalar_arithmetic_format_name(format); + let source_format = match format { + ScalarArithmeticFormat::F32 => "f", + ScalarArithmeticFormat::F64 => "d", + }; + + let mut modifier_names = vec![rounding_name]; + if subnormal == ScalarArithmeticSubnormal::Ftz { + modifier_names.push("ftz"); + } + if saturation == ScalarArithmeticSaturation::Sat { + modifier_names.push("sat"); + } + + let modifier_id = modifier_names.join("_"); + let modifier_symbol = modifier_names.join("."); + let id = format!("{operation_name}_{modifier_id}_{format_name}"); + let source_record = format!("int_nvvm_{operation_name}_{modifier_id}_{source_format}"); + let llvm_symbol = format!("llvm.nvvm.{operation_name}.{modifier_symbol}.{source_format}"); + let mut ptx_modifier_names = modifier_names.clone(); + if operation == ScalarArithmeticOperation::Add + && subnormal == ScalarArithmeticSubnormal::Ftz + && saturation == ScalarArithmeticSaturation::Sat + { + ptx_modifier_names.swap(1, 2); + } + let ptx_modifiers = ptx_modifier_names + .iter() + .copied() + .chain(std::iter::once(format_name)) + .map(str::to_owned) + .collect::>(); + let selection_record = match operation { + ScalarArithmeticOperation::Fma => { + format!("INT_NVVM_FMA_{modifier_id}_{format_name}") + } + ScalarArithmeticOperation::Mul | ScalarArithmeticOperation::Div => format!( + "INT_NVVM_{}_{}_{}", + operation_name.to_ascii_uppercase(), + modifier_id.to_ascii_uppercase(), + source_format.to_ascii_uppercase() + ), + ScalarArithmeticOperation::Add => { + let selection_modifier = ptx_modifier_names.join("_").to_ascii_uppercase(); + format!( + "INT_NVVM_ADD_{selection_modifier}_{}", + source_format.to_ascii_uppercase() + ) + } + }; + let argument_count = if operation == ScalarArithmeticOperation::Fma { + 3 + } else { + 2 + }; + let operands = (0..argument_count) + .map(|argument| format!("$src{argument}")) + .collect::>() + .join(", "); + let ptx_modifier_symbol = ptx_modifier_names.join("."); + let selection_asm = + format!("{operation_name}.{ptx_modifier_symbol}.{format_name} \t$dst, {operands};"); + let properties = match (operation, format) { + (ScalarArithmeticOperation::Mul | ScalarArithmeticOperation::Add, _) => { + vec!["Commutative", "IntrNoMem", "IntrSpeculatable"] + } + (ScalarArithmeticOperation::Div, _) => vec!["IntrNoMem"], + (ScalarArithmeticOperation::Fma, ScalarArithmeticFormat::F32) => { + vec!["IntrNoMem", "IntrSpeculatable"] + } + (ScalarArithmeticOperation::Fma, ScalarArithmeticFormat::F64) => { + vec!["IntrNoCreateUndefOrPoison", "IntrNoMem", "IntrSpeculatable"] + } + }; + let (ptx_isa_section, ptx_isa_url) = match operation { + ScalarArithmeticOperation::Mul => ( + "9.7.3.5 Floating Point Instructions: mul", + "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-mul", + ), + ScalarArithmeticOperation::Fma => ( + "9.7.3.6 Floating Point Instructions: fma", + "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-fma", + ), + ScalarArithmeticOperation::Div => ( + "9.7.3.8 Floating Point Instructions: div", + "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-div", + ), + ScalarArithmeticOperation::Add => ( + "9.7.3.3 Floating Point Instructions: add", + "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-add", + ), + }; + + Some(ScalarArithmeticRecipe { + id, + operation_key: format!( + "scalar.arithmetic.{operation_name}.{modifier_symbol}.{format_name}" + ), + source_record, + llvm_symbol, + selection_record, + selection_asm, + rust_type: format_name, + argument_count, + properties, + ptx_modifiers, + ptx_isa_section, + ptx_isa_url, + }) +} + +pub(in crate::resolve) fn expand_scalar_arithmetic_admission( + admission: &ScalarArithmeticAdmission, +) -> Result> { + ensure!( + admission.runtime_validation == RuntimeValidation::Unexecuted, + "scalar-arithmetic runtime may be marked executed only with GPU evidence" + ); + let expected = canonical_scalar_arithmetic_variants(); + let actual = admission + .variants + .iter() + .map(|variant| { + ( + variant.format, + variant.operation, + variant.rounding, + variant.subnormal, + variant.saturation, + ) + }) + .collect::>(); + ensure!( + actual == expected, + "compact scalar-arithmetic admission must list the canonical 64 variants" + ); + + admission + .variants + .iter() + .map(|variant| { + validate_abi_id(&variant.abi_id)?; + let identity = ( + variant.format, + variant.operation, + variant.rounding, + variant.subnormal, + variant.saturation, + ); + let recipe = scalar_arithmetic_recipe(identity) + .context("scalar arithmetic is outside the closed recipe set")?; + scalar_arithmetic_overlay_record(recipe, admission, identity, &variant.abi_id) + }) + .collect() +} + +pub(in crate::resolve) fn scalar_arithmetic_overlay_record( + recipe: ScalarArithmeticRecipe, + admission: &ScalarArithmeticAdmission, + variant: ScalarArithmeticVariant, + abi_id: &str, +) -> Result { + let (format, operation, rounding, subnormal, saturation) = variant; + let rust_arguments = vec![recipe.rust_type.into(); recipe.argument_count]; + let ptx_operands = vec![OperandPattern::Register; recipe.argument_count + 1]; + let summary = format!( + "Computes scalar {} with explicit {} rounding.", + scalar_arithmetic_operation_name(operation), + scalar_arithmetic_rounding_name(rounding) + ); + Ok(OverlayIntrinsic { + id: recipe.id.clone(), + abi_id: abi_id.into(), + operation_key: recipe.operation_key, + family: "scalar_arithmetic".into(), + source: None, + source_record: Some(recipe.source_record), + rust_module: "float".into(), + rust_name: recipe.id.clone(), + rust_arguments: rust_arguments.clone(), + rust_result: recipe.rust_type.into(), + safe: true, + must_use: true, + safe_allowlist_reason: Some("Scalar arithmetic has no caller obligations.".into()), + public_rust_path: format!("cuda_intrinsics::float::{}", recipe.id), + compatibility_rust_paths: vec![format!("cuda_device::float::{}", recipe.id)], + dialect_op_type: "ScalarArithmeticOp".into(), + dialect_op_name: "nvvm.scalar_arithmetic".into(), + dialect_operands: rust_arguments.clone(), + dialect_results: vec![recipe.rust_type.into()], + llvm_symbol: Some(recipe.llvm_symbol), + resolved_llvm_symbol: None, + llvm_arguments: rust_arguments, + llvm_results: vec![recipe.rust_type.into()], + pure: operation != ScalarArithmeticOperation::Div, + memory: "none".into(), + convergent: false, + execution_scope: "thread".into(), + minimum_ptx: "7.0".into(), + minimum_sm: Some("sm_80".into()), + ptx_result: recipe.rust_type.into(), + targets: "all".into(), + ptx_isa_version: "9.3".into(), + ptx_isa_section: recipe.ptx_isa_section.into(), + ptx_isa_url: recipe.ptx_isa_url.into(), + lowering: "generated_scalar_arithmetic".into(), + backend_lowerings: [ + ( + IntrinsicBackend::LlvmNvptx, + &admission.llvm_evidence_profile, + ), + ( + IntrinsicBackend::LibNvvm, + &admission.libnvvm_evidence_profile, + ), + ] + .into_iter() + .map(|(backend, evidence_profile)| OverlayBackendLowering { + backend, + mechanism: match backend { + IntrinsicBackend::LlvmNvptx if saturation == ScalarArithmeticSaturation::Sat => { + // LLVM 21 has no typed intrinsic for these LLVM 22 forms. + BackendLoweringMechanism::InlinePtx + } + IntrinsicBackend::LlvmNvptx => BackendLoweringMechanism::TypedNvvm, + IntrinsicBackend::LibNvvm => BackendLoweringMechanism::InlinePtx, + }, + evidence_profile: evidence_profile.clone(), + targets: None, + minimum_ptx: Some("7.0".into()), + minimum_sm: Some("sm_80".into()), + }) + .collect(), + packed_atomic: None, + redux: None, + vote: None, + active_mask: None, + warp_match: None, + warp_barrier: None, + warp_shuffle: None, + dot_product: None, + packed_alu: None, + integer_minmax: None, + packed_conversion: None, + scalar_conversion: None, + scalar_arithmetic: Some(ScalarArithmetic { + format, + operation, + rounding, + subnormal, + saturation, + runtime_validation: admission.runtime_validation, + }), + scalar_math: None, + extended_minmax: None, + cp_async_copy: None, + cp_async_control: None, + cp_async_mbarrier: None, + mbarrier_basic: None, + movmatrix: None, + mbarrier_extended: None, + register_mma: None, + sparse_mma: None, + prmt: None, + cluster_barrier: None, + wgmma_control: None, + special_register: None, + debug_control: None, + cluster_memory: None, + clc: None, + tma: None, + tcgen05: None, + ldmatrix_variant: None, + ldmatrix_safety: None, + ldmatrix_adapter: None, + selected_address_space: None, + expected_ptx: InstructionPattern { + mnemonic: scalar_arithmetic_operation_name(operation).into(), + modifiers: recipe.ptx_modifiers, + operands: ptx_operands, + }, + summary, + }) +} + +pub(in crate::resolve) fn validate_scalar_arithmetic_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, +) -> Result<()> { + let arithmetic = policy + .scalar_arithmetic + .as_ref() + .with_context(|| format!("{} has no scalar-arithmetic contract", policy.id))?; + let variant = ( + arithmetic.format, + arithmetic.operation, + arithmetic.rounding, + arithmetic.subnormal, + arithmetic.saturation, + ); + let recipe = scalar_arithmetic_recipe(variant).with_context(|| { + format!( + "{} is outside the closed scalar-arithmetic recipe", + policy.id + ) + })?; + ensure!( + arithmetic.runtime_validation == RuntimeValidation::Unexecuted, + "{} scalar-arithmetic runtime may be executed only with GPU evidence", + policy.id + ); + let signature = vec![recipe.rust_type.to_owned(); recipe.argument_count]; + ensure!( + policy.id == recipe.id + && policy.operation_key == recipe.operation_key + && policy.source_record.as_deref() == Some(recipe.source_record.as_str()) + && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol.as_str()) + && policy.resolved_llvm_symbol.is_none() + && policy.llvm_arguments == signature + && policy.llvm_results == [recipe.rust_type], + "{} scalar-arithmetic identity or LLVM source changed", + policy.id + ); + let expected_properties = recipe + .properties + .iter() + .map(|property| (*property).to_owned()) + .collect::>(); + ensure!( + declaration.properties == expected_properties, + "{} imported scalar-arithmetic properties changed", + policy.id + ); + let direct = declaration + .selections + .first() + .with_context(|| format!("{} has no imported scalar selection", policy.id))?; + ensure!( + direct.source_record == recipe.selection_record + && direct.asm == recipe.selection_asm + && direct.predicates.is_empty() + && direct.constraints.is_empty(), + "{} direct scalar-arithmetic selection changed", + policy.id + ); + if arithmetic.operation == ScalarArithmeticOperation::Add { + let rounding = scalar_arithmetic_rounding_name(arithmetic.rounding); + let source_modifiers = [ + Some(rounding), + (arithmetic.subnormal == ScalarArithmeticSubnormal::Ftz).then_some("ftz"), + (arithmetic.saturation == ScalarArithmeticSaturation::Sat).then_some("sat"), + ] + .into_iter() + .flatten() + .collect::>() + .join("_"); + let mut expected = Vec::new(); + if arithmetic.format == ScalarArithmeticFormat::F32 + && arithmetic.subnormal == ScalarArithmeticSubnormal::Preserve + { + let saturation = if arithmetic.saturation == ScalarArithmeticSaturation::Sat { + "_sat" + } else { + "" + }; + for instruction in ["ADD", "SUB"] { + for input_format in ["bf16", "f16"] { + expected.push(( + format!( + "INT_NVVM_MIXED_{instruction}_{rounding}{saturation}_f32_{input_format}" + ), + format!( + "{}.{rounding}{}.f32.{input_format} \t$dst, $a, $b;", + instruction.to_ascii_lowercase(), + if arithmetic.saturation == ScalarArithmeticSaturation::Sat { + ".sat" + } else { + "" + } + ), + vec![ + "Subtarget->getSmVersion() >= 100".to_owned(), + "Subtarget->getPTXVersion() >= 86".to_owned(), + ], + )); + } + } + } + let source_format = match arithmetic.format { + ScalarArithmeticFormat::F32 => "F", + ScalarArithmeticFormat::F64 => "D", + }; + expected.push(( + format!("INT_NVVM_SUB_{source_modifiers}_{source_format}"), + format!("sub.{} \t$dst, $a, $b;", recipe.ptx_modifiers.join(".")), + Vec::new(), + )); + ensure!( + declaration.selections.len() == expected.len() + 1, + "{} must retain one direct add and the reviewed non-add alternatives", + policy.id + ); + for (selection, (source_record, asm, predicates)) in + declaration.selections[1..].iter().zip(expected) + { + ensure!( + selection.source_record == source_record + && selection.asm == asm + && selection.predicates == predicates + && selection.constraints.is_empty(), + "{} non-add scalar selection changed", + policy.id + ); + } + } else if arithmetic.format == ScalarArithmeticFormat::F32 + && arithmetic.operation == ScalarArithmeticOperation::Fma + && arithmetic.subnormal == ScalarArithmeticSubnormal::Preserve + { + ensure!( + declaration.selections.len() == 3, + "{} must retain one direct and two mixed-input selections", + policy.id + ); + let rounding = scalar_arithmetic_rounding_name(arithmetic.rounding); + let saturation = if arithmetic.saturation == ScalarArithmeticSaturation::Sat { + "_sat" + } else { + "" + }; + for (selection, input_format) in declaration.selections[1..].iter().zip(["bf16", "f16"]) { + ensure!( + selection.source_record + == format!("INT_NVVM_MIXED_FMA_{rounding}{saturation}_f32_{input_format}") + && selection.asm + == format!( + "fma.{rounding}{}.f32.{input_format} \t$dst, $a, $b, $c;", + if arithmetic.saturation == ScalarArithmeticSaturation::Sat { + ".sat" + } else { + "" + } + ) + && selection.predicates + == [ + "Subtarget->getSmVersion() >= 100", + "Subtarget->getPTXVersion() >= 86", + ] + && selection.constraints.is_empty(), + "{} mixed-input scalar-arithmetic selection changed", + policy.id + ); + } + } else { + ensure!( + declaration.selections.len() == 1, + "{} gained an unreviewed scalar-arithmetic selection", + policy.id + ); + } + let mut selected = Vec::new(); + for selection in &declaration.selections { + if selection_matches_policy(policy, selection)? { + selected.push(selection); + } + } + ensure!( + selected.len() == 1 && selected[0].source_record == recipe.selection_record, + "{} must select only its direct scalar arithmetic instruction", + policy.id + ); + ensure!( + policy.rust_module == "float" + && policy.rust_name == recipe.id + && policy.rust_arguments == signature + && policy.rust_result == recipe.rust_type + && policy.safe + && policy.must_use + && policy.compatibility_rust_paths == [format!("cuda_device::float::{}", recipe.id)] + && policy.dialect_op_type == "ScalarArithmeticOp" + && policy.dialect_op_name == "nvvm.scalar_arithmetic" + && policy.dialect_operands == signature + && policy.dialect_results == [recipe.rust_type] + && policy.lowering == "generated_scalar_arithmetic", + "{} changed its scalar-arithmetic API, carrier, or lowering", + policy.id + ); + ensure!( + policy.pure == (arithmetic.operation != ScalarArithmeticOperation::Div) + && policy.memory == "none" + && !policy.convergent + && policy.execution_scope == "thread" + && policy.minimum_ptx == "7.0" + && policy.minimum_sm.as_deref() == Some("sm_80") + && policy.ptx_result == recipe.rust_type + && policy.targets == "all" + && policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section == recipe.ptx_isa_section + && policy.ptx_isa_url == recipe.ptx_isa_url, + "{} scalar-arithmetic effects, provenance, or target floor changed", + policy.id + ); + ensure!( + policy.expected_ptx + == InstructionPattern { + mnemonic: scalar_arithmetic_operation_name(arithmetic.operation).into(), + modifiers: recipe.ptx_modifiers.clone(), + operands: vec![OperandPattern::Register; recipe.argument_count + 1], + }, + "{} expected scalar-arithmetic PTX changed", + policy.id + ); + let llvm_mechanism = if arithmetic.saturation == ScalarArithmeticSaturation::Sat { + BackendLoweringMechanism::InlinePtx + } else { + BackendLoweringMechanism::TypedNvvm + }; + let expected_backends = [ + (IntrinsicBackend::LlvmNvptx, llvm_mechanism), + ( + IntrinsicBackend::LibNvvm, + BackendLoweringMechanism::InlinePtx, + ), + ]; + ensure!( + policy.backend_lowerings.len() == 2 + && expected_backends.into_iter().all(|(backend, mechanism)| { + policy.backend_lowerings.iter().any(|lowering| { + lowering.backend == backend + && lowering.mechanism == mechanism + && lowering.minimum_ptx.as_deref() == Some("7.0") + && lowering.minimum_sm.as_deref() == Some("sm_80") + && !lowering.evidence_profile.trim().is_empty() + }) + }), + "{} has the wrong reviewed scalar-arithmetic backend routes", + policy.id + ); + validate_selected_target_predicates(policy, direct)?; + ensure_no_other_family_contract(policy, "scalar arithmetic")?; + Ok(()) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/scalar_conversion.rs b/crates/cuda-intrinsics-gen/src/resolve/families/scalar_conversion.rs new file mode 100644 index 0000000000..b683d5147e --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/scalar_conversion.rs @@ -0,0 +1,490 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, ImportedIntrinsic, IntrinsicBackend, OverlayBackendLowering, + OverlayIntrinsic, RuntimeValidation, ScalarConversion, ScalarConversionAdapter, + ScalarConversionAdmission, ScalarConversionDestinationFormat, + ScalarConversionResultRepresentation, ScalarConversionRounding, ScalarConversionSaturation, + ScalarConversionSourceFormat, +}; +use crate::ptx::{InstructionPattern, OperandPattern}; +use anyhow::{Context, Result, ensure}; + +use crate::resolve::guards::*; +use crate::resolve::targets::*; + +#[derive(Clone, Copy)] +pub(in crate::resolve) struct ScalarConversionRecipe { + id: &'static str, + abi_id: &'static str, + operation_key: &'static str, + source_record: &'static str, + llvm_symbol: &'static str, + rust_name: &'static str, + selection_record: &'static str, + selection_asm: &'static str, + minimum_ptx: &'static str, + minimum_sm: &'static str, + ptx_modifiers: &'static [&'static str], +} + +pub(in crate::resolve) fn scalar_conversion_recipe( + rounding: ScalarConversionRounding, + saturation: ScalarConversionSaturation, +) -> Option { + use ScalarConversionRounding::{NearestAway, NearestEven, TowardZero}; + use ScalarConversionSaturation::{None, Relu, ReluSatfinite, Satfinite}; + + Some(match (rounding, saturation) { + (NearestAway, None) => ScalarConversionRecipe { + id: "cvt_rna_tf32_f32", + abi_id: "i0368", + operation_key: "convert.f32.tf32.rna", + source_record: "int_nvvm_f2tf32_rna", + llvm_symbol: "llvm.nvvm.f2tf32.rna", + rust_name: "cvt_rna_tf32_f32", + selection_record: "CVT_to_tf32_rna", + selection_asm: "cvt.rna.tf32.f32 \t$dst, $src;", + minimum_ptx: "7.0", + minimum_sm: "sm_80", + ptx_modifiers: &["rna", "tf32", "f32"], + }, + (NearestAway, Satfinite) => ScalarConversionRecipe { + id: "cvt_rna_satfinite_tf32_f32", + abi_id: "i0369", + operation_key: "convert.f32.tf32.rna.satfinite", + source_record: "int_nvvm_f2tf32_rna_satfinite", + llvm_symbol: "llvm.nvvm.f2tf32.rna.satfinite", + rust_name: "cvt_rna_satfinite_tf32_f32", + selection_record: "CVT_to_tf32_rna_satf", + selection_asm: "cvt.rna.satfinite.tf32.f32 \t$dst, $src;", + minimum_ptx: "8.1", + minimum_sm: "sm_80", + ptx_modifiers: &["rna", "satfinite", "tf32", "f32"], + }, + (NearestEven, None) => ScalarConversionRecipe { + id: "cvt_rn_tf32_f32", + abi_id: "i0370", + operation_key: "convert.f32.tf32.rn", + source_record: "int_nvvm_f2tf32_rn", + llvm_symbol: "llvm.nvvm.f2tf32.rn", + rust_name: "cvt_rn_tf32_f32", + selection_record: "CVT_to_tf32_rn", + selection_asm: "cvt.rn.tf32.f32 \t$dst, $src;", + minimum_ptx: "7.8", + minimum_sm: "sm_90", + ptx_modifiers: &["rn", "tf32", "f32"], + }, + (NearestEven, Relu) => ScalarConversionRecipe { + id: "cvt_rn_relu_tf32_f32", + abi_id: "i0371", + operation_key: "convert.f32.tf32.rn.relu", + source_record: "int_nvvm_f2tf32_rn_relu", + llvm_symbol: "llvm.nvvm.f2tf32.rn.relu", + rust_name: "cvt_rn_relu_tf32_f32", + selection_record: "CVT_to_tf32_rn_relu", + selection_asm: "cvt.rn.relu.tf32.f32 \t$dst, $src;", + minimum_ptx: "7.8", + minimum_sm: "sm_90", + ptx_modifiers: &["rn", "relu", "tf32", "f32"], + }, + (NearestEven, Satfinite) => ScalarConversionRecipe { + id: "cvt_rn_satfinite_tf32_f32", + abi_id: "i0372", + operation_key: "convert.f32.tf32.rn.satfinite", + source_record: "int_nvvm_f2tf32_rn_satfinite", + llvm_symbol: "llvm.nvvm.f2tf32.rn.satfinite", + rust_name: "cvt_rn_satfinite_tf32_f32", + selection_record: "CVT_to_tf32_rn_satf", + selection_asm: "cvt.rn.satfinite.tf32.f32 \t$dst, $src;", + minimum_ptx: "8.6", + minimum_sm: "sm_100", + ptx_modifiers: &["rn", "satfinite", "tf32", "f32"], + }, + (NearestEven, ReluSatfinite) => ScalarConversionRecipe { + id: "cvt_rn_relu_satfinite_tf32_f32", + abi_id: "i0373", + operation_key: "convert.f32.tf32.rn.relu.satfinite", + source_record: "int_nvvm_f2tf32_rn_relu_satfinite", + llvm_symbol: "llvm.nvvm.f2tf32.rn.relu.satfinite", + rust_name: "cvt_rn_relu_satfinite_tf32_f32", + selection_record: "CVT_to_tf32_rn_relu_satf", + selection_asm: "cvt.rn.relu.satfinite.tf32.f32 \t$dst, $src;", + minimum_ptx: "8.6", + minimum_sm: "sm_100", + ptx_modifiers: &["rn", "relu", "satfinite", "tf32", "f32"], + }, + (TowardZero, None) => ScalarConversionRecipe { + id: "cvt_rz_tf32_f32", + abi_id: "i0374", + operation_key: "convert.f32.tf32.rz", + source_record: "int_nvvm_f2tf32_rz", + llvm_symbol: "llvm.nvvm.f2tf32.rz", + rust_name: "cvt_rz_tf32_f32", + selection_record: "CVT_to_tf32_rz", + selection_asm: "cvt.rz.tf32.f32 \t$dst, $src;", + minimum_ptx: "7.8", + minimum_sm: "sm_90", + ptx_modifiers: &["rz", "tf32", "f32"], + }, + (TowardZero, Relu) => ScalarConversionRecipe { + id: "cvt_rz_relu_tf32_f32", + abi_id: "i0375", + operation_key: "convert.f32.tf32.rz.relu", + source_record: "int_nvvm_f2tf32_rz_relu", + llvm_symbol: "llvm.nvvm.f2tf32.rz.relu", + rust_name: "cvt_rz_relu_tf32_f32", + selection_record: "CVT_to_tf32_rz_relu", + selection_asm: "cvt.rz.relu.tf32.f32 \t$dst, $src;", + minimum_ptx: "7.8", + minimum_sm: "sm_90", + ptx_modifiers: &["rz", "relu", "tf32", "f32"], + }, + (TowardZero, Satfinite) => ScalarConversionRecipe { + id: "cvt_rz_satfinite_tf32_f32", + abi_id: "i0376", + operation_key: "convert.f32.tf32.rz.satfinite", + source_record: "int_nvvm_f2tf32_rz_satfinite", + llvm_symbol: "llvm.nvvm.f2tf32.rz.satfinite", + rust_name: "cvt_rz_satfinite_tf32_f32", + selection_record: "CVT_to_tf32_rz_satf", + selection_asm: "cvt.rz.satfinite.tf32.f32 \t$dst, $src;", + minimum_ptx: "8.6", + minimum_sm: "sm_100", + ptx_modifiers: &["rz", "satfinite", "tf32", "f32"], + }, + (TowardZero, ReluSatfinite) => ScalarConversionRecipe { + id: "cvt_rz_relu_satfinite_tf32_f32", + abi_id: "i0377", + operation_key: "convert.f32.tf32.rz.relu.satfinite", + source_record: "int_nvvm_f2tf32_rz_relu_satfinite", + llvm_symbol: "llvm.nvvm.f2tf32.rz.relu.satfinite", + rust_name: "cvt_rz_relu_satfinite_tf32_f32", + selection_record: "CVT_to_tf32_rz_relu_satf", + selection_asm: "cvt.rz.relu.satfinite.tf32.f32 \t$dst, $src;", + minimum_ptx: "8.6", + minimum_sm: "sm_100", + ptx_modifiers: &["rz", "relu", "satfinite", "tf32", "f32"], + }, + _ => return Option::None, + }) +} + +pub(in crate::resolve) fn expand_scalar_conversion_admission( + admission: &ScalarConversionAdmission, +) -> Result> { + ensure!( + admission.runtime_validation == RuntimeValidation::Unexecuted, + "scalar-conversion runtime may be marked executed only with GPU evidence" + ); + let expected = [ + ( + ScalarConversionRounding::NearestAway, + ScalarConversionSaturation::None, + ), + ( + ScalarConversionRounding::NearestAway, + ScalarConversionSaturation::Satfinite, + ), + ( + ScalarConversionRounding::NearestEven, + ScalarConversionSaturation::None, + ), + ( + ScalarConversionRounding::NearestEven, + ScalarConversionSaturation::Relu, + ), + ( + ScalarConversionRounding::NearestEven, + ScalarConversionSaturation::Satfinite, + ), + ( + ScalarConversionRounding::NearestEven, + ScalarConversionSaturation::ReluSatfinite, + ), + ( + ScalarConversionRounding::TowardZero, + ScalarConversionSaturation::None, + ), + ( + ScalarConversionRounding::TowardZero, + ScalarConversionSaturation::Relu, + ), + ( + ScalarConversionRounding::TowardZero, + ScalarConversionSaturation::Satfinite, + ), + ( + ScalarConversionRounding::TowardZero, + ScalarConversionSaturation::ReluSatfinite, + ), + ]; + let actual = admission + .variants + .iter() + .map(|variant| (variant.rounding, variant.saturation)) + .collect::>(); + ensure!( + actual == expected, + "compact scalar-conversion admission must list the canonical ten variants" + ); + + admission + .variants + .iter() + .map(|variant| { + let recipe = scalar_conversion_recipe(variant.rounding, variant.saturation) + .context("scalar conversion is outside the closed recipe set")?; + ensure!( + variant.abi_id == recipe.abi_id, + "{} must reserve ABI ID {}", + recipe.id, + recipe.abi_id + ); + scalar_conversion_overlay_record( + recipe, + admission, + variant.rounding, + variant.saturation, + ) + }) + .collect() +} + +pub(in crate::resolve) fn scalar_conversion_overlay_record( + recipe: ScalarConversionRecipe, + admission: &ScalarConversionAdmission, + rounding: ScalarConversionRounding, + saturation: ScalarConversionSaturation, +) -> Result { + Ok(OverlayIntrinsic { + id: recipe.id.into(), + abi_id: recipe.abi_id.into(), + operation_key: recipe.operation_key.into(), + family: "scalar_conversion".into(), + source: None, + source_record: Some(recipe.source_record.into()), + rust_module: "convert".into(), + rust_name: recipe.rust_name.into(), + rust_arguments: vec!["f32".into()], + rust_result: "u32".into(), + safe: true, + must_use: true, + safe_allowlist_reason: Some("This conversion has no caller obligations.".into()), + public_rust_path: format!("cuda_intrinsics::convert::{}", recipe.rust_name), + compatibility_rust_paths: vec![format!( + "cuda_device::convert::{}", + recipe.rust_name + )], + dialect_op_type: "ScalarConversionOp".into(), + dialect_op_name: "nvvm.scalar_conversion".into(), + dialect_operands: vec!["f32".into()], + dialect_results: vec!["i32".into()], + llvm_symbol: Some(recipe.llvm_symbol.into()), + resolved_llvm_symbol: None, + llvm_arguments: vec!["f32".into()], + llvm_results: vec!["i32".into()], + pure: true, + memory: "none".into(), + convergent: false, + execution_scope: "thread".into(), + minimum_ptx: recipe.minimum_ptx.into(), + minimum_sm: Some(recipe.minimum_sm.into()), + ptx_result: "u32".into(), + targets: "all".into(), + ptx_isa_version: "9.3".into(), + ptx_isa_section: "9.7.9.22 Data Movement and Conversion Instructions: cvt".into(), + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cvt".into(), + lowering: "generated_scalar_conversion".into(), + backend_lowerings: [ + (IntrinsicBackend::LlvmNvptx, &admission.llvm_evidence_profile), + (IntrinsicBackend::LibNvvm, &admission.libnvvm_evidence_profile), + ] + .into_iter() + .map(|(backend, evidence_profile)| OverlayBackendLowering { + backend, + mechanism: match backend { + IntrinsicBackend::LlvmNvptx => BackendLoweringMechanism::TypedNvvm, + IntrinsicBackend::LibNvvm => BackendLoweringMechanism::InlinePtx, + }, + evidence_profile: evidence_profile.clone(), + targets: None, + minimum_ptx: Some(recipe.minimum_ptx.into()), + minimum_sm: Some(recipe.minimum_sm.into()), + }) + .collect(), + packed_atomic: None, + redux: None, + vote: None, + active_mask: None, + warp_match: None, + warp_barrier: None, + warp_shuffle: None, + dot_product: None, + packed_alu: None, + integer_minmax: None, + packed_conversion: None, + scalar_conversion: Some(ScalarConversion { + source_format: ScalarConversionSourceFormat::F32, + destination_format: ScalarConversionDestinationFormat::Tf32, + rounding, + saturation, + result_representation: ScalarConversionResultRepresentation::RawU32Bits, + adapter: ScalarConversionAdapter::DirectF32ToRawU32Bits, + runtime_validation: admission.runtime_validation, + }), + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cp_async_copy: None, + cp_async_control: None, + cp_async_mbarrier: None, + mbarrier_basic: None, + movmatrix: None, + mbarrier_extended: None, + register_mma: None, + sparse_mma: None, + prmt: None, + cluster_barrier: None, + wgmma_control: None, + special_register: None, + debug_control: None, + cluster_memory: None, + clc: None, + tma: None, + tcgen05: None, + ldmatrix_variant: None, + ldmatrix_safety: None, + ldmatrix_adapter: None, + selected_address_space: None, + expected_ptx: InstructionPattern { + mnemonic: "cvt".into(), + modifiers: recipe.ptx_modifiers.iter().map(|value| (*value).into()).collect(), + operands: vec![OperandPattern::Register, OperandPattern::Register], + }, + summary: format!( + "Converts one f32 value with {} and returns raw TF32 bits.", + recipe.ptx_modifiers[0] + ), + }) +} + +pub(in crate::resolve) fn validate_scalar_conversion_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, +) -> Result<()> { + let conversion = policy + .scalar_conversion + .as_ref() + .with_context(|| format!("{} has no scalar-conversion contract", policy.id))?; + let recipe = scalar_conversion_recipe(conversion.rounding, conversion.saturation) + .with_context(|| { + format!( + "{} is outside the closed scalar-conversion recipe", + policy.id + ) + })?; + ensure!( + conversion.source_format == ScalarConversionSourceFormat::F32 + && conversion.destination_format == ScalarConversionDestinationFormat::Tf32 + && conversion.result_representation == ScalarConversionResultRepresentation::RawU32Bits + && conversion.adapter == ScalarConversionAdapter::DirectF32ToRawU32Bits + && conversion.runtime_validation == RuntimeValidation::Unexecuted, + "{} changed its scalar-conversion representation or adapter", + policy.id + ); + ensure!( + policy.id == recipe.id + && policy.abi_id == recipe.abi_id + && policy.operation_key == recipe.operation_key + && policy.source_record.as_deref() == Some(recipe.source_record) + && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) + && policy.resolved_llvm_symbol.is_none() + && policy.llvm_arguments == ["f32"] + && policy.llvm_results == ["i32"], + "{} scalar-conversion identity or LLVM source changed", + policy.id + ); + ensure!( + declaration.properties == ["IntrNoMem", "IntrSpeculatable"] + && declaration.selections.len() == 1 + && declaration.selections[0].source_record == recipe.selection_record + && declaration.selections[0].asm == recipe.selection_asm + && declaration.selections[0].constraints.is_empty(), + "{} imported scalar-conversion declaration or selection changed", + policy.id + ); + ensure!( + policy.rust_module == "convert" + && policy.rust_name == recipe.rust_name + && policy.rust_arguments == ["f32"] + && policy.rust_result == "u32" + && policy.safe + && policy.must_use + && policy.compatibility_rust_paths + == [format!("cuda_device::convert::{}", recipe.rust_name)] + && policy.dialect_op_type == "ScalarConversionOp" + && policy.dialect_op_name == "nvvm.scalar_conversion" + && policy.dialect_operands == ["f32"] + && policy.dialect_results == ["i32"] + && policy.lowering == "generated_scalar_conversion", + "{} changed its scalar-conversion API, carrier, or lowering", + policy.id + ); + ensure!( + policy.pure + && policy.memory == "none" + && !policy.convergent + && policy.execution_scope == "thread" + && policy.minimum_ptx == recipe.minimum_ptx + && policy.minimum_sm.as_deref() == Some(recipe.minimum_sm) + && policy.ptx_result == "u32" + && policy.targets == "all", + "{} scalar-conversion effects or target floor changed", + policy.id + ); + ensure!( + policy.expected_ptx + == InstructionPattern { + mnemonic: "cvt".into(), + modifiers: recipe + .ptx_modifiers + .iter() + .map(|value| (*value).into()) + .collect(), + operands: vec![OperandPattern::Register, OperandPattern::Register], + }, + "{} expected PTX changed", + policy.id + ); + let expected_backends = [ + ( + IntrinsicBackend::LlvmNvptx, + BackendLoweringMechanism::TypedNvvm, + ), + ( + IntrinsicBackend::LibNvvm, + BackendLoweringMechanism::InlinePtx, + ), + ]; + ensure!( + policy.backend_lowerings.len() == 2 + && expected_backends.into_iter().all(|(backend, mechanism)| { + policy.backend_lowerings.iter().any(|lowering| { + lowering.backend == backend + && lowering.mechanism == mechanism + && lowering.minimum_ptx.as_deref() == Some(recipe.minimum_ptx) + && lowering.minimum_sm.as_deref() == Some(recipe.minimum_sm) + && !lowering.evidence_profile.trim().is_empty() + }) + }), + "{} must define the typed LLVM and inline-PTX libNVVM routes", + policy.id + ); + validate_selected_target_predicates(policy, &declaration.selections[0])?; + ensure_no_other_family_contract(policy, "scalar conversion")?; + Ok(()) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/scalar_math.rs b/crates/cuda-intrinsics-gen/src/resolve/families/scalar_math.rs new file mode 100644 index 0000000000..54b0c0f3cb --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/scalar_math.rs @@ -0,0 +1,692 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, ImportedIntrinsic, IntrinsicBackend, IntrinsicSource, + OverlayBackendLowering, OverlayIntrinsic, RuntimeValidation, ScalarMath, ScalarMathAdmission, + ScalarMathFormat, ScalarMathOperation, ScalarMathPrecision, ScalarMathSubnormal, +}; +use crate::ptx::{InstructionPattern, OperandPattern}; +use anyhow::{Context, Result, ensure}; + +use crate::resolve::abi_ledger::*; +use crate::resolve::guards::*; +use crate::resolve::targets::*; + +pub(in crate::resolve) type ScalarMathVariant = ( + ScalarMathFormat, + ScalarMathOperation, + ScalarMathPrecision, + ScalarMathSubnormal, +); + +#[derive(Clone, PartialEq, Eq)] +pub(in crate::resolve) enum ScalarMathRecipeSource { + /// Bound to a monomorphic tblgen record in the pinned import. + Imported { + source_record: String, + llvm_symbol: String, + }, + /// Bound to an overloaded (polymorphic) tblgen record: the import + /// carries `anonymous_14`/`anyfloat` signature tokens and the concrete + /// f32 instantiation is recorded as the resolved symbol. Only ex2 uses + /// this (LLVM 22 models it as `int_nvvm_ex2_approx{,_ftz}` without a + /// per-format record). + ImportedOverloaded { + source_record: String, + llvm_symbol: String, + resolved_llvm_symbol: String, + }, + /// No record exists in the pinned tblgen import at all. Only tanh uses + /// this: llc selects `llvm.nvvm.tanh.approx.f32` via NVVMIntrinsic-class + /// matching, but the import exports no record for it, so the op is + /// admitted directly against the PTX instruction. + PtxNative { instruction: String }, +} + +pub(in crate::resolve) struct ScalarMathRecipe { + id: String, + operation_key: String, + source: ScalarMathRecipeSource, + rust_type: &'static str, + dialect_type: &'static str, + minimum_ptx: &'static str, + minimum_sm: &'static str, + properties: Vec<&'static str>, + ptx_modifiers: Vec, + ptx_isa_section: &'static str, + ptx_isa_url: &'static str, + /// The generator's tblgen import found no DAG selection pattern for the + /// intrinsic, so the lowering routes it through inline PTX. Note this is + /// an import limitation, not an llc one: llc still selects these + /// intrinsics through NVVMIntrinsic-class pattern matching, which the + /// evidence import cannot see. Promoting them to typed calls once the + /// import understands those patterns would reopen them to LLVM + /// optimization. + force_inline_ptx: bool, +} + +pub(in crate::resolve) fn scalar_math_format_name(format: ScalarMathFormat) -> &'static str { + match format { + ScalarMathFormat::F16 => "f16", + ScalarMathFormat::F32 => "f32", + ScalarMathFormat::F64 => "f64", + } +} + +pub(in crate::resolve) fn scalar_math_operation_name( + operation: ScalarMathOperation, +) -> &'static str { + match operation { + ScalarMathOperation::Sin => "sin", + ScalarMathOperation::Cos => "cos", + ScalarMathOperation::Ex2 => "ex2", + ScalarMathOperation::Lg2 => "lg2", + ScalarMathOperation::Rcp => "rcp", + ScalarMathOperation::Rsqrt => "rsqrt", + ScalarMathOperation::Sqrt => "sqrt", + ScalarMathOperation::Tanh => "tanh", + } +} + +pub(in crate::resolve) fn scalar_math_precision_name( + precision: ScalarMathPrecision, +) -> &'static str { + match precision { + ScalarMathPrecision::Approx => "approx", + ScalarMathPrecision::Rn => "rn", + ScalarMathPrecision::Rz => "rz", + ScalarMathPrecision::Rm => "rm", + ScalarMathPrecision::Rp => "rp", + } +} + +pub(in crate::resolve) fn canonical_scalar_math_variants() -> Vec { + use ScalarMathFormat::{F16, F32, F64}; + use ScalarMathOperation::{Cos, Ex2, Lg2, Rcp, Rsqrt, Sin, Sqrt, Tanh}; + use ScalarMathPrecision::{Approx, Rm, Rn, Rp, Rz}; + use ScalarMathSubnormal::{Ftz, Preserve}; + + vec![ + // sin: approx f32 only + (F32, Sin, Approx, Preserve), + (F32, Sin, Approx, Ftz), + // cos: approx f32 only + (F32, Cos, Approx, Preserve), + (F32, Cos, Approx, Ftz), + // lg2: approx f32 only (f64 approx produces invalid PTX) + (F32, Lg2, Approx, Preserve), + (F32, Lg2, Approx, Ftz), + // rcp: approx only with ftz, rounded for both formats + (F32, Rcp, Approx, Ftz), + (F32, Rcp, Rn, Preserve), + (F32, Rcp, Rn, Ftz), + (F32, Rcp, Rz, Preserve), + (F32, Rcp, Rz, Ftz), + (F32, Rcp, Rm, Preserve), + (F32, Rcp, Rm, Ftz), + (F32, Rcp, Rp, Preserve), + (F32, Rcp, Rp, Ftz), + (F64, Rcp, Approx, Ftz), + (F64, Rcp, Rn, Preserve), + (F64, Rcp, Rz, Preserve), + (F64, Rcp, Rm, Preserve), + (F64, Rcp, Rp, Preserve), + // rsqrt: approx only (PTX has no rounded rsqrt). The f64+ftz variant + // is valid PTX (`rsqrt.approx.ftz.f64` assembles for sm_80+, and LLVM + // selects `llvm.nvvm.rsqrt.approx.ftz.d` directly); it is deferred + // only because it has not been probed under the pinned evidence + // profile yet, not because the instruction is invalid. + (F32, Rsqrt, Approx, Preserve), + (F32, Rsqrt, Approx, Ftz), + (F64, Rsqrt, Approx, Preserve), + // sqrt: approx f32 only, rounded for both formats + (F32, Sqrt, Approx, Preserve), + (F32, Sqrt, Approx, Ftz), + (F32, Sqrt, Rn, Preserve), + (F32, Sqrt, Rn, Ftz), + (F32, Sqrt, Rz, Preserve), + (F32, Sqrt, Rz, Ftz), + (F32, Sqrt, Rm, Preserve), + (F32, Sqrt, Rm, Ftz), + (F32, Sqrt, Rp, Preserve), + (F32, Sqrt, Rp, Ftz), + (F64, Sqrt, Rn, Preserve), + (F64, Sqrt, Rz, Preserve), + (F64, Sqrt, Rm, Preserve), + (F64, Sqrt, Rp, Preserve), + // ex2: approx f32 only (PTX has no ex2.approx.f64). Appended after + // the original 37 so existing ABI ids stay stable. + (F32, Ex2, Approx, Preserve), + (F32, Ex2, Approx, Ftz), + // tanh: approx f32 only; the instruction has no ftz form (PTX ISA + // Table 29) and no rounded variants. Hardware floor is sm_75; the + // family contract gates it at the attested sm_80 evidence floor. + (F32, Tanh, Approx, Preserve), + // ex2.approx.f16 starts at PTX 7.0 / sm_75. LLVM 22 rejects the + // .ftz.f16 spelling; .ftz.bf16 is a separate PTX 7.8 / sm_90 variant. + (F16, Ex2, Approx, Preserve), + ] +} + +pub(in crate::resolve) fn scalar_math_recipe( + variant: ScalarMathVariant, +) -> Option { + if !canonical_scalar_math_variants().contains(&variant) { + return None; + } + let (format, operation, precision, subnormal) = variant; + let operation_name = scalar_math_operation_name(operation); + let precision_name = scalar_math_precision_name(precision); + let format_name = scalar_math_format_name(format); + let source_format = match format { + ScalarMathFormat::F16 => "f16", + ScalarMathFormat::F32 => "f", + ScalarMathFormat::F64 => "d", + }; + + let mut modifier_names = vec![precision_name]; + if subnormal == ScalarMathSubnormal::Ftz { + modifier_names.push("ftz"); + } + + let modifier_id = modifier_names.join("_"); + let modifier_symbol = modifier_names.join("."); + let id = format!("{operation_name}_{modifier_id}_{format_name}"); + let source = match operation { + // LLVM 22 models ex2 as overloaded records without a per-format + // suffix (`int_nvvm_ex2_approx{,_ftz}` over anyfloat); bind those + // for declaration-level checks and record the concrete f32 + // instantiation as the resolved symbol. + ScalarMathOperation::Ex2 => ScalarMathRecipeSource::ImportedOverloaded { + source_record: format!("int_nvvm_{operation_name}_{modifier_id}"), + llvm_symbol: format!("llvm.nvvm.{operation_name}.{modifier_symbol}"), + resolved_llvm_symbol: format!( + "llvm.nvvm.{operation_name}.{modifier_symbol}.{format_name}" + ), + }, + // LLVM 22.1.2's tblgen export has no record for tanh at all, so the + // op is admitted directly against the PTX instruction. + ScalarMathOperation::Tanh => ScalarMathRecipeSource::PtxNative { + instruction: format!("{operation_name}.{modifier_symbol}.{format_name}"), + }, + _ => ScalarMathRecipeSource::Imported { + source_record: format!("int_nvvm_{operation_name}_{modifier_id}_{source_format}"), + llvm_symbol: format!("llvm.nvvm.{operation_name}.{modifier_symbol}.{source_format}"), + }, + }; + let ptx_modifiers = modifier_names + .iter() + .copied() + .chain(std::iter::once(format_name)) + .map(str::to_owned) + .collect::>(); + + // Operations whose llvm.nvvm.* records carry no *imported* DAG selection + // pattern in the pinned evidence (sel=0). llc itself still selects them + // (via NVVMIntrinsic-class patterns the tblgen import cannot see), but + // the evidence-driven contract only admits typed calls backed by an + // imported selection, so these route through inline PTX for now. + let force_inline_ptx = matches!( + operation, + ScalarMathOperation::Sin + | ScalarMathOperation::Cos + | ScalarMathOperation::Ex2 + | ScalarMathOperation::Lg2 + | ScalarMathOperation::Rsqrt + | ScalarMathOperation::Tanh + ); + + let (ptx_isa_section, ptx_isa_url) = match (operation, format) { + (ScalarMathOperation::Ex2, ScalarMathFormat::F16) => ( + "9.7.4.10 Half Precision Floating Point Instructions: ex2", + "https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-ex2", + ), + (ScalarMathOperation::Sin, _) => ( + "9.7.3.10 Floating Point Instructions: sin", + "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-sin", + ), + (ScalarMathOperation::Cos, _) => ( + "9.7.3.11 Floating Point Instructions: cos", + "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-cos", + ), + (ScalarMathOperation::Lg2, _) => ( + "9.7.3.12 Floating Point Instructions: lg2", + "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-lg2", + ), + (ScalarMathOperation::Rcp, _) => ( + "9.7.3.7 Floating Point Instructions: rcp", + "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-rcp", + ), + (ScalarMathOperation::Rsqrt, _) => ( + "9.7.3.14 Floating Point Instructions: rsqrt", + "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-rsqrt", + ), + (ScalarMathOperation::Sqrt, _) => ( + "9.7.3.9 Floating Point Instructions: sqrt", + "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-sqrt", + ), + (ScalarMathOperation::Ex2, _) => ( + "9.7.3.13 Floating Point Instructions: ex2", + "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-ex2", + ), + (ScalarMathOperation::Tanh, _) => ( + "9.7.3.15 Floating Point Instructions: tanh", + "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-tanh", + ), + }; + + let (rust_type, dialect_type, minimum_ptx, minimum_sm) = match format { + ScalarMathFormat::F16 => ("u16", "i16", "7.0", "sm_75"), + ScalarMathFormat::F32 => ("f32", "f32", "7.0", "sm_80"), + ScalarMathFormat::F64 => ("f64", "f64", "7.0", "sm_80"), + }; + Some(ScalarMathRecipe { + id, + operation_key: format!("scalar.math.{operation_name}.{modifier_symbol}.{format_name}"), + source, + rust_type, + dialect_type, + minimum_ptx, + minimum_sm, + properties: vec!["IntrNoMem"], + ptx_modifiers, + ptx_isa_section, + ptx_isa_url, + force_inline_ptx, + }) +} + +pub(in crate::resolve) fn expand_scalar_math_admission( + admission: &ScalarMathAdmission, +) -> Result> { + ensure!( + admission.runtime_validation == RuntimeValidation::Unexecuted, + "scalar-math runtime may be marked executed only with GPU evidence" + ); + let expected = canonical_scalar_math_variants(); + let actual = admission + .variants + .iter() + .map(|variant| { + ( + variant.format, + variant.operation, + variant.precision, + variant.subnormal, + ) + }) + .collect::>(); + ensure!( + actual == expected, + "compact scalar-math admission must list the canonical 41 variants" + ); + + admission + .variants + .iter() + .map(|variant| { + validate_abi_id(&variant.abi_id)?; + let identity = ( + variant.format, + variant.operation, + variant.precision, + variant.subnormal, + ); + let recipe = scalar_math_recipe(identity) + .context("scalar math is outside the closed recipe set")?; + let libnvvm_evidence_profile = variant + .libnvvm_evidence_profile + .as_ref() + .unwrap_or(&admission.libnvvm_evidence_profile); + ensure!( + !libnvvm_evidence_profile.trim().is_empty(), + "scalar-math libNVVM evidence profile must not be empty" + ); + scalar_math_overlay_record( + recipe, + admission, + identity, + &variant.abi_id, + libnvvm_evidence_profile, + ) + }) + .collect() +} + +pub(in crate::resolve) fn scalar_math_overlay_record( + recipe: ScalarMathRecipe, + admission: &ScalarMathAdmission, + variant: ScalarMathVariant, + abi_id: &str, + libnvvm_evidence_profile: &str, +) -> Result { + let (format, operation, precision, subnormal) = variant; + let ptx_operands = vec![OperandPattern::Register; 2]; // 1 result + 1 operand + let summary = format!( + "Computes unary {} with {} precision.", + scalar_math_operation_name(operation), + scalar_math_precision_name(precision), + ); + let (source, source_record, llvm_symbol, resolved_llvm_symbol, llvm_arguments, llvm_results) = + match &recipe.source { + ScalarMathRecipeSource::Imported { + source_record, + llvm_symbol, + } => ( + None, + Some(source_record.clone()), + Some(llvm_symbol.clone()), + None, + vec![recipe.rust_type.to_owned()], + vec![recipe.rust_type.to_owned()], + ), + // The polymorphic signature tokens mirror the imported + // overloaded record verbatim (anyfloat over anonymous_14), the + // same shape packed_alu uses for llvm.nvvm.fabs. + ScalarMathRecipeSource::ImportedOverloaded { + source_record, + llvm_symbol, + resolved_llvm_symbol, + } => ( + None, + Some(source_record.clone()), + Some(llvm_symbol.clone()), + Some(resolved_llvm_symbol.clone()), + vec!["anonymous_14".to_owned()], + vec!["anyfloat".to_owned()], + ), + ScalarMathRecipeSource::PtxNative { instruction } => ( + Some(IntrinsicSource::PtxNative { + instruction: instruction.clone(), + }), + None, + None, + None, + Vec::new(), + Vec::new(), + ), + }; + Ok(OverlayIntrinsic { + id: recipe.id.clone(), + abi_id: abi_id.into(), + operation_key: recipe.operation_key, + family: "scalar_math".into(), + source, + source_record, + rust_module: "float".into(), + rust_name: recipe.id.clone(), + rust_arguments: vec![recipe.rust_type.into()], + rust_result: recipe.rust_type.into(), + safe: true, + must_use: true, + safe_allowlist_reason: Some("Scalar math has no caller obligations.".into()), + public_rust_path: format!("cuda_intrinsics::float::{}", recipe.id), + compatibility_rust_paths: vec![format!("cuda_device::float::{}", recipe.id)], + dialect_op_type: "ScalarMathOp".into(), + dialect_op_name: "nvvm.scalar_math".into(), + dialect_operands: vec![recipe.dialect_type.into()], + dialect_results: vec![recipe.dialect_type.into()], + llvm_symbol, + resolved_llvm_symbol, + llvm_arguments, + llvm_results, + pure: true, + memory: "none".into(), + convergent: false, + execution_scope: "thread".into(), + minimum_ptx: recipe.minimum_ptx.into(), + minimum_sm: Some(recipe.minimum_sm.into()), + ptx_result: recipe.rust_type.into(), + targets: "all".into(), + ptx_isa_version: "9.3".into(), + ptx_isa_section: recipe.ptx_isa_section.into(), + ptx_isa_url: recipe.ptx_isa_url.into(), + lowering: "generated_scalar_math".into(), + backend_lowerings: [ + ( + IntrinsicBackend::LlvmNvptx, + admission.llvm_evidence_profile.as_str(), + ), + (IntrinsicBackend::LibNvvm, libnvvm_evidence_profile), + ] + .into_iter() + .map( + |(backend, evidence_profile): (IntrinsicBackend, &str)| OverlayBackendLowering { + backend, + mechanism: match backend { + IntrinsicBackend::LlvmNvptx if recipe.force_inline_ptx => { + BackendLoweringMechanism::InlinePtx + } + IntrinsicBackend::LlvmNvptx => BackendLoweringMechanism::TypedNvvm, + IntrinsicBackend::LibNvvm => BackendLoweringMechanism::InlinePtx, + }, + evidence_profile: evidence_profile.to_owned(), + targets: None, + minimum_ptx: Some(recipe.minimum_ptx.into()), + minimum_sm: Some(recipe.minimum_sm.into()), + }, + ) + .collect(), + packed_atomic: None, + redux: None, + vote: None, + active_mask: None, + warp_match: None, + warp_barrier: None, + warp_shuffle: None, + dot_product: None, + packed_alu: None, + integer_minmax: None, + packed_conversion: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: Some(ScalarMath { + format, + operation, + precision, + subnormal, + runtime_validation: admission.runtime_validation, + }), + extended_minmax: None, + cp_async_copy: None, + cp_async_control: None, + cp_async_mbarrier: None, + mbarrier_basic: None, + movmatrix: None, + mbarrier_extended: None, + register_mma: None, + sparse_mma: None, + prmt: None, + cluster_barrier: None, + wgmma_control: None, + special_register: None, + debug_control: None, + cluster_memory: None, + clc: None, + tma: None, + tcgen05: None, + ldmatrix_variant: None, + ldmatrix_safety: None, + ldmatrix_adapter: None, + selected_address_space: None, + expected_ptx: InstructionPattern { + mnemonic: scalar_math_operation_name(operation).into(), + modifiers: recipe.ptx_modifiers, + operands: ptx_operands, + }, + summary, + }) +} + +pub(in crate::resolve) fn validate_scalar_math_policy( + policy: &OverlayIntrinsic, + declaration: Option<&ImportedIntrinsic>, +) -> Result<()> { + let math = policy + .scalar_math + .as_ref() + .with_context(|| format!("{} has no scalar-math contract", policy.id))?; + let variant = (math.format, math.operation, math.precision, math.subnormal); + let recipe = scalar_math_recipe(variant) + .with_context(|| format!("{} is outside the closed scalar-math recipe", policy.id))?; + ensure!( + math.runtime_validation == RuntimeValidation::Unexecuted, + "{} scalar-math runtime may be executed only with GPU evidence", + policy.id + ); + let signature = vec![recipe.rust_type.to_owned()]; + let source_matches = match &recipe.source { + ScalarMathRecipeSource::Imported { + source_record, + llvm_symbol, + } => { + policy.source.is_none() + && policy.source_record.as_deref() == Some(source_record.as_str()) + && policy.llvm_symbol.as_deref() == Some(llvm_symbol.as_str()) + && policy.resolved_llvm_symbol.is_none() + && policy.llvm_arguments == signature + && policy.llvm_results == [recipe.rust_type] + } + ScalarMathRecipeSource::ImportedOverloaded { + source_record, + llvm_symbol, + resolved_llvm_symbol, + } => { + policy.source.is_none() + && policy.source_record.as_deref() == Some(source_record.as_str()) + && policy.llvm_symbol.as_deref() == Some(llvm_symbol.as_str()) + && policy.resolved_llvm_symbol.as_deref() == Some(resolved_llvm_symbol.as_str()) + && policy.llvm_arguments == ["anonymous_14"] + && policy.llvm_results == ["anyfloat"] + } + ScalarMathRecipeSource::PtxNative { instruction } => { + policy.source + == Some(IntrinsicSource::PtxNative { + instruction: instruction.clone(), + }) + && policy.source_record.is_none() + && policy.llvm_symbol.is_none() + && policy.resolved_llvm_symbol.is_none() + && policy.llvm_arguments.is_empty() + && policy.llvm_results.is_empty() + } + }; + ensure!( + policy.id == recipe.id && policy.operation_key == recipe.operation_key && source_matches, + "{} scalar-math identity or LLVM source changed", + policy.id + ); + ensure!( + declaration.is_none() == matches!(recipe.source, ScalarMathRecipeSource::PtxNative { .. }), + "{} scalar-math source kind and imported declaration disagree", + policy.id + ); + if let Some(declaration) = declaration { + let expected_properties = recipe + .properties + .iter() + .map(|property| (*property).to_owned()) + .collect::>(); + ensure!( + declaration.properties == expected_properties, + "{} imported scalar-math properties changed", + policy.id + ); + ensure!( + declaration.selections.len() <= 1, + "{} gained an unreviewed scalar-math selection", + policy.id + ); + if let Some(direct) = declaration.selections.first() { + ensure!( + direct.predicates.is_empty() && direct.constraints.is_empty(), + "{} direct scalar-math selection changed", + policy.id + ); + } + } + ensure!( + policy.rust_module == "float" + && policy.rust_name == recipe.id + && policy.rust_arguments == signature + && policy.rust_result == recipe.rust_type + && policy.safe + && policy.must_use + && policy.compatibility_rust_paths == [format!("cuda_device::float::{}", recipe.id)] + && policy.dialect_op_type == "ScalarMathOp" + && policy.dialect_op_name == "nvvm.scalar_math" + && policy.dialect_operands == [recipe.dialect_type] + && policy.dialect_results == [recipe.dialect_type] + && policy.lowering == "generated_scalar_math", + "{} changed its scalar-math API, carrier, or lowering", + policy.id + ); + ensure!( + policy.pure + && policy.memory == "none" + && !policy.convergent + && policy.execution_scope == "thread" + && policy.minimum_ptx == recipe.minimum_ptx + && policy.minimum_sm.as_deref() == Some(recipe.minimum_sm) + && policy.ptx_result == recipe.rust_type + && policy.targets == "all" + && policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section == recipe.ptx_isa_section + && policy.ptx_isa_url == recipe.ptx_isa_url, + "{} scalar-math effects, provenance, or target floor changed", + policy.id + ); + ensure!( + policy.expected_ptx + == InstructionPattern { + mnemonic: scalar_math_operation_name(math.operation).into(), + modifiers: recipe.ptx_modifiers.clone(), + operands: vec![OperandPattern::Register; 2], + }, + "{} expected scalar-math PTX changed", + policy.id + ); + let llvm_mechanism = if recipe.force_inline_ptx { + BackendLoweringMechanism::InlinePtx + } else { + BackendLoweringMechanism::TypedNvvm + }; + let expected_backends = [ + ( + IntrinsicBackend::LlvmNvptx, + llvm_mechanism, + recipe.minimum_sm, + ), + ( + IntrinsicBackend::LibNvvm, + BackendLoweringMechanism::InlinePtx, + recipe.minimum_sm, + ), + ]; + ensure!( + policy.backend_lowerings.len() == 2 + && expected_backends + .into_iter() + .all(|(backend, mechanism, minimum_sm)| { + policy.backend_lowerings.iter().any(|lowering| { + lowering.backend == backend + && lowering.mechanism == mechanism + && lowering.minimum_ptx.as_deref() == Some(recipe.minimum_ptx) + && lowering.minimum_sm.as_deref() == Some(minimum_sm) + && !lowering.evidence_profile.trim().is_empty() + }) + }), + "{} has the wrong reviewed scalar-math backend routes (expected {} / {})", + policy.id, + recipe.minimum_ptx, + recipe.minimum_sm + ); + if let Some(direct) = declaration.and_then(|declaration| declaration.selections.first()) { + validate_selected_target_predicates(policy, direct)?; + } + ensure_no_other_family_contract(policy, "scalar math")?; + Ok(()) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/special_register.rs b/crates/cuda-intrinsics-gen/src/resolve/families/special_register.rs new file mode 100644 index 0000000000..be45a4cdab --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/special_register.rs @@ -0,0 +1,1459 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, ClusterSregAdmission, ImportedIntrinsic, IntrinsicBackend, + IntrinsicSource, OverlayBackendLowering, OverlayIntrinsic, RuntimeValidation, SpecialRegister, + SpecialRegisterAdmission, SpecialRegisterKind, SpecialRegisterLlvmExclusion, + SpecialRegisterLlvmExclusionReason, SpecialRegisterObservation, + SpecialRegisterOutputConstraint, SpecialRegisterPtxType, SpecialRegisterWidth, +}; +use crate::ptx::{InstructionPattern, OperandPattern}; +use anyhow::{Context, Result, bail, ensure}; +use std::collections::BTreeMap; + +use crate::resolve::guards::*; + +pub(in crate::resolve) const REVIEWED_SPECIAL_REGISTERS: [SpecialRegisterKind; 12] = [ + SpecialRegisterKind::Clock, + SpecialRegisterKind::Clock64, + SpecialRegisterKind::Globaltimer, + SpecialRegisterKind::Envreg1, + SpecialRegisterKind::Envreg2, + SpecialRegisterKind::Smid, + SpecialRegisterKind::Nsmid, + SpecialRegisterKind::Gridid, + SpecialRegisterKind::Warpid, + SpecialRegisterKind::Nwarpid, + SpecialRegisterKind::DynamicSmemSize, + SpecialRegisterKind::TotalSmemSize, +]; + +#[derive(Clone, Copy)] +pub(in crate::resolve) struct SpecialRegisterRecipe { + kind: SpecialRegisterKind, + id: &'static str, + operation_key: &'static str, + source_record: Option<&'static str>, + llvm_symbol: Option<&'static str>, + rust_module: &'static str, + compatibility_paths: &'static [&'static str], + dialect_op_type: &'static str, + dialect_op_name: &'static str, + register_spelling: &'static str, + observation: SpecialRegisterObservation, + result_width: SpecialRegisterWidth, + ptx_type: SpecialRegisterPtxType, + output_constraint: SpecialRegisterOutputConstraint, + llvm_mechanism: BackendLoweringMechanism, + libnvvm_mechanism: BackendLoweringMechanism, + minimum_ptx: &'static str, + minimum_sm: Option<&'static str>, + execution_scope: &'static str, + ptx_isa_section: &'static str, + ptx_isa_url: &'static str, + selection_record: Option<&'static str>, + selection_asm: Option<&'static str>, + summary: &'static str, +} + +pub(in crate::resolve) fn special_register_recipe( + kind: SpecialRegisterKind, +) -> SpecialRegisterRecipe { + use BackendLoweringMechanism::{InlinePtx, TypedNvvm}; + use SpecialRegisterKind::*; + use SpecialRegisterObservation::{StablePure, VolatileObservation}; + use SpecialRegisterOutputConstraint::{Register32, Register64}; + use SpecialRegisterPtxType::{B32, U32, U64}; + use SpecialRegisterWidth::{B32 as Width32, B64 as Width64}; + + match kind { + Clock => SpecialRegisterRecipe { + kind, + id: "clock", + operation_key: "debug.clock", + source_record: Some("int_nvvm_read_ptx_sreg_clock"), + llvm_symbol: Some("llvm.nvvm.read.ptx.sreg.clock"), + rust_module: "debug", + compatibility_paths: &["cuda_device::debug::clock"], + dialect_op_type: "ReadPtxSregClockOp", + dialect_op_name: "nvvm.read_ptx_sreg_clock", + register_spelling: "%clock", + observation: VolatileObservation, + result_width: Width32, + ptx_type: U32, + output_constraint: Register32, + llvm_mechanism: TypedNvvm, + libnvvm_mechanism: TypedNvvm, + minimum_ptx: "1.0", + minimum_sm: None, + execution_scope: "sm", + ptx_isa_section: "10.23 Special Registers: %clock, %clock_hi", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-clock-clock-hi", + selection_record: Some("SREG_CLOCK"), + selection_asm: Some("mov.u32 \t$d, %clock;"), + summary: "Samples the current SM's 32-bit clock counter.", + }, + Clock64 => SpecialRegisterRecipe { + kind, + id: "clock64", + operation_key: "debug.clock64", + source_record: Some("int_nvvm_read_ptx_sreg_clock64"), + llvm_symbol: Some("llvm.nvvm.read.ptx.sreg.clock64"), + rust_module: "debug", + compatibility_paths: &["cuda_device::debug::clock64"], + dialect_op_type: "ReadPtxSregClock64Op", + dialect_op_name: "nvvm.read_ptx_sreg_clock64", + register_spelling: "%clock64", + observation: VolatileObservation, + result_width: Width64, + ptx_type: U64, + output_constraint: Register64, + llvm_mechanism: TypedNvvm, + libnvvm_mechanism: TypedNvvm, + minimum_ptx: "2.0", + minimum_sm: Some("sm_20"), + execution_scope: "sm", + ptx_isa_section: "10.24 Special Registers: %clock64", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-clock64", + selection_record: Some("SREG_CLOCK64"), + selection_asm: Some("mov.u64 \t$d, %clock64;"), + summary: "Samples the current SM's 64-bit clock counter.", + }, + Globaltimer => SpecialRegisterRecipe { + kind, + id: "globaltimer", + operation_key: "debug.global_timer", + source_record: Some("int_nvvm_read_ptx_sreg_globaltimer"), + llvm_symbol: Some("llvm.nvvm.read.ptx.sreg.globaltimer"), + rust_module: "debug", + compatibility_paths: &["cuda_device::debug::globaltimer"], + dialect_op_type: "ReadPtxSregGlobaltimerOp", + dialect_op_name: "nvvm.read_ptx_sreg_globaltimer", + register_spelling: "%globaltimer", + observation: VolatileObservation, + result_width: Width64, + ptx_type: U64, + output_constraint: Register64, + llvm_mechanism: TypedNvvm, + libnvvm_mechanism: InlinePtx, + minimum_ptx: "3.1", + minimum_sm: Some("sm_30"), + execution_scope: "device", + ptx_isa_section: "10.28 Special Registers: %globaltimer, %globaltimer_lo, %globaltimer_hi", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-globaltimer", + selection_record: Some("SREG_GLOBALTIMER"), + selection_asm: Some("mov.u64 \t$d, %globaltimer;"), + summary: "Samples the device-wide 64-bit global timer.", + }, + Envreg1 => SpecialRegisterRecipe { + kind, + id: "envreg1", + operation_key: "grid.environment_register.1", + source_record: Some("int_nvvm_read_ptx_sreg_envreg1"), + llvm_symbol: Some("llvm.nvvm.read.ptx.sreg.envreg1"), + rust_module: "grid", + compatibility_paths: &["cuda_device::grid::envreg1"], + dialect_op_type: "ReadPtxSregEnvReg1Op", + dialect_op_name: "nvvm.read_ptx_sreg_envreg1", + register_spelling: "%envreg1", + observation: StablePure, + result_width: Width32, + ptx_type: B32, + output_constraint: Register32, + llvm_mechanism: TypedNvvm, + libnvvm_mechanism: TypedNvvm, + minimum_ptx: "2.1", + minimum_sm: None, + execution_scope: "grid", + ptx_isa_section: "10.27 Special Registers: %envreg<32>", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-envreg", + selection_record: None, + selection_asm: None, + summary: "Reads PTX environment register 1.", + }, + Envreg2 => SpecialRegisterRecipe { + kind, + id: "envreg2", + operation_key: "grid.environment_register.2", + source_record: Some("int_nvvm_read_ptx_sreg_envreg2"), + llvm_symbol: Some("llvm.nvvm.read.ptx.sreg.envreg2"), + rust_module: "grid", + compatibility_paths: &["cuda_device::grid::envreg2"], + dialect_op_type: "ReadPtxSregEnvReg2Op", + dialect_op_name: "nvvm.read_ptx_sreg_envreg2", + register_spelling: "%envreg2", + observation: StablePure, + result_width: Width32, + ptx_type: B32, + output_constraint: Register32, + llvm_mechanism: TypedNvvm, + libnvvm_mechanism: TypedNvvm, + minimum_ptx: "2.1", + minimum_sm: None, + execution_scope: "grid", + ptx_isa_section: "10.27 Special Registers: %envreg<32>", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-envreg", + selection_record: None, + selection_asm: None, + summary: "Reads PTX environment register 2.", + }, + Smid => SpecialRegisterRecipe { + kind, + id: "smid", + operation_key: "execution.sm_identifier", + source_record: Some("int_nvvm_read_ptx_sreg_smid"), + llvm_symbol: Some("llvm.nvvm.read.ptx.sreg.smid"), + rust_module: "thread", + compatibility_paths: &["cuda_device::thread::smid", "cuda_device::smid"], + dialect_op_type: "ReadPtxSregSmIdOp", + dialect_op_name: "nvvm.read_ptx_sreg_smid", + register_spelling: "%smid", + observation: VolatileObservation, + result_width: Width32, + ptx_type: U32, + output_constraint: Register32, + llvm_mechanism: InlinePtx, + libnvvm_mechanism: InlinePtx, + minimum_ptx: "1.3", + minimum_sm: None, + execution_scope: "thread", + ptx_isa_section: "10.8 Special Registers: %smid", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-smid", + selection_record: Some("SREG_SMID"), + selection_asm: Some("mov.u32 \t$d, %smid;"), + summary: "Samples the SM currently executing this thread.", + }, + Nsmid => SpecialRegisterRecipe { + kind, + id: "nsmid", + operation_key: "execution.sm_identifier_bound", + source_record: Some("int_nvvm_read_ptx_sreg_nsmid"), + llvm_symbol: Some("llvm.nvvm.read.ptx.sreg.nsmid"), + rust_module: "thread", + compatibility_paths: &["cuda_device::thread::nsmid", "cuda_device::nsmid"], + dialect_op_type: "ReadPtxSregNsmIdOp", + dialect_op_name: "nvvm.read_ptx_sreg_nsmid", + register_spelling: "%nsmid", + observation: StablePure, + result_width: Width32, + ptx_type: U32, + output_constraint: Register32, + llvm_mechanism: TypedNvvm, + libnvvm_mechanism: TypedNvvm, + minimum_ptx: "2.0", + minimum_sm: Some("sm_20"), + execution_scope: "device", + ptx_isa_section: "10.9 Special Registers: %nsmid", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-nsmid", + selection_record: Some("SREG_NSMID"), + selection_asm: Some("mov.u32 \t$d, %nsmid;"), + summary: "Returns the upper bound for SM identifiers.", + }, + Gridid => SpecialRegisterRecipe { + kind, + id: "gridid", + operation_key: "launch.grid_identifier", + source_record: None, + llvm_symbol: None, + rust_module: "thread", + compatibility_paths: &["cuda_device::thread::gridid", "cuda_device::gridid"], + dialect_op_type: "ReadPtxSregGridIdOp", + dialect_op_name: "nvvm.read_ptx_sreg_gridid", + register_spelling: "%gridid", + observation: StablePure, + result_width: Width64, + ptx_type: U64, + output_constraint: Register64, + llvm_mechanism: InlinePtx, + libnvvm_mechanism: InlinePtx, + minimum_ptx: "3.0", + minimum_sm: Some("sm_30"), + execution_scope: "grid", + ptx_isa_section: "10.10 Special Registers: %gridid", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-gridid", + selection_record: None, + selection_asm: None, + summary: "Returns the full 64-bit temporal grid identifier.", + }, + Warpid => SpecialRegisterRecipe { + kind, + id: "warpid", + operation_key: "warp.hardware_identifier", + source_record: Some("int_nvvm_read_ptx_sreg_warpid"), + llvm_symbol: Some("llvm.nvvm.read.ptx.sreg.warpid"), + rust_module: "warp", + compatibility_paths: &["cuda_device::warp::warpid"], + dialect_op_type: "ReadPtxSregWarpIdOp", + dialect_op_name: "nvvm.read_ptx_sreg_warpid", + register_spelling: "%warpid", + observation: VolatileObservation, + result_width: Width32, + ptx_type: U32, + output_constraint: Register32, + llvm_mechanism: InlinePtx, + libnvvm_mechanism: InlinePtx, + minimum_ptx: "1.3", + minimum_sm: None, + execution_scope: "cta", + ptx_isa_section: "10.4 Special Registers: %warpid", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-warpid", + selection_record: Some("SREG_WARPID"), + selection_asm: Some("mov.u32 \t$d, %warpid;"), + summary: "Samples the hardware warp currently executing this thread.", + }, + Nwarpid => SpecialRegisterRecipe { + kind, + id: "nwarpid", + operation_key: "warp.hardware_identifier_bound", + source_record: Some("int_nvvm_read_ptx_sreg_nwarpid"), + llvm_symbol: Some("llvm.nvvm.read.ptx.sreg.nwarpid"), + rust_module: "warp", + compatibility_paths: &["cuda_device::warp::nwarpid"], + dialect_op_type: "ReadPtxSregNwarpIdOp", + dialect_op_name: "nvvm.read_ptx_sreg_nwarpid", + register_spelling: "%nwarpid", + observation: StablePure, + result_width: Width32, + ptx_type: U32, + output_constraint: Register32, + llvm_mechanism: TypedNvvm, + libnvvm_mechanism: TypedNvvm, + minimum_ptx: "2.0", + minimum_sm: Some("sm_20"), + execution_scope: "cta", + ptx_isa_section: "10.5 Special Registers: %nwarpid", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-nwarpid", + selection_record: Some("SREG_NWARPID"), + selection_asm: Some("mov.u32 \t$d, %nwarpid;"), + summary: "Returns the upper bound for hardware warp identifiers.", + }, + DynamicSmemSize => SpecialRegisterRecipe { + kind, + id: "dynamic_smem_size", + operation_key: "shared.dynamic_size", + source_record: Some("int_nvvm_read_ptx_sreg_dynamic_smem_size"), + llvm_symbol: Some("llvm.nvvm.read.ptx.sreg.dynamic_smem_size"), + rust_module: "shared", + compatibility_paths: &["cuda_device::shared::dynamic_smem_size"], + dialect_op_type: "ReadPtxSregDynamicSmemSizeOp", + dialect_op_name: "nvvm.read_ptx_sreg_dynamic_smem_size", + register_spelling: "%dynamic_smem_size", + observation: StablePure, + result_width: Width32, + ptx_type: U32, + output_constraint: Register32, + llvm_mechanism: InlinePtx, + libnvvm_mechanism: InlinePtx, + minimum_ptx: "4.1", + minimum_sm: Some("sm_20"), + execution_scope: "cta", + ptx_isa_section: "10.32 Special Registers: %dynamic_smem_size", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-dynamic-smem-size", + selection_record: Some("INT_PTX_SREG_DYNAMIC_SMEM_SIZE"), + selection_asm: Some("mov.u32 \t$d, %dynamic_smem_size;"), + summary: "Returns the launch-time dynamic shared-memory size in bytes.", + }, + TotalSmemSize => SpecialRegisterRecipe { + kind, + id: "total_smem_size", + operation_key: "shared.total_size", + source_record: Some("int_nvvm_read_ptx_sreg_total_smem_size"), + llvm_symbol: Some("llvm.nvvm.read.ptx.sreg.total_smem_size"), + rust_module: "shared", + compatibility_paths: &["cuda_device::shared::total_smem_size"], + dialect_op_type: "ReadPtxSregTotalSmemSizeOp", + dialect_op_name: "nvvm.read_ptx_sreg_total_smem_size", + register_spelling: "%total_smem_size", + observation: StablePure, + result_width: Width32, + ptx_type: U32, + output_constraint: Register32, + llvm_mechanism: InlinePtx, + libnvvm_mechanism: InlinePtx, + minimum_ptx: "4.1", + minimum_sm: Some("sm_20"), + execution_scope: "cta", + ptx_isa_section: "10.30 Special Registers: %total_smem_size", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-total-smem-size", + selection_record: Some("INT_PTX_SREG_TOTAL_SMEM_SIZE"), + selection_asm: Some("mov.u32 \t$d, %total_smem_size;"), + summary: "Returns the total user shared-memory allocation in bytes.", + }, + } +} + +pub(in crate::resolve) fn special_register_ptx_type( + ptx_type: SpecialRegisterPtxType, +) -> &'static str { + match ptx_type { + SpecialRegisterPtxType::B32 => "b32", + SpecialRegisterPtxType::U32 => "u32", + SpecialRegisterPtxType::U64 => "u64", + } +} + +pub(in crate::resolve) fn special_register_backend_floor( + recipe: SpecialRegisterRecipe, + backend: IntrinsicBackend, +) -> (Option<&'static str>, Option<&'static str>) { + match backend { + IntrinsicBackend::LlvmNvptx => { + let minimum_ptx = if matches!(recipe.minimum_ptx, "4.1") { + "4.1" + } else { + "3.2" + }; + let minimum_sm = if recipe.minimum_sm == Some("sm_30") { + "sm_30" + } else { + "sm_20" + }; + (Some(minimum_ptx), Some(minimum_sm)) + } + IntrinsicBackend::LibNvvm => (None, Some("sm_75")), + } +} + +pub(in crate::resolve) fn special_register_contract( + recipe: SpecialRegisterRecipe, +) -> SpecialRegister { + let llvm_exclusion = + (recipe.kind == SpecialRegisterKind::Gridid).then(|| SpecialRegisterLlvmExclusion { + source_record: "int_nvvm_read_ptx_sreg_gridid".into(), + llvm_symbol: "llvm.nvvm.read.ptx.sreg.gridid".into(), + imported_result_width: SpecialRegisterWidth::B32, + reason: SpecialRegisterLlvmExclusionReason::ResultWidthMismatch, + }); + SpecialRegister { + register: recipe.kind, + observation: recipe.observation, + result_width: recipe.result_width, + ptx_type: recipe.ptx_type, + output_constraint: recipe.output_constraint, + llvm_exclusion, + } +} + +pub(in crate::resolve) fn expand_special_register_admission( + admission: &SpecialRegisterAdmission, +) -> Result> { + ensure!( + admission.runtime_validation == RuntimeValidation::Unexecuted, + "special-register runtime validation may be marked executed only with GPU evidence" + ); + ensure!( + !admission.llvm_evidence_profile.trim().is_empty() + && !admission.libnvvm_evidence_profile.trim().is_empty(), + "compact special-register admission requires both backend evidence profiles" + ); + ensure!( + admission.registers == REVIEWED_SPECIAL_REGISTERS + && admission.product_count == REVIEWED_SPECIAL_REGISTERS.len(), + "compact special-register admission must list the canonical 12 registers exactly once and in order" + ); + + admission + .registers + .iter() + .copied() + .map(|kind| { + let recipe = special_register_recipe(kind); + let width = recipe.result_width.bits(); + let rust_result = format!("u{width}"); + let dialect_result = format!("i{width}"); + let source = match recipe.source_record { + Some(_) => None, + None => Some(IntrinsicSource::PtxNative { + instruction: format!( + "mov.{} {}", + special_register_ptx_type(recipe.ptx_type), + recipe.register_spelling + ), + }), + }; + Ok(OverlayIntrinsic { + id: recipe.id.into(), + abi_id: String::new(), + operation_key: recipe.operation_key.into(), + family: "sreg".into(), + source, + source_record: recipe.source_record.map(str::to_owned), + rust_module: recipe.rust_module.into(), + rust_name: recipe.id.into(), + rust_arguments: vec![], + rust_result: rust_result.clone(), + safe: true, + must_use: false, + safe_allowlist_reason: Some( + "reading this special register has no caller obligations.".into(), + ), + public_rust_path: format!("cuda_intrinsics::{}::{}", recipe.rust_module, recipe.id), + compatibility_rust_paths: recipe + .compatibility_paths + .iter() + .map(|path| (*path).into()) + .collect(), + dialect_op_type: recipe.dialect_op_type.into(), + dialect_op_name: recipe.dialect_op_name.into(), + dialect_operands: vec![], + dialect_results: vec![dialect_result.clone()], + llvm_symbol: recipe.llvm_symbol.map(str::to_owned), + resolved_llvm_symbol: None, + llvm_arguments: vec![], + llvm_results: recipe + .source_record + .map(|_| dialect_result) + .into_iter() + .collect(), + pure: recipe.observation == SpecialRegisterObservation::StablePure, + memory: if matches!( + recipe.kind, + SpecialRegisterKind::Clock + | SpecialRegisterKind::Clock64 + | SpecialRegisterKind::Globaltimer + ) { + "inaccessible_read_write".into() + } else { + "none".into() + }, + convergent: false, + execution_scope: recipe.execution_scope.into(), + minimum_ptx: recipe.minimum_ptx.into(), + minimum_sm: recipe.minimum_sm.map(str::to_owned), + ptx_result: rust_result, + targets: "all".into(), + ptx_isa_version: "9.3".into(), + ptx_isa_section: recipe.ptx_isa_section.into(), + ptx_isa_url: recipe.ptx_isa_url.into(), + lowering: "generated_special_register".into(), + backend_lowerings: [ + ( + IntrinsicBackend::LlvmNvptx, + recipe.llvm_mechanism, + &admission.llvm_evidence_profile, + ), + ( + IntrinsicBackend::LibNvvm, + recipe.libnvvm_mechanism, + &admission.libnvvm_evidence_profile, + ), + ] + .into_iter() + .map( + |(backend, mechanism, evidence_profile)| OverlayBackendLowering { + minimum_ptx: special_register_backend_floor(recipe, backend) + .0 + .map(str::to_owned), + minimum_sm: special_register_backend_floor(recipe, backend) + .1 + .map(str::to_owned), + backend, + mechanism, + evidence_profile: evidence_profile.clone(), + targets: None, + }, + ) + .collect(), + packed_atomic: None, + redux: None, + vote: None, + active_mask: None, + warp_match: None, + warp_barrier: None, + warp_shuffle: None, + dot_product: None, + packed_alu: None, + integer_minmax: None, + packed_conversion: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cp_async_copy: None, + cp_async_control: None, + cp_async_mbarrier: None, + mbarrier_basic: None, + movmatrix: None, + mbarrier_extended: None, + register_mma: None, + sparse_mma: None, + prmt: None, + cluster_barrier: None, + wgmma_control: None, + special_register: Some(special_register_contract(recipe)), + debug_control: None, + cluster_memory: None, + clc: None, + tma: None, + tcgen05: None, + ldmatrix_variant: None, + ldmatrix_safety: None, + ldmatrix_adapter: None, + selected_address_space: None, + expected_ptx: InstructionPattern { + mnemonic: "mov".into(), + modifiers: vec![special_register_ptx_type(recipe.ptx_type).into()], + operands: vec![ + OperandPattern::Register, + OperandPattern::Exact { + value: recipe.register_spelling.into(), + }, + ], + }, + summary: recipe.summary.into(), + }) + }) + .collect() +} + +pub(in crate::resolve) fn validate_special_register_policy( + policy: &OverlayIntrinsic, + source: &IntrinsicSource, + declaration: Option<&ImportedIntrinsic>, +) -> Result<()> { + let special = policy + .special_register + .as_ref() + .with_context(|| format!("{} has no closed special-register contract", policy.id))?; + let recipe = special_register_recipe(special.register); + ensure!( + special == &special_register_contract(recipe), + "{} special-register width, PTX type, constraint, observation, or LLVM exclusion changed", + policy.id + ); + let expected_source = match recipe.source_record { + Some(source_record) => IntrinsicSource::LlvmImported { + source_record: source_record.into(), + }, + None => IntrinsicSource::PtxNative { + instruction: format!( + "mov.{} {}", + special_register_ptx_type(recipe.ptx_type), + recipe.register_spelling + ), + }, + }; + ensure!( + policy.id == recipe.id + && policy.operation_key == recipe.operation_key + && source == &expected_source + && policy.source_record.as_deref() == recipe.source_record + && policy.llvm_symbol.as_deref() == recipe.llvm_symbol + && policy.resolved_llvm_symbol.is_none() + && policy.llvm_arguments.is_empty(), + "{} special-register identity or source changed", + policy.id + ); + let width = recipe.result_width.bits(); + let rust_result = format!("u{width}"); + let dialect_result = format!("i{width}"); + let expected_llvm_results = recipe + .source_record + .map(|_| dialect_result.clone()) + .into_iter() + .collect::>(); + ensure!( + policy.rust_module == recipe.rust_module + && policy.rust_name == recipe.id + && policy.rust_arguments.is_empty() + && policy.rust_result == rust_result + && policy.safe + && !policy.must_use + && policy.public_rust_path + == format!("cuda_intrinsics::{}::{}", recipe.rust_module, recipe.id) + && policy.compatibility_rust_paths + == recipe + .compatibility_paths + .iter() + .map(|path| (*path).to_owned()) + .collect::>(), + "{} special-register Rust API changed", + policy.id + ); + ensure!( + policy.dialect_op_type == recipe.dialect_op_type + && policy.dialect_op_name == recipe.dialect_op_name + && policy.dialect_operands.is_empty() + && policy.dialect_results == [dialect_result.as_str()] + && policy.llvm_results == expected_llvm_results + && policy.lowering == "generated_special_register", + "{} special-register carrier, result width, or lowering changed", + policy.id + ); + let expected_pure = recipe.observation == SpecialRegisterObservation::StablePure; + let expected_memory = if matches!( + recipe.kind, + SpecialRegisterKind::Clock + | SpecialRegisterKind::Clock64 + | SpecialRegisterKind::Globaltimer + ) { + "inaccessible_read_write" + } else { + "none" + }; + ensure!( + policy.pure == expected_pure + && policy.memory == expected_memory + && !policy.convergent + && policy.execution_scope == recipe.execution_scope, + "{} special-register observation or effects changed", + policy.id + ); + ensure!( + policy.minimum_ptx == recipe.minimum_ptx + && policy.minimum_sm.as_deref() == recipe.minimum_sm + && policy.ptx_result == rust_result + && policy.targets == "all" + && policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section == recipe.ptx_isa_section + && policy.ptx_isa_url == recipe.ptx_isa_url, + "{} special-register target floor or PTX provenance changed", + policy.id + ); + ensure!( + policy.expected_ptx + == InstructionPattern { + mnemonic: "mov".into(), + modifiers: vec![special_register_ptx_type(recipe.ptx_type).into()], + operands: vec![ + OperandPattern::Register, + OperandPattern::Exact { + value: recipe.register_spelling.into(), + }, + ], + }, + "{} special-register PTX shape changed", + policy.id + ); + let expected_routes = [ + ( + IntrinsicBackend::LlvmNvptx, + recipe.llvm_mechanism, + special_register_backend_floor(recipe, IntrinsicBackend::LlvmNvptx), + ), + ( + IntrinsicBackend::LibNvvm, + recipe.libnvvm_mechanism, + special_register_backend_floor(recipe, IntrinsicBackend::LibNvvm), + ), + ]; + ensure!( + policy.backend_lowerings.len() == expected_routes.len() + && policy.backend_lowerings.iter().zip(expected_routes).all( + |(actual, (backend, mechanism, (minimum_ptx, minimum_sm)))| { + actual.backend == backend + && actual.mechanism == mechanism + && !actual.evidence_profile.trim().is_empty() + && actual.minimum_ptx.as_deref() == minimum_ptx + && actual.minimum_sm.as_deref() == minimum_sm + } + ), + "{} special-register backend routes changed", + policy.id + ); + ensure!( + policy.ldmatrix_variant.is_none() + && policy.ldmatrix_safety.is_none() + && policy.ldmatrix_adapter.is_none() + && policy.packed_atomic.is_none() + && policy.redux.is_none() + && policy.vote.is_none() + && policy.active_mask.is_none() + && policy.warp_match.is_none() + && policy.warp_barrier.is_none() + && policy.warp_shuffle.is_none() + && policy.dot_product.is_none() + && policy.packed_alu.is_none() + && policy.packed_conversion.is_none() + && policy.cp_async_copy.is_none() + && policy.cp_async_control.is_none() + && policy.cp_async_mbarrier.is_none() + && policy.mbarrier_basic.is_none() + && policy.register_mma.is_none() + && policy.sparse_mma.is_none() + && policy.prmt.is_none() + && policy.selected_address_space.is_none(), + "{} mixes another generated-family contract with a special register", + policy.id + ); + + match (recipe.source_record, declaration) { + (None, None) => {} + (Some(_), Some(declaration)) => { + let timer = matches!( + recipe.kind, + SpecialRegisterKind::Clock + | SpecialRegisterKind::Clock64 + | SpecialRegisterKind::Globaltimer + ); + let expected_properties: &[&str] = if timer { + &[ + "IntrInaccessibleMemOnly", + "IntrNoCallback", + "IntrNoFree", + "IntrWillReturn", + "NoUndef", + ] + } else { + &["IntrNoMem", "IntrSpeculatable", "NoUndef"] + }; + ensure!( + declaration.arguments.is_empty() + && declaration.results == [dialect_result.as_str()] + && declaration.properties + == expected_properties + .iter() + .map(|property| (*property).to_owned()) + .collect::>() + && if timer { + declaration + .classes + .iter() + .any(|class| class == "PTXReadNCSRegIntrinsic") + } else { + declaration + .classes + .iter() + .any(|class| class == "NVVMPureIntrinsic") + }, + "{} imported special-register signature, class, or properties changed", + policy.id + ); + match (recipe.selection_record, recipe.selection_asm) { + (None, None) => ensure!( + declaration.selections.is_empty(), + "{} selectionless environment-register contract changed", + policy.id + ), + (Some(selection_record), Some(selection_asm)) => ensure!( + declaration.selections.len() == 1 + && declaration.selections[0].source_record == selection_record + && declaration.selections[0].asm == selection_asm + && declaration.selections[0].predicates.is_empty() + && declaration.selections[0].constraints.is_empty(), + "{} imported special-register selection changed", + policy.id + ), + _ => unreachable!("closed special-register selection recipe"), + } + } + _ => bail!( + "{} special-register source and imported declaration disagree", + policy.id + ), + } + Ok(()) +} + +pub(in crate::resolve) fn validate_special_register_llvm_exclusion( + policy: &OverlayIntrinsic, + imported_by_record: &BTreeMap<&str, &ImportedIntrinsic>, +) -> Result<()> { + let Some(exclusion) = policy + .special_register + .as_ref() + .and_then(|special| special.llvm_exclusion.as_ref()) + else { + return Ok(()); + }; + let declaration = imported_by_record + .get(exclusion.source_record.as_str()) + .with_context(|| { + format!( + "{} excludes missing imported LLVM record {}", + policy.id, exclusion.source_record + ) + })?; + ensure!( + policy.id == "gridid" + && exclusion.reason == SpecialRegisterLlvmExclusionReason::ResultWidthMismatch + && exclusion.imported_result_width == SpecialRegisterWidth::B32 + && declaration.llvm_name == exclusion.llvm_symbol + && declaration.arguments.is_empty() + && declaration.results == ["i32"] + && declaration.properties == ["IntrNoMem", "IntrSpeculatable", "NoUndef"] + && declaration.selections.len() == 1 + && declaration.selections[0].source_record == "SREG_GRIDID" + && declaration.selections[0].asm == "mov.u32 \t$d, %gridid;" + && declaration.selections[0].predicates.is_empty() + && declaration.selections[0].constraints.is_empty() + && policy.rust_result == "u64" + && policy.dialect_results == ["i64"], + "{} LLVM exclusion no longer proves the reviewed i32-to-u64 gridid width mismatch", + policy.id + ); + Ok(()) +} + +pub(in crate::resolve) fn validate_sreg_policy( + policy: &OverlayIntrinsic, + source: &IntrinsicSource, + declaration: Option<&ImportedIntrinsic>, +) -> Result<()> { + if policy.special_register.is_some() { + return validate_special_register_policy(policy, source, declaration); + } + let declaration = declaration.context("sreg requires imported LLVM declaration")?; + ensure!( + policy.rust_arguments.is_empty() && policy.llvm_arguments.is_empty(), + "{} is not a zero-operand intrinsic; the sreg recipe cannot lower it", + policy.id + ); + ensure!( + matches!(policy.rust_result.as_str(), "u32" | "u64"), + "{} has unsupported raw scalar result {}", + policy.id, + policy.rust_result + ); + let expected_llvm_result = match policy.rust_result.as_str() { + "u32" => "i32", + "u64" => "i64", + _ => unreachable!(), + }; + ensure!( + policy.llvm_results == [expected_llvm_result] + && policy.ptx_result == policy.rust_result + && policy.lowering == "direct_nvvm", + "{} has a signature or lowering outside the scalar direct-NVVM sreg recipe", + policy.id + ); + ensure!( + policy.resolved_llvm_symbol.is_none() + && policy.backend_lowerings.is_empty() + && policy.special_register.is_none(), + "{} uses a backend contract outside the direct-NVVM sreg recipe", + policy.id + ); + ensure_no_other_family_contract(policy, "sreg")?; + if policy.id.starts_with("lanemask_") { + validate_lanemask_policy(policy, declaration)?; + } + if is_cluster_sreg_source(&declaration.source_record) { + validate_cluster_sreg_policy(policy, declaration)?; + } + Ok(()) +} + +pub(in crate::resolve) fn is_cluster_sreg_source(source_record: &str) -> bool { + source_record == "int_nvvm_read_ptx_sreg_cluster_ctarank" + || source_record == "int_nvvm_read_ptx_sreg_cluster_nctarank" + || source_record.starts_with("int_nvvm_read_ptx_sreg_cluster_ctaid_") + || source_record.starts_with("int_nvvm_read_ptx_sreg_cluster_nctaid_") + || source_record.starts_with("int_nvvm_read_ptx_sreg_clusterid_") + || source_record.starts_with("int_nvvm_read_ptx_sreg_nclusterid_") +} + +#[derive(Clone)] +pub(in crate::resolve) struct ClusterSregRecipe { + id: String, + operation_key: String, + source_suffix: String, + llvm_suffix: String, + selection_record: String, + ptx_register: String, + compatibility_path: Option, + op_type: String, + scope: &'static str, + section: &'static str, + anchor: &'static str, + range: Option<&'static str>, + safe_reason: String, + summary: String, +} + +#[derive(Clone, Copy)] +pub(in crate::resolve) struct ClusterSregXyzFamilyRecipe { + id_prefix: &'static str, + operation_key_prefix: &'static str, + source_prefix: &'static str, + llvm_prefix: &'static str, + selection_prefix: &'static str, + ptx_prefix: &'static str, + compatibility_prefix: Option<&'static str>, + op_type_prefix: &'static str, + scope: &'static str, + section: &'static str, + anchor: &'static str, + x_range: &'static str, + yz_range: &'static str, + safe_reason: &'static str, + summary: &'static str, +} + +pub(in crate::resolve) const CLUSTER_SREG_AXES: [&str; 3] = ["x", "y", "z"]; + +pub(in crate::resolve) const CLUSTER_SREG_XYZ_FAMILIES: [ClusterSregXyzFamilyRecipe; 4] = [ + ClusterSregXyzFamilyRecipe { + id_prefix: "cluster_block_idx", + operation_key_prefix: "launch.cluster.block_index", + source_prefix: "cluster_ctaid_", + llvm_prefix: "cluster.ctaid.", + selection_prefix: "INT_PTX_SREG_CLUSTER_CTAID_", + ptx_prefix: "%cluster_ctaid.", + compatibility_prefix: Some("cuda_device::cluster::cluster_ctaid"), + op_type_prefix: "ReadPtxSregClusterCtaid", + scope: "cta", + section: "10.14 Special Registers: %cluster_ctaid", + anchor: "cluster-ctaid", + x_range: "Range", + yz_range: "Range", + safe_reason: "reading the read-only block index within its cluster has no caller obligations", + summary: "Returns the block's {axis} index within its thread block cluster.", + }, + ClusterSregXyzFamilyRecipe { + id_prefix: "cluster_dim", + operation_key_prefix: "launch.cluster.dimension", + source_prefix: "cluster_nctaid_", + llvm_prefix: "cluster.nctaid.", + selection_prefix: "INT_PTX_SREG_CLUSTER_NCTAID_", + ptx_prefix: "%cluster_nctaid.", + compatibility_prefix: Some("cuda_device::cluster::cluster_nctaid"), + op_type_prefix: "ReadPtxSregClusterNctaid", + scope: "cluster", + section: "10.15 Special Registers: %cluster_nctaid", + anchor: "cluster-nctaid", + x_range: "Range", + yz_range: "Range", + safe_reason: "reading the read-only cluster dimension has no caller obligations", + summary: "Returns the number of blocks in the cluster's {axis} dimension.", + }, + ClusterSregXyzFamilyRecipe { + id_prefix: "cluster_idx", + operation_key_prefix: "launch.cluster.index", + source_prefix: "clusterid_", + llvm_prefix: "clusterid.", + selection_prefix: "INT_PTX_SREG_CLUSTERID_", + ptx_prefix: "%clusterid.", + compatibility_prefix: Some("cuda_device::cluster::__cluster_idx"), + op_type_prefix: "ReadPtxSregClusterId", + scope: "cluster", + section: "10.12 Special Registers: %clusterid", + anchor: "clusterid", + x_range: "Range", + yz_range: "Range", + safe_reason: "reading the read-only cluster index has no caller obligations", + summary: "Returns the cluster's {axis} index within the grid.", + }, + ClusterSregXyzFamilyRecipe { + id_prefix: "cluster_grid_dim", + operation_key_prefix: "launch.cluster.grid_dimension", + source_prefix: "nclusterid_", + llvm_prefix: "nclusterid.", + selection_prefix: "INT_PTX_SREG_NCLUSTERID_", + ptx_prefix: "%nclusterid.", + compatibility_prefix: Some("cuda_device::cluster::__cluster_grid_dim"), + op_type_prefix: "ReadPtxSregNclusterId", + scope: "grid", + section: "10.13 Special Registers: %nclusterid", + anchor: "nclusterid", + x_range: "Range", + yz_range: "Range", + safe_reason: "reading the read-only cluster-grid dimension has no caller obligations", + summary: "Returns the number of clusters in the grid's {axis} dimension.", + }, +]; + +pub(in crate::resolve) fn cluster_sreg_recipes() -> Vec { + let mut recipes = Vec::with_capacity(14); + for family in CLUSTER_SREG_XYZ_FAMILIES { + for axis in CLUSTER_SREG_AXES { + let axis_upper = axis.to_ascii_uppercase(); + recipes.push(ClusterSregRecipe { + id: format!("{}_{axis}", family.id_prefix), + operation_key: format!("{}.{axis}", family.operation_key_prefix), + source_suffix: format!("{}{axis}", family.source_prefix), + llvm_suffix: format!("{}{axis}", family.llvm_prefix), + selection_record: format!("{}{axis}", family.selection_prefix), + ptx_register: format!("{}{axis}", family.ptx_prefix), + compatibility_path: family + .compatibility_prefix + .map(|prefix| format!("{prefix}{axis_upper}")), + op_type: format!("{}{axis_upper}Op", family.op_type_prefix), + scope: family.scope, + section: family.section, + anchor: family.anchor, + range: Some(if axis == "x" { + family.x_range + } else { + family.yz_range + }), + safe_reason: family.safe_reason.into(), + summary: family.summary.replace("{axis}", &axis_upper), + }); + } + } + recipes.extend([ + ClusterSregRecipe { + id: "cluster_block_rank".into(), + operation_key: "launch.cluster.block_rank".into(), + source_suffix: "cluster_ctarank".into(), + llvm_suffix: "cluster.ctarank".into(), + selection_record: "INT_PTX_SREG_CLUSTER_CTARANK".into(), + ptx_register: "%cluster_ctarank".into(), + compatibility_path: None, + op_type: "ReadPtxSregClusterCtarankOp".into(), + scope: "cta", + section: "10.16 Special Registers: %cluster_ctarank", + anchor: "cluster-ctarank", + range: None, + safe_reason: + "reading the read-only block rank within its cluster has no caller obligations" + .into(), + summary: "Returns the block's linear rank within its thread block cluster.".into(), + }, + ClusterSregRecipe { + id: "cluster_block_count".into(), + operation_key: "launch.cluster.block_count".into(), + source_suffix: "cluster_nctarank".into(), + llvm_suffix: "cluster.nctarank".into(), + selection_record: "INT_PTX_SREG_CLUSTER_NCTARANK".into(), + ptx_register: "%cluster_nctarank".into(), + compatibility_path: None, + op_type: "ReadPtxSregClusterNctarankOp".into(), + scope: "cluster", + section: "10.17 Special Registers: %cluster_nctarank", + anchor: "cluster-nctarank", + range: None, + safe_reason: "reading the read-only block count has no caller obligations".into(), + summary: "Returns the total number of blocks in the thread block cluster.".into(), + }, + ]); + recipes +} + +pub(in crate::resolve) fn expand_cluster_sreg_admission( + admission: &ClusterSregAdmission, +) -> Result> { + ensure!( + admission.axes == CLUSTER_SREG_AXES, + "cluster-sreg axes must be exactly x, y, z" + ); + ensure!( + admission.xyz_product_count == 12 && admission.record_count == 14, + "cluster-sreg admission must expand to 12 xyz and 14 total records" + ); + let recipes = cluster_sreg_recipes(); + ensure!( + recipes.len() == admission.record_count, + "cluster-sreg recipe count disagrees with its admission" + ); + Ok(recipes.into_iter().map(cluster_sreg_policy).collect()) +} + +pub(in crate::resolve) fn cluster_sreg_policy(recipe: ClusterSregRecipe) -> OverlayIntrinsic { + let compatibility_rust_paths = recipe.compatibility_path.iter().cloned().collect(); + OverlayIntrinsic { + id: recipe.id.clone(), + abi_id: String::new(), + operation_key: recipe.operation_key, + family: "sreg".into(), + source: None, + source_record: Some(format!("int_nvvm_read_ptx_sreg_{}", recipe.source_suffix)), + rust_module: "sreg".into(), + rust_name: recipe.id.clone(), + rust_arguments: vec![], + rust_result: "u32".into(), + safe: true, + must_use: false, + safe_allowlist_reason: Some(recipe.safe_reason), + public_rust_path: format!("cuda_intrinsics::sreg::{}", recipe.id), + compatibility_rust_paths, + dialect_op_type: recipe.op_type, + dialect_op_name: format!("nvvm.read_ptx_sreg_{}", recipe.source_suffix), + dialect_operands: vec![], + dialect_results: vec![], + llvm_symbol: Some(format!("llvm.nvvm.read.ptx.sreg.{}", recipe.llvm_suffix)), + resolved_llvm_symbol: None, + llvm_arguments: vec![], + llvm_results: vec!["i32".into()], + pure: true, + memory: "none".into(), + convergent: false, + execution_scope: recipe.scope.into(), + minimum_ptx: "7.8".into(), + minimum_sm: Some("sm_90".into()), + ptx_result: "u32".into(), + targets: "all".into(), + ptx_isa_version: "9.3".into(), + ptx_isa_section: recipe.section.into(), + ptx_isa_url: format!( + "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-{}", + recipe.anchor + ), + lowering: "direct_nvvm".into(), + backend_lowerings: vec![], + packed_atomic: None, + redux: None, + vote: None, + active_mask: None, + warp_match: None, + warp_barrier: None, + warp_shuffle: None, + dot_product: None, + packed_alu: None, + integer_minmax: None, + packed_conversion: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cp_async_copy: None, + cp_async_control: None, + cp_async_mbarrier: None, + mbarrier_basic: None, + movmatrix: None, + mbarrier_extended: None, + register_mma: None, + sparse_mma: None, + prmt: None, + cluster_barrier: None, + wgmma_control: None, + special_register: None, + debug_control: None, + cluster_memory: None, + clc: None, + tma: None, + tcgen05: None, + ldmatrix_variant: None, + ldmatrix_safety: None, + ldmatrix_adapter: None, + selected_address_space: None, + expected_ptx: InstructionPattern { + mnemonic: "mov".into(), + modifiers: vec!["u32".into()], + operands: vec![ + OperandPattern::Register, + OperandPattern::Exact { + value: recipe.ptx_register, + }, + ], + }, + summary: recipe.summary, + } +} + +pub(in crate::resolve) fn validate_cluster_sreg_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, +) -> Result<()> { + ensure!( + !declaration.source_record.ends_with("_w"), + "{} selects an unused always-zero fourth cluster-register component", + policy.id + ); + let recipe = cluster_sreg_recipes() + .into_iter() + .find(|recipe| recipe.id == policy.id) + .with_context(|| format!("{} is not a reviewed cluster special register", policy.id))?; + + let source_record = format!("int_nvvm_read_ptx_sreg_{}", recipe.source_suffix); + let llvm_symbol = format!("llvm.nvvm.read.ptx.sreg.{}", recipe.llvm_suffix); + let compatibility_paths = recipe + .compatibility_path + .iter() + .cloned() + .collect::>(); + ensure!( + policy.operation_key == recipe.operation_key + && policy.source.is_none() + && policy.source_record.as_deref() == Some(source_record.as_str()) + && policy.llvm_symbol.as_deref() == Some(llvm_symbol.as_str()) + && policy.resolved_llvm_symbol.is_none(), + "{} cluster-register identity does not match its closed recipe", + policy.id + ); + ensure!( + policy.rust_module == "sreg" + && policy.rust_name == policy.id + && policy.rust_arguments.is_empty() + && policy.rust_result == "u32" + && policy.safe + && !policy.must_use + && policy + .safe_allowlist_reason + .as_deref() + .is_some_and(|reason| !reason.is_empty()) + && policy.public_rust_path == format!("cuda_intrinsics::sreg::{}", policy.id) + && policy.compatibility_rust_paths == compatibility_paths, + "{} must preserve its reviewed raw and compatibility APIs", + policy.id + ); + ensure!( + policy.dialect_op_type == recipe.op_type + && policy.dialect_op_name == format!("nvvm.read_ptx_sreg_{}", recipe.source_suffix) + && policy.dialect_operands.is_empty() + && policy.dialect_results.is_empty() + && policy.llvm_arguments.is_empty() + && policy.llvm_results == ["i32"] + && policy.lowering == "direct_nvvm", + "{} is outside the closed cluster-register lowering recipe", + policy.id + ); + ensure!( + policy.pure + && policy.memory == "none" + && !policy.convergent + && policy.execution_scope == recipe.scope + && policy.minimum_ptx == "7.8" + && policy.minimum_sm.as_deref() == Some("sm_90") + && policy.ptx_result == "u32" + && policy.targets == "all", + "{} cluster-register effects or target floor disagree with PTX", + policy.id + ); + ensure!( + policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section == recipe.section + && policy.ptx_isa_url + == format!( + "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-{}", + recipe.anchor + ), + "{} cluster-register PTX provenance disagrees with the reviewed recipe", + policy.id + ); + let mut properties = vec!["IntrNoMem", "IntrSpeculatable", "NoUndef"]; + properties.extend(recipe.range); + ensure!( + declaration.arguments.is_empty() + && declaration.results == ["i32"] + && declaration.classes + == [ + "SDPatternOperator", + "Intrinsic", + "DefaultAttrsIntrinsic", + "NVVMPureIntrinsic", + "PTXReadSRegIntrinsicNB_r32", + ] + && declaration.properties == properties, + "{} declaration shape or properties disagree with LLVM TableGen", + policy.id + ); + let [selection] = declaration.selections.as_slice() else { + bail!("{} must have exactly one LLVM selection", policy.id); + }; + ensure!( + selection.source_record == recipe.selection_record + && selection.asm == format!("mov.u32 \t$d, {};", recipe.ptx_register) + && selection.predicates + == [ + "Subtarget->getSmVersion() >= 90", + "Subtarget->getPTXVersion() >= 78", + ] + && selection.constraints.address_space.is_none() + && selection.constraints.immediate_bindings.is_empty(), + "{} selector disagrees with LLVM TableGen", + policy.id + ); + ensure!( + policy.expected_ptx.mnemonic == "mov" + && policy.expected_ptx.modifiers == ["u32"] + && policy.expected_ptx.operands + == [ + OperandPattern::Register, + OperandPattern::Exact { + value: recipe.ptx_register, + }, + ], + "{} expected PTX does not match its cluster register", + policy.id + ); + Ok(()) +} + +pub(in crate::resolve) fn validate_lanemask_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, +) -> Result<()> { + let (suffix, abi_id, section, op_type) = match policy.id.as_str() { + "lanemask_lt" => ("lt", "i0035", "10.13", "ReadPtxSregLanemaskLtOp"), + "lanemask_le" => ("le", "i0036", "10.12", "ReadPtxSregLanemaskLeOp"), + "lanemask_eq" => ("eq", "i0037", "10.11", "ReadPtxSregLanemaskEqOp"), + "lanemask_ge" => ("ge", "i0038", "10.14", "ReadPtxSregLanemaskGeOp"), + "lanemask_gt" => ("gt", "i0039", "10.15", "ReadPtxSregLanemaskGtOp"), + _ => bail!("{} is not a reviewed lane-mask special register", policy.id), + }; + ensure!( + policy.abi_id == abi_id + && policy.operation_key == format!("warp.lane_mask.{suffix}") + && policy.source.is_none() + && policy.source_record.as_deref() + == Some(format!("int_nvvm_read_ptx_sreg_lanemask_{suffix}").as_str()) + && policy.llvm_symbol.as_deref() + == Some(format!("llvm.nvvm.read.ptx.sreg.lanemask.{suffix}").as_str()) + && policy.resolved_llvm_symbol.is_none(), + "{} lane-mask identity does not match its closed recipe", + policy.id + ); + ensure!( + policy.rust_module == "sreg" + && policy.rust_name == policy.id + && policy.rust_arguments.is_empty() + && policy.rust_result == "u32" + && policy.safe + && policy.must_use + && policy + .safe_allowlist_reason + .as_deref() + .is_some_and(|reason| !reason.is_empty()) + && policy.public_rust_path == format!("cuda_intrinsics::sreg::{}", policy.id) + && policy.compatibility_rust_paths == [format!("cuda_device::warp::{}", policy.id)], + "{} must preserve its safe must-use raw and compatibility APIs", + policy.id + ); + ensure!( + policy.dialect_op_type == op_type + && policy.dialect_op_name == format!("nvvm.read_ptx_sreg_lanemask_{suffix}") + && policy.dialect_operands.is_empty() + && policy.dialect_results.is_empty() + && policy.llvm_arguments.is_empty() + && policy.llvm_results == ["i32"] + && policy.lowering == "direct_nvvm", + "{} is outside the closed lane-mask lowering recipe", + policy.id + ); + ensure!( + policy.pure + && policy.memory == "none" + && !policy.convergent + && policy.execution_scope == "thread" + && policy.minimum_ptx == "2.0" + && policy.minimum_sm.as_deref() == Some("sm_20") + && policy.ptx_result == "u32" + && policy.targets == "all", + "{} lane-mask effects or target floor disagree with PTX", + policy.id + ); + ensure!( + policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section == format!("{section} Special Registers: %lanemask_{suffix}") + && policy.ptx_isa_url + == format!( + "https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-lanemask-{suffix}" + ), + "{} lane-mask PTX provenance disagrees with the reviewed recipe", + policy.id + ); + ensure!( + declaration.properties == ["IntrNoMem", "IntrSpeculatable", "NoUndef"], + "{} lane-mask properties disagree with the imported LLVM declaration", + policy.id + ); + ensure!( + policy.expected_ptx.mnemonic == "mov" + && policy.expected_ptx.modifiers == ["u32"] + && policy.expected_ptx.operands + == [ + OperandPattern::Register, + OperandPattern::Exact { + value: format!("%lanemask_{suffix}"), + }, + ], + "{} expected PTX does not match its lane-mask register", + policy.id + ); + Ok(()) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/tcgen05/admission.rs b/crates/cuda-intrinsics-gen/src/resolve/families/tcgen05/admission.rs new file mode 100644 index 0000000000..a30936a1b9 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/tcgen05/admission.rs @@ -0,0 +1,850 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, ImportedIntrinsic, IntrinsicBackend, OverlayBackendLowering, + OverlayIntrinsic, RuntimeValidation, TargetSelectorBinding, Tcgen05, Tcgen05Admission, + Tcgen05CpGroup, Tcgen05MmaBUsage, Tcgen05MmaForm, Tcgen05Operation, Tcgen05SourceContract, +}; +use crate::ptx::InstructionPattern; +use anyhow::{Context, Result, ensure}; +use std::collections::BTreeSet; + +use super::*; +use crate::resolve::abi_ledger::*; +use crate::resolve::guards::*; +use crate::resolve::targets::*; + +pub(in crate::resolve) fn expand_tcgen05_admission( + admission: &Tcgen05Admission, +) -> Result> { + ensure!( + admission.runtime_validation == RuntimeValidation::Unexecuted, + "tcgen05 runtime validation may be marked executed only with GPU evidence" + ); + ensure!( + !admission.llvm_evidence_profile.trim().is_empty() + && !admission.libnvvm_evidence_profile.trim().is_empty(), + "compact tcgen05 admission requires both backend evidence profiles" + ); + let expected = [ + Tcgen05Operation::Alloc, + Tcgen05Operation::Dealloc, + Tcgen05Operation::RelinquishAllocPermit, + Tcgen05Operation::FenceBeforeThreadSync, + Tcgen05Operation::FenceAfterThreadSync, + Tcgen05Operation::Commit, + Tcgen05Operation::CommitSharedCluster, + Tcgen05Operation::MmaWsF16, + Tcgen05Operation::MmaF16, + Tcgen05Operation::MmaWsBf16, + Tcgen05Operation::MmaWsTf32, + Tcgen05Operation::CpSmemToTmem, + Tcgen05Operation::Ld16x256bX8Pure, + Tcgen05Operation::Ld16x256bPure, + Tcgen05Operation::LoadWait, + Tcgen05Operation::StoreWait, + Tcgen05Operation::AllocCg2, + Tcgen05Operation::DeallocCg2, + Tcgen05Operation::RelinquishAllocPermitCg2, + Tcgen05Operation::MmaF16Cg2, + Tcgen05Operation::CommitCg2, + Tcgen05Operation::CommitSharedClusterCg2, + Tcgen05Operation::CommitMulticastCg2, + Tcgen05Operation::CpSmemToTmemCg2, + Tcgen05Operation::CommitMulticast, + Tcgen05Operation::ShiftDown, + Tcgen05Operation::ShiftDownCg2, + ]; + let has_control_variants = admission.variants.iter().any(|variant| { + matches!( + variant.operation, + Tcgen05Operation::CommitMulticast + | Tcgen05Operation::ShiftDown + | Tcgen05Operation::ShiftDownCg2 + ) + }); + if has_control_variants { + ensure!( + admission + .control_llvm_evidence_profile + .as_deref() + .is_some_and(|profile| !profile.trim().is_empty()) + && admission + .control_libnvvm_evidence_profile + .as_deref() + .is_some_and(|profile| !profile.trim().is_empty()), + "compact tcgen05 control admission requires both backend evidence profiles" + ); + } + let expected = if has_control_variants { + &expected[..] + } else { + &expected[..24] + }; + ensure!( + admission + .variants + .iter() + .map(|variant| variant.operation) + .eq(expected.iter().copied()), + "compact tcgen05 admission must list all 24 base operations or all 27 current operations in canonical order" + ); + + let mut records = admission + .variants + .iter() + .map(|variant| { + let recipe = tcgen05_recipe(variant.operation); + let control = matches!( + variant.operation, + Tcgen05Operation::CommitMulticast + | Tcgen05Operation::ShiftDown + | Tcgen05Operation::ShiftDownCg2 + ); + ensure!( + variant.abi_id == recipe.abi_id, + "{} must keep reserved ABI ID {}", + recipe.id, + recipe.abi_id + ); + Ok(OverlayIntrinsic { + id: recipe.id.into(), + abi_id: variant.abi_id.clone(), + operation_key: recipe.operation_key.into(), + family: "tcgen05".into(), + source: None, + source_record: Some(recipe.source_record.into()), + rust_module: "tcgen05".into(), + rust_name: recipe.id.into(), + rust_arguments: recipe + .rust_arguments + .iter() + .map(|value| (*value).into()) + .collect(), + rust_result: recipe.rust_result.into(), + safe: recipe.safe, + must_use: false, + safe_allowlist_reason: recipe.safe_reason.map(Into::into), + public_rust_path: format!("cuda_intrinsics::tcgen05::{}", recipe.id), + compatibility_rust_paths: vec![format!("cuda_device::tcgen05::{}", recipe.id)], + dialect_op_type: recipe.dialect_op_type.into(), + dialect_op_name: recipe.dialect_op_name.into(), + dialect_operands: recipe + .dialect_operands + .iter() + .map(|value| (*value).into()) + .collect(), + dialect_results: recipe + .dialect_results + .iter() + .map(|value| (*value).into()) + .collect(), + llvm_symbol: Some(recipe.llvm_symbol.into()), + resolved_llvm_symbol: None, + llvm_arguments: recipe + .llvm_arguments + .iter() + .map(|value| (*value).into()) + .collect(), + llvm_results: recipe + .llvm_results + .iter() + .map(|value| (*value).into()) + .collect(), + pure: false, + memory: recipe.memory.into(), + convergent: true, + execution_scope: recipe.operation.execution_scope().into(), + minimum_ptx: "8.6".into(), + minimum_sm: None, + ptx_result: recipe.rust_result.into(), + targets: TCGEN05_LLVM_TARGETS.into(), + ptx_isa_version: "8.6".into(), + ptx_isa_section: "Tensor Memory tcgen05 instructions".into(), + ptx_isa_url: + "https://docs.nvidia.com/cuda/parallel-thread-execution/#tensor-memory".into(), + lowering: "generated_tcgen05".into(), + backend_lowerings: vec![ + OverlayBackendLowering { + backend: IntrinsicBackend::LlvmNvptx, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: if control { + admission + .control_llvm_evidence_profile + .as_ref() + .expect("validated tcgen05 control LLVM evidence profile") + .clone() + } else { + admission.llvm_evidence_profile.clone() + }, + targets: None, + minimum_ptx: Some("8.6".into()), + minimum_sm: None, + }, + OverlayBackendLowering { + backend: IntrinsicBackend::LibNvvm, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: if control { + admission + .control_libnvvm_evidence_profile + .as_ref() + .expect("validated tcgen05 control libNVVM evidence profile") + .clone() + } else { + admission.libnvvm_evidence_profile.clone() + }, + targets: Some(TCGEN05_LIBNVVM_TARGETS.into()), + minimum_ptx: Some("8.6".into()), + minimum_sm: None, + }, + ], + packed_atomic: None, + redux: None, + vote: None, + active_mask: None, + warp_match: None, + warp_barrier: None, + warp_shuffle: None, + dot_product: None, + packed_alu: None, + integer_minmax: None, + packed_conversion: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cp_async_copy: None, + cp_async_control: None, + cp_async_mbarrier: None, + mbarrier_basic: None, + movmatrix: None, + mbarrier_extended: None, + register_mma: None, + sparse_mma: None, + prmt: None, + cluster_barrier: None, + wgmma_control: None, + special_register: None, + debug_control: None, + cluster_memory: None, + clc: None, + tma: None, + tcgen05: Some(Tcgen05 { + operation: recipe.operation, + cp: None, + ld: None, + st: None, + mma: None, + adapter: recipe.adapter, + source_contract: recipe.source_contract, + runtime_validation: admission.runtime_validation, + }), + ldmatrix_variant: None, + ldmatrix_safety: None, + ldmatrix_adapter: None, + selected_address_space: None, + expected_ptx: InstructionPattern { + mnemonic: "tcgen05".into(), + modifiers: recipe.modifiers, + operands: recipe.operands, + }, + summary: recipe.summary.into(), + }) + }) + .collect::>>()?; + + if admission.cp_variants.is_empty() { + ensure!( + admission.cp_llvm_evidence_profile.is_none() + && admission.cp_libnvvm_evidence_profile.is_none(), + "tcgen05 copy evidence profiles require admitted copy variants" + ); + } else { + ensure!( + [ + Tcgen05MmaBUsage::Discard, + Tcgen05MmaBUsage::LastUse, + Tcgen05MmaBUsage::Fill, + Tcgen05MmaBUsage::Use, + ] + .map(Tcgen05MmaBUsage::selector_value) + == [0, 1, 2, 3], + "tcgen05 MMA B-usage selector mapping changed" + ); + ensure!( + admission + .cp_llvm_evidence_profile + .as_deref() + .is_some_and(|profile| !profile.trim().is_empty()) + && admission + .cp_libnvvm_evidence_profile + .as_deref() + .is_some_and(|profile| !profile.trim().is_empty()), + "compact tcgen05 copy admission requires both backend evidence profiles" + ); + + use Tcgen05CpGroup::{Cg1, Cg2}; + let expected_cp = TCGEN05_CP_MEMBERS + .into_iter() + .flat_map(|member| [Cg1, Cg2].into_iter().map(move |group| (member, group))) + .collect::>(); + ensure!( + admission + .cp_variants + .iter() + .map(|variant| (variant.member, variant.group)) + .eq(expected_cp), + "compact tcgen05 copy admission must list all 34 variants in canonical order" + ); + for variant in &admission.cp_variants { + validate_abi_id(&variant.abi_id)?; + let base_id = if variant.group == Cg1 { + "tcgen05_cp_smem_to_tmem" + } else { + "tcgen05_cp_smem_to_tmem_cg2" + }; + let base = records + .iter() + .find(|record| record.id == base_id) + .unwrap() + .clone(); + records.push(materialize_tcgen05_cp_variant(&base, admission, variant)); + } + } + + if admission.ld_variants.is_empty() { + ensure!( + admission.ld_llvm_evidence_profile.is_none() + && admission.ld_libnvvm_evidence_profile.is_none(), + "tcgen05 load evidence profiles require admitted load variants" + ); + } else { + ensure!( + admission + .ld_llvm_evidence_profile + .as_deref() + .is_some_and(|profile| !profile.trim().is_empty()) + && admission + .ld_libnvvm_evidence_profile + .as_deref() + .is_some_and(|profile| !profile.trim().is_empty()), + "compact tcgen05 load admission requires both backend evidence profiles" + ); + let expected_ld = TCGEN05_LD_VARIANTS + .into_iter() + .flat_map(|(shape, multiplicity)| { + [false, true] + .into_iter() + .map(move |pack16| (shape, multiplicity, pack16)) + }) + .collect::>(); + ensure!( + admission + .ld_variants + .iter() + .map(|variant| (variant.shape, variant.multiplicity, variant.pack16)) + .eq(expected_ld), + "compact tcgen05 load admission must list all 58 variants in canonical order" + ); + let base = records + .iter() + .find(|record| record.id == "tcgen05_ld_16x256b_pure") + .expect("closed tcgen05 load base") + .clone(); + let first_load = records.len(); + for variant in &admission.ld_variants { + validate_abi_id(&variant.abi_id)?; + records.push(materialize_tcgen05_ld_variant(&base, admission, variant)); + } + for pair in records[first_load..].chunks_exact(2) { + let raw = pair[0].tcgen05.as_ref().and_then(|tcgen05| tcgen05.ld); + let packed = pair[1].tcgen05.as_ref().and_then(|tcgen05| tcgen05.ld); + ensure!( + raw.is_some_and(|ld| !ld.pack16) + && packed.is_some_and(|ld| ld.pack16) + && raw.map(|mut ld| { + ld.pack16 = true; + ld + }) == packed + && pair[0].source_record == pair[1].source_record + && pair[0].llvm_symbol == pair[1].llvm_symbol, + "tcgen05 load source sharing must pair raw and pack16 leaves" + ); + } + } + + if admission.st_variants.is_empty() { + ensure!( + admission.st_llvm_evidence_profile.is_none() + && admission.st_libnvvm_evidence_profile.is_none(), + "tcgen05 store evidence profiles require admitted store variants" + ); + } else { + ensure!( + admission + .st_llvm_evidence_profile + .as_deref() + .is_some_and(|profile| !profile.trim().is_empty()) + && admission + .st_libnvvm_evidence_profile + .as_deref() + .is_some_and(|profile| !profile.trim().is_empty()), + "compact tcgen05 store admission requires both backend evidence profiles" + ); + let expected_st = TCGEN05_ST_VARIANTS + .into_iter() + .flat_map(|(shape, multiplicity)| { + [false, true] + .into_iter() + .map(move |unpack16| (shape, multiplicity, unpack16)) + }) + .collect::>(); + ensure!( + admission + .st_variants + .iter() + .map(|variant| (variant.shape, variant.multiplicity, variant.unpack16)) + .eq(expected_st), + "compact tcgen05 store admission must list all 58 variants in canonical order" + ); + let base = records + .iter() + .find(|record| record.id == "tcgen05_ld_16x256b_pure") + .expect("closed tcgen05 store base") + .clone(); + let first_store = records.len(); + for variant in &admission.st_variants { + validate_abi_id(&variant.abi_id)?; + records.push(materialize_tcgen05_st_variant(&base, admission, variant)); + } + for pair in records[first_store..].chunks_exact(2) { + let raw = pair[0].tcgen05.as_ref().and_then(|tcgen05| tcgen05.st); + let unpacked = pair[1].tcgen05.as_ref().and_then(|tcgen05| tcgen05.st); + ensure!( + raw.is_some_and(|st| !st.unpack16) + && unpacked.is_some_and(|st| st.unpack16) + && raw.map(|mut st| { + st.unpack16 = true; + st + }) == unpacked + && pair[0].source_record == pair[1].source_record + && pair[0].llvm_symbol == pair[1].llvm_symbol, + "tcgen05 store source sharing must pair raw and unpack16 leaves" + ); + } + } + + if admission.ld_offset_variants.is_empty() && admission.st_offset_variants.is_empty() { + ensure!( + admission.offset_llvm_evidence_profile.is_none() + && admission.offset_libnvvm_evidence_profile.is_none(), + "tcgen05 offset evidence profiles require admitted offset load/store variants" + ); + } else { + ensure!( + admission + .offset_llvm_evidence_profile + .as_deref() + .is_some_and(|profile| !profile.trim().is_empty()) + && admission + .offset_libnvvm_evidence_profile + .as_deref() + .is_some_and(|profile| !profile.trim().is_empty()), + "compact tcgen05 offset admission requires both backend evidence profiles" + ); + let expected_ld = TCGEN05_OFFSET_LDST_VARIANTS + .into_iter() + .flat_map(|(shape, multiplicity)| { + [false, true] + .into_iter() + .map(move |pack16| (shape, multiplicity, pack16)) + }) + .collect::>(); + ensure!( + admission + .ld_offset_variants + .iter() + .map(|variant| (variant.shape, variant.multiplicity, variant.pack16)) + .eq(expected_ld), + "compact tcgen05 offset load admission must list all 16 variants in canonical order" + ); + let expected_st = TCGEN05_OFFSET_LDST_VARIANTS + .into_iter() + .flat_map(|(shape, multiplicity)| { + [false, true] + .into_iter() + .map(move |unpack16| (shape, multiplicity, unpack16)) + }) + .collect::>(); + ensure!( + admission + .st_offset_variants + .iter() + .map(|variant| (variant.shape, variant.multiplicity, variant.unpack16)) + .eq(expected_st), + "compact tcgen05 offset store admission must list all 16 variants in canonical order" + ); + + let base = records + .iter() + .find(|record| record.id == "tcgen05_ld_16x256b_pure") + .expect("closed tcgen05 offset load/store base") + .clone(); + let first_load = records.len(); + for variant in &admission.ld_offset_variants { + validate_abi_id(&variant.abi_id)?; + records.push(materialize_tcgen05_ld_variant(&base, admission, variant)); + } + for pair in records[first_load..].chunks_exact(2) { + let raw = pair[0].tcgen05.as_ref().and_then(|tcgen05| tcgen05.ld); + let packed = pair[1].tcgen05.as_ref().and_then(|tcgen05| tcgen05.ld); + ensure!( + raw.is_some_and(|ld| !ld.pack16) + && packed.is_some_and(|ld| ld.pack16) + && raw.map(|mut ld| { + ld.pack16 = true; + ld + }) == packed + && pair[0].source_record == pair[1].source_record + && pair[0].llvm_symbol == pair[1].llvm_symbol, + "tcgen05 offset load source sharing must pair raw and pack16 leaves" + ); + } + + let first_store = records.len(); + for variant in &admission.st_offset_variants { + validate_abi_id(&variant.abi_id)?; + records.push(materialize_tcgen05_st_variant(&base, admission, variant)); + } + for pair in records[first_store..].chunks_exact(2) { + let raw = pair[0].tcgen05.as_ref().and_then(|tcgen05| tcgen05.st); + let unpacked = pair[1].tcgen05.as_ref().and_then(|tcgen05| tcgen05.st); + ensure!( + raw.is_some_and(|st| !st.unpack16) + && unpacked.is_some_and(|st| st.unpack16) + && raw.map(|mut st| { + st.unpack16 = true; + st + }) == unpacked + && pair[0].source_record == pair[1].source_record + && pair[0].llvm_symbol == pair[1].llvm_symbol, + "tcgen05 offset store source sharing must pair raw and unpack16 leaves" + ); + } + } + + if admission.mma_variants.is_empty() { + ensure!( + admission.mma_llvm_evidence_profile.is_none() + && admission.mma_libnvvm_evidence_profile.is_none() + && admission.mma_llvm_target_contracts.is_empty() + && admission.mma_libnvvm_target_contracts.is_empty(), + "tcgen05 MMA profiles and target contracts require admitted variants" + ); + } else { + ensure!( + admission + .mma_llvm_evidence_profile + .as_deref() + .is_some_and(|profile| !profile.trim().is_empty()) + && admission + .mma_libnvvm_evidence_profile + .as_deref() + .is_some_and(|profile| !profile.trim().is_empty()), + "compact tcgen05 MMA admission requires both backend evidence profiles" + ); + let expected_variants = TCGEN05_MMA_FORMS + .into_iter() + .map(|form| (form, None)) + .chain( + TCGEN05_MMA_ALIASES + .into_iter() + .map(|alias| (Tcgen05MmaForm::WsTensor, Some(alias))), + ) + .chain( + TCGEN05_MMA_ALIASES + .into_iter() + .map(|alias| (Tcgen05MmaForm::Shared, Some(alias))), + ) + .collect::>(); + ensure!( + admission + .mma_variants + .iter() + .map(|variant| (variant.form, variant.alias)) + .eq(expected_variants), + "compact tcgen05 MMA admission must list all 24 APIs in canonical order" + ); + for variant in &admission.mma_variants { + validate_abi_id(&variant.abi_id)?; + } + let expected_llvm = expected_tcgen05_mma_target_contracts(IntrinsicBackend::LlvmNvptx); + let expected_libnvvm = expected_tcgen05_mma_target_contracts(IntrinsicBackend::LibNvvm); + ensure!( + admission.mma_llvm_target_contracts == expected_llvm + && admission.mma_libnvvm_target_contracts == expected_libnvvm, + "tcgen05 MMA target contracts changed from the reviewed backend matrices" + ); + let llvm_target = + resolve_target_contracts("tcgen05 MMA LLVM", &admission.mma_llvm_target_contracts)?; + let libnvvm_target = resolve_target_contracts( + "tcgen05 MMA libNVVM", + &admission.mma_libnvvm_target_contracts, + )?; + let fixed_selector = [TargetSelectorBinding { + name: "kind".into(), + value: "f8f6f4".into(), + }]; + let fixed_llvm = resolve_target_contract( + "tcgen05 MMA LLVM alias", + &fixed_selector, + &admission.mma_llvm_target_contracts, + )?; + let fixed_libnvvm = resolve_target_contract( + "tcgen05 MMA libNVVM alias", + &fixed_selector, + &admission.mma_libnvvm_target_contracts, + )?; + for variant in &admission.mma_variants { + let (llvm, libnvvm) = if variant.alias.is_some() { + (&fixed_llvm, &fixed_libnvvm) + } else { + (&llvm_target, &libnvvm_target) + }; + records.push(materialize_tcgen05_mma_variant( + admission, variant, llvm, libnvvm, + )); + } + } + Ok(records) +} + +pub(in crate::resolve) fn validate_tcgen05_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, +) -> Result<()> { + let tcgen05 = policy + .tcgen05 + .as_ref() + .with_context(|| format!("{} has no closed tcgen05 contract", policy.id))?; + if let Some(mma) = &tcgen05.mma { + return validate_tcgen05_mma_policy(policy, declaration, tcgen05, mma); + } + if let Some(ld) = tcgen05.ld { + return validate_tcgen05_ld_policy(policy, declaration, tcgen05, ld); + } + if let Some(st) = tcgen05.st { + return validate_tcgen05_st_policy(policy, declaration, tcgen05, st); + } + if let Some(cp) = tcgen05.cp { + return validate_tcgen05_cp_policy(policy, declaration, tcgen05, cp); + } + ensure!( + !matches!( + tcgen05.operation, + Tcgen05Operation::Ld | Tcgen05Operation::St | Tcgen05Operation::Mma + ), + "{} has a tcgen05 load/store/MMA operation without its closed identity", + policy.id + ); + let recipe = tcgen05_recipe(tcgen05.operation); + ensure!( + policy.id == recipe.id + && policy.abi_id == recipe.abi_id + && policy.operation_key == recipe.operation_key + && policy.source.is_none() + && policy.source_record.as_deref() == Some(recipe.source_record) + && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) + && policy.resolved_llvm_symbol.is_none() + && declaration.source_record == recipe.source_record + && declaration.llvm_name == recipe.llvm_symbol, + "{} tcgen05 identity changed", + policy.id + ); + ensure!( + policy.rust_module == "tcgen05" + && policy.rust_name == recipe.id + && policy.rust_arguments == recipe.rust_arguments + && policy.rust_result == recipe.rust_result + && policy.safe == recipe.safe + && !policy.must_use + && policy.safe_allowlist_reason.as_deref() == recipe.safe_reason + && policy.public_rust_path == format!("cuda_intrinsics::tcgen05::{}", recipe.id) + && policy.compatibility_rust_paths == [format!("cuda_device::tcgen05::{}", recipe.id)], + "{} tcgen05 Rust API changed", + policy.id + ); + ensure!( + policy.dialect_op_type == recipe.dialect_op_type + && policy.dialect_op_name == recipe.dialect_op_name + && policy.dialect_operands == recipe.dialect_operands + && policy.dialect_results == recipe.dialect_results + && policy.llvm_arguments == recipe.llvm_arguments + && policy.llvm_results == recipe.llvm_results + && declaration.arguments == recipe.llvm_arguments + && declaration.results == recipe.llvm_results + && declaration.classes == recipe.imported_classes + && declaration.properties == recipe.imported_properties + && policy.lowering == "generated_tcgen05", + "{} tcgen05 carrier or imported declaration changed", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == recipe.memory + && policy.convergent + && policy.execution_scope == recipe.operation.execution_scope() + && tcgen05.ld.is_none() + && tcgen05.st.is_none() + && tcgen05.adapter == recipe.adapter + && tcgen05.source_contract == recipe.source_contract + && tcgen05.runtime_validation == RuntimeValidation::Unexecuted, + "{} tcgen05 semantics changed", + policy.id + ); + ensure!( + policy.minimum_ptx == "8.6" + && policy.minimum_sm.is_none() + && policy.targets == TCGEN05_LLVM_TARGETS + && policy.ptx_isa_version == "8.6" + && policy.ptx_result == recipe.rust_result + && policy.expected_ptx.mnemonic == "tcgen05" + && policy.expected_ptx.modifiers == recipe.modifiers + && policy.expected_ptx.operands == recipe.operands, + "{} tcgen05 target or PTX contract changed", + policy.id + ); + validate_tcgen05_backend_routes(policy, "tcgen05")?; + validate_tcgen05_source_contract(&recipe, declaration)?; + ensure_no_other_family_contract(policy, "tcgen05")?; + Ok(()) +} + +pub(in crate::resolve) fn validate_tcgen05_backend_routes( + policy: &OverlayIntrinsic, + family: &str, +) -> Result<()> { + ensure!( + policy.backend_lowerings.len() == 2 + && policy.backend_lowerings[0].backend == IntrinsicBackend::LlvmNvptx + && policy.backend_lowerings[1].backend == IntrinsicBackend::LibNvvm + && policy.backend_lowerings[0].targets.is_none() + && policy.backend_lowerings[1].targets.as_deref() == Some(TCGEN05_LIBNVVM_TARGETS) + && policy.backend_lowerings.iter().all(|route| { + route.mechanism == BackendLoweringMechanism::InlinePtx + && route.minimum_ptx.as_deref() == Some("8.6") + && route.minimum_sm.is_none() + && !route.evidence_profile.trim().is_empty() + }), + "{} {family} backend route changed", + policy.id + ); + Ok(()) +} + +pub(in crate::resolve) fn validate_tcgen05_source_contract( + recipe: &Tcgen05Recipe, + declaration: &ImportedIntrinsic, +) -> Result<()> { + match recipe.source_contract { + Tcgen05SourceContract::ExactTablegenSelection => { + ensure!( + declaration.selections.len() == 1, + "{} must keep one exact tcgen05 selection", + recipe.id + ); + let selection = &declaration.selections[0]; + let expected_predicate = if matches!( + recipe.operation, + Tcgen05Operation::ShiftDown | Tcgen05Operation::ShiftDownCg2 + ) { + "Subtarget->hasTcgen05ShiftSupport()" + } else { + "Subtarget->hasTcgen05InstSupport()" + }; + ensure!( + recipe.selection_record == Some(selection.source_record.as_str()) + && recipe.selection_asm == Some(selection.asm.as_str()) + && selection.predicates == [expected_predicate] + && selection.constraints.is_empty(), + "{} exact tcgen05 selection changed", + recipe.id + ); + } + Tcgen05SourceContract::TablegenSelectionChangesPtx + if matches!( + recipe.operation, + Tcgen05Operation::Commit | Tcgen05Operation::CommitCg2 + ) => + { + ensure!( + declaration.selections.len() == 1, + "{} must keep one canonical commit selection", + recipe.id + ); + let selection = &declaration.selections[0]; + ensure!( + recipe.selection_record == Some(selection.source_record.as_str()) + && recipe.selection_asm == Some(selection.asm.as_str()) + && selection.predicates == ["Subtarget->hasTcgen05InstSupport()"] + && selection.constraints.is_empty() + && !recipe.operands.is_empty(), + "{} canonical commit selection changed", + recipe.id + ); + } + Tcgen05SourceContract::TablegenSelectionChangesPtx => { + ensure!( + declaration.selections.len() == 64, + "{} must keep the 64 collector selections", + recipe.id + ); + let actual = declaration + .selections + .iter() + .map(|selection| { + ensure!( + selection.constraints.is_empty(), + "{} collector selection gained constraints", + recipe.id + ); + let predicate = if selection.asm.contains(".kind::i8.") { + "Subtarget->hasTcgen05MMAI8Kind()" + } else { + "Subtarget->hasTcgen05InstSupport()" + }; + ensure!( + selection.predicates == [predicate], + "{} collector predicate changed", + recipe.id + ); + Ok(selection.asm.clone()) + }) + .collect::>>()?; + let expected = ["f16", "tf32", "f8f6f4", "i8"] + .into_iter() + .flat_map(|kind| { + ["b0", "b1", "b2", "b3"].into_iter().flat_map(move |collector| { + ["discard", "fill", "use", "lastuse"].into_iter().map(move |action| { + format!("tcgen05.mma.ws.cta_group::1.kind::{kind}.collector::{collector}::{action} [$dtmem], [$a], $b, $idesc, $enable_inp_d;") + }) + }) + }) + .collect::>(); + ensure!( + actual == expected, + "{} collector selection matrix changed", + recipe.id + ); + } + Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection => ensure!( + declaration.selections.is_empty(), + "{} unexpectedly gained a TableGen selection; review the backend route", + recipe.id + ), + } + Ok(()) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/tcgen05/ldst_cp.rs b/crates/cuda-intrinsics-gen/src/resolve/families/tcgen05/ldst_cp.rs new file mode 100644 index 0000000000..ae64dc496b --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/tcgen05/ldst_cp.rs @@ -0,0 +1,1081 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + ImportedIntrinsic, OverlayIntrinsic, RuntimeValidation, Tcgen05, Tcgen05Adapter, + Tcgen05Admission, Tcgen05Cp, Tcgen05CpAdmissionVariant, Tcgen05CpGroup, Tcgen05CpMember, + Tcgen05Ld, Tcgen05LdAdmissionVariant, Tcgen05LdMultiplicity, Tcgen05LdShape, Tcgen05Operation, + Tcgen05SourceContract, Tcgen05St, Tcgen05StAdmissionVariant, +}; +use crate::ptx::OperandPattern; +use anyhow::{Result, ensure}; + +use super::*; +use crate::resolve::guards::*; + +pub(in crate::resolve) struct Tcgen05CpMemberRecipe { + pub(in crate::resolve) llvm_suffix: &'static str, + pub(in crate::resolve) ptx_suffix: &'static str, + pub(in crate::resolve) op_suffix: &'static str, + pub(in crate::resolve) selection_stem: &'static str, +} + +pub(in crate::resolve) const TCGEN05_CP_MEMBERS: [Tcgen05CpMember; 17] = [ + Tcgen05CpMember::M128x128bB4x16P64, + Tcgen05CpMember::M128x128bB6x16P32, + Tcgen05CpMember::M128x128b, + Tcgen05CpMember::M128x256bB4x16P64, + Tcgen05CpMember::M128x256bB6x16P32, + Tcgen05CpMember::M32x128bWarpx4B4x16P64, + Tcgen05CpMember::M32x128bWarpx4B6x16P32, + Tcgen05CpMember::M32x128bWarpx4, + Tcgen05CpMember::M4x256bB4x16P64, + Tcgen05CpMember::M4x256bB6x16P32, + Tcgen05CpMember::M4x256b, + Tcgen05CpMember::M64x128bWarpx2Pair0123B4x16P64, + Tcgen05CpMember::M64x128bWarpx2Pair0123B6x16P32, + Tcgen05CpMember::M64x128bWarpx2Pair0123, + Tcgen05CpMember::M64x128bWarpx2Pair0213B4x16P64, + Tcgen05CpMember::M64x128bWarpx2Pair0213B6x16P32, + Tcgen05CpMember::M64x128bWarpx2Pair0213, +]; + +pub(in crate::resolve) fn tcgen05_cp_member_recipe( + member: Tcgen05CpMember, +) -> Tcgen05CpMemberRecipe { + use Tcgen05CpMember::*; + let (llvm_suffix, ptx_suffix, op_suffix, selection_stem) = match member { + M128x128bB4x16P64 => ( + "128x128b.b4x16_p64", + "128x128b.b8x16.b4x16_p64", + "128x128bB4x16P64", + "128x128bb4x16_p64", + ), + M128x128bB6x16P32 => ( + "128x128b.b6x16_p32", + "128x128b.b8x16.b6x16_p32", + "128x128bB6x16P32", + "128x128bb6x16_p32", + ), + M128x128b => ("128x128b", "128x128b", "128x128b", "128x128b"), + M128x256bB4x16P64 => ( + "128x256b.b4x16_p64", + "128x256b.b8x16.b4x16_p64", + "128x256bB4x16P64", + "128x256bb4x16_p64", + ), + M128x256bB6x16P32 => ( + "128x256b.b6x16_p32", + "128x256b.b8x16.b6x16_p32", + "128x256bB6x16P32", + "128x256bb6x16_p32", + ), + M32x128bWarpx4B4x16P64 => ( + "32x128b_warpx4.b4x16_p64", + "32x128b.warpx4.b8x16.b4x16_p64", + "32x128bWarpx4B4x16P64", + "32x128b4x16_p64", + ), + M32x128bWarpx4B6x16P32 => ( + "32x128b_warpx4.b6x16_p32", + "32x128b.warpx4.b8x16.b6x16_p32", + "32x128bWarpx4B6x16P32", + "32x128b6x16_p32", + ), + M32x128bWarpx4 => ( + "32x128b_warpx4", + "32x128b.warpx4", + "32x128bWarpx4", + "32x128", + ), + M4x256bB4x16P64 => ( + "4x256b.b4x16_p64", + "4x256b.b8x16.b4x16_p64", + "4x256bB4x16P64", + "4x256bb4x16_p64", + ), + M4x256bB6x16P32 => ( + "4x256b.b6x16_p32", + "4x256b.b8x16.b6x16_p32", + "4x256bB6x16P32", + "4x256bb6x16_p32", + ), + M4x256b => ("4x256b", "4x256b", "4x256b", "4x256b"), + M64x128bWarpx2Pair0123B4x16P64 => ( + "64x128b_warpx2_01_23.b4x16_p64", + "64x128b.warpx2::01_23.b8x16.b4x16_p64", + "64x128bWarpx2Pair0123B4x16P64", + "64x128_2b4x16_p64", + ), + M64x128bWarpx2Pair0123B6x16P32 => ( + "64x128b_warpx2_01_23.b6x16_p32", + "64x128b.warpx2::01_23.b8x16.b6x16_p32", + "64x128bWarpx2Pair0123B6x16P32", + "64x128_2b6x16_p32", + ), + M64x128bWarpx2Pair0123 => ( + "64x128b_warpx2_01_23", + "64x128b.warpx2::01_23", + "64x128bWarpx2Pair0123", + "64x128_2", + ), + M64x128bWarpx2Pair0213B4x16P64 => ( + "64x128b_warpx2_02_13.b4x16_p64", + "64x128b.warpx2::02_13.b8x16.b4x16_p64", + "64x128bWarpx2Pair0213B4x16P64", + "64x128_1b4x16_p64", + ), + M64x128bWarpx2Pair0213B6x16P32 => ( + "64x128b_warpx2_02_13.b6x16_p32", + "64x128b.warpx2::02_13.b8x16.b6x16_p32", + "64x128bWarpx2Pair0213B6x16P32", + "64x128_1b6x16_p32", + ), + M64x128bWarpx2Pair0213 => ( + "64x128b_warpx2_02_13", + "64x128b.warpx2::02_13", + "64x128bWarpx2Pair0213", + "64x128_1", + ), + }; + Tcgen05CpMemberRecipe { + llvm_suffix, + ptx_suffix, + op_suffix, + selection_stem, + } +} + +pub(in crate::resolve) fn materialize_tcgen05_cp_variant( + base: &OverlayIntrinsic, + admission: &Tcgen05Admission, + variant: &Tcgen05CpAdmissionVariant, +) -> OverlayIntrinsic { + let recipe = tcgen05_cp_member_recipe(variant.member); + let group = match variant.group { + Tcgen05CpGroup::Cg1 => 1, + Tcgen05CpGroup::Cg2 => 2, + }; + let group_suffix = if group == 1 { "" } else { "_cg2" }; + let id_suffix = recipe.llvm_suffix.replace('.', "_"); + let id = format!("tcgen05_cp_{id_suffix}{group_suffix}"); + let mut record = base.clone(); + record.id = id.clone(); + record.abi_id = variant.abi_id.clone(); + record.operation_key = format!("tcgen05.cp.{}.cg{group}", recipe.llvm_suffix); + record.source_record = Some(format!("int_nvvm_tcgen05_cp_{}_cg{group}", id_suffix)); + record.rust_name = id.clone(); + record.public_rust_path = format!("cuda_intrinsics::tcgen05::{id}"); + record.compatibility_rust_paths = vec![format!("cuda_device::tcgen05::{id}")]; + record.dialect_op_type = format!( + "Tcgen05Cp{}{}Op", + recipe.op_suffix, + if group == 1 { "" } else { "Cg2" } + ); + record.dialect_op_name = format!("nvvm.{id}"); + record.llvm_symbol = Some(format!( + "llvm.nvvm.tcgen05.cp.{}.cg{group}", + recipe.llvm_suffix + )); + record.backend_lowerings[0].evidence_profile = admission + .cp_llvm_evidence_profile + .as_ref() + .expect("validated tcgen05 copy LLVM evidence profile") + .clone(); + record.backend_lowerings[1].evidence_profile = admission + .cp_libnvvm_evidence_profile + .as_ref() + .expect("validated tcgen05 copy libNVVM evidence profile") + .clone(); + record.tcgen05 = Some(Tcgen05 { + operation: if group == 1 { + Tcgen05Operation::CpSmemToTmem + } else { + Tcgen05Operation::CpSmemToTmemCg2 + }, + cp: Some(Tcgen05Cp { + member: variant.member, + group: variant.group, + }), + ld: None, + st: None, + mma: None, + adapter: Tcgen05Adapter::TmemDescriptorToVoid, + source_contract: Tcgen05SourceContract::ExactTablegenSelection, + runtime_validation: admission.runtime_validation, + }); + record.expected_ptx.modifiers = std::iter::once("cp".into()) + .chain(std::iter::once(format!("cta_group::{group}"))) + .chain(recipe.ptx_suffix.split('.').map(Into::into)) + .collect(); + record.summary = format!( + "Copies one {} tile from shared memory to tensor memory.", + recipe.ptx_suffix + ); + record +} + +pub(in crate::resolve) const TCGEN05_LD_VARIANTS: [(Tcgen05LdShape, Tcgen05LdMultiplicity); 29] = { + use Tcgen05LdMultiplicity::*; + use Tcgen05LdShape::*; + [ + (M16x64b, X1), + (M16x64b, X2), + (M16x64b, X4), + (M16x64b, X8), + (M16x64b, X16), + (M16x64b, X32), + (M16x64b, X64), + (M16x64b, X128), + (M16x128b, X1), + (M16x128b, X2), + (M16x128b, X4), + (M16x128b, X8), + (M16x128b, X16), + (M16x128b, X32), + (M16x128b, X64), + (M16x256b, X1), + (M16x256b, X2), + (M16x256b, X4), + (M16x256b, X8), + (M16x256b, X16), + (M16x256b, X32), + (M32x32b, X1), + (M32x32b, X2), + (M32x32b, X4), + (M32x32b, X8), + (M32x32b, X16), + (M32x32b, X32), + (M32x32b, X64), + (M32x32b, X128), + ] +}; + +pub(in crate::resolve) const TCGEN05_OFFSET_LDST_VARIANTS: [( + Tcgen05LdShape, + Tcgen05LdMultiplicity, +); 8] = { + use Tcgen05LdMultiplicity::*; + use Tcgen05LdShape::*; + [ + (M16x32bx2, X1), + (M16x32bx2, X2), + (M16x32bx2, X4), + (M16x32bx2, X8), + (M16x32bx2, X16), + (M16x32bx2, X32), + (M16x32bx2, X64), + (M16x32bx2, X128), + ] +}; + +pub(in crate::resolve) fn tcgen05_ld_shape_name(shape: Tcgen05LdShape) -> &'static str { + match shape { + Tcgen05LdShape::M16x32bx2 => "16x32bx2", + Tcgen05LdShape::M16x64b => "16x64b", + Tcgen05LdShape::M16x128b => "16x128b", + Tcgen05LdShape::M16x256b => "16x256b", + Tcgen05LdShape::M32x32b => "32x32b", + } +} + +pub(in crate::resolve) fn tcgen05_ld_multiplicity_name( + multiplicity: Tcgen05LdMultiplicity, +) -> &'static str { + match multiplicity { + Tcgen05LdMultiplicity::X1 => "x1", + Tcgen05LdMultiplicity::X2 => "x2", + Tcgen05LdMultiplicity::X4 => "x4", + Tcgen05LdMultiplicity::X8 => "x8", + Tcgen05LdMultiplicity::X16 => "x16", + Tcgen05LdMultiplicity::X32 => "x32", + Tcgen05LdMultiplicity::X64 => "x64", + Tcgen05LdMultiplicity::X128 => "x128", + } +} + +pub(in crate::resolve) fn tcgen05_ld_register_count(ld: Tcgen05Ld) -> usize { + ld.shape.register_multiplier() * ld.multiplicity.count() +} + +pub(in crate::resolve) fn tcgen05_ld_id(ld: Tcgen05Ld) -> String { + format!( + "tcgen05_ld_{}_{}_{}", + tcgen05_ld_shape_name(ld.shape), + tcgen05_ld_multiplicity_name(ld.multiplicity), + if ld.pack16 { "pack16" } else { "raw" } + ) +} + +pub(in crate::resolve) fn tcgen05_ld_source_record(ld: Tcgen05Ld) -> String { + format!( + "int_nvvm_tcgen05_ld_{}_{}", + tcgen05_ld_shape_name(ld.shape), + tcgen05_ld_multiplicity_name(ld.multiplicity) + ) +} + +pub(in crate::resolve) fn tcgen05_ld_llvm_symbol(ld: Tcgen05Ld) -> String { + format!( + "llvm.nvvm.tcgen05.ld.{}.{}", + tcgen05_ld_shape_name(ld.shape), + tcgen05_ld_multiplicity_name(ld.multiplicity) + ) +} + +pub(in crate::resolve) fn tcgen05_ld_rust_result(ld: Tcgen05Ld) -> String { + match tcgen05_ld_register_count(ld) { + 1 => "u32".into(), + count => format!("[u32; {count}]"), + } +} + +pub(in crate::resolve) fn tcgen05_ld_llvm_result(ld: Tcgen05Ld) -> String { + match tcgen05_ld_register_count(ld) { + 1 => "i32".into(), + count => format!("v{count}i32"), + } +} + +pub(in crate::resolve) fn tcgen05_ld_op_type(ld: Tcgen05Ld) -> String { + let shape = tcgen05_ld_shape_name(ld.shape); + let multiplicity = tcgen05_ld_multiplicity_name(ld.multiplicity) + .strip_prefix('x') + .unwrap(); + format!( + "Tcgen05Ld{shape}X{multiplicity}{}Op", + if ld.pack16 { "Pack16" } else { "Raw" } + ) +} + +pub(in crate::resolve) fn tcgen05_ld_modifiers(ld: Tcgen05Ld) -> Vec { + let mut modifiers = vec![ + "ld".into(), + "sync".into(), + "aligned".into(), + tcgen05_ld_shape_name(ld.shape).into(), + tcgen05_ld_multiplicity_name(ld.multiplicity).into(), + ]; + if ld.pack16 { + modifiers.push("pack::16b".into()); + } + modifiers.push("b32".into()); + modifiers +} + +pub(in crate::resolve) fn tcgen05_ld_operands(ld: Tcgen05Ld) -> Vec { + let result = OperandPattern::RegisterList { + length: tcgen05_ld_register_count(ld), + }; + let mut operands = vec![result, OperandPattern::Address]; + if ld.shape == Tcgen05LdShape::M16x32bx2 { + operands.push(OperandPattern::Immediate); + } + operands +} + +pub(in crate::resolve) fn materialize_tcgen05_ld_variant( + base: &OverlayIntrinsic, + admission: &Tcgen05Admission, + variant: &Tcgen05LdAdmissionVariant, +) -> OverlayIntrinsic { + let ld = Tcgen05Ld { + shape: variant.shape, + multiplicity: variant.multiplicity, + pack16: variant.pack16, + }; + let id = tcgen05_ld_id(ld); + let register_count = tcgen05_ld_register_count(ld); + let rust_result = tcgen05_ld_rust_result(ld); + let llvm_result = tcgen05_ld_llvm_result(ld); + let mode = if ld.pack16 { "pack16" } else { "raw" }; + let mut record = base.clone(); + record.id = id.clone(); + record.abi_id = variant.abi_id.clone(); + record.operation_key = format!( + "tcgen05.ld.{}.{}.{}", + tcgen05_ld_shape_name(ld.shape), + tcgen05_ld_multiplicity_name(ld.multiplicity), + mode + ); + record.source_record = Some(tcgen05_ld_source_record(ld)); + record.rust_name = id.clone(); + let has_half_split_offset = ld.shape == Tcgen05LdShape::M16x32bx2; + record.rust_arguments = if has_half_split_offset { + vec!["u32".into(), "i64".into()] + } else { + vec!["u32".into()] + }; + record.rust_result = rust_result.clone(); + record.must_use = true; + record.public_rust_path = format!("cuda_intrinsics::tcgen05::{id}"); + record.compatibility_rust_paths = vec![format!( + "cuda_device::tcgen05::{}", + if has_half_split_offset { + format!("__{id}") + } else { + id.clone() + } + )]; + record.dialect_op_type = tcgen05_ld_op_type(ld); + record.dialect_op_name = format!("nvvm.{id}"); + record.dialect_operands = if has_half_split_offset { + vec!["i32".into(), "i64".into()] + } else { + vec!["i32".into()] + }; + record.dialect_results = vec!["i32".into(); register_count]; + record.llvm_symbol = Some(tcgen05_ld_llvm_symbol(ld)); + record.resolved_llvm_symbol = None; + record.llvm_arguments = if has_half_split_offset { + vec!["tmem_ptr".into(), "i64".into(), "i1".into()] + } else { + vec!["tmem_ptr".into(), "i1".into()] + }; + record.llvm_results = vec![llvm_result]; + record.ptx_result = rust_result; + record.execution_scope = Tcgen05Operation::Ld.execution_scope().into(); + record.backend_lowerings[0].evidence_profile = if has_half_split_offset { + &admission.offset_llvm_evidence_profile + } else { + &admission.ld_llvm_evidence_profile + } + .as_ref() + .expect("validated tcgen05 load LLVM evidence profile") + .clone(); + record.backend_lowerings[1].evidence_profile = if has_half_split_offset { + &admission.offset_libnvvm_evidence_profile + } else { + &admission.ld_libnvvm_evidence_profile + } + .as_ref() + .expect("validated tcgen05 load libNVVM evidence profile") + .clone(); + record.tcgen05 = Some(Tcgen05 { + operation: Tcgen05Operation::Ld, + cp: None, + ld: Some(ld), + st: None, + mma: None, + adapter: if has_half_split_offset { + Tcgen05Adapter::TmemHalfSplitOffsetInjectPack16ToU32Registers + } else { + Tcgen05Adapter::TmemInjectPack16ToU32Registers + }, + source_contract: Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection, + runtime_validation: admission.runtime_validation, + }); + record.expected_ptx.modifiers = tcgen05_ld_modifiers(ld); + record.expected_ptx.operands = tcgen05_ld_operands(ld); + record.summary = format!( + "Loads {register_count} {} 32-bit register value{} from tensor memory.", + if ld.pack16 { "packed" } else { "raw" }, + if register_count == 1 { "" } else { "s" } + ); + record +} + +pub(in crate::resolve) const TCGEN05_ST_VARIANTS: [(Tcgen05LdShape, Tcgen05LdMultiplicity); 29] = + TCGEN05_LD_VARIANTS; + +pub(in crate::resolve) fn tcgen05_st_register_count(st: Tcgen05St) -> usize { + st.shape.register_multiplier() * st.multiplicity.count() +} + +pub(in crate::resolve) fn tcgen05_st_id(st: Tcgen05St) -> String { + format!( + "tcgen05_st_{}_{}_{}", + tcgen05_ld_shape_name(st.shape), + tcgen05_ld_multiplicity_name(st.multiplicity), + if st.unpack16 { "unpack16" } else { "raw" } + ) +} + +pub(in crate::resolve) fn tcgen05_st_source_record(st: Tcgen05St) -> String { + format!( + "int_nvvm_tcgen05_st_{}_{}", + tcgen05_ld_shape_name(st.shape), + tcgen05_ld_multiplicity_name(st.multiplicity) + ) +} + +pub(in crate::resolve) fn tcgen05_st_llvm_symbol(st: Tcgen05St) -> String { + format!( + "llvm.nvvm.tcgen05.st.{}.{}", + tcgen05_ld_shape_name(st.shape), + tcgen05_ld_multiplicity_name(st.multiplicity) + ) +} + +pub(in crate::resolve) fn tcgen05_st_rust_data(st: Tcgen05St) -> String { + match tcgen05_st_register_count(st) { + 1 => "u32".into(), + count => format!("[u32; {count}]"), + } +} + +pub(in crate::resolve) fn tcgen05_st_llvm_data(st: Tcgen05St) -> String { + match tcgen05_st_register_count(st) { + 1 => "i32".into(), + count => format!("v{count}i32"), + } +} + +pub(in crate::resolve) fn tcgen05_st_op_type(st: Tcgen05St) -> String { + let shape = tcgen05_ld_shape_name(st.shape); + let multiplicity = tcgen05_ld_multiplicity_name(st.multiplicity) + .strip_prefix('x') + .unwrap(); + format!( + "Tcgen05St{shape}X{multiplicity}{}Op", + if st.unpack16 { "Unpack16" } else { "Raw" } + ) +} + +pub(in crate::resolve) fn tcgen05_st_modifiers(st: Tcgen05St) -> Vec { + let mut modifiers = vec![ + "st".into(), + "sync".into(), + "aligned".into(), + tcgen05_ld_shape_name(st.shape).into(), + tcgen05_ld_multiplicity_name(st.multiplicity).into(), + ]; + if st.unpack16 { + modifiers.push("unpack::16b".into()); + } + modifiers.push("b32".into()); + modifiers +} + +pub(in crate::resolve) fn tcgen05_st_operands(st: Tcgen05St) -> Vec { + let data = OperandPattern::RegisterList { + length: tcgen05_st_register_count(st), + }; + let mut operands = vec![OperandPattern::Address]; + if st.shape == Tcgen05LdShape::M16x32bx2 { + operands.push(OperandPattern::Immediate); + } + operands.push(data); + operands +} + +pub(in crate::resolve) fn materialize_tcgen05_st_variant( + base: &OverlayIntrinsic, + admission: &Tcgen05Admission, + variant: &Tcgen05StAdmissionVariant, +) -> OverlayIntrinsic { + let st = Tcgen05St { + shape: variant.shape, + multiplicity: variant.multiplicity, + unpack16: variant.unpack16, + }; + let id = tcgen05_st_id(st); + let register_count = tcgen05_st_register_count(st); + let rust_data = tcgen05_st_rust_data(st); + let llvm_data = tcgen05_st_llvm_data(st); + let mode = if st.unpack16 { "unpack16" } else { "raw" }; + let mut record = base.clone(); + record.id = id.clone(); + record.abi_id = variant.abi_id.clone(); + record.operation_key = format!( + "tcgen05.st.{}.{}.{}", + tcgen05_ld_shape_name(st.shape), + tcgen05_ld_multiplicity_name(st.multiplicity), + mode + ); + record.source_record = Some(tcgen05_st_source_record(st)); + record.rust_name = id.clone(); + let has_half_split_offset = st.shape == Tcgen05LdShape::M16x32bx2; + record.rust_arguments = if has_half_split_offset { + vec!["u32".into(), "i64".into(), rust_data] + } else { + vec!["u32".into(), rust_data] + }; + record.rust_result = "()".into(); + record.must_use = false; + record.public_rust_path = format!("cuda_intrinsics::tcgen05::{id}"); + record.compatibility_rust_paths = vec![format!( + "cuda_device::tcgen05::{}", + if has_half_split_offset { + format!("__{id}") + } else { + id.clone() + } + )]; + record.dialect_op_type = tcgen05_st_op_type(st); + record.dialect_op_name = format!("nvvm.{id}"); + record.dialect_operands = std::iter::once("i32".into()) + .chain(has_half_split_offset.then(|| "i64".into())) + .chain(std::iter::repeat_n("i32".into(), register_count)) + .collect(); + record.dialect_results.clear(); + record.llvm_symbol = Some(tcgen05_st_llvm_symbol(st)); + record.resolved_llvm_symbol = None; + record.llvm_arguments = if has_half_split_offset { + vec!["tmem_ptr".into(), "i64".into(), llvm_data, "i1".into()] + } else { + vec!["tmem_ptr".into(), llvm_data, "i1".into()] + }; + record.llvm_results.clear(); + record.memory = "write".into(); + record.ptx_result = "()".into(); + record.execution_scope = Tcgen05Operation::St.execution_scope().into(); + record.backend_lowerings[0].evidence_profile = if has_half_split_offset { + &admission.offset_llvm_evidence_profile + } else { + &admission.st_llvm_evidence_profile + } + .as_ref() + .expect("validated tcgen05 store LLVM evidence profile") + .clone(); + record.backend_lowerings[1].evidence_profile = if has_half_split_offset { + &admission.offset_libnvvm_evidence_profile + } else { + &admission.st_libnvvm_evidence_profile + } + .as_ref() + .expect("validated tcgen05 store libNVVM evidence profile") + .clone(); + record.tcgen05 = Some(Tcgen05 { + operation: Tcgen05Operation::St, + cp: None, + ld: None, + st: Some(st), + mma: None, + adapter: if has_half_split_offset { + Tcgen05Adapter::TmemHalfSplitOffsetU32RegistersInjectUnpack16ToVoid + } else { + Tcgen05Adapter::TmemU32RegistersInjectUnpack16ToVoid + }, + source_contract: Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection, + runtime_validation: admission.runtime_validation, + }); + record.expected_ptx.modifiers = tcgen05_st_modifiers(st); + record.expected_ptx.operands = tcgen05_st_operands(st); + record.summary = format!( + "Stores {register_count} {} 32-bit register value{} to tensor memory.", + if st.unpack16 { "unpacked" } else { "raw" }, + if register_count == 1 { "" } else { "s" } + ); + record +} + +pub(in crate::resolve) fn validate_tcgen05_ld_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, + tcgen05: &Tcgen05, + ld: Tcgen05Ld, +) -> Result<()> { + let has_half_split_offset = ld.shape == Tcgen05LdShape::M16x32bx2; + let variants: &[(Tcgen05LdShape, Tcgen05LdMultiplicity)] = if has_half_split_offset { + &TCGEN05_OFFSET_LDST_VARIANTS + } else { + &TCGEN05_LD_VARIANTS + }; + ensure!( + variants.contains(&(ld.shape, ld.multiplicity)), + "{} has an unsupported tcgen05 load identity", + policy.id + ); + let id = tcgen05_ld_id(ld); + let source_record = tcgen05_ld_source_record(ld); + let llvm_symbol = tcgen05_ld_llvm_symbol(ld); + let rust_result = tcgen05_ld_rust_result(ld); + let llvm_result = tcgen05_ld_llvm_result(ld); + let register_count = tcgen05_ld_register_count(ld); + let mode = if ld.pack16 { "pack16" } else { "raw" }; + let expected_rust_arguments = if has_half_split_offset { + vec!["u32", "i64"] + } else { + vec!["u32"] + }; + let compatibility_name = if has_half_split_offset { + format!("__{id}") + } else { + id.clone() + }; + let expected_dialect_operands = if has_half_split_offset { + vec!["i32", "i64"] + } else { + vec!["i32"] + }; + let expected_llvm_arguments = if has_half_split_offset { + vec!["tmem_ptr", "i64", "i1"] + } else { + vec!["tmem_ptr", "i1"] + }; + let expected_properties = if has_half_split_offset { + vec![ + "ImmArg", + "ImmArg", + "IntrArgMemOnly", + "IntrConvergent", + "NoCapture", + ] + } else { + vec![ + "ImmArg", + "IntrArgMemOnly", + "IntrConvergent", + "NoCapture", + ] + }; + ensure!( + policy.id == id + && policy.operation_key + == format!( + "tcgen05.ld.{}.{}.{}", + tcgen05_ld_shape_name(ld.shape), + tcgen05_ld_multiplicity_name(ld.multiplicity), + mode + ) + && policy.source.is_none() + && policy.source_record.as_deref() == Some(source_record.as_str()) + && policy.llvm_symbol.as_deref() == Some(llvm_symbol.as_str()) + && policy.resolved_llvm_symbol.is_none() + && declaration.source_record == source_record + && declaration.llvm_name == llvm_symbol, + "{} tcgen05 load identity changed", + policy.id + ); + ensure!( + policy.rust_module == "tcgen05" + && policy.rust_name == id + && policy.rust_arguments == expected_rust_arguments + && policy.rust_result == rust_result + && !policy.safe + && policy.must_use + && policy.safe_allowlist_reason.is_none() + && policy.public_rust_path == format!("cuda_intrinsics::tcgen05::{id}") + && policy.compatibility_rust_paths + == [format!("cuda_device::tcgen05::{compatibility_name}")], + "{} tcgen05 load Rust API changed", + policy.id + ); + ensure!( + policy.dialect_op_type == tcgen05_ld_op_type(ld) + && policy.dialect_op_name == format!("nvvm.{id}") + && policy.dialect_operands == expected_dialect_operands + && policy.dialect_results == vec!["i32"; register_count] + && policy.llvm_arguments == expected_llvm_arguments + && policy.llvm_results == [llvm_result.as_str()] + && declaration.arguments == expected_llvm_arguments + && declaration.results == [llvm_result.as_str()] + && declaration.classes == ["SDPatternOperator", "Intrinsic", "NVVM_TCGEN05_LD"] + && declaration.properties == expected_properties + && declaration.selections.is_empty() + && policy.lowering == "generated_tcgen05", + "{} tcgen05 load carrier or imported declaration changed", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == "read" + && policy.convergent + && policy.execution_scope == Tcgen05Operation::Ld.execution_scope() + && tcgen05.operation == Tcgen05Operation::Ld + && tcgen05.cp.is_none() + && tcgen05.ld == Some(ld) + && tcgen05.st.is_none() + && tcgen05.adapter + == if has_half_split_offset { + Tcgen05Adapter::TmemHalfSplitOffsetInjectPack16ToU32Registers + } else { + Tcgen05Adapter::TmemInjectPack16ToU32Registers + } + && tcgen05.source_contract == Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection + && tcgen05.runtime_validation == RuntimeValidation::Unexecuted, + "{} tcgen05 load semantics changed", + policy.id + ); + ensure!( + policy.minimum_ptx == "8.6" + && policy.minimum_sm.is_none() + && policy.targets == TCGEN05_LLVM_TARGETS + && policy.ptx_isa_version == "8.6" + && policy.ptx_result == rust_result + && policy.expected_ptx.mnemonic == "tcgen05" + && policy.expected_ptx.modifiers == tcgen05_ld_modifiers(ld) + && policy.expected_ptx.operands == tcgen05_ld_operands(ld), + "{} tcgen05 load target or PTX contract changed", + policy.id + ); + validate_tcgen05_backend_routes(policy, "tcgen05 load")?; + ensure_no_other_family_contract(policy, "tcgen05 load")?; + Ok(()) +} + +pub(in crate::resolve) fn validate_tcgen05_st_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, + tcgen05: &Tcgen05, + st: Tcgen05St, +) -> Result<()> { + let has_half_split_offset = st.shape == Tcgen05LdShape::M16x32bx2; + let variants: &[(Tcgen05LdShape, Tcgen05LdMultiplicity)] = if has_half_split_offset { + &TCGEN05_OFFSET_LDST_VARIANTS + } else { + &TCGEN05_ST_VARIANTS + }; + ensure!( + variants.contains(&(st.shape, st.multiplicity)), + "{} has an unsupported tcgen05 store identity", + policy.id + ); + let id = tcgen05_st_id(st); + let source_record = tcgen05_st_source_record(st); + let llvm_symbol = tcgen05_st_llvm_symbol(st); + let rust_data = tcgen05_st_rust_data(st); + let llvm_data = tcgen05_st_llvm_data(st); + let register_count = tcgen05_st_register_count(st); + let mode = if st.unpack16 { "unpack16" } else { "raw" }; + let expected_rust_arguments = if has_half_split_offset { + vec!["u32", "i64", rust_data.as_str()] + } else { + vec!["u32", rust_data.as_str()] + }; + let compatibility_name = if has_half_split_offset { + format!("__{id}") + } else { + id.clone() + }; + let expected_dialect_operands = std::iter::once("i32".to_owned()) + .chain(has_half_split_offset.then(|| "i64".to_owned())) + .chain(std::iter::repeat_n("i32".to_owned(), register_count)) + .collect::>(); + let expected_llvm_arguments = if has_half_split_offset { + vec!["tmem_ptr", "i64", llvm_data.as_str(), "i1"] + } else { + vec!["tmem_ptr", llvm_data.as_str(), "i1"] + }; + let expected_properties = if has_half_split_offset { + vec![ + "ImmArg", + "ImmArg", + "IntrArgMemOnly", + "IntrConvergent", + "NoCapture", + ] + } else { + vec![ + "ImmArg", + "IntrArgMemOnly", + "IntrConvergent", + "NoCapture", + ] + }; + ensure!( + policy.id == id + && policy.operation_key + == format!( + "tcgen05.st.{}.{}.{}", + tcgen05_ld_shape_name(st.shape), + tcgen05_ld_multiplicity_name(st.multiplicity), + mode + ) + && policy.source.is_none() + && policy.source_record.as_deref() == Some(source_record.as_str()) + && policy.llvm_symbol.as_deref() == Some(llvm_symbol.as_str()) + && policy.resolved_llvm_symbol.is_none() + && declaration.source_record == source_record + && declaration.llvm_name == llvm_symbol, + "{} tcgen05 store identity changed", + policy.id + ); + ensure!( + policy.rust_module == "tcgen05" + && policy.rust_name == id + && policy.rust_arguments == expected_rust_arguments + && policy.rust_result == "()" + && !policy.safe + && !policy.must_use + && policy.safe_allowlist_reason.is_none() + && policy.public_rust_path == format!("cuda_intrinsics::tcgen05::{id}") + && policy.compatibility_rust_paths + == [format!("cuda_device::tcgen05::{compatibility_name}")], + "{} tcgen05 store Rust API changed", + policy.id + ); + ensure!( + policy.dialect_op_type == tcgen05_st_op_type(st) + && policy.dialect_op_name == format!("nvvm.{id}") + && policy.dialect_operands == expected_dialect_operands + && policy.dialect_results.is_empty() + && policy.llvm_arguments == expected_llvm_arguments + && policy.llvm_results.is_empty() + && declaration.arguments == expected_llvm_arguments + && declaration.results.is_empty() + && declaration.classes == ["SDPatternOperator", "Intrinsic", "NVVM_TCGEN05_ST"] + && declaration.properties == expected_properties + && declaration.selections.is_empty() + && policy.lowering == "generated_tcgen05", + "{} tcgen05 store carrier or imported declaration changed", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == "write" + && policy.convergent + && policy.execution_scope == Tcgen05Operation::St.execution_scope() + && tcgen05.operation == Tcgen05Operation::St + && tcgen05.cp.is_none() + && tcgen05.ld.is_none() + && tcgen05.st == Some(st) + && tcgen05.adapter + == if has_half_split_offset { + Tcgen05Adapter::TmemHalfSplitOffsetU32RegistersInjectUnpack16ToVoid + } else { + Tcgen05Adapter::TmemU32RegistersInjectUnpack16ToVoid + } + && tcgen05.source_contract == Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection + && tcgen05.runtime_validation == RuntimeValidation::Unexecuted, + "{} tcgen05 store semantics changed", + policy.id + ); + ensure!( + policy.minimum_ptx == "8.6" + && policy.minimum_sm.is_none() + && policy.targets == TCGEN05_LLVM_TARGETS + && policy.ptx_isa_version == "8.6" + && policy.ptx_result == "()" + && policy.expected_ptx.mnemonic == "tcgen05" + && policy.expected_ptx.modifiers == tcgen05_st_modifiers(st) + && policy.expected_ptx.operands == tcgen05_st_operands(st), + "{} tcgen05 store target or PTX contract changed", + policy.id + ); + validate_tcgen05_backend_routes(policy, "tcgen05 store")?; + ensure_no_other_family_contract(policy, "tcgen05 store")?; + Ok(()) +} + +pub(in crate::resolve) fn validate_tcgen05_cp_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, + tcgen05: &Tcgen05, + cp: Tcgen05Cp, +) -> Result<()> { + ensure!( + TCGEN05_CP_MEMBERS.contains(&cp.member), + "{} has an unsupported tcgen05 copy identity", + policy.id + ); + let recipe = tcgen05_cp_member_recipe(cp.member); + let group = match cp.group { + Tcgen05CpGroup::Cg1 => 1, + Tcgen05CpGroup::Cg2 => 2, + }; + let group_suffix = if group == 1 { "" } else { "_cg2" }; + let id_suffix = recipe.llvm_suffix.replace('.', "_"); + let id = format!("tcgen05_cp_{id_suffix}{group_suffix}"); + let operation = if group == 1 { + Tcgen05Operation::CpSmemToTmem + } else { + Tcgen05Operation::CpSmemToTmemCg2 + }; + let op_type = format!( + "Tcgen05Cp{}{}Op", + recipe.op_suffix, + if group == 1 { "" } else { "Cg2" } + ); + let source_record = format!("int_nvvm_tcgen05_cp_{}_cg{group}", id_suffix); + let llvm_symbol = format!("llvm.nvvm.tcgen05.cp.{}.cg{group}", recipe.llvm_suffix); + let modifiers = std::iter::once("cp".into()) + .chain(std::iter::once(format!("cta_group::{group}"))) + .chain(recipe.ptx_suffix.split('.').map(Into::into)) + .collect::>(); + ensure!( + policy.id == id + && policy.operation_key == format!("tcgen05.cp.{}.cg{group}", recipe.llvm_suffix) + && policy.source.is_none() + && policy.source_record.as_deref() == Some(source_record.as_str()) + && policy.llvm_symbol.as_deref() == Some(llvm_symbol.as_str()) + && policy.resolved_llvm_symbol.is_none() + && declaration.source_record == source_record + && declaration.llvm_name == llvm_symbol, + "{} tcgen05 copy identity changed", + policy.id + ); + ensure!( + policy.rust_module == "tcgen05" + && policy.rust_name == id + && policy.rust_arguments == ["u32", "u64"] + && policy.rust_result == "()" + && !policy.safe + && !policy.must_use + && policy.safe_allowlist_reason.is_none() + && policy.public_rust_path == format!("cuda_intrinsics::tcgen05::{id}") + && policy.compatibility_rust_paths == [format!("cuda_device::tcgen05::{id}")], + "{} tcgen05 copy Rust API changed", + policy.id + ); + ensure!( + policy.dialect_op_type == op_type + && policy.dialect_op_name == format!("nvvm.{id}") + && policy.dialect_operands == ["i32", "i64"] + && policy.dialect_results.is_empty() + && policy.llvm_arguments == ["tmem_ptr", "i64"] + && policy.llvm_results.is_empty() + && declaration.arguments == ["tmem_ptr", "i64"] + && declaration.results.is_empty() + && declaration.classes == ["SDPatternOperator", "Intrinsic"] + && declaration.properties + == [ + "IntrConvergent", + "IntrInaccessibleMemOrArgMemOnly", + "NoCapture", + ] + && policy.lowering == "generated_tcgen05", + "{} tcgen05 copy carrier or declaration changed", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == "read_write" + && policy.convergent + && policy.execution_scope == "thread" + && tcgen05.operation == operation + && tcgen05.ld.is_none() + && tcgen05.st.is_none() + && tcgen05.adapter == Tcgen05Adapter::TmemDescriptorToVoid + && tcgen05.source_contract == Tcgen05SourceContract::ExactTablegenSelection + && tcgen05.runtime_validation == RuntimeValidation::Unexecuted, + "{} tcgen05 copy semantics changed", + policy.id + ); + ensure!( + policy.minimum_ptx == "8.6" + && policy.minimum_sm.is_none() + && policy.targets == TCGEN05_LLVM_TARGETS + && policy.ptx_isa_version == "8.6" + && policy.ptx_result == "()" + && policy.expected_ptx.mnemonic == "tcgen05" + && policy.expected_ptx.modifiers == modifiers + && policy.expected_ptx.operands == [OperandPattern::Address, OperandPattern::Register], + "{} tcgen05 copy target or PTX contract changed", + policy.id + ); + validate_tcgen05_backend_routes(policy, "tcgen05 copy")?; + ensure!( + declaration.selections.len() == 1, + "{} must keep one exact tcgen05 copy selection", + policy.id + ); + let selection = &declaration.selections[0]; + ensure!( + selection.source_record == format!("TCGEN05_CP_{}_cg{group}", recipe.selection_stem) + && selection.asm + == format!( + "tcgen05.cp.cta_group::{group}.{} \t[$tmem_addr], $sdesc;", + recipe.ptx_suffix + ) + && selection.predicates == ["Subtarget->hasTcgen05InstSupport()"] + && selection.constraints.is_empty(), + "{} exact tcgen05 copy selection changed", + policy.id + ); + ensure_no_other_family_contract(policy, "tcgen05 copy")?; + Ok(()) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/tcgen05/mma.rs b/crates/cuda-intrinsics-gen/src/resolve/families/tcgen05/mma.rs new file mode 100644 index 0000000000..d81f70fc2d --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/tcgen05/mma.rs @@ -0,0 +1,937 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, CatalogTargetRequirement, ImportedIntrinsic, IntrinsicBackend, + OverlayBackendLowering, OverlayIntrinsic, RuntimeValidation, TargetContract, + TargetSelectorBinding, Tcgen05, Tcgen05Adapter, Tcgen05Admission, Tcgen05Mma, + Tcgen05MmaAdmissionVariant, Tcgen05MmaAlias, Tcgen05MmaBUsage, Tcgen05MmaFixedSelectors, + Tcgen05MmaForm, Tcgen05MmaKind, Tcgen05MmaSelectorLayout, Tcgen05Operation, + Tcgen05SourceContract, +}; +use crate::ptx::{InstructionPattern, OperandPattern}; +use anyhow::{Result, ensure}; +use std::collections::BTreeSet; + +use super::*; +use crate::resolve::guards::*; +use crate::resolve::targets::*; + +pub(in crate::resolve) const TCGEN05_MMA_FORMS: [Tcgen05MmaForm; 14] = [ + Tcgen05MmaForm::Shared, + Tcgen05MmaForm::Tensor, + Tcgen05MmaForm::TensorAshift, + Tcgen05MmaForm::SpShared, + Tcgen05MmaForm::SpTensor, + Tcgen05MmaForm::SpTensorAshift, + Tcgen05MmaForm::WsShared, + Tcgen05MmaForm::WsSharedZeroColMask, + Tcgen05MmaForm::WsSpShared, + Tcgen05MmaForm::WsSpSharedZeroColMask, + Tcgen05MmaForm::WsSpTensor, + Tcgen05MmaForm::WsSpTensorZeroColMask, + Tcgen05MmaForm::WsTensor, + Tcgen05MmaForm::WsTensorZeroColMask, +]; +pub(in crate::resolve) const TCGEN05_MMA_ALIASES: [Tcgen05MmaAlias; 5] = [ + Tcgen05MmaAlias::E4m3, + Tcgen05MmaAlias::E5m2, + Tcgen05MmaAlias::E2m3, + Tcgen05MmaAlias::E3m2, + Tcgen05MmaAlias::E2m1, +]; +pub(in crate::resolve) const TCGEN05_MMA_KINDS: [Tcgen05MmaKind; 4] = [ + Tcgen05MmaKind::F16, + Tcgen05MmaKind::F8f6f4, + Tcgen05MmaKind::I8, + Tcgen05MmaKind::Tf32, +]; +pub(in crate::resolve) const TCGEN05_MMA_DIALECT_OP_TYPE: &str = "Tcgen05MmaOp"; +pub(in crate::resolve) const TCGEN05_MMA_DIALECT_OP_NAME: &str = "nvvm.tcgen05_mma"; +pub(in crate::resolve) fn tcgen05_mma_form_name(form: Tcgen05MmaForm) -> &'static str { + match form { + Tcgen05MmaForm::Shared => "shared", + Tcgen05MmaForm::Tensor => "tensor", + Tcgen05MmaForm::TensorAshift => "tensor_ashift", + Tcgen05MmaForm::SpShared => "sp_shared", + Tcgen05MmaForm::SpTensor => "sp_tensor", + Tcgen05MmaForm::SpTensorAshift => "sp_tensor_ashift", + Tcgen05MmaForm::WsShared => "ws_shared", + Tcgen05MmaForm::WsSharedZeroColMask => "ws_shared_zero_col_mask", + Tcgen05MmaForm::WsSpShared => "ws_sp_shared", + Tcgen05MmaForm::WsSpSharedZeroColMask => "ws_sp_shared_zero_col_mask", + Tcgen05MmaForm::WsSpTensor => "ws_sp_tensor", + Tcgen05MmaForm::WsSpTensorZeroColMask => "ws_sp_tensor_zero_col_mask", + Tcgen05MmaForm::WsTensor => "ws_tensor", + Tcgen05MmaForm::WsTensorZeroColMask => "ws_tensor_zero_col_mask", + } +} + +pub(in crate::resolve) fn tcgen05_mma_alias_name(alias: Tcgen05MmaAlias) -> &'static str { + match alias { + Tcgen05MmaAlias::E4m3 => "e4m3", + Tcgen05MmaAlias::E5m2 => "e5m2", + Tcgen05MmaAlias::E2m3 => "e2m3", + Tcgen05MmaAlias::E3m2 => "e3m2", + Tcgen05MmaAlias::E2m1 => "e2m1", + } +} + +pub(in crate::resolve) fn tcgen05_mma_kind_name(kind: Tcgen05MmaKind) -> &'static str { + match kind { + Tcgen05MmaKind::F16 => "f16", + Tcgen05MmaKind::Tf32 => "tf32", + Tcgen05MmaKind::F8f6f4 => "f8f6f4", + Tcgen05MmaKind::I8 => "i8", + } +} + +pub(in crate::resolve) fn tcgen05_mma_b_usage_name(usage: Tcgen05MmaBUsage) -> &'static str { + match usage { + Tcgen05MmaBUsage::Discard => "discard", + Tcgen05MmaBUsage::Fill => "fill", + Tcgen05MmaBUsage::Use => "use", + Tcgen05MmaBUsage::LastUse => "lastuse", + } +} + +pub(in crate::resolve) fn tcgen05_mma_is_ws(form: Tcgen05MmaForm) -> bool { + matches!( + form, + Tcgen05MmaForm::WsShared + | Tcgen05MmaForm::WsSharedZeroColMask + | Tcgen05MmaForm::WsSpShared + | Tcgen05MmaForm::WsSpSharedZeroColMask + | Tcgen05MmaForm::WsSpTensor + | Tcgen05MmaForm::WsSpTensorZeroColMask + | Tcgen05MmaForm::WsTensor + | Tcgen05MmaForm::WsTensorZeroColMask + ) +} + +pub(in crate::resolve) fn tcgen05_mma_is_sparse(form: Tcgen05MmaForm) -> bool { + matches!( + form, + Tcgen05MmaForm::SpShared + | Tcgen05MmaForm::SpTensor + | Tcgen05MmaForm::SpTensorAshift + | Tcgen05MmaForm::WsSpShared + | Tcgen05MmaForm::WsSpSharedZeroColMask + | Tcgen05MmaForm::WsSpTensor + | Tcgen05MmaForm::WsSpTensorZeroColMask + ) +} + +pub(in crate::resolve) fn tcgen05_mma_is_tensor_a(form: Tcgen05MmaForm) -> bool { + matches!( + form, + Tcgen05MmaForm::Tensor + | Tcgen05MmaForm::TensorAshift + | Tcgen05MmaForm::SpTensor + | Tcgen05MmaForm::SpTensorAshift + | Tcgen05MmaForm::WsSpTensor + | Tcgen05MmaForm::WsSpTensorZeroColMask + | Tcgen05MmaForm::WsTensor + | Tcgen05MmaForm::WsTensorZeroColMask + ) +} + +pub(in crate::resolve) fn tcgen05_mma_is_ashift(form: Tcgen05MmaForm) -> bool { + matches!( + form, + Tcgen05MmaForm::TensorAshift | Tcgen05MmaForm::SpTensorAshift + ) +} + +pub(in crate::resolve) fn tcgen05_mma_has_zero_col_mask(form: Tcgen05MmaForm) -> bool { + matches!( + form, + Tcgen05MmaForm::WsSharedZeroColMask + | Tcgen05MmaForm::WsSpSharedZeroColMask + | Tcgen05MmaForm::WsSpTensorZeroColMask + | Tcgen05MmaForm::WsTensorZeroColMask + ) +} + +pub(in crate::resolve) fn tcgen05_mma_source_record(form: Tcgen05MmaForm) -> String { + format!("int_nvvm_tcgen05_mma_{}", tcgen05_mma_form_name(form)) +} + +pub(in crate::resolve) fn tcgen05_mma_llvm_symbol(form: Tcgen05MmaForm) -> String { + let suffix = match form { + Tcgen05MmaForm::Shared => "shared", + Tcgen05MmaForm::Tensor => "tensor", + Tcgen05MmaForm::TensorAshift => "tensor.ashift", + Tcgen05MmaForm::SpShared => "sp.shared", + Tcgen05MmaForm::SpTensor => "sp.tensor", + Tcgen05MmaForm::SpTensorAshift => "sp.tensor.ashift", + Tcgen05MmaForm::WsShared => "ws.shared", + Tcgen05MmaForm::WsSharedZeroColMask => "ws.shared.zero_col_mask", + Tcgen05MmaForm::WsSpShared => "ws.sp.shared", + Tcgen05MmaForm::WsSpSharedZeroColMask => "ws.sp.shared.zero_col_mask", + Tcgen05MmaForm::WsSpTensor => "ws.sp.tensor", + Tcgen05MmaForm::WsSpTensorZeroColMask => "ws.sp.tensor.zero_col_mask", + Tcgen05MmaForm::WsTensor => "ws.tensor", + Tcgen05MmaForm::WsTensorZeroColMask => "ws.tensor.zero_col_mask", + }; + format!("llvm.nvvm.tcgen05.mma.{suffix}") +} + +pub(in crate::resolve) fn tcgen05_mma_llvm_arguments(form: Tcgen05MmaForm) -> Vec { + let mut arguments = vec![ + "tmem_ptr".into(), + if tcgen05_mma_is_tensor_a(form) { + "tmem_ptr".into() + } else { + "i64".into() + }, + "i64".into(), + "i32".into(), + "i1".into(), + ]; + if tcgen05_mma_is_sparse(form) { + arguments.push("tmem_ptr".into()); + } + if tcgen05_mma_has_zero_col_mask(form) { + arguments.push("i64".into()); + } + arguments.extend(["i32".into(), "i32".into(), "i32".into()]); + arguments +} + +pub(in crate::resolve) fn tcgen05_mma_selector_layout( + form: Tcgen05MmaForm, +) -> Tcgen05MmaSelectorLayout { + let first = + 5 + u8::from(tcgen05_mma_is_sparse(form)) + u8::from(tcgen05_mma_has_zero_col_mask(form)); + if tcgen05_mma_is_ws(form) { + Tcgen05MmaSelectorLayout::WarpSpecialized { + kind_argument: first, + b_buffer_argument: first + 1, + b_usage_argument: first + 2, + } + } else { + Tcgen05MmaSelectorLayout::Base { + kind_argument: first, + cta_group_argument: first + 1, + collector_a_argument: first + 2, + collector_a_upper_exclusive: if tcgen05_mma_is_ashift(form) { 2 } else { 4 }, + } + } +} + +pub(in crate::resolve) fn tcgen05_mma_imported_properties(form: Tcgen05MmaForm) -> Vec { + let layout = tcgen05_mma_selector_layout(form); + let (kind, second, third, second_lower, second_upper, third_upper) = match layout { + Tcgen05MmaSelectorLayout::Base { + kind_argument, + cta_group_argument, + collector_a_argument, + collector_a_upper_exclusive, + } => ( + kind_argument, + cta_group_argument, + collector_a_argument, + 1, + 3, + collector_a_upper_exclusive, + ), + Tcgen05MmaSelectorLayout::WarpSpecialized { + kind_argument, + b_buffer_argument, + b_usage_argument, + } => (kind_argument, b_buffer_argument, b_usage_argument, 0, 4, 4), + }; + let mut properties = vec![ + format!("ImmArg"), + format!("ImmArg"), + format!("ImmArg"), + "IntrArgMemOnly".into(), + format!("Range"), + format!("Range"), + format!("Range"), + "WriteOnly".into(), + ]; + if tcgen05_mma_is_tensor_a(form) { + properties.push("ReadOnly".into()); + } + if form == Tcgen05MmaForm::Tensor { + properties.extend([ + "ArgInfo".into(), + "ArgInfo".into(), + "ArgInfo".into(), + ]); + } + properties.sort(); + properties +} + +pub(in crate::resolve) fn tcgen05_mma_selection_asm( + form: Tcgen05MmaForm, + kind: Tcgen05MmaKind, + cta_group: u8, + collector_a: Option<&str>, + b_buffer: Option, + b_usage: Option, +) -> String { + let mut head = "tcgen05.mma".to_owned(); + if tcgen05_mma_is_ws(form) { + head.push_str(".ws"); + } + if tcgen05_mma_is_sparse(form) { + head.push_str(".sp"); + } + head.push_str(&format!( + ".cta_group::{cta_group}.kind::{}", + tcgen05_mma_kind_name(kind) + )); + if tcgen05_mma_is_ws(form) { + head.push_str(&format!( + ".collector::b{}::{}", + b_buffer.expect("warp-specialized B buffer"), + tcgen05_mma_b_usage_name(b_usage.expect("warp-specialized B usage")) + )); + } else { + head.push_str(&format!( + ".collector::a::{}", + collector_a.expect("base collector A usage") + )); + if tcgen05_mma_is_ashift(form) { + head.push_str(".ashift"); + } + } + + let a = if tcgen05_mma_is_tensor_a(form) { + "[$a]" + } else { + "$a" + }; + let mut operands = format!("[$dtmem], {a}, $b"); + if tcgen05_mma_is_sparse(form) { + operands.push_str(", [$spmetadata]"); + } + operands.push_str(", $idesc, $enable_inp_d"); + if tcgen05_mma_has_zero_col_mask(form) { + operands.push_str(", $zero_col_mask"); + } + format!("{head} {operands};") +} + +pub(in crate::resolve) fn tcgen05_mma_all_selection_asms(form: Tcgen05MmaForm) -> BTreeSet { + if tcgen05_mma_is_ws(form) { + return [ + Tcgen05MmaBUsage::Discard, + Tcgen05MmaBUsage::LastUse, + Tcgen05MmaBUsage::Fill, + Tcgen05MmaBUsage::Use, + ] + .into_iter() + .flat_map(|usage| { + TCGEN05_MMA_KINDS.into_iter().flat_map(move |kind| { + (0..4).map(move |buffer| { + tcgen05_mma_selection_asm(form, kind, 1, None, Some(buffer), Some(usage)) + }) + }) + }) + .collect(); + } + TCGEN05_MMA_KINDS + .into_iter() + .flat_map(|kind| { + (1..=2).flat_map(move |group| { + ["discard", "lastuse", "fill", "use"].map(move |usage| { + tcgen05_mma_selection_asm(form, kind, group, Some(usage), None, None) + }) + }) + }) + .collect() +} + +pub(in crate::resolve) fn tcgen05_mma_valid_selection_asms( + form: Tcgen05MmaForm, +) -> BTreeSet { + if !tcgen05_mma_is_ashift(form) { + return tcgen05_mma_all_selection_asms(form); + } + TCGEN05_MMA_KINDS + .into_iter() + .flat_map(|kind| { + (1..=2).flat_map(move |group| { + ["discard", "lastuse"].map(move |usage| { + tcgen05_mma_selection_asm(form, kind, group, Some(usage), None, None) + }) + }) + }) + .collect() +} + +pub(in crate::resolve) fn tcgen05_mma_target_contract( + kind: Tcgen05MmaKind, + alternatives: &[(&str, &str)], +) -> TargetContract { + TargetContract { + selectors: vec![TargetSelectorBinding { + name: "kind".into(), + value: tcgen05_mma_kind_name(kind).into(), + }], + alternatives: alternatives + .iter() + .map( + |(target, minimum_ptx)| crate::model::TargetContractAlternative { + target: (*target).into(), + minimum_ptx: (*minimum_ptx).into(), + }, + ) + .collect(), + } +} + +pub(in crate::resolve) fn expected_tcgen05_mma_target_contracts( + backend: IntrinsicBackend, +) -> Vec { + const LLVM_COMMON: &[(&str, &str)] = &[ + ("sm_100a", "8.6"), + ("sm_100f", "8.8"), + ("sm_101a", "8.6"), + ("sm_101f", "8.8"), + ("sm_103a", "8.8"), + ("sm_103f", "8.8"), + ("sm_110a", "9.0"), + ("sm_110f", "9.0"), + ]; + const LLVM_I8: &[(&str, &str)] = &[("sm_100a", "8.6"), ("sm_101a", "8.6"), ("sm_110a", "9.0")]; + const LIBNVVM_COMMON: &[(&str, &str)] = &[ + ("sm_100a", "8.6"), + ("sm_100f", "8.8"), + ("sm_103a", "8.8"), + ("sm_103f", "8.8"), + ("sm_110a", "9.0"), + ("sm_110f", "9.0"), + ]; + const LIBNVVM_I8: &[(&str, &str)] = &[("sm_100a", "8.6"), ("sm_110a", "9.0")]; + TCGEN05_MMA_KINDS + .into_iter() + .map(|kind| { + let alternatives = match (backend, kind) { + (IntrinsicBackend::LlvmNvptx, Tcgen05MmaKind::I8) => LLVM_I8, + (IntrinsicBackend::LlvmNvptx, _) => LLVM_COMMON, + (IntrinsicBackend::LibNvvm, Tcgen05MmaKind::I8) => LIBNVVM_I8, + (IntrinsicBackend::LibNvvm, _) => LIBNVVM_COMMON, + }; + tcgen05_mma_target_contract(kind, alternatives) + }) + .collect() +} + +pub(in crate::resolve) fn tcgen05_mma_expected_ptx( + form: Tcgen05MmaForm, + alias: Option, +) -> InstructionPattern { + let kind = if alias.is_some() { + Tcgen05MmaKind::F8f6f4 + } else { + Tcgen05MmaKind::F16 + }; + let asm = if tcgen05_mma_is_ws(form) { + tcgen05_mma_selection_asm( + form, + kind, + 1, + None, + Some(0), + Some(Tcgen05MmaBUsage::Discard), + ) + } else { + tcgen05_mma_selection_asm(form, kind, 1, Some("discard"), None, None) + }; + let head = asm.split_whitespace().next().unwrap(); + let mut components = head.split('.'); + let mnemonic = components.next().unwrap().into(); + let modifiers = components.map(str::to_owned).collect(); + let mut operands = vec![OperandPattern::Address]; + operands.push(if tcgen05_mma_is_tensor_a(form) { + OperandPattern::Address + } else { + OperandPattern::Register + }); + operands.push(OperandPattern::Register); + if tcgen05_mma_is_sparse(form) { + operands.push(OperandPattern::Address); + } + operands.extend([ + OperandPattern::Register, + OperandPattern::Exact { + value: "%enable_pred".into(), + }, + ]); + if tcgen05_mma_has_zero_col_mask(form) { + operands.push(OperandPattern::Register); + } + InstructionPattern { + mnemonic, + modifiers, + operands, + } +} + +pub(in crate::resolve) fn tcgen05_mma_rust_arguments( + form: Tcgen05MmaForm, + alias: Option, +) -> Vec { + if alias.is_some() && form == Tcgen05MmaForm::WsTensor { + return ["u32", "u32", "u64", "u64", "u32", "bool"] + .into_iter() + .map(str::to_owned) + .collect(); + } + let mut arguments = tcgen05_mma_llvm_arguments(form); + if alias.is_some() { + arguments.truncate(arguments.len() - 3); + } + arguments + .iter() + .map(|argument| match argument.as_str() { + "tmem_ptr" | "i32" => "u32", + "i64" => "u64", + "i1" => "bool", + _ => unreachable!("closed tcgen05 MMA LLVM type"), + }) + .map(str::to_owned) + .collect() +} + +pub(in crate::resolve) fn tcgen05_mma_dialect_operands( + form: Tcgen05MmaForm, + _alias: Option, +) -> Vec { + let mut arguments = tcgen05_mma_llvm_arguments(form); + arguments.truncate(arguments.len() - 3); + arguments + .iter() + .map(|argument| match argument.as_str() { + "tmem_ptr" | "i32" => "i32", + "i64" => "i64", + "i1" => "i1", + _ => unreachable!("closed tcgen05 MMA LLVM type"), + }) + .map(str::to_owned) + .collect() +} + +pub(in crate::resolve) fn tcgen05_mma_public_id( + form: Tcgen05MmaForm, + alias: Option, +) -> String { + alias.map_or_else( + || format!("tcgen05_mma_{}", tcgen05_mma_form_name(form)), + |alias| { + let alias = tcgen05_mma_alias_name(alias); + if form == Tcgen05MmaForm::WsTensor { + format!("tcgen05_mma_ws_{alias}") + } else { + format!("tcgen05_mma_{alias}") + } + }, + ) +} + +pub(in crate::resolve) fn tcgen05_mma_operation_key( + form: Tcgen05MmaForm, + alias: Option, +) -> String { + let base = tcgen05_mma_form_name(form).replace('_', "."); + alias.map_or_else( + || format!("tcgen05.mma.{base}"), + |alias| format!("tcgen05.mma.{base}.{}", tcgen05_mma_alias_name(alias)), + ) +} + +pub(in crate::resolve) fn tcgen05_mma_adapter( + form: Tcgen05MmaForm, + alias: Option, +) -> Tcgen05Adapter { + match alias { + Some(_) if form == Tcgen05MmaForm::WsTensor => { + Tcgen05Adapter::MmaWsFixedSelectorsDropLegacyADescriptor + } + _ => Tcgen05Adapter::MmaDirectSelectors, + } +} + +pub(in crate::resolve) fn materialize_tcgen05_mma_variant( + admission: &Tcgen05Admission, + variant: &Tcgen05MmaAdmissionVariant, + llvm_target: &CatalogTargetRequirement, + libnvvm_target: &CatalogTargetRequirement, +) -> OverlayIntrinsic { + let form = variant.form; + let alias = variant.alias; + let id = tcgen05_mma_public_id(form, alias); + let operation_key = tcgen05_mma_operation_key(form, alias); + let llvm_arguments = tcgen05_mma_llvm_arguments(form); + let rust_arguments = tcgen05_mma_rust_arguments(form, alias); + let dialect_operands = tcgen05_mma_dialect_operands(form, alias); + let fixed_selectors = alias.map(|_| Tcgen05MmaFixedSelectors { + kind: Tcgen05MmaKind::F8f6f4, + b_buffer: 0, + b_usage: Tcgen05MmaBUsage::Discard, + }); + + OverlayIntrinsic { + id: id.clone(), + abi_id: variant.abi_id.clone(), + operation_key, + family: "tcgen05".into(), + source: None, + source_record: Some(tcgen05_mma_source_record(form)), + rust_module: "tcgen05".into(), + rust_name: id.clone(), + rust_arguments, + rust_result: "()".into(), + safe: false, + must_use: false, + safe_allowlist_reason: None, + public_rust_path: format!("cuda_intrinsics::tcgen05::{id}"), + compatibility_rust_paths: vec![if alias.is_some() { + format!("cuda_device::tcgen05::{id}") + } else { + format!("cuda_device::tcgen05::__{id}") + }], + dialect_op_type: TCGEN05_MMA_DIALECT_OP_TYPE.into(), + dialect_op_name: TCGEN05_MMA_DIALECT_OP_NAME.into(), + dialect_operands, + dialect_results: vec![], + llvm_symbol: Some(tcgen05_mma_llvm_symbol(form)), + resolved_llvm_symbol: None, + llvm_arguments, + llvm_results: vec![], + pure: false, + memory: "read_write".into(), + // TableGen omits IntrConvergent, but collective MMA must not move across control flow. + convergent: true, + execution_scope: "thread".into(), + minimum_ptx: "8.6".into(), + minimum_sm: None, + ptx_result: "()".into(), + targets: TCGEN05_LLVM_TARGETS.into(), + ptx_isa_version: "8.6".into(), + ptx_isa_section: "Tensor Memory tcgen05 instructions".into(), + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#tensor-memory".into(), + lowering: "generated_tcgen05_mma".into(), + backend_lowerings: vec![ + OverlayBackendLowering { + backend: IntrinsicBackend::LlvmNvptx, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: admission + .mma_llvm_evidence_profile + .as_ref() + .expect("validated tcgen05 MMA LLVM evidence profile") + .clone(), + targets: None, + minimum_ptx: None, + minimum_sm: None, + }, + OverlayBackendLowering { + backend: IntrinsicBackend::LibNvvm, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: admission + .mma_libnvvm_evidence_profile + .as_ref() + .expect("validated tcgen05 MMA libNVVM evidence profile") + .clone(), + targets: None, + minimum_ptx: None, + minimum_sm: None, + }, + ], + packed_atomic: None, + redux: None, + vote: None, + active_mask: None, + warp_match: None, + warp_barrier: None, + warp_shuffle: None, + dot_product: None, + packed_alu: None, + integer_minmax: None, + packed_conversion: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cp_async_copy: None, + cp_async_control: None, + cp_async_mbarrier: None, + mbarrier_basic: None, + movmatrix: None, + mbarrier_extended: None, + register_mma: None, + sparse_mma: None, + prmt: None, + cluster_barrier: None, + wgmma_control: None, + special_register: None, + debug_control: None, + cluster_memory: None, + clc: None, + tma: None, + tcgen05: Some(Tcgen05 { + operation: Tcgen05Operation::Mma, + cp: None, + ld: None, + st: None, + mma: Some(Tcgen05Mma { + form, + selector_layout: tcgen05_mma_selector_layout(form), + fixed_selectors, + alias, + llvm_target: llvm_target.clone(), + libnvvm_target: libnvvm_target.clone(), + }), + adapter: tcgen05_mma_adapter(form, alias), + source_contract: Tcgen05SourceContract::TablegenSelectionChangesPtx, + runtime_validation: admission.runtime_validation, + }), + ldmatrix_variant: None, + ldmatrix_safety: None, + ldmatrix_adapter: None, + selected_address_space: None, + expected_ptx: tcgen05_mma_expected_ptx(form, alias), + summary: match (form, alias) { + (Tcgen05MmaForm::WsTensor, Some(_)) => { + "Issues one f8f6f4 warp-specialized tensor-memory MMA.".into() + } + (Tcgen05MmaForm::Shared, Some(_)) => { + "Issues one f8f6f4 standard tensor-memory MMA.".into() + } + _ => "Issues one selector-controlled tensor-memory MMA.".into(), + }, + } +} + +pub(in crate::resolve) fn validate_tcgen05_mma_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, + tcgen05: &Tcgen05, + mma: &Tcgen05Mma, +) -> Result<()> { + let form = mma.form; + let alias = mma.alias; + ensure!( + match alias { + Some(alias) => { + matches!(form, Tcgen05MmaForm::Shared | Tcgen05MmaForm::WsTensor) + && TCGEN05_MMA_ALIASES.contains(&alias) + } + None => TCGEN05_MMA_FORMS.contains(&form), + }, + "{} has an unsupported tcgen05 MMA identity", + policy.id + ); + let id = tcgen05_mma_public_id(form, alias); + let operation_key = tcgen05_mma_operation_key(form, alias); + let llvm_arguments = tcgen05_mma_llvm_arguments(form); + let expected_fixed = alias.map(|_| Tcgen05MmaFixedSelectors { + kind: Tcgen05MmaKind::F8f6f4, + b_buffer: 0, + b_usage: Tcgen05MmaBUsage::Discard, + }); + ensure!( + policy.id == id + && policy.operation_key == operation_key + && policy.source.is_none() + && policy.source_record.as_deref() == Some(tcgen05_mma_source_record(form).as_str()) + && policy.llvm_symbol.as_deref() == Some(tcgen05_mma_llvm_symbol(form).as_str()) + && policy.resolved_llvm_symbol.is_none() + && declaration.source_record == tcgen05_mma_source_record(form) + && declaration.llvm_name == tcgen05_mma_llvm_symbol(form), + "{} tcgen05 MMA identity changed", + policy.id + ); + ensure!( + policy.rust_module == "tcgen05" + && policy.rust_name == id + && policy.rust_arguments == tcgen05_mma_rust_arguments(form, alias) + && policy.rust_result == "()" + && !policy.safe + && !policy.must_use + && policy.safe_allowlist_reason.is_none() + && policy.public_rust_path == format!("cuda_intrinsics::tcgen05::{id}") + && policy.compatibility_rust_paths + == [if alias.is_some() { + format!("cuda_device::tcgen05::{id}") + } else { + format!("cuda_device::tcgen05::__{id}") + }], + "{} tcgen05 MMA Rust API changed", + policy.id + ); + ensure!( + policy.dialect_op_type == TCGEN05_MMA_DIALECT_OP_TYPE + && policy.dialect_op_name == TCGEN05_MMA_DIALECT_OP_NAME + && policy.dialect_operands == tcgen05_mma_dialect_operands(form, alias) + && policy.dialect_results.is_empty() + && policy.llvm_arguments == llvm_arguments + && policy.llvm_results.is_empty() + && declaration.arguments == llvm_arguments + && declaration.results.is_empty() + && declaration.classes + == [ + "SDPatternOperator", + "Intrinsic", + "DefaultAttrsIntrinsic", + "DefaultAttrsIntrinsicFlags", + ] + && declaration.properties == tcgen05_mma_imported_properties(form) + && policy.lowering == "generated_tcgen05_mma", + "{} tcgen05 MMA carrier or imported declaration changed", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == "read_write" + && policy.convergent + && policy.execution_scope == "thread" + && tcgen05.operation == Tcgen05Operation::Mma + && tcgen05.cp.is_none() + && tcgen05.ld.is_none() + && tcgen05.st.is_none() + && mma.selector_layout == tcgen05_mma_selector_layout(form) + && mma.fixed_selectors == expected_fixed + && tcgen05.adapter == tcgen05_mma_adapter(form, alias) + && tcgen05.source_contract == Tcgen05SourceContract::TablegenSelectionChangesPtx + && tcgen05.runtime_validation == RuntimeValidation::Unexecuted, + "{} tcgen05 MMA semantics or selector contract changed", + policy.id + ); + + let expected_llvm_contracts = + expected_tcgen05_mma_target_contracts(IntrinsicBackend::LlvmNvptx); + let expected_libnvvm_contracts = + expected_tcgen05_mma_target_contracts(IntrinsicBackend::LibNvvm); + let selected = [TargetSelectorBinding { + name: "kind".into(), + value: "f8f6f4".into(), + }]; + let expected_llvm = if alias.is_some() { + resolve_target_contract( + "tcgen05 MMA LLVM alias", + &selected, + &expected_llvm_contracts, + )? + } else { + resolve_target_contracts("tcgen05 MMA LLVM", &expected_llvm_contracts)? + }; + let expected_libnvvm = if alias.is_some() { + resolve_target_contract( + "tcgen05 MMA libNVVM alias", + &selected, + &expected_libnvvm_contracts, + )? + } else { + resolve_target_contracts("tcgen05 MMA libNVVM", &expected_libnvvm_contracts)? + }; + ensure!( + mma.llvm_target == expected_llvm + && mma.libnvvm_target == expected_libnvvm + && policy.minimum_ptx == "8.6" + && policy.minimum_sm.is_none() + && policy.targets == TCGEN05_LLVM_TARGETS + && policy.ptx_result == "()" + && policy.ptx_isa_version == "8.6" + && policy.expected_ptx == tcgen05_mma_expected_ptx(form, alias), + "{} tcgen05 MMA target or PTX contract changed", + policy.id + ); + let backend_pairs = policy + .backend_lowerings + .iter() + .map(|route| (route.backend, route.mechanism)) + .collect::>(); + ensure!( + policy.backend_lowerings.len() == 2 + && backend_pairs + == BTreeSet::from([ + ( + IntrinsicBackend::LlvmNvptx, + BackendLoweringMechanism::InlinePtx, + ), + ( + IntrinsicBackend::LibNvvm, + BackendLoweringMechanism::InlinePtx, + ), + ]) + && policy.backend_lowerings.iter().all(|route| { + !route.evidence_profile.trim().is_empty() + && route.targets.is_none() + && route.minimum_ptx.is_none() + && route.minimum_sm.is_none() + }), + "{} tcgen05 MMA backend routes changed", + policy.id + ); + + let expected_all = tcgen05_mma_all_selection_asms(form); + let expected_valid = if alias.is_some() { + BTreeSet::from([if tcgen05_mma_is_ws(form) { + tcgen05_mma_selection_asm( + form, + Tcgen05MmaKind::F8f6f4, + 1, + None, + Some(0), + Some(Tcgen05MmaBUsage::Discard), + ) + } else { + tcgen05_mma_selection_asm(form, Tcgen05MmaKind::F8f6f4, 1, Some("discard"), None, None) + }]) + } else { + tcgen05_mma_valid_selection_asms(form) + }; + let mut source_records = BTreeSet::new(); + let actual_all = declaration + .selections + .iter() + .map(|selection| { + ensure!( + !selection.source_record.is_empty() + && source_records.insert(selection.source_record.as_str()) + && selection.constraints.is_empty(), + "{} tcgen05 MMA selection provenance changed", + policy.id + ); + let predicate = if selection.asm.contains(".kind::i8.") { + "Subtarget->hasTcgen05MMAI8Kind()" + } else { + "Subtarget->hasTcgen05InstSupport()" + }; + ensure!( + selection.predicates == [predicate], + "{} tcgen05 MMA selection predicate changed", + policy.id + ); + Ok(selection.asm.clone()) + }) + .collect::>>()?; + ensure!( + actual_all == expected_all + && declaration.selections.len() == expected_all.len() + && expected_valid.is_subset(&expected_all) + && expected_valid.len() + == if alias.is_some() { + 1 + } else if tcgen05_mma_is_ws(form) { + 64 + } else if tcgen05_mma_is_ashift(form) { + 16 + } else { + 32 + }, + "{} tcgen05 MMA selection matrix changed", + policy.id + ); + ensure_no_other_family_contract(policy, "tcgen05 MMA")?; + Ok(()) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/tcgen05/mod.rs b/crates/cuda-intrinsics-gen/src/resolve/families/tcgen05/mod.rs new file mode 100644 index 0000000000..4792a0e0a0 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/tcgen05/mod.rs @@ -0,0 +1,14 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +pub(in crate::resolve) mod admission; +pub(in crate::resolve) mod ldst_cp; +pub(in crate::resolve) mod mma; +pub(in crate::resolve) mod recipes; + +pub(in crate::resolve) use admission::*; +pub(in crate::resolve) use ldst_cp::*; +pub(in crate::resolve) use mma::*; +pub(in crate::resolve) use recipes::*; diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/tcgen05/recipes.rs b/crates/cuda-intrinsics-gen/src/resolve/families/tcgen05/recipes.rs new file mode 100644 index 0000000000..8e4752757c --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/tcgen05/recipes.rs @@ -0,0 +1,914 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{Tcgen05Adapter, Tcgen05Operation, Tcgen05SourceContract}; +use crate::ptx::OperandPattern; + +pub(in crate::resolve) const TCGEN05_LLVM_TARGETS: &str = "sm_100a|sm_101a|sm_103a|sm_110a"; +pub(in crate::resolve) const TCGEN05_LIBNVVM_TARGETS: &str = "sm_100a|sm_103a|sm_110a"; +pub(in crate::resolve) struct Tcgen05Recipe { + pub(in crate::resolve) operation: Tcgen05Operation, + pub(in crate::resolve) abi_id: &'static str, + pub(in crate::resolve) id: &'static str, + pub(in crate::resolve) operation_key: &'static str, + pub(in crate::resolve) source_record: &'static str, + pub(in crate::resolve) llvm_symbol: &'static str, + pub(in crate::resolve) rust_arguments: &'static [&'static str], + pub(in crate::resolve) rust_result: &'static str, + pub(in crate::resolve) dialect_op_type: &'static str, + pub(in crate::resolve) dialect_op_name: &'static str, + pub(in crate::resolve) dialect_operands: &'static [&'static str], + pub(in crate::resolve) dialect_results: &'static [&'static str], + pub(in crate::resolve) llvm_arguments: &'static [&'static str], + pub(in crate::resolve) llvm_results: &'static [&'static str], + pub(in crate::resolve) imported_classes: &'static [&'static str], + pub(in crate::resolve) imported_properties: &'static [&'static str], + pub(in crate::resolve) adapter: Tcgen05Adapter, + pub(in crate::resolve) source_contract: Tcgen05SourceContract, + pub(in crate::resolve) safe: bool, + pub(in crate::resolve) safe_reason: Option<&'static str>, + pub(in crate::resolve) memory: &'static str, + pub(in crate::resolve) modifiers: Vec, + pub(in crate::resolve) operands: Vec, + pub(in crate::resolve) selection_record: Option<&'static str>, + pub(in crate::resolve) selection_asm: Option<&'static str>, + pub(in crate::resolve) summary: &'static str, +} + +pub(in crate::resolve) fn tcgen05_recipe(operation: Tcgen05Operation) -> Tcgen05Recipe { + const EMPTY: &[&str] = &[]; + const F32_X4: &[&str] = &["f32", "f32", "f32", "f32"]; + const F32_X32: &[&str] = &["f32"; 32]; + const BASE_CLASSES: &[&str] = &["SDPatternOperator", "Intrinsic"]; + const MMA_CLASSES: &[&str] = &[ + "SDPatternOperator", + "Intrinsic", + "DefaultAttrsIntrinsic", + "DefaultAttrsIntrinsicFlags", + ]; + const LOAD_CLASSES: &[&str] = &["SDPatternOperator", "Intrinsic", "NVVM_TCGEN05_LD"]; + const CONVERGENT_ARG_MEMORY: &[&str] = &["IntrArgMemOnly", "IntrConvergent", "NoCapture"]; + const CONVERGENT_INACCESSIBLE_ARG_MEMORY: &[&str] = &[ + "IntrConvergent", + "IntrInaccessibleMemOrArgMemOnly", + "NoCapture", + ]; + const CONVERGENT_INACCESSIBLE_MEMORY: &[&str] = &["IntrConvergent", "IntrInaccessibleMemOnly"]; + const ALLOC_PROPERTIES: &[&str] = &[ + "IntrConvergent", + "IntrInaccessibleMemOrArgMemOnly", + "NoCapture", + "WriteOnly", + ]; + const FENCE_PROPERTIES: &[&str] = &["IntrHasSideEffects", "IntrNoMem"]; + const MMA_WS_PROPERTIES: &[&str] = &[ + "ImmArg", + "ImmArg", + "ImmArg", + "IntrArgMemOnly", + "Range", + "Range", + "Range", + "ReadOnly", + "WriteOnly", + ]; + const MMA_CG1_PROPERTIES: &[&str] = &[ + "ImmArg", + "ImmArg", + "IntrArgMemOnly", + "Range", + "Range", + "WriteOnly", + ]; + const LOAD_PROPERTIES: &[&str] = &[ + "ImmArg", + "IntrArgMemOnly", + "IntrConvergent", + "NoCapture", + ]; + + let ( + abi_id, + id, + operation_key, + source_record, + llvm_symbol, + dialect_op_type, + dialect_op_name, + adapter, + source_contract, + ) = match operation { + Tcgen05Operation::Alloc => ( + "i0343", + "tcgen05_alloc", + "tcgen05.alloc.cg1", + "int_nvvm_tcgen05_alloc_shared_cg1", + "llvm.nvvm.tcgen05.alloc.shared.cg1", + "Tcgen05AllocOp", + "nvvm.tcgen05_alloc", + Tcgen05Adapter::SharedPointerColumnsToVoid, + Tcgen05SourceContract::ExactTablegenSelection, + ), + Tcgen05Operation::Dealloc => ( + "i0344", + "tcgen05_dealloc", + "tcgen05.dealloc.cg1", + "int_nvvm_tcgen05_dealloc_cg1", + "llvm.nvvm.tcgen05.dealloc.cg1", + "Tcgen05DeallocOp", + "nvvm.tcgen05_dealloc", + Tcgen05Adapter::TmemAddressColumnsToVoid, + Tcgen05SourceContract::ExactTablegenSelection, + ), + Tcgen05Operation::RelinquishAllocPermit => ( + "i0345", + "tcgen05_relinquish_alloc_permit", + "tcgen05.relinquish_alloc_permit.cg1", + "int_nvvm_tcgen05_relinq_alloc_permit_cg1", + "llvm.nvvm.tcgen05.relinq.alloc.permit.cg1", + "Tcgen05RelinquishAllocPermitOp", + "nvvm.tcgen05_relinquish_alloc_permit", + Tcgen05Adapter::NoOperands, + Tcgen05SourceContract::ExactTablegenSelection, + ), + Tcgen05Operation::FenceBeforeThreadSync => ( + "i0346", + "tcgen05_fence_before_thread_sync", + "tcgen05.fence.before_thread_sync", + "int_nvvm_tcgen05_fence_before_thread_sync", + "llvm.nvvm.tcgen05.fence.before.thread.sync", + "Tcgen05FenceBeforeThreadSyncOp", + "nvvm.tcgen05_fence_before_thread_sync", + Tcgen05Adapter::NoOperands, + Tcgen05SourceContract::ExactTablegenSelection, + ), + Tcgen05Operation::FenceAfterThreadSync => ( + "i0347", + "tcgen05_fence_after_thread_sync", + "tcgen05.fence.after.thread.sync", + "int_nvvm_tcgen05_fence_after_thread_sync", + "llvm.nvvm.tcgen05.fence.after.thread.sync", + "Tcgen05FenceAfterThreadSyncOp", + "nvvm.tcgen05_fence_after_thread_sync", + Tcgen05Adapter::NoOperands, + Tcgen05SourceContract::ExactTablegenSelection, + ), + Tcgen05Operation::Commit => ( + "i0348", + "tcgen05_commit", + "tcgen05.commit.cg1", + "int_nvvm_tcgen05_commit_cg1", + "llvm.nvvm.tcgen05.commit.cg1", + "Tcgen05CommitOp", + "nvvm.tcgen05_commit", + Tcgen05Adapter::BarrierPointerToVoid, + Tcgen05SourceContract::TablegenSelectionChangesPtx, + ), + Tcgen05Operation::CommitSharedCluster => ( + "i0349", + "tcgen05_commit_shared_cluster", + "tcgen05.commit.shared_cluster.cg1", + "int_nvvm_tcgen05_commit_shared_cg1", + "llvm.nvvm.tcgen05.commit.shared.cg1", + "Tcgen05CommitSharedClusterOp", + "nvvm.tcgen05_commit_shared_cluster", + Tcgen05Adapter::BarrierPointerToVoid, + Tcgen05SourceContract::ExactTablegenSelection, + ), + Tcgen05Operation::MmaWsF16 => ( + "i0350", + "tcgen05_mma_ws_f16", + "tcgen05.mma.ws.f16.cg1", + "int_nvvm_tcgen05_mma_ws_tensor", + "llvm.nvvm.tcgen05.mma.ws.tensor", + "Tcgen05MmaWsF16Op", + "nvvm.tcgen05_mma_ws_f16", + Tcgen05Adapter::MmaWsDropLegacyADescriptor, + Tcgen05SourceContract::TablegenSelectionChangesPtx, + ), + Tcgen05Operation::MmaF16 => ( + "i0351", + "tcgen05_mma_f16", + "tcgen05.mma.f16.cg1", + "int_nvvm_tcgen05_mma_shared_disable_output_lane_cg1", + "llvm.nvvm.tcgen05.mma.shared.disable_output_lane.cg1", + "Tcgen05MmaF16Op", + "nvvm.tcgen05_mma_f16", + Tcgen05Adapter::MmaInjectZeroDisableLanes, + Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection, + ), + Tcgen05Operation::MmaWsBf16 => ( + "i0352", + "tcgen05_mma_ws_bf16", + "tcgen05.mma.ws.bf16.cg1", + "int_nvvm_tcgen05_mma_ws_tensor", + "llvm.nvvm.tcgen05.mma.ws.tensor", + "Tcgen05MmaWsBf16Op", + "nvvm.tcgen05_mma_ws_bf16", + Tcgen05Adapter::MmaWsDropLegacyADescriptor, + Tcgen05SourceContract::TablegenSelectionChangesPtx, + ), + Tcgen05Operation::MmaWsTf32 => ( + "i0353", + "tcgen05_mma_ws_tf32", + "tcgen05.mma.ws.tf32.cg1", + "int_nvvm_tcgen05_mma_ws_tensor", + "llvm.nvvm.tcgen05.mma.ws.tensor", + "Tcgen05MmaWsTf32Op", + "nvvm.tcgen05_mma_ws_tf32", + Tcgen05Adapter::MmaWsDropLegacyADescriptor, + Tcgen05SourceContract::TablegenSelectionChangesPtx, + ), + Tcgen05Operation::CpSmemToTmem => ( + "i0354", + "tcgen05_cp_smem_to_tmem", + "tcgen05.cp.128x256b.cg1", + "int_nvvm_tcgen05_cp_128x256b_cg1", + "llvm.nvvm.tcgen05.cp.128x256b.cg1", + "Tcgen05CpSmemToTmemOp", + "nvvm.tcgen05_cp_smem_to_tmem", + Tcgen05Adapter::TmemDescriptorToVoid, + Tcgen05SourceContract::ExactTablegenSelection, + ), + Tcgen05Operation::Ld16x256bX8Pure => ( + "i0355", + "tcgen05_ld_16x256b_x8_pure", + "tcgen05.ld.16x256b.x8", + "int_nvvm_tcgen05_ld_16x256b_x8", + "llvm.nvvm.tcgen05.ld.16x256b.x8", + "Tcgen05Ld16x256bX8PureOp", + "nvvm.tcgen05_ld_16x256b_x8_pure", + Tcgen05Adapter::TmemToF32x32, + Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection, + ), + Tcgen05Operation::Ld16x256bPure => ( + "i0356", + "tcgen05_ld_16x256b_pure", + "tcgen05.ld.16x256b.x1", + "int_nvvm_tcgen05_ld_16x256b_x1", + "llvm.nvvm.tcgen05.ld.16x256b.x1", + "Tcgen05Ld16x256bPureOp", + "nvvm.tcgen05_ld_16x256b_pure", + Tcgen05Adapter::TmemToF32x4, + Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection, + ), + Tcgen05Operation::LoadWait => ( + "i0357", + "tcgen05_load_wait", + "tcgen05.wait.ld", + "int_nvvm_tcgen05_wait_ld", + "llvm.nvvm.tcgen05.wait.ld", + "Tcgen05LoadWaitOp", + "nvvm.tcgen05_load_wait", + Tcgen05Adapter::NoOperands, + Tcgen05SourceContract::ExactTablegenSelection, + ), + Tcgen05Operation::StoreWait => ( + "i0358", + "tcgen05_store_wait", + "tcgen05.wait.st", + "int_nvvm_tcgen05_wait_st", + "llvm.nvvm.tcgen05.wait.st", + "Tcgen05StoreWaitOp", + "nvvm.tcgen05_store_wait", + Tcgen05Adapter::NoOperands, + Tcgen05SourceContract::ExactTablegenSelection, + ), + Tcgen05Operation::AllocCg2 => ( + "i0359", + "tcgen05_alloc_cg2", + "tcgen05.alloc.cg2", + "int_nvvm_tcgen05_alloc_shared_cg2", + "llvm.nvvm.tcgen05.alloc.shared.cg2", + "Tcgen05AllocCg2Op", + "nvvm.tcgen05_alloc_cg2", + Tcgen05Adapter::SharedPointerColumnsToVoid, + Tcgen05SourceContract::ExactTablegenSelection, + ), + Tcgen05Operation::DeallocCg2 => ( + "i0360", + "tcgen05_dealloc_cg2", + "tcgen05.dealloc.cg2", + "int_nvvm_tcgen05_dealloc_cg2", + "llvm.nvvm.tcgen05.dealloc.cg2", + "Tcgen05DeallocCg2Op", + "nvvm.tcgen05_dealloc_cg2", + Tcgen05Adapter::TmemAddressColumnsToVoid, + Tcgen05SourceContract::ExactTablegenSelection, + ), + Tcgen05Operation::RelinquishAllocPermitCg2 => ( + "i0361", + "tcgen05_relinquish_alloc_permit_cg2", + "tcgen05.relinquish_alloc_permit.cg2", + "int_nvvm_tcgen05_relinq_alloc_permit_cg2", + "llvm.nvvm.tcgen05.relinq.alloc.permit.cg2", + "Tcgen05RelinquishAllocPermitCg2Op", + "nvvm.tcgen05_relinquish_alloc_permit_cg2", + Tcgen05Adapter::NoOperands, + Tcgen05SourceContract::ExactTablegenSelection, + ), + Tcgen05Operation::MmaF16Cg2 => ( + "i0362", + "tcgen05_mma_f16_cg2", + "tcgen05.mma.f16.cg2", + "int_nvvm_tcgen05_mma_shared_disable_output_lane_cg2", + "llvm.nvvm.tcgen05.mma.shared.disable_output_lane.cg2", + "Tcgen05MmaF16Cg2Op", + "nvvm.tcgen05_mma_f16_cg2", + Tcgen05Adapter::MmaInjectZeroDisableLanes, + Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection, + ), + Tcgen05Operation::CommitCg2 => ( + "i0363", + "tcgen05_commit_cg2", + "tcgen05.commit.cg2", + "int_nvvm_tcgen05_commit_cg2", + "llvm.nvvm.tcgen05.commit.cg2", + "Tcgen05CommitCg2Op", + "nvvm.tcgen05_commit_cg2", + Tcgen05Adapter::BarrierPointerToVoid, + Tcgen05SourceContract::TablegenSelectionChangesPtx, + ), + Tcgen05Operation::CommitSharedClusterCg2 => ( + "i0364", + "tcgen05_commit_shared_cluster_cg2", + "tcgen05.commit.shared_cluster.cg2", + "int_nvvm_tcgen05_commit_shared_cg2", + "llvm.nvvm.tcgen05.commit.shared.cg2", + "Tcgen05CommitSharedClusterCg2Op", + "nvvm.tcgen05_commit_shared_cluster_cg2", + Tcgen05Adapter::BarrierPointerToVoid, + Tcgen05SourceContract::ExactTablegenSelection, + ), + Tcgen05Operation::CommitMulticastCg2 => ( + "i0365", + "tcgen05_commit_multicast_cg2", + "tcgen05.commit.multicast.cg2", + "int_nvvm_tcgen05_commit_mc_shared_cg2", + "llvm.nvvm.tcgen05.commit.mc.shared.cg2", + "Tcgen05CommitMulticastCg2Op", + "nvvm.tcgen05_commit_multicast_cg2", + Tcgen05Adapter::BarrierPointerMaskToVoid, + Tcgen05SourceContract::ExactTablegenSelection, + ), + Tcgen05Operation::CpSmemToTmemCg2 => ( + "i0366", + "tcgen05_cp_smem_to_tmem_cg2", + "tcgen05.cp.128x256b.cg2", + "int_nvvm_tcgen05_cp_128x256b_cg2", + "llvm.nvvm.tcgen05.cp.128x256b.cg2", + "Tcgen05CpSmemToTmemCg2Op", + "nvvm.tcgen05_cp_smem_to_tmem_cg2", + Tcgen05Adapter::TmemDescriptorToVoid, + Tcgen05SourceContract::ExactTablegenSelection, + ), + Tcgen05Operation::CommitMulticast => ( + "i0760", + "tcgen05_commit_multicast", + "tcgen05.commit.multicast.cg1", + "int_nvvm_tcgen05_commit_mc_shared_cg1", + "llvm.nvvm.tcgen05.commit.mc.shared.cg1", + "Tcgen05CommitMulticastOp", + "nvvm.tcgen05_commit_multicast", + Tcgen05Adapter::BarrierPointerMaskToVoid, + Tcgen05SourceContract::ExactTablegenSelection, + ), + Tcgen05Operation::ShiftDown => ( + "i0761", + "tcgen05_shift_down", + "tcgen05.shift.down.cg1", + "int_nvvm_tcgen05_shift_down_cg1", + "llvm.nvvm.tcgen05.shift.down.cg1", + "Tcgen05ShiftDownOp", + "nvvm.tcgen05_shift_down", + Tcgen05Adapter::TmemAddressToVoid, + Tcgen05SourceContract::ExactTablegenSelection, + ), + Tcgen05Operation::ShiftDownCg2 => ( + "i0762", + "tcgen05_shift_down_cg2", + "tcgen05.shift.down.cg2", + "int_nvvm_tcgen05_shift_down_cg2", + "llvm.nvvm.tcgen05.shift.down.cg2", + "Tcgen05ShiftDownCg2Op", + "nvvm.tcgen05_shift_down_cg2", + Tcgen05Adapter::TmemAddressToVoid, + Tcgen05SourceContract::ExactTablegenSelection, + ), + Tcgen05Operation::Ld | Tcgen05Operation::St | Tcgen05Operation::Mma => { + unreachable!("tcgen05 load/store variants use their compact recipes") + } + }; + + let cg2 = matches!( + operation, + Tcgen05Operation::AllocCg2 + | Tcgen05Operation::DeallocCg2 + | Tcgen05Operation::RelinquishAllocPermitCg2 + | Tcgen05Operation::MmaF16Cg2 + | Tcgen05Operation::CommitCg2 + | Tcgen05Operation::CommitSharedClusterCg2 + | Tcgen05Operation::CommitMulticastCg2 + | Tcgen05Operation::CpSmemToTmemCg2 + | Tcgen05Operation::ShiftDownCg2 + ); + let group = if cg2 { "cta_group::2" } else { "cta_group::1" }; + let ( + rust_arguments, + rust_result, + dialect_operands, + dialect_results, + llvm_arguments, + llvm_results, + imported_classes, + imported_properties, + safe, + safe_reason, + memory, + modifiers, + operands, + selection_record, + selection_asm, + summary, + ) = match operation { + Tcgen05Operation::Alloc | Tcgen05Operation::AllocCg2 => ( + &["*mut u32", "u32"] as &[_], + "()", + &["ptr", "i32"] as &[_], + EMPTY, + &["shared_ptr", "i32"] as &[_], + EMPTY, + BASE_CLASSES, + ALLOC_PROPERTIES, + false, + None, + "write", + vec![ + "alloc".into(), + group.into(), + "sync".into(), + "aligned".into(), + "shared::cta".into(), + "b32".into(), + ], + vec![OperandPattern::Address, OperandPattern::Register], + Some(if cg2 { + "TCGEN05_ALLOC_S64_CG2" + } else { + "TCGEN05_ALLOC_S64_CG1" + }), + Some(if cg2 { + "tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 \t[$dst], $ncols;" + } else { + "tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 \t[$dst], $ncols;" + }), + "Allocates tensor-memory columns and writes their address to shared memory.", + ), + Tcgen05Operation::Dealloc | Tcgen05Operation::DeallocCg2 => ( + &["u32", "u32"] as &[_], + "()", + &["i32", "i32"] as &[_], + EMPTY, + &["tmem_ptr", "i32"] as &[_], + EMPTY, + BASE_CLASSES, + CONVERGENT_ARG_MEMORY, + false, + None, + "read_write", + vec![ + "dealloc".into(), + group.into(), + "sync".into(), + "aligned".into(), + "b32".into(), + ], + vec![OperandPattern::Register, OperandPattern::Register], + Some(if cg2 { + "TCGEN05_DEALLOC_CG2" + } else { + "TCGEN05_DEALLOC_CG1" + }), + Some(if cg2 { + "tcgen05.dealloc.cta_group::2.sync.aligned.b32 \t$tmem_addr, $ncols;" + } else { + "tcgen05.dealloc.cta_group::1.sync.aligned.b32 \t$tmem_addr, $ncols;" + }), + "Releases tensor-memory columns allocated to the CTA group.", + ), + Tcgen05Operation::RelinquishAllocPermit | Tcgen05Operation::RelinquishAllocPermitCg2 => ( + EMPTY, + "()", + EMPTY, + EMPTY, + EMPTY, + EMPTY, + BASE_CLASSES, + CONVERGENT_INACCESSIBLE_MEMORY, + true, + Some( + "The operation has no Rust memory operand and only releases this CTA group's allocation permit.", + ), + "read_write", + vec![ + "relinquish_alloc_permit".into(), + group.into(), + "sync".into(), + "aligned".into(), + ], + vec![], + Some(if cg2 { + "TCGEN05_RELINQ_CG2" + } else { + "TCGEN05_RELINQ_CG1" + }), + Some(if cg2 { + "tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;" + } else { + "tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;" + }), + "Releases the CTA group's tensor-memory allocation permit.", + ), + Tcgen05Operation::FenceBeforeThreadSync | Tcgen05Operation::FenceAfterThreadSync => { + let before = operation == Tcgen05Operation::FenceBeforeThreadSync; + ( + EMPTY, + "()", + EMPTY, + EMPTY, + EMPTY, + EMPTY, + BASE_CLASSES, + FENCE_PROPERTIES, + true, + Some("The operation only orders the calling thread's tcgen05 accesses."), + "none", + vec![if before { + "fence::before_thread_sync".into() + } else { + "fence::after_thread_sync".into() + }], + vec![], + Some(if before { + "tcgen05_fence_before_thread_sync" + } else { + "tcgen05_fence_after_thread_sync" + }), + Some(if before { + "tcgen05.fence::before_thread_sync;" + } else { + "tcgen05.fence::after_thread_sync;" + }), + if before { + "Orders prior tcgen05 accesses before thread synchronization." + } else { + "Orders later tcgen05 accesses after thread synchronization." + }, + ) + } + Tcgen05Operation::Commit + | Tcgen05Operation::CommitCg2 + | Tcgen05Operation::CommitSharedCluster + | Tcgen05Operation::CommitSharedClusterCg2 => { + let shared = matches!( + operation, + Tcgen05Operation::CommitSharedCluster | Tcgen05Operation::CommitSharedClusterCg2 + ); + let llvm_args = if shared { + &["shared_ptr"] as &[_] + } else { + &["ptr"] as &[_] + }; + let record = match (cg2, shared) { + (false, false) => "TCGEN05_COMMIT_CG1", + (true, false) => "TCGEN05_COMMIT_CG2", + (false, true) => "TCGEN05_COMMIT_S64_CG1", + (true, true) => "TCGEN05_COMMIT_S64_CG2", + }; + let asm = match cg2 { + false => { + "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 \t[$mbar];" + } + true => { + "tcgen05.commit.cta_group::2.mbarrier::arrive::one.shared::cluster.b64 \t[$mbar];" + } + }; + let mut modifiers = vec![ + "commit".into(), + group.into(), + "mbarrier::arrive::one".into(), + ]; + if shared { + modifiers.push("shared::cluster".into()); + } + modifiers.push("b64".into()); + ( + &["*mut u64"] as &[_], + "()", + &["ptr"] as &[_], + EMPTY, + llvm_args, + EMPTY, + BASE_CLASSES, + CONVERGENT_INACCESSIBLE_ARG_MEMORY, + false, + None, + "read_write", + modifiers, + vec![OperandPattern::Address], + Some(record), + Some(asm), + if shared { + "Commits tcgen05 completion to a cluster-shared mbarrier." + } else { + "Commits tcgen05 completion to an mbarrier." + }, + ) + } + Tcgen05Operation::MmaWsF16 | Tcgen05Operation::MmaWsBf16 | Tcgen05Operation::MmaWsTf32 => { + let kind = if operation == Tcgen05Operation::MmaWsTf32 { + "kind::tf32" + } else { + "kind::f16" + }; + ( + &["u32", "u32", "u64", "u64", "u32", "bool"] as &[_], + "()", + &["i32", "i32", "i64", "i64", "i32", "i1"] as &[_], + EMPTY, + &[ + "tmem_ptr", "tmem_ptr", "i64", "i32", "i1", "i32", "i32", "i32", + ] as &[_], + EMPTY, + MMA_CLASSES, + MMA_WS_PROPERTIES, + false, + None, + "read_write", + vec![ + "mma".into(), + "ws".into(), + "cta_group::1".into(), + kind.into(), + ], + vec![ + OperandPattern::Address, + OperandPattern::Address, + OperandPattern::Register, + OperandPattern::Register, + OperandPattern::Exact { + value: "%enable_pred".into(), + }, + ], + None, + None, + match operation { + Tcgen05Operation::MmaWsF16 => "Issues warp-specialized f16 tensor-memory MMA.", + Tcgen05Operation::MmaWsBf16 => { + "Issues warp-specialized bf16 tensor-memory MMA using the f16 instruction class." + } + _ => "Issues warp-specialized tf32 tensor-memory MMA.", + }, + ) + } + Tcgen05Operation::MmaF16 | Tcgen05Operation::MmaF16Cg2 => ( + &["u32", "u64", "u64", "u32", "bool"] as &[_], + "()", + &["i32", "i64", "i64", "i32", "i1"] as &[_], + EMPTY, + if cg2 { + &["tmem_ptr", "i64", "i64", "i32", "i1", "v8i32", "i32", "i32"] as &[_] + } else { + &["tmem_ptr", "i64", "i64", "i32", "i1", "v4i32", "i32", "i32"] as &[_] + }, + EMPTY, + MMA_CLASSES, + MMA_CG1_PROPERTIES, + false, + None, + "write", + vec!["mma".into(), group.into(), "kind::f16".into()], + vec![ + OperandPattern::Address, + OperandPattern::Register, + OperandPattern::Register, + OperandPattern::Register, + OperandPattern::Exact { + value: if cg2 { + "{%z, %z, %z, %z, %z, %z, %z, %z}" + } else { + "{%z, %z, %z, %z}" + } + .into(), + }, + OperandPattern::Exact { + value: "%enable_pred".into(), + }, + ], + None, + None, + "Issues f16 tensor-memory MMA with zeroed disable-output-lane controls.", + ), + Tcgen05Operation::CpSmemToTmem | Tcgen05Operation::CpSmemToTmemCg2 => ( + &["u32", "u64"] as &[_], + "()", + &["i32", "i64"] as &[_], + EMPTY, + &["tmem_ptr", "i64"] as &[_], + EMPTY, + BASE_CLASSES, + CONVERGENT_INACCESSIBLE_ARG_MEMORY, + false, + None, + "read_write", + vec!["cp".into(), group.into(), "128x256b".into()], + vec![OperandPattern::Address, OperandPattern::Register], + Some(if cg2 { + "TCGEN05_CP_128x256b_cg2" + } else { + "TCGEN05_CP_128x256b_cg1" + }), + Some(if cg2 { + "tcgen05.cp.cta_group::2.128x256b \t[$tmem_addr], $sdesc;" + } else { + "tcgen05.cp.cta_group::1.128x256b \t[$tmem_addr], $sdesc;" + }), + "Copies one 128x256-bit tile from shared memory to tensor memory.", + ), + Tcgen05Operation::Ld16x256bX8Pure | Tcgen05Operation::Ld16x256bPure => { + let x8 = operation == Tcgen05Operation::Ld16x256bX8Pure; + ( + &["u32"] as &[_], + if x8 { "[f32; 32]" } else { "[f32; 4]" }, + &["i32"] as &[_], + if x8 { F32_X32 } else { F32_X4 }, + &["tmem_ptr", "i1"] as &[_], + if x8 { + &["v32i32"] as &[_] + } else { + &["v4i32"] as &[_] + }, + LOAD_CLASSES, + LOAD_PROPERTIES, + false, + None, + "read", + vec![ + "ld".into(), + "sync".into(), + "aligned".into(), + "16x256b".into(), + if x8 { "x8".into() } else { "x1".into() }, + "b32".into(), + ], + vec![ + OperandPattern::RegisterList { + length: if x8 { 32 } else { 4 }, + }, + OperandPattern::Address, + ], + None, + None, + if x8 { + "Loads 32 f32 register values from tensor memory." + } else { + "Loads four f32 register values from tensor memory." + }, + ) + } + Tcgen05Operation::LoadWait | Tcgen05Operation::StoreWait => { + let load = operation == Tcgen05Operation::LoadWait; + ( + EMPTY, + "()", + EMPTY, + EMPTY, + EMPTY, + EMPTY, + BASE_CLASSES, + CONVERGENT_INACCESSIBLE_MEMORY, + true, + Some("The operation only waits for the calling thread's prior tcgen05 access."), + "read_write", + vec![ + if load { + "wait::ld".into() + } else { + "wait::st".into() + }, + "sync".into(), + "aligned".into(), + ], + vec![], + Some(if load { + "tcgen05_wait_ld" + } else { + "tcgen05_wait_st" + }), + Some(if load { + "tcgen05.wait::ld.sync.aligned;" + } else { + "tcgen05.wait::st.sync.aligned;" + }), + if load { + "Waits for prior tensor-memory loads to complete." + } else { + "Waits for prior tensor-memory stores to complete." + }, + ) + } + Tcgen05Operation::CommitMulticast | Tcgen05Operation::CommitMulticastCg2 => ( + &["*mut u64", "u16"] as &[_], + "()", + &["ptr", "i16"] as &[_], + EMPTY, + &["shared_ptr", "i16"] as &[_], + EMPTY, + BASE_CLASSES, + CONVERGENT_INACCESSIBLE_ARG_MEMORY, + false, + None, + "read_write", + vec![ + "commit".into(), + group.into(), + "mbarrier::arrive::one".into(), + "shared::cluster".into(), + "multicast::cluster".into(), + "b64".into(), + ], + vec![OperandPattern::Address, OperandPattern::Register], + Some(if cg2 { + "TCGEN05_COMMIT_S64_CG2_MC" + } else { + "TCGEN05_COMMIT_S64_CG1_MC" + }), + Some(if cg2 { + "tcgen05.commit.cta_group::2.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 \t[$mbar], $mc;" + } else { + "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 \t[$mbar], $mc;" + }), + "Commits tcgen05 completion to the selected cluster mbarriers.", + ), + Tcgen05Operation::ShiftDown | Tcgen05Operation::ShiftDownCg2 => ( + &["u32"] as &[_], + "()", + &["i32"] as &[_], + EMPTY, + &["tmem_ptr"] as &[_], + EMPTY, + BASE_CLASSES, + CONVERGENT_ARG_MEMORY, + false, + None, + "read_write", + vec!["shift".into(), group.into(), "down".into()], + vec![OperandPattern::Address], + Some(if cg2 { + "TCGEN05_SHIFT_CG2" + } else { + "TCGEN05_SHIFT_CG1" + }), + Some(if cg2 { + "tcgen05.shift.cta_group::2.down \t[$tmem_addr];" + } else { + "tcgen05.shift.cta_group::1.down \t[$tmem_addr];" + }), + "Shifts tensor-memory rows down by one row.", + ), + Tcgen05Operation::Ld | Tcgen05Operation::St | Tcgen05Operation::Mma => { + unreachable!("tcgen05 load/store variants use their compact recipes") + } + }; + + Tcgen05Recipe { + operation, + abi_id, + id, + operation_key, + source_record, + llvm_symbol, + rust_arguments, + rust_result, + dialect_op_type, + dialect_op_name, + dialect_operands, + dialect_results, + llvm_arguments, + llvm_results, + imported_classes, + imported_properties, + adapter, + source_contract, + safe, + safe_reason, + memory, + modifiers, + operands, + selection_record, + selection_asm, + summary, + } +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/threadfence.rs b/crates/cuda-intrinsics-gen/src/resolve/families/threadfence.rs new file mode 100644 index 0000000000..d83e1741ac --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/threadfence.rs @@ -0,0 +1,363 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, ImportedIntrinsic, IntrinsicBackend, OverlayBackendLowering, + OverlayIntrinsic, RuntimeValidation, ThreadfenceAdmission, ThreadfenceScope, +}; +use crate::ptx::InstructionPattern; +use anyhow::{Result, ensure}; + +#[derive(Clone, Copy)] +pub(in crate::resolve) struct ThreadfenceRecipe { + pub(in crate::resolve) scope: ThreadfenceScope, + pub(in crate::resolve) abi_id: &'static str, + pub(in crate::resolve) id: &'static str, + pub(in crate::resolve) operation_key: &'static str, + pub(in crate::resolve) source_record: &'static str, + pub(in crate::resolve) selection_record: &'static str, + pub(in crate::resolve) llvm_symbol: &'static str, + pub(in crate::resolve) dialect_op_type: &'static str, + pub(in crate::resolve) dialect_op_name: &'static str, + pub(in crate::resolve) ptx_level: &'static str, + pub(in crate::resolve) execution_scope: &'static str, + pub(in crate::resolve) minimum_ptx: &'static str, + pub(in crate::resolve) minimum_sm: Option<&'static str>, + pub(in crate::resolve) summary: &'static str, +} + +pub(in crate::resolve) fn threadfence_recipe(scope: ThreadfenceScope) -> ThreadfenceRecipe { + match scope { + ThreadfenceScope::Cta => ThreadfenceRecipe { + scope, + abi_id: "i0298", + id: "threadfence_block", + operation_key: "memory.fence.cta.sc", + source_record: "int_nvvm_membar_cta", + selection_record: "INT_MEMBAR_CTA", + llvm_symbol: "llvm.nvvm.membar.cta", + dialect_op_type: "ThreadfenceBlockOp", + dialect_op_name: "nvvm.threadfence_block", + ptx_level: "cta", + execution_scope: "cta", + minimum_ptx: "1.4", + minimum_sm: None, + summary: "Orders this thread's memory operations for observers in its CTA.", + }, + ThreadfenceScope::Device => ThreadfenceRecipe { + scope, + abi_id: "i0299", + id: "threadfence", + operation_key: "memory.fence.device.sc", + source_record: "int_nvvm_membar_gl", + selection_record: "INT_MEMBAR_GL", + llvm_symbol: "llvm.nvvm.membar.gl", + dialect_op_type: "ThreadfenceOp", + dialect_op_name: "nvvm.threadfence", + ptx_level: "gl", + execution_scope: "device", + minimum_ptx: "1.4", + minimum_sm: None, + summary: "Orders this thread's memory operations for observers on its GPU.", + }, + ThreadfenceScope::System => ThreadfenceRecipe { + scope, + abi_id: "i0300", + id: "threadfence_system", + operation_key: "memory.fence.system.sc", + source_record: "int_nvvm_membar_sys", + selection_record: "INT_MEMBAR_SYS", + llvm_symbol: "llvm.nvvm.membar.sys", + dialect_op_type: "ThreadfenceSystemOp", + dialect_op_name: "nvvm.threadfence_system", + ptx_level: "sys", + execution_scope: "system", + minimum_ptx: "2.0", + minimum_sm: Some("sm_20"), + summary: "Orders this thread's memory operations for system-wide observers.", + }, + } +} + +pub(in crate::resolve) fn threadfence_scope_for_id(id: &str) -> Option { + match id { + "threadfence_block" => Some(ThreadfenceScope::Cta), + "threadfence" => Some(ThreadfenceScope::Device), + "threadfence_system" => Some(ThreadfenceScope::System), + _ => None, + } +} + +pub(in crate::resolve) fn expand_threadfence_admission( + admission: &ThreadfenceAdmission, +) -> Result> { + ensure!( + admission.runtime_validation == RuntimeValidation::Unexecuted, + "threadfence runtime validation may be marked executed only with GPU evidence" + ); + ensure!( + !admission.llvm_evidence_profile.trim().is_empty() + && !admission.libnvvm_evidence_profile.trim().is_empty(), + "compact threadfence admission requires both backend evidence profiles" + ); + let expected_scopes = [ + ThreadfenceScope::Cta, + ThreadfenceScope::Device, + ThreadfenceScope::System, + ]; + let actual_scopes = admission + .variants + .iter() + .map(|variant| variant.scope) + .collect::>(); + ensure!( + actual_scopes == expected_scopes, + "compact threadfence admission must contain each reviewed scope exactly once in canonical order" + ); + + admission + .variants + .iter() + .map(|variant| { + let recipe = threadfence_recipe(variant.scope); + ensure!( + variant.abi_id == recipe.abi_id, + "{} must keep reserved ABI ID {}", + recipe.id, + recipe.abi_id + ); + Ok(OverlayIntrinsic { + id: recipe.id.into(), + abi_id: variant.abi_id.clone(), + operation_key: recipe.operation_key.into(), + family: "sync".into(), + source: None, + source_record: Some(recipe.source_record.into()), + rust_module: "fence".into(), + rust_name: recipe.id.into(), + rust_arguments: vec![], + rust_result: "()".into(), + safe: true, + must_use: false, + safe_allowlist_reason: Some( + "a fence only orders the calling thread's memory operations and has no caller preconditions" + .into(), + ), + public_rust_path: format!("cuda_intrinsics::fence::{}", recipe.id), + compatibility_rust_paths: vec![ + format!("cuda_device::fence::{}", recipe.id), + format!("cuda_device::{}", recipe.id), + ], + dialect_op_type: recipe.dialect_op_type.into(), + dialect_op_name: recipe.dialect_op_name.into(), + dialect_operands: vec![], + dialect_results: vec![], + llvm_symbol: Some(recipe.llvm_symbol.into()), + resolved_llvm_symbol: None, + llvm_arguments: vec![], + llvm_results: vec![], + pure: false, + memory: "read_write".into(), + convergent: false, + execution_scope: recipe.execution_scope.into(), + minimum_ptx: recipe.minimum_ptx.into(), + minimum_sm: recipe.minimum_sm.map(Into::into), + ptx_result: "()".into(), + targets: "all".into(), + ptx_isa_version: "9.3".into(), + ptx_isa_section: + "9.7.14.4 Parallel Synchronization and Communication Instructions: membar / fence" + .into(), + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-membar-fence".into(), + lowering: "direct_nvvm".into(), + backend_lowerings: vec![ + OverlayBackendLowering { + backend: IntrinsicBackend::LlvmNvptx, + mechanism: BackendLoweringMechanism::TypedNvvm, + evidence_profile: admission.llvm_evidence_profile.clone(), + targets: None, + minimum_ptx: Some("3.2".into()), + minimum_sm: Some("sm_20".into()), + }, + OverlayBackendLowering { + backend: IntrinsicBackend::LibNvvm, + mechanism: BackendLoweringMechanism::TypedNvvm, + evidence_profile: admission.libnvvm_evidence_profile.clone(), + targets: None, + minimum_ptx: Some("7.0".into()), + minimum_sm: Some("sm_80".into()), + }, + ], + packed_atomic: None, + redux: None, + vote: None, + active_mask: None, + warp_match: None, + warp_barrier: None, + warp_shuffle: None, + dot_product: None, + packed_alu: None, + integer_minmax: None, + packed_conversion: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cp_async_copy: None, + cp_async_control: None, + cp_async_mbarrier: None, + mbarrier_basic: None, + movmatrix: None, + mbarrier_extended: None, + register_mma: None, + sparse_mma: None, + prmt: None, + cluster_barrier: None, + wgmma_control: None, + special_register: None, + debug_control: None, + cluster_memory: None, + clc: None, + tma: None, + tcgen05: None, + ldmatrix_variant: None, + ldmatrix_safety: None, + ldmatrix_adapter: None, + selected_address_space: None, + expected_ptx: InstructionPattern { + mnemonic: "membar".into(), + modifiers: vec![recipe.ptx_level.into()], + operands: vec![], + }, + summary: recipe.summary.into(), + }) + }) + .collect() +} + +pub(in crate::resolve) fn validate_threadfence_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, + scope: ThreadfenceScope, +) -> Result<()> { + let recipe = threadfence_recipe(scope); + ensure!( + recipe.scope == scope + && policy.id == recipe.id + && policy.abi_id == recipe.abi_id + && policy.operation_key == recipe.operation_key + && policy.source.is_none() + && policy.source_record.as_deref() == Some(recipe.source_record) + && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) + && policy.resolved_llvm_symbol.is_none(), + "{} threadfence identity does not match its closed recipe", + policy.id + ); + ensure!( + policy.rust_module == "fence" + && policy.rust_name == recipe.id + && policy.rust_arguments.is_empty() + && policy.rust_result == "()" + && policy.safe + && !policy.must_use + && policy.public_rust_path == format!("cuda_intrinsics::fence::{}", recipe.id) + && policy.compatibility_rust_paths + == [ + format!("cuda_device::fence::{}", recipe.id), + format!("cuda_device::{}", recipe.id), + ], + "{} threadfence Rust API does not match its closed recipe", + policy.id + ); + ensure!( + policy.dialect_op_type == recipe.dialect_op_type + && policy.dialect_op_name == recipe.dialect_op_name + && policy.dialect_operands.is_empty() + && policy.dialect_results.is_empty() + && policy.llvm_arguments.is_empty() + && policy.llvm_results.is_empty() + && policy.lowering == "direct_nvvm", + "{} threadfence carrier or lowering does not match its closed recipe", + policy.id + ); + ensure!( + declaration.properties == ["IntrNoCallback"] + && !policy.pure + && policy.memory == "read_write" + && !policy.convergent + && policy.execution_scope == recipe.execution_scope, + "{} threadfence effects disagree with the imported declaration", + policy.id + ); + ensure!( + policy.minimum_ptx == recipe.minimum_ptx + && policy.minimum_sm.as_deref() == recipe.minimum_sm + && policy.ptx_result == "()" + && policy.targets == "all" + && policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section + == "9.7.14.4 Parallel Synchronization and Communication Instructions: membar / fence" + && policy.ptx_isa_url + == "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-membar-fence", + "{} threadfence target floor or PTX provenance changed", + policy.id + ); + ensure!( + policy.expected_ptx + == (InstructionPattern { + mnemonic: "membar".into(), + modifiers: vec![recipe.ptx_level.into()], + operands: vec![], + }), + "{} expected PTX does not match its closed threadfence scope", + policy.id + ); + ensure!( + policy.backend_lowerings.len() == 2 + && policy.backend_lowerings.iter().any(|route| { + route.backend == IntrinsicBackend::LlvmNvptx + && route.mechanism == BackendLoweringMechanism::TypedNvvm + && route.minimum_ptx.as_deref() == Some("3.2") + && route.minimum_sm.as_deref() == Some("sm_20") + }) + && policy.backend_lowerings.iter().any(|route| { + route.backend == IntrinsicBackend::LibNvvm + && route.mechanism == BackendLoweringMechanism::TypedNvvm + && route.minimum_ptx.as_deref() == Some("7.0") + && route.minimum_sm.as_deref() == Some("sm_80") + }), + "{} must keep both reviewed typed threadfence routes", + policy.id + ); + ensure!( + policy.ldmatrix_variant.is_none() + && policy.ldmatrix_safety.is_none() + && policy.ldmatrix_adapter.is_none() + && policy.packed_atomic.is_none() + && policy.redux.is_none() + && policy.vote.is_none() + && policy.active_mask.is_none() + && policy.warp_match.is_none() + && policy.warp_barrier.is_none() + && policy.warp_shuffle.is_none() + && policy.dot_product.is_none() + && policy.packed_alu.is_none() + && policy.packed_conversion.is_none() + && policy.cp_async_copy.is_none() + && policy.cp_async_control.is_none() + && policy.cp_async_mbarrier.is_none() + && policy.mbarrier_basic.is_none() + && policy.register_mma.is_none() + && policy.sparse_mma.is_none() + && policy.prmt.is_none() + && policy.cluster_barrier.is_none() + && policy.special_register.is_none() + && policy.debug_control.is_none() + && policy.selected_address_space.is_none(), + "{} mixes another generated-family contract with threadfence", + policy.id + ); + Ok(()) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/tma.rs b/crates/cuda-intrinsics-gen/src/resolve/families/tma.rs new file mode 100644 index 0000000000..0f7d8a1d09 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/tma.rs @@ -0,0 +1,1930 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, ImportedAddressSpace, ImportedIntrinsic, IntrinsicBackend, + OverlayBackendLowering, OverlayIntrinsic, RuntimeValidation, Tma, TmaAdapter, TmaAdmission, + TmaOperation, TmaReduction, TmaReductionAdmissionVariant, TmaReductionLoadMode, + TmaReductionOperation, +}; +use crate::ptx::{InstructionPattern, OperandPattern}; +use anyhow::{Context, Result, ensure}; + +use super::*; +use crate::resolve::abi_ledger::*; +use crate::resolve::guards::*; + +pub(in crate::resolve) const TMA_BLACKWELL_TARGETS: &str = "sm_100a|sm_101a|sm_103a|sm_110a"; +pub(in crate::resolve) const TENSOR_MAP_REPLACE_TARGETS: &str = + "sm_100a|sm_100f|sm_103a|sm_103f|sm_110a|sm_110f|sm_120a|sm_120f|sm_121a|sm_121f|sm_90a"; +pub(in crate::resolve) const TMA_OPERATIONS: [TmaOperation; 47] = [ + TmaOperation::G2sTile1d, + TmaOperation::G2sTile2d, + TmaOperation::G2sTile2dMulticast, + TmaOperation::G2sTile2dMulticastCg2, + TmaOperation::G2sTile3d, + TmaOperation::G2sTile4d, + TmaOperation::G2sTile5d, + TmaOperation::S2gTile1d, + TmaOperation::S2gTile2d, + TmaOperation::S2gTile3d, + TmaOperation::S2gTile4d, + TmaOperation::S2gTile5d, + TmaOperation::CommitGroup, + TmaOperation::WaitGroup, + TmaOperation::WaitGroupRead, + TmaOperation::PrefetchTensorMap, + TmaOperation::PrefetchTile1d, + TmaOperation::PrefetchTile2d, + TmaOperation::PrefetchTile3d, + TmaOperation::PrefetchTile4d, + TmaOperation::PrefetchTile5d, + TmaOperation::PrefetchTileGather4TwoDimensional, + TmaOperation::ReplaceBoxDim, + TmaOperation::ReplaceElementStride, + TmaOperation::ReplaceElementType, + TmaOperation::ReplaceFillMode, + TmaOperation::ReplaceGlobalAddress, + TmaOperation::ReplaceGlobalDim, + TmaOperation::ReplaceGlobalStride, + TmaOperation::ReplaceInterleaveLayout, + TmaOperation::ReplaceRank, + TmaOperation::ReplaceSwizzleAtomicity, + TmaOperation::ReplaceSwizzleMode, + TmaOperation::FenceProxyTensorMapAcquireCluster, + TmaOperation::FenceProxyTensorMapAcquireCta, + TmaOperation::FenceProxyTensorMapAcquireGpu, + TmaOperation::FenceProxyTensorMapAcquireSystem, + TmaOperation::FenceProxyTensorMapReleaseCluster, + TmaOperation::FenceProxyTensorMapReleaseCta, + TmaOperation::FenceProxyTensorMapReleaseGpu, + TmaOperation::FenceProxyTensorMapReleaseSystem, + TmaOperation::PrefetchTile1dCacheHint, + TmaOperation::PrefetchTile2dCacheHint, + TmaOperation::PrefetchTile3dCacheHint, + TmaOperation::PrefetchTile4dCacheHint, + TmaOperation::PrefetchTile5dCacheHint, + TmaOperation::PrefetchTileGather4TwoDimensionalCacheHint, +]; + +pub(in crate::resolve) const TMA_REDUCTION_OPERATIONS: [TmaReductionOperation; 8] = [ + TmaReductionOperation::Add, + TmaReductionOperation::And, + TmaReductionOperation::Dec, + TmaReductionOperation::Inc, + TmaReductionOperation::Max, + TmaReductionOperation::Min, + TmaReductionOperation::Or, + TmaReductionOperation::Xor, +]; + +pub(in crate::resolve) fn tma_reduction_matrix() -> Vec { + let mut reductions = Vec::with_capacity(64); + for operation in TMA_REDUCTION_OPERATIONS { + for dimensions in 1..=5 { + reductions.push(TmaReduction { + operation, + load_mode: TmaReductionLoadMode::Tile, + dimensions, + }); + } + for dimensions in 3..=5 { + reductions.push(TmaReduction { + operation, + load_mode: TmaReductionLoadMode::Im2col, + dimensions, + }); + } + } + reductions +} + +#[cfg(test)] +pub(in crate::resolve) fn tma_reduction_admission_variants() -> Vec { + tma_reduction_matrix() + .into_iter() + .enumerate() + .map(|(index, reduction)| TmaReductionAdmissionVariant { + // Preserve the repository's current ledger assignments in this test fixture. + abi_id: format!("i{:04}", 923 + index), + operation: reduction.operation, + load_mode: reduction.load_mode, + dimensions: reduction.dimensions, + }) + .collect() +} + +pub(in crate::resolve) fn tma_reduction_operation_name( + operation: TmaReductionOperation, +) -> (&'static str, &'static str) { + match operation { + TmaReductionOperation::Add => ("add", "Add"), + TmaReductionOperation::And => ("and", "And"), + TmaReductionOperation::Dec => ("dec", "Dec"), + TmaReductionOperation::Inc => ("inc", "Inc"), + TmaReductionOperation::Max => ("max", "Max"), + TmaReductionOperation::Min => ("min", "Min"), + TmaReductionOperation::Or => ("or", "Or"), + TmaReductionOperation::Xor => ("xor", "Xor"), + } +} + +pub(in crate::resolve) fn tma_reduction_load_mode_name( + load_mode: TmaReductionLoadMode, +) -> (&'static str, &'static str, &'static str) { + match load_mode { + TmaReductionLoadMode::Tile => ("tile", "Tile", "tile"), + TmaReductionLoadMode::Im2col => ("im2col", "Im2col", "im2col_no_offs"), + } +} + +pub(in crate::resolve) struct TmaReductionRecipe { + id: String, + operation_key: String, + source_record: String, + llvm_symbol: String, + rust_arguments: Vec<&'static str>, + dialect_op_type: String, + dialect_op_name: String, + dialect_operands: Vec<&'static str>, + llvm_arguments: Vec<&'static str>, + modifiers: Vec, + operands: Vec, + summary: String, +} + +pub(in crate::resolve) fn tma_reduction_recipe( + reduction: TmaReduction, +) -> Result { + ensure!( + (1..=5).contains(&reduction.dimensions), + "TMA reduction dimensionality must be in 1..=5" + ); + ensure!( + reduction.load_mode != TmaReductionLoadMode::Im2col || reduction.dimensions >= 3, + "TMA reduction im2col mode requires at least three dimensions" + ); + + let (operation, operation_camel) = tma_reduction_operation_name(reduction.operation); + let (load_mode, load_mode_camel, ptx_load_mode) = + tma_reduction_load_mode_name(reduction.load_mode); + let dimensions = reduction.dimensions as usize; + let id = format!( + "cp_async_bulk_tensor_reduce_{operation}_{load_mode}_{}d", + reduction.dimensions + ); + let mut rust_arguments = vec!["*const u8", "*const u8"]; + rust_arguments.extend(std::iter::repeat_n("i32", dimensions)); + let mut dialect_operands = vec!["ptr", "ptr"]; + dialect_operands.extend(std::iter::repeat_n("i32", dimensions)); + let mut llvm_arguments = vec!["shared_ptr", "ptr"]; + llvm_arguments.extend(std::iter::repeat_n("i32", dimensions)); + llvm_arguments.extend(["i64", "i1"]); + + Ok(TmaReductionRecipe { + operation_key: format!( + "memory.reduce.async.bulk.tensor.{operation}.{load_mode}.{}d", + reduction.dimensions + ), + source_record: format!( + "int_nvvm_cp_async_bulk_tensor_reduce_{operation}_{load_mode}_{}d", + reduction.dimensions + ), + llvm_symbol: format!( + "llvm.nvvm.cp.async.bulk.tensor.reduce.{operation}.{load_mode}.{}d", + reduction.dimensions + ), + dialect_op_type: format!( + "CpAsyncBulkTensorReduce{operation_camel}{load_mode_camel}{}dOp", + reduction.dimensions + ), + dialect_op_name: format!( + "nvvm.cp_async_bulk_tensor_reduce_{operation}_{load_mode}_{}d", + reduction.dimensions + ), + rust_arguments, + dialect_operands, + llvm_arguments, + modifiers: vec![ + "reduce".into(), + "async".into(), + "bulk".into(), + "tensor".into(), + format!("{}d", reduction.dimensions), + "global".into(), + "shared::cta".into(), + operation.into(), + ptx_load_mode.into(), + "bulk_group".into(), + ], + operands: vec![OperandPattern::Address, OperandPattern::Address], + summary: format!( + "Starts a TMA tensor {operation} reduction from shared to global memory in {load_mode} mode." + ), + id, + }) +} + +pub(in crate::resolve) fn expand_tma_reduction_variant( + admission: &TmaAdmission, + variant: &TmaReductionAdmissionVariant, +) -> Result { + validate_abi_id(&variant.abi_id)?; + let reduction = TmaReduction { + operation: variant.operation, + load_mode: variant.load_mode, + dimensions: variant.dimensions, + }; + let recipe = tma_reduction_recipe(reduction)?; + + Ok(OverlayIntrinsic { + id: recipe.id.clone(), + abi_id: variant.abi_id.clone(), + operation_key: recipe.operation_key.clone(), + family: "tma".into(), + source: None, + source_record: Some(recipe.source_record.clone()), + rust_module: "tma".into(), + rust_name: recipe.id.clone(), + rust_arguments: recipe + .rust_arguments + .iter() + .map(|value| (*value).into()) + .collect(), + rust_result: "()".into(), + safe: false, + must_use: false, + safe_allowlist_reason: None, + public_rust_path: format!("cuda_intrinsics::tma::{}", recipe.id), + compatibility_rust_paths: vec![format!("cuda_device::tma::{}", recipe.id)], + dialect_op_type: recipe.dialect_op_type.clone(), + dialect_op_name: recipe.dialect_op_name.clone(), + dialect_operands: recipe + .dialect_operands + .iter() + .map(|value| (*value).into()) + .collect(), + dialect_results: vec![], + llvm_symbol: Some(recipe.llvm_symbol.clone()), + resolved_llvm_symbol: None, + llvm_arguments: recipe + .llvm_arguments + .iter() + .map(|value| (*value).into()) + .collect(), + llvm_results: vec![], + pure: false, + memory: "read_write".into(), + convergent: true, + execution_scope: "thread".into(), + minimum_ptx: "8.0".into(), + minimum_sm: Some("sm_90".into()), + ptx_result: "()".into(), + targets: "all".into(), + ptx_isa_version: "9.3".into(), + ptx_isa_section: + "9.7.9.26.5.3 Data Movement and Conversion Instructions: cp.reduce.async.bulk.tensor" + .into(), + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-reduce-async-bulk-tensor".into(), + lowering: "generated_tma".into(), + backend_lowerings: vec![ + OverlayBackendLowering { + backend: IntrinsicBackend::LlvmNvptx, + mechanism: BackendLoweringMechanism::TypedNvvm, + evidence_profile: admission + .reduce_llvm_evidence_profile + .as_ref() + .expect("validated TMA reduction LLVM evidence profile") + .clone(), + targets: None, + minimum_ptx: Some("8.0".into()), + minimum_sm: Some("sm_90".into()), + }, + OverlayBackendLowering { + backend: IntrinsicBackend::LibNvvm, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: admission + .reduce_libnvvm_evidence_profile + .as_ref() + .expect("validated TMA reduction libNVVM evidence profile") + .clone(), + targets: None, + minimum_ptx: Some("8.0".into()), + minimum_sm: Some("sm_90".into()), + }, + ], + packed_atomic: None, + redux: None, + vote: None, + active_mask: None, + warp_match: None, + warp_barrier: None, + warp_shuffle: None, + dot_product: None, + packed_alu: None, + integer_minmax: None, + packed_conversion: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cp_async_copy: None, + cp_async_control: None, + cp_async_mbarrier: None, + mbarrier_basic: None, + movmatrix: None, + mbarrier_extended: None, + register_mma: None, + sparse_mma: None, + prmt: None, + cluster_barrier: None, + wgmma_control: None, + special_register: None, + debug_control: None, + cluster_memory: None, + clc: None, + tma: Some(Tma { + operation: TmaOperation::Reduce, + reduction: Some(reduction), + adapter: TmaAdapter::ReductionPointersCoordinatesInjectDefaults, + runtime_validation: admission.runtime_validation, + }), + tcgen05: None, + ldmatrix_variant: None, + ldmatrix_safety: None, + ldmatrix_adapter: None, + selected_address_space: None, + expected_ptx: InstructionPattern { + mnemonic: "cp".into(), + modifiers: recipe.modifiers, + operands: recipe.operands, + }, + summary: recipe.summary, + }) +} + +pub(in crate::resolve) struct TmaRecipe { + pub(in crate::resolve) operation: TmaOperation, + pub(in crate::resolve) abi_id: &'static str, + pub(in crate::resolve) id: &'static str, + pub(in crate::resolve) operation_key: &'static str, + pub(in crate::resolve) source_record: &'static str, + pub(in crate::resolve) llvm_symbol: &'static str, + pub(in crate::resolve) resolved_llvm_symbol: Option, + pub(in crate::resolve) selected_address_space: Option, + pub(in crate::resolve) llvm_mechanism: BackendLoweringMechanism, + pub(in crate::resolve) rust_arguments: Vec<&'static str>, + pub(in crate::resolve) dialect_op_type: &'static str, + pub(in crate::resolve) dialect_op_name: &'static str, + pub(in crate::resolve) dialect_operands: Vec<&'static str>, + pub(in crate::resolve) llvm_arguments: Vec<&'static str>, + pub(in crate::resolve) adapter: TmaAdapter, + pub(in crate::resolve) safe: bool, + pub(in crate::resolve) safe_reason: Option<&'static str>, + pub(in crate::resolve) convergent: bool, + pub(in crate::resolve) minimum_ptx: &'static str, + pub(in crate::resolve) minimum_sm: Option<&'static str>, + pub(in crate::resolve) targets: &'static str, + pub(in crate::resolve) memory: &'static str, + pub(in crate::resolve) ptx_isa_section: &'static str, + pub(in crate::resolve) ptx_isa_url: &'static str, + pub(in crate::resolve) mnemonic: &'static str, + pub(in crate::resolve) modifiers: Vec, + pub(in crate::resolve) operands: Vec, + pub(in crate::resolve) summary: &'static str, +} + +pub(in crate::resolve) fn tma_recipe(operation: TmaOperation) -> TmaRecipe { + let (abi_id, id, operation_key, source_record, llvm_symbol, op_type, op_name) = match operation + { + TmaOperation::Reduce => unreachable!("TMA reductions use tma_reduction_recipe"), + TmaOperation::G2sTile1d => ( + "i0328", + "cp_async_bulk_tensor_1d_g2s", + "memory.copy.async.bulk.tensor.g2s.tile.1d", + "int_nvvm_cp_async_bulk_tensor_g2s_tile_1d", + "llvm.nvvm.cp.async.bulk.tensor.g2s.tile.1d", + "CpAsyncBulkTensorG2sTile1dOp", + "nvvm.cp_async_bulk_tensor_g2s_tile_1d", + ), + TmaOperation::G2sTile2d => ( + "i0329", + "cp_async_bulk_tensor_2d_g2s", + "memory.copy.async.bulk.tensor.g2s.tile.2d", + "int_nvvm_cp_async_bulk_tensor_g2s_tile_2d", + "llvm.nvvm.cp.async.bulk.tensor.g2s.tile.2d", + "CpAsyncBulkTensorG2sTile2dOp", + "nvvm.cp_async_bulk_tensor_g2s_tile_2d", + ), + TmaOperation::G2sTile2dMulticast => ( + "i0330", + "cp_async_bulk_tensor_2d_g2s_multicast", + "memory.copy.async.bulk.tensor.g2s.tile.2d.multicast", + "int_nvvm_cp_async_bulk_tensor_g2s_tile_2d", + "llvm.nvvm.cp.async.bulk.tensor.g2s.tile.2d", + "CpAsyncBulkTensorG2sTile2dMulticastOp", + "nvvm.cp_async_bulk_tensor_g2s_tile_2d_multicast", + ), + TmaOperation::G2sTile2dMulticastCg2 => ( + "i0331", + "cp_async_bulk_tensor_2d_g2s_multicast_cg2", + "memory.copy.async.bulk.tensor.g2s.tile.2d.multicast.cta_group_2", + "int_nvvm_cp_async_bulk_tensor_g2s_tile_2d", + "llvm.nvvm.cp.async.bulk.tensor.g2s.tile.2d", + "CpAsyncBulkTensorG2sTile2dMulticastCg2Op", + "nvvm.cp_async_bulk_tensor_g2s_tile_2d_multicast_cg2", + ), + TmaOperation::G2sTile3d => ( + "i0332", + "cp_async_bulk_tensor_3d_g2s", + "memory.copy.async.bulk.tensor.g2s.tile.3d", + "int_nvvm_cp_async_bulk_tensor_g2s_tile_3d", + "llvm.nvvm.cp.async.bulk.tensor.g2s.tile.3d", + "CpAsyncBulkTensorG2sTile3dOp", + "nvvm.cp_async_bulk_tensor_g2s_tile_3d", + ), + TmaOperation::G2sTile4d => ( + "i0333", + "cp_async_bulk_tensor_4d_g2s", + "memory.copy.async.bulk.tensor.g2s.tile.4d", + "int_nvvm_cp_async_bulk_tensor_g2s_tile_4d", + "llvm.nvvm.cp.async.bulk.tensor.g2s.tile.4d", + "CpAsyncBulkTensorG2sTile4dOp", + "nvvm.cp_async_bulk_tensor_g2s_tile_4d", + ), + TmaOperation::G2sTile5d => ( + "i0334", + "cp_async_bulk_tensor_5d_g2s", + "memory.copy.async.bulk.tensor.g2s.tile.5d", + "int_nvvm_cp_async_bulk_tensor_g2s_tile_5d", + "llvm.nvvm.cp.async.bulk.tensor.g2s.tile.5d", + "CpAsyncBulkTensorG2sTile5dOp", + "nvvm.cp_async_bulk_tensor_g2s_tile_5d", + ), + TmaOperation::S2gTile1d => ( + "i0335", + "cp_async_bulk_tensor_1d_s2g", + "memory.copy.async.bulk.tensor.s2g.tile.1d", + "int_nvvm_cp_async_bulk_tensor_s2g_tile_1d", + "llvm.nvvm.cp.async.bulk.tensor.s2g.tile.1d", + "CpAsyncBulkTensorS2gTile1dOp", + "nvvm.cp_async_bulk_tensor_s2g_tile_1d", + ), + TmaOperation::S2gTile2d => ( + "i0336", + "cp_async_bulk_tensor_2d_s2g", + "memory.copy.async.bulk.tensor.s2g.tile.2d", + "int_nvvm_cp_async_bulk_tensor_s2g_tile_2d", + "llvm.nvvm.cp.async.bulk.tensor.s2g.tile.2d", + "CpAsyncBulkTensorS2gTile2dOp", + "nvvm.cp_async_bulk_tensor_s2g_tile_2d", + ), + TmaOperation::S2gTile3d => ( + "i0337", + "cp_async_bulk_tensor_3d_s2g", + "memory.copy.async.bulk.tensor.s2g.tile.3d", + "int_nvvm_cp_async_bulk_tensor_s2g_tile_3d", + "llvm.nvvm.cp.async.bulk.tensor.s2g.tile.3d", + "CpAsyncBulkTensorS2gTile3dOp", + "nvvm.cp_async_bulk_tensor_s2g_tile_3d", + ), + TmaOperation::S2gTile4d => ( + "i0338", + "cp_async_bulk_tensor_4d_s2g", + "memory.copy.async.bulk.tensor.s2g.tile.4d", + "int_nvvm_cp_async_bulk_tensor_s2g_tile_4d", + "llvm.nvvm.cp.async.bulk.tensor.s2g.tile.4d", + "CpAsyncBulkTensorS2gTile4dOp", + "nvvm.cp_async_bulk_tensor_s2g_tile_4d", + ), + TmaOperation::S2gTile5d => ( + "i0339", + "cp_async_bulk_tensor_5d_s2g", + "memory.copy.async.bulk.tensor.s2g.tile.5d", + "int_nvvm_cp_async_bulk_tensor_s2g_tile_5d", + "llvm.nvvm.cp.async.bulk.tensor.s2g.tile.5d", + "CpAsyncBulkTensorS2gTile5dOp", + "nvvm.cp_async_bulk_tensor_s2g_tile_5d", + ), + TmaOperation::CommitGroup => ( + "i0340", + "cp_async_bulk_commit_group", + "memory.copy.async.bulk.group.commit", + "int_nvvm_cp_async_bulk_commit_group", + "llvm.nvvm.cp.async.bulk.commit.group", + "CpAsyncBulkCommitGroupOp", + "nvvm.cp_async_bulk_commit_group", + ), + TmaOperation::WaitGroup => ( + "i0341", + "cp_async_bulk_wait_group", + "memory.copy.async.bulk.group.wait_max_pending", + "int_nvvm_cp_async_bulk_wait_group", + "llvm.nvvm.cp.async.bulk.wait.group", + "CpAsyncBulkWaitGroupOp", + "nvvm.cp_async_bulk_wait_group", + ), + TmaOperation::WaitGroupRead => ( + "i0342", + "cp_async_bulk_wait_group_read", + "memory.copy.async.bulk.group.wait_read_max_pending", + "int_nvvm_cp_async_bulk_wait_group_read", + "llvm.nvvm.cp.async.bulk.wait.group.read", + "CpAsyncBulkWaitGroupReadOp", + "nvvm.cp_async_bulk_wait_group_read", + ), + TmaOperation::PrefetchTensorMap => ( + "i0887", + "prefetch_tma_descriptor", + "memory.prefetch.tensor_map", + "int_nvvm_prefetch_tensormap", + "llvm.nvvm.prefetch.tensormap", + "PrefetchTensorMapOp", + "nvvm.prefetch_tensormap", + ), + TmaOperation::PrefetchTile1d => ( + "i0888", + "cp_async_bulk_prefetch_tensor_1d_l2", + "memory.prefetch.async.bulk.tensor.global.tile.1d.l2", + "int_nvvm_cp_async_bulk_tensor_prefetch_tile_1d", + "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.1d", + "CpAsyncBulkPrefetchTensor1dL2Op", + "nvvm.cp_async_bulk_prefetch_tensor_1d_l2", + ), + TmaOperation::PrefetchTile2d => ( + "i0889", + "cp_async_bulk_prefetch_tensor_2d_l2", + "memory.prefetch.async.bulk.tensor.global.tile.2d.l2", + "int_nvvm_cp_async_bulk_tensor_prefetch_tile_2d", + "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.2d", + "CpAsyncBulkPrefetchTensor2dL2Op", + "nvvm.cp_async_bulk_prefetch_tensor_2d_l2", + ), + TmaOperation::PrefetchTile3d => ( + "i0890", + "cp_async_bulk_prefetch_tensor_3d_l2", + "memory.prefetch.async.bulk.tensor.global.tile.3d.l2", + "int_nvvm_cp_async_bulk_tensor_prefetch_tile_3d", + "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.3d", + "CpAsyncBulkPrefetchTensor3dL2Op", + "nvvm.cp_async_bulk_prefetch_tensor_3d_l2", + ), + TmaOperation::PrefetchTile4d => ( + "i0891", + "cp_async_bulk_prefetch_tensor_4d_l2", + "memory.prefetch.async.bulk.tensor.global.tile.4d.l2", + "int_nvvm_cp_async_bulk_tensor_prefetch_tile_4d", + "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.4d", + "CpAsyncBulkPrefetchTensor4dL2Op", + "nvvm.cp_async_bulk_prefetch_tensor_4d_l2", + ), + TmaOperation::PrefetchTile5d => ( + "i0892", + "cp_async_bulk_prefetch_tensor_5d_l2", + "memory.prefetch.async.bulk.tensor.global.tile.5d.l2", + "int_nvvm_cp_async_bulk_tensor_prefetch_tile_5d", + "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.5d", + "CpAsyncBulkPrefetchTensor5dL2Op", + "nvvm.cp_async_bulk_prefetch_tensor_5d_l2", + ), + TmaOperation::PrefetchTileGather4TwoDimensional => ( + "i0893", + "cp_async_bulk_prefetch_tensor_gather4_2d_l2", + "memory.prefetch.async.bulk.tensor.global.tile.gather4.2d.l2", + "int_nvvm_cp_async_bulk_tensor_prefetch_tile_gather4_2d", + "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.gather4.2d", + "CpAsyncBulkPrefetchTensorGather4TwoDimensionalL2Op", + "nvvm.cp_async_bulk_prefetch_tensor_gather4_2d_l2", + ), + TmaOperation::PrefetchTile1dCacheHint => ( + "i0917", + "cp_async_bulk_prefetch_tensor_1d_l2_cache_hint", + "memory.prefetch.async.bulk.tensor.global.tile.1d.l2.cache_hint", + "int_nvvm_cp_async_bulk_tensor_prefetch_tile_1d", + "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.1d", + "CpAsyncBulkPrefetchTensor1dL2CacheHintOp", + "nvvm.cp_async_bulk_prefetch_tensor_1d_l2_cache_hint", + ), + TmaOperation::PrefetchTile2dCacheHint => ( + "i0918", + "cp_async_bulk_prefetch_tensor_2d_l2_cache_hint", + "memory.prefetch.async.bulk.tensor.global.tile.2d.l2.cache_hint", + "int_nvvm_cp_async_bulk_tensor_prefetch_tile_2d", + "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.2d", + "CpAsyncBulkPrefetchTensor2dL2CacheHintOp", + "nvvm.cp_async_bulk_prefetch_tensor_2d_l2_cache_hint", + ), + TmaOperation::PrefetchTile3dCacheHint => ( + "i0919", + "cp_async_bulk_prefetch_tensor_3d_l2_cache_hint", + "memory.prefetch.async.bulk.tensor.global.tile.3d.l2.cache_hint", + "int_nvvm_cp_async_bulk_tensor_prefetch_tile_3d", + "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.3d", + "CpAsyncBulkPrefetchTensor3dL2CacheHintOp", + "nvvm.cp_async_bulk_prefetch_tensor_3d_l2_cache_hint", + ), + TmaOperation::PrefetchTile4dCacheHint => ( + "i0920", + "cp_async_bulk_prefetch_tensor_4d_l2_cache_hint", + "memory.prefetch.async.bulk.tensor.global.tile.4d.l2.cache_hint", + "int_nvvm_cp_async_bulk_tensor_prefetch_tile_4d", + "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.4d", + "CpAsyncBulkPrefetchTensor4dL2CacheHintOp", + "nvvm.cp_async_bulk_prefetch_tensor_4d_l2_cache_hint", + ), + TmaOperation::PrefetchTile5dCacheHint => ( + "i0921", + "cp_async_bulk_prefetch_tensor_5d_l2_cache_hint", + "memory.prefetch.async.bulk.tensor.global.tile.5d.l2.cache_hint", + "int_nvvm_cp_async_bulk_tensor_prefetch_tile_5d", + "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.5d", + "CpAsyncBulkPrefetchTensor5dL2CacheHintOp", + "nvvm.cp_async_bulk_prefetch_tensor_5d_l2_cache_hint", + ), + TmaOperation::PrefetchTileGather4TwoDimensionalCacheHint => ( + "i0922", + "cp_async_bulk_prefetch_tensor_gather4_2d_l2_cache_hint", + "memory.prefetch.async.bulk.tensor.global.tile.gather4.2d.l2.cache_hint", + "int_nvvm_cp_async_bulk_tensor_prefetch_tile_gather4_2d", + "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.gather4.2d", + "CpAsyncBulkPrefetchTensorGather4TwoDimensionalL2CacheHintOp", + "nvvm.cp_async_bulk_prefetch_tensor_gather4_2d_l2_cache_hint", + ), + TmaOperation::ReplaceBoxDim => ( + "i0894", + "tensormap_replace_box_dim", + "memory.tensor_map.replace.box_dim", + "int_nvvm_tensormap_replace_box_dim", + "llvm.nvvm.tensormap.replace.box.dim", + "ReplaceTensorMapBoxDimOp", + "nvvm.tensormap_replace_box_dim", + ), + TmaOperation::ReplaceElementStride => ( + "i0895", + "tensormap_replace_element_stride", + "memory.tensor_map.replace.element_stride", + "int_nvvm_tensormap_replace_element_stride", + "llvm.nvvm.tensormap.replace.element.stride", + "ReplaceTensorMapElementStrideOp", + "nvvm.tensormap_replace_element_stride", + ), + TmaOperation::ReplaceElementType => ( + "i0896", + "tensormap_replace_element_type", + "memory.tensor_map.replace.element_type", + "int_nvvm_tensormap_replace_elemtype", + "llvm.nvvm.tensormap.replace.elemtype", + "ReplaceTensorMapElementTypeOp", + "nvvm.tensormap_replace_element_type", + ), + TmaOperation::ReplaceFillMode => ( + "i0897", + "tensormap_replace_fill_mode", + "memory.tensor_map.replace.fill_mode", + "int_nvvm_tensormap_replace_fill_mode", + "llvm.nvvm.tensormap.replace.fill.mode", + "ReplaceTensorMapFillModeOp", + "nvvm.tensormap_replace_fill_mode", + ), + TmaOperation::ReplaceGlobalAddress => ( + "i0898", + "tensormap_replace_global_address", + "memory.tensor_map.replace.global_address", + "int_nvvm_tensormap_replace_global_address", + "llvm.nvvm.tensormap.replace.global.address", + "ReplaceTensorMapGlobalAddressOp", + "nvvm.tensormap_replace_global_address", + ), + TmaOperation::ReplaceGlobalDim => ( + "i0899", + "tensormap_replace_global_dim", + "memory.tensor_map.replace.global_dim", + "int_nvvm_tensormap_replace_global_dim", + "llvm.nvvm.tensormap.replace.global.dim", + "ReplaceTensorMapGlobalDimOp", + "nvvm.tensormap_replace_global_dim", + ), + TmaOperation::ReplaceGlobalStride => ( + "i0900", + "tensormap_replace_global_stride", + "memory.tensor_map.replace.global_stride", + "int_nvvm_tensormap_replace_global_stride", + "llvm.nvvm.tensormap.replace.global.stride", + "ReplaceTensorMapGlobalStrideOp", + "nvvm.tensormap_replace_global_stride", + ), + TmaOperation::ReplaceInterleaveLayout => ( + "i0901", + "tensormap_replace_interleave_layout", + "memory.tensor_map.replace.interleave_layout", + "int_nvvm_tensormap_replace_interleave_layout", + "llvm.nvvm.tensormap.replace.interleave.layout", + "ReplaceTensorMapInterleaveLayoutOp", + "nvvm.tensormap_replace_interleave_layout", + ), + TmaOperation::ReplaceRank => ( + "i0902", + "tensormap_replace_rank", + "memory.tensor_map.replace.rank", + "int_nvvm_tensormap_replace_rank", + "llvm.nvvm.tensormap.replace.rank", + "ReplaceTensorMapRankOp", + "nvvm.tensormap_replace_rank", + ), + TmaOperation::ReplaceSwizzleAtomicity => ( + "i0903", + "tensormap_replace_swizzle_atomicity", + "memory.tensor_map.replace.swizzle_atomicity", + "int_nvvm_tensormap_replace_swizzle_atomicity", + "llvm.nvvm.tensormap.replace.swizzle.atomicity", + "ReplaceTensorMapSwizzleAtomicityOp", + "nvvm.tensormap_replace_swizzle_atomicity", + ), + TmaOperation::ReplaceSwizzleMode => ( + "i0904", + "tensormap_replace_swizzle_mode", + "memory.tensor_map.replace.swizzle_mode", + "int_nvvm_tensormap_replace_swizzle_mode", + "llvm.nvvm.tensormap.replace.swizzle.mode", + "ReplaceTensorMapSwizzleModeOp", + "nvvm.tensormap_replace_swizzle_mode", + ), + TmaOperation::FenceProxyTensorMapAcquireCluster => ( + "i0905", + "fence_proxy_tensormap_generic_acquire_cluster", + "memory.fence.proxy.tensor_map.generic.acquire.cluster", + "int_nvvm_fence_proxy_tensormap_generic_acquire_cluster", + "llvm.nvvm.fence.proxy.tensormap_generic.acquire.cluster", + "FenceProxyTensorMapGenericAcquireClusterOp", + "nvvm.fence_proxy_tensormap_generic_acquire_cluster", + ), + TmaOperation::FenceProxyTensorMapAcquireCta => ( + "i0906", + "fence_proxy_tensormap_generic_acquire_cta", + "memory.fence.proxy.tensor_map.generic.acquire.cta", + "int_nvvm_fence_proxy_tensormap_generic_acquire_cta", + "llvm.nvvm.fence.proxy.tensormap_generic.acquire.cta", + "FenceProxyTensorMapGenericAcquireCtaOp", + "nvvm.fence_proxy_tensormap_generic_acquire_cta", + ), + TmaOperation::FenceProxyTensorMapAcquireGpu => ( + "i0907", + "fence_proxy_tensormap_generic_acquire_gpu", + "memory.fence.proxy.tensor_map.generic.acquire.gpu", + "int_nvvm_fence_proxy_tensormap_generic_acquire_gpu", + "llvm.nvvm.fence.proxy.tensormap_generic.acquire.gpu", + "FenceProxyTensorMapGenericAcquireGpuOp", + "nvvm.fence_proxy_tensormap_generic_acquire_gpu", + ), + TmaOperation::FenceProxyTensorMapAcquireSystem => ( + "i0908", + "fence_proxy_tensormap_generic_acquire_system", + "memory.fence.proxy.tensor_map.generic.acquire.system", + "int_nvvm_fence_proxy_tensormap_generic_acquire_sys", + "llvm.nvvm.fence.proxy.tensormap_generic.acquire.sys", + "FenceProxyTensorMapGenericAcquireSystemOp", + "nvvm.fence_proxy_tensormap_generic_acquire_system", + ), + TmaOperation::FenceProxyTensorMapReleaseCluster => ( + "i0909", + "fence_proxy_tensormap_generic_release_cluster", + "memory.fence.proxy.tensor_map.generic.release.cluster", + "int_nvvm_fence_proxy_tensormap_generic_release_cluster", + "llvm.nvvm.fence.proxy.tensormap_generic.release.cluster", + "FenceProxyTensorMapGenericReleaseClusterOp", + "nvvm.fence_proxy_tensormap_generic_release_cluster", + ), + TmaOperation::FenceProxyTensorMapReleaseCta => ( + "i0910", + "fence_proxy_tensormap_generic_release_cta", + "memory.fence.proxy.tensor_map.generic.release.cta", + "int_nvvm_fence_proxy_tensormap_generic_release_cta", + "llvm.nvvm.fence.proxy.tensormap_generic.release.cta", + "FenceProxyTensorMapGenericReleaseCtaOp", + "nvvm.fence_proxy_tensormap_generic_release_cta", + ), + TmaOperation::FenceProxyTensorMapReleaseGpu => ( + "i0911", + "fence_proxy_tensormap_generic_release_gpu", + "memory.fence.proxy.tensor_map.generic.release.gpu", + "int_nvvm_fence_proxy_tensormap_generic_release_gpu", + "llvm.nvvm.fence.proxy.tensormap_generic.release.gpu", + "FenceProxyTensorMapGenericReleaseGpuOp", + "nvvm.fence_proxy_tensormap_generic_release_gpu", + ), + TmaOperation::FenceProxyTensorMapReleaseSystem => ( + "i0912", + "fence_proxy_tensormap_generic_release_system", + "memory.fence.proxy.tensor_map.generic.release.system", + "int_nvvm_fence_proxy_tensormap_generic_release_sys", + "llvm.nvvm.fence.proxy.tensormap_generic.release.sys", + "FenceProxyTensorMapGenericReleaseSystemOp", + "nvvm.fence_proxy_tensormap_generic_release_system", + ), + }; + + let dimensions = operation.dimensions(); + let is_g2s = matches!( + operation, + TmaOperation::G2sTile1d + | TmaOperation::G2sTile2d + | TmaOperation::G2sTile2dMulticast + | TmaOperation::G2sTile2dMulticastCg2 + | TmaOperation::G2sTile3d + | TmaOperation::G2sTile4d + | TmaOperation::G2sTile5d + ); + let is_s2g = matches!( + operation, + TmaOperation::S2gTile1d + | TmaOperation::S2gTile2d + | TmaOperation::S2gTile3d + | TmaOperation::S2gTile4d + | TmaOperation::S2gTile5d + ); + let multicast = matches!( + operation, + TmaOperation::G2sTile2dMulticast | TmaOperation::G2sTile2dMulticastCg2 + ); + let cg2 = operation == TmaOperation::G2sTile2dMulticastCg2; + let prefetch_coordinates = operation.prefetch_coordinate_count(); + let is_release_fence = matches!( + operation, + TmaOperation::FenceProxyTensorMapReleaseCluster + | TmaOperation::FenceProxyTensorMapReleaseCta + | TmaOperation::FenceProxyTensorMapReleaseGpu + | TmaOperation::FenceProxyTensorMapReleaseSystem + ); + let is_acquire_fence = matches!( + operation, + TmaOperation::FenceProxyTensorMapAcquireCluster + | TmaOperation::FenceProxyTensorMapAcquireCta + | TmaOperation::FenceProxyTensorMapAcquireGpu + | TmaOperation::FenceProxyTensorMapAcquireSystem + ); + + let mut rust_arguments = Vec::new(); + let mut dialect_operands = Vec::new(); + let mut llvm_arguments = Vec::new(); + let (adapter, safe, safe_reason, convergent, summary) = if is_g2s { + rust_arguments.extend(["*mut u8", "*const u8"]); + rust_arguments.extend(std::iter::repeat_n("i32", dimensions.unwrap())); + rust_arguments.push("*mut u64"); + if multicast { + rust_arguments.push("u16"); + } + dialect_operands.extend(["ptr", "ptr", "ptr"]); + dialect_operands.extend(std::iter::repeat_n("i32", dimensions.unwrap())); + dialect_operands.extend(["i16", "i64"]); + llvm_arguments.extend(["shared_cluster_ptr", "shared_ptr", "ptr"]); + llvm_arguments.extend(std::iter::repeat_n("i32", dimensions.unwrap())); + llvm_arguments.extend(["i16", "i64", "i1", "i1", "i32"]); + ( + if multicast { + TmaAdapter::G2sPointersCoordinatesBarrierMaskInjectDefaults + } else { + TmaAdapter::G2sPointersCoordinatesBarrierInjectDefaults + }, + false, + None, + true, + if multicast { + "Starts a multicast TMA tile copy from global to cluster shared memory." + } else { + "Starts a TMA tile copy from global to cluster shared memory." + }, + ) + } else if is_s2g { + rust_arguments.extend(["*const u8", "*const u8"]); + rust_arguments.extend(std::iter::repeat_n("i32", dimensions.unwrap())); + dialect_operands.extend(["ptr", "ptr"]); + dialect_operands.extend(std::iter::repeat_n("i32", dimensions.unwrap())); + llvm_arguments.extend(["shared_ptr", "ptr"]); + llvm_arguments.extend(std::iter::repeat_n("i32", dimensions.unwrap())); + llvm_arguments.extend(["i64", "i1"]); + ( + TmaAdapter::S2gPointersCoordinatesInjectDefaults, + false, + None, + true, + "Starts a TMA tile copy from shared to global memory.", + ) + } else if operation == TmaOperation::CommitGroup || is_release_fence { + ( + TmaAdapter::NoOperands, + true, + Some(if operation == TmaOperation::CommitGroup { + "committing this thread's pending bulk-copy group has no Rust memory-safety precondition." + } else { + "publishing prior tensor-map writes has no Rust memory-safety precondition." + }), + false, + if operation == TmaOperation::CommitGroup { + "Commits this thread's pending asynchronous bulk copies as one group." + } else { + "Publishes prior generic-proxy tensor-map writes to the selected tensor-map proxy scope." + }, + ) + } else if operation == TmaOperation::PrefetchTensorMap { + rust_arguments.push("*const u8"); + dialect_operands.push("ptr"); + llvm_arguments.push("anyptr"); + ( + TmaAdapter::DescriptorPointer, + false, + None, + false, + "Prefetches a live tensor-map descriptor into the tensor-map cache.", + ) + } else if let Some(coordinate_count) = prefetch_coordinates { + rust_arguments.push("*const u8"); + rust_arguments.extend(std::iter::repeat_n("i32", coordinate_count)); + dialect_operands.push("ptr"); + dialect_operands.extend(std::iter::repeat_n("i32", coordinate_count)); + if operation.uses_prefetch_cache_hint() { + rust_arguments.push("u64"); + dialect_operands.push("i64"); + } + llvm_arguments.push("ptr"); + llvm_arguments.extend(std::iter::repeat_n("i32", coordinate_count)); + llvm_arguments.extend(["i64", "i1"]); + ( + if operation.uses_prefetch_cache_hint() { + TmaAdapter::DescriptorCoordinatesCacheHintInjectFlag + } else { + TmaAdapter::DescriptorCoordinatesInjectDefaults + }, + false, + None, + true, + if operation.uses_prefetch_cache_hint() { + "Prefetches one tensor tile through a tensor map into L2 using an explicit cache hint." + } else { + "Prefetches one tensor tile through a tensor map into L2." + }, + ) + } else if operation == TmaOperation::ReplaceGlobalAddress { + rust_arguments.extend(["*mut u8", "*const u8"]); + dialect_operands.extend(["ptr", "ptr"]); + llvm_arguments.extend(["anyptr", "i64"]); + ( + TmaAdapter::DescriptorAndAddressPointers, + false, + None, + false, + "Replaces the global base address stored in a writable tensor-map descriptor.", + ) + } else if matches!( + operation, + TmaOperation::ReplaceBoxDim + | TmaOperation::ReplaceElementStride + | TmaOperation::ReplaceGlobalDim + ) { + rust_arguments.extend(["*mut u8", "u32", "u32"]); + dialect_operands.extend(["ptr", "i32", "i32"]); + llvm_arguments.extend(["anyptr", "i32", "i32"]); + ( + TmaAdapter::DescriptorOrdinalAndU32, + false, + None, + false, + "Replaces one indexed 32-bit field in a writable tensor-map descriptor.", + ) + } else if operation == TmaOperation::ReplaceGlobalStride { + rust_arguments.extend(["*mut u8", "u32", "u64"]); + dialect_operands.extend(["ptr", "i32", "i64"]); + llvm_arguments.extend(["anyptr", "i32", "i64"]); + ( + TmaAdapter::DescriptorOrdinalAndU64, + false, + None, + false, + "Replaces one indexed 64-bit stride in a writable tensor-map descriptor.", + ) + } else if matches!( + operation, + TmaOperation::ReplaceElementType + | TmaOperation::ReplaceFillMode + | TmaOperation::ReplaceInterleaveLayout + | TmaOperation::ReplaceSwizzleAtomicity + | TmaOperation::ReplaceSwizzleMode + ) { + rust_arguments.extend(["*mut u8", "u32"]); + dialect_operands.extend(["ptr", "i32"]); + llvm_arguments.extend(["anyptr", "i32"]); + ( + TmaAdapter::DescriptorAndImmediateU32, + false, + None, + false, + "Replaces one immediate-valued field in a writable tensor-map descriptor.", + ) + } else if operation == TmaOperation::ReplaceRank { + rust_arguments.extend(["*mut u8", "u32"]); + dialect_operands.extend(["ptr", "i32"]); + llvm_arguments.extend(["anyptr", "i32"]); + ( + TmaAdapter::DescriptorAndRuntimeU32, + false, + None, + false, + "Replaces the rank stored in a writable tensor-map descriptor.", + ) + } else if is_acquire_fence { + rust_arguments.push("*const u8"); + dialect_operands.push("ptr"); + llvm_arguments.extend(["ptr", "i32"]); + ( + TmaAdapter::DescriptorPointerInjectBytes, + false, + None, + false, + "Acquires a published tensor-map descriptor into the selected tensor-map proxy scope.", + ) + } else { + rust_arguments.push("u32"); + dialect_operands.push("i32"); + llvm_arguments.push("i32"); + ( + TmaAdapter::CompileTimeConstantMaxPending, + true, + Some( + "waiting on this thread's bulk-copy groups has no Rust memory-safety precondition.", + ), + false, + if operation == TmaOperation::WaitGroupRead { + "Waits for bulk-copy groups and completes their reads." + } else { + "Waits until at most the requested bulk-copy groups remain pending." + }, + ) + }; + + let is_replace = matches!( + operation, + TmaOperation::ReplaceBoxDim + | TmaOperation::ReplaceElementStride + | TmaOperation::ReplaceElementType + | TmaOperation::ReplaceFillMode + | TmaOperation::ReplaceGlobalAddress + | TmaOperation::ReplaceGlobalDim + | TmaOperation::ReplaceGlobalStride + | TmaOperation::ReplaceInterleaveLayout + | TmaOperation::ReplaceRank + | TmaOperation::ReplaceSwizzleAtomicity + | TmaOperation::ReplaceSwizzleMode + ); + let is_fence = is_acquire_fence || is_release_fence; + let descriptor_control = operation == TmaOperation::PrefetchTensorMap + || prefetch_coordinates.is_some() + || is_replace + || is_fence; + let blackwell_tma = cg2 + || matches!( + operation, + TmaOperation::PrefetchTileGather4TwoDimensional + | TmaOperation::PrefetchTileGather4TwoDimensionalCacheHint + ); + + let (mnemonic, modifiers, operands) = if is_g2s { + let mut modifiers = vec![ + "async".into(), + "bulk".into(), + "tensor".into(), + format!("{}d", dimensions.unwrap()), + "shared::cluster".into(), + "global".into(), + "tile".into(), + "mbarrier::complete_tx::bytes".into(), + ]; + if multicast { + modifiers.push("multicast::cluster".into()); + } + if cg2 { + modifiers.push("cta_group::2".into()); + } + let mut operands = vec![ + OperandPattern::Address, + OperandPattern::Address, + OperandPattern::Address, + ]; + if multicast { + operands.push(OperandPattern::Register); + } + ("cp", modifiers, operands) + } else if is_s2g { + ( + "cp", + vec![ + "async".into(), + "bulk".into(), + "tensor".into(), + format!("{}d", dimensions.unwrap()), + "global".into(), + "shared::cta".into(), + "tile".into(), + "bulk_group".into(), + ], + vec![OperandPattern::Address, OperandPattern::Address], + ) + } else if operation == TmaOperation::CommitGroup { + ( + "cp", + vec!["async".into(), "bulk".into(), "commit_group".into()], + vec![], + ) + } else if matches!( + operation, + TmaOperation::WaitGroup | TmaOperation::WaitGroupRead + ) { + let mut modifiers = vec!["async".into(), "bulk".into(), "wait_group".into()]; + if operation == TmaOperation::WaitGroupRead { + modifiers.push("read".into()); + } + ("cp", modifiers, vec![OperandPattern::Immediate]) + } else if operation == TmaOperation::PrefetchTensorMap { + ( + "prefetch", + vec!["tensormap".into()], + vec![OperandPattern::Address], + ) + } else if prefetch_coordinates.is_some() { + let (dimensionality, tile) = if matches!( + operation, + TmaOperation::PrefetchTileGather4TwoDimensional + | TmaOperation::PrefetchTileGather4TwoDimensionalCacheHint + ) { + ("2d", "tile::gather4") + } else { + ( + match prefetch_coordinates.unwrap() { + 1 => "1d", + 2 => "2d", + 3 => "3d", + 4 => "4d", + 5 => "5d", + _ => unreachable!("closed TMA prefetch dimensionality"), + }, + "tile", + ) + }; + let mut modifiers = vec![ + "async".into(), + "bulk".into(), + "prefetch".into(), + "tensor".into(), + dimensionality.into(), + "L2".into(), + "global".into(), + tile.into(), + ]; + let mut operands = vec![OperandPattern::Address]; + if operation.uses_prefetch_cache_hint() { + modifiers.push("L2::cache_hint".into()); + operands.push(OperandPattern::Register); + } + ("cp", modifiers, operands) + } else if is_replace { + let (field, width, operands) = match operation { + TmaOperation::ReplaceGlobalAddress => ( + "global_address", + "b64", + vec![OperandPattern::Address, OperandPattern::Register], + ), + TmaOperation::ReplaceRank => ( + "rank", + "b32", + vec![OperandPattern::Address, OperandPattern::Register], + ), + TmaOperation::ReplaceBoxDim => ( + "box_dim", + "b32", + vec![ + OperandPattern::Address, + OperandPattern::Immediate, + OperandPattern::Register, + ], + ), + TmaOperation::ReplaceElementStride => ( + "element_stride", + "b32", + vec![ + OperandPattern::Address, + OperandPattern::Immediate, + OperandPattern::Register, + ], + ), + TmaOperation::ReplaceGlobalDim => ( + "global_dim", + "b32", + vec![ + OperandPattern::Address, + OperandPattern::Immediate, + OperandPattern::Register, + ], + ), + TmaOperation::ReplaceGlobalStride => ( + "global_stride", + "b64", + vec![ + OperandPattern::Address, + OperandPattern::Immediate, + OperandPattern::Register, + ], + ), + TmaOperation::ReplaceElementType => ( + "elemtype", + "b32", + vec![OperandPattern::Address, OperandPattern::Immediate], + ), + TmaOperation::ReplaceFillMode => ( + "fill_mode", + "b32", + vec![OperandPattern::Address, OperandPattern::Immediate], + ), + TmaOperation::ReplaceInterleaveLayout => ( + "interleave_layout", + "b32", + vec![OperandPattern::Address, OperandPattern::Immediate], + ), + TmaOperation::ReplaceSwizzleAtomicity => ( + "swizzle_atomicity", + "b32", + vec![OperandPattern::Address, OperandPattern::Immediate], + ), + TmaOperation::ReplaceSwizzleMode => ( + "swizzle_mode", + "b32", + vec![OperandPattern::Address, OperandPattern::Immediate], + ), + _ => unreachable!("TMA tensor-map replace operation was matched"), + }; + ( + "tensormap", + vec![ + "replace".into(), + "tile".into(), + field.into(), + "global".into(), + "b1024".into(), + width.into(), + ], + operands, + ) + } else if is_fence { + let (semantics, scope) = match operation { + TmaOperation::FenceProxyTensorMapAcquireCluster => ("acquire", "cluster"), + TmaOperation::FenceProxyTensorMapAcquireCta => ("acquire", "cta"), + TmaOperation::FenceProxyTensorMapAcquireGpu => ("acquire", "gpu"), + TmaOperation::FenceProxyTensorMapAcquireSystem => ("acquire", "sys"), + TmaOperation::FenceProxyTensorMapReleaseCluster => ("release", "cluster"), + TmaOperation::FenceProxyTensorMapReleaseCta => ("release", "cta"), + TmaOperation::FenceProxyTensorMapReleaseGpu => ("release", "gpu"), + TmaOperation::FenceProxyTensorMapReleaseSystem => ("release", "sys"), + _ => unreachable!("TMA tensor-map fence operation was matched"), + }; + ( + "fence", + vec![ + "proxy".into(), + "tensormap::generic".into(), + semantics.into(), + scope.into(), + ], + if is_acquire_fence { + vec![ + OperandPattern::Address, + OperandPattern::Exact { + value: "128".into(), + }, + ] + } else { + vec![] + }, + ) + } else { + unreachable!("TMA operation category was matched") + }; + + let minimum_ptx = if blackwell_tma || operation == TmaOperation::ReplaceSwizzleAtomicity { + "8.6" + } else if is_replace || is_fence { + "8.3" + } else { + "8.0" + }; + let (minimum_sm, targets) = if blackwell_tma { + (None, TMA_BLACKWELL_TARGETS) + } else if operation == TmaOperation::ReplaceSwizzleAtomicity { + (None, BLACKWELL_LDMATRIX_LLVM_TARGETS) + } else if is_replace { + (None, TENSOR_MAP_REPLACE_TARGETS) + } else { + (Some("sm_90"), "all") + }; + let memory = if operation == TmaOperation::PrefetchTensorMap || prefetch_coordinates.is_some() { + "read" + } else if is_replace { + "write" + } else { + "read_write" + }; + let polymorphic_descriptor = operation == TmaOperation::PrefetchTensorMap; + + TmaRecipe { + operation, + abi_id, + id, + operation_key, + source_record, + llvm_symbol, + resolved_llvm_symbol: polymorphic_descriptor.then(|| format!("{llvm_symbol}.p0")), + selected_address_space: polymorphic_descriptor.then_some(ImportedAddressSpace::Generic), + llvm_mechanism: if is_replace { + BackendLoweringMechanism::InlinePtx + } else { + BackendLoweringMechanism::TypedNvvm + }, + rust_arguments, + dialect_op_type: op_type, + dialect_op_name: op_name, + dialect_operands, + llvm_arguments, + adapter, + safe, + safe_reason, + convergent, + minimum_ptx, + minimum_sm, + targets, + memory, + ptx_isa_section: if descriptor_control { + "Tensor-map descriptor and asynchronous bulk tensor operations" + } else { + "9.7.9.26.5 Asynchronous bulk tensor copy" + }, + ptx_isa_url: if descriptor_control { + "https://docs.nvidia.com/cuda/parallel-thread-execution/#tensor-map" + } else { + "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async-bulk-tensor" + }, + mnemonic, + modifiers, + operands, + summary, + } +} + +pub(in crate::resolve) fn expand_tma_admission( + admission: &TmaAdmission, +) -> Result> { + ensure!( + admission.runtime_validation == RuntimeValidation::Unexecuted, + "TMA runtime validation may be marked executed only with GPU evidence" + ); + ensure!( + !admission.llvm_evidence_profile.trim().is_empty() + && !admission.libnvvm_evidence_profile.trim().is_empty(), + "compact TMA admission requires both backend evidence profiles" + ); + ensure!( + admission + .variants + .iter() + .map(|variant| variant.operation) + .eq(TMA_OPERATIONS), + "compact TMA admission must list all {} operations in canonical order", + TMA_OPERATIONS.len() + ); + + let mut records = admission + .variants + .iter() + .map(|variant| { + let recipe = tma_recipe(variant.operation); + ensure!( + variant.abi_id == recipe.abi_id, + "{} must keep reserved ABI ID {}", + recipe.id, + recipe.abi_id + ); + Ok(OverlayIntrinsic { + id: recipe.id.into(), + abi_id: variant.abi_id.clone(), + operation_key: recipe.operation_key.into(), + family: "tma".into(), + source: None, + source_record: Some(recipe.source_record.into()), + rust_module: "tma".into(), + rust_name: recipe.id.into(), + rust_arguments: recipe + .rust_arguments + .iter() + .map(|value| (*value).into()) + .collect(), + rust_result: "()".into(), + safe: recipe.safe, + must_use: false, + safe_allowlist_reason: recipe.safe_reason.map(Into::into), + public_rust_path: format!("cuda_intrinsics::tma::{}", recipe.id), + compatibility_rust_paths: vec![format!("cuda_device::tma::{}", recipe.id)], + dialect_op_type: recipe.dialect_op_type.into(), + dialect_op_name: recipe.dialect_op_name.into(), + dialect_operands: recipe + .dialect_operands + .iter() + .map(|value| (*value).into()) + .collect(), + dialect_results: vec![], + llvm_symbol: Some(recipe.llvm_symbol.into()), + resolved_llvm_symbol: recipe.resolved_llvm_symbol.clone(), + llvm_arguments: recipe + .llvm_arguments + .iter() + .map(|value| (*value).into()) + .collect(), + llvm_results: vec![], + pure: false, + memory: recipe.memory.into(), + convergent: recipe.convergent, + execution_scope: "thread".into(), + minimum_ptx: recipe.minimum_ptx.into(), + minimum_sm: recipe.minimum_sm.map(Into::into), + ptx_result: "()".into(), + targets: recipe.targets.into(), + ptx_isa_version: "9.3".into(), + ptx_isa_section: recipe.ptx_isa_section.into(), + ptx_isa_url: recipe.ptx_isa_url.into(), + lowering: "generated_tma".into(), + backend_lowerings: vec![ + OverlayBackendLowering { + backend: IntrinsicBackend::LlvmNvptx, + mechanism: recipe.llvm_mechanism, + evidence_profile: admission.llvm_evidence_profile.clone(), + targets: None, + minimum_ptx: Some(recipe.minimum_ptx.into()), + minimum_sm: recipe.minimum_sm.map(Into::into), + }, + OverlayBackendLowering { + backend: IntrinsicBackend::LibNvvm, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: admission.libnvvm_evidence_profile.clone(), + targets: None, + minimum_ptx: Some(recipe.minimum_ptx.into()), + minimum_sm: recipe.minimum_sm.map(Into::into), + }, + ], + packed_atomic: None, + redux: None, + vote: None, + active_mask: None, + warp_match: None, + warp_barrier: None, + warp_shuffle: None, + dot_product: None, + packed_alu: None, + integer_minmax: None, + packed_conversion: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cp_async_copy: None, + cp_async_control: None, + cp_async_mbarrier: None, + mbarrier_basic: None, + movmatrix: None, + mbarrier_extended: None, + register_mma: None, + sparse_mma: None, + prmt: None, + cluster_barrier: None, + wgmma_control: None, + special_register: None, + debug_control: None, + cluster_memory: None, + clc: None, + tma: Some(Tma { + operation: recipe.operation, + reduction: None, + adapter: recipe.adapter, + runtime_validation: admission.runtime_validation, + }), + tcgen05: None, + ldmatrix_variant: None, + ldmatrix_safety: None, + ldmatrix_adapter: None, + selected_address_space: recipe.selected_address_space, + expected_ptx: InstructionPattern { + mnemonic: recipe.mnemonic.into(), + modifiers: recipe.modifiers, + operands: recipe.operands, + }, + summary: recipe.summary.into(), + }) + }) + .collect::>>()?; + + if !admission.reduce_variants.is_empty() { + ensure!( + admission.reduce_llvm_evidence_profile.is_some(), + "compact TMA reduction admission requires reduce_llvm_evidence_profile" + ); + ensure!( + admission.reduce_libnvvm_evidence_profile.is_some(), + "compact TMA reduction admission requires reduce_libnvvm_evidence_profile" + ); + + let expected_reductions = tma_reduction_matrix(); + ensure!( + admission.reduce_variants.len() == expected_reductions.len() + && admission + .reduce_variants + .iter() + .zip(expected_reductions.iter()) + .all(|(variant, expected)| { + variant.operation == expected.operation + && variant.load_mode == expected.load_mode + && variant.dimensions == expected.dimensions + }), + "compact TMA reduction admission must list all 64 operations in canonical order" + ); + for variant in &admission.reduce_variants { + records.push(expand_tma_reduction_variant(admission, variant)?); + } + } + Ok(records) +} + +pub(in crate::resolve) fn validate_tma_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, +) -> Result<()> { + let tma = policy + .tma + .as_ref() + .with_context(|| format!("{} has no closed TMA contract", policy.id))?; + if tma.operation == TmaOperation::Reduce { + return validate_tma_reduction_policy(policy, declaration, tma); + } + ensure!( + tma.reduction.is_none(), + "{} non-reduction TMA operation carries a reduction contract", + policy.id + ); + let recipe = tma_recipe(tma.operation); + ensure!( + policy.id == recipe.id + && policy.abi_id == recipe.abi_id + && policy.operation_key == recipe.operation_key + && policy.source.is_none() + && policy.source_record.as_deref() == Some(recipe.source_record) + && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) + && policy.resolved_llvm_symbol == recipe.resolved_llvm_symbol + && declaration.source_record == recipe.source_record + && declaration.llvm_name == recipe.llvm_symbol, + "{} TMA identity changed", + policy.id + ); + ensure!( + policy.rust_module == "tma" + && policy.rust_name == recipe.id + && policy.rust_arguments == recipe.rust_arguments + && policy.rust_result == "()" + && policy.safe == recipe.safe + && !policy.must_use + && policy.safe_allowlist_reason.as_deref() == recipe.safe_reason + && policy.public_rust_path == format!("cuda_intrinsics::tma::{}", recipe.id) + && policy.compatibility_rust_paths == [format!("cuda_device::tma::{}", recipe.id)], + "{} TMA Rust API changed", + policy.id + ); + ensure!( + policy.dialect_op_type == recipe.dialect_op_type + && policy.dialect_op_name == recipe.dialect_op_name + && policy.dialect_operands == recipe.dialect_operands + && policy.dialect_results.is_empty() + && policy.llvm_arguments == recipe.llvm_arguments + && policy.llvm_results.is_empty() + && declaration.arguments == recipe.llvm_arguments + && declaration.results.is_empty() + && policy.selected_address_space == recipe.selected_address_space + && policy.lowering == "generated_tma", + "{} TMA carrier or LLVM adapter changed", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == recipe.memory + && policy.convergent == recipe.convergent + && policy.execution_scope == "thread" + && tma.adapter == recipe.adapter + && tma.runtime_validation == RuntimeValidation::Unexecuted, + "{} TMA semantics changed", + policy.id + ); + ensure!( + policy.minimum_ptx == recipe.minimum_ptx + && policy.minimum_sm.as_deref() == recipe.minimum_sm + && policy.targets == recipe.targets + && policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section == recipe.ptx_isa_section + && policy.ptx_isa_url == recipe.ptx_isa_url + && policy.ptx_result == "()" + && policy.expected_ptx.mnemonic == recipe.mnemonic + && policy.expected_ptx.modifiers == recipe.modifiers + && policy.expected_ptx.operands == recipe.operands, + "{} TMA target or PTX contract changed", + policy.id + ); + let valid_route = |backend, mechanism| { + policy.backend_lowerings.iter().any(|route| { + route.backend == backend + && route.mechanism == mechanism + && route.minimum_ptx.as_deref() == Some(recipe.minimum_ptx) + && route.minimum_sm.as_deref() == recipe.minimum_sm + && !route.evidence_profile.trim().is_empty() + }) + }; + ensure!( + policy.backend_lowerings.len() == 2 + && valid_route(IntrinsicBackend::LlvmNvptx, recipe.llvm_mechanism,) + && valid_route( + IntrinsicBackend::LibNvvm, + BackendLoweringMechanism::InlinePtx, + ), + "{} TMA backend route changed", + policy.id + ); + let expected_properties = tma_imported_properties(tma.operation); + ensure!( + declaration.properties == expected_properties, + "{} imported TMA declaration changed", + policy.id + ); + ensure_no_other_family_contract(policy, "TMA")?; + Ok(()) +} + +pub(in crate::resolve) fn validate_tma_reduction_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, + tma: &Tma, +) -> Result<()> { + ensure!( + tma.operation == TmaOperation::Reduce, + "{} is not a TMA reduction", + policy.id + ); + let reduction = tma + .reduction + .with_context(|| format!("{} has no TMA reduction contract", policy.id))?; + let recipe = tma_reduction_recipe(reduction)?; + ensure!( + policy.id == recipe.id + && policy.operation_key == recipe.operation_key + && policy.source.is_none() + && policy.source_record.as_deref() == Some(recipe.source_record.as_str()) + && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol.as_str()) + && policy.resolved_llvm_symbol.is_none() + && declaration.source_record == recipe.source_record + && declaration.llvm_name == recipe.llvm_symbol, + "{} TMA reduction identity changed", + policy.id + ); + ensure!( + policy.rust_module == "tma" + && policy.rust_name == recipe.id + && policy.rust_arguments == recipe.rust_arguments + && policy.rust_result == "()" + && !policy.safe + && !policy.must_use + && policy.safe_allowlist_reason.is_none() + && policy.public_rust_path == format!("cuda_intrinsics::tma::{}", recipe.id) + && policy.compatibility_rust_paths == [format!("cuda_device::tma::{}", recipe.id)], + "{} TMA reduction Rust API changed", + policy.id + ); + ensure!( + policy.dialect_op_type == recipe.dialect_op_type + && policy.dialect_op_name == recipe.dialect_op_name + && policy.dialect_operands == recipe.dialect_operands + && policy.dialect_results.is_empty() + && policy.llvm_arguments == recipe.llvm_arguments + && policy.llvm_results.is_empty() + && declaration.arguments == recipe.llvm_arguments + && declaration.results.is_empty() + && declaration.selections.is_empty() + && policy.lowering == "generated_tma", + "{} TMA reduction carrier or LLVM adapter changed", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == "read_write" + && policy.convergent + && policy.execution_scope == "thread" + && tma.adapter == TmaAdapter::ReductionPointersCoordinatesInjectDefaults + && tma.runtime_validation == RuntimeValidation::Unexecuted, + "{} TMA reduction semantics changed", + policy.id + ); + ensure!( + policy.minimum_ptx == "8.0" + && policy.minimum_sm.as_deref() == Some("sm_90") + && policy.targets == "all" + && policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section + == "9.7.9.26.5.3 Data Movement and Conversion Instructions: cp.reduce.async.bulk.tensor" + && policy.ptx_isa_url + == "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-reduce-async-bulk-tensor" + && policy.ptx_result == "()" + && policy.expected_ptx.mnemonic == "cp" + && policy.expected_ptx.modifiers == recipe.modifiers + && policy.expected_ptx.operands == recipe.operands, + "{} TMA reduction target or PTX contract changed", + policy.id + ); + let valid_route = |backend, mechanism| { + policy.backend_lowerings.iter().any(|route| { + route.backend == backend + && route.mechanism == mechanism + && route.minimum_ptx.as_deref() == Some("8.0") + && route.minimum_sm.as_deref() == Some("sm_90") + && !route.evidence_profile.trim().is_empty() + }) + }; + ensure!( + policy.backend_lowerings.len() == 2 + && valid_route( + IntrinsicBackend::LlvmNvptx, + BackendLoweringMechanism::TypedNvvm, + ) + && valid_route( + IntrinsicBackend::LibNvvm, + BackendLoweringMechanism::InlinePtx, + ), + "{} TMA reduction backend route changed", + policy.id + ); + ensure!( + declaration.properties == tma_reduction_imported_properties(reduction.dimensions as usize), + "{} imported TMA reduction declaration changed", + policy.id + ); + ensure_no_other_family_contract(policy, "TMA")?; + Ok(()) +} + +pub(in crate::resolve) fn tma_imported_properties(operation: TmaOperation) -> Vec { + assert_ne!(operation, TmaOperation::Reduce); + let dimensions = operation.dimensions(); + if matches!( + operation, + TmaOperation::G2sTile1d + | TmaOperation::G2sTile2d + | TmaOperation::G2sTile2dMulticast + | TmaOperation::G2sTile2dMulticastCg2 + | TmaOperation::G2sTile3d + | TmaOperation::G2sTile4d + | TmaOperation::G2sTile5d + ) { + let dimensions = dimensions.unwrap(); + let mut properties = vec![ + format!("ImmArg", dimensions + 5), + format!("ImmArg", dimensions + 6), + format!("ImmArg", dimensions + 7), + "IntrConvergent".into(), + format!("Range", dimensions + 7), + "ReadOnly".into(), + "WriteOnly".into(), + ]; + properties.sort(); + return properties; + } + if matches!( + operation, + TmaOperation::S2gTile1d + | TmaOperation::S2gTile2d + | TmaOperation::S2gTile3d + | TmaOperation::S2gTile4d + | TmaOperation::S2gTile5d + ) { + return vec![ + format!("ImmArg", dimensions.unwrap() + 3), + "IntrConvergent".into(), + "ReadOnly".into(), + "ReadOnly".into(), + ]; + } + if operation == TmaOperation::CommitGroup { + return vec![]; + } + if matches!( + operation, + TmaOperation::WaitGroup | TmaOperation::WaitGroupRead + ) { + return vec!["ImmArg".into()]; + } + if operation == TmaOperation::PrefetchTensorMap { + return vec![ + "IntrArgMemOnly".into(), + "NoCapture".into(), + "ReadOnly".into(), + ]; + } + if let Some(coordinate_count) = operation.prefetch_coordinate_count() { + return vec![ + format!("ImmArg", coordinate_count + 2), + "IntrConvergent".into(), + "ReadOnly".into(), + ]; + } + if matches!( + operation, + TmaOperation::FenceProxyTensorMapAcquireCluster + | TmaOperation::FenceProxyTensorMapAcquireCta + | TmaOperation::FenceProxyTensorMapAcquireGpu + | TmaOperation::FenceProxyTensorMapAcquireSystem + ) { + return vec![ + "ImmArg".into(), + "IntrArgMemOnly".into(), + "IntrNoCallback".into(), + "Range".into(), + ]; + } + if matches!( + operation, + TmaOperation::FenceProxyTensorMapReleaseCluster + | TmaOperation::FenceProxyTensorMapReleaseCta + | TmaOperation::FenceProxyTensorMapReleaseGpu + | TmaOperation::FenceProxyTensorMapReleaseSystem + ) { + return vec!["IntrNoCallback".into()]; + } + + let mut properties = vec![ + "IntrArgMemOnly".into(), + "IntrWriteMem".into(), + "NoCapture".into(), + ]; + match operation { + TmaOperation::ReplaceBoxDim + | TmaOperation::ReplaceElementStride + | TmaOperation::ReplaceGlobalDim + | TmaOperation::ReplaceGlobalStride => { + properties.push("ImmArg".into()); + } + TmaOperation::ReplaceElementType => { + properties.extend([ + "ArgInfo".into(), + "ImmArg".into(), + "Range".into(), + ]); + } + TmaOperation::ReplaceFillMode => { + properties.extend([ + "ArgInfo".into(), + "ImmArg".into(), + "Range".into(), + ]); + } + TmaOperation::ReplaceInterleaveLayout => { + properties.extend([ + "ArgInfo".into(), + "ImmArg".into(), + "Range".into(), + ]); + } + TmaOperation::ReplaceSwizzleAtomicity => { + properties.extend([ + "ArgInfo".into(), + "ImmArg".into(), + "Range".into(), + ]); + } + TmaOperation::ReplaceSwizzleMode => { + properties.extend([ + "ArgInfo".into(), + "ImmArg".into(), + "Range".into(), + ]); + } + TmaOperation::ReplaceGlobalAddress | TmaOperation::ReplaceRank => {} + _ => unreachable!("closed TMA operation property contract"), + } + properties.sort(); + properties +} + +pub(in crate::resolve) fn tma_reduction_imported_properties(dimensions: usize) -> Vec { + vec![ + format!("ImmArg", dimensions + 3), + "IntrConvergent".into(), + "ReadOnly".into(), + "ReadOnly".into(), + ] +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/families/warp.rs b/crates/cuda-intrinsics-gen/src/resolve/families/warp.rs new file mode 100644 index 0000000000..e9292c33b9 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/families/warp.rs @@ -0,0 +1,1702 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + ActiveMaskAdapter, ActiveMaskObservation, BackendLoweringMechanism, ImportedIntrinsic, + IntrinsicBackend, IntrinsicSource, MaskEncoding, MatchOperandEncoding, OverlayIntrinsic, + PreSm70MemberMaskRule, VoteAdapter, VoteMode, VoteParticipation, WarpBarrierAdapter, + WarpBarrierMaskEncoding, WarpBarrierMemoryOrdering, WarpBarrierParticipation, WarpMatchAdapter, + WarpMatchMode, WarpMatchParticipation, WarpMatchValueWidth, WarpShuffleAdapter, + WarpShuffleMode, WarpShuffleOperandEncoding, WarpShuffleParticipation, WarpShuffleSourceLane, + WarpShuffleValueKind, +}; +use crate::ptx::OperandPattern; +use anyhow::{Context, Result, ensure}; +use std::collections::BTreeSet; + +use super::*; + +pub(in crate::resolve) fn validate_sync_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, +) -> Result<()> { + if let Some(scope) = threadfence_scope_for_id(&policy.id) { + return validate_threadfence_policy(policy, declaration, scope); + } + ensure!( + policy.id == "sync_threads" + && policy.abi_id == "i0034" + && policy.operation_key == "synchronization.cta.barrier.aligned.all" + && policy.source.is_none() + && policy.source_record.as_deref() == Some("int_nvvm_barrier_cta_sync_aligned_all") + && policy.llvm_symbol.as_deref() == Some("llvm.nvvm.barrier.cta.sync.aligned.all") + && policy.resolved_llvm_symbol.is_none(), + "{} sync identity does not match the closed sync_threads recipe", + policy.id + ); + ensure!( + policy.rust_module == "thread" + && policy.rust_name == "sync_threads" + && policy.rust_arguments.is_empty() + && policy.rust_result == "()" + && !policy.safe + && !policy.must_use + && policy.safe_allowlist_reason.is_none() + && policy.public_rust_path == "cuda_intrinsics::thread::sync_threads" + && policy.compatibility_rust_paths + == [ + "cuda_device::thread::sync_threads", + "cuda_device::sync_threads", + ], + "{} must preserve the unsafe sync_threads raw API and both cuda-device compatibility paths", + policy.id + ); + ensure!( + policy.dialect_op_type == "Barrier0Op" + && policy.dialect_op_name == "nvvm.barrier0" + && policy.dialect_operands.is_empty() + && policy.dialect_results.is_empty() + && policy.llvm_arguments == ["i32"] + && policy.llvm_results.is_empty() + && policy.lowering == "generated_sync_threads", + "{} is outside the fixed-zero sync_threads lowering recipe", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == "read_write" + && policy.convergent + && policy.execution_scope == "cta" + && policy.minimum_ptx == "1.0" + && policy.minimum_sm.is_none() + && policy.ptx_result == "()" + && policy.targets == "all", + "{} sync effects or native target floor disagree with the closed recipe", + policy.id + ); + ensure!( + policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section + == "9.7.14.1 Parallel Synchronization and Communication Instructions: bar, barrier" + && policy.ptx_isa_url + == "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-bar-barrier", + "{} sync PTX provenance disagrees with the reviewed recipe", + policy.id + ); + ensure!( + declaration.properties == ["IntrConvergent", "IntrNoCallback"], + "{} sync properties disagree with the imported LLVM declaration", + policy.id + ); + ensure!( + policy.ldmatrix_variant.is_none() + && policy.ldmatrix_safety.is_none() + && policy.ldmatrix_adapter.is_none() + && policy.packed_atomic.is_none() + && policy.redux.is_none() + && policy.vote.is_none() + && policy.active_mask.is_none() + && policy.warp_match.is_none() + && policy.warp_barrier.is_none() + && policy.warp_shuffle.is_none() + && policy.dot_product.is_none() + && policy.packed_alu.is_none() + && policy.packed_conversion.is_none() + && policy.cp_async_copy.is_none() + && policy.cp_async_control.is_none() + && policy.mbarrier_basic.is_none() + && policy.selected_address_space.is_none(), + "{} mixes another generated-family contract with sync", + policy.id + ); + ensure!( + policy.expected_ptx.mnemonic == "bar" + && policy.expected_ptx.modifiers == ["sync"] + && policy.expected_ptx.operands == [OperandPattern::Exact { value: "0".into() }], + "{} expected PTX does not match literal bar.sync 0", + policy.id + ); + + let backend_pairs: BTreeSet<_> = policy + .backend_lowerings + .iter() + .map(|lowering| (lowering.backend, lowering.mechanism)) + .collect(); + ensure!( + policy.backend_lowerings.len() == 2 + && backend_pairs + == BTreeSet::from([ + ( + IntrinsicBackend::LlvmNvptx, + BackendLoweringMechanism::TypedNvvm, + ), + ( + IntrinsicBackend::LibNvvm, + BackendLoweringMechanism::InlinePtx, + ), + ]), + "{} must define exactly the reviewed LLVM typed and libNVVM inline-PTX routes", + policy.id + ); + for lowering in &policy.backend_lowerings { + let floor_matches = match lowering.backend { + IntrinsicBackend::LlvmNvptx => { + lowering.mechanism == BackendLoweringMechanism::TypedNvvm + && lowering.minimum_ptx.as_deref() == Some("3.2") + && lowering.minimum_sm.as_deref() == Some("sm_20") + } + IntrinsicBackend::LibNvvm => { + lowering.mechanism == BackendLoweringMechanism::InlinePtx + && lowering.minimum_ptx.is_none() + && lowering.minimum_sm.as_deref() == Some("sm_75") + } + }; + ensure!( + floor_matches && !lowering.evidence_profile.trim().is_empty(), + "{} backend {:?} does not carry its reviewed sync profile floor", + policy.id, + lowering.backend + ); + } + Ok(()) +} + +pub(in crate::resolve) fn validate_vote_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, +) -> Result<()> { + let vote = policy + .vote + .as_ref() + .with_context(|| format!("{} has no closed vote contract", policy.id))?; + let recipe = vote_recipe(vote.mode); + ensure!( + vote.participation + == VoteParticipation::ExecutingLaneNamedAllNamedLanesSameInstructionAndMask + && vote.legacy_pre_sm70 + == PreSm70MemberMaskRule::AllNamedLanesConvergedAndOnlyNamedLanesActive + && vote.adapter == VoteAdapter::DirectMaskPredicate + && vote.mask_encoding == MaskEncoding::RegisterOrImmediate, + "{} requests an unsupported vote participation, pre-sm70 rule, adapter, or mask encoding", + policy.id + ); + ensure!( + policy.id == recipe.id + && policy.abi_id == recipe.abi_id + && policy.operation_key == recipe.operation_key + && policy.source.is_none() + && policy.source_record.as_deref() == Some(recipe.source_record) + && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) + && policy.resolved_llvm_symbol.is_none(), + "{} vote identity does not match its closed mode recipe", + policy.id + ); + let expected_compatibility_paths: Vec = if recipe.has_compatibility_path { + vec![format!("cuda_device::warp::{}", recipe.rust_name)] + } else { + vec![] + }; + ensure!( + policy.rust_module == "warp" + && policy.rust_name == recipe.rust_name + && policy.rust_arguments == ["u32", "bool"] + && policy.rust_result == recipe.rust_result + && !policy.safe + && policy.must_use + && policy.safe_allowlist_reason.is_none() + && policy.public_rust_path == format!("cuda_intrinsics::warp::{}", recipe.rust_name) + && policy.compatibility_rust_paths == expected_compatibility_paths, + "{} must preserve its unsafe must-use vote raw API and reviewed compatibility path", + policy.id + ); + ensure!( + policy.dialect_op_type == recipe.dialect_op_type + && policy.dialect_op_name == recipe.dialect_op_name + && policy.dialect_operands == ["i32", "i1"] + && policy.dialect_results == [recipe.llvm_result] + && policy.llvm_arguments == ["i32", "i1"] + && policy.llvm_results == [recipe.llvm_result] + && policy.lowering == "generated_vote", + "{} is outside the closed two-operand vote lowering recipe", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == "inaccessible_read_write" + && policy.convergent + && policy.execution_scope == "warp" + && policy.minimum_ptx == "6.0" + && policy.minimum_sm.as_deref() == Some("sm_30") + && policy.ptx_result == recipe.rust_result + && policy.targets == "all", + "{} vote effects, carrier, or target floor disagree with its mode recipe", + policy.id + ); + ensure!( + policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section == "9.7.14.10 Warp Vote Instructions: vote.sync" + && policy.ptx_isa_url + == "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-vote-sync", + "{} vote PTX provenance disagrees with the reviewed recipe", + policy.id + ); + ensure!( + declaration.properties + == [ + "IntrConvergent", + "IntrInaccessibleMemOnly", + "IntrNoCallback", + ], + "{} vote memory and convergence effects disagree with the imported declaration", + policy.id + ); + ensure!( + policy.backend_lowerings.is_empty() + && policy.packed_atomic.is_none() + && policy.ldmatrix_variant.is_none() + && policy.ldmatrix_safety.is_none() + && policy.ldmatrix_adapter.is_none() + && policy.redux.is_none() + && policy.active_mask.is_none() + && policy.warp_match.is_none() + && policy.warp_barrier.is_none() + && policy.warp_shuffle.is_none() + && policy.dot_product.is_none() + && policy.packed_alu.is_none() + && policy.packed_conversion.is_none() + && policy.cp_async_copy.is_none() + && policy.cp_async_control.is_none() + && policy.mbarrier_basic.is_none() + && policy.selected_address_space.is_none(), + "{} mixes another generated-family contract with vote", + policy.id + ); + ensure!( + policy.expected_ptx.mnemonic == "vote" + && policy.expected_ptx.modifiers == ["sync", recipe.ptx_mode, recipe.ptx_type] + && policy.expected_ptx.operands + == [ + OperandPattern::Register, + OperandPattern::Register, + OperandPattern::RegisterOrImmediate, + ], + "{} expected PTX does not match its closed vote mode recipe", + policy.id + ); + + let expected_selection_records = + BTreeSet::from([recipe.immediate_selection, recipe.register_selection]); + let actual_selection_records: BTreeSet<_> = declaration + .selections + .iter() + .map(|selection| selection.source_record.as_str()) + .collect(); + ensure!( + declaration.selections.len() == 2 && actual_selection_records == expected_selection_records, + "{} vote declaration must contain exactly its immediate/register selection pair", + policy.id + ); + let expected_asm = format!( + "vote.sync.{}.{} \t$dest, $pred, $mask;", + recipe.ptx_mode, recipe.ptx_type + ); + for selection in &declaration.selections { + ensure!( + selection.asm == expected_asm + && selection.predicates + == [ + "Subtarget->getPTXVersion() >= 60", + "Subtarget->getSmVersion() >= 30", + ] + && selection.constraints.is_empty(), + "{} vote immediate/register selections disagree on PTX shape, target predicates, or constraints", + policy.id + ); + } + Ok(()) +} + +pub(in crate::resolve) struct VoteRecipe { + pub(in crate::resolve) id: &'static str, + pub(in crate::resolve) abi_id: &'static str, + pub(in crate::resolve) operation_key: &'static str, + pub(in crate::resolve) source_record: &'static str, + pub(in crate::resolve) llvm_symbol: &'static str, + pub(in crate::resolve) rust_name: &'static str, + pub(in crate::resolve) rust_result: &'static str, + pub(in crate::resolve) dialect_op_type: &'static str, + pub(in crate::resolve) dialect_op_name: &'static str, + pub(in crate::resolve) llvm_result: &'static str, + pub(in crate::resolve) ptx_mode: &'static str, + pub(in crate::resolve) ptx_type: &'static str, + pub(in crate::resolve) immediate_selection: &'static str, + pub(in crate::resolve) register_selection: &'static str, + pub(in crate::resolve) has_compatibility_path: bool, +} + +pub(in crate::resolve) fn vote_recipe(mode: VoteMode) -> VoteRecipe { + match mode { + VoteMode::All => VoteRecipe { + id: "all_sync", + abi_id: "i0040", + operation_key: "warp.vote.sync.all.pred", + source_record: "int_nvvm_vote_all_sync", + llvm_symbol: "llvm.nvvm.vote.all.sync", + rust_name: "all_sync", + rust_result: "bool", + dialect_op_type: "VoteSyncAllOp", + dialect_op_name: "nvvm.vote_sync_all", + llvm_result: "i1", + ptx_mode: "all", + ptx_type: "pred", + immediate_selection: "VOTE_SYNC_ALLi", + register_selection: "VOTE_SYNC_ALLr", + has_compatibility_path: true, + }, + VoteMode::Any => VoteRecipe { + id: "any_sync", + abi_id: "i0041", + operation_key: "warp.vote.sync.any.pred", + source_record: "int_nvvm_vote_any_sync", + llvm_symbol: "llvm.nvvm.vote.any.sync", + rust_name: "any_sync", + rust_result: "bool", + dialect_op_type: "VoteSyncAnyOp", + dialect_op_name: "nvvm.vote_sync_any", + llvm_result: "i1", + ptx_mode: "any", + ptx_type: "pred", + immediate_selection: "VOTE_SYNC_ANYi", + register_selection: "VOTE_SYNC_ANYr", + has_compatibility_path: true, + }, + VoteMode::Ballot => VoteRecipe { + id: "ballot_sync", + abi_id: "i0042", + operation_key: "warp.vote.sync.ballot.b32", + source_record: "int_nvvm_vote_ballot_sync", + llvm_symbol: "llvm.nvvm.vote.ballot.sync", + rust_name: "ballot_sync", + rust_result: "u32", + dialect_op_type: "VoteSyncBallotOp", + dialect_op_name: "nvvm.vote_sync_ballot", + llvm_result: "i32", + ptx_mode: "ballot", + ptx_type: "b32", + immediate_selection: "VOTE_SYNC_BALLOTi", + register_selection: "VOTE_SYNC_BALLOTr", + has_compatibility_path: true, + }, + VoteMode::Uni => VoteRecipe { + id: "uni_sync", + abi_id: "i0043", + operation_key: "warp.vote.sync.uni.pred", + source_record: "int_nvvm_vote_uni_sync", + llvm_symbol: "llvm.nvvm.vote.uni.sync", + rust_name: "uni_sync", + rust_result: "bool", + dialect_op_type: "VoteSyncUniOp", + dialect_op_name: "nvvm.vote_sync_uni", + llvm_result: "i1", + ptx_mode: "uni", + ptx_type: "pred", + immediate_selection: "VOTE_SYNC_UNIi", + register_selection: "VOTE_SYNC_UNIr", + has_compatibility_path: false, + }, + } +} + +pub(in crate::resolve) fn validate_active_mask_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, +) -> Result<()> { + let active_mask = policy + .active_mask + .as_ref() + .with_context(|| format!("{} has no closed active-mask contract", policy.id))?; + ensure!( + active_mask.observation == ActiveMaskObservation::ExecutingLanesAtInstruction + && active_mask.adapter == ActiveMaskAdapter::DirectZeroOperandMask, + "{} requests an unsupported active-mask observation or adapter", + policy.id + ); + ensure!( + policy.id == "active_mask" + && policy.abi_id == "i0044" + && policy.operation_key == "warp.active_mask" + && policy.source.is_none() + && policy.source_record.as_deref() == Some("int_nvvm_activemask") + && policy.llvm_symbol.as_deref() == Some("llvm.nvvm.activemask") + && policy.resolved_llvm_symbol.is_none(), + "{} active-mask identity does not match the closed recipe", + policy.id + ); + ensure!( + policy.rust_module == "warp" + && policy.rust_name == "active_mask" + && policy.rust_arguments.is_empty() + && policy.rust_result == "u32" + && policy.safe + && policy.must_use + && policy + .safe_allowlist_reason + .as_deref() + .is_some_and(|reason| !reason.is_empty()) + && policy.public_rust_path == "cuda_intrinsics::warp::active_mask" + && policy.compatibility_rust_paths == ["cuda_device::warp::active_mask"], + "{} must preserve its safe must-use raw and compatibility APIs", + policy.id + ); + ensure!( + policy.dialect_op_type == "ActiveMaskOp" + && policy.dialect_op_name == "nvvm.activemask" + && policy.dialect_operands.is_empty() + && policy.dialect_results == ["i32"] + && policy.llvm_arguments.is_empty() + && policy.llvm_results == ["i32"] + && policy.lowering == "generated_active_mask", + "{} is outside the closed zero-operand active-mask lowering recipe", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == "inaccessible_read_write" + && policy.convergent + && policy.execution_scope == "warp" + && policy.minimum_ptx == "6.2" + && policy.minimum_sm.as_deref() == Some("sm_30") + && policy.ptx_result == "u32" + && policy.targets == "all", + "{} active-mask effects or target floor disagree with the closed recipe", + policy.id + ); + ensure!( + policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section + == "9.7.14.12 Parallel Synchronization and Communication Instructions: activemask" + && policy.ptx_isa_url + == "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-activemask", + "{} active-mask PTX provenance disagrees with the reviewed recipe", + policy.id + ); + ensure!( + declaration.properties + == [ + "IntrConvergent", + "IntrHasSideEffects", + "IntrInaccessibleMemOnly", + "IntrNoCallback", + ] + && declaration.selections.len() == 1 + && declaration.selections[0].source_record == "ACTIVEMASK" + && declaration.selections[0].asm == "activemask.b32 \t$dest;" + && declaration.selections[0].predicates + == [ + "Subtarget->getSmVersion() >= 30", + "Subtarget->getPTXVersion() >= 62", + ] + && declaration.selections[0].constraints.is_empty(), + "{} active-mask declaration or selection facts changed", + policy.id + ); + ensure!( + policy.expected_ptx.mnemonic == "activemask" + && policy.expected_ptx.modifiers == ["b32"] + && policy.expected_ptx.operands == [OperandPattern::Register], + "{} expected PTX does not match activemask.b32", + policy.id + ); + ensure!( + policy.packed_atomic.is_none() + && policy.redux.is_none() + && policy.vote.is_none() + && policy.warp_match.is_none() + && policy.warp_barrier.is_none() + && policy.warp_shuffle.is_none() + && policy.dot_product.is_none() + && policy.packed_alu.is_none() + && policy.packed_conversion.is_none() + && policy.cp_async_copy.is_none() + && policy.cp_async_control.is_none() + && policy.mbarrier_basic.is_none() + && policy.ldmatrix_variant.is_none() + && policy.ldmatrix_safety.is_none() + && policy.ldmatrix_adapter.is_none() + && policy.selected_address_space.is_none(), + "{} mixes another generated-family contract with active_mask", + policy.id + ); + + let backend_pairs: BTreeSet<_> = policy + .backend_lowerings + .iter() + .map(|lowering| (lowering.backend, lowering.mechanism)) + .collect(); + ensure!( + policy.backend_lowerings.len() == 2 + && backend_pairs + == BTreeSet::from([ + ( + IntrinsicBackend::LlvmNvptx, + BackendLoweringMechanism::TypedNvvm, + ), + ( + IntrinsicBackend::LibNvvm, + BackendLoweringMechanism::InlinePtx, + ), + ]), + "{} must keep the LLVM typed and libNVVM inline-PTX routes explicit", + policy.id + ); + for lowering in &policy.backend_lowerings { + let floor_matches = match lowering.backend { + IntrinsicBackend::LlvmNvptx => { + lowering.mechanism == BackendLoweringMechanism::TypedNvvm + && lowering.minimum_ptx.as_deref() == Some("6.2") + && lowering.minimum_sm.as_deref() == Some("sm_30") + } + IntrinsicBackend::LibNvvm => { + lowering.mechanism == BackendLoweringMechanism::InlinePtx + && lowering.minimum_ptx.is_none() + && lowering.minimum_sm.as_deref() == Some("sm_75") + } + }; + ensure!( + floor_matches && !lowering.evidence_profile.trim().is_empty(), + "{} backend {:?} does not carry its reviewed active-mask floor", + policy.id, + lowering.backend + ); + } + Ok(()) +} + +pub(in crate::resolve) struct WarpMatchRecipe { + pub(in crate::resolve) id: &'static str, + pub(in crate::resolve) abi_id: &'static str, + pub(in crate::resolve) operation_key: &'static str, + pub(in crate::resolve) source_record: &'static str, + pub(in crate::resolve) llvm_symbol: &'static str, + pub(in crate::resolve) rust_name: &'static str, + pub(in crate::resolve) rust_value: &'static str, + pub(in crate::resolve) llvm_value: &'static str, + pub(in crate::resolve) dialect_op_type: &'static str, + pub(in crate::resolve) dialect_op_name: &'static str, + pub(in crate::resolve) ptx_mode: &'static str, + pub(in crate::resolve) ptx_type: &'static str, + pub(in crate::resolve) selections: [&'static str; 4], + pub(in crate::resolve) adapter: WarpMatchAdapter, +} + +pub(in crate::resolve) fn warp_match_recipe( + mode: WarpMatchMode, + width: WarpMatchValueWidth, +) -> WarpMatchRecipe { + match (mode, width) { + (WarpMatchMode::Any, WarpMatchValueWidth::B32) => WarpMatchRecipe { + id: "match_any_sync", + abi_id: "i0045", + operation_key: "warp.match.sync.any.b32", + source_record: "int_nvvm_match_any_sync_i32", + llvm_symbol: "llvm.nvvm.match.any.sync.i32", + rust_name: "match_any_sync", + rust_value: "u32", + llvm_value: "i32", + dialect_op_type: "MatchAnySyncI32Op", + dialect_op_name: "nvvm.match_any_sync_i32", + ptx_mode: "any", + ptx_type: "b32", + selections: [ + "MATCH_ANY_SYNC_32ii", + "MATCH_ANY_SYNC_32ir", + "MATCH_ANY_SYNC_32ri", + "MATCH_ANY_SYNC_32rr", + ], + adapter: WarpMatchAdapter::DirectMask, + }, + (WarpMatchMode::Any, WarpMatchValueWidth::B64) => WarpMatchRecipe { + id: "match_any_i64_sync", + abi_id: "i0046", + operation_key: "warp.match.sync.any.b64", + source_record: "int_nvvm_match_any_sync_i64", + llvm_symbol: "llvm.nvvm.match.any.sync.i64", + rust_name: "match_any_i64_sync", + rust_value: "u64", + llvm_value: "i64", + dialect_op_type: "MatchAnySyncI64Op", + dialect_op_name: "nvvm.match_any_sync_i64", + ptx_mode: "any", + ptx_type: "b64", + selections: [ + "MATCH_ANY_SYNC_64ii", + "MATCH_ANY_SYNC_64ir", + "MATCH_ANY_SYNC_64ri", + "MATCH_ANY_SYNC_64rr", + ], + adapter: WarpMatchAdapter::DirectMask, + }, + (WarpMatchMode::All, WarpMatchValueWidth::B32) => WarpMatchRecipe { + id: "match_all_sync", + abi_id: "i0047", + operation_key: "warp.match.sync.all.b32", + source_record: "int_nvvm_match_all_sync_i32p", + llvm_symbol: "llvm.nvvm.match.all.sync.i32p", + rust_name: "match_all_sync", + rust_value: "u32", + llvm_value: "i32", + dialect_op_type: "MatchAllSyncI32Op", + dialect_op_name: "nvvm.match_all_sync_i32", + ptx_mode: "all", + ptx_type: "b32", + selections: [ + "MATCH_ALLP_SYNC_32ii", + "MATCH_ALLP_SYNC_32ir", + "MATCH_ALLP_SYNC_32ri", + "MATCH_ALLP_SYNC_32rr", + ], + adapter: WarpMatchAdapter::ProjectMaskDiscardPredicate, + }, + (WarpMatchMode::All, WarpMatchValueWidth::B64) => WarpMatchRecipe { + id: "match_all_i64_sync", + abi_id: "i0048", + operation_key: "warp.match.sync.all.b64", + source_record: "int_nvvm_match_all_sync_i64p", + llvm_symbol: "llvm.nvvm.match.all.sync.i64p", + rust_name: "match_all_i64_sync", + rust_value: "u64", + llvm_value: "i64", + dialect_op_type: "MatchAllSyncI64Op", + dialect_op_name: "nvvm.match_all_sync_i64", + ptx_mode: "all", + ptx_type: "b64", + selections: [ + "MATCH_ALLP_SYNC_64ii", + "MATCH_ALLP_SYNC_64ir", + "MATCH_ALLP_SYNC_64ri", + "MATCH_ALLP_SYNC_64rr", + ], + adapter: WarpMatchAdapter::ProjectMaskDiscardPredicate, + }, + } +} + +pub(in crate::resolve) fn validate_warp_match_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, +) -> Result<()> { + let warp_match = policy + .warp_match + .as_ref() + .with_context(|| format!("{} has no closed warp-match contract", policy.id))?; + let recipe = warp_match_recipe(warp_match.mode, warp_match.value_width); + ensure!( + warp_match.participation + == WarpMatchParticipation::ExecutingLaneNamedAllNamedLanesSameInstructionAndMask + && warp_match.adapter == recipe.adapter + && warp_match.value_encoding == MatchOperandEncoding::RegisterOrImmediate + && warp_match.mask_encoding == MatchOperandEncoding::RegisterOrImmediate, + "{} requests an unsupported warp-match participation, adapter, or encoding", + policy.id + ); + ensure!( + policy.id == recipe.id + && policy.abi_id == recipe.abi_id + && policy.operation_key == recipe.operation_key + && policy.source.is_none() + && policy.source_record.as_deref() == Some(recipe.source_record) + && policy.llvm_symbol.as_deref() == Some(recipe.llvm_symbol) + && policy.resolved_llvm_symbol.is_none(), + "{} warp-match identity does not match its closed recipe", + policy.id + ); + ensure!( + policy.rust_module == "warp" + && policy.rust_name == recipe.rust_name + && policy.rust_arguments == ["u32", recipe.rust_value] + && policy.rust_result == "u32" + && !policy.safe + && policy.must_use + && policy.safe_allowlist_reason.is_none() + && policy.public_rust_path == format!("cuda_intrinsics::warp::{}", recipe.rust_name) + && policy.compatibility_rust_paths + == [format!("cuda_device::warp::{}", recipe.rust_name)], + "{} must preserve its unsafe raw and stable compatibility paths", + policy.id + ); + let expected_llvm_results: &[&str] = match warp_match.mode { + WarpMatchMode::Any => &["i32"], + WarpMatchMode::All => &["i32", "i1"], + }; + ensure!( + policy.dialect_op_type == recipe.dialect_op_type + && policy.dialect_op_name == recipe.dialect_op_name + && policy.dialect_operands == ["i32", recipe.llvm_value] + && policy.dialect_results == ["i32"] + && policy.llvm_arguments == ["i32", recipe.llvm_value] + && policy.llvm_results == expected_llvm_results + && policy.lowering == "generated_warp_match", + "{} is outside the closed two-operand warp-match lowering recipe", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == "inaccessible_read_write" + && policy.convergent + && policy.execution_scope == "warp" + && policy.minimum_ptx == "6.0" + && policy.minimum_sm.as_deref() == Some("sm_70") + && policy.ptx_result == "u32" + && policy.targets == "all", + "{} warp-match effects, carrier, or target floor disagree with its recipe", + policy.id + ); + ensure!( + policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section + == "9.7.14.11 Parallel Synchronization and Communication Instructions: match.sync" + && policy.ptx_isa_url + == "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-match-sync", + "{} warp-match PTX provenance disagrees with the reviewed recipe", + policy.id + ); + ensure!( + declaration.properties + == [ + "IntrConvergent", + "IntrInaccessibleMemOnly", + "IntrNoCallback", + ], + "{} warp-match effects disagree with the imported declaration", + policy.id + ); + ensure!( + policy.backend_lowerings.is_empty() + && policy.packed_atomic.is_none() + && policy.redux.is_none() + && policy.vote.is_none() + && policy.active_mask.is_none() + && policy.dot_product.is_none() + && policy.packed_alu.is_none() + && policy.packed_conversion.is_none() + && policy.cp_async_copy.is_none() + && policy.cp_async_control.is_none() + && policy.mbarrier_basic.is_none() + && policy.warp_barrier.is_none() + && policy.warp_shuffle.is_none() + && policy.ldmatrix_variant.is_none() + && policy.ldmatrix_safety.is_none() + && policy.ldmatrix_adapter.is_none() + && policy.selected_address_space.is_none(), + "{} mixes another generated-family contract with warp_match", + policy.id + ); + let destination = match warp_match.mode { + WarpMatchMode::Any => OperandPattern::Register, + WarpMatchMode::All => OperandPattern::RegisterPredicatePair, + }; + ensure!( + policy.expected_ptx.mnemonic == "match" + && policy.expected_ptx.modifiers == [recipe.ptx_mode, "sync", recipe.ptx_type] + && policy.expected_ptx.operands + == [ + destination, + OperandPattern::RegisterOrImmediate, + OperandPattern::RegisterOrImmediate, + ], + "{} expected PTX does not match its closed match.sync recipe", + policy.id + ); + let actual_selection_records: BTreeSet<_> = declaration + .selections + .iter() + .map(|selection| selection.source_record.as_str()) + .collect(); + ensure!( + declaration.selections.len() == 4 + && actual_selection_records == BTreeSet::from(recipe.selections), + "{} warp-match declaration must contain exactly ii/ir/ri/rr selections", + policy.id + ); + let destination = if warp_match.mode == WarpMatchMode::All { + "$dest|$pred" + } else { + "$dest" + }; + let expected_asm = format!( + "match.{}.sync.{} \t{}, $value, $mask;", + recipe.ptx_mode, recipe.ptx_type, destination + ); + for selection in &declaration.selections { + ensure!( + selection.asm == expected_asm + && selection.predicates + == [ + "Subtarget->getPTXVersion() >= 60", + "Subtarget->getSmVersion() >= 70", + ] + && selection.constraints.is_empty(), + "{} warp-match selections disagree on PTX shape, predicates, or constraints", + policy.id + ); + } + Ok(()) +} + +pub(in crate::resolve) fn validate_elect_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, +) -> Result<()> { + ensure!( + policy.id == "elect_sync" + && policy.abi_id == "i0367" + && policy.operation_key == "warp.elect.sync" + && policy.source.is_none() + && policy.source_record.as_deref() == Some("int_nvvm_elect_sync") + && policy.llvm_symbol.as_deref() == Some("llvm.nvvm.elect.sync") + && policy.resolved_llvm_symbol.is_none(), + "{} elect identity does not match the closed recipe", + policy.id + ); + ensure!( + policy.rust_module == "warp" + && policy.rust_name == "elect_sync" + && policy.rust_arguments == ["u32"] + && policy.rust_result == "(u32, bool)" + && !policy.safe + && policy.must_use + && policy.safe_allowlist_reason.is_none() + && policy.public_rust_path == "cuda_intrinsics::warp::elect_sync" + && policy.compatibility_rust_paths == ["cuda_device::warp::elect_sync"], + "{} must keep its unsafe raw API and stable compatibility path", + policy.id + ); + ensure!( + policy.dialect_op_type == "ElectSyncOp" + && policy.dialect_op_name == "nvvm.elect_sync" + && policy.dialect_operands == ["i32"] + && policy.dialect_results == ["i32", "i1"] + && policy.llvm_arguments == ["i32"] + && policy.llvm_results == ["i32", "i1"] + && policy.lowering == "generated_elect", + "{} is outside the closed elect lowering recipe", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == "inaccessible_read_write" + && policy.convergent + && policy.execution_scope == "warp" + && policy.minimum_ptx == "8.0" + && policy.minimum_sm.as_deref() == Some("sm_90") + && policy.ptx_result == "(u32, bool)" + && policy.targets == "all", + "{} elect effects or target floor disagree with the closed recipe", + policy.id + ); + ensure!( + policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section + == "9.7.14.15 Parallel Synchronization and Communication Instructions: elect.sync" + && policy.ptx_isa_url + == "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-elect-sync", + "{} elect PTX provenance disagrees with the reviewed recipe", + policy.id + ); + ensure!( + declaration.classes == ["SDPatternOperator", "Intrinsic", "DefaultAttrsIntrinsic"] + && declaration.properties == ["IntrConvergent", "IntrInaccessibleMemOnly"] + && declaration.arguments == ["i32"] + && declaration.results == ["i32", "i1"], + "{} elect declaration facts changed", + policy.id + ); + ensure!( + declaration.selections.len() == 2 + && declaration.selections.iter().all(|selection| { + ["INT_ELECT_SYNC_I", "INT_ELECT_SYNC_R"].contains(&selection.source_record.as_str()) + && selection.asm == "elect.sync \t$dest|$pred, $mask;" + && selection.predicates + == [ + "Subtarget->getPTXVersion() >= 80", + "Subtarget->getSmVersion() >= 90", + ] + && selection.constraints.is_empty() + }), + "{} elect selections changed", + policy.id + ); + ensure!( + policy.expected_ptx.mnemonic == "elect" + && policy.expected_ptx.modifiers == ["sync"] + && policy.expected_ptx.operands + == [ + OperandPattern::RegisterPredicatePair, + OperandPattern::RegisterOrImmediate, + ], + "{} expected PTX does not match elect.sync", + policy.id + ); + ensure!( + policy.backend_lowerings.len() == 2 + && policy.backend_lowerings.iter().any(|route| { + route.backend == IntrinsicBackend::LlvmNvptx + && route.mechanism == BackendLoweringMechanism::TypedNvvm + && route.minimum_ptx.as_deref() == Some("8.0") + && route.minimum_sm.as_deref() == Some("sm_90") + && !route.evidence_profile.trim().is_empty() + }) + && policy.backend_lowerings.iter().any(|route| { + route.backend == IntrinsicBackend::LibNvvm + && route.mechanism == BackendLoweringMechanism::InlinePtx + && route.minimum_ptx.as_deref() == Some("8.0") + && route.minimum_sm.as_deref() == Some("sm_90") + && !route.evidence_profile.trim().is_empty() + }), + "{} must keep the LLVM typed and libNVVM inline-PTX routes explicit", + policy.id + ); + ensure!( + policy.packed_atomic.is_none() + && policy.redux.is_none() + && policy.vote.is_none() + && policy.active_mask.is_none() + && policy.warp_match.is_none() + && policy.warp_barrier.is_none() + && policy.warp_shuffle.is_none() + && policy.dot_product.is_none() + && policy.packed_alu.is_none() + && policy.packed_conversion.is_none() + && policy.cp_async_copy.is_none() + && policy.cp_async_control.is_none() + && policy.cp_async_mbarrier.is_none() + && policy.mbarrier_basic.is_none() + && policy.movmatrix.is_none() + && policy.mbarrier_extended.is_none() + && policy.register_mma.is_none() + && policy.sparse_mma.is_none() + && policy.prmt.is_none() + && policy.cluster_barrier.is_none() + && policy.wgmma_control.is_none() + && policy.special_register.is_none() + && policy.debug_control.is_none() + && policy.cluster_memory.is_none() + && policy.clc.is_none() + && policy.ldmatrix_variant.is_none() + && policy.ldmatrix_safety.is_none() + && policy.ldmatrix_adapter.is_none() + && policy.selected_address_space.is_none(), + "{} mixes another generated-family contract with elect", + policy.id + ); + Ok(()) +} + +pub(in crate::resolve) fn validate_warp_barrier_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, +) -> Result<()> { + let barrier = policy + .warp_barrier + .as_ref() + .with_context(|| format!("{} has no closed warp-barrier contract", policy.id))?; + ensure!( + barrier.participation + == WarpBarrierParticipation::ExecutingLaneNamedAllNamedLanesSameInstructionAndMask + && barrier.legacy_pre_sm70 + == PreSm70MemberMaskRule::AllNamedLanesConvergedAndOnlyNamedLanesActive + && barrier.adapter == WarpBarrierAdapter::DirectMemberMask + && barrier.mask_encoding == WarpBarrierMaskEncoding::RegisterOrImmediate + && barrier.memory_ordering == WarpBarrierMemoryOrdering::ParticipatingLanes, + "{} requests an unsupported warp-barrier participation, legacy rule, adapter, mask encoding, or memory ordering", + policy.id + ); + ensure!( + policy.id == "sync_mask" + && policy.abi_id == "i0049" + && policy.operation_key == "warp.barrier.sync.masked" + && policy.source.is_none() + && policy.source_record.as_deref() == Some("int_nvvm_bar_warp_sync") + && policy.llvm_symbol.as_deref() == Some("llvm.nvvm.bar.warp.sync") + && policy.resolved_llvm_symbol.is_none(), + "{} warp-barrier identity does not match the closed sync_mask recipe", + policy.id + ); + ensure!( + policy.rust_module == "warp" + && policy.rust_name == "sync_mask" + && policy.rust_arguments == ["u32"] + && policy.rust_result == "()" + && !policy.safe + && !policy.must_use + && policy.safe_allowlist_reason.is_none() + && policy.public_rust_path == "cuda_intrinsics::warp::sync_mask" + && policy.compatibility_rust_paths == ["cuda_device::warp::sync_mask"], + "{} must keep its unsafe raw API and safe cuda-device compatibility path distinct", + policy.id + ); + ensure!( + policy.dialect_op_type == "BarWarpSyncOp" + && policy.dialect_op_name == "nvvm.bar_warp_sync" + && policy.dialect_operands == ["i32"] + && policy.dialect_results.is_empty() + && policy.llvm_arguments == ["i32"] + && policy.llvm_results.is_empty() + && policy.lowering == "generated_warp_barrier", + "{} is outside the closed one-mask warp-barrier lowering recipe", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == "read_write" + && policy.convergent + && policy.execution_scope == "warp" + && policy.minimum_ptx == "6.0" + && policy.minimum_sm.as_deref() == Some("sm_30") + && policy.ptx_result == "()" + && policy.targets == "all", + "{} warp-barrier effects or target floor disagree with the closed recipe", + policy.id + ); + ensure!( + policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section + == "9.7.14.2 Parallel Synchronization and Communication Instructions: bar.warp.sync" + && policy.ptx_isa_url + == "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-bar-warp-sync", + "{} warp-barrier PTX provenance disagrees with the reviewed recipe", + policy.id + ); + ensure!( + declaration.properties == ["IntrConvergent", "IntrNoCallback"], + "{} warp-barrier effects disagree with the imported declaration", + policy.id + ); + ensure!( + policy.packed_atomic.is_none() + && policy.redux.is_none() + && policy.vote.is_none() + && policy.active_mask.is_none() + && policy.warp_match.is_none() + && policy.warp_shuffle.is_none() + && policy.dot_product.is_none() + && policy.packed_alu.is_none() + && policy.packed_conversion.is_none() + && policy.cp_async_copy.is_none() + && policy.cp_async_control.is_none() + && policy.mbarrier_basic.is_none() + && policy.ldmatrix_variant.is_none() + && policy.ldmatrix_safety.is_none() + && policy.ldmatrix_adapter.is_none() + && policy.selected_address_space.is_none(), + "{} mixes another generated-family contract with warp_barrier", + policy.id + ); + ensure!( + policy.expected_ptx.mnemonic == "bar" + && policy.expected_ptx.modifiers == ["warp", "sync"] + && policy.expected_ptx.operands == [OperandPattern::RegisterOrImmediate], + "{} expected PTX does not match bar.warp.sync mask", + policy.id + ); + + let backend_pairs: BTreeSet<_> = policy + .backend_lowerings + .iter() + .map(|lowering| (lowering.backend, lowering.mechanism)) + .collect(); + ensure!( + policy.backend_lowerings.len() == 2 + && backend_pairs + == BTreeSet::from([ + ( + IntrinsicBackend::LlvmNvptx, + BackendLoweringMechanism::TypedNvvm, + ), + ( + IntrinsicBackend::LibNvvm, + BackendLoweringMechanism::TypedNvvm, + ), + ]), + "{} must define exactly the reviewed typed LLVM and libNVVM routes", + policy.id + ); + for lowering in &policy.backend_lowerings { + let floor_matches = match lowering.backend { + IntrinsicBackend::LlvmNvptx => { + lowering.mechanism == BackendLoweringMechanism::TypedNvvm + && lowering.minimum_ptx.as_deref() == Some("6.0") + && lowering.minimum_sm.as_deref() == Some("sm_30") + } + IntrinsicBackend::LibNvvm => { + lowering.mechanism == BackendLoweringMechanism::TypedNvvm + && lowering.minimum_ptx.as_deref() == Some("6.0") + && lowering.minimum_sm.as_deref() == Some("sm_75") + } + }; + ensure!( + floor_matches && !lowering.evidence_profile.trim().is_empty(), + "{} backend {:?} does not carry its reviewed warp-barrier profile floor", + policy.id, + lowering.backend + ); + } + + let expected_selection_records = BTreeSet::from(["INT_BAR_WARP_SYNC_I", "INT_BAR_WARP_SYNC_R"]); + let actual_selection_records: BTreeSet<_> = declaration + .selections + .iter() + .map(|selection| selection.source_record.as_str()) + .collect(); + ensure!( + declaration.selections.len() == 2 && actual_selection_records == expected_selection_records, + "{} warp-barrier declaration must contain exactly its immediate/register selection pair", + policy.id + ); + for selection in &declaration.selections { + ensure!( + selection.asm == "bar.warp.sync \t$i;" + && selection.predicates + == [ + "Subtarget->getPTXVersion() >= 60", + "Subtarget->getSmVersion() >= 30", + ] + && selection.constraints.is_empty(), + "{} warp-barrier selections disagree on PTX shape, target predicates, or constraints", + policy.id + ); + } + Ok(()) +} + +pub(in crate::resolve) fn validate_warp_shuffle_policy( + policy: &OverlayIntrinsic, + declaration: Option<&ImportedIntrinsic>, +) -> Result<()> { + let shuffle = policy + .warp_shuffle + .as_ref() + .with_context(|| format!("{} has no closed warp-shuffle contract", policy.id))?; + let recipe = warp_shuffle_recipe(shuffle.mode, shuffle.value_kind); + ensure!( + shuffle.participation + == WarpShuffleParticipation::ExecutingLaneNamedAllNamedLanesSameInstructionAndMask + && shuffle.legacy_pre_sm70 + == PreSm70MemberMaskRule::AllNamedLanesConvergedAndOnlyNamedLanesActive + && shuffle.source_lane + == WarpShuffleSourceLane::InRangeSourceActiveAndNamedOutOfRangeCopiesSelf + && shuffle.adapter == recipe.adapter + && shuffle.clamp == recipe.clamp + && shuffle.lane_encoding == recipe.operand_encoding + && shuffle.mask_encoding == recipe.operand_encoding, + "{} requests an unsupported warp-shuffle semantic or operand contract", + policy.id + ); + + let source_matches = match recipe.source { + WarpShuffleRecipeSource::LlvmImported { + source_record, + llvm_symbol, + } => { + policy.source.is_none() + && policy.source_record.as_deref() == Some(source_record) + && policy.llvm_symbol.as_deref() == Some(llvm_symbol) + && policy.resolved_llvm_symbol.is_none() + } + WarpShuffleRecipeSource::PtxNative { instruction } => { + policy.source + == Some(IntrinsicSource::PtxNative { + instruction: instruction.into(), + }) + && policy.source_record.is_none() + && policy.llvm_symbol.is_none() + && policy.resolved_llvm_symbol.is_none() + && policy.llvm_arguments.is_empty() + && policy.llvm_results.is_empty() + } + }; + ensure!( + policy.id == recipe.id + && policy.abi_id == recipe.abi_id + && policy.operation_key == recipe.operation_key + && source_matches, + "{} warp-shuffle identity does not match its closed mode and value recipe", + policy.id + ); + ensure!( + policy.rust_module == "warp" + && policy.rust_name == recipe.rust_name + && policy.rust_arguments == ["u32", recipe.rust_value, "u32"] + && policy.rust_result == recipe.rust_value + && !policy.safe + && policy.must_use + && policy.safe_allowlist_reason.is_none() + && policy.public_rust_path == format!("cuda_intrinsics::warp::{}", recipe.rust_name) + && policy.compatibility_rust_paths + == [format!("cuda_device::warp::{}", recipe.rust_name)], + "{} must preserve its unsafe must-use warp-shuffle raw API and compatibility path", + policy.id + ); + ensure!( + policy.dialect_op_type == recipe.dialect_op_type + && policy.dialect_op_name == recipe.dialect_op_name + && policy.dialect_operands == ["i32", recipe.dialect_value, "i32"] + && policy.dialect_results == [recipe.dialect_value] + && policy.lowering == recipe.lowering + && match recipe.source { + WarpShuffleRecipeSource::LlvmImported { .. } => { + policy.llvm_arguments == ["i32", recipe.dialect_value, "i32", "i32"] + && policy.llvm_results == [recipe.dialect_value] + } + WarpShuffleRecipeSource::PtxNative { .. } => { + policy.llvm_arguments.is_empty() && policy.llvm_results.is_empty() + } + }, + "{} is outside the closed warp-shuffle lowering recipe", + policy.id + ); + ensure!( + !policy.pure + && policy.memory == "inaccessible_read_write" + && policy.convergent + && policy.execution_scope == "warp" + && policy.minimum_ptx == "6.0" + && policy.minimum_sm.as_deref() == Some("sm_30") + && policy.ptx_result == recipe.rust_value + && policy.targets == "all", + "{} warp-shuffle effects, carrier, or target floor disagree with its recipe", + policy.id + ); + ensure!( + policy.ptx_isa_version == "9.3" + && policy.ptx_isa_section + == "9.7.9.6 Data Movement and Conversion Instructions: shfl.sync" + && policy.ptx_isa_url + == "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-shfl-sync", + "{} warp-shuffle PTX provenance disagrees with the reviewed recipe", + policy.id + ); + if let Some(declaration) = declaration { + ensure!( + matches!(recipe.source, WarpShuffleRecipeSource::LlvmImported { .. }) + && declaration.classes + == [ + "ClangBuiltin", + "NVVMBuiltin", + "SDPatternOperator", + "Intrinsic" + ] + && declaration.properties + == [ + "IntrConvergent", + "IntrInaccessibleMemOnly", + "IntrNoCallback", + ], + "{} warp-shuffle class or effects disagree with the imported declaration", + policy.id + ); + } else { + ensure!( + matches!(recipe.source, WarpShuffleRecipeSource::PtxNative { .. }), + "{} imported warp shuffle is missing its LLVM declaration", + policy.id + ); + } + ensure!( + policy.packed_atomic.is_none() + && policy.redux.is_none() + && policy.vote.is_none() + && policy.active_mask.is_none() + && policy.warp_match.is_none() + && policy.warp_barrier.is_none() + && policy.dot_product.is_none() + && policy.packed_alu.is_none() + && policy.packed_conversion.is_none() + && policy.cp_async_copy.is_none() + && policy.cp_async_control.is_none() + && policy.mbarrier_basic.is_none() + && policy.ldmatrix_variant.is_none() + && policy.ldmatrix_safety.is_none() + && policy.ldmatrix_adapter.is_none() + && policy.selected_address_space.is_none(), + "{} mixes another generated-family contract with warp_shuffle", + policy.id + ); + let expected_operands = match recipe.adapter { + WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp => vec![ + OperandPattern::Register, + OperandPattern::Register, + OperandPattern::RegisterOrImmediate, + OperandPattern::Exact { + value: recipe.clamp.to_string(), + }, + OperandPattern::RegisterOrImmediate, + ], + WarpShuffleAdapter::MaskValueLaneOrDeltaSplitI64LowHighB32InsertClampReassemble => vec![ + OperandPattern::Exact { value: "lo".into() }, + OperandPattern::Exact { value: "lo".into() }, + OperandPattern::Register, + OperandPattern::Exact { + value: recipe.clamp.to_string(), + }, + OperandPattern::Register, + ], + }; + ensure!( + policy.expected_ptx.mnemonic == "shfl" + && policy.expected_ptx.modifiers == ["sync", recipe.ptx_mode, "b32"] + && policy.expected_ptx.operands == expected_operands, + "{} expected PTX does not match its closed shfl.sync recipe", + policy.id + ); + + let backend_pairs: BTreeSet<_> = policy + .backend_lowerings + .iter() + .map(|lowering| (lowering.backend, lowering.mechanism)) + .collect(); + ensure!( + policy.backend_lowerings.len() == 2 + && backend_pairs + == BTreeSet::from([ + (IntrinsicBackend::LlvmNvptx, recipe.backend_mechanism), + (IntrinsicBackend::LibNvvm, recipe.backend_mechanism), + ]), + "{} must define exactly the reviewed LLVM and libNVVM routes", + policy.id + ); + for lowering in &policy.backend_lowerings { + let floor_matches = match lowering.backend { + IntrinsicBackend::LlvmNvptx => { + lowering.mechanism == recipe.backend_mechanism + && lowering.minimum_ptx.as_deref() == Some("6.0") + && lowering.minimum_sm.as_deref() == Some("sm_30") + } + IntrinsicBackend::LibNvvm => { + lowering.mechanism == recipe.backend_mechanism + && lowering.minimum_ptx.as_deref() == Some("6.0") + && lowering.minimum_sm.as_deref() == Some("sm_75") + } + }; + ensure!( + floor_matches && !lowering.evidence_profile.trim().is_empty(), + "{} backend {:?} does not carry its reviewed warp-shuffle profile floor", + policy.id, + lowering.backend + ); + } + + if let Some(declaration) = declaration { + let selection_records: BTreeSet<_> = declaration + .selections + .iter() + .map(|selection| selection.source_record.as_str()) + .collect(); + ensure!( + declaration.selections.len() == 8 + && selection_records.len() == 8 + && selection_records + .iter() + .all(|source_record| !source_record.trim().is_empty()), + "{} warp-shuffle declaration must contain exactly eight distinct operand-encoding selections", + policy.id + ); + let expected_asm = format!( + "shfl.sync.{}.b32 \t$dst, $src, $offset, $mask, $threadmask;", + recipe.ptx_mode + ); + for selection in &declaration.selections { + ensure!( + selection.asm == expected_asm + && selection.predicates + == [ + "Subtarget->getPTXVersion() >= 60", + "Subtarget->getSmVersion() >= 30", + ] + && selection.constraints.is_empty(), + "{} warp-shuffle selections disagree on PTX shape, target predicates, or constraints", + policy.id + ); + } + } + Ok(()) +} + +#[derive(Clone, Copy)] +pub(in crate::resolve) enum WarpShuffleRecipeSource { + LlvmImported { + source_record: &'static str, + llvm_symbol: &'static str, + }, + PtxNative { + instruction: &'static str, + }, +} + +pub(in crate::resolve) struct WarpShuffleRecipe { + pub(in crate::resolve) id: &'static str, + pub(in crate::resolve) abi_id: &'static str, + pub(in crate::resolve) operation_key: &'static str, + pub(in crate::resolve) source: WarpShuffleRecipeSource, + pub(in crate::resolve) rust_name: &'static str, + pub(in crate::resolve) rust_value: &'static str, + pub(in crate::resolve) dialect_value: &'static str, + pub(in crate::resolve) dialect_op_type: &'static str, + pub(in crate::resolve) dialect_op_name: &'static str, + pub(in crate::resolve) ptx_mode: &'static str, + pub(in crate::resolve) clamp: u32, + pub(in crate::resolve) adapter: WarpShuffleAdapter, + pub(in crate::resolve) operand_encoding: WarpShuffleOperandEncoding, + pub(in crate::resolve) lowering: &'static str, + pub(in crate::resolve) backend_mechanism: BackendLoweringMechanism, +} + +pub(in crate::resolve) fn warp_shuffle_recipe( + mode: WarpShuffleMode, + value_kind: WarpShuffleValueKind, +) -> WarpShuffleRecipe { + match (mode, value_kind) { + (WarpShuffleMode::Idx, WarpShuffleValueKind::I32) => WarpShuffleRecipe { + id: "shuffle_sync", + abi_id: "i0050", + operation_key: "warp.shuffle.sync.idx.i32", + source: WarpShuffleRecipeSource::LlvmImported { + source_record: "int_nvvm_shfl_sync_idx_i32", + llvm_symbol: "llvm.nvvm.shfl.sync.idx.i32", + }, + rust_name: "shuffle_sync", + rust_value: "u32", + dialect_value: "i32", + dialect_op_type: "ShflSyncIdxI32Op", + dialect_op_name: "nvvm.shfl_sync_idx_i32", + ptx_mode: "idx", + clamp: 31, + adapter: WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp, + operand_encoding: WarpShuffleOperandEncoding::RegisterOrImmediate, + lowering: "generated_warp_shuffle", + backend_mechanism: BackendLoweringMechanism::TypedNvvm, + }, + (WarpShuffleMode::Bfly, WarpShuffleValueKind::I32) => WarpShuffleRecipe { + id: "shuffle_xor_sync", + abi_id: "i0051", + operation_key: "warp.shuffle.sync.bfly.i32", + source: WarpShuffleRecipeSource::LlvmImported { + source_record: "int_nvvm_shfl_sync_bfly_i32", + llvm_symbol: "llvm.nvvm.shfl.sync.bfly.i32", + }, + rust_name: "shuffle_xor_sync", + rust_value: "u32", + dialect_value: "i32", + dialect_op_type: "ShflSyncBflyI32Op", + dialect_op_name: "nvvm.shfl_sync_bfly_i32", + ptx_mode: "bfly", + clamp: 31, + adapter: WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp, + operand_encoding: WarpShuffleOperandEncoding::RegisterOrImmediate, + lowering: "generated_warp_shuffle", + backend_mechanism: BackendLoweringMechanism::TypedNvvm, + }, + (WarpShuffleMode::Down, WarpShuffleValueKind::I32) => WarpShuffleRecipe { + id: "shuffle_down_sync", + abi_id: "i0052", + operation_key: "warp.shuffle.sync.down.i32", + source: WarpShuffleRecipeSource::LlvmImported { + source_record: "int_nvvm_shfl_sync_down_i32", + llvm_symbol: "llvm.nvvm.shfl.sync.down.i32", + }, + rust_name: "shuffle_down_sync", + rust_value: "u32", + dialect_value: "i32", + dialect_op_type: "ShflSyncDownI32Op", + dialect_op_name: "nvvm.shfl_sync_down_i32", + ptx_mode: "down", + clamp: 31, + adapter: WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp, + operand_encoding: WarpShuffleOperandEncoding::RegisterOrImmediate, + lowering: "generated_warp_shuffle", + backend_mechanism: BackendLoweringMechanism::TypedNvvm, + }, + (WarpShuffleMode::Up, WarpShuffleValueKind::I32) => WarpShuffleRecipe { + id: "shuffle_up_sync", + abi_id: "i0053", + operation_key: "warp.shuffle.sync.up.i32", + source: WarpShuffleRecipeSource::LlvmImported { + source_record: "int_nvvm_shfl_sync_up_i32", + llvm_symbol: "llvm.nvvm.shfl.sync.up.i32", + }, + rust_name: "shuffle_up_sync", + rust_value: "u32", + dialect_value: "i32", + dialect_op_type: "ShflSyncUpI32Op", + dialect_op_name: "nvvm.shfl_sync_up_i32", + ptx_mode: "up", + clamp: 0, + adapter: WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp, + operand_encoding: WarpShuffleOperandEncoding::RegisterOrImmediate, + lowering: "generated_warp_shuffle", + backend_mechanism: BackendLoweringMechanism::TypedNvvm, + }, + (WarpShuffleMode::Idx, WarpShuffleValueKind::F32) => WarpShuffleRecipe { + id: "shuffle_f32_sync", + abi_id: "i0054", + operation_key: "warp.shuffle.sync.idx.f32", + source: WarpShuffleRecipeSource::LlvmImported { + source_record: "int_nvvm_shfl_sync_idx_f32", + llvm_symbol: "llvm.nvvm.shfl.sync.idx.f32", + }, + rust_name: "shuffle_f32_sync", + rust_value: "f32", + dialect_value: "f32", + dialect_op_type: "ShflSyncIdxF32Op", + dialect_op_name: "nvvm.shfl_sync_idx_f32", + ptx_mode: "idx", + clamp: 31, + adapter: WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp, + operand_encoding: WarpShuffleOperandEncoding::RegisterOrImmediate, + lowering: "generated_warp_shuffle", + backend_mechanism: BackendLoweringMechanism::TypedNvvm, + }, + (WarpShuffleMode::Bfly, WarpShuffleValueKind::F32) => WarpShuffleRecipe { + id: "shuffle_xor_f32_sync", + abi_id: "i0055", + operation_key: "warp.shuffle.sync.bfly.f32", + source: WarpShuffleRecipeSource::LlvmImported { + source_record: "int_nvvm_shfl_sync_bfly_f32", + llvm_symbol: "llvm.nvvm.shfl.sync.bfly.f32", + }, + rust_name: "shuffle_xor_f32_sync", + rust_value: "f32", + dialect_value: "f32", + dialect_op_type: "ShflSyncBflyF32Op", + dialect_op_name: "nvvm.shfl_sync_bfly_f32", + ptx_mode: "bfly", + clamp: 31, + adapter: WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp, + operand_encoding: WarpShuffleOperandEncoding::RegisterOrImmediate, + lowering: "generated_warp_shuffle", + backend_mechanism: BackendLoweringMechanism::TypedNvvm, + }, + (WarpShuffleMode::Down, WarpShuffleValueKind::F32) => WarpShuffleRecipe { + id: "shuffle_down_f32_sync", + abi_id: "i0056", + operation_key: "warp.shuffle.sync.down.f32", + source: WarpShuffleRecipeSource::LlvmImported { + source_record: "int_nvvm_shfl_sync_down_f32", + llvm_symbol: "llvm.nvvm.shfl.sync.down.f32", + }, + rust_name: "shuffle_down_f32_sync", + rust_value: "f32", + dialect_value: "f32", + dialect_op_type: "ShflSyncDownF32Op", + dialect_op_name: "nvvm.shfl_sync_down_f32", + ptx_mode: "down", + clamp: 31, + adapter: WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp, + operand_encoding: WarpShuffleOperandEncoding::RegisterOrImmediate, + lowering: "generated_warp_shuffle", + backend_mechanism: BackendLoweringMechanism::TypedNvvm, + }, + (WarpShuffleMode::Up, WarpShuffleValueKind::F32) => WarpShuffleRecipe { + id: "shuffle_up_f32_sync", + abi_id: "i0057", + operation_key: "warp.shuffle.sync.up.f32", + source: WarpShuffleRecipeSource::LlvmImported { + source_record: "int_nvvm_shfl_sync_up_f32", + llvm_symbol: "llvm.nvvm.shfl.sync.up.f32", + }, + rust_name: "shuffle_up_f32_sync", + rust_value: "f32", + dialect_value: "f32", + dialect_op_type: "ShflSyncUpF32Op", + dialect_op_name: "nvvm.shfl_sync_up_f32", + ptx_mode: "up", + clamp: 0, + adapter: WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp, + operand_encoding: WarpShuffleOperandEncoding::RegisterOrImmediate, + lowering: "generated_warp_shuffle", + backend_mechanism: BackendLoweringMechanism::TypedNvvm, + }, + (WarpShuffleMode::Idx, WarpShuffleValueKind::I64) => WarpShuffleRecipe { + id: "shuffle_u64_sync", + abi_id: "i0058", + operation_key: "warp.shuffle.sync.idx.i64", + source: WarpShuffleRecipeSource::PtxNative { + instruction: "shfl.sync.idx.b32", + }, + rust_name: "shuffle_u64_sync", + rust_value: "u64", + dialect_value: "i64", + dialect_op_type: "ShflSyncIdxI64Op", + dialect_op_name: "nvvm.shfl_sync_idx_i64", + ptx_mode: "idx", + clamp: 31, + adapter: + WarpShuffleAdapter::MaskValueLaneOrDeltaSplitI64LowHighB32InsertClampReassemble, + operand_encoding: WarpShuffleOperandEncoding::RegisterOnly, + lowering: "generated_warp_shuffle_i64_inline_ptx", + backend_mechanism: BackendLoweringMechanism::InlinePtx, + }, + (WarpShuffleMode::Bfly, WarpShuffleValueKind::I64) => WarpShuffleRecipe { + id: "shuffle_xor_u64_sync", + abi_id: "i0059", + operation_key: "warp.shuffle.sync.bfly.i64", + source: WarpShuffleRecipeSource::PtxNative { + instruction: "shfl.sync.bfly.b32", + }, + rust_name: "shuffle_xor_u64_sync", + rust_value: "u64", + dialect_value: "i64", + dialect_op_type: "ShflSyncBflyI64Op", + dialect_op_name: "nvvm.shfl_sync_bfly_i64", + ptx_mode: "bfly", + clamp: 31, + adapter: + WarpShuffleAdapter::MaskValueLaneOrDeltaSplitI64LowHighB32InsertClampReassemble, + operand_encoding: WarpShuffleOperandEncoding::RegisterOnly, + lowering: "generated_warp_shuffle_i64_inline_ptx", + backend_mechanism: BackendLoweringMechanism::InlinePtx, + }, + (WarpShuffleMode::Down, WarpShuffleValueKind::I64) => WarpShuffleRecipe { + id: "shuffle_down_u64_sync", + abi_id: "i0060", + operation_key: "warp.shuffle.sync.down.i64", + source: WarpShuffleRecipeSource::PtxNative { + instruction: "shfl.sync.down.b32", + }, + rust_name: "shuffle_down_u64_sync", + rust_value: "u64", + dialect_value: "i64", + dialect_op_type: "ShflSyncDownI64Op", + dialect_op_name: "nvvm.shfl_sync_down_i64", + ptx_mode: "down", + clamp: 31, + adapter: + WarpShuffleAdapter::MaskValueLaneOrDeltaSplitI64LowHighB32InsertClampReassemble, + operand_encoding: WarpShuffleOperandEncoding::RegisterOnly, + lowering: "generated_warp_shuffle_i64_inline_ptx", + backend_mechanism: BackendLoweringMechanism::InlinePtx, + }, + (WarpShuffleMode::Up, WarpShuffleValueKind::I64) => WarpShuffleRecipe { + id: "shuffle_up_u64_sync", + abi_id: "i0061", + operation_key: "warp.shuffle.sync.up.i64", + source: WarpShuffleRecipeSource::PtxNative { + instruction: "shfl.sync.up.b32", + }, + rust_name: "shuffle_up_u64_sync", + rust_value: "u64", + dialect_value: "i64", + dialect_op_type: "ShflSyncUpI64Op", + dialect_op_name: "nvvm.shfl_sync_up_i64", + ptx_mode: "up", + clamp: 0, + adapter: + WarpShuffleAdapter::MaskValueLaneOrDeltaSplitI64LowHighB32InsertClampReassemble, + operand_encoding: WarpShuffleOperandEncoding::RegisterOnly, + lowering: "generated_warp_shuffle_i64_inline_ptx", + backend_mechanism: BackendLoweringMechanism::InlinePtx, + }, + } +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/guards.rs b/crates/cuda-intrinsics-gen/src/resolve/guards.rs new file mode 100644 index 0000000000..9ee475cbf6 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/guards.rs @@ -0,0 +1,602 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, CpAsyncControlOperation, DotProductAdapter, + ExecutionControlOperation, IntrinsicBackend, OverlayIntrinsic, Tcgen05CpGroup, + Tcgen05MmaBUsage, Tcgen05MmaKind, Tcgen05SourceContract, TmaOperation, WgmmaControlMode, +}; +use crate::ptx::OperandPattern; +use anyhow::{Result, ensure}; +use ptx_parse::ParseError; +use std::collections::{BTreeMap, BTreeSet}; + +use super::families::*; + +pub(super) fn insert_unique(set: &mut BTreeSet, value: &str, kind: &str) -> Result<()> { + ensure!(set.insert(value.to_owned()), "duplicate {kind}: {value}"); + Ok(()) +} + +pub(super) fn selection_matches_policy( + policy: &OverlayIntrinsic, + selection: &crate::model::ImportedSelection, +) -> Result { + if policy.family == "wgmma_control" { + let Some(control) = &policy.wgmma_control else { + return Ok(false); + }; + let recipe = wgmma_control_recipe(control.mode); + let selection_shape_matches = if control.mode == WgmmaControlMode::WaitGroup { + selection.asm == "wgmma.wait_group.sync.aligned \t$n;" + } else { + policy.expected_ptx.matches(&selection.asm)? + }; + return Ok(selection.source_record == recipe.selection_record + && selection_shape_matches + && selection.predicates == ["Subtarget->getPTXVersion() >= 80", "hasSM90a"] + && selection.constraints.is_empty()); + } + + if policy.family == "tma" { + return Ok(selection_matches_tma_policy(policy, selection)); + } + if matches!( + policy.family.as_str(), + "counted_barrier" | "grid_dependency" | "register_control" + ) { + let Some(operation) = ExecutionControlOperation::from_catalog_id(&policy.id) else { + return Ok(false); + }; + let recipe = execution_control_recipe(operation); + return Ok(recipe + .selection_records + .contains(&selection.source_record.as_str()) + && selection.asm == recipe.selection_asm + && selection.predicates == recipe.selection_predicates + && selection.constraints.is_empty()); + } + if policy.family == "tcgen05" { + let Some(tcgen05) = &policy.tcgen05 else { + return Ok(false); + }; + if let Some(mma) = &tcgen05.mma { + let expected = if mma.alias.is_some() { + BTreeSet::from([if tcgen05_mma_is_ws(mma.form) { + tcgen05_mma_selection_asm( + mma.form, + Tcgen05MmaKind::F8f6f4, + 1, + None, + Some(0), + Some(Tcgen05MmaBUsage::Discard), + ) + } else { + tcgen05_mma_selection_asm( + mma.form, + Tcgen05MmaKind::F8f6f4, + 1, + Some("discard"), + None, + None, + ) + }]) + } else { + tcgen05_mma_valid_selection_asms(mma.form) + }; + return Ok(expected.contains(&selection.asm) + && selection.predicates + == [if selection.asm.contains(".kind::i8.") { + "Subtarget->hasTcgen05MMAI8Kind()" + } else { + "Subtarget->hasTcgen05InstSupport()" + }] + && selection.constraints.is_empty()); + } + if tcgen05.source_contract != Tcgen05SourceContract::ExactTablegenSelection { + return Ok(false); + } + if let Some(cp) = tcgen05.cp { + let recipe = tcgen05_cp_member_recipe(cp.member); + let group = match cp.group { + Tcgen05CpGroup::Cg1 => 1, + Tcgen05CpGroup::Cg2 => 2, + }; + return Ok(selection.source_record + == format!("TCGEN05_CP_{}_cg{group}", recipe.selection_stem) + && selection.asm + == format!( + "tcgen05.cp.cta_group::{group}.{} \t[$tmem_addr], $sdesc;", + recipe.ptx_suffix + ) + && selection.predicates == ["Subtarget->hasTcgen05InstSupport()"] + && selection.constraints.is_empty()); + } + let recipe = tcgen05_recipe(tcgen05.operation); + return Ok( + recipe.selection_record == Some(selection.source_record.as_str()) + && recipe.selection_asm == Some(selection.asm.as_str()) + && selection.constraints.is_empty(), + ); + } + if policy.family == "sparse_mma" { + let Some(last) = policy.expected_ptx.operands.last() else { + return Ok(false); + }; + if *last != OperandPattern::Immediate { + return Ok(false); + } + let mut selection_shape = policy.expected_ptx.clone(); + *selection_shape.operands.last_mut().unwrap() = OperandPattern::RegisterOrImmediate; + return Ok(selection_shape.matches(&selection.asm)? && selection.constraints.is_empty()); + } + if policy.family == "sync" { + if policy.id == "sync_threads" { + return Ok(selection.source_record == "BARRIER_CTA_SYNC_ALIGNED_ALL_i" + && selection.asm == "bar.sync \t$i;" + && selection.predicates.is_empty() + && selection.constraints.is_empty()); + } + let Some(scope) = threadfence_scope_for_id(&policy.id) else { + return Ok(false); + }; + let recipe = threadfence_recipe(scope); + return Ok(selection.source_record == recipe.selection_record + && selection.asm == format!("membar.{};", recipe.ptx_level) + && selection.predicates.is_empty() + && selection.constraints.is_empty()); + } + + if policy.family == "vote" { + let Some(vote) = &policy.vote else { + return Ok(false); + }; + let recipe = vote_recipe(vote.mode); + return Ok([recipe.immediate_selection, recipe.register_selection] + .contains(&selection.source_record.as_str()) + && policy.expected_ptx.matches(&selection.asm)? + && selection.constraints.address_space.is_none() + && selection.constraints.immediate_bindings.is_empty()); + } + + if policy.family == "warp_match" { + let Some(warp_match) = &policy.warp_match else { + return Ok(false); + }; + let recipe = warp_match_recipe(warp_match.mode, warp_match.value_width); + return Ok(recipe + .selections + .contains(&selection.source_record.as_str()) + && policy.expected_ptx.matches(&selection.asm)? + && selection.constraints.is_empty()); + } + + if policy.family == "elect" { + return Ok(["INT_ELECT_SYNC_I", "INT_ELECT_SYNC_R"] + .contains(&selection.source_record.as_str()) + && selection.asm == "elect.sync \t$dest|$pred, $mask;" + && selection.constraints.is_empty()); + } + + if policy.family == "warp_barrier" { + return Ok(policy.warp_barrier.is_some() + && ["INT_BAR_WARP_SYNC_I", "INT_BAR_WARP_SYNC_R"] + .contains(&selection.source_record.as_str()) + && policy.expected_ptx.matches(&selection.asm)? + && selection.constraints.is_empty()); + } + + if policy.family == "warp_shuffle" { + let Some(shuffle) = &policy.warp_shuffle else { + return Ok(false); + }; + let recipe = warp_shuffle_recipe(shuffle.mode, shuffle.value_kind); + return Ok(selection.asm + == format!( + "shfl.sync.{}.b32 \t$dst, $src, $offset, $mask, $threadmask;", + recipe.ptx_mode + ) + && selection.constraints.is_empty()); + } + + if policy.family == "cp_async_copy" { + let Some(copy) = &policy.cp_async_copy else { + return Ok(false); + }; + let Some(recipe) = cp_async_copy_recipe(copy) else { + return Ok(false); + }; + return Ok(recipe + .selections + .contains(&selection.source_record.as_str()) + && policy.expected_ptx.matches(&selection.asm)? + && selection.constraints.is_empty()); + } + + if policy.family == "cp_async_control" { + let Some(control) = &policy.cp_async_control else { + return Ok(false); + }; + let recipe = cp_async_control_recipe(control.operation); + let instruction_matches = if control.operation == CpAsyncControlOperation::WaitGroup { + selection.asm == "cp.async.wait_group \t$n;" + } else { + policy.expected_ptx.matches(&selection.asm)? + }; + return Ok(selection.source_record == recipe.selection + && instruction_matches + && selection.constraints.is_empty()); + } + + if policy.family == "cp_async_mbarrier" { + let Some(bridge) = &policy.cp_async_mbarrier else { + return Ok(false); + }; + let recipe = cp_async_mbarrier_recipe(bridge.operation, bridge.state_space); + return Ok(selection.source_record == recipe.selection + && selection.asm == recipe.selection_asm + && selection.constraints.is_empty()); + } + + if policy.family == "mbarrier_basic" { + let Some(mbarrier) = &policy.mbarrier_basic else { + return Ok(false); + }; + let recipe = mbarrier_basic_recipe(mbarrier.operation); + return Ok(selection.source_record == recipe.selection + && policy.expected_ptx.matches(&selection.asm)? + && selection.constraints.is_empty()); + } + + if !policy.expected_ptx.matches(&selection.asm)? + || policy + .selected_address_space + .is_some_and(|address_space| selection.constraints.address_space != Some(address_space)) + { + return Ok(false); + } + + let Some(dot_product) = &policy.dot_product else { + return Ok(true); + }; + if selection.constraints.address_space.is_some() { + return Ok(false); + } + Ok(match dot_product.adapter { + DotProductAdapter::DirectThreeOperands => { + selection.constraints.immediate_bindings.is_empty() + } + DotProductAdapter::InsertLowHalfFalse => { + selection.constraints.immediate_bindings.len() == 1 + && selection.constraints.immediate_bindings[0].argument_index == 2 + && selection.constraints.immediate_bindings[0].value == 0 + } + }) +} + +pub(super) fn selection_matches_tma_policy( + policy: &OverlayIntrinsic, + selection: &crate::model::ImportedSelection, +) -> bool { + let Some(tma) = &policy.tma else { + return false; + }; + let operation = tma.operation; + if matches!( + operation, + TmaOperation::Reduce + | TmaOperation::PrefetchTensorMap + | TmaOperation::ReplaceBoxDim + | TmaOperation::ReplaceElementStride + | TmaOperation::ReplaceElementType + | TmaOperation::ReplaceFillMode + | TmaOperation::ReplaceGlobalAddress + | TmaOperation::ReplaceGlobalDim + | TmaOperation::ReplaceGlobalStride + | TmaOperation::ReplaceInterleaveLayout + | TmaOperation::ReplaceRank + | TmaOperation::ReplaceSwizzleAtomicity + | TmaOperation::ReplaceSwizzleMode + ) { + return false; + } + let (source_record, asm) = match operation { + TmaOperation::G2sTile1d => ( + "TMA_G2S_TILE_CG0_1D", + "cp.async.bulk.tensor.1d.shared::cluster.global.tile.mbarrier::complete_tx::bytes$cg [$dst], [$tmap, {{$d0}}], [$mbar];", + ), + TmaOperation::G2sTile2d => ( + "TMA_G2S_TILE_CG0_2D", + "cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes$cg [$dst], [$tmap, {{$d0, $d1}}], [$mbar];", + ), + TmaOperation::G2sTile2dMulticast => ( + "TMA_G2S_TILE_CG0_2D_MC", + "cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.multicast::cluster$cg [$dst], [$tmap, {{$d0, $d1}}], [$mbar], $mc;", + ), + TmaOperation::G2sTile2dMulticastCg2 => ( + "TMA_G2S_TILE_2D_MC", + "cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.multicast::cluster$cg [$dst], [$tmap, {{$d0, $d1}}], [$mbar], $mc;", + ), + TmaOperation::G2sTile3d => ( + "TMA_G2S_TILE_CG0_3D", + "cp.async.bulk.tensor.3d.shared::cluster.global.tile.mbarrier::complete_tx::bytes$cg [$dst], [$tmap, {{$d0, $d1, $d2}}], [$mbar];", + ), + TmaOperation::G2sTile4d => ( + "TMA_G2S_TILE_CG0_4D", + "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes$cg [$dst], [$tmap, {{$d0, $d1, $d2, $d3}}], [$mbar];", + ), + TmaOperation::G2sTile5d => ( + "TMA_G2S_TILE_CG0_5D", + "cp.async.bulk.tensor.5d.shared::cluster.global.tile.mbarrier::complete_tx::bytes$cg [$dst], [$tmap, {{$d0, $d1, $d2, $d3, $d4}}], [$mbar];", + ), + TmaOperation::S2gTile1d => ( + "TMA_TENSOR_S2G_TILE_1D", + "cp.async.bulk.tensor.1d.global.shared::cta.tile.bulk_group [$tmap, {{$d0}}], [$src];", + ), + TmaOperation::S2gTile2d => ( + "TMA_TENSOR_S2G_TILE_2D", + "cp.async.bulk.tensor.2d.global.shared::cta.tile.bulk_group [$tmap, {{$d0, $d1}}], [$src];", + ), + TmaOperation::S2gTile3d => ( + "TMA_TENSOR_S2G_TILE_3D", + "cp.async.bulk.tensor.3d.global.shared::cta.tile.bulk_group [$tmap, {{$d0, $d1, $d2}}], [$src];", + ), + TmaOperation::S2gTile4d => ( + "TMA_TENSOR_S2G_TILE_4D", + "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group [$tmap, {{$d0, $d1, $d2, $d3}}], [$src];", + ), + TmaOperation::S2gTile5d => ( + "TMA_TENSOR_S2G_TILE_5D", + "cp.async.bulk.tensor.5d.global.shared::cta.tile.bulk_group [$tmap, {{$d0, $d1, $d2, $d3, $d4}}], [$src];", + ), + TmaOperation::CommitGroup => ("CP_ASYNC_BULK_COMMIT_GROUP", "cp.async.bulk.commit_group;"), + TmaOperation::WaitGroup => ("CP_ASYNC_BULK_WAIT_GROUP", "cp.async.bulk.wait_group \t$n;"), + TmaOperation::WaitGroupRead => ( + "CP_ASYNC_BULK_WAIT_GROUP_READ", + "cp.async.bulk.wait_group.read \t$n;", + ), + TmaOperation::PrefetchTile1d => ( + "TMA_TENSOR_PF_TILE_1D", + "cp.async.bulk.prefetch.tensor.1d.L2.global.tile [$tmap, {{$d0}}];", + ), + TmaOperation::PrefetchTile2d => ( + "TMA_TENSOR_PF_TILE_2D", + "cp.async.bulk.prefetch.tensor.2d.L2.global.tile [$tmap, {{$d0, $d1}}];", + ), + TmaOperation::PrefetchTile3d => ( + "TMA_TENSOR_PF_TILE_3D", + "cp.async.bulk.prefetch.tensor.3d.L2.global.tile [$tmap, {{$d0, $d1, $d2}}];", + ), + TmaOperation::PrefetchTile4d => ( + "TMA_TENSOR_PF_TILE_4D", + "cp.async.bulk.prefetch.tensor.4d.L2.global.tile [$tmap, {{$d0, $d1, $d2, $d3}}];", + ), + TmaOperation::PrefetchTile5d => ( + "TMA_TENSOR_PF_TILE_5D", + "cp.async.bulk.prefetch.tensor.5d.L2.global.tile [$tmap, {{$d0, $d1, $d2, $d3, $d4}}];", + ), + TmaOperation::PrefetchTileGather4TwoDimensional => ( + "TMA_TENSOR_PF_TILE_GATHER4_2D", + "cp.async.bulk.prefetch.tensor.2d.L2.global.tile::gather4 [$tmap, {{$d0, $d1, $d2, $d3, $d4}}];", + ), + TmaOperation::PrefetchTile1dCacheHint => ( + "TMA_TENSOR_PF_TILE_1D_CH", + "cp.async.bulk.prefetch.tensor.1d.L2.global.tile.L2::cache_hint [$tmap, {{$d0}}], $ch;", + ), + TmaOperation::PrefetchTile2dCacheHint => ( + "TMA_TENSOR_PF_TILE_2D_CH", + "cp.async.bulk.prefetch.tensor.2d.L2.global.tile.L2::cache_hint [$tmap, {{$d0, $d1}}], $ch;", + ), + TmaOperation::PrefetchTile3dCacheHint => ( + "TMA_TENSOR_PF_TILE_3D_CH", + "cp.async.bulk.prefetch.tensor.3d.L2.global.tile.L2::cache_hint [$tmap, {{$d0, $d1, $d2}}], $ch;", + ), + TmaOperation::PrefetchTile4dCacheHint => ( + "TMA_TENSOR_PF_TILE_4D_CH", + "cp.async.bulk.prefetch.tensor.4d.L2.global.tile.L2::cache_hint [$tmap, {{$d0, $d1, $d2, $d3}}], $ch;", + ), + TmaOperation::PrefetchTile5dCacheHint => ( + "TMA_TENSOR_PF_TILE_5D_CH", + "cp.async.bulk.prefetch.tensor.5d.L2.global.tile.L2::cache_hint [$tmap, {{$d0, $d1, $d2, $d3, $d4}}], $ch;", + ), + TmaOperation::PrefetchTileGather4TwoDimensionalCacheHint => ( + "TMA_TENSOR_PF_TILE_GATHER4_2D_CH", + "cp.async.bulk.prefetch.tensor.2d.L2.global.tile::gather4.L2::cache_hint [$tmap, {{$d0, $d1, $d2, $d3, $d4}}], $ch;", + ), + TmaOperation::FenceProxyTensorMapAcquireCluster => ( + "INT_FENCE_PROXY_TENSORMAP_GENERIC_ACQUIRE_CLUSTER", + "fence.proxy.tensormap::generic.acquire.cluster [$addr], 128;", + ), + TmaOperation::FenceProxyTensorMapAcquireCta => ( + "INT_FENCE_PROXY_TENSORMAP_GENERIC_ACQUIRE_CTA", + "fence.proxy.tensormap::generic.acquire.cta [$addr], 128;", + ), + TmaOperation::FenceProxyTensorMapAcquireGpu => ( + "INT_FENCE_PROXY_TENSORMAP_GENERIC_ACQUIRE_GPU", + "fence.proxy.tensormap::generic.acquire.gpu [$addr], 128;", + ), + TmaOperation::FenceProxyTensorMapAcquireSystem => ( + "INT_FENCE_PROXY_TENSORMAP_GENERIC_ACQUIRE_SYS", + "fence.proxy.tensormap::generic.acquire.sys [$addr], 128;", + ), + TmaOperation::FenceProxyTensorMapReleaseCluster => ( + "INT_FENCE_PROXY_TENSORMAP_GENERIC_RELEASE_CLUSTER", + "fence.proxy.tensormap::generic.release.cluster;", + ), + TmaOperation::FenceProxyTensorMapReleaseCta => ( + "INT_FENCE_PROXY_TENSORMAP_GENERIC_RELEASE_CTA", + "fence.proxy.tensormap::generic.release.cta;", + ), + TmaOperation::FenceProxyTensorMapReleaseGpu => ( + "INT_FENCE_PROXY_TENSORMAP_GENERIC_RELEASE_GPU", + "fence.proxy.tensormap::generic.release.gpu;", + ), + TmaOperation::FenceProxyTensorMapReleaseSystem => ( + "INT_FENCE_PROXY_TENSORMAP_GENERIC_RELEASE_SYS", + "fence.proxy.tensormap::generic.release.sys;", + ), + TmaOperation::Reduce + | TmaOperation::PrefetchTensorMap + | TmaOperation::ReplaceBoxDim + | TmaOperation::ReplaceElementStride + | TmaOperation::ReplaceElementType + | TmaOperation::ReplaceFillMode + | TmaOperation::ReplaceGlobalAddress + | TmaOperation::ReplaceGlobalDim + | TmaOperation::ReplaceGlobalStride + | TmaOperation::ReplaceInterleaveLayout + | TmaOperation::ReplaceRank + | TmaOperation::ReplaceSwizzleAtomicity + | TmaOperation::ReplaceSwizzleMode => unreachable!("selectionless TMA operation"), + }; + + let mut immediate_bindings = Vec::new(); + if matches!( + operation, + TmaOperation::G2sTile1d + | TmaOperation::G2sTile2d + | TmaOperation::G2sTile2dMulticast + | TmaOperation::G2sTile2dMulticastCg2 + | TmaOperation::G2sTile3d + | TmaOperation::G2sTile4d + | TmaOperation::G2sTile5d + ) { + let dimensions = operation.dimensions().unwrap(); + immediate_bindings.push(crate::model::ImportedImmediateBinding { + argument_index: dimensions + 5, + value: if matches!( + operation, + TmaOperation::G2sTile2dMulticast | TmaOperation::G2sTile2dMulticastCg2 + ) { + -1 + } else { + 0 + }, + }); + immediate_bindings.push(crate::model::ImportedImmediateBinding { + argument_index: dimensions + 6, + value: 0, + }); + } else if matches!( + operation, + TmaOperation::S2gTile1d + | TmaOperation::S2gTile2d + | TmaOperation::S2gTile3d + | TmaOperation::S2gTile4d + | TmaOperation::S2gTile5d + ) { + immediate_bindings.push(crate::model::ImportedImmediateBinding { + argument_index: operation.dimensions().unwrap() + 3, + value: 0, + }); + } else if let Some(coordinate_count) = operation.prefetch_coordinate_count() { + immediate_bindings.push(crate::model::ImportedImmediateBinding { + argument_index: coordinate_count + 2, + value: if operation.uses_prefetch_cache_hint() { + -1 + } else { + 0 + }, + }); + } + + selection.source_record == source_record + && selection.asm == asm + && selection.constraints + == crate::model::ImportedSelectionConstraints { + address_space: None, + immediate_bindings, + } +} + +pub(super) fn ensure_exact_inline_ptx_backends( + policy: &OverlayIntrinsic, + requirements: [(IntrinsicBackend, &str, Option<&str>); 2], + family: &str, +) -> Result<()> { + let backend_pairs: BTreeSet<_> = policy + .backend_lowerings + .iter() + .map(|lowering| (lowering.backend, lowering.mechanism)) + .collect(); + ensure!( + policy.backend_lowerings.len() == 2 + && backend_pairs + == BTreeSet::from([ + ( + IntrinsicBackend::LlvmNvptx, + BackendLoweringMechanism::InlinePtx, + ), + ( + IntrinsicBackend::LibNvvm, + BackendLoweringMechanism::InlinePtx, + ), + ]), + "{} must define exactly two reviewed {family} inline-PTX routes", + policy.id + ); + let requirements: BTreeMap<_, _> = requirements + .into_iter() + .map(|(backend, ptx, minimum_sm)| (backend, (ptx, minimum_sm))) + .collect(); + for lowering in &policy.backend_lowerings { + let (minimum_ptx, minimum_sm) = requirements[&lowering.backend]; + ensure!( + lowering.minimum_ptx.as_deref() == Some(minimum_ptx) + && lowering.minimum_sm.as_deref() == minimum_sm + && !lowering.evidence_profile.trim().is_empty(), + "{} backend {:?} does not carry its exact {family} floor", + policy.id, + lowering.backend + ); + } + Ok(()) +} + +pub(super) fn ensure_no_other_family_contract( + policy: &OverlayIntrinsic, + family: &str, +) -> Result<()> { + ensure!( + policy.packed_atomic.is_none() + && policy.redux.is_none() + && policy.vote.is_none() + && policy.active_mask.is_none() + && policy.warp_match.is_none() + && policy.warp_barrier.is_none() + && policy.warp_shuffle.is_none() + && policy.dot_product.is_none() + && policy.ldmatrix_variant.is_none() + && policy.ldmatrix_safety.is_none() + && policy.ldmatrix_adapter.is_none() + && (policy.family == "tma" || policy.selected_address_space.is_none()) + && (policy.family == "packed_alu") == policy.packed_alu.is_some() + && (policy.family == "integer_minmax") == policy.integer_minmax.is_some() + && (policy.family == "packed_conversion") == policy.packed_conversion.is_some() + && (policy.family == "scalar_conversion") == policy.scalar_conversion.is_some() + && (policy.family == "scalar_arithmetic") == policy.scalar_arithmetic.is_some() + && (policy.family == "scalar_math") == policy.scalar_math.is_some() + && (policy.family == "extended_minmax") == policy.extended_minmax.is_some() + && (policy.family == "cp_async_copy") == policy.cp_async_copy.is_some() + && (policy.family == "cp_async_control") == policy.cp_async_control.is_some() + && (policy.family == "cp_async_mbarrier") == policy.cp_async_mbarrier.is_some() + && (policy.family == "mbarrier_basic") == policy.mbarrier_basic.is_some() + && (policy.family == "movmatrix") == policy.movmatrix.is_some() + && (policy.family == "mbarrier_extended") == policy.mbarrier_extended.is_some() + && (policy.family == "register_mma") == policy.register_mma.is_some() + && (policy.family == "sparse_mma") == policy.sparse_mma.is_some() + && (policy.family == "prmt") == policy.prmt.is_some() + && (policy.family == "cluster_barrier") == policy.cluster_barrier.is_some() + && (policy.family == "wgmma_control") == policy.wgmma_control.is_some() + && (policy.family == "debug_control") == policy.debug_control.is_some() + && (policy.family == "cluster_memory") == policy.cluster_memory.is_some() + && (policy.family == "clc") == policy.clc.is_some() + && (policy.family == "tma") == policy.tma.is_some() + && (policy.family == "tcgen05") == policy.tcgen05.is_some(), + "{} mixes another generated-family contract with {family}", + policy.id + ); + Ok(()) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/materialize.rs b/crates/cuda-intrinsics-gen/src/resolve/materialize.rs new file mode 100644 index 0000000000..2a5af0aa82 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/materialize.rs @@ -0,0 +1,490 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + CatalogBackend, CatalogBackendLowering, CatalogDialect, CatalogHalfOpenRange, CatalogIntrinsic, + CatalogLdmatrix, CatalogLlvm, CatalogLlvmResultFacts, CatalogRust, CatalogSelection, + CatalogSemantics, CatalogTarget, CatalogTargetRequirement, EvidenceRecord, + ExecutionControlOperation, ImportedIntrinsic, IntrinsicSource, OverlayIntrinsic, + RuntimeValidation, +}; +use anyhow::{Context, Result, ensure}; +use std::collections::BTreeMap; + +use super::abi_ledger::*; +use super::evidence::*; +use super::guards::*; +use super::targets::*; + +pub(super) fn resolve_backend_lowerings( + policy: &OverlayIntrinsic, + evidence_by_profile_id: &BTreeMap<(&str, &str), IndexedEvidence<'_>>, +) -> Result> { + let mut resolved = Vec::with_capacity(policy.backend_lowerings.len()); + let mut runtime_states = Vec::with_capacity(policy.backend_lowerings.len()); + for lowering in &policy.backend_lowerings { + let evidence = evidence_by_profile_id + .get(&(lowering.evidence_profile.as_str(), policy.id.as_str())) + .with_context(|| { + format!( + "{} has no evidence in backend profile {}", + policy.id, lowering.evidence_profile + ) + })?; + validate_evidence(policy, evidence, Some(lowering))?; + runtime_states.push(evidence.record.runtime_validation); + resolved.push(CatalogBackendLowering { + backend: lowering.backend, + mechanism: lowering.mechanism, + evidence_profile: lowering.evidence_profile.clone(), + target: backend_target_requirement(policy, lowering)?, + version: evidence.backend_version.to_owned(), + sha256: evidence.backend_sha256.to_owned(), + artifact_path: evidence.file.artifact_path.clone(), + build_id_prefix: evidence.file.build_id_prefix.clone(), + status: evidence.record.status.clone(), + stages: evidence.record.stages.clone(), + }); + } + if let Some(safety) = &policy.ldmatrix_safety { + match safety.runtime_validation { + RuntimeValidation::Unexecuted => ensure!( + runtime_states + .iter() + .all(|state| *state == Some(RuntimeValidation::Unexecuted)), + "{} overlay says runtime is unexecuted but backend evidence disagrees", + policy.id + ), + RuntimeValidation::Executed => ensure!( + runtime_states.contains(&Some(RuntimeValidation::Executed)), + "{} overlay says runtime is executed but no backend evidence has an executed runtime stage", + policy.id + ), + } + } + if policy.family == "stmatrix" { + ensure!( + runtime_states + .iter() + .all(|state| *state == Some(RuntimeValidation::Unexecuted)), + "{} stmatrix source contract is unexecuted but backend evidence disagrees", + policy.id + ); + } + if let Some(packed) = &policy.packed_atomic { + match packed.runtime_validation { + RuntimeValidation::Unexecuted => ensure!( + runtime_states + .iter() + .all(|state| *state == Some(RuntimeValidation::Unexecuted)), + "{} packed-atomic runtime is unexecuted but backend evidence disagrees", + policy.id + ), + RuntimeValidation::Executed => ensure!( + runtime_states.contains(&Some(RuntimeValidation::Executed)), + "{} packed-atomic runtime is executed but no backend evidence records execution", + policy.id + ), + } + } + if let Some(mbarrier) = &policy.mbarrier_basic { + match mbarrier.runtime_validation { + RuntimeValidation::Unexecuted => ensure!( + runtime_states + .iter() + .all(|state| *state == Some(RuntimeValidation::Unexecuted)), + "{} mbarrier runtime is unexecuted but backend evidence disagrees", + policy.id + ), + RuntimeValidation::Executed => ensure!( + runtime_states.contains(&Some(RuntimeValidation::Executed)), + "{} mbarrier runtime is executed but no backend evidence records execution", + policy.id + ), + } + } + if let Some(movmatrix) = &policy.movmatrix { + match movmatrix.runtime_validation { + RuntimeValidation::Unexecuted => ensure!( + runtime_states + .iter() + .all(|state| *state == Some(RuntimeValidation::Unexecuted)), + "{} movmatrix runtime is unexecuted but backend evidence disagrees", + policy.id + ), + RuntimeValidation::Executed => ensure!( + runtime_states.contains(&Some(RuntimeValidation::Executed)), + "{} movmatrix runtime is executed but no backend evidence records execution", + policy.id + ), + } + } + if let Some(mbarrier) = &policy.mbarrier_extended { + match mbarrier.runtime_validation { + RuntimeValidation::Unexecuted => ensure!( + runtime_states + .iter() + .all(|state| *state == Some(RuntimeValidation::Unexecuted)), + "{} extended-mbarrier runtime is unexecuted but backend evidence disagrees", + policy.id + ), + RuntimeValidation::Executed => ensure!( + runtime_states.contains(&Some(RuntimeValidation::Executed)), + "{} extended-mbarrier runtime is executed but no backend evidence records execution", + policy.id + ), + } + } + if let Some(bridge) = &policy.cp_async_mbarrier { + match bridge.runtime_validation { + RuntimeValidation::Unexecuted => ensure!( + runtime_states + .iter() + .all(|state| *state == Some(RuntimeValidation::Unexecuted)), + "{} cp.async mbarrier runtime is unexecuted but backend evidence disagrees", + policy.id + ), + RuntimeValidation::Executed => ensure!( + runtime_states.contains(&Some(RuntimeValidation::Executed)), + "{} cp.async mbarrier runtime is executed but no backend evidence records execution", + policy.id + ), + } + } + if let Some(mma) = &policy.register_mma { + match mma.runtime_validation { + RuntimeValidation::Unexecuted => ensure!( + runtime_states + .iter() + .all(|state| *state == Some(RuntimeValidation::Unexecuted)), + "{} register-MMA runtime is unexecuted but backend evidence disagrees", + policy.id + ), + RuntimeValidation::Executed => ensure!( + runtime_states.contains(&Some(RuntimeValidation::Executed)), + "{} register-MMA runtime is executed but no backend evidence records execution", + policy.id + ), + } + } + if let Some(mma) = &policy.sparse_mma { + match mma.runtime_validation { + RuntimeValidation::Unexecuted => ensure!( + runtime_states + .iter() + .all(|state| *state == Some(RuntimeValidation::Unexecuted)), + "{} sparse-MMA runtime is unexecuted but backend evidence disagrees", + policy.id + ), + RuntimeValidation::Executed => ensure!( + runtime_states.contains(&Some(RuntimeValidation::Executed)), + "{} sparse-MMA runtime is executed but no backend evidence records execution", + policy.id + ), + } + } + if let Some(debug) = &policy.debug_control { + match debug.runtime_validation { + RuntimeValidation::Unexecuted => ensure!( + runtime_states + .iter() + .all(|state| *state == Some(RuntimeValidation::Unexecuted)), + "{} debug-control runtime is unexecuted but backend evidence disagrees", + policy.id + ), + RuntimeValidation::Executed => ensure!( + runtime_states.contains(&Some(RuntimeValidation::Executed)), + "{} debug-control runtime is executed but no backend evidence records execution", + policy.id + ), + } + } + if let Some(cluster_memory) = &policy.cluster_memory { + match cluster_memory.runtime_validation { + RuntimeValidation::Unexecuted => ensure!( + runtime_states + .iter() + .all(|state| *state == Some(RuntimeValidation::Unexecuted)), + "{} cluster-memory runtime is unexecuted but backend evidence disagrees", + policy.id + ), + RuntimeValidation::Executed => ensure!( + runtime_states.contains(&Some(RuntimeValidation::Executed)), + "{} cluster-memory runtime is executed but no backend evidence records execution", + policy.id + ), + } + } + if let Some(clc) = &policy.clc { + match clc.runtime_validation { + RuntimeValidation::Unexecuted => ensure!( + runtime_states + .iter() + .all(|state| *state == Some(RuntimeValidation::Unexecuted)), + "{} CLC runtime is unexecuted but backend evidence disagrees", + policy.id + ), + RuntimeValidation::Executed => ensure!( + runtime_states.contains(&Some(RuntimeValidation::Executed)), + "{} CLC runtime is executed but no backend evidence records execution", + policy.id + ), + } + } + if let Some(conversion) = &policy.scalar_conversion { + match conversion.runtime_validation { + RuntimeValidation::Unexecuted => ensure!( + runtime_states + .iter() + .all(|state| *state == Some(RuntimeValidation::Unexecuted)), + "{} scalar-conversion runtime is unexecuted but backend evidence disagrees", + policy.id + ), + RuntimeValidation::Executed => ensure!( + runtime_states.contains(&Some(RuntimeValidation::Executed)), + "{} scalar-conversion runtime is executed but no backend evidence records execution", + policy.id + ), + } + } + if let Some(arithmetic) = &policy.scalar_arithmetic { + match arithmetic.runtime_validation { + RuntimeValidation::Unexecuted => ensure!( + runtime_states + .iter() + .all(|state| *state == Some(RuntimeValidation::Unexecuted)), + "{} scalar-arithmetic runtime is unexecuted but backend evidence disagrees", + policy.id + ), + RuntimeValidation::Executed => ensure!( + runtime_states.contains(&Some(RuntimeValidation::Executed)), + "{} scalar-arithmetic runtime is executed but no backend evidence records execution", + policy.id + ), + } + } + resolved.sort_by_key(|lowering| lowering.backend); + Ok(resolved) +} + +#[allow(clippy::too_many_arguments)] +pub(super) fn resolve_record( + policy: &OverlayIntrinsic, + source: IntrinsicSource, + declaration: Option<&ImportedIntrinsic>, + evidence: &EvidenceRecord, + backend_profile: &str, + backend_version: &str, + backend_sha256: &str, + backend_lowerings: Vec, + intrinsic_abi: u32, +) -> Result { + materialize_record( + policy, + source, + declaration, + CatalogBackend { + profile: backend_profile.to_owned(), + version: backend_version.to_owned(), + sha256: backend_sha256.to_owned(), + status: evidence.status.clone(), + target_triple: evidence.target_triple.clone(), + gpu_target: evidence.gpu_target.clone(), + ptx_feature: evidence.ptx_feature.clone(), + }, + backend_lowerings, + intrinsic_abi, + ) +} + +pub(super) fn materialize_record( + policy: &OverlayIntrinsic, + source: IntrinsicSource, + declaration: Option<&ImportedIntrinsic>, + backend: CatalogBackend, + backend_lowerings: Vec, + intrinsic_abi: u32, +) -> Result { + let native_target = if let Some(mma) = policy + .tcgen05 + .as_ref() + .and_then(|tcgen05| tcgen05.mma.as_ref()) + { + mma.llvm_target.clone() + } else { + CatalogTargetRequirement { + minimum_ptx: parse_ptx_version(&policy.minimum_ptx, &policy.id)?, + hardware: parse_hardware_target(policy)?, + } + }; + let llvm = if let Some(declaration) = declaration { + Some(CatalogLlvm { + symbol: policy + .llvm_symbol + .clone() + .expect("validated imported LLVM symbol"), + resolved_symbol: policy.resolved_llvm_symbol.clone(), + arguments: policy.llvm_arguments.clone(), + results: policy.llvm_results.clone(), + properties: declaration.properties.clone(), + result_facts: imported_result_facts(&declaration.properties)?, + }) + } else { + None + }; + let preserves_empty_dialect_signature = (policy.family == "sync" + && policy.id == "sync_threads") + || matches!( + ExecutionControlOperation::from_catalog_id(&policy.id), + Some( + ExecutionControlOperation::SetMaxNRegInc | ExecutionControlOperation::SetMaxNRegDec + ) + ); + let dialect_operands = + if policy.dialect_operands.is_empty() && !preserves_empty_dialect_signature { + policy.llvm_arguments.clone() + } else { + policy.dialect_operands.clone() + }; + let dialect_results = if policy.dialect_results.is_empty() && !preserves_empty_dialect_signature + { + policy.llvm_results.clone() + } else { + policy.dialect_results.clone() + }; + let mut selections = Vec::new(); + for selection in declaration + .into_iter() + .flat_map(|declaration| declaration.selections.iter()) + { + if selection_matches_policy(policy, selection)? { + selections.push(CatalogSelection { + source_record: selection.source_record.clone(), + asm: selection.asm.clone(), + predicates: selection.predicates.clone(), + constraints: selection.constraints.clone(), + }); + } + } + Ok(CatalogIntrinsic { + id: policy.id.clone(), + operation_key: policy.operation_key.clone(), + family: policy.family.clone(), + source, + selections, + rust: CatalogRust { + abi_id: policy.abi_id.clone(), + module: policy.rust_module.clone(), + name: policy.rust_name.clone(), + arguments: policy.rust_arguments.clone(), + result: policy.rust_result.clone(), + safe: policy.safe, + must_use: policy.must_use, + safe_allowlist_reason: policy.safe_allowlist_reason.clone(), + canonical_path: canonical_rust_path(intrinsic_abi, &policy.abi_id), + public_path: policy.public_rust_path.clone(), + compatibility_paths: policy.compatibility_rust_paths.clone(), + }, + dialect: CatalogDialect { + op_type: policy.dialect_op_type.clone(), + op_name: policy.dialect_op_name.clone(), + operands: dialect_operands, + results: dialect_results, + }, + llvm, + semantics: CatalogSemantics { + pure: policy.pure, + memory: policy.memory.clone(), + convergent: policy.convergent, + execution_scope: policy.execution_scope.clone(), + }, + target: CatalogTarget { + minimum_ptx: native_target.minimum_ptx, + hardware: native_target.hardware, + ptx_result: policy.ptx_result.clone(), + targets: policy.targets.clone(), + ptx_isa_version: policy.ptx_isa_version.clone(), + ptx_isa_section: policy.ptx_isa_section.clone(), + ptx_isa_url: policy.ptx_isa_url.clone(), + }, + backend, + backend_lowerings, + packed_atomic: policy.packed_atomic.clone(), + redux: policy.redux.clone(), + vote: policy.vote.clone(), + active_mask: policy.active_mask.clone(), + warp_match: policy.warp_match.clone(), + warp_barrier: policy.warp_barrier.clone(), + warp_shuffle: policy.warp_shuffle.clone(), + dot_product: policy.dot_product.clone(), + packed_alu: policy.packed_alu.clone(), + integer_minmax: policy.integer_minmax.clone(), + packed_conversion: policy.packed_conversion.clone(), + scalar_conversion: policy.scalar_conversion.clone(), + scalar_arithmetic: policy.scalar_arithmetic.clone(), + scalar_math: policy.scalar_math.clone(), + extended_minmax: policy.extended_minmax.clone(), + cp_async_copy: policy.cp_async_copy.clone(), + cp_async_control: policy.cp_async_control.clone(), + cp_async_mbarrier: policy.cp_async_mbarrier.clone(), + mbarrier_basic: policy.mbarrier_basic.clone(), + movmatrix: policy.movmatrix.clone(), + mbarrier_extended: policy.mbarrier_extended.clone(), + register_mma: policy.register_mma.clone(), + sparse_mma: policy.sparse_mma.clone(), + prmt: policy.prmt.clone(), + cluster_barrier: policy.cluster_barrier.clone(), + wgmma_control: policy.wgmma_control.clone(), + special_register: policy.special_register.clone(), + debug_control: policy.debug_control.clone(), + cluster_memory: policy.cluster_memory.clone(), + clc: policy.clc.clone(), + tma: policy.tma.clone(), + tcgen05: policy.tcgen05.clone(), + ldmatrix: policy + .ldmatrix_variant + .clone() + .map(|variant| CatalogLdmatrix { + variant, + safety: policy + .ldmatrix_safety + .clone() + .expect("validated ldmatrix safety"), + adapter: policy.ldmatrix_adapter.expect("validated ldmatrix adapter"), + selected_address_space: policy + .selected_address_space + .expect("validated ldmatrix address space"), + }), + lowering: policy.lowering.clone(), + expected_ptx: policy.expected_ptx.clone(), + summary: policy.summary.clone(), + }) +} + +pub(super) fn imported_result_facts(properties: &[String]) -> Result { + let no_undef = properties.iter().any(|property| property == "NoUndef"); + let mut range = None; + for property in properties { + let Some(bounds) = property + .strip_prefix("Range')) + else { + continue; + }; + let (lower, upper_exclusive) = bounds + .split_once(',') + .with_context(|| format!("malformed return range property {property:?}"))?; + ensure!( + !lower.is_empty() && !upper_exclusive.is_empty(), + "malformed return range property {property:?}" + ); + ensure!(range.is_none(), "duplicate return range properties"); + range = Some(CatalogHalfOpenRange { + lower: lower.to_owned(), + upper_exclusive: upper_exclusive.to_owned(), + }); + } + Ok(CatalogLlvmResultFacts { no_undef, range }) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/mod.rs b/crates/cuda-intrinsics-gen/src/resolve/mod.rs new file mode 100644 index 0000000000..d35455dfa5 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/mod.rs @@ -0,0 +1,32 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +mod abi_ledger; +mod driver; +mod evidence; +mod families; +mod guards; +mod materialize; +mod overlay; +mod policy; +mod targets; +#[cfg(test)] +mod tests; + +pub(crate) use abi_ledger::validate_operation_key; +pub use driver::resolve; +pub(crate) use driver::resolve_candidate; +// The remaining pre-split surface: nothing outside `resolve` names these +// today, so the re-exports would otherwise trip `unused_imports`. +#[allow(unused_imports)] +pub(crate) use driver::CandidateResolution; +#[cfg(test)] +pub(crate) use driver::{test_catalog_with_clc, test_catalog_with_tcgen05, test_catalog_with_tma}; +pub(crate) use families::cluster::cluster_memory_inline_recipe; +pub(crate) use families::mbarrier_extended::mbarrier_extended_inline_recipe; +#[cfg(test)] +pub(crate) use overlay::CATALOG_SCHEMA; +#[allow(unused_imports)] +pub(crate) use targets::{resolve_target_contract, resolve_target_contracts}; diff --git a/crates/cuda-intrinsics-gen/src/resolve/overlay.rs b/crates/cuda-intrinsics-gen/src/resolve/overlay.rs new file mode 100644 index 0000000000..88d9f637c7 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/overlay.rs @@ -0,0 +1,858 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + OverlayFile, OverlayIntrinsic, OverlayShardFile, RegisterMmaAccumulator, Tcgen05Operation, + TmaOperation, +}; +use crate::util::sha256_bytes; +use anyhow::{Context, Result, ensure}; +use std::collections::{BTreeMap, BTreeSet}; +use std::fs; +use std::path::{Component, Path}; + +use super::abi_ledger::*; +use super::families::*; +use super::guards::*; + +pub(super) const OVERLAY_SCHEMA: u32 = 44; +pub(super) const MINIMUM_OVERLAY_SHARD_SCHEMA: u32 = 26; +pub(super) const OVERLAY_SHARD_SCHEMA: u32 = 62; +pub(super) const REGISTER_MMA_F8F6F4_SHARD_SCHEMA: u32 = 46; +pub(super) const REGISTER_MMA_F8F6F4_F16_SHARD_SCHEMA: u32 = 47; +pub(super) const REGISTER_MMA_MXF8F6F4_SHARD_SCHEMA: u32 = 60; +pub(super) const REGISTER_MMA_FP8_SHARD_SCHEMA: u32 = 48; +pub(super) const REGISTER_MMA_AMPERE_FLOAT_SHARD_SCHEMA: u32 = 49; +pub(super) const SPARSE_MMA_F8F6F4_SHARD_SCHEMA: u32 = 27; +pub(super) const SPARSE_MMA_F8F6F4_F16_SHARD_SCHEMA: u32 = 50; +pub(super) const PRMT_SHARD_SCHEMA: u32 = 28; +pub(super) const PACKED_CONVERSION_FP8_SHARD_SCHEMA: u32 = 29; +pub(super) const PACKED_CONVERSION_FP8_F16X2_SHARD_SCHEMA: u32 = 59; +pub(super) const CLUSTER_SREG_SHARD_SCHEMA: u32 = 30; +pub(super) const CLUSTER_BARRIER_SHARD_SCHEMA: u32 = 31; +pub(super) const SPECIAL_REGISTER_SHARD_SCHEMA: u32 = 32; +pub(super) const DEBUG_CONTROL_SHARD_SCHEMA: u32 = 33; +pub(super) const THREADFENCE_SHARD_SCHEMA: u32 = 34; +pub(super) const STMATRIX_SHARD_SCHEMA: u32 = 35; +pub(super) const CLUSTER_MEMORY_SHARD_SCHEMA: u32 = 39; +pub(super) const CLC_SHARD_SCHEMA: u32 = 40; +pub(super) const TMA_SHARD_SCHEMA: u32 = 61; +pub(super) const TMA_REDUCTION_SHARD_SCHEMA: u32 = 62; +pub(super) const MBARRIER_EXTENDED_SHARD_SCHEMA: u32 = 40; +pub(super) const WGMMA_CONTROL_SHARD_SCHEMA: u32 = 38; +pub(super) const TCGEN05_SHARD_SCHEMA: u32 = 42; +pub(super) const SCALAR_CONVERSION_SHARD_SCHEMA: u32 = 43; +pub(super) const SCALAR_ARITHMETIC_SHARD_SCHEMA: u32 = 45; +pub(super) const EXTENDED_MINMAX_SHARD_SCHEMA: u32 = 60; +pub(super) const TCGEN05_CP_SHARD_SCHEMA: u32 = 52; +pub(super) const TCGEN05_LD_SHARD_SCHEMA: u32 = 53; +pub(super) const TCGEN05_ST_SHARD_SCHEMA: u32 = 54; +pub(super) const TCGEN05_OFFSET_LDST_SHARD_SCHEMA: u32 = 55; +pub(super) const TCGEN05_CONTROL_SHARD_SCHEMA: u32 = 56; +pub(super) const TCGEN05_MMA_SHARD_SCHEMA: u32 = 57; +pub(super) const SCALAR_MATH_SHARD_SCHEMA: u32 = 58; +pub(crate) const CATALOG_SCHEMA: u32 = 46; +pub(super) fn read_overlay( + repo_root: &Path, + manifest_path: &Path, +) -> Result<(OverlayFile, String)> { + let manifest_bytes = + fs::read(manifest_path).with_context(|| format!("read {}", manifest_path.display()))?; + let mut overlay: OverlayFile = toml::from_slice(&manifest_bytes) + .with_context(|| format!("parse {}", manifest_path.display()))?; + + ensure!( + overlay.intrinsics.is_empty(), + "overlay.toml must list family shards instead of inline intrinsic records" + ); + ensure!( + !overlay.shards.is_empty(), + "overlay.toml must list at least one family shard" + ); + + let mut previous = None; + let mut seen = BTreeSet::new(); + let mut hash_input = Vec::new(); + append_overlay_hash_input(&mut hash_input, "intrinsics/overlay.toml", &manifest_bytes); + + for shard_name in &overlay.shards { + validate_overlay_shard_path(shard_name)?; + ensure!( + seen.insert(shard_name.as_str()), + "overlay.toml lists duplicate shard {shard_name}" + ); + if let Some(previous) = previous { + ensure!( + previous < shard_name.as_str(), + "overlay.toml shards must be sorted" + ); + } + previous = Some(shard_name.as_str()); + + let relative = Path::new("intrinsics").join(shard_name); + let path = repo_root.join(&relative); + let bytes = fs::read(&path).with_context(|| format!("read {}", path.display()))?; + let mut shard: OverlayShardFile = + toml::from_slice(&bytes).with_context(|| format!("parse {}", path.display()))?; + validate_overlay_shard_schema(&shard, &path)?; + let int4_mma_admission = shard.register_mma_int4.take(); + let int8_mma_admission = shard.register_mma_int8.take(); + let binary_mma_admission = shard.register_mma_b1.take(); + let f8f6f4_f32_mma_admission = shard.register_mma_f8f6f4_f32.take(); + let f8f6f4_f16_mma_admission = shard.register_mma_f8f6f4_f16.take(); + let mxf8f6f4_f32_mma_admission = shard.register_mma_mxf8f6f4_f32.take(); + let fp8_mma_admission = shard.register_mma_fp8.take(); + let ampere_float_mma_admission = shard.register_mma_ampere_float.take(); + let sparse_mma_admission = shard.sparse_mma_integer.take(); + let sparse_mma_f8f6f4_admission = shard.sparse_mma_f8f6f4_f32.take(); + let sparse_mma_f8f6f4_f16_admission = shard.sparse_mma_f8f6f4_f16.take(); + let prmt_admission = shard.prmt.take(); + let packed_conversion_fp8_admission = shard.packed_conversion_fp8.take(); + let packed_conversion_fp8_f16x2_admission = shard.packed_conversion_fp8_f16x2.take(); + let scalar_conversion_admission = shard.scalar_conversion.take(); + let scalar_arithmetic_admission = shard.scalar_arithmetic.take(); + let scalar_math_admission = shard.scalar_math.take(); + let extended_minmax_admission = shard.extended_minmax.take(); + let cluster_sreg_admission = shard.cluster_sreg.take(); + let cluster_barrier_admission = shard.cluster_barrier.take(); + let mbarrier_extended_admission = shard.mbarrier_extended.take(); + let special_register_admission = shard.special_registers.take(); + let debug_control_admission = shard.debug_control.take(); + let threadfence_admission = shard.threadfence.take(); + let cluster_memory_admission = shard.cluster_memory.take(); + let stmatrix_admission = shard.stmatrix.take(); + let clc_admission = shard.clc.take(); + let wgmma_control_admission = shard.wgmma_controls.take(); + let tma_admission = shard.tma.take(); + let tcgen05_admission = shard.tcgen05.take(); + let compact_mma_count = usize::from(int4_mma_admission.is_some()) + + usize::from(int8_mma_admission.is_some()) + + usize::from(binary_mma_admission.is_some()) + + usize::from(f8f6f4_f32_mma_admission.is_some()) + + usize::from(f8f6f4_f16_mma_admission.is_some()) + + usize::from(mxf8f6f4_f32_mma_admission.is_some()) + + usize::from(fp8_mma_admission.is_some()) + + usize::from(ampere_float_mma_admission.is_some()) + + usize::from(sparse_mma_admission.is_some()) + + usize::from(sparse_mma_f8f6f4_admission.is_some()) + + usize::from(sparse_mma_f8f6f4_f16_admission.is_some()); + ensure!( + compact_mma_count <= 1, + "overlay shard {} contains more than one compact MMA admission", + path.display() + ); + let integer_mma_admission = int4_mma_admission + .map(|admission| (RegisterMmaIntegerKind::Int4, admission)) + .or_else(|| { + int8_mma_admission.map(|admission| (RegisterMmaIntegerKind::Int8, admission)) + }); + if let Some((kind, admission)) = integer_mma_admission { + ensure!( + shard.family == "register_mma" && shard.intrinsics.is_empty(), + "compact integer MMA admission must be the only content of a register_mma shard" + ); + shard.intrinsics = expand_register_mma_integer_admission(kind, &admission)?; + } + if let Some(admission) = binary_mma_admission { + ensure!( + shard.family == "register_mma" && shard.intrinsics.is_empty(), + "compact binary MMA admission must be the only content of a register_mma shard" + ); + shard.intrinsics = expand_register_mma_binary_admission(&admission)?; + } + if let Some(admission) = f8f6f4_f32_mma_admission { + ensure!( + shard.family == "register_mma" && shard.intrinsics.is_empty(), + "compact dense f8f6f4 MMA admission must be the only content of a register_mma shard" + ); + shard.intrinsics = + expand_register_mma_f8f6f4_admission(&admission, RegisterMmaAccumulator::F32)?; + } + if let Some(admission) = f8f6f4_f16_mma_admission { + ensure!( + shard.family == "register_mma" && shard.intrinsics.is_empty(), + "compact dense f8f6f4 MMA admission must be the only content of a register_mma shard" + ); + shard.intrinsics = + expand_register_mma_f8f6f4_admission(&admission, RegisterMmaAccumulator::F16)?; + } + if let Some(admission) = mxf8f6f4_f32_mma_admission { + ensure!( + shard.family == "register_mma" && shard.intrinsics.is_empty(), + "compact dense mxf8f6f4 MMA admission must be the only content of a register_mma shard" + ); + shard.intrinsics = expand_register_mma_mxf8f6f4_admission(&admission)?; + } + if let Some(admission) = fp8_mma_admission { + ensure!( + shard.family == "register_mma" && shard.intrinsics.is_empty(), + "compact standard FP8 MMA admission must be the only content of a register_mma shard" + ); + shard.intrinsics = expand_register_mma_fp8_admission(&admission)?; + } + if let Some(admission) = ampere_float_mma_admission { + ensure!( + shard.family == "register_mma" && shard.intrinsics.is_empty(), + "compact Ampere floating-point MMA admission must be the only content of a register_mma shard" + ); + shard.intrinsics = expand_register_mma_ampere_float_admission(&admission)?; + } + if let Some(admission) = sparse_mma_admission { + ensure!( + shard.family == "sparse_mma" && shard.intrinsics.is_empty(), + "compact sparse MMA admission must be the only content of a sparse_mma shard" + ); + shard.intrinsics = expand_sparse_mma_integer_admission(&admission)?; + } + if let Some(admission) = sparse_mma_f8f6f4_admission { + ensure!( + shard.family == "sparse_mma" && shard.intrinsics.is_empty(), + "compact sparse f8f6f4 MMA admission must be the only content of a sparse_mma shard" + ); + shard.intrinsics = expand_sparse_mma_f8f6f4_admission(&admission)?; + } + if let Some(admission) = sparse_mma_f8f6f4_f16_admission { + ensure!( + shard.family == "sparse_mma" && shard.intrinsics.is_empty(), + "compact sparse f8f6f4 F16 MMA admission must be the only content of a sparse_mma shard" + ); + shard.intrinsics = expand_sparse_mma_f8f6f4_f16_admission(&admission)?; + } + if let Some(admission) = prmt_admission { + ensure!( + shard.family == "prmt" && shard.intrinsics.is_empty(), + "compact prmt admission must be the only content of a prmt shard" + ); + shard.intrinsics = expand_prmt_admission(&admission)?; + } + if let Some(admission) = packed_conversion_fp8_admission { + ensure!( + shard.family == "packed_conversion" && shard.intrinsics.is_empty(), + "compact FP8 conversion admission must be the only content of a packed_conversion shard" + ); + shard.intrinsics = expand_packed_conversion_fp8_admission(&admission)?; + } + if let Some(admission) = packed_conversion_fp8_f16x2_admission { + ensure!( + shard.family == "packed_conversion" && shard.intrinsics.is_empty(), + "compact FP8 f16x2 conversion admission must be the only content of a packed_conversion shard" + ); + shard.intrinsics = expand_packed_conversion_fp8_f16x2_admission(&admission)?; + } + if let Some(admission) = scalar_conversion_admission { + ensure!( + shard.family == "scalar_conversion" && shard.intrinsics.is_empty(), + "compact scalar-conversion admission must be the only content of its shard" + ); + shard.intrinsics = expand_scalar_conversion_admission(&admission)?; + } + if let Some(admission) = scalar_arithmetic_admission { + ensure!( + shard.family == "scalar_arithmetic" && shard.intrinsics.is_empty(), + "compact scalar-arithmetic admission must be the only content of its shard" + ); + shard.intrinsics = expand_scalar_arithmetic_admission(&admission)?; + } + if let Some(admission) = scalar_math_admission { + ensure!( + shard.family == "scalar_math" && shard.intrinsics.is_empty(), + "compact scalar-math admission must be the only content of its shard" + ); + shard.intrinsics = expand_scalar_math_admission(&admission)?; + } + if let Some(admission) = extended_minmax_admission { + ensure!( + shard.family == "extended_minmax" && shard.intrinsics.is_empty(), + "compact extended-minmax admission must be the only content of its shard" + ); + shard.intrinsics = expand_extended_minmax_admission(&admission)?; + } + if let Some(admission) = cluster_sreg_admission { + ensure!( + shard.family == "sreg" && shard.intrinsics.is_empty() && compact_mma_count == 0, + "compact cluster-sreg admission must be the only content of an sreg shard" + ); + shard.intrinsics = expand_cluster_sreg_admission(&admission)?; + } + if let Some(admission) = cluster_barrier_admission { + ensure!( + shard.family == "cluster_barrier" && shard.intrinsics.is_empty(), + "compact cluster-barrier admission must be the only content of its shard" + ); + shard.intrinsics = expand_cluster_barrier_admission(&admission)?; + } + if let Some(admission) = mbarrier_extended_admission { + ensure!( + shard.family == "mbarrier_extended" && shard.intrinsics.is_empty(), + "compact extended-mbarrier admission must be the only content of its shard" + ); + shard.intrinsics = expand_mbarrier_extended_admission(&admission)?; + } + if let Some(admission) = special_register_admission { + ensure!( + shard.family == "sreg" && shard.intrinsics.is_empty(), + "compact special-register admission must be the only content of an sreg shard" + ); + shard.intrinsics = expand_special_register_admission(&admission)?; + } + if let Some(admission) = debug_control_admission { + ensure!( + shard.family == "debug_control" && shard.intrinsics.is_empty(), + "compact debug-control admission must be the only content of a debug_control shard" + ); + shard.intrinsics = expand_debug_control_admission(&admission)?; + } + if let Some(admission) = threadfence_admission { + ensure!( + shard.family == "sync" && shard.intrinsics.is_empty(), + "compact threadfence admission must be the only content of a sync shard" + ); + shard.intrinsics = expand_threadfence_admission(&admission)?; + } + if let Some(admission) = cluster_memory_admission { + ensure!( + shard.family == "cluster_memory" && shard.intrinsics.is_empty(), + "compact cluster-memory admission must be the only content of its shard" + ); + shard.intrinsics = expand_cluster_memory_admission(&admission)?; + } + if let Some(admission) = stmatrix_admission { + ensure!( + shard.family == "stmatrix" && shard.intrinsics.is_empty(), + "compact stmatrix admission must be the only content of its shard" + ); + shard.intrinsics = expand_stmatrix_admission(&admission)?; + } + if let Some(admission) = clc_admission { + ensure!( + shard.family == "clc" && shard.intrinsics.is_empty(), + "compact CLC admission must be the only content of a clc shard" + ); + shard.intrinsics = expand_clc_admission(&admission)?; + } + if let Some(admission) = wgmma_control_admission { + ensure!( + shard.family == "wgmma_control" && shard.intrinsics.is_empty(), + "compact WGMMA-control admission must be the only content of its shard" + ); + shard.intrinsics = expand_wgmma_control_admission(&admission)?; + } + if let Some(admission) = tma_admission { + ensure!( + shard.family == "tma" && shard.intrinsics.is_empty(), + "compact TMA admission must be the only content of a tma shard" + ); + shard.intrinsics = expand_tma_admission(&admission)?; + } + if let Some(admission) = tcgen05_admission { + ensure!( + shard.family == "tcgen05" && shard.intrinsics.is_empty(), + "compact tcgen05 admission must be the only content of a tcgen05 shard" + ); + shard.intrinsics = expand_tcgen05_admission(&admission)?; + } + ensure!( + !shard.intrinsics.is_empty(), + "overlay shard {} contains no intrinsic records", + path.display() + ); + for record in &shard.intrinsics { + ensure!( + record.family == shard.family, + "overlay shard {} declares family {}, but intrinsic {} uses family {}", + path.display(), + shard.family, + record.id, + record.family + ); + } + + append_overlay_hash_input( + &mut hash_input, + relative + .to_str() + .context("overlay shard path is not valid UTF-8")?, + &bytes, + ); + overlay.intrinsics.extend(shard.intrinsics); + } + + Ok((overlay, sha256_bytes(&hash_input))) +} + +pub(super) fn validate_overlay_shard_schema(shard: &OverlayShardFile, path: &Path) -> Result<()> { + validate_overlay_shard_schema_with_max(shard, path, OVERLAY_SHARD_SCHEMA) +} + +pub(super) fn validate_overlay_shard_schema_with_max( + shard: &OverlayShardFile, + path: &Path, + maximum_schema: u32, +) -> Result<()> { + ensure!( + (MINIMUM_OVERLAY_SHARD_SCHEMA..=maximum_schema).contains(&shard.schema), + "unsupported overlay shard schema {} in {}", + shard.schema, + path.display() + ); + ensure!( + shard.sparse_mma_f8f6f4_f32.is_none() || shard.schema >= SPARSE_MMA_F8F6F4_SHARD_SCHEMA, + "compact sparse f8f6f4 MMA admission requires overlay shard schema {}", + SPARSE_MMA_F8F6F4_SHARD_SCHEMA + ); + ensure!( + shard.sparse_mma_f8f6f4_f16.is_none() || shard.schema >= SPARSE_MMA_F8F6F4_F16_SHARD_SCHEMA, + "compact sparse f8f6f4 F16 MMA admission requires overlay shard schema {}", + SPARSE_MMA_F8F6F4_F16_SHARD_SCHEMA + ); + ensure!( + shard.register_mma_f8f6f4_f32.is_none() || shard.schema >= REGISTER_MMA_F8F6F4_SHARD_SCHEMA, + "compact dense f8f6f4 MMA admission requires overlay shard schema {}", + REGISTER_MMA_F8F6F4_SHARD_SCHEMA + ); + ensure!( + shard.register_mma_f8f6f4_f16.is_none() + || shard.schema >= REGISTER_MMA_F8F6F4_F16_SHARD_SCHEMA, + "compact dense f8f6f4 F16 MMA admission requires overlay shard schema {}", + REGISTER_MMA_F8F6F4_F16_SHARD_SCHEMA + ); + ensure!( + shard.register_mma_mxf8f6f4_f32.is_none() + || shard.schema >= REGISTER_MMA_MXF8F6F4_SHARD_SCHEMA, + "compact dense mxf8f6f4 MMA admission requires overlay shard schema {}", + REGISTER_MMA_MXF8F6F4_SHARD_SCHEMA + ); + ensure!( + shard.register_mma_fp8.is_none() || shard.schema >= REGISTER_MMA_FP8_SHARD_SCHEMA, + "compact standard FP8 MMA admission requires overlay shard schema {}", + REGISTER_MMA_FP8_SHARD_SCHEMA + ); + ensure!( + shard.register_mma_ampere_float.is_none() + || shard.schema >= REGISTER_MMA_AMPERE_FLOAT_SHARD_SCHEMA, + "compact Ampere floating-point MMA admission requires overlay shard schema {}", + REGISTER_MMA_AMPERE_FLOAT_SHARD_SCHEMA + ); + ensure!( + shard.prmt.is_none() || shard.schema >= PRMT_SHARD_SCHEMA, + "compact prmt admission requires overlay shard schema {}", + PRMT_SHARD_SCHEMA + ); + ensure!( + shard.packed_conversion_fp8.is_none() || shard.schema >= PACKED_CONVERSION_FP8_SHARD_SCHEMA, + "compact FP8 conversion admission requires overlay shard schema {}", + PACKED_CONVERSION_FP8_SHARD_SCHEMA + ); + ensure!( + shard.packed_conversion_fp8_f16x2.is_none() + || shard.schema >= PACKED_CONVERSION_FP8_F16X2_SHARD_SCHEMA, + "compact FP8 f16x2 conversion admission requires overlay shard schema {}", + PACKED_CONVERSION_FP8_F16X2_SHARD_SCHEMA + ); + ensure!( + shard.scalar_conversion.is_none() || shard.schema >= SCALAR_CONVERSION_SHARD_SCHEMA, + "compact scalar-conversion admission requires overlay shard schema {}", + SCALAR_CONVERSION_SHARD_SCHEMA + ); + ensure!( + shard.scalar_arithmetic.is_none() || shard.schema >= SCALAR_ARITHMETIC_SHARD_SCHEMA, + "compact scalar-arithmetic admission requires overlay shard schema {}", + SCALAR_ARITHMETIC_SHARD_SCHEMA + ); + ensure!( + shard.scalar_math.is_none() || shard.schema >= SCALAR_MATH_SHARD_SCHEMA, + "compact scalar-math admission requires overlay shard schema {}", + SCALAR_MATH_SHARD_SCHEMA + ); + ensure!( + shard.extended_minmax.is_none() || shard.schema >= EXTENDED_MINMAX_SHARD_SCHEMA, + "compact extended-minmax admission requires overlay shard schema {}", + EXTENDED_MINMAX_SHARD_SCHEMA + ); + ensure!( + shard.cluster_sreg.is_none() || shard.schema >= CLUSTER_SREG_SHARD_SCHEMA, + "compact cluster-sreg admission requires overlay shard schema {}", + CLUSTER_SREG_SHARD_SCHEMA + ); + ensure!( + shard.cluster_barrier.is_none() || shard.schema >= CLUSTER_BARRIER_SHARD_SCHEMA, + "compact cluster-barrier admission requires overlay shard schema {}", + CLUSTER_BARRIER_SHARD_SCHEMA + ); + ensure!( + shard.special_registers.is_none() || shard.schema >= SPECIAL_REGISTER_SHARD_SCHEMA, + "compact special-register admission requires overlay shard schema {}", + SPECIAL_REGISTER_SHARD_SCHEMA + ); + ensure!( + shard.debug_control.is_none() || shard.schema >= DEBUG_CONTROL_SHARD_SCHEMA, + "compact debug-control admission requires overlay shard schema {}", + DEBUG_CONTROL_SHARD_SCHEMA + ); + ensure!( + shard.threadfence.is_none() || shard.schema >= THREADFENCE_SHARD_SCHEMA, + "compact threadfence admission requires overlay shard schema {}", + THREADFENCE_SHARD_SCHEMA + ); + ensure!( + shard.cluster_memory.is_none() || shard.schema >= CLUSTER_MEMORY_SHARD_SCHEMA, + "compact cluster-memory admission requires overlay shard schema {}", + CLUSTER_MEMORY_SHARD_SCHEMA + ); + ensure!( + shard.stmatrix.is_none() || shard.schema >= STMATRIX_SHARD_SCHEMA, + "compact stmatrix admission requires overlay shard schema {}", + STMATRIX_SHARD_SCHEMA + ); + ensure!( + shard.clc.is_none() || shard.schema >= CLC_SHARD_SCHEMA, + "compact CLC admission requires overlay shard schema {}", + CLC_SHARD_SCHEMA + ); + ensure!( + shard.tma.is_none() || shard.schema >= TMA_SHARD_SCHEMA, + "compact TMA admission requires overlay shard schema {}", + TMA_SHARD_SCHEMA + ); + ensure!( + shard + .tma + .as_ref() + .is_none_or(|admission| admission.reduce_variants.is_empty()) + || shard.schema >= TMA_REDUCTION_SHARD_SCHEMA, + "compact TMA reduction admission requires overlay shard schema {}", + TMA_REDUCTION_SHARD_SCHEMA + ); + ensure!( + shard.mbarrier_extended.is_none() || shard.schema >= MBARRIER_EXTENDED_SHARD_SCHEMA, + "compact extended-mbarrier admission requires overlay shard schema {}", + MBARRIER_EXTENDED_SHARD_SCHEMA + ); + ensure!( + shard.wgmma_controls.is_none() || shard.schema >= WGMMA_CONTROL_SHARD_SCHEMA, + "compact WGMMA-control admission requires overlay shard schema {}", + WGMMA_CONTROL_SHARD_SCHEMA + ); + ensure!( + shard.tcgen05.is_none() || shard.schema >= TCGEN05_SHARD_SCHEMA, + "compact tcgen05 admission requires overlay shard schema {}", + TCGEN05_SHARD_SCHEMA + ); + ensure!( + shard + .tcgen05 + .as_ref() + .is_none_or(|admission| admission.cp_variants.is_empty()) + || shard.schema >= TCGEN05_CP_SHARD_SCHEMA, + "compact tcgen05 copy admission requires overlay shard schema {}", + TCGEN05_CP_SHARD_SCHEMA + ); + ensure!( + shard + .tcgen05 + .as_ref() + .is_none_or(|admission| admission.ld_variants.is_empty()) + || shard.schema >= TCGEN05_LD_SHARD_SCHEMA, + "compact tcgen05 load admission requires overlay shard schema {}", + TCGEN05_LD_SHARD_SCHEMA + ); + ensure!( + shard + .tcgen05 + .as_ref() + .is_none_or(|admission| admission.st_variants.is_empty()) + || shard.schema >= TCGEN05_ST_SHARD_SCHEMA, + "compact tcgen05 store admission requires overlay shard schema {}", + TCGEN05_ST_SHARD_SCHEMA + ); + ensure!( + shard.tcgen05.as_ref().is_none_or(|admission| { + admission.ld_offset_variants.is_empty() && admission.st_offset_variants.is_empty() + }) || shard.schema >= TCGEN05_OFFSET_LDST_SHARD_SCHEMA, + "compact tcgen05 offset load/store admission requires overlay shard schema {}", + TCGEN05_OFFSET_LDST_SHARD_SCHEMA + ); + ensure!( + shard.tcgen05.as_ref().is_none_or(|admission| { + admission.control_llvm_evidence_profile.is_none() + && admission.control_libnvvm_evidence_profile.is_none() + && !admission.variants.iter().any(|variant| { + matches!( + variant.operation, + Tcgen05Operation::CommitMulticast + | Tcgen05Operation::ShiftDown + | Tcgen05Operation::ShiftDownCg2 + ) + }) + }) || shard.schema >= TCGEN05_CONTROL_SHARD_SCHEMA, + "compact tcgen05 control admission requires overlay shard schema {}", + TCGEN05_CONTROL_SHARD_SCHEMA + ); + if let Some(admission) = &shard.tcgen05 + && shard.schema >= TCGEN05_CONTROL_SHARD_SCHEMA + { + ensure!( + admission + .control_llvm_evidence_profile + .as_deref() + .is_some_and(|profile| !profile.trim().is_empty()) + && admission + .control_libnvvm_evidence_profile + .as_deref() + .is_some_and(|profile| !profile.trim().is_empty()) + && [ + Tcgen05Operation::CommitMulticast, + Tcgen05Operation::ShiftDown, + Tcgen05Operation::ShiftDownCg2, + ] + .into_iter() + .all(|operation| admission + .variants + .iter() + .any(|variant| variant.operation == operation)), + "compact tcgen05 schema {} requires all three control variants and both backend evidence profiles", + TCGEN05_CONTROL_SHARD_SCHEMA + ); + } + ensure!( + shard.tcgen05.as_ref().is_none_or(|admission| { + admission.mma_variants.is_empty() + && admission.mma_llvm_evidence_profile.is_none() + && admission.mma_libnvvm_evidence_profile.is_none() + && admission.mma_llvm_target_contracts.is_empty() + && admission.mma_libnvvm_target_contracts.is_empty() + }) || shard.schema >= TCGEN05_MMA_SHARD_SCHEMA, + "compact tcgen05 MMA admission requires overlay shard schema {}", + TCGEN05_MMA_SHARD_SCHEMA + ); + Ok(()) +} + +pub(super) fn validate_overlay_shard_path(path: &str) -> Result<()> { + let path = Path::new(path); + ensure!( + path.extension().and_then(|extension| extension.to_str()) == Some("toml"), + "overlay shard path must name a TOML file: {}", + path.display() + ); + let components: Vec<_> = path.components().collect(); + ensure!( + components.len() >= 2 && components[0] == Component::Normal("overlay".as_ref()), + "overlay shard path must stay under intrinsics/overlay: {}", + path.display() + ); + ensure!( + components + .iter() + .all(|component| matches!(component, Component::Normal(_))), + "overlay shard path contains a non-normal component: {}", + path.display() + ); + Ok(()) +} + +pub(super) fn append_overlay_hash_input(output: &mut Vec, path: &str, contents: &[u8]) { + output.extend_from_slice(&(path.len() as u64).to_le_bytes()); + output.extend_from_slice(path.as_bytes()); + output.extend_from_slice(&(contents.len() as u64).to_le_bytes()); + output.extend_from_slice(contents); +} + +pub(super) fn shares_tma_2d_g2s_symbol(record: &OverlayIntrinsic, symbol: &str) -> bool { + symbol == "llvm.nvvm.cp.async.bulk.tensor.g2s.tile.2d" + && record.tma.as_ref().is_some_and(|tma| { + matches!( + tma.operation, + TmaOperation::G2sTile2d + | TmaOperation::G2sTile2dMulticast + | TmaOperation::G2sTile2dMulticastCg2 + ) + }) +} + +pub(super) fn shares_tma_prefetch_tile_symbol(record: &OverlayIntrinsic, symbol: &str) -> bool { + let Some(tma) = record.tma.as_ref() else { + return false; + }; + let expected_symbol = match tma.operation { + TmaOperation::PrefetchTile1d | TmaOperation::PrefetchTile1dCacheHint => { + "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.1d" + } + TmaOperation::PrefetchTile2d | TmaOperation::PrefetchTile2dCacheHint => { + "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.2d" + } + TmaOperation::PrefetchTile3d | TmaOperation::PrefetchTile3dCacheHint => { + "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.3d" + } + TmaOperation::PrefetchTile4d | TmaOperation::PrefetchTile4dCacheHint => { + "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.4d" + } + TmaOperation::PrefetchTile5d | TmaOperation::PrefetchTile5dCacheHint => { + "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.5d" + } + TmaOperation::PrefetchTileGather4TwoDimensional + | TmaOperation::PrefetchTileGather4TwoDimensionalCacheHint => { + "llvm.nvvm.cp.async.bulk.tensor.prefetch.tile.gather4.2d" + } + _ => return false, + }; + symbol == expected_symbol +} + +pub(super) fn shares_tcgen05_mma_symbol(record: &OverlayIntrinsic, symbol: &str) -> bool { + record.tcgen05.as_ref().is_some_and(|tcgen05| { + (symbol == "llvm.nvvm.tcgen05.mma.ws.tensor" + && matches!( + tcgen05.operation, + Tcgen05Operation::MmaWsF16 + | Tcgen05Operation::MmaWsBf16 + | Tcgen05Operation::MmaWsTf32 + )) + || tcgen05 + .mma + .as_ref() + .is_some_and(|mma| symbol == tcgen05_mma_llvm_symbol(mma.form)) + }) +} + +pub(super) fn shares_tcgen05_ld_symbol(record: &OverlayIntrinsic, symbol: &str) -> bool { + let Some(tcgen05) = record.tcgen05.as_ref() else { + return false; + }; + if let Some(ld) = tcgen05.ld { + return tcgen05.operation == Tcgen05Operation::Ld + && record.source_record.as_deref() == Some(tcgen05_ld_source_record(ld).as_str()) + && symbol == tcgen05_ld_llvm_symbol(ld); + } + matches!( + (tcgen05.operation, record.source_record.as_deref(), symbol), + ( + Tcgen05Operation::Ld16x256bPure, + Some("int_nvvm_tcgen05_ld_16x256b_x1"), + "llvm.nvvm.tcgen05.ld.16x256b.x1" + ) | ( + Tcgen05Operation::Ld16x256bX8Pure, + Some("int_nvvm_tcgen05_ld_16x256b_x8"), + "llvm.nvvm.tcgen05.ld.16x256b.x8" + ) + ) +} + +pub(super) fn shares_tcgen05_st_symbol(record: &OverlayIntrinsic, symbol: &str) -> bool { + record.tcgen05.as_ref().is_some_and(|tcgen05| { + tcgen05.st.is_some_and(|st| { + tcgen05.operation == Tcgen05Operation::St + && record.source_record.as_deref() == Some(tcgen05_st_source_record(st).as_str()) + && symbol == tcgen05_st_llvm_symbol(st) + }) + }) +} + +pub(super) fn validate_unique_overlay( + records: &[OverlayIntrinsic], + intrinsic_abi: u32, +) -> Result<()> { + let mut ids = BTreeSet::new(); + let mut abi_ids = BTreeSet::new(); + let mut operation_keys = BTreeSet::new(); + let mut paths = BTreeSet::new(); + let mut op_variants = BTreeSet::new(); + let mut op_type_names = BTreeMap::new(); + let mut symbol_bases = BTreeMap::new(); + let mut symbols = BTreeSet::new(); + let mut rust_items = BTreeSet::new(); + for record in records { + insert_unique(&mut ids, &record.id, "catalog ID")?; + validate_abi_id(&record.abi_id)?; + insert_unique(&mut abi_ids, &record.abi_id, "intrinsic ABI ID")?; + validate_operation_key(&record.operation_key)?; + insert_unique( + &mut operation_keys, + &record.operation_key, + "intrinsic operation key", + )?; + if let Some(previous_name) = op_type_names.insert( + record.dialect_op_type.as_str(), + record.dialect_op_name.as_str(), + ) { + ensure!( + previous_name == record.dialect_op_name, + "dialect op type {} maps to both {} and {}", + record.dialect_op_type, + previous_name, + record.dialect_op_name + ); + } + let op_variant = format!( + "{}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}:{:?}", + record.dialect_op_name, + record.ldmatrix_variant, + record.packed_atomic, + record.redux, + record.vote, + record.active_mask, + record.warp_match, + record.warp_barrier, + record.warp_shuffle, + record.dot_product, + record.cp_async_copy, + record.cp_async_control, + record.cp_async_mbarrier, + record.mbarrier_basic, + record.movmatrix, + record.mbarrier_extended, + record.register_mma, + record.sparse_mma, + record.prmt, + record.cluster_barrier, + record.cluster_memory, + record.scalar_conversion, + record.scalar_arithmetic, + record.scalar_math, + record.extended_minmax, + record.tcgen05, + ); + insert_unique(&mut op_variants, &op_variant, "dialect op variant")?; + if let Some(symbol) = &record.llvm_symbol { + let is_resolved = record.resolved_llvm_symbol.is_some(); + let shares_reviewed_symbol = shares_tma_2d_g2s_symbol(record, symbol) + || shares_tma_prefetch_tile_symbol(record, symbol) + || shares_tcgen05_mma_symbol(record, symbol) + || shares_tcgen05_ld_symbol(record, symbol) + || shares_tcgen05_st_symbol(record, symbol); + if let Some((previous_was_resolved, previous_shared_symbol)) = + symbol_bases.insert(symbol, (is_resolved, shares_reviewed_symbol)) + { + ensure!( + (previous_was_resolved && is_resolved) + || (previous_shared_symbol && shares_reviewed_symbol), + "duplicate LLVM symbol {symbol} is reused without a resolved symbol" + ); + } + if !shares_reviewed_symbol { + insert_unique( + &mut symbols, + record.resolved_llvm_symbol.as_ref().unwrap_or(symbol), + "resolved LLVM symbol", + )?; + } + } + insert_unique( + &mut rust_items, + &format!("{}::{}", record.rust_module, record.rust_name), + "raw Rust item", + )?; + insert_unique( + &mut paths, + &canonical_rust_path(intrinsic_abi, &record.abi_id), + "canonical Rust path", + )?; + insert_unique(&mut paths, &record.public_rust_path, "public Rust path")?; + for path in &record.compatibility_rust_paths { + insert_unique(&mut paths, path, "compatibility Rust path")?; + } + } + Ok(()) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/policy.rs b/crates/cuda-intrinsics-gen/src/resolve/policy.rs new file mode 100644 index 0000000000..77131f0bb7 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/policy.rs @@ -0,0 +1,612 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, ClcOperation, CpAsyncSourceSize, ImportedIntrinsic, IntrinsicSource, + MbarrierExtendedOperation, OverlayIntrinsic, ScalarMathOperation, SpecialRegisterKind, + Tcgen05SourceContract, TmaOperation, +}; +use anyhow::{Context, Result, bail, ensure}; + +use super::abi_ledger::*; +use super::families::*; +use super::guards::*; +use super::targets::*; + +pub(super) fn resolve_policy_source(policy: &OverlayIntrinsic) -> Result { + match (&policy.source, &policy.source_record) { + (None, Some(source_record)) => Ok(IntrinsicSource::LlvmImported { + source_record: source_record.clone(), + }), + (Some(source @ IntrinsicSource::PtxNative { .. }), None) => Ok(source.clone()), + (Some(IntrinsicSource::LlvmImported { source_record }), None) => { + ensure!( + !source_record.trim().is_empty(), + "{} has an empty imported LLVM source record", + policy.id + ); + Ok(IntrinsicSource::LlvmImported { + source_record: source_record.clone(), + }) + } + (Some(_), Some(_)) => bail!( + "{} mixes tagged source provenance with the legacy source_record field", + policy.id + ), + (None, None) => bail!("{} has no intrinsic source provenance", policy.id), + } +} + +pub(super) fn validate_policy( + policy: &OverlayIntrinsic, + source: &IntrinsicSource, + declaration: Option<&ImportedIntrinsic>, + intrinsic_abi: u32, +) -> Result<()> { + validate_abi_id(&policy.abi_id)?; + parse_ptx_version(&policy.minimum_ptx, &policy.id)?; + parse_hardware_target(policy)?; + policy.expected_ptx.validate().map_err(|reason| { + anyhow::anyhow!( + "{} has an invalid expected PTX pattern: {reason}", + policy.id + ) + })?; + let public_path = format!( + "cuda_intrinsics::{}::{}", + policy.rust_module, policy.rust_name + ); + ensure!( + policy.public_rust_path == public_path, + "{} public Rust path must be {}", + policy.id, + public_path + ); + let canonical_path = canonical_rust_path(intrinsic_abi, &policy.abi_id); + ensure!( + canonical_path != policy.public_rust_path + && !policy + .compatibility_rust_paths + .iter() + .any(|path| path == &canonical_path || path == &policy.public_rust_path), + "{} must keep canonical, public, and compatibility Rust paths distinct", + policy.id + ); + match (source, declaration) { + (IntrinsicSource::LlvmImported { .. }, Some(declaration)) => { + ensure!( + policy.llvm_symbol.as_deref() == Some(declaration.llvm_name.as_str()), + "{} LLVM symbol mismatch: imported {}, overlay {:?}", + policy.id, + declaration.llvm_name, + policy.llvm_symbol + ); + ensure!( + declaration.arguments == policy.llvm_arguments, + "{} LLVM argument signature mismatch: imported {:?}, overlay {:?}", + policy.id, + declaration.arguments, + policy.llvm_arguments + ); + ensure!( + declaration.results == policy.llvm_results, + "{} LLVM result signature mismatch: imported {:?}, overlay {:?}", + policy.id, + declaration.results, + policy.llvm_results + ); + } + (IntrinsicSource::PtxNative { instruction }, None) => ensure!( + !instruction.trim().is_empty() + && policy.llvm_symbol.is_none() + && policy.resolved_llvm_symbol.is_none() + && policy.llvm_arguments.is_empty() + && policy.llvm_results.is_empty(), + "{} PTX-native source must not invent LLVM source facts", + policy.id + ), + _ => bail!( + "{} source kind and imported declaration disagree", + policy.id + ), + } + match policy.family.as_str() { + "sreg" => validate_sreg_policy(policy, source, declaration)?, + "ldmatrix" => validate_ldmatrix_policy( + policy, + declaration.context("ldmatrix requires imported LLVM declaration")?, + )?, + "stmatrix" => validate_stmatrix_policy( + policy, + declaration.context("stmatrix requires imported LLVM declaration")?, + )?, + "packed_atomic" => validate_packed_atomic_policy(policy, source)?, + "redux" => validate_redux_policy( + policy, + declaration.context("redux requires imported LLVM declaration")?, + )?, + "dotprod" => validate_dot_product_policy( + policy, + declaration.context("dotprod requires imported LLVM declaration")?, + )?, + "sync" => validate_sync_policy( + policy, + declaration.context("sync requires imported LLVM declaration")?, + )?, + "vote" => validate_vote_policy( + policy, + declaration.context("vote requires imported LLVM declaration")?, + )?, + "active_mask" => validate_active_mask_policy( + policy, + declaration.context("active_mask requires imported LLVM declaration")?, + )?, + "warp_match" => validate_warp_match_policy( + policy, + declaration.context("warp_match requires imported LLVM declaration")?, + )?, + "elect" => validate_elect_policy( + policy, + declaration.context("elect requires imported LLVM declaration")?, + )?, + "warp_barrier" => validate_warp_barrier_policy( + policy, + declaration.context("warp_barrier requires imported LLVM declaration")?, + )?, + "warp_shuffle" => validate_warp_shuffle_policy(policy, declaration)?, + "packed_alu" => validate_packed_alu_policy(policy, source, declaration)?, + "integer_minmax" => validate_integer_minmax_policy(policy, source, declaration)?, + "packed_conversion" => validate_packed_conversion_policy(policy, source, declaration)?, + "scalar_conversion" => validate_scalar_conversion_policy( + policy, + declaration.context("scalar_conversion requires imported LLVM declaration")?, + )?, + "scalar_arithmetic" => validate_scalar_arithmetic_policy( + policy, + declaration.context("scalar_arithmetic requires imported LLVM declaration")?, + )?, + "scalar_math" => validate_scalar_math_policy(policy, declaration)?, + "extended_minmax" => validate_extended_minmax_policy( + policy, + declaration.context("extended_minmax requires imported LLVM declaration")?, + )?, + "cp_async_copy" => validate_cp_async_copy_policy( + policy, + declaration.context("cp_async_copy requires imported LLVM declaration")?, + )?, + "cp_async_control" => validate_cp_async_control_policy( + policy, + declaration.context("cp_async_control requires imported LLVM declaration")?, + )?, + "cp_async_mbarrier" => validate_cp_async_mbarrier_policy( + policy, + declaration.context("cp_async_mbarrier requires imported LLVM declaration")?, + )?, + "mbarrier_basic" => validate_mbarrier_basic_policy( + policy, + declaration.context("mbarrier_basic requires imported LLVM declaration")?, + )?, + "movmatrix" => validate_movmatrix_policy(policy, source)?, + "mbarrier_extended" => validate_mbarrier_extended_policy(policy, source, declaration)?, + "register_mma" => validate_register_mma_policy( + policy, + declaration.context("register_mma requires imported LLVM declaration")?, + )?, + "sparse_mma" => validate_sparse_mma_policy( + policy, + declaration.context("sparse_mma requires imported LLVM declaration")?, + )?, + "prmt" => validate_prmt_policy( + policy, + declaration.context("prmt requires imported LLVM declaration")?, + )?, + "cluster_barrier" => validate_cluster_barrier_policy( + policy, + declaration.context("cluster_barrier requires imported LLVM declaration")?, + )?, + "debug_control" => validate_debug_control_policy(policy, source)?, + "cluster_memory" => validate_cluster_memory_policy(policy, source, declaration)?, + "clc" => validate_clc_policy( + policy, + declaration.context("clc requires imported LLVM declaration")?, + )?, + "wgmma_control" => validate_wgmma_control_policy( + policy, + declaration.context("wgmma_control requires imported LLVM declaration")?, + )?, + "tma" => validate_tma_policy( + policy, + declaration.context("tma requires imported LLVM declaration")?, + )?, + "counted_barrier" | "grid_dependency" | "register_control" => { + validate_execution_control_policy( + policy, + declaration.context("execution-control requires imported LLVM declaration")?, + )? + } + "tcgen05" => validate_tcgen05_policy( + policy, + declaration.context("tcgen05 requires imported LLVM declaration")?, + )?, + family => bail!("{} uses unsupported generated family {family:?}", policy.id), + } + ensure!( + (policy.family == "movmatrix") == policy.movmatrix.is_some(), + "{} mixes the movmatrix contract with another generated family", + policy.id + ); + ensure!( + (policy.family == "mbarrier_extended") == policy.mbarrier_extended.is_some(), + "{} mixes the extended-mbarrier contract with another generated family", + policy.id + ); + ensure!( + (policy.family == "register_mma") == policy.register_mma.is_some(), + "{} mixes the register-MMA contract with another generated family", + policy.id + ); + ensure!( + (policy.family == "sparse_mma") == policy.sparse_mma.is_some(), + "{} mixes the sparse-MMA contract with another generated family", + policy.id + ); + ensure!( + (policy.family == "prmt") == policy.prmt.is_some(), + "{} mixes the prmt contract with another generated family", + policy.id + ); + ensure!( + (policy.family == "cluster_barrier") == policy.cluster_barrier.is_some(), + "{} mixes the cluster-barrier contract with another generated family", + policy.id + ); + ensure!( + (policy.family == "wgmma_control") == policy.wgmma_control.is_some(), + "{} mixes the WGMMA-control contract with another generated family", + policy.id + ); + ensure!( + (policy.family == "tma") == policy.tma.is_some(), + "{} mixes the TMA contract with another generated family", + policy.id + ); + ensure!( + policy.special_register.is_none() || policy.family == "sreg", + "{} mixes the special-register contract with another generated family", + policy.id + ); + ensure!( + (policy.family == "debug_control") == policy.debug_control.is_some(), + "{} mixes the debug-control contract with another generated family", + policy.id + ); + ensure!( + (policy.family == "cluster_memory") == policy.cluster_memory.is_some(), + "{} mixes the cluster-memory contract with another generated family", + policy.id + ); + ensure!( + (policy.family == "clc") == policy.clc.is_some(), + "{} mixes the CLC contract with another generated family", + policy.id + ); + ensure!( + (policy.family == "tma") == policy.tma.is_some(), + "{} mixes the TMA contract with another generated family", + policy.id + ); + ensure!( + (policy.family == "tcgen05") == policy.tcgen05.is_some(), + "{} mixes the tcgen05 contract with another generated family", + policy.id + ); + ensure!( + (policy.family == "scalar_conversion") == policy.scalar_conversion.is_some(), + "{} mixes the scalar-conversion contract with another generated family", + policy.id + ); + ensure!( + (policy.family == "scalar_arithmetic") == policy.scalar_arithmetic.is_some(), + "{} mixes the scalar-arithmetic contract with another generated family", + policy.id + ); + ensure!( + (policy.family == "scalar_math") == policy.scalar_math.is_some(), + "{} mixes the scalar-math contract with another generated family", + policy.id + ); + ensure!( + (policy.family == "extended_minmax") == policy.extended_minmax.is_some(), + "{} mixes the extended-minmax contract with another generated family", + policy.id + ); + ensure!( + !policy.execution_scope.trim().is_empty(), + "{} has no execution scope", + policy.id + ); + ensure!( + !policy.ptx_isa_version.trim().is_empty() + && !policy.ptx_isa_section.trim().is_empty() + && policy.ptx_isa_url.starts_with("https://docs.nvidia.com/"), + "{} has incomplete or non-authoritative PTX ISA provenance", + policy.id + ); + match (policy.safe, policy.safe_allowlist_reason.as_deref()) { + (true, Some(reason)) if !reason.trim().is_empty() => {} + (true, _) => bail!( + "{} is safe but has no nonempty safe_allowlist_reason", + policy.id + ), + (false, Some(reason)) if !reason.trim().is_empty() => bail!( + "{} is unsafe but has a safe_allowlist_reason; safe exceptions apply only to safe items", + policy.id + ), + (false, _) => {} + } + if let Some(declaration) = declaration { + if policy.pure { + ensure!( + declaration + .classes + .iter() + .any(|class| class == "NVVMPureIntrinsic") + || (matches!( + policy.family.as_str(), + "packed_alu" | "scalar_arithmetic" | "extended_minmax" + ) && declaration + .properties + .iter() + .any(|property| property == "IntrNoMem") + && declaration + .properties + .iter() + .any(|property| property == "IntrSpeculatable")) + || (policy.family == "scalar_math" + && declaration + .properties + .iter() + .any(|property| property == "IntrNoMem")), + "{} is marked pure, but its imported declaration is not an NVVMPureIntrinsic", + policy.id + ); + } + if policy.memory == "none" { + ensure!( + declaration + .properties + .iter() + .any(|property| property == "IntrNoMem"), + "{} is marked no-memory, but its imported declaration lacks IntrNoMem", + policy.id + ); + } + let imported_convergent = declaration + .properties + .iter() + .any(|property| property == "IntrConvergent"); + let convergence_supplied_by_ptx = + (matches!(policy.family.as_str(), "register_mma" | "sparse_mma") + && (policy.register_mma.is_some() || policy.sparse_mma.is_some()) + || (policy.family == "cluster_memory" + && policy.cluster_memory.is_some() + && policy.backend_lowerings.iter().all(|lowering| { + lowering.mechanism == BackendLoweringMechanism::InlinePtx + })) + || (policy.family == "mbarrier_extended" + && policy.mbarrier_extended.is_some() + && policy.backend_lowerings.iter().all(|lowering| { + lowering.mechanism == BackendLoweringMechanism::InlinePtx + })) + || (policy.family == "tcgen05" + && policy.tcgen05.is_some() + && policy.backend_lowerings.iter().all(|lowering| { + lowering.mechanism == BackendLoweringMechanism::InlinePtx + }))) + && policy.convergent + && !imported_convergent; + ensure!( + imported_convergent == policy.convergent || convergence_supplied_by_ptx, + "{} convergence mismatch: imported {}, overlay {}", + policy.id, + imported_convergent, + policy.convergent + ); + let selectionless_closed_family = (policy.family == "packed_conversion" + && policy.packed_conversion.is_some()) + || (policy.family == "prmt" && policy.prmt.is_some()) + || (policy.family == "clc" + && policy.clc.as_ref().is_some_and(|clc| { + matches!( + clc.operation, + ClcOperation::QueryIsCanceled + | ClcOperation::QueryGetFirstCtaidX + | ClcOperation::QueryGetFirstCtaidY + | ClcOperation::QueryGetFirstCtaidZ + ) + })) + || (policy.family == "sreg" + && policy.special_register.as_ref().is_some_and(|special| { + matches!( + special.register, + SpecialRegisterKind::Envreg1 | SpecialRegisterKind::Envreg2 + ) + })) + || policy.family == "stmatrix" + || (policy.family == "tma" + && policy.tma.as_ref().is_some_and(|tma| { + matches!( + tma.operation, + TmaOperation::Reduce + | TmaOperation::PrefetchTensorMap + | TmaOperation::ReplaceBoxDim + | TmaOperation::ReplaceElementStride + | TmaOperation::ReplaceElementType + | TmaOperation::ReplaceFillMode + | TmaOperation::ReplaceGlobalAddress + | TmaOperation::ReplaceGlobalDim + | TmaOperation::ReplaceGlobalStride + | TmaOperation::ReplaceInterleaveLayout + | TmaOperation::ReplaceRank + | TmaOperation::ReplaceSwizzleAtomicity + | TmaOperation::ReplaceSwizzleMode + ) + })) + || (policy.family == "tcgen05" + && policy.tcgen05.as_ref().is_some_and(|tcgen05| { + tcgen05.source_contract + == Tcgen05SourceContract::LlvmCustomLoweringWithoutSelection + })) + || (policy.family == "scalar_math" + && policy.scalar_math.as_ref().is_some_and(|sm| { + // Tanh is absent here only because it is PTX-native (no + // imported declaration), so this check never sees it. + matches!( + sm.operation, + ScalarMathOperation::Sin + | ScalarMathOperation::Cos + | ScalarMathOperation::Ex2 + | ScalarMathOperation::Lg2 + | ScalarMathOperation::Rsqrt + ) + })); + ensure!( + !declaration.selections.is_empty() || selectionless_closed_family, + "{} has a declaration but no NVPTX TableGen selection record", + policy.id + ); + let mut matching_selections = Vec::new(); + for selection in &declaration.selections { + if selection_matches_policy(policy, selection)? { + matching_selections.push(selection); + } + } + let expected_selection_count = match policy.family.as_str() { + "vote" | "warp_barrier" | "elect" => 2, + "warp_match" => 4, + "warp_shuffle" => 8, + "counted_barrier" => 4, + "packed_conversion" | "prmt" | "stmatrix" => 0, + "tma" + if policy.tma.as_ref().is_some_and(|tma| { + matches!( + tma.operation, + TmaOperation::Reduce + | TmaOperation::PrefetchTensorMap + | TmaOperation::ReplaceBoxDim + | TmaOperation::ReplaceElementStride + | TmaOperation::ReplaceElementType + | TmaOperation::ReplaceFillMode + | TmaOperation::ReplaceGlobalAddress + | TmaOperation::ReplaceGlobalDim + | TmaOperation::ReplaceGlobalStride + | TmaOperation::ReplaceInterleaveLayout + | TmaOperation::ReplaceRank + | TmaOperation::ReplaceSwizzleAtomicity + | TmaOperation::ReplaceSwizzleMode + ) + }) => + { + 0 + } + "tcgen05" + if policy + .tcgen05 + .as_ref() + .and_then(|tcgen05| tcgen05.mma.as_ref()) + .is_some() => + { + let mma = policy.tcgen05.as_ref().unwrap().mma.as_ref().unwrap(); + if mma.alias.is_some() { + 1 + } else if tcgen05_mma_is_ws(mma.form) { + 64 + } else if tcgen05_mma_is_ashift(mma.form) { + 16 + } else { + 32 + } + } + "tcgen05" + if policy.tcgen05.as_ref().is_some_and(|tcgen05| { + tcgen05.source_contract != Tcgen05SourceContract::ExactTablegenSelection + }) => + { + 0 + } + "clc" + if policy.clc.as_ref().is_some_and(|clc| { + matches!( + clc.operation, + ClcOperation::QueryIsCanceled + | ClcOperation::QueryGetFirstCtaidX + | ClcOperation::QueryGetFirstCtaidY + | ClcOperation::QueryGetFirstCtaidZ + ) + }) => + { + 0 + } + "sreg" + if policy.special_register.as_ref().is_some_and(|special| { + matches!( + special.register, + SpecialRegisterKind::Envreg1 | SpecialRegisterKind::Envreg2 + ) + }) => + { + 0 + } + "mbarrier_extended" + if policy.mbarrier_extended.as_ref().is_some_and(|mbarrier| { + matches!( + mbarrier.operation, + MbarrierExtendedOperation::ArriveExpectTxCta + | MbarrierExtendedOperation::ArriveExpectTxCluster + | MbarrierExtendedOperation::TryWaitParityCta + | MbarrierExtendedOperation::TryWaitParityCluster + ) + }) => + { + 0 + } + "mbarrier_extended" if policy.id == "nanosleep" => 2, + "cp_async_copy" + if policy + .cp_async_copy + .as_ref() + .is_some_and(|copy| copy.source_size == CpAsyncSourceSize::Runtime) => + { + 2 + } + "cluster_memory" if policy.id == "map_shared_rank" => 2, + "scalar_math" + if policy.scalar_math.as_ref().is_some_and(|sm| { + matches!( + sm.operation, + ScalarMathOperation::Sin + | ScalarMathOperation::Cos + | ScalarMathOperation::Ex2 + | ScalarMathOperation::Lg2 + | ScalarMathOperation::Rsqrt + ) + }) => + { + 0 + } + _ => 1, + }; + ensure!( + matching_selections.len() == expected_selection_count, + "{} expected PTX {:?} does not agree with its closed imported selection set", + policy.id, + policy.expected_ptx + ); + for selection in matching_selections { + validate_selected_target_predicates(policy, selection)?; + } + } + Ok(()) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/targets.rs b/crates/cuda-intrinsics-gen/src/resolve/targets.rs new file mode 100644 index 0000000000..9289264d70 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/targets.rs @@ -0,0 +1,852 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + CatalogHardwareAlternative, CatalogHardwareTarget, CatalogTargetAlternative, + CatalogTargetContract, CatalogTargetRequirement, ClcOperation, ExecutionControlOperation, + IntrinsicBackend, LdmatrixElement, LdmatrixShape, OverlayIntrinsic, PtxVersion, + SparseMmaAccumulator, TargetContract, TargetSelectorBinding, Tcgen05Operation, TmaOperation, +}; +use anyhow::{Context, Result, bail, ensure}; + +use super::families::*; + +pub(super) fn validate_selected_target_predicates( + policy: &OverlayIntrinsic, + selection: &crate::model::ImportedSelection, +) -> Result<()> { + if policy.family == "wgmma_control" { + ensure!( + selection.predicates == ["Subtarget->getPTXVersion() >= 80", "hasSM90a"] + && parse_ptx_version(&policy.minimum_ptx, &policy.id)?.encoded() == 80 + && parse_hardware_target(policy)? + == CatalogHardwareTarget::AnyOf { + alternatives: vec![CatalogHardwareAlternative::ExactArchitecture { + sm: 90, + }], + }, + "{} WGMMA-control selection must retain the PTX 8.0 and sm_90a gates", + policy.id + ); + return Ok(()); + } + if policy.family == "register_control" { + let operation = ExecutionControlOperation::from_catalog_id(&policy.id) + .with_context(|| format!("{} has no closed register-control operation", policy.id))?; + ensure!( + matches!( + operation, + ExecutionControlOperation::SetMaxNRegInc | ExecutionControlOperation::SetMaxNRegDec + ) && selection.predicates == ["Subtarget->hasSetMaxNRegSupport()"] + && parse_ptx_version(&policy.minimum_ptx, &policy.id)?.encoded() == 80 + && policy.minimum_sm.is_none() + && policy.targets == TENSOR_MAP_REPLACE_TARGETS, + "{} setmaxnreg selection must retain the helper predicate and reviewed PTX 8.0 target matrix", + policy.id + ); + return Ok(()); + } + + let mma_family = matches!(policy.family.as_str(), "register_mma" | "sparse_mma"); + let tcgen05_mma = policy + .tcgen05 + .as_ref() + .and_then(|tcgen05| tcgen05.mma.as_ref()); + let mut imported_ptx: Option = None; + let mut imported_sm: Option = None; + let mut has_dot_instructions = false; + let mut has_clc_multicast_support = false; + let mut has_tma_blackwell_support = false; + let mut has_tcgen05_support = false; + let mut has_tcgen05_shift_support = false; + let mut has_tcgen05_mma_i8_support = false; + let mut has_ldstmatrix_blackwell_support = false; + let mut has_mma_block_scale_support = false; + let mut has_redux_sync_f32_support = false; + for predicate in &selection.predicates { + if let Some(value) = predicate.strip_prefix("Subtarget->getPTXVersion() >= ") { + let value = value.parse::().with_context(|| { + format!("{} has malformed PTX predicate {predicate:?}", policy.id) + })?; + ensure!( + imported_ptx.is_none() || mma_family, + "{} has duplicate PTX predicates", + policy.id + ); + imported_ptx = Some(imported_ptx.unwrap_or_default().max(value)); + } else if let Some(value) = predicate.strip_prefix("Subtarget->getSmVersion() >= ") { + let value = value.parse::().with_context(|| { + format!("{} has malformed SM predicate {predicate:?}", policy.id) + })?; + ensure!( + imported_sm.is_none() || mma_family, + "{} has duplicate SM predicates", + policy.id + ); + imported_sm = Some(imported_sm.unwrap_or_default().max(value)); + } else if predicate == "hasDotInstructions" { + ensure!( + policy.family == "dotprod", + "{} selected instruction uses dot-product target gating outside the dotprod family", + policy.id + ); + ensure!( + !has_dot_instructions && imported_ptx.is_none() && imported_sm.is_none(), + "{} has duplicate or conflicting dot-product target predicates", + policy.id + ); + has_dot_instructions = true; + imported_ptx = Some(50); + imported_sm = Some(61); + } else if predicate == "Subtarget->hasClusterLaunchControlTryCancelMulticastSupport()" { + ensure!( + policy.family == "clc" + && policy + .clc + .as_ref() + .is_some_and(|clc| { clc.operation == ClcOperation::TryCancelMulticast }), + "{} uses the CLC multicast target predicate outside that operation", + policy.id + ); + ensure!( + !has_clc_multicast_support + && imported_ptx.is_none() + && imported_sm.is_none() + && !has_dot_instructions, + "{} has duplicate or conflicting CLC multicast target predicates", + policy.id + ); + has_clc_multicast_support = true; + } else if predicate == "Subtarget->hasTMABlackwellSupport()" { + ensure!( + policy.family == "tma" + && policy.tma.as_ref().is_some_and(|tma| { + matches!( + tma.operation, + TmaOperation::G2sTile2dMulticastCg2 + | TmaOperation::PrefetchTileGather4TwoDimensional + | TmaOperation::PrefetchTileGather4TwoDimensionalCacheHint + ) + }), + "{} uses the Blackwell TMA target predicate outside the cta_group::2 operation", + policy.id + ); + ensure!( + !has_tma_blackwell_support + && imported_ptx.is_none() + && imported_sm.is_none() + && !has_dot_instructions + && !has_clc_multicast_support, + "{} has duplicate or conflicting Blackwell TMA target predicates", + policy.id + ); + has_tma_blackwell_support = true; + } else if predicate == "Subtarget->hasTcgen05InstSupport()" { + ensure!( + policy.family == "tcgen05" && policy.tcgen05.is_some(), + "{} uses the tcgen05 target predicate outside that family", + policy.id + ); + ensure!( + !has_tcgen05_support + && !has_tcgen05_shift_support + && !has_tcgen05_mma_i8_support + && imported_ptx.is_none() + && imported_sm.is_none() + && !has_dot_instructions + && !has_clc_multicast_support + && !has_tma_blackwell_support, + "{} has duplicate or conflicting tcgen05 target predicates", + policy.id + ); + has_tcgen05_support = true; + } else if predicate == "Subtarget->hasTcgen05MMAI8Kind()" { + ensure!( + tcgen05_mma.is_some() && selection.asm.contains(".kind::i8."), + "{} uses the tcgen05 I8 predicate outside an I8 MMA selection", + policy.id + ); + ensure!( + !has_tcgen05_mma_i8_support + && !has_tcgen05_support + && !has_tcgen05_shift_support + && imported_ptx.is_none() + && imported_sm.is_none() + && !has_dot_instructions + && !has_clc_multicast_support + && !has_tma_blackwell_support, + "{} has duplicate or conflicting tcgen05 I8 predicates", + policy.id + ); + has_tcgen05_mma_i8_support = true; + } else if predicate == "Subtarget->hasTcgen05ShiftSupport()" { + ensure!( + policy.family == "tcgen05" + && policy.tcgen05.as_ref().is_some_and(|tcgen05| { + matches!( + tcgen05.operation, + Tcgen05Operation::ShiftDown | Tcgen05Operation::ShiftDownCg2 + ) + }), + "{} uses the tcgen05 shift target predicate outside a shift operation", + policy.id + ); + ensure!( + !has_tcgen05_shift_support + && !has_tcgen05_support + && !has_tcgen05_mma_i8_support + && imported_ptx.is_none() + && imported_sm.is_none() + && !has_dot_instructions + && !has_clc_multicast_support + && !has_tma_blackwell_support, + "{} has duplicate or conflicting tcgen05 shift target predicates", + policy.id + ); + has_tcgen05_shift_support = true; + } else if predicate == "Subtarget->hasLdStmatrixBlackwellSupport()" { + ensure!( + policy.family == "ldmatrix" + && policy.ldmatrix_variant.as_ref().is_some_and(|variant| { + variant.shape != LdmatrixShape::M8n8 + && variant.element != LdmatrixElement::B16 + }), + "{} uses the Blackwell ld/stmatrix predicate outside a Blackwell ldmatrix variant", + policy.id + ); + ensure!( + !has_ldstmatrix_blackwell_support + && imported_ptx.is_none() + && imported_sm.is_none() + && !has_dot_instructions + && !has_clc_multicast_support + && !has_tma_blackwell_support + && !has_tcgen05_support + && !has_tcgen05_shift_support, + "{} has duplicate or conflicting Blackwell ldmatrix target predicates", + policy.id + ); + has_ldstmatrix_blackwell_support = true; + } else if predicate == "Subtarget->hasMMABlockScale()" { + ensure!( + matches!(policy.family.as_str(), "register_mma" | "sparse_mma") + && (policy.register_mma.is_some() || policy.sparse_mma.is_some()), + "{} uses the MMA block-scale target predicate outside a closed MMA family", + policy.id + ); + ensure!( + !has_mma_block_scale_support + && imported_ptx.is_none() + && imported_sm.is_none() + && !has_dot_instructions + && !has_clc_multicast_support + && !has_tma_blackwell_support + && !has_tcgen05_support + && !has_tcgen05_shift_support + && !has_ldstmatrix_blackwell_support, + "{} has duplicate or conflicting MMA block-scale target predicates", + policy.id + ); + has_mma_block_scale_support = true; + } else if predicate == "Subtarget->hasReduxSyncF32()" { + ensure!( + policy.family == "redux" + && policy + .redux + .as_ref() + .is_some_and(|redux| redux_recipe(redux.operation).value_type == "f32"), + "{} uses the f32 redux target predicate outside an f32 redux operation", + policy.id + ); + ensure!( + !has_redux_sync_f32_support && imported_ptx.is_none() && imported_sm.is_none(), + "{} has duplicate or conflicting f32 redux target predicates", + policy.id + ); + has_redux_sync_f32_support = true; + } else { + bail!( + "{} selected instruction has unsupported target predicate {predicate:?}; target gates must fail closed", + policy.id + ); + } + } + let overlay_ptx = parse_ptx_version(&policy.minimum_ptx, &policy.id)?.encoded(); + if let Some(mma) = tcgen05_mma { + let i8_selection = selection.asm.contains(".kind::i8."); + ensure!( + selection.predicates.len() == 1 + && ((i8_selection && has_tcgen05_mma_i8_support && !has_tcgen05_support) + || (!i8_selection && has_tcgen05_support && !has_tcgen05_mma_i8_support)) + && !has_tcgen05_shift_support, + "{} tcgen05 MMA selection does not retain its exact kind predicate", + policy.id + ); + let contracts = expected_tcgen05_mma_target_contracts(IntrinsicBackend::LlvmNvptx); + let expected_target = if let Some(fixed) = mma.fixed_selectors { + resolve_target_contract( + "tcgen05 MMA selected predicate", + &[TargetSelectorBinding { + name: "kind".into(), + value: tcgen05_mma_kind_name(fixed.kind).into(), + }], + &contracts, + )? + } else { + resolve_target_contracts("tcgen05 MMA selected predicate", &contracts)? + }; + ensure!( + mma.llvm_target == expected_target, + "{} tcgen05 MMA predicate does not map to its closed LLVM target matrix", + policy.id + ); + return Ok(()); + } + if has_mma_block_scale_support { + let (minimum_ptx_matches, target_matches) = match policy.family.as_str() { + "register_mma" => ( + overlay_ptx == 87, + policy.targets == "sm_120a|sm_120f|sm_121a|sm_121f" && policy.minimum_sm.is_none(), + ), + "sparse_mma" => ( + overlay_ptx == 87, + policy.minimum_sm.is_none() + && policy.sparse_mma.as_ref().is_some_and(|mma| { + policy.targets + == match mma.accumulator { + SparseMmaAccumulator::F16 | SparseMmaAccumulator::F32 => { + SPARSE_MMA_F8F6F4_TARGETS + } + SparseMmaAccumulator::S32 => return false, + } + }), + ), + _ => (false, false), + }; + ensure!( + minimum_ptx_matches && target_matches, + "{} MMA block-scale predicate requires its reviewed Blackwell target matrix", + policy.id + ); + return Ok(()); + } + if has_redux_sync_f32_support { + ensure!( + selection.predicates.len() == 1 + && overlay_ptx == 86 + && policy.targets == REDUX_F32_TARGETS + && policy.minimum_sm.is_none(), + "{} f32 redux selection must carry only the hasReduxSyncF32 predicate and its reviewed Blackwell target matrix", + policy.id + ); + return Ok(()); + } + // MMA uses reviewed inline PTX. Its imported predicates gate LLVM's typed + // selection, while the closed recipe and terminal evidence set the native + // PTX floor. + if mma_family { + return Ok(()); + } + if let Some(imported_ptx) = imported_ptx { + ensure!( + overlay_ptx == imported_ptx, + "{} minimum PTX {} disagrees with selected instruction predicate PTX {}", + policy.id, + policy.minimum_ptx, + format_args!("{}.{}", imported_ptx / 10, imported_ptx % 10) + ); + } + if let Some(imported_sm) = imported_sm { + if let Some(packed) = &policy.packed_alu { + ensure!( + packed.native_minimum_sm == imported_sm, + "{} native minimum SM {} disagrees with selected instruction predicate sm_{}", + policy.id, + packed.native_minimum_sm, + imported_sm + ); + } else { + let overlay_target = parse_hardware_target(policy)?; + let target_matches = overlay_target + == CatalogHardwareTarget::AnyOf { + alternatives: vec![CatalogHardwareAlternative::MinimumSm { sm: imported_sm }], + }; + ensure!( + target_matches, + "{} minimum SM {:?} disagrees with selected instruction predicate sm_{}", + policy.id, + policy.minimum_sm, + imported_sm + ); + } + } + if policy.family == "ldmatrix" { + if policy + .ldmatrix_variant + .as_ref() + .is_some_and(|variant| variant.shape == LdmatrixShape::M8n8) + { + ensure!( + imported_ptx.is_some() && imported_sm.is_some() && selection.predicates.len() == 2, + "{} classic ldmatrix selection must carry exactly its PTX and SM predicates", + policy.id + ); + } else { + ensure!( + has_ldstmatrix_blackwell_support + && selection.predicates.len() == 1 + && policy.targets == BLACKWELL_LDMATRIX_LLVM_TARGETS, + "{} Blackwell ldmatrix selection must retain its helper predicate and reviewed exact targets", + policy.id + ); + } + } else if policy.family == "dotprod" { + ensure!( + has_dot_instructions && selection.predicates.len() == 1, + "{} dotprod selection must carry only the hasDotInstructions predicate", + policy.id + ); + } else if policy.family == "clc" { + match policy.clc.as_ref().map(|clc| clc.operation) { + Some(ClcOperation::TryCancel) => ensure!( + imported_ptx.is_some() && imported_sm.is_some() && selection.predicates.len() == 2, + "{} selection must carry exactly its PTX and SM predicates", + policy.id + ), + Some(ClcOperation::TryCancelMulticast) => ensure!( + has_clc_multicast_support + && selection.predicates.len() == 1 + && parse_hardware_target(policy)? + == CatalogHardwareTarget::AnyOf { + alternatives: vec![ + CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 101 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 103 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 120 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 121 }, + ], + }, + "{} multicast target predicate must map to the reviewed exact architectures", + policy.id + ), + _ => bail!( + "{} query operation unexpectedly has an instruction selection", + policy.id + ), + } + } else if policy.family == "tma" { + if policy.tma.as_ref().is_some_and(|tma| { + matches!( + tma.operation, + TmaOperation::G2sTile2dMulticastCg2 + | TmaOperation::PrefetchTileGather4TwoDimensional + | TmaOperation::PrefetchTileGather4TwoDimensionalCacheHint + ) + }) { + ensure!( + has_tma_blackwell_support + && selection.predicates.len() == 1 + && parse_hardware_target(policy)? + == CatalogHardwareTarget::AnyOf { + alternatives: vec![ + CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 101 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 103 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, + ], + }, + "{} Blackwell TMA predicate must map to the reviewed exact architectures", + policy.id + ); + } else { + ensure!( + imported_ptx.is_some() && imported_sm.is_some() && selection.predicates.len() == 2, + "{} TMA selection must carry exactly its PTX and SM predicates", + policy.id + ); + } + } else if policy.family == "tcgen05" { + let shift = policy.tcgen05.as_ref().is_some_and(|tcgen05| { + matches!( + tcgen05.operation, + Tcgen05Operation::ShiftDown | Tcgen05Operation::ShiftDownCg2 + ) + }); + ensure!( + ((shift && has_tcgen05_shift_support && !has_tcgen05_support) + || (!shift && has_tcgen05_support && !has_tcgen05_shift_support)) + && selection.predicates.len() == 1 + && parse_hardware_target(policy)? + == CatalogHardwareTarget::AnyOf { + alternatives: vec![ + CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 101 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 103 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, + ], + }, + "{} tcgen05 selection must use the reviewed accelerated architectures", + policy.id + ); + } else if matches!( + policy.family.as_str(), + "vote" + | "active_mask" + | "warp_match" + | "elect" + | "warp_barrier" + | "warp_shuffle" + | "cp_async_copy" + | "cp_async_control" + | "mbarrier_basic" + | "cluster_barrier" + | "scalar_conversion" + | "grid_dependency" + ) { + ensure!( + imported_ptx.is_some() && imported_sm.is_some() && selection.predicates.len() == 2, + "{} selection must carry exactly its PTX and SM predicates", + policy.id + ); + } + Ok(()) +} + +pub(super) fn parse_ptx_version(value: &str, intrinsic_id: &str) -> Result { + value + .parse() + .map_err(|reason: String| anyhow::anyhow!("{intrinsic_id} minimum_ptx {value:?}: {reason}")) +} + +pub(super) fn parse_hardware_target(policy: &OverlayIntrinsic) -> Result { + parse_hardware_target_fields(&policy.id, &policy.targets, policy.minimum_sm.as_deref()) +} + +/// Resolve every selector contract without merging target sets. +pub(crate) fn resolve_target_contracts( + intrinsic_id: &str, + contracts: &[TargetContract], +) -> Result { + ensure!( + !contracts.is_empty(), + "{intrinsic_id} has no reviewed target contracts" + ); + + let expected_selector_names = contracts[0] + .selectors + .iter() + .map(|selector| selector.name.as_str()) + .collect::>(); + let mut resolved_contracts = Vec::with_capacity(contracts.len()); + for contract in contracts { + validate_selector_bindings( + intrinsic_id, + "target-contract selectors", + &contract.selectors, + )?; + ensure!( + contract + .selectors + .iter() + .map(|selector| selector.name.as_str()) + .eq(expected_selector_names.iter().copied()), + "{intrinsic_id} target contracts must use one closed selector schema" + ); + ensure!( + !contract.alternatives.is_empty(), + "{intrinsic_id} target contract {:?} has no alternatives", + contract.selectors + ); + + let alternatives = contract + .alternatives + .iter() + .map(|alternative| { + Ok(CatalogTargetAlternative { + minimum_ptx: parse_ptx_version(&alternative.minimum_ptx, intrinsic_id)?, + hardware: parse_target_contract_hardware(intrinsic_id, &alternative.target)?, + }) + }) + .collect::>>()?; + ensure!( + alternatives + .windows(2) + .all(|pair| target_hardware_sort_key(pair[0].hardware) + < target_hardware_sort_key(pair[1].hardware)), + "{intrinsic_id} target contract {:?} alternatives must have unique, sorted hardware targets", + contract.selectors + ); + let minimum_count = alternatives + .iter() + .filter(|alternative| { + matches!( + alternative.hardware, + CatalogHardwareAlternative::MinimumSm { .. } + ) + }) + .count(); + ensure!( + minimum_count == 0 || alternatives.len() == 1, + "{intrinsic_id} target contract {:?} cannot mix a minimum-SM range with other alternatives", + contract.selectors + ); + resolved_contracts.push(CatalogTargetContract { + selectors: contract.selectors.clone(), + alternatives, + }); + } + ensure!( + resolved_contracts + .windows(2) + .all(|pair| pair[0].selectors < pair[1].selectors), + "{intrinsic_id} target contracts must have unique, sorted selector bindings" + ); + + let minimum_ptx = resolved_contracts + .iter() + .flat_map(|contract| contract.alternatives.iter()) + .map(|alternative| alternative.minimum_ptx) + .min() + .unwrap(); + Ok(CatalogTargetRequirement { + minimum_ptx, + hardware: CatalogHardwareTarget::TargetMatrix { + contracts: resolved_contracts, + }, + }) +} + +/// Resolve one fixed selector tuple from a closed target matrix. +pub(crate) fn resolve_target_contract( + intrinsic_id: &str, + selected: &[TargetSelectorBinding], + contracts: &[TargetContract], +) -> Result { + validate_selector_bindings(intrinsic_id, "selected target selectors", selected)?; + let resolved = resolve_target_contracts(intrinsic_id, contracts)?; + let CatalogHardwareTarget::TargetMatrix { contracts } = resolved.hardware else { + unreachable!("target-contract resolution always returns a matrix") + }; + + let matching = contracts + .into_iter() + .filter(|contract| contract.selectors == selected) + .collect::>(); + ensure!( + matching.len() == 1, + "{intrinsic_id} selected target selectors {:?} must match exactly one reviewed contract", + selected + ); + let contract = matching.into_iter().next().unwrap(); + let minimum_ptx = contract + .alternatives + .iter() + .map(|alternative| alternative.minimum_ptx) + .min() + .unwrap(); + Ok(CatalogTargetRequirement { + minimum_ptx, + hardware: CatalogHardwareTarget::TargetMatrix { + contracts: vec![contract], + }, + }) +} + +pub(super) fn validate_selector_bindings( + intrinsic_id: &str, + label: &str, + selectors: &[TargetSelectorBinding], +) -> Result<()> { + for selector in selectors { + ensure!( + is_target_selector_token(&selector.name) && is_target_selector_token(&selector.value), + "{intrinsic_id} {label} must use lowercase snake-case names and values" + ); + } + ensure!( + selectors.windows(2).all(|pair| pair[0].name < pair[1].name), + "{intrinsic_id} {label} must have unique, sorted names" + ); + Ok(()) +} + +pub(super) fn is_target_selector_token(value: &str) -> bool { + value + .bytes() + .next() + .is_some_and(|byte| byte.is_ascii_lowercase()) + && value + .bytes() + .all(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'_') + && !value.ends_with('_') + && !value.contains("__") +} + +pub(super) fn parse_target_contract_hardware( + intrinsic_id: &str, + target: &str, +) -> Result { + if let Some(minimum) = target.strip_suffix('+') { + let sm = parse_sm_spelling(intrinsic_id, "target contract", minimum, None)?; + return Ok(CatalogHardwareAlternative::MinimumSm { sm }); + } + parse_exact_hardware_alternative(intrinsic_id, target) +} + +pub(super) fn target_hardware_sort_key(hardware: CatalogHardwareAlternative) -> (u16, u8) { + match hardware { + CatalogHardwareAlternative::MinimumSm { sm } => (sm, 0), + CatalogHardwareAlternative::ExactArchitecture { sm } => (sm, 1), + CatalogHardwareAlternative::FamilyTarget { sm } => (sm, 2), + } +} + +pub(super) fn parse_hardware_target_fields( + intrinsic_id: &str, + targets: &str, + minimum_sm: Option<&str>, +) -> Result { + if targets == "all" { + let Some(minimum_sm) = minimum_sm else { + return Ok(CatalogHardwareTarget::All); + }; + let sm = parse_sm_spelling(intrinsic_id, "minimum_sm", minimum_sm, None)?; + return Ok(CatalogHardwareTarget::AnyOf { + alternatives: vec![CatalogHardwareAlternative::MinimumSm { sm }], + }); + } + + if targets.contains('|') { + ensure!( + minimum_sm.is_none(), + "{} target alternatives {:?} cannot be combined with minimum_sm", + intrinsic_id, + targets + ); + let spellings = targets.split('|').collect::>(); + ensure!( + spellings.len() >= 2, + "{} target alternatives must contain at least two targets", + intrinsic_id + ); + ensure!( + spellings.windows(2).all(|pair| pair[0] < pair[1]), + "{} target alternatives must be unique and sorted", + intrinsic_id + ); + return Ok(CatalogHardwareTarget::AnyOf { + alternatives: spellings + .into_iter() + .map(|spelling| parse_exact_hardware_alternative(intrinsic_id, spelling)) + .collect::>>()?, + }); + } + + ensure!( + minimum_sm.is_none(), + "{} exact targets {:?} cannot be combined with minimum_sm", + intrinsic_id, + targets + ); + Ok(CatalogHardwareTarget::AnyOf { + alternatives: vec![parse_exact_hardware_alternative(intrinsic_id, targets)?], + }) +} + +pub(super) fn parse_exact_hardware_alternative( + intrinsic_id: &str, + target: &str, +) -> Result { + let suffix = target + .chars() + .last() + .filter(|suffix| matches!(suffix, 'a' | 'f')); + let Some(suffix) = suffix else { + bail!( + "{} targets {:?} must be `all`, exact `sm_Na`, or family `sm_Nf`", + intrinsic_id, + target + ); + }; + let sm = parse_sm_spelling(intrinsic_id, "targets", target, Some(suffix))?; + Ok(match suffix { + 'a' => CatalogHardwareAlternative::ExactArchitecture { sm }, + 'f' => CatalogHardwareAlternative::FamilyTarget { sm }, + _ => unreachable!(), + }) +} + +pub(super) fn parse_sm_spelling( + intrinsic_id: &str, + field: &str, + value: &str, + suffix: Option, +) -> Result { + let body = value.strip_prefix("sm_").with_context(|| { + format!("{intrinsic_id} {field} {value:?} must use canonical sm_NN spelling") + })?; + let digits = match suffix { + Some(suffix) => body.strip_suffix(suffix).with_context(|| { + format!("{intrinsic_id} {field} {value:?} has the wrong target suffix") + })?, + None => body, + }; + ensure!( + matches!(digits.len(), 2 | 3) && digits.bytes().all(|byte| byte.is_ascii_digit()), + "{} {} {:?} must use canonical sm_NN{} spelling", + intrinsic_id, + field, + value, + suffix.map_or("", |suffix| if suffix == 'a' { "a" } else { "f" }) + ); + let sm: u16 = digits + .parse() + .with_context(|| format!("{intrinsic_id} {field} target is too large"))?; + let canonical = match suffix { + Some(suffix) => format!("sm_{sm}{suffix}"), + None => format!("sm_{sm}"), + }; + ensure!( + sm > 0 && canonical == value, + "{} {} {:?} is not canonical", + intrinsic_id, + field, + value + ); + Ok(sm) +} + +pub(super) fn backend_target_requirement( + policy: &OverlayIntrinsic, + lowering: &crate::model::OverlayBackendLowering, +) -> Result { + if let Some(mma) = policy + .tcgen05 + .as_ref() + .and_then(|tcgen05| tcgen05.mma.as_ref()) + { + return Ok(match lowering.backend { + IntrinsicBackend::LlvmNvptx => mma.llvm_target.clone(), + IntrinsicBackend::LibNvvm => mma.libnvvm_target.clone(), + }); + } + let minimum_ptx = lowering + .minimum_ptx + .as_deref() + .unwrap_or(&policy.minimum_ptx); + let targets = lowering.targets.as_deref().unwrap_or(&policy.targets); + let minimum_sm = lowering.minimum_sm.as_deref().or_else(|| { + if lowering.targets.is_none() { + policy.minimum_sm.as_deref() + } else { + None + } + }); + Ok(CatalogTargetRequirement { + minimum_ptx: parse_ptx_version(minimum_ptx, &policy.id)?, + hardware: parse_hardware_target_fields(&policy.id, targets, minimum_sm)?, + }) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/tests/evidence.rs b/crates/cuda-intrinsics-gen/src/resolve/tests/evidence.rs new file mode 100644 index 0000000000..c3ca1bead2 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/tests/evidence.rs @@ -0,0 +1,1129 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, CatalogHardwareAlternative, CatalogTargetAlternative, + DotProductOperation, DotProductSignedness, EvidenceArtifactKind, EvidenceFile, EvidenceRecord, + EvidenceStageKind, IntrinsicBackend, LdmatrixElement, LdmatrixLayout, LdmatrixMultiplicity, + LdmatrixShape, LdmatrixStateSpace, OverlayBackendLowering, PackedConversionDestinationFormat, + PackedConversionRounding, PackedConversionSaturation, RuntimeValidation, +}; +use crate::ptx::{InstructionPattern, OperandPattern}; + +use super::fixtures::*; +use crate::resolve::evidence::*; +use crate::resolve::families::*; + +#[test] +fn legacy_evidence_schema_is_unchanged_and_rejects_matrix_fields() { + let legacy = EvidenceFile { + schema: 5, + backend_profile: "legacy".into(), + backend_kind: None, + llvm_revision: "test".into(), + backend_version: "LLVM legacy test".into(), + backend_sha256: "0123456789abcdef".into(), + artifact_path: None, + build_id_prefix: None, + nvvm_ir_version: None, + debug_ir_version: None, + records: vec![evidence()], + }; + let bytes = serde_json::to_vec(&legacy).unwrap(); + assert_eq!(parse_evidence_bytes(&bytes, "legacy").unwrap(), legacy); + + let mut with_matrix_field = serde_json::to_value(&legacy).unwrap(); + with_matrix_field["matrices"] = serde_json::json!([]); + let error = parse_synthetic_evidence(&with_matrix_field).unwrap_err(); + assert!(error.to_string().contains("legacy evidence")); + assert!(format!("{error:#}").contains("unknown field")); +} + +#[test] +fn compact_evidence_matrix_equals_explicit_records() { + let expanded = parse_synthetic_evidence(&synthetic_matrix_json()).unwrap(); + let mut expected = vec![EvidenceRecord { + id: "synthetic_explicit".into(), + source: None, + source_record: Some("int_synthetic_explicit".into()), + llvm_symbol: Some("llvm.synthetic.explicit".into()), + resolved_llvm_symbol: None, + llvm_arguments: vec!["i32".into()], + llvm_results: vec!["i32".into()], + concrete_llvm_arguments: vec![], + concrete_llvm_results: vec![], + target_triple: "nvptx64-nvidia-cuda".into(), + gpu_target: "sm_80".into(), + ptx_feature: "+ptx71".into(), + status: "lowered".into(), + stages: vec![], + declaration_attributes_canonicalized: None, + runtime_validation: None, + expected_ptx: InstructionPattern { + mnemonic: "mma".into(), + modifiers: vec!["sync".into(), "explicit".into()], + operands: vec![OperandPattern::Register], + }, + }]; + expected.extend(["s8", "u8"].into_iter().map(|element| EvidenceRecord { + id: format!("synthetic_{element}"), + source: None, + source_record: Some(format!("int_synthetic_{element}")), + llvm_symbol: Some(format!("llvm.synthetic.{element}")), + resolved_llvm_symbol: None, + llvm_arguments: vec!["i32".into()], + llvm_results: vec!["i32".into()], + concrete_llvm_arguments: vec![], + concrete_llvm_results: vec![], + target_triple: "nvptx64-nvidia-cuda".into(), + gpu_target: "sm_80".into(), + ptx_feature: "+ptx71".into(), + status: "lowered".into(), + stages: vec![shared_matrix_stage()], + declaration_attributes_canonicalized: None, + runtime_validation: None, + expected_ptx: InstructionPattern { + mnemonic: "mma".into(), + modifiers: vec!["sync".into(), element.into()], + operands: vec![OperandPattern::Register], + }, + })); + assert_eq!(expanded.schema, 6); + assert_eq!(expanded.records, expected); +} + +#[test] +fn matrix_identity_mutations_reach_existing_evidence_validation() { + let mut expanded = parse_synthetic_evidence(&policy_matrix_json()).unwrap(); + let record = expanded.records.pop().unwrap(); + validate_test_evidence(&policy(), record.clone()).unwrap(); + + let mut wrong_source = record.clone(); + wrong_source.source_record = Some("int_nvvm_read_ptx_sreg_tid_y".into()); + assert!( + validate_test_evidence(&policy(), wrong_source) + .unwrap_err() + .to_string() + .contains("source provenance mismatch") + ); + + let mut wrong_symbol = record.clone(); + wrong_symbol.llvm_symbol = Some("llvm.nvvm.read.ptx.sreg.tid.y".into()); + assert!( + validate_test_evidence(&policy(), wrong_symbol) + .unwrap_err() + .to_string() + .contains("signature mismatch") + ); + + let mut wrong_signature = record.clone(); + wrong_signature.llvm_arguments.push("i32".into()); + assert!( + validate_test_evidence(&policy(), wrong_signature) + .unwrap_err() + .to_string() + .contains("signature mismatch") + ); + + let mut wrong_ptx = record; + wrong_ptx.expected_ptx.modifiers.push("changed".into()); + assert!( + validate_test_evidence(&policy(), wrong_ptx) + .unwrap_err() + .to_string() + .contains("PTX expectation mismatch") + ); +} + +#[test] +fn evidence_matrix_rejects_bad_counts_fixtures_placeholders_and_collisions() { + let base = synthetic_matrix_json(); + + let mut bad_product = base.clone(); + bad_product["matrices"][0]["product_count"] = 3.into(); + assert_synthetic_evidence_error(&bad_product, "expands to 2 records"); + + let mut unknown_fixture = base.clone(); + unknown_fixture["matrices"][0]["fixtures"][0] = "missing".into(); + assert_synthetic_evidence_error(&unknown_fixture, "unknown fixture"); + + let mut uncovered_fixture = base.clone(); + let extra = uncovered_fixture["fixtures"][0].clone(); + uncovered_fixture["fixtures"] + .as_array_mut() + .unwrap() + .push(extra); + uncovered_fixture["fixtures"][1]["id"] = "unused".into(); + assert_synthetic_evidence_error(&uncovered_fixture, "not referenced"); + + let mut wrong_coverage = base.clone(); + wrong_coverage["fixtures"][0]["coverage_count"] = 1.into(); + assert_synthetic_evidence_error(&wrong_coverage, "covers 2 expanded records"); + + let mut malformed = base.clone(); + malformed["matrices"][0]["template"]["id"] = "synthetic_$element".into(); + assert_synthetic_evidence_error(&malformed, "malformed matrix placeholder"); + + let mut unknown_axis = base.clone(); + unknown_axis["matrices"][0]["template"]["id"] = "synthetic_${other}".into(); + assert_synthetic_evidence_error(&unknown_axis, "unknown matrix axis"); + + let mut collision = base.clone(); + collision["matrices"][0]["template"]["id"] = "synthetic".into(); + assert_synthetic_evidence_error(&collision, "duplicate expanded evidence ID"); +} + +#[test] +fn exact_operand_matrix_placeholders_fail_closed() { + let base = policy_matrix_json(); + + let mut unknown = base.clone(); + unknown["matrices"][0]["template"]["expected_ptx"]["operands"][1]["value"] = + "%tid.${other}".into(); + assert_synthetic_evidence_error(&unknown, "unknown matrix axis other"); + + let mut unterminated = base.clone(); + unterminated["matrices"][0]["template"]["expected_ptx"]["operands"][1]["value"] = + "%tid.${axis".into(); + assert_synthetic_evidence_error(&unterminated, "unterminated matrix placeholder"); + + let mut disallowed = base; + disallowed["matrices"][0]["template"]["expected_ptx"]["mnemonic"] = "mov.${axis}".into(); + assert_synthetic_evidence_error(&disallowed, "PTX mnemonic cannot contain"); +} + +#[test] +fn evidence_matrix_rejects_bad_axes_fixture_ids_and_stage_conflicts() { + let base = synthetic_matrix_json(); + + let mut no_fixture = base.clone(); + no_fixture["matrices"][0]["fixtures"] = serde_json::json!([]); + assert_synthetic_evidence_error(&no_fixture, "references no shared fixture"); + + let mut empty_axes = base.clone(); + empty_axes["matrices"][0]["axes"] = serde_json::json!([]); + assert_synthetic_evidence_error(&empty_axes, "has no axes"); + + let mut duplicate_axis = base.clone(); + let axis = duplicate_axis["matrices"][0]["axes"][0].clone(); + duplicate_axis["matrices"][0]["axes"] + .as_array_mut() + .unwrap() + .push(axis); + duplicate_axis["matrices"][0]["product_count"] = 4.into(); + assert_synthetic_evidence_error(&duplicate_axis, "axes must be unique and sorted"); + + let mut empty_values = base.clone(); + empty_values["matrices"][0]["axes"][0]["values"] = serde_json::json!([]); + assert_synthetic_evidence_error(&empty_values, "has no values"); + + let mut empty_axis_name = base.clone(); + empty_axis_name["matrices"][0]["axes"][0]["name"] = "".into(); + assert_synthetic_evidence_error(&empty_axis_name, "is not a safe token"); + + let mut empty_value = base.clone(); + empty_value["matrices"][0]["axes"][0]["values"][0] = "".into(); + assert_synthetic_evidence_error(&empty_value, "unsafe value"); + + let mut duplicate_value = base.clone(); + duplicate_value["matrices"][0]["axes"][0]["values"][1] = "s8".into(); + assert_synthetic_evidence_error(&duplicate_value, "duplicate value"); + + let mut unsafe_value = base.clone(); + unsafe_value["matrices"][0]["axes"][0]["values"][0] = "../s8".into(); + assert_synthetic_evidence_error(&unsafe_value, "unsafe value"); + + let mut unused_axis = base.clone(); + unused_axis["matrices"][0]["axes"] + .as_array_mut() + .unwrap() + .push(serde_json::json!({"name": "other", "values": ["x"]})); + assert_synthetic_evidence_error(&unused_axis, "unused axis"); + + let mut duplicate_fixture = base.clone(); + let fixture = duplicate_fixture["fixtures"][0].clone(); + duplicate_fixture["fixtures"] + .as_array_mut() + .unwrap() + .push(fixture); + assert_synthetic_evidence_error(&duplicate_fixture, "duplicate evidence fixture ID"); + + let mut fixture_placeholder = base.clone(); + fixture_placeholder["fixtures"][0]["stages"][0]["detail"] = "covers ${element}".into(); + assert_synthetic_evidence_error(&fixture_placeholder, "cannot contain matrix placeholders"); + + let mut missing_symbol = base.clone(); + missing_symbol["matrices"][0]["template"] + .as_object_mut() + .unwrap() + .remove("llvm_symbol"); + assert_synthetic_evidence_error(&missing_symbol, "missing field `llvm_symbol`"); + + let mut conflicting_stage = base; + conflicting_stage["matrices"][0]["template"]["facts"]["stages"] = + conflicting_stage["fixtures"][0]["stages"].clone(); + assert_synthetic_evidence_error(&conflicting_stage, "conflicting duplicate"); +} + +#[test] +fn typed_evidence_accepts_direct_scalar_intrinsic_signatures() { + let policy = dot_product_policy(DotProductOperation::Dp2a, DotProductSignedness::Signed); + let mut record = dot_product_evidence(&policy); + validate_typed_llvm_evidence(&policy, &record).unwrap(); + + record.concrete_llvm_arguments.remove(2); + let error = validate_typed_llvm_evidence(&policy, &record).unwrap_err(); + assert!(error.to_string().contains("resolved signature")); +} + +#[test] +fn typed_evidence_maps_cluster_shared_pointers_to_address_space_seven() { + let mut policy = dot_product_policy(DotProductOperation::Dp2a, DotProductSignedness::Signed); + let mut record = dot_product_evidence(&policy); + policy.llvm_arguments = vec![ + "shared_cluster_ptr".into(), + "shared_ptr".into(), + "ptr".into(), + ]; + record.concrete_llvm_arguments = vec![ + "ptr addrspace(7)".into(), + "ptr addrspace(3)".into(), + "ptr".into(), + ]; + + validate_typed_llvm_evidence(&policy, &record).unwrap(); +} + +#[test] +fn packed_conversion_evidence_separates_llvm_declaration_facts_from_libnvvm() { + for (destination, result) in [ + (PackedConversionDestinationFormat::Bf16x2, "v2bf16"), + (PackedConversionDestinationFormat::F16x2, "v2f16"), + ] { + let policy = packed_conversion_policy( + destination, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::None, + ); + let llvm = policy + .backend_lowerings + .iter() + .find(|lowering| lowering.backend == IntrinsicBackend::LlvmNvptx) + .unwrap(); + let mut record = packed_conversion_evidence(&policy); + record.status = "validated".into(); + record.stages = [ + EvidenceStageKind::DeclarationCanonicalization, + EvidenceStageKind::BackendCodegen, + EvidenceStageKind::PtxAssembly, + ] + .into_iter() + .map(|stage| { + evidence_stage( + stage, + BackendLoweringMechanism::TypedNvvm, + &["sm_80", "ptx70"], + ) + }) + .collect(); + let assembly = record + .stages + .iter_mut() + .find(|stage| stage.stage == EvidenceStageKind::PtxAssembly) + .unwrap(); + assembly.tool_path = Some("/usr/local/cuda/bin/ptxas".into()); + assembly.tool_version = Some("CUDA 13.3 V13.3.33".into()); + assembly.tool_sha256 = + Some("7fdd01a4cf50e30746da98989c9272a907f491e6fd7fecfda14642e4375f88fb".into()); + assert_eq!(record.concrete_llvm_results, [result]); + validate_packed_conversion_backend_evidence(&policy, &record, llvm).unwrap(); + + let mut lowered = record.clone(); + lowered.status = "lowered".into(); + let error = + validate_packed_conversion_backend_evidence(&policy, &lowered, llvm).unwrap_err(); + assert!( + error.to_string().contains("validated evidence status"), + "{error:#}" + ); + + for required in [ + EvidenceStageKind::DeclarationCanonicalization, + EvidenceStageKind::BackendCodegen, + EvidenceStageKind::PtxAssembly, + ] { + let mut missing = record.clone(); + missing.stages.retain(|stage| stage.stage != required); + let error = + validate_packed_conversion_backend_evidence(&policy, &missing, llvm).unwrap_err(); + assert!( + error + .to_string() + .contains("successful auxiliary typed-NVVM"), + "{error:#}" + ); + + let mut failed = record.clone(); + failed + .stages + .iter_mut() + .find(|stage| stage.stage == required) + .unwrap() + .outcome = "failed".into(); + let error = + validate_packed_conversion_backend_evidence(&policy, &failed, llvm).unwrap_err(); + assert!( + error + .to_string() + .contains("successful auxiliary typed-NVVM"), + "{error:#}" + ); + + let mut wrong_mechanism = record.clone(); + wrong_mechanism + .stages + .iter_mut() + .find(|stage| stage.stage == required) + .unwrap() + .mechanism = Some(BackendLoweringMechanism::InlinePtx); + let error = + validate_packed_conversion_backend_evidence(&policy, &wrong_mechanism, llvm) + .unwrap_err(); + assert!( + error + .to_string() + .contains("successful auxiliary typed-NVVM"), + "{error:#}" + ); + } + + let mut missing_tool_identity = record.clone(); + missing_tool_identity + .stages + .iter_mut() + .find(|stage| stage.stage == EvidenceStageKind::PtxAssembly) + .unwrap() + .tool_sha256 = None; + let error = + validate_packed_conversion_backend_evidence(&policy, &missing_tool_identity, llvm) + .unwrap_err(); + assert!( + error.to_string().contains("exact tool identity"), + "{error:#}" + ); + + for stage_kind in [ + EvidenceStageKind::BackendCodegen, + EvidenceStageKind::PtxAssembly, + ] { + let mut wrong_floor = record.clone(); + wrong_floor + .stages + .iter_mut() + .find(|stage| stage.stage == stage_kind) + .unwrap() + .targets = vec!["sm_75".into(), "ptx70".into()]; + let error = validate_packed_conversion_backend_evidence(&policy, &wrong_floor, llvm) + .unwrap_err(); + assert!( + error.to_string().contains("catalog floor sm_80"), + "{error:#}" + ); + } + + record.declaration_attributes_canonicalized = None; + let error = + validate_packed_conversion_backend_evidence(&policy, &record, llvm).unwrap_err(); + assert!( + error + .to_string() + .contains("canonical declaration attributes") + ); + } + + let policy = packed_conversion_policy( + PackedConversionDestinationFormat::Bf16x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::None, + ); + let libnvvm = policy + .backend_lowerings + .iter() + .find(|lowering| lowering.backend == IntrinsicBackend::LibNvvm) + .unwrap(); + let mut record = packed_conversion_evidence(&policy); + record.concrete_llvm_arguments.clear(); + record.concrete_llvm_results.clear(); + record.declaration_attributes_canonicalized = None; + validate_packed_conversion_backend_evidence(&policy, &record, libnvvm).unwrap(); + + record.concrete_llvm_arguments = policy.llvm_arguments.clone(); + let error = validate_packed_conversion_backend_evidence(&policy, &record, libnvvm).unwrap_err(); + assert!( + error + .to_string() + .contains("must not claim typed LLVM support") + ); + + record.concrete_llvm_arguments.clear(); + record.stages.push(evidence_stage( + EvidenceStageKind::BackendCodegen, + BackendLoweringMechanism::TypedNvvm, + &["sm_80", "ptx70"], + )); + let error = validate_packed_conversion_backend_evidence(&policy, &record, libnvvm).unwrap_err(); + assert!( + error + .to_string() + .contains("must not claim typed LLVM support") + ); +} + +#[test] +fn signature_and_evidence_mismatches_are_rejected() { + let mut imported = declaration(); + imported.results = vec!["i64".into()]; + assert!( + validate_imported_policy(&policy(), &imported) + .unwrap_err() + .to_string() + .contains("LLVM result signature mismatch") + ); + + let mut backend_evidence = evidence(); + backend_evidence.llvm_results = vec!["i64".into()]; + assert!( + validate_test_evidence(&policy(), backend_evidence) + .unwrap_err() + .to_string() + .contains("evidence signature mismatch") + ); + + let mut backend_evidence = evidence(); + backend_evidence.expected_ptx = sreg_pattern("%tid.y"); + assert!( + validate_test_evidence(&policy(), backend_evidence) + .unwrap_err() + .to_string() + .contains("evidence PTX expectation mismatch") + ); +} + +#[test] +fn validated_llvm_evidence_requires_exact_ptxas_identity() { + let mut record = evidence(); + record.status = "validated".into(); + record.stages.push(crate::model::EvidenceStage { + targets: vec!["sm_75".into()], + representation: "probe PTX".into(), + stage: EvidenceStageKind::PtxAssembly, + mechanism: Some(BackendLoweringMechanism::TypedNvvm), + outcome: "succeeded".into(), + detail: "accepted".into(), + artifact_kind: None, + tool_path: Some("/usr/local/cuda/bin/ptxas".into()), + tool_version: Some("CUDA 13.3 V13.3.33".into()), + tool_sha256: Some( + "7fdd01a4cf50e30746da98989c9272a907f491e6fd7fecfda14642e4375f88fb".into(), + ), + }); + assert!(has_valid_ptx_assembly_stage( + &record, + BackendLoweringMechanism::TypedNvvm + )); + + let stage = record.stages.last_mut().unwrap(); + stage.tool_path = None; + assert!(!has_valid_ptx_assembly_stage( + &record, + BackendLoweringMechanism::TypedNvvm + )); + record.stages.clear(); + assert!(!has_valid_ptx_assembly_stage( + &record, + BackendLoweringMechanism::TypedNvvm + )); +} + +#[test] +fn validated_libnvvm_evidence_requires_a_real_cubin_terminal() { + let mut record = evidence(); + record.stages.push(crate::model::EvidenceStage { + targets: vec!["sm_90".into(), "ptx78".into()], + representation: "linked output".into(), + stage: EvidenceStageKind::DeviceLink, + mechanism: Some(BackendLoweringMechanism::InlinePtx), + outcome: "succeeded".into(), + detail: "test".into(), + artifact_kind: None, + tool_path: Some("/usr/local/cuda-13.3/lib64/libnvJitLink.so.13.3.33".into()), + tool_version: Some("V13.3.33".into()), + tool_sha256: Some( + "3ba1e744347cd68617b862eccfd98b125482e882b7a6319f42abc9a768513db8".into(), + ), + }); + assert!(!has_valid_cubin_device_link_stage( + &record, + BackendLoweringMechanism::InlinePtx + )); + record.stages[0].artifact_kind = Some(EvidenceArtifactKind::Cubin); + assert!(has_valid_cubin_device_link_stage( + &record, + BackendLoweringMechanism::InlinePtx + )); +} + +fn evidence_stage( + stage: EvidenceStageKind, + mechanism: BackendLoweringMechanism, + targets: &[&str], +) -> crate::model::EvidenceStage { + crate::model::EvidenceStage { + targets: targets.iter().map(|target| (*target).into()).collect(), + representation: "test".into(), + stage, + mechanism: Some(mechanism), + outcome: "succeeded".into(), + detail: "test".into(), + artifact_kind: None, + tool_path: None, + tool_version: None, + tool_sha256: None, + } +} + +#[test] +fn backend_stage_targets_and_executed_status_are_monotonic() { + let mut target_policy = policy(); + target_policy.minimum_ptx = "6.5".into(); + target_policy.minimum_sm = Some("sm_75".into()); + let lowering = crate::model::OverlayBackendLowering { + backend: IntrinsicBackend::LlvmNvptx, + mechanism: BackendLoweringMechanism::TypedNvvm, + evidence_profile: "test".into(), + targets: None, + minimum_ptx: None, + minimum_sm: None, + }; + let mut record = evidence(); + record.status = "validated".into(); + record.runtime_validation = Some(RuntimeValidation::Unexecuted); + record.stages = vec![ + evidence_stage( + EvidenceStageKind::BackendCodegen, + BackendLoweringMechanism::TypedNvvm, + &["sm_75", "ptx65"], + ), + evidence_stage( + EvidenceStageKind::PtxAssembly, + BackendLoweringMechanism::TypedNvvm, + &["sm_75", "ptx65"], + ), + ]; + validate_selected_stage_targets(&target_policy, &record, &lowering).unwrap(); + + record.stages[0].targets = vec!["sm_75a".into(), "ptx65".into()]; + assert!(validate_selected_stage_targets(&target_policy, &record, &lowering).is_err()); + record.stages[0].targets = vec!["sm_75".into(), "ptx65".into()]; + + record.stages[1].targets = vec!["sm_80".into(), "ptx65".into()]; + validate_selected_stage_targets(&target_policy, &record, &lowering).unwrap(); + + record.stages[1].targets = vec!["sm_90a".into(), "ptx65".into()]; + validate_selected_stage_targets(&target_policy, &record, &lowering).unwrap(); + + record.stages[1].targets = vec!["sm_74".into(), "ptx65".into()]; + assert!( + validate_selected_stage_targets(&target_policy, &record, &lowering) + .unwrap_err() + .to_string() + .contains("catalog floor sm_75") + ); + + record.stages[1].targets = vec!["sm_75".into(), "ptx65".into()]; + record.status = "executed".into(); + record.runtime_validation = Some(RuntimeValidation::Executed); + assert!( + validate_selected_stage_targets(&target_policy, &record, &lowering) + .unwrap_err() + .to_string() + .contains("runtime stage") + ); +} + +#[test] +fn exact_and_family_evidence_targets_match_at_every_stage() { + let lowering = crate::model::OverlayBackendLowering { + backend: IntrinsicBackend::LlvmNvptx, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: "test".into(), + targets: None, + minimum_ptx: None, + minimum_sm: None, + }; + + for (target, wrong_targets) in [ + ("sm_120a", ["sm_120", "sm_120f", "sm_121a"]), + ("sm_120f", ["sm_120", "sm_120a", "sm_121f"]), + ] { + let mut target_policy = policy(); + target_policy.minimum_ptx = "8.7".into(); + target_policy.targets = target.into(); + let mut record = evidence(); + record.status = "validated".into(); + record.runtime_validation = Some(RuntimeValidation::Unexecuted); + record.stages = vec![ + evidence_stage( + EvidenceStageKind::BackendCodegen, + BackendLoweringMechanism::InlinePtx, + &[target, "ptx87"], + ), + evidence_stage( + EvidenceStageKind::PtxAssembly, + BackendLoweringMechanism::InlinePtx, + &[target, "ptx87"], + ), + ]; + validate_selected_stage_targets(&target_policy, &record, &lowering).unwrap(); + + for wrong in wrong_targets { + record.stages[0].targets = vec![wrong.into(), "ptx87".into()]; + let error = validate_selected_stage_targets(&target_policy, &record, &lowering) + .unwrap_err() + .to_string(); + assert!(error.contains(target), "{error}"); + } + record.stages[0].targets = vec![target.into(), "ptx87".into()]; + for wrong in wrong_targets { + record.stages[1].targets = vec![wrong.into(), "ptx87".into()]; + let error = validate_selected_stage_targets(&target_policy, &record, &lowering) + .unwrap_err() + .to_string(); + assert!(error.contains(target), "{error}"); + } + } +} + +#[test] +fn suffixed_evidence_target_spellings_are_normalized() { + for target in ["sm_120a", "compute_120a", "sm_120f", "compute_120f"] { + assert!(is_normalized_stage_target(target), "{target}"); + } + for target in ["sm_120", "compute_120", "ptx87"] { + assert!(is_normalized_stage_target(target), "{target}"); + } + for target in ["sm_0120a", "sm_120af", "sm_120x", "compute_120A"] { + assert!(!is_normalized_stage_target(target), "{target}"); + } +} + +#[test] +fn libnvvm_stage_may_report_newer_ptx_than_the_native_instruction_floor() { + let mut target_policy = policy(); + target_policy.minimum_ptx = "1.0".into(); + target_policy.minimum_sm = None; + let lowering = crate::model::OverlayBackendLowering { + backend: IntrinsicBackend::LibNvvm, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: "test".into(), + targets: None, + minimum_ptx: None, + minimum_sm: Some("sm_75".into()), + }; + let mut record = evidence(); + record.stages = vec![evidence_stage( + EvidenceStageKind::BackendCodegen, + BackendLoweringMechanism::InlinePtx, + &["sm_75", "ptx93"], + )]; + validate_selected_stage_targets(&target_policy, &record, &lowering).unwrap(); + + record.stages[0].targets = vec!["sm_75".into(), "ptx09".into()]; + assert!( + validate_selected_stage_targets(&target_policy, &record, &lowering) + .unwrap_err() + .to_string() + .contains("catalog floor sm_75 / PTX 1.0") + ); + + let llvm_lowering = crate::model::OverlayBackendLowering { + backend: IntrinsicBackend::LlvmNvptx, + mechanism: BackendLoweringMechanism::TypedNvvm, + evidence_profile: "test".into(), + targets: None, + minimum_ptx: Some("3.2".into()), + minimum_sm: Some("sm_20".into()), + }; + record.stages = vec![evidence_stage( + EvidenceStageKind::BackendCodegen, + BackendLoweringMechanism::TypedNvvm, + &["sm_20", "ptx93"], + )]; + assert!( + validate_selected_stage_targets(&target_policy, &record, &llvm_lowering) + .unwrap_err() + .to_string() + .contains("catalog floor sm_20 / PTX 3.2") + ); +} + +#[test] +fn imported_selection_must_match_the_full_ptx_shape() { + let mut imported = declaration(); + imported.selections[0].asm = "mov.u32 $d, %tid.xy;".into(); + let error = validate_imported_policy(&policy(), &imported).unwrap_err(); + assert!(error.to_string().contains("does not agree")); + + imported.selections[0].asm = "mov.u32.relaxed $d, %tid.x;".into(); + let error = validate_imported_policy(&policy(), &imported).unwrap_err(); + assert!(error.to_string().contains("does not agree")); + + imported.selections[0].asm = "mov.u32 $d, %tid.x;".into(); + validate_imported_policy(&policy(), &imported).unwrap(); +} + +#[test] +fn blackwell_ldmatrix_evidence_covers_every_target_with_its_effective_ptx_floor() { + let mut policy = policy(); + policy.family = "ldmatrix".into(); + policy.minimum_ptx = "8.6".into(); + policy.minimum_sm = None; + policy.targets = BLACKWELL_LDMATRIX_LLVM_TARGETS.into(); + policy.ldmatrix_variant = Some(crate::model::LdmatrixVariant { + shape: LdmatrixShape::M16n16, + multiplicity: LdmatrixMultiplicity::X1, + layout: LdmatrixLayout::Transposed, + element: LdmatrixElement::B8, + state_space: LdmatrixStateSpace::Shared, + }); + let lowering = OverlayBackendLowering { + backend: IntrinsicBackend::LlvmNvptx, + mechanism: BackendLoweringMechanism::TypedNvvm, + evidence_profile: "test".into(), + targets: None, + minimum_ptx: None, + minimum_sm: None, + }; + let mut record = evidence(); + record.status = "validated".into(); + record.stages.clear(); + for (target, ptx) in [ + ("sm_100a", "ptx86"), + ("sm_100f", "ptx88"), + ("sm_103a", "ptx88"), + ("sm_103f", "ptx88"), + ("sm_110a", "ptx90"), + ("sm_110f", "ptx90"), + ("sm_120a", "ptx87"), + ("sm_120f", "ptx88"), + ("sm_121a", "ptx88"), + ("sm_121f", "ptx88"), + ] { + record.stages.push(evidence_stage( + EvidenceStageKind::BackendCodegen, + BackendLoweringMechanism::TypedNvvm, + &[target, ptx], + )); + let mut assembly = evidence_stage( + EvidenceStageKind::PtxAssembly, + BackendLoweringMechanism::TypedNvvm, + &[target, ptx], + ); + assembly.tool_path = Some("/tool/ptxas".into()); + assembly.tool_version = Some("test".into()); + assembly.tool_sha256 = Some("0".repeat(64)); + record.stages.push(assembly); + } + validate_selected_stage_targets(&policy, &record, &lowering).unwrap(); + + record.stages.retain(|stage| { + !(stage.stage == EvidenceStageKind::PtxAssembly + && stage.targets.iter().any(|target| target == "sm_121f")) + }); + assert!( + validate_selected_stage_targets(&policy, &record, &lowering) + .unwrap_err() + .to_string() + .contains("one structured stage for each") + ); +} + +#[test] +fn paired_target_evidence_checks_backend_and_runtime_floors() { + let policy = policy(); + let hardware = vec![ + CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 103 }, + ]; + let paired = vec![ + CatalogTargetAlternative { + minimum_ptx: "8.6".parse().unwrap(), + hardware: hardware[0], + }, + CatalogTargetAlternative { + minimum_ptx: "8.8".parse().unwrap(), + hardware: hardware[1], + }, + ]; + let llvm = OverlayBackendLowering { + backend: IntrinsicBackend::LlvmNvptx, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: "test".into(), + targets: None, + minimum_ptx: None, + minimum_sm: None, + }; + let mut llvm_record = evidence(); + llvm_record.status = "validated".into(); + llvm_record.stages.clear(); + for (target, ptx) in [("sm_100a", "ptx86"), ("sm_103a", "ptx88")] { + llvm_record.stages.push(evidence_stage( + EvidenceStageKind::BackendCodegen, + llvm.mechanism, + &[target, ptx], + )); + let mut assembly = evidence_stage( + EvidenceStageKind::PtxAssembly, + llvm.mechanism, + &[target, ptx], + ); + assembly.tool_path = Some("/tool/ptxas".into()); + assembly.tool_version = Some("test".into()); + assembly.tool_sha256 = Some("0".repeat(64)); + llvm_record.stages.push(assembly); + } + validate_target_matrix_stage_targets( + &policy, + &llvm_record, + &llvm, + EvidenceStageKind::PtxAssembly, + &hardware, + 86, + Some(&paired), + ) + .unwrap(); + + let mut wrong_floor = llvm_record.clone(); + wrong_floor + .stages + .iter_mut() + .find(|stage| { + stage.stage == EvidenceStageKind::BackendCodegen + && stage.targets.iter().any(|target| target == "sm_103a") + }) + .unwrap() + .targets = vec!["sm_103a".into(), "ptx86".into()]; + assert!( + validate_target_matrix_stage_targets( + &policy, + &wrong_floor, + &llvm, + EvidenceStageKind::PtxAssembly, + &hardware, + 86, + Some(&paired), + ) + .unwrap_err() + .to_string() + .contains("wrong PTX floor") + ); + + let libnvvm = OverlayBackendLowering { + backend: IntrinsicBackend::LibNvvm, + ..llvm.clone() + }; + let mut libnvvm_record = evidence(); + libnvvm_record.status = "validated".into(); + libnvvm_record.stages.clear(); + for (target, ptx) in [("sm_100a", "ptx88"), ("sm_103a", "ptx90")] { + libnvvm_record.stages.push(evidence_stage( + EvidenceStageKind::BackendCodegen, + libnvvm.mechanism, + &[target, ptx], + )); + let mut assembly = evidence_stage( + EvidenceStageKind::PtxAssembly, + libnvvm.mechanism, + &[target, ptx], + ); + assembly.tool_path = Some("/tool/ptxas".into()); + assembly.tool_version = Some("test".into()); + assembly.tool_sha256 = Some("0".repeat(64)); + libnvvm_record.stages.push(assembly); + let mut link = evidence_stage( + EvidenceStageKind::DeviceLink, + libnvvm.mechanism, + &[target, ptx], + ); + link.artifact_kind = Some(EvidenceArtifactKind::Cubin); + link.tool_path = Some("/tool/nvlink".into()); + link.tool_version = Some("test".into()); + link.tool_sha256 = Some("0".repeat(64)); + libnvvm_record.stages.push(link); + } + validate_target_matrix_stage_targets( + &policy, + &libnvvm_record, + &libnvvm, + EvidenceStageKind::DeviceLink, + &hardware, + 86, + Some(&paired), + ) + .unwrap(); + + let mut executed = llvm_record; + executed.status = "executed".into(); + executed.stages.push(evidence_stage( + EvidenceStageKind::Runtime, + llvm.mechanism, + &["sm_103a", "ptx88"], + )); + validate_target_matrix_stage_targets( + &policy, + &executed, + &llvm, + EvidenceStageKind::PtxAssembly, + &hardware, + 86, + Some(&paired), + ) + .unwrap(); + + let mut wrong_runtime_ptx = executed.clone(); + wrong_runtime_ptx.stages.last_mut().unwrap().targets = vec!["sm_103a".into(), "ptx87".into()]; + assert!( + validate_target_matrix_stage_targets( + &policy, + &wrong_runtime_ptx, + &llvm, + EvidenceStageKind::PtxAssembly, + &hardware, + 86, + Some(&paired), + ) + .unwrap_err() + .to_string() + .contains("paired floor") + ); + + let mut wrong_runtime_hardware = executed; + wrong_runtime_hardware.stages.last_mut().unwrap().targets = + vec!["sm_110a".into(), "ptx90".into()]; + assert!( + validate_target_matrix_stage_targets( + &policy, + &wrong_runtime_hardware, + &llvm, + EvidenceStageKind::PtxAssembly, + &hardware, + 86, + Some(&paired), + ) + .unwrap_err() + .to_string() + .contains("outside its target matrix") + ); +} + +#[test] +fn sparse_f8f6f4_f16_evidence_covers_every_target_and_floor() { + let policy = expand_sparse_mma_f8f6f4_f16_admission(&test_sparse_mma_f8f6f4_f16_admission()) + .unwrap() + .remove(0); + let lowering = policy + .backend_lowerings + .iter() + .find(|lowering| lowering.backend == IntrinsicBackend::LlvmNvptx) + .unwrap(); + let mut record = evidence(); + record.status = "validated".into(); + record.stages.clear(); + for (target, ptx) in [ + ("sm_120a", "ptx87"), + ("sm_120f", "ptx88"), + ("sm_121a", "ptx88"), + ("sm_121f", "ptx88"), + ] { + record.stages.push(evidence_stage( + EvidenceStageKind::BackendCodegen, + BackendLoweringMechanism::InlinePtx, + &[target, ptx], + )); + let mut assembly = evidence_stage( + EvidenceStageKind::PtxAssembly, + BackendLoweringMechanism::InlinePtx, + &[target, ptx], + ); + assembly.artifact_kind = Some(EvidenceArtifactKind::Cubin); + assembly.tool_path = Some("/tool/ptxas".into()); + assembly.tool_version = Some("test".into()); + assembly.tool_sha256 = Some("0".repeat(64)); + record.stages.push(assembly); + } + validate_selected_stage_targets(&policy, &record, lowering).unwrap(); + + let valid = record.clone(); + record.stages.retain(|stage| { + !(stage.stage == EvidenceStageKind::PtxAssembly + && stage.targets.iter().any(|target| target == "sm_121f")) + }); + assert!( + validate_selected_stage_targets(&policy, &record, lowering) + .unwrap_err() + .to_string() + .contains("one structured stage for each") + ); + + record = valid; + record + .stages + .iter_mut() + .find(|stage| { + stage.stage == EvidenceStageKind::BackendCodegen + && stage.targets.iter().any(|target| target == "sm_120f") + }) + .unwrap() + .targets = vec!["sm_120f".into(), "ptx87".into()]; + assert!( + validate_selected_stage_targets(&policy, &record, lowering) + .unwrap_err() + .to_string() + .contains("wrong PTX floor") + ); +} + +#[test] +fn non_blackwell_evidence_rejects_multiple_target_specific_stage_pairs() { + let policy = policy(); + let lowering = OverlayBackendLowering { + backend: IntrinsicBackend::LlvmNvptx, + mechanism: BackendLoweringMechanism::TypedNvvm, + evidence_profile: "test".into(), + targets: None, + minimum_ptx: None, + minimum_sm: None, + }; + let mut record = evidence(); + record.stages = vec![ + evidence_stage( + EvidenceStageKind::BackendCodegen, + BackendLoweringMechanism::TypedNvvm, + &["sm_20", "ptx20"], + ), + evidence_stage( + EvidenceStageKind::BackendCodegen, + BackendLoweringMechanism::TypedNvvm, + &["sm_21", "ptx20"], + ), + ]; + + assert!( + validate_selected_stage_targets(&policy, &record, &lowering) + .unwrap_err() + .to_string() + .contains("outside reviewed target-matrix evidence") + ); +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/tests/fixtures.rs b/crates/cuda-intrinsics-gen/src/resolve/tests/fixtures.rs new file mode 100644 index 0000000000..44e749fd0e --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/tests/fixtures.rs @@ -0,0 +1,1904 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + AbiLedgerEntry, AbiLedgerFile, BackendLoweringMechanism, CatalogHardwareAlternative, + CatalogHardwareTarget, ClcAdmission, ClcOperation, ClusterBarrierAdmission, ClusterBarrierMode, + ClusterMemoryAdmission, ClusterMemoryOperation, DebugControlAdmission, DebugControlOperation, + DotProductOperation, DotProductSignedness, EvidenceFile, EvidenceRecord, EvidenceStage, + EvidenceStageKind, ImportedFile, ImportedIntrinsic, IntrinsicBackend, IntrinsicSource, + MaskEncoding, MbarrierExtendedAdmission, MbarrierExtendedOperation, MovmatrixAdapter, + MovmatrixParticipation, OverlayBackendLowering, OverlayFile, OverlayIntrinsic, + PackedAluAdapter, PackedAluFormat, PackedAluOperation, PackedConversionAdapter, + PackedConversionDestinationFormat, PackedConversionFp8Admission, PackedConversionFp8Direction, + PackedConversionFp8F16x2Admission, PackedConversionFp8Format, PackedConversionRounding, + PackedConversionSaturation, PackedConversionSourceFormat, PreSm70MemberMaskRule, PrmtAdmission, + PrmtMode, RegisterMmaAccumulator, RegisterMmaAmpereFloatAdmission, RegisterMmaF8F6F4Admission, + RegisterMmaFp8Admission, RuntimeValidation, SparseMmaElement, SparseMmaF8F6F4Admission, + SparseMmaF8F6F4F16Admission, SpecialRegisterAdmission, StmatrixAdmission, StmatrixLayout, + StmatrixMultiplicity, Tcgen05Admission, Tcgen05CpAdmissionVariant, Tcgen05CpGroup, + Tcgen05LdAdmissionVariant, Tcgen05MmaAdmissionVariant, Tcgen05MmaForm, Tcgen05Operation, + Tcgen05StAdmissionVariant, ThreadfenceAdmission, ThreadfenceScope, TmaAdmission, VoteAdapter, + VoteMode, VoteParticipation, WarpShuffleAdapter, WarpShuffleMode, WarpShuffleParticipation, + WarpShuffleSourceLane, WarpShuffleValueKind, WgmmaControlAdmission, WgmmaControlMode, +}; +use crate::ptx::{InstructionPattern, OperandPattern}; +use crate::util::read_json; +use anyhow::Result; +use std::collections::BTreeMap; +use std::fs; +use std::path::{Path, PathBuf}; + +use crate::model::ImportedSelection; +use crate::resolve::abi_ledger::*; +use crate::resolve::evidence::*; +use crate::resolve::families::*; +use crate::resolve::overlay::*; +use crate::resolve::policy::*; +use crate::resolve::targets::*; + +pub(super) fn sreg_pattern(special_register: &str) -> InstructionPattern { + InstructionPattern::new( + "mov", + &["u32"], + vec![ + OperandPattern::Register, + OperandPattern::Exact { + value: special_register.into(), + }, + ], + ) +} + +pub(super) fn policy() -> OverlayIntrinsic { + OverlayIntrinsic { + id: "thread_idx_x".into(), + abi_id: "i0001".into(), + operation_key: "launch.thread_index.x".into(), + family: "sreg".into(), + source: None, + source_record: Some("int_nvvm_read_ptx_sreg_tid_x".into()), + rust_module: "sreg".into(), + rust_name: "thread_idx_x".into(), + rust_arguments: vec![], + rust_result: "u32".into(), + safe: true, + must_use: false, + safe_allowlist_reason: Some("no caller obligations".into()), + public_rust_path: "cuda_intrinsics::sreg::thread_idx_x".into(), + compatibility_rust_paths: vec!["cuda_device::thread::threadIdx_x".into()], + dialect_op_type: "ReadPtxSregTidXOp".into(), + dialect_op_name: "nvvm.read_ptx_sreg_tid_x".into(), + dialect_operands: vec![], + dialect_results: vec![], + llvm_symbol: Some("llvm.nvvm.read.ptx.sreg.tid.x".into()), + resolved_llvm_symbol: None, + llvm_arguments: vec![], + llvm_results: vec!["i32".into()], + pure: true, + memory: "none".into(), + convergent: false, + execution_scope: "thread".into(), + minimum_ptx: "2.0".into(), + minimum_sm: None, + ptx_result: "u32".into(), + targets: "all".into(), + ptx_isa_version: "9.3".into(), + ptx_isa_section: "10.1 Special Registers: %tid".into(), + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/".into(), + lowering: "direct_nvvm".into(), + backend_lowerings: vec![], + packed_atomic: None, + redux: None, + vote: None, + active_mask: None, + warp_match: None, + warp_barrier: None, + warp_shuffle: None, + dot_product: None, + packed_alu: None, + integer_minmax: None, + packed_conversion: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cp_async_copy: None, + cp_async_control: None, + cp_async_mbarrier: None, + mbarrier_basic: None, + movmatrix: None, + mbarrier_extended: None, + register_mma: None, + sparse_mma: None, + prmt: None, + cluster_barrier: None, + wgmma_control: None, + special_register: None, + debug_control: None, + cluster_memory: None, + clc: None, + tma: None, + tcgen05: None, + ldmatrix_variant: None, + ldmatrix_safety: None, + ldmatrix_adapter: None, + selected_address_space: None, + expected_ptx: sreg_pattern("%tid.x"), + summary: "thread index".into(), + } +} + +pub(super) fn distinct_policy() -> OverlayIntrinsic { + let mut record = policy(); + record.id = "thread_idx_y".into(); + record.abi_id = "i0002".into(); + record.operation_key = "launch.thread_index.y".into(); + record.source_record = Some("int_nvvm_read_ptx_sreg_tid_y".into()); + record.rust_name = "thread_idx_y".into(); + record.public_rust_path = "cuda_intrinsics::sreg::thread_idx_y".into(); + record.compatibility_rust_paths = vec!["cuda_device::thread::threadIdx_y".into()]; + record.dialect_op_type = "ReadPtxSregTidYOp".into(); + record.dialect_op_name = "nvvm.read_ptx_sreg_tid_y".into(); + record.llvm_symbol = Some("llvm.nvvm.read.ptx.sreg.tid.y".into()); + record.expected_ptx = sreg_pattern("%tid.y"); + record +} + +pub(super) fn movmatrix_policy() -> OverlayIntrinsic { + let mut record = policy(); + record.id = "movmatrix_trans_b16".into(); + record.abi_id = "i0305".into(); + record.operation_key = "movmatrix.m8n8.trans.b16".into(); + record.family = "movmatrix".into(); + record.source = Some(IntrinsicSource::PtxNative { + instruction: "movmatrix.sync.aligned.m8n8.trans.b16".into(), + }); + record.source_record = None; + record.rust_module = "matrix".into(); + record.rust_name = "movmatrix_trans_b16".into(); + record.rust_arguments = vec!["u32".into()]; + record.rust_result = "u32".into(); + record.safe = false; + record.must_use = true; + record.safe_allowlist_reason = None; + record.public_rust_path = "cuda_intrinsics::matrix::movmatrix_trans_b16".into(); + record.compatibility_rust_paths = vec!["cuda_device::wmma::movmatrix_trans_b16".into()]; + record.dialect_op_type = "MovmatrixTransB16Op".into(); + record.dialect_op_name = "nvvm.movmatrix_trans_b16".into(); + record.dialect_operands = vec!["i32".into()]; + record.dialect_results = vec!["i32".into()]; + record.llvm_symbol = None; + record.resolved_llvm_symbol = None; + record.llvm_arguments.clear(); + record.llvm_results.clear(); + record.pure = false; + record.memory = "inaccessible_read_write".into(); + record.convergent = true; + record.execution_scope = "warp".into(); + record.minimum_ptx = "7.8".into(); + record.minimum_sm = Some("sm_75".into()); + record.ptx_result = "u32".into(); + record.targets = "all".into(); + record.ptx_isa_section = + "9.7.15.5.17 Warp-level matrix transpose instruction: movmatrix".into(); + record.lowering = "generated_movmatrix_inline_ptx".into(); + record.backend_lowerings = [IntrinsicBackend::LlvmNvptx, IntrinsicBackend::LibNvvm] + .into_iter() + .map(|backend| OverlayBackendLowering { + backend, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: match backend { + IntrinsicBackend::LlvmNvptx => "llvm-test", + IntrinsicBackend::LibNvvm => "libnvvm-test", + } + .into(), + targets: None, + minimum_ptx: Some("7.8".into()), + minimum_sm: Some("sm_75".into()), + }) + .collect(); + record.movmatrix = Some(crate::model::Movmatrix { + participation: MovmatrixParticipation::AllWarpLanesSameInstructionNoExitedLanes, + adapter: MovmatrixAdapter::PackedB16x2U32ToPackedB16x2U32, + runtime_validation: RuntimeValidation::Unexecuted, + }); + record.expected_ptx = InstructionPattern::new( + "movmatrix", + &["sync", "aligned", "m8n8", "trans", "b16"], + vec![OperandPattern::Register, OperandPattern::Register], + ); + record.summary = "Transposes one packed b16 matrix fragment across a warp.".into(); + record +} + +pub(super) fn declaration() -> ImportedIntrinsic { + ImportedIntrinsic { + source_record: "int_nvvm_read_ptx_sreg_tid_x".into(), + llvm_name: "llvm.nvvm.read.ptx.sreg.tid.x".into(), + arguments: vec![], + results: vec!["i32".into()], + classes: vec!["NVVMPureIntrinsic".into()], + properties: vec![ + "IntrNoMem".into(), + "IntrSpeculatable".into(), + "NoUndef".into(), + "Range".into(), + ], + selections: vec![ImportedSelection { + source_record: "INT_PTX_SREG_TID_x".into(), + asm: "mov.u32 $d, %tid.x;".into(), + predicates: vec![], + constraints: Default::default(), + }], + } +} + +pub(super) fn evidence() -> EvidenceRecord { + EvidenceRecord { + id: "thread_idx_x".into(), + source: None, + source_record: Some("int_nvvm_read_ptx_sreg_tid_x".into()), + llvm_symbol: Some("llvm.nvvm.read.ptx.sreg.tid.x".into()), + resolved_llvm_symbol: None, + llvm_arguments: vec![], + llvm_results: vec!["i32".into()], + concrete_llvm_arguments: vec![], + concrete_llvm_results: vec![], + target_triple: "nvptx64-nvidia-cuda".into(), + gpu_target: "sm_70".into(), + ptx_feature: "+ptx60".into(), + status: "lowered".into(), + stages: vec![], + declaration_attributes_canonicalized: None, + runtime_validation: None, + expected_ptx: sreg_pattern("%tid.x"), + } +} + +pub(super) fn validate_test_evidence( + policy: &OverlayIntrinsic, + record: EvidenceRecord, +) -> Result<()> { + let file = EvidenceFile { + schema: 3, + backend_profile: "test".into(), + backend_kind: None, + llvm_revision: "test".into(), + backend_version: "LLVM version test".into(), + backend_sha256: "0123456789abcdef".into(), + artifact_path: None, + build_id_prefix: None, + nvvm_ir_version: None, + debug_ir_version: None, + records: vec![record], + }; + let indexed = IndexedEvidence { + file: &file, + record: &file.records[0], + backend_version: &file.backend_version, + backend_sha256: &file.backend_sha256, + }; + validate_evidence(policy, &indexed, None) +} + +pub(super) fn shared_matrix_stage() -> EvidenceStage { + EvidenceStage { + targets: vec!["sm_80".into(), "ptx71".into()], + representation: "shared fixture".into(), + stage: EvidenceStageKind::BackendCodegen, + mechanism: Some(BackendLoweringMechanism::InlinePtx), + outcome: "succeeded".into(), + detail: "$dst remains fixture text".into(), + artifact_kind: None, + tool_path: None, + tool_version: None, + tool_sha256: None, + } +} + +pub(super) fn synthetic_matrix_json() -> serde_json::Value { + serde_json::json!({ + "schema": 6, + "backend_profile": "matrix-test", + "backend_kind": "llvm_nvptx", + "llvm_revision": "test", + "backend_version": "LLVM matrix test", + "backend_sha256": "0123456789abcdef", + "defaults": { + "llvm_arguments": ["i32"], + "llvm_results": ["i32"], + "target_triple": "nvptx64-nvidia-cuda", + "gpu_target": "sm_80", + "ptx_feature": "+ptx71", + "status": "lowered" + }, + "fixtures": [{ + "id": "shared", + "coverage_count": 2, + "stages": [{ + "targets": ["sm_80", "ptx71"], + "representation": "shared fixture", + "stage": "backend_codegen", + "mechanism": "inline_ptx", + "outcome": "succeeded", + "detail": "$dst remains fixture text" + }] + }], + "matrices": [{ + "axes": [{ + "name": "element", + "values": ["s8", "u8"] + }], + "product_count": 2, + "fixtures": ["shared"], + "template": { + "id": "synthetic_${element}", + "source_record": "int_synthetic_${element}", + "llvm_symbol": "llvm.synthetic.${element}", + "expected_ptx": { + "mnemonic": "mma", + "modifiers": ["sync", "${element}"], + "operands": [{"kind": "register"}] + } + } + }], + "records": [{ + "id": "synthetic_explicit", + "source_record": "int_synthetic_explicit", + "llvm_symbol": "llvm.synthetic.explicit", + "llvm_arguments": ["i32"], + "llvm_results": ["i32"], + "target_triple": "nvptx64-nvidia-cuda", + "gpu_target": "sm_80", + "ptx_feature": "+ptx71", + "status": "lowered", + "expected_ptx": { + "mnemonic": "mma", + "modifiers": ["sync", "explicit"], + "operands": [{"kind": "register"}] + } + }] + }) +} + +pub(super) fn policy_matrix_json() -> serde_json::Value { + serde_json::json!({ + "schema": 6, + "backend_profile": "matrix-test", + "llvm_revision": "test", + "backend_version": "LLVM matrix test", + "backend_sha256": "0123456789abcdef", + "defaults": { + "llvm_arguments": [], + "llvm_results": ["i32"], + "target_triple": "nvptx64-nvidia-cuda", + "gpu_target": "sm_70", + "ptx_feature": "+ptx60", + "status": "lowered" + }, + "fixtures": [{ + "id": "policy_fixture", + "coverage_count": 1, + "stages": [{ + "targets": ["sm_70", "ptx60"], + "representation": "policy fixture", + "stage": "backend_codegen", + "mechanism": "typed_nvvm", + "outcome": "succeeded", + "detail": "shared policy fixture" + }] + }], + "matrices": [{ + "axes": [{ + "name": "axis", + "values": ["x"] + }], + "product_count": 1, + "fixtures": ["policy_fixture"], + "template": { + "id": "thread_idx_${axis}", + "source_record": "int_nvvm_read_ptx_sreg_tid_${axis}", + "llvm_symbol": "llvm.nvvm.read.ptx.sreg.tid.${axis}", + "expected_ptx": { + "mnemonic": "mov", + "modifiers": ["u32"], + "operands": [ + {"kind": "register"}, + {"kind": "exact", "value": "%tid.${axis}"} + ] + } + } + }] + }) +} + +pub(super) fn parse_synthetic_evidence(value: &serde_json::Value) -> Result { + parse_evidence_bytes(&serde_json::to_vec(value).unwrap(), "synthetic evidence") +} + +pub(super) fn assert_synthetic_evidence_error(value: &serde_json::Value, expected: &str) { + let error = parse_synthetic_evidence(value).unwrap_err(); + let message = format!("{error:#}"); + assert!( + message.contains(expected), + "expected {expected:?} in {message:?}" + ); +} + +pub(super) fn overlay_file(records: Vec) -> OverlayFile { + OverlayFile { + schema: OVERLAY_SCHEMA, + catalog_version: "test".into(), + intrinsic_abi: 1, + backend_profile: "test".into(), + shards: vec![], + intrinsics: records, + } +} + +pub(super) fn bind_pinned_abi_ids(repo_root: &Path, overlay: &mut OverlayFile) { + let ledger_path = repo_root.join(format!("intrinsics/abi-v{}.toml", overlay.intrinsic_abi)); + let ledger: AbiLedgerFile = + toml::from_str(&std::fs::read_to_string(ledger_path).unwrap()).unwrap(); + bind_generated_abi_ids(overlay, &ledger).unwrap(); +} + +pub(super) fn validate_imported_policy( + policy: &OverlayIntrinsic, + declaration: &ImportedIntrinsic, +) -> Result<()> { + let source = resolve_policy_source(policy)?; + validate_policy(policy, &source, Some(declaration), 1) +} + +pub(super) fn pinned_active_mask_and_warp_match_records() +-> BTreeMap { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let (overlay, _) = + read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")).unwrap(); + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let declarations: BTreeMap<_, _> = imported + .intrinsics + .into_iter() + .map(|record| (record.source_record.clone(), record)) + .collect(); + + overlay + .intrinsics + .into_iter() + .filter(|record| matches!(record.family.as_str(), "active_mask" | "warp_match")) + .map(|policy| { + let declaration = declarations[policy.source_record.as_deref().unwrap()].clone(); + (policy.id.clone(), (policy, declaration)) + }) + .collect() +} + +pub(super) fn pinned_mbarrier_basic_records() +-> BTreeMap { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let (overlay, _) = + read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")).unwrap(); + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let declarations: BTreeMap<_, _> = imported + .intrinsics + .into_iter() + .map(|record| (record.source_record.clone(), record)) + .collect(); + + overlay + .intrinsics + .into_iter() + .filter(|record| record.family == "mbarrier_basic") + .map(|policy| { + let declaration = declarations[policy.source_record.as_deref().unwrap()].clone(); + (policy.id.clone(), (policy, declaration)) + }) + .collect() +} + +pub(super) fn pinned_cp_async_mbarrier_records() +-> BTreeMap { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let (overlay, _) = + read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")).unwrap(); + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let declarations: BTreeMap<_, _> = imported + .intrinsics + .into_iter() + .map(|record| (record.source_record.clone(), record)) + .collect(); + + overlay + .intrinsics + .into_iter() + .filter(|record| record.family == "cp_async_mbarrier") + .map(|policy| { + let declaration = declarations[policy.source_record.as_deref().unwrap()].clone(); + (policy.id.clone(), (policy, declaration)) + }) + .collect() +} + +pub(super) fn packed_policy(id: &str) -> OverlayIntrinsic { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")) + .unwrap() + .0 + .intrinsics + .into_iter() + .find(|record| record.id == id) + .unwrap() +} + +pub(super) fn packed_alu_policy( + format: PackedAluFormat, + operation: PackedAluOperation, +) -> OverlayIntrinsic { + let recipe = packed_alu_recipe(format, operation).expect("test recipe pair"); + let (rust_module, rust_type, dialect_type, adapter) = match format { + PackedAluFormat::Bf16x2 => ("bf16x2", "u32", "i32", PackedAluAdapter::DirectPackedU32), + PackedAluFormat::F16x2 => ("f16x2", "u32", "i32", PackedAluAdapter::DirectPackedU32), + PackedAluFormat::F32x2 => ("f32x2", "u64", "i64", PackedAluAdapter::DirectPackedU64), + }; + let mut record = policy(); + record.id = recipe.id.into(); + record.abi_id = recipe.abi_id.into(); + record.operation_key = recipe.operation_key.into(); + record.family = "packed_alu".into(); + match &recipe.source { + PackedAluRecipeSource::Imported { + record: source_record, + symbol, + resolved_symbol, + arguments, + results, + .. + } => { + record.source = None; + record.source_record = Some((*source_record).into()); + record.llvm_symbol = Some((*symbol).into()); + record.resolved_llvm_symbol = resolved_symbol.map(str::to_owned); + record.llvm_arguments = arguments.iter().map(|value| (*value).into()).collect(); + record.llvm_results = results.iter().map(|value| (*value).into()).collect(); + } + PackedAluRecipeSource::PtxNative => { + record.source = Some(IntrinsicSource::PtxNative { + instruction: recipe.ptx_mnemonic.into(), + }); + record.source_record = None; + record.llvm_symbol = None; + record.resolved_llvm_symbol = None; + record.llvm_arguments.clear(); + record.llvm_results.clear(); + } + } + record.rust_module = rust_module.into(); + record.rust_name = recipe.rust_name.into(); + record.rust_arguments = vec![rust_type.into(); recipe.arity]; + record.rust_result = rust_type.into(); + record.safe = true; + record.must_use = recipe.must_use; + record.safe_allowlist_reason = Some("the operation has no caller obligations".into()); + record.public_rust_path = format!("cuda_intrinsics::{rust_module}::{}", recipe.rust_name); + record.compatibility_rust_paths = + vec![format!("cuda_device::{rust_module}::{}", recipe.rust_name)]; + record.dialect_op_type = recipe.dialect_op_type.into(); + record.dialect_op_name = recipe.dialect_op_name.into(); + record.dialect_operands = vec![dialect_type.into(); recipe.arity]; + record.dialect_results = vec![dialect_type.into()]; + record.pure = true; + record.memory = "none".into(); + record.convergent = false; + record.execution_scope = "thread".into(); + record.minimum_ptx = recipe.minimum_ptx.into(); + record.minimum_sm = Some(recipe.minimum_sm.into()); + record.ptx_result = rust_type.into(); + record.ptx_isa_section = recipe.ptx_isa_section.into(); + record.ptx_isa_url = recipe.ptx_isa_url.into(); + record.lowering = "generated_packed_alu_inline_ptx".into(); + record.backend_lowerings = [IntrinsicBackend::LlvmNvptx, IntrinsicBackend::LibNvvm] + .into_iter() + .map(|backend| { + let (minimum_ptx, minimum_sm) = + packed_alu_backend_floor(&recipe, format, operation, backend); + crate::model::OverlayBackendLowering { + backend, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: format!("{backend:?}-test"), + targets: None, + minimum_ptx: Some(minimum_ptx.into()), + minimum_sm: Some(minimum_sm.into()), + } + }) + .collect(); + record.packed_alu = Some(crate::model::PackedAlu { + format, + native_minimum_sm: recipe.native_minimum_sm, + operation, + adapter, + }); + record.expected_ptx = InstructionPattern::new( + recipe.ptx_mnemonic.split('.').next().unwrap(), + recipe.modifiers, + vec![OperandPattern::Register; recipe.arity + 1], + ); + record.summary = format!("packed {rust_module} arithmetic"); + record +} + +pub(super) fn packed_alu_declaration( + format: PackedAluFormat, + operation: PackedAluOperation, +) -> Option { + let recipe = packed_alu_recipe(format, operation).expect("test recipe pair"); + let PackedAluRecipeSource::Imported { + record, + symbol, + arguments, + results, + properties, + selection, + selection_asm, + .. + } = recipe.source + else { + return None; + }; + let classes = if matches!(operation, PackedAluOperation::Min | PackedAluOperation::Max) { + vec!["Intrinsic".into()] + } else { + vec!["Intrinsic".into(), "NVVMPureIntrinsic".into()] + }; + let mut selections = vec![ImportedSelection { + source_record: selection.into(), + asm: selection_asm.into(), + predicates: vec![ + format!("Subtarget->getSmVersion() >= {}", recipe.native_minimum_sm), + format!( + "Subtarget->getPTXVersion() >= {}", + recipe.minimum_ptx.replace('.', "") + ), + ], + constraints: Default::default(), + }]; + if operation == PackedAluOperation::Abs { + selections.extend((0..5).map(|index| ImportedSelection { + source_record: format!("OTHER_ABS_{index}"), + asm: "abs.f32 $dst, $src0;".into(), + predicates: vec![], + constraints: Default::default(), + })); + } + Some(ImportedIntrinsic { + source_record: record.into(), + llvm_name: symbol.into(), + arguments: arguments.iter().map(|value| (*value).into()).collect(), + results: results.iter().map(|value| (*value).into()).collect(), + classes, + properties: properties.iter().map(|value| (*value).into()).collect(), + selections, + }) +} + +pub(super) fn packed_conversion_policy( + destination_format: PackedConversionDestinationFormat, + rounding: PackedConversionRounding, + saturation: PackedConversionSaturation, +) -> OverlayIntrinsic { + let conversion = crate::model::PackedConversion { + source_format: PackedConversionSourceFormat::F32x2, + destination_format, + rounding, + saturation, + adapter: PackedConversionAdapter::ReverseHighLowOperands, + }; + let recipe = packed_conversion_recipe(&conversion).expect("test packed-conversion recipe"); + let mut record = policy(); + record.id = recipe.id.into(); + record.abi_id = recipe.abi_id.into(); + record.operation_key = recipe.operation_key.into(); + record.family = "packed_conversion".into(); + record.source_record = Some(recipe.source_record.into()); + record.rust_module = "convert".into(); + record.rust_name = recipe.rust_name.into(); + record.rust_arguments = vec!["f32".into(), "f32".into()]; + let result_width = packed_conversion_result_width(&conversion); + record.rust_result = format!("u{result_width}"); + record.safe = true; + record.must_use = false; + record.safe_allowlist_reason = Some("the operation has no caller obligations".into()); + record.public_rust_path = format!("cuda_intrinsics::convert::{}", recipe.rust_name); + record.compatibility_rust_paths = vec![recipe.compatibility_path.into()]; + record.dialect_op_type = recipe.dialect_op_type.into(); + record.dialect_op_name = recipe.dialect_op_name.into(); + record.dialect_operands = vec!["f32".into(), "f32".into()]; + record.dialect_results = vec![format!("i{result_width}")]; + record.llvm_symbol = Some(recipe.llvm_symbol.into()); + record.llvm_arguments = vec!["f32".into(), "f32".into()]; + record.llvm_results = vec![recipe.llvm_result.into()]; + record.pure = true; + record.memory = "none".into(); + record.convergent = false; + record.execution_scope = "thread".into(); + let (minimum_ptx, minimum_sm) = packed_conversion_floor(&conversion); + record.minimum_ptx = minimum_ptx.into(); + record.minimum_sm = Some(minimum_sm.into()); + record.ptx_result = format!("u{result_width}"); + record.ptx_isa_section = "9.7.9.22 Data Movement and Conversion Instructions: cvt".into(); + record.ptx_isa_url = "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cvt".into(); + record.lowering = packed_conversion_lowering(&conversion).into(); + record.backend_lowerings = [IntrinsicBackend::LlvmNvptx, IntrinsicBackend::LibNvvm] + .into_iter() + .map(|backend| OverlayBackendLowering { + backend, + mechanism: packed_conversion_backend_mechanism(&conversion, backend), + evidence_profile: "test".into(), + targets: None, + minimum_ptx: Some(minimum_ptx.into()), + minimum_sm: Some(minimum_sm.into()), + }) + .collect(); + let modifiers = packed_conversion_ptx_modifiers(&conversion); + record.packed_conversion = Some(conversion); + record.expected_ptx = + InstructionPattern::new("cvt", &modifiers, vec![OperandPattern::Register; 3]); + record.summary = recipe.summary.into(); + record +} + +pub(super) fn packed_conversion_declaration(policy: &OverlayIntrinsic) -> ImportedIntrinsic { + ImportedIntrinsic { + source_record: policy.source_record.clone().unwrap(), + llvm_name: policy.llvm_symbol.clone().unwrap(), + arguments: policy.llvm_arguments.clone(), + results: policy.llvm_results.clone(), + classes: vec!["Intrinsic".into(), "NVVMPureIntrinsic".into()], + properties: vec!["IntrNoMem".into(), "IntrSpeculatable".into()], + selections: vec![], + } +} + +pub(super) fn packed_conversion_evidence(policy: &OverlayIntrinsic) -> EvidenceRecord { + let mut record = evidence(); + record.id = policy.id.clone(); + record.source_record = policy.source_record.clone(); + record.llvm_symbol = policy.llvm_symbol.clone(); + record.resolved_llvm_symbol = policy.resolved_llvm_symbol.clone(); + record.llvm_arguments = policy.llvm_arguments.clone(); + record.llvm_results = policy.llvm_results.clone(); + record.concrete_llvm_arguments = policy.llvm_arguments.clone(); + record.concrete_llvm_results = policy.llvm_results.clone(); + record.declaration_attributes_canonicalized = Some(true); + record.expected_ptx = policy.expected_ptx.clone(); + record +} + +pub(super) fn redux_policy() -> OverlayIntrinsic { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")) + .unwrap() + .0 + .intrinsics + .into_iter() + .find(|record| record.id == "redux_sync_add") + .unwrap() +} + +pub(super) fn redux_declaration() -> ImportedIntrinsic { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let text = std::fs::read_to_string(repo_root.join("intrinsics/imported.json")).unwrap(); + serde_json::from_str::(&text) + .unwrap() + .intrinsics + .into_iter() + .find(|record| record.source_record == "int_nvvm_redux_sync_add") + .unwrap() +} + +pub(super) fn sync_policy() -> OverlayIntrinsic { + packed_policy("sync_threads") +} + +pub(super) fn sync_declaration() -> ImportedIntrinsic { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let text = std::fs::read_to_string(repo_root.join("intrinsics/imported.json")).unwrap(); + serde_json::from_str::(&text) + .unwrap() + .intrinsics + .into_iter() + .find(|record| record.source_record == "int_nvvm_barrier_cta_sync_aligned_all") + .unwrap() +} + +pub(super) fn warp_barrier_policy() -> OverlayIntrinsic { + packed_policy("sync_mask") +} + +pub(super) fn warp_barrier_declaration() -> ImportedIntrinsic { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let text = std::fs::read_to_string(repo_root.join("intrinsics/imported.json")).unwrap(); + serde_json::from_str::(&text) + .unwrap() + .intrinsics + .into_iter() + .find(|record| record.source_record == "int_nvvm_bar_warp_sync") + .unwrap() +} + +pub(super) fn vote_policy(mode: VoteMode) -> OverlayIntrinsic { + let recipe = vote_recipe(mode); + let mut record = policy(); + record.id = recipe.id.into(); + record.abi_id = recipe.abi_id.into(); + record.operation_key = recipe.operation_key.into(); + record.family = "vote".into(); + record.source_record = Some(recipe.source_record.into()); + record.rust_module = "warp".into(); + record.rust_name = recipe.rust_name.into(); + record.rust_arguments = vec!["u32".into(), "bool".into()]; + record.rust_result = recipe.rust_result.into(); + record.safe = false; + record.must_use = true; + record.safe_allowlist_reason = None; + record.public_rust_path = format!("cuda_intrinsics::warp::{}", recipe.rust_name); + record.compatibility_rust_paths = if recipe.has_compatibility_path { + vec![format!("cuda_device::warp::{}", recipe.rust_name)] + } else { + vec![] + }; + record.dialect_op_type = recipe.dialect_op_type.into(); + record.dialect_op_name = recipe.dialect_op_name.into(); + record.dialect_operands = vec!["i32".into(), "i1".into()]; + record.dialect_results = vec![recipe.llvm_result.into()]; + record.llvm_symbol = Some(recipe.llvm_symbol.into()); + record.llvm_arguments = vec!["i32".into(), "i1".into()]; + record.llvm_results = vec![recipe.llvm_result.into()]; + record.pure = false; + record.memory = "inaccessible_read_write".into(); + record.convergent = true; + record.execution_scope = "warp".into(); + record.minimum_ptx = "6.0".into(); + record.minimum_sm = Some("sm_30".into()); + record.ptx_result = recipe.rust_result.into(); + record.ptx_isa_section = "9.7.14.10 Warp Vote Instructions: vote.sync".into(); + record.ptx_isa_url = "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-vote-sync".into(); + record.lowering = "generated_vote".into(); + record.vote = Some(crate::model::Vote { + mode, + participation: VoteParticipation::ExecutingLaneNamedAllNamedLanesSameInstructionAndMask, + legacy_pre_sm70: PreSm70MemberMaskRule::AllNamedLanesConvergedAndOnlyNamedLanesActive, + adapter: VoteAdapter::DirectMaskPredicate, + mask_encoding: MaskEncoding::RegisterOrImmediate, + }); + record.expected_ptx = InstructionPattern::new( + "vote", + &["sync", recipe.ptx_mode, recipe.ptx_type], + vec![ + OperandPattern::Register, + OperandPattern::Register, + OperandPattern::RegisterOrImmediate, + ], + ); + record.summary = "warp vote".into(); + record +} + +pub(super) fn vote_declaration(mode: VoteMode) -> ImportedIntrinsic { + let recipe = vote_recipe(mode); + let selection = |source_record: &str| ImportedSelection { + source_record: source_record.into(), + asm: format!( + "vote.sync.{}.{} \t$dest, $pred, $mask;", + recipe.ptx_mode, recipe.ptx_type + ), + predicates: vec![ + "Subtarget->getPTXVersion() >= 60".into(), + "Subtarget->getSmVersion() >= 30".into(), + ], + constraints: Default::default(), + }; + ImportedIntrinsic { + source_record: recipe.source_record.into(), + llvm_name: recipe.llvm_symbol.into(), + arguments: vec!["i32".into(), "i1".into()], + results: vec![recipe.llvm_result.into()], + classes: vec![ + "ClangBuiltin".into(), + "NVVMBuiltin".into(), + "SDPatternOperator".into(), + "Intrinsic".into(), + ], + properties: vec![ + "IntrConvergent".into(), + "IntrInaccessibleMemOnly".into(), + "IntrNoCallback".into(), + ], + selections: vec![ + selection(recipe.immediate_selection), + selection(recipe.register_selection), + ], + } +} + +pub(super) fn warp_shuffle_policy( + mode: WarpShuffleMode, + value_kind: WarpShuffleValueKind, +) -> OverlayIntrinsic { + let recipe = warp_shuffle_recipe(mode, value_kind); + let mut record = policy(); + record.id = recipe.id.into(); + record.abi_id = recipe.abi_id.into(); + record.operation_key = recipe.operation_key.into(); + record.family = "warp_shuffle".into(); + match recipe.source { + WarpShuffleRecipeSource::LlvmImported { + source_record, + llvm_symbol, + } => { + record.source_record = Some(source_record.into()); + record.llvm_symbol = Some(llvm_symbol.into()); + record.llvm_arguments = vec![ + "i32".into(), + recipe.dialect_value.into(), + "i32".into(), + "i32".into(), + ]; + record.llvm_results = vec![recipe.dialect_value.into()]; + } + WarpShuffleRecipeSource::PtxNative { instruction } => { + record.source = Some(IntrinsicSource::PtxNative { + instruction: instruction.into(), + }); + record.source_record = None; + record.llvm_symbol = None; + record.llvm_arguments.clear(); + record.llvm_results.clear(); + } + } + record.rust_module = "warp".into(); + record.rust_name = recipe.rust_name.into(); + record.rust_arguments = vec!["u32".into(), recipe.rust_value.into(), "u32".into()]; + record.rust_result = recipe.rust_value.into(); + record.safe = false; + record.must_use = true; + record.safe_allowlist_reason = None; + record.public_rust_path = format!("cuda_intrinsics::warp::{}", recipe.rust_name); + record.compatibility_rust_paths = vec![format!("cuda_device::warp::{}", recipe.rust_name)]; + record.dialect_op_type = recipe.dialect_op_type.into(); + record.dialect_op_name = recipe.dialect_op_name.into(); + record.dialect_operands = vec!["i32".into(), recipe.dialect_value.into(), "i32".into()]; + record.dialect_results = vec![recipe.dialect_value.into()]; + record.pure = false; + record.memory = "inaccessible_read_write".into(); + record.convergent = true; + record.execution_scope = "warp".into(); + record.minimum_ptx = "6.0".into(); + record.minimum_sm = Some("sm_30".into()); + record.ptx_result = recipe.rust_value.into(); + record.targets = "all".into(); + record.ptx_isa_section = "9.7.9.6 Data Movement and Conversion Instructions: shfl.sync".into(); + record.ptx_isa_url = "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-shfl-sync".into(); + record.lowering = recipe.lowering.into(); + record.backend_lowerings = vec![ + crate::model::OverlayBackendLowering { + backend: IntrinsicBackend::LlvmNvptx, + mechanism: recipe.backend_mechanism, + evidence_profile: "llvm-test".into(), + targets: None, + minimum_ptx: Some("6.0".into()), + minimum_sm: Some("sm_30".into()), + }, + crate::model::OverlayBackendLowering { + backend: IntrinsicBackend::LibNvvm, + mechanism: recipe.backend_mechanism, + evidence_profile: "libnvvm-test".into(), + targets: None, + minimum_ptx: Some("6.0".into()), + minimum_sm: Some("sm_75".into()), + }, + ]; + record.warp_shuffle = Some(crate::model::WarpShuffle { + mode, + value_kind, + participation: + WarpShuffleParticipation::ExecutingLaneNamedAllNamedLanesSameInstructionAndMask, + legacy_pre_sm70: PreSm70MemberMaskRule::AllNamedLanesConvergedAndOnlyNamedLanesActive, + source_lane: WarpShuffleSourceLane::InRangeSourceActiveAndNamedOutOfRangeCopiesSelf, + adapter: recipe.adapter, + clamp: recipe.clamp, + lane_encoding: recipe.operand_encoding, + mask_encoding: recipe.operand_encoding, + }); + let operands = match recipe.adapter { + WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp => vec![ + OperandPattern::Register, + OperandPattern::Register, + OperandPattern::RegisterOrImmediate, + OperandPattern::Exact { + value: recipe.clamp.to_string(), + }, + OperandPattern::RegisterOrImmediate, + ], + WarpShuffleAdapter::MaskValueLaneOrDeltaSplitI64LowHighB32InsertClampReassemble => { + vec![ + OperandPattern::Exact { value: "lo".into() }, + OperandPattern::Exact { value: "lo".into() }, + OperandPattern::Register, + OperandPattern::Exact { + value: recipe.clamp.to_string(), + }, + OperandPattern::Register, + ] + } + }; + record.expected_ptx = + InstructionPattern::new("shfl", &["sync", recipe.ptx_mode, "b32"], operands); + record.summary = "synchronized warp shuffle".into(); + record +} + +pub(super) fn warp_shuffle_declaration( + mode: WarpShuffleMode, + value_kind: WarpShuffleValueKind, +) -> ImportedIntrinsic { + let recipe = warp_shuffle_recipe(mode, value_kind); + let WarpShuffleRecipeSource::LlvmImported { + source_record, + llvm_symbol, + } = recipe.source + else { + panic!("PTX-native i64 shuffles have no imported declaration"); + }; + let selections = (0..8) + .map(|index| ImportedSelection { + source_record: format!("anonymous_test_{index}"), + asm: format!( + "shfl.sync.{}.b32 \t$dst, $src, $offset, $mask, $threadmask;", + recipe.ptx_mode + ), + predicates: vec![ + "Subtarget->getPTXVersion() >= 60".into(), + "Subtarget->getSmVersion() >= 30".into(), + ], + constraints: Default::default(), + }) + .collect(); + ImportedIntrinsic { + source_record: source_record.into(), + llvm_name: llvm_symbol.into(), + arguments: vec![ + "i32".into(), + recipe.dialect_value.into(), + "i32".into(), + "i32".into(), + ], + results: vec![recipe.dialect_value.into()], + classes: vec![ + "ClangBuiltin".into(), + "NVVMBuiltin".into(), + "SDPatternOperator".into(), + "Intrinsic".into(), + ], + properties: vec![ + "IntrConvergent".into(), + "IntrInaccessibleMemOnly".into(), + "IntrNoCallback".into(), + ], + selections, + } +} + +pub(super) fn sync_evidence(policy: &OverlayIntrinsic) -> EvidenceRecord { + let mut record = evidence(); + record.id = policy.id.clone(); + record.source_record = policy.source_record.clone(); + record.llvm_symbol = policy.llvm_symbol.clone(); + record.llvm_arguments = policy.llvm_arguments.clone(); + record.llvm_results = policy.llvm_results.clone(); + record.expected_ptx = policy.expected_ptx.clone(); + record +} + +pub(super) fn dot_product_policy( + operation: DotProductOperation, + signedness: DotProductSignedness, +) -> OverlayIntrinsic { + let recipe = dot_product_recipe(operation, signedness); + let mut record = policy(); + record.id = recipe.id.into(); + record.abi_id = match (operation, signedness) { + (DotProductOperation::Dp4a, DotProductSignedness::Signed) => "i0030", + (DotProductOperation::Dp4a, DotProductSignedness::Unsigned) => "i0031", + (DotProductOperation::Dp2a, DotProductSignedness::Signed) => "i0032", + (DotProductOperation::Dp2a, DotProductSignedness::Unsigned) => "i0033", + } + .into(); + record.operation_key = recipe.operation_key.into(); + record.family = "dotprod".into(); + record.source = None; + record.source_record = Some(recipe.source_record.into()); + record.rust_module = "dotprod".into(); + record.rust_name = recipe.rust_name.into(); + record.rust_arguments = vec!["u32".into(), "u32".into(), recipe.rust_value.into()]; + record.rust_result = recipe.rust_value.into(); + record.safe = true; + record.must_use = false; + record.safe_allowlist_reason = Some( + "per-thread integer arithmetic has no memory, pointer, or participation obligations".into(), + ); + record.public_rust_path = format!("cuda_intrinsics::dotprod::{}", recipe.rust_name); + record.compatibility_rust_paths = vec![format!("cuda_device::dotprod::{}", recipe.rust_name)]; + record.dialect_op_type = recipe.dialect_op_type.into(); + record.dialect_op_name = recipe.dialect_op_name.into(); + record.dialect_operands = vec!["i32".into(), "i32".into(), "i32".into()]; + record.dialect_results = vec!["i32".into()]; + record.llvm_symbol = Some(recipe.llvm_symbol.into()); + record.resolved_llvm_symbol = None; + record.llvm_arguments = recipe + .llvm_arguments + .iter() + .map(|argument| (*argument).into()) + .collect(); + record.llvm_results = vec!["i32".into()]; + record.pure = true; + record.memory = "none".into(); + record.convergent = false; + record.execution_scope = "thread".into(); + record.minimum_ptx = "5.0".into(); + record.minimum_sm = Some("sm_61".into()); + record.ptx_result = recipe.rust_value.into(); + record.targets = "all".into(); + record.lowering = "generated_dotprod".into(); + record.backend_lowerings = vec![ + crate::model::OverlayBackendLowering { + backend: IntrinsicBackend::LlvmNvptx, + mechanism: BackendLoweringMechanism::TypedNvvm, + evidence_profile: "llvm-test".into(), + targets: None, + minimum_ptx: None, + minimum_sm: None, + }, + crate::model::OverlayBackendLowering { + backend: IntrinsicBackend::LibNvvm, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: "libnvvm-test".into(), + targets: None, + minimum_ptx: None, + minimum_sm: Some("sm_75".into()), + }, + ]; + record.dot_product = Some(crate::model::DotProduct { + operation, + signedness, + adapter: recipe.adapter, + }); + record.expected_ptx = InstructionPattern::new( + recipe.ptx_mnemonic, + recipe.ptx_modifiers, + vec![OperandPattern::Register; 4], + ); + record.summary = "packed integer dot product".into(); + record +} + +pub(super) fn dot_product_declaration( + operation: DotProductOperation, + signedness: DotProductSignedness, +) -> ImportedIntrinsic { + let recipe = dot_product_recipe(operation, signedness); + let selection = |source_record: &str, half: Option<(&str, i64)>| ImportedSelection { + source_record: source_record.into(), + asm: format!( + "{}.{} $dst, $a, $b, $c;", + recipe.ptx_mnemonic, + match half { + Some((name, _)) => { + let types = &recipe.ptx_modifiers[1..]; + format!("{name}.{}", types.join(".")) + } + None => recipe.ptx_modifiers.join("."), + } + ), + predicates: vec!["hasDotInstructions".into()], + constraints: crate::model::ImportedSelectionConstraints { + address_space: None, + immediate_bindings: half + .map(|(_, value)| { + vec![crate::model::ImportedImmediateBinding { + argument_index: 2, + value, + }] + }) + .unwrap_or_default(), + }, + }; + let selections = match operation { + DotProductOperation::Dp4a => vec![selection("DOT4", None)], + DotProductOperation::Dp2a => vec![ + selection("DOT2_hi", Some(("hi", -1))), + selection("DOT2_lo", Some(("lo", 0))), + ], + }; + ImportedIntrinsic { + source_record: recipe.source_record.into(), + llvm_name: recipe.llvm_symbol.into(), + arguments: recipe + .llvm_arguments + .iter() + .map(|argument| (*argument).into()) + .collect(), + results: vec!["i32".into()], + classes: vec!["NVVMPureIntrinsic".into()], + properties: recipe + .llvm_properties + .iter() + .map(|property| (*property).into()) + .collect(), + selections, + } +} + +pub(super) fn dot_product_evidence(policy: &OverlayIntrinsic) -> EvidenceRecord { + let mut record = evidence(); + record.id = policy.id.clone(); + record.source_record = policy.source_record.clone(); + record.llvm_symbol = policy.llvm_symbol.clone(); + record.llvm_arguments = policy.llvm_arguments.clone(); + record.llvm_results = policy.llvm_results.clone(); + record.concrete_llvm_arguments = policy.llvm_arguments.clone(); + record.concrete_llvm_results = policy.llvm_results.clone(); + record.declaration_attributes_canonicalized = Some(true); + record.gpu_target = "sm_61".into(); + record.ptx_feature = "+ptx50".into(); + record.expected_ptx = policy.expected_ptx.clone(); + record +} + +pub(super) fn validate_ptx_native_policy(policy: &OverlayIntrinsic) -> Result<()> { + let source = resolve_policy_source(policy)?; + validate_policy(policy, &source, None, 1) +} + +pub(super) fn ledger_entry(record: &OverlayIntrinsic) -> AbiLedgerEntry { + AbiLedgerEntry { + abi_id: record.abi_id.clone(), + status: "active".into(), + catalog_id: record.id.clone(), + operation_key: record.operation_key.clone(), + raw_rust_signature: raw_rust_signature(record), + } +} + +pub(super) fn ledger(entries: Vec) -> AbiLedgerFile { + AbiLedgerFile { + schema: 1, + intrinsic_abi: 1, + entries, + } +} + +pub(super) fn test_f8f6f4_admission() -> SparseMmaF8F6F4Admission { + let formats = vec![ + SparseMmaElement::E2m1, + SparseMmaElement::E2m3, + SparseMmaElement::E3m2, + SparseMmaElement::E4m3, + SparseMmaElement::E5m2, + ]; + SparseMmaF8F6F4Admission { + llvm_evidence_profile: "llvm-test".into(), + libnvvm_evidence_profile: "libnvvm-test".into(), + runtime_validation: RuntimeValidation::Unexecuted, + a_elements: formats.clone(), + b_elements: formats, + product_count: 25, + } +} + +pub(super) fn test_sparse_mma_f8f6f4_f16_admission() -> SparseMmaF8F6F4F16Admission { + SparseMmaF8F6F4F16Admission { + llvm_evidence_profile: "llvm-test".into(), + libnvvm_evidence_profile: "libnvvm-test".into(), + runtime_validation: RuntimeValidation::Unexecuted, + _legacy_first_abi_id: Some("i0525".into()), + a_elements: SPARSE_MMA_F8F6F4_ELEMENTS.into(), + b_elements: SPARSE_MMA_F8F6F4_ELEMENTS.into(), + product_count: 25, + } +} + +pub(super) fn test_register_mma_f8f6f4_admission( + _accumulator: RegisterMmaAccumulator, +) -> RegisterMmaF8F6F4Admission { + let formats = REGISTER_MMA_F8F6F4_ELEMENTS.to_vec(); + RegisterMmaF8F6F4Admission { + llvm_evidence_profile: "llvm-test".into(), + libnvvm_evidence_profile: "libnvvm-test".into(), + runtime_validation: RuntimeValidation::Unexecuted, + _legacy_first_abi_id: None, + a_elements: formats.clone(), + b_elements: formats, + product_count: 25, + targets: ["sm_120a", "sm_120f", "sm_121a", "sm_121f"] + .map(Into::into) + .into(), + } +} + +pub(super) fn test_register_mma_fp8_admission() -> RegisterMmaFp8Admission { + RegisterMmaFp8Admission { + llvm_evidence_profile: "llvm-fp8-test".into(), + libnvvm_evidence_profile: "libnvvm-fp8-test".into(), + runtime_validation: RuntimeValidation::Unexecuted, + _legacy_first_abi_id: Some("i0504".into()), + shapes: REGISTER_MMA_FP8_SHAPES.into(), + accumulators: REGISTER_MMA_FP8_ACCUMULATORS.into(), + a_elements: REGISTER_MMA_FP8_ELEMENTS.into(), + b_elements: REGISTER_MMA_FP8_ELEMENTS.into(), + product_count: 16, + } +} + +pub(super) fn test_register_mma_ampere_float_admission() -> RegisterMmaAmpereFloatAdmission { + RegisterMmaAmpereFloatAdmission { + llvm_evidence_profile: "llvm-ampere-float-test".into(), + libnvvm_evidence_profile: "libnvvm-ampere-float-test".into(), + runtime_validation: RuntimeValidation::Unexecuted, + _legacy_first_abi_id: Some("i0520".into()), + product_count: 5, + variants: REGISTER_MMA_AMPERE_FLOAT_VARIANTS.into(), + } +} + +pub(super) fn test_prmt_admission() -> PrmtAdmission { + let variants = [ + PrmtMode::Generic, + PrmtMode::F4e, + PrmtMode::B4e, + PrmtMode::Rc8, + PrmtMode::Ecl, + PrmtMode::Ecr, + PrmtMode::Rc16, + ] + .map(|mode| crate::model::PrmtAdmissionVariant { + abi_id: prmt_recipe(mode).abi_id.into(), + mode, + }) + .into(); + PrmtAdmission { + llvm_evidence_profile: "llvm-test".into(), + libnvvm_evidence_profile: "libnvvm-test".into(), + runtime_validation: RuntimeValidation::Unexecuted, + variants, + } +} + +pub(super) fn test_fp8_conversion_admission() -> PackedConversionFp8Admission { + PackedConversionFp8Admission { + llvm_evidence_profile: "llvm-test".into(), + libnvvm_evidence_profile: "libnvvm-test".into(), + runtime_validation: RuntimeValidation::Unexecuted, + destination_formats: vec![ + PackedConversionDestinationFormat::E4m3x2, + PackedConversionDestinationFormat::E5m2x2, + ], + saturations: vec![ + PackedConversionSaturation::Satfinite, + PackedConversionSaturation::SatfiniteRelu, + ], + product_count: 4, + } +} + +pub(super) fn test_fp8_f16x2_conversion_admission() -> PackedConversionFp8F16x2Admission { + PackedConversionFp8F16x2Admission { + llvm_evidence_profile: "llvm-test".into(), + libnvvm_evidence_profile: "libnvvm-test".into(), + runtime_validation: RuntimeValidation::Unexecuted, + fp8_formats: vec![ + PackedConversionFp8Format::E4m3x2, + PackedConversionFp8Format::E5m2x2, + ], + directions: vec![ + PackedConversionFp8Direction::Pack, + PackedConversionFp8Direction::Unpack, + ], + relu_variants: true, + product_count: 8, + } +} + +pub(super) fn test_cluster_barrier_admission() -> ClusterBarrierAdmission { + let variants = [ + ClusterBarrierMode::Arrive, + ClusterBarrierMode::ArriveAligned, + ClusterBarrierMode::ArriveRelaxed, + ClusterBarrierMode::ArriveRelaxedAligned, + ClusterBarrierMode::Wait, + ClusterBarrierMode::WaitAligned, + ] + .map(|mode| crate::model::ClusterBarrierAdmissionVariant { + abi_id: cluster_barrier_recipe(mode).abi_id.into(), + mode, + }) + .into(); + ClusterBarrierAdmission { + llvm_evidence_profile: "llvm-test".into(), + libnvvm_evidence_profile: "libnvvm-test".into(), + runtime_validation: RuntimeValidation::Unexecuted, + variants, + } +} + +pub(super) fn test_wgmma_control_admission() -> WgmmaControlAdmission { + let variants = [ + WgmmaControlMode::Fence, + WgmmaControlMode::CommitGroup, + WgmmaControlMode::WaitGroup, + ] + .map(|mode| crate::model::WgmmaControlAdmissionVariant { + abi_id: wgmma_control_recipe(mode).abi_id.into(), + mode, + }) + .into(); + WgmmaControlAdmission { + llvm_evidence_profile: "llvm-test".into(), + libnvvm_evidence_profile: "libnvvm-test".into(), + runtime_validation: RuntimeValidation::Unexecuted, + variants, + } +} + +pub(super) fn test_special_register_admission() -> SpecialRegisterAdmission { + SpecialRegisterAdmission { + llvm_evidence_profile: "rust-llvm-22.1.2-1cb4e383".into(), + libnvvm_evidence_profile: "cuda-13.3-libnvvm-13.3.33".into(), + runtime_validation: RuntimeValidation::Unexecuted, + registers: REVIEWED_SPECIAL_REGISTERS.into(), + product_count: REVIEWED_SPECIAL_REGISTERS.len(), + } +} + +pub(super) fn test_debug_control_admission() -> DebugControlAdmission { + DebugControlAdmission { + llvm_evidence_profile: "llvm-debug-test".into(), + libnvvm_evidence_profile: "libnvvm-debug-test".into(), + runtime_validation: RuntimeValidation::Unexecuted, + operations: vec![ + DebugControlOperation::Trap, + DebugControlOperation::Breakpoint, + DebugControlOperation::Pmevent, + ], + abi_ids: vec!["i9001".into(), "i9002".into(), "i9003".into()], + } +} + +pub(super) fn test_clc_admission() -> ClcAdmission { + let operations = [ + ClcOperation::TryCancel, + ClcOperation::TryCancelMulticast, + ClcOperation::QueryIsCanceled, + ClcOperation::QueryGetFirstCtaidX, + ClcOperation::QueryGetFirstCtaidY, + ClcOperation::QueryGetFirstCtaidZ, + ]; + ClcAdmission { + llvm_evidence_profile: "llvm-clc-test".into(), + libnvvm_evidence_profile: "libnvvm-clc-test".into(), + runtime_validation: RuntimeValidation::Unexecuted, + variants: operations + .into_iter() + .map(|operation| crate::model::ClcAdmissionVariant { + abi_id: clc_recipe(operation).abi_id.into(), + operation, + }) + .collect(), + } +} + +pub(super) fn test_tma_admission() -> TmaAdmission { + TmaAdmission { + llvm_evidence_profile: "llvm-tma-test".into(), + libnvvm_evidence_profile: "libnvvm-tma-test".into(), + reduce_llvm_evidence_profile: Some("llvm-tma-reduce-test".into()), + reduce_libnvvm_evidence_profile: Some("libnvvm-tma-reduce-test".into()), + runtime_validation: RuntimeValidation::Unexecuted, + variants: TMA_OPERATIONS + .into_iter() + .map(|operation| crate::model::TmaAdmissionVariant { + abi_id: tma_recipe(operation).abi_id.into(), + operation, + }) + .collect(), + reduce_variants: tma_reduction_admission_variants(), + } +} + +pub(super) fn test_tcgen05_admission() -> Tcgen05Admission { + let operations = [ + Tcgen05Operation::Alloc, + Tcgen05Operation::Dealloc, + Tcgen05Operation::RelinquishAllocPermit, + Tcgen05Operation::FenceBeforeThreadSync, + Tcgen05Operation::FenceAfterThreadSync, + Tcgen05Operation::Commit, + Tcgen05Operation::CommitSharedCluster, + Tcgen05Operation::MmaWsF16, + Tcgen05Operation::MmaF16, + Tcgen05Operation::MmaWsBf16, + Tcgen05Operation::MmaWsTf32, + Tcgen05Operation::CpSmemToTmem, + Tcgen05Operation::Ld16x256bX8Pure, + Tcgen05Operation::Ld16x256bPure, + Tcgen05Operation::LoadWait, + Tcgen05Operation::StoreWait, + Tcgen05Operation::AllocCg2, + Tcgen05Operation::DeallocCg2, + Tcgen05Operation::RelinquishAllocPermitCg2, + Tcgen05Operation::MmaF16Cg2, + Tcgen05Operation::CommitCg2, + Tcgen05Operation::CommitSharedClusterCg2, + Tcgen05Operation::CommitMulticastCg2, + Tcgen05Operation::CpSmemToTmemCg2, + Tcgen05Operation::CommitMulticast, + Tcgen05Operation::ShiftDown, + Tcgen05Operation::ShiftDownCg2, + ]; + Tcgen05Admission { + llvm_evidence_profile: "llvm-tcgen05-test".into(), + libnvvm_evidence_profile: "libnvvm-tcgen05-test".into(), + cp_llvm_evidence_profile: None, + cp_libnvvm_evidence_profile: None, + ld_llvm_evidence_profile: None, + ld_libnvvm_evidence_profile: None, + st_llvm_evidence_profile: None, + st_libnvvm_evidence_profile: None, + offset_llvm_evidence_profile: None, + offset_libnvvm_evidence_profile: None, + control_llvm_evidence_profile: Some("llvm-tcgen05-control-test".into()), + control_libnvvm_evidence_profile: Some("libnvvm-tcgen05-control-test".into()), + mma_llvm_evidence_profile: None, + mma_libnvvm_evidence_profile: None, + mma_llvm_target_contracts: vec![], + mma_libnvvm_target_contracts: vec![], + runtime_validation: RuntimeValidation::Unexecuted, + variants: operations + .into_iter() + .map(|operation| crate::model::Tcgen05AdmissionVariant { + abi_id: tcgen05_recipe(operation).abi_id.into(), + operation, + }) + .collect(), + cp_variants: vec![], + ld_variants: vec![], + st_variants: vec![], + ld_offset_variants: vec![], + st_offset_variants: vec![], + mma_variants: vec![], + } +} + +pub(super) fn without_tcgen05_control(mut admission: Tcgen05Admission) -> Tcgen05Admission { + admission.variants.truncate(24); + admission.control_llvm_evidence_profile = None; + admission.control_libnvvm_evidence_profile = None; + admission +} + +pub(super) fn test_tcgen05_cp_admission() -> Tcgen05Admission { + let mut admission = test_tcgen05_admission(); + admission.cp_llvm_evidence_profile = Some("llvm-tcgen05-cp-test".into()); + admission.cp_libnvvm_evidence_profile = Some("libnvvm-tcgen05-cp-test".into()); + admission.cp_variants = TCGEN05_CP_MEMBERS + .into_iter() + .flat_map(|member| { + [Tcgen05CpGroup::Cg1, Tcgen05CpGroup::Cg2] + .into_iter() + .map(move |group| (member, group)) + }) + .enumerate() + .map(|(index, (member, group))| Tcgen05CpAdmissionVariant { + abi_id: format!("i{:04}", 578 + index), + member, + group, + }) + .collect(); + admission +} + +pub(super) fn test_tcgen05_ld_admission() -> Tcgen05Admission { + let mut admission = test_tcgen05_cp_admission(); + admission.ld_llvm_evidence_profile = Some("llvm-tcgen05-ld-test".into()); + admission.ld_libnvvm_evidence_profile = Some("libnvvm-tcgen05-ld-test".into()); + admission.ld_variants = TCGEN05_LD_VARIANTS + .into_iter() + .flat_map(|(shape, multiplicity)| { + [false, true] + .into_iter() + .map(move |pack16| (shape, multiplicity, pack16)) + }) + .enumerate() + .map( + |(index, (shape, multiplicity, pack16))| Tcgen05LdAdmissionVariant { + abi_id: format!("i{:04}", 612 + index), + shape, + multiplicity, + pack16, + }, + ) + .collect(); + admission +} + +pub(super) fn test_tcgen05_st_admission() -> Tcgen05Admission { + let mut admission = test_tcgen05_ld_admission(); + admission.st_llvm_evidence_profile = Some("llvm-tcgen05-st-test".into()); + admission.st_libnvvm_evidence_profile = Some("libnvvm-tcgen05-st-test".into()); + admission.st_variants = TCGEN05_ST_VARIANTS + .into_iter() + .flat_map(|(shape, multiplicity)| { + [false, true] + .into_iter() + .map(move |unpack16| (shape, multiplicity, unpack16)) + }) + .enumerate() + .map( + |(index, (shape, multiplicity, unpack16))| Tcgen05StAdmissionVariant { + abi_id: format!("i{:04}", 670 + index), + shape, + multiplicity, + unpack16, + }, + ) + .collect(); + admission +} + +pub(super) fn test_tcgen05_offset_admission() -> Tcgen05Admission { + let mut admission = test_tcgen05_st_admission(); + admission.offset_llvm_evidence_profile = Some("llvm-tcgen05-offset-test".into()); + admission.offset_libnvvm_evidence_profile = Some("libnvvm-tcgen05-offset-test".into()); + admission.ld_offset_variants = TCGEN05_OFFSET_LDST_VARIANTS + .into_iter() + .flat_map(|(shape, multiplicity)| { + [false, true] + .into_iter() + .map(move |pack16| (shape, multiplicity, pack16)) + }) + .enumerate() + .map( + |(index, (shape, multiplicity, pack16))| Tcgen05LdAdmissionVariant { + abi_id: format!("i{:04}", 728 + index), + shape, + multiplicity, + pack16, + }, + ) + .collect(); + admission.st_offset_variants = TCGEN05_OFFSET_LDST_VARIANTS + .into_iter() + .flat_map(|(shape, multiplicity)| { + [false, true] + .into_iter() + .map(move |unpack16| (shape, multiplicity, unpack16)) + }) + .enumerate() + .map( + |(index, (shape, multiplicity, unpack16))| Tcgen05StAdmissionVariant { + abi_id: format!("i{:04}", 744 + index), + shape, + multiplicity, + unpack16, + }, + ) + .collect(); + admission +} + +pub(super) fn test_tcgen05_mma_admission() -> Tcgen05Admission { + let mut admission = test_tcgen05_admission(); + admission.mma_llvm_evidence_profile = Some("llvm-tcgen05-mma-test".into()); + admission.mma_libnvvm_evidence_profile = Some("libnvvm-tcgen05-mma-test".into()); + admission.mma_llvm_target_contracts = + expected_tcgen05_mma_target_contracts(IntrinsicBackend::LlvmNvptx); + admission.mma_libnvvm_target_contracts = + expected_tcgen05_mma_target_contracts(IntrinsicBackend::LibNvvm); + admission.mma_variants = TCGEN05_MMA_FORMS + .into_iter() + .map(|form| (form, None)) + .chain( + TCGEN05_MMA_ALIASES + .into_iter() + .map(|alias| (Tcgen05MmaForm::WsTensor, Some(alias))), + ) + .enumerate() + .map(|(index, (form, alias))| Tcgen05MmaAdmissionVariant { + abi_id: format!("i{:04}", 763 + index), + form, + alias, + }) + .chain( + TCGEN05_MMA_ALIASES + .into_iter() + .enumerate() + .map(|(index, alias)| Tcgen05MmaAdmissionVariant { + abi_id: format!("i{:04}", 1011 + index), + form: Tcgen05MmaForm::Shared, + alias: Some(alias), + }), + ) + .collect(); + admission +} + +pub(super) fn assert_tcgen05_backend_target_split(record: &OverlayIntrinsic) { + let llvm = &record.backend_lowerings[0]; + let libnvvm = &record.backend_lowerings[1]; + assert_eq!(record.targets, TCGEN05_LLVM_TARGETS); + assert_eq!(llvm.backend, IntrinsicBackend::LlvmNvptx); + assert_eq!(llvm.targets, None); + assert_eq!(libnvvm.backend, IntrinsicBackend::LibNvvm); + assert_eq!(libnvvm.targets.as_deref(), Some(TCGEN05_LIBNVVM_TARGETS)); + assert_eq!( + backend_target_requirement(record, llvm).unwrap().hardware, + CatalogHardwareTarget::AnyOf { + alternatives: vec![ + CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 101 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 103 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, + ], + } + ); + assert_eq!( + backend_target_requirement(record, libnvvm) + .unwrap() + .hardware, + CatalogHardwareTarget::AnyOf { + alternatives: vec![ + CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 103 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, + ], + } + ); +} + +pub(super) fn test_threadfence_admission() -> ThreadfenceAdmission { + ThreadfenceAdmission { + llvm_evidence_profile: "llvm-test".into(), + libnvvm_evidence_profile: "libnvvm-test".into(), + runtime_validation: RuntimeValidation::Unexecuted, + variants: vec![ + crate::model::ThreadfenceAdmissionVariant { + abi_id: "i0298".into(), + scope: ThreadfenceScope::Cta, + }, + crate::model::ThreadfenceAdmissionVariant { + abi_id: "i0299".into(), + scope: ThreadfenceScope::Device, + }, + crate::model::ThreadfenceAdmissionVariant { + abi_id: "i0300".into(), + scope: ThreadfenceScope::System, + }, + ], + } +} + +pub(super) fn test_cluster_memory_admission() -> ClusterMemoryAdmission { + ClusterMemoryAdmission { + llvm_evidence_profile: "llvm-cluster-memory-test".into(), + libnvvm_evidence_profile: "libnvvm-cluster-memory-test".into(), + runtime_validation: RuntimeValidation::Unexecuted, + variants: vec![ + crate::model::ClusterMemoryAdmissionVariant { + abi_id: "i0320".into(), + operation: ClusterMemoryOperation::MapSharedRank, + }, + crate::model::ClusterMemoryAdmissionVariant { + abi_id: "i0321".into(), + operation: ClusterMemoryOperation::ReadU32, + }, + ], + } +} + +pub(super) fn test_stmatrix_admission() -> StmatrixAdmission { + let variants = [ + (StmatrixMultiplicity::X2, StmatrixLayout::Normal, "i0301"), + ( + StmatrixMultiplicity::X2, + StmatrixLayout::Transposed, + "i0302", + ), + (StmatrixMultiplicity::X4, StmatrixLayout::Normal, "i0303"), + ( + StmatrixMultiplicity::X4, + StmatrixLayout::Transposed, + "i0304", + ), + ] + .map( + |(multiplicity, layout, abi_id)| crate::model::StmatrixAdmissionVariant { + abi_id: abi_id.into(), + multiplicity, + layout, + }, + ) + .into(); + StmatrixAdmission { + llvm_evidence_profile: "llvm-test".into(), + libnvvm_evidence_profile: "libnvvm-test".into(), + runtime_validation: RuntimeValidation::Unexecuted, + variants, + } +} + +pub(super) fn test_mbarrier_extended_admission() -> MbarrierExtendedAdmission { + let variants = [ + MbarrierExtendedOperation::ArriveExpectTxCta, + MbarrierExtendedOperation::ArriveExpectTxCluster, + MbarrierExtendedOperation::ArriveRemoteCluster, + MbarrierExtendedOperation::TryWaitTokenCta, + MbarrierExtendedOperation::TryWaitParityCta, + MbarrierExtendedOperation::TryWaitParityCluster, + MbarrierExtendedOperation::FenceProxyAsyncSharedCta, + MbarrierExtendedOperation::FenceMbarrierInitReleaseCluster, + MbarrierExtendedOperation::FenceProxyAsyncGenericReleaseSharedCtaCluster, + MbarrierExtendedOperation::FenceProxyAsyncGenericAcquireSharedClusterCluster, + MbarrierExtendedOperation::Nanosleep, + ] + .map(|operation| crate::model::MbarrierExtendedAdmissionVariant { + abi_id: mbarrier_extended_recipe(operation).abi_id.into(), + operation, + }) + .into(); + MbarrierExtendedAdmission { + llvm_evidence_profile: "llvm-test".into(), + libnvvm_evidence_profile: "libnvvm-test".into(), + runtime_validation: RuntimeValidation::Unexecuted, + variants, + } +} + +pub(super) struct CandidateTestRepo(pub(super) PathBuf); + +impl Drop for CandidateTestRepo { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.0); + } +} + +pub(super) fn repo_without_evidence() -> CandidateTestRepo { + use std::sync::atomic::{AtomicU64, Ordering}; + + static NEXT: AtomicU64 = AtomicU64::new(0); + let source = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let root = std::env::temp_dir().join(format!( + "cuda-intrinsics-candidate-{}-{}", + std::process::id(), + NEXT.fetch_add(1, Ordering::Relaxed) + )); + let input = root.join("intrinsics"); + fs::create_dir_all(input.join("overlay")).unwrap(); + for name in [ + "upstream.lock", + "imported.json", + "overlay.toml", + "abi-v1.toml", + ] { + fs::copy(source.join("intrinsics").join(name), input.join(name)).unwrap(); + } + for entry in fs::read_dir(source.join("intrinsics/overlay")).unwrap() { + let entry = entry.unwrap(); + if entry.path().extension().and_then(|value| value.to_str()) == Some("toml") { + fs::copy(entry.path(), input.join("overlay").join(entry.file_name())).unwrap(); + } + } + CandidateTestRepo(root) +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/tests/misc_families.rs b/crates/cuda-intrinsics-gen/src/resolve/tests/misc_families.rs new file mode 100644 index 0000000000..42777630a3 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/tests/misc_families.rs @@ -0,0 +1,1123 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::extract::IMPORTED_SCHEMA; +use crate::model::{ + BackendLoweringMechanism, CpAsyncAdapter, CpAsyncCachePolicy, CpAsyncControlAdapter, + CpAsyncControlOperation, CpAsyncCopySize, CpAsyncMbarrierAdapter, CpAsyncMbarrierOperation, + CpAsyncMbarrierStateSpace, CpAsyncSourceSize, DebugControlAdapter, DebugControlOperation, + DotProductAdapter, DotProductOperation, DotProductSignedness, ImportedFile, ImportedIntrinsic, + IntrinsicBackend, IntrinsicSource, MbarrierBasicAdapter, MbarrierBasicOperation, + MbarrierExtendedAdapter, MbarrierExtendedOperation, MbarrierExtendedSourceContract, + OverlayIntrinsic, OverlayShardFile, ReduxAdapter, ReduxOperation, RuntimeValidation, + WgmmaControlAdapter, WgmmaControlMode, +}; +use crate::ptx::OperandPattern; +use crate::util::read_json; +use std::collections::BTreeMap; +use std::fs; +use std::path::Path; + +use super::fixtures::*; +use crate::resolve::evidence::*; +use crate::resolve::families::*; +use crate::resolve::guards::*; +use crate::resolve::materialize::*; +use crate::resolve::overlay::*; +use crate::resolve::policy::*; +use crate::resolve::targets::*; + +#[test] +fn f32_redux_recipes_match_pinned_llvm_records() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let declarations: BTreeMap<_, _> = imported + .intrinsics + .iter() + .map(|record| (record.source_record.as_str(), record)) + .collect(); + let cases = [ + ( + ReduxOperation::Fmin, + "redux.sync.min.f32 \t$dst, $src, $mask;", + &["sync", "min", "f32"][..], + ), + ( + ReduxOperation::FminNan, + "redux.sync.min.NaN.f32 \t$dst, $src, $mask;", + &["sync", "min", "NaN", "f32"], + ), + ( + ReduxOperation::FminAbs, + "redux.sync.min.abs.f32 \t$dst, $src, $mask;", + &["sync", "min", "abs", "f32"], + ), + ( + ReduxOperation::FminAbsNan, + "redux.sync.min.abs.NaN.f32 \t$dst, $src, $mask;", + &["sync", "min", "abs", "NaN", "f32"], + ), + ( + ReduxOperation::Fmax, + "redux.sync.max.f32 \t$dst, $src, $mask;", + &["sync", "max", "f32"], + ), + ( + ReduxOperation::FmaxNan, + "redux.sync.max.NaN.f32 \t$dst, $src, $mask;", + &["sync", "max", "NaN", "f32"], + ), + ( + ReduxOperation::FmaxAbs, + "redux.sync.max.abs.f32 \t$dst, $src, $mask;", + &["sync", "max", "abs", "f32"], + ), + ( + ReduxOperation::FmaxAbsNan, + "redux.sync.max.abs.NaN.f32 \t$dst, $src, $mask;", + &["sync", "max", "abs", "NaN", "f32"], + ), + ]; + + for (operation, expected_asm, expected_modifiers) in cases { + let recipe = redux_recipe(operation); + let declaration = declarations[recipe.source_record]; + assert_eq!(declaration.llvm_name, recipe.llvm_symbol); + assert_eq!(declaration.arguments, ["f32", "i32"]); + assert_eq!(declaration.results, ["f32"]); + assert_eq!(declaration.selections.len(), 1); + assert_eq!(declaration.selections[0].asm, expected_asm); + assert_eq!( + declaration.selections[0].predicates, + ["Subtarget->hasReduxSyncF32()"] + ); + assert_eq!(recipe.ptx_modifiers, expected_modifiers); + assert_eq!(recipe.minimum_ptx, "8.6"); + assert_eq!(recipe.minimum_sm, None); + assert_eq!(recipe.targets, REDUX_F32_TARGETS); + } +} + +#[test] +fn compact_debug_control_admission_is_closed() { + let records = expand_debug_control_admission(&test_debug_control_admission()).unwrap(); + assert_eq!(records.len(), 3); + assert_eq!(records[0].id, "trap"); + assert_eq!(records[1].id, "breakpoint"); + assert_eq!(records[2].id, "pmevent"); + for record in &records { + validate_ptx_native_policy(record).unwrap(); + assert_eq!(record.backend_lowerings.len(), 2); + assert!( + record + .debug_control + .as_ref() + .is_some_and(|debug| debug.runtime_validation == RuntimeValidation::Unexecuted) + ); + } + assert_eq!(records[0].minimum_ptx, "1.0"); + assert_eq!(records[0].minimum_sm, None); + assert_eq!(records[1].minimum_ptx, "1.0"); + assert_eq!(records[1].minimum_sm.as_deref(), Some("sm_11")); + assert_eq!(records[2].minimum_ptx, "1.4"); + assert_eq!( + records[2].expected_ptx.operands, + [OperandPattern::Immediate] + ); + + let mut pending = test_debug_control_admission(); + pending.abi_ids.clear(); + assert!(expand_debug_control_admission(&pending).is_err()); + + let mut missing = test_debug_control_admission(); + missing.operations.pop(); + assert!(expand_debug_control_admission(&missing).is_err()); + + let mut duplicate_operation = test_debug_control_admission(); + duplicate_operation.operations[2] = DebugControlOperation::Breakpoint; + assert!(expand_debug_control_admission(&duplicate_operation).is_err()); + + let mut duplicate_id = test_debug_control_admission(); + duplicate_id.abi_ids[2] = duplicate_id.abi_ids[1].clone(); + assert!(expand_debug_control_admission(&duplicate_id).is_err()); + + let mut malformed_id = test_debug_control_admission(); + malformed_id.abi_ids[0] = "debug1".into(); + assert!(expand_debug_control_admission(&malformed_id).is_err()); + + let mut executed = test_debug_control_admission(); + executed.runtime_validation = RuntimeValidation::Executed; + assert!(expand_debug_control_admission(&executed).is_err()); + + let mut wrong_source = records[0].clone(); + wrong_source.source = Some(IntrinsicSource::LlvmImported { + source_record: "invented".into(), + }); + assert!(validate_ptx_native_policy(&wrong_source).is_err()); + + let mut wrong_adapter = records[2].clone(); + wrong_adapter.debug_control.as_mut().unwrap().adapter = DebugControlAdapter::Direct; + assert!(validate_ptx_native_policy(&wrong_adapter).is_err()); + + let mut wrong_immediate = records[2].clone(); + wrong_immediate.expected_ptx.operands = vec![OperandPattern::Register]; + assert!(validate_ptx_native_policy(&wrong_immediate).is_err()); + + let mut wrong_floor = records[1].clone(); + wrong_floor.backend_lowerings[0].minimum_sm = Some("sm_75".into()); + assert!(validate_ptx_native_policy(&wrong_floor).is_err()); +} + +#[test] +fn debug_control_compact_schema_is_reserved_for_aggregation() { + let shard = |schema| OverlayShardFile { + schema, + family: "debug_control".into(), + intrinsics: vec![], + register_mma_int4: None, + register_mma_int8: None, + register_mma_b1: None, + register_mma_f8f6f4_f32: None, + register_mma_f8f6f4_f16: None, + register_mma_mxf8f6f4_f32: None, + register_mma_fp8: None, + register_mma_ampere_float: None, + sparse_mma_integer: None, + sparse_mma_f8f6f4_f32: None, + sparse_mma_f8f6f4_f16: None, + prmt: None, + packed_conversion_fp8: None, + packed_conversion_fp8_f16x2: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cluster_sreg: None, + cluster_barrier: None, + mbarrier_extended: None, + special_registers: None, + debug_control: Some(test_debug_control_admission()), + threadfence: None, + cluster_memory: None, + stmatrix: None, + clc: None, + wgmma_controls: None, + tma: None, + tcgen05: None, + }; + let path = Path::new("intrinsics/overlay/debug_control.toml"); + validate_overlay_shard_schema_with_max(&shard(33), path, 33).unwrap(); + assert!(validate_overlay_shard_schema_with_max(&shard(33), path, 32).is_err()); + let error = validate_overlay_shard_schema_with_max(&shard(32), path, 33).unwrap_err(); + assert!( + error + .to_string() + .contains("requires overlay shard schema 33") + ); +} + +#[test] +fn active_debug_control_sources_parse_and_prove_both_backend_routes() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let shard_path = repo_root.join("intrinsics/overlay/debug_control.toml"); + let shard: OverlayShardFile = + toml::from_str(&fs::read_to_string(&shard_path).unwrap()).unwrap(); + validate_overlay_shard_schema_with_max(&shard, &shard_path, 33).unwrap(); + let admission = shard.debug_control.unwrap(); + assert_eq!(admission.abi_ids, ["i0295", "i0296", "i0297"]); + let records = expand_debug_control_admission(&admission).unwrap(); + + let evidence = vec![ + read_evidence_file( + &repo_root.join("intrinsics/evidence/rust-llvm-22.1.2-1cb4e383-debug-control.json"), + ) + .unwrap(), + read_evidence_file( + &repo_root.join("intrinsics/evidence/cuda-13.3-libnvvm-13.3.33-debug-control.json"), + ) + .unwrap(), + ]; + let indexed = index_evidence(&evidence, "1cb4e3833c1919c2e6fb579a23ac0e2b22587b7e").unwrap(); + for record in &records { + let routes = resolve_backend_lowerings(record, &indexed).unwrap(); + assert_eq!(routes.len(), 2); + assert!(routes.iter().all(|route| { + route.mechanism == BackendLoweringMechanism::InlinePtx && route.status == "validated" + })); + } +} + +#[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 + .iter() + .map(|record| record.abi_id.as_str()) + .collect::>(), + (306..=316) + .map(|id| format!("i{id:04}")) + .collect::>() + ); + + let imported: ImportedFile = read_json( + &Path::new(env!("CARGO_MANIFEST_DIR")) + .join("../..") + .join("intrinsics/imported.json"), + ) + .unwrap(); + for record in &records { + let contract = record.mbarrier_extended.as_ref().unwrap(); + let (template, constraints) = mbarrier_extended_inline_recipe(contract.operation); + assert!(template.ends_with(';') || template.ends_with("; }")); + assert!(constraints.contains("~{memory}")); + assert!(record.convergent && record.memory == "read_write"); + assert!( + record + .backend_lowerings + .iter() + .all(|lowering| { lowering.mechanism == BackendLoweringMechanism::InlinePtx }) + ); + match contract.source_contract { + MbarrierExtendedSourceContract::LlvmImported => { + let declaration = imported + .intrinsics + .iter() + .find(|declaration| { + Some(declaration.source_record.as_str()) == record.source_record.as_deref() + }) + .unwrap(); + validate_imported_policy(record, declaration).unwrap(); + } + MbarrierExtendedSourceContract::PtxNativeRawClusterAddress => { + validate_ptx_native_policy(record).unwrap(); + assert_eq!(record.rust_arguments, ["u64"]); + assert_eq!(record.dialect_operands, ["i64"]); + } + } + } + + let base = records + .iter() + .find(|record| record.id == "mbarrier_arrive_expect_tx") + .unwrap(); + let declaration = imported + .intrinsics + .iter() + .find(|declaration| { + Some(declaration.source_record.as_str()) == base.source_record.as_deref() + }) + .unwrap(); + let mut wrong_adapter = base.clone(); + wrong_adapter.mbarrier_extended.as_mut().unwrap().adapter = + MbarrierExtendedAdapter::PointerTokenToPredicate; + assert!(validate_imported_policy(&wrong_adapter, declaration).is_err()); + let mut wrong_floor = base.clone(); + wrong_floor.minimum_ptx = "8.6".into(); + assert!(validate_imported_policy(&wrong_floor, declaration).is_err()); + let mut lost_clobber = base.clone(); + lost_clobber.memory = "none".into(); + assert!(validate_imported_policy(&lost_clobber, declaration).is_err()); + + let remote = records + .iter() + .find(|record| record.id == "mbarrier_arrive_cluster") + .unwrap(); + let incompatible = imported + .intrinsics + .iter() + .find(|declaration| { + declaration.source_record == "int_nvvm_mbarrier_arrive_scope_cluster_space_cluster" + }) + .unwrap(); + assert_eq!(incompatible.arguments, ["shared_cluster_ptr", "i32"]); + assert!( + validate_mbarrier_extended_policy( + remote, + &IntrinsicSource::LlvmImported { + source_record: incompatible.source_record.clone(), + }, + Some(incompatible), + ) + .is_err() + ); + + let mut missing = test_mbarrier_extended_admission(); + missing.variants.pop(); + assert!(expand_mbarrier_extended_admission(&missing).is_err()); + let mut duplicate = test_mbarrier_extended_admission(); + duplicate.variants[10].operation = MbarrierExtendedOperation::ArriveExpectTxCta; + assert!(expand_mbarrier_extended_admission(&duplicate).is_err()); + let mut wrong_abi = test_mbarrier_extended_admission(); + wrong_abi.variants[0].abi_id = "i9999".into(); + assert!(expand_mbarrier_extended_admission(&wrong_abi).is_err()); +} + +#[test] +fn compact_wgmma_control_admission_and_semantics_fail_closed() { + let records = expand_wgmma_control_admission(&test_wgmma_control_admission()).unwrap(); + assert_eq!(records.len(), 3); + + let imported: ImportedFile = read_json( + &Path::new(env!("CARGO_MANIFEST_DIR")) + .join("../..") + .join("intrinsics/imported.json"), + ) + .unwrap(); + let declaration_for = |record: &OverlayIntrinsic| { + imported + .intrinsics + .iter() + .find(|declaration| { + Some(declaration.source_record.as_str()) == record.source_record.as_deref() + }) + .unwrap() + }; + for record in &records { + validate_imported_policy(record, declaration_for(record)).unwrap(); + assert_eq!(record.targets, "sm_90a"); + assert_eq!(record.minimum_ptx, "8.0"); + assert_eq!(record.backend_lowerings.len(), 2); + } + + let wait = records + .iter() + .find(|record| { + record + .wgmma_control + .as_ref() + .is_some_and(|control| control.mode == WgmmaControlMode::WaitGroup) + }) + .unwrap(); + assert_eq!(wait.rust_arguments, ["u64"]); + assert_eq!(wait.dialect_operands, ["i64"]); + assert_eq!( + wait.compatibility_rust_paths, + ["cuda_device::wgmma::__wgmma_wait_group"] + ); + + let mut wrong_adapter = wait.clone(); + wrong_adapter.wgmma_control.as_mut().unwrap().adapter = WgmmaControlAdapter::NoArguments; + assert!(validate_imported_policy(&wrong_adapter, declaration_for(wait)).is_err()); + + let mut wrong_participation = wait.clone(); + wrong_participation.execution_scope = "warp".into(); + assert!(validate_imported_policy(&wrong_participation, declaration_for(wait)).is_err()); + + let mut wrong_route = wait.clone(); + wrong_route.backend_lowerings[1].mechanism = BackendLoweringMechanism::TypedNvvm; + assert!(validate_imported_policy(&wrong_route, declaration_for(wait)).is_err()); + + let mut wrong_target = wait.clone(); + wrong_target.targets = "all".into(); + wrong_target.minimum_sm = Some("sm_90".into()); + assert!(validate_imported_policy(&wrong_target, declaration_for(wait)).is_err()); + + let mut missing = test_wgmma_control_admission(); + missing.variants.pop(); + assert!(expand_wgmma_control_admission(&missing).is_err()); + + let mut reversed = test_wgmma_control_admission(); + reversed.variants.reverse(); + assert!(expand_wgmma_control_admission(&reversed).is_err()); + + let mut duplicate = test_wgmma_control_admission(); + duplicate.variants[2].mode = WgmmaControlMode::Fence; + assert!(expand_wgmma_control_admission(&duplicate).is_err()); + + let mut wrong_abi = test_wgmma_control_admission(); + wrong_abi.variants[0].abi_id = "i9999".into(); + assert!(expand_wgmma_control_admission(&wrong_abi).is_err()); + + let mut executed = test_wgmma_control_admission(); + executed.runtime_validation = RuntimeValidation::Executed; + assert!(expand_wgmma_control_admission(&executed).is_err()); +} + +#[test] +fn wgmma_control_compact_schema_is_reserved_for_aggregation() { + let shard = OverlayShardFile { + schema: WGMMA_CONTROL_SHARD_SCHEMA, + family: "wgmma_control".into(), + intrinsics: vec![], + register_mma_int4: None, + register_mma_int8: None, + register_mma_b1: None, + register_mma_f8f6f4_f32: None, + register_mma_f8f6f4_f16: None, + register_mma_mxf8f6f4_f32: None, + register_mma_fp8: None, + register_mma_ampere_float: None, + sparse_mma_integer: None, + sparse_mma_f8f6f4_f32: None, + sparse_mma_f8f6f4_f16: None, + prmt: None, + packed_conversion_fp8: None, + packed_conversion_fp8_f16x2: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cluster_sreg: None, + cluster_barrier: None, + mbarrier_extended: None, + special_registers: None, + debug_control: None, + threadfence: None, + cluster_memory: None, + stmatrix: None, + clc: None, + wgmma_controls: Some(test_wgmma_control_admission()), + tma: None, + tcgen05: None, + }; + let path = Path::new("intrinsics/overlay/wgmma_control.toml"); + validate_overlay_shard_schema_with_max(&shard, path, WGMMA_CONTROL_SHARD_SCHEMA).unwrap(); + let mut old = shard; + old.schema -= 1; + assert!( + validate_overlay_shard_schema_with_max(&old, path, WGMMA_CONTROL_SHARD_SCHEMA) + .unwrap_err() + .to_string() + .contains("requires overlay shard schema 38") + ); +} + +#[test] +fn cp_async_copy_recipe_admits_only_classic_llvm_forms() { + let cases = [ + ( + CpAsyncCachePolicy::Ca, + CpAsyncCopySize::B4, + CpAsyncSourceSize::Full, + Some("cp_async_ca_4"), + ), + ( + CpAsyncCachePolicy::Ca, + CpAsyncCopySize::B4, + CpAsyncSourceSize::Runtime, + Some("cp_async_ca_zfill_4"), + ), + ( + CpAsyncCachePolicy::Ca, + CpAsyncCopySize::B8, + CpAsyncSourceSize::Full, + Some("cp_async_ca_8"), + ), + ( + CpAsyncCachePolicy::Ca, + CpAsyncCopySize::B8, + CpAsyncSourceSize::Runtime, + Some("cp_async_ca_zfill_8"), + ), + ( + CpAsyncCachePolicy::Ca, + CpAsyncCopySize::B16, + CpAsyncSourceSize::Full, + Some("cp_async_ca_16"), + ), + ( + CpAsyncCachePolicy::Ca, + CpAsyncCopySize::B16, + CpAsyncSourceSize::Runtime, + Some("cp_async_ca_zfill_16"), + ), + ( + CpAsyncCachePolicy::Cg, + CpAsyncCopySize::B4, + CpAsyncSourceSize::Full, + None, + ), + ( + CpAsyncCachePolicy::Cg, + CpAsyncCopySize::B4, + CpAsyncSourceSize::Runtime, + None, + ), + ( + CpAsyncCachePolicy::Cg, + CpAsyncCopySize::B8, + CpAsyncSourceSize::Full, + None, + ), + ( + CpAsyncCachePolicy::Cg, + CpAsyncCopySize::B8, + CpAsyncSourceSize::Runtime, + None, + ), + ( + CpAsyncCachePolicy::Cg, + CpAsyncCopySize::B16, + CpAsyncSourceSize::Full, + Some("cp_async_cg_16"), + ), + ( + CpAsyncCachePolicy::Cg, + CpAsyncCopySize::B16, + CpAsyncSourceSize::Runtime, + Some("cp_async_cg_zfill_16"), + ), + ]; + + for (cache_policy, copy_size, source_size, expected) in cases { + let copy = crate::model::CpAsyncCopy { + cache_policy, + copy_size, + source_size, + adapter: if source_size == CpAsyncSourceSize::Runtime { + CpAsyncAdapter::DirectPointersAndSourceSize + } else { + CpAsyncAdapter::DirectPointers + }, + runtime_validation: RuntimeValidation::Unexecuted, + }; + assert_eq!( + cp_async_copy_recipe(©).map(|recipe| recipe.id), + expected + ); + } +} + +#[test] +fn pinned_cp_async_records_match_the_closed_recipes() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let (overlay, _) = + read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")).unwrap(); + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let declarations: BTreeMap<_, _> = imported + .intrinsics + .iter() + .map(|record| (record.source_record.as_str(), record)) + .collect(); + let policies: Vec<_> = overlay + .intrinsics + .iter() + .filter(|record| matches!(record.family.as_str(), "cp_async_copy" | "cp_async_control")) + .collect(); + + assert_eq!(policies.len(), 11); + for policy in policies { + let declaration = declarations[policy.source_record.as_deref().unwrap()]; + validate_imported_policy(policy, declaration).unwrap(); + } +} + +#[test] +fn pinned_cp_async_mbarrier_records_match_the_closed_recipes() { + let records = pinned_cp_async_mbarrier_records(); + assert_eq!(records.len(), 4); + + for (policy, declaration) in records.values() { + validate_imported_policy(policy, declaration).unwrap(); + } +} + +#[test] +fn cp_async_mbarrier_recipes_fail_closed() { + let records = pinned_cp_async_mbarrier_records(); + let (arrive, declaration) = &records["cp_async_mbarrier_arrive"]; + let reject = |policy: &OverlayIntrinsic, declaration: &ImportedIntrinsic, expected: &str| { + let error = validate_imported_policy(policy, declaration).unwrap_err(); + let message = error.to_string(); + assert!(message.contains(expected), "unexpected error: {message}"); + }; + + let mut wrong_symbol = arrive.clone(); + wrong_symbol.llvm_symbol = Some("llvm.nvvm.cp.async.mbarrier.changed".into()); + reject(&wrong_symbol, declaration, "LLVM symbol mismatch"); + + let mut wrong_signature = arrive.clone(); + wrong_signature.rust_arguments = vec!["*const u64".into()]; + reject( + &wrong_signature, + declaration, + "closed cp.async mbarrier Rust API", + ); + + let mut wrong_operation = arrive.clone(); + wrong_operation + .cp_async_mbarrier + .as_mut() + .unwrap() + .operation = CpAsyncMbarrierOperation::ArriveNoInc; + reject(&wrong_operation, declaration, "identity does not match"); + + let mut wrong_state_space = arrive.clone(); + wrong_state_space + .cp_async_mbarrier + .as_mut() + .unwrap() + .state_space = CpAsyncMbarrierStateSpace::Shared; + reject(&wrong_state_space, declaration, "identity does not match"); + + let mut wrong_adapter = arrive.clone(); + wrong_adapter.cp_async_mbarrier.as_mut().unwrap().adapter = + CpAsyncMbarrierAdapter::PointerToVoid; + wrong_adapter.rust_result = "u64".into(); + reject( + &wrong_adapter, + declaration, + "closed cp.async mbarrier Rust API", + ); + + let mut executed_without_evidence = arrive.clone(); + executed_without_evidence + .cp_async_mbarrier + .as_mut() + .unwrap() + .runtime_validation = RuntimeValidation::Executed; + reject( + &executed_without_evidence, + declaration, + "unrecorded cp.async mbarrier runtime validation", + ); + + let mut wrong_properties = declaration.clone(); + wrong_properties.properties.pop(); + reject(arrive, &wrong_properties, "cp.async mbarrier properties"); + + let mut wrong_selection = declaration.clone(); + wrong_selection.selections[0].source_record = "CP_ASYNC_MBARRIER_CHANGED".into(); + reject( + arrive, + &wrong_selection, + "imported cp.async mbarrier selection changed", + ); + + let mut wrong_floor = arrive.clone(); + wrong_floor.minimum_sm = Some("sm_90".into()); + reject(&wrong_floor, declaration, "effects or target floor"); + + let mut wrong_llvm_route = arrive.clone(); + wrong_llvm_route + .backend_lowerings + .iter_mut() + .find(|lowering| lowering.backend == IntrinsicBackend::LlvmNvptx) + .unwrap() + .mechanism = BackendLoweringMechanism::InlinePtx; + reject(&wrong_llvm_route, declaration, "reviewed typed-LLVM"); + + let mut wrong_lib_route = arrive.clone(); + wrong_lib_route + .backend_lowerings + .iter_mut() + .find(|lowering| lowering.backend == IntrinsicBackend::LibNvvm) + .unwrap() + .mechanism = BackendLoweringMechanism::TypedNvvm; + reject(&wrong_lib_route, declaration, "reviewed typed-LLVM"); + + let mut mixed_family = arrive.clone(); + mixed_family.cp_async_control = Some(crate::model::CpAsyncControl { + operation: CpAsyncControlOperation::CommitGroup, + adapter: CpAsyncControlAdapter::NoOperands, + runtime_validation: RuntimeValidation::Unexecuted, + }); + reject( + &mixed_family, + declaration, + "mixes another generated-family contract", + ); +} + +#[test] +fn pinned_mbarrier_basic_records_match_the_closed_recipes() { + let records = pinned_mbarrier_basic_records(); + assert_eq!(records.len(), 4); + + for (policy, declaration) in records.values() { + validate_imported_policy(policy, declaration).unwrap(); + } +} + +#[test] +fn mbarrier_basic_recipes_fail_closed() { + let records = pinned_mbarrier_basic_records(); + let (init, init_declaration) = &records["mbarrier_init"]; + let reject = |policy: &OverlayIntrinsic, declaration: &ImportedIntrinsic, expected: &str| { + let error = validate_imported_policy(policy, declaration).unwrap_err(); + let message = error.to_string(); + assert!(message.contains(expected), "unexpected error: {message}"); + }; + + let mut wrong_symbol = init.clone(); + wrong_symbol.llvm_symbol = Some("llvm.nvvm.mbarrier.init.changed".into()); + reject(&wrong_symbol, init_declaration, "LLVM symbol mismatch"); + + let mut wrong_signature = init.clone(); + wrong_signature.rust_arguments = vec!["*mut u32".into(), "u32".into()]; + reject( + &wrong_signature, + init_declaration, + "unsafe mbarrier raw and compatibility API", + ); + + let mut wrong_operation = init.clone(); + wrong_operation.mbarrier_basic.as_mut().unwrap().operation = MbarrierBasicOperation::Arrive; + reject( + &wrong_operation, + init_declaration, + "operation, state space, and adapter disagree", + ); + + let mut wrong_adapter = init.clone(); + wrong_adapter.mbarrier_basic.as_mut().unwrap().adapter = + MbarrierBasicAdapter::InvalPointerToVoid; + reject( + &wrong_adapter, + init_declaration, + "operation, state space, and adapter disagree", + ); + + let mut executed_without_evidence = init.clone(); + executed_without_evidence + .mbarrier_basic + .as_mut() + .unwrap() + .runtime_validation = RuntimeValidation::Executed; + reject( + &executed_without_evidence, + init_declaration, + "unrecorded mbarrier runtime validation", + ); + + let mut wrong_properties = init_declaration.clone(); + wrong_properties.properties.pop(); + reject(init, &wrong_properties, "mbarrier properties"); + + let mut wrong_selection = init_declaration.clone(); + wrong_selection.selections[0].source_record = "MBARRIER_INIT_CHANGED".into(); + reject( + init, + &wrong_selection, + "imported mbarrier selection changed", + ); + + let mut wrong_ptx_floor = init.clone(); + wrong_ptx_floor.minimum_ptx = "7.1".into(); + reject( + &wrong_ptx_floor, + init_declaration, + "effects or target floor", + ); + + let mut wrong_sm_floor = init.clone(); + wrong_sm_floor.minimum_sm = Some("sm_90".into()); + reject(&wrong_sm_floor, init_declaration, "effects or target floor"); + + let mut wrong_llvm_route = init.clone(); + wrong_llvm_route + .backend_lowerings + .iter_mut() + .find(|lowering| lowering.backend == IntrinsicBackend::LlvmNvptx) + .unwrap() + .mechanism = BackendLoweringMechanism::InlinePtx; + reject( + &wrong_llvm_route, + init_declaration, + "reviewed mbarrier backend routes", + ); + + let mut wrong_lib_nvvm_route = init.clone(); + wrong_lib_nvvm_route + .backend_lowerings + .iter_mut() + .find(|lowering| lowering.backend == IntrinsicBackend::LibNvvm) + .unwrap() + .mechanism = BackendLoweringMechanism::TypedNvvm; + reject( + &wrong_lib_nvvm_route, + init_declaration, + "reviewed mbarrier backend routes", + ); + + let mut route_with_unreviewed_floor = init.clone(); + route_with_unreviewed_floor.backend_lowerings[0].minimum_sm = Some("sm_90".into()); + reject( + &route_with_unreviewed_floor, + init_declaration, + "reviewed mbarrier backend routes", + ); + + let mut mixed_family = init.clone(); + mixed_family.cp_async_control = Some(crate::model::CpAsyncControl { + operation: CpAsyncControlOperation::CommitGroup, + adapter: CpAsyncControlAdapter::NoOperands, + runtime_validation: RuntimeValidation::Unexecuted, + }); + reject( + &mixed_family, + init_declaration, + "mixes another generated-family contract", + ); +} + +#[test] +fn redux_contract_validates_effects_participation_and_operand_adapter() { + let valid = redux_policy(); + let imported = redux_declaration(); + validate_imported_policy(&valid, &imported).unwrap(); + + assert_eq!( + valid.redux.as_ref().unwrap().adapter, + ReduxAdapter::MaskValueToSourceMemberMask + ); + + let mut missing_contract = valid.clone(); + missing_contract.redux = None; + assert!( + validate_imported_policy(&missing_contract, &imported) + .unwrap_err() + .to_string() + .contains("closed redux contract") + ); + + let mut wrong_effect = valid.clone(); + wrong_effect.memory = "none".into(); + assert!( + validate_imported_policy(&wrong_effect, &imported) + .unwrap_err() + .to_string() + .contains("redux effects") + ); + + let mut missing_imported_effect = imported; + missing_imported_effect + .properties + .retain(|property| property != "IntrInaccessibleMemOnly"); + assert!( + validate_imported_policy(&valid, &missing_imported_effect) + .unwrap_err() + .to_string() + .contains("memory and convergence effects") + ); +} + +#[test] +fn every_redux_variant_matches_its_closed_recipe() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let (overlay, _) = + read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")).unwrap(); + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + assert_eq!(imported.schema, IMPORTED_SCHEMA); + let declarations: BTreeMap<_, _> = imported + .intrinsics + .iter() + .map(|record| (record.source_record.as_str(), record)) + .collect(); + let redux: Vec<_> = overlay + .intrinsics + .iter() + .filter(|record| record.family == "redux") + .collect(); + assert_eq!(redux.len(), 16); + + for policy in redux { + let declaration = declarations + .get(policy.source_record.as_deref().unwrap()) + .unwrap(); + validate_imported_policy(policy, declaration).unwrap(); + } + + let mut mismatched = packed_policy("redux_sync_min_u32"); + mismatched.redux.as_mut().unwrap().operation = ReduxOperation::Umax; + let declaration = declarations["int_nvvm_redux_sync_umin"]; + assert!( + validate_imported_policy(&mismatched, declaration) + .unwrap_err() + .to_string() + .contains("closed operation recipe") + ); +} + +#[test] +fn every_dot_product_variant_matches_its_closed_recipe() { + let variants = [ + ( + DotProductOperation::Dp4a, + DotProductSignedness::Signed, + "dp4a_s32", + "int_nvvm_idp4a_s_s", + "integer.dot_product.dp4a.s32", + ), + ( + DotProductOperation::Dp4a, + DotProductSignedness::Unsigned, + "dp4a_u32", + "int_nvvm_idp4a_u_u", + "integer.dot_product.dp4a.u32", + ), + ( + DotProductOperation::Dp2a, + DotProductSignedness::Signed, + "dp2a_s32", + "int_nvvm_idp2a_s_s", + "integer.dot_product.dp2a.lo.s32", + ), + ( + DotProductOperation::Dp2a, + DotProductSignedness::Unsigned, + "dp2a_u32", + "int_nvvm_idp2a_u_u", + "integer.dot_product.dp2a.lo.u32", + ), + ]; + + for (operation, signedness, id, source_record, operation_key) in variants { + let policy = dot_product_policy(operation, signedness); + let declaration = dot_product_declaration(operation, signedness); + assert_eq!(policy.id, id); + assert_eq!(policy.source_record.as_deref(), Some(source_record)); + assert_eq!(policy.operation_key, operation_key); + validate_imported_policy(&policy, &declaration).unwrap(); + } +} + +#[test] +fn pinned_dot_product_records_match_the_reviewed_overlay() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let (overlay, _) = + read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")).unwrap(); + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let declarations: BTreeMap<_, _> = imported + .intrinsics + .iter() + .map(|record| (record.source_record.as_str(), record)) + .collect(); + let dot_products: Vec<_> = overlay + .intrinsics + .iter() + .filter(|record| record.family == "dotprod") + .collect(); + assert_eq!(dot_products.len(), 4); + + for policy in dot_products { + let declaration = declarations[policy.source_record.as_deref().unwrap()]; + validate_imported_policy(policy, declaration).unwrap(); + let selected: Vec<_> = declaration + .selections + .iter() + .filter(|selection| selection_matches_policy(policy, selection).unwrap()) + .collect(); + assert_eq!(selected.len(), 1); + if policy.id.starts_with("dp2a") { + assert_eq!(selected[0].constraints.immediate_bindings[0].value, 0); + } + } +} + +#[test] +fn dp2a_selects_only_the_reviewed_low_half_binding() { + let policy = dot_product_policy(DotProductOperation::Dp2a, DotProductSignedness::Signed); + let declaration = + dot_product_declaration(DotProductOperation::Dp2a, DotProductSignedness::Signed); + let resolved = resolve_record( + &policy, + resolve_policy_source(&policy).unwrap(), + Some(&declaration), + &dot_product_evidence(&policy), + "test", + "LLVM version test", + "0123456789abcdef", + vec![], + 1, + ) + .unwrap(); + + assert_eq!(resolved.selections.len(), 1); + assert_eq!(resolved.selections[0].source_record, "DOT2_lo"); + assert_eq!( + resolved.selections[0].constraints.immediate_bindings, + [crate::model::ImportedImmediateBinding { + argument_index: 2, + value: 0, + }] + ); + assert_eq!( + resolved.dot_product.as_ref().unwrap().adapter, + DotProductAdapter::InsertLowHalfFalse + ); + + let mut wrong_binding = declaration; + wrong_binding.selections[1].constraints.immediate_bindings[0].value = -1; + let error = validate_imported_policy(&policy, &wrong_binding).unwrap_err(); + assert!(error.to_string().contains("does not agree")); +} + +#[test] +fn dot_product_recipe_rejects_unreviewed_api_and_adapter_changes() { + let valid = dot_product_policy(DotProductOperation::Dp2a, DotProductSignedness::Unsigned); + let declaration = + dot_product_declaration(DotProductOperation::Dp2a, DotProductSignedness::Unsigned); + + let mut wrong_adapter = valid.clone(); + wrong_adapter.dot_product.as_mut().unwrap().adapter = DotProductAdapter::DirectThreeOperands; + assert!( + validate_imported_policy(&wrong_adapter, &declaration) + .unwrap_err() + .to_string() + .contains("source adapter") + ); + + let mut must_use = valid.clone(); + must_use.must_use = true; + assert!( + validate_imported_policy(&must_use, &declaration) + .unwrap_err() + .to_string() + .contains("non-must-use") + ); + + let mut wrong_llvm_signature = valid; + wrong_llvm_signature.llvm_arguments = vec!["i32".into(); 3]; + assert!( + validate_imported_policy(&wrong_llvm_signature, &declaration) + .unwrap_err() + .to_string() + .contains("LLVM argument signature mismatch") + ); +} + +#[test] +fn dot_product_target_predicate_is_closed_to_ptx50_and_sm61() { + let policy = dot_product_policy(DotProductOperation::Dp4a, DotProductSignedness::Signed); + let selection = + &dot_product_declaration(DotProductOperation::Dp4a, DotProductSignedness::Signed) + .selections[0]; + validate_selected_target_predicates(&policy, selection).unwrap(); + + let mut wrong_ptx = policy.clone(); + wrong_ptx.minimum_ptx = "5.1".into(); + assert!( + validate_selected_target_predicates(&wrong_ptx, selection) + .unwrap_err() + .to_string() + .contains("minimum PTX") + ); + + let mut wrong_sm = policy; + wrong_sm.minimum_sm = Some("sm_60".into()); + assert!( + validate_selected_target_predicates(&wrong_sm, selection) + .unwrap_err() + .to_string() + .contains("minimum SM") + ); +} + +#[test] +fn return_range_properties_are_half_open_and_unique() { + let facts = + imported_result_facts(&["NoUndef".into(), "Range".into()]).unwrap(); + assert!(facts.no_undef); + let range = facts.range.unwrap(); + assert_eq!(range.lower, "1"); + assert_eq!(range.upper_exclusive, "1025"); + + let duplicate = + imported_result_facts(&["Range".into(), "Range".into()]).unwrap_err(); + assert!(duplicate.to_string().contains("duplicate return range")); +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/tests/mma.rs b/crates/cuda-intrinsics-gen/src/resolve/tests/mma.rs new file mode 100644 index 0000000000..11d336b15e --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/tests/mma.rs @@ -0,0 +1,2064 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, CatalogHardwareAlternative, CatalogHardwareTarget, + EvidenceArtifactKind, EvidenceStageKind, ImportedFile, IntrinsicBackend, OverlayIntrinsic, + OverlayShardFile, PreSm70MemberMaskRule, RegisterMmaAccumulator, RegisterMmaAdapter, + RegisterMmaCompatibilitySource, RegisterMmaElement, RegisterMmaKind, RegisterMmaOperation, + RegisterMmaOverflow, RegisterMmaShape, RuntimeValidation, SparseMmaAccumulator, + SparseMmaAdapter, SparseMmaElement, SparseMmaLlvmAdapter, SparseMmaMetadata, SparseMmaOverflow, + SparseMmaSelector, SparseMmaShape, WarpBarrierAdapter, WarpBarrierMaskEncoding, + WarpBarrierMemoryOrdering, WarpBarrierParticipation, +}; +use crate::ptx::OperandPattern; +use crate::util::read_json; +use std::collections::{BTreeMap, BTreeSet}; +use std::fs; +use std::path::Path; + +use super::fixtures::*; +use crate::resolve::driver::*; +use crate::resolve::families::*; +use crate::resolve::guards::*; +use crate::resolve::overlay::*; + +#[test] +fn stmatrix_admission_is_closed_and_uses_schema_35() { + let shard = |schema| { + toml::from_str::(&format!( + r#" +schema = {schema} +family = "stmatrix" + +[stmatrix] +llvm_evidence_profile = "llvm-test" +libnvvm_evidence_profile = "libnvvm-test" +runtime_validation = "unexecuted" + +[[stmatrix.variant]] +abi_id = "i0301" +multiplicity = "x2" +layout = "normal" + +[[stmatrix.variant]] +abi_id = "i0302" +multiplicity = "x2" +layout = "transposed" + +[[stmatrix.variant]] +abi_id = "i0303" +multiplicity = "x4" +layout = "normal" + +[[stmatrix.variant]] +abi_id = "i0304" +multiplicity = "x4" +layout = "transposed" +"# + )) + .unwrap() + }; + let path = Path::new("intrinsics/overlay/stmatrix.toml"); + + let old = shard(STMATRIX_SHARD_SCHEMA - 1); + assert!( + validate_overlay_shard_schema(&old, path) + .unwrap_err() + .to_string() + .contains("requires overlay shard schema 35") + ); + + let current = shard(STMATRIX_SHARD_SCHEMA); + validate_overlay_shard_schema(¤t, path).unwrap(); + let admission = current.stmatrix.unwrap(); + let records = expand_stmatrix_admission(&admission).unwrap(); + assert_eq!( + records + .iter() + .map(|record| (record.abi_id.as_str(), record.id.as_str())) + .collect::>(), + [ + ("i0301", "stmatrix_m8n8_x2_b16"), + ("i0302", "stmatrix_m8n8_x2_trans_b16"), + ("i0303", "stmatrix_m8n8_x4_b16"), + ("i0304", "stmatrix_m8n8_x4_trans_b16"), + ] + ); + + let mut reordered = admission.clone(); + reordered.variants.swap(0, 1); + assert!(expand_stmatrix_admission(&reordered).is_err()); + + let mut wrong_id = admission.clone(); + wrong_id.variants[0].abi_id = "i0302".into(); + assert!(expand_stmatrix_admission(&wrong_id).is_err()); + + let mut executed = admission; + executed.runtime_validation = RuntimeValidation::Executed; + assert!(expand_stmatrix_admission(&executed).is_err()); +} + +#[test] +fn pinned_stmatrix_records_match_llvm_and_reject_contract_drift() { + let records = expand_stmatrix_admission(&test_stmatrix_admission()).unwrap(); + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let declarations = imported + .intrinsics + .iter() + .map(|record| (record.source_record.as_str(), record)) + .collect::>(); + + for record in &records { + let declaration = declarations[record.source_record.as_deref().unwrap()]; + assert!(declaration.selections.is_empty()); + validate_imported_policy(record, declaration).unwrap(); + } + + let declaration = declarations["int_nvvm_stmatrix_sync_aligned_m8n8_x2_b16"]; + let mut changed = records[0].clone(); + changed.memory = "read_write".into(); + assert!(validate_imported_policy(&changed, declaration).is_err()); + + let mut changed = records[0].clone(); + changed.convergent = false; + assert!(validate_imported_policy(&changed, declaration).is_err()); + + let mut changed = records[0].clone(); + changed.backend_lowerings[0].mechanism = BackendLoweringMechanism::InlinePtx; + assert!(validate_imported_policy(&changed, declaration).is_err()); + + let mut changed = records[0].clone(); + changed.minimum_sm = Some("sm_80".into()); + assert!(validate_imported_policy(&changed, declaration).is_err()); +} + +#[test] +fn compact_f8f6f4_axes_require_the_exact_canonical_matrix() { + let records = expand_sparse_mma_f8f6f4_admission(&test_f8f6f4_admission()).unwrap(); + assert_eq!(records.len(), 25); + assert_eq!( + records[0].id, + "mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m1_e2m1_f32" + ); + assert_eq!(records[0].rust_arguments[0], "[f32; 4]"); + assert_eq!(records[0].dialect_results, ["f32"; 4]); + assert_eq!(records[0].llvm_results, ["f32"; 4]); + + let mut missing = test_f8f6f4_admission(); + missing.a_elements.pop(); + assert!(expand_sparse_mma_f8f6f4_admission(&missing).is_err()); + + let mut duplicate = test_f8f6f4_admission(); + duplicate.a_elements[4] = SparseMmaElement::E4m3; + assert!(expand_sparse_mma_f8f6f4_admission(&duplicate).is_err()); + + let mut extra = test_f8f6f4_admission(); + extra.b_elements.push(SparseMmaElement::S4); + assert!(expand_sparse_mma_f8f6f4_admission(&extra).is_err()); + + let mut unsorted = test_f8f6f4_admission(); + unsorted.b_elements.swap(0, 1); + assert!(expand_sparse_mma_f8f6f4_admission(&unsorted).is_err()); + + let mut wrong_count = test_f8f6f4_admission(); + wrong_count.product_count = 24; + assert!(expand_sparse_mma_f8f6f4_admission(&wrong_count).is_err()); +} + +#[test] +fn compact_sparse_f8f6f4_f16_admission_is_closed_and_ordered() { + let admission = test_sparse_mma_f8f6f4_f16_admission(); + let records = expand_sparse_mma_f8f6f4_f16_admission(&admission).unwrap(); + assert_eq!(records.len(), 25); + assert!(records.iter().all(|record| record.abi_id.is_empty())); + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let mut bound = overlay_file(records.clone()); + bind_pinned_abi_ids(&repo_root, &mut bound); + assert_eq!( + bound + .intrinsics + .iter() + .map(|record| record.abi_id.clone()) + .collect::>(), + (525..=549) + .map(|id| format!("i{id:04}")) + .collect::>() + ); + assert_eq!( + records[0].id, + "mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e2m1_f16" + ); + assert_eq!( + records[24].id, + "mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e5m2_e5m2_f16" + ); + let first = &records[0]; + assert_eq!( + first.source_record.as_deref(), + Some("int_nvvm_mma_sp_ordered_metadata_m16n8k64_row_col_kind_f8f6f4_f16_e2m1_e2m1_f16") + ); + assert_eq!( + first.llvm_symbol.as_deref(), + Some("llvm.nvvm.mma.sp.ordered.metadata.m16n8k64.row.col.kind.f8f6f4.f16.e2m1.e2m1.f16") + ); + assert_eq!( + first.rust_arguments, + ["[u32; 2]", "[u32; 4]", "[u32; 4]", "u32", "u32"] + ); + assert_eq!(first.rust_result, "[u32; 2]"); + assert_eq!(first.dialect_results, ["i32"; 2]); + assert_eq!(first.llvm_results, ["v2f16"; 2]); + assert_eq!( + first.llvm_arguments, + [ + "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", "v2f16", "v2f16", "i32", "i32" + ] + ); + assert_eq!(first.minimum_ptx, "8.7"); + assert_eq!(first.minimum_sm, None); + assert_eq!(first.targets, SPARSE_MMA_F8F6F4_TARGETS); + assert!(first.convergent && !first.pure); + assert!(first.backend_lowerings.iter().all(|route| { + route.mechanism == BackendLoweringMechanism::InlinePtx + && route.minimum_ptx.as_deref() == Some("8.7") + && route.minimum_sm.is_none() + })); + let mma = first.sparse_mma.as_ref().unwrap(); + assert_eq!(mma.accumulator, SparseMmaAccumulator::F16); + assert_eq!(mma.selector, SparseMmaSelector::ImmediateZero); + assert_eq!( + mma.adapter, + SparseMmaAdapter::C2U32A4U32B4U32MetadataU32SelectorU32ToD2U32 + ); + assert_eq!( + mma.llvm_adapter, + SparseMmaLlvmAdapter::A4I32B4I32C2V2F16MetadataI32SelectorI32ToD2V2F16 + ); + assert_eq!( + first.compatibility_rust_paths, + ["cuda_device::wmma::mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e2m1_f16"] + ); + assert_eq!( + first.expected_ptx.operands, + [ + OperandPattern::RegisterList { length: 2 }, + OperandPattern::RegisterList { length: 4 }, + OperandPattern::RegisterList { length: 4 }, + OperandPattern::RegisterList { length: 2 }, + OperandPattern::Register, + OperandPattern::Immediate, + ] + ); + + let mut missing = admission.clone(); + missing.a_elements.pop(); + assert!(expand_sparse_mma_f8f6f4_f16_admission(&missing).is_err()); + + let mut reordered = admission.clone(); + reordered.b_elements.swap(0, 1); + assert!(expand_sparse_mma_f8f6f4_f16_admission(&reordered).is_err()); + + let mut wrong_count = admission.clone(); + wrong_count.product_count = 24; + assert!(expand_sparse_mma_f8f6f4_f16_admission(&wrong_count).is_err()); + + let mut legacy_first_id = admission.clone(); + legacy_first_id._legacy_first_abi_id = Some("i9999".into()); + assert!( + expand_sparse_mma_f8f6f4_f16_admission(&legacy_first_id) + .unwrap() + .iter() + .all(|record| record.abi_id.is_empty()) + ); + + let mut missing_evidence = admission.clone(); + missing_evidence.llvm_evidence_profile.clear(); + assert!(expand_sparse_mma_f8f6f4_f16_admission(&missing_evidence).is_err()); + + let mut executed = admission; + executed.runtime_validation = RuntimeValidation::Executed; + assert!(expand_sparse_mma_f8f6f4_f16_admission(&executed).is_err()); +} + +#[test] +fn sparse_f8f6f4_f16_candidate_floor_uses_the_resolved_policy() { + let policy = expand_sparse_mma_f8f6f4_f16_admission(&test_sparse_mma_f8f6f4_f16_admission()) + .unwrap() + .remove(0); + let (_, requirement) = candidate_llvm_route(&policy).unwrap(); + + validate_candidate_target(&policy, &requirement, "sm_120a", "+ptx87").unwrap(); + validate_candidate_target(&policy, &requirement, "sm_120f", "+ptx88").unwrap(); + validate_candidate_target(&policy, &requirement, "sm_121a", "+ptx88").unwrap(); + validate_candidate_target(&policy, &requirement, "sm_121f", "+ptx88").unwrap(); + for target in ["sm_120f", "sm_121a", "sm_121f"] { + assert!( + validate_candidate_target(&policy, &requirement, target, "+ptx87").is_err(), + "{target} must require PTX 8.8" + ); + } +} + +#[test] +fn compact_dense_f8f6f4_admission_is_closed() { + for (accumulator, first, first_id, last_id, rust_arguments, rust_result, adapter) in [ + ( + RegisterMmaAccumulator::F32, + 454, + "mma_m16n8k32_f32_e2m1_e2m1", + "mma_m16n8k32_f32_e5m2_e5m2", + ["[f32; 4]", "[u32; 4]", "[u32; 2]"], + "[f32; 4]", + RegisterMmaAdapter::C4F32A4U32B2U32ToD4F32, + ), + ( + RegisterMmaAccumulator::F16, + 479, + "mma_m16n8k32_f16_e2m1_e2m1", + "mma_m16n8k32_f16_e5m2_e5m2", + ["[u32; 2]", "[u32; 4]", "[u32; 2]"], + "[u32; 2]", + RegisterMmaAdapter::C2U32A4U32B2U32ToD2U32, + ), + ] { + let admission = test_register_mma_f8f6f4_admission(accumulator); + let records = expand_register_mma_f8f6f4_admission(&admission, accumulator).unwrap(); + assert_eq!(records.len(), 25); + assert!(records.iter().all(|record| record.abi_id.is_empty())); + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let mut bound = overlay_file(records.clone()); + bind_pinned_abi_ids(&repo_root, &mut bound); + assert_eq!( + bound + .intrinsics + .iter() + .map(|record| record.abi_id.clone()) + .collect::>(), + (first..first + 25) + .map(|id| format!("i{id:04}")) + .collect::>() + ); + assert_eq!(records[0].id, first_id); + assert_eq!(records[24].id, last_id); + assert!(records.iter().all(|record| { + record.rust_arguments == rust_arguments + && record.rust_result == rust_result + && record.register_mma.as_ref().unwrap().adapter == adapter + && record.targets == REGISTER_MMA_F8F6F4_TARGETS + && record.minimum_ptx == "8.7" + && record.minimum_sm.is_none() + && record.backend_lowerings.len() == 2 + && record + .backend_lowerings + .iter() + .all(|route| route.mechanism == BackendLoweringMechanism::InlinePtx) + })); + + let mut missing = admission.clone(); + missing.a_elements.pop(); + assert!(expand_register_mma_f8f6f4_admission(&missing, accumulator).is_err()); + + let mut reordered = admission.clone(); + reordered.b_elements.swap(0, 1); + assert!(expand_register_mma_f8f6f4_admission(&reordered, accumulator).is_err()); + + let mut wrong_count = admission.clone(); + wrong_count.product_count = 24; + assert!(expand_register_mma_f8f6f4_admission(&wrong_count, accumulator).is_err()); + + let mut legacy_first_id = admission.clone(); + legacy_first_id._legacy_first_abi_id = Some("i9999".into()); + assert!( + expand_register_mma_f8f6f4_admission(&legacy_first_id, accumulator) + .unwrap() + .iter() + .all(|record| record.abi_id.is_empty()) + ); + + let mut reordered_targets = admission.clone(); + reordered_targets.targets.swap(0, 1); + assert!(expand_register_mma_f8f6f4_admission(&reordered_targets, accumulator).is_err()); + + let mut missing_evidence = admission.clone(); + missing_evidence.llvm_evidence_profile.clear(); + assert!(expand_register_mma_f8f6f4_admission(&missing_evidence, accumulator).is_err()); + + let mut executed = admission; + executed.runtime_validation = RuntimeValidation::Executed; + assert!(expand_register_mma_f8f6f4_admission(&executed, accumulator).is_err()); + } + + let admission = test_register_mma_f8f6f4_admission(RegisterMmaAccumulator::F32); + assert!(expand_register_mma_f8f6f4_admission(&admission, RegisterMmaAccumulator::F64).is_err()); +} + +#[test] +fn compact_standard_fp8_admission_is_closed_and_ordered() { + let admission = test_register_mma_fp8_admission(); + let records = expand_register_mma_fp8_admission(&admission).unwrap(); + assert_eq!(records.len(), 16); + assert!(records.iter().all(|record| record.abi_id.is_empty())); + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let mut bound = overlay_file(records.clone()); + bind_pinned_abi_ids(&repo_root, &mut bound); + assert_eq!( + bound + .intrinsics + .iter() + .map(|record| record.abi_id.clone()) + .collect::>(), + (504..=519) + .map(|id| format!("i{id:04}")) + .collect::>() + ); + + let mut index = 0; + for &shape in ®ISTER_MMA_FP8_SHAPES { + let (shape_name, a_count, b_count) = register_mma_fp8_shape_contract(shape).unwrap(); + for &accumulator in ®ISTER_MMA_FP8_ACCUMULATORS { + let (scalar, arguments, result, adapter, result_count) = match accumulator { + RegisterMmaAccumulator::F16 => ( + "f16", + vec![ + "[u32; 2]".to_owned(), + format!("[u32; {a_count}]"), + if b_count == 1 { + "u32".to_owned() + } else { + "[u32; 2]".to_owned() + }, + ], + "[u32; 2]", + if shape == RegisterMmaShape::M16n8k16 { + RegisterMmaAdapter::C2U32A2U32B1U32ToD2U32 + } else { + RegisterMmaAdapter::C2U32A4U32B2U32ToD2U32 + }, + 2, + ), + RegisterMmaAccumulator::F32 => ( + "f32", + vec![ + "[f32; 4]".to_owned(), + format!("[u32; {a_count}]"), + if b_count == 1 { + "u32".to_owned() + } else { + "[u32; 2]".to_owned() + }, + ], + "[f32; 4]", + if shape == RegisterMmaShape::M16n8k16 { + RegisterMmaAdapter::C4F32A2U32B1U32ToD4F32 + } else { + RegisterMmaAdapter::C4F32A4U32B2U32ToD4F32 + }, + 4, + ), + _ => unreachable!(), + }; + for &a_element in ®ISTER_MMA_FP8_ELEMENTS { + for &b_element in ®ISTER_MMA_FP8_ELEMENTS { + let record = &records[index]; + let a = register_mma_fp8_element_name(a_element).unwrap(); + let b = register_mma_fp8_element_name(b_element).unwrap(); + assert_eq!(record.id, format!("mma_{shape_name}_fp8_{scalar}_{a}_{b}")); + assert_eq!( + record.operation_key, + format!( + "matrix.mma.{shape_name}.row.col.standard_fp8.{scalar}.{a}.{b}.{scalar}" + ) + ); + assert_eq!( + record.source_record.as_deref(), + Some( + format!("int_nvvm_mma_{shape_name}_row_col_{scalar}_{a}_{b}_{scalar}") + .as_str() + ) + ); + assert_eq!(record.rust_arguments, arguments); + assert_eq!(record.rust_result, result); + assert_eq!( + record.minimum_ptx, + register_mma_fp8_minimum_ptx(shape, accumulator) + ); + assert_eq!(record.minimum_sm.as_deref(), Some("sm_89")); + assert_eq!( + record.expected_ptx.operands, + [result_count, a_count, b_count, result_count] + .map(|length| OperandPattern::RegisterList { length }) + ); + let mma = record.register_mma.as_ref().unwrap(); + assert_eq!(mma.kind, Some(RegisterMmaKind::Standard)); + assert_eq!(mma.adapter, adapter); + index += 1; + } + } + } + } + assert_eq!(index, records.len()); + + let mut reordered = admission.clone(); + reordered.shapes.swap(0, 1); + assert!(expand_register_mma_fp8_admission(&reordered).is_err()); + let mut reordered = admission.clone(); + reordered.accumulators.swap(0, 1); + assert!(expand_register_mma_fp8_admission(&reordered).is_err()); + let mut reordered = admission.clone(); + reordered.a_elements.swap(0, 1); + assert!(expand_register_mma_fp8_admission(&reordered).is_err()); + let mut reordered = admission.clone(); + reordered.b_elements.swap(0, 1); + assert!(expand_register_mma_fp8_admission(&reordered).is_err()); + let mut wrong = admission.clone(); + wrong.product_count = 15; + assert!(expand_register_mma_fp8_admission(&wrong).is_err()); + let mut legacy_first_id = admission.clone(); + legacy_first_id._legacy_first_abi_id = Some("i9999".into()); + assert!( + expand_register_mma_fp8_admission(&legacy_first_id) + .unwrap() + .iter() + .all(|record| record.abi_id.is_empty()) + ); + let mut wrong = admission; + wrong.runtime_validation = RuntimeValidation::Executed; + assert!(expand_register_mma_fp8_admission(&wrong).is_err()); +} + +#[test] +fn compact_ampere_float_mma_admission_is_closed_and_ordered() { + let admission = test_register_mma_ampere_float_admission(); + let records = expand_register_mma_ampere_float_admission(&admission).unwrap(); + assert_eq!(records.len(), 5); + assert!(records.iter().all(|record| record.abi_id.is_empty())); + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let mut bound = overlay_file(records.clone()); + bind_pinned_abi_ids(&repo_root, &mut bound); + assert_eq!( + bound + .intrinsics + .iter() + .map(|record| (record.abi_id.as_str(), record.id.as_str())) + .collect::>(), + [ + ("i0520", "mma_m16n8k4_f32_tf32"), + ("i0521", "mma_m16n8k8_f16_f16"), + ("i0522", "mma_m16n8k8_f32_bf16"), + ("i0523", "mma_m16n8k8_f32_f16"), + ("i0524", "mma_m16n8k16_f16_f16"), + ] + ); + for record in &records { + assert_eq!(record.minimum_ptx, "7.0"); + assert_eq!(record.minimum_sm.as_deref(), Some("sm_80")); + assert_eq!(record.backend_lowerings.len(), 2); + assert!(record.backend_lowerings.iter().all(|lowering| { + lowering.mechanism == BackendLoweringMechanism::InlinePtx + && lowering.minimum_ptx.as_deref() == Some("7.0") + && lowering.minimum_sm.as_deref() == Some("sm_80") + })); + let mma = record.register_mma.as_ref().unwrap(); + assert_eq!(mma.kind, None); + assert_eq!( + mma.compatibility_source, + RegisterMmaCompatibilitySource::GeneratedStub + ); + } + + let mut reordered = admission.clone(); + reordered.variants.swap(0, 1); + assert!(expand_register_mma_ampere_float_admission(&reordered).is_err()); + let mut legacy_first_id = admission.clone(); + legacy_first_id._legacy_first_abi_id = Some("i9999".into()); + assert!( + expand_register_mma_ampere_float_admission(&legacy_first_id) + .unwrap() + .iter() + .all(|record| record.abi_id.is_empty()) + ); + let mut wrong = admission.clone(); + wrong.product_count = 4; + assert!(expand_register_mma_ampere_float_admission(&wrong).is_err()); + let mut wrong = admission.clone(); + wrong.llvm_evidence_profile.clear(); + assert!(expand_register_mma_ampere_float_admission(&wrong).is_err()); + let mut wrong = admission; + wrong.runtime_validation = RuntimeValidation::Executed; + assert!(expand_register_mma_ampere_float_admission(&wrong).is_err()); + + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let path = repo_root.join("intrinsics/overlay/register_mma_ampere_float.toml"); + let bytes = fs::read(&path).unwrap(); + let mut shard: OverlayShardFile = toml::from_slice(&bytes).unwrap(); + validate_overlay_shard_schema(&shard, &path).unwrap(); + shard.schema = REGISTER_MMA_AMPERE_FLOAT_SHARD_SCHEMA - 1; + assert!( + validate_overlay_shard_schema(&shard, &path) + .unwrap_err() + .to_string() + .contains("requires overlay shard schema 49") + ); +} + +#[test] +fn ampere_float_mma_policies_match_llvm_and_fail_closed() { + let records = + expand_register_mma_ampere_float_admission(&test_register_mma_ampere_float_admission()) + .unwrap(); + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let mut overlay = overlay_file(records); + bind_pinned_abi_ids(&repo_root, &mut overlay); + let records = overlay.intrinsics; + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let declarations = imported + .intrinsics + .iter() + .map(|record| (record.source_record.as_str(), record)) + .collect::>(); + + for record in &records { + let declaration = declarations[record.source_record.as_deref().unwrap()]; + validate_imported_policy(record, declaration).unwrap(); + let expected_predicates = match record.id.as_str() { + "mma_m16n8k4_f32_tf32" | "mma_m16n8k16_f16_f16" => [ + "Subtarget->getSmVersion() >= 80", + "Subtarget->getPTXVersion() >= 70", + ], + _ => [ + "Subtarget->getPTXVersion() >= 65", + "Subtarget->getSmVersion() >= 75", + ], + }; + assert_eq!(declaration.selections[0].predicates, expected_predicates); + assert_eq!(record.minimum_ptx, "7.0"); + assert_eq!(record.minimum_sm.as_deref(), Some("sm_80")); + } + + let valid = &records[2]; + assert_eq!(valid.id, "mma_m16n8k8_f32_bf16"); + assert_eq!( + declarations[valid.source_record.as_deref().unwrap()].selections[0].predicates, + [ + "Subtarget->getPTXVersion() >= 65", + "Subtarget->getSmVersion() >= 75", + ] + ); + assert_eq!(valid.minimum_ptx, "7.0"); + assert_eq!(valid.minimum_sm.as_deref(), Some("sm_80")); + assert!(valid.backend_lowerings.iter().all(|lowering| { + lowering.minimum_ptx.as_deref() == Some("7.0") + && lowering.minimum_sm.as_deref() == Some("sm_80") + })); + let declaration = declarations[valid.source_record.as_deref().unwrap()]; + let mut wrong = valid.clone(); + wrong.minimum_ptx = "6.5".into(); + assert!(validate_imported_policy(&wrong, declaration).is_err()); + let mut wrong = valid.clone(); + wrong.register_mma.as_mut().unwrap().adapter = RegisterMmaAdapter::C4F32A4U32B2U32ToD4F32; + assert!(validate_imported_policy(&wrong, declaration).is_err()); + let mut wrong = valid.clone(); + wrong.register_mma.as_mut().unwrap().kind = Some(RegisterMmaKind::Standard); + assert!(validate_imported_policy(&wrong, declaration).is_err()); + let mut wrong_declaration = declaration.clone(); + wrong_declaration.selections[0].predicates[0] = "Subtarget->getPTXVersion() >= 70".into(); + assert!(validate_imported_policy(valid, &wrong_declaration).is_err()); +} + +#[test] +fn standard_fp8_policies_match_llvm_and_fail_closed() { + let records = expand_register_mma_fp8_admission(&test_register_mma_fp8_admission()).unwrap(); + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let mut overlay = overlay_file(records); + bind_pinned_abi_ids(&repo_root, &mut overlay); + let records = overlay.intrinsics; + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let declarations = imported + .intrinsics + .iter() + .map(|record| (record.source_record.as_str(), record)) + .collect::>(); + + for record in &records { + let declaration = declarations[record.source_record.as_deref().unwrap()]; + validate_imported_policy(record, declaration).unwrap(); + let mma = record.register_mma.as_ref().unwrap(); + let source_floor = declaration.selections[0] + .predicates + .iter() + .find(|predicate| predicate.contains("getPTXVersion")) + .unwrap(); + if mma.shape == RegisterMmaShape::M16n8k32 { + assert!(source_floor.ends_with(">= 84")); + if mma.accumulator == RegisterMmaAccumulator::F16 { + assert_eq!(record.minimum_ptx, "8.7"); + } else { + assert_eq!(record.minimum_ptx, "8.4"); + } + } else { + assert!(source_floor.ends_with(">= 87")); + assert_eq!(record.minimum_ptx, "8.7"); + } + } + + let valid = &records[0]; + let declaration = declarations[valid.source_record.as_deref().unwrap()]; + let mut wrong = valid.clone(); + wrong.register_mma.as_mut().unwrap().kind = None; + assert!(validate_imported_policy(&wrong, declaration).is_err()); + let mut wrong = valid.clone(); + wrong.register_mma.as_mut().unwrap().adapter = RegisterMmaAdapter::C2U32A4U32B2U32ToD2U32; + assert!(validate_imported_policy(&wrong, declaration).is_err()); + let mut wrong_declaration = declaration.clone(); + wrong_declaration.selections[0].predicates[1] = "Subtarget->getPTXVersion() >= 84".into(); + assert!(validate_imported_policy(valid, &wrong_declaration).is_err()); + let mut wrong = valid.clone(); + wrong.minimum_ptx = "8.4".into(); + assert!(validate_imported_policy(&wrong, declaration).is_err()); + + let mut old = expand_register_mma_f8f6f4_admission( + &test_register_mma_f8f6f4_admission(RegisterMmaAccumulator::F32), + RegisterMmaAccumulator::F32, + ) + .unwrap() + .remove(18); + let mut standard = records[12].clone(); + assert_ne!(old.id, standard.id); + standard.id = old.id.clone(); + assert!(validate_unique_overlay(&[old.clone(), standard], 1).is_err()); + old.register_mma.as_mut().unwrap().kind = Some(RegisterMmaKind::F8f6f4); + assert_eq!( + old.register_mma.as_ref().unwrap().kind, + Some(RegisterMmaKind::F8f6f4) + ); +} + +#[test] +fn standard_fp8_resolves_exact_routes_and_target_floors() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = resolve(&repo_root).unwrap(); + let records = catalog + .intrinsics + .iter() + .filter(|record| { + record.rust.abi_id.as_str() >= "i0504" && record.rust.abi_id.as_str() <= "i0519" + }) + .collect::>(); + let expected = [ + ( + "mma_m16n8k16_fp8_f16_e4m3_e4m3", + "i0504", + RegisterMmaShape::M16n8k16, + RegisterMmaAccumulator::F16, + 87, + ), + ( + "mma_m16n8k16_fp8_f16_e4m3_e5m2", + "i0505", + RegisterMmaShape::M16n8k16, + RegisterMmaAccumulator::F16, + 87, + ), + ( + "mma_m16n8k16_fp8_f16_e5m2_e4m3", + "i0506", + RegisterMmaShape::M16n8k16, + RegisterMmaAccumulator::F16, + 87, + ), + ( + "mma_m16n8k16_fp8_f16_e5m2_e5m2", + "i0507", + RegisterMmaShape::M16n8k16, + RegisterMmaAccumulator::F16, + 87, + ), + ( + "mma_m16n8k16_fp8_f32_e4m3_e4m3", + "i0508", + RegisterMmaShape::M16n8k16, + RegisterMmaAccumulator::F32, + 87, + ), + ( + "mma_m16n8k16_fp8_f32_e4m3_e5m2", + "i0509", + RegisterMmaShape::M16n8k16, + RegisterMmaAccumulator::F32, + 87, + ), + ( + "mma_m16n8k16_fp8_f32_e5m2_e4m3", + "i0510", + RegisterMmaShape::M16n8k16, + RegisterMmaAccumulator::F32, + 87, + ), + ( + "mma_m16n8k16_fp8_f32_e5m2_e5m2", + "i0511", + RegisterMmaShape::M16n8k16, + RegisterMmaAccumulator::F32, + 87, + ), + ( + "mma_m16n8k32_fp8_f16_e4m3_e4m3", + "i0512", + RegisterMmaShape::M16n8k32, + RegisterMmaAccumulator::F16, + 87, + ), + ( + "mma_m16n8k32_fp8_f16_e4m3_e5m2", + "i0513", + RegisterMmaShape::M16n8k32, + RegisterMmaAccumulator::F16, + 87, + ), + ( + "mma_m16n8k32_fp8_f16_e5m2_e4m3", + "i0514", + RegisterMmaShape::M16n8k32, + RegisterMmaAccumulator::F16, + 87, + ), + ( + "mma_m16n8k32_fp8_f16_e5m2_e5m2", + "i0515", + RegisterMmaShape::M16n8k32, + RegisterMmaAccumulator::F16, + 87, + ), + ( + "mma_m16n8k32_fp8_f32_e4m3_e4m3", + "i0516", + RegisterMmaShape::M16n8k32, + RegisterMmaAccumulator::F32, + 84, + ), + ( + "mma_m16n8k32_fp8_f32_e4m3_e5m2", + "i0517", + RegisterMmaShape::M16n8k32, + RegisterMmaAccumulator::F32, + 84, + ), + ( + "mma_m16n8k32_fp8_f32_e5m2_e4m3", + "i0518", + RegisterMmaShape::M16n8k32, + RegisterMmaAccumulator::F32, + 84, + ), + ( + "mma_m16n8k32_fp8_f32_e5m2_e5m2", + "i0519", + RegisterMmaShape::M16n8k32, + RegisterMmaAccumulator::F32, + 84, + ), + ]; + assert_eq!(records.len(), expected.len()); + let sm_89 = CatalogHardwareTarget::AnyOf { + alternatives: vec![CatalogHardwareAlternative::MinimumSm { sm: 89 }], + }; + let expected_routes = [ + ( + IntrinsicBackend::LlvmNvptx, + BackendLoweringMechanism::InlinePtx, + "rust-llvm-22.1.2-1cb4e383", + ), + ( + IntrinsicBackend::LibNvvm, + BackendLoweringMechanism::InlinePtx, + "cuda-13.3-libnvvm-13.3.33", + ), + ]; + let mut floor_groups = BTreeMap::new(); + + for (record, &(id, abi_id, shape, accumulator, minimum_ptx)) in records.iter().zip(&expected) { + assert_eq!(record.id, id); + assert_eq!(record.rust.abi_id, abi_id); + let mma = record.register_mma.as_ref().unwrap(); + assert_eq!((mma.shape, mma.accumulator), (shape, accumulator)); + assert_eq!(record.target.minimum_ptx.encoded(), minimum_ptx); + assert_eq!(record.target.hardware, sm_89); + assert_eq!( + record + .backend_lowerings + .iter() + .map(|route| ( + route.backend, + route.mechanism, + route.evidence_profile.as_str() + )) + .collect::>(), + expected_routes + ); + for route in &record.backend_lowerings { + assert_eq!(route.target.minimum_ptx.encoded(), minimum_ptx); + assert_eq!(route.target.hardware, sm_89); + } + *floor_groups + .entry((shape, accumulator, minimum_ptx)) + .or_insert(0) += 1; + + if shape == RegisterMmaShape::M16n8k32 && accumulator == RegisterMmaAccumulator::F16 { + assert_eq!( + record + .selections + .iter() + .flat_map(|selection| &selection.predicates) + .filter(|predicate| predicate.contains("getPTXVersion")) + .map(String::as_str) + .collect::>(), + ["Subtarget->getPTXVersion() >= 84"] + ); + let llvm_route = &record.backend_lowerings[0]; + assert!(llvm_route.stages.iter().any(|stage| { + stage.stage == EvidenceStageKind::PtxAssembly + && stage.mechanism == Some(BackendLoweringMechanism::InlinePtx) + && stage.outcome == "succeeded" + && stage.targets == ["sm_89", "ptx87"] + && stage.artifact_kind == Some(EvidenceArtifactKind::Cubin) + })); + assert!(llvm_route.stages.iter().any(|stage| { + stage.stage == EvidenceStageKind::PtxAssembly + && stage.mechanism == Some(BackendLoweringMechanism::InlinePtx) + && stage.outcome == "failed" + && stage.targets == ["sm_89", "ptx86"] + })); + } + } + assert_eq!( + floor_groups, + BTreeMap::from([ + ( + (RegisterMmaShape::M16n8k16, RegisterMmaAccumulator::F16, 87), + 4 + ), + ( + (RegisterMmaShape::M16n8k16, RegisterMmaAccumulator::F32, 87), + 4 + ), + ( + (RegisterMmaShape::M16n8k32, RegisterMmaAccumulator::F16, 87), + 4 + ), + ( + (RegisterMmaShape::M16n8k32, RegisterMmaAccumulator::F32, 84), + 4 + ), + ]) + ); +} + +#[test] +fn dense_f8f6f4_candidate_floor_uses_the_resolved_policy() { + for accumulator in [RegisterMmaAccumulator::F16, RegisterMmaAccumulator::F32] { + let policy = expand_register_mma_f8f6f4_admission( + &test_register_mma_f8f6f4_admission(accumulator), + accumulator, + ) + .unwrap() + .remove(0); + let (_, requirement) = candidate_llvm_route(&policy).unwrap(); + + validate_candidate_target(&policy, &requirement, "sm_120a", "+ptx87").unwrap(); + validate_candidate_target(&policy, &requirement, "sm_120f", "+ptx88").unwrap(); + validate_candidate_target(&policy, &requirement, "sm_121a", "+ptx88").unwrap(); + validate_candidate_target(&policy, &requirement, "sm_121f", "+ptx88").unwrap(); + for target in ["sm_120f", "sm_121a", "sm_121f"] { + assert!( + validate_candidate_target(&policy, &requirement, target, "+ptx87").is_err(), + "{target} must require PTX 8.8" + ); + } + } +} + +#[test] +fn pinned_register_mma_records_match_the_closed_recipes_and_fail_closed() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let (mut overlay, _) = + read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")).unwrap(); + bind_pinned_abi_ids(&repo_root, &mut overlay); + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let declarations: BTreeMap<_, _> = imported + .intrinsics + .iter() + .map(|record| (record.source_record.as_str(), record)) + .collect(); + let records: Vec<_> = overlay + .intrinsics + .iter() + .filter(|record| record.family == "register_mma") + .collect(); + assert_eq!(records.len(), 154); + + let dense_f8f6f4_records = records + .iter() + .copied() + .filter(|record| { + record.register_mma.as_ref().is_some_and(|mma| { + matches!(mma.kind, None | Some(RegisterMmaKind::F8f6f4)) + && register_mma_f8f6f4_element_name(mma.a_element).is_some() + && register_mma_f8f6f4_element_name(mma.b_element).is_some() + }) + }) + .collect::>(); + assert_eq!(dense_f8f6f4_records.len(), 50); + for (accumulator, first_abi, last_abi, adapter) in [ + ( + RegisterMmaAccumulator::F32, + "i0454", + "i0478", + RegisterMmaAdapter::C4F32A4U32B2U32ToD4F32, + ), + ( + RegisterMmaAccumulator::F16, + "i0479", + "i0503", + RegisterMmaAdapter::C2U32A4U32B2U32ToD2U32, + ), + ] { + let family = dense_f8f6f4_records + .iter() + .copied() + .filter(|record| record.register_mma.as_ref().unwrap().accumulator == accumulator) + .collect::>(); + assert_eq!(family.len(), 25); + assert_eq!(family[0].abi_id, first_abi); + assert_eq!(family[24].abi_id, last_abi); + assert!( + family + .iter() + .all(|record| record.register_mma.as_ref().unwrap().adapter == adapter) + ); + } + assert!(dense_f8f6f4_records.iter().all(|record| { + record.minimum_ptx == "8.7" + && record.minimum_sm.is_none() + && record.targets == REGISTER_MMA_F8F6F4_TARGETS + })); + + let mxf8f6f4_records = records + .iter() + .copied() + .filter(|record| { + record + .register_mma + .as_ref() + .is_some_and(|mma| mma.kind == Some(RegisterMmaKind::Mxf8f6f4)) + }) + .collect::>(); + assert_eq!(mxf8f6f4_records.len(), 25); + assert_eq!(mxf8f6f4_records[0].abi_id, "i0858"); + assert_eq!(mxf8f6f4_records[24].abi_id, "i0882"); + assert!(mxf8f6f4_records.iter().all(|record| { + let mma = record.register_mma.as_ref().unwrap(); + mma.shape == RegisterMmaShape::M16n8k32 + && mma.accumulator == RegisterMmaAccumulator::F32 + && mma.adapter == RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 + && record.rust_arguments + == [ + "[f32; 4]", "[u32; 4]", "[u32; 2]", "u32", "u16", "u16", "u32", "u16", "u16", + ] + && record.rust_result == "[f32; 4]" + && record.minimum_ptx == "8.7" + && record.minimum_sm.is_none() + && record.targets == REGISTER_MMA_F8F6F4_TARGETS + && record + .expected_ptx + .modifiers + .iter() + .any(|modifier| modifier == "kind::mxf8f6f4") + })); + + let integer_records: Vec<_> = records + .iter() + .copied() + .filter(|record| { + record.register_mma.as_ref().is_some_and(|mma| { + mma.operation == RegisterMmaOperation::Multiply + && mma.accumulator == RegisterMmaAccumulator::S32 + }) + }) + .collect(); + assert_eq!(integer_records.len(), 48); + let binary_records = records + .iter() + .copied() + .filter(|record| { + record + .register_mma + .as_ref() + .is_some_and(|mma| mma.operation != RegisterMmaOperation::Multiply) + }) + .collect::>(); + assert_eq!(binary_records.len(), 6); + let int8_records = integer_records + .iter() + .copied() + .filter(|record| { + let mma = record.register_mma.as_ref().unwrap(); + matches!( + mma.a_element, + RegisterMmaElement::S8 | RegisterMmaElement::U8 + ) && matches!( + mma.b_element, + RegisterMmaElement::S8 | RegisterMmaElement::U8 + ) + }) + .collect::>(); + assert_eq!(int8_records.len(), 24); + let int4_records = integer_records + .iter() + .copied() + .filter(|record| { + let mma = record.register_mma.as_ref().unwrap(); + matches!( + mma.a_element, + RegisterMmaElement::S4 | RegisterMmaElement::U4 + ) && matches!( + mma.b_element, + RegisterMmaElement::S4 | RegisterMmaElement::U4 + ) + }) + .collect::>(); + assert_eq!(int4_records.len(), 24); + let actual_variants = integer_records + .iter() + .map(|record| { + let mma = record.register_mma.as_ref().unwrap(); + (mma.shape, mma.a_element, mma.b_element, mma.overflow) + }) + .collect::>(); + let expected_int8_variants = [ + RegisterMmaShape::M8n8k16, + RegisterMmaShape::M16n8k16, + RegisterMmaShape::M16n8k32, + ] + .into_iter() + .flat_map(|shape| { + [RegisterMmaElement::S8, RegisterMmaElement::U8] + .into_iter() + .flat_map(move |a_element| { + [RegisterMmaElement::S8, RegisterMmaElement::U8] + .into_iter() + .flat_map(move |b_element| { + [ + RegisterMmaOverflow::Wrapping, + RegisterMmaOverflow::Satfinite, + ] + .into_iter() + .map(move |overflow| (shape, a_element, b_element, overflow)) + }) + }) + }) + .collect::>(); + let expected_int4_variants = [ + RegisterMmaShape::M8n8k32, + RegisterMmaShape::M16n8k32, + RegisterMmaShape::M16n8k64, + ] + .into_iter() + .flat_map(|shape| { + [RegisterMmaElement::S4, RegisterMmaElement::U4] + .into_iter() + .flat_map(move |a_element| { + [RegisterMmaElement::S4, RegisterMmaElement::U4] + .into_iter() + .flat_map(move |b_element| { + [ + RegisterMmaOverflow::Wrapping, + RegisterMmaOverflow::Satfinite, + ] + .into_iter() + .map(move |overflow| (shape, a_element, b_element, overflow)) + }) + }) + }) + .collect::>(); + let expected_variants = expected_int8_variants + .union(&expected_int4_variants) + .copied() + .collect::>(); + assert_eq!(actual_variants, expected_variants); + assert_eq!( + integer_records + .iter() + .filter(|record| { + record.register_mma.as_ref().unwrap().compatibility_source + == RegisterMmaCompatibilitySource::GeneratedStub + }) + .count(), + 47 + ); + + let actual_binary_variants = binary_records + .iter() + .map(|record| { + let mma = record.register_mma.as_ref().unwrap(); + (mma.shape, mma.operation) + }) + .collect::>(); + let expected_binary_variants = [ + RegisterMmaShape::M8n8k128, + RegisterMmaShape::M16n8k128, + RegisterMmaShape::M16n8k256, + ] + .into_iter() + .flat_map(|shape| { + [RegisterMmaOperation::XorPopc, RegisterMmaOperation::AndPopc] + .into_iter() + .map(move |operation| (shape, operation)) + }) + .collect::>(); + assert_eq!(actual_binary_variants, expected_binary_variants); + assert!(binary_records.iter().all(|record| { + let mma = record.register_mma.as_ref().unwrap(); + mma.accumulator == RegisterMmaAccumulator::S32 + && mma.a_element == RegisterMmaElement::B1 + && mma.b_element == RegisterMmaElement::B1 + && mma.overflow == RegisterMmaOverflow::Wrapping + && mma.compatibility_source == RegisterMmaCompatibilitySource::GeneratedStub + && record.expected_ptx.modifiers.ends_with(&[ + match mma.operation { + RegisterMmaOperation::XorPopc => "xor".into(), + RegisterMmaOperation::AndPopc => "and".into(), + RegisterMmaOperation::Multiply => unreachable!(), + }, + "popc".into(), + ]) + })); + + for record in &binary_records { + let mma = record.register_mma.as_ref().unwrap(); + let (arguments, result, adapter) = match mma.shape { + RegisterMmaShape::M8n8k128 => ( + &["[i32; 2]", "u32", "u32"] as &[_], + "[i32; 2]", + RegisterMmaAdapter::C2I32A1U32B1U32ToD2I32, + ), + RegisterMmaShape::M16n8k128 => ( + &["[i32; 4]", "[u32; 2]", "u32"] as &[_], + "[i32; 4]", + RegisterMmaAdapter::C4I32A2U32B1U32ToD4I32, + ), + RegisterMmaShape::M16n8k256 => ( + &["[i32; 4]", "[u32; 4]", "[u32; 2]"] as &[_], + "[i32; 4]", + RegisterMmaAdapter::C4I32A4U32B2U32ToD4I32, + ), + _ => unreachable!(), + }; + assert_eq!(record.rust_arguments, arguments); + assert_eq!(record.rust_result, result); + assert_eq!(mma.adapter, adapter); + let expected_floor = match (mma.shape, mma.operation) { + (RegisterMmaShape::M8n8k128, RegisterMmaOperation::XorPopc) => ("7.0", "sm_75"), + (_, RegisterMmaOperation::XorPopc) => ("7.0", "sm_80"), + (_, RegisterMmaOperation::AndPopc) => ("7.1", "sm_80"), + _ => unreachable!(), + }; + assert_eq!(record.minimum_ptx, expected_floor.0); + assert_eq!(record.minimum_sm.as_deref(), Some(expected_floor.1)); + } + + for record in integer_records.iter().filter(|record| { + matches!( + record.register_mma.as_ref().unwrap().shape, + RegisterMmaShape::M8n8k16 | RegisterMmaShape::M8n8k32 + ) + }) { + assert_eq!(record.rust_arguments, ["[i32; 2]", "u32", "u32"]); + assert_eq!(record.rust_result, "[i32; 2]"); + assert_eq!(record.minimum_ptx, "6.5"); + assert_eq!(record.minimum_sm.as_deref(), Some("sm_75")); + assert_eq!( + record.register_mma.as_ref().unwrap().adapter, + RegisterMmaAdapter::C2I32A1U32B1U32ToD2I32 + ); + } + + for record in int4_records + .iter() + .filter(|record| record.register_mma.as_ref().unwrap().shape == RegisterMmaShape::M16n8k32) + { + assert_eq!(record.rust_arguments, ["[i32; 4]", "[u32; 2]", "u32"]); + assert_eq!(record.rust_result, "[i32; 4]"); + assert_eq!(record.minimum_ptx, "7.0"); + assert_eq!(record.minimum_sm.as_deref(), Some("sm_80")); + assert_eq!( + record.register_mma.as_ref().unwrap().adapter, + RegisterMmaAdapter::C4I32A2U32B1U32ToD4I32 + ); + } + + for record in int4_records + .iter() + .filter(|record| record.register_mma.as_ref().unwrap().shape == RegisterMmaShape::M16n8k64) + { + assert_eq!(record.rust_arguments, ["[i32; 4]", "[u32; 4]", "[u32; 2]"]); + assert_eq!(record.rust_result, "[i32; 4]"); + assert_eq!(record.minimum_ptx, "7.0"); + assert_eq!(record.minimum_sm.as_deref(), Some("sm_80")); + assert_eq!( + record.register_mma.as_ref().unwrap().adapter, + RegisterMmaAdapter::C4I32A4U32B2U32ToD4I32 + ); + } + + let actual_int4_abi_ids = int4_records + .iter() + .map(|record| record.abi_id.as_str()) + .collect::>(); + let expected_int4_abi_ids = (133..=156) + .map(|id| format!("i{id:04}")) + .collect::>(); + assert_eq!( + actual_int4_abi_ids, + expected_int4_abi_ids + .iter() + .map(String::as_str) + .collect::>() + ); + + let int8_k32 = int8_records + .iter() + .find(|record| record.register_mma.as_ref().unwrap().shape == RegisterMmaShape::M16n8k32) + .unwrap(); + let int4_k32 = int4_records + .iter() + .find(|record| record.register_mma.as_ref().unwrap().shape == RegisterMmaShape::M16n8k32) + .unwrap(); + assert_eq!( + int8_k32.register_mma.as_ref().unwrap().adapter, + RegisterMmaAdapter::C4I32A4U32B2U32ToD4I32 + ); + assert_eq!( + int4_k32.register_mma.as_ref().unwrap().adapter, + RegisterMmaAdapter::C4I32A2U32B1U32ToD4I32 + ); + + for policy in &records { + let declaration = declarations[policy.source_record.as_deref().unwrap()]; + assert_eq!(declaration.selections.len(), 1); + assert!( + selection_matches_policy(policy, &declaration.selections[0]).unwrap(), + "{}", + policy.id + ); + validate_imported_policy(policy, declaration).unwrap(); + } + + let valid = records[0]; + let declaration = declarations[valid.source_record.as_deref().unwrap()]; + + let mut non_convergent = valid.clone(); + non_convergent.convergent = false; + assert!( + validate_imported_policy(&non_convergent, declaration) + .unwrap_err() + .to_string() + .contains("effects") + ); + + let mut typed_route = valid.clone(); + typed_route.backend_lowerings[0].mechanism = BackendLoweringMechanism::TypedNvvm; + assert!(validate_imported_policy(&typed_route, declaration).is_err()); + + let mut selectionless = declaration.clone(); + selectionless.selections.clear(); + assert!(validate_imported_policy(valid, &selectionless).is_err()); + + let mut crossed_variant = valid.clone(); + crossed_variant.register_mma.as_mut().unwrap().a_element = RegisterMmaElement::F16; + assert!(validate_imported_policy(&crossed_variant, declaration).is_err()); + + let generated = int8_records + .iter() + .copied() + .find(|record| record.id == "mma_m16n8k16_s32_s8_u8_satfinite") + .unwrap(); + let generated_declaration = declarations[generated.source_record.as_deref().unwrap()]; + + let mut wrong_stub_owner = generated.clone(); + wrong_stub_owner + .register_mma + .as_mut() + .unwrap() + .compatibility_source = RegisterMmaCompatibilitySource::ExistingStub; + assert!(validate_imported_policy(&wrong_stub_owner, generated_declaration).is_err()); + + let mut wrong_b_element = generated.clone(); + wrong_b_element.register_mma.as_mut().unwrap().b_element = RegisterMmaElement::S8; + assert!(validate_imported_policy(&wrong_b_element, generated_declaration).is_err()); + + let mut wrong_overflow = generated.clone(); + wrong_overflow.register_mma.as_mut().unwrap().overflow = RegisterMmaOverflow::Wrapping; + assert!(validate_imported_policy(&wrong_overflow, generated_declaration).is_err()); + + let mut wrong_shape = generated.clone(); + wrong_shape.register_mma.as_mut().unwrap().shape = RegisterMmaShape::M16n8k32; + assert!(validate_imported_policy(&wrong_shape, generated_declaration).is_err()); + + let mut wrong_adapter = generated.clone(); + wrong_adapter.register_mma.as_mut().unwrap().adapter = + RegisterMmaAdapter::C4I32A4U32B2U32ToD4I32; + assert!(validate_imported_policy(&wrong_adapter, generated_declaration).is_err()); + + let binary = binary_records + .iter() + .copied() + .find(|record| record.id == "mma_m8n8k128_s32_b1_xor_popc") + .unwrap(); + let binary_declaration = declarations[binary.source_record.as_deref().unwrap()]; + + let mut wrong_binary_operation = binary.clone(); + wrong_binary_operation + .register_mma + .as_mut() + .unwrap() + .operation = RegisterMmaOperation::AndPopc; + assert!(validate_imported_policy(&wrong_binary_operation, binary_declaration).is_err()); + + let mut wrong_binary_floor = binary.clone(); + wrong_binary_floor.minimum_sm = Some("sm_80".into()); + assert!(validate_imported_policy(&wrong_binary_floor, binary_declaration).is_err()); + + let mut wrong_binary_element = binary.clone(); + wrong_binary_element + .register_mma + .as_mut() + .unwrap() + .a_element = RegisterMmaElement::U4; + assert!(validate_imported_policy(&wrong_binary_element, binary_declaration).is_err()); +} + +#[test] +fn pinned_sparse_mma_records_close_shape_specific_selectors_and_ranges() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let (mut overlay, _) = + read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")).unwrap(); + bind_pinned_abi_ids(&repo_root, &mut overlay); + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let declarations: BTreeMap<_, _> = imported + .intrinsics + .iter() + .map(|record| (record.source_record.as_str(), record)) + .collect(); + let records = overlay + .intrinsics + .iter() + .filter(|record| record.family == "sparse_mma") + .collect::>(); + assert_eq!(records.len(), 114); + assert_eq!( + records + .iter() + .map(|record| record.abi_id.clone()) + .collect::>(), + (163..=251) + .chain(525..=549) + .map(|id| format!("i{id:04}")) + .collect::>() + ); + + let mut derived_ids = BTreeSet::new(); + let mut derived_operation_keys = BTreeSet::new(); + let mut derived_source_records = BTreeSet::new(); + let mut derived_llvm_symbols = BTreeSet::new(); + for record in &records { + let identity = &sparse_mma_recipe(record.sparse_mma.as_ref().unwrap()) + .unwrap() + .identity; + assert_eq!(record.id, identity.id); + assert_eq!(record.operation_key, identity.operation_key); + assert_eq!( + record.source_record.as_deref(), + Some(identity.source_record.as_str()) + ); + assert_eq!( + record.llvm_symbol.as_deref(), + Some(identity.llvm_symbol.as_str()) + ); + assert_eq!(record.expected_ptx.modifiers, identity.ptx_modifiers); + assert!(derived_ids.insert(identity.id.clone())); + assert!(derived_operation_keys.insert(identity.operation_key.clone())); + assert!(derived_source_records.insert(identity.source_record.clone())); + assert!(derived_llvm_symbols.insert(identity.llvm_symbol.clone())); + } + assert_eq!(derived_ids.len(), 114); + assert_eq!(derived_operation_keys.len(), 114); + assert_eq!(derived_source_records.len(), 114); + assert_eq!(derived_llvm_symbols.len(), 114); + + let integer_records = records + .iter() + .copied() + .filter(|record| { + record.sparse_mma.as_ref().unwrap().accumulator == SparseMmaAccumulator::S32 + }) + .collect::>(); + assert_eq!(integer_records.len(), 64); + let f32_records = records + .iter() + .copied() + .filter(|record| { + record.sparse_mma.as_ref().unwrap().accumulator == SparseMmaAccumulator::F32 + }) + .collect::>(); + assert_eq!(f32_records.len(), 25); + let f16_records = records + .iter() + .copied() + .filter(|record| { + record.sparse_mma.as_ref().unwrap().accumulator == SparseMmaAccumulator::F16 + }) + .collect::>(); + assert_eq!(f16_records.len(), 25); + + let variants = integer_records + .iter() + .map(|record| { + let mma = record.sparse_mma.as_ref().unwrap(); + let carrier = + sparse_mma_carrier_recipe(mma.shape, mma.a_element, mma.b_element).unwrap(); + assert_eq!(mma.accumulator, SparseMmaAccumulator::S32); + assert_eq!(mma.selector, carrier.selector); + assert_eq!(mma.adapter, carrier.adapter); + assert_eq!(mma.llvm_adapter, carrier.llvm_adapter); + assert_eq!(record.rust_arguments, carrier.rust_arguments()); + assert_eq!(record.dialect_operands, carrier.dialect_operands()); + assert_eq!(record.llvm_arguments, carrier.llvm_arguments()); + assert_eq!( + record.expected_ptx.operands, + carrier.expected_ptx_operands() + ); + assert_eq!(record.minimum_ptx, sparse_mma_minimum_ptx(mma)); + assert_eq!(record.minimum_sm.as_deref(), Some("sm_80")); + assert_eq!( + record.expected_ptx.operands.last(), + Some(&OperandPattern::Immediate) + ); + assert_eq!( + record.expected_ptx.modifiers.first().map(String::as_str), + Some(match mma.metadata { + SparseMmaMetadata::Standard => "sp", + SparseMmaMetadata::Ordered => "sp::ordered_metadata", + }) + ); + ( + mma.shape, + mma.a_element, + mma.b_element, + mma.overflow, + mma.metadata, + ) + }) + .collect::>(); + let mut expected_variants = BTreeSet::new(); + for shape in [SparseMmaShape::M16n8k32, SparseMmaShape::M16n8k64] { + let metadata = match shape { + SparseMmaShape::M16n8k32 => [ + Some(SparseMmaMetadata::Standard), + Some(SparseMmaMetadata::Ordered), + ], + SparseMmaShape::M16n8k64 => [ + Some(SparseMmaMetadata::Standard), + Some(SparseMmaMetadata::Ordered), + ], + SparseMmaShape::M16n8k128 => [None, None], + }; + for a_element in [SparseMmaElement::S8, SparseMmaElement::U8] { + for b_element in [SparseMmaElement::S8, SparseMmaElement::U8] { + for overflow in [SparseMmaOverflow::Wrapping, SparseMmaOverflow::Satfinite] { + for metadata in metadata.into_iter().flatten() { + expected_variants.insert((shape, a_element, b_element, overflow, metadata)); + } + } + } + } + } + for shape in [SparseMmaShape::M16n8k64, SparseMmaShape::M16n8k128] { + for a_element in [SparseMmaElement::S4, SparseMmaElement::U4] { + for b_element in [SparseMmaElement::S4, SparseMmaElement::U4] { + for overflow in [SparseMmaOverflow::Wrapping, SparseMmaOverflow::Satfinite] { + for metadata in [SparseMmaMetadata::Standard, SparseMmaMetadata::Ordered] { + expected_variants.insert((shape, a_element, b_element, overflow, metadata)); + } + } + } + } + } + assert_eq!(variants, expected_variants); + + let f8f6f4_formats = [ + SparseMmaElement::E2m1, + SparseMmaElement::E2m3, + SparseMmaElement::E3m2, + SparseMmaElement::E4m3, + SparseMmaElement::E5m2, + ]; + assert_eq!( + f32_records + .iter() + .map(|record| { + let mma = record.sparse_mma.as_ref().unwrap(); + (mma.a_element, mma.b_element) + }) + .collect::>(), + f8f6f4_formats + .into_iter() + .flat_map(|a| f8f6f4_formats.into_iter().map(move |b| (a, b))) + .collect() + ); + assert_eq!( + f32_records + .iter() + .map(|record| record.abi_id.as_str()) + .collect::>(), + (227..=251) + .map(|id| format!("i{id:04}")) + .collect::>() + .iter() + .map(String::as_str) + .collect() + ); + assert_eq!( + f16_records + .iter() + .map(|record| { + let mma = record.sparse_mma.as_ref().unwrap(); + (mma.a_element, mma.b_element) + }) + .collect::>(), + f8f6f4_formats + .into_iter() + .flat_map(|a| f8f6f4_formats.into_iter().map(move |b| (a, b))) + .collect() + ); + assert_eq!( + f16_records + .iter() + .map(|record| record.abi_id.as_str()) + .collect::>(), + (525..=549) + .map(|id| format!("i{id:04}")) + .collect::>() + .iter() + .map(String::as_str) + .collect() + ); + for record in &f32_records { + let mma = record.sparse_mma.as_ref().unwrap(); + assert_eq!(mma.shape, SparseMmaShape::M16n8k64); + assert_eq!(mma.accumulator, SparseMmaAccumulator::F32); + assert_eq!(mma.overflow, SparseMmaOverflow::NotApplicable); + assert_eq!(mma.metadata, SparseMmaMetadata::Ordered); + assert_eq!(mma.selector, SparseMmaSelector::ImmediateZero); + assert_eq!( + mma.adapter, + SparseMmaAdapter::C4F32A4U32B4U32MetadataU32SelectorU32ToD4F32 + ); + assert_eq!( + mma.llvm_adapter, + SparseMmaLlvmAdapter::A4I32B4I32C4F32MetadataI32SelectorI32ToD4F32 + ); + assert_eq!( + record.rust_arguments, + ["[f32; 4]", "[u32; 4]", "[u32; 4]", "u32", "u32"] + ); + assert_eq!(record.rust_result, "[f32; 4]"); + assert_eq!( + record.dialect_operands, + [ + "f32", "f32", "f32", "f32", "u32", "u32", "u32", "u32", "u32", "u32", "u32", "u32", + "u32", "u32" + ] + ); + assert_eq!(record.dialect_results, ["f32", "f32", "f32", "f32"]); + assert_eq!( + record.llvm_arguments, + [ + "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", "f32", "f32", "f32", "f32", + "i32", "i32" + ] + ); + assert_eq!(record.llvm_results, ["f32", "f32", "f32", "f32"]); + assert_eq!(record.minimum_ptx, "8.7"); + assert_eq!(record.minimum_sm, None); + // Same contract as the F16 accumulator: both float forms are gated + // on `hasMMABlockScale()`, so neither is narrower than the other. + assert_eq!(record.targets, SPARSE_MMA_F8F6F4_TARGETS); + assert_eq!(record.backend_lowerings.len(), 2); + assert!(record.backend_lowerings.iter().all(|lowering| { + lowering.mechanism == BackendLoweringMechanism::InlinePtx + && lowering.minimum_ptx.as_deref() == Some("8.7") + && lowering.minimum_sm.is_none() + })); + assert_eq!( + record.expected_ptx.operands, + [ + OperandPattern::RegisterList { length: 4 }, + OperandPattern::RegisterList { length: 4 }, + OperandPattern::RegisterList { length: 4 }, + OperandPattern::RegisterList { length: 4 }, + OperandPattern::Register, + OperandPattern::Immediate, + ] + ); + } + + for record in &f16_records { + let mma = record.sparse_mma.as_ref().unwrap(); + assert_eq!(mma.shape, SparseMmaShape::M16n8k64); + assert_eq!(mma.accumulator, SparseMmaAccumulator::F16); + assert_eq!(mma.overflow, SparseMmaOverflow::NotApplicable); + assert_eq!(mma.metadata, SparseMmaMetadata::Ordered); + assert_eq!(mma.selector, SparseMmaSelector::ImmediateZero); + assert_eq!( + mma.adapter, + SparseMmaAdapter::C2U32A4U32B4U32MetadataU32SelectorU32ToD2U32 + ); + assert_eq!( + mma.llvm_adapter, + SparseMmaLlvmAdapter::A4I32B4I32C2V2F16MetadataI32SelectorI32ToD2V2F16 + ); + assert_eq!( + record.rust_arguments, + ["[u32; 2]", "[u32; 4]", "[u32; 4]", "u32", "u32"] + ); + assert_eq!(record.rust_result, "[u32; 2]"); + assert_eq!( + record.dialect_operands, + [ + "i32", "i32", "u32", "u32", "u32", "u32", "u32", "u32", "u32", "u32", "u32", "u32" + ] + ); + assert_eq!(record.dialect_results, ["i32", "i32"]); + assert_eq!( + record.llvm_arguments, + [ + "i32", "i32", "i32", "i32", "i32", "i32", "i32", "i32", "v2f16", "v2f16", "i32", + "i32" + ] + ); + assert_eq!(record.llvm_results, ["v2f16", "v2f16"]); + assert_eq!(record.minimum_ptx, "8.7"); + assert_eq!(record.minimum_sm, None); + assert_eq!(record.targets, SPARSE_MMA_F8F6F4_TARGETS); + assert!(record.convergent && !record.pure); + assert!(record.backend_lowerings.iter().all(|lowering| { + lowering.mechanism == BackendLoweringMechanism::InlinePtx + && lowering.minimum_ptx.as_deref() == Some("8.7") + && lowering.minimum_sm.is_none() + })); + assert_eq!( + record.expected_ptx.operands, + [ + OperandPattern::RegisterList { length: 2 }, + OperandPattern::RegisterList { length: 4 }, + OperandPattern::RegisterList { length: 4 }, + OperandPattern::RegisterList { length: 2 }, + OperandPattern::Register, + OperandPattern::Immediate, + ] + ); + } + + for policy in &records { + let declaration = declarations[policy.source_record.as_deref().unwrap()]; + assert_eq!(declaration.selections.len(), 1); + assert!( + selection_matches_policy(policy, &declaration.selections[0]).unwrap(), + "{}", + policy.id + ); + validate_imported_policy(policy, declaration).unwrap(); + } + + let mut selectionless = declarations[records[0].source_record.as_deref().unwrap()].clone(); + selectionless.selections.clear(); + assert!(validate_imported_policy(records[0], &selectionless).is_err()); + + for (id, range_prefix, wrong_range) in [ + ("mma_sp_m16n8k32_s32_s8", "Range"), + ("mma_sp_m16n8k64_s32_s8", "Range"), + ( + "mma_sp_ordered_metadata_m16n8k64_s32_s4", + "Range", + ), + ("mma_sp_m16n8k64_s32_s4", "Range"), + ( + "mma_sp_ordered_metadata_m16n8k128_s32_s4", + "Range", + ), + ("mma_sp_m16n8k128_s32_s4", "Range"), + ( + "mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f32_e2m1_e2m1_f32", + "Range", + ), + ( + "mma_sp_ordered_metadata_m16n8k64_kind_f8f6f4_f16_e2m1_e2m1_f16", + "Range", + ), + ] { + let valid = records + .iter() + .copied() + .find(|record| record.id == id) + .unwrap(); + let declaration = declarations[valid.source_record.as_deref().unwrap()]; + + let mut runtime_selector = valid.clone(); + *runtime_selector.expected_ptx.operands.last_mut().unwrap() = OperandPattern::Register; + assert!( + validate_imported_policy(&runtime_selector, declaration) + .unwrap_err() + .to_string() + .contains("exact selection") + ); + + let mut wrong_declaration = declaration.clone(); + *wrong_declaration + .properties + .iter_mut() + .find(|property| property.starts_with(range_prefix)) + .unwrap() = wrong_range.into(); + assert!( + validate_imported_policy(valid, &wrong_declaration) + .unwrap_err() + .to_string() + .contains("immediate range") + ); + } + + let k64 = records + .iter() + .copied() + .find(|record| record.id == "mma_sp_m16n8k64_s32_s8") + .unwrap(); + let k64_declaration = declarations[k64.source_record.as_deref().unwrap()]; + assert_eq!(k64.minimum_ptx, "7.1"); + let ordered_k64 = records + .iter() + .copied() + .find(|record| record.id == "mma_sp_ordered_metadata_m16n8k64_s32_s8") + .unwrap(); + assert_eq!(ordered_k64.minimum_ptx, "8.5"); + + let f8f6f4 = f32_records[0]; + let f8f6f4_declaration = declarations[f8f6f4.source_record.as_deref().unwrap()]; + let mut widened_family = f8f6f4.clone(); + widened_family.targets = "sm_120f".into(); + assert!(validate_imported_policy(&widened_family, f8f6f4_declaration).is_err()); + let mut widened_architecture = f8f6f4.clone(); + widened_architecture.targets = "sm_121a".into(); + assert!(validate_imported_policy(&widened_architecture, f8f6f4_declaration).is_err()); + + let f8f6f4_f16 = f16_records[0]; + for block_scale in [f8f6f4, f8f6f4_f16] { + let mut missing_predicate = + declarations[block_scale.source_record.as_deref().unwrap()].clone(); + missing_predicate.selections[0].predicates.clear(); + assert!( + validate_imported_policy(block_scale, &missing_predicate) + .unwrap_err() + .to_string() + .contains("exact selection") + ); + } + + let ordered_k64_int4 = records + .iter() + .copied() + .find(|record| record.id == "mma_sp_ordered_metadata_m16n8k64_s32_s4") + .unwrap(); + let ordered_k64_int4_declaration = + declarations[ordered_k64_int4.source_record.as_deref().unwrap()]; + assert_eq!(ordered_k64_int4.minimum_ptx, "8.5"); + assert_eq!(ordered_k64_int4.rust_arguments[1], "[u32; 2]"); + assert_eq!(ordered_k64_int4.rust_arguments[2], "[u32; 2]"); + assert_eq!(ordered_k64_int4.llvm_arguments.len(), 10); + assert_eq!( + ordered_k64_int4.sparse_mma.as_ref().unwrap().selector, + SparseMmaSelector::ImmediateZeroOrOne + ); + assert_eq!( + ordered_k64_int4.sparse_mma.as_ref().unwrap().adapter, + SparseMmaAdapter::C4I32A2U32B2U32MetadataU32SelectorU32ToD4I32 + ); + assert_eq!( + ordered_k64_int4.sparse_mma.as_ref().unwrap().llvm_adapter, + SparseMmaLlvmAdapter::A2I32B2I32C4I32MetadataI32SelectorI32ToD4I32 + ); + let standard_k64_int4 = records + .iter() + .copied() + .find(|record| record.id == "mma_sp_m16n8k64_s32_s4") + .unwrap(); + assert_eq!(standard_k64_int4.minimum_ptx, "7.1"); + assert_eq!( + standard_k64_int4.rust_arguments, + ordered_k64_int4.rust_arguments + ); + assert_eq!( + standard_k64_int4.llvm_arguments, + ordered_k64_int4.llvm_arguments + ); + + let ordered_k128_int4 = records + .iter() + .copied() + .find(|record| record.id == "mma_sp_ordered_metadata_m16n8k128_s32_s4") + .unwrap(); + let ordered_k128_int4_declaration = + declarations[ordered_k128_int4.source_record.as_deref().unwrap()]; + assert_eq!(ordered_k128_int4.minimum_ptx, "8.5"); + assert_eq!(ordered_k128_int4.rust_arguments[1], "[u32; 4]"); + assert_eq!(ordered_k128_int4.rust_arguments[2], "[u32; 4]"); + assert_eq!(ordered_k128_int4.llvm_arguments.len(), 14); + assert_eq!( + ordered_k128_int4.sparse_mma.as_ref().unwrap().selector, + SparseMmaSelector::ImmediateZero + ); + assert_eq!( + ordered_k128_int4.sparse_mma.as_ref().unwrap().adapter, + SparseMmaAdapter::C4I32A4U32B4U32MetadataU32SelectorU32ToD4I32 + ); + assert_eq!( + ordered_k128_int4.sparse_mma.as_ref().unwrap().llvm_adapter, + SparseMmaLlvmAdapter::A4I32B4I32C4I32MetadataI32SelectorI32ToD4I32 + ); + let standard_k128_int4 = records + .iter() + .copied() + .find(|record| record.id == "mma_sp_m16n8k128_s32_s4") + .unwrap(); + assert_eq!(standard_k128_int4.minimum_ptx, "7.1"); + assert_eq!( + standard_k128_int4.rust_arguments, + ordered_k128_int4.rust_arguments + ); + assert_eq!( + standard_k128_int4.llvm_arguments, + ordered_k128_int4.llvm_arguments + ); + + let mut wrong_k128_selector = ordered_k128_int4.clone(); + wrong_k128_selector.sparse_mma.as_mut().unwrap().selector = + SparseMmaSelector::ImmediateZeroOrOne; + assert!(validate_imported_policy(&wrong_k128_selector, ordered_k128_int4_declaration).is_err()); + + let mut mixed_k128_width = ordered_k128_int4.clone(); + mixed_k128_width.sparse_mma.as_mut().unwrap().b_element = SparseMmaElement::U8; + assert!( + validate_imported_policy(&mixed_k128_width, ordered_k128_int4_declaration) + .unwrap_err() + .to_string() + .contains("unsupported sparse-MMA variant") + ); + + let mut mixed_width = ordered_k64_int4.clone(); + mixed_width.sparse_mma.as_mut().unwrap().b_element = SparseMmaElement::U8; + assert!( + validate_imported_policy(&mixed_width, ordered_k64_int4_declaration) + .unwrap_err() + .to_string() + .contains("unsupported sparse-MMA variant") + ); + + let mut wrong_k64_selector = k64.clone(); + wrong_k64_selector.sparse_mma.as_mut().unwrap().selector = + SparseMmaSelector::ImmediateZeroOrOne; + assert!(validate_imported_policy(&wrong_k64_selector, k64_declaration).is_err()); + + let mut wrong_k64_adapter = k64.clone(); + wrong_k64_adapter.sparse_mma.as_mut().unwrap().adapter = + SparseMmaAdapter::C4I32A2U32B2U32MetadataU32SelectorU32ToD4I32; + assert!(validate_imported_policy(&wrong_k64_adapter, k64_declaration).is_err()); + + let mut wrong_k64_llvm_adapter = k64.clone(); + wrong_k64_llvm_adapter + .sparse_mma + .as_mut() + .unwrap() + .llvm_adapter = SparseMmaLlvmAdapter::A2I32B2I32C4I32MetadataI32SelectorI32ToD4I32; + assert!(validate_imported_policy(&wrong_k64_llvm_adapter, k64_declaration).is_err()); + + let mut wrong_k64_shape = k64.clone(); + wrong_k64_shape.sparse_mma.as_mut().unwrap().shape = SparseMmaShape::M16n8k32; + assert!(validate_imported_policy(&wrong_k64_shape, k64_declaration).is_err()); + + let mut wrong_k64_carriers = k64.clone(); + wrong_k64_carriers.dialect_operands.pop(); + assert!(validate_imported_policy(&wrong_k64_carriers, k64_declaration).is_err()); + + let mut wrong_k64_lowering = k64.clone(); + wrong_k64_lowering.lowering = "generated_register_mma".into(); + assert!(validate_imported_policy(&wrong_k64_lowering, k64_declaration).is_err()); + + let mut mismatched_metadata_identity = k64.clone(); + mismatched_metadata_identity + .sparse_mma + .as_mut() + .unwrap() + .metadata = SparseMmaMetadata::Ordered; + assert!(validate_imported_policy(&mismatched_metadata_identity, k64_declaration).is_err()); +} + +#[test] +fn movmatrix_recipe_is_exact_and_fails_closed() { + let valid = movmatrix_policy(); + validate_ptx_native_policy(&valid).unwrap(); + + let reject = |policy: &OverlayIntrinsic, expected: &str| { + let message = validate_ptx_native_policy(policy).unwrap_err().to_string(); + assert!(message.contains(expected), "unexpected error: {message}"); + }; + + let mut wrong_shape = valid.clone(); + wrong_shape.expected_ptx.modifiers[2] = "m16n8".into(); + reject(&wrong_shape, "closed movmatrix recipe"); + + let mut wrong_participation = valid.clone(); + wrong_participation.convergent = false; + reject(&wrong_participation, "closed movmatrix recipe"); + + // A warp collective is not a function of its own operand, so the + // pure/no-memory pair every other collective avoids must stay rejected + // here too. + let mut wrong_purity = valid.clone(); + wrong_purity.pure = true; + reject(&wrong_purity, "closed movmatrix recipe"); + + let mut wrong_memory = valid.clone(); + wrong_memory.memory = "none".into(); + reject(&wrong_memory, "closed movmatrix recipe"); + + let mut wrong_floor = valid.clone(); + wrong_floor.backend_lowerings[0].minimum_ptx = Some("8.0".into()); + reject(&wrong_floor, "exact movmatrix floor"); + + let mut mixed = valid; + mixed.warp_barrier = Some(crate::model::WarpBarrier { + participation: + WarpBarrierParticipation::ExecutingLaneNamedAllNamedLanesSameInstructionAndMask, + legacy_pre_sm70: PreSm70MemberMaskRule::AllNamedLanesConvergedAndOnlyNamedLanesActive, + adapter: WarpBarrierAdapter::DirectMemberMask, + mask_encoding: WarpBarrierMaskEncoding::RegisterOrImmediate, + memory_ordering: WarpBarrierMemoryOrdering::ParticipatingLanes, + }); + reject(&mixed, "mixes another generated-family contract"); +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/tests/mod.rs b/crates/cuda-intrinsics-gen/src/resolve/tests/mod.rs new file mode 100644 index 0000000000..d98fc552a5 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/tests/mod.rs @@ -0,0 +1,16 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +mod evidence; +mod fixtures; +mod misc_families; +mod mma; +mod overlay_schema; +mod scalar_packed; +mod sreg_threadfence; +mod targets_abi; +mod tcgen05; +mod tma_clc_cluster; +mod warp; diff --git a/crates/cuda-intrinsics-gen/src/resolve/tests/overlay_schema.rs b/crates/cuda-intrinsics-gen/src/resolve/tests/overlay_schema.rs new file mode 100644 index 0000000000..946333f7cf --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/tests/overlay_schema.rs @@ -0,0 +1,736 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + CatalogFile, OverlayShardFile, PackedAluFormat, PackedAluOperation, RegisterMmaAccumulator, + RegisterMmaKind, +}; +use crate::util::read_json; +use std::collections::{BTreeMap, BTreeSet}; +use std::path::Path; + +use super::fixtures::*; +use crate::resolve::abi_ledger::*; +use crate::resolve::driver::*; +use crate::resolve::guards::*; +use crate::resolve::overlay::*; + +#[test] +fn duplicate_values_are_rejected() { + let mut values = BTreeSet::new(); + insert_unique(&mut values, "thread_idx_x", "catalog ID").unwrap(); + let error = insert_unique(&mut values, "thread_idx_x", "catalog ID").unwrap_err(); + assert!(error.to_string().contains("duplicate catalog ID")); +} + +#[test] +fn overloaded_symbols_require_distinct_resolved_identities() { + let bf16 = packed_alu_policy(PackedAluFormat::Bf16x2, PackedAluOperation::Abs); + let f16 = packed_alu_policy(PackedAluFormat::F16x2, PackedAluOperation::Abs); + validate_unique_overlay(&[bf16.clone(), f16.clone()], 1).unwrap(); + + let mut unresolved = f16.clone(); + unresolved.resolved_llvm_symbol = None; + let error = validate_unique_overlay(&[bf16.clone(), unresolved], 1).unwrap_err(); + assert!(error.to_string().contains("without a resolved symbol")); + + let mut duplicate = f16; + duplicate.resolved_llvm_symbol = bf16.resolved_llvm_symbol.clone(); + let error = validate_unique_overlay(&[bf16, duplicate], 1).unwrap_err(); + assert!(error.to_string().contains("duplicate resolved LLVM symbol")); +} + +#[test] +fn overlay_manifest_loads_sorted_family_shards() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let (overlay, hash) = + read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")).unwrap(); + assert_eq!(overlay.schema, OVERLAY_SCHEMA); + assert_eq!(overlay.shards.len(), 64); + assert_eq!(overlay.intrinsics.len(), 1016); + assert_eq!( + overlay + .intrinsics + .iter() + .filter(|record| record.family == "scalar_arithmetic") + .count(), + 64 + ); + assert_eq!( + overlay + .intrinsics + .iter() + .filter(|record| record.family == "scalar_conversion") + .count(), + 10 + ); + assert_eq!( + overlay + .intrinsics + .iter() + .filter(|record| record.family == "prmt") + .count(), + 7 + ); + assert_eq!( + overlay + .intrinsics + .iter() + .filter(|record| record.family == "debug_control") + .count(), + 3 + ); + assert_eq!( + overlay + .intrinsics + .iter() + .filter(|record| record.family == "stmatrix") + .count(), + 4 + ); + assert_eq!( + overlay + .intrinsics + .iter() + .filter(|record| record.family == "packed_alu") + .count(), + 30 + ); + assert_eq!( + overlay + .intrinsics + .iter() + .filter(|record| record.family == "packed_conversion") + .count(), + 18 + ); + assert_eq!( + overlay + .intrinsics + .iter() + .filter(|record| record.family == "active_mask") + .count(), + 1 + ); + assert_eq!( + overlay + .intrinsics + .iter() + .filter(|record| record.family == "dotprod") + .count(), + 4 + ); + assert_eq!( + overlay + .intrinsics + .iter() + .filter(|record| record.family == "ldmatrix") + .count(), + 18 + ); + assert_eq!( + overlay + .intrinsics + .iter() + .filter(|record| record.family == "register_mma") + .count(), + 154 + ); + assert_eq!( + overlay + .intrinsics + .iter() + .filter(|record| record.family == "sparse_mma") + .count(), + 114 + ); + assert_eq!( + overlay + .intrinsics + .iter() + .filter(|record| record.family == "cp_async_copy") + .count(), + 8 + ); + assert_eq!( + overlay + .intrinsics + .iter() + .filter(|record| record.family == "cp_async_control") + .count(), + 3 + ); + assert_eq!( + overlay + .intrinsics + .iter() + .filter(|record| record.family == "cp_async_mbarrier") + .count(), + 4 + ); + assert_eq!( + overlay + .intrinsics + .iter() + .filter(|record| record.family == "mbarrier_basic") + .count(), + 4 + ); + assert_eq!( + overlay + .intrinsics + .iter() + .filter(|record| record.family == "cluster_memory") + .count(), + 2 + ); + assert_eq!( + overlay + .intrinsics + .iter() + .filter(|record| record.family == "clc") + .count(), + 6 + ); + assert_eq!( + overlay + .intrinsics + .iter() + .filter(|record| record.family == "sync") + .count(), + 4 + ); + assert_eq!( + overlay + .intrinsics + .iter() + .filter(|record| record.family == "vote") + .count(), + 4 + ); + assert_eq!( + overlay + .intrinsics + .iter() + .filter(|record| record.family == "warp_barrier") + .count(), + 1 + ); + assert_eq!( + overlay + .intrinsics + .iter() + .filter(|record| record.family == "warp_match") + .count(), + 4 + ); + assert_eq!( + overlay + .intrinsics + .iter() + .filter(|record| record.family == "warp_shuffle") + .count(), + 12 + ); + assert_eq!(hash.len(), 64); + + for invalid in [ + "../outside.toml", + "/absolute.toml", + "other/family.toml", + "overlay/../outside.toml", + "overlay/not-toml.json", + ] { + assert!(validate_overlay_shard_path(invalid).is_err(), "{invalid}"); + } +} + +#[test] +fn overlay_shard_schema_range_is_composable_and_new_fields_fail_closed() { + let shard = |schema, sparse_mma_f8f6f4_f32, prmt| OverlayShardFile { + schema, + family: "sparse_mma".into(), + intrinsics: vec![], + register_mma_int4: None, + register_mma_int8: None, + register_mma_b1: None, + register_mma_f8f6f4_f32: None, + register_mma_f8f6f4_f16: None, + register_mma_mxf8f6f4_f32: None, + register_mma_fp8: None, + register_mma_ampere_float: None, + sparse_mma_integer: None, + sparse_mma_f8f6f4_f32, + sparse_mma_f8f6f4_f16: None, + prmt, + packed_conversion_fp8: None, + packed_conversion_fp8_f16x2: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cluster_sreg: None, + cluster_barrier: None, + mbarrier_extended: None, + special_registers: None, + debug_control: None, + threadfence: None, + cluster_memory: None, + stmatrix: None, + clc: None, + wgmma_controls: None, + tma: None, + tcgen05: None, + }; + let path = Path::new("intrinsics/overlay/test.toml"); + validate_overlay_shard_schema(&shard(26, None, None), path).unwrap(); + validate_overlay_shard_schema(&shard(27, None, None), path).unwrap(); + validate_overlay_shard_schema(&shard(28, None, None), path).unwrap(); + validate_overlay_shard_schema(&shard(29, None, None), path).unwrap(); + validate_overlay_shard_schema(&shard(30, None, None), path).unwrap(); + validate_overlay_shard_schema(&shard(31, None, None), path).unwrap(); + validate_overlay_shard_schema(&shard(32, None, None), path).unwrap(); + validate_overlay_shard_schema(&shard(33, None, None), path).unwrap(); + validate_overlay_shard_schema(&shard(34, None, None), path).unwrap(); + validate_overlay_shard_schema(&shard(35, None, None), path).unwrap(); + validate_overlay_shard_schema(&shard(27, Some(test_f8f6f4_admission()), None), path).unwrap(); + validate_overlay_shard_schema_with_max( + &shard(27, Some(test_f8f6f4_admission()), None), + path, + 30, + ) + .unwrap(); + validate_overlay_shard_schema(&shard(28, None, Some(test_prmt_admission())), path).unwrap(); + validate_overlay_shard_schema_with_max(&shard(28, None, Some(test_prmt_admission())), path, 30) + .unwrap(); + + assert!(validate_overlay_shard_schema(&shard(25, None, None), path).is_err()); + for schema in 35..=OVERLAY_SHARD_SCHEMA { + validate_overlay_shard_schema(&shard(schema, None, None), path).unwrap(); + } + assert!( + validate_overlay_shard_schema(&shard(OVERLAY_SHARD_SCHEMA + 1, None, None), path).is_err() + ); + let error = + validate_overlay_shard_schema(&shard(26, Some(test_f8f6f4_admission()), None), path) + .unwrap_err(); + assert!( + error + .to_string() + .contains("requires overlay shard schema 27") + ); + let error = validate_overlay_shard_schema(&shard(27, None, Some(test_prmt_admission())), path) + .unwrap_err(); + assert!( + error + .to_string() + .contains("requires overlay shard schema 28") + ); + let mut f16_mma = shard(REGISTER_MMA_F8F6F4_F16_SHARD_SCHEMA, None, None); + f16_mma.family = "register_mma".into(); + f16_mma.register_mma_f8f6f4_f16 = Some(test_register_mma_f8f6f4_admission( + RegisterMmaAccumulator::F16, + )); + validate_overlay_shard_schema(&f16_mma, path).unwrap(); + f16_mma.schema -= 1; + assert!( + validate_overlay_shard_schema(&f16_mma, path) + .unwrap_err() + .to_string() + .contains("requires overlay shard schema 47") + ); + + let mut sparse_f16_mma = shard(SPARSE_MMA_F8F6F4_F16_SHARD_SCHEMA, None, None); + sparse_f16_mma.sparse_mma_f8f6f4_f16 = Some(test_sparse_mma_f8f6f4_f16_admission()); + validate_overlay_shard_schema(&sparse_f16_mma, path).unwrap(); + sparse_f16_mma.schema -= 1; + assert!( + validate_overlay_shard_schema(&sparse_f16_mma, path) + .unwrap_err() + .to_string() + .contains("requires overlay shard schema 50") + ); + + let mut standard_fp8_mma = shard(REGISTER_MMA_FP8_SHARD_SCHEMA, None, None); + standard_fp8_mma.family = "register_mma".into(); + standard_fp8_mma.register_mma_fp8 = Some(test_register_mma_fp8_admission()); + validate_overlay_shard_schema(&standard_fp8_mma, path).unwrap(); + standard_fp8_mma.schema -= 1; + assert!( + validate_overlay_shard_schema(&standard_fp8_mma, path) + .unwrap_err() + .to_string() + .contains("requires overlay shard schema 48") + ); + + let fp8_shard = |schema| OverlayShardFile { + schema, + family: "packed_conversion".into(), + intrinsics: vec![], + register_mma_int4: None, + register_mma_int8: None, + register_mma_b1: None, + register_mma_f8f6f4_f32: None, + register_mma_f8f6f4_f16: None, + register_mma_mxf8f6f4_f32: None, + register_mma_fp8: None, + register_mma_ampere_float: None, + sparse_mma_integer: None, + sparse_mma_f8f6f4_f32: None, + sparse_mma_f8f6f4_f16: None, + prmt: None, + packed_conversion_fp8: Some(test_fp8_conversion_admission()), + packed_conversion_fp8_f16x2: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cluster_sreg: None, + cluster_barrier: None, + mbarrier_extended: None, + special_registers: None, + debug_control: None, + threadfence: None, + cluster_memory: None, + stmatrix: None, + clc: None, + wgmma_controls: None, + tma: None, + tcgen05: None, + }; + validate_overlay_shard_schema(&fp8_shard(29), path).unwrap(); + validate_overlay_shard_schema_with_max(&fp8_shard(29), path, 30).unwrap(); + let error = validate_overlay_shard_schema(&fp8_shard(28), path).unwrap_err(); + assert!( + error + .to_string() + .contains("requires overlay shard schema 29") + ); + + let cluster_shard = OverlayShardFile { + schema: 31, + family: "cluster_barrier".into(), + intrinsics: vec![], + register_mma_int4: None, + register_mma_int8: None, + register_mma_b1: None, + register_mma_f8f6f4_f32: None, + register_mma_f8f6f4_f16: None, + register_mma_mxf8f6f4_f32: None, + register_mma_fp8: None, + register_mma_ampere_float: None, + sparse_mma_integer: None, + sparse_mma_f8f6f4_f32: None, + sparse_mma_f8f6f4_f16: None, + prmt: None, + packed_conversion_fp8: None, + packed_conversion_fp8_f16x2: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cluster_sreg: None, + cluster_barrier: Some(test_cluster_barrier_admission()), + mbarrier_extended: None, + special_registers: None, + debug_control: None, + threadfence: None, + cluster_memory: None, + stmatrix: None, + clc: None, + wgmma_controls: None, + tma: None, + tcgen05: None, + }; + validate_overlay_shard_schema_with_max(&cluster_shard, path, 31).unwrap(); + let mut old_cluster_shard = cluster_shard; + old_cluster_shard.schema = 30; + assert!( + validate_overlay_shard_schema_with_max(&old_cluster_shard, path, 31) + .unwrap_err() + .to_string() + .contains("requires overlay shard schema 31") + ); + + let extended_shard = OverlayShardFile { + schema: MBARRIER_EXTENDED_SHARD_SCHEMA, + family: "mbarrier_extended".into(), + intrinsics: vec![], + register_mma_int4: None, + register_mma_int8: None, + register_mma_b1: None, + register_mma_f8f6f4_f32: None, + register_mma_f8f6f4_f16: None, + register_mma_mxf8f6f4_f32: None, + register_mma_fp8: None, + register_mma_ampere_float: None, + sparse_mma_integer: None, + sparse_mma_f8f6f4_f32: None, + sparse_mma_f8f6f4_f16: None, + prmt: None, + packed_conversion_fp8: None, + packed_conversion_fp8_f16x2: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cluster_sreg: None, + cluster_barrier: None, + mbarrier_extended: Some(test_mbarrier_extended_admission()), + special_registers: None, + debug_control: None, + threadfence: None, + cluster_memory: None, + stmatrix: None, + clc: None, + wgmma_controls: None, + tma: None, + tcgen05: None, + }; + validate_overlay_shard_schema_with_max(&extended_shard, path, MBARRIER_EXTENDED_SHARD_SCHEMA) + .unwrap(); + let mut old_extended_shard = extended_shard; + old_extended_shard.schema = MBARRIER_EXTENDED_SHARD_SCHEMA - 1; + assert!( + validate_overlay_shard_schema_with_max( + &old_extended_shard, + path, + MBARRIER_EXTENDED_SHARD_SCHEMA, + ) + .unwrap_err() + .to_string() + .contains("requires overlay shard schema 40") + ); +} + +#[test] +fn existing_catalog_intrinsics_keep_their_json_shape_without_kind() { + fn abi_number(abi_id: &str) -> u16 { + validate_abi_id(abi_id).unwrap(); + abi_id[1..].parse().unwrap() + } + + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let original: serde_json::Value = + read_json(&repo_root.join("intrinsics/catalog.json")).unwrap(); + let catalog: CatalogFile = read_json(&repo_root.join("intrinsics/catalog.json")).unwrap(); + let serialized = serde_json::to_value(&catalog).unwrap(); + let original_by_abi = original["intrinsics"] + .as_array() + .unwrap() + .iter() + .map(|record| { + let abi_id = record["rust"]["abi_id"].as_str().unwrap(); + (abi_number(abi_id), record) + }) + .collect::>(); + let serialized_by_abi = serialized["intrinsics"] + .as_array() + .unwrap() + .iter() + .map(|record| { + let abi_id = record["rust"]["abi_id"].as_str().unwrap(); + (abi_number(abi_id), record) + }) + .collect::>(); + for abi in 1..=503 { + assert_eq!(original_by_abi.get(&abi), serialized_by_abi.get(&abi)); + } + + let existing = catalog + .intrinsics + .iter() + .filter(|record| abi_number(&record.rust.abi_id) <= 503) + .collect::>(); + assert_eq!(existing.len(), 503); + assert!(existing.iter().all(|record| { + record + .register_mma + .as_ref() + .is_none_or(|mma| mma.kind.is_none()) + })); + + let standard_fp8 = catalog + .intrinsics + .iter() + .filter(|record| (504..=519).contains(&abi_number(&record.rust.abi_id))) + .collect::>(); + assert!(standard_fp8.is_empty() || standard_fp8.len() == 16); + if standard_fp8.len() == 16 { + assert!(standard_fp8.iter().all(|record| { + record + .register_mma + .as_ref() + .is_some_and(|mma| mma.kind == Some(RegisterMmaKind::Standard)) + })); + for abi in 504..=519 { + assert_eq!(original_by_abi[&abi]["register_mma"]["kind"], "standard"); + } + } +} + +#[test] +fn duplicate_identity_surfaces_are_rejected_independently() { + let first = policy(); + + let mut second = distinct_policy(); + second.abi_id = first.abi_id.clone(); + assert!( + validate_unique_overlay(&[first.clone(), second], 1) + .unwrap_err() + .to_string() + .contains("duplicate intrinsic ABI ID") + ); + + let mut second = distinct_policy(); + second.operation_key = first.operation_key.clone(); + assert!( + validate_unique_overlay(&[first.clone(), second], 1) + .unwrap_err() + .to_string() + .contains("duplicate intrinsic operation key") + ); + + let mut second = distinct_policy(); + second.public_rust_path = first.public_rust_path.clone(); + assert!( + validate_unique_overlay(&[first.clone(), second], 1) + .unwrap_err() + .to_string() + .contains("duplicate public Rust path") + ); + + let mut second = distinct_policy(); + second.dialect_op_name = first.dialect_op_name.clone(); + assert!( + validate_unique_overlay(&[first.clone(), second], 1) + .unwrap_err() + .to_string() + .contains("duplicate dialect op variant") + ); + + let mut second = distinct_policy(); + second.llvm_symbol = first.llvm_symbol.clone(); + assert!( + validate_unique_overlay(&[first, second], 1) + .unwrap_err() + .to_string() + .contains("duplicate LLVM symbol") + ); +} + +#[test] +fn safe_record_requires_an_allowlist_reason() { + let mut record = policy(); + record.safe_allowlist_reason = None; + assert!( + validate_imported_policy(&record, &declaration()) + .unwrap_err() + .to_string() + .contains("safe_allowlist_reason") + ); +} + +#[test] +fn candidate_resolution_is_the_only_path_that_can_omit_evidence() { + let repo = repo_without_evidence(); + let candidate = resolve_candidate( + &repo.0, + "thread_idx_x", + "LLVM version candidate", + &"a".repeat(64), + "sm_80", + "+ptx70", + ) + .unwrap(); + assert_eq!(candidate.catalog.intrinsics.len(), 1); + assert_eq!(candidate.catalog.intrinsics[0].id, "thread_idx_x"); + assert_eq!(candidate.catalog.intrinsics[0].backend.status, "candidate"); + + let scalar = resolve_candidate( + &repo.0, + "i0390", + "LLVM version candidate", + &"a".repeat(64), + "sm_80", + "+ptx70", + ) + .unwrap(); + assert_eq!(scalar.catalog.intrinsics[0].id, "mul_rn_f64"); + assert!(candidate.catalog.inputs.evidence_sha256.is_empty()); + + for (dense_id, abi_id) in [ + ("mma_m16n8k32_f32_e2m1_e2m1", "i0454"), + ("mma_m16n8k32_f16_e2m1_e2m1", "i0479"), + ] { + let dense_by_name = resolve_candidate( + &repo.0, + dense_id, + "LLVM version candidate", + &"a".repeat(64), + "sm_120f", + "+ptx88", + ) + .unwrap(); + let dense_by_abi = resolve_candidate( + &repo.0, + abi_id, + "LLVM version candidate", + &"a".repeat(64), + "sm_120f", + "+ptx88", + ) + .unwrap(); + assert_eq!( + dense_by_name.catalog.intrinsics, + dense_by_abi.catalog.intrinsics + ); + for lookup in [dense_id, abi_id] { + let error = resolve_candidate( + &repo.0, + lookup, + "LLVM version candidate", + &"a".repeat(64), + "sm_120f", + "+ptx87", + ) + .unwrap_err(); + assert!(error.to_string().contains(dense_id), "{error:#}"); + } + } + + let error = resolve(&repo.0).unwrap_err(); + assert!( + error.to_string().contains("intrinsics/evidence"), + "{error:#}" + ); + let error = resolve_candidate( + &repo.0, + "not_an_intrinsic", + "LLVM version candidate", + &"a".repeat(64), + "sm_80", + "+ptx70", + ) + .unwrap_err(); + assert!(error.to_string().contains("unknown overlay intrinsic")); +} + +#[test] +fn candidate_resolution_cannot_change_normal_catalog_bytes() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let before = crate::util::pretty_json(&resolve(&repo_root).unwrap()).unwrap(); + resolve_candidate( + &repo_root, + "thread_idx_x", + "LLVM version candidate", + &"a".repeat(64), + "sm_80", + "+ptx70", + ) + .unwrap(); + let after = crate::util::pretty_json(&resolve(&repo_root).unwrap()).unwrap(); + assert_eq!(before.as_bytes(), after.as_bytes()); +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/tests/scalar_packed.rs b/crates/cuda-intrinsics-gen/src/resolve/tests/scalar_packed.rs new file mode 100644 index 0000000000..213107132c --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/tests/scalar_packed.rs @@ -0,0 +1,1041 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, CatalogFile, ExtendedMinMaxAdmission, ExtendedMinMaxFormat, + ExtendedMinMaxNan, ExtendedMinMaxOperation, ExtendedMinMaxSubnormal, ImportedFile, + ImportedIntrinsic, IntrinsicBackend, IntrinsicSource, OverlayIntrinsic, OverlayShardFile, + PackedAluFormat, PackedAluOperation, PackedConversionDestinationFormat, + PackedConversionRounding, PackedConversionSaturation, PrmtMode, RuntimeValidation, + ScalarArithmeticAdmission, ScalarArithmeticFormat, ScalarArithmeticOperation, + ScalarArithmeticRounding, ScalarArithmeticSaturation, ScalarArithmeticSubnormal, + ScalarConversionAdmission, ScalarConversionRounding, ScalarConversionSaturation, + ScalarMathAdmission, +}; +use crate::util::read_json; +use std::collections::BTreeMap; +use std::path::Path; + +use super::fixtures::*; +use crate::model::{ImportedSelection, ScalarConversionAdmissionVariant}; +use crate::resolve::driver::*; +use crate::resolve::families::*; +use crate::resolve::guards::*; +use crate::resolve::overlay::*; +use crate::resolve::policy::*; + +#[test] +fn compact_prmt_admission_requires_every_mode_and_reserved_abi_id() { + assert_eq!( + expand_prmt_admission(&test_prmt_admission()).unwrap().len(), + 7 + ); + + let mut missing = test_prmt_admission(); + missing.variants.pop(); + assert!(expand_prmt_admission(&missing).is_err()); + + let mut duplicate = test_prmt_admission(); + duplicate.variants[6].mode = PrmtMode::Rc8; + assert!(expand_prmt_admission(&duplicate).is_err()); + + let mut wrong_abi = test_prmt_admission(); + wrong_abi.variants[0].abi_id = "i9999".into(); + assert!(expand_prmt_admission(&wrong_abi).is_err()); +} + +#[test] +fn compact_fp8_conversion_axes_require_the_exact_closed_product() { + let records = expand_packed_conversion_fp8_admission(&test_fp8_conversion_admission()).unwrap(); + assert_eq!(records.len(), 4); + assert_eq!(records[0].id, "cvt_rn_satfinite_e4m3x2_f32"); + assert_eq!(records[1].id, "cvt_rn_satfinite_relu_e4m3x2_f32"); + assert_eq!(records[2].id, "cvt_rn_satfinite_e5m2x2_f32"); + assert_eq!(records[3].id, "cvt_rn_satfinite_relu_e5m2x2_f32"); + assert!(records.iter().all(|record| { + record.rust_result == "u16" + && record.dialect_results == ["i16"] + && record.llvm_results == ["i16"] + && record.minimum_ptx == "8.1" + && record.minimum_sm.as_deref() == Some("sm_89") + && record.pure + && !record.convergent + })); + + let mut missing_format = test_fp8_conversion_admission(); + missing_format.destination_formats.pop(); + assert!(expand_packed_conversion_fp8_admission(&missing_format).is_err()); + + let mut reversed_formats = test_fp8_conversion_admission(); + reversed_formats.destination_formats.reverse(); + assert!(expand_packed_conversion_fp8_admission(&reversed_formats).is_err()); + + let mut missing_saturation = test_fp8_conversion_admission(); + missing_saturation.saturations.pop(); + assert!(expand_packed_conversion_fp8_admission(&missing_saturation).is_err()); + + let mut reversed_saturations = test_fp8_conversion_admission(); + reversed_saturations.saturations.reverse(); + assert!(expand_packed_conversion_fp8_admission(&reversed_saturations).is_err()); + + let mut wrong_count = test_fp8_conversion_admission(); + wrong_count.product_count = 3; + assert!(expand_packed_conversion_fp8_admission(&wrong_count).is_err()); + + let mut executed = test_fp8_conversion_admission(); + executed.runtime_validation = RuntimeValidation::Executed; + assert!(expand_packed_conversion_fp8_admission(&executed).is_err()); +} + +#[test] +fn compact_fp8_f16x2_conversion_axes_require_the_exact_closed_product() { + let records = + expand_packed_conversion_fp8_f16x2_admission(&test_fp8_f16x2_conversion_admission()) + .unwrap(); + assert_eq!(records.len(), 8); + assert_eq!(records[0].id, "cvt_rn_satfinite_e4m3x2_f16x2"); + assert_eq!(records[1].id, "cvt_rn_satfinite_relu_e4m3x2_f16x2"); + assert_eq!(records[2].id, "cvt_rn_f16x2_e4m3x2"); + assert_eq!(records[3].id, "cvt_rn_relu_f16x2_e4m3x2"); + assert_eq!(records[4].id, "cvt_rn_satfinite_e5m2x2_f16x2"); + assert_eq!(records[5].id, "cvt_rn_satfinite_relu_e5m2x2_f16x2"); + assert_eq!(records[6].id, "cvt_rn_f16x2_e5m2x2"); + assert_eq!(records[7].id, "cvt_rn_relu_f16x2_e5m2x2"); + + // Every form takes one packed operand and carries the Ada floor, + // including the unpacks whose destination is the older f16x2. + assert!(records.iter().all(|record| { + record.rust_arguments.len() == 1 + && record.dialect_operands.len() == 1 + && record.llvm_arguments.len() == 1 + && record.expected_ptx.operands.len() == 2 + && record.minimum_ptx == "8.1" + && record.minimum_sm.as_deref() == Some("sm_89") + && record.pure + && !record.convergent + })); + + // Packing narrows f16x2 to a 16-bit FP8 pair; unpacking widens back. + assert!( + records + .iter() + .filter(|record| record.id.ends_with("_f16x2")) + .all(|record| { + record.rust_arguments == ["u32"] + && record.rust_result == "u16" + && record.llvm_arguments == ["v2f16"] + && record.llvm_results == ["i16"] + }) + ); + assert!( + records + .iter() + .filter(|record| !record.id.ends_with("_f16x2")) + .all(|record| { + record.rust_arguments == ["u16"] + && record.rust_result == "u32" + && record.llvm_arguments == ["i16"] + && record.llvm_results == ["v2f16"] + }) + ); + + let mut missing_format = test_fp8_f16x2_conversion_admission(); + missing_format.fp8_formats.pop(); + assert!(expand_packed_conversion_fp8_f16x2_admission(&missing_format).is_err()); + + let mut reversed_formats = test_fp8_f16x2_conversion_admission(); + reversed_formats.fp8_formats.reverse(); + assert!(expand_packed_conversion_fp8_f16x2_admission(&reversed_formats).is_err()); + + let mut missing_direction = test_fp8_f16x2_conversion_admission(); + missing_direction.directions.pop(); + assert!(expand_packed_conversion_fp8_f16x2_admission(&missing_direction).is_err()); + + let mut reversed_directions = test_fp8_f16x2_conversion_admission(); + reversed_directions.directions.reverse(); + assert!(expand_packed_conversion_fp8_f16x2_admission(&reversed_directions).is_err()); + + let mut without_relu = test_fp8_f16x2_conversion_admission(); + without_relu.relu_variants = false; + assert!(expand_packed_conversion_fp8_f16x2_admission(&without_relu).is_err()); + + let mut wrong_count = test_fp8_f16x2_conversion_admission(); + wrong_count.product_count = 4; + assert!(expand_packed_conversion_fp8_f16x2_admission(&wrong_count).is_err()); + + let mut executed = test_fp8_f16x2_conversion_admission(); + executed.runtime_validation = RuntimeValidation::Executed; + assert!(expand_packed_conversion_fp8_f16x2_admission(&executed).is_err()); +} + +#[test] +fn packed_atomic_closed_semantics_reject_every_unreviewed_mutation() { + let valid = packed_policy("packed_atomic_add_f16x2"); + validate_ptx_native_policy(&valid).unwrap(); + + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let overlay = + std::fs::read_to_string(repo_root.join("intrinsics/overlay/packed_atomic.toml")).unwrap(); + for (field, accepted, rejected) in [ + ( + "state space", + "state_space = \"global\"", + "state_space = \"shared\"", + ), + ( + "ordering", + "ordering = \"relaxed\"", + "ordering = \"acquire\"", + ), + ("scope", "scope = \"gpu\"", "scope = \"system\""), + ( + "rounding", + "rounding = \"nearest_even\"", + "rounding = \"toward_zero\"", + ), + ( + "subnormal", + "subnormal = \"preserve\"", + "subnormal = \"flush\"", + ), + ( + "atomicity", + "atomicity = \"per_element\"", + "atomicity = \"coherent_pair\"", + ), + ( + "pointer safety", + "pointer_contract = \"mutable_global_u32_aligned4\"", + "pointer_contract = \"unaligned\"", + ), + ( + "mixed access safety", + "access_contract = \"no_mixed_whole_word_or_non_atomic_access\"", + "access_contract = \"mixed_access_allowed\"", + ), + ( + "scope safety", + "scope_contract = \"racing_atomics_mutually_inclusive\"", + "scope_contract = \"scope_unchecked\"", + ), + ( + "codegen", + "codegen_contract = \"exact_native_instruction\"", + "codegen_contract = \"semantic_equivalence\"", + ), + ] { + let mutated = overlay.replacen(accepted, rejected, 1); + let error = toml::from_str::(&mutated).unwrap_err(); + assert!( + error + .to_string() + .contains(rejected.split(" = ").next().unwrap()), + "{field} mutation did not fail closed: {error}" + ); + } + + let mut safe = valid; + safe.safe = true; + safe.safe_allowlist_reason = Some("incorrectly claims no caller obligations".into()); + assert!( + validate_ptx_native_policy(&safe) + .unwrap_err() + .to_string() + .contains("unsafe must-use packed atomic") + ); +} + +#[test] +fn packed_alu_recipes_accept_only_the_reviewed_source_shape_and_floor() { + let operations = [ + PackedAluOperation::Add, + PackedAluOperation::Sub, + PackedAluOperation::Mul, + PackedAluOperation::Fma, + PackedAluOperation::FmaRelu, + PackedAluOperation::Min, + PackedAluOperation::Max, + PackedAluOperation::Neg, + PackedAluOperation::Abs, + ]; + for format in [PackedAluFormat::Bf16x2, PackedAluFormat::F16x2] { + for operation in operations { + let policy = packed_alu_policy(format, operation); + match packed_alu_declaration(format, operation) { + Some(declaration) => validate_imported_policy(&policy, &declaration).unwrap(), + None => validate_ptx_native_policy(&policy).unwrap(), + } + } + } + + let declaration = + packed_alu_declaration(PackedAluFormat::Bf16x2, PackedAluOperation::Fma).unwrap(); + let reject_imported = |policy: &OverlayIntrinsic, message: &str| { + let error = validate_imported_policy(policy, &declaration).unwrap_err(); + assert!(error.to_string().contains(message), "{error:#}"); + }; + + let valid = packed_alu_policy(PackedAluFormat::Bf16x2, PackedAluOperation::Fma); + let mut wrong_source = valid.clone(); + wrong_source.source_record = Some("int_nvvm_fma_rn_bf16".into()); + reject_imported(&wrong_source, "source"); + + let mut wrong_format = valid.clone(); + wrong_format.packed_alu.as_mut().unwrap().format = PackedAluFormat::F16x2; + reject_imported(&wrong_format, "identity"); + + let mut wrong_operation = valid.clone(); + wrong_operation.packed_alu.as_mut().unwrap().operation = PackedAluOperation::Max; + reject_imported(&wrong_operation, "identity"); + + let mut wrong_floor = valid.clone(); + wrong_floor.minimum_sm = Some("sm_90".into()); + reject_imported(&wrong_floor, "target floor"); + + let mut wrong_effects = valid.clone(); + wrong_effects.memory = "read".into(); + reject_imported(&wrong_effects, "effects"); + + let mut wrong_section = valid.clone(); + wrong_section.ptx_isa_section = "9.7.4 Floating Point Instructions".into(); + reject_imported(&wrong_section, "PTX provenance"); + + let mut wrong_url = valid.clone(); + wrong_url.ptx_isa_url = + "https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions" + .into(); + reject_imported(&wrong_url, "PTX provenance"); + + let mut wrong_adapter = valid.clone(); + wrong_adapter.lowering = "direct_nvvm".into(); + reject_imported(&wrong_adapter, "lowering recipe"); + + let mut wrong_backend = valid; + wrong_backend.backend_lowerings[0].mechanism = BackendLoweringMechanism::TypedNvvm; + reject_imported(&wrong_backend, "inline-PTX routes"); + + let mut wrong_native = packed_alu_policy(PackedAluFormat::Bf16x2, PackedAluOperation::Add); + wrong_native.source = Some(IntrinsicSource::PtxNative { + instruction: "add.bf16x2".into(), + }); + let error = validate_ptx_native_policy(&wrong_native).unwrap_err(); + assert!(error.to_string().contains("PTX-native recipe")); + + let mut invented_llvm = packed_alu_policy(PackedAluFormat::F16x2, PackedAluOperation::Add); + invented_llvm.llvm_symbol = Some("llvm.nvvm.add.rn.f16x2".into()); + let error = validate_ptx_native_policy(&invented_llvm).unwrap_err(); + assert!( + error + .to_string() + .contains("must not invent LLVM source facts") + ); + + let mut unreviewed_modifier = + packed_alu_policy(PackedAluFormat::F16x2, PackedAluOperation::Add); + unreviewed_modifier.expected_ptx.modifiers = vec!["rn".into(), "ftz".into(), "f16x2".into()]; + let error = validate_ptx_native_policy(&unreviewed_modifier).unwrap_err(); + assert!(error.to_string().contains("exact packed-ALU instruction")); + + let mut wrong_arity = packed_alu_policy(PackedAluFormat::F16x2, PackedAluOperation::Add); + wrong_arity.expected_ptx.operands.pop(); + let error = validate_ptx_native_policy(&wrong_arity).unwrap_err(); + assert!(error.to_string().contains("exact packed-ALU instruction")); + + let f16_declaration = + packed_alu_declaration(PackedAluFormat::F16x2, PackedAluOperation::Fma).unwrap(); + let reject_f16 = |policy: &OverlayIntrinsic, message: &str| { + let error = validate_imported_policy(policy, &f16_declaration).unwrap_err(); + assert!(error.to_string().contains(message), "{error:#}"); + }; + let f16 = packed_alu_policy(PackedAluFormat::F16x2, PackedAluOperation::Fma); + + let mut wrong_signature = f16.clone(); + wrong_signature.llvm_arguments = vec!["v2bf16".into(); 3]; + reject_f16(&wrong_signature, "LLVM argument signature mismatch"); + + let mut missing_must_use = f16.clone(); + missing_must_use.must_use = false; + reject_f16(&missing_must_use, "reviewed safe packed-ALU API"); + + let mut wrong_native_floor = f16.clone(); + wrong_native_floor + .packed_alu + .as_mut() + .unwrap() + .native_minimum_sm = 70; + reject_f16(&wrong_native_floor, "target floor"); + + let mut wrong_backend_floor = f16; + wrong_backend_floor.backend_lowerings[0].minimum_ptx = Some("4.2".into()); + reject_f16(&wrong_backend_floor, "exact packed-ALU floor"); + + let abs_declaration = + packed_alu_declaration(PackedAluFormat::F16x2, PackedAluOperation::Abs).unwrap(); + let mut wrong_abs = packed_alu_policy(PackedAluFormat::F16x2, PackedAluOperation::Abs); + wrong_abs.resolved_llvm_symbol = Some("llvm.nvvm.fabs.v2bf16".into()); + let error = validate_imported_policy(&wrong_abs, &abs_declaration).unwrap_err(); + assert!(error.to_string().contains("LLVM source or signature")); +} + +#[test] +fn pinned_packed_alu_records_match_the_closed_recipes() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let (overlay, _) = + read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")).unwrap(); + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let declarations: BTreeMap<_, _> = imported + .intrinsics + .iter() + .map(|record| (record.source_record.as_str(), record)) + .collect(); + let packed: Vec<_> = overlay + .intrinsics + .iter() + .filter(|record| record.family == "packed_alu") + .collect(); + assert_eq!(packed.len(), 30); + for policy in packed { + let source = resolve_policy_source(policy).unwrap(); + let declaration = policy + .source_record + .as_deref() + .map(|record| declarations[record]); + validate_policy(policy, &source, declaration, 1).unwrap(); + } +} + +#[test] +fn pinned_packed_conversion_records_match_the_closed_recipes() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let overlay = load_resolution_base(&repo_root).unwrap().overlay; + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let declarations: BTreeMap<_, _> = imported + .intrinsics + .iter() + .map(|record| (record.source_record.as_str(), record)) + .collect(); + let packed: Vec<_> = overlay + .intrinsics + .iter() + .filter(|record| record.family == "packed_conversion") + .collect(); + assert_eq!(packed.len(), 18); + for policy in packed { + let source = resolve_policy_source(policy).unwrap(); + let declaration = policy + .source_record + .as_deref() + .map(|record| declarations[record]); + validate_policy(policy, &source, declaration, 1).unwrap(); + } +} + +#[test] +fn selectionless_packed_conversion_is_admitted_only_by_its_closed_recipe() { + let cases = [ + ( + PackedConversionDestinationFormat::Bf16x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::None, + ), + ( + PackedConversionDestinationFormat::F16x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::None, + ), + ( + PackedConversionDestinationFormat::F16x2, + PackedConversionRounding::TowardZero, + PackedConversionSaturation::None, + ), + ( + PackedConversionDestinationFormat::F16x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::Relu, + ), + ( + PackedConversionDestinationFormat::Bf16x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::Relu, + ), + ( + PackedConversionDestinationFormat::Bf16x2, + PackedConversionRounding::TowardZero, + PackedConversionSaturation::None, + ), + ( + PackedConversionDestinationFormat::E4m3x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::Satfinite, + ), + ( + PackedConversionDestinationFormat::E4m3x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::SatfiniteRelu, + ), + ( + PackedConversionDestinationFormat::E5m2x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::Satfinite, + ), + ( + PackedConversionDestinationFormat::E5m2x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::SatfiniteRelu, + ), + ]; + for (destination, rounding, saturation) in cases { + let policy = packed_conversion_policy(destination, rounding, saturation); + let declaration = packed_conversion_declaration(&policy); + validate_imported_policy(&policy, &declaration).unwrap(); + } + + let valid = packed_conversion_policy( + PackedConversionDestinationFormat::Bf16x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::None, + ); + let declaration = packed_conversion_declaration(&valid); + + let reject = |policy: &OverlayIntrinsic, declaration: &ImportedIntrinsic, message: &str| { + let error = validate_imported_policy(policy, declaration).unwrap_err(); + assert!(error.to_string().contains(message), "{error:#}"); + }; + + let mut wrong_source = valid.clone(); + wrong_source.source_record = Some("int_nvvm_ff2bf16x2_rz".into()); + reject(&wrong_source, &declaration, "identity or LLVM source"); + + let mut wrong_floor = valid.clone(); + wrong_floor.minimum_ptx = "7.8".into(); + reject(&wrong_floor, &declaration, "target floor"); + + let mut wrong_effect = valid.clone(); + wrong_effect.pure = false; + reject(&wrong_effect, &declaration, "effects"); + + let mut wrong_section = valid.clone(); + wrong_section.ptx_isa_section = "9.7.9 Data Movement and Conversion Instructions".into(); + reject(&wrong_section, &declaration, "PTX provenance"); + + let mut wrong_url = valid.clone(); + wrong_url.ptx_isa_url = + "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions" + .into(); + reject(&wrong_url, &declaration, "PTX provenance"); + + let mut wrong_shape = valid.clone(); + wrong_shape.expected_ptx.modifiers.swap(1, 2); + reject(&wrong_shape, &declaration, "expected PTX does not match"); + + let mut wrong_identity = valid.clone(); + wrong_identity.id = "cvt_f16x2_f32".into(); + reject(&wrong_identity, &declaration, "identity or LLVM source"); + + let mut unsupported = valid.clone(); + let conversion = unsupported.packed_conversion.as_mut().unwrap(); + conversion.rounding = PackedConversionRounding::TowardZero; + conversion.saturation = PackedConversionSaturation::Relu; + reject( + &unsupported, + &declaration, + "unsupported packed-conversion source, destination", + ); + + let mut wrong_compatibility = valid.clone(); + wrong_compatibility.compatibility_rust_paths = + vec!["cuda_device::tcgen05::cvt_f32x2_bf16x2".into()]; + reject(&wrong_compatibility, &declaration, "conversion API"); + + let mut wrong_result = valid.clone(); + wrong_result.llvm_results = vec!["v2f16".into()]; + reject(&wrong_result, &declaration, "result signature mismatch"); + + let mut selected = declaration.clone(); + selected.selections.push(ImportedSelection { + source_record: "UNREVIEWED".into(), + asm: "cvt.rn.bf16x2.f32 $d, $a, $b;".into(), + predicates: vec![], + constraints: Default::default(), + }); + reject(&valid, &selected, "selectionless"); +} + +#[test] +fn scalar_conversion_admission_is_closed_and_backend_specific() { + let variants = [ + ( + "i0368", + ScalarConversionRounding::NearestAway, + ScalarConversionSaturation::None, + ), + ( + "i0369", + ScalarConversionRounding::NearestAway, + ScalarConversionSaturation::Satfinite, + ), + ( + "i0370", + ScalarConversionRounding::NearestEven, + ScalarConversionSaturation::None, + ), + ( + "i0371", + ScalarConversionRounding::NearestEven, + ScalarConversionSaturation::Relu, + ), + ( + "i0372", + ScalarConversionRounding::NearestEven, + ScalarConversionSaturation::Satfinite, + ), + ( + "i0373", + ScalarConversionRounding::NearestEven, + ScalarConversionSaturation::ReluSatfinite, + ), + ( + "i0374", + ScalarConversionRounding::TowardZero, + ScalarConversionSaturation::None, + ), + ( + "i0375", + ScalarConversionRounding::TowardZero, + ScalarConversionSaturation::Relu, + ), + ( + "i0376", + ScalarConversionRounding::TowardZero, + ScalarConversionSaturation::Satfinite, + ), + ( + "i0377", + ScalarConversionRounding::TowardZero, + ScalarConversionSaturation::ReluSatfinite, + ), + ]; + let admission = ScalarConversionAdmission { + llvm_evidence_profile: "llvm-scalar".into(), + libnvvm_evidence_profile: "libnvvm-scalar".into(), + runtime_validation: RuntimeValidation::Unexecuted, + variants: variants + .iter() + .map( + |(abi_id, rounding, saturation)| ScalarConversionAdmissionVariant { + abi_id: (*abi_id).into(), + rounding: *rounding, + saturation: *saturation, + }, + ) + .collect(), + }; + + let records = expand_scalar_conversion_admission(&admission).unwrap(); + assert_eq!(records.len(), 10); + for (record, (abi_id, rounding, saturation)) in records.iter().zip(variants) { + assert_eq!(record.abi_id, abi_id); + assert_eq!(record.rust_arguments, ["f32"]); + assert_eq!(record.rust_result, "u32"); + assert_eq!( + record.compatibility_rust_paths, + [format!("cuda_device::convert::{}", record.rust_name)] + ); + let conversion = record.scalar_conversion.as_ref().unwrap(); + assert_eq!(conversion.rounding, rounding); + assert_eq!(conversion.saturation, saturation); + assert!(record.backend_lowerings.iter().any(|lowering| { + lowering.backend == IntrinsicBackend::LlvmNvptx + && lowering.mechanism == BackendLoweringMechanism::TypedNvvm + })); + assert!(record.backend_lowerings.iter().any(|lowering| { + lowering.backend == IntrinsicBackend::LibNvvm + && lowering.mechanism == BackendLoweringMechanism::InlinePtx + })); + } + + let mut reordered = admission.clone(); + reordered.variants.swap(0, 1); + assert!( + expand_scalar_conversion_admission(&reordered) + .unwrap_err() + .to_string() + .contains("canonical ten variants") + ); +} + +#[test] +fn scalar_arithmetic_admission_is_closed_and_selects_only_direct_ptx() { + let variants = canonical_scalar_arithmetic_variants(); + let admission = ScalarArithmeticAdmission { + llvm_evidence_profile: "llvm-scalar-arithmetic".into(), + libnvvm_evidence_profile: "libnvvm-scalar-arithmetic".into(), + runtime_validation: RuntimeValidation::Unexecuted, + variants: variants + .iter() + .copied() + .enumerate() + .map( + |(index, variant)| crate::model::ScalarArithmeticAdmissionVariant { + abi_id: format!("i{:04}", 390 + index), + format: variant.0, + operation: variant.1, + rounding: variant.2, + subnormal: variant.3, + saturation: variant.4, + }, + ) + .collect(), + }; + let records = expand_scalar_arithmetic_admission(&admission).unwrap(); + assert_eq!(records.len(), 64); + assert_eq!(records.first().unwrap().id, "mul_rn_f64"); + assert_eq!(records.first().unwrap().abi_id, "i0390"); + assert_eq!(records.last().unwrap().id, "add_rp_ftz_sat_f32"); + assert_eq!(records.last().unwrap().abi_id, "i0453"); + validate_unique_overlay(&records, 1).unwrap(); + let llvm_inline_count = records + .iter() + .filter(|record| { + record.backend_lowerings.iter().any(|lowering| { + lowering.backend == IntrinsicBackend::LlvmNvptx + && lowering.mechanism == BackendLoweringMechanism::InlinePtx + }) + }) + .count(); + assert_eq!(llvm_inline_count, 16); + assert!(records.iter().all(|record| { + let expected = if record + .scalar_arithmetic + .as_ref() + .is_some_and(|arithmetic| arithmetic.saturation == ScalarArithmeticSaturation::Sat) + { + BackendLoweringMechanism::InlinePtx + } else { + BackendLoweringMechanism::TypedNvvm + }; + record.backend_lowerings.iter().any(|lowering| { + lowering.backend == IntrinsicBackend::LlvmNvptx && lowering.mechanism == expected + }) + })); + + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let declarations = imported + .intrinsics + .into_iter() + .map(|record| (record.source_record.clone(), record)) + .collect::>(); + let mut mixed_count = 0; + let mut add_two_selection_count = 0; + let mut add_six_selection_count = 0; + for record in &records { + let declaration = &declarations[record.source_record.as_ref().unwrap()]; + validate_imported_policy(record, declaration).unwrap(); + let selected = declaration + .selections + .iter() + .filter(|selection| selection_matches_policy(record, selection).unwrap()) + .collect::>(); + assert_eq!(selected.len(), 1, "{}", record.id); + if declaration.selections.len() == 3 { + mixed_count += 1; + } + if record + .scalar_arithmetic + .as_ref() + .is_some_and(|contract| contract.operation == ScalarArithmeticOperation::Add) + { + match declaration.selections.len() { + 2 => add_two_selection_count += 1, + 6 => add_six_selection_count += 1, + count => panic!("unexpected add selection count {count}"), + } + } + } + assert_eq!(mixed_count, 8); + assert_eq!(add_two_selection_count, 12); + assert_eq!(add_six_selection_count, 8); + let add_ftz_sat = records + .iter() + .find(|record| record.id == "add_rn_ftz_sat_f32") + .unwrap(); + assert_eq!( + add_ftz_sat.expected_ptx.modifiers, + ["rn", "sat", "ftz", "f32"] + ); + + let mut non_contiguous_abi = admission.clone(); + non_contiguous_abi.variants[0].abi_id = "i9999".into(); + let non_contiguous_records = expand_scalar_arithmetic_admission(&non_contiguous_abi).unwrap(); + assert_eq!(non_contiguous_records[0].abi_id, "i9999"); + + let mut reordered = admission.clone(); + reordered.variants.swap(0, 1); + assert!( + expand_scalar_arithmetic_admission(&reordered) + .unwrap_err() + .to_string() + .contains("canonical 64 variants") + ); + assert!( + scalar_arithmetic_recipe(( + ScalarArithmeticFormat::F64, + ScalarArithmeticOperation::Fma, + ScalarArithmeticRounding::Rn, + ScalarArithmeticSubnormal::Ftz, + ScalarArithmeticSaturation::None, + )) + .is_none() + ); +} + +#[test] +fn scalar_math_admission_keeps_semantic_order_but_accepts_non_contiguous_abi_ids() { + let variants = canonical_scalar_math_variants(); + let admission = ScalarMathAdmission { + llvm_evidence_profile: "llvm-scalar-math".into(), + libnvvm_evidence_profile: "libnvvm-scalar-math".into(), + runtime_validation: RuntimeValidation::Unexecuted, + variants: variants + .iter() + .copied() + .enumerate() + .map( + |(index, variant)| crate::model::ScalarMathAdmissionVariant { + abi_id: format!("i{:04}", 782 + index), + libnvvm_evidence_profile: (index == 40).then(|| "libnvvm-ex2-f16".into()), + format: variant.0, + operation: variant.1, + precision: variant.2, + subnormal: variant.3, + }, + ) + .collect(), + }; + + let records = expand_scalar_math_admission(&admission).unwrap(); + assert_eq!(records.len(), 41); + assert_eq!(records.first().unwrap().abi_id, "i0782"); + assert_eq!(records.last().unwrap().abi_id, "i0822"); + assert!( + records + .iter() + .all(|record| record.backend_lowerings.len() == 2) + ); + assert_eq!( + records.last().unwrap().backend_lowerings[1].evidence_profile, + "libnvvm-ex2-f16" + ); + assert!( + records[..40].iter().all(|record| { + record.backend_lowerings[1].evidence_profile == "libnvvm-scalar-math" + }) + ); + + let mut non_contiguous_abi = admission.clone(); + non_contiguous_abi.variants[40].abi_id = "i9999".into(); + let records = expand_scalar_math_admission(&non_contiguous_abi).unwrap(); + assert_eq!(records.last().unwrap().abi_id, "i9999"); + + let mut reordered = admission; + reordered.variants.swap(0, 1); + assert!( + expand_scalar_math_admission(&reordered) + .unwrap_err() + .to_string() + .contains("canonical 41 variants") + ); +} + +#[test] +fn checked_in_scalar_math_overlay_pins_the_current_libnvvm_override_set() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let (overlay, _) = + read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")).unwrap(); + let scalar_math = overlay + .intrinsics + .iter() + .filter(|record| record.family == "scalar_math") + .collect::>(); + assert_eq!(scalar_math.len(), 41); + assert!( + scalar_math + .iter() + .all(|record| record.backend_lowerings.len() == 2) + ); + + let family_profile = "cuda-13.3-libnvvm-13.3.33-scalar-math"; + let overrides = scalar_math + .iter() + .filter_map(|record| { + let libnvvm = record + .backend_lowerings + .iter() + .find(|route| route.backend == IntrinsicBackend::LibNvvm) + .unwrap(); + (libnvvm.evidence_profile != family_profile).then_some(record.id.as_str()) + }) + .collect::>(); + assert_eq!(overrides, ["ex2_approx_f16"]); +} + +#[test] +fn extended_minmax_admission_is_exact_and_fail_closed() { + let variants = canonical_extended_minmax_variants(); + let admission = ExtendedMinMaxAdmission { + llvm_evidence_profile: "llvm-extended-minmax".into(), + libnvvm_evidence_profile: "libnvvm-extended-minmax".into(), + runtime_validation: RuntimeValidation::Unexecuted, + variants: variants + .iter() + .copied() + .enumerate() + .map( + |(index, variant)| crate::model::ExtendedMinMaxAdmissionVariant { + abi_id: if index < 28 { + format!("i{:04}", 550 + index) + } else { + format!("i{:04}", 830 + index - 28) + }, + format: variant.0, + operation: variant.1, + subnormal: variant.2, + nan: variant.3, + xorsign_abs: variant.4, + }, + ) + .collect(), + }; + let records = expand_extended_minmax_admission(&admission).unwrap(); + assert_eq!(records.len(), 52); + assert_eq!(records.first().unwrap().id, "min_ftz_f16x2"); + assert_eq!(records.first().unwrap().abi_id, "i0550"); + // The first reserved block ends here; the scalar 16-bit forms admitted + // afterwards continue at a second base rather than at i0578. + assert_eq!(records[27].id, "max_xorsign_abs_f16x2"); + assert_eq!(records[27].abi_id, "i0577"); + assert_eq!(records[28].id, "min_f16"); + assert_eq!(records[28].abi_id, "i0830"); + assert_eq!(records.last().unwrap().id, "max_nan_xorsign_abs_bf16"); + assert_eq!(records.last().unwrap().abi_id, "i0853"); + validate_unique_overlay(&records, 1).unwrap(); + assert_eq!( + records + .iter() + .filter(|record| record.minimum_sm.as_deref() == Some("sm_80")) + .count(), + 20 + ); + assert_eq!( + records + .iter() + .filter(|record| record.minimum_sm.as_deref() == Some("sm_86")) + .count(), + 32 + ); + assert!(records.iter().all(|record| { + record.backend_lowerings.len() == 2 + && record + .backend_lowerings + .iter() + .all(|lowering| lowering.mechanism == BackendLoweringMechanism::InlinePtx) + })); + + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let declarations = imported + .intrinsics + .into_iter() + .map(|record| (record.source_record.clone(), record)) + .collect::>(); + for record in &records { + let declaration = &declarations[record.source_record.as_ref().unwrap()]; + assert_eq!(declaration.selections.len(), 1, "{}", record.id); + validate_imported_policy(record, declaration).unwrap(); + assert_eq!( + declaration + .selections + .iter() + .filter(|selection| selection_matches_policy(record, selection).unwrap()) + .count(), + 1, + "{}", + record.id + ); + } + let packed_max = records + .iter() + .find(|record| record.id == "max_nan_bf16x2") + .unwrap(); + assert_eq!( + declarations[packed_max.source_record.as_ref().unwrap()].selections[0].source_record, + "INT_NVVM_FMAN_NaN_bf16x2" + ); + assert_eq!( + declarations["int_nvvm_fmax_ftz_nan_xorsign_abs_f"].selections[0].source_record, + "INT_NVVM_FMAX_FTZ_NAN_XORSIGN_ABS_F" + ); + assert!( + declarations["int_nvvm_fmax_ftz_nan_xorsign_abs_f"].selections[0] + .asm + .contains(".NaN.") + ); + for excluded in [ + "int_nvvm_fmin_ftz_bf16x2", + "int_nvvm_fmax_ftz_nan_bf16x2", + "int_nvvm_fmin_f", + "int_nvvm_fmax_nan_f", + ] { + assert!(declarations[excluded].selections.is_empty(), "{excluded}"); + assert!( + records + .iter() + .all(|record| record.source_record.as_deref() != Some(excluded)), + "{excluded} was admitted without an instruction selection" + ); + } + + let catalog: CatalogFile = read_json(&repo_root.join("intrinsics/catalog.json")).unwrap(); + for id in ["min_f16x2", "max_f16x2", "min_bf16x2", "max_bf16x2"] { + assert_eq!( + catalog + .intrinsics + .iter() + .find(|record| record.id == id) + .unwrap() + .family, + "packed_alu" + ); + assert!(records.iter().all(|record| record.id != id)); + } + + let mut non_contiguous_abi = admission.clone(); + non_contiguous_abi.variants[0].abi_id = "i9999".into(); + let non_contiguous_records = expand_extended_minmax_admission(&non_contiguous_abi).unwrap(); + assert_eq!(non_contiguous_records[0].abi_id, "i9999"); + + let mut reordered = admission.clone(); + reordered.variants.swap(0, 1); + assert!( + expand_extended_minmax_admission(&reordered) + .unwrap_err() + .to_string() + .contains("exact canonical 28 variants") + ); + assert!( + extended_minmax_recipe(( + ExtendedMinMaxFormat::Bf16x2, + ExtendedMinMaxOperation::Min, + ExtendedMinMaxSubnormal::Ftz, + ExtendedMinMaxNan::Nan, + false, + )) + .is_none() + ); + let mut changed = declarations["int_nvvm_fmin_ftz_f16x2"].clone(); + changed.selections[0].predicates.swap(0, 1); + assert!(validate_imported_policy(&records[0], &changed).is_err()); +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/tests/sreg_threadfence.rs b/crates/cuda-intrinsics-gen/src/resolve/tests/sreg_threadfence.rs new file mode 100644 index 0000000000..b5ae466941 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/tests/sreg_threadfence.rs @@ -0,0 +1,486 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, EvidenceStageKind, ImportedFile, IntrinsicBackend, IntrinsicSource, + OverlayShardFile, RuntimeValidation, SpecialRegisterOutputConstraint, +}; +use crate::util::read_json; +use std::collections::{BTreeMap, BTreeSet}; +use std::path::Path; + +use super::fixtures::*; +use crate::resolve::driver::*; +use crate::resolve::evidence::*; +use crate::resolve::families::*; +use crate::resolve::guards::*; +use crate::resolve::overlay::*; +use crate::resolve::policy::*; + +#[test] +fn cluster_sreg_admission_uses_its_fixed_introduction_schema() { + let shard = |schema| { + toml::from_str::(&format!( + r#" +schema = {schema} +family = "sreg" + +[cluster_sreg] +axes = ["x", "y", "z"] +xyz_product_count = 12 +record_count = 14 +"# + )) + .unwrap() + }; + let path = Path::new("intrinsics/overlay/sreg_cluster.toml"); + + let old = shard(CLUSTER_SREG_SHARD_SCHEMA - 1); + assert!( + validate_overlay_shard_schema(&old, path) + .unwrap_err() + .to_string() + .contains("requires overlay shard schema 30") + ); + + let current = shard(CLUSTER_SREG_SHARD_SCHEMA); + validate_overlay_shard_schema(¤t, path).unwrap(); + let admission = current.cluster_sreg.unwrap(); + assert_eq!(expand_cluster_sreg_admission(&admission).unwrap().len(), 14); + + let mut wrong_axes = admission.clone(); + wrong_axes.axes.swap(0, 1); + assert!(expand_cluster_sreg_admission(&wrong_axes).is_err()); + + let mut wrong_count = admission; + wrong_count.record_count = 13; + assert!(expand_cluster_sreg_admission(&wrong_count).is_err()); +} + +#[test] +fn special_register_admission_is_closed_and_schema_gated() { + let admission = test_special_register_admission(); + let records = expand_special_register_admission(&admission).unwrap(); + assert_eq!(records.len(), 12); + assert_eq!( + records + .iter() + .map(|record| record.id.as_str()) + .collect::>(), + [ + "clock", + "clock64", + "globaltimer", + "envreg1", + "envreg2", + "smid", + "nsmid", + "gridid", + "warpid", + "nwarpid", + "dynamic_smem_size", + "total_smem_size", + ] + ); + + let mut reordered = admission.clone(); + reordered.registers.swap(0, 1); + assert!(expand_special_register_admission(&reordered).is_err()); + + let mut wrong_count = admission.clone(); + wrong_count.product_count -= 1; + assert!(expand_special_register_admission(&wrong_count).is_err()); + + let mut executed = admission.clone(); + executed.runtime_validation = RuntimeValidation::Executed; + assert!(expand_special_register_admission(&executed).is_err()); + + let shard = |schema| OverlayShardFile { + schema, + family: "sreg".into(), + intrinsics: vec![], + register_mma_int4: None, + register_mma_int8: None, + register_mma_b1: None, + register_mma_f8f6f4_f32: None, + register_mma_f8f6f4_f16: None, + register_mma_mxf8f6f4_f32: None, + register_mma_fp8: None, + register_mma_ampere_float: None, + sparse_mma_integer: None, + sparse_mma_f8f6f4_f32: None, + sparse_mma_f8f6f4_f16: None, + prmt: None, + packed_conversion_fp8: None, + packed_conversion_fp8_f16x2: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cluster_sreg: None, + cluster_barrier: None, + mbarrier_extended: None, + special_registers: Some(admission.clone()), + debug_control: None, + threadfence: None, + cluster_memory: None, + stmatrix: None, + clc: None, + wgmma_controls: None, + tma: None, + tcgen05: None, + }; + let path = Path::new("intrinsics/overlay/sreg_special.toml"); + validate_overlay_shard_schema(&shard(SPECIAL_REGISTER_SHARD_SCHEMA), path).unwrap(); + assert!( + validate_overlay_shard_schema(&shard(SPECIAL_REGISTER_SHARD_SCHEMA - 1), path) + .unwrap_err() + .to_string() + .contains("requires overlay shard schema 32") + ); +} + +#[test] +fn pinned_special_registers_preserve_apis_widths_and_backend_routes() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let base = load_resolution_base(&repo_root).unwrap(); + let declarations = index_imported_intrinsics(&base.imported).unwrap(); + let records = base + .overlay + .intrinsics + .iter() + .filter(|record| record.special_register.is_some()) + .collect::>(); + assert_eq!(records.len(), 12); + assert_eq!( + records + .iter() + .map(|record| record.abi_id.clone()) + .collect::>(), + (283..=294) + .map(|id| format!("i{id:04}")) + .collect::>() + ); + + for record in &records { + let source = resolve_policy_source(record).unwrap(); + let declaration = resolve_imported_declaration(record, &source, &declarations).unwrap(); + validate_special_register_policy(record, &source, declaration).unwrap(); + validate_special_register_llvm_exclusion(record, &declarations).unwrap(); + } + + for (id, section, anchor) in [ + ("clock", "10.23", "special-registers-clock-clock-hi"), + ("clock64", "10.24", "special-registers-clock64"), + ("globaltimer", "10.28", "special-registers-globaltimer"), + ("envreg1", "10.27", "special-registers-envreg"), + ("envreg2", "10.27", "special-registers-envreg"), + ("smid", "10.8", "special-registers-smid"), + ("nsmid", "10.9", "special-registers-nsmid"), + ("gridid", "10.10", "special-registers-gridid"), + ("warpid", "10.4", "special-registers-warpid"), + ("nwarpid", "10.5", "special-registers-nwarpid"), + ( + "dynamic_smem_size", + "10.32", + "special-registers-dynamic-smem-size", + ), + ( + "total_smem_size", + "10.30", + "special-registers-total-smem-size", + ), + ] { + let record = records.iter().find(|record| record.id == id).unwrap(); + assert!(record.ptx_isa_section.starts_with(section)); + assert!(record.ptx_isa_url.ends_with(anchor)); + } + + let gridid = records.iter().find(|record| record.id == "gridid").unwrap(); + assert_eq!(gridid.rust_result, "u64"); + assert_eq!(gridid.dialect_results, ["i64"]); + assert!(matches!( + resolve_policy_source(gridid).unwrap(), + IntrinsicSource::PtxNative { .. } + )); + assert!( + gridid + .special_register + .as_ref() + .unwrap() + .llvm_exclusion + .is_some() + ); + + let clock = records.iter().find(|record| record.id == "clock").unwrap(); + let source = resolve_policy_source(clock).unwrap(); + let declaration = resolve_imported_declaration(clock, &source, &declarations) + .unwrap() + .unwrap(); + + let mut wrong_effects = (*clock).clone(); + wrong_effects.memory = "none".into(); + assert!(validate_special_register_policy(&wrong_effects, &source, Some(declaration)).is_err()); + + let mut wrong_contract = (*clock).clone(); + wrong_contract + .special_register + .as_mut() + .unwrap() + .output_constraint = SpecialRegisterOutputConstraint::Register64; + assert!(validate_special_register_policy(&wrong_contract, &source, Some(declaration)).is_err()); + + let mut wrong_route = (*clock).clone(); + wrong_route.backend_lowerings[0].mechanism = BackendLoweringMechanism::InlinePtx; + assert!(validate_special_register_policy(&wrong_route, &source, Some(declaration)).is_err()); +} + +#[test] +fn special_register_evidence_validates_both_backend_routes() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let policies = expand_special_register_admission(&test_special_register_admission()).unwrap(); + let mut evidence_files = vec![ + read_evidence_file( + &repo_root.join("intrinsics/evidence/rust-llvm-22.1.2-1cb4e383-special-registers.json"), + ) + .unwrap(), + read_evidence_file( + &repo_root.join("intrinsics/evidence/cuda-13.3-libnvvm-13.3.33-special-registers.json"), + ) + .unwrap(), + ]; + let llvm_revision = "1cb4e3833c1919c2e6fb579a23ac0e2b22587b7e"; + let indexed = index_evidence(&evidence_files, llvm_revision).unwrap(); + for policy in &policies { + for lowering in &policy.backend_lowerings { + let evidence = indexed + .get(&(lowering.evidence_profile.as_str(), policy.id.as_str())) + .unwrap(); + validate_evidence(policy, evidence, Some(lowering)).unwrap(); + } + } + + let libnvvm = evidence_files + .iter_mut() + .find(|file| file.backend_kind == Some(IntrinsicBackend::LibNvvm)) + .unwrap(); + libnvvm + .records + .iter_mut() + .find(|record| record.id == "gridid") + .unwrap() + .stages + .retain(|stage| stage.stage != EvidenceStageKind::DeviceLink); + let indexed = index_evidence(&evidence_files, llvm_revision).unwrap(); + let policy = policies + .iter() + .find(|policy| policy.id == "gridid") + .unwrap(); + let lowering = policy + .backend_lowerings + .iter() + .find(|lowering| lowering.backend == IntrinsicBackend::LibNvvm) + .unwrap(); + let evidence = indexed + .get(&(lowering.evidence_profile.as_str(), policy.id.as_str())) + .unwrap(); + assert!(validate_evidence(policy, evidence, Some(lowering)).is_err()); +} + +#[test] +fn threadfence_admission_is_closed_and_uses_schema_34() { + let shard = |schema| { + toml::from_str::(&format!( + r#" +schema = {schema} +family = "sync" + +[threadfence] +llvm_evidence_profile = "llvm-test" +libnvvm_evidence_profile = "libnvvm-test" +runtime_validation = "unexecuted" + +[[threadfence.variant]] +abi_id = "i0298" +scope = "cta" + +[[threadfence.variant]] +abi_id = "i0299" +scope = "device" + +[[threadfence.variant]] +abi_id = "i0300" +scope = "system" +"# + )) + .unwrap() + }; + let path = Path::new("intrinsics/overlay/threadfence.toml"); + + let old = shard(THREADFENCE_SHARD_SCHEMA - 1); + assert!( + validate_overlay_shard_schema(&old, path) + .unwrap_err() + .to_string() + .contains("requires overlay shard schema 34") + ); + + let current = shard(THREADFENCE_SHARD_SCHEMA); + validate_overlay_shard_schema(¤t, path).unwrap(); + let admission = current.threadfence.unwrap(); + let records = expand_threadfence_admission(&admission).unwrap(); + assert_eq!(records.len(), 3); + assert_eq!( + records + .iter() + .map(|record| (record.abi_id.as_str(), record.id.as_str())) + .collect::>(), + [ + ("i0298", "threadfence_block"), + ("i0299", "threadfence"), + ("i0300", "threadfence_system"), + ] + ); + + let mut reordered = admission.clone(); + reordered.variants.swap(0, 1); + assert!(expand_threadfence_admission(&reordered).is_err()); + + let mut wrong_id = admission.clone(); + wrong_id.variants[0].abi_id = "i0300".into(); + assert!(expand_threadfence_admission(&wrong_id).is_err()); + + let mut executed = admission.clone(); + executed.runtime_validation = RuntimeValidation::Executed; + assert!(expand_threadfence_admission(&executed).is_err()); +} + +#[test] +fn pinned_threadfences_match_llvm_and_reject_contract_drift() { + let records = expand_threadfence_admission(&test_threadfence_admission()).unwrap(); + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let declarations = imported + .intrinsics + .iter() + .map(|record| (record.source_record.as_str(), record)) + .collect::>(); + + for record in &records { + let declaration = declarations[record.source_record.as_deref().unwrap()]; + validate_imported_policy(record, declaration).unwrap(); + assert_eq!( + declaration + .selections + .iter() + .filter(|selection| selection_matches_policy(record, selection).unwrap()) + .count(), + 1 + ); + } + + let declaration = declarations["int_nvvm_membar_cta"]; + let mut changed = records[0].clone(); + changed.memory = "none".into(); + assert!(validate_imported_policy(&changed, declaration).is_err()); + + let mut changed = records[0].clone(); + changed.convergent = true; + assert!(validate_imported_policy(&changed, declaration).is_err()); + + let mut changed = records[0].clone(); + changed.backend_lowerings[0].mechanism = BackendLoweringMechanism::InlinePtx; + assert!(validate_imported_policy(&changed, declaration).is_err()); + + let mut changed = records[0].clone(); + changed.minimum_ptx = "2.0".into(); + assert!(validate_imported_policy(&changed, declaration).is_err()); +} + +#[test] +fn pinned_cluster_sregs_preserve_helpers_and_reject_unused_w_components() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let (mut overlay, _) = + read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")).unwrap(); + bind_pinned_abi_ids(&repo_root, &mut overlay); + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let declarations = imported + .intrinsics + .iter() + .map(|record| (record.source_record.as_str(), record)) + .collect::>(); + let records = overlay + .intrinsics + .iter() + .filter(|record| { + record + .source_record + .as_deref() + .is_some_and(is_cluster_sreg_source) + }) + .collect::>(); + + assert_eq!(records.len(), 14); + let actual_abi_ids = records + .iter() + .map(|record| record.abi_id.clone()) + .collect::>(); + let expected_abi_ids = (263..=276) + .map(|id| format!("i{id:04}")) + .collect::>(); + assert_eq!(actual_abi_ids, expected_abi_ids); + for record in &records { + let declaration = declarations[record.source_record.as_deref().unwrap()]; + validate_imported_policy(record, declaration).unwrap(); + assert!(!declaration.source_record.ends_with("_w")); + } + + let compatibility_paths = records + .iter() + .flat_map(|record| record.compatibility_rust_paths.iter().map(String::as_str)) + .collect::>(); + assert_eq!( + compatibility_paths, + [ + "cuda_device::cluster::cluster_ctaidX", + "cuda_device::cluster::cluster_ctaidY", + "cuda_device::cluster::cluster_ctaidZ", + "cuda_device::cluster::cluster_nctaidX", + "cuda_device::cluster::cluster_nctaidY", + "cuda_device::cluster::cluster_nctaidZ", + "cuda_device::cluster::__cluster_grid_dimX", + "cuda_device::cluster::__cluster_grid_dimY", + "cuda_device::cluster::__cluster_grid_dimZ", + "cuda_device::cluster::__cluster_idxX", + "cuda_device::cluster::__cluster_idxY", + "cuda_device::cluster::__cluster_idxZ", + ] + .into_iter() + .collect() + ); + + let source = resolve_policy_source(records[0]).unwrap(); + let error = validate_sreg_policy( + records[0], + &source, + Some(declarations["int_nvvm_read_ptx_sreg_cluster_ctaid_w"]), + ) + .unwrap_err(); + assert!(error.to_string().contains("unused always-zero")); + + let mut mixed = (*records[0]).clone(); + mixed.sparse_mma = overlay + .intrinsics + .iter() + .find_map(|record| record.sparse_mma.clone()); + let source = resolve_policy_source(&mixed).unwrap(); + let error = validate_sreg_policy( + &mixed, + &source, + Some(declarations[mixed.source_record.as_deref().unwrap()]), + ) + .unwrap_err(); + assert!(error.to_string().contains("mixes another generated-family")); +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/tests/targets_abi.rs b/crates/cuda-intrinsics-gen/src/resolve/tests/targets_abi.rs new file mode 100644 index 0000000000..29bb145d4d --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/tests/targets_abi.rs @@ -0,0 +1,813 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, CatalogHalfOpenRange, CatalogHardwareAlternative, + CatalogHardwareTarget, CatalogTargetAlternative, CatalogTargetContract, + CatalogTargetRequirement, IntrinsicBackend, LdmatrixElement, LdmatrixLayout, + LdmatrixMultiplicity, LdmatrixShape, LdmatrixStateSpace, OverlayIntrinsic, PtxVersion, + ReduxAdapter, ReduxOperation, ReduxParticipation, TargetContract, TargetSelectorBinding, +}; + +use super::fixtures::*; +use crate::model::ImportedSelection; +use crate::resolve::abi_ledger::*; +use crate::resolve::driver::*; +use crate::resolve::evidence::*; +use crate::resolve::families::*; +use crate::resolve::guards::*; +use crate::resolve::materialize::*; +use crate::resolve::overlay::*; +use crate::resolve::policy::*; +use crate::resolve::targets::*; + +#[test] +fn selected_target_predicates_fail_closed() { + let selection = ImportedSelection { + source_record: "selection".into(), + asm: "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {{$r0, $r1, $r2, $r3}}, [$src];".into(), + predicates: vec![ + "Subtarget->getPTXVersion() >= 65".into(), + "Subtarget->getSmVersion() >= 75".into(), + ], + constraints: Default::default(), + }; + + let mut too_low_ptx = policy(); + too_low_ptx.minimum_ptx = "6.4".into(); + too_low_ptx.minimum_sm = Some("sm_75".into()); + assert!( + validate_selected_target_predicates(&too_low_ptx, &selection) + .unwrap_err() + .to_string() + .contains("minimum PTX") + ); + + let mut too_low_sm = policy(); + too_low_sm.minimum_ptx = "6.5".into(); + too_low_sm.minimum_sm = Some("sm_74".into()); + assert!( + validate_selected_target_predicates(&too_low_sm, &selection) + .unwrap_err() + .to_string() + .contains("minimum SM") + ); + + let mut unknown = selection; + unknown + .predicates + .push("Subtarget->hasMysteryFeature()".into()); + assert!( + validate_selected_target_predicates(&too_low_sm, &unknown) + .unwrap_err() + .to_string() + .contains("fail closed") + ); +} + +#[test] +fn blackwell_ldmatrix_predicate_requires_the_reviewed_exact_target_set() { + let mut record = policy(); + record.family = "ldmatrix".into(); + record.minimum_ptx = "8.6".into(); + record.minimum_sm = None; + record.targets = BLACKWELL_LDMATRIX_LLVM_TARGETS.into(); + record.ldmatrix_variant = Some(crate::model::LdmatrixVariant { + shape: LdmatrixShape::M16n16, + multiplicity: LdmatrixMultiplicity::X1, + layout: LdmatrixLayout::Transposed, + element: LdmatrixElement::B8, + state_space: LdmatrixStateSpace::Shared, + }); + let selection = ImportedSelection { + source_record: "selection".into(), + asm: "ldmatrix.sync.aligned.m16n16.x1.trans.shared.b8 {{$r0, $r1}}, [$src];".into(), + predicates: vec!["Subtarget->hasLdStmatrixBlackwellSupport()".into()], + constraints: Default::default(), + }; + + validate_selected_target_predicates(&record, &selection).unwrap(); + + record.targets = "sm_100a|sm_101a".into(); + assert!(validate_selected_target_predicates(&record, &selection).is_err()); + + record.targets = BLACKWELL_LDMATRIX_LLVM_TARGETS.into(); + let mut conflicting = selection; + conflicting + .predicates + .push("Subtarget->getPTXVersion() >= 86".into()); + assert!(validate_selected_target_predicates(&record, &conflicting).is_err()); +} + +#[test] +fn f32_redux_predicate_requires_the_reviewed_exact_target_matrix() { + let mut record = policy(); + record.family = "redux".into(); + record.minimum_ptx = "8.6".into(); + record.minimum_sm = None; + record.targets = REDUX_F32_TARGETS.into(); + record.redux = Some(crate::model::Redux { + operation: ReduxOperation::Fmin, + participation: ReduxParticipation::ExecutingLaneNamedAllNamedLanesSameInstructionAndMask, + adapter: ReduxAdapter::MaskValueToSourceMemberMask, + }); + let selection = ImportedSelection { + source_record: "selection".into(), + asm: "redux.sync.min.f32 \t$dst, $src, $mask;".into(), + predicates: vec!["Subtarget->hasReduxSyncF32()".into()], + constraints: Default::default(), + }; + + validate_selected_target_predicates(&record, &selection).unwrap(); + + record.targets = + "sm_100a|sm_100f|sm_103a|sm_103f|sm_110a|sm_110f|sm_120a|sm_120f|sm_121a|sm_121f".into(); + assert!(validate_selected_target_predicates(&record, &selection).is_err()); + + record.targets = REDUX_F32_TARGETS.into(); + record.minimum_ptx = "8.8".into(); + assert!(validate_selected_target_predicates(&record, &selection).is_err()); + + record.minimum_ptx = "8.6".into(); + let mut conflicting = selection; + conflicting + .predicates + .push("Subtarget->getPTXVersion() >= 88".into()); + assert!(validate_selected_target_predicates(&record, &conflicting).is_err()); +} + +#[test] +fn intrinsic_abi_identity_is_stable_and_explicit() { + let policy = policy(); + let declaration = declaration(); + let resolved = resolve_record( + &policy, + resolve_policy_source(&policy).unwrap(), + Some(&declaration), + &evidence(), + "test", + "LLVM version test", + "0123456789abcdef", + vec![], + 1, + ) + .unwrap(); + + assert_eq!(resolved.rust.abi_id, "i0001"); + assert_eq!( + resolved.rust.canonical_path, + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0001" + ); + assert_eq!( + resolved.rust.public_path, + "cuda_intrinsics::sreg::thread_idx_x" + ); + assert_eq!( + resolved.rust.compatibility_paths, + ["cuda_device::thread::threadIdx_x"] + ); + assert_eq!( + resolved.llvm.as_ref().unwrap().properties, + [ + "IntrNoMem", + "IntrSpeculatable", + "NoUndef", + "Range" + ] + ); + assert!(resolved.llvm.as_ref().unwrap().result_facts.no_undef); + assert_eq!( + resolved.llvm.as_ref().unwrap().result_facts.range, + Some(CatalogHalfOpenRange { + lower: "0".into(), + upper_exclusive: "1024".into(), + }) + ); + assert_eq!(resolved.backend.version, "LLVM version test"); + assert_eq!(resolved.backend.sha256, "0123456789abcdef"); +} + +#[test] +fn malformed_intrinsic_abi_ids_are_rejected() { + for abi_id in ["thread_idx_x", "i1", "x0001", "i00a1"] { + let mut record = policy(); + record.abi_id = abi_id.into(); + let error = validate_unique_overlay(&[record], 1).unwrap_err(); + assert!(error.to_string().contains("stable `iNNNN` form")); + } +} + +#[test] +fn ptx_versions_are_parsed_once_and_serialize_compatibly() { + for (text, encoded) in [("2.0", 20), ("6.5", 65), ("10.0", 100)] { + let version = parse_ptx_version(text, "test").unwrap(); + assert_eq!(version.encoded(), encoded); + assert_eq!( + serde_json::to_string(&version).unwrap(), + format!("\"{text}\"") + ); + assert_eq!( + serde_json::from_str::(&format!("\"{text}\"")).unwrap(), + version + ); + } + for malformed in ["6", "6.05", " 6.5", "06.5", "6.5 "] { + assert!(parse_ptx_version(malformed, "test").is_err(), "{malformed}"); + } +} + +#[test] +fn hardware_targets_are_parsed_without_losing_suffix_semantics() { + let all = policy(); + assert_eq!( + parse_hardware_target(&all).unwrap(), + CatalogHardwareTarget::All + ); + + let mut minimum = policy(); + minimum.minimum_sm = Some("sm_75".into()); + assert_eq!( + parse_hardware_target(&minimum).unwrap(), + CatalogHardwareTarget::AnyOf { + alternatives: vec![CatalogHardwareAlternative::MinimumSm { sm: 75 }], + } + ); + + let mut exact = policy(); + exact.targets = "sm_120a".into(); + assert_eq!( + parse_hardware_target(&exact).unwrap(), + CatalogHardwareTarget::AnyOf { + alternatives: vec![CatalogHardwareAlternative::ExactArchitecture { sm: 120 }], + } + ); + + let mut family = policy(); + family.targets = "sm_120f".into(); + assert_eq!( + parse_hardware_target(&family).unwrap(), + CatalogHardwareTarget::AnyOf { + alternatives: vec![CatalogHardwareAlternative::FamilyTarget { sm: 120 }], + } + ); +} + +#[test] +fn stage_hardware_shared_parser_preserves_canonical_evidence_language() { + for (target, expected) in [ + ("sm_75", CatalogHardwareAlternative::MinimumSm { sm: 75 }), + ( + "sm_120a", + CatalogHardwareAlternative::ExactArchitecture { sm: 120 }, + ), + ( + "compute_120f", + CatalogHardwareAlternative::FamilyTarget { sm: 120 }, + ), + ] { + assert_eq!(parse_stage_hardware(target), Some(expected), "{target}"); + } + for target in ["sm_090", "sm_1000"] { + assert_eq!(parse_stage_hardware(target), None, "{target}"); + } +} + +#[test] +fn reviewed_target_floors_equal_shared_backend_derivation() { + for (target, floor) in [ + ("sm_120a", 87), + ("sm_120f", 88), + ("sm_121a", 88), + ("sm_121f", 88), + ] { + let hardware = parse_stage_hardware(target).unwrap(); + assert_eq!(f8f6f4_llvm_ptx_floor(hardware).unwrap(), floor, "{target}"); + } + for (target, floor) in [ + ("sm_100a", 86), + ("sm_100f", 88), + ("sm_103a", 88), + ("sm_103f", 88), + ("sm_110a", 90), + ("sm_110f", 90), + ("sm_120a", 87), + ("sm_120f", 88), + ("sm_121a", 88), + ("sm_121f", 88), + ] { + let hardware = parse_stage_hardware(target).unwrap(); + assert_eq!( + blackwell_ldmatrix_llvm_ptx_floor(hardware).unwrap(), + floor, + "{target}" + ); + } +} + +fn selector(name: &str, value: &str) -> TargetSelectorBinding { + TargetSelectorBinding { + name: name.into(), + value: value.into(), + } +} + +fn target_contract( + selectors: Vec, + alternatives: &[(&str, &str)], +) -> TargetContract { + TargetContract { + selectors, + alternatives: alternatives + .iter() + .map( + |(target, minimum_ptx)| crate::model::TargetContractAlternative { + target: (*target).into(), + minimum_ptx: (*minimum_ptx).into(), + }, + ) + .collect(), + } +} + +#[test] +fn target_contracts_keep_selector_specific_ptx_hardware_pairs() { + let contracts = [ + target_contract( + vec![selector("kind", "f16")], + &[ + ("sm_100a", "8.6"), + ("sm_101a", "8.6"), + ("sm_103a", "8.8"), + ("sm_110a", "9.0"), + ], + ), + target_contract( + vec![selector("kind", "i8")], + &[("sm_100a", "8.6"), ("sm_101a", "8.6"), ("sm_110a", "9.0")], + ), + ]; + let full_requirement = resolve_target_contracts("tcgen05_mma", &contracts).unwrap(); + assert_eq!(full_requirement.minimum_ptx.encoded(), 86); + let CatalogHardwareTarget::TargetMatrix { + contracts: full_contracts, + } = &full_requirement.hardware + else { + panic!("full target contracts must remain a matrix") + }; + assert_eq!(full_contracts.len(), 2); + assert_eq!(full_contracts[0].selectors, [selector("kind", "f16")]); + assert_eq!(full_contracts[1].selectors, [selector("kind", "i8")]); + + let requirement = + resolve_target_contract("tcgen05_mma", &[selector("kind", "i8")], &contracts).unwrap(); + + assert_eq!(requirement.minimum_ptx.encoded(), 86); + assert_eq!( + requirement.hardware, + CatalogHardwareTarget::TargetMatrix { + contracts: vec![CatalogTargetContract { + selectors: vec![selector("kind", "i8")], + alternatives: vec![ + CatalogTargetAlternative { + minimum_ptx: "8.6".parse().unwrap(), + hardware: CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, + }, + CatalogTargetAlternative { + minimum_ptx: "8.6".parse().unwrap(), + hardware: CatalogHardwareAlternative::ExactArchitecture { sm: 101 }, + }, + CatalogTargetAlternative { + minimum_ptx: "9.0".parse().unwrap(), + hardware: CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, + }, + ], + }], + } + ); + + let mut policy = policy(); + policy.id = "tcgen05_mma".into(); + validate_candidate_target(&policy, &requirement, "sm_100a", "+ptx86").unwrap(); + assert!(validate_candidate_target(&policy, &requirement, "sm_103a", "+ptx88").is_err()); + assert!(validate_candidate_target(&policy, &requirement, "sm_110a", "+ptx88").is_err()); + validate_candidate_target(&policy, &requirement, "sm_110a", "+ptx90").unwrap(); + + let libnvvm = [target_contract( + vec![selector("kind", "i8")], + &[("sm_100a", "8.6"), ("sm_110a", "9.0")], + )]; + let libnvvm_requirement = + resolve_target_contract("tcgen05_mma", &[selector("kind", "i8")], &libnvvm).unwrap(); + assert_ne!(requirement, libnvvm_requirement); + validate_candidate_target(&policy, &requirement, "sm_101a", "+ptx86").unwrap(); + assert!(validate_candidate_target(&policy, &libnvvm_requirement, "sm_101a", "+ptx86").is_err()); +} + +#[test] +fn target_contract_selection_and_shape_fail_closed() { + resolve_target_contract( + "tcgen05_mma", + &[selector("kind", "f16")], + &[target_contract( + vec![selector("kind", "f16")], + &[("sm_100a", "8.6"), ("sm_100f", "8.8"), ("sm_103a", "8.8")], + )], + ) + .unwrap(); + + let valid = target_contract( + vec![selector("kind", "f16"), selector("scale_d", "false")], + &[("sm_100a", "8.6"), ("sm_103a", "8.8")], + ); + assert!( + resolve_target_contract( + "tcgen05_mma", + &[selector("kind", "i8"), selector("scale_d", "false")], + std::slice::from_ref(&valid), + ) + .unwrap_err() + .to_string() + .contains("exactly one reviewed contract") + ); + + for invalid in [ + target_contract( + vec![selector("scale_d", "false"), selector("kind", "f16")], + &[("sm_100a", "8.6")], + ), + target_contract(vec![selector("kind", "F16")], &[("sm_100a", "8.6")]), + target_contract(vec![selector("kind_", "f16")], &[("sm_100a", "8.6")]), + target_contract( + vec![selector("kind", "f16")], + &[("sm_103a", "8.8"), ("sm_100a", "8.6")], + ), + target_contract( + vec![selector("kind", "f16")], + &[("sm_100+", "8.6"), ("sm_103a", "8.8")], + ), + target_contract(vec![selector("kind", "f16")], &[("sm_100a", "8.60")]), + ] { + assert!( + resolve_target_contract( + "tcgen05_mma", + &invalid.selectors, + std::slice::from_ref(&invalid) + ) + .is_err(), + "{invalid:?}" + ); + } +} + +#[test] +fn legacy_target_requirement_json_is_unchanged() { + let requirement = CatalogTargetRequirement { + minimum_ptx: "8.6".parse().unwrap(), + hardware: CatalogHardwareTarget::AnyOf { + alternatives: vec![CatalogHardwareAlternative::ExactArchitecture { sm: 100 }], + }, + }; + assert_eq!( + serde_json::to_value(requirement).unwrap(), + serde_json::json!({ + "minimum_ptx": "8.6", + "hardware": { + "kind": "any_of", + "alternatives": [{ "kind": "exact_architecture", "sm": 100 }] + } + }) + ); +} + +#[test] +fn malformed_or_conflicting_hardware_targets_are_rejected() { + for malformed in [ + "sm_120", + "sm_120af", + "sm_120A", + "sm_0120a", + "sm_0a", + "sm_120+", + "compute_120a", + "all ", + ] { + let mut record = policy(); + record.targets = malformed.into(); + assert!(parse_hardware_target(&record).is_err(), "{malformed}"); + } + + let mut suffixed_minimum = policy(); + suffixed_minimum.minimum_sm = Some("sm_90a".into()); + assert!(parse_hardware_target(&suffixed_minimum).is_err()); + + for target in ["sm_120a", "sm_120f"] { + let mut conflicting = policy(); + conflicting.targets = target.into(); + conflicting.minimum_sm = Some("sm_120".into()); + let error = parse_hardware_target(&conflicting).unwrap_err().to_string(); + assert!(error.contains("cannot be combined"), "{error}"); + } +} + +#[test] +fn exact_inline_ptx_routes_can_inherit_exact_or_family_targets() { + for target in ["sm_120a", "sm_120f"] { + let mut record = policy(); + record.minimum_ptx = "8.7".into(); + record.targets = target.into(); + record.backend_lowerings = [IntrinsicBackend::LlvmNvptx, IntrinsicBackend::LibNvvm] + .into_iter() + .map(|backend| crate::model::OverlayBackendLowering { + backend, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: "test".into(), + targets: None, + minimum_ptx: Some("8.7".into()), + minimum_sm: None, + }) + .collect(); + + ensure_exact_inline_ptx_backends( + &record, + [ + (IntrinsicBackend::LlvmNvptx, "8.7", None), + (IntrinsicBackend::LibNvvm, "8.7", None), + ], + "test", + ) + .unwrap(); + for lowering in &record.backend_lowerings { + assert_eq!( + backend_target_requirement(&record, lowering) + .unwrap() + .hardware, + parse_hardware_target(&record).unwrap() + ); + } + + record.backend_lowerings[0].minimum_sm = Some("sm_120".into()); + assert!( + ensure_exact_inline_ptx_backends( + &record, + [ + (IntrinsicBackend::LlvmNvptx, "8.7", None), + (IntrinsicBackend::LibNvvm, "8.7", None), + ], + "test", + ) + .is_err() + ); + } +} + +#[test] +fn backend_route_target_override_is_exact_and_does_not_inherit_the_record_sm_floor() { + let mut record = policy(); + record.minimum_ptx = "8.6".into(); + record.minimum_sm = Some("sm_75".into()); + record.targets = "all".into(); + let lowering = crate::model::OverlayBackendLowering { + backend: IntrinsicBackend::LibNvvm, + mechanism: BackendLoweringMechanism::InlinePtx, + evidence_profile: "test".into(), + targets: Some("sm_100a|sm_120a".into()), + minimum_ptx: None, + minimum_sm: None, + }; + + assert_eq!( + backend_target_requirement(&record, &lowering).unwrap(), + CatalogTargetRequirement { + minimum_ptx: parse_ptx_version("8.6", &record.id).unwrap(), + hardware: CatalogHardwareTarget::AnyOf { + alternatives: vec![ + CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 120 }, + ], + }, + } + ); +} + +#[test] +fn abi_ledger_requires_exact_active_identity() { + let record = policy(); + let frozen_entry = ledger_entry(&record); + validate_abi_ledger( + &overlay_file(vec![record.clone()]), + &ledger(vec![frozen_entry.clone()]), + ) + .unwrap(); + + let mut reassigned = record.clone(); + reassigned.id = "different_catalog_id".into(); + let error = validate_abi_ledger( + &overlay_file(vec![reassigned]), + &ledger(vec![frozen_entry.clone()]), + ) + .unwrap_err(); + assert!(error.to_string().contains("catalog ID mismatch")); + + let mut reassigned = record.clone(); + reassigned.operation_key = "launch.block_index.x".into(); + let error = validate_abi_ledger( + &overlay_file(vec![reassigned]), + &ledger(vec![frozen_entry.clone()]), + ) + .unwrap_err(); + assert!(error.to_string().contains("operation key mismatch")); + + for mutate in [ + |record: &mut OverlayIntrinsic| record.safe = false, + |record: &mut OverlayIntrinsic| record.rust_arguments.push("u32".into()), + |record: &mut OverlayIntrinsic| record.rust_result = "u64".into(), + ] { + let mut changed_signature = record.clone(); + mutate(&mut changed_signature); + let error = validate_abi_ledger( + &overlay_file(vec![changed_signature]), + &ledger(vec![frozen_entry.clone()]), + ) + .unwrap_err(); + assert!(error.to_string().contains("raw Rust signature mismatch")); + } +} + +#[test] +fn generated_abi_binding_uses_catalog_identity_not_axis_position() { + let mut record = policy(); + let mut frozen = ledger_entry(&record); + frozen.abi_id = "i9001".into(); + record.abi_id.clear(); + let mut overlay = overlay_file(vec![record]); + + bind_generated_abi_ids(&mut overlay, &ledger(vec![frozen])).unwrap(); + + assert_eq!(overlay.intrinsics[0].abi_id, "i9001"); +} + +#[test] +fn generated_abi_binding_rejects_missing_tombstoned_or_ambiguous_identity() { + let record = policy(); + let mut unbound = record.clone(); + unbound.abi_id.clear(); + + let error = bind_generated_abi_ids( + &mut overlay_file(vec![unbound.clone()]), + &ledger(vec![ledger_entry(&distinct_policy())]), + ) + .unwrap_err(); + assert!(error.to_string().contains("has no ABI ledger entry")); + + let mut tombstone = ledger_entry(&record); + tombstone.status = "tombstone".into(); + let error = bind_generated_abi_ids( + &mut overlay_file(vec![unbound.clone()]), + &ledger(vec![tombstone]), + ) + .unwrap_err(); + assert!(error.to_string().contains("non-active ABI ledger entry")); + + let first = ledger_entry(&record); + let mut duplicate = first.clone(); + duplicate.abi_id = "i9002".into(); + let error = bind_generated_abi_ids( + &mut overlay_file(vec![unbound]), + &ledger(vec![first, duplicate]), + ) + .unwrap_err(); + assert!( + error + .to_string() + .contains("duplicate ABI ledger catalog ID") + ); +} + +#[test] +fn generated_abi_binding_checks_derived_operation_and_raw_signature() { + let record = policy(); + let mut unbound = record.clone(); + unbound.abi_id.clear(); + + let mut wrong_operation = ledger_entry(&record); + wrong_operation.operation_key = "launch.block_index.x".into(); + let error = bind_generated_abi_ids( + &mut overlay_file(vec![unbound.clone()]), + &ledger(vec![wrong_operation]), + ) + .unwrap_err(); + assert!(error.to_string().contains("operation key mismatch")); + + let mut wrong_signature = ledger_entry(&record); + wrong_signature + .raw_rust_signature + .arguments + .push("u32".into()); + let error = bind_generated_abi_ids( + &mut overlay_file(vec![unbound]), + &ledger(vec![wrong_signature]), + ) + .unwrap_err(); + assert!(error.to_string().contains("raw Rust signature mismatch")); +} + +#[test] +fn abi_ledger_does_not_freeze_public_or_backend_implementation_details() { + let record = policy(); + let frozen_entry = ledger_entry(&record); + let mut evolved = record.clone(); + evolved.rust_module = "coordinates".into(); + evolved.rust_name = "thread_x".into(); + evolved.public_rust_path = "cuda_intrinsics::coordinates::thread_x".into(); + evolved.llvm_symbol = Some("llvm.nvvm.backend.v2.tid.x".into()); + evolved.llvm_arguments = vec!["i8".into()]; + evolved.llvm_results = vec!["i64".into()]; + evolved.dialect_op_type = "ReadThreadIndexXOpV2".into(); + evolved.dialect_op_name = "nvvm.read_thread_index_x_v2".into(); + evolved.lowering = "backend_v2_adapter".into(); + + validate_abi_ledger(&overlay_file(vec![evolved]), &ledger(vec![frozen_entry])).unwrap(); +} + +#[test] +fn tombstoned_or_unlisted_abi_ids_cannot_reappear() { + let record = policy(); + let mut tombstone = ledger_entry(&record); + tombstone.status = "tombstone".into(); + let error = validate_abi_ledger( + &overlay_file(vec![record.clone()]), + &ledger(vec![tombstone]), + ) + .unwrap_err(); + assert!(error.to_string().contains("cannot reappear")); + + let error = validate_abi_ledger(&overlay_file(vec![record]), &ledger(vec![])).unwrap_err(); + assert!(error.to_string().contains("contains no entries")); +} + +#[test] +fn every_active_ledger_entry_requires_an_overlay_record() { + let record = policy(); + let error = validate_abi_ledger(&overlay_file(vec![]), &ledger(vec![ledger_entry(&record)])) + .unwrap_err(); + assert!(error.to_string().contains("has no overlay record")); +} + +#[test] +fn candidate_targets_are_canonical_and_satisfy_every_floor() { + let policy = policy(); + let requirement = CatalogTargetRequirement { + minimum_ptx: "7.0".parse().unwrap(), + hardware: CatalogHardwareTarget::AnyOf { + alternatives: vec![CatalogHardwareAlternative::MinimumSm { sm: 80 }], + }, + }; + validate_candidate_target(&policy, &requirement, "sm_80", "+ptx70").unwrap(); + validate_candidate_target(&policy, &requirement, "sm_90a", "+ptx86").unwrap(); + assert!( + validate_candidate_target(&policy, &requirement, "sm_75", "+ptx70") + .unwrap_err() + .to_string() + .contains("hardware requirement") + ); + assert!( + validate_candidate_target(&policy, &requirement, "sm_80", "+ptx69") + .unwrap_err() + .to_string() + .contains("PTX floor") + ); + for malformed in ["compute_80", "sm_080", "sm_80x"] { + assert!( + validate_candidate_target(&policy, &requirement, malformed, "+ptx70").is_err(), + "{malformed}" + ); + } + for malformed in ["ptx70", "+ptx7", "+ptx070"] { + assert!( + validate_candidate_target(&policy, &requirement, "sm_80", malformed).is_err(), + "{malformed}" + ); + } + + let exact = CatalogTargetRequirement { + minimum_ptx: "8.7".parse().unwrap(), + hardware: CatalogHardwareTarget::AnyOf { + alternatives: vec![CatalogHardwareAlternative::ExactArchitecture { sm: 120 }], + }, + }; + validate_candidate_target(&policy, &exact, "sm_120a", "+ptx87").unwrap(); + assert!(validate_candidate_target(&policy, &exact, "sm_120a", "+ptx86").is_err()); + assert!(validate_candidate_target(&policy, &exact, "sm_120", "+ptx87").is_err()); + assert!(validate_candidate_target(&policy, &exact, "sm_120f", "+ptx87").is_err()); + + let family = CatalogTargetRequirement { + minimum_ptx: "8.7".parse().unwrap(), + hardware: CatalogHardwareTarget::AnyOf { + alternatives: vec![CatalogHardwareAlternative::FamilyTarget { sm: 120 }], + }, + }; + validate_candidate_target(&policy, &family, "sm_120f", "+ptx87").unwrap(); + assert!(validate_candidate_target(&policy, &family, "sm_120a", "+ptx87").is_err()); +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/tests/tcgen05.rs b/crates/cuda-intrinsics-gen/src/resolve/tests/tcgen05.rs new file mode 100644 index 0000000000..90766e5124 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/tests/tcgen05.rs @@ -0,0 +1,1366 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + CatalogHardwareAlternative, CatalogHardwareTarget, ImportedFile, OverlayIntrinsic, + OverlayShardFile, RuntimeValidation, Tcgen05Adapter, Tcgen05MmaBUsage, + Tcgen05MmaFixedSelectors, Tcgen05MmaForm, Tcgen05MmaKind, Tcgen05MmaSelectorLayout, + Tcgen05Operation, Tcgen05SourceContract, +}; +use crate::ptx::OperandPattern; +use crate::util::read_json; +use std::collections::{BTreeMap, BTreeSet}; +use std::path::Path; + +use super::fixtures::*; +use crate::resolve::families::*; +use crate::resolve::guards::*; +use crate::resolve::overlay::*; +use crate::resolve::targets::*; + +#[test] +fn compact_tcgen05_admission_matches_llvm_and_fails_closed() { + let records = expand_tcgen05_admission(&test_tcgen05_admission()).unwrap(); + assert_eq!(records.len(), 27); + assert_eq!( + records + .iter() + .map(|record| (record.abi_id.as_str(), record.id.as_str())) + .collect::>(), + [ + ("i0343", "tcgen05_alloc"), + ("i0344", "tcgen05_dealloc"), + ("i0345", "tcgen05_relinquish_alloc_permit"), + ("i0346", "tcgen05_fence_before_thread_sync"), + ("i0347", "tcgen05_fence_after_thread_sync"), + ("i0348", "tcgen05_commit"), + ("i0349", "tcgen05_commit_shared_cluster"), + ("i0350", "tcgen05_mma_ws_f16"), + ("i0351", "tcgen05_mma_f16"), + ("i0352", "tcgen05_mma_ws_bf16"), + ("i0353", "tcgen05_mma_ws_tf32"), + ("i0354", "tcgen05_cp_smem_to_tmem"), + ("i0355", "tcgen05_ld_16x256b_x8_pure"), + ("i0356", "tcgen05_ld_16x256b_pure"), + ("i0357", "tcgen05_load_wait"), + ("i0358", "tcgen05_store_wait"), + ("i0359", "tcgen05_alloc_cg2"), + ("i0360", "tcgen05_dealloc_cg2"), + ("i0361", "tcgen05_relinquish_alloc_permit_cg2"), + ("i0362", "tcgen05_mma_f16_cg2"), + ("i0363", "tcgen05_commit_cg2"), + ("i0364", "tcgen05_commit_shared_cluster_cg2"), + ("i0365", "tcgen05_commit_multicast_cg2"), + ("i0366", "tcgen05_cp_smem_to_tmem_cg2"), + ("i0760", "tcgen05_commit_multicast"), + ("i0761", "tcgen05_shift_down"), + ("i0762", "tcgen05_shift_down_cg2"), + ] + ); + + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let declarations = imported + .intrinsics + .iter() + .map(|record| (record.source_record.as_str(), record)) + .collect::>(); + for record in &records { + let declaration = declarations[record.source_record.as_deref().unwrap()]; + validate_imported_policy(record, declaration).unwrap(); + assert_tcgen05_backend_target_split(record); + assert_eq!( + record.execution_scope, + record.tcgen05.as_ref().unwrap().operation.execution_scope() + ); + } + + let bf16 = records + .iter() + .find(|record| record.id == "tcgen05_mma_ws_bf16") + .unwrap(); + assert!(bf16.expected_ptx.modifiers.contains(&"kind::f16".into())); + assert!(!bf16.expected_ptx.modifiers.contains(&"kind::bf16".into())); + assert_eq!( + bf16.tcgen05.as_ref().unwrap().source_contract, + Tcgen05SourceContract::TablegenSelectionChangesPtx + ); + + let mma_f16 = records + .iter() + .find(|record| record.id == "tcgen05_mma_f16") + .unwrap(); + assert!(mma_f16.expected_ptx.matches( + "tcgen05.mma.cta_group::1.kind::f16 [%r1], %rd1, %rd2, %r2, {%z, %z, %z, %z}, %enable_pred;" + ).unwrap()); + assert!(!mma_f16.expected_ptx.matches( + "tcgen05.mma.cta_group::1.kind::f16 [%r1], %rd1, %rd2, %r2, {%clock64, %z, %z, %z}, %enable_pred;" + ).unwrap()); + assert!(!mma_f16.expected_ptx.matches( + "tcgen05.mma.cta_group::1.kind::f16 [%r1], %rd1, %rd2, %r2, {%z, %z, %z, %z}, %other_pred;" + ).unwrap()); + + let mma_f16_cg2 = records + .iter() + .find(|record| record.id == "tcgen05_mma_f16_cg2") + .unwrap(); + assert!(mma_f16_cg2.expected_ptx.matches( + "tcgen05.mma.cta_group::2.kind::f16 [%r1], %rd1, %rd2, %r2, {%z, %z, %z, %z, %z, %z, %z, %z}, %enable_pred;" + ).unwrap()); + assert!( + !records + .iter() + .any(|record| { record.rust_name == "tcgen05_mma_ws_f16_with_collector" }) + ); + assert_eq!( + parse_hardware_target(&records[0]).unwrap(), + CatalogHardwareTarget::AnyOf { + alternatives: vec![ + CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 101 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 103 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, + ], + } + ); + + let legacy_records = + expand_tcgen05_admission(&without_tcgen05_control(test_tcgen05_admission())).unwrap(); + assert_eq!(legacy_records.len(), 24); + assert!(legacy_records.iter().all(|record| { + record + .backend_lowerings + .iter() + .map(|lowering| lowering.evidence_profile.as_str()) + .eq(["llvm-tcgen05-test", "libnvvm-tcgen05-test"]) + })); + + let multicast = records + .iter() + .find(|record| record.id == "tcgen05_commit_multicast") + .unwrap(); + assert_eq!(multicast.abi_id, "i0760"); + assert_eq!(multicast.rust_arguments, ["*mut u64", "u16"]); + assert_eq!(multicast.llvm_arguments, ["shared_ptr", "i16"]); + assert_eq!(multicast.execution_scope, "thread"); + assert_eq!( + multicast.backend_lowerings[0].evidence_profile, + "llvm-tcgen05-control-test" + ); + assert_eq!( + multicast.backend_lowerings[1].evidence_profile, + "libnvvm-tcgen05-control-test" + ); + + for (id, group) in [ + ("tcgen05_shift_down", "cta_group::1"), + ("tcgen05_shift_down_cg2", "cta_group::2"), + ] { + let shift = records.iter().find(|record| record.id == id).unwrap(); + assert_eq!(shift.rust_arguments, ["u32"]); + assert_eq!(shift.llvm_arguments, ["tmem_ptr"]); + assert_eq!(shift.execution_scope, "thread"); + assert_eq!(shift.memory, "read_write"); + assert_eq!(shift.expected_ptx.modifiers, ["shift", group, "down"]); + assert_eq!(shift.expected_ptx.operands, [OperandPattern::Address]); + } + + let mut missing = test_tcgen05_admission(); + missing.variants.pop(); + assert!(expand_tcgen05_admission(&missing).is_err()); + + let mut reordered = test_tcgen05_admission(); + reordered.variants.swap(0, 1); + assert!(expand_tcgen05_admission(&reordered).is_err()); + + let mut wrong_abi = test_tcgen05_admission(); + wrong_abi.variants[0].abi_id = "i9999".into(); + assert!(expand_tcgen05_admission(&wrong_abi).is_err()); + + let mut executed = test_tcgen05_admission(); + executed.runtime_validation = RuntimeValidation::Executed; + assert!(expand_tcgen05_admission(&executed).is_err()); + + let mut missing_control_evidence = test_tcgen05_admission(); + missing_control_evidence.control_llvm_evidence_profile = None; + assert!(expand_tcgen05_admission(&missing_control_evidence).is_err()); + + let mut partial_control = without_tcgen05_control(test_tcgen05_admission()); + partial_control + .variants + .push(test_tcgen05_admission().variants[24].clone()); + assert!(expand_tcgen05_admission(&partial_control).is_err()); + + let shift = records + .iter() + .find(|record| record.id == "tcgen05_shift_down") + .unwrap(); + let mut wrong_shift_declaration = declarations[shift.source_record.as_deref().unwrap()].clone(); + wrong_shift_declaration.selections[0].predicates = + vec!["Subtarget->hasTcgen05InstSupport()".into()]; + assert!(validate_imported_policy(shift, &wrong_shift_declaration).is_err()); + + let declaration = declarations[records[0].source_record.as_deref().unwrap()]; + let mut wrong_adapter = records[0].clone(); + wrong_adapter.tcgen05.as_mut().unwrap().adapter = Tcgen05Adapter::NoOperands; + assert!(validate_imported_policy(&wrong_adapter, declaration).is_err()); + + let mut changed_declaration = declaration.clone(); + changed_declaration.properties.pop(); + assert!(validate_imported_policy(&records[0], &changed_declaration).is_err()); + + let mut broadened_libnvvm = records[0].clone(); + broadened_libnvvm.backend_lowerings[1].targets = None; + assert!(validate_imported_policy(&broadened_libnvvm, declaration).is_err()); + + let mut wrong_alloc_scope = records[0].clone(); + wrong_alloc_scope.execution_scope = "thread".into(); + assert!(validate_imported_policy(&wrong_alloc_scope, declaration).is_err()); + + let pure_load = &records[12]; + let pure_declaration = declarations[pure_load.source_record.as_deref().unwrap()]; + let mut wrong_pure_load_scope = pure_load.clone(); + wrong_pure_load_scope.execution_scope = "thread".into(); + assert!(validate_imported_policy(&wrong_pure_load_scope, pure_declaration).is_err()); +} + +#[test] +fn compact_tcgen05_mma_admission_closes_sources_selectors_and_targets() { + let records = expand_tcgen05_admission(&test_tcgen05_mma_admission()).unwrap(); + assert_eq!(records.len(), 51); + let mma = &records[27..]; + assert_eq!(mma.len(), 24); + assert_eq!( + mma.iter() + .map(|record| (record.abi_id.as_str(), record.id.as_str())) + .collect::>(), + [ + ("i0763", "tcgen05_mma_shared"), + ("i0764", "tcgen05_mma_tensor"), + ("i0765", "tcgen05_mma_tensor_ashift"), + ("i0766", "tcgen05_mma_sp_shared"), + ("i0767", "tcgen05_mma_sp_tensor"), + ("i0768", "tcgen05_mma_sp_tensor_ashift"), + ("i0769", "tcgen05_mma_ws_shared"), + ("i0770", "tcgen05_mma_ws_shared_zero_col_mask"), + ("i0771", "tcgen05_mma_ws_sp_shared"), + ("i0772", "tcgen05_mma_ws_sp_shared_zero_col_mask"), + ("i0773", "tcgen05_mma_ws_sp_tensor"), + ("i0774", "tcgen05_mma_ws_sp_tensor_zero_col_mask"), + ("i0775", "tcgen05_mma_ws_tensor"), + ("i0776", "tcgen05_mma_ws_tensor_zero_col_mask"), + ("i0777", "tcgen05_mma_ws_e4m3"), + ("i0778", "tcgen05_mma_ws_e5m2"), + ("i0779", "tcgen05_mma_ws_e2m3"), + ("i0780", "tcgen05_mma_ws_e3m2"), + ("i0781", "tcgen05_mma_ws_e2m1"), + ("i1011", "tcgen05_mma_e4m3"), + ("i1012", "tcgen05_mma_e5m2"), + ("i1013", "tcgen05_mma_e2m3"), + ("i1014", "tcgen05_mma_e3m2"), + ("i1015", "tcgen05_mma_e2m1"), + ] + ); + assert!(mma.iter().all(|record| { + record.dialect_op_type == TCGEN05_MMA_DIALECT_OP_TYPE + && record.dialect_op_name == TCGEN05_MMA_DIALECT_OP_NAME + })); + assert_eq!( + [ + Tcgen05MmaBUsage::Discard, + Tcgen05MmaBUsage::LastUse, + Tcgen05MmaBUsage::Fill, + Tcgen05MmaBUsage::Use, + ] + .map(Tcgen05MmaBUsage::selector_value), + [0, 1, 2, 3] + ); + + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let declarations = imported + .intrinsics + .iter() + .map(|record| (record.source_record.as_str(), record)) + .collect::>(); + for record in mma { + let declaration = declarations[record.source_record.as_deref().unwrap()]; + validate_imported_policy(record, declaration).unwrap(); + assert_eq!(record.memory, "read_write"); + assert!(record.convergent); + assert_eq!(record.execution_scope, "thread"); + assert!(!record.safe); + assert_eq!( + record + .expected_ptx + .operands + .iter() + .filter(|operand| { + **operand + == (OperandPattern::Exact { + value: "%enable_pred".into(), + }) + }) + .count(), + 1 + ); + } + + let dense = &mma[0]; + assert!(dense.expected_ptx.matches( + "tcgen05.mma.cta_group::1.kind::f16.collector::a::discard [%r1], %rd1, %rd2, %r2, %enable_pred;" + ).unwrap()); + assert!(!dense.expected_ptx.matches( + "tcgen05.mma.cta_group::1.kind::f16.collector::a::discard [%r1], %rd1, %rd2, %r2, %clock64;" + ).unwrap()); + assert!(mma[4].expected_ptx.matches( + "tcgen05.mma.sp.cta_group::1.kind::f16.collector::a::discard [%r1], [%r2], %rd1, [%r3], %r4, %enable_pred;" + ).unwrap()); + assert!(mma[7].expected_ptx.matches( + "tcgen05.mma.ws.cta_group::1.kind::f16.collector::b0::discard [%r1], %rd1, %rd2, %r2, %enable_pred, %rd3;" + ).unwrap()); + assert!(mma[14].expected_ptx.matches( + "tcgen05.mma.ws.cta_group::1.kind::f8f6f4.collector::b0::discard [%r1], [%r2], %rd1, %r3, %enable_pred;" + ).unwrap()); + assert!(mma[19].expected_ptx.matches( + "tcgen05.mma.cta_group::1.kind::f8f6f4.collector::a::discard [%r1], %rd1, %rd2, %r2, %enable_pred;" + ).unwrap()); + + let canonical = &mma[..14]; + assert_eq!( + canonical + .iter() + .map(|record| record.source_record.as_deref().unwrap()) + .collect::>() + .len(), + 14 + ); + let selected_count = |record: &OverlayIntrinsic| { + declarations[record.source_record.as_deref().unwrap()] + .selections + .iter() + .filter(|selection| selection_matches_policy(record, selection).unwrap()) + .count() + }; + let newly_matched = canonical + .iter() + .filter(|record| record.id != "tcgen05_mma_ws_tensor") + .map(selected_count) + .sum::(); + assert_eq!(newly_matched, 608); + assert_eq!(selected_count(&canonical[12]), 64); + assert_eq!(canonical.iter().map(selected_count).sum::(), 672); + assert_eq!(mma[14..].iter().map(selected_count).sum::(), 10); + assert_eq!( + [2, 5] + .into_iter() + .map(|index| { + declarations[canonical[index].source_record.as_deref().unwrap()] + .selections + .len() + - selected_count(&canonical[index]) + }) + .sum::(), + 32 + ); + + for record in canonical { + assert_eq!( + record.compatibility_rust_paths, + [format!("cuda_device::tcgen05::__{}", record.id)] + ); + assert_eq!( + &record.rust_arguments[record.rust_arguments.len() - 3..], + ["u32", "u32", "u32"] + ); + let contract = record.tcgen05.as_ref().unwrap().mma.as_ref().unwrap(); + let CatalogHardwareTarget::TargetMatrix { contracts } = &contract.llvm_target.hardware + else { + panic!("generic LLVM tcgen05 MMA target must be a matrix") + }; + assert_eq!(contracts.len(), 4); + assert_eq!( + contracts + .iter() + .map(|contract| contract.alternatives.len()) + .collect::>(), + [8, 8, 3, 8] + ); + let CatalogHardwareTarget::TargetMatrix { contracts } = &contract.libnvvm_target.hardware + else { + panic!("generic libNVVM tcgen05 MMA target must be a matrix") + }; + assert_eq!(contracts.len(), 4); + assert_eq!( + contracts + .iter() + .map(|contract| contract.alternatives.len()) + .collect::>(), + [6, 6, 2, 6] + ); + } + for record in &mma[14..19] { + assert_eq!( + record.compatibility_rust_paths, + [format!("cuda_device::tcgen05::{}", record.id)] + ); + assert_eq!( + record.rust_arguments, + ["u32", "u32", "u64", "u64", "u32", "bool"] + ); + assert_eq!(record.dialect_operands, ["i32", "i32", "i64", "i32", "i1"]); + let contract = record.tcgen05.as_ref().unwrap().mma.as_ref().unwrap(); + assert_eq!( + contract.fixed_selectors, + Some(Tcgen05MmaFixedSelectors { + kind: Tcgen05MmaKind::F8f6f4, + b_buffer: 0, + b_usage: Tcgen05MmaBUsage::Discard, + }) + ); + let CatalogHardwareTarget::TargetMatrix { contracts } = &contract.llvm_target.hardware + else { + panic!("fixed LLVM tcgen05 MMA target must be a matrix") + }; + assert_eq!(contracts.len(), 1); + assert_eq!(contracts[0].selectors[0].value, "f8f6f4"); + } + for record in &mma[19..] { + assert_eq!( + record.compatibility_rust_paths, + [format!("cuda_device::tcgen05::{}", record.id)] + ); + assert_eq!(record.rust_arguments, ["u32", "u64", "u64", "u32", "bool"]); + assert_eq!(record.dialect_operands, ["i32", "i64", "i64", "i32", "i1"]); + let contract = record.tcgen05.as_ref().unwrap().mma.as_ref().unwrap(); + assert_eq!(contract.form, Tcgen05MmaForm::Shared); + assert_eq!( + contract.fixed_selectors, + Some(Tcgen05MmaFixedSelectors { + kind: Tcgen05MmaKind::F8f6f4, + b_buffer: 0, + b_usage: Tcgen05MmaBUsage::Discard, + }) + ); + assert_eq!( + record.tcgen05.as_ref().unwrap().adapter, + Tcgen05Adapter::MmaDirectSelectors + ); + let CatalogHardwareTarget::TargetMatrix { contracts } = &contract.llvm_target.hardware + else { + panic!("fixed LLVM tcgen05 MMA target must be a matrix") + }; + assert_eq!(contracts.len(), 1); + assert_eq!(contracts[0].selectors[0].value, "f8f6f4"); + } + + for legacy in [ + "tcgen05_mma_ws_f16", + "tcgen05_mma_ws_bf16", + "tcgen05_mma_ws_tf32", + ] { + let old = &records[..27] + .iter() + .find(|record| record.id == legacy) + .unwrap(); + assert!(old.tcgen05.as_ref().unwrap().mma.is_none()); + assert!(old.abi_id.starts_with("i035")); + assert_eq!( + old.expected_ptx.operands.last(), + Some(&OperandPattern::Exact { + value: "%enable_pred".into(), + }) + ); + } + let legacy_f16 = records + .iter() + .find(|record| record.id == "tcgen05_mma_ws_f16") + .unwrap(); + assert!( + legacy_f16 + .expected_ptx + .matches("tcgen05.mma.ws.cta_group::1.kind::f16 [%r1], [%r2], %rd1, %r3, %enable_pred;") + .unwrap() + ); + assert!( + !legacy_f16 + .expected_ptx + .matches("tcgen05.mma.ws.cta_group::1.kind::f16 [%r1], [%r2], %rd1, %r3, %other_pred;") + .unwrap() + ); +} + +#[test] +fn compact_tcgen05_mma_admission_fails_closed_on_drift() { + let mut missing = test_tcgen05_mma_admission(); + missing.mma_variants.pop(); + assert!(expand_tcgen05_admission(&missing).is_err()); + + let mut reordered = test_tcgen05_mma_admission(); + reordered.mma_variants.swap(0, 1); + assert!(expand_tcgen05_admission(&reordered).is_err()); + + let mut non_contiguous_abi = test_tcgen05_mma_admission(); + non_contiguous_abi.mma_variants[0].abi_id = "i9999".into(); + let records = expand_tcgen05_admission(&non_contiguous_abi).unwrap(); + assert!(records.iter().any(|record| record.abi_id == "i9999")); + + let mut missing_target = test_tcgen05_mma_admission(); + missing_target.mma_llvm_target_contracts[0] + .alternatives + .pop(); + assert!(expand_tcgen05_admission(&missing_target).is_err()); + + let mut broadened_i8 = test_tcgen05_mma_admission(); + broadened_i8.mma_libnvvm_target_contracts[2] + .alternatives + .insert( + 1, + crate::model::TargetContractAlternative { + target: "sm_103a".into(), + minimum_ptx: "8.8".into(), + }, + ); + assert!(expand_tcgen05_admission(&broadened_i8).is_err()); + + let records = expand_tcgen05_admission(&test_tcgen05_mma_admission()).unwrap(); + let record = records + .iter() + .find(|record| record.id == "tcgen05_mma_tensor_ashift") + .unwrap(); + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let declaration = imported + .intrinsics + .iter() + .find(|declaration| { + Some(declaration.source_record.as_str()) == record.source_record.as_deref() + }) + .unwrap(); + + let mut wrong_layout = record.clone(); + let mma = wrong_layout.tcgen05.as_mut().unwrap().mma.as_mut().unwrap(); + mma.selector_layout = Tcgen05MmaSelectorLayout::Base { + kind_argument: 6, + cta_group_argument: 5, + collector_a_argument: 7, + collector_a_upper_exclusive: 2, + }; + assert!(validate_imported_policy(&wrong_layout, declaration).is_err()); + + let mut wrong_properties = declaration.clone(); + wrong_properties.properties.pop(); + assert!(validate_imported_policy(record, &wrong_properties).is_err()); + + let mut wrong_range = declaration.clone(); + let range = wrong_range + .properties + .iter_mut() + .find(|property| property.as_str() == "Range") + .unwrap(); + *range = "Range".into(); + assert!(validate_imported_policy(record, &wrong_range).is_err()); + + let mut wrong_predicate = declaration.clone(); + wrong_predicate.selections[0].predicates = vec!["hasSM100a".into()]; + assert!(validate_imported_policy(record, &wrong_predicate).is_err()); + + let mut wrong_memory = record.clone(); + wrong_memory.memory = "write".into(); + assert!(validate_imported_policy(&wrong_memory, declaration).is_err()); + + let mut wrong_convergence = record.clone(); + wrong_convergence.convergent = false; + assert!(validate_imported_policy(&wrong_convergence, declaration).is_err()); + + let mut split_dialect_op = record.clone(); + split_dialect_op.dialect_op_type = "Tcgen05MmaTensorAshiftOp".into(); + split_dialect_op.dialect_op_name = "nvvm.tcgen05_mma_tensor_ashift".into(); + assert!(validate_imported_policy(&split_dialect_op, declaration).is_err()); +} + +#[test] +fn compact_tcgen05_copy_admission_matches_all_llvm_records_and_fails_closed() { + let records = expand_tcgen05_admission(&test_tcgen05_cp_admission()).unwrap(); + assert_eq!(records.len(), 61); + let copies = &records[27..]; + assert_eq!(copies.len(), 34); + assert_eq!( + (copies[0].abi_id.as_str(), copies[0].id.as_str()), + ("i0578", "tcgen05_cp_128x128b_b4x16_p64") + ); + assert_eq!( + (copies[33].abi_id.as_str(), copies[33].id.as_str()), + ("i0611", "tcgen05_cp_64x128b_warpx2_02_13_cg2") + ); + + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let declarations = imported + .intrinsics + .iter() + .map(|record| (record.source_record.as_str(), record)) + .collect::>(); + for record in copies { + let declaration = declarations[record.source_record.as_deref().unwrap()]; + validate_imported_policy(record, declaration).unwrap(); + assert_tcgen05_backend_target_split(record); + assert_eq!(record.rust_arguments, ["u32", "u64"]); + assert_eq!(record.rust_result, "()"); + assert!(!record.safe); + assert!(record.convergent); + assert_eq!(record.memory, "read_write"); + assert_eq!(record.targets, TCGEN05_LLVM_TARGETS); + } + + let packed = copies + .iter() + .find(|record| record.id == "tcgen05_cp_128x128b_b4x16_p64") + .unwrap(); + assert_eq!( + packed.expected_ptx.modifiers, + ["cp", "cta_group::1", "128x128b", "b8x16", "b4x16_p64"] + ); + let warpx4 = copies + .iter() + .find(|record| record.id == "tcgen05_cp_32x128b_warpx4") + .unwrap(); + assert_eq!( + warpx4.expected_ptx.modifiers, + ["cp", "cta_group::1", "32x128b", "warpx4"] + ); + let pair_01_23 = copies + .iter() + .find(|record| record.id == "tcgen05_cp_64x128b_warpx2_01_23_b6x16_p32_cg2") + .unwrap(); + assert_eq!( + pair_01_23.expected_ptx.modifiers, + [ + "cp", + "cta_group::2", + "64x128b", + "warpx2::01_23", + "b8x16", + "b6x16_p32" + ] + ); + let pair_02_13 = copies + .iter() + .find(|record| record.id == "tcgen05_cp_64x128b_warpx2_02_13") + .unwrap(); + assert_eq!( + pair_02_13.expected_ptx.modifiers, + ["cp", "cta_group::1", "64x128b", "warpx2::02_13"] + ); + + let mut missing = test_tcgen05_cp_admission(); + missing.cp_variants.pop(); + assert!(expand_tcgen05_admission(&missing).is_err()); + + let mut reordered = test_tcgen05_cp_admission(); + reordered.cp_variants.swap(0, 1); + assert!(expand_tcgen05_admission(&reordered).is_err()); + + let mut non_contiguous_abi = test_tcgen05_cp_admission(); + non_contiguous_abi.cp_variants[0].abi_id = "i9999".into(); + let records = expand_tcgen05_admission(&non_contiguous_abi).unwrap(); + assert!(records.iter().any(|record| record.abi_id == "i9999")); + + let mut missing_evidence = test_tcgen05_cp_admission(); + missing_evidence.cp_llvm_evidence_profile = None; + assert!(expand_tcgen05_admission(&missing_evidence).is_err()); + + let declaration = declarations[copies[0].source_record.as_deref().unwrap()]; + let mut wrong_spelling = copies[0].clone(); + wrong_spelling.expected_ptx.modifiers.remove(3); + assert!(validate_imported_policy(&wrong_spelling, declaration).is_err()); + + let mut broadened_libnvvm = copies[0].clone(); + broadened_libnvvm.backend_lowerings[1].targets = None; + assert!(validate_imported_policy(&broadened_libnvvm, declaration).is_err()); +} + +#[test] +fn compact_tcgen05_load_admission_matches_all_llvm_records_and_fails_closed() { + let records = expand_tcgen05_admission(&test_tcgen05_ld_admission()).unwrap(); + assert_eq!(records.len(), 119); + let loads = &records[61..]; + assert_eq!(loads.len(), 58); + assert_eq!( + (loads[0].abi_id.as_str(), loads[0].id.as_str()), + ("i0612", "tcgen05_ld_16x64b_x1_raw") + ); + assert_eq!( + (loads[1].abi_id.as_str(), loads[1].id.as_str()), + ("i0613", "tcgen05_ld_16x64b_x1_pack16") + ); + assert_eq!( + (loads[57].abi_id.as_str(), loads[57].id.as_str()), + ("i0669", "tcgen05_ld_32x32b_x128_pack16") + ); + + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let declarations = imported + .intrinsics + .iter() + .map(|record| (record.source_record.as_str(), record)) + .collect::>(); + let mut source_counts = BTreeMap::new(); + for record in loads { + let source_record = record.source_record.as_deref().unwrap(); + *source_counts.entry(source_record).or_insert(0) += 1; + let declaration = declarations[source_record]; + validate_imported_policy(record, declaration).unwrap(); + assert_tcgen05_backend_target_split(record); + assert_eq!(record.rust_arguments, ["u32"]); + assert!(record.must_use); + assert!(!record.safe); + assert!(!record.pure); + assert_eq!(record.memory, "read"); + assert!(record.convergent); + assert_eq!(record.execution_scope, "warp"); + assert!( + !record + .rust_arguments + .iter() + .any(|argument| argument == "bool") + ); + assert_eq!(record.dialect_operands, ["i32"]); + assert_eq!(record.llvm_arguments, ["tmem_ptr", "i1"]); + assert!(declaration.selections.is_empty()); + } + assert_eq!(source_counts.len(), 29); + assert!(source_counts.values().all(|count| *count == 2)); + assert_eq!( + records + .iter() + .filter(|record| { + record.source_record.as_deref() == Some("int_nvvm_tcgen05_ld_16x256b_x1") + }) + .count(), + 3 + ); + assert_eq!( + records + .iter() + .filter(|record| { + record.source_record.as_deref() == Some("int_nvvm_tcgen05_ld_16x256b_x8") + }) + .count(), + 3 + ); + validate_unique_overlay(&records, 1).unwrap(); + + let scalar_raw = &loads[0]; + assert_eq!(scalar_raw.rust_result, "u32"); + assert_eq!(scalar_raw.dialect_results, ["i32"]); + assert_eq!(scalar_raw.llvm_results, ["i32"]); + assert_eq!( + scalar_raw.expected_ptx.modifiers, + ["ld", "sync", "aligned", "16x64b", "x1", "b32"] + ); + assert_eq!( + scalar_raw.expected_ptx.operands, + [ + OperandPattern::RegisterList { length: 1 }, + OperandPattern::Address, + ] + ); + assert!( + scalar_raw + .expected_ptx + .matches("tcgen05.ld.sync.aligned.16x64b.x1.b32 {%r1}, [%r2];") + .unwrap() + ); + assert!( + !scalar_raw + .expected_ptx + .matches("tcgen05.ld.sync.aligned.16x64b.x1.b32 %r1, [%r2];") + .unwrap() + ); + let scalar_pack = &loads[1]; + assert_eq!( + scalar_pack.expected_ptx.modifiers, + ["ld", "sync", "aligned", "16x64b", "x1", "pack::16b", "b32",] + ); + let largest = &loads[57]; + assert_eq!(largest.rust_result, "[u32; 128]"); + assert_eq!(largest.dialect_results.len(), 128); + assert_eq!(largest.llvm_results, ["v128i32"]); + assert_eq!( + largest.expected_ptx.operands, + [ + OperandPattern::RegisterList { length: 128 }, + OperandPattern::Address, + ] + ); + + let mut missing = test_tcgen05_ld_admission(); + missing.ld_variants.pop(); + assert!(expand_tcgen05_admission(&missing).is_err()); + + let mut reordered = test_tcgen05_ld_admission(); + reordered.ld_variants.swap(0, 1); + assert!(expand_tcgen05_admission(&reordered).is_err()); + + let mut non_contiguous_abi = test_tcgen05_ld_admission(); + non_contiguous_abi.ld_variants[0].abi_id = "i9999".into(); + let records = expand_tcgen05_admission(&non_contiguous_abi).unwrap(); + assert!(records.iter().any(|record| record.abi_id == "i9999")); + + let mut missing_evidence = test_tcgen05_ld_admission(); + missing_evidence.ld_llvm_evidence_profile = None; + assert!(expand_tcgen05_admission(&missing_evidence).is_err()); + + let declaration = declarations[loads[0].source_record.as_deref().unwrap()]; + let mut wrong_selector = loads[0].clone(); + wrong_selector + .tcgen05 + .as_mut() + .unwrap() + .ld + .as_mut() + .unwrap() + .pack16 = true; + assert!(validate_imported_policy(&wrong_selector, declaration).is_err()); + + let mut changed_declaration = declaration.clone(); + changed_declaration.properties.pop(); + assert!(validate_imported_policy(&loads[0], &changed_declaration).is_err()); + + let mut wrong_scope = loads[0].clone(); + wrong_scope.execution_scope = "thread".into(); + assert!(validate_imported_policy(&wrong_scope, declaration).is_err()); + + let mut unreviewed_sharing = records.clone(); + unreviewed_sharing[62].tcgen05.as_mut().unwrap().operation = Tcgen05Operation::Alloc; + assert!(validate_unique_overlay(&unreviewed_sharing, 1).is_err()); +} + +#[test] +fn compact_tcgen05_store_admission_matches_all_llvm_records_and_fails_closed() { + let records = expand_tcgen05_admission(&test_tcgen05_st_admission()).unwrap(); + assert_eq!(records.len(), 177); + let stores = &records[119..]; + assert_eq!(stores.len(), 58); + assert_eq!( + (stores[0].abi_id.as_str(), stores[0].id.as_str()), + ("i0670", "tcgen05_st_16x64b_x1_raw") + ); + assert_eq!( + (stores[1].abi_id.as_str(), stores[1].id.as_str()), + ("i0671", "tcgen05_st_16x64b_x1_unpack16") + ); + assert_eq!( + (stores[57].abi_id.as_str(), stores[57].id.as_str()), + ("i0727", "tcgen05_st_32x32b_x128_unpack16") + ); + + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let declarations = imported + .intrinsics + .iter() + .map(|record| (record.source_record.as_str(), record)) + .collect::>(); + let mut source_counts = BTreeMap::new(); + for record in stores { + let source_record = record.source_record.as_deref().unwrap(); + *source_counts.entry(source_record).or_insert(0) += 1; + let declaration = declarations[source_record]; + validate_imported_policy(record, declaration).unwrap(); + assert_tcgen05_backend_target_split(record); + assert_eq!(record.rust_arguments.len(), 2); + assert_eq!(record.rust_result, "()"); + assert!(!record.must_use); + assert!(!record.safe); + assert!(!record.pure); + assert_eq!(record.memory, "write"); + assert!(record.convergent); + assert_eq!(record.execution_scope, "warp"); + assert_eq!(record.dialect_results, Vec::::new()); + assert_eq!(record.llvm_results, Vec::::new()); + assert_eq!(record.llvm_arguments, declaration.arguments); + assert_eq!( + declaration.properties, + [ + "ImmArg", + "IntrArgMemOnly", + "IntrConvergent", + "NoCapture", + ] + ); + assert!(declaration.selections.is_empty()); + } + assert_eq!(source_counts.len(), 29); + assert!(source_counts.values().all(|count| *count == 2)); + validate_unique_overlay(&records, 1).unwrap(); + + let scalar_raw = &stores[0]; + assert_eq!(scalar_raw.rust_arguments, ["u32", "u32"]); + assert_eq!(scalar_raw.dialect_operands, ["i32", "i32"]); + assert_eq!(scalar_raw.llvm_arguments, ["tmem_ptr", "i32", "i1"]); + assert_eq!( + scalar_raw.expected_ptx.modifiers, + ["st", "sync", "aligned", "16x64b", "x1", "b32"] + ); + assert_eq!( + scalar_raw.expected_ptx.operands, + [ + OperandPattern::Address, + OperandPattern::RegisterList { length: 1 }, + ] + ); + assert!( + scalar_raw + .expected_ptx + .matches("tcgen05.st.sync.aligned.16x64b.x1.b32 [%r1], {%r2};") + .unwrap() + ); + assert!( + !scalar_raw + .expected_ptx + .matches("tcgen05.st.sync.aligned.16x64b.x1.b32 [%r1], %r2;") + .unwrap() + ); + assert_eq!( + stores[1].expected_ptx.modifiers, + [ + "st", + "sync", + "aligned", + "16x64b", + "x1", + "unpack::16b", + "b32", + ] + ); + let largest = &stores[57]; + assert_eq!(largest.rust_arguments, ["u32", "[u32; 128]"]); + assert_eq!(largest.dialect_operands.len(), 129); + assert_eq!(largest.llvm_arguments, ["tmem_ptr", "v128i32", "i1"]); + assert_eq!( + largest.expected_ptx.operands, + [ + OperandPattern::Address, + OperandPattern::RegisterList { length: 128 }, + ] + ); + + let mut missing = test_tcgen05_st_admission(); + missing.st_variants.pop(); + assert!(expand_tcgen05_admission(&missing).is_err()); + + let mut reordered = test_tcgen05_st_admission(); + reordered.st_variants.swap(0, 1); + assert!(expand_tcgen05_admission(&reordered).is_err()); + + let mut non_contiguous_abi = test_tcgen05_st_admission(); + non_contiguous_abi.st_variants[0].abi_id = "i9999".into(); + let records = expand_tcgen05_admission(&non_contiguous_abi).unwrap(); + assert!(records.iter().any(|record| record.abi_id == "i9999")); + + let mut missing_evidence = test_tcgen05_st_admission(); + missing_evidence.st_llvm_evidence_profile = None; + assert!(expand_tcgen05_admission(&missing_evidence).is_err()); + + let declaration = declarations[stores[0].source_record.as_deref().unwrap()]; + let mut wrong_selector = stores[0].clone(); + wrong_selector + .tcgen05 + .as_mut() + .unwrap() + .st + .as_mut() + .unwrap() + .unpack16 = true; + assert!(validate_imported_policy(&wrong_selector, declaration).is_err()); + + let mut changed_declaration = declaration.clone(); + changed_declaration.properties.pop(); + assert!(validate_imported_policy(&stores[0], &changed_declaration).is_err()); + + let mut wrong_scope = stores[0].clone(); + wrong_scope.execution_scope = "thread".into(); + assert!(validate_imported_policy(&wrong_scope, declaration).is_err()); + + let mut unreviewed_sharing = records.clone(); + unreviewed_sharing[120].tcgen05.as_mut().unwrap().operation = Tcgen05Operation::Alloc; + assert!(validate_unique_overlay(&unreviewed_sharing, 1).is_err()); +} + +#[test] +fn compact_tcgen05_offset_admission_is_exact_and_fails_closed() { + let records = expand_tcgen05_admission(&test_tcgen05_offset_admission()).unwrap(); + assert_eq!(records.len(), 209); + let loads = &records[177..193]; + let stores = &records[193..209]; + assert_eq!( + (loads[0].abi_id.as_str(), loads[0].id.as_str()), + ("i0728", "tcgen05_ld_16x32bx2_x1_raw") + ); + assert_eq!( + (loads[15].abi_id.as_str(), loads[15].id.as_str()), + ("i0743", "tcgen05_ld_16x32bx2_x128_pack16") + ); + assert_eq!( + (stores[0].abi_id.as_str(), stores[0].id.as_str()), + ("i0744", "tcgen05_st_16x32bx2_x1_raw") + ); + assert_eq!( + (stores[15].abi_id.as_str(), stores[15].id.as_str()), + ("i0759", "tcgen05_st_16x32bx2_x128_unpack16") + ); + + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let declarations = imported + .intrinsics + .iter() + .map(|record| (record.source_record.as_str(), record)) + .collect::>(); + let mut load_sources = BTreeMap::new(); + for record in loads { + let source = record.source_record.as_deref().unwrap(); + *load_sources.entry(source).or_insert(0) += 1; + let declaration = declarations[source]; + validate_imported_policy(record, declaration).unwrap(); + assert_tcgen05_backend_target_split(record); + assert_eq!(record.rust_arguments, ["u32", "i64"]); + assert_eq!(record.dialect_operands, ["i32", "i64"]); + assert_eq!(record.llvm_arguments, ["tmem_ptr", "i64", "i1"]); + assert_eq!(record.execution_scope, "warp"); + assert!(record.must_use); + assert_eq!( + record.compatibility_rust_paths, + [format!("cuda_device::tcgen05::__{}", record.id)] + ); + assert_eq!( + record.expected_ptx.operands.last(), + Some(&OperandPattern::Immediate) + ); + assert_eq!( + declaration.properties, + [ + "ImmArg", + "ImmArg", + "IntrArgMemOnly", + "IntrConvergent", + "NoCapture", + ] + ); + } + assert_eq!(load_sources.len(), 8); + assert!(load_sources.values().all(|count| *count == 2)); + + let mut store_sources = BTreeMap::new(); + for record in stores { + let source = record.source_record.as_deref().unwrap(); + *store_sources.entry(source).or_insert(0) += 1; + let declaration = declarations[source]; + validate_imported_policy(record, declaration).unwrap(); + assert_tcgen05_backend_target_split(record); + assert_eq!(record.rust_arguments[0..2], ["u32", "i64"]); + assert_eq!(record.dialect_operands[0..2], ["i32", "i64"]); + assert_eq!(record.llvm_arguments[0..2], ["tmem_ptr", "i64"]); + assert_eq!(record.execution_scope, "warp"); + assert!(!record.must_use); + assert_eq!( + record.compatibility_rust_paths, + [format!("cuda_device::tcgen05::__{}", record.id)] + ); + assert_eq!( + record.expected_ptx.operands.get(1), + Some(&OperandPattern::Immediate) + ); + assert_eq!( + declaration.properties, + [ + "ImmArg", + "ImmArg", + "IntrArgMemOnly", + "IntrConvergent", + "NoCapture", + ] + ); + } + assert_eq!(store_sources.len(), 8); + assert!(store_sources.values().all(|count| *count == 2)); + validate_unique_overlay(&records, 1).unwrap(); + + let scalar_pack = &loads[1]; + assert_eq!( + scalar_pack.expected_ptx.operands, + [ + OperandPattern::RegisterList { length: 1 }, + OperandPattern::Address, + OperandPattern::Immediate, + ] + ); + assert!( + scalar_pack + .expected_ptx + .matches("tcgen05.ld.sync.aligned.16x32bx2.x1.pack::16b.b32 {%r1}, [%r2], 16;") + .unwrap() + ); + + let scalar_unpack = &stores[1]; + assert_eq!( + scalar_unpack.expected_ptx.operands, + [ + OperandPattern::Address, + OperandPattern::Immediate, + OperandPattern::RegisterList { length: 1 }, + ] + ); + assert!( + scalar_unpack + .expected_ptx + .matches("tcgen05.st.sync.aligned.16x32bx2.x1.unpack::16b.b32 [%r1], 16, {%r2};") + .unwrap() + ); + + let mut missing = test_tcgen05_offset_admission(); + missing.ld_offset_variants.pop(); + assert!(expand_tcgen05_admission(&missing).is_err()); + + let mut reordered = test_tcgen05_offset_admission(); + reordered.st_offset_variants.swap(0, 1); + assert!(expand_tcgen05_admission(&reordered).is_err()); + + let mut non_contiguous_abi = test_tcgen05_offset_admission(); + non_contiguous_abi.ld_offset_variants[0].abi_id = "i9999".into(); + let records = expand_tcgen05_admission(&non_contiguous_abi).unwrap(); + assert!(records.iter().any(|record| record.abi_id == "i9999")); + + let mut missing_evidence = test_tcgen05_offset_admission(); + missing_evidence.offset_llvm_evidence_profile = None; + assert!(expand_tcgen05_admission(&missing_evidence).is_err()); + + let load_declaration = declarations[loads[0].source_record.as_deref().unwrap()]; + let mut wrong_offset_type = loads[0].clone(); + wrong_offset_type.rust_arguments[1] = "i32".into(); + assert!(validate_imported_policy(&wrong_offset_type, load_declaration).is_err()); + + let mut wrong_scope = loads[0].clone(); + wrong_scope.execution_scope = "thread".into(); + assert!(validate_imported_policy(&wrong_scope, load_declaration).is_err()); + + let mut missing_immediate = loads[0].clone(); + missing_immediate.expected_ptx.operands.pop(); + assert!(validate_imported_policy(&missing_immediate, load_declaration).is_err()); + + let store_declaration = declarations[stores[0].source_record.as_deref().unwrap()]; + let mut wrong_adapter = stores[0].clone(); + wrong_adapter.tcgen05.as_mut().unwrap().adapter = + Tcgen05Adapter::TmemU32RegistersInjectUnpack16ToVoid; + assert!(validate_imported_policy(&wrong_adapter, store_declaration).is_err()); + + let mut changed_declaration = store_declaration.clone(); + changed_declaration.properties.remove(0); + assert!(validate_imported_policy(&stores[0], &changed_declaration).is_err()); +} + +#[test] +fn tcgen05_compact_schema_is_reserved_for_aggregation() { + let shard = |schema, admission| OverlayShardFile { + schema, + family: "tcgen05".into(), + intrinsics: vec![], + register_mma_int4: None, + register_mma_int8: None, + register_mma_b1: None, + register_mma_f8f6f4_f32: None, + register_mma_f8f6f4_f16: None, + register_mma_mxf8f6f4_f32: None, + register_mma_fp8: None, + register_mma_ampere_float: None, + sparse_mma_integer: None, + sparse_mma_f8f6f4_f32: None, + sparse_mma_f8f6f4_f16: None, + prmt: None, + packed_conversion_fp8: None, + packed_conversion_fp8_f16x2: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cluster_sreg: None, + cluster_barrier: None, + mbarrier_extended: None, + special_registers: None, + debug_control: None, + threadfence: None, + cluster_memory: None, + stmatrix: None, + clc: None, + wgmma_controls: None, + tma: None, + tcgen05: Some(admission), + }; + let path = Path::new("intrinsics/overlay/tcgen05.toml"); + validate_overlay_shard_schema_with_max( + &shard( + TCGEN05_SHARD_SCHEMA, + without_tcgen05_control(test_tcgen05_admission()), + ), + path, + TCGEN05_CP_SHARD_SCHEMA, + ) + .unwrap(); + assert!( + validate_overlay_shard_schema_with_max( + &shard( + TCGEN05_SHARD_SCHEMA - 1, + without_tcgen05_control(test_tcgen05_admission()), + ), + path, + TCGEN05_CP_SHARD_SCHEMA, + ) + .unwrap_err() + .to_string() + .contains("requires overlay shard schema 42") + ); + validate_overlay_shard_schema_with_max( + &shard( + TCGEN05_CP_SHARD_SCHEMA, + without_tcgen05_control(test_tcgen05_cp_admission()), + ), + path, + TCGEN05_CP_SHARD_SCHEMA, + ) + .unwrap(); + assert!( + validate_overlay_shard_schema_with_max( + &shard( + TCGEN05_CP_SHARD_SCHEMA - 1, + without_tcgen05_control(test_tcgen05_cp_admission()), + ), + path, + TCGEN05_CP_SHARD_SCHEMA, + ) + .unwrap_err() + .to_string() + .contains("requires overlay shard schema 52") + ); + validate_overlay_shard_schema_with_max( + &shard( + TCGEN05_LD_SHARD_SCHEMA, + without_tcgen05_control(test_tcgen05_ld_admission()), + ), + path, + TCGEN05_LD_SHARD_SCHEMA, + ) + .unwrap(); + assert!( + validate_overlay_shard_schema_with_max( + &shard( + TCGEN05_LD_SHARD_SCHEMA - 1, + without_tcgen05_control(test_tcgen05_ld_admission()), + ), + path, + TCGEN05_LD_SHARD_SCHEMA, + ) + .unwrap_err() + .to_string() + .contains("requires overlay shard schema 53") + ); + validate_overlay_shard_schema_with_max( + &shard( + TCGEN05_ST_SHARD_SCHEMA, + without_tcgen05_control(test_tcgen05_st_admission()), + ), + path, + TCGEN05_ST_SHARD_SCHEMA, + ) + .unwrap(); + assert!( + validate_overlay_shard_schema_with_max( + &shard( + TCGEN05_ST_SHARD_SCHEMA - 1, + without_tcgen05_control(test_tcgen05_st_admission()), + ), + path, + TCGEN05_ST_SHARD_SCHEMA, + ) + .unwrap_err() + .to_string() + .contains("requires overlay shard schema 54") + ); + validate_overlay_shard_schema_with_max( + &shard( + TCGEN05_OFFSET_LDST_SHARD_SCHEMA, + without_tcgen05_control(test_tcgen05_offset_admission()), + ), + path, + TCGEN05_OFFSET_LDST_SHARD_SCHEMA, + ) + .unwrap(); + assert!( + validate_overlay_shard_schema_with_max( + &shard( + TCGEN05_OFFSET_LDST_SHARD_SCHEMA - 1, + without_tcgen05_control(test_tcgen05_offset_admission()), + ), + path, + TCGEN05_OFFSET_LDST_SHARD_SCHEMA, + ) + .unwrap_err() + .to_string() + .contains("requires overlay shard schema 55") + ); + validate_overlay_shard_schema_with_max( + &shard( + TCGEN05_CONTROL_SHARD_SCHEMA, + test_tcgen05_offset_admission(), + ), + path, + TCGEN05_CONTROL_SHARD_SCHEMA, + ) + .unwrap(); + assert!( + validate_overlay_shard_schema_with_max( + &shard( + TCGEN05_CONTROL_SHARD_SCHEMA, + without_tcgen05_control(test_tcgen05_offset_admission()), + ), + path, + TCGEN05_CONTROL_SHARD_SCHEMA, + ) + .unwrap_err() + .to_string() + .contains("requires all three control variants and both backend evidence profiles") + ); + assert!( + validate_overlay_shard_schema_with_max( + &shard( + TCGEN05_CONTROL_SHARD_SCHEMA - 1, + test_tcgen05_offset_admission(), + ), + path, + TCGEN05_CONTROL_SHARD_SCHEMA, + ) + .unwrap_err() + .to_string() + .contains("requires overlay shard schema 56") + ); + validate_overlay_shard_schema_with_max( + &shard(TCGEN05_MMA_SHARD_SCHEMA, test_tcgen05_mma_admission()), + path, + TCGEN05_MMA_SHARD_SCHEMA, + ) + .unwrap(); + assert!( + validate_overlay_shard_schema_with_max( + &shard(TCGEN05_MMA_SHARD_SCHEMA - 1, test_tcgen05_mma_admission(),), + path, + TCGEN05_MMA_SHARD_SCHEMA, + ) + .unwrap_err() + .to_string() + .contains("requires overlay shard schema 57") + ); +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/tests/tma_clc_cluster.rs b/crates/cuda-intrinsics-gen/src/resolve/tests/tma_clc_cluster.rs new file mode 100644 index 0000000000..9301eeb28c --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/tests/tma_clc_cluster.rs @@ -0,0 +1,680 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, CatalogHardwareAlternative, CatalogHardwareTarget, ClcAdapter, + ClusterBarrierMode, ClusterBarrierOrdering, ClusterMemoryAdapter, ClusterMemoryOperation, + EvidenceStageKind, ImportedFile, IntrinsicBackend, IntrinsicSource, OverlayIntrinsic, + OverlayShardFile, RuntimeValidation, TmaAdapter, TmaOperation, +}; +use crate::util::read_json; +use std::collections::{BTreeMap, BTreeSet}; +use std::path::Path; + +use super::fixtures::*; +use crate::resolve::evidence::*; +use crate::resolve::families::*; +use crate::resolve::overlay::*; +use crate::resolve::policy::*; +use crate::resolve::targets::*; + +#[test] +fn compact_clc_admission_matches_llvm_and_fails_closed() { + let records = expand_clc_admission(&test_clc_admission()).unwrap(); + assert_eq!(records.len(), 6); + assert_eq!( + records + .iter() + .map(|record| (record.abi_id.as_str(), record.id.as_str())) + .collect::>(), + [ + ("i0322", "clc_try_cancel"), + ("i0323", "clc_try_cancel_multicast"), + ("i0324", "clc_query_is_canceled"), + ("i0325", "clc_query_get_first_ctaid_x"), + ("i0326", "clc_query_get_first_ctaid_y"), + ("i0327", "clc_query_get_first_ctaid_z"), + ] + ); + + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let declarations = imported + .intrinsics + .iter() + .map(|record| (record.source_record.as_str(), record)) + .collect::>(); + for record in &records { + let declaration = declarations[record.source_record.as_deref().unwrap()]; + validate_imported_policy(record, declaration).unwrap(); + } + + assert_eq!( + parse_hardware_target(&records[1]).unwrap(), + CatalogHardwareTarget::AnyOf { + alternatives: vec![ + CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 101 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 103 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 120 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 121 }, + ], + } + ); + + let mut missing = test_clc_admission(); + missing.variants.pop(); + assert!(expand_clc_admission(&missing).is_err()); + + let mut reordered = test_clc_admission(); + reordered.variants.swap(0, 1); + assert!(expand_clc_admission(&reordered).is_err()); + + let mut wrong_abi = test_clc_admission(); + wrong_abi.variants[0].abi_id = "i9999".into(); + assert!(expand_clc_admission(&wrong_abi).is_err()); + + let mut executed = test_clc_admission(); + executed.runtime_validation = RuntimeValidation::Executed; + assert!(expand_clc_admission(&executed).is_err()); + + let declaration = declarations[records[2].source_record.as_deref().unwrap()]; + let mut wrong_adapter = records[2].clone(); + wrong_adapter.clc.as_mut().unwrap().adapter = ClcAdapter::PairU64ToI128U32; + assert!(validate_imported_policy(&wrong_adapter, declaration).is_err()); + + let mut unsorted_targets = records[1].clone(); + unsorted_targets.targets = "sm_120a|sm_100a".into(); + assert!(parse_hardware_target(&unsorted_targets).is_err()); + + let mut duplicate_targets = records[1].clone(); + duplicate_targets.targets = "sm_100a|sm_100a".into(); + assert!(parse_hardware_target(&duplicate_targets).is_err()); +} + +#[test] +fn clc_compact_schema_is_reserved_for_aggregation() { + let shard = |schema| OverlayShardFile { + schema, + family: "clc".into(), + intrinsics: vec![], + register_mma_int4: None, + register_mma_int8: None, + register_mma_b1: None, + register_mma_f8f6f4_f32: None, + register_mma_f8f6f4_f16: None, + register_mma_mxf8f6f4_f32: None, + register_mma_fp8: None, + register_mma_ampere_float: None, + sparse_mma_integer: None, + sparse_mma_f8f6f4_f32: None, + sparse_mma_f8f6f4_f16: None, + prmt: None, + packed_conversion_fp8: None, + packed_conversion_fp8_f16x2: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cluster_sreg: None, + cluster_barrier: None, + mbarrier_extended: None, + special_registers: None, + debug_control: None, + threadfence: None, + cluster_memory: None, + stmatrix: None, + clc: Some(test_clc_admission()), + wgmma_controls: None, + tma: None, + tcgen05: None, + }; + let path = Path::new("intrinsics/overlay/clc.toml"); + validate_overlay_shard_schema_with_max(&shard(CLC_SHARD_SCHEMA), path, CLC_SHARD_SCHEMA) + .unwrap(); + assert!( + validate_overlay_shard_schema_with_max( + &shard(CLC_SHARD_SCHEMA - 1), + path, + CLC_SHARD_SCHEMA, + ) + .unwrap_err() + .to_string() + .contains("requires overlay shard schema 40") + ); +} + +#[test] +fn compact_tma_admission_matches_llvm_and_fails_closed() { + let records = expand_tma_admission(&test_tma_admission()).unwrap(); + assert_eq!( + records.len(), + TMA_OPERATIONS.len() + tma_reduction_matrix().len() + ); + assert!(records.iter().take(TMA_OPERATIONS.len()).all(|record| { + let operation = record.tma.as_ref().unwrap().operation; + let recipe = tma_recipe(operation); + record.backend_lowerings.iter().any(|route| { + route.backend == IntrinsicBackend::LlvmNvptx && route.mechanism == recipe.llvm_mechanism + }) && record.backend_lowerings.iter().any(|route| { + route.backend == IntrinsicBackend::LibNvvm + && route.mechanism == BackendLoweringMechanism::InlinePtx + }) + })); + assert_eq!( + records + .iter() + .filter(|record| { + record + .tma + .as_ref() + .unwrap() + .operation + .prefetch_coordinate_count() + .is_some() + }) + .count(), + 12, + "the tile-prefetch subfamily must contain every plain/cache-hint pair" + ); + assert_eq!( + records + .iter() + .take(TMA_OPERATIONS.len()) + .map(|record| (record.abi_id.as_str(), record.id.as_str())) + .collect::>(), + [ + ("i0328", "cp_async_bulk_tensor_1d_g2s"), + ("i0329", "cp_async_bulk_tensor_2d_g2s"), + ("i0330", "cp_async_bulk_tensor_2d_g2s_multicast"), + ("i0331", "cp_async_bulk_tensor_2d_g2s_multicast_cg2"), + ("i0332", "cp_async_bulk_tensor_3d_g2s"), + ("i0333", "cp_async_bulk_tensor_4d_g2s"), + ("i0334", "cp_async_bulk_tensor_5d_g2s"), + ("i0335", "cp_async_bulk_tensor_1d_s2g"), + ("i0336", "cp_async_bulk_tensor_2d_s2g"), + ("i0337", "cp_async_bulk_tensor_3d_s2g"), + ("i0338", "cp_async_bulk_tensor_4d_s2g"), + ("i0339", "cp_async_bulk_tensor_5d_s2g"), + ("i0340", "cp_async_bulk_commit_group"), + ("i0341", "cp_async_bulk_wait_group"), + ("i0342", "cp_async_bulk_wait_group_read"), + ("i0887", "prefetch_tma_descriptor"), + ("i0888", "cp_async_bulk_prefetch_tensor_1d_l2"), + ("i0889", "cp_async_bulk_prefetch_tensor_2d_l2"), + ("i0890", "cp_async_bulk_prefetch_tensor_3d_l2"), + ("i0891", "cp_async_bulk_prefetch_tensor_4d_l2"), + ("i0892", "cp_async_bulk_prefetch_tensor_5d_l2"), + ("i0893", "cp_async_bulk_prefetch_tensor_gather4_2d_l2"), + ("i0894", "tensormap_replace_box_dim"), + ("i0895", "tensormap_replace_element_stride"), + ("i0896", "tensormap_replace_element_type"), + ("i0897", "tensormap_replace_fill_mode"), + ("i0898", "tensormap_replace_global_address"), + ("i0899", "tensormap_replace_global_dim"), + ("i0900", "tensormap_replace_global_stride"), + ("i0901", "tensormap_replace_interleave_layout"), + ("i0902", "tensormap_replace_rank"), + ("i0903", "tensormap_replace_swizzle_atomicity"), + ("i0904", "tensormap_replace_swizzle_mode"), + ("i0905", "fence_proxy_tensormap_generic_acquire_cluster"), + ("i0906", "fence_proxy_tensormap_generic_acquire_cta"), + ("i0907", "fence_proxy_tensormap_generic_acquire_gpu"), + ("i0908", "fence_proxy_tensormap_generic_acquire_system"), + ("i0909", "fence_proxy_tensormap_generic_release_cluster"), + ("i0910", "fence_proxy_tensormap_generic_release_cta"), + ("i0911", "fence_proxy_tensormap_generic_release_gpu"), + ("i0912", "fence_proxy_tensormap_generic_release_system"), + ("i0917", "cp_async_bulk_prefetch_tensor_1d_l2_cache_hint"), + ("i0918", "cp_async_bulk_prefetch_tensor_2d_l2_cache_hint"), + ("i0919", "cp_async_bulk_prefetch_tensor_3d_l2_cache_hint"), + ("i0920", "cp_async_bulk_prefetch_tensor_4d_l2_cache_hint"), + ("i0921", "cp_async_bulk_prefetch_tensor_5d_l2_cache_hint"), + ( + "i0922", + "cp_async_bulk_prefetch_tensor_gather4_2d_l2_cache_hint" + ), + ] + ); + + let reductions = &records[TMA_OPERATIONS.len()..]; + assert_eq!(reductions.len(), 64); + assert_eq!( + (reductions[0].abi_id.as_str(), reductions[0].id.as_str()), + ("i0923", "cp_async_bulk_tensor_reduce_add_tile_1d") + ); + assert_eq!( + (reductions[7].abi_id.as_str(), reductions[7].id.as_str()), + ("i0930", "cp_async_bulk_tensor_reduce_add_im2col_5d") + ); + assert_eq!( + (reductions[63].abi_id.as_str(), reductions[63].id.as_str()), + ("i0986", "cp_async_bulk_tensor_reduce_xor_im2col_5d") + ); + assert!(reductions.iter().all(|record| { + record.tma.as_ref().is_some_and(|tma| { + tma.operation == TmaOperation::Reduce + && tma.reduction.is_some() + && tma.adapter == TmaAdapter::ReductionPointersCoordinatesInjectDefaults + }) + })); + + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let declarations = imported + .intrinsics + .iter() + .map(|record| (record.source_record.as_str(), record)) + .collect::>(); + for record in &records { + let declaration = declarations[record.source_record.as_deref().unwrap()]; + validate_imported_policy(record, declaration).unwrap(); + } + + assert_eq!( + parse_hardware_target(&records[3]).unwrap(), + CatalogHardwareTarget::AnyOf { + alternatives: vec![ + CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 101 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 103 }, + CatalogHardwareAlternative::ExactArchitecture { sm: 110 }, + ], + } + ); + + let mut missing = test_tma_admission(); + missing.variants.pop(); + assert!(expand_tma_admission(&missing).is_err()); + + let mut reordered = test_tma_admission(); + reordered.variants.swap(0, 1); + assert!(expand_tma_admission(&reordered).is_err()); + + let mut wrong_abi = test_tma_admission(); + wrong_abi.variants[0].abi_id = "i9999".into(); + assert!(expand_tma_admission(&wrong_abi).is_err()); + + let mut missing_reduction = test_tma_admission(); + missing_reduction.reduce_variants.pop(); + assert!(expand_tma_admission(&missing_reduction).is_err()); + + let mut reordered_reduction = test_tma_admission(); + reordered_reduction.reduce_variants.swap(0, 1); + assert!(expand_tma_admission(&reordered_reduction).is_err()); + + let mut non_contiguous_reduction_abi = test_tma_admission(); + non_contiguous_reduction_abi.reduce_variants[0].abi_id = "i9999".into(); + let records = expand_tma_admission(&non_contiguous_reduction_abi).unwrap(); + assert_eq!(records[TMA_OPERATIONS.len()].abi_id, "i9999"); + + let mut executed = test_tma_admission(); + executed.runtime_validation = RuntimeValidation::Executed; + assert!(expand_tma_admission(&executed).is_err()); + + let declaration = declarations[records[0].source_record.as_deref().unwrap()]; + let mut wrong_adapter = records[0].clone(); + wrong_adapter.tma.as_mut().unwrap().adapter = TmaAdapter::NoOperands; + assert!(validate_imported_policy(&wrong_adapter, declaration).is_err()); +} + +#[test] +fn tma_compact_schema_is_reserved_for_aggregation() { + let shard = |schema: u32, include_reductions: bool| { + let mut admission = test_tma_admission(); + if !include_reductions { + admission.reduce_variants.clear(); + } + OverlayShardFile { + schema, + family: "tma".into(), + intrinsics: vec![], + register_mma_int4: None, + register_mma_int8: None, + register_mma_b1: None, + register_mma_f8f6f4_f32: None, + register_mma_f8f6f4_f16: None, + register_mma_mxf8f6f4_f32: None, + register_mma_fp8: None, + register_mma_ampere_float: None, + sparse_mma_integer: None, + sparse_mma_f8f6f4_f32: None, + sparse_mma_f8f6f4_f16: None, + prmt: None, + packed_conversion_fp8: None, + packed_conversion_fp8_f16x2: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cluster_sreg: None, + cluster_barrier: None, + mbarrier_extended: None, + special_registers: None, + debug_control: None, + threadfence: None, + cluster_memory: None, + stmatrix: None, + clc: None, + wgmma_controls: None, + tma: Some(admission), + tcgen05: None, + } + }; + let path = Path::new("intrinsics/overlay/tma.toml"); + + validate_overlay_shard_schema_with_max( + &shard(TMA_SHARD_SCHEMA, false), + path, + TMA_REDUCTION_SHARD_SCHEMA, + ) + .unwrap(); + assert!( + validate_overlay_shard_schema_with_max( + &shard(TMA_SHARD_SCHEMA - 1, false), + path, + TMA_REDUCTION_SHARD_SCHEMA, + ) + .unwrap_err() + .to_string() + .contains("compact TMA admission requires overlay shard schema 61") + ); + + validate_overlay_shard_schema_with_max( + &shard(TMA_REDUCTION_SHARD_SCHEMA, true), + path, + TMA_REDUCTION_SHARD_SCHEMA, + ) + .unwrap(); + assert!( + validate_overlay_shard_schema_with_max( + &shard(TMA_REDUCTION_SHARD_SCHEMA - 1, true), + path, + TMA_REDUCTION_SHARD_SCHEMA, + ) + .unwrap_err() + .to_string() + .contains("compact TMA reduction admission requires overlay shard schema 62") + ); +} + +#[test] +fn cluster_memory_compact_schema_is_reserved_and_fail_closed() { + let shard = |schema| OverlayShardFile { + schema, + family: "cluster_memory".into(), + intrinsics: vec![], + register_mma_int4: None, + register_mma_int8: None, + register_mma_b1: None, + register_mma_f8f6f4_f32: None, + register_mma_f8f6f4_f16: None, + register_mma_mxf8f6f4_f32: None, + register_mma_fp8: None, + register_mma_ampere_float: None, + sparse_mma_integer: None, + sparse_mma_f8f6f4_f32: None, + sparse_mma_f8f6f4_f16: None, + prmt: None, + packed_conversion_fp8: None, + packed_conversion_fp8_f16x2: None, + scalar_conversion: None, + scalar_arithmetic: None, + scalar_math: None, + extended_minmax: None, + cluster_sreg: None, + cluster_barrier: None, + mbarrier_extended: None, + special_registers: None, + debug_control: None, + threadfence: None, + cluster_memory: Some(test_cluster_memory_admission()), + stmatrix: None, + clc: None, + wgmma_controls: None, + tma: None, + tcgen05: None, + }; + let path = Path::new("intrinsics/overlay/cluster_memory.toml"); + validate_overlay_shard_schema_with_max( + &shard(CLUSTER_MEMORY_SHARD_SCHEMA), + path, + CLUSTER_MEMORY_SHARD_SCHEMA, + ) + .unwrap(); + assert!( + validate_overlay_shard_schema_with_max( + &shard(CLUSTER_MEMORY_SHARD_SCHEMA), + path, + CLUSTER_MEMORY_SHARD_SCHEMA - 1, + ) + .is_err() + ); + let error = validate_overlay_shard_schema_with_max( + &shard(CLUSTER_MEMORY_SHARD_SCHEMA - 1), + path, + CLUSTER_MEMORY_SHARD_SCHEMA, + ) + .unwrap_err(); + assert!( + error + .to_string() + .contains("requires overlay shard schema 39") + ); +} + +#[test] +fn cluster_memory_admission_preserves_mapa_identity_and_ptx_native_read() { + let admission = test_cluster_memory_admission(); + let records = expand_cluster_memory_admission(&admission).unwrap(); + assert_eq!(records.len(), 2); + assert_eq!(records[0].id, "map_shared_rank"); + assert_eq!(records[1].id, "dsmem_read_u32"); + assert_eq!(records[0].abi_id, "i0320"); + assert_eq!(records[1].abi_id, "i0321"); + assert_eq!( + cluster_memory_inline_recipe(ClusterMemoryOperation::MapSharedRank), + ("mapa.shared::cluster.u64 $0, $1, $2;", "=l,l,r") + ); + assert_eq!( + cluster_memory_inline_recipe(ClusterMemoryOperation::ReadU32), + ( + "{ .reg .u64 %mapped; mapa.shared::cluster.u64 %mapped, $1, $2; ld.shared::cluster.u32 $0, [%mapped]; }", + "=r,l,r,~{memory}" + ) + ); + + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let mapa = &records[0]; + let declaration = imported + .intrinsics + .iter() + .find(|declaration| declaration.source_record == "int_nvvm_mapa_shared_cluster") + .unwrap(); + assert_eq!(declaration.arguments, ["shared_ptr", "i32"]); + assert_eq!(declaration.results, ["shared_cluster_ptr"]); + assert_eq!( + declaration.properties, + ["IntrNoMem", "IntrSpeculatable", "NoCapture"] + ); + validate_imported_policy(mapa, declaration).unwrap(); + assert_eq!( + declaration + .selections + .iter() + .filter(|selection| selection.asm.starts_with("mapa.shared::cluster.u64")) + .map(|selection| selection.source_record.as_str()) + .collect::>(), + BTreeSet::from(["mapa_shared_cluster_64", "mapa_shared_cluster_64i"]) + ); + + let read = &records[1]; + validate_ptx_native_policy(read).unwrap(); + assert!(read.source_record.is_none()); + assert!(read.llvm_symbol.is_none()); + assert!(matches!( + resolve_policy_source(read).unwrap(), + IntrinsicSource::PtxNative { .. } + )); + assert_eq!(read.memory, "read"); + + let mut wrong_adapter = mapa.clone(); + wrong_adapter.cluster_memory.as_mut().unwrap().adapter = + ClusterMemoryAdapter::ConstU32PointerRankToU32; + assert!(validate_imported_policy(&wrong_adapter, declaration).is_err()); + + let mut typed_as3 = mapa.clone(); + typed_as3.llvm_results = vec!["shared_ptr".into()]; + assert!(validate_imported_policy(&typed_as3, declaration).is_err()); + + let mut wrong_route = mapa.clone(); + wrong_route.backend_lowerings[0].mechanism = BackendLoweringMechanism::TypedNvvm; + assert!(validate_imported_policy(&wrong_route, declaration).is_err()); + + let mut wrong_floor = read.clone(); + wrong_floor.minimum_sm = Some("sm_80".into()); + assert!(validate_ptx_native_policy(&wrong_floor).is_err()); + + let mut missing = admission.clone(); + missing.variants.pop(); + assert!(expand_cluster_memory_admission(&missing).is_err()); + + let mut duplicate = admission.clone(); + duplicate.variants[1].operation = ClusterMemoryOperation::MapSharedRank; + assert!(expand_cluster_memory_admission(&duplicate).is_err()); + + let mut wrong_abi = admission.clone(); + wrong_abi.variants[0].abi_id = "i9999".into(); + assert!(expand_cluster_memory_admission(&wrong_abi).is_err()); + + let mut executed = admission; + executed.runtime_validation = RuntimeValidation::Executed; + assert!(expand_cluster_memory_admission(&executed).is_err()); +} + +#[test] +fn compact_cluster_barrier_admission_and_semantics_fail_closed() { + let records = expand_cluster_barrier_admission(&test_cluster_barrier_admission()).unwrap(); + assert_eq!(records.len(), 6); + + let imported: ImportedFile = read_json( + &Path::new(env!("CARGO_MANIFEST_DIR")) + .join("../..") + .join("intrinsics/imported.json"), + ) + .unwrap(); + for record in &records { + let declaration = imported + .intrinsics + .iter() + .find(|declaration| { + Some(declaration.source_record.as_str()) == record.source_record.as_deref() + }) + .unwrap(); + validate_imported_policy(record, declaration).unwrap(); + } + + let declaration_for = |record: &OverlayIntrinsic| { + imported + .intrinsics + .iter() + .find(|declaration| { + Some(declaration.source_record.as_str()) == record.source_record.as_deref() + }) + .unwrap() + }; + let base = records + .iter() + .find(|record| { + record + .cluster_barrier + .as_ref() + .is_some_and(|barrier| barrier.mode == ClusterBarrierMode::ArriveAligned) + }) + .unwrap(); + + let mut wrong_mode = base.clone(); + wrong_mode.cluster_barrier.as_mut().unwrap().mode = ClusterBarrierMode::WaitAligned; + assert!(validate_imported_policy(&wrong_mode, declaration_for(base)).is_err()); + + let mut wrong_order = base.clone(); + wrong_order.cluster_barrier.as_mut().unwrap().ordering = ClusterBarrierOrdering::Relaxed; + assert!(validate_imported_policy(&wrong_order, declaration_for(base)).is_err()); + + let mut wrong_alignment = base.clone(); + wrong_alignment.cluster_barrier.as_mut().unwrap().aligned = false; + assert!(validate_imported_policy(&wrong_alignment, declaration_for(base)).is_err()); + + let mut missing = test_cluster_barrier_admission(); + missing.variants.pop(); + assert!(expand_cluster_barrier_admission(&missing).is_err()); + + let mut duplicate = test_cluster_barrier_admission(); + duplicate.variants[5].mode = ClusterBarrierMode::Arrive; + assert!(expand_cluster_barrier_admission(&duplicate).is_err()); + + let mut wrong_abi = test_cluster_barrier_admission(); + wrong_abi.variants[0].abi_id = "i9999".into(); + assert!(expand_cluster_barrier_admission(&wrong_abi).is_err()); +} + +#[test] +fn cluster_barrier_evidence_validates_both_backend_routes() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let mut admission = test_cluster_barrier_admission(); + admission.llvm_evidence_profile = "rust-llvm-22.1.2-1cb4e383".into(); + admission.libnvvm_evidence_profile = "cuda-13.3-libnvvm-13.3.33-cluster-barrier".into(); + let policies = expand_cluster_barrier_admission(&admission).unwrap(); + let evidence_files = vec![ + read_evidence_file( + &repo_root.join("intrinsics/evidence/rust-llvm-22.1.2-1cb4e383-cluster-barrier.json"), + ) + .unwrap(), + read_evidence_file( + &repo_root.join("intrinsics/evidence/cuda-13.3-libnvvm-13.3.33-cluster-barrier.json"), + ) + .unwrap(), + ]; + let indexed = + index_evidence(&evidence_files, "1cb4e3833c1919c2e6fb579a23ac0e2b22587b7e").unwrap(); + + for policy in &policies { + for lowering in &policy.backend_lowerings { + let evidence = indexed + .get(&(lowering.evidence_profile.as_str(), policy.id.as_str())) + .unwrap(); + validate_evidence(policy, evidence, Some(lowering)).unwrap(); + } + } + + let mut missing_typed_failure = evidence_files.clone(); + let libnvvm = missing_typed_failure + .iter_mut() + .find(|file| file.backend_kind == Some(IntrinsicBackend::LibNvvm)) + .unwrap(); + for record in &mut libnvvm.records { + record.stages.retain(|stage| { + stage.mechanism != Some(BackendLoweringMechanism::TypedNvvm) + || stage.stage != EvidenceStageKind::DeviceLink + }); + } + let indexed = index_evidence( + &missing_typed_failure, + "1cb4e3833c1919c2e6fb579a23ac0e2b22587b7e", + ) + .unwrap(); + let policy = &policies[0]; + let lowering = policy + .backend_lowerings + .iter() + .find(|lowering| lowering.backend == IntrinsicBackend::LibNvvm) + .unwrap(); + let evidence = indexed + .get(&(lowering.evidence_profile.as_str(), policy.id.as_str())) + .unwrap(); + assert!(validate_evidence(policy, evidence, Some(lowering)).is_err()); +} diff --git a/crates/cuda-intrinsics-gen/src/resolve/tests/warp.rs b/crates/cuda-intrinsics-gen/src/resolve/tests/warp.rs new file mode 100644 index 0000000000..65c4260068 --- /dev/null +++ b/crates/cuda-intrinsics-gen/src/resolve/tests/warp.rs @@ -0,0 +1,1008 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::model::{ + BackendLoweringMechanism, ImportedFile, ImportedIntrinsic, IntrinsicBackend, IntrinsicSource, + OverlayIntrinsic, PreSm70MemberMaskRule, VoteMode, WarpMatchAdapter, WarpShuffleAdapter, + WarpShuffleMode, WarpShuffleOperandEncoding, WarpShuffleValueKind, +}; +use crate::ptx::{InstructionPattern, OperandPattern}; +use crate::util::read_json; +use std::collections::{BTreeMap, BTreeSet}; +use std::path::Path; + +use super::fixtures::*; +use crate::resolve::families::*; +use crate::resolve::guards::*; +use crate::resolve::materialize::*; +use crate::resolve::overlay::*; +use crate::resolve::policy::*; + +#[test] +fn pinned_active_mask_and_warp_match_recipes_resolve() { + let records = pinned_active_mask_and_warp_match_records(); + assert_eq!(records.len(), 5); + + for (policy, declaration) in records.values() { + validate_imported_policy(policy, declaration).unwrap(); + } +} + +#[test] +fn active_mask_and_warp_match_recipes_fail_closed() { + let records = pinned_active_mask_and_warp_match_records(); + let reject = |policy: &OverlayIntrinsic, declaration: &ImportedIntrinsic, expected: &str| { + let error = validate_imported_policy(policy, declaration).unwrap_err(); + let message = error.to_string(); + assert!(message.contains(expected), "unexpected error: {message}"); + }; + + let (active_mask, active_mask_declaration) = &records["active_mask"]; + let mut wrong_identity = active_mask.clone(); + wrong_identity.operation_key = "warp.active_mask.changed".into(); + reject( + &wrong_identity, + active_mask_declaration, + "active-mask identity", + ); + + let mut wrong_effects = active_mask.clone(); + wrong_effects.memory = "none".into(); + reject( + &wrong_effects, + active_mask_declaration, + "active-mask effects or target floor", + ); + + let (match_any, match_any_declaration) = &records["match_any_sync"]; + let mut wrong_adapter = match_any.clone(); + wrong_adapter.warp_match.as_mut().unwrap().adapter = + WarpMatchAdapter::ProjectMaskDiscardPredicate; + reject( + &wrong_adapter, + match_any_declaration, + "warp-match participation, adapter, or encoding", + ); + + let (match_all, match_all_declaration) = &records["match_all_sync"]; + let mut wrong_projection = match_all.clone(); + wrong_projection.dialect_results.push("i1".into()); + reject( + &wrong_projection, + match_all_declaration, + "two-operand warp-match lowering recipe", + ); + + let (match_any_i64, match_any_i64_declaration) = &records["match_any_i64_sync"]; + let mut incomplete_selections = match_any_i64_declaration.clone(); + incomplete_selections.selections.pop(); + reject( + match_any_i64, + &incomplete_selections, + "exactly ii/ir/ri/rr selections", + ); + + let (match_all_i64, match_all_i64_declaration) = &records["match_all_i64_sync"]; + let mut wrong_predicates = match_all_i64_declaration.clone(); + wrong_predicates.selections[0].predicates[0] = "Subtarget->getPTXVersion() >= 61".into(); + reject( + match_all_i64, + &wrong_predicates, + "PTX shape, predicates, or constraints", + ); +} + +#[test] +fn ptx_native_source_provenance_fails_closed() { + let mut mixed = packed_policy("packed_atomic_add_f16x2"); + mixed.source_record = Some("invented_llvm_record".into()); + assert!( + resolve_policy_source(&mixed) + .unwrap_err() + .to_string() + .contains("mixes tagged source provenance") + ); + + let mut fake_llvm = packed_policy("packed_atomic_add_f16x2"); + fake_llvm.llvm_symbol = Some("llvm.fake.packed.atomic".into()); + fake_llvm.llvm_arguments = vec!["ptr".into(), "i32".into()]; + fake_llvm.llvm_results = vec!["i32".into()]; + assert!( + validate_ptx_native_policy(&fake_llvm) + .unwrap_err() + .to_string() + .contains("must not invent LLVM source facts") + ); + + let mut wrong_instruction = packed_policy("packed_atomic_add_f16x2"); + wrong_instruction.source = Some(IntrinsicSource::PtxNative { + instruction: "atom.global.add.noftz.bf16x2".into(), + }); + assert!( + validate_ptx_native_policy(&wrong_instruction) + .unwrap_err() + .to_string() + .contains("does not match its packed format") + ); + + let mut wrong_kind = packed_policy("packed_atomic_add_f16x2"); + wrong_kind.source = Some(IntrinsicSource::LlvmImported { + source_record: "invented_llvm_record".into(), + }); + assert!( + validate_ptx_native_policy(&wrong_kind) + .unwrap_err() + .to_string() + .contains("source kind and imported declaration disagree") + ); +} + +#[test] +fn vote_modes_keep_exact_abi_identity_and_both_selection_encodings() { + for mode in [ + VoteMode::All, + VoteMode::Any, + VoteMode::Ballot, + VoteMode::Uni, + ] { + let policy = vote_policy(mode); + let declaration = vote_declaration(mode); + validate_imported_policy(&policy, &declaration).unwrap(); + assert_eq!( + policy.vote.as_ref().unwrap().legacy_pre_sm70, + PreSm70MemberMaskRule::AllNamedLanesConvergedAndOnlyNamedLanesActive + ); + + let selected: Vec<_> = declaration + .selections + .iter() + .filter(|selection| selection_matches_policy(&policy, selection).unwrap()) + .collect(); + assert_eq!(selected.len(), 2); + assert!( + selected.iter().any(|selection| { + selection.source_record == vote_recipe(mode).immediate_selection + }) + ); + assert!( + selected.iter().any(|selection| { + selection.source_record == vote_recipe(mode).register_selection + }) + ); + + let mut record = evidence(); + record.id = policy.id.clone(); + record.source_record = policy.source_record.clone(); + record.llvm_symbol = policy.llvm_symbol.clone(); + record.llvm_arguments = policy.llvm_arguments.clone(); + record.llvm_results = policy.llvm_results.clone(); + record.expected_ptx = policy.expected_ptx.clone(); + let resolved = resolve_record( + &policy, + resolve_policy_source(&policy).unwrap(), + Some(&declaration), + &record, + "test", + "LLVM version test", + "0123456789abcdef", + vec![], + 1, + ) + .unwrap(); + assert_eq!(resolved.selections.len(), 2); + assert_eq!(resolved.vote, policy.vote); + } +} + +#[test] +fn vote_contract_rejects_unreviewed_identity_effect_and_selection_changes() { + let valid = vote_policy(VoteMode::All); + let declaration = vote_declaration(VoteMode::All); + + let mut wrong_abi = valid.clone(); + wrong_abi.abi_id = "i0041".into(); + assert!( + validate_imported_policy(&wrong_abi, &declaration) + .unwrap_err() + .to_string() + .contains("vote identity") + ); + + let mut safe = valid.clone(); + safe.safe = true; + safe.safe_allowlist_reason = Some("incorrectly hides participation obligations".into()); + assert!( + validate_imported_policy(&safe, &declaration) + .unwrap_err() + .to_string() + .contains("unsafe must-use vote") + ); + + let mut wrong_memory = valid.clone(); + wrong_memory.memory = "none".into(); + assert!( + validate_imported_policy(&wrong_memory, &declaration) + .unwrap_err() + .to_string() + .contains("vote effects") + ); + + let mut register_only_mask = valid.clone(); + register_only_mask.expected_ptx.operands[2] = OperandPattern::Register; + assert!( + validate_imported_policy(®ister_only_mask, &declaration) + .unwrap_err() + .to_string() + .contains("expected PTX") + ); + + let mut one_selection = declaration.clone(); + one_selection.selections.pop(); + assert!( + validate_imported_policy(&valid, &one_selection) + .unwrap_err() + .to_string() + .contains("immediate/register selection pair") + ); + + let mut different_predicates = declaration; + different_predicates.selections[1].predicates[0] = "Subtarget->getPTXVersion() >= 61".into(); + assert!( + validate_imported_policy(&valid, &different_predicates) + .unwrap_err() + .to_string() + .contains("disagree on PTX shape") + ); +} + +#[test] +fn uni_vote_is_raw_only_while_existing_votes_keep_compatibility_paths() { + for mode in [VoteMode::All, VoteMode::Any, VoteMode::Ballot] { + assert_eq!(vote_policy(mode).compatibility_rust_paths.len(), 1); + } + let uni = vote_policy(VoteMode::Uni); + assert!(uni.compatibility_rust_paths.is_empty()); + + let mut invented_compatibility_path = uni.clone(); + invented_compatibility_path.compatibility_rust_paths = + vec!["cuda_device::warp::uni_sync".into()]; + assert!( + validate_imported_policy( + &invented_compatibility_path, + &vote_declaration(VoteMode::Uni), + ) + .unwrap_err() + .to_string() + .contains("reviewed compatibility path") + ); +} + +#[test] +fn warp_shuffle_variants_keep_exact_identity_clamp_and_eight_selections() { + for (mode, value_kind, clamp) in [ + (WarpShuffleMode::Idx, WarpShuffleValueKind::I32, 31), + (WarpShuffleMode::Bfly, WarpShuffleValueKind::I32, 31), + (WarpShuffleMode::Down, WarpShuffleValueKind::I32, 31), + (WarpShuffleMode::Up, WarpShuffleValueKind::I32, 0), + (WarpShuffleMode::Idx, WarpShuffleValueKind::F32, 31), + (WarpShuffleMode::Bfly, WarpShuffleValueKind::F32, 31), + (WarpShuffleMode::Down, WarpShuffleValueKind::F32, 31), + (WarpShuffleMode::Up, WarpShuffleValueKind::F32, 0), + ] { + let policy = warp_shuffle_policy(mode, value_kind); + let declaration = warp_shuffle_declaration(mode, value_kind); + validate_imported_policy(&policy, &declaration).unwrap(); + + assert_eq!(policy.warp_shuffle.as_ref().unwrap().clamp, clamp); + assert_eq!( + declaration + .selections + .iter() + .filter(|selection| selection_matches_policy(&policy, selection).unwrap()) + .count(), + 8 + ); + + let mut record = evidence(); + record.id = policy.id.clone(); + record.source_record = policy.source_record.clone(); + record.llvm_symbol = policy.llvm_symbol.clone(); + record.llvm_arguments = policy.llvm_arguments.clone(); + record.llvm_results = policy.llvm_results.clone(); + record.expected_ptx = policy.expected_ptx.clone(); + let resolved = resolve_record( + &policy, + resolve_policy_source(&policy).unwrap(), + Some(&declaration), + &record, + "test", + "LLVM version test", + "0123456789abcdef", + vec![], + 1, + ) + .unwrap(); + assert_eq!(resolved.selections.len(), 8); + assert_eq!(resolved.warp_shuffle, policy.warp_shuffle); + } +} + +#[test] +fn pinned_warp_shuffle_records_match_the_closed_recipes() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let (overlay, _) = + read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")).unwrap(); + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let declarations: BTreeMap<_, _> = imported + .intrinsics + .iter() + .map(|record| (record.source_record.as_str(), record)) + .collect(); + let all_policies: Vec<_> = overlay + .intrinsics + .iter() + .filter(|record| record.family == "warp_shuffle") + .collect(); + assert_eq!(all_policies.len(), 12); + let native_policies: Vec<_> = all_policies + .iter() + .copied() + .filter(|record| record.source.is_some()) + .collect(); + assert_eq!(native_policies.len(), 4); + for policy in native_policies { + validate_ptx_native_policy(policy).unwrap(); + } + let policies: Vec<_> = overlay + .intrinsics + .iter() + .filter(|record| record.family == "warp_shuffle" && record.source_record.is_some()) + .collect(); + assert_eq!(policies.len(), 8); + for policy in policies { + let declaration = declarations[policy.source_record.as_deref().unwrap()]; + validate_imported_policy(policy, declaration).unwrap(); + } +} + +#[test] +fn pinned_llvm_has_no_direct_i64_or_f64_shuffle_record() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let direct_64: Vec<_> = imported + .intrinsics + .iter() + .filter(|record| { + record.llvm_name.starts_with("llvm.nvvm.shfl") + && (record.llvm_name.contains(".i64") || record.llvm_name.contains(".f64")) + }) + .map(|record| record.llvm_name.as_str()) + .collect(); + assert!( + direct_64.is_empty(), + "unexpected LLVM records: {direct_64:?}" + ); +} + +#[test] +fn i64_warp_shuffle_recipes_are_exact_ptx_native_pairs() { + let cases = [ + ( + WarpShuffleMode::Idx, + "shuffle_u64_sync", + "i0058", + "warp.shuffle.sync.idx.i64", + "shfl.sync.idx.b32", + "idx", + 31, + "ShflSyncIdxI64Op", + "nvvm.shfl_sync_idx_i64", + ), + ( + WarpShuffleMode::Bfly, + "shuffle_xor_u64_sync", + "i0059", + "warp.shuffle.sync.bfly.i64", + "shfl.sync.bfly.b32", + "bfly", + 31, + "ShflSyncBflyI64Op", + "nvvm.shfl_sync_bfly_i64", + ), + ( + WarpShuffleMode::Down, + "shuffle_down_u64_sync", + "i0060", + "warp.shuffle.sync.down.i64", + "shfl.sync.down.b32", + "down", + 31, + "ShflSyncDownI64Op", + "nvvm.shfl_sync_down_i64", + ), + ( + WarpShuffleMode::Up, + "shuffle_up_u64_sync", + "i0061", + "warp.shuffle.sync.up.i64", + "shfl.sync.up.b32", + "up", + 0, + "ShflSyncUpI64Op", + "nvvm.shfl_sync_up_i64", + ), + ]; + + for (mode, id, abi_id, operation_key, instruction, ptx_mode, clamp, op_type, op_name) in cases { + let policy = warp_shuffle_policy(mode, WarpShuffleValueKind::I64); + validate_ptx_native_policy(&policy).unwrap(); + + assert_eq!(policy.id, id); + assert_eq!(policy.abi_id, abi_id); + assert_eq!(policy.operation_key, operation_key); + assert_eq!( + policy.source, + Some(IntrinsicSource::PtxNative { + instruction: instruction.into(), + }) + ); + assert!(policy.source_record.is_none()); + assert!(policy.llvm_symbol.is_none()); + assert!(policy.resolved_llvm_symbol.is_none()); + assert!(policy.llvm_arguments.is_empty()); + assert!(policy.llvm_results.is_empty()); + assert_eq!(policy.rust_arguments, ["u32", "u64", "u32"]); + assert_eq!(policy.rust_result, "u64"); + assert!(!policy.safe); + assert!(policy.must_use); + assert_eq!(policy.dialect_op_type, op_type); + assert_eq!(policy.dialect_op_name, op_name); + assert_eq!(policy.dialect_operands, ["i32", "i64", "i32"]); + assert_eq!(policy.dialect_results, ["i64"]); + assert_eq!(policy.lowering, "generated_warp_shuffle_i64_inline_ptx"); + + let shuffle = policy.warp_shuffle.as_ref().unwrap(); + assert_eq!(shuffle.clamp, clamp); + assert_eq!( + shuffle.adapter, + WarpShuffleAdapter::MaskValueLaneOrDeltaSplitI64LowHighB32InsertClampReassemble + ); + assert_eq!( + shuffle.lane_encoding, + WarpShuffleOperandEncoding::RegisterOnly + ); + assert_eq!( + shuffle.mask_encoding, + WarpShuffleOperandEncoding::RegisterOnly + ); + assert_eq!( + policy.expected_ptx, + InstructionPattern::new( + "shfl", + &["sync", ptx_mode, "b32"], + vec![ + OperandPattern::Exact { value: "lo".into() }, + OperandPattern::Exact { value: "lo".into() }, + OperandPattern::Register, + OperandPattern::Exact { + value: clamp.to_string(), + }, + OperandPattern::Register, + ], + ) + ); + + let routes: BTreeMap<_, _> = policy + .backend_lowerings + .iter() + .map(|route| (route.backend, route)) + .collect(); + assert_eq!(routes.len(), 2); + for backend in [IntrinsicBackend::LlvmNvptx, IntrinsicBackend::LibNvvm] { + assert_eq!( + routes[&backend].mechanism, + BackendLoweringMechanism::InlinePtx + ); + assert_eq!(routes[&backend].minimum_ptx.as_deref(), Some("6.0")); + } + assert_eq!( + routes[&IntrinsicBackend::LlvmNvptx].minimum_sm.as_deref(), + Some("sm_30") + ); + assert_eq!( + routes[&IntrinsicBackend::LibNvvm].minimum_sm.as_deref(), + Some("sm_75") + ); + + let mut record = evidence(); + record.id = policy.id.clone(); + record.source = policy.source.clone(); + record.source_record = None; + record.llvm_symbol = None; + record.llvm_arguments.clear(); + record.llvm_results.clear(); + record.expected_ptx = policy.expected_ptx.clone(); + let resolved = resolve_record( + &policy, + resolve_policy_source(&policy).unwrap(), + None, + &record, + "test", + "LLVM version test", + "0123456789abcdef", + vec![], + 1, + ) + .unwrap(); + assert!(resolved.llvm.is_none()); + assert!(resolved.selections.is_empty()); + assert_eq!(resolved.warp_shuffle, policy.warp_shuffle); + } +} + +#[test] +fn i64_warp_shuffle_contract_rejects_unreviewed_changes() { + let valid = warp_shuffle_policy(WarpShuffleMode::Idx, WarpShuffleValueKind::I64); + validate_ptx_native_policy(&valid).unwrap(); + + let reject = |policy: &OverlayIntrinsic, expected: &str| { + let error = validate_ptx_native_policy(policy).unwrap_err(); + let message = error.to_string(); + assert!(message.contains(expected), "unexpected error: {message}"); + }; + + let mut fabricated_llvm = valid.clone(); + fabricated_llvm.source = None; + fabricated_llvm.source_record = Some("int_nvvm_shfl_sync_idx_i64".into()); + fabricated_llvm.llvm_symbol = Some("llvm.nvvm.shfl.sync.idx.i64".into()); + fabricated_llvm.llvm_arguments = vec!["i32".into(), "i64".into(), "i32".into(), "i32".into()]; + fabricated_llvm.llvm_results = vec!["i64".into()]; + reject( + &fabricated_llvm, + "source kind and imported declaration disagree", + ); + + let mut wrong_source = valid.clone(); + wrong_source.source = Some(IntrinsicSource::PtxNative { + instruction: "shfl.sync.down.b32".into(), + }); + reject(&wrong_source, "warp-shuffle identity"); + + let mut wrong_adapter = valid.clone(); + wrong_adapter.warp_shuffle.as_mut().unwrap().adapter = + WarpShuffleAdapter::MaskValueLaneOrDeltaInsertClamp; + reject(&wrong_adapter, "semantic or operand contract"); + + let mut wrong_mode = valid.clone(); + wrong_mode.warp_shuffle.as_mut().unwrap().mode = WarpShuffleMode::Up; + wrong_mode.warp_shuffle.as_mut().unwrap().clamp = 0; + wrong_mode.expected_ptx.modifiers[1] = "up".into(); + wrong_mode.expected_ptx.operands[3] = OperandPattern::Exact { value: "0".into() }; + reject(&wrong_mode, "warp-shuffle identity"); + + let mut wrong_clamp = valid.clone(); + wrong_clamp.warp_shuffle.as_mut().unwrap().clamp = 0; + reject(&wrong_clamp, "semantic or operand contract"); + + let mut broad_encoding = valid.clone(); + broad_encoding.warp_shuffle.as_mut().unwrap().lane_encoding = + WarpShuffleOperandEncoding::RegisterOrImmediate; + reject(&broad_encoding, "semantic or operand contract"); + + let mut typed_backend = valid.clone(); + typed_backend.backend_lowerings[0].mechanism = BackendLoweringMechanism::TypedNvvm; + reject(&typed_backend, "reviewed LLVM and libNVVM routes"); + + let mut wrong_native_floor = valid.clone(); + wrong_native_floor.minimum_sm = Some("sm_70".into()); + reject(&wrong_native_floor, "target floor"); + + let mut wrong_profile_floor = valid.clone(); + wrong_profile_floor + .backend_lowerings + .iter_mut() + .find(|route| route.backend == IntrinsicBackend::LibNvvm) + .unwrap() + .minimum_sm = Some("sm_80".into()); + reject(&wrong_profile_floor, "profile floor"); + + let mut safe = valid.clone(); + safe.safe = true; + safe.safe_allowlist_reason = Some("incorrectly hides participation obligations".into()); + reject(&safe, "unsafe must-use warp-shuffle"); + + let mut wrong_ptx = valid; + wrong_ptx.expected_ptx.operands[0] = OperandPattern::Register; + reject(&wrong_ptx, "closed shfl.sync recipe"); +} + +#[test] +fn warp_shuffle_contract_rejects_unreviewed_policy_changes() { + let valid = warp_shuffle_policy(WarpShuffleMode::Idx, WarpShuffleValueKind::I32); + let declaration = warp_shuffle_declaration(WarpShuffleMode::Idx, WarpShuffleValueKind::I32); + + let reject_policy = |policy: &OverlayIntrinsic, expected: &str| { + let error = match validate_imported_policy(policy, &declaration) { + Ok(()) => panic!("{expected} mutation was accepted"), + Err(error) => error, + }; + let message = error.to_string(); + assert!(message.contains(expected), "unexpected error: {message}"); + }; + + let mut wrong_identity = valid.clone(); + wrong_identity.operation_key = "warp.shuffle.sync.idx.changed".into(); + reject_policy(&wrong_identity, "warp-shuffle identity"); + + let mut safe = valid.clone(); + safe.safe = true; + safe.safe_allowlist_reason = Some("incorrectly hides participation obligations".into()); + reject_policy(&safe, "unsafe must-use warp-shuffle"); + + let mut wrong_signature = valid.clone(); + wrong_signature.dialect_operands.pop(); + reject_policy(&wrong_signature, "closed warp-shuffle lowering recipe"); + + let mut wrong_clamp = valid.clone(); + wrong_clamp.warp_shuffle.as_mut().unwrap().clamp = 0; + reject_policy(&wrong_clamp, "semantic or operand contract"); + + let mut missing_contract = valid.clone(); + missing_contract.warp_shuffle = None; + reject_policy(&missing_contract, "closed warp-shuffle contract"); + + let mut mixed_contract = valid.clone(); + mixed_contract.vote = vote_policy(VoteMode::All).vote; + reject_policy(&mixed_contract, "mixes another generated-family contract"); + + let mut wrong_backend_floor = valid.clone(); + wrong_backend_floor + .backend_lowerings + .iter_mut() + .find(|lowering| lowering.backend == IntrinsicBackend::LibNvvm) + .unwrap() + .minimum_sm = Some("sm_80".into()); + reject_policy(&wrong_backend_floor, "profile floor"); +} + +#[test] +fn warp_shuffle_contract_rejects_selection_drift() { + let valid = warp_shuffle_policy(WarpShuffleMode::Down, WarpShuffleValueKind::F32); + let declaration = warp_shuffle_declaration(WarpShuffleMode::Down, WarpShuffleValueKind::F32); + let reject = |declaration: &ImportedIntrinsic, expected: &str| { + let error = validate_imported_policy(&valid, declaration).unwrap_err(); + let message = error.to_string(); + assert!(message.contains(expected), "unexpected error: {message}"); + }; + + let mut missing_selection = declaration.clone(); + missing_selection.selections.pop(); + reject( + &missing_selection, + "eight distinct operand-encoding selections", + ); + + let mut duplicate_selection = declaration.clone(); + duplicate_selection.selections[7].source_record = + duplicate_selection.selections[0].source_record.clone(); + reject( + &duplicate_selection, + "eight distinct operand-encoding selections", + ); + + let mut empty_selection_name = declaration.clone(); + empty_selection_name.selections[7].source_record.clear(); + reject( + &empty_selection_name, + "eight distinct operand-encoding selections", + ); + + let mut wrong_asm = declaration.clone(); + wrong_asm.selections[0].asm = + "shfl.sync.up.b32 \t$dst, $src, $offset, $mask, $threadmask;".into(); + reject(&wrong_asm, "selections disagree on PTX shape"); + + let mut wrong_predicate = declaration.clone(); + wrong_predicate.selections[0].predicates[0] = "Subtarget->getPTXVersion() >= 61".into(); + reject(&wrong_predicate, "selections disagree on PTX shape"); + + let mut constrained = declaration; + constrained.selections[0] + .constraints + .immediate_bindings + .push(crate::model::ImportedImmediateBinding { + argument_index: 2, + value: 1, + }); + reject(&constrained, "selections disagree on PTX shape"); + + let mut wrong_classes = + warp_shuffle_declaration(WarpShuffleMode::Down, WarpShuffleValueKind::F32); + wrong_classes.classes.pop(); + reject(&wrong_classes, "class or effects"); +} + +#[test] +fn sync_threads_selects_only_the_fixed_immediate_barrier_recipe() { + let policy = sync_policy(); + let declaration = sync_declaration(); + validate_imported_policy(&policy, &declaration).unwrap(); + + let selected: Vec<_> = declaration + .selections + .iter() + .filter(|selection| selection_matches_policy(&policy, selection).unwrap()) + .collect(); + assert_eq!(selected.len(), 1); + assert_eq!(selected[0].source_record, "BARRIER_CTA_SYNC_ALIGNED_ALL_i"); + assert_eq!(selected[0].asm, "bar.sync \t$i;"); + assert!(policy.expected_ptx.matches("bar.sync 0;").unwrap()); + assert!(!policy.expected_ptx.matches(&selected[0].asm).unwrap()); + assert_eq!(policy.minimum_ptx, "1.0"); + assert!(policy.minimum_sm.is_none()); + let llvm_route = policy + .backend_lowerings + .iter() + .find(|lowering| lowering.backend == IntrinsicBackend::LlvmNvptx) + .unwrap(); + assert_eq!(llvm_route.minimum_ptx.as_deref(), Some("3.2")); + assert_eq!(llvm_route.minimum_sm.as_deref(), Some("sm_20")); + + let resolved = resolve_record( + &policy, + resolve_policy_source(&policy).unwrap(), + Some(&declaration), + &sync_evidence(&policy), + "test", + "LLVM version test", + "0123456789abcdef", + vec![], + 1, + ) + .unwrap(); + assert!(resolved.dialect.operands.is_empty()); + assert!(resolved.dialect.results.is_empty()); + assert_eq!(resolved.selections.len(), 1); + assert_eq!( + resolved.selections[0].source_record, + "BARRIER_CTA_SYNC_ALIGNED_ALL_i" + ); +} + +#[test] +fn sync_threads_recipe_rejects_unreviewed_selection_effect_and_floor_changes() { + let valid = sync_policy(); + let declaration = sync_declaration(); + + let mut register_only = declaration.clone(); + register_only + .selections + .retain(|selection| selection.source_record.ends_with("_r")); + assert!( + validate_imported_policy(&valid, ®ister_only) + .unwrap_err() + .to_string() + .contains("does not agree") + ); + + let mut wrong_properties = declaration.clone(); + wrong_properties.properties.pop(); + assert!( + validate_imported_policy(&valid, &wrong_properties) + .unwrap_err() + .to_string() + .contains("sync properties") + ); + + let mut wrong_source = valid.clone(); + wrong_source.source_record = Some("int_nvvm_barrier0".into()); + assert!( + validate_imported_policy(&wrong_source, &declaration) + .unwrap_err() + .to_string() + .contains("sync identity") + ); + + let mut wrong_signature = valid.clone(); + wrong_signature.llvm_arguments.clear(); + assert!( + validate_imported_policy(&wrong_signature, &declaration) + .unwrap_err() + .to_string() + .contains("LLVM argument signature mismatch") + ); + + let mut wrong_path = valid.clone(); + wrong_path.compatibility_rust_paths.swap(0, 1); + assert!( + validate_imported_policy(&wrong_path, &declaration) + .unwrap_err() + .to_string() + .contains("both cuda-device compatibility paths") + ); + + let mut safe = valid.clone(); + safe.safe = true; + safe.safe_allowlist_reason = Some("incorrectly hides the participation contract".into()); + assert!( + validate_imported_policy(&safe, &declaration) + .unwrap_err() + .to_string() + .contains("unsafe sync_threads raw API") + ); + + let mut wrong_effect = valid.clone(); + wrong_effect.memory = "none".into(); + assert!( + validate_imported_policy(&wrong_effect, &declaration) + .unwrap_err() + .to_string() + .contains("sync effects") + ); + + let mut native_floor = valid.clone(); + native_floor.minimum_sm = Some("sm_75".into()); + assert!( + validate_imported_policy(&native_floor, &declaration) + .unwrap_err() + .to_string() + .contains("native target floor") + ); + + let mut missing_profile_floor = valid; + missing_profile_floor + .backend_lowerings + .iter_mut() + .find(|lowering| lowering.backend == IntrinsicBackend::LibNvvm) + .unwrap() + .minimum_sm = None; + assert!( + validate_imported_policy(&missing_profile_floor, &declaration) + .unwrap_err() + .to_string() + .contains("profile floor") + ); + + let mut wrong_llvm_floor = sync_policy(); + wrong_llvm_floor + .backend_lowerings + .iter_mut() + .find(|lowering| lowering.backend == IntrinsicBackend::LlvmNvptx) + .unwrap() + .minimum_ptx = None; + assert!( + validate_imported_policy(&wrong_llvm_floor, &declaration) + .unwrap_err() + .to_string() + .contains("profile floor") + ); +} + +#[test] +fn sync_mask_matches_the_closed_warp_barrier_recipe() { + let policy = warp_barrier_policy(); + let declaration = warp_barrier_declaration(); + validate_imported_policy(&policy, &declaration).unwrap(); + + let selected: Vec<_> = declaration + .selections + .iter() + .filter(|selection| selection_matches_policy(&policy, selection).unwrap()) + .collect(); + assert_eq!(selected.len(), 2); + assert_eq!( + selected + .iter() + .map(|selection| selection.source_record.as_str()) + .collect::>(), + BTreeSet::from(["INT_BAR_WARP_SYNC_I", "INT_BAR_WARP_SYNC_R"]) + ); + + let mut record = evidence(); + record.id = policy.id.clone(); + record.source_record = policy.source_record.clone(); + record.llvm_symbol = policy.llvm_symbol.clone(); + record.llvm_arguments = policy.llvm_arguments.clone(); + record.llvm_results = policy.llvm_results.clone(); + record.expected_ptx = policy.expected_ptx.clone(); + let resolved = resolve_record( + &policy, + resolve_policy_source(&policy).unwrap(), + Some(&declaration), + &record, + "test", + "LLVM version test", + "0123456789abcdef", + vec![], + 1, + ) + .unwrap(); + assert_eq!(resolved.selections.len(), 2); + assert_eq!(resolved.warp_barrier, policy.warp_barrier); +} + +#[test] +fn sync_mask_recipe_rejects_unreviewed_contract_and_selection_changes() { + let valid = warp_barrier_policy(); + let declaration = warp_barrier_declaration(); + + let mut wrong_identity = valid.clone(); + wrong_identity.id = "bar_warp_sync".into(); + assert!( + validate_imported_policy(&wrong_identity, &declaration) + .unwrap_err() + .to_string() + .contains("warp-barrier identity") + ); + + let mut missing_contract = valid.clone(); + missing_contract.warp_barrier = None; + assert!( + validate_imported_policy(&missing_contract, &declaration) + .unwrap_err() + .to_string() + .contains("closed warp-barrier contract") + ); + + let mut safe_raw_api = valid.clone(); + safe_raw_api.safe = true; + safe_raw_api.safe_allowlist_reason = Some("incorrectly hides participation rules".into()); + assert!( + validate_imported_policy(&safe_raw_api, &declaration) + .unwrap_err() + .to_string() + .contains("unsafe raw API") + ); + + let mut wrong_memory = valid.clone(); + wrong_memory.memory = "none".into(); + assert!( + validate_imported_policy(&wrong_memory, &declaration) + .unwrap_err() + .to_string() + .contains("effects or target floor") + ); + + let mut register_only = valid.clone(); + register_only.expected_ptx.operands[0] = OperandPattern::Register; + assert!( + validate_imported_policy(®ister_only, &declaration) + .unwrap_err() + .to_string() + .contains("expected PTX") + ); + + let mut one_selection = declaration.clone(); + one_selection.selections.pop(); + assert!( + validate_imported_policy(&valid, &one_selection) + .unwrap_err() + .to_string() + .contains("immediate/register selection pair") + ); + + let mut wrong_predicate = declaration.clone(); + wrong_predicate.selections[1].predicates[0] = "Subtarget->getPTXVersion() >= 61".into(); + assert!( + validate_imported_policy(&valid, &wrong_predicate) + .unwrap_err() + .to_string() + .contains("selections disagree") + ); + + let mut missing_libnvvm_floor = valid; + missing_libnvvm_floor + .backend_lowerings + .iter_mut() + .find(|lowering| lowering.backend == IntrinsicBackend::LibNvvm) + .unwrap() + .minimum_sm = None; + assert!( + validate_imported_policy(&missing_libnvvm_floor, &declaration) + .unwrap_err() + .to_string() + .contains("profile floor") + ); +} From b595ad4a9302993ce344018dc5e452b1008c9c2f Mon Sep 17 00:00:00 2001 From: nihalpasham Date: Fri, 28 Aug 2026 22:19:00 +0530 Subject: [PATCH 11/18] refactor(cargo-oxide): split commands.rs into commands/ modules --- crates/cargo-oxide/src/commands.rs | 11297 ---------------- crates/cargo-oxide/src/commands/artifacts.rs | 166 + crates/cargo-oxide/src/commands/build_run.rs | 472 + crates/cargo-oxide/src/commands/clean.rs | 234 + .../cargo-oxide/src/commands/codegen_env.rs | 641 + crates/cargo-oxide/src/commands/context.rs | 511 + crates/cargo-oxide/src/commands/doctor.rs | 702 + .../cargo-oxide/src/commands/examples_list.rs | 620 + .../cargo-oxide/src/commands/fingerprint.rs | 368 + crates/cargo-oxide/src/commands/fmt.rs | 164 + crates/cargo-oxide/src/commands/host_cargo.rs | 528 + crates/cargo-oxide/src/commands/interop.rs | 938 ++ crates/cargo-oxide/src/commands/ltoir.rs | 220 + .../cargo-oxide/src/commands/materialize.rs | 334 + crates/cargo-oxide/src/commands/mod.rs | 50 + .../cargo-oxide/src/commands/passthrough.rs | 311 + .../src/commands/pipeline_debug.rs | 306 + crates/cargo-oxide/src/commands/scaffold.rs | 333 + .../cargo-oxide/src/commands/setup_update.rs | 141 + crates/cargo-oxide/src/commands/tests.rs | 4227 ++++++ crates/cargo-oxide/src/commands/toolkit.rs | 259 + 21 files changed, 11525 insertions(+), 11297 deletions(-) delete mode 100644 crates/cargo-oxide/src/commands.rs create mode 100644 crates/cargo-oxide/src/commands/artifacts.rs create mode 100644 crates/cargo-oxide/src/commands/build_run.rs create mode 100644 crates/cargo-oxide/src/commands/clean.rs create mode 100644 crates/cargo-oxide/src/commands/codegen_env.rs create mode 100644 crates/cargo-oxide/src/commands/context.rs create mode 100644 crates/cargo-oxide/src/commands/doctor.rs create mode 100644 crates/cargo-oxide/src/commands/examples_list.rs create mode 100644 crates/cargo-oxide/src/commands/fingerprint.rs create mode 100644 crates/cargo-oxide/src/commands/fmt.rs create mode 100644 crates/cargo-oxide/src/commands/host_cargo.rs create mode 100644 crates/cargo-oxide/src/commands/interop.rs create mode 100644 crates/cargo-oxide/src/commands/ltoir.rs create mode 100644 crates/cargo-oxide/src/commands/materialize.rs create mode 100644 crates/cargo-oxide/src/commands/mod.rs create mode 100644 crates/cargo-oxide/src/commands/passthrough.rs create mode 100644 crates/cargo-oxide/src/commands/pipeline_debug.rs create mode 100644 crates/cargo-oxide/src/commands/scaffold.rs create mode 100644 crates/cargo-oxide/src/commands/setup_update.rs create mode 100644 crates/cargo-oxide/src/commands/tests.rs create mode 100644 crates/cargo-oxide/src/commands/toolkit.rs diff --git a/crates/cargo-oxide/src/commands.rs b/crates/cargo-oxide/src/commands.rs deleted file mode 100644 index 4e36ec4afa..0000000000 --- a/crates/cargo-oxide/src/commands.rs +++ /dev/null @@ -1,11297 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -//! Command implementations for cargo-oxide. -//! -//! These port the xtask commands with improvements: -//! - Backend path resolved via discovery chain instead of hardcoded relative path -//! - Workspace root resolved by walking up from CWD instead of assuming CWD - -use crate::backend; -use sha2::Digest as _; -use std::collections::BTreeMap; -use std::fs; -use std::path::{Path, PathBuf}; -use std::process::Command; - -const MATERIALIZE_ENV: &str = reserved_oxide_symbols::MATERIALIZE_CUBIN_ENV; -const EXPECTED_PROVENANCE_ENV: &str = reserved_oxide_symbols::MATERIALIZER_PROVENANCE_ENV; -const MATERIALIZER_HANDSHAKE_ENV: &str = reserved_oxide_symbols::MATERIALIZER_HANDSHAKE_ENV; -const CODEGEN_FINGERPRINT_ENV: &str = reserved_oxide_symbols::CODEGEN_FINGERPRINT_ENV; -const DEVICE_CODEGEN_CRATE_ENV: &str = reserved_oxide_symbols::DEVICE_CODEGEN_CRATE_ENV; -const BACKEND_IDENTITY_CFG: &str = "cuda_oxide_internal_backend_identity"; -const LEGACY_CODEGEN_FINGERPRINT_CFG: &str = "cuda_oxide_internal_codegen_env"; -const LEGACY_MATERIALIZER_PROVENANCE_CFG: &str = "cuda_oxide_internal_materializer_provenance"; -const MATERIALIZER_HANDSHAKE_CACHE: &str = ".oxide-artifacts/materializer-handshake/v1.json"; - -#[derive(Clone, Debug, Default, Eq, PartialEq)] -struct MaterializationMode { - prepared: Option, -} - -#[derive(Clone, Debug, Eq, PartialEq)] -struct PreparedMaterialization { - provenance_sha256_hex: String, - tool_identity_handshake_json: String, -} - -impl MaterializationMode { - fn enabled(&self) -> bool { - self.prepared.is_some() - } - - fn apply(&self, cmd: &mut Command) { - if let Some(prepared) = &self.prepared { - // These override inherited/project values: they are a single - // wrapper-generated handshake tied to this Cargo invocation. - cmd.env(MATERIALIZE_ENV, "1") - .env(EXPECTED_PROVENANCE_ENV, &prepared.provenance_sha256_hex) - .env( - MATERIALIZER_HANDSHAKE_ENV, - &prepared.tool_identity_handshake_json, - ) - .env("CUDA_OXIDE_EMIT_NVVM_IR", "1"); - } - } -} - -fn prepare_materialization( - ctx: &Context, - cli_requested: bool, - cli_arch: Option<&str>, - emit_nvvm_ir: bool, -) -> MaterializationMode { - prepare_materialization_result(ctx, cli_requested, cli_arch, emit_nvvm_ir).unwrap_or_else( - |error| { - eprintln!("Error: {error}"); - std::process::exit(2); - }, - ) -} - -/// `prepare_materialization` with the ambient `CUDA_OXIDE_MATERIALIZE_CUBIN` -/// injected, so `cargo_passthrough_command_with_env` can reach -/// `materialization_requested_with_env`. -/// -/// Note this still exits the process on error, which inside a unit test aborts -/// the whole test binary rather than failing one case -- a further reason tests -/// must not reach the ambient read. -fn prepare_materialization_with_env( - ctx: &Context, - cli_requested: bool, - cli_arch: Option<&str>, - emit_nvvm_ir: bool, - materialize_env: Option, -) -> MaterializationMode { - prepare_materialization_result_with_env( - ctx, - cli_requested, - cli_arch, - emit_nvvm_ir, - materialize_env, - ) - .unwrap_or_else(|error| { - eprintln!("Error: {error}"); - std::process::exit(2); - }) -} - -const EMIT_NVVM_IR_ENV: &str = "CUDA_OXIDE_EMIT_NVVM_IR"; - -fn nvvm_ir_requested(ctx: &Context) -> Result { - nvvm_ir_requested_with_env(ctx, std::env::var_os(EMIT_NVVM_IR_ENV)) -} - -/// `nvvm_ir_requested` with the ambient `CUDA_OXIDE_EMIT_NVVM_IR` injected. -/// -/// The process value outranks project config, so resolution has to be -/// injectable for unit tests: an exported `CUDA_OXIDE_EMIT_NVVM_IR` would -/// otherwise decide the answer before the configured value is consulted. -fn nvvm_ir_requested_with_env( - ctx: &Context, - env_value: Option, -) -> Result { - if let Some(value) = env_value { - let value = value - .into_string() - .map_err(|_| format!("{EMIT_NVVM_IR_ENV} is not valid Unicode"))?; - return parse_strict_bool(EMIT_NVVM_IR_ENV, &value); - } - - if let Some(value) = project_config_env(ctx, EMIT_NVVM_IR_ENV) { - return parse_strict_bool(EMIT_NVVM_IR_ENV, value); - } - - Ok(false) -} - -fn materialization_requested(ctx: &Context, cli_requested: bool) -> Result { - materialization_requested_with_env(ctx, cli_requested, std::env::var_os(MATERIALIZE_ENV)) -} - -/// `materialization_requested` with the ambient `CUDA_OXIDE_MATERIALIZE_CUBIN` -/// injected. -/// -/// The process value outranks project config, so resolution has to be -/// injectable for unit tests: an exported value would otherwise turn -/// materialization on for tests that pass `materialize_cubin: false`, sending -/// them into `discover_materializer_provenance`. Same rationale as -/// `nvvm_ir_requested_with_env`. -fn materialization_requested_with_env( - ctx: &Context, - cli_requested: bool, - env_value: Option, -) -> Result { - if cli_requested { - return Ok(true); - } - - if let Some(value) = env_value { - let value = value - .into_string() - .map_err(|_| format!("{MATERIALIZE_ENV} is not valid Unicode"))?; - return parse_strict_bool(MATERIALIZE_ENV, &value); - } - - if let Some(value) = project_config_env(ctx, MATERIALIZE_ENV) { - return parse_strict_bool(MATERIALIZE_ENV, value); - } - - Ok(false) -} - -fn prepare_materialization_result( - ctx: &Context, - cli_requested: bool, - cli_arch: Option<&str>, - emit_nvvm_ir: bool, -) -> Result { - prepare_materialization_result_with_env( - ctx, - cli_requested, - cli_arch, - emit_nvvm_ir, - std::env::var_os(MATERIALIZE_ENV), - ) -} - -/// `prepare_materialization_result` with the ambient -/// `CUDA_OXIDE_MATERIALIZE_CUBIN` injected, forwarded to -/// `materialization_requested_with_env`. -fn prepare_materialization_result_with_env( - ctx: &Context, - cli_requested: bool, - cli_arch: Option<&str>, - emit_nvvm_ir: bool, - materialize_env: Option, -) -> Result { - let enabled = materialization_requested_with_env(ctx, cli_requested, materialize_env)?; - if !enabled { - return Ok(MaterializationMode::default()); - } - if emit_nvvm_ir { - return Err( - "--materialize-cubin cannot be combined with --emit-nvvm-ir; one requests a final cubin and the other requests NVVM IR" - .to_string(), - ); - } - - let arch = configured_arch_label(ctx, cli_arch).ok_or_else(|| { - "--materialize-cubin requires --arch, CUDA_OXIDE_TARGET, or a configured default-arch" - .to_string() - })?; - let _: cuda_artifact_finalizer::CudaArch = arch - .parse() - .map_err(|error| format!("invalid materialization target {arch:?}: {error}"))?; - - let handshake = discover_materializer_handshake(ctx)?; - let handshake_json = serde_json::to_string(&handshake) - .map_err(|error| format!("could not encode materializer handshake: {error}"))?; - Ok(MaterializationMode { - prepared: Some(PreparedMaterialization { - provenance_sha256_hex: digest_hex(&handshake.provenance_sha256), - tool_identity_handshake_json: handshake_json, - }), - }) -} - -fn discover_materializer_handshake( - ctx: &Context, -) -> Result { - let executable = std::env::current_exe() - .map_err(|error| format!("could not locate cargo-oxide executable: {error}"))?; - let mut command = materializer_discovery_command(ctx, &executable); - if let Some(cached) = read_materializer_handshake_cache(ctx) { - command.env(MATERIALIZER_HANDSHAKE_ENV, cached); - } - let output = command - .output() - .map_err(|error| format!("could not start CUDA materializer discovery: {error}"))?; - if !output.status.success() { - let stderr = String::from_utf8_lossy(&output.stderr); - return Err(format!( - "CUDA materializer discovery failed: {}", - stderr.trim() - )); - } - let handshake = String::from_utf8(output.stdout) - .map_err(|_| "CUDA materializer discovery returned non-UTF-8 output".to_string())?; - let handshake: cuda_artifact_finalizer::MaterializerHandshakeV1 = - serde_json::from_str(handshake.trim()).map_err(|error| { - format!("CUDA materializer discovery returned an invalid v1 handshake: {error}") - })?; - if !handshake.has_consistent_provenance() { - return Err(format!( - "CUDA materializer discovery returned an inconsistent handshake version {}", - handshake.version - )); - } - write_materializer_handshake_cache(ctx, &handshake); - Ok(handshake) -} - -fn materializer_discovery_command(ctx: &Context, executable: &Path) -> Command { - let mut command = Command::new(executable); - command.arg("__materializer-handshake"); - apply_config_env(&mut command, ctx); - apply_ld_library_path(&mut command, ctx); - // Only the local cache explicitly installed by the caller may seed the - // helper; never consume an inherited or project-provided internal value. - command.env_remove(MATERIALIZER_HANDSHAKE_ENV); - command -} - -fn materializer_handshake_cache_path(ctx: &Context) -> PathBuf { - ctx.workspace_root.join(MATERIALIZER_HANDSHAKE_CACHE) -} - -fn read_materializer_handshake_cache(ctx: &Context) -> Option { - let json = fs::read_to_string(materializer_handshake_cache_path(ctx)).ok()?; - let handshake: cuda_artifact_finalizer::MaterializerHandshakeV1 = - serde_json::from_str(json.trim()).ok()?; - handshake.has_consistent_provenance().then_some(json) -} - -fn write_materializer_handshake_cache( - ctx: &Context, - handshake: &cuda_artifact_finalizer::MaterializerHandshakeV1, -) { - let path = materializer_handshake_cache_path(ctx); - let Some(parent) = path.parent() else { - return; - }; - if fs::create_dir_all(parent).is_err() { - return; - } - let Ok(json) = serde_json::to_string(handshake) else { - return; - }; - let temporary = parent.join(format!("v1.{}.tmp", std::process::id())); - if fs::write(&temporary, json).is_ok() { - let _ = fs::rename(&temporary, &path); - } -} - -pub fn print_materializer_handshake() { - let cached = std::env::var(MATERIALIZER_HANDSHAKE_ENV) - .ok() - .and_then(|json| serde_json::from_str(&json).ok()) - .filter(cuda_artifact_finalizer::MaterializerHandshakeV1::has_consistent_provenance); - let finalizer = cached - .as_ref() - .map_or_else( - cuda_artifact_finalizer::Finalizer::discover, - cuda_artifact_finalizer::Finalizer::discover_with_handshake, - ) - .unwrap_or_else(|error| { - eprintln!("could not discover CUDA artifact finalizer: {error}"); - std::process::exit(1); - }); - let handshake = finalizer.materializer_handshake().unwrap_or_else(|| { - eprintln!( - "the loaded libNVVM or nvJitLink library cannot be tied to an exact file; refusing materialization because Cargo could not fingerprint the compiler inputs" - ); - std::process::exit(1); - }); - println!( - "{}", - serde_json::to_string(&handshake).unwrap_or_else(|error| { - eprintln!("could not encode CUDA materializer handshake: {error}"); - std::process::exit(1); - }) - ); -} - -fn parse_strict_bool(name: &str, value: &str) -> Result { - match value.trim().to_ascii_lowercase().as_str() { - "1" | "true" | "yes" | "on" => Ok(true), - "0" | "false" | "no" | "off" => Ok(false), - _ => Err(format!( - "{name} must be a boolean (accepted true values: 1, true, yes, on; false values: 0, false, no, off), got {value:?}" - )), - } -} - -fn digest_hex(digest: &[u8; 32]) -> String { - use std::fmt::Write; - let mut hex = String::with_capacity(64); - for byte in digest { - write!(&mut hex, "{byte:02x}").expect("writing to String cannot fail"); - } - hex -} - -/// Project-local cuda-oxide defaults loaded from `.cargo/cuda-oxide.toml`. -#[derive(Debug, Clone, Default, PartialEq, Eq)] -pub struct OxideConfig { - /// Explicit backend shared object path. - pub backend: Option, - /// Default CUDA architecture for codegen commands. - pub default_arch: Option, - /// Additional rustflags appended after cuda-oxide's required flags. - pub extra_rustflags: Vec, - /// Environment variables applied to child Cargo invocations. - pub env: Vec<(String, String)>, -} - -/// Pre-resolved context shared across all commands. -/// -/// Built once at startup by [`resolve_context`] and passed by reference to -/// every command handler. Avoids repeated filesystem walks and backend builds. -pub struct Context { - /// Absolute path to the workspace root (contains top-level `Cargo.toml`). - pub workspace_root: PathBuf, - /// Path to `crates/rustc-codegen-cuda` (backend source tree). - pub codegen_crate: PathBuf, - /// Path to `crates/rustc-codegen-cuda/examples/`. - pub examples_dir: PathBuf, - /// Path to the built `librustc_codegen_cuda.so` shared object. - pub backend_so: PathBuf, - /// True when running from inside the cuda-oxide workspace; false for - /// standalone projects scaffolded by `cargo oxide new`. - pub is_workspace: bool, - /// Project-local cuda-oxide defaults. - pub config: OxideConfig, -} - -/// Resolve the workspace root and backend, or exit with a helpful error. -/// -/// Supports two modes: -/// - **Workspace mode**: CWD is inside the cuda-oxide repo (detected by -/// `crates/rustc-codegen-cuda` directory). Examples are resolved from the -/// workspace examples directory. -/// - **Standalone mode**: CWD has a `Cargo.toml` but is not inside the -/// workspace. The backend is located via cache or auto-fetch. Commands -/// like `run` operate on the current directory directly. -pub fn resolve_context() -> Context { - if let Some(workspace_root) = backend::find_workspace_root() { - let codegen_crate = workspace_root.join("crates/rustc-codegen-cuda"); - let examples_dir = codegen_crate.join("examples"); - let config = load_oxide_config(&workspace_root); - let backend_so = backend::find_or_build_backend(&workspace_root, config.backend.as_deref()); - return Context { - workspace_root, - codegen_crate, - examples_dir, - backend_so, - is_workspace: true, - config, - }; - } - - let cwd = std::env::current_dir().unwrap_or_else(|e| { - eprintln!("Error: cannot determine current directory: {}", e); - std::process::exit(1); - }); - - if cwd.join("Cargo.toml").is_file() { - let config = load_oxide_config(&cwd); - let backend_so = backend::find_or_build_backend(&cwd, config.backend.as_deref()); - return Context { - workspace_root: cwd.clone(), - codegen_crate: cwd.clone(), - examples_dir: cwd.clone(), - backend_so, - is_workspace: false, - config, - }; - } - - eprintln!("Error: Could not find cuda-oxide workspace or a standalone Cargo.toml."); - eprintln!(); - eprintln!("Run from inside the cuda-oxide repository, or from a project created"); - eprintln!("with `cargo oxide new `."); - std::process::exit(1); -} - -/// Resolve a context for commands that must not build or fetch the backend. -/// -/// Identical discovery to [`resolve_context`], except the backend `.so` is -/// only located via [`backend::backend_so_candidate`], never built and never -/// cloned, and an invalid `.cargo/cuda-oxide.toml` degrades to defaults with -/// a warning instead of exiting (so `doctor` can report it as a failed -/// check). Passive commands such as `doctor` and `clean` must remain usable -/// without triggering backend setup or network access. -/// `run`/`build`/`pipeline`/`setup` still build the backend on demand. -pub fn resolve_passive_context() -> Context { - if let Some(workspace_root) = backend::find_workspace_root() { - let codegen_crate = workspace_root.join("crates/rustc-codegen-cuda"); - let examples_dir = codegen_crate.join("examples"); - let config = load_oxide_config_lenient(&workspace_root); - let backend_so = backend::backend_so_candidate(&workspace_root, config.backend.as_deref()); - return Context { - workspace_root, - codegen_crate, - examples_dir, - backend_so, - is_workspace: true, - config, - }; - } - - let cwd = std::env::current_dir().unwrap_or_else(|e| { - eprintln!("Error: cannot determine current directory: {}", e); - std::process::exit(1); - }); - - if cwd.join("Cargo.toml").is_file() { - let config = load_oxide_config_lenient(&cwd); - let backend_so = backend::backend_so_candidate(&cwd, config.backend.as_deref()); - return Context { - workspace_root: cwd.clone(), - codegen_crate: cwd.clone(), - examples_dir: cwd.clone(), - backend_so, - is_workspace: false, - config, - }; - } - - eprintln!("Error: Could not find cuda-oxide workspace or a standalone Cargo.toml."); - eprintln!(); - eprintln!("Run from inside the cuda-oxide repository, or from a project created"); - eprintln!("with `cargo oxide new `."); - std::process::exit(1); -} - -#[derive(Debug, Clone, Eq, PartialEq)] -struct ExampleInfo { - name: String, - title: String, - description: String, - requirements: Vec, -} - -#[derive(Debug, Default, Eq, PartialEq)] -struct ParsedReadme { - title: Option, - description: Option, - requirements: Vec, -} - -#[derive(Debug, Default, Eq, PartialEq)] -struct ManifestInfo { - description: Option, -} - -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/."); - std::process::exit(1); - } - - let examples = discover_examples(&ctx.examples_dir).unwrap_or_else(|error| { - eprintln!("Error: {error}"); - std::process::exit(1); - }); - - let output = if json { - format_examples_json(&examples).unwrap_or_else(|error| { - eprintln!("Error: could not serialize example list: {error}"); - std::process::exit(1); - }) - } else { - format_examples_human(&examples) - }; - - print!("{output}"); -} - -fn discover_examples(examples_dir: &Path) -> Result, String> { - let entries = fs::read_dir(examples_dir).map_err(|error| { - format!( - "could not read examples directory {}: {error}", - examples_dir.display() - ) - })?; - - let mut examples = Vec::new(); - - for entry in entries { - let entry = entry.map_err(|error| { - format!( - "could not read an entry under {}: {error}", - examples_dir.display() - ) - })?; - - let file_type = entry - .file_type() - .map_err(|error| format!("could not inspect {}: {error}", entry.path().display()))?; - - if !file_type.is_dir() { - continue; - } - - let example_dir = entry.path(); - let name = entry.file_name().into_string().map_err(|name| { - format!( - "example directory name is not valid UTF-8: {}", - name.to_string_lossy() - ) - })?; - - // A directory without a manifest is not an example (scratch dirs, - // checked-out tooling, ...). Skip it instead of failing the listing. - let manifest_path = example_dir.join("Cargo.toml"); - if !manifest_path.is_file() { - eprintln!( - "Warning: skipping {}: no top-level Cargo.toml", - example_dir.display() - ); - continue; - } - - let manifest = parse_example_manifest(&manifest_path)?; - - let readme_path = example_dir.join("README.md"); - let parsed_readme = if readme_path.is_file() { - let contents = fs::read_to_string(&readme_path) - .map_err(|error| format!("could not read {}: {error}", readme_path.display()))?; - parse_example_readme(&name, &contents) - } else { - ParsedReadme::default() - }; - - let ParsedReadme { - title, - description, - requirements, - } = parsed_readme; - - let title = title.unwrap_or_else(|| name.clone()); - - let description = description.or(manifest.description).unwrap_or_else(|| { - if title != name { - title.clone() - } else { - "No description documented.".to_string() - } - }); - - examples.push(ExampleInfo { - name, - title, - description, - requirements, - }); - } - - examples.sort_by(|left, right| left.name.cmp(&right.name)); - - if examples.is_empty() { - return Err(format!( - "no examples found under {}", - examples_dir.display() - )); - } - - Ok(examples) -} - -fn parse_example_manifest(path: &Path) -> Result { - let contents = fs::read_to_string(path) - .map_err(|error| format!("could not read {}: {error}", path.display()))?; - - let manifest: toml::Value = toml::from_str(&contents) - .map_err(|error| format!("could not parse {}: {error}", path.display()))?; - - let package = manifest - .get("package") - .and_then(toml::Value::as_table) - .ok_or_else(|| format!("{} has no [package] table", path.display()))?; - - let description = package - .get("description") - .and_then(toml::Value::as_str) - .map(str::trim) - .filter(|description| !description.is_empty()) - .map(str::to_owned); - - Ok(ManifestInfo { description }) -} - -fn parse_example_readme(crate_name: &str, contents: &str) -> ParsedReadme { - let lines: Vec<&str> = contents.lines().collect(); - let mut headings = Vec::new(); - let mut in_code_fence = false; - - for (index, line) in lines.iter().enumerate() { - if is_code_fence(line) { - in_code_fence = !in_code_fence; - continue; - } - - if in_code_fence { - continue; - } - - if let Some((level, title)) = parse_markdown_heading(line) { - headings.push((index, level, title)); - } - } - - let crate_heading = normalize_heading(crate_name); - - let selected_heading = match headings.first() { - Some(first) if normalize_heading(&first.2) == crate_heading => headings - .get(1) - .filter(|heading| { - heading.1 == 2 - && first_prose_paragraph_in_range(&lines, first.0 + 1, heading.0).is_none() - && !is_generic_section_heading(&normalize_heading(&heading.2)) - }) - .or(Some(first)), - Some(first) => Some(first), - None => None, - }; - - let title = selected_heading - .map(|(_, _, title)| strip_inline_markdown(title)) - .filter(|title| !title.is_empty()); - - let description_start = selected_heading.map(|(index, _, _)| index + 1).unwrap_or(0); - - let description_end = headings - .iter() - .find(|(index, _, _)| *index >= description_start) - .map(|(index, _, _)| *index) - .unwrap_or(lines.len()); - - let description = first_prose_paragraph_in_range(&lines, description_start, description_end); - let requirements = extract_requirements(&lines); - - ParsedReadme { - title, - description, - requirements, - } -} - -fn parse_markdown_heading(line: &str) -> Option<(usize, String)> { - let trimmed = line.trim_start(); - let level = trimmed.bytes().take_while(|byte| *byte == b'#').count(); - - if !(1..=6).contains(&level) { - return None; - } - - let remainder = &trimmed[level..]; - if !remainder.chars().next().is_some_and(char::is_whitespace) { - return None; - } - - let title = remainder.trim().trim_end_matches('#').trim(); - - if title.is_empty() { - None - } else { - Some((level, title.to_string())) - } -} - -fn is_code_fence(line: &str) -> bool { - let trimmed = line.trim_start(); - trimmed.starts_with("```") || trimmed.starts_with("~~~") -} - -fn strip_inline_markdown(value: &str) -> String { - value - .replace("**", "") - .replace("__", "") - .replace('`', "") - .trim() - .to_string() -} - -fn normalize_heading(value: &str) -> String { - strip_inline_markdown(value) - .trim_matches(|character: char| { - character == ':' || character == '-' || character.is_whitespace() - }) - .to_ascii_lowercase() -} - -fn is_generic_section_heading(heading: &str) -> bool { - matches!( - heading, - "overview" - | "what this example does" - | "key concepts" - | "key concepts demonstrated" - | "build" - | "build and run" - | "usage" - | "expected output" - | "requirements" - | "hardware requirements" - | "prerequisites" - | "potential errors" - | "how it works" - | "how it works under the hood" - | "generated ptx" - | "run" - | "test" - | "tests" - | "correctness" - | "trigger" - | "kernels" - | "features tested" - | "what this tests" - | "what it tests" - | "what this demonstrates" - | "why this exists" - | "the bug" - | "final design" - ) -} - -fn first_prose_paragraph_in_range(lines: &[&str], start: usize, end: usize) -> Option { - let end = end.min(lines.len()); - let start = start.min(end); - let mut paragraph = Vec::new(); - let mut in_code_fence = false; - - for line in &lines[start..end] { - if is_code_fence(line) { - if !paragraph.is_empty() { - break; - } - in_code_fence = !in_code_fence; - continue; - } - - if in_code_fence { - continue; - } - - let trimmed = line.trim(); - - if parse_markdown_heading(trimmed).is_some() { - break; - } - - if trimmed.is_empty() { - if !paragraph.is_empty() { - break; - } - continue; - } - - if is_non_prose_markdown(trimmed) { - if !paragraph.is_empty() { - break; - } - continue; - } - - paragraph.push(trimmed); - } - - if paragraph.is_empty() { - None - } else { - Some(paragraph.join(" ")) - } -} - -fn is_non_prose_markdown(line: &str) -> bool { - line.starts_with("- ") - || line.starts_with("* ") - || line.starts_with("+ ") - || line.starts_with('>') - || line.starts_with('|') - || line.starts_with("![") - || line.starts_with("