feat(gguf,dsv4): GGUF block-dequant + DSv4 architecture kernel surface

crates/metaltile-codegen/src/emit.rs

@@ -176,6 +176,7 @@ pub fn render_swift_wrappers(kernels: &[Kernel]) -> String {
     );
     for k in kernels {
         emit_swift_wrapper(&mut out, k);
+        emit_swift_wrapper_threadgroups(&mut out, k);
         if k.wants_indirect_variant {
             emit_swift_wrapper_indirect(&mut out, k);
         }
@@ -242,6 +243,66 @@ fn emit_swift_wrapper(out: &mut String, k: &Kernel) {
     writeln!(out, "    }}\n").ok();
 }
 
+/// `dispatchThreadgroups` variant of `emit_swift_wrapper`. IDENTICAL bindings,
+/// but `gridSize` is in THREADGROUPS and dispatch uses
+/// `dispatchThreadgroups(_:threadsPerThreadgroup:)`. REQUIRED for coop_tile /
+/// simdgroup_matrix kernels: `dispatchThreads` puts Metal in non-uniform-
+/// threadgroup mode, which silently produces WRONG results for cooperative-
+/// matrix ops even when the grid is a multiple of the threadgroup size. Plain
+/// simd_sum/elementwise kernels are fine on either path.
+fn emit_swift_wrapper_threadgroups(out: &mut String, k: &Kernel) {
+    use std::fmt::Write as _;
+    let fn_name = format!("{}_threadgroups", swift_safe_name(&k.name));
+    writeln!(
+        out,
+        "    /// dispatchThreadgroups variant of `{}` (gridSize in THREADGROUPS).",
+        k.name
+    )
+    .ok();
+    writeln!(out, "    /// Use for coop_tile / simdgroup-matrix kernels.").ok();
+    writeln!(out, "    public static func {fn_name}(").ok();
+    for p in &k.params {
+        let label = swift_safe_name(&p.name);
+        writeln!(out, "        {label}: MTLBuffer, {label}Offset: Int = 0,").ok();
+    }
+    for c in &k.constexprs {
+        let label = swift_safe_name(c.name.name());
+        let swift_ty = swift_scalar_type(dtype_suffix(c.dtype));
+        writeln!(out, "        {label}: {swift_ty},").ok();
+    }
+    writeln!(out, "        gridSize: MTLSize,").ok();
+    writeln!(out, "        threadgroupSize: MTLSize,").ok();
+    writeln!(out, "        on commandBuffer: MTLCommandBuffer").ok();
+    writeln!(out, "    ) {{").ok();
+    writeln!(out, "        let pso = PSOCache.shared.pipelineState(for: \"{}\")", k.name).ok();
+    writeln!(
+        out,
+        "        guard let enc = commandBuffer.makeComputeCommandEncoder() else {{ return }}"
+    )
+    .ok();
+    writeln!(out, "        enc.setComputePipelineState(pso)").ok();
+    let mut slot = 0usize;
+    for p in &k.params {
+        let label = swift_safe_name(&p.name);
+        writeln!(out, "        enc.setBuffer({label}, offset: {label}Offset, index: {slot})").ok();
+        slot += 1;
+    }
+    for c in &k.constexprs {
+        let label = swift_safe_name(c.name.name());
+        let len = swift_scalar_size(dtype_suffix(c.dtype));
+        writeln!(out, "        var {label}_v = {label}").ok();
+        writeln!(out, "        enc.setBytes(&{label}_v, length: {len}, index: {slot})").ok();
+        slot += 1;
+    }
+    writeln!(
+        out,
+        "        enc.dispatchThreadgroups(gridSize, threadsPerThreadgroup: threadgroupSize)"
+    )
+    .ok();
+    writeln!(out, "        enc.endEncoding()").ok();
+    writeln!(out, "    }}\n").ok();
+}
+
 /// Indirect-dispatch variant of `emit_swift_wrapper`. Same buffer +
 /// constexpr bindings, same PSO (the underlying kernel is unchanged),
 /// but the dispatch shape comes from an `MTLBuffer` carrying
@@ -490,6 +551,7 @@ pub fn dtype_suffix(dt: DType) -> &'static str {
         DType::U32 => "u32",
         DType::I8 => "i8",
         DType::U8 => "u8",
+        DType::U16 => "u16",
         DType::I64 => "i64",
         DType::U64 => "u64",
         DType::I4 => "i4",
@@ -590,4 +652,48 @@ mod tests {
         assert!(swift.contains("func mt_add_f32("));
         assert!(!swift.contains("_indirect("));
     }
+
+    /// The `dispatchThreadgroups` wrapper (REQUIRED for coop_tile /
+    /// simdgroup-matrix kernels) is a distinct emit path: a `_threadgroups`
+    /// function taking `gridSize` + `threadgroupSize` in THREADGROUPS and
+    /// calling `dispatchThreadgroups(_:threadsPerThreadgroup:)`. Pins its
+    /// signature + buffer/constexpr binding + the dispatch call so a codegen
+    /// change to it surfaces as a reviewable diff (per docs/testing.md, a new
+    /// emit path lands with a fixture exercising it).
+    #[test]
+    fn emits_threadgroups_wrapper_with_bindings_and_dispatch() {
+        let mut k = dummy_kernel("mt_moe_bgemm_mma");
+        // A second buffer + a constexpr so both binding loops are exercised.
+        k.params.push(Param {
+            name: "x".into(),
+            dtype: DType::F32,
+            shape: Shape::scalar(),
+            is_output: false,
+            kind: ParamKind::Tensor,
+        });
+        k.constexprs.push(metaltile_core::ir::ConstExprDecl {
+            name: metaltile_core::constexpr::ConstExpr::new("k_in"),
+            dtype: DType::U32,
+            value: None,
+        });
+        let swift = render_swift_wrappers(&[k]);
+
+        // The threadgroups variant exists alongside the direct wrapper.
+        assert!(swift.contains("func mt_moe_bgemm_mma("));
+        assert!(swift.contains("public static func mt_moe_bgemm_mma_threadgroups("));
+        // gridSize is in THREADGROUPS and dispatch uses dispatchThreadgroups.
+        assert!(swift.contains("gridSize: MTLSize,"));
+        assert!(swift.contains("threadgroupSize: MTLSize,"));
+        assert!(
+            swift
+                .contains("dispatchThreadgroups(gridSize, threadsPerThreadgroup: threadgroupSize)")
+        );
+        // Buffers bind by slot; the constexpr is set as inline bytes.
+        assert!(swift.contains("enc.setBuffer(out, offset: outOffset, index: 0)"));
+        assert!(swift.contains("enc.setBuffer(x, offset: xOffset, index: 1)"));
+        assert!(swift.contains("var k_in_v = k_in"));
+        assert!(swift.contains("enc.setBytes(&k_in_v, length: 4, index: 2)"));
+        // PSO lookup uses the base kernel name (no `_threadgroups` PSO exists).
+        assert!(swift.contains("pipelineState(for: \"mt_moe_bgemm_mma\")"));
+    }
 }

