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

feat(mlkem_top): KeyGen stage 4 - H(ek) + full dk, end-to-end KAT pass

Browse Source 
Add ST_H stage: second sha3_top (mb_en=1) computes H(ek) over 800B ek as 6
pre-padded SHA3-256 rate blocks. Per block: assemble 136 bytes (h_padbyte
applies 0x06...0x80 padding on final block) into h_block_r, feed (hold valid
until mb_ready drops), wait permute; capture digest on last block into hek_r.

Full dk readback tap: dk = dk_pke(768) || ek(800) || H(ek)(32) || z(32) = 1632B.

End-to-end TB (tb_mlkem_kg_kat_xsim, no force/release): drive KAT count=0 d/z,
run full KeyGen FSM (IDLE->G->A->C->N->M->E->H->DONE), verify:
  ek  == KAT pk (800B)  byte-exact
  dk  == KAT sk (1632B) byte-exact
Done in 21403 cycles. ML-KEM-512 KeyGen complete and KAT-verified.
Prior stage TBs (2c/2e/2f) still pass (no regression).
This commit is contained in:
FallenSigh
2026-06-28 02:18:52 +08:00
parent 17914911c3
commit 9824ed8f2c
3 changed files with 1819 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

68
sync_rtl/top/TB/tb_mlkem_kg_kat_xsim.v Normal file
View File

