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

feat(dec): Decaps D3+D4 - w = v'-INTT(s.u_hat) + m' recovery

Browse Source 
Completes K-PKE.Decrypt (FIPS 203 Alg 15) in hardware: m' is recovered.

D3 (ST_DEC_W) reuses the Encaps V MAC/INTT machine (u_row tied to 0):
- MAC s_hat[j] (bank_a slot j*K) o u_hat[j] (bank_se rel j) -> psum bank_t[UPSUM]
  -- identical addressing to Encaps V (t_hat[j] o y_hat[j]), so free reuse.
- INTT(psum) in place.
- SUB: w = v' - psum mod Q (negative -> +Q), written to bank_t[UPSUM].
  To read v' and psum in parallel during SUB (one read port per bank), D1's v'
  write was relocated from bank_t to bank_a slot DEC_VASLOT=1 (always free:
  s_hat occupies j*K, slot 1 is unused for K>=2). This mirrors V-ADD reading
  psum (bank_t) + e2 (bank_a) simultaneously.

D4 (ST_DEC_MENC): m' = byteEncode_1(Compress_1(w)). Compress_1(w)=1 iff
832 < w <= 2496 (Q=3329); bits packed LSB-first into mprime_r, exposed on
dbg_mprime_o (was a placeholder tied to m_r).

Added ST_DEC_W to the u_* machine muxes/sub-phases and the FSM chain
NTT->W->MENC->DONE. TB verify_d3 checks w (bank_t UPSUM); verify_d4 checks the
32-byte m' against golden (== the KAT-decrypted m == original message).

Verified: dec D1-D4 K=2/3/4 all cases PASS; KeyGen + Encaps unregressed.
This commit is contained in:
FallenSigh
2026-06-29 18:57:29 +08:00
parent 940946f30c
commit 7f519fe826
3 changed files with 166 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

137
sync_rtl/top/mlkem_top.v
View File