crates/metaltile-codegen/src/passes/algebraic_simplify.rs

@@ -59,9 +59,8 @@
 //!   Tools", 2nd ed., §8.4–8.5.  Classic treatment of algebraic identities and
 //!   reduction in strength.
 
-use std::collections::{BTreeMap, BTreeSet};
-
 use metaltile_core::ir::{BinOpKind, Block, BlockId, Kernel, Op, UnaryOpKind, ValueId};
+use rustc_hash::{FxHashMap, FxHashSet};
 
 use super::remap;
 use crate::error::{Error, Result};
@@ -93,11 +92,11 @@ impl super::Pass for AlgebraicSimplifyPass {
 // ---------------------------------------------------------------------------
 
 /// Resolve transitive replacement chains: {v2→v1, v1→v0} becomes {v2→v0, v1→v0}.
-fn resolve_transitive(map: &BTreeMap<ValueId, ValueId>) -> BTreeMap<ValueId, ValueId> {
-    let mut resolved = BTreeMap::new();
+fn resolve_transitive(map: &FxHashMap<ValueId, ValueId>) -> FxHashMap<ValueId, ValueId> {
+    let mut resolved = FxHashMap::default();
     for (&key, &val) in map.iter() {
         let mut terminal = val;
-        let mut visited = BTreeSet::new();
+        let mut visited = FxHashSet::default();
         visited.insert(key);
         while let Some(&next) = map.get(&terminal) {
             if !visited.insert(terminal) {
@@ -122,10 +121,11 @@ fn simplify_block_once(block: &mut Block) -> bool {
     let n = block.ops.len();
     let mut const_overwrites: Vec<(usize, i64)> = Vec::new();
     let mut op_replacements: Vec<(usize, Op)> = Vec::new();
-    let mut vid_replacements: BTreeMap<ValueId, ValueId> = BTreeMap::new();
+    let mut vid_replacements: FxHashMap<ValueId, ValueId> =
+        FxHashMap::with_capacity_and_hasher(block.ops.len(), Default::default());
 
     // Build a map for peephole lookups.
-    let vid_to_op_pos: BTreeMap<ValueId, usize> = block
+    let vid_to_op_pos: FxHashMap<ValueId, usize> = block
         .results
         .iter()
         .enumerate()
@@ -180,7 +180,7 @@ fn simplify_block_once(block: &mut Block) -> bool {
     // Remove dead ops whose results were redirected via ReplaceWithVid.
     // Without this, the same pattern re-matches on the next iteration,
     // producing the same replacement and causing an infinite fixpoint loop.
-    let dead_vids: BTreeSet<ValueId> = vid_replacements.keys().copied().collect();
+    let dead_vids: FxHashSet<ValueId> = vid_replacements.keys().copied().collect();
     if !dead_vids.is_empty() {
         let mut new_ops = Vec::new();
         let mut new_results = Vec::new();
@@ -217,7 +217,7 @@ fn try_simplify(
     op: &Op,
     pos: usize,
     block: &Block,
-    vid_to_pos: &BTreeMap<ValueId, usize>,
+    vid_to_pos: &FxHashMap<ValueId, usize>,
 ) -> Option<SimpResult> {
     match op {
         // ---- BinOp patterns ----
@@ -292,7 +292,7 @@ fn simplify_binop(
     rhs: ValueId,
     _pos: usize,
     block: &Block,
-    vid_to_pos: &BTreeMap<ValueId, usize>,
+    vid_to_pos: &FxHashMap<ValueId, usize>,
 ) -> Option<SimpResult> {
     let lv = find_const_in_block(block, lhs);
     let rv = find_const_in_block(block, rhs);
@@ -383,7 +383,7 @@ fn simplify_select(
     on_true: ValueId,
     on_false: ValueId,
     block: &Block,
-    vid_to_pos: &BTreeMap<ValueId, usize>,
+    vid_to_pos: &FxHashMap<ValueId, usize>,
 ) -> Option<SimpResult> {
     // Select(true, a, b) → a
     if let Some(1) = find_const_in_block(block, cond) {
@@ -427,7 +427,7 @@ fn find_const_in_block(block: &Block, vid: ValueId) -> Option<i64> {
 fn get_defining_op<'a>(
     vid: ValueId,
     block: &'a Block,
-    vid_to_pos: &BTreeMap<ValueId, usize>,
+    vid_to_pos: &FxHashMap<ValueId, usize>,
 ) -> Option<(usize, &'a Op)> {
     let &pos = vid_to_pos.get(&vid)?;
     Some((pos, &block.ops[pos]))
@@ -437,7 +437,7 @@ fn get_defining_op<'a>(
 fn get_neg_arg(
     vid: ValueId,
     block: &Block,
-    vid_to_pos: &BTreeMap<ValueId, usize>,
+    vid_to_pos: &FxHashMap<ValueId, usize>,
 ) -> Option<ValueId> {
     let (_pos, op) = get_defining_op(vid, block, vid_to_pos)?;
     if let Op::UnaryOp { op: UnaryOpKind::Neg, value } = op { Some(*value) } else { None }
@@ -447,7 +447,7 @@ fn get_neg_arg(
 fn get_not_arg(
     vid: ValueId,
     block: &Block,
-    vid_to_pos: &BTreeMap<ValueId, usize>,
+    vid_to_pos: &FxHashMap<ValueId, usize>,
 ) -> Option<ValueId> {
     let (_pos, op) = get_defining_op(vid, block, vid_to_pos)?;
     match op {
@@ -476,7 +476,7 @@ fn get_not_arg(
 }
 
 /// Remap all ValueId references in an op.
-fn remap_values_in_op(op: &mut Op, map: &BTreeMap<ValueId, ValueId>) {
+fn remap_values_in_op(op: &mut Op, map: &FxHashMap<ValueId, ValueId>) {
     op.for_each_value_id_mut(&mut |v| {
         if let Some(&nv) = map.get(v) {
             *v = nv;

crates/metaltile-codegen/src/passes/block_util.rs

@@ -6,9 +6,8 @@
 //! within a [`Block`].  Used by Unroll, LICM, DSE, IfConversion, ValueSink, and
 //! any pass that rewrites the op/results arrays of a block.
 
-use std::collections::BTreeSet;
-
 use metaltile_core::ir::{Block, Op, ValueId};
+use rustc_hash::FxHashSet;
 
 /// Remove ops at the given indices from a block.
 ///
@@ -18,7 +17,7 @@ pub fn remove_ops(block: &mut Block, indices: &[usize]) {
     if indices.is_empty() {
         return;
     }
-    let skip: BTreeSet<usize> = indices.iter().copied().collect();
+    let skip: FxHashSet<usize> = indices.iter().copied().collect();
     let old_ops = std::mem::take(&mut block.ops);
     let old_results = std::mem::take(&mut block.results);
     let mut new_ops = Vec::with_capacity(old_ops.len().saturating_sub(indices.len()));

crates/metaltile-codegen/src/passes/copy_prop.rs

@@ -31,12 +31,11 @@
 //!   ACM TOPLAS 13(2):181–210.  Sparse conditional constant propagation framework
 //!   that subsumes copy propagation.
 
-use std::collections::{BTreeMap, BTreeSet};
-
 use metaltile_core::{
     dtype::DType,
     ir::{Block, BlockId, Kernel, Op, ValueId},
 };
+use rustc_hash::{FxHashMap, FxHashSet};
 
 use super::remap;
 use crate::error::{Error, Result};
@@ -64,11 +63,11 @@ impl super::Pass for CopyPropPass {
 }
 
 /// Resolve transitive replacement chains: {v2→v1, v1→v0} becomes {v2→v0, v1→v0}.
-fn resolve_transitive(map: &BTreeMap<ValueId, ValueId>) -> BTreeMap<ValueId, ValueId> {
-    let mut resolved = BTreeMap::new();
+fn resolve_transitive(map: &FxHashMap<ValueId, ValueId>) -> FxHashMap<ValueId, ValueId> {
+    let mut resolved = FxHashMap::default();
     for (&key, &val) in map.iter() {
         let mut terminal = val;
-        let mut visited = BTreeSet::new();
+        let mut visited = FxHashSet::default();
         visited.insert(key);
         while let Some(&next) = map.get(&terminal) {
             if !visited.insert(terminal) {
@@ -91,7 +90,8 @@ fn copy_prop_block_fixpoint(block: &mut Block) {
 
 fn copy_prop_block_once(block: &mut Block) -> bool {
     let n = block.ops.len();
-    let mut vid_replacements: BTreeMap<ValueId, ValueId> = BTreeMap::new();
+    let mut vid_replacements: FxHashMap<ValueId, ValueId> =
+        FxHashMap::with_capacity_and_hasher(block.ops.len(), Default::default());
 
     for i in 0..n {
         let op = &block.ops[i];
@@ -117,7 +117,7 @@ fn copy_prop_block_once(block: &mut Block) -> bool {
     // Remove dead ops whose results were redirected via identity propagation.
     // Without this, the same identity pattern re-matches on the next iteration,
     // producing the same replacement and causing an infinite fixpoint loop.
-    let dead_vids: BTreeSet<ValueId> = vid_replacements.keys().copied().collect();
+    let dead_vids: FxHashSet<ValueId> = vid_replacements.keys().copied().collect();
     if !dead_vids.is_empty() {
         let mut new_ops = Vec::new();
         let mut new_results = Vec::new();

crates/metaltile-codegen/src/passes/licm.rs

@@ -38,7 +38,7 @@
 use std::collections::{BTreeMap, BTreeSet};
 
 use metaltile_core::ir::{Block, BlockId, Kernel, Op, ParamKind, ValueId};
-use rustc_hash::FxHashMap;
+use rustc_hash::{FxHashMap, FxHashSet};
 
 use super::remap;
 use crate::error::{Error, Result};
@@ -58,7 +58,7 @@ impl super::Pass for LicmPass {
             .collect();
 
         // Build a definition map: ValueId -> BlockId where it's defined.
-        let mut def_block: BTreeMap<ValueId, BlockId> = BTreeMap::new();
+        let mut def_block: FxHashMap<ValueId, BlockId> = FxHashMap::default();
         for vid in kernel.body.results.iter().flatten() {
             def_block.insert(*vid, kernel.body.id);
         }
@@ -101,7 +101,7 @@ impl super::Pass for LicmPass {
 fn licm_block(
     block: &mut Block,
     blocks: &mut FxHashMap<BlockId, Block>,
-    def_block: &BTreeMap<ValueId, BlockId>,
+    def_block: &FxHashMap<ValueId, BlockId>,
     read_only: &BTreeSet<String>,
 ) {
     let n = block.ops.len();
@@ -126,7 +126,7 @@ fn licm_block(
 
             // Build the initial invariant set: ValueIds defined before position `i`
             // in the parent block, plus any from ancestor blocks.
-            let mut invariant: BTreeSet<ValueId> = BTreeSet::new();
+            let mut invariant: FxHashSet<ValueId> = FxHashSet::default();
             for j in 0..i {
                 if let Some(Some(vid)) = block.results.get(j) {
                     invariant.insert(*vid);
@@ -264,7 +264,7 @@ fn licm_block(
 
 /// Remove ops at given indices from a block. Indices must be sorted ascending.
 fn remove_ops_from_block(block: &mut Block, indices: &[usize]) {
-    let skip: BTreeSet<usize> = indices.iter().copied().collect();
+    let skip: FxHashSet<usize> = indices.iter().copied().collect();
     let old_ops = std::mem::take(&mut block.ops);
     let old_results = std::mem::take(&mut block.results);
     let mut new_ops = Vec::new();

crates/metaltile-codegen/src/passes/remap.rs

@@ -30,11 +30,27 @@ use smallvec::SmallVec;
 // remap_value_ids — mutate ValueId references in an Op
 // ---------------------------------------------------------------------------
 
+/// `ValueId → ValueId` lookup abstraction so [`remap_value_ids`] accepts either
+/// a `BTreeMap` or an `FxHashMap`. Callers pick the map type for performance;
+/// the remap itself is a pure per-reference substitution, so iteration order is
+/// irrelevant and either map yields identical output.
+pub trait VidRemap {
+    fn remap(&self, v: ValueId) -> Option<ValueId>;
+}
+
+impl VidRemap for BTreeMap<ValueId, ValueId> {
+    fn remap(&self, v: ValueId) -> Option<ValueId> { self.get(&v).copied() }
+}
+
+impl<S: std::hash::BuildHasher> VidRemap for std::collections::HashMap<ValueId, ValueId, S> {
+    fn remap(&self, v: ValueId) -> Option<ValueId> { self.get(&v).copied() }
+}
+
 /// Remap all `ValueId` references in `op` according to `map`.
 /// References not present in `map` are left unchanged.
-pub fn remap_value_ids(op: &mut Op, map: &BTreeMap<ValueId, ValueId>) {
+pub fn remap_value_ids<M: VidRemap>(op: &mut Op, map: &M) {
     op.for_each_value_id_mut(&mut |v| {
-        if let Some(&nv) = map.get(v) {
+        if let Some(nv) = map.remap(*v) {
             *v = nv;
         }
     });