refactor(kg): registered read-ahead for ST_N (bank_se NTT load)

First per-consumer step of stage 2b. n_ridx becomes a read-ahead pointer
leading the consume index by 1; bank_se read is registered into n_rd_data and
fed to ntt_core one cycle later (n_valid delayed 1 cyc). Cores hold ready_o
high through LOAD so a fixed 1-cycle skew suffices. Matches sd_bram registered
timing for this read port. 11/11 KAT PASS, byte-exact.
This commit is contained in:
2026-06-28 16:17:30 +08:00
parent 4d3adc6b57
commit 45e07c28e8

View File

@@ -83,15 +83,19 @@ module mlkem_top #(
localparam SAW = 5; // slot-address width (>=clog2(28))
// ================================================================
// Coefficient storage as 3 BRAM banks (Phase 2). Was one async-read
// reg array polymem[0:28*256-1]; now split by polynomial so each bank
// is a registered-read sd_bram (infers BRAM / ASIC compiled SRAM).
// Coefficient storage as 3 banks (Phase 2). Was one async-read reg
// array polymem[0:28*256-1]; split by polynomial so each becomes a
// registered-read sd_bram (infers BRAM / ASIC SRAM) in the final step.
// bank_a : A_hat[i][j] slots 0..K^2-1 -> D=KMAX*KMAX*256=4096
// bank_se: s_hat||e_hat slots slot_s.. -> D=2*KMAX*256=2048 (rel 0..2K-1)
// bank_t : t_hat[i] slots slot_t.. -> D=KMAX*256=1024 (rel 0..K-1)
// Addresses are base-relative: bank index = abs_slot - base_slot.
// ----- Stage 2a: async behavioural banks (timing-identical refactor).
// Stage 2b converts these to registered sd_bram + read-ahead. -----
// Stage 2b (in progress): per-consumer registered read-ahead. Banks are
// still async reg arrays here; converted consumers read via their own
// 1-cycle pipeline reg (== sd_bram timing). Once every consumer of a
// bank is registered, the array is replaced by an sd_bram with its read
// port muxed across the (phase-disjoint) consumers.
// ================================================================
localparam PA_AW = 12; // bank_a addr width (4096)
localparam PSE_AW = 11; // bank_se addr width (2048)
localparam PT_AW = 10; // bank_t addr width (1024)
@@ -402,19 +406,22 @@ module mlkem_top #(
// ---- ntt_core: forward NTT (mode=0, no scaling) of s[i],e[i] in place ----
// N-stage bookkeeping: process slots S0,S1,E0,E1 (= SLOT_S0 + n_slot).
reg [4:0] n_slot; // 0..2K (process s_hat[0..K-1] then e_hat[0..K-1])
reg [8:0] n_ridx; // load read index 0..256
reg [8:0] n_ridx; // load read-AHEAD pointer 0..256 (leads consume by 1)
reg [7:0] n_widx; // output write index 0..255
reg n_valid; // feeding coeffs to ntt_core
reg n_valid; // feeding coeffs to ntt_core (delayed 1 cyc vs n_ridx)
reg n_loading; // 1 while presenting load addresses to bank_se
reg n_pending; // waiting for ntt_core IDLE to start next slot
reg [11:0] n_rd_data; // registered bank_se read (== sd_bram timing)
wire [SAW-1:0] n_slot_addr = slot_s_rt + n_slot; // s_hat then e_hat contiguous
wire ntt_ready;
wire [11:0] ntt_coeff;
wire ntt_vo;
wire ntt_done;
// bank_se: NTT slot relative index = n_slot_addr - slot_s_rt == n_slot
// bank_se read addr for the NTT load (relative slot = n_slot); registered
// into n_rd_data, which feeds ntt_core 1 cycle later.
wire [13:0] ntt_rd_full = n_slot*256 + n_ridx[7:0];
wire [11:0] ntt_in = bank_se[ntt_rd_full[PSE_AW-1:0]];
wire [11:0] ntt_in = n_rd_data;
ntt_core u_ntt (
.clk(clk), .rst_n(rst_n),
@@ -544,6 +551,8 @@ module mlkem_top #(
n_ridx <= 9'd0;
n_widx <= 8'd0;
n_valid <= 1'b0;
n_loading <= 1'b0;
n_rd_data <= 12'd0;
n_pending <= 1'b0;
m_i <= 2'd0;
m_j <= 2'd0;
@@ -666,24 +675,33 @@ module mlkem_top #(
end
end
// Arm N stage when C finishes: start NTT on slot S0
// Arm N stage when C finishes: prime load of slot S0. n_ridx is a
// read-ahead pointer; bank_se read is registered into n_rd_data and
// fed to ntt_core one cycle later, so valid starts low (priming).
if (st == ST_C && st_next == ST_N) begin
n_slot <= 3'd0;
n_ridx <= 9'd0;
n_widx <= 8'd0;
n_valid <= 1'b1; // begin loading first poly
n_valid <= 1'b0;
n_loading <= 1'b1; // begin presenting load addresses
n_pending <= 1'b0;
end
// ---- ST_N: forward NTT each of S0,S1,E0,E1 in place ----
if (st == ST_N) begin
// LOAD phase: stream 256 coeffs into ntt_core
if (n_valid && ntt_ready) begin
if (n_ridx == 9'd255) begin
n_valid <= 1'b0; // last coeff presented this cycle
n_ridx <= 9'd0;
// LOAD phase: present read-ahead addr to bank_se; the value
// registered last cycle (n_rd_data) is consumed by ntt_core
// this cycle (n_valid). Cores hold ready high through LOAD, so
// a fixed 1-cycle skew suffices (no backpressure gating).
if (n_loading) begin
if (n_ridx == 9'd256) begin
// 256th coeff (bank_se[255]) consumed this cycle; stop
n_loading <= 1'b0;
n_valid <= 1'b0;
end else begin
n_ridx <= n_ridx + 9'd1;
n_rd_data <= bank_se[ntt_rd_full[PSE_AW-1:0]];
n_ridx <= n_ridx + 9'd1;
n_valid <= 1'b1; // data presented last cycle is valid
end
end
@@ -705,10 +723,11 @@ module mlkem_top #(
end
end
// Kick next slot's load once core is back IDLE
// Kick next slot's load once core is back IDLE (re-prime)
if (n_pending && ntt_ready && !ntt_done) begin
n_valid <= 1'b1;
n_ridx <= 9'd0;
n_valid <= 1'b0;
n_loading <= 1'b1;
n_pending <= 1'b0;
end
end