// vim: ft=arm
.text
.align 4

/* Z: 32x1
 z0[0] ..  z0[15] z1[0] .. z1[15]
*/
    

.global {{G}}apple_amx_mmm_f32_32x1_{{suffix}}
{{G}}apple_amx_mmm_f32_32x1_{{suffix}}:

{{ AMX_SET }}

    // set x1 to a 128 bytes aligned block for loads
    mov x1, sp
    lsr x1, x1, #7
    lsl x1, x1, #7
    sub x1, x1, 128

{% include "dispatcher.j2" %}

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

.add_mat_mul:

    ldr         x2, [x0, #24]       // b
    ldp         x3, x4, [x0, #8]    // k, a

    cmp         x3, #0
    beq         .non_linear_loop

    orr x4, x4, {{ 0 | setting(62) }}  // load a pair of A

    mov x5, {{ 0 | setting(44) }}      // f32
    orr x5, x5, {{ 0 | setting(38) }}  // Broadcast Y

    orr x6, x5, {{ 0 | setting(20) }}  // z offset
    orr x6, x6, {{ 0 | setting(16) }}  // x offset

    cmp         x3, #16
    blt         .packed_packed_loop_1

    mov x9, {{ 0 | setting(32) }}        // Y broadcast offset += 1

    .packed_packed_loop_16:
        mov x7, x5
        mov x8, x6
        {{ amx("ldy", 2) }}
        {% for k in range(0, 16) %}
            {{ amx("ldx", 4) }}
            add x4, x4, 128
            {{ amx("vecfp", 7) }}
            {{ amx("vecfp", 8) }}
            add x7, x7, x9
            add x8, x8, x9
        {% endfor %}
        add x2, x2, #64
        sub x3, x3, #16
        cmp x3, #16
    bge .packed_packed_loop_16

    cmp x3, #0
    beq .non_linear_loop

    .packed_packed_loop_1:
        ldr w7, [x2], #4
        str w7, [x1]
        {{ amx("ldx", 4) }}
        {{ amx("ldy", 1) }}
        {{ amx("vecfp", 5) }}
        {{ amx("vecfp", 6) }}
        add x4, x4, 128
        subs x3, x3, #1
    bne .packed_packed_loop_1

    b .non_linear_loop

.clear:
    // top left
    eor x2, x2, x2
    orr x2, x2, {{ 0 | setting(27) }}
    orr x2, x2, {{ 0 | setting(28) }}
    orr x2, x2, {{ 0 | setting(29) }}  // Z = 0
    {{ amx("fma32", 2) }}

    // top right
    orr x2, x2, {{ 0 | setting(20) }}  // Z row = 1
    {{ amx("fma32", 2) }}

    // bottom right
    orr x2, x2, {{ 0 | setting(21) }}  // Z row = 3
    {{ amx("fma32", 2) }}

    // bottom left
    eor x2, x2, {{ 0 | setting(20) }}  // Z row = 2
    {{ amx("fma32", 2) }}

    b .non_linear_loop

.per_col_sub:

    // performs a unary neg on Z
    eor x2, x2, x2                      // X[0] = Z[0]
    // extr[hxyz] is suport confusing

    mov x4, {{ 0 | setting(63) }}          // vector mode
    orr x4, x4, {{ 0 | setting(28) }}
    orr x4, x4, {{ 0 | setting(27) }}      // Z=-X

    {{ amx("extrx", 2) }}
    {{ amx("fms32", 4) }}
    add x2, x2, {{ 0 | setting(20) }}    // next Z row
    add x4, x4, {{ 0 | setting(20) }}    // next Z row
    {{ amx("extrx", 2) }}              // extr[hxyz] is confusing
    {{ amx("fms32", 4) }}

    // continue

.per_col_add:
    ldr         x2, [x0, #8]

    // broadcast value to x0
    ld1         { v0.s }[0], [x2]
    dup         v0.4s, v0.s[0]
    st1         { v0.4s }, [x1], #16
    st1         { v0.4s }, [x1], #16
    st1         { v0.4s }, [x1], #16
    st1         { v0.4s }, [x1], #16
    sub         x1, x1, #64

    {{ amx("ldx", 1) }} // load into x0 by default

    mov x2, {{ 0 | setting(28) }}      // z += y
    {{ amx("fma32", 2) }}

    orr x2, x2, {{ 0 | setting(20) }}  // target is now z1
    {{ amx("fma32", 2) }}

    b .non_linear_loop

.per_col_sub_flipped:
    ldr         x2, [x0, #8]

    // broadcast value to x0
    ld1         { v0.s }[0], [x2]
    dup         v0.4s, v0.s[0]
    st1         { v0.4s }, [x1], #16
    st1         { v0.4s }, [x1], #16
    st1         { v0.4s }, [x1], #16
    st1         { v0.4s }, [x1], #16
    sub x1, x1, #64

    {{ amx("ldx", 1) }} // load into x0 by default

    mov x2, {{ 0 | setting(28) }}      // z += y
    {{ amx("fms32", 2) }}

    orr x2, x2, {{ 0 | setting(20) }}  // target is now z1
    {{ amx("fms32", 2) }}

    b .non_linear_loop

.per_row_sub_flipped:
    ldr         x2, [x0, #8]

    ld1         { v0.4s, v1.4s, v2.4s, v3.4s }, [x2], #64
    st1         { v0.4s, v1.4s, v2.4s, v3.4s }, [x1], #64
    ld1         { v0.4s, v1.4s, v2.4s, v3.4s }, [x2]
    st1         { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
    sub x1, x1, #64

    orr x2, x1, {{ 0 | setting(62) }}  // load a pair
    {{ amx("ldy", 2) }}

    mov x2, {{ 0 | setting(63) }}      // vector mode
    orr x2, x2, {{ 0 | setting(29) }}  // z -= y

    // top left
    {{ amx("fms32", 2) }}

    // bottom left
    orr x2, x2, {{ 0 | setting(20) }}  // Z row = 1
    orr x2, x2, {{ 0 | setting(6) }}   // Y offset
    {{ amx("fms32", 2) }}

    b .non_linear_loop

.per_row_sub:
    // performs a unary neg on Z
    eor x2, x2, x2                      // X[0] = Z[0]

    mov x4, {{ 0 | setting(63) }}          // vector mode
    orr x4, x4, {{ 0 | setting(28) }}
    orr x4, x4, {{ 0 | setting(27) }}      // Z=-X

    {{ amx("extrx", 2) }}
    {{ amx("fms32", 4) }}
    add x2, x2, {{ 0 | setting(20) }}    // next Z row
    add x4, x4, {{ 0 | setting(20) }}    // next Z row
    {{ amx("extrx", 2) }}
    {{ amx("fms32", 4) }}

    // continue

.per_row_add:
    ldr         x2, [x0, #8]

    ld1         { v0.4s, v1.4s, v2.4s, v3.4s }, [x2], #64
    st1         { v0.4s, v1.4s, v2.4s, v3.4s }, [x1], #64
    ld1         { v0.4s, v1.4s, v2.4s, v3.4s }, [x2]
    st1         { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
    sub x1, x1, #64

    orr x2, x1, {{ 0 | setting(62) }}  // load a pair
    {{ amx("ldy", 2) }}

    mov x2, {{ 0 | setting(63) }}      // vector mode
    orr x2, x2, {{ 0 | setting(29) }}  // z += y

    // top left
    {{ amx("fma32", 2) }}

    // bottom left
    orr x2, x2, {{ 0 | setting(20) }}  // Z row = 1
    orr x2, x2, {{ 0 | setting(6) }}   // Y offset
    {{ amx("fma32", 2) }}

    b .non_linear_loop

.per_row_min:
    mov x2, 5
    b .per_row_min_max
.per_row_max:
    mov x2, 7
.per_row_min_max:
    ldr         x5, [x0, #8]

    ld1         { v0.4s, v1.4s, v2.4s, v3.4s }, [x5], #64
    st1         { v0.4s, v1.4s, v2.4s, v3.4s }, [x1], #64
    ld1         { v0.4s, v1.4s, v2.4s, v3.4s }, [x5]
    st1         { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
    sub x1, x1, #64

    orr x5, x1, {{ 0 | setting(62) }}  // load a pair
    {{ amx("ldx", 5) }}

    lsl x2, x2, 47                  // max(x,z) (or min)
    orr x2, x2, {{ 0 | setting(44) }}  // f32
    {{ amx("vecfp", 2) }}

    orr x2, x2, {{ 0 | setting(16) }}  // x1
    orr x2, x2, {{ 0 | setting(20) }}  // z1
    {{ amx("vecfp", 2) }}

    b .non_linear_loop

.per_col_min:
    mov x2, 5
    b .per_col_min_max
.per_col_max:
    mov x2, 7
.per_col_min_max:
    ldr         x4, [x0, #8]

    // broadcast value to x0
    ld1         { v0.s }[0], [x4]
    dup         v0.4s, v0.s[0]
    st1         { v0.4s }, [x1], #16
    st1         { v0.4s }, [x1], #16
    st1         { v0.4s }, [x1], #16
    st1         { v0.4s }, [x1], #16
    sub         x1, x1, #64

    {{ amx("ldx", 1) }}

    lsl x2, x2, 47                  // max(x,z) (or min)
    orr x2, x2, {{ 0 | setting(44) }}  // f32

    {{ amx("vecfp", 2) }}
    orr x2, x2, {{ 0 | setting(20) }}  // z offset
    {{ amx("vecfp", 2) }}

    b .non_linear_loop

.per_col_mul:
    ldr         x4, [x0, #8]

    // broadcast value to y0
    ld1         { v0.s }[0], [x4]
    dup         v0.4s, v0.s[0]
    st1         { v0.4s }, [x1], #16
    st1         { v0.4s }, [x1], #16
    st1         { v0.4s }, [x1], #16
    st1         { v0.4s }, [x1], #16
    sub         x1, x1, #64

    {{ amx("ldy", 1) }}

    eor x2, x2, x2                      // X[0] = Z[0]
    {{ amx("extrx", 2) }}
    mov x4, {{ 0 | setting(63) }}          // vector mode
    orr x4, x4, {{ 0 | setting(27) }}      // Z=X*Y
    {{ amx("fma32", 4) }}
    orr x2, x2, {{ 0 | setting(20) }}      // Z1
    {{ amx("extrx", 2) }}
    orr x4, x4, {{ 0 | setting(20) }}      // Z1
    {{ amx("fma32", 4) }}

    b .non_linear_loop

.per_row_mul:
    ldr         x2, [x0, #8]

    ld1         { v0.4s, v1.4s, v2.4s, v3.4s }, [x2], #64
    st1         { v0.4s, v1.4s, v2.4s, v3.4s }, [x1], #64
    ld1         { v0.4s, v1.4s, v2.4s, v3.4s }, [x2]
    st1         { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
    sub x1, x1, #64

    orr x2, x1, {{ 0 | setting(62) }}      // pair
    {{ amx("ldy", 2) }}

    eor x2, x2, x2                      // X[0] = Z[0]
    {{ amx("extrx", 2) }}
    mov x4, {{ 0 | setting(63) }}          // vector mode
    orr x4, x4, {{ 0 | setting(27) }}      // Z=X*Y
    {{ amx("fma32", 4) }}
    orr x2, x2, {{ 0 | setting(20) }}      // Z1
    {{ amx("extrx", 2) }}
    orr x4, x4, {{ 0 | setting(20) }}      // Z1
    orr x4, x4, {{ 0 | setting(6) }}       // Y1
    {{ amx("fma32", 4) }}

    b .non_linear_loop

.scalar_sub:
    // performs a unary neg on Z, then go to scalar_add
    eor x2, x2, x2                      // X[0] = Z[0]

    mov x4, {{ 0 | setting(63) }}          // vector mode
    orr x4, x4, {{ 0 | setting(28) }}
    orr x4, x4, {{ 0 | setting(27) }}      // Z=-X
    {{ amx("extrx", 2) }}
    {{ amx("fms32", 4) }}
    add x2, x2, {{ 0 | setting(20) }}    // next Z row
    add x4, x4, {{ 0 | setting(20) }}    // next Z row
    {{ amx("extrx", 2) }}
    {{ amx("fms32", 4) }}

    // continue on purpose

.scalar_add:
    ldr         w5, [x0, #8]

    fmov        s0, w5
    dup         v0.4s, v0.s[0]
    dup         v1.4s, v0.s[0]
    dup         v2.4s, v0.s[0]
    dup         v3.4s, v0.s[0]

    st1         { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
    {{ amx("ldx", 1) }}    // load 16 values

    mov x2, {{ 0 | setting(28) }}          // Z+=X 
    {{ amx("fma32", 2) }}
    add x2, x2, {{ 0 | setting(20) }}    // next Z row
    {{ amx("fma32", 2) }}
    b .non_linear_loop

.scalar_sub_flipped:
    ldr         w5, [x0, #8]
    fmov        s0, w5
    dup         v0.4s, v0.s[0]
    dup         v1.4s, v0.s[0]
    dup         v2.4s, v0.s[0]
    dup         v3.4s, v0.s[0]

    st1         { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
    {{ amx("ldx", 1) }}    // load 16 values

    mov x2, {{ 0 | setting(28) }}          // Z-=X 
    {{ amx("fms32", 2) }}
    add x2, x2, {{ 0 | setting(20) }}    // next Z row
    {{ amx("fms32", 2) }}
    b .non_linear_loop

.scalar_mul:
    ldr         w5, [x0, #8]
    fmov        s0, w5
    dup         v0.4s, v0.s[0]
    dup         v1.4s, v0.s[0]
    dup         v2.4s, v0.s[0]
    dup         v3.4s, v0.s[0]

    st1         { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
    {{ amx("ldy", 1) }}    // load 16 values

    eor x2, x2, x2                      // X[0] = Z[0]

    mov x4, {{ 0 | setting(63) }}          // vector mode
    orr x4, x4, {{ 0 | setting(27) }}      // Z=X*Y

    {{ amx("extrx", 2) }}
    {{ amx("fma32", 4) }}
    add x2, x2, {{ 0 | setting(20) }}    // next Z row
    add x4, x4, {{ 0 | setting(20) }}    // next Z row
    {{ amx("extrx", 2) }}
    {{ amx("fma32", 4) }}

    b .non_linear_loop

.scalar_min:
    mov x2, 5
    b .scalar_min_max
.scalar_max:
    mov x2, 7
.scalar_min_max:
    ldr         w5, [x0, #8]
    fmov        s0, w5
    dup         v0.4s, v0.s[0]
    dup         v1.4s, v0.s[0]
    dup         v2.4s, v0.s[0]
    dup         v3.4s, v0.s[0]

    st1         { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
    {{ amx("ldx", 1) }}    // load 16 values

    lsl x2, x2, 47
    orr x2, x2, {{ 0 | setting(44) }} // f32

    {{ amx("vecfp", 2) }}
    add x2, x2, {{ 0 | setting(20) }} // next Z
    {{ amx("vecfp", 2) }}

    b .non_linear_loop

.add_unicast:
    ldp         x5, x6, [x0, #8]            // c base ptr, rsc
    ldp         x7, x8, [x0, #24]           // csc, item_size

    {% for neon in range(0, 8) %}
       {% for lane in range(0, 4) %}
           ld1 { v{{neon}}.s }[{{lane}}], [x5], x6
       {% endfor %}
    {% endfor %}
    mov x8, x1
    st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x8], #64
    st1 { v4.4s, v5.4s, v6.4s, v7.4s }, [x8], #64

    orr x8, x1, {{ 0 | setting(62) }}          // pair
    {{ amx("ldy", 8) }}

    eor x2, x2, x2
    orr x2, x2, {{ 0 | setting(63) }}  // vector mode
    orr x2, x2, {{ 0 | setting(29) }}  // perform Z0+=Y0
    {{ amx("fma32", 2) }}
    orr x2, x2, {{ 0 | setting(20) }}  // Z1
    orr x2, x2, 64                  // offset Y by 16 values
    {{ amx("fma32", 2) }}
    
    b .non_linear_loop

.add_row_col_products:
    ldp         x5, x6, [x0, #8]            // a base ptr, b base ptr

    ld1         { v0.s }[0], [x6]
    st1         { v0.s }[0], [x1]
    {{ amx("ldy", 1) }}

    ld1         { v0.4s, v1.4s, v2.4s, v3.4s }, [x5], #64
    st1         { v0.4s, v1.4s, v2.4s, v3.4s }, [x1], #64
    ld1         { v0.4s, v1.4s, v2.4s, v3.4s }, [x5]
    st1         { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
    sub x1, x1, #64

    orr x2, x1, {{ 0 | setting(62) }}  // load a pair
    {{ amx("ldx", 2) }}

    mov x2, {{ 0 | setting(44) }}      // f32
    orr x2, x2, {{ 0 | setting(38) }}  // Broadcast Y
    {{ amx("vecfp", 2) }}

    orr x2, x2, {{ 0 | setting(20) }}  // Z row = 1
    orr x2, x2, {{ 0 | setting(16) }}  // X offset
    {{ amx("vecfp", 2) }}

    b .non_linear_loop

.store:
    ldp         x5, x6, [x0, #8]            // c base ptr, rsc
    ldp         x7, x8, [x0, #24]           // csc, item_size

    ands        x8, x5, 0x7f
    bne         .store_generic
    cmp         x6, 4
    bne         .store_generic
    cmp         x7, 4
    bne         .store_generic
 
    orr x5, x5, {{ 0 | setting(62) }}          // pair
    {{ amx("stz", 5) }}
    b .non_linear_loop

 .store_generic:

    orr x8, x1, {{ 0 | setting(62) }}          // pair
    {{ amx("stz", 8) }}

    mov x8, x1
    ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x8], #64
    ld1 { v4.4s, v5.4s, v6.4s, v7.4s }, [x8], #64
    {% for neon in range(0, 8) %}
       {% for lane in range(0, 4) %}
           st1 { v{{neon}}.s }[{{lane}}], [x5], x6
       {% endfor %}
    {% endfor %}
    
    b .non_linear_loop

.load_tile:
    ldr  x2, [x0, #16]                      // row major ptr
    orr  x2, x2, {{ 0 | setting(62) }}           // load pairs
    {{ amx("ldz", 2) }}
    b .non_linear_loop

.return:
{{ AMX_CLR }}
ret
