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

76
sync_rtl/top/TB/tb_mlkem_kg_katK_xsim.v Normal file
View File

@@ -0,0 +1,76 @@
// tb_mlkem_kg_katK_xsim.v - ML-KEM KeyGen vs NIST KAT, parametric K (KP) + CASE.
// xelab -generic_top KP=2|3|4 ; xsim -testplusarg CASE=n
// KP=2 -> k2 vectors (ML-KEM-512), KP=3 -> k3 (768), KP=4 -> k4 (1024).
// ek = 384*KP+32 bytes (==KAT pk), dk = 768*KP+96 bytes (==KAT sk).
`timescale 1ns/1ps
module tb_mlkem_kg_katK_xsim;
parameter KP = 2;
localparam EKB = 384*KP + 32;
localparam DKB = 768*KP + 96;
reg clk=0, rst_n=0, start_i=0;
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;
reg dbg_byte_sel_i=0; reg [10:0] dbg_byte_idx_i=0; wire [7:0] dbg_byte_o;
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),
.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),
.dbg_dk_idx_i(dbg_dk_idx_i), .dbg_dk_o(dbg_dk_o),
.dbg_rho_o(dbg_rho_o), .dbg_sigma_o(dbg_sigma_o)
);
always #5 clk = ~clk;
reg [255:0] dmem [0:0];
reg [255:0] zmem [0:0];
reg [7:0] ek_gold [0:EKB-1];
reg [7:0] dk_gold [0:DKB-1];
integer c, i, errors, casenum;
reg [8*80-1:0] tag, dfile, zfile, ekfile, dkfile;
initial begin
if (!$value$plusargs("CASE=%d", casenum)) casenum = 0;
$sformat(tag, "k%0d", KP);
$sformat(dfile, "sync_rtl/top/TB/vectors/kat_%0s_c%0d_d.hex", tag, casenum);
$sformat(zfile, "sync_rtl/top/TB/vectors/kat_%0s_c%0d_z.hex", tag, casenum);
$sformat(ekfile, "sync_rtl/top/TB/vectors/kat_%0s_c%0d_ek.hex", tag, casenum);
$sformat(dkfile, "sync_rtl/top/TB/vectors/kat_%0s_c%0d_dk.hex", tag, casenum);
$readmemh(dfile, dmem);
$readmemh(zfile, zmem);
$readmemh(ekfile, ek_gold);
$readmemh(dkfile, dk_gold);
d_i = dmem[0]; z_i = zmem[0];
rst_n=0; repeat(4) @(posedge clk); rst_n=1; @(posedge clk);
start_i=1; @(posedge clk); start_i=0;
c=0; while(!done_o && c<2000000) begin @(posedge clk); c=c+1; end
if(!done_o) begin $display("FAIL K=%0d case %0d: timeout", KP, casenum); $finish; end
$display("=== ML-KEM K=%0d KAT case %0d: KeyGen done in %0d cyc ===", KP, casenum, c);
errors = 0;
dbg_byte_sel_i = 1'b0;
for (i = 0; i < EKB; i = i + 1) begin
dbg_byte_idx_i = i[10:0]; @(posedge clk); @(posedge clk);
if (dbg_byte_o !== ek_gold[i]) begin
if (errors < 8) $display(" EK[%0d] got=%02x exp=%02x", i, dbg_byte_o, ek_gold[i]);
errors = errors + 1;
end
end
for (i = 0; i < DKB; i = i + 1) begin
dbg_dk_idx_i = i[11:0]; @(posedge clk); @(posedge clk);
if (dbg_dk_o !== dk_gold[i]) begin
if (errors < 8) $display(" DK[%0d] got=%02x exp=%02x", i, dbg_dk_o, dk_gold[i]);
errors = errors + 1;
end
end
if (errors == 0) $display("K=%0d CASE %0d PASS: ek (%0dB)==pk, dk (%0dB)==sk", KP, casenum, EKB, DKB);
else $display("K=%0d CASE %0d FAIL: %0d mismatches", KP, casenum, errors);
$finish;
end
initial begin #120000000; $display("FAIL: global timeout"); $finish; end
endmodule

1
sync_rtl/top/TB/vectors/kat_k3_c0_d.hex Normal file
View File

@@ -0,0 +1 @@
2426f1941779574d3f1b163bd57f7e173e229e630ec7f7073bdf365137c4bb6d

2400
sync_rtl/top/TB/vectors/kat_k3_c0_dk.hex Normal file
View File

File diff suppressed because it is too large Load Diff

1184
sync_rtl/top/TB/vectors/kat_k3_c0_ek.hex Normal file
View File

File diff suppressed because it is too large Load Diff

1
sync_rtl/top/TB/vectors/kat_k3_c0_z.hex Normal file
View File

@@ -0,0 +1 @@
687acf9406694974d383032f7579378f449c75d0560af56cf921ec48404896f6

1
sync_rtl/top/TB/vectors/kat_k3_c1_d.hex Normal file
View File

@@ -0,0 +1 @@
9b68d9d2390fa18739b239a58649393a28f97f6b164ec5e4334f4df864fc9cd6

2400
sync_rtl/top/TB/vectors/kat_k3_c1_dk.hex Normal file
View File

File diff suppressed because it is too large Load Diff

1184
sync_rtl/top/TB/vectors/kat_k3_c1_ek.hex Normal file
View File

File diff suppressed because it is too large Load Diff

1
sync_rtl/top/TB/vectors/kat_k3_c1_z.hex Normal file
View File

@@ -0,0 +1 @@
a02e20d538b412ff074d121a80b0583e0b54e85b267da0789c5ca565342ee66d

1
sync_rtl/top/TB/vectors/kat_k3_c2_d.hex Normal file
View File

@@ -0,0 +1 @@
905eb1f02ade542733c35132f5a998d3ec80d20ec8ed235bf228081157034763

2400
sync_rtl/top/TB/vectors/kat_k3_c2_dk.hex Normal file
View File

File diff suppressed because it is too large Load Diff

1184
sync_rtl/top/TB/vectors/kat_k3_c2_ek.hex Normal file
View File

File diff suppressed because it is too large Load Diff

1
sync_rtl/top/TB/vectors/kat_k3_c2_z.hex Normal file
View File

@@ -0,0 +1 @@
d026bb15f52e8df87779ff953ccf42a9d5401111c402c448d77cb291bbe6aa1e

1
sync_rtl/top/TB/vectors/kat_k4_c0_d.hex Normal file
View File

@@ -0,0 +1 @@
2426f1941779574d3f1b163bd57f7e173e229e630ec7f7073bdf365137c4bb6d

3168
sync_rtl/top/TB/vectors/kat_k4_c0_dk.hex Normal file
View File

File diff suppressed because it is too large Load Diff

1568
sync_rtl/top/TB/vectors/kat_k4_c0_ek.hex Normal file
View File

File diff suppressed because it is too large Load Diff

1
sync_rtl/top/TB/vectors/kat_k4_c0_z.hex Normal file
View File

@@ -0,0 +1 @@
687acf9406694974d383032f7579378f449c75d0560af56cf921ec48404896f6

1
sync_rtl/top/TB/vectors/kat_k4_c1_d.hex Normal file
View File

@@ -0,0 +1 @@
9b68d9d2390fa18739b239a58649393a28f97f6b164ec5e4334f4df864fc9cd6

3168
sync_rtl/top/TB/vectors/kat_k4_c1_dk.hex Normal file
View File

File diff suppressed because it is too large Load Diff

1568
sync_rtl/top/TB/vectors/kat_k4_c1_ek.hex Normal file
View File

File diff suppressed because it is too large Load Diff

1
sync_rtl/top/TB/vectors/kat_k4_c1_z.hex Normal file
View File

@@ -0,0 +1 @@
a02e20d538b412ff074d121a80b0583e0b54e85b267da0789c5ca565342ee66d

1
sync_rtl/top/TB/vectors/kat_k4_c2_d.hex Normal file
View File

@@ -0,0 +1 @@
905eb1f02ade542733c35132f5a998d3ec80d20ec8ed235bf228081157034763

3168
sync_rtl/top/TB/vectors/kat_k4_c2_dk.hex Normal file
View File

File diff suppressed because it is too large Load Diff

1568
sync_rtl/top/TB/vectors/kat_k4_c2_ek.hex Normal file
View File

File diff suppressed because it is too large Load Diff

1
sync_rtl/top/TB/vectors/kat_k4_c2_z.hex Normal file
View File

@@ -0,0 +1 @@
d026bb15f52e8df87779ff953ccf42a9d5401111c402c448d77cb291bbe6aa1e

48
sync_rtl/top/TB/xsim_run.tcl
View File

@@ -1,16 +1,12 @@
# xsim_run.tcl - Vivado XSIM for mlkem_top ML-KEM-512 KeyGen.
# xsim_run.tcl - Vivado XSIM for mlkem_top ML-KEM KeyGen, all K in {2,3,4}.
#
# Compiles the full KeyGen datapath + all leaf modules, then runs the
# end-to-end NIST KAT testbench for count=0..4 (each verifies ek==KAT.pk
# 800B and dk==KAT.sk 1632B byte-exact).
# Compiles the full KeyGen datapath + leaf modules, then runs the parametric
# NIST KAT testbench for ML-KEM-512 (K=2), 768 (K=3), 1024 (K=4). Each case
# verifies ek==KAT.pk and dk==KAT.sk byte-exact.
#
# Run from project root:
# ./run_tb.sh top
# (or: vivado-style step-by-step with the commands below)
# ================================================================
# Step 1: Compile RTL sources (leaf modules + top)
# ================================================================
# ---- Step 1: compile RTL ----
xvlog -sv --relax -i . sync_rtl/sha3/keccak_round.v
xvlog -sv --relax -i . sync_rtl/sha3/keccak_core.v
xvlog -sv --relax -i . sync_rtl/sha3/sha3_top.v
@@ -25,21 +21,23 @@ xvlog -sv --relax -i . sync_rtl/poly_mul/poly_mul_zeta_rom.v
xvlog -sv --relax -i . sync_rtl/poly_mul/poly_mul_sync.v
xvlog -sv --relax -i . sync_rtl/top/mlkem_top.v
# ================================================================
# Step 2: Compile KAT testbench (parameterized by +CASE=n)
# ================================================================
xvlog -sv --relax sync_rtl/top/TB/tb_mlkem_kg_katN_xsim.v
# ---- Step 2: compile parametric KAT testbench ----
xvlog -sv --relax sync_rtl/top/TB/tb_mlkem_kg_katK_xsim.v
# ================================================================
# Step 3: Elaborate snapshot
# ================================================================
xelab tb_mlkem_kg_katN_xsim -s mlkem_kg_kat --timescale 1ns/1ps
# ---- Step 3: elaborate one snapshot per K ----
xelab tb_mlkem_kg_katK_xsim -generic_top KP=2 -s mlkem_kg_k2 --timescale 1ns/1ps
xelab tb_mlkem_kg_katK_xsim -generic_top KP=3 -s mlkem_kg_k3 --timescale 1ns/1ps
xelab tb_mlkem_kg_katK_xsim -generic_top KP=4 -s mlkem_kg_k4 --timescale 1ns/1ps
# ================================================================
# Step 4: Run all 5 KAT cases (count=0..4)
# ================================================================
xsim mlkem_kg_kat -R -testplusarg CASE=0
xsim mlkem_kg_kat -R -testplusarg CASE=1
xsim mlkem_kg_kat -R -testplusarg CASE=2
xsim mlkem_kg_kat -R -testplusarg CASE=3
xsim mlkem_kg_kat -R -testplusarg CASE=4
# ---- Step 4: run KAT cases (K=2: 0..4, K=3/4: 0..2) ----
xsim mlkem_kg_k2 -R -testplusarg CASE=0
xsim mlkem_kg_k2 -R -testplusarg CASE=1
xsim mlkem_kg_k2 -R -testplusarg CASE=2
xsim mlkem_kg_k2 -R -testplusarg CASE=3
xsim mlkem_kg_k2 -R -testplusarg CASE=4
xsim mlkem_kg_k3 -R -testplusarg CASE=0
xsim mlkem_kg_k3 -R -testplusarg CASE=1
xsim mlkem_kg_k3 -R -testplusarg CASE=2
xsim mlkem_kg_k4 -R -testplusarg CASE=0
xsim mlkem_kg_k4 -R -testplusarg CASE=1
xsim mlkem_kg_k4 -R -testplusarg CASE=2

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;