From 157c445764407da0f9ef5cdf8c4cb34e6f3ea7bd Mon Sep 17 00:00:00 2001
From: Hanno Becker <beckphan@amazon.co.uk>
Date: Mon, 9 Sep 2024 14:45:41 +0100
Subject: [PATCH] Add cpucap.h header and detect AArch64 systems

Also, allow `FORCE_AARCH64` to double-check that a system is recognized
as AArch64. Use this in all AArch64-based CI builds.

Signed-off-by: Hanno Becker <beckphan@amazon.co.uk>
---
 .github/workflows/bench.yml         |  6 +++---
 .github/workflows/bench_ec2_all.yml |  2 ++
 mk/schemes.mk                       |  2 +-
 mlkem/params.h                      |  2 ++
 mlkem/sys/cpucap.h                  | 16 ++++++++++++++++
 5 files changed, 24 insertions(+), 4 deletions(-)
 create mode 100644 mlkem/sys/cpucap.h

diff --git a/.github/workflows/bench.yml b/.github/workflows/bench.yml
index a477a027f..dfc9362a6 100644
--- a/.github/workflows/bench.yml
+++ b/.github/workflows/bench.yml
@@ -18,19 +18,19 @@ jobs:
           name: Arm Cortex-A72 (Raspberry Pi 4) benchmarks
           bench_pmu: PMU
           archflags: -mcpu=cortex-a72
-          cflags:
+          cflags: -DFORCE_AARCH64
           bench_extra_args:
         - system: rpi5
           name: Arm Cortex-A76 (Raspberry Pi 5) benchmarks
           bench_pmu: PERF
           archflags: "-mcpu=cortex-a76 -march=armv8.2-a"
-          cflags:
+          cflags: -DFORCE_AARCH64
           bench_extra_args:
         - system: a55
           name: Arm Cortex-A55 (Snapdragon 888) benchmarks
           bench_pmu: PERF
           archflags: "-mcpu=cortex-a55 -march=armv8.2-a"
-          cflags: -static
+          cflags: "-static -DFORCE_AARCH64"
           bench_extra_args: -w exec-on-a55
     runs-on: self-hosted-${{ matrix.target.system }}
     defaults:
diff --git a/.github/workflows/bench_ec2_all.yml b/.github/workflows/bench_ec2_all.yml
index 889fa81e2..f4eb50611 100644
--- a/.github/workflows/bench_ec2_all.yml
+++ b/.github/workflows/bench_ec2_all.yml
@@ -18,6 +18,7 @@ jobs:
       ec2_instance_type: t4g.small
       ec2_ami_id: ami-096ea6a12ea24a797
       archflags: -mcpu=cortex-a76 -march=armv8.2-a
+      cflags: -DFORCE_AARCH64
       store_results: 'true'
       name: Graviton2
     secrets: inherit
@@ -32,6 +33,7 @@ jobs:
       ec2_instance_type: c7g.medium
       ec2_ami_id: ami-096ea6a12ea24a797
       archflags: -mcpu=neoverse-v1 -march=armv8.4-a
+      cflags: -DFORCE_AARCH64
       store_results: 'true'
       name: Graviton3
     secrets: inherit
diff --git a/mk/schemes.mk b/mk/schemes.mk
index 24fc1f6c1..dd4d25321 100644
--- a/mk/schemes.mk
+++ b/mk/schemes.mk
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 SOURCES = $(wildcard mlkem/*.c)
 
-CPPFLAGS += -Imlkem
+CPPFLAGS += -Imlkem -Imlkem/sys
 TESTS = test_kyber bench_kyber gen_NISTKAT gen_KAT
 
 MLKEM512_DIR = $(BUILD_DIR)/mlkem512
diff --git a/mlkem/params.h b/mlkem/params.h
index 36c2a538c..951f95fe1 100644
--- a/mlkem/params.h
+++ b/mlkem/params.h
@@ -2,6 +2,8 @@
 #ifndef PARAMS_H
 #define PARAMS_H
 
+#include "cpucap.h"
+
 #ifndef KYBER_K
 #define KYBER_K 3   /* Change this for different security strengths */
 #endif
diff --git a/mlkem/sys/cpucap.h b/mlkem/sys/cpucap.h
new file mode 100644
index 000000000..22b9c8029
--- /dev/null
+++ b/mlkem/sys/cpucap.h
@@ -0,0 +1,16 @@
+// SPDX-License-Identifier: Apache-2.0
+
+#ifndef CPUCAP_H
+#define CPUCAP_H
+
+/* Check if we're running on an AArch64 system. _M_ARM64 is set by MSVC. */
+#if defined(__aarch64__) || defined(_M_ARM64)
+#define SYS_AARCH64
+#endif
+
+/* If FORCE_AARCH64 is set, assert that we're indeed on an AArch64 system. */
+#if defined(FORCE_AARCH64) && !defined(SYS_AARCH64)
+#error "FORCE_AARCH64 is set, but we don't seem to be on an AArch64 system.
+#endif
+
+#endif