From 3a539937546f089a6d0f0860c60a8c340acb1e10 Mon Sep 17 00:00:00 2001 From: FallenSigh Date: Sun, 28 Jun 2026 03:24:58 +0800 Subject: [PATCH] refactor(kg): make ML-KEM K a runtime input k_i instead of a parameter 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). --- sync_rtl/top/TB/tb_mlkem_kg_katK_xsim.v | 7 +- sync_rtl/top/mlkem_top.v | 165 +++++++++++++----------- 2 files changed, 95 insertions(+), 77 deletions(-) diff --git a/sync_rtl/top/TB/tb_mlkem_kg_katK_xsim.v b/sync_rtl/top/TB/tb_mlkem_kg_katK_xsim.v index 6c9a5cc..dad77b5 100644 --- a/sync_rtl/top/TB/tb_mlkem_kg_katK_xsim.v +++ b/sync_rtl/top/TB/tb_mlkem_kg_katK_xsim.v @@ -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; diff --git a/sync_rtl/top/mlkem_top.v b/sync_rtl/top/mlkem_top.v index bdb078a..6d10944 100644 --- a/sync_rtl/top/mlkem_top.v +++ b/sync_rtl/top/mlkem_top.v @@ -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;