Address review comments on PR #1014

Reviewer-requested cleanup for the x86_64 rej_uniform assembly and
HOL Light proof:

Contract tightening (dev and mldsa copies of arith_native_x86_64.h):
  - requires(memory_no_alias(buf, 840)) instead of
    memory_no_alias(buf, MLD_AVX2_REJ_UNIFORM_BUFLEN) so the literal
    matches the HOL Light spec exactly.
  - requires(table == (const uint8_t *)mld_rej_uniform_table) pinning
    the table to the exported rejection-sampling table, replacing the
    looser memory_no_alias(table, 256 * sizeof(uint64_t)).
  - Clarify sync comment.

vzeroupper removal: none of the other asm routines issue vzeroupper;
drop it from rej_uniform for consistency. This shifts the function
length by 3 bytes, so the HOL Light proof's nonoverlapping 246 / pc+245
references in mldsa_rej_uniform.ml become 243 / pc+242 accordingly, and
the two X86_STEPS_TAC invocations that stepped the vzeroupper byte are
removed. Bytecode regenerated via autogen --update-hol-light-bytecode.

Autogen plumbing: register rej_uniform_avx2_asm.S in the x86_64 HOL
Light asm joblist so the proofs/hol_light/x86_64/mldsa/ copy is
regenerated by scripts/autogen. Add gen_avx2_hol_light_rej_uniform_table
to regenerate proofs/hol_light/x86_64/proofs/mldsa_rej_uniform_table.ml
alongside the C/aarch64 lookup tables (matches mlkem-native's pattern).

Cross-reference comment in proofs/hol_light/x86_64/proofs/
rej_uniform_avx2_asm.ml pointing at the CBMC contract.

Proof runtime: ~5-6 min in the CI native build.

Signed-off-by: Jake Massimo <jakemas@amazon.com>
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: 2 people

dev/x86_64/src/arith_native_x86_64.h

@@ -74,16 +74,16 @@ __contract__(
 );
 
 #define mld_rej_uniform_avx2_asm MLD_NAMESPACE(rej_uniform_avx2_asm)
-/* This must be kept in sync with the HOL-Light specification
+/* This contract must be kept in sync with the HOL-Light specification
  * in proofs/hol_light/x86_64/proofs/rej_uniform_avx2_asm.ml */
 MLD_MUST_CHECK_RETURN_VALUE
 unsigned mld_rej_uniform_avx2_asm(
     int32_t *r, const uint8_t buf[MLD_AVX2_REJ_UNIFORM_BUFLEN],
     const uint8_t *table)
 __contract__(
   requires(memory_no_alias(r, sizeof(int32_t) * MLDSA_N))
-  requires(memory_no_alias(buf, MLD_AVX2_REJ_UNIFORM_BUFLEN))
-  requires(memory_no_alias(table, 256 * sizeof(uint64_t)))
+  requires(memory_no_alias(buf, 840))
+  requires(table == (const uint8_t *)mld_rej_uniform_table)
   assigns(memory_slice(r, sizeof(int32_t) * MLDSA_N))
   ensures(return_value <= MLDSA_N)
   ensures(array_bound(r, 0, return_value, 0, MLDSA_Q))

dev/x86_64/src/rej_uniform_avx2_asm.S

@@ -149,7 +149,6 @@ rej_uniform_avx2_asm_scalar:
         jmp     rej_uniform_avx2_asm_scalar
 
 rej_uniform_avx2_asm_done:
-        vzeroupper
         ret
 
 /* To facilitate single-compilation-unit (SCU) builds, undefine all macros.

mldsa/src/native/x86_64/src/arith_native_x86_64.h

@@ -74,16 +74,16 @@ __contract__(
 );
 
 #define mld_rej_uniform_avx2_asm MLD_NAMESPACE(rej_uniform_avx2_asm)
-/* This must be kept in sync with the HOL-Light specification
+/* This contract must be kept in sync with the HOL-Light specification
  * in proofs/hol_light/x86_64/proofs/rej_uniform_avx2_asm.ml */
 MLD_MUST_CHECK_RETURN_VALUE
 unsigned mld_rej_uniform_avx2_asm(
     int32_t *r, const uint8_t buf[MLD_AVX2_REJ_UNIFORM_BUFLEN],
     const uint8_t *table)
 __contract__(
   requires(memory_no_alias(r, sizeof(int32_t) * MLDSA_N))
-  requires(memory_no_alias(buf, MLD_AVX2_REJ_UNIFORM_BUFLEN))
-  requires(memory_no_alias(table, 256 * sizeof(uint64_t)))
+  requires(memory_no_alias(buf, 840))
+  requires(table == (const uint8_t *)mld_rej_uniform_table)
   assigns(memory_slice(r, sizeof(int32_t) * MLDSA_N))
   ensures(return_value <= MLDSA_N)
   ensures(array_bound(r, 0, return_value, 0, MLDSA_Q))

mldsa/src/native/x86_64/src/rej_uniform_avx2_asm.S

@@ -88,7 +88,6 @@ Lrej_uniform_avx2_asm_scalar:
         jmp Lrej_uniform_avx2_asm_scalar
 
 Lrej_uniform_avx2_asm_done:
-        vzeroupper
         retq
         .cfi_endproc
 

proofs/hol_light/x86_64/mldsa/rej_uniform_avx2_asm.S

@@ -20,85 +20,80 @@
 
 /*
  * WARNING: This file is auto-derived from the mldsa-native source file
- *   dev/x86_64/src/rej_uniform_avx2.S using scripts/simpasm. Do not modify it directly.
+ *   dev/x86_64/src/rej_uniform_avx2_asm.S using scripts/simpasm. Do not modify it directly.
  */
 
-#if defined(__ELF__)
-.section .note.GNU-stack,"",@progbits
-#endif
-
 .text
 .balign 4
 #ifdef __APPLE__
-.global _PQCP_MLDSA_NATIVE_MLDSA44_mld_rej_uniform_avx2
-_PQCP_MLDSA_NATIVE_MLDSA44_mld_rej_uniform_avx2:
+.global _PQCP_MLDSA_NATIVE_MLDSA44_rej_uniform_avx2_asm
+_PQCP_MLDSA_NATIVE_MLDSA44_rej_uniform_avx2_asm:
 #else
-.global PQCP_MLDSA_NATIVE_MLDSA44_mld_rej_uniform_avx2
-PQCP_MLDSA_NATIVE_MLDSA44_mld_rej_uniform_avx2:
+.global PQCP_MLDSA_NATIVE_MLDSA44_rej_uniform_avx2_asm
+PQCP_MLDSA_NATIVE_MLDSA44_rej_uniform_avx2_asm:
 #endif
 
         .cfi_startproc
         endbr64
-        movabs $0xff050403ff020100,%r10
-
-        vmovq  %r10,%xmm0
-        movabs $0xff0b0a09ff080706,%r10
-
-        vpinsrq $0x1,%r10,%xmm0,%xmm0
-        movabs $0xff090807ff060504,%r10
+        movabsq $-0xfafbfc00fdff00, %r10 # imm = 0xFF050403FF020100
+        vmovq %r10, %xmm0
+        movabsq $-0xf4f5f600f7f8fa, %r10 # imm = 0xFF0B0A09FF080706
+        vpinsrq $0x1, %r10, %xmm0, %xmm0
+        movabsq $-0xf6f7f800f9fafc, %r10 # imm = 0xFF090807FF060504
+        vmovq %r10, %xmm3
+        movabsq $-0xf0f1f200f3f4f6, %r10 # imm = 0xFF0F0E0DFF0C0B0A
+        vpinsrq $0x1, %r10, %xmm3, %xmm3
+        vinserti128 $0x1, %xmm3, %ymm0, %ymm0
+        movl $0x7fffff, %r8d         # imm = 0x7FFFFF
+        vmovd %r8d, %xmm1
+        vpbroadcastd %xmm1, %ymm1
+        movl $0x7fe001, %r8d         # imm = 0x7FE001
+        vmovd %r8d, %xmm2
+        vpbroadcastd %xmm2, %ymm2
+        xorl %eax, %eax
+        xorl %ecx, %ecx
 
-        vmovq  %r10,%xmm3
-        movabs $0xff0f0e0dff0c0b0a,%r10
+Lrej_uniform_avx2_asm_loop:
+        cmpl $0xf8, %eax
+        ja Lrej_uniform_avx2_asm_scalar
+        cmpl $0x328, %ecx            # imm = 0x328
+        ja Lrej_uniform_avx2_asm_scalar
+        vmovdqu (%rsi,%rcx), %ymm3
+        addl $0x18, %ecx
+        vpermq $0x94, %ymm3, %ymm3     # ymm3 = ymm3[0,1,1,2]
+        vpshufb %ymm0, %ymm3, %ymm3
+        vpand %ymm1, %ymm3, %ymm3
+        vpsubd %ymm2, %ymm3, %ymm4
+        vmovmskps %ymm4, %r8d
+        popcntl %r8d, %r9d
+        vmovq (%rdx,%r8,8), %xmm4
+        vpmovzxbd %xmm4, %ymm4    # ymm4 = xmm4[0],zero,zero,zero,xmm4[1],zero,zero,zero,xmm4[2],zero,zero,zero,xmm4[3],zero,zero,zero,xmm4[4],zero,zero,zero,xmm4[5],zero,zero,zero,xmm4[6],zero,zero,zero,xmm4[7],zero,zero,zero
+        vpermd %ymm3, %ymm4, %ymm3
+        vmovdqu %ymm3, (%rdi,%rax,4)
+        addl %r9d, %eax
+        jmp Lrej_uniform_avx2_asm_loop
 
-        vpinsrq $0x1,%r10,%xmm3,%xmm3
-        vinserti128 $0x1,%xmm3,%ymm0,%ymm0
-        mov    $0x7fffff,%r8d
-        vmovd  %r8d,%xmm1
-        vpbroadcastd %xmm1,%ymm1
-        mov    $0x7fe001,%r8d
-        vmovd  %r8d,%xmm2
-        vpbroadcastd %xmm2,%ymm2
-        xor    %eax,%eax
-        xor    %ecx,%ecx
+Lrej_uniform_avx2_asm_scalar:
+        cmpl $0x100, %eax            # imm = 0x100
+        jae Lrej_uniform_avx2_asm_done
+        cmpl $0x345, %ecx            # imm = 0x345
+        ja Lrej_uniform_avx2_asm_done
+        movzwl (%rsi,%rcx), %r8d
+        movzbl 0x2(%rsi,%rcx), %r9d
+        shll $0x10, %r9d
+        orl %r9d, %r8d
+        andl $0x7fffff, %r8d         # imm = 0x7FFFFF
+        addl $0x3, %ecx
+        cmpl $0x7fe001, %r8d         # imm = 0x7FE001
+        jae Lrej_uniform_avx2_asm_scalar
+        movl %r8d, (%rdi,%rax,4)
+        addl $0x1, %eax
+        jmp Lrej_uniform_avx2_asm_scalar
 
-Lmld_rej_uniform_avx2_loop:
-        cmp    $0xf8,%eax
-        ja     Lmld_rej_uniform_avx2_scalar
-        cmp    $0x328,%ecx
-        ja     Lmld_rej_uniform_avx2_scalar
-        vmovdqu (%rsi,%rcx,1),%ymm3
-        add    $0x18,%ecx
-        vpermq $0x94,%ymm3,%ymm3
-        vpshufb %ymm0,%ymm3,%ymm3
-        vpand  %ymm1,%ymm3,%ymm3
-        vpsubd %ymm2,%ymm3,%ymm4
-        vmovmskps %ymm4,%r8d
-        popcnt %r8d,%r9d
-        vmovq  (%rdx,%r8,8),%xmm4
-        vpmovzxbd %xmm4,%ymm4
-        vpermd %ymm3,%ymm4,%ymm3
-        vmovdqu %ymm3,(%rdi,%rax,4)
-        add    %r9d,%eax
-        jmp    Lmld_rej_uniform_avx2_loop
-
-Lmld_rej_uniform_avx2_scalar:
-        cmp    $0x100,%eax
-        jae    Lmld_rej_uniform_avx2_done
-        cmp    $0x345,%ecx
-        ja     Lmld_rej_uniform_avx2_done
-        movzwl (%rsi,%rcx,1),%r8d
-        movzbl 0x2(%rsi,%rcx,1),%r9d
-        shl    $0x10,%r9d
-        or     %r9d,%r8d
-        and    $0x7fffff,%r8d
-        add    $0x3,%ecx
-        cmp    $0x7fe001,%r8d
-        jae    Lmld_rej_uniform_avx2_scalar
-        mov    %r8d,(%rdi,%rax,4)
-        add    $0x1,%eax
-        jmp    Lmld_rej_uniform_avx2_scalar
-
-Lmld_rej_uniform_avx2_done:
-        vzeroupper
-        ret
+Lrej_uniform_avx2_asm_done:
+        retq
         .cfi_endproc
+
+#if defined(__ELF__)
+.section .note.GNU-stack,"",%progbits
+#endif

proofs/hol_light/x86_64/proofs/mldsa_rej_uniform_table.ml

@@ -1,11 +1,18 @@
 (*
- * Copyright (c) The mldsa-native project authors
  * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
  * SPDX-License-Identifier: Apache-2.0 OR ISC OR MIT-0
  *)
 
-(* Lookup table for ML-DSA rejection uniform sampling. *)
-(* Each entry is 8 bytes: permutation indices for VPERMD. *)
+(*
+ * WARNING: This file is auto-generated from scripts/autogen
+ *          in the mldsa-native repository.
+ *          Do not modify it directly.
+ *)
+
+(*
+ * Lookup table used by rejection sampling in the x86_64 AVX2
+ * implementation. See autogen for details.
+ *)
 
 let mldsa_rej_uniform_table = (REWRITE_RULE[MAP] o define)
   `mldsa_rej_uniform_table:byte list = MAP word [
@@ -264,5 +271,5 @@ let mldsa_rej_uniform_table = (REWRITE_RULE[MAP] o define)
     2;   3;   4;   5;   6;   7;   0;   0;
     0;   2;   3;   4;   5;   6;   7;   0;
     1;   2;   3;   4;   5;   6;   7;   0;
-    0;   1;   2;   3;   4;   5;   6;   7]`
-;;
+    0;   1;   2;   3;   4;   5;   6;   7
+]`;;

proofs/hol_light/x86_64/proofs/rej_uniform_avx2_asm.ml

@@ -101,7 +101,6 @@ let mldsa_rej_uniform_mc = define_assert_from_elf
   0x44; 0x89; 0x04; 0x87;  (* MOV (Memop Doubleword (%%% (rdi,2,rax))) (% r8d) *)
   0x83; 0xc0; 0x01;        (* ADD (% eax) (Imm8 (word 1)) *)
   0xeb; 0xc3;              (* JMP (Imm8 (word 195)) *)
-  0xc5; 0xf8; 0x77;        (* VZEROUPPER *)
   0xc3                     (* RET *)
 ];;
 (*** BYTECODE END ***)
@@ -406,6 +405,31 @@ let CMP_MASK_CORRECT = prove(
   CONV_TAC(DEPTH_CONV WORD_NUM_RED_CONV) THEN
   CONV_TAC NUM_REDUCE_CONV);;
 
+(* Pre-compute the 256 table entry values for VPERMD brute force.
+   Each entry is an int64 value: 8 bytes from the table at offset 8*mask. *)
+let TABLE_ENTRY_VALS =
+  let table_expanded =
+    (REWRITE_CONV[mldsa_rej_uniform_table; num_of_wordlist; DIMINDEX_8] THENC
+     DEPTH_CONV WORD_NUM_RED_CONV THENC NUM_REDUCE_CONV)
+    `num_of_wordlist mldsa_rej_uniform_table` in
+  let table_num = rhs(concl table_expanded) in
+  let entries = Array.init 256 (fun m ->
+    let tm = mk_comb(mk_comb(`(MOD)`,
+      mk_comb(mk_comb(`(DIV)`, table_num),
+      mk_comb(mk_comb(`(EXP)`, `2`), mk_numeral(Num.num_of_int(64*m))))),
+      mk_comb(mk_comb(`(EXP)`, `2`), `64`)) in
+    let th = NUM_REDUCE_CONV tm in
+    let rhs_val = rhs(concl th) in
+    (* Prove: (num_of_wordlist table DIV 2^(64*m)) MOD 2^64 = entry_m *)
+    let lhs_tm = mk_comb(mk_comb(`(MOD)`,
+      mk_comb(mk_comb(`(DIV)`,
+        `num_of_wordlist mldsa_rej_uniform_table`),
+      mk_comb(mk_comb(`(EXP)`, `2`), mk_numeral(Num.num_of_int(64*m))))),
+      mk_comb(mk_comb(`(EXP)`, `2`), `64`)) in
+    let eq = mk_eq(lhs_tm, rhs_val) in
+    EQT_ELIM((REWRITE_CONV[table_expanded] THENC NUM_REDUCE_CONV) eq)) in
+  entries;;
+
 (* TABLE_ENTRY_FROM_MEMORY: connect bytes64 memory read at table+8k to
    (table_num DIV 2^(64k)) MOD 2^64 via bigdigit/bignum_from_memory *)
 let TABLE_ENTRY_FROM_MEMORY = prove(
@@ -1416,9 +1440,9 @@ let VAL_RCX_ADD3_ZX = prove
 let SCALAR_BODY_LEMMA = prove
  (`!res buf table (inlist:(24 word)list) pc stackpointer N K i.
     LENGTH inlist = 280 /\
-    nonoverlapping (word pc, 246) (res, 1024) /\
-    nonoverlapping (word pc, 246) (buf, 840) /\
-    nonoverlapping (word pc, 246) (table, 2048) /\
+    nonoverlapping (word pc, 243) (res, 1024) /\
+    nonoverlapping (word pc, 243) (buf, 840) /\
+    nonoverlapping (word pc, 243) (table, 2048) /\
     nonoverlapping (res, 1024) (buf, 840) /\
     nonoverlapping (res, 1024) (table, 2048) /\
     24 * N <= 832 /\
@@ -2284,9 +2308,9 @@ let SCALAR_BODY_LEMMA = prove
 let MLDSA_REJ_UNIFORM_CORRECT = prove
  (`!res buf table (inlist:(24 word)list) pc.
     LENGTH inlist = 280 /\
-    nonoverlapping (word pc, 246) (res, 1024) /\
-    nonoverlapping (word pc, 246) (buf, 840) /\
-    nonoverlapping (word pc, 246) (table, 2048) /\
+    nonoverlapping (word pc, 243) (res, 1024) /\
+    nonoverlapping (word pc, 243) (buf, 840) /\
+    nonoverlapping (word pc, 243) (table, 2048) /\
     nonoverlapping (res, 1024) (buf, 840) /\
     nonoverlapping (res, 1024) (table, 2048)
     ==> ensures x86
@@ -2296,7 +2320,7 @@ let MLDSA_REJ_UNIFORM_CORRECT = prove
               read(memory :> bytes(buf,840)) s = num_of_wordlist inlist /\
               read(memory :> bytes(table,2048)) s =
                 num_of_wordlist(mldsa_rej_uniform_table:byte list))
-         (\s. read RIP s = word(pc + 245) /\
+         (\s. read RIP s = word(pc + 242) /\
               let outlist = SUB_LIST(0,256) (REJ_SAMPLE inlist) in
               let outlen = LENGTH outlist in
               C_RETURN s = word outlen /\
@@ -3775,7 +3799,7 @@ let MLDSA_REJ_UNIFORM_CORRECT = prove
            is_eq(concl th)
         then ASSUME_TAC(CONV_RULE(RAND_CONV(DEPTH_CONV WORD_NUM_RED_CONV)) th)
         else failwith "not RIP") THEN
-      X86_STEPS_TAC MLDSA_REJ_UNIFORM_EXEC [55] THEN
+      (* vzeroupper removed (was step 55); RIP is already at the RET. *)
       ENSURES_FINAL_STATE_TAC THEN ASM_REWRITE_TAC[] THEN
       CONV_TAC(TOP_DEPTH_CONV let_CONV) THEN
       SUBGOAL_THEN `SUB_LIST (0,256) (REJ_SAMPLE (inlist:(24 word)list)) =
@@ -3843,8 +3867,7 @@ let MLDSA_REJ_UNIFORM_CORRECT = prove
           let c = concl th in
           if is_conj c && (try can (find_term ((=) `LENGTH (REJ_SAMPLE (SUB_LIST (0,8 * N + K) (inlist:(24 word)list)))`)) c with _ -> false)
           then STRIP_ASSUME_TAC th else failwith "not inv") THEN
-        (* VZEROUPPER *)
-        X86_STEPS_TAC MLDSA_REJ_UNIFORM_EXEC [55] THEN
+        (* vzeroupper removed (was step 55); RIP is already at the RET. *)
         ENSURES_FINAL_STATE_TAC THEN ASM_REWRITE_TAC[] THEN
         CONV_TAC(TOP_DEPTH_CONV let_CONV) THEN
         (* The disjunct at K: either count-exit (256 <= outlen_K) or offset-exit (837 < 24*N+3*K) *)
@@ -3969,6 +3992,10 @@ let MLDSA_REJ_UNIFORM_CORRECT = prove
 
 (* ========================================================================= *)
 (* SUBROUTINE_CORRECT variants (standard x86_64 ABI).                        *)
+(*                                                                           *)
+(* These specifications must be kept in sync with the CBMC contract in       *)
+(* dev/x86_64/src/arith_native_x86_64.h / mldsa/src/native/x86_64/src/       *)
+(* arith_native_x86_64.h for mld_rej_uniform_avx2_asm.                       *)
 (* ========================================================================= *)
 
 let MLDSA_REJ_UNIFORM_NOIBT_SUBROUTINE_CORRECT = prove