feat: add ML-KEM accelerator with AXI slave wrapper
Wrap mlkem_top in an AXI4 burst slave (mlkem_axi_slave) mapped to the freed FFT slot 6 at 0x1f400000 (virtual 0xbf400000). The core runs on sys_clk/2 (25 MHz); the wrapper handles the clock-domain crossing: start toggle->pulse, wait-stated byte-preload handshake, busy/done sync, and direct readback taps (ek/dk/ct/ss). CPU drives it via a register map (CTRL/STATUS/CONFIG, d/z/msg, preload + readback windows). Wire into soc_top on the existing axiOut_6 nets (no crossbar change), add mlkem_intr on intrpt[3]. Add sdk mlkem_test example running a KeyGen->Encaps->Decaps self-consistency round trip.
This commit is contained in:
404
rtl/ip/mlkem/sample_ntt/sample_ntt_sync.v
Normal file
404
rtl/ip/mlkem/sample_ntt/sample_ntt_sync.v
Normal file
@@ -0,0 +1,404 @@
|
||||
// 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 "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
|
||||
Reference in New Issue
Block a user