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

Pipeline ML-KEM datapath bottlenecks

Browse Source
This commit is contained in:
FallenSigh
2026-07-08 00:23:46 +08:00
parent 8c3f4317f5
commit 372a90e601
17 changed files with 776 additions and 176 deletions
Download Patch File Download Diff File Expand all files Collapse all files

73
sync_rtl/top/TB/tb_mlkem_dec_katK_xsim.v
View File

@@ -13,6 +13,8 @@ module tb_mlkem_dec_katK_xsim;
localparam EKB = 384*KP + 32; // ek_pke bytes within dk
localparam DKPB = 384*KP; // dk_pke bytes
localparam CTB = (KP==4) ? 1568 : (32*(10*KP+4)); // ct bytes: 768/1088/1568
localparam integer CLK_PERIOD_NS = 100; // 10 MHz
localparam integer CLK_HALF_NS = CLK_PERIOD_NS / 2;
reg clk=0, rst_n=0, start_i=0;
reg [2:0] k_i;
@@ -47,7 +49,7 @@ module tb_mlkem_dec_katK_xsim;
.dbg_mprime_o(dbg_mprime_o), .dbg_kbar_o(dbg_kbar_o),
.dbg_decz_o(dbg_decz_o), .dbg_dech_o(dbg_dech_o)
);
always #5 clk = ~clk;
always #(CLK_HALF_NS) clk = ~clk;
reg [7:0] dk_b [0:DKB-1];
reg [7:0] ct_b [0:CTB-1];
@@ -55,8 +57,72 @@ module tb_mlkem_dec_katK_xsim;
reg [7:0] ss_b [0:31];
reg [7:0] ssn_b [0:31];
integer c, i, j, errors, casenum;
integer state_cyc [0:31];
integer si;
reg [8*80-1:0] tag, dkfile, ctfile, ssfile, ctnfile, ssnfile;
function [8*16-1:0] state_name;
input [4:0] s;
begin
case (s)
5'd0: state_name = "IDLE";
5'd1: state_name = "G";
5'd2: state_name = "A";
5'd3: state_name = "C";
5'd4: state_name = "N";
5'd5: state_name = "M";
5'd6: state_name = "E";
5'd7: state_name = "H";
5'd8: state_name = "ENC_LOAD";
5'd9: state_name = "ENC_H";
5'd10: state_name = "ENC_G";
5'd11: state_name = "ENC_A";
5'd12: state_name = "ENC_TDEC";
5'd13: state_name = "ENC_C";
5'd14: state_name = "ENC_N";
5'd15: state_name = "ENC_U";
5'd16: state_name = "ENC_C1";
5'd17: state_name = "ENC_V";
5'd18: state_name = "ENC_C2";
5'd19: state_name = "ENC_E2MV";
5'd20: state_name = "DEC_LOAD";
5'd21: state_name = "DEC_DECOMP";
5'd22: state_name = "DEC_SDEC";
5'd23: state_name = "DEC_NTT";
5'd24: state_name = "DEC_W";
5'd25: state_name = "DEC_MENC";
5'd26: state_name = "DEC_G";
5'd27: state_name = "DEC_J";
5'd28: state_name = "DEC_CMP";
5'd31: state_name = "DONE";
default: state_name = "UNKNOWN";
endcase
end
endfunction
task reset_timing;
begin
for (si = 0; si < 32; si = si + 1) state_cyc[si] = 0;
end
endtask
task print_timing;
input [8*16-1:0] op_name;
input integer cycles;
integer runtime_ns;
begin
runtime_ns = cycles * CLK_PERIOD_NS;
$display("TIME K=%0d CASE=%0d OP=%0s cycles=%0d clk=10MHz period=%0dns runtime=%0d ns (%0d.%03d us, %0d.%03d ms)",
KP, casenum, op_name, cycles, CLK_PERIOD_NS, runtime_ns,
runtime_ns/1000, runtime_ns%1000,
runtime_ns/1000000, (runtime_ns%1000000)/1000);
$display("TIME_BREAKDOWN K=%0d CASE=%0d OP=%0s", KP, casenum, op_name);
for (si = 0; si < 32; si = si + 1)
if (state_cyc[si] != 0)
$display(" STATE %-12s cycles=%0d time_ns=%0d", state_name(si[4:0]), state_cyc[si], state_cyc[si]*CLK_PERIOD_NS);
end
endtask
initial begin
if (!$value$plusargs("CASE=%d", casenum)) casenum = 0;
$sformat(tag, "k%0d", KP);
@@ -130,9 +196,12 @@ module tb_mlkem_dec_katK_xsim;
end
c_in_we = 1'b0; @(posedge clk);
start_i=1; @(posedge clk); start_i=0;
c=0; while(!done_o && c<2000000) begin @(posedge clk); c=c+1; end
#1;
reset_timing;
c=0; while(!done_o && c<2000000) begin state_cyc[dut.st] = state_cyc[dut.st] + 1; @(posedge clk); c=c+1; end
if(!done_o) begin $display("FAIL K=%0d case %0d: timeout (sel=%0d)", KP, casenum, sel); $finish; end
$display("=== Decaps run (sel=%0d) done in %0d cyc ===", sel, c);
print_timing(sel ? "DecapsReject" : "DecapsAccept", c);
// settle so dbg taps reflect the finished run
repeat(2) @(posedge clk);
end

