Skip to content

Commit 73ab094

Browse files
jakemaswillieyz
authored andcommitted
x86_64 + HOL-Light: Replace poly_use_hint AVX2 intrinsics with hand-written assembly and HOL-Light proofs
Replaces the poly_use_hint_32 (ML-DSA-65/87) and poly_use_hint_88 (ML-DSA-44) AVX2 intrinsics with hand-written assembly, and adds HOL-Light functional correctness and constant-time/memory-safety proofs for both x86_64 routines. Signed-off-by: Jake Massimo <jakemas@amazon.com> Co-authored-by: willieyz <willieyz@users.noreply.github.com>
1 parent e67942f commit 73ab094

33 files changed

Lines changed: 4333 additions & 438 deletions

.github/workflows/hol_light.yml

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -200,6 +200,10 @@ jobs:
200200
needs: ["mldsa_specs.ml", "mldsa_utils.ml", "subroutine_signatures.ml"]
201201
- name: polyz_unpack_19_avx2_asm
202202
needs: ["mldsa_specs.ml", "mldsa_utils.ml", "subroutine_signatures.ml"]
203+
- name: poly_use_hint_32_avx2_asm
204+
needs: ["mldsa_specs.ml", "mldsa_utils.ml", "subroutine_signatures.ml"]
205+
- name: poly_use_hint_88_avx2_asm
206+
needs: ["mldsa_specs.ml", "mldsa_utils.ml", "subroutine_signatures.ml"]
203207
- name: ntt_avx2_asm
204208
needs: ["mldsa_specs.ml", "mldsa_utils.ml", "mldsa_zetas.ml", "subroutine_signatures.ml"]
205209
- name: intt_avx2_asm

