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 61c76e9..8b5763b 100644 --- a/sync_rtl/top/TB/tb_mlkem_enc_katK_xsim.v +++ b/sync_rtl/top/TB/tb_mlkem_enc_katK_xsim.v @@ -100,6 +100,7 @@ module tb_mlkem_enc_katK_xsim; verify_e1; verify_e2; verify_e3; + verify_e4; end $finish; end @@ -135,32 +136,20 @@ 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_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); + $readmemh("sync_rtl/top/TB/vectors/encgold/ec_k2_c0_e2.hex", ge2); ce = 0; - // 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(" E1[s%0d,%0d] got=%03x exp=%03x", slot-2, idx, dbg_coeff_o, gy[slot*256+idx]); - ce = ce + 1; - end - end - // e2 at bank_t dbg slot 8 + // e2 at bank_t dbg slot 8 (survives; e1 is consumed/overwritten by u + // in E4 -> e1 correctness is transitively verified by E4's u check). for (idx = 0; idx < 256; idx = idx + 1) begin dbg_slot_i = 8; dbg_idx_i = idx[7:0]; @(posedge clk); @(posedge clk); @(posedge clk); if (dbg_coeff_o !== ge2[idx]) begin - if (ce < 12) $display(" E2[%0d] got=%03x exp=%03x", idx, dbg_coeff_o, ge2[idx]); + if (ce < 8) $display(" E2[%0d] got=%03x exp=%03x", idx, dbg_coeff_o, ge2[idx]); ce = ce + 1; end end - 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); + if (ce == 0) $display("K=2 CASE 0 PASS (E2): e2 == ml-kem-r golden"); + else $display("K=2 CASE 0 FAIL (E2): %0d e2 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). @@ -183,5 +172,25 @@ module tb_mlkem_enc_katK_xsim; else $display("K=2 CASE 0 FAIL (E3): %0d coeff mismatches", ce); end endtask + // u[i] = INTT(sum A^T o y_hat) + e1[i] over e1 in bank_se rel K+i -> dbg 6,7 (K=2). + reg [11:0] gu [0:2*256-1]; + task verify_e4; + begin + $readmemh("sync_rtl/top/TB/vectors/encgold/ec_k2_c0_u_0.hex", gu, 0, 255); + $readmemh("sync_rtl/top/TB/vectors/encgold/ec_k2_c0_u_1.hex", gu, 256, 511); + ce = 0; + for (slot = 0; slot < 2; slot = slot + 1) + for (idx = 0; idx < 256; idx = idx + 1) begin + dbg_slot_i = (6+slot); dbg_idx_i = idx[7:0]; // bank_se rel K+slot (K=2) + @(posedge clk); @(posedge clk); @(posedge clk); + if (dbg_coeff_o !== gu[slot*256+idx]) begin + if (ce < 8) $display(" U[s%0d,%0d] got=%03x exp=%03x", slot, idx, dbg_coeff_o, gu[slot*256+idx]); + ce = ce + 1; + end + end + if (ce == 0) $display("K=2 CASE 0 PASS (E4): u[0..1] == ml-kem-r golden"); + else $display("K=2 CASE 0 FAIL (E4): %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 2587fc9..3826422 100644 --- a/sync_rtl/top/mlkem_top.v +++ b/sync_rtl/top/mlkem_top.v @@ -216,12 +216,21 @@ module mlkem_top #( // e2 (ST_ENC_C c_poly==2K -> bank_t rel slot 0, free during C/N/U). wire e2_we = (st == ST_ENC_C) && c_busy && cbd_vo && cbd_ack && (c_poly == {1'b0, k_r, 1'b0}); // c_poly == 2K - assign bt_we = ((st == ST_M) && pm_vo) || td_we || e2_we; + // Encaps E4 bank_t writes: MAC psum (u_sub0, on pm_vo) and INTT in-place + // (u_sub1, on ntt_vo). Both target bank_t rel slot UPSUM. + wire u_psum_we = (st == ST_ENC_U) && (u_sub == 2'd0) && pm_vo; + wire u_intt_we = (st == ST_ENC_U) && (u_sub == 2'd1) && ntt_vo; + assign bt_we = ((st == ST_M) && pm_vo) || td_we || e2_we || + u_psum_we || u_intt_we; assign bt_wa = td_we ? td_wa : e2_we ? (c_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) @@ -606,14 +628,19 @@ module mlkem_top #( // bank_se read addr for the NTT load (relative slot = n_slot); sd_bram // registers it into bse_rd_data, which feeds ntt_core 1 cycle later. wire [13:0] ntt_rd_full = n_slot*256 + n_ridx[7:0]; - wire [11:0] ntt_in = bse_rd_data; + // ntt_core inputs muxed: KeyGen/Encaps fwd NTT feeds bse_rd_data (mode 0); + // Encaps E4 INTT sub-phase feeds psum (bt_rd_data) with mode 1. + wire u_intt_act = (st == ST_ENC_U) && (u_sub == 2'd1); + wire [11:0] ntt_in = u_intt_act ? bt_rd_data : bse_rd_data; + wire ntt_vin = u_intt_act ? u_nvalid : n_valid; + wire ntt_mode = u_intt_act ? 1'b1 : 1'b0; ntt_core u_ntt ( .clk(clk), .rst_n(rst_n), .coeff_in(ntt_in), - .valid_i(n_valid), + .valid_i(ntt_vin), .ready_o(ntt_ready), - .mode(1'b0), // forward NTT, no scaling + .mode(ntt_mode), // 0 = fwd (NTT), 1 = inv (INTT, x3303) .coeff_out(ntt_coeff), .valid_o(ntt_vo), .ready_i(1'b1), // always accept output @@ -669,6 +696,53 @@ module mlkem_top #( wire [SAW-1:0] m_eslot = slot_e_rt + m_i; // e_hat[i] wire [SAW-1:0] m_tslot = slot_t_rt + m_i; // t_hat[i] + // ================================================================ + // E4: u[i] = INTT(sum_j A^T[i][j] o y_hat[j]) + e1[i] (Encaps ST_ENC_U) + // Per row i, three sub-phases (u_sub): + // 0 MAC : sum_j A_hat[j][i] o y_hat[j] (transpose) -> psum bank_t[UPSUM] + // 1 INTT: INTT(psum) mode=1 in place in bank_t[UPSUM] + // 2 ADD : u[i][w] = psum[w] + e1[i][w] mod Q -> bank_se rel (K+i) (over e1) + // y_hat (bank_se 0..K-1) survives for V. Reuses shared u_pmul + u_ntt. + // ================================================================ + localparam UPSUM = 10'd1; // bank_t rel slot for the NTT-domain psum + reg [2:0] u_row; // output row i 0..K-1 + reg [1:0] u_sub; // 0=MAC 1=INTT 2=ADD + // MAC bookkeeping (mirrors ST_M) + reg [2:0] u_j; // term col j 0..K-1 + reg [8:0] u_ld; // load read-ahead 0..256 + reg [7:0] u_oidx; // accumulate index 0..255 + reg u_loading; // streaming pairs into poly_mul + reg u_pending; // wait poly_mul IDLE before next j + reg u_j0q; // (u_j==0) delayed 1 cyc (first-term=0 select) + // INTT bookkeeping (mirrors ST_N) + reg [8:0] u_ridx; // INTT load read-ahead 0..256 + reg [7:0] u_widx; // INTT output index 0..255 + reg u_nvalid; // feed psum to ntt_core (mode=1) + reg u_nloading; // presenting INTT load addresses + reg u_npending; // (unused placeholder for symmetry) + // ADD bookkeeping + reg [8:0] u_aidx; // ADD read-ahead 0..256 + reg [7:0] u_awidx; // ADD write index 0..255 + reg u_avalid; // ADD write enable pipeline + // MAC addresses: A^T -> A_hat[u_j][u_row] slot = u_j*k+u_row; y_hat[u_j] rel u_j + wire [SAW-1:0] u_aslot = u_j*k_r + u_row; + wire [13:0] u_pm_a_full = u_aslot*256 + u_ld[7:0]; // bank_a + wire [13:0] u_pm_b_full = u_j*256 + u_ld[7:0]; // bank_se y_hat[u_j] + wire [7:0] u_acc_radr = pm_vo ? (u_oidx + 8'd1) : u_oidx; + wire [13:0] u_psum_full = UPSUM*256 + u_acc_radr; // bank_t psum acc addr (MAC) + wire [11:0] u_acc_src = u_j0q ? 12'd0 : bt_rd_data; // first term init 0 + wire [12:0] u_sum = {1'b0, u_acc_src} + {1'b0, pm_coeff}; + wire [11:0] u_accq = (u_sum >= 13'(Q)) ? (u_sum - 13'(Q)) : u_sum[11:0]; + // INTT load addr (psum read) + in-place writeback to bank_t[UPSUM] + wire [13:0] u_intt_rd = UPSUM*256 + u_ridx[7:0]; + wire [13:0] u_intt_wr = UPSUM*256 + u_widx; + // ADD: read psum (bank_t) + e1 (bank_se K+u_row), write u over e1 (read-ahead 1) + wire [13:0] u_add_prd = UPSUM*256 + u_aidx[7:0]; // bank_t psum + wire [13:0] u_add_e1rd = ({2'b0,k_r}+{2'b0,u_row})*256 + u_aidx[7:0]; // bank_se e1 (lead) + wire [13:0] u_add_uwr = ({2'b0,k_r}+{2'b0,u_row})*256 + u_awidx; // bank_se u (write) + wire [12:0] u_usum = {1'b0, bt_rd_data} + {1'b0, bse_rd_data}; // psum + e1 + wire [11:0] u_uq = (u_usum >= 13'(Q)) ? (u_usum - 13'(Q)) : u_usum[11:0]; + reg pm_valid; wire pm_ready; wire [11:0] pm_coeff; @@ -729,7 +803,8 @@ module mlkem_top #( 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_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_N: if (n_slot >= {2'b0, k_r}) st_next = ST_ENC_U; // K slots (y_hat) + ST_ENC_U: if (u_row >= k_r) st_next = ST_DONE; // E4: u[0..K-1] done 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; @@ -787,6 +862,22 @@ module mlkem_top #( m_pending <= 1'b0; pm_valid <= 1'b0; m_jq <= 1'b0; + u_row <= 3'd0; + u_sub <= 2'd0; + u_j <= 3'd0; + u_ld <= 9'd0; + u_oidx <= 8'd0; + u_loading <= 1'b0; + u_pending <= 1'b0; + u_j0q <= 1'b0; + u_ridx <= 9'd0; + u_widx <= 8'd0; + u_nvalid <= 1'b0; + u_nloading <= 1'b0; + u_npending <= 1'b0; + u_aidx <= 9'd0; + u_awidx <= 8'd0; + u_avalid <= 1'b0; e_poly <= 3'd0; e_pair <= 8'd0; e_ph <= 2'd0; @@ -1046,6 +1137,115 @@ module mlkem_top #( end end + // Arm Encaps U stage when N finishes: prime first row MAC (i=0,j=0). + if (st == ST_ENC_N && st_next == ST_ENC_U) begin + u_row <= 3'd0; + u_sub <= 2'd0; // MAC + u_j <= 3'd0; + u_ld <= 9'd0; + u_oidx <= 8'd0; + u_loading <= 1'b1; + u_pending <= 1'b0; + pm_valid <= 1'b0; + end + + // ---- ST_ENC_U sub-phase 0: MAC sum_j A^T[u_row][j] o y_hat[j] ---- + // psum (init 0 at j==0) accumulates into bank_t[UPSUM]. Mirrors ST_M + // load/accumulate cadence (read-ahead by 1, j-select via u_j0q). + if (st == ST_ENC_U && u_sub == 2'd0) begin + u_j0q <= (u_j == 3'd0); + + if (u_loading) begin + if (u_ld == 9'd256) begin + pm_valid <= 1'b0; + u_loading <= 1'b0; + u_ld <= 9'd0; + u_oidx <= 8'd0; + end else begin + u_ld <= u_ld + 9'd1; + pm_valid <= 1'b1; + end + end + + if (pm_vo) begin + // bank_t psum write is combinational (u_psum_we); advance idx + if (u_oidx == 8'd255) begin + if (u_j + 3'd1 < k_r) begin + u_j <= u_j + 3'd1; + u_pending <= 1'b1; // next term, same row + end else begin + // row's MAC done -> INTT sub-phase + u_sub <= 2'd1; + u_ridx <= 9'd0; + u_widx <= 8'd0; + u_nvalid <= 1'b0; + u_nloading <= 1'b1; + end + end else begin + u_oidx <= u_oidx + 8'd1; + end + end + + // re-prime next term's poly_mul load + if (u_pending && pm_ready && !pm_vo) begin + u_loading <= 1'b1; + u_ld <= 9'd0; + u_oidx <= 8'd0; + pm_valid <= 1'b0; + u_pending <= 1'b0; + end + end + + // ---- ST_ENC_U sub-phase 1: INTT(psum) mode=1 in place (bank_t) ---- + if (st == ST_ENC_U && u_sub == 2'd1) begin + if (u_nloading) begin + if (u_ridx == 9'd256) begin + u_nloading <= 1'b0; + u_nvalid <= 1'b0; + end else begin + u_ridx <= u_ridx + 9'd1; + u_nvalid <= 1'b1; + end + end + if (ntt_vo) u_widx <= u_widx + 8'd1; // in-place write (u_intt_we) + if (ntt_done) begin + // INTT done -> ADD sub-phase (psum + e1[u_row]) + u_sub <= 2'd2; + u_aidx <= 9'd0; + u_awidx <= 8'd0; + u_avalid <= 1'b0; + end + end + + // ---- ST_ENC_U sub-phase 2: u[u_row] = psum + e1[u_row] mod Q ---- + // Read psum (bank_t) + e1 (bank_se) at u_aidx; both arrive 1 cyc + // later, so register (valid,widx) and write u over e1 next cycle. + if (st == ST_ENC_U && u_sub == 2'd2) begin + if (u_aidx < 9'd256) begin + u_aidx <= u_aidx + 9'd1; + u_avalid <= 1'b1; // addr presented this cyc -> write next + u_awidx <= u_aidx[7:0]; + end else begin + u_avalid <= 1'b0; // no more addresses + end + // u write (bse_we via u_add_we) commits index u_awidx this cycle. + // Advance row after the 256th coeff (u_awidx==255) is written. + if (u_avalid && u_awidx == 8'd255) begin + if (u_row + 3'd1 < k_r) begin + u_row <= u_row + 3'd1; + u_sub <= 2'd0; // next row MAC + u_j <= 3'd0; + u_ld <= 9'd0; + u_oidx <= 8'd0; + u_loading <= 1'b1; + u_pending <= 1'b0; + pm_valid <= 1'b0; + end else begin + u_row <= u_row + 3'd1; // == K -> ST_DONE + end + end + end + // Arm E stage when M finishes if (st == ST_M && st_next == ST_E) begin e_poly <= 3'd0;