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

feat(enc): Encaps E6 - v = INTT(sum t_hat o y_hat) + e2 + mu

Browse Source 
Storage choreography (per plan risk notes):
- TDEC now decodes t_hat[j] into bank_a slot j*K (was bank_t). This makes
  V's MAC reuse E4's u_aslot=u_j*K+u_row addressing with u_row=0, no mux
  change. bank_t has no room for K=4 (t_hat would fill all 4 slots vs
  psum's UPSUM slot), hence bank_a (16 slots, A_hat dead after E4).
- New ST_ENC_E2MV state relocates e2 (bank_t[0]) -> bank_a[1] so V-ADD
  reads psum (bank_t) and e2 (bank_a) from different banks (no port
  conflict). bank_se (y_hat + u) stays intact -> verify_e2/e3/e4 unaffected.
- V reuses the u_* MAC/INTT/ADD machine with u_row tied to 0 (u_row_max=1).
  ADD computes psum + e2 + mu mod Q -> bank_t[UPSUM] in place;
  mu[w] = m_r[w] ? 1665 : 0 (Decompress_1). FSM: C1->TDEC->E2MV->V->DONE.

Bug found+fixed during bring-up: e2 relocation was off-by-one (wrote
e2[i+1] into slot i) because em_we/em_widx were registered an extra cycle
past the bram read. Fixed: em_widx==em_ridx, write scheduled for the cycle
bt_rd_data presents e2[em_ridx].

TB: verify_e6 compares v (bank_t dbg slot 9, K=2) to ml-kem-r golden.
verify_e1 dropped (TDEC overwrites bank_a A_hat slots; A_hat transitively
verified by E4). Verified: K=2 E2/E3/E4/E6 == golden, E5 c1 == KAT prefix;
K=3/4 E0+E5 pass; KeyGen K=2 unregressed.
This commit is contained in:
FallenSigh
2026-06-29 11:03:33 +08:00
parent 4fee8bded3
commit e114bec5ee
3 changed files with 169 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
.claude/plans/encaps_plan.md
View File

@@ -72,7 +72,7 @@ ct_bytes_rt = c1_bytes_rt + c2_bytes_rt; // 768 / 1088 / 1568
- **E3 — ŷ = NTT(y)**:ST_ENC_N,mode=0,就地。dbg 对 y_hat。
- **E4 — u = INTT(Σ Â[j][i]∘ŷ[j]) + e1**:ST_ENC_U,poly_mul + 累加 + INTT(mode=1) + 加 e1。**转置寻址 slot=j·K+i**。dbg 对 u。
- **E5 — Compress_du + byteEncode_du → c1** ✅:comp_decomp(mode0,d=du) + 通用 LSB-first 打包器(ST_ENC_C1)写 ct_bram c1 区。dbg_ct tap 比 ct[0..c1_bytes]==KAT.ct 前缀,K=2/3/4 全过(含 K=4 du=11 跨字节)。runner = `./run_enc.sh K CASE`。
- **E6 — v = INTT(Σ t̂[j]∘ŷ[j]) + e2 + mu**:ST_ENC_V,mu 流内由 m bit 生成。dbg 对 v。
- **E6 — v = INTT(Σ t̂[j]∘ŷ[j]) + e2 + mu** ✅:ST_ENC_C1→ST_ENC_TDEC(t̂ 解码进 **bank_a** slot j·K)→ST_ENC_E2MV(e2 从 bank_t[0] 搬到 **bank_a[1]**)→ST_ENC_V(复用 u_* MAC/INTT/ADD,u_row≡0)。mu[w]=m_r[w]?1665:0。v 落 bank_t[UPSUM],dbg slot 9(K=2)。dbg 对 v==ml-kem-r golden 全过。**踩坑:e2 搬迁曾 off-by-one(写 e2[i+1] 到 slot i),已修(em_widx==em_ridx,写延后 1 拍对齐 bram 读)。** verify_e1 已废弃(TDEC 覆盖 bank_a 的 A_hat;A_hat 由 E4 转置 MAC 间接验证)。
- **E7 — Compress_dv + byteEncode_dv → c2 + 端到端 KAT**:写 ct c2 区。干净 TB 喂 ek/m,比 ct(全长)==KAT.ct 且 ss==KAT.ss,K=2/3/4 各 count=0..N。
## 验证

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

@@ -102,10 +102,13 @@ module tb_mlkem_enc_katK_xsim;
// bank_t during C/N/U. A_hat equals KeyGen golden (K=2 c0). ----
// ---- E2: verify y[i], e1[i] (bank_se), e2 (bank_t slot_t) vs ml-kem-r.
if (KP == 2 && casenum == 0) begin
verify_e1;
// E1 (verify_e1) dropped: E6's TDEC overwrites bank_a (A_hat slots)
// with t_hat, so a post-run A_hat readback is invalid. A_hat is
// transitively verified by E4 (transpose MAC) and E6 (v uses t_hat).
verify_e2;
verify_e3;
verify_e4;
verify_e6;
end
// E5: c1 = byteEncode_du(Compress_du(u)) must equal KAT.ct[0..C1B-1].
// Runs for every K/case (ct_b is the full KAT ciphertext).
@@ -221,4 +224,24 @@ module tb_mlkem_enc_katK_xsim;
else $display("K=%0d CASE %0d FAIL (E5): %0d c1 byte mismatches", KP, casenum, be);
end
endtask
// E6: v = INTT(sum_j t_hat[j] o y_hat[j]) + e2 + mu lives in bank_t rel slot
// UPSUM=1 -> dbg slot slot_t_rt+1 = 9 (K=2). Compare to ml-kem-r golden v.
reg [11:0] gv [0:255];
task verify_e6;
begin
$readmemh("sync_rtl/top/TB/vectors/encgold/ec_k2_c0_v.hex", gv);
ce = 0;
for (idx = 0; idx < 256; idx = idx + 1) begin
dbg_slot_i = 9; dbg_idx_i = idx[7:0]; // bank_t rel UPSUM (K=2)
@(posedge clk); @(posedge clk); @(posedge clk);
if (dbg_coeff_o !== gv[idx]) begin
if (ce < 8) $display(" V[%0d] got=%03x exp=%03x", idx, dbg_coeff_o, gv[idx]);
ce = ce + 1;
end
end
if (ce == 0) $display("K=2 CASE 0 PASS (E6): v == ml-kem-r golden");
else $display("K=2 CASE 0 FAIL (E6): %0d coeff mismatches", ce);
end
endtask
endmodule

187
sync_rtl/top/mlkem_top.v
View File

@@ -124,9 +124,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]};
// bank_t write for TDEC (registered); muxed with ST_M write below.
// 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;
reg [9:0] td_wa; // PT_AW=10 (declared below; literal here)
reg [11:0] td_wa; // bank_a address (PA_AW=12)
reg [11:0] td_wd;
// ---- ct_bram: ciphertext byte buffer (<=1568 B). Written by E5/E7
@@ -189,9 +190,17 @@ module mlkem_top #(
.wr_en(ba_we), .wr_addr(ba_wa), .wr_data(ba_wd)
);
// ST_A / ST_ENC_A write: commit snt_coeff to A_hat[a_slot] when accepted.
assign ba_we = (st == ST_A || st == ST_ENC_A) && a_busy && snt_vo && snt_ack;
assign ba_wa = (a_slot*256 + a_widx) & ((1<<PA_AW)-1);
assign ba_wd = snt_coeff;
// 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).
assign ba_we = ((st == ST_A || st == ST_ENC_A) && a_busy && snt_vo && snt_ack)
|| td_we || em_we;
assign ba_wa = td_we ? td_wa[PA_AW-1:0]
: em_we ? em_wa[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
snt_coeff;
// bank_se: s_hat || e_hat. Readers = ST_N load, ST_M load (pm_b, s_hat[j]),
// ST_M acc (e_hat, j==0), ST_E dk-half, dbg. Writers = ST_C (CBD), ST_N
@@ -222,25 +231,26 @@ module mlkem_top #(
.clk(clk), .rd_addr(bt_rd_addr), .rd_data(bt_rd_data),
.wr_en(bt_we), .wr_addr(bt_wa), .wr_data(bt_wd)
);
// bank_t write: KeyGen ST_M accumulate, Encaps TDEC byteDecode12, or Encaps
// e2 (ST_ENC_C c_poly==2K -> bank_t rel slot 0, free during C/N/U).
// bank_t write: KeyGen ST_M accumulate, or Encaps e2 (ST_ENC_C c_poly==2K
// -> bank_t rel slot 0, free during C/N/U). (TDEC writes bank_a now.)
wire e2_we = (st == ST_ENC_C) && c_busy && cbd_vo && cbd_ack &&
(c_poly == {1'b0, k_r, 1'b0}); // c_poly == 2K
// 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)) :
// 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;
assign bt_we = ((st == ST_M) && pm_vo) || e2_we ||
u_psum_we || u_intt_we || u_v_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]) :
((m_i*256 + m_oidx) & ((1<<PT_AW)-1));
assign bt_wd = td_we ? td_wd :
e2_we ? cbd_modq :
assign bt_wd = e2_we ? cbd_modq :
u_psum_we ? u_accq :
u_intt_we ? ntt_coeff :
u_v_we ? u_vq :
m_accq;
// Debug readback (registered for timing)
@@ -251,18 +261,24 @@ 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 / Encaps ST_ENC_U MAC drive pm_a addr; else dbg.
// bank_a read port: ST_M / Encaps ST_ENC_U MAC drive pm_a addr; ST_ENC_V
// MAC reads t_hat[j] (u_pm_a_full, u_row=0) and ADD reads e2 (E2_ASLOT); 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] :
(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]
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.
// 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.
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_U || st == ST_ENC_V) ?
((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
(st == ST_ENC_E2MV) ? em_rd[PT_AW-1:0] : // e2 relocate 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.
@@ -273,6 +289,7 @@ module mlkem_top #(
(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
(st == ST_ENC_V) ? u_pm_b_full[PSE_AW-1:0] : // V MAC: y_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;
@@ -364,13 +381,14 @@ module mlkem_top #(
localparam ST_ENC_H = 5'd9; // H(ek) via multi-block SHA3-256 (reuses ST_H logic)
localparam ST_ENC_G = 5'd10; // (K,r) = G(m || H(ek)), 64-byte single block
localparam ST_ENC_A = 5'd11; // regenerate A_hat (transpose used in U)
localparam ST_ENC_TDEC = 5'd12; // byteDecode12: ek -> t_hat (bank_t)
localparam ST_ENC_TDEC = 5'd12; // byteDecode12: ek -> t_hat (bank_a slot j*K)
localparam ST_ENC_C = 5'd13; // sample y,e1,e2 via CBD (eta1/eta2)
localparam ST_ENC_N = 5'd14; // forward NTT of y in place
localparam ST_ENC_U = 5'd15; // u = INTT(sum A^T o y_hat) + e1
localparam ST_ENC_C1 = 5'd16; // Compress_du + byteEncode_du -> ct c1
localparam ST_ENC_V = 5'd17; // v = INTT(sum t_hat o y_hat) + e2 + mu
localparam ST_ENC_C2 = 5'd18; // Compress_dv + byteEncode_dv -> ct c2
localparam ST_ENC_E2MV = 5'd19; // relocate e2 bank_t[0] -> bank_a[E2_ASLOT]
localparam ST_DONE = 5'd31;
reg [4:0] st, st_next;
@@ -640,8 +658,8 @@ module mlkem_top #(
// 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];
// 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);
// 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 [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;
@@ -754,6 +772,38 @@ module mlkem_top #(
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];
// ================================================================
// E6: v = INTT(sum_j t_hat[j] o y_hat[j]) + e2 + mu (Encaps ST_ENC_V).
// Reuses the u_* MAC/INTT/ADD machine with u_row tied to 0 (single output
// poly). Differences from E4:
// - MAC reads t_hat[j] from bank_a slot j*K (== u_aslot=u_j*K, u_row=0),
// so the bank_a MAC address is unchanged.
// - ADD computes psum + e2 + mu (not psum + e1); e2 lives in bank_a slot
// E2_ASLOT (relocated from bank_t[0] in ST_ENC_E2MV), v writes back to
// bank_t[UPSUM] in place.
// - 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;
// 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)
wire [13:0] u_vadd_wr = UPSUM*256 + u_awidx; // bank_t v (write)
wire [11:0] u_mu = m_r[u_awidx] ? 12'd1665 : 12'd0; // mu by msg bit
wire [13:0] u_vsum = {2'b0, bt_rd_data} + {2'b0, ba_rd_data} + {2'b0, u_mu}; // psum+e2+mu
wire [13:0] u_vsub1 = (u_vsum >= 14'(Q)) ? (u_vsum - 14'(Q)) : u_vsum;
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;
// 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
reg [7:0] em_widx; // bank_a write index
reg em_we; // a coeff is being written this cycle
reg em_done;
wire [13:0] em_rd = 14'd0 + em_ridx[7:0]; // bank_t rel slot 0 (e2)
wire [11:0] em_wa = E2_ASLOT*256 + em_widx; // bank_a E2_ASLOT
// ================================================================
// E5/E7: Compress_d + byteEncode_d -> ciphertext (Encaps ST_ENC_C1/C2).
// Per coeff: read poly coeff -> comp_decomp (mode 0 compress, d=du/dv) ->
@@ -855,8 +905,10 @@ module mlkem_top #(
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_ENC_U; // K slots (y_hat)
ST_ENC_U: if (u_row >= k_r) st_next = ST_ENC_C1; // u[0..K-1] done
ST_ENC_C1: if (cp_done) st_next = ST_DONE; // E5: c1 packed
ST_ENC_TDEC: if (td_done) st_next = ST_DONE; // (TDEC deferred to V-prep later)
ST_ENC_C1: if (cp_done) st_next = ST_ENC_TDEC; // E5 done -> E6
ST_ENC_TDEC: if (td_done) st_next = ST_ENC_E2MV; // t_hat decoded -> relocate e2
ST_ENC_E2MV: if (em_done) st_next = ST_ENC_V; // e2 relocated -> compute v
ST_ENC_V: if (u_row >= 3'd1) st_next = ST_DONE; // E6: v done (1 row)
ST_DONE: st_next = ST_IDLE;
default: st_next = ST_IDLE;
endcase
@@ -881,7 +933,7 @@ module mlkem_top #(
td_b2 <= 8'd0;
td_done <= 1'b0;
td_we <= 1'b0;
td_wa <= 10'd0;
td_wa <= 12'd0;
td_wd <= 12'd0;
rho_r <= 256'd0;
sigma_r <= 256'd0;
@@ -929,6 +981,10 @@ module mlkem_top #(
u_aidx <= 9'd0;
u_awidx <= 8'd0;
u_avalid <= 1'b0;
em_ridx <= 9'd0;
em_widx <= 8'd0;
em_we <= 1'b0;
em_done <= 1'b0;
cd_coeff <= 12'd0;
cd_valid <= 1'b0;
cp_poly <= 3'd0;
@@ -973,8 +1029,9 @@ module mlkem_top #(
// BRAM write-enables default low; pulsed where a byte is written.
ek_we <= 1'b0;
dkp_we <= 1'b0;
td_we <= 1'b0; // TDEC bank_t write default low
td_we <= 1'b0; // TDEC bank_a write default low
ct_we <= 1'b0; // ct_bram byte write default low (E5/E7)
em_we <= 1'b0; // e2-relocate bank_a write default low (E6)
// Kick off when entering from IDLE: KeyGen starts G; Encaps captures
// op/m and arms the H(ek) machinery (ST_ENC_H reuses the ST_H FSM).
@@ -1213,10 +1270,11 @@ module mlkem_top #(
pm_valid <= 1'b0;
end
// ---- ST_ENC_U sub-phase 0: MAC sum_j A^T[u_row][j] o y_hat[j] ----
// ---- ST_ENC_U/V sub-phase 0: MAC sum_j (A^T or t_hat)[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: 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
u_j0q <= (u_j == 3'd0);
if (u_loading) begin
@@ -1260,8 +1318,8 @@ module mlkem_top #(
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
// ---- 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 (u_nloading) begin
if (u_ridx == 9'd256) begin
u_nloading <= 1'b0;
@@ -1281,10 +1339,13 @@ module mlkem_top #(
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
// ---- ST_ENC_U/V sub-phase 2: ADD + 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
// 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
if (u_aidx < 9'd256) begin
u_aidx <= u_aidx + 9'd1;
u_avalid <= 1'b1; // addr presented this cyc -> write next
@@ -1292,10 +1353,10 @@ module mlkem_top #(
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.
// writeback (u_add_we / u_v_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
if (u_row + 3'd1 < u_row_max) begin
u_row <= u_row + 3'd1;
u_sub <= 2'd0; // next row MAC
u_j <= 3'd0;
@@ -1305,7 +1366,7 @@ module mlkem_top #(
u_pending <= 1'b0;
pm_valid <= 1'b0;
end else begin
u_row <= u_row + 3'd1; // == K -> ST_ENC_C1
u_row <= u_row + 3'd1; // U: ==K ->C1; V: ==1 ->DONE
end
end
end
@@ -1536,8 +1597,10 @@ module mlkem_top #(
a_busy <= 1'b0;
end
// Arm ST_ENC_TDEC (byteDecode12: ek -> t_hat in bank_t).
if (st == ST_ENC_A && st_next == ST_ENC_TDEC) begin
// Arm ST_ENC_TDEC (byteDecode12: ek -> t_hat in bank_a). Now follows
// E5 (ST_ENC_C1) since e2 occupied bank_t during C/N/U and t_hat is
// only needed for V. (Was armed off ST_ENC_A in the E1 scaffold.)
if (st == ST_ENC_C1 && st_next == ST_ENC_TDEC) begin
td_poly <= 3'd0;
td_trip <= 8'd0;
td_ph <= 3'd0;
@@ -1554,15 +1617,16 @@ module mlkem_top #(
if (td_ph == 3'd2) td_b1 <= ek_rd_data;
if (td_ph == 3'd3) td_b2 <= ek_rd_data;
// write decoded coeffs into bank_t (rel slot = td_poly)
// 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).
if (td_ph == 3'd3) begin
td_we <= 1'b1;
td_wa <= (td_poly*256 + {td_trip, 1'b0}) & ((1<<PT_AW)-1); // 2*trip
td_wa <= (({4'd0,td_poly}*{9'd0,k_r})*256 + {td_trip, 1'b0}) & ((1<<PA_AW)-1); // 2*trip
td_wd <= td_c0;
end
if (td_ph == 3'd4) begin
td_we <= 1'b1;
td_wa <= (td_poly*256 + {td_trip, 1'b0} + 1) & ((1<<PT_AW)-1); // 2*trip+1
td_wa <= (({4'd0,td_poly}*{9'd0,k_r})*256 + {td_trip, 1'b0} + 1) & ((1<<PA_AW)-1); // 2*trip+1
td_wd <= td_c1;
end
@@ -1580,6 +1644,43 @@ module mlkem_top #(
td_ph <= td_ph + 3'd1;
end
end
// Arm ST_ENC_E2MV when TDEC finishes: copy e2 (bank_t rel 0) into
// bank_a[E2_ASLOT]. em_ridx is a read-ahead pointer (bank_t read is
// registered, 1-cyc latency); em_we/em_widx trail it by 1.
if (st == ST_ENC_TDEC && st_next == ST_ENC_E2MV) begin
em_ridx <= 9'd0;
em_widx <= 8'd0;
em_we <= 1'b0;
em_done <= 1'b0;
end
// ---- ST_ENC_E2MV: e2 bank_t[0] -> bank_a[E2_ASLOT], 256 coeffs ----
// Present read addr em_ridx (bank_t rel 0); bank_t read is registered
// so bt_rd_data next cycle = e2[em_ridx]. Schedule the bank_a write
// (em_we/em_widx) for that next cycle with em_widx == em_ridx, so the
// write commits e2[R] into slot[R] (no off-by-one).
if (st == ST_ENC_E2MV && !em_done) begin
if (em_ridx <= 9'd255) begin // presenting a valid e2 addr this cycle
em_we <= 1'b1; // -> write next cycle (ba_wd = bt_rd_data)
em_widx <= em_ridx[7:0];
end
if (em_ridx == 9'd256) em_done <= 1'b1; // last write (e2[255]) issued
else em_ridx <= em_ridx + 9'd1;
end
// Arm ST_ENC_V when e2 relocated: prime first (only) row MAC (i=0,j=0).
// Reuses the u_* MAC/INTT/ADD machine with u_row tied to 0.
if (st == ST_ENC_E2MV && st_next == ST_ENC_V) 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
end
end