BITS 64

SECTION .note.GNU-stack noalloc noexec nowrite progbits

SECTION .text

GLOBAL dft_avx2

; rdi -> output
; rsi -> input
; rdx -> costab
; ecx -> fft_size
dft_avx2:
    push rbx
    push r12
    push r13
    push r14
    push r15

    vxorpd xmm0, xmm0
    xor rbx, rbx
dft_avx2_memset:
    vmovapd [rdi + rbx * 8], xmm0
    vmovapd [rdi + rbx * 8 + 16], xmm0
    vmovapd [rdi + rbx * 8 + 32], xmm0
    vmovapd [rdi + rbx * 8 + 48], xmm0
    add rbx, 8
    cmp rbx, rcx
    jb dft_avx2_memset

    mov r12d, ecx
    shl r12, 2

    mov eax, ecx
    shl rax, 32
    add rax, rcx
    vmovq xmm15, rax
    vmovq xmm14, rax
    vpunpcklqdq xmm15, xmm14

    mov eax, ecx
    shr eax, 2
    mov r8d, eax
    shl rax, 32
    add rax, r8
    vmovq xmm14, rax
    vmovq xmm13, rax
    vpunpcklqdq xmm14, xmm13

    mov eax, ecx
    dec eax
    shl rax, 32
    add rax, rcx
    dec rax
    vmovq xmm13, rax
    vmovq xmm12, rax
    vpunpcklqdq xmm13, xmm12

    xor ebx, ebx
dft_avx2_loop1:
    mov eax, ebx
    shl rax, 32
    vmovq xmm0, rax
    add rax, rbx
    vmovq xmm1, rax
    vpslld xmm1, 1
    vpaddd xmm1, xmm0
    mov r8, rbx
    vpunpcklqdq xmm0, xmm1
    shl r8, 3
    vmovdqa xmm1, xmm0
    vmovapd xmm4, [rsi + r8]
    vpaddd xmm1, xmm14
    vmovapd xmm5, [rsi + r8 + 16]
    vmovdqa xmm2, xmm0

    xor r11, r11
dft_avx2_loop2:
    vmovdqa xmm6, xmm0
    vmovdqa xmm7, xmm1
    vpcmpgtd xmm6, xmm13
    vpcmpgtd xmm7, xmm13
    vpand xmm6, xmm15
    vpand xmm7, xmm15
    vpsubd xmm0, xmm6
    vpsubd xmm1, xmm7
    vmovdqa xmm6, xmm0
    vmovdqa xmm7, xmm1
    vpcmpgtd xmm6, xmm13
    vpcmpgtd xmm7, xmm13
    vpand xmm6, xmm15
    vpand xmm7, xmm15
    vpsubd xmm0, xmm6
    vpsubd xmm1, xmm7
    vmovdqa xmm6, xmm0
    vmovdqa xmm7, xmm1
    vpcmpgtd xmm6, xmm13
    vpcmpgtd xmm7, xmm13
    vpand xmm6, xmm15
    vpand xmm7, xmm15
    vpsubd xmm0, xmm6
    vpsubd xmm1, xmm7
    vmovdqa xmm6, xmm0
    vmovdqa xmm7, xmm1
    vpcmpgtd xmm6, xmm13
    vpcmpgtd xmm7, xmm13
    vpand xmm6, xmm15
    vpand xmm7, xmm15
    vpsubd xmm0, xmm6
    vpsubd xmm1, xmm7

    vpcmpeqd xmm12, xmm12
    vmovdqa xmm11, xmm0
    vgatherdpd xmm6, [rdx+8*xmm0], xmm12
    vpcmpeqd xmm12, xmm12
    vpunpckhqdq xmm11, xmm0
    vgatherdpd xmm8, [rdx+8*xmm11], xmm12

    vmulpd xmm6, xmm4
    vmovapd xmm7, [rdi + r11]
    vmulpd xmm8, xmm5
    vmovapd xmm9, [rdi + r11 + 16]
    vaddpd xmm7, xmm6
    vaddpd xmm9, xmm8
    vmovapd [rdi + r11], xmm7
    vmovapd [rdi + r11 + 16], xmm9

    vpcmpeqd xmm12, xmm12
    vmovdqa xmm11, xmm1
    vgatherdpd xmm6, [rdx+8*xmm1], xmm12
    vpcmpeqd xmm12, xmm12
    vpunpckhqdq xmm11, xmm1
    vgatherdpd xmm8, [rdx+8*xmm11], xmm12

    add r11, r12

    vmulpd xmm6, xmm4
    vmovapd xmm7, [rdi + r11]
    vmulpd xmm8, xmm5
    vmovapd xmm9, [rdi + r11 + 16]
    vaddpd xmm7, xmm6
    vaddpd xmm9, xmm8
    vmovapd [rdi + r11], xmm7
    vmovapd [rdi + r11 + 16], xmm9

    sub r11, r12
    vpaddd xmm0, xmm2
    add r11, 32
    vpaddd xmm1, xmm2
    cmp r11, r12
    jb dft_avx2_loop2

    add rbx, 4
    cmp rbx, rcx
    jb dft_avx2_loop1

dft_avx2_end:
    pop r15
    pop r14
    pop r13
    pop r12
    pop rbx
    ret
