feat(enc): Encaps E1 - rho load + A regen + byteDecode12 t_hat

ST_ENC_LOAD: stream rho (32B) from ek_bram[384k..] into rho_r (read-ahead,
1-cyc bram latency). ST_ENC_A: regenerate A_hat via SampleNTT into bank_a
(reuses ST_A datapath, gated on st==ST_A||ST_ENC_A). ST_ENC_TDEC:
byteDecode12 ek[i*384..] -> t_hat[i] into bank_t (5-cycle micro-phase per
3-byte/2-coeff triple; bt write port muxed with ST_M).

Verified (K=2 c0) A_hat (1024 coeffs) + t_hat (512) == KeyGen golden via
dbg_coeff_o; E0 ss==KAT.ss still passes all K/cases (no timeout).
This commit is contained in:
2026-06-29 01:44:50 +08:00
parent 0a8b3dae69
commit 31c967c8a4
2 changed files with 173 additions and 9 deletions

View File

@@ -91,7 +91,50 @@ module tb_mlkem_enc_katK_xsim;
if (errors == 0) $display("K=%0d CASE %0d PASS (E0): ss == KAT.ss", KP, casenum);
else $display("K=%0d CASE %0d FAIL (E0): %0d ss mismatches", KP, casenum, errors);
// ---- E1: verify A_hat (slots 0..K^2-1) and t_hat (slots slot_t..) ----
// The KAT case-0 Encaps reuses ek=pk from KeyGen case-0, so regenerated
// A_hat and decoded t_hat equal KeyGen's golden vectors (K=2 c0 only).
if (KP == 2 && casenum == 0) verify_e1;
$finish;
end
// golden coeff arrays (K=2 c0): A_hat[i][j] and t_hat[i], 256 coeffs each
reg [11:0] ga [0:4*256-1]; // A00,A01,A10,A11 concatenated (slot*256+idx)
reg [11:0] gt [0:2*256-1]; // t_hat0, t_hat1
integer ce, slot, idx;
task verify_e1;
begin
$readmemh("test_framework/modules/mlkem_keygen/golden/c000_Ahat_0_0.hex", ga, 0, 255);
$readmemh("test_framework/modules/mlkem_keygen/golden/c000_Ahat_0_1.hex", ga, 256, 511);
$readmemh("test_framework/modules/mlkem_keygen/golden/c000_Ahat_1_0.hex", ga, 512, 767);
$readmemh("test_framework/modules/mlkem_keygen/golden/c000_Ahat_1_1.hex", ga, 768, 1023);
$readmemh("test_framework/modules/mlkem_keygen/golden/c000_that_0.hex", gt, 0, 255);
$readmemh("test_framework/modules/mlkem_keygen/golden/c000_that_1.hex", gt, 256, 511);
ce = 0;
// A_hat: slots 0..3 (i*k+j)
for (slot = 0; slot < 4; slot = slot + 1)
for (idx = 0; idx < 256; idx = idx + 1) begin
dbg_slot_i = slot[3:0]; dbg_idx_i = idx[7:0];
@(posedge clk); @(posedge clk); @(posedge clk);
if (dbg_coeff_o !== ga[slot*256+idx]) begin
if (ce < 8) $display(" A[s%0d,%0d] got=%03x exp=%03x", slot, idx, dbg_coeff_o, ga[slot*256+idx]);
ce = ce + 1;
end
end
// t_hat: slots slot_t_rt(=8) .. +1
for (slot = 0; slot < 2; slot = slot + 1)
for (idx = 0; idx < 256; idx = idx + 1) begin
dbg_slot_i = (8+slot); dbg_idx_i = idx[7:0];
@(posedge clk); @(posedge clk); @(posedge clk);
if (dbg_coeff_o !== gt[slot*256+idx]) begin
if (ce < 12) $display(" T[%0d,%0d] got=%03x exp=%03x", slot, idx, dbg_coeff_o, gt[slot*256+idx]);
ce = ce + 1;
end
end
if (ce == 0) $display("K=2 CASE 0 PASS (E1): A_hat (1024) + t_hat (512) == KeyGen golden");
else $display("K=2 CASE 0 FAIL (E1): %0d coeff mismatches", ce);
end
endtask
initial begin #120000000; $display("FAIL: global timeout"); $finish; end
endmodule

