feat(mlkem_top): KeyGen stages 2a-2c (G, SampleNTT A_hat, CBD s/e)

Fresh valid/ready KeyGen FSM for ML-KEM-512 (K=2, eta1=3). Independent
keccak per consumer (no arbiter). Verified stage-by-stage vs ml-kem-r golden:
- 2a G(d||K): rho/sigma exact (d byte0-low, K at byte32, no reversal).
- 2b SampleNTT: A_hat[i][j] from seed rho||j||i, 1024/1024 coeffs exact.
- 2c CBD: s[i]=CBD3(PRF(sigma,i)), e[i]=CBD3(PRF(sigma,K+i)); signed->mod-q
  (+Q when negative); 2048/2048 (A+s+e) coeffs exact.

polymem register array (10 slots x 256), debug readback tap (dbg_slot/idx ->
coeff, rho/sigma taps) for stage TBs. a_busy/c_busy guards (defensive after
sample_ntt fix). FSM: IDLE->G->A->C->DONE (datapath extended in later stages).

Plan + progress doc in .claude/plans/keygen_plan.md.
This commit is contained in:
2026-06-28 01:41:44 +08:00
parent 6db3c7cc5e
commit 2f206a6bc5
4 changed files with 2430 additions and 0 deletions

View File

@@ -0,0 +1,45 @@
// tb_mlkem_kg_2a_xsim.v - Stage 2a: verify G(d||K) -> rho/sigma in mlkem_top.
`timescale 1ns/1ps
module tb_mlkem_kg_2a_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;
localparam [255:0] RHO_EXP = 256'h15f74355ca862c3cdf3dab780c35cf24b88bf144706090a1c17e41205f9f1379;
localparam [255:0] SIG_EXP = 256'h69b042001b5630b1a039116cbfd29f62c0bde5a6b571504a9fcce68bed667fd5;
integer c;
initial begin
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<1000) begin @(posedge clk); c=c+1; end
if(!done_o) begin $display("FAIL: timeout"); $finish; end
$display("=== Stage 2a: G(d||K) ===");
if (dbg_rho_o===RHO_EXP && dbg_sigma_o===SIG_EXP) begin
$display("PASS: rho = %064x", dbg_rho_o);
$display("PASS: sigma = %064x", dbg_sigma_o);
$display("ALL TESTS PASSED");
end else begin
$display("FAIL:");
$display(" rho got=%064x", dbg_rho_o);
$display(" rho exp=%064x", RHO_EXP);
$display(" sig got=%064x", dbg_sigma_o);
$display(" sig exp=%064x", SIG_EXP);
end
$finish;
end
initial begin #50000; $display("FAIL: global timeout"); $finish; end
endmodule

View File

@@ -0,0 +1,60 @@
// tb_mlkem_kg_2c_xsim.v - Stage 2c: verify A_hat + s + e stored in mlkem_top.
// Reads golden kg_c000_AsE.hex (8 polys x 256 = 2048 lines, mod-q) and checks
// polymem slots A00,A01,A10,A11,S0,S1,E0,E1 via the debug readback tap.
`timescale 1ns/1ps
module tb_mlkem_kg_2c_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:2047];
// slot order matches golden file order
reg [3:0] slot_of [0:7];
integer c, p, idx, errors, gi;
initial begin
$readmemh("sync_rtl/top/TB/vectors/kg_c000_AsE.hex", gold);
slot_of[0]=4'd0; slot_of[1]=4'd1; slot_of[2]=4'd2; slot_of[3]=4'd3; // A00..A11
slot_of[4]=4'd4; slot_of[5]=4'd5; // S0,S1
slot_of[6]=4'd6; slot_of[7]=4'd7; // E0,E1
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<200000) begin @(posedge clk); c=c+1; end
if(!done_o) begin $display("FAIL: timeout after %0d cyc", c); $finish; end
$display("=== Stage 2c: A_hat + s + e (8 polys) === done in %0d cyc", c);
errors = 0;
for (p = 0; p < 8; 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); // 2 cyc for registered readback
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 (2048/2048 coeffs)");
else $display("TESTS FAILED: %0d mismatches", errors);
$finish;
end
initial begin #5000000; $display("FAIL: global timeout"); $finish; end
endmodule

