#include "params.h"

.include "fq.h"
.include "shuffle.h"

.macro butterfly l0,l1,l2,l3,h0,h1,h2,h3 z0=3,z1=3,z2=3,z3=3
vpsubd		%ymm\h0,%ymm\l0,%ymm12
vpsubd		%ymm\h1,%ymm\l1,%ymm13
vpaddd		%ymm\h0,%ymm\l0,%ymm\l0

vpmuldq		%ymm\z0,%ymm12,%ymm12
vpmuldq		%ymm\z1,%ymm13,%ymm13
vpsubd		%ymm\h2,%ymm\l2,%ymm14

vpmuldq		%ymm\z2,%ymm14,%ymm14
vpsubd		%ymm\h3,%ymm\l3,%ymm15
vpaddd		%ymm\h1,%ymm\l1,%ymm\l1

vpmuldq		%ymm\z3,%ymm15,%ymm15
vpaddd		%ymm\h2,%ymm\l2,%ymm\l2
vpaddd		%ymm\h3,%ymm\l3,%ymm\l3

vpmuldq		%ymm0,%ymm12,%ymm\h0
vpmuldq		%ymm0,%ymm13,%ymm\h1
vpmuldq		%ymm0,%ymm14,%ymm\h2
vpmuldq		%ymm0,%ymm15,%ymm\h3

vpmuldq		%ymm1,%ymm\h0,%ymm\h0
vpmuldq		%ymm1,%ymm\h1,%ymm\h1
vpmuldq		%ymm1,%ymm\h2,%ymm\h2
vpmuldq		%ymm1,%ymm\h3,%ymm\h3

vpsubd		%ymm12,%ymm\h0,%ymm\h0
vpsubd		%ymm13,%ymm\h1,%ymm\h1

vpsrlq		$32,%ymm\h0,%ymm\h0
vpsrlq		$32,%ymm\h1,%ymm\h1
vpsubd		%ymm14,%ymm\h2,%ymm\h2

vpsrlq		$32,%ymm\h2,%ymm\h2
vpsubd		%ymm15,%ymm\h3,%ymm\h3

vpsrlq		$32,%ymm\h3,%ymm\h3
.endm

.global invntt_levels0t2_avx
invntt_levels0t2_avx:
#consts
vpbroadcastd	4*PINV(%rdx),%ymm0
vpbroadcastd	4*P(%rdx),%ymm1

#load
vpmovzxdq	(%rsi),%ymm4
vpmovzxdq	16(%rsi),%ymm5
vpmovzxdq	32(%rsi),%ymm6
vpmovzxdq	48(%rsi),%ymm7
vpmovzxdq	64(%rsi),%ymm8
vpmovzxdq	80(%rsi),%ymm9
vpmovzxdq	96(%rsi),%ymm10
vpmovzxdq	112(%rsi),%ymm11

scale:
#factor
vpbroadcastd	4*DIV(%rdx),%ymm3

#mul
vpmuldq		%ymm3,%ymm4,%ymm4
vpmuldq		%ymm3,%ymm5,%ymm5
vpmuldq		%ymm3,%ymm6,%ymm6
vpmuldq		%ymm3,%ymm7,%ymm7
vpmuldq		%ymm3,%ymm8,%ymm8
vpmuldq		%ymm3,%ymm9,%ymm9
vpmuldq		%ymm3,%ymm10,%ymm10
vpmuldq		%ymm3,%ymm11,%ymm11

#reduce
vpmuldq		%ymm0,%ymm4,%ymm12
vpmuldq		%ymm0,%ymm5,%ymm13
vpmuldq		%ymm0,%ymm6,%ymm14
vpmuldq		%ymm0,%ymm7,%ymm15
vpmuldq		%ymm1,%ymm12,%ymm12
vpmuldq		%ymm1,%ymm13,%ymm13
vpmuldq		%ymm1,%ymm14,%ymm14
vpmuldq		%ymm1,%ymm15,%ymm15
vpsubd		%ymm12,%ymm4,%ymm4
vpsubd		%ymm13,%ymm5,%ymm5
vpsubd		%ymm14,%ymm6,%ymm6
vpsubd		%ymm15,%ymm7,%ymm7
vpsrlq		$32,%ymm4,%ymm4
vpsrlq		$32,%ymm5,%ymm5
vpsrlq		$32,%ymm6,%ymm6
vpsrlq		$32,%ymm7,%ymm7

