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

feat(dec): Decaps D2 - s_hat=byteDecode12(dk_pke) + u_hat=NTT(u')

Browse Source 
K-PKE.Decrypt steps 2-3 (FIPS 203 Alg 15), both by reusing Encaps machines:

- ST_DEC_SDEC reuses the Encaps TDEC (byteDecode12) machine: only the byte
  source changes (td_byte mux -> dkp_rd_data; dkp_rd_addr driven by td_ekaddr
  in SDEC). Decodes dk_pke -> s_hat[j] into bank_a slot j*K, the same layout
  t_hat uses, so the D3 MAC can read s_hat[j] with the existing addressing.
- ST_DEC_NTT reuses the forward-NTT machine (n_slot_max=k_r) to transform
  u'[i] in place in bank_se rel slots 0..K-1 -> u_hat[i]. Added ST_DEC_NTT to
  the bank_se read/write muxes and the NTT load/process/arm blocks alongside
  ST_N/ST_ENC_N.
- FSM: DECOMP -> SDEC -> NTT -> DONE.

TB verify_d2 checks s_hat[i] (bank_a slot i*K) and u_hat[i] (bank_se rel i)
against golden. verify_d1 narrowed to v' only: D2's in-place NTT overwrites u'
in bank_se, so u' correctness is now proven transitively via u_hat==NTT(u').

Verified: dec D2 K=2/3/4 all cases PASS; KeyGen + Encaps unregressed.
This commit is contained in:
FallenSigh
2026-06-29 18:00:44 +08:00
parent ecc00d6dd5
commit 940946f30c
3 changed files with 92 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

59
sync_rtl/top/mlkem_top.v
View File

