From 5a7d5d6a47d6422e3925c77c52c020590e91b48b Mon Sep 17 00:00:00 2001 From: FallenSigh Date: Sun, 28 Jun 2026 14:49:05 +0800 Subject: [PATCH] refactor(kg): move ek/dk_pke byte storage into BRAM (sd_bram) 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. --- sync_rtl/top/TB/xsim_run.tcl | 1 + sync_rtl/top/mlkem_top.v | 176 +++++++++++++++++++++++++---------- 2 files changed, 130 insertions(+), 47 deletions(-) diff --git a/sync_rtl/top/TB/xsim_run.tcl b/sync_rtl/top/TB/xsim_run.tcl index 621b149..863ef8c 100644 --- a/sync_rtl/top/TB/xsim_run.tcl +++ b/sync_rtl/top/TB/xsim_run.tcl @@ -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 ---- diff --git a/sync_rtl/top/mlkem_top.v b/sync_rtl/top/mlkem_top.v index 6d10944..2e2db77 100644 --- a/sync_rtl/top/mlkem_top.v +++ b/sync_rtl/top/mlkem_top.v @@ -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)