mlkem_native.h: Introduce MLK_TOTAL_ALLOC constants

For consumers implementing custom allocations (in cases where the stack is
not large enough), it may be important to know how much space has to be
reserved for allocations of large structures.

This commit adds per-operation allocation constants:
 - MLK_TOTAL_ALLOC_{512,768,1024}_{KEYPAIR,ENCAPS,DECAPS} for each parameter
   set and operation, indicating peak memory consumption.
 - MLK_MAX_TOTAL_ALLOC_{KEYPAIR,ENCAPS,DECAPS} for maximum across parameter
   sets per operation.
 - MLK_TOTAL_ALLOC_{512,768,1024} for maximum across operations per parameter
   set.
 - MLK_MAX_TOTAL_ALLOC for maximum across all parameter sets and operations.

The alloc_test is extended to check that the constants indeed match the peak
memory consumption of the corresponding functions.

 - Ported from pq-code-package/mldsa-native#850
 - Resolves #1442

Signed-off-by: Matthias J. Kannwischer <matthias@kannwischer.eu>

Author: mkannwischer

.github/workflows/baremetal.yml

@@ -46,4 +46,4 @@ jobs:
           acvp: ${{ matrix.target.acvp }}
           examples: false
           stack: false
-          alloc: false
+          alloc: true

mlkem/mlkem_native.S

@@ -168,6 +168,23 @@
 #undef MLK_ERR_FAIL
 #undef MLK_ERR_OUT_OF_MEMORY
 #undef MLK_H
+#undef MLK_MAX3_
+#undef MLK_MAX_TOTAL_ALLOC
+#undef MLK_MAX_TOTAL_ALLOC_DECAPS
+#undef MLK_MAX_TOTAL_ALLOC_ENCAPS
+#undef MLK_MAX_TOTAL_ALLOC_KEYPAIR
+#undef MLK_TOTAL_ALLOC_1024
+#undef MLK_TOTAL_ALLOC_1024_DECAPS
+#undef MLK_TOTAL_ALLOC_1024_ENCAPS
+#undef MLK_TOTAL_ALLOC_1024_KEYPAIR
+#undef MLK_TOTAL_ALLOC_512
+#undef MLK_TOTAL_ALLOC_512_DECAPS
+#undef MLK_TOTAL_ALLOC_512_ENCAPS
+#undef MLK_TOTAL_ALLOC_512_KEYPAIR
+#undef MLK_TOTAL_ALLOC_768
+#undef MLK_TOTAL_ALLOC_768_DECAPS
+#undef MLK_TOTAL_ALLOC_768_ENCAPS
+#undef MLK_TOTAL_ALLOC_768_KEYPAIR
 #undef crypto_kem_check_pk
 #undef crypto_kem_check_sk
 #undef crypto_kem_dec

mlkem/mlkem_native.c

@@ -157,6 +157,23 @@
 #undef MLK_ERR_FAIL
 #undef MLK_ERR_OUT_OF_MEMORY
 #undef MLK_H
+#undef MLK_MAX3_
+#undef MLK_MAX_TOTAL_ALLOC
+#undef MLK_MAX_TOTAL_ALLOC_DECAPS
+#undef MLK_MAX_TOTAL_ALLOC_ENCAPS
+#undef MLK_MAX_TOTAL_ALLOC_KEYPAIR
+#undef MLK_TOTAL_ALLOC_1024
+#undef MLK_TOTAL_ALLOC_1024_DECAPS
+#undef MLK_TOTAL_ALLOC_1024_ENCAPS
+#undef MLK_TOTAL_ALLOC_1024_KEYPAIR
+#undef MLK_TOTAL_ALLOC_512
+#undef MLK_TOTAL_ALLOC_512_DECAPS
+#undef MLK_TOTAL_ALLOC_512_ENCAPS
+#undef MLK_TOTAL_ALLOC_512_KEYPAIR
+#undef MLK_TOTAL_ALLOC_768
+#undef MLK_TOTAL_ALLOC_768_DECAPS
+#undef MLK_TOTAL_ALLOC_768_ENCAPS
+#undef MLK_TOTAL_ALLOC_768_KEYPAIR
 #undef crypto_kem_check_pk
 #undef crypto_kem_check_sk
 #undef crypto_kem_dec

mlkem/mlkem_native.h