@@ -245,13 +245,17 @@ module mlkem_top #(
// Encaps E6 also writes t_hat here via byteDecode12 (TDEC): t_hat[j] lands
// in bank_a at slot j*K so V's MAC reuses E4's u_aslot=u_j*K (u_row=0) addr.
// ST_ENC_E2MV writes the relocated e2 into bank_a[E2_ASLOT] (em_we).
// D1 DECOMP writes v' into bank_a[DEC_VASLOT] (dec_v_we) so D3 can read it
// in parallel with the bank_t psum during the w-SUB.
assign ba_we = ((st == ST_A || st == ST_ENC_A) && a_busy && snt_vo && snt_ack)
|| td_we || em_we;
|| td_we || em_we || dec_v_we;
assign ba_wa = td_we ? td_wa[PA_AW-1:0]
: em_we ? em_wa[PA_AW-1:0]
: dec_v_we ? dec_v_wr[PA_AW-1:0]
: (a_slot*256 + a_widx) & ((1<<PA_AW)-1);
assign ba_wd = td_we ? td_wd :
em_we ? bt_rd_data : // e2 read from bank_t rel slot 0
dec_v_we ? cd_out : // v' decompressed coeff
snt_coeff;
// bank_se: s_hat || e_hat. Readers = ST_N load, ST_M load (pm_b, s_hat[j]),
@@ -290,24 +294,24 @@ module mlkem_top #(
// Encaps E4/E6 bank_t writes: MAC psum (u_sub0, on pm_vo) and INTT in-place
// (u_sub1, on ntt_vo). Both target bank_t rel slot UPSUM. ST_ENC_V (E6) adds
// the V-ADD writeback (psum + e2 + mu) at UPSUM via u_v_we below.
wire u_psum_we = (st == ST_ENC_U || st == ST_ENC_V) && (u_sub == 2'd0) && pm_vo;
wire u_intt_we = (st == ST_ENC_U || st == ST_ENC_V) && (u_sub == 2'd1) && ntt_vo;
// D1 DECOMP writeback: cd_vo while decoding c2 writes v' to bank_t[DEC_VSLOT].
wire u_psum_we = (st == ST_ENC_U || st == ST_ENC_V || st == ST_DEC_W) && (u_sub == 2'd0) && pm_vo;
wire u_intt_we = (st == ST_ENC_U || st == ST_ENC_V || st == ST_DEC_W) && (u_sub == 2'd1) && ntt_vo;
// D1 DECOMP writeback: cd_vo while decoding c2 writes v' to bank_a[DEC_VASLOT].
wire dec_v_we = (st == ST_DEC_DECOMP) && dec_in_c2 && cd_vo;
wire [13:0] dec_v_wr = DEC_VSLOT*256 + dec_widx;
wire [13:0] dec_v_wr = DEC_VASLOT*256 + dec_widx;
assign bt_we = ((st == ST_M) && pm_vo) || e2_we ||
u_psum_we || u_intt_we || u_v_we || dec_v_we;
u_psum_we || u_intt_we || u_v_we || u_w_we;
assign bt_wa = e2_we ? (c_widx & ((1<<PT_AW)-1)) :
u_psum_we ? ((UPSUM*256 + u_oidx) & ((1<<PT_AW)-1)) :
u_intt_we ? (u_intt_wr[PT_AW-1:0]) :
u_v_we ? (u_vadd_wr[PT_AW-1:0]) :
dec_v_we ? (dec_v_wr[PT_AW-1:0]) :
u_w_we ? (u_wsub_wr[PT_AW-1:0]) :
((m_i*256 + m_oidx) & ((1<<PT_AW)-1));
assign bt_wd = e2_we ? cbd_modq :
u_psum_we ? u_accq :
u_intt_we ? ntt_coeff :
u_v_we ? u_vq :
dec_v_we ? cd_out :
u_w_we ? u_wq :
m_accq;
// Debug readback (registered for timing)
@@ -325,18 +329,22 @@ module mlkem_top #(
(st == ST_ENC_V) ?
((u_sub == 2'd2) ? u_v_e2rd[PA_AW-1:0] // ADD: e2 read
: u_pm_a_full[PA_AW-1:0]) : // MAC: t_hat[u_j]
(st == ST_DEC_W) ?
((u_sub == 2'd2) ? u_w_vrd[PA_AW-1:0] // SUB: v' read
: u_pm_a_full[PA_AW-1:0]) : // MAC: s_hat[u_j]
dbg_a_addr[PA_AW-1:0];
// bank_t read port: ST_M acc (t_hat), ST_E ek-half (byteEncode), Encaps
// ST_ENC_U/V (MAC psum acc, INTT load, or ADD psum read by sub-phase),
// ST_ENC_E2MV (e2 read from rel slot 0); else dbg.
// ST_ENC_U/V + Decaps ST_DEC_W (MAC psum acc, INTT load, or ADD/SUB psum
// read by sub-phase), ST_ENC_E2MV (e2 read from rel slot 0); else dbg.
assign bt_rd_addr = (st == ST_M) ? m_tacc_full[PT_AW-1:0] :
(st == ST_E) ? e_rd_full[PT_AW-1:0] :
(st == ST_ENC_U || st == ST_ENC_V) ?
(st == ST_ENC_U || st == ST_ENC_V || st == ST_DEC_W) ?
((u_sub == 2'd0) ? u_psum_full[PT_AW-1:0] : // MAC acc
(u_sub == 2'd1) ? u_intt_rd[PT_AW-1:0] : // INTT load
u_add_prd[PT_AW-1:0]) : // ADD psum read
u_add_prd[PT_AW-1:0]) : // ADD/SUB psum read
(st == ST_ENC_E2MV) ? em_rd[PT_AW-1:0] : // e2 relocate read
(st == ST_ENC_C2) ? cp_bt_full[PT_AW-1:0] : // C2: v[cp_idx]
(st == ST_DEC_MENC) ? men_rd[PT_AW-1:0] : // D4: w[men_idx]
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.
@@ -348,6 +356,7 @@ module mlkem_top #(
((u_sub == 2'd0) ? u_pm_b_full[PSE_AW-1:0] : // MAC: y_hat[u_j]
u_add_e1rd[PSE_AW-1:0]) : // ADD: e1[u_row] read
(st == ST_ENC_V) ? u_pm_b_full[PSE_AW-1:0] : // V MAC: y_hat[u_j]
(st == ST_DEC_W && u_sub == 2'd0) ? u_pm_b_full[PSE_AW-1:0] : // D3 MAC: u_hat[u_j]
(st == ST_ENC_C1) ? cp_se_full[PSE_AW-1:0] : // C1: u[cp_poly]
dbg_se_addr[PSE_AW-1:0];
// bank_se write: KeyGen ST_C CBD (s/e), Encaps ST_ENC_C CBD (y/e1, c_poly<2K;
@@ -471,6 +480,8 @@ module mlkem_top #(
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_DEC_W = 5'd24; // D3: w = v' - INTT(sum_j s_hat[j] o u_hat[j])
localparam ST_DEC_MENC = 5'd25; // D4: m' = byteEncode_1(Compress_1(w))
localparam ST_DONE = 5'd31;
reg [4:0] st, st_next;
@@ -499,8 +510,8 @@ module mlkem_top #(
assign dbg_sigma_o = sigma_r;
assign dbg_r_o = r_r;
assign dbg_hek_o = hek_r;
// Decaps taps: m' reuses m_r (Decrypt writes it), z/h parsed from dk at load.
assign dbg_mprime_o = m_r;
// Decaps taps: m' from D4 (mprime_r); z/h parsed from dk at load.
assign dbg_mprime_o = mprime_r;
assign dbg_kbar_o = kbar_r;
assign dbg_decz_o = z_r;
assign dbg_dech_o = hek_r; // Decaps parses dk's H(ek) into hek_r at load
@@ -748,7 +759,7 @@ module mlkem_top #(
wire [13:0] ntt_rd_full = n_slot*256 + n_ridx[7:0];
// ntt_core inputs muxed: KeyGen/Encaps fwd NTT feeds bse_rd_data (mode 0);
// Encaps E4/E6 INTT sub-phase feeds psum (bt_rd_data) with mode 1.
wire u_intt_act = (st == ST_ENC_U || st == ST_ENC_V) && (u_sub == 2'd1);
wire u_intt_act = (st == ST_ENC_U || st == ST_ENC_V || st == ST_DEC_W) && (u_sub == 2'd1);
wire [11:0] ntt_in = u_intt_act ? bt_rd_data : bse_rd_data;
wire ntt_vin = u_intt_act ? u_nvalid : n_valid;
wire ntt_mode = u_intt_act ? 1'b1 : 1'b0;
@@ -873,7 +884,7 @@ module mlkem_top #(
// - mu[w] = Decompress_1(m bit w) = m_r[w] ? 1665 : 0 (1665 = round(Q/2)).
// u_row_max bounds the row loop: K for U, 1 for V.
localparam E2_ASLOT = 12'd1; // bank_a slot holding relocated e2 (never a t_hat slot, K>=2)
wire [2:0] u_row_max = (st == ST_ENC_V) ? 3'd1 : k_r;
wire [2:0] u_row_max = (st == ST_ENC_V || st == ST_DEC_W) ? 3'd1 : k_r;
// V-ADD: read psum (bank_t UPSUM) + e2 (bank_a E2_ASLOT) at u_aidx (lead),
// add mu (by write index u_awidx), write v to bank_t[UPSUM] at u_awidx.
wire [13:0] u_v_e2rd = E2_ASLOT*256 + u_aidx[7:0]; // bank_a e2 (lead)
@@ -884,6 +895,34 @@ module mlkem_top #(
wire [11:0] u_vq = (u_vsub1 >= 14'(Q)) ? (u_vsub1 - 14'(Q)) : u_vsub1[11:0];
wire u_v_we = (st == ST_ENC_V) && (u_sub == 2'd2) && u_avalid;
// ================================================================
// D3: w = v' - INTT(sum_j s_hat[j] o u_hat[j]) (Decaps ST_DEC_W).
// Reuses the u_* MAC/INTT machine (u_row tied to 0, single poly):
// MAC : s_hat[j] (bank_a slot j*K, == u_aslot with u_row=0) o u_hat[j]
// (bank_se rel j, == u_pm_b_full) -> psum bank_t[UPSUM]. (same
// addressing as Encaps V MAC, which reads t_hat[j] o y_hat[j].)
// INTT: INTT(psum) in place in bank_t[UPSUM].
// SUB : w[i] = v'[i] - psum[i] mod Q. v' is in bank_a[DEC_VASLOT] (D1),
// psum in bank_t[UPSUM]; read both in parallel (like V-ADD reads
// psum + e2). w written back to bank_t[UPSUM] in place.
wire [13:0] u_w_vrd = DEC_VASLOT*256 + u_aidx[7:0]; // bank_a v' (lead)
wire [13:0] u_wsub_wr = UPSUM*256 + u_awidx; // bank_t w (write)
wire [12:0] u_wdiff = {1'b0, ba_rd_data} - {1'b0, bt_rd_data}; // v' - psum
wire [11:0] u_wq = u_wdiff[12] ? (u_wdiff[11:0] + 12'(Q)) : u_wdiff[11:0];
wire u_w_we = (st == ST_DEC_W) && (u_sub == 2'd2) && u_avalid;
// ================================================================
// D4: m' = byteEncode_1(Compress_1(w)) (Decaps ST_DEC_MENC).
// Compress_1(x) = round(2x/Q) mod 2 = 1 iff x in (Q/4, 3Q/4], i.e.
// 832 < x <= 2496 for Q=3329. Pack 256 bits LSB-first into 32 bytes.
// w is read from bank_t[UPSUM] (in place from D3), 1 coeff per 2 cycles.
reg [7:0] men_idx; // coeff 0..255
reg [1:0] men_ph; // micro-phase
reg men_done;
reg [255:0] mprime_r; // recovered message m' (32 bytes, bit-packed LSB-first)
wire [13:0] men_rd = UPSUM*256 + men_idx; // bank_t w read addr
wire men_w_bit = (bt_rd_data > 12'd832) && (bt_rd_data <= 12'd2496); // Compress_1
// ST_ENC_E2MV: copy e2 (bank_t rel slot 0, 256 coeffs) into bank_a[E2_ASLOT].
// bank_t read leads (em_ridx), bank_a write trails 1 cycle (em_we/em_widx).
reg [8:0] em_ridx; // 0..256 read-ahead over e2
@@ -953,12 +992,15 @@ module mlkem_top #(
// decompress), and write each result (mod q) to a bank.
// c1 = K polys, d=du -> u'[i] in bank_se rel i (0..K-1)
// c2 = 1 poly, d=dv -> v' in bank_t rel slot DEC_VSLOT
// c1 = K polys, d=du -> u'[i] in bank_se rel i (0..K-1)
// c2 = 1 poly, d=dv -> v' in bank_a slot DEC_VASLOT (so D3's w-SUB can
// read psum (bank_t) and v' (bank_a) in parallel, like Encaps V-ADD).
// micro-phase dec_ph:
// 0 present c_in byte addr; 1 capture byte, append 8 bits to buf;
// while (nbits>=d): extract d-bit symbol -> comp_decomp (decompress);
// 2 wait pipe; 3 cd_vo -> write coeff to bank, advance coeff index.
// ================================================================
localparam DEC_VSLOT = 10'd2; // bank_t rel slot holding v' (avoid UPSUM=1)
localparam DEC_VASLOT = 12'd1; // bank_a slot for v' (free: s_hat at j*K, slot 1 unused K>=2)
reg [2:0] dec_poly; // c1: 0..K-1; c2: single (=K marker)
reg [7:0] dec_cidx; // coeff 0..255 within poly (feed side)
reg [7:0] dec_widx; // coeff 0..255 within poly (writeback side, lags pipe)
@@ -1029,7 +1071,11 @@ module mlkem_top #(
// 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_DEC_NTT: if (n_slot >= {2'b0, k_r}) st_next = ST_DEC_W;
// D3: w = v' - INTT(sum_j s_hat[j] o u_hat[j]). Single output poly
// (u_row 0..0); MAC->INTT->SUB, then D4 encodes m'.
ST_DEC_W: if (u_row >= 3'd1) st_next = ST_DEC_MENC;
ST_DEC_MENC: if (men_done) 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;
@@ -1170,6 +1216,10 @@ module mlkem_top #(
dec_dc_coeff <= 12'd0;
dec_dc_valid <= 1'b0;
dec_dc_done <= 1'b0;
men_idx <= 8'd0;
men_ph <= 2'd0;
men_done <= 1'b0;
mprime_r <= 256'd0;
h_blk <= 3'd0;
h_byte <= 8'd0;
h_phase <= 2'd0;
@@ -1445,7 +1495,7 @@ module mlkem_top #(
// 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).
if ((st == ST_ENC_U || st == ST_ENC_V) && u_sub == 2'd0) begin
if ((st == ST_ENC_U || st == ST_ENC_V || st == ST_DEC_W) && u_sub == 2'd0) begin
u_j0q <= (u_j == 3'd0);
if (u_loading) begin
@@ -1490,7 +1540,7 @@ module mlkem_top #(
end
// ---- ST_ENC_U/V sub-phase 1: INTT(psum) mode=1 in place (bank_t) ----
if ((st == ST_ENC_U || st == ST_ENC_V) && u_sub == 2'd1) begin
if ((st == ST_ENC_U || st == ST_ENC_V || st == ST_DEC_W) && u_sub == 2'd1) begin
if (u_nloading) begin
if (u_ridx == 9'd256) begin
u_nloading <= 1'b0;
@@ -1510,13 +1560,14 @@ module mlkem_top #(
end
end
// ---- ST_ENC_U/V sub-phase 2: ADD + writeback ----
// ---- ST_ENC_U/V + ST_DEC_W sub-phase 2: ADD/SUB + writeback ----
// U: u[u_row] = psum + e1[u_row] mod Q -> bank_se (u_add_we).
// V: v = psum + e2 + mu mod Q -> bank_t[UPSUM] (u_v_we).
// Read psum (bank_t) + e1/e2 at u_aidx; both arrive 1 cyc later, so
// D3 w = v' - psum mod Q -> bank_t[UPSUM] (u_w_we).
// Read psum (bank_t) + e1/e2/v' at u_aidx; both arrive 1 cyc later, so
// register (valid,widx) and write next cycle. u_row_max bounds the
// row loop (K for U, 1 for V).
if ((st == ST_ENC_U || st == ST_ENC_V) && u_sub == 2'd2) begin
// row loop (K for U, 1 for V/D3).
if ((st == ST_ENC_U || st == ST_ENC_V || st == ST_DEC_W) && u_sub == 2'd2) begin
if (u_aidx < 9'd256) begin
u_aidx <= u_aidx + 9'd1;
u_avalid <= 1'b1; // addr presented this cyc -> write next
@@ -1647,7 +1698,7 @@ module mlkem_top #(
// ph2: drop dc_valid (pulse captured) -> ph3.
// ph3: cd_vo -> dec_u_we/dec_v_we writes cd_out; advance coeff.
// c1 = K polys (d=du) -> u'[poly] bank_se; then c2 = 1 poly (d=dv)
// -> v' bank_t[DEC_VSLOT]. Byte stream is contiguous in c_in_bram.
// -> v' bank_a[DEC_VASLOT]. Byte stream is contiguous in c_in_bram.
if (st == ST_DEC_DECOMP && !dec_dc_done) begin
case (dec_ph)
2'd0: begin
@@ -1952,6 +2003,42 @@ module mlkem_top #(
u_pending <= 1'b0;
pm_valid <= 1'b0;
end
// Arm D3 (ST_DEC_W) when D2 NTT finishes: same u_* MAC/INTT machine,
// single row (u_row 0..0). MAC s_hat[j] o u_hat[j], INTT, then SUB.
if (st == ST_DEC_NTT && st_next == ST_DEC_W) begin
u_row <= 3'd0;
u_sub <= 2'd0; // MAC
u_j <= 3'd0;
u_ld <= 9'd0;
u_oidx <= 8'd0;
u_loading <= 1'b1;
u_pending <= 1'b0;
pm_valid <= 1'b0;
end
// Arm D4 (ST_DEC_MENC) when D3 w done: byteEncode_1(Compress_1(w)).
if (st == ST_DEC_W && st_next == ST_DEC_MENC) begin
men_idx <= 8'd0;
men_ph <= 2'd0;
men_done <= 1'b0;
end
// ---- ST_DEC_MENC (D4): m' = byteEncode_1(Compress_1(w)) ----
// Compress_1(w)=1 iff 832 < w <= 2496 (Q=3329). Pack 256 bits
// LSB-first into mprime_r (bit men_idx). Read w from bank_t[UPSUM].
// ph0: present w[men_idx] addr; ph1: bt_rd_data valid -> set bit.
if (st == ST_DEC_MENC && !men_done) begin
case (men_ph)
2'd0: men_ph <= 2'd1; // addr presented; wait read
default: begin // ph1: bt_rd_data = w[men_idx]
mprime_r[men_idx] <= men_w_bit;
if (men_idx == 8'd255) men_done <= 1'b1;
else men_idx <= men_idx + 8'd1;
men_ph <= 2'd0;
end
endcase
end
end
end