// vim: ft=arm

// C tile regs
// 
//      q8[0]    q10[0]   q12[0]    q14[0]
//      q8[1]    q10[1]   q12[1]    q14[1]
//      q8[2]    q10[2]   q12[2]    q14[2]
//      q8[3]    q10[3]   q12[3]    q14[3]
//
//      q9[0]    q11[0]   q13[0]    q15[0]
//      q9[1]    q11[1]   q13[1]    q15[1]
//      q9[2]    q11[2]   q13[2]    q15[2]
//      q9[3]    q11[3]   q13[3]    q15[3]

    .arm
    .text
    .global armv7neon_mmm_i32_8x4_{{suffix}}
    .type armv7neon_mmm_i32_8x4_{{suffix}}, %function

armv7neon_mmm_i32_8x4_{{suffix}}:

    pld     [r0]
    push    { r4-r12 }
    vpush   { q4-q7 }

{% include "dispatcher.j2" %}

.add_mat_mul:
    // r3 r4 r5 r6
    // k  a  b  packing
    cmp     r3, #0
    beq     .non_linear_loop

    mov     r1, r4 // packed A ptr
    pld     [r3]
    pld     [r5]

    cmp     r6, #1
    beq     .packed_packed_i8i8

    .packed_packed_loop_1:

    vldmia          r1!, { q0, q1 }
    vldmia          r5!, { q2 }

    vmla.s32       q8, q0, d4[0]
    vmla.s32       q9, q1, d4[0]

    vmla.s32       q10, q0, d4[1]
    vmla.s32       q11, q1, d4[1]

    vmla.s32       q12, q0, d5[0]
    vmla.s32       q13, q1, d5[0]

    vmla.s32       q14, q0, d5[1]
    vmla.s32       q15, q1, d5[1]

    subs r3, r3, #1
    bne .packed_packed_loop_1

    b   .non_linear_loop

    .packed_packed_i8i8:
    pld     [r5]                           // packed B ptr       

    cmp r3, #4
    blt .packed_packed_loop_i8i8_1

    .packed_packed_loop_i8i8_4:
    pld             [r1, #64]
    pld             [r5, #64]

    // q2: d4 -> d4,d5 A even cols (from r1)
    // q3: d6 -> d6,d7 A odd cols (from r1)
    // q0: s0 -> d0 : B even lines (from r5)
    // q1: s4 -> d2 : B odd lines (from r5)

    // 0
    vldmia          r1!, { d4 }
    vldmia          r5!, { s0 }

    vmovl.s8        q2, d4
    vmovl.s8        q0, d0

    vmlal.s16       q8, d4, d0[0]
    vmlal.s16       q9, d5, d0[0]

    vldmia          r1!, { d6 }

    vmlal.s16       q10, d4, d0[1]
    vmlal.s16       q11, d5, d0[1]

    vldmia          r5!, { s4 }

    vmlal.s16       q12, d4, d0[2]
    vmlal.s16       q13, d5, d0[2]

    vmlal.s16       q14, d4, d0[3]
    vmlal.s16       q15, d5, d0[3]

    // 1
    vmovl.s8        q3, d6
    vmovl.s8        q1, d2

    vmlal.s16       q8, d6, d2[0]
    vldmia          r1!, { d4 }
    vmlal.s16       q9, d7, d2[0]
    vldmia          r5!, { s0 }

    vmlal.s16       q10, d6, d2[1]
    vmlal.s16       q11, d7, d2[1]

    vmlal.s16       q12, d6, d2[2]
    vmlal.s16       q13, d7, d2[2]

    vmlal.s16       q14, d6, d2[3]
    vmlal.s16       q15, d7, d2[3]

    // 2
    vmovl.s8        q2, d4
    vmovl.s8        q0, d0

    vmlal.s16       q8, d4, d0[0]
    vmlal.s16       q9, d5, d0[0]

    vldmia          r1!, { d6 }

    vmlal.s16       q10, d4, d0[1]
    vmlal.s16       q11, d5, d0[1]

    vldmia          r5!, { s4 }

    vmlal.s16       q12, d4, d0[2]
    vmlal.s16       q13, d5, d0[2]

    vmlal.s16       q14, d4, d0[3]
    vmlal.s16       q15, d5, d0[3]

    // 3
    vmovl.s8        q3, d6
    vmovl.s8        q1, d2

    vmlal.s16       q8, d6, d2[0]
    vmlal.s16       q9, d7, d2[0]

    vmlal.s16       q10, d6, d2[1]
    vmlal.s16       q11, d7, d2[1]

    vmlal.s16       q12, d6, d2[2]
    vmlal.s16       q13, d7, d2[2]

    vmlal.s16       q14, d6, d2[3]
    vmlal.s16       q15, d7, d2[3]

    sub r3, r3, #4
    cmp r3, #4
    bge .packed_packed_loop_i8i8_4

    cmp r3, #0
    beq .non_linear_loop

    .packed_packed_loop_i8i8_1:

    vldmia          r1!, { s0, s1 }
    vmovl.s8        q0, d0
    vldmia          r5!, { s4 }
    vmovl.s8        q1, d2

    vmlal.s16       q8, d0, d2[0]
    vmlal.s16       q9, d1, d2[0]

    vmlal.s16       q10, d0, d2[1]
    vmlal.s16       q11, d1, d2[1]

    vmlal.s16       q12, d0, d2[2]
    vmlal.s16       q13, d1, d2[2]

    vmlal.s16       q14, d0, d2[3]
    vmlal.s16       q15, d1, d2[3]

    subs r3, r3, #1
    bne .packed_packed_loop_i8i8_1
    b   .non_linear_loop

{% set from = 8 %}{% set to = 15 %}{% include "armv7neon_mmm_i32_scalars.j2" %}
{% set mr = 8 %}{% set from = 8 %}{% set to = 15 %}{% include "armv7neon_mmm_i32_per_rows.j2" %}
{% set mr = 8 %}{% set from = 8 %}{% set to = 15 %}{% include "armv7neon_mmm_i32_per_cols.j2" %}

.add_unicast:
    // r3, r4, r5, r6 <- ptr, rsc, csc, size
    cmp     r6, #4
    beq     .non_linear_addc_i32

    {% for col in range(0, 4) %}
        mov         r8, r3
        {% for reg in range(0, 4) %}
            vld1.s8     d0[0], [ r8 ], r4
            vld1.s8     d0[1], [ r8 ], r4
            vmovl.s8    q0, d0
            vmovl.s16   q0, d0
            vadd.i32    d{{ col * 4 + reg + 16 }}, d0
        {% endfor %}
        add r3, r3, r5
    {% endfor %}

    b .non_linear_loop

.non_linear_addc_i32:

    {% for col in range(0, 4) %}
        mov         r8, r3
        {% for reg in range(0, 4) %}
            vld1.s32    d0[0], [ r8 ], r4
            vld1.s32    d0[1], [ r8 ], r4
            vadd.i32    d{{ col * 4 + reg + 16 }}, d0
        {% endfor %}
        {% if col < 3 %}
            add r3, r3, r5
        {% endif %}
    {% endfor %}

b .non_linear_loop

.add_row_col_products:
    vldmia          r3!, { q0, q1 }
    vldmia          r4!, { q4 }

    vmla.s32        q8, q0, d8[0]
    vmla.s32        q9, q1, d8[0]

    vmla.s32        q10, q0, d8[1]
    vmla.s32        q11, q1, d8[1]

    vmla.s32        q12, q0, d9[0]
    vmla.s32        q13, q1, d9[0]

    vmla.s32        q14, q0, d9[1]
    vmla.s32        q15, q1, d9[1]

    b .non_linear_loop

    {% include "armv7neon_mmm_i32_scale_q8_q15.j2" %}

.store:
    // r3, r4, r5, r6 <- ptr, rsc, csc, size
    cmp     r6, #4
    beq     .store_strides_i32

    {% for reg in range(8, 16) %}
        vmovn.s32 d{{ reg * 2 }}, q{{reg}}
        vmovn.s16 d{{ reg * 2 }}, q{{reg}}
    {% endfor %}
    {% for col in range(0, 4) %}
        mov         r8, r3
        {% for reg in range(0, 2) %}
            {% set d = (col * 2 + reg) * 2 + 16 %}
            vst1.s8     d{{d}}[0], [ r8 ], r4
            vst1.s8     d{{d}}[1], [ r8 ], r4
            vst1.s8     d{{d}}[2], [ r8 ], r4
            vst1.s8     d{{d}}[3], [ r8 ], r4
        {% endfor %}
        {% if col < 3 %}
            add r3, r3, r5
        {% endif %}
    {% endfor %}

    b .non_linear_loop

.store_strides_i32:

    {% for col in range(0, 4) %}
        mov         r8, r3
        {% for reg in range(0, 4) %}
            {% for lane in range(0, 2) %}
                vst1.s32     d{{ col * 4 + reg + 16 }}[{{lane}}], [ r8 ], r4
            {% endfor %}
        {% endfor %}
        {% if col < 3 %}
            add r3, r3, r5
        {% endif %}
    {% endfor %}

    b .non_linear_loop

.load_tile:
    vldmia          r3!, { q8-q15 }
    b .non_linear_loop

.return:
    vpop        { q4-q7 }
    pop         { r4-r12 }

    bx          lr

