// vim: ft=arm
//
// SME2 i32 32x32 quantized matmul kernel.
//
// ZA tile layout (4 .S tiles, 16x16 i32 each):
//   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 (K decrements by 4 per iter, since SMOPA at i8 reduces 4):
//   ld1b {z0, z1}, pn8/z, [A]    ; 32 M × 4 K = 128 i8 of A
//   ld1b {z2, z3}, pn8/z, [B]    ; 32 N × 4 K = 128 i8 of B
//   smopa za0.s, p0/m, p0/m, z0.b, z2.b   ; ZA0 += A[0..16] × B[0..16]
//   smopa za1.s, p0/m, p0/m, z0.b, z3.b
//   smopa za2.s, p0/m, p0/m, z1.b, z2.b
//   smopa za3.s, p0/m, p0/m, z1.b, z3.b
//
// SMOPA at i8 throughput: 4-way K reduction per insn × 16x16 cells = 1024
// MACs per insn. With 4-tile rotation we approach 4 SMOPAs/cycle = 4 K
// reduction × 16x16 = 4096 MACs/cycle ≈ ~16 TOPS theoretical peak.
//
// Calling convention (extern "C", AAPCS64):
//   x0 = const *FusedKerSpec<i32>, advanced 40 B per dispatcher iteration
//   x1 = 4 KiB scratch buffer for tile spills (used by store-generic / q_scale)
//
// Tract packing requirement: i8 inputs packed with K_alignment=4 (SMOPA
// requires K%4=0). The PackedFormat::with_k_alignment(4) handles this.

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

