FallenSigh
a9e50ebc0c
Add ST_M stage: t_hat[i] = e_hat[i] + sum_j A_hat[i][j] o s_hat[j] via poly_mul_sync + inline mod-add accumulation. Per (i,j): stream 256 (A,shat) pairs into poly_mul, then accumulate 256 products into T_i (seeded from E_i when j==0, else running T_i). m_pending waits for poly_mul IDLE between terms. Verified vs ml-kem-r golden: 512/512 t_hat coeffs exact (19885 cyc).
460 lines
18 KiB
Verilog
460 lines
18 KiB
Verilog
// 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 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;
|
|
|
|
// ================================================================
|
|
// 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_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)
|
|
|
|
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_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;
|
|
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
|
|
end
|
|
end
|
|
|
|
endmodule
|