x86_64: Replace rej_uniform intrinsics with assembly

Replace the AVX2 intrinsics implementation of rej_uniform with
hand-written x86_64 assembly, resolving #926.

The assembly follows the same algorithmic structure as the intrinsics
version: load 32 bytes, vpermq to rearrange 64-bit lanes, vpshufb to
extract 8x 3-byte groups, mask to 23 bits, compare against MLDSA_Q,
then use the lookup table to compact valid coefficients.

Key design decisions:
- Table passed as parameter (consistent with aarch64 approach),
  avoiding external symbol references for simpasm compatibility
- All constants constructed from immediates (no .rodata section),
  enabling future HOL-Light formal verification
- Register name #defines with #undef cleanup for SCU builds
- CBMC contract on assembly function declaration (following mlkem-native)
- vzeroupper at function exit to avoid AVX-SSE transition penalties

Also adds poly_uniform to the component benchmark.

Signed-off-by: jakemas <jakemas@amazon.com>

Author: jakemas

BIBLIOGRAPHY.md

@@ -227,7 +227,7 @@ source code and documentation.
   - [dev/x86_64/src/poly_use_hint_88_avx2.c](dev/x86_64/src/poly_use_hint_88_avx2.c)
   - [dev/x86_64/src/polyz_unpack_17_avx2.c](dev/x86_64/src/polyz_unpack_17_avx2.c)
   - [dev/x86_64/src/polyz_unpack_19_avx2.c](dev/x86_64/src/polyz_unpack_19_avx2.c)
-  - [dev/x86_64/src/rej_uniform_avx2.c](dev/x86_64/src/rej_uniform_avx2.c)
+  - [dev/x86_64/src/rej_uniform_avx2.S](dev/x86_64/src/rej_uniform_avx2.S)
   - [dev/x86_64/src/rej_uniform_eta2_avx2.c](dev/x86_64/src/rej_uniform_eta2_avx2.c)
   - [dev/x86_64/src/rej_uniform_eta4_avx2.c](dev/x86_64/src/rej_uniform_eta4_avx2.c)
   - [mldsa/src/native/x86_64/src/intt.S](mldsa/src/native/x86_64/src/intt.S)
@@ -245,11 +245,12 @@ source code and documentation.
   - [mldsa/src/native/x86_64/src/poly_use_hint_88_avx2.c](mldsa/src/native/x86_64/src/poly_use_hint_88_avx2.c)
   - [mldsa/src/native/x86_64/src/polyz_unpack_17_avx2.c](mldsa/src/native/x86_64/src/polyz_unpack_17_avx2.c)
   - [mldsa/src/native/x86_64/src/polyz_unpack_19_avx2.c](mldsa/src/native/x86_64/src/polyz_unpack_19_avx2.c)
-  - [mldsa/src/native/x86_64/src/rej_uniform_avx2.c](mldsa/src/native/x86_64/src/rej_uniform_avx2.c)
+  - [mldsa/src/native/x86_64/src/rej_uniform_avx2.S](mldsa/src/native/x86_64/src/rej_uniform_avx2.S)
   - [mldsa/src/native/x86_64/src/rej_uniform_eta2_avx2.c](mldsa/src/native/x86_64/src/rej_uniform_eta2_avx2.c)
   - [mldsa/src/native/x86_64/src/rej_uniform_eta4_avx2.c](mldsa/src/native/x86_64/src/rej_uniform_eta4_avx2.c)
   - [proofs/hol_light/x86_64/mldsa/mldsa_intt.S](proofs/hol_light/x86_64/mldsa/mldsa_intt.S)
   - [proofs/hol_light/x86_64/mldsa/mldsa_ntt.S](proofs/hol_light/x86_64/mldsa/mldsa_ntt.S)
+  - [proofs/hol_light/x86_64/mldsa/mldsa_rej_uniform.S](proofs/hol_light/x86_64/mldsa/mldsa_rej_uniform.S)
 
 ### `Round3_Spec`
 

dev/x86_64/meta.h

@@ -80,7 +80,8 @@ static MLD_INLINE int mld_rej_uniform_native(int32_t *r, unsigned len,
   }
 
   /* Safety: outlen is at most MLDSA_N and, hence, this cast is safe. */
-  return (int)mld_rej_uniform_avx2(r, buf);
+  return (int)mld_rej_uniform_avx2(r, buf,
+                                   (const uint8_t *)mld_rej_uniform_table);
 }
 
 #if defined(MLD_CONFIG_MULTILEVEL_WITH_SHARED) || MLDSA_ETA == 2

dev/x86_64/src/arith_native_x86_64.h

