Skip to content

Commit 3506c1f

Browse files
siemen11siemen11
committed
[crypto/mlkem] Add main subroutines for ml-kem1024
Add the main arithmetic and subroutines for the ml-kem1024 in the otbn. Signed-off-by: Siemen Dhooghe <sdhooghe@google.com>
1 parent 58f8fdd commit 3506c1f

54 files changed

Lines changed: 5084 additions & 0 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎sw/otbn/crypto/mlkem1024/BUILD‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
# Copyright lowRISC contributors (OpenTitan project).
2+
# Licensed under the Apache License, Version 2.0, see LICENSE for details.
3+
# SPDX-License-Identifier: Apache-2.0
4+
5+
load("//rules:otbn.bzl", "otbn_library")
6+
7+
package(default_visibility = ["//visibility:public"])
8+
9+
otbn_library(
10+
name = "mlkem1024_arith",
11+
srcs = [
12+
"mlkem1024_arith.s",
13+
],
14+
)
15+
16+
otbn_library(
17+
name = "mlkem1024_sample",
18+
srcs = [
19+
"mlkem1024_sample.s",
20+
],
21+
)
22+
23+
otbn_library(
24+
name = "mlkem1024_expand",
25+
srcs = [
26+
"mlkem1024_expand.s",
27+
],
28+
)
29+
30+
otbn_library(
31+
name = "mlkem1024_ntt",
32+
srcs = [
33+
"mlkem1024_ntt.s",
34+
],
35+
)
36+
37+
otbn_library(
38+
name = "mlkem1024_encoding",
39+
srcs = [
40+
"mlkem1024_encoding.s",
41+
],
42+
)
43+
44+
otbn_library(
45+
name = "mlkem1024_decoding",
46+
srcs = [
47+
"mlkem1024_decoding.s",
48+
],
49+
)

‎sw/otbn/crypto/mlkem1024/mlkem1024_arith.s‎