BIBLIOGRAPHY.md

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -304,8 +304,8 @@ source code and documentation.
304304
- [dev/x86_64/src/poly_chknorm_avx2_asm.S](dev/x86_64/src/poly_chknorm_avx2_asm.S)
305305
- [dev/x86_64/src/poly_decompose_32_avx2_asm.S](dev/x86_64/src/poly_decompose_32_avx2_asm.S)
306306
- [dev/x86_64/src/poly_decompose_88_avx2_asm.S](dev/x86_64/src/poly_decompose_88_avx2_asm.S)
307-
- [dev/x86_64/src/poly_use_hint_32_avx2.c](dev/x86_64/src/poly_use_hint_32_avx2.c)
308-
- [dev/x86_64/src/poly_use_hint_88_avx2.c](dev/x86_64/src/poly_use_hint_88_avx2.c)
307+
- [dev/x86_64/src/poly_use_hint_32_avx2_asm.S](dev/x86_64/src/poly_use_hint_32_avx2_asm.S)
308+
- [dev/x86_64/src/poly_use_hint_88_avx2_asm.S](dev/x86_64/src/poly_use_hint_88_avx2_asm.S)
309309
- [dev/x86_64/src/polyz_unpack_17_avx2_asm.S](dev/x86_64/src/polyz_unpack_17_avx2_asm.S)
310310
- [dev/x86_64/src/polyz_unpack_19_avx2_asm.S](dev/x86_64/src/polyz_unpack_19_avx2_asm.S)
311311
- [dev/x86_64/src/rej_uniform_avx2.c](dev/x86_64/src/rej_uniform_avx2.c)
@@ -322,8 +322,8 @@ source code and documentation.
322322
- [mldsa/src/native/x86_64/src/poly_chknorm_avx2_asm.S](mldsa/src/native/x86_64/src/poly_chknorm_avx2_asm.S)
323323
- [mldsa/src/native/x86_64/src/poly_decompose_32_avx2_asm.S](mldsa/src/native/x86_64/src/poly_decompose_32_avx2_asm.S)
324324
- [mldsa/src/native/x86_64/src/poly_decompose_88_avx2_asm.S](mldsa/src/native/x86_64/src/poly_decompose_88_avx2_asm.S)
325-
- [mldsa/src/native/x86_64/src/poly_use_hint_32_avx2.c](mldsa/src/native/x86_64/src/poly_use_hint_32_avx2.c)
326-
- [mldsa/src/native/x86_64/src/poly_use_hint_88_avx2.c](mldsa/src/native/x86_64/src/poly_use_hint_88_avx2.c)
325+
- [mldsa/src/native/x86_64/src/poly_use_hint_32_avx2_asm.S](mldsa/src/native/x86_64/src/poly_use_hint_32_avx2_asm.S)
326+
- [mldsa/src/native/x86_64/src/poly_use_hint_88_avx2_asm.S](mldsa/src/native/x86_64/src/poly_use_hint_88_avx2_asm.S)
327327
- [mldsa/src/native/x86_64/src/polyz_unpack_17_avx2_asm.S](mldsa/src/native/x86_64/src/polyz_unpack_17_avx2_asm.S)
328328
- [mldsa/src/native/x86_64/src/polyz_unpack_19_avx2_asm.S](mldsa/src/native/x86_64/src/polyz_unpack_19_avx2_asm.S)
329329
- [mldsa/src/native/x86_64/src/rej_uniform_avx2.c](mldsa/src/native/x86_64/src/rej_uniform_avx2.c)
@@ -340,6 +340,8 @@ source code and documentation.
340340
- [proofs/hol_light/x86_64/mldsa/poly_chknorm_avx2_asm.S](proofs/hol_light/x86_64/mldsa/poly_chknorm_avx2_asm.S)
341341
- [proofs/hol_light/x86_64/mldsa/poly_decompose_32_avx2_asm.S](proofs/hol_light/x86_64/mldsa/poly_decompose_32_avx2_asm.S)
342342
- [proofs/hol_light/x86_64/mldsa/poly_decompose_88_avx2_asm.S](proofs/hol_light/x86_64/mldsa/poly_decompose_88_avx2_asm.S)
343+
- [proofs/hol_light/x86_64/mldsa/poly_use_hint_32_avx2_asm.S](proofs/hol_light/x86_64/mldsa/poly_use_hint_32_avx2_asm.S)
344+
- [proofs/hol_light/x86_64/mldsa/poly_use_hint_88_avx2_asm.S](proofs/hol_light/x86_64/mldsa/poly_use_hint_88_avx2_asm.S)
343345
- [proofs/hol_light/x86_64/mldsa/polyz_unpack_17_avx2_asm.S](proofs/hol_light/x86_64/mldsa/polyz_unpack_17_avx2_asm.S)
344346
- [proofs/hol_light/x86_64/mldsa/polyz_unpack_19_avx2_asm.S](proofs/hol_light/x86_64/mldsa/polyz_unpack_19_avx2_asm.S)
345347

dev/x86_64/meta.h

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -195,7 +195,7 @@ static MLD_INLINE int mld_poly_use_hint_32_native(int32_t *a, const int32_t *h)
195195
{
196196
return MLD_NATIVE_FUNC_FALLBACK;
197197
}
198-
mld_poly_use_hint_32_avx2(a, h);
198+
mld_poly_use_hint_32_avx2_asm(a, h);
199199
return MLD_NATIVE_FUNC_SUCCESS;
200200
}
201201
#endif /* MLD_CONFIG_MULTILEVEL_WITH_SHARED || MLD_CONFIG_PARAMETER_SET == 65 \
@@ -209,7 +209,7 @@ static MLD_INLINE int mld_poly_use_hint_88_native(int32_t *a, const int32_t *h)
209209
{
210210
return MLD_NATIVE_FUNC_FALLBACK;
211211
}
212-
mld_poly_use_hint_88_avx2(a, h);
212+
mld_poly_use_hint_88_avx2_asm(a, h);
213213
return MLD_NATIVE_FUNC_SUCCESS;
214214
}
215215
#endif /* MLD_CONFIG_MULTILEVEL_WITH_SHARED || MLD_CONFIG_PARAMETER_SET == 44 \