.global {{G}}sme_qmmm_i32_32x32_{{suffix}}
{{G}}sme_qmmm_i32_32x32_{{suffix}}:

    stp     q8,  q9,  [sp, #-128]!
    stp     q10, q11, [sp, #32]
    stp     q12, q13, [sp, #64]
    stp     q14, q15, [sp, #96]

    sub     sp, sp, #4096
    mov     x1, sp

    smstart
    ptrue   p0.b
    ptrue   pn8.b
    mov     w8, #0

{% include "dispatcher.j2" %}

// -------- AddMatMul: ZA += A·B at i8 with K=4 reduction per SMOPA ----------

.add_mat_mul:
    ldr     x9, [x0, #32]       // packing index
    ldr     x2, [x0, #24]       // b ptr
    ldp     x3, x4, [x0, #8]    // k, a ptr
    cmp     x3, #0
    b.eq    .non_linear_loop
    cmp     x9, #1
    b.eq    .Lmatmul_loop
// i32i32 fallback (packing != 1, auto-test path): ZA += A[:,k] (x) B[k,:], one
// K-step at a time via predicated MLA rank-1 updates. One instruction per line:
// the Apple/LLVM AArch64 assembler treats `;` as a COMMENT, so semicolon-packed
// statements silently drop everything after the first `;`.
.Lk32:
    ld1w    {z2.s}, p0/z, [x2]               // B[k, 0..16]
    ld1w    {z3.s}, p0/z, [x2, #1, mul vl]   // B[k, 16..32]
    mov     w12, #0
.Lkt:
    ldr     w10, [x4, w12, uxtw #2]          // A[k, w12]
    dup     z4.s, w10
    mov     z16.s, p0/m, za0h.s[w12, 0]
    mov     z17.s, p0/m, za1h.s[w12, 0]
    mla     z16.s, p0/m, z2.s, z4.s          // C[w12, 0..16]  += A[w12]   * B[0..16]
    mla     z17.s, p0/m, z3.s, z4.s          // C[w12, 16..32] += A[w12]   * B[16..32]
    mov     za0h.s[w12, 0], p0/m, z16.s
    mov     za1h.s[w12, 0], p0/m, z17.s
    add     w10, w12, #16
    ldr     w10, [x4, w10, uxtw #2]          // A[k, w12+16]
    dup     z4.s, w10
    mov     z18.s, p0/m, za2h.s[w12, 0]
    mov     z19.s, p0/m, za3h.s[w12, 0]
    mla     z18.s, p0/m, z2.s, z4.s          // C[w12+16, 0..16]  += A[w12+16] * B[0..16]
    mla     z19.s, p0/m, z3.s, z4.s          // C[w12+16, 16..32] += A[w12+16] * B[16..32]
    mov     za2h.s[w12, 0], p0/m, z18.s
    mov     za3h.s[w12, 0], p0/m, z19.s
    add     w12, w12, #1
    cmp     w12, #16
    b.lt    .Lkt
    add     x4, x4, #128
    add     x2, x2, #128
    subs    x3, x3, #1
    b.ne    .Lk32
    b       .non_linear_loop

.Lmatmul_loop:
    ld1b    {z0.b, z1.b}, pn8/z, [x4]
    ld1b    {z2.b, z3.b}, pn8/z, [x2]
    add     x4, x4, #128
    add     x2, x2, #128
    smopa   za0.s, p0/m, p0/m, z0.b, z2.b
    smopa   za1.s, p0/m, p0/m, z0.b, z3.b
    smopa   za2.s, p0/m, p0/m, z1.b, z2.b
    smopa   za3.s, p0/m, p0/m, z1.b, z3.b
    subs    x3, x3, #4
    b.gt    .Lmatmul_loop
    b       .non_linear_loop

.clear:
    zero    {za}
    b       .non_linear_loop

// -------- Store: i32 tile -> memory (port of Phase 1 f32 store) -----------

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

    cmp     x7, #4
    b.ne    .Lstore_generic
    cmp     x9, #4
    b.ne    .Lstore_generic

    add     x4, x5, #64
    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:
    mov     x13, x9             // preserve item_size before x9 is reused as a ptr
    mov     x4, x1
    add     x9, x1, #64
    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

    mov     x3, #0
.Lstore_row:
    mov     x4, x5
    mov     x10, #0
    lsl     x9, x3, #7
    add     x11, x1, x9
.Lstore_col:
    ldr     w9, [x11], #4
    cmp     x13, #1             // item_size: 1 -> strb, 2 -> strh, else (4) -> str
    b.eq    .Lstore_b1
    cmp     x13, #2
    b.eq    .Lstore_b2
    str     w9, [x4]
    b       .Lstore_cnext
.Lstore_b1:
    strb    w9, [x4]
    b       .Lstore_cnext
.Lstore_b2:
    strh    w9, [x4]
.Lstore_cnext:
    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

// -------- LoadTile: ZA := row-major i32 tile from memory -------------------

.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

// -------- AddUnicast: ZA += C (strided load + add) ------------------------

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

    cmp     x7, #4
    b.ne    .Laddu_generic
    cmp     x9, #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]
    add     z6.s, p0/m, z6.s, z8.s
    add     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]
    add     z6.s, p0/m, z6.s, z8.s
    add     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:
    // Strided gather to scratch, then contig accumulate (mirrors Phase 1).
    mov     x3, #0
    mov     x10, x1
.Laddu_gather_row:
    mov     x11, x5
    mov     x4, #0
.Laddu_gather_col:
    ldr     w9, [x11]
    str     w9, [x10], #4
    add     x11, x11, x7
    add     x4, x4, #1
    cmp     x4, #32
    b.lt    .Laddu_gather_col
    add     x5, x5, x6
    add     x3, x3, #1
    cmp     x3, #32
    b.lt    .Laddu_gather_row

    mov     x4, x1
    add     x9, x1, #64
    mov     w12, #0
.Laddu_apply_top:
    ld1w    {z8.s}, p0/z, [x4]
    ld1w    {z10.s}, p0/z, [x9]
    mov     z6.s, p0/m, za0h.s[w12, 0]
    mov     z7.s, p0/m, za1h.s[w12, 0]
    add     z6.s, p0/m, z6.s, z8.s
    add     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     x4, x4, #128
    add     x9, x9, #128
    add     w12, w12, #1
    cmp     w12, #16
    b.lt    .Laddu_apply_top
    mov     w12, #0
.Laddu_apply_bot:
    ld1w    {z8.s}, p0/z, [x4]
    ld1w    {z10.s}, p0/z, [x9]
    mov     z6.s, p0/m, za2h.s[w12, 0]
    mov     z7.s, p0/m, za3h.s[w12, 0]
    add     z6.s, p0/m, z6.s, z8.s
    add     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     x4, x4, #128
    add     x9, x9, #128
    add     w12, w12, #1
    cmp     w12, #16
    b.lt    .Laddu_apply_bot
    b       .non_linear_loop

// -------- AddRowColProducts: ZA += rows ⊗ cols (i32 outer product) --------
//
// rows: 32 i32 (broadcast per M-row), cols: 32 i32 (lane vector per N-col).
// Per ZA row, we need: ZA[i, j] += rows[i] * cols[j]. Slice-by-slice.

.add_row_col_products:
    ldp     x2, x3, [x0, #8]    // rows ptr, cols ptr
    ld1w    {z4.s}, p0/z, [x3]            // cols[0..16]
    ld1w    {z5.s}, p0/z, [x3, #1, mul vl] // cols[16..32]

    // Top 16 rows
    mov     w12, #0
.Larcp_top:
    ldr     w9, [x2], #4
    dup     z16.s, w9                      // broadcast rows[i] to z16
    mov     z6.s, p0/m, za0h.s[w12, 0]
    mov     z7.s, p0/m, za1h.s[w12, 0]
    mla     z6.s, p0/m, z16.s, z4.s        // z6 += z16 * cols[0..16]
    mla     z7.s, p0/m, z16.s, z5.s        // z7 += z16 * cols[16..32]
    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    .Larcp_top
    // Bottom 16 rows
    mov     w12, #0
.Larcp_bot:
    ldr     w9, [x2], #4
    dup     z16.s, w9
    mov     z6.s, p0/m, za2h.s[w12, 0]
    mov     z7.s, p0/m, za3h.s[w12, 0]
    mla     z6.s, p0/m, z16.s, z4.s
    mla     z7.s, p0/m, z16.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    .Larcp_bot
    b       .non_linear_loop

// -------- scalar fuse ops: broadcast scalar, apply lane-wise --------------
//
// Sub vs SubF (matches Phase 1's f32 convention):
//   ScalarSub   → result = scalar - z   (mnemonic: subr)
//   ScalarSubF  → result = z - scalar   (mnemonic: sub)

{% macro scalar_op_i32(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_i32('.scalar_add', 'add') }}
{{ scalar_op_i32('.scalar_mul', 'mul') }}
{{ scalar_op_i32('.scalar_sub', 'subr') }}
{{ scalar_op_i32('.scalar_sub_flipped', 'sub') }}
{{ scalar_op_i32('.scalar_min', 'smin') }}
{{ scalar_op_i32('.scalar_max', 'smax') }}

// -------- per_col fuse ops: 32-elem vector, broadcast across M rows ------

{% macro per_col_op_i32(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_i32('.per_col_add', 'add') }}
{{ per_col_op_i32('.per_col_mul', 'mul') }}
{{ per_col_op_i32('.per_col_sub', 'subr') }}
{{ per_col_op_i32('.per_col_sub_flipped', 'sub') }}
{{ per_col_op_i32('.per_col_min', 'smin') }}
{{ per_col_op_i32('.per_col_max', 'smax') }}

// -------- per_row fuse ops: 32-elem vector, one scalar per M row ---------

{% macro per_row_op_i32(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_i32('.per_row_add', 'add') }}
{{ per_row_op_i32('.per_row_mul', 'mul') }}
{{ per_row_op_i32('.per_row_sub', 'subr') }}
{{ per_row_op_i32('.per_row_sub_flipped', 'sub') }}
{{ per_row_op_i32('.per_row_min', 'smin') }}
{{ per_row_op_i32('.per_row_max', 'smax') }}

// -------- Quantization fuse ops (bit-exact port of generic/rounding.rs) ----
//
// Strategy: spill the 32x32 i32 ZA tile to the 4 KiB scratch (x1), quantize
// element-wise in SCALAR GP registers (streaming-mode legal: smull/lsr/asr/
// cneg/cset/... are base A64 and unaffected by PSTATE.SM), then reload to ZA.
// Quant is not the hot path; this mirrors the scalar approach already proven
// in arm64/sve/sve_mmm_i32.c. Everything is inlined (no `bl` — a nested call
// would clobber x30 and corrupt the final `ret`).
//
// Bit-exactness: the reference forms the FULL i64 product (mult*v) and does a
// single magnitude-rounding shift by (shift+31) with a per-policy nudge. A
// vector sqdmulh+srshl truncates the low 31 bits before the second shift, so
// it is NOT equivalent — hence the i64 scalar port.
//
// RoundingPolicy: Native=0 Zero=1 Away=2 MinusInf=3 PlusInf=4 Even=5 Odd=6.

// Spill ZA0..ZA3 -> scratch[x1] as a contiguous 32x32 row-major i32 matrix
// (same layout the generic store path uses). Clobbers x4, x9, w12.
{% macro za_spill(sfx) %}
    mov     x4, x1
    add     x9, x1, #64
    mov     w12, #0
.Lspt_{{sfx}}:
    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    .Lspt_{{sfx}}
    mov     w12, #0
.Lspb_{{sfx}}:
    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    .Lspb_{{sfx}}
{% endmacro %}

// Reload scratch[x1] (32x32 row-major i32) -> ZA0..ZA3. Clobbers x4,x9,w12,z6,z7.
{% macro za_reload(sfx) %}
    mov     x4, x1
    add     x9, x1, #64
    mov     w12, #0
.Lrlt_{{sfx}}:
    ld1w    {z6.s}, p0/z, [x4]
    ld1w    {z7.s}, p0/z, [x9]
    mov     za0h.s[w12, 0], p0/m, z6.s
    mov     za1h.s[w12, 0], p0/m, z7.s
    add     x4, x4, #128
    add     x9, x9, #128
    add     w12, w12, #1
    cmp     w12, #16
    b.lt    .Lrlt_{{sfx}}
    mov     w12, #0
.Lrlb_{{sfx}}:
    ld1w    {z6.s}, p0/z, [x4]
    ld1w    {z7.s}, p0/z, [x9]
    mov     za2h.s[w12, 0], p0/m, z6.s
    mov     za3h.s[w12, 0], p0/m, z7.s
    add     x4, x4, #128
    add     x9, x9, #128
    add     w12, w12, #1
    cmp     w12, #16
    b.lt    .Lrlb_{{sfx}}
{% endmacro %}

// Magnitude-rounding shared by q_scale and q_shr (mirrors `Mul<i32> for Scaler`
// / `i32::q_shr`). In: x12 = val (i64), x5 = shift, x6 = policy. Out: w14 (i32).
// Clobbers x13,x15,x16,x17. Preserves x5,x6,x7,x10,x11,x12.
{% macro round_mag(sfx) %}
    cmp     x5, #0
    b.gt    .Lrpos_{{sfx}}
    neg     x13, x5
    lsl     x13, x12, x13                   // val << (-shift)
    mov     w14, w13
    b       .Lrend_{{sfx}}
.Lrpos_{{sfx}}:
    cmp     x12, #0
    cneg    x15, x12, mi                    // x15 = |val|
    sub     x13, x5, #1
    mov     x16, #1
    lsl     x16, x16, x13                   // x16 = half = 1 << (shift-1)
    cmp     x6, #2                          // Away -> nudge 0
    b.eq    .Lrn0_{{sfx}}
    cmp     x6, #1                          // Zero -> nudge -1
    b.ne    .Lrna_{{sfx}}
    mov     x17, #-1
    b       .Lrnd_{{sfx}}
.Lrna_{{sfx}}:
    cmp     x6, #3                          // MinusInf -> -(val >= 0)
    b.ne    .Lrnb_{{sfx}}
    cmp     x12, #0
    cset    x17, ge
    neg     x17, x17
    b       .Lrnd_{{sfx}}
.Lrnb_{{sfx}}:
    cmp     x6, #4                          // PlusInf -> -(val <= 0)
    b.ne    .Lrnc_{{sfx}}
    cmp     x12, #0
    cset    x17, le
    neg     x17, x17
    b       .Lrnd_{{sfx}}
.Lrnc_{{sfx}}:
    cmp     x6, #5                          // Even -> ((|val|>>shift)&1) - 1
    b.ne    .Lrno_{{sfx}}
    lsr     x17, x15, x5
    and     x17, x17, #1
    sub     x17, x17, #1
    b       .Lrnd_{{sfx}}
.Lrno_{{sfx}}:                              // Odd -> -((|val|>>shift)&1)
    lsr     x17, x15, x5
    and     x17, x17, #1
    neg     x17, x17
    b       .Lrnd_{{sfx}}
.Lrn0_{{sfx}}:
    mov     x17, #0
.Lrnd_{{sfx}}:
    add     x15, x15, x16
    add     x15, x15, x17
    lsr     x15, x15, x5                    // (|val| + half + nudge) >> shift
    cmp     x12, #0
    cneg    x14, x15, mi                    // signum(val) * mag
.Lrend_{{sfx}}:
{% endmacro %}

// QScale(shift, policy, mult): val = mult*v (i64); shift += 31; magnitude round.
.q_scale:
    ldr     x5, [x0, #8]                    // shift (isize)
    ldr     x6, [x0, #16]                   // policy
    ldr     w7, [x0, #24]                   // mult (i32)
    add     x5, x5, #31
    {{ za_spill('qsc') }}
    mov     x10, x1
    mov     x11, #1024
.Lqsc_loop:
    ldr     w9, [x10]
    smull   x12, w7, w9                     // val = (i64)mult * (i64)v
    {{ round_mag('qsc') }}
    str     w14, [x10], #4
    subs    x11, x11, #1
    b.ne    .Lqsc_loop
    {{ za_reload('qsc') }}
    b       .non_linear_loop

// RoundingShiftRight(shift, policy): val = v (i64); magnitude round (shift>0).
.q_shr:
    ldr     x5, [x0, #8]                    // shift (usize, >= 1)
    ldr     x6, [x0, #16]                   // policy
    {{ za_spill('qsr') }}
    mov     x10, x1
    mov     x11, #1024
.Lqsr_loop:
    ldr     w9, [x10]
    sxtw    x12, w9                         // val = (i64)v
    {{ round_mag('qsr') }}
    str     w14, [x10], #4
    subs    x11, x11, #1
    b.ne    .Lqsr_loop
    {{ za_reload('qsr') }}
    b       .non_linear_loop

// ShiftLeft(shift): result = v << shift (32-bit wrapping, matches i32::q_shl).
.q_shl:
    ldr     x5, [x0, #8]                    // shift (usize)
    {{ za_spill('qsl') }}
    mov     x10, x1
    mov     x11, #1024
.Lqsl_loop:
    ldr     w9, [x10]
    lsl     w9, w9, w5
    str     w9, [x10], #4
    subs    x11, x11, #1
    b.ne    .Lqsl_loop
    {{ za_reload('qsl') }}
    b       .non_linear_loop

// -------- LeakyRelu (excluded via CAN_FUSE_I32) ---------------------------

.leaky_relu:
    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