@@ -428,4 +428,65 @@ int MLK_API_NAMESPACE(check_sk)(
 #endif /* MLK_CONFIG_API_NO_SUPERCOP */
 #endif /* !MLK_CONFIG_API_CONSTANTS_ONLY */
 
+
+/***************************** Memory Usage **********************************/
+
+/*
+ * By default mlkem-native performs all memory allocations on the stack.
+ * Alternatively, mlkem-native supports custom allocation of large structures
+ * through the `MLK_CONFIG_CUSTOM_ALLOC_FREE` configuration option.
+ * See mlkem_native_config.h for details.
+ *
+ * `MLK_TOTAL_ALLOC_{512,768,1024}_{KEYPAIR,ENCAPS,DECAPS}` indicates the
+ * maximum (accumulative) allocation via MLK_ALLOC for each parameter set and
+ * operation. Note that some stack allocation remains even when using custom
+ * allocators, so these values are lower than total stack usage with the default
+ * stack-only allocation.
+ *
+ * These constants may be used to implement custom allocations using a
+ * fixed-sized buffer and a simple allocator (e.g., bump allocator).
+ */
+/* check-magic: off */
+#define MLK_TOTAL_ALLOC_512_KEYPAIR 10048
+#define MLK_TOTAL_ALLOC_512_ENCAPS 8384
+#define MLK_TOTAL_ALLOC_512_DECAPS 9152
+#define MLK_TOTAL_ALLOC_768_KEYPAIR 15552
+#define MLK_TOTAL_ALLOC_768_ENCAPS 13248
+#define MLK_TOTAL_ALLOC_768_DECAPS 14336
+#define MLK_TOTAL_ALLOC_1024_KEYPAIR 22400
+#define MLK_TOTAL_ALLOC_1024_ENCAPS 19136
+#define MLK_TOTAL_ALLOC_1024_DECAPS 20704
+/* check-magic: on */
+
+/*
+ * `MLK_MAX_TOTAL_ALLOC_{KEYPAIR,ENCAPS,DECAPS}` is the maximum across all
+ * parameter sets for each operation.
+ * `MLK_MAX_TOTAL_ALLOC` is the maximum across all parameter sets and
+ * operations.
+ */
+#define MLK_MAX_TOTAL_ALLOC_KEYPAIR MLK_TOTAL_ALLOC_1024_KEYPAIR
+#define MLK_MAX_TOTAL_ALLOC_ENCAPS MLK_TOTAL_ALLOC_1024_ENCAPS
+#define MLK_MAX_TOTAL_ALLOC_DECAPS MLK_TOTAL_ALLOC_1024_DECAPS
+
+#define MLK_MAX3_(a, b, c) \
+  ((a) > (b) ? ((a) > (c) ? (a) : (c)) : ((b) > (c) ? (b) : (c)))
+
+/*
+ * `MLK_TOTAL_ALLOC_{512,768,1024}` is the maximum across all operations for
+ * each parameter set.
+ * `MLK_MAX_TOTAL_ALLOC` is the maximum across all parameter sets and
+ * operations.
+ */
+#define MLK_TOTAL_ALLOC_512                                          \
+  MLK_MAX3_(MLK_TOTAL_ALLOC_512_KEYPAIR, MLK_TOTAL_ALLOC_512_ENCAPS, \
+            MLK_TOTAL_ALLOC_512_DECAPS)
+#define MLK_TOTAL_ALLOC_768                                          \
+  MLK_MAX3_(MLK_TOTAL_ALLOC_768_KEYPAIR, MLK_TOTAL_ALLOC_768_ENCAPS, \
+            MLK_TOTAL_ALLOC_768_DECAPS)
+#define MLK_TOTAL_ALLOC_1024                                           \
+  MLK_MAX3_(MLK_TOTAL_ALLOC_1024_KEYPAIR, MLK_TOTAL_ALLOC_1024_ENCAPS, \
+            MLK_TOTAL_ALLOC_1024_DECAPS)
+
+#define MLK_MAX_TOTAL_ALLOC MLK_TOTAL_ALLOC_1024
+
 #endif /* !MLK_H */

scripts/autogen

@@ -2019,6 +2019,10 @@ def check_macro_typos():
         if m in macros:
             return True
 
+        # Ignore alloc macros only defined in mlkem_native.h
+        if m.startswith("MLK_TOTAL_ALLOC") or m.startswith("MLK_MAX_TOTAL_ALLOC"):
+            return True
+
         is_autogen = filename == "scripts/autogen"
 
         # Exclude system capabilities, which are enum values