From b1278a96952923d54c19ef67ac95b178da0dd4c2 Mon Sep 17 00:00:00 2001 From: czoli1976 <64466170+czoli1976@users.noreply.github.com> Date: Fri, 22 May 2026 13:37:04 +0100 Subject: [PATCH] linalg/arm64/sme: EXPERIMENTAL f16 SME kernels for Apple M5/A19 (FEAT_SME_F16F16) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds half-precision SME kernels for the mmm_f16 / mmv_f16 slots, the f16 companions to the merged f32 SME backend (#2230). They use the non-widening half-precision SME path (`fmopa za.h`, FEAT_SME_F16F16): f16 inputs, f16 accumulate in ZA.H, consuming tract's native K-major f16 packing directly. - sme_mmm_f16_32x32 (GEMM): one 32x32 ZA.H tile, `fmopa za.h` per K-step (at SVL=512, one f16 FMOPA covers the whole 32x32 tile). where(SME_F16F16). - sme_mmv_f16_64x1 (GEMV, N==1): vgx2 ZA.H group, SME2 multi-vec `fmla za.h[w8,0,vgx2]`. where(SME2 && SME_F16F16). EXPERIMENTAL / effectively Apple M5 + A19 only, and unvalidated on real hardware. FEAT_SME_F16F16 is an optional SME2 feature that, among shipping silicon, only the Apple M5 / A19 implement. The "smoking gun" is upstream LLVM's CPU definition (llvm/llvm-project commit f85494f6afeb, "Define apple-m5/a19"): def TuneAppleM5 : SubtargetFeature<"apple-m5", ..., FeatureSME, FeatureSME2, FeatureSMEF64F64, FeatureSMEI16I64, FeatureSME2p1, FeatureSMEB16B16, FeatureSMEF16F16, ...> In LLVM's whole AArch64 CPU table FeatureSMEF16F16 appears for apple-m5/a19 and nothing else: the Apple M4/A18 report it as 0 (verified by sysctl on an M4), and the newest non-Apple SME2 cores (Arm C1/Lumex in Exynos 2600, Cortex-X925, Qualcomm Oryon) have SME/SME2 but not FEAT_SME_F16F16. So this needs community testing on an actual M5 / A19 (iPhone 17) — the maintainers' M4 cannot exercise it (it falls back to AMX f16). Build gating (so the f32 SME backend is never regressed): the f16 unit uses the `+sme-f16f16` assembler extension. A new dummy_sme_f16f16.S probe + assembler_supports_sme_f16f16() compiles the f16 kernels as a separate object gated on the `tract_sme_f16f16` cfg; the f32 SME kernels keep building on toolchains that have base SME but not f16f16. The Rust f16 registrations, detection (HWCAP2_SME_F16F16 bit 42 on Linux / sysctl on macOS, plus the existing 512-bit SVL check), and plug() wiring are all behind that cfg. Validated under QEMU only (no SME_F16F16 hardware available): with `qemu-aarch64 -cpu max,sme512=on` the full SME auto-test surface passes 220/220 (the two new f16 kernels: matmul proptest, every fuse op, store layouts, frame; plus the existing f32 GEMM/GEMV, no regression). Builds clean on macOS (Apple clang) and debian:sid (gcc 15). Co-Authored-By: Claude Opus 4.7 (1M context) --- linalg/arm64/sme/dummy_sme_f16f16.S | 16 ++ linalg/arm64/sme/sme_mmm_f16_32x32.S.j2 | 291 ++++++++++++++++++++++++ linalg/arm64/sme/sme_mmv_f16_64x1.S.j2 | 220 ++++++++++++++++++ linalg/build.rs | 29 ++- linalg/src/arm64/sme.rs | 96 ++++++++ 5 files changed, 651 insertions(+), 1 deletion(-) create mode 100644 linalg/arm64/sme/dummy_sme_f16f16.S create mode 100644 linalg/arm64/sme/sme_mmm_f16_32x32.S.j2 create mode 100644 linalg/arm64/sme/sme_mmv_f16_64x1.S.j2 diff --git a/linalg/arm64/sme/dummy_sme_f16f16.S b/linalg/arm64/sme/dummy_sme_f16f16.S new file mode 100644 index 0000000000..12be2b8d67 --- /dev/null +++ b/linalg/arm64/sme/dummy_sme_f16f16.S @@ -0,0 +1,16 @@ +// Build-time capability probe for FEAT_SME_F16F16, used by build.rs +// (assembler_supports_sme_f16f16). The experimental f16 SME kernels use the +// non-widening half-precision outer product `fmopa za.h`, which needs the +// +sme-f16f16 assembler extension. An assembler that supports base SME but not +// FEAT_SME_F16F16 (so the basic dummy_sme.S probe passes but this one fails) +// must still build the f32 SME kernels — so this is a SEPARATE probe gating +// only the f16 unit + the `tract_sme_f16f16` cfg. Not linked into anything. +.arch armv9-a+sme2+sme-f16f16 +.text +.globl tract_sme_f16f16_probe +tract_sme_f16f16_probe: + smstart + ptrue p0.h + fmopa za0.h, p0/m, p0/m, z0.h, z1.h + smstop + ret diff --git a/linalg/arm64/sme/sme_mmm_f16_32x32.S.j2 b/linalg/arm64/sme/sme_mmm_f16_32x32.S.j2 new file mode 100644 index 0000000000..00505d433b --- /dev/null +++ b/linalg/arm64/sme/sme_mmm_f16_32x32.S.j2 @@ -0,0 +1,291 @@ +// vim: ft=arm +// +// SME f16 32x32 matmul kernel (FEAT_SME_F16F16). +// +// Uses the non-widening half-precision outer product `fmopa za.h` (f16 inputs, +// f16 accumulate), which mirrors the f32 32x32 kernel's structure exactly but +// in a single ZA.H tile: at SVL=512 a ZA.H tile is 32x32 f16, so one FMOPA +// covers the whole MR=32 x NR=32 output (1024 MACs/insn = the f32 kernel's +// 4-tile MAC count). K-major packing is consumed directly (one ld1h = 32 f16 = +// one K-step of A or B), so no custom packer is needed. +// +// ZA0.H : C[0..32, 0..32] (the entire 32x32 f16 tile) +// +// Calling convention (extern "C", AAPCS64): +// x0 = const *FusedKerSpec, advanced 40 B per dispatcher iteration. +// x1 = stack-resident 2 KiB scratch buffer for tile spills (strided store). +// +// Streaming mode: PSTATE.SM=1 from prologue smstart to epilogue smstop. +// v8..v15 saved/restored across the streaming region per AAPCS. + +.arch armv9-a+sme2+sme-f16f16 +.text +.align 4 + +.global {{G}}sme_mmm_f16_32x32_{{suffix}} +{{G}}sme_mmm_f16_32x32_{{suffix}}: + + stp q8, q9, [sp, #-128]! + stp q10, q11, [sp, #32] + stp q12, q13, [sp, #64] + stp q14, q15, [sp, #96] + + // 2 KiB tile-spill scratch (32x32 f16 = 2048 B), kept live across the call. + sub sp, sp, #2048 + mov x1, sp + + smstart + ptrue p0.h + +{% include "dispatcher.j2" %} + +// -------- supported fuse ops --------------------------------------------- + +.add_mat_mul: + ldr x2, [x0, #24] // b + ldp x3, x4, [x0, #8] // k, a + + cmp x3, #0 + b.eq .non_linear_loop + +.Lmatmul_loop: + ld1h {z0.h}, p0/z, [x4] // 32 f16 of A column (one K-step) + ld1h {z1.h}, p0/z, [x2] // 32 f16 of B row + add x4, x4, #64 // advance 32 f16 = 64 B + add x2, x2, #64 + fmopa za0.h, p0/m, p0/m, z0.h, z1.h + subs x3, x3, #1 + b.ne .Lmatmul_loop + b .non_linear_loop + +.clear: + zero {za} + b .non_linear_loop + +.store: + // FusedKerSpec::Store(OutputStoreKer { ptr, row_byte_stride, + // col_byte_stride, item_size }) + ldp x5, x6, [x0, #8] // ptr, row_byte_stride + ldp x7, x8, [x0, #24] // col_byte_stride, item_size + + // Fast path: contiguous f16 columns (col_stride == 2, item_size == 2). + cmp x7, #2 + b.ne .Lstore_generic + cmp x8, #2 + b.ne .Lstore_generic + + mov w12, #0 +.Lstore_fast: + st1h {za0h.h[w12, 0]}, p0, [x5] + add x5, x5, x6 + add w12, w12, #1 + cmp w12, #32 + b.lt .Lstore_fast + b .non_linear_loop + +.Lstore_generic: + // Spill ZA0.H → scratch (32x32 f16, 64 B/row), then per-element scatter. + mov x4, x1 + mov w12, #0 +.Lstore_spill: + st1h {za0h.h[w12, 0]}, p0, [x4] + add x4, x4, #64 + add w12, w12, #1 + cmp w12, #32 + b.lt .Lstore_spill + + mov x3, #0 +.Lstore_row: + mov x4, x5 + mov x10, #0 + lsl x9, x3, #6 // row*64 byte offset in scratch + add x11, x1, x9 +.Lstore_col: + ldrh w9, [x11], #2 + strh w9, [x4] + add x4, x4, x7 + add x10, x10, #1 + cmp x10, #32 + b.lt .Lstore_col + add x5, x5, x6 + add x3, x3, #1 + cmp x3, #32 + b.lt .Lstore_row + b .non_linear_loop + +// -------- scalar ops ------------------------------------------------------ +// +// ScalarSub → result = scalar - z (fsubr) +// ScalarSubF → result = z - scalar (fsub) + +{% macro scalar_op(label, op) %} +{{label}}: + ldrh w2, [x0, #8] + dup z4.h, w2 + mov w12, #0 +.L{{label|replace('.', '')}}_loop: + mov z6.h, p0/m, za0h.h[w12, 0] + {{op}} z6.h, p0/m, z6.h, z4.h + mov za0h.h[w12, 0], p0/m, z6.h + add w12, w12, #1 + cmp w12, #32 + b.lt .L{{label|replace('.', '')}}_loop + b .non_linear_loop +{% endmacro %} + +{{ scalar_op('.scalar_add', 'fadd') }} +{{ scalar_op('.scalar_mul', 'fmul') }} +{{ scalar_op('.scalar_sub', 'fsubr') }} +{{ scalar_op('.scalar_sub_flipped', 'fsub') }} +{{ scalar_op('.scalar_min', 'fmin') }} +{{ scalar_op('.scalar_max', 'fmax') }} + +// -------- per-col ops ----------------------------------------------------- +// +// 32-element column vector → z4.h, applied to every row. + +{% macro per_col_op(label, op) %} +{{label}}: + ldr x2, [x0, #8] + ld1h {z4.h}, p0/z, [x2] + mov w12, #0 +.L{{label|replace('.', '')}}_loop: + mov z6.h, p0/m, za0h.h[w12, 0] + {{op}} z6.h, p0/m, z6.h, z4.h + mov za0h.h[w12, 0], p0/m, z6.h + add w12, w12, #1 + cmp w12, #32 + b.lt .L{{label|replace('.', '')}}_loop + b .non_linear_loop +{% endmacro %} + +{{ per_col_op('.per_col_add', 'fadd') }} +{{ per_col_op('.per_col_mul', 'fmul') }} +{{ per_col_op('.per_col_sub', 'fsubr') }} +{{ per_col_op('.per_col_sub_flipped', 'fsub') }} +{{ per_col_op('.per_col_min', 'fmin') }} +{{ per_col_op('.per_col_max', 'fmax') }} + +// -------- per-row ops ----------------------------------------------------- +// +// 32-element row vector at x2; broadcast one f16 per slice (row). + +{% macro per_row_op(label, op) %} +{{label}}: + ldr x2, [x0, #8] + mov w12, #0 +.L{{label|replace('.', '')}}_loop: + ldrh w4, [x2], #2 + dup z4.h, w4 + mov z6.h, p0/m, za0h.h[w12, 0] + {{op}} z6.h, p0/m, z6.h, z4.h + mov za0h.h[w12, 0], p0/m, z6.h + add w12, w12, #1 + cmp w12, #32 + b.lt .L{{label|replace('.', '')}}_loop + b .non_linear_loop +{% endmacro %} + +{{ per_row_op('.per_row_add', 'fadd') }} +{{ per_row_op('.per_row_mul', 'fmul') }} +{{ per_row_op('.per_row_sub', 'fsubr') }} +{{ per_row_op('.per_row_sub_flipped', 'fsub') }} +{{ per_row_op('.per_row_min', 'fmin') }} +{{ per_row_op('.per_row_max', 'fmax') }} + +// -------- AddRowColProducts: ZA += rows ⊗ cols (rank-1 K=1) --------------- + +.add_row_col_products: + ldp x2, x3, [x0, #8] // rows ptr, cols ptr + ld1h {z0.h}, p0/z, [x2] + ld1h {z1.h}, p0/z, [x3] + fmopa za0.h, p0/m, p0/m, z0.h, z1.h + b .non_linear_loop + +// -------- AddUnicast: ZA += C[i][j] from strided buffer ------------------- + +.add_unicast: + ldp x5, x6, [x0, #8] // ptr, row_byte_stride + ldp x7, x8, [x0, #24] // col_byte_stride, item_size + + cmp x7, #2 + b.ne .Laddu_generic + cmp x8, #2 + b.ne .Laddu_generic + + mov w12, #0 +.Laddu_fast: + ld1h {z8.h}, p0/z, [x5] + mov z6.h, p0/m, za0h.h[w12, 0] + fadd z6.h, p0/m, z6.h, z8.h + mov za0h.h[w12, 0], p0/m, z6.h + add x5, x5, x6 + add w12, w12, #1 + cmp w12, #32 + b.lt .Laddu_fast + b .non_linear_loop + +.Laddu_generic: + // Strided gather → scratch (32x32 f16), then contiguous accumulate. + mov x3, #0 + mov x9, x1 +.Laddu_gen_row: + mov x10, #0 + mov x11, x5 +.Laddu_gen_col: + ldrh w4, [x11] + strh w4, [x9], #2 + add x11, x11, x7 + add x10, x10, #1 + cmp x10, #32 + b.lt .Laddu_gen_col + add x5, x5, x6 + add x3, x3, #1 + cmp x3, #32 + b.lt .Laddu_gen_row + + mov x9, x1 + mov w12, #0 +.Laddu_gen_apply: + ld1h {z8.h}, p0/z, [x9] + mov z6.h, p0/m, za0h.h[w12, 0] + fadd z6.h, p0/m, z6.h, z8.h + mov za0h.h[w12, 0], p0/m, z6.h + add x9, x9, #64 + add w12, w12, #1 + cmp w12, #32 + b.lt .Laddu_gen_apply + b .non_linear_loop + +// -------- LoadTile: ZA := row-major 32x32 f16 tile (64 B/row) ------------- + +.load_tile: + ldr x2, [x0, #16] // row-major ptr + mov w12, #0 +.Lloadtile_loop: + ld1h {z6.h}, p0/z, [x2] + mov za0h.h[w12, 0], p0/m, z6.h + add x2, x2, #64 + add w12, w12, #1 + cmp w12, #32 + b.lt .Lloadtile_loop + b .non_linear_loop + +// -------- not implemented for f16 ----------------------------------------- + +.leaky_relu: +.q_scale: +.q_shl: +.q_shr: + b .unsupported + +// -------- epilogue -------------------------------------------------------- + +.return: + smstop + add sp, sp, #2048 + ldp q14, q15, [sp, #96] + ldp q12, q13, [sp, #64] + ldp q10, q11, [sp, #32] + ldp q8, q9, [sp], #128 + ret diff --git a/linalg/arm64/sme/sme_mmv_f16_64x1.S.j2 b/linalg/arm64/sme/sme_mmv_f16_64x1.S.j2 new file mode 100644 index 0000000000..706617bfe6 --- /dev/null +++ b/linalg/arm64/sme/sme_mmv_f16_64x1.S.j2 @@ -0,0 +1,220 @@ +// vim: ft=arm +// +// SME2 f16 64x1 GEMV kernel (FEAT_SME2 + FEAT_SME_F16F16). +// +// Mirrors the f32 64x1 GEMV but in half precision: the 64-element f16 output +// column maps to a vgx2 group of 2 ZA.H slices (2 x 32 f16 at SVL=512). The +// inner K-step loads 64 f16 of A's column via one multi-vec LD1H, broadcasts +// B[k] with LD1RH, and issues one multi-vec vgx2 FMLA-into-ZA.H. f16 accumulate. +// +// Calling convention (extern "C", AAPCS64): +// x0 = const *FusedKerSpec, advanced 40 B per dispatcher iteration. +// x1 = stack-resident 128 B scratch buffer (strided-store / AddUnicast). +// w8 = 0 throughout (vgx-group base index). + +.arch armv9-a+sme2+sme-f16f16 +.text +.align 4 + +.global {{G}}sme_mmv_f16_64x1_{{suffix}} +{{G}}sme_mmv_f16_64x1_{{suffix}}: + + stp q8, q9, [sp, #-128]! + stp q10, q11, [sp, #32] + stp q12, q13, [sp, #64] + stp q14, q15, [sp, #96] + + // 128 B = 64 f16 spill buffer. + sub sp, sp, #128 + mov x1, sp + + smstart + ptrue p0.h + ptrue pn8.b + mov w8, #0 + +{% include "dispatcher.j2" %} + +// -------- supported fuse ops ----------------------------------------------- + +.add_mat_mul: + ldr x2, [x0, #24] // b ptr + ldp x3, x4, [x0, #8] // k, a ptr + cmp x3, #0 + b.eq .non_linear_loop + +.Lmmv_loop: + ld1h {z0.h-z1.h}, pn8/z, [x4] + add x4, x4, #128 // 64 f16 = 128 B + ld1rh {z4.h}, p0/z, [x2] + add x2, x2, #2 + fmla za.h[w8, 0, vgx2], {z0.h-z1.h}, z4.h[0] + subs x3, x3, #1 + b.ne .Lmmv_loop + b .non_linear_loop + +.clear: + zero {za} + b .non_linear_loop + +.store: + // FusedKerSpec::Store(OutputStoreKer { ptr, row_byte_stride, + // col_byte_stride, item_size }) + ldp x5, x6, [x0, #8] // ptr, row_byte_stride + ldp x7, x9, [x0, #24] // col_byte_stride, item_size + + // Fast path: 64 outputs contiguous (row_byte_stride==2 AND item_size==2). + cmp x6, #2 + b.ne .Lstore_generic + cmp x9, #2 + b.ne .Lstore_generic + + mov {z0.h-z1.h}, za.h[w8, 0, vgx2] + st1h {z0.h}, p0, [x5] + st1h {z1.h}, p0, [x5, #1, mul vl] + b .non_linear_loop + +.Lstore_generic: + mov {z0.h-z1.h}, za.h[w8, 0, vgx2] + st1h {z0.h}, p0, [x1] + st1h {z1.h}, p0, [x1, #1, mul vl] + + mov x3, #0 + mov x9, x1 +.Lstore_scatter: + ldrh w10, [x9], #2 + strh w10, [x5] + add x5, x5, x6 + add x3, x3, #1 + cmp x3, #64 + b.lt .Lstore_scatter + b .non_linear_loop + +// -------- LoadTile: ZA := 64-element f16 vec from row-major source --------- + +.load_tile: + ldr x2, [x0, #16] + ld1h {z0.h-z1.h}, pn8/z, [x2] + mov za.h[w8, 0, vgx2], {z0.h-z1.h} + b .non_linear_loop + +// -------- AddRowColProducts: ZA += rows ⊗ cols (rank-1 K=1) --------------- + +.add_row_col_products: + ldp x2, x3, [x0, #8] // rows ptr, cols ptr + ld1h {z0.h-z1.h}, pn8/z, [x2] + ld1rh {z4.h}, p0/z, [x3] + fmla za.h[w8, 0, vgx2], {z0.h-z1.h}, z4.h[0] + b .non_linear_loop + +// -------- AddUnicast: ZA += C from strided buffer -------------------------- + +.add_unicast: + ldp x5, x6, [x0, #8] // ptr, row_byte_stride + ldp x7, x9, [x0, #24] // col_byte_stride, item_size + + cmp x6, #2 + b.ne .Laddu_generic + cmp x9, #2 + b.ne .Laddu_generic + + ld1h {z16.h-z17.h}, pn8/z, [x5] + mov {z0.h-z1.h}, za.h[w8, 0, vgx2] + fadd z0.h, p0/m, z0.h, z16.h + fadd z1.h, p0/m, z1.h, z17.h + mov za.h[w8, 0, vgx2], {z0.h-z1.h} + b .non_linear_loop + +.Laddu_generic: + mov x3, #0 + mov x9, x1 +.Laddu_gather: + ldrh w10, [x5] + strh w10, [x9], #2 + add x5, x5, x6 + add x3, x3, #1 + cmp x3, #64 + b.lt .Laddu_gather + ld1h {z16.h-z17.h}, pn8/z, [x1] + mov {z0.h-z1.h}, za.h[w8, 0, vgx2] + fadd z0.h, p0/m, z0.h, z16.h + fadd z1.h, p0/m, z1.h, z17.h + mov za.h[w8, 0, vgx2], {z0.h-z1.h} + b .non_linear_loop + +// -------- scalar / per_col ops (degenerate at NR=1) ------------------------ + +{% macro scalar_op(label, op) %} +{{label}}: + ldrh w2, [x0, #8] + dup z4.h, w2 + mov {z0.h-z1.h}, za.h[w8, 0, vgx2] + {{op}} z0.h, p0/m, z0.h, z4.h + {{op}} z1.h, p0/m, z1.h, z4.h + mov za.h[w8, 0, vgx2], {z0.h-z1.h} + b .non_linear_loop +{% endmacro %} + +{{ scalar_op('.scalar_add', 'fadd') }} +{{ scalar_op('.scalar_mul', 'fmul') }} +{{ scalar_op('.scalar_sub', 'fsubr') }} +{{ scalar_op('.scalar_sub_flipped', 'fsub') }} +{{ scalar_op('.scalar_min', 'fmin') }} +{{ scalar_op('.scalar_max', 'fmax') }} + +{% macro per_col_op(label, op) %} +{{label}}: + ldr x2, [x0, #8] + ld1rh {z4.h}, p0/z, [x2] + mov {z0.h-z1.h}, za.h[w8, 0, vgx2] + {{op}} z0.h, p0/m, z0.h, z4.h + {{op}} z1.h, p0/m, z1.h, z4.h + mov za.h[w8, 0, vgx2], {z0.h-z1.h} + b .non_linear_loop +{% endmacro %} + +{{ per_col_op('.per_col_add', 'fadd') }} +{{ per_col_op('.per_col_mul', 'fmul') }} +{{ per_col_op('.per_col_sub', 'fsubr') }} +{{ per_col_op('.per_col_sub_flipped', 'fsub') }} +{{ per_col_op('.per_col_min', 'fmin') }} +{{ per_col_op('.per_col_max', 'fmax') }} + +// -------- per_row ops: 64-element f16 bias, lane-wise ----------------------- + +{% macro per_row_op(label, op) %} +{{label}}: + ldr x2, [x0, #8] + ld1h {z16.h-z17.h}, pn8/z, [x2] + mov {z0.h-z1.h}, za.h[w8, 0, vgx2] + {{op}} z0.h, p0/m, z0.h, z16.h + {{op}} z1.h, p0/m, z1.h, z17.h + mov za.h[w8, 0, vgx2], {z0.h-z1.h} + b .non_linear_loop +{% endmacro %} + +{{ per_row_op('.per_row_add', 'fadd') }} +{{ per_row_op('.per_row_mul', 'fmul') }} +{{ per_row_op('.per_row_sub', 'fsubr') }} +{{ per_row_op('.per_row_sub_flipped', 'fsub') }} +{{ per_row_op('.per_row_min', 'fmin') }} +{{ per_row_op('.per_row_max', 'fmax') }} + +// -------- not implemented -------------------------------------------------- + +.leaky_relu: +.q_scale: +.q_shl: +.q_shr: + b .unsupported + +// -------- epilogue --------------------------------------------------------- + +.return: + smstop + add sp, sp, #128 + ldp q14, q15, [sp, #96] + ldp q12, q13, [sp, #64] + ldp q10, q11, [sp, #32] + ldp q8, q9, [sp], #128 + ret diff --git a/linalg/build.rs b/linalg/build.rs index 51e8ceecbf..b680a8dc6a 100644 --- a/linalg/build.rs +++ b/linalg/build.rs @@ -37,6 +37,21 @@ fn assembler_supports_sme() -> bool { .is_ok() } +// Probe whether the assembler additionally supports FEAT_SME_F16F16 (the +// non-widening `fmopa za.h`). This is a SEPARATE probe from assembler_supports_sme +// so that a toolchain with base SME but not +sme-f16f16 still builds the f32 SME +// kernels — only the experimental f16 unit + the `tract_sme_f16f16` cfg are gated +// off. (FEAT_SME_F16F16 is, as of 2025-26 silicon, effectively Apple M5/A19-only.) +fn assembler_supports_sme_f16f16() -> bool { + cc::Build::new() + .file("arm64/sme/dummy_sme_f16f16.S") + .cargo_metadata(false) + .cargo_warnings(false) + .warnings(false) + .try_compile("tract_sme_f16f16_probe") + .is_ok() +} + fn jump_table() -> Vec { println!("cargo:rerun-if-changed=src/frame/mmm/fuse.rs"); std::fs::read_to_string("src/frame/mmm/fuse.rs") @@ -104,6 +119,8 @@ fn main() { // `tract_sme` is set below only when both include_sme() and the assembler // SME probe succeed; declare it so rustc's unexpected-cfg lint stays quiet. println!("cargo:rustc-check-cfg=cfg(tract_sme)"); + // Set below only when the f16 SME unit compiles (assembler has +sme-f16f16). + println!("cargo:rustc-check-cfg=cfg(tract_sme_f16f16)"); match arch.as_ref() { "x86_64" => { @@ -184,8 +201,18 @@ fn main() { } if include_sme() && assembler_supports_sme() { let files = preprocess_files("arm64/sme", &[], &suffix, false); - cc::Build::new().files(files).compile("sme"); + // The experimental f16 kernels use `fmopa za.h` (+sme-f16f16), + // which an assembler may lack even when base SME assembles. + // Build them as a separate unit gated on a dedicated probe so the + // f32 SME kernels still build on f16f16-less toolchains. + let (f16_files, base_files): (Vec<_>, Vec<_>) = + files.into_iter().partition(|f| f.to_string_lossy().contains("f16")); + cc::Build::new().files(base_files).compile("sme"); println!("cargo:rustc-cfg=tract_sme"); + if !f16_files.is_empty() && assembler_supports_sme_f16f16() { + cc::Build::new().files(f16_files).compile("sme_f16f16"); + println!("cargo:rustc-cfg=tract_sme_f16f16"); + } } if std::env::var("CARGO_FEATURE_NO_FP16").is_err() { let config = diff --git a/linalg/src/arm64/sme.rs b/linalg/src/arm64/sme.rs index 5b8b194d6c..c39a9569e4 100644 --- a/linalg/src/arm64/sme.rs +++ b/linalg/src/arm64/sme.rs @@ -1,6 +1,11 @@ use crate::Ops; use crate::frame::mmm::ImplementationQuality::ManuallyOptimized; use crate::mmm::*; +// rustc has a builtin primitive `f16` that a `use crate::mmm::*` glob does not +// shadow; without this explicit import `` resolves to the primitive (which +// is not a LADatum) and the f16 kernel registration fails to compile. +#[cfg(tract_sme_f16f16)] +use tract_data::prelude::f16; // CAN_FUSE: everything except LeakyRelu / QScale / RoundingShiftRight / // ShiftLeft. LoadTile, AddUnicast, AddRowColProducts, per-row/col/scalar @@ -18,6 +23,12 @@ const CAN_FUSE: fn(&FusedSpec) -> bool = |f| { const SME: fn() -> bool = has_sme; const SME2: fn() -> bool = has_sme2; +#[cfg(tract_sme_f16f16)] +const SME_F16F16: fn() -> bool = has_sme_f16f16; +// The f16 GEMV uses an SME2 multi-vec FMLA into a ZA.H group, so it needs both +// FEAT_SME2 and FEAT_SME_F16F16. +#[cfg(tract_sme_f16f16)] +const SME2_F16F16: fn() -> bool = || has_sme2() && has_sme_f16f16(); // Streaming vector length in bytes, read via `RDSVL x0, #1` (encoding // 0x04bf5820). RDSVL is legal in non-streaming mode, but is UNDEFINED @@ -65,6 +76,27 @@ MMMExternKernel!( quality(ManuallyOptimized) ); +// f16 GEMM via the non-widening half-precision outer product `fmopa za.h` +// (FEAT_SME_F16F16): one 32x32 ZA.H tile, f16 accumulate, gated on +// FEAT_SME_F16F16 + the 512-bit SVL geometry. Compiled only when the assembler +// supports +sme-f16f16 (the `tract_sme_f16f16` cfg from build.rs). +#[cfg(tract_sme_f16f16)] +MMMExternKernel!( + sme_mmm_f16_32x32(32, 32)@(128, 128) + where(SME_F16F16) + can_fuse(CAN_FUSE) + quality(ManuallyOptimized) +); + +// f16 GEMV (N==1) via the SME2 multi-vec FMLA into a vgx2 ZA.H group. +#[cfg(tract_sme_f16f16)] +MMMExternKernel!( + sme_mmv_f16_64x1(64, 1)@(128, 128) + where(SME2_F16F16) + can_fuse(CAN_FUSE) + quality(ManuallyOptimized) +); + #[cfg(target_os = "macos")] pub fn has_sme() -> bool { // TRACT_SME_DISABLE=1 forces the SME path off so callers can A/B @@ -179,6 +211,58 @@ pub fn has_sme2() -> bool { false } +#[cfg(all(target_os = "macos", tract_sme_f16f16))] +pub fn has_sme_f16f16() -> bool { + if std::env::var_os("TRACT_SME_DISABLE").is_some() { + return false; + } + use std::ffi::{CString, c_char, c_int, c_void}; + use std::ptr::null_mut; + unsafe extern "C" { + fn sysctlbyname( + name: *const c_char, + oldp: *mut c_void, + oldlenp: *mut usize, + newp: *mut c_void, + newlen: usize, + ) -> c_int; + } + let Ok(name) = CString::new("hw.optional.arm.FEAT_SME_F16F16") else { + return false; + }; + let mut value: u64 = 0; + let mut len: usize = std::mem::size_of::(); + unsafe { + if sysctlbyname(name.as_ptr(), &mut value as *mut _ as *mut c_void, &mut len, null_mut(), 0) + != 0 + { + return false; + } + } + // FEAT_SME_F16F16 present AND the 512-bit streaming geometry our f16 kernel + // assumes. (M4 may lack the sysctl, in which case we fall back to AMX f16.) + value != 0 && sme_geometry_supported() +} + +#[cfg(all(target_os = "linux", tract_sme_f16f16))] +pub fn has_sme_f16f16() -> bool { + // HWCAP2_SME_F16F16 = 1 << 42 on aarch64 (kernel ABI). + const HWCAP2_SME_F16F16: u64 = 1 << 42; + unsafe extern "C" { + fn getauxval(t: u64) -> u64; + } + const AT_HWCAP2: u64 = 26; + let feat = unsafe { (getauxval(AT_HWCAP2) & HWCAP2_SME_F16F16) != 0 }; + // FEAT_SME_F16F16 present AND the 512-bit streaming geometry (rejects + // qemu-user unless run with -cpu max,sme512=on). + feat && sme_geometry_supported() +} + +#[cfg(all(not(any(target_os = "macos", target_os = "linux")), tract_sme_f16f16))] +pub fn has_sme_f16f16() -> bool { + false +} + pub fn plug(ops: &mut Ops) { if has_sme() { log::info!("SME optimisation activated"); @@ -190,6 +274,18 @@ pub fn plug(ops: &mut Ops) { ops.mmv_f32 = Box::new(|_, _| sme_mmv_f32_64x1.mmm()); ops.mmm_impls.extend_from_slice(&[sme_mmv_f32_64x1.mmm()]); } + #[cfg(tract_sme_f16f16)] + if has_sme_f16f16() { + log::info!("SME f16 optimisation activated (experimental, M5/A19)"); + ops.mmm_f16 = Box::new(|_, _, _| sme_mmm_f16_32x32.mmm()); + ops.mmm_impls.extend_from_slice(&[sme_mmm_f16_32x32.mmm()]); + } + #[cfg(tract_sme_f16f16)] + if has_sme2() && has_sme_f16f16() { + log::info!("SME2 f16 GEMV optimisation activated (experimental, M5/A19)"); + ops.mmv_f16 = Box::new(|_, _| sme_mmv_f16_64x1.mmm()); + ops.mmm_impls.extend_from_slice(&[sme_mmv_f16_64x1.mmm()]); + } if !has_sme() && !has_sme2() { log::info!("No SME optimisation"); }