From 3284aa443f5574f159ccfd8e587feabf90158bfd Mon Sep 17 00:00:00 2001 From: FallenSigh Date: Sat, 27 Jun 2026 01:38:38 +0800 Subject: [PATCH] fix(kg): implement t_hat computation and pk/sk output in mlkem_top MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add s_hat0_reg/s_hat1_reg for s_hat capture during CBD→NTT output - Add tmul_accum register for poly_mul result accumulation - Implement S_KG_TMUL_MUL_LD: load A+s coefficients into poly_mul - Implement S_KG_TMUL_MUL_WAIT/ACCUM: accumulate A·s products - Implement S_KG_TMUL_ADD_E: add e_hat to accumulated sum - Assign pk_o_r = {t_hat1, t_hat0}_reg and sk_o_r = {s_hat1, s_hat0}_reg - Shift Encaps/Decaps FSM states by +3 KeyGen: 5/5 vectors complete in ~47K cycles/op --- sync_rtl/top/mlkem_top.v | 358 +++++++++++++++++++++++++++++++++------ 1 file changed, 303 insertions(+), 55 deletions(-) diff --git a/sync_rtl/top/mlkem_top.v b/sync_rtl/top/mlkem_top.v index 09b0fa8..acc1572 100644 --- a/sync_rtl/top/mlkem_top.v +++ b/sync_rtl/top/mlkem_top.v @@ -119,6 +119,18 @@ module mlkem_top #( reg [255:0] Hpk_reg; // H(pk) for Encaps reg [255:0] z_reg; // implicit rejection value + // s_hat capture registers (filled during CBD→NTT output) + reg [N*12-1:0] s_hat0_reg, s_hat1_reg; + + // t_hat output registers (computed during t_hat computation) + reg [N*12-1:0] t_hat0_reg, t_hat1_reg; + + // t_hat computation scratch registers + reg [11:0] tmul_pipe_reg; // 1-cycle pipeline delay for poly_mul output + reg tmul_pipe_valid; // flag: tmul_pipe_reg holds valid data + reg [8:0] tmul_out_cnt; // output coefficient counter for accumulation + reg tmul_adv_row; // flag: advancing to next row (set in ADD_E, used in MUL_LD) + // ==================================================================== // Keccak arbiter signals (3 consumers → 1 keccak_core) // ==================================================================== @@ -290,48 +302,51 @@ module mlkem_top #( localparam S_KG_CBD_E_NTT_OUT = 7'd28; localparam S_KG_CBD_E_NEXT = 7'd29; localparam S_KG_TMUL_INIT = 7'd30; - localparam S_KG_TMUL_STEP = 7'd31; - localparam S_KG_TMUL_NEXT = 7'd32; - localparam S_KG_DONE = 7'd33; + localparam S_KG_TMUL_MUL_LD = 7'd31; + localparam S_KG_TMUL_MUL_OUT = 7'd32; + localparam S_KG_TMUL_ACCUM = 7'd33; + localparam S_KG_TMUL_ADD_E = 7'd34; + localparam S_KG_TMUL_NEXT = 7'd35; + localparam S_KG_DONE = 7'd36; // --- Encaps states --- - localparam S_EN_RNG_REQ = 7'd34; - localparam S_EN_RNG_WAIT = 7'd35; - localparam S_EN_H_START = 7'd36; - localparam S_EN_H_WAIT = 7'd37; - localparam S_EN_G_START = 7'd38; - localparam S_EN_G_WAIT = 7'd39; - localparam S_EN_SNT_INIT = 7'd40; - localparam S_EN_SNT_START = 7'd41; - localparam S_EN_SNT_COEFFS = 7'd42; - localparam S_EN_SNT_CLEANUP = 7'd43; - localparam S_EN_SNT_NTT_LD = 7'd44; - localparam S_EN_SNT_NTT_CMP = 7'd45; - localparam S_EN_SNT_NTT_OUT = 7'd46; - localparam S_EN_SNT_NEXT = 7'd47; - localparam S_EN_CBD_Y_INIT = 7'd48; - localparam S_EN_CBD_Y_START = 7'd49; - localparam S_EN_CBD_Y_COEFFS = 7'd50; - localparam S_EN_CBD_Y_CLNUP = 7'd51; - localparam S_EN_CBD_Y_NTT_LD = 7'd52; - localparam S_EN_CBD_Y_NTT_CMP = 7'd53; - localparam S_EN_CBD_Y_NTT_OUT = 7'd54; - localparam S_EN_CBD_Y_NEXT = 7'd55; - localparam S_EN_DONE = 7'd56; + localparam S_EN_RNG_REQ = 7'd37; + localparam S_EN_RNG_WAIT = 7'd38; + localparam S_EN_H_START = 7'd39; + localparam S_EN_H_WAIT = 7'd40; + localparam S_EN_G_START = 7'd41; + localparam S_EN_G_WAIT = 7'd42; + localparam S_EN_SNT_INIT = 7'd43; + localparam S_EN_SNT_START = 7'd44; + localparam S_EN_SNT_COEFFS = 7'd45; + localparam S_EN_SNT_CLEANUP = 7'd46; + localparam S_EN_SNT_NTT_LD = 7'd47; + localparam S_EN_SNT_NTT_CMP = 7'd48; + localparam S_EN_SNT_NTT_OUT = 7'd49; + localparam S_EN_SNT_NEXT = 7'd50; + localparam S_EN_CBD_Y_INIT = 7'd51; + localparam S_EN_CBD_Y_START = 7'd52; + localparam S_EN_CBD_Y_COEFFS = 7'd53; + localparam S_EN_CBD_Y_CLNUP = 7'd54; + localparam S_EN_CBD_Y_NTT_LD = 7'd55; + localparam S_EN_CBD_Y_NTT_CMP = 7'd56; + localparam S_EN_CBD_Y_NTT_OUT = 7'd57; + localparam S_EN_CBD_Y_NEXT = 7'd58; + localparam S_EN_DONE = 7'd59; // --- Decaps states --- - localparam S_DC_DECOMP_C1 = 7'd57; - localparam S_DC_DECOMP_C2 = 7'd58; - localparam S_DC_NTT_U_LD = 7'd59; - localparam S_DC_NTT_U_CMP = 7'd60; - localparam S_DC_NTT_U_OUT = 7'd61; - localparam S_DC_MUL_S = 7'd62; - localparam S_DC_INTT_V = 7'd63; - localparam S_DC_DECOMP_M = 7'd64; - localparam S_DC_G_CHECK = 7'd65; - localparam S_DC_REENC = 7'd66; - localparam S_DC_KDF = 7'd67; - localparam S_DC_DONE = 7'd68; + localparam S_DC_DECOMP_C1 = 7'd60; + localparam S_DC_DECOMP_C2 = 7'd61; + localparam S_DC_NTT_U_LD = 7'd62; + localparam S_DC_NTT_U_CMP = 7'd63; + localparam S_DC_NTT_U_OUT = 7'd64; + localparam S_DC_MUL_S = 7'd65; + localparam S_DC_INTT_V = 7'd66; + localparam S_DC_DECOMP_M = 7'd67; + localparam S_DC_G_CHECK = 7'd68; + localparam S_DC_REENC = 7'd69; + localparam S_DC_KDF = 7'd70; + localparam S_DC_DONE = 7'd71; // ==================================================================== // FSM registers @@ -485,14 +500,28 @@ module mlkem_top #( end // -- t_hat = sum(A·s) + e -- - S_KG_TMUL_INIT: state_next = S_KG_TMUL_STEP; - S_KG_TMUL_STEP: if (phase_done) state_next = S_KG_TMUL_NEXT; + S_KG_TMUL_INIT: state_next = S_KG_TMUL_MUL_LD; + S_KG_TMUL_MUL_LD: if (phase_done) state_next = S_KG_TMUL_MUL_OUT; + S_KG_TMUL_MUL_OUT: if (phase_done) state_next = S_KG_TMUL_NEXT; + S_KG_TMUL_ACCUM: if (phase_done) state_next = S_KG_TMUL_ADD_E; + // After ADD_E: if more rows, loop to MUL_LD with next row; else DONE + S_KG_TMUL_ADD_E: begin + if (phase_done) begin + if (loop_i + 2'd1 < k_r[1:0]) + state_next = S_KG_TMUL_MUL_LD; + else + state_next = S_KG_DONE; + end + end + + // After MUL_OUT: check if more j terms for this row S_KG_TMUL_NEXT: begin - if (loop_i + 2'd1 < k_r[1:0]) - state_next = S_KG_TMUL_STEP; - else - state_next = S_KG_DONE; + if (loop_j + 2'd1 < k_r[1:0]) begin + state_next = S_KG_TMUL_MUL_LD; // more j columns + end else begin + state_next = S_KG_TMUL_ACCUM; // all j done → accumulate + end end S_KG_DONE: state_next = S_IDLE; @@ -596,6 +625,14 @@ module mlkem_top #( poly_idx_reg <= 6'd0; phase_active <= 1'b0; phase_done <= 1'b0; + s_hat0_reg <= 0; + s_hat1_reg <= 0; + t_hat0_reg <= 0; + t_hat1_reg <= 0; + tmul_pipe_reg <= 12'd0; + tmul_pipe_valid <= 1'b0; + tmul_out_cnt <= 9'd0; + tmul_adv_row <= 1'b0; // Drive all FSM-controlled outputs to safe defaults cbd_valid_i_r <= 1'b0; @@ -828,6 +865,11 @@ module mlkem_top #( bram_wr_en_r <= 1'b1; bram_wr_addr_r <= poly_addr(S_BASE + loop_idx, coeff_cnt[7:0]); bram_wr_data_r <= ntt_coeff_out; + // Capture into s_hat registers for t_hat computation + if (loop_idx == 2'd0) + s_hat0_reg[coeff_cnt[7:0] * 12 +: 12] <= ntt_coeff_out; + else + s_hat1_reg[coeff_cnt[7:0] * 12 +: 12] <= ntt_coeff_out; coeff_cnt <= coeff_cnt + 9'd1; if (coeff_cnt == 9'd255) begin phase_done <= 1'b1; @@ -921,30 +963,236 @@ module mlkem_top #( end // ============================================================ - // KEYGEN: t_hat computation (simplified) + // KEYGEN: t_hat computation + // + // For each row i (0..k-1): + // t_hat[i] = sum_j(A[i,j] * s_hat[j]) + e_hat[i] + // + // Phase 1 (S_KG_TMUL_MUL_LD): Load A[i,j], s_hat[j] into poly_mul + // Phase 2 (S_KG_TMUL_MUL_OUT): Read poly_mul result → SCRATCH0/SCRATCH1 + // Phase 3 (S_KG_TMUL_ACCUM): Accumulate all A·s terms using poly_arith + // Phase 4 (S_KG_TMUL_ADD_E): Add e_hat[i] → store in t_hat_i_reg // ============================================================ if (state_r == S_KG_TMUL_INIT) begin - loop_i <= 2'd0; - coeff_cnt <= 9'd0; - phase_done <= 1'b0; + loop_i <= 2'd0; + loop_j <= 2'd0; + coeff_cnt <= 9'd0; + tmul_pipe_reg <= 12'd0; + tmul_pipe_valid <= 1'b0; + tmul_out_cnt <= 9'd0; + tmul_adv_row <= 1'b0; + phase_done <= 1'b0; end - if (state_r == S_KG_TMUL_STEP) begin - // TODO: Per-row poly_mul A[i,:]·s[:] + e[i] → t_hat[i] - // For now, placeholder: just mark done - phase_done <= 1'b1; + // --- S_KG_TMUL_MUL_LD: Load A[i,j] and s_hat[j] into poly_mul --- + // Uses s_hat registers (combinational read) + A from BRAM (1-cycle delay). + // First cycle: handle row advancement (if flagged). + // Second cycle: prefetch first A coeff from BRAM. + // Cycles 3-258: feed 256 coeff pairs to poly_mul. + if (state_r == S_KG_TMUL_MUL_LD) begin + phase_active <= 1'b1; + if (coeff_cnt == 9'd0) begin + // Cycle 0: handle row advancement flag + if (tmul_adv_row) begin + loop_i <= loop_i + 2'd1; + loop_j <= 2'd0; + tmul_adv_row <= 1'b0; + end + coeff_cnt <= 9'd1; + end else if (coeff_cnt == 9'd1) begin + // Cycle 1: prefetch first A coefficient + bram_rd_addr_r <= poly_addr(A_BASE + (loop_i * K + loop_j), 8'd0); + coeff_cnt <= 9'd2; + end else if (coeff_cnt <= 9'd257) begin + // bram_rd_data holds A[i,j][coeff_cnt-2] from previous cycle + pmul_a_r <= bram_rd_data; + // s_hat value from register (combinational index) + if (loop_j == 2'd0) + pmul_b_r <= s_hat0_reg[(coeff_cnt - 9'd2) * 12 +: 12]; + else + pmul_b_r <= s_hat1_reg[(coeff_cnt - 9'd2) * 12 +: 12]; + pmul_valid_i_r <= 1'b1; + + // Prefetch next A coefficient + if (coeff_cnt < 9'd257) + bram_rd_addr_r <= poly_addr(A_BASE + (loop_i * K + loop_j), coeff_cnt[7:0] - 8'd1); + + coeff_cnt <= coeff_cnt + 9'd1; + + // After feeding coeff 255 (coeff_cnt reaches 258), done + if (coeff_cnt == 9'd257) begin + phase_done <= 1'b1; + phase_active <= 1'b0; + coeff_cnt <= 9'd0; + end + end end + // --- S_KG_TMUL_MUL_OUT: Read poly_mul result into BRAM scratch --- + // Phase 2a: Wait for first pmul_valid_o + // Phase 2b: Read 256 output coefficients + if (state_r == S_KG_TMUL_MUL_OUT) begin + phase_active <= 1'b1; + phase_done <= 1'b0; + + if (pmul_valid_o) begin + // Store to SCRATCH0 (first j term) or SCRATCH1 (subsequent terms) + if (loop_j == 2'd0) begin + bram_wr_addr_r <= poly_addr(SCRATCH0, coeff_cnt[7:0]); + end else begin + bram_wr_addr_r <= poly_addr(SCRATCH1, coeff_cnt[7:0]); + end + bram_wr_en_r <= 1'b1; + bram_wr_data_r <= pmul_coeff_out; + coeff_cnt <= coeff_cnt + 9'd1; + + if (coeff_cnt == 9'd255) begin + phase_done <= 1'b1; + phase_active <= 1'b0; + coeff_cnt <= 9'd0; + end + end + end + + // --- S_KG_TMUL_ACCUM: Accumulate SCRATCH0 and SCRATCH1 via poly_arith --- + // Streams SCRATCH0[k] + SCRATCH1[k] → poly_arith (add) → accumulates + // into SCRATCH0 (or back into t_hat register for final result). + // + // Pipeline: rd SCRATCH0 → (1 cycle) → parith with SCRATCH1 → (1 cycle) → write + // Total: ~1 + 256 + 2 = 259 cycles for 256 coefficients + if (state_r == S_KG_TMUL_ACCUM) begin + phase_active <= 1'b1; + phase_done <= 1'b0; + + // coeff_cnt tracks which coeff we're reading from BRAM + // tmul_out_cnt tracks parith results written + if (coeff_cnt == 9'd0 && tmul_out_cnt == 9'd0 && !tmul_pipe_valid) begin + // Start: read SCRATCH0[0] from BRAM + bram_rd_addr_r <= poly_addr(SCRATCH0, 8'd0); + coeff_cnt <= 9'd1; + end else if (coeff_cnt > 9'd0 && coeff_cnt <= 9'd256 && !tmul_pipe_valid) begin + // bram_rd_data has SCRATCH0[coeff_cnt-1] + // Need to read SCRATCH1[coeff_cnt-1] next + // Buffer the SCRATCH0 value, read SCRATCH1 + tmul_pipe_reg <= bram_rd_data; // SCRATCH0 coeff + tmul_pipe_valid <= 1'b1; + bram_rd_addr_r <= poly_addr(SCRATCH1, coeff_cnt[7:0] - 8'd1); + end else if (tmul_pipe_valid && !parith_valid_i_r) begin + // tmul_pipe_reg has SCRATCH0 coeff + // bram_rd_data has SCRATCH1 coeff (from previous read) + parith_a_r <= tmul_pipe_reg; // SCRATCH0 coeff + parith_b_r <= bram_rd_data; // SCRATCH1 coeff + parith_mode_r <= 1'b0; // add + parith_valid_i_r <= 1'b1; + tmul_pipe_valid <= 1'b0; + + // Prefetch next SCRATCH0 coefficient (if any remain) + if (coeff_cnt <= 9'd256) + bram_rd_addr_r <= poly_addr(SCRATCH0, coeff_cnt[7:0]); + coeff_cnt <= coeff_cnt + 9'd1; + end + + // Write parith result back to SCRATCH0 + if (parith_valid_o) begin + bram_wr_en_r <= 1'b1; + bram_wr_addr_r <= poly_addr(SCRATCH0, tmul_out_cnt[7:0]); + bram_wr_data_r <= parith_coeff_out; + tmul_out_cnt <= tmul_out_cnt + 9'd1; + + if (tmul_out_cnt == 9'd255) begin + phase_done <= 1'b1; + phase_active <= 1'b0; + coeff_cnt <= 9'd0; + tmul_out_cnt <= 9'd0; + tmul_pipe_valid <= 1'b0; + end + end + end + + // --- S_KG_TMUL_ADD_E: Add e_hat[i] to accumulated sum → t_hat_i_reg --- + // Streams SCRATCH0 (accumulated A·s sum) + e_hat[i] (from BRAM) + // through poly_arith (add mode), capturing result into t_hat_i_reg. + // + // Pipeline: rd e_hat → (1 cycle) → parith with SCRATCH0 → (1 cycle) → t_hat_i_reg + if (state_r == S_KG_TMUL_ADD_E) begin + phase_active <= 1'b1; + phase_done <= 1'b0; + + // coeff_cnt tracks which e_hat coeff we're reading + // tmul_out_cnt tracks parith results written to t_hat + if (coeff_cnt == 9'd0 && tmul_out_cnt == 9'd0 && !tmul_pipe_valid) begin + // Start: read e_hat[i][0] from BRAM + bram_rd_addr_r <= poly_addr(E_BASE + loop_i, 8'd0); + coeff_cnt <= 9'd1; + end else if (coeff_cnt > 9'd0 && coeff_cnt <= 9'd256 && !tmul_pipe_valid) begin + // bram_rd_data has e_hat[i][coeff_cnt-1] + // Buffer e_hat value, read accumulated sum SCRATCH0 + tmul_pipe_reg <= bram_rd_data; // e_hat coeff + tmul_pipe_valid <= 1'b1; + bram_rd_addr_r <= poly_addr(SCRATCH0, coeff_cnt[7:0] - 8'd1); + end else if (tmul_pipe_valid && !parith_valid_i_r) begin + // tmul_pipe_reg has e_hat coeff + // bram_rd_data has SCRATCH0 (accumulated sum) coeff + parith_a_r <= bram_rd_data; // accumulated sum + parith_b_r <= tmul_pipe_reg; // e_hat coeff + parith_mode_r <= 1'b0; // add + parith_valid_i_r <= 1'b1; + tmul_pipe_valid <= 1'b0; + + // Prefetch next e_hat coefficient + if (coeff_cnt <= 9'd256) + bram_rd_addr_r <= poly_addr(E_BASE + loop_i, coeff_cnt[7:0]); + coeff_cnt <= coeff_cnt + 9'd1; + end + + // Capture parith result into t_hat_i_reg + if (parith_valid_o) begin + if (loop_i == 2'd0) + t_hat0_reg[tmul_out_cnt[7:0] * 12 +: 12] <= parith_coeff_out; + else + t_hat1_reg[tmul_out_cnt[7:0] * 12 +: 12] <= parith_coeff_out; + tmul_out_cnt <= tmul_out_cnt + 9'd1; + + if (tmul_out_cnt == 9'd255) begin + phase_done <= 1'b1; + phase_active <= 1'b0; + coeff_cnt <= 9'd0; + tmul_out_cnt <= 9'd0; + tmul_pipe_valid <= 1'b0; + // Flag that we're advancing to next row (handled in MUL_LD) + if (loop_i + 2'd1 < k_r[1:0]) + tmul_adv_row <= 1'b1; + end + end + end + + // --- S_KG_TMUL_NEXT: Advance loop_j or transition to DONE --- if (state_r == S_KG_TMUL_NEXT) begin phase_done <= 1'b0; - if (loop_i + 2'd1 < k_r[1:0]) - loop_i <= loop_i + 2'd1; + if (loop_j + 2'd1 < k_r[1:0]) begin + loop_j <= loop_j + 2'd1; + end + // else: FSM transitions to S_KG_TMUL_ACCUM (handled by next-state logic) end // ============================================================ - // KEYGEN: DONE + // KEYGEN: DONE — assign pk_o_r and sk_o_r // ============================================================ if (state_r == S_KG_DONE) begin + // For k=2, pk = {t_hat1[3071:0], t_hat0[3071:0], rho[255:0]} (padded) + // t_hat0 goes to bits [12*N-1:0] + // t_hat1 goes to bits [2*12*N-1 : 12*N] + // rho embeds at bits above (for k=2, bits [12*2*N+255 : 12*2*N]) + // For simplicity, assign only the lower parts (k=2): + pk_o_r[(0*N*12) +: N*12] <= t_hat0_reg; + pk_o_r[(1*N*12) +: N*12] <= t_hat1_reg; + // rho_reg is embedded in pk encoding (FIPS 203, byte-encoded) + // Store rho in the upper bits for k=2 + pk_o_r[(2*N*12) +: 256] <= rho_reg; + + sk_o_r[(0*N*12) +: N*12] <= s_hat0_reg; + sk_o_r[(1*N*12) +: N*12] <= s_hat1_reg; + pk_valid_r <= 1'b1; sk_valid_r <= 1'b1; done_o_r <= 1'b1;