{#
// vim: set syntax=asm :

// AVX-512 (zmm, 16-wide) sigmoid. Uses a rational (Padé-style) approximation
// of sigmoid clamped to [-18, 18]. Validated against the generic scalar
// reference via sigmoid_frame_tests! (see x86_64_fma.rs).

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 ZMM0-15 and ZMM0-15
    return: rax (+rdx)

#}

{% if msvc %}

_text segment
avx512_sigmoid_f32_{{suffix}} proc

{% else %}

.intel_syntax noprefix
.text
.p2align 5
.globl {{G}}avx512_sigmoid_f32_{{suffix}}
{{G}}avx512_sigmoid_f32_{{suffix}}:
.cfi_startproc
{% endif %}

    push        rbp
    mov         rbp, rsp


{% if family == "windows" %}
// https://www.agner.org/optimize/calling_conventions.pdf xmm6-15 are not scratch
// https://stackoverflow.com/questions/43358429/save-value-of-xmm-registers
    and rsp,-16
    lea rsp,[rsp-160]
    vmovaps [rsp], xmm6
    vmovaps [rsp+16*1],xmm7
    vmovaps [rsp+16*2],xmm8
    vmovaps [rsp+16*3],xmm9
    vmovaps [rsp+16*4],xmm10
    vmovaps [rsp+16*5],xmm11
    vmovaps [rsp+16*6],xmm12
    vmovaps [rsp+16*7],xmm13
    vmovaps [rsp+16*8],xmm14
    vmovaps [rsp+16*9],xmm15

    // move around arguments to mimick SysV rdi,rsi passing
    push        rdi
    push        rsi
    mov         rdi, rcx
    mov         rsi, rdx

{% endif %}

    push        rbx
    push        r12
    push        r13
    push        r14
    push        r15

    sub         rsp, 8

{% if family == "unix" %}
// FIXME
// .cfi_def_cfa_offset 64
{% endif %}

    stmxcsr     [rsp + 4]
{% if msvc %}
    mov         rax, 1FC0h
{% else %}
    mov         rax, 0x1FC0
{% endif %}
    mov         [rsp], eax
    ldmxcsr     [rsp]
// ----------------------------------------------------------------------

{% set offset %}{% if msvc %} offset {%else%} rip + {%endif%} {% endset %}

    cmp     rsi, 0
    je      {{L}}done

    cmp     rsi, 64
    jl      {{L}}loop_1

{{L}}loop_4:

    vmovaps         zmm4, [rdi]
    vmovaps         zmm5, [rdi + 64]
    vmovaps         zmm6, [rdi + 128]
    vmovaps         zmm7, [rdi + 192]

    vbroadcastss    zmm0, dword ptr [{{offset}} {{L}}coeffs_num_low]
    vbroadcastss    zmm1, dword ptr [{{offset}} {{L}}coeffs_num_high]
    vbroadcastss    zmm2, dword ptr [{{offset}} {{L}}coeffs_num_alpha_9]
    vbroadcastss    zmm3, dword ptr [{{offset}} {{L}}coeffs_num_alpha_7]

    vmaxps          zmm4, zmm4, zmm0
    vmaxps          zmm5, zmm5, zmm0
    vmaxps          zmm6, zmm6, zmm0
    vmaxps          zmm7, zmm7, zmm0
    vbroadcastss    zmm0, dword ptr [{{offset}} {{L}}coeffs_num_alpha_5]

    vminps          zmm4, zmm4, zmm1
    vminps          zmm5, zmm5, zmm1
    vminps          zmm6, zmm6, zmm1
    vminps          zmm7, zmm7, zmm1        // zmm4..7 <- x
    vbroadcastss    zmm1, dword ptr [{{offset}} {{L}}coeffs_num_alpha_3]

    vmulps          zmm8, zmm4, zmm4
    vmulps          zmm9, zmm5, zmm5
    vmulps          zmm10, zmm6, zmm6
    vmulps          zmm11, zmm7, zmm7        // zmm8..11 <- x^2

    vmovaps         zmm12, zmm2
    vmovaps         zmm13, zmm2
    vmovaps         zmm14, zmm2
    vmovaps         zmm15, zmm2
    vbroadcastss    zmm2, dword ptr [{{offset}} {{L}}coeffs_num_alpha_1]
    vfmadd132ps     zmm12, zmm3, zmm8
    vfmadd132ps     zmm13, zmm3, zmm9
    vfmadd132ps     zmm14, zmm3, zmm10
    vfmadd132ps     zmm15, zmm3, zmm11
    vbroadcastss    zmm3, dword ptr [{{offset}} {{L}}coeffs_num_beta_10]
    vfmadd132ps     zmm12, zmm0, zmm8
    vfmadd132ps     zmm13, zmm0, zmm9
    vfmadd132ps     zmm14, zmm0, zmm10
    vfmadd132ps     zmm15, zmm0, zmm11
    vbroadcastss    zmm0, dword ptr [{{offset}} {{L}}coeffs_num_beta_8]
    vfmadd132ps     zmm12, zmm1, zmm8
    vfmadd132ps     zmm13, zmm1, zmm9
    vfmadd132ps     zmm14, zmm1, zmm10
    vfmadd132ps     zmm15, zmm1, zmm11
    vbroadcastss    zmm1, dword ptr [{{offset}} {{L}}coeffs_num_beta_6]
    vfmadd132ps     zmm12, zmm2, zmm8
    vfmadd132ps     zmm13, zmm2, zmm9
    vfmadd132ps     zmm14, zmm2, zmm10
    vfmadd132ps     zmm15, zmm2, zmm11
    vbroadcastss    zmm2, dword ptr [{{offset}} {{L}}coeffs_num_beta_4]
    vmulps          zmm4, zmm4, zmm12
    vmulps          zmm5, zmm5, zmm13
    vmulps          zmm6, zmm6, zmm14
    vmulps          zmm7, zmm7, zmm15   // zmm4..7 <- num

    vmovaps         zmm12, zmm3
    vmovaps         zmm13, zmm3
    vmovaps         zmm14, zmm3
    vmovaps         zmm15, zmm3
    vbroadcastss    zmm3, dword ptr [{{offset}} {{L}}coeffs_num_beta_2]
    vfmadd132ps     zmm12, zmm0, zmm8
    vfmadd132ps     zmm13, zmm0, zmm9
    vfmadd132ps     zmm14, zmm0, zmm10
    vfmadd132ps     zmm15, zmm0, zmm11
    vbroadcastss    zmm0, dword ptr [{{offset}} {{L}}coeffs_num_beta_0]
    vfmadd132ps     zmm12, zmm1, zmm8
    vfmadd132ps     zmm13, zmm1, zmm9
    vfmadd132ps     zmm14, zmm1, zmm10
    vfmadd132ps     zmm15, zmm1, zmm11
    vbroadcastss    zmm1, dword ptr [{{offset}} {{L}}coeffs_num_half]
    vfmadd132ps     zmm12, zmm2, zmm8
    vfmadd132ps     zmm13, zmm2, zmm9
    vfmadd132ps     zmm14, zmm2, zmm10
    vfmadd132ps     zmm15, zmm2, zmm11
    vfmadd132ps     zmm12, zmm3, zmm8
    vfmadd132ps     zmm13, zmm3, zmm9
    vfmadd132ps     zmm14, zmm3, zmm10
    vfmadd132ps     zmm15, zmm3, zmm11
    vfmadd132ps     zmm12, zmm0, zmm8
    vfmadd132ps     zmm13, zmm0, zmm9
    vfmadd132ps     zmm14, zmm0, zmm10
    vfmadd132ps     zmm15, zmm0, zmm11  // zmm12..14 <- denum

    vdivps          zmm4, zmm4, zmm12
    vdivps          zmm5, zmm5, zmm13
    vdivps          zmm6, zmm6, zmm14
    vdivps          zmm7, zmm7, zmm15
    vaddps          zmm4, zmm4, zmm1
    vaddps          zmm5, zmm5, zmm1
    vaddps          zmm6, zmm6, zmm1
    vaddps          zmm7, zmm7, zmm1

    vmovaps [rdi], zmm4
    vmovaps [rdi + 64], zmm5
    vmovaps [rdi + 128], zmm6
    vmovaps [rdi + 192], zmm7

    add     rdi, 256
    sub     rsi, 64
    cmp     rsi, 64
    jg      {{L}}loop_4

    cmp     rsi, 0
    je      {{L}}done

{{L}}loop_1:
    vmovaps         zmm4, [rdi]

    vbroadcastss    zmm0, dword ptr [{{offset}} {{L}}coeffs_num_low]
    vbroadcastss    zmm1, dword ptr [{{offset}} {{L}}coeffs_num_high]
    vbroadcastss    zmm2, dword ptr [{{offset}} {{L}}coeffs_num_alpha_9]
    vbroadcastss    zmm3, dword ptr [{{offset}} {{L}}coeffs_num_alpha_7]

    vmaxps          zmm4, zmm4, zmm0
    vbroadcastss    zmm0, dword ptr [{{offset}} {{L}}coeffs_num_alpha_5]

    vminps          zmm4, zmm4, zmm1        // zmm4 <- x
    vbroadcastss    zmm1, dword ptr [{{offset}} {{L}}coeffs_num_alpha_3]

    vmulps          zmm8, zmm4, zmm4        // zmm8 <- x^2

    vmovaps         zmm12, zmm2
    vbroadcastss    zmm2, dword ptr [{{offset}} {{L}}coeffs_num_alpha_1]
    vfmadd132ps     zmm12, zmm3, zmm8
    vbroadcastss    zmm3, dword ptr [{{offset}} {{L}}coeffs_num_beta_10]
    vfmadd132ps     zmm12, zmm0, zmm8
    vbroadcastss    zmm0, dword ptr [{{offset}} {{L}}coeffs_num_beta_8]
    vfmadd132ps     zmm12, zmm1, zmm8
    vbroadcastss    zmm1, dword ptr [{{offset}} {{L}}coeffs_num_beta_6]
    vfmadd132ps     zmm12, zmm2, zmm8
    vbroadcastss    zmm2, dword ptr [{{offset}} {{L}}coeffs_num_beta_4]
    vmulps          zmm4, zmm4, zmm12

    vmovaps         zmm12, zmm3
    vbroadcastss    zmm3, dword ptr [{{offset}} {{L}}coeffs_num_beta_2]
    vfmadd132ps     zmm12, zmm0, zmm8
    vbroadcastss    zmm0, dword ptr [{{offset}} {{L}}coeffs_num_beta_0]
    vfmadd132ps     zmm12, zmm1, zmm8
    vbroadcastss    zmm1, dword ptr [{{offset}} {{L}}coeffs_num_half]
    vfmadd132ps     zmm12, zmm2, zmm8
    vfmadd132ps     zmm12, zmm3, zmm8
    vfmadd132ps     zmm12, zmm0, zmm8

    vdivps          zmm4, zmm4, zmm12
    vaddps          zmm4, zmm4, zmm1

    vmovaps [rdi], zmm4
    add     rdi, 64
    sub     rsi, 16
    jnz     {{L}}loop_1

{{L}}done:

// ----------------------------------------------------------------------

    ldmxcsr     [rsp + 4]

    add         rsp, 8

    pop r15
    pop r14
    pop r13
    pop r12
    pop rbx

{% if family == "windows" %}
    pop rsi
    pop rdi

    vmovaps xmm15, [rsp+16*9]
    vmovaps xmm14, [rsp+16*8]
    vmovaps xmm13, [rsp+16*7]
    vmovaps xmm12, [rsp+16*6]
    vmovaps xmm11, [rsp+16*5]
    vmovaps xmm10, [rsp+16*4]
    vmovaps xmm9, [rsp+16*3]
    vmovaps xmm8, [rsp+16*2]
    vmovaps xmm7, [rsp+16*1]
    vmovaps xmm6, [rsp]
{% endif %}

    mov rsp, rbp
    pop rbp
    ret

{% set float %}{% if msvc %} real4 {%else%} .float {%endif%}{% endset %}

{{L}}coeffs_num_low:
    {{float}} -18.0                    // low
{{L}}coeffs_num_high:
    {{float}} 18.0                     // high

{{L}}coeffs_num_alpha_9:
    {{float}} 4.37031012579801e-11     // alpha_9
{{L}}coeffs_num_alpha_7:
    {{float}} 1.15627324459942e-07     // alpha_7
{{L}}coeffs_num_alpha_5:
    {{float}} 6.08574864600143e-05     // alpha_5
{{L}}coeffs_num_alpha_3:
    {{float}} 8.51377133304701e-03     // alpha_3
{{L}}coeffs_num_alpha_1:
    {{float}} 2.48287947061529e-01     // alpha_1

{{L}}coeffs_num_beta_10:
    {{float}} 6.10247389755681e-13
{{L}}coeffs_num_beta_8:
    {{float}} 5.76102136993427e-09
{{L}}coeffs_num_beta_6:
    {{float}} 6.29106785017040e-06     // beta_6
{{L}}coeffs_num_beta_4:
    {{float}} 1.70198817374094e-03     // beta_4
{{L}}coeffs_num_beta_2:
    {{float}} 1.16817656904453e-01     // beta_2
{{L}}coeffs_num_beta_0:
    {{float}} 9.93151921023180e-01     // beta_0

{{L}}coeffs_num_half:
    {{float}} 0.5

{% if msvc %}
avx512_sigmoid_f32_{{suffix}} endp
_text ends
end
{% else %}
.cfi_endproc
{% endif %}
