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

feat(mlkem_top): parameterize K in {2,3,4} (ML-KEM 512/768/1024)

Browse Source 
Generalize KeyGen from K=2-hardcoded to compile-time parameter K:
- eta1 derived (3 for K=2, else 2); slot layout SLOT_S/E/T = K*K+{0,K,2K},
  NUM_SLOTS = K*K+3K; SAW=5 slot-addr width.
- A-stage: explicit a_i/a_j row-major counters (slot = i*K+j) instead of
  K=2 bit-tricks. C/N stages: parametric slot bases, 2K polys.
- M-stage: m_i/m_j widened to 3-bit (must reach K=4); slots i*K+j etc.
- E-stage: 2K polys, e_is_dk split, rho offset 384*K.
- H(ek): H_NBLK=ceil((EK_BYTES+1)/136), H_LAST padding generalized;
  h_blk 4-bit. Byte mems sized EK_BYTES/DK_BYTES.
- Widen dbg_byte_idx_i to [10:0] (ek up to 1568B for K=4).

Parametric TB (tb_mlkem_kg_katK, KP generic + CASE plusarg). Verified
byte-exact vs NIST KAT:
  K=2 (512):  cases 0..4  ek 800B  / dk 1632B
  K=3 (768):  cases 0..2  ek 1184B / dk 2400B  (~36k cyc)
  K=4 (1024): cases 0..2  ek 1568B / dk 3168B  (~54k cyc)
run_tb.sh top runs all three parameter sets.
This commit is contained in:
FallenSigh
2026-06-28 02:59:58 +08:00
parent 2f46c0790f
commit b2bf798454
27 changed files with 25166 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files

175
sync_rtl/top/mlkem_top.v
View File

