FallenSigh
5941fee980
Phase 1 complete — all 4 leaf modules verified: - rng_sync.v: 256-bit Galois LFSR PRNG (10/10 PASS) - sample_cbd_sync.v: CBD sampler with keccak_core PRF (2560/2560 PASS) - sample_ntt_sync.v: SHAKE-128 rejection sampling for A matrix (1536/1536 PASS) - xsim Verilog TBs for sha3 module (tb_sha3_xsim.v, tb_sha3_xsim_simple.v, tb_keccak_core_xsim.v)
363 lines
14 KiB
Verilog
363 lines
14 KiB
Verilog
// 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.
|
||
//
|
||
// 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
|
||
|
||
`include "sync_rtl/common/defines.vh"
|
||
|
||
/* verilator lint_off UNUSEDPARAM */
|
||
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
|
||
|
||
// ================================================================
|
||
// 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 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 during ABSORB and first phase of SQUEEZE.
|
||
// Keccak captures it on the transition (when ready_o=1).
|
||
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
|