feat(mlkem_top): KeyGen stage 2e (matrix accumulate t_hat)
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:
57
sync_rtl/top/TB/tb_mlkem_kg_2e_xsim.v
Normal file
57
sync_rtl/top/TB/tb_mlkem_kg_2e_xsim.v
Normal 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
512
sync_rtl/top/TB/vectors/kg_c000_that.hex
Normal 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
|
||||
@@ -73,6 +73,7 @@ module mlkem_top #(
|
||||
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_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;
|
||||
|
||||
reg [3:0] st, st_next;
|
||||
@@ -195,6 +196,46 @@ module mlkem_top #(
|
||||
.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
|
||||
st_next = st;
|
||||
case (st)
|
||||
@@ -202,7 +243,8 @@ module mlkem_top #(
|
||||
ST_G: if (sha3_vo) st_next = ST_A;
|
||||
ST_A: if (a_pair >= K*K) st_next = ST_C;
|
||||
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;
|
||||
default: st_next = ST_IDLE;
|
||||
endcase
|
||||
@@ -230,6 +272,13 @@ module mlkem_top #(
|
||||
n_widx <= 8'd0;
|
||||
n_valid <= 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
|
||||
st <= st_next;
|
||||
|
||||
@@ -351,6 +400,59 @@ module mlkem_top #(
|
||||
n_pending <= 1'b0;
|
||||
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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user