fallensigh/mlkem-sync
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
b2bf798454854d33dece8244d33efa3670ecb0ef
mlkem-sync/sync_rtl/top/mlkem_top.v
FallenSigh b2bf798454 feat(mlkem_top): parameterize K in {2,3,4} (ML-KEM 512/768/1024) 
Generalize KeyGen from K=2-hardcoded to compile-time parameter K:
- eta1 derived (3 for K=2, else 2); slot layout SLOT_S/E/T = K*K+{0,K,2K},
  NUM_SLOTS = K*K+3K; SAW=5 slot-addr width.
- A-stage: explicit a_i/a_j row-major counters (slot = i*K+j) instead of
  K=2 bit-tricks. C/N stages: parametric slot bases, 2K polys.
- M-stage: m_i/m_j widened to 3-bit (must reach K=4); slots i*K+j etc.
- E-stage: 2K polys, e_is_dk split, rho offset 384*K.
- H(ek): H_NBLK=ceil((EK_BYTES+1)/136), H_LAST padding generalized;
  h_blk 4-bit. Byte mems sized EK_BYTES/DK_BYTES.
- Widen dbg_byte_idx_i to [10:0] (ek up to 1568B for K=4).

Parametric TB (tb_mlkem_kg_katK, KP generic + CASE plusarg). Verified
byte-exact vs NIST KAT:
  K=2 (512):  cases 0..4  ek 800B  / dk 1632B
  K=3 (768):  cases 0..2  ek 1184B / dk 2400B  (~36k cyc)
  K=4 (1024): cases 0..2  ek 1568B / dk 3168B  (~54k cyc)
run_tb.sh top runs all three parameter sets.
2026-06-28 02:59:58 +08:00

672 lines
27 KiB
Verilog
Raw Blame History

