fallensigh/mlkem-sync
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(mlkem_top): KeyGen stage 2e (matrix accumulate t_hat)

Browse Source 
Add ST_M stage: t_hat[i] = e_hat[i] + sum_j A_hat[i][j] o s_hat[j] via
poly_mul_sync + inline mod-add accumulation. Per (i,j): stream 256 (A,shat)
pairs into poly_mul, then accumulate 256 products into T_i (seeded from E_i
when j==0, else running T_i). m_pending waits for poly_mul IDLE between terms.

Verified vs ml-kem-r golden: 512/512 t_hat coeffs exact (19885 cyc).
This commit is contained in:
FallenSigh
2026-06-28 01:53:23 +08:00
parent 4c692e570a
commit a9e50ebc0c
3 changed files with 672 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

57
sync_rtl/top/TB/tb_mlkem_kg_2e_xsim.v Normal file
View File

@@ -0,0 +1,57 @@
// tb_mlkem_kg_2e_xsim.v - Stage 2e: verify t_hat = e_hat + sum_j A[i][j]*s_hat[j].
// After ST_M, slots T0,T1 must hold that_0,that_1.
// Golden: kg_c000_that.hex (2 polys x 256 = 512 lines, mod-q).
`timescale 1ns/1ps
module tb_mlkem_kg_2e_xsim;
reg clk=0, rst_n=0, start_i=0;
reg [255:0] d_i, z_i=0;
wire busy_o, done_o;
reg [3:0] dbg_slot_i=0; reg [7:0] dbg_idx_i=0; wire [11:0] dbg_coeff_o;
wire [255:0] dbg_rho_o, dbg_sigma_o;
mlkem_top #(.K(2)) dut (
.clk(clk), .rst_n(rst_n), .d_i(d_i), .z_i(z_i), .start_i(start_i),
.busy_o(busy_o), .done_o(done_o),
.dbg_slot_i(dbg_slot_i), .dbg_idx_i(dbg_idx_i), .dbg_coeff_o(dbg_coeff_o),
.dbg_rho_o(dbg_rho_o), .dbg_sigma_o(dbg_sigma_o)
);
always #5 clk = ~clk;
localparam [255:0] D_LIT = 256'h2426f1941779574d3f1b163bd57f7e173e229e630ec7f7073bdf365137c4bb6d;
reg [11:0] gold [0:511];
reg [3:0] slot_of [0:1]; // T0,T1
integer c, p, idx, errors, gi;
initial begin
$readmemh("sync_rtl/top/TB/vectors/kg_c000_that.hex", gold);
slot_of[0]=4'd8; slot_of[1]=4'd9;
d_i = D_LIT;
rst_n=0; repeat(4) @(posedge clk); rst_n=1; @(posedge clk);
start_i=1; @(posedge clk); start_i=0;
c=0; while(!done_o && c<400000) begin @(posedge clk); c=c+1; end
if(!done_o) begin $display("FAIL: timeout after %0d cyc", c); $finish; end
$display("=== Stage 2e: t_hat = e_hat + sum A o s_hat === done in %0d cyc", c);
errors = 0;
for (p = 0; p < 2; p = p + 1) begin
for (idx = 0; idx < 256; idx = idx + 1) begin
dbg_slot_i = slot_of[p];
dbg_idx_i = idx[7:0];
@(posedge clk); @(posedge clk);
gi = p*256 + idx;
if (dbg_coeff_o !== gold[gi]) begin
if (errors < 8)
$display(" MISMATCH slot%0d[%0d]: got=%03x exp=%03x",
slot_of[p], idx, dbg_coeff_o, gold[gi]);
errors = errors + 1;
end
end
end
if (errors == 0) $display("ALL TESTS PASSED (512/512 t_hat coeffs)");
else $display("TESTS FAILED: %0d mismatches", errors);
$finish;
end
initial begin #20000000; $display("FAIL: global timeout"); $finish; end
endmodule

512
sync_rtl/top/TB/vectors/kg_c000_that.hex Normal file
View File

@@ -0,0 +1,512 @@
780
01b
300
1d8
17e
68e
3f4
927
148
193
8f9
633
0ec
bcb
360
125
7d3
285
b87
006
5f5
7af
8f4
5ec
7da
289
71d
435
c90
72b
956
361
c68
98d
362
545
2ab
b95
7bc
81e
7c5
1c7
0c2
996
bf6
acf
1be
78e
94d
481
387
1f7
c09
877
2e0
cb8
5ea
7dc
730
b9b
3a2
879
42c
148
266
24f
64c
5fc
af4
65b
73d
516
2df
20a
551
53d
936
429
cc3
291
3f9
071
7af
325
692
97d
3a8
16c
012
7be
553
cf3
208
9f3
6d2
bf4
22a
203
c0a
b08
0c4
3bb
99d
c50
cba
391
319
9e5
992
c62
296
439
848
2f2
b69
156
525
60e
638
b40
1db
8b6
83d
63c
20c
86b
1e4
173
c21
50b
231
57f
2aa
5b9
9f8
600
a44
a85
6d8
97e
988
9de
ca9
1d5
c10
3bc
8e8
525
08d
cfb
0ea
3c0
974
376
680
5aa
636
28e
343
beb
015
654
243
2dd
083
7c5
287
78b
267
c89
50f
62a
484
a35
910
4d3
cda
2a4
4b5
939
285
570
3d6
bcd
110
867
6d9
270
307
54f
ae8
1a1
c10
b7f
363
1ad
44e
09c
4aa
3e2
5d2
9be
9b1
858
15e
65b
662
bb6
037
075
a85
08e
8cd
c28
822
c5a
73a
778
4ca
2cb
39c
8a2
500
93d
269
504
926
9de
62a
04e
60a
70f
990
ca2
987
c88
397
29b
7fb
48a
cd6
8a9
871
c94
b5a
0cb
871
7c7
158
2ff
221
21f
83e
923
a42
7fa
2bd
0c3
c12
42a
1c6
601
86c
c1b
4d1
62b
658
a18
545
35a
5bd
bd6
90d
6b9
c33
160
657
1ad
987
04e
017
320
525
987
320
bba
0ec
5c5
1da
11f
293
9a2
7ae
3c1
bd1
ad2
722
a3a
60a
a97
6fa
367
5c3
aad
148
78e
240
a01
112
41f
768
85b
a3d
162
cf9
a00
32c
3bc
620
c00
68f
c12
a5f
29c
c61
c72
192
cfe
102
32e
51a
a87
910
198
686
7d6
841
af4
636
5ee
91a
255
289
960
317
4c7
c31
0e5
739
51b
68c
608
2cd
65c
8d0
13c
0ef
a02
143
7f3
bae
8e5
440
6e5
7ec
1df
56c
4f0
ba6
496
340
091
3be
0ff
ad5
551
519
5d7
0a9
461
289
6fb
908
153
44c
7b8
3ef
652
c94
023
849
3a4
aad
739
967
80c
1a3
ced
246
00d
58a
271
35c
bca
3a4
4e8
a6f
309
036
cdc
a80
22b
48d
b93
278
440
0d5
cbc
7e9
c60
3fb
824
c4a
079
746
164
b39
86c
c18
608
2c3
53b
c21
c95
216
bc9
9d1
3d4
46a
4d4
453
466
9e0
682
013
018
8cc
7e2
10e
8f5
23a
5da
b8c
953
3b3
5cc
459
81f
89d
494
276
9a0
af0
265
53c
333
66b
1f2
277
cdb
3c2
899
4f0
c08
b4c
37b
80c
564
b2b
505
c77
0f1
6b8
901
98e
8d1
5d6
3fd
476
8be
2f0
c24
8b1
2af
03a
b84
5da
1ff
746
6fd
ca4
c9b
8d3
94c
5ff
0c8
a0f
32a
ce0
33a
a90
4c1

104
sync_rtl/top/mlkem_top.v
View File

@@ -73,6 +73,7 @@ module mlkem_top #(
localparam ST_A = 4'd2; // generate A_hat[i][j] via SampleNTT localparam ST_A = 4'd2; // generate A_hat[i][j] via SampleNTT
localparam ST_C = 4'd3; // generate s[i],e[i] via CBD localparam ST_C = 4'd3; // generate s[i],e[i] via CBD
localparam ST_N = 4'd4; // forward NTT of s[i],e[i] in place localparam ST_N = 4'd4; // forward NTT of s[i],e[i] in place
localparam ST_M = 4'd5; // matrix accumulate t_hat = e_hat + sum A o s_hat
localparam ST_DONE = 4'd15; localparam ST_DONE = 4'd15;
reg [3:0] st, st_next; reg [3:0] st, st_next;
@@ -195,6 +196,46 @@ module mlkem_top #(
.done_o(ntt_done) .done_o(ntt_done)
); );
// ---- poly_mul_sync: t_hat[i] = e_hat[i] + sum_j A_hat[i][j] o s_hat[j] ----
// M-stage bookkeeping. For each (i,j): LOAD 256 (A,shat) pairs, then accumulate
// 256 products into T_i (init from E_i when j==0, else from running T_i).
reg [1:0] m_i; // row 0..K
reg [1:0] m_j; // col 0..K
reg [8:0] m_ld; // load index 0..256
reg [7:0] m_oidx; // output/accum index 0..255
reg m_loading; // 1 while streaming pairs into poly_mul
reg m_pending; // wait for poly_mul IDLE before next (i,j)
wire [3:0] m_aslot = {2'b0, m_i[0], m_j[0]}; // A_hat[i][j] slot = i*2+j (0..3)
wire [3:0] m_sslot = SLOT_S0 + {3'b0, m_j[0]}; // s_hat[j]
wire [3:0] m_eslot = SLOT_E0 + {3'b0, m_i[0]}; // e_hat[i]
wire [3:0] m_tslot = SLOT_T0 + {3'b0, m_i[0]}; // t_hat[i]
reg pm_valid;
wire pm_ready;
wire [11:0] pm_coeff;
wire pm_vo;
wire [11:0] pm_a_in = polymem[m_aslot*256 + m_ld[7:0]];
wire [11:0] pm_b_in = polymem[m_sslot*256 + m_ld[7:0]];
poly_mul_sync u_pmul (
.clk(clk), .rst_n(rst_n),
.coeff_a_in(pm_a_in),
.coeff_b_in(pm_b_in),
.valid_i(pm_valid),
.ready_o(pm_ready),
.coeff_out(pm_coeff),
.valid_o(pm_vo),
.ready_i(1'b1)
);
// accumulator source: e_hat[i] for first term (j==0), else running t_hat[i]
wire [11:0] m_acc_src = (m_j == 2'd0) ? polymem[m_eslot*256 + m_oidx]
: polymem[m_tslot*256 + m_oidx];
// (a + b) mod Q (both < Q, sum < 2Q): one conditional subtract
wire [12:0] m_sum = {1'b0, m_acc_src} + {1'b0, pm_coeff};
wire [11:0] m_accq = (m_sum >= 13'(Q)) ? (m_sum - 13'(Q)) : m_sum[11:0];
always @(*) begin always @(*) begin
st_next = st; st_next = st;
case (st) case (st)
@@ -202,7 +243,8 @@ module mlkem_top #(
ST_G: if (sha3_vo) st_next = ST_A; ST_G: if (sha3_vo) st_next = ST_A;
ST_A: if (a_pair >= K*K) st_next = ST_C; ST_A: if (a_pair >= K*K) st_next = ST_C;
ST_C: if (c_poly >= 2*K) st_next = ST_N; ST_C: if (c_poly >= 2*K) st_next = ST_N;
ST_N: if (n_slot >= 2*K) st_next = ST_DONE; ST_N: if (n_slot >= 2*K) st_next = ST_M;
ST_M: if (m_i >= K) st_next = ST_DONE;
ST_DONE: st_next = ST_IDLE; ST_DONE: st_next = ST_IDLE;
default: st_next = ST_IDLE; default: st_next = ST_IDLE;
endcase endcase
@@ -230,6 +272,13 @@ module mlkem_top #(
n_widx <= 8'd0; n_widx <= 8'd0;
n_valid <= 1'b0; n_valid <= 1'b0;
n_pending <= 1'b0; n_pending <= 1'b0;
m_i <= 2'd0;
m_j <= 2'd0;
m_ld <= 9'd0;
m_oidx <= 8'd0;
m_loading <= 1'b0;
m_pending <= 1'b0;
pm_valid <= 1'b0;
end else begin end else begin
st <= st_next; st <= st_next;
@@ -351,6 +400,59 @@ module mlkem_top #(
n_pending <= 1'b0; n_pending <= 1'b0;
end end
end end
// Arm M stage when N finishes: start first (i=0,j=0) poly_mul load
if (st == ST_N && st_next == ST_M) begin
m_i <= 2'd0;
m_j <= 2'd0;
m_ld <= 9'd0;
m_oidx <= 8'd0;
m_loading <= 1'b1;
m_pending <= 1'b0;
pm_valid <= 1'b1;
end
// ---- ST_M: t_hat[i] = e_hat[i] + sum_j A[i][j] o s_hat[j] ----
if (st == ST_M) begin
// LOAD: stream 256 (A,shat) pairs into poly_mul
if (m_loading && pm_valid && pm_ready) begin
if (m_ld == 9'd255) begin
pm_valid <= 1'b0; // last pair presented
m_loading <= 1'b0;
m_ld <= 9'd0;
m_oidx <= 8'd0;
end else begin
m_ld <= m_ld + 9'd1;
end
end
// ACCUMULATE: each product coeff += e_hat (j==0) or running t_hat
if (pm_vo) begin
polymem[m_tslot*256 + m_oidx] <= m_accq;
if (m_oidx == 8'd255) begin
// finished this (i,j) term; advance
if (m_j + 2'd1 < K) begin
m_j <= m_j + 2'd1;
m_pending <= 1'b1; // next term, same row
end else begin
m_j <= 2'd0;
m_i <= m_i + 2'd1; // next row (or == K -> DONE)
if (m_i + 2'd1 < K) m_pending <= 1'b1;
end
end else begin
m_oidx <= m_oidx + 8'd1;
end
end
// Start next (i,j) poly_mul load once core is IDLE again
if (m_pending && pm_ready && !pm_vo) begin
pm_valid <= 1'b1;
m_loading <= 1'b1;
m_ld <= 9'd0;
m_oidx <= 8'd0;
m_pending <= 1'b0;
end
end
end end
end end