Use global optimization flag, use AArch64 ASM only on AArch64 system

Signed-off-by: Hanno Becker <[email protected]>
pq-code-package · Sep 12, 2024 · deddf53 · deddf53
1 parent c326bb5
commit deddf53
Show file tree

Hide file tree

Showing 10 changed files with 52 additions and 13 deletions.
diff --git a/mk/config.mk b/mk/config.mk
@@ -69,7 +69,7 @@ endif
 RNG ?=
 BENCH :=
 CYCLES ?=
-OPT ?= REF # (one of REF, AARCH64)
+OPT ?= 1
 RETAINED_VARS := RNG BENCH CYCLES OPT
 
 BUILD_DIR := test/build

diff --git a/mk/schemes.mk b/mk/schemes.mk
@@ -1,8 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 SOURCES = $(wildcard mlkem/*.c)
-ifeq ($(OPT),AARCH64)
+ifeq ($(OPT),1)
 	SOURCES += $(wildcard mlkem/asm/aarch64/*.S)
-	CPPFLAGS += -DMLKEM_OPT_AARCH64
+	CPPFLAGS += -DMLKEM_USE_ASM
 endif
 
 CPPFLAGS += -Imlkem -Imlkem/sys

diff --git a/mlkem/asm/aarch64/macro.S → mlkem/asm/aarch64/common.i b/mlkem/asm/aarch64/macro.S → mlkem/asm/aarch64/common.i
diff --git a/mlkem/asm/aarch64/intt_123_4567.S b/mlkem/asm/aarch64/intt_123_4567.S
@@ -23,8 +23,11 @@
 /// SOFTWARE.
 ///
 
+#include "config.h"
+#if defined(MLKEM_USE_AARCH64_ASM)
+
 // Needed to provide ASM_LOAD directive
-#include <hal_env.h>
+#include "common.i"
 
 .macro mulmodq dst, src, const, idx0, idx1
         sqrdmulh t2.8h,      \src\().8h, \const\().h[\idx1\()]
@@ -354,3 +357,5 @@ layer123_start:
 
         pop_stack
         ret
+
+#endif /* MLKEM_USE_AARCH64_ASM */
diff --git a/mlkem/asm/aarch64/intt_kyber_123_45_67_twiddles.S b/mlkem/asm/aarch64/intt_kyber_123_45_67_twiddles.S
@@ -23,6 +23,9 @@
 /// SOFTWARE.
 ///
 
+#include "config.h"
+#if defined(MLKEM_USE_AARCH64_ASM)
+
 roots_l56:
 .short -910
 .short -910
@@ -491,3 +494,5 @@ roots_l012:
 .short 6762
 .short 0
 .short 0
+
+#endif /* MLKEM_USE_AARCH64_ASM */
diff --git a/mlkem/asm/aarch64/ntt_123_4567.S b/mlkem/asm/aarch64/ntt_123_4567.S
@@ -23,8 +23,11 @@
 /// SOFTWARE.
 ///
 
+#include "config.h"
+#if defined(MLKEM_USE_AARCH64_ASM)
+
 // Needed to provide ASM_LOAD directive
-#include <hal_env.h>
+#include "common.i"
 
 .macro mulmodq dst, src, const, idx0, idx1
         sqrdmulh t2.8h,   \src\().8h, \const\().h[\idx1]
@@ -311,3 +314,5 @@ layer4567_start:
 
         pop_stack
         ret
+
+#endif /* MLKEM_USE_AARCH64_ASM */
diff --git a/mlkem/asm/aarch64/ntt_kyber_123_45_67_twiddles.S b/mlkem/asm/aarch64/ntt_kyber_123_45_67_twiddles.S
@@ -23,6 +23,9 @@
 /// SOFTWARE.
 ///
 
+#include "config.h"
+#if defined(MLKEM_USE_AARCH64_ASM)
+
 .p2align 2
 roots_l012:
 .short -1600
@@ -491,3 +494,5 @@ roots_l56:
 .short -3878
 .short -11566
 .short -11566
+
+#endif /* MLKEM_USE_AARCH64_ASM */
diff --git a/mlkem/asm/asm.h b/mlkem/asm/asm.h
@@ -4,10 +4,11 @@
 
 #include <stdint.h>
 #include "params.h"
+#include "config.h"
 
-#ifdef MLKEM_OPT_AARCH64
+#ifdef MLKEM_USE_AARCH64_ASM
 void ntt_kyber_123_4567(int16_t *);
 void intt_kyber_123_4567(int16_t *);
-#endif
+#endif /* MLKEM_USE_AARCH64_ASM */
 
 #endif
diff --git a/mlkem/ntt.c b/mlkem/ntt.c
@@ -81,9 +81,9 @@ static int16_t fqmul(int16_t a, int16_t b)
  **************************************************/
 void ntt(int16_t r[256])
 {
-    #ifdef MLKEM_OPT_AARCH64
+    #ifdef MLKEM_USE_AARCH64_ASM
     ntt_kyber_123_4567(r);
-    #else /* OPT_AARCH64 */
+    #else /* MLKEM_USE_AARCH64_ASM */
     unsigned int len, start, j, k;
     int16_t t, zeta;
 
@@ -101,7 +101,7 @@ void ntt(int16_t r[256])
             }
         }
     }
-    #endif /* OPT_AARCH64 */
+    #endif /* MLKEM_USE_AARCH64_ASM */
 }
 
 /*************************************************
@@ -116,9 +116,9 @@ void ntt(int16_t r[256])
  **************************************************/
 void invntt(int16_t r[256])
 {
-    #ifdef MLKEM_OPT_AARCH64
+    #ifdef MLKEM_USE_AARCH64_ASM
     intt_kyber_123_4567(r);
-    #else /* OPT_AARCH64 */
+    #else /* MLKEM_USE_AARCH64_ASM */
     unsigned int start, len, j, k;
     int16_t t, zeta;
     const int16_t f = 1441; // mont^2/128
@@ -143,7 +143,7 @@ void invntt(int16_t r[256])
     {
         r[j] = fqmul(r[j], f);
     }
-    #endif /* OPT_AARCH64 */
+    #endif /* MLKEM_USE_AARCH64_ASM */
 }
 
 /*************************************************

diff --git a/mlkem/sys/config.h b/mlkem/sys/config.h
@@ -0,0 +1,18 @@
+// SPDX-License-Identifier: Apache-2.0
+
+#ifndef CONFIG_H
+#define CONFIG_H
+
+#include "cpucap.h"
+
+#if defined(MLKEM_USE_ASM)
+
+#if defined(SYS_AARCH64)
+#define MLKEM_USE_AARCH64_ASM
+#else /* SYS_AARCH64 */
+/* Check x86_64 at some point */
+#warning "Selected optimized build, but no platform-specific assembly present"
+#endif /* SYS_AARCH64 */
+
+#endif /* MLKEM_USE_ASM */
+#endif /* CONFIG_H */