73
sync_rtl/top/TB/tb_mlkem_enc_katK_xsim.v
View File

@@ -12,6 +12,8 @@ module tb_mlkem_enc_katK_xsim;
localparam CTB = (KP==4) ? 1568 : (32*(10*KP+4)); // ct bytes: K2 768,K3 1088,K4 1568
localparam DU = (KP==4) ? 11 : 10; // compression du
localparam C1B = 32*DU*KP; // c1 byte count: K2 640,K3 960,K4 1408
localparam integer CLK_PERIOD_NS = 100; // 10 MHz
localparam integer CLK_HALF_NS = CLK_PERIOD_NS / 2;
reg clk=0, rst_n=0, start_i=0;
reg [2:0] k_i;
@@ -41,15 +43,79 @@ module tb_mlkem_enc_katK_xsim;
.dbg_r_o(dbg_r_o), .dbg_hek_o(dbg_hek_o),
.dbg_mprime_o(), .dbg_kbar_o(), .dbg_decz_o(), .dbg_dech_o()
);
always #5 clk = ~clk;
always #(CLK_HALF_NS) clk = ~clk;
reg [7:0] ek_b [0:EKB-1];
reg [7:0] m_b [0:31];
reg [7:0] ss_b [0:31];
reg [7:0] ct_b [0:CTB-1];
integer c, i, errors, casenum, j;
integer state_cyc [0:31];
integer si;
reg [8*80-1:0] tag, ekfile, mfile, ssfile, ctfile;
function [8*16-1:0] state_name;
input [4:0] s;
begin
case (s)
5'd0: state_name = "IDLE";
5'd1: state_name = "G";
5'd2: state_name = "A";
5'd3: state_name = "C";
5'd4: state_name = "N";
5'd5: state_name = "M";
5'd6: state_name = "E";
5'd7: state_name = "H";
5'd8: state_name = "ENC_LOAD";
5'd9: state_name = "ENC_H";
5'd10: state_name = "ENC_G";
5'd11: state_name = "ENC_A";
5'd12: state_name = "ENC_TDEC";
5'd13: state_name = "ENC_C";
5'd14: state_name = "ENC_N";
5'd15: state_name = "ENC_U";
5'd16: state_name = "ENC_C1";
5'd17: state_name = "ENC_V";
5'd18: state_name = "ENC_C2";
5'd19: state_name = "ENC_E2MV";
5'd20: state_name = "DEC_LOAD";
5'd21: state_name = "DEC_DECOMP";
5'd22: state_name = "DEC_SDEC";
5'd23: state_name = "DEC_NTT";
5'd24: state_name = "DEC_W";
5'd25: state_name = "DEC_MENC";
5'd26: state_name = "DEC_G";
5'd27: state_name = "DEC_J";
5'd28: state_name = "DEC_CMP";
5'd31: state_name = "DONE";
default: state_name = "UNKNOWN";
endcase
end
endfunction
task reset_timing;
begin
for (si = 0; si < 32; si = si + 1) state_cyc[si] = 0;
end
endtask
task print_timing;
input [8*16-1:0] op_name;
input integer cycles;
integer runtime_ns;
begin
runtime_ns = cycles * CLK_PERIOD_NS;
$display("TIME K=%0d CASE=%0d OP=%0s cycles=%0d clk=10MHz period=%0dns runtime=%0d ns (%0d.%03d us, %0d.%03d ms)",
KP, casenum, op_name, cycles, CLK_PERIOD_NS, runtime_ns,
runtime_ns/1000, runtime_ns%1000,
runtime_ns/1000000, (runtime_ns%1000000)/1000);
$display("TIME_BREAKDOWN K=%0d CASE=%0d OP=%0s", KP, casenum, op_name);
for (si = 0; si < 32; si = si + 1)
if (state_cyc[si] != 0)
$display(" STATE %-12s cycles=%0d time_ns=%0d", state_name(si[4:0]), state_cyc[si], state_cyc[si]*CLK_PERIOD_NS);
end
endtask
initial begin
if (!$value$plusargs("CASE=%d", casenum)) casenum = 0;
$sformat(tag, "k%0d", KP);
@@ -81,9 +147,12 @@ module tb_mlkem_enc_katK_xsim;
// ---- run Encaps ----
start_i=1; @(posedge clk); start_i=0;
c=0; while(!done_o && c<2000000) begin @(posedge clk); c=c+1; end
#1;
reset_timing;
c=0; while(!done_o && c<2000000) begin state_cyc[dut.st] = state_cyc[dut.st] + 1; @(posedge clk); c=c+1; end
if(!done_o) begin $display("FAIL K=%0d case %0d: timeout", KP, casenum); $finish; end
$display("=== Encaps E0 done in %0d cyc ===", c);
print_timing("Encaps", c);
$write(" H(ek) = "); for (j=0;j<32;j=j+1) $write("%02x", dbg_hek_o[8*j +: 8]); $write("\n");
$write(" r = "); for (j=0;j<32;j=j+1) $write("%02x", dbg_r_o[8*j +: 8]); $write("\n");

