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

feat(enc): Encaps E3 - y_hat = NTT(y) in place

Browse Source 
ST_ENC_N reuses the ST_N forward-NTT datapath (mode=0, read-ahead load,
in-place writeback to bank_se). Slot count parameterized: KeyGen 2K (s,e),
Encaps K (y only; e1/e2 stay time-domain) via n_slot_max. bse rd/wr muxes
and the C->N arming extended for ST_ENC_N.

Verified (K=2 c0) y_hat[0..1] == ml-kem-r golden; E0/E1/E2 still pass
(E2 y check dropped since NTT overwrites y in place -> covered by E3);
K=3/4 run through E3 without timeout.
This commit is contained in:
FallenSigh
2026-06-29 02:01:37 +08:00
parent cdc5ce25b1
commit 8ed4d59546
2 changed files with 44 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
sync_rtl/top/TB/tb_mlkem_enc_katK_xsim.v
View File

@@ -99,6 +99,7 @@ module tb_mlkem_enc_katK_xsim;
if (KP == 2 && casenum == 0) begin
verify_e1;
verify_e2;
verify_e3;
end
$finish;
end
@@ -134,19 +135,18 @@ module tb_mlkem_enc_katK_xsim;
// e2 lives in bank_t rel slot 0 -> dbg slot = slot_t_rt = 8.
task verify_e2;
begin
$readmemh("sync_rtl/top/TB/vectors/encgold/ec_k2_c0_y_0.hex", gy, 0, 255);
$readmemh("sync_rtl/top/TB/vectors/encgold/ec_k2_c0_y_1.hex", gy, 256, 511);
$readmemh("sync_rtl/top/TB/vectors/encgold/ec_k2_c0_e1_0.hex", gy, 512, 767);
$readmemh("sync_rtl/top/TB/vectors/encgold/ec_k2_c0_e1_1.hex", gy, 768, 1023);
$readmemh("sync_rtl/top/TB/vectors/encgold/ec_k2_c0_e2.hex", ge2);
ce = 0;
// y0,y1,e1_0,e1_1 at bank_se dbg slots 4..7
for (slot = 0; slot < 4; slot = slot + 1)
// e1_0,e1_1 at bank_se dbg slots 6,7 (time-domain, untouched by NTT).
// (y at slots 4,5 is overwritten in place by y_hat -> checked in E3.)
for (slot = 2; slot < 4; slot = slot + 1)
for (idx = 0; idx < 256; idx = idx + 1) begin
dbg_slot_i = (4+slot); dbg_idx_i = idx[7:0];
@(posedge clk); @(posedge clk); @(posedge clk);
if (dbg_coeff_o !== gy[slot*256+idx]) begin
if (ce < 8) $display(" Y[s%0d,%0d] got=%03x exp=%03x", slot, idx, dbg_coeff_o, gy[slot*256+idx]);
if (ce < 8) $display(" E1[s%0d,%0d] got=%03x exp=%03x", slot-2, idx, dbg_coeff_o, gy[slot*256+idx]);
ce = ce + 1;
end
end
@@ -159,9 +159,29 @@ module tb_mlkem_enc_katK_xsim;
ce = ce + 1;
end
end
if (ce == 0) $display("K=2 CASE 0 PASS (E2): y[0..1],e1[0..1],e2 == ml-kem-r golden");
if (ce == 0) $display("K=2 CASE 0 PASS (E2): e1[0..1],e2 == ml-kem-r golden");
else $display("K=2 CASE 0 FAIL (E2): %0d coeff mismatches", ce);
end
endtask
// y_hat[i] = NTT(y[i]) in place at bank_se rel slots 0..K-1 -> dbg slots 4..5 (K=2).
reg [11:0] gyh [0:2*256-1];
task verify_e3;
begin
$readmemh("sync_rtl/top/TB/vectors/encgold/ec_k2_c0_yhat_0.hex", gyh, 0, 255);
$readmemh("sync_rtl/top/TB/vectors/encgold/ec_k2_c0_yhat_1.hex", gyh, 256, 511);
ce = 0;
for (slot = 0; slot < 2; slot = slot + 1)
for (idx = 0; idx < 256; idx = idx + 1) begin
dbg_slot_i = (4+slot); dbg_idx_i = idx[7:0];
@(posedge clk); @(posedge clk); @(posedge clk);
if (dbg_coeff_o !== gyh[slot*256+idx]) begin
if (ce < 8) $display(" YH[s%0d,%0d] got=%03x exp=%03x", slot, idx, dbg_coeff_o, gyh[slot*256+idx]);
ce = ce + 1;
end
end
if (ce == 0) $display("K=2 CASE 0 PASS (E3): y_hat[0..1] == ml-kem-r golden");
else $display("K=2 CASE 0 FAIL (E3): %0d coeff mismatches", ce);
end
endtask
initial begin #120000000; $display("FAIL: global timeout"); $finish; end
endmodule

35
sync_rtl/top/mlkem_top.v
View File

