From ffce9c0e53322d2b023c4d667e57a0aa2b0be36b Mon Sep 17 00:00:00 2001 From: Raul Estrada Date: Sat, 22 Aug 2026 22:22:08 -0600 Subject: [PATCH] feat(wgmma): add two-slot F16 partial-wait pipeline lowering Signed-off-by: Raul Estrada --- crates/cuda-device/src/wgmma.rs | 28 ++-- .../tests/wgmma_pipeline_carrier_ptx.rs | 50 ++++-- crates/dialect-nvvm/src/ops/wgmma.rs | 116 ++++++++++++- crates/dialect-nvvm/tests/ops_test.rs | 27 ++- .../translator/terminator/intrinsics/wgmma.rs | 9 +- .../mir-lower/src/convert/interface_impls.rs | 19 ++- .../mir-lower/src/convert/intrinsics/wgmma.rs | 60 +++++-- .../src/wgmma_deferred_accumulator.rs | 79 ++++++--- crates/mir-lower/tests/lowering_test.rs | 158 +++++++++++++++++- .../advanced/matrix-multiply-accelerators.md | 19 ++- .../appendix/supported-features.md | 2 +- cuda-oxide-book/compiler/mlir-dialects.md | 6 +- 12 files changed, 488 insertions(+), 85 deletions(-) diff --git a/crates/cuda-device/src/wgmma.rs b/crates/cuda-device/src/wgmma.rs index aa9ed3afb0..8d1cbe8a87 100644 --- a/crates/cuda-device/src/wgmma.rs +++ b/crates/cuda-device/src/wgmma.rs @@ -19,7 +19,8 @@ //! `N + 1` independent accumulator slots. //! //! The `m64n64k16.f32.f16.f16` variant supports canonical linear full-drain -//! regions and the canonical counted K-loop. The `m64n64k8.f32.tf32.tf32` +//! regions, the canonical counted K-loop, and a two-slot straight-line +//! `wait_group<1>` pipeline. The `m64n64k8.f32.tf32.tf32` //! variant uses the canonical linear full-drain carrier only. The //! `m64n128k16.f32.bf16.bf16` variant supports canonical linear full-drain //! regions with a 64-value `[[f32; 8]; 8]` accumulator. Every accepted @@ -83,8 +84,7 @@ //! no pointer fallback, counted-loop lowering, or partial-wait pipeline for //! this shape. //! - **F16 linear full drain:** the m64n64 canonical accumulator and full-drain -//! lifetime are supported, but there is no pointer fallback or partial-wait -//! pipeline for F16. +//! lifetime are supported, with no pointer fallback. //! - **TF32 linear full drain:** `m64n64k8.f32.tf32.tf32` uses the same //! canonical accumulator and full-drain lifetime. TF32 has no pointer //! fallback, counted-loop lowering, or partial-wait pipeline. @@ -92,16 +92,19 @@ //! compile-time trip count, and `u64` descriptor recurrences of the form `desc + const`. The //! fence-to-final-wait lifetime is fused so the accumulator is loaded and //! stored once for the whole loop rather than once per iteration. -//! - **Partial-wait pipeline:** a static `wgmma_wait_group::()` with -//! `1 <= N <= 7`, exactly `N + 1` distinct canonical accumulator slots, one -//! MMA per committed group, round-robin slot reuse only after the matching -//! partial wait, and a mandatory final `wgmma_wait_group::<0>()`. +//! - **Partial-wait pipeline:** BF16 supports a static `wgmma_wait_group::()` +//! with `1 <= N <= 7` and exactly `N + 1` distinct canonical accumulator +//! slots. F16 supports only `wgmma_wait_group::<1>()` with exactly two slots. +//! Both forms require one MMA per committed group, round-robin slot reuse only +//! after the matching partial wait, and a mandatory final +//! `wgmma_wait_group::<0>()`. //! - Accumulator values must not be read or written while their asynchronous //! lifetime is pending. //! - Unsupported BF16 full-drain accumulator pointer shapes retain the deferred //! pointer-form lowering. Dynamic partial waits, unsupported control flow, -//! malformed pipeline schedules, F16 partial-wait/pipelined shapes, TF32 -//! counted-loop or partial-wait shapes, and the legacy K=16 TF32 +//! malformed pipeline schedules, F16 partial-wait depths other than one, +//! F16 counted pipelines, TF32 counted-loop or partial-wait shapes, and the +//! legacy K=16 TF32 //! compatibility entry point are rejected. //! //! # Hardware Support @@ -282,10 +285,11 @@ pub unsafe fn wgmma_mma_m64n128k16_f32_bf16(acc: &mut [[f32; 8]; 8], desc_a: u64 /// WGMMA with f32 accumulator and f16 inputs. /// -/// This variant supports canonical linear full-drain regions and canonical +/// This variant supports canonical linear full-drain regions, canonical /// counted K-loops with a compile-time trip count and affine `u64` descriptor -/// recurrences. The accumulator must have the public `[[f32; 8]; 4]` shape. -/// Partial waits and non-canonical pointer fallback remain unsupported for F16. +/// recurrences, and the straight-line two-slot `wait_group<1>` pipeline. The +/// accumulator must have the public `[[f32; 8]; 4]` shape. Deeper partial waits, +/// counted pipelines, and non-canonical pointer fallback remain unsupported. /// /// # PTX /// diff --git a/crates/cuda-oxide-codegen/tests/wgmma_pipeline_carrier_ptx.rs b/crates/cuda-oxide-codegen/tests/wgmma_pipeline_carrier_ptx.rs index 1e0eb8b99e..32d7b85f54 100644 --- a/crates/cuda-oxide-codegen/tests/wgmma_pipeline_carrier_ptx.rs +++ b/crates/cuda-oxide-codegen/tests/wgmma_pipeline_carrier_ptx.rs @@ -3,8 +3,9 @@ * SPDX-License-Identifier: Apache-2.0 */ -//! Proves that a two-slot BF16 WGMMA pipeline with `wait_group<1>` can keep -//! sixty-four tied `f32` accumulator values spill-free through code generation. +//! Proves that two-slot BF16 and F16 WGMMA pipelines with `wait_group<1>` can +//! keep sixty-four tied `f32` accumulator values spill-free through code +//! generation. #![cfg(unix)] @@ -43,7 +44,11 @@ fn accumulator_register_list(slot: usize) -> String { .join(", ") } -fn build_wgmma_pipeline_carrier_kernel(module: &mut CodegenModule) { +fn build_wgmma_pipeline_carrier_kernel( + module: &mut CodegenModule, + kernel_name: &str, + mnemonic: &str, +) { module.edit(|ctx, module| { let module_region = module.get_operation().deref(ctx).get_region(0); let module_block = module_region @@ -75,7 +80,7 @@ fn build_wgmma_pipeline_carrier_kernel(module: &mut CodegenModule) { 1, ); let function = MirFuncOp::new(ctx, function_op, TypeAttr::new(function_type.into())); - function.set_symbol_name(ctx, "wgmma_pipeline_carrier_kernel".try_into().unwrap()); + function.set_symbol_name(ctx, kernel_name.try_into().unwrap()); let function_region = function.get_operation().deref(ctx).get_region(0); let entry = BasicBlock::new(ctx, None, argument_types); @@ -95,10 +100,10 @@ fn build_wgmma_pipeline_carrier_kernel(module: &mut CodegenModule) { let template = format!( "{{\n\ \x20 wgmma.fence.sync.aligned;\n\ - \x20 wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16 \ + \x20 {mnemonic} \ {{{slot0}}}, ${descriptor_base}, ${}, 1, 1, 1, 0, 0;\n\ \x20 wgmma.commit_group.sync.aligned;\n\ - \x20 wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16 \ + \x20 {mnemonic} \ {{{slot1}}}, ${}, ${}, 1, 1, 1, 0, 0;\n\ \x20 wgmma.commit_group.sync.aligned;\n\ \x20 wgmma.wait_group.sync.aligned 1;\n\ @@ -185,7 +190,7 @@ fn build_wgmma_pipeline_carrier_kernel(module: &mut CodegenModule) { }); module - .mark_kernel_entry("wgmma_pipeline_carrier_kernel") + .mark_kernel_entry(kernel_name) .expect("kernel entry must be marked"); } @@ -211,10 +216,9 @@ fn used_register_count(ptxas_stderr: &str) -> Option { }) } -#[test] -fn bf16_wgmma_wait_group_one_keeps_two_accumulator_slots_without_spills() { - let mut module = CodegenModule::new("wgmma_pipeline_carrier").unwrap(); - build_wgmma_pipeline_carrier_kernel(&mut module); +fn run_wgmma_wait_group_one_pipeline_probe(kernel_name: &str, mnemonic: &str) { + let mut module = CodegenModule::new(kernel_name).unwrap(); + build_wgmma_pipeline_carrier_kernel(&mut module, kernel_name, mnemonic); let compiler = Compiler::discover().expect("LLVM 21+ llc/opt must be installed"); let options = CompileOptions::new(Target::parse("sm_90a").unwrap()); @@ -229,15 +233,11 @@ fn bf16_wgmma_wait_group_one_keeps_two_accumulator_slots_without_spills() { "PTX must target sm_90a:\n{text}" ); assert!( - text.contains("wgmma_pipeline_carrier_kernel"), + text.contains(kernel_name), "pipeline carrier kernel is missing:\n{text}" ); assert_eq!(text.matches("wgmma.fence.sync.aligned").count(), 1); - assert_eq!( - text.matches("wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16") - .count(), - 2 - ); + assert_eq!(text.matches(mnemonic).count(), 2); assert_eq!(text.matches("wgmma.commit_group.sync.aligned").count(), 2); assert_eq!(text.matches("wgmma.wait_group.sync.aligned 1").count(), 1); assert_eq!(text.matches("wgmma.wait_group.sync.aligned 0").count(), 1); @@ -299,3 +299,19 @@ fn bf16_wgmma_wait_group_one_keeps_two_accumulator_slots_without_spills() { eprintln!("{stderr}"); } + +#[test] +fn bf16_wgmma_wait_group_one_keeps_two_accumulator_slots_without_spills() { + run_wgmma_wait_group_one_pipeline_probe( + "wgmma_bf16_pipeline_carrier_kernel", + "wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16", + ); +} + +#[test] +fn f16_wgmma_wait_group_one_keeps_two_accumulator_slots_without_spills() { + run_wgmma_wait_group_one_pipeline_probe( + "wgmma_f16_pipeline_carrier_kernel", + "wgmma.mma_async.sync.aligned.m64n64k16.f32.f16.f16", + ); +} diff --git a/crates/dialect-nvvm/src/ops/wgmma.rs b/crates/dialect-nvvm/src/ops/wgmma.rs index bd8f3faf7b..04844973c7 100644 --- a/crates/dialect-nvvm/src/ops/wgmma.rs +++ b/crates/dialect-nvvm/src/ops/wgmma.rs @@ -22,10 +22,11 @@ //! accumulator to LLVM. The pointer-form group remains the deferred fallback. //! //! F16 uses the same 32-value accumulator model for canonical linear -//! full-drain regions and the canonical counted K-loop. TF32 uses the same -//! carrier only for canonical linear full-drain regions, with the hardware -//! `m64n64k8` shape. Neither F16 nor TF32 has a deferred pointer-form group or -//! partial-wait pipeline carrier, and TF32 has no counted-loop carrier. +//! full-drain regions, the canonical counted K-loop, and the narrow two-slot +//! straight-line `wait_group<1>` pipeline. TF32 uses the same carrier only for +//! canonical linear full-drain regions, with the hardware `m64n64k8` shape. +//! F16 and TF32 have no deferred pointer-form group, TF32 has no partial-wait +//! or counted-loop carrier, and F16 counted pipelines remain unsupported. use dialect_mir::types::{MirPtrType, address_space}; use pliron::{ @@ -1335,6 +1336,112 @@ impl Verify for WgmmaMmaPipelineValuesM64N64K16F32Bf16Op { } } +/// Value-form F16 WGMMA pipeline for the supported two-slot `wait_group<1>` shape. +/// +/// This carrier deliberately models only the first F16 partial-wait topology: +/// two independent 32-value accumulator slots, one or more committed groups +/// issued round-robin across those slots, `wait_group<1>` before every reuse, +/// and a final `wait_group<0>` before results become visible to LLVM. +/// +/// Operand layout: +/// +/// ```text +/// [ +/// slot0_acc_0, ..., slot0_acc_31, +/// slot1_acc_0, ..., slot1_acc_31, +/// desc_a_0, desc_b_0, ..., desc_a_g, desc_b_g, +/// ] +/// ``` +/// +/// Result layout mirrors the 64 flattened accumulator inputs. +#[pliron_op(name = "nvvm.wgmma_mma_pipeline_values_m64n64k16_f32_f16", format)] +pub struct WgmmaMmaPipelineValuesM64N64K16F32F16Op; + +impl WgmmaMmaPipelineValuesM64N64K16F32F16Op { + /// Wrap an existing operation pointer. + pub fn new(op: Ptr) -> Self { + Self { op } + } + + /// Build the supported two-slot F16 partial-wait pipeline. + pub fn build( + ctx: &mut Context, + accumulators: Vec, + descriptors: Vec, + ) -> Ptr { + let f32_ty = FP32Type::get(ctx); + let result_count = accumulators.len(); + let mut operands = Vec::with_capacity(result_count + descriptors.len()); + operands.extend(accumulators); + operands.extend(descriptors); + + Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![f32_ty.into(); result_count], + operands, + vec![], + 0, + ) + } +} + +impl Verify for WgmmaMmaPipelineValuesM64N64K16F32F16Op { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + let result_count = op.get_num_results(); + let expected_results = 2 * WGMMA_M64N64_F32_ACCUMULATOR_COUNT; + if result_count != expected_results { + return verify_err!( + op.loc(), + "nvvm.wgmma_mma_pipeline_values_m64n64k16_f32_f16 requires exactly two 32-value accumulator slots" + ); + } + + let operand_count = op.get_num_operands(); + if operand_count < result_count + 4 { + return verify_err!( + op.loc(), + "nvvm.wgmma_mma_pipeline_values_m64n64k16_f32_f16 requires at least two committed groups" + ); + } + let descriptor_count = operand_count - result_count; + if !descriptor_count.is_multiple_of(2) { + return verify_err!( + op.loc(), + "nvvm.wgmma_mma_pipeline_values_m64n64k16_f32_f16 descriptors must form pairs" + ); + } + if descriptor_count / 2 < 2 { + return verify_err!( + op.loc(), + "nvvm.wgmma_mma_pipeline_values_m64n64k16_f32_f16 requires at least two committed groups" + ); + } + + for accumulator_index in 0..result_count { + if !is_f32(ctx, op.get_operand(accumulator_index).get_type(ctx)) + || !is_f32(ctx, op.get_result(accumulator_index).get_type(ctx)) + { + return verify_err!( + op.loc(), + "nvvm.wgmma_mma_pipeline_values_m64n64k16_f32_f16 accumulator operands and results must be f32" + ); + } + } + for descriptor_index in result_count..operand_count { + if !is_u64(ctx, op.get_operand(descriptor_index).get_type(ctx)) { + return verify_err!( + op.loc(), + "nvvm.wgmma_mma_pipeline_values_m64n64k16_f32_f16 descriptors must be u64" + ); + } + } + + Ok(()) + } +} + /// Register WGMMA operations with the context. pub(super) fn register(ctx: &mut Context) { WgmmaMaxPendingAttr::register(ctx); @@ -1353,4 +1460,5 @@ pub(super) fn register(ctx: &mut Context) { WgmmaMmaLoopValuesM64N64K16F32F16Op::register(ctx); WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op::register(ctx); WgmmaMmaPipelineValuesM64N64K16F32Bf16Op::register(ctx); + WgmmaMmaPipelineValuesM64N64K16F32F16Op::register(ctx); } diff --git a/crates/dialect-nvvm/tests/ops_test.rs b/crates/dialect-nvvm/tests/ops_test.rs index aaf2be5014..22957253f0 100644 --- a/crates/dialect-nvvm/tests/ops_test.rs +++ b/crates/dialect-nvvm/tests/ops_test.rs @@ -44,7 +44,7 @@ use dialect_nvvm::ops::{ WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op, WgmmaMmaLoopValuesM64N64K16F32Bf16Op, WgmmaMmaLoopValuesM64N64K16F32F16Op, WgmmaMmaM64N64K8F32Tf32Op, WgmmaMmaM64N64K16F32Bf16Op, WgmmaMmaM64N64K16F32F16Op, WgmmaMmaM64N128K16F32Bf16Op, - WgmmaMmaPipelineValuesM64N64K16F32Bf16Op, + WgmmaMmaPipelineValuesM64N64K16F32Bf16Op, WgmmaMmaPipelineValuesM64N64K16F32F16Op, }; #[test] @@ -106,6 +106,7 @@ fn handwritten_ops_match_reviewed_allowlist() { ("wgmma.rs", "WgmmaMmaLoopValuesM64N64K16F32F16Op"), ("wgmma.rs", "WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op"), ("wgmma.rs", "WgmmaMmaPipelineValuesM64N64K16F32Bf16Op"), + ("wgmma.rs", "WgmmaMmaPipelineValuesM64N64K16F32F16Op"), ]; expected.sort_unstable(); found.sort_unstable(); @@ -5190,6 +5191,30 @@ fn handwritten_ffi_and_wgmma_carriers_verify_exact_shapes() { .is_err() ); + let f16_pipeline_group = WgmmaMmaPipelineValuesM64N64K16F32F16Op::build( + &mut ctx, + vec![f32_value; 64], + vec![u64_value; 8], + ); + { + let pipeline_ref = f16_pipeline_group.deref(&ctx); + assert_eq!(pipeline_ref.get_num_operands(), 72); + assert_eq!(pipeline_ref.get_num_results(), 64); + } + let f16_pipeline = WgmmaMmaPipelineValuesM64N64K16F32F16Op::new(f16_pipeline_group); + assert!(f16_pipeline.verify(&ctx).is_ok()); + + let f16_too_few_pipeline_slots = WgmmaMmaPipelineValuesM64N64K16F32F16Op::build( + &mut ctx, + vec![f32_value; 32], + vec![u64_value; 4], + ); + assert!( + WgmmaMmaPipelineValuesM64N64K16F32F16Op::new(f16_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()); diff --git a/crates/mir-importer/src/translator/terminator/intrinsics/wgmma.rs b/crates/mir-importer/src/translator/terminator/intrinsics/wgmma.rs index daa4762c11..026b036706 100644 --- a/crates/mir-importer/src/translator/terminator/intrinsics/wgmma.rs +++ b/crates/mir-importer/src/translator/terminator/intrinsics/wgmma.rs @@ -272,9 +272,10 @@ pub fn emit_wgmma_mma_m64n128k16_f32_bf16( /// Emit F16 m64n64k16 WGMMA pointer form. /// -/// `mir-lower` accepts this variant only in a canonical linear full-drain -/// region ending in `wait_group<0>`. Counted loops and partial-wait pipelines -/// remain BF16-only. +/// `mir-lower` accepts this variant in a canonical linear full-drain region, +/// the canonical counted K-loop, or the narrow two-slot straight-line +/// `wait_group<1>` pipeline. Deeper partial waits and counted pipelines remain +/// unsupported for F16. #[allow(clippy::too_many_arguments)] pub fn emit_wgmma_mma_m64n64k16_f32_f16( ctx: &mut Context, @@ -315,7 +316,7 @@ pub fn emit_wgmma_mma_m64n64k16_f32_f16( /// /// `mir-lower` accepts this variant only in a canonical linear full-drain /// region ending in `wait_group<0>`. Counted loops, partial-wait pipelines, -/// and pointer fallback remain BF16-only. +/// and pointer fallback remain unsupported for TF32. pub fn emit_wgmma_mma_m64n64k8_f32_tf32( ctx: &mut Context, body: &mir::Body, diff --git a/crates/mir-lower/src/convert/interface_impls.rs b/crates/mir-lower/src/convert/interface_impls.rs index 4fc76f9a73..940616acdd 100644 --- a/crates/mir-lower/src/convert/interface_impls.rs +++ b/crates/mir-lower/src/convert/interface_impls.rs @@ -45,7 +45,7 @@ use dialect_nvvm::ops::{ WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op, WgmmaMmaLoopValuesM64N64K16F32Bf16Op, WgmmaMmaLoopValuesM64N64K16F32F16Op, WgmmaMmaM64N64K8F32Tf32Op, WgmmaMmaM64N64K16F32Bf16Op, WgmmaMmaM64N64K16F32F16Op, WgmmaMmaM64N128K16F32Bf16Op, - WgmmaMmaPipelineValuesM64N64K16F32Bf16Op, + WgmmaMmaPipelineValuesM64N64K16F32Bf16Op, WgmmaMmaPipelineValuesM64N64K16F32F16Op, }; // ---- Arithmetic ops -------------------------------------------------------- @@ -1263,6 +1263,23 @@ impl MirToLlvmConversion for WgmmaMmaPipelineValuesM64N64K16F32Bf16Op { } } +#[op_interface_impl] +impl MirToLlvmConversion for WgmmaMmaPipelineValuesM64N64K16F32F16Op { + fn convert( + &self, + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + operands_info: &OperandsInfo, + ) -> Result<()> { + super::intrinsics::wgmma::convert_mma_pipeline_values( + ctx, + rewriter, + self.get_operation(), + operands_info, + ) + } +} + // ---- NVVM WMMA ops --------------------------------------------------------- // ---- NVVM Atomic ops ------------------------------------------------------- diff --git a/crates/mir-lower/src/convert/intrinsics/wgmma.rs b/crates/mir-lower/src/convert/intrinsics/wgmma.rs index 29eeb935f8..5898838de6 100644 --- a/crates/mir-lower/src/convert/intrinsics/wgmma.rs +++ b/crates/mir-lower/src/convert/intrinsics/wgmma.rs @@ -8,6 +8,7 @@ use crate::convert::intrinsics::common::*; use dialect_nvvm::ops::{ WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op, WgmmaMmaPipelineValuesM64N64K16F32Bf16Op, + WgmmaMmaPipelineValuesM64N64K16F32F16Op, }; use llvm_export::ops as llvm; use llvm_export::types::{self as llvm_types, VoidType}; @@ -39,7 +40,7 @@ impl WgmmaInputKind { match self { Self::Bf16 => "bf16.bf16", Self::F16 => "f16.f16", - Self::Tf32 => unreachable!("TF32 has no counted K-loop lowering"), + Self::Tf32 => unreachable!("TF32 has no m64n64k16 WGMMA lowering"), } } } @@ -298,9 +299,15 @@ fn pipeline_accumulator_operand_list(slot: usize) -> String { .join(", ") } -fn pipeline_template(slot_count: usize, group_count: usize, max_pending_groups: u8) -> String { +fn pipeline_template( + slot_count: usize, + group_count: usize, + max_pending_groups: u8, + input_kind: WgmmaInputKind, +) -> String { let result_count = slot_count * VALUE_ACCUMULATOR_COUNT; let descriptor_base = result_count * 2; + let input_types = input_kind.ptx_suffix(); let mut template = String::from("{\n wgmma.fence.sync.aligned;\n"); for group_index in 0..group_count { @@ -309,7 +316,7 @@ fn pipeline_template(slot_count: usize, group_count: usize, max_pending_groups: let desc_a = descriptor_base + group_index * 2; let desc_b = desc_a + 1; template.push_str(&format!( - " wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16 \ + " wgmma.mma_async.sync.aligned.m64n64k16.f32.{input_types} \ {{{accumulators}}}, ${desc_a}, ${desc_b}, 1, 1, 1, 0, 0;\n" )); template.push_str(" wgmma.commit_group.sync.aligned;\n"); @@ -668,12 +675,14 @@ pub(crate) fn convert_mma_loop_pipeline_values( Ok(()) } -/// Lower a multi-slot BF16 WGMMA pipeline to one convergent inline-PTX scope. +/// Lower a supported multi-slot BF16 or F16 WGMMA pipeline to one convergent inline-PTX scope. /// /// Each accumulator slot owns 32 tied `f32` registers. Groups are committed /// independently and issued round-robin across `N + 1` slots for /// `wait_group`, ensuring a slot is not reused until its previous group has -/// completed. A final `wait_group<0>` occurs before any result escapes to LLVM. +/// completed. BF16 supports the existing static depths; F16 is deliberately +/// restricted to the two-slot `wait_group<1>` form. A final `wait_group<0>` +/// occurs before any result escapes to LLVM. pub(crate) fn convert_mma_pipeline_values( ctx: &mut Context, rewriter: &mut DialectConversionRewriter, @@ -699,10 +708,17 @@ pub(crate) fn convert_mma_pipeline_values( return pliron::input_err_noloc!("pipeline WGMMA descriptors must form pairs"); } - let pipeline = Operation::get_op::(op, ctx) - .expect("pipeline conversion must be invoked for the pipeline WGMMA op"); - let Some(max_pending_groups) = pipeline.max_pending_groups(ctx) else { - return pliron::input_err_noloc!("pipeline WGMMA is missing max_pending_groups"); + let (max_pending_groups, input_kind) = if let Some(pipeline) = + Operation::get_op::(op, ctx) + { + let Some(max_pending_groups) = pipeline.max_pending_groups(ctx) else { + return pliron::input_err_noloc!("pipeline WGMMA is missing max_pending_groups"); + }; + (max_pending_groups, WgmmaInputKind::Bf16) + } else if Operation::get_op::(op, ctx).is_some() { + (1, WgmmaInputKind::F16) + } else { + return pliron::input_err_noloc!("unsupported pipeline WGMMA carrier"); }; if !(1..=7).contains(&max_pending_groups) { return pliron::input_err_noloc!("pipeline WGMMA max_pending_groups must be in 1..=7"); @@ -720,7 +736,7 @@ pub(crate) fn convert_mma_pipeline_values( ); } - let template = pipeline_template(slot_count, group_count, max_pending_groups); + let template = pipeline_template(slot_count, group_count, max_pending_groups, input_kind); let constraints = pipeline_constraints(result_count, descriptor_count); let f32_ty = FP32Type::get(ctx); let struct_ty: TypeHandle = llvm_types::StructType::get_unnamed( @@ -1073,7 +1089,7 @@ mod tests { #[test] fn pipeline_template_throttles_groups_and_finishes_with_wait_zero() { - let template = pipeline_template(2, 4, 1); + let template = pipeline_template(2, 4, 1, WgmmaInputKind::Bf16); assert_eq!(template.matches("wgmma.mma_async").count(), 4); assert_eq!( template.matches("wgmma.commit_group.sync.aligned").count(), @@ -1092,6 +1108,28 @@ mod tests { assert!(!template.contains("ld.f32")); assert!(!template.contains("st.f32")); assert!(!template.contains(".reg .f32")); + assert!(template.contains("m64n64k16.f32.bf16.bf16")); + assert!(!template.contains("m64n64k16.f32.f16.f16")); + + let f16_template = pipeline_template(2, 4, 1, WgmmaInputKind::F16); + assert_eq!(f16_template.matches("wgmma.mma_async").count(), 4); + assert_eq!( + f16_template + .matches("wgmma.wait_group.sync.aligned 1") + .count(), + 3 + ); + assert_eq!( + f16_template + .matches("wgmma.wait_group.sync.aligned 0") + .count(), + 1 + ); + assert!(f16_template.contains("m64n64k16.f32.f16.f16")); + assert!(!f16_template.contains("m64n64k16.f32.bf16.bf16")); + assert!(!f16_template.contains("ld.f32")); + assert!(!f16_template.contains("st.f32")); + assert!(!f16_template.contains(".reg .f32")); let constraints = pipeline_constraints(64, 8); assert_eq!( diff --git a/crates/mir-lower/src/wgmma_deferred_accumulator.rs b/crates/mir-lower/src/wgmma_deferred_accumulator.rs index 4dbf80bbee..c95ab2c76d 100644 --- a/crates/mir-lower/src/wgmma_deferred_accumulator.rs +++ b/crates/mir-lower/src/wgmma_deferred_accumulator.rs @@ -16,11 +16,12 @@ //! accumulator and 64 scalar SSA values. Unsupported m64n64 BF16 accumulator //! shapes retain the existing deferred pointer fallback. //! F16 m64n64 is accepted for the canonical accumulator in linear full-drain -//! regions and the canonical single-slot counted K-loop. TF32 is accepted only -//! for the canonical m64n64 accumulator in a linear full-drain region, using -//! the hardware `m64n64k8.f32.tf32.tf32` shape. BF16 m64n128 is accepted only -//! for canonical linear full-drain regions. Partial waits, counted pipelines, -//! and pointer fallback remain m64n64-BF16-only. +//! regions, the canonical single-slot counted K-loop, and the narrow two-slot +//! straight-line `wait_group<1>` pipeline. TF32 is accepted only for the +//! canonical m64n64 accumulator in a linear full-drain region, using the +//! hardware `m64n64k8.f32.tf32.tf32` shape. BF16 m64n128 is accepted only for +//! canonical linear full-drain regions. Counted pipelines and pointer fallback +//! remain m64n64-BF16-only. //! //! Straight-line regions keep the existing shape: //! @@ -40,12 +41,12 @@ //! The complete asynchronous lifetime is then represented by one value-form loop //! operation so LLVM never sees an in-flight accumulator between iterations. //! -//! Straight-line BF16 partial-wait pipelines are recognized separately. A -//! static -//! `wait_group` with `N > 0` requires `N + 1` canonical accumulator slots. -//! Groups are committed and the slots are reused round-robin only after the -//! corresponding partial wait has made the oldest slot safe. Every accepted -//! pipeline ends with `wait_group<0>` before any accumulator value escapes. +//! Straight-line partial-wait pipelines are recognized separately. BF16 keeps +//! the existing static `wait_group` support, where `N > 0` requires `N + 1` +//! canonical accumulator slots. F16 deliberately admits only `wait_group<1>` +//! with two slots. Groups are committed and slots are reused round-robin only +//! after the corresponding partial wait has made the oldest slot safe. Every +//! accepted pipeline ends with `wait_group<0>` before any accumulator escapes. //! //! A second narrow loop form composes these properties for the production //! counted-pipeline depths validated on Hopper: two canonical accumulator slots @@ -70,7 +71,8 @@ use dialect_nvvm::ops::{ WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op, WgmmaMmaLoopValuesM64N64K16F32Bf16Op, WgmmaMmaLoopValuesM64N64K16F32F16Op, WgmmaMmaM64N64K8F32Tf32Op, WgmmaMmaM64N64K16F32Bf16Op, WgmmaMmaM64N64K16F32F16Op, WgmmaMmaM64N128K16F32Bf16Op, - WgmmaMmaPipelineValuesM64N64K16F32Bf16Op, WgmmaWaitGroupSyncAlignedOp, + WgmmaMmaPipelineValuesM64N64K16F32Bf16Op, WgmmaMmaPipelineValuesM64N64K16F32F16Op, + WgmmaWaitGroupSyncAlignedOp, }; use mir_transforms::analyses::{induction, loop_info::LoopInfo}; use pliron::{ @@ -133,6 +135,7 @@ struct PipelinePlan { accumulators: Vec, descriptors: Vec, max_pending_groups: u8, + kind: WgmmaMmaKind, } #[derive(Clone, Copy)] @@ -142,6 +145,7 @@ enum PipelineEvent { accumulator: Value, desc_a: Value, desc_b: Value, + kind: WgmmaMmaKind, }, Commit(Ptr), Wait { @@ -1587,7 +1591,15 @@ fn collect_pipeline_events( ); } - if Operation::get_op::(operation, ctx).is_some() { + let pipeline_kind = + if Operation::get_op::(operation, ctx).is_some() { + Some(WgmmaMmaKind::Bf16M64N64) + } else if Operation::get_op::(operation, ctx).is_some() { + Some(WgmmaMmaKind::F16M64N64) + } else { + None + }; + if let Some(kind) = pipeline_kind { require_pointer_mma_shape(ctx, operation)?; let operation_ref = operation.deref(ctx); let accumulator = operation_ref.get_operand(0); @@ -1597,6 +1609,7 @@ fn collect_pipeline_events( accumulator, desc_a: operation_ref.get_operand(1), desc_b: operation_ref.get_operand(2), + kind, }); continue; } @@ -1675,6 +1688,7 @@ fn validate_pipeline_events( let mut waits = Vec::new(); let mut partial_wait_positions = Vec::new(); let mut partial_wait_value = None; + let mut pipeline_kind = None; let mut index = 0usize; let mut saw_final_wait = false; @@ -1684,12 +1698,22 @@ fn validate_pipeline_events( accumulator, desc_a, desc_b, + kind, } = events[index] else { return pliron::input_err_noloc!( "WGMMA pipelined region requires exactly one MMA before each commit_group" ); }; + match pipeline_kind { + Some(expected) if expected != kind => { + return pliron::input_err_noloc!( + "one WGMMA partial-wait pipeline cannot mix MMA variants or shapes" + ); + } + None => pipeline_kind = Some(kind), + _ => {} + } index += 1; let Some(PipelineEvent::Commit(commit)) = events.get(index).copied() else { @@ -1761,6 +1785,12 @@ fn validate_pipeline_events( ); } let max_pending_groups = partial_wait_value.expect("pipeline collector saw a partial wait"); + let kind = pipeline_kind.expect("pipeline validation saw at least one MMA"); + if kind == WgmmaMmaKind::F16M64N64 && max_pending_groups != 1 { + return pliron::input_err_noloc!( + "F16 WGMMA partial-wait pipeline supports only wait_group<1> with two accumulator slots" + ); + } let slot_count = usize::try_from(max_pending_groups + 1).expect("wait_group depth in 1..=7 must fit usize"); if group_accumulators.len() < slot_count { @@ -1812,6 +1842,7 @@ fn validate_pipeline_events( descriptors, max_pending_groups: u8::try_from(max_pending_groups) .expect("wait_group depth in 1..=7 must fit u8"), + kind, }) } @@ -1980,6 +2011,7 @@ fn apply_pipeline_plan(ctx: &mut Context, plan: PipelinePlan) { accumulators, descriptors, max_pending_groups, + kind, } = plan; let final_wait = *waits @@ -2004,12 +2036,21 @@ fn apply_pipeline_plan(ctx: &mut Context, plan: PipelinePlan) { all_accumulator_values.extend(accumulator_values); } - let pipeline = WgmmaMmaPipelineValuesM64N64K16F32Bf16Op::build( - ctx, - all_accumulator_values, - descriptors, - max_pending_groups, - ); + let pipeline = match kind { + WgmmaMmaKind::Bf16M64N64 => WgmmaMmaPipelineValuesM64N64K16F32Bf16Op::build( + ctx, + all_accumulator_values, + descriptors, + max_pending_groups, + ), + WgmmaMmaKind::F16M64N64 => { + debug_assert_eq!(max_pending_groups, 1); + WgmmaMmaPipelineValuesM64N64K16F32F16Op::build(ctx, all_accumulator_values, descriptors) + } + WgmmaMmaKind::Tf32M64N64 | WgmmaMmaKind::Bf16M64N128 => { + unreachable!("pipeline validation admits only m64n64 BF16/F16") + } + }; pipeline.deref_mut(ctx).set_loc(loc.clone()); let result_count = accumulators.len() * ACCUMULATOR_LEN; let accumulator_results = (0..result_count) diff --git a/crates/mir-lower/tests/lowering_test.rs b/crates/mir-lower/tests/lowering_test.rs index 06a47809bc..369db46c43 100644 --- a/crates/mir-lower/tests/lowering_test.rs +++ b/crates/mir-lower/tests/lowering_test.rs @@ -10107,6 +10107,117 @@ fn test_pointer_form_wgmma_partial_wait_pipeline_keeps_multiple_groups_in_flight Ok(()) } +#[test] +fn test_pointer_form_f16_wgmma_two_slot_wait_one_pipeline_lowers() -> 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 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_f16( + &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, 1); + } + } + 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)) + .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, + "expected one fused F16 two-slot partial-wait WGMMA pipeline" + ); + + let asm = &matching[0]; + let template = asm + .get_attr_inline_asm_template(&ctx) + .map(|value| String::from((*value).clone())) + .expect("F16 pipeline 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(), + GROUP_COUNT + ); + assert!(!template.contains("m64n64k16.f32.bf16.bf16")); + 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")); + assert!(template.contains("{$32, $33")); + assert!( + !template.contains(".reg .f32") + && !template.contains("ld.f32") + && !template.contains("st.f32"), + "F16 pipeline must keep accumulators value-threaded: {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); + + let load_count = body + .iter() + .filter(|operation| Operation::get_op::(**operation, &ctx).is_some()) + .count(); + let store_count = body + .iter() + .filter(|operation| Operation::get_op::(**operation, &ctx).is_some()) + .count(); + assert_eq!(load_count, RESULT_COUNT); + assert_eq!(store_count, RESULT_COUNT); + + Ok(()) +} + fn build_pointer_form_wgmma_counted_pipeline_case( ctx: &mut Context, slot_count: usize, @@ -11600,13 +11711,44 @@ fn test_linear_wgmma_full_drain_rejects_mixed_f16_and_tf32() -> Result<(), anyho } #[test] -fn test_f16_wgmma_partial_wait_remains_unsupported() -> Result<(), anyhow::Error> { +fn test_f16_wgmma_partial_wait_two_remains_unsupported() -> Result<(), anyhow::Error> { + const SLOT_COUNT: usize = 3; + const GROUP_COUNT: usize = 3; let mut ctx = make_test_ctx(); let (module_ptr, entry, accumulators, descriptors) = - build_wgmma_canonical_pointer_test_kernel(&mut ctx, 1, 2); + build_wgmma_canonical_pointer_test_kernel(&mut ctx, SLOT_COUNT, GROUP_COUNT * 2); nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); - append_pointer_wgmma_mma_f16( + for group in 0..GROUP_COUNT { + append_pointer_wgmma_mma_f16( + &mut ctx, + entry, + accumulators[group], + descriptors[group * 2], + descriptors[group * 2 + 1], + ); + nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + } + append_wgmma_wait_group_constant(&mut ctx, entry, 2); + append_wgmma_wait_group_constant(&mut ctx, entry, 0); + append_return(&mut ctx, entry); + + assert_wgmma_lowering_rejected( + &mut ctx, + module_ptr, + "F16 WGMMA partial-wait pipeline supports only wait_group<1> with two accumulator slots", + ); + Ok(()) +} + +#[test] +fn test_wgmma_partial_wait_pipeline_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, 2, 4); + + nvvm::WgmmaFenceSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_pointer_wgmma_mma( &mut ctx, entry, accumulators[0], @@ -11614,6 +11756,14 @@ fn test_f16_wgmma_partial_wait_remains_unsupported() -> Result<(), anyhow::Error descriptors[1], ); nvvm::WgmmaCommitGroupSyncAlignedOp::build(&mut ctx).insert_at_back(entry, &ctx); + append_pointer_wgmma_mma_f16( + &mut ctx, + entry, + accumulators[1], + 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); @@ -11621,7 +11771,7 @@ fn test_f16_wgmma_partial_wait_remains_unsupported() -> Result<(), anyhow::Error assert_wgmma_lowering_rejected( &mut ctx, module_ptr, - "WGMMA deferred accumulator lowering requires wait_group<0>", + "one WGMMA partial-wait pipeline cannot mix MMA variants or shapes", ); Ok(()) } diff --git a/cuda-oxide-book/advanced/matrix-multiply-accelerators.md b/cuda-oxide-book/advanced/matrix-multiply-accelerators.md index 2c6643377e..28e110d07a 100644 --- a/cuda-oxide-book/advanced/matrix-multiply-accelerators.md +++ b/cuda-oxide-book/advanced/matrix-multiply-accelerators.md @@ -124,8 +124,9 @@ hardware shape uses K=8. ### Current cuda-oxide WGMMA lowering The compiler supports `m64n64k16.f32.bf16.bf16` across three conservative -lowering shapes, `m64n64k16.f32.f16.f16` for canonical linear full-drain and -counted K-loop regions, `m64n64k8.f32.tf32.tf32` for canonical linear +lowering shapes, `m64n64k16.f32.f16.f16` for canonical linear full-drain, +counted K-loop, and two-slot `wait_group<1>` straight-line pipeline regions, +`m64n64k8.f32.tf32.tf32` for canonical linear full-drain regions only, and `m64n128k16.f32.bf16.bf16` for canonical linear full-drain regions. In every accepted case, the entire asynchronous accumulator lifetime stays inside one convergent inline-PTX statement so LLVM @@ -151,10 +152,11 @@ convergent PTX region as the WGMMA instruction. The accumulator is therefore loaded once before the K-loop and stored once after its final wait rather than round-tripped every iteration. -**Static partial-wait pipeline.** A BF16 m64n64 `wait_group` with `N > 0` -uses `N + 1` independent accumulator slots. Groups are committed separately and -slots are reused round-robin only after the partial wait has made the oldest -slot safe. For example, two groups can remain overlapped with: +**Static partial-wait pipeline.** BF16 m64n64 supports static `wait_group` +with `N > 0`, using `N + 1` independent accumulator slots. F16 deliberately +supports only the first topology: `wait_group<1>` with two slots. Groups are +committed separately and slots are reused round-robin only after the partial +wait has made the oldest slot safe. For example, two groups can remain overlapped with: ```text wgmma_fence @@ -171,8 +173,9 @@ accumulator. A final `wait_group<0>` is mandatory before any accumulator value escapes the fused region. Selection is intentionally fail-closed. Dynamic partial waits, unsupported -control flow, malformed accumulator schedules, F16 partial-wait shapes, TF32 -counted-loop or partial-wait shapes, m64n128 counted-loop or partial-wait +control flow, malformed accumulator schedules, F16 partial-wait depths other +than one, F16 counted pipelines, TF32 counted-loop or partial-wait shapes, +m64n128 counted-loop or partial-wait shapes, F16/TF32/m64n128 non-canonical accumulators, and the legacy K=16 TF32 compatibility entry point are rejected instead of exposing an in-flight accumulator to LLVM. diff --git a/cuda-oxide-book/appendix/supported-features.md b/cuda-oxide-book/appendix/supported-features.md index 9506a531d0..73f58bb00f 100644 --- a/cuda-oxide-book/appendix/supported-features.md +++ b/cuda-oxide-book/appendix/supported-features.md @@ -271,7 +271,7 @@ Anonymous promoted allocations remain unsupported. |:--------|:-------|:------------| | Warp-Level MMA (`wmma`) | **Full** | Register-only `mma.sync` shapes, `movmatrix`, and warp-cooperative `ldmatrix` loads. | | Sparse MMA | **Full** | Structured-sparsity `mma.sp` shapes alongside the dense ones, in the same `wmma` module. | -| Warpgroup MMA (`wgmma`) | **Partial** | Hopper `sm_90a`: fence/commit/wait pipeline, shared-memory descriptors, and `m64n64k16` MMA with `bf16`/`f16` inputs and `f32` accumulate. Gap: the lowering covers specific proven loop patterns (`bf16` works in straight-line code, counted K-loops, and partial-wait pipelines; `f16` in straight-line code only), and `tf32` calls are rejected pending [#1076](https://github.com/NVlabs/cuda-oxide/issues/1076). | +| Warpgroup MMA (`wgmma`) | **Partial** | Hopper `sm_90a`: fence/commit/wait pipeline, shared-memory descriptors, `m64n64k16` BF16/F16 and `m64n64k8` TF32 with `f32` accumulation, plus m64n128 BF16 full-drain lowering. BF16 m64n64 supports full-drain, counted K-loops, and static partial waits; F16 m64n64 supports full-drain, counted K-loops, and the two-slot `wait_group<1>` straight-line pipeline; TF32 and m64n128 BF16 currently support full-drain only. | | Tensor Core Gen 5 (`tcgen05`) | **Full** | Blackwell sm_100+: TMEM alloc/dealloc, MMA, `stmatrix`, CTA-pair (cg2) variants. | | Accumulator Fragment Algebra (`mma_frag`) | **Full** | Index algebra for the `m16n8k16` accumulator, so a lane can address its own slots of the `[f32; 4]` fragment. | | FP8 / FP6 / FP4 Formats | **Partial** | Conversions (`convert::cvt_*` for `e4m3`/`e5m2`) and the matrix path (FP8 `mma.sync` shapes, the `mxf8f6f4` shapes, tcgen05 descriptors) ship; the `mma_mxf8f6f4` example compiles them. Gap: no *arithmetic* on these formats. There's no add/mul/min/max the way `f16` and `bf16` have, so values are carried as packed bit patterns and converted before use. | diff --git a/cuda-oxide-book/compiler/mlir-dialects.md b/cuda-oxide-book/compiler/mlir-dialects.md index 02cdc46832..b71a7bb225 100644 --- a/cuda-oxide-book/compiler/mlir-dialects.md +++ b/cuda-oxide-book/compiler/mlir-dialects.md @@ -297,7 +297,7 @@ they become `call` instructions to `@llvm.nvvm.*` intrinsics. ### Architecture Coverage At catalog SHA-256 `00372dfe` (the stamp in every `ops/generated/` file -header), the dialect holds 575 operations across 42 modules, and they come +header), the dialect holds 576 operations across 42 modules, and they come from two different places. The split is the first thing to know about it, because it decides where -- and whether -- you would add one. If the header stamp no longer starts with `00372dfe`, the counts on this page predate the @@ -305,7 +305,7 @@ catalog you are reading. **Hand-written**, directly under `crates/dialect-nvvm/src/ops/`. These are the ops with bespoke verification or lowering that the intrinsic catalog does not -describe. There are seven modules and 26 operations: +describe. There are seven modules and 27 operations: | Module | Description | Ops | | :-------- | :---------------------------------------------------------- | --: | @@ -315,7 +315,7 @@ describe. There are seven modules and 26 operations: | `debug` | `assertfail`, `vprintf` | 2 | | `grid` | Cooperative `grid_sync` | 1 | | `memory` | Generic-to-shared address conversion with a byte offset | 1 | -| `wgmma` | Warpgroup MMA descriptors; bf16/f16 at m64n64k16, bf16 at m64n128k16, tf32 at m64n64k8 | 14 | +| `wgmma` | Warpgroup MMA descriptors; bf16/f16 at m64n64k16, bf16 at m64n128k16, tf32 at m64n64k8 | 15 | **Generated**, under `ops/generated/`, from `intrinsics/catalog.json` by `cuda-intrinsics-gen`. Every file there opens with `// @generated ... DO NOT