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

feat(enc): Encaps E4 - u = INTT(sum A^T o y_hat) + e1

Browse Source 
ST_ENC_U computes u[i] per row in 3 sub-phases reusing shared u_pmul + u_ntt:
  sub0 MAC : sum_j A_hat[j][i] o y_hat[j] (TRANSPOSE: slot=j*K+i) -> NTT-domain
             psum in bank_t rel slot UPSUM=1 (e2 in slot 0), init 0 at j==0
  sub1 INTT: INTT(psum) mode=1 (built-in x3303) in place in bank_t[UPSUM]
  sub2 ADD : u[i][w] = psum[w] + e1[i][w] mod Q -> bank_se rel (K+i), over e1
y_hat (bank_se 0..K-1) preserved for V. ntt_core mode + input muxed for the
INTT sub-phase; bank_a/se/t read+write ports extended for all 3 sub-phases.

Fixed a duplicate 'assign bse_we' (stale + new both present -> ADD writes
X-dropped); collapsed to one. Verified (K=2 c0) u[0..1] == ml-kem-r golden
(transpose + INTT + e1 all correct); E0/E1/E3 pass, E2 trimmed to e2 (e1
consumed into u, transitively checked by E4); K=3/4 no timeout.
This commit is contained in:
FallenSigh
2026-06-29 02:26:01 +08:00
parent 8ed4d59546
commit ee875d2ff7
2 changed files with 242 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

45
sync_rtl/top/TB/tb_mlkem_enc_katK_xsim.v
View File

@@ -100,6 +100,7 @@ module tb_mlkem_enc_katK_xsim;
verify_e1;
verify_e2;
verify_e3;
verify_e4;
end
$finish;
end
@@ -135,32 +136,20 @@ module tb_mlkem_enc_katK_xsim;
// e2 lives in bank_t rel slot 0 -> dbg slot = slot_t_rt = 8.
task verify_e2;
begin
$readmemh("sync_rtl/top/TB/vectors/encgold/ec_k2_c0_e1_0.hex", gy, 512, 767);
$readmemh("sync_rtl/top/TB/vectors/encgold/ec_k2_c0_e1_1.hex", gy, 768, 1023);
$readmemh("sync_rtl/top/TB/vectors/encgold/ec_k2_c0_e2.hex", ge2);
$readmemh("sync_rtl/top/TB/vectors/encgold/ec_k2_c0_e2.hex", ge2);
ce = 0;
// e1_0,e1_1 at bank_se dbg slots 6,7 (time-domain, untouched by NTT).
// (y at slots 4,5 is overwritten in place by y_hat -> checked in E3.)
for (slot = 2; slot < 4; slot = slot + 1)
for (idx = 0; idx < 256; idx = idx + 1) begin
dbg_slot_i = (4+slot); dbg_idx_i = idx[7:0];
@(posedge clk); @(posedge clk); @(posedge clk);
if (dbg_coeff_o !== gy[slot*256+idx]) begin
if (ce < 8) $display(" E1[s%0d,%0d] got=%03x exp=%03x", slot-2, idx, dbg_coeff_o, gy[slot*256+idx]);
ce = ce + 1;
end
end
// e2 at bank_t dbg slot 8
// e2 at bank_t dbg slot 8 (survives; e1 is consumed/overwritten by u
// in E4 -> e1 correctness is transitively verified by E4's u check).
for (idx = 0; idx < 256; idx = idx + 1) begin
dbg_slot_i = 8; dbg_idx_i = idx[7:0];
@(posedge clk); @(posedge clk); @(posedge clk);
if (dbg_coeff_o !== ge2[idx]) begin
if (ce < 12) $display(" E2[%0d] got=%03x exp=%03x", idx, dbg_coeff_o, ge2[idx]);
if (ce < 8) $display(" E2[%0d] got=%03x exp=%03x", idx, dbg_coeff_o, ge2[idx]);
ce = ce + 1;
end
end
if (ce == 0) $display("K=2 CASE 0 PASS (E2): e1[0..1],e2 == ml-kem-r golden");
else $display("K=2 CASE 0 FAIL (E2): %0d coeff mismatches", ce);
if (ce == 0) $display("K=2 CASE 0 PASS (E2): e2 == ml-kem-r golden");
else $display("K=2 CASE 0 FAIL (E2): %0d e2 mismatches", ce);
end
endtask
// y_hat[i] = NTT(y[i]) in place at bank_se rel slots 0..K-1 -> dbg slots 4..5 (K=2).
@@ -183,5 +172,25 @@ module tb_mlkem_enc_katK_xsim;
else $display("K=2 CASE 0 FAIL (E3): %0d coeff mismatches", ce);
end
endtask
// u[i] = INTT(sum A^T o y_hat) + e1[i] over e1 in bank_se rel K+i -> dbg 6,7 (K=2).
reg [11:0] gu [0:2*256-1];
task verify_e4;
begin
$readmemh("sync_rtl/top/TB/vectors/encgold/ec_k2_c0_u_0.hex", gu, 0, 255);
$readmemh("sync_rtl/top/TB/vectors/encgold/ec_k2_c0_u_1.hex", gu, 256, 511);
ce = 0;
for (slot = 0; slot < 2; slot = slot + 1)
for (idx = 0; idx < 256; idx = idx + 1) begin
dbg_slot_i = (6+slot); dbg_idx_i = idx[7:0]; // bank_se rel K+slot (K=2)
@(posedge clk); @(posedge clk); @(posedge clk);
if (dbg_coeff_o !== gu[slot*256+idx]) begin
if (ce < 8) $display(" U[s%0d,%0d] got=%03x exp=%03x", slot, idx, dbg_coeff_o, gu[slot*256+idx]);
ce = ce + 1;
end
end
if (ce == 0) $display("K=2 CASE 0 PASS (E4): u[0..1] == ml-kem-r golden");
else $display("K=2 CASE 0 FAIL (E4): %0d coeff mismatches", ce);
end
endtask
initial begin #120000000; $display("FAIL: global timeout"); $finish; end
endmodule

