// 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 // ML-KEM key pair per FIPS 203 Algorithm 16 (KeyGen_internal): // (rho,sigma) = G(d || K) // A_hat[i][j] = SampleNTT(rho || j || i) i,j in 0..K-1 // s[i] = CBD3(PRF(sigma, i)), e[i] = CBD3(PRF(sigma, K+i)) // s_hat[i] = NTT(s[i]); e_hat[i] = NTT(e[i]) // t_hat[i] = e_hat[i] + sum_j A_hat[i][j] o s_hat[j] // ek = byteEncode12(t_hat[0..K-1]) || rho // dk = byteEncode12(s_hat[0..K-1]) || ek || H(ek) || z // // Built incrementally and verified stage-by-stage against ml-kem-r golden // vectors (test_framework/modules/mlkem_keygen/golden) and NIST KAT. // // Uses independent (verified) leaf modules, each with its own keccak_core: // sha3_top, sample_ntt_sync, sample_cbd_sync, ntt_core, poly_mul_sync, // mod_add_sync. No shared-keccak arbiter. `include "sync_rtl/common/defines.vh" module mlkem_top #( parameter KMAX = 4 // storage sizing (worst case = ML-KEM-1024) ) ( input clk, 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] z_i, // implicit-rejection seed z input start_i, // pulse to begin KeyGen output busy_o, // high while running output done_o, // pulse when ek/dk ready // Debug readback tap: read one stored coefficient by (poly slot, index). // Lets stage TBs verify intermediates without wide buses. input [3:0] dbg_slot_i, // poly slot (see localparams below) input [7:0] dbg_idx_i, // coefficient index 0..255 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 [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. input [11:0] dbg_dk_idx_i, output [7:0] dbg_dk_o, // Debug taps for hash outputs output [255:0] dbg_rho_o, output [255:0] dbg_sigma_o ); localparam Q = `Q; // 3329 // 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, sized for KMAX (worst case). Runtime k uses a // sub-range. 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_rt .. +k-1 // e_hat[i] : slots slot_e_rt .. +k-1 // t_hat[i] : slots slot_t_rt .. +k-1 // NUM_SLOTS = KMAX*KMAX + 3*KMAX = 28 for KMAX=4. // ================================================================ localparam NUM_SLOTS = KMAX*KMAX + 3*KMAX; localparam SAW = 5; // slot-address width (>=clog2(28)) // ================================================================ // Coefficient storage as 3 banks (Phase 2). Was one async-read reg // array polymem[0:28*256-1]; split by polynomial so each becomes a // registered-read sd_bram (infers BRAM / ASIC SRAM) in the final step. // bank_a : A_hat[i][j] slots 0..K^2-1 -> D=KMAX*KMAX*256=4096 // bank_se: s_hat||e_hat slots slot_s.. -> D=2*KMAX*256=2048 (rel 0..2K-1) // bank_t : t_hat[i] slots slot_t.. -> D=KMAX*256=1024 (rel 0..K-1) // Addresses are base-relative: bank index = abs_slot - base_slot. // Stage 2b (in progress): per-consumer registered read-ahead. Banks are // still async reg arrays here; converted consumers read via their own // 1-cycle pipeline reg (== sd_bram timing). Once every consumer of a // bank is registered, the array is replaced by an sd_bram with its read // port muxed across the (phase-disjoint) consumers. // ================================================================ localparam PA_AW = 12; // bank_a addr width (4096) localparam PSE_AW = 11; // bank_se addr width (2048) localparam PT_AW = 10; // bank_t addr width (1024) reg [11:0] bank_a [0:(1<= slot_t_rt) dbg_coeff_r <= bank_t [dbg_t_addr[PT_AW-1:0]]; else if (dbg_slot_i >= slot_s_rt) dbg_coeff_r <= bank_se[dbg_se_addr[PSE_AW-1:0]]; else dbg_coeff_r <= bank_a [dbg_a_addr[PA_AW-1:0]]; end assign dbg_coeff_o = dbg_coeff_r; // ek and dk_pke byte memories sized for KMAX. localparam EK_MAX = 384*KMAX + 32; // 1568 localparam DK_MAX = 384*KMAX; // 1536 // 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 // 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. // ================================================================ localparam ST_IDLE = 4'd0; localparam ST_G = 4'd1; // run G(d||K), capture rho/sigma localparam ST_A = 4'd2; // generate A_hat[i][j] via SampleNTT 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/dk BRAM localparam ST_H = 4'd7; // H(ek) via multi-block SHA3-256 localparam ST_DONE = 4'd15; reg [3:0] st, st_next; reg [255:0] rho_r, sigma_r; // 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_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]) 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 < {2'b0, k_r}) ? (slot_s_rt + c_poly) : (slot_e_rt + (c_poly - {2'b0, k_r})); assign busy_o = (st != ST_IDLE); assign done_o = (st == ST_DONE); assign dbg_rho_o = rho_r; assign dbg_sigma_o = sigma_r; // ---- sha3_top in G mode: data_i = {K_byte, d} (d byte0 in [7:0]) ---- reg sha3_valid; wire sha3_ready; wire [511:0] sha3_hash; wire sha3_vo; reg sha3_ack; // consumer ready for hash wire [511:0] g_data = {248'b0, 5'b0, k_r, d_i}; // data_i[263:256]=k, [255:0]=d // ================================================================ // Shared keccak_core + phase mux (3 consumers -> 1 core) // // G/H (u_sha3), SampleNTT (u_snt) and CBD (u_cbd) each need a Keccak // permutation, but they run in DISJOINT top-FSM phases (ST_G/ST_H, ST_A, // ST_C respectively), so at most one is ever active. One keccak_core // serves all three: kc_state_i/kc_valid_i are muxed by phase; kc_valid_o // is broadcast but GATED per consumer so an inactive consumer never // latches a permutation result meant for another (the samplers latch // squeeze state unconditionally on their kc_valid_o input). // ================================================================ wire [1599:0] kc_state_o; // shared core output (to all consumers) wire kc_valid_o; // shared core output-valid wire kc_ready_o; // shared core input-ready // per-consumer drives toward the core wire [1599:0] sha3_kc_state_i, snt_kc_state_i, cbd_kc_state_i; wire sha3_kc_valid_i, snt_kc_valid_i, cbd_kc_valid_i; /* verilator lint_off UNUSEDSIGNAL */ wire sha3_kc_ready_i, snt_kc_ready_i, cbd_kc_ready_i; // all 1'b1 /* verilator lint_on UNUSEDSIGNAL */ // phase selects (mutually exclusive) wire sel_sha3 = (st == ST_G) || (st == ST_H); wire sel_snt = (st == ST_A); wire sel_cbd = (st == ST_C); // gated output-valid: only the active consumer sees kc_valid_o wire kc_valid_o_sha3 = kc_valid_o & sel_sha3; wire kc_valid_o_snt = kc_valid_o & sel_snt; wire kc_valid_o_cbd = kc_valid_o & sel_cbd; // input mux: route the active consumer's request to the core wire [1599:0] kc_state_i_mux = sel_snt ? snt_kc_state_i : sel_cbd ? cbd_kc_state_i : sha3_kc_state_i; wire kc_valid_i_mux = sel_snt ? snt_kc_valid_i : sel_cbd ? cbd_kc_valid_i : sel_sha3 ? sha3_kc_valid_i : 1'b0; keccak_core #(.ROUNDS(24)) u_keccak ( .clk(clk), .rst_n(rst_n), .state_i(kc_state_i_mux), .valid_i(kc_valid_i_mux), .ready_o(kc_ready_o), .state_o(kc_state_o), .valid_o(kc_valid_o), .ready_i(1'b1) // consumers always accept (kc_ready_i=1) ); // ---- single shared sha3_top serving BOTH G and H ---- // G (ST_G) uses single-block mode (mb_en=0); H(ek) (ST_H) uses the // multi-block absorb path (mb_en=1). These phases are disjoint in the // top FSM, so one sha3_top (one keccak_core) is sufficient. mb_en and // ready_i are muxed by phase; data_i/mode only matter while mb_en=0. wire sha3_mb_en = (st == ST_H); sha3_top_shared u_sha3 ( .clk(clk), .rst_n(rst_n), .mode(2'b00), // G = SHA3-512 (only used when mb_en=0) .data_i(g_data), .valid_i(sha3_valid), .ready_o(sha3_ready), .hash_o(sha3_hash), .valid_o(sha3_vo), .ready_i(sha3_mb_en ? h_ack : sha3_ack), .mb_en(sha3_mb_en), .mb_block_i(h_block_r), .mb_valid_i(h_mbvalid), .mb_last_i(h_mblast), .mb_ready_o(h_mbready), // shared keccak_core interface (gated by phase in the mux below) .kc_state_o(kc_state_o), .kc_valid_o(kc_valid_o_sha3), .kc_ready_o(kc_ready_o), .kc_state_i(sha3_kc_state_i), .kc_valid_i(sha3_kc_valid_i), .kc_ready_i(sha3_kc_ready_i) ); // ---- multi-block H(ek) state (SHA3-256, 6/9/12 blocks); fed to shared u_sha3 ---- reg [1087:0] h_block_r; // current pre-padded rate block reg h_mbvalid; reg h_mblast; wire h_mbready; wire [511:0] h_hash; wire h_vo; 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..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) // h_hash / h_vo are now served by the shared u_sha3 above (mb_en=1 during // ST_H). They alias the single core's outputs; the H consumer logic below // already gates on st==ST_H, and u_sha3's valid_o/hash_o are mb-selected. assign h_hash = sha3_hash; assign h_vo = sha3_vo; // 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. // 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 == 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; wire [11:0] snt_coeff; wire snt_vo; wire snt_last; reg snt_ack; // we accept coeffs sample_ntt_sync_shared #(.K(KMAX)) u_snt ( .clk(clk), .rst_n(rst_n), .rho_i(rho_r), .k_i(k_r), .i_idx(a_i[1:0]), .j_idx(a_j[1:0]), .valid_i(snt_valid), .ready_o(snt_ready), .coeff_o(snt_coeff), .valid_o(snt_vo), .ready_i(snt_ack), .last_o(snt_last), // shared keccak_core interface .kc_state_o(kc_state_o), .kc_valid_o(kc_valid_o_snt), .kc_ready_o(kc_ready_o), .kc_state_i(snt_kc_state_i), .kc_valid_i(snt_kc_valid_i), .kc_ready_i(snt_kc_ready_i) ); // ---- sample_cbd_sync: s[i]=CBD3(PRF(sigma,i)), e[i]=CBD3(PRF(sigma,K+i)) ---- reg cbd_valid; wire cbd_ready; wire [11:0] cbd_coeff; // 12-bit signed (two's complement) wire cbd_vo; wire cbd_last; reg cbd_ack; sample_cbd_sync_shared u_cbd ( .clk(clk), .rst_n(rst_n), .seed_i(sigma_r), .nonce_i(c_nonce), .eta_i(eta1_rt), .valid_i(cbd_valid), .ready_o(cbd_ready), .coeff_o(cbd_coeff), .valid_o(cbd_vo), .ready_i(cbd_ack), .last_o(cbd_last), // shared keccak_core interface .kc_state_o(kc_state_o), .kc_valid_o(kc_valid_o_cbd), .kc_ready_o(kc_ready_o), .kc_state_i(cbd_kc_state_i), .kc_valid_i(cbd_kc_valid_i), .kc_ready_i(cbd_kc_ready_i) ); // signed (two's complement) -> [0,Q): add Q when negative wire [11:0] cbd_modq = cbd_coeff[11] ? (cbd_coeff + 12'(Q)) : cbd_coeff; // ---- 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 [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-AHEAD pointer 0..256 (leads consume by 1) reg [7:0] n_widx; // output write index 0..255 reg n_valid; // feeding coeffs to ntt_core (delayed 1 cyc vs n_ridx) reg n_loading; // 1 while presenting load addresses to bank_se reg n_pending; // waiting for ntt_core IDLE to start next slot reg [11:0] n_rd_data; // registered bank_se read (== sd_bram timing) wire [SAW-1:0] n_slot_addr = slot_s_rt + n_slot; // s_hat then e_hat contiguous wire ntt_ready; wire [11:0] ntt_coeff; wire ntt_vo; wire ntt_done; // bank_se read addr for the NTT load (relative slot = n_slot); registered // into n_rd_data, which feeds ntt_core 1 cycle later. wire [13:0] ntt_rd_full = n_slot*256 + n_ridx[7:0]; wire [11:0] ntt_in = n_rd_data; ntt_core u_ntt ( .clk(clk), .rst_n(rst_n), .coeff_in(ntt_in), .valid_i(n_valid), .ready_o(ntt_ready), .mode(1'b0), // forward NTT, no scaling .coeff_out(ntt_coeff), .valid_o(ntt_vo), .ready_i(1'b1), // always accept output .done_o(ntt_done) ); // ---- 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 [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 reg m_pending; // wait for poly_mul IDLE before next (i,j) // ---- 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 [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 wire e_is_dk = (e_poly >= {1'b0, k_r}); 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_rt + e_pidx) : (slot_t_rt + e_pidx); // two coeffs of the current pair: ek half reads t_hat (bank_t), dk half // reads s_hat (bank_se). Relative index within the bank = e_pidx. wire [13:0] e_rd0_full = e_pidx*256 + {e_pair, 1'b0}; wire [13:0] e_rd1_full = e_pidx*256 + {e_pair, 1'b1}; wire [11:0] e_c0 = e_is_dk ? bank_se[e_rd0_full[PSE_AW-1:0]] : bank_t[e_rd0_full[PT_AW-1:0]]; wire [11:0] e_c1 = e_is_dk ? bank_se[e_rd1_full[PSE_AW-1:0]] : bank_t[e_rd1_full[PT_AW-1:0]]; // 3 packed bytes wire [7:0] e_b0 = e_c0[7:0]; 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 [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] wire [SAW-1:0] m_eslot = slot_e_rt + m_i; // e_hat[i] wire [SAW-1:0] m_tslot = slot_t_rt + m_i; // t_hat[i] reg pm_valid; wire pm_ready; wire [11:0] pm_coeff; wire pm_vo; // pm_a: A_hat[i][j] in bank_a (abs slot m_aslot). pm_b: s_hat[j] in // bank_se (relative slot = m_sslot - slot_s_rt = m_j). wire [13:0] pm_a_full = m_aslot*256 + m_ld[7:0]; wire [13:0] pm_b_full = m_j*256 + m_ld[7:0]; wire [11:0] pm_a_in = bank_a [pm_a_full[PA_AW-1:0]]; wire [11:0] pm_b_in = bank_se[pm_b_full[PSE_AW-1:0]]; poly_mul_sync u_pmul ( .clk(clk), .rst_n(rst_n), .coeff_a_in(pm_a_in), .coeff_b_in(pm_b_in), .valid_i(pm_valid), .ready_o(pm_ready), .coeff_out(pm_coeff), .valid_o(pm_vo), .ready_i(1'b1) ); // accumulator source: e_hat[i] for first term (j==0), else running t_hat[i]. // e_hat[i] lives in bank_se at relative slot (slot_e_rt-slot_s_rt + m_i) = K+m_i. // t_hat[i] lives in bank_t at relative slot m_i. wire [13:0] m_eacc_full = ({2'b0, k_r} + {2'b0, m_i})*256 + m_oidx; // K+m_i wire [13:0] m_tacc_full = m_i*256 + m_oidx; wire [11:0] m_acc_src = (m_j == 2'd0) ? bank_se[m_eacc_full[PSE_AW-1:0]] : bank_t [m_tacc_full[PT_AW-1:0]]; // (a + b) mod Q (both < Q, sum < 2Q): one conditional subtract wire [12:0] m_sum = {1'b0, m_acc_src} + {1'b0, pm_coeff}; wire [11:0] m_accq = (m_sum >= 13'(Q)) ? (m_sum - 13'(Q)) : m_sum[11:0]; always @(*) begin st_next = st; case (st) ST_IDLE: if (start_i) st_next = ST_G; ST_G: if (sha3_vo) st_next = ST_A; ST_A: if (a_pair >= kk_rt) st_next = ST_C; ST_C: if (c_poly >= {1'b0, k_r, 1'b0}) st_next = ST_N; ST_N: if (n_slot >= {1'b0, k_r, 1'b0}) st_next = ST_M; ST_M: if (m_i >= k_r) st_next = ST_E; ST_E: if (e_done) st_next = ST_H; ST_H: if (h_phase == 2'd3) st_next = ST_DONE; ST_DONE: st_next = ST_IDLE; default: st_next = ST_IDLE; endcase end always @(posedge clk or negedge rst_n) begin if (!rst_n) begin st <= ST_IDLE; k_r <= 3'd0; rho_r <= 256'd0; sigma_r <= 256'd0; sha3_valid <= 1'b0; sha3_ack <= 1'b0; snt_valid <= 1'b0; snt_ack <= 1'b0; a_pair <= 5'd0; a_i <= 3'd0; a_j <= 3'd0; a_widx <= 8'd0; a_busy <= 1'b0; cbd_valid <= 1'b0; cbd_ack <= 1'b0; c_poly <= 3'd0; c_widx <= 8'd0; c_busy <= 1'b0; n_slot <= 3'd0; n_ridx <= 9'd0; n_widx <= 8'd0; n_valid <= 1'b0; n_loading <= 1'b0; n_rd_data <= 12'd0; n_pending <= 1'b0; m_i <= 2'd0; m_j <= 2'd0; m_ld <= 9'd0; m_oidx <= 8'd0; m_loading <= 1'b0; m_pending <= 1'b0; 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; h_ack <= 1'b0; hek_r <= 256'd0; 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 sha3_valid <= 1'b1; sha3_ack <= 1'b1; end // Drop valid once accepted if (sha3_valid && sha3_ready) sha3_valid <= 1'b0; // Capture rho/sigma when G completes; arm A stage if (st == ST_G && sha3_vo) begin rho_r <= sha3_hash[255:0]; // rho = G output bytes 0..31 sigma_r <= sha3_hash[511:256]; // sigma = bytes 32..63 sha3_ack <= 1'b0; snt_valid <= 1'b1; // start first SampleNTT snt_ack <= 1'b1; a_pair <= 5'd0; a_i <= 3'd0; a_j <= 3'd0; a_widx <= 8'd0; a_busy <= 1'b0; end // ---- ST_A: drive SampleNTT, store 256 coeffs per pair ---- if (st == ST_A) begin // mark busy once this pair's request accepted if (snt_valid && snt_ready) begin snt_valid <= 1'b0; a_busy <= 1'b1; end // store each output coefficient only while busy (ignore stale last coeff from prior poly) if (a_busy && snt_vo && snt_ack) begin bank_a[(a_slot*256 + a_widx) & ((1<0 after last end // Slot complete when ntt_core returns to DONE if (ntt_done) begin if (n_slot + 3'd1 < {1'b0, k_r, 1'b0}) begin n_slot <= n_slot + 3'd1; n_widx <= 8'd0; n_pending <= 1'b1; // wait one cycle for core IDLE end else begin n_slot <= n_slot + 3'd1; // == 2K -> ST_DONE end end // Kick next slot's load once core is back IDLE (re-prime) if (n_pending && ntt_ready && !ntt_done) begin n_ridx <= 9'd0; n_valid <= 1'b0; n_loading <= 1'b1; n_pending <= 1'b0; end end // Arm M stage when N finishes: start first (i=0,j=0) poly_mul load if (st == ST_N && st_next == ST_M) begin m_i <= 2'd0; m_j <= 2'd0; m_ld <= 9'd0; m_oidx <= 8'd0; m_loading <= 1'b1; m_pending <= 1'b0; pm_valid <= 1'b1; end // ---- ST_M: t_hat[i] = e_hat[i] + sum_j A[i][j] o s_hat[j] ---- if (st == ST_M) begin // LOAD: stream 256 (A,shat) pairs into poly_mul if (m_loading && pm_valid && pm_ready) begin if (m_ld == 9'd255) begin pm_valid <= 1'b0; // last pair presented m_loading <= 1'b0; m_ld <= 9'd0; m_oidx <= 8'd0; end else begin m_ld <= m_ld + 9'd1; end end // ACCUMULATE: each product coeff += e_hat (j==0) or running t_hat if (pm_vo) begin // bank_t relative slot = m_tslot - slot_t_rt = m_i bank_t[(m_i*256 + m_oidx) & ((1< DONE) if (m_i + 2'd1 < k_r) m_pending <= 1'b1; end end else begin m_oidx <= m_oidx + 8'd1; end end // Start next (i,j) poly_mul load once core is IDLE again if (m_pending && pm_ready && !pm_vo) begin pm_valid <= 1'b1; m_loading <= 1'b1; m_ld <= 9'd0; m_oidx <= 8'd0; m_pending <= 1'b0; end end // Arm E stage when M finishes 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, s_hat -> dk_pke, 1 byte/cycle ---- if (st == ST_E && !e_done) begin if (e_poly < {1'b0, k_r, 1'b0}) begin // write current byte (e_byte_d) to target BRAM if (!e_is_dk) begin ek_we <= 1'b1; ek_wa <= e_boff[10:0] + {9'd0, e_byte}; ek_wd <= e_byte_d; end else begin dkp_we <= 1'b1; dkp_wa <= e_boff[10:0] + {9'd0, e_byte}; dkp_wd <= e_byte_d; end // 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_byte <= e_byte + 2'd1; end end else begin // 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 end // Arm H stage when E finishes if (st == ST_E && st_next == ST_H) begin h_blk <= 3'd0; h_byte <= 8'd0; h_phase <= 2'd0; // assemble 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/9/12 pre-padded blocks) ---- if (st == ST_H) begin case (h_phase) // 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 // 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 end // feed: hold valid until accepted (mb_ready drops) 2'd1: begin if (h_mbvalid && !h_mbready) begin h_mbvalid <= 1'b0; h_mblast <= 1'b0; h_phase <= 2'd2; // wait permute end end // wait permute done: ready again (more blocks) or digest valid (last) 2'd2: begin if (h_vo) begin hek_r <= h_hash[255:0]; h_phase <= 2'd3; // done end else if (h_mbready) begin h_blk <= h_blk + 3'd1; h_phase <= 2'd0; // assemble next block end end default: ; // 2'd3 done: hold endcase end end end endmodule