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

refactor(kg): make ML-KEM K a runtime input k_i instead of a parameter

Browse Source 
mlkem_top now sizes storage for KMAX=4 (worst case) and selects the
active ML-KEM parameter set at start_i via the k_i input. All K-derived
quantities (eta1, slot bases, ek/dk byte counts, H(ek) block count, FSM
bounds) are computed at runtime from the captured k_r.

Verified byte-exact against NIST KAT for all three parameter sets:
  K=2 (512)  cases 0-4, K=3 (768) cases 0-2, K=4 (1024) cases 0-2
  -> 11/11 PASS (ek==pk, dk==sk).
This commit is contained in:
FallenSigh
2026-06-28 03:24:58 +08:00
parent b7e4fd9323
commit 3a53993754
2 changed files with 95 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
sync_rtl/top/TB/tb_mlkem_kg_katK_xsim.v
View File

@@ -9,6 +9,7 @@ module tb_mlkem_kg_katK_xsim;
localparam DKB = 768*KP + 96;
reg clk=0, rst_n=0, start_i=0;
reg [2:0] k_i;
reg [255:0] d_i, z_i;
wire busy_o, done_o;
reg [3:0] dbg_slot_i=0; reg [7:0] dbg_idx_i=0; wire [11:0] dbg_coeff_o;
@@ -16,8 +17,9 @@ module tb_mlkem_kg_katK_xsim;
reg [11:0] dbg_dk_idx_i=0; wire [7:0] dbg_dk_o;
wire [255:0] dbg_rho_o, dbg_sigma_o;
mlkem_top #(.K(KP)) dut (
.clk(clk), .rst_n(rst_n), .d_i(d_i), .z_i(z_i), .start_i(start_i),
// KMAX defaults to 4 (worst-case sizing); KP selects the runtime k value.
mlkem_top dut (
.clk(clk), .rst_n(rst_n), .k_i(k_i), .d_i(d_i), .z_i(z_i), .start_i(start_i),
.busy_o(busy_o), .done_o(done_o),
.dbg_slot_i(dbg_slot_i), .dbg_idx_i(dbg_idx_i), .dbg_coeff_o(dbg_coeff_o),
.dbg_byte_sel_i(dbg_byte_sel_i), .dbg_byte_idx_i(dbg_byte_idx_i), .dbg_byte_o(dbg_byte_o),
@@ -45,6 +47,7 @@ module tb_mlkem_kg_katK_xsim;
$readmemh(ekfile, ek_gold);
$readmemh(dkfile, dk_gold);
d_i = dmem[0]; z_i = zmem[0];
k_i = KP[2:0];
rst_n=0; repeat(4) @(posedge clk); rst_n=1; @(posedge clk);
start_i=1; @(posedge clk); start_i=0;

165
sync_rtl/top/mlkem_top.v
View File

@@ -1,4 +1,7 @@
// mlkem_top.v - ML-KEM-512 KeyGen top-level integration (K=2, eta1=3).
// mlkem_top.v - ML-KEM KeyGen top-level integration. Runtime-selectable
// parameter set via k_i: k=2 (ML-KEM-512, eta1=3), k=3 (768), k=4 (1024).
// Storage is sized for KMAX (worst case = ML-KEM-1024); k_i picks the
// active sub-range at start_i.
//
// Streaming valid/ready interface. Given seeds d and z, computes the
// ML-KEM key pair per FIPS 203 Algorithm 16 (KeyGen_internal):
@@ -20,10 +23,11 @@
`include "sync_rtl/common/defines.vh"
module mlkem_top #(
parameter K = 2 // ML-KEM-512=2, 768=3, 1024=4 (eta1 derived)
parameter KMAX = 4 // storage sizing (worst case = ML-KEM-1024)
) (
input clk,
input rst_n,
input [2:0] k_i, // RUNTIME ML-KEM param: 2=512, 3=768, 4=1024
input [255:0] d_i, // KeyGen seed d (byte 0 in d_i[7:0])
input [255:0] z_i, // implicit-rejection seed z
input start_i, // pulse to begin KeyGen
@@ -49,22 +53,33 @@ 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;
// Runtime ML-KEM parameter, captured at start_i.
reg [2:0] k_r;
// FIPS 203: eta1 = 3 for ML-KEM-512 (k=2), else 2 (k=3/4).
wire [1:0] eta1_rt = (k_r == 3'd2) ? 2'd3 : 2'd2;
// Runtime-derived sizes (k_r in {2,3,4}). Small multiplies are cheap.
wire [5:0] kk_rt = k_r * k_r; // 4/9/16
wire [5:0] slot_s_rt = kk_rt; // s_hat base slot
wire [5:0] slot_e_rt = kk_rt + k_r; // e_hat base slot
wire [5:0] slot_t_rt = kk_rt + {1'b0, k_r} + {1'b0, k_r}; // t_hat base = kk+2k
wire [11:0] ek_bytes_rt = 12'd384 * {9'b0, k_r} + 12'd32; // 800/1184/1568
wire [11:0] dk_bytes_rt = 12'd384 * {9'b0, k_r}; // 768/1152/1536
// H(ek) block count = ceil((ek_bytes+1)/136): 6/9/12 for k=2/3/4 (table)
wire [3:0] h_nblk_rt = (k_r == 3'd2) ? 4'd6 : (k_r == 3'd3) ? 4'd9 : 4'd12;
wire [11:0] h_last_rt = {6'b0, h_nblk_rt} * 12'd136 - 12'd1; // final padded byte index
// ================================================================
// 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)
// Polynomial storage, sized for KMAX (worst case). Runtime k uses a
// sub-range. 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_rt .. +k-1
// e_hat[i] : slots slot_e_rt .. +k-1
// t_hat[i] : slots slot_t_rt .. +k-1
// NUM_SLOTS = KMAX*KMAX + 3*KMAX = 28 for KMAX=4.
// ================================================================
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 NUM_SLOTS = KMAX*KMAX + 3*KMAX;
localparam SAW = 5; // slot-address width (>=clog2(28))
reg [11:0] polymem [0:NUM_SLOTS*256-1];
@@ -74,31 +89,30 @@ 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 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];
// ek and dk_pke byte memories sized for KMAX.
localparam EK_MAX = 384*KMAX + 32; // 1568
localparam DK_MAX = 384*KMAX; // 1536
reg [7:0] ek_mem [0:EK_MAX-1];
reg [7:0] dkp_mem [0:DK_MAX-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;
// 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
// full dk = dk_pke(dk_bytes) || ek(ek_bytes) || H(ek)(32) || z(32)
wire [11:0] dk_ek_end = dk_bytes_rt + ek_bytes_rt; // ek region end
wire [11:0] dk_hek_end = dk_ek_end + 12'd32; // H(ek) region end
reg [7:0] dbg_dk_r;
always @(posedge clk) begin
if (dbg_dk_idx_i < DK_BYTES[11:0])
if (dbg_dk_idx_i < dk_bytes_rt)
dbg_dk_r <= dkp_mem[dbg_dk_idx_i];
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 if (dbg_dk_idx_i < dk_ek_end)
dbg_dk_r <= ek_mem[dbg_dk_idx_i - dk_bytes_rt];
else if (dbg_dk_idx_i < dk_hek_end)
dbg_dk_r <= hek_r[(dbg_dk_idx_i - dk_ek_end)*8 +: 8];
else
dbg_dk_r <= z_i[(dbg_dk_idx_i - DK_HEK_END[11:0])*8 +: 8];
dbg_dk_r <= z_i[(dbg_dk_idx_i - dk_hek_end)*8 +: 8];
end
assign dbg_dk_o = dbg_dk_r;
@@ -119,21 +133,21 @@ module mlkem_top #(
reg [255:0] rho_r, sigma_r;
// 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 [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
wire [SAW-1:0] a_slot = a_i*k_r + a_j; // A_hat[i][j] slot = i*k + j
// C-generation bookkeeping: 2*K polys (s[0..K-1] then e[0..K-1])
reg [4:0] c_poly; // 0..2K
// 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));
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 < {2'b0, k_r}) ? (slot_s_rt + c_poly)
: (slot_e_rt + (c_poly - {2'b0, k_r}));
assign busy_o = (st != ST_IDLE);
assign done_o = (st == ST_DONE);
@@ -146,7 +160,7 @@ module mlkem_top #(
wire [511:0] sha3_hash;
wire sha3_vo;
reg sha3_ack; // consumer ready for hash
wire [511:0] g_data = {248'b0, 8'(K), d_i}; // data_i[263:256]=K, [255:0]=d
wire [511:0] g_data = {248'b0, 5'b0, k_r, d_i}; // data_i[263:256]=k, [255:0]=d
sha3_top u_sha3 (
.clk(clk), .rst_n(rst_n),
@@ -190,20 +204,19 @@ module mlkem_top #(
.mb_ready_o(h_mbready)
);
// 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
// SHA3-256 over ek (ek_bytes_rt bytes): rate=136. Padded length = h_nblk_rt*136.
// pad: byte ek_bytes_rt = 0x06 (domain + first pad bit), last byte |= 0x80.
// byte b (0..135) of block blk: global g = blk*136 + b.
// Reads runtime ek_bytes_rt / h_last_rt (stable during ST_H).
function [7:0] h_padbyte(input [3:0] blk, input [7:0] b);
integer g;
begin
g = blk*136 + b;
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;
if (g < ek_bytes_rt) h_padbyte = ek_mem[g];
else if (g == h_last_rt && g == ek_bytes_rt) h_padbyte = 8'h86; // 0x06|0x80
else if (g == ek_bytes_rt) h_padbyte = 8'h06;
else if (g == h_last_rt) h_padbyte = 8'h80;
else h_padbyte = 8'h00;
end
endfunction
@@ -215,10 +228,10 @@ module mlkem_top #(
wire snt_last;
reg snt_ack; // we accept coeffs
sample_ntt_sync #(.K(K)) u_snt (
sample_ntt_sync #(.K(KMAX)) u_snt (
.clk(clk), .rst_n(rst_n),
.rho_i(rho_r),
.k_i(3'(K)),
.k_i(k_r),
.i_idx(a_i[1:0]),
.j_idx(a_j[1:0]),
.valid_i(snt_valid),
@@ -241,7 +254,7 @@ module mlkem_top #(
.clk(clk), .rst_n(rst_n),
.seed_i(sigma_r),
.nonce_i(c_nonce),
.eta_i(2'(ETA1)),
.eta_i(eta1_rt),
.valid_i(cbd_valid),
.ready_o(cbd_ready),
.coeff_o(cbd_coeff),
@@ -260,7 +273,7 @@ module mlkem_top #(
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 [SAW-1:0] n_slot_addr = SLOT_S + n_slot; // s_hat then e_hat contiguous
wire [SAW-1:0] n_slot_addr = slot_s_rt + n_slot; // s_hat then e_hat contiguous
wire ntt_ready;
wire [11:0] ntt_coeff;
@@ -298,9 +311,9 @@ module mlkem_top #(
reg [9:0] e_rho; // 0..31 rho byte copy index (ek tail)
reg e_done; // serialization complete
// 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);
wire e_is_dk = (e_poly >= {1'b0, k_r});
wire [4:0] e_pidx = e_is_dk ? (e_poly - {1'b0, k_r}) : e_poly; // index within target
wire [SAW-1:0] e_slot = e_is_dk ? (slot_s_rt + e_pidx) : (slot_t_rt + 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}];
@@ -312,10 +325,10 @@ module mlkem_top #(
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 [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]
wire [SAW-1:0] m_aslot = m_i*k_r + m_j; // A_hat[i][j] slot = i*k + j
wire [SAW-1:0] m_sslot = slot_s_rt + m_j; // s_hat[j]
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]
reg pm_valid;
wire pm_ready;
@@ -347,10 +360,10 @@ module mlkem_top #(
case (st)
ST_IDLE: if (start_i) st_next = ST_G;
ST_G: if (sha3_vo) st_next = ST_A;
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_E;
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;
ST_N: if (n_slot >= {1'b0, k_r, 1'b0}) st_next = ST_M;
ST_M: if (m_i >= k_r) st_next = ST_E;
ST_E: if (e_done) st_next = ST_H;
ST_H: if (h_phase == 2'd3) st_next = ST_DONE;
ST_DONE: st_next = ST_IDLE;
@@ -361,6 +374,7 @@ module mlkem_top #(
always @(posedge clk or negedge rst_n) begin
if (!rst_n) begin
st <= ST_IDLE;
k_r <= 3'd0;
rho_r <= 256'd0;
sigma_r <= 256'd0;
sha3_valid <= 1'b0;
@@ -406,6 +420,7 @@ module mlkem_top #(
// Kick off G when entering ST_G
if (st == ST_IDLE && start_i) begin
k_r <= k_i; // capture runtime ML-KEM param
sha3_valid <= 1'b1;
sha3_ack <= 1'b1;
end
@@ -442,14 +457,14 @@ module mlkem_top #(
a_pair <= a_pair + 5'd1;
a_widx <= 8'd0;
a_busy <= 1'b0;
if (a_j + 3'd1 < K) begin
if (a_j + 3'd1 < k_r) 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 + 5'd1 < K*K) snt_valid <= 1'b1;
if (a_pair + 5'd1 < kk_rt) snt_valid <= 1'b1;
end else begin
a_widx <= a_widx + 8'd1;
end
@@ -478,7 +493,7 @@ module mlkem_top #(
c_poly <= c_poly + 3'd1;
c_widx <= 8'd0;
c_busy <= 1'b0;
if (c_poly + 3'd1 < 2*K) cbd_valid <= 1'b1;
if (c_poly + 3'd1 < {1'b0, k_r, 1'b0}) cbd_valid <= 1'b1;
end else begin
c_widx <= c_widx + 8'd1;
end
@@ -514,7 +529,7 @@ module mlkem_top #(
// Slot complete when ntt_core returns to DONE
if (ntt_done) begin
if (n_slot + 3'd1 < 2*K) begin
if (n_slot + 3'd1 < {1'b0, k_r, 1'b0}) begin
n_slot <= n_slot + 3'd1;
n_widx <= 8'd0;
n_pending <= 1'b1; // wait one cycle for core IDLE
@@ -561,13 +576,13 @@ module mlkem_top #(
polymem[m_tslot*256 + m_oidx] <= m_accq;
if (m_oidx == 8'd255) begin
// finished this (i,j) term; advance
if (m_j + 2'd1 < K) begin
if (m_j + 2'd1 < k_r) begin
m_j <= m_j + 2'd1;
m_pending <= 1'b1; // next term, same row
end else begin
m_j <= 2'd0;
m_i <= m_i + 2'd1; // next row (or == K -> DONE)
if (m_i + 2'd1 < K) m_pending <= 1'b1;
if (m_i + 2'd1 < k_r) m_pending <= 1'b1;
end
end else begin
m_oidx <= m_oidx + 8'd1;
@@ -594,7 +609,7 @@ 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 < 2*K) begin
if (e_poly < {1'b0, k_r, 1'b0}) 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;
@@ -613,7 +628,7 @@ module mlkem_top #(
end
end else begin
// 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];
ek_mem[dk_bytes_rt + 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
@@ -638,7 +653,7 @@ module mlkem_top #(
if (h_byte == 8'd135) begin
h_byte <= 8'd0;
h_mbvalid <= 1'b1;
h_mblast <= (h_blk == H_NBLK-1);
h_mblast <= (h_blk == h_nblk_rt - 4'd1);
h_phase <= 2'd1; // feed
end else begin
h_byte <= h_byte + 8'd1;