vpmuldq		%ymm0,%ymm8,%ymm12
vpmuldq		%ymm0,%ymm9,%ymm13
vpmuldq		%ymm0,%ymm10,%ymm14
vpmuldq		%ymm0,%ymm11,%ymm15
vpmuldq		%ymm1,%ymm12,%ymm12
vpmuldq		%ymm1,%ymm13,%ymm13
vpmuldq		%ymm1,%ymm14,%ymm14
vpmuldq		%ymm1,%ymm15,%ymm15
vpsubd		%ymm12,%ymm8,%ymm8
vpsubd		%ymm13,%ymm9,%ymm9
vpsubd		%ymm14,%ymm10,%ymm10
vpsubd		%ymm15,%ymm11,%ymm11
vpsrlq		$32,%ymm8,%ymm8
vpsrlq		$32,%ymm9,%ymm9
vpsrlq		$32,%ymm10,%ymm10
vpsrlq		$32,%ymm11,%ymm11

level0:
#zetas
vmovdqu		48(%rcx),%ymm3
vmovdqu		80(%rcx),%ymm15
vpshufd		$0x1B,%ymm15,%ymm15
vpshufd		$0x1B,%ymm3,%ymm3
vextracti128	$1,%ymm15,%xmm14
vextracti128	$1,%ymm3,%xmm2
vpmovzxdq	%xmm3,%ymm3
vpmovzxdq	%xmm15,%ymm15
vpmovzxdq	%xmm2,%ymm2
vpmovzxdq	%xmm14,%ymm14

butterfly	4,6,8,10,5,7,9,11 14,15,2,3

#2,1,2,1,2,1,2,1

level1:
#zetas
vmovdqu		16(%rcx),%ymm3
vpshufd		$0x1B,%ymm3,%ymm3
vextracti128	$1,%ymm3,%xmm15
vpmovzxdq	%xmm3,%ymm3
vpmovzxdq	%xmm15,%ymm15

#consts
vpbroadcastd	4*V(%rdx),%ymm2

butterfly	4,5,8,9,6,7,10,11 15,15,3,3

#4,2,1,1,4,2,1,1

vpxor		%ymm3,%ymm3,%ymm3
red32p		4 12
red32n		8 13

#1,2,1,1,-1,2,1,1

level2:
#zetas
vmovdqu		(%rcx),%xmm3
vpshufd		$0x1B,%xmm3,%xmm3
vpmovzxdq	%xmm3,%ymm3

butterfly	4,5,6,7,8,9,10,11 3,3,3,3

#1,4,2,2,1,1,1,1

red32p		5 12
red32p		6 13
red32p		7 14

#1,1,1,1,1,1,1,1

#consts
pmovzxbd	4*MASK(%rdx),%xmm2

#shuffle
shuffle4	4,5,3,5
shuffle4	6,7,4,7
shuffle4	8,9,6,9
shuffle4	10,11,8,11

shuffle8	3,4,10,4
shuffle8	6,8,3,8
shuffle8	5,7,6,7
shuffle8	9,11,5,11

#store
vpermd          %ymm10,%ymm2,%ymm10
vpermd          %ymm3,%ymm2,%ymm3
vpermd          %ymm6,%ymm2,%ymm6
vpermd          %ymm5,%ymm2,%ymm5
vpermd          %ymm4,%ymm2,%ymm4
vpermd          %ymm8,%ymm2,%ymm8
vpermd          %ymm7,%ymm2,%ymm7
vpermd          %ymm11,%ymm2,%ymm11
vmovdqa         %xmm10,(%rdi)
vmovdqa         %xmm3,16(%rdi)
vmovdqa         %xmm6,32(%rdi)
vmovdqa         %xmm5,48(%rdi)
vmovdqa         %xmm4,64(%rdi)
vmovdqa         %xmm8,80(%rdi)
vmovdqa         %xmm7,96(%rdi)
vmovdqa         %xmm11,112(%rdi)