File diff suppressed because it is too large Load Diff

277
sync_rtl/top/mlkem_top.v Normal file
View File

@@ -0,0 +1,277 @@
// mlkem_top.v - ML-KEM-512 KeyGen top-level integration (K=2, eta1=3).
//
// Streaming valid/ready interface. Given seeds d and z, computes the
// ML-KEM key pair per FIPS 203 Algorithm 16 (KeyGen_internal):
// (rho,sigma) = G(d || K)
// A_hat[i][j] = SampleNTT(rho || j || i) i,j in 0..K-1
// s[i] = CBD3(PRF(sigma, i)), e[i] = CBD3(PRF(sigma, K+i))
// s_hat[i] = NTT(s[i]); e_hat[i] = NTT(e[i])
// t_hat[i] = e_hat[i] + sum_j A_hat[i][j] o s_hat[j]
// ek = byteEncode12(t_hat[0..K-1]) || rho
// dk = byteEncode12(s_hat[0..K-1]) || ek || H(ek) || z
//
// Built incrementally and verified stage-by-stage against ml-kem-r golden
// vectors (test_framework/modules/mlkem_keygen/golden) and NIST KAT.
//
// Uses independent (verified) leaf modules, each with its own keccak_core:
// sha3_top, sample_ntt_sync, sample_cbd_sync, ntt_core, poly_mul_sync,
// mod_add_sync. No shared-keccak arbiter.
`include "sync_rtl/common/defines.vh"
module mlkem_top #(
parameter K = 2, // ML-KEM-512
parameter ETA1 = 3
) (
input clk,
input rst_n,
input [255:0] d_i, // KeyGen seed d (byte 0 in d_i[7:0])
input [255:0] z_i, // implicit-rejection seed z
input start_i, // pulse to begin KeyGen
output busy_o, // high while running
output done_o, // pulse when ek/dk ready
// Debug readback tap: read one stored coefficient by (poly slot, index).
// Lets stage TBs verify intermediates without wide buses.
input [3:0] dbg_slot_i, // poly slot (see localparams below)
input [7:0] dbg_idx_i, // coefficient index 0..255
output [11:0] dbg_coeff_o,
// Debug taps for hash outputs
output [255:0] dbg_rho_o,
output [255:0] dbg_sigma_o
);
localparam Q = `Q; // 3329
// ================================================================
// Polynomial storage: K=2 needs A_hat[2][2]=4, s/s_hat[2], e/e_hat[2],
// t_hat[2]. Reuse slots: s and s_hat share (NTT in place), same for e.
// Slot map:
// 0..3 : A_hat[0][0],A_hat[0][1],A_hat[1][0],A_hat[1][1]
// 4..5 : s_hat[0], s_hat[1] (s[i] then overwritten by NTT)
// 6..7 : e_hat[0], e_hat[1] (e[i] then overwritten by NTT)
// 8..9 : t_hat[0], t_hat[1]
// ================================================================
localparam SLOT_A00 = 4'd0, SLOT_A01 = 4'd1, SLOT_A10 = 4'd2, SLOT_A11 = 4'd3;
localparam SLOT_S0 = 4'd4, SLOT_S1 = 4'd5;
localparam SLOT_E0 = 4'd6, SLOT_E1 = 4'd7;
localparam SLOT_T0 = 4'd8, SLOT_T1 = 4'd9;
localparam NUM_SLOTS = 10;
reg [11:0] polymem [0:NUM_SLOTS*256-1];
// Debug readback (registered for timing)
reg [11:0] dbg_coeff_r;
always @(posedge clk) dbg_coeff_r <= polymem[dbg_slot_i*256 + dbg_idx_i];
assign dbg_coeff_o = dbg_coeff_r;
// ================================================================
// Top-level FSM (built incrementally). Stage 2a: G only.
// ================================================================
localparam ST_IDLE = 4'd0;
localparam ST_G = 4'd1; // run G(d||K), capture rho/sigma
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_DONE = 4'd15;
reg [3:0] st, st_next;
reg [255:0] rho_r, sigma_r;
// A-generation bookkeeping
reg [2:0] a_pair; // 0..K*K (=4) pairs done
reg [7:0] a_widx; // write index 0..255 within current poly
reg a_busy; // 1 once current pair's request accepted (gates collect)
wire [1:0] a_i = a_pair[1] ? 2'd1 : 2'd0; // pair/K (K=2)
wire [1:0] a_j = a_pair[0] ? 2'd1 : 2'd0; // pair%K
wire [3:0] a_slot = {2'b0, a_pair[1], a_pair[0]}; // SLOT_A00..A11 = pair index
// C-generation bookkeeping: 2*K polys = s0,s1,e0,e1 (idx 0..3)
reg [2:0] c_poly; // 0..2K
reg [7:0] c_widx;
reg c_busy; // 1 once current poly's request accepted (gates collect)
wire [7:0] c_nonce = {5'b0, c_poly}; // s:0,1 e:2,3 == nonce
// slot: c_poly 0->S0,1->S1,2->E0,3->E1
wire [3:0] c_slot = (c_poly == 3'd0) ? SLOT_S0 :
(c_poly == 3'd1) ? SLOT_S1 :
(c_poly == 3'd2) ? SLOT_E0 : SLOT_E1;
assign busy_o = (st != ST_IDLE);
assign done_o = (st == ST_DONE);
assign dbg_rho_o = rho_r;
assign dbg_sigma_o = sigma_r;
// ---- sha3_top in G mode: data_i = {K_byte, d} (d byte0 in [7:0]) ----
reg sha3_valid;
wire sha3_ready;
wire [511:0] sha3_hash;
wire sha3_vo;
reg sha3_ack; // consumer ready for hash
wire [511:0] g_data = {248'b0, 8'(K), d_i}; // data_i[263:256]=K, [255:0]=d
sha3_top u_sha3 (
.clk(clk), .rst_n(rst_n),
.mode(2'b00), // G = SHA3-512
.data_i(g_data),
.valid_i(sha3_valid),
.ready_o(sha3_ready),
.hash_o(sha3_hash),
.valid_o(sha3_vo),
.ready_i(sha3_ack),
.mb_en(1'b0), .mb_block_i(1088'b0), .mb_valid_i(1'b0),
.mb_last_i(1'b0), .mb_ready_o()
);
// ---- sample_ntt_sync: Â[i][j] = SampleNTT(rho || j || i) ----
reg snt_valid;
wire snt_ready;
wire [11:0] snt_coeff;
wire snt_vo;
wire snt_last;
reg snt_ack; // we accept coeffs
sample_ntt_sync #(.K(K)) u_snt (
.clk(clk), .rst_n(rst_n),
.rho_i(rho_r),
.k_i(3'(K)),
.i_idx(a_i),
.j_idx(a_j),
.valid_i(snt_valid),
.ready_o(snt_ready),
.coeff_o(snt_coeff),
.valid_o(snt_vo),
.ready_i(snt_ack),
.last_o(snt_last)
);
// ---- sample_cbd_sync: s[i]=CBD3(PRF(sigma,i)), e[i]=CBD3(PRF(sigma,K+i)) ----
reg cbd_valid;
wire cbd_ready;
wire [11:0] cbd_coeff; // 12-bit signed (two's complement)
wire cbd_vo;
wire cbd_last;
reg cbd_ack;
sample_cbd_sync u_cbd (
.clk(clk), .rst_n(rst_n),
.seed_i(sigma_r),
.nonce_i(c_nonce),
.eta_i(2'(ETA1)),
.valid_i(cbd_valid),
.ready_o(cbd_ready),
.coeff_o(cbd_coeff),
.valid_o(cbd_vo),
.ready_i(cbd_ack),
.last_o(cbd_last)
);
// signed (two's complement) -> [0,Q): add Q when negative
wire [11:0] cbd_modq = cbd_coeff[11] ? (cbd_coeff + 12'(Q)) : cbd_coeff;
always @(*) begin
st_next = st;
case (st)
ST_IDLE: if (start_i) st_next = ST_G;
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_DONE;
ST_DONE: st_next = ST_IDLE;
default: st_next = ST_IDLE;
endcase
end
always @(posedge clk or negedge rst_n) begin
if (!rst_n) begin
st <= ST_IDLE;
rho_r <= 256'd0;
sigma_r <= 256'd0;
sha3_valid <= 1'b0;
sha3_ack <= 1'b0;
snt_valid <= 1'b0;
snt_ack <= 1'b0;
a_pair <= 3'd0;
a_widx <= 8'd0;
a_busy <= 1'b0;
cbd_valid <= 1'b0;
cbd_ack <= 1'b0;
c_poly <= 3'd0;
c_widx <= 8'd0;
c_busy <= 1'b0;
end else begin
st <= st_next;
// Kick off G when entering ST_G
if (st == ST_IDLE && start_i) begin
sha3_valid <= 1'b1;
sha3_ack <= 1'b1;
end
// Drop valid once accepted
if (sha3_valid && sha3_ready) sha3_valid <= 1'b0;
// Capture rho/sigma when G completes; arm A stage
if (st == ST_G && sha3_vo) begin
rho_r <= sha3_hash[255:0]; // rho = G output bytes 0..31
sigma_r <= sha3_hash[511:256]; // sigma = bytes 32..63
sha3_ack <= 1'b0;
snt_valid <= 1'b1; // start first SampleNTT
snt_ack <= 1'b1;
a_pair <= 3'd0;
a_widx <= 8'd0;
a_busy <= 1'b0;
end
// ---- ST_A: drive SampleNTT, store 256 coeffs per pair ----
if (st == ST_A) begin
// mark busy once this pair's request accepted
if (snt_valid && snt_ready) begin
snt_valid <= 1'b0;
a_busy <= 1'b1;
end
// store each output coefficient only while busy (ignore stale last coeff from prior poly)
if (a_busy && snt_vo && snt_ack) begin
polymem[a_slot*256 + a_widx] <= snt_coeff;
if (snt_last) begin
// finished this poly; advance to next pair
a_pair <= a_pair + 3'd1;
a_widx <= 8'd0;
a_busy <= 1'b0;
// start next SampleNTT if more pairs remain
if (a_pair + 3'd1 < K*K) snt_valid <= 1'b1;
end else begin
a_widx <= a_widx + 8'd1;
end
end
end
// Arm C stage when A finishes
if (st == ST_A && st_next == ST_C) begin
cbd_valid <= 1'b1;
cbd_ack <= 1'b1;
c_poly <= 3'd0;
c_widx <= 8'd0;
c_busy <= 1'b0;
end
// ---- ST_C: drive CBD, store 256 mod-q coeffs per poly ----
if (st == ST_C) begin
if (cbd_valid && cbd_ready) begin
cbd_valid <= 1'b0;
c_busy <= 1'b1;
end
if (c_busy && cbd_vo && cbd_ack) begin
polymem[c_slot*256 + c_widx] <= cbd_modq;
if (cbd_last) begin
c_poly <= c_poly + 3'd1;
c_widx <= 8'd0;
c_busy <= 1'b0;
if (c_poly + 3'd1 < 2*K) cbd_valid <= 1'b1;
end else begin
c_widx <= c_widx + 8'd1;
end
end
end
end
end
endmodule