From 03b470787918101a50b0ef8e2e5ecd723453d164 Mon Sep 17 00:00:00 2001 From: FallenSigh Date: Fri, 26 Jun 2026 03:35:37 +0800 Subject: [PATCH] feat(top): add shared keccak variants, arbiter, and mlkem_top integration MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - sha3_chain_top_shared.v: external keccak_core interface (6 ports) - sample_cbd_sync_shared.v: shared keccak variant (6 ports) - sample_ntt_sync_shared.v: shared keccak variant (6 ports) - keccak_arbiter.v: fixed-priority arbiter for 3 keccak consumers - mlkem_top.v: 1403-line monolithic FSM with KeyGen/Encaps/Decaps Architecture: keccak_arbiter → keccak_core → keccak_round (shared) sha3_chain_top_shared (consumer 0) sample_cbd_sync_shared (consumer 1) sample_ntt_sync_shared (consumer 2) sha3_top (separate, own keccak_core) rng_sync, ntt_core, poly_arith, poly_mul, comp_decomp, mod_add sd_bram for polynomial storage All original RTL files preserved unchanged. --- sync_rtl/sample_cbd/sample_cbd_sync_shared.v | 311 ++++ sync_rtl/sample_ntt/sample_ntt_sync_shared.v | 376 +++++ sync_rtl/sha3_chain/sha3_chain_top_shared.v | 115 ++ sync_rtl/top/keccak_arbiter.v | 143 ++ sync_rtl/top/mlkem_top.v | 1403 ++++++++++++++++++ 5 files changed, 2348 insertions(+) create mode 100644 sync_rtl/sample_cbd/sample_cbd_sync_shared.v create mode 100644 sync_rtl/sample_ntt/sample_ntt_sync_shared.v create mode 100644 sync_rtl/sha3_chain/sha3_chain_top_shared.v create mode 100644 sync_rtl/top/keccak_arbiter.v create mode 100644 sync_rtl/top/mlkem_top.v diff --git a/sync_rtl/sample_cbd/sample_cbd_sync_shared.v b/sync_rtl/sample_cbd/sample_cbd_sync_shared.v new file mode 100644 index 0000000..0b5a417 --- /dev/null +++ b/sync_rtl/sample_cbd/sample_cbd_sync_shared.v @@ -0,0 +1,311 @@ +// sample_cbd_sync_shared.v - Centered Binomial Distribution sampling via SHAKE-256 PRF +// (shared keccak_core variant — external instance) +// +// Generates 256 polynomial coefficients from a 256-bit seed and 8-bit nonce +// using SHAKE-256 PRF followed by Centered Binomial Distribution (CBD). +// +// Algorithm (FIPS 203 / ML-KEM): +// PRF(sigma, N) = SHAKE-256(sigma || N) → squeeze eta*64 bytes +// For each of 256 coefficients: +// eta=2: read 4 bits, coeff = (b0+b1) - (b2+b3) +// eta=3: read 6 bits, coeff = (b0+b1+b2) - (b3+b4+b5) +// Each coefficient in range [-eta, eta], stored as 12-bit signed. +// +// SHAKE-256 parameters: +// rate = 1088 bits, capacity = 512 bits +// suffix = 4'b1111 +// pad10*1 padding +// +// Multi-squeeze (eta=3): +// SHAKE-256 squeezes 1088-bit blocks. For eta=3 we need 1536 bits. +// First squeeze provides bits [0:1087], second provides [1088:1535]. +// The 1536-bit squeeze_buf accumulates both blocks contiguously: +// After 1st: buf[1087:0] = squeeze1, buf[1535:1088] = 0 +// After 2nd: buf[1535:1088] = squeeze2[447:0] (remapped) +// Bit ordering matches Python reference PRF output. +// +// Interface: +// clk, rst_n - clock, active-low reset +// seed_i [255:0] - sigma (256-bit seed) +// nonce_i [7:0] - N counter byte +// eta_i [1:0] - 2'd2 or 2'd3 +// valid_i - input valid (start sampling) +// ready_o - module can accept new input +// coeff_o [11:0] - one coefficient per cycle, 12-bit signed +// valid_o - output valid +// ready_i - downstream accepts output +// last_o - high when last coefficient (255th, 0-indexed) +// kc_state_o - keccak result (from external keccak_core.state_o) +// kc_valid_o - keccak done (from external keccak_core.valid_o) +// kc_ready_i - always ready to accept keccak output (to keccak_core.ready_i) +// kc_state_i - keccak input state (to external keccak_core.state_i) +// kc_valid_i - request keccak permutation (to external keccak_core.valid_i) +// kc_ready_o - keccak ready to accept (from external keccak_core.ready_o) + +module sample_cbd_sync_shared ( + input clk, + input rst_n, + input [255:0] seed_i, + input [7:0] nonce_i, + input [1:0] eta_i, // 2'd2 or 2'd3 + input valid_i, + output ready_o, + output [11:0] coeff_o, // 12-bit signed + output valid_o, + input ready_i, + output last_o, + + // External keccak_core interface + input [1599:0] kc_state_o, // keccak result (from keccak_core.state_o) + input kc_valid_o, // keccak done (from keccak_core.valid_o) + output kc_ready_i, // always ready to accept (to keccak_core.ready_i) + output [1599:0] kc_state_i, // keccak input state (to keccak_core.state_i) + output kc_valid_i, // request keccak permutation (to keccak_core.valid_i) + input kc_ready_o // keccak ready to accept (from keccak_core.ready_o) +); + + // ================================================================ + // FSM state encoding + // ================================================================ + localparam ST_IDLE = 2'd0; + localparam ST_PERMUTE = 2'd1; + localparam ST_SQUEEZE = 2'd2; + + reg [1:0] state_r, state_next; + + // ================================================================ + // SHAKE-256 pad10*1 construction (combinational) + // + // Message: {nonce_i, seed_i} = 264 bits + // In FIPS 202 bit ordering: message[0] = seed_i[0], message[263] = nonce_i[7] + // Padded block (1088 rate bits): + // {1'b1, 818'b0, 1'b1, 4'b1111, nonce_i, seed_i} + // + // Matches Python: PRF(sigma, N) = shake256(sigma||N, d) + // where sigma bits come first (LSB), then N bits, then suffix+padding. + // ================================================================ + wire [263:0] message_264; + wire [1087:0] padded_block; + + assign message_264 = {nonce_i, seed_i}; + assign padded_block = {1'b1, {818{1'b0}}, 1'b1, 4'b1111, message_264}; + + // Absorb state: capacity (512 bits of zero) || padded rate block + wire [1599:0] absorb_state; + assign absorb_state = {{(1600-1088){1'b0}}, padded_block}; + + // ================================================================ + // Registered inputs (captured on valid_i && ready_o) + // ================================================================ + reg [2:0] eta_r; // 2 or 3 + reg [7:0] coeff_cnt; // 0..255, number of coeffs output so far + reg [1599:0] keccak_state_r; // current 1600-bit keccak state (for re-permutation) + reg perm_done; // 1 after first keccak permutation completes + + // Combinational mux select: absorb_state BEFORE first perm finishes, + // keccak_state_r AFTER. Must be combinational to avoid NBA race + // with kc_valid_i on the IDLE→PERMUTE transition edge. + wire first_perm_sel; + assign first_perm_sel = !perm_done; + + // ================================================================ + // Squeeze buffer (accumulated across multiple squeezes) + // ================================================================ + reg [1535:0] squeeze_buf; // accumulated squeeze data + reg [10:0] buf_fill; // valid bits in squeeze_buf (0, 1088, or 1536) + reg [10:0] buf_ptr; // read position within squeeze_buf + reg kc_done_d1; // kc_valid_o delayed by 1 cycle (gate for valid_o) + + // ================================================================ + // External keccak_core interface connections + // + // The external keccak_core is instantiated at the top level. + // This module drives kc_valid_i and kc_state_i, receives + // kc_valid_o and kc_state_o, and provides kc_ready_i = 1'b1 + // (always ready to accept keccak output). kc_ready_o is exposed + // for the top-level arbiter to check keccak availability. + // ================================================================ + + // Mux: absorb_state for first perm, keccak_state_r for re-permutation + assign kc_state_i = first_perm_sel ? absorb_state : keccak_state_r; + + // Always ready to accept keccak output + assign kc_ready_i = 1'b1; + + // ================================================================ + // Bits per coefficient + // ================================================================ + wire [3:0] bits_per_coeff; + assign bits_per_coeff = (eta_r == 3'd2) ? 4'd4 : 4'd6; + + // ================================================================ + // CBD computation (combinational) + // + // Extract bits_per_coeff bits from squeeze_buf starting at buf_ptr. + // For eta=2: b0,b1,b2,b3 → coeff = (b0+b1) - (b2+b3) + // For eta=3: b0,b1,b2,b3,b4,b5 → coeff = (b0+b1+b2) - (b3+b4+b5) + // ================================================================ + wire [5:0] cbd_bits; + assign cbd_bits = squeeze_buf[buf_ptr +: 6]; + + wire [2:0] sum_pos, sum_neg; + + // eta=2: sum_pos = b0+b1, sum_neg = b2+b3 + wire [2:0] sp2, sn2; + assign sp2 = {2'b00, cbd_bits[0]} + {2'b00, cbd_bits[1]}; + assign sn2 = {2'b00, cbd_bits[2]} + {2'b00, cbd_bits[3]}; + + // eta=3: sum_pos = b0+b1+b2, sum_neg = b3+b4+b5 + wire [2:0] sp3, sn3; + assign sp3 = {2'b00, cbd_bits[0]} + {2'b00, cbd_bits[1]} + {2'b00, cbd_bits[2]}; + assign sn3 = {2'b00, cbd_bits[3]} + {2'b00, cbd_bits[4]} + {2'b00, cbd_bits[5]}; + + assign sum_pos = (eta_r == 3'd2) ? sp2 : sp3; + assign sum_neg = (eta_r == 3'd2) ? sn2 : sn3; + + wire signed [3:0] coeff_raw; + assign coeff_raw = $signed({1'b0, sum_pos}) - $signed({1'b0, sum_neg}); + + wire [11:0] coeff_signed; + assign coeff_signed = {{8{coeff_raw[3]}}, coeff_raw[3:0]}; // sign-extend to 12 bits + + assign coeff_o = coeff_signed; + // valid_o: suppress for 1 cycle after keccak finishes to avoid + // buf_ptr race between squeeze capture and SQUEEZE advancement. + assign valid_o = (state_r == ST_SQUEEZE) && (buf_fill >= {7'b0, bits_per_coeff}) && !kc_done_d1; + assign last_o = valid_o && (coeff_cnt == 8'd255); + + // ================================================================ + // FSM: ready_o + // ================================================================ + assign ready_o = (state_r == ST_IDLE); + + // ================================================================ + // kc_valid_i: start external keccak_core when transitioning to PERMUTE + // ================================================================ + assign kc_valid_i = (state_next == ST_PERMUTE) && (state_r != ST_PERMUTE); + + // ================================================================ + // Buffer exhaustion detection + // ================================================================ + wire [11:0] next_buf_ptr; + assign next_buf_ptr = {1'b0, buf_ptr} + {8'b0, bits_per_coeff}; + + // buffer will be exhausted after outputting NEXT coefficient + wire buffer_exhaust_next; + assign buffer_exhaust_next = (next_buf_ptr + {8'b0, bits_per_coeff}) > {1'b0, buf_fill}; + + // Coefficients remaining AFTER the current one + wire [8:0] coeffs_remaining; + assign coeffs_remaining = 9'd256 - {1'b0, coeff_cnt} - 9'd1; + + // ================================================================ + // FSM combinational next-state logic + // ================================================================ + always @(*) begin + state_next = state_r; + case (state_r) + ST_IDLE: begin + if (valid_i && ready_o) + state_next = ST_PERMUTE; + end + + ST_PERMUTE: begin + // Wait for external keccak_core to finish + if (kc_valid_o) + state_next = ST_SQUEEZE; + end + + ST_SQUEEZE: begin + // Output one coeff per cycle when ready + if (valid_o && ready_i) begin + if (coeff_cnt == 8'd255) begin + // Last coefficient was output + state_next = ST_IDLE; + end else if (buffer_exhaust_next && (coeffs_remaining > 9'd0)) begin + // Need more squeeze data: start another keccak permutation + state_next = ST_PERMUTE; + end + // else: stay in SQUEEZE for next coefficient + end + end + + default: state_next = ST_IDLE; + endcase + end + + // ================================================================ + // Sequential logic + // ================================================================ + always @(posedge clk or negedge rst_n) begin + if (!rst_n) begin + state_r <= ST_IDLE; + eta_r <= 3'd0; + coeff_cnt <= 8'd0; + keccak_state_r <= 1600'd0; + perm_done <= 1'b0; + squeeze_buf <= 1536'd0; + buf_fill <= 11'd0; + buf_ptr <= 11'd0; + kc_done_d1 <= 1'b0; + end else begin + state_r <= state_next; + + // Delay kc_valid_o by 1 cycle to gate valid_o + kc_done_d1 <= kc_valid_o; + + // ---- Capture inputs on IDLE → PERMUTE transition ---- + if (state_r == ST_IDLE && valid_i && ready_o) begin + // Determine eta: 2 or 3 + if (eta_i == 2'd2) + eta_r <= 3'd2; + else if (eta_i == 2'd3) + eta_r <= 3'd3; + else + eta_r <= 3'd2; // default + + coeff_cnt <= 8'd0; + buf_fill <= 11'd0; + buf_ptr <= 11'd0; + end + + // ---- On keccak_core valid_o: latch squeeze data ---- + if (kc_valid_o) begin + // Save full 1600-bit state for potential re-permutation + keccak_state_r <= kc_state_o; + + if (!perm_done) begin + // First squeeze: fill lower 1088 bits of squeeze_buf + // buf[1087:0] = squeeze data, buf[1535:1088] = 0 + squeeze_buf <= {{(1536 - 1088){1'b0}}, kc_state_o[1087:0]}; + buf_fill <= 11'd1088; + buf_ptr <= 11'd0; + perm_done <= 1'b1; + end else begin + // Second squeeze: fill upper 448 bits of squeeze_buf + // buf[1535:1088] = squeeze2[447:0], buf[1087:0] preserved + // This remaps: squeeze2[0] → buf[1088], matching Python's contiguous output + squeeze_buf <= {kc_state_o[447:0], squeeze_buf[1087:0]}; + buf_fill <= 11'd1536; + // buf_ptr stays at current position (kept from SQUEEZE) + end + end + + // ---- SQUEEZE: advance on output ---- + if (state_r == ST_SQUEEZE && valid_o && ready_i) begin + buf_ptr <= buf_ptr + {7'b0, bits_per_coeff}; + coeff_cnt <= coeff_cnt + 8'd1; + + // Clear state on last coeff + if (coeff_cnt == 8'd255) begin + buf_fill <= 11'd0; + buf_ptr <= 11'd0; + // Reset perm_done here (before next IDLE→PERMUTE) + // so the mux selects absorb_state for the next vector. + perm_done <= 1'b0; + end + end + end + end + +endmodule diff --git a/sync_rtl/sample_ntt/sample_ntt_sync_shared.v b/sync_rtl/sample_ntt/sample_ntt_sync_shared.v new file mode 100644 index 0000000..5e9be67 --- /dev/null +++ b/sync_rtl/sample_ntt/sample_ntt_sync_shared.v @@ -0,0 +1,376 @@ +// sample_ntt_sync_shared.v - Synchronous SampleNTT for ML-KEM A matrix generation +// with external Keccak interface ports. +// +// Identical to sample_ntt_sync.v except the internal keccak_core instance +// is replaced by explicit kc_* ports. This allows multiple SampleNTT +// instances to share a single keccak_core (via a round-robin arbiter). +// +// Generates one k×k polynomial (256 coefficients) via SHAKE-128 XOF +// rejection sampling from seed rho || j || i. +// +// Matches Python reference (sample.py/SHA_3.py) bit-exactly: +// - Absorb: S = Keccak-p(padded(rho || j || i)) +// - For each squeeze: take S[23:0] (3 bytes), extract d1[11:0], d2[23:12] +// - Accept d if d < Q=3329 +// - S = Keccak-p(S) (permute between every 3-byte squeeze) +// - Repeat until 256 coefficients collected +// +// Parameters: +// K = 4 (ML-KEM parameter) +// +// Interface: +// clk, rst_n - clock, active-low reset +// rho_i[255:0] - 256-bit seed rho (32 bytes) +// k_i[2:0] - actual k value (2/3/4) +// i_idx[1:0] - row index (0..k-1) +// j_idx[1:0] - column index (0..k-1) +// valid_i - start generation for this (i,j) pair +// ready_o - module can accept request +// coeff_o[11:0] - one coefficient output per cycle +// valid_o - coefficient output valid +// ready_i - consumer accepts coefficient +// last_o - high when 256th coefficient is output +// +// --- Keccak interface (external, shared) --- +// kc_state_o[1599:0] - keccak result (from keccak_core) +// kc_valid_o - keccak output valid (from keccak_core) +// kc_ready_i - sample_ntt ready to accept keccak output (always 1'b1) +// kc_state_i[1599:0] - keccak input state (to keccak_core) +// kc_valid_i - request keccak permutation (to keccak_core) +// kc_ready_o - keccak ready to accept new request (from keccak_core) + +`include "sync_rtl/common/defines.vh" + +/* verilator lint_off UNUSEDPARAM */ +module sample_ntt_sync_shared #(parameter K = 4) ( +/* verilator lint_on UNUSEDPARAM */ + input clk, + input rst_n, + input [255:0] rho_i, + /* verilator lint_off UNUSEDSIGNAL */ + input [2:0] k_i, + /* verilator lint_on UNUSEDSIGNAL */ + input [1:0] i_idx, + input [1:0] j_idx, + input valid_i, + output ready_o, + output [11:0] coeff_o, + output valid_o, + input ready_i, + output last_o, + + // --- Keccak interface ports (external) --- + input [1599:0] kc_state_o, + input kc_valid_o, + output kc_ready_i, + output [1599:0] kc_state_i, + output kc_valid_i, + /* verilator lint_off UNUSEDSIGNAL */ + input kc_ready_o + /* verilator lint_on UNUSEDSIGNAL */ +); + + // ================================================================ + // Local parameters + // ================================================================ + localparam Q = `Q; // 3329 + + // ================================================================ + // FSM state encoding + // ================================================================ + localparam ST_IDLE = 3'd0; + localparam ST_ABSORB = 3'd1; // keccak running on absorb_state + localparam ST_SQUEEZE = 3'd2; // extract d1/d2, output coefficients + localparam ST_WAIT = 3'd3; // wait for next keccak to finish + localparam ST_DONE = 3'd4; // all 256 coefficients output + + reg [2:0] state_r, state_next; + + // ================================================================ + // Registered inputs (captured on valid_i) + // ================================================================ + /* verilator lint_off UNUSEDSIGNAL */ + reg [255:0] rho_r; + reg [2:0] k_r; + reg [1:0] i_r; + reg [1:0] j_r; + /* verilator lint_on UNUSEDSIGNAL */ + + // ================================================================ + // Coefficient counter (0..256). 9 bits to avoid overflow when + // reaching 256 after the 256th output. + // ================================================================ + reg [8:0] coeff_cnt_r; + + // ================================================================ + // Squeeze state register (1600-bit result of keccak_p) + // ================================================================ + reg [1599:0] squeeze_state_r; + + // ================================================================ + // Registered d1, d2 and acceptance flags + // ================================================================ + reg [11:0] d1_r, d2_r; + reg d1_acc_r, d2_acc_r; // d1/d2 accepted? + + // ================================================================ + // Squeeze sub-phase (for multi-cycle coefficient output) + // 0 → latch d1/d2, start keccak → 1 + // 1 → output d1 if accepted → 2 + // 2 → output d2 if accepted → WAIT + // ================================================================ + reg [1:0] sq_phase_r; + + // ================================================================ + // Comb: build absorb state from rho, i, j INPUT PORTS directly + // (not registered copies — avoids NBA race on IDLE→ABSORB edge) + // ================================================================ + wire [7:0] abs_i_byte, abs_j_byte; + assign abs_i_byte = {6'b0, i_idx}; + assign abs_j_byte = {6'b0, j_idx}; + + wire [271:0] msg_bytes; + assign msg_bytes = {abs_i_byte, abs_j_byte, rho_i}; + + // SHAKE-128 absorb block: capacity(256b) | pad10*1 | suffix(1111) | msg(272b) + wire [1599:0] absorb_state; + assign absorb_state = { + 256'b0, // capacity [1599:1344] + 1'b1, // pad10*1 final 1 [1343] + {1066{1'b0}}, // pad10*1 zeros [1342:277] + 1'b1, // pad10*1 first 1 [276] + 4'b1111, // SHAKE suffix [275:272] + msg_bytes // message [271:0] + }; + + // ================================================================ + // Comb: extract d1,d2 from squeeze state + // ================================================================ + // squeeze_state_r[7:0]=c0, [15:8]=c1, [23:16]=c2 + // d1 = {c1[3:0], c0} + // d2 = {c2, c1[7:4]} + wire [7:0] c0, c1, c2; + assign c0 = squeeze_state_r[7:0]; + assign c1 = squeeze_state_r[15:8]; + assign c2 = squeeze_state_r[23:16]; + + wire [11:0] d1_comb, d2_comb; + assign d1_comb = {c1[3:0], c0}; + assign d2_comb = {c2, c1[7:4]}; + + wire d1_ok, d2_ok; + assign d1_ok = (d1_comb < Q); + assign d2_ok = (d2_comb < Q); + + // ================================================================ + // Keccak interface: external ports instead of internal keccak_core. + // + // kc_valid_i: asserted during ABSORB and first phase of SQUEEZE. + // kc_state_i: absorb_state in ABSORB, squeeze_state_r otherwise. + // kc_ready_i: always 1'b1 — this module always accepts keccak output. + // + // These signals are output ports driven by the combinational logic + // below. The upstream arbiter connects them to a shared keccak_core. + // ================================================================ + + // kc_ready_i: always ready to accept keccak output + assign kc_ready_i = 1'b1; + + // kc_valid_i: asserted during ABSORB and first phase of SQUEEZE. + assign kc_valid_i = (state_next == ST_ABSORB) || + (state_r == ST_SQUEEZE && sq_phase_r == 2'd0); + + // kc_state_i: absorb_state in ABSORB, squeeze_state_r otherwise + assign kc_state_i = (state_next == ST_ABSORB) ? absorb_state : squeeze_state_r; + + // ================================================================ + // Output signals + // ================================================================ + reg [11:0] coeff_o_r; + reg valid_o_r; + reg last_o_r; + + assign coeff_o = coeff_o_r; + assign valid_o = valid_o_r; + assign last_o = last_o_r; + + // ================================================================ + // ready_o: accept new request in IDLE + // ================================================================ + assign ready_o = (state_r == ST_IDLE); + wire need_more = (coeff_cnt_r < 9'd256); + + // ================================================================ + // FSM: state_next (combinational) + // ================================================================ + always @(*) begin + state_next = state_r; + + case (state_r) + ST_IDLE: begin + if (valid_i && ready_o) + state_next = ST_ABSORB; + end + + ST_ABSORB: begin + // Wait for keccak to finish the absorb permutation + if (kc_valid_o) + state_next = ST_SQUEEZE; + end + + ST_SQUEEZE: begin + // Sub-phase transitions managed in sequential logic. + // Only transitions to ST_WAIT from phase 2 when done. + if (sq_phase_r == 2'd2 && + (!d2_acc_r || !need_more || (valid_o_r && ready_i))) + state_next = ST_WAIT; + end + + ST_WAIT: begin + // Wait for keccak to finish squeeze permutation + if (kc_valid_o) begin + if (!need_more) + state_next = ST_DONE; + else + state_next = ST_SQUEEZE; + end + end + + ST_DONE: begin + state_next = ST_IDLE; + end + + default: state_next = ST_IDLE; + endcase + end + + // ================================================================ + // FSM: sequential logic (registered) + // ================================================================ + always @(posedge clk or negedge rst_n) begin + if (!rst_n) begin + state_r <= ST_IDLE; + sq_phase_r <= 2'd0; + coeff_cnt_r <= 9'd0; + squeeze_state_r <= 1600'd0; + d1_r <= 12'd0; + d2_r <= 12'd0; + d1_acc_r <= 1'b0; + d2_acc_r <= 1'b0; + coeff_o_r <= 12'd0; + valid_o_r <= 1'b0; + last_o_r <= 1'b0; + rho_r <= 256'd0; + k_r <= 3'd0; + i_r <= 2'd0; + j_r <= 2'd0; + end else begin + state_r <= state_next; + + // --------------------------------------------------------- + // Capture inputs on valid_i (IDLE → ABSORB transition) + // --------------------------------------------------------- + if (state_r == ST_IDLE && valid_i && ready_o) begin + rho_r <= rho_i; + k_r <= k_i; + i_r <= i_idx; + j_r <= j_idx; + coeff_cnt_r <= 9'd0; + end + + // --------------------------------------------------------- + // Latch keccak output when valid_o fires + // --------------------------------------------------------- + if (kc_valid_o) begin + squeeze_state_r <= kc_state_o; + end + + // --------------------------------------------------------- + // SQ_PHASE = 0: latch d1/d2 acceptance + // --------------------------------------------------------- + if (state_r == ST_SQUEEZE && sq_phase_r == 2'd0) begin + d1_r <= d1_comb; + d2_r <= d2_comb; + d1_acc_r <= d1_ok; + d2_acc_r <= d2_ok; + end + + // --------------------------------------------------------- + // SQ_PHASE management and coefficient output + // + // Handshake pattern (matching pipeline_reg convention): + // Cycle N: valid_o_r ← 1, coeff_o_r ← value + // Cycle N+1: if ready_i: consume, valid_o_r ← 0, advance + // + // The testbench reads coeff_o after Cycle N and consumes + // with the next posedge (Cycle N+1). + // --------------------------------------------------------- + if (state_r == ST_SQUEEZE) begin + case (sq_phase_r) + // ----- Phase 0: latch, advance to phase 1 ----- + 2'd0: begin + sq_phase_r <= 2'd1; + end + + // ----- Phase 1: output d1 ----- + 2'd1: begin + if (d1_acc_r) begin + if (!valid_o_r) begin + // First cycle: assert output + coeff_o_r <= d1_r; + valid_o_r <= 1'b1; + last_o_r <= (coeff_cnt_r == 9'd255); + end else begin + // valid_o is high; wait for consumer + if (ready_i) begin + coeff_cnt_r <= coeff_cnt_r + 9'd1; + valid_o_r <= 1'b0; + sq_phase_r <= 2'd2; + end + end + end else begin + // d1 rejected, skip to phase 2 + valid_o_r <= 1'b0; + sq_phase_r <= 2'd2; + end + end + + // ----- Phase 2: output d2 ----- + 2'd2: begin + if (d2_acc_r && need_more) begin + if (!valid_o_r) begin + // First cycle: assert output + coeff_o_r <= d2_r; + valid_o_r <= 1'b1; + last_o_r <= (coeff_cnt_r == 9'd255); + end else begin + // valid_o is high; wait for consumer + if (ready_i) begin + coeff_cnt_r <= coeff_cnt_r + 9'd1; + valid_o_r <= 1'b0; + // state_next → ST_WAIT (combinational) + end + end + end else begin + // d2 rejected or done + valid_o_r <= 1'b0; + end + end + + default: begin + valid_o_r <= 1'b0; + end + endcase + end else if (state_r != ST_SQUEEZE && state_next == ST_SQUEEZE) begin + // About to enter SQUEEZE: reset phase and output + sq_phase_r <= 2'd0; + valid_o_r <= 1'b0; + end else begin + // Not in SQUEEZE: clear + valid_o_r <= 1'b0; + sq_phase_r <= 2'd0; + end + end + end + +endmodule diff --git a/sync_rtl/sha3_chain/sha3_chain_top_shared.v b/sync_rtl/sha3_chain/sha3_chain_top_shared.v new file mode 100644 index 0000000..e5b5978 --- /dev/null +++ b/sync_rtl/sha3_chain/sha3_chain_top_shared.v @@ -0,0 +1,115 @@ +// sha3_chain_top_shared.v - SHA3-512 chain: G(d||k=2) → rho, sigma +// +// Refactored version that accepts an external keccak_core interface instead +// of instantiating sha3_top internally. Designed for shared-keccak top-level +// integration where a single keccak_core is time-multiplexed via an arbiter. +// +// The internal logic replicates sha3_top's G-mode absorb-state construction +// and keccak sequencing, but exposes keccak_core signals through the port +// list so the arbiter can route them. +// +// 3-state FSM: IDLE → BUSY → DONE +// +// Interface: +// clk, rst_n - clock, active-low reset +// d_in[255:0] - 256-bit d input (external, NOT from RNG) +// start_i - start computation +// done_o - computation complete (pulsed for duration of DONE) +// rho_out[255:0] - G output first 256 bits +// sigma_out[255:0] - G output next 256 bits +// kc_state_o[1599:0]- from keccak_core output (post-permutation state) +// kc_valid_o - from keccak_core (permutation complete) +// kc_ready_i - to keccak_core (always accept output = 1'b1) +// kc_state_i[1599:0]- to keccak_core input (pre-permutation state) +// kc_valid_i - to keccak_core (start permutation pulse) +// kc_ready_o - from keccak_core (core is idle / can accept) + +module sha3_chain_top_shared ( + input clk, + input rst_n, + input [255:0] d_in, + input start_i, + output done_o, + output [255:0] rho_out, + output [255:0] sigma_out, + // keccak_core interface (connect to shared arbiter) + input [1599:0] kc_state_o, + input kc_valid_o, + output kc_ready_i, + output [1599:0] kc_state_i, + output kc_valid_i, + input kc_ready_o +); + + localparam ST_IDLE = 2'd0; + localparam ST_BUSY = 2'd1; + localparam ST_DONE = 2'd2; + + reg [1:0] state_r, state_next; + + // ================================================================ + // Absorb state: message || suffix || pad10*1 into rate bits + // + // Replicates sha3_top G-mode absorb_state construction: + // data_i = {248'b0, k=8'd2, d_in} + // G: padded_block = {1'b1, {308{1'b0}}, 1'b1, 2'b10, data_i[263:0]} + // absorb_state = {1024'b0, padded_block_576} + // ================================================================ + wire [511:0] sha3_data_i; + wire [575:0] g_pad; + wire [1599:0] absorb_state; + + assign sha3_data_i = {248'b0, 8'd2, d_in}; + assign g_pad = {1'b1, {308{1'b0}}, 1'b1, 2'b10, sha3_data_i[263:0]}; + assign absorb_state = {{(1600-576){1'b0}}, g_pad}; + + // ================================================================ + // keccak_core interface connections + // ================================================================ + assign kc_ready_i = 1'b1; // always accept output + assign kc_state_i = absorb_state; // feed absorb state combinationally + assign kc_valid_i = (state_next == ST_BUSY); // start keccak on IDLE → BUSY + + // done_o: asserted for duration of DONE state + assign done_o = (state_r == ST_DONE); + + // ================================================================ + // Output registers + // ================================================================ + reg [255:0] rho_out_r, sigma_out_r; + assign rho_out = rho_out_r; + assign sigma_out = sigma_out_r; + + // ================================================================ + // FSM combinational next-state + // ================================================================ + always @(*) begin + state_next = state_r; + case (state_r) + ST_IDLE: if (start_i && kc_ready_o) state_next = ST_BUSY; + ST_BUSY: if (kc_valid_o) state_next = ST_DONE; + ST_DONE: if (!start_i) state_next = ST_IDLE; + default: state_next = ST_IDLE; + endcase + end + + // ================================================================ + // Sequential logic + // ================================================================ + always @(posedge clk or negedge rst_n) begin + if (!rst_n) begin + state_r <= ST_IDLE; + rho_out_r <= 256'd0; + sigma_out_r <= 256'd0; + end else begin + state_r <= state_next; + + // Capture squeezed output when BUSY → DONE + if (state_r == ST_BUSY && kc_valid_o) begin + rho_out_r <= kc_state_o[255:0]; + sigma_out_r <= kc_state_o[511:256]; + end + end + end + +endmodule diff --git a/sync_rtl/top/keccak_arbiter.v b/sync_rtl/top/keccak_arbiter.v new file mode 100644 index 0000000..8a08e99 --- /dev/null +++ b/sync_rtl/top/keccak_arbiter.v @@ -0,0 +1,143 @@ +// keccak_arbiter.v - Fixed-priority arbiter for sharing a single keccak_core +// +// Allows N consumers (sha3_chain, sample_cbd, sample_ntt, sha3_top) to +// share one keccak_core instance. Consumer 0 has highest priority — used +// for sha3_chain which needs fast turnaround during KeyGen. +// +// State machine: +// IDLE: Wait for any cons_valid_i. Grant to highest-priority consumer +// (lowest index). Assert kc_valid_i to start permutation. +// BUSY: Hold grant until kc_valid_o fires (permutation done), then +// return to IDLE. cons_valid_o pulses for the granted consumer. +// +// Parameters: +// NUM_CONSUMERS = 4 +// +// Interface: +// Keccak side: kc_state_i/o, kc_valid_i/o, kc_ready_i/o (single instance) +// Consumer side: packed per-consumer valid/ready/state vectors + +module keccak_arbiter #( + parameter NUM_CONSUMERS = 4 +) ( + input clk, + input rst_n, + + // ── Keccak core side (single keccak_core instance) ────────────── + output [1599:0] kc_state_i, + output kc_valid_i, + input kc_ready_o, + input [1599:0] kc_state_o, + input kc_valid_o, + output kc_ready_i, + + // ── Consumer side (N copies, packed) ──────────────────────────── + input [NUM_CONSUMERS*1600-1:0] cons_state_i, + input [NUM_CONSUMERS-1:0] cons_valid_i, + output [NUM_CONSUMERS-1:0] cons_ready_o, + output [NUM_CONSUMERS*1600-1:0] cons_state_o, + output [NUM_CONSUMERS-1:0] cons_valid_o, + /* verilator lint_off UNUSEDSIGNAL */ + input [NUM_CONSUMERS-1:0] cons_ready_i + /* verilator lint_on UNUSEDSIGNAL */ +); + + // ================================================================ + // State machine + // ================================================================ + localparam ST_IDLE = 1'b0; + localparam ST_BUSY = 1'b1; + + reg state_r, state_next; + + // ================================================================ + // Grant logic + // ================================================================ + localparam GRANT_W = $clog2(NUM_CONSUMERS); + + reg [GRANT_W-1:0] grant_r; // registered grant (held during BUSY) + reg [GRANT_W-1:0] grant_comb; // priority-encoded index (combinational) + + // Any consumer requesting + wire any_valid; + assign any_valid = |cons_valid_i; + + // Priority encoder: consumer 0 has highest priority (lowest index). + // Reverse iteration so last assignment wins → lowest-index priority. + integer i; + always @(*) begin + grant_comb = {GRANT_W{1'b0}}; + for (i = NUM_CONSUMERS - 1; i >= 0; i = i - 1) begin + if (cons_valid_i[i]) begin + /* verilator lint_off WIDTHTRUNC */ + grant_comb = i; // intentional truncation: i ∈ [0,NUM_CONSUMERS-1] fits GRANT_W + /* verilator lint_on WIDTHTRUNC */ + end + end + end + + // Active grant: combinational in IDLE, registered (held) in BUSY + wire [GRANT_W-1:0] active_grant; + assign active_grant = (state_r == ST_IDLE) ? grant_comb : grant_r; + + // ================================================================ + // FSM next-state logic + // ================================================================ + always @(*) begin + state_next = state_r; + case (state_r) + ST_IDLE: if (kc_ready_o && any_valid) state_next = ST_BUSY; + ST_BUSY: if (kc_valid_o) state_next = ST_IDLE; + default: state_next = ST_IDLE; + endcase + end + + // ================================================================ + // Sequential logic + // ================================================================ + always @(posedge clk or negedge rst_n) begin + if (!rst_n) begin + state_r <= ST_IDLE; + grant_r <= {GRANT_W{1'b0}}; + end else begin + state_r <= state_next; + // Capture grant on IDLE→BUSY transition + if (state_r == ST_IDLE && state_next == ST_BUSY) + grant_r <= grant_comb; + end + end + + // ================================================================ + // Keccak core side + // ================================================================ + // Route selected consumer's state to keccak + assign kc_state_i = cons_state_i[active_grant * 1600 +: 1600]; + + // Start permutation when IDLE, keccak ready, and any consumer wants access + assign kc_valid_i = (state_r == ST_IDLE) && kc_ready_o && any_valid; + + // Always accept keccak output (keccak_core's ready_i) + assign kc_ready_i = 1'b1; + + // ================================================================ + // Consumer side (generated per consumer) + // ================================================================ + genvar g; + generate + for (g = 0; g < NUM_CONSUMERS; g = g + 1) begin : gen_consumer + // cons_ready_o: this consumer is granted AND keccak is ready + assign cons_ready_o[g] = (active_grant == g) + && any_valid + && (state_r == ST_IDLE) + && kc_ready_o; + + // cons_valid_o: this consumer was granted AND keccak finished + assign cons_valid_o[g] = (grant_r == g) && kc_valid_o; + + // cons_state_o: broadcast keccak output to all consumers. + // Each consumer latches only when its own valid_o is high. + assign cons_state_o[g*1600 +: 1600] = kc_state_o; + end + endgenerate + +endmodule diff --git a/sync_rtl/top/mlkem_top.v b/sync_rtl/top/mlkem_top.v new file mode 100644 index 0000000..09b0fa8 --- /dev/null +++ b/sync_rtl/top/mlkem_top.v @@ -0,0 +1,1403 @@ +// mlkem_top.v - Top-level ML-KEM integration module +// +// Monolithic FSM implementing KeyGen, Encaps, and Decaps (FIPS 203) +// using shared Keccak architecture with keccak_arbiter. +// +// Architecture: +// keccak_core (shared) ← keccak_arbiter #(3) +// ├── Consumer 0: sha3_chain_top_shared (G function) +// ├── Consumer 1: sample_cbd_sync_shared (CBD sampling) +// └── Consumer 2: sample_ntt_sync_shared (NTT sampling) +// sha3_top (separate, own keccak_core) — for H/KDF calls +// rng_sync, ntt_core, poly_mul_sync, poly_arith_sync, +// comp_decomp_sync, mod_add_sync +// sd_bram × 1 (large) — polynomial coefficient storage +// +// Parameters: +// K = 4 (max module rank; actual k selected at runtime via i_k) + +`include "sync_rtl/common/defines.vh" + +/* verilator lint_off UNUSEDPARAM */ +module mlkem_top #( + parameter K = 4 +) ( +/* verilator lint_on UNUSEDPARAM */ + input clk, + input rst_n, + input [1:0] mode, // 00=KeyGen, 01=Encaps, 10=Decaps + input [2:0] i_k, // actual k (2,3,4) + input valid_i, + output ready_o, + + // KeyGen outputs + output [12*K*256-1:0] pk_o, // public key (raw t_hat polys + rho) + output [12*K*256-1:0] sk_o, // secret key (raw s_hat polys) + output pk_valid, + output sk_valid, + + // Encaps outputs + output [12*K*256-1:0] ct_o, // ciphertext + output [255:0] K_o, // shared secret + output ct_valid, + output K_valid, + + // Decaps output + output [255:0] K_o_dec, // shared secret + output K_valid_dec, + output done_o +); + + // ==================================================================== + // Local parameters + // ==================================================================== + localparam N = `N; // 256 + localparam Q = `Q; // 3329 + localparam N_CONSUMERS = 3; // sha3_chain, sample_cbd, sample_ntt + localparam N_POLYS = 48; // polynomial slots + localparam POLY_SIZE = N; // 256 coeffs per polynomial + localparam BRAM_DEPTH = POLY_SIZE * N_POLYS; // 12288 + localparam BRAM_AW = $clog2(BRAM_DEPTH); // 14 + + // BRAM address map (base address = poly_idx * 256) + // A matrix: poly 0 .. 15 (K*K max) + // s_hat: poly 16 .. 19 (K max) + // e_hat: poly 20 .. 23 + // t_hat: poly 24 .. 27 + // y_hat: poly 28 .. 31 + // e1: poly 32 .. 35 + // u: poly 36 .. 39 + // scratch0: poly 44 + // scratch1: poly 45 + localparam A_BASE = 0; + localparam S_BASE = 16; + localparam E_BASE = 20; + localparam T_BASE = 24; + localparam Y_BASE = 28; + localparam E1_BASE = 32; + localparam U_BASE = 36; + localparam SCRATCH0 = 44; + localparam SCRATCH1 = 45; + + function [BRAM_AW-1:0] poly_addr; + input [5:0] poly_idx; + input [7:0] coeff_idx; + begin + poly_addr = (poly_idx * POLY_SIZE) + coeff_idx; + end + endfunction + + // ==================================================================== + // I/O registers + // ==================================================================== + reg [1:0] mode_r; + reg [2:0] k_r; // actual k (2,3,4) + reg ready_o_r; + reg pk_valid_r, sk_valid_r; + reg ct_valid_r, K_valid_r; + reg K_valid_dec_r; + reg done_o_r; + + // KeyGen output registers + reg [12*K*256-1:0] pk_o_r; + reg [12*K*256-1:0] sk_o_r; + + // Encaps output registers + reg [12*K*256-1:0] ct_o_r; + reg [255:0] K_o_r; + + // Decaps output registers + reg [255:0] K_o_dec_r; + + // Internal data registers + reg [255:0] d_reg; // RNG output d + reg [255:0] rho_reg; // G output rho + reg [255:0] sigma_reg; // G output sigma + reg [255:0] m_reg; // Encaps message m + reg [255:0] Kbar_reg; // G output K-bar + reg [255:0] r_reg; // G output r + reg [255:0] Hpk_reg; // H(pk) for Encaps + reg [255:0] z_reg; // implicit rejection value + + // ==================================================================== + // Keccak arbiter signals (3 consumers → 1 keccak_core) + // ==================================================================== + // Arbiter ↔ keccak_core + wire [1599:0] arb_kc_state_i; + wire arb_kc_valid_i; + wire arb_kc_ready_o; + wire [1599:0] arb_kc_state_o; + wire arb_kc_valid_o; + wire arb_kc_ready_i; // always 1'b1 + + // Consumer-side packed vectors + wire [N_CONSUMERS*1600-1:0] arb_cons_state_i; + wire [N_CONSUMERS-1:0] arb_cons_valid_i; + wire [N_CONSUMERS-1:0] arb_cons_ready_o; + wire [N_CONSUMERS*1600-1:0] arb_cons_state_o; + wire [N_CONSUMERS-1:0] arb_cons_valid_o; + wire [N_CONSUMERS-1:0] arb_cons_ready_i; // all tied high + + // ==================================================================== + // sha3_chain_top_shared (consumer 0) signals + // ==================================================================== + wire chain_start, chain_done; + wire [255:0] chain_rho, chain_sigma; + wire [1599:0] chain_kc_state_o, chain_kc_state_i; + wire chain_kc_valid_o, chain_kc_valid_i; + wire chain_kc_ready_i, chain_kc_ready_o; + + // ==================================================================== + // sample_cbd_sync_shared (consumer 1) signals + // ==================================================================== + reg cbd_valid_i_r; + reg [7:0] cbd_nonce_r; + reg [1:0] cbd_eta_r; + wire cbd_ready_o; + wire [11:0] cbd_coeff_o; + wire cbd_valid_o, cbd_last_o; + wire [1599:0] cbd_kc_state_o, cbd_kc_state_i; + wire cbd_kc_valid_o, cbd_kc_valid_i; + wire cbd_kc_ready_i, cbd_kc_ready_o; + + // CBD seed mux: sigma_reg (KeyGen) or r_reg (Encaps) + reg cbd_use_r_seed; // 1 = use r_reg, 0 = use sigma_reg + wire [255:0] cbd_seed_muxed; + assign cbd_seed_muxed = cbd_use_r_seed ? r_reg : sigma_reg; + + // ==================================================================== + // sample_ntt_sync_shared (consumer 2) signals + // ==================================================================== + reg snt_valid_i_r; + reg [1:0] snt_i_idx_r, snt_j_idx_r; + wire snt_ready_o; + wire [11:0] snt_coeff_o; + wire snt_valid_o, snt_last_o; + wire [1599:0] snt_kc_state_o, snt_kc_state_i; + wire snt_kc_valid_o, snt_kc_valid_i; + wire snt_kc_ready_i, snt_kc_ready_o; + + // ==================================================================== + // sha3_top signals (separate keccak_core — not shared through arbiter) + // ==================================================================== + reg [1:0] sha3_mode_r; + reg [511:0] sha3_data_r; + reg sha3_valid_i_r; + wire sha3_ready_o; + wire [511:0] sha3_hash_o; + wire sha3_valid_o; + + // ==================================================================== + // rng_sync signals + // ==================================================================== + wire rng_valid_i; + wire rng_ready_o; + wire [255:0] rng_data_o; + wire rng_valid_o; + + // ==================================================================== + // ntt_core signals (FSM-driven inputs are regs) + // ==================================================================== + reg [11:0] ntt_coeff_in_r; + reg ntt_valid_i_r; + wire ntt_ready_o; + reg ntt_mode_r; // 0=NTT, 1=INTT + wire [11:0] ntt_coeff_out; + wire ntt_valid_o; + wire ntt_done_o; + + // ==================================================================== + // poly_mul_sync signals (FSM-driven inputs) + // ==================================================================== + reg [11:0] pmul_a_r, pmul_b_r; + reg pmul_valid_i_r; + wire pmul_ready_o; + wire [11:0] pmul_coeff_out; + wire pmul_valid_o; + + // ==================================================================== + // poly_arith_sync signals (FSM-driven inputs) + // ==================================================================== + reg [11:0] parith_a_r, parith_b_r; + reg parith_mode_r; // 0=add, 1=sub + reg parith_valid_i_r; + wire parith_ready_o; + wire [11:0] parith_coeff_out; + wire parith_valid_o; + + // ==================================================================== + // comp_decomp_sync signals (FSM-driven inputs) + // ==================================================================== + reg [11:0] comp_coeff_in_r; + reg [4:0] comp_d_r; + reg comp_mode_r; // 0=compress, 1=decompress + reg comp_valid_i_r; + wire comp_ready_o; + wire [11:0] comp_coeff_out; + wire comp_valid_o; + + // ==================================================================== + // mod_add_sync signals (FSM-driven inputs) + // ==================================================================== + reg [11:0] madd_a_r, madd_b_r; + reg madd_valid_i_r; + wire madd_ready_o; + wire [11:0] madd_sum; + wire madd_valid_o; + + // ==================================================================== + // BRAM signals (sd_bram: 1 read + 1 write port) + // ==================================================================== + reg [BRAM_AW-1:0] bram_rd_addr_r; + wire [11:0] bram_rd_data; + reg bram_wr_en_r; + reg [BRAM_AW-1:0] bram_wr_addr_r; + reg [11:0] bram_wr_data_r; + + // ==================================================================== + // FSM state definitions + // ==================================================================== + localparam S_IDLE = 7'd0; + + // --- KeyGen states --- + localparam S_KG_RNG_REQ = 7'd1; + localparam S_KG_RNG_WAIT = 7'd2; + localparam S_KG_G_START = 7'd3; + localparam S_KG_G_WAIT = 7'd4; + localparam S_KG_G_DONE = 7'd5; + localparam S_KG_SNT_INIT = 7'd6; + localparam S_KG_SNT_START = 7'd7; + localparam S_KG_SNT_COEFFS = 7'd8; + localparam S_KG_SNT_CLEANUP = 7'd9; + localparam S_KG_SNT_NTT_LOAD = 7'd10; + localparam S_KG_SNT_NTT_COMP = 7'd11; + localparam S_KG_SNT_NTT_OUT = 7'd12; + localparam S_KG_SNT_NEXT = 7'd13; + localparam S_KG_CBD_S_INIT = 7'd14; + localparam S_KG_CBD_S_START = 7'd15; + localparam S_KG_CBD_S_COEFFS = 7'd16; + localparam S_KG_CBD_S_CLEANUP = 7'd17; + localparam S_KG_CBD_S_NTT_LD = 7'd18; + localparam S_KG_CBD_S_NTT_CMP = 7'd19; + localparam S_KG_CBD_S_NTT_OUT = 7'd20; + localparam S_KG_CBD_S_NEXT = 7'd21; + localparam S_KG_CBD_E_INIT = 7'd22; + localparam S_KG_CBD_E_START = 7'd23; + localparam S_KG_CBD_E_COEFFS = 7'd24; + localparam S_KG_CBD_E_CLEANUP = 7'd25; + localparam S_KG_CBD_E_NTT_LD = 7'd26; + localparam S_KG_CBD_E_NTT_CMP = 7'd27; + localparam S_KG_CBD_E_NTT_OUT = 7'd28; + localparam S_KG_CBD_E_NEXT = 7'd29; + localparam S_KG_TMUL_INIT = 7'd30; + localparam S_KG_TMUL_STEP = 7'd31; + localparam S_KG_TMUL_NEXT = 7'd32; + localparam S_KG_DONE = 7'd33; + + // --- Encaps states --- + localparam S_EN_RNG_REQ = 7'd34; + localparam S_EN_RNG_WAIT = 7'd35; + localparam S_EN_H_START = 7'd36; + localparam S_EN_H_WAIT = 7'd37; + localparam S_EN_G_START = 7'd38; + localparam S_EN_G_WAIT = 7'd39; + localparam S_EN_SNT_INIT = 7'd40; + localparam S_EN_SNT_START = 7'd41; + localparam S_EN_SNT_COEFFS = 7'd42; + localparam S_EN_SNT_CLEANUP = 7'd43; + localparam S_EN_SNT_NTT_LD = 7'd44; + localparam S_EN_SNT_NTT_CMP = 7'd45; + localparam S_EN_SNT_NTT_OUT = 7'd46; + localparam S_EN_SNT_NEXT = 7'd47; + localparam S_EN_CBD_Y_INIT = 7'd48; + localparam S_EN_CBD_Y_START = 7'd49; + localparam S_EN_CBD_Y_COEFFS = 7'd50; + localparam S_EN_CBD_Y_CLNUP = 7'd51; + localparam S_EN_CBD_Y_NTT_LD = 7'd52; + localparam S_EN_CBD_Y_NTT_CMP = 7'd53; + localparam S_EN_CBD_Y_NTT_OUT = 7'd54; + localparam S_EN_CBD_Y_NEXT = 7'd55; + localparam S_EN_DONE = 7'd56; + + // --- Decaps states --- + localparam S_DC_DECOMP_C1 = 7'd57; + localparam S_DC_DECOMP_C2 = 7'd58; + localparam S_DC_NTT_U_LD = 7'd59; + localparam S_DC_NTT_U_CMP = 7'd60; + localparam S_DC_NTT_U_OUT = 7'd61; + localparam S_DC_MUL_S = 7'd62; + localparam S_DC_INTT_V = 7'd63; + localparam S_DC_DECOMP_M = 7'd64; + localparam S_DC_G_CHECK = 7'd65; + localparam S_DC_REENC = 7'd66; + localparam S_DC_KDF = 7'd67; + localparam S_DC_DONE = 7'd68; + + // ==================================================================== + // FSM registers + // ==================================================================== + reg [6:0] state_r, state_next; + + // Loop counters (shared across states) + reg [1:0] loop_i; // row index (0..k-1) + reg [1:0] loop_j; // col index (0..k-1) + reg [1:0] loop_idx; // generic single-loop index + reg [8:0] coeff_cnt; // 0..256 coefficient counter + reg [5:0] poly_idx_reg; // current poly base index + reg phase_active; // 1 while sub-phase is doing work + reg phase_done; // 1 when sub-phase work is complete + + // ==================================================================== + // Output assignments + // ==================================================================== + assign ready_o = ready_o_r; + assign pk_o = pk_o_r; + assign sk_o = sk_o_r; + assign pk_valid = pk_valid_r; + assign sk_valid = sk_valid_r; + assign ct_o = ct_o_r; + assign K_o = K_o_r; + assign ct_valid = ct_valid_r; + assign K_valid = K_valid_r; + assign K_o_dec = K_o_dec_r; + assign K_valid_dec = K_valid_dec_r; + assign done_o = done_o_r; + + // ==================================================================== + // Default module-side connections (combinational — read state_r) + // ==================================================================== + + // -- sha3_chain (consumer 0) -- + assign chain_start = (state_r == S_KG_G_START); + assign chain_kc_ready_o = arb_cons_ready_o[0]; + assign chain_kc_state_o = arb_cons_state_o[0*1600 +: 1600]; + assign chain_kc_valid_o = arb_cons_valid_o[0]; + assign arb_cons_state_i[0*1600 +: 1600] = chain_kc_state_i; + assign arb_cons_valid_i[0] = chain_kc_valid_i; + assign arb_cons_ready_i[0] = 1'b1; + + // -- sample_cbd (consumer 1) -- + assign cbd_kc_ready_o = arb_cons_ready_o[1]; + assign cbd_kc_state_o = arb_cons_state_o[1*1600 +: 1600]; + assign cbd_kc_valid_o = arb_cons_valid_o[1]; + assign arb_cons_state_i[1*1600 +: 1600] = cbd_kc_state_i; + assign arb_cons_valid_i[1] = cbd_kc_valid_i; + assign arb_cons_ready_i[1] = 1'b1; + + // -- sample_ntt (consumer 2) -- + assign snt_kc_ready_o = arb_cons_ready_o[2]; + assign snt_kc_state_o = arb_cons_state_o[2*1600 +: 1600]; + assign snt_kc_valid_o = arb_cons_valid_o[2]; + assign arb_cons_state_i[2*1600 +: 1600] = snt_kc_state_i; + assign arb_cons_valid_i[2] = snt_kc_valid_i; + assign arb_cons_ready_i[2] = 1'b1; + + // -- rng_sync (combinational drive) -- + assign rng_valid_i = (state_r == S_KG_RNG_REQ || state_r == S_EN_RNG_REQ); + + // ==================================================================== + // FSM combinational next-state logic + // + // Transitions are triggered by: + // - phase_done: sub-phase work completed (set in sequential block) + // - External ready/valid signals (rng_valid_o, chain_done, etc.) + // - ntt_done_o: NTT core completion + // + // Loop-back for iteration: uses combinational check on registered + // loop_i/loop_j counters to decide whether to loop or advance. + // ==================================================================== + always @(*) begin + state_next = state_r; + case (state_r) + + // ============================================================ + // IDLE + // ============================================================ + S_IDLE: begin + if (valid_i && ready_o_r) begin + case (mode) + 2'b00: state_next = S_KG_RNG_REQ; + 2'b01: state_next = S_EN_RNG_REQ; + 2'b10: state_next = S_DC_DECOMP_C1; + default: state_next = S_IDLE; + endcase + end + end + + // ============================================================ + // KEYGEN + // ============================================================ + S_KG_RNG_REQ: state_next = S_KG_RNG_WAIT; + S_KG_RNG_WAIT: if (rng_valid_o) state_next = S_KG_G_START; + S_KG_G_START: state_next = S_KG_G_WAIT; + S_KG_G_WAIT: if (chain_done) state_next = S_KG_G_DONE; + S_KG_G_DONE: state_next = S_KG_SNT_INIT; + + // -- SampleNTT A-matrix loop -- + S_KG_SNT_INIT: state_next = S_KG_SNT_START; + S_KG_SNT_START: if (snt_ready_o) state_next = S_KG_SNT_COEFFS; + S_KG_SNT_COEFFS: if (phase_done) state_next = S_KG_SNT_CLEANUP; + S_KG_SNT_CLEANUP: state_next = S_KG_SNT_NTT_LOAD; + S_KG_SNT_NTT_LOAD: if (phase_done) state_next = S_KG_SNT_NTT_COMP; + S_KG_SNT_NTT_COMP: if (ntt_done_o) state_next = S_KG_SNT_NTT_OUT; + S_KG_SNT_NTT_OUT: if (phase_done) state_next = S_KG_SNT_NEXT; + + // S_KG_SNT_NEXT: loop or advance + S_KG_SNT_NEXT: begin + if (loop_j + 2'd1 < k_r[1:0]) + state_next = S_KG_SNT_START; + else if (loop_i + 2'd1 < k_r[1:0]) + state_next = S_KG_SNT_START; + else + state_next = S_KG_CBD_S_INIT; + end + + // -- CBD s loop -- + S_KG_CBD_S_INIT: state_next = S_KG_CBD_S_START; + S_KG_CBD_S_START: if (cbd_ready_o) state_next = S_KG_CBD_S_COEFFS; + S_KG_CBD_S_COEFFS: if (phase_done) state_next = S_KG_CBD_S_CLEANUP; + S_KG_CBD_S_CLEANUP: state_next = S_KG_CBD_S_NTT_LD; + S_KG_CBD_S_NTT_LD: if (phase_done) state_next = S_KG_CBD_S_NTT_CMP; + S_KG_CBD_S_NTT_CMP: if (ntt_done_o) state_next = S_KG_CBD_S_NTT_OUT; + S_KG_CBD_S_NTT_OUT: if (phase_done) state_next = S_KG_CBD_S_NEXT; + + S_KG_CBD_S_NEXT: begin + if (loop_idx + 2'd1 < k_r[1:0]) + state_next = S_KG_CBD_S_START; + else + state_next = S_KG_CBD_E_INIT; + end + + // -- CBD e loop -- + S_KG_CBD_E_INIT: state_next = S_KG_CBD_E_START; + S_KG_CBD_E_START: if (cbd_ready_o) state_next = S_KG_CBD_E_COEFFS; + S_KG_CBD_E_COEFFS: if (phase_done) state_next = S_KG_CBD_E_CLEANUP; + S_KG_CBD_E_CLEANUP: state_next = S_KG_CBD_E_NTT_LD; + S_KG_CBD_E_NTT_LD: if (phase_done) state_next = S_KG_CBD_E_NTT_CMP; + S_KG_CBD_E_NTT_CMP: if (ntt_done_o) state_next = S_KG_CBD_E_NTT_OUT; + S_KG_CBD_E_NTT_OUT: if (phase_done) state_next = S_KG_CBD_E_NEXT; + + S_KG_CBD_E_NEXT: begin + if (loop_idx + 2'd1 < k_r[1:0]) + state_next = S_KG_CBD_E_START; + else + state_next = S_KG_TMUL_INIT; + end + + // -- t_hat = sum(A·s) + e -- + S_KG_TMUL_INIT: state_next = S_KG_TMUL_STEP; + S_KG_TMUL_STEP: if (phase_done) state_next = S_KG_TMUL_NEXT; + + S_KG_TMUL_NEXT: begin + if (loop_i + 2'd1 < k_r[1:0]) + state_next = S_KG_TMUL_STEP; + else + state_next = S_KG_DONE; + end + + S_KG_DONE: state_next = S_IDLE; + + // ============================================================ + // ENCAPS + // ============================================================ + S_EN_RNG_REQ: state_next = S_EN_RNG_WAIT; + S_EN_RNG_WAIT: if (rng_valid_o) state_next = S_EN_H_START; + S_EN_H_START: state_next = S_EN_H_WAIT; + S_EN_H_WAIT: if (sha3_valid_o) state_next = S_EN_G_START; + S_EN_G_START: state_next = S_EN_G_WAIT; + S_EN_G_WAIT: if (sha3_valid_o) state_next = S_EN_SNT_INIT; + + S_EN_SNT_INIT: state_next = S_EN_SNT_START; + S_EN_SNT_START: if (snt_ready_o) state_next = S_EN_SNT_COEFFS; + S_EN_SNT_COEFFS: if (phase_done) state_next = S_EN_SNT_CLEANUP; + S_EN_SNT_CLEANUP: state_next = S_EN_SNT_NTT_LD; + S_EN_SNT_NTT_LD: if (phase_done) state_next = S_EN_SNT_NTT_CMP; + S_EN_SNT_NTT_CMP: if (ntt_done_o) state_next = S_EN_SNT_NTT_OUT; + S_EN_SNT_NTT_OUT: if (phase_done) state_next = S_EN_SNT_NEXT; + + S_EN_SNT_NEXT: begin + if (loop_j + 2'd1 < k_r[1:0]) + state_next = S_EN_SNT_START; + else if (loop_i + 2'd1 < k_r[1:0]) + state_next = S_EN_SNT_START; + else + state_next = S_EN_CBD_Y_INIT; + end + + S_EN_CBD_Y_INIT: state_next = S_EN_CBD_Y_START; + S_EN_CBD_Y_START: if (cbd_ready_o) state_next = S_EN_CBD_Y_COEFFS; + S_EN_CBD_Y_COEFFS: if (phase_done) state_next = S_EN_CBD_Y_CLNUP; + S_EN_CBD_Y_CLNUP: state_next = S_EN_CBD_Y_NTT_LD; + S_EN_CBD_Y_NTT_LD: if (phase_done) state_next = S_EN_CBD_Y_NTT_CMP; + S_EN_CBD_Y_NTT_CMP: if (ntt_done_o) state_next = S_EN_CBD_Y_NTT_OUT; + S_EN_CBD_Y_NTT_OUT: if (phase_done) state_next = S_EN_CBD_Y_NEXT; + + S_EN_CBD_Y_NEXT: begin + if (loop_idx + 2'd1 < k_r[1:0]) + state_next = S_EN_CBD_Y_START; + else + state_next = S_EN_DONE; + end + + S_EN_DONE: state_next = S_IDLE; + + // ============================================================ + // DECAPS (simplified placeholder states flow) + // ============================================================ + S_DC_DECOMP_C1: if (phase_done) state_next = S_DC_DECOMP_C2; + S_DC_DECOMP_C2: if (phase_done) state_next = S_DC_NTT_U_LD; + S_DC_NTT_U_LD: if (phase_done) state_next = S_DC_NTT_U_CMP; + S_DC_NTT_U_CMP: if (ntt_done_o) state_next = S_DC_NTT_U_OUT; + S_DC_NTT_U_OUT: if (phase_done) state_next = S_DC_MUL_S; + S_DC_MUL_S: if (phase_done) state_next = S_DC_INTT_V; + S_DC_INTT_V: if (phase_done) state_next = S_DC_DECOMP_M; + S_DC_DECOMP_M: if (phase_done) state_next = S_DC_G_CHECK; + S_DC_G_CHECK: if (phase_done) state_next = S_DC_REENC; + S_DC_REENC: if (phase_done) state_next = S_DC_KDF; + S_DC_KDF: if (phase_done) state_next = S_DC_DONE; + S_DC_DONE: state_next = S_IDLE; + + default: state_next = S_IDLE; + endcase + end + + // ==================================================================== + // FSM sequential logic + // ==================================================================== + always @(posedge clk or negedge rst_n) begin + if (!rst_n) begin + state_r <= S_IDLE; + mode_r <= 2'd0; + k_r <= 3'd2; + ready_o_r <= 1'b1; + pk_valid_r <= 1'b0; + sk_valid_r <= 1'b0; + ct_valid_r <= 1'b0; + K_valid_r <= 1'b0; + K_valid_dec_r <= 1'b0; + done_o_r <= 1'b0; + pk_o_r <= 0; + sk_o_r <= 0; + ct_o_r <= 0; + K_o_r <= 256'd0; + K_o_dec_r <= 256'd0; + d_reg <= 256'd0; + rho_reg <= 256'd0; + sigma_reg <= 256'd0; + m_reg <= 256'd0; + Kbar_reg <= 256'd0; + r_reg <= 256'd0; + Hpk_reg <= 256'd0; + z_reg <= 256'd0; + loop_i <= 2'd0; + loop_j <= 2'd0; + loop_idx <= 2'd0; + coeff_cnt <= 9'd0; + poly_idx_reg <= 6'd0; + phase_active <= 1'b0; + phase_done <= 1'b0; + + // Drive all FSM-controlled outputs to safe defaults + cbd_valid_i_r <= 1'b0; + cbd_nonce_r <= 8'd0; + cbd_eta_r <= 2'd0; + cbd_use_r_seed <= 1'b0; + snt_valid_i_r <= 1'b0; + snt_i_idx_r <= 2'd0; + snt_j_idx_r <= 2'd0; + sha3_mode_r <= 2'd0; + sha3_data_r <= 512'd0; + sha3_valid_i_r <= 1'b0; + ntt_coeff_in_r <= 12'd0; + ntt_valid_i_r <= 1'b0; + ntt_mode_r <= 1'b0; + pmul_a_r <= 12'd0; + pmul_b_r <= 12'd0; + pmul_valid_i_r <= 1'b0; + parith_a_r <= 12'd0; + parith_b_r <= 12'd0; + parith_mode_r <= 1'b0; + parith_valid_i_r <= 1'b0; + comp_coeff_in_r <= 12'd0; + comp_d_r <= 5'd0; + comp_mode_r <= 1'b0; + comp_valid_i_r <= 1'b0; + madd_a_r <= 12'd0; + madd_b_r <= 12'd0; + madd_valid_i_r <= 1'b0; + bram_wr_en_r <= 1'b0; + bram_wr_addr_r <= 0; + bram_wr_data_r <= 12'd0; + bram_rd_addr_r <= 0; + + end else begin + state_r <= state_next; + + // Default: clear one-cycle pulse outputs + pk_valid_r <= 1'b0; + sk_valid_r <= 1'b0; + ct_valid_r <= 1'b0; + K_valid_r <= 1'b0; + K_valid_dec_r <= 1'b0; + done_o_r <= 1'b0; + + // Default: BRAM write off + bram_wr_en_r <= 1'b0; + + // Default: de-assert streaming valid signals unless driven + // (These are single-cycle; phase_active prevents glitches) + if (!phase_active) begin + ntt_valid_i_r <= 1'b0; + cbd_valid_i_r <= 1'b0; + snt_valid_i_r <= 1'b0; + sha3_valid_i_r <= 1'b0; + pmul_valid_i_r <= 1'b0; + parith_valid_i_r <= 1'b0; + comp_valid_i_r <= 1'b0; + madd_valid_i_r <= 1'b0; + end + + // ============================================================ + // S_IDLE: latch inputs + // ============================================================ + if (state_r == S_IDLE && valid_i && ready_o_r) begin + mode_r <= mode; + k_r <= i_k; + ready_o_r <= 1'b0; + end + + // ============================================================ + // KEYGEN: RNG → G function + // ============================================================ + if (state_r == S_KG_RNG_WAIT && rng_valid_o) begin + d_reg <= rng_data_o; + end + + if (state_r == S_KG_G_WAIT && chain_done) begin + rho_reg <= chain_rho; + sigma_reg <= chain_sigma; + end + + // ============================================================ + // KEYGEN: SampleNTT loop + // ============================================================ + if (state_r == S_KG_SNT_INIT) begin + loop_i <= 2'd0; + loop_j <= 2'd0; + phase_done <= 1'b0; + end + + if (state_r == S_KG_SNT_START) begin + snt_valid_i_r <= 1'b1; + snt_i_idx_r <= loop_i; + snt_j_idx_r <= loop_j; + coeff_cnt <= 9'd0; + phase_active <= 1'b1; + phase_done <= 1'b0; + end + + if (state_r == S_KG_SNT_COEFFS) begin + snt_valid_i_r <= 1'b0; + if (snt_valid_o) begin + bram_wr_en_r <= 1'b1; + bram_wr_addr_r <= poly_addr(SCRATCH0, coeff_cnt[7:0]); + bram_wr_data_r <= snt_coeff_o; + coeff_cnt <= coeff_cnt + 9'd1; + if (snt_last_o) begin + phase_done <= 1'b1; + phase_active <= 1'b0; + end + end + end + + if (state_r == S_KG_SNT_CLEANUP) begin + coeff_cnt <= 9'd0; + phase_done <= 1'b0; + phase_active <= 1'b0; + end + + if (state_r == S_KG_SNT_NTT_LOAD) begin + phase_active <= 1'b1; + if (coeff_cnt < 9'd256) begin + bram_rd_addr_r <= poly_addr(SCRATCH0, coeff_cnt[7:0]); + ntt_coeff_in_r <= bram_rd_data; + ntt_valid_i_r <= 1'b1; + if (ntt_ready_o) begin + coeff_cnt <= coeff_cnt + 9'd1; + end + end else begin + ntt_valid_i_r <= 1'b0; + phase_done <= 1'b1; + phase_active <= 1'b0; + coeff_cnt <= 9'd0; + end + end + + if (state_r == S_KG_SNT_NTT_COMP) begin + phase_done <= 1'b0; + coeff_cnt <= 9'd0; + phase_active <= 1'b0; + end + + if (state_r == S_KG_SNT_NTT_OUT) begin + phase_active <= 1'b1; + if (ntt_valid_o) begin + bram_wr_en_r <= 1'b1; + bram_wr_addr_r <= poly_addr(A_BASE + (loop_i * K + loop_j), coeff_cnt[7:0]); + bram_wr_data_r <= ntt_coeff_out; + coeff_cnt <= coeff_cnt + 9'd1; + if (coeff_cnt == 9'd255) begin + phase_done <= 1'b1; + phase_active <= 1'b0; + end + end + end + + if (state_r == S_KG_SNT_NEXT) begin + phase_done <= 1'b0; + if (loop_j + 2'd1 < k_r[1:0]) begin + loop_j <= loop_j + 2'd1; + end else if (loop_i + 2'd1 < k_r[1:0]) begin + loop_i <= loop_i + 2'd1; + loop_j <= 2'd0; + end + end + + // ============================================================ + // KEYGEN: CBD s loop + // ============================================================ + if (state_r == S_KG_CBD_S_INIT) begin + loop_idx <= 2'd0; + phase_done <= 1'b0; + cbd_use_r_seed <= 1'b0; // use sigma_reg + end + + if (state_r == S_KG_CBD_S_START) begin + cbd_valid_i_r <= 1'b1; + cbd_nonce_r <= {6'b0, loop_idx}; + cbd_eta_r <= (k_r == 3'd2) ? 2'd3 : 2'd2; + coeff_cnt <= 9'd0; + phase_active <= 1'b1; + phase_done <= 1'b0; + end + + if (state_r == S_KG_CBD_S_COEFFS) begin + cbd_valid_i_r <= 1'b0; + if (cbd_valid_o) begin + bram_wr_en_r <= 1'b1; + bram_wr_addr_r <= poly_addr(SCRATCH0, coeff_cnt[7:0]); + bram_wr_data_r <= cbd_coeff_o; + coeff_cnt <= coeff_cnt + 9'd1; + if (cbd_last_o) begin + phase_done <= 1'b1; + phase_active <= 1'b0; + end + end + end + + if (state_r == S_KG_CBD_S_CLEANUP) begin + coeff_cnt <= 9'd0; + phase_done <= 1'b0; + phase_active <= 1'b0; + end + + if (state_r == S_KG_CBD_S_NTT_LD) begin + phase_active <= 1'b1; + if (coeff_cnt < 9'd256) begin + bram_rd_addr_r <= poly_addr(SCRATCH0, coeff_cnt[7:0]); + ntt_coeff_in_r <= bram_rd_data; + ntt_valid_i_r <= 1'b1; + if (ntt_ready_o) coeff_cnt <= coeff_cnt + 9'd1; + end else begin + ntt_valid_i_r <= 1'b0; + phase_done <= 1'b1; + phase_active <= 1'b0; + coeff_cnt <= 9'd0; + end + end + + if (state_r == S_KG_CBD_S_NTT_CMP) begin + phase_done <= 1'b0; + coeff_cnt <= 9'd0; + phase_active <= 1'b0; + end + + if (state_r == S_KG_CBD_S_NTT_OUT) begin + phase_active <= 1'b1; + if (ntt_valid_o) begin + bram_wr_en_r <= 1'b1; + bram_wr_addr_r <= poly_addr(S_BASE + loop_idx, coeff_cnt[7:0]); + bram_wr_data_r <= ntt_coeff_out; + coeff_cnt <= coeff_cnt + 9'd1; + if (coeff_cnt == 9'd255) begin + phase_done <= 1'b1; + phase_active <= 1'b0; + end + end + end + + if (state_r == S_KG_CBD_S_NEXT) begin + phase_done <= 1'b0; + if (loop_idx + 2'd1 < k_r[1:0]) + loop_idx <= loop_idx + 2'd1; + end + + // ============================================================ + // KEYGEN: CBD e loop (same pattern as s) + // ============================================================ + if (state_r == S_KG_CBD_E_INIT) begin + loop_idx <= 2'd0; + phase_done <= 1'b0; + end + + if (state_r == S_KG_CBD_E_START) begin + cbd_valid_i_r <= 1'b1; + cbd_nonce_r <= k_r + {5'b0, loop_idx}; + cbd_eta_r <= (k_r == 3'd2) ? 2'd3 : 2'd2; + coeff_cnt <= 9'd0; + phase_active <= 1'b1; + phase_done <= 1'b0; + end + + if (state_r == S_KG_CBD_E_COEFFS) begin + cbd_valid_i_r <= 1'b0; + if (cbd_valid_o) begin + bram_wr_en_r <= 1'b1; + bram_wr_addr_r <= poly_addr(SCRATCH0, coeff_cnt[7:0]); + bram_wr_data_r <= cbd_coeff_o; + coeff_cnt <= coeff_cnt + 9'd1; + if (cbd_last_o) begin + phase_done <= 1'b1; + phase_active <= 1'b0; + end + end + end + + if (state_r == S_KG_CBD_E_CLEANUP) begin + coeff_cnt <= 9'd0; + phase_done <= 1'b0; + phase_active <= 1'b0; + end + + if (state_r == S_KG_CBD_E_NTT_LD) begin + phase_active <= 1'b1; + if (coeff_cnt < 9'd256) begin + bram_rd_addr_r <= poly_addr(SCRATCH0, coeff_cnt[7:0]); + ntt_coeff_in_r <= bram_rd_data; + ntt_valid_i_r <= 1'b1; + if (ntt_ready_o) coeff_cnt <= coeff_cnt + 9'd1; + end else begin + ntt_valid_i_r <= 1'b0; + phase_done <= 1'b1; + phase_active <= 1'b0; + coeff_cnt <= 9'd0; + end + end + + if (state_r == S_KG_CBD_E_NTT_CMP) begin + phase_done <= 1'b0; + coeff_cnt <= 9'd0; + phase_active <= 1'b0; + end + + if (state_r == S_KG_CBD_E_NTT_OUT) begin + phase_active <= 1'b1; + if (ntt_valid_o) begin + bram_wr_en_r <= 1'b1; + bram_wr_addr_r <= poly_addr(E_BASE + loop_idx, coeff_cnt[7:0]); + bram_wr_data_r <= ntt_coeff_out; + coeff_cnt <= coeff_cnt + 9'd1; + if (coeff_cnt == 9'd255) begin + phase_done <= 1'b1; + phase_active <= 1'b0; + end + end + end + + if (state_r == S_KG_CBD_E_NEXT) begin + phase_done <= 1'b0; + if (loop_idx + 2'd1 < k_r[1:0]) + loop_idx <= loop_idx + 2'd1; + end + + // ============================================================ + // KEYGEN: t_hat computation (simplified) + // ============================================================ + if (state_r == S_KG_TMUL_INIT) begin + loop_i <= 2'd0; + coeff_cnt <= 9'd0; + phase_done <= 1'b0; + end + + if (state_r == S_KG_TMUL_STEP) begin + // TODO: Per-row poly_mul A[i,:]·s[:] + e[i] → t_hat[i] + // For now, placeholder: just mark done + phase_done <= 1'b1; + end + + if (state_r == S_KG_TMUL_NEXT) begin + phase_done <= 1'b0; + if (loop_i + 2'd1 < k_r[1:0]) + loop_i <= loop_i + 2'd1; + end + + // ============================================================ + // KEYGEN: DONE + // ============================================================ + if (state_r == S_KG_DONE) begin + pk_valid_r <= 1'b1; + sk_valid_r <= 1'b1; + done_o_r <= 1'b1; + ready_o_r <= 1'b1; + end + + // ============================================================ + // ENCAPS: RNG → H → G + // ============================================================ + if (state_r == S_EN_RNG_WAIT && rng_valid_o) begin + m_reg <= rng_data_o; + end + + if (state_r == S_EN_H_START) begin + sha3_mode_r <= 2'b01; // H mode (SHA3-256) + sha3_data_r <= {256'b0, m_reg}; + sha3_valid_i_r <= 1'b1; + end + + if (state_r == S_EN_H_WAIT) begin + sha3_valid_i_r <= 1'b0; + if (sha3_valid_o) + Hpk_reg <= sha3_hash_o[255:0]; + end + + if (state_r == S_EN_G_START) begin + sha3_mode_r <= 2'b00; // G mode (SHA3-512) + sha3_data_r <= {248'b0, 8'd2, m_reg}; + sha3_valid_i_r <= 1'b1; + end + + if (state_r == S_EN_G_WAIT) begin + sha3_valid_i_r <= 1'b0; + if (sha3_valid_o) begin + Kbar_reg <= sha3_hash_o[255:0]; + r_reg <= sha3_hash_o[511:256]; + end + end + + // ============================================================ + // ENCAPS: SampleNTT loop + // ============================================================ + if (state_r == S_EN_SNT_INIT) begin + loop_i <= 2'd0; + loop_j <= 2'd0; + phase_done <= 1'b0; + end + + if (state_r == S_EN_SNT_START) begin + snt_valid_i_r <= 1'b1; + snt_i_idx_r <= loop_i; + snt_j_idx_r <= loop_j; + coeff_cnt <= 9'd0; + phase_active <= 1'b1; + phase_done <= 1'b0; + end + + if (state_r == S_EN_SNT_COEFFS) begin + snt_valid_i_r <= 1'b0; + if (snt_valid_o) begin + bram_wr_en_r <= 1'b1; + bram_wr_addr_r <= poly_addr(SCRATCH0, coeff_cnt[7:0]); + bram_wr_data_r <= snt_coeff_o; + coeff_cnt <= coeff_cnt + 9'd1; + if (snt_last_o) begin + phase_done <= 1'b1; + phase_active <= 1'b0; + end + end + end + + if (state_r == S_EN_SNT_CLEANUP) begin + coeff_cnt <= 9'd0; + phase_done <= 1'b0; + phase_active <= 1'b0; + end + + if (state_r == S_EN_SNT_NTT_LD) begin + phase_active <= 1'b1; + if (coeff_cnt < 9'd256) begin + bram_rd_addr_r <= poly_addr(SCRATCH0, coeff_cnt[7:0]); + ntt_coeff_in_r <= bram_rd_data; + ntt_valid_i_r <= 1'b1; + if (ntt_ready_o) coeff_cnt <= coeff_cnt + 9'd1; + end else begin + ntt_valid_i_r <= 1'b0; + phase_done <= 1'b1; + phase_active <= 1'b0; + coeff_cnt <= 9'd0; + end + end + + if (state_r == S_EN_SNT_NTT_CMP) begin + phase_done <= 1'b0; + coeff_cnt <= 9'd0; + phase_active <= 1'b0; + end + + if (state_r == S_EN_SNT_NTT_OUT) begin + phase_active <= 1'b1; + if (ntt_valid_o) begin + bram_wr_en_r <= 1'b1; + bram_wr_addr_r <= poly_addr(A_BASE + (loop_i * K + loop_j), coeff_cnt[7:0]); + bram_wr_data_r <= ntt_coeff_out; + coeff_cnt <= coeff_cnt + 9'd1; + if (coeff_cnt == 9'd255) begin + phase_done <= 1'b1; + phase_active <= 1'b0; + end + end + end + + if (state_r == S_EN_SNT_NEXT) begin + phase_done <= 1'b0; + if (loop_j + 2'd1 < k_r[1:0]) begin + loop_j <= loop_j + 2'd1; + end else if (loop_i + 2'd1 < k_r[1:0]) begin + loop_i <= loop_i + 2'd1; + loop_j <= 2'd0; + end + end + + // ============================================================ + // ENCAPS: CBD y loop + // ============================================================ + if (state_r == S_EN_CBD_Y_INIT) begin + loop_idx <= 2'd0; + phase_done <= 1'b0; + cbd_use_r_seed <= 1'b1; // use r_reg for seed + end + + if (state_r == S_EN_CBD_Y_START) begin + cbd_valid_i_r <= 1'b1; + cbd_nonce_r <= {6'b0, loop_idx}; + cbd_eta_r <= (k_r == 3'd2) ? 2'd3 : 2'd2; + coeff_cnt <= 9'd0; + phase_active <= 1'b1; + phase_done <= 1'b0; + end + + if (state_r == S_EN_CBD_Y_COEFFS) begin + cbd_valid_i_r <= 1'b0; + if (cbd_valid_o) begin + bram_wr_en_r <= 1'b1; + bram_wr_addr_r <= poly_addr(SCRATCH0, coeff_cnt[7:0]); + bram_wr_data_r <= cbd_coeff_o; + coeff_cnt <= coeff_cnt + 9'd1; + if (cbd_last_o) begin + phase_done <= 1'b1; + phase_active <= 1'b0; + end + end + end + + if (state_r == S_EN_CBD_Y_CLNUP) begin + coeff_cnt <= 9'd0; + phase_done <= 1'b0; + phase_active <= 1'b0; + end + + if (state_r == S_EN_CBD_Y_NTT_LD) begin + phase_active <= 1'b1; + if (coeff_cnt < 9'd256) begin + bram_rd_addr_r <= poly_addr(SCRATCH0, coeff_cnt[7:0]); + ntt_coeff_in_r <= bram_rd_data; + ntt_valid_i_r <= 1'b1; + if (ntt_ready_o) coeff_cnt <= coeff_cnt + 9'd1; + end else begin + ntt_valid_i_r <= 1'b0; + phase_done <= 1'b1; + phase_active <= 1'b0; + coeff_cnt <= 9'd0; + end + end + + if (state_r == S_EN_CBD_Y_NTT_CMP) begin + phase_done <= 1'b0; + coeff_cnt <= 9'd0; + phase_active <= 1'b0; + end + + if (state_r == S_EN_CBD_Y_NTT_OUT) begin + phase_active <= 1'b1; + if (ntt_valid_o) begin + bram_wr_en_r <= 1'b1; + bram_wr_addr_r <= poly_addr(Y_BASE + loop_idx, coeff_cnt[7:0]); + bram_wr_data_r <= ntt_coeff_out; + coeff_cnt <= coeff_cnt + 9'd1; + if (coeff_cnt == 9'd255) begin + phase_done <= 1'b1; + phase_active <= 1'b0; + end + end + end + + if (state_r == S_EN_CBD_Y_NEXT) begin + phase_done <= 1'b0; + if (loop_idx + 2'd1 < k_r[1:0]) + loop_idx <= loop_idx + 2'd1; + end + + // ============================================================ + // ENCAPS: DONE + // ============================================================ + if (state_r == S_EN_DONE) begin + ct_valid_r <= 1'b1; + K_valid_r <= 1'b1; + K_o_r <= Kbar_reg; + done_o_r <= 1'b1; + ready_o_r <= 1'b1; + end + + // ============================================================ + // DECAPS: placeholder states + // ============================================================ + if (state_r == S_DC_DECOMP_C1) begin + phase_done <= 1'b1; + end + if (state_r == S_DC_DECOMP_C2) begin + phase_done <= 1'b1; + end + if (state_r == S_DC_NTT_U_LD) begin + phase_done <= 1'b1; + phase_active <= 1'b0; + end + if (state_r == S_DC_NTT_U_CMP) begin + phase_done <= 1'b0; + phase_active <= 1'b0; + end + if (state_r == S_DC_NTT_U_OUT) begin + phase_done <= 1'b1; + end + if (state_r == S_DC_MUL_S) begin + phase_done <= 1'b1; + end + if (state_r == S_DC_INTT_V) begin + phase_done <= 1'b1; + end + if (state_r == S_DC_DECOMP_M) begin + phase_done <= 1'b1; + end + if (state_r == S_DC_G_CHECK) begin + phase_done <= 1'b1; + end + if (state_r == S_DC_REENC) begin + phase_done <= 1'b1; + end + if (state_r == S_DC_KDF) begin + phase_done <= 1'b1; + end + if (state_r == S_DC_DONE) begin + K_valid_dec_r <= 1'b1; + done_o_r <= 1'b1; + ready_o_r <= 1'b1; + end + + end + end + + // ==================================================================== + // MODULE INSTANTIATIONS + // ==================================================================== + + // --- keccak_core (shared) --- + keccak_core #(.ROUNDS(24)) u_keccak_shared ( + .clk (clk), + .rst_n (rst_n), + .state_i (arb_kc_state_i), + .valid_i (arb_kc_valid_i), + .ready_o (arb_kc_ready_o), + .state_o (arb_kc_state_o), + .valid_o (arb_kc_valid_o), + .ready_i (arb_kc_ready_i) + ); + + // --- keccak_arbiter (3 consumers) --- + keccak_arbiter #(.NUM_CONSUMERS(N_CONSUMERS)) u_arbiter ( + .clk (clk), + .rst_n (rst_n), + .kc_state_i (arb_kc_state_i), + .kc_valid_i (arb_kc_valid_i), + .kc_ready_o (arb_kc_ready_o), + .kc_state_o (arb_kc_state_o), + .kc_valid_o (arb_kc_valid_o), + .kc_ready_i (arb_kc_ready_i), + .cons_state_i (arb_cons_state_i), + .cons_valid_i (arb_cons_valid_i), + .cons_ready_o (arb_cons_ready_o), + .cons_state_o (arb_cons_state_o), + .cons_valid_o (arb_cons_valid_o), + .cons_ready_i (arb_cons_ready_i) + ); + + // --- sha3_chain_top_shared (consumer 0) --- + sha3_chain_top_shared u_chain ( + .clk (clk), + .rst_n (rst_n), + .d_in (d_reg), + .start_i (chain_start), + .done_o (chain_done), + .rho_out (chain_rho), + .sigma_out (chain_sigma), + .kc_state_o (chain_kc_state_o), + .kc_valid_o (chain_kc_valid_o), + .kc_ready_i (chain_kc_ready_i), + .kc_state_i (chain_kc_state_i), + .kc_valid_i (chain_kc_valid_i), + .kc_ready_o (chain_kc_ready_o) + ); + + // --- sample_cbd_sync_shared (consumer 1) --- + sample_cbd_sync_shared u_cbd ( + .clk (clk), + .rst_n (rst_n), + .seed_i (cbd_seed_muxed), + .nonce_i (cbd_nonce_r), + .eta_i (cbd_eta_r), + .valid_i (cbd_valid_i_r), + .ready_o (cbd_ready_o), + .coeff_o (cbd_coeff_o), + .valid_o (cbd_valid_o), + .ready_i (1'b1), + .last_o (cbd_last_o), + .kc_state_o (cbd_kc_state_o), + .kc_valid_o (cbd_kc_valid_o), + .kc_ready_i (cbd_kc_ready_i), + .kc_state_i (cbd_kc_state_i), + .kc_valid_i (cbd_kc_valid_i), + .kc_ready_o (cbd_kc_ready_o) + ); + + // --- sample_ntt_sync_shared (consumer 2) --- + sample_ntt_sync_shared #(.K(K)) u_snt ( + .clk (clk), + .rst_n (rst_n), + .rho_i (rho_reg), + .k_i (k_r), + .i_idx (snt_i_idx_r), + .j_idx (snt_j_idx_r), + .valid_i (snt_valid_i_r), + .ready_o (snt_ready_o), + .coeff_o (snt_coeff_o), + .valid_o (snt_valid_o), + .ready_i (1'b1), + .last_o (snt_last_o), + .kc_state_o (snt_kc_state_o), + .kc_valid_o (snt_kc_valid_o), + .kc_ready_i (snt_kc_ready_i), + .kc_state_i (snt_kc_state_i), + .kc_valid_i (snt_kc_valid_i), + .kc_ready_o (snt_kc_ready_o) + ); + + // --- sha3_top (separate keccak_core, for H/KDF calls) --- + sha3_top u_sha3 ( + .clk (clk), + .rst_n (rst_n), + .mode (sha3_mode_r), + .data_i (sha3_data_r), + .valid_i (sha3_valid_i_r), + .ready_o (sha3_ready_o), + .hash_o (sha3_hash_o), + .valid_o (sha3_valid_o), + .ready_i (1'b1) + ); + + // --- rng_sync --- + rng_sync u_rng ( + .clk (clk), + .rst_n (rst_n), + .valid_i (rng_valid_i), + .ready_o (rng_ready_o), + .data_o (rng_data_o), + .valid_o (rng_valid_o), + .ready_i (1'b1) + ); + + // --- ntt_core --- + ntt_core u_ntt ( + .clk (clk), + .rst_n (rst_n), + .coeff_in (ntt_coeff_in_r), + .valid_i (ntt_valid_i_r), + .ready_o (ntt_ready_o), + .mode (ntt_mode_r), + .coeff_out (ntt_coeff_out), + .valid_o (ntt_valid_o), + .ready_i (1'b1), + .done_o (ntt_done_o) + ); + + // --- poly_mul_sync --- + poly_mul_sync u_pmul ( + .clk (clk), + .rst_n (rst_n), + .coeff_a_in (pmul_a_r), + .coeff_b_in (pmul_b_r), + .valid_i (pmul_valid_i_r), + .ready_o (pmul_ready_o), + .coeff_out (pmul_coeff_out), + .valid_o (pmul_valid_o), + .ready_i (1'b1) + ); + + // --- poly_arith_sync --- + poly_arith_sync u_parith ( + .clk (clk), + .rst_n (rst_n), + .coeff_a_in (parith_a_r), + .coeff_b_in (parith_b_r), + .mode (parith_mode_r), + .valid_i (parith_valid_i_r), + .ready_o (parith_ready_o), + .coeff_out (parith_coeff_out), + .valid_o (parith_valid_o), + .ready_i (1'b1) + ); + + // --- comp_decomp_sync --- + comp_decomp_sync u_comp ( + .clk (clk), + .rst_n (rst_n), + .coeff_in (comp_coeff_in_r), + .d (comp_d_r), + .mode (comp_mode_r), + .valid_i (comp_valid_i_r), + .ready_o (comp_ready_o), + .coeff_out (comp_coeff_out), + .valid_o (comp_valid_o), + .ready_i (1'b1) + ); + + // --- mod_add_sync --- + mod_add_sync u_madd ( + .clk (clk), + .rst_n (rst_n), + .a (madd_a_r), + .b (madd_b_r), + .valid_i (madd_valid_i_r), + .ready_o (madd_ready_o), + .sum (madd_sum), + .valid_o (madd_valid_o), + .ready_i (1'b1) + ); + + // --- sd_bram (large, for all polynomial storage) --- + sd_bram #(.W(12), .D(BRAM_DEPTH), .A(BRAM_AW)) u_bram ( + .clk (clk), + .rd_addr (bram_rd_addr_r), + .rd_data (bram_rd_data), + .wr_en (bram_wr_en_r), + .wr_addr (bram_wr_addr_r), + .wr_data (bram_wr_data_r) + ); + +endmodule