Move AArch64 ASM into asm/aarch64/*, rename OPT values

Use OPT=AARCH64 at the toplevel to enable AArch64 optimizations. Refined (internal) configuration options, e.g. for switching between NTT implementations, to follow. Signed-off-by: Hanno Becker <[email protected]>
pq-code-package · Sep 9, 2024 · 807cf4e · 807cf4e
1 parent 53d884f
commit 807cf4e
Show file tree

Hide file tree

Showing 9 changed files with 26 additions and 16 deletions.
diff --git a/mk/schemes.mk b/mk/schemes.mk
@@ -1,8 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 SOURCES = $(wildcard mlkem/*.c)
-ifneq ($(OPT),REF)
-	SOURCES += $(wildcard mlkem/asm/*.S)
-	CPPFLAGS += -D$(OPT)
+ifeq ($(OPT),AARCH64)
+	SOURCES += $(wildcard mlkem/asm/aarch64/*.S)
+	CPPFLAGS += -DMLKEM_OPT_AARCH64
 endif
 
 CPPFLAGS += -Imlkem

diff --git a/mlkem/asm/intt_123_4567.S → mlkem/asm/aarch64/intt_123_4567.S b/mlkem/asm/intt_123_4567.S → mlkem/asm/aarch64/intt_123_4567.S
diff --git a/mlkem/asm/intt_kyber_123_45_67_twiddles.S → ...m/aarch64/intt_kyber_123_45_67_twiddles.S b/mlkem/asm/intt_kyber_123_45_67_twiddles.S → ...m/aarch64/intt_kyber_123_45_67_twiddles.S
diff --git a/mlkem/asm/macro.S → mlkem/asm/aarch64/macro.S b/mlkem/asm/macro.S → mlkem/asm/aarch64/macro.S
diff --git a/mlkem/asm/ntt_123_4567.S → mlkem/asm/aarch64/ntt_123_4567.S b/mlkem/asm/ntt_123_4567.S → mlkem/asm/aarch64/ntt_123_4567.S
diff --git a/mlkem/asm/ntt_kyber_123_45_67_twiddles.S → ...sm/aarch64/ntt_kyber_123_45_67_twiddles.S b/mlkem/asm/ntt_kyber_123_45_67_twiddles.S → ...sm/aarch64/ntt_kyber_123_45_67_twiddles.S
diff --git a/mlkem/asm/asm.h b/mlkem/asm/asm.h
@@ -0,0 +1,13 @@
+// SPDX-License-Identifier: Apache-2.0
+#ifndef ASM_H
+#define ASM_H
+
+#include <stdint.h>
+#include "params.h"
+
+#ifdef MLKEM_OPT_AARCH64
+void ntt_kyber_123_4567(int16_t *);
+void intt_kyber_123_4567(int16_t *);
+#endif
+
+#endif
diff --git a/mlkem/ntt.c b/mlkem/ntt.c
@@ -4,6 +4,8 @@
 #include "reduce.h"
 #include <stdint.h>
 
+#include "asm/asm.h"
+
 /* Code to generate zetas and zetas_inv used in the number-theoretic transform:
 
 #define KYBER_ROOT_OF_UNITY 17
@@ -37,11 +39,6 @@ void init_ntt() {
 }
 */
 
-#ifdef NTT123_4567
-void ntt_kyber_123_4567(int16_t *);
-void intt_kyber_123_4567(int16_t *);
-#endif
-
 const int16_t zetas[128] =
 {
     -1044, -758,  -359,  -1517, 1493,  1422,  287,   202,  -171,  622,   1577,
@@ -84,9 +81,9 @@ static int16_t fqmul(int16_t a, int16_t b)
  **************************************************/
 void ntt(int16_t r[256])
 {
-    #ifdef NTT123_4567
+    #ifdef MLKEM_OPT_AARCH64
     ntt_kyber_123_4567(r);
-    #else
+    #else /* OPT_AARCH64 */
     unsigned int len, start, j, k;
     int16_t t, zeta;
 
@@ -104,7 +101,7 @@ void ntt(int16_t r[256])
             }
         }
     }
-    #endif
+    #endif /* OPT_AARCH64 */
 }
 
 /*************************************************
@@ -119,9 +116,9 @@ void ntt(int16_t r[256])
  **************************************************/
 void invntt(int16_t r[256])
 {
-    #ifdef NTT123_4567
+    #ifdef MLKEM_OPT_AARCH64
     intt_kyber_123_4567(r);
-    #else
+    #else /* OPT_AARCH64 */
     unsigned int start, len, j, k;
     int16_t t, zeta;
     const int16_t f = 1441; // mont^2/128
@@ -146,7 +143,7 @@ void invntt(int16_t r[256])
     {
         r[j] = fqmul(r[j], f);
     }
-    #endif
+    #endif /* OPT_AARCH64 */
 }
 
 /*************************************************

diff --git a/scripts/tests b/scripts/tests
@@ -292,7 +292,7 @@ _shared_options = [
     ),
     click.option(
         "--opt",
-        type=click.Choice(["REF", "NTT123_4567"], case_sensitive=False),
+        type=click.Choice(["REF", "AARCH64"], case_sensitive=False),
         help="Choose optimized version",
     ),
 ]
@@ -455,7 +455,7 @@ def bench(
         extra_make_envs=process_make_envs(cflags, arch_flags),
         extra_make_args=[
             f"CYCLES={cycles}",
-            *(f"OPT={opt}" if opt is not None else ""),
+            f"OPT={opt}" if opt is not None else "",
         ],
     )