; Copyright (C) 2026 Kiyotsugu Arai
; SPDX-License-Identifier: LGPL-3.0-or-later
;
; fft_batch_avx2.asm — AVX2 batched FFT butterfly kernel (float)
;
; Function:
;   batch4_butterfly_avx2_float(data, n, half_len, tw)
;
; Performs one stage of butterflies on the interleaved data of 4 signals.
; By broadcasting the same twiddle factor to all 4 signals, processes 4 FFTs simultaneously.
;
; Data layout (position k):
;   data[k*8 .. k*8+7] = {s0_re, s0_im, s1_re, s1_im, s2_re, s2_im, s3_re, s3_im}
;   Each position is 32 bytes = 1 x YMM register
;
; Twiddle table:
;   tw[j] = complex<float> = {re, im}  (8 bytes)
;   Broadcast: {re,im,re,im,re,im,re,im} → 1 x YMM
;
; Windows x64 calling convention:
;   rcx = float* data          (4-signal interleaved, 32-byte aligned recommended)
;   edx = int n                (FFT size, power of 2)
;   r8d = int half_len         (butterfly half-length, power of 2)
;   r9  = const float* tw      (twiddle table: complex<float> array)
;   return value: none
;   non-clobbered: rbx, rbp, rdi, rsi, r12-r15, xmm6-xmm15

.code

; sign_mask: {+1,-1,+1,-1,+1,-1,+1,-1} for addsub pattern
; (used for sign negation of the im component in complex multiplication)

batch4_butterfly_avx2_float PROC
    ; Save non-volatile registers
    push    rbx
    push    rdi
    push    rsi

    ; rcx = data, edx = n, r8d = half_len, r9 = tw
    mov     rdi, rcx                    ; rdi = data
    movsxd  rax, edx                    ; rax = n
    movsxd  rbx, r8d                    ; rbx = half_len
    mov     rsi, r9                     ; rsi = tw

    ; len = half_len * 2
    lea     r10, [rbx*2]               ; r10 = len = 2 * half_len

    ; sign_mask = {-1,+1,-1,+1,-1,+1,-1,+1} (element order)
    ; complex multiplication: re positions (even) are -1, im positions (odd) are +1
    mov     r11d, 0BF800000h            ; -1.0f IEEE 754
    vmovd   xmm4, r11d
    mov     r11d, 03F800000h            ; +1.0f IEEE 754
    vmovd   xmm3, r11d
    vpunpckldq xmm5, xmm4, xmm3        ; {-1.0f, +1.0f, 0, 0}
    vbroadcastsd ymm5, xmm5            ; {-1,+1,-1,+1,-1,+1,-1,+1}

    ; outer loop: i = 0, i < n, i += len
    xor     ecx, ecx                    ; ecx = i = 0

outer_loop:
    cmp     rcx, rax                    ; i < n ?
    jge     done

    ; inner loop: j = 0, j < half_len
    xor     edx, edx                    ; edx = j = 0

inner_loop:
    cmp     rdx, rbx                    ; j < half_len ?
    jge     inner_done

    ; k = i + j
    lea     r11, [rcx + rdx]            ; r11 = k = i + j

    ; Compute byte offsets (each position = 8 floats = 32 bytes)
    shl     r11, 5                      ; r11 = k * 32
    lea     r10, [rcx + rdx]
    add     r10, rbx                    ; r10 = k + half_len
    shl     r10, 5                      ; r10 = (k + half_len) * 32

    ; Load u = data[k] (4 signals' complex values at position k)
    vmovups ymm0, YMMWORD PTR [rdi + r11]

    ; Load d = data[k + half_len] (4 signals at position k+half_len)
    vmovups ymm1, YMMWORD PTR [rdi + r10]

    ; Broadcast twiddle factor tw[j]
    ; tw[j] = {re, im} = 8 bytes at tw + j*8
    lea     r8, [rdx * 8]              ; j * sizeof(complex<float>)
    vbroadcastsd ymm2, QWORD PTR [rsi + r8]
    ; ymm2 = {w_re, w_im, w_re, w_im, w_re, w_im, w_re, w_im}

    ; Complex multiply: v = d * w
    ; d = {d0r,d0i, d1r,d1i, d2r,d2i, d3r,d3i}  (ymm1)
    ; w = {wr, wi,  wr, wi,  wr, wi,  wr, wi }  (ymm2)
    ; use only ymm0-5 (ymm6+ are non-clobbered under Windows x64)

    ; d_re = {d0r,d0r, d1r,d1r, d2r,d2r, d3r,d3r}
    vshufps ymm3, ymm1, ymm1, 0A0h

    ; d_im = {d0i,d0i, d1i,d1i, d2i,d2i, d3i,d3i}
    vshufps ymm4, ymm1, ymm1, 0F5h     ; ymm1 no longer needed

    ; w_flip = {wi,wr, wi,wr, wi,wr, wi,wr}
    vshufps ymm1, ymm2, ymm2, 0B1h     ; reuse ymm1

    ; p1 = d_re * w
    vmulps  ymm3, ymm3, ymm2           ; ymm2 no longer needed

    ; p2 = d_im * w_flip
    vmulps  ymm4, ymm4, ymm1           ; ymm1 no longer needed

    ; v = p1 + p2 * sign_mask
    vmulps  ymm4, ymm4, ymm5           ; p2 * sign_mask
    vaddps  ymm1, ymm3, ymm4           ; v = p1 + p2*sign

    ; Butterfly: store u+v, u-v
    vaddps  ymm3, ymm0, ymm1           ; u + v
    vsubps  ymm4, ymm0, ymm1           ; u - v

    vmovups YMMWORD PTR [rdi + r11], ymm3   ; data[k] = u + v
    vmovups YMMWORD PTR [rdi + r10], ymm4   ; data[k+half] = u - v

    ; j++
    inc     edx
    jmp     inner_loop

inner_done:
    ; i += len (= 2 * half_len)
    lea     r10, [rbx*2]
    add     rcx, r10
    jmp     outer_loop

done:
    vzeroupper
    pop     rsi
    pop     rdi
    pop     rbx
    ret

batch4_butterfly_avx2_float ENDP

END