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

Pipeline ML-KEM datapath bottlenecks

Browse Source
This commit is contained in:
FallenSigh
2026-07-08 00:23:46 +08:00
parent 8c3f4317f5
commit 372a90e601
17 changed files with 776 additions and 176 deletions
Download Patch File Download Diff File Expand all files Collapse all files

53
sync_rtl/top/mlkem_top.v
View File

@@ -22,7 +22,7 @@
`include "sync_rtl/common/defines.vh"
(* use_dsp = "no" *)
(* use_dsp = "yes" *)
module mlkem_top #(
parameter KMAX = 4 // storage sizing (worst case = ML-KEM-1024)
) (
@@ -194,6 +194,7 @@ module mlkem_top #(
wire [10:0] cin_rd_addr;
wire [7:0] cin_rd_data;
reg [10:0] cin_rd_addr_r; // D5/D7 read address (tied 0 until then)
reg dj_active; // background Decaps J(z||c) owns c_in read port
sd_bram #(.W(8), .D(2048), .A(11)) u_c_in_bram (
.clk(clk),
.rd_addr(cin_rd_addr), .rd_data(cin_rd_data),
@@ -203,7 +204,7 @@ module mlkem_top #(
// stage) so the registered BRAM read yields the byte 1 cycle later, matching
// the assemble/writeback pipeline. Otherwise the registered cin_rd_addr_r
// (D1 walker / idle) drives it.
assign cin_rd_addr = (st == ST_DEC_J) ? dj_c_idx[10:0] :
assign cin_rd_addr = dj_active ? dj_c_idx[10:0] :
(st == ST_DEC_CMP) ? cmp_idx : cin_rd_addr_r;
// ================================================================
@@ -567,7 +568,7 @@ module mlkem_top #(
// ST_ENC_A (snt), ST_ENC_C (cbd).
wire sel_sha3 = (st == ST_G) || (st == ST_H) ||
(st == ST_ENC_H) || (st == ST_ENC_G) ||
(st == ST_DEC_G) || (st == ST_DEC_J);
(st == ST_DEC_G) || dj_active;
wire sel_snt = (st == ST_A) || (st == ST_ENC_A);
wire sel_cbd = (st == ST_C) || (st == ST_ENC_C);
@@ -597,14 +598,14 @@ module mlkem_top #(
// G/H/J share one sha3_top. Single-block (mb_en=0): KeyGen G (ST_G, mode 00),
// Encaps G (ST_ENC_G, mode 11), Decaps D5 G (ST_DEC_G, mode 11). Multi-block
// absorb (mb_en=1): H(ek) (ST_H/ST_ENC_H, SHA3-256 pad) and Decaps D5 J
// (ST_DEC_J, SHAKE256 pad over z||c). mb_* inputs are muxed by phase.
wire sha3_mb_en = (st == ST_H) || (st == ST_ENC_H) || (st == ST_DEC_J);
// (background dj_active, SHAKE256 pad over z||c). mb_* inputs are muxed by phase.
wire sha3_mb_en = (st == ST_H) || (st == ST_ENC_H) || dj_active;
wire [1:0] sha3_mode = (st == ST_ENC_G || st == ST_DEC_G) ? 2'b11 : 2'b00;
// multi-block feed mux: J (ST_DEC_J) drives dj_*, else H drives h_*.
wire [1087:0] mb_block_mux = (st == ST_DEC_J) ? dj_block_r : h_block_r;
wire mb_valid_mux = (st == ST_DEC_J) ? dj_mbvalid : h_mbvalid;
wire mb_last_mux = (st == ST_DEC_J) ? dj_mblast : h_mblast;
wire mb_ack_mux = (st == ST_DEC_J) ? dj_ack : h_ack;
// multi-block feed mux: background J drives dj_*, else H drives h_*.
wire [1087:0] mb_block_mux = dj_active ? dj_block_r : h_block_r;
wire mb_valid_mux = dj_active ? dj_mbvalid : h_mbvalid;
wire mb_last_mux = dj_active ? dj_mblast : h_mblast;
wire mb_ack_mux = dj_active ? dj_ack : h_ack;
sha3_top_shared u_sha3 (
.clk(clk), .rst_n(rst_n),
.mode(sha3_mode), // G = SHA3-512 (only used when mb_en=0)
@@ -1145,12 +1146,12 @@ module mlkem_top #(
// D3: w = v' - INTT(sum_j s_hat[j] o u_hat[j]). Single output poly
// (u_row 0..0); MAC->INTT->SUB, then D4 encodes m'.
ST_DEC_W: if (u_row >= 3'd1) st_next = ST_DEC_MENC;
ST_DEC_MENC: if (men_done) st_next = ST_DEC_G;
// D5: (K',r') = G(m'||h) single-block, then K-bar = J(z||c) multi-block.
ST_DEC_G: if (sha3_vo) st_next = ST_DEC_J;
// D5 J done -> D6 re-encrypt: c' = K-PKE.Encrypt(ek_pke, m', r').
// Reuse the entire Encaps pipeline (rho load -> A -> C -> ... -> C2).
// r' is in r_r (CBD seed), m' in m_r (V/mu), ek_pke in ek_bram.
ST_DEC_MENC: if (men_done && dj_done) st_next = ST_DEC_G;
// D5: (K',r') = G(m'||h) single-block. K-bar = J(z||c) runs in a
// background SHAKE256 walker after D1, so the main FSM only waits
// for it at the D4->D5 boundary if it somehow has not finished.
ST_DEC_G: if (sha3_vo) st_next = ST_ENC_LOAD;
// Legacy state kept unused; background J normally skips it.
ST_DEC_J: if (dj_done) st_next = ST_ENC_LOAD;
ST_G: if (sha3_vo) st_next = ST_A;
ST_A: if (a_pair >= kk_rt) st_next = ST_C;
@@ -1303,6 +1304,7 @@ module mlkem_top #(
cmp_neq <= 1'b0;
cmp_done <= 1'b0;
dec_reject <= 1'b0;
dj_active <= 1'b0;
dj_blk <= 4'd0;
dj_byte <= 8'd0;
dj_phase <= 2'd0;
@@ -1349,6 +1351,8 @@ module mlkem_top #(
end else if (op_i == 2'd2) begin
// Decaps: dk_pke/ek_pke already in BRAM; z/H(ek) captured into
// z_r/hek_r during load (below). Nothing else to arm in D0.
dj_active <= 1'b0;
dj_done <= 1'b0;
end else begin
sha3_valid <= 1'b1;
sha3_ack <= 1'b1;
@@ -1993,9 +1997,11 @@ module mlkem_top #(
sha3_ack <= 1'b0;
end
// Arm D5 J(z||c) when G completes (ST_DEC_G -> ST_DEC_J): assemble the
// first 136-byte block. dj_ack high to consume the final digest.
if (st == ST_DEC_G && st_next == ST_DEC_J) begin
// Arm D5 J(z||c) as soon as D1 has stopped reading c_in_bram. It
// then runs in the background during D2/D3/D4 using the otherwise
// idle SHAKE256 path and c_in read port.
if (st == ST_DEC_DECOMP && st_next == ST_DEC_SDEC) begin
dj_active <= 1'b1;
dj_blk <= 4'd0;
dj_byte <= 8'd0;
dj_phase <= 2'd0;
@@ -2007,14 +2013,14 @@ module mlkem_top #(
cin_rd_addr_r <= 11'd0; // present c byte 0 (g=32 -> c_idx 0)
end
// ---- ST_DEC_J: K-bar = J(z||c) multi-block SHAKE256 ----
// ---- Background Decaps J: K-bar = J(z||c) multi-block SHAKE256 ----
// Mirror of the H(ek) multi-block machine. Byte source by global g:
// g < 32 -> z_r byte g
// 32 <= g < msglen -> c_in_bram byte (g-32), registered read
// else -> SHAKE256 pad constant
// c_in_bram read is registered: present c_idx for dj_byte this cycle,
// write back the byte that arrived for the addr presented last cycle.
if (st == ST_DEC_J && !dj_done) begin
if (dj_active && !dj_done) begin
case (dj_phase)
2'd0: begin
// writeback the byte read for the previous address
@@ -2034,7 +2040,7 @@ module mlkem_top #(
dj_wb_zidx <= dj_g[4:0];
dj_wb_pad <= dj_padconst(dj_g);
// c_in addr is presented combinationally (cin_rd_addr
// mux uses dj_c_idx during ST_DEC_J); data lands next
// mux uses dj_c_idx while dj_active); data lands next
// cycle, matching this byte's writeback.
dj_byte <= dj_byte + 8'd1;
end else begin
@@ -2056,6 +2062,7 @@ module mlkem_top #(
if (sha3_vo) begin
kbar_r <= sha3_hash[255:0];
dj_done <= 1'b1;
dj_active <= 1'b0;
dj_phase<= 2'd3;
end else if (h_mbready) begin
dj_blk <= dj_blk + 4'd1;
@@ -2095,7 +2102,7 @@ module mlkem_top #(
end
if (st_next == ST_ENC_LOAD &&
(st == ST_ENC_G || st == ST_DEC_J)) begin
(st == ST_ENC_G || st == ST_DEC_G || st == ST_DEC_J)) begin
rl_idx <= 6'd0;
rl_widx <= 6'd0;
rl_vld <= 1'b0;