// mlkem_top.v - ML-KEM-512 KeyGen top-level integration (K=2, eta1=3). // // 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 K = 2, // ML-KEM-512 parameter ETA1 = 3 ) ( input clk, input rst_n, 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 [9:0] dbg_byte_idx_i, output [7:0] dbg_byte_o, // Debug taps for hash outputs output [255:0] dbg_rho_o, output [255:0] dbg_sigma_o ); localparam Q = `Q; // 3329 // ================================================================ // 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] // ================================================================ 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; reg [11:0] polymem [0:NUM_SLOTS*256-1]; // Debug readback (registered for timing) reg [11:0] dbg_coeff_r; 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 reg [7:0] ek_mem [0:EK_BYTES-1]; reg [7:0] dkp_mem [0:DK_BYTES-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; // ================================================================ // 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_mem, dkp_mem localparam ST_DONE = 4'd15; 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 // 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; 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, 8'(K), d_i}; // data_i[263:256]=K, [255:0]=d sha3_top u_sha3 ( .clk(clk), .rst_n(rst_n), .mode(2'b00), // G = SHA3-512 .data_i(g_data), .valid_i(sha3_valid), .ready_o(sha3_ready), .hash_o(sha3_hash), .valid_o(sha3_vo), .ready_i(sha3_ack), .mb_en(1'b0), .mb_block_i(1088'b0), .mb_valid_i(1'b0), .mb_last_i(1'b0), .mb_ready_o() ); // ---- 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 #(.K(K)) u_snt ( .clk(clk), .rst_n(rst_n), .rho_i(rho_r), .k_i(3'(K)), .i_idx(a_i), .j_idx(a_j), .valid_i(snt_valid), .ready_o(snt_ready), .coeff_o(snt_coeff), .valid_o(snt_vo), .ready_i(snt_ack), .last_o(snt_last) ); // ---- 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 u_cbd ( .clk(clk), .rst_n(rst_n), .seed_i(sigma_r), .nonce_i(c_nonce), .eta_i(2'(ETA1)), .valid_i(cbd_valid), .ready_o(cbd_ready), .coeff_o(cbd_coeff), .valid_o(cbd_vo), .ready_i(cbd_ack), .last_o(cbd_last) ); // 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 [2:0] n_slot; // 0..2K (4 polys) 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 ntt_ready; wire [11:0] ntt_coeff; wire ntt_vo; wire ntt_done; wire [11:0] ntt_in = polymem[n_slot_addr*256 + n_ridx[7:0]]; 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 [1:0] m_i; // row 0..K reg [1:0] m_j; // col 0..K 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 [2:0] e_poly; // which source poly: 0,1 = t_hat0,t_hat1 -> ek // 2,3 = s_hat0,s_hat1 -> 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; // 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}]; // 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 [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 [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] reg pm_valid; wire pm_ready; wire [11:0] pm_coeff; wire pm_vo; wire [11:0] pm_a_in = polymem[m_aslot*256 + m_ld[7:0]]; wire [11:0] pm_b_in = polymem[m_sslot*256 + m_ld[7: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] wire [11:0] m_acc_src = (m_j == 2'd0) ? polymem[m_eslot*256 + m_oidx] : polymem[m_tslot*256 + m_oidx]; // (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 >= K*K) st_next = ST_C; ST_C: if (c_poly >= 2*K) st_next = ST_N; ST_N: if (n_slot >= 2*K) st_next = ST_M; ST_M: if (m_i >= K) st_next = ST_E; ST_E: if (e_done) 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; 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 <= 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_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_rho <= 10'd0; e_done <= 1'b0; end else begin st <= st_next; // Kick off G when entering ST_G if (st == ST_IDLE && start_i) begin 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 <= 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 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; a_widx <= 8'd0; a_busy <= 1'b0; // start next SampleNTT if more pairs remain if (a_pair + 3'd1 < K*K) snt_valid <= 1'b1; end else begin a_widx <= a_widx + 8'd1; end end end // Arm C stage when A finishes if (st == ST_A && st_next == ST_C) begin cbd_valid <= 1'b1; cbd_ack <= 1'b1; c_poly <= 3'd0; c_widx <= 8'd0; c_busy <= 1'b0; end // ---- ST_C: drive CBD, store 256 mod-q coeffs per poly ---- if (st == ST_C) begin if (cbd_valid && cbd_ready) begin cbd_valid <= 1'b0; c_busy <= 1'b1; end if (c_busy && cbd_vo && cbd_ack) begin polymem[c_slot*256 + c_widx] <= cbd_modq; if (cbd_last) begin c_poly <= c_poly + 3'd1; c_widx <= 8'd0; c_busy <= 1'b0; if (c_poly + 3'd1 < 2*K) cbd_valid <= 1'b1; end else begin c_widx <= c_widx + 8'd1; end end end // Arm N stage when C finishes: start NTT on slot S0 if (st == ST_C && st_next == ST_N) begin n_slot <= 3'd0; n_ridx <= 9'd0; n_widx <= 8'd0; n_valid <= 1'b1; // begin loading first poly n_pending <= 1'b0; end // ---- ST_N: forward NTT each of S0,S1,E0,E1 in place ---- if (st == ST_N) begin // LOAD phase: stream 256 coeffs into ntt_core if (n_valid && ntt_ready) begin if (n_ridx == 9'd255) begin n_valid <= 1'b0; // last coeff presented this cycle n_ridx <= 9'd0; end else begin n_ridx <= n_ridx + 9'd1; end end // OUTPUT phase: collect 256 results, write back to same slot if (ntt_vo) begin polymem[n_slot_addr*256 + n_widx] <= ntt_coeff; n_widx <= n_widx + 8'd1; // wraps 255->0 after last end // Slot complete when ntt_core returns to DONE if (ntt_done) begin if (n_slot + 3'd1 < 2*K) 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 if (n_pending && ntt_ready && !ntt_done) begin n_valid <= 1'b1; n_ridx <= 9'd0; 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 polymem[m_tslot*256 + m_oidx] <= m_accq; if (m_oidx == 8'd255) begin // finished this (i,j) term; advance if (m_j + 2'd1 < K) begin m_j <= m_j + 2'd1; m_pending <= 1'b1; // next term, same row end else begin m_j <= 2'd0; m_i <= m_i + 2'd1; // next row (or == K -> DONE) if (m_i + 2'd1 < K) 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_rho <= 10'd0; e_done <= 1'b0; end // ---- 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 ek_mem[e_boff] <= e_b0; ek_mem[e_boff + 1] <= e_b1; ek_mem[e_boff + 2] <= e_b2; end else begin dkp_mem[e_boff] <= e_b0; dkp_mem[e_boff + 1] <= e_b1; dkp_mem[e_boff + 2] <= e_b2; end if (e_pair == 8'd127) begin e_pair <= 8'd0; e_poly <= e_poly + 3'd1; // next poly (or ->4 = 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]; if (e_rho == 10'd31) e_done <= 1'b1; else e_rho <= e_rho + 10'd1; end end end end endmodule