Lines changed: 254 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,254 @@
1+
/* Copyright lowRISC contributors (OpenTitan project). */
2+
/* Licensed under the Apache License, Version 2.0, see LICENSE for details. */
3+
/* SPDX-License-Identifier: Apache-2.0 */
4+
5+
/* Polynomial arithmetic operations for ML-KEM-1024 (q = 3329). */
6+
7+
.globl poly_add
8+
.globl poly_sub
9+
.globl poly_mul
10+
.globl poly_mul_add
11+
12+
.text
13+
14+
/**
15+
* Pointwise Addition of polynomials for ML-KEM-1024.
16+
*
17+
* Computes c(X) = a(X) + b(X) mod 3329.
18+
*
19+
* Preconditions:
20+
* - The OTBN `MOD` CSR must be loaded with q = 3329 (0x00000D01).
21+
* - DMEM buffers at x2, x3, and x4 must be 32-byte aligned.
22+
* - x31 must point to a valid stack area in DMEM for register preservation.
23+
*
24+
* @param[in] x2: DMEM address of polynomial a(X) (256 32-bit words, 1024 bytes).
25+
* @param[in] x3: DMEM address of polynomial b(X) (256 32-bit words, 1024 bytes).
26+
* @param[out] x4: DMEM address of polynomial c(X) (256 32-bit words, 1024 bytes).
27+
*
28+
* Clobbered registers: w0, w1, w2.
29+
*/
30+
.balign 4
31+
poly_add:
32+
/* Push clobbered general-purpose registers onto the stack. */
33+
.irp reg, x2, x3, x4, x10, x11, x12
34+
sw \reg, 0(x31)
35+
addi x31, x31, 4
36+
.endr
37+
38+
addi x10, x0, 0
39+
addi x11, x0, 1
40+
addi x12, x0, 2
41+
42+
/* Loop 32 iterations (32 * 8 = 256 coefficients). */
43+
loopi 32, 4
44+
bn.lid x10, 0(x2++)
45+
bn.lid x11, 0(x3++)
46+
bn.addvm.8S w2, w0, w1
47+
bn.sid x12, 0(x4++)
48+
/* End loop (4 instructions) */
49+
50+
/* Restore registers from stack. */
51+
.irp reg, x12, x11, x10, x4, x3, x2
52+
addi x31, x31, -4
53+
lw \reg, 0(x31)
54+
.endr
55+
56+
ret
57+
58+
/**
59+
* Pointwise Subtraction of polynomials for ML-KEM-1024.
60+
*
61+
* Computes c(X) = a(X) - b(X) mod 3329.
62+
*
63+
* Preconditions:
64+
* - The OTBN `MOD` CSR must be loaded with q = 3329 (0x00000D01).
65+
* - DMEM buffers at x2, x3, and x4 must be 32-byte aligned.
66+
* - x31 must point to a valid stack area in DMEM for register preservation.
67+
*
68+
* @param[in] x2: DMEM address of polynomial a(X) (256 32-bit words, 1024 bytes).
69+
* @param[in] x3: DMEM address of polynomial b(X) (256 32-bit words, 1024 bytes).
70+
* @param[out] x4: DMEM address of polynomial c(X) (256 32-bit words, 1024 bytes).
71+
*
72+
* Clobbered registers: w0, w1, w2.
73+
*/
74+
.balign 4
75+
poly_sub:
76+
/* Push clobbered general-purpose registers onto the stack. */
77+
.irp reg, x2, x3, x4, x10, x11, x12
78+
sw \reg, 0(x31)
79+
addi x31, x31, 4
80+
.endr
81+
82+
addi x10, x0, 0
83+
addi x11, x0, 1
84+
addi x12, x0, 2
85+
86+
/* Loop 32 iterations (32 * 8 = 256 coefficients). */
87+
loopi 32, 4
88+
bn.lid x10, 0(x2++)
89+
bn.lid x11, 0(x3++)
90+
bn.subvm.8S w2, w0, w1
91+
bn.sid x12, 0(x4++)
92+
/* End loop (4 instructions) */
93+
94+
/* Restore registers from stack. */
95+
.irp reg, x12, x11, x10, x4, x3, x2
96+
addi x31, x31, -4
97+
lw \reg, 0(x31)
98+
.endr
99+
100+
ret
101+
102+
/**
103+
* Pointwise Montgomery Multiplication of NTT polynomials for ML-KEM-1024.
104+
*
105+
* Computes c(X) = a(X) * b(X) mod 3329 in the NTT domain. Evaluates 128 degree-1
106+
* polynomial multiplications modulo (X^2 - gamma_i).
107+
*
108+
* Preconditions:
109+
* - The OTBN `MOD` CSR must be loaded with q = 3329 (0x00000D01) and mu = 2488732927 (0x94570CFF).
110+
* - DMEM buffers at x2, x3, x4, and x5 must be 32-byte aligned.
111+
* - x31 must point to a valid stack area in DMEM for register preservation.
112+
*
113+
* @param[in] x2: DMEM address of polynomial a(X) (256 32-bit words).
114+
* @param[in] x3: DMEM address of polynomial b(X) (256 32-bit words).
115+
* @param[in] x4: DMEM address of gamma twiddles (128 32-bit words, 512 bytes total).
116+
* @param[out] x5: DMEM output address for c(X) = a(X) * b(X) (256 32-bit words).
117+
*
118+
* Clobbered registers: w0, w1, w2, w3, w4, w5, w6, w7, w8, w9, w14.
119+
*/
120+
.balign 4
121+
poly_mul:
122+
/* Push clobbered general-purpose registers onto the stack. */
123+
.irp reg, x2, x3, x4, x5, x10, x11, x12
124+
sw \reg, 0(x31)
125+
addi x31, x31, 4
126+
.endr
127+
128+
/* Setup WDR index registers for loop LID operations */
129+
addi x10, x0, 0
130+
addi x11, x0, 1
131+
addi x12, x0, 2
132+
133+
/* Generate mask w9 = (0x00000000_ffffffff, ...) for selecting even slots */
134+
bn.not w9, w31
135+
bn.trn1.8S w9, w9, w31
136+
137+
/* Loop 32 iterations processing 4 quadratic pairs (8 coefficients) per iteration. */
138+
loopi 32, 18
139+
/* Load 8 coefficients of a into w0, b into w1, and 8 twiddles into w2. */
140+
bn.lid x10, 0(x2++)
141+
bn.lid x11, 0(x3++)
142+
bn.lid x12, 0(x4++)
143+
144+
/* Basecase even products c0 = a0*b0 + a1*b1*gamma */
145+
bn.mulvm.8S w3, w0, w1
146+
bn.mulvm.8S w4, w3, w2
147+
bn.rshi w5, w31, w4 >> 32
148+
bn.addvm.8S w5, w3, w5 /* w5[even] = c0 = a0*b0 + a1*b1*gamma */
149+
150+
/* Basecase odd products c1 = a0*b1 + a1*b0 */
151+
bn.rshi w6, w1, w1 >> 32 /* w6[even] = b1 */
152+
bn.mulvm.8S w7, w0, w6 /* w7[even] = a0*b1 */
153+
bn.rshi w8, w31, w0 >> 32 /* w8[even] = a1 */
154+
bn.mulvm.8S w8, w8, w1 /* w8[even] = a1*b0 */
155+
bn.addvm.8S w8, w7, w8 /* w8[even] = c1 = a0*b1 + a1*b0 */
156+
157+
/* Combine c0 (even) and c1 (odd) */
158+
bn.and w5, w5, w9 /* clear odd slots of c0 */
159+
bn.rshi w8, w8, w31 >> 224 /* move c1 to odd slots */
160+
bn.not w14, w9 /* mask for odd slots */
161+
bn.and w8, w8, w14 /* clear even slots of c1 */
162+
bn.or w0, w5, w8 /* w0 = [c7, c6, c5, c4, c3, c2, c1, c0] */
163+
164+
/* Store 8 result coefficients back to DMEM[x5] */
165+
bn.sid x10, 0(x5++)
166+
/* End of loop body (18 instructions) */
167+
168+
/* Restore registers from stack. */
169+
.irp reg, x12, x11, x10, x5, x4, x3, x2
170+
addi x31, x31, -4
171+
lw \reg, 0(x31)
172+
.endr
173+
174+
ret
175+
176+
/**
177+
* Pointwise Montgomery Multiply-Accumulate of NTT polynomials for ML-KEM-1024.
178+
*
179+
* Computes c_out(X) = c_in(X) + a(X) * b(X) mod 3329 in the NTT domain.
180+
* Avoids extra memory passes during matrix-vector multiplication (A * s + e).
181+
*
182+
* Preconditions:
183+
* - The OTBN `MOD` CSR must be loaded with q = 3329 (0x00000D01) and mu = 2488732927 (0x94570CFF).
184+
* - DMEM buffers at x2, x3, x4, x5, and x6 must be 32-byte aligned.
185+
* - x31 must point to a valid stack area in DMEM for register preservation.
186+
*
187+
* @param[in] x2: DMEM address of polynomial a(X) (256 32-bit words).
188+
* @param[in] x3: DMEM address of polynomial b(X) (256 32-bit words).
189+
* @param[in] x4: DMEM address of gamma twiddles (128 32-bit words, 512 bytes total).
190+
* @param[in] x5: DMEM address of input accumulator polynomial c_in(X) (256 32-bit words).
191+
* @param[out] x6: DMEM output address for c_out(X) = c_in(X) + a(X) * b(X) mod 3329.
192+
*
193+
* Clobbered registers: w0, w1, w2, w3, w4, w5, w6, w7, w8, w9, w13, w14.
194+
*/
195+
.balign 4
196+
poly_mul_add:
197+
/* Push clobbered general-purpose registers onto the stack. */
198+
.irp reg, x2, x3, x4, x5, x6, x10, x11, x12, x13
199+
sw \reg, 0(x31)
200+
addi x31, x31, 4
201+
.endr
202+
203+
/* Setup WDR index registers for loop LID operations */
204+
addi x10, x0, 0
205+
addi x11, x0, 1
206+
addi x12, x0, 2
207+
addi x13, x0, 13
208+
209+
/* Generate mask w9 = (0x00000000_ffffffff, ...) for selecting even slots */
210+
bn.not w9, w31
211+
bn.trn1.8S w9, w9, w31
212+
213+
/* Loop 32 iterations processing 4 quadratic pairs (8 coefficients) per iteration. */
214+
loopi 32, 22
215+
/* Load a into w0, b into w1, gamma into w2, and accumulator c_in into w13. */
216+
bn.lid x10, 0(x2++)
217+
bn.lid x11, 0(x3++)
218+
bn.lid x12, 0(x4++)
219+
bn.lid x13, 0(x5++)
220+
221+
/* Basecase even products c0 = a0*b0 + a1*b1*gamma */
222+
bn.mulvm.8S w3, w0, w1
223+
bn.mulvm.8S w4, w3, w2
224+
bn.rshi w5, w31, w4 >> 32
225+
bn.addvm.8S w5, w3, w5 /* w5[even] = a0*b0 + a1*b1*gamma */
226+
bn.addvm.8S w5, w5, w13 /* w5[even] = c0_out = (a0*b0 + a1*b1*gamma) + c0_in */
227+
228+
/* Basecase odd products c1 = a0*b1 + a1*b0 */
229+
bn.rshi w6, w1, w1 >> 32 /* w6[even] = b1 */
230+
bn.mulvm.8S w7, w0, w6 /* w7[even] = a0*b1 */
231+
bn.rshi w8, w31, w0 >> 32 /* w8[even] = a1 */
232+
bn.mulvm.8S w8, w8, w1 /* w8[even] = a1*b0 */
233+
bn.addvm.8S w8, w7, w8 /* w8[even] = a0*b1 + a1*b0 */
234+
bn.rshi w14, w31, w13 >> 32 /* w14[even] = c1_in */
235+
bn.addvm.8S w8, w8, w14 /* w8[even] = c1_out = (a0*b1 + a1*b0) + c1_in */
236+
237+
/* Combine c0_out (even) and c1_out (odd) */
238+
bn.and w5, w5, w9 /* clear odd slots of c0_out */
239+
bn.rshi w8, w8, w31 >> 224 /* move c1_out to odd slots */
240+
bn.not w14, w9 /* mask for odd slots */
241+
bn.and w8, w8, w14 /* clear even slots of c1_out */
242+
bn.or w0, w5, w8 /* w0 = [c7, c6, c5, c4, c3, c2, c1, c0] */
243+
244+
/* Store 8 result coefficients back to DMEM[x6] */
245+
bn.sid x10, 0(x6++)
246+
/* End of loop body (22 instructions) */
247+
248+
/* Restore registers from stack. */
249+
.irp reg, x13, x12, x11, x10, x6, x5, x4, x3, x2
250+
addi x31, x31, -4
251+
lw \reg, 0(x31)
252+
.endr
253+
254+
ret

0 commit comments

Comments
 (0)