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

feat(enc): Encaps E3 - y_hat = NTT(y) in place

Browse Source 
ST_ENC_N reuses the ST_N forward-NTT datapath (mode=0, read-ahead load,
in-place writeback to bank_se). Slot count parameterized: KeyGen 2K (s,e),
Encaps K (y only; e1/e2 stay time-domain) via n_slot_max. bse rd/wr muxes
and the C->N arming extended for ST_ENC_N.

Verified (K=2 c0) y_hat[0..1] == ml-kem-r golden; E0/E1/E2 still pass
(E2 y check dropped since NTT overwrites y in place -> covered by E3);
K=3/4 run through E3 without timeout.
This commit is contained in:
FallenSigh
2026-06-29 02:01:37 +08:00
parent cdc5ce25b1
commit 8ed4d59546
2 changed files with 44 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
sync_rtl/top/mlkem_top.v
View File

@@ -241,7 +241,7 @@ module mlkem_top #(
dbg_t_addr[PT_AW-1:0];
// bank_se read port: ST_N load (s/e NTT), ST_M load (pm_b s_hat[j]) vs
// accumulate (e_hat, mutually exclusive via m_loading), ST_E dk-half, dbg.
assign bse_rd_addr = (st == ST_N) ? ntt_rd_full[PSE_AW-1:0] :
assign bse_rd_addr = (st == ST_N || st == ST_ENC_N) ? ntt_rd_full[PSE_AW-1:0] :
(st == ST_M) ? (m_loading ? pm_b_full[PSE_AW-1:0]
: m_eacc_full[PSE_AW-1:0]) :
(st == ST_E) ? e_rd_full[PSE_AW-1:0] :
@@ -249,14 +249,14 @@ module mlkem_top #(
// bank_se write port: ST_C CBD store (rel slot c_poly), ST_N NTT writeback
// (rel slot n_slot). Disjoint states.
// bank_se write: KeyGen ST_C CBD (s/e), Encaps ST_ENC_C CBD (y/e1, c_poly<2K;
// e2 at c_poly==2K goes to bank_t instead), ST_N NTT writeback. Disjoint.
// e2 at c_poly==2K goes to bank_t instead), ST_N/ST_ENC_N NTT writeback.
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})) ||
((st == ST_N) && ntt_vo);
assign bse_wa = (st == ST_N) ? ((n_slot*256 + n_widx) & ((1<<PSE_AW)-1))
(((st == ST_N) || (st == ST_ENC_N)) && ntt_vo);
assign bse_wa = (st == ST_N || st == ST_ENC_N) ? ((n_slot*256 + n_widx) & ((1<<PSE_AW)-1))
: ((c_poly*256 + c_widx) & ((1<<PSE_AW)-1));
assign bse_wd = (st == ST_N) ? ntt_coeff : cbd_modq;
assign bse_wd = (st == ST_N || st == ST_ENC_N) ? ntt_coeff : cbd_modq;
always @(posedge clk) begin
if (dbg_slot_i >= slot_t_rt) dbg_coeff_r <= bt_rd_data; // bank_t (sd_bram)
else if (dbg_slot_i >= slot_s_rt) dbg_coeff_r <= bse_rd_data; // bank_se (sd_bram)
@@ -596,6 +596,8 @@ module mlkem_top #(
reg n_loading; // 1 while presenting load addresses to bank_se
reg n_pending; // waiting for ntt_core IDLE to start next slot
wire [SAW-1:0] n_slot_addr = slot_s_rt + n_slot; // s_hat then e_hat contiguous
// NTT slot count: KeyGen processes 2K (s,e); Encaps processes K (y only).
wire [4:0] n_slot_max = (st == ST_ENC_N) ? {2'b0, k_r} : {1'b0, k_r, 1'b0};
wire ntt_ready;
wire [11:0] ntt_coeff;
@@ -726,7 +728,8 @@ module mlkem_top #(
ST_ENC_G: if (sha3_vo) st_next = ST_ENC_LOAD; // (K,r) captured
ST_ENC_LOAD: if (rl_idx >= 6'd33) st_next = ST_ENC_A; // rho loaded
ST_ENC_A: if (a_pair >= kk_rt) st_next = ST_ENC_C;
ST_ENC_C: if (c_poly >= {k_r, 1'b1}) st_next = ST_DONE; // E2: stop (2K+1 polys: y,e1,e2)
ST_ENC_C: if (c_poly >= {k_r, 1'b1}) st_next = ST_ENC_N; // 2K+1 polys done
ST_ENC_N: if (n_slot >= {2'b0, k_r}) st_next = ST_DONE; // E3: K slots (y_hat)
ST_ENC_TDEC: if (td_done) st_next = ST_DONE; // (TDEC deferred to V-prep later)
ST_DONE: st_next = ST_IDLE;
default: st_next = ST_IDLE;
@@ -926,7 +929,8 @@ module mlkem_top #(
// Arm N stage when C finishes: prime load of slot S0. n_ridx is a
// read-ahead pointer; bank_se read is registered inside sd_bram (bse_rd_data)
// and fed to ntt_core one cycle later, so valid starts low (priming).
if (st == ST_C && st_next == ST_N) begin
if ((st == ST_C && st_next == ST_N) ||
(st == ST_ENC_C && st_next == ST_ENC_N)) begin
n_slot <= 3'd0;
n_ridx <= 9'd0;
n_widx <= 8'd0;
@@ -935,12 +939,11 @@ module mlkem_top #(
n_pending <= 1'b0;
end
// ---- ST_N: forward NTT each of S0,S1,E0,E1 in place ----
if (st == ST_N) begin
// LOAD phase: present read-ahead addr to bank_se (bse_rd_addr);
// sd_bram registers it, so bse_rd_data is consumed by ntt_core
// one cycle later (n_valid). Cores hold ready high through LOAD,
// so a fixed 1-cycle skew suffices (no backpressure gating).
// ---- ST_N / ST_ENC_N: forward NTT in place. KeyGen: 2K slots
// (s,e). Encaps: K slots (y only; e1/e2 stay time-domain). ----
if (st == ST_N || st == ST_ENC_N) begin
// slot-count bound: 2K for KeyGen, K for Encaps
// (n_slot_max below); same LOAD/OUTPUT cadence either way.
if (n_loading) begin
if (n_ridx == 9'd256) begin
// 256th coeff consumed this cycle; stop presenting addr
@@ -953,20 +956,18 @@ module mlkem_top #(
end
// OUTPUT phase: collect 256 results, write back to same slot.
// The bank_se write is the combinational bse_we/bse_wa/bse_wd
// assigns (rel slot n_slot); here only advance the write index.
if (ntt_vo) begin
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 < {1'b0, k_r, 1'b0}) begin
if (n_slot + 3'd1 < n_slot_max) 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
n_slot <= n_slot + 3'd1; // == n_slot_max -> exit
end
end