linalg/arm64/sme: experimental f16 SME kernels for Apple M5/A19

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

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<f16>, 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