dev/x86_64/src/arith_native_x86_64.h

Lines changed: 28 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -143,11 +143,35 @@ __contract__(
143143
);
144144

145145
#if !defined(MLD_CONFIG_NO_VERIFY_API)
146-
#define mld_poly_use_hint_32_avx2 MLD_NAMESPACE(mld_poly_use_hint_32_avx2)
147-
void mld_poly_use_hint_32_avx2(int32_t *a, const int32_t *h);
146+
#define mld_poly_use_hint_32_avx2_asm MLD_NAMESPACE(poly_use_hint_32_avx2_asm)
147+
MLD_SYSV_ABI
148+
void mld_poly_use_hint_32_avx2_asm(int32_t *a, const int32_t *h)
149+
/* This must be kept in sync with the HOL-Light specification
150+
* in proofs/hol_light/x86_64/proofs/poly_use_hint_32_avx2_asm.ml */
151+
__contract__(
152+
requires(memory_no_alias(a, sizeof(int32_t) * MLDSA_N))
153+
requires(memory_no_alias(h, sizeof(int32_t) * MLDSA_N))
154+
requires(array_bound(a, 0, MLDSA_N, 0, MLDSA_Q))
155+
requires(array_bound(h, 0, MLDSA_N, 0, 2))
156+
assigns(memory_slice(a, sizeof(int32_t) * MLDSA_N))
157+
/* check-magic: 16 == (MLDSA_Q - 1) / (2 * ((MLDSA_Q - 1) / 32)) */
158+
ensures(array_bound(a, 0, MLDSA_N, 0, 16))
159+
);
148160

149-
#define mld_poly_use_hint_88_avx2 MLD_NAMESPACE(mld_poly_use_hint_88_avx2)
150-
void mld_poly_use_hint_88_avx2(int32_t *a, const int32_t *h);
161+
#define mld_poly_use_hint_88_avx2_asm MLD_NAMESPACE(poly_use_hint_88_avx2_asm)
162+
MLD_SYSV_ABI
163+
void mld_poly_use_hint_88_avx2_asm(int32_t *a, const int32_t *h)
164+
/* This must be kept in sync with the HOL-Light specification
165+
* in proofs/hol_light/x86_64/proofs/poly_use_hint_88_avx2_asm.ml */
166+
__contract__(
167+
requires(memory_no_alias(a, sizeof(int32_t) * MLDSA_N))
168+
requires(memory_no_alias(h, sizeof(int32_t) * MLDSA_N))
169+
requires(array_bound(a, 0, MLDSA_N, 0, MLDSA_Q))
170+
requires(array_bound(h, 0, MLDSA_N, 0, 2))
171+
assigns(memory_slice(a, sizeof(int32_t) * MLDSA_N))
172+
/* check-magic: 44 == (MLDSA_Q - 1) / (2 * ((MLDSA_Q - 1) / 88)) */
173+
ensures(array_bound(a, 0, MLDSA_N, 0, 44))
174+
);
151175
#endif /* !MLD_CONFIG_NO_VERIFY_API */
152176

153177
#define mld_poly_chknorm_avx2_asm MLD_NAMESPACE(poly_chknorm_avx2_asm)

dev/x86_64/src/poly_use_hint_32_avx2.c