20
sync_rtl/top/TB/tb_mlkem_hello_world_xsim.v
View File

@@ -1,8 +1,8 @@
// tb_mlkem_hello_world_xsim.v - Hardware run of the ML-KEM hello_world demo.
//
// Mirrors the Rust example end-to-end on the mlkem_top DUT (ML-KEM-512, K=2):
// 1. Alice: KeyGen(d=0x42..., z=0x77...) -> (ek 800B, dk 1632B)
// 2. Bob: Encaps(ek, m=0xDE...) -> (shared_key, kem_ct 768B)
// Mirrors the Rust example end-to-end on the mlkem_top DUT (default ML-KEM-512, K=2):
// 1. Alice: KeyGen(d=0x42..., z=0x77...) -> (ek, dk)
// 2. Bob: Encaps(ek, m=0xDE...) -> (shared_key, kem_ct)
// 3. Bob: XOR-encrypt "hello world" with key -> encrypted
// 4. Alice: Decaps(dk, kem_ct) -> recovered_key
// 5. Alice: XOR-decrypt encrypted with key -> "hello world"
@@ -17,10 +17,12 @@
// xelab tb_mlkem_hello_world_xsim ; xsim
`timescale 1ns/1ps
module tb_mlkem_hello_world_xsim;
localparam KP = 2; // ML-KEM-512
localparam EKB = 384*KP + 32; // 800
localparam DKB = 768*KP + 96; // 1632
localparam CTB = 32*(10*KP + 4); // 768
parameter KP = 2; // 2/3/4 = ML-KEM-512/768/1024
localparam DU = (KP == 4) ? 11 : 10;
localparam DV = (KP == 4) ? 5 : 4;
localparam EKB = 384*KP + 32;
localparam DKB = 768*KP + 96;
localparam CTB = 32*(DU*KP + DV);
localparam MLEN = 11; // "hello world"
reg clk=0, rst_n=0, start_i=0;
@@ -93,7 +95,7 @@ module tb_mlkem_hello_world_xsim;
msg_b[0]="h"; msg_b[1]="e"; msg_b[2]="l"; msg_b[3]="l"; msg_b[4]="o";
msg_b[5]=" "; msg_b[6]="w"; msg_b[7]="o"; msg_b[8]="r"; msg_b[9]="l"; msg_b[10]="d";
$display("=== ML-KEM hello_world (hardware, ML-KEM-512) ===");
$display("=== ML-KEM hello_world (hardware, K=%0d) ===", KP);
$write("Original: \"");
for (i = 0; i < MLEN; i = i + 1) $write("%c", msg_b[i]);
$display("\"\n");
@@ -189,7 +191,7 @@ module tb_mlkem_hello_world_xsim;
if (dec_b[i] !== msg_b[i]) errors = errors + 1;
if (errors == 0)
$display("Success: keys match, message recovered. (hardware ML-KEM-512)");
$display("Success: keys match, message recovered. (hardware K=%0d)", KP);
else
$display("FAILURE: %0d mismatches", errors);
$finish;

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

@@ -7,6 +7,8 @@ module tb_mlkem_kg_katK_xsim;
parameter KP = 2;
localparam EKB = 384*KP + 32;
localparam DKB = 768*KP + 96;
localparam integer CLK_PERIOD_NS = 100; // 10 MHz
localparam integer CLK_HALF_NS = CLK_PERIOD_NS / 2;
reg clk=0, rst_n=0, start_i=0;
reg [2:0] k_i;
@@ -33,7 +35,7 @@ module tb_mlkem_kg_katK_xsim;
.dbg_r_o(), .dbg_hek_o(),
.dbg_mprime_o(), .dbg_kbar_o(), .dbg_decz_o(), .dbg_dech_o()
);
always #5 clk = ~clk;
always #(CLK_HALF_NS) clk = ~clk;
reg [255:0] dmem [0:0];
reg [255:0] zmem [0:0];
@@ -42,8 +44,72 @@ module tb_mlkem_kg_katK_xsim;
reg [7:0] ek_got [0:EKB-1]; // ek bytes read back from DUT
reg [7:0] dk_got [0:DKB-1]; // dk bytes read back from DUT
integer c, i, errors, casenum, j;
integer state_cyc [0:31];
integer si;
reg [8*80-1:0] tag, dfile, zfile, ekfile, dkfile;
function [8*16-1:0] state_name;
input [4:0] s;
begin
case (s)
5'd0: state_name = "IDLE";
5'd1: state_name = "G";
5'd2: state_name = "A";
5'd3: state_name = "C";
5'd4: state_name = "N";
5'd5: state_name = "M";
5'd6: state_name = "E";
5'd7: state_name = "H";
5'd8: state_name = "ENC_LOAD";
5'd9: state_name = "ENC_H";
5'd10: state_name = "ENC_G";
5'd11: state_name = "ENC_A";
5'd12: state_name = "ENC_TDEC";
5'd13: state_name = "ENC_C";
5'd14: state_name = "ENC_N";
5'd15: state_name = "ENC_U";
5'd16: state_name = "ENC_C1";
5'd17: state_name = "ENC_V";
5'd18: state_name = "ENC_C2";
5'd19: state_name = "ENC_E2MV";
5'd20: state_name = "DEC_LOAD";
5'd21: state_name = "DEC_DECOMP";
5'd22: state_name = "DEC_SDEC";
5'd23: state_name = "DEC_NTT";
5'd24: state_name = "DEC_W";
5'd25: state_name = "DEC_MENC";
5'd26: state_name = "DEC_G";
5'd27: state_name = "DEC_J";
5'd28: state_name = "DEC_CMP";
5'd31: state_name = "DONE";
default: state_name = "UNKNOWN";
endcase
end
endfunction
task reset_timing;
begin
for (si = 0; si < 32; si = si + 1) state_cyc[si] = 0;
end
endtask
task print_timing;
input [8*16-1:0] op_name;
input integer cycles;
integer runtime_ns;
begin
runtime_ns = cycles * CLK_PERIOD_NS;
$display("TIME K=%0d CASE=%0d OP=%0s cycles=%0d clk=10MHz period=%0dns runtime=%0d ns (%0d.%03d us, %0d.%03d ms)",
KP, casenum, op_name, cycles, CLK_PERIOD_NS, runtime_ns,
runtime_ns/1000, runtime_ns%1000,
runtime_ns/1000000, (runtime_ns%1000000)/1000);
$display("TIME_BREAKDOWN K=%0d CASE=%0d OP=%0s", KP, casenum, op_name);
for (si = 0; si < 32; si = si + 1)
if (state_cyc[si] != 0)
$display(" STATE %-12s cycles=%0d time_ns=%0d", state_name(si[4:0]), state_cyc[si], state_cyc[si]*CLK_PERIOD_NS);
end
endtask
// Dump a byte array as offset-prefixed hex, 32 bytes/line.
task dump_bytes(input [8*16-1:0] name, input integer n);
integer a, b;
@@ -80,9 +146,12 @@ module tb_mlkem_kg_katK_xsim;
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
#1;
reset_timing;
c=0; while(!done_o && c<2000000) begin state_cyc[dut.st] = state_cyc[dut.st] + 1; @(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);
print_timing("KeyGen", c);
errors = 0;
dbg_byte_sel_i = 1'b0;

12
sync_rtl/top/TB/tb_mlkem_two_inst_xsim.v
View File

@@ -14,10 +14,12 @@
// xelab tb_mlkem_two_inst_xsim ; xsim
`timescale 1ns/1ps
module tb_mlkem_two_inst_xsim;
localparam KP = 2; // ML-KEM-512
localparam EKB = 384*KP + 32; // 800
localparam DKB = 768*KP + 96; // 1632
localparam CTB = 32*(10*KP + 4); // 768
parameter KP = 2; // 2/3/4 = ML-KEM-512/768/1024
localparam DU = (KP == 4) ? 11 : 10;
localparam DV = (KP == 4) ? 5 : 4;
localparam EKB = 384*KP + 32;
localparam DKB = 768*KP + 96;
localparam CTB = 32*(DU*KP + DV);
localparam MLEN = 11; // "hello world"
reg clk=0;
@@ -100,7 +102,7 @@ module tb_mlkem_two_inst_xsim;
msg_b[0]="h"; msg_b[1]="e"; msg_b[2]="l"; msg_b[3]="l"; msg_b[4]="o";
msg_b[5]=" "; msg_b[6]="w"; msg_b[7]="o"; msg_b[8]="r"; msg_b[9]="l"; msg_b[10]="d";
$display("=== ML-KEM hello_world (TWO instances: genenc + dec, ML-KEM-512) ===");
$display("=== ML-KEM hello_world (TWO instances: genenc + dec, K=%0d) ===", KP);
$write("Original: \""); for(i=0;i<MLEN;i=i+1) $write("%c", msg_b[i]); $display("\"\n");
a_rst_n=0; b_rst_n=0; repeat(4) @(posedge clk); a_rst_n=1; b_rst_n=1; @(posedge clk);

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

