refactor(kg): split polymem into 3 banks {a, se, t} (async, stage 2a)
Replace the single async-read polymem[0:28*256-1] with 3 polynomial-indexed banks (bank_a A_hat / bank_se s_hat||e_hat / bank_t t_hat), addressed by abs_slot - base_slot. Still async-read here -- a pure refactor that validates bank sizing and base-relative addressing with zero timing change before stage 2b converts them to registered sd_bram + read-ahead pipelines. 11/11 KAT PASS, byte-exact, 0 file-not-found.
This commit is contained in:
47
.claude/plans/phase2_polymem_bram.md
Normal file
47
.claude/plans/phase2_polymem_bram.md
Normal file
@@ -0,0 +1,47 @@
|
||||
# Phase 2: polymem -> BRAM banks
|
||||
|
||||
## Goal
|
||||
Replace the single multi-port async-read `reg [11:0] polymem [0:28*256-1]` with
|
||||
registered-read `sd_bram` banks that infer real BRAM (ASIC: compiled SRAM).
|
||||
|
||||
## Bank split (all sd_bram, W=12, 1R+1W registered)
|
||||
- bank_a : A_hat[i][j], slots 0..K^2-1 -> D=4096 (KMAX^2*256), A=12
|
||||
- bank_se: s_hat[i] then e_hat[i] -> D=2048 (2*KMAX*256), A=11
|
||||
relative slot = abs_slot - slot_s_rt (s: 0..K-1, e: K..2K-1)
|
||||
- bank_t : t_hat[i], slots slot_t_rt.. -> D=1024 (KMAX*256), A=10
|
||||
relative slot = abs_slot - slot_t_rt
|
||||
|
||||
## Port budget (verified disjoint)
|
||||
- poly_mul LOAD vs ACCUMULATE never overlap; ST_N/ST_M/ST_E are disjoint
|
||||
top-states. => each bank needs only 1R+1W. sd_bram fits.
|
||||
- bank_t ST_M-acc reads t_hat[idx] (j>0 source) and writes t_hat[idx] result
|
||||
same cycle: 1R+1W, read-old/write-new. Prefetch keeps read AHEAD of write
|
||||
(monotonic idx), so no same-pass RAW hazard.
|
||||
|
||||
## Read sites (all need 1-cycle read-ahead vs async today)
|
||||
- ntt_in (ST_N load) <- bank_se[n_slot]
|
||||
- e_c0,e_c1 (ST_E) <- bank_se (dk) or bank_t (ek); SERIALIZE the
|
||||
two reads across the 3-byte/pair window
|
||||
- pm_a_in (ST_M load) <- bank_a[m_aslot]
|
||||
- pm_b_in (ST_M load) <- bank_se[m_j]
|
||||
- m_acc_src (ST_M acc) <- bank_se[K+m_i] (j==0) or bank_t[m_i] (j>0)
|
||||
- dbg_coeff_o (not on KAT path) <- route by slot range (compile-only)
|
||||
|
||||
## Write sites
|
||||
- ST_A: bank_a[a_slot] <= snt_coeff
|
||||
- ST_C: bank_se[c_slot-slot_s] <= cbd_modq
|
||||
- ST_N: bank_se[n_slot] <= ntt_coeff (write-back, same slot)
|
||||
- ST_M-acc: bank_t[m_i] <= m_accq
|
||||
|
||||
## Read-ahead strategy
|
||||
Cores hold ready_o high through entire LOAD (no mid-stream backpressure) =>
|
||||
fixed 1-cycle skew. Pattern: advance a read-address pointer 1 cycle ahead of
|
||||
the consumer index; delay the consumer's valid by 1 cycle ("prime the pipe").
|
||||
For ST_M-acc (irregular pm_vo cadence 0,1,1): 2-entry prefetch skid buffer,
|
||||
read pointer runs monotonically ahead of write pointer.
|
||||
|
||||
## Checkpoints (KAT-gated: 11 cases, byte-exact, 0 file-not-found)
|
||||
- 2a: split polymem -> 4 ASYNC-read banks (pure refactor, zero timing change).
|
||||
Validates bank sizing + base-relative addressing + debug mux. COMMIT.
|
||||
- 2b: convert banks to registered sd_bram + add read-ahead pipelines to every
|
||||
consumer FSM (ST_N, ST_E, ST_M load, ST_M acc). COMMIT.
|
||||
@@ -82,11 +82,36 @@ module mlkem_top #(
|
||||
localparam NUM_SLOTS = KMAX*KMAX + 3*KMAX;
|
||||
localparam SAW = 5; // slot-address width (>=clog2(28))
|
||||
|
||||
reg [11:0] polymem [0:NUM_SLOTS*256-1];
|
||||
// ================================================================
|
||||
// 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).
|
||||
// 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. -----
|
||||
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)
|
||||
reg [11:0] bank_a [0:(1<<PA_AW)-1];
|
||||
reg [11:0] bank_se [0:(1<<PSE_AW)-1];
|
||||
reg [11:0] bank_t [0:(1<<PT_AW)-1];
|
||||
|
||||
// Debug readback (registered for timing)
|
||||
reg [11:0] dbg_coeff_r;
|
||||
always @(posedge clk) dbg_coeff_r <= polymem[dbg_slot_i*256 + dbg_idx_i];
|
||||
// Route debug read by absolute slot range (not on the KAT correctness
|
||||
// path — TB reads ek/dk BRAMs only). A: 0..K^2-1; s/e: slot_s..slot_t-1;
|
||||
// t: slot_t.. Index within bank is (abs_slot-base)*256 + idx.
|
||||
wire [13:0] dbg_a_addr = dbg_slot_i*256 + dbg_idx_i;
|
||||
wire [13:0] dbg_se_addr = (dbg_slot_i - slot_s_rt)*256 + dbg_idx_i;
|
||||
wire [13:0] dbg_t_addr = (dbg_slot_i - slot_t_rt)*256 + dbg_idx_i;
|
||||
always @(posedge clk) begin
|
||||
if (dbg_slot_i >= slot_t_rt) dbg_coeff_r <= bank_t [dbg_t_addr[PT_AW-1:0]];
|
||||
else if (dbg_slot_i >= slot_s_rt) dbg_coeff_r <= bank_se[dbg_se_addr[PSE_AW-1:0]];
|
||||
else dbg_coeff_r <= bank_a [dbg_a_addr[PA_AW-1:0]];
|
||||
end
|
||||
assign dbg_coeff_o = dbg_coeff_r;
|
||||
|
||||
// ek and dk_pke byte memories sized for KMAX.
|
||||
@@ -387,7 +412,9 @@ module mlkem_top #(
|
||||
wire [11:0] ntt_coeff;
|
||||
wire ntt_vo;
|
||||
wire ntt_done;
|
||||
wire [11:0] ntt_in = polymem[n_slot_addr*256 + n_ridx[7:0]];
|
||||
// bank_se: NTT slot relative index = n_slot_addr - slot_s_rt == n_slot
|
||||
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]];
|
||||
|
||||
ntt_core u_ntt (
|
||||
.clk(clk), .rst_n(rst_n),
|
||||
@@ -423,9 +450,12 @@ module mlkem_top #(
|
||||
wire e_is_dk = (e_poly >= {1'b0, k_r});
|
||||
wire [4:0] e_pidx = e_is_dk ? (e_poly - {1'b0, k_r}) : e_poly; // index within target
|
||||
wire [SAW-1:0] e_slot = e_is_dk ? (slot_s_rt + e_pidx) : (slot_t_rt + e_pidx);
|
||||
// two coeffs of the current pair
|
||||
wire [11:0] e_c0 = polymem[e_slot*256 + {e_pair, 1'b0}];
|
||||
wire [11:0] e_c1 = polymem[e_slot*256 + {e_pair, 1'b1}];
|
||||
// two coeffs of the current pair: ek half reads t_hat (bank_t), dk half
|
||||
// reads s_hat (bank_se). Relative index within the bank = e_pidx.
|
||||
wire [13:0] e_rd0_full = e_pidx*256 + {e_pair, 1'b0};
|
||||
wire [13:0] e_rd1_full = e_pidx*256 + {e_pair, 1'b1};
|
||||
wire [11:0] e_c0 = e_is_dk ? bank_se[e_rd0_full[PSE_AW-1:0]] : bank_t[e_rd0_full[PT_AW-1:0]];
|
||||
wire [11:0] e_c1 = e_is_dk ? bank_se[e_rd1_full[PSE_AW-1:0]] : bank_t[e_rd1_full[PT_AW-1:0]];
|
||||
// 3 packed bytes
|
||||
wire [7:0] e_b0 = e_c0[7:0];
|
||||
wire [7:0] e_b1 = {e_c1[3:0], e_c0[11:8]};
|
||||
@@ -445,8 +475,12 @@ module mlkem_top #(
|
||||
wire pm_ready;
|
||||
wire [11:0] pm_coeff;
|
||||
wire pm_vo;
|
||||
wire [11:0] pm_a_in = polymem[m_aslot*256 + m_ld[7:0]];
|
||||
wire [11:0] pm_b_in = polymem[m_sslot*256 + m_ld[7:0]];
|
||||
// pm_a: A_hat[i][j] in bank_a (abs slot m_aslot). pm_b: s_hat[j] in
|
||||
// bank_se (relative slot = m_sslot - slot_s_rt = m_j).
|
||||
wire [13:0] pm_a_full = m_aslot*256 + m_ld[7:0];
|
||||
wire [13:0] pm_b_full = m_j*256 + m_ld[7:0];
|
||||
wire [11:0] pm_a_in = bank_a [pm_a_full[PA_AW-1:0]];
|
||||
wire [11:0] pm_b_in = bank_se[pm_b_full[PSE_AW-1:0]];
|
||||
|
||||
poly_mul_sync u_pmul (
|
||||
.clk(clk), .rst_n(rst_n),
|
||||
@@ -459,9 +493,13 @@ module mlkem_top #(
|
||||
.ready_i(1'b1)
|
||||
);
|
||||
|
||||
// accumulator source: e_hat[i] for first term (j==0), else running t_hat[i]
|
||||
wire [11:0] m_acc_src = (m_j == 2'd0) ? polymem[m_eslot*256 + m_oidx]
|
||||
: polymem[m_tslot*256 + m_oidx];
|
||||
// accumulator source: e_hat[i] for first term (j==0), else running t_hat[i].
|
||||
// e_hat[i] lives in bank_se at relative slot (slot_e_rt-slot_s_rt + m_i) = K+m_i.
|
||||
// t_hat[i] lives in bank_t at relative slot m_i.
|
||||
wire [13:0] m_eacc_full = ({2'b0, k_r} + {2'b0, m_i})*256 + m_oidx; // K+m_i
|
||||
wire [13:0] m_tacc_full = m_i*256 + m_oidx;
|
||||
wire [11:0] m_acc_src = (m_j == 2'd0) ? bank_se[m_eacc_full[PSE_AW-1:0]]
|
||||
: bank_t [m_tacc_full[PT_AW-1:0]];
|
||||
// (a + b) mod Q (both < Q, sum < 2Q): one conditional subtract
|
||||
wire [12:0] m_sum = {1'b0, m_acc_src} + {1'b0, pm_coeff};
|
||||
wire [11:0] m_accq = (m_sum >= 13'(Q)) ? (m_sum - 13'(Q)) : m_sum[11:0];
|
||||
@@ -578,7 +616,7 @@ module mlkem_top #(
|
||||
|
||||
// store each output coefficient only while busy (ignore stale last coeff from prior poly)
|
||||
if (a_busy && snt_vo && snt_ack) begin
|
||||
polymem[a_slot*256 + a_widx] <= snt_coeff;
|
||||
bank_a[(a_slot*256 + a_widx) & ((1<<PA_AW)-1)] <= snt_coeff;
|
||||
if (snt_last) begin
|
||||
// finished this poly; advance (i,j) in row-major order
|
||||
a_pair <= a_pair + 5'd1;
|
||||
@@ -615,7 +653,8 @@ module mlkem_top #(
|
||||
end
|
||||
|
||||
if (c_busy && cbd_vo && cbd_ack) begin
|
||||
polymem[c_slot*256 + c_widx] <= cbd_modq;
|
||||
// bank_se relative slot = c_slot - slot_s_rt = c_poly
|
||||
bank_se[(c_poly*256 + c_widx) & ((1<<PSE_AW)-1)] <= cbd_modq;
|
||||
if (cbd_last) begin
|
||||
c_poly <= c_poly + 3'd1;
|
||||
c_widx <= 8'd0;
|
||||
@@ -650,7 +689,8 @@ module mlkem_top #(
|
||||
|
||||
// OUTPUT phase: collect 256 results, write back to same slot
|
||||
if (ntt_vo) begin
|
||||
polymem[n_slot_addr*256 + n_widx] <= ntt_coeff;
|
||||
// bank_se relative slot = n_slot_addr - slot_s_rt = n_slot
|
||||
bank_se[(n_slot*256 + n_widx) & ((1<<PSE_AW)-1)] <= ntt_coeff;
|
||||
n_widx <= n_widx + 8'd1; // wraps 255->0 after last
|
||||
end
|
||||
|
||||
@@ -700,7 +740,8 @@ module mlkem_top #(
|
||||
|
||||
// ACCUMULATE: each product coeff += e_hat (j==0) or running t_hat
|
||||
if (pm_vo) begin
|
||||
polymem[m_tslot*256 + m_oidx] <= m_accq;
|
||||
// bank_t relative slot = m_tslot - slot_t_rt = m_i
|
||||
bank_t[(m_i*256 + m_oidx) & ((1<<PT_AW)-1)] <= m_accq;
|
||||
if (m_oidx == 8'd255) begin
|
||||
// finished this (i,j) term; advance
|
||||
if (m_j + 2'd1 < k_r) begin
|
||||
|
||||
Reference in New Issue
Block a user