// vim: ft=arm
//
// SME f32 32x32 matmul kernel.
//
// ZA tile layout (4 .S tiles, 16x16 each, indexed left/right x top/bottom):
//
//   ZA0.S : C[0..16,  0..16]   (top-left)
//   ZA1.S : C[0..16,  16..32]  (top-right)
//   ZA2.S : C[16..32, 0..16]   (bottom-left)
//   ZA3.S : C[16..32, 16..32]  (bottom-right)
//
// Inner K-step: load 32 f32 of A (split z0+z2) and 32 of B (split z1+z3),
// issue 4 FMOPAs (one per tile). All 4 tiles are independent → SME unit
// reaches 1 fmopa/cycle = ~2 TFLOPS on M4.
//
// Calling convention (extern "C", AAPCS64):
//   x0 = const *FusedKerSpec<f32>, advanced 40 B per dispatcher iteration.
//   x1 = stack-resident 4 KiB scratch buffer for tile spills (Phase 1B+).
//
// Streaming mode: PSTATE.SM=1 from prologue smstart to epilogue smstop.
// V0..V31 (low 128 bits = Z0..Z31 low) are destroyed by the smstart/smstop
// pair; v8..v15 are saved/restored to stack across the streaming region per
// AAPCS callee-save rules.

.arch armv9-a+sme2
.text
.align 4