View File

@@ -101,6 +101,33 @@ module mlkem_top #(
wire [11:0] c2_bytes_rt = 12'd32 * {7'b0, dv_rt}; // 128/128/160
wire [11:0] ct_bytes_rt = c1_bytes_rt + c2_bytes_rt; // 768/1088/1568
assign ss_o = ss_r;
// ---- E1: rho-load + byteDecode12 bookkeeping ----
// rho load: read 32 bytes from ek_bram at offset 384*k into rho_r.
reg [5:0] rl_idx; // 0..33 read-ahead pointer over rho bytes
reg [5:0] rl_widx; // 0..31 byte being written (1 cyc behind rl_idx)
reg rl_vld; // a byte is arriving this cycle
wire [11:0] rl_ek_off = dk_bytes_rt + {6'b0, rl_idx}; // ek addr = 384k + idx
// byteDecode12: read 3 ek bytes -> 2 coeffs into bank_t. Walk K polys * 128
// triples. 5-cycle micro-phase per triple: ph0..2 fetch b0,b1,b2 (read-ahead,
// 1-cycle bram latency), ph3 write c0, ph4 write c1.
reg [2:0] td_poly; // 0..k-1 t_hat poly index
reg [7:0] td_trip; // 0..127 triple within poly
reg [2:0] td_ph; // 0..4 micro-phase
reg [7:0] td_b0, td_b1, td_b2; // captured ek bytes
reg td_done;
// ek byte address for the triple: poly*384 + trip*3 + (ph for ph<3)
wire [11:0] td_base = {td_poly, 7'd0} * 12'd3; // poly*384 (=poly*256*1.5)... see below
// poly*384 + trip*3 + byteidx; byteidx = ph (0,1,2) during fetch phases
wire [11:0] td_byteidx = {2'd0, td_trip, 1'b0} + {3'd0, td_trip}; // trip*3
wire [11:0] td_ekaddr = ({4'd0,td_poly}*12'd384) + td_byteidx + {9'd0, td_ph};
// decoded coeffs: c0 = b0 | ((b1&0xF)<<8); c1 = (b1>>4) | (b2<<4)
wire [11:0] td_c0 = {td_b1[3:0], td_b0};
wire [11:0] td_c1 = {td_b2, td_b1[7:4]};
// bank_t write for TDEC (registered); muxed with ST_M write below.
reg td_we;
reg [9:0] td_wa; // PT_AW=10 (declared below; literal here)
reg [11:0] td_wd;
// ct readback tap (ct_bram added in E5/E7); tied off until then.
/* verilator lint_off UNUSEDSIGNAL */
wire [10:0] dbg_ct_idx_unused = dbg_ct_idx_i;
@@ -151,8 +178,8 @@ module mlkem_top #(
.clk(clk), .rd_addr(ba_rd_addr), .rd_data(ba_rd_data),
.wr_en(ba_we), .wr_addr(ba_wa), .wr_data(ba_wd)
);
// ST_A write: commit snt_coeff to A_hat[a_slot] the cycle it is accepted.
assign ba_we = (st == ST_A) && a_busy && snt_vo && snt_ack;
// ST_A / ST_ENC_A write: commit snt_coeff to A_hat[a_slot] when accepted.
assign ba_we = (st == ST_A || st == ST_ENC_A) && a_busy && snt_vo && snt_ack;
assign ba_wa = (a_slot*256 + a_widx) & ((1<<PA_AW)-1);
assign ba_wd = snt_coeff;
@@ -185,10 +212,12 @@ module mlkem_top #(
.clk(clk), .rd_addr(bt_rd_addr), .rd_data(bt_rd_data),
.wr_en(bt_we), .wr_addr(bt_wa), .wr_data(bt_wd)
);
// ST_M accumulate write: t_hat[m_i][m_oidx] <= (acc + product) mod Q when pm_vo.
assign bt_we = (st == ST_M) && pm_vo;
assign bt_wa = (m_i*256 + m_oidx) & ((1<<PT_AW)-1);
assign bt_wd = m_accq;
// bank_t write: KeyGen ST_M accumulate, or Encaps TDEC byteDecode12.
// ST_M: t_hat[m_i][m_oidx] <= (acc+product) mod Q when pm_vo.
// ST_ENC_TDEC: t_hat[td_poly][2*td_trip(+1)] <= decoded coeff (td_we/wa/wd).
assign bt_we = ((st == ST_M) && pm_vo) || td_we;
assign bt_wa = td_we ? td_wa : ((m_i*256 + m_oidx) & ((1<<PT_AW)-1));
assign bt_wd = td_we ? td_wd : m_accq;
// Debug readback (registered for timing)
reg [11:0] dbg_coeff_r;
@@ -474,6 +503,8 @@ module mlkem_top #(
// single read port serves dbg_byte_o (ek loop) then dbg_dk_o (dk loop).
wire [11:0] dbgdk_ek_off = dbg_dk_idx_i - dk_bytes_rt; // offset within ek region
assign ek_rd_addr = (st == ST_H || st == ST_ENC_H) ? h_g_addr[10:0] :
(st == ST_ENC_LOAD) ? rl_ek_off[10:0] : // rho load
(st == ST_ENC_TDEC) ? td_ekaddr[10:0] : // byteDecode12
dbgdk_in_ek ? dbgdk_ek_off[10:0]
: dbg_byte_idx_i;
// dkp BRAM read-address mux: dbg_byte (sel=1) or dbg_dk (dkp region).
@@ -675,7 +706,10 @@ module mlkem_top #(
ST_H: if (h_phase == 2'd3) st_next = ST_DONE;
// ---- Encaps ----
ST_ENC_H: if (h_phase == 2'd3) st_next = ST_ENC_G;
ST_ENC_G: if (sha3_vo) st_next = ST_DONE; // E0: stop here (K,r captured)
ST_ENC_G: if (sha3_vo) st_next = ST_ENC_LOAD; // (K,r) captured
ST_ENC_LOAD: if (rl_idx >= 6'd33) st_next = ST_ENC_A; // rho loaded
ST_ENC_A: if (a_pair >= kk_rt) st_next = ST_ENC_TDEC;
ST_ENC_TDEC: if (td_done) st_next = ST_DONE; // E1: stop here (A, t_hat ready)
ST_DONE: st_next = ST_IDLE;
default: st_next = ST_IDLE;
endcase
@@ -689,6 +723,19 @@ module mlkem_top #(
m_r <= 256'd0;
ss_r <= 256'd0;
r_r <= 256'd0;
rl_idx <= 6'd0;
rl_widx <= 6'd0;
rl_vld <= 1'b0;
td_poly <= 3'd0;
td_trip <= 8'd0;
td_ph <= 3'd0;
td_b0 <= 8'd0;
td_b1 <= 8'd0;
td_b2 <= 8'd0;
td_done <= 1'b0;
td_we <= 1'b0;
td_wa <= 10'd0;
td_wd <= 12'd0;
rho_r <= 256'd0;
sigma_r <= 256'd0;
sha3_valid <= 1'b0;
@@ -751,6 +798,7 @@ module mlkem_top #(
// BRAM write-enables default low; pulsed where a byte is written.
ek_we <= 1'b0;
dkp_we <= 1'b0;
td_we <= 1'b0; // TDEC bank_t write default low
// Kick off when entering from IDLE: KeyGen starts G; Encaps captures
// op/m and arms the H(ek) machinery (ST_ENC_H reuses the ST_H FSM).
@@ -789,8 +837,8 @@ module mlkem_top #(
a_busy <= 1'b0;
end
// ---- ST_A: drive SampleNTT, store 256 coeffs per pair ----
if (st == ST_A) begin
// ---- ST_A / ST_ENC_A: drive SampleNTT, store 256 coeffs per pair ----
if (st == ST_A || st == ST_ENC_A) begin
// mark busy once this pair's request accepted
if (snt_valid && snt_ready) begin
snt_valid <= 1'b0;
@@ -1106,6 +1154,79 @@ module mlkem_top #(
r_r <= sha3_hash[511:256]; // r = G bytes 32..63
sha3_ack <= 1'b0;
end
// Arm rho-load when entering ST_ENC_LOAD. rho = ek[384k .. 384k+31].
if (st == ST_ENC_G && st_next == ST_ENC_LOAD) begin
rl_idx <= 6'd0;
rl_widx <= 6'd0;
rl_vld <= 1'b0;
end
// ---- ST_ENC_LOAD: stream 32 rho bytes from ek_bram into rho_r ----
// ek read is registered (1-cyc latency): the data arriving this cycle
// is for the address presented last cycle (rl_idx-1). Present addr
// rl_idx (valid 0..31); write rho_r[rl_idx-1] when rl_idx in 1..32.
if (st == ST_ENC_LOAD) begin
if (rl_idx >= 6'd1 && rl_idx <= 6'd32)
rho_r[(rl_idx-6'd1)*8 +: 8] <= ek_rd_data;
rl_idx <= rl_idx + 6'd1; // exit when rl_idx >= 33
end
// Arm ST_ENC_A (regenerate A_hat via SampleNTT, same as ST_A).
if (st == ST_ENC_LOAD && st_next == ST_ENC_A) begin
snt_valid <= 1'b1;
snt_ack <= 1'b1;
a_pair <= 5'd0;
a_i <= 3'd0;
a_j <= 3'd0;
a_widx <= 8'd0;
a_busy <= 1'b0;
end
// Arm ST_ENC_TDEC (byteDecode12: ek -> t_hat in bank_t).
if (st == ST_ENC_A && st_next == ST_ENC_TDEC) begin
td_poly <= 3'd0;
td_trip <= 8'd0;
td_ph <= 3'd0;
td_done <= 1'b0;
end
// ---- ST_ENC_TDEC: byteDecode12 ek -> t_hat[0..k-1] in bank_t ----
// 5-cycle micro-phase per triple (read-ahead, 1-cyc bram latency):
// ph0: present b0 addr; ph1: capture b0, present b1; ph2: capture
// b1, present b2; ph3: capture b2 + write c0; ph4: write c1, advance.
if (st == ST_ENC_TDEC && !td_done) begin
// capture the byte that arrived for the address presented last cycle
if (td_ph == 3'd1) td_b0 <= ek_rd_data;
if (td_ph == 3'd2) td_b1 <= ek_rd_data;
if (td_ph == 3'd3) td_b2 <= ek_rd_data;
// write decoded coeffs into bank_t (rel slot = td_poly)
if (td_ph == 3'd3) begin
td_we <= 1'b1;
td_wa <= (td_poly*256 + {td_trip, 1'b0}) & ((1<<PT_AW)-1); // 2*trip
td_wd <= td_c0;
end
if (td_ph == 3'd4) begin
td_we <= 1'b1;
td_wa <= (td_poly*256 + {td_trip, 1'b0} + 1) & ((1<<PT_AW)-1); // 2*trip+1
td_wd <= td_c1;
end
// advance micro-phase / triple / poly
if (td_ph == 3'd4) begin
td_ph <= 3'd0;
if (td_trip == 8'd127) begin
td_trip <= 8'd0;
if (td_poly + 3'd1 < k_r) td_poly <= td_poly + 3'd1;
else td_done <= 1'b1;
end else begin
td_trip <= td_trip + 8'd1;
end
end else begin
td_ph <= td_ph + 3'd1;
end
end
end
end