Rewrote tomont.S with a macro instead of the function call

Signed-off-by: Dusan Kostic <[email protected]>

Author: dkostic

dev/x86_64/src/tomont.S

@@ -18,7 +18,8 @@
  * Changes:
  * - Add call to csub in reduce128_avx to produce outputs
  *   in [0,1,...,q-1] rather than [0,1,...,q], matching the
- *   semantics of mlk_poly_reduce().
+ *   semantics of mlk_poly_reduce(),
+ * - Use a macro instead of a local function call.
  */
 
 #include "../../../common.h"
@@ -29,29 +30,16 @@
 #include "consts.h"
 #include "fq.inc"
 
-.text
-.global MLK_ASM_NAMESPACE(tomont_avx2)
-.balign 4
-MLK_ASM_FN_SYMBOL(tomont_avx2)
-#consts
-vmovdqa		MLK_AVX2_BACKEND_DATA_OFFSET_16XQ*2(%rsi),%ymm0
-vmovdqa		MLK_AVX2_BACKEND_DATA_OFFSET_16XMONTSQLO*2(%rsi),%ymm1
-vmovdqa		MLK_AVX2_BACKEND_DATA_OFFSET_16XMONTSQHI*2(%rsi),%ymm2
-call		tomont_avx2_core
-add		$256,%rdi
-call		tomont_avx2_core
-ret
-
-tomont_avx2_core:
+.macro tomont_128_coefficients offset
 #load
-vmovdqa		(%rdi),%ymm3
-vmovdqa		32(%rdi),%ymm4
-vmovdqa		64(%rdi),%ymm5
-vmovdqa		96(%rdi),%ymm6
-vmovdqa		128(%rdi),%ymm7
-vmovdqa		160(%rdi),%ymm8
-vmovdqa		192(%rdi),%ymm9
-vmovdqa		224(%rdi),%ymm10
+vmovdqa	(\offset +  0)(%rdi), %ymm3
+vmovdqa	(\offset + 32)(%rdi), %ymm4
+vmovdqa	(\offset + 64)(%rdi), %ymm5
+vmovdqa	(\offset + 96)(%rdi), %ymm6
+vmovdqa	(\offset +128)(%rdi), %ymm7
+vmovdqa	(\offset +160)(%rdi), %ymm8
+vmovdqa	(\offset +192)(%rdi), %ymm9
+vmovdqa	(\offset +224)(%rdi), %ymm10
 
 fqmulprecomp	1,2,3,11
 fqmulprecomp	1,2,4,12
@@ -63,14 +51,27 @@ fqmulprecomp	1,2,9,12
 fqmulprecomp	1,2,10,13
 
 #store
-vmovdqa		%ymm3,(%rdi)
-vmovdqa		%ymm4,32(%rdi)
-vmovdqa		%ymm5,64(%rdi)
-vmovdqa		%ymm6,96(%rdi)
-vmovdqa		%ymm7,128(%rdi)
-vmovdqa		%ymm8,160(%rdi)
-vmovdqa		%ymm9,192(%rdi)
-vmovdqa		%ymm10,224(%rdi)
+vmovdqa	%ymm3,  (\offset +  0)(%rdi)
+vmovdqa	%ymm4,  (\offset + 32)(%rdi)
+vmovdqa	%ymm5,  (\offset + 64)(%rdi)
+vmovdqa	%ymm6,  (\offset + 96)(%rdi)
+vmovdqa	%ymm7,  (\offset +128)(%rdi)
+vmovdqa	%ymm8,  (\offset +160)(%rdi)
+vmovdqa	%ymm9,  (\offset +192)(%rdi)
+vmovdqa	%ymm10, (\offset +224)(%rdi)
+.endm
+
+.text
+.global MLK_ASM_NAMESPACE(tomont_avx2)
+.balign 4
+MLK_ASM_FN_SYMBOL(tomont_avx2)
+#consts
+vmovdqa		MLK_AVX2_BACKEND_DATA_OFFSET_16XQ*2(%rsi),%ymm0
+vmovdqa		MLK_AVX2_BACKEND_DATA_OFFSET_16XMONTSQLO*2(%rsi),%ymm1
+vmovdqa		MLK_AVX2_BACKEND_DATA_OFFSET_16XMONTSQHI*2(%rsi),%ymm2
+
+tomont_128_coefficients 0
+tomont_128_coefficients 2*128
 
 ret
 

mlkem/src/native/x86_64/src/tomont.S

