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

feat(mlkem_top): KeyGen stage 2d (forward NTT of s/e)

Browse Source 
Add ST_N stage: forward NTT (ntt_core mode=0, no scaling) of s[0],s[1],
e[0],e[1] in place (slots S0,S1,E0,E1). Per slot: stream 256 coeffs into
ntt_core during LOAD, collect 256 outputs back to same slot. n_pending
waits for core IDLE between slots.

Verified vs ml-kem-r golden: 1024/1024 shat/ehat coeffs exact (17318 cyc).
This commit is contained in:
FallenSigh
2026-06-28 01:47:54 +08:00
parent 2f206a6bc5
commit 4c692e570a
3 changed files with 1162 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

82
sync_rtl/top/mlkem_top.v
View File

@@ -72,6 +72,7 @@ module mlkem_top #(
localparam ST_G = 4'd1; // run G(d||K), capture rho/sigma
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_DONE = 4'd15;
reg [3:0] st, st_next;
@@ -167,13 +168,41 @@ 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;
// ---- 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).
reg [2:0] n_slot; // 0..2K (4 polys)
reg [8:0] n_ridx; // load read index 0..256
reg [7:0] n_widx; // output write index 0..255
reg n_valid; // feeding coeffs to ntt_core
reg n_pending; // waiting for ntt_core IDLE to start next slot
wire [3:0] n_slot_addr = SLOT_S0 + {1'b0, n_slot};
wire ntt_ready;
wire [11:0] ntt_coeff;
wire ntt_vo;
wire ntt_done;
wire [11:0] ntt_in = polymem[n_slot_addr*256 + n_ridx[7:0]];
ntt_core u_ntt (
.clk(clk), .rst_n(rst_n),
.coeff_in(ntt_in),
.valid_i(n_valid),
.ready_o(ntt_ready),
.mode(1'b0), // forward NTT, no scaling
.coeff_out(ntt_coeff),
.valid_o(ntt_vo),
.ready_i(1'b1), // always accept output
.done_o(ntt_done)
);
always @(*) begin
st_next = st;
case (st)
ST_IDLE: if (start_i) st_next = ST_G;
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_DONE;
ST_C: if (c_poly >= 2*K) st_next = ST_N;
ST_N: if (n_slot >= 2*K) st_next = ST_DONE;
ST_DONE: st_next = ST_IDLE;
default: st_next = ST_IDLE;
endcase
@@ -196,6 +225,11 @@ module mlkem_top #(
c_poly <= 3'd0;
c_widx <= 8'd0;
c_busy <= 1'b0;
n_slot <= 3'd0;
n_ridx <= 9'd0;
n_widx <= 8'd0;
n_valid <= 1'b0;
n_pending <= 1'b0;
end else begin
st <= st_next;
@@ -271,6 +305,52 @@ module mlkem_top #(
end
end
end
// Arm N stage when C finishes: start NTT on slot S0
if (st == ST_C && st_next == ST_N) begin
n_slot <= 3'd0;
n_ridx <= 9'd0;
n_widx <= 8'd0;
n_valid <= 1'b1; // begin loading first poly
n_pending <= 1'b0;
end
// ---- ST_N: forward NTT each of S0,S1,E0,E1 in place ----
if (st == ST_N) begin
// LOAD phase: stream 256 coeffs into ntt_core
if (n_valid && ntt_ready) begin
if (n_ridx == 9'd255) begin
n_valid <= 1'b0; // last coeff presented this cycle
n_ridx <= 9'd0;
end else begin
n_ridx <= n_ridx + 9'd1;
end
end
// OUTPUT phase: collect 256 results, write back to same slot
if (ntt_vo) begin
polymem[n_slot_addr*256 + n_widx] <= ntt_coeff;
n_widx <= n_widx + 8'd1; // wraps 255->0 after last
end
// Slot complete when ntt_core returns to DONE
if (ntt_done) begin
if (n_slot + 3'd1 < 2*K) begin
n_slot <= n_slot + 3'd1;
n_widx <= 8'd0;
n_pending <= 1'b1; // wait one cycle for core IDLE
end else begin
n_slot <= n_slot + 3'd1; // == 2K -> ST_DONE
end
end
// Kick next slot's load once core is back IDLE
if (n_pending && ntt_ready && !ntt_done) begin
n_valid <= 1'b1;
n_ridx <= 9'd0;
n_pending <= 1'b0;
end
end
end
end