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; localparam DKB = 768*KP + 96;
reg clk=0, rst_n=0, start_i=0; reg clk=0, rst_n=0, start_i=0;
reg [2:0] k_i;
reg [255:0] d_i, z_i; reg [255:0] d_i, z_i;
wire busy_o, done_o; 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; 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; reg [11:0] dbg_dk_idx_i=0; wire [7:0] dbg_dk_o;
wire [255:0] dbg_rho_o, dbg_sigma_o; wire [255:0] dbg_rho_o, dbg_sigma_o;
mlkem_top #(.K(KP)) dut ( // KMAX defaults to 4 (worst-case sizing); KP selects the runtime k value.
.clk(clk), .rst_n(rst_n), .d_i(d_i), .z_i(z_i), .start_i(start_i), 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), .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_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), .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(ekfile, ek_gold);
$readmemh(dkfile, dk_gold); $readmemh(dkfile, dk_gold);
d_i = dmem[0]; z_i = zmem[0]; 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); rst_n=0; repeat(4) @(posedge clk); rst_n=1; @(posedge clk);
start_i=1; @(posedge clk); start_i=0; start_i=1; @(posedge clk); start_i=0;

163
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 // Streaming valid/ready interface. Given seeds d and z, computes the
// ML-KEM key pair per FIPS 203 Algorithm 16 (KeyGen_internal): // ML-KEM key pair per FIPS 203 Algorithm 16 (KeyGen_internal):
@@ -20,10 +23,11 @@
`include "sync_rtl/common/defines.vh" `include "sync_rtl/common/defines.vh"
module mlkem_top #( 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 clk,
input rst_n, 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] d_i, // KeyGen seed d (byte 0 in d_i[7:0])
input [255:0] z_i, // implicit-rejection seed z input [255:0] z_i, // implicit-rejection seed z
input start_i, // pulse to begin KeyGen input start_i, // pulse to begin KeyGen
@@ -49,22 +53,33 @@ module mlkem_top #(
); );
localparam Q = `Q; // 3329 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}. // Polynomial storage, sized for KMAX (worst case). Runtime k uses a
// Slot layout (each slot = 256 coeffs): // sub-range. Slot layout (each slot = 256 coeffs):
// A_hat[i][j] : slots 0 .. K*K-1 at index i*K + j // 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) // s_hat[i] : slots slot_s_rt .. +k-1
// e_hat[i] : slots SLOT_E .. +K-1 // e_hat[i] : slots slot_e_rt .. +k-1
// t_hat[i] : slots SLOT_T .. +K-1 // t_hat[i] : slots slot_t_rt .. +k-1
// NUM_SLOTS = K*K + 3*K (10 / 24 / 28 for K=2/3/4) // NUM_SLOTS = KMAX*KMAX + 3*KMAX = 28 for KMAX=4.
// ================================================================ // ================================================================
localparam SLOT_S = K*K; // s_hat base slot localparam NUM_SLOTS = KMAX*KMAX + 3*KMAX;
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)) localparam SAW = 5; // slot-address width (>=clog2(28))
reg [11:0] polymem [0:NUM_SLOTS*256-1]; 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]; always @(posedge clk) dbg_coeff_r <= polymem[dbg_slot_i*256 + dbg_idx_i];
assign dbg_coeff_o = dbg_coeff_r; assign dbg_coeff_o = dbg_coeff_r;
// ek and dk_pke byte memories (byteEncode12 output). // ek and dk_pke byte memories sized for KMAX.
// ek = 384*K + 32 bytes (== KAT pk), dk_pke = 384*K bytes (== KAT sk prefix) localparam EK_MAX = 384*KMAX + 32; // 1568
localparam EK_BYTES = 384*K + 32; // 800 / 1184 / 1568 localparam DK_MAX = 384*KMAX; // 1536
localparam DK_BYTES = 384*K; // 768 / 1152 / 1536 reg [7:0] ek_mem [0:EK_MAX-1];
reg [7:0] ek_mem [0:EK_BYTES-1]; reg [7:0] dkp_mem [0:DK_MAX-1];
reg [7:0] dkp_mem [0:DK_BYTES-1];
reg [7:0] dbg_byte_r; reg [7:0] dbg_byte_r;
always @(posedge clk) always @(posedge clk)
dbg_byte_r <= dbg_byte_sel_i ? dkp_mem[dbg_byte_idx_i] : ek_mem[dbg_byte_idx_i]; 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; assign dbg_byte_o = dbg_byte_r;
// full dk = dk_pke(DK_BYTES) || ek(EK_BYTES) || H(ek)(32) || z(32) // 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 wire [11:0] dk_ek_end = dk_bytes_rt + ek_bytes_rt; // ek region end
localparam DK_HEK_END = DK_EK_END + 32; // H(ek) region end wire [11:0] dk_hek_end = dk_ek_end + 12'd32; // H(ek) region end
reg [7:0] dbg_dk_r; reg [7:0] dbg_dk_r;
always @(posedge clk) begin 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]; dbg_dk_r <= dkp_mem[dbg_dk_idx_i];
else if (dbg_dk_idx_i < DK_EK_END[11:0]) else if (dbg_dk_idx_i < dk_ek_end)
dbg_dk_r <= ek_mem[dbg_dk_idx_i - DK_BYTES[11:0]]; dbg_dk_r <= ek_mem[dbg_dk_idx_i - dk_bytes_rt];
else if (dbg_dk_idx_i < DK_HEK_END[11:0]) else if (dbg_dk_idx_i < dk_hek_end)
dbg_dk_r <= hek_r[(dbg_dk_idx_i - DK_EK_END[11:0])*8 +: 8]; dbg_dk_r <= hek_r[(dbg_dk_idx_i - dk_ek_end)*8 +: 8];
else 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 end
assign dbg_dk_o = dbg_dk_r; assign dbg_dk_o = dbg_dk_r;
@@ -119,21 +133,21 @@ module mlkem_top #(
reg [255:0] rho_r, sigma_r; reg [255:0] rho_r, sigma_r;
// A-generation bookkeeping: explicit i/j counters (avoid runtime divide) // A-generation bookkeeping: explicit i/j counters (avoid runtime divide)
reg [2:0] a_i; // row 0..K-1 reg [2:0] a_i; // row 0..k-1
reg [2:0] a_j; // col 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 [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 [7:0] a_widx; // write index 0..255 within current poly
reg a_busy; // 1 once current pair's request accepted (gates collect) 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]) // C-generation bookkeeping: 2*k polys (s[0..k-1] then e[0..k-1])
reg [4:0] c_poly; // 0..2K reg [4:0] c_poly; // 0..2k
reg [7:0] c_widx; reg [7:0] c_widx;
reg c_busy; reg c_busy;
wire [7:0] c_nonce = {3'b0, c_poly}; // s:0..K-1 e:K..2K-1 == nonce 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] // 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) wire [SAW-1:0] c_slot = (c_poly < {2'b0, k_r}) ? (slot_s_rt + c_poly)
: (SLOT_E + (c_poly - K)); : (slot_e_rt + (c_poly - {2'b0, k_r}));
assign busy_o = (st != ST_IDLE); assign busy_o = (st != ST_IDLE);
assign done_o = (st == ST_DONE); assign done_o = (st == ST_DONE);
@@ -146,7 +160,7 @@ module mlkem_top #(
wire [511:0] sha3_hash; wire [511:0] sha3_hash;
wire sha3_vo; wire sha3_vo;
reg sha3_ack; // consumer ready for hash 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 ( sha3_top u_sha3 (
.clk(clk), .rst_n(rst_n), .clk(clk), .rst_n(rst_n),
@@ -190,19 +204,18 @@ module mlkem_top #(
.mb_ready_o(h_mbready) .mb_ready_o(h_mbready)
); );
// SHA3-256 over EK_BYTES-byte ek: rate=136. Padded length = H_NBLK*136. // SHA3-256 over ek (ek_bytes_rt bytes): rate=136. Padded length = h_nblk_rt*136.
// pad: byte EK_BYTES = 0x06 (domain + first pad bit), last byte |= 0x80. // pad: byte ek_bytes_rt = 0x06 (domain + first pad bit), last byte |= 0x80.
localparam H_NBLK = (EK_BYTES + 136) / 136; // ceil((EK_BYTES+1)/136): 6/9/12 // byte b (0..135) of block blk: global g = blk*136 + b.
localparam H_LAST = H_NBLK*136 - 1; // index of final padded byte // Reads runtime ek_bytes_rt / h_last_rt (stable during ST_H).
// 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); function [7:0] h_padbyte(input [3:0] blk, input [7:0] b);
integer g; integer g;
begin begin
g = blk*136 + b; g = blk*136 + b;
if (g < EK_BYTES) h_padbyte = ek_mem[g]; if (g < ek_bytes_rt) h_padbyte = ek_mem[g];
else if (g == H_LAST && g == EK_BYTES) h_padbyte = 8'h86; // 0x06|0x80 else if (g == h_last_rt && g == ek_bytes_rt) h_padbyte = 8'h86; // 0x06|0x80
else if (g == EK_BYTES) h_padbyte = 8'h06; else if (g == ek_bytes_rt) h_padbyte = 8'h06;
else if (g == H_LAST) h_padbyte = 8'h80; else if (g == h_last_rt) h_padbyte = 8'h80;
else h_padbyte = 8'h00; else h_padbyte = 8'h00;
end end
endfunction endfunction
@@ -215,10 +228,10 @@ module mlkem_top #(
wire snt_last; wire snt_last;
reg snt_ack; // we accept coeffs 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), .clk(clk), .rst_n(rst_n),
.rho_i(rho_r), .rho_i(rho_r),
.k_i(3'(K)), .k_i(k_r),
.i_idx(a_i[1:0]), .i_idx(a_i[1:0]),
.j_idx(a_j[1:0]), .j_idx(a_j[1:0]),
.valid_i(snt_valid), .valid_i(snt_valid),
@@ -241,7 +254,7 @@ module mlkem_top #(
.clk(clk), .rst_n(rst_n), .clk(clk), .rst_n(rst_n),
.seed_i(sigma_r), .seed_i(sigma_r),
.nonce_i(c_nonce), .nonce_i(c_nonce),
.eta_i(2'(ETA1)), .eta_i(eta1_rt),
.valid_i(cbd_valid), .valid_i(cbd_valid),
.ready_o(cbd_ready), .ready_o(cbd_ready),
.coeff_o(cbd_coeff), .coeff_o(cbd_coeff),
@@ -260,7 +273,7 @@ module mlkem_top #(
reg [7:0] n_widx; // output write index 0..255 reg [7:0] n_widx; // output write index 0..255
reg n_valid; // feeding coeffs to ntt_core reg n_valid; // feeding coeffs to ntt_core
reg n_pending; // waiting for ntt_core IDLE to start next slot 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 ntt_ready;
wire [11:0] ntt_coeff; 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 [9:0] e_rho; // 0..31 rho byte copy index (ek tail)
reg e_done; // serialization complete reg e_done; // serialization complete
// source poly slot: t_hat[e_poly] for ek half, s_hat[e_poly-K] for dk half // 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 e_is_dk = (e_poly >= {1'b0, k_r});
wire [4:0] e_pidx = e_is_dk ? (e_poly - K) : e_poly; // index within target 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 + e_pidx) : (SLOT_T + e_pidx); 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 // two coeffs of the current pair
wire [11:0] e_c0 = polymem[e_slot*256 + {e_pair, 1'b0}]; 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}]; 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_base = e_pidx * 12'd384;
wire [11:0] e_boff = e_base + {e_pair, 1'b0} + {2'b0, e_pair}; // pair*3 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_aslot = m_i*k_r + 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_sslot = slot_s_rt + m_j; // s_hat[j]
wire [SAW-1:0] m_eslot = SLOT_E + m_i; // e_hat[i] wire [SAW-1:0] m_eslot = slot_e_rt + m_i; // e_hat[i]
wire [SAW-1:0] m_tslot = SLOT_T + m_i; // t_hat[i] wire [SAW-1:0] m_tslot = slot_t_rt + m_i; // t_hat[i]
reg pm_valid; reg pm_valid;
wire pm_ready; wire pm_ready;
@@ -347,10 +360,10 @@ module mlkem_top #(
case (st) case (st)
ST_IDLE: if (start_i) st_next = ST_G; ST_IDLE: if (start_i) st_next = ST_G;
ST_G: if (sha3_vo) st_next = ST_A; ST_G: if (sha3_vo) st_next = ST_A;
ST_A: if (a_pair >= K*K) st_next = ST_C; ST_A: if (a_pair >= kk_rt) st_next = ST_C;
ST_C: if (c_poly >= 2*K) st_next = ST_N; ST_C: if (c_poly >= {1'b0, k_r, 1'b0}) st_next = ST_N;
ST_N: if (n_slot >= 2*K) st_next = ST_M; ST_N: if (n_slot >= {1'b0, k_r, 1'b0}) st_next = ST_M;
ST_M: if (m_i >= K) st_next = ST_E; ST_M: if (m_i >= k_r) st_next = ST_E;
ST_E: if (e_done) st_next = ST_H; ST_E: if (e_done) st_next = ST_H;
ST_H: if (h_phase == 2'd3) st_next = ST_DONE; ST_H: if (h_phase == 2'd3) st_next = ST_DONE;
ST_DONE: st_next = ST_IDLE; ST_DONE: st_next = ST_IDLE;
@@ -361,6 +374,7 @@ module mlkem_top #(
always @(posedge clk or negedge rst_n) begin always @(posedge clk or negedge rst_n) begin
if (!rst_n) begin if (!rst_n) begin
st <= ST_IDLE; st <= ST_IDLE;
k_r <= 3'd0;
rho_r <= 256'd0; rho_r <= 256'd0;
sigma_r <= 256'd0; sigma_r <= 256'd0;
sha3_valid <= 1'b0; sha3_valid <= 1'b0;
@@ -406,6 +420,7 @@ module mlkem_top #(
// Kick off G when entering ST_G // Kick off G when entering ST_G
if (st == ST_IDLE && start_i) begin if (st == ST_IDLE && start_i) begin
k_r <= k_i; // capture runtime ML-KEM param
sha3_valid <= 1'b1; sha3_valid <= 1'b1;
sha3_ack <= 1'b1; sha3_ack <= 1'b1;
end end
@@ -442,14 +457,14 @@ module mlkem_top #(
a_pair <= a_pair + 5'd1; a_pair <= a_pair + 5'd1;
a_widx <= 8'd0; a_widx <= 8'd0;
a_busy <= 1'b0; 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; a_j <= a_j + 3'd1;
end else begin end else begin
a_j <= 3'd0; a_j <= 3'd0;
a_i <= a_i + 3'd1; a_i <= a_i + 3'd1;
end end
// start next SampleNTT if more pairs remain // 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 end else begin
a_widx <= a_widx + 8'd1; a_widx <= a_widx + 8'd1;
end end
@@ -478,7 +493,7 @@ module mlkem_top #(
c_poly <= c_poly + 3'd1; c_poly <= c_poly + 3'd1;
c_widx <= 8'd0; c_widx <= 8'd0;
c_busy <= 1'b0; 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 end else begin
c_widx <= c_widx + 8'd1; c_widx <= c_widx + 8'd1;
end end
@@ -514,7 +529,7 @@ module mlkem_top #(
// Slot complete when ntt_core returns to DONE // Slot complete when ntt_core returns to DONE
if (ntt_done) begin 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_slot <= n_slot + 3'd1;
n_widx <= 8'd0; n_widx <= 8'd0;
n_pending <= 1'b1; // wait one cycle for core IDLE 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; polymem[m_tslot*256 + m_oidx] <= m_accq;
if (m_oidx == 8'd255) begin if (m_oidx == 8'd255) begin
// finished this (i,j) term; advance // 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_j <= m_j + 2'd1;
m_pending <= 1'b1; // next term, same row m_pending <= 1'b1; // next term, same row
end else begin end else begin
m_j <= 2'd0; m_j <= 2'd0;
m_i <= m_i + 2'd1; // next row (or == K -> DONE) 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
end else begin end else begin
m_oidx <= m_oidx + 8'd1; 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 ---- // ---- ST_E: byteEncode12 t_hat -> ek_mem, s_hat -> dkp_mem, ek tail = rho ----
if (st == ST_E && !e_done) begin 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 // pack current coeff-pair (3 bytes): [0,K)=ek, [K,2K)=dk_pke
if (!e_is_dk) begin if (!e_is_dk) begin
ek_mem[e_boff] <= e_b0; ek_mem[e_boff] <= e_b0;
@@ -613,7 +628,7 @@ module mlkem_top #(
end end
end else begin end else begin
// rho copy: ek_mem[384*K + r] = rho byte r (r = 0..31) // 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; if (e_rho == 10'd31) e_done <= 1'b1;
else e_rho <= e_rho + 10'd1; else e_rho <= e_rho + 10'd1;
end end
@@ -638,7 +653,7 @@ module mlkem_top #(
if (h_byte == 8'd135) begin if (h_byte == 8'd135) begin
h_byte <= 8'd0; h_byte <= 8'd0;
h_mbvalid <= 1'b1; 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 h_phase <= 2'd1; // feed
end else begin end else begin
h_byte <= h_byte + 8'd1; h_byte <= h_byte + 8'd1;