@@ -65,7 +65,8 @@ void mld_nttunpack_avx2(int32_t *r);
 #define mld_rej_uniform_avx2 MLD_NAMESPACE(mld_rej_uniform_avx2)
 MLD_MUST_CHECK_RETURN_VALUE
 unsigned mld_rej_uniform_avx2(int32_t *r,
-                              const uint8_t buf[MLD_AVX2_REJ_UNIFORM_BUFLEN]);
+                              const uint8_t buf[MLD_AVX2_REJ_UNIFORM_BUFLEN],
+                              const uint8_t *table);
 
 #define mld_rej_uniform_eta2_avx2 MLD_NAMESPACE(mld_rej_uniform_eta2_avx2)
 MLD_MUST_CHECK_RETURN_VALUE

dev/x86_64/src/rej_uniform_avx2.S

@@ -0,0 +1,173 @@
+/*
+ * 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)
+/* simpasm: header-end */
+
+#define out %rdi
+#define in %rsi
+#define tab %rdx
+
+#define ctr %eax
+#define pos %ecx
+
+#define good %r8d
+#define cnt %r9d
+#define tmp %r10
+
+#define idx8 %ymm0
+#define mask %ymm1
+#define bound %ymm2
+#define data %ymm3
+#define cmp_result %ymm4
+
+        .text
+
+/*
+ * unsigned mld_rej_uniform_avx2(int32_t *r, const uint8_t *buf,
+ *                               const uint8_t *table)
+ *
+ * Rejection sampling of uniform polynomial coefficients.
+ * Extracts 23-bit values from a byte buffer and accepts those < MLDSA_Q.
+ *
+ * Arguments:   out (rdi): pointer to output coefficient array r (256 x int32_t)
+ *              in  (rsi): pointer to input byte buffer buf (840 bytes)
+ *              tab (rdx): pointer to rejection sampling lookup table (256x8 bytes)
+ *
+ * Returns:     ctr (eax): number of valid coefficients written to r
+ */
+        .balign 4
+        .global MLD_ASM_NAMESPACE(mld_rej_uniform_avx2)
+MLD_ASM_FN_SYMBOL(mld_rej_uniform_avx2)
+
+/*
+ * Construct the shuffle mask for extracting 8 x 23-bit values from 24 bytes.
+ *
+ * After vpermq with 0x94, the 32 loaded bytes are rearranged as:
+ *   Low  128 bits: bytes [0..15]  (original 64-bit lanes 0, 1)
+ *   High 128 bits: bytes [8..23]  (original 64-bit lanes 1, 2)
+ *
+ * vpshufb then picks 3-byte groups and zero-pads each to a 32-bit lane:
+ *   Low  half: {0,1,2,FF, 3,4,5,FF, 6,7,8,FF, 9,10,11,FF}
+ *   High half: {4,5,6,FF, 7,8,9,FF, 10,11,12,FF, 13,14,15,FF}
+ *
+ * This extracts 8 non-overlapping 3-byte windows from the first 24 input bytes.
+ */
+        movq    $0xFF050403FF020100, tmp
+        vmovq   tmp, %xmm0
+        movq    $0xFF0B0A09FF080706, tmp
+        vpinsrq $1, tmp, %xmm0, %xmm0
+        movq    $0xFF090807FF060504, tmp
+        vmovq   tmp, %xmm3
+        movq    $0xFF0F0E0DFF0C0B0A, tmp
+        vpinsrq $1, tmp, %xmm3, %xmm3
+        vinserti128 $1, %xmm3, idx8, idx8
+
+// Construct broadcast constants
+        movl    $0x7FFFFF, good
+        vmovd   good, %xmm1
+        vpbroadcastd %xmm1, mask              // mask: 23-bit extraction
+
+        movl    $8380417, good                 // MLDSA_Q
+        vmovd   good, %xmm2
+        vpbroadcastd %xmm2, bound             // bound: rejection threshold
+
+// Initialize counters
+        xorl    ctr, ctr                       // ctr = 0
+        xorl    pos, pos                       // pos = 0
+
+/*
+ * Main SIMD loop: process 24 input bytes into up to 8 coefficients
+ * per iteration. Loops while ctr <= MLDSA_N - 8 and pos <= BUFLEN - 32.
+ */
+mld_rej_uniform_avx2_loop:
+        cmpl    $248, ctr                      // MLDSA_N - 8
+        ja      mld_rej_uniform_avx2_scalar
+        cmpl    $808, pos                      // MLD_AVX2_REJ_UNIFORM_BUFLEN - 32
+        ja      mld_rej_uniform_avx2_scalar
+
+        vmovdqu (in, %rcx), data               // load 32 bytes from buf[pos]
+        addl    $24, pos                       // advance pos
+        vpermq  $0x94, data, data              // rearrange 64-bit lanes: [2,1,1,0]
+        vpshufb idx8, data, data               // extract 8 x 3-byte groups
+        vpand   mask, data, data               // mask to 23 bits
+
+        vpsubd  bound, data, cmp_result        // d - Q: negative if d < Q (valid)
+        vmovmskps cmp_result, good             // extract sign bits as 8-bit mask
+
+        popcntl good, cnt                      // count valid coefficients
+
+        vmovq   (tab, %r8, 8), %xmm4          // load permutation from table[good]
+        vpmovzxbd %xmm4, cmp_result            // zero-extend to 8 dword indices
+        vpermd  data, cmp_result, data          // compact valid coefficients to front
+
+        vmovdqu data, (out, %rax, 4)           // store at r[ctr]
+        addl    cnt, ctr                       // ctr += popcount(good)
+
+        jmp     mld_rej_uniform_avx2_loop
+
+/*
+ * Scalar fallback loop: process remaining bytes one coefficient at a time.
+ * Loops while ctr < MLDSA_N and pos <= BUFLEN - 3.
+ */
+mld_rej_uniform_avx2_scalar:
+        cmpl    $256, ctr                      // MLDSA_N
+        jae     mld_rej_uniform_avx2_done
+        cmpl    $837, pos                      // MLD_AVX2_REJ_UNIFORM_BUFLEN - 3
+        ja      mld_rej_uniform_avx2_done
+
+        movzwl  (in, %rcx), good               // load 2 bytes at buf[pos]
+        movzbl  2(in, %rcx), cnt               // load third byte
+        shll    $16, cnt
+        orl     cnt, good
+        andl    $0x7FFFFF, good                // mask to 23 bits
+        addl    $3, pos                        // advance pos
+
+        cmpl    $8380417, good                 // MLDSA_Q
+        jae     mld_rej_uniform_avx2_scalar    // reject if >= Q
+
+        movl    good, (out, %rax, 4)           // store valid coefficient
+        addl    $1, ctr                        // ctr++
+        jmp     mld_rej_uniform_avx2_scalar
+
+mld_rej_uniform_avx2_done:
+        vzeroupper
+        ret
+
+/* To facilitate single-compilation-unit (SCU) builds, undefine all macros.
+ * Don't modify by hand -- this is auto-generated by scripts/autogen. */
+#undef out
+#undef in
+#undef tab
+#undef ctr
+#undef pos
+#undef good
+#undef cnt
+#undef tmp
+#undef idx8
+#undef mask
+#undef bound
+#undef data
+#undef cmp_result
+
+/* simpasm: footer-start */
+#endif /* MLD_ARITH_BACKEND_X86_64_DEFAULT && !MLD_CONFIG_MULTILEVEL_NO_SHARED \
+        */

