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

feat(mlkem_top): KeyGen stage 2f (byteEncode12 -> ek, dk_pke)

Browse Source 
Add ST_E stage: serialize t_hat[0..1] -> ek_mem[0..767], s_hat[0..1] ->
dkp_mem[0..767] via byteEncode12 (2 coeffs -> 3 bytes, LSB-first 12-bit:
b0=c0[7:0], b1={c1[3:0],c0[11:8]}, b2=c1[11:4]), then copy rho into
ek_mem[768..799]. Byte readback tap (dbg_byte_sel/idx -> dbg_byte_o).

Verified vs KAT-derived golden: ek 800B (== KAT pk) + dk_pke 768B
(== KAT sk prefix) byte-exact (20430 cyc). Completes Stage 2 datapath.
This commit is contained in:
FallenSigh
2026-06-28 02:03:03 +08:00
parent a9e50ebc0c
commit 17914911c3
4 changed files with 1714 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

81
sync_rtl/top/mlkem_top.v
View File

@@ -36,6 +36,10 @@ module mlkem_top #(
input [3:0] dbg_slot_i, // poly slot (see localparams below)
input [7:0] dbg_idx_i, // coefficient index 0..255
output [11:0] dbg_coeff_o,
// Debug byte readback: ek (sel=0, 0..799) / dk_pke (sel=1, 0..767)
input dbg_byte_sel_i,
input [9:0] dbg_byte_idx_i,
output [7:0] dbg_byte_o,
// Debug taps for hash outputs
output [255:0] dbg_rho_o,
output [255:0] dbg_sigma_o
@@ -65,6 +69,17 @@ module mlkem_top #(
always @(posedge clk) dbg_coeff_r <= polymem[dbg_slot_i*256 + dbg_idx_i];
assign dbg_coeff_o = dbg_coeff_r;
// ek (800B) and dk_pke (768B) byte memories (byteEncode12 output)
localparam EK_BYTES = 384*K + 32; // 800 for K=2
localparam DK_BYTES = 384*K; // 768 for K=2
reg [7:0] ek_mem [0:EK_BYTES-1];
reg [7:0] dkp_mem [0:DK_BYTES-1];
reg [7:0] dbg_byte_r;
always @(posedge clk)
dbg_byte_r <= dbg_byte_sel_i ? dkp_mem[dbg_byte_idx_i] : ek_mem[dbg_byte_idx_i];
assign dbg_byte_o = dbg_byte_r;
// ================================================================
// Top-level FSM (built incrementally). Stage 2a: G only.
// ================================================================
@@ -74,6 +89,7 @@ module mlkem_top #(
localparam ST_C = 4'd3; // generate s[i],e[i] via CBD
localparam ST_N = 4'd4; // forward NTT of s[i],e[i] in place
localparam ST_M = 4'd5; // matrix accumulate t_hat = e_hat + sum A o s_hat
localparam ST_E = 4'd6; // byteEncode12 -> ek_mem, dkp_mem
localparam ST_DONE = 4'd15;
reg [3:0] st, st_next;
@@ -206,6 +222,29 @@ module mlkem_top #(
reg m_loading; // 1 while streaming pairs into poly_mul
reg m_pending; // wait for poly_mul IDLE before next (i,j)
// ---- Stage 2f: byteEncode12 serializer ----
// Pack each poly (2 coeffs -> 3 bytes, LSB-first 12-bit). ek = t_hat[0..K-1]
// bytes || rho; dk_pke = s_hat[0..K-1] bytes. Walk coeff pairs per poly.
reg [2:0] e_poly; // which source poly: 0,1 = t_hat0,t_hat1 -> ek
// 2,3 = s_hat0,s_hat1 -> dk_pke
reg [7:0] e_pair; // 0..127 coeff-pair within poly
reg [9:0] e_rho; // 0..31 rho byte copy index (ek tail)
reg e_done; // serialization complete
// source poly slot for current e_poly
wire [3:0] e_slot = (e_poly == 3'd0) ? SLOT_T0 :
(e_poly == 3'd1) ? SLOT_T1 :
(e_poly == 3'd2) ? SLOT_S0 : SLOT_S1;
// two coeffs of the current pair
wire [11:0] e_c0 = polymem[e_slot*256 + {e_pair, 1'b0}];
wire [11:0] e_c1 = polymem[e_slot*256 + {e_pair, 1'b1}];
// 3 packed bytes
wire [7:0] e_b0 = e_c0[7:0];
wire [7:0] e_b1 = {e_c1[3:0], e_c0[11:8]};
wire [7:0] e_b2 = e_c1[11:4];
// byte base offset within target memory: poly index *384 (= 128 pairs *3)
wire [9:0] e_base = (e_poly[0]) ? 10'd384 : 10'd0; // poly0->0, poly1->384
wire [9:0] e_boff = e_base + {e_pair, 1'b0} + {2'b0, e_pair}; // pair*3
wire [3:0] m_aslot = {2'b0, m_i[0], m_j[0]}; // A_hat[i][j] slot = i*2+j (0..3)
wire [3:0] m_sslot = SLOT_S0 + {3'b0, m_j[0]}; // s_hat[j]
wire [3:0] m_eslot = SLOT_E0 + {3'b0, m_i[0]}; // e_hat[i]
@@ -244,7 +283,8 @@ module mlkem_top #(
ST_A: if (a_pair >= K*K) st_next = ST_C;
ST_C: if (c_poly >= 2*K) st_next = ST_N;
ST_N: if (n_slot >= 2*K) st_next = ST_M;
ST_M: if (m_i >= K) st_next = ST_DONE;
ST_M: if (m_i >= K) st_next = ST_E;
ST_E: if (e_done) st_next = ST_DONE;
ST_DONE: st_next = ST_IDLE;
default: st_next = ST_IDLE;
endcase
@@ -279,6 +319,10 @@ module mlkem_top #(
m_loading <= 1'b0;
m_pending <= 1'b0;
pm_valid <= 1'b0;
e_poly <= 3'd0;
e_pair <= 8'd0;
e_rho <= 10'd0;
e_done <= 1'b0;
end else begin
st <= st_next;
@@ -453,6 +497,41 @@ module mlkem_top #(
m_pending <= 1'b0;
end
end
// Arm E stage when M finishes
if (st == ST_M && st_next == ST_E) begin
e_poly <= 3'd0;
e_pair <= 8'd0;
e_rho <= 10'd0;
e_done <= 1'b0;
end
// ---- ST_E: byteEncode12 t_hat -> ek_mem, s_hat -> dkp_mem, ek tail = rho ----
if (st == ST_E && !e_done) begin
if (e_poly < 3'd4) begin
// pack current coeff-pair (3 bytes)
if (e_poly < 3'd2) begin
ek_mem[e_boff] <= e_b0;
ek_mem[e_boff + 1] <= e_b1;
ek_mem[e_boff + 2] <= e_b2;
end else begin
dkp_mem[e_boff] <= e_b0;
dkp_mem[e_boff + 1] <= e_b1;
dkp_mem[e_boff + 2] <= e_b2;
end
if (e_pair == 8'd127) begin
e_pair <= 8'd0;
e_poly <= e_poly + 3'd1; // next poly (or ->4 = rho phase)
end else begin
e_pair <= e_pair + 8'd1;
end
end else begin
// rho copy: ek_mem[768 + r] = rho byte r (r = 0..31)
ek_mem[10'd768 + e_rho] <= rho_r[e_rho*8 +: 8];
if (e_rho == 10'd31) e_done <= 1'b1;
else e_rho <= e_rho + 10'd1;
end
end
end
end