ret

.global invntt_levels3t5_avx
invntt_levels3t5_avx:
#consts
vpbroadcastd	4*PINV(%rdx),%ymm0
vpbroadcastd	4*P(%rdx),%ymm1

#load
vpmovzxdq	(%rsi),%ymm4
vpmovzxdq	32(%rsi),%ymm5
vpmovzxdq	64(%rsi),%ymm6
vpmovzxdq	96(%rsi),%ymm7
vpmovzxdq	128(%rsi),%ymm8
vpmovzxdq	160(%rsi),%ymm9
vpmovzxdq	192(%rsi),%ymm10
vpmovzxdq	224(%rsi),%ymm11

level3:
#zetas
vpbroadcastd	12(%rcx),%ymm3
vpbroadcastd	16(%rcx),%ymm2
vpbroadcastd	20(%rcx),%ymm15
vpbroadcastd	24(%rcx),%ymm14

butterfly	4,6,8,10,5,7,9,11 14,15,2,3

level4:
#zetas
vpbroadcastd	4(%rcx),%ymm3
vpbroadcastd	8(%rcx),%ymm15

#consts
vpbroadcastd	4*V(%rdx),%ymm2

butterfly	4,5,8,9,6,7,10,11 15,15,3,3

vpxor		%ymm3,%ymm3,%ymm3
red32p		4 12
red32n		8 13

level5:
#zetas
vpbroadcastd	(%rcx),%ymm3

butterfly	4,5,6,7,8,9,10,11 3,3,3,3

#consts
pmovzxbd	4*MASK(%rdx),%xmm3

red32p		5 12
red32p		6 13
red32p		7 14

#store
vpermd          %ymm4,%ymm3,%ymm4
vpermd          %ymm5,%ymm3,%ymm5
vpermd          %ymm6,%ymm3,%ymm6
vpermd          %ymm7,%ymm3,%ymm7
vpermd          %ymm8,%ymm3,%ymm8
vpermd          %ymm9,%ymm3,%ymm9
vpermd          %ymm10,%ymm3,%ymm10
vpermd          %ymm11,%ymm3,%ymm11
vmovdqa         %xmm4,(%rdi)
vmovdqa         %xmm5,32(%rdi)
vmovdqa         %xmm6,64(%rdi)
vmovdqa         %xmm7,96(%rdi)
vmovdqa         %xmm8,128(%rdi)
vmovdqa         %xmm9,160(%rdi)
vmovdqa         %xmm10,192(%rdi)
vmovdqa         %xmm11,224(%rdi)

ret

.global invntt_levels6t8_avx
invntt_levels6t8_avx:
#consts
vpbroadcastd	4*PINV(%rdx),%ymm0
vpbroadcastd	4*P(%rdx),%ymm1

#load
vpmovzxdq	(%rsi),%ymm4
vpmovzxdq	256(%rsi),%ymm5
vpmovzxdq	512(%rsi),%ymm6
vpmovzxdq	768(%rsi),%ymm7
vpmovzxdq	1024(%rsi),%ymm8
vpmovzxdq	1280(%rsi),%ymm9
vpmovzxdq	1536(%rsi),%ymm10
vpmovzxdq	1792(%rsi),%ymm11

level6:
#zetas
vpbroadcastd	12(%rcx),%ymm3
vpbroadcastd	16(%rcx),%ymm2
vpbroadcastd	20(%rcx),%ymm15
vpbroadcastd	24(%rcx),%ymm14

butterfly	4,6,8,10,5,7,9,11 14,15,2,3

level7:
#zetas
vpbroadcastd	4(%rcx),%ymm3
vpbroadcastd	8(%rcx),%ymm15

