// sample_ntt_sync.v - Synchronous SampleNTT for ML-KEM A matrix generation // // Generates one k×k polynomial (256 coefficients) via SHAKE-128 XOF // rejection sampling from seed rho || j || i. // // FIPS 202/203 conformant SHAKE-128 squeeze: // - Absorb: S = Keccak-p(padded(rho || j || i)) // - Squeeze the full 1344-bit (168-byte) rate as 56 groups of 3 bytes, // each group g read from S[24*g +: 24]; extract d1[11:0], d2[23:12] // - Accept d if d < Q=3329 // - Only after all 56 groups of the block are consumed: S = Keccak-p(S) // (re-permute once per rate block, NOT per 3-byte group) // - Repeat until 256 coefficients collected (~3 blocks) // // 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 `include "sync_rtl/common/defines.vh" /* verilator lint_off UNUSEDPARAM */ (* use_dsp = "no" *) module sample_ntt_sync #(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 ); // ================================================================ // Local parameters // ================================================================ localparam Q = `Q; // 3329 // SHAKE-128 rate = 1344 bits = 168 bytes = 56 groups of 3 bytes. // After consuming all 56 groups of a block, re-permute the state. localparam GRP_MAX = 6'd55; // ================================================================ // 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; // ================================================================ // Group pointer: which 3-byte group within the 1344-bit rate // block is currently being consumed (0..GRP_MAX). Re-permute when // it would exceed GRP_MAX. // ================================================================ reg [5:0] grp_ptr_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 the current 3-byte group of squeeze state // ================================================================ // group g occupies bits [24*g +: 24]: c0=byte0, c1=byte1, c2=byte2 // d1 = {c1[3:0], c0} // d2 = {c2, c1[7:4]} wire [10:0] grp_bit_off; assign grp_bit_off = grp_ptr_r * 11'd24; // 0..1320, +24 ≤ 1344 (rate) wire [23:0] grp_bits; assign grp_bits = squeeze_state_r[ grp_bit_off +: 24 ]; wire [7:0] c0, c1, c2; assign c0 = grp_bits[7:0]; assign c1 = grp_bits[15:8]; assign c2 = grp_bits[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 core instantiation // ================================================================ wire kc_valid_i; wire [1599:0] kc_state_i; /* verilator lint_off UNUSEDSIGNAL */ wire kc_ready_o; /* verilator lint_on UNUSEDSIGNAL */ wire [1599:0] kc_state_o; wire kc_valid_o; keccak_core #(.ROUNDS(24)) u_keccak ( .clk (clk), .rst_n (rst_n), .state_i (kc_state_i), .valid_i (kc_valid_i), .ready_o (kc_ready_o), .state_o (kc_state_o), .valid_o (kc_valid_o), .ready_i (1'b1) ); // kc_valid_i: asserted on the ABSORB load, and for one cycle when the // squeeze block is exhausted (SQUEEZE → WAIT) to re-permute the state. assign kc_valid_i = (state_next == ST_ABSORB) || (state_r == ST_SQUEEZE && state_next == ST_WAIT); // 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); // grp_done: current 3-byte group fully consumed (phase 2 resolved): // d2 rejected, no longer need coefficients, or d2 was just accepted out. wire grp_done = (state_r == ST_SQUEEZE) && (sq_phase_r == 2'd2) && (!d2_acc_r || !need_more || (valid_o_r && ready_i)); // ================================================================ // 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 // A group is fully consumed once phase 2 resolves (d2 output, // rejected, or no longer needed). Then either advance to the // next group in this block, re-permute (block exhausted), or // finish. if (grp_done) begin if (!need_more) state_next = ST_DONE; else if (grp_ptr_r < GRP_MAX) state_next = ST_SQUEEZE; // next group, no re-permute else state_next = ST_WAIT; // block exhausted: re-permute end 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; grp_ptr_r <= 6'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. A fresh block // starts at group 0 (ABSORB load or WAIT re-permute result). // --------------------------------------------------------- if (kc_valid_o) begin squeeze_state_r <= kc_state_o; grp_ptr_r <= 6'd0; 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 && need_more) 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 or no longer needed: 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 // Group consumed but block not exhausted: advance to the // next 3-byte group within the same block (no re-permute). if (grp_done && need_more && grp_ptr_r < GRP_MAX) begin grp_ptr_r <= grp_ptr_r + 6'd1; sq_phase_r <= 2'd0; end 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