#include "params.h"

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

.macro butterfly rl0,rl1,rl2,rl3,rh0,rh1,rh2,rh3 z0=3,z1=3,z2=3,z3=3
#mul
vpmuldq		%ymm\z0,%ymm\rh0,%ymm\rh0
vpmuldq		%ymm\z1,%ymm\rh1,%ymm\rh1
vpmuldq		%ymm\z2,%ymm\rh2,%ymm\rh2
vpmuldq		%ymm\z3,%ymm\rh3,%ymm\rh3

#reduce
vpmuldq		%ymm0,%ymm\rh0,%ymm12
vpmuldq		%ymm0,%ymm\rh1,%ymm13
vpmuldq		%ymm0,%ymm\rh2,%ymm14
vpmuldq		%ymm0,%ymm\rh3,%ymm15
vpmuldq		%ymm1,%ymm12,%ymm12
vpmuldq		%ymm1,%ymm13,%ymm13
vpmuldq		%ymm1,%ymm14,%ymm14
vpmuldq		%ymm1,%ymm15,%ymm15
vpsubd		%ymm12,%ymm\rh0,%ymm12
vpsubd		%ymm13,%ymm\rh1,%ymm13
vpsubd		%ymm14,%ymm\rh2,%ymm14
vpsubd		%ymm15,%ymm\rh3,%ymm15
vpsrlq		$32,%ymm12,%ymm12
vpsrlq		$32,%ymm13,%ymm13
vpsrlq		$32,%ymm14,%ymm14
vpsrlq		$32,%ymm15,%ymm15

#update
vpsubd		%ymm12,%ymm\rl0,%ymm\rh0
vpsubd		%ymm13,%ymm\rl1,%ymm\rh1
vpsubd		%ymm14,%ymm\rl2,%ymm\rh2
vpsubd		%ymm15,%ymm\rl3,%ymm\rh3
vpaddd		%ymm12,%ymm\rl0,%ymm\rl0
vpaddd		%ymm13,%ymm\rl1,%ymm\rl1
vpaddd		%ymm14,%ymm\rl2,%ymm\rl2
vpaddd		%ymm15,%ymm\rl3,%ymm\rl3
.endm

.global ntt_levels0t2_avx
ntt_levels0t2_avx:
#consts
vpbroadcastd	4*PINV(%rdx),%ymm0
vpbroadcastd	4*P(%rdx),%ymm1
pmovzxbd	4*MASK(%rdx),%xmm2

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

#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

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

level1:
#zetas
vpbroadcastd	4(%rcx),%ymm12
vpbroadcastd	8(%rcx),%ymm13

butterfly	4,5,8,9,6,7,10,11 12,12,13,13

level2:
#zetas
vpbroadcastd	12(%rcx),%ymm12
vpbroadcastd	16(%rcx),%ymm13
vpbroadcastd	20(%rcx),%ymm14
vpbroadcastd	24(%rcx),%ymm15

butterfly	4,6,8,10,5,7,9,11 12,13,14,15

store:
#store
vpermd		%ymm4,%ymm2,%ymm4
vpermd		%ymm5,%ymm2,%ymm5
vpermd		%ymm6,%ymm2,%ymm6
vpermd		%ymm7,%ymm2,%ymm7
vpermd		%ymm8,%ymm2,%ymm8
vpermd		%ymm9,%ymm2,%ymm9
vpermd		%ymm10,%ymm2,%ymm10
vpermd		%ymm11,%ymm2,%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

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

#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

#reduce
red32p		4 12
red32p		5 13
red32p		6 14
red32p		7 15

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

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

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

level4:
#zetas
vpbroadcastd	4(%rcx),%ymm12
vpbroadcastd	8(%rcx),%ymm13

butterfly	4,5,8,9,6,7,10,11 12,12,13,13

level5:
#zetas
vpbroadcastd	12(%rcx),%ymm12
vpbroadcastd	16(%rcx),%ymm13
vpbroadcastd	20(%rcx),%ymm14
vpbroadcastd	24(%rcx),%ymm15

butterfly	4,6,8,10,5,7,9,11 12,13,14,15

#store
vpermd		%ymm4,%ymm2,%ymm4
vpermd		%ymm5,%ymm2,%ymm5
vpermd		%ymm6,%ymm2,%ymm6
vpermd		%ymm7,%ymm2,%ymm7
vpermd		%ymm8,%ymm2,%ymm8
vpermd		%ymm9,%ymm2,%ymm9
vpermd		%ymm10,%ymm2,%ymm10
vpermd		%ymm11,%ymm2,%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 ntt_levels6t8_avx
ntt_levels6t8_avx:
#consts
vpbroadcastd	4*PINV(%rdx),%ymm0
vpbroadcastd	4*P(%rdx),%ymm1
vpbroadcastd	4*V(%rdx),%ymm2

#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

#reduce
red32p		4 12
red32p		5 13
red32p		6 14
red32p		7 15

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

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

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

level7:
#zetas
vpbroadcastd	4(%rcx),%ymm12
vpbroadcastd	8(%rcx),%ymm13

butterfly	4,5,8,9,6,7,10,11 12,12,13,13

level8:
#zetas
vpbroadcastd	12(%rcx),%ymm12
vpbroadcastd	16(%rcx),%ymm13
vpbroadcastd	20(%rcx),%ymm14
vpbroadcastd	24(%rcx),%ymm15

butterfly	4,6,8,10,5,7,9,11 12,13,14,15

#store
vpermd		%ymm4,%ymm2,%ymm4
vpermd		%ymm5,%ymm2,%ymm5
vpermd		%ymm6,%ymm2,%ymm6
vpermd		%ymm7,%ymm2,%ymm7
vpermd		%ymm8,%ymm2,%ymm8
vpermd		%ymm9,%ymm2,%ymm9
vpermd		%ymm10,%ymm2,%ymm10
vpermd		%ymm11,%ymm2,%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 ntt_levels9t11_avx
ntt_levels9t11_avx:
#consts
vpbroadcastd	4*PINV(%rdx),%ymm0
vpbroadcastd	4*P(%rdx),%ymm1
vpbroadcastd	4*V(%rdx),%ymm2

#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

#reduce
red32p		4 12
red32p		5 13
red32p		6 14
red32p		7 15

level9:
#zetas
vpmovzxdq	(%rcx),%ymm12

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

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

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

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

level10:
#zetas
vpmovzxdq	16(%rcx),%ymm12
vpmovzxdq	32(%rcx),%ymm13

butterfly	7,5,8,6,3,10,4,11 12,12,13,13

level11:
#zetas
vpmovzxdq	48(%rcx),%ymm12
vpmovzxdq	64(%rcx),%ymm13
vpmovzxdq	80(%rcx),%ymm14
vpmovzxdq	96(%rcx),%ymm15

butterfly	7,3,8,4,5,10,6,11 12,13,14,15

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

ret