@@ -20,8 +20,7 @@
`include "sync_rtl/common/defines.vh"
module mlkem_top #(
parameter K = 2, // ML-KEM-512
parameter ETA1 = 3
parameter K = 2 // ML-KEM-512=2, 768=3, 1024=4 (eta1 derived)
) (
input clk,
input rst_n,
@@ -38,7 +37,7 @@ module mlkem_top #(
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,
input [10:0] dbg_byte_idx_i,
output [7:0] dbg_byte_o,
// Debug full-dk readback: dk = dk_pke(768) || ek(800) || H(ek)(32) || z(32)
// = 1632 bytes. Index 0..1631.
@@ -50,21 +49,23 @@ module mlkem_top #(
);
localparam Q = `Q; // 3329
// FIPS 203: eta1 = 3 for ML-KEM-512 (K=2), else 2 (K=3/4).
localparam ETA1 = (K == 2) ? 3 : 2;
// ================================================================
// Polynomial storage: K=2 needs A_hat[2][2]=4, s/s_hat[2], e/e_hat[2],
// t_hat[2]. Reuse slots: s and s_hat share (NTT in place), same for e.
// Slot map:
// 0..3 : A_hat[0][0],A_hat[0][1],A_hat[1][0],A_hat[1][1]
// 4..5 : s_hat[0], s_hat[1] (s[i] then overwritten by NTT)
// 6..7 : e_hat[0], e_hat[1] (e[i] then overwritten by NTT)
// 8..9 : t_hat[0], t_hat[1]
// Polynomial storage, generalized for K in {2,3,4}.
// Slot layout (each slot = 256 coeffs):
// A_hat[i][j] : slots 0 .. K*K-1 at index i*K + j
// s_hat[i] : slots SLOT_S .. +K-1 (s[i] then overwritten by NTT)
// e_hat[i] : slots SLOT_E .. +K-1
// t_hat[i] : slots SLOT_T .. +K-1
// NUM_SLOTS = K*K + 3*K (10 / 24 / 28 for K=2/3/4)
// ================================================================
localparam SLOT_A00 = 4'd0, SLOT_A01 = 4'd1, SLOT_A10 = 4'd2, SLOT_A11 = 4'd3;
localparam SLOT_S0 = 4'd4, SLOT_S1 = 4'd5;
localparam SLOT_E0 = 4'd6, SLOT_E1 = 4'd7;
localparam SLOT_T0 = 4'd8, SLOT_T1 = 4'd9;
localparam NUM_SLOTS = 10;
localparam SLOT_S = K*K; // s_hat base slot
localparam SLOT_E = K*K + K; // e_hat base slot
localparam SLOT_T = K*K + 2*K; // t_hat base slot
localparam NUM_SLOTS = K*K + 3*K;
localparam SAW = 5; // slot-address width (>=clog2(28))
reg [11:0] polymem [0:NUM_SLOTS*256-1];
@@ -73,9 +74,10 @@ 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
// ek and dk_pke byte memories (byteEncode12 output).
// ek = 384*K + 32 bytes (== KAT pk), dk_pke = 384*K bytes (== KAT sk prefix)
localparam EK_BYTES = 384*K + 32; // 800 / 1184 / 1568
localparam DK_BYTES = 384*K; // 768 / 1152 / 1536
reg [7:0] ek_mem [0:EK_BYTES-1];
reg [7:0] dkp_mem [0:DK_BYTES-1];
@@ -84,17 +86,19 @@ module mlkem_top #(
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;
// full dk = dk_pke(0..767) || ek(768..1567) || H(ek)(1568..1599) || z(1600..1631)
// full dk = dk_pke(DK_BYTES) || ek(EK_BYTES) || H(ek)(32) || z(32)
localparam DK_EK_END = DK_BYTES + EK_BYTES; // ek region end
localparam DK_HEK_END = DK_EK_END + 32; // H(ek) region end
reg [7:0] dbg_dk_r;
always @(posedge clk) begin
if (dbg_dk_idx_i < 12'd768)
if (dbg_dk_idx_i < DK_BYTES[11:0])
dbg_dk_r <= dkp_mem[dbg_dk_idx_i];
else if (dbg_dk_idx_i < 12'd1568)
dbg_dk_r <= ek_mem[dbg_dk_idx_i - 12'd768];
else if (dbg_dk_idx_i < 12'd1600)
dbg_dk_r <= hek_r[(dbg_dk_idx_i - 12'd1568)*8 +: 8];
else if (dbg_dk_idx_i < DK_EK_END[11:0])
dbg_dk_r <= ek_mem[dbg_dk_idx_i - DK_BYTES[11:0]];
else if (dbg_dk_idx_i < DK_HEK_END[11:0])
dbg_dk_r <= hek_r[(dbg_dk_idx_i - DK_EK_END[11:0])*8 +: 8];
else
dbg_dk_r <= z_i[(dbg_dk_idx_i - 12'd1600)*8 +: 8];
dbg_dk_r <= z_i[(dbg_dk_idx_i - DK_HEK_END[11:0])*8 +: 8];
end
assign dbg_dk_o = dbg_dk_r;
@@ -114,23 +118,22 @@ module mlkem_top #(
reg [3:0] st, st_next;
reg [255:0] rho_r, sigma_r;
// A-generation bookkeeping
reg [2:0] a_pair; // 0..K*K (=4) pairs done
reg [7:0] a_widx; // write index 0..255 within current poly
reg a_busy; // 1 once current pair's request accepted (gates collect)
wire [1:0] a_i = a_pair[1] ? 2'd1 : 2'd0; // pair/K (K=2)
wire [1:0] a_j = a_pair[0] ? 2'd1 : 2'd0; // pair%K
wire [3:0] a_slot = {2'b0, a_pair[1], a_pair[0]}; // SLOT_A00..A11 = pair index
// A-generation bookkeeping: explicit i/j counters (avoid runtime divide)
reg [2:0] a_i; // row 0..K-1
reg [2:0] a_j; // col 0..K-1
reg [4:0] a_pair; // 0..K*K pairs done (for done test)
reg [7:0] a_widx; // write index 0..255 within current poly
reg a_busy; // 1 once current pair's request accepted (gates collect)
wire [SAW-1:0] a_slot = a_i*K + a_j; // A_hat[i][j] slot = i*K + j
// C-generation bookkeeping: 2*K polys = s0,s1,e0,e1 (idx 0..3)
reg [2:0] c_poly; // 0..2K
reg [7:0] c_widx;
reg c_busy; // 1 once current poly's request accepted (gates collect)
wire [7:0] c_nonce = {5'b0, c_poly}; // s:0,1 e:2,3 == nonce
// slot: c_poly 0->S0,1->S1,2->E0,3->E1
wire [3:0] c_slot = (c_poly == 3'd0) ? SLOT_S0 :
(c_poly == 3'd1) ? SLOT_S1 :
(c_poly == 3'd2) ? SLOT_E0 : SLOT_E1;
// C-generation bookkeeping: 2*K polys (s[0..K-1] then e[0..K-1])
reg [4:0] c_poly; // 0..2K
reg [7:0] c_widx;
reg c_busy;
wire [7:0] c_nonce = {3'b0, c_poly}; // s:0..K-1 e:K..2K-1 == nonce
// slot: c_poly < K -> s_hat[c_poly], else e_hat[c_poly-K]
wire [SAW-1:0] c_slot = (c_poly < K) ? (SLOT_S + c_poly)
: (SLOT_E + (c_poly - K));
assign busy_o = (st != ST_IDLE);
assign done_o = (st == ST_DONE);
@@ -167,7 +170,7 @@ module mlkem_top #(
wire h_vo;
reg h_ack;
reg [255:0] hek_r; // captured H(ek)
reg [2:0] h_blk; // 0..5 block index
reg [3:0] h_blk; // 0..H_NBLK-1 block index (up to 11 for K=4)
reg [7:0] h_byte; // 0..135 byte within block being assembled
reg [1:0] h_phase; // 0=assemble 1=feed 2=wait-perm 3=done
@@ -187,17 +190,20 @@ module mlkem_top #(
.mb_ready_o(h_mbready)
);
// byte b (0..135) of block blk for SHA3-256 over 800-byte ek:
// global byte g = blk*136 + b; ek_mem[g] if g<800;
// g==800 -> 0x06 (domain+first pad bit); g==815 -> |0x80 (last block); else 0
function [7:0] h_padbyte(input [2:0] blk, input [7:0] b);
// SHA3-256 over EK_BYTES-byte ek: rate=136. Padded length = H_NBLK*136.
// pad: byte EK_BYTES = 0x06 (domain + first pad bit), last byte |= 0x80.
localparam H_NBLK = (EK_BYTES + 136) / 136; // ceil((EK_BYTES+1)/136): 6/9/12
localparam H_LAST = H_NBLK*136 - 1; // index of final padded byte
// byte b (0..135) of block blk: global g = blk*136 + b
function [7:0] h_padbyte(input [3:0] blk, input [7:0] b);
integer g;
begin
g = blk*136 + b;
if (g < 800) h_padbyte = ek_mem[g];
else if (g == 800) h_padbyte = (g == 815) ? 8'h86 : 8'h06;
else if (g == 815) h_padbyte = 8'h80; // 6th block (815 = 5*136+135)
else h_padbyte = 8'h00;
if (g < EK_BYTES) h_padbyte = ek_mem[g];
else if (g == H_LAST && g == EK_BYTES) h_padbyte = 8'h86; // 0x06|0x80
else if (g == EK_BYTES) h_padbyte = 8'h06;
else if (g == H_LAST) h_padbyte = 8'h80;
else h_padbyte = 8'h00;
end
endfunction
@@ -213,8 +219,8 @@ module mlkem_top #(
.clk(clk), .rst_n(rst_n),
.rho_i(rho_r),
.k_i(3'(K)),
.i_idx(a_i),
.j_idx(a_j),
.i_idx(a_i[1:0]),
.j_idx(a_j[1:0]),
.valid_i(snt_valid),
.ready_o(snt_ready),
.coeff_o(snt_coeff),
@@ -249,12 +255,12 @@ module mlkem_top #(
// ---- ntt_core: forward NTT (mode=0, no scaling) of s[i],e[i] in place ----
// N-stage bookkeeping: process slots S0,S1,E0,E1 (= SLOT_S0 + n_slot).
reg [2:0] n_slot; // 0..2K (4 polys)
reg [4:0] n_slot; // 0..2K (process s_hat[0..K-1] then e_hat[0..K-1])
reg [8:0] n_ridx; // load read index 0..256
reg [7:0] n_widx; // output write index 0..255
reg n_valid; // feeding coeffs to ntt_core
reg n_pending; // waiting for ntt_core IDLE to start next slot
wire [3:0] n_slot_addr = SLOT_S0 + {1'b0, n_slot};
wire [SAW-1:0] n_slot_addr = SLOT_S + n_slot; // s_hat then e_hat contiguous
wire ntt_ready;
wire [11:0] ntt_coeff;
@@ -277,8 +283,8 @@ module mlkem_top #(
// ---- poly_mul_sync: t_hat[i] = e_hat[i] + sum_j A_hat[i][j] o s_hat[j] ----
// M-stage bookkeeping. For each (i,j): LOAD 256 (A,shat) pairs, then accumulate
// 256 products into T_i (init from E_i when j==0, else from running T_i).
reg [1:0] m_i; // row 0..K
reg [1:0] m_j; // col 0..K
reg [2:0] m_i; // row 0..K (needs to reach K to exit)
reg [2:0] m_j; // col 0..K-1
reg [8:0] m_ld; // load index 0..256
reg [7:0] m_oidx; // output/accum index 0..255
reg m_loading; // 1 while streaming pairs into poly_mul
@@ -287,15 +293,14 @@ module mlkem_top #(
// ---- 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 [4:0] e_poly; // 0..2K-1: [0,K) = t_hat -> ek; [K,2K) = s_hat -> 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;
// source poly slot: t_hat[e_poly] for ek half, s_hat[e_poly-K] for dk half
wire e_is_dk = (e_poly >= K);
wire [4:0] e_pidx = e_is_dk ? (e_poly - K) : e_poly; // index within target
wire [SAW-1:0] e_slot = e_is_dk ? (SLOT_S + e_pidx) : (SLOT_T + e_pidx);
// 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}];
@@ -304,13 +309,13 @@ module mlkem_top #(
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 [11:0] e_base = e_pidx * 12'd384;
wire [11: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]
wire [3:0] m_tslot = SLOT_T0 + {3'b0, m_i[0]}; // t_hat[i]
wire [SAW-1:0] m_aslot = m_i*K + m_j; // A_hat[i][j] slot = i*K + j
wire [SAW-1:0] m_sslot = SLOT_S + m_j; // s_hat[j]
wire [SAW-1:0] m_eslot = SLOT_E + m_i; // e_hat[i]
wire [SAW-1:0] m_tslot = SLOT_T + m_i; // t_hat[i]
reg pm_valid;
wire pm_ready;
@@ -362,7 +367,9 @@ module mlkem_top #(
sha3_ack <= 1'b0;
snt_valid <= 1'b0;
snt_ack <= 1'b0;
a_pair <= 3'd0;
a_pair <= 5'd0;
a_i <= 3'd0;
a_j <= 3'd0;
a_widx <= 8'd0;
a_busy <= 1'b0;
cbd_valid <= 1'b0;
@@ -412,7 +419,9 @@ module mlkem_top #(
sha3_ack <= 1'b0;
snt_valid <= 1'b1; // start first SampleNTT
snt_ack <= 1'b1;
a_pair <= 3'd0;
a_pair <= 5'd0;
a_i <= 3'd0;
a_j <= 3'd0;
a_widx <= 8'd0;
a_busy <= 1'b0;
end
@@ -429,12 +438,18 @@ module mlkem_top #(
if (a_busy && snt_vo && snt_ack) begin
polymem[a_slot*256 + a_widx] <= snt_coeff;
if (snt_last) begin
// finished this poly; advance to next pair
a_pair <= a_pair + 3'd1;
// finished this poly; advance (i,j) in row-major order
a_pair <= a_pair + 5'd1;
a_widx <= 8'd0;
a_busy <= 1'b0;
if (a_j + 3'd1 < K) begin
a_j <= a_j + 3'd1;
end else begin
a_j <= 3'd0;
a_i <= a_i + 3'd1;
end
// start next SampleNTT if more pairs remain
if (a_pair + 3'd1 < K*K) snt_valid <= 1'b1;
if (a_pair + 5'd1 < K*K) snt_valid <= 1'b1;
end else begin
a_widx <= a_widx + 8'd1;
end
@@ -579,9 +594,9 @@ module mlkem_top #(
// ---- 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
if (e_poly < 2*K) begin
// pack current coeff-pair (3 bytes): [0,K)=ek, [K,2K)=dk_pke
if (!e_is_dk) begin
ek_mem[e_boff] <= e_b0;
ek_mem[e_boff + 1] <= e_b1;
ek_mem[e_boff + 2] <= e_b2;
@@ -592,13 +607,13 @@ module mlkem_top #(
end
if (e_pair == 8'd127) begin
e_pair <= 8'd0;
e_poly <= e_poly + 3'd1; // next poly (or ->4 = rho phase)
e_poly <= e_poly + 5'd1; // next poly (or ->2K = 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];
// rho copy: ek_mem[384*K + r] = rho byte r (r = 0..31)
ek_mem[12'(384*K) + 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
@@ -623,7 +638,7 @@ module mlkem_top #(
if (h_byte == 8'd135) begin
h_byte <= 8'd0;
h_mbvalid <= 1'b1;
h_mblast <= (h_blk == 3'd5);
h_mblast <= (h_blk == H_NBLK-1);
h_phase <= 2'd1; // feed
end else begin
h_byte <= h_byte + 8'd1;