@@ -158,6 +158,10 @@ module mlkem_top #(
// decoded coeffs: c0 = b0 | ((b1&0xF)<<8); c1 = (b1>>4) | (b2<<4)
wire [11:0] td_c0 = {td_b1[3:0], td_b0};
wire [11:0] td_c1 = {td_b2, td_b1[7:4]};
// TDEC byte source: Encaps reads ek (ek_rd_data); Decaps D2 (ST_DEC_SDEC)
// reuses the same machine to byteDecode12 dk_pke -> s_hat, reading dkp_bram.
wire td_dec_s = (st == ST_DEC_SDEC);
wire [7:0] td_byte = td_dec_s ? dkp_rd_data : ek_rd_data;
// bank_a write for TDEC (registered); t_hat[td_poly] -> bank_a slot
// td_poly*K. td_wa is a bank_a byte address (PA_AW=12).
reg td_we;
@@ -336,7 +340,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 || st == ST_ENC_N) ? ntt_rd_full[PSE_AW-1:0] :
assign bse_rd_addr = (st == ST_N || st == ST_ENC_N || st == ST_DEC_NTT) ? 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] :
@@ -358,15 +362,15 @@ module mlkem_top #(
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) || (st == ST_ENC_N)) && ntt_vo) ||
(((st == ST_N) || (st == ST_ENC_N) || (st == ST_DEC_NTT)) && ntt_vo) ||
u_add_we || dec_u_we;
assign bse_wa = u_add_we ? u_add_uwr[PSE_AW-1:0] :
dec_u_we ? dec_u_wr[PSE_AW-1:0] :
(st == ST_N || st == ST_ENC_N) ? ((n_slot*256 + n_widx) & ((1<<PSE_AW)-1))
(st == ST_N || st == ST_ENC_N || st == ST_DEC_NTT) ? ((n_slot*256 + n_widx) & ((1<<PSE_AW)-1))
: ((c_poly*256 + c_widx) & ((1<<PSE_AW)-1));
assign bse_wd = u_add_we ? u_uq :
dec_u_we ? cd_out :
(st == ST_N || st == ST_ENC_N) ? ntt_coeff : cbd_modq;
(st == ST_N || st == ST_ENC_N || st == ST_DEC_NTT) ? 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)
@@ -465,6 +469,8 @@ module mlkem_top #(
// ---- Decaps states ----
localparam ST_DEC_LOAD = 5'd20; // dk/c already streamed in; parse/settle (D0)
localparam ST_DEC_DECOMP = 5'd21; // D1: byteDecode_d + Decompress c1->u', c2->v'
localparam ST_DEC_SDEC = 5'd22; // D2: byteDecode12 dk_pke -> s_hat (bank_a slot j*K)
localparam ST_DEC_NTT = 5'd23; // D2: u_hat[i] = NTT(u'[i]) in place (bank_se rel i)
localparam ST_DONE = 5'd31;
reg [4:0] st, st_next;
@@ -648,8 +654,10 @@ module mlkem_top #(
(st == ST_ENC_TDEC) ? td_ekaddr[10:0] : // byteDecode12
dbgdk_in_ek ? dbgdk_ek_off[10:0]
: dbg_byte_idx_i;
// dkp BRAM read-address mux: dbg_byte (sel=1) or dbg_dk (dkp region).
assign dkp_rd_addr = dbg_byte_sel_i ? dbg_byte_idx_i : dbg_dk_idx_i[10:0];
// dkp BRAM read-address mux: D2 byteDecode12 (ST_DEC_SDEC) walks dk_pke via
// td_ekaddr (poly*384 + trip*3 + ph); else dbg_byte (sel=1) / dbg_dk.
assign dkp_rd_addr = (st == ST_DEC_SDEC) ? td_ekaddr[10:0] :
dbg_byte_sel_i ? dbg_byte_idx_i : dbg_dk_idx_i[10:0];
// ---- sample_ntt_sync: Â[i][j] = SampleNTT(rho || j || i) ----
reg snt_valid;
@@ -729,7 +737,7 @@ module mlkem_top #(
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 [4:0] n_slot_max = (st == ST_ENC_N || st == ST_DEC_NTT) ? {2'b0, k_r} : {1'b0, k_r, 1'b0};
wire ntt_ready;
wire [11:0] ntt_coeff;
@@ -1017,7 +1025,11 @@ module mlkem_top #(
// D0: settle after dk/c parse, then D1 Decompress. D1 lands in DONE
// once u'/v' are computed so the stage can be dbg-checked.
ST_DEC_LOAD: st_next = ST_DEC_DECOMP;
ST_DEC_DECOMP: if (dec_dc_done) st_next = ST_DONE;
ST_DEC_DECOMP: if (dec_dc_done) st_next = ST_DEC_SDEC;
// D2: s_hat decode (reuse TDEC machine), then u_hat = NTT(u'). D2
// lands in DONE once u_hat is ready so the stage can be dbg-checked.
ST_DEC_SDEC: if (td_done) st_next = ST_DEC_NTT;
ST_DEC_NTT: if (n_slot >= {2'b0, k_r}) st_next = ST_DONE;
ST_G: if (sha3_vo) st_next = ST_A;
ST_A: if (a_pair >= kk_rt) st_next = ST_C;
ST_C: if (c_poly >= {1'b0, k_r, 1'b0}) st_next = ST_N;
@@ -1300,7 +1312,8 @@ module mlkem_top #(
// 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) ||
(st == ST_ENC_C && st_next == ST_ENC_N)) begin
(st == ST_ENC_C && st_next == ST_ENC_N) ||
(st == ST_DEC_SDEC && st_next == ST_DEC_NTT)) begin
n_slot <= 3'd0;
n_ridx <= 9'd0;
n_widx <= 8'd0;
@@ -1309,9 +1322,9 @@ module mlkem_top #(
n_pending <= 1'b0;
end
// ---- 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
// ---- ST_N / ST_ENC_N / ST_DEC_NTT: forward NTT in place. KeyGen: 2K
// slots (s,e). Encaps: K slots (y). Decaps D2: K slots (u' -> u_hat). ----
if (st == ST_N || st == ST_ENC_N || st == ST_DEC_NTT) begin
// slot-count bound: 2K for KeyGen, K for Encaps
// (n_slot_max below); same LOAD/OUTPUT cadence either way.
if (n_loading) begin
@@ -1858,16 +1871,22 @@ module mlkem_top #(
td_ph <= 3'd0;
td_done <= 1'b0;
end
// Arm D2 s_hat decode (ST_DEC_SDEC) when DECOMP finishes: same TDEC
// machine, byteDecode12 dk_pke -> s_hat in bank_a slot j*K.
if (st == ST_DEC_DECOMP && st_next == ST_DEC_SDEC) begin
td_poly <= 3'd0;
td_trip <= 8'd0;
td_ph <= 3'd0;
td_done <= 1'b0;
end
// ---- ST_ENC_TDEC: byteDecode12 ek -> t_hat[0..k-1] in bank_t ----
// 5-cycle micro-phase per triple (read-ahead, 1-cyc bram latency):
// ph0: present b0 addr; ph1: capture b0, present b1; ph2: capture
// b1, present b2; ph3: capture b2 + write c0; ph4: write c1, advance.
if (st == ST_ENC_TDEC && !td_done) begin
// Encaps: ek -> t_hat. Decaps D2: dk_pke -> s_hat. Same 5-cycle
// micro-phase per triple; only the byte source (td_byte) differs.
if ((st == ST_ENC_TDEC || st == ST_DEC_SDEC) && !td_done) begin
// capture the byte that arrived for the address presented last cycle
if (td_ph == 3'd1) td_b0 <= ek_rd_data;
if (td_ph == 3'd2) td_b1 <= ek_rd_data;
if (td_ph == 3'd3) td_b2 <= ek_rd_data;
if (td_ph == 3'd1) td_b0 <= td_byte;
if (td_ph == 3'd2) td_b1 <= td_byte;
if (td_ph == 3'd3) td_b2 <= td_byte;
// write decoded coeffs into bank_a at slot td_poly*K (so V's MAC
// reads t_hat[j] via the same u_aslot=u_j*K addressing, u_row=0).