From a9e50ebc0c64d4d7f71740f63c3f28d76c63784a Mon Sep 17 00:00:00 2001 From: FallenSigh Date: Sun, 28 Jun 2026 01:53:23 +0800 Subject: [PATCH] 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). --- sync_rtl/top/TB/tb_mlkem_kg_2e_xsim.v | 57 +++ sync_rtl/top/TB/vectors/kg_c000_that.hex | 512 +++++++++++++++++++++++ sync_rtl/top/mlkem_top.v | 104 ++++- 3 files changed, 672 insertions(+), 1 deletion(-) create mode 100644 sync_rtl/top/TB/tb_mlkem_kg_2e_xsim.v create mode 100644 sync_rtl/top/TB/vectors/kg_c000_that.hex diff --git a/sync_rtl/top/TB/tb_mlkem_kg_2e_xsim.v b/sync_rtl/top/TB/tb_mlkem_kg_2e_xsim.v new file mode 100644 index 0000000..d9cbe06 --- /dev/null +++ b/sync_rtl/top/TB/tb_mlkem_kg_2e_xsim.v @@ -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 diff --git a/sync_rtl/top/TB/vectors/kg_c000_that.hex b/sync_rtl/top/TB/vectors/kg_c000_that.hex new file mode 100644 index 0000000..88d7f5d --- /dev/null +++ b/sync_rtl/top/TB/vectors/kg_c000_that.hex @@ -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 diff --git a/sync_rtl/top/mlkem_top.v b/sync_rtl/top/mlkem_top.v index a25846a..70dcf44 100644 --- a/sync_rtl/top/mlkem_top.v +++ b/sync_rtl/top/mlkem_top.v @@ -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