Merge pull request #348 from pq-code-package/aarch64-rej-eta-implementation

Add native imlementation of rej_eta (AArch64 + AVX2)

Author: mkannwischer

mldsa/native/aarch64/meta.h

@@ -11,6 +11,8 @@
 #define MLD_USE_NATIVE_NTT
 #define MLD_USE_NATIVE_INTT
 #define MLD_USE_NATIVE_REJ_UNIFORM
+#define MLD_USE_NATIVE_REJ_UNIFORM_ETA2
+#define MLD_USE_NATIVE_REJ_UNIFORM_ETA4
 
 /* Identifier for this backend so that source and assembly files
  * in the build can be appropriately guarded. */
@@ -43,6 +45,34 @@ static MLD_INLINE int mld_rej_uniform_native(int32_t *r, unsigned len,
   return (int)mld_rej_uniform_asm(r, buf, buflen, mld_rej_uniform_table);
 }
 
+static MLD_INLINE int mld_rej_uniform_eta2_native(int32_t *r, unsigned len,
+                                                  const uint8_t *buf,
+                                                  unsigned buflen)
+{
+  /* AArch64 implementation assumes specific buffer lengths */
+  if (len != MLDSA_N || buflen != MLD_AARCH64_REJ_UNIFORM_ETA2_BUFLEN)
+  {
+    return -1;
+  }
+
+  return (int)mld_rej_uniform_eta2_asm(r, buf, buflen,
+                                       mld_rej_uniform_eta_table);
+}
+
+static MLD_INLINE int mld_rej_uniform_eta4_native(int32_t *r, unsigned len,
+                                                  const uint8_t *buf,
+                                                  unsigned buflen)
+{
+  /* AArch64 implementation assumes specific buffer lengths */
+  if (len != MLDSA_N || buflen != MLD_AARCH64_REJ_UNIFORM_ETA4_BUFLEN)
+  {
+    return -1;
+  }
+
+  return (int)mld_rej_uniform_eta4_asm(r, buf, buflen,
+                                       mld_rej_uniform_eta_table);
+}
+
 #endif /* !__ASSEMBLER__ */
 
 #endif /* !MLD_NATIVE_AARCH64_META_H */

mldsa/native/aarch64/src/arith_native_aarch64.h

@@ -27,6 +27,22 @@ extern const int32_t mld_aarch64_intt_zetas_layer78[];
 extern const int32_t mld_aarch64_intt_zetas_layer123456[];
 
 extern const uint8_t mld_rej_uniform_table[];
+extern const uint8_t mld_rej_uniform_eta_table[];
+
+
+/*
+ * Sampling 256 coefficients mod 15 using rejection sampling from 4 bits.
+ * Expected number of required bytes: (256 * (16/15))/2 = 136.5 bytes.
+ * We sample 1 block (=136 bytes) of SHAKE256_RATE output initially.
+ * Sampling 2 blocks initially results in slightly worse performance.
+ */
+#define MLD_AARCH64_REJ_UNIFORM_ETA2_BUFLEN (1 * 136)
+/*
+ * Sampling 256 coefficients mod 9 using rejection sampling from 4 bits.
+ * Expected number of required bytes: (256 * (16/9))/2 = 227.5 bytes.
+ * We sample 2 blocks (=272 bytes) of SHAKE256_RATE output initially.
+ */
+#define MLD_AARCH64_REJ_UNIFORM_ETA4_BUFLEN (2 * 136)
 
 #define mld_ntt_asm MLD_NAMESPACE(ntt_asm)
 void mld_ntt_asm(int32_t *, const int32_t *, const int32_t *);
@@ -38,4 +54,12 @@ void mld_intt_asm(int32_t *, const int32_t *, const int32_t *);
 uint64_t mld_rej_uniform_asm(int32_t *r, const uint8_t *buf, unsigned buflen,
                              const uint8_t *table);
 
+#define mld_rej_uniform_eta2_asm MLD_NAMESPACE(rej_uniform_eta2_asm)
+unsigned mld_rej_uniform_eta2_asm(int32_t *r, const uint8_t *buf,
+                                  unsigned buflen, const uint8_t *table);
+
+#define mld_rej_uniform_eta4_asm MLD_NAMESPACE(rej_uniform_eta4_asm)
+unsigned mld_rej_uniform_eta4_asm(int32_t *r, const uint8_t *buf,
+                                  unsigned buflen, const uint8_t *table);
+
 #endif /* !MLD_NATIVE_AARCH64_SRC_ARITH_NATIVE_AARCH64_H */