230
sync_rtl/top/mlkem_top.v
View File

@@ -216,12 +216,21 @@ module mlkem_top #(
// e2 (ST_ENC_C c_poly==2K -> bank_t rel slot 0, free during C/N/U).
wire e2_we = (st == ST_ENC_C) && c_busy && cbd_vo && cbd_ack &&
(c_poly == {1'b0, k_r, 1'b0}); // c_poly == 2K
assign bt_we = ((st == ST_M) && pm_vo) || td_we || e2_we;
// Encaps E4 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.
wire u_psum_we = (st == ST_ENC_U) && (u_sub == 2'd0) && pm_vo;
wire u_intt_we = (st == ST_ENC_U) && (u_sub == 2'd1) && ntt_vo;
assign bt_we = ((st == ST_M) && pm_vo) || td_we || e2_we ||
u_psum_we || u_intt_we;
assign bt_wa = td_we ? td_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]) :
((m_i*256 + m_oidx) & ((1<<PT_AW)-1));
assign bt_wd = td_we ? td_wd :
e2_we ? cbd_modq :
u_psum_we ? u_accq :
u_intt_we ? ntt_coeff :
m_accq;
// Debug readback (registered for timing)
@@ -232,12 +241,18 @@ module mlkem_top #(
wire [13:0] dbg_a_addr = dbg_slot_i*256 + dbg_idx_i;
wire [13:0] dbg_se_addr = (dbg_slot_i - slot_s_rt)*256 + dbg_idx_i;
wire [13:0] dbg_t_addr = (dbg_slot_i - slot_t_rt)*256 + dbg_idx_i;
// bank_a read port: ST_M load drives pm_a_full; otherwise debug index.
assign ba_rd_addr = (st == ST_M) ? pm_a_full[PA_AW-1:0] : dbg_a_addr[PA_AW-1:0];
// bank_t read port: ST_M acc drives the t_hat accumulate addr; ST_E ek-half
// drives the byteEncode addr; otherwise debug index.
// bank_a read port: ST_M / Encaps ST_ENC_U MAC drive pm_a addr; else dbg.
assign ba_rd_addr = (st == ST_M) ? pm_a_full[PA_AW-1:0] :
(st == ST_ENC_U) ? u_pm_a_full[PA_AW-1:0] :
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 (MAC psum acc, or INTT psum load by sub-phase); 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) ?
((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
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.
@@ -245,18 +260,25 @@ module mlkem_top #(
(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] :
(st == ST_ENC_U) ?
((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
dbg_se_addr[PSE_AW-1:0];
// 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/ST_ENC_N NTT writeback.
// e2 at c_poly==2K goes to bank_t instead), ST_N/ST_ENC_N NTT writeback,
// Encaps E4 ADD u-writeback (u_sub==2).
// Encaps E4 ADD: write u[u_row] over e1 in bank_se rel slot (K+u_row).
wire u_add_we = (st == ST_ENC_U) && (u_sub == 2'd2) && u_avalid;
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);
assign bse_wa = (st == ST_N || st == ST_ENC_N) ? ((n_slot*256 + n_widx) & ((1<<PSE_AW)-1))
(((st == ST_N) || (st == ST_ENC_N)) && ntt_vo) ||
u_add_we;
assign bse_wa = u_add_we ? u_add_uwr[PSE_AW-1:0] :
(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 || st == ST_ENC_N) ? ntt_coeff : cbd_modq;
assign bse_wd = u_add_we ? u_uq :
(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)
@@ -606,14 +628,19 @@ module mlkem_top #(
// bank_se read addr for the NTT load (relative slot = n_slot); sd_bram
// registers it into bse_rd_data, which feeds ntt_core 1 cycle later.
wire [13:0] ntt_rd_full = n_slot*256 + n_ridx[7:0];
wire [11:0] ntt_in = bse_rd_data;
// ntt_core inputs muxed: KeyGen/Encaps fwd NTT feeds bse_rd_data (mode 0);
// Encaps E4 INTT sub-phase feeds psum (bt_rd_data) with mode 1.
wire u_intt_act = (st == ST_ENC_U) && (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;
ntt_core u_ntt (
.clk(clk), .rst_n(rst_n),
.coeff_in(ntt_in),
.valid_i(n_valid),
.valid_i(ntt_vin),
.ready_o(ntt_ready),
.mode(1'b0), // forward NTT, no scaling
.mode(ntt_mode), // 0 = fwd (NTT), 1 = inv (INTT, x3303)
.coeff_out(ntt_coeff),
.valid_o(ntt_vo),
.ready_i(1'b1), // always accept output
@@ -669,6 +696,53 @@ module mlkem_top #(
wire [SAW-1:0] m_eslot = slot_e_rt + m_i; // e_hat[i]
wire [SAW-1:0] m_tslot = slot_t_rt + m_i; // t_hat[i]
// ================================================================
// E4: u[i] = INTT(sum_j A^T[i][j] o y_hat[j]) + e1[i] (Encaps ST_ENC_U)
// Per row i, three sub-phases (u_sub):
// 0 MAC : sum_j A_hat[j][i] o y_hat[j] (transpose) -> psum bank_t[UPSUM]
// 1 INTT: INTT(psum) mode=1 in place in bank_t[UPSUM]
// 2 ADD : u[i][w] = psum[w] + e1[i][w] mod Q -> bank_se rel (K+i) (over e1)
// y_hat (bank_se 0..K-1) survives for V. Reuses shared u_pmul + u_ntt.
// ================================================================
localparam UPSUM = 10'd1; // bank_t rel slot for the NTT-domain psum
reg [2:0] u_row; // output row i 0..K-1
reg [1:0] u_sub; // 0=MAC 1=INTT 2=ADD
// MAC bookkeeping (mirrors ST_M)
reg [2:0] u_j; // term col j 0..K-1
reg [8:0] u_ld; // load read-ahead 0..256
reg [7:0] u_oidx; // accumulate index 0..255
reg u_loading; // streaming pairs into poly_mul
reg u_pending; // wait poly_mul IDLE before next j
reg u_j0q; // (u_j==0) delayed 1 cyc (first-term=0 select)
// INTT bookkeeping (mirrors ST_N)
reg [8:0] u_ridx; // INTT load read-ahead 0..256
reg [7:0] u_widx; // INTT output index 0..255
reg u_nvalid; // feed psum to ntt_core (mode=1)
reg u_nloading; // presenting INTT load addresses
reg u_npending; // (unused placeholder for symmetry)
// ADD bookkeeping
reg [8:0] u_aidx; // ADD read-ahead 0..256
reg [7:0] u_awidx; // ADD write index 0..255
reg u_avalid; // ADD write enable pipeline
// MAC addresses: A^T -> A_hat[u_j][u_row] slot = u_j*k+u_row; y_hat[u_j] rel u_j
wire [SAW-1:0] u_aslot = u_j*k_r + u_row;
wire [13:0] u_pm_a_full = u_aslot*256 + u_ld[7:0]; // bank_a
wire [13:0] u_pm_b_full = u_j*256 + u_ld[7:0]; // bank_se y_hat[u_j]
wire [7:0] u_acc_radr = pm_vo ? (u_oidx + 8'd1) : u_oidx;
wire [13:0] u_psum_full = UPSUM*256 + u_acc_radr; // bank_t psum acc addr (MAC)
wire [11:0] u_acc_src = u_j0q ? 12'd0 : bt_rd_data; // first term init 0
wire [12:0] u_sum = {1'b0, u_acc_src} + {1'b0, pm_coeff};
wire [11:0] u_accq = (u_sum >= 13'(Q)) ? (u_sum - 13'(Q)) : u_sum[11:0];
// INTT load addr (psum read) + in-place writeback to bank_t[UPSUM]
wire [13:0] u_intt_rd = UPSUM*256 + u_ridx[7:0];
wire [13:0] u_intt_wr = UPSUM*256 + u_widx;
// ADD: read psum (bank_t) + e1 (bank_se K+u_row), write u over e1 (read-ahead 1)
wire [13:0] u_add_prd = UPSUM*256 + u_aidx[7:0]; // bank_t psum
wire [13:0] u_add_e1rd = ({2'b0,k_r}+{2'b0,u_row})*256 + u_aidx[7:0]; // bank_se e1 (lead)
wire [13:0] u_add_uwr = ({2'b0,k_r}+{2'b0,u_row})*256 + u_awidx; // bank_se u (write)
wire [12:0] u_usum = {1'b0, bt_rd_data} + {1'b0, bse_rd_data}; // psum + e1
wire [11:0] u_uq = (u_usum >= 13'(Q)) ? (u_usum - 13'(Q)) : u_usum[11:0];
reg pm_valid;
wire pm_ready;
wire [11:0] pm_coeff;
@@ -729,7 +803,8 @@ module mlkem_top #(
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_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_N: if (n_slot >= {2'b0, k_r}) st_next = ST_ENC_U; // K slots (y_hat)
ST_ENC_U: if (u_row >= k_r) st_next = ST_DONE; // E4: u[0..K-1] done
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;
@@ -787,6 +862,22 @@ module mlkem_top #(
m_pending <= 1'b0;
pm_valid <= 1'b0;
m_jq <= 1'b0;
u_row <= 3'd0;
u_sub <= 2'd0;
u_j <= 3'd0;
u_ld <= 9'd0;
u_oidx <= 8'd0;
u_loading <= 1'b0;
u_pending <= 1'b0;
u_j0q <= 1'b0;
u_ridx <= 9'd0;
u_widx <= 8'd0;
u_nvalid <= 1'b0;
u_nloading <= 1'b0;
u_npending <= 1'b0;
u_aidx <= 9'd0;
u_awidx <= 8'd0;
u_avalid <= 1'b0;
e_poly <= 3'd0;
e_pair <= 8'd0;
e_ph <= 2'd0;
@@ -1046,6 +1137,115 @@ module mlkem_top #(
end
end
// Arm Encaps U stage when N finishes: prime first row MAC (i=0,j=0).
if (st == ST_ENC_N && st_next == ST_ENC_U) 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
// ---- ST_ENC_U sub-phase 0: MAC sum_j A^T[u_row][j] o y_hat[j] ----
// 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).
if (st == ST_ENC_U && u_sub == 2'd0) begin
u_j0q <= (u_j == 3'd0);
if (u_loading) begin
if (u_ld == 9'd256) begin
pm_valid <= 1'b0;
u_loading <= 1'b0;
u_ld <= 9'd0;
u_oidx <= 8'd0;
end else begin
u_ld <= u_ld + 9'd1;
pm_valid <= 1'b1;
end
end
if (pm_vo) begin
// bank_t psum write is combinational (u_psum_we); advance idx
if (u_oidx == 8'd255) begin
if (u_j + 3'd1 < k_r) begin
u_j <= u_j + 3'd1;
u_pending <= 1'b1; // next term, same row
end else begin
// row's MAC done -> INTT sub-phase
u_sub <= 2'd1;
u_ridx <= 9'd0;
u_widx <= 8'd0;
u_nvalid <= 1'b0;
u_nloading <= 1'b1;
end
end else begin
u_oidx <= u_oidx + 8'd1;
end
end
// re-prime next term's poly_mul load
if (u_pending && pm_ready && !pm_vo) begin
u_loading <= 1'b1;
u_ld <= 9'd0;
u_oidx <= 8'd0;
pm_valid <= 1'b0;
u_pending <= 1'b0;
end
end
// ---- ST_ENC_U sub-phase 1: INTT(psum) mode=1 in place (bank_t) ----
if (st == ST_ENC_U && u_sub == 2'd1) begin
if (u_nloading) begin
if (u_ridx == 9'd256) begin
u_nloading <= 1'b0;
u_nvalid <= 1'b0;
end else begin
u_ridx <= u_ridx + 9'd1;
u_nvalid <= 1'b1;
end
end
if (ntt_vo) u_widx <= u_widx + 8'd1; // in-place write (u_intt_we)
if (ntt_done) begin
// INTT done -> ADD sub-phase (psum + e1[u_row])
u_sub <= 2'd2;
u_aidx <= 9'd0;
u_awidx <= 8'd0;
u_avalid <= 1'b0;
end
end
// ---- ST_ENC_U sub-phase 2: u[u_row] = psum + e1[u_row] mod Q ----
// Read psum (bank_t) + e1 (bank_se) at u_aidx; both arrive 1 cyc
// later, so register (valid,widx) and write u over e1 next cycle.
if (st == ST_ENC_U && 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
u_awidx <= u_aidx[7:0];
end else begin
u_avalid <= 1'b0; // no more addresses
end
// u write (bse_we via u_add_we) commits index u_awidx this cycle.
// Advance row after the 256th coeff (u_awidx==255) is written.
if (u_avalid && u_awidx == 8'd255) begin
if (u_row + 3'd1 < k_r) begin
u_row <= u_row + 3'd1;
u_sub <= 2'd0; // next row 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 else begin
u_row <= u_row + 3'd1; // == K -> ST_DONE
end
end
end
// Arm E stage when M finishes
if (st == ST_M && st_next == ST_E) begin
e_poly <= 3'd0;