#consts
vpbroadcastd	4*V(%rdx),%ymm2

butterfly	4,5,8,9,6,7,10,11 15,15,3,3

vpxor		%ymm3,%ymm3,%ymm3
red32p		4 12
red32n		8 13

level8:
#zetas
vpbroadcastd	(%rcx),%ymm3

butterfly	4,5,6,7,8,9,10,11 3,3,3,3

#consts
pmovzxbd	4*MASK(%rdx),%xmm3

red32p		5 12
red32p		6 13
red32p		7 14

#store
vpermd          %ymm4,%ymm3,%ymm4
vpermd          %ymm5,%ymm3,%ymm5
vpermd          %ymm6,%ymm3,%ymm6
vpermd          %ymm7,%ymm3,%ymm7
vpermd          %ymm8,%ymm3,%ymm8
vpermd          %ymm9,%ymm3,%ymm9
vpermd          %ymm10,%ymm3,%ymm10
vpermd          %ymm11,%ymm3,%ymm11
vmovdqa         %xmm4,(%rdi)
vmovdqa         %xmm5,256(%rdi)
vmovdqa         %xmm6,512(%rdi)
vmovdqa         %xmm7,768(%rdi)
vmovdqa         %xmm8,1024(%rdi)
vmovdqa         %xmm9,1280(%rdi)
vmovdqa         %xmm10,1536(%rdi)
vmovdqa         %xmm11,1792(%rdi)

ret

.global invntt_levels9t11_avx
invntt_levels9t11_avx:
#consts
vpbroadcastd	4*PINV(%rdx),%ymm0
vpbroadcastd	4*P(%rdx),%ymm1

#load
vpmovzxdq	(%rsi),%ymm4
vpmovzxdq	2048(%rsi),%ymm5
vpmovzxdq	4096(%rsi),%ymm6
vpmovzxdq	6144(%rsi),%ymm7
vpmovzxdq	8192(%rsi),%ymm8
vpmovzxdq	10240(%rsi),%ymm9
vpmovzxdq	12288(%rsi),%ymm10
vpmovzxdq	14336(%rsi),%ymm11

level9:
#zetas
vpbroadcastd	12(%rcx),%ymm3
vpbroadcastd	16(%rcx),%ymm2
vpbroadcastd	20(%rcx),%ymm15
vpbroadcastd	24(%rcx),%ymm14

butterfly	4,6,8,10,5,7,9,11 14,15,2,3

level10:
#zetas
vpbroadcastd	4(%rcx),%ymm3
vpbroadcastd	8(%rcx),%ymm15

#consts
vpbroadcastd	4*V(%rdx),%ymm2

butterfly	4,5,8,9,6,7,10,11 15,15,3,3

vpxor		%ymm3,%ymm3,%ymm3
red32p		4 12
red32n		8 13

level11:
#zetas
vpbroadcastd	(%rcx),%ymm3

butterfly	4,5,6,7,8,9,10,11 3,3,3,3

#consts
pmovzxbd	4*MASK(%rdx),%xmm3

#red32p		5 12
#red32p		6 13
#red32p		7 14

#store
vpermd          %ymm4,%ymm3,%ymm4
vpermd          %ymm5,%ymm3,%ymm5
vpermd          %ymm6,%ymm3,%ymm6
vpermd          %ymm7,%ymm3,%ymm7
vpermd          %ymm8,%ymm3,%ymm8
vpermd          %ymm9,%ymm3,%ymm9
vpermd          %ymm10,%ymm3,%ymm10
vpermd          %ymm11,%ymm3,%ymm11
vmovdqa         %xmm4,(%rdi)
vmovdqa         %xmm5,2048(%rdi)
vmovdqa         %xmm6,4096(%rdi)
vmovdqa         %xmm7,6144(%rdi)
vmovdqa         %xmm8,8192(%rdi)
vmovdqa         %xmm9,10240(%rdi)
vmovdqa         %xmm10,12288(%rdi)
vmovdqa         %xmm11,14336(%rdi)

ret
