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

fix(kg): implement t_hat computation and pk/sk output in mlkem_top

Browse Source 
- 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
This commit is contained in:
FallenSigh
2026-06-27 01:38:38 +08:00
parent 880e87daad
commit 3284aa443f
1 changed files with 303 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

346
sync_rtl/top/mlkem_top.v
View File

@@ -119,6 +119,18 @@ module mlkem_top #(
reg [255:0] Hpk_reg; // H(pk) for Encaps reg [255:0] Hpk_reg; // H(pk) for Encaps
reg [255:0] z_reg; // implicit rejection value 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) // 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_NTT_OUT = 7'd28;
localparam S_KG_CBD_E_NEXT = 7'd29; localparam S_KG_CBD_E_NEXT = 7'd29;
localparam S_KG_TMUL_INIT = 7'd30; localparam S_KG_TMUL_INIT = 7'd30;
localparam S_KG_TMUL_STEP = 7'd31; localparam S_KG_TMUL_MUL_LD = 7'd31;
localparam S_KG_TMUL_NEXT = 7'd32; localparam S_KG_TMUL_MUL_OUT = 7'd32;
localparam S_KG_DONE = 7'd33; 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 --- // --- Encaps states ---
localparam S_EN_RNG_REQ = 7'd34; localparam S_EN_RNG_REQ = 7'd37;
localparam S_EN_RNG_WAIT = 7'd35; localparam S_EN_RNG_WAIT = 7'd38;
localparam S_EN_H_START = 7'd36; localparam S_EN_H_START = 7'd39;
localparam S_EN_H_WAIT = 7'd37; localparam S_EN_H_WAIT = 7'd40;
localparam S_EN_G_START = 7'd38; localparam S_EN_G_START = 7'd41;
localparam S_EN_G_WAIT = 7'd39; localparam S_EN_G_WAIT = 7'd42;
localparam S_EN_SNT_INIT = 7'd40; localparam S_EN_SNT_INIT = 7'd43;
localparam S_EN_SNT_START = 7'd41; localparam S_EN_SNT_START = 7'd44;
localparam S_EN_SNT_COEFFS = 7'd42; localparam S_EN_SNT_COEFFS = 7'd45;
localparam S_EN_SNT_CLEANUP = 7'd43; localparam S_EN_SNT_CLEANUP = 7'd46;
localparam S_EN_SNT_NTT_LD = 7'd44; localparam S_EN_SNT_NTT_LD = 7'd47;
localparam S_EN_SNT_NTT_CMP = 7'd45; localparam S_EN_SNT_NTT_CMP = 7'd48;
localparam S_EN_SNT_NTT_OUT = 7'd46; localparam S_EN_SNT_NTT_OUT = 7'd49;
localparam S_EN_SNT_NEXT = 7'd47; localparam S_EN_SNT_NEXT = 7'd50;
localparam S_EN_CBD_Y_INIT = 7'd48; localparam S_EN_CBD_Y_INIT = 7'd51;
localparam S_EN_CBD_Y_START = 7'd49; localparam S_EN_CBD_Y_START = 7'd52;
localparam S_EN_CBD_Y_COEFFS = 7'd50; localparam S_EN_CBD_Y_COEFFS = 7'd53;
localparam S_EN_CBD_Y_CLNUP = 7'd51; localparam S_EN_CBD_Y_CLNUP = 7'd54;
localparam S_EN_CBD_Y_NTT_LD = 7'd52; localparam S_EN_CBD_Y_NTT_LD = 7'd55;
localparam S_EN_CBD_Y_NTT_CMP = 7'd53; localparam S_EN_CBD_Y_NTT_CMP = 7'd56;
localparam S_EN_CBD_Y_NTT_OUT = 7'd54; localparam S_EN_CBD_Y_NTT_OUT = 7'd57;
localparam S_EN_CBD_Y_NEXT = 7'd55; localparam S_EN_CBD_Y_NEXT = 7'd58;
localparam S_EN_DONE = 7'd56; localparam S_EN_DONE = 7'd59;
// --- Decaps states --- // --- Decaps states ---
localparam S_DC_DECOMP_C1 = 7'd57; localparam S_DC_DECOMP_C1 = 7'd60;
localparam S_DC_DECOMP_C2 = 7'd58; localparam S_DC_DECOMP_C2 = 7'd61;
localparam S_DC_NTT_U_LD = 7'd59; localparam S_DC_NTT_U_LD = 7'd62;
localparam S_DC_NTT_U_CMP = 7'd60; localparam S_DC_NTT_U_CMP = 7'd63;
localparam S_DC_NTT_U_OUT = 7'd61; localparam S_DC_NTT_U_OUT = 7'd64;
localparam S_DC_MUL_S = 7'd62; localparam S_DC_MUL_S = 7'd65;
localparam S_DC_INTT_V = 7'd63; localparam S_DC_INTT_V = 7'd66;
localparam S_DC_DECOMP_M = 7'd64; localparam S_DC_DECOMP_M = 7'd67;
localparam S_DC_G_CHECK = 7'd65; localparam S_DC_G_CHECK = 7'd68;
localparam S_DC_REENC = 7'd66; localparam S_DC_REENC = 7'd69;
localparam S_DC_KDF = 7'd67; localparam S_DC_KDF = 7'd70;
localparam S_DC_DONE = 7'd68; localparam S_DC_DONE = 7'd71;
// ==================================================================== // ====================================================================
// FSM registers // FSM registers
@@ -485,15 +500,29 @@ module mlkem_top #(
end end
// -- t_hat = sum(A·s) + e -- // -- t_hat = sum(A·s) + e --
S_KG_TMUL_INIT: state_next = S_KG_TMUL_STEP; S_KG_TMUL_INIT: state_next = S_KG_TMUL_MUL_LD;
S_KG_TMUL_STEP: if (phase_done) state_next = S_KG_TMUL_NEXT; 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;
S_KG_TMUL_NEXT: begin // 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]) if (loop_i + 2'd1 < k_r[1:0])
state_next = S_KG_TMUL_STEP; state_next = S_KG_TMUL_MUL_LD;
else else
state_next = S_KG_DONE; state_next = S_KG_DONE;
end end
end
// After MUL_OUT: check if more j terms for this row
S_KG_TMUL_NEXT: begin
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; S_KG_DONE: state_next = S_IDLE;
@@ -596,6 +625,14 @@ module mlkem_top #(
poly_idx_reg <= 6'd0; poly_idx_reg <= 6'd0;
phase_active <= 1'b0; phase_active <= 1'b0;
phase_done <= 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 // Drive all FSM-controlled outputs to safe defaults
cbd_valid_i_r <= 1'b0; cbd_valid_i_r <= 1'b0;
@@ -828,6 +865,11 @@ module mlkem_top #(
bram_wr_en_r <= 1'b1; bram_wr_en_r <= 1'b1;
bram_wr_addr_r <= poly_addr(S_BASE + loop_idx, coeff_cnt[7:0]); bram_wr_addr_r <= poly_addr(S_BASE + loop_idx, coeff_cnt[7:0]);
bram_wr_data_r <= ntt_coeff_out; 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; coeff_cnt <= coeff_cnt + 9'd1;
if (coeff_cnt == 9'd255) begin if (coeff_cnt == 9'd255) begin
phase_done <= 1'b1; phase_done <= 1'b1;
@@ -921,30 +963,236 @@ module mlkem_top #(
end 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 if (state_r == S_KG_TMUL_INIT) begin
loop_i <= 2'd0; loop_i <= 2'd0;
loop_j <= 2'd0;
coeff_cnt <= 9'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; phase_done <= 1'b0;
end end
if (state_r == S_KG_TMUL_STEP) begin // --- S_KG_TMUL_MUL_LD: Load A[i,j] and s_hat[j] into poly_mul ---
// TODO: Per-row poly_mul A[i,:]·s[:] + e[i] → t_hat[i] // Uses s_hat registers (combinational read) + A from BRAM (1-cycle delay).
// For now, placeholder: just mark done // 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_done <= 1'b1;
phase_active <= 1'b0;
coeff_cnt <= 9'd0;
end
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 if (state_r == S_KG_TMUL_NEXT) begin
phase_done <= 1'b0; phase_done <= 1'b0;
if (loop_i + 2'd1 < k_r[1:0]) if (loop_j + 2'd1 < k_r[1:0]) begin
loop_i <= loop_i + 2'd1; loop_j <= loop_j + 2'd1;
end
// else: FSM transitions to S_KG_TMUL_ACCUM (handled by next-state logic)
end end
// ============================================================ // ============================================================
// KEYGEN: DONE // KEYGEN: DONE — assign pk_o_r and sk_o_r
// ============================================================ // ============================================================
if (state_r == S_KG_DONE) begin 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; pk_valid_r <= 1'b1;
sk_valid_r <= 1'b1; sk_valid_r <= 1'b1;
done_o_r <= 1'b1; done_o_r <= 1'b1;