.global {{G}}sme_mmm_f32_32x32_{{suffix}}
{{G}}sme_mmm_f32_32x32_{{suffix}}:

    // Save callee-saved q8..q15 (AAPCS preserves low 64 bits of v8..v15;
    // we save the full 128-bit Q to keep the stack layout simple).
    stp     q8,  q9,  [sp, #-128]!
    stp     q10, q11, [sp, #32]
    stp     q12, q13, [sp, #64]
    stp     q14, q15, [sp, #96]

    // Allocate 4 KiB tile-spill scratch (kept live across the whole call).
    sub     sp, sp, #4096
    mov     x1, sp

    smstart
    ptrue   p0.b

{% 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:
    ld1w    {z0.s}, p0/z, [x4]
    ld1w    {z2.s}, p0/z, [x4, #1, mul vl]
    ld1w    {z1.s}, p0/z, [x2]
    ld1w    {z3.s}, p0/z, [x2, #1, mul vl]
    add     x4, x4, #128
    add     x2, x2, #128

    fmopa   za0.s, p0/m, p0/m, z0.s, z1.s    // C[0..16,  0..16]
    fmopa   za1.s, p0/m, p0/m, z0.s, z3.s    // C[0..16,  16..32]
    fmopa   za2.s, p0/m, p0/m, z2.s, z1.s    // C[16..32, 0..16]
    fmopa   za3.s, p0/m, p0/m, z2.s, z3.s    // C[16..32, 16..32]

    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 })
    //   [x0, #8]  = ptr           [x0, #16] = row_byte_stride
    //   [x0, #24] = col_byte_stride  [x0, #32] = item_size
    ldp     x5, x6, [x0, #8]    // x5 = ptr, x6 = row_byte_stride
    ldp     x7, x8, [x0, #24]   // x7 = col_byte_stride, x8 = item_size

    // Fast path: contiguous f32 columns (col_stride == 4) → direct ZA→user.
    // st1w-from-ZA does not accept "[Xn, #imm, MUL VL]" offsets, so we keep
    // two parallel base pointers for the left and right halves of each row.
    cmp     x7, #4
    b.ne    .Lstore_generic
    cmp     x8, #4
    b.ne    .Lstore_generic

    add     x4, x5, #64         // right-half base
    mov     w12, #0
.Lstore_top:
    st1w    {za0h.s[w12, 0]}, p0, [x5]
    st1w    {za1h.s[w12, 0]}, p0, [x4]
    add     x5, x5, x6
    add     x4, x4, x6
    add     w12, w12, #1
    cmp     w12, #16
    b.lt    .Lstore_top
    mov     w12, #0
.Lstore_bot:
    st1w    {za2h.s[w12, 0]}, p0, [x5]
    st1w    {za3h.s[w12, 0]}, p0, [x4]
    add     x5, x5, x6
    add     x4, x4, x6
    add     w12, w12, #1
    cmp     w12, #16
    b.lt    .Lstore_bot
    b       .non_linear_loop

.Lstore_generic:
    // Slow path: spill ZA → x1 (scratch buffer, 32x32 row-major, 128 B/row)
    // using two parallel pointers, then per-element strided scatter.
    mov     x4, x1              // left-half pointer
    add     x9, x1, #64         // right-half pointer
    mov     w12, #0
.Lstore_spill_top:
    st1w    {za0h.s[w12, 0]}, p0, [x4]
    st1w    {za1h.s[w12, 0]}, p0, [x9]
    add     x4, x4, #128
    add     x9, x9, #128
    add     w12, w12, #1
    cmp     w12, #16
    b.lt    .Lstore_spill_top
    mov     w12, #0
.Lstore_spill_bot:
    st1w    {za2h.s[w12, 0]}, p0, [x4]
    st1w    {za3h.s[w12, 0]}, p0, [x9]
    add     x4, x4, #128
    add     x9, x9, #128
    add     w12, w12, #1
    cmp     w12, #16
    b.lt    .Lstore_spill_bot

    // Strided f32 scatter: 32 rows × 32 cols.
    mov     x3, #0
.Lstore_row:
    mov     x4, x5
    mov     x10, #0
    lsl     x9, x3, #7          // row*128 byte offset in scratch
    add     x11, x1, x9
.Lstore_col:
    ldr     w9, [x11], #4
    str     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 ------------------------------------------------------
//
// FusedKerSpec::Scalar{Add,Mul,Sub,SubF,Min,Max}(TI) — broadcast scalar from
// [x0, #8] to all lanes, apply elementwise across the 4-tile 32x32 grid.
//
// Sub vs SubF semantics (matching apple_amx + tests/fuse.rs):
//   ScalarSub   → result = scalar - z   (mnemonic fsubr)
//   ScalarSubF  → result = z - scalar   (mnemonic fsub)
//
// Slice-op loop: for each slice index w12, extract ZA tile slice → Z reg,
// op with broadcast-scalar in z4, insert Z back. Two halves × 4 tiles total.

{% macro scalar_op(label, op) %}
{{label}}:
    ldr     w2, [x0, #8]
    dup     z4.s, w2
    mov     w12, #0
.L{{label|replace('.', '')}}_top:
    mov     z6.s, p0/m, za0h.s[w12, 0]
    mov     z7.s, p0/m, za1h.s[w12, 0]
    {{op}}  z6.s, p0/m, z6.s, z4.s
    {{op}}  z7.s, p0/m, z7.s, z4.s
    mov     za0h.s[w12, 0], p0/m, z6.s
    mov     za1h.s[w12, 0], p0/m, z7.s
    add     w12, w12, #1
    cmp     w12, #16
    b.lt    .L{{label|replace('.', '')}}_top
    mov     w12, #0
.L{{label|replace('.', '')}}_bot:
    mov     z6.s, p0/m, za2h.s[w12, 0]
    mov     z7.s, p0/m, za3h.s[w12, 0]
    {{op}}  z6.s, p0/m, z6.s, z4.s
    {{op}}  z7.s, p0/m, z7.s, z4.s
    mov     za2h.s[w12, 0], p0/m, z6.s
    mov     za3h.s[w12, 0], p0/m, z7.s
    add     w12, w12, #1
    cmp     w12, #16
    b.lt    .L{{label|replace('.', '')}}_bot
    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 (cols 0-15) + z5 (cols 16-31).
// Same z4/z5 is applied to every row across the 4-tile grid.

{% macro per_col_op(label, op) %}
{{label}}:
    ldr     x2, [x0, #8]
    ld1w    {z4.s}, p0/z, [x2]
    ld1w    {z5.s}, p0/z, [x2, #1, mul vl]
    mov     w12, #0
.L{{label|replace('.', '')}}_top:
    mov     z6.s, p0/m, za0h.s[w12, 0]
    mov     z7.s, p0/m, za1h.s[w12, 0]
    {{op}}  z6.s, p0/m, z6.s, z4.s
    {{op}}  z7.s, p0/m, z7.s, z5.s
    mov     za0h.s[w12, 0], p0/m, z6.s
    mov     za1h.s[w12, 0], p0/m, z7.s
    add     w12, w12, #1
    cmp     w12, #16
    b.lt    .L{{label|replace('.', '')}}_top
    mov     w12, #0
.L{{label|replace('.', '')}}_bot:
    mov     z6.s, p0/m, za2h.s[w12, 0]
    mov     z7.s, p0/m, za3h.s[w12, 0]
    {{op}}  z6.s, p0/m, z6.s, z4.s
    {{op}}  z7.s, p0/m, z7.s, z5.s
    mov     za2h.s[w12, 0], p0/m, z6.s
    mov     za3h.s[w12, 0], p0/m, z7.s
    add     w12, w12, #1
    cmp     w12, #16
    b.lt    .L{{label|replace('.', '')}}_bot
    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 (top 16 rows) and x2+64 (bottom 16 rows).
// Load one f32 per iteration and broadcast (no SVE indexed-broadcast for
// arbitrary i across 16 lanes, so we just walk the bias pointer).

{% macro per_row_op(label, op) %}
{{label}}:
    ldr     x2, [x0, #8]
    add     x3, x2, #64
    mov     w12, #0
.L{{label|replace('.', '')}}_top:
    ldr     w4, [x2], #4
    dup     z4.s, w4
    mov     z6.s, p0/m, za0h.s[w12, 0]
    mov     z7.s, p0/m, za1h.s[w12, 0]
    {{op}}  z6.s, p0/m, z6.s, z4.s
    {{op}}  z7.s, p0/m, z7.s, z4.s
    mov     za0h.s[w12, 0], p0/m, z6.s
    mov     za1h.s[w12, 0], p0/m, z7.s
    add     w12, w12, #1
    cmp     w12, #16
    b.lt    .L{{label|replace('.', '')}}_top
    mov     w12, #0
.L{{label|replace('.', '')}}_bot:
    ldr     w4, [x3], #4
    dup     z4.s, w4
    mov     z6.s, p0/m, za2h.s[w12, 0]
    mov     z7.s, p0/m, za3h.s[w12, 0]
    {{op}}  z6.s, p0/m, z6.s, z4.s
    {{op}}  z7.s, p0/m, z7.s, z4.s
    mov     za2h.s[w12, 0], p0/m, z6.s
    mov     za3h.s[w12, 0], p0/m, z7.s
    add     w12, w12, #1
    cmp     w12, #16
    b.lt    .L{{label|replace('.', '')}}_bot
    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 update) ------------
//
// Same shape as a K=1 matmul step: load 32 f32 of rows + 32 of cols, four
// FMOPAs into the 2x2 ZA grid.

.add_row_col_products:
    ldp     x2, x3, [x0, #8]            // rows ptr, cols ptr
    ld1w    {z0.s}, p0/z, [x2]
    ld1w    {z2.s}, p0/z, [x2, #1, mul vl]
    ld1w    {z1.s}, p0/z, [x3]
    ld1w    {z3.s}, p0/z, [x3, #1, mul vl]
    fmopa   za0.s, p0/m, p0/m, z0.s, z1.s
    fmopa   za1.s, p0/m, p0/m, z0.s, z3.s
    fmopa   za2.s, p0/m, p0/m, z2.s, z1.s
    fmopa   za3.s, p0/m, p0/m, z2.s, z3.s
    b       .non_linear_loop

// -------- AddUnicast: ZA += C[i][j] from strided buffer -------------------
//
// FusedKerSpec::AddUnicast(OutputStoreKer { ptr, row_byte_stride,
//                                           col_byte_stride, item_size })
// Fast path: contiguous f32 cols (col_stride == 4) — load each row via
// ld1w then in-place fadd to ZA slice.

.add_unicast:
    ldp     x5, x6, [x0, #8]    // ptr, row_byte_stride
    ldp     x7, x8, [x0, #24]   // col_byte_stride, item_size

    cmp     x7, #4
    b.ne    .Laddu_generic
    cmp     x8, #4
    b.ne    .Laddu_generic

    add     x4, x5, #64
    mov     w12, #0
.Laddu_top:
    ld1w    {z8.s}, p0/z, [x5]
    ld1w    {z9.s}, p0/z, [x4]
    mov     z6.s, p0/m, za0h.s[w12, 0]
    mov     z7.s, p0/m, za1h.s[w12, 0]
    fadd    z6.s, p0/m, z6.s, z8.s
    fadd    z7.s, p0/m, z7.s, z9.s
    mov     za0h.s[w12, 0], p0/m, z6.s
    mov     za1h.s[w12, 0], p0/m, z7.s
    add     x5, x5, x6
    add     x4, x4, x6
    add     w12, w12, #1
    cmp     w12, #16
    b.lt    .Laddu_top
    mov     w12, #0
.Laddu_bot:
    ld1w    {z8.s}, p0/z, [x5]
    ld1w    {z9.s}, p0/z, [x4]
    mov     z6.s, p0/m, za2h.s[w12, 0]
    mov     z7.s, p0/m, za3h.s[w12, 0]
    fadd    z6.s, p0/m, z6.s, z8.s
    fadd    z7.s, p0/m, z7.s, z9.s
    mov     za2h.s[w12, 0], p0/m, z6.s
    mov     za3h.s[w12, 0], p0/m, z7.s
    add     x5, x5, x6
    add     x4, x4, x6
    add     w12, w12, #1
    cmp     w12, #16
    b.lt    .Laddu_bot
    b       .non_linear_loop

.Laddu_generic:
    // Generic strided gather: walk 32 rows × 32 cols, accumulate one
    // element at a time into the spill scratch, then re-load each slice
    // into ZA and fadd.
    //
    // Phase 1B keeps this slow but correct — it triggers only for non-
    // contiguous AddUnicast which auto-tests don't exercise.
    mov     x3, #0                  // row idx
    mov     x9, x1                  // scratch ptr
.Laddu_gen_row:
    mov     x10, #0
    mov     x11, x5
.Laddu_gen_col:
    ldr     w4, [x11]
    str     w4, [x9], #4
    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

    // Now scratch holds 32x32 f32 row-major. Same loop as fast path but
    // reading from scratch (contiguous).
    mov     x9, x1
    add     x4, x9, #64
    mov     w12, #0
.Laddu_gen_apply_top:
    ld1w    {z8.s}, p0/z, [x9]
    ld1w    {z10.s}, p0/z, [x4]
    mov     z6.s, p0/m, za0h.s[w12, 0]
    mov     z7.s, p0/m, za1h.s[w12, 0]
    fadd    z6.s, p0/m, z6.s, z8.s
    fadd    z7.s, p0/m, z7.s, z10.s
    mov     za0h.s[w12, 0], p0/m, z6.s
    mov     za1h.s[w12, 0], p0/m, z7.s
    add     x9, x9, #128
    add     x4, x4, #128
    add     w12, w12, #1
    cmp     w12, #16
    b.lt    .Laddu_gen_apply_top
    mov     w12, #0
.Laddu_gen_apply_bot:
    ld1w    {z8.s}, p0/z, [x9]
    ld1w    {z10.s}, p0/z, [x4]
    mov     z6.s, p0/m, za2h.s[w12, 0]
    mov     z7.s, p0/m, za3h.s[w12, 0]
    fadd    z6.s, p0/m, z6.s, z8.s
    fadd    z7.s, p0/m, z7.s, z10.s
    mov     za2h.s[w12, 0], p0/m, z6.s
    mov     za3h.s[w12, 0], p0/m, z7.s
    add     x9, x9, #128
    add     x4, x4, #128
    add     w12, w12, #1
    cmp     w12, #16
    b.lt    .Laddu_gen_apply_bot
    b       .non_linear_loop

// -------- LoadTile: ZA := row-major tile from memory ----------------------
//
// FusedKerSpec::LoadTile(col_major_ptr, row_major_ptr):
//   [x0, #8]  = col-major ptr (unused; AMX prefers this for its layout)
//   [x0, #16] = row-major ptr (32x32 f32, 128 B per row)
//
// We use the row-major pointer because the ZA H-tile store path is itself
// row-major and matches naturally.

.load_tile:
    ldr     x2, [x0, #16]
    add     x4, x2, #64
    mov     w12, #0
.Lloadtile_top:
    ld1w    {z6.s}, p0/z, [x2]
    ld1w    {z7.s}, p0/z, [x4]
    mov     za0h.s[w12, 0], p0/m, z6.s
    mov     za1h.s[w12, 0], p0/m, z7.s
    add     x2, x2, #128
    add     x4, x4, #128
    add     w12, w12, #1
    cmp     w12, #16
    b.lt    .Lloadtile_top
    mov     w12, #0
.Lloadtile_bot:
    ld1w    {z6.s}, p0/z, [x2]
    ld1w    {z7.s}, p0/z, [x4]
    mov     za2h.s[w12, 0], p0/m, z6.s
    mov     za3h.s[w12, 0], p0/m, z7.s
    add     x2, x2, #128
    add     x4, x4, #128
    add     w12, w12, #1
    cmp     w12, #16
    b.lt    .Lloadtile_bot
    b       .non_linear_loop

// -------- still not implemented (low priority for Phase 1) ----------------

.leaky_relu:
.q_scale:
.q_shl:
.q_shr:
    b       .unsupported

// -------- epilogue --------------------------------------------------------

.return:
    smstop
    add     sp, sp, #4096
    ldp     q14, q15, [sp, #96]
    ldp     q12, q13, [sp, #64]
    ldp     q10, q11, [sp, #32]
    ldp     q8,  q9,  [sp], #128
    ret
