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

feat(mlkem_top): KeyGen stage 2e (matrix accumulate t_hat)

Browse Source 
Add ST_M stage: t_hat[i] = e_hat[i] + sum_j A_hat[i][j] o s_hat[j] via
poly_mul_sync + inline mod-add accumulation. Per (i,j): stream 256 (A,shat)
pairs into poly_mul, then accumulate 256 products into T_i (seeded from E_i
when j==0, else running T_i). m_pending waits for poly_mul IDLE between terms.

Verified vs ml-kem-r golden: 512/512 t_hat coeffs exact (19885 cyc).
This commit is contained in:
FallenSigh
2026-06-28 01:53:23 +08:00
parent 4c692e570a
commit a9e50ebc0c
3 changed files with 672 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

104
sync_rtl/top/mlkem_top.v
View File

@@ -73,6 +73,7 @@ module mlkem_top #(
localparam ST_A = 4'd2; // generate A_hat[i][j] via SampleNTT
localparam ST_C = 4'd3; // generate s[i],e[i] via CBD
localparam ST_N = 4'd4; // forward NTT of s[i],e[i] in place
localparam ST_M = 4'd5; // matrix accumulate t_hat = e_hat + sum A o s_hat
localparam ST_DONE = 4'd15;
reg [3:0] st, st_next;
@@ -195,6 +196,46 @@ module mlkem_top #(
.done_o(ntt_done)
);
// ---- poly_mul_sync: t_hat[i] = e_hat[i] + sum_j A_hat[i][j] o s_hat[j] ----
// M-stage bookkeeping. For each (i,j): LOAD 256 (A,shat) pairs, then accumulate
// 256 products into T_i (init from E_i when j==0, else from running T_i).
reg [1:0] m_i; // row 0..K
reg [1:0] m_j; // col 0..K
reg [8:0] m_ld; // load index 0..256
reg [7:0] m_oidx; // output/accum index 0..255
reg m_loading; // 1 while streaming pairs into poly_mul
reg m_pending; // wait for poly_mul IDLE before next (i,j)
wire [3:0] m_aslot = {2'b0, m_i[0], m_j[0]}; // A_hat[i][j] slot = i*2+j (0..3)
wire [3:0] m_sslot = SLOT_S0 + {3'b0, m_j[0]}; // s_hat[j]
wire [3:0] m_eslot = SLOT_E0 + {3'b0, m_i[0]}; // e_hat[i]
wire [3:0] m_tslot = SLOT_T0 + {3'b0, m_i[0]}; // t_hat[i]
reg pm_valid;
wire pm_ready;
wire [11:0] pm_coeff;
wire pm_vo;
wire [11:0] pm_a_in = polymem[m_aslot*256 + m_ld[7:0]];
wire [11:0] pm_b_in = polymem[m_sslot*256 + m_ld[7:0]];
poly_mul_sync u_pmul (
.clk(clk), .rst_n(rst_n),
.coeff_a_in(pm_a_in),
.coeff_b_in(pm_b_in),
.valid_i(pm_valid),
.ready_o(pm_ready),
.coeff_out(pm_coeff),
.valid_o(pm_vo),
.ready_i(1'b1)
);
// accumulator source: e_hat[i] for first term (j==0), else running t_hat[i]
wire [11:0] m_acc_src = (m_j == 2'd0) ? polymem[m_eslot*256 + m_oidx]
: polymem[m_tslot*256 + m_oidx];
// (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];
always @(*) begin
st_next = st;
case (st)
@@ -202,7 +243,8 @@ module mlkem_top #(
ST_G: if (sha3_vo) st_next = ST_A;
ST_A: if (a_pair >= K*K) st_next = ST_C;
ST_C: if (c_poly >= 2*K) st_next = ST_N;
ST_N: if (n_slot >= 2*K) st_next = ST_DONE;
ST_N: if (n_slot >= 2*K) st_next = ST_M;
ST_M: if (m_i >= K) st_next = ST_DONE;
ST_DONE: st_next = ST_IDLE;
default: st_next = ST_IDLE;
endcase
@@ -230,6 +272,13 @@ module mlkem_top #(
n_widx <= 8'd0;
n_valid <= 1'b0;
n_pending <= 1'b0;
m_i <= 2'd0;
m_j <= 2'd0;
m_ld <= 9'd0;
m_oidx <= 8'd0;
m_loading <= 1'b0;
m_pending <= 1'b0;
pm_valid <= 1'b0;
end else begin
st <= st_next;
@@ -351,6 +400,59 @@ module mlkem_top #(
n_pending <= 1'b0;
end
end
// Arm M stage when N finishes: start first (i=0,j=0) poly_mul load
if (st == ST_N && st_next == ST_M) begin
m_i <= 2'd0;
m_j <= 2'd0;
m_ld <= 9'd0;
m_oidx <= 8'd0;
m_loading <= 1'b1;
m_pending <= 1'b0;
pm_valid <= 1'b1;
end
// ---- ST_M: t_hat[i] = e_hat[i] + sum_j A[i][j] o s_hat[j] ----
if (st == ST_M) begin
// LOAD: stream 256 (A,shat) pairs into poly_mul
if (m_loading && pm_valid && pm_ready) begin
if (m_ld == 9'd255) begin
pm_valid <= 1'b0; // last pair presented
m_loading <= 1'b0;
m_ld <= 9'd0;
m_oidx <= 8'd0;
end else begin
m_ld <= m_ld + 9'd1;
end
end
// ACCUMULATE: each product coeff += e_hat (j==0) or running t_hat
if (pm_vo) begin
polymem[m_tslot*256 + m_oidx] <= m_accq;
if (m_oidx == 8'd255) begin
// finished this (i,j) term; advance
if (m_j + 2'd1 < K) begin
m_j <= m_j + 2'd1;
m_pending <= 1'b1; // next term, same row
end else begin
m_j <= 2'd0;
m_i <= m_i + 2'd1; // next row (or == K -> DONE)
if (m_i + 2'd1 < K) m_pending <= 1'b1;
end
end else begin
m_oidx <= m_oidx + 8'd1;
end
end
// Start next (i,j) poly_mul load once core is IDLE again
if (m_pending && pm_ready && !pm_vo) begin
pm_valid <= 1'b1;
m_loading <= 1'b1;
m_ld <= 9'd0;
m_oidx <= 8'd0;
m_pending <= 1'b0;
end
end
end
end