diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index eb3895a48..bad8ce1a0 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -55,7 +55,16 @@ applies to that file. weights, reduced to the one instantiation the `metal_gemm_time` probe launches as its reference arm. It is taken from [ml-explore/mlx](https://github.com/ml-explore/mlx) v0.32.1 (`mlx/backend/metal/kernels/quantized.h`, `steel/gemm/mma.h`, `steel/gemm/loader.h`); the file's header names what was taken out. It is compiled by the probe -only and never enters the engine's kernels. MLX is under the MIT licence, Copyright (c) 2023 Apple Inc.: +only and never enters the engine's kernels. + +`crates/metal/kernels/gemm_tensor.metal` (the prompt pass's GEMMs on Metal 4's tensor ops) follows the design of MLX's +NAX quantized matmul, v0.32.1 (`mlx/backend/metal/kernels/quantized_nax.h`, `steel/gemm/nax.h`): the weights +dequantized into threadgroup memory a K step, the 64-by-64 tile over a 2-by-2 grid of SIMD groups, the `matmul2d` +descriptor with the weight operand transposed. It is written anew over this crate's own staging and quant owners and +copies no MLX code; the rule that picks it, `MatrixUnits` in `crates/metal/src/library.rs`, is MLX's +`is_nax_available` (`mlx/backend/metal/device.cpp`). Both files' comments name the MLX lines they follow. + +MLX is under the MIT licence, Copyright (c) 2023 Apple Inc.: > Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated > documentation files (the "Software"), to deal in the Software without restriction, including without limitation the diff --git a/crates/metal/build.rs b/crates/metal/build.rs index b6aa25132..d9e658f6c 100644 --- a/crates/metal/build.rs +++ b/crates/metal/build.rs @@ -1,7 +1,7 @@ //! The build of the Metal backend: the kernels' one source, and the //! metallib when the Metal Toolchain is installed. //! -//! Two outputs land in `OUT_DIR`: +//! Three outputs land in `OUT_DIR`: //! //! - `bloomery.metal`: the kernels' source with the fault site's table //! written in front of it, from `bloomery-device`'s table — the one owner @@ -10,6 +10,9 @@ //! - `bloomery.metallib`: the same source compiled by `xcrun metal`, when //! the Metal Toolchain is installed. Its presence is the `metal_lib` //! cfg, and the runtime loads it instead of the source. +//! - `bloomery_tensor.metal`: the tensor-op GEMMs' source, the same table in +//! front of the files it reads, which the runtime compiles at Metal 4 on a +//! GPU that takes them (`library.rs`). use std::path::{Path, PathBuf}; use std::process::Command; @@ -262,6 +265,7 @@ fn main() { "probe", "norm", "quant", + "gemm_shared", "gemm", "matvec_q4k", "matvec_q6k", @@ -289,6 +293,24 @@ fn main() { let source_path = out_dir.join("bloomery.metal"); std::fs::write(&source_path, &combined).expect("the source is written"); + // The tensor-op GEMMs read the fault's raise, the quant owners and the + // shared GEMM ABI, the engine source's own files, in its order. + let tensor: Vec<(&str, String)> = kernels + .iter() + .filter(|(name, _)| ["fault", "quant", "gemm_shared"].contains(name)) + .cloned() + .chain(std::iter::once(( + "gemm_tensor", + std::fs::read_to_string(kernels_dir.join("gemm_tensor.metal")) + .expect("kernels/gemm_tensor.metal is read"), + ))) + .collect(); + std::fs::write( + out_dir.join("bloomery_tensor.metal"), + combined_source(&header, &tensor), + ) + .expect("the tensor source is written"); + match build_metallib(&source_path, &out_dir) { Some(_) => println!("cargo:rustc-cfg=metal_lib=\"metallib\""), None => println!("cargo:rustc-cfg=metal_lib=\"source\""), diff --git a/crates/metal/kernels/experts.metal b/crates/metal/kernels/experts.metal index d62745cb2..74c000a97 100644 --- a/crates/metal/kernels/experts.metal +++ b/crates/metal/kernels/experts.metal @@ -48,7 +48,8 @@ // in every row of the threadgroup, because every row reads every slot. // // This file comes after matvec_q6k.metal and quant.metal (the decode owners), -// norm.metal (rms_scale, rms_usable) and fault.metal (raise). +// norm.metal (rms_scale, rms_usable), gemm_shared.metal (pair_swiglu) and +// fault.metal (raise). #include @@ -97,11 +98,6 @@ inline device const char * expert_base(bool shared_slot, return shared_slot ? shexp : stack; } -// silu(gate) * up. No route weight: the fold multiplies by it. -inline float pair_swiglu(float gate, float up) { - return gate / (1.0f + exp(-gate)) * up; -} - // True when the routed slot's id is past the stack. Slot 8 is the shared // expert and leaves `expert` at 0; the caller reads the shared tensors. inline bool pair_expert(device const int * ids, uint slot, uint n_experts, diff --git a/crates/metal/kernels/gemm.metal b/crates/metal/kernels/gemm.metal index 447534ae6..ea316c8d7 100644 --- a/crates/metal/kernels/gemm.metal +++ b/crates/metal/kernels/gemm.metal @@ -63,8 +63,9 @@ // no word past the output's last row or column is written; the residual // epilogue reads, adds and writes back per element. // -// This file comes after quant.metal: it reads the quant owners and the -// fault site's table. +// This file comes after quant.metal and gemm_shared.metal: it reads the +// quant owners, the fault site's table, the tile's geometry, the argument +// blocks and the activations' checked cast (gemm_stage_x). #include @@ -72,12 +73,6 @@ using namespace metal; #pragma METAL fp contract(off) -// The tile's geometry: the kernel's own constants, the ones the host's -// launch module (kernels/gemm.rs) mirrors. -#define GEMM_ROWS 64 // weight rows a threadgroup owns -#define GEMM_ACTS 64 // activation rows a threadgroup owns -#define GEMM_THREADS 128 // four SIMD groups - // The 8x8 fragments one SIMD group owns along each side: two SIMD groups // share a side, in 8-row blocks sixteen rows apart. #define GEMM_FRAGS_A (GEMM_ACTS / 16) @@ -87,67 +82,6 @@ static_assert(GEMM_THREADS == 4 * 32, "four SIMD groups of 32 threads"); static_assert(GEMM_ROWS % 16 == 0 && GEMM_ACTS % 16 == 0, "the tile splits into 8x8 fragments, two SIMD groups to a side"); -// The K values one step covers: 32, a whole number of every owner's -// 16-value staging chunks and half the smallest group (A4's 64). -template -constexpr int gemm_bk() { - return 32; -} - -// The staged rows' pitch in elements: the K step plus 16 bytes, steel's -// padding. -template -constexpr int gemm_lda() { - return gemm_bk() + 16 / (int)sizeof(T); -} - -// The threadgroup bytes one tile stages: the f32 arm's 18432, the f16 -// arm's 10240. The allocation itself lives in the entry kernel (a -// threadgroup variable may be declared only there), and the tile body takes -// it as a pointer. -template -constexpr int gemm_shmem_bytes() { - return (GEMM_ROWS + GEMM_ACTS) * gemm_lda() * (int)sizeof(T); -} - -// The argument block of a plain GEMM, bound by value at slot 3. -struct GemmArgs { - uint cols; // the input columns K, a whole number of the type's blocks - uint rows; // the weight rows N, the output's columns - uint acts; // the activation rows R, the output's rows - uint x_stride; // the f32 elements between x's rows, at least cols - uint y_stride; // the f32 elements between y's rows, at least rows - uint layer; // the layer a raise is recorded in -}; - -static_assert(sizeof(GemmArgs) == 24, "GemmArgs is 24 bytes: the host packs six words"); - -// One segment of a segmented GEMM: the weight rows of one tensor, and where -// they sit in the weights and in the output. The mat-vec's segment table's -// own four fields (norm.metal's MatvecSeg), the GEMM's meaning: the output -// offset is a column offset in the R-row output, and the first group is a -// weight-tile threadgroup. -struct GemmSeg { - uint rows; // the segment's weight rows - uint wblocks; // the segment's first row's block offset in the weights - uint out; // the segment's first column in the output - uint first_group; // the segment's first weight-tile threadgroup -}; - -// The argument block of a segmented GEMM, bound by value at slot 3. -struct GemmSegArgs { - uint cols; - uint acts; - uint x_stride; - uint y_stride; - uint layer; - uint nseg; - GemmSeg segs[BLOOMERY_MATVEC_MAX_SEGMENTS]; -}; - -static_assert(sizeof(GemmSegArgs) == 24 + 16 * BLOOMERY_MATVEC_MAX_SEGMENTS, - "GemmSegArgs is the header's six words and the segments' tables"); - // The float fragment multiply: a `vec` is the 8x8 simdgroup // matrix's per-lane storage, so the fragment aliases the matrix through // `thread_elements()` (mma.h's BaseMMAFrag::mma, the reinterpret is MLX's). @@ -275,30 +209,12 @@ inline void gemm_tile(threadgroup char * shmem, } } - // B: whole 16-value runs of the activation rows, loaded as scalars - // (no window alignment is assumed) and checked as float4s. The f16 - // arm's cast is the activation quantizer's: a value that is not - // finite or past the half range (one test: not at most the half - // maximum in magnitude, which a NaN fails too) raises the - // quantizer's site and stages NaN, so the row's outputs are NaN — - // never a silent inf. + // B: whole 16-value runs of the activation rows, each through the + // checked cast (gemm_stage_x). auto stage_b = [&](int k0) { FOR_UNROLL(short b = 0; b < B_CHUNKS; ++b) { device const float * xr = b_src[b] + k0; - FOR_UNROLL(short g = 0; g < 4; ++g) { - const float4 v = float4(xr[4 * g], xr[4 * g + 1], xr[4 * g + 2], xr[4 * g + 3]); - if (sizeof(T) == 2) { - const bool4 bad = !(fabs(v) <= 65504.f); - if (any(bad)) { - raise(faults, layer, BLOOMERY_FAULT_SITE_QUANT_COLUMN); - b_dst[b][g] = vec(select(v, float4(NAN), bad)); - } else { - b_dst[b][g] = vec(v); - } - } else { - b_dst[b][g] = vec(v); - } - } + gemm_stage_x(xr, b_dst[b], faults, layer); } }; diff --git a/crates/metal/kernels/gemm_routed.metal b/crates/metal/kernels/gemm_routed.metal index b375e5506..9ddb4c2e9 100644 --- a/crates/metal/kernels/gemm_routed.metal +++ b/crates/metal/kernels/gemm_routed.metal @@ -35,8 +35,9 @@ // K loop's alone; the tile's height and its split over the SIMD groups do // not move it. // -// This file comes after gemm.metal (the helpers and the instantiated tile -// pieces), experts.metal (pair_swiglu) and route.metal (the table's shape). +// This file comes after gemm_shared.metal (the routed tile's geometry, its +// argument block, pair_swiglu and gemm_stage_x), gemm.metal (the helpers +// and the instantiated tile pieces) and route.metal (the table's shape). #include @@ -69,8 +70,6 @@ struct F32W { GEMM_ENTRY(f32w_store_f32, F32W, float, false, gemm_kernel) -#define ROUTED_ROWS 64 // weight rows of a tile -#define ROUTED_ACTS 64 // activation rows of a tile #define ROUTED_FRAGS_A (ROUTED_ACTS / 8) #define ROUTED_FRAGS_W 2 // a SIMD group's 8-row weight fragments: its 16 weight rows @@ -86,20 +85,6 @@ constexpr int routed_shmem_bytes() { return (ROUTED_ROWS + ROUTED_ACTS) * gemm_lda() * (int)sizeof(T); } -// Eight words, bound by value. -struct RoutedArgs { - uint cols; // the input columns K - uint x_stride; // f32 elements between the activation rows - uint y_stride; // f32 elements between the output rows - uint layer; - uint n_experts; // the stack's experts; the table's group n_experts is shared - uint stack_blocks; // blocks of one expert's rows in the stack - uint gather; // 1: an activation row is the slot's pass row; 0: the place's - uint rows; // outputs per expert: the pair's rows of gate, the down's rows -}; - -static_assert(sizeof(RoutedArgs) == 32, "RoutedArgs is eight words"); - // One tile's K loop and store over its first NA 8-row A fragments: the rows // under 8 NA, which hold the run. Every SIMD group multiplies the same NA // fragments against its own 16 weight rows, so every fragment the loop holds @@ -160,20 +145,7 @@ inline void routed_tile(threadgroup T * sa, if (b_stage) { device const float * xr = b_src + k0; threadgroup vec * bd = (threadgroup vec *)((threadgroup T *)b_dst + at); - FOR_UNROLL(short g = 0; g < 4; ++g) { - const float4 v = float4(xr[4 * g], xr[4 * g + 1], xr[4 * g + 2], xr[4 * g + 3]); - if (sizeof(T) == 2) { - const bool4 over = !(fabs(v) <= 65504.f); - if (any(over)) { - raise(faults, args.layer, BLOOMERY_FAULT_SITE_QUANT_COLUMN); - bd[g] = vec(select(v, float4(NAN), over)); - } else { - bd[g] = vec(v); - } - } else { - bd[g] = vec(v); - } - } + gemm_stage_x(xr, bd, faults, args.layer); } const int pos = k0 + a_koff; const int ib = pos / Q::values_per_block(); diff --git a/crates/metal/kernels/gemm_shared.metal b/crates/metal/kernels/gemm_shared.metal new file mode 100644 index 000000000..a53609822 --- /dev/null +++ b/crates/metal/kernels/gemm_shared.metal @@ -0,0 +1,130 @@ +// What the two GEMM libraries share: the engine library's simdgroup GEMMs +// (gemm.metal, gemm_routed.metal) and the tensor-op GEMMs of the library +// compiled at Metal 4 (gemm_tensor.metal). The tile's geometry and the +// argument blocks the host packs are one program's ABI whichever library +// runs it (kernels/gemm.rs and kernels/routed.rs launch both), and the +// expert pair's SwiGLU is the decode pair's (experts.metal) and both +// grouped GEMMs' epilogue. The activations' checked cast is every GEMM's. + +#include + +using namespace metal; + +#pragma METAL fp contract(off) + +// The tile's geometry: the kernel's own constants, the ones the host's +// launch module (kernels/gemm.rs) mirrors. +#define GEMM_ROWS 64 // weight rows a threadgroup owns +#define GEMM_ACTS 64 // activation rows a threadgroup owns +#define GEMM_THREADS 128 // four SIMD groups + +// The K values one step covers: 32, a whole number of every owner's +// 16-value staging chunks and half the smallest group (A4's 64). +template +constexpr int gemm_bk() { + return 32; +} + +// The staged rows' pitch in elements: the K step plus 16 bytes, steel's +// padding. +template +constexpr int gemm_lda() { + return gemm_bk() + 16 / (int)sizeof(T); +} + +// The threadgroup bytes one tile stages: the f32 arm's 18432, the f16 +// arm's 10240. The allocation itself lives in the entry kernel (a +// threadgroup variable may be declared only there), and the tile body takes +// it as a pointer. +template +constexpr int gemm_shmem_bytes() { + return (GEMM_ROWS + GEMM_ACTS) * gemm_lda() * (int)sizeof(T); +} + +// The argument block of a plain GEMM, bound by value at slot 3. +struct GemmArgs { + uint cols; // the input columns K, a whole number of the type's blocks + uint rows; // the weight rows N, the output's columns + uint acts; // the activation rows R, the output's rows + uint x_stride; // the f32 elements between x's rows, at least cols + uint y_stride; // the f32 elements between y's rows, at least rows + uint layer; // the layer a raise is recorded in +}; + +static_assert(sizeof(GemmArgs) == 24, "GemmArgs is 24 bytes: the host packs six words"); + +// One segment of a segmented GEMM: the weight rows of one tensor, and where +// they sit in the weights and in the output. The mat-vec's segment table's +// own four fields (norm.metal's MatvecSeg), the GEMM's meaning: the output +// offset is a column offset in the R-row output, and the first group is a +// weight-tile threadgroup. +struct GemmSeg { + uint rows; // the segment's weight rows + uint wblocks; // the segment's first row's block offset in the weights + uint out; // the segment's first column in the output + uint first_group; // the segment's first weight-tile threadgroup +}; + +// The argument block of a segmented GEMM, bound by value at slot 3. +struct GemmSegArgs { + uint cols; + uint acts; + uint x_stride; + uint y_stride; + uint layer; + uint nseg; + GemmSeg segs[BLOOMERY_MATVEC_MAX_SEGMENTS]; +}; + +static_assert(sizeof(GemmSegArgs) == 24 + 16 * BLOOMERY_MATVEC_MAX_SEGMENTS, + "GemmSegArgs is the header's six words and the segments' tables"); + +// The routed tile's geometry (gemm_routed.metal, gemm_tensor.metal): the +// host's launch module (kernels/routed.rs) mirrors it. +#define ROUTED_ROWS 64 // weight rows of a tile +#define ROUTED_ACTS 64 // activation rows of a tile + +// Eight words, bound by value. +struct RoutedArgs { + uint cols; // the input columns K + uint x_stride; // f32 elements between the activation rows + uint y_stride; // f32 elements between the output rows + uint layer; + uint n_experts; // the stack's experts; the table's group n_experts is shared + uint stack_blocks; // blocks of one expert's rows in the stack + uint gather; // 1: an activation row is the slot's pass row; 0: the place's + uint rows; // outputs per expert: the pair's rows of gate, the down's rows +}; + +static_assert(sizeof(RoutedArgs) == 32, "RoutedArgs is eight words"); + +// silu(gate) * up. No route weight: the fold multiplies by it. +inline float pair_swiglu(float gate, float up) { + return gate / (1.0f + exp(-gate)) * up; +} + +// One 16-value run of an activation row, staged for a GEMM's tile: loaded +// as scalars (no window alignment is assumed) and checked as float4s. The +// f16 arm's cast is the activation quantizer's: a value that is not finite +// or past the half range (one test: not at most the half maximum in +// magnitude, which a NaN fails too) raises the quantizer's site and stages +// NaN, so the row's outputs are NaN — never a silent inf. Every GEMM of both +// libraries stages its activations through this one cast. +template +inline void gemm_stage_x(device const float * xr, threadgroup vec * dst, + device atomic_uint * faults, uint layer) { + FOR_UNROLL(short g = 0; g < 4; ++g) { + const float4 v = float4(xr[4 * g], xr[4 * g + 1], xr[4 * g + 2], xr[4 * g + 3]); + if (sizeof(T) == 2) { + const bool4 bad = !(fabs(v) <= 65504.f); + if (any(bad)) { + raise(faults, layer, BLOOMERY_FAULT_SITE_QUANT_COLUMN); + dst[g] = vec(select(v, float4(NAN), bad)); + } else { + dst[g] = vec(v); + } + } else { + dst[g] = vec(v); + } + } +} diff --git a/crates/metal/kernels/gemm_tensor.metal b/crates/metal/kernels/gemm_tensor.metal new file mode 100644 index 000000000..ff7ef951f --- /dev/null +++ b/crates/metal/kernels/gemm_tensor.metal @@ -0,0 +1,517 @@ +// The prompt pass's GEMMs on Metal 4's tensor ops (MetalPerformancePrimitives' +// `matmul2d`), the path a GPU whose cores carry matrix units takes. The +// entry points, under the host names of the simdgroup kernels they stand in +// for (one library or the other is the GEMMs' library; `Library` picks): +// +// gemm_{q4k,q6k,a4}_{store,resid}[_seg]_{f16,f32} gemm.metal's +// gemm_routed_{q4k,q6k,a4}_{pair,plain}_{f16,f32} gemm_routed.metal's +// +// Same arguments, buffers, grid and meanings as those (their headers hold +// them): the epilogues, the segment table, the route table's tiles, the f16 +// operand's checked cast and its fault, the expert-id fault. What differs is +// the multiply and its sum order. +// +// The design is MLX's NAX quantized matmul (MIT licence; see +// THIRD_PARTY_NOTICES.md), the path MLX takes where `is_nax_available`: +// mlx/backend/metal/kernels/quantized_nax.h:935-1078 qmm_t_nax_tgp_impl: +// the weights dequantized into threadgroup memory a K step, the 64x64 +// tile split over a 2x2 grid of SIMD groups (WM = WN = 2), each SIMD +// group's 32x32 block of the output its own op's +// mlx/backend/metal/kernels/steel/gemm/nax.h:387-411 the descriptor: +// matmul2d_descriptor(m, n, k, transpose_left = false, +// transpose_right = true, ..., mode::multiply_accumulate) at +// execution_simdgroup scope, the weight operand [n][k] (transposed) +// What differs from MLX: +// - the activations are staged too, in the simdgroup kernels' layout (the +// same K step of 32 and pitch): an f32 activation is cast to half with +// the range check that raises the activation quantizer's site, which +// needs the value in hand; MLX loads its already-half activations from +// device memory straight into registers. MLX's K step is 64 with only +// the weights staged (9 KB); ours is 32 with both (10 KB for f16), the +// same threadgroup bytes, so as many tiles sit on a core; +// - the op's operands are tensors over the staged rows and its +// destination a cooperative tensor read back through +// `get_multidimensional_index`, the layout the API defines; MLX loads +// its fragments into the cooperative tensors' storage in the layout its +// NAX fragment assumes (`BaseNAXFrag::get_coord`), which only the +// hardware it targets defines. A dense SIMD group's op is the whole +// 32x32x32 step, not MLX's 16x32x16 pieces; a routed group's ops are +// 16 places high, MLX's fragment height, so a short run wastes less; +// - relaxed precision is off (MLX's descriptor sets it): the clauses' +// band is derived for the multiply at its defined precision. +// Where llama.cpp's tensor `kernel_mul_mm` (ggml-metal/kernels/mul_mm.metal, +// under GGML_METAL_HAS_TENSOR) differs: one op over the whole threadgroup +// (execution_simdgroups<4>, a 64-weight by 128-activation tile, K dynamic), +// the f32 activations read from device memory as the op's left operand with +// no cast, relaxed precision on, and the destination stored through +// `cooperative_tensor::store`. +// +// Geometry: GEMM_THREADS (128) threads in four SIMD groups; a tile is +// GEMM_ROWS (64) weight rows by GEMM_ACTS (64) activation rows, a K step +// gemm_bk (32) values, both staged [row][k] at gemm_lda's pitch (the +// weights at the top, the activations under them). SIMD group g owns the +// activation rows 32 (g >> 1) on by the weight rows 32 (g & 1) on: one +// matmul2d a K step, accumulating into its cooperative tensor. A thread +// stages one 16-value chunk of each side a K step (one `dequant16`, one +// checked float4 run); rows past the tile's last live row stage that row's +// values (a defined access) and are not stored. The store is per element, +// under each element's own row and column bound. A SIMD group whose +// activation rows are all past the tile's live rows runs no op and stores +// nothing. +// +// The routed tile (gemm_routed.metal's): ROUTED_ROWS (64) weight rows by +// ROUTED_ACTS (64) slot places, the table's tile a threadgroup. SIMD group g +// owns the places 32 (g >> 1) on, as two blocks of TENSOR_RM (16) places +// (MLX's fragment height), each its own op, run only when the run reaches +// the block: a run's places are multiplied in 16-place steps, so a short +// run (the routed tiles' common case) wastes at most 15 places, and only +// the reached blocks' activation rows are staged. The pair's weight rows +// are the gate of 32 outputs, then their up; group g multiplies its places +// against 16 outputs' gate and the same 16 outputs' up, two ops of one +// type, so the two destinations share one layout and element i of each is +// one output's gate and up. The plain tile's group multiplies its places by +// 32 down rows. +// +// The sum order of an output is the op's: the clauses hold these kernels to +// a band derived for any order of its K products, not to the simdgroup +// kernels' bits. +// +// This file comes after quant.metal (the quant owners), fault.metal (raise) +// and gemm_shared.metal (the geometry, the argument blocks, pair_swiglu, +// the activations' checked cast gemm_stage_x). + +#include +#include +#include + +using namespace metal; +using namespace mpp::tensor_ops; + +#pragma METAL fp contract(off) + +// The tensor op's precision: relaxed precision lets the implementation trade +// accuracy for speed by its own rule, which no band is derived for. +#define TENSOR_RELAXED false + +// A SIMD group's block of the dense tile: half of each side (MLX's WM = WN = 2). +#define TENSOR_SM (GEMM_ACTS / 2) +#define TENSOR_SN (GEMM_ROWS / 2) +// A pair group's outputs: half of the tile's 32 outputs. +#define TENSOR_PAIR_N (ROUTED_ROWS / 4) +// A routed op's places: half of a SIMD group's 32. +#define TENSOR_RM (TENSOR_SM / 2) + +static_assert(GEMM_THREADS == 4 * 32, "four SIMD groups of 32 threads"); +static_assert(GEMM_ROWS * 2 == GEMM_THREADS && GEMM_ACTS * 2 == GEMM_THREADS, + "a thread stages one 16-value chunk of a weight row and of an activation row"); +static_assert(ROUTED_ROWS == GEMM_ROWS && ROUTED_ACTS == GEMM_ACTS, + "the routed tile is the dense tile's shape"); +static_assert(TENSOR_RM * 2 == TENSOR_SM, "a routed SIMD group's places are two ops' blocks"); +static_assert(gemm_bk() == 32 && gemm_bk() == 32, + "a K step is two 16-value chunks, the op's static K"); + +// The op of one SIMD group: an m-by-n block of the output over a K step, +// the left operand (the activations) [m][k], the right (the weights) [n][k]. +template +using tensor_mm = matmul2d; + +// `rows` staged rows of one K step at `at`, as the op's operand: K the +// contiguous extent, the rows at the staged pitch. +template +inline auto tensor_rows(threadgroup T * at) { + return tensor(at, extents(), array({1, gemm_lda()})); +} + +// One 16-value chunk of a weight row: one `dequant16` of the quant owner, +// staged as four contiguous vectors. +template +inline void tensor_stage_w(device const typename Q::Block * row, int pos, + threadgroup vec * dst) { + const int ib = pos / Q::values_per_block(); + const short il = (short)((pos / 16) % (Q::values_per_block() / 16)); + matrix reg; + Q::template dequant16(row[ib], il, reg); + FOR_UNROLL(short i = 0; i < 4; ++i) { + dst[i] = vec(reg[i][0], reg[i][1], reg[i][2], reg[i][3]); + } +} + +// The cooperative tensor's elements a lane holds, zeroed. +template +inline void tensor_zero(thread C & ct) { + FOR_UNROLL(ushort i = 0; i < ct.get_capacity(); ++i) { + if (ct.is_valid_element(i)) { + ct[i] = 0.f; + } + } +} + +// One threadgroup's tile of the dense GEMM: the K loop and the store, +// shared by the plain and the segmented entries (gemm.metal's gemm_tile's +// place: `r0` the first weight row, `r1` the first activation row). +template +inline void tensor_tile(threadgroup char * shmem, + device const typename Q::Block * W, + device const float * x, + device float * y, + int rows, + int cols, + int acts, + uint x_stride, + uint y_stride, + uint layer, + device atomic_uint * faults, + int r0, + int r1, + ushort tiitg, + ushort sgitg) { + constexpr int BK = gemm_bk(); + constexpr int LDA = gemm_lda(); + + threadgroup T * sa = (threadgroup T *)shmem; + threadgroup T * sb = sa + GEMM_ROWS * LDA; + + // The tile's live rows. The staging rows past them clamp to the last + // live row — a defined access, never stored. + const short nr0 = (short)min(GEMM_ROWS, rows - r0); + const short nr1 = (short)min(GEMM_ACTS, acts - r1); + const uint nb = (uint)cols / (uint)Q::values_per_block(); + + // The thread's chunk of each side: row tiitg / 2, the K step's half + // tiitg % 2. + const short row = (short)(tiitg / 2); + const short koff = (short)((tiitg % 2) * 16); + device const typename Q::Block * a_src = W + (ulong)(r0 + min(row, (short)(nr0 - 1))) * nb; + threadgroup vec * a_dst = (threadgroup vec *)(sa + row * LDA + koff); + device const float * b_src = x + (ulong)(r1 + min(row, (short)(nr1 - 1))) * x_stride + koff; + threadgroup vec * b_dst = (threadgroup vec *)(sb + row * LDA + koff); + + // The SIMD group's block: activation rows tm on, weight rows tn on. + const short tm = TENSOR_SM * (short)(sgitg >> 1); + const short tn = TENSOR_SN * (short)(sgitg & 1); + const bool live = tm < nr1; + + tensor_mm mm; + auto ta = tensor_rows(sb + tm * LDA); + auto tw = tensor_rows(sa + tn * LDA); + auto ct = mm.template get_destination_cooperative_tensor(); + tensor_zero(ct); + + for (int k0 = 0; k0 < cols; k0 += BK) { + // The previous step's operands are read: the tiles may be written + // over. + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_stage_x(b_src + k0, b_dst, faults, layer); + tensor_stage_w(a_src, k0 + koff, a_dst); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (live) { + mm.run(ta, tw, ct); + } + } + + // The store, per element under its row and column bound: no word past + // the output's last row or column is written. The destination's index + // is (weight row, activation row) in the block. + if (!live) { + return; + } + FOR_UNROLL(ushort i = 0; i < ct.get_capacity(); ++i) { + if (!ct.is_valid_element(i)) { + continue; + } + const auto at = ct.get_multidimensional_index(i); + const short n = tn + (short)at[0]; + const short m = tm + (short)at[1]; + if (n < nr0 && m < nr1) { + device float * p = y + (ulong)(r1 + m) * y_stride + (uint)(r0 + n); + *p = RESID ? *p + ct[i] : ct[i]; + } + } +} + +// The plain GEMM, storing or adding into the output. +template +kernel void tensor_gemm_kernel(device const typename Q::Block * W [[buffer(0)]], + device const float * x [[buffer(1)]], + device float * y [[buffer(2)]], + constant GemmArgs & args [[buffer(3)]], + device atomic_uint * faults [[buffer(4)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiitg [[thread_index_in_threadgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + threadgroup char shmem[gemm_shmem_bytes()]; + tensor_tile(shmem, W, x, y, (int)args.rows, (int)args.cols, (int)args.acts, + args.x_stride, args.y_stride, args.layer, faults, + (int)tgpig.y * GEMM_ROWS, (int)tgpig.x * GEMM_ACTS, tiitg, sgitg); +} + +// The weight rows of up to four tensors of one type and one width +// (gemm.metal's gemm_seg_kernel's table and meaning). +template +kernel void tensor_gemm_seg_kernel(device const typename Q::Block * W [[buffer(0)]], + device const float * x [[buffer(1)]], + device float * y [[buffer(2)]], + constant GemmSegArgs & args [[buffer(3)]], + device atomic_uint * faults [[buffer(4)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiitg [[thread_index_in_threadgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + threadgroup char shmem[gemm_shmem_bytes()]; + uint s = 0; + for (uint i = args.nseg; i > 0; --i) { + if (tgpig.y >= args.segs[i - 1u].first_group) { + s = i - 1u; + break; + } + } + const uint local = tgpig.y - args.segs[s].first_group; + tensor_tile(shmem, W + (ulong)args.segs[s].wblocks, x, y + args.segs[s].out, + (int)args.segs[s].rows, (int)args.cols, (int)args.acts, + args.x_stride, args.y_stride, args.layer, faults, + (int)local * GEMM_ROWS, (int)tgpig.x * GEMM_ACTS, tiitg, sgitg); +} + +// The pair's store of one block of places from `m0` on: element i of the +// gate's and the up's destinations is one output's, SwiGLU'd and weighted +// by its slot's route weight, under the run's length. `y_tile` is the +// tile's first place's row at the tile's first output. +template +inline void tensor_pair_store(thread C & cg, thread C & cu, short m0, short tn, short count, + const device uint * places, device const float * route_w, + device float * y_tile, uint y_stride) { + FOR_UNROLL(ushort i = 0; i < cg.get_capacity(); ++i) { + if (!cg.is_valid_element(i)) { + continue; + } + const auto at = cg.get_multidimensional_index(i); + const short m = m0 + (short)at[1]; + if (m < count) { + y_tile[(ulong)m * y_stride + (uint)(tn + (short)at[0])] = + pair_swiglu(cg[i], cu[i]) * route_w[places[m]]; + } + } +} + +// The plain tile's store of one block of places from `m0` on, under the +// run's length. +template +inline void tensor_plain_store(thread C & ct, short m0, short tn, short count, + device float * y_tile, uint y_stride) { + FOR_UNROLL(ushort i = 0; i < ct.get_capacity(); ++i) { + if (!ct.is_valid_element(i)) { + continue; + } + const auto at = ct.get_multidimensional_index(i); + const short m = m0 + (short)at[1]; + if (m < count) { + y_tile[(ulong)m * y_stride + (uint)(tn + (short)at[0])] = ct[i]; + } + } +} + +// The grouped GEMM of the routed experts over the route table's tiles +// (gemm_routed.metal's gemm_routed_kernel's tile, table and meaning). +template +kernel void tensor_routed_kernel(device const typename Q::Block * W [[buffer(0)]], + device const typename Q::Block * W_up [[buffer(1)]], + device const typename Q::Block * W_sh [[buffer(2)]], + device const typename Q::Block * W_sh_up [[buffer(3)]], + device const float * x [[buffer(4)]], + device float * y [[buffer(5)]], + device const uint4 * tiles [[buffer(6)]], + device const uint * sorted [[buffer(7)]], + device const float * route_w [[buffer(8)]], + constant RoutedArgs & args [[buffer(9)]], + device atomic_uint * faults [[buffer(10)]], + uint2 tgpig [[threadgroup_position_in_grid]], + ushort tiitg [[thread_index_in_threadgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + threadgroup char shmem[gemm_shmem_bytes()]; + constexpr int BK = gemm_bk(); + constexpr int LDA = gemm_lda(); + + if (tgpig.x >= tiles[0].x) { + return; + } + const uint4 tile = tiles[1 + tgpig.x]; + const uint expert = tile.x; + const uint first = tile.y; + const short count = (short)tile.z; + // A tile the table cannot write (an id past the shared group, a run + // empty or longer than the tile) raises the expert-id site and stores + // nothing. + if (expert > args.n_experts || count < 1 || count > ROUTED_ACTS) { + if (tiitg == 0) { + raise(faults, args.layer, BLOOMERY_FAULT_SITE_EXPERT_ID); + } + return; + } + const bool shared = expert == args.n_experts; + device const typename Q::Block * wg = shared ? W_sh : W + (ulong)expert * args.stack_blocks; + device const typename Q::Block * wu = PAIR ? (shared ? W_sh_up : W_up + (ulong)expert * args.stack_blocks) : wg; + // The tile's first output: the pair's tile is 32 outputs (64 weight rows + // of two tensors), the plain tile 64 rows of one. + const uint r0 = tgpig.y * (PAIR ? ROUTED_ROWS / 2 : ROUTED_ROWS); + const uint nb = args.cols / (uint)Q::values_per_block(); + + threadgroup T * sa = (threadgroup T *)shmem; + threadgroup T * sb = sa + ROUTED_ROWS * LDA; + + // A thread's chunk of each side: row tiitg / 2, the K step's half + // tiitg % 2. The pair's staged rows: 0 to 31 the gate of outputs r0 on, + // 32 to 63 their up. + const short row = (short)(tiitg / 2); + const short koff = (short)((tiitg % 2) * 16); + device const typename Q::Block * a_src = + PAIR ? (row < ROUTED_ROWS / 2 ? wg : wu) + (ulong)(r0 + (uint)(row % (ROUTED_ROWS / 2))) * nb + : wg + (ulong)(r0 + (uint)row) * nb; + threadgroup vec * a_dst = (threadgroup vec *)(sa + row * LDA + koff); + + // The activation rows the tile multiplies: the run's 16-row blocks. A + // row past the run's length takes the run's last (a defined access). + const short live_rows = TENSOR_RM * ((count + TENSOR_RM - 1) / TENSOR_RM); + const bool b_stage = row < live_rows; + const uint b_place = first + (uint)min(row, (short)(count - 1)); + const ulong src_row = + args.gather != 0u ? (ulong)(sorted[b_place] / BLOOMERY_EXPERT_SLOTS) : (ulong)b_place; + device const float * b_src = x + src_row * args.x_stride + koff; + threadgroup vec * b_dst = (threadgroup vec *)(sb + row * LDA + koff); + + // The SIMD group's two blocks of places, tm and tm + TENSOR_RM on, each + // live when the run reaches it. + const short tm = TENSOR_SM * (short)(sgitg >> 1); + const bool live0 = tm < count; + const bool live1 = tm + TENSOR_RM < count; + auto ta0 = tensor_rows(sb + tm * LDA); + auto ta1 = tensor_rows(sb + (tm + TENSOR_RM) * LDA); + const device uint * places = sorted + first; + device float * y_tile = y + (ulong)first * args.y_stride + r0; + + if (PAIR) { + // The group's 16 outputs' gate rows and up rows: one op type, so the + // two destinations share one layout. + const short tn = TENSOR_PAIR_N * (short)(sgitg & 1); + tensor_mm mm; + auto tg = tensor_rows(sa + tn * LDA); + auto tu = tensor_rows(sa + (ROUTED_ROWS / 2 + tn) * LDA); + auto cg0 = mm.template get_destination_cooperative_tensor(); + auto cu0 = mm.template get_destination_cooperative_tensor(); + auto cg1 = mm.template get_destination_cooperative_tensor(); + auto cu1 = mm.template get_destination_cooperative_tensor(); + tensor_zero(cg0); + tensor_zero(cu0); + tensor_zero(cg1); + tensor_zero(cu1); + for (int k0 = 0; k0 < (int)args.cols; k0 += BK) { + threadgroup_barrier(mem_flags::mem_threadgroup); + if (b_stage) { + gemm_stage_x(b_src + k0, b_dst, faults, args.layer); + } + tensor_stage_w(a_src, k0 + koff, a_dst); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (live0) { + mm.run(ta0, tg, cg0); + mm.run(ta0, tu, cu0); + } + if (live1) { + mm.run(ta1, tg, cg1); + mm.run(ta1, tu, cu1); + } + } + if (live0) { + tensor_pair_store(cg0, cu0, tm, tn, count, places, route_w, y_tile, args.y_stride); + } + if (live1) { + tensor_pair_store(cg1, cu1, tm + TENSOR_RM, tn, count, places, route_w, y_tile, + args.y_stride); + } + } else { + const short tn = TENSOR_SN * (short)(sgitg & 1); + tensor_mm mm; + auto tw = tensor_rows(sa + tn * LDA); + auto c0 = mm.template get_destination_cooperative_tensor(); + auto c1 = mm.template get_destination_cooperative_tensor(); + tensor_zero(c0); + tensor_zero(c1); + for (int k0 = 0; k0 < (int)args.cols; k0 += BK) { + threadgroup_barrier(mem_flags::mem_threadgroup); + if (b_stage) { + gemm_stage_x(b_src + k0, b_dst, faults, args.layer); + } + tensor_stage_w(a_src, k0 + koff, a_dst); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (live0) { + mm.run(ta0, tw, c0); + } + if (live1) { + mm.run(ta1, tw, c1); + } + } + if (live0) { + tensor_plain_store(c0, tm, tn, count, y_tile, args.y_stride); + } + if (live1) { + tensor_plain_store(c1, tm + TENSOR_RM, tn, count, y_tile, args.y_stride); + } + } +} + +// The entries, under the simdgroup kernels' host names. + +#define TENSOR_GEMM_ENTRY(SUFFIX, Q, T, RESID, SEG, ARGS) \ + template [[host_name("gemm_" #SUFFIX)]] kernel void SEG( \ + device const Q::Block * W [[buffer(0)]], \ + device const float * x [[buffer(1)]], \ + device float * y [[buffer(2)]], \ + constant ARGS & args [[buffer(3)]], \ + device atomic_uint * faults [[buffer(4)]], \ + uint3 tgpig [[threadgroup_position_in_grid]], \ + ushort tiitg [[thread_index_in_threadgroup]], \ + ushort sgitg [[simdgroup_index_in_threadgroup]]); + +#define TENSOR_GEMM_FAMILY(STEM, Q) \ + TENSOR_GEMM_ENTRY(STEM##_store_f16, Q, half, false, tensor_gemm_kernel, GemmArgs) \ + TENSOR_GEMM_ENTRY(STEM##_store_f32, Q, float, false, tensor_gemm_kernel, GemmArgs) \ + TENSOR_GEMM_ENTRY(STEM##_resid_f16, Q, half, true, tensor_gemm_kernel, GemmArgs) \ + TENSOR_GEMM_ENTRY(STEM##_resid_f32, Q, float, true, tensor_gemm_kernel, GemmArgs) \ + TENSOR_GEMM_ENTRY(STEM##_seg_store_f16, Q, half, false, tensor_gemm_seg_kernel, GemmSegArgs) \ + TENSOR_GEMM_ENTRY(STEM##_seg_store_f32, Q, float, false, tensor_gemm_seg_kernel, GemmSegArgs) \ + TENSOR_GEMM_ENTRY(STEM##_seg_resid_f16, Q, half, true, tensor_gemm_seg_kernel, GemmSegArgs) \ + TENSOR_GEMM_ENTRY(STEM##_seg_resid_f32, Q, float, true, tensor_gemm_seg_kernel, GemmSegArgs) + +TENSOR_GEMM_FAMILY(q4k, Q4K) +TENSOR_GEMM_FAMILY(q6k, Q6K) +TENSOR_GEMM_FAMILY(a4, A4) + +#define TENSOR_ROUTED_ENTRY(NAME, Q, T, PAIR) \ + template [[host_name("gemm_routed_" #NAME)]] \ + kernel void tensor_routed_kernel( \ + device const Q::Block * W [[buffer(0)]], \ + device const Q::Block * W_up [[buffer(1)]], \ + device const Q::Block * W_sh [[buffer(2)]], \ + device const Q::Block * W_sh_up [[buffer(3)]], \ + device const float * x [[buffer(4)]], \ + device float * y [[buffer(5)]], \ + device const uint4 * tiles [[buffer(6)]], \ + device const uint * sorted [[buffer(7)]], \ + device const float * route_w [[buffer(8)]], \ + constant RoutedArgs & args [[buffer(9)]], \ + device atomic_uint * faults [[buffer(10)]], \ + uint2 tgpig [[threadgroup_position_in_grid]], \ + ushort tiitg [[thread_index_in_threadgroup]], \ + ushort sgitg [[simdgroup_index_in_threadgroup]]); + +TENSOR_ROUTED_ENTRY(q4k_pair_f16, Q4K, half, true) +TENSOR_ROUTED_ENTRY(q4k_pair_f32, Q4K, float, true) +TENSOR_ROUTED_ENTRY(q4k_plain_f16, Q4K, half, false) +TENSOR_ROUTED_ENTRY(q4k_plain_f32, Q4K, float, false) +TENSOR_ROUTED_ENTRY(q6k_pair_f16, Q6K, half, true) +TENSOR_ROUTED_ENTRY(q6k_pair_f32, Q6K, float, true) +TENSOR_ROUTED_ENTRY(q6k_plain_f16, Q6K, half, false) +TENSOR_ROUTED_ENTRY(q6k_plain_f32, Q6K, float, false) +TENSOR_ROUTED_ENTRY(a4_pair_f16, A4, half, true) +TENSOR_ROUTED_ENTRY(a4_pair_f32, A4, float, true) +TENSOR_ROUTED_ENTRY(a4_plain_f16, A4, half, false) +TENSOR_ROUTED_ENTRY(a4_plain_f32, A4, float, false) diff --git a/crates/metal/src/bin/gate_metal_qwen35moe_e2e/mac.rs b/crates/metal/src/bin/gate_metal_qwen35moe_e2e/mac.rs index a8e7376a8..05a830abd 100644 --- a/crates/metal/src/bin/gate_metal_qwen35moe_e2e/mac.rs +++ b/crates/metal/src/bin/gate_metal_qwen35moe_e2e/mac.rs @@ -81,6 +81,8 @@ //! the profiled pass's bits; a pass's copies at rows inside it move none of its //! bits and are its state there, and a prompt cut back to a mark inside its //! pass continues the same every time, within the teacher-forced band. +//! - `q_tensor`: `q_free` with the pass's GEMMs forced onto the tensor-op +//! kernels, on any device that builds them (`pass.rs`). //! //! - `r_resident`: the session's kernel buffers are held in one residency set //! (`resident.rs`): the set holds every byte the open and the first pass @@ -110,7 +112,7 @@ use bloomery_device::{Device, DeviceError}; use bloomery_metal::arch::qwen3moe::count::{LiveKeys, Ticket, token_count}; use bloomery_metal::arch::qwen3moe::session::{Session, SessionError, StepMode}; use bloomery_metal::device::MetalDevice; -use bloomery_metal::library::Library; +use bloomery_metal::library::{Library, TENSOR_SOURCE}; #[path = "ahead.rs"] mod ahead; @@ -435,19 +437,29 @@ fn e_token_id(s: &mut Session) -> Clause { } /// The library the session's kernels come from: the device's, or the -/// source with a mutant's edits in. +/// source with a mutant's edits in. A tensor-op mutant leaves it the +/// device's: its edits are in the program `q_tensor` forces. fn library_of( device: &MetalDevice, - mutant: Option<&mutants::Mutant>, + mutant: Option<(&mutants::Mutant, mutants::Program)>, ) -> Result, String> { match mutant { - None => Ok(None), - Some(m) => { - let text = mutants::source_of(m)?; + Some((m, mutants::Program::Main)) => { + let text = mutants::source_of(m, mutants::Program::Main)?; Library::from_source_text(device, &text) .map(Some) .map_err(why) } + _ => Ok(None), + } +} + +/// The tensor-op program `q_tensor` forces: the shipped one, or a tensor-op +/// mutant's. +fn tensor_text_of(mutant: Option<(&mutants::Mutant, mutants::Program)>) -> Result { + match mutant { + Some((m, mutants::Program::Tensor)) => mutants::source_of(m, mutants::Program::Tensor), + _ => Ok(TENSOR_SOURCE.to_string()), } } @@ -460,7 +472,7 @@ pub fn main() -> ExitCode { println!("{GATE}: {name} takes a value"); return ExitCode::from(64); } - const NAMES: [&str; 23] = [ + const NAMES: [&str; 24] = [ "d_count", "e_token_id", "b_chain", @@ -479,6 +491,7 @@ pub fn main() -> ExitCode { "q_chunk", "q_ref", "q_free", + "q_tensor", "q_fault", "q_layers", "q_mark", @@ -523,10 +536,17 @@ pub fn main() -> ExitCode { } }; let library = mutated.as_ref().unwrap_or_else(|| device.library()); - if let Some(m) = mutant { + let tensor_text = match tensor_text_of(mutant) { + Ok(text) => text, + Err(e) => { + println!("{GATE}: the mutant does not build: {e}"); + return ExitCode::from(64); + } + }; + if let Some((m, _)) = mutant { println!("mutant: {} ({})", m.name, m.why); } - let only = mutant.map(|m| m.clause).or(flag(&args, "--only")); + let only = mutant.map(|(m, _)| m.clause).or(flag(&args, "--only")); if let Some(name) = only && !NAMES.contains(&name) { @@ -607,6 +627,7 @@ pub fn main() -> ExitCode { "q_chunk" => pass::q_chunk(&device, library, &path), "q_ref" => pass::q_ref(&device, library, &path), "q_free" => pass::q_free(&device, library, &path), + "q_tensor" => pass::q_tensor(&device, library, &tensor_text, &path), "q_fault" => pass::q_fault(&device, library, &path), "q_layers" => pass::q_layers(&device, library, &path), "q_mark" => pass::q_mark(&device, library, &path), @@ -617,7 +638,7 @@ pub fn main() -> ExitCode { report.clause(name, result); } let scope = match (mutant, only) { - (Some(m), _) => format!(" (mutant {})", m.name), + (Some((m, _)), _) => format!(" (mutant {})", m.name), (None, Some(name)) => format!(" (only {name})"), (None, None) => String::new(), }; @@ -635,7 +656,9 @@ fn mutants_all() -> ExitCode { } }; let mut green = Vec::new(); - for mutant in mutants::MUTANTS { + let mut count = 0usize; + for (mutant, _) in mutants::all() { + count += 1; let out = match Command::new(&exe).args(["--mutant", mutant.name]).output() { Ok(out) => out, Err(e) => { @@ -657,10 +680,7 @@ fn mutants_all() -> ExitCode { } } if green.is_empty() { - println!( - "{GATE} --mutants-all: PASS ({} mutants, each red on its clause)", - mutants::MUTANTS.len() - ); + println!("{GATE} --mutants-all: PASS ({count} mutants, each red on its clause)"); return ExitCode::SUCCESS; } println!("{GATE} --mutants-all: FAIL (not red: {})", green.join(", ")); diff --git a/crates/metal/src/bin/gate_metal_qwen35moe_e2e/mutants.rs b/crates/metal/src/bin/gate_metal_qwen35moe_e2e/mutants.rs index 683fbf79e..cc7b6a550 100644 --- a/crates/metal/src/bin/gate_metal_qwen35moe_e2e/mutants.rs +++ b/crates/metal/src/bin/gate_metal_qwen35moe_e2e/mutants.rs @@ -22,15 +22,19 @@ //! no kernel side. The host edits themselves are one-off edits recorded in the //! round that proves the clause. -use bloomery_metal::library::SOURCE; +#[path = "../shared/mutants.rs"] +mod program; -/// A defect and the clause that must be red on it. -pub struct Mutant { - pub name: &'static str, - pub clause: &'static str, - pub why: &'static str, - /// `(find, replace, occurrences of find)`, applied in order. - pub edits: &'static [(&'static str, &'static str, usize)], +pub use program::{Mutant, Program, source_of}; + +/// Every mutant of both tables, with the program it edits. +pub fn all() -> impl Iterator { + program::all(MUTANTS, TENSOR_MUTANTS) +} + +/// The mutant of this name, and the program it edits. +pub fn find(name: &str) -> Option<(&'static Mutant, Program)> { + all().find(|(m, _)| m.name == name) } pub const MUTANTS: &[Mutant] = &[ @@ -145,22 +149,11 @@ pub const MUTANTS: &[Mutant] = &[ }, ]; -pub fn find(name: &str) -> Option<&'static Mutant> { - MUTANTS.iter().find(|m| m.name == name) -} - -/// The kernels' source with the mutant's edits applied. -pub fn source_of(mutant: &Mutant) -> Result { - let mut text = SOURCE.to_string(); - for (find, replace, count) in mutant.edits { - let found = text.matches(find).count(); - if found != *count { - return Err(format!( - "mutant {}: {find:?} occurs {found} times, not {count}", - mutant.name - )); - } - text = text.replace(find, replace); - } - Ok(text) -} +/// The tensor-op GEMMs' defects `q_tensor` must catch: the program it forces +/// is the one its pass runs. +pub const TENSOR_MUTANTS: &[Mutant] = &[Mutant { + name: "tensor_pass_pair_swapped", + clause: "q_tensor", + why: "the routed pair's gate and up accumulators swap places in the tensor-op SwiGLU", + edits: &[("pair_swiglu(cg[i], cu[i])", "pair_swiglu(cu[i], cg[i])", 1)], +}]; diff --git a/crates/metal/src/bin/gate_metal_qwen35moe_e2e/pass.rs b/crates/metal/src/bin/gate_metal_qwen35moe_e2e/pass.rs index de8c3cb5a..b3ec6af37 100644 --- a/crates/metal/src/bin/gate_metal_qwen35moe_e2e/pass.rs +++ b/crates/metal/src/bin/gate_metal_qwen35moe_e2e/pass.rs @@ -38,6 +38,12 @@ //! to each mark and continued to the end, twice: the same pick, logits row //! and stores both times, and the tokens after it within `q_free`'s band of //! the stepped run's. +//! - `q_tensor`: `q_free` with the pass's GEMMs on the tensor-op kernels +//! (`Library::with_tensor_gemms`), forced on any device that builds them, +//! held to the same band and cap: the tensor path's dot band is `q_free`'s +//! f16 operand roundings plus an f32 sum of at most K − 1 additions a term +//! (the kernels gate's `gemm_band`), under 2^-11 of a dot's magnitude at +//! K = 2,048, so the pass is another realization of the same rule. use std::path::Path; @@ -49,7 +55,7 @@ use bloomery_metal::arch::qwen3moe::session::{ }; use bloomery_metal::device::MetalDevice; use bloomery_metal::kernels::{CONV_TAPS, HEAD_COLS, HEAD_DIM, KV_HEADS, RING_ROWS, VALUE_HEADS}; -use bloomery_metal::library::Library; +use bloomery_metal::library::{GemmPath, Library}; use runtime::prompt::PrefillMode; use runtime::{Target, Want}; @@ -607,6 +613,31 @@ pub fn q_ref(device: &MetalDevice, library: &Library, path: &Path) -> Clause { /// `q_free`: the greedy tokens of the stepped run teacher-forced through the /// pass run. pub fn q_free(device: &MetalDevice, library: &Library, path: &Path) -> Clause { + free_run(device, library, path, "q_free") +} + +/// `q_tensor`: `q_free` with the pass's GEMMs from the tensor-op program +/// `tensor_text` (the shipped one, or a mutant's). +pub fn q_tensor(device: &MetalDevice, library: &Library, tensor_text: &str, path: &Path) -> Clause { + let tensor = Library::tensor_from_source_text(device, tensor_text) + .map_err(|e| format!("the tensor-op GEMMs do not build here: {e}"))?; + let forced = library.with_tensor_gemms(tensor); + if forced.gemm_path() != GemmPath::TensorOps { + return Err(format!( + "the forced library's GEMMs are {}, not the tensor-op ones", + forced.gemm_path().name() + )); + } + let line = free_run(device, &forced, path, "q_tensor")?; + match line { + Verdict::Noted(line) => Ok(Verdict::Noted(format!("tensor-op GEMMs: {line}"))), + other => Ok(other), + } +} + +/// The greedy tokens of the stepped run teacher-forced through the pass run, +/// both sessions on `library`'s kernels. +fn free_run(device: &MetalDevice, library: &Library, path: &Path, clause: &str) -> Clause { let ids = prompt_ids(1024)?; let ctx = ids.len() + FREE_STEPS + 1; let mut pass = open(device, library, path, ctx, 2048, PrefillMode::Batch)?; @@ -626,7 +657,7 @@ pub fn q_free(device: &MetalDevice, library: &Library, path: &Path) -> Clause { if a.top == b.top { agree += 1; } - println!("q_free token {t}: picks pass={} steps={}", a.top, b.top); + println!("{clause} token {t}: picks pass={} steps={}", a.top, b.top); tok = b.top; } let line = format!( diff --git a/crates/metal/src/bin/gate_metal_qwen35moe_kernels/clauses_gemm.rs b/crates/metal/src/bin/gate_metal_qwen35moe_kernels/clauses_gemm.rs index f53bb4193..d8dd392a3 100644 --- a/crates/metal/src/bin/gate_metal_qwen35moe_kernels/clauses_gemm.rs +++ b/crates/metal/src/bin/gate_metal_qwen35moe_kernels/clauses_gemm.rs @@ -27,7 +27,20 @@ //! //! So the f32 arm's factor is `G(base + 1 + K/8 + 8)` and the f16 arm's is //! `(1 + 2^-11)² (1 + G(base + K/8 + 8)) − 1`, both over the oracle's -//! magnitude sum. The f16 arm carries two absolute terms besides: a cast +//! magnitude sum. +//! +//! The tensor-op kernels (`gemm_tensor.metal`, [`GemmPath::TensorOps`]) sum +//! in the op's order, which no document fixes, at the op's defined +//! precision (relaxed precision off): f32 accumulation, the f16 operands' +//! products exact in it and the f32 operands' rounded once. Any order of a +//! sum of K terms passes a term through at most K − 1 additions (every +//! addition on its path to the root of the summation tree), so the +//! accumulation term is K, whatever the order, in place of K/8 + 8: the f32 +//! arm's factor is `G(base + 1 + K)`, the f16 arm's +//! `(1 + 2^-11)² (1 + G(base + K)) − 1` — at K = 2048 a G of 2052 terms, +//! 1.2e-4 of the magnitude sum, at K = 512 3.1e-5. The casts, the subnormal +//! terms, the staged round-trip and the residual's rounding are the +//! simdgroup kernels' own. The f16 arm carries two absolute terms besides: a cast //! into the subnormal range rounds on the subnormal grid (spacing 2^-24), //! so its error is at most 2^-25 whatever the magnitude — the weights' side //! sums to `2^-25 Σ_c |x_c|` and the activations' to `2^-25 Σ_c |w_c|`. The @@ -42,6 +55,7 @@ use bloomery_device::{Buffer, Device, DeviceError, Elems}; use bloomery_metal::kernels::{ GEMM_ACTS, GemmMode, GemmParams, GemmSegParams, Operand, Quant, Segment, }; +use bloomery_metal::library::GemmPath; use gguf::quant::dequant_row; use super::common::{self, BadArg, Edit, LAYER}; @@ -251,14 +265,22 @@ pub(super) fn weight_abs_sums( } /// The GEMM's band: the relative factor over the oracle's magnitude sum and -/// the operations of the longest path (the absolute floor's count). The -/// derivation is this module's doc's. -pub(super) fn gemm_band(quant: Quant, cols: usize, operand: Operand) -> (f64, usize) { +/// the operations of the longest path (the absolute floor's count), for the +/// program `path` runs. The derivation is this module's doc's. +pub(super) fn gemm_band( + quant: Quant, + cols: usize, + operand: Operand, + path: GemmPath, +) -> (f64, usize) { let base = match quant { Quant::Q4K => 4, Quant::Q6K | Quant::A4 => 2, }; - let acc = cols / 8 + 8; + let acc = match path { + GemmPath::Simdgroup => cols / 8 + 8, + GemmPath::TensorOps => cols, + }; match operand { Operand::F32 => { let n = base + 1 + acc; @@ -280,7 +302,7 @@ pub(super) fn gemm_band(quant: Quant, cols: usize, operand: Operand) -> (f64, us fn run_checked(h: &Harness<'_>, launch: &GemmLaunch<'_>, w_abs: &[f64]) -> Result { let got = launch_gemm(h, launch)?; common::expect_clean(h, "the gemm launch")?; - check_rows(launch, &got, w_abs, None) + check_rows(launch, &got, w_abs, None, h.kernels.gemm_path()) } /// The band check of `got` against the oracle, row by row: `skip` names the @@ -291,8 +313,9 @@ fn check_rows( got: &[f32], w_abs: &[f64], skip: Option, + path: GemmPath, ) -> Result { - let (factor, ops) = gemm_band(launch.quant, launch.cols, launch.operand); + let (factor, ops) = gemm_band(launch.quant, launch.cols, launch.operand, path); let floor = ops as f64 * f64::from(f32::MIN_POSITIVE); let mut worst = 0f64; for r in 0..launch.acts { @@ -633,6 +656,36 @@ pub fn gemm_a4_resid(h: &Harness<'_>, mlx: &Fixture) -> Clause { run_cases(h, GemmMode::Residual, &a4_cases(mlx)?) } +/// Every dense arm of one weight type on `h`'s kernels (the tensor-op +/// kernels, for the `tensor_gemm_*` clauses): both epilogues in both +/// operands, every case within the band of `h`'s GEMM path. +fn tensor_arms(h: &Harness<'_>, cases: &[Case<'_>]) -> Clause { + let mut notes = Vec::new(); + for mode in [GemmMode::Store, GemmMode::Residual] { + let Verdict::Noted(note) = run_cases(h, mode, cases)? else { + return Err(format!("{mode:?}: the cases read no band")); + }; + notes.push(format!("{mode:?} {note}")); + } + Ok(Verdict::Noted(format!( + "{} kernels: {}", + h.kernels.gemm_path().name(), + notes.join("; ") + ))) +} + +pub fn tensor_gemm_q4k(h: &Harness<'_>, fixture: &Fixture) -> Clause { + tensor_arms(h, &q4k_cases(fixture)?) +} + +pub fn tensor_gemm_q6k(h: &Harness<'_>, fixture: &Fixture) -> Clause { + tensor_arms(h, &q6k_cases(fixture)?) +} + +pub fn tensor_gemm_a4(h: &Harness<'_>, mlx: &Fixture) -> Clause { + tensor_arms(h, &a4_cases(mlx)?) +} + /// Every segment of a segmented dispatch is the same bits as the same tensor /// run alone: the kernel does not change the order of a row's sum, and a /// segment's tiles are the ones a launch of it alone would run. One group @@ -866,7 +919,7 @@ pub fn gemm_range(h: &Harness<'_>, fixture: &Fixture) -> Clause { } } // The other rows are finite and in band. - check_rows(&launch, &got, &w_abs, Some(bad_row)) + check_rows(&launch, &got, &w_abs, Some(bad_row), h.kernels.gemm_path()) .map_err(|e| format!("{what}, another row: {e}"))?; } // The boundary: the largest finite half is taken, not faulted. diff --git a/crates/metal/src/bin/gate_metal_qwen35moe_kernels/clauses_route.rs b/crates/metal/src/bin/gate_metal_qwen35moe_kernels/clauses_route.rs index 45ba0f567..84c23c0ec 100644 --- a/crates/metal/src/bin/gate_metal_qwen35moe_kernels/clauses_route.rs +++ b/crates/metal/src/bin/gate_metal_qwen35moe_kernels/clauses_route.rs @@ -41,7 +41,7 @@ use bloomery_metal::kernels::{ LOGIT_STRIDE, Operand, PARTIALS_PER_HEAD, Quant, ROUTE_TILE, RouteSelectParams, RouteTableParams, RoutedKind, RoutedParams, VALUE_HEADS, route_table_words, route_tiles, }; -use bloomery_metal::library::{Library, SOURCE}; +use bloomery_metal::library::{GemmPath, Library, SOURCE}; use gguf::quant::GgmlType; use super::clauses_gemm::{HALF_SUBNORMAL, gemm_band, weight_abs_sums}; @@ -733,10 +733,11 @@ fn check_grouped( operand: Operand, stride: usize, every: Option, + path: GemmPath, ) -> Result<(f64, f64), String> { let slots = sc.rows * PER_ROW; - let (f_pair, ops_pair) = gemm_band(ex.quant, K, operand); - let (f_down, ops_down) = gemm_band(ex.quant, EXPERT_K, operand); + let (f_pair, ops_pair) = gemm_band(ex.quant, K, operand, path); + let (f_down, ops_down) = gemm_band(ex.quant, EXPERT_K, operand, path); // The f32 arm has no cast, so no subnormal term; its band is the dense // GEMM's own (`gemm_band`), and the subnormal terms are the f16 arm's. let sub = match operand { @@ -843,32 +844,28 @@ fn range_fault(h: &Harness<'_>, ex: &Experts) -> Result<(), String> { Ok(()) } -fn grouped_quant( - h: &Harness<'_>, - fix: &Fixture, - quant: Quant, - operands: &[Operand], -) -> Result { - let ex = experts_of(fix, quant)?; +fn grouped_quant(h: &Harness<'_>, ex: &Experts, operands: &[Operand]) -> Result { + let quant = ex.quant; let mut parts = Vec::new(); for &operand in operands { // A tile and eight rows: every row picks the hot expert, so its // group and the shared group both take two tiles. let sc = scene(ROUTE_TILE + 8, Some(5), 0x6A0); - let run = grouped_run(h, &ex, &sc, operand)?; + let run = grouped_run(h, ex, &sc, operand)?; for (what, fault) in [("pair", run.pair_fault), ("down", run.down_fault)] { if let Some(f) = fault { return Err(format!("the clean {what} launch raised {f}")); } } - let (pair, down) = check_grouped(&ex, &sc, &run, operand, 7, Some(5))?; + let (pair, down) = + check_grouped(ex, &sc, &run, operand, 7, Some(5), h.kernels.gemm_path())?; parts.push(format!( "{} {operand:?} pair {pair:.3} down {down:.3}", quant.name() )); } if operands.contains(&Operand::F16) { - range_fault(h, &ex)?; + range_fault(h, ex)?; parts.push(format!( "{} an activation past the half range is the fault and NaN", quant.name() @@ -882,12 +879,15 @@ fn grouped_quant( pub fn grouped_gemm(h: &Harness<'_>, fixture: &Fixture, mlx: &Result) -> Clause { let mut notes = vec![grouped_quant( h, - fixture, - Quant::Q4K, + &experts_of(fixture, Quant::Q4K)?, &[Operand::F16, Operand::F32], )?]; match mlx { - Ok(m) => notes.push(grouped_quant(h, m, Quant::A4, &[Operand::F16])?), + Ok(m) => notes.push(grouped_quant( + h, + &experts_of(m, Quant::A4)?, + &[Operand::F16], + )?), Err(e) => return Err(format!("the A4 fixture: {e}")), } Ok(Verdict::Noted(format!( @@ -896,9 +896,149 @@ pub fn grouped_gemm(h: &Harness<'_>, fixture: &Fixture, mlx: &Result Result { + let ty = GgmlType::Q6_K; + let qkv_rows = 8192; + let qkv = fix.rows("blk.2.attn_qkv.weight", ty, K, qkv_rows)?; + let row = qkv.len() / qkv_rows; + let cycled = |first: usize, rows: usize| { + let mut out = Vec::with_capacity(rows * row); + for r in 0..rows { + let at = (first + r) % qkv_rows; + out.extend_from_slice(&qkv[at * row..(at + 1) * row]); + } + out + }; + Ok(Experts { + quant: Quant::Q6K, + gate: cycled(0, EXPERTS * EXPERT_K), + up: cycled(qkv_rows / 2, EXPERTS * EXPERT_K), + down: stack(fix, "blk.2.ffn_down_exps.weight", ty, EXPERT_K, HIDDEN)?, + shared_gate: cycled(1000, EXPERT_K), + shared_up: cycled(5000, EXPERT_K), + shared_down: fix + .rows("blk.2.ffn_down_shexp.weight", ty, EXPERT_K, HIDDEN)? + .to_vec(), + }) +} + +/// A tile the table cannot write — its run one slot longer than a tile — +/// raises the expert-id site in the pair and in the down, and neither +/// stores a word of its places; the table's other tiles are stored. +fn bad_tile(h: &Harness<'_>, ex: &Experts) -> Result<(), String> { + let sc = scene(ROUTE_TILE + 8, Some(5), 0x6A1); + let table = table_run(h, &sc.ids)?; + if let Some(f) = table.fault { + return Err(format!("the table raised {f}")); + } + // Tile 0's words are (group, first place, run, unused). + let mut tiles = table.tiles.clone(); + let first = tiles[4 + 1] as usize; + let count = tiles[4 + 2] as usize; + tiles[4 + 2] = (ROUTE_TILE + 1) as u32; + let over = OverTable { + kernels: h.kernels, + sorted: &table.sorted, + tiles: &tiles, + max_tiles: route_tiles(sc.rows * PER_ROW, EXPERTS), + }; + let out = grouped_launch(h, ex, &sc, Operand::F16, &over)?; + expect_site( + out.pair_fault, + "the pair on a tile past the tile's slots", + FaultSite::ExpertId, + )?; + expect_site( + out.down_fault, + "the down on a tile past the tile's slots", + FaultSite::ExpertId, + )?; + for p in first..first + count { + for (what, row) in [ + ("pair", &out.h[p * EXPERT_K..(p + 1) * EXPERT_K]), + ("down", &out.y[p * HIDDEN..(p + 1) * HIDDEN]), + ] { + if let Some(v) = row.iter().find(|v| v.to_bits() != SENTINEL_BITS) { + return Err(format!( + "the {what} stored {v:?} at place {p} of the bad tile" + )); + } + } + } + let after = first + count; + if out.h[after * EXPERT_K..(after + 1) * EXPERT_K] + .iter() + .any(|v| v.to_bits() == SENTINEL_BITS) + { + return Err(format!( + "place {after}, the next tile's first, was not stored" + )); + } + Ok(()) +} + +/// The grouped GEMM's every arm of one weight type on `h`'s kernels (the +/// tensor-op kernels, for the `tensor_grouped_*` clauses): the pair and the +/// down in both operands, every checked place within the band of `h`'s +/// GEMM path; runs of 20 and 40 places (the hot group's and the shared +/// group's), which end inside a SIMD group's second block of places and +/// inside the second group's first; an activation past the half range the +/// fault and NaN; a tile past the tile's slots the expert-id site with +/// nothing stored. +fn grouped_arms(h: &Harness<'_>, ex: &Experts) -> Clause { + let note = grouped_quant(h, ex, &[Operand::F16, Operand::F32])?; + let mut runs = Vec::new(); + for rows in [20usize, 40] { + let sc = scene(rows, Some(9), 0x6A2 + rows as u64); + let run = grouped_run(h, ex, &sc, Operand::F16)?; + for (what, fault) in [("pair", run.pair_fault), ("down", run.down_fault)] { + if let Some(f) = fault { + return Err(format!( + "the clean {what} launch on runs of {rows} raised {f}" + )); + } + } + let (pair, down) = check_grouped( + ex, + &sc, + &run, + Operand::F16, + 1, + Some(9), + h.kernels.gemm_path(), + ) + .map_err(|e| format!("runs of {rows}: {e}"))?; + runs.push(format!("runs of {rows} pair {pair:.3} down {down:.3}")); + } + bad_tile(h, ex)?; + Ok(Verdict::Noted(format!( + "{} kernels, worst error over its band: {note}, {}; a tile past the tile's slots raises the expert-id site and stores nothing", + h.kernels.gemm_path().name(), + runs.join(", ") + ))) +} + +pub fn tensor_grouped_q4k(h: &Harness<'_>, fixture: &Fixture) -> Clause { + grouped_arms(h, &experts_of(fixture, Quant::Q4K)?) +} + +pub fn tensor_grouped_q6k(h: &Harness<'_>, fixture: &Fixture) -> Clause { + grouped_arms(h, &q6k_experts(fixture)?) +} + +pub fn tensor_grouped_a4(h: &Harness<'_>, mlx: &Fixture) -> Clause { + grouped_arms(h, &experts_of(mlx, Quant::A4)?) +} + // ------------------------------------------------------------ grouped oracle -/// The routed tile the gate keeps as the oracle, from `#define ROUTED_ROWS` on. +/// The routed tile the gate keeps as the oracle, from `#define ROUTED_FRAGS_A` +/// on. const ORACLE: &str = include_str!("routed_oracle.metal"); /// The shipped routed GEMM's file, whose routed tile the oracle replaces. @@ -924,15 +1064,22 @@ const LANES: [&[usize]; TOP_K] = [ ]; /// The kernels compiled from the shipped source with the routed tile swapped -/// for the oracle's and the table's tile at the oracle's. +/// for the oracle's, and the table's tile and the routed tile's activation +/// rows at the oracle's. fn oracle_kernels(h: &Harness<'_>) -> Result { let at = SHIPPED - .find("#define ROUTED_ROWS") - .ok_or("gemm_routed.metal has no `#define ROUTED_ROWS`; the tile's start moved")?; + .find("#define ROUTED_FRAGS_A") + .ok_or("gemm_routed.metal has no `#define ROUTED_FRAGS_A`; the tile's start moved")?; let from = format!("#define ROUTE_TILE {ROUTE_TILE}\n"); let to = format!("#define ROUTE_TILE {ORACLE_TILE}\n"); + let acts_from = format!("#define ROUTED_ACTS {ROUTE_TILE} //"); + let acts_to = format!("#define ROUTED_ACTS {ORACLE_TILE} //"); let mut text = SOURCE.to_string(); - for (find, with) in [(&SHIPPED[at..], ORACLE), (from.as_str(), to.as_str())] { + for (find, with) in [ + (&SHIPPED[at..], ORACLE), + (from.as_str(), to.as_str()), + (acts_from.as_str(), acts_to.as_str()), + ] { let found = text.matches(find).count(); if found != 1 { return Err(format!( diff --git a/crates/metal/src/bin/gate_metal_qwen35moe_kernels/clauses_tensor.rs b/crates/metal/src/bin/gate_metal_qwen35moe_kernels/clauses_tensor.rs new file mode 100644 index 000000000..f5a898342 --- /dev/null +++ b/crates/metal/src/bin/gate_metal_qwen35moe_kernels/clauses_tensor.rs @@ -0,0 +1,79 @@ +//! The tensor-op GEMMs (`kernels/gemm_tensor.metal`), forced: every clause +//! here runs them on the device that runs the gate, whatever its matrix-unit +//! rule picks for the engine, through a library whose GEMMs come from the +//! tensor-op program (`Library::with_tensor_gemms`). A device whose runtime +//! does not build that program fails each clause by name; none is skipped. +//! Each clause starts from a clear fault word, whatever the clause before +//! it left raised. +//! +//! - `tensor_gemm_q4k`, `tensor_gemm_q6k`, `tensor_gemm_a4`: the dense GEMM's +//! cases of `gemm_*_store` and `gemm_*_resid`, both epilogues in both +//! operands, within the tensor path's band (`clauses_gemm`'s doc). +//! - `tensor_gemm_seg`: `gemm_seg` on these kernels — a segment is the same +//! bits as the same tensor run alone. +//! - `tensor_gemm_range`: `gemm_range` on these kernels — the f16 operand's +//! range check raises the quantizer's site and makes the row NaN. +//! - `tensor_grouped_q4k`, `tensor_grouped_q6k`, `tensor_grouped_a4`: the +//! routed pair and down over the route table's tiles, both operands, +//! within the tensor path's band; an activation past the half range the +//! fault and NaN; a tile past the tile's slots the expert-id site with +//! nothing stored. + +use bloomery_device::conformance::Clause; + +use super::fixture::Fixture; +use super::{ClauseRun, Harness, clauses_gemm, clauses_route}; + +/// The tensor-op clauses over `tensor`, the harness of the tensor-op GEMMs +/// (or why this device has none). +pub fn table<'a>( + tensor: &'a Result, String>, + fixture: &'a Result, + mlx: &'a Result, +) -> Vec<(&'static str, ClauseRun<'a>)> { + let on = |f: fn(&Harness<'_>, &Fixture) -> Clause, fix: &'a Result| { + Box::new(move || match (tensor, fix) { + (Err(e), _) => Err(format!("the tensor-op GEMMs do not build here: {e}")), + (_, Err(e)) => Err(e.clone()), + (Ok(h), Ok(fix)) => { + h.clear_fault()?; + f(h, fix) + } + }) as ClauseRun<'a> + }; + vec![ + ( + "tensor_gemm_q4k", + on(clauses_gemm::tensor_gemm_q4k, fixture), + ), + ( + "tensor_gemm_q6k", + on(clauses_gemm::tensor_gemm_q6k, fixture), + ), + ("tensor_gemm_a4", on(clauses_gemm::tensor_gemm_a4, mlx)), + ( + "tensor_gemm_seg", + Box::new(move || match (tensor, fixture, mlx) { + (Err(e), _, _) => Err(format!("the tensor-op GEMMs do not build here: {e}")), + (_, Err(e), _) | (_, _, Err(e)) => Err(e.clone()), + (Ok(h), Ok(fix), Ok(split)) => { + h.clear_fault()?; + clauses_gemm::gemm_seg(h, fix, split) + } + }), + ), + ("tensor_gemm_range", on(clauses_gemm::gemm_range, fixture)), + ( + "tensor_grouped_q4k", + on(clauses_route::tensor_grouped_q4k, fixture), + ), + ( + "tensor_grouped_q6k", + on(clauses_route::tensor_grouped_q6k, fixture), + ), + ( + "tensor_grouped_a4", + on(clauses_route::tensor_grouped_a4, mlx), + ), + ] +} diff --git a/crates/metal/src/bin/gate_metal_qwen35moe_kernels/mac.rs b/crates/metal/src/bin/gate_metal_qwen35moe_kernels/mac.rs index ebcb24a7f..7e958ab70 100644 --- a/crates/metal/src/bin/gate_metal_qwen35moe_kernels/mac.rs +++ b/crates/metal/src/bin/gate_metal_qwen35moe_kernels/mac.rs @@ -7,7 +7,13 @@ //! same blocks with an f64 dot, held to the band `oracle.rs` derives. //! `--list` prints the clause names in the order the gate runs them: the //! mat-vec clauses, then the A4 ones, then the token's kernels (delta, -//! router, experts, attention, head). +//! router, experts, attention, head), then the prompt pass's, then the +//! tensor-op GEMMs' (`clauses_tensor.rs`). +//! +//! The GEMM clauses outside `clauses_tensor.rs` run the simdgroup GEMMs on +//! every device; the `tensor_*` clauses run the tensor-op GEMMs on every +//! device that builds them. The header names the path the device's own rule +//! picks for the engine. //! //! `--fixture PATH` names the fixture's trunk file (the Mac's cache copy //! otherwise); `--mlx-fixture PATH` names the MLX-format fixture split the @@ -30,7 +36,7 @@ use bloomery_metal::device::MetalDevice; use bloomery_metal::kernels::{ HeadNormParams, Kernels, MatvecParams, Mode, NormOnLoad, Quant, SegParams, Segment, }; -use bloomery_metal::library::{Library, SOURCE}; +use bloomery_metal::library::{Library, SOURCE, TENSOR_SOURCE}; use bloomery_metal::stream::MetalStream; #[path = "clauses_a4.rs"] @@ -59,6 +65,8 @@ mod clauses_router; mod clauses_rows; #[path = "clauses_seg.rs"] mod clauses_seg; +#[path = "clauses_tensor.rs"] +mod clauses_tensor; #[path = "common.rs"] mod common; #[path = "fixture.rs"] @@ -401,6 +409,18 @@ fn kernels_of_text(device: &MetalDevice, text: &str) -> Result<(Library, Kernels Ok((library, kernels)) } +/// `base` with its GEMMs from the tensor-op program `text`, and its kernels. +fn tensor_kernels_of_text( + device: &MetalDevice, + base: &Library, + text: &str, +) -> Result<(Library, Kernels), String> { + let tensor = Library::tensor_from_source_text(device, text).map_err(why)?; + let library = base.with_tensor_gemms(tensor); + let kernels = Kernels::from_library(&library).map_err(why)?; + Ok((library, kernels)) +} + /// The fixture path: `--fixture PATH`, else the qwen35moe fixture under the /// fixture root (`BLOOMERY_FIXTURE_ROOT`), as the reference sets resolve it. fn fixture_path(args: &[String]) -> Result { @@ -433,7 +453,9 @@ fn mutants_all(args: &[String]) -> ExitCode { } }; let mut green = Vec::new(); - for mutant in mutants::MUTANTS { + let mut count = 0usize; + for (mutant, _) in mutants::all() { + count += 1; let mut child = Command::new(&exe); child.args(["--mutant", mutant.name]); if let Some(path) = flag(args, "--fixture") { @@ -464,8 +486,7 @@ fn mutants_all(args: &[String]) -> ExitCode { } if green.is_empty() { println!( - "gate_metal_qwen35moe_kernels --mutants-all: PASS ({} mutants, each red on its clause)", - mutants::MUTANTS.len() + "gate_metal_qwen35moe_kernels --mutants-all: PASS ({count} mutants, each red on its clause)" ); return ExitCode::SUCCESS; } @@ -506,7 +527,10 @@ pub fn main() -> ExitCode { } }; let stream = device.stream(); - let kernels = match Kernels::new(&device) { + // The simdgroup GEMMs, on every device: the tensor-op ones have clauses + // of their own. + let simdgroup = device.library().with_simdgroup_gemms(); + let kernels = match Kernels::from_library(&simdgroup) { Ok(kernels) => kernels, Err(e) => { println!("clause open: FAIL ({e})"); @@ -526,6 +550,14 @@ pub fn main() -> ExitCode { info.name, ::BACKEND ); + let units = device.library().matrix_units(); + println!( + "gemm: architecture={} metal4={} matrix-units={} engine-path={}", + units.architecture, + units.metal4, + units.present(), + device.library().gemm_path().name() + ); let mutant = match flag(&args, "--mutant") { Some(name) => match mutants::find(name) { Some(mutant) => Some(mutant), @@ -537,32 +569,63 @@ pub fn main() -> ExitCode { None => None, }; let (run_source, mutated) = match mutant { - Some(mutant) => match mutants::source_of(mutant) { - Ok(text) => match kernels_of_text(&device, &text) { - Ok(mutated) => (text, Some(mutated)), + Some((mutant, mutants::Program::Main)) => { + match mutants::source_of(mutant, mutants::Program::Main) { + Ok(text) => match kernels_of_text(&device, &text) { + Ok(mutated) => (text, Some(mutated)), + Err(e) => { + println!("gate_metal_qwen35moe_kernels: the mutant does not build: {e}"); + return ExitCode::from(64); + } + }, Err(e) => { println!("gate_metal_qwen35moe_kernels: the mutant does not build: {e}"); return ExitCode::from(64); } - }, - Err(e) => { - println!("gate_metal_qwen35moe_kernels: the mutant does not build: {e}"); - return ExitCode::from(64); } - }, - None => (SOURCE.to_string(), None), + } + _ => (SOURCE.to_string(), None), + }; + let tensor_source = match mutant { + Some((mutant, mutants::Program::Tensor)) => { + match mutants::source_of(mutant, mutants::Program::Tensor) { + Ok(text) => text, + Err(e) => { + println!("gate_metal_qwen35moe_kernels: the mutant does not build: {e}"); + return ExitCode::from(64); + } + } + } + _ => TENSOR_SOURCE.to_string(), }; let harness = Harness { device: &device, stream, kernels: mutated.as_ref().map_or(&kernels, |(_, k)| k), - library: mutated.as_ref().map_or(device.library(), |(l, _)| l), + library: mutated.as_ref().map_or(&simdgroup, |(l, _)| l), }; - if let Some(mutant) = mutant { + let tensor_built = tensor_kernels_of_text(&device, harness.library, &tensor_source); + if let (Some((mutant, mutants::Program::Tensor)), Err(e)) = (mutant, &tensor_built) { + println!( + "gate_metal_qwen35moe_kernels: the mutant {} does not build: {e}", + mutant.name + ); + return ExitCode::from(64); + } + let tensor = tensor_built + .as_ref() + .map_err(Clone::clone) + .map(|(l, k)| Harness { + device: &device, + stream, + kernels: k, + library: l, + }); + if let Some((mutant, _)) = mutant { println!("mutant: {} ({})", mutant.name, mutant.why); } let only = flag(&args, "--child") - .or(mutant.map(|m| m.clause)) + .or(mutant.map(|(m, _)| m.clause)) .or(flag(&args, "--only")); let fixture = fixture_path(&args).and_then(|path| fixture::Fixture::open(&path)); match &fixture { @@ -647,6 +710,7 @@ pub fn main() -> ExitCode { .chain(clauses_head::table(&ctx)) .chain(clauses_rows::table(&ctx)) .chain(clauses_route::table(&ctx, &mlx)) + .chain(clauses_tensor::table(&tensor, &fixture, &mlx)) .collect(); if args.iter().any(|a| a == "--list") { for (name, _) in &clauses { @@ -667,7 +731,7 @@ pub fn main() -> ExitCode { } let scope = match (mutant, only) { - (Some(mutant), _) => format!(" (mutant {})", mutant.name), + (Some((mutant, _)), _) => format!(" (mutant {})", mutant.name), (None, Some(name)) => format!(" (only {name})"), (None, None) => String::new(), }; diff --git a/crates/metal/src/bin/gate_metal_qwen35moe_kernels/mutants.rs b/crates/metal/src/bin/gate_metal_qwen35moe_kernels/mutants.rs index adfa6d380..bf45e8e70 100644 --- a/crates/metal/src/bin/gate_metal_qwen35moe_kernels/mutants.rs +++ b/crates/metal/src/bin/gate_metal_qwen35moe_kernels/mutants.rs @@ -14,16 +14,24 @@ //! `attn_bounds` also have a kernel-side contract, and the kernel mutants //! below pin that side. The gate's refusal loop turns red on a deleted launch //! check; a table that stayed green on one would be wrong. +//! +//! The tensor-op GEMMs' mutants ([`TENSOR_MUTANTS`]) edit their own program's +//! text ([`TENSOR_SOURCE`]) and are caught by the `tensor_*` clauses, which +//! run those kernels on any device that compiles them. + +#[path = "../shared/mutants.rs"] +mod program; -use bloomery_metal::library::SOURCE; +pub use program::{Mutant, Program, source_of}; + +/// Every mutant of both tables, with the program it edits. +pub fn all() -> impl Iterator { + program::all(MUTANTS, TENSOR_MUTANTS) +} -/// A defect and the clause that must be red on it. -pub struct Mutant { - pub name: &'static str, - pub clause: &'static str, - pub why: &'static str, - /// `(find, replace, occurrences of find)`, applied in order. - pub edits: &'static [(&'static str, &'static str, usize)], +/// The mutant of this name, and the program it edits. +pub fn find(name: &str) -> Option<(&'static Mutant, Program)> { + all().find(|(m, _)| m.name == name) } const TAIL_STORE_FIND: &str = @@ -1025,8 +1033,8 @@ pub const MUTANTS: &[Mutant] = &[ clause: "grouped_gemm", why: "the grouped GEMM's f16 range check is gone, so an activation past the half range casts to inf", edits: &[( - "const bool4 over = !(fabs(v) <= 65504.f);", - "const bool4 over = bool4(false);", + "const bool4 bad = !(fabs(v) <= 65504.f);", + "const bool4 bad = bool4(false);", 1, )], }, @@ -1088,23 +1096,94 @@ pub const MUTANTS: &[Mutant] = &[ }, ]; -/// The mutant of this name. -pub fn find(name: &str) -> Option<&'static Mutant> { - MUTANTS.iter().find(|m| m.name == name) -} - -/// The kernels' source with the mutant's edits applied. -pub fn source_of(mutant: &Mutant) -> Result { - let mut text = SOURCE.to_string(); - for (find, replace, count) in mutant.edits { - let found = text.matches(find).count(); - if found != *count { - return Err(format!( - "mutant {}: {find:?} occurs {found} times, not {count}", - mutant.name - )); - } - text = text.replace(find, replace); - } - Ok(text) -} +/// The tensor-op GEMMs' defects, at least one for each `tensor_*` clause. +pub const TENSOR_MUTANTS: &[Mutant] = &[ + Mutant { + name: "tensor_k_step_skipped", + clause: "tensor_gemm_q4k", + why: "the dense tile's K loop stops one step short: the last 32 columns are never multiplied", + edits: &[( + " for (int k0 = 0; k0 < cols; k0 += BK) {", + " for (int k0 = 0; k0 + BK < cols; k0 += BK) {", + 1, + )], + }, + Mutant { + name: "tensor_rows_offset", + clause: "tensor_gemm_q6k", + why: "the dense op's weight fragment starts one staged row late: every output takes the next weight row", + edits: &[( + "auto ta = tensor_rows(sb + tm * LDA);\n auto tw = tensor_rows(sa + tn * LDA);", + "auto ta = tensor_rows(sb + tm * LDA);\n auto tw = tensor_rows(sa + (tn + 1) * LDA);", + 1, + )], + }, + Mutant { + name: "tensor_resid_stores", + clause: "tensor_gemm_a4", + why: "the residual epilogue stores the product over the output instead of adding to it", + edits: &[("*p = RESID ? *p + ct[i] : ct[i];", "*p = ct[i];", 1)], + }, + Mutant { + name: "tensor_seg_out", + clause: "tensor_gemm_seg", + why: "every segment stores at the first segment's output offset", + edits: &[("y + args.segs[s].out", "y", 1)], + }, + Mutant { + name: "tensor_range_no_raise", + clause: "tensor_gemm_range", + why: "an activation past the half range stages NaN but raises no quantizer site", + edits: &[( + " if (any(bad)) {\n raise(faults, layer, BLOOMERY_FAULT_SITE_QUANT_COLUMN);", + " if (any(bad)) {", + 1, + )], + }, + Mutant { + name: "tensor_pair_swapped", + clause: "tensor_grouped_q4k", + why: "the pair's gate and up accumulators swap places in the SwiGLU", + edits: &[("pair_swiglu(cg[i], cu[i])", "pair_swiglu(cu[i], cg[i])", 1)], + }, + Mutant { + name: "tensor_routed_short_block", + clause: "tensor_grouped_q4k", + why: "a SIMD group's second block of places runs only when the run passes it by more than 8: a run ending 1 to 8 places into the block leaves them unstored", + edits: &[( + "const bool live1 = tm + TENSOR_RM < count;", + "const bool live1 = tm + TENSOR_RM + 8 < count;", + 1, + )], + }, + Mutant { + name: "tensor_routed_block_rows", + clause: "tensor_grouped_a4", + why: "a SIMD group's second block of places reads the first block's staged rows", + edits: &[( + "auto ta1 = tensor_rows(sb + (tm + TENSOR_RM) * LDA);", + "auto ta1 = tensor_rows(sb + tm * LDA);", + 1, + )], + }, + Mutant { + name: "tensor_routed_k_step_skipped", + clause: "tensor_grouped_q6k", + why: "the routed tile's K loops (pair and plain) stop one step short", + edits: &[( + " for (int k0 = 0; k0 < (int)args.cols; k0 += BK) {", + " for (int k0 = 0; k0 + BK < (int)args.cols; k0 += BK) {", + 2, + )], + }, + Mutant { + name: "tensor_tile_no_raise", + clause: "tensor_grouped_a4", + why: "a tile the table cannot write stores nothing but raises no expert-id site", + edits: &[( + " raise(faults, args.layer, BLOOMERY_FAULT_SITE_EXPERT_ID);", + "", + 1, + )], + }, +]; diff --git a/crates/metal/src/bin/gate_metal_qwen35moe_kernels/routed_oracle.metal b/crates/metal/src/bin/gate_metal_qwen35moe_kernels/routed_oracle.metal index 1fa5cdedd..996078f3e 100644 --- a/crates/metal/src/bin/gate_metal_qwen35moe_kernels/routed_oracle.metal +++ b/crates/metal/src/bin/gate_metal_qwen35moe_kernels/routed_oracle.metal @@ -2,13 +2,12 @@ // at a 32-slot table tile, each SIMD group multiplying two 16-row fragments // of activation rows (8 rows each, a test a fragment for the run's length) // against two 16-row fragments of weight rows. The gate compiles it in place -// of gemm_routed.metal's routed tile, from `#define ROUTED_ROWS` on, with the -// table's tile at 32, and holds the shipped tile to its outputs bit for bit: -// a tile's height and its split over the SIMD groups leave every output's K -// order as it is. +// of gemm_routed.metal's routed tile, from `#define ROUTED_FRAGS_A` on, with +// the table's tile and gemm_shared.metal's ROUTED_ACTS at 32 (the tile's +// weight rows and its argument block are gemm_shared.metal's), and holds the +// shipped tile to its outputs bit for bit: a tile's height and its split over +// the SIMD groups leave every output's K order as it is. -#define ROUTED_ROWS 64 // weight rows of a tile -#define ROUTED_ACTS 32 // activation rows of a tile #define ROUTED_FRAGS_A (ROUTED_ACTS / 16) #define ROUTED_FRAGS_W (ROUTED_ROWS / 16) @@ -21,20 +20,6 @@ constexpr int routed_shmem_bytes() { return (ROUTED_ROWS + ROUTED_ACTS) * gemm_lda() * (int)sizeof(T); } -// Eight words, bound by value. -struct RoutedArgs { - uint cols; // the input columns K - uint x_stride; // f32 elements between the activation rows - uint y_stride; // f32 elements between the output rows - uint layer; - uint n_experts; // the stack's experts; the table's group n_experts is shared - uint stack_blocks; // blocks of one expert's rows in the stack - uint gather; // 1: an activation row is the slot's pass row; 0: the place's - uint rows; // outputs per expert: the pair's rows of gate, the down's rows -}; - -static_assert(sizeof(RoutedArgs) == 32, "RoutedArgs is eight words"); - template kernel void gemm_routed_kernel(device const typename Q::Block * W [[buffer(0)]], device const typename Q::Block * W_up [[buffer(1)]], diff --git a/crates/metal/src/bin/generate_metal.rs b/crates/metal/src/bin/generate_metal.rs index 5a88bb879..5d039fa18 100644 --- a/crates/metal/src/bin/generate_metal.rs +++ b/crates/metal/src/bin/generate_metal.rs @@ -22,12 +22,13 @@ //! `--tokens` spelling, so a run's output feeds back as an input), for //! gates. The summary lines go to stderr, `key=value` one a line: //! -//! model= open_s= open_load_s= open_rows_ms= -//! open_residency_ms= open_pipelines_ms= device_allocated_bytes= rows_arena_bytes= -//! prompt_ids= +//! model= gemm_path= open_s= +//! open_load_s= open_rows_ms= open_residency_ms= +//! open_pipelines_ms= device_allocated_bytes= rows_arena_bytes= prompt_ids= //! prompt_s= prompt_tok_s= decode_tokens= //! decode_passes= decode_s= decode_tok_s= //! stop= decode_first_ms= @@ -789,6 +790,7 @@ generate_metal (--prompt | --tokens a,b,c) [-n N] [--ctx C] [--print-ids] let decode_s = out.decode_wall.as_secs_f64(); eprintln!("# development-loop numbers: no lease, no witness"); eprintln!("model={}", a.model.display()); + eprintln!("gemm_path={}", device.library().gemm_path().name()); eprintln!("open_s={open_s:.3}"); let split = session.open_split(); eprintln!("open_load_s={:.3}", split.load.as_secs_f64()); diff --git a/crates/metal/src/bin/metal_gemm_time/mac.rs b/crates/metal/src/bin/metal_gemm_time/mac.rs index b0277b9e3..9b006c7e8 100644 --- a/crates/metal/src/bin/metal_gemm_time/mac.rs +++ b/crates/metal/src/bin/metal_gemm_time/mac.rs @@ -23,7 +23,14 @@ //! //! The arms (`--arm NAME`, repeatable; `f16` and `f32` by default): //! -//! - `f16`, `f32`: the crate's GEMM, launched through `Kernels::gemm`. +//! - `f16`, `f32`: the crate's simdgroup GEMM, launched through +//! `Kernels::gemm`, on every device (the GEMM path the device's own rule +//! picks is printed in the header, not timed by these arms). +//! - `tensor16`, `tensor32`: the tensor-op GEMM (`gemm_tensor.metal`, Metal +//! 4's `matmul2d`), launched through `Kernels::gemm` on kernels whose GEMMs +//! come from that program, on any device that builds it. +//! - `tensor16_relaxed`, `tensor32_relaxed`: the same with the op's relaxed +//! precision on (MLX's and llama.cpp's setting; the engine keeps it off). //! - `mlx16`, `mlx32`: MLX's `qmm_t` for 4-bit group-64 weights, A4 only //! (`probe_mlx_qmm.metal`, its own library), reading MLX's layout; its //! output is checked against a CPU product. @@ -57,6 +64,10 @@ //! `--routed` times the experts' grouped GEMM over a route table instead of //! the dense shapes (`routed.rs`, with its own `--rows`, `--skew` and //! `--hot`). +//! +//! The header's `gemm:` line names the device's GPU architecture, whether +//! it reports Metal 4's family, the matrix-unit rule's verdict and the GEMM +//! path the engine takes on it. use std::process::ExitCode; use std::time::{Duration, Instant}; @@ -67,7 +78,7 @@ use bloomery_metal::encode::{Binding, Bindings, ComputeEncoder}; use bloomery_metal::kernels::{ GEMM_ACTS, GEMM_ROWS, GEMM_THREADS, GemmMode, GemmParams, Kernels, Operand, Quant, }; -use bloomery_metal::library::SOURCE; +use bloomery_metal::library::{SOURCE, TENSOR_SOURCE}; use bloomery_metal::{GpuInterval, Library, MetalBuffer, MetalStream}; #[path = "../shared/oracle.rs"] @@ -178,8 +189,11 @@ struct Method { /// What is timed at a configuration. #[derive(Clone, Copy)] enum Arm { - /// The crate's GEMM, through `Kernels::gemm`. + /// The crate's simdgroup GEMM, through `Kernels::gemm`. Engine(Operand), + /// The tensor-op GEMM, through `Kernels::gemm`; the op's relaxed + /// precision on when the flag is. + Tensor(Operand, bool), /// MLX's `qmm_t`, A4 only. Mlx(Operand), /// The shipped source with a variant's edits. @@ -198,18 +212,28 @@ impl Arm { fn name(self) -> String { match self { Self::Engine(o) => operand_name(o).to_string(), + Self::Tensor(o, relaxed) => format!( + "tensor{}{}", + &operand_name(o)[1..], + if relaxed { "_relaxed" } else { "" } + ), Self::Mlx(Operand::F16) => "mlx16".to_string(), Self::Mlx(Operand::F32) => "mlx32".to_string(), Self::Variant(v, o) => format!("{}_{}", v.name, operand_name(o)), } } - /// The arm named `name`: `f16`, `f32`, `mlx16`, `mlx32` or + /// The arm named `name`: `f16`, `f32`, `tensor16`, `tensor32`, + /// `tensor16_relaxed`, `tensor32_relaxed`, `mlx16`, `mlx32` or /// `_f16` / `_f32`. fn parse(name: &str) -> Result { match name { "f16" => return Ok(Self::Engine(Operand::F16)), "f32" => return Ok(Self::Engine(Operand::F32)), + "tensor16" => return Ok(Self::Tensor(Operand::F16, false)), + "tensor32" => return Ok(Self::Tensor(Operand::F32, false)), + "tensor16_relaxed" => return Ok(Self::Tensor(Operand::F16, true)), + "tensor32_relaxed" => return Ok(Self::Tensor(Operand::F32, true)), "mlx16" => return Ok(Self::Mlx(Operand::F16)), "mlx32" => return Ok(Self::Mlx(Operand::F32)), _ => {} @@ -226,7 +250,7 @@ impl Arm { .map(|v| v.name.to_string()) .collect(); Err(format!( - "no arm named {name}: f16, f32, mlx16, mlx32, or _f16 / _f32 of {}", + "no arm named {name}: f16, f32, tensor16, tensor32, tensor16_relaxed, tensor32_relaxed, mlx16, mlx32, or _f16 / _f32 of {}", known.join(", ") )) } @@ -236,9 +260,46 @@ impl Arm { struct Libs { mlx: Option, variants: Vec<(&'static str, Library)>, + /// The tensor-op program and the kernels over it, strict then relaxed. + tensor: [Option<(Library, Kernels)>; 2], +} + +/// The tensor-op GEMMs' program, with the op's relaxed precision on when +/// `relaxed` is. +pub(crate) fn tensor_library(device: &MetalDevice, relaxed: bool) -> Result { + const STRICT: &str = "#define TENSOR_RELAXED false"; + let text = if relaxed { + if TENSOR_SOURCE.matches(STRICT).count() != 1 { + return Err(format!( + "the tensor-op source holds {STRICT:?} other than once" + )); + } + TENSOR_SOURCE.replace(STRICT, "#define TENSOR_RELAXED true") + } else { + TENSOR_SOURCE.to_string() + }; + Library::tensor_from_source_text(device, &text) + .map_err(|e| format!("the tensor-op GEMMs do not build here: {e}")) } impl Libs { + /// The tensor-op program and the kernels whose GEMMs come from it. + fn tensor( + &mut self, + device: &MetalDevice, + relaxed: bool, + ) -> Result<&(Library, Kernels), String> { + let slot = &mut self.tensor[usize::from(relaxed)]; + if slot.is_none() { + let library = tensor_library(device, relaxed)?; + let kernels = + Kernels::from_library(&device.library().with_tensor_gemms(library.clone())) + .map_err(why)?; + *slot = Some((library, kernels)); + } + Ok(slot.as_ref().expect("the kernels were just made")) + } + fn mlx(&mut self, device: &MetalDevice) -> Result<&Library, String> { if self.mlx.is_none() { self.mlx = Some(Library::from_source_text(device, mlx::SOURCE).map_err(why)?); @@ -361,9 +422,18 @@ fn run() -> Result<(), String> { info.name, device.library().path().name() ); + let units = device.library().matrix_units(); + println!( + "gemm: architecture={} metal4={} matrix-units={} engine-path={}", + units.architecture, + units.metal4, + units.present(), + device.library().gemm_path().name() + ); let mut libs = Libs { mlx: None, variants: Vec::new(), + tensor: [None, None], }; for arm in &arms { if let Arm::Variant(v, _) = arm { @@ -385,7 +455,8 @@ fn run() -> Result<(), String> { method.gflop, method.copies ); - let kernels = Kernels::new(&device).map_err(why)?; + // The `f16` and `f32` arms' GEMMs are the simdgroup ones on every device. + let kernels = Kernels::from_library(&device.library().with_simdgroup_gemms()).map_err(why)?; let ctx = Ctx { device: &device, stream, @@ -474,6 +545,9 @@ fn pipelines(device: &MetalDevice, libs: &mut Libs, arms: &[Arm]) -> Result<(), for quant in [Quant::Q4K, Quant::Q6K, Quant::A4] { let (library, name): (&Library, String) = match arm { Arm::Engine(o) => (device.library(), engine_kernel(quant, o)), + Arm::Tensor(o, relaxed) => { + (&libs.tensor(device, relaxed)?.0, engine_kernel(quant, o)) + } Arm::Mlx(o) => { if quant != Quant::A4 { continue; @@ -504,7 +578,7 @@ fn peaks(device: &MetalDevice, stream: &MetalStream, method: Method) -> Result<( let ctx = Ctx { device, stream, - kernels: &Kernels::new(device).map_err(why)?, + kernels: &Kernels::from_library(&device.library().with_simdgroup_gemms()).map_err(why)?, method, }; let floats = vec![0.0625f32; 512]; @@ -672,28 +746,34 @@ fn one(ctx: &Ctx<'_>, libs: &mut Libs, arm: Arm, cfg: &Config) -> Result<(), Str let fault = ctx.device.fault_word(); let timed = match arm { - Arm::Engine(operand) => time_dispatches(ctx, pool.copies, dispatches, |encoder, copy| { - ctx.kernels - .gemm( - encoder, - GemmParams { - quant, - rows: n, - cols: k, - acts: r, - weights: pool.window(copy, 0)?, - x: x_f32.window(Elems(0), Elems(r * k)).map_err(why)?, - x_stride: k, - y: y.window_mut(Elems(0), Elems(r * n)).map_err(why)?, - y_stride: n, - mode: GemmMode::Store, - operand, - fault, - layer: LAYER_NONE, - }, - ) - .map_err(|e| format!("the dispatch did not encode: {e}")) - })?, + Arm::Engine(operand) | Arm::Tensor(operand, _) => { + let kernels = match arm { + Arm::Tensor(_, relaxed) => &libs.tensor(ctx.device, relaxed)?.1, + _ => ctx.kernels, + }; + time_dispatches(ctx, pool.copies, dispatches, |encoder, copy| { + kernels + .gemm( + encoder, + GemmParams { + quant, + rows: n, + cols: k, + acts: r, + weights: pool.window(copy, 0)?, + x: x_f32.window(Elems(0), Elems(r * k)).map_err(why)?, + x_stride: k, + y: y.window_mut(Elems(0), Elems(r * n)).map_err(why)?, + y_stride: n, + mode: GemmMode::Store, + operand, + fault, + layer: LAYER_NONE, + }, + ) + .map_err(|e| format!("the dispatch did not encode: {e}")) + })? + } Arm::Variant(v, operand) => { let library = libs.variant(ctx.device, v)?; let pipeline = library diff --git a/crates/metal/src/bin/metal_gemm_time/routed.rs b/crates/metal/src/bin/metal_gemm_time/routed.rs index ea4398749..215c3a9f8 100644 --- a/crates/metal/src/bin/metal_gemm_time/routed.rs +++ b/crates/metal/src/bin/metal_gemm_time/routed.rs @@ -15,8 +15,11 @@ //! pass does. //! //! The arms (`--arm NAME`, repeatable; `ship` by default): `ship`, the -//! shipped source, and every routed variant (`variants::ROUTED`), each -//! launched with the engine's bindings and argument block. `--hot` sets the +//! shipped source's simdgroup kernels, `tensor` and `tensor_relaxed`, the +//! tensor-op program's (`gemm_tensor.metal`, the second with the op's +//! relaxed precision on) over the shipped table, and every routed variant +//! (`variants::ROUTED`), each launched with the engine's bindings and +//! argument block. `--hot` sets the //! stack's per-expert stride to zero, so every group reads expert 0's rows: //! the same instructions over cache-resident weights. //! @@ -60,10 +63,11 @@ pub struct Opts { pub hot: bool, } -/// An arm: its name and its library. +/// An arm: its name, the library of its table and of its grouped GEMMs. struct RArm<'l> { name: String, library: &'l Library, + gemms: &'l Library, } /// The weights of a layer: the stacks and the shared tensors, uploaded once. @@ -385,12 +389,18 @@ fn hash(ctx: &Ctx<'_>, p: &Pass) -> Result<(u64, u64), String> { /// Time every arm at every row count and skew. pub fn run(ctx: &Ctx<'_>, opts: &Opts) -> Result<(), String> { let mut built: Vec<(&'static Variant, Library)> = Vec::new(); + let mut tensors: Vec<(&str, Library)> = Vec::new(); for name in &opts.arms { - if name != "ship" { + if name == "tensor" || name == "tensor_relaxed" { + tensors.push(( + name, + super::tensor_library(ctx.device, name == "tensor_relaxed")?, + )); + } else if name != "ship" { let v = variants::find_routed(name).ok_or_else(|| { let known: Vec<&str> = variants::ROUTED.iter().map(|v| v.name).collect(); format!( - "no routed arm named {name}: ship, or one of {}", + "no routed arm named {name}: ship, tensor, tensor_relaxed, or one of {}", known.join(", ") ) })?; @@ -407,6 +417,13 @@ pub fn run(ctx: &Ctx<'_>, opts: &Opts) -> Result<(), String> { arms.push(RArm { name: name.clone(), library: ctx.device.library(), + gemms: ctx.device.library(), + }); + } else if let Some((_, library)) = tensors.iter().find(|(n, _)| n == name) { + arms.push(RArm { + name: name.clone(), + library: ctx.device.library(), + gemms: library, }); } else { let (_, library) = built @@ -416,6 +433,7 @@ pub fn run(ctx: &Ctx<'_>, opts: &Opts) -> Result<(), String> { arms.push(RArm { name: name.clone(), library, + gemms: library, }); } } @@ -426,12 +444,9 @@ pub fn run(ctx: &Ctx<'_>, opts: &Opts) -> Result<(), String> { let ids = picks(rows, skew, 0x7AB1E ^ rows as u64); for arm in &arms { let mut p = pass(ctx, arm, rows, &ids)?; - let pair_p = arm - .library - .pipeline("gemm_routed_a4_pair_f16") - .map_err(why)?; + let pair_p = arm.gemms.pipeline("gemm_routed_a4_pair_f16").map_err(why)?; let down_p = arm - .library + .gemms .pipeline("gemm_routed_a4_plain_f16") .map_err(why)?; // One pair and one down first, so the down reads the pair's diff --git a/crates/metal/src/bin/shared/mutants.rs b/crates/metal/src/bin/shared/mutants.rs new file mode 100644 index 000000000..7b9231c17 --- /dev/null +++ b/crates/metal/src/bin/shared/mutants.rs @@ -0,0 +1,55 @@ +// The mutant support both Metal gates share: a module of each gate's +// `mutants.rs`, whose own tables `MUTANTS` and `TENSOR_MUTANTS` it walks. + +use bloomery_metal::library::{SOURCE, TENSOR_SOURCE}; + +/// The program a mutant's edits go into. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Program { + /// The kernels' library ([`SOURCE`]). + Main, + /// The tensor-op GEMMs' library ([`TENSOR_SOURCE`]), which the tensor + /// clauses force. + Tensor, +} + +/// A defect and the clause that must be red on it. +pub struct Mutant { + pub name: &'static str, + pub clause: &'static str, + pub why: &'static str, + /// `(find, replace, occurrences of find)`, applied in order. + pub edits: &'static [(&'static str, &'static str, usize)], +} + +/// Every mutant of a gate's two tables, with the program it edits: `main` +/// edits the kernels' library, `tensor` the tensor-op GEMMs'. +pub fn all( + main: &'static [Mutant], + tensor: &'static [Mutant], +) -> impl Iterator { + main.iter() + .map(|m| (m, Program::Main)) + .chain(tensor.iter().map(|m| (m, Program::Tensor))) +} + +/// The source of `program` with the mutant's edits applied; an edit whose +/// find text occurs another number of times than it names is refused. +pub fn source_of(mutant: &Mutant, program: Program) -> Result { + let mut text = match program { + Program::Main => SOURCE, + Program::Tensor => TENSOR_SOURCE, + } + .to_string(); + for (find, replace, count) in mutant.edits { + let found = text.matches(find).count(); + if found != *count { + return Err(format!( + "mutant {}: {find:?} occurs {found} times, not {count}", + mutant.name + )); + } + text = text.replace(find, replace); + } + Ok(text) +} diff --git a/crates/metal/src/kernels/mod.rs b/crates/metal/src/kernels/mod.rs index 9aff4d868..47fea7ea5 100644 --- a/crates/metal/src/kernels/mod.rs +++ b/crates/metal/src/kernels/mod.rs @@ -84,7 +84,7 @@ pub use seg::{SegParams, Segment}; use bloomery_device::DeviceError; use crate::device::MetalDevice; -use crate::library::{Library, Pipeline}; +use crate::library::{GemmPath, Library, Pipeline}; /// The pipelines of one mat-vec type: the plain entry, the entry that norms its /// input as it loads it, the one that adds into its output, and the segmented @@ -118,10 +118,14 @@ pub struct Kernels { q6k: MatvecPipelines, a4: MatvecPipelines, /// The dense GEMM of each weight type: every epilogue and operand, - /// plain and segmented (`gemm.metal`). + /// plain and segmented (`gemm.metal`, or `gemm_tensor.metal` on the + /// tensor-op path). q4k_gemm: gemm::GemmPipelines, q6k_gemm: gemm::GemmPipelines, a4_gemm: gemm::GemmPipelines, + /// The program the dense and the grouped GEMMs run: the library's + /// [`Library::gemm_path`]. + gemm_path: GemmPath, /// The gated delta net's out projection: the Q4_K entry (A4's is /// `headnorm_a4`; Q6_K has none). headnorm: Pipeline, @@ -282,9 +286,10 @@ impl Kernels { ], count, )?, - q4k_gemm: gemm::GemmPipelines::new(library, "gemm_q4k")?, - q6k_gemm: gemm::GemmPipelines::new(library, "gemm_q6k")?, - a4_gemm: gemm::GemmPipelines::new(library, "gemm_a4")?, + q4k_gemm: gemm::GemmPipelines::new(library.gemm_library(), "gemm_q4k")?, + q6k_gemm: gemm::GemmPipelines::new(library.gemm_library(), "gemm_q6k")?, + a4_gemm: gemm::GemmPipelines::new(library.gemm_library(), "gemm_a4")?, + gemm_path: library.gemm_path(), headnorm: library.pipeline_with_simdgroups("q4k_matvec_headnorm_resid", count)?, headnorm_a4: library.pipeline_with_simdgroups("a4_matvec_headnorm_resid", count)?, embed_q4k: library.pipeline("embed_q4k")?, @@ -317,6 +322,12 @@ impl Kernels { }) } + /// The program the dense and the grouped GEMMs run. + #[must_use] + pub fn gemm_path(&self) -> GemmPath { + self.gemm_path + } + /// The SIMD groups in a mat-vec's threadgroup. #[must_use] pub fn simdgroups(&self) -> usize { diff --git a/crates/metal/src/kernels/routed.rs b/crates/metal/src/kernels/routed.rs index e423a0364..33e2d9d65 100644 --- a/crates/metal/src/kernels/routed.rs +++ b/crates/metal/src/kernels/routed.rs @@ -92,11 +92,14 @@ pub(crate) struct RoutedPipelines { } impl RoutedPipelines { + /// The pipelines of `library`, the grouped GEMM's from its GEMMs' library + /// ([`Library::gemm_library`]). pub(crate) fn new(library: &Library) -> Result { + let gemm = library.gemm_library(); Ok(Self { - q4k: RoutedFamily::new(library, "q4k")?, - q6k: RoutedFamily::new(library, "q6k")?, - a4: RoutedFamily::new(library, "a4")?, + q4k: RoutedFamily::new(gemm, "q4k")?, + q6k: RoutedFamily::new(gemm, "q6k")?, + a4: RoutedFamily::new(gemm, "a4")?, f32w: library.pipeline("gemm_f32w_store_f32")?, select: library.pipeline("route_rows_select")?, table: library.pipeline("route_table")?, diff --git a/crates/metal/src/library.rs b/crates/metal/src/library.rs index 68ab67bf3..e087e1b61 100644 --- a/crates/metal/src/library.rs +++ b/crates/metal/src/library.rs @@ -13,6 +13,13 @@ //! arithmetic kernel compares them bit for bit when both exist. A machine with //! neither — no metallib and a runtime that refuses the source — is a named //! error. +//! +//! The prompt pass's GEMMs have a second program: the tensor-op GEMMs +//! (`kernels/gemm_tensor.metal`), compiled by the runtime from +//! [`TENSOR_SOURCE`] at language version 4.0 with the same math options, on +//! a GPU whose cores carry matrix units ([`MatrixUnits`]). The library holds +//! it beside its own, and the GEMMs' pipelines come from it +//! ([`Library::gemm_path`]); every other kernel stays the 3.2 library's. use core::ptr::NonNull; @@ -20,8 +27,8 @@ use objc2::runtime::ProtocolObject; use objc2_foundation::NSString; use objc2_metal::{ MTLCompileOptions, MTLComputePipelineState, MTLDataType, MTLDevice, MTLFunction, - MTLFunctionConstantValues, MTLLanguageVersion, MTLLibrary, MTLMathFloatingPointFunctions, - MTLMathMode, + MTLFunctionConstantValues, MTLGPUFamily, MTLLanguageVersion, MTLLibrary, + MTLMathFloatingPointFunctions, MTLMathMode, }; use bloomery_device::DeviceError; @@ -34,6 +41,11 @@ use crate::handle::Handle; /// into. pub const SOURCE: &str = include_str!(concat!(env!("OUT_DIR"), "/bloomery.metal")); +/// The tensor-op GEMMs' source (`kernels/gemm_tensor.metal` after the files it +/// reads), with the same table in front: the text the runtime compiles at +/// Metal 4, and the text the kernel gate edits its tensor mutants into. +pub const TENSOR_SOURCE: &str = include_str!(concat!(env!("OUT_DIR"), "/bloomery_tensor.metal")); + /// The metallib `build.rs` compiled, when the Metal Toolchain was there. #[cfg(metal_lib = "metallib")] const METALLIB: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/bloomery.metallib")); @@ -58,6 +70,86 @@ impl LibraryPath { } } +/// Which program the prompt pass's GEMMs run. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum GemmPath { + /// The `simdgroup_matrix` kernels of the library itself. + Simdgroup, + /// The Metal 4 tensor-op kernels (`matmul2d`). + TensorOps, +} + +impl GemmPath { + /// The path's name, as a record prints it. + #[must_use] + pub fn name(self) -> &'static str { + match self { + Self::Simdgroup => "simdgroup", + Self::TensorOps => "tensor", + } + } +} + +/// Whether a GPU's cores carry the matrix units Metal 4's tensor ops reach: +/// the one fact the GEMM path is chosen by. +/// +/// The rule is MLX's `is_nax_available` (mlx/backend/metal/device.cpp, +/// v0.32.1, :947–966 and the architecture parse at :588–601): the +/// architecture name ends in two generation digits and a class letter +/// (`applegpu_g13s`), and the units are there from generation 17 on, 18 on +/// for the phone class `p`. The macOS floor it also asks for, 26.2, is this +/// backend's own ([`crate::device`]). Metal 4's family is asked as well: +/// the tensor ops compile on no device without it. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct MatrixUnits { + /// `MTLDevice.architecture.name`. + pub architecture: String, + /// `supportsFamily(MTLGPUFamilyMetal4)`. + pub metal4: bool, +} + +impl MatrixUnits { + /// The facts of `device`. + pub(crate) fn of(device: &ProtocolObject) -> Self { + Self { + architecture: device.architecture().name().to_string(), + metal4: device.supportsFamily(MTLGPUFamily::Metal4), + } + } + + /// The architecture's generation and class, when its name ends in two + /// digits and a letter. + #[must_use] + pub fn generation(&self) -> Option<(u32, char)> { + let tail: Vec = self.architecture.chars().rev().take(3).collect(); + match tail.as_slice() { + [class, ones, tens] if class.is_ascii_alphabetic() => { + Some((tens.to_digit(10)? * 10 + ones.to_digit(10)?, *class)) + } + _ => None, + } + } + + /// The rule: Metal 4, and a generation that carries the units. + #[must_use] + pub fn present(&self) -> bool { + self.metal4 + && self + .generation() + .is_some_and(|(generation, class)| generation >= if class == 'p' { 18 } else { 17 }) + } + + /// The GEMM path this device takes. + #[must_use] + pub fn path(&self) -> GemmPath { + if self.present() { + GemmPath::TensorOps + } else { + GemmPath::Simdgroup + } + } +} + /// A compute pipeline and the most threads a group of it may hold. pub struct Pipeline { state: Handle, @@ -85,6 +177,8 @@ pub struct Library { path: LibraryPath, rows_per_simd_group: Option, router_sg: Option, + /// The tensor-op GEMMs' library, when the GEMMs run on them. + tensor: Option>, } /// The value of `#define NAME ` in `text`, when one such line names a @@ -103,21 +197,26 @@ fn define_of(text: &str, name: &str) -> Option { impl Library { /// The library this machine's Metal device takes: the metallib when the - /// build made one, the source otherwise. + /// build made one, the source otherwise, and the tensor-op GEMMs beside + /// it on a GPU whose cores carry matrix units. pub(crate) fn compile(device: &Handle) -> Result { #[cfg(metal_lib = "metallib")] - { - Self::metallib_of(device) - } + let library = Self::metallib_of(device)?; #[cfg(not(metal_lib = "metallib"))] - { - Self::source_of(device, SOURCE) + let library = Self::source_of(device, SOURCE, MTLLanguageVersion::Version3_2)?; + match MatrixUnits::of(device).path() { + GemmPath::Simdgroup => Ok(library), + GemmPath::TensorOps => { + let tensor = + Self::source_of(device, TENSOR_SOURCE, MTLLanguageVersion::Version4_0)?; + Ok(library.with_tensor_gemms(tensor)) + } } } /// The library the runtime compiles from the source. pub fn from_source(device: &MetalDevice) -> Result { - Self::source_of(device.handle(), SOURCE) + Self::source_of(device.handle(), SOURCE, MTLLanguageVersion::Version3_2) } /// The library the runtime compiles from the given source text. For the @@ -125,15 +224,69 @@ impl Library { /// edited into the shipped source's text) and its gate-only kernels /// appended to that text. The engine takes [`Library::from_source`]. pub fn from_source_text(device: &MetalDevice, text: &str) -> Result { - Self::source_of(device.handle(), text) + Self::source_of(device.handle(), text, MTLLanguageVersion::Version3_2) + } + + /// The tensor-op GEMMs' library the runtime compiles at Metal 4 from + /// `text` ([`TENSOR_SOURCE`], or the gate's mutant of it), on any device + /// that compiles it: the gate and the probe run it whatever + /// [`MatrixUnits`] says. A device whose runtime refuses it is a named + /// error. + pub fn tensor_from_source_text(device: &MetalDevice, text: &str) -> Result { + Self::source_of(device.handle(), text, MTLLanguageVersion::Version4_0) + } + + /// This library with its GEMMs taken from `tensor`, a library of + /// [`Library::tensor_from_source_text`]: every other kernel stays this + /// library's. + #[must_use] + pub fn with_tensor_gemms(&self, tensor: Library) -> Self { + Self { + tensor: Some(Box::new(tensor)), + ..self.clone() + } + } + + /// This library with the simdgroup GEMMs, whatever its device picked: + /// the kernel gate's clauses of those kernels run them on every device. + #[must_use] + pub fn with_simdgroup_gemms(&self) -> Self { + Self { + tensor: None, + ..self.clone() + } + } + + /// The library the GEMMs' pipelines come from. + pub(crate) fn gemm_library(&self) -> &Library { + self.tensor.as_deref().unwrap_or(self) } - fn source_of(device: &Handle, text: &str) -> Result { + /// Which program the GEMMs run. + #[must_use] + pub fn gemm_path(&self) -> GemmPath { + match self.tensor { + Some(_) => GemmPath::TensorOps, + None => GemmPath::Simdgroup, + } + } + + /// The facts the GEMM path of this library's device is chosen by. + #[must_use] + pub fn matrix_units(&self) -> MatrixUnits { + MatrixUnits::of(&self.device) + } + + fn source_of( + device: &Handle, + text: &str, + version: MTLLanguageVersion, + ) -> Result { let source = NSString::from_str(text); let options = MTLCompileOptions::new(); options.setMathMode(MTLMathMode::Safe); options.setMathFloatingPointFunctions(MTLMathFloatingPointFunctions::Precise); - options.setLanguageVersion(MTLLanguageVersion::Version3_2); + options.setLanguageVersion(version); let library = device .newLibraryWithSource_options_error(&source, Some(&options)) .map_err(|e| DeviceError::backend("library", e.localizedDescription().to_string()))?; @@ -143,6 +296,7 @@ impl Library { path: LibraryPath::Source, rows_per_simd_group: define_of(text, "BLOOMERY_MATVEC_ROWS_PER_SIMD_GROUP"), router_sg: define_of(text, "BLOOMERY_ROUTER_SG"), + tensor: None, }) } @@ -164,6 +318,7 @@ impl Library { path: LibraryPath::Metallib, rows_per_simd_group: define_of(SOURCE, "BLOOMERY_MATVEC_ROWS_PER_SIMD_GROUP"), router_sg: define_of(SOURCE, "BLOOMERY_ROUTER_SG"), + tensor: None, }) } @@ -242,3 +397,31 @@ impl Library { }) } } + +#[cfg(test)] +mod tests { + use super::*; + + fn units(architecture: &str, metal4: bool) -> MatrixUnits { + MatrixUnits { + architecture: architecture.to_string(), + metal4, + } + } + + #[test] + fn the_matrix_unit_rule_is_mlx_s() { + // The M1 Pro and the M4 Pro this backend is gated on: no units. + assert_eq!(units("applegpu_g13s", true).path(), GemmPath::Simdgroup); + assert_eq!(units("applegpu_g16s", true).path(), GemmPath::Simdgroup); + // Generation 17 carries them outside the phone class, 18 inside it. + assert_eq!(units("applegpu_g17s", true).path(), GemmPath::TensorOps); + assert_eq!(units("applegpu_g17p", true).path(), GemmPath::Simdgroup); + assert_eq!(units("applegpu_g18p", true).path(), GemmPath::TensorOps); + // No Metal 4, or a name the rule cannot read: no units. + assert_eq!(units("applegpu_g17s", false).path(), GemmPath::Simdgroup); + assert_eq!(units("applegpu", true).generation(), None); + assert_eq!(units("applegpu", true).path(), GemmPath::Simdgroup); + assert_eq!(units("applegpu_g17s", true).generation(), Some((17, 's'))); + } +}