Lines changed: 0 additions & 103 deletions
This file was deleted.
Lines changed: 168 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,168 @@
1+
/*
2+
* Copyright (c) The mldsa-native project authors
3+
* SPDX-License-Identifier: Apache-2.0 OR ISC OR MIT
4+
*/
5+
6+
/* References
7+
* ==========
8+
*
9+
* - [REF_AVX2]
10+
* CRYSTALS-Dilithium optimized AVX2 implementation
11+
* Bai, Ducas, Kiltz, Lepoint, Lyubashevsky, Schwabe, Seiler, Stehlé
12+
* https://github.com/pq-crystals/dilithium/tree/master/avx2
13+
*/
14+
15+
/*
16+
* This file is derived from the public domain
17+
* AVX2 Dilithium implementation @[REF_AVX2].
18+
*/
19+
20+
21+
/*yaml
22+
Name: poly_use_hint_32_avx2_asm
23+
Description: x86_64 AVX2 hint application (alpha = (Q-1)/32).
24+
Use the hint polynomial h to correct the high bits of the polynomial a,
25+
in place. Variant for parameter sets ML-DSA-65 and ML-DSA-87
26+
(GAMMA2 = (Q-1)/32).
27+
Signature: void mld_poly_use_hint_32_avx2_asm(int32_t *a, const int32_t *h)
28+
ABI:
29+
Architecture: x86_64
30+
CallingConvention: SysV
31+
Features: [AVX2]
32+
rdi:
33+
type: buffer
34+
size_bytes: 1024
35+
permissions: read/write
36+
c_parameter: int32_t *a
37+
description: Input/output polynomial (256 x int32_t)
38+
rsi:
39+
type: buffer
40+
size_bytes: 1024
41+
permissions: read-only
42+
c_parameter: const int32_t *h
43+
description: Hint polynomial (256 x int32_t)
44+
*/
45+
46+
#include "../../../common.h"
47+
48+
#if defined(MLD_ARITH_BACKEND_X86_64_DEFAULT) && \
49+
!defined(MLD_CONFIG_NO_VERIFY_API) && \
50+
!defined(MLD_CONFIG_MULTILEVEL_NO_SHARED) && \
51+
(defined(MLD_CONFIG_MULTILEVEL_WITH_SHARED) || \
52+
(MLD_CONFIG_PARAMETER_SET == 65 || MLD_CONFIG_PARAMETER_SET == 87))
53+
54+
/* simpasm: header-end */
55+
56+
/* Reference:
57+
* - @[REF_AVX2] calls poly_decompose to compute all a1, a0 before the loop.
58+
* - Our implementation of decompose() is slightly different from that in
59+
* @[REF_AVX2]. See poly_decompose_32_avx2.c for more information.
60+
*/
61+
62+
// a aliased with a0
63+
.macro decompose32_avx2 a1, a, temp1, temp2, temp3
64+
// Compute a1 = round-(a / 523776), where 523776 = 2*GAMMA2 = 128*B
65+
// with B = 4092, and round-() denotes "round half down". This is
66+
// exact for 0 <= a < Q. It is computed as round-(ceil(a / 128) / B)
67+
// ≈ round(ceil(a / 128) * 1025 / 2^22), with 1025 = floor(2^22 / B):
68+
// temp1 = ceil(a / 128) = floor((a + 127) / 2^7) (vpaddd 127; vpsrld 7)
69+
// temp1 = floor(temp1 * 1025 / 2^16) (vpmulhuw by 1025)
70+
// temp1 = round(temp1 / 2^6) via mulhrs(temp1, 2^9) (vpmulhrsw)
71+
// See mld_poly_decompose_32_avx2() for the detailed argument.
72+
vpaddd \a, %ymm5, \temp1
73+
vpsrld $7, \temp1, \temp1
74+
vpmulhuw %ymm8, \temp1, \temp1
75+
vpmulhrsw %ymm7, \temp1, \temp1
76+
77+
// If a1 = 16, i.e. a > 31*GAMMA2, proceed as if a' = a - Q was
78+
// given instead. (For a = 31*GAMMA2 + 1 thus a' = -GAMMA2, we
79+
// still round it to 0 like other "wrapped around" cases.)
80+
81+
// Check for wrap-around
82+
vpcmpgtd %ymm4, \a, \temp2
83+
vpandn \temp1, \temp2, \a1
84+
85+
// Compute remainder a0 = a - a1 * 2 * GAMMA2 = a - a1 * 523776
86+
vpslld $10, \temp1, \temp3
87+
vpsubd \temp1, \temp3, \temp1
88+
vpslld $9, \temp1, \temp1
89+
vpsubd \temp1, \a, \a
90+
91+
// If wrap-around is required, adjust a0 by -1
92+
vpaddd \temp2, \a, \a
93+
.endm
94+
95+
/* We follow the specification and the reference C implementation and check
96+
* a0 > 0. The @[REF_AVX2] implementation used to check a0 >= 0 instead, but
97+
* this was corrected to match the specification in pq-crystals/dilithium
98+
* commit bba9534a3c5cbef66481cd6518d786e23ae601b5.
99+
*/
100+
101+
// a aliased with delta
102+
.macro use_hint32_avx2 b, a, h, a1, temp1, temp2, temp3
103+
decompose32_avx2 \a1, \a, \temp1, \temp2, \temp3
104+
105+
// delta = (a0 <= 0) ? -1 : 1
106+
vpcmpgtd %ymm6, \a, \a
107+
vpandn \h, \a, \a
108+
vpslld $1, \a, \a
109+
vpsubd \a, \h, \h
110+
111+
// b = (b + delta * h) % 16
112+
vpaddd \a1, \h, \b
113+
vpand %ymm3, \b, \b
114+
.endm
115+
116+
.text
117+
.balign 16
118+
.global MLD_ASM_NAMESPACE(poly_use_hint_32_avx2_asm)
119+
MLD_ASM_FN_SYMBOL(poly_use_hint_32_avx2_asm)
120+
121+
// Initialize constants
122+
movl $127, %ecx
123+
124+
/* check-magic: 1025 == floor(2^22 / 4092) */
125+
movl $1025, %r8d
126+
vmovd %r8d, %xmm8
127+
vpbroadcastd %xmm8, %ymm8
128+
129+
xorl %eax, %eax
130+
vpxor %xmm6, %xmm6, %xmm6
131+
vmovd %ecx, %xmm5
132+
133+
/* 31 * ((Q-1) / 32) == 31 * GAMMA2, wrap-around threshold */
134+
movl $8118528, %ecx
135+
136+
/* round(x * 2^9 / 2^15) => round(x / 2^6), for f1 = round(f1'' / 2^6) */
137+
movl $512, %r9d
138+
vmovd %r9d, %xmm7
139+
vpbroadcastd %xmm7, %ymm7
140+
141+
vmovd %ecx, %xmm4
142+
movl $15, %ecx
143+
vpbroadcastd %xmm5, %ymm5
144+
vmovd %ecx, %xmm3
145+
vpbroadcastd %xmm4, %ymm4
146+
vpbroadcastd %xmm3, %ymm3
147+
148+
149+
poly_use_hint_32_avx2_asm_loop:
150+
vmovdqa (%rdi), %ymm0
151+
vmovdqa (%rsi), %ymm2
152+
153+
use_hint32_avx2 %ymm2, %ymm0, %ymm2, %ymm9, %ymm1, %ymm11, %ymm10
154+
155+
vmovdqa %ymm2, (%rdi)
156+
addq $32, %rdi
157+
addq $32, %rsi
158+
addq $32, %rax
159+
cmpq $1024, %rax
160+
jne poly_use_hint_32_avx2_asm_loop
161+
ret
162+
163+
/* simpasm: footer-start */
164+
165+
#endif /* MLD_ARITH_BACKEND_X86_64_DEFAULT && !MLD_CONFIG_NO_VERIFY_API && \
166+
!MLD_CONFIG_MULTILEVEL_NO_SHARED && \
167+
(MLD_CONFIG_MULTILEVEL_WITH_SHARED || MLD_CONFIG_PARAMETER_SET == 65 \
168+
|| MLD_CONFIG_PARAMETER_SET == 87) */

0 commit comments

Comments
 (0)