feat(wgmma): add two-slot F16 partial-wait pipeline lowering

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,25 +84,27 @@
 //!   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.
 //! - **Canonical counted K-loop:** one BF16 or F16 m64n64 MMA per iteration,
 //!   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::<N>()` 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::<N>()`
+//!   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
 ///

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<u32> {
     })
 }
 
-#[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",
+    );
+}

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<Operation>) -> Self {
+        Self { op }
+    }
+
+    /// Build the supported two-slot F16 partial-wait pipeline.
+    pub fn build(
+        ctx: &mut Context,
+        accumulators: Vec<Value>,
+        descriptors: Vec<Value>,
+    ) -> Ptr<Operation> {
+        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);
 }

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

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,