AVX2: Add native implementation of poly_chknorm

This adds the AVX2 intrinsics implementation of poly_chknorm from
https://github.com/pq-crystals/dilithium/blob/master/avx2/poly.c.

Signed-off-by: jammychiou1 <[email protected]>

Author: jammychiou1

BIBLIOGRAPHY.md

@@ -149,6 +149,7 @@ source code and documentation.
   - [mldsa/native/x86_64/src/ntt.S](mldsa/native/x86_64/src/ntt.S)
   - [mldsa/native/x86_64/src/nttunpack.S](mldsa/native/x86_64/src/nttunpack.S)
   - [mldsa/native/x86_64/src/poly_caddq_avx2.c](mldsa/native/x86_64/src/poly_caddq_avx2.c)
+  - [mldsa/native/x86_64/src/poly_chknorm_avx2.c](mldsa/native/x86_64/src/poly_chknorm_avx2.c)
   - [mldsa/native/x86_64/src/poly_decompose_32_avx2.c](mldsa/native/x86_64/src/poly_decompose_32_avx2.c)
   - [mldsa/native/x86_64/src/poly_decompose_88_avx2.c](mldsa/native/x86_64/src/poly_decompose_88_avx2.c)
   - [mldsa/native/x86_64/src/poly_use_hint_32_avx2.c](mldsa/native/x86_64/src/poly_use_hint_32_avx2.c)

mldsa/native/api.h

@@ -254,4 +254,20 @@ static MLD_INLINE void mld_poly_use_hint_88_native(int32_t *b, const int32_t *a,
                                                    const int32_t *h);
 #endif /* MLD_USE_NATIVE_POLY_USE_HINT_88 */
 
+#if defined(MLD_USE_NATIVE_POLY_CHKNORM)
+/*************************************************
+ * Name:        mld_poly_chknorm_native
+ *
+ * Description: Check infinity norm of polynomial against given bound.
+ *              Assumes input coefficients were reduced by mld_reduce32().
+ *
+ * Arguments:   - const int32_t *a: pointer to polynomial
+ *              - int32_t B: norm bound
+ *
+ * Returns 0 if norm is strictly smaller than B <= (MLDSA_Q-1)/8 and 0xFFFFFFFF
+ * otherwise.
+ **************************************************/
+static MLD_INLINE uint32_t mld_poly_chknorm_native(const int32_t *a, int32_t B);
+#endif /* MLD_USE_NATIVE_POLY_CHKNORM */
+
 #endif /* !MLD_NATIVE_API_H */

mldsa/native/x86_64/meta.h

@@ -22,6 +22,7 @@
 #define MLD_USE_NATIVE_POLY_CADDQ
 #define MLD_USE_NATIVE_POLY_USE_HINT_32
 #define MLD_USE_NATIVE_POLY_USE_HINT_88
+#define MLD_USE_NATIVE_POLY_CHKNORM
 
 #if !defined(__ASSEMBLER__)
 #include <string.h>
@@ -133,6 +134,11 @@ static MLD_INLINE void mld_poly_use_hint_88_native(int32_t *b, const int32_t *a,
                             (const __m256i *)h);
 }
 
+static MLD_INLINE uint32_t mld_poly_chknorm_native(const int32_t *a, int32_t B)
+{
+  return mld_poly_chknorm_avx2((const __m256i *)a, B);
+}
+
 #endif /* !__ASSEMBLER__ */
 
 #endif /* !MLD_NATIVE_X86_64_META_H */

mldsa/native/x86_64/src/arith_native_x86_64.h

@@ -69,4 +69,7 @@ void mld_poly_use_hint_32_avx2(__m256i *b, const __m256i *a, const __m256i *h);
 #define mld_poly_use_hint_88_avx2 MLD_NAMESPACE(mld_poly_use_hint_88_avx2)
 void mld_poly_use_hint_88_avx2(__m256i *b, const __m256i *a, const __m256i *h);
 
+#define mld_poly_chknorm_avx2 MLD_NAMESPACE(mld_poly_chknorm_avx2)
+uint32_t mld_poly_chknorm_avx2(const __m256i *a, int32_t B);
+
 #endif /* !MLD_NATIVE_X86_64_SRC_ARITH_NATIVE_X86_64_H */

mldsa/native/x86_64/src/poly_chknorm_avx2.c

@@ -0,0 +1,53 @@
+/*
+ * Copyright (c) The mldsa-native project authors
+ * SPDX-License-Identifier: Apache-2.0 OR ISC OR MIT
+ */
+
+/* References
+ * ==========
+ *
+ * - [REF_AVX2]
+ *   CRYSTALS-Dilithium optimized AVX2 implementation
+ *   Bai, Ducas, Kiltz, Lepoint, Lyubashevsky, Schwabe, Seiler, Stehlé
+ *   https://github.com/pq-crystals/dilithium/tree/master/avx2
+ */
+
+/*
+ * This file is derived from the public domain
+ * AVX2 Dilithium implementation @[REF_AVX2].
+ */
+
+#include "../../../common.h"
+
+#if defined(MLD_ARITH_BACKEND_X86_64_DEFAULT) && \
+    !defined(MLD_CONFIG_MULTILEVEL_NO_SHARED)
+
+#include <immintrin.h>
+#include <stdint.h>
+#include "arith_native_x86_64.h"
+
+uint32_t mld_poly_chknorm_avx2(const __m256i *a, int32_t B)
+{
+  unsigned int i;
+  __m256i f, t;
+  const __m256i bound = _mm256_set1_epi32(B - 1);
+
+  t = _mm256_setzero_si256();
+  for (i = 0; i < MLDSA_N / 8; i++)
+  {
+    f = _mm256_load_si256(&a[i]);
+    f = _mm256_abs_epi32(f);
+    f = _mm256_cmpgt_epi32(f, bound);
+    t = _mm256_or_si256(t, f);
+  }
+
+  return _mm256_testz_si256(t, t) - 1;
+}
+
+#else /* MLD_ARITH_BACKEND_X86_64_DEFAULT && !MLD_CONFIG_MULTILEVEL_NO_SHARED \
+       */
+
+MLD_EMPTY_CU(avx2_poly_chknorm)
+
+#endif /* !(MLD_ARITH_BACKEND_X86_64_DEFAULT && \
+          !MLD_CONFIG_MULTILEVEL_NO_SHARED) */

mldsa/poly.c

@@ -324,6 +324,11 @@ void mld_poly_use_hint(mld_poly *b, const mld_poly *a, const mld_poly *h)
 MLD_INTERNAL_API
 uint32_t mld_poly_chknorm(const mld_poly *a, int32_t B)
 {
+#if defined(MLD_USE_NATIVE_POLY_CHKNORM)
+  /* TODO: proof */
+  mld_assert_bound(a->coeffs, MLDSA_N, -REDUCE32_RANGE_MAX, REDUCE32_RANGE_MAX);
+  return mld_poly_chknorm_native(a->coeffs, B);
+#else
   unsigned int i;
   uint32_t t = 0;
   mld_assert_bound(a->coeffs, MLDSA_N, -REDUCE32_RANGE_MAX, REDUCE32_RANGE_MAX);
@@ -346,6 +351,7 @@ uint32_t mld_poly_chknorm(const mld_poly *a, int32_t B)
   }
 
   return t;
+#endif /* !MLD_USE_NATIVE_POLY_CHKNORM */
 }
 
 /*************************************************