// 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