@@ -0,0 +1,68 @@
// tb_mlkem_kg_kat_xsim.v - Stage 4 end-to-end: full ML-KEM-512 KeyGen vs NIST KAT.
// Drives d/z (KAT count=0), runs KeyGen, verifies:
// ek (800B, sel=0) == KAT pk
// dk (1632B, dk tap) == KAT sk (= dk_pke || ek || H(ek) || z)
// No force/release — pure valid/ready via start_i/done_o.
`timescale 1ns/1ps
module tb_mlkem_kg_kat_xsim;
reg clk=0, rst_n=0, start_i=0;
reg [255:0] d_i, z_i;
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;
reg dbg_byte_sel_i=0; reg [9:0] dbg_byte_idx_i=0; wire [7:0] dbg_byte_o;
reg [11:0] dbg_dk_idx_i=0; wire [7:0] dbg_dk_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_byte_sel_i(dbg_byte_sel_i), .dbg_byte_idx_i(dbg_byte_idx_i), .dbg_byte_o(dbg_byte_o),
.dbg_dk_idx_i(dbg_dk_idx_i), .dbg_dk_o(dbg_dk_o),
.dbg_rho_o(dbg_rho_o), .dbg_sigma_o(dbg_sigma_o)
);
always #5 clk = ~clk;
// KAT count=0 (byte0-low literals)
localparam [255:0] D_LIT = 256'h2426f1941779574d3f1b163bd57f7e173e229e630ec7f7073bdf365137c4bb6d;
localparam [255:0] Z_LIT = 256'h687acf9406694974d383032f7579378f449c75d0560af56cf921ec48404896f6;
reg [7:0] ek_gold [0:799];
reg [7:0] dk_gold [0:1631];
integer c, i, errors;
initial begin
$readmemh("sync_rtl/top/TB/vectors/c000_ek_bytes.hex", ek_gold);
$readmemh("sync_rtl/top/TB/vectors/c000_dk_full_bytes.hex", dk_gold);
d_i = D_LIT; z_i = Z_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<600000) begin @(posedge clk); c=c+1; end
if(!done_o) begin $display("FAIL: timeout after %0d cyc", c); $finish; end
$display("=== Stage 4: ML-KEM-512 KeyGen end-to-end vs NIST KAT === done in %0d cyc", c);
errors = 0;
// ek == KAT pk (800B)
dbg_byte_sel_i = 1'b0;
for (i = 0; i < 800; i = i + 1) begin
dbg_byte_idx_i = i[9:0]; @(posedge clk); @(posedge clk);
if (dbg_byte_o !== ek_gold[i]) begin
if (errors < 8) $display(" EK[%0d] got=%02x exp=%02x", i, dbg_byte_o, ek_gold[i]);
errors = errors + 1;
end
end
// dk == KAT sk (1632B)
for (i = 0; i < 1632; i = i + 1) begin
dbg_dk_idx_i = i[11:0]; @(posedge clk); @(posedge clk);
if (dbg_dk_o !== dk_gold[i]) begin
if (errors < 8) $display(" DK[%0d] got=%02x exp=%02x", i, dbg_dk_o, dk_gold[i]);
errors = errors + 1;
end
end
if (errors == 0) $display("ALL TESTS PASSED: ek==KAT.pk (800B), dk==KAT.sk (1632B)");
else $display("TESTS FAILED: %0d byte mismatches", errors);
$finish;
end
initial begin #40000000; $display("FAIL: global timeout"); $finish; end
endmodule

1632
sync_rtl/top/TB/vectors/c000_dk_full_bytes.hex Normal file
View File

File diff suppressed because it is too large Load Diff

120
sync_rtl/top/mlkem_top.v
View File

@@ -40,6 +40,10 @@ module mlkem_top #(
input dbg_byte_sel_i,
input [9:0] dbg_byte_idx_i,
output [7:0] dbg_byte_o,
// Debug full-dk readback: dk = dk_pke(768) || ek(800) || H(ek)(32) || z(32)
// = 1632 bytes. Index 0..1631.
input [11:0] dbg_dk_idx_i,
output [7:0] dbg_dk_o,
// Debug taps for hash outputs
output [255:0] dbg_rho_o,
output [255:0] dbg_sigma_o
@@ -80,6 +84,20 @@ module mlkem_top #(
dbg_byte_r <= dbg_byte_sel_i ? dkp_mem[dbg_byte_idx_i] : ek_mem[dbg_byte_idx_i];
assign dbg_byte_o = dbg_byte_r;
// full dk = dk_pke(0..767) || ek(768..1567) || H(ek)(1568..1599) || z(1600..1631)
reg [7:0] dbg_dk_r;
always @(posedge clk) begin
if (dbg_dk_idx_i < 12'd768)
dbg_dk_r <= dkp_mem[dbg_dk_idx_i];
else if (dbg_dk_idx_i < 12'd1568)
dbg_dk_r <= ek_mem[dbg_dk_idx_i - 12'd768];
else if (dbg_dk_idx_i < 12'd1600)
dbg_dk_r <= hek_r[(dbg_dk_idx_i - 12'd1568)*8 +: 8];
else
dbg_dk_r <= z_i[(dbg_dk_idx_i - 12'd1600)*8 +: 8];
end
assign dbg_dk_o = dbg_dk_r;
// ================================================================
// Top-level FSM (built incrementally). Stage 2a: G only.
// ================================================================
@@ -90,6 +108,7 @@ module mlkem_top #(
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_E = 4'd6; // byteEncode12 -> ek_mem, dkp_mem
localparam ST_H = 4'd7; // H(ek) via multi-block SHA3-256
localparam ST_DONE = 4'd15;
reg [3:0] st, st_next;
@@ -139,6 +158,49 @@ module mlkem_top #(
.mb_last_i(1'b0), .mb_ready_o()
);
// ---- second sha3_top dedicated to multi-block H(ek) (SHA3-256, 800B->6 blk) ----
reg [1087:0] h_block_r; // current pre-padded rate block
reg h_mbvalid;
reg h_mblast;
wire h_mbready;
wire [511:0] h_hash;
wire h_vo;
reg h_ack;
reg [255:0] hek_r; // captured H(ek)
reg [2:0] h_blk; // 0..5 block index
reg [7:0] h_byte; // 0..135 byte within block being assembled
reg [1:0] h_phase; // 0=assemble 1=feed 2=wait-perm 3=done
sha3_top u_sha3_h (
.clk(clk), .rst_n(rst_n),
.mode(2'b01), // unused in mb mode
.data_i(512'b0),
.valid_i(1'b0),
.ready_o(),
.hash_o(h_hash),
.valid_o(h_vo),
.ready_i(h_ack),
.mb_en(1'b1),
.mb_block_i(h_block_r),
.mb_valid_i(h_mbvalid),
.mb_last_i(h_mblast),
.mb_ready_o(h_mbready)
);
// byte b (0..135) of block blk for SHA3-256 over 800-byte ek:
// global byte g = blk*136 + b; ek_mem[g] if g<800;
// g==800 -> 0x06 (domain+first pad bit); g==815 -> |0x80 (last block); else 0
function [7:0] h_padbyte(input [2:0] blk, input [7:0] b);
integer g;
begin
g = blk*136 + b;
if (g < 800) h_padbyte = ek_mem[g];
else if (g == 800) h_padbyte = (g == 815) ? 8'h86 : 8'h06;
else if (g == 815) h_padbyte = 8'h80; // 6th block (815 = 5*136+135)
else h_padbyte = 8'h00;
end
endfunction
// ---- sample_ntt_sync: Â[i][j] = SampleNTT(rho || j || i) ----
reg snt_valid;
wire snt_ready;
@@ -284,7 +346,8 @@ module mlkem_top #(
ST_C: if (c_poly >= 2*K) st_next = ST_N;
ST_N: if (n_slot >= 2*K) st_next = ST_M;
ST_M: if (m_i >= K) st_next = ST_E;
ST_E: if (e_done) st_next = ST_DONE;
ST_E: if (e_done) st_next = ST_H;
ST_H: if (h_phase == 2'd3) st_next = ST_DONE;
ST_DONE: st_next = ST_IDLE;
default: st_next = ST_IDLE;
endcase
@@ -323,6 +386,14 @@ module mlkem_top #(
e_pair <= 8'd0;
e_rho <= 10'd0;
e_done <= 1'b0;
h_block_r <= 1088'd0;
h_mbvalid <= 1'b0;
h_mblast <= 1'b0;
h_ack <= 1'b0;
hek_r <= 256'd0;
h_blk <= 3'd0;
h_byte <= 8'd0;
h_phase <= 2'd0;
end else begin
st <= st_next;
@@ -532,6 +603,53 @@ module mlkem_top #(
else e_rho <= e_rho + 10'd1;
end
end
// Arm H stage when E finishes
if (st == ST_E && st_next == ST_H) begin
h_blk <= 3'd0;
h_byte <= 8'd0;
h_phase <= 2'd0; // assemble
h_mbvalid<= 1'b0;
h_mblast <= 1'b0;
h_ack <= 1'b1; // ready to consume final digest
end
// ---- ST_H: H(ek) via multi-block SHA3-256 (6 pre-padded blocks) ----
if (st == ST_H) begin
case (h_phase)
// assemble 136 bytes of block h_blk into h_block_r
2'd0: begin
h_block_r[h_byte*8 +: 8] <= h_padbyte(h_blk, h_byte);
if (h_byte == 8'd135) begin
h_byte <= 8'd0;
h_mbvalid <= 1'b1;
h_mblast <= (h_blk == 3'd5);
h_phase <= 2'd1; // feed
end else begin
h_byte <= h_byte + 8'd1;
end
end
// feed: hold valid until accepted (mb_ready drops)
2'd1: begin
if (h_mbvalid && !h_mbready) begin
h_mbvalid <= 1'b0;
h_mblast <= 1'b0;
h_phase <= 2'd2; // wait permute
end
end
// wait permute done: ready again (more blocks) or digest valid (last)
2'd2: begin
if (h_vo) begin
hek_r <= h_hash[255:0];
h_phase <= 2'd3; // done
end else if (h_mbready) begin
h_blk <= h_blk + 3'd1;
h_phase <= 2'd0; // assemble next block
end
end
default: ; // 2'd3 done: hold
endcase
end
end
end