// vim: ft=arm

// x20..x27 are used, callee-preserved

// C tile regs: v8 to v31, (scratch)
// - x19-x29 to preserve (but x19, x28, x29 not used) 
// - d8..d15 to preserve
// - v16 to v31, no need to preserve

// v8 is used, d8 (lower half) must preserved
// v0-v7 (scratch registers)
//  packed A buffering (2x8 values): alternating v0, v1 with v2, v3
//  packed B buffering (2x8 values): alternating v4, v5 with v6, v7

.text
.align 4

.cpu generic+fp+simd
.global {{G}}arm64simd_mmm_f32_24x4_{{core}}_{{suffix}}
{{G}}arm64simd_mmm_f32_24x4_{{core}}_{{suffix}}:

    stp         x20, x21, [sp, #-16]!
    stp         x22, x23, [sp, #-16]!
    stp         x24, x25, [sp, #-16]!
    stp         x26, x27, [sp, #-16]!

    stp         d8, d9, [sp, #-16]!
    stp         d10, d11, [sp, #-16]!
    stp         d12, d13, [sp, #-16]!
    stp         d14, d15, [sp, #-16]!

{% include "dispatcher.tmpliq" %}

.add_mat_mul:
    ldr         x2, [x0, #24]       // b
    ldp         x3, x1, [x0, #8]    // k, a

    cmp         x3, #0
    beq         .non_linear_loop

    ld1         { v0.4s, v1.4s, v2.4s, v3.4s }, [ x1 ], #64
    ld1         { v4.4s, v5.4s }, [ x1 ], #32
    ld1         { v7.4s }, [ x2 ], #16

{% capture packed_packed_loop1 %}
    {% if core == "a53" %}
        {% include "arm64simd_mmm_f32_24x4/packed_packed_loop1/cortex_a53.tmpli" %}
    {% elsif core == "a55" %}
        {% include "arm64simd_mmm_f32_24x4/packed_packed_loop1/cortex_a55.tmpli" %}
    {% else %}
        {% include "arm64simd_mmm_f32_24x4/packed_packed_loop1/naive.tmpli" %}
    {% endif %}
{% endcapture %}

    cmp         x3, #4
    blt         .packed_packed_loop_1

.p2align 4
.packed_packed_loop_4:
    {{ packed_packed_loop1 }}
    {{ packed_packed_loop1 }}
    {{ packed_packed_loop1 }}
    {{ packed_packed_loop1 }}

    sub x3, x3, #4
    cmp x3, #4
    bge .packed_packed_loop_4

    cmp x3, #0
    beq .non_linear_loop

.p2align 4
.packed_packed_loop_1:
    {{ packed_packed_loop1 }}
    subs        x3, x3, #1
    bne .packed_packed_loop_1

    b   .non_linear_loop

{% include "arm64simd_mmm_f32_scalars.tmpliq" from:8, to:31 %}
{% include "arm64simd_mmm_f32_per_rows.tmpliq" mr:24, from:8, to:31 %}
{% include "arm64simd_mmm_f32_per_cols.tmpliq" mr:24, from:8, to:31 %}

.add_unicast:
    ldp         x5, x6, [x0, #8]
    ldp         x7, x8, [x0, #24]

    {% for col in (0..3) %}
        mov x4, x5
        {% for reg in (0..5) %}
            {% for lane in (0..3) %}
                ld1 {v0.s}[{{lane}}], [ x4 ], x6
            {% endfor %}
            fadd v{{col | times:6 | plus: 8 | plus: reg}}.4s, v{{col | times:6 | plus: 8 | plus: reg}}.4s, v0.4s
        {% endfor %}
        add x5, x5, x7
    {% endfor %}

    b           .non_linear_loop

.add_row_col_products:
    ldr     x2, [x0, #8]
    ldr     x3, [x0, #16]

    ld1         { v0.4s, v1.4s, v2.4s, v3.4s }, [ x2 ], #64
    ld1         { v7.4s }, [ x3 ]
    ld1         { v4.4s, v5.4s }, [ x2 ]

    fmla        v8.4s, v0.4s, v7.s[0]
    fmla        v9.4s, v1.4s, v7.s[0]
    fmla        v10.4s, v2.4s, v7.s[0]
    fmla        v11.4s, v3.4s, v7.s[0]
    fmla        v12.4s, v4.4s, v7.s[0]
    fmla        v13.4s, v5.4s, v7.s[0]

    fmla        v14.4s, v0.4s, v7.s[1]
    fmla        v15.4s, v1.4s, v7.s[1]
    fmla        v16.4s, v2.4s, v7.s[1]
    fmla        v17.4s, v3.4s, v7.s[1]
    fmla        v18.4s, v4.4s, v7.s[1]
    fmla        v19.4s, v5.4s, v7.s[1]

    fmla        v20.4s, v0.4s, v7.s[2]
    fmla        v21.4s, v1.4s, v7.s[2]
    fmla        v22.4s, v2.4s, v7.s[2]
    fmla        v23.4s, v3.4s, v7.s[2]
    fmla        v24.4s, v4.4s, v7.s[2]
    fmla        v25.4s, v5.4s, v7.s[2]

    fmla        v26.4s, v0.4s, v7.s[3]
    fmla        v27.4s, v1.4s, v7.s[3]
    fmla        v28.4s, v2.4s, v7.s[3]
    fmla        v29.4s, v3.4s, v7.s[3]
    fmla        v30.4s, v4.4s, v7.s[3]
    fmla        v31.4s, v5.4s, v7.s[3]

    b           .non_linear_loop

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

    cmp         x6, #4
    bne           .store_strides_generic

    {% for col in (0..3) %}
        str q{{col | times:6 | plus:8 | plus: 0}}, [ x5 ]
        str q{{col | times:6 | plus:8 | plus: 1}}, [ x5, #16 ]
        str q{{col | times:6 | plus:8 | plus: 2}}, [ x5, #32 ]
        str q{{col | times:6 | plus:8 | plus: 3}}, [ x5, #48 ]
        str q{{col | times:6 | plus:8 | plus: 4}}, [ x5, #64 ]
        str q{{col | times:6 | plus:8 | plus: 5}}, [ x5, #80 ]
        add x5, x5, x7
    {% endfor %}

    b           .non_linear_loop

.store_strides_generic:

    {% for col in (0..3) %}
        mov x4, x5
        {% for reg in (0..5) %}
            {% for lane in (0..3) %}
                st1 { v{{col | times:6 | plus:8 | plus: reg}}.s }[{{lane}}], [ x4 ], x6
            {% endfor %}
        {% endfor %}
        add x5, x5, x7
    {% endfor %}

    b           .non_linear_loop

.return:
    ldp         d14, d15, [sp], #16
    ldp         d12, d13, [sp], #16
    ldp         d10, d11, [sp], #16
    ldp         d8, d9, [sp], #16

    ldp         x26, x27, [sp], #16
    ldp         x24, x25, [sp], #16
    ldp         x22, x23, [sp], #16
    ldp         x20, x21, [sp], #16

    ret