@@ -241,7 +241,7 @@ module mlkem_top #(
dbg_t_addr[PT_AW-1:0];
// bank_se read port: ST_N load (s/e NTT), ST_M load (pm_b s_hat[j]) vs
// accumulate (e_hat, mutually exclusive via m_loading), ST_E dk-half, dbg.
assign bse_rd_addr = (st == ST_N) ? ntt_rd_full[PSE_AW-1:0] :
assign bse_rd_addr = (st == ST_N || st == ST_ENC_N) ? ntt_rd_full[PSE_AW-1:0] :
(st == ST_M) ? (m_loading ? pm_b_full[PSE_AW-1:0]
: m_eacc_full[PSE_AW-1:0]) :
(st == ST_E) ? e_rd_full[PSE_AW-1:0] :
@@ -249,14 +249,14 @@ module mlkem_top #(
// bank_se write port: ST_C CBD store (rel slot c_poly), ST_N NTT writeback
// (rel slot n_slot). Disjoint states.
// bank_se write: KeyGen ST_C CBD (s/e), Encaps ST_ENC_C CBD (y/e1, c_poly<2K;
// e2 at c_poly==2K goes to bank_t instead), ST_N NTT writeback. Disjoint.
// e2 at c_poly==2K goes to bank_t instead), ST_N/ST_ENC_N NTT writeback.
assign bse_we = ((st == ST_C) && c_busy && cbd_vo && cbd_ack) ||
((st == ST_ENC_C) && c_busy && cbd_vo && cbd_ack &&
(c_poly < {1'b0, k_r, 1'b0})) ||
((st == ST_N) && ntt_vo);
assign bse_wa = (st == ST_N) ? ((n_slot*256 + n_widx) & ((1<<PSE_AW)-1))
(((st == ST_N) || (st == ST_ENC_N)) && ntt_vo);
assign bse_wa = (st == ST_N || st == ST_ENC_N) ? ((n_slot*256 + n_widx) & ((1<<PSE_AW)-1))
: ((c_poly*256 + c_widx) & ((1<<PSE_AW)-1));
assign bse_wd = (st == ST_N) ? ntt_coeff : cbd_modq;
assign bse_wd = (st == ST_N || st == ST_ENC_N) ? ntt_coeff : cbd_modq;
always @(posedge clk) begin
if (dbg_slot_i >= slot_t_rt) dbg_coeff_r <= bt_rd_data; // bank_t (sd_bram)
else if (dbg_slot_i >= slot_s_rt) dbg_coeff_r <= bse_rd_data; // bank_se (sd_bram)
@@ -596,6 +596,8 @@ module mlkem_top #(
reg n_loading; // 1 while presenting load addresses to bank_se
reg n_pending; // waiting for ntt_core IDLE to start next slot
wire [SAW-1:0] n_slot_addr = slot_s_rt + n_slot; // s_hat then e_hat contiguous
// NTT slot count: KeyGen processes 2K (s,e); Encaps processes K (y only).
wire [4:0] n_slot_max = (st == ST_ENC_N) ? {2'b0, k_r} : {1'b0, k_r, 1'b0};
wire ntt_ready;
wire [11:0] ntt_coeff;
@@ -726,7 +728,8 @@ module mlkem_top #(
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_C;
ST_ENC_C: if (c_poly >= {k_r, 1'b1}) st_next = ST_DONE; // E2: stop (2K+1 polys: y,e1,e2)
ST_ENC_C: if (c_poly >= {k_r, 1'b1}) st_next = ST_ENC_N; // 2K+1 polys done
ST_ENC_N: if (n_slot >= {2'b0, k_r}) st_next = ST_DONE; // E3: K slots (y_hat)
ST_ENC_TDEC: if (td_done) st_next = ST_DONE; // (TDEC deferred to V-prep later)
ST_DONE: st_next = ST_IDLE;
default: st_next = ST_IDLE;
@@ -926,7 +929,8 @@ module mlkem_top #(
// Arm N stage when C finishes: prime load of slot S0. n_ridx is a
// read-ahead pointer; bank_se read is registered inside sd_bram (bse_rd_data)
// and fed to ntt_core one cycle later, so valid starts low (priming).
if (st == ST_C && st_next == ST_N) begin
if ((st == ST_C && st_next == ST_N) ||
(st == ST_ENC_C && st_next == ST_ENC_N)) begin
n_slot <= 3'd0;
n_ridx <= 9'd0;
n_widx <= 8'd0;
@@ -935,12 +939,11 @@ module mlkem_top #(
n_pending <= 1'b0;
end
// ---- ST_N: forward NTT each of S0,S1,E0,E1 in place ----
if (st == ST_N) begin
// LOAD phase: present read-ahead addr to bank_se (bse_rd_addr);
// sd_bram registers it, so bse_rd_data is consumed by ntt_core
// one cycle later (n_valid). Cores hold ready high through LOAD,
// so a fixed 1-cycle skew suffices (no backpressure gating).
// ---- ST_N / ST_ENC_N: forward NTT in place. KeyGen: 2K slots
// (s,e). Encaps: K slots (y only; e1/e2 stay time-domain). ----
if (st == ST_N || st == ST_ENC_N) begin
// slot-count bound: 2K for KeyGen, K for Encaps
// (n_slot_max below); same LOAD/OUTPUT cadence either way.
if (n_loading) begin
if (n_ridx == 9'd256) begin
// 256th coeff consumed this cycle; stop presenting addr
@@ -953,20 +956,18 @@ module mlkem_top #(
end
// OUTPUT phase: collect 256 results, write back to same slot.
// The bank_se write is the combinational bse_we/bse_wa/bse_wd
// assigns (rel slot n_slot); here only advance the write index.
if (ntt_vo) begin
n_widx <= n_widx + 8'd1; // wraps 255->0 after last
end
// Slot complete when ntt_core returns to DONE
if (ntt_done) begin
if (n_slot + 3'd1 < {1'b0, k_r, 1'b0}) begin
if (n_slot + 3'd1 < n_slot_max) begin
n_slot <= n_slot + 3'd1;
n_widx <= 8'd0;
n_pending <= 1'b1; // wait one cycle for core IDLE
end else begin
n_slot <= n_slot + 3'd1; // == 2K -> ST_DONE
n_slot <= n_slot + 3'd1; // == n_slot_max -> exit
end
end