{#
// vim: set syntax=asm :

/* mmm 64 x 3:

    zmm0 zmm4 zmm8
    zmm1 zmm5 zmm9
    zmm2 zmm6 zmm10
    zmm3 zmm7 zmm11

System V ABI:
    args: rdi, rsi, rdx, rcx, r8, r9
    preserve: rbx, rsp, rbp, r12, r13, r14, r15
    scratch: rax, rdi, rsi, rdx, rcx, r8, r9, r10, r11
    return: rax (+rdx)

Windows ABI:
    args: RCX, RDX, R8, R9
    preserve: RBX, RBP, RDI, RSI, RSP, R12, R13, R14, R15, and XMM6-15
    scratch: RAX, RCX, RDX, R8, R9, R10, R11, XMM0-5, and the upper portions of YMM0-15 and ZMM0-15
    return: rax (+rdx)
*/
#}

{% set size = "64x3" %}{% set suffix = suffix %}{% set G = G %}{% set arch = "avx512" %}{% include "preamble.j2" %}

{{L}}clear:
    vzeroall
    jmp     {{L}}non_linear_loop

{{L}}add_mat_mul:
    mov     rcx,    [rdi + 24]   // B
    mov     rax,    [rdi + 16]   // A

    mov     rbx,    [rdi + 8]    // k
    test    rbx,    rbx
    jz      {{L}}non_linear_loop

{{L}}main_loop_packed_packed:
	{% include "4x3/packed_packed_loop1/avx-512.S.raw" %}

    dec             rbx
    jnz             {{L}}main_loop_packed_packed

    jmp             {{L}}non_linear_loop

{% set from = 0 %}{% set to = 11 %}{% include "f32_scalars.j2" %}
{% set mr = 64 %}{% set from = 0 %}{% set to = 11 %}{% include "f32_per_rows.j2" %}
{% set mr = 64 %}{% set from = 0 %}{% set to = 11 %}{% include "f32_per_cols.j2" %}
{% set from = 0 %}{% set to = 11 %}{% include "avx512_mmm_load_tile.j2" %}

{{L}}add_unicast:

    mov     r10,    [rdi + 8]           // c ptr
    mov     rsi,    [rdi + 16]          // row stride
    mov     rbx,    [rdi + 24]          // col stride

    mov     eax,    0

{% for i in range(0, 4) %}
    pinsrd  xmm14, eax, {{i}}
    add     eax,    esi
{% endfor %}
{% for i in range(0, 4) %}
    pinsrd  xmm15, eax, {{i}}
    add     eax,    esi
{% endfor %}
{% for i in range(0, 4) %}
    pinsrd  xmm12, eax, {{i}}
    add     eax,    esi
{% endfor %}
{% for i in range(0, 4) %}
    pinsrd  xmm13, eax, {{i}}
    add     eax,    esi
{% endfor %}

    vperm2f128      ymm14,  ymm14, ymm15,         32 // ymm14 <- xmm14::xmm15
    vperm2f128      ymm13,  ymm12, ymm13,         32 // ymm12 <- xmm12::xmm13
    vinsertf32x8    zmm14, zmm14, ymm13, 1

{% for i in range(0, 3) %}
    kxnorw k1,k1,k1
    vgatherdps      zmm12{k1},  [ r10 + zmm14 ]
    add     r10, rbx
    vaddps          zmm{{ i * 4 }},   zmm{{ i * 4 }},   zmm12
{% endfor %}

    imul    esi,    16
    vpbroadcastd    zmm15, esi

{% for j in range(1, 4) %}
    mov     r10,    [rdi + 8]
    vpaddd          zmm14, zmm14, zmm15

    {% for i in range(0, 3) %}
        kxnorw k1,k1,k1
        vgatherdps      zmm12{k1},  [ r10 + zmm14 ]
        add     r10, rbx
        vaddps          zmm{{ i * 4 + j }},   zmm{{ i * 4 + j }},   zmm12
    {% endfor %}
{% endfor %}

    jmp    {{L}}non_linear_loop

{{L}}add_row_col_products:
    mov             rax, [ rdi + 8 ]
    mov             rbx, [ rdi + 16 ]

    vbroadcastss    zmm13, dword ptr [rbx]
    vbroadcastss    zmm14, dword ptr [rbx+4]
    vbroadcastss    zmm15, dword ptr [rbx+8]

{% for i in range(0, 4) %}
    vmovups         zmm12, zmmword ptr [rax+{{ i * 64 }}]
    vfmadd231ps     zmm{{i}}, zmm12, zmm13
    vfmadd231ps     zmm{{ i + 4 }}, zmm12, zmm14
    vfmadd231ps     zmm{{ i + 8 }}, zmm12, zmm15
{% endfor %}

    jmp    {{L}}non_linear_loop

{{L}}store:
    mov     r8,     [rdi + 8]           // c ptr
    mov     rsi,    [rdi + 16]          // row stride
    mov     rbx,    [rdi + 24]          // col stride

    // tops of cols
    lea     r9,     [ r8 + rbx ]
    lea     r10,    [ r8 + 2 * rbx ]
    lea     r11,    [ r10 + rbx ]

    {% for word in range(0, 4) %}
        {% for quarter in range(0, 4) %}
            {% for r in range(0, 3) %}
                vextractf32x4 xmm{{ r + 12 }}, zmm{{ r * 4 + word }}, {{quarter}}
            {% endfor %}
            {% for row in range(0, 4) %}
                {% for i in range(0, 3) %}
                    vextractps  dword ptr [r{{ i + 8 }}], xmm{{ i + 12 }}, {{row}}
                    add         r{{ i + 8 }}, rsi
                {% endfor %}
            {% endfor %}
        {% endfor %}
    {% endfor %}    

    jmp     {{L}}non_linear_loop

{% set size = "64x3" %}{% set suffix = suffix %}{% set G = G %}{% set L = L %}{% set arch = "avx512" %}{% include "postamble.j2" %}