@@ -13,8 +13,10 @@ xvlog -sv --relax -i . sync_rtl/sha3/sha3_top_shared.v
xvlog -sv --relax -i . sync_rtl/sample_ntt/sample_ntt_sync_shared.v
xvlog -sv --relax -i . sync_rtl/sample_cbd/sample_cbd_sync_shared.v
xvlog -sv --relax -i . sync_rtl/ntt/barrett_mul.v
xvlog -sv --relax -i . sync_rtl/ntt/barrett_mul_pipe.v
xvlog -sv --relax -i . sync_rtl/ntt/zeta_rom.v
xvlog -sv --relax -i . sync_rtl/ntt/butterfly_unit.v
xvlog -sv --relax -i . sync_rtl/ntt/butterfly_unit_pipe.v
xvlog -sv --relax -i . sync_rtl/ntt/ntt_core.v
xvlog -sv --relax -i . sync_rtl/poly_mul/basecase_mul.v
xvlog -sv --relax -i . sync_rtl/poly_mul/basecase_mul_pipe.v

53
sync_rtl/top/mlkem_top.v
View File

@@ -22,7 +22,7 @@
`include "sync_rtl/common/defines.vh"
(* use_dsp = "no" *)
(* use_dsp = "yes" *)
module mlkem_top #(
parameter KMAX = 4 // storage sizing (worst case = ML-KEM-1024)
) (
@@ -194,6 +194,7 @@ module mlkem_top #(
wire [10:0] cin_rd_addr;
wire [7:0] cin_rd_data;
reg [10:0] cin_rd_addr_r; // D5/D7 read address (tied 0 until then)
reg dj_active; // background Decaps J(z||c) owns c_in read port
sd_bram #(.W(8), .D(2048), .A(11)) u_c_in_bram (
.clk(clk),
.rd_addr(cin_rd_addr), .rd_data(cin_rd_data),
@@ -203,7 +204,7 @@ module mlkem_top #(
// stage) so the registered BRAM read yields the byte 1 cycle later, matching
// the assemble/writeback pipeline. Otherwise the registered cin_rd_addr_r
// (D1 walker / idle) drives it.
assign cin_rd_addr = (st == ST_DEC_J) ? dj_c_idx[10:0] :
assign cin_rd_addr = dj_active ? dj_c_idx[10:0] :
(st == ST_DEC_CMP) ? cmp_idx : cin_rd_addr_r;
// ================================================================
@@ -567,7 +568,7 @@ module mlkem_top #(
// ST_ENC_A (snt), ST_ENC_C (cbd).
wire sel_sha3 = (st == ST_G) || (st == ST_H) ||
(st == ST_ENC_H) || (st == ST_ENC_G) ||
(st == ST_DEC_G) || (st == ST_DEC_J);
(st == ST_DEC_G) || dj_active;
wire sel_snt = (st == ST_A) || (st == ST_ENC_A);
wire sel_cbd = (st == ST_C) || (st == ST_ENC_C);
@@ -597,14 +598,14 @@ module mlkem_top #(
// G/H/J share one sha3_top. Single-block (mb_en=0): KeyGen G (ST_G, mode 00),
// Encaps G (ST_ENC_G, mode 11), Decaps D5 G (ST_DEC_G, mode 11). Multi-block
// absorb (mb_en=1): H(ek) (ST_H/ST_ENC_H, SHA3-256 pad) and Decaps D5 J
// (ST_DEC_J, SHAKE256 pad over z||c). mb_* inputs are muxed by phase.
wire sha3_mb_en = (st == ST_H) || (st == ST_ENC_H) || (st == ST_DEC_J);
// (background dj_active, SHAKE256 pad over z||c). mb_* inputs are muxed by phase.
wire sha3_mb_en = (st == ST_H) || (st == ST_ENC_H) || dj_active;
wire [1:0] sha3_mode = (st == ST_ENC_G || st == ST_DEC_G) ? 2'b11 : 2'b00;
// multi-block feed mux: J (ST_DEC_J) drives dj_*, else H drives h_*.
wire [1087:0] mb_block_mux = (st == ST_DEC_J) ? dj_block_r : h_block_r;
wire mb_valid_mux = (st == ST_DEC_J) ? dj_mbvalid : h_mbvalid;
wire mb_last_mux = (st == ST_DEC_J) ? dj_mblast : h_mblast;
wire mb_ack_mux = (st == ST_DEC_J) ? dj_ack : h_ack;
// multi-block feed mux: background J drives dj_*, else H drives h_*.
wire [1087:0] mb_block_mux = dj_active ? dj_block_r : h_block_r;
wire mb_valid_mux = dj_active ? dj_mbvalid : h_mbvalid;
wire mb_last_mux = dj_active ? dj_mblast : h_mblast;
wire mb_ack_mux = dj_active ? dj_ack : h_ack;
sha3_top_shared u_sha3 (
.clk(clk), .rst_n(rst_n),
.mode(sha3_mode), // G = SHA3-512 (only used when mb_en=0)
@@ -1145,12 +1146,12 @@ module mlkem_top #(
// D3: w = v' - INTT(sum_j s_hat[j] o u_hat[j]). Single output poly
// (u_row 0..0); MAC->INTT->SUB, then D4 encodes m'.
ST_DEC_W: if (u_row >= 3'd1) st_next = ST_DEC_MENC;
ST_DEC_MENC: if (men_done) st_next = ST_DEC_G;
// D5: (K',r') = G(m'||h) single-block, then K-bar = J(z||c) multi-block.
ST_DEC_G: if (sha3_vo) st_next = ST_DEC_J;
// D5 J done -> D6 re-encrypt: c' = K-PKE.Encrypt(ek_pke, m', r').
// Reuse the entire Encaps pipeline (rho load -> A -> C -> ... -> C2).
// r' is in r_r (CBD seed), m' in m_r (V/mu), ek_pke in ek_bram.
ST_DEC_MENC: if (men_done && dj_done) st_next = ST_DEC_G;
// D5: (K',r') = G(m'||h) single-block. K-bar = J(z||c) runs in a
// background SHAKE256 walker after D1, so the main FSM only waits
// for it at the D4->D5 boundary if it somehow has not finished.
ST_DEC_G: if (sha3_vo) st_next = ST_ENC_LOAD;
// Legacy state kept unused; background J normally skips it.
ST_DEC_J: if (dj_done) st_next = ST_ENC_LOAD;
ST_G: if (sha3_vo) st_next = ST_A;
ST_A: if (a_pair >= kk_rt) st_next = ST_C;
@@ -1303,6 +1304,7 @@ module mlkem_top #(
cmp_neq <= 1'b0;
cmp_done <= 1'b0;
dec_reject <= 1'b0;
dj_active <= 1'b0;
dj_blk <= 4'd0;
dj_byte <= 8'd0;
dj_phase <= 2'd0;
@@ -1349,6 +1351,8 @@ module mlkem_top #(
end else if (op_i == 2'd2) begin
// Decaps: dk_pke/ek_pke already in BRAM; z/H(ek) captured into
// z_r/hek_r during load (below). Nothing else to arm in D0.
dj_active <= 1'b0;
dj_done <= 1'b0;
end else begin
sha3_valid <= 1'b1;
sha3_ack <= 1'b1;
@@ -1993,9 +1997,11 @@ module mlkem_top #(
sha3_ack <= 1'b0;
end
// Arm D5 J(z||c) when G completes (ST_DEC_G -> ST_DEC_J): assemble the
// first 136-byte block. dj_ack high to consume the final digest.
if (st == ST_DEC_G && st_next == ST_DEC_J) begin
// Arm D5 J(z||c) as soon as D1 has stopped reading c_in_bram. It
// then runs in the background during D2/D3/D4 using the otherwise
// idle SHAKE256 path and c_in read port.
if (st == ST_DEC_DECOMP && st_next == ST_DEC_SDEC) begin
dj_active <= 1'b1;
dj_blk <= 4'd0;
dj_byte <= 8'd0;
dj_phase <= 2'd0;
@@ -2007,14 +2013,14 @@ module mlkem_top #(
cin_rd_addr_r <= 11'd0; // present c byte 0 (g=32 -> c_idx 0)
end
// ---- ST_DEC_J: K-bar = J(z||c) multi-block SHAKE256 ----
// ---- Background Decaps J: K-bar = J(z||c) multi-block SHAKE256 ----
// Mirror of the H(ek) multi-block machine. Byte source by global g:
// g < 32 -> z_r byte g
// 32 <= g < msglen -> c_in_bram byte (g-32), registered read
// else -> SHAKE256 pad constant
// c_in_bram read is registered: present c_idx for dj_byte this cycle,
// write back the byte that arrived for the addr presented last cycle.
if (st == ST_DEC_J && !dj_done) begin
if (dj_active && !dj_done) begin
case (dj_phase)
2'd0: begin
// writeback the byte read for the previous address
@@ -2034,7 +2040,7 @@ module mlkem_top #(
dj_wb_zidx <= dj_g[4:0];
dj_wb_pad <= dj_padconst(dj_g);
// c_in addr is presented combinationally (cin_rd_addr
// mux uses dj_c_idx during ST_DEC_J); data lands next
// mux uses dj_c_idx while dj_active); data lands next
// cycle, matching this byte's writeback.
dj_byte <= dj_byte + 8'd1;
end else begin
@@ -2056,6 +2062,7 @@ module mlkem_top #(
if (sha3_vo) begin
kbar_r <= sha3_hash[255:0];
dj_done <= 1'b1;
dj_active <= 1'b0;
dj_phase<= 2'd3;
end else if (h_mbready) begin
dj_blk <= dj_blk + 4'd1;
@@ -2095,7 +2102,7 @@ module mlkem_top #(
end
if (st_next == ST_ENC_LOAD &&
(st == ST_ENC_G || st == ST_DEC_J)) begin
(st == ST_ENC_G || st == ST_DEC_G || st == ST_DEC_J)) begin
rl_idx <= 6'd0;
rl_widx <= 6'd0;
rl_vld <= 1'b0;