// vim: ft=arm

// C tile regs: v16 to v31, no need to preserve

// no preservation either for v0-v7...
// v8..v15 are callee-preserved
// 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

.global {{G}}arm64simd_mmm_f32_32x3_{{core}}_{{suffix}}
{{G}}arm64simd_mmm_f32_32x3_{{core}}_{{suffix}}:

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

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

{% include "dispatcher.j2" %}

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

    cmp         x3, #0
    beq         .non_linear_loop

    cmp         x4, #1
    beq         .f32f16
    cmp         x4, #2
    beq         .f16f32
    cmp         x4, #3
    beq         .f16f16

.p2align 4
.packed_packed_loop_1:
    ld1         { v7.4s }, [ x2 ]
    ld1         { v0.4s, v1.4s, v2.4s, v3.4s }, [ x1 ], #64
    ld1         { v4.4s, v5.4s, v6.4s }, [ x1 ], #48
    add         x2, x2, #12

{% for col in range(0, 3) %}
    fmla        v{{ col * 8 + 8 }}.4s, v0.4s, v7.s[{{ col }}]
{% endfor %}

    ld1         { v0.4s }, [ x1 ], #16

{% for row in range(1, 7) %}
    {% for col in range(0, 3) %}
        fmla        v{{ col * 8 + 8 + row }}.4s, v{{row}}.4s, v7.s[{{ col }}]
    {% endfor %}
{% endfor %}

{% for col in range(0, 3) %}
    fmla        v{{ col * 8 + 15 }}.4s, v0.4s, v7.s[{{ col }}]
{% endfor %}

    subs        x3, x3, #1
    bne         .packed_packed_loop_1

    b           .non_linear_loop

.p2align 4
.f32f16:
    ld1         { v7.4h }, [ x2 ]
    ld1         { v0.4s, v1.4s, v2.4s, v3.4s }, [ x1 ], #64
    ld1         { v4.4s, v5.4s, v6.4s }, [ x1 ], #48
    fcvtl       v7.4s, v7.4h
    add         x2, x2, #6

{% for col in range(0, 3) %}
    fmla        v{{ col * 8 + 8 }}.4s, v0.4s, v7.s[{{ col }}]
{% endfor %}

    ld1         { v0.4s }, [ x1 ], #16

{% for row in range(1, 7) %}
    {% for col in range(0, 3) %}
        fmla        v{{ col * 8 + 8 + row }}.4s, v{{row}}.4s, v7.s[{{ col }}]
    {% endfor %}
{% endfor %}

{% for col in range(0, 3) %}
    fmla        v{{ col * 8 + 15 }}.4s, v0.4s, v7.s[{{ col }}]
{% endfor %}

    subs        x3, x3, #1
    bne         .f32f16

    b           .non_linear_loop

.p2align 4
.f16f32:
    ld1         { v7.4s }, [ x2 ]
    ld1         { v0.8h, v1.8h, v2.8h, v3.8h }, [ x1 ], #64
    add         x2, x2, #12

    fcvtl       v4.4s, v0.4h
    fcvtl2      v5.4s, v0.8h
    fcvtl       v6.4s, v1.4h
    fcvtl2      v0.4s, v1.8h

    {% for col in range(0, 3) %}
        fmla        v{{ col * 8 + 8 }}.4s, v4.4s, v7.s[{{ col }}]
        fmla        v{{ col * 8 + 9 }}.4s, v5.4s, v7.s[{{ col }}]
        fmla        v{{ col * 8 + 10 }}.4s, v6.4s, v7.s[{{ col }}]
        fmla        v{{ col * 8 + 11 }}.4s, v0.4s, v7.s[{{ col }}]
    {% endfor %}

    fcvtl       v4.4s, v2.4h
    fcvtl2      v5.4s, v2.8h
    fcvtl       v6.4s, v3.4h
    fcvtl2      v1.4s, v3.8h
    
    {% for col in range(0, 3) %}
        fmla        v{{ col * 8 + 12 }}.4s, v4.4s, v7.s[{{ col }}]
        fmla        v{{ col * 8 + 13 }}.4s, v5.4s, v7.s[{{ col }}]
        fmla        v{{ col * 8 + 14 }}.4s, v6.4s, v7.s[{{ col }}]
        fmla        v{{ col * 8 + 15 }}.4s, v1.4s, v7.s[{{ col }}]
    {% endfor %}

    subs        x3, x3, #1
    bne         .f16f32

    b           .non_linear_loop

