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

optimize

Browse Source
This commit is contained in:
FallenSigh
2026-07-11 17:21:21 +08:00
parent 47fbcd31b6
commit d8d57db82a
9 changed files with 504 additions and 316 deletions
Download Patch File Download Diff File Expand all files Collapse all files

231
sync_rtl/top/mlkem_top.v
View File

@@ -22,7 +22,7 @@
`include "sync_rtl/common/defines.vh"
(* use_dsp = "yes" *)
module mlkem_top #(
parameter KMAX = 4 // storage sizing (worst case = ML-KEM-1024)
) (
@@ -90,6 +90,9 @@ module mlkem_top #(
);
localparam Q = `Q; // 3329
localparam [11:0] Q12 = `Q;
localparam [12:0] Q13 = `Q;
localparam [13:0] Q14 = `Q;
// Runtime ML-KEM parameter, captured at start_i.
reg [2:0] k_r;
@@ -309,15 +312,15 @@ module mlkem_top #(
// D1 DECOMP writeback: cd_vo while decoding c2 writes v' to bank_a[DEC_VASLOT].
wire dec_v_we = (st == ST_DEC_DECOMP) && dec_in_c2 && cd_vo;
wire [13:0] dec_v_wr = DEC_VASLOT*256 + dec_widx;
assign bt_we = ((st == ST_M) && pm_vo) || e2_we ||
assign bt_we = ((st == ST_M) && pm_vo) || cbdw_bt_en ||
u_psum_we || u_intt_we || u_v_we || u_w_we;
assign bt_wa = e2_we ? (c_widx & ((1<<PT_AW)-1)) :
assign bt_wa = cbdw_bt_en ? ({6'd0, cbdw_widx} & ((1<<PT_AW)-1)) :
u_psum_we ? ((UPSUM*256 + u_oidx) & ((1<<PT_AW)-1)) :
u_intt_we ? (u_intt_wr[PT_AW-1:0]) :
u_v_we ? (u_vadd_wr[PT_AW-1:0]) :
u_w_we ? (u_wsub_wr[PT_AW-1:0]) :
((m_i*256 + m_oidx) & ((1<<PT_AW)-1));
assign bt_wd = e2_we ? cbd_modq :
assign bt_wd = cbdw_bt_en ? cbdw_data :
u_psum_we ? u_accq :
u_intt_we ? ntt_coeff :
u_v_we ? u_vq :
@@ -378,16 +381,16 @@ module mlkem_top #(
// bank_se rel slot dec_poly (0..K-1), at coeff index dec_widx.
wire dec_u_we = (st == ST_DEC_DECOMP) && !dec_in_c2 && cd_vo;
wire [13:0] dec_u_wr = {2'b0, dec_poly}*256 + dec_widx;
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})) ||
assign bse_we = cbdw_se_en ||
(((st == ST_N) || (st == ST_ENC_N) || (st == ST_DEC_NTT)) && ntt_vo) ||
u_add_we || dec_u_we;
assign bse_wa = u_add_we ? u_add_uwr[PSE_AW-1:0] :
assign bse_wa = cbdw_se_en ? (({1'b0, cbdw_poly}*256 + cbdw_widx) & ((1<<PSE_AW)-1)) :
u_add_we ? u_add_uwr[PSE_AW-1:0] :
dec_u_we ? dec_u_wr[PSE_AW-1:0] :
(st == ST_N || st == ST_ENC_N || st == ST_DEC_NTT) ? ((n_slot*256 + n_widx) & ((1<<PSE_AW)-1))
: ((c_poly*256 + c_widx) & ((1<<PSE_AW)-1));
assign bse_wd = u_add_we ? u_uq :
assign bse_wd = cbdw_se_en ? cbdw_data :
u_add_we ? u_uq :
dec_u_we ? cd_out :
(st == ST_N || st == ST_ENC_N || st == ST_DEC_NTT) ? ntt_coeff : cbd_modq;
always @(posedge clk) begin
@@ -560,9 +563,9 @@ module mlkem_top #(
// per-consumer drives toward the core
wire [1599:0] sha3_kc_state_i, snt_kc_state_i, cbd_kc_state_i;
wire sha3_kc_valid_i, snt_kc_valid_i, cbd_kc_valid_i;
/* verilator lint_off UNUSEDSIGNAL */
/* verilator lint_off UNUSED */
wire sha3_kc_ready_i, snt_kc_ready_i, cbd_kc_ready_i; // all 1'b1
/* verilator lint_on UNUSEDSIGNAL */
/* verilator lint_on UNUSED */
// phase selects (mutually exclusive). Encaps adds ST_ENC_H/G (sha3),
// ST_ENC_A (snt), ST_ENC_C (cbd).
@@ -792,7 +795,35 @@ module mlkem_top #(
);
// signed (two's complement) -> [0,Q): add Q when negative
wire [11:0] cbd_modq = cbd_coeff[11] ? (cbd_coeff + 12'(Q)) : cbd_coeff;
wire [11:0] cbd_modq = cbd_coeff[11] ? (cbd_coeff + 12'd3329) : cbd_coeff;
// Pipeline the CBD -> bank write by one cycle. The squeeze_buf[buf_ptr]
// extraction -> CBD add/sub -> mod-Q -> BRAM-write route is long and
// routing-dominated (the mlkem_clk critical path at 100 MHz). Register the
// write data/addr/enable here; the bank write commits one cycle later. No
// consumer reads these coeffs until the following stage, which begins more
// than one cycle after the last CBD coeff, so there is no read-after-write
// hazard. The registered write takes priority in the bank_se/bank_t mux so
// the 1-cycle spillover into the next state uses the correct addr/data.
reg cbdw_se_en, cbdw_bt_en;
reg [4:0] cbdw_poly;
reg [7:0] cbdw_widx;
reg [11:0] cbdw_data;
always @(posedge clk or negedge rst_n) begin
if (!rst_n) begin
cbdw_se_en <= 1'b0; cbdw_bt_en <= 1'b0;
cbdw_poly <= 5'd0; cbdw_widx <= 8'd0; cbdw_data <= 12'd0;
end else begin
cbdw_se_en <= ((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}));
cbdw_bt_en <= (st == ST_ENC_C) && c_busy && cbd_vo && cbd_ack &&
(c_poly == {1'b0, k_r, 1'b0});
cbdw_poly <= c_poly;
cbdw_widx <= c_widx;
cbdw_data <= cbd_modq;
end
end
// ---- ntt_core: forward NTT (mode=0, no scaling) of s[i],e[i] in place ----
// N-stage bookkeeping: process slots S0,S1,E0,E1 (= SLOT_S0 + n_slot).
@@ -917,7 +948,7 @@ module mlkem_top #(
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];
wire [11:0] u_accq = (u_sum >= Q13) ? (u_sum - Q13) : 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;
@@ -926,7 +957,7 @@ module mlkem_top #(
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];
wire [11:0] u_uq = (u_usum >= Q13) ? (u_usum - Q13) : u_usum[11:0];
// ================================================================
// E6: v = INTT(sum_j t_hat[j] o y_hat[j]) + e2 + mu (Encaps ST_ENC_V).
@@ -947,8 +978,8 @@ module mlkem_top #(
wire [13:0] u_vadd_wr = UPSUM*256 + u_awidx; // bank_t v (write)
wire [11:0] u_mu = m_r[u_awidx] ? 12'd1665 : 12'd0; // mu by msg bit
wire [13:0] u_vsum = {2'b0, bt_rd_data} + {2'b0, ba_rd_data} + {2'b0, u_mu}; // psum+e2+mu
wire [13:0] u_vsub1 = (u_vsum >= 14'(Q)) ? (u_vsum - 14'(Q)) : u_vsum;
wire [11:0] u_vq = (u_vsub1 >= 14'(Q)) ? (u_vsub1 - 14'(Q)) : u_vsub1[11:0];
wire [13:0] u_vsub1 = (u_vsum >= Q14) ? (u_vsum - Q14) : u_vsum;
wire [11:0] u_vq = (u_vsub1 >= Q14) ? (u_vsub1 - Q14) : u_vsub1[11:0];
wire u_v_we = (st == ST_ENC_V) && (u_sub == 2'd2) && u_avalid;
// ================================================================
@@ -964,7 +995,7 @@ module mlkem_top #(
wire [13:0] u_w_vrd = DEC_VASLOT*256 + u_aidx[7:0]; // bank_a v' (lead)
wire [13:0] u_wsub_wr = UPSUM*256 + u_awidx; // bank_t w (write)
wire [12:0] u_wdiff = {1'b0, ba_rd_data} - {1'b0, bt_rd_data}; // v' - psum
wire [11:0] u_wq = u_wdiff[12] ? (u_wdiff[11:0] + 12'(Q)) : u_wdiff[11:0];
wire [11:0] u_wq = u_wdiff[12] ? (u_wdiff[11:0] + Q12) : u_wdiff[11:0];
wire u_w_we = (st == ST_DEC_W) && (u_sub == 2'd2) && u_avalid;
// ================================================================
@@ -1038,13 +1069,18 @@ module mlkem_top #(
.ready_i(1'b1)
);
// bit-packer / ct walk bookkeeping
reg [2:0] cp_poly; // c1: 0..K-1 (u rows); c2: single v
reg [7:0] cp_idx; // coeff 0..255 within poly
reg [2:0] cp_ph; // micro-phase
reg [24:0] cp_buf; // bit accumulator (LSB-first)
reg [5:0] cp_nbits; // valid bits in cp_buf
reg [2:0] cp_poly; // c1: 0..K-1 (u rows); c2: single v (feed side)
reg [7:0] cp_idx; // coeff 0..255 within poly (feed side, drives bank addr)
reg [2:0] cp_ph; // feed micro-phase toggle (0=addr settle, 1=feed)
reg [63:0] cp_buf; // bit accumulator (LSB-first); widened for streaming
reg [6:0] cp_nbits; // valid bits in cp_buf
reg [11:0] cp_wa; // ct_bram byte write address (runs c1 then c2)
reg cp_done; // serialization complete (this region)
reg cp_fed_done; // all coeffs for this region have been fed in
reg [10:0] cp_fcnt; // coeffs fed into comp_decomp this region
reg [10:0] cp_ccnt; // compressed coeffs collected this region
// total coeffs this region: c1 = K*256 (u rows), c2 = 256 (single v)
wire [10:0] cp_ntot = (st == ST_ENC_C2) ? 11'd256 : ({8'd0, k_r} << 8);
// coeff source: c1 (ST_ENC_C1) reads u[cp_poly] in bank_se rel (K+cp_poly);
// c2 (ST_ENC_C2) reads v in bank_t rel slot UPSUM (single poly, cp_poly=0).
wire [13:0] cp_se_full = ({2'b0,k_r}+{2'b0,cp_poly})*256 + cp_idx; // bank_se u[cp_poly]
@@ -1128,7 +1164,7 @@ module mlkem_top #(
wire [11:0] m_acc_src = m_jq ? bse_rd_data : bt_rd_data;
// (a + b) mod Q (both < Q, sum < 2Q): one conditional subtract
wire [12:0] m_sum = {1'b0, m_acc_src} + {1'b0, pm_coeff};
wire [11:0] m_accq = (m_sum >= 13'(Q)) ? (m_sum - 13'(Q)) : m_sum[11:0];
wire [11:0] m_accq = (m_sum >= Q13) ? (m_sum - Q13) : m_sum[11:0];
always @(*) begin
st_next = st;
@@ -1257,10 +1293,13 @@ module mlkem_top #(
cp_poly <= 3'd0;
cp_idx <= 8'd0;
cp_ph <= 3'd0;
cp_buf <= 25'd0;
cp_nbits <= 6'd0;
cp_buf <= 64'd0;
cp_nbits <= 7'd0;
cp_wa <= 12'd0;
cp_done <= 1'b0;
cp_fed_done<= 1'b0;
cp_fcnt <= 11'd0;
cp_ccnt <= 11'd0;
ct_we <= 1'b0;
ct_wa <= 11'd0;
ct_wd <= 8'd0;
@@ -1692,69 +1731,86 @@ module mlkem_top #(
// Arm E5 (ST_ENC_C1) when U finishes: c1 = byteEncode_du(Compress_du(u)).
// Walk K polys * 256 coeffs; reset bit-packer + ct write pointer.
if (st == ST_ENC_U && st_next == ST_ENC_C1) begin
cp_poly <= 3'd0;
cp_idx <= 8'd0;
cp_ph <= 3'd0;
cp_buf <= 25'd0;
cp_nbits <= 6'd0;
cp_wa <= 12'd0;
cp_done <= 1'b0;
cd_valid <= 1'b0;
cp_poly <= 3'd0;
cp_idx <= 8'd0;
cp_ph <= 3'd0;
cp_buf <= 64'd0;
cp_nbits <= 7'd0;
cp_wa <= 12'd0;
cp_done <= 1'b0;
cp_fed_done <= 1'b0;
cp_fcnt <= 11'd0;
cp_ccnt <= 11'd0;
cd_valid <= 1'b0;
end
// ---- ST_ENC_C1/C2: Compress_d -> byteEncode_d -> ct region ----
// C1 (E5): Compress_du(u[0..K-1]) from bank_se -> ct[0..c1_bytes).
// C2 (E7): Compress_dv(v) from bank_t[UPSUM] -> ct[c1_bytes..ct_bytes).
// Per coeff, 5-phase micro-sequence (read-ahead 1 cyc bram plus
// variable-latency comp_decomp), then a drain sub-phase emitting whole bytes:
// ph0: present coeff addr (cp_se_full / cp_bt_full by state).
// ph1: coeff arrives (cp_coeff_src) -> latch into cd_coeff, pulse cd_valid.
// ph2: drop cd_valid (1-cyc pulse); comp_decomp captures.
// ph3: cd_vo high -> cd_out (low cp_d bits) valid; append LSB-first.
// ph4: drain: while >=8 bits buffered, emit one ct byte/cycle; then advance.
// Each poly = 256 coeffs = 32*d bytes (whole), so cp_buf empties at
// each poly boundary (no carry across polys / regions).
//
// Streaming (B1): the comp_decomp pipeline is fully pipelined and its
// sole output consumer never backpressures (ready_i=1'b1), so instead
// of feeding one coeff and draining the whole pipe per coeff (~12 cyc),
// we run a decoupled producer + collector:
// PRODUCER (2-cycle cadence): ph0 lets the bank read address settle
// (1-cyc sd_bram latency), ph1 latches the coeff and pulses
// cd_valid, then advances the address. One coeff enters the pipe
// every 2 cycles = <=5.5 bits/cyc (du<=11), strictly under the
// single-port ct drain rate (8 bits/cyc), so cp_buf stays bounded.
// COLLECTOR (every cycle): when cd_vo, append the compressed cp_d
// bits LSB-first; independently emit one ct byte/cycle whenever
// >=8 bits are buffered. Coeffs come out in-order, one per pipe
// output. Each poly is 256*d bits (multiple of 8), so the buffer
// empties to 0 at region end (cp_done when all collected+drained).
if ((st == ST_ENC_C1 || st == ST_ENC_C2) && !cp_done) begin
case (cp_ph)
3'd0: cp_ph <= 3'd1; // addr presented; wait read
3'd1: begin
cd_coeff <= cp_coeff_src; // u (bank_se) or v (bank_t)
cd_valid <= 1'b1; // feed comp_decomp (1-cyc pulse)
cp_ph <= 3'd2;
end
3'd2: begin
cd_valid <= 1'b0; // comp_decomp captured this cyc
cp_ph <= 3'd3;
end
3'd3: begin
if (cd_vo) begin
// Append compressed low cp_d bits LSB-first at bit cp_nbits.
cp_buf <= cp_buf | (({13'd0, cd_out} & ((25'd1 << cp_d) - 25'd1)) << cp_nbits);
cp_nbits <= cp_nbits + {1'b0, cp_d};
cp_ph <= 3'd4;
end
end
default: begin // 3'd4: drain whole bytes
if (cp_nbits >= 6'd8) begin
ct_we <= 1'b1;
ct_wa <= cp_wa;
ct_wd <= cp_buf[7:0];
cp_wa <= cp_wa + 12'd1;
cp_buf <= cp_buf >> 8;
cp_nbits <= cp_nbits - 6'd8;
// ---- PRODUCER: feed one coeff every 2 cycles ----
cd_valid <= 1'b0;
if (!cp_fed_done) begin
if (cp_ph == 3'd0) begin
cp_ph <= 3'd1; // address settling
end else begin
cd_coeff <= cp_coeff_src; // u (bank_se) or v (bank_t)
cd_valid <= 1'b1; // 1-cyc feed pulse
cp_fcnt <= cp_fcnt + 11'd1;
if (cp_fcnt + 11'd1 >= cp_ntot) cp_fed_done <= 1'b1;
if (cp_idx == 8'd255) begin
cp_idx <= 8'd0;
cp_poly <= cp_poly + 3'd1; // next u row (c1); c2 unused
end else begin
// coeff fully packed; advance coeff / poly
if (cp_idx == 8'd255) begin
cp_idx <= 8'd0;
if (cp_poly + 3'd1 < cp_poly_max) cp_poly <= cp_poly + 3'd1;
else cp_done <= 1'b1; // region complete
end else begin
cp_idx <= cp_idx + 8'd1;
end
cp_ph <= 3'd0;
cp_idx <= cp_idx + 8'd1;
end
cp_ph <= 3'd0;
end
endcase
end
// ---- COLLECTOR: append compressed bits + drain ct bytes ----
if (cd_vo) begin
// append cp_d bits at cp_nbits, then drain one byte if >=8.
if ((cp_nbits + {1'b0, cp_d}) >= 7'd8) begin
ct_we <= 1'b1;
ct_wa <= cp_wa;
ct_wd <= (cp_buf | (({52'd0, cd_out} & ((64'd1 << cp_d) - 64'd1)) << cp_nbits)) & 64'hFF;
cp_wa <= cp_wa + 12'd1;
cp_buf <= (cp_buf | (({52'd0, cd_out} & ((64'd1 << cp_d) - 64'd1)) << cp_nbits)) >> 8;
cp_nbits <= (cp_nbits + {1'b0, cp_d}) - 7'd8;
end else begin
cp_buf <= cp_buf | (({52'd0, cd_out} & ((64'd1 << cp_d) - 64'd1)) << cp_nbits);
cp_nbits <= cp_nbits + {1'b0, cp_d};
end
cp_ccnt <= cp_ccnt + 11'd1; // one compressed coeff collected
end else if (cp_nbits >= 7'd8) begin
// no new coeff this cycle: keep draining buffered whole bytes.
ct_we <= 1'b1;
ct_wa <= cp_wa;
ct_wd <= cp_buf[7:0];
cp_wa <= cp_wa + 12'd1;
cp_buf <= cp_buf >> 8;
cp_nbits <= cp_nbits - 7'd8;
end
// region complete once all coeffs collected and buffer drained.
if (cp_ccnt >= cp_ntot && cp_nbits < 7'd8)
cp_done <= 1'b1;
end
// Arm E7 (ST_ENC_C2) when V finishes: c2 = byteEncode_dv(Compress_dv(v)).
@@ -1762,13 +1818,16 @@ module mlkem_top #(
// from c1_bytes (NOT reset) so c2 lands right after c1. cp_buf/nbits
// are 0 here (c1 ended on a poly/byte boundary) but reset for safety.
if (st == ST_ENC_V && st_next == ST_ENC_C2) begin
cp_poly <= 3'd0;
cp_idx <= 8'd0;
cp_ph <= 3'd0;
cp_buf <= 25'd0;
cp_nbits <= 6'd0;
cp_done <= 1'b0;
cd_valid <= 1'b0;
cp_poly <= 3'd0;
cp_idx <= 8'd0;
cp_ph <= 3'd0;
cp_buf <= 64'd0;
cp_nbits <= 7'd0;
cp_done <= 1'b0;
cp_fed_done <= 1'b0;
cp_fcnt <= 11'd0;
cp_ccnt <= 11'd0;
cd_valid <= 1'b0;
// cp_wa intentionally preserved (= c1_bytes_rt from C1).
end