// 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=2, 768=3, 1024=4 (eta1 derived)
) (
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 [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
// FIPS 203: eta1 = 3 for ML-KEM-512 (K=2), else 2 (K=3/4).
localparam ETA1 = (K == 2) ? 3 : 2;
// ================================================================
// Polynomial storage, generalized for K in {2,3,4}.
// 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 .. +K-1 (s[i] then overwritten by NTT)
// e_hat[i] : slots SLOT_E .. +K-1
// t_hat[i] : slots SLOT_T .. +K-1
// NUM_SLOTS = K*K + 3*K (10 / 24 / 28 for K=2/3/4)
// ================================================================
localparam SLOT_S = K*K; // s_hat base slot
localparam SLOT_E = K*K + K; // e_hat base slot
localparam SLOT_T = K*K + 2*K; // t_hat base slot
localparam NUM_SLOTS = K*K + 3*K;
localparam SAW = 5; // slot-address width (>=clog2(28))
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 and dk_pke byte memories (byteEncode12 output).
// ek = 384*K + 32 bytes (== KAT pk), dk_pke = 384*K bytes (== KAT sk prefix)
localparam EK_BYTES = 384*K + 32; // 800 / 1184 / 1568
localparam DK_BYTES = 384*K; // 768 / 1152 / 1536
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;
// full dk = dk_pke(DK_BYTES) || ek(EK_BYTES) || H(ek)(32) || z(32)
localparam DK_EK_END = DK_BYTES + EK_BYTES; // ek region end
localparam DK_HEK_END = DK_EK_END + 32; // H(ek) region end
reg [7:0] dbg_dk_r;
always @(posedge clk) begin
if (dbg_dk_idx_i < DK_BYTES[11:0])
dbg_dk_r <= dkp_mem[dbg_dk_idx_i];
else if (dbg_dk_idx_i < DK_EK_END[11:0])
dbg_dk_r <= ek_mem[dbg_dk_idx_i - DK_BYTES[11:0]];
else if (dbg_dk_idx_i < DK_HEK_END[11:0])
dbg_dk_r <= hek_r[(dbg_dk_idx_i - DK_EK_END[11:0])*8 +: 8];
else
dbg_dk_r <= z_i[(dbg_dk_idx_i - DK_HEK_END[11:0])*8 +: 8];
end
assign dbg_dk_o = dbg_dk_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_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 + 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 < K) ? (SLOT_S + c_poly)
: (SLOT_E + (c_poly - K));
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()
);
// ---- second sha3_top dedicated to multi-block H(ek) (SHA3-256, 800B->6 blk) ----
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..135 byte within block being assembled
reg [1:0] h_phase; // 0=assemble 1=feed 2=wait-perm 3=done
sha3_top u_sha3_h (
.clk(clk), .rst_n(rst_n),
.mode(2'b01), // unused in mb mode
.data_i(512'b0),
.valid_i(1'b0),
.ready_o(),
.hash_o(h_hash),
.valid_o(h_vo),
.ready_i(h_ack),
.mb_en(1'b1),
.mb_block_i(h_block_r),
.mb_valid_i(h_mbvalid),
.mb_last_i(h_mblast),
.mb_ready_o(h_mbready)
);
// SHA3-256 over EK_BYTES-byte ek: rate=136. Padded length = H_NBLK*136.
// pad: byte EK_BYTES = 0x06 (domain + first pad bit), last byte |= 0x80.
localparam H_NBLK = (EK_BYTES + 136) / 136; // ceil((EK_BYTES+1)/136): 6/9/12
localparam H_LAST = H_NBLK*136 - 1; // index of final padded byte
// byte b (0..135) of block blk: global g = blk*136 + b
function [7:0] h_padbyte(input [3:0] blk, input [7:0] b);
integer g;
begin
g = blk*136 + b;
if (g < EK_BYTES) h_padbyte = ek_mem[g];
else if (g == H_LAST && g == EK_BYTES) h_padbyte = 8'h86; // 0x06|0x80
else if (g == EK_BYTES) h_padbyte = 8'h06;
else if (g == H_LAST) h_padbyte = 8'h80;
else h_padbyte = 8'h00;
end
endfunction
// ---- 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[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)
);
// ---- 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 [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 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 [SAW-1:0] n_slot_addr = SLOT_S + n_slot; // s_hat then e_hat contiguous
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 [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 [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 >= K);
wire [4:0] e_pidx = e_is_dk ? (e_poly - K) : e_poly; // index within target
wire [SAW-1:0] e_slot = e_is_dk ? (SLOT_S + e_pidx) : (SLOT_T + e_pidx);
// 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 [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
wire [SAW-1:0] m_aslot = m_i*K + m_j; // A_hat[i][j] slot = i*K + j
wire [SAW-1:0] m_sslot = SLOT_S + m_j; // s_hat[j]
wire [SAW-1:0] m_eslot = SLOT_E + m_i; // e_hat[i]
wire [SAW-1:0] m_tslot = SLOT_T + m_i; // 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_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;
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_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;
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;
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 <= 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
polymem[a_slot*256 + a_widx] <= snt_coeff;
if (snt_last) begin
// finished this poly; advance (i,j) in row-major order
a_pair <= a_pair + 5'd1;
a_widx <= 8'd0;
a_busy <= 1'b0;
if (a_j + 3'd1 < K) begin
a_j <= a_j + 3'd1;
end else begin
a_j <= 3'd0;
a_i <= a_i + 3'd1;
end
// start next SampleNTT if more pairs remain
if (a_pair + 5'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 < 2*K) begin
// pack current coeff-pair (3 bytes): [0,K)=ek, [K,2K)=dk_pke
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;
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 + 5'd1; // next poly (or ->2K = rho phase)
end else begin
e_pair <= e_pair + 8'd1;
end
end else begin
// rho copy: ek_mem[384*K + r] = rho byte r (r = 0..31)
ek_mem[12'(384*K) + 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
// 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
end
// ---- ST_H: H(ek) via multi-block SHA3-256 (6 pre-padded blocks) ----
if (st == ST_H) begin
case (h_phase)
// assemble 136 bytes of block h_blk into h_block_r
2'd0: begin
h_block_r[h_byte*8 +: 8] <= h_padbyte(h_blk, h_byte);
if (h_byte == 8'd135) begin
h_byte <= 8'd0;
h_mbvalid <= 1'b1;
h_mblast <= (h_blk == H_NBLK-1);
h_phase <= 2'd1; // feed
end else begin
h_byte <= h_byte + 8'd1;
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
Reference in New Issue View Git Blame Copy Permalink