From 8ed4d59546018e705b1c47c725f8908b02043dcd Mon Sep 17 00:00:00 2001 From: FallenSigh Date: Mon, 29 Jun 2026 02:01:37 +0800 Subject: [PATCH] feat(enc): Encaps E3 - y_hat = NTT(y) in place 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. --- sync_rtl/top/TB/tb_mlkem_enc_katK_xsim.v | 32 ++++++++++++++++++---- sync_rtl/top/mlkem_top.v | 35 ++++++++++++------------ 2 files changed, 44 insertions(+), 23 deletions(-) diff --git a/sync_rtl/top/TB/tb_mlkem_enc_katK_xsim.v b/sync_rtl/top/TB/tb_mlkem_enc_katK_xsim.v index 4025cbf..61c76e9 100644 --- a/sync_rtl/top/TB/tb_mlkem_enc_katK_xsim.v +++ b/sync_rtl/top/TB/tb_mlkem_enc_katK_xsim.v @@ -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 diff --git a/sync_rtl/top/mlkem_top.v b/sync_rtl/top/mlkem_top.v index 78744b4..2587fc9 100644 --- a/sync_rtl/top/mlkem_top.v +++ b/sync_rtl/top/mlkem_top.v @@ -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<= 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