@@ -18,7 +18,8 @@
  * Changes:
  * - Add call to csub in reduce128_avx to produce outputs
  *   in [0,1,...,q-1] rather than [0,1,...,q], matching the
- *   semantics of mlk_poly_reduce().
+ *   semantics of mlk_poly_reduce(),
+ * - Use a macro instead of a local function call.
  */
 
 #include "../../../common.h"
@@ -40,14 +41,6 @@ MLK_ASM_FN_SYMBOL(tomont_avx2)
         vmovdqa	(%rsi), %ymm0
         vmovdqa	0xa0(%rsi), %ymm1
         vmovdqa	0xc0(%rsi), %ymm2
-        callq	Ltomont_avx2_core
-        addq	$0x100, %rdi            # imm = 0x100
-        callq	Ltomont_avx2_core
-        retq
-        .cfi_endproc
-
-Ltomont_avx2_core:
-        .cfi_startproc
         vmovdqa	(%rdi), %ymm3
         vmovdqa	0x20(%rdi), %ymm4
         vmovdqa	0x40(%rdi), %ymm5
@@ -96,6 +89,54 @@ Ltomont_avx2_core:
         vmovdqa	%ymm8, 0xa0(%rdi)
         vmovdqa	%ymm9, 0xc0(%rdi)
         vmovdqa	%ymm10, 0xe0(%rdi)
+        vmovdqa	0x100(%rdi), %ymm3
+        vmovdqa	0x120(%rdi), %ymm4
+        vmovdqa	0x140(%rdi), %ymm5
+        vmovdqa	0x160(%rdi), %ymm6
+        vmovdqa	0x180(%rdi), %ymm7
+        vmovdqa	0x1a0(%rdi), %ymm8
+        vmovdqa	0x1c0(%rdi), %ymm9
+        vmovdqa	0x1e0(%rdi), %ymm10
+        vpmullw	%ymm1, %ymm3, %ymm11
+        vpmulhw	%ymm2, %ymm3, %ymm3
+        vpmulhw	%ymm0, %ymm11, %ymm11
+        vpsubw	%ymm11, %ymm3, %ymm3
+        vpmullw	%ymm1, %ymm4, %ymm12
+        vpmulhw	%ymm2, %ymm4, %ymm4
+        vpmulhw	%ymm0, %ymm12, %ymm12
+        vpsubw	%ymm12, %ymm4, %ymm4
+        vpmullw	%ymm1, %ymm5, %ymm13
+        vpmulhw	%ymm2, %ymm5, %ymm5
+        vpmulhw	%ymm0, %ymm13, %ymm13
+        vpsubw	%ymm13, %ymm5, %ymm5
+        vpmullw	%ymm1, %ymm6, %ymm14
+        vpmulhw	%ymm2, %ymm6, %ymm6
+        vpmulhw	%ymm0, %ymm14, %ymm14
+        vpsubw	%ymm14, %ymm6, %ymm6
+        vpmullw	%ymm1, %ymm7, %ymm15
+        vpmulhw	%ymm2, %ymm7, %ymm7
+        vpmulhw	%ymm0, %ymm15, %ymm15
+        vpsubw	%ymm15, %ymm7, %ymm7
+        vpmullw	%ymm1, %ymm8, %ymm11
+        vpmulhw	%ymm2, %ymm8, %ymm8
+        vpmulhw	%ymm0, %ymm11, %ymm11
+        vpsubw	%ymm11, %ymm8, %ymm8
+        vpmullw	%ymm1, %ymm9, %ymm12
+        vpmulhw	%ymm2, %ymm9, %ymm9
+        vpmulhw	%ymm0, %ymm12, %ymm12
+        vpsubw	%ymm12, %ymm9, %ymm9
+        vpmullw	%ymm1, %ymm10, %ymm13
+        vpmulhw	%ymm2, %ymm10, %ymm10
+        vpmulhw	%ymm0, %ymm13, %ymm13
+        vpsubw	%ymm13, %ymm10, %ymm10
+        vmovdqa	%ymm3, 0x100(%rdi)
+        vmovdqa	%ymm4, 0x120(%rdi)
+        vmovdqa	%ymm5, 0x140(%rdi)
+        vmovdqa	%ymm6, 0x160(%rdi)
+        vmovdqa	%ymm7, 0x180(%rdi)
+        vmovdqa	%ymm8, 0x1a0(%rdi)
+        vmovdqa	%ymm9, 0x1c0(%rdi)
+        vmovdqa	%ymm10, 0x1e0(%rdi)
         retq
         .cfi_endproc