perf(mlkem): ping-pong poly_mul + decoupled MAC load/accumulate (B2)

Double-buffer poly_mul_sync so the next MAC term's 256-cycle load overlaps
the current term's accumulate. Decouple the top-level MAC load index (u_ldj)
from the accumulate index (u_j) in the shared ENC_U/ENC_V/DEC_W engine; load
reads bank_a/bank_se while accumulate reads bank_t (disjoint ports, no
conflict). Accumulate stays strictly serial (psum order-dependent); only the
load is hidden.

Verified byte-exact vs NIST KAT in ~/Dev/mlkem: Encaps 9/9, Decaps 9/9
(incl. implicit-reject + DECOMP), KeyGen all, hello_world end-to-end.

ENC_U 5585->5073, ENC_V 2793->2537 (K=2). Larger relative win at K=4
(K^2 terms). poly_mul contract unchanged (standalone TB 11/11).
This commit is contained in:
2026-07-11 02:37:28 +08:00
parent ce40a6df42
commit 975903000e
2 changed files with 173 additions and 162 deletions

View File

@@ -896,10 +896,12 @@ module mlkem_top #(
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 [2:0] u_j; // ACCUMULATE term col j 0..K-1 (psum order)
reg [2:0] u_ldj; // LOAD term col j 0..K-1 (runs ahead of u_j)
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_ldwait; // load-side waiting for a free poly_mul buffer
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)
@@ -912,10 +914,15 @@ module mlkem_top #(
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;
// MAC addresses (B2 decoupled): the LOAD side runs ahead on u_ldj (term being
// streamed into poly_mul's idle buffer); the ACCUMULATE side stays on u_j
// (term whose products are being summed into psum). A^T -> A_hat[u_ldj][u_row]
// slot = u_ldj*k+u_row; y_hat[u_ldj] rel u_ldj. Load reads bank_a/bank_se;
// accumulate reads bank_t (psum) -> no read-port conflict, so the next term's
// load overlaps the current term's accumulate.
wire [SAW-1:0] u_aslot = u_ldj*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 [13:0] u_pm_b_full = u_ldj*256 + u_ld[7:0]; // bank_se y_hat[u_ldj]
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
@@ -1243,6 +1250,8 @@ module mlkem_top #(
u_row <= 3'd0;
u_sub <= 2'd0;
u_j <= 3'd0;
u_ldj <= 3'd0;
u_ldwait <= 1'b0;
u_ld <= 9'd0;
u_oidx <= 8'd0;
u_loading <= 1'b0;
@@ -1591,6 +1600,8 @@ module mlkem_top #(
u_row <= 3'd0;
u_sub <= 2'd0; // MAC
u_j <= 3'd0;
u_ldj <= 3'd0;
u_ldwait <= 1'b0;
u_ld <= 9'd0;
u_oidx <= 8'd0;
u_loading <= 1'b1;
@@ -1599,30 +1610,53 @@ module mlkem_top #(
end
// ---- ST_ENC_U/V sub-phase 0: MAC sum_j (A^T or t_hat)[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).
// U: A^T[u_row][j] (bank_a u_j*K+u_row). V: t_hat[j] (bank_a u_j*K, u_row=0).
// psum (init 0 at j==0) accumulates into bank_t[UPSUM].
// B2: decoupled load/accumulate. poly_mul is double-buffered, so the
// LOAD of term u_ldj can run AHEAD, overlapping the ACCUMULATE of an
// earlier term u_j (psum is order-dependent, so accumulate stays
// strictly serial; only the 256-cyc load is hidden). Load reads
// bank_a/bank_se; accumulate reads bank_t (UPSUM) — disjoint ports.
// u_ldj = term being streamed into poly_mul (0..K-1, then "done");
// u_j = term whose products are being accumulated (via u_j0q sel).
if ((st == ST_ENC_U || st == ST_ENC_V || st == ST_DEC_W) && u_sub == 2'd0) begin
u_j0q <= (u_j == 3'd0);
// ---- LOAD side: stream term u_ldj into poly_mul's free buffer ----
// Within a term poly_mul holds ready high for all 256 pairs, so the
// fixed 1-cyc addr->valid skew is kept. Between terms we wait for a
// free buffer (pm_ready) before priming the next (u_ldwait).
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;
if (u_ldj + 3'd1 < k_r) begin
u_ldj <= u_ldj + 3'd1; // more terms to load
u_ldwait <= 1'b1; // wait for a free buffer
end
end else begin
u_ld <= u_ld + 9'd1;
pm_valid <= 1'b1;
end
end else if (u_ldwait) begin
// Next term's load: start as soon as poly_mul has a free buffer.
// pm_ready high => buf[ld_ptr] free; begin presenting addr now,
// assert pm_valid next cycle (same skew as term start).
if (pm_ready) begin
u_loading <= 1'b1;
u_ld <= 9'd0;
u_ldwait <= 1'b0;
pm_valid <= 1'b0;
end
end
// ---- ACCUMULATE side: consume poly_mul outputs in load order ----
if (pm_vo) begin
// bank_t psum write is combinational (u_psum_we); advance idx
if (u_oidx == 8'd255) begin
u_oidx <= 8'd0;
if (u_j + 3'd1 < k_r) begin
u_j <= u_j + 3'd1;
u_pending <= 1'b1; // next term, same row
u_j <= u_j + 3'd1; // next term's products
end else begin
// row's MAC done -> INTT sub-phase
u_sub <= 2'd1;
@@ -1635,15 +1669,6 @@ module mlkem_top #(
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/V sub-phase 1: INTT(psum) mode=1 in place (bank_t) ----
@@ -1689,6 +1714,8 @@ module mlkem_top #(
u_row <= u_row + 3'd1;
u_sub <= 2'd0; // next row MAC
u_j <= 3'd0;
u_ldj <= 3'd0;
u_ldwait <= 1'b0;
u_ld <= 9'd0;
u_oidx <= 8'd0;
u_loading <= 1'b1;
@@ -2244,6 +2271,8 @@ module mlkem_top #(
u_row <= 3'd0;
u_sub <= 2'd0; // MAC
u_j <= 3'd0;
u_ldj <= 3'd0;
u_ldwait <= 1'b0;
u_ld <= 9'd0;
u_oidx <= 8'd0;
u_loading <= 1'b1;
@@ -2257,9 +2286,11 @@ module mlkem_top #(
u_row <= 3'd0;
u_sub <= 2'd0; // MAC
u_j <= 3'd0;
u_ldj <= 3'd0;
u_ld <= 9'd0;
u_oidx <= 8'd0;
u_loading <= 1'b1;
u_ldwait <= 1'b0;
u_pending <= 1'b0;
pm_valid <= 1'b0;
end