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

refactor(kg): move ek/dk_pke byte storage into BRAM (sd_bram)

Browse Source 
Phase 1 of migrating mlkem_top's large arrays to inferable RAM. ek_mem
and dkp_mem reg arrays are replaced by two sd_bram instances (1R/1W,
registered read). Datapath changes to fit single-port-per-cycle BRAM:

  - ST_E writes 1 byte/cycle (was 3): added e_byte sub-counter; ST_E
    length ~3x (K=2 KeyGen 21403->22433 cyc, ~5%).
  - ST_H ek read is now registered: assemble phase presents the address
    one cycle ahead and writes back the byte that arrived (h_wb_* pipe),
    h_byte runs 0..136 to flush the final byte. Pad bytes via h_padconst.
  - dbg_byte_o/dbg_dk_o read combinationally off the BRAM registered
    output (net 1-cycle latency, within the TB's 2-cycle read wait);
    region decode for dk readback unchanged.

Add sd_bram.v to the top TB compile list. Verified byte-exact vs NIST
KAT: K=2 c0-4, K=3 c0-2, K=4 c0-2 -> 11/11 PASS, 0 file-not-found.
This commit is contained in:
FallenSigh
2026-06-28 14:49:05 +08:00
parent 8774e03a0e
commit 5a7d5d6a47
2 changed files with 130 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -19,6 +19,7 @@ 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/poly_mul_zeta_rom.v
xvlog -sv --relax -i . sync_rtl/poly_mul/poly_mul_sync.v
xvlog -sv --relax -i . sync_rtl/storage/sd_bram.v
xvlog -sv --relax -i . sync_rtl/top/mlkem_top.v
# ---- Step 2: compile parametric KAT testbench ----

176
sync_rtl/top/mlkem_top.v
View File

@@ -92,29 +92,46 @@ module mlkem_top #(
// 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;
// ek / dk_pke byte storage in BRAM (sd_bram: 1R/1W, 1-cycle read latency).
// Each single read port is shared time-disjointly: ST_H assemble reads ek
// (in ST_H), debug readback reads ek/dk (in ST_DONE). Write port driven by
// ST_E (1 byte/cycle) and the rho-tail copy.
wire [10:0] ek_rd_addr, dkp_rd_addr;
wire [7:0] ek_rd_data, dkp_rd_data;
reg ek_we, dkp_we;
reg [10:0] ek_wa, dkp_wa;
reg [7:0] ek_wd, dkp_wd;
sd_bram #(.W(8), .D(2048), .A(11)) u_ek_bram (
.clk(clk),
.rd_addr(ek_rd_addr), .rd_data(ek_rd_data),
.wr_en(ek_we), .wr_addr(ek_wa), .wr_data(ek_wd)
);
sd_bram #(.W(8), .D(2048), .A(11)) u_dkp_bram (
.clk(clk),
.rd_addr(dkp_rd_addr), .rd_data(dkp_rd_data),
.wr_en(dkp_we), .wr_addr(dkp_wa), .wr_data(dkp_wd)
);
// 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_rt)
dbg_dk_r <= dkp_mem[dbg_dk_idx_i];
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)*8 +: 8];
end
assign dbg_dk_o = dbg_dk_r;
// Debug-region selects for dk readback (combinational region decode).
wire dbgdk_in_dkp = (dbg_dk_idx_i < dk_bytes_rt);
wire dbgdk_in_ek = (dbg_dk_idx_i >= dk_bytes_rt) && (dbg_dk_idx_i < dk_ek_end);
// Read-address muxes (declared here; depend on ST_H h_g defined below via
// the function-free expression h_blk*136 + h_byte). See assigns after the
// H-stage register declarations.
// dbg_byte_o / dbg_dk_o are combinational taps over the (registered) BRAM
// read data; net 1-cycle latency, well within the TB's 2-cycle read wait.
assign dbg_byte_o = dbg_byte_sel_i ? dkp_rd_data : ek_rd_data;
assign dbg_dk_o = dbgdk_in_dkp ? dkp_rd_data :
dbgdk_in_ek ? ek_rd_data :
(dbg_dk_idx_i < dk_hek_end) ? hek_r[(dbg_dk_idx_i - dk_ek_end)*8 +: 8]
: z_i[(dbg_dk_idx_i - dk_hek_end)*8 +: 8];
// ================================================================
// Top-level FSM (built incrementally). Stage 2a: G only.
@@ -125,7 +142,7 @@ module mlkem_top #(
localparam ST_C = 4'd3; // generate s[i],e[i] via CBD
localparam ST_N = 4'd4; // forward NTT of s[i],e[i] in place
localparam ST_M = 4'd5; // matrix accumulate t_hat = e_hat + sum A o s_hat
localparam ST_E = 4'd6; // byteEncode12 -> ek_mem, dkp_mem
localparam ST_E = 4'd6; // byteEncode12 -> ek/dk BRAM
localparam ST_H = 4'd7; // H(ek) via multi-block SHA3-256
localparam ST_DONE = 4'd15;
@@ -185,8 +202,16 @@ module mlkem_top #(
reg h_ack;
reg [255:0] hek_r; // captured H(ek)
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 [7:0] h_byte; // 0..136 byte address being presented (assemble)
reg [1:0] h_phase; // 0=assemble 1=feed 2=wait-perm 3=done
// ek BRAM read is registered (1-cycle latency); assemble presents the
// address for h_byte this cycle and writes the byte that arrived (for the
// address presented last cycle) into h_block_r. Writeback pipeline regs:
reg h_wb_vld; // a byte is arriving this cycle
reg [7:0] h_wb_idx; // its position within the 136-byte block
reg [11:0] h_wb_g; // its global ek byte index
reg [7:0] h_wb_pad; // pad constant to use if g is out of ek range
reg h_wb_inek; // 1 if g is within ek range (use BRAM data)
sha3_top u_sha3_h (
.clk(clk), .rst_n(rst_n),
@@ -207,19 +232,32 @@ module mlkem_top #(
// 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);
// ek bytes come from the ek BRAM (registered read); only the pad/zero bytes
// are constants, returned by h_padconst. Reads runtime ek_bytes_rt/h_last_rt.
function [7:0] h_padconst(input [3:0] blk, input [7:0] b);
integer g;
begin
g = blk*136 + b;
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;
if (g == h_last_rt && g == ek_bytes_rt) h_padconst = 8'h86; // 0x06|0x80
else if (g == ek_bytes_rt) h_padconst = 8'h06;
else if (g == h_last_rt) h_padconst = 8'h80;
else h_padconst = 8'h00;
end
endfunction
// Global ek byte index for the address currently presented in ST_H assemble.
wire [11:0] h_g_addr = {8'd0, h_blk} * 12'd136 + {4'd0, h_byte};
// ek BRAM read-address mux: ST_H assemble drives it; otherwise debug readback.
// dbg_dk ek-region takes priority (its loop runs after the ek loop), so the
// single read port serves dbg_byte_o (ek loop) then dbg_dk_o (dk loop).
wire [11:0] dbgdk_ek_off = dbg_dk_idx_i - dk_bytes_rt; // offset within ek region
assign ek_rd_addr = (st == ST_H) ? h_g_addr[10:0] :
dbgdk_in_ek ? dbgdk_ek_off[10:0]
: dbg_byte_idx_i;
// dkp BRAM read-address mux: dbg_byte (sel=1) or dbg_dk (dkp region).
assign dkp_rd_addr = dbg_byte_sel_i ? dbg_byte_idx_i : dbg_dk_idx_i[10:0];
// ---- sample_ntt_sync: Â[i][j] = SampleNTT(rho || j || i) ----
reg snt_valid;
wire snt_ready;
@@ -308,6 +346,7 @@ module mlkem_top #(
// bytes || rho; dk_pke = s_hat[0..K-1] bytes. Walk coeff pairs per poly.
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 [1:0] e_byte; // 0..2 byte within the current pair (BRAM: 1 write/cycle)
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
@@ -324,6 +363,8 @@ module mlkem_top #(
// byte base offset within target memory: poly index *384 (= 128 pairs *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
// BRAM write is 1 byte/cycle: select packed byte by e_byte sub-counter
wire [7:0] e_byte_d = (e_byte == 2'd0) ? e_b0 : (e_byte == 2'd1) ? e_b1 : e_b2;
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]
@@ -405,8 +446,15 @@ module mlkem_top #(
pm_valid <= 1'b0;
e_poly <= 3'd0;
e_pair <= 8'd0;
e_byte <= 2'd0;
e_rho <= 10'd0;
e_done <= 1'b0;
ek_we <= 1'b0;
dkp_we <= 1'b0;
ek_wa <= 11'd0;
dkp_wa <= 11'd0;
ek_wd <= 8'd0;
dkp_wd <= 8'd0;
h_block_r <= 1088'd0;
h_mbvalid <= 1'b0;
h_mblast <= 1'b0;
@@ -415,9 +463,18 @@ module mlkem_top #(
h_blk <= 3'd0;
h_byte <= 8'd0;
h_phase <= 2'd0;
h_wb_vld <= 1'b0;
h_wb_idx <= 8'd0;
h_wb_g <= 12'd0;
h_wb_pad <= 8'd0;
h_wb_inek <= 1'b0;
end else begin
st <= st_next;
// BRAM write-enables default low; pulsed where a byte is written.
ek_we <= 1'b0;
dkp_we <= 1'b0;
// Kick off G when entering ST_G
if (st == ST_IDLE && start_i) begin
k_r <= k_i; // capture runtime ML-KEM param
@@ -603,32 +660,41 @@ module mlkem_top #(
if (st == ST_M && st_next == ST_E) begin
e_poly <= 3'd0;
e_pair <= 8'd0;
e_byte <= 2'd0;
e_rho <= 10'd0;
e_done <= 1'b0;
end
// ---- ST_E: byteEncode12 t_hat -> ek_mem, s_hat -> dkp_mem, ek tail = rho ----
// ---- ST_E: byteEncode12 t_hat -> ek, s_hat -> dk_pke, 1 byte/cycle ----
if (st == ST_E && !e_done) begin
if (e_poly < {1'b0, k_r, 1'b0}) begin
// pack current coeff-pair (3 bytes): [0,K)=ek, [K,2K)=dk_pke
// write current byte (e_byte_d) to target BRAM
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;
ek_we <= 1'b1;
ek_wa <= e_boff[10:0] + {9'd0, e_byte};
ek_wd <= e_byte_d;
end else begin
dkp_mem[e_boff] <= e_b0;
dkp_mem[e_boff + 1] <= e_b1;
dkp_mem[e_boff + 2] <= e_b2;
dkp_we <= 1'b1;
dkp_wa <= e_boff[10:0] + {9'd0, e_byte};
dkp_wd <= e_byte_d;
end
if (e_pair == 8'd127) begin
e_pair <= 8'd0;
e_poly <= e_poly + 5'd1; // next poly (or ->2K = rho phase)
// advance sub-counters
if (e_byte == 2'd2) begin
e_byte <= 2'd0;
if (e_pair == 8'd127) begin
e_pair <= 8'd0;
e_poly <= e_poly + 5'd1; // next poly or -> rho phase
end else begin
e_pair <= e_pair + 8'd1;
end
end else begin
e_pair <= e_pair + 8'd1;
e_byte <= e_byte + 2'd1;
end
end else begin
// rho copy: ek_mem[384*K + r] = rho byte r (r = 0..31)
ek_mem[dk_bytes_rt + e_rho] <= rho_r[e_rho*8 +: 8];
// rho copy: ek[384*K + r] = rho byte r (r = 0..31), 1 byte/cycle
ek_we <= 1'b1;
ek_wa <= dk_bytes_rt[10:0] + {1'b0, e_rho};
ek_wd <= rho_r[e_rho*8 +: 8];
if (e_rho == 10'd31) e_done <= 1'b1;
else e_rho <= e_rho + 10'd1;
end
@@ -642,21 +708,37 @@ module mlkem_top #(
h_mbvalid<= 1'b0;
h_mblast <= 1'b0;
h_ack <= 1'b1; // ready to consume final digest
h_wb_vld <= 1'b0; // no pending writeback yet
end
// ---- ST_H: H(ek) via multi-block SHA3-256 (6 pre-padded blocks) ----
// ---- ST_H: H(ek) via multi-block SHA3-256 (6/9/12 pre-padded blocks) ----
if (st == ST_H) begin
case (h_phase)
// assemble 136 bytes of block h_blk into h_block_r
// assemble 136 bytes of block h_blk into h_block_r.
// ek BRAM read is registered: present addr for h_byte this
// cycle, write back the byte that arrived for the addr we
// presented last cycle (h_wb_*). h_byte runs 0..136 (one
// extra cycle to flush the final writeback).
2'd0: begin
h_block_r[h_byte*8 +: 8] <= h_padbyte(h_blk, h_byte);
if (h_byte == 8'd135) begin
// writeback the byte read for the previous address
if (h_wb_vld)
h_block_r[h_wb_idx*8 +: 8] <= h_wb_inek ? ek_rd_data : h_wb_pad;
if (h_byte <= 8'd135) begin
// set up writeback for the address presented this cycle
h_wb_vld <= 1'b1;
h_wb_idx <= h_byte;
h_wb_g <= h_g_addr;
h_wb_inek <= (h_g_addr < ek_bytes_rt);
h_wb_pad <= h_padconst(h_blk, h_byte);
h_byte <= h_byte + 8'd1;
end else begin
// h_byte==136: final writeback (for byte 135) done above
h_wb_vld <= 1'b0;
h_byte <= 8'd0;
h_mbvalid <= 1'b1;
h_mblast <= (h_blk == h_nblk_rt - 4'd1);
h_phase <= 2'd1; // feed
end else begin
h_byte <= h_byte + 8'd1;
end
end
// feed: hold valid until accepted (mb_ready drops)