FallenSigh
6db3c7cc5e
Phase-1 (d1 output) lacked the 'need_more' guard that phase-2 (d2) had, so when the 256th accepted coefficient was a d1 whose group then advanced, the FSM could fire one extra valid_o after last_o for certain seeds (e.g. KAT count=0 rho, seed i=0/j=1 emitted 257 pulses). In mlkem_top KeyGen this trailing pulse leaked into the next poly's index 0, shifting the stream. Fix: gate phase-1 d1 output with 'd1_acc_r && need_more' (mirrors phase-2). Applied to both sample_ntt_sync and sample_ntt_sync_shared. Standalone TB had a coverage blind spot (stopped reading at 256, never checked valid_o stayed low after last_o). Added a regression assertion: counts spurious post-last_o pulses and fails if any. Verified the assertion catches the bug on a reverted-fix copy (3 spurious) and passes on the fix. Verified: - 40-seed audit (sync) + 24-seed audit (shared): all exactly 256 pulses, last_o@256, zero post-last pulses. - Verilator vs hashlib oracle: 1536/1536 (no real coeff dropped). - Full framework regression: 4334/4334. - mlkem_top KeyGen Stage 2c: 2048/2048 A_hat+s+e coeffs exact.
418 lines
17 KiB
Verilog
418 lines
17 KiB
Verilog
// 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.
|
||
//
|
||
// 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
|
||
//
|
||
// --- 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
|
||
|
||
// 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).
|
||
// ================================================================
|
||
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 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 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
|