dev/x86_64/src/rej_uniform_avx2.c

@@ -90,7 +90,6 @@
 #include "src/native/x86_64/src/poly_use_hint_88_avx2.c"
 #include "src/native/x86_64/src/polyz_unpack_17_avx2.c"
 #include "src/native/x86_64/src/polyz_unpack_19_avx2.c"
-#include "src/native/x86_64/src/rej_uniform_avx2.c"
 #include "src/native/x86_64/src/rej_uniform_eta2_avx2.c"
 #include "src/native/x86_64/src/rej_uniform_eta4_avx2.c"
 #include "src/native/x86_64/src/rej_uniform_table.c"

mldsa/mldsa_native.c

@@ -86,6 +86,7 @@
 #include "src/native/x86_64/src/pointwise_acc_l4.S"
 #include "src/native/x86_64/src/pointwise_acc_l5.S"
 #include "src/native/x86_64/src/pointwise_acc_l7.S"
+#include "src/native/x86_64/src/rej_uniform_avx2.S"
 #endif /* MLD_SYS_X86_64 */
 #endif /* MLD_CONFIG_USE_NATIVE_BACKEND_ARITH */
 

mldsa/mldsa_native_asm.S

@@ -80,7 +80,8 @@ static MLD_INLINE int mld_rej_uniform_native(int32_t *r, unsigned len,
   }
 
   /* Safety: outlen is at most MLDSA_N and, hence, this cast is safe. */
-  return (int)mld_rej_uniform_avx2(r, buf);
+  return (int)mld_rej_uniform_avx2(r, buf,
+                                   (const uint8_t *)mld_rej_uniform_table);
 }
 
 #if defined(MLD_CONFIG_MULTILEVEL_WITH_SHARED) || MLDSA_ETA == 2

mldsa/src/native/x86_64/meta.h

@@ -65,7 +65,8 @@ void mld_nttunpack_avx2(int32_t *r);
 #define mld_rej_uniform_avx2 MLD_NAMESPACE(mld_rej_uniform_avx2)
 MLD_MUST_CHECK_RETURN_VALUE
 unsigned mld_rej_uniform_avx2(int32_t *r,
-                              const uint8_t buf[MLD_AVX2_REJ_UNIFORM_BUFLEN]);
+                              const uint8_t buf[MLD_AVX2_REJ_UNIFORM_BUFLEN],
+                              const uint8_t *table);
 
 #define mld_rej_uniform_eta2_avx2 MLD_NAMESPACE(mld_rej_uniform_eta2_avx2)
 MLD_MUST_CHECK_RETURN_VALUE