.p2align 4
.f16f16:
    ld1         { v7.4h }, [ x2 ]
    ld1         { v0.8h, v1.8h, v2.8h, v3.8h }, [ x1 ], #64
    add         x2, x2, #6

    fcvtl       v7.4s, v7.4h

    fcvtl       v4.4s, v0.4h
    fcvtl2      v5.4s, v0.8h
    fcvtl       v6.4s, v1.4h
    fcvtl2      v0.4s, v1.8h

    {% for col in range(0, 3) %}
        fmla        v{{ col * 8 + 8 }}.4s, v4.4s, v7.s[{{ col }}]
        fmla        v{{ col * 8 + 9 }}.4s, v5.4s, v7.s[{{ col }}]
        fmla        v{{ col * 8 + 10 }}.4s, v6.4s, v7.s[{{ col }}]
        fmla        v{{ col * 8 + 11 }}.4s, v0.4s, v7.s[{{ col }}]
    {% endfor %}

    fcvtl       v4.4s, v2.4h
    fcvtl2      v5.4s, v2.8h
    fcvtl       v6.4s, v3.4h
    fcvtl2      v1.4s, v3.8h
    
    {% for col in range(0, 3) %}
        fmla        v{{ col * 8 + 12 }}.4s, v4.4s, v7.s[{{ col }}]
        fmla        v{{ col * 8 + 13 }}.4s, v5.4s, v7.s[{{ col }}]
        fmla        v{{ col * 8 + 14 }}.4s, v6.4s, v7.s[{{ col }}]
        fmla        v{{ col * 8 + 15 }}.4s, v1.4s, v7.s[{{ col }}]
    {% endfor %}

    subs        x3, x3, #1
    bne         .f16f16

    b           .non_linear_loop


{% set from = 8 %}{% set to = 31 %}{% include "arm64simd_mmm_f32_scalars.j2" %}
{% set mr = 32 %}{% set from = 8 %}{% set to = 31 %}{% include "arm64simd_mmm_f32_per_rows.j2" %}
{% set mr = 32 %}{% set from = 8 %}{% set to = 31 %}{% include "arm64simd_mmm_f32_per_cols.j2" %}
{% set from = 8 %}{% set to = 31 %}{% include "arm64simd_mmm_load_tile.j2" %}

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

    {% for col in range(0, 3) %}
        mov x4, x5
        {% for reg in range(0, 8) %}
            {% for lane in range(0, 4) %}
                ld1 {v0.s}[{{lane}}], [ x4 ], x6
            {% endfor %}
            fadd v{{ col * 8 + 8 + reg }}.4s, v{{ col * 8 + 8 + reg }}.4s, v0.4s
        {% endfor %}
        add x5, x5, x7
    {% endfor %}

    b           .non_linear_loop

.add_row_col_products:
    ldp         x2, x3, [x0, #8]

    ld1         { v7.d }[0], [ x3 ], #8
    ld1         { v7.s }[2], [ x3 ], #4
    ld1         { v0.4s, v1.4s, v2.4s, v3.4s }, [ x2 ], #64
    ld1         { v4.4s, v5.4s, v6.4s }, [ x2 ], #48

{% for col in range(0, 3) %}
    fmla        v{{ col * 8 + 8 }}.4s, v0.4s, v7.s[{{ col }}]
{% endfor %}

    ld1         { v0.4s }, [ x2 ], #16

{% for row in range(1, 7) %}
    {% for col in range(0, 3) %}
        fmla        v{{ col * 8 + 8 + row }}.4s, v{{row}}.4s, v7.s[{{ col }}]
    {% endfor %}
{% endfor %}

{% for col in range(0, 3) %}
    fmla        v{{ col * 8 + 15 }}.4s, v0.4s, v7.s[{{ col }}]
{% endfor %}

    b           .non_linear_loop

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

    cmp         x8, #2
    beq         .store_f16

    cmp         x6, #4
    beq         .store_strides_contig


    {% for col in range(0, 3) %}
        mov x4, x5
        {% for reg in range(0, 8) %}
            {% for lane in range(0, 4) %}
                st1 { v{{ col * 8 + 8 + reg }}.s }[{{lane}}], [ x4 ], x6
            {% endfor %}
        {% endfor %}
        add x5, x5, x7
    {% endfor %}
    b           .non_linear_loop

.store_strides_contig:

    {% for col in range(0, 3) %}
        mov x4, x5
        {% for r in range(0, 8) %}
            st1 { v{{ col * 8 + 8 + r }}.4s }, [ x4 ], 16
        {% endfor %}
        add x5, x5, x7
    {% endfor %}

    b           .non_linear_loop

.store_f16:

    cmp         x6, #2
    beq         .store_strides_contig_f16

    {% for col in range(0, 3) %}
        {% for reg in range(0, 4) %}
            fcvtn  v{{reg}}.4h, v{{ (col * 4 + reg) * 2 + 8 }}.4s
            fcvtn2 v{{reg}}.8h, v{{ (col * 4 + reg) * 2 + 9 }}.4s
        {% endfor %}

        mov x4, x5
        {% for reg in range(0, 4) %}
            {% for lane in range(0, 8) %}
                st1 { v{{reg}}.h }[{{lane}}], [ x4 ], x6
            {% endfor %}
        {% endfor %}
        add x5, x5, x7

    {% endfor %}


    b           .non_linear_loop

.store_strides_contig_f16:

    {% for col in range(0, 3) %}
        {% for reg in range(0, 4) %}
            fcvtn  v{{reg}}.4h, v{{ (col * 4 + reg) * 2 + 8 }}.4s
            fcvtn2 v{{reg}}.8h, v{{ (col * 4 + reg) * 2 + 9 }}.4s
        {% endfor %}

        mov x4, x5
        {% for reg in range(0, 4) %}
            st1 { v{{reg}}.4s }, [ x4 ], #16
        {% 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         x24, x25, [sp], #16
    ldp         x22, x23, [sp], #16
    ldp         